diff --git a/flashinfer/experimental/b12x/comm/pcie/README.md b/flashinfer/experimental/b12x/comm/pcie/README.md index 40e6125a136..fe3ceafad0a 100644 --- a/flashinfer/experimental/b12x/comm/pcie/README.md +++ b/flashinfer/experimental/b12x/comm/pcie/README.md @@ -159,6 +159,59 @@ work before entering this method with identical arguments. A profile change cannot occur during capture or overlap graph replay. No topology is inferred from the profile; callers choose the prepared profile explicitly. +### Cluster-scoped phased profiles + +Experimental cluster profiles opt in through +`AllReduce.plan(..., joint_only=True, joint_cluster_size=8)`. The default is +`joint_cluster_size=1`, which preserves ordinary queries and launch policy. +Cluster plans require SM120/SM121, at least 64 SMs, 512 threads per CTA, and a +fixed 64-CTA grid. Preparation checks that the compiled kernel can keep this +grid resident with eight CTAs per cluster; unsupported geometry is rejected. + +`joint_latency` additionally admits `full_phased` and `payload_rsag_phased`; +`joint_bandwidth` additionally admits `payload_rsag_phased`. These modes require +exactly eight active CTAs: one complete cluster. All other clusters return +before accessing IPC state. Each writer issues a system fence before the +active cluster synchronizes and advances to its next peer destination. +Ordinary modes retain active counts 1–64 and execute without cluster barriers. +They still use the cluster launch geometry when selected on a cluster plan; +this geometry can affect scheduling and does not guarantee zero overhead. +Both paths use the same workspace layout, BF16 arithmetic, and aliasing contract. + +For an established joint channel, declare the geometry with the existing +configuration type and submit the plan through the normal preparation batch: + +```python +plan = runtime.plan( + inp, + out=out, + joint_only=True, + joint_cluster_size=8, + override=CollectiveConfig( + implementation="joint_bandwidth", + threads=512, + blocks=64, + submode="payload_rsag_phased", + active_blocks=8, + ), +) +``` + +After preparation, a startup specification can retain ordinary alternatives +alongside the phased profile: + +```python +spec = joint_startup_spec( + plan, + profiles=(("payload_rsag_phased", 8), ("bulk_pull", 11), ("payload_rsag", 42)), +) +``` + +The cluster dimension is immutable plan metadata, while the selected mode and +active count remain device data. Startup selection and serving use the same +compiled launcher. Cluster launches use ordinary CUDA graph kernel attributes, +without cooperative launch or additional protocol counters. + ### Captured startup mode selection `AllReduce.plan(...)` excludes joint kernels by default and retains static transport diff --git a/flashinfer/experimental/b12x/comm/pcie/_collective_tuning.py b/flashinfer/experimental/b12x/comm/pcie/_collective_tuning.py index bfeb46ad1b5..2c3019ea9c1 100644 --- a/flashinfer/experimental/b12x/comm/pcie/_collective_tuning.py +++ b/flashinfer/experimental/b12x/comm/pcie/_collective_tuning.py @@ -14,7 +14,7 @@ from b12x.preparation._efficiency import capture_autotune_effort from ._tuning import PcieQuery, TUNING as NATIVE -from ._joint_profile import JOINT_FAMILIES +from ._joint_profile import JOINT_FAMILIES, profile_records @dataclass(frozen=True, kw_only=True) @@ -42,6 +42,23 @@ def _configs(query): blocks = int(call.get("blocks", call.get("block_limit", 16))) threads = int(call["threads"]) if name in JOINT_FAMILIES: + cluster_size = call.get("joint_cluster_size", 1) + if cluster_size != 1: + from ._joint_profile import JOINT_MODES, validate_cluster_geometry + + validate_cluster_geometry(blocks, threads, cluster_size) + modes = {value: mode for mode, value in JOINT_MODES.items()} + for mode, active in profile_records( + name, blocks, cluster_size=cluster_size + ): + yield CollectiveConfig( + implementation=name, + threads=threads, + blocks=blocks, + submode=modes[mode], + active_blocks=active, + ) + continue for count_threads in (32, 64, 128, 256, 512): for mode in JOINT_FAMILIES[name]: units = int(call["size_packs"]) @@ -319,6 +336,8 @@ def _family_bucket(query): if field in setup }, ) + if "joint_cluster_size" in call: + alternatives[name]["joint_cluster_size"] = call["joint_cluster_size"] if len(sizes) != 1: raise ValueError("collective alternatives disagree on payload size") size = sizes.pop() diff --git a/flashinfer/experimental/b12x/comm/pcie/_joint_cluster_cute.py b/flashinfer/experimental/b12x/comm/pcie/_joint_cluster_cute.py new file mode 100644 index 00000000000..398282d302f --- /dev/null +++ b/flashinfer/experimental/b12x/comm/pcie/_joint_cluster_cute.py @@ -0,0 +1,382 @@ +"""Cluster-scoped phased transfers alongside the ordinary joint protocols. + +A fixed 64-CTA grid consists of eight portable clusters. Phased profiles admit +only cluster zero; ordinary profiles use their selected CTA count without a +cluster barrier. The signal and payload regions retain the ordinary layout. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64, Uint32 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import dsl_user_op + +from b12x._lib.compiler import KernelCompileSpec, compile as b12x_compile +from b12x._lib.compile_plan import attach_programs +from b12x._lib.program_cache import program_cache +from b12x._lib.intrinsics.memory import ld_global_v4_u32, st_global_v4_u32 +from b12x._lib.runtime_control import raise_if_kernel_resolution_frozen +from b12x._lib.utils.device import current_cuda_stream +from b12x._lib.utils.pointers import make_ptr +from ._cute_intrinsics import ( + graph_epoch_arrive_serialized, + ld_relaxed_gpu_u32, + ld_relaxed_sys_global_v4_u32, + pack_f32x2_to_bf16x2, + st_relaxed_sys_global_v4_u32, +) +from ._joint_cute import _JointLaunch, _EPOCH_OFFSET +from ._joint_profile import validate_cluster_geometry +from ._oneshot_cute import _dummy + + +@dsl_user_op +def _threadfence_sys(*, loc=None, ip=None): + llvm.inline_asm( + None, + [], + "fence.sc.sys;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +class _JointClusterLaunch(_JointLaunch): + @cute.jit + def __call__( + self, + peers: cute.Pointer, + signals: cute.Pointer, + local_signal: cute.Pointer, + selector: cute.Pointer, + payload: cute.Pointer, + output: cute.Pointer, + size: Int64, + capacity: Int64, + stream: cuda.CUstream, + ): + self.kernel( + peers, signals, local_signal, selector, payload, output, size, capacity + ).launch(grid=(64, 1, 1), block=(512, 1, 1), cluster=(8, 1, 1), stream=stream) + + @cute.jit + def _finish_write_phase(self): + # Every thread in the active cluster participates, including empty tails. + _threadfence_sys() + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + @cute.jit + def _ready_phased( + self, + peers: cute.Pointer, + payload: Int64, + output: Int64, + size: Int64, + capacity: Int64, + flat: Int64, + stride: Int64, + mode: Uint32, + ): + local = Int64(peers[self._rank]) + capacity * Int64(16) + # The final peer owns the remainder, including every pack when part is zero. + part = size // Int64(self._world_size) + for scatter_phase in cutlass.range_constexpr(self._world_size): + target = (self._rank + scatter_phase) % self._world_size + begin, end = Int64(0), size + if mode == Uint32(1): + begin = Int64(target) * part + end = begin + part + if cutlass.const_expr(target == self._world_size - 1): + end = size + index = begin + flat + while index < end: + words = ld_global_v4_u32(payload + index * Int64(16)) + dest = Int64(peers[target]) + ( + capacity + Int64(self._rank) * capacity + index + ) * Int64(16) + self._publish_bf16_words(dest, words) + index += stride + if cutlass.const_expr(scatter_phase != 0): + self._finish_write_phase() + begin, end = Int64(0), size + if mode == Uint32(1): + begin = Int64(self._rank) * part + end = begin + part + if cutlass.const_expr(self._rank == self._world_size - 1): + end = size + index = begin + flat + while index < end: + acc = cute.make_rmem_tensor((8,), Float32) + self._receive_bf16_sum(local, capacity, index, acc) + result = cute.make_rmem_tensor((4,), Uint32) + for word in cutlass.range_constexpr(4): + result[word] = pack_f32x2_to_bf16x2(acc[word * 2], acc[word * 2 + 1]) + st_global_v4_u32( + output + index * Int64(16), + result[0], + result[1], + result[2], + result[3], + ) + index += stride + if mode == Uint32(1): + for gather_phase in cutlass.range_constexpr(1, self._world_size): + target = (self._rank + gather_phase) % self._world_size + # Match the reduction's thread map so each thread reads its own + # output packs. Empty intervals still enter every phase barrier. + index = begin + flat + while index < end: + words = ld_global_v4_u32(output + index * Int64(16)) + dest = Int64(peers[target]) + ( + capacity + Int64(self._rank) * capacity + index + ) * Int64(16) + self._publish_bf16_words(dest, words) + index += stride + self._finish_write_phase() + for owner in cutlass.range_constexpr(self._world_size): + if cutlass.const_expr(owner != self._rank): + remote_begin = Int64(owner) * part + remote_end = remote_begin + part + if cutlass.const_expr(owner == self._world_size - 1): + remote_end = size + index = remote_begin + flat + while index < remote_end: + address = local + (Int64(owner) * capacity + index) * Int64(16) + gathered = cute.make_rmem_tensor((4,), Uint32) + waiting = cutlass.Boolean(True) + while waiting: + words = ld_relaxed_sys_global_v4_u32(address) + waiting = cutlass.Boolean(False) + for word in cutlass.range_constexpr(4): + gathered[word] = words[word] + waiting = ( + waiting + | ((words[word] & Uint32(0xFFFF)) == Uint32(0)) + | ((words[word] & Uint32(0xFFFF0000)) == Uint32(0)) + ) + st_global_v4_u32( + output + index * Int64(16), + gathered[0], + gathered[1], + gathered[2], + gathered[3], + ) + st_relaxed_sys_global_v4_u32( + address, Uint32(0), Uint32(0), Uint32(0), Uint32(0) + ) + index += stride + + @cute.kernel + def kernel( + self, + peers: cute.Pointer, + signals: cute.Pointer, + local_signal: cute.Pointer, + selector: cute.Pointer, + payload: cute.Pointer, + output: cute.Pointer, + size: Int64, + capacity: Int64, + ): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + profile = selector[0] + mode = profile >> Uint32(16) + active = profile & Uint32(0xFFFF) + if (mode == Uint32(5)) | (mode == Uint32(6)): + # A partial cluster cannot enter a cluster barrier, even if a caller + # corrupts the selector after host-side profile validation. + if active != Uint32(8): + active = Uint32(0) + if cutlass.const_expr(self._family == "bandwidth"): + if mode == Uint32(5): + active = Uint32(0) + if Uint32(bid) < active: + # Inactive clusters return before any IPC state read or epoch arrival. + self_signal = Int64(local_signal.toint()) + generation = ld_relaxed_gpu_u32(self_signal + Int64(_EPOCH_OFFSET)) + slot_peers = peers + Int64(generation % Uint32(2)) * Int64(16) + cute.arch.sync_threads() + if tid == Int32(self._threads - 1): + graph_epoch_arrive_serialized( + self_signal + Int64(_EPOCH_OFFSET), + self_signal + Int64(_EPOCH_OFFSET + 4), + active, + ) + flat = Int64(bid) * Int64(self._threads) + Int64(tid) + stride = Int64(active) * Int64(self._threads) + src, dst = Int64(payload.toint()), Int64(output.toint()) + if mode == Uint32(6): + self._ready_phased( + slot_peers, + src, + dst, + size, + capacity, + flat, + stride, + Uint32(1), + ) + elif mode == Uint32(4): + self._pull(slot_peers, signals, src, dst, size, flat, stride) + elif mode == Uint32(1): + self._ready(slot_peers, src, dst, size, capacity, flat, stride, mode) + else: + if cutlass.const_expr(self._family == "latency"): + if mode == Uint32(5): + self._ready_phased( + slot_peers, + src, + dst, + size, + capacity, + flat, + stride, + Uint32(0), + ) + else: + self._ready( + slot_peers, src, dst, size, capacity, flat, stride, mode + ) + else: + if mode == Uint32(2): + self._bulk( + slot_peers, + signals, + src, + dst, + size, + capacity, + flat, + stride, + True, + ) + else: + self._bulk( + slot_peers, + signals, + src, + dst, + size, + capacity, + flat, + stride, + False, + ) + + +@program_cache +def get_joint_cluster_launcher( + family, world_size, rank, threads, max_blocks, cluster_size, device_index +): + validate_cluster_geometry(max_blocks, threads, cluster_size) + if cluster_size != 8 or family not in ("latency", "bandwidth"): + raise ValueError( + "clustered joint launch requires cluster eight and a joint family" + ) + if world_size not in (2, 3, 4, 6, 8, 10) or not 0 <= rank < world_size: + raise ValueError("unsupported joint collective rank group") + key = (family, world_size, rank, threads, max_blocks, cluster_size, device_index) + launch = _JointClusterLaunch(family, world_size, rank, threads, max_blocks) + raise_if_kernel_resolution_frozen("cute.compile", target=launch, cache_key=key) + raw = b12x_compile( + launch, + _dummy(cutlass.Uint64, 8), + _dummy(cutlass.Uint64, 8), + _dummy(Uint32, 4), + _dummy(Uint32, 4), + _dummy(Uint32, 16), + _dummy(Uint32, 16), + 1, + 1, + current_cuda_stream(), + compile_spec=KernelCompileSpec.from_key("comm.pcie.joint_cluster", 2, key[:-1]), + ) + + def run( + peer_table_address, + signal_table_address, + local_signal_address, + selector_address, + input_address, + output_address, + size_packs, + capacity_packs, + ): + if not 0 < size_packs <= capacity_packs: + raise ValueError("joint payload must fit its prepared capacity") + raw( + make_ptr( + cutlass.Uint64, + peer_table_address, + cute.AddressSpace.gmem, + assumed_align=8, + ), + make_ptr( + cutlass.Uint64, + signal_table_address, + cute.AddressSpace.gmem, + assumed_align=8, + ), + make_ptr( + Uint32, local_signal_address, cute.AddressSpace.gmem, assumed_align=4 + ), + make_ptr(Uint32, selector_address, cute.AddressSpace.gmem, assumed_align=4), + make_ptr(Uint32, input_address, cute.AddressSpace.gmem, assumed_align=16), + make_ptr(Uint32, output_address, cute.AddressSpace.gmem, assumed_align=16), + int(size_packs), + int(capacity_packs), + current_cuda_stream(), + ) + + return attach_programs(run, raw) + + +def validate_cluster_residency(launcher, blocks, threads, cluster_size): + """Reject unsupported compiled geometry without changing the captured grid.""" + from b12x.preparation.types import CandidateResourceIneligibleError + + (raw,) = launcher.__b12x_dependencies__ + executable = raw._executable + if len(executable.kernel_info) != 1: + raise RuntimeError("clustered joint launcher must retain one CUDA kernel") + (name,) = executable.kernel_info + status, kernel = cuda.cuLibraryGetKernel( + cuda.CUlibrary(int(executable.library)), name.encode() + ) + if status != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError("cannot resolve clustered joint kernel") + status, function = cuda.cuKernelGetFunction(kernel) + if status != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError("cannot query clustered joint kernel") + attribute = cuda.CUlaunchAttribute() + attribute.id = cuda.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION + attribute.value.clusterDim.x = cluster_size + attribute.value.clusterDim.y = attribute.value.clusterDim.z = 1 + config = cuda.CUlaunchConfig() + config.gridDimX, config.gridDimY, config.gridDimZ = blocks, 1, 1 + config.blockDimX, config.blockDimY, config.blockDimZ = threads, 1, 1 + config.sharedMemBytes = 0 + config.attrs, config.numAttrs = [attribute], 1 + status, potential = cuda.cuOccupancyMaxPotentialClusterSize(function, config) + if status != cuda.CUresult.CUDA_SUCCESS: + raise CandidateResourceIneligibleError( + "cannot support the prepared joint cluster geometry" + ) + status, resident = cuda.cuOccupancyMaxActiveClusters(function, config) + if ( + status != cuda.CUresult.CUDA_SUCCESS + or potential < cluster_size + or resident * cluster_size < blocks + ): + raise CandidateResourceIneligibleError( + "joint cluster grid exceeds compiled kernel residency" + ) diff --git a/flashinfer/experimental/b12x/comm/pcie/_joint_preparation.py b/flashinfer/experimental/b12x/comm/pcie/_joint_preparation.py index bf1e38d1bdc..6ab0f52211a 100644 --- a/flashinfer/experimental/b12x/comm/pcie/_joint_preparation.py +++ b/flashinfer/experimental/b12x/comm/pcie/_joint_preparation.py @@ -114,6 +114,17 @@ def compile_joint(family, world, rank, threads, max_blocks, ordinal): return get_joint_launcher(family, world, rank, threads, max_blocks, ordinal) +@program_cache(scope="preparation") +def compile_joint_cluster( + family, world, rank, threads, max_blocks, cluster_size, ordinal +): + from ._joint_cluster_cute import get_joint_cluster_launcher + + return get_joint_cluster_launcher( + family, world, rank, threads, max_blocks, cluster_size, ordinal + ) + + class _JointExecutionState(_OneshotExecutionState): def __init__(self, query, runtime, launcher, workspace): self.query, self.runtime, self.launcher = query, runtime, launcher @@ -121,9 +132,16 @@ def __init__(self, query, runtime, launcher, workspace): self._bound_input = None self.workspace = workspace call = query.call - self.profiles = profile_records(call["transport"], call["blocks"]) + self.cluster_size = call.get("joint_cluster_size", 1) + self.profiles = profile_records( + call["transport"], call["blocks"], cluster_size=self.cluster_size + ) selected = selected_profile( - call["transport"], call["blocks"], call["submode"], call["active_blocks"] + call["transport"], + call["blocks"], + call["submode"], + call["active_blocks"], + cluster_size=self.cluster_size, ) self.selector = torch.tensor( [selected], dtype=torch.uint32, device=runtime.device @@ -151,7 +169,11 @@ def select_profile(self, submode, active_blocks): owner="joint PCIe profile validation", exchange_group=self.runtime.exchange_group, setup=lambda: selected_profile( - call["transport"], call["blocks"], submode, active_blocks + call["transport"], + call["blocks"], + submode, + active_blocks, + cluster_size=self.cluster_size, ), ) torch.cuda.synchronize(self.runtime.device) @@ -164,6 +186,9 @@ def select_profile(self, submode, active_blocks): call["blocks"], submode, active_blocks, + ) + + ( + () if self.cluster_size == 1 else (("cluster_size", self.cluster_size),) ), ) self.selector.fill_(selected) @@ -193,8 +218,16 @@ def run_plain(self, inp, out): def plan(query, *, runtime, override=None): call = query.call + cluster_size = call.get("joint_cluster_size", 1) + from ._joint_profile import validate_cluster_geometry + + validate_cluster_geometry(call["blocks"], call["threads"], cluster_size) profile_index( - call["transport"], call["blocks"], call["submode"], call["active_blocks"] + call["transport"], + call["blocks"], + call["submode"], + call["active_blocks"], + cluster_size=cluster_size, ) if call["dtype"] != "bfloat16" or not call["stage_input"]: raise ValueError("joint PCIe all-reduce requires staged BF16 input") @@ -210,7 +243,20 @@ def plan(query, *, runtime, override=None): profiles_bytes = 4 def materialize(selection, device): - launcher = load_programs(compile_joint(*args, device.ordinal)) + if cluster_size == 1: + launcher = load_programs(compile_joint(*args, device.ordinal)) + else: + from ._joint_cluster_cute import validate_cluster_residency + + validate_cluster_geometry( + call["blocks"], call["threads"], cluster_size, device.identity + ) + launcher = load_programs( + compile_joint_cluster(*args, cluster_size, device.ordinal) + ) + validate_cluster_residency( + launcher, call["blocks"], call["threads"], cluster_size + ) return _JointExecutionState(query, runtime, launcher, _workspace(runtime)) return Plan( @@ -219,8 +265,11 @@ def materialize(selection, device): override=override, _compile_jobs=lambda config, device: ( CompileJob.create( - "b12x.comm.pcie._joint_preparation:compile_joint", + "b12x.comm.pcie._joint_preparation:compile_joint" + if cluster_size == 1 + else "b12x.comm.pcie._joint_preparation:compile_joint_cluster", *args, + *((cluster_size,) if cluster_size != 1 else ()), CompileJob.DEVICE_ORDINAL, ), ), diff --git a/flashinfer/experimental/b12x/comm/pcie/_joint_profile.py b/flashinfer/experimental/b12x/comm/pcie/_joint_profile.py index 2c1d015d089..dc227d22cda 100644 --- a/flashinfer/experimental/b12x/comm/pcie/_joint_profile.py +++ b/flashinfer/experimental/b12x/comm/pcie/_joint_profile.py @@ -1,13 +1,50 @@ """Prepared profile geometry for joint BF16 PCIe all-reduce kernels.""" -JOINT_MODES = {"full": 0, "payload_rsag": 1, "bulk_push": 2, "bulk_pull": 3, "pull": 4} +JOINT_MODES = { + "full": 0, + "payload_rsag": 1, + "bulk_push": 2, + "bulk_pull": 3, + "pull": 4, + "full_phased": 5, + "payload_rsag_phased": 6, +} JOINT_FAMILIES = { "joint_latency": ("full", "payload_rsag", "pull"), "joint_bandwidth": ("payload_rsag", "bulk_push", "bulk_pull", "pull"), } +JOINT_CLUSTER_PHASES = { + "joint_latency": ("full_phased", "payload_rsag_phased"), + "joint_bandwidth": ("payload_rsag_phased",), +} + + +def validate_cluster_geometry(max_blocks, threads, cluster_size, device=None): + """Admit one portable active cluster inside the fixed joint launch grid.""" + if type(cluster_size) is not int or cluster_size not in (1, 8): + raise ValueError("joint cluster size must be 1 or 8") + if cluster_size == 1: + return + if max_blocks != 64 or threads != 512: + raise ValueError("clustered joint kernels require 64 CTAs and 512 threads") + if device is not None and ( + device.compute_capability not in ((12, 0), (12, 1)) or device.sm_count < 64 + ): + raise ValueError( + "clustered joint kernels require SM120/SM121 with at least 64 SMs" + ) -def profile_records(implementation, max_blocks): +def joint_modes(implementation, *, cluster_size=1): + if type(cluster_size) is not int or cluster_size not in (1, 8): + raise ValueError("joint cluster size must be 1 or 8") + modes = JOINT_FAMILIES.get(implementation, ()) + return modes + ( + JOINT_CLUSTER_PHASES.get(implementation, ()) if cluster_size == 8 else () + ) + + +def profile_records(implementation, max_blocks, *, cluster_size=1): if ( implementation not in JOINT_FAMILIES or type(max_blocks) is not int @@ -16,27 +53,46 @@ def profile_records(implementation, max_blocks): raise ValueError( "joint profiles require a latency/bandwidth family and 1–64 CTAs" ) + modes = joint_modes(implementation, cluster_size=cluster_size) + if cluster_size == 8 and max_blocks != 64: + raise ValueError("clustered joint profiles require a 64-CTA physical grid") return tuple( (JOINT_MODES[mode], blocks) - for mode in JOINT_FAMILIES[implementation] - for blocks in range(1, max_blocks + 1) + for mode in modes + for blocks in ( + (8,) + if mode in JOINT_CLUSTER_PHASES[implementation] + else range(1, max_blocks + 1) + ) ) -def profile_index(implementation, max_blocks, submode, active_blocks): +def profile_index( + implementation, max_blocks, submode, active_blocks, *, cluster_size=1 +): if type(active_blocks) is not int or not 1 <= active_blocks <= max_blocks: raise ValueError("active CTA count must fit the prepared maximum grid") - if submode not in JOINT_FAMILIES.get(implementation, ()): + if submode not in joint_modes(implementation, cluster_size=cluster_size): raise ValueError("submode does not belong to the prepared joint family") - return profile_records(implementation, max_blocks).index( + if submode in JOINT_CLUSTER_PHASES.get(implementation, ()) and active_blocks != 8: + raise ValueError( + "phased joint profiles require exactly one active eight-CTA cluster" + ) + return profile_records(implementation, max_blocks, cluster_size=cluster_size).index( (JOINT_MODES[submode], active_blocks) ) -def selected_profile(implementation, max_blocks, submode, active_blocks): +def selected_profile( + implementation, max_blocks, submode, active_blocks, *, cluster_size=1 +): """Encode an admitted profile for one device load without changing its identity.""" - index = profile_index(implementation, max_blocks, submode, active_blocks) - mode, blocks = profile_records(implementation, max_blocks)[index] + index = profile_index( + implementation, max_blocks, submode, active_blocks, cluster_size=cluster_size + ) + mode, blocks = profile_records( + implementation, max_blocks, cluster_size=cluster_size + )[index] return (mode << 16) | blocks diff --git a/flashinfer/experimental/b12x/comm/pcie/_joint_startup.py b/flashinfer/experimental/b12x/comm/pcie/_joint_startup.py index 9f7124c7339..cc369c0cde1 100644 --- a/flashinfer/experimental/b12x/comm/pcie/_joint_startup.py +++ b/flashinfer/experimental/b12x/comm/pcie/_joint_startup.py @@ -5,7 +5,12 @@ import torch from b12x.preparation import FinalistPolicy -from ._joint_profile import JOINT_FAMILIES, JOINT_MODES, selected_profile +from ._joint_profile import ( + JOINT_MODES, + joint_modes, + selected_profile, + validate_cluster_geometry, +) def joint_finalist_policy(): @@ -41,6 +46,8 @@ def joint_startup_spec(plan, *, profiles=None): call = state.query.call family = call["transport"] geometry = (family, call["threads"], call["blocks"]) + cluster_size = call.get("joint_cluster_size", 1) + validate_cluster_geometry(call["blocks"], call["threads"], cluster_size) if profiles is not None: profiles = tuple(profiles) if not profiles or any( @@ -52,7 +59,9 @@ def joint_startup_spec(plan, *, profiles=None): ) encoded = [] for mode, active in profiles: - selected_profile(family, call["blocks"], mode, active) + selected_profile( + family, call["blocks"], mode, active, cluster_size=cluster_size + ) encoded.append((JOINT_MODES[mode], active)) if len(set(encoded)) != len(encoded): raise ValueError("runtime profiles must be distinct") @@ -69,9 +78,11 @@ def joint_startup_spec(plan, *, profiles=None): ): mode = assignment["submode"] active = assignment["active_blocks"] - selected_profile(family, call["blocks"], mode, active) + selected_profile( + family, call["blocks"], mode, active, cluster_size=cluster_size + ) measured[mode] = (JOINT_MODES[mode], active) - modes = JOINT_FAMILIES[family] + modes = joint_modes(family, cluster_size=cluster_size) if set(measured) != set(modes): raise ValueError( "startup tuning requires measured finalists for every mode of the selected geometry" @@ -94,6 +105,8 @@ def joint_startup_spec(plan, *, profiles=None): inplace, profiles, ) + if cluster_size != 1: + key += (("cluster_size", cluster_size),) return JointStartupSpec( key, profiles, state.selector, state, inplace, state.query.rank ) @@ -363,6 +376,7 @@ def prepare_joint_startup_tuner(specs, *, cache_policy="cold"): call["blocks"], call["submode"], call["active_blocks"], + cluster_size=call.get("joint_cluster_size", 1), ) ) producer.status.fill_(-1) diff --git a/flashinfer/experimental/b12x/comm/pcie/_tuning.py b/flashinfer/experimental/b12x/comm/pcie/_tuning.py index 69e18b9b2fb..a1a73408680 100644 --- a/flashinfer/experimental/b12x/comm/pcie/_tuning.py +++ b/flashinfer/experimental/b12x/comm/pcie/_tuning.py @@ -72,6 +72,31 @@ def _validate_query(query: PcieQuery, device) -> None: raise ValueError(f"{query.surface} supports world sizes {worlds}") if not 0 <= query.rank < query.world_size: raise ValueError("rank must belong to the caller's process group") + if "joint_cluster_size" in query.call: + from ._joint_profile import profile_index, validate_cluster_geometry + + call = query.call + if ( + not query.surface.endswith(".all_reduce") + or call.get("transport") not in ("joint_latency", "joint_bandwidth") + or call.get("dtype") != "bfloat16" + or not call.get("stage_input") + ): + raise ValueError( + "joint cluster geometry requires a staged joint BF16 all-reduce" + ) + validate_cluster_geometry( + call["blocks"], call["threads"], call["joint_cluster_size"], device + ) + if call["blocks"] > query.setup.get("joint_max_blocks", 0): + raise ValueError("joint cluster grid exceeds the prepared channel capacity") + profile_index( + call["transport"], + call["blocks"], + call["submode"], + call["active_blocks"], + cluster_size=call["joint_cluster_size"], + ) if query.surface.endswith(".all_gather_heads"): transport = query.call.get("gather_transport", "barrier") from ._gather_preparation import GATHER_MODES, validate_query as validate_gather @@ -114,6 +139,6 @@ def _validate_query(query: PcieQuery, device) -> None: ) -TUNING = replace(TUNING, query_schema_version=41, validate_query=_validate_query) +TUNING = replace(TUNING, query_schema_version=42, validate_query=_validate_query) __all__ = ["PcieQuery", "PcieConfig", "SURFACES", "TUNING"] diff --git a/flashinfer/experimental/b12x/comm/pcie/pcie_allreduce.py b/flashinfer/experimental/b12x/comm/pcie/pcie_allreduce.py index d1bdea65eae..6278a61bf9c 100644 --- a/flashinfer/experimental/b12x/comm/pcie/pcie_allreduce.py +++ b/flashinfer/experimental/b12x/comm/pcie/pcie_allreduce.py @@ -487,6 +487,7 @@ def plan( invocation: FrozenMapping = FrozenMapping(), override: PcieConfig | CollectiveConfig | None = None, joint_only: bool = False, + joint_cluster_size: int = 1, **call, ) -> Plan: """Select an existing native channel once, then declare its plan. @@ -494,9 +495,18 @@ def plan( Semantic channels must already exist. This does not construct transports, register inputs, allocate workspace, or compile any program. Joint kernels require explicit ``joint_only=True``; ordinary plans use static transports. + ``joint_cluster_size=8`` admits phased profiles in a 64-CTA, 512-thread + cluster launch. The default retains the ordinary query and launch policy. """ from ._preparation import plan, query_from_runtime + if type(joint_cluster_size) is not int or joint_cluster_size not in (1, 8): + raise ValueError("joint cluster size must be 1 or 8") + if joint_cluster_size != 1 and (not joint_only or operation != "all_reduce"): + raise ValueError( + "clustered joint planning requires joint_only=True and plain all_reduce" + ) + if operation not in ( "all_reduce", "all_reduce_fused_add_rms_norm", @@ -556,6 +566,26 @@ def plan( ) if default not in alternatives: default = next(iter(alternatives)) + if joint_cluster_size != 1: + from dataclasses import replace + + alternatives = { + name: ( + replace( + native, + call=FrozenMapping( + { + **native.call, + "joint_cluster_size": joint_cluster_size, + "threads": 512, + "blocks": 64, + } + ), + ), + runtime, + ) + for name, (native, runtime) in alternatives.items() + } return collective_plan( alternatives, default=default, diff --git a/tests/experimental/b12x/comm/test_joint_cluster_gpu.py b/tests/experimental/b12x/comm/test_joint_cluster_gpu.py new file mode 100644 index 00000000000..e34e1ce36ad --- /dev/null +++ b/tests/experimental/b12x/comm/test_joint_cluster_gpu.py @@ -0,0 +1,312 @@ +"""Production cluster plans preserve tails, profile transitions and graph ownership.""" + +import ctypes +from datetime import timedelta +import gc +import os +from pathlib import Path + +import pytest + +from .test_pcie_joint_gpu import ( + _operands, + _poison_scratch, + _reference, + _assert_epoch_wrapped, +) + +SIZES = ( + 16, + 48, + 4112, + 162784, + 163840, + 169984, + 325568, + 327680, + 337920, + 339952, + 1310720, + 2293760, +) +CAPACITY_BYTES = 2293760 +PROFILES = { + "joint_latency": ( + ("payload_rsag_phased", 8), + ("payload_rsag", 64), + ("pull", 2), + ("pull", 8), + ("full_phased", 8), + ("full", 25), + ("pull", 7), + ), + "joint_bandwidth": ( + ("payload_rsag_phased", 8), + ("bulk_pull", 11), + ("payload_rsag", 42), + ("payload_rsag", 64), + ("bulk_push", 8), + ("pull", 7), + ), +} + + +def _geometry(graph, cluster): + import cuda.bindings.driver as cuda + + raw = cuda.CUgraph(graph.raw_cuda_graph()) + status, _, count = cuda.cuGraphGetNodes(raw) + assert status == cuda.CUresult.CUDA_SUCCESS and count == 1 + status, nodes, _ = cuda.cuGraphGetNodes(raw, count) + assert status == cuda.CUresult.CUDA_SUCCESS + status, params = cuda.cuGraphKernelNodeGetParams(nodes[0]) + assert status == cuda.CUresult.CUDA_SUCCESS + assert (params.gridDimX, params.gridDimY, params.gridDimZ) == (64, 1, 1) + assert (params.blockDimX, params.blockDimY, params.blockDimZ) == (512, 1, 1) + status, attribute = cuda.cuGraphKernelNodeGetAttribute( + nodes[0], cuda.CUkernelNodeAttrID.CU_LAUNCH_ATTRIBUTE_COOPERATIVE + ) + assert status == cuda.CUresult.CUDA_SUCCESS and attribute.cooperative == 0 + status, attribute = cuda.cuGraphKernelNodeGetAttribute( + nodes[0], cuda.CUkernelNodeAttrID.CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION + ) + assert status == cuda.CUresult.CUDA_SUCCESS + actual = (attribute.clusterDim.x, attribute.clusterDim.y, attribute.clusterDim.z) + assert actual == (8, 1, 1) if cluster else actual in ((0, 0, 0), (1, 1, 1)) + + +def _read_signal(state): + import torch + + data = (ctypes.c_ubyte * state.workspace.signal_bytes)() + stream = torch.cuda.current_stream() + state.runtime._ipc.cudaMemcpyAsync( + ctypes.addressof(data), + state.workspace.shared.local_ptr, + len(data), + stream.cuda_stream, + ) + stream.synchronize() + return bytes(data) + + +def _exercise(rank, family, alias): + import torch + import torch.distributed as dist + from b12x.comm.pcie import ( + CollectiveConfig, + joint_startup_spec, + prepare_joint_startup_tuner, + ) + from b12x.comm.pcie._joint_profile import JOINT_MODES + from b12x.comm.pcie._oneshot_preparation import _prepare_plain_call + from b12x.comm.pcie.pcie_allreduce import PCIeAllReduce + from b12x.preparation import CollectiveRequirement, PreparationSession + + device = torch.device("cuda", rank) + manager = PCIeAllReduce.from_exchange_group( + exchange_group=dist.group.WORLD, + device=device, + single_channel=True, + joint=True, + max_size=CAPACITY_BYTES, + eager_buffer_bytes=CAPACITY_BYTES, + ) + manager.prepare_channels(("cluster-profile-transitions",)) + session = PreparationSession(device=device, autotune=False, compile_workers=2) + cases, graphs, startup_graph = [], [], None + requests = [] + try: + for nbytes in SIZES: + source = torch.full( + (nbytes // 2,), rank + 1, dtype=torch.bfloat16, device=device + ) + out = source if alias else torch.empty_like(source) + plans = {} + for clustered in (False, True): + mode = ( + "payload_rsag_phased" + if clustered + else ("full" if family == "joint_latency" else "payload_rsag") + ) + plan = manager.plan( + source, + out=out, + joint_only=True, + joint_cluster_size=8 if clustered else 1, + override=CollectiveConfig( + implementation=family, + threads=512, + blocks=64, + submode=mode, + active_blocks=8 if clustered else 1, + ), + ) + name = f"{family}-{nbytes}-{alias}-{clustered}" + requests.append( + plan.request( + name=name, + collective=CollectiveRequirement(name, tuple(range(4))), + prepare_call=lambda state, source=source, out=out: ( + _prepare_plain_call(state, inp=source, out=out) + ), + ) + ) + plans[clustered] = plan + cases.append((nbytes, source, out, plans)) + session.prepare( + tuple(requests), + coordinator=lambda progress: ( + progress.ready_collectives[0].key + if progress.ready_collectives + else None + ), + ) + for nbytes, source, out, plans in cases: + for clustered, plan in plans.items(): + graph = torch.cuda.CUDAGraph(keep_graph=True) + dist.barrier() + with session.capture(), manager.capture(), torch.cuda.graph(graph): + manager.all_reduce(source, out=out, plan=plan) + graph.instantiate() + _geometry(graph, clustered) + graphs.append( + (nbytes, clustered, graph, source, out, plan.prepared.state) + ) + assert ( + plan.prepared.state.workspace.capacity_packs == CAPACITY_BYTES // 16 + ) + race_case = next(case for case in cases if case[0] == 327680) + race_plan = race_case[3][True] + race_profiles = PROFILES[family][: 4 if family == "joint_latency" else 3] + spec = joint_startup_spec(race_plan, profiles=race_profiles) + producer = prepare_joint_startup_tuner((spec,), cache_policy="warm") + startup_graph = torch.cuda.CUDAGraph() + dist.barrier() + with session.capture(), manager.capture(), torch.cuda.graph(startup_graph): + producer.run() + session.freeze() + for nbytes, clustered, graph, source, out, state in graphs: + profiles = ( + PROFILES[family] + if clustered + else tuple(p for p in PROFILES[family] if not p[0].endswith("_phased")) + ) + for position, (mode, active) in enumerate(profiles): + state.select_profile(mode, active) + rollover = position == 0 + _poison_scratch(state, rollover=rollover) + values = _operands(4, source.numel(), position + SIZES.index(nbytes)) + expected = _reference(values, rank, mode.removesuffix("_phased")) + for _ in range(3): + source.copy_(values[rank]) + if not alias: + out.fill_(float("nan")) + torch.cuda.synchronize() + dist.barrier() + before = torch.cuda.memory_stats() + graph.replay() + torch.cuda.synchronize() + after = torch.cuda.memory_stats() + for counter in ( + "allocation.all.allocated", + "allocation.all.freed", + "segment.all.allocated", + "segment.all.freed", + ): + assert before[counter] == after[counter] + torch.testing.assert_close( + out.cpu(), expected, rtol=0, atol=0, equal_nan=True + ) + if not alias: + torch.testing.assert_close( + source.cpu(), values[rank], rtol=0, atol=0, equal_nan=True + ) + if rollover: + _assert_epoch_wrapped(state) + if clustered: + for invalid in (1, 7, 9, 16, 64): + with pytest.raises( + RuntimeError, + match="phased joint profiles require exactly one active eight-CTA cluster", + ): + state.select_profile("payload_rsag_phased", invalid) + # Even corrupted device selectors must not partially enter a barrier. + source.fill_(3) + out.fill_(3) + state.selector.fill_((6 << 16) | 7) + torch.cuda.synchronize() + dist.barrier() + before = _read_signal(state) + graph.replay() + torch.cuda.synchronize() + assert _read_signal(state) == before + assert torch.equal(out, torch.full_like(out, 3)) + startup_graph.replay() + torch.cuda.synchronize() + assert producer.status.item() == 0 + winner = int(race_plan.prepared.state.selector.item()) + allowed = {(JOINT_MODES[mode] << 16) | active for mode, active in race_profiles} + assert winner in allowed + winners = [None] * 4 + dist.all_gather_object(winners, winner) + assert len(set(winners)) == 1 + _, _, graph, source, out, _ = next(g for g in graphs if g[0] == 327680 and g[1]) + mode = next( + mode + for mode, active in race_profiles + if winner == (JOINT_MODES[mode] << 16) | active + ) + values = _operands(4, source.numel(), 173) + expected = _reference(values, rank, mode.removesuffix("_phased")) + for _ in range(3): + source.copy_(values[rank]) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + out.cpu(), expected, rtol=0, atol=0, equal_nan=True + ) + finally: + if startup_graph is not None: + startup_graph.reset() + for _, _, graph, *_ in graphs: + graph.reset() + session.close() + manager.close() + + +def _worker(rank, family, alias, directory): + import torch + import torch.distributed as dist + + torch.set_num_threads(1) + torch.cuda.set_device(rank) + dist.init_process_group( + "gloo", + init_method=(Path(directory) / "group").as_uri(), + rank=rank, + world_size=4, + timeout=timedelta(seconds=240), + ) + try: + _exercise(rank, family, alias) + finally: + gc.collect() + dist.destroy_process_group() + gc.collect() + + +@pytest.mark.skipif( + os.environ.get("B12X_RUN_JOINT_CLUSTER_TEST") != "1", + reason="requires an assigned TP4 SM120/SM121 group", +) +@pytest.mark.parametrize("family", ("joint_latency", "joint_bandwidth")) +@pytest.mark.parametrize("alias", (False, True)) +def test_joint_cluster_prepared_graph_transitions(family, alias, tmp_path): + import torch + import torch.multiprocessing as mp + + assert torch.cuda.device_count() == 4 + mp.spawn(_worker, args=(family, alias, str(tmp_path)), nprocs=4, join=True) diff --git a/tests/experimental/b12x/preparation/test_joint_cluster_intervals.py b/tests/experimental/b12x/preparation/test_joint_cluster_intervals.py new file mode 100644 index 00000000000..694999c5686 --- /dev/null +++ b/tests/experimental/b12x/preparation/test_joint_cluster_intervals.py @@ -0,0 +1,89 @@ +"""Owner-interval traversal preserves the payload-ready partition and coverage.""" + +from collections import Counter +import random +import pytest + + +def owner(index, size, world): + part = size // world + return min(index // part, world - 1) if part else world - 1 + + +def interval(peer, size, world): + part = size // world + return peer * part, size if peer == world - 1 else (peer + 1) * part + + +def work(peer, lane, size, world, stride): + begin, end = interval(peer, size, world) + return range(begin + lane, end, stride) + + +@pytest.mark.parametrize("world", (2, 3, 4, 6, 8, 10)) +def test_owner_intervals_preserve_coverage_and_thread_local_handoff(world): + stride = 8 * 512 + cases = sorted( + { + 1, + world - 1, + world, + world + 1, + stride - 1, + stride, + stride + 1, + 81408 // 16, + 81920 // 16, + 84992 // 16, + 162784 // 16, + 163840 // 16, + 169984 // 16, + 325568 // 16, + 327680 // 16, + 337920 // 16, + 339952 // 16, + 1310720 // 16, + 2293760 // 16, + *(random.Random(9151 + i).randrange(1, 50000) for i in range(12)), + } + ) + for size in cases: + all_indices = Counter() + for peer in range(world): + expected = { + index for index in range(size) if owner(index, size, world) == peer + } + assigned = Counter( + index + for lane in range(stride) + for index in work(peer, lane, size, world, stride) + ) + assert set(assigned) == expected and all( + count == 1 for count in assigned.values() + ) + all_indices.update(assigned) + # Scatter to this owner, its reduction, and all of its outgoing + # gathers use this exact per-thread map. Other owners cannot alias it. + begin, end = interval(peer, size, world) + assert 0 <= begin <= end <= size + for other in range(world): + if peer == other: + continue + a, b = interval(other, size, world) + assert max(begin, a) >= min(end, b) + assert all_indices == Counter(range(size)) + + +def test_zero_sized_owner_regions_do_not_remove_phase_barriers(): + for world in (2, 3, 4, 6, 8, 10): + for size in range(1, world): + assert [interval(peer, size, world) for peer in range(world - 1)] == [ + (0, 0) + ] * (world - 1) + assert interval(world - 1, size, world) == (0, size) + for rank in range(world): + scatter = [(rank + phase) % world for phase in range(world)] + gather = [(rank + phase) % world for phase in range(1, world)] + assert len(scatter[1:]) == len(gather) == world - 1 + assert set(scatter) == set(range(world)) + assert set(gather) == set(range(world)) - {rank} diff --git a/tests/experimental/b12x/preparation/test_joint_clusters.py b/tests/experimental/b12x/preparation/test_joint_clusters.py new file mode 100644 index 00000000000..6a446137505 --- /dev/null +++ b/tests/experimental/b12x/preparation/test_joint_clusters.py @@ -0,0 +1,339 @@ +"""Cluster geometry is opt-in metadata; ordinary profiles retain their identity.""" + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from b12x.preparation import FrozenMapping +from b12x.comm.pcie._collective_tuning import CollectiveConfig, CollectiveQuery, TUNING +from b12x.comm.pcie._joint_profile import ( + JOINT_FAMILIES, + JOINT_SIGNAL_BYTES, + joint_modes, + profile_records, + selected_profile, + validate_cluster_geometry, +) +from b12x.comm.pcie._tuning import PcieQuery, TUNING as NATIVE + + +def native_query(family="joint_latency", *, cluster_size=1, mode="full", active=1): + call = dict( + transport=family, + threads=512, + blocks=64, + submode=mode, + active_blocks=active, + dtype="bfloat16", + stage_input=True, + size_packs=10240, + ) + if cluster_size != 1: + call["joint_cluster_size"] = cluster_size + return PcieQuery( + surface="OneshotAllReduce.all_reduce", + world_size=4, + rank=0, + topology="pcie_ipc", + call=FrozenMapping(call), + setup=FrozenMapping( + dict( + shape=(81920,), + joint_max_blocks=64, + rank_data_nbytes=256, + owned_slab_nbytes=512, + owned_slab_count=1, + joint_workspace_nbytes=4096, + ) + ), + ) + + +def collective_query(native): + c = native.call + return CollectiveQuery( + operation="all_reduce", + world_size=4, + rank=0, + alternatives=FrozenMapping( + { + c["transport"]: dict( + surface=native.surface, + world_size=native.world_size, + rank=native.rank, + topology=native.topology, + call=native.call.to_dict(), + setup=native.setup.to_dict(), + ) + } + ), + default=CollectiveConfig( + implementation=c["transport"], + threads=c["threads"], + blocks=c["blocks"], + submode=c["submode"], + active_blocks=c["active_blocks"], + ), + ) + + +def test_ordinary_contract_and_exact_profile_records_are_unchanged(): + assert ( + TUNING.query_schema_version, + TUNING.candidate_contract_version, + TUNING.config_schema_version, + ) == (7, 3, 2) + assert NATIVE.query_schema_version == 42 + assert tuple(CollectiveConfig.__dataclass_fields__) == ( + "implementation", + "threads", + "blocks", + "submode", + "active_blocks", + ) + assert JOINT_FAMILIES == { + "joint_latency": ("full", "payload_rsag", "pull"), + "joint_bandwidth": ("payload_rsag", "bulk_push", "bulk_pull", "pull"), + } + for family, modes in { + "joint_latency": (0, 1, 4), + "joint_bandwidth": (1, 2, 3, 4), + }.items(): + expected = tuple((mode, active) for mode in modes for active in range(1, 65)) + assert profile_records(family, 64) == expected + assert profile_records(family, 64, cluster_size=1) == expected + assert profile_records(family, 64, cluster_size=8)[: len(expected)] == expected + + +@pytest.mark.parametrize( + "family,mode,active", + [ + ("joint_latency", "full", 25), + ("joint_latency", "pull", 7), + ("joint_latency", "payload_rsag", 64), + ("joint_bandwidth", "bulk_pull", 11), + ("joint_bandwidth", "payload_rsag", 42), + ("joint_bandwidth", "bulk_pull", 37), + ("joint_latency", "full_phased", 8), + ("joint_latency", "payload_rsag_phased", 8), + ("joint_bandwidth", "payload_rsag_phased", 8), + ], +) +def test_cluster_config_admits_runtime_profiles_without_candidate_axis_loss( + family, mode, active +): + q = collective_query(native_query(family, cluster_size=8, mode=mode, active=active)) + TUNING.validate_query(q, None) + TUNING.validate_config(q, q.default, None) + assert selected_profile(family, 64, mode, active, cluster_size=8) & 0xFFFF == active + + +@pytest.mark.parametrize("active", [0, 1, 7, 9, 16, 32, 64, True]) +def test_partial_or_multiple_phase_clusters_are_rejected(active): + for family in JOINT_FAMILIES: + with pytest.raises(ValueError): + selected_profile(family, 64, "payload_rsag_phased", active, cluster_size=8) + + +def test_phase_profiles_are_opt_in_and_family_specific(): + for family in JOINT_FAMILIES: + with pytest.raises(ValueError): + selected_profile(family, 64, "payload_rsag_phased", 8) + with pytest.raises(ValueError): + selected_profile("joint_bandwidth", 64, "full_phased", 8, cluster_size=8) + assert joint_modes("joint_latency", cluster_size=8) == ( + *JOINT_FAMILIES["joint_latency"], + "full_phased", + "payload_rsag_phased", + ) + + +@pytest.mark.parametrize( + "blocks,threads,cluster", + [ + (8, 512, 8), + (32, 512, 8), + (63, 512, 8), + (64, 256, 8), + (64, 512, 16), + (64, 512, True), + ], +) +def test_unsupported_cluster_geometry_rejected(blocks, threads, cluster): + with pytest.raises(ValueError): + validate_cluster_geometry(blocks, threads, cluster) + + +@pytest.mark.parametrize("cc,sms", [((9, 0), 132), ((12, 0), 63), ((12, 2), 188)]) +def test_unsupported_cluster_devices_rejected(cc, sms): + with pytest.raises(ValueError): + validate_cluster_geometry( + 64, 512, 8, SimpleNamespace(compute_capability=cc, sm_count=sms) + ) + + +@pytest.mark.parametrize( + "surface,transport", + [ + ("OneshotAllReduce.all_reduce_fused_add", "joint_latency"), + ("OneshotAllReduce.all_reduce", "pull"), + ], +) +def test_cluster_metadata_cannot_escape_joint_plain_surface(surface, transport): + q = native_query(cluster_size=8) + q = replace( + q, surface=surface, call=FrozenMapping({**q.call, "transport": transport}) + ) + with pytest.raises(ValueError, match="staged joint BF16"): + NATIVE.validate_query(q, None) + + +def test_cluster_compile_job_and_workspace_keep_standard_preparation_ownership(): + from b12x.comm.pcie import _joint_preparation as preparation + from b12x._lib.compile_pool import CompileJob + + runtime = SimpleNamespace(device=torch.device("cpu")) + ordinary = preparation.plan(native_query(), runtime=runtime) + clustered = preparation.plan(native_query(cluster_size=8), runtime=runtime) + (old,) = ordinary._compile_jobs(None, None) + (new,) = clustered._compile_jobs(None, None) + assert old == CompileJob.create( + "b12x.comm.pcie._joint_preparation:compile_joint", + "latency", + 4, + 0, + 512, + 64, + CompileJob.DEVICE_ORDINAL, + ) + assert new == CompileJob.create( + "b12x.comm.pcie._joint_preparation:compile_joint_cluster", + "latency", + 4, + 0, + 512, + 64, + 8, + CompileJob.DEVICE_ORDINAL, + ) + assert [ + x.resident_nbytes for x in ordinary._memory_requirements(None, None).persistent + ] == [ + x.resident_nbytes for x in clustered._memory_requirements(None, None).persistent + ] + signals, slot, total = preparation._layout(2293760 // 16, 4) + assert JOINT_SIGNAL_BYTES == 266248 + assert signals == 266496 and total == signals + 2 * slot + + +def test_cluster_materialization_unwraps_detected_device_identity(monkeypatch): + from b12x.preparation import DetectedDevice, DeviceIdentity + from b12x.comm.pcie import _joint_preparation as preparation + from b12x.comm.pcie import _joint_cluster_cute as kernels + + identity = DeviceIdentity("nvidia", (12, 0), 188, "SM120 test device") + device = DetectedDevice(ordinal=3, identity=identity) + workspace, launcher = object(), object() + runtime = SimpleNamespace(device=torch.device("cpu"), _joint_workspace=workspace) + compiled, checked = [], [] + monkeypatch.setattr( + preparation, + "compile_joint_cluster", + lambda *args: compiled.append(args) or launcher, + ) + monkeypatch.setattr(preparation, "load_programs", lambda value: value) + monkeypatch.setattr( + kernels, "validate_cluster_residency", lambda *args: checked.append(args) + ) + monkeypatch.setattr(preparation, "_JointExecutionState", lambda *args: args) + query = native_query(cluster_size=8) + plan = preparation.plan(query, runtime=runtime) + assert plan._materialize(None, device) == (query, runtime, launcher, workspace) + assert compiled == [("latency", 4, 0, 512, 64, 8, 3)] + assert checked == [(launcher, 64, 512, 8)] + + +def test_startup_keys_and_phase_admission_use_immutable_cluster_geometry(monkeypatch): + from b12x.comm.pcie._joint_preparation import _JointExecutionState + from b12x.comm.pcie import joint_startup_spec + + monkeypatch.setattr( + torch.distributed, "get_process_group_ranks", lambda group: [0, 1, 2, 3] + ) + + def prepared(cluster): + state = _JointExecutionState.__new__(_JointExecutionState) + state.query = native_query(cluster_size=cluster) + state.workspace = SimpleNamespace(capacity_packs=143360) + state.selector = torch.zeros(1, dtype=torch.uint32) + state.runtime = SimpleNamespace(exchange_group=object()) + return SimpleNamespace( + prepared=SimpleNamespace(state=state, finalists=()), + query=SimpleNamespace(inplace=False), + ) + + plain = prepared(1) + clustered = prepared(8) + a = joint_startup_spec(plain, profiles=(("pull", 7),)) + b = joint_startup_spec(clustered, profiles=(("pull", 7),)) + assert a.key == ( + "joint_startup", + 1, + 4, + (0, 1, 2, 3), + "joint_latency", + 512, + 64, + (81920,), + "bfloat16", + 10240, + 143360, + False, + ((4, 7),), + ) + assert b.key == (*a.key, ("cluster_size", 8)) + assert joint_startup_spec( + clustered, profiles=(("payload_rsag_phased", 8), ("payload_rsag", 64)) + ).profiles == ((6, 8), (1, 64)) + with pytest.raises(ValueError): + joint_startup_spec(clustered, profiles=(("payload_rsag_phased", 7),)) + + +def test_public_plan_omits_default_cluster_field(monkeypatch): + from b12x.comm.pcie import _preparation, _collective_tuning + from b12x.comm.pcie.pcie_allreduce import PCIeAllReduce + + native = native_query() + runtime = SimpleNamespace(device=torch.device("cpu")) + initial = replace(native, call=FrozenMapping({**native.call, "transport": "pull"})) + monkeypatch.setattr(_preparation, "query_from_runtime", lambda *a, **k: initial) + monkeypatch.setattr( + _collective_tuning, + "oneshot_alternatives", + lambda *a: {"joint_latency": (native, runtime)}, + ) + monkeypatch.setattr( + _collective_tuning, "plan", lambda alternatives, **kwargs: alternatives + ) + manager = PCIeAllReduce.__new__(PCIeAllReduce) + manager.algorithm = "oneshot" + manager._bf16_twoshot = None + manager._runtime = SimpleNamespace( + single_channel=True, _prepared_channel_for_stream=lambda *a: runtime + ) + inp = torch.zeros(16, dtype=torch.bfloat16) + omitted = manager.plan(inp, joint_only=True) + explicit = manager.plan(inp, joint_only=True, joint_cluster_size=1) + assert omitted == explicit == {"joint_latency": (native, runtime)} + clustered = manager.plan(inp, joint_only=True, joint_cluster_size=8)[ + "joint_latency" + ][0] + assert clustered.call == FrozenMapping({**native.call, "joint_cluster_size": 8}) + for kwargs in [ + dict(joint_only=False), + dict(joint_only=True, operation="all_reduce_fused_add"), + ]: + with pytest.raises(ValueError, match="plain all_reduce"): + manager.plan(inp, joint_cluster_size=8, **kwargs)