[GG] fix MLA query BMM cuBLAS read-ahead without query copies

[voipmonitor]

Summary

• Add a stable ABI CUDA op, torch.ops._C.safe_mla_query_bmm, for the MLA query absorption BMM shape used by B12X sparse MLA decode.
• Route B12X sparse MLA query absorption through that op instead of materializing the head-major query view with .contiguous().
• Keep the fallback conservative: if the CUDA op is unavailable or the tensors are not CUDA BF16, the safe path materializes the query before torch.bmm.
• Supersedes [GG] fix(mla): restore safe query BMM layouts #170. That PR fixed the Xid/read-ahead issue by copying the query view, but it regressed DCP1 decode throughput.

Root Cause

The B12X sparse MLA query absorption path feeds a head-major, non-contiguous BF16 query view into torch.bmm with shape [heads, tokens, q_dim] x [heads, q_dim, latent] -> [heads, tokens, latent]. Tight custom/DCP allocations can expose cuBLAS BF16 kernels reading past the logical tensor tail for alignment. The broad .contiguous() workaround made the input safe but added a per-layer copy on the hot decode path.

This PR keeps the same math but calls cuBLAS directly with the row-major mapping for this exact layout, using CUBLAS_COMPUTE_32F_PEDANTIC and CUBLAS_GEMM_DEFAULT. It does not mutate the global cuBLAS math mode and does not add tail padding or reduce KV cache.

Validation

• ruff format tests/kernels/test_safe_mla_query_bmm.py tests/v1/attention/test_mla_backends.py vllm/model_executor/layers/attention/mla_attention.py vllm/v1/attention/backends/mla/b12x_mla_sparse.py
• ruff check tests/kernels/test_safe_mla_query_bmm.py tests/v1/attention/test_mla_backends.py vllm/model_executor/layers/attention/mla_attention.py vllm/v1/attention/backends/mla/b12x_mla_sparse.py
• git diff --check lil/dev/gilded-gnosis...HEAD
• Built v20 candidate image with this op: voipmonitor/vllm:gilded-gnosis-v20-safeqbmm-vllm9fdb155-si055f839-fi801d57a-cu132-20260723
• Built-image GPU smoke: safe_mla_query_bmm matches torch.bmm(query.contiguous(), weight) for (heads,tokens) = (8,1), (8,2), (8,6), (8,11), (11,6), (16,6) and replays under CUDA graph capture.
• DCP1 gate on GPUs 0-7, TP8/DCP1/MTP0/A16/orig Luke NVFP4, MAX_NUM_SEQS=1, GRAPH=6: aggregate decode 87.46 and 87.41 tok/s, KV budget 570,688 tokens.

Host pytest note: targeted CPU pytest in this checkout is blocked by missing fastapi in the host environment during tests/conftest.py import, before these tests run.

Summary by CodeRabbit

• New Features

  • Added a CUDA-optimized BF16 matrix multiplication path for MLA attention.
  • MLA attention can now automatically use the optimized path when supported, with a compatible fallback.
  • Added validation for tensor layouts, devices, shapes, and data types.

• Bug Fixes

  • Prevented potential out-of-bounds reads during BF16 CUDA matrix multiplication.

• Tests

  • Added coverage for numerical correctness, CUDA Graph replay, backend selection, and fallback behavior.

[coderabbitai]

Warning

Review limit reached

@voipmonitor, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 56 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: da69700f-29d8-4c03-ae4e-0af6168228a7

📥 Commits

Reviewing files that changed from the base of the PR and between c7312da and a9a6c62.

📒 Files selected for processing (5)

• CMakeLists.txt
• csrc/libtorch_stable/attention/mla/safe_query_bmm.cu
• tests/kernels/test_safe_mla_query_bmm.py
• tests/v1/attention/test_mla_backends.py
• vllm/model_executor/layers/attention/mla_attention.py

📝 Walkthrough

Walkthrough

Adds a CUDA BF16 safe_mla_query_bmm operator using cuBLAS, registers and builds it in the stable extension, routes supported MLA attention paths through it, and adds CUDA and CPU tests for correctness, graph replay, flag propagation, fallback behavior, and opt-outs.

Changes

Safe MLA Query BMM

Layer / File(s)	Summary
CUDA operator contract and implementation csrc/libtorch_stable/ops.h, csrc/libtorch_stable/torch_bindings.cpp, csrc/libtorch_stable/attention/mla/safe_query_bmm.cu, CMakeLists.txt	Defines and registers safe_mla_query_bmm, validates CUDA BF16 tensor shapes and strides, executes batched cuBLAS GEMM, and links the required CUDA library.
MLA attention dispatch vllm/model_executor/layers/attention/mla_attention.py, vllm/v1/attention/backends/mla/b12x_mla_sparse.py	Propagates the backend capability flag and routes supported MQA query projections through the safe operator, with a contiguous torch.bmm fallback.
Validation coverage tests/kernels/test_safe_mla_query_bmm.py, tests/v1/attention/test_mla_backends.py	Tests numerical equivalence, CUDA Graph replay, capability propagation, fallback contiguity, and explicit safe-path opt-outs.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant B12xMLASparseImpl
  participant MLAAttention
  participant _run_mla_query_bmm
  participant safe_mla_query_bmm
  participant cuBLAS
  B12xMLASparseImpl->>MLAAttention: provide safe BMM capability
  MLAAttention->>_run_mla_query_bmm: pass MQA query projection tensors
  _run_mla_query_bmm->>safe_mla_query_bmm: dispatch CUDA BF16 operation
  safe_mla_query_bmm->>cuBLAS: execute strided batched GEMM
  cuBLAS-->>safe_mla_query_bmm: write output tensor
  safe_mla_query_bmm-->>MLAAttention: return populated output

