From df604cfe10a7ada011c5674ca3c18cee21387bf2 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Wed, 2 Sep 2026 03:00:10 -0400 Subject: [PATCH] feat(comm): add a lossless bf16 PCIe two-shot all-reduce Add PCIeTwoShotBF16, a pull-mode two-shot all-reduce for TP decode payloads above the one-shot ceiling and below the DMA ring floor, where the NCCL ring is the incumbent. Each rank stages its bf16 payload in its IPC slab, a per-CTA flag barrier follows, every rank pulls the peers' packs of its quarter and accumulates them in fp32 in a fixed rank order, and after a second barrier the reduced quarters are pulled into place. The result is rounded to bf16 once, so it lies within one bf16 rounding of the exact sum (the bf16 ring rounds after every hop) and is deterministic across calls and graph replays. Four RTX PRO 6000 Blackwell over PCIe, graph replay, two-shot vs NCCL ring: 128 KB 15.0 vs 16.7 us, 256 KB 20.3 vs 25.6, 384 KB 26.3 vs 33.5, 512 KB 32.8 vs 41.7, 768 KB 44.5 vs 54.5; equal from 1 MB up. Co-Authored-By: Claude Fable 5.1 --- b12x/comm/pcie/_twoshot_bf16_cute.py | 1059 ++++++++++++++++++++++++++ b12x/comm/pcie/pcie_twoshot_bf16.py | 707 +++++++++++++++++ tests/comm/test_pcie_twoshot_bf16.py | 152 ++++ 3 files changed, 1918 insertions(+) create mode 100644 b12x/comm/pcie/_twoshot_bf16_cute.py create mode 100644 b12x/comm/pcie/pcie_twoshot_bf16.py create mode 100644 tests/comm/test_pcie_twoshot_bf16.py diff --git a/b12x/comm/pcie/_twoshot_bf16_cute.py b/b12x/comm/pcie/_twoshot_bf16_cute.py new file mode 100644 index 000000000..50935fa63 --- /dev/null +++ b/b12x/comm/pcie/_twoshot_bf16_cute.py @@ -0,0 +1,1059 @@ +"""CuTeDSL kernels for lossless BF16 PCIe two-shot collectives. + +Structure mirrors :mod:`_twoshot_cute` (the fp8-transport variant): phase one +pushes this rank's shard into every peer's IPC staging slot (posted PCIe +writes), a per-CTA flag barrier follows, and phase two either reduces the +local quarter (reduce_scatter) or copies the staged shards into place +(all_gather). Payload packs are 16 bytes = 8 bf16 values; the reduction +accumulates in fp32 in a fixed rank order (local rank first, then +``(local + i) % world`` for ``i = 1..world-1``) and rounds once to bf16, so a +given rank's output is deterministic across runs. +""" + +from __future__ import annotations + +import functools +from collections.abc import Callable, Sequence + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64, Uint32 + +from b12x._lib.compiler import KernelCompileSpec +from b12x._lib.compiler import compile as b12x_compile +from b12x._lib.intrinsics import ld_global_nc_v4_u32, ld_global_v4_u32, st_global_v4_u32 +from b12x._lib.runtime_control import raise_if_kernel_resolution_frozen +from b12x._lib.utils import current_cuda_stream, make_ptr + +from ._cute_intrinsics import ( + graph_epoch_arrive_serialized, + ld_relaxed_gpu_u32, + pack_f32x2_to_bf16x2, + unpack_bf16x2, +) +from ._twoshot_cute import ( + _GRAPH_BLOCKS_ARRIVED_OFFSET, + _GRAPH_EPOCH_OFFSET, + _MAX_BLOCKS, + _MAX_RANKS, + _FLAG_STRIDE, + _SELF_COUNTER_BYTES, + _fence_sc_sys, + _ld_generic_v4_u32, + _ld_global_u32, + _ld_relaxed_sys_u32, + _st_generic_v4_u32, + _st_global_u32, + _st_relaxed_sys_u32, +) + +_PREPARED_BF16_LAUNCHERS: set[tuple[object, ...]] = set() +_PACK_ELEMS = 8 # bf16 values per 16-byte pack + + +class _TwoShotBf16Launch: + def __init__( + self, + operation: str, + world_size: int, + rank: int, + device_slot_selection: bool, + slot_bias: int, + threads: int, + row_elems: int, + ) -> None: + if operation not in ("reduce_scatter", "all_gather"): + raise ValueError(f"invalid two-shot operation {operation!r}") + self._operation = operation + self._world_size = int(world_size) + self._rank = int(rank) + self._device_slot_selection = bool(device_slot_selection) + self._slot_bias = int(slot_bias) & 1 + self._threads = int(threads) + self._row_elems = int(row_elems) + + @cute.jit + def __call__( + self, + payload: cute.Pointer, + staging0: cute.Pointer, + staging1: cute.Pointer, + staging2: cute.Pointer, + staging3: cute.Pointer, + staging4: cute.Pointer, + staging5: cute.Pointer, + staging6: cute.Pointer, + staging7: cute.Pointer, + signal0: cute.Pointer, + signal1: cute.Pointer, + signal2: cute.Pointer, + signal3: cute.Pointer, + signal4: cute.Pointer, + signal5: cute.Pointer, + signal6: cute.Pointer, + signal7: cute.Pointer, + output: cute.Pointer, + rank: Int32, + pack_stride: Int64, + slot_bytes: Int64, + rows_per_rank: Int32, + grid_x: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + payload, + staging0, + staging1, + staging2, + staging3, + staging4, + staging5, + staging6, + staging7, + signal0, + signal1, + signal2, + signal3, + signal4, + signal5, + signal6, + signal7, + output, + rank, + pack_stride, + slot_bytes, + rows_per_rank, + ).launch( + grid=(grid_x, 1, 1), + block=[self._threads, 1, 1], + max_number_threads=(512, 1, 1), + min_blocks_per_mp=1, + cluster=(1, 1, 1), + stream=stream, + ) + + @cute.jit + def _select_address( + self, + pointers: Sequence[cute.Pointer], + index: Int32, + ) -> Int64: + """Select one scalar launch pointer without unrolling peer work.""" + + address = Int64(pointers[0].toint()) + if cutlass.const_expr(self._world_size == 2): + if index == Int32(1): + address = Int64(pointers[1].toint()) + return address + if index < Int32(4): + if index < Int32(2): + address = Int64(pointers[0].toint()) + if index == Int32(1): + address = Int64(pointers[1].toint()) + else: + address = Int64(pointers[2].toint()) + if index == Int32(3): + address = Int64(pointers[3].toint()) + else: + if index < Int32(6): + address = Int64(pointers[4].toint()) + if index == Int32(5): + address = Int64(pointers[5].toint()) + else: + address = Int64(pointers[6].toint()) + if cutlass.const_expr(self._world_size == 8): + if index == Int32(7): + address = Int64(pointers[7].toint()) + return address + + @cute.jit + def _barrier( + self, + signals: Sequence[cute.Pointer], + rank: Int32, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + cute.arch.barrier() + if tidx < Int32(self._world_size): + _fence_sc_sys() + if cutlass.const_expr(self._operation == "all_gather"): + self_base = Int64(signals[self._rank].toint()) + else: + self_base = self._select_address(signals, rank) + self_counter_address = self_base + ( + Int64(bidx) * Int64(_MAX_RANKS) + Int64(tidx) + ) * Int64(4) + value = _ld_global_u32(self_counter_address) + Uint32(1) + _st_global_u32(self_counter_address, value) + + flag_slot = Int64(value % Uint32(2)) + peer_base = self._select_address(signals, tidx) + peer_counter_address = ( + peer_base + + Int64(_SELF_COUNTER_BYTES) + + ( + (flag_slot * Int64(_MAX_BLOCKS) + Int64(bidx)) + * Int64(_MAX_RANKS * _FLAG_STRIDE) + + Int64(rank) * Int64(_FLAG_STRIDE) + ) + * Int64(4) + ) + self_counter_address = ( + self_base + + Int64(_SELF_COUNTER_BYTES) + + ( + (flag_slot * Int64(_MAX_BLOCKS) + Int64(bidx)) + * Int64(_MAX_RANKS * _FLAG_STRIDE) + + Int64(tidx) * Int64(_FLAG_STRIDE) + ) + * Int64(4) + ) + _st_relaxed_sys_u32(peer_counter_address, value) + observed = _ld_relaxed_sys_u32(self_counter_address) + while observed != value: + observed = _ld_relaxed_sys_u32(self_counter_address) + cute.arch.barrier() + + @cute.jit + def _accumulate_words( + self, + accumulator: cute.Tensor, + words, + ) -> None: + for word_index in cutlass.range_constexpr(4): + lo, hi = unpack_bf16x2(words[word_index]) + element = word_index * 2 + accumulator[element] = accumulator[element] + lo + accumulator[element + 1] = accumulator[element + 1] + hi + + @cute.jit + def _load_accumulate_pack_global_nc( + self, + accumulator: cute.Tensor, + address: Int64, + ) -> None: + words = ld_global_nc_v4_u32(address) + self._accumulate_words(accumulator, words) + + @cute.jit + def _load_accumulate_pack_generic( + self, + accumulator: cute.Tensor, + address: Int64, + ) -> None: + words = _ld_generic_v4_u32(address) + self._accumulate_words(accumulator, words) + + @cute.jit + def _store_pack(self, output_address: Int64, accumulator: cute.Tensor) -> None: + st_global_v4_u32( + output_address, + pack_f32x2_to_bf16x2(accumulator[0], accumulator[1]), + pack_f32x2_to_bf16x2(accumulator[2], accumulator[3]), + pack_f32x2_to_bf16x2(accumulator[4], accumulator[5]), + pack_f32x2_to_bf16x2(accumulator[6], accumulator[7]), + ) + + @cute.kernel + def kernel( + self, + payload: cute.Pointer, + staging0: cute.Pointer, + staging1: cute.Pointer, + staging2: cute.Pointer, + staging3: cute.Pointer, + staging4: cute.Pointer, + staging5: cute.Pointer, + staging6: cute.Pointer, + staging7: cute.Pointer, + signal0: cute.Pointer, + signal1: cute.Pointer, + signal2: cute.Pointer, + signal3: cute.Pointer, + signal4: cute.Pointer, + signal5: cute.Pointer, + signal6: cute.Pointer, + signal7: cute.Pointer, + output: cute.Pointer, + rank: Int32, + pack_stride: Int64, + slot_bytes: Int64, + rows_per_rank: Int32, + ) -> None: + staging = ( + staging0, + staging1, + staging2, + staging3, + staging4, + staging5, + staging6, + staging7, + ) + signals = ( + signal0, + signal1, + signal2, + signal3, + signal4, + signal5, + signal6, + signal7, + ) + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + gdim, _, _ = cute.arch.grid_dim() + block_threads, _, _ = cute.arch.block_dim() + if cutlass.const_expr(self._operation == "all_gather"): + local_rank = Int32(self._rank) + else: + local_rank = rank + packs_per_row = Int32(self._row_elems // _PACK_ELEMS) + shard_packs = Int64(rows_per_rank) * Int64(packs_per_row) + chunk = (shard_packs + Int64(gdim) - Int64(1)) // Int64(gdim) + begin = Int64(bidx) * chunk + end = begin + chunk + if end > shard_packs: + end = shard_packs + + payload_address = Int64(payload.toint()) + staging_slot_offset = Int64(0) + if cutlass.const_expr(self._device_slot_selection): + if cutlass.const_expr(self._operation == "all_gather"): + self_signal = Int64(signals[self._rank].toint()) + else: + self_signal = self._select_address(signals, local_rank) + generation = ld_relaxed_gpu_u32(self_signal + Int64(_GRAPH_EPOCH_OFFSET)) + staging_slot_offset = ( + Int64((generation + Uint32(self._slot_bias)) % Uint32(2)) * slot_bytes + ) + cute.arch.barrier() + if Int32(tidx) == Int32(self._threads - 1): + graph_epoch_arrive_serialized( + self_signal + Int64(_GRAPH_EPOCH_OFFSET), + self_signal + Int64(_GRAPH_BLOCKS_ARRIVED_OFFSET), + Uint32(gdim), + ) + + # Phase one: push remote shards (posted PCIe writes), rank-staggered. + peer_index = Int32(1) + while peer_index < Int32(self._world_size): + destination = (local_rank + peer_index) % Int32(self._world_size) + destination_base = ( + self._select_address(staging, destination) + staging_slot_offset + ) + destination_payload = destination_base + ( + Int64(local_rank) * pack_stride * Int64(16) + ) + if cutlass.const_expr(self._operation == "reduce_scatter"): + source_pack = Int64(destination) * shard_packs + else: + source_pack = Int64(0) + + index = begin + Int64(tidx) + while index < end: + words = ld_global_nc_v4_u32( + payload_address + (source_pack + index) * Int64(16) + ) + _st_generic_v4_u32( + destination_payload + index * Int64(16), + words[0], + words[1], + words[2], + words[3], + ) + index += Int64(block_threads) + peer_index += Int32(1) + + self._barrier(signals, local_rank) + + if cutlass.const_expr(self._operation == "all_gather"): + self_base = Int64(staging[self._rank].toint()) + else: + self_base = self._select_address(staging, local_rank) + self_base += staging_slot_offset + output_address = Int64(output.toint()) + if cutlass.const_expr(self._operation == "reduce_scatter"): + index = begin + Int64(tidx) + while index < end: + accumulator = cute.make_rmem_tensor((_PACK_ELEMS,), cutlass.Float32) + for lane in cutlass.range_constexpr(_PACK_ELEMS): + accumulator[lane] = Float32(0.0) + + local_pack = Int64(local_rank) * shard_packs + index + self._load_accumulate_pack_global_nc( + accumulator, + payload_address + local_pack * Int64(16), + ) + + for peer_index in cutlass.range( + Int32(1), + Int32(self._world_size), + Int32(1), + unroll=1, + ): + source_rank = (local_rank + peer_index) % Int32(self._world_size) + staged_pack = ( + self_base + + Int64(source_rank) * pack_stride * Int64(16) + + index * Int64(16) + ) + self._load_accumulate_pack_generic(accumulator, staged_pack) + + self._store_pack(output_address + index * Int64(16), accumulator) + index += Int64(self._threads) + else: + first_index = begin + Int64(tidx) + iteration_count = Int32( + (end - first_index + Int64(self._threads - 1)) // Int64(self._threads) + ) + peer_index = Int32(0) + index = Int64(0) + while peer_index < Int32(self._world_size): + source_rank = (local_rank + peer_index) % Int32(self._world_size) + source_payload_base = Int64(0) + if source_rank == local_rank: + source_payload_base = payload_address + else: + source_payload_base = self_base + Int64( + source_rank + ) * pack_stride * Int64(16) + destination_base = output_address + Int64( + source_rank + ) * shard_packs * Int64(16) + for iteration in cutlass.range( + Int32(0), + iteration_count, + Int32(1), + unroll=1, + ): + index = first_index + Int64(iteration) * Int64(self._threads) + # Generic addressing serves both the local payload and the + # IPC-mapped peer slabs. + words = _ld_generic_v4_u32(source_payload_base + index * Int64(16)) + st_global_v4_u32( + destination_base + index * Int64(16), + words[0], + words[1], + words[2], + words[3], + ) + peer_index += Int32(1) + + +def _bf16_process_key( + operation: str, + world_size: int, + rank: int, + device_slot_selection: bool, + slot_bias: int, + threads: int, + row_elems: int, + device_index: int, +) -> tuple[object, ...]: + return ( + "bf16", + str(operation), + int(world_size), + int(rank), + bool(device_slot_selection), + int(slot_bias) & 1 if device_slot_selection else 0, + int(threads), + int(row_elems), + int(device_index), + ) + + +def is_twoshot_bf16_launcher_prepared( + operation: str, + world_size: int, + rank: int, + device_slot_selection: bool, + slot_bias: int, + threads: int, + row_elems: int, + device_index: int, +) -> bool: + return ( + _bf16_process_key( + operation, + world_size, + rank, + device_slot_selection, + slot_bias, + threads, + row_elems, + device_index, + ) + in _PREPARED_BF16_LAUNCHERS + ) + + +@functools.cache +def get_twoshot_bf16_launcher( + operation: str, + world_size: int, + rank: int, + device_slot_selection: bool, + slot_bias: int, + threads: int, + row_elems: int, + device_index: int, +) -> Callable[..., None]: + """Compile and return one static world/operation/thread specialization.""" + process_key = _bf16_process_key( + operation, + world_size, + rank, + device_slot_selection, + slot_bias, + threads, + row_elems, + device_index, + ) + del device_index # part of the process-local cache key + if world_size not in (2, 4, 8): + raise ValueError(f"unsupported world size {world_size}") + if rank < 0 or rank >= world_size: + raise ValueError(f"rank {rank} is outside world size {world_size}") + if threads <= 0 or threads > 512 or threads % 32 != 0: + raise ValueError("threads must be a warp-aligned value in [32, 512]") + if row_elems <= 0 or row_elems % _PACK_ELEMS != 0: + raise ValueError("row_elems must be a positive multiple of 8") + slot_bias = int(slot_bias) & 1 + launch = _TwoShotBf16Launch( + operation, + world_size, + rank, + device_slot_selection, + slot_bias, + threads, + row_elems, + ) + cache_key = ( + "bf16", + operation, + int(world_size), + int(rank), + bool(device_slot_selection), + slot_bias, + int(threads), + int(row_elems), + ) + raise_if_kernel_resolution_frozen( + "cute.compile", target=launch, cache_key=cache_key + ) + raw = b12x_compile( + launch, + make_ptr(cutlass.Uint32, 16, cute.AddressSpace.gmem, assumed_align=16), + *( + make_ptr( + cutlass.Uint32, + 16, + cute.AddressSpace.gmem, + assumed_align=16, + ) + for _ in range(8) + ), + *( + make_ptr( + cutlass.Uint32, + 16, + cute.AddressSpace.gmem, + assumed_align=4, + ) + for _ in range(8) + ), + make_ptr(cutlass.Uint32, 16, cute.AddressSpace.gmem, assumed_align=16), + 0, + 1, + 1, + 1, + 1, + current_cuda_stream(), + compile_spec=KernelCompileSpec.from_key( + f"comm.pcie.twoshot_bf16.{operation}", + 1, + cache_key, + ), + ) + + def run( + payload_address: int, + staging_addresses: Sequence[int], + signal_addresses: Sequence[int], + output_address: int, + rank: int, + pack_stride: int, + slot_bytes: int, + rows_per_rank: int, + grid_x: int, + ) -> None: + if len(staging_addresses) != 8 or len(signal_addresses) != 8: + raise ValueError("two-shot scalar pointer ABI requires eight peers") + raw_args = ( + make_ptr( + cutlass.Uint32, + payload_address, + cute.AddressSpace.gmem, + assumed_align=16, + ), + *( + make_ptr( + cutlass.Uint32, + address, + cute.AddressSpace.gmem, + assumed_align=16, + ) + for address in staging_addresses + ), + *( + make_ptr( + cutlass.Uint32, + address, + cute.AddressSpace.gmem, + assumed_align=4, + ) + for address in signal_addresses + ), + make_ptr( + cutlass.Uint32, + output_address, + cute.AddressSpace.gmem, + assumed_align=16, + ), + rank, + pack_stride, + slot_bytes, + rows_per_rank, + grid_x, + current_cuda_stream(), + ) + raw(*raw_args) + + _PREPARED_BF16_LAUNCHERS.add(process_key) + return run + + +class _TwoShotPullAllReduceLaunch(_TwoShotBf16Launch): + """Single-launch lossless bf16 all-reduce built on remote READS only. + + Phase A stages this rank's full payload into its own IPC slab (local + copy). After barrier #1 every rank pulls its quarter of every peer's + staged payload, accumulates in fp32 (rank order: self, then + ``(rank + i) % world``), rounds once to bf16 and writes the reduced quarter + both to the output and to its slab's reduced region. After barrier #2 + every rank pulls the other ranks' reduced quarters into the output. + Wire bytes per rank: 1.5P of PCIe reads (same as ring), two barriers. + """ + + def __init__( + self, + world_size: int, + rank: int, + device_slot_selection: bool, + slot_bias: int, + threads: int, + row_elems: int, + ) -> None: + super().__init__( + "reduce_scatter", # barrier addressing uses the dynamic-rank path + world_size, + rank, + device_slot_selection, + slot_bias, + threads, + row_elems, + ) + + @cute.jit + def __call__( + self, + payload: cute.Pointer, + staging0: cute.Pointer, + staging1: cute.Pointer, + staging2: cute.Pointer, + staging3: cute.Pointer, + staging4: cute.Pointer, + staging5: cute.Pointer, + staging6: cute.Pointer, + staging7: cute.Pointer, + signal0: cute.Pointer, + signal1: cute.Pointer, + signal2: cute.Pointer, + signal3: cute.Pointer, + signal4: cute.Pointer, + signal5: cute.Pointer, + signal6: cute.Pointer, + signal7: cute.Pointer, + output: cute.Pointer, + rank: Int32, + reduced_offset: Int64, + slot_bytes: Int64, + rows_per_rank: Int32, + grid_x: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + payload, + staging0, + staging1, + staging2, + staging3, + staging4, + staging5, + staging6, + staging7, + signal0, + signal1, + signal2, + signal3, + signal4, + signal5, + signal6, + signal7, + output, + rank, + reduced_offset, + slot_bytes, + rows_per_rank, + ).launch( + grid=(grid_x, 1, 1), + block=[self._threads, 1, 1], + max_number_threads=(512, 1, 1), + min_blocks_per_mp=1, + cluster=(1, 1, 1), + stream=stream, + ) + + @cute.jit + def _load_accumulate_pack_remote( + self, + accumulator: cute.Tensor, + address: Int64, + ) -> None: + words = ld_global_v4_u32(address) + self._accumulate_words(accumulator, words) + + @cute.kernel + def kernel( + self, + payload: cute.Pointer, + staging0: cute.Pointer, + staging1: cute.Pointer, + staging2: cute.Pointer, + staging3: cute.Pointer, + staging4: cute.Pointer, + staging5: cute.Pointer, + staging6: cute.Pointer, + staging7: cute.Pointer, + signal0: cute.Pointer, + signal1: cute.Pointer, + signal2: cute.Pointer, + signal3: cute.Pointer, + signal4: cute.Pointer, + signal5: cute.Pointer, + signal6: cute.Pointer, + signal7: cute.Pointer, + output: cute.Pointer, + rank: Int32, + reduced_offset: Int64, + slot_bytes: Int64, + rows_per_rank: Int32, + ) -> None: + staging = ( + staging0, + staging1, + staging2, + staging3, + staging4, + staging5, + staging6, + staging7, + ) + signals = ( + signal0, + signal1, + signal2, + signal3, + signal4, + signal5, + signal6, + signal7, + ) + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + gdim, _, _ = cute.arch.grid_dim() + local_rank = rank + packs_per_row = Int32(self._row_elems // _PACK_ELEMS) + quarter_packs = Int64(rows_per_rank) * Int64(packs_per_row) + full_packs = quarter_packs * Int64(self._world_size) + threads = Int64(self._threads) + grid_threads = Int64(gdim) * threads + flat = Int64(bidx) * threads + Int64(tidx) + + payload_address = Int64(payload.toint()) + output_address = Int64(output.toint()) + staging_slot_offset = Int64(0) + self_signal = self._select_address(signals, local_rank) + if cutlass.const_expr(self._device_slot_selection): + generation = ld_relaxed_gpu_u32(self_signal + Int64(_GRAPH_EPOCH_OFFSET)) + staging_slot_offset = ( + Int64((generation + Uint32(self._slot_bias)) % Uint32(2)) * slot_bytes + ) + cute.arch.barrier() + if Int32(tidx) == Int32(self._threads - 1): + graph_epoch_arrive_serialized( + self_signal + Int64(_GRAPH_EPOCH_OFFSET), + self_signal + Int64(_GRAPH_BLOCKS_ARRIVED_OFFSET), + Uint32(gdim), + ) + self_base = self._select_address(staging, local_rank) + staging_slot_offset + + # Phase A: stage the full local payload into this rank's own slab. + index = flat + while index < full_packs: + words = ld_global_nc_v4_u32(payload_address + index * Int64(16)) + st_global_v4_u32( + self_base + index * Int64(16), + words[0], + words[1], + words[2], + words[3], + ) + index += grid_threads + + self._barrier(signals, local_rank) + + # Phase B: pull my quarter from every peer's staged payload, reduce. + quarter_base = Int64(local_rank) * quarter_packs + index = flat + while index < quarter_packs: + accumulator = cute.make_rmem_tensor((_PACK_ELEMS,), cutlass.Float32) + for lane in cutlass.range_constexpr(_PACK_ELEMS): + accumulator[lane] = Float32(0.0) + # Issue the local and all remote pack loads back-to-back (the peer + # loop is unrolled at compile time), then accumulate in fixed + # rank order: self, then (rank + i) % world for i = 1..world-1. + pack_offset = (quarter_base + index) * Int64(16) + local_words = ld_global_nc_v4_u32(payload_address + pack_offset) + peer_words = [] + for peer_index in cutlass.range_constexpr(1, self._world_size): + source_rank = (local_rank + Int32(peer_index)) % Int32(self._world_size) + peer_base = ( + self._select_address(staging, source_rank) + staging_slot_offset + ) + peer_words.append(ld_global_v4_u32(peer_base + pack_offset)) + self._accumulate_words(accumulator, local_words) + for peer_index in cutlass.range_constexpr(self._world_size - 1): + self._accumulate_words(accumulator, peer_words[peer_index]) + self._store_pack( + output_address + (quarter_base + index) * Int64(16), accumulator + ) + self._store_pack( + self_base + reduced_offset + index * Int64(16), accumulator + ) + index += grid_threads + + self._barrier(signals, local_rank) + + # Phase C: pull the other ranks' reduced quarters into the output. + for peer_index in cutlass.range_constexpr(1, self._world_size): + source_rank = (local_rank + Int32(peer_index)) % Int32(self._world_size) + peer_reduced = ( + self._select_address(staging, source_rank) + + staging_slot_offset + + reduced_offset + ) + destination = output_address + Int64(source_rank) * quarter_packs * Int64( + 16 + ) + index = flat + while index < quarter_packs: + words = ld_global_v4_u32(peer_reduced + index * Int64(16)) + st_global_v4_u32( + destination + index * Int64(16), + words[0], + words[1], + words[2], + words[3], + ) + index += grid_threads + + +@functools.cache +def get_twoshot_bf16_allreduce_launcher( + world_size: int, + rank: int, + device_slot_selection: bool, + slot_bias: int, + threads: int, + row_elems: int, + device_index: int, +) -> Callable[..., None]: + """Compile the single-launch pull-based bf16 all-reduce specialization.""" + process_key = _bf16_process_key( + "all_reduce_pull", + world_size, + rank, + device_slot_selection, + slot_bias, + threads, + row_elems, + device_index, + ) + del device_index + if world_size not in (2, 4, 8): + raise ValueError(f"unsupported world size {world_size}") + if rank < 0 or rank >= world_size: + raise ValueError(f"rank {rank} is outside world size {world_size}") + if threads <= 0 or threads > 512 or threads % 32 != 0: + raise ValueError("threads must be a warp-aligned value in [32, 512]") + if row_elems <= 0 or row_elems % _PACK_ELEMS != 0: + raise ValueError("row_elems must be a positive multiple of 8") + slot_bias = int(slot_bias) & 1 + launch = _TwoShotPullAllReduceLaunch( + world_size, + rank, + device_slot_selection, + slot_bias, + threads, + row_elems, + ) + cache_key = ( + "bf16", + "all_reduce_pull", + int(world_size), + int(rank), + bool(device_slot_selection), + slot_bias, + int(threads), + int(row_elems), + ) + raise_if_kernel_resolution_frozen( + "cute.compile", target=launch, cache_key=cache_key + ) + raw = b12x_compile( + launch, + make_ptr(cutlass.Uint32, 16, cute.AddressSpace.gmem, assumed_align=16), + *( + make_ptr( + cutlass.Uint32, + 16, + cute.AddressSpace.gmem, + assumed_align=16, + ) + for _ in range(8) + ), + *( + make_ptr( + cutlass.Uint32, + 16, + cute.AddressSpace.gmem, + assumed_align=4, + ) + for _ in range(8) + ), + make_ptr(cutlass.Uint32, 16, cute.AddressSpace.gmem, assumed_align=16), + 0, + 1, + 1, + 1, + 1, + current_cuda_stream(), + compile_spec=KernelCompileSpec.from_key( + "comm.pcie.twoshot_bf16.all_reduce_pull", + 1, + cache_key, + ), + ) + + def run( + payload_address: int, + staging_addresses: Sequence[int], + signal_addresses: Sequence[int], + output_address: int, + rank: int, + reduced_offset: int, + slot_bytes: int, + rows_per_rank: int, + grid_x: int, + ) -> None: + if len(staging_addresses) != 8 or len(signal_addresses) != 8: + raise ValueError("two-shot scalar pointer ABI requires eight peers") + raw_args = ( + make_ptr( + cutlass.Uint32, + payload_address, + cute.AddressSpace.gmem, + assumed_align=16, + ), + *( + make_ptr( + cutlass.Uint32, + address, + cute.AddressSpace.gmem, + assumed_align=16, + ) + for address in staging_addresses + ), + *( + make_ptr( + cutlass.Uint32, + address, + cute.AddressSpace.gmem, + assumed_align=4, + ) + for address in signal_addresses + ), + make_ptr( + cutlass.Uint32, + output_address, + cute.AddressSpace.gmem, + assumed_align=16, + ), + rank, + reduced_offset, + slot_bytes, + rows_per_rank, + grid_x, + current_cuda_stream(), + ) + raw(*raw_args) + + _PREPARED_BF16_LAUNCHERS.add(process_key) + return run + + +def is_twoshot_bf16_allreduce_launcher_prepared( + world_size: int, + rank: int, + device_slot_selection: bool, + slot_bias: int, + threads: int, + row_elems: int, + device_index: int, +) -> bool: + return ( + _bf16_process_key( + "all_reduce_pull", + world_size, + rank, + device_slot_selection, + slot_bias, + threads, + row_elems, + device_index, + ) + in _PREPARED_BF16_LAUNCHERS + ) + + +__all__ = [ + "get_twoshot_bf16_launcher", + "is_twoshot_bf16_launcher_prepared", + "get_twoshot_bf16_allreduce_launcher", + "is_twoshot_bf16_allreduce_launcher_prepared", +] diff --git a/b12x/comm/pcie/pcie_twoshot_bf16.py b/b12x/comm/pcie/pcie_twoshot_bf16.py new file mode 100644 index 000000000..6f84adebd --- /dev/null +++ b/b12x/comm/pcie/pcie_twoshot_bf16.py @@ -0,0 +1,707 @@ +"""Lossless BF16 PCIe two-shot all-reduce runtime (reduce_scatter + all_gather). + +Host-side twin of :mod:`pcie_twoshot` without the fp8 wire codec: payloads +travel as bf16 packs, are accumulated in fp32 in a fixed rank order and +rounded once. Intended for TP decode all-reduces above the one-shot +ceiling (tens of KB) and below the DMA ring floor (MB), where NCCL ring is +the incumbent. Graph capture follows the two-shot contract: enter +``runtime.capture()`` around ``torch.cuda.graph`` and keep replays of graphs +captured from one instance serialized. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Optional, Sequence + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from ._cuda_ipc import CudaRTLibrary +from ._twoshot_bf16_cute import ( + get_twoshot_bf16_allreduce_launcher, + get_twoshot_bf16_launcher, + is_twoshot_bf16_allreduce_launcher_prepared, + is_twoshot_bf16_launcher_prepared, +) +from .pcie_oneshot import ( + _ABANDONED_PCIE_RUNTIME_QUARANTINE, + IPC_SLAB_ALIGNMENT, + PCIeOneshotAllReduce, + _finish_collective_runtime_setup, + _raise_local_cleanup_errors, + _align_up, + _coordinated_close_channels, + _cuda_device_index, + _device_guard, + _is_current_stream_capturing, + _normalize_device, + _OwnedSharedBuffer, + _require_collective_contract, + _require_full_grid_residency, + _run_collective_preallocation_setup, +) +from .pcie_twoshot import ( + SUPPORTED_WORLD_SIZES, + TWOSHOT_REQUIRED_SMS, + _MAX_BLOCKS, + _SIGNAL_BYTES, + _pad_scalar_peer_ptrs, +) + +_PACK_ELEMS = 8 + + +@dataclass(frozen=True) +class _TwoShotBf16Layout: + signal_bytes: int + pack_stride: int + reduced_offset: int + slot_bytes: int + slab_bytes: int + + +def _make_layout(max_rows: int, row_elems: int, world_size: int) -> _TwoShotBf16Layout: + if world_size not in SUPPORTED_WORLD_SIZES: + raise ValueError(f"unsupported world size {world_size}") + if max_rows <= 0 or max_rows % world_size != 0: + raise ValueError("max_rows must be positive and divisible by world size") + if row_elems <= 0 or row_elems % _PACK_ELEMS != 0: + raise ValueError("row_elems must be a positive multiple of 8") + max_rows_per_rank = max_rows // world_size + packs_per_row = row_elems // _PACK_ELEMS + pack_stride = _align_up(max_rows_per_rank * packs_per_row, 16) + payload_bytes = world_size * pack_stride * 16 + # The pull all-reduce keeps one reduced quarter per slot after the + # full staged payload. + reduced_offset = _align_up(payload_bytes, IPC_SLAB_ALIGNMENT) + slot_bytes = _align_up(reduced_offset + pack_stride * 16, IPC_SLAB_ALIGNMENT) + signal_bytes = _align_up(_SIGNAL_BYTES, IPC_SLAB_ALIGNMENT) + return _TwoShotBf16Layout( + signal_bytes=signal_bytes, + pack_stride=pack_stride, + reduced_offset=reduced_offset, + slot_bytes=slot_bytes, + slab_bytes=signal_bytes + 2 * slot_bytes, + ) + + +class PCIeTwoShotBF16: + """Two-shot lossless bf16 reduce_scatter / all_gather / all_reduce runtime.""" + + def __init__(self, *args, **kwargs) -> None: + raise RuntimeError("use PCIeTwoShotBF16.from_exchange_group()") + + @classmethod + def _from_prepared_factory( + cls, + *, + rank: int, + world_size: int, + device: torch.device, + signal_ptrs: Sequence[int], + staging_ptrs: Sequence[Sequence[int]], + owned_buffers: Sequence[_OwnedSharedBuffer], + ipc: CudaRTLibrary, + exchange_group: ProcessGroup, + max_rows: int, + row_elems: int, + pack_stride: int, + reduced_offset: int, + slot_bytes: int, + ) -> "PCIeTwoShotBF16": + self = object.__new__(cls) + self.rank = rank + self.world_size = world_size + self.device = _normalize_device(device) + self.exchange_group = exchange_group + self._signal_ptrs = tuple(int(pointer) for pointer in signal_ptrs) + self._staging_ptrs = tuple( + tuple(int(pointer) for pointer in slot) for slot in staging_ptrs + ) + if len(self._signal_ptrs) != 8 or ( + len(self._staging_ptrs) != 2 + or any(len(slot) != 8 for slot in self._staging_ptrs) + ): + raise ValueError("two-shot scalar pointer ABI requires 8 peers and 2 slots") + self._owned_buffers = list(owned_buffers) + self._ipc = ipc + self.max_rows = max_rows + self.row_elems = row_elems + self._pack_stride = int(pack_stride) + self._reduced_offset = int(reduced_offset) + self._slot_bytes = int(slot_bytes) + self._slot = 0 + self._device_slot_selection = False + self._device_slot_bias = 0 + self._capture_context_depth = 0 + self._closed = False + self._ipc_imports_closed = False + self._ipc_exports_freed = False + self._coordinated_close_complete = False + self._closed_ipc_import_indices: set[tuple[int, int]] = set() + # Persistent reduce-scatter shard for all_reduce (stable address for + # CUDA graphs; sized for the largest supported message). + self._shard = torch.empty( + max_rows // world_size, row_elems, dtype=torch.bfloat16, device=self.device + ) + return self + + @classmethod + def from_exchange_group( + cls, + *, + exchange_group: ProcessGroup, + device: torch.device | int | str, + max_rows: int, + row_elems: int, + ) -> "PCIeTwoShotBF16": + rank = dist.get_rank(group=exchange_group) + world_size = dist.get_world_size(group=exchange_group) + + def validate_factory_arguments(): + device_obj = _normalize_device(device) + normalized_max_rows = int(max_rows) + normalized_row_elems = int(row_elems) + if world_size not in SUPPORTED_WORLD_SIZES: + raise ValueError(f"unsupported world size {world_size}") + if device_obj.type != "cuda": + raise ValueError("PCIe twoshot requires a CUDA device") + if normalized_max_rows <= 0: + raise ValueError("max_rows must be positive") + if normalized_row_elems <= 0 or normalized_row_elems % _PACK_ELEMS != 0: + raise ValueError("row_elems must be a positive multiple of 8") + if normalized_max_rows % world_size != 0: + raise ValueError("max_rows must be divisible by world size") + return device_obj, normalized_max_rows, normalized_row_elems + + device_obj, max_rows, row_elems = _run_collective_preallocation_setup( + owner="PCIe twoshot-bf16 argument validation", + exchange_group=exchange_group, + setup=validate_factory_arguments, + ) + _require_full_grid_residency( + owner="PCIe twoshot-bf16", + required_sms=TWOSHOT_REQUIRED_SMS, + device=device_obj, + exchange_group=exchange_group, + ) + + def prepare(): + prepared_ipc = CudaRTLibrary() + prepared_ipc.cudaSetDevice(_cuda_device_index(device_obj)) + return prepared_ipc, _make_layout(max_rows, row_elems, world_size) + + ipc, layout = _run_collective_preallocation_setup( + owner="PCIe twoshot-bf16", + exchange_group=exchange_group, + setup=prepare, + ) + _require_collective_contract( + owner="PCIe twoshot-bf16 channel layout", + exchange_group=exchange_group, + contract=(int(max_rows), int(row_elems), layout), + ) + shared = PCIeOneshotAllReduce._allocate_shared_buffer( + exchange_group, + layout.slab_bytes, + zero_fill=True, + ipc=ipc, + ) + peer_ptrs = list(shared.peer_ptrs) + signal_ptrs = _pad_scalar_peer_ptrs(peer_ptrs, rank=rank, world_size=world_size) + staging_ptrs = ( + _pad_scalar_peer_ptrs( + [p + layout.signal_bytes for p in peer_ptrs], + rank=rank, + world_size=world_size, + ), + _pad_scalar_peer_ptrs( + [p + layout.signal_bytes + layout.slot_bytes for p in peer_ptrs], + rank=rank, + world_size=world_size, + ), + ) + runtime: Optional[PCIeTwoShotBF16] = None + init_error: BaseException | None = None + try: + runtime = cls._from_prepared_factory( + rank=rank, + world_size=world_size, + device=device_obj, + signal_ptrs=signal_ptrs, + staging_ptrs=staging_ptrs, + owned_buffers=[shared], + ipc=ipc, + exchange_group=exchange_group, + max_rows=max_rows, + row_elems=row_elems, + pack_stride=layout.pack_stride, + reduced_offset=layout.reduced_offset, + slot_bytes=layout.slot_bytes, + ) + except Exception as exc: + init_error = exc + + def detach_shared_ownership() -> None: + if runtime is not None: + runtime._owned_buffers.clear() + + _finish_collective_runtime_setup( + owner="PCIe twoshot-bf16", + exchange_group=exchange_group, + ipc=ipc, + shared=shared, + local_error=init_error, + detach_shared_ownership=detach_shared_ownership, + ) + assert runtime is not None + return runtime + + # ---- checks --------------------------------------------------------- + + def _check(self, payload: torch.Tensor, rows: int) -> None: + if self._closed: + raise RuntimeError("PCIeTwoShotBF16 is closed") + if payload.shape != (rows, self.row_elems): + raise ValueError( + f"payload shape {tuple(payload.shape)} != ({rows}, {self.row_elems})" + ) + if rows > self.max_rows: + raise ValueError("pcie_twoshot_bf16 staging capacity exceeded") + if payload.device != self.device: + raise ValueError("payload must be on the runtime CUDA device") + if payload.dtype != torch.bfloat16: + raise TypeError("payload must be bfloat16") + if not payload.is_contiguous(): + raise ValueError("payload must be contiguous") + + def _device_index(self) -> int: + return ( + self.device.index + if self.device.index is not None + else torch.cuda.current_device() + ) + + def accepts(self, inp: torch.Tensor) -> bool: + """True when ``all_reduce`` can serve this tensor.""" + if self._closed or inp.dtype != torch.bfloat16 or not inp.is_contiguous(): + return False + if inp.device != self.device: + return False + numel = inp.numel() + if numel == 0 or numel % (self.row_elems * self.world_size) != 0: + return False + return numel // self.row_elems <= self.max_rows + + # ---- graph plumbing --------------------------------------------------- + + def prepare_graph( + self, + *, + operations: Sequence[str] = ("reduce_scatter", "all_gather"), + threads: int = 512, + ) -> None: + if self._closed: + raise RuntimeError("PCIeTwoShotBF16 is closed") + if _is_current_stream_capturing(self.device): + raise RuntimeError( + "prepare_graph() must be called before CUDA graph capture" + ) + threads = int(threads) + if threads <= 0 or threads > 512 or threads % 32 != 0: + raise ValueError("threads must be a warp-aligned value in [32, 512]") + requested = tuple(str(operation) for operation in operations) + device_index = self._device_index() + with torch.cuda.device(self.device): + for operation in dict.fromkeys(requested): + for slot_bias in (0, 1): + get_twoshot_bf16_launcher( + operation, + self.world_size, + self.rank, + True, + slot_bias, + threads, + self.row_elems, + device_index, + ) + for slot_bias in (0, 1): + get_twoshot_bf16_allreduce_launcher( + self.world_size, + self.rank, + True, + slot_bias, + threads, + self.row_elems, + device_index, + ) + + @contextmanager + def capture( + self, + *, + operations: Sequence[str] = ("reduce_scatter", "all_gather"), + threads: int = 512, + ): + if self._capture_context_depth: + raise RuntimeError( + "overlapping PCIe twoshot-bf16 capture contexts are not allowed" + ) + self.prepare_graph(operations=operations, threads=threads) + self._capture_context_depth = 1 + try: + yield self + finally: + self._capture_context_depth = 0 + + # ---- launch ----------------------------------------------------------- + + def _launch( + self, + operation: str, + payload: torch.Tensor, + out: torch.Tensor, + *, + rows_per_rank: int, + threads: int, + block_limit: int, + ) -> None: + threads = int(threads) + if threads <= 0 or threads > 512 or threads % 32 != 0: + raise ValueError("threads must be a warp-aligned value in [32, 512]") + shard_packs = rows_per_rank * (self.row_elems // _PACK_ELEMS) + if shard_packs > self._pack_stride: + raise ValueError("pcie_twoshot_bf16 staging capacity exceeded") + if block_limit <= 0 or block_limit > _MAX_BLOCKS: + raise ValueError(f"block_limit must be in [1, {_MAX_BLOCKS}]") + blocks = max(1, min(int(block_limit), (shard_packs + threads - 1) // threads)) + capturing = _is_current_stream_capturing(self.device) + device_index = self._device_index() + if capturing: + if self._capture_context_depth <= 0: + raise RuntimeError( + "cold PCIe twoshot-bf16 CUDA graph capture is not allowed; " + "enter runtime.capture() before torch.cuda.graph()" + ) + graph_slot_bias = ( + self._device_slot_bias + if self._device_slot_selection + else self._slot & 1 + ) + if not is_twoshot_bf16_launcher_prepared( + operation, + self.world_size, + self.rank, + True, + graph_slot_bias, + threads, + self.row_elems, + device_index, + ): + raise RuntimeError( + "cold PCIe twoshot-bf16 CUDA graph capture is not allowed; " + "enter runtime.capture() before torch.cuda.graph()" + ) + if capturing and not self._device_slot_selection: + self._device_slot_bias = self._slot & 1 + self._device_slot_selection = True + if self._device_slot_selection: + slot = 0 + else: + slot = self._slot % 2 + self._slot += 1 + with torch.cuda.device(self.device): + launcher = get_twoshot_bf16_launcher( + operation, + self.world_size, + self.rank, + self._device_slot_selection, + self._device_slot_bias, + threads, + self.row_elems, + device_index, + ) + if not self._device_slot_selection: + for slot_bias in (0, 1): + get_twoshot_bf16_launcher( + operation, + self.world_size, + self.rank, + True, + slot_bias, + threads, + self.row_elems, + device_index, + ) + launcher( + payload.data_ptr(), + self._staging_ptrs[slot], + self._signal_ptrs, + out.data_ptr(), + self.rank, + self._pack_stride, + self._slot_bytes, + rows_per_rank, + blocks, + ) + + # ---- public collectives --------------------------------------------- + + def reduce_scatter( + self, + payload: torch.Tensor, + out: Optional[torch.Tensor] = None, + *, + threads: int = 512, + block_limit: int = 64, + ) -> torch.Tensor: + with _device_guard(self.device): + rows = payload.shape[0] + self._check(payload, rows) + if rows % self.world_size != 0: + raise ValueError("rows must be divisible by world size") + if out is None: + out = torch.empty( + rows // self.world_size, + self.row_elems, + dtype=torch.bfloat16, + device=self.device, + ) + self._launch( + "reduce_scatter", + payload, + out, + rows_per_rank=rows // self.world_size, + threads=threads, + block_limit=block_limit, + ) + return out + + def all_gather( + self, + payload: torch.Tensor, + out: Optional[torch.Tensor] = None, + *, + threads: int = 512, + block_limit: int = 64, + ) -> torch.Tensor: + with _device_guard(self.device): + rows = payload.shape[0] + self._check(payload, rows) + if out is None: + out = torch.empty( + rows * self.world_size, + self.row_elems, + dtype=torch.bfloat16, + device=self.device, + ) + self._launch( + "all_gather", + payload, + out, + rows_per_rank=rows, + threads=threads, + block_limit=block_limit, + ) + return out + + def _launch_pull_all_reduce( + self, + payload: torch.Tensor, + out: torch.Tensor, + *, + rows_per_rank: int, + threads: int, + block_limit: int, + ) -> None: + threads = int(threads) + if threads <= 0 or threads > 512 or threads % 32 != 0: + raise ValueError("threads must be a warp-aligned value in [32, 512]") + quarter_packs = rows_per_rank * (self.row_elems // _PACK_ELEMS) + if quarter_packs > self._pack_stride: + raise ValueError("pcie_twoshot_bf16 staging capacity exceeded") + if block_limit <= 0 or block_limit > _MAX_BLOCKS: + raise ValueError(f"block_limit must be in [1, {_MAX_BLOCKS}]") + blocks = max(1, min(int(block_limit), (quarter_packs + threads - 1) // threads)) + capturing = _is_current_stream_capturing(self.device) + device_index = self._device_index() + if capturing: + if self._capture_context_depth <= 0: + raise RuntimeError( + "cold PCIe twoshot-bf16 CUDA graph capture is not allowed; " + "enter runtime.capture() before torch.cuda.graph()" + ) + graph_slot_bias = ( + self._device_slot_bias + if self._device_slot_selection + else self._slot & 1 + ) + if not is_twoshot_bf16_allreduce_launcher_prepared( + self.world_size, + self.rank, + True, + graph_slot_bias, + threads, + self.row_elems, + device_index, + ): + raise RuntimeError( + "cold PCIe twoshot-bf16 CUDA graph capture is not allowed; " + "enter runtime.capture() before torch.cuda.graph()" + ) + if capturing and not self._device_slot_selection: + self._device_slot_bias = self._slot & 1 + self._device_slot_selection = True + if self._device_slot_selection: + slot = 0 + else: + slot = self._slot % 2 + self._slot += 1 + with torch.cuda.device(self.device): + launcher = get_twoshot_bf16_allreduce_launcher( + self.world_size, + self.rank, + self._device_slot_selection, + self._device_slot_bias, + threads, + self.row_elems, + device_index, + ) + if not self._device_slot_selection: + for slot_bias in (0, 1): + get_twoshot_bf16_allreduce_launcher( + self.world_size, + self.rank, + True, + slot_bias, + threads, + self.row_elems, + device_index, + ) + launcher( + payload.data_ptr(), + self._staging_ptrs[slot], + self._signal_ptrs, + out.data_ptr(), + self.rank, + self._reduced_offset, + self._slot_bytes, + rows_per_rank, + blocks, + ) + + def all_reduce( + self, + inp: torch.Tensor, + out: Optional[torch.Tensor] = None, + *, + threads: int = 512, + block_limit: int = 64, + ) -> torch.Tensor: + """Lossless bf16 all-reduce: one pull-based launch (2 barriers).""" + if not self.accepts(inp): + raise ValueError("input not accepted by PCIeTwoShotBF16.all_reduce") + rows = inp.numel() // self.row_elems + payload = inp.view(rows, self.row_elems) + with _device_guard(self.device): + if out is None: + out = torch.empty_like(inp) + out_view = out.view(rows, self.row_elems) + self._launch_pull_all_reduce( + payload, + out_view, + rows_per_rank=rows // self.world_size, + threads=threads, + block_limit=block_limit, + ) + return out + + # ---- teardown (mirrors pcie_twoshot) ----------------------------------- + + def _closed_import_indices(self) -> set[tuple[int, int]]: + closed = getattr(self, "_closed_ipc_import_indices", None) + if closed is None: + closed = set() + self._closed_ipc_import_indices = closed + return closed + + def _all_python_ipc_imports_closed(self, closed: set[tuple[int, int]]) -> bool: + return all( + (buffer_index, remote_index) in closed + for buffer_index, shared in enumerate(self._owned_buffers) + for remote_index, _ in enumerate(shared.remote_ptrs) + ) + + def _close_ipc_imports_strict(self) -> None: + if self._ipc_imports_closed: + return + self._closed = True + failures: list[tuple[str, Exception]] = [] + closed = self._closed_import_indices() + for buffer_index, shared in enumerate(self._owned_buffers): + for remote_index, ptr in enumerate(shared.remote_ptrs): + key = (buffer_index, remote_index) + if key in closed: + continue + try: + self._ipc.cudaIpcCloseMemHandle(ptr) + except Exception as exc: + failures.append((f"CUDA IPC import {ptr}", exc)) + else: + closed.add(key) + if not failures and self._all_python_ipc_imports_closed(closed): + self._ipc_imports_closed = True + if failures: + _raise_local_cleanup_errors( + "PCIe twoshot-bf16", "IPC import close", failures + ) + + def _free_ipc_exports_strict(self) -> None: + if self._ipc_exports_freed: + return + self._close_ipc_imports_strict() + failures: list[tuple[str, Exception]] = [] + remaining = [] + for shared in self._owned_buffers: + try: + self._ipc.cudaFree(shared.local_ptr) + except Exception as exc: + remaining.append(shared) + failures.append((f"CUDA IPC export {shared.local_ptr}", exc)) + self._owned_buffers = remaining + if not remaining: + self._ipc_exports_freed = True + if failures: + _raise_local_cleanup_errors( + "PCIe twoshot-bf16", "IPC export free", failures + ) + + def close(self) -> None: + if getattr(self, "_coordinated_close_complete", False): + return + _coordinated_close_channels( + (self,), + exchange_group=self.exchange_group, + device=self.device, + ) + + def __enter__(self) -> "PCIeTwoShotBF16": + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + self.close() + + def __del__( + self, + _quarantine: dict[int, object] = _ABANDONED_PCIE_RUNTIME_QUARANTINE, + ) -> None: + if getattr(self, "_coordinated_close_complete", False): + return + if getattr(self, "_owned_buffers", ()): + _quarantine[id(self)] = self + + +__all__ = ["PCIeTwoShotBF16"] diff --git a/tests/comm/test_pcie_twoshot_bf16.py b/tests/comm/test_pcie_twoshot_bf16.py new file mode 100644 index 000000000..888fa4a5a --- /dev/null +++ b/tests/comm/test_pcie_twoshot_bf16.py @@ -0,0 +1,152 @@ +"""Correctness for the lossless BF16 PCIe two-shot all-reduce. + +Run with torchrun on 2, 4 or 8 GPUs: + + python -m torch.distributed.run --nproc-per-node=4 \ + tests/comm/test_pcie_twoshot_bf16.py + +Every all-reduce is checked against the exact fp32 sum of the gathered +inputs: the kernel accumulates in fp32 in a fixed rank order and rounds once, +so each output element must lie within one bf16 rounding of the exact sum, +and repeated calls (eager and graph replay) must be bitwise identical. +""" + +from __future__ import annotations + +import os + +import torch +import torch.distributed as dist + +from b12x.comm.pcie.pcie_twoshot_bf16 import PCIeTwoShotBF16, _make_layout + +ROW_ELEMS = int(os.getenv("B12X_TEST_TWOSHOT_BF16_ROW_ELEMS", "4096")) +MAX_ROWS = int(os.getenv("B12X_TEST_TWOSHOT_BF16_MAX_ROWS", "512")) +ROWS = (8, 16, 32, 64, 96, 128, 192, 256) + + +def test_layout_scales_with_rows_and_ranks() -> None: + base = _make_layout(64, ROW_ELEMS, 4) + assert base.pack_stride > 0 and base.slot_bytes > 0 + assert base.reduced_offset >= 0 and base.slab_bytes >= base.slot_bytes + taller = _make_layout(128, ROW_ELEMS, 4) + assert taller.slab_bytes > base.slab_bytes + wider = _make_layout(64, 2 * ROW_ELEMS, 4) + assert wider.slab_bytes > base.slab_bytes + + +def _payload(seed: int, rows: int, device: torch.device) -> torch.Tensor: + gen = torch.Generator(device="cpu").manual_seed(seed) + x = torch.randn(rows, ROW_ELEMS, generator=gen, dtype=torch.float32) * 3.0 + return x.to(device=device, dtype=torch.bfloat16) + + +def _exact_sum(x: torch.Tensor, world: int) -> torch.Tensor: + gathered = [torch.empty_like(x) for _ in range(world)] + dist.all_gather(gathered, x) + return sum(g.float() for g in gathered) + + +def _assert_one_rounding(out: torch.Tensor, exact: torch.Tensor) -> None: + err = (out.float() - exact).abs() + bound = exact.abs() * 2.0**-8 + 1e-5 + worst = (err - bound).max().item() + assert bool((err <= bound).all()), ( + f"output exceeds one bf16 rounding by {worst:.3e}" + ) + + +def _check_all_reduce( + pool: PCIeTwoShotBF16, rank: int, world: int, rows: int, step: int +) -> None: + x = _payload(1000 * step + 7 * rows + rank, rows, pool.device) + assert pool.accepts(x) + exact = _exact_sum(x, world) + out = pool.all_reduce(x) + torch.cuda.synchronize() + assert out.shape == x.shape and out.dtype == torch.bfloat16 + _assert_one_rounding(out, exact) + again = pool.all_reduce(x) + torch.cuda.synchronize() + assert torch.equal(out, again), "two-shot all-reduce must be deterministic" + # The bf16 NCCL ring rounds after every hop; the two-shot rounds once. + ref = x.clone() + dist.all_reduce(ref) + assert (out.float() - exact).abs().max() <= (ref.float() - exact).abs().max() + 1e-5 + + +def _check_graph_capture(pool: PCIeTwoShotBF16, rank: int, world: int) -> None: + rows = 64 + static = _payload(4242 + rank, rows, pool.device) + exact = _exact_sum(static, world) + eager = pool.all_reduce(static) + torch.cuda.synchronize() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with pool.capture(): + with torch.cuda.stream(stream): + for _ in range(3): + pool.all_reduce(static) + torch.cuda.current_stream().wait_stream(stream) + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = pool.all_reduce(static) + torch.cuda.synchronize() + dist.barrier() + for _ in range(3): + graph.replay() + torch.cuda.synchronize() + dist.barrier() + assert torch.equal(captured, eager), "graph replay must match the eager result" + _assert_one_rounding(captured, exact) + + +def _check_rejects_unsupported(pool: PCIeTwoShotBF16, world: int) -> None: + device = pool.device + assert not pool.accepts( + torch.empty(MAX_ROWS + world, ROW_ELEMS, dtype=torch.bfloat16, device=device) + ) + assert not pool.accepts( + torch.empty(world, ROW_ELEMS, dtype=torch.float16, device=device) + ) + assert not pool.accepts( + torch.empty(world, ROW_ELEMS + 8, dtype=torch.bfloat16, device=device) + ) + if world > 1: + assert not pool.accepts( + torch.empty(world + 1, ROW_ELEMS, dtype=torch.bfloat16, device=device) + ) + + +def main() -> None: + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl") + device = torch.device("cuda", local_rank) + + pool = PCIeTwoShotBF16.from_exchange_group( + exchange_group=dist.group.WORLD, + device=device, + max_rows=MAX_ROWS, + row_elems=ROW_ELEMS, + ) + pool.prepare_graph() + _check_rejects_unsupported(pool, world) + for step in range(3): # exercises the double-buffered staging slots + for rows in ROWS: + if rows % world == 0 and rows <= MAX_ROWS: + _check_all_reduce(pool, rank, world, rows, step) + _check_graph_capture(pool, rank, world) + dist.barrier() + if rank == 0: + print(f"pcie_twoshot_bf16 correctness OK ({world} ranks, rows {ROWS})") + pool.close() + dist.destroy_process_group() + + +if __name__ == "__main__": + main()