fix(pcie): support head-major DCP reduce outputs - #54
Merged
lukealonso merged 2 commits intoJul 21, 2026
Merged
Conversation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThe change adds head-major output support to sparse MLA scratch buffers and PCIe LSE reduce-scatter. It introduces stride-aware buffer construction and CUDA indexing, validates token-major or head-major layouts, and expands Python and GPU tests for output placement and batch coverage. ChangesHead-major layout support
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant lse_reduce_scatter
participant PCIeDCPA2A
participant dcp_lse_reduce_kernel
Caller->>lse_reduce_scatter: provide partial output and output tensors
lse_reduce_scatter->>lse_reduce_scatter: validate layouts and derive strides
lse_reduce_scatter->>PCIeDCPA2A: pass tensors and stride parameters
PCIeDCPA2A->>dcp_lse_reduce_kernel: launch stride-aware reduction
dcp_lse_reduce_kernel->>Caller: write reduced output tensor
Suggested reviewers: 🚥 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 |
![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: 1
🧹 Nitpick comments (1)
sparkinfer/attention/sparse_mla/_scratch.py (1)
434-448: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate head-major stride formula.
The literal
(v_head_dim, max_total_q * v_head_dim, 1)here duplicates the same head-major stride computation already expressed (generically, fromoutput.shape) inworkspace.py's_split_output_buffer_from_tmp. If the head-major convention ever changes, only one of the two spots is likely to get updated.♻️ Suggested consolidation
- if caps.head_major_output: - output_buffer, _ = materialize_scratch_strided_view( - scratch_storage, - offset_bytes=layout.output_offset_bytes, - shape=(max_total_q, num_q_heads, v_head_dim), - stride=(v_head_dim, max_total_q * v_head_dim, 1), - dtype=caps.dtype, - ) + if caps.head_major_output: + output_buffer, _ = materialize_scratch_strided_view( + scratch_storage, + offset_bytes=layout.output_offset_bytes, + shape=(max_total_q, num_q_heads, v_head_dim), + stride=_split_tmp_output_stride( + max_total_q=max_total_q, + num_q_heads=num_q_heads, + max_chunks_per_row=1, + v_head_dim=v_head_dim, + head_major_output=True, + )[:3] + (1,), + dtype=caps.dtype, + )(or extract a small shared 3-tuple helper in
workspace.pyreused by both call 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/attention/sparse_mla/_scratch.py` around lines 434 - 448, Consolidate the head-major stride calculation used by the `caps.head_major_output` branch in the surrounding output-buffer setup with the existing generic computation in `_split_output_buffer_from_tmp`. Reuse a shared helper or established workspace logic instead of hardcoding `(v_head_dim, max_total_q * v_head_dim, 1)`, while preserving the current shape and layout behavior.
🤖 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_a2a.py`:
- Around line 33-47: Update _is_supported_bhd_layout so
capacity_strided_head_major additionally requires stride_head to be divisible by
8, matching the CUDA extension’s alignment requirement. Keep the existing
dimensional, contiguous-last-dimension, minimum-capacity-stride, and
packed-token-major checks unchanged so _validate can reject unaligned head-major
tensors cleanly.
---
Nitpick comments:
In `@sparkinfer/attention/sparse_mla/_scratch.py`:
- Around line 434-448: Consolidate the head-major stride calculation used by the
`caps.head_major_output` branch in the surrounding output-buffer setup with the
existing generic computation in `_split_output_buffer_from_tmp`. Reuse a shared
helper or established workspace logic instead of hardcoding `(v_head_dim,
max_total_q * v_head_dim, 1)`, while preserving the current shape and layout
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: c6a95f66-ffe1-498f-8cb7-15aaf5e26926
📒 Files selected for processing (7)
sparkinfer/attention/_shared/workspace.pysparkinfer/attention/sparse_mla/_scratch.pysparkinfer/comm/pcie/pcie_dcp_a2a.cusparkinfer/comm/pcie/pcie_dcp_a2a.pytests/attention/test_compressed_scratch_bindings.pytests/comm/test_pcie_dcp_a2a.pytests/comm/test_pcie_dcp_a2a_gpu.py
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
35 tasks
This was referenced Jul 30, 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
2 participants
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
Fixes the DCP PCIe LSE reduce output contract so callers can provide a BHD tensor whose logical layout is token-major but whose physical storage is head-major. This avoids handing cuBLAS a strided view backed by a tight IPC/custom allocation in the downstream MLA V up-projection path.
Root Cause
The GLM MLA
_v_up_projcall uses a strided batched GEMM layout. Guarded repros showed cuBLAS can legally read ahead to the next 64 KiB boundary for this pattern. Normal PyTorch allocator segments make that safe, but tight DCP IPC/pool outputs can fault on the first unmapped page after the logical tensor.Padding every allocation or cloning pool outputs fixes correctness but costs either KV capacity or decode latency. This PR instead lets the DCP producer write directly into caller-provided head-major storage while preserving the public BHD shape.
Changes
head_major_outputscratch planning for sparse MLA output buffers.Validation
python3 -m pytest tests/comm/test_pcie_dcp_a2a.py::test_runtime_accepts_head_major_input_and_output tests/comm/test_pcie_dcp_a2a.py::test_runtime_validates_and_dispatches_to_extension -q-> 2 passedSPARKINFERSparseMLAScratchCaps(head_major_output=True)-> passedSummary by CodeRabbit