[GG] Fuse MXFP8 and BF16 MLA query assembly

[voipmonitor]

What

Integrates SparkInfer's fused MLA query projection into the GLM MLA path for
both native MXFP8 and materialized BF16 W_UK_T weights.

• Native MXFP8 uses the operation merged in
  local-inference-lab/sparkinfer#74.
• BF16 weights use the extension in local-inference-lab/sparkinfer#75.
• The fused epilogue writes [M,H,576] directly and appends q_pe without a
  separate concat kernel.
• BF16 output omits the unused scale argument; E4M3 output passes _q_scale.
• H=8/11/16 and M=1..32 are capability-gated. Unsupported cases keep the
  staged path.

The integration imports the final public API,
sparkinfer.gemm.mla_query_projection; the obsolete
sparkinfer.attention.mla_query path is not restored.

Layout and topology behavior

• DCP1 B12X writes directly into its final query workspace when the backend
  confirms an unpadded BF16 layout.
• DCP>1 and padded virtual-TP layouts use a graph-owned local query tensor and
  continue through their established gather/copy path.
• TP6 is covered by the BF16 H=11 specialization.
• Existing UK BMM warmup is skipped only when fused query projection is
  qualified; UV warmup is unchanged.
• Prewarm is limited to graph-visible M regimes to avoid unnecessary retained
  allocations.
• With FP8 KV, only an actual E4M3 query tensor requires backend
  supports_quant_query_input; the fused BF16 query remains valid for DCP.

Safety scope and PR 173

This PR removes the query-side torch.bmm from qualified H=8/11/16,
M=1..32 execution, but it does not supersede #173. Unsupported shapes,
missing SparkInfer support, or failed capability checks still use the staged
query BMM. #173 remains the no-copy cuBLAS read-ahead fix for that fallback.
Merge #173 first, then refresh this PR so its safe operation is retained only
behind the fused-path fallback.

The downstream V up-projection is a separate issue already fixed by #147 and
local-inference-lab/sparkinfer#54.

Performance evidence

Exact paired TP8/DCP1/MTP0 profiling, averaged over eight ranks:

Metric	Staged	Fused	Delta
Query BMM + concat per layer	4.047940 us	2.880912 us	-28.8%
CUDA kernel events, four steps	10,200	9,888	-312
GPU trace span	50,940.968 us	50,450.708 us	-0.96%

No extra hot-path synchronization appeared. Final TP8/DCP1/MTP0 duration
validation on GPUs 0-7 produced aggregate 87.658/87.610 tok/s (mean 87.634)
and active per-user 88.359/88.313 tok/s (mean 88.336), with zero errors.

E2E topology matrix

All values are aggregate CC1 output tok/s. Every 64k case includes the full
prefill-to-decode transition. Models were fully loaded before benchmarks, and
parallel instances were measured sequentially.

Model path	TP	DCP	MTP	ctx0	ctx64k	Errors
NF3 MXFP8	4	1	0	64.44	61.66	0
NF3 MXFP8	4	2	0	58.11	57.51	0
NF3 MXFP8	4	4	0	57.61	56.24	0
Luke BF16	6	1	0	68.69	67.11	0
Luke BF16	6	2	0	58.11	57.27	0
Luke BF16	6	3	0	52.38	50.82	0
Luke BF16	6	6	0	42.66	41.98	0
Luke BF16	8	1	0	87.63	-	0
Luke BF16	8	2	0	73.39	71.33	0
Luke BF16	8	4	0	71.79	69.47	0
Luke BF16	8	8	0	66.43	65.83	0
NF3 MXFP8	4	4	3	111.60	102.13	0
Luke BF16	6	3	3	89.15	79.04	0
Luke BF16	8	4	3	124.27	111.77	0

All workers logged fused-query warmup. The final log audit found no traceback,
Xid, illegal-memory-access, CUBLAS_STATUS, FAULT_PDE, invalid-layout, or
engine-initialization failure.

Validation

• SparkInfer BMM/query/registry suites: 93 passed on GPU 0.
• Focused vLLM dispatch, warmup, workspace, fallback, BF16/MXFP8, and output
  tests: 14 passed.
• Final DCP/FP8 guard regression selection: 9 passed, 2325 deselected.
• Direct H=11 custom-op smoke passed against the final SparkInfer API.
• TP4/TP6/TP8, every MTP0 DCP divisor, representative MTP3, and 64k
  prefill-to-decode transitions passed.
• Ruff, Ruff format, compileall, and git diff --check passed.

Dependencies

• merged: local-inference-lab/sparkinfer#74
• required: local-inference-lab/sparkinfer#75
• fallback safety and merge first: [GG] fix MLA query BMM cuBLAS read-ahead without query copies #173

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes

Fused BF16 and MXFP8 MLA query kernels are added through Sparkinfer-backed custom ops, integrated into MLA decode with workspace and fallback handling, and included in CUDA-graph-aware warmup. B12X sparse attention gains compatible workspace views and extensive CPU test coverage.

Fused MLA query support

Layer / File(s)	Summary
Fused MLA kernel API vllm/model_executor/kernels/attention/b12x_mxfp8_bmm.py	Adds cached Sparkinfer loading, capability checks, custom ops, runtime wrappers, module inspection, and fused-query execution contracts.
Decode dispatch and workspace integration vllm/model_executor/layers/attention/mla_attention.py, vllm/v1/attention/backends/mla/b12x_mla_sparse.py, tests/v1/attention/test_mla_backends.py	Adds conditional fused query dispatch, BF16/FP8 output selection, backend workspace handling, B12X workspace views, fallbacks, and dispatch tests.
Fused query warmup vllm/model_executor/warmup/kernel_warmup.py, vllm/model_executor/kernels/attention/b12x_mxfp8_bmm.py, tests/model_executor/test_flashinfer_autotune_cache.py	Prewarms deduplicated fused-query signatures using graph capture sizes and skips replaced B12X warmup work where applicable.

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

Sequence Diagram(s)

sequenceDiagram
  participant MLAAttention
  participant B12xMLASparseImpl
  participant FusedMLAQueryRunner
  participant Sparkinfer
  MLAAttention->>B12xMLASparseImpl: request compatible output workspace
  B12xMLASparseImpl-->>MLAAttention: return workspace view or None
  MLAAttention->>FusedMLAQueryRunner: run BF16 or MXFP8 fused query
  FusedMLAQueryRunner->>Sparkinfer: call mla_query.run
  Sparkinfer-->>MLAAttention: produce fused query tensor

Loading

Possibly related PRs

• local-inference-lab/vllm#148: Both changes update MLA weight-processing code in mla_attention.py.

Suggested reviewers: lukealonso

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.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
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.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: adding fused MXFP8 and BF16 MLA query assembly support.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch perf/gg-fused-mla-query-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.}