fix(pcie): retire capture channel aliases - #103
Conversation
📝 WalkthroughWalkthroughPCIe DCP A2A, one-shot, and two-shot collectives now select alternating staging slots on the device. Capture lifecycle logic restores eager channel mappings after nested graph capture. GPU tests inspect staging markers and validate CUDA Graph replay alternation. ChangesPCIe staging and capture behavior
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Capture
participant PCIeCollective
participant DeviceSelector
participant CUDAGraph
Capture->>PCIeCollective: Capture collective with dual staging metadata
PCIeCollective->>CUDAGraph: Record kernel arguments
CUDAGraph->>DeviceSelector: Replay kernel launch
DeviceSelector->>DeviceSelector: Coordinate generation and select slot
DeviceSelector-->>PCIeCollective: Publish selected staging pointers
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
sparkinfer/comm/pcie/pcie_dcp_a2a.cu (1)
176-186: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winIdentical staging-slot-selection block duplicated across both kernels.
The 7-line block (shared
stagingdecl, thread-0 parity load + peer-pointer copy,__syncthreads()) is byte-for-byte identical indcp_lse_reduce_kernel(Lines 176-186) andall_gather_heads_kernel(Lines 317-327). Since this logic encodes the correctness-critical parity contract (which slot a given execution must use), keeping it in one place would reduce the risk of the two copies silently diverging under future edits.♻️ Proposed extraction into a shared device helper
+template <int world_size> +DINLINE void select_staging(RankStaging &staging, + const DoubleStaging &staging_options, + Signal *self) { + if (threadIdx.x == 0) { + const int slot = + int(load_flag(&self->self_counter[blockIdx.x][0]) & FlagType{1}); +#pragma unroll + for (int peer = 0; peer < world_size; ++peer) { + staging.ptrs[peer] = staging_options.slots[slot].ptrs[peer]; + } + } + __syncthreads(); +}Then in each kernel:
__shared__ RankStaging staging; - if (threadIdx.x == 0) { - const int slot = - int(load_flag(&self->self_counter[blockIdx.x][0]) & FlagType{1}); -#pragma unroll - for (int peer = 0; peer < world_size; ++peer) { - staging.ptrs[peer] = staging_options.slots[slot].ptrs[peer]; - } - } - __syncthreads(); + select_staging<world_size>(staging, staging_options, self);Also applies to: 317-327
🤖 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_dcp_a2a.cu` around lines 176 - 186, Extract the duplicated staging-slot-selection logic from dcp_lse_reduce_kernel and all_gather_heads_kernel into one shared device helper, including the shared staging declaration, parity-based slot selection, peer-pointer copy, and synchronization. Update both kernels to call the helper while preserving the existing staging_options, self_counter, blockIdx.x, and world_size behavior.sparkinfer/comm/pcie/pcie_dcp_a2a.py (1)
767-805: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd nested-capture restoration coverage. Sequential capture tests already validate eager mapping restoration, but
capture()allows recursive use fromfor_stream()and its restore logic removes every matching alias independently; add a regression test that nestswith pool.capture(...):and asserts the inner alias is fully unwound before the outer one restores.🤖 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_dcp_a2a.py` around lines 767 - 805, The capture restoration tests lack coverage for recursively nested pool.capture contexts and alias cleanup. Add a regression test using nested with pool.capture(...) blocks that verifies the inner capture’s stream alias is removed or restored before the outer context exits, while preserving the outer mapping until its own restoration; target the capture method and existing eager mapping restoration test setup.
🤖 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.
Nitpick comments:
In `@sparkinfer/comm/pcie/pcie_dcp_a2a.cu`:
- Around line 176-186: Extract the duplicated staging-slot-selection logic from
dcp_lse_reduce_kernel and all_gather_heads_kernel into one shared device helper,
including the shared staging declaration, parity-based slot selection,
peer-pointer copy, and synchronization. Update both kernels to call the helper
while preserving the existing staging_options, self_counter, blockIdx.x, and
world_size behavior.
In `@sparkinfer/comm/pcie/pcie_dcp_a2a.py`:
- Around line 767-805: The capture restoration tests lack coverage for
recursively nested pool.capture contexts and alias cleanup. Add a regression
test using nested with pool.capture(...) blocks that verifies the inner
capture’s stream alias is removed or restored before the outer context exits,
while preserving the outer mapping until its own restoration; target the capture
method and existing eager mapping restoration test setup.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 49521d6b-65fb-4f1f-8cea-973dda78003d
📒 Files selected for processing (4)
sparkinfer/comm/pcie/pcie_dcp_a2a.cusparkinfer/comm/pcie/pcie_dcp_a2a.pytests/comm/test_pcie_dcp_a2a.pytests/comm/test_pcie_dcp_a2a_gpu.py
…60730' into HEAD # Conflicts: # sparkinfer/comm/pcie/pcie_dcp_a2a.cu
|
Addressed both CodeRabbit nitpicks in
Current-head verification:
No live multi-GPU replay result is claimed; the production service was not rebuilt or restarted. |
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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_oneshot.cu`:
- Line 691: Replace host-side staging-slot selection via
dbuf_slot_selector_.take() at sparkinfer/comm/pcie/pcie_oneshot.cu:691 and :768,
and the reduce-scatter/all-gather selections at
sparkinfer/comm/pcie/pcie_twoshot.cu:319 and :335, with replay-safe device-owned
parity selection while passing both eager slabs; alternatively reject graph
capture for these paths. Preserve CUDA graph capture/replay support as the
required behavior.
🪄 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: 4cb63097-54d1-497b-a464-be28fb05da3b
📒 Files selected for processing (7)
sparkinfer/comm/pcie/pcie_dcp_a2a.cusparkinfer/comm/pcie/pcie_oneshot.cusparkinfer/comm/pcie/pcie_twoshot.cusparkinfer/comm/pcie/two_slot_selector.htests/comm/test_pcie_dcp_a2a.pytests/comm/test_pcie_dcp_a2a_gpu.pytests/comm/test_pcie_two_slot_selector.py
|
Final stacked-head update at
Verification on the current stack: 44 passed, 2 opt-in GPU skips; DCP A2A, one-shot, and two-shot CUDA extensions compiled and loaded together from a fresh cache; six changed Python files pass Ruff format/lint. Multi-GPU execution is owned by the separate field-test agent; no rental/production result is claimed here. @coderabbitai review |
|
✅ Action performedReview finished.
|
Rental-machine validation — FAIL (strict safety review supersedes earlier PASS)Public PR head: The previous comment proved that the corrected scratch-reuse oracle passes, but a strict synchronization/lifecycle review found P1 defects outside that oracle:
Useful positive evidence remains valid but is not acceptance evidence: DCP2/DCP4 basic numerical smoke, fused add+RMS graph at 2/4 ranks, fresh first-use capture at 2/4 ranks, 1,025 four-rank opposite-order and 17-collective replay tests, and two-shot correctness. The first fused graph run also exposed and retained a CUDA 13.3 Python-binding harness incompatibility; the version-neutral edge/geometry inspector then passed. Immutable evidence and chronology:
A replacement now uses one operation-wide same-stream DCP control node plus collective setup rollback, retryable teardown, GC retention and exact group ordering. I will not push it until fresh-cache deterministic large↔small skewed DCP2/DCP4 and independent review both pass. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
tests/comm/test_pcie_twoshot.py (2)
65-80: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd assertion messages to the slot checks.
Both assertions fail without context.
tests/comm/test_pcie_oneshot_torture.pyline 151 includes the observed slot list in the message. Includesnapshotsandchanged_slotshere for the same diagnostic value.♻️ Proposed message additions
- assert all(len(changed) == 1 for changed in changed_slots) - assert all( - changed_slots[index] != changed_slots[index + 1] - for index in range(len(changed_slots) - 1) - ) + assert all(len(changed) == 1 for changed in changed_slots), ( + f"each operation must change exactly one staging slot: {changed_slots} " + f"from {snapshots}" + ) + assert all( + changed_slots[index] != changed_slots[index + 1] + for index in range(len(changed_slots) - 1) + ), f"staging slots did not alternate: {changed_slots}"🤖 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_twoshot.py` around lines 65 - 80, Update _assert_alternating_slots so both assertions include diagnostic messages containing snapshots and changed_slots, matching the contextual style used by the oneshot torture test; preserve the existing assertion conditions.
25-62: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDerive the staging layout from the production module.
_local_staging_wordsrecomputespack_stride,scale_offset,scale_stride,slot_bytes, andsignal_byteswith the same formulas assparkinfer/comm/pcie/pcie_twoshot.pyfrom_exchange_group. The test also hardcodes the alignment256while the module usesIPC_SLAB_ALIGNMENT. If the layout or the alignment constant changes, this helper reads the wrong offset and the slot assertions become meaningless instead of failing.Expose the layout computation from
sparkinfer/comm/pcie/pcie_twoshot.pyand call it here, or at least importIPC_SLAB_ALIGNMENTand_align_up.🤖 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_twoshot.py` around lines 25 - 62, Update _local_staging_words to reuse the production staging-layout computation from from_exchange_group in pcie_twoshot.py, including IPC_SLAB_ALIGNMENT and _align_up, instead of duplicating pack_stride, scale_offset, scale_stride, slot_bytes, and signal_bytes formulas. Expose a suitable layout helper if needed, then use its returned offsets and sizes so the test remains synchronized with production changes.tests/comm/test_pcie_oneshot_fused_rmsnorm_gpu.py (1)
75-91: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueShare
_local_eager_wordsinstead of copying it.
tests/comm/test_pcie_oneshot_torture.pylines 46-60 contain the same function body. Move it into a shared test helper module so both tests read the staging markers through one implementation.🤖 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_oneshot_fused_rmsnorm_gpu.py` around lines 75 - 91, The helper _local_eager_words is duplicated across the PCIe oneshot tests. Move its implementation into a shared test helper module, then import and reuse that helper in both test_pcie_oneshot_fused_rmsnorm_gpu.py and test_pcie_oneshot_torture.py, removing the local definitions while preserving its current behavior.tests/comm/test_pcie_oneshot_torture.py (1)
120-151: 🩺 Stability & Availability | 🔵 TrivialThe reported deterministic failure is in this function.
The PR description reports that one-shot eager, graph, and multistream replay tests fail deterministically in
_run_graph_scratch_reuseon rented multi-GPU hardware, and that the head should not merge as-is. The new marker logic here is self-consistent: 17 layers alternate slots, so the last two layer markers must exchange positions on each replay.Record the failing assertion text and the observed
final_slotsorsnapshotvalues against the current head. That evidence separates a kernel parity defect from a test expectation defect. I can prepare a reduced two-layer repro that isolates one device selection per replay, if that helps.🤖 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_oneshot_torture.py` around lines 120 - 151, Run the failing one-shot eager, graph, and multistream replay cases on the current head, especially _run_graph_scratch_reuse, and capture the exact assertion text plus observed final_slots and snapshot values from the marker checks. Do not change the alternating-slot expectation yet; use the collected evidence to determine whether the failure is in kernel parity or the test expectation.
🤖 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_oneshot.cu`:
- Around line 199-220: The staged selectors require full grid residency and a
zero-entry staging counter. In sparkinfer/comm/pcie/pcie_oneshot.cu:199-220, add
a host-side occupancy-based residency limit using
cudaOccupancyMaxActiveBlocksPerMultiprocessor and multiProcessorCount, then
clamp blocks for the SPARKINFER_PCIE_ONESHOT_BLOCK_LIMIT and rows * ctas_per_row
launch paths. In sparkinfer/comm/pcie/pcie_twoshot.cu:97-118, apply the same
clamp in reduce_scatter and all_gather, including the default block_limit of 64
with 512 threads, and document that staging_arrive requires single-stream
ownership.
---
Nitpick comments:
In `@tests/comm/test_pcie_oneshot_fused_rmsnorm_gpu.py`:
- Around line 75-91: The helper _local_eager_words is duplicated across the PCIe
oneshot tests. Move its implementation into a shared test helper module, then
import and reuse that helper in both test_pcie_oneshot_fused_rmsnorm_gpu.py and
test_pcie_oneshot_torture.py, removing the local definitions while preserving
its current behavior.
In `@tests/comm/test_pcie_oneshot_torture.py`:
- Around line 120-151: Run the failing one-shot eager, graph, and multistream
replay cases on the current head, especially _run_graph_scratch_reuse, and
capture the exact assertion text plus observed final_slots and snapshot values
from the marker checks. Do not change the alternating-slot expectation yet; use
the collected evidence to determine whether the failure is in kernel parity or
the test expectation.
In `@tests/comm/test_pcie_twoshot.py`:
- Around line 65-80: Update _assert_alternating_slots so both assertions include
diagnostic messages containing snapshots and changed_slots, matching the
contextual style used by the oneshot torture test; preserve the existing
assertion conditions.
- Around line 25-62: Update _local_staging_words to reuse the production
staging-layout computation from from_exchange_group in pcie_twoshot.py,
including IPC_SLAB_ALIGNMENT and _align_up, instead of duplicating pack_stride,
scale_offset, scale_stride, slot_bytes, and signal_bytes formulas. Expose a
suitable layout helper if needed, then use its returned offsets and sizes so the
test remains synchronized with production changes.
🪄 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: 68d2ebcd-2210-413b-95a2-16c2c640e949
📒 Files selected for processing (8)
sparkinfer/comm/pcie/pcie_dcp_a2a.pysparkinfer/comm/pcie/pcie_oneshot.cusparkinfer/comm/pcie/pcie_twoshot.cutests/comm/test_pcie_dcp_a2a.pytests/comm/test_pcie_dcp_a2a_gpu.pytests/comm/test_pcie_oneshot_fused_rmsnorm_gpu.pytests/comm/test_pcie_oneshot_torture.pytests/comm/test_pcie_twoshot.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tests/comm/test_pcie_dcp_a2a_gpu.py
- tests/comm/test_pcie_dcp_a2a.py
- sparkinfer/comm/pcie/pcie_dcp_a2a.py
| template <int ngpus> | ||
| DINLINE void select_rank_data(RankData& selected, const DoubleRankData& options, Signal* self_sg) { | ||
| if (threadIdx.x == 0) { | ||
| // Every supported staging grid fits concurrently on the target Blackwell | ||
| // GPU. The last arriving block publishes one launch-wide generation, so | ||
| // variable grid sizes cannot split a launch across two slabs. | ||
| const FlagType generation = ld_flag_acquire_gpu(&self_sg->staging_generation); | ||
| const FlagType prior = atomicAdd(&self_sg->staging_arrive, FlagType{1}); | ||
| if (prior == static_cast<FlagType>(gridDim.x - 1)) { | ||
| self_sg->staging_arrive = 0; | ||
| __threadfence(); | ||
| atomicAdd(&self_sg->staging_generation, FlagType{1}); | ||
| } else { | ||
| while (ld_flag_acquire_gpu(&self_sg->staging_generation) == generation) { | ||
| } | ||
| } | ||
| const RankData* source = options.slots[generation & FlagType{1}]; | ||
| #pragma unroll | ||
| for (int peer = 0; peer < ngpus; ++peer) selected.ptrs[peer] = source->ptrs[peer]; | ||
| } | ||
| __syncthreads(); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Unbounded launch-wide rendezvous in both selectors. Both files implement the same protocol: each block's thread 0 increments staging_arrive, and every block except the last arriver spins on staging_generation with no exit path. The shared root cause is that correctness depends on two unchecked preconditions — full grid residency and a zero staging_arrive at entry — asserted only in comments.
sparkinfer/comm/pcie/pcie_oneshot.cu#L199-L220: add a host-side residency check before the staged launches (cudaOccupancyMaxActiveBlocksPerMultiprocessortimesmultiProcessorCount) and clampblocks, including theSPARKINFER_PCIE_ONESHOT_BLOCK_LIMITandrows * ctas_per_rowpaths.sparkinfer/comm/pcie/pcie_twoshot.cu#L97-L118: apply the same residency clamp inreduce_scatterandall_gather, whereblock_limitdefaults to 64 with 512 threads, and document the single-stream ownership requirement for the sharedstaging_arrivecounter.
📍 Affects 2 files
sparkinfer/comm/pcie/pcie_oneshot.cu#L199-L220(this comment)sparkinfer/comm/pcie/pcie_twoshot.cu#L97-L118
🤖 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_oneshot.cu` around lines 199 - 220, The staged
selectors require full grid residency and a zero-entry staging counter. In
sparkinfer/comm/pcie/pcie_oneshot.cu:199-220, add a host-side occupancy-based
residency limit using cudaOccupancyMaxActiveBlocksPerMultiprocessor and
multiProcessorCount, then clamp blocks for the
SPARKINFER_PCIE_ONESHOT_BLOCK_LIMIT and rows * ctas_per_row launch paths. In
sparkinfer/comm/pcie/pcie_twoshot.cu:97-118, apply the same clamp in
reduce_scatter and all_gather, including the default block_limit of 64 with 512
threads, and document that staging_arrive requires single-stream ownership.
Source: Coding guidelines
Frozen replacement qualification — PASS (candidate only)This is not a PASS for the current public #103 head Frozen source: Host: Vast.ai Exact commands and resultsWith
The first combined rejection invocation accidentally omitted Paired vLLM integration used exact vLLM head Collect-only found exactly 64 nodes; execution was 64 passed in 102.17 s. Raw logs remain on the active rental under
Conclusion: the frozen replacement passes the relevant live DCP4, |
|
Superseding field result: the frozen semantic-channel successor is not yet The focused four-rank eager/capture/torture and atomic vLLM/SparkInfer suites The engine failed before KV allocation/health, so there is no performance or The complete call trace makes the repair direction deliberately narrow:
Re-qualification will include the same full-model boot/traffic/shutdown gate, Complete immutable server log and analysis: |
|
Closing as superseded, not merged. Public head 8fd90a4 remains rejected for standalone promotion. The reviewed replacement is #105 at bc62980, paired atomically with field vLLM f99e1e7b8636ca3811ab6d23084ac6da63420dc3 and its patch-identical clean upstream PR local-inference-lab/vllm#216. The final corrected pair passed the complete GLM gate; evidence and exact reproduction contract: https://github.com/malaiwah/glm52-exl3-vast/blob/5c76a2536e7fc9a5f1cb6bf182531889f5385e65/docs/field-review-results/2026-07-30-vast-46335896/UPSTREAM-REPAIR-CAMPAIGN.md and https://github.com/malaiwah/glm52-exl3-vast/blob/5c76a2536e7fc9a5f1cb6bf182531889f5385e65/docs/field-review-results/2026-07-30-vast-46335896/COUNTER-VALIDATION.md. Do not cherry-pick #103 alone and call it equivalent. |
1 participant
Summary
Retire graph-owned DCP channel aliases when capture exits and derive DCP staging-slot ownership from device execution state so CUDA graph replay advances the slot.
Fixes #95. Also incorporates #101 so both reviewed PCIe corrections coexist on one stack.
Scope after re-review
Issue #95's original odd-operation corruption claim was retracted: MLA's mandatory TP reduction closes that proposed graph-boundary window, and AG/RS pairing is structural. Two defects remain actionable:
_channelsretained aliases written for torch-owned nested capture streams after the context exited. A recycled stream key could hand a later unwrapped capture a channel owned by a previous graph.Change
_all_channelsfor captured-kernel replay lifetime.Invariants
Verification
No rental or production-GPU result is claimed here. Production was not rebuilt or restarted.
Review follow-up
22c11df: extracted the correctness-critical DCP staging selection into one shared device helper and added nestedpool.capture()restoration coverage.7e17517: stacked fix(pcie): prevent two-slot selector overflow #101's execution-owned one-shot/two-shot selection.25b7810: removed the obsolete host selector after DCP migration and applied changed-file formatting.5c4f8b0: folded reusable selection into existing kernels, removing the reviewed extra GPU launch.Stack
Merge #101 first. This branch contains #101 plus the DCP execution-owned parity and channel-lifetime commits; after #101 lands, GitHub will narrow this PR to the DCP-specific delta.
Summary by CodeRabbit
Bug Fixes
Tests