From 65ebe0d19b7e004d27141baad873bad24a477c92 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Thu, 30 Jul 2026 20:01:03 +0000 Subject: [PATCH 01/11] fix(pcie): prevent DCP staging slot overflow --- sparkinfer/comm/pcie/pcie_dcp_a2a.cu | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/sparkinfer/comm/pcie/pcie_dcp_a2a.cu b/sparkinfer/comm/pcie/pcie_dcp_a2a.cu index 2458198a8..5f372b54a 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_a2a.cu +++ b/sparkinfer/comm/pcie/pcie_dcp_a2a.cu @@ -343,7 +343,16 @@ public: int64_t output_capacity_elems_; int64_t lse_offset_; int64_t lse_capacity_; - int slot_ = 0; + uint32_t next_slot_ = 0; + + int take_slot() { + // Host-side selection executes once during CUDA graph capture. Graph + // replay intentionally reuses that capture-stable address; eager calls + // toggle slabs without an overflowing counter. + const int slot = static_cast(next_slot_); + next_slot_ ^= uint32_t{1}; + return slot; + } PCIeDCPA2A(Signal **signals, const std::vector> &staging, @@ -396,7 +405,7 @@ public: // Staging happens inside the kernel (warp-per-row, before the start // barrier); no host staging memcpys are issued. - const int slot = slot_++ % 2; + const int slot = take_slot(); const int heads_per_rank = total_heads / world_size_; const int rows = batch * heads_per_rank; const int warps_per_block = threads / 32; @@ -457,7 +466,7 @@ public: // Staging happens inside the kernel (warp-per-row, before the start // barrier); no host staging memcpy is issued. - const int slot = slot_++ % 2; + const int slot = take_slot(); const int rows = batch * total_heads; const int warps_per_block = threads / 32; const int blocks = std::max( From 941c8c8ec6af00b1841dee6c4327f08a58194c7c Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Thu, 30 Jul 2026 20:27:43 +0000 Subject: [PATCH 02/11] fix(pcie): retire capture channel aliases --- sparkinfer/comm/pcie/pcie_dcp_a2a.py | 15 +++++++++ tests/comm/test_pcie_dcp_a2a.py | 48 +++++++++++++++++++++++++--- 2 files changed, 59 insertions(+), 4 deletions(-) diff --git a/sparkinfer/comm/pcie/pcie_dcp_a2a.py b/sparkinfer/comm/pcie/pcie_dcp_a2a.py index 49f41ce26..1fe8a4067 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_a2a.py +++ b/sparkinfer/comm/pcie/pcie_dcp_a2a.py @@ -766,6 +766,7 @@ def all_gather_heads( @contextmanager def capture(self, stream: object = None): """Bind nested CUDA captures to the enclosing stream's channel.""" + previous_channels: Optional[dict[int, PCIeDCPA2A]] = None if self.single_channel: channel = self.for_stream(stream) else: @@ -774,6 +775,7 @@ def capture(self, stream: object = None): "PCIe DCP A2A capture context must be entered before CUDA " "graph capture starts" ) + previous_channels = dict(self._channels) stream_key = _current_stream_key(self.device, stream) key = 0 if stream_key is None else int(stream_key) # Keep a graph-owned channel even when CUDA recycles the enclosing @@ -787,6 +789,19 @@ def capture(self, stream: object = None): popped = self._capture_channel_stack.pop() if popped is not channel: raise RuntimeError("PCIe DCP A2A capture channel stack corrupted") + if previous_channels is not None: + # Captured graph nodes retain the channel through + # ``_all_channels``. Restore eager mappings and discard every + # nested capture-stream alias so a recycled stream key cannot + # hand this graph-owned channel to a later unwrapped capture. + for key, mapped in tuple(self._channels.items()): + if mapped is not channel: + continue + previous = previous_channels.get(key) + if previous is None: + del self._channels[key] + else: + self._channels[key] = previous def close(self) -> None: if self._closed: diff --git a/tests/comm/test_pcie_dcp_a2a.py b/tests/comm/test_pcie_dcp_a2a.py index 50088b56a..4f1e01cfc 100644 --- a/tests/comm/test_pcie_dcp_a2a.py +++ b/tests/comm/test_pcie_dcp_a2a.py @@ -269,10 +269,16 @@ def make_channel(stream_key): capturing[0] = False assert target_channel is not draft_channel - assert pool._channels[70] is target_channel - assert pool._channels[80] is draft_channel + assert pool._channels == {} + assert target_channel in pool._all_channels + assert draft_channel in pool._all_channels assert [entry[0] for entry in created] == [7, 8] + capturing[0] = True + current_stream[0] = 70 + with pytest.raises(RuntimeError, match="no channel during CUDA graph capture"): + pool.for_stream() + def test_pool_isolates_reused_capture_stream_keys(monkeypatch): created = [] @@ -317,8 +323,7 @@ def make_channel(stream_key): capturing[0] = False assert target_channel is not draft_channel - assert pool._channels[7] is draft_channel - assert pool._channels[70] is draft_channel + assert pool._channels == {} assert target_channel in pool._all_channels assert draft_channel in pool._all_channels assert [entry[0] for entry in created] == [7, 7] @@ -328,6 +333,41 @@ def make_channel(stream_key): assert draft_channel._ext.dispose_calls == [1234] +def test_pool_restores_eager_mapping_after_capture(monkeypatch): + current_stream = [7] + capturing = [False] + pool = PCIeDCPA2APool( + rank=0, + world_size=2, + device=torch.device("cpu"), + max_batch_size=4, + total_heads=32, + head_dim=64, + channel_factory=lambda stream_key: _make_runtime(), + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_dcp_a2a._current_stream_key", + lambda device, stream=None: ( + current_stream[0] if stream is None else int(stream) + ), + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_dcp_a2a._is_current_stream_capturing", + lambda device: capturing[0], + ) + + eager_channel = pool.for_stream() + with pool.capture(7) as graph_channel: + capturing[0] = True + current_stream[0] = 70 + assert pool.for_stream() is graph_channel + capturing[0] = False + + assert graph_channel is not eager_channel + assert pool._channels == {7: eager_channel} + assert graph_channel in pool._all_channels + + def test_pool_rolls_back_throwaway_capture_channels(monkeypatch): created = [] From 3e5ef42468db9f23044dbf87ece3b1b8ccf4c564 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Thu, 30 Jul 2026 22:27:34 +0000 Subject: [PATCH 03/11] fix(pcie): derive DCP staging slot on device --- sparkinfer/comm/pcie/pcie_dcp_a2a.cu | 55 +++++++++---- tests/comm/test_pcie_dcp_a2a.py | 21 +++++ tests/comm/test_pcie_dcp_a2a_gpu.py | 115 +++++++++++++++++++++++++++ 3 files changed, 175 insertions(+), 16 deletions(-) diff --git a/sparkinfer/comm/pcie/pcie_dcp_a2a.cu b/sparkinfer/comm/pcie/pcie_dcp_a2a.cu index 2458198a8..7e153981a 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_a2a.cu +++ b/sparkinfer/comm/pcie/pcie_dcp_a2a.cu @@ -78,6 +78,10 @@ struct RankStaging { void *ptrs[kMaxRanks]; }; +struct DoubleStaging { + RankStaging slots[2]; +}; + template struct __align__(16) Pack { T values[8]; }; @@ -149,7 +153,7 @@ template __global__ void __launch_bounds__(512, 1) dcp_lse_reduce_kernel(const T *__restrict__ local_output, const float *__restrict__ local_lse, - RankStaging staging, int64_t lse_offset, + DoubleStaging staging_options, int64_t lse_offset, RankSignals signals, Signal *self, T *__restrict__ output, int rank, int batch, int total_heads, int head_dim, @@ -169,6 +173,17 @@ __global__ void __launch_bounds__(512, 1) const auto *local_packs = reinterpret_cast *>(local_output); auto *output_packs = reinterpret_cast *>(output); + __shared__ RankStaging staging; + if (threadIdx.x == 0) { + const int slot = + int(load_flag(&self->self_counter[blockIdx.x][0]) & FlagType{1}); +#pragma unroll + for (int peer = 0; peer < world_size; ++peer) { + staging.ptrs[peer] = staging_options.slots[slot].ptrs[peer]; + } + } + __syncthreads(); + auto *staging_out = reinterpret_cast *>(staging.ptrs[rank]); auto *staging_lse = reinterpret_cast( reinterpret_cast(staging.ptrs[rank]) + lse_offset); @@ -286,7 +301,7 @@ __global__ void __launch_bounds__(512, 1) template __global__ void __launch_bounds__(512, 1) all_gather_heads_kernel(const T *__restrict__ local_input, - RankStaging staging, RankSignals signals, + DoubleStaging staging_options, RankSignals signals, Signal *self, T *__restrict__ output, int rank, int batch, int local_heads, int head_dim) { constexpr int kPackElems = 8; @@ -299,6 +314,17 @@ __global__ void __launch_bounds__(512, 1) const int warp_stride = gridDim.x * warps_per_block; const auto *local_packs = reinterpret_cast *>(local_input); + __shared__ RankStaging staging; + if (threadIdx.x == 0) { + const int slot = + int(load_flag(&self->self_counter[blockIdx.x][0]) & FlagType{1}); +#pragma unroll + for (int peer = 0; peer < world_size; ++peer) { + staging.ptrs[peer] = staging_options.slots[slot].ptrs[peer]; + } + } + __syncthreads(); + auto *staging_out = reinterpret_cast *>(staging.ptrs[rank]); for (int row = warp_first; row < rows; row += warp_stride) { const int batch_index = row / total_heads; @@ -339,11 +365,10 @@ public: int world_size_; RankSignals signals_{}; Signal *self_signal_; - RankStaging staging_[2]{}; + DoubleStaging staging_{}; int64_t output_capacity_elems_; int64_t lse_offset_; int64_t lse_capacity_; - int slot_ = 0; PCIeDCPA2A(Signal **signals, const std::vector> &staging, @@ -354,8 +379,8 @@ public: lse_capacity_(lse_capacity) { for (int peer = 0; peer < world_size_; ++peer) { signals_.signals[peer] = signals[peer]; - staging_[0].ptrs[peer] = staging[peer][0]; - staging_[1].ptrs[peer] = staging[peer][1]; + staging_.slots[0].ptrs[peer] = staging[peer][0]; + staging_.slots[1].ptrs[peer] = staging[peer][1]; } } @@ -394,9 +419,8 @@ public: throw std::runtime_error("threads must be a multiple of 32"); } - // Staging happens inside the kernel (warp-per-row, before the start - // barrier); no host staging memcpys are issued. - const int slot = slot_++ % 2; + // Both staging slabs are kernel arguments. Each block selects from its + // pre-increment execution counter, so capture/replay cannot bake a slot. const int heads_per_rank = total_heads / world_size_; const int rows = batch * heads_per_rank; const int warps_per_block = threads / 32; @@ -405,8 +429,8 @@ public: #define LAUNCH(world) \ dcp_lse_reduce_kernel<<>>( \ - partial_output, partial_lse, staging_[slot], lse_offset_, signals_, \ - self_signal_, output, rank_, batch, total_heads, head_dim, \ + partial_output, partial_lse, staging_, lse_offset_, signals_, \ + self_signal_, output, rank_, batch, total_heads, head_dim, \ input_stride_batch, input_stride_head, output_stride_batch, \ output_stride_head, natural_log) switch (world_size_) { @@ -455,9 +479,8 @@ public: throw std::runtime_error("threads must be a multiple of 32"); } - // Staging happens inside the kernel (warp-per-row, before the start - // barrier); no host staging memcpy is issued. - const int slot = slot_++ % 2; + // Slot ownership is selected from device execution state in the kernel, + // including for back-to-back eager launches. const int rows = batch * total_heads; const int warps_per_block = threads / 32; const int blocks = std::max( @@ -465,8 +488,8 @@ public: #define LAUNCH(world) \ all_gather_heads_kernel<<>>( \ - local_input, staging_[slot], signals_, self_signal_, output, rank_, \ - batch, local_heads, head_dim) + local_input, staging_, signals_, self_signal_, output, rank_, batch, \ + local_heads, head_dim) switch (world_size_) { case 2: LAUNCH(2); diff --git a/tests/comm/test_pcie_dcp_a2a.py b/tests/comm/test_pcie_dcp_a2a.py index 4f1e01cfc..0fc1dd8f0 100644 --- a/tests/comm/test_pcie_dcp_a2a.py +++ b/tests/comm/test_pcie_dcp_a2a.py @@ -1,5 +1,7 @@ from __future__ import annotations +from pathlib import Path + import pytest import torch @@ -97,6 +99,25 @@ def _make_runtime(ext: _FakeExt | None = None) -> PCIeDCPA2A: ) +def test_cuda_staging_slot_is_selected_from_execution_state(): + source = ( + Path(__file__).parents[2] + / "sparkinfer" + / "comm" + / "pcie" + / "pcie_dcp_a2a.cu" + ).read_text() + + assert "int slot_" not in source + assert "slot_++" not in source + assert ( + source.count("load_flag(&self->self_counter[blockIdx.x][0])") == 2 + ) + assert "select_staging_slot_kernel" not in source + assert source.count("DoubleStaging staging_options") == 2 + assert source.count("__shared__ RankStaging staging;") == 2 + + def test_staging_layout_has_aligned_disjoint_slots(): layout = _staging_layout( signal_bytes=12345, diff --git a/tests/comm/test_pcie_dcp_a2a_gpu.py b/tests/comm/test_pcie_dcp_a2a_gpu.py index 48cc0acf8..332eddbc4 100644 --- a/tests/comm/test_pcie_dcp_a2a_gpu.py +++ b/tests/comm/test_pcie_dcp_a2a_gpu.py @@ -173,6 +173,73 @@ def _check_eager( torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) +def _check_eager_adjacency( + pool: PCIeDCPA2APool, + rank: int, + world_size: int, + device: torch.device, +) -> None: + channel = pool.for_stream() + first_query = _rank_query( + 700, rank, world_size, MAX_BATCH, torch.bfloat16, device + ) + second_query = _rank_query( + 701, rank, world_size, MAX_BATCH, torch.bfloat16, device + ) + partial_output, partial_lse = _rank_inputs( + 702, rank, 1, torch.bfloat16, device + ) + first_gather = torch.empty( + MAX_BATCH, + TOTAL_HEADS, + QUERY_HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + second_gather = torch.empty_like(first_gather) + reduced = torch.empty( + 1, + TOTAL_HEADS // world_size, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + + # Issue large-grid AG -> small-grid RS -> large-grid AG without a host + # synchronization so adjacent eager executions exercise slot advancement. + channel.all_gather_heads(first_query, first_gather) + channel.lse_reduce_scatter(partial_output, partial_lse, reduced) + channel.all_gather_heads(second_query, second_gather) + torch.cuda.synchronize(device) + + expected_first = torch.cat( + [ + _rank_query( + 700, source, world_size, MAX_BATCH, torch.bfloat16, device + ) + for source in range(world_size) + ], + dim=1, + ) + expected_second = torch.cat( + [ + _rank_query( + 701, source, world_size, MAX_BATCH, torch.bfloat16, device + ) + for source in range(world_size) + ], + dim=1, + ) + torch.testing.assert_close(first_gather, expected_first, rtol=0, atol=0) + torch.testing.assert_close(second_gather, expected_second, rtol=0, atol=0) + torch.testing.assert_close( + reduced, + _reference(702, rank, world_size, 1, torch.bfloat16, device), + rtol=2e-2, + atol=2e-2, + ) + + def _check_graph( pool: PCIeDCPA2APool, rank: int, @@ -302,6 +369,52 @@ def _check_graph( for out, reference in zip(gathered_queries, expected_queries, strict=True): torch.testing.assert_close(out, reference, rtol=0, atol=0) + # One captured operation has odd capture-time parity. Adjacent A -> A + # replays must advance the device-owned slot at execution time. + odd_input = torch.empty( + 1, + TOTAL_HEADS // world_size, + QUERY_HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + odd_output = torch.empty( + 1, + TOTAL_HEADS, + QUERY_HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + with torch.cuda.stream(stream): + channel.all_gather_heads(odd_input, odd_output) + stream.synchronize() + dist.barrier() + + odd_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(odd_graph, stream=stream): + channel.all_gather_heads(odd_input, odd_output) + stream.synchronize() + + for replay in range(8): + step = 3000 + replay + odd_input.copy_( + _rank_query(step, rank, world_size, 1, torch.bfloat16, device) + ) + expected = torch.cat( + [ + _rank_query( + step, source, world_size, 1, torch.bfloat16, device + ) + for source in range(world_size) + ], + dim=1, + ) + stream.wait_stream(torch.cuda.current_stream(device)) + odd_graph.replay() + odd_graph.replay() + stream.synchronize() + torch.testing.assert_close(odd_output, expected, rtol=0, atol=0) + def _worker(rank: int, world_size: int, port: int) -> None: torch.cuda.set_device(rank) @@ -323,6 +436,8 @@ def _worker(rank: int, world_size: int, port: int) -> None: try: _check_eager(pool, rank, world_size, device) dist.barrier() + _check_eager_adjacency(pool, rank, world_size, device) + dist.barrier() _check_graph(pool, rank, world_size, device) torch.cuda.synchronize(device) finally: From 1b9a37b935a4868e4dec693c3b718b093aca5bb4 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Thu, 30 Jul 2026 22:28:22 +0000 Subject: [PATCH 04/11] fix(pcie): bound all two-slot host selectors --- sparkinfer/comm/pcie/pcie_dcp_a2a.cu | 8 ++-- sparkinfer/comm/pcie/pcie_oneshot.cu | 12 +++--- sparkinfer/comm/pcie/pcie_twoshot.cu | 12 +++--- sparkinfer/comm/pcie/two_slot_selector.h | 38 +++++++++++++++++++ tests/comm/test_pcie_two_slot_selector.py | 46 +++++++++++++++++++++++ 5 files changed, 101 insertions(+), 15 deletions(-) create mode 100644 sparkinfer/comm/pcie/two_slot_selector.h create mode 100644 tests/comm/test_pcie_two_slot_selector.py diff --git a/sparkinfer/comm/pcie/pcie_dcp_a2a.cu b/sparkinfer/comm/pcie/pcie_dcp_a2a.cu index 5f372b54a..6145b7986 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_a2a.cu +++ b/sparkinfer/comm/pcie/pcie_dcp_a2a.cu @@ -26,6 +26,8 @@ #include #include +#include "two_slot_selector.h" + #define CHECK_CUDA_SUCCESS(cmd) \ do { \ cudaError_t e = cmd; \ @@ -343,15 +345,13 @@ public: int64_t output_capacity_elems_; int64_t lse_offset_; int64_t lse_capacity_; - uint32_t next_slot_ = 0; + sparkinfer::pcie::TwoSlotSelector slot_selector_; int take_slot() { // Host-side selection executes once during CUDA graph capture. Graph // replay intentionally reuses that capture-stable address; eager calls // toggle slabs without an overflowing counter. - const int slot = static_cast(next_slot_); - next_slot_ ^= uint32_t{1}; - return slot; + return slot_selector_.take(); } PCIeDCPA2A(Signal **signals, diff --git a/sparkinfer/comm/pcie/pcie_oneshot.cu b/sparkinfer/comm/pcie/pcie_oneshot.cu index fc01c8273..5512a0d48 100644 --- a/sparkinfer/comm/pcie/pcie_oneshot.cu +++ b/sparkinfer/comm/pcie/pcie_oneshot.cu @@ -23,6 +23,8 @@ #include #include +#include "two_slot_selector.h" + #define CHECK_CUDA_SUCCESS(cmd) \ do { \ cudaError_t e = cmd; \ @@ -566,7 +568,7 @@ class PCIeAllreduce { std::map ipc_handles_; bool dbuf_enabled_ = false; - int dbuf_slot_ = 0; + sparkinfer::pcie::TwoSlotSelector dbuf_slot_selector_; void* dbuf_raw_[2][kMaxRanks] = {}; RankData* dbuf_rd_[2] = {}; @@ -626,7 +628,7 @@ class PCIeAllreduce { buffers_[ptrs0[rank_]] = dbuf_rd_[0]; buffers_[ptrs1[rank_]] = dbuf_rd_[1]; dbuf_enabled_ = true; - dbuf_slot_ = 0; + dbuf_slot_selector_ = sparkinfer::pcie::TwoSlotSelector{}; } void register_buffer(void** ptrs) { @@ -686,8 +688,7 @@ class PCIeAllreduce { if (dbuf_enabled_) { // The kernel stages the input into the eager slot itself; no host // staging memcpy is issued. - int slot = dbuf_slot_ % 2; - dbuf_slot_++; + const int slot = dbuf_slot_selector_.take(); ptrs = dbuf_rd_[slot]; staging = reinterpret_cast(dbuf_raw_[slot][rank_]); } else if (status == cudaStreamCaptureStatusActive) { @@ -764,8 +765,7 @@ class PCIeAllreduce { if (dbuf_enabled_) { // The kernel stages the input into the eager slot itself; no host // staging memcpy is issued. - int slot = dbuf_slot_ % 2; - dbuf_slot_++; + const int slot = dbuf_slot_selector_.take(); ptrs = dbuf_rd_[slot]; staging = reinterpret_cast(dbuf_raw_[slot][rank_]); } else if (status == cudaStreamCaptureStatusActive) { diff --git a/sparkinfer/comm/pcie/pcie_twoshot.cu b/sparkinfer/comm/pcie/pcie_twoshot.cu index 0df42ec8a..e4ec112fa 100644 --- a/sparkinfer/comm/pcie/pcie_twoshot.cu +++ b/sparkinfer/comm/pcie/pcie_twoshot.cu @@ -36,6 +36,8 @@ #include #include +#include "two_slot_selector.h" + #define CHECK_CUDA_SUCCESS(cmd) \ do { \ cudaError_t e = cmd; \ @@ -261,7 +263,9 @@ class PCIeTwoShot { int64_t scale_offset_; // bytes int64_t scale_stride_; // floats per src region int64_t max_shard_packs_; - int slot_ = 0; + sparkinfer::pcie::TwoSlotSelector slot_selector_; + + int take_slot() { return slot_selector_.take(); } PCIeTwoShot(Signal** signals, const std::vector>& staging, int64_t pack_stride, int64_t scale_offset, int64_t scale_stride, int rank, @@ -312,8 +316,7 @@ class PCIeTwoShot { throw std::runtime_error("num_rows must be divisible by world size"); const int64_t rows_per_rank = num_rows / world_size_; check_shard(rows_per_rank, row_elems); - const int s = slot_ % 2; - slot_++; + const int s = take_slot(); const int64_t shard_packs = rows_per_rank * (row_elems / 16); int blocks = std::max(1, std::min(block_limit, (shard_packs + threads - 1) / threads)); @@ -329,8 +332,7 @@ class PCIeTwoShot { void all_gather(cudaStream_t stream, const void* payload, const void* scale, void* out, int64_t rows_per_rank, int64_t row_elems, int threads, int block_limit) { check_shard(rows_per_rank, row_elems); - const int s = slot_ % 2; - slot_++; + const int s = take_slot(); const int64_t shard_packs = rows_per_rank * (row_elems / 16); int blocks = std::max(1, std::min(block_limit, (shard_packs + threads - 1) / threads)); diff --git a/sparkinfer/comm/pcie/two_slot_selector.h b/sparkinfer/comm/pcie/two_slot_selector.h new file mode 100644 index 000000000..21d1c48c9 --- /dev/null +++ b/sparkinfer/comm/pcie/two_slot_selector.h @@ -0,0 +1,38 @@ +#pragma once + +#include + +namespace sparkinfer::pcie { + +// A double-buffer selector with an explicitly finite state space. Keeping the +// phase as an enum makes signed overflow and out-of-range remainders impossible. +class TwoSlotSelector { + public: + constexpr int take() noexcept { + if (next_ == Next::kFirst) { + next_ = Next::kSecond; + return 0; + } + next_ = Next::kFirst; + return 1; + } + + private: + enum class Next : std::uint32_t { kFirst, kSecond }; + Next next_ = Next::kFirst; +}; + +namespace detail { +constexpr bool two_slot_sequence_is_preserved() { + TwoSlotSelector selector; + return selector.take() == 0 && selector.take() == 1 && + selector.take() == 0 && selector.take() == 1; +} +} // namespace detail + +static_assert(detail::two_slot_sequence_is_preserved(), + "two-slot selection must start at zero and alternate"); +static_assert(sizeof(TwoSlotSelector) == sizeof(std::uint32_t), + "two-slot selector storage must remain 32-bit"); + +} // namespace sparkinfer::pcie diff --git a/tests/comm/test_pcie_two_slot_selector.py b/tests/comm/test_pcie_two_slot_selector.py new file mode 100644 index 000000000..485a3e9d0 --- /dev/null +++ b/tests/comm/test_pcie_two_slot_selector.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import subprocess +from pathlib import Path + + +def test_two_slot_selector_preserves_phase_for_long_run(tmp_path: Path) -> None: + repo_root = Path(__file__).resolve().parents[2] + source = tmp_path / "two_slot_selector_test.cpp" + binary = tmp_path / "two_slot_selector_test" + source.write_text( + r""" +#include +#include "sparkinfer/comm/pcie/two_slot_selector.h" + +int main() { + sparkinfer::pcie::TwoSlotSelector selector; + constexpr std::uint64_t kTransitions = 10'000'003; + for (std::uint64_t call = 0; call < kTransitions; ++call) { + if (selector.take() != static_cast(call & 1u)) return 1; + } + return selector.take() == static_cast(kTransitions & 1u) ? 0 : 2; +} +""", + encoding="utf-8", + ) + + subprocess.run( + [ + "c++", + "-std=c++17", + "-O3", + "-Wall", + "-Wextra", + "-Werror", + "-I", + str(repo_root), + str(source), + "-o", + str(binary), + ], + check=True, + capture_output=True, + text=True, + ) + subprocess.run([str(binary)], check=True) From 6b1d9eb4819fe46b159aabee55cf997582febd3c Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Thu, 30 Jul 2026 22:31:52 +0000 Subject: [PATCH 05/11] test(pcie): observe replay slot advancement --- sparkinfer/comm/pcie/pcie_dcp_a2a.cu | 23 +++++++++++++++++ tests/comm/test_pcie_dcp_a2a.py | 21 --------------- tests/comm/test_pcie_dcp_a2a_gpu.py | 38 ++++++++++++++++------------ 3 files changed, 45 insertions(+), 37 deletions(-) diff --git a/sparkinfer/comm/pcie/pcie_dcp_a2a.cu b/sparkinfer/comm/pcie/pcie_dcp_a2a.cu index 7e153981a..05424808c 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_a2a.cu +++ b/sparkinfer/comm/pcie/pcie_dcp_a2a.cu @@ -658,6 +658,27 @@ static void all_gather_heads(fptr_t pointer, torch::Tensor &local_input, } } +// Narrow synchronization hook for the multi-rank CUDA regression. It can +// reset and observe the first local staging word in each slot. +static std::vector debug_staging_words(fptr_t pointer, int64_t fill) { + auto *runtime = reinterpret_cast(pointer); + TORCH_CHECK(fill >= -1 && fill <= 0xffff); + std::vector words(2); + for (int slot = 0; slot < 2; ++slot) { + auto *address = runtime->staging_.slots[slot].ptrs[runtime->rank_]; + if (fill >= 0) { + const auto fill_word = static_cast(fill); + CHECK_CUDA_SUCCESS(cudaMemcpy(address, &fill_word, sizeof(fill_word), + cudaMemcpyHostToDevice)); + } + unsigned short word; + CHECK_CUDA_SUCCESS( + cudaMemcpy(&word, address, sizeof(word), cudaMemcpyDeviceToHost)); + words[slot] = int64_t(word); + } + return words; +} + static void dispose(fptr_t pointer) { delete reinterpret_cast(pointer); } @@ -670,6 +691,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { "fused PCIe DCP LSE reduce-scatter"); module.def("all_gather_heads", &all_gather_heads, "PCIe DCP head-dimension all-gather"); + module.def("_debug_staging_words", &debug_staging_words, + "test-only local staging words"); module.def("dispose", &dispose, "dispose PCIe DCP A2A"); module.def("meta_size", &meta_size, "signal metadata size"); } diff --git a/tests/comm/test_pcie_dcp_a2a.py b/tests/comm/test_pcie_dcp_a2a.py index 0fc1dd8f0..4f1e01cfc 100644 --- a/tests/comm/test_pcie_dcp_a2a.py +++ b/tests/comm/test_pcie_dcp_a2a.py @@ -1,7 +1,5 @@ from __future__ import annotations -from pathlib import Path - import pytest import torch @@ -99,25 +97,6 @@ def _make_runtime(ext: _FakeExt | None = None) -> PCIeDCPA2A: ) -def test_cuda_staging_slot_is_selected_from_execution_state(): - source = ( - Path(__file__).parents[2] - / "sparkinfer" - / "comm" - / "pcie" - / "pcie_dcp_a2a.cu" - ).read_text() - - assert "int slot_" not in source - assert "slot_++" not in source - assert ( - source.count("load_flag(&self->self_counter[blockIdx.x][0])") == 2 - ) - assert "select_staging_slot_kernel" not in source - assert source.count("DoubleStaging staging_options") == 2 - assert source.count("__shared__ RankStaging staging;") == 2 - - def test_staging_layout_has_aligned_disjoint_slots(): layout = _staging_layout( signal_bytes=12345, diff --git a/tests/comm/test_pcie_dcp_a2a_gpu.py b/tests/comm/test_pcie_dcp_a2a_gpu.py index 332eddbc4..f30a9e75b 100644 --- a/tests/comm/test_pcie_dcp_a2a_gpu.py +++ b/tests/comm/test_pcie_dcp_a2a_gpu.py @@ -395,25 +395,31 @@ def _check_graph( channel.all_gather_heads(odd_input, odd_output) stream.synchronize() - for replay in range(8): - step = 3000 + replay - odd_input.copy_( - _rank_query(step, rank, world_size, 1, torch.bfloat16, device) - ) - expected = torch.cat( - [ - _rank_query( - step, source, world_size, 1, torch.bfloat16, device - ) - for source in range(world_size) - ], - dim=1, - ) + snapshots = [tuple(channel._ext._debug_staging_words(channel._ptr, 0))] + for value in (1.0, 2.0): + odd_input.fill_(value) stream.wait_stream(torch.cuda.current_stream(device)) odd_graph.replay() - odd_graph.replay() stream.synchronize() - torch.testing.assert_close(odd_output, expected, rtol=0, atol=0) + snapshots.append( + tuple(channel._ext._debug_staging_words(channel._ptr, -1)) + ) + torch.testing.assert_close( + odd_output, torch.full_like(odd_output, value), rtol=0, atol=0 + ) + + changed_slots = [ + { + slot + for slot, (before, after) in enumerate( + zip(snapshots[index], snapshots[index + 1], strict=True) + ) + if before != after + } + for index in range(2) + ] + assert all(len(changed) == 1 for changed in changed_slots) + assert changed_slots[0] != changed_slots[1] def _worker(rank: int, world_size: int, port: int) -> None: From d3ab9bdf15ea8cd4c0ce0a16e229667a2831b7d6 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Thu, 30 Jul 2026 22:35:54 +0000 Subject: [PATCH 06/11] test(pcie): keep slot observation test-local --- sparkinfer/comm/pcie/pcie_dcp_a2a.cu | 23 ------------- tests/comm/test_pcie_dcp_a2a_gpu.py | 49 +++++++++++++++++++++++----- 2 files changed, 41 insertions(+), 31 deletions(-) diff --git a/sparkinfer/comm/pcie/pcie_dcp_a2a.cu b/sparkinfer/comm/pcie/pcie_dcp_a2a.cu index 05424808c..7e153981a 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_a2a.cu +++ b/sparkinfer/comm/pcie/pcie_dcp_a2a.cu @@ -658,27 +658,6 @@ static void all_gather_heads(fptr_t pointer, torch::Tensor &local_input, } } -// Narrow synchronization hook for the multi-rank CUDA regression. It can -// reset and observe the first local staging word in each slot. -static std::vector debug_staging_words(fptr_t pointer, int64_t fill) { - auto *runtime = reinterpret_cast(pointer); - TORCH_CHECK(fill >= -1 && fill <= 0xffff); - std::vector words(2); - for (int slot = 0; slot < 2; ++slot) { - auto *address = runtime->staging_.slots[slot].ptrs[runtime->rank_]; - if (fill >= 0) { - const auto fill_word = static_cast(fill); - CHECK_CUDA_SUCCESS(cudaMemcpy(address, &fill_word, sizeof(fill_word), - cudaMemcpyHostToDevice)); - } - unsigned short word; - CHECK_CUDA_SUCCESS( - cudaMemcpy(&word, address, sizeof(word), cudaMemcpyDeviceToHost)); - words[slot] = int64_t(word); - } - return words; -} - static void dispose(fptr_t pointer) { delete reinterpret_cast(pointer); } @@ -691,8 +670,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { "fused PCIe DCP LSE reduce-scatter"); module.def("all_gather_heads", &all_gather_heads, "PCIe DCP head-dimension all-gather"); - module.def("_debug_staging_words", &debug_staging_words, - "test-only local staging words"); module.def("dispose", &dispose, "dispose PCIe DCP A2A"); module.def("meta_size", &meta_size, "signal metadata size"); } diff --git a/tests/comm/test_pcie_dcp_a2a_gpu.py b/tests/comm/test_pcie_dcp_a2a_gpu.py index f30a9e75b..8cb3314b7 100644 --- a/tests/comm/test_pcie_dcp_a2a_gpu.py +++ b/tests/comm/test_pcie_dcp_a2a_gpu.py @@ -1,5 +1,6 @@ from __future__ import annotations +import ctypes import os import socket @@ -11,6 +12,7 @@ from sparkinfer.comm.pcie.pcie_dcp_a2a import ( PCIeDCPA2APool, _load_extension, + _staging_layout, lse_reduce_scatter_reference, ) @@ -99,6 +101,33 @@ def _reference( ) +def _local_staging_words(channel, stream: torch.cuda.Stream) -> tuple[int, int]: + assert channel._ipc is not None + assert len(channel._owned_buffers) == 1 + layout = _staging_layout( + signal_bytes=int(channel._ext.meta_size()), + world_size=channel.world_size, + max_batch_size=channel.max_batch_size, + total_heads=channel.total_heads, + head_dim=channel.head_dim, + query_head_dim=channel.query_head_dim, + ) + words = (ctypes.c_uint16(), ctypes.c_uint16()) + local_ptr = channel._owned_buffers[0].local_ptr + for word, offset in zip( + words, + (layout.staging0_offset, layout.staging1_offset), + strict=True, + ): + channel._ipc.cudaMemcpyAsync( + ctypes.addressof(word), + local_ptr + offset, + ctypes.sizeof(word), + int(stream.cuda_stream), + ) + stream.synchronize() + return words[0].value, words[1].value + def _check_eager( pool: PCIeDCPA2APool, rank: int, @@ -395,15 +424,19 @@ def _check_graph( channel.all_gather_heads(odd_input, odd_output) stream.synchronize() - snapshots = [tuple(channel._ext._debug_staging_words(channel._ptr, 0))] + # Prime both slots after capture, then observe them read-only. A graph with + # a host-baked slot changes the same word on both replays; execution-owned + # parity changes alternate words. + with torch.cuda.stream(stream): + for value in (11.0, 12.0): + odd_input.fill_(value) + channel.all_gather_heads(odd_input, odd_output) + snapshots = [_local_staging_words(channel, stream)] for value in (1.0, 2.0): - odd_input.fill_(value) - stream.wait_stream(torch.cuda.current_stream(device)) - odd_graph.replay() - stream.synchronize() - snapshots.append( - tuple(channel._ext._debug_staging_words(channel._ptr, -1)) - ) + with torch.cuda.stream(stream): + odd_input.fill_(value) + odd_graph.replay() + snapshots.append(_local_staging_words(channel, stream)) torch.testing.assert_close( odd_output, torch.full_like(odd_output, value), rtol=0, atol=0 ) From 22c11dfba8451395ea0133b767937bac332c0be2 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Fri, 31 Jul 2026 00:54:57 +0000 Subject: [PATCH 07/11] refactor(pcie): centralize DCP staging selection --- sparkinfer/comm/pcie/pcie_dcp_a2a.cu | 35 ++++++++++--------- tests/comm/test_pcie_dcp_a2a.py | 51 ++++++++++++++++++++++++++++ 2 files changed, 68 insertions(+), 18 deletions(-) diff --git a/sparkinfer/comm/pcie/pcie_dcp_a2a.cu b/sparkinfer/comm/pcie/pcie_dcp_a2a.cu index 48b6973be..ed5995aee 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_a2a.cu +++ b/sparkinfer/comm/pcie/pcie_dcp_a2a.cu @@ -101,6 +101,21 @@ static DINLINE FlagType load_flag(FlagType *address) { return value; } +template +DINLINE void select_staging(RankStaging &staging, + const DoubleStaging &staging_options, + Signal *self) { + if (threadIdx.x == 0) { + const int slot = + int(load_flag(&self->self_counter[blockIdx.x][0]) & FlagType{1}); +#pragma unroll + for (int peer = 0; peer < world_size; ++peer) { + staging.ptrs[peer] = staging_options.slots[slot].ptrs[peer]; + } + } + __syncthreads(); +} + // pre_sync orders in-kernel staging stores (from every warp of the block) // before the flag post; without staging the flags can go out immediately. template @@ -174,15 +189,7 @@ __global__ void __launch_bounds__(512, 1) auto *output_packs = reinterpret_cast *>(output); __shared__ RankStaging staging; - if (threadIdx.x == 0) { - const int slot = - int(load_flag(&self->self_counter[blockIdx.x][0]) & FlagType{1}); -#pragma unroll - for (int peer = 0; peer < world_size; ++peer) { - staging.ptrs[peer] = staging_options.slots[slot].ptrs[peer]; - } - } - __syncthreads(); + select_staging(staging, staging_options, self); auto *staging_out = reinterpret_cast *>(staging.ptrs[rank]); auto *staging_lse = reinterpret_cast( @@ -315,15 +322,7 @@ __global__ void __launch_bounds__(512, 1) const auto *local_packs = reinterpret_cast *>(local_input); __shared__ RankStaging staging; - if (threadIdx.x == 0) { - const int slot = - int(load_flag(&self->self_counter[blockIdx.x][0]) & FlagType{1}); -#pragma unroll - for (int peer = 0; peer < world_size; ++peer) { - staging.ptrs[peer] = staging_options.slots[slot].ptrs[peer]; - } - } - __syncthreads(); + select_staging(staging, staging_options, self); auto *staging_out = reinterpret_cast *>(staging.ptrs[rank]); for (int row = warp_first; row < rows; row += warp_stride) { diff --git a/tests/comm/test_pcie_dcp_a2a.py b/tests/comm/test_pcie_dcp_a2a.py index 4f1e01cfc..ac8ceca4c 100644 --- a/tests/comm/test_pcie_dcp_a2a.py +++ b/tests/comm/test_pcie_dcp_a2a.py @@ -368,6 +368,57 @@ def test_pool_restores_eager_mapping_after_capture(monkeypatch): assert graph_channel in pool._all_channels +def test_pool_restores_nested_capture_mappings(monkeypatch): + current_stream = [7] + capturing = [False] + pool = PCIeDCPA2APool( + rank=0, + world_size=2, + device=torch.device("cpu"), + max_batch_size=4, + total_heads=32, + head_dim=64, + channel_factory=lambda stream_key: _make_runtime(), + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_dcp_a2a._current_stream_key", + lambda device, stream=None: ( + current_stream[0] if stream is None else int(stream) + ), + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_dcp_a2a._is_current_stream_capturing", + lambda device: capturing[0], + ) + + eager_channel = pool.for_stream() + with pool.capture(7) as outer_channel: + assert pool._channels == {7: outer_channel} + + current_stream[0] = 8 + with pool.capture() as inner_channel: + capturing[0] = True + current_stream[0] = 80 + assert pool.for_stream() is inner_channel + assert pool._channels == { + 7: outer_channel, + 8: inner_channel, + 80: inner_channel, + } + capturing[0] = False + + assert pool._channels == {7: outer_channel} + capturing[0] = True + current_stream[0] = 70 + assert pool.for_stream() is outer_channel + capturing[0] = False + + assert eager_channel is not outer_channel + assert outer_channel is not inner_channel + assert pool._channels == {7: eager_channel} + assert pool._capture_channel_stack == [] + + def test_pool_rolls_back_throwaway_capture_channels(monkeypatch): created = [] From 7e17517542c818cc65f140143a02eafa37c1d761 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Fri, 31 Jul 2026 01:23:29 +0000 Subject: [PATCH 08/11] fix(pcie): advance reusable staging slots on replay --- sparkinfer/comm/pcie/pcie_oneshot.cu | 120 +++++++++++------- sparkinfer/comm/pcie/pcie_twoshot.cu | 56 +++++--- .../test_pcie_oneshot_fused_rmsnorm_gpu.py | 45 ++++++- tests/comm/test_pcie_oneshot_torture.py | 61 ++++++++- tests/comm/test_pcie_twoshot.py | 109 ++++++++++++++++ 5 files changed, 324 insertions(+), 67 deletions(-) diff --git a/sparkinfer/comm/pcie/pcie_oneshot.cu b/sparkinfer/comm/pcie/pcie_oneshot.cu index 5512a0d48..a6e78f081 100644 --- a/sparkinfer/comm/pcie/pcie_oneshot.cu +++ b/sparkinfer/comm/pcie/pcie_oneshot.cu @@ -23,7 +23,6 @@ #include #include -#include "two_slot_selector.h" #define CHECK_CUDA_SUCCESS(cmd) \ do { \ @@ -78,6 +77,8 @@ static bool oneshot_push_enabled() { } struct Signal { + alignas(128) FlagType staging_generation; + FlagType active_staging_slot; alignas(128) FlagType self_counter[kMaxBlocks][kMaxRanks]; alignas(128) FlagType peer_counter[2][kMaxBlocks][kMaxRanks * kFlagStride]; alignas(128) FlagType rms_arrive[kMaxBlocks]; @@ -89,6 +90,10 @@ struct __align__(16) RankData { const void* __restrict__ ptrs[kMaxRanks]; }; +struct DoubleRankData { + RankData* slots[2]; +}; + struct __align__(16) RankSignals { Signal* signals[kMaxRanks]; }; @@ -179,6 +184,25 @@ static DINLINE FlagType ld_flag_relaxed(FlagType* flag_addr) { return flag; } +__global__ void advance_staging_slot(Signal* self_sg) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + const FlagType generation = self_sg->staging_generation; + self_sg->active_staging_slot = generation & FlagType{1}; + self_sg->staging_generation = generation + FlagType{1}; + } +} + +template +DINLINE void select_rank_data(RankData& selected, const DoubleRankData& options, Signal* self_sg) { + if (threadIdx.x == 0) { + const int slot = int(self_sg->active_staging_slot); + const RankData* source = options.slots[slot]; +#pragma unroll + for (int peer = 0; peer < ngpus; ++peer) selected.ptrs[peer] = source->ptrs[peer]; + } + __syncthreads(); +} + static DINLINE FlagType ld_flag_relaxed_gpu(FlagType* flag_addr) { FlagType flag; asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(flag) : "l"(flag_addr) : "memory"); @@ -255,17 +279,23 @@ DINLINE float block_reduce_sum(float value, float* warp_sums) { // sufficient for staging visibility. template __global__ void __launch_bounds__(512, 1) pcie_allreduce_kernel( - RankData* _dp, + DoubleRankData data_options, RankSignals sg, Signal* self_sg, const T* __restrict__ input, T* __restrict__ result, - T* __restrict__ self_staging, int rank, int size) { using P = typename packed_t::P; using A = typename packed_t::A; - auto dp = *_dp; + RankData dp; + if constexpr (stage_input) { + __shared__ RankData selected; + select_rank_data(selected, data_options, self_sg); + dp = selected; + } else { + dp = *data_options.slots[0]; + } const P* rotated[ngpus]; #pragma unroll for (int i = 0; i < ngpus; i++) { @@ -273,7 +303,8 @@ __global__ void __launch_bounds__(512, 1) pcie_allreduce_kernel( } if constexpr (stage_input) { const P* input_p = reinterpret_cast(input); - P* staging_p = reinterpret_cast(self_staging); + P* staging_p = + reinterpret_cast(const_cast(dp.ptrs[rank])); for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; idx += gridDim.x * blockDim.x) { staging_p[idx] = input_p[idx]; } @@ -330,7 +361,7 @@ constexpr int kModeStagePush = 2; // resident. template __global__ void __launch_bounds__(512, 1) pcie_allreduce_fused_add_rms_norm_kernel( - RankData* _dp, + DoubleRankData data_options, RankSignals sg, Signal* self_sg, const T* __restrict__ input, @@ -338,7 +369,6 @@ __global__ void __launch_bounds__(512, 1) pcie_allreduce_fused_add_rms_norm_kern const T* __restrict__ weight, T* __restrict__ output, T* __restrict__ residual_output, - T* __restrict__ self_staging, int rank, int hidden_packs, int ctas_per_row, @@ -348,6 +378,14 @@ __global__ void __launch_bounds__(512, 1) pcie_allreduce_fused_add_rms_norm_kern using A = typename packed_t::A; __shared__ float warp_sums[32]; __shared__ float s_inv_rms; + RankData dp; + if constexpr (mode == kModeRegistered) { + dp = *data_options.slots[0]; + } else { + __shared__ RankData selected; + select_rank_data(selected, data_options, self_sg); + dp = selected; + } const int row = single_cta ? blockIdx.x : blockIdx.x / ctas_per_row; const int row_cta = single_cta ? 0 : blockIdx.x - row * ctas_per_row; @@ -357,7 +395,6 @@ __global__ void __launch_bounds__(512, 1) pcie_allreduce_fused_add_rms_norm_kern const int max_packs = (hidden_packs + col_stride - 1) / col_stride; const bool reg_norm = max_packs <= kRegPacks; - auto dp = *_dp; const P* rotated[ngpus]; #pragma unroll for (int i = 0; i < ngpus; i++) { @@ -369,7 +406,7 @@ __global__ void __launch_bounds__(512, 1) pcie_allreduce_fused_add_rms_norm_kern const P* weight_p = reinterpret_cast(weight); P* output_p = reinterpret_cast(output); P* residual_output_p = reinterpret_cast(residual_output); - P* staging_p = reinterpret_cast(self_staging); + P* staging_p = reinterpret_cast(const_cast(dp.ptrs[rank])); // Sense-reversing row barrier: latch the generation before arriving. // Launches are stream-serialized, so this read cannot race a peer CTA's @@ -568,9 +605,7 @@ class PCIeAllreduce { std::map ipc_handles_; bool dbuf_enabled_ = false; - sparkinfer::pcie::TwoSlotSelector dbuf_slot_selector_; - void* dbuf_raw_[2][kMaxRanks] = {}; - RankData* dbuf_rd_[2] = {}; + DoubleRankData dbuf_data_ = {}; PCIeAllreduce(Signal** signals, void* rank_data, size_t rank_data_sz, int rank, int world_size) : rank_(rank), @@ -618,17 +653,14 @@ class PCIeAllreduce { for (int s = 0; s < 2; s++) { void** ptrs = s == 0 ? ptrs0 : ptrs1; RankData data; - for (int i = 0; i < world_size_; i++) { - data.ptrs[i] = ptrs[i]; - dbuf_raw_[s][i] = ptrs[i]; - } - dbuf_rd_[s] = d_rank_data_base_++; - CHECK_CUDA_SUCCESS(cudaMemcpy(dbuf_rd_[s], &data, sizeof(RankData), cudaMemcpyHostToDevice)); + for (int i = 0; i < world_size_; i++) data.ptrs[i] = ptrs[i]; + dbuf_data_.slots[s] = d_rank_data_base_++; + CHECK_CUDA_SUCCESS( + cudaMemcpy(dbuf_data_.slots[s], &data, sizeof(RankData), cudaMemcpyHostToDevice)); } - buffers_[ptrs0[rank_]] = dbuf_rd_[0]; - buffers_[ptrs1[rank_]] = dbuf_rd_[1]; + buffers_[ptrs0[rank_]] = dbuf_data_.slots[0]; + buffers_[ptrs1[rank_]] = dbuf_data_.slots[1]; dbuf_enabled_ = true; - dbuf_slot_selector_ = sparkinfer::pcie::TwoSlotSelector{}; } void register_buffer(void** ptrs) { @@ -680,17 +712,16 @@ class PCIeAllreduce { if (block_limit > kMaxBlocks) throw std::runtime_error("max supported block limit is " + std::to_string(kMaxBlocks)); - RankData* ptrs; - T* staging = nullptr; + DoubleRankData data_options = {}; + bool stage_input = false; cudaStreamCaptureStatus status; CHECK_CUDA_SUCCESS(cudaStreamIsCapturing(stream, &status)); if (dbuf_enabled_) { - // The kernel stages the input into the eager slot itself; no host - // staging memcpy is issued. - const int slot = dbuf_slot_selector_.take(); - ptrs = dbuf_rd_[slot]; - staging = reinterpret_cast(dbuf_raw_[slot][rank_]); + // A captured selector kernel advances the active eager slab on every + // execution, including CUDA graph replay. + data_options = dbuf_data_; + stage_input = true; } else if (status == cudaStreamCaptureStatusActive) { throw std::runtime_error( "PCIe oneshot graph capture requires eager IPC buffers; construct the runtime with eager buffers or use " @@ -700,22 +731,24 @@ class PCIeAllreduce { if (it == buffers_.end()) throw std::runtime_error( "buffer address " + std::to_string(reinterpret_cast(input)) + " is not registered!"); - ptrs = it->second; + data_options.slots[0] = it->second; + data_options.slots[1] = it->second; } size /= d; int blocks = std::min(block_limit, (size + threads - 1) / threads); blocks = std::max(blocks, 1); + if (stage_input) advance_staging_slot<<<1, 1, 0, stream>>>(self_sg_); #define KL(ngpus) \ do { \ - if (staging != nullptr) { \ + if (stage_input) { \ pcie_allreduce_kernel \ - <<>>(ptrs, sg_, self_sg_, input, output, staging, rank_, size); \ + <<>>(data_options, sg_, self_sg_, input, output, rank_, size); \ } else { \ pcie_allreduce_kernel \ - <<>>(ptrs, sg_, self_sg_, input, output, staging, rank_, size); \ - } \ + <<>>(data_options, sg_, self_sg_, input, output, rank_, size); \ + } \ } while (0) switch (world_size_) { case 2: @@ -758,16 +791,15 @@ class PCIeAllreduce { throw std::runtime_error( "fused allreduce RMSNorm requires hidden size to be a multiple of " + std::to_string(pack_size)); - RankData* ptrs; - T* staging = nullptr; + DoubleRankData data_options = {}; + bool use_eager_staging = false; cudaStreamCaptureStatus status; CHECK_CUDA_SUCCESS(cudaStreamIsCapturing(stream, &status)); if (dbuf_enabled_) { - // The kernel stages the input into the eager slot itself; no host - // staging memcpy is issued. - const int slot = dbuf_slot_selector_.take(); - ptrs = dbuf_rd_[slot]; - staging = reinterpret_cast(dbuf_raw_[slot][rank_]); + // A captured selector kernel advances the active eager slab on every + // execution, including CUDA graph replay. + data_options = dbuf_data_; + use_eager_staging = true; } else if (status == cudaStreamCaptureStatusActive) { throw std::runtime_error( "PCIe oneshot graph capture requires eager IPC buffers; construct the runtime with eager buffers or use " @@ -777,7 +809,8 @@ class PCIeAllreduce { if (it == buffers_.end()) throw std::runtime_error( "buffer address " + std::to_string(reinterpret_cast(input)) + " is not registered!"); - ptrs = it->second; + data_options.slots[0] = it->second; + data_options.slots[1] = it->second; } int rows = size / hidden_size; @@ -800,13 +833,14 @@ class PCIeAllreduce { const int blocks = rows * ctas_per_row; const bool single = ctas_per_row == 1; const int shard_packs = size / pack_size; - const int mode = staging == nullptr ? kModeRegistered + const int mode = !use_eager_staging ? kModeRegistered : oneshot_push_enabled() ? kModeStagePush : kModeStagePull; + if (use_eager_staging) advance_staging_slot<<<1, 1, 0, stream>>>(self_sg_); #define KL(ngpus, SINGLE, MODE) \ pcie_allreduce_fused_add_rms_norm_kernel<<>>( \ - ptrs, sg_, self_sg_, input, residual, weight, output, residual_output, staging, rank_, hidden_packs, \ + data_options, sg_, self_sg_, input, residual, weight, output, residual_output, rank_, hidden_packs, \ ctas_per_row, shard_packs, epsilon); #define DISPATCH_MODE(ngpus, SINGLE) \ do { \ diff --git a/sparkinfer/comm/pcie/pcie_twoshot.cu b/sparkinfer/comm/pcie/pcie_twoshot.cu index e4ec112fa..c12e11880 100644 --- a/sparkinfer/comm/pcie/pcie_twoshot.cu +++ b/sparkinfer/comm/pcie/pcie_twoshot.cu @@ -36,7 +36,6 @@ #include #include -#include "two_slot_selector.h" #define CHECK_CUDA_SUCCESS(cmd) \ do { \ @@ -59,6 +58,8 @@ using FlagType = uint32_t; constexpr int kFlagStride = 32; struct Signal { + alignas(128) FlagType staging_generation; + FlagType active_staging_slot; alignas(128) FlagType self_counter[kMaxBlocks][kMaxRanks]; alignas(128) FlagType peer_counter[2][kMaxBlocks][kMaxRanks * kFlagStride]; }; @@ -67,6 +68,10 @@ struct __align__(16) RankPtrs { void* __restrict__ ptrs[kMaxRanks]; }; +struct DoubleRankPtrs { + RankPtrs slots[2]; +}; + struct RankSignals { Signal* signals[kMaxRanks]; }; @@ -83,6 +88,24 @@ static DINLINE FlagType ld_flag_relaxed(FlagType* flag_addr) { return flag; } +__global__ void advance_staging_slot(Signal* self_sg) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + const FlagType generation = self_sg->staging_generation; + self_sg->active_staging_slot = generation & FlagType{1}; + self_sg->staging_generation = generation + FlagType{1}; + } +} + +template +DINLINE void select_rank_ptrs(RankPtrs& selected, const DoubleRankPtrs& options, Signal* self_sg) { + if (threadIdx.x == 0) { + const int slot = int(self_sg->active_staging_slot); +#pragma unroll + for (int peer = 0; peer < ngpus; ++peer) selected.ptrs[peer] = options.slots[slot].ptrs[peer]; + } + __syncthreads(); +} + // Block-pairwise barrier: after it returns, all prior writes from every // rank's block blockIdx.x are visible to this block. template @@ -148,9 +171,11 @@ DINLINE void block_range(int64_t total, int64_t& begin, int64_t& end) { template __global__ void __launch_bounds__(512, 1) rs_fp8_kernel( const Fp8Pack* __restrict__ src_payload, const float* __restrict__ src_scale, - RankPtrs staging, int64_t pack_stride, int64_t scale_offset, int64_t scale_stride, - RankSignals sg, Signal* self_sg, nv_bfloat16* __restrict__ out, int rank, - int rows_per_rank, int row_elems) { + DoubleRankPtrs staging_options, int64_t pack_stride, int64_t scale_offset, + int64_t scale_stride, RankSignals sg, Signal* self_sg, + nv_bfloat16* __restrict__ out, int rank, int rows_per_rank, int row_elems) { + __shared__ RankPtrs staging; + select_rank_ptrs(staging, staging_options, self_sg); const int packs_per_row = row_elems / 16; const int64_t shard_packs = int64_t(rows_per_rank) * packs_per_row; int64_t begin, end; @@ -202,9 +227,11 @@ __global__ void __launch_bounds__(512, 1) rs_fp8_kernel( template __global__ void __launch_bounds__(512, 1) ag_fp8_kernel( const Fp8Pack* __restrict__ src_payload, const float* __restrict__ src_scale, - RankPtrs staging, int64_t pack_stride, int64_t scale_offset, int64_t scale_stride, - RankSignals sg, Signal* self_sg, nv_bfloat16* __restrict__ out, int rank, - int rows_per_rank, int row_elems) { + DoubleRankPtrs staging_options, int64_t pack_stride, int64_t scale_offset, + int64_t scale_stride, RankSignals sg, Signal* self_sg, + nv_bfloat16* __restrict__ out, int rank, int rows_per_rank, int row_elems) { + __shared__ RankPtrs staging; + select_rank_ptrs(staging, staging_options, self_sg); const int packs_per_row = row_elems / 16; const int64_t shard_packs = int64_t(rows_per_rank) * packs_per_row; int64_t begin, end; @@ -258,14 +285,11 @@ class PCIeTwoShot { RankSignals sg_; Signal* self_sg_; - RankPtrs staging_[2]; + DoubleRankPtrs staging_; int64_t pack_stride_; // Fp8Packs per src region int64_t scale_offset_; // bytes int64_t scale_stride_; // floats per src region int64_t max_shard_packs_; - sparkinfer::pcie::TwoSlotSelector slot_selector_; - - int take_slot() { return slot_selector_.take(); } PCIeTwoShot(Signal** signals, const std::vector>& staging, int64_t pack_stride, int64_t scale_offset, int64_t scale_stride, int rank, @@ -279,7 +303,7 @@ class PCIeTwoShot { max_shard_packs_(pack_stride) { for (int i = 0; i < world_size_; i++) { sg_.signals[i] = signals[i]; - for (int s = 0; s < 2; s++) staging_[s].ptrs[i] = staging[i][s]; + for (int s = 0; s < 2; s++) staging_.slots[s].ptrs[i] = staging[i][s]; } } @@ -316,15 +340,15 @@ class PCIeTwoShot { throw std::runtime_error("num_rows must be divisible by world size"); const int64_t rows_per_rank = num_rows / world_size_; check_shard(rows_per_rank, row_elems); - const int s = take_slot(); const int64_t shard_packs = rows_per_rank * (row_elems / 16); int blocks = std::max(1, std::min(block_limit, (shard_packs + threads - 1) / threads)); + advance_staging_slot<<<1, 1, 0, stream>>>(self_sg_); dispatch([&](auto ng) { constexpr int ngpus = decltype(ng)::value; rs_fp8_kernel<<>>( reinterpret_cast(payload), reinterpret_cast(scale), - staging_[s], pack_stride_, scale_offset_, scale_stride_, sg_, self_sg_, + staging_, pack_stride_, scale_offset_, scale_stride_, sg_, self_sg_, reinterpret_cast(out), rank_, int(rows_per_rank), int(row_elems)); }); } @@ -332,15 +356,15 @@ class PCIeTwoShot { void all_gather(cudaStream_t stream, const void* payload, const void* scale, void* out, int64_t rows_per_rank, int64_t row_elems, int threads, int block_limit) { check_shard(rows_per_rank, row_elems); - const int s = take_slot(); const int64_t shard_packs = rows_per_rank * (row_elems / 16); int blocks = std::max(1, std::min(block_limit, (shard_packs + threads - 1) / threads)); + advance_staging_slot<<<1, 1, 0, stream>>>(self_sg_); dispatch([&](auto ng) { constexpr int ngpus = decltype(ng)::value; ag_fp8_kernel<<>>( reinterpret_cast(payload), reinterpret_cast(scale), - staging_[s], pack_stride_, scale_offset_, scale_stride_, sg_, self_sg_, + staging_, pack_stride_, scale_offset_, scale_stride_, sg_, self_sg_, reinterpret_cast(out), rank_, int(rows_per_rank), int(row_elems)); }); } diff --git a/tests/comm/test_pcie_oneshot_fused_rmsnorm_gpu.py b/tests/comm/test_pcie_oneshot_fused_rmsnorm_gpu.py index 34f9e5e04..508895978 100644 --- a/tests/comm/test_pcie_oneshot_fused_rmsnorm_gpu.py +++ b/tests/comm/test_pcie_oneshot_fused_rmsnorm_gpu.py @@ -1,5 +1,6 @@ from __future__ import annotations +import ctypes import os import socket @@ -71,6 +72,23 @@ def _assert_close( torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) +def _local_eager_words( + channel, stream: torch.cuda.Stream, *, offset: int = 0 +) -> tuple[int, int]: + assert channel._ipc is not None + assert channel._eager_ptrs is not None + words = (ctypes.c_uint64(), ctypes.c_uint64()) + for word, slot_ptrs in zip(words, channel._eager_ptrs, strict=True): + channel._ipc.cudaMemcpyAsync( + ctypes.addressof(word), + slot_ptrs[channel.rank] + offset, + ctypes.sizeof(word), + int(stream.cuda_stream), + ) + stream.synchronize() + return words[0].value, words[1].value + + def _cuda_graph_kernel_count(graph: torch.cuda.CUDAGraph) -> int: graph_handle = graph.raw_cuda_graph() result, _, num_nodes = cudart.cudaGraphGetNodes(graph_handle) @@ -155,7 +173,7 @@ def _run_graph( pool.for_stream() graph = torch.cuda.CUDAGraph(keep_graph=True) - with pool.capture(), torch.cuda.graph(graph): + with pool.capture() as channel, torch.cuda.graph(graph): pool.all_reduce_fused_add_rms_norm( inp, residual, @@ -164,7 +182,13 @@ def _run_graph( out=out, residual_out=residual, ) - assert _cuda_graph_kernel_count(graph) == 1 + assert _cuda_graph_kernel_count(graph) == 2 + stream = torch.cuda.current_stream(device) + offset = 0 + if os.getenv("SPARKINFER_PCIE_ONESHOT_PUSH", "0") not in ("", "0"): + remote_source = (rank + 1) % dist.get_world_size() + offset = remote_source * inp.numel() * inp.element_size() + snapshots = [_local_eager_words(channel, stream, offset=offset)] for iteration in range(3): next_inp, next_residual, _ = _make_inputs( @@ -187,6 +211,23 @@ def _run_graph( torch.cuda.synchronize(device) _assert_close(out, expected_out, dtype) _assert_close(residual, expected_residual, dtype) + snapshots.append(_local_eager_words(channel, stream, offset=offset)) + + changed_slots = [ + { + slot + for slot, (before, after) in enumerate( + zip(snapshots[index], snapshots[index + 1], strict=True) + ) + if before != after + } + for index in range(3) + ] + assert all(len(changed) == 1 for changed in changed_slots) + assert all( + changed_slots[index] != changed_slots[index + 1] + for index in range(len(changed_slots) - 1) + ) def _worker(rank: int, world_size: int, port: int) -> None: diff --git a/tests/comm/test_pcie_oneshot_torture.py b/tests/comm/test_pcie_oneshot_torture.py index a5ab20a8d..a5ac844a3 100644 --- a/tests/comm/test_pcie_oneshot_torture.py +++ b/tests/comm/test_pcie_oneshot_torture.py @@ -1,5 +1,6 @@ from __future__ import annotations +import ctypes import os import socket @@ -16,9 +17,15 @@ reason="set SPARKINFER_RUN_PCIE_ONESHOT_TORTURE=1 to run PCIe oneshot CUDA torture tests", ) -TORTURE_EAGER_ITERS = int(os.getenv("SPARKINFER_PCIE_ONESHOT_TORTURE_EAGER_ITERS", "256")) -TORTURE_GRAPH_REPLAYS = int(os.getenv("SPARKINFER_PCIE_ONESHOT_TORTURE_GRAPH_REPLAYS", "256")) -TORTURE_MULTISTREAM_ITERS = int(os.getenv("SPARKINFER_PCIE_ONESHOT_TORTURE_MULTISTREAM_ITERS", "256")) +TORTURE_EAGER_ITERS = int( + os.getenv("SPARKINFER_PCIE_ONESHOT_TORTURE_EAGER_ITERS", "256") +) +TORTURE_GRAPH_REPLAYS = int( + os.getenv("SPARKINFER_PCIE_ONESHOT_TORTURE_GRAPH_REPLAYS", "256") +) +TORTURE_MULTISTREAM_ITERS = int( + os.getenv("SPARKINFER_PCIE_ONESHOT_TORTURE_MULTISTREAM_ITERS", "256") +) def _free_port() -> int: @@ -36,7 +43,26 @@ def _rank_sum(world_size: int) -> int: return world_size * (world_size - 1) // 2 -def _run_eager(pool: PCIeOneshotAllReducePool, device: torch.device, rank: int, world_size: int) -> None: +def _local_eager_words( + channel, stream: torch.cuda.Stream, *, offset: int = 0 +) -> tuple[int, int]: + assert channel._ipc is not None + assert channel._eager_ptrs is not None + words = (ctypes.c_uint64(), ctypes.c_uint64()) + for word, slot_ptrs in zip(words, channel._eager_ptrs, strict=True): + channel._ipc.cudaMemcpyAsync( + ctypes.addressof(word), + slot_ptrs[channel.rank] + offset, + ctypes.sizeof(word), + int(stream.cuda_stream), + ) + stream.synchronize() + return words[0].value, words[1].value + + +def _run_eager( + pool: PCIeOneshotAllReducePool, device: torch.device, rank: int, world_size: int +) -> None: dtypes = (torch.float16, torch.bfloat16, torch.float32) numels = (8, 256, 4096, 32768) rank_sum = _rank_sum(world_size) @@ -60,7 +86,6 @@ def _run_graph_scratch_reuse( world_size: int, ) -> None: stream = torch.cuda.Stream(device=device) - channel = pool.for_stream(stream) rank_sum = _rank_sum(world_size) layers = 17 numel = 4096 @@ -77,7 +102,7 @@ def fill_sources(iteration: int) -> None: torch.cuda.synchronize(device) graph = torch.cuda.CUDAGraph() - with pool.capture(stream), torch.cuda.graph(graph, stream=stream): + with pool.capture(stream) as channel, torch.cuda.graph(graph, stream=stream): for layer in range(layers): scratch.copy_(sources[layer]) channel.all_reduce(scratch, out=outs[layer]) @@ -91,6 +116,30 @@ def fill_sources(iteration: int) -> None: base = float((iteration % 32) * 5 + layer) _assert_constant(out, world_size * base + rank_sum) + # A graph with an odd number of collectives must change its selected + # staging slot on every replay. Host-side selection would freeze here. + snapshots = [_local_eager_words(channel, stream)] + for iteration in (101, 102): + fill_sources(iteration) + graph.replay() + stream.synchronize() + snapshots.append(_local_eager_words(channel, stream)) + for layer, out in enumerate(outs): + base = float((iteration % 32) * 5 + layer) + _assert_constant(out, world_size * base + rank_sum) + changed_slots = [ + { + slot + for slot, (before, after) in enumerate( + zip(snapshots[index], snapshots[index + 1], strict=True) + ) + if before != after + } + for index in range(2) + ] + assert all(len(changed) == 1 for changed in changed_slots) + assert changed_slots[0] != changed_slots[1] + def _run_multistream( pool: PCIeOneshotAllReducePool, diff --git a/tests/comm/test_pcie_twoshot.py b/tests/comm/test_pcie_twoshot.py index 0862089c5..e45c01bb6 100644 --- a/tests/comm/test_pcie_twoshot.py +++ b/tests/comm/test_pcie_twoshot.py @@ -6,6 +6,7 @@ --nproc-per-node=8 tests/distributed/test_pcie_twoshot.py """ +import ctypes import os import time @@ -21,6 +22,64 @@ ROW_ELEMS = 6144 +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def _local_staging_words( + pool: PCIeTwoShotSP, stream: torch.cuda.Stream +) -> tuple[int, int]: + assert len(pool._owned_buffers) == 1 + max_rows_per_rank = pool.max_rows // pool.world_size + pack_stride = _align_up(max_rows_per_rank * (pool.row_elems // 16), 16) + payload_bytes = pool.world_size * pack_stride * 16 + scale_offset = _align_up(payload_bytes, 256) + scale_stride = _align_up(max_rows_per_rank, 64) + slot_bytes = _align_up( + scale_offset + pool.world_size * scale_stride * 4, + 256, + ) + signal_bytes = _align_up(int(pool._ext.meta_size()), 256) + remote_source = (pool.rank + 1) % pool.world_size + source_offset = remote_source * pack_stride * 16 + words = (ctypes.c_uint64(), ctypes.c_uint64()) + local_ptr = pool._owned_buffers[0].local_ptr + for word, offset in zip( + words, + ( + signal_bytes + source_offset, + signal_bytes + slot_bytes + source_offset, + ), + strict=True, + ): + pool._ipc.cudaMemcpyAsync( + ctypes.addressof(word), + local_ptr + offset, + ctypes.sizeof(word), + int(stream.cuda_stream), + ) + stream.synchronize() + return words[0].value, words[1].value + + +def _assert_alternating_slots(snapshots: list[tuple[int, int]]) -> None: + changed_slots = [ + { + slot + for slot, (before, after) in enumerate( + zip(snapshots[index], snapshots[index + 1], strict=True) + ) + if before != after + } + for index in range(len(snapshots) - 1) + ] + assert all(len(changed) == 1 for changed in changed_slots) + assert all( + changed_slots[index] != changed_slots[index + 1] + for index in range(len(changed_slots) - 1) + ) + + def _partial(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) @@ -127,6 +186,56 @@ def _check_graph_capture(pool: PCIeTwoShotSP, rank: int, world: int) -> None: ) assert torch.equal(ag_out, ag_ref) + # Capture a single collective so adjacent replays must alternate the + # device-selected staging slot. The two-op graph above has even parity. + odd_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(odd_graph): + pool.all_gather_fp8(ag_q, ag_s, ag_out) + + for value in (11, 12): + ag_q.fill_(float(value + rank)) + ag_s.fill_(1.0) + pool.all_gather_fp8(ag_q, ag_s, ag_out) + torch.cuda.synchronize() + snapshots = [_local_staging_words(pool, torch.cuda.current_stream(device))] + + for value in (1, 2): + ag_q.fill_(float(value + rank)) + ag_s.fill_(1.0) + dist.barrier() + odd_graph.replay() + torch.cuda.synchronize() + snapshots.append(_local_staging_words(pool, torch.cuda.current_stream(device))) + for source_rank in range(world): + lo = source_rank * rows_per_rank + hi = lo + rows_per_rank + assert torch.all(ag_out[lo:hi] == float(value + source_rank)) + + _assert_alternating_slots(snapshots) + + rs_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(rs_graph): + pool.reduce_scatter_fp8(q_in, s_in, rs_out) + + for value in (21, 22): + q_in.fill_(float(value + rank)) + s_in.fill_(1.0) + pool.reduce_scatter_fp8(q_in, s_in, rs_out) + torch.cuda.synchronize() + snapshots = [_local_staging_words(pool, torch.cuda.current_stream(device))] + + rank_sum = world * (world - 1) // 2 + for value in (3, 4): + q_in.fill_(float(value + rank)) + s_in.fill_(1.0) + dist.barrier() + rs_graph.replay() + torch.cuda.synchronize() + snapshots.append(_local_staging_words(pool, torch.cuda.current_stream(device))) + assert torch.all(rs_out == float(world * value + rank_sum)) + + _assert_alternating_slots(snapshots) + def _bench(pool: PCIeTwoShotSP, rank: int, world: int) -> None: device = pool.device From 25b7810e0f21cc80cd30e9d85728193b2515f585 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Fri, 31 Jul 2026 01:27:09 +0000 Subject: [PATCH 09/11] refactor(pcie): retire host slot selector after DCP migration --- sparkinfer/comm/pcie/pcie_dcp_a2a.py | 10 ++--- sparkinfer/comm/pcie/two_slot_selector.h | 38 ------------------- tests/comm/test_pcie_dcp_a2a.py | 4 +- tests/comm/test_pcie_dcp_a2a_gpu.py | 21 +++-------- tests/comm/test_pcie_two_slot_selector.py | 46 ----------------------- 5 files changed, 11 insertions(+), 108 deletions(-) delete mode 100644 sparkinfer/comm/pcie/two_slot_selector.h delete mode 100644 tests/comm/test_pcie_two_slot_selector.py diff --git a/sparkinfer/comm/pcie/pcie_dcp_a2a.py b/sparkinfer/comm/pcie/pcie_dcp_a2a.py index 1fe8a4067..85f49be5f 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_a2a.py +++ b/sparkinfer/comm/pcie/pcie_dcp_a2a.py @@ -37,15 +37,15 @@ def _is_supported_bhd_layout(tensor: torch.Tensor) -> bool: return False batch, heads, head_dim = (int(value) for value in tensor.shape) stride_batch, stride_head, _ = (int(value) for value in tensor.stride()) - packed_token_major = ( - stride_batch == heads * head_dim and stride_head == head_dim - ) + packed_token_major = stride_batch == heads * head_dim and stride_head == head_dim capacity_strided_head_major = ( stride_batch == head_dim and stride_head >= batch * head_dim and stride_head % 8 == 0 ) return packed_token_major or capacity_strided_head_major + + @dataclass(frozen=True) class _StagingLayout: signal_bytes: int @@ -386,9 +386,7 @@ def _validate( f"output shape must be {expected_out}, got {tuple(out.shape)}" ) if not _is_supported_bhd_layout(partial_output): - raise ValueError( - "partial_output must be packed token-major or head-major" - ) + raise ValueError("partial_output must be packed token-major or head-major") if not partial_lse.is_contiguous(): raise ValueError("partial_lse must be contiguous") if not _is_supported_bhd_layout(out): diff --git a/sparkinfer/comm/pcie/two_slot_selector.h b/sparkinfer/comm/pcie/two_slot_selector.h deleted file mode 100644 index 21d1c48c9..000000000 --- a/sparkinfer/comm/pcie/two_slot_selector.h +++ /dev/null @@ -1,38 +0,0 @@ -#pragma once - -#include - -namespace sparkinfer::pcie { - -// A double-buffer selector with an explicitly finite state space. Keeping the -// phase as an enum makes signed overflow and out-of-range remainders impossible. -class TwoSlotSelector { - public: - constexpr int take() noexcept { - if (next_ == Next::kFirst) { - next_ = Next::kSecond; - return 0; - } - next_ = Next::kFirst; - return 1; - } - - private: - enum class Next : std::uint32_t { kFirst, kSecond }; - Next next_ = Next::kFirst; -}; - -namespace detail { -constexpr bool two_slot_sequence_is_preserved() { - TwoSlotSelector selector; - return selector.take() == 0 && selector.take() == 1 && - selector.take() == 0 && selector.take() == 1; -} -} // namespace detail - -static_assert(detail::two_slot_sequence_is_preserved(), - "two-slot selection must start at zero and alternate"); -static_assert(sizeof(TwoSlotSelector) == sizeof(std::uint32_t), - "two-slot selector storage must remain 32-bit"); - -} // namespace sparkinfer::pcie diff --git a/tests/comm/test_pcie_dcp_a2a.py b/tests/comm/test_pcie_dcp_a2a.py index ac8ceca4c..445a174a8 100644 --- a/tests/comm/test_pcie_dcp_a2a.py +++ b/tests/comm/test_pcie_dcp_a2a.py @@ -187,9 +187,7 @@ def test_runtime_validates_and_dispatches_to_extension(): def test_runtime_accepts_head_major_input_and_output(): ext = _FakeExt() runtime = _make_runtime(ext) - input_storage = torch.arange( - 32 * 4 * 64, dtype=torch.bfloat16 - ).reshape(32, 4, 64) + input_storage = torch.arange(32 * 4 * 64, dtype=torch.bfloat16).reshape(32, 4, 64) partial_output = input_storage.transpose(0, 1)[:2] partial_lse = torch.zeros(2, 32, dtype=torch.float32) output_storage = torch.empty(16, 2, 64, dtype=torch.bfloat16) diff --git a/tests/comm/test_pcie_dcp_a2a_gpu.py b/tests/comm/test_pcie_dcp_a2a_gpu.py index 8cb3314b7..e8d48faf7 100644 --- a/tests/comm/test_pcie_dcp_a2a_gpu.py +++ b/tests/comm/test_pcie_dcp_a2a_gpu.py @@ -128,6 +128,7 @@ def _local_staging_words(channel, stream: torch.cuda.Stream) -> tuple[int, int]: stream.synchronize() return words[0].value, words[1].value + def _check_eager( pool: PCIeDCPA2APool, rank: int, @@ -209,15 +210,9 @@ def _check_eager_adjacency( device: torch.device, ) -> None: channel = pool.for_stream() - first_query = _rank_query( - 700, rank, world_size, MAX_BATCH, torch.bfloat16, device - ) - second_query = _rank_query( - 701, rank, world_size, MAX_BATCH, torch.bfloat16, device - ) - partial_output, partial_lse = _rank_inputs( - 702, rank, 1, torch.bfloat16, device - ) + first_query = _rank_query(700, rank, world_size, MAX_BATCH, torch.bfloat16, device) + second_query = _rank_query(701, rank, world_size, MAX_BATCH, torch.bfloat16, device) + partial_output, partial_lse = _rank_inputs(702, rank, 1, torch.bfloat16, device) first_gather = torch.empty( MAX_BATCH, TOTAL_HEADS, @@ -243,18 +238,14 @@ def _check_eager_adjacency( expected_first = torch.cat( [ - _rank_query( - 700, source, world_size, MAX_BATCH, torch.bfloat16, device - ) + _rank_query(700, source, world_size, MAX_BATCH, torch.bfloat16, device) for source in range(world_size) ], dim=1, ) expected_second = torch.cat( [ - _rank_query( - 701, source, world_size, MAX_BATCH, torch.bfloat16, device - ) + _rank_query(701, source, world_size, MAX_BATCH, torch.bfloat16, device) for source in range(world_size) ], dim=1, diff --git a/tests/comm/test_pcie_two_slot_selector.py b/tests/comm/test_pcie_two_slot_selector.py deleted file mode 100644 index 485a3e9d0..000000000 --- a/tests/comm/test_pcie_two_slot_selector.py +++ /dev/null @@ -1,46 +0,0 @@ -from __future__ import annotations - -import subprocess -from pathlib import Path - - -def test_two_slot_selector_preserves_phase_for_long_run(tmp_path: Path) -> None: - repo_root = Path(__file__).resolve().parents[2] - source = tmp_path / "two_slot_selector_test.cpp" - binary = tmp_path / "two_slot_selector_test" - source.write_text( - r""" -#include -#include "sparkinfer/comm/pcie/two_slot_selector.h" - -int main() { - sparkinfer::pcie::TwoSlotSelector selector; - constexpr std::uint64_t kTransitions = 10'000'003; - for (std::uint64_t call = 0; call < kTransitions; ++call) { - if (selector.take() != static_cast(call & 1u)) return 1; - } - return selector.take() == static_cast(kTransitions & 1u) ? 0 : 2; -} -""", - encoding="utf-8", - ) - - subprocess.run( - [ - "c++", - "-std=c++17", - "-O3", - "-Wall", - "-Wextra", - "-Werror", - "-I", - str(repo_root), - str(source), - "-o", - str(binary), - ], - check=True, - capture_output=True, - text=True, - ) - subprocess.run([str(binary)], check=True) From 5c4f8b01962c42d3aab8ac36b1e663974ad2537b Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Fri, 31 Jul 2026 01:57:21 +0000 Subject: [PATCH 10/11] perf(pcie): fold staging selection into collectives --- sparkinfer/comm/pcie/pcie_oneshot.cu | 48 ++++++++++--------- sparkinfer/comm/pcie/pcie_twoshot.cu | 35 +++++++++----- .../test_pcie_oneshot_fused_rmsnorm_gpu.py | 2 +- 3 files changed, 50 insertions(+), 35 deletions(-) diff --git a/sparkinfer/comm/pcie/pcie_oneshot.cu b/sparkinfer/comm/pcie/pcie_oneshot.cu index a6e78f081..aeab3387d 100644 --- a/sparkinfer/comm/pcie/pcie_oneshot.cu +++ b/sparkinfer/comm/pcie/pcie_oneshot.cu @@ -77,8 +77,8 @@ static bool oneshot_push_enabled() { } struct Signal { - alignas(128) FlagType staging_generation; - FlagType active_staging_slot; + alignas(128) FlagType staging_arrive; + FlagType staging_generation; alignas(128) FlagType self_counter[kMaxBlocks][kMaxRanks]; alignas(128) FlagType peer_counter[2][kMaxBlocks][kMaxRanks * kFlagStride]; alignas(128) FlagType rms_arrive[kMaxBlocks]; @@ -184,25 +184,6 @@ static DINLINE FlagType ld_flag_relaxed(FlagType* flag_addr) { return flag; } -__global__ void advance_staging_slot(Signal* self_sg) { - if (blockIdx.x == 0 && threadIdx.x == 0) { - const FlagType generation = self_sg->staging_generation; - self_sg->active_staging_slot = generation & FlagType{1}; - self_sg->staging_generation = generation + FlagType{1}; - } -} - -template -DINLINE void select_rank_data(RankData& selected, const DoubleRankData& options, Signal* self_sg) { - if (threadIdx.x == 0) { - const int slot = int(self_sg->active_staging_slot); - const RankData* source = options.slots[slot]; -#pragma unroll - for (int peer = 0; peer < ngpus; ++peer) selected.ptrs[peer] = source->ptrs[peer]; - } - __syncthreads(); -} - static DINLINE FlagType ld_flag_relaxed_gpu(FlagType* flag_addr) { FlagType flag; asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(flag) : "l"(flag_addr) : "memory"); @@ -215,6 +196,29 @@ static DINLINE FlagType ld_flag_acquire_gpu(FlagType* flag_addr) { return flag; } +template +DINLINE void select_rank_data(RankData& selected, const DoubleRankData& options, Signal* self_sg) { + if (threadIdx.x == 0) { + // Every supported staging grid fits concurrently on the target Blackwell + // GPU. The last arriving block publishes one launch-wide generation, so + // variable grid sizes cannot split a launch across two slabs. + const FlagType generation = ld_flag_acquire_gpu(&self_sg->staging_generation); + const FlagType prior = atomicAdd(&self_sg->staging_arrive, FlagType{1}); + if (prior == static_cast(gridDim.x - 1)) { + self_sg->staging_arrive = 0; + __threadfence(); + atomicAdd(&self_sg->staging_generation, FlagType{1}); + } else { + while (ld_flag_acquire_gpu(&self_sg->staging_generation) == generation) { + } + } + const RankData* source = options.slots[generation & FlagType{1}]; +#pragma unroll + for (int peer = 0; peer < ngpus; ++peer) selected.ptrs[peer] = source->ptrs[peer]; + } + __syncthreads(); +} + template DINLINE void multi_gpu_barrier(const RankSignals& sg, Signal* self_sg, int rank) { if constexpr (!is_start) __syncthreads(); @@ -738,7 +742,6 @@ class PCIeAllreduce { size /= d; int blocks = std::min(block_limit, (size + threads - 1) / threads); blocks = std::max(blocks, 1); - if (stage_input) advance_staging_slot<<<1, 1, 0, stream>>>(self_sg_); #define KL(ngpus) \ do { \ @@ -836,7 +839,6 @@ class PCIeAllreduce { const int mode = !use_eager_staging ? kModeRegistered : oneshot_push_enabled() ? kModeStagePush : kModeStagePull; - if (use_eager_staging) advance_staging_slot<<<1, 1, 0, stream>>>(self_sg_); #define KL(ngpus, SINGLE, MODE) \ pcie_allreduce_fused_add_rms_norm_kernel<<>>( \ diff --git a/sparkinfer/comm/pcie/pcie_twoshot.cu b/sparkinfer/comm/pcie/pcie_twoshot.cu index c12e11880..b753f0e1b 100644 --- a/sparkinfer/comm/pcie/pcie_twoshot.cu +++ b/sparkinfer/comm/pcie/pcie_twoshot.cu @@ -58,8 +58,8 @@ using FlagType = uint32_t; constexpr int kFlagStride = 32; struct Signal { - alignas(128) FlagType staging_generation; - FlagType active_staging_slot; + alignas(128) FlagType staging_arrive; + FlagType staging_generation; alignas(128) FlagType self_counter[kMaxBlocks][kMaxRanks]; alignas(128) FlagType peer_counter[2][kMaxBlocks][kMaxRanks * kFlagStride]; }; @@ -88,18 +88,29 @@ static DINLINE FlagType ld_flag_relaxed(FlagType* flag_addr) { return flag; } -__global__ void advance_staging_slot(Signal* self_sg) { - if (blockIdx.x == 0 && threadIdx.x == 0) { - const FlagType generation = self_sg->staging_generation; - self_sg->active_staging_slot = generation & FlagType{1}; - self_sg->staging_generation = generation + FlagType{1}; - } +static DINLINE FlagType ld_flag_acquire_gpu(FlagType* flag_addr) { + FlagType flag; + asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(flag) : "l"(flag_addr) : "memory"); + return flag; } template DINLINE void select_rank_ptrs(RankPtrs& selected, const DoubleRankPtrs& options, Signal* self_sg) { if (threadIdx.x == 0) { - const int slot = int(self_sg->active_staging_slot); + // kMaxBlocks fits concurrently on the target Blackwell GPU. Rendezvous + // before publishing the generation so every block in this launch selects + // one slab even when the preceding launch used a different grid size. + const FlagType generation = ld_flag_acquire_gpu(&self_sg->staging_generation); + const FlagType prior = atomicAdd(&self_sg->staging_arrive, FlagType{1}); + if (prior == static_cast(gridDim.x - 1)) { + self_sg->staging_arrive = 0; + __threadfence(); + atomicAdd(&self_sg->staging_generation, FlagType{1}); + } else { + while (ld_flag_acquire_gpu(&self_sg->staging_generation) == generation) { + } + } + const int slot = int(generation & FlagType{1}); #pragma unroll for (int peer = 0; peer < ngpus; ++peer) selected.ptrs[peer] = options.slots[slot].ptrs[peer]; } @@ -340,10 +351,11 @@ class PCIeTwoShot { throw std::runtime_error("num_rows must be divisible by world size"); const int64_t rows_per_rank = num_rows / world_size_; check_shard(rows_per_rank, row_elems); + if (block_limit > kMaxBlocks) + throw std::runtime_error("pcie_twoshot block limit exceeds signal capacity"); const int64_t shard_packs = rows_per_rank * (row_elems / 16); int blocks = std::max(1, std::min(block_limit, (shard_packs + threads - 1) / threads)); - advance_staging_slot<<<1, 1, 0, stream>>>(self_sg_); dispatch([&](auto ng) { constexpr int ngpus = decltype(ng)::value; rs_fp8_kernel<<>>( @@ -356,10 +368,11 @@ class PCIeTwoShot { void all_gather(cudaStream_t stream, const void* payload, const void* scale, void* out, int64_t rows_per_rank, int64_t row_elems, int threads, int block_limit) { check_shard(rows_per_rank, row_elems); + if (block_limit > kMaxBlocks) + throw std::runtime_error("pcie_twoshot block limit exceeds signal capacity"); const int64_t shard_packs = rows_per_rank * (row_elems / 16); int blocks = std::max(1, std::min(block_limit, (shard_packs + threads - 1) / threads)); - advance_staging_slot<<<1, 1, 0, stream>>>(self_sg_); dispatch([&](auto ng) { constexpr int ngpus = decltype(ng)::value; ag_fp8_kernel<<>>( diff --git a/tests/comm/test_pcie_oneshot_fused_rmsnorm_gpu.py b/tests/comm/test_pcie_oneshot_fused_rmsnorm_gpu.py index 508895978..dfa3f52a2 100644 --- a/tests/comm/test_pcie_oneshot_fused_rmsnorm_gpu.py +++ b/tests/comm/test_pcie_oneshot_fused_rmsnorm_gpu.py @@ -182,7 +182,7 @@ def _run_graph( out=out, residual_out=residual, ) - assert _cuda_graph_kernel_count(graph) == 2 + assert _cuda_graph_kernel_count(graph) == 1 stream = torch.cuda.current_stream(device) offset = 0 if os.getenv("SPARKINFER_PCIE_ONESHOT_PUSH", "0") not in ("", "0"): From 8fd90a4f687895e705bdf52fcff84ef653a5abf0 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Thu, 30 Jul 2026 23:35:24 -0400 Subject: [PATCH 11/11] test(pcie): fix replay slot observation --- tests/comm/test_pcie_oneshot_torture.py | 46 +++++++++++++++---------- 1 file changed, 28 insertions(+), 18 deletions(-) diff --git a/tests/comm/test_pcie_oneshot_torture.py b/tests/comm/test_pcie_oneshot_torture.py index a5ac844a3..8512a0c8a 100644 --- a/tests/comm/test_pcie_oneshot_torture.py +++ b/tests/comm/test_pcie_oneshot_torture.py @@ -109,36 +109,46 @@ def fill_sources(iteration: int) -> None: stream.synchronize() for iteration in range(TORTURE_GRAPH_REPLAYS): - fill_sources(iteration) - graph.replay() + with torch.cuda.stream(stream): + fill_sources(iteration) + graph.replay() stream.synchronize() for layer, out in enumerate(outs): base = float((iteration % 32) * 5 + layer) _assert_constant(out, world_size * base + rank_sum) - # A graph with an odd number of collectives must change its selected - # staging slot on every replay. Host-side selection would freeze here. + # A graph with an odd number of collectives must swap which slot receives + # the final layer on every replay. Both slots are overwritten by this + # 17-layer graph, so merely counting changed slots cannot identify the + # selected parity: the final two layer markers must exchange positions. snapshots = [_local_eager_words(channel, stream)] + expected_words = [ + tuple(int(source.view(torch.uint64)[0].item()) for source in sources[-2:]) + ] for iteration in (101, 102): - fill_sources(iteration) - graph.replay() + with torch.cuda.stream(stream): + fill_sources(iteration) + graph.replay() stream.synchronize() snapshots.append(_local_eager_words(channel, stream)) + expected_words.append( + tuple(int(source.view(torch.uint64)[0].item()) for source in sources[-2:]) + ) for layer, out in enumerate(outs): base = float((iteration % 32) * 5 + layer) _assert_constant(out, world_size * base + rank_sum) - changed_slots = [ - { - slot - for slot, (before, after) in enumerate( - zip(snapshots[index], snapshots[index + 1], strict=True) - ) - if before != after - } - for index in range(2) - ] - assert all(len(changed) == 1 for changed in changed_slots) - assert changed_slots[0] != changed_slots[1] + + final_slots = [] + for snapshot, expected in zip(snapshots, expected_words, strict=True): + assert set(snapshot) == set(expected), ( + f"staging slots {snapshot} do not contain the final layer markers " + f"{expected}" + ) + final_slots.append(snapshot.index(expected[-1])) + assert all( + final_slots[index] != final_slots[index + 1] + for index in range(len(final_slots) - 1) + ), f"odd-collective graph did not alternate final staging slots: {final_slots}" def _run_multistream(