[GG] fix(mla): tail-pad cuBLAS BMM inputs - #140
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review (1)
Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds cuBLAS tail-padding utilities and propagates tail-padding requirements through unquantized and MXFP8 linear layers, MLA attention, and DCP A2A reductions. Tests cover value preservation, contiguous buffers, CUDA graphs, storage-tail contracts, and parameter forwarding. ChangescuBLAS storage and linear execution
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant B12xMLASparseImpl
participant MLAAttention
participant CUBLASTailPadding
participant DCPReduction
B12xMLASparseImpl->>MLAAttention: configure MLA tail padding
MLAAttention->>CUBLASTailPadding: pad BMM inputs when required
MLAAttention->>DCPReduction: request tail-padded reduction output
DCPReduction-->>MLAAttention: return tail-padded output
Possibly related PRs
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 |
433b4ac to
371085e
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
vllm/utils/cublas.py (1)
30-41: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
strict=Trueto thezip()call.
tensor.shapeandtensor.stride()are always equal length, so this is functionally safe, but Ruff (B905) flags it and it's a one-line fix.🔧 Proposed fix
last_item = tensor.storage_offset() + sum( - (size - 1) * stride for size, stride in zip(tensor.shape, tensor.stride()) + (size - 1) * stride + for size, stride in zip(tensor.shape, tensor.stride(), strict=True) )🤖 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 `@vllm/utils/cublas.py` around lines 30 - 41, Update the zip call in storage_tail_bytes to pass strict=True when pairing tensor.shape with tensor.stride(), leaving the surrounding storage-tail calculation unchanged.Source: Linters/SAST tools
vllm/v1/attention/backends/mla/b12x_mla_sparse.py (1)
976-982: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winTail-padding byte count is decoupled from the configured value across files. The MLA backend configures
mla_bmm_tail_padding_bytes(a literal64 * 1024inb12x_mla_sparse.py), but only a booleantail_pad_outputflag is propagated throughdcp_a2a_lse_reduce/cp_lse_ag_out_rs; every consumer independently re-derives the padding size from the module-levelCUBLAS_BMM_TAIL_PADDING_BYTESdefault invllm/utils/cublas.py. This is correct today only because both literals coincidentally match — any future divergence (e.g. a backend needing a different tail size) would silently apply the wrong padding amount, risking exactly the cuBLAS unmapped-read/Xid hazard this PR sets out to fix.
vllm/v1/attention/backends/mla/b12x_mla_sparse.py#L976-L982: import and reuseCUBLAS_BMM_TAIL_PADDING_BYTESfromvllm.utils.cublasinstead of the literal64 * 1024, and consider threading the actual byte value (not just a bool) through the downstream calls.vllm/v1/attention/ops/dcp_alltoall.py#L294-L304: accept the actual tail-padding byte count from the caller (or a module-shared constant import) instead of hardcodingCUBLAS_BMM_TAIL_PADDING_BYTESin thetail_padded_emptycall.vllm/v1/attention/ops/dcp_alltoall.py#L876-L892: same fix for the_dcp_a2a_unpack_combineallocation.vllm/v1/attention/ops/common.py#L227-L247: pass the configured byte count intoensure_cublas_tail_paddinginstead of relying on its default.🤖 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 `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py` around lines 976 - 982, Use one configured tail-padding byte count throughout the MLA path instead of independently re-deriving the default. In vllm/v1/attention/backends/mla/b12x_mla_sparse.py#L976-L982, initialize mla_bmm_tail_padding_bytes from CUBLAS_BMM_TAIL_PADDING_BYTES and propagate the value through downstream calls. Update vllm/v1/attention/ops/dcp_alltoall.py#L294-L304 and `#L876-L892` so tail_padded_empty and _dcp_a2a_unpack_combine receive and use the caller’s byte count, and update vllm/v1/attention/ops/common.py#L227-L247 to pass that configured value to ensure_cublas_tail_padding instead of relying on its default.
🤖 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 `@vllm/utils/cublas.py`:
- Around line 30-41: Update the zip call in storage_tail_bytes to pass
strict=True when pairing tensor.shape with tensor.stride(), leaving the
surrounding storage-tail calculation unchanged.
In `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py`:
- Around line 976-982: Use one configured tail-padding byte count throughout the
MLA path instead of independently re-deriving the default. In
vllm/v1/attention/backends/mla/b12x_mla_sparse.py#L976-L982, initialize
mla_bmm_tail_padding_bytes from CUBLAS_BMM_TAIL_PADDING_BYTES and propagate the
value through downstream calls. Update
vllm/v1/attention/ops/dcp_alltoall.py#L294-L304 and `#L876-L892` so
tail_padded_empty and _dcp_a2a_unpack_combine receive and use the caller’s byte
count, and update vllm/v1/attention/ops/common.py#L227-L247 to pass that
configured value to ensure_cublas_tail_padding instead of relying on its
default.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 669869ef-c8f8-4ea6-9edf-6d51ed408b20
📒 Files selected for processing (12)
tests/distributed/test_dcp_a2a.pytests/model_executor/kernels/test_b12x_mxfp8_linear.pytests/model_executor/layers/test_unquantized_linear_tail_padding.pyvllm/model_executor/kernels/linear/mxfp8/b12x.pyvllm/model_executor/layers/attention/mla_attention.pyvllm/model_executor/layers/linear.pyvllm/model_executor/layers/mla.pyvllm/model_executor/layers/utils.pyvllm/utils/cublas.pyvllm/v1/attention/backends/mla/b12x_mla_sparse.pyvllm/v1/attention/ops/common.pyvllm/v1/attention/ops/dcp_alltoall.py
(cherry picked from commit 1905e2e)
1905e2e to
4a14a60
Compare
|
Superseded by the head-major output contract in #147 together with local-inference-lab/sparkinfer#54. That solution passed the guarded cuBLAS reproduction and DCP1/DCP2/DCP4 validation without the KV-capacity and hot-path costs of generic tail padding. |
1 participant
Summary
Fix the GLM MLA cuBLAS BMM layout hazard without reintroducing a blanket
.contiguous()copy on the hot path.This PR adds an explicit 64 KiB tail-padding storage contract for tensors that feed the GLM MLA cuBLAS BMMs, and wires that contract through the B12X sparse MLA / DCP producers.
Companion b12x PR for native MXFP8 linear output padding: local-inference-lab/sparkinfer#51
Stack Position
This PR is stacked on the clean GG reapply chain:
Do not merge until explicitly approved.
Root cause
The guarded CUDA repro confirmed that the relevant SM120 cuBLAS strided BMM pattern can read past the logical end of the BF16 input tensor up to the next 64 KiB boundary. With ordinary PyTorch caching-allocator outputs this is usually masked by allocator slack. With tight producer-owned storage, especially DCP/B12X IPC/NCCL outputs or custom linear outputs, the read can cross into an unmapped page and trigger Xid31 /
FAULT_PDE/CUBLAS_STATUS_INTERNAL_ERRORin_v_up_proj.Earlier safety fixes that force
.contiguous()made the crash go away, but they added unconditional copies and caused DCP1 decode regressions. This PR instead fixes the producer/output-storage contract.Changes
vllm.utils.cublashelpers:CUBLAS_BMM_TAIL_PADDING_BYTES = 64 * 1024tail_padded_empty(...)storage_tail_bytes(...)ensure_cublas_tail_padding(...)output_tail_padding_byteson MLA Q projection layers so unquantized and MXFP8 linear producers can write directly into safe storage.ensure_cublas_tail_padding()at the two BMM call sites, but the intended path is producer-padded zero-copy output.Validation
Unit / focused tests:
tests/model_executor/layers/test_unquantized_linear_tail_padding.py: 2 passed.tests/distributed/test_dcp_a2a.py -k 'test_cublas_tail_padding_preserves_values_and_storage_contract': 1 passed.tests/distributed/test_dcp_a2a.py -k 'test_ag_rs_can_tail_pad_output_for_mla_bmm or test_b12x_lse_reduce_makes_views_contiguous or test_pack_unpack_combine_matches_reference': 26 passed.tests/model_executor/kernels/test_b12x_mxfp8_linear.py -k 'test_b12x_mxfp8_apply_uses_packed_weight or test_b12x_mxfp8_custom_op_body_uses_forward_context': 2 passed.git diff --check: clean.E2E validation on the diagnostic image
local/vllm:gilded-gnosis-v20-tailpad-producers-20260720(sha256:8eb1ef5e416d0c9757b80d146b22e987e70511d9b3d59f7dff5ee777ce8fd56a), with native CUDA allocator and no.clone()workaround:B12X_PCIE_DMA=1,F8_DMA=ring, long prompt no-prefix-cache run: passed, no new Xid indmesg.82.71 tok/s, prefill64k5897 tok/s, KV546752, no Xid.69.17 tok/s, prefill64k5493 tok/s, KV1038976, no Xid.118.58 tok/s, prefill64k5295 tok/s, KV2097152, no Xid.76.57 tok/s, prefill64k2372 tok/s, KV1459736, no Xid.The DCP1 check is the important regression gate for the previous
.contiguous()fix: the producer-tail-padding path avoids that hot-path slowdown.