From 560fef16648e4953bf39b701a226c322a0ab6c7c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 2 Nov 2026 19:50:43 +0000 Subject: [PATCH] Add explicit clustered phased PCIe all-reduce profiles Prepare fixed 64-CTA, 512-thread joint launches with one active eight-CTA cluster for phased BF16 transfers. Preserve ordinary profile identities and IPC layouts while validating cluster eligibility and residency. Validate profile admission and owner intervals on CPU, and exact CUDA graph replay, startup selection, aliasing, allocation stability, and profile transitions on TP4 SM120. AI-assisted implementation and review. --- .../experimental/b12x/comm/pcie/README.md | 53 +++ .../b12x/comm/pcie/_collective_tuning.py | 21 +- .../b12x/comm/pcie/_joint_cluster_cute.py | 382 ++++++++++++++++++ .../b12x/comm/pcie/_joint_preparation.py | 61 ++- .../b12x/comm/pcie/_joint_profile.py | 76 +++- .../b12x/comm/pcie/_joint_startup.py | 22 +- .../experimental/b12x/comm/pcie/_tuning.py | 27 +- .../b12x/comm/pcie/pcie_allreduce.py | 30 ++ .../b12x/comm/test_joint_cluster_gpu.py | 312 ++++++++++++++ .../test_joint_cluster_intervals.py | 89 ++++ .../b12x/preparation/test_joint_clusters.py | 339 ++++++++++++++++ 11 files changed, 1390 insertions(+), 22 deletions(-) create mode 100644 flashinfer/experimental/b12x/comm/pcie/_joint_cluster_cute.py create mode 100644 tests/experimental/b12x/comm/test_joint_cluster_gpu.py create mode 100644 tests/experimental/b12x/preparation/test_joint_cluster_intervals.py create mode 100644 tests/experimental/b12x/preparation/test_joint_clusters.py 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)