Loading

Possibly related PRs

• local-inference-lab/vllm#136: Both modify MLA BMM operand handling for cuBLAS BF16 correctness.
• local-inference-lab/vllm#138: Both change MLA BMM dispatch and operand preparation in mla_attention.py.
• local-inference-lab/vllm#31: Both modify initialization behavior in B12xMLASparseImpl.

Suggested reviewers: yewentao256, lukealonso, matthewbonanni

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately captures the main change: fixing MLA query BMM read-ahead via a safe cuBLAS path without query copies.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/gg-mla-query-bmm-pedantic-standalone-20260723

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

CMakeLists.txt (1)

384-416: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Add the BLAS dependency for the stable MLA BMM helper.

csrc/libtorch_stable/attention/mla/safe_query_bmm.cu is unconditionally added to _C_stable_libtorch, but it uses cublasGemmStridedBatchedEx and the current cuBLAS discovery/linking only runs when VLLM_GPU_LANG == "CUDA". This needs a matching BLAS link for ROCm builds that cover the HIP-generated symbols, otherwise the stable extension will fail to resolve libblas/hipblas/rocblas during link 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 `@CMakeLists.txt` around lines 384 - 416, Update the _C_stable_libtorch build
configuration to discover and link the appropriate BLAS dependency for HIP
builds as well as CUDA builds, covering the cuBLAS/hipBLAS/rocBLAS symbols used
by safe_query_bmm.cu. Ensure the stable extension target links that dependency
whenever VLLM_GPU_LANG is CUDA or HIP, without changing the source list.

🤖 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 `@tests/v1/attention/test_mla_backends.py`:
- Around line 865-877: The sparse profile test double used by
MLAAttention.forward_impl must define use_safe_mla_query_bmm before execution.
Update the manual MLAAttention setup around ProfileSparseImpl to initialize this
attribute consistently with the existing sparse unit-path test double.

---

Outside diff comments:
In `@CMakeLists.txt`:
- Around line 384-416: Update the _C_stable_libtorch build configuration to
discover and link the appropriate BLAS dependency for HIP builds as well as CUDA
builds, covering the cuBLAS/hipBLAS/rocBLAS symbols used by safe_query_bmm.cu.
Ensure the stable extension target links that dependency whenever VLLM_GPU_LANG
is CUDA or HIP, without changing the source list.

🪄 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: a9163e0e-a5f7-48c3-91cd-f0d4fdbfecfc

📥 Commits

Reviewing files that changed from the base of the PR and between 4a4299c and c7312da.

📒 Files selected for processing (8)

• CMakeLists.txt
• csrc/libtorch_stable/attention/mla/safe_query_bmm.cu
• csrc/libtorch_stable/ops.h
• csrc/libtorch_stable/torch_bindings.cpp
• tests/kernels/test_safe_mla_query_bmm.py
• tests/v1/attention/test_mla_backends.py
• vllm/model_executor/layers/attention/mla_attention.py
• vllm/v1/attention/backends/mla/b12x_mla_sparse.py

[voipmonitor]

Review follow-up pushed in 9fcbe21071.

• Added the missing use_safe_mla_query_bmm field to the sparse profile fixture.
• Added the fixture backend identity required by current GG profiling and the output-dtype field required by the dependent [GG] Fuse MXFP8 and BF16 MLA query assembly #174 integration.
• The suggested HIP BLAS link is not applicable: safe_query_bmm.cu calls cuBLAS directly. Instead, the source is now included only for CUDA builds and runtime dispatch is explicitly CUDA-gated. HIP retains the existing contiguous torch.bmm fallback.

Validation: both affected CPU tests pass in the CUDA test image; Ruff and git diff --check pass. The clean release build will validate the CUDA extension link.

[voipmonitor]

Added a9a6c62dec to keep the safe operand/layout contract while using CUBLAS_COMPUTE_32F instead of CUBLAS_COMPUTE_32F_PEDANTIC. The pedantic mode forced a slower non-tensor-core path on production prefill shapes; it is not part of the tail-read safety contract. The expanded GPU test now includes the 8,192 x 192 production prefill shape. Validation on the rebuilt implementation: 7/7 safe-query BMM tests passed, including CUDA graph replay.