From 04d8019e34a36f504862df31ebf87a55f7d75c81 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 25 Jul 2026 01:36:02 +0000 Subject: [PATCH 1/3] perf(pcie): add exact DCP top-k owner exchange --- sparkinfer/comm/pcie/__init__.py | 7 +- sparkinfer/comm/pcie/api.py | 6 + sparkinfer/comm/pcie/pcie_dcp_topk.cu | 309 +++++++++++++++++++ sparkinfer/comm/pcie/pcie_dcp_topk.py | 419 ++++++++++++++++++++++++++ tests/comm/test_pcie_dcp_topk.py | 151 ++++++++++ tests/comm/test_pcie_dcp_topk_gpu.py | 186 ++++++++++++ tests/test_packaging.py | 1 + 7 files changed, 1078 insertions(+), 1 deletion(-) create mode 100644 sparkinfer/comm/pcie/pcie_dcp_topk.cu create mode 100644 sparkinfer/comm/pcie/pcie_dcp_topk.py create mode 100644 tests/comm/test_pcie_dcp_topk.py create mode 100644 tests/comm/test_pcie_dcp_topk_gpu.py diff --git a/sparkinfer/comm/pcie/__init__.py b/sparkinfer/comm/pcie/__init__.py index 892e8b982..61270bc10 100644 --- a/sparkinfer/comm/pcie/__init__.py +++ b/sparkinfer/comm/pcie/__init__.py @@ -11,6 +11,7 @@ - ``TwoShotReduceScatter``: two-shot sequence-parallel collectives with per-token FP8-e4m3 transport. - ``DcpAllToAll``: DCP attention exchange with fused LSE reduce-scatter. +- ``DcpTopKOwnerExchange``: exact DCP candidate owner staging. Raw CUDA (not CuTe): each class JIT-builds its colocated ``.cu`` via torch.utils.cpp_extension, so nvcc must be available at runtime. @@ -33,12 +34,14 @@ "TwoShotReduceScatter", "DcpAllToAll", "DcpAllToAllPool", + "DcpTopKOwnerExchange", "autotune_dma_crossovers", "parse_oneshot_max_size", "lse_reduce_scatter_reference", + "owner_stage_reference", "is_supported", ), - dtypes=("bf16", "fp32", "fp8_e4m3"), + dtypes=("bf16", "fp32", "fp8_e4m3", "int32"), requires=("multi_gpu",), provenance=Provenance( repo="https://github.com/lukealonso/sparkinfer", @@ -54,6 +57,7 @@ from .api import ( # noqa: F401 DcpAllToAll, DcpAllToAllPool, + DcpTopKOwnerExchange, DmaAllReduce, OneshotAllReduce, OneshotAllReducePool, @@ -61,6 +65,7 @@ autotune_dma_crossovers, is_supported, lse_reduce_scatter_reference, + owner_stage_reference, parse_oneshot_max_size, ) diff --git a/sparkinfer/comm/pcie/api.py b/sparkinfer/comm/pcie/api.py index 3189a2823..b362c0338 100644 --- a/sparkinfer/comm/pcie/api.py +++ b/sparkinfer/comm/pcie/api.py @@ -12,6 +12,10 @@ from .pcie_dcp_a2a import ( lse_reduce_scatter_reference, ) +from .pcie_dcp_topk import ( + PCIeDCPTopKOwnerExchange as DcpTopKOwnerExchange, + owner_stage_reference, +) from .pcie_dma import ( PCIeDmaAllReduce as DmaAllReduce, ) @@ -50,8 +54,10 @@ def is_supported(device=None) -> bool: "TwoShotReduceScatter", "DcpAllToAll", "DcpAllToAllPool", + "DcpTopKOwnerExchange", "autotune_dma_crossovers", "parse_oneshot_max_size", "lse_reduce_scatter_reference", + "owner_stage_reference", "is_supported", ] diff --git a/sparkinfer/comm/pcie/pcie_dcp_topk.cu b/sparkinfer/comm/pcie/pcie_dcp_topk.cu new file mode 100644 index 000000000..410a01e5f --- /dev/null +++ b/sparkinfer/comm/pcie/pcie_dcp_topk.cu @@ -0,0 +1,309 @@ +// Exact owner-sharded transport for DCP sparse top-k. +// +// Candidate exchange runs on the DCP group. Each source writes directly into +// the destination owner's row-major CUDA-IPC slab, which the exact row-top-k +// kernel can consume without a pack, NCCL all-to-all, or unpack. + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#define CHECK_CUDA_SUCCESS(cmd) \ + do { \ + cudaError_t e = cmd; \ + if (e != cudaSuccess) { \ + std::stringstream message; \ + message << cudaGetErrorString(e) << "\\n" << __FILE__ << ':' << __LINE__; \ + throw std::runtime_error(message.str()); \ + } \ + } while (0) + +namespace pcie_dcp_topk { + +constexpr int kMaxBlocks = 128; +constexpr int kMaxRanks = 8; +constexpr int kFlagStride = 32; +using FlagType = uint32_t; + +struct Signal { + alignas(128) FlagType self_counter[kMaxBlocks][kMaxRanks]; + alignas(128) + FlagType peer_counter[2][kMaxBlocks][kMaxRanks * kFlagStride]; +}; + +struct RankSignals { + Signal *signals[kMaxRanks]; +}; + +struct RankStaging { + void *ptrs[kMaxRanks]; +}; + +#define DINLINE __device__ __forceinline__ + +static DINLINE void store_flag(FlagType *address, FlagType value) { + asm volatile("st.relaxed.sys.global.u32 [%1], %0;" : : "r"(value), + "l"(address)); +} +static DINLINE FlagType load_flag(FlagType *address) { + FlagType value; + asm volatile("ld.relaxed.sys.global.u32 %0, [%1];" + : "=r"(value) + : "l"(address)); + return value; +} + +template +DINLINE void block_pair_barrier(const RankSignals &signals, Signal *self, + int rank) { + __syncthreads(); + if (threadIdx.x < world_size) { + __threadfence_system(); + const auto value = + self->self_counter[blockIdx.x][threadIdx.x] += FlagType{1}; + auto *peer = &signals.signals[threadIdx.x] + ->peer_counter[value % 2][blockIdx.x][rank * kFlagStride]; + auto *mine = + &self->peer_counter[value % 2][blockIdx.x][threadIdx.x * kFlagStride]; + store_flag(peer, value); + while (load_flag(mine) != value) { + } + } + __syncthreads(); +} + +DINLINE void block_range(int64_t total, int64_t &begin, int64_t &end) { + const int64_t chunk = (total + gridDim.x - 1) / gridDim.x; + begin = int64_t(blockIdx.x) * chunk; + end = min(begin + chunk, total); +} + +// Destination layout is two independent row-major planes: +// +// indices[owner_row, source_rank, topk] +// scores [owner_row, source_rank, topk] +// +// Every source writes its disjoint source-rank column directly. The owner sees +// the exact rank-major table expected by the NCCL oracle. +template +__global__ void __launch_bounds__(512, 1) stage_owner_candidates_kernel( + const int4 *__restrict__ local_indices, + const int4 *__restrict__ local_scores, RankStaging staging, + RankSignals signals, Signal *self, int rank, int rows, int topk, + int64_t candidate_plane_packs) { + const int owner_rows = rows / world_size; + const int packs_per_row = topk / 4; + const int64_t owner_packs = int64_t(owner_rows) * packs_per_row; + const int64_t output_row_packs = int64_t(world_size) * packs_per_row; + int64_t begin, end; + block_range(owner_packs, begin, end); + +#pragma unroll 1 + for (int step = 0; step < world_size; ++step) { + const int destination = (rank + step) % world_size; + auto *destination_indices = + reinterpret_cast(staging.ptrs[destination]); + auto *destination_scores = destination_indices + candidate_plane_packs; + const int64_t input_offset = int64_t(destination) * owner_packs; + for (int64_t pack = begin + threadIdx.x; pack < end; + pack += blockDim.x) { + const int64_t owner_row = pack / packs_per_row; + const int64_t column_pack = pack - owner_row * packs_per_row; + const int64_t output_offset = + owner_row * output_row_packs + int64_t(rank) * packs_per_row + + column_pack; + destination_indices[output_offset] = local_indices[input_offset + pack]; + destination_scores[output_offset] = local_scores[input_offset + pack]; + } + } + + block_pair_barrier(signals, self, rank); +} + +static void validate_launch(int world_size, int threads, int block_limit) { + if (threads < world_size || threads > 512 || threads % 32 != 0) { + throw std::runtime_error( + "threads must be a multiple of 32 in [32, 512]"); + } + if (block_limit <= 0 || block_limit > kMaxBlocks) { + throw std::runtime_error("invalid block limit"); + } +} + +class PCIeDCPTopKOwnerExchange { + public: + int rank_; + int world_size_; + int max_rows_; + int max_owner_rows_; + int topk_; + int64_t candidate_plane_elems_; + RankSignals signals_{}; + Signal *self_signal_; + RankStaging candidates_[2]{}; + int slot_ = 0; + + PCIeDCPTopKOwnerExchange( + Signal **signals, + const std::vector> &candidate_staging, + int max_rows, int topk, int rank, int world_size) + : rank_(rank), world_size_(world_size), max_rows_(max_rows), + max_owner_rows_(max_rows / world_size), topk_(topk), + candidate_plane_elems_(int64_t(max_owner_rows_) * world_size * topk), + self_signal_(signals[rank]) { + for (int peer = 0; peer < world_size_; ++peer) { + signals_.signals[peer] = signals[peer]; + candidates_[0].ptrs[peer] = candidate_staging[peer][0]; + candidates_[1].ptrs[peer] = candidate_staging[peer][1]; + } + } + + int stage(cudaStream_t stream, const int *local_indices, + const float *local_scores, int rows, int threads, + int block_limit) { + if (rows <= 0 || rows > max_rows_ || rows % world_size_ != 0) { + throw std::runtime_error( + "rows must fit capacity and be divisible by world size"); + } + validate_launch(world_size_, threads, block_limit); + const int64_t owner_packs = + int64_t(rows / world_size_) * (topk_ / 4); + const int blocks = int(std::max( + 1, std::min(block_limit, + (owner_packs + threads - 1) / threads))); + const int slot = slot_++ % 2; + const int64_t candidate_plane_packs = candidate_plane_elems_ / 4; + +#define LAUNCH(world) \ + stage_owner_candidates_kernel<<>>( \ + reinterpret_cast(local_indices), \ + reinterpret_cast(local_scores), candidates_[slot], \ + signals_, self_signal_, rank_, rows, topk_, candidate_plane_packs) + switch (world_size_) { + case 2: + LAUNCH(2); + break; + case 3: + LAUNCH(3); + break; + case 4: + LAUNCH(4); + break; + case 6: + LAUNCH(6); + break; + case 8: + LAUNCH(8); + break; + default: + throw std::runtime_error("unsupported DCP top-k world size"); + } +#undef LAUNCH + CHECK_CUDA_SUCCESS(cudaGetLastError()); + return slot; + } + + void *local_candidate_indices(int slot) const { + return candidates_[slot].ptrs[rank_]; + } + + void *local_candidate_scores(int slot) const { + return static_cast( + static_cast(candidates_[slot].ptrs[rank_]) + + candidate_plane_elems_ * sizeof(int32_t)); + } +}; + +} // namespace pcie_dcp_topk + +using fptr_t = int64_t; + +static fptr_t init_dcp_topk_owner_exchange( + const std::vector &signal_ptrs, + const std::vector &candidate0_ptrs, + const std::vector &candidate1_ptrs, int64_t max_rows, + int64_t topk, int64_t rank) { + const int world_size = int(signal_ptrs.size()); + TORCH_CHECK(world_size == 2 || world_size == 3 || world_size == 4 || + world_size == 6 || world_size == 8); + TORCH_CHECK_EQ(candidate0_ptrs.size(), signal_ptrs.size()); + TORCH_CHECK_EQ(candidate1_ptrs.size(), signal_ptrs.size()); + TORCH_CHECK(rank >= 0 && rank < world_size); + TORCH_CHECK(max_rows > 0 && max_rows % world_size == 0); + TORCH_CHECK(topk > 0 && topk % 4 == 0); + + pcie_dcp_topk::Signal *signals[pcie_dcp_topk::kMaxRanks]; + std::vector> candidates(world_size); + for (int peer = 0; peer < world_size; ++peer) { + signals[peer] = + reinterpret_cast(signal_ptrs[peer]); + candidates[peer] = {reinterpret_cast(candidate0_ptrs[peer]), + reinterpret_cast(candidate1_ptrs[peer])}; + } + return reinterpret_cast( + new pcie_dcp_topk::PCIeDCPTopKOwnerExchange( + signals, candidates, int(max_rows), int(topk), int(rank), + world_size)); +} + +static std::vector stage_owner_candidates( + fptr_t pointer, torch::Tensor &local_indices, torch::Tensor &local_scores, + int64_t threads, int64_t block_limit) { + auto *runtime = + reinterpret_cast(pointer); + const at::cuda::OptionalCUDAGuard device_guard(device_of(local_indices)); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + + TORCH_CHECK(local_indices.is_cuda() && local_scores.is_cuda()); + TORCH_CHECK(local_scores.device() == local_indices.device()); + TORCH_CHECK(local_indices.is_contiguous() && local_scores.is_contiguous()); + TORCH_CHECK_EQ(local_indices.scalar_type(), at::ScalarType::Int); + TORCH_CHECK_EQ(local_scores.scalar_type(), at::ScalarType::Float); + TORCH_CHECK_EQ(local_indices.dim(), 2); + TORCH_CHECK_EQ(local_scores.sizes(), local_indices.sizes()); + TORCH_CHECK_EQ(local_indices.size(1), runtime->topk_); + const int64_t rows = local_indices.size(0); + TORCH_CHECK_GT(rows, 0); + TORCH_CHECK_EQ(rows % runtime->world_size_, 0); + + const int slot = runtime->stage( + stream, reinterpret_cast(local_indices.data_ptr()), + reinterpret_cast(local_scores.data_ptr()), int(rows), + int(threads), int(block_limit)); + const int64_t owner_rows = rows / runtime->world_size_; + const int64_t candidate_width = runtime->world_size_ * runtime->topk_; + auto no_delete = [](void *) {}; + auto indices = torch::from_blob( + runtime->local_candidate_indices(slot), {owner_rows, candidate_width}, + no_delete, local_indices.options()); + auto scores = torch::from_blob( + runtime->local_candidate_scores(slot), {owner_rows, candidate_width}, + no_delete, local_scores.options()); + return {indices, scores}; +} + +static void dispose_owner_exchange(fptr_t pointer) { + delete reinterpret_cast(pointer); +} + +static int64_t meta_size() { return sizeof(pcie_dcp_topk::Signal); } + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("init_dcp_topk_owner_exchange", &init_dcp_topk_owner_exchange, + "initialize exact DCP top-k owner exchange"); + module.def("stage_owner_candidates", &stage_owner_candidates, + "stage exact rank-major candidates on each row owner"); + module.def("dispose_owner_exchange", &dispose_owner_exchange, + "dispose exact DCP top-k owner exchange"); + module.def("meta_size", &meta_size, "signal metadata size"); +} diff --git a/sparkinfer/comm/pcie/pcie_dcp_topk.py b/sparkinfer/comm/pcie/pcie_dcp_topk.py new file mode 100644 index 000000000..7dae26471 --- /dev/null +++ b/sparkinfer/comm/pcie/pcie_dcp_topk.py @@ -0,0 +1,419 @@ +"""Exact owner-sharded PCIe transport for DCP sparse top-k.""" + +from __future__ import annotations + +import os +from contextlib import suppress +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path +from typing import Any, Optional, Sequence + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup +from torch.utils.cpp_extension import load + +from ._cuda_ipc import CudaRTLibrary +from .pcie_oneshot import ( + IPC_SLAB_ALIGNMENT, + PCIeOneshotAllReduce, + _align_up, + _coordinated_close_channels, + _current_stream_key, + _normalize_device, + _OwnedSharedBuffer, +) + + +SUPPORTED_WORLD_SIZES = (2, 3, 4, 6, 8) + + +def _validate_launch_config(*, threads: int, block_limit: int, world_size: int) -> None: + if threads < world_size or threads > 512 or threads % 32 != 0: + raise ValueError("threads must be a multiple of 32 in [32, 512]") + if block_limit <= 0 or block_limit > 128: + raise ValueError("block_limit must be in [1, 128]") + + +@dataclass(frozen=True) +class _DoubleBufferLayout: + signal_bytes: int + staging0_offset: int + staging1_offset: int + slot_bytes: int + slab_bytes: int + plane_bytes: int = 0 + + +def _candidate_staging_layout( + *, + signal_bytes: int, + max_rows: int, + topk: int, + world_size: int, +) -> _DoubleBufferLayout: + if signal_bytes <= 0: + raise ValueError("signal_bytes must be positive") + 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 topk <= 0 or topk % 4 != 0: + raise ValueError("topk must be a positive multiple of 4") + + max_owner_rows = max_rows // world_size + plane_bytes = max_owner_rows * world_size * topk * 4 + slot_bytes = _align_up(plane_bytes * 2, IPC_SLAB_ALIGNMENT) + staging0_offset = _align_up(signal_bytes, IPC_SLAB_ALIGNMENT) + staging1_offset = staging0_offset + slot_bytes + return _DoubleBufferLayout( + signal_bytes=signal_bytes, + staging0_offset=staging0_offset, + staging1_offset=staging1_offset, + slot_bytes=slot_bytes, + slab_bytes=staging1_offset + slot_bytes, + plane_bytes=plane_bytes, + ) + + +@lru_cache(maxsize=1) +def _load_extension(): + source = Path(__file__).with_name("pcie_dcp_topk.cu") + verbose = os.getenv("SPARKINFER_PCIE_DCP_TOPK_VERBOSE_BUILD", "0") == "1" + return load( + name="sparkinfer_pcie_dcp_topk_ext", + sources=[str(source)], + extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr"], + extra_ldflags=["-lcuda"], + verbose=verbose, + ) + + +def owner_stage_reference( + rank_indices: torch.Tensor, + rank_scores: torch.Tensor, + owner_rank: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build one DCP owner's exact row-major candidate planes.""" + if rank_indices.shape != rank_scores.shape or rank_indices.ndim != 3: + raise ValueError( + "rank candidates must have matching [world, rows, topk] shapes" + ) + world_size, rows, _ = rank_indices.shape + if rows % world_size != 0: + raise ValueError("rows must be divisible by world size") + if not 0 <= owner_rank < world_size: + raise ValueError("invalid owner rank") + owner_rows = rows // world_size + row_slice = slice(owner_rank * owner_rows, (owner_rank + 1) * owner_rows) + return ( + rank_indices[:, row_slice].transpose(0, 1).flatten(1).contiguous(), + rank_scores[:, row_slice].transpose(0, 1).flatten(1).contiguous(), + ) + + +class _IPCChannel: + def _init_channel( + self, + *, + device: torch.device | int | str, + exchange_group: Optional[ProcessGroup], + ipc: Optional[CudaRTLibrary], + owned_buffers: Optional[Sequence[_OwnedSharedBuffer]], + ext_module: Any, + stream_affine: bool, + ) -> None: + self.device = _normalize_device(device) + self.exchange_group = exchange_group + self._ipc = ipc + self._owned_buffers = list(owned_buffers or ()) + self._ext = ext_module + self._stream_affine = bool(stream_affine) + self._owner_stream_key: Optional[int] = None + self._closed = False + self._ipc_imports_closed = False + self._ipc_exports_freed = False + self._ptr = 0 + + def _bind_stream(self) -> None: + if not self._stream_affine or self.device.type != "cuda": + return + stream_key = _current_stream_key(self.device) + if stream_key is None: + return + if self._owner_stream_key is None: + self._owner_stream_key = int(stream_key) + elif self._owner_stream_key != int(stream_key): + raise RuntimeError( + f"{type(self).__name__} is stream-affine; create one instance " + "per CUDA stream" + ) + + def _dispose(self) -> None: + raise NotImplementedError + + def _close_ipc_imports(self) -> None: + if self._ipc_imports_closed: + return + self._closed = True + if self._ptr: + with suppress(Exception): + self._dispose() + self._ptr = 0 + if self._ipc is not None: + for shared in self._owned_buffers: + for ptr in shared.remote_ptrs: + with suppress(Exception): + self._ipc.cudaIpcCloseMemHandle(ptr) + self._ipc_imports_closed = True + + def _free_ipc_exports(self) -> None: + if self._ipc_exports_freed: + return + self._close_ipc_imports() + if self._ipc is not None: + for shared in self._owned_buffers: + with suppress(Exception): + self._ipc.cudaFree(shared.local_ptr) + self._owned_buffers.clear() + self._ipc_exports_freed = True + + def close(self) -> None: + self._close_ipc_imports() + self._free_ipc_exports() + + def close_coordinated(self) -> None: + """Collectively close peer mappings before freeing exported storage.""" + _coordinated_close_channels( + (self,), + exchange_group=self.exchange_group, + device=self.device, + ) + + def __del__(self) -> None: + with suppress(Exception): + self.close() + + +class PCIeDCPTopKOwnerExchange(_IPCChannel): + """Write exact DCP candidates directly into each row owner's IPC slab.""" + + def __init__( + self, + *, + rank: int, + world_size: int, + device: torch.device | int | str, + signal_ptrs: Sequence[int], + staging0_ptrs: Sequence[int], + staging1_ptrs: Sequence[int], + max_rows: int, + topk: int, + exchange_group: Optional[ProcessGroup] = None, + ipc: Optional[CudaRTLibrary] = None, + owned_buffers: Optional[Sequence[_OwnedSharedBuffer]] = None, + ext_module=None, + stream_affine: bool = True, + ) -> None: + if world_size not in SUPPORTED_WORLD_SIZES: + raise ValueError(f"unsupported world size {world_size}") + if not 0 <= rank < world_size: + raise ValueError(f"invalid rank {rank} for 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 topk <= 0 or topk % 4 != 0: + raise ValueError("topk must be a positive multiple of 4") + if not ( + len(signal_ptrs) == len(staging0_ptrs) == len(staging1_ptrs) == world_size + ): + raise ValueError("signal and staging pointers must match world size") + + ext = ext_module or _load_extension() + self._init_channel( + device=device, + exchange_group=exchange_group, + ipc=ipc, + owned_buffers=owned_buffers, + ext_module=ext, + stream_affine=stream_affine, + ) + self.rank = int(rank) + self.world_size = int(world_size) + self.max_rows = int(max_rows) + self.max_owner_rows = self.max_rows // self.world_size + self.topk = int(topk) + self._ptr = self._ext.init_dcp_topk_owner_exchange( + list(signal_ptrs), + list(staging0_ptrs), + list(staging1_ptrs), + self.max_rows, + self.topk, + self.rank, + ) + + @classmethod + def from_exchange_group( + cls, + *, + exchange_group: ProcessGroup, + device: torch.device | int | str, + max_rows: int, + topk: int, + ext_module=None, + stream_affine: bool = True, + ) -> "PCIeDCPTopKOwnerExchange": + rank = dist.get_rank(group=exchange_group) + world_size = dist.get_world_size(group=exchange_group) + device_obj = _normalize_device(device) + if device_obj.type != "cuda": + raise ValueError("DCP top-k owner exchange requires a CUDA device") + ipc = CudaRTLibrary() + ipc.cudaSetDevice(device_obj.index or 0) + ext = ext_module or _load_extension() + layout = _candidate_staging_layout( + signal_bytes=int(ext.meta_size()), + max_rows=max_rows, + topk=topk, + world_size=world_size, + ) + owned: list[_OwnedSharedBuffer] = [] + try: + slab = PCIeOneshotAllReduce._allocate_shared_buffer( + exchange_group, + layout.slab_bytes, + zero_fill=True, + ipc=ipc, + ) + owned.append(slab) + return cls( + rank=rank, + world_size=world_size, + device=device_obj, + signal_ptrs=slab.peer_ptrs, + staging0_ptrs=tuple( + ptr + layout.staging0_offset for ptr in slab.peer_ptrs + ), + staging1_ptrs=tuple( + ptr + layout.staging1_offset for ptr in slab.peer_ptrs + ), + max_rows=max_rows, + topk=topk, + exchange_group=exchange_group, + ipc=ipc, + owned_buffers=owned, + ext_module=ext, + stream_affine=stream_affine, + ) + except Exception: + _release_failed_allocations(owned, ipc) + raise + + @classmethod + def from_process_group( + cls, + *, + process_group: ProcessGroup, + device: torch.device | int | str, + max_rows: int, + topk: int, + ext_module=None, + stream_affine: bool = True, + ) -> "PCIeDCPTopKOwnerExchange": + return cls.from_exchange_group( + exchange_group=process_group, + device=device, + max_rows=max_rows, + topk=topk, + ext_module=ext_module, + stream_affine=stream_affine, + ) + + def stage_candidates( + self, + local_indices: torch.Tensor, + local_scores: torch.Tensor, + *, + threads: int = 512, + block_limit: int = 128, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self._closed: + raise RuntimeError("PCIeDCPTopKOwnerExchange is closed") + if local_indices.device != self.device or local_scores.device != self.device: + raise ValueError("inputs must be on the runtime device") + if local_indices.dtype != torch.int32: + raise ValueError("local_indices must be int32") + if local_scores.dtype != torch.float32: + raise ValueError("local_scores must be float32") + if local_indices.shape != local_scores.shape or local_indices.ndim != 2: + raise ValueError("inputs must have matching [rows, topk] shapes") + rows, topk = (int(value) for value in local_indices.shape) + if rows <= 0 or rows > self.max_rows or rows % self.world_size != 0: + raise ValueError( + f"rows must be in (0, {self.max_rows}] and divisible by world_size" + ) + if topk != self.topk: + raise ValueError(f"topk {topk} does not match configured topk {self.topk}") + if not local_indices.is_contiguous() or not local_scores.is_contiguous(): + raise ValueError("inputs must be contiguous") + _validate_launch_config( + threads=int(threads), + block_limit=int(block_limit), + world_size=self.world_size, + ) + self._bind_stream() + candidate_indices, candidate_scores = self._ext.stage_owner_candidates( + self._ptr, + local_indices, + local_scores, + int(threads), + int(block_limit), + ) + expected = (rows // self.world_size, self.world_size * self.topk) + _validate_candidate_view( + candidate_indices, + expected=expected, + dtype=torch.int32, + device=self.device, + label="index", + ) + _validate_candidate_view( + candidate_scores, + expected=expected, + dtype=torch.float32, + device=self.device, + label="score", + ) + return candidate_indices, candidate_scores + + def _dispose(self) -> None: + self._ext.dispose_owner_exchange(self._ptr) + + +def _validate_candidate_view( + tensor: torch.Tensor, + *, + expected: tuple[int, int], + dtype: torch.dtype, + device: torch.device, + label: str, +) -> None: + if ( + tuple(tensor.shape) != expected + or tensor.dtype != dtype + or tensor.device != device + or not tensor.is_contiguous() + ): + raise RuntimeError(f"extension returned an invalid candidate {label} view") + + +def _release_failed_allocations( + owned: Sequence[_OwnedSharedBuffer], + ipc: CudaRTLibrary, +) -> None: + for shared in owned: + for ptr in shared.remote_ptrs: + with suppress(Exception): + ipc.cudaIpcCloseMemHandle(ptr) + with suppress(Exception): + ipc.cudaFree(shared.local_ptr) diff --git a/tests/comm/test_pcie_dcp_topk.py b/tests/comm/test_pcie_dcp_topk.py new file mode 100644 index 000000000..e1c73d7f2 --- /dev/null +++ b/tests/comm/test_pcie_dcp_topk.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import pytest +import torch + +from sparkinfer.comm.pcie.pcie_dcp_topk import ( + PCIeDCPTopKOwnerExchange, + _candidate_staging_layout, + owner_stage_reference, +) + + +class _FakeExt: + def __init__(self) -> None: + self.stage_calls: list[tuple] = [] + self.dispose_owner_calls: list[int] = [] + + def init_dcp_topk_owner_exchange(self, *args): + return 1234 + + def stage_owner_candidates( + self, pointer, local_indices, local_scores, threads, block_limit + ): + world_size = 2 + rank = 1 + owner_rows = local_indices.shape[0] // world_size + row_slice = slice(rank * owner_rows, (rank + 1) * owner_rows) + indices = local_indices[row_slice].repeat(1, world_size) + scores = local_scores[row_slice].repeat(1, world_size) + self.stage_calls.append((pointer, threads, block_limit)) + return indices, scores + + def dispose_owner_exchange(self, pointer): + self.dispose_owner_calls.append(pointer) + + +def _make_owner(ext: _FakeExt | None = None) -> PCIeDCPTopKOwnerExchange: + return PCIeDCPTopKOwnerExchange( + rank=1, + world_size=2, + device="cpu", + signal_ptrs=(10, 20), + staging0_ptrs=(30, 40), + staging1_ptrs=(50, 60), + max_rows=8, + topk=4, + ext_module=ext or _FakeExt(), + ) + + +def test_candidate_staging_layout_is_aligned(): + owner = _candidate_staging_layout( + signal_bytes=513, + max_rows=8, + topk=4, + world_size=2, + ) + assert owner.staging0_offset == 768 + assert owner.plane_bytes == 128 + assert owner.slot_bytes == 256 + assert owner.staging1_offset == 1024 + assert owner.slab_bytes == 1280 + + +def test_owner_stage_reference_preserves_rank_major_order_and_score_bits(): + world_size, rows, topk = 4, 8, 4 + indices = torch.empty(world_size, rows, topk, dtype=torch.int32) + score_bits = torch.empty(world_size, rows, topk, dtype=torch.int32) + for rank in range(world_size): + for row in range(rows): + indices[rank, row].fill_(1000 * rank + 10 * row) + score_bits[rank, row].fill_(0x3F000000 + rank * 16 + row) + scores = score_bits.view(torch.float32) + + owner_indices, owner_scores = owner_stage_reference(indices, scores, 2) + + assert owner_indices.shape == (2, world_size * topk) + assert owner_scores.shape == owner_indices.shape + for owner_row, global_row in enumerate((4, 5)): + for rank in range(world_size): + rank_slice = slice(rank * topk, (rank + 1) * topk) + assert torch.equal( + owner_indices[owner_row, rank_slice], indices[rank, global_row] + ) + assert torch.equal( + owner_scores[owner_row, rank_slice].view(torch.int32), + score_bits[rank, global_row], + ) + + +def test_owner_dispatches_and_disposes(): + ext = _FakeExt() + owner = _make_owner(ext) + indices = torch.arange(32, dtype=torch.int32).reshape(8, 4) + scores = torch.arange(32, dtype=torch.float32).reshape(8, 4) / 10 + + candidate_indices, candidate_scores = owner.stage_candidates( + indices, + scores, + threads=128, + block_limit=32, + ) + assert torch.equal(candidate_indices, indices[4:].repeat(1, 2)) + assert torch.equal(candidate_scores, scores[4:].repeat(1, 2)) + assert ext.stage_calls == [(1234, 128, 32)] + + owner.close() + assert ext.dispose_owner_calls == [1234] + + +def test_owner_rejects_invalid_contracts(): + owner = _make_owner() + indices = torch.zeros(8, 4, dtype=torch.int32) + scores = torch.zeros(8, 4, dtype=torch.float32) + + with pytest.raises(ValueError, match="matching"): + owner.stage_candidates(indices, scores[:4]) + with pytest.raises(ValueError, match="matching"): + owner.stage_candidates(indices.flatten(), scores.flatten()) + with pytest.raises(ValueError, match="local_indices must be int32"): + owner.stage_candidates(indices.float(), scores) + with pytest.raises(ValueError, match="local_scores must be float32"): + owner.stage_candidates(indices, scores.bfloat16()) + with pytest.raises(ValueError, match="divisible"): + owner.stage_candidates(indices[:7], scores[:7]) + with pytest.raises(ValueError, match="threads"): + owner.stage_candidates(indices, scores, threads=31) + with pytest.raises(ValueError, match="block_limit"): + owner.stage_candidates(indices, scores, block_limit=129) + + +def test_configuration_rejects_invalid_capacity_and_topk(): + with pytest.raises(ValueError, match="multiple of 4"): + _candidate_staging_layout( + signal_bytes=256, + max_rows=8, + topk=6, + world_size=2, + ) + with pytest.raises(ValueError, match="divisible"): + PCIeDCPTopKOwnerExchange( + rank=0, + world_size=4, + device="cpu", + signal_ptrs=(1, 2, 3, 4), + staging0_ptrs=(5, 6, 7, 8), + staging1_ptrs=(9, 10, 11, 12), + max_rows=6, + topk=4, + ext_module=_FakeExt(), + ) diff --git a/tests/comm/test_pcie_dcp_topk_gpu.py b/tests/comm/test_pcie_dcp_topk_gpu.py new file mode 100644 index 000000000..b31720471 --- /dev/null +++ b/tests/comm/test_pcie_dcp_topk_gpu.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +import os +import socket + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from sparkinfer.comm.pcie.pcie_dcp_topk import ( + PCIeDCPTopKOwnerExchange, + _load_extension, + owner_stage_reference, +) + + +pytestmark = pytest.mark.skipif( + os.getenv("SPARKINFER_RUN_PCIE_DCP_TOPK_TEST") != "1", + reason="set SPARKINFER_RUN_PCIE_DCP_TOPK_TEST=1 to run GPU tests", +) + +MAX_ROWS = 64 +TOPK = 2048 + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _inputs( + step: int, global_rank: int, rows: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + base = torch.arange(rows * TOPK, dtype=torch.int32).reshape(rows, TOPK) + indices = (base + global_rank * 1_000_000 + step * 10_000).to(device) + score_bits = ( + torch.arange(rows * TOPK, dtype=torch.int32).reshape(rows, TOPK) + + 0x3E800000 + + global_rank * 4096 + + step * 64 + ) + return indices, score_bits.to(device).view(torch.float32) + + +def _dcp_groups(tp_world_size: int, dcp_world_size: int): + return [ + dist.new_group( + list(range(start, start + dcp_world_size)), + backend="nccl", + ) + for start in range(0, tp_world_size, dcp_world_size) + ] + + +def _worker( + rank: int, + tp_world_size: int, + dcp_world_size: int, + port: int, +) -> None: + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + "nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=tp_world_size, + ) + groups = _dcp_groups(tp_world_size, dcp_world_size) + dcp_partition = rank // dcp_world_size + dcp_rank = rank % dcp_world_size + dcp_group = groups[dcp_partition] + dcp_global_ranks = list( + range( + dcp_partition * dcp_world_size, + (dcp_partition + 1) * dcp_world_size, + ) + ) + + max_rows = (MAX_ROWS // dcp_world_size) * dcp_world_size + test_rows = sorted( + { + dcp_world_size, + 2 * dcp_world_size, + 4 * dcp_world_size, + max_rows, + } + ) + owner = PCIeDCPTopKOwnerExchange.from_process_group( + process_group=dcp_group, + device=device, + max_rows=max_rows, + topk=TOPK, + ) + graph_owner = PCIeDCPTopKOwnerExchange.from_process_group( + process_group=dcp_group, + device=device, + max_rows=max_rows, + topk=TOPK, + ) + try: + for step, rows in enumerate(test_rows): + local_indices, local_scores = _inputs(step, rank, rows, device) + candidate_indices, candidate_scores = owner.stage_candidates( + local_indices, local_scores + ) + torch.cuda.synchronize(device) + + rank_inputs = [ + _inputs(step, source, rows, device) for source in dcp_global_ranks + ] + expected_indices, expected_scores = owner_stage_reference( + torch.stack([item[0] for item in rank_inputs]), + torch.stack([item[1] for item in rank_inputs]), + dcp_rank, + ) + torch.testing.assert_close( + candidate_indices, expected_indices, rtol=0, atol=0 + ) + assert torch.equal( + candidate_scores.view(torch.int32), + expected_scores.view(torch.int32), + ) + graph_rows = max_rows + graph_indices, graph_scores = _inputs(10, rank, graph_rows, device) + graph = torch.cuda.CUDAGraph() + dist.barrier() + with torch.cuda.graph(graph): + graph_candidate_indices, graph_candidate_scores = ( + graph_owner.stage_candidates(graph_indices, graph_scores) + ) + + for replay_step in (20, 21): + next_indices, next_scores = _inputs(replay_step, rank, graph_rows, device) + graph_indices.copy_(next_indices) + graph_scores.copy_(next_scores) + dist.barrier() + graph.replay() + torch.cuda.synchronize(device) + + expected_rank_inputs = [ + _inputs(replay_step, source, graph_rows, device) + for source in dcp_global_ranks + ] + expected_indices, expected_scores = owner_stage_reference( + torch.stack([item[0] for item in expected_rank_inputs]), + torch.stack([item[1] for item in expected_rank_inputs]), + dcp_rank, + ) + torch.testing.assert_close( + graph_candidate_indices, expected_indices, rtol=0, atol=0 + ) + assert torch.equal( + graph_candidate_scores.view(torch.int32), + expected_scores.view(torch.int32), + ) + dist.barrier() + torch.cuda.synchronize(device) + finally: + graph_owner.close_coordinated() + owner.close_coordinated() + dist.destroy_process_group() + + +def test_pcie_dcp_topk_exact_owner_exchange(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + tp_world_size = int(os.getenv("SPARKINFER_PCIE_DCP_TOPK_TP", "8")) + dcp_world_size = int(os.getenv("SPARKINFER_PCIE_DCP_TOPK_DCP", "4")) + if ( + tp_world_size not in (2, 3, 4, 6, 8) + or dcp_world_size not in (2, 3, 4, 6, 8) + or tp_world_size % dcp_world_size != 0 + ): + pytest.skip("unsupported TP/DCP top-k geometry") + if torch.cuda.device_count() < tp_world_size: + pytest.skip(f"need {tp_world_size} CUDA devices") + _load_extension() + mp.spawn( + _worker, + args=(tp_world_size, dcp_world_size, _free_port()), + nprocs=tp_world_size, + join=True, + ) diff --git a/tests/test_packaging.py b/tests/test_packaging.py index a8b04a103..ce1cb089e 100644 --- a/tests/test_packaging.py +++ b/tests/test_packaging.py @@ -12,6 +12,7 @@ PCIE_PACKAGE = "sparkinfer.comm.pcie" RUNTIME_CUDA_SOURCES = { "pcie_dcp_a2a.cu", + "pcie_dcp_topk.cu", "pcie_dma.cu", "pcie_oneshot.cu", "pcie_twoshot.cu", From 19dca9a5c3682c01a75938a0da3d420dea444232 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 25 Jul 2026 06:25:03 +0000 Subject: [PATCH 2/3] fix(pcie): define top-k staging slot lifetime --- sparkinfer/comm/pcie/pcie_dcp_topk.cu | 10 ++++- sparkinfer/comm/pcie/pcie_dcp_topk.py | 6 +++ tests/comm/test_pcie_dcp_topk_gpu.py | 60 ++++++++++++++++++--------- 3 files changed, 54 insertions(+), 22 deletions(-) diff --git a/sparkinfer/comm/pcie/pcie_dcp_topk.cu b/sparkinfer/comm/pcie/pcie_dcp_topk.cu index 410a01e5f..cb0db6c73 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_topk.cu +++ b/sparkinfer/comm/pcie/pcie_dcp_topk.cu @@ -151,7 +151,7 @@ class PCIeDCPTopKOwnerExchange { RankSignals signals_{}; Signal *self_signal_; RankStaging candidates_[2]{}; - int slot_ = 0; + uint32_t next_slot_ = 0; PCIeDCPTopKOwnerExchange( Signal **signals, @@ -181,7 +181,13 @@ class PCIeDCPTopKOwnerExchange { const int blocks = int(std::max( 1, std::min(block_limit, (owner_packs + threads - 1) / threads))); - const int slot = slot_++ % 2; + // Host-side slot selection executes once during CUDA graph capture. Graph + // replay intentionally reuses that capture-stable address: the channel is + // stream-affine and the owner consumer is ordered after this kernel in the + // same graph/stream. Eager calls toggle slabs without an overflowing + // counter so one call cannot overwrite the immediately preceding result. + const int slot = static_cast(next_slot_); + next_slot_ ^= uint32_t{1}; const int64_t candidate_plane_packs = candidate_plane_elems_ / 4; #define LAUNCH(world) \ diff --git a/sparkinfer/comm/pcie/pcie_dcp_topk.py b/sparkinfer/comm/pcie/pcie_dcp_topk.py index 7dae26471..2a1456a69 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_topk.py +++ b/sparkinfer/comm/pcie/pcie_dcp_topk.py @@ -337,6 +337,12 @@ def stage_candidates( threads: int = 512, block_limit: int = 128, ) -> tuple[torch.Tensor, torch.Tensor]: + """Stage exact candidates and return channel-owned owner views. + + Consumers must be enqueued on this channel's stream before another + stage call. CUDA graph replay deliberately reuses its capture-time + staging address; stream ordering retires the prior consumer first. + """ if self._closed: raise RuntimeError("PCIeDCPTopKOwnerExchange is closed") if local_indices.device != self.device or local_scores.device != self.device: diff --git a/tests/comm/test_pcie_dcp_topk_gpu.py b/tests/comm/test_pcie_dcp_topk_gpu.py index b31720471..3af5ff30f 100644 --- a/tests/comm/test_pcie_dcp_topk_gpu.py +++ b/tests/comm/test_pcie_dcp_topk_gpu.py @@ -101,11 +101,13 @@ def _worker( topk=TOPK, ) try: + eager_output_ptrs = [] for step, rows in enumerate(test_rows): local_indices, local_scores = _inputs(step, rank, rows, device) candidate_indices, candidate_scores = owner.stage_candidates( local_indices, local_scores ) + eager_output_ptrs.append(candidate_indices.data_ptr()) torch.cuda.synchronize(device) rank_inputs = [ @@ -123,39 +125,57 @@ def _worker( candidate_scores.view(torch.int32), expected_scores.view(torch.int32), ) + assert eager_output_ptrs[0] != eager_output_ptrs[1] + assert eager_output_ptrs[0] == eager_output_ptrs[2] + assert eager_output_ptrs[1] == eager_output_ptrs[3] graph_rows = max_rows graph_indices, graph_scores = _inputs(10, rank, graph_rows, device) + graph_owner_rows = graph_rows // dcp_world_size + graph_candidate_width = dcp_world_size * TOPK + consumed_indices = torch.empty( + (graph_owner_rows, graph_candidate_width), + dtype=torch.int32, + device=device, + ) + consumed_scores = torch.empty( + (graph_owner_rows, graph_candidate_width), + dtype=torch.float32, + device=device, + ) graph = torch.cuda.CUDAGraph() dist.barrier() with torch.cuda.graph(graph): graph_candidate_indices, graph_candidate_scores = ( graph_owner.stage_candidates(graph_indices, graph_scores) ) - - for replay_step in (20, 21): + consumed_indices.copy_(graph_candidate_indices) + consumed_scores.copy_(graph_candidate_scores) + + # Queue several replays without host/device synchronization. Every + # replay writes the capture-stable slab, then consumes it on the same + # stream before the next replay can overwrite it. + final_replay_step = 23 + for replay_step in range(20, final_replay_step + 1): next_indices, next_scores = _inputs(replay_step, rank, graph_rows, device) graph_indices.copy_(next_indices) graph_scores.copy_(next_scores) - dist.barrier() graph.replay() - torch.cuda.synchronize(device) + torch.cuda.synchronize(device) - expected_rank_inputs = [ - _inputs(replay_step, source, graph_rows, device) - for source in dcp_global_ranks - ] - expected_indices, expected_scores = owner_stage_reference( - torch.stack([item[0] for item in expected_rank_inputs]), - torch.stack([item[1] for item in expected_rank_inputs]), - dcp_rank, - ) - torch.testing.assert_close( - graph_candidate_indices, expected_indices, rtol=0, atol=0 - ) - assert torch.equal( - graph_candidate_scores.view(torch.int32), - expected_scores.view(torch.int32), - ) + expected_rank_inputs = [ + _inputs(final_replay_step, source, graph_rows, device) + for source in dcp_global_ranks + ] + expected_indices, expected_scores = owner_stage_reference( + torch.stack([item[0] for item in expected_rank_inputs]), + torch.stack([item[1] for item in expected_rank_inputs]), + dcp_rank, + ) + torch.testing.assert_close(consumed_indices, expected_indices, rtol=0, atol=0) + assert torch.equal( + consumed_scores.view(torch.int32), + expected_scores.view(torch.int32), + ) dist.barrier() torch.cuda.synchronize(device) finally: From 88a04184956ca6e6a23cb54ebafff640856a6bfb Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 25 Jul 2026 06:41:53 +0000 Subject: [PATCH 3/3] fix(pcie): guard graph staging reuse across ranks --- sparkinfer/comm/pcie/pcie_dcp_topk.cu | 17 +++++++-- sparkinfer/comm/pcie/pcie_dcp_topk.py | 3 +- tests/comm/test_pcie_dcp_topk_gpu.py | 50 ++++++++++++++++++--------- 3 files changed, 50 insertions(+), 20 deletions(-) diff --git a/sparkinfer/comm/pcie/pcie_dcp_topk.cu b/sparkinfer/comm/pcie/pcie_dcp_topk.cu index cb0db6c73..0f04bf2f8 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_topk.cu +++ b/sparkinfer/comm/pcie/pcie_dcp_topk.cu @@ -100,7 +100,7 @@ __global__ void __launch_bounds__(512, 1) stage_owner_candidates_kernel( const int4 *__restrict__ local_indices, const int4 *__restrict__ local_scores, RankStaging staging, RankSignals signals, Signal *self, int rank, int rows, int topk, - int64_t candidate_plane_packs) { + int64_t candidate_plane_packs, bool wait_for_prior_consumer) { const int owner_rows = rows / world_size; const int packs_per_row = topk / 4; const int64_t owner_packs = int64_t(owner_rows) * packs_per_row; @@ -108,6 +108,14 @@ __global__ void __launch_bounds__(512, 1) stage_owner_candidates_kernel( int64_t begin, end; block_range(owner_packs, begin, end); + // A captured graph reuses one fixed staging address on every replay. Each + // rank reaches this kernel only after its previous same-stream owner + // consumer, so this group barrier prevents a faster peer from overwriting a + // slower owner's slab before that consumer retires. + if (wait_for_prior_consumer) { + block_pair_barrier(signals, self, rank); + } + #pragma unroll 1 for (int step = 0; step < world_size; ++step) { const int destination = (rank + step) % world_size; @@ -189,12 +197,17 @@ class PCIeDCPTopKOwnerExchange { const int slot = static_cast(next_slot_); next_slot_ ^= uint32_t{1}; const int64_t candidate_plane_packs = candidate_plane_elems_ / 4; + cudaStreamCaptureStatus capture_status; + CHECK_CUDA_SUCCESS(cudaStreamIsCapturing(stream, &capture_status)); + const bool wait_for_prior_consumer = + capture_status != cudaStreamCaptureStatusNone; #define LAUNCH(world) \ stage_owner_candidates_kernel<<>>( \ reinterpret_cast(local_indices), \ reinterpret_cast(local_scores), candidates_[slot], \ - signals_, self_signal_, rank_, rows, topk_, candidate_plane_packs) + signals_, self_signal_, rank_, rows, topk_, candidate_plane_packs, \ + wait_for_prior_consumer) switch (world_size_) { case 2: LAUNCH(2); diff --git a/sparkinfer/comm/pcie/pcie_dcp_topk.py b/sparkinfer/comm/pcie/pcie_dcp_topk.py index 2a1456a69..bacf5cfbe 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_topk.py +++ b/sparkinfer/comm/pcie/pcie_dcp_topk.py @@ -341,7 +341,8 @@ def stage_candidates( Consumers must be enqueued on this channel's stream before another stage call. CUDA graph replay deliberately reuses its capture-time - staging address; stream ordering retires the prior consumer first. + staging address; a group barrier at the next replay prevents peer + writers from racing a slower rank's prior same-stream consumer. """ if self._closed: raise RuntimeError("PCIeDCPTopKOwnerExchange is closed") diff --git a/tests/comm/test_pcie_dcp_topk_gpu.py b/tests/comm/test_pcie_dcp_topk_gpu.py index 3af5ff30f..fab48caf9 100644 --- a/tests/comm/test_pcie_dcp_topk_gpu.py +++ b/tests/comm/test_pcie_dcp_topk_gpu.py @@ -148,34 +148,50 @@ def _worker( graph_candidate_indices, graph_candidate_scores = ( graph_owner.stage_candidates(graph_indices, graph_scores) ) + # Force one owner per DCP group to remain in its consumer while + # faster peers are ready to replay and write its staging slab. + if dcp_rank == 0: + torch.cuda._sleep(20_000_000) consumed_indices.copy_(graph_candidate_indices) consumed_scores.copy_(graph_candidate_scores) # Queue several replays without host/device synchronization. Every # replay writes the capture-stable slab, then consumes it on the same # stream before the next replay can overwrite it. - final_replay_step = 23 - for replay_step in range(20, final_replay_step + 1): - next_indices, next_scores = _inputs(replay_step, rank, graph_rows, device) + replay_steps = tuple(range(20, 24)) + replay_inputs = [ + _inputs(replay_step, rank, graph_rows, device) + for replay_step in replay_steps + ] + replayed_indices = [torch.empty_like(consumed_indices) for _ in replay_steps] + replayed_scores = [torch.empty_like(consumed_scores) for _ in replay_steps] + torch.cuda.synchronize(device) + for replay_idx in range(len(replay_steps)): + next_indices, next_scores = replay_inputs[replay_idx] graph_indices.copy_(next_indices) graph_scores.copy_(next_scores) graph.replay() + replayed_indices[replay_idx].copy_(consumed_indices) + replayed_scores[replay_idx].copy_(consumed_scores) torch.cuda.synchronize(device) - expected_rank_inputs = [ - _inputs(final_replay_step, source, graph_rows, device) - for source in dcp_global_ranks - ] - expected_indices, expected_scores = owner_stage_reference( - torch.stack([item[0] for item in expected_rank_inputs]), - torch.stack([item[1] for item in expected_rank_inputs]), - dcp_rank, - ) - torch.testing.assert_close(consumed_indices, expected_indices, rtol=0, atol=0) - assert torch.equal( - consumed_scores.view(torch.int32), - expected_scores.view(torch.int32), - ) + for replay_idx, replay_step in enumerate(replay_steps): + expected_rank_inputs = [ + _inputs(replay_step, source, graph_rows, device) + for source in dcp_global_ranks + ] + expected_indices, expected_scores = owner_stage_reference( + torch.stack([item[0] for item in expected_rank_inputs]), + torch.stack([item[1] for item in expected_rank_inputs]), + dcp_rank, + ) + torch.testing.assert_close( + replayed_indices[replay_idx], expected_indices, rtol=0, atol=0 + ) + assert torch.equal( + replayed_scores[replay_idx].view(torch.int32), + expected_scores.view(torch.int32), + ) dist.barrier() torch.cuda.synchronize(device) finally: