[GG] Add BF16 MLA query projection and assembly#75
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughAdds a Triton BF16 MLA query projection backend, unified BF16/MXFP8 API dispatch, CUDA-graph and capability tests, and benchmark support for selecting weight formats and reporting expanded correctness metadata. ChangesMLA query projection weight formats
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Benchmark
participant mla_query_projection_api
participant BF16Backend
participant TritonKernel
Benchmark->>mla_query_projection_api: prewarm BF16 projection
mla_query_projection_api->>BF16Backend: prewarm weight and M regimes
BF16Backend->>TritonKernel: compile and launch warmup shapes
Benchmark->>mla_query_projection_api: run projection
mla_query_projection_api->>BF16Backend: dispatch BF16 tensors
BF16Backend->>TritonKernel: project q_nope and append q_pe
TritonKernel-->>Benchmark: write output tensor
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
voipmonitor
marked this pull request as ready for review
July 23, 2026 16:42
![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
🤖 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/gemm/mla_query_projection/api.py`:
- Around line 92-117: Update can_implement() so the package-wide is_supported()
gate applies only to the MXFP8 path, while the BF16 path uses its own capability
check matching _bf16.can_implement() requirements (SM12x/Triton without
requiring CUTLASS DSL). Preserve the existing recipe-specific argument handling
and ensure shared support metadata no longer incorrectly marks BF16 unsupported
on Triton-only Blackwell installations.
🪄 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: 2cc7459b-3103-4d64-a8ec-6c6d98a21d9b
📒 Files selected for processing (5)
benchmarks/benchmark_mla_query_projection.pysparkinfer/gemm/mla_query_projection/__init__.pysparkinfer/gemm/mla_query_projection/_bf16.pysparkinfer/gemm/mla_query_projection/api.pytests/gemm/test_mla_query_projection.py
Comment on lines
+92
to
+117
| weight_format: Literal["bf16", "mxfp8"] = "mxfp8", | ||
| device=None, | ||
| ) -> bool: | ||
| """Return whether the qualified fused-query specialization covers a plan.""" | ||
| return is_supported(device) and _mxfp8.can_implement_mla_query_projection( | ||
| batch=num_heads, | ||
| max_m=max_m, | ||
| n=latent_dim, | ||
| k=nope_dim, | ||
| output_dtype=output_dtype, | ||
| b_major="n", | ||
| sf_axis="n", | ||
| ) | ||
| if not is_supported(device): | ||
| return False | ||
| if weight_format == "bf16": | ||
| return _bf16.can_implement( | ||
| num_heads=num_heads, | ||
| max_m=max_m, | ||
| nope_dim=nope_dim, | ||
| latent_dim=latent_dim, | ||
| output_dtype=output_dtype, | ||
| device=device, | ||
| ) | ||
| if weight_format == "mxfp8": | ||
| return _mxfp8.can_implement_mla_query_projection( | ||
| batch=num_heads, | ||
| max_m=max_m, | ||
| n=latent_dim, | ||
| k=nope_dim, | ||
| output_dtype=output_dtype, | ||
| b_major="n", | ||
| sf_axis="n", | ||
| ) | ||
| return False |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check whether the BF16 kernel module has any CUTLASS dependency
rg -n "cutlass" sparkinfer/gemm/mla_query_projection/_bf16.py
rg -n "requires" sparkinfer/gemm/mla_query_projection/__init__.py
rg -n "def has_cutlass_dsl" -A5 sparkinfer/_lib/gating.pyRepository: local-inference-lab/sparkinfer
Length of output: 436
🏁 Script executed:
#!/bin/bash
set -eu
echo "== api.py outline/snippet =="
wc -l sparkinfer/gemm/mla_query_projection/api.py
sed -n '1,150p' sparkinfer/gemm/mla_query_projection/api.py
echo "== _bf16.py imports/snippet =="
head -160 sparkinfer/gemm/mla_query_projection/_bf16.py
echo "== _mxfp8.py cutslass references =="
rg -n "cutlass|requires|can_implement_mla_query_projection" sparkinfer/gemm/mla_query_projection/_mxfp8.py sparkinfer/gemm/mla_query_projection/__init__.py sparkinfer/gemm/mla_query_projection/_bf16.py
echo "== gating.py relevant symbols =="
sed -n '1,130p' sparkinfer/_lib/gating.py
echo "== search default_is_supported definitions/usages =="
rg -n "def default_is_supported|default_is_supported|require_sparkinfer|META =|requires" sparkinfer -g '*.py' | head -200Repository: local-inference-lab/sparkinfer
Length of output: 9662
🏁 Script executed:
#!/bin/bash
set -eu
echo "== locate mxfp8 shared file =="
fd -a '_shared' . | sed 's#^\./##'
echo "== files under gemm/mla_query_projection =="
git ls-files sparkinfer/gemm/mla_query_projection
echo "== META and required entry points =="
rg -n "class META|META|requires|entry_points|requires_sparkinfer|require_sparkinfer|is_supported" sparkinfer/gemm/mla_query_projection sparkinfer -g '*.py' | head -300
echo "== gating.py full relevant section =="
wc -l sparkinfer/_lib/gating.py
sed -n '1,140p' sparkinfer/_lib/gating.py
echo "== exact bf16 imports and can support section =="
sed -n '100,260p' sparkinfer/gemm/mla_query_projection/_bf16.pyRepository: local-inference-lab/sparkinfer
Length of output: 40881
🏁 Script executed:
#!/bin/bash
set -eu
echo "== __init__.py META =="
cat -n sparkinfer/gemm/mla_query_projection/__init__.py
echo "== search mxfp8_bmm imports/declaration =="
rg -n "mxfp8_bmm|_mxfp8|mla_query_projection.*mxfp8|mxfp8.*mla_query|can_implement_mla|require" sparkinfer/gemm -g '*.py'
echo "== search mxfp8 modules =="
git ls-files sparkinfer | rg '(^|/)mxfp8|mla_query_projection|_shared'Repository: local-inference-lab/sparkinfer
Length of output: 28294
🏁 Script executed:
#!/bin/bash
set -eu
echo "== mxfp8_bmm outline relevant =="
wc -l sparkinfer/gemm/_shared/mxfp8_bmm.py
sed -n '1,120p' sparkinfer/gemm/_shared/mxfp8_bmm.py
rg -n "def can_implement_mla_query_projection|def mla_query_projection|def prewarm_mla_query_projection|cutlass|cute|cuTe|requires|has_cutlass|require_cutlass|default_is_supported" sparkinfer/gemm/_shared/mxfp8_bmm.py
echo "== cutlass/cute imports in module =="
rg -n "^(from|import) .*cute|from .*cutlass|import .*cutlass|cute|cutlass" sparkinfer/gemm/_shared/mxfp8_bmm.py sparkinfer/gemm/mla_query_projection/__init__.py sparkinfer/gemm/mla_query_projection/api.py sparkinfer/gemm/mla_query_projection/_bf16.pyRepository: local-inference-lab/sparkinfer
Length of output: 28394
Weight the shared is_supported() gate to MXFP8, not BF16.
can_implement() shares the package-wide gate for weight_format="bf16" even though _bf16.py is the only BF16 code path and only needs SM12x/Triton. default_is_supported() also requires CUTLASS DSL unconditionally, and this shared metadata says both bf16/mxfp8 recipes are supported. Add a per-recipe capability gate or split metadata so BF16 plans are not reported unsupported on Triton-only Blackwell installs.
🤖 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/gemm/mla_query_projection/api.py` around lines 92 - 117, Update
can_implement() so the package-wide is_supported() gate applies only to the
MXFP8 path, while the BF16 path uses its own capability check matching
_bf16.can_implement() requirements (SM12x/Triton without requiring CUTLASS DSL).
Preserve the existing recipe-specific argument handling and ensure shared
support metadata no longer incorrectly marks BF16 unsupported on Triton-only
Blackwell installations.
11 tasks
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.
What
Extends the merged
sparkinfer.gemm.mla_query_projectionoperation with asmall-M BF16-weight specialization for GLM MLA query absorption.
q_nope [H,M,192]x BF16W_UK_T [H,192,512].[M,H,576]output with the 64-wide RoPE suffix appendedin the same launch.
This is a second weight-format dispatch inside the operation introduced by
#74. It does not restore the removed
sparkinfer.attention.mla_querypackageand does not modify the merged MXFP8 lowering.
Contract
gemm.mla_query_projection.gemm/_shared/mxfp8_bmm.py.q_scale; E4M3 output requires one.preserved.
--weight-format bf16instead of creating aparallel API.
The kernel accumulates in FP32 and rounds the 512-wide projection to BF16
before optional E4M3 conversion, matching the staged operation boundary.
Triton and cuBLAS may reduce K in a different order, so BF16 projection is not
promised bitwise identical. Across H=8/11/16 and M=1/4/6/17/32:
absolute difference;
absolute difference.
Native MXFP8 behavior remains bitwise checked and unchanged.
Performance evidence
Exact paired TP8/DCP1/MTP0 profiling, averaged over eight ranks:
The 312 removed events are exactly 78 layers x 4 steps. Both traces contain 12
cudaStreamWaitEvent, 8 event records, and 4 graph launches. FinalTP8/DCP1/MTP0 duration validation produced aggregate 87.658/87.610 tok/s
(mean 87.634), active per-user 88.359/88.313 tok/s (mean 88.336), and zero
errors.
The standalone graph-replay benchmark remains launch-floor dominated for this
small BF16 operation. Across H=8/11/16 and M=1/4/32, its fused/staged median
ratio was 0.994 to 1.024 over 1,000 balanced AB/BA samples per shape.
E2E integration matrix
Paired vLLM PR:
local-inference-lab/vllm#174at1f2bf43c6b. All values areaggregate CC1 output tok/s; 64k cases include prefill-to-decode transition.
Every worker logged fused-query warmup. No final server log contained a
traceback, Xid, illegal-memory-access,
CUBLAS_STATUS,FAULT_PDE,invalid-layout, or engine-initialization failure.
Validation
tests/gemm/test_bmm.py,tests/gemm/test_mla_query_projection.py, andtests/test_registry.py: 93 passed on GPU 0.M=1/6/9/17/32; pitched outputs; scale validation; and CUDA graph replay with
mutated inputs are covered.
selection at 9 passed and 2325 deselected.
compileall, andgit diff --checkpassed.Paired vLLM integration:
local-inference-lab/vllm#174