From 0a3081884b7ed61d4e6eef709207d91ff5917bc5 Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Mon, 20 Jul 2026 12:26:08 -0500 Subject: [PATCH 1/6] feat(comm): add selected-record PCIe exchange Add a direct CUDA-IPC exchange for destination-selected fixed-width records, expose it through the Sparkinfer PCIe facade, and preserve runtime-JIT CUDA sources in built wheels. Co-authored-by: OpenAI Codex Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> --- sparkinfer/comm/pcie/__init__.py | 8 +- sparkinfer/comm/pcie/api.py | 8 + sparkinfer/comm/pcie/pcie_selected_records.cu | 319 ++++++++ sparkinfer/comm/pcie/pcie_selected_records.py | 745 ++++++++++++++++++ tests/comm/test_pcie_selected_records.py | 459 +++++++++++ tests/comm/test_pcie_selected_records_gpu.py | 447 +++++++++++ 6 files changed, 1985 insertions(+), 1 deletion(-) create mode 100644 sparkinfer/comm/pcie/pcie_selected_records.cu create mode 100644 sparkinfer/comm/pcie/pcie_selected_records.py create mode 100644 tests/comm/test_pcie_selected_records.py create mode 100644 tests/comm/test_pcie_selected_records_gpu.py diff --git a/sparkinfer/comm/pcie/__init__.py b/sparkinfer/comm/pcie/__init__.py index 892e8b98..55117276 100644 --- a/sparkinfer/comm/pcie/__init__.py +++ b/sparkinfer/comm/pcie/__init__.py @@ -11,6 +11,8 @@ - ``TwoShotReduceScatter``: two-shot sequence-parallel collectives with per-token FP8-e4m3 transport. - ``DcpAllToAll``: DCP attention exchange with fused LSE reduce-scatter. +- ``SelectedRecordExchange``: direct peer writes for destination-selected, + fixed-width records. 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 +35,14 @@ "TwoShotReduceScatter", "DcpAllToAll", "DcpAllToAllPool", + "SelectedRecordExchange", + "SelectedRecordExchangeInitializationError", "autotune_dma_crossovers", "parse_oneshot_max_size", "lse_reduce_scatter_reference", "is_supported", ), - dtypes=("bf16", "fp32", "fp8_e4m3"), + dtypes=("bf16", "fp32", "fp8_e4m3", "uint8"), requires=("multi_gpu",), provenance=Provenance( repo="https://github.com/lukealonso/sparkinfer", @@ -57,6 +61,8 @@ DmaAllReduce, OneshotAllReduce, OneshotAllReducePool, + SelectedRecordExchange, + SelectedRecordExchangeInitializationError, TwoShotReduceScatter, autotune_dma_crossovers, is_supported, diff --git a/sparkinfer/comm/pcie/api.py b/sparkinfer/comm/pcie/api.py index 3189a282..4f436445 100644 --- a/sparkinfer/comm/pcie/api.py +++ b/sparkinfer/comm/pcie/api.py @@ -27,6 +27,12 @@ from .pcie_oneshot import ( parse_pcie_oneshot_max_size as parse_oneshot_max_size, ) +from .pcie_selected_records import ( + PCIeSelectedRecordExchange as SelectedRecordExchange, +) +from .pcie_selected_records import ( + PCIeSelectedRecordExchangeInitializationError as SelectedRecordExchangeInitializationError, +) from .pcie_twoshot import ( PCIeTwoShotSP as TwoShotReduceScatter, ) @@ -50,6 +56,8 @@ def is_supported(device=None) -> bool: "TwoShotReduceScatter", "DcpAllToAll", "DcpAllToAllPool", + "SelectedRecordExchange", + "SelectedRecordExchangeInitializationError", "autotune_dma_crossovers", "parse_oneshot_max_size", "lse_reduce_scatter_reference", diff --git a/sparkinfer/comm/pcie/pcie_selected_records.cu b/sparkinfer/comm/pcie/pcie_selected_records.cu new file mode 100644 index 00000000..770a2f1f --- /dev/null +++ b/sparkinfer/comm/pcie/pcie_selected_records.cu @@ -0,0 +1,319 @@ +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace { + +constexpr int kThreads = 256; +constexpr int kMaxBlocks = 65535; +constexpr int kMaxWorldSize = 32; + +__global__ void barrier_all_peers_kernel( + const int64_t* __restrict__ publish_flag_ptrs, + const int64_t* __restrict__ wait_flag_ptrs, + int32_t* __restrict__ send_counters, + int32_t* __restrict__ wait_counters, + int world_size, + uint64_t timeout_cycles, + int phase) { + const int peer = static_cast(threadIdx.x); + if (peer >= world_size) { + return; + } + + const int64_t counter_offset = + static_cast(phase) * world_size + peer; + const int32_t publish_value = send_counters[counter_offset] + 1; + send_counters[counter_offset] = publish_value; + __threadfence_system(); + auto* publish_flag = reinterpret_cast(static_cast( + publish_flag_ptrs[counter_offset])); + asm volatile( + "st.relaxed.sys.global.u32 [%1], %0;" + : + : "r"(publish_value), "l"(publish_flag)); + + const int32_t expected = wait_counters[counter_offset] + 1; + wait_counters[counter_offset] = expected; + const auto* wait_flag = reinterpret_cast( + static_cast(wait_flag_ptrs[counter_offset])); + int32_t observed; + const uint64_t started = clock64(); + uint32_t spins = 0; + do { + asm volatile( + "ld.acquire.sys.global.u32 %0, [%1];" + : "=r"(observed) + : "l"(wait_flag)); + if (((++spins & 0x3ffU) == 0U) && + (clock64() - started > timeout_cycles)) { + printf( + "B12X selected-record barrier timed out: phase=%d peer=%d " + "expected=%d observed=%d\n", + phase, + peer, + expected, + observed); + asm volatile("trap;"); + return; + } + } while (static_cast(observed - expected) < 0); +} + +__host__ __device__ __forceinline__ int64_t record_byte_offset( + int64_t record_index, + int64_t record_bytes) { + return record_index * record_bytes; +} + +template +__global__ void scatter_records_kernel( + const uint8_t* __restrict__ records, + const index_t* __restrict__ local_indices_by_destination, + const int64_t* __restrict__ peer_payload_ptrs, + int64_t selected_records, + int64_t units_per_record, + int64_t record_bytes, + int world_size) { + const int destination = static_cast(blockIdx.y); + if (destination >= world_size) { + return; + } + + const int64_t payload_units = selected_records * units_per_record; + const int64_t grid_stride = + static_cast(blockDim.x) * gridDim.x; + for (int64_t payload_unit = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + payload_unit < payload_units; + payload_unit += grid_stride) { + const int64_t selected = payload_unit / units_per_record; + const int64_t unit_in_record = payload_unit - selected * units_per_record; + const int64_t byte_in_record = unit_in_record * sizeof(copy_t); + const int64_t map_offset = + static_cast(destination) * selected_records + selected; + const int64_t local_record = static_cast( + local_indices_by_destination[map_offset]); + if (local_record < 0) { + continue; + } + + // Keep both pool- and output-scaled products in Int64. A valid record + // index can cross the 2 GiB byte boundary even when the index is Int32. + const int64_t source_offset = + record_byte_offset(local_record, record_bytes) + byte_in_record; + const int64_t destination_offset = + record_byte_offset(selected, record_bytes) + byte_in_record; + const auto* source = reinterpret_cast( + records + source_offset); + auto* destination_ptr = reinterpret_cast( + static_cast(peer_payload_ptrs[destination]) + + destination_offset); + *destination_ptr = *source; + } +} + +void validate_tensor( + const torch::Tensor& tensor, + const char* name, + torch::ScalarType scalar_type) { + TORCH_CHECK(tensor.is_cuda(), name, " must be CUDA"); + TORCH_CHECK(tensor.scalar_type() == scalar_type, name, " has wrong dtype"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +int64_t record_byte_offset_for_test( + int64_t record_index, + int64_t record_bytes) { + TORCH_CHECK(record_index >= 0, "record_index must be non-negative"); + TORCH_CHECK(record_bytes > 0, "record_bytes must be positive"); + TORCH_CHECK( + record_index <= INT64_MAX / record_bytes, + "record byte offset exceeds int64 capacity"); + return record_byte_offset(record_index, record_bytes); +} + +void exchange( + torch::Tensor records, + torch::Tensor local_indices_by_destination, + torch::Tensor peer_payload_ptrs, + int64_t local_payload_ptr, + torch::Tensor barrier_publish_ptrs, + torch::Tensor barrier_wait_ptrs, + torch::Tensor send_counters, + torch::Tensor wait_counters, + torch::Tensor output, + int64_t record_bytes, + int64_t timeout_cycles) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(records)); + validate_tensor(records, "records", torch::kUInt8); + TORCH_CHECK( + local_indices_by_destination.is_cuda(), "local indices must be CUDA"); + TORCH_CHECK( + local_indices_by_destination.scalar_type() == torch::kInt32 || + local_indices_by_destination.scalar_type() == torch::kInt64, + "local indices must be int32 or int64"); + TORCH_CHECK( + local_indices_by_destination.is_contiguous(), + "local indices must be contiguous"); + validate_tensor(peer_payload_ptrs, "peer payload pointers", torch::kInt64); + validate_tensor( + barrier_publish_ptrs, "barrier publish pointers", torch::kInt64); + validate_tensor(barrier_wait_ptrs, "barrier wait pointers", torch::kInt64); + validate_tensor(send_counters, "send counters", torch::kInt32); + validate_tensor(wait_counters, "wait counters", torch::kInt32); + validate_tensor(output, "output", torch::kUInt8); + + TORCH_CHECK(record_bytes > 0, "record_bytes must be positive"); + TORCH_CHECK(timeout_cycles > 0, "timeout_cycles must be positive"); + TORCH_CHECK(local_payload_ptr != 0, "local payload pointer must be nonzero"); + TORCH_CHECK( + local_indices_by_destination.dim() >= 2, + "local indices must have rank at least 2"); + const int64_t world_size = local_indices_by_destination.size(0); + TORCH_CHECK( + world_size >= 1 && world_size <= kMaxWorldSize, + "world size must be in [1, 32]"); + TORCH_CHECK( + peer_payload_ptrs.numel() == world_size, + "peer payload pointer count must match world size"); + TORCH_CHECK( + barrier_publish_ptrs.numel() == 2 * world_size, + "publish pointer count must cover two barrier phases"); + TORCH_CHECK( + barrier_wait_ptrs.numel() == 2 * world_size, + "wait pointer count must cover two barrier phases"); + TORCH_CHECK( + send_counters.numel() == 2 * world_size, + "send counter count must cover two barrier phases"); + TORCH_CHECK( + wait_counters.numel() == 2 * world_size, + "wait counter count must cover two barrier phases"); + TORCH_CHECK( + records.dim() >= 2 && records.size(-1) == record_bytes, + "records must end in record_bytes"); + TORCH_CHECK( + output.dim() >= 2 && output.size(-1) == record_bytes, + "output must end in record_bytes"); + + const int64_t selected_records = + local_indices_by_destination.numel() / world_size; + TORCH_CHECK( + selected_records <= INT64_MAX / record_bytes, + "selected payload exceeds int64 capacity"); + const int64_t payload_bytes = selected_records * record_bytes; + TORCH_CHECK( + output.numel() == payload_bytes, + "output size must exactly match the selected payload"); + + const auto stream = c10::cuda::getCurrentCUDAStream().stream(); + if (payload_bytes > 0) { + const bool vectorized = + record_bytes % static_cast(sizeof(uint4)) == 0 && + reinterpret_cast(records.data_ptr()) % + alignof(uint4) == + 0; + const int64_t unit_bytes = vectorized ? sizeof(uint4) : sizeof(uint8_t); + const int64_t units_per_record = record_bytes / unit_bytes; + const int64_t payload_units = selected_records * units_per_record; + const int64_t requested_blocks = + (payload_units + kThreads - 1) / kThreads; + const int blocks = static_cast( + std::min(requested_blocks, kMaxBlocks)); + const dim3 grid( + static_cast(blocks), + static_cast(world_size)); + if (local_indices_by_destination.scalar_type() == torch::kInt32) { + if (vectorized) { + scatter_records_kernel<<>>( + records.data_ptr(), + local_indices_by_destination.data_ptr(), + peer_payload_ptrs.data_ptr(), + selected_records, + units_per_record, + record_bytes, + static_cast(world_size)); + } else { + scatter_records_kernel<<>>( + records.data_ptr(), + local_indices_by_destination.data_ptr(), + peer_payload_ptrs.data_ptr(), + selected_records, + units_per_record, + record_bytes, + static_cast(world_size)); + } + } else { + if (vectorized) { + scatter_records_kernel<<>>( + records.data_ptr(), + local_indices_by_destination.data_ptr(), + peer_payload_ptrs.data_ptr(), + selected_records, + units_per_record, + record_bytes, + static_cast(world_size)); + } else { + scatter_records_kernel<<>>( + records.data_ptr(), + local_indices_by_destination.data_ptr(), + peer_payload_ptrs.data_ptr(), + selected_records, + units_per_record, + record_bytes, + static_cast(world_size)); + } + } + AT_CUDA_CHECK(cudaGetLastError()); + } + + barrier_all_peers_kernel<<<1, 32, 0, stream>>>( + barrier_publish_ptrs.data_ptr(), + barrier_wait_ptrs.data_ptr(), + send_counters.data_ptr(), + wait_counters.data_ptr(), + static_cast(world_size), + static_cast(timeout_cycles), + 0); + AT_CUDA_CHECK(cudaGetLastError()); + + if (payload_bytes > 0) { + AT_CUDA_CHECK(cudaMemcpyAsync( + output.data_ptr(), + reinterpret_cast( + static_cast(local_payload_ptr)), + static_cast(payload_bytes), + cudaMemcpyDeviceToDevice, + stream)); + } + + barrier_all_peers_kernel<<<1, 32, 0, stream>>>( + barrier_publish_ptrs.data_ptr(), + barrier_wait_ptrs.data_ptr(), + send_counters.data_ptr(), + wait_counters.data_ptr(), + static_cast(world_size), + static_cast(timeout_cycles), + 1); + AT_CUDA_CHECK(cudaGetLastError()); +} + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def( + "record_byte_offset_for_test", + &record_byte_offset_for_test, + "Compute the Int64 byte offset used by selected-record scatter"); + module.def( + "exchange", + &exchange, + "Direct-scatter selected records with two device-side barriers"); +} diff --git a/sparkinfer/comm/pcie/pcie_selected_records.py b/sparkinfer/comm/pcie/pcie_selected_records.py new file mode 100644 index 00000000..7ff032d3 --- /dev/null +++ b/sparkinfer/comm/pcie/pcie_selected_records.py @@ -0,0 +1,745 @@ +"""Direct PCIe exchange for destination-selected fixed-width records.""" + +from __future__ import annotations + +import ctypes +from contextlib import suppress +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path +from typing import 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, cudaIpcMemHandle_t +from .pcie_dma import FLAG_STRIDE +from .pcie_oneshot import ( + IPC_SLAB_ALIGNMENT, + _current_stream_key, + _is_current_stream_capturing, + _normalize_device, + _OwnedSharedBuffer, + _align_up, +) + + +MAX_WORLD_SIZE = 32 +DEFAULT_BARRIER_TIMEOUT_CYCLES = 8_000_000_000 +_MAX_INT64 = (1 << 63) - 1 +_IPC_HANDLE_BYTES = ctypes.sizeof(cudaIpcMemHandle_t) +_CONFIG_VERSION = 1 +_CONFIG_FIELDS = ( + "version", + "world_size", + "max_records", + "record_bytes", + "flags_bytes", + "payload_offset", + "payload_bytes", + "slab_bytes", + "flag_stride", + "slab_alignment", +) + + +class PCIeSelectedRecordExchangeInitializationError(RuntimeError): + """The selected-record channel is unavailable and callers should fall back.""" + + +@dataclass(frozen=True) +class _SelectedRecordLayout: + flags_bytes: int + payload_offset: int + payload_bytes: int + slab_bytes: int + + +def _selected_record_layout( + *, world_size: int, max_records: int, record_bytes: int +) -> _SelectedRecordLayout: + world_size = int(world_size) + max_records = int(max_records) + record_bytes = int(record_bytes) + if not 1 <= world_size <= MAX_WORLD_SIZE: + raise ValueError( + f"world_size must be in [1, {MAX_WORLD_SIZE}], got {world_size}" + ) + if max_records <= 0: + raise ValueError("max_records must be positive") + if record_bytes <= 0: + raise ValueError("record_bytes must be positive") + if max_records > _MAX_INT64 // record_bytes: + raise ValueError("max_records * record_bytes exceeds int64 capacity") + + flags_bytes = _align_up( + 2 * world_size * FLAG_STRIDE, + IPC_SLAB_ALIGNMENT, + ) + payload_offset = flags_bytes + payload_bytes = max_records * record_bytes + slab_bytes = payload_offset + payload_bytes + if slab_bytes > _MAX_INT64: + raise ValueError("selected-record IPC slab exceeds int64 capacity") + return _SelectedRecordLayout( + flags_bytes=flags_bytes, + payload_offset=payload_offset, + payload_bytes=payload_bytes, + slab_bytes=slab_bytes, + ) + + +@lru_cache(maxsize=1) +def _load_extension(): + source = Path(__file__).with_name("pcie_selected_records.cu") + return load( + name="sparkinfer_pcie_selected_records_ext", + sources=[str(source)], + extra_cuda_cflags=["-O3"], + extra_ldflags=["-lcuda"], + verbose=False, + ) + + +def _all_ranks_succeeded( + status: torch.Tensor, + local_success: bool, + process_group: ProcessGroup, +) -> bool: + status.fill_(1 if local_success else 0) + dist.all_reduce(status, op=dist.ReduceOp.MIN, group=process_group) + return bool(status.item()) + + +def _configuration_values( + *, + world_size: int, + max_records: int, + record_bytes: int, + layout: _SelectedRecordLayout, +) -> tuple[int, ...]: + return ( + _CONFIG_VERSION, + int(world_size), + int(max_records), + int(record_bytes), + layout.flags_bytes, + layout.payload_offset, + layout.payload_bytes, + layout.slab_bytes, + FLAG_STRIDE, + IPC_SLAB_ALIGNMENT, + ) + + +def _validate_rank_configuration( + *, + process_group: ProcessGroup, + device: torch.device, + status: torch.Tensor, + world_size: int, + max_records: int, + record_bytes: int, + layout: _SelectedRecordLayout, +) -> None: + local_config: Optional[torch.Tensor] = None + gathered_configs: Optional[torch.Tensor] = None + local_error: Optional[Exception] = None + try: + local_config = torch.tensor( + _configuration_values( + world_size=world_size, + max_records=max_records, + record_bytes=record_bytes, + layout=layout, + ), + dtype=torch.int64, + device=device, + ) + gathered_configs = torch.empty( + world_size * len(_CONFIG_FIELDS), + dtype=torch.int64, + device=device, + ) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record configuration handshake allocation failed on at least " + "one rank" + ) from local_error + + assert local_config is not None + assert gathered_configs is not None + try: + dist.all_gather_into_tensor( + gathered_configs, + local_config, + group=process_group, + ) + except Exception as exc: + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record configuration handshake collective failed" + ) from exc + + rows = gathered_configs.view(world_size, len(_CONFIG_FIELDS)) + local_matches = False + local_error = None + try: + local_matches = bool(torch.all(rows == local_config.unsqueeze(0)).item()) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record configuration handshake comparison failed on at least " + "one rank" + ) from local_error + if not _all_ranks_succeeded(status, local_matches, process_group): + try: + rank_configs = rows.cpu().tolist() + except Exception: + rank_configs = "unavailable" + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record configuration mismatch across ranks " + f"(fields={_CONFIG_FIELDS}, values={rank_configs})" + ) + + +def _gather_ipc_handles( + *, + process_group: ProcessGroup, + device: torch.device, + world_size: int, + local_ptr: int, + ipc: CudaRTLibrary, + status: torch.Tensor, +) -> tuple[bytes, ...]: + local_handle_tensor: Optional[torch.Tensor] = None + gathered_handle_tensor: Optional[torch.Tensor] = None + local_error: Optional[Exception] = None + try: + local_handle_tensor = torch.empty( + _IPC_HANDLE_BYTES, + dtype=torch.uint8, + device=device, + ) + gathered_handle_tensor = torch.empty( + world_size * _IPC_HANDLE_BYTES, + dtype=torch.uint8, + device=device, + ) + local_handle = ipc.cudaIpcGetMemHandleBytes(local_ptr) + if len(local_handle) != _IPC_HANDLE_BYTES: + raise ValueError( + f"CUDA IPC handle has {len(local_handle)} bytes, " + f"expected {_IPC_HANDLE_BYTES}" + ) + local_handle_tensor.copy_( + torch.frombuffer(bytearray(local_handle), dtype=torch.uint8) + ) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record IPC handle preparation failed on at least one rank" + ) from local_error + + assert local_handle_tensor is not None + assert gathered_handle_tensor is not None + try: + dist.all_gather_into_tensor( + gathered_handle_tensor, + local_handle_tensor, + group=process_group, + ) + except Exception as exc: + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record IPC handle tensor exchange failed" + ) from exc + + handles: Optional[tuple[bytes, ...]] = None + local_error = None + try: + gathered_cpu = gathered_handle_tensor.view(world_size, _IPC_HANDLE_BYTES).cpu() + handles = tuple(bytes(row.tolist()) for row in gathered_cpu) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record IPC handle materialization failed on at least one rank" + ) from local_error + assert handles is not None + return handles + + +def _free_shared_buffer(ipc: CudaRTLibrary, shared: _OwnedSharedBuffer) -> None: + for ptr in shared.remote_ptrs: + with suppress(Exception): + ipc.cudaIpcCloseMemHandle(ptr) + with suppress(Exception): + ipc.cudaFree(shared.local_ptr) + + +def _allocate_shared_buffer_rank_consistent( + *, + process_group: ProcessGroup, + rank: int, + world_size: int, + device: torch.device, + size_in_bytes: int, + ipc: CudaRTLibrary, + status: torch.Tensor, +) -> _OwnedSharedBuffer: + local_ptr: Optional[int] = None + local_error: Optional[Exception] = None + try: + local_ptr = ipc.cudaMalloc(size_in_bytes) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + if local_ptr is not None: + with suppress(Exception): + ipc.cudaFree(local_ptr) + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record IPC slab allocation failed on at least one rank" + ) from local_error + + assert local_ptr is not None + try: + ipc.cudaMemset(local_ptr, 0, size_in_bytes) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + with suppress(Exception): + ipc.cudaFree(local_ptr) + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record IPC slab initialization failed on at least one rank" + ) from local_error + + try: + handles = _gather_ipc_handles( + process_group=process_group, + device=device, + world_size=world_size, + local_ptr=local_ptr, + ipc=ipc, + status=status, + ) + except Exception: + with suppress(Exception): + ipc.cudaFree(local_ptr) + raise + + peer_ptrs: list[int] = [] + remote_ptrs: list[int] = [] + local_error = None + try: + for peer_rank, handle in enumerate(handles): + if peer_rank == rank: + peer_ptrs.append(local_ptr) + else: + remote_ptr = ipc.cudaIpcOpenMemHandleBytes(handle) + peer_ptrs.append(remote_ptr) + remote_ptrs.append(remote_ptr) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + shared = _OwnedSharedBuffer( + local_ptr=local_ptr, + peer_ptrs=tuple(peer_ptrs), + remote_ptrs=tuple(remote_ptrs), + ) + _free_shared_buffer(ipc, shared) + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record IPC handle open failed on at least one rank" + ) from local_error + opened_every_rank = len(peer_ptrs) == world_size + if not _all_ranks_succeeded(status, opened_every_rank, process_group): + shared = _OwnedSharedBuffer( + local_ptr=local_ptr, + peer_ptrs=tuple(peer_ptrs), + remote_ptrs=tuple(remote_ptrs), + ) + _free_shared_buffer(ipc, shared) + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record IPC setup did not open every rank on all processes" + ) + shared = _OwnedSharedBuffer( + local_ptr=local_ptr, + peer_ptrs=tuple(peer_ptrs), + remote_ptrs=tuple(remote_ptrs), + ) + + local_error = None + try: + stream = torch.cuda.current_stream(device) + stream_ptr = int(stream.cuda_stream) + for remote_ptr in shared.remote_ptrs: + ipc.cudaMemcpyAsync(local_ptr, remote_ptr, 1, stream_ptr) + ipc.cudaMemcpyAsync(remote_ptr, local_ptr, 1, stream_ptr) + stream.synchronize() + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + _free_shared_buffer(ipc, shared) + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record CUDA IPC peer open/copy probe failed on at least one rank" + ) from local_error + return shared + + +class PCIeSelectedRecordExchange: + """One ordered direct-scatter channel for selected byte records. + + ``local_indices_by_destination`` is destination-major. Its trailing + dimensions flatten to the output record order. Each entry is either a + non-negative index into ``records`` or negative when this source does not + own that selected record. Across all ranks, exactly one source must own + each destination/output position. Every rank must invoke the channel in + the same order with the same selected-record count. + """ + + def __init__( + self, + *, + rank: int, + world_size: int, + device: torch.device | int | str, + peer_slab_ptrs: Sequence[int], + payload_offset: int, + max_records: int, + record_bytes: int, + process_group: Optional[ProcessGroup] = None, + ipc: Optional[CudaRTLibrary] = None, + owned_buffer: Optional[_OwnedSharedBuffer] = None, + ext_module=None, + barrier_timeout_cycles: int = DEFAULT_BARRIER_TIMEOUT_CYCLES, + stream_affine: bool = True, + ) -> None: + layout = _selected_record_layout( + world_size=world_size, + max_records=max_records, + record_bytes=record_bytes, + ) + if not 0 <= int(rank) < int(world_size): + raise ValueError(f"invalid rank {rank} for world size {world_size}") + if len(peer_slab_ptrs) != int(world_size): + raise ValueError("peer_slab_ptrs must match world_size") + if int(payload_offset) < layout.flags_bytes: + raise ValueError("payload_offset overlaps selected-record barrier flags") + if int(barrier_timeout_cycles) <= 0: + raise ValueError("barrier_timeout_cycles must be positive") + + self.rank = int(rank) + self.world_size = int(world_size) + self.device = _normalize_device(device) + self.process_group = process_group + self.max_records = int(max_records) + self.record_bytes = int(record_bytes) + self.payload_offset = int(payload_offset) + self.barrier_timeout_cycles = int(barrier_timeout_cycles) + self._ipc = ipc + self._owned_buffer = owned_buffer + self._ext = ext_module or _load_extension() + self._stream_affine = bool(stream_affine) + self._owner_stream_key: Optional[int] = None + self._closed = False + + slab_ptrs = tuple(int(ptr) for ptr in peer_slab_ptrs) + self._local_payload_ptr = slab_ptrs[self.rank] + self.payload_offset + self._peer_payload_ptrs = torch.tensor( + [ptr + self.payload_offset for ptr in slab_ptrs], + dtype=torch.int64, + device=self.device, + ) + self._barrier_publish_ptrs = torch.tensor( + [ + [ + slab_ptrs[destination] + + (phase * self.world_size + self.rank) * FLAG_STRIDE + for destination in range(self.world_size) + ] + for phase in range(2) + ], + dtype=torch.int64, + device=self.device, + ) + self._barrier_wait_ptrs = torch.tensor( + [ + [ + slab_ptrs[self.rank] + + (phase * self.world_size + source) * FLAG_STRIDE + for source in range(self.world_size) + ] + for phase in range(2) + ], + dtype=torch.int64, + device=self.device, + ) + self._send_counters = torch.zeros( + (2, self.world_size), + dtype=torch.int32, + device=self.device, + ) + self._wait_counters = torch.zeros_like(self._send_counters) + + @classmethod + def from_process_group( + cls, + *, + process_group: ProcessGroup, + device: torch.device | int | str, + max_records: int, + record_bytes: int, + barrier_timeout_cycles: int = DEFAULT_BARRIER_TIMEOUT_CYCLES, + ext_module=None, + stream_affine: bool = True, + ) -> "PCIeSelectedRecordExchange": + rank = dist.get_rank(group=process_group) + world_size = dist.get_world_size(group=process_group) + device_obj = _normalize_device(device) + if device_obj.type != "cuda": + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record exchange requires a CUDA device" + ) + device_index = ( + torch.cuda.current_device() + if device_obj.index is None + else int(device_obj.index) + ) + device_obj = torch.device("cuda", device_index) + + try: + backend = dist.get_backend(group=process_group) + except TypeError: + backend = dist.get_backend(process_group) + if "nccl" not in str(backend).lower(): + raise PCIeSelectedRecordExchangeInitializationError( + f"selected-record exchange requires an NCCL process group, got {backend}" + ) + + status = torch.empty(1, dtype=torch.int32, device=device_obj) + + layout: Optional[_SelectedRecordLayout] = None + local_error: Optional[Exception] = None + try: + layout = _selected_record_layout( + world_size=world_size, + max_records=max_records, + record_bytes=record_bytes, + ) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record configuration is invalid on at least one rank" + ) from local_error + assert layout is not None + _validate_rank_configuration( + process_group=process_group, + device=device_obj, + status=status, + world_size=world_size, + max_records=max_records, + record_bytes=record_bytes, + layout=layout, + ) + + ext = None + local_error = None + try: + ext = ext_module or _load_extension() + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record CUDA extension failed to load on at least one rank" + ) from local_error + assert ext is not None + + ipc: Optional[CudaRTLibrary] = None + local_error = None + try: + ipc = CudaRTLibrary() + ipc.cudaSetDevice(device_index) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record CUDA IPC setup failed on at least one rank" + ) from local_error + assert ipc is not None + + shared = _allocate_shared_buffer_rank_consistent( + process_group=process_group, + rank=rank, + world_size=world_size, + device=device_obj, + size_in_bytes=layout.slab_bytes, + ipc=ipc, + status=status, + ) + exchange: Optional[PCIeSelectedRecordExchange] = None + local_error = None + try: + exchange = cls( + rank=rank, + world_size=world_size, + device=device_obj, + peer_slab_ptrs=shared.peer_ptrs, + payload_offset=layout.payload_offset, + max_records=max_records, + record_bytes=record_bytes, + process_group=process_group, + ipc=ipc, + owned_buffer=None, + ext_module=ext, + barrier_timeout_cycles=barrier_timeout_cycles, + stream_affine=stream_affine, + ) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + if exchange is not None: + exchange._closed = True + _free_shared_buffer(ipc, shared) + raise PCIeSelectedRecordExchangeInitializationError( + "selected-record runtime initialization failed on at least one rank" + ) from local_error + assert exchange is not None + exchange._owned_buffer = shared + return exchange + + def _bind_stream_key(self, stream_key: Optional[int]) -> None: + if not self._stream_affine or stream_key is None: + return + if self._owner_stream_key is None: + self._owner_stream_key = int(stream_key) + return + if self._owner_stream_key != int(stream_key): + raise RuntimeError( + "PCIe selected-record channels are stream-affine; create a " + "separate channel for each CUDA stream" + ) + + def _check_stream(self) -> None: + if self.device.type != "cuda": + return + stream_key = _current_stream_key(self.device) + if ( + _is_current_stream_capturing(self.device) + and self._stream_affine + and self._owner_stream_key is None + ): + raise RuntimeError( + "PCIe selected-record channels must be bound to their CUDA stream " + "before first-use graph capture" + ) + self._bind_stream_key(stream_key) + + def _validate( + self, + records: torch.Tensor, + local_indices_by_destination: torch.Tensor, + out: torch.Tensor, + ) -> int: + if self._closed: + raise RuntimeError("PCIeSelectedRecordExchange is closed") + if records.device != self.device: + raise ValueError("records must be on the exchange device") + if local_indices_by_destination.device != self.device: + raise ValueError("local indices must be on the exchange device") + if out.device != self.device: + raise ValueError("output must be on the exchange device") + if records.dtype != torch.uint8 or not records.is_contiguous(): + raise ValueError("records must be contiguous uint8") + if records.ndim < 2 or records.shape[-1] != self.record_bytes: + raise ValueError( + f"records must end in the configured {self.record_bytes}-byte width" + ) + if local_indices_by_destination.dtype not in (torch.int32, torch.int64): + raise ValueError("local indices must be int32 or int64") + if not local_indices_by_destination.is_contiguous(): + raise ValueError("local indices must be contiguous") + if ( + local_indices_by_destination.ndim < 2 + or local_indices_by_destination.shape[0] != self.world_size + ): + raise ValueError( + "local indices must have shape [world_size, ...selected records]" + ) + active_records = local_indices_by_destination.numel() // self.world_size + if active_records > self.max_records: + raise ValueError( + f"selected record count {active_records} exceeds capacity " + f"{self.max_records}" + ) + if out.dtype != torch.uint8 or not out.is_contiguous(): + raise ValueError("output must be contiguous uint8") + if out.ndim < 2 or out.shape[-1] != self.record_bytes: + raise ValueError( + f"output must end in the configured {self.record_bytes}-byte width" + ) + if out.numel() != active_records * self.record_bytes: + raise ValueError( + "output must contain exactly one record per selected position" + ) + return active_records + + def exchange( + self, + records: torch.Tensor, + local_indices_by_destination: torch.Tensor, + out: torch.Tensor, + ) -> torch.Tensor: + """Scatter owned records to each destination and fill ``out`` exactly.""" + self._check_stream() + self._validate(records, local_indices_by_destination, out) + self._ext.exchange( + records, + local_indices_by_destination, + self._peer_payload_ptrs, + self._local_payload_ptr, + self._barrier_publish_ptrs, + self._barrier_wait_ptrs, + self._send_counters, + self._wait_counters, + out, + self.record_bytes, + self.barrier_timeout_cycles, + ) + return out + + def _release_owned_buffer(self, *, synchronize: bool) -> None: + if self._owned_buffer is None or self._ipc is None: + return + if self.device.type == "cuda": + device_index = ( + torch.cuda.current_device() + if self.device.index is None + else int(self.device.index) + ) + self._ipc.cudaSetDevice(device_index) + if synchronize: + with suppress(Exception): + torch.cuda.synchronize(self.device) + _free_shared_buffer(self._ipc, self._owned_buffer) + self._owned_buffer = None + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._release_owned_buffer(synchronize=True) + + def __del__(self) -> None: + with suppress(Exception): + self.close() + + +__all__ = [ + "PCIeSelectedRecordExchange", + "PCIeSelectedRecordExchangeInitializationError", +] diff --git a/tests/comm/test_pcie_selected_records.py b/tests/comm/test_pcie_selected_records.py new file mode 100644 index 00000000..cf6e2607 --- /dev/null +++ b/tests/comm/test_pcie_selected_records.py @@ -0,0 +1,459 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +from sparkinfer.comm.pcie import ( + SelectedRecordExchange as PCIeSelectedRecordExchange, + SelectedRecordExchangeInitializationError as PCIeSelectedRecordExchangeInitializationError, +) +from sparkinfer.comm.pcie.pcie_dma import FLAG_STRIDE +from sparkinfer.comm.pcie.pcie_oneshot import _OwnedSharedBuffer +from sparkinfer.comm.pcie.pcie_selected_records import ( + _CONFIG_FIELDS, + _IPC_HANDLE_BYTES, + _gather_ipc_handles, + _selected_record_layout, + _validate_rank_configuration, +) + + +class _FakeExt: + def __init__(self, rank: int) -> None: + self.rank = rank + self.calls = [] + + def exchange( + self, + records, + local_indices_by_destination, + peer_payload_ptrs, + local_payload_ptr, + barrier_publish_ptrs, + barrier_wait_ptrs, + send_counters, + wait_counters, + out, + record_bytes, + timeout_cycles, + ): + self.calls.append( + { + "peer_payload_ptrs": peer_payload_ptrs.clone(), + "local_payload_ptr": local_payload_ptr, + "publish_ptrs": barrier_publish_ptrs.clone(), + "wait_ptrs": barrier_wait_ptrs.clone(), + "send_counters_ptr": send_counters.data_ptr(), + "wait_counters_ptr": wait_counters.data_ptr(), + "record_bytes": record_bytes, + "timeout_cycles": timeout_cycles, + } + ) + flat_records = records.reshape(-1, record_bytes) + flat_indices = local_indices_by_destination[self.rank].reshape(-1) + flat_out = out.reshape(-1, record_bytes) + for selected, local_index in enumerate(flat_indices.tolist()): + if local_index >= 0: + flat_out[selected].copy_(flat_records[local_index]) + send_counters.add_(1) + wait_counters.add_(1) + + +def _make_runtime( + *, + rank: int = 1, + world_size: int = 3, + max_records: int = 4, + record_bytes: int = 37, +) -> tuple[PCIeSelectedRecordExchange, _FakeExt, tuple[int, ...]]: + ext = _FakeExt(rank) + layout = _selected_record_layout( + world_size=world_size, + max_records=max_records, + record_bytes=record_bytes, + ) + slab_ptrs = tuple(10_000 * (peer + 1) for peer in range(world_size)) + runtime = PCIeSelectedRecordExchange( + rank=rank, + world_size=world_size, + device=torch.device("cpu"), + peer_slab_ptrs=slab_ptrs, + payload_offset=layout.payload_offset, + max_records=max_records, + record_bytes=record_bytes, + ext_module=ext, + ) + return runtime, ext, slab_ptrs + + +def test_layout_uses_runtime_world_size_and_pool_scaled_int64_capacity(): + layout = _selected_record_layout( + world_size=7, + max_records=32_769, + record_bytes=65_536, + ) + + assert layout.flags_bytes >= 2 * 7 * FLAG_STRIDE + assert layout.flags_bytes % 256 == 0 + assert layout.payload_offset == layout.flags_bytes + assert layout.payload_bytes == 32_769 * 65_536 + assert layout.payload_bytes > 2**31 + assert layout.slab_bytes == layout.payload_offset + layout.payload_bytes + + +def test_layout_rejects_invalid_and_int64_overflowing_capacities(): + with pytest.raises(ValueError, match="world_size"): + _selected_record_layout(world_size=33, max_records=1, record_bytes=1) + with pytest.raises(ValueError, match="max_records"): + _selected_record_layout(world_size=2, max_records=0, record_bytes=1) + with pytest.raises(ValueError, match="record_bytes"): + _selected_record_layout(world_size=2, max_records=1, record_bytes=0) + with pytest.raises(ValueError, match="int64"): + _selected_record_layout( + world_size=2, + max_records=2**62, + record_bytes=4, + ) + + +def test_exchange_dispatches_exact_odd_width_records_with_persistent_state(): + runtime, ext, slab_ptrs = _make_runtime() + records = torch.arange(5 * 37, dtype=torch.uint8).reshape(5, 37) + local_indices = torch.full((3, 1, 2), -1, dtype=torch.int32) + local_indices[1, 0] = torch.tensor([4, 1], dtype=torch.int32) + out = torch.zeros(1, 2, 37, dtype=torch.uint8) + + returned = runtime.exchange(records, local_indices, out) + runtime.exchange(records.flip(0).contiguous(), local_indices, out) + + assert returned is out + assert torch.equal(out[0, 0], records[0]) + assert torch.equal(out[0, 1], records[3]) + assert len(ext.calls) == 2 + assert ext.calls[0]["record_bytes"] == 37 + assert ext.calls[0]["send_counters_ptr"] == ext.calls[1]["send_counters_ptr"] + assert ext.calls[0]["wait_counters_ptr"] == ext.calls[1]["wait_counters_ptr"] + + layout = _selected_record_layout( + world_size=3, + max_records=4, + record_bytes=37, + ) + assert ext.calls[0]["peer_payload_ptrs"].tolist() == [ + ptr + layout.payload_offset for ptr in slab_ptrs + ] + assert ext.calls[0]["local_payload_ptr"] == slab_ptrs[1] + layout.payload_offset + assert ext.calls[0]["publish_ptrs"].tolist() == [ + [slab_ptrs[destination] + FLAG_STRIDE for destination in range(3)], + [slab_ptrs[destination] + 4 * FLAG_STRIDE for destination in range(3)], + ] + assert ext.calls[0]["wait_ptrs"].tolist() == [ + [slab_ptrs[1] + source * FLAG_STRIDE for source in range(3)], + [slab_ptrs[1] + (3 + source) * FLAG_STRIDE for source in range(3)], + ] + assert torch.equal(runtime._send_counters, torch.full((2, 3), 2, dtype=torch.int32)) + assert torch.equal(runtime._wait_counters, torch.full((2, 3), 2, dtype=torch.int32)) + + +def test_exchange_accepts_int64_indices_and_an_empty_selection(): + runtime, ext, _ = _make_runtime(max_records=2) + records = torch.empty((0, 37), dtype=torch.uint8) + local_indices = torch.empty((3, 0), dtype=torch.int64) + out = torch.empty((0, 37), dtype=torch.uint8) + + assert runtime.exchange(records, local_indices, out) is out + assert len(ext.calls) == 1 + + +def test_exchange_rejects_invalid_shapes_dtypes_capacity_and_closed_state(): + runtime, _, _ = _make_runtime(max_records=2) + records = torch.zeros(3, 37, dtype=torch.uint8) + indices = torch.full((3, 2), -1, dtype=torch.int32) + out = torch.zeros(2, 37, dtype=torch.uint8) + + with pytest.raises(ValueError, match="configured 37-byte width"): + runtime.exchange(records[:, :-1].contiguous(), indices, out) + with pytest.raises(ValueError, match="int32 or int64"): + runtime.exchange(records, indices.to(torch.int16), out) + with pytest.raises(ValueError, match="world_size"): + runtime.exchange(records, indices[:2], out) + with pytest.raises(ValueError, match="exceeds capacity"): + runtime.exchange( + records, + torch.full((3, 3), -1, dtype=torch.int32), + torch.zeros(3, 37, dtype=torch.uint8), + ) + with pytest.raises(ValueError, match="exactly one record"): + runtime.exchange(records, indices, torch.zeros(3, 37, dtype=torch.uint8)) + + runtime.close() + with pytest.raises(RuntimeError, match="closed"): + runtime.exchange(records, indices, out) + + +def test_cpu_process_group_initialization_raises_the_fallback_exception(monkeypatch): + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records.dist.get_rank", + lambda group=None: 0, + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records.dist.get_world_size", + lambda group=None: 2, + ) + + with pytest.raises( + PCIeSelectedRecordExchangeInitializationError, + match="requires a CUDA device", + ): + PCIeSelectedRecordExchange.from_process_group( + process_group=object(), + device="cpu", + max_records=2, + record_bytes=37, + ) + + +@pytest.mark.parametrize( + ("field", "replacement"), + ( + ("world_size", 3), + ("max_records", 5), + ("record_bytes", 38), + ("payload_offset", 4096), + ), +) +def test_rank_configuration_handshake_rejects_mismatches( + monkeypatch, + field: str, + replacement: int, +): + layout = _selected_record_layout( + world_size=2, + max_records=4, + record_bytes=37, + ) + field_index = _CONFIG_FIELDS.index(field) + + def fake_all_gather_into_tensor(output, local, group=None): + rows = local.repeat(2, 1) + rows[1, field_index] = replacement + output.copy_(rows.reshape(-1)) + + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records.dist.all_gather_into_tensor", + fake_all_gather_into_tensor, + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records._all_ranks_succeeded", + lambda status, local_success, process_group: local_success, + ) + + with pytest.raises( + PCIeSelectedRecordExchangeInitializationError, + match="configuration mismatch", + ): + _validate_rank_configuration( + process_group=object(), + device=torch.device("cpu"), + status=torch.empty(1, dtype=torch.int32), + world_size=2, + max_records=4, + record_bytes=37, + layout=layout, + ) + + +def test_ipc_handle_preflight_stops_before_tensor_collective(monkeypatch): + collective_calls = [] + + class _SerializationFailureIPC: + def cudaIpcGetMemHandleBytes(self, pointer): + raise RuntimeError(f"cannot serialize pointer {pointer}") + + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records.dist.all_gather_into_tensor", + lambda *args, **kwargs: collective_calls.append((args, kwargs)), + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records._all_ranks_succeeded", + lambda status, local_success, process_group: local_success, + ) + + with pytest.raises( + PCIeSelectedRecordExchangeInitializationError, + match="handle preparation failed", + ): + _gather_ipc_handles( + process_group=object(), + device=torch.device("cpu"), + world_size=2, + local_ptr=1234, + ipc=_SerializationFailureIPC(), + status=torch.empty(1, dtype=torch.int32), + ) + assert collective_calls == [] + + +def test_ipc_handle_allocation_preflight_stops_before_tensor_collective(monkeypatch): + collective_calls = [] + status = torch.empty(1, dtype=torch.int32) + + def fail_empty(*args, **kwargs): + raise RuntimeError("intentional handle tensor allocation failure") + + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records.torch.empty", + fail_empty, + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records.dist.all_gather_into_tensor", + lambda *args, **kwargs: collective_calls.append((args, kwargs)), + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records._all_ranks_succeeded", + lambda status, local_success, process_group: local_success, + ) + + with pytest.raises( + PCIeSelectedRecordExchangeInitializationError, + match="handle preparation failed", + ): + _gather_ipc_handles( + process_group=object(), + device=torch.device("cpu"), + world_size=2, + local_ptr=1234, + ipc=object(), + status=status, + ) + assert collective_calls == [] + + +def test_ipc_handles_use_one_fixed_uint8_tensor_collective(monkeypatch): + local_handle = bytes(range(_IPC_HANDLE_BYTES)) + peer_handle = bytes(reversed(range(_IPC_HANDLE_BYTES))) + collective_shapes = [] + + class _IPC: + def cudaIpcGetMemHandleBytes(self, pointer): + assert pointer == 1234 + return local_handle + + def fake_all_gather_into_tensor(output, local, group=None): + collective_shapes.append((tuple(output.shape), tuple(local.shape), local.dtype)) + output.view(2, _IPC_HANDLE_BYTES)[0].copy_(local) + output.view(2, _IPC_HANDLE_BYTES)[1].copy_( + torch.tensor(tuple(peer_handle), dtype=torch.uint8) + ) + + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records.dist.all_gather_into_tensor", + fake_all_gather_into_tensor, + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records._all_ranks_succeeded", + lambda status, local_success, process_group: local_success, + ) + + handles = _gather_ipc_handles( + process_group=object(), + device=torch.device("cpu"), + world_size=2, + local_ptr=1234, + ipc=_IPC(), + status=torch.empty(1, dtype=torch.int32), + ) + + assert handles == (local_handle, peer_handle) + assert collective_shapes == [ + ((2 * _IPC_HANDLE_BYTES,), (_IPC_HANDLE_BYTES,), torch.uint8) + ] + + +def test_first_use_graph_capture_requires_a_prebound_matching_stream(monkeypatch): + runtime, _, _ = _make_runtime() + runtime.device = torch.device("cuda", 0) + stream_key = [77] + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records._is_current_stream_capturing", + lambda device: True, + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records._current_stream_key", + lambda device: stream_key[0], + ) + + with pytest.raises(RuntimeError, match="before first-use graph capture"): + runtime._check_stream() + + runtime._bind_stream_key(77) + runtime._check_stream() + stream_key[0] = 78 + with pytest.raises(RuntimeError, match="stream-affine"): + runtime._check_stream() + + +def test_cleanup_rebinds_the_configured_cuda_device(monkeypatch): + operations = [] + + class _CleanupIPC: + def cudaSetDevice(self, device): + operations.append(("set_device", device)) + + def cudaIpcCloseMemHandle(self, pointer): + operations.append(("close", pointer)) + + def cudaFree(self, pointer): + operations.append(("free", pointer)) + + runtime, _, _ = _make_runtime() + runtime.device = torch.device("cuda", 3) + runtime._ipc = _CleanupIPC() + runtime._owned_buffer = _OwnedSharedBuffer( + local_ptr=100, + peer_ptrs=(100, 200), + remote_ptrs=(200,), + ) + monkeypatch.setattr( + torch.cuda, + "synchronize", + lambda device: operations.append(("synchronize", device.index)), + ) + + runtime.close() + + assert operations == [ + ("set_device", 3), + ("synchronize", 3), + ("close", 200), + ("free", 100), + ] + + +def test_cuda_source_uses_int64_pool_scaled_offsets_and_no_domain_policy(): + module_path = ( + Path(__file__).parents[2] + / "sparkinfer" + / "comm" + / "pcie" + / "pcie_selected_records.cu" + ) + source = module_path.read_text(encoding="utf-8") + python_source = module_path.with_suffix(".py").read_text(encoding="utf-8") + + assert "const int64_t local_record" in source + assert "record_byte_offset(local_record, record_bytes)" in source + assert "record_byte_offset(selected, record_bytes)" in source + assert "const int64_t destination_offset" in source + assert "destination_ptr" in source + assert "scatter_records_kernel" in source + assert "scatter_records_kernel" in source + assert "barrier_all_peers_kernel" in source + assert source.count("barrier_all_peers_kernel<<<") == 2 + assert "cudaMalloc" not in source + assert "compact" not in source.lower() + assert "remap" not in source.lower() + assert "_broadcast_gather_object" not in python_source + assert python_source.count("dist.all_gather_into_tensor") == 2 diff --git a/tests/comm/test_pcie_selected_records_gpu.py b/tests/comm/test_pcie_selected_records_gpu.py new file mode 100644 index 00000000..8dc28654 --- /dev/null +++ b/tests/comm/test_pcie_selected_records_gpu.py @@ -0,0 +1,447 @@ +from __future__ import annotations + +import os +import socket +from contextlib import suppress + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from sparkinfer.comm.pcie.pcie_selected_records import ( + PCIeSelectedRecordExchange, + PCIeSelectedRecordExchangeInitializationError, + _load_extension, +) + + +pytestmark = pytest.mark.skipif( + os.getenv("SPARKINFER_RUN_PCIE_SELECTED_RECORDS_TEST") != "1", + reason=( + "set SPARKINFER_RUN_PCIE_SELECTED_RECORDS_TEST=1 to run selected-record GPU tests" + ), +) + +MAX_RECORDS = 31 +RECORD_WIDTHS = (16, 37, 432) +ODD_RECORD_BYTES = 37 +POOL_RECORDS_PER_RANK = 67 + + +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 _record_values( + global_indices: torch.Tensor, + record_bytes: int, + iteration: int, +) -> torch.Tensor: + byte = torch.arange(record_bytes, device=global_indices.device, dtype=torch.int64) + values = ( + global_indices.to(torch.int64).unsqueeze(1) * 37 + + byte.unsqueeze(0) * 13 + + iteration * 7 + ) + return values.remainder(256).to(torch.uint8) + + +def _local_records( + rank: int, + world_size: int, + device: torch.device, + iteration: int, + record_bytes: int, +) -> torch.Tensor: + global_indices = ( + torch.arange(POOL_RECORDS_PER_RANK, device=device, dtype=torch.int64) + * world_size + + rank + ) + return _record_values(global_indices, record_bytes, iteration) + + +def _selection(destination: int, total_records: int, count: int, iteration: int): + selected = torch.arange(count, dtype=torch.int64) + return (selected * 17 + destination * 11 + iteration * 23).remainder(total_records) + + +def _local_maps( + rank: int, + world_size: int, + count: int, + iteration: int, + device: torch.device, +) -> torch.Tensor: + total_records = POOL_RECORDS_PER_RANK * world_size + maps = torch.full((world_size, count), -1, dtype=torch.int32) + for destination in range(world_size): + selected = _selection(destination, total_records, count, iteration) + owned = selected.remainder(world_size) == rank + maps[destination, owned] = ( + selected[owned].div(world_size, rounding_mode="floor").to(torch.int32) + ) + return maps.to(device) + + +def _expected( + rank: int, + world_size: int, + count: int, + iteration: int, + device: torch.device, + record_bytes: int, +) -> torch.Tensor: + selected = _selection( + rank, + POOL_RECORDS_PER_RANK * world_size, + count, + iteration, + ).to(device) + return _record_values(selected, record_bytes, iteration) + + +def _internal_addresses(exchange: PCIeSelectedRecordExchange) -> tuple[object, ...]: + shared = exchange._owned_buffer + assert shared is not None + return ( + shared.local_ptr, + shared.peer_ptrs, + shared.remote_ptrs, + exchange._local_payload_ptr, + exchange._peer_payload_ptrs.data_ptr(), + exchange._barrier_publish_ptrs.data_ptr(), + exchange._barrier_wait_ptrs.data_ptr(), + exchange._send_counters.data_ptr(), + exchange._wait_counters.data_ptr(), + ) + + +def _check_eager( + exchange: PCIeSelectedRecordExchange, + rank: int, + world_size: int, + device: torch.device, + record_bytes: int, +) -> None: + addresses = _internal_addresses(exchange) + for iteration, count in enumerate((0, 1, 13, MAX_RECORDS), start=1): + records = _local_records(rank, world_size, device, iteration, record_bytes) + maps = _local_maps(rank, world_size, count, iteration, device) + out = torch.empty((count, record_bytes), dtype=torch.uint8, device=device) + returned = exchange.exchange(records, maps, out) + torch.cuda.synchronize(device) + assert returned is out + assert torch.equal( + out, + _expected(rank, world_size, count, iteration, device, record_bytes), + ) + assert _internal_addresses(exchange) == addresses + + iteration = 50 + records = _local_records(rank, world_size, device, iteration, record_bytes) + maps = _local_maps(rank, world_size, MAX_RECORDS, iteration, device) + out = torch.empty((MAX_RECORDS, record_bytes), dtype=torch.uint8, device=device) + exchange.exchange(records, maps, out) + torch.cuda.synchronize(device) + allocated_before = torch.cuda.memory_allocated(device) + torch.cuda.reset_peak_memory_stats(device) + for _ in range(8): + exchange.exchange(records, maps, out) + torch.cuda.synchronize(device) + assert torch.cuda.memory_allocated(device) == allocated_before + assert torch.cuda.max_memory_allocated(device) == allocated_before + assert _internal_addresses(exchange) == addresses + assert torch.equal( + out, + _expected( + rank, + world_size, + MAX_RECORDS, + iteration, + device, + record_bytes, + ), + ) + + +def _check_graph( + exchange: PCIeSelectedRecordExchange, + rank: int, + world_size: int, + device: torch.device, + record_bytes: int, +) -> None: + records = _local_records(rank, world_size, device, 100, record_bytes) + maps = _local_maps(rank, world_size, MAX_RECORDS, 100, device) + out = torch.empty((MAX_RECORDS, record_bytes), dtype=torch.uint8, device=device) + stream = torch.cuda.Stream(device=device) + with torch.cuda.stream(stream): + exchange.exchange(records, maps, out) + stream.synchronize() + dist.barrier() + addresses = _internal_addresses(exchange) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + exchange.exchange(records, maps, out) + graph.replay() + stream.synchronize() + assert _internal_addresses(exchange) == addresses + + iterations = tuple(range(101, 109)) + record_updates = [ + _local_records(rank, world_size, device, iteration, record_bytes) + for iteration in iterations + ] + map_updates = [ + _local_maps(rank, world_size, MAX_RECORDS, iteration, device) + for iteration in iterations + ] + expected = torch.stack( + [ + _expected( + rank, + world_size, + MAX_RECORDS, + iteration, + device, + record_bytes, + ) + for iteration in iterations + ] + ) + observed = torch.empty_like(expected) + allocated_before = torch.cuda.memory_allocated(device) + torch.cuda.reset_peak_memory_stats(device) + for replay in range(len(iterations)): + records.copy_(record_updates[replay]) + maps.copy_(map_updates[replay]) + stream.wait_stream(torch.cuda.current_stream(device)) + graph.replay() + stream.synchronize() + observed[replay].copy_(out) + torch.cuda.synchronize(device) + assert torch.cuda.memory_allocated(device) == allocated_before + assert torch.cuda.max_memory_allocated(device) == allocated_before + assert _internal_addresses(exchange) == addresses + assert torch.equal(observed, expected) + + +def _check_big_pool_offset( + rank: int, + world_size: int, + device: torch.device, +) -> None: + record_bytes = 65_536 + high_record = 2**31 // record_bytes + 1 + required_bytes = (high_record + 1) * record_bytes + torch.cuda.empty_cache() + free_bytes, _ = torch.cuda.mem_get_info(device) + available = torch.tensor( + [1 if rank != 0 or free_bytes >= required_bytes + 512 * 1024**2 else 0], + dtype=torch.int32, + device=device, + ) + dist.all_reduce(available, op=dist.ReduceOp.MIN) + if not bool(available.item()): + raise RuntimeError( + "selected-record big-offset gate requires at least " + f"{required_bytes + 512 * 1024**2} free bytes on rank 0" + ) + + exchange = PCIeSelectedRecordExchange.from_process_group( + process_group=dist.group.WORLD, + device=device, + max_records=1, + record_bytes=record_bytes, + ) + try: + rows = high_record + 1 if rank == 0 else 1 + records = torch.empty((rows, record_bytes), dtype=torch.uint8, device=device) + expected = torch.arange(record_bytes, dtype=torch.int64, device=device) + expected = expected.remainder(251).to(torch.uint8).reshape(1, record_bytes) + if rank == 0: + records[high_record].copy_(expected[0]) + maps = torch.full((world_size, 1), -1, dtype=torch.int32, device=device) + if rank == 0: + maps.fill_(high_record) + out = torch.empty_like(expected) + exchange.exchange(records, maps, out) + torch.cuda.synchronize(device) + assert high_record * record_bytes > 2**31 + assert torch.equal(out, expected) + finally: + exchange.close() + with suppress(Exception): + del records + torch.cuda.empty_cache() + dist.barrier() + + +def _check_large_offset_arithmetic() -> None: + extension = _load_extension() + high_source_record = 2**31 // 65_536 + 1 + high_destination_record = 2**31 // 432 + 1 + assert extension.record_byte_offset_for_test(high_source_record, 65_536) > 2**31 + assert extension.record_byte_offset_for_test(high_destination_record, 432) > 2**31 + + +def _check_rank_configuration_mismatch( + rank: int, + device: torch.device, +) -> None: + from sparkinfer.comm.pcie import pcie_selected_records as module + + malloc_calls = [] + original_malloc = module.CudaRTLibrary.cudaMalloc + + def unexpected_malloc(_self, size): + malloc_calls.append(size) + raise AssertionError("configuration mismatch reached IPC allocation") + + module.CudaRTLibrary.cudaMalloc = unexpected_malloc + try: + with pytest.raises( + PCIeSelectedRecordExchangeInitializationError, + match="configuration mismatch", + ): + PCIeSelectedRecordExchange.from_process_group( + process_group=dist.group.WORLD, + device=device, + max_records=MAX_RECORDS + rank % 2, + record_bytes=ODD_RECORD_BYTES, + ) + assert malloc_calls == [] + finally: + module.CudaRTLibrary.cudaMalloc = original_malloc + dist.barrier() + + +def _check_rank_consistent_handle_serialization_failure( + rank: int, + device: torch.device, +) -> None: + from sparkinfer.comm.pcie import pcie_selected_records as module + + original_get_handle = module.CudaRTLibrary.cudaIpcGetMemHandleBytes + if rank == 0: + + def fail_get_handle(_self, _pointer): + raise RuntimeError("intentional IPC handle serialization failure") + + module.CudaRTLibrary.cudaIpcGetMemHandleBytes = fail_get_handle + try: + with pytest.raises( + PCIeSelectedRecordExchangeInitializationError, + match="handle preparation failed", + ): + PCIeSelectedRecordExchange.from_process_group( + process_group=dist.group.WORLD, + device=device, + max_records=1, + record_bytes=ODD_RECORD_BYTES, + ) + finally: + module.CudaRTLibrary.cudaIpcGetMemHandleBytes = original_get_handle + dist.barrier() + + +def _check_rank_consistent_initialization_failure( + rank: int, + device: torch.device, +) -> None: + from sparkinfer.comm.pcie import pcie_selected_records as module + + original_malloc = module.CudaRTLibrary.cudaMalloc + if rank == 0: + + def fail_malloc(_self, _size): + raise RuntimeError("intentional selected-record allocation failure") + + module.CudaRTLibrary.cudaMalloc = fail_malloc + try: + with pytest.raises(PCIeSelectedRecordExchangeInitializationError): + PCIeSelectedRecordExchange.from_process_group( + process_group=dist.group.WORLD, + device=device, + max_records=1, + record_bytes=ODD_RECORD_BYTES, + ) + finally: + module.CudaRTLibrary.cudaMalloc = original_malloc + dist.barrier() + + +def _worker(rank: int, 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=world_size, + ) + _check_large_offset_arithmetic() + _check_rank_configuration_mismatch(rank, device) + for record_bytes in RECORD_WIDTHS: + exchange = PCIeSelectedRecordExchange.from_process_group( + process_group=dist.group.WORLD, + device=device, + max_records=MAX_RECORDS, + record_bytes=record_bytes, + ) + try: + _check_eager(exchange, rank, world_size, device, record_bytes) + torch.cuda.synchronize(device) + finally: + exchange.close() + dist.barrier() + + graph_exchange = PCIeSelectedRecordExchange.from_process_group( + process_group=dist.group.WORLD, + device=device, + max_records=MAX_RECORDS, + record_bytes=record_bytes, + ) + try: + _check_graph( + graph_exchange, + rank, + world_size, + device, + record_bytes, + ) + torch.cuda.synchronize(device) + finally: + graph_exchange.close() + dist.barrier() + + _check_big_pool_offset(rank, world_size, device) + _check_rank_consistent_handle_serialization_failure(rank, device) + _check_rank_consistent_initialization_failure(rank, device) + dist.destroy_process_group() + + +@pytest.mark.parametrize("world_size", (2, 4), ids=("world2", "world4")) +def test_pcie_selected_records_eager_graph_and_big_offset_correctness( + world_size: int, +) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + if ( + world_size == 4 + and os.getenv("SPARKINFER_RUN_PCIE_SELECTED_RECORDS_WORLD4_TEST") != "1" + ): + pytest.skip( + "set SPARKINFER_RUN_PCIE_SELECTED_RECORDS_WORLD4_TEST=1 for the 4-GPU gate" + ) + if torch.cuda.device_count() < world_size: + pytest.skip( + f"need {world_size} CUDA devices, found {torch.cuda.device_count()}" + ) + _load_extension() + mp.spawn(_worker, args=(world_size, _free_port()), nprocs=world_size, join=True) From b2ace2635b79baa24c375a8975d2dcb1e5e6fc15 Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Mon, 20 Jul 2026 22:50:25 -0500 Subject: [PATCH 2/6] feat(comm): add selected-record copy-engine exchange Add single-plane and three-plane CUDA-IPC copy-engine exchange with rotated peer scheduling, bounded slab overlays, graph-safe deferred release, and large-offset coverage. Co-authored-by: OpenAI Codex Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> --- README.md | 17 + sparkinfer/comm/pcie/__init__.py | 4 + sparkinfer/comm/pcie/api.py | 4 + .../comm/pcie/pcie_selected_records_ce.cu | 1132 ++++++++++++++++ .../comm/pcie/pcie_selected_records_ce.py | 1148 +++++++++++++++++ tests/comm/test_pcie_selected_records_ce.py | 872 +++++++++++++ .../comm/test_pcie_selected_records_ce_gpu.py | 825 ++++++++++++ 7 files changed, 4002 insertions(+) create mode 100644 sparkinfer/comm/pcie/pcie_selected_records_ce.cu create mode 100644 sparkinfer/comm/pcie/pcie_selected_records_ce.py create mode 100644 tests/comm/test_pcie_selected_records_ce.py create mode 100644 tests/comm/test_pcie_selected_records_ce_gpu.py diff --git a/README.md b/README.md index e0ed708c..d2fac462 100644 --- a/README.md +++ b/README.md @@ -120,6 +120,23 @@ Set `SPARKINFER_PRINT_COMPILE_PROGRESS=1` to log each compiler invocation with i cache-key parameters and duration — useful for figuring out what warmup actually covered. `SPARKINFER_TIMING=1` enables per-kernel timing logs. +### Selected-record copy exchange + +`sparkinfer.comm.pcie.SelectedRecordCopyExchange` moves sparse, fixed-width +records between CUDA-IPC peers through the PCIe copy engine. `exchange()` +handles one record plane. `exchange_layers()` packs three planes with one +destination map, issues one rotated DMA exchange and arrival barrier, then +unpacks the planes into caller-owned outputs. The layered layout overlays the +same IPC slab; `layered_max_records` reports its bounded capacity so callers can +fall back to `exchange()` when needed. + +All ranks must call exchanges in the same order and with rank-consistent active +record counts. One exchange object is stream-affine. Run one eager exchange on +that stream before CUDA graph capture, then replay it on the same stream. The +balanced deferred-release protocol makes single-layer, three-layer, and empty +exchanges safe to alternate, but `close()` must run before the process group is +destroyed. + ## Where to look next - `tests/` is the executable spec — per-group API and numerical-reference diff --git a/sparkinfer/comm/pcie/__init__.py b/sparkinfer/comm/pcie/__init__.py index 55117276..b283ac08 100644 --- a/sparkinfer/comm/pcie/__init__.py +++ b/sparkinfer/comm/pcie/__init__.py @@ -13,6 +13,8 @@ - ``DcpAllToAll``: DCP attention exchange with fused LSE reduce-scatter. - ``SelectedRecordExchange``: direct peer writes for destination-selected, fixed-width records. +- ``SelectedRecordCopyExchange``: compact, copy-engine transfer, and unpack + for destination-selected fixed-width records. Raw CUDA (not CuTe): each class JIT-builds its colocated ``.cu`` via torch.utils.cpp_extension, so nvcc must be available at runtime. @@ -37,6 +39,7 @@ "DcpAllToAllPool", "SelectedRecordExchange", "SelectedRecordExchangeInitializationError", + "SelectedRecordCopyExchange", "autotune_dma_crossovers", "parse_oneshot_max_size", "lse_reduce_scatter_reference", @@ -63,6 +66,7 @@ OneshotAllReducePool, SelectedRecordExchange, SelectedRecordExchangeInitializationError, + SelectedRecordCopyExchange, TwoShotReduceScatter, autotune_dma_crossovers, is_supported, diff --git a/sparkinfer/comm/pcie/api.py b/sparkinfer/comm/pcie/api.py index 4f436445..848deb27 100644 --- a/sparkinfer/comm/pcie/api.py +++ b/sparkinfer/comm/pcie/api.py @@ -33,6 +33,9 @@ from .pcie_selected_records import ( PCIeSelectedRecordExchangeInitializationError as SelectedRecordExchangeInitializationError, ) +from .pcie_selected_records_ce import ( + PCIeSelectedRecordCopyExchange as SelectedRecordCopyExchange, +) from .pcie_twoshot import ( PCIeTwoShotSP as TwoShotReduceScatter, ) @@ -58,6 +61,7 @@ def is_supported(device=None) -> bool: "DcpAllToAllPool", "SelectedRecordExchange", "SelectedRecordExchangeInitializationError", + "SelectedRecordCopyExchange", "autotune_dma_crossovers", "parse_oneshot_max_size", "lse_reduce_scatter_reference", diff --git a/sparkinfer/comm/pcie/pcie_selected_records_ce.cu b/sparkinfer/comm/pcie/pcie_selected_records_ce.cu new file mode 100644 index 00000000..b6e1b5ce --- /dev/null +++ b/sparkinfer/comm/pcie/pcie_selected_records_ce.cu @@ -0,0 +1,1132 @@ +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace { + +constexpr int kThreads = 256; +constexpr int kWarpSize = 32; +constexpr int kMaxBlocks = 65535; +constexpr int kMaxWorldSize = 32; +constexpr int kLayerPlanes = 3; + +__host__ __device__ __forceinline__ int64_t record_byte_offset( + int64_t record_index, + int64_t record_bytes) { + return record_index * record_bytes; +} + +__global__ void barrier_all_peers_kernel( + const int64_t* __restrict__ publish_flag_ptrs, + const int64_t* __restrict__ wait_flag_ptrs, + int32_t* __restrict__ send_counters, + int32_t* __restrict__ wait_counters, + int world_size, + uint64_t timeout_cycles, + int phase) { + const int peer = static_cast(threadIdx.x); + if (peer >= world_size) { + return; + } + + const int64_t counter_offset = + static_cast(phase) * world_size + peer; + const int32_t publish_value = send_counters[counter_offset] + 1; + send_counters[counter_offset] = publish_value; + __threadfence_system(); + auto* publish_flag = reinterpret_cast(static_cast( + publish_flag_ptrs[counter_offset])); + asm volatile( + "st.relaxed.sys.global.u32 [%1], %0;" + : + : "r"(publish_value), "l"(publish_flag)); + + const int32_t expected = wait_counters[counter_offset] + 1; + wait_counters[counter_offset] = expected; + const auto* wait_flag = reinterpret_cast( + static_cast(wait_flag_ptrs[counter_offset])); + int32_t observed; + const uint64_t started = clock64(); + uint32_t spins = 0; + do { + asm volatile( + "ld.acquire.sys.global.u32 %0, [%1];" + : "=r"(observed) + : "l"(wait_flag)); + if (((++spins & 0x3ffU) == 0U) && + (clock64() - started > timeout_cycles)) { + printf( + "B12X copy-engine selected-record barrier timed out: phase=%d " + "peer=%d expected=%d observed=%d\n", + phase, + peer, + expected, + observed); + asm volatile("trap;"); + return; + } + } while (static_cast(observed - expected) < 0); +} + +__global__ void publish_all_peers_kernel( + const int64_t* __restrict__ publish_flag_ptrs, + int32_t* __restrict__ send_counters, + int world_size, + int phase) { + const int peer = static_cast(threadIdx.x); + if (peer >= world_size) { + return; + } + + const int64_t counter_offset = + static_cast(phase) * world_size + peer; + const int32_t publish_value = send_counters[counter_offset] + 1; + send_counters[counter_offset] = publish_value; + __threadfence_system(); + auto* publish_flag = reinterpret_cast(static_cast( + publish_flag_ptrs[counter_offset])); + asm volatile( + "st.relaxed.sys.global.u32 [%1], %0;" + : + : "r"(publish_value), "l"(publish_flag)); +} + +__global__ void wait_all_peers_kernel( + const int64_t* __restrict__ wait_flag_ptrs, + int32_t* __restrict__ wait_counters, + int world_size, + uint64_t timeout_cycles, + int phase) { + const int peer = static_cast(threadIdx.x); + if (peer >= world_size) { + return; + } + + const int64_t counter_offset = + static_cast(phase) * world_size + peer; + const int32_t expected = wait_counters[counter_offset] + 1; + wait_counters[counter_offset] = expected; + const auto* wait_flag = reinterpret_cast( + static_cast(wait_flag_ptrs[counter_offset])); + int32_t observed; + const uint64_t started = clock64(); + uint32_t spins = 0; + do { + asm volatile( + "ld.acquire.sys.global.u32 %0, [%1];" + : "=r"(observed) + : "l"(wait_flag)); + if (((++spins & 0x3ffU) == 0U) && + (clock64() - started > timeout_cycles)) { + printf( + "B12X copy-engine selected-record release wait timed out: phase=%d " + "peer=%d expected=%d observed=%d\n", + phase, + peer, + expected, + observed); + asm volatile("trap;"); + return; + } + } while (static_cast(observed - expected) < 0); +} + +__device__ __forceinline__ uint64_t broadcast_lane_zero(uint64_t value) { + uint32_t low = static_cast(value); + uint32_t high = static_cast(value >> 32); + low = __shfl_sync(0xffffffff, low, 0); + high = __shfl_sync(0xffffffff, high, 0); + return (static_cast(high) << 32) | low; +} + +template +__global__ void pack_compact_records_kernel( + const uint8_t* __restrict__ records, + const index_t* __restrict__ local_indices, + uint8_t* __restrict__ primary, + uint8_t* __restrict__ overflow, + int64_t selected_records, + int64_t record_bytes, + int64_t primary_capacity, + int64_t primary_positions_offset, + int64_t primary_records_offset, + int64_t overflow_positions_offset, + int64_t overflow_records_offset) { + const int lane = threadIdx.x % kWarpSize; + const int64_t first_warp = + (static_cast(blockIdx.x) * blockDim.x + threadIdx.x) / + kWarpSize; + const int64_t warp_stride = + static_cast(blockDim.x) * gridDim.x / kWarpSize; + + for (int64_t selected = first_warp; selected < selected_records; + selected += warp_stride) { + const int64_t local_record = static_cast(local_indices[selected]); + if (local_record < 0) { + continue; + } + + uint64_t ordinal = 0; + if (lane == 0) { + ordinal = atomicAdd( + reinterpret_cast(primary), + static_cast(1)); + } + ordinal = broadcast_lane_zero(ordinal); + const bool use_primary = ordinal < static_cast(primary_capacity); + const int64_t compact_index = use_primary + ? static_cast(ordinal) + : static_cast(ordinal) - primary_capacity; + uint8_t* payload = use_primary ? primary : overflow; + const int64_t positions_offset = + use_primary ? primary_positions_offset : overflow_positions_offset; + const int64_t records_offset = + use_primary ? primary_records_offset : overflow_records_offset; + + if (lane == 0) { + reinterpret_cast(payload + positions_offset)[compact_index] = + selected; + } + + const int64_t units_per_record = record_bytes / sizeof(copy_t); + const auto* source = reinterpret_cast( + records + record_byte_offset(local_record, record_bytes)); + auto* destination = reinterpret_cast( + payload + records_offset + + record_byte_offset(compact_index, record_bytes)); + for (int64_t unit = lane; unit < units_per_record; unit += kWarpSize) { + destination[unit] = source[unit]; + } + } +} + +template +__global__ void unpack_compact_records_kernel( + const uint8_t* __restrict__ primary_base, + int64_t primary_stride, + const int64_t* __restrict__ peer_overflow_ptrs, + uint8_t* __restrict__ output, + int64_t selected_records, + int64_t record_bytes, + int64_t primary_capacity, + int64_t primary_positions_offset, + int64_t primary_records_offset, + int64_t overflow_positions_offset, + int64_t overflow_records_offset, + int world_size) { + const int lane = threadIdx.x % kWarpSize; + const int64_t first_warp = + (static_cast(blockIdx.x) * blockDim.x + threadIdx.x) / + kWarpSize; + const int64_t warp_stride = + static_cast(blockDim.x) * gridDim.x / kWarpSize; + const int64_t total_compact_records = + static_cast(world_size) * selected_records; + + for (int64_t compact_record = first_warp; + compact_record < total_compact_records; + compact_record += warp_stride) { + const int source = static_cast(compact_record / selected_records); + const int64_t ordinal = compact_record % selected_records; + const uint8_t* primary = primary_base + source * primary_stride; + const uint64_t total_records = + *reinterpret_cast(primary); + if (ordinal >= static_cast(total_records)) { + continue; + } + + const bool use_primary = ordinal < primary_capacity; + const int64_t compact_index = + use_primary ? ordinal : ordinal - primary_capacity; + const uint8_t* payload = use_primary + ? primary + : reinterpret_cast( + static_cast(peer_overflow_ptrs[source])); + const int64_t positions_offset = + use_primary ? primary_positions_offset : overflow_positions_offset; + const int64_t records_offset = + use_primary ? primary_records_offset : overflow_records_offset; + const int64_t selected = + reinterpret_cast(payload + positions_offset)[compact_index]; + if (selected < 0 || selected >= selected_records) { + continue; + } + + const int64_t units_per_record = record_bytes / sizeof(copy_t); + const auto* source_record = reinterpret_cast( + payload + records_offset + + record_byte_offset(compact_index, record_bytes)); + auto* destination = reinterpret_cast( + output + record_byte_offset(selected, record_bytes)); + for (int64_t unit = lane; unit < units_per_record; unit += kWarpSize) { + destination[unit] = source_record[unit]; + } + } +} + +template +__global__ void pack_compact_record_layers_kernel( + const uint8_t* __restrict__ records0, + const uint8_t* __restrict__ records1, + const uint8_t* __restrict__ records2, + const index_t* __restrict__ local_indices, + uint8_t* __restrict__ primary, + uint8_t* __restrict__ overflow, + int64_t selected_records, + int64_t record_bytes, + int64_t primary_capacity, + int64_t primary_positions_offset, + int64_t primary_records_offset, + int64_t overflow_positions_offset, + int64_t overflow_records_offset) { + const int lane = threadIdx.x % kWarpSize; + const int64_t first_warp = + (static_cast(blockIdx.x) * blockDim.x + threadIdx.x) / + kWarpSize; + const int64_t warp_stride = + static_cast(blockDim.x) * gridDim.x / kWarpSize; + const int64_t packet_record_bytes = record_bytes * kLayerPlanes; + + for (int64_t selected = first_warp; selected < selected_records; + selected += warp_stride) { + const int64_t local_record = static_cast(local_indices[selected]); + if (local_record < 0) { + continue; + } + + uint64_t ordinal = 0; + if (lane == 0) { + ordinal = atomicAdd( + reinterpret_cast(primary), + static_cast(1)); + } + ordinal = broadcast_lane_zero(ordinal); + const bool use_primary = ordinal < static_cast(primary_capacity); + const int64_t compact_index = use_primary + ? static_cast(ordinal) + : static_cast(ordinal) - primary_capacity; + uint8_t* payload = use_primary ? primary : overflow; + const int64_t positions_offset = + use_primary ? primary_positions_offset : overflow_positions_offset; + const int64_t records_offset = + use_primary ? primary_records_offset : overflow_records_offset; + + if (lane == 0) { + reinterpret_cast(payload + positions_offset)[compact_index] = + selected; + } + + const int64_t units_per_record = record_bytes / sizeof(copy_t); + const uint8_t* source_bytes[kLayerPlanes] = {records0, records1, records2}; + uint8_t* packet = payload + records_offset + + record_byte_offset(compact_index, packet_record_bytes); +#pragma unroll + for (int layer = 0; layer < kLayerPlanes; ++layer) { + const auto* source = reinterpret_cast( + source_bytes[layer] + record_byte_offset(local_record, record_bytes)); + auto* destination = reinterpret_cast( + packet + static_cast(layer) * record_bytes); + for (int64_t unit = lane; unit < units_per_record; unit += kWarpSize) { + destination[unit] = source[unit]; + } + } + } +} + +template +__global__ void unpack_compact_record_layers_kernel( + const uint8_t* __restrict__ primary_base, + int64_t primary_stride, + const int64_t* __restrict__ peer_overflow_ptrs, + uint8_t* __restrict__ output0, + uint8_t* __restrict__ output1, + uint8_t* __restrict__ output2, + int64_t selected_records, + int64_t record_bytes, + int64_t primary_capacity, + int64_t primary_positions_offset, + int64_t primary_records_offset, + int64_t overflow_positions_offset, + int64_t overflow_records_offset, + int world_size) { + const int lane = threadIdx.x % kWarpSize; + const int64_t first_warp = + (static_cast(blockIdx.x) * blockDim.x + threadIdx.x) / + kWarpSize; + const int64_t warp_stride = + static_cast(blockDim.x) * gridDim.x / kWarpSize; + const int64_t total_compact_records = + static_cast(world_size) * selected_records; + const int64_t packet_record_bytes = record_bytes * kLayerPlanes; + + for (int64_t compact_record = first_warp; + compact_record < total_compact_records; + compact_record += warp_stride) { + const int source_rank = + static_cast(compact_record / selected_records); + const int64_t ordinal = compact_record % selected_records; + const uint8_t* primary = primary_base + source_rank * primary_stride; + const uint64_t total_records = + *reinterpret_cast(primary); + if (ordinal >= static_cast(total_records)) { + continue; + } + + const bool use_primary = ordinal < primary_capacity; + const int64_t compact_index = + use_primary ? ordinal : ordinal - primary_capacity; + const uint8_t* payload = use_primary + ? primary + : reinterpret_cast( + static_cast(peer_overflow_ptrs[source_rank])); + const int64_t positions_offset = + use_primary ? primary_positions_offset : overflow_positions_offset; + const int64_t records_offset = + use_primary ? primary_records_offset : overflow_records_offset; + const int64_t selected = + reinterpret_cast(payload + positions_offset)[compact_index]; + if (selected < 0 || selected >= selected_records) { + continue; + } + + const int64_t units_per_record = record_bytes / sizeof(copy_t); + const uint8_t* packet = payload + records_offset + + record_byte_offset(compact_index, packet_record_bytes); + uint8_t* output_bytes[kLayerPlanes] = {output0, output1, output2}; +#pragma unroll + for (int layer = 0; layer < kLayerPlanes; ++layer) { + const auto* source = reinterpret_cast( + packet + static_cast(layer) * record_bytes); + auto* destination = reinterpret_cast( + output_bytes[layer] + record_byte_offset(selected, record_bytes)); + for (int64_t unit = lane; unit < units_per_record; unit += kWarpSize) { + destination[unit] = source[unit]; + } + } + } +} + +void validate_tensor( + const torch::Tensor& tensor, + const char* name, + torch::ScalarType scalar_type) { + TORCH_CHECK(tensor.is_cuda(), name, " must be CUDA"); + TORCH_CHECK(tensor.scalar_type() == scalar_type, name, " has wrong dtype"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +int launch_blocks_for_warps(int64_t warps) { + const int64_t requested = + (warps * kWarpSize + kThreads - 1) / kThreads; + return static_cast( + std::max(1, std::min(requested, kMaxBlocks))); +} + +template +void launch_pack( + const torch::Tensor& records, + const torch::Tensor& local_indices, + int64_t primary_ptr, + int64_t overflow_ptr, + int64_t record_bytes, + int64_t primary_capacity, + int64_t primary_positions_offset, + int64_t primary_records_offset, + int64_t overflow_positions_offset, + int64_t overflow_records_offset, + cudaStream_t stream) { + const int64_t selected_records = local_indices.numel(); + const int blocks = launch_blocks_for_warps(selected_records); + pack_compact_records_kernel<<>>( + records.data_ptr(), + local_indices.data_ptr(), + reinterpret_cast(static_cast(primary_ptr)), + reinterpret_cast(static_cast(overflow_ptr)), + selected_records, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset); +} + +void pack_compact_records( + torch::Tensor records, + torch::Tensor local_indices, + int64_t primary_ptr, + int64_t overflow_ptr, + int64_t record_bytes, + int64_t primary_capacity, + int64_t primary_positions_offset, + int64_t primary_records_offset, + int64_t overflow_positions_offset, + int64_t overflow_records_offset) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(records)); + validate_tensor(records, "records", torch::kUInt8); + TORCH_CHECK(local_indices.is_cuda(), "local indices must be CUDA"); + TORCH_CHECK( + local_indices.scalar_type() == torch::kInt32 || + local_indices.scalar_type() == torch::kInt64, + "local indices must be int32 or int64"); + TORCH_CHECK(local_indices.is_contiguous(), "local indices must be contiguous"); + TORCH_CHECK(records.dim() >= 2, "records must have rank at least 2"); + TORCH_CHECK(record_bytes > 0, "record_bytes must be positive"); + TORCH_CHECK( + records.size(-1) == record_bytes, + "records must end in record_bytes"); + TORCH_CHECK(primary_ptr != 0, "primary pointer must be nonzero"); + TORCH_CHECK(overflow_ptr != 0, "overflow pointer must be nonzero"); + TORCH_CHECK(primary_capacity > 0, "primary capacity must be positive"); + TORCH_CHECK( + primary_capacity <= local_indices.numel(), + "primary capacity cannot exceed selected records"); + + const auto stream = c10::cuda::getCurrentCUDAStream().stream(); + AT_CUDA_CHECK(cudaMemsetAsync( + reinterpret_cast(static_cast(primary_ptr)), + 0, + sizeof(uint64_t), + stream)); + if (local_indices.numel() == 0) { + return; + } + + const bool vectorized = + record_bytes % static_cast(sizeof(uint4)) == 0 && + reinterpret_cast(records.data_ptr()) % alignof(uint4) == 0 && + (static_cast(primary_ptr) + primary_records_offset) % + alignof(uint4) == + 0 && + (static_cast(overflow_ptr) + overflow_records_offset) % + alignof(uint4) == + 0; + if (local_indices.scalar_type() == torch::kInt32) { + if (vectorized) { + launch_pack( + records, + local_indices, + primary_ptr, + overflow_ptr, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + stream); + } else { + launch_pack( + records, + local_indices, + primary_ptr, + overflow_ptr, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + stream); + } + } else if (vectorized) { + launch_pack( + records, + local_indices, + primary_ptr, + overflow_ptr, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + stream); + } else { + launch_pack( + records, + local_indices, + primary_ptr, + overflow_ptr, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + stream); + } + AT_CUDA_CHECK(cudaGetLastError()); +} + +template +void launch_pack_layers( + const torch::Tensor& records0, + const torch::Tensor& records1, + const torch::Tensor& records2, + const torch::Tensor& local_indices, + int64_t primary_ptr, + int64_t overflow_ptr, + int64_t record_bytes, + int64_t primary_capacity, + int64_t primary_positions_offset, + int64_t primary_records_offset, + int64_t overflow_positions_offset, + int64_t overflow_records_offset, + cudaStream_t stream) { + const int64_t selected_records = local_indices.numel(); + const int blocks = launch_blocks_for_warps(selected_records); + pack_compact_record_layers_kernel + <<>>( + records0.data_ptr(), + records1.data_ptr(), + records2.data_ptr(), + local_indices.data_ptr(), + reinterpret_cast(static_cast(primary_ptr)), + reinterpret_cast(static_cast(overflow_ptr)), + selected_records, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset); +} + +void pack_compact_record_layers( + torch::Tensor records0, + torch::Tensor records1, + torch::Tensor records2, + torch::Tensor local_indices, + int64_t primary_ptr, + int64_t overflow_ptr, + int64_t record_bytes, + int64_t primary_capacity, + int64_t primary_positions_offset, + int64_t primary_records_offset, + int64_t overflow_positions_offset, + int64_t overflow_records_offset) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(records0)); + validate_tensor(records0, "layer 0 records", torch::kUInt8); + validate_tensor(records1, "layer 1 records", torch::kUInt8); + validate_tensor(records2, "layer 2 records", torch::kUInt8); + TORCH_CHECK( + records1.device() == records0.device() && + records2.device() == records0.device(), + "all layer records must use the same CUDA device"); + TORCH_CHECK( + records0.sizes() == records1.sizes() && + records0.sizes() == records2.sizes(), + "all layer record tensors must have the same shape"); + TORCH_CHECK(local_indices.is_cuda(), "local indices must be CUDA"); + TORCH_CHECK( + local_indices.device() == records0.device(), + "local indices and records must use the same CUDA device"); + TORCH_CHECK( + local_indices.scalar_type() == torch::kInt32 || + local_indices.scalar_type() == torch::kInt64, + "local indices must be int32 or int64"); + TORCH_CHECK(local_indices.is_contiguous(), "local indices must be contiguous"); + TORCH_CHECK(records0.dim() >= 2, "records must have rank at least 2"); + TORCH_CHECK(record_bytes > 0, "record_bytes must be positive"); + TORCH_CHECK( + record_bytes <= std::numeric_limits::max() / kLayerPlanes, + "layered record width exceeds int64 capacity"); + TORCH_CHECK( + records0.size(-1) == record_bytes, + "layer records must end in record_bytes"); + TORCH_CHECK(primary_ptr != 0, "primary pointer must be nonzero"); + TORCH_CHECK(overflow_ptr != 0, "overflow pointer must be nonzero"); + TORCH_CHECK(primary_capacity > 0, "primary capacity must be positive"); + TORCH_CHECK( + primary_capacity <= local_indices.numel(), + "primary capacity cannot exceed selected records"); + + const auto stream = c10::cuda::getCurrentCUDAStream().stream(); + AT_CUDA_CHECK(cudaMemsetAsync( + reinterpret_cast(static_cast(primary_ptr)), + 0, + sizeof(uint64_t), + stream)); + if (local_indices.numel() == 0) { + return; + } + + const bool vectorized = + record_bytes % static_cast(sizeof(uint4)) == 0 && + reinterpret_cast(records0.data_ptr()) % + alignof(uint4) == + 0 && + reinterpret_cast(records1.data_ptr()) % + alignof(uint4) == + 0 && + reinterpret_cast(records2.data_ptr()) % + alignof(uint4) == + 0 && + (static_cast(primary_ptr) + primary_records_offset) % + alignof(uint4) == + 0 && + (static_cast(overflow_ptr) + overflow_records_offset) % + alignof(uint4) == + 0; + if (local_indices.scalar_type() == torch::kInt32) { + if (vectorized) { + launch_pack_layers( + records0, + records1, + records2, + local_indices, + primary_ptr, + overflow_ptr, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + stream); + } else { + launch_pack_layers( + records0, + records1, + records2, + local_indices, + primary_ptr, + overflow_ptr, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + stream); + } + } else if (vectorized) { + launch_pack_layers( + records0, + records1, + records2, + local_indices, + primary_ptr, + overflow_ptr, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + stream); + } else { + launch_pack_layers( + records0, + records1, + records2, + local_indices, + primary_ptr, + overflow_ptr, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + stream); + } + AT_CUDA_CHECK(cudaGetLastError()); +} + +template +void launch_unpack( + int64_t primary_base_ptr, + int64_t primary_stride, + const torch::Tensor& peer_overflow_ptrs, + torch::Tensor& output, + int64_t selected_records, + int64_t record_bytes, + int64_t primary_capacity, + int64_t primary_positions_offset, + int64_t primary_records_offset, + int64_t overflow_positions_offset, + int64_t overflow_records_offset, + cudaStream_t stream) { + const int world_size = static_cast(peer_overflow_ptrs.numel()); + const int64_t warps = static_cast(world_size) * selected_records; + const int blocks = launch_blocks_for_warps(warps); + unpack_compact_records_kernel<<>>( + reinterpret_cast( + static_cast(primary_base_ptr)), + primary_stride, + peer_overflow_ptrs.data_ptr(), + output.data_ptr(), + selected_records, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + world_size); +} + +void unpack_compact_records( + int64_t primary_base_ptr, + int64_t primary_stride, + torch::Tensor peer_overflow_ptrs, + torch::Tensor output, + int64_t selected_records, + int64_t record_bytes, + int64_t primary_capacity, + int64_t primary_positions_offset, + int64_t primary_records_offset, + int64_t overflow_positions_offset, + int64_t overflow_records_offset) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(output)); + validate_tensor(peer_overflow_ptrs, "peer overflow pointers", torch::kInt64); + validate_tensor(output, "output", torch::kUInt8); + TORCH_CHECK(primary_base_ptr != 0, "primary base pointer must be nonzero"); + TORCH_CHECK(primary_stride > 0, "primary stride must be positive"); + TORCH_CHECK(record_bytes > 0, "record_bytes must be positive"); + TORCH_CHECK(selected_records > 0, "selected_records must be positive"); + TORCH_CHECK(primary_capacity > 0, "primary capacity must be positive"); + const int64_t world_size = peer_overflow_ptrs.numel(); + TORCH_CHECK( + world_size >= 2 && world_size <= kMaxWorldSize, + "world size must be in [2, 32]"); + TORCH_CHECK( + selected_records <= std::numeric_limits::max() / record_bytes, + "selected payload exceeds int64 capacity"); + TORCH_CHECK( + output.numel() == selected_records * record_bytes, + "output size must exactly match the selected payload"); + + const auto stream = c10::cuda::getCurrentCUDAStream().stream(); + const bool vectorized = + record_bytes % static_cast(sizeof(uint4)) == 0 && + reinterpret_cast(output.data_ptr()) % alignof(uint4) == 0 && + (static_cast(primary_base_ptr) + primary_records_offset) % + alignof(uint4) == + 0; + if (vectorized) { + launch_unpack( + primary_base_ptr, + primary_stride, + peer_overflow_ptrs, + output, + selected_records, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + stream); + } else { + launch_unpack( + primary_base_ptr, + primary_stride, + peer_overflow_ptrs, + output, + selected_records, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + stream); + } + AT_CUDA_CHECK(cudaGetLastError()); +} + +template +void launch_unpack_layers( + int64_t primary_base_ptr, + int64_t primary_stride, + const torch::Tensor& peer_overflow_ptrs, + torch::Tensor& output0, + torch::Tensor& output1, + torch::Tensor& output2, + int64_t selected_records, + int64_t record_bytes, + int64_t primary_capacity, + int64_t primary_positions_offset, + int64_t primary_records_offset, + int64_t overflow_positions_offset, + int64_t overflow_records_offset, + cudaStream_t stream) { + const int world_size = static_cast(peer_overflow_ptrs.numel()); + const int64_t warps = static_cast(world_size) * selected_records; + const int blocks = launch_blocks_for_warps(warps); + unpack_compact_record_layers_kernel<<>>( + reinterpret_cast( + static_cast(primary_base_ptr)), + primary_stride, + peer_overflow_ptrs.data_ptr(), + output0.data_ptr(), + output1.data_ptr(), + output2.data_ptr(), + selected_records, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + world_size); +} + +void unpack_compact_record_layers( + int64_t primary_base_ptr, + int64_t primary_stride, + torch::Tensor peer_overflow_ptrs, + torch::Tensor output0, + torch::Tensor output1, + torch::Tensor output2, + int64_t selected_records, + int64_t record_bytes, + int64_t primary_capacity, + int64_t primary_positions_offset, + int64_t primary_records_offset, + int64_t overflow_positions_offset, + int64_t overflow_records_offset) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(output0)); + validate_tensor(peer_overflow_ptrs, "peer overflow pointers", torch::kInt64); + validate_tensor(output0, "layer 0 output", torch::kUInt8); + validate_tensor(output1, "layer 1 output", torch::kUInt8); + validate_tensor(output2, "layer 2 output", torch::kUInt8); + TORCH_CHECK( + output1.device() == output0.device() && + output2.device() == output0.device() && + peer_overflow_ptrs.device() == output0.device(), + "layer outputs and peer pointers must use the same CUDA device"); + TORCH_CHECK( + output0.sizes() == output1.sizes() && + output0.sizes() == output2.sizes(), + "all layer output tensors must have the same shape"); + TORCH_CHECK(primary_base_ptr != 0, "primary base pointer must be nonzero"); + TORCH_CHECK(primary_stride > 0, "primary stride must be positive"); + TORCH_CHECK(record_bytes > 0, "record_bytes must be positive"); + TORCH_CHECK( + record_bytes <= std::numeric_limits::max() / kLayerPlanes, + "layered record width exceeds int64 capacity"); + TORCH_CHECK(selected_records > 0, "selected_records must be positive"); + TORCH_CHECK(primary_capacity > 0, "primary capacity must be positive"); + const int64_t world_size = peer_overflow_ptrs.numel(); + TORCH_CHECK( + world_size >= 2 && world_size <= kMaxWorldSize, + "world size must be in [2, 32]"); + TORCH_CHECK( + selected_records <= std::numeric_limits::max() / record_bytes, + "selected payload exceeds int64 capacity"); + TORCH_CHECK( + output0.numel() == selected_records * record_bytes, + "each output size must exactly match the selected payload"); + + const auto stream = c10::cuda::getCurrentCUDAStream().stream(); + const bool vectorized = + record_bytes % static_cast(sizeof(uint4)) == 0 && + reinterpret_cast(output0.data_ptr()) % + alignof(uint4) == + 0 && + reinterpret_cast(output1.data_ptr()) % + alignof(uint4) == + 0 && + reinterpret_cast(output2.data_ptr()) % + alignof(uint4) == + 0 && + (static_cast(primary_base_ptr) + primary_records_offset) % + alignof(uint4) == + 0; + if (vectorized) { + launch_unpack_layers( + primary_base_ptr, + primary_stride, + peer_overflow_ptrs, + output0, + output1, + output2, + selected_records, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + stream); + } else { + launch_unpack_layers( + primary_base_ptr, + primary_stride, + peer_overflow_ptrs, + output0, + output1, + output2, + selected_records, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + stream); + } + AT_CUDA_CHECK(cudaGetLastError()); +} + +void barrier_all_peers( + torch::Tensor publish_flag_ptrs, + torch::Tensor wait_flag_ptrs, + torch::Tensor send_counters, + torch::Tensor wait_counters, + int64_t phase, + int64_t timeout_cycles) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(send_counters)); + validate_tensor(publish_flag_ptrs, "publish flag pointers", torch::kInt64); + validate_tensor(wait_flag_ptrs, "wait flag pointers", torch::kInt64); + validate_tensor(send_counters, "send counters", torch::kInt32); + validate_tensor(wait_counters, "wait counters", torch::kInt32); + TORCH_CHECK(phase == 0 || phase == 1, "phase must be zero or one"); + TORCH_CHECK(timeout_cycles > 0, "timeout_cycles must be positive"); + TORCH_CHECK( + publish_flag_ptrs.numel() == wait_flag_ptrs.numel(), + "publish and wait pointer counts must match"); + TORCH_CHECK( + send_counters.numel() == wait_counters.numel(), + "send and wait counter counts must match"); + TORCH_CHECK( + publish_flag_ptrs.numel() == send_counters.numel(), + "barrier pointer and counter counts must match"); + TORCH_CHECK( + publish_flag_ptrs.dim() == 2 && publish_flag_ptrs.size(0) == 2, + "barrier pointers must have shape [2, world_size]"); + const int world_size = static_cast(publish_flag_ptrs.size(1)); + TORCH_CHECK( + world_size >= 2 && world_size <= kMaxWorldSize, + "world size must be in [2, 32]"); + + const auto stream = c10::cuda::getCurrentCUDAStream().stream(); + barrier_all_peers_kernel<<<1, kMaxWorldSize, 0, stream>>>( + publish_flag_ptrs.data_ptr(), + wait_flag_ptrs.data_ptr(), + send_counters.data_ptr(), + wait_counters.data_ptr(), + world_size, + static_cast(timeout_cycles), + static_cast(phase)); + AT_CUDA_CHECK(cudaGetLastError()); +} + +void publish_all_peers( + torch::Tensor publish_flag_ptrs, + torch::Tensor send_counters, + int64_t phase) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(send_counters)); + validate_tensor(publish_flag_ptrs, "publish flag pointers", torch::kInt64); + validate_tensor(send_counters, "send counters", torch::kInt32); + TORCH_CHECK(phase == 0 || phase == 1, "phase must be zero or one"); + TORCH_CHECK( + publish_flag_ptrs.numel() == send_counters.numel(), + "publish pointer and counter counts must match"); + TORCH_CHECK( + publish_flag_ptrs.dim() == 2 && publish_flag_ptrs.size(0) == 2, + "publish pointers must have shape [2, world_size]"); + const int world_size = static_cast(publish_flag_ptrs.size(1)); + TORCH_CHECK( + world_size >= 2 && world_size <= kMaxWorldSize, + "world size must be in [2, 32]"); + + const auto stream = c10::cuda::getCurrentCUDAStream().stream(); + publish_all_peers_kernel<<<1, kMaxWorldSize, 0, stream>>>( + publish_flag_ptrs.data_ptr(), + send_counters.data_ptr(), + world_size, + static_cast(phase)); + AT_CUDA_CHECK(cudaGetLastError()); +} + +void wait_all_peers( + torch::Tensor wait_flag_ptrs, + torch::Tensor wait_counters, + int64_t phase, + int64_t timeout_cycles) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(wait_counters)); + validate_tensor(wait_flag_ptrs, "wait flag pointers", torch::kInt64); + validate_tensor(wait_counters, "wait counters", torch::kInt32); + TORCH_CHECK(phase == 0 || phase == 1, "phase must be zero or one"); + TORCH_CHECK(timeout_cycles > 0, "timeout_cycles must be positive"); + TORCH_CHECK( + wait_flag_ptrs.numel() == wait_counters.numel(), + "wait pointer and counter counts must match"); + TORCH_CHECK( + wait_flag_ptrs.dim() == 2 && wait_flag_ptrs.size(0) == 2, + "wait pointers must have shape [2, world_size]"); + const int world_size = static_cast(wait_flag_ptrs.size(1)); + TORCH_CHECK( + world_size >= 2 && world_size <= kMaxWorldSize, + "world size must be in [2, 32]"); + + const auto stream = c10::cuda::getCurrentCUDAStream().stream(); + wait_all_peers_kernel<<<1, kMaxWorldSize, 0, stream>>>( + wait_flag_ptrs.data_ptr(), + wait_counters.data_ptr(), + world_size, + static_cast(timeout_cycles), + static_cast(phase)); + AT_CUDA_CHECK(cudaGetLastError()); +} + +int64_t record_byte_offset_for_test( + int64_t record_index, + int64_t record_bytes) { + TORCH_CHECK(record_index >= 0, "record_index must be non-negative"); + TORCH_CHECK(record_bytes > 0, "record_bytes must be positive"); + TORCH_CHECK( + record_index <= std::numeric_limits::max() / record_bytes, + "record byte offset exceeds int64 capacity"); + return record_byte_offset(record_index, record_bytes); +} + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def( + "pack_compact_records", + &pack_compact_records, + "Pack locally owned selected records into copy-engine staging packets"); + module.def( + "pack_compact_record_layers", + &pack_compact_record_layers, + "Pack three record planes into one copy-engine staging packet"); + module.def( + "unpack_compact_records", + &unpack_compact_records, + "Unpack source packets into destination-selected record order"); + module.def( + "unpack_compact_record_layers", + &unpack_compact_record_layers, + "Unpack one packet directly into three destination record planes"); + module.def( + "barrier_all_peers", + &barrier_all_peers, + "Publish and wait for every selected-record peer"); + module.def( + "publish_all_peers", + &publish_all_peers, + "Publish one selected-record generation to every peer"); + module.def( + "wait_all_peers", + &wait_all_peers, + "Wait for one selected-record generation from every peer"); + module.def( + "record_byte_offset_for_test", + &record_byte_offset_for_test, + "Compute the Int64 byte offset used by selected-record copy exchange"); +} diff --git a/sparkinfer/comm/pcie/pcie_selected_records_ce.py b/sparkinfer/comm/pcie/pcie_selected_records_ce.py new file mode 100644 index 00000000..511959da --- /dev/null +++ b/sparkinfer/comm/pcie/pcie_selected_records_ce.py @@ -0,0 +1,1148 @@ +"""Copy-engine PCIe exchange for destination-selected fixed-width records.""" + +from __future__ import annotations + +from contextlib import suppress +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path +from typing import Optional, Sequence + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup +from torch.utils.cpp_extension import load + +from . import pcie_dma as _pcie_dma +from ._cuda_ipc import CudaRTLibrary +from .pcie_dma import FLAG_STRIDE +from .pcie_oneshot import ( + IPC_SLAB_ALIGNMENT, + _current_stream_key, + _is_current_stream_capturing, + _normalize_device, + _OwnedSharedBuffer, + _align_up, +) +from .pcie_selected_records import ( + DEFAULT_BARRIER_TIMEOUT_CYCLES, + MAX_WORLD_SIZE, + PCIeSelectedRecordExchangeInitializationError, + _all_ranks_succeeded, + _allocate_shared_buffer_rank_consistent, + _free_shared_buffer, +) + + +_MAX_INT64 = (1 << 63) - 1 +_CONFIG_VERSION = 2 +_COUNT_BYTES = 8 +_POSITION_BYTES = 8 +_PACKET_HEADER_BYTES = 2 * _COUNT_BYTES +_LAYER_PLANES = 3 +_CONFIG_FIELDS = ( + "version", + "world_size", + "max_records", + "record_bytes", + "layered_max_records", + "layered_record_bytes", + "layered_slab_bytes", + "primary_capacity", + "overflow_capacity", + "flags_bytes", + "receive_primary_base", + "staging_primary_base", + "staging_overflow_base", + "primary_positions_offset", + "primary_records_offset", + "primary_stride", + "overflow_positions_offset", + "overflow_records_offset", + "overflow_stride", + "slab_bytes", + "flag_stride", + "slab_alignment", +) + + +@dataclass(frozen=True) +class _CopyExchangeLayout: + flags_bytes: int + receive_primary_base: int + staging_primary_base: int + staging_overflow_base: int + primary_capacity: int + overflow_capacity: int + primary_positions_offset: int + primary_records_offset: int + primary_stride: int + overflow_positions_offset: int + overflow_records_offset: int + overflow_stride: int + slab_bytes: int + + +def _checked_region_end(base: int, count: int, stride: int, name: str) -> int: + if count < 0 or stride < 0 or base < 0: + raise ValueError(f"{name} has a negative layout component") + if count and stride > (_MAX_INT64 - base) // count: + raise ValueError(f"{name} exceeds int64 capacity") + return base + count * stride + + +def _copy_exchange_layout( + *, + world_size: int, + max_records: int, + record_bytes: int, + primary_capacity: Optional[int] = None, +) -> _CopyExchangeLayout: + world_size = int(world_size) + max_records = int(max_records) + record_bytes = int(record_bytes) + if not 2 <= world_size <= MAX_WORLD_SIZE: + raise ValueError( + f"world_size must be in [2, {MAX_WORLD_SIZE}], got {world_size}" + ) + if max_records <= 0: + raise ValueError("max_records must be positive") + if record_bytes <= 0: + raise ValueError("record_bytes must be positive") + if max_records > _MAX_INT64 // record_bytes: + raise ValueError("max_records * record_bytes exceeds int64 capacity") + + if primary_capacity is None: + primary_capacity = (max_records + world_size - 1) // world_size + primary_capacity = int(primary_capacity) + if not 1 <= primary_capacity <= max_records: + raise ValueError("primary_capacity must be in [1, max_records]") + overflow_capacity = max_records - primary_capacity + + flags_bytes = _align_up( + 2 * world_size * FLAG_STRIDE, + IPC_SLAB_ALIGNMENT, + ) + primary_positions_offset = _PACKET_HEADER_BYTES + primary_records_offset = _align_up( + primary_positions_offset + primary_capacity * _POSITION_BYTES, + 16, + ) + primary_stride = _align_up( + _checked_region_end( + primary_records_offset, + primary_capacity, + record_bytes, + "primary packet", + ), + IPC_SLAB_ALIGNMENT, + ) + + # Keep a valid, distinct overflow pointer even when primary_capacity covers + # the complete pool. No overflow bytes are consumed in that configuration. + overflow_storage_records = max(1, overflow_capacity) + overflow_positions_offset = 0 + overflow_records_offset = _align_up( + overflow_storage_records * _POSITION_BYTES, + 16, + ) + overflow_stride = _align_up( + _checked_region_end( + overflow_records_offset, + overflow_storage_records, + record_bytes, + "overflow packet", + ), + IPC_SLAB_ALIGNMENT, + ) + + receive_primary_base = flags_bytes + staging_primary_base = _align_up( + _checked_region_end( + receive_primary_base, + world_size, + primary_stride, + "primary receive area", + ), + IPC_SLAB_ALIGNMENT, + ) + staging_overflow_base = _align_up( + _checked_region_end( + staging_primary_base, + world_size, + primary_stride, + "primary staging area", + ), + IPC_SLAB_ALIGNMENT, + ) + slab_bytes = _checked_region_end( + staging_overflow_base, + world_size, + overflow_stride, + "overflow staging area", + ) + if slab_bytes > _MAX_INT64: + raise ValueError("copy-engine selected-record slab exceeds int64 capacity") + + return _CopyExchangeLayout( + flags_bytes=flags_bytes, + receive_primary_base=receive_primary_base, + staging_primary_base=staging_primary_base, + staging_overflow_base=staging_overflow_base, + primary_capacity=primary_capacity, + overflow_capacity=overflow_capacity, + primary_positions_offset=primary_positions_offset, + primary_records_offset=primary_records_offset, + primary_stride=primary_stride, + overflow_positions_offset=overflow_positions_offset, + overflow_records_offset=overflow_records_offset, + overflow_stride=overflow_stride, + slab_bytes=slab_bytes, + ) + + +def _largest_layered_capacity_for_slab( + *, + world_size: int, + max_records: int, + record_bytes: int, + slab_bytes: int, +) -> tuple[int, _CopyExchangeLayout]: + max_records = int(max_records) + if max_records <= 0: + raise ValueError("max_records must be positive") + if record_bytes > _MAX_INT64 // _LAYER_PLANES: + raise ValueError("layered record width exceeds int64 capacity") + layered_record_bytes = record_bytes * _LAYER_PLANES + low = 1 + high = max_records + first_layout = _copy_exchange_layout( + world_size=world_size, + max_records=1, + record_bytes=layered_record_bytes, + ) + if first_layout.slab_bytes > int(slab_bytes): + raise ValueError("base slab cannot hold one three-layer record") + while low < high: + candidate = (low + high + 1) // 2 + candidate_layout = _copy_exchange_layout( + world_size=world_size, + max_records=candidate, + record_bytes=layered_record_bytes, + ) + if candidate_layout.slab_bytes <= int(slab_bytes): + low = candidate + else: + high = candidate - 1 + layered_layout = _copy_exchange_layout( + world_size=world_size, + max_records=low, + record_bytes=layered_record_bytes, + ) + return low, layered_layout + + +def _layered_layout_for_slab( + *, + world_size: int, + max_records: int, + record_bytes: int, + slab_bytes: int, + layered_max_records: Optional[int], + fallback_layout: _CopyExchangeLayout, +) -> tuple[int, _CopyExchangeLayout]: + if layered_max_records is None: + try: + return _largest_layered_capacity_for_slab( + world_size=world_size, + max_records=max_records, + record_bytes=record_bytes, + slab_bytes=slab_bytes, + ) + except ValueError as exc: + if str(exc) != "base slab cannot hold one three-layer record": + raise + return 0, fallback_layout + + layered_max_records = int(layered_max_records) + if layered_max_records == 0: + return 0, fallback_layout + if not 1 <= layered_max_records <= int(max_records): + raise ValueError("layered_max_records must be in [0, max_records]") + if record_bytes > _MAX_INT64 // _LAYER_PLANES: + raise ValueError("layered record width exceeds int64 capacity") + layered_layout = _copy_exchange_layout( + world_size=world_size, + max_records=layered_max_records, + record_bytes=record_bytes * _LAYER_PLANES, + ) + if layered_layout.slab_bytes > int(slab_bytes): + raise ValueError( + "three-layer layout exceeds the existing selected-record slab " + f"({layered_layout.slab_bytes} > {slab_bytes} bytes)" + ) + return layered_max_records, layered_layout + + +@lru_cache(maxsize=1) +def _load_extension(): + source = Path(__file__).with_name("pcie_selected_records_ce.cu") + return load( + name="sparkinfer_pcie_selected_records_ce_ext", + sources=[str(source)], + extra_cuda_cflags=["-O3"], + extra_ldflags=["-lcuda"], + verbose=False, + ) + + +def _configuration_values( + *, + world_size: int, + max_records: int, + record_bytes: int, + layout: _CopyExchangeLayout, + layered_max_records: int, + layered_layout: _CopyExchangeLayout, +) -> tuple[int, ...]: + return ( + _CONFIG_VERSION, + int(world_size), + int(max_records), + int(record_bytes), + int(layered_max_records), + int(record_bytes) * _LAYER_PLANES, + layered_layout.slab_bytes, + layout.primary_capacity, + layout.overflow_capacity, + layout.flags_bytes, + layout.receive_primary_base, + layout.staging_primary_base, + layout.staging_overflow_base, + layout.primary_positions_offset, + layout.primary_records_offset, + layout.primary_stride, + layout.overflow_positions_offset, + layout.overflow_records_offset, + layout.overflow_stride, + layout.slab_bytes, + FLAG_STRIDE, + IPC_SLAB_ALIGNMENT, + ) + + +def _validate_rank_configuration( + *, + process_group: ProcessGroup, + device: torch.device, + status: torch.Tensor, + world_size: int, + max_records: int, + record_bytes: int, + layout: _CopyExchangeLayout, + layered_max_records: int, + layered_layout: _CopyExchangeLayout, +) -> None: + local_config: Optional[torch.Tensor] = None + gathered_configs: Optional[torch.Tensor] = None + local_error: Optional[Exception] = None + try: + local_config = torch.tensor( + _configuration_values( + world_size=world_size, + max_records=max_records, + record_bytes=record_bytes, + layout=layout, + layered_max_records=layered_max_records, + layered_layout=layered_layout, + ), + dtype=torch.int64, + device=device, + ) + gathered_configs = torch.empty( + world_size * len(_CONFIG_FIELDS), + dtype=torch.int64, + device=device, + ) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + raise PCIeSelectedRecordExchangeInitializationError( + "copy-engine selected-record configuration allocation failed on at " + "least one rank" + ) from local_error + + assert local_config is not None + assert gathered_configs is not None + try: + dist.all_gather_into_tensor( + gathered_configs, + local_config, + group=process_group, + ) + except Exception as exc: + raise PCIeSelectedRecordExchangeInitializationError( + "copy-engine selected-record configuration collective failed" + ) from exc + + rows = gathered_configs.view(world_size, len(_CONFIG_FIELDS)) + local_matches = False + local_error = None + try: + local_matches = bool(torch.all(rows == local_config.unsqueeze(0)).item()) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + raise PCIeSelectedRecordExchangeInitializationError( + "copy-engine selected-record configuration comparison failed on at " + "least one rank" + ) from local_error + if not _all_ranks_succeeded(status, local_matches, process_group): + try: + rank_configs = rows.cpu().tolist() + except Exception: + rank_configs = "unavailable" + raise PCIeSelectedRecordExchangeInitializationError( + "copy-engine selected-record configuration mismatch across ranks " + f"(fields={_CONFIG_FIELDS}, values={rank_configs})" + ) + + +class PCIeSelectedRecordCopyExchange: + """Ordered pack/copy-engine/unpack exchange for selected byte records. + + ``local_indices_by_destination`` is destination-major. Its trailing + dimensions flatten to the output record order. Each entry is either a + non-negative index into ``records`` or negative when this source does not + own that selected record. Across all ranks, exactly one source must own + each destination/output position. Every rank must invoke the channel in + the same order with the same selected-record count. + + Records are compacted per destination. A balanced primary packet is moved + with the CUDA copy engine; owner skew beyond that packet remains in the + source IPC slab and is read during unpack. Both exchange modes publish a + release after unpack and defer the matching wait until the next staging + reuse. Keeping that protocol identical makes separately captured normal and + layered CUDA graphs safe to alternate. ``exchange_layers`` packs exactly + three record planes into one packet. The layered packet layout overlays the + same IPC slab with a smaller ``layered_max_records`` capacity; callers must + use the single-layer fallback above that bound. Configurations whose slab + cannot hold even one layered record expose a zero layered capacity while + retaining the normal exchange. + """ + + def __init__( + self, + *, + rank: int, + world_size: int, + device: torch.device | int | str, + peer_slab_ptrs: Sequence[int], + max_records: int, + record_bytes: int, + primary_capacity: Optional[int] = None, + layered_max_records: Optional[int] = None, + process_group: Optional[ProcessGroup] = None, + ipc: Optional[CudaRTLibrary] = None, + owned_buffer: Optional[_OwnedSharedBuffer] = None, + ext_module=None, + dma_ext_module=None, + barrier_timeout_cycles: int = DEFAULT_BARRIER_TIMEOUT_CYCLES, + stream_affine: bool = True, + ) -> None: + layout = _copy_exchange_layout( + world_size=world_size, + max_records=max_records, + record_bytes=record_bytes, + primary_capacity=primary_capacity, + ) + layered_max_records, layered_layout = _layered_layout_for_slab( + world_size=world_size, + max_records=max_records, + record_bytes=record_bytes, + slab_bytes=layout.slab_bytes, + layered_max_records=layered_max_records, + fallback_layout=layout, + ) + if not 0 <= int(rank) < int(world_size): + raise ValueError(f"invalid rank {rank} for world size {world_size}") + if len(peer_slab_ptrs) != int(world_size): + raise ValueError("peer_slab_ptrs must match world_size") + if int(barrier_timeout_cycles) <= 0: + raise ValueError("barrier_timeout_cycles must be positive") + + self.rank = int(rank) + self.world_size = int(world_size) + self.device = _normalize_device(device) + self.process_group = process_group + self.max_records = int(max_records) + self.record_bytes = int(record_bytes) + self.layered_max_records = layered_max_records + self.primary_capacity = layout.primary_capacity + self.overflow_capacity = layout.overflow_capacity + self.barrier_timeout_cycles = int(barrier_timeout_cycles) + self._layout = layout + self._layer_layout = layered_layout + self._ipc = ipc + self._owned_buffer = owned_buffer + self._ext = ext_module or _load_extension() + self._dma_ext = dma_ext_module or _pcie_dma._load_extension() + self._stream_affine = bool(stream_affine) + self._owner_stream_key: Optional[int] = None + self._closed = False + self._deferred_release_pending = False + + slab_ptrs = tuple(int(pointer) for pointer in peer_slab_ptrs) + self._slab_ptrs = slab_ptrs + self._local_slab_ptr = slab_ptrs[self.rank] + self._peer_overflow_ptrs = torch.tensor( + [ + [ + slab_ptrs[source] + + layout.staging_overflow_base + + destination * layout.overflow_stride + for source in range(self.world_size) + ] + for destination in range(self.world_size) + ], + dtype=torch.int64, + device=self.device, + ) + self._peer_layer_overflow_ptrs = torch.tensor( + [ + [ + slab_ptrs[source] + + layered_layout.staging_overflow_base + + destination * layered_layout.overflow_stride + for source in range(self.world_size) + ] + for destination in range(self.world_size) + ], + dtype=torch.int64, + device=self.device, + ) + self._barrier_publish_ptrs = torch.tensor( + [ + [ + slab_ptrs[destination] + + (phase * self.world_size + self.rank) * FLAG_STRIDE + for destination in range(self.world_size) + ] + for phase in range(2) + ], + dtype=torch.int64, + device=self.device, + ) + self._barrier_wait_ptrs = torch.tensor( + [ + [ + slab_ptrs[self.rank] + + (phase * self.world_size + source) * FLAG_STRIDE + for source in range(self.world_size) + ] + for phase in range(2) + ], + dtype=torch.int64, + device=self.device, + ) + self._send_counters = torch.zeros( + (2, self.world_size), + dtype=torch.int32, + device=self.device, + ) + self._wait_counters = torch.zeros_like(self._send_counters) + + @classmethod + def from_process_group( + cls, + *, + process_group: ProcessGroup, + device: torch.device | int | str, + max_records: int, + record_bytes: int, + barrier_timeout_cycles: int = DEFAULT_BARRIER_TIMEOUT_CYCLES, + ext_module=None, + dma_ext_module=None, + stream_affine: bool = True, + primary_capacity: Optional[int] = None, + layered_max_records: Optional[int] = None, + ) -> "PCIeSelectedRecordCopyExchange": + rank = dist.get_rank(group=process_group) + world_size = dist.get_world_size(group=process_group) + device_obj = _normalize_device(device) + if device_obj.type != "cuda": + raise PCIeSelectedRecordExchangeInitializationError( + "copy-engine selected-record exchange requires a CUDA device" + ) + device_index = ( + torch.cuda.current_device() + if device_obj.index is None + else int(device_obj.index) + ) + device_obj = torch.device("cuda", device_index) + + try: + backend = dist.get_backend(group=process_group) + except TypeError: + backend = dist.get_backend(process_group) + if "nccl" not in str(backend).lower(): + raise PCIeSelectedRecordExchangeInitializationError( + "copy-engine selected-record exchange requires an NCCL process " + f"group, got {backend}" + ) + + status = torch.empty(1, dtype=torch.int32, device=device_obj) + layout: Optional[_CopyExchangeLayout] = None + layered_layout: Optional[_CopyExchangeLayout] = None + resolved_layered_max_records: Optional[int] = None + local_error: Optional[Exception] = None + try: + layout = _copy_exchange_layout( + world_size=world_size, + max_records=max_records, + record_bytes=record_bytes, + primary_capacity=primary_capacity, + ) + resolved_layered_max_records, layered_layout = _layered_layout_for_slab( + world_size=world_size, + max_records=max_records, + record_bytes=record_bytes, + slab_bytes=layout.slab_bytes, + layered_max_records=layered_max_records, + fallback_layout=layout, + ) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + raise PCIeSelectedRecordExchangeInitializationError( + "copy-engine selected-record configuration is invalid on at " + "least one rank" + ) from local_error + assert layout is not None + assert layered_layout is not None + assert resolved_layered_max_records is not None + _validate_rank_configuration( + process_group=process_group, + device=device_obj, + status=status, + world_size=world_size, + max_records=max_records, + record_bytes=record_bytes, + layout=layout, + layered_max_records=resolved_layered_max_records, + layered_layout=layered_layout, + ) + + ext = None + dma_ext = None + local_error = None + try: + ext = ext_module or _load_extension() + dma_ext = dma_ext_module or _pcie_dma._load_extension() + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + raise PCIeSelectedRecordExchangeInitializationError( + "copy-engine selected-record CUDA extensions failed to load on " + "at least one rank" + ) from local_error + assert ext is not None + assert dma_ext is not None + + ipc: Optional[CudaRTLibrary] = None + local_error = None + try: + ipc = CudaRTLibrary() + ipc.cudaSetDevice(device_index) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + raise PCIeSelectedRecordExchangeInitializationError( + "copy-engine selected-record CUDA IPC setup failed on at least one rank" + ) from local_error + assert ipc is not None + + shared = _allocate_shared_buffer_rank_consistent( + process_group=process_group, + rank=rank, + world_size=world_size, + device=device_obj, + size_in_bytes=layout.slab_bytes, + ipc=ipc, + status=status, + ) + exchange: Optional[PCIeSelectedRecordCopyExchange] = None + local_error = None + try: + exchange = cls( + rank=rank, + world_size=world_size, + device=device_obj, + peer_slab_ptrs=shared.peer_ptrs, + max_records=max_records, + record_bytes=record_bytes, + primary_capacity=layout.primary_capacity, + layered_max_records=resolved_layered_max_records, + process_group=process_group, + ipc=ipc, + owned_buffer=None, + ext_module=ext, + dma_ext_module=dma_ext, + barrier_timeout_cycles=barrier_timeout_cycles, + stream_affine=stream_affine, + ) + except Exception as exc: + local_error = exc + if not _all_ranks_succeeded(status, local_error is None, process_group): + if exchange is not None: + exchange._closed = True + _free_shared_buffer(ipc, shared) + raise PCIeSelectedRecordExchangeInitializationError( + "copy-engine selected-record runtime initialization failed on at " + "least one rank" + ) from local_error + assert exchange is not None + exchange._owned_buffer = shared + return exchange + + def _bind_stream_key(self, stream_key: Optional[int]) -> None: + if not self._stream_affine or stream_key is None: + return + if self._owner_stream_key is None: + self._owner_stream_key = int(stream_key) + return + if self._owner_stream_key != int(stream_key): + raise RuntimeError( + "PCIe selected-record copy channels are stream-affine; create a " + "separate channel for each CUDA stream" + ) + + def _check_stream(self) -> None: + if self.device.type != "cuda": + return + stream_key = _current_stream_key(self.device) + if ( + _is_current_stream_capturing(self.device) + and self._stream_affine + and self._owner_stream_key is None + ): + raise RuntimeError( + "PCIe selected-record copy channels must be bound to their CUDA " + "stream before first-use graph capture" + ) + self._bind_stream_key(stream_key) + + def _validate( + self, + records: torch.Tensor, + local_indices_by_destination: torch.Tensor, + out: torch.Tensor, + ) -> int: + if self._closed: + raise RuntimeError("PCIeSelectedRecordCopyExchange is closed") + if records.device != self.device: + raise ValueError("records must be on the exchange device") + if local_indices_by_destination.device != self.device: + raise ValueError("local indices must be on the exchange device") + if out.device != self.device: + raise ValueError("output must be on the exchange device") + if records.dtype != torch.uint8 or not records.is_contiguous(): + raise ValueError("records must be contiguous uint8") + if records.ndim < 2 or records.shape[-1] != self.record_bytes: + raise ValueError( + f"records must end in the configured {self.record_bytes}-byte width" + ) + if local_indices_by_destination.dtype not in (torch.int32, torch.int64): + raise ValueError("local indices must be int32 or int64") + if not local_indices_by_destination.is_contiguous(): + raise ValueError("local indices must be contiguous") + if ( + local_indices_by_destination.ndim < 2 + or local_indices_by_destination.shape[0] != self.world_size + ): + raise ValueError( + "local indices must have shape [world_size, ...selected records]" + ) + active_records = local_indices_by_destination.numel() // self.world_size + if active_records > self.max_records: + raise ValueError( + f"selected record count {active_records} exceeds capacity " + f"{self.max_records}" + ) + if active_records > _MAX_INT64 // self.record_bytes: + raise ValueError("active selected-record payload exceeds int64 capacity") + if out.dtype != torch.uint8 or not out.is_contiguous(): + raise ValueError("output must be contiguous uint8") + if out.ndim < 2 or out.shape[-1] != self.record_bytes: + raise ValueError( + f"output must end in the configured {self.record_bytes}-byte width" + ) + if out.numel() != active_records * self.record_bytes: + raise ValueError( + "output must contain exactly one record per selected position" + ) + return active_records + + def _active_primary_capacity(self, active_records: int) -> int: + if active_records == 0: + return 0 + scaled = ( + self.primary_capacity * active_records + self.max_records - 1 + ) // self.max_records + active_primary = min(scaled, self.primary_capacity, active_records) + active_overflow = active_records - active_primary + if active_overflow > self.overflow_capacity: + raise ValueError("active selected-record overflow exceeds channel capacity") + return active_primary + + def _active_layer_primary_capacity(self, active_records: int) -> int: + if active_records == 0: + return 0 + layout = self._layer_layout + scaled = ( + layout.primary_capacity * active_records + self.layered_max_records - 1 + ) // self.layered_max_records + active_primary = min( + scaled, + layout.primary_capacity, + active_records, + ) + active_overflow = active_records - active_primary + if active_overflow > layout.overflow_capacity: + raise ValueError( + "active three-layer selected-record overflow exceeds channel capacity" + ) + return active_primary + + def _receive_primary_ptr(self, destination: int, source: int) -> int: + return ( + self._slab_ptrs[destination] + + self._layout.receive_primary_base + + source * self._layout.primary_stride + ) + + def _local_staging_primary_ptr(self, destination: int) -> int: + return ( + self._local_slab_ptr + + self._layout.staging_primary_base + + destination * self._layout.primary_stride + ) + + def _local_staging_overflow_ptr(self, destination: int) -> int: + return ( + self._local_slab_ptr + + self._layout.staging_overflow_base + + destination * self._layout.overflow_stride + ) + + def _layer_receive_primary_ptr(self, destination: int, source: int) -> int: + return ( + self._slab_ptrs[destination] + + self._layer_layout.receive_primary_base + + source * self._layer_layout.primary_stride + ) + + def _local_layer_staging_primary_ptr(self, destination: int) -> int: + return ( + self._local_slab_ptr + + self._layer_layout.staging_primary_base + + destination * self._layer_layout.primary_stride + ) + + def _local_layer_staging_overflow_ptr(self, destination: int) -> int: + return ( + self._local_slab_ptr + + self._layer_layout.staging_overflow_base + + destination * self._layer_layout.overflow_stride + ) + + def _barrier(self, phase: int) -> None: + self._ext.barrier_all_peers( + self._barrier_publish_ptrs, + self._barrier_wait_ptrs, + self._send_counters, + self._wait_counters, + int(phase), + self.barrier_timeout_cycles, + ) + + def _publish_deferred_release(self) -> None: + self._ext.publish_all_peers( + self._barrier_publish_ptrs, + self._send_counters, + 1, + ) + self._deferred_release_pending = True + + def _wait_for_deferred_release(self) -> None: + if not self._deferred_release_pending: + return + self._ext.wait_all_peers( + self._barrier_wait_ptrs, + self._wait_counters, + 1, + self.barrier_timeout_cycles, + ) + self._deferred_release_pending = False + + def _require_eager_warmup_for_capture(self) -> None: + if ( + self.device.type == "cuda" + and _is_current_stream_capturing(self.device) + and not self._deferred_release_pending + ): + raise RuntimeError( + "copy-engine selected-record exchange requires one eager " + "warmup before CUDA graph capture" + ) + + def _validate_layers( + self, + records_by_layer: Sequence[torch.Tensor], + local_indices_by_destination: torch.Tensor, + outputs_by_layer: Sequence[torch.Tensor], + ) -> int: + if len(records_by_layer) != _LAYER_PLANES: + raise ValueError(f"records_by_layer must contain {_LAYER_PLANES} tensors") + if len(outputs_by_layer) != _LAYER_PLANES: + raise ValueError(f"outputs_by_layer must contain {_LAYER_PLANES} tensors") + + active_records: Optional[int] = None + for records, out in zip(records_by_layer, outputs_by_layer, strict=True): + current = self._validate(records, local_indices_by_destination, out) + if active_records is None: + active_records = current + elif current != active_records: + raise ValueError("all layer outputs must have the same record count") + + assert active_records is not None + if active_records: + output_ranges = sorted( + ( + int(output.data_ptr()), + int(output.data_ptr()) + output.numel(), + ) + for output in outputs_by_layer + ) + if any( + output_ranges[index][1] > output_ranges[index + 1][0] + for index in range(len(output_ranges) - 1) + ): + raise ValueError("outputs_by_layer must use non-overlapping storage") + if active_records > self.layered_max_records: + raise ValueError( + f"three-layer selected record count {active_records} exceeds " + f"bounded capacity {self.layered_max_records}; use the " + "single-layer exchange fallback" + ) + record_shapes = {tuple(records.shape) for records in records_by_layer} + if len(record_shapes) != 1: + raise ValueError("all layer record tensors must have the same shape") + return active_records + + def exchange( + self, + records: torch.Tensor, + local_indices_by_destination: torch.Tensor, + out: torch.Tensor, + ) -> torch.Tensor: + """Pack owned records, copy them to peers, and reconstruct ``out``.""" + self._check_stream() + self._require_eager_warmup_for_capture() + active_records = self._validate( + records, + local_indices_by_destination, + out, + ) + self._wait_for_deferred_release() + if active_records == 0: + self._barrier(0) + out.zero_() + self._publish_deferred_release() + return out + + active_primary = self._active_primary_capacity(active_records) + primary_copy_bytes = _align_up( + self._layout.primary_records_offset + active_primary * self.record_bytes, + 16, + ) + + # Rotate both pack and CE issue order by source rank. At phase zero, + # rank R targets destination R instead of every rank targeting rank 0. + for destination_phase in range(self.world_size): + destination = (self.rank + destination_phase) % self.world_size + self._ext.pack_compact_records( + records, + local_indices_by_destination[destination], + self._local_staging_primary_ptr(destination), + self._local_staging_overflow_ptr(destination), + self.record_bytes, + active_primary, + self._layout.primary_positions_offset, + self._layout.primary_records_offset, + self._layout.overflow_positions_offset, + self._layout.overflow_records_offset, + ) + + for destination_phase in range(self.world_size): + destination = (self.rank + destination_phase) % self.world_size + self._dma_ext.dma_copy( + self._receive_primary_ptr(destination, self.rank), + self._local_staging_primary_ptr(destination), + primary_copy_bytes, + ) + + self._barrier(0) + out.zero_() + self._ext.unpack_compact_records( + self._local_slab_ptr + self._layout.receive_primary_base, + self._layout.primary_stride, + self._peer_overflow_ptrs[self.rank], + out, + active_records, + self.record_bytes, + active_primary, + self._layout.primary_positions_offset, + self._layout.primary_records_offset, + self._layout.overflow_positions_offset, + self._layout.overflow_records_offset, + ) + self._publish_deferred_release() + return out + + def exchange_layers( + self, + records_by_layer: Sequence[torch.Tensor], + local_indices_by_destination: torch.Tensor, + outputs_by_layer: Sequence[torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Exchange three layer planes with one shared destination-position table. + + Each selected slot carries three native records in one packet. The + arrival barrier protects packet visibility. Unpack then publishes a + release generation without waiting; the matching wait is issued only + before this channel reuses staging for its next exchange. + """ + self._check_stream() + self._require_eager_warmup_for_capture() + active_records = self._validate_layers( + records_by_layer, + local_indices_by_destination, + outputs_by_layer, + ) + outputs = ( + outputs_by_layer[0], + outputs_by_layer[1], + outputs_by_layer[2], + ) + self._wait_for_deferred_release() + + if active_records == 0: + self._barrier(0) + for output in outputs: + output.zero_() + self._publish_deferred_release() + return outputs + + active_primary = self._active_layer_primary_capacity(active_records) + layout = self._layer_layout + packet_record_bytes = self.record_bytes * _LAYER_PLANES + primary_copy_bytes = _align_up( + layout.primary_records_offset + active_primary * packet_record_bytes, + 16, + ) + + for destination_phase in range(self.world_size): + destination = (self.rank + destination_phase) % self.world_size + self._ext.pack_compact_record_layers( + records_by_layer[0], + records_by_layer[1], + records_by_layer[2], + local_indices_by_destination[destination], + self._local_layer_staging_primary_ptr(destination), + self._local_layer_staging_overflow_ptr(destination), + self.record_bytes, + active_primary, + layout.primary_positions_offset, + layout.primary_records_offset, + layout.overflow_positions_offset, + layout.overflow_records_offset, + ) + + for destination_phase in range(self.world_size): + destination = (self.rank + destination_phase) % self.world_size + self._dma_ext.dma_copy( + self._layer_receive_primary_ptr(destination, self.rank), + self._local_layer_staging_primary_ptr(destination), + primary_copy_bytes, + ) + + self._barrier(0) + for output in outputs: + output.zero_() + self._ext.unpack_compact_record_layers( + self._local_slab_ptr + layout.receive_primary_base, + layout.primary_stride, + self._peer_layer_overflow_ptrs[self.rank], + outputs[0], + outputs[1], + outputs[2], + active_records, + self.record_bytes, + active_primary, + layout.primary_positions_offset, + layout.primary_records_offset, + layout.overflow_positions_offset, + layout.overflow_records_offset, + ) + self._publish_deferred_release() + return outputs + + def _finish_deferred_release_for_close(self) -> None: + if not self._deferred_release_pending: + return + if self.device.type == "cuda": + device_index = ( + torch.cuda.current_device() + if self.device.index is None + else int(self.device.index) + ) + if self._ipc is not None: + self._ipc.cudaSetDevice(device_index) + torch.cuda.synchronize(self.device) + self._wait_for_deferred_release() + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + + def _release_owned_buffer(self, *, synchronize: bool) -> None: + if self._owned_buffer is None or self._ipc is None: + return + if self.device.type == "cuda": + device_index = ( + torch.cuda.current_device() + if self.device.index is None + else int(self.device.index) + ) + self._ipc.cudaSetDevice(device_index) + if synchronize: + with suppress(Exception): + torch.cuda.synchronize(self.device) + _free_shared_buffer(self._ipc, self._owned_buffer) + self._owned_buffer = None + + def close(self) -> None: + if self._closed: + return + self._finish_deferred_release_for_close() + self._closed = True + self._release_owned_buffer(synchronize=True) + + def __del__(self) -> None: + with suppress(Exception): + self.close() + + +__all__ = [ + "PCIeSelectedRecordCopyExchange", + "PCIeSelectedRecordExchangeInitializationError", +] diff --git a/tests/comm/test_pcie_selected_records_ce.py b/tests/comm/test_pcie_selected_records_ce.py new file mode 100644 index 00000000..b4328f12 --- /dev/null +++ b/tests/comm/test_pcie_selected_records_ce.py @@ -0,0 +1,872 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +from sparkinfer.comm.pcie.pcie_dma import FLAG_STRIDE +from sparkinfer.comm.pcie.pcie_oneshot import _OwnedSharedBuffer +from sparkinfer.comm.pcie.pcie_selected_records import ( + PCIeSelectedRecordExchangeInitializationError, +) +from sparkinfer.comm.pcie.pcie_selected_records_ce import ( + _CONFIG_FIELDS, + PCIeSelectedRecordCopyExchange, + _copy_exchange_layout, + _layered_layout_for_slab, + _largest_layered_capacity_for_slab, + _validate_rank_configuration, +) + + +class _FakeExt: + def __init__(self, rank: int) -> None: + self.rank = rank + self.pack_calls = [] + self.layer_pack_calls = [] + self.barrier_calls = [] + self.unpack_calls = [] + self.layer_unpack_calls = [] + self.publish_calls = [] + self.wait_calls = [] + self.operations = [] + self._packets = {} + + def pack_compact_records( + self, + records, + local_indices, + primary_ptr, + overflow_ptr, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + ): + call = { + "records": records.clone(), + "indices": local_indices.clone(), + "primary_ptr": primary_ptr, + "overflow_ptr": overflow_ptr, + "record_bytes": record_bytes, + "primary_capacity": primary_capacity, + "primary_positions_offset": primary_positions_offset, + "primary_records_offset": primary_records_offset, + "overflow_positions_offset": overflow_positions_offset, + "overflow_records_offset": overflow_records_offset, + } + self.pack_calls.append(call) + self.operations.append(("pack", primary_ptr)) + self._packets[primary_ptr] = call + + def pack_compact_record_layers( + self, + records0, + records1, + records2, + local_indices, + primary_ptr, + overflow_ptr, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + ): + call = { + "records": (records0.clone(), records1.clone(), records2.clone()), + "indices": local_indices.clone(), + "primary_ptr": primary_ptr, + "overflow_ptr": overflow_ptr, + "record_bytes": record_bytes, + "primary_capacity": primary_capacity, + "primary_positions_offset": primary_positions_offset, + "primary_records_offset": primary_records_offset, + "overflow_positions_offset": overflow_positions_offset, + "overflow_records_offset": overflow_records_offset, + } + self.layer_pack_calls.append(call) + self.operations.append(("pack_layers", primary_ptr)) + self._packets[primary_ptr] = call + + def barrier_all_peers( + self, + publish_ptrs, + wait_ptrs, + send_counters, + wait_counters, + phase, + timeout_cycles, + ): + self.barrier_calls.append( + { + "phase": phase, + "timeout_cycles": timeout_cycles, + "publish_ptrs": publish_ptrs.clone(), + "wait_ptrs": wait_ptrs.clone(), + } + ) + self.operations.append(("barrier", phase)) + send_counters[phase].add_(1) + wait_counters[phase].add_(1) + + def publish_all_peers(self, publish_ptrs, send_counters, phase): + self.publish_calls.append( + {"phase": phase, "publish_ptrs": publish_ptrs.clone()} + ) + self.operations.append(("publish", phase)) + send_counters[phase].add_(1) + + def wait_all_peers( + self, + wait_ptrs, + wait_counters, + phase, + timeout_cycles, + ): + self.wait_calls.append( + { + "phase": phase, + "timeout_cycles": timeout_cycles, + "wait_ptrs": wait_ptrs.clone(), + } + ) + self.operations.append(("wait", phase)) + wait_counters[phase].add_(1) + + def unpack_compact_records( + self, + primary_base_ptr, + primary_stride, + peer_overflow_ptrs, + out, + selected_records, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + ): + self.unpack_calls.append( + { + "primary_base_ptr": primary_base_ptr, + "primary_stride": primary_stride, + "peer_overflow_ptrs": peer_overflow_ptrs.clone(), + "selected_records": selected_records, + "record_bytes": record_bytes, + "primary_capacity": primary_capacity, + } + ) + self.operations.append(("unpack", selected_records)) + # A single-process fake can reconstruct the records this source owns + # for its own destination. Distributed GPU tests cover peer packets. + # Rank rotation always packs this rank's own destination at phase 0. + local_packet = self.pack_calls[0] + flat_records = local_packet["records"].reshape(-1, record_bytes) + flat_indices = local_packet["indices"].reshape(-1) + flat_out = out.reshape(-1, record_bytes) + for selected, local_index in enumerate(flat_indices.tolist()): + if local_index >= 0: + flat_out[selected].copy_(flat_records[local_index]) + + def unpack_compact_record_layers( + self, + primary_base_ptr, + primary_stride, + peer_overflow_ptrs, + out0, + out1, + out2, + selected_records, + record_bytes, + primary_capacity, + primary_positions_offset, + primary_records_offset, + overflow_positions_offset, + overflow_records_offset, + ): + self.layer_unpack_calls.append( + { + "primary_base_ptr": primary_base_ptr, + "primary_stride": primary_stride, + "peer_overflow_ptrs": peer_overflow_ptrs.clone(), + "selected_records": selected_records, + "record_bytes": record_bytes, + "primary_capacity": primary_capacity, + } + ) + self.operations.append(("unpack_layers", selected_records)) + local_packet = self.layer_pack_calls[0] + flat_indices = local_packet["indices"].reshape(-1) + for records, out in zip( + local_packet["records"], + (out0, out1, out2), + strict=True, + ): + flat_records = records.reshape(-1, record_bytes) + flat_out = out.reshape(-1, record_bytes) + for selected, local_index in enumerate(flat_indices.tolist()): + if local_index >= 0: + flat_out[selected].copy_(flat_records[local_index]) + + +class _FakeDma: + def __init__(self) -> None: + self.calls = [] + + def dma_copy(self, destination, source, size): + self.calls.append((destination, source, size)) + + +def _make_runtime( + *, + rank: int = 1, + world_size: int = 3, + max_records: int = 5, + record_bytes: int = 37, + primary_capacity: int | None = None, + layered_max_records: int | None = None, +): + ext = _FakeExt(rank) + dma = _FakeDma() + slab_ptrs = tuple(1_000_000 * (peer + 1) for peer in range(world_size)) + runtime = PCIeSelectedRecordCopyExchange( + rank=rank, + world_size=world_size, + device=torch.device("cpu"), + peer_slab_ptrs=slab_ptrs, + max_records=max_records, + record_bytes=record_bytes, + primary_capacity=primary_capacity, + layered_max_records=layered_max_records, + ext_module=ext, + dma_ext_module=dma, + ) + return runtime, ext, dma, slab_ptrs + + +def test_layout_supports_runtime_world_sizes_balanced_packets_and_int64_offsets(): + layout = _copy_exchange_layout( + world_size=7, + max_records=32_769, + record_bytes=65_537, + ) + + assert layout.primary_capacity == (32_769 + 6) // 7 + assert layout.overflow_capacity == 32_769 - layout.primary_capacity + assert layout.flags_bytes >= 2 * 7 * FLAG_STRIDE + assert layout.flags_bytes % 256 == 0 + assert layout.primary_records_offset % 16 == 0 + assert layout.overflow_records_offset % 16 == 0 + assert layout.primary_stride % 256 == 0 + assert layout.overflow_stride % 256 == 0 + assert layout.slab_bytes > 2**31 + + base = _copy_exchange_layout( + world_size=4, + max_records=65_536, + record_bytes=432, + ) + naive_layered = _copy_exchange_layout( + world_size=4, + max_records=65_536, + record_bytes=1_296, + ) + layered_max_records, layered = _largest_layered_capacity_for_slab( + world_size=4, + max_records=65_536, + record_bytes=432, + slab_bytes=base.slab_bytes, + ) + assert base.slab_bytes == 144_182_272 + assert naive_layered.slab_bytes == base.slab_bytes + 270 * 1024**2 + assert naive_layered.slab_bytes - base.flags_bytes == 427_296_768 + assert layered_max_records == 22_112 + assert layered_max_records >= 2 * 8_192 + assert layered_max_records < 3 * 8_192 + assert layered.slab_bytes == 144_173_056 + assert layered.slab_bytes <= base.slab_bytes + + +@pytest.mark.parametrize("world_size", range(2, 9)) +def test_layered_layout_is_deterministic_and_bounded_for_dcp2_through_dcp8( + world_size, +): + base = _copy_exchange_layout( + world_size=world_size, + max_records=65_536, + record_bytes=432, + ) + layered_max_records, layered = _largest_layered_capacity_for_slab( + world_size=world_size, + max_records=65_536, + record_bytes=432, + slab_bytes=base.slab_bytes, + ) + + assert 1 <= layered_max_records <= 65_536 + assert layered.slab_bytes <= base.slab_bytes + + +def test_layered_exchange_reuses_the_exact_production_slab_allocation(): + runtime, _, _, _ = _make_runtime( + rank=0, + world_size=4, + max_records=65_536, + record_bytes=432, + ) + + legacy_layout = _copy_exchange_layout( + world_size=4, + max_records=65_536, + record_bytes=432, + ) + assert runtime._layout.slab_bytes == legacy_layout.slab_bytes == 144_182_272 + assert runtime._layout.slab_bytes < 150 * 1024**2 + assert runtime.layered_max_records == 22_112 + assert runtime._layer_layout.slab_bytes == 144_173_056 + assert runtime._layer_layout.slab_bytes <= runtime._layout.slab_bytes + + +def test_single_layer_exchange_remains_available_when_layered_packet_cannot_fit(): + base = _copy_exchange_layout( + world_size=4, + max_records=1, + record_bytes=65_536, + primary_capacity=1, + ) + + layered_max_records, layered = _layered_layout_for_slab( + world_size=4, + max_records=1, + record_bytes=65_536, + slab_bytes=base.slab_bytes, + layered_max_records=None, + fallback_layout=base, + ) + + assert layered_max_records == 0 + assert layered is base + + +@pytest.mark.parametrize("world_size", [1, 33]) +def test_layout_rejects_unsupported_world_sizes(world_size): + with pytest.raises(ValueError, match="world_size"): + _copy_exchange_layout( + world_size=world_size, + max_records=4, + record_bytes=37, + ) + + +def test_layout_rejects_invalid_and_int64_overflowing_capacities(): + with pytest.raises(ValueError, match="max_records"): + _copy_exchange_layout(world_size=2, max_records=0, record_bytes=1) + with pytest.raises(ValueError, match="record_bytes"): + _copy_exchange_layout(world_size=2, max_records=1, record_bytes=0) + with pytest.raises(ValueError, match="primary_capacity"): + _copy_exchange_layout( + world_size=2, + max_records=4, + record_bytes=1, + primary_capacity=5, + ) + with pytest.raises(ValueError, match="int64"): + _copy_exchange_layout( + world_size=2, + max_records=2**62, + record_bytes=4, + ) + with pytest.raises(ValueError, match="max_records"): + _largest_layered_capacity_for_slab( + world_size=2, + max_records=0, + record_bytes=1, + slab_bytes=1, + ) + + +def test_exchange_uses_rank_rotated_pack_and_dma_order_with_odd_width_records(): + runtime, ext, dma, slab_ptrs = _make_runtime() + records = torch.arange(7 * 37, dtype=torch.uint8).reshape(7, 37) + local_indices = torch.full((3, 1, 4), -1, dtype=torch.int64) + local_indices[1, 0, :3] = torch.tensor([6, 2, 4], dtype=torch.int64) + out = torch.full((1, 4, 37), 255, dtype=torch.uint8) + + returned = runtime.exchange(records, local_indices, out) + + assert returned is out + assert torch.equal(out[0, 0], records[6]) + assert torch.equal(out[0, 1], records[2]) + assert torch.equal(out[0, 2], records[4]) + assert torch.equal(out[0, 3], torch.zeros(37, dtype=torch.uint8)) + layout = runtime._layout + expected_destinations = [1, 2, 0] + assert [ + (call["primary_ptr"] - layout.staging_primary_base - slab_ptrs[1]) + // layout.primary_stride + for call in ext.pack_calls + ] == expected_destinations + assert [ + (destination - layout.receive_primary_base - slab_ptrs[dest]) + // layout.primary_stride + for (destination, _, _), dest in zip( + dma.calls, expected_destinations, strict=True + ) + ] == [1, 1, 1] + assert [ + (source - layout.staging_primary_base - slab_ptrs[1]) // layout.primary_stride + for _, source, _ in dma.calls + ] == expected_destinations + assert [call["phase"] for call in ext.barrier_calls] == [0] + assert [call["phase"] for call in ext.publish_calls] == [1] + assert ext.wait_calls == [] + assert torch.equal( + runtime._send_counters, + torch.ones((2, 3), dtype=torch.int32), + ) + assert torch.equal( + runtime._wait_counters, + torch.stack( + ( + torch.ones(3, dtype=torch.int32), + torch.zeros(3, dtype=torch.int32), + ) + ), + ) + runtime.close() + assert [call["phase"] for call in ext.wait_calls] == [1] + + +def test_exchange_preserves_persistent_state_and_accepts_int32_indices(): + runtime, ext, dma, _ = _make_runtime(record_bytes=32) + records = torch.arange(6 * 32, dtype=torch.uint8).reshape(6, 32) + local_indices = torch.full((3, 2), -1, dtype=torch.int32) + local_indices[1] = torch.tensor([5, 1], dtype=torch.int32) + out = torch.zeros(2, 32, dtype=torch.uint8) + + runtime.exchange(records, local_indices, out) + runtime.exchange(records.flip(0).contiguous(), local_indices, out) + + assert len(ext.pack_calls) == 6 + assert len(dma.calls) == 6 + assert [call["phase"] for call in ext.barrier_calls] == [0, 0] + assert [call["phase"] for call in ext.publish_calls] == [1, 1] + assert [call["phase"] for call in ext.wait_calls] == [1] + assert torch.equal( + runtime._send_counters, + torch.full((2, 3), 2, dtype=torch.int32), + ) + runtime.close() + assert [call["phase"] for call in ext.wait_calls] == [1, 1] + + +def test_exchange_layers_uses_one_packet_per_destination_and_deferred_release(): + runtime, ext, dma, _ = _make_runtime(record_bytes=32) + base = torch.arange(7 * 32, dtype=torch.int64).remainder(256).to(torch.uint8) + records = tuple( + (base.reshape(7, 32) + layer * 41).to(torch.uint8) for layer in range(3) + ) + local_indices = torch.full((3, 1, 4), -1, dtype=torch.int64) + local_indices[1, 0, :3] = torch.tensor([6, 2, 4], dtype=torch.int64) + outputs = tuple(torch.full((1, 4, 32), 255, dtype=torch.uint8) for _ in range(3)) + + returned = runtime.exchange_layers(records, local_indices, outputs) + + assert all( + actual is expected for actual, expected in zip(returned, outputs, strict=True) + ) + for layer, output in enumerate(outputs): + assert torch.equal(output[0, 0], records[layer][6]) + assert torch.equal(output[0, 1], records[layer][2]) + assert torch.equal(output[0, 2], records[layer][4]) + assert torch.equal(output[0, 3], torch.zeros(32, dtype=torch.uint8)) + assert len(ext.layer_pack_calls) == runtime.world_size + assert len(ext.layer_unpack_calls) == 1 + assert len(dma.calls) == runtime.world_size + assert [call["phase"] for call in ext.barrier_calls] == [0] + assert [call["phase"] for call in ext.publish_calls] == [1] + assert ext.wait_calls == [] + expected_copy_bytes = ( + ( + runtime._layer_layout.primary_records_offset + + runtime._active_layer_primary_capacity(4) * 3 * runtime.record_bytes + + 15 + ) + // 16 + * 16 + ) + assert {size for _, _, size in dma.calls} == {expected_copy_bytes} + assert runtime._deferred_release_pending + + operation_count = len(ext.operations) + runtime.exchange_layers(records, local_indices, outputs) + second_operations = ext.operations[operation_count:] + assert second_operations[0] == ("wait", 1) + assert second_operations[-1] == ("publish", 1) + assert [call["phase"] for call in ext.barrier_calls] == [0, 0] + assert [call["phase"] for call in ext.wait_calls] == [1] + assert [call["phase"] for call in ext.publish_calls] == [1, 1] + + runtime.close() + assert [call["phase"] for call in ext.wait_calls] == [1, 1] + assert not runtime._deferred_release_pending + + +def test_exchange_layers_empty_selection_keeps_order_and_defers_release(): + runtime, ext, dma, _ = _make_runtime(max_records=2) + records = tuple(torch.empty((0, 37), dtype=torch.uint8) for _ in range(3)) + indices = torch.empty((3, 0), dtype=torch.int64) + outputs = tuple(torch.empty((0, 37), dtype=torch.uint8) for _ in range(3)) + + returned = runtime.exchange_layers(records, indices, outputs) + + assert all( + actual is expected for actual, expected in zip(returned, outputs, strict=True) + ) + assert ext.layer_pack_calls == [] + assert dma.calls == [] + assert [call["phase"] for call in ext.barrier_calls] == [0] + assert [call["phase"] for call in ext.publish_calls] == [1] + runtime.close() + assert [call["phase"] for call in ext.wait_calls] == [1] + + +def test_single_layer_exchange_empty_selection_keeps_order_and_defers_release(): + runtime, ext, dma, _ = _make_runtime(max_records=2) + records = torch.empty((0, 37), dtype=torch.uint8) + indices = torch.empty((3, 0), dtype=torch.int64) + output = torch.empty((0, 37), dtype=torch.uint8) + + returned = runtime.exchange(records, indices, output) + + assert returned is output + assert ext.pack_calls == [] + assert dma.calls == [] + assert [call["phase"] for call in ext.barrier_calls] == [0] + assert [call["phase"] for call in ext.publish_calls] == [1] + assert runtime._deferred_release_pending + runtime.close() + assert [call["phase"] for call in ext.wait_calls] == [1] + + +def test_single_layer_exchange_drains_layered_release_before_reusing_slab(): + runtime, ext, _, _ = _make_runtime(record_bytes=32) + records = tuple(torch.zeros(4, 32, dtype=torch.uint8) for _ in range(3)) + indices = torch.full((3, 2), -1, dtype=torch.int32) + outputs = tuple(torch.zeros(2, 32, dtype=torch.uint8) for _ in range(3)) + runtime.exchange_layers(records, indices, outputs) + + operation_count = len(ext.operations) + runtime.exchange(records[0], indices, outputs[0]) + + operations = ext.operations[operation_count:] + assert operations[0] == ("wait", 1) + assert [operation for operation in operations if operation[0] == "barrier"] == [ + ("barrier", 0), + ] + assert operations[-1] == ("publish", 1) + assert runtime._deferred_release_pending + runtime.close() + assert not runtime._deferred_release_pending + + +def test_layered_and_single_exchanges_keep_one_balanced_release_chain(): + runtime, ext, _, _ = _make_runtime(record_bytes=32) + records = tuple(torch.zeros(4, 32, dtype=torch.uint8) for _ in range(3)) + indices = torch.full((3, 2), -1, dtype=torch.int32) + outputs = tuple(torch.zeros(2, 32, dtype=torch.uint8) for _ in range(3)) + + runtime.exchange_layers(records, indices, outputs) + operation_count = len(ext.operations) + calls = ( + lambda: runtime.exchange(records[0], indices, outputs[0]), + lambda: runtime.exchange(records[0], indices, outputs[0]), + lambda: runtime.exchange_layers(records, indices, outputs), + ) + for call in calls: + call() + operations = ext.operations[operation_count:] + assert operations[0] == ("wait", 1) + assert operations[-1] == ("publish", 1) + assert [op for op in operations if op[0] == "barrier"] == [("barrier", 0)] + operation_count = len(ext.operations) + + assert len(ext.publish_calls) == 4 + assert len(ext.wait_calls) == 3 + runtime.close() + assert len(ext.wait_calls) == 4 + + +def test_exchange_layers_rejects_capacity_invalid_or_aliased_planes(): + records = tuple(torch.zeros(3, 37, dtype=torch.uint8) for _ in range(3)) + indices = torch.full((3, 2), -1, dtype=torch.int32) + outputs = tuple(torch.zeros(2, 37, dtype=torch.uint8) for _ in range(3)) + + layered, _, _, _ = _make_runtime(max_records=2) + with pytest.raises(ValueError, match="records_by_layer"): + layered.exchange_layers(records[:2], indices, outputs) + with pytest.raises(ValueError, match="outputs_by_layer"): + layered.exchange_layers(records, indices, outputs[:2]) + with pytest.raises(ValueError, match="non-overlapping storage"): + layered.exchange_layers(records, indices, (outputs[0], outputs[0], outputs[2])) + overlapping = torch.zeros(3, 37, dtype=torch.uint8) + with pytest.raises(ValueError, match="non-overlapping storage"): + layered.exchange_layers( + records, + indices, + (overlapping[:2], overlapping[1:], outputs[2]), + ) + + +def test_exchange_layers_rejects_counts_above_same_slab_capacity(): + runtime, _, _, _ = _make_runtime( + max_records=5, + layered_max_records=2, + ) + count = runtime.layered_max_records + 1 + records = tuple(torch.zeros(count, 37, dtype=torch.uint8) for _ in range(3)) + indices = torch.full((3, count), -1, dtype=torch.int32) + outputs = tuple(torch.zeros(count, 37, dtype=torch.uint8) for _ in range(3)) + + with pytest.raises(ValueError, match="single-layer exchange fallback"): + runtime.exchange_layers(records, indices, outputs) + + +def test_explicit_layered_capacity_cannot_grow_the_base_slab(): + with pytest.raises(ValueError, match="exceeds the existing selected-record slab"): + _make_runtime( + world_size=4, + max_records=65_536, + record_bytes=432, + layered_max_records=22_113, + ) + + +def test_explicit_full_primary_capacity_scales_without_overflow(): + runtime, _, _, _ = _make_runtime( + world_size=2, + max_records=8, + primary_capacity=8, + ) + + assert runtime._active_primary_capacity(1) == 1 + assert runtime._active_primary_capacity(5) == 5 + assert runtime._active_primary_capacity(8) == 8 + + +def test_exchange_accepts_empty_selection_but_keeps_collective_order(): + runtime, ext, dma, _ = _make_runtime(max_records=2) + records = torch.empty((0, 37), dtype=torch.uint8) + local_indices = torch.empty((3, 0), dtype=torch.int64) + out = torch.empty((0, 37), dtype=torch.uint8) + + assert runtime.exchange(records, local_indices, out) is out + assert ext.pack_calls == [] + assert dma.calls == [] + assert [call["phase"] for call in ext.barrier_calls] == [0] + assert [call["phase"] for call in ext.publish_calls] == [1] + assert runtime._deferred_release_pending + + +def test_exchange_rejects_invalid_shapes_dtypes_capacity_and_closed_state(): + runtime, _, _, _ = _make_runtime(max_records=2) + records = torch.zeros(3, 37, dtype=torch.uint8) + indices = torch.full((3, 2), -1, dtype=torch.int32) + out = torch.zeros(2, 37, dtype=torch.uint8) + + with pytest.raises(ValueError, match="configured 37-byte width"): + runtime.exchange(records[:, :-1].contiguous(), indices, out) + with pytest.raises(ValueError, match="int32 or int64"): + runtime.exchange(records, indices.to(torch.int16), out) + with pytest.raises(ValueError, match="world_size"): + runtime.exchange(records, indices[:2], out) + with pytest.raises(ValueError, match="exceeds capacity"): + runtime.exchange( + records, + torch.full((3, 3), -1, dtype=torch.int32), + torch.zeros(3, 37, dtype=torch.uint8), + ) + with pytest.raises(ValueError, match="exactly one record"): + runtime.exchange(records, indices, torch.zeros(3, 37, dtype=torch.uint8)) + + runtime.close() + with pytest.raises(RuntimeError, match="closed"): + runtime.exchange(records, indices, out) + + +def test_cpu_process_group_initialization_raises_fallback_exception(monkeypatch): + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records_ce.dist.get_rank", + lambda group=None: 0, + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records_ce.dist.get_world_size", + lambda group=None: 2, + ) + + with pytest.raises( + PCIeSelectedRecordExchangeInitializationError, + match="requires a CUDA device", + ): + PCIeSelectedRecordCopyExchange.from_process_group( + process_group=object(), + device="cpu", + max_records=2, + record_bytes=37, + ) + + +@pytest.mark.parametrize( + "mismatched_field", + ("record_bytes", "layered_max_records", "layered_slab_bytes"), +) +def test_rank_configuration_handshake_rejects_mismatches( + monkeypatch, + mismatched_field, +): + layout = _copy_exchange_layout( + world_size=2, + max_records=4, + record_bytes=37, + ) + layered_max_records, layered_layout = _largest_layered_capacity_for_slab( + world_size=2, + max_records=4, + record_bytes=37, + slab_bytes=layout.slab_bytes, + ) + field_index = _CONFIG_FIELDS.index(mismatched_field) + + def fake_all_gather_into_tensor(output, local, group=None): + rows = local.repeat(2, 1) + rows[1, field_index].add_(1) + output.copy_(rows.reshape(-1)) + + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records_ce.dist.all_gather_into_tensor", + fake_all_gather_into_tensor, + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records_ce._all_ranks_succeeded", + lambda status, local_success, process_group: local_success, + ) + + with pytest.raises( + PCIeSelectedRecordExchangeInitializationError, + match="configuration mismatch", + ): + _validate_rank_configuration( + process_group=object(), + device=torch.device("cpu"), + status=torch.empty(1, dtype=torch.int32), + world_size=2, + max_records=4, + record_bytes=37, + layout=layout, + layered_max_records=layered_max_records, + layered_layout=layered_layout, + ) + + +def test_first_use_graph_capture_requires_prebound_matching_stream(monkeypatch): + runtime, _, _, _ = _make_runtime() + runtime.device = torch.device("cuda", 0) + stream_key = [77] + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records_ce._is_current_stream_capturing", + lambda device: True, + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_selected_records_ce._current_stream_key", + lambda device: stream_key[0], + ) + + with pytest.raises(RuntimeError, match="before first-use graph capture"): + runtime._check_stream() + runtime._bind_stream_key(77) + runtime._check_stream() + stream_key[0] = 78 + with pytest.raises(RuntimeError, match="stream-affine"): + runtime._check_stream() + + +def test_cleanup_rebinds_configured_cuda_device(monkeypatch): + operations = [] + + class _CleanupIPC: + def cudaSetDevice(self, device): + operations.append(("set_device", device)) + + def cudaIpcCloseMemHandle(self, pointer): + operations.append(("close", pointer)) + + def cudaFree(self, pointer): + operations.append(("free", pointer)) + + runtime, _, _, _ = _make_runtime() + runtime.device = torch.device("cuda", 3) + runtime._ipc = _CleanupIPC() + runtime._owned_buffer = _OwnedSharedBuffer( + local_ptr=100, + peer_ptrs=(100, 200), + remote_ptrs=(200,), + ) + monkeypatch.setattr( + torch.cuda, + "synchronize", + lambda device: operations.append(("synchronize", device.index)), + ) + + runtime.close() + + assert operations == [ + ("set_device", 3), + ("synchronize", 3), + ("close", 200), + ("free", 100), + ] + + +def test_cuda_source_uses_ce_pack_unpack_rotation_and_generic_contract(): + module_path = ( + Path(__file__).parents[2] + / "sparkinfer" + / "comm" + / "pcie" + / "pcie_selected_records_ce.cu" + ) + source = module_path.read_text(encoding="utf-8") + python_source = module_path.with_suffix(".py").read_text(encoding="utf-8") + + assert "pack_compact_records_kernel" in source + assert "unpack_compact_records_kernel" in source + assert "pack_compact_record_layers_kernel" in source + assert "unpack_compact_record_layers_kernel" in source + assert "publish_all_peers_kernel" in source + assert "wait_all_peers_kernel" in source + assert "launch_pack" in source + assert "launch_pack" in source + assert "launch_pack" in source + assert "launch_pack" in source + assert "const int64_t local_record" in source + assert "record_byte_offset(local_record, record_bytes)" in source + assert "barrier_all_peers_kernel" in source + assert "cudaMemcpy" not in source + assert "cudaMalloc" not in source + assert "_pcie_dma._load_extension()" in python_source + assert "self._dma_ext.dma_copy(" in python_source + assert "(self.rank + destination_phase) % self.world_size" in python_source + assert "def exchange_layers(" in python_source + assert "self._publish_deferred_release()" in python_source + assert "self._wait_for_deferred_release()" in python_source + assert "torch.cat" not in python_source + lowered = (source + python_source).lower() + assert "topk" not in lowered + assert "glm" not in lowered + assert "mtp" not in lowered diff --git a/tests/comm/test_pcie_selected_records_ce_gpu.py b/tests/comm/test_pcie_selected_records_ce_gpu.py new file mode 100644 index 00000000..6050d6a3 --- /dev/null +++ b/tests/comm/test_pcie_selected_records_ce_gpu.py @@ -0,0 +1,825 @@ +from __future__ import annotations + +import os +import socket +from contextlib import suppress + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from sparkinfer.comm.pcie import pcie_dma as _pcie_dma +from sparkinfer.comm.pcie.pcie_selected_records_ce import ( + PCIeSelectedRecordCopyExchange, + _load_extension, +) + + +pytestmark = pytest.mark.skipif( + os.getenv("SPARKINFER_RUN_PCIE_SELECTED_RECORDS_CE_TEST") != "1", + reason=( + "set SPARKINFER_RUN_PCIE_SELECTED_RECORDS_CE_TEST=1 to run copy-engine " + "selected-record GPU tests" + ), +) + +MAX_RECORDS = 31 +PRIMARY_CAPACITY = 1 +POOL_RECORDS_PER_RANK = 67 +BIG_RECORD_BYTES = 65_536 + + +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 _record_values( + global_indices: torch.Tensor, + record_bytes: int, + iteration: int, +) -> torch.Tensor: + byte = torch.arange( + record_bytes, + device=global_indices.device, + dtype=torch.int64, + ) + values = ( + global_indices.to(torch.int64).unsqueeze(1) * 37 + + byte.unsqueeze(0) * 13 + + iteration * 7 + ) + return values.remainder(256).to(torch.uint8) + + +def _local_records( + rank: int, + world_size: int, + device: torch.device, + iteration: int, + record_bytes: int, +) -> torch.Tensor: + global_indices = ( + torch.arange( + POOL_RECORDS_PER_RANK, + device=device, + dtype=torch.int64, + ) + * world_size + + rank + ) + return _record_values(global_indices, record_bytes, iteration) + + +def _selection( + destination: int, + total_records: int, + count: int, + iteration: int, +) -> torch.Tensor: + selected = torch.arange(count, dtype=torch.int64) + return (selected * 17 + destination * 11 + iteration * 23).remainder(total_records) + + +def _local_maps( + rank: int, + world_size: int, + count: int, + iteration: int, + device: torch.device, +) -> torch.Tensor: + total_records = POOL_RECORDS_PER_RANK * world_size + maps = torch.full((world_size, count), -1, dtype=torch.int32) + for destination in range(world_size): + selected = _selection(destination, total_records, count, iteration) + owned = selected.remainder(world_size) == rank + maps[destination, owned] = ( + selected[owned].div(world_size, rounding_mode="floor").to(torch.int32) + ) + return maps.to(device) + + +def _expected( + rank: int, + world_size: int, + count: int, + iteration: int, + device: torch.device, + record_bytes: int, +) -> torch.Tensor: + selected = _selection( + rank, + POOL_RECORDS_PER_RANK * world_size, + count, + iteration, + ).to(device) + return _record_values(selected, record_bytes, iteration) + + +def _layer_records( + rank: int, + world_size: int, + device: torch.device, + iteration: int, + record_bytes: int, + layer: int, +) -> torch.Tensor: + records = _local_records( + rank, + world_size, + device, + iteration, + record_bytes, + ) + return (records.to(torch.int16) + layer * 53).remainder(256).to(torch.uint8) + + +def _skewed_local_maps( + rank: int, + world_size: int, + count: int, + device: torch.device, +) -> torch.Tensor: + maps = torch.full( + (world_size, count), + -1, + dtype=torch.int32, + device=device, + ) + if rank == 0: + maps.copy_( + torch.arange(count, dtype=torch.int32, device=device).expand( + world_size, + count, + ) + ) + return maps + + +def _layer_expected( + world_size: int, + count: int, + iteration: int, + device: torch.device, + record_bytes: int, + layer: int, +) -> torch.Tensor: + selected = torch.arange(count, dtype=torch.int64, device=device) * world_size + expected = _record_values(selected, record_bytes, iteration) + return (expected.to(torch.int16) + layer * 53).remainder(256).to(torch.uint8) + + +def _internal_addresses( + exchange: PCIeSelectedRecordCopyExchange, +) -> tuple[object, ...]: + shared = exchange._owned_buffer + assert shared is not None + return ( + shared.local_ptr, + shared.peer_ptrs, + shared.remote_ptrs, + exchange._local_slab_ptr, + exchange._peer_overflow_ptrs.data_ptr(), + exchange._peer_layer_overflow_ptrs.data_ptr(), + exchange._barrier_publish_ptrs.data_ptr(), + exchange._barrier_wait_ptrs.data_ptr(), + exchange._send_counters.data_ptr(), + exchange._wait_counters.data_ptr(), + ) + + +def _assert_forced_overflow( + exchange: PCIeSelectedRecordCopyExchange, + maps: torch.Tensor, +) -> None: + active_records = maps.numel() // exchange.world_size + active_primary = exchange._active_primary_capacity(active_records) + valid_by_destination = maps.ge(0).sum(dim=1) + assert bool(torch.all(valid_by_destination > active_primary).item()) + + +def _check_eager( + exchange: PCIeSelectedRecordCopyExchange, + rank: int, + world_size: int, + device: torch.device, + record_bytes: int, +) -> None: + addresses = _internal_addresses(exchange) + for iteration, count in enumerate((0, 1, 13, MAX_RECORDS), start=1): + records = _local_records( + rank, + world_size, + device, + iteration, + record_bytes, + ) + maps = _local_maps(rank, world_size, count, iteration, device) + if count == MAX_RECORDS: + _assert_forced_overflow(exchange, maps) + out = torch.empty( + (count, record_bytes), + dtype=torch.uint8, + device=device, + ) + returned = exchange.exchange(records, maps, out) + torch.cuda.synchronize(device) + assert returned is out + assert torch.equal( + out, + _expected( + rank, + world_size, + count, + iteration, + device, + record_bytes, + ), + ) + assert _internal_addresses(exchange) == addresses + + iteration = 50 + records = _local_records( + rank, + world_size, + device, + iteration, + record_bytes, + ) + maps = _local_maps(rank, world_size, MAX_RECORDS, iteration, device) + _assert_forced_overflow(exchange, maps) + out = torch.empty( + (MAX_RECORDS, record_bytes), + dtype=torch.uint8, + device=device, + ) + exchange.exchange(records, maps, out) + torch.cuda.synchronize(device) + allocated_before = torch.cuda.memory_allocated(device) + torch.cuda.reset_peak_memory_stats(device) + for _ in range(8): + exchange.exchange(records, maps, out) + torch.cuda.synchronize(device) + assert torch.cuda.memory_allocated(device) == allocated_before + assert torch.cuda.max_memory_allocated(device) == allocated_before + assert _internal_addresses(exchange) == addresses + assert torch.equal( + out, + _expected( + rank, + world_size, + MAX_RECORDS, + iteration, + device, + record_bytes, + ), + ) + + +def _check_graph( + exchange: PCIeSelectedRecordCopyExchange, + rank: int, + world_size: int, + device: torch.device, + record_bytes: int, +) -> None: + records = _local_records(rank, world_size, device, 100, record_bytes) + maps = _local_maps(rank, world_size, MAX_RECORDS, 100, device) + _assert_forced_overflow(exchange, maps) + out = torch.empty( + (MAX_RECORDS, record_bytes), + dtype=torch.uint8, + device=device, + ) + stream = torch.cuda.Stream(device=device) + with torch.cuda.stream(stream): + exchange.exchange(records, maps, out) + stream.synchronize() + dist.barrier() + addresses = _internal_addresses(exchange) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + exchange.exchange(records, maps, out) + dist.barrier() + graph.replay() + stream.synchronize() + assert _internal_addresses(exchange) == addresses + + iterations = tuple(range(101, 109)) + record_updates = [ + _local_records(rank, world_size, device, iteration, record_bytes) + for iteration in iterations + ] + map_updates = [ + _local_maps(rank, world_size, MAX_RECORDS, iteration, device) + for iteration in iterations + ] + expected = torch.stack( + [ + _expected( + rank, + world_size, + MAX_RECORDS, + iteration, + device, + record_bytes, + ) + for iteration in iterations + ] + ) + observed = torch.empty_like(expected) + allocated_before = torch.cuda.memory_allocated(device) + torch.cuda.reset_peak_memory_stats(device) + for replay in range(len(iterations)): + records.copy_(record_updates[replay]) + maps.copy_(map_updates[replay]) + stream.wait_stream(torch.cuda.current_stream(device)) + graph.replay() + stream.synchronize() + observed[replay].copy_(out) + torch.cuda.synchronize(device) + assert torch.cuda.memory_allocated(device) == allocated_before + assert torch.cuda.max_memory_allocated(device) == allocated_before + assert _internal_addresses(exchange) == addresses + assert torch.equal(observed, expected) + + +def _check_layers_eager( + exchange: PCIeSelectedRecordCopyExchange, + rank: int, + world_size: int, + device: torch.device, + record_bytes: int, +) -> None: + count = min(8, exchange.layered_max_records) + assert count == 8 + maps = _skewed_local_maps(rank, world_size, count, device) + active_primary = exchange._active_layer_primary_capacity(count) + if rank == 0: + assert bool(torch.all(maps.ge(0).sum(dim=1) > active_primary).item()) + addresses = _internal_addresses(exchange) + + for iteration in (200, 201): + records = tuple( + _layer_records( + rank, + world_size, + device, + iteration, + record_bytes, + layer, + ) + for layer in range(3) + ) + outputs = tuple( + torch.empty((count, record_bytes), dtype=torch.uint8, device=device) + for _ in range(3) + ) + returned = exchange.exchange_layers(records, maps, outputs) + torch.cuda.synchronize(device) + for layer, (actual, output) in enumerate(zip(returned, outputs, strict=True)): + assert actual is output + assert torch.equal( + output, + _layer_expected( + world_size, + count, + iteration, + device, + record_bytes, + layer, + ), + ) + assert _internal_addresses(exchange) == addresses + + +def _check_layers_graph( + exchange: PCIeSelectedRecordCopyExchange, + rank: int, + world_size: int, + device: torch.device, + record_bytes: int, +) -> None: + count = min(8, exchange.layered_max_records) + assert count == 8 + iteration = 300 + records = tuple( + _layer_records( + rank, + world_size, + device, + iteration, + record_bytes, + layer, + ) + for layer in range(3) + ) + maps = _skewed_local_maps(rank, world_size, count, device) + outputs = tuple( + torch.empty((count, record_bytes), dtype=torch.uint8, device=device) + for _ in range(3) + ) + stream = torch.cuda.Stream(device=device) + with torch.cuda.stream(stream): + exchange.exchange_layers(records, maps, outputs) + stream.synchronize() + dist.barrier() + addresses = _internal_addresses(exchange) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + exchange.exchange_layers(records, maps, outputs) + dist.barrier() + graph.replay() + stream.synchronize() + assert _internal_addresses(exchange) == addresses + + iterations = tuple(range(301, 305)) + updates = [ + tuple( + _layer_records( + rank, + world_size, + device, + current_iteration, + record_bytes, + layer, + ) + for layer in range(3) + ) + for current_iteration in iterations + ] + observed = tuple( + torch.empty( + (len(iterations), count, record_bytes), + dtype=torch.uint8, + device=device, + ) + for _ in range(3) + ) + expected = tuple( + torch.stack( + [ + _layer_expected( + world_size, + count, + current_iteration, + device, + record_bytes, + layer, + ) + for current_iteration in iterations + ] + ) + for layer in range(3) + ) + stream.wait_stream(torch.cuda.current_stream(device)) + stream.synchronize() + allocated_before = torch.cuda.memory_allocated(device) + torch.cuda.reset_peak_memory_stats(device) + for replay, _ in enumerate(iterations): + for records_plane, update in zip(records, updates[replay], strict=True): + records_plane.copy_(update) + stream.wait_stream(torch.cuda.current_stream(device)) + graph.replay() + stream.synchronize() + for layer, output in enumerate(outputs): + observed[layer][replay].copy_(output) + torch.cuda.synchronize(device) + assert torch.cuda.memory_allocated(device) == allocated_before + assert torch.cuda.max_memory_allocated(device) == allocated_before + assert _internal_addresses(exchange) == addresses + assert all( + torch.equal(actual, wanted) + for actual, wanted in zip(observed, expected, strict=True) + ) + + +def _check_mixed_graph_modes( + exchange: PCIeSelectedRecordCopyExchange, + rank: int, + world_size: int, + device: torch.device, + record_bytes: int, +) -> None: + layer_count = min(8, exchange.layered_max_records) + layer_maps = _skewed_local_maps(rank, world_size, layer_count, device) + single_maps = _local_maps(rank, world_size, MAX_RECORDS, 400, device) + single_records = _local_records(rank, world_size, device, 400, record_bytes) + single_out = torch.empty( + (MAX_RECORDS, record_bytes), dtype=torch.uint8, device=device + ) + layer_records = tuple( + _layer_records( + rank, + world_size, + device, + 500, + record_bytes, + layer, + ) + for layer in range(3) + ) + layer_outputs = tuple( + torch.empty((layer_count, record_bytes), dtype=torch.uint8, device=device) + for _ in range(3) + ) + empty_records = torch.empty((0, record_bytes), dtype=torch.uint8, device=device) + empty_maps = torch.empty((world_size, 0), dtype=torch.int32, device=device) + empty_out = torch.empty((0, record_bytes), dtype=torch.uint8, device=device) + stream = torch.cuda.Stream(device=device) + + with torch.cuda.stream(stream): + exchange.exchange(single_records, single_maps, single_out) + exchange.exchange_layers(layer_records, layer_maps, layer_outputs) + stream.synchronize() + dist.barrier() + + single_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(single_graph, stream=stream): + exchange.exchange(single_records, single_maps, single_out) + dist.barrier() + layer_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(layer_graph, stream=stream): + exchange.exchange_layers(layer_records, layer_maps, layer_outputs) + dist.barrier() + empty_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(empty_graph, stream=stream): + exchange.exchange(empty_records, empty_maps, empty_out) + dist.barrier() + + # Reproduce production switching and preserve the graph-balanced release + # chain even when one rank-consistent exchange has no selected records. + replay_modes = ("layer", "empty", "empty", "single", "single", "layer") + for replay, mode in enumerate(replay_modes): + iteration = 600 + replay + if mode == "single": + single_records.copy_( + _local_records(rank, world_size, device, iteration, record_bytes) + ) + single_maps.copy_( + _local_maps(rank, world_size, MAX_RECORDS, iteration, device) + ) + elif mode == "layer": + for layer, records in enumerate(layer_records): + records.copy_( + _layer_records( + rank, + world_size, + device, + iteration, + record_bytes, + layer, + ) + ) + stream.wait_stream(torch.cuda.current_stream(device)) + dist.barrier() + if mode == "single": + single_graph.replay() + elif mode == "empty": + empty_graph.replay() + else: + layer_graph.replay() + stream.synchronize() + + if mode == "single": + assert torch.equal( + single_out, + _expected( + rank, + world_size, + MAX_RECORDS, + iteration, + device, + record_bytes, + ), + ) + elif mode == "layer": + for layer, output in enumerate(layer_outputs): + assert torch.equal( + output, + _layer_expected( + world_size, + layer_count, + iteration, + device, + record_bytes, + layer, + ), + ) + + +def _check_big_pool_offset( + rank: int, + world_size: int, + device: torch.device, +) -> None: + high_record = 2**31 // BIG_RECORD_BYTES + 1 + required_bytes = (high_record + 1) * BIG_RECORD_BYTES + torch.cuda.empty_cache() + free_bytes, _ = torch.cuda.mem_get_info(device) + available = torch.tensor( + [1 if rank != 0 or free_bytes >= required_bytes + 512 * 1024**2 else 0], + dtype=torch.int32, + device=device, + ) + dist.all_reduce(available, op=dist.ReduceOp.MIN) + if not bool(available.item()): + raise RuntimeError( + "copy-engine selected-record big-offset gate requires at least " + f"{required_bytes + 512 * 1024**2} free bytes on rank 0" + ) + + exchange = PCIeSelectedRecordCopyExchange.from_process_group( + process_group=dist.group.WORLD, + device=device, + max_records=1, + record_bytes=BIG_RECORD_BYTES, + primary_capacity=1, + ) + records = None + try: + rows = high_record + 1 if rank == 0 else 1 + records = torch.empty( + (rows, BIG_RECORD_BYTES), + dtype=torch.uint8, + device=device, + ) + expected = torch.arange( + BIG_RECORD_BYTES, + dtype=torch.int64, + device=device, + ) + expected = ( + expected.remainder(251) + .to(torch.uint8) + .reshape( + 1, + BIG_RECORD_BYTES, + ) + ) + if rank == 0: + records[high_record].copy_(expected[0]) + maps = torch.full( + (world_size, 1), + -1, + dtype=torch.int32, + device=device, + ) + if rank == 0: + maps.fill_(high_record) + out = torch.empty_like(expected) + exchange.exchange(records, maps, out) + torch.cuda.synchronize(device) + assert high_record * BIG_RECORD_BYTES > 2**31 + assert torch.equal(out, expected) + finally: + exchange.close() + with suppress(Exception): + del records + torch.cuda.empty_cache() + dist.barrier() + + +def _worker( + rank: int, + world_size: int, + port: int, + record_bytes: 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=world_size, + ) + try: + eager_exchange = PCIeSelectedRecordCopyExchange.from_process_group( + process_group=dist.group.WORLD, + device=device, + max_records=MAX_RECORDS, + record_bytes=record_bytes, + primary_capacity=PRIMARY_CAPACITY, + ) + try: + _check_eager( + eager_exchange, + rank, + world_size, + device, + record_bytes, + ) + _check_layers_eager( + eager_exchange, + rank, + world_size, + device, + record_bytes, + ) + torch.cuda.synchronize(device) + finally: + eager_exchange.close() + dist.barrier() + + graph_exchange = PCIeSelectedRecordCopyExchange.from_process_group( + process_group=dist.group.WORLD, + device=device, + max_records=MAX_RECORDS, + record_bytes=record_bytes, + primary_capacity=PRIMARY_CAPACITY, + ) + try: + _check_graph( + graph_exchange, + rank, + world_size, + device, + record_bytes, + ) + torch.cuda.synchronize(device) + finally: + graph_exchange.close() + dist.barrier() + + layer_graph_exchange = PCIeSelectedRecordCopyExchange.from_process_group( + process_group=dist.group.WORLD, + device=device, + max_records=MAX_RECORDS, + record_bytes=record_bytes, + primary_capacity=PRIMARY_CAPACITY, + ) + try: + _check_layers_graph( + layer_graph_exchange, + rank, + world_size, + device, + record_bytes, + ) + torch.cuda.synchronize(device) + finally: + layer_graph_exchange.close() + dist.barrier() + + mixed_graph_exchange = PCIeSelectedRecordCopyExchange.from_process_group( + process_group=dist.group.WORLD, + device=device, + max_records=MAX_RECORDS, + record_bytes=record_bytes, + primary_capacity=PRIMARY_CAPACITY, + ) + try: + _check_mixed_graph_modes( + mixed_graph_exchange, + rank, + world_size, + device, + record_bytes, + ) + torch.cuda.synchronize(device) + finally: + mixed_graph_exchange.close() + dist.barrier() + + _check_big_pool_offset(rank, world_size, device) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize( + ("world_size", "record_bytes"), + ((4, 432), (8, 656)), + ids=("dcp4-record432", "dcp8-record656"), +) +def test_pcie_selected_record_ce_eager_graph_overflow_and_big_offset( + world_size: int, + record_bytes: int, +) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + if ( + world_size == 8 + and os.getenv("SPARKINFER_RUN_PCIE_SELECTED_RECORDS_CE_WORLD8_TEST") != "1" + ): + pytest.skip( + "set SPARKINFER_RUN_PCIE_SELECTED_RECORDS_CE_WORLD8_TEST=1 for the " + "DCP8/656-byte gate" + ) + if torch.cuda.device_count() < world_size: + pytest.skip( + f"need {world_size} CUDA devices, found {torch.cuda.device_count()}" + ) + + _load_extension() + _pcie_dma._load_extension() + mp.spawn( + _worker, + args=(world_size, _free_port(), record_bytes), + nprocs=world_size, + join=True, + ) From 33b4e4c92c11105646492d98ecc0c8a576ad6929 Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Mon, 20 Jul 2026 11:54:44 -0500 Subject: [PATCH 3/6] fix(dcp): support disposable A2A channel profiling --- sparkinfer/comm/pcie/pcie_dcp_a2a.py | 41 +++++++++++ tests/comm/test_pcie_dcp_a2a.py | 102 +++++++++++++++++++++++++++ 2 files changed, 143 insertions(+) diff --git a/sparkinfer/comm/pcie/pcie_dcp_a2a.py b/sparkinfer/comm/pcie/pcie_dcp_a2a.py index 81c421f2..bdfc86d7 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_a2a.py +++ b/sparkinfer/comm/pcie/pcie_dcp_a2a.py @@ -497,6 +497,12 @@ def __del__(self) -> None: self.close() +@dataclass(frozen=True) +class _ChannelCheckpoint: + pool_id: int + channels: tuple[tuple[int, PCIeDCPA2A], ...] + + class PCIeDCPA2APool: """Create an independent DCP collective channel for each CUDA stream.""" @@ -702,6 +708,41 @@ def capture(self, stream: object = None): if popped is not channel: raise RuntimeError("PCIe DCP A2A capture channel stack corrupted") + def checkpoint_channels(self) -> _ChannelCheckpoint: + """Snapshot channels before a disposable CUDA graph capture.""" + if self._closed: + raise RuntimeError("PCIeDCPA2APool is closed") + if self._capture_channel_stack: + raise RuntimeError("cannot checkpoint channels during CUDA graph capture") + return _ChannelCheckpoint( + pool_id=id(self), + channels=tuple(self._channels.items()), + ) + + def rollback_channels(self, checkpoint: _ChannelCheckpoint) -> None: + """Release channels created after ``checkpoint`` and restore aliases.""" + if self._closed: + raise RuntimeError("PCIeDCPA2APool is closed") + if self._capture_channel_stack: + raise RuntimeError("cannot roll back channels during CUDA graph capture") + if checkpoint.pool_id != id(self): + raise ValueError("channel checkpoint belongs to a different pool") + + saved_channels = dict(checkpoint.channels) + saved_channel_ids = {id(channel) for channel in saved_channels.values()} + new_channels: list[PCIeDCPA2A] = [] + seen: set[int] = set() + for channel in self._channels.values(): + channel_id = id(channel) + if channel_id not in saved_channel_ids and channel_id not in seen: + seen.add(channel_id) + new_channels.append(channel) + + self._channels.clear() + self._channels.update(saved_channels) + for channel in new_channels: + channel.close() + def close(self) -> None: if self._closed: return diff --git a/tests/comm/test_pcie_dcp_a2a.py b/tests/comm/test_pcie_dcp_a2a.py index 7e9b30f7..751c67d0 100644 --- a/tests/comm/test_pcie_dcp_a2a.py +++ b/tests/comm/test_pcie_dcp_a2a.py @@ -5,6 +5,7 @@ from sparkinfer.comm.pcie.pcie_dcp_a2a import ( PCIeDCPA2A, + PCIeDCPA2APool, _staging_layout, lse_reduce_scatter_reference, ) @@ -223,3 +224,104 @@ def test_runtime_rejects_shape_dtype_and_capacity_mismatches(): runtime.all_gather_heads(good_output[:, :8]) with pytest.raises(ValueError, match="exceeds configured capacity"): runtime.all_gather_heads(torch.zeros(5, 16, 64, dtype=torch.bfloat16)) + + +def test_pool_rollback_closes_only_channels_created_after_checkpoint(monkeypatch): + created = [] + current_stream = [7] + capturing = [False] + + def make_channel(stream_key): + ext = _FakeExt() + runtime = _make_runtime(ext) + created.append((stream_key, runtime, ext)) + return runtime + + pool = PCIeDCPA2APool( + rank=0, + world_size=2, + device=torch.device("cpu"), + max_batch_size=4, + total_heads=32, + head_dim=64, + channel_factory=make_channel, + ) + 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], + ) + + target_channel = pool.for_stream(7) + checkpoint = pool.checkpoint_channels() + + with pool.capture(8) as draft_channel: + capturing[0] = True + current_stream[0] = 80 + assert pool.for_stream() is draft_channel + capturing[0] = False + + pool.rollback_channels(checkpoint) + + assert pool._channels == {7: target_channel} + assert created[0][2].dispose_calls == [] + assert created[1][2].dispose_calls == [1234] + + +def test_pool_rollback_removes_new_alias_without_closing_saved_channel(monkeypatch): + ext = _FakeExt() + 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(ext), + ) + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_dcp_a2a._current_stream_key", + lambda device, stream=None: 70 if stream is None else int(stream), + ) + capturing = [False] + monkeypatch.setattr( + "sparkinfer.comm.pcie.pcie_dcp_a2a._is_current_stream_capturing", + lambda device: capturing[0], + ) + + channel = pool.for_stream(7) + checkpoint = pool.checkpoint_channels() + with pool.capture(7): + capturing[0] = True + assert pool.for_stream() is channel + capturing[0] = False + + assert pool._channels == {7: channel, 70: channel} + pool.rollback_channels(checkpoint) + + assert pool._channels == {7: channel} + assert ext.dispose_calls == [] + + +def test_pool_rejects_checkpoint_from_another_pool(): + def make_pool(): + return 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(), + ) + + first = make_pool() + second = make_pool() + + with pytest.raises(ValueError, match="different pool"): + second.rollback_channels(first.checkpoint_channels()) From f69aad7650f3e810b8f496c73fdaa0081313a35c Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Tue, 21 Jul 2026 12:44:13 -0500 Subject: [PATCH 4/6] ci(fork): validate selected-record transport on hosted runners Assisted-by: OpenAI Codex Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> --- .github/workflows/fork-ci.yml | 42 +++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 .github/workflows/fork-ci.yml diff --git a/.github/workflows/fork-ci.yml b/.github/workflows/fork-ci.yml new file mode 100644 index 00000000..2299ae73 --- /dev/null +++ b/.github/workflows/fork-ci.yml @@ -0,0 +1,42 @@ +name: fork-ci + +on: + pull_request: + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + selected-record-transport: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0 + with: + python-version: "3.12" + cache: pip + - name: Install CPU test dependencies + run: | + python -m pip install --upgrade pip + python -m pip install pytest ruff torch + python -m pip install --no-deps -e . + - name: Ruff + run: | + ruff check \ + sparkinfer/comm/pcie/pcie_selected_records.py \ + sparkinfer/comm/pcie/pcie_selected_records_ce.py \ + sparkinfer/comm/pcie/pcie_dcp_a2a.py \ + tests/comm/test_pcie_selected_records.py \ + tests/comm/test_pcie_selected_records_ce.py \ + tests/comm/test_pcie_dcp_a2a.py + - name: Selected-record transport tests + run: | + pytest -q \ + tests/comm/test_pcie_selected_records.py \ + tests/comm/test_pcie_selected_records_ce.py \ + tests/comm/test_pcie_dcp_a2a.py From 225f79062ff06498d6094b2dadd1ba7b26753262 Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:52:09 -0500 Subject: [PATCH 5/6] test(comm): cover 368-byte selected-record exchange --- tests/comm/test_pcie_selected_records_ce_gpu.py | 8 ++++---- tests/comm/test_pcie_selected_records_gpu.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/comm/test_pcie_selected_records_ce_gpu.py b/tests/comm/test_pcie_selected_records_ce_gpu.py index 6050d6a3..6e28c552 100644 --- a/tests/comm/test_pcie_selected_records_ce_gpu.py +++ b/tests/comm/test_pcie_selected_records_ce_gpu.py @@ -793,8 +793,8 @@ def _worker( @pytest.mark.parametrize( ("world_size", "record_bytes"), - ((4, 432), (8, 656)), - ids=("dcp4-record432", "dcp8-record656"), + ((4, 368), (4, 432), (8, 368), (8, 656)), + ids=("dcp4-record368", "dcp4-record432", "dcp8-record368", "dcp8-record656"), ) def test_pcie_selected_record_ce_eager_graph_overflow_and_big_offset( world_size: int, @@ -807,8 +807,8 @@ def test_pcie_selected_record_ce_eager_graph_overflow_and_big_offset( and os.getenv("SPARKINFER_RUN_PCIE_SELECTED_RECORDS_CE_WORLD8_TEST") != "1" ): pytest.skip( - "set SPARKINFER_RUN_PCIE_SELECTED_RECORDS_CE_WORLD8_TEST=1 for the " - "DCP8/656-byte gate" + "set SPARKINFER_RUN_PCIE_SELECTED_RECORDS_CE_WORLD8_TEST=1 for " + "the DCP8 selected-record gate" ) if torch.cuda.device_count() < world_size: pytest.skip( diff --git a/tests/comm/test_pcie_selected_records_gpu.py b/tests/comm/test_pcie_selected_records_gpu.py index 8dc28654..68daed9e 100644 --- a/tests/comm/test_pcie_selected_records_gpu.py +++ b/tests/comm/test_pcie_selected_records_gpu.py @@ -24,7 +24,7 @@ ) MAX_RECORDS = 31 -RECORD_WIDTHS = (16, 37, 432) +RECORD_WIDTHS = (16, 37, 368, 432) ODD_RECORD_BYTES = 37 POOL_RECORDS_PER_RANK = 67 From 31d296cc35f7653057fe26445aa45a454535fef3 Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:52:42 -0500 Subject: [PATCH 6/6] fix(indexer): bound long-context prefill scratch --- sparkinfer/attention/nsa_indexer/paged.py | 107 +++++++++++++++--- .../test_paged_prefill_topk_long_context.py | 59 ++++++++++ 2 files changed, 148 insertions(+), 18 deletions(-) diff --git a/sparkinfer/attention/nsa_indexer/paged.py b/sparkinfer/attention/nsa_indexer/paged.py index e0a88b69..c7250b78 100644 --- a/sparkinfer/attention/nsa_indexer/paged.py +++ b/sparkinfer/attention/nsa_indexer/paged.py @@ -34,7 +34,86 @@ # cap that keeps the candidate buffers capacity-independent (~topk*8B*cap per # row) at very long contexts. _TWO_LEVEL_SLICE_TOKENS = 16384 -_TWO_LEVEL_MAX_SLICES = 32 +_TWO_LEVEL_MAX_SLICES_ENV = "SPARKINFER_PAGED_INDEX_TWO_LEVEL_MAX_SLICES" + + +def _read_positive_int_env(name: str, default: int) -> int: + raw_value = os.environ.get(name) + if raw_value is None: + return default + try: + value = int(raw_value) + except ValueError as exc: + raise ValueError( + f"{name} must be a positive integer, got {raw_value!r}" + ) from exc + if value < 1: + raise ValueError(f"{name} must be a positive integer, got {raw_value!r}") + return value + + +_TWO_LEVEL_MAX_SLICES = _read_positive_int_env( + _TWO_LEVEL_MAX_SLICES_ENV, + 32, +) + + +def _two_level_slice_width( + width_tokens: int, + page_size: int, + *, + max_slices: int = _TWO_LEVEL_MAX_SLICES, +) -> int: + """Choose an aligned level-1 width while bounding candidate-buffer slices.""" + if max_slices < 1: + raise ValueError(f"max_slices must be positive, got {max_slices}") + slice_tokens = _TWO_LEVEL_SLICE_TOKENS + min_slice = -(-width_tokens // max_slices) + if min_slice > slice_tokens: + slice_tokens = -(-min_slice // page_size) * page_size + return slice_tokens + + +def _plan_two_level_slices( + page_table_width: int, + page_size: int, + supertile_pages: int, + *, + max_slices: int = _TWO_LEVEL_MAX_SLICES, +) -> list[tuple[int, int]]: + """Plan bounded two-level candidates or select the streaming-fold fallback.""" + if page_table_width < 1 or page_size < 1 or supertile_pages < 1: + raise ValueError( + "page_table_width, page_size, and supertile_pages must be positive" + ) + width_tokens = page_table_width * page_size + if width_tokens < 2 * _TWO_LEVEL_SLICE_TOKENS: + return [] + + slice_tokens = _two_level_slice_width( + width_tokens, + page_size, + max_slices=max_slices, + ) + num_chunks = -(-page_table_width // supertile_pages) + slices: list[tuple[int, int]] = [] + base = 0 + for chunk_idx in range(num_chunks): + chunk_pages = ( + min((chunk_idx + 1) * supertile_pages, page_table_width) + - chunk_idx * supertile_pages + ) + chunk_tokens = chunk_pages * page_size + chunk_slices = max(1, -(-chunk_tokens // slice_tokens)) + if base + chunk_slices > max_slices: + # A wider slice cannot span separate supertile launches. Use the + # bounded two-buffer streaming fold instead of exceeding the cap. + return [] + slices.append((chunk_slices, base)) + base += chunk_slices + return slices + + from sparkinfer.attention.nsa_indexer.reference import ( pack_index_k_cache_reference, paged_decode_logits_reference, @@ -779,27 +858,19 @@ def index_topk_fp8( # the last chunk. The per-chunk tile-logits scratch stays at the supertile # ceiling. Physical-slot output keeps the legacy carry chain (the fold # gather emits logical indices). - width_tokens = page_table_width * page_size - two_level_slices: list[tuple[int, int]] = [] + two_level_slices = _plan_two_level_slices( + page_table_width, + page_size, + supertile_pages, + ) total_slices = 0 fold_values = None fold_indices = None fold_lengths = None - if not output_physical_slots and width_tokens >= 2 * _TWO_LEVEL_SLICE_TOKENS: - slice_tokens = _TWO_LEVEL_SLICE_TOKENS - min_slice = -(-width_tokens // _TWO_LEVEL_MAX_SLICES) - if min_slice > slice_tokens: - slice_tokens = -(-min_slice // page_size) * page_size - base = 0 - for c in range(num_chunks): - c_pages = ( - min((c + 1) * supertile_pages, page_table_width) - c * supertile_pages - ) - c_tokens = c_pages * page_size - splits_c = max(1, -(-c_tokens // slice_tokens)) - two_level_slices.append((splits_c, base)) - base += splits_c - total_slices = base + if output_physical_slots: + two_level_slices = [] + if two_level_slices: + total_slices = sum(chunk_slices for chunk_slices, _ in two_level_slices) fold_values = torch.empty( (q_rows * total_slices, topk), dtype=torch.float32, device=q_fp8.device ) diff --git a/tests/attention/test_paged_prefill_topk_long_context.py b/tests/attention/test_paged_prefill_topk_long_context.py index 3ff68d73..381b82a2 100644 --- a/tests/attention/test_paged_prefill_topk_long_context.py +++ b/tests/attention/test_paged_prefill_topk_long_context.py @@ -28,6 +28,8 @@ from sparkinfer.attention.nsa_indexer._impl import clear_indexer_caches from sparkinfer.attention.nsa_indexer.paged import ( + _plan_two_level_slices, + _two_level_slice_width, index_topk_fp8, pack_paged_index_k_cache_reference, prepare_paged_indexer_metadata, @@ -45,6 +47,63 @@ _PAGE_START = 3 +@pytest.mark.parametrize( + ("width_tokens", "max_slices", "expected_width", "expected_slices"), + [ + (131072, 32, 16384, 8), + (262144, 32, 16384, 16), + (262144, 8, 32768, 8), + (300000, 8, 37504, 8), + ], +) +def test_two_level_slice_width_bounds_candidate_workspace( + width_tokens: int, + max_slices: int, + expected_width: int, + expected_slices: int, +) -> None: + slice_width = _two_level_slice_width( + width_tokens, + _PAGE, + max_slices=max_slices, + ) + assert slice_width == expected_width + assert -(-width_tokens // slice_width) == expected_slices + + +@pytest.mark.parametrize( + ( + "width_tokens", + "supertile_tokens", + "max_slices", + "expected_chunks", + "expected_slices", + ), + [ + (131072, 16384, 8, 8, 8), + (167936, 16384, 8, 0, 0), + (262144, 32768, 8, 8, 8), + (262144, 16384, 8, 0, 0), + (300032, 16384, 32, 19, 19), + ], +) +def test_two_level_slice_plan_falls_back_before_exceeding_cap( + width_tokens: int, + supertile_tokens: int, + max_slices: int, + expected_chunks: int, + expected_slices: int, +) -> None: + plan = _plan_two_level_slices( + width_tokens // _PAGE, + _PAGE, + supertile_tokens // _PAGE, + max_slices=max_slices, + ) + assert len(plan) == expected_chunks + assert sum(chunk_slices for chunk_slices, _ in plan) == expected_slices + + def _build_scene(device: torch.device, seq_len: int, scores: str) -> dict: """Build a paged prefill scene plus its fp32 reference top-k threshold.