integration: validate selected-record CKV transports on latest Sparkinfer

.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

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

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
         )

sparkinfer/comm/pcie/__init__.py

@@ -11,6 +11,10 @@
 - ``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.
+- ``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.
@@ -33,12 +37,15 @@
         "TwoShotReduceScatter",
         "DcpAllToAll",
         "DcpAllToAllPool",
+        "SelectedRecordExchange",
+        "SelectedRecordExchangeInitializationError",
+        "SelectedRecordCopyExchange",
         "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 +64,9 @@
         DmaAllReduce,
         OneshotAllReduce,
         OneshotAllReducePool,
+        SelectedRecordExchange,
+        SelectedRecordExchangeInitializationError,
+        SelectedRecordCopyExchange,
         TwoShotReduceScatter,
         autotune_dma_crossovers,
         is_supported,

sparkinfer/comm/pcie/api.py

@@ -27,6 +27,15 @@
 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_selected_records_ce import (
+    PCIeSelectedRecordCopyExchange as SelectedRecordCopyExchange,
+)
 from .pcie_twoshot import (
     PCIeTwoShotSP as TwoShotReduceScatter,
 )
@@ -50,6 +59,9 @@ def is_supported(device=None) -> bool:
     "TwoShotReduceScatter",
     "DcpAllToAll",
     "DcpAllToAllPool",
+    "SelectedRecordExchange",
+    "SelectedRecordExchangeInitializationError",
+    "SelectedRecordCopyExchange",
     "autotune_dma_crossovers",
     "parse_oneshot_max_size",
     "lse_reduce_scatter_reference",

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