[SparkInfer] Account persistent PCIe DMA output in KV profiling#76
Open
voipmonitor wants to merge 1 commit into
Open
[SparkInfer] Account persistent PCIe DMA output in KV profiling#76voipmonitor wants to merge 1 commit into
voipmonitor wants to merge 1 commit into
Conversation
Contributor
📝 WalkthroughWalkthroughChangesPersistent PCIe DMA output workspace
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Caller
participant PCIeDmaAllReduce
participant OutputStorage
participant PersistentOutputView
Caller->>PCIeDmaAllReduce: all_reduce(inp, out=None)
PCIeDmaAllReduce->>OutputStorage: _ensure_output_storage()
OutputStorage-->>PCIeDmaAllReduce: persistent uint8 storage
PCIeDmaAllReduce->>PersistentOutputView: create output view
PersistentOutputView-->>Caller: input-shaped output tensor
🚥 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
sparkinfer/comm/pcie/pcie_dma.py (1)
145-160: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winValidate the actual storage before creating the typed view.
The helper checks
size_bytes <= max_bytes, but never verifies thatstorageactually containsmax_bytesbytes. It also views the entire byte tensor, so a storage length not divisible by the target dtype can fail unexpectedly. Validate the storage capacity and create the view fromstorage[:size_bytes].Because this relies on PyTorch’s dtype-view rules, verify the behavior against the declared torch 2.12.0 dependency.
🤖 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_dma.py` around lines 145 - 160, Update _persistent_output_view to validate that storage contains at least max_bytes bytes before creating the typed view, while retaining the existing size_bytes limit check and device/dtype validation. Create the typed view from storage[:size_bytes] rather than the entire storage tensor, ensuring non-dtype-aligned excess capacity cannot cause view failures. Confirm the implementation is compatible with the declared torch 2.12.0 dependency.
🤖 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_dma.py`:
- Around line 215-220: Update close() cleanup to release the owner reference to
_output_storage after callers and captured graphs have stopped using the ring,
allowing the persistent CUDA allocation to be reclaimed. Preserve the existing
cleanup ordering and behavior for the rest of the ring.
- Around line 380-386: Move the _ensure_output_storage() call inside the `if out
is None` branch in `all_reduce`, and pass its result to
`_persistent_output_view` there. Keep explicit-output calls free of workspace
allocation while preserving the existing default-output behavior.
---
Nitpick comments:
In `@sparkinfer/comm/pcie/pcie_dma.py`:
- Around line 145-160: Update _persistent_output_view to validate that storage
contains at least max_bytes bytes before creating the typed view, while
retaining the existing size_bytes limit check and device/dtype validation.
Create the typed view from storage[:size_bytes] rather than the entire storage
tensor, ensuring non-dtype-aligned excess capacity cannot cause view failures.
Confirm the implementation is compatible with the declared torch 2.12.0
dependency.
🪄 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: 8a0b5371-a849-45e4-994c-f4fdedbc3a1c
📒 Files selected for processing (3)
sparkinfer/comm/pcie/pcie_dma.pytests/comm/test_pcie_dma_gpu.pytests/comm/test_pcie_dma_output.py
11 tasks
Contributor
Author
|
Addressed both workspace-lifetime review findings in
Validation: |
17 tasks
voipmonitor
force-pushed
the
fix/pcie-dma-persistent-output-20260723
branch
from
July 26, 2026 08:21
3f4b657 to
0c98d6b
Compare
Contributor
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_dma.py`:
- Around line 367-375: Update _ensure_output_storage in
sparkinfer/comm/pcie/pcie_dma.py (lines 367-375) to allocate max_bytes plus the
required 64 KiB workspace tail, while keeping returned output views limited to
the input footprint. Update tests/comm/test_pcie_dma_gpu.py (lines 101-102) to
assert the full backing capacity and verify returned output views exclude the
tail bytes.
🪄 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: 1f75a394-58f4-4295-b8b8-57f7bbad2e98
📒 Files selected for processing (3)
sparkinfer/comm/pcie/pcie_dma.pytests/comm/test_pcie_dma_gpu.pytests/comm/test_pcie_dma_output.py
This was referenced Jul 26, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
reusable, exact-capacity workspace;
observes the persistent allocation before KV-cache sizing;
out=calls allocation-free;close().The allocation is exactly
max_bytes. This PR does not add a mapped tail orpermit any downstream kernel to access bytes outside the logical allocation.
Root cause and scope
A long prefill could be the first call to
PCIeDmaAllReduce.all_reduce(..., out=None)after vLLM had already assignedremaining memory to KV blocks. The late default-output allocation could then
fail, and allocating a fresh output per call also violates the stable-address
contract required by CUDA graph replay.
Constructor allocation is too early for this integration: the ring exists
before vLLM's initial memory snapshot, so constructor-owned bytes are not
charged by the later profile delta. The first profiled default-output dispatch
is both early enough and visible to KV sizing.
The separately reproduced cuBLAS read-ahead issue is not solved by allocator
padding here. Its consumer path must use a legal GEMM/layout contract; carrying
a 64 KiB tail in a general communication allocator would hide an invalid
consumer and permanently waste memory per pool.
Validation
test_pcie_dma_output.pyandtest_pcie_dma_gpu.pywith the two-GPU opt-in:18 passed;
max_bytesand is released onclose.
Complete
rtx6kpro#34
E2E validation used current GG, vLLM #172 plus the MRV2 pool-lifecycle
successor to #168, and this PR:
AI assistance disclosure
The implementation review, tests, E2E validation, and this write-up were
prepared with OpenAI Codex assistance and reviewed by the human submitter.