perf(pcie): add exact DCP top-k owner exchange#79
Conversation
📝 WalkthroughWalkthroughAdds a PCIe/CUDA-IPC owner exchange for exact DCP sparse top-k candidates, including public exports, shared-buffer lifecycle management, synchronized CUDA staging, validation, CPU tests, GPU integration tests, and packaging updates. ChangesDCP TopK owner exchange
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant PCIeDCPTopKOwnerExchange
participant stage_candidates
participant stage_owner_candidates
participant IPCStagingBuffers
PCIeDCPTopKOwnerExchange->>stage_candidates: submit local indices and scores
stage_candidates->>stage_owner_candidates: launch CUDA staging
stage_owner_candidates->>IPCStagingBuffers: write owner candidate planes
stage_owner_candidates-->>stage_candidates: return staging slot
stage_candidates-->>PCIeDCPTopKOwnerExchange: return candidate tensors
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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.
Actionable comments posted: 1
🧹 Nitpick comments (3)
sparkinfer/comm/pcie/pcie_dcp_topk.py (1)
274-288: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winConsider asserting cross-rank agreement on
max_rows/topk.
layout.slab_bytesis derived from rank-localmax_rows/topk, and the slab allocation is collective. If one rank is constructed with a different geometry, every rank's peer pointers still map successfully but the owner planes disagree, which surfaces later as silent candidate corruption rather than a construction error. A small all-reduce/broadcast check of(max_rows, topk, ext.meta_size())before allocating would fail fast.🤖 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_topk.py` around lines 274 - 288, Before the collective allocation in the candidate staging setup, add a cross-rank agreement check for max_rows, topk, and ext.meta_size(). Use the existing exchange_group collective communication utilities to validate these values match on every rank, and fail construction immediately with a clear error before calling _allocate_shared_buffer when they differ.sparkinfer/comm/pcie/pcie_dcp_topk.cu (1)
66-89: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winDocument (or enforce) the grid-symmetry invariant this barrier depends on.
block_pair_barrieronly completes if every rank in the DCP group launched the identical number of blocks, since blockiwaits on peer blocki's counter slot.blocksis computed per-rank from caller-suppliedrows/threads/block_limit(Line 181), so any asymmetry silently deadlocks the whole group with no timeout or diagnostic. Consider stating the invariant in the header comment and having the Python layer assert agreed-uponrows/threads/block_limitacross the exchange group at construction time.🤖 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_topk.cu` around lines 66 - 89, Document the grid-symmetry invariant required by block_pair_barrier: every rank in the DCP group must launch the same number of blocks because block indices synchronize directly. Add a clear header comment near block_pair_barrier, and update the Python construction/exchange-group validation to assert matching rows, threads, and block_limit across all ranks before launching work.tests/comm/test_pcie_dcp_topk_gpu.py (1)
126-160: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftAdd a large-offset case and a back-to-back replay case.
With
MAX_ROWS=64/TOPK=2048the largest staged element offset is ~1.3e5, so the kernel's Int64 offset math (owner_row * output_row_packs,int64_t(destination) * owner_packs) is never exercised anywhere near the2^31 / strideboundary — a future 32-bit narrowing would pass this suite. A capacity-scaled case (largetopk/max_rows, gated on available memory) would close that gap.Separately,
torch.cuda.synchronize(device)after everygraph.replay()means two consecutive replays never overlap, which is precisely the pattern that would expose the pinned-slot hazard noted insparkinfer/comm/pcie/pcie_dcp_topk.cuLine 184. Consider a replay pair without an intervening full sync.As per coding guidelines: "Every test or repro for a kernel indexing a paged pool must include a big-pid case with live pages beyond the
2^31 / strideboundary."🤖 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_dcp_topk_gpu.py` around lines 126 - 160, Add coverage in the graph replay test around `_inputs` and `owner_stage_reference` for a memory-gated, capacity-scaled large TOPK/MAX_ROWS configuration whose live page IDs exceed the 2^31/stride boundary, while preserving the existing assertions. Also add a back-to-back replay pair that calls `graph.replay()` twice without an intervening `torch.cuda.synchronize(device)`, then synchronize and validate both outputs to exercise the pinned-slot hazard.Source: Coding guidelines
🤖 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_dcp_topk.cu`:
- Around line 181-185: Update the slot selection around slot_++ so the counter
uses a non-overflowing representation with explicit two-slot wrapping instead of
signed modulo arithmetic. Ensure captured CUDA graph replays obtain a fresh
staging slot only after the owner consumer retires the prior output, preserving
the two-slab handoff and preventing writes from clobbering live staging data.
---
Nitpick comments:
In `@sparkinfer/comm/pcie/pcie_dcp_topk.cu`:
- Around line 66-89: Document the grid-symmetry invariant required by
block_pair_barrier: every rank in the DCP group must launch the same number of
blocks because block indices synchronize directly. Add a clear header comment
near block_pair_barrier, and update the Python construction/exchange-group
validation to assert matching rows, threads, and block_limit across all ranks
before launching work.
In `@sparkinfer/comm/pcie/pcie_dcp_topk.py`:
- Around line 274-288: Before the collective allocation in the candidate staging
setup, add a cross-rank agreement check for max_rows, topk, and ext.meta_size().
Use the existing exchange_group collective communication utilities to validate
these values match on every rank, and fail construction immediately with a clear
error before calling _allocate_shared_buffer when they differ.
In `@tests/comm/test_pcie_dcp_topk_gpu.py`:
- Around line 126-160: Add coverage in the graph replay test around `_inputs`
and `owner_stage_reference` for a memory-gated, capacity-scaled large
TOPK/MAX_ROWS configuration whose live page IDs exceed the 2^31/stride boundary,
while preserving the existing assertions. Also add a back-to-back replay pair
that calls `graph.replay()` twice without an intervening
`torch.cuda.synchronize(device)`, then synchronize and validate both outputs to
exercise the pinned-slot hazard.
🪄 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: f9b9e671-002e-4146-b387-d6a3c1a78201
📒 Files selected for processing (7)
sparkinfer/comm/pcie/__init__.pysparkinfer/comm/pcie/api.pysparkinfer/comm/pcie/pcie_dcp_topk.cusparkinfer/comm/pcie/pcie_dcp_topk.pytests/comm/test_pcie_dcp_topk.pytests/comm/test_pcie_dcp_topk_gpu.pytests/test_packaging.py
|
Review follow-ups
Validation:
|
1 participant
Summary
Add a graph-safe CUDA-IPC transport that stages exact sparse-indexer candidates directly into each DCP row owner's rank-major input table.
Each source rank writes its disjoint
(indices, FP32 score bits)column into the destination owner's double-buffered slab. The owner can run the existing exact row-top-k without a pack tensor, NCCL all-to-all, or unpack kernel.Supported DCP world sizes are
2, 3, 4, 6, 8.Scope
This PR intentionally contains only candidate owner exchange.
An exact int24 final-index gather was implemented and measured during development, but it took 3.260 ms versus 1.254 ms for NCCL and regressed 64k E2E prefill by 5.45%. It has therefore been removed rather than carried as unused code. Final selected indices remain a standard NCCL all-gather in the vLLM integration.
Results
TP8/DCP4 production-shape candidate phase:
After adding the cross-rank graph-replay reuse guard, the same eager shape
measured 1.0338 ms versus 1.0336 ms before the guard (+0.013%, noise). The
guard therefore has no measurable eager-path cost.
GLM-5.2 64k E2E prefill was 5,766.3 tok/s with NCCL owner staging and 5,784.7 tok/s with this path (+0.32%, within run variance). The persistent double buffer reduced KV capacity by 10,240 tokens (0.47%). This is therefore a lower-traffic optional transport, not an E2E throughput claim.
Double buffering is required: a single payload can be overwritten by a faster producer while a peer still consumes the previous layer.
Validation
6 passedCPU contract and packaging tests.rtol=0, atol=0.queues four replays without host synchronization; every generation remains
exact.
ruff checkandruff format --checkpass.The corresponding vLLM owner algorithm is tracked independently in local-inference-lab/vllm#178.