Add NVFP4/FP8 Q/K/P/V + 2:4 attention quantization for MLA (vLLM TRITON_MLA) - #2244
Draft
kaix-nv wants to merge 5 commits into
Draft
Add NVFP4/FP8 Q/K/P/V + 2:4 attention quantization for MLA (vLLM TRITON_MLA)#2244kaix-nv wants to merge 5 commits into
kaix-nv wants to merge 5 commits into
Conversation
Contributor
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: trueComment |
MLA kernels (modelopt/torch/kernels/quantization/attention/mla/): - mla_prefill: varlen prefill with asymmetric head dims, fused Q/K/P/V fake quant (NVFP4 1x16-along-contraction / per-tensor FP8), token-exact N:M score sparsity, natural-log LSE for chunked-context merging. - mla_decode: split-K absorbed decode over the paged latent cache with grouped query heads and fused P QDQ. The write-once quantized latent (module-level kv_c/k_pe QDQ before the cache write) is the single representation both BMM1-K and BMM2-V consume - no on-read re-quant. - reference: independent eager oracles replicating the kernel schedules. vLLM integration (TRITON_MLA, vLLM >= 0.26): - ModelOptMLAImpl reclasses TritonMLAImpl; forward_mha swaps a per-layer prefill backend around the inherited kv_b_proj/chunked-context plumbing, forward_mqa runs the FP32-carrier absorbed-Q QDQ + decode kernel. - Installer discovers MLAAttention on the quantize path with MLA gates (fp8 cache, sparse indexer, q padding, skip-softmax rejected) and the existing validate-then-publish/rollback flow. - configure_vllm_nvfp4_mla_quantizers: k_format governs kv_c/k_pe/k_mha, v_format the prefill projected V (k_mha/v_mha named to stay outside the *[kv]_bmm_quantizer calibration patterns). - Fix stale vllm.attention.layer MLA import in vllm_ptq_utils; extend the reload key rewrite to kv_c/k_pe quantizers. Verified in vllm/vllm-openai:v0.26.0 on RTX A5000 (SM86): MLA kernel golden tests 28 passed, dense kernel regression 33 passed, runtime and worker suites 84 passed, dynamic-modules e2e 10 passed. E4M3 quant-path tests are capability-gated and still need an SM89+ run. Signed-off-by: Kai Xu <kaix@nvidia.com>
The module-level kv_c/k_pe latent QDQ was applied in _QuantVLLMMLAAttention.forward before super().forward, so vLLM's prefill kv_b_proj projected an already-quantized latent, which the prefill kernel then quantized again (k_mha/v_mha) — double-quantizing the new-tokens prefill K/V operands. Decode was unaffected (single write-once latent, consumed as-is). Move the latent QDQ off the module forward and into a ModelOptMLAImpl.do_kv_cache_update override that quantizes out-of-place at cache-write time. vLLM calls do_kv_cache_update before forward_impl and passes forward_impl the same tensors, so the quantized copy reaches the cache (read by decode) while prefill projects the bf16 latent — making the new-tokens prefill operands single-quant. Gated by a _kv_quant_in_cache_write flag so the fakequant/calibration path (fused impl not installed) still quantizes module-side. Residual, documented in the README: cached-context prefill chunks gather from the quantized paged cache and re-quantize, so those chunks stay double-quant — inherent to reading a stored quantized latent, and absent for single-chunk prompts. Verified: MLA + dense vLLM runtime synthetic suites 35 passed in vllm/vllm-openai:v0.26.0; numeric single-quant behavior exercised by the SM89+ tiny-DeepSeek install+generate e2e. Signed-off-by: Kai Xu <kaix@nvidia.com>
The decode kernel consumes the write-once quantized latent as-is for both BMMs (V = trans(k_nope)), so decode V reuses the latent's K-side feature-axis quantization — it does not honor v_format, which applies to prefill only. A shared latent cannot be stored quantized along both K's feature axis and V's token axis; independent token-axis decode V would require an on-read re-quant (double-quant) or a raw-cache on-read model, either of which forfeits the write-once step/split stability. Correct the README and decode kernel docstring to state the actual contract. Signed-off-by: Kai Xu <kaix@nvidia.com>
…uant) Switch the fused MLA quant model from a compressed write-once NVFP4 latent cache to the faithful quantized-BMM-operand model used by the MNI reference and the dense Q/K/P/V path: the latent cache stays RAW (bf16), and every BMM operand is fake-quantized once, in-kernel, along its own contraction axis. Decode now honors k_format and v_format independently — it reads the raw latent and quantizes K along the feature axis and V along the token axis separately (plus the existing in-kernel P and caller-side Q). This fixes the prior contract gap where decode V inherited k_format and v_format was ignored. Kernel (mla_decode.py): add K_QDQ/V_QDQ constexprs + k_qdq/v_qdq host args; capture raw V before the K on-read quant; independent feature-axis K and token-axis V QDQ via _v_qdq_nvfp4/fp8_scalar_qdq; IEEE fp32 dots when an operand is NVFP4. V groups are block-16 at absolute token positions (split determinism); the open tail block re-quantizes as the sequence grows (the on-read tradeoff vs write-once, documented). Integration: keep the latent cache raw — _QuantVLLMMLAAttention.forward skips the module kv_c/k_pe QDQ under the renamed _skip_module_kv_quant flag, and ModelOptMLAImpl no longer overrides do_kv_cache_update (removed the write-once/double-quant machinery from the prior two commits). quant_kw.decode carries k_qdq/v_qdq from k_format/v_format; forward_mqa passes them. Reference oracle + tests updated for independent on-read K/V quant; README rewritten for the quantized-BMM model. Verified in vllm/vllm-openai:v0.26.0 on RTX A5000 (SM86): decode/prefill kernel baselines, MLA + dense vLLM runtime suites, and dense kernel regression all pass (94 passed); the NVFP4/FP8 K/V/P decode branches compile past all structural stages and gate only at the SM89 fp8e4nv cast, so their numerics validate on the pending SM89 cluster run. Signed-off-by: Kai Xu <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Contributor
|
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2244 +/- ##
==========================================
- Coverage 78.95% 77.76% -1.20%
==========================================
Files 522 528 +6
Lines 60550 62468 +1918
==========================================
+ Hits 47810 48576 +766
- Misses 12740 13892 +1152
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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
1 participant
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 does this PR do?
Type of change: new feature
Adds fused NVFP4/FP8 Q/K/P/V and 2:4 attention quantization for MLA models
(DeepSeek-family) on the vLLM
TRITON_MLAbackend, closing the gap where thecompact attention-quant path (
install_vllm_nvfp4_attention) supported onlyregular GQA/MHA attention.
The design follows the quantized-BMM-operand model (matching the dense
Q/K/P/V path and the MNI reference): the latent KV cache stays raw, and every
attention BMM operand is fake-quantized once, in-kernel, along its own
contraction axis.
mla_decode.py, split-K absorbed MQA): reads the raw latent andquantizes K along the feature axis (
k_format), V along the token axis(
v_format) — independently — plus P after the softmax row-sum and acaller-side FP32-carrier Q. V groups are block-16 at absolute token positions
(split-deterministic).
mla_prefill.py, varlen, asymmetric 192/128 head dims): quantizesthe projected Q/K/P/V operands once from the bf16 latent, with optional
token-exact N:M (2:4) score sparsity on new-token attention.
ModelOptMLAImplreclassesTritonMLAImpland overridesforward_mha/forward_mqa;_QuantVLLMMLAAttentionskips the module-levellatent QDQ (
_skip_module_kv_quant) so the cache stays raw; the installerdiscovers
MLAAttentionwith MLA-specific gates (fp8 cache, sparse indexer,q-padding, skip-softmax rejected) under the existing validate-then-publish
flow.
Contract notes: NVFP4 is 1×16 along the contraction axis (E4M3 block scale,
amax/(6·448)); the softmax denominator stays unquantized; on-read V trades thewrite-once cache's tail-block stability for an independent
v_format— thefaithful quantized-BMM behavior.
Usage
Testing
Verified in the CI-pinned
vllm/vllm-openai:v0.26.0container on an RTX A5000(SM86):
runtime installer/adapter suites, and the dense attention kernel regression:
94 passed (E4M3 quant-path tests are compute-capability-gated and skipped
on SM86).
structural stages and gate only at the SM89
fp8e4nvcast, so their numericsvalidate on an SM89+ run.
Pending on SM89+ hardware (draft): the E4M3 numeric golden tests
(decode/prefill K/V/P quant vs. eager oracles) and the tiny-DeepSeek
install→generate end-to-end.
Before your PR is "Ready for review"
previously unsupported by the compact attention installer)
you follow guidance in
CONTRIBUTING.md: N/Aoracles + synthetic vLLM runtime/adapter tests)
Additional Information
Draft: functionally complete and structurally verified on SM86; kept as draft
until the E4M3 numeric suite runs on SM89+ hardware.