From f3d5c14bac183fdff60877f6edc7dad377dd6e0f 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] 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> --- pyproject.toml | 2 +- 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 +++++++++++ 7 files changed, 1986 insertions(+), 2 deletions(-) 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/pyproject.toml b/pyproject.toml index 0f9cf4f4..af2d94a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ where = ["."] include = ["sparkinfer*"] [tool.setuptools.package-data] -"sparkinfer.distributed" = ["*.cu"] +"sparkinfer.comm.pcie" = ["*.cu"] [tool.ruff.lint] select = ["E", "F", "B", "SIM"] 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)