feat(comm): add copy-engine selected-record exchange#65
Conversation
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 <codex@openai.com> Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
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 <codex@openai.com> Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
📝 WalkthroughWalkthroughAdds public selected-record PCIe exchange APIs, direct-scatter and copy-engine CUDA implementations, layered three-plane transfers, IPC slab coordination, stream and lifecycle rules, packaging updates, documentation, and extensive CPU, CUDA graph, and distributed NCCL tests. ChangesSelected-record PCIe exchange
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant PCIeSelectedRecordCopyExchange
participant CUDAKernels
participant PCIeDMA
participant PeerSlabs
Caller->>PCIeSelectedRecordCopyExchange: exchange or exchange_layers
PCIeSelectedRecordCopyExchange->>CUDAKernels: pack selected records
PCIeSelectedRecordCopyExchange->>PCIeDMA: copy packed payloads
PCIeDMA->>PeerSlabs: write peer staging regions
PCIeSelectedRecordCopyExchange->>CUDAKernels: synchronize and unpack
PCIeSelectedRecordCopyExchange-->>Caller: return output tensors
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
The native 368-byte direct and CE/bulk transports have now passed full TP4/DCP4/MTP3 FP8-RoPE integration on the current GG/v20 stack. Runtime logs confirmed The transport is functionally correct, but the dependent vLLM selected-record path does not currently improve the already optimized Sparkinfer/B12X DCP baseline: C1 is approximately neutral and C4 regresses by roughly 5-11%. Full numbers and the release disposition are recorded in local-inference-lab/vllm#161. Recommendation: retain this generic transport PR for review and future work, but keep the vLLM integration opt-in/default-off and do not make it a release-image dependency yet. |
Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (7)
sparkinfer/comm/pcie/pcie_selected_records.cu (1)
100-106: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winUnbounded
local_recordallows OOB reads ofrecords.Only the negative sentinel is filtered; a stale or malformed map entry
>= records.numel() / record_bytesreads past the source tensor (the Python_validatealso checks only dtype/shape, not index range). A record-count parameter plus one compare keeps the fast path intact and turns a silent OOB device read into a skip.🛡️ Proposed guard
template <typename index_t, typename copy_t> __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, + int64_t source_records, int world_size) { @@ const int64_t local_record = static_cast<int64_t>( local_indices_by_destination[map_offset]); - if (local_record < 0) { + if (local_record < 0 || local_record >= source_records) { continue; }Pass
records.numel() / record_bytesat each of the four launch sites.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sparkinfer/comm/pcie/pcie_selected_records.cu` around lines 100 - 106, Add a record-count parameter to the relevant CUDA kernel and pass records.numel() / record_bytes from all four launch sites. In the local_record handling near local_indices_by_destination, skip entries when local_record is negative or greater than or equal to the record count before accessing records, preserving the existing fast path for valid indices.sparkinfer/comm/pcie/pcie_selected_records_ce.py (3)
367-368: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMark the deliberate broad catches so the lint gate stays clean.
Ruff flags BLE001 at Lines 367, 393, 403, 614, 642, 657 and 694. Catching everything is correct here — any local failure has to be converted into a rank-agreed status before the next collective, otherwise ranks diverge — so add a targeted
# noqa: BLE001with that rationale rather than narrowing the exception types.♻️ Example
- except Exception as exc: + except Exception as exc: # noqa: BLE001 - must reach the group status collective local_error = exc🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sparkinfer/comm/pcie/pcie_selected_records_ce.py` around lines 367 - 368, Add a targeted `# noqa: BLE001` annotation to each deliberate broad `except Exception` handler at the referenced locations, including the handlers assigning to `local_error`, and include the rationale that local failures must become rank-agreed status before the next collective. Do not narrow the exception types or alter the existing error-handling flow.Source: Linters/SAST tools
254-265: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDon't key the fallback path on an exception message string.
Rewording the message at Line 225 silently converts "slab cannot hold a layered record → zero layered capacity" into a hard failure, and nothing in the type system or the tests catches it. A dedicated exception subclass makes the contract explicit.
♻️ Proposed change
+class _LayeredPacketDoesNotFit(ValueError): + """The base slab cannot hold even one three-layer record.""" + + def _largest_layered_capacity_for_slab( @@ if first_layout.slab_bytes > int(slab_bytes): - raise ValueError("base slab cannot hold one three-layer record") + raise _LayeredPacketDoesNotFit( + "base slab cannot hold one three-layer record" + ) @@ - except ValueError as exc: - if str(exc) != "base slab cannot hold one three-layer record": - raise + except _LayeredPacketDoesNotFit: return 0, fallback_layout🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sparkinfer/comm/pcie/pcie_selected_records_ce.py` around lines 254 - 265, Introduce a dedicated exception subclass for the “base slab cannot hold one three-layer record” condition in _largest_layered_capacity_for_slab, raise it for that case, and catch that subclass in the layered_max_records fallback instead of comparing ValueError message text. Preserve the existing (0, fallback_layout) result for this condition while re-raising unrelated ValueError instances.
786-857: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSingle-plane and layered paths are parallel copies that differ only by layout. Every helper and the whole exchange orchestration exist twice, distinguished solely by which
_CopyExchangeLayoutis read and how wide a packet record is. The class docstring's guarantee that separately captured normal and layered CUDA graphs are safe to alternate depends on both copies of the barrier/release chain staying identical, and nothing structurally enforces that today.
sparkinfer/comm/pcie/pcie_selected_records_ce.py#L786-L857: fold_active_layer_primary_capacityinto_active_primary_capacityand the three_layer_*pointer helpers into their single-plane counterparts by takinglayout(and the capacity bound) as arguments.sparkinfer/comm/pcie/pcie_selected_records_ce.py#L943-L1099: express the shared eight-step sequence (stream check, warmup gate, validate, drain, empty-selection short circuit, rotated pack loop, rotated DMA loop, barrier(0), zero, unpack, publish) once, injecting the layout, packet width, and pack/unpack entry points per mode.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sparkinfer/comm/pcie/pcie_selected_records_ce.py` around lines 786 - 857, Unify the duplicated single-plane and layered exchange logic in sparkinfer/comm/pcie/pcie_selected_records_ce.py:786-857 by folding _active_layer_primary_capacity and the _layer_* pointer helpers into parameterized counterparts that accept the layout and capacity bound, and refactor sparkinfer/comm/pcie/pcie_selected_records_ce.py:943-1099 so one shared orchestration implements the eight-step exchange sequence while receiving mode-specific layout, packet width, and pack/unpack entry points; update both callers to use this shared path without changing their behavior.tests/comm/test_pcie_selected_records_ce_gpu.py (1)
620-634: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winA VRAM shortfall on rank 0 fails the build instead of skipping.
The big-offset gate is the guideline-mandated
2^31 / striderepro and should stay, but a machine that simply lacksrequired_bytes + 512 MiBfree now reports a red run for something the code did not break.pytest.skipis unavailable inside themp.spawnworker, so evaluate the capacity gate before spawning (or return a distinguishable skip signal from the worker) rather than raisingRuntimeErroron every rank.Separately,
with suppress(Exception): del recordsat Lines 681-682 is dead — deleting a local bound to either a tensor orNonecannot raise.♻️ Cleanup for the `finally` block
finally: exchange.close() - with suppress(Exception): - del records + del records torch.cuda.empty_cache()Also applies to: 679-683
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/comm/test_pcie_selected_records_ce_gpu.py` around lines 620 - 634, Update the big-offset capacity gate in the test setup around the mp.spawn worker so insufficient rank-0 VRAM skips the test rather than raising RuntimeError on every rank; evaluate capacity before spawning or propagate a distinguishable worker skip result to the parent. Preserve the guideline-mandated 2^31 / BIG_RECORD_BYTES repro when capacity is available, and remove the dead suppress(Exception) deletion of records in the cleanup/finally block.sparkinfer/comm/pcie/pcie_selected_records_ce.cu (2)
27-140: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the publish/wait halves into
__device__helpers.
barrier_all_peers_kernelis a verbatim concatenation ofpublish_all_peers_kernelandwait_all_peers_kernel, so the fence/relaxed-store sequence and the spin/timeout/trap loop each exist three times. The deferred-release protocol only works if the barrier's wait half stays bit-identical towait_all_peers_kernel; keeping one copy of each removes that drift risk without changing math or synchronization semantics.♻️ Sketch: shared device helpers
+__device__ __forceinline__ void publish_peer_flag( + const int64_t* __restrict__ publish_flag_ptrs, + int32_t* __restrict__ send_counters, + int64_t counter_offset) { + const int32_t publish_value = send_counters[counter_offset] + 1; + send_counters[counter_offset] = publish_value; + __threadfence_system(); + auto* publish_flag = reinterpret_cast<int32_t*>( + static_cast<uintptr_t>(publish_flag_ptrs[counter_offset])); + asm volatile( + "st.relaxed.sys.global.u32 [%1], %0;" + : + : "r"(publish_value), "l"(publish_flag)); +} + +__device__ __forceinline__ void wait_peer_flag( + const int64_t* __restrict__ wait_flag_ptrs, + int32_t* __restrict__ wait_counters, + int64_t counter_offset, + uint64_t timeout_cycles, + int phase, + int peer, + const char* what) { /* single copy of the spin/timeout/trap loop */ }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sparkinfer/comm/pcie/pcie_selected_records_ce.cu` around lines 27 - 140, Extract the publish sequence from publish_all_peers_kernel and the wait/timeout sequence from wait_all_peers_kernel into shared __device__ helpers, preserving their existing counter, fence, relaxed-store, acquire-load, timeout, and trap semantics. Update barrier_all_peers_kernel to invoke both helpers in order, and have the standalone kernels reuse the same helpers so the barrier wait path remains identical to wait_all_peers_kernel.
503-566: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCollapse the hand-rolled index/vector dispatch.
The
int32/int64xuint4/uint8_tfan-out is written out four times here and again inpack_compact_record_layers(Lines 680-742), with the full 11-argument call list duplicated in each branch. A small dispatch helper (or anAT_DISPATCH-style macro over the two index types wrapping avectorizedbranch) would keep the argument list in one place and make it harder to pass a mismatched offset to one branch only.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sparkinfer/comm/pcie/pcie_selected_records_ce.cu` around lines 503 - 566, Refactor the dispatch around launch_pack in the current selection path to eliminate the four duplicated int32/int64 × uint4/uint8_t branches, using a small helper or AT_DISPATCH-style structure that keeps the 11 launch arguments in one call site. Apply the same dispatch consolidation to pack_compact_record_layers, preserving the existing vectorized selection and index-type behavior while reusing the correct offsets.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@sparkinfer/comm/pcie/pcie_selected_records_ce.cu`:
- Around line 170-208: The pack kernels accept oversized record indices and can
read beyond the records tensor. In
sparkinfer/comm/pcie/pcie_selected_records_ce.cu lines 170-208, add a shared
record_rows parameter to pack_compact_records_kernel, reject local_record values
below zero or at least record_rows, and derive record_rows from records.numel()
/ record_bytes in launch_pack. Apply the same parameter and guard in
pack_compact_record_layers_kernel at lines 298-341, passing the value from
launch_pack_layers before indexing any plane.
In `@sparkinfer/comm/pcie/pcie_selected_records_ce.py`:
- Around line 1101-1116: Restrict _finish_deferred_release_for_close to explicit
close() calls and remove its invocation from __del__. Make __del__ release only
local resources, while emitting a warning when _deferred_release_pending remains
true; preserve explicit close()’s existing collective drain behavior before
process-group or peer teardown.
In `@sparkinfer/comm/pcie/pcie_selected_records.py`:
- Around line 375-390: Update the CUDA IPC peer probe in the initialization flow
around _gather_ipc_handles to use a dedicated payload-region byte rather than
byte 0 of the barrier flag region for both peer copies. Ensure the remote
cudaMemset initialization is explicitly synchronized before the probe reads
shared memory, while preserving the existing stream synchronization,
rank-success check, cleanup, and error propagation.
In `@tests/comm/test_pcie_selected_records_ce_gpu.py`:
- Around line 303-309: Update the CUDA graph replay blocks in the test,
including the corresponding block around the second occurrence, to invoke
graph.replay() within a torch.cuda.stream(stream) context. Keep
stream.synchronize() and the address assertions unchanged so replay executes on
the channel’s capture stream.
In `@tests/comm/test_pcie_selected_records_ce.py`:
- Around line 166-175: Update the fake unpack reconstruction near the
local_packet selection to use the pack from the current in-flight exchange
rather than pack_calls[0]. Track or reset pack_calls and layer_pack_calls at
each exchange boundary, and apply the same correction to the corresponding logic
around the second referenced region, so repeated exchanges reconstruct from
their own records and indices.
---
Nitpick comments:
In `@sparkinfer/comm/pcie/pcie_selected_records_ce.cu`:
- Around line 27-140: Extract the publish sequence from publish_all_peers_kernel
and the wait/timeout sequence from wait_all_peers_kernel into shared __device__
helpers, preserving their existing counter, fence, relaxed-store, acquire-load,
timeout, and trap semantics. Update barrier_all_peers_kernel to invoke both
helpers in order, and have the standalone kernels reuse the same helpers so the
barrier wait path remains identical to wait_all_peers_kernel.
- Around line 503-566: Refactor the dispatch around launch_pack in the current
selection path to eliminate the four duplicated int32/int64 × uint4/uint8_t
branches, using a small helper or AT_DISPATCH-style structure that keeps the 11
launch arguments in one call site. Apply the same dispatch consolidation to
pack_compact_record_layers, preserving the existing vectorized selection and
index-type behavior while reusing the correct offsets.
In `@sparkinfer/comm/pcie/pcie_selected_records_ce.py`:
- Around line 367-368: Add a targeted `# noqa: BLE001` annotation to each
deliberate broad `except Exception` handler at the referenced locations,
including the handlers assigning to `local_error`, and include the rationale
that local failures must become rank-agreed status before the next collective.
Do not narrow the exception types or alter the existing error-handling flow.
- Around line 254-265: Introduce a dedicated exception subclass for the “base
slab cannot hold one three-layer record” condition in
_largest_layered_capacity_for_slab, raise it for that case, and catch that
subclass in the layered_max_records fallback instead of comparing ValueError
message text. Preserve the existing (0, fallback_layout) result for this
condition while re-raising unrelated ValueError instances.
- Around line 786-857: Unify the duplicated single-plane and layered exchange
logic in sparkinfer/comm/pcie/pcie_selected_records_ce.py:786-857 by folding
_active_layer_primary_capacity and the _layer_* pointer helpers into
parameterized counterparts that accept the layout and capacity bound, and
refactor sparkinfer/comm/pcie/pcie_selected_records_ce.py:943-1099 so one shared
orchestration implements the eight-step exchange sequence while receiving
mode-specific layout, packet width, and pack/unpack entry points; update both
callers to use this shared path without changing their behavior.
In `@sparkinfer/comm/pcie/pcie_selected_records.cu`:
- Around line 100-106: Add a record-count parameter to the relevant CUDA kernel
and pass records.numel() / record_bytes from all four launch sites. In the
local_record handling near local_indices_by_destination, skip entries when
local_record is negative or greater than or equal to the record count before
accessing records, preserving the existing fast path for valid indices.
In `@tests/comm/test_pcie_selected_records_ce_gpu.py`:
- Around line 620-634: Update the big-offset capacity gate in the test setup
around the mp.spawn worker so insufficient rank-0 VRAM skips the test rather
than raising RuntimeError on every rank; evaluate capacity before spawning or
propagate a distinguishable worker skip result to the parent. Preserve the
guideline-mandated 2^31 / BIG_RECORD_BYTES repro when capacity is available, and
remove the dead suppress(Exception) deletion of records in the cleanup/finally
block.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 16230c3e-187e-4f8b-971b-40dbeb5b3915
📒 Files selected for processing (12)
README.mdpyproject.tomlsparkinfer/comm/pcie/__init__.pysparkinfer/comm/pcie/api.pysparkinfer/comm/pcie/pcie_selected_records.cusparkinfer/comm/pcie/pcie_selected_records.pysparkinfer/comm/pcie/pcie_selected_records_ce.cusparkinfer/comm/pcie/pcie_selected_records_ce.pytests/comm/test_pcie_selected_records.pytests/comm/test_pcie_selected_records_ce.pytests/comm/test_pcie_selected_records_ce_gpu.pytests/comm/test_pcie_selected_records_gpu.py
| for (int64_t selected = first_warp; selected < selected_records; | ||
| selected += warp_stride) { | ||
| const int64_t local_record = static_cast<int64_t>(local_indices[selected]); | ||
| if (local_record < 0) { | ||
| continue; | ||
| } | ||
|
|
||
| uint64_t ordinal = 0; | ||
| if (lane == 0) { | ||
| ordinal = atomicAdd( | ||
| reinterpret_cast<unsigned long long*>(primary), | ||
| static_cast<unsigned long long>(1)); | ||
| } | ||
| ordinal = broadcast_lane_zero(ordinal); | ||
| const bool use_primary = ordinal < static_cast<uint64_t>(primary_capacity); | ||
| const int64_t compact_index = use_primary | ||
| ? static_cast<int64_t>(ordinal) | ||
| : static_cast<int64_t>(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<int64_t*>(payload + positions_offset)[compact_index] = | ||
| selected; | ||
| } | ||
|
|
||
| const int64_t units_per_record = record_bytes / sizeof(copy_t); | ||
| const auto* source = reinterpret_cast<const copy_t*>( | ||
| records + record_byte_offset(local_record, record_bytes)); | ||
| auto* destination = reinterpret_cast<copy_t*>( | ||
| 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]; | ||
| } | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Neither pack kernel bounds the caller-supplied record index, allowing out-of-bounds device reads. Both kernels reject only local_record < 0; the record row count is never passed in, and the host cannot validate device-resident indices (it checks only dtype, contiguity and size(-1) == record_bytes). A stale or oversized index therefore reads past the records tensor and either faults or packs garbage into the peer packet. One shared record_rows parameter plus a single upper-bound guard fixes both sites.
sparkinfer/comm/pcie/pcie_selected_records_ce.cu#L170-L208: add arecord_rowsparameter topack_compact_records_kernel, change the guard tolocal_record < 0 || local_record >= record_rows, and deriverecord_rows = records.numel() / record_bytesinlaunch_pack.sparkinfer/comm/pcie/pcie_selected_records_ce.cu#L298-L341: apply the samerecord_rowsparameter and guard inpack_compact_record_layers_kernelbefore indexing any of the three planes, passing it fromlaunch_pack_layers.
📍 Affects 1 file
sparkinfer/comm/pcie/pcie_selected_records_ce.cu#L170-L208(this comment)sparkinfer/comm/pcie/pcie_selected_records_ce.cu#L298-L341
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@sparkinfer/comm/pcie/pcie_selected_records_ce.cu` around lines 170 - 208, The
pack kernels accept oversized record indices and can read beyond the records
tensor. In sparkinfer/comm/pcie/pcie_selected_records_ce.cu lines 170-208, add a
shared record_rows parameter to pack_compact_records_kernel, reject local_record
values below zero or at least record_rows, and derive record_rows from
records.numel() / record_bytes in launch_pack. Apply the same parameter and
guard in pack_compact_record_layers_kernel at lines 298-341, passing the value
from launch_pack_layers before indexing any plane.
| 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) | ||
|
|
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
__del__ can execute a peer-dependent spin-wait that traps the CUDA context.
_finish_deferred_release_for_close launches the device-side wait_all_peers spin, which retires only after every peer publishes phase 1. Reaching it from __del__ — arbitrary GC points, interpreter shutdown, or after a peer already exited / dist.destroy_process_group() ran / peer IPC slabs were freed — stalls the process for the whole barrier_timeout_cycles budget and then hits asm volatile("trap;") (pcie_selected_records_ce.cu Line 136), aborting the context. suppress(Exception) in __del__ cannot recover from that, and the crash surfaces far from its cause. Every current test drains via an explicit close() before process-group teardown, so this path is unexercised.
Suggest restricting the collective drain to explicit close() and having __del__ release only local resources (plus a warning when a release is still pending).
🛡️ Sketch
- def close(self) -> None:
+ def close(self, *, drain_peers: bool = True) -> None:
if self._closed:
return
- self._finish_deferred_release_for_close()
+ if drain_peers:
+ self._finish_deferred_release_for_close()
self._closed = True
self._release_owned_buffer(synchronize=True)
def __del__(self) -> None:
with suppress(Exception):
- self.close()
+ self.close(drain_peers=False)Also applies to: 1140-1142
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@sparkinfer/comm/pcie/pcie_selected_records_ce.py` around lines 1101 - 1116,
Restrict _finish_deferred_release_for_close to explicit close() calls and remove
its invocation from __del__. Make __del__ release only local resources, while
emitting a warning when _deferred_release_pending remains true; preserve
explicit close()’s existing collective drain behavior before process-group or
peer teardown.
| 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 |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Peer probe reads/writes byte 0 of the flag region without an explicit fence against the remote cudaMemset.
cudaMemcpyAsync(local_ptr, remote_ptr, 1, ...) pulls the peer's byte 0 into this rank's byte 0, which is barrier flag slot (phase 0, source 0). Correctness relies on every peer's cudaMemset at Line 311 having completed before any peer reads it — which currently holds only implicitly via legacy-default-stream/NCCL ordering inside _gather_ipc_handles. If that memset ever lands on a non-blocking or per-thread default stream, a non-zero byte can be latched into the flag slot and the very first barrier_all_peers_kernel wait (observed - expected < 0) can pass prematurely. Probing into the payload region and syncing after the memset removes the coupling entirely.
🛡️ Suggested hardening
try:
ipc.cudaMemset(local_ptr, 0, size_in_bytes)
+ if device.type == "cuda":
+ torch.cuda.synchronize(device)
except Exception as exc:
local_error = exc local_error = None
try:
stream = torch.cuda.current_stream(device)
stream_ptr = int(stream.cuda_stream)
+ # Probe outside the barrier-flag region so a torn probe cannot latch a
+ # stale flag value.
+ probe_offset = size_in_bytes - 1
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)
+ ipc.cudaMemcpyAsync(
+ local_ptr + probe_offset, remote_ptr + probe_offset, 1, stream_ptr
+ )
+ ipc.cudaMemcpyAsync(
+ remote_ptr + probe_offset, local_ptr + probe_offset, 1, stream_ptr
+ )
stream.synchronize()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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 | |
| local_error = None | |
| try: | |
| stream = torch.cuda.current_stream(device) | |
| stream_ptr = int(stream.cuda_stream) | |
| # Probe outside the barrier-flag region so a torn probe cannot latch a | |
| # stale flag value. | |
| probe_offset = size_in_bytes - 1 | |
| for remote_ptr in shared.remote_ptrs: | |
| ipc.cudaMemcpyAsync( | |
| local_ptr + probe_offset, remote_ptr + probe_offset, 1, stream_ptr | |
| ) | |
| ipc.cudaMemcpyAsync( | |
| remote_ptr + probe_offset, local_ptr + probe_offset, 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 |
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 383-383: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@sparkinfer/comm/pcie/pcie_selected_records.py` around lines 375 - 390, Update
the CUDA IPC peer probe in the initialization flow around _gather_ipc_handles to
use a dedicated payload-region byte rather than byte 0 of the barrier flag
region for both peer copies. Ensure the remote cudaMemset initialization is
explicitly synchronized before the probe reads shared memory, while preserving
the existing stream synchronization, rank-success check, cleanup, and error
propagation.
| 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 |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
torch.cuda.CUDAGraph replay which stream does it launch on
💡 Result:
When you call torch.cuda.CUDAGraph.replay, the graph is launched on the stream that is current at the time of the replay call [1][2]. In the PyTorch implementation, the replay method retrieves the current stream using at::cuda::getCurrentCUDAStream and passes it to the underlying cudaGraphLaunch API [1][2]. Unlike the capture phase, which explicitly requires a non-default stream, the replay phase can be executed on any stream, including the default stream [1][2]. This design allows you to record a sequence of operations into a graph once and then replay that graph on different streams or at different points in your program's execution, provided you ensure appropriate synchronization if necessary.
Citations:
- 1: https://github.com/pytorch/pytorch/blob/e9ebbd3b/aten/src/ATen/cuda/CUDAGraph.cpp
- 2: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/cuda/CUDAGraph.cpp
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== git status =="
git status --short
echo "== locate file =="
fd -a 'test_pcie_selected_records_ce_gpu.py' . || true
echo "== relevant file excerpt =="
if [ -f tests/comm/test_pcie_selected_records_ce_gpu.py ]; then
wc -l tests/comm/test_pcie_selected_records_ce_gpu.py
sed -n '260,360p' tests/comm/test_pcie_selected_records_ce_gpu.py | cat -n
fi
echo "== search channel stream affinity patterns =="
rg -n "current_stream|CUDAGraph|graph.replay|stream.wait_stream|exchange\\.exchange|set_stream|cuda\\.stream" tests/comm test_pcie_selected_records_ce_gpu.py py spark infer . -g '*.py' 2>/dev/null | head -n 200 || trueRepository: local-inference-lab/sparkinfer
Length of output: 2059
🌐 Web query:
PyTorch CUDAGraph replay stream wait_stream synchronize documentation
💡 Result:
In PyTorch, CUDA Graph operations rely on strict stream and synchronization rules to ensure captured work is deterministic and self-contained [1][2]. The use of wait_stream and synchronization varies significantly between the capture and replay phases. Capture Phase: 1. Stream Requirement: CUDA graphs must be captured on a non-default stream (e.g., a side stream) [3][2][4]. The default stream (stream 0) is prohibited because of its implicit synchronization behavior [2]. 2. Dependency Management: All operations performed during capture must form a tree structure branching from the capturing stream [1][2]. If work is offloaded to a side stream, it must explicitly wait for the capture stream via side_stream.wait_stream(capture_stream) to satisfy CUDA graph’s dependency requirements [2]. 3. Prohibited Synchronizations: You cannot perform CPU-GPU synchronization (e.g., torch.cuda.synchronize,.item, or.cpu) during capture [1][2]. These operations block the CPU and cannot be recorded in the graph, as the graph expects to execute entirely on the GPU during replay [2]. Replay Phase: 1. Stream Flexibility: After capture, a CUDAGraph can be replayed on any stream, including the default stream [4][5]. 2. Asynchronous Execution: The replay call is asynchronous; it enqueues the captured graph work onto the current stream and returns immediately [6][7]. 3. Synchronization Hooks: Because replay is asynchronous, if you need to perform actions (such as freeing device memory) that depend on the graph completion, you should use registration hooks like register_replay_end_hook(hook) [6][7]. 4. Teardown Safety: PyTorch documentation notes that teardown does not automatically synchronize CUDA [6]. When destroying a graph, if you have in-flight replays, you must ensure all relevant streams are synchronized (e.g., pass synchronize_before_release=True if using specific APIs) to avoid use-after-free errors [6][7]. In summary, wait_stream is essential during the capture phase to define dependencies between streams, while synchronization (torch.cuda.synchronize) is generally forbidden during capture but necessary to manage memory and execution order safely during or after the asynchronous replay process [1][6][2].
Citations:
- 1: https://docs.nvidia.com/dl-cuda-graph/cuda-graph-basics/constraints.html
- 2: https://docs.nvidia.com/dl-cuda-graph/latest/troubleshooting/capture-failures.html
- 3: https://docs.nvidia.com/dl-cuda-graph/torch-cuda-graph/torch-integration.html
- 4: https://github.com/pytorch/pytorch/blob/e9ebbd3b/aten/src/ATen/cuda/CUDAGraph.cpp
- 5: > if graph capture is thread local pytorch/pytorch#137844
- 6: http://docs.pytorch.org/docs/main/generated/torch.cuda.graphs.CUDAGraph.html
- 7: http://docs.pytorch.org/docs/main/generated/torch.cuda.CUDAGraph.html
🌐 Web query:
CUDAGraph replay current stream PyTorch implementation at::cuda::getCurrentCUDAStream
💡 Result:
In PyTorch, the relationship between CUDA Graph replay and the "current stream" as retrieved by at::cuda::getCurrentCUDAStream is governed by the underlying CUDA Graph execution API [1][2]. When you call the replay method on a CUDAGraph instance, PyTorch executes the captured graph using the stream returned by at::cuda::getCurrentCUDAStream at the time of the replay call [1][2]. Specifically, the implementation in the PyTorch C++ ATen library explicitly calls cudaGraphLaunch(graph_exec_, at::cuda::getCurrentCUDAStream) to enqueue the graph for execution on that stream [1][2]. Key technical points: 1. Capture vs. Replay Stream: While CUDA graphs must be captured on a non-default stream [1][2], they are designed to be replayed on any valid CUDA stream [1]. 2. Stream Dependency: Because replay relies on the current stream of the calling thread, it is essential to ensure that the thread-local current stream is correctly set—often using a StreamGuard or similar mechanism—before invoking replay if you intend to execute the graph on a specific stream [3][4]. 3. Execution Context: The use of at::cuda::getCurrentCUDAStream within CUDAGraph::replay ensures that the graph is launched asynchronously on the currently active stream, maintaining compatibility with existing stream-based synchronization patterns in your application [1][2]. Developers should be aware that if the current stream is not explicitly managed, it will default to the device's current stream (often the default stream) [3][4], which may not be appropriate for all high-performance or concurrent workloads [1][2].
Citations:
- 1: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/cuda/CUDAGraph.cpp
- 2: https://github.com/pytorch/pytorch/blob/e9ebbd3b/aten/src/ATen/cuda/CUDAGraph.cpp
- 3: https://docs.pytorch.org/cppdocs/api/cuda/utilities.html
- 4: https://github.com/pytorch/pytorch/blob/main/c10/cuda/CUDAStream.h
Replay under the channel’s CUDA stream
CUDAGraph.replay() launches on the current stream and does not inherit the capture stream. Since the channel affinity is tied to stream, replaying on the default stream leaves stream.wait_stream() empty for the graph execution. Replay in a torch.cuda.stream(stream) context, or otherwise establish the proper wait_stream(...) dependency before replay.
Also applies to: 336-342.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/comm/test_pcie_selected_records_ce_gpu.py` around lines 303 - 309,
Update the CUDA graph replay blocks in the test, including the corresponding
block around the second occurrence, to invoke graph.replay() within a
torch.cuda.stream(stream) context. Keep stream.synchronize() and the address
assertions unchanged so replay executes on the channel’s capture stream.
| # 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]) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
The fake unpack always reconstructs from the first exchange's pack, not the current one.
pack_calls and layer_pack_calls are never cleared, so pack_calls[0] / layer_pack_calls[0] keep pointing at the very first exchange. test_exchange_preserves_persistent_state_and_accepts_int32_indices (Line 454) re-runs with records.flip(0) and would be reconstructed from the original records; it only passes because the second output is never asserted. Track the packs of the in-flight exchange instead so repeated-exchange reconstruction is actually covered.
💚 Proposed fix
self.operations = []
self._packets = {}
+ self._pending_packs = []
+ self._pending_layer_packs = []
@@
self.pack_calls.append(call)
+ self._pending_packs.append(call)
self.operations.append(("pack", primary_ptr))
@@
self.layer_pack_calls.append(call)
+ self._pending_layer_packs.append(call)
self.operations.append(("pack_layers", primary_ptr))
@@
- local_packet = self.pack_calls[0]
+ local_packet = self._pending_packs[0]
+ self._pending_packs.clear()
@@
- local_packet = self.layer_pack_calls[0]
+ local_packet = self._pending_layer_packs[0]
+ self._pending_layer_packs.clear()Also applies to: 204-215
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/comm/test_pcie_selected_records_ce.py` around lines 166 - 175, Update
the fake unpack reconstruction near the local_packet selection to use the pack
from the current in-flight exchange rather than pack_calls[0]. Track or reset
pack_calls and layer_pack_calls at each exchange boundary, and apply the same
correction to the corresponding logic around the second referenced region, so
repeated exchanges reconstruct from their own records and indices.
2 participants
Summary
Add a copy-engine selected-record exchange and a three-layer bulk packet API.
The CE path compacts locally owned records, preserves final destination slots, rank-rotates destination order, transfers contiguous primary packets with asynchronous device-to-device copies, and unpacks at the destination. Owner skew beyond the primary capacity uses a bounded IPC overflow region.
exchange_layers()moves exactly three record planes with one destination map and one synchronization/release chain, matching Full -> S1/S2/S3 sparse-MLA lookahead.Stack note
This PR is stacked on #64. GitHub therefore shows the direct exchange prerequisite until that PR merges. The new review unit is
0495cbe.Contract and fallback
Validation
Tests cover:
The focused direct/CE CPU transport suite passes 49 tests. Native 368-byte CE transport passed the four-GPU eager/graph byte-equivalence gate and the FP8-RoPE model integration.
Design
GLM-5.2 Sparse CKV under Decode Context Parallelism
Credits
The sparse CKV and bulk shared-layer design originates from Koush's work and builds on Luke Alonso's B12X/Sparkinfer infrastructure.
Implementation, tests, and documentation were developed with OpenAI Codex assistance and manually reviewed and GPU-validated.
Summary by CodeRabbit
New Features
uint8record data.Documentation
Tests