Add NVFP4/FP8 Q/K/P/V + 2:4 attention quantization for MLA (vLLM TRITON_MLA)

.gitignore

@@ -20,6 +20,7 @@ htmlcov/
 .coverage.*
 coverage.xml
 .pytest_cache/
+.cache/
 
 # Sphinx documentation
 docs/build

CHANGELOG.rst

@@ -10,6 +10,8 @@ Changelog
 
 - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint.
 
+- Add fused NVFP4/FP8 Q/K/P/V and 2:4 attention quantization for MLA models (DeepSeek-family) on the vLLM ``TRITON_MLA`` backend, served through ``examples/vllm_serve`` (``QuantSparseAttnWorker`` / ``install_vllm_nvfp4_attention``). Each attention BMM operand is fake-quantized in-kernel from a raw latent cache — decode quantizes K (feature axis) and V (token axis) on read and P after the softmax row-sum; prefill quantizes the projected operands. Optional checkpoint N:M score sparsity applies to prefill new-token attention. Requires ``--attention-backend TRITON_MLA`` and ``--enforce-eager``.
+
 *Megatron Framework (M-LM / M-Bridge)*
 
 - Add SFT-masked data support to ``examples/megatron_bridge/distill.py``: ``--sft --sft_dataset_root <dir>`` distills on raw prompt-completion JSONL (``{"input", "output"}`` records) with the loss masked to the response tokens, using Megatron-Bridge's ``FinetuningDatasetConfig`` and the model's own HuggingFace tokenizer instead of the pretraining ``GPTDataset`` and ``NullTokenizer``.

examples/vllm_serve/README.md

@@ -219,7 +219,28 @@ K is QDQ before its cache write, while V is written pristine. Complete 16-token
 
 Supported configurations are regular decoder self-attention with FlashInfer or FlashAttention, fp16/bf16 model and KV cache, equal Q/K/V head dimensions that are multiples of 16, and DCP 1. The FlashInfer adapter preserves both NHD and HND cache strides and separates mixed decode/prefill launches so each phase keeps its own kernel contract. The default `FULL_AND_PIECEWISE` mode remains enabled for fixed N:M and attention-only NVFP4; checkpoints with calibrated decode `threshold_scale_factor` must use a non-`FULL` decode graph mode such as `--enforce-eager` because the live sequence length is not replayed as a Python scalar.
 
-Unsupported features are sliding window, ALiBi, softcap, sinks, FP8 KV cache, cross/encoder/MLA attention, KV sharing or transfer, prefix caching, speculative decoding, DBO/ubatching, and `FULL` mixed/prefill CUDA graphs.
+Unsupported features are sliding window, ALiBi, softcap, sinks, FP8 KV cache, cross/encoder attention, KV sharing or transfer, prefix caching, speculative decoding, DBO/ubatching, and `FULL` mixed/prefill CUDA graphs.
+
+#### MLA attention (DeepSeek-family)
+
+The quantized worker also supports dense MLA text attention on the `TRITON_MLA` backend (pass `--attention-backend TRITON_MLA` when it is not the platform default):
+
+```bash
+python vllm_serve_sparse_attn.py <MODEL_PATH> -tp 8 \
+  --no-enable-prefix-caching --enforce-eager \
+  --attention-backend TRITON_MLA \
+  --worker-cls sparse_attn_worker.QuantSparseAttnWorker
+```
+
+The MLA operand mapping differs from regular attention because the latent cache is shared by both BMMs:
+
+- `q_format`: prefill quantizes the projected 192-d query in-kernel; decode quantizes the absorbed `kv_lora_rank + rope`-d query (FP32 QDQ carrier). With `q_format=fp8` the module-level quantizer QDQs the pre-projection query instead.
+- `k_format`: governs the write-once latent-cache QDQ (`kv_c`, `k_pe`, applied at cache-write time by the impl) and the prefill projected K (in-kernel).
+- `v_format`: governs the prefill projected V (in-kernel) **only**. **Decode V is not independently quantized** — decode BMM2 consumes the write-once quantized latent as-is, so decode V inherits `k_format` (the latent's feature-axis quantization used for BMM1-K), not `v_format`. This is a deliberate consequence of the write-once, no-on-read-re-quant design (stable across steps/splits): a shared latent cannot be stored quantized along both K's feature axis and V's token axis at once. Reading it independently as a token-axis-quantized V would require an on-read re-quant (the MNI-style raw-cache decode model), trading away that stability. So for decode, treat the contract as "V inherits k_format," not independent V quantization.
+- The latent QDQ is applied at the cache-write hook (not the module forward), so the **new-tokens** prefill projection reads the bf16 latent and its K/V operands are quantized exactly once. Note one residual: cached-**context** prefill chunks gather from the already-quantized paged cache and re-quantize the projected operands, so those chunks are double-quantized — inherent to reading a stored quantized latent, and absent for short prompts that fit a single chunk.
+- `p_format`: fused into both the prefill and decode kernels; the softmax denominator stays unquantized and P amax defaults to 1.0.
+
+MLA decode uses a fixed 32-split, 32-key-tile schedule with tile boundaries at absolute token positions, so quantized decode results are stable as the sequence grows and reproducible across batch shapes and devices. Optional checkpoint N:M sparsity applies to prefill new-token attention only (cached-context chunks run dense); skip-softmax is rejected on MLA layers. Sparse-only installation ignores MLA layers. DeepSeek V3.2-style sparse-indexer MLA and FP8 latent caches are unsupported.
 
 ## Known Problems
 

examples/vllm_serve/vllm_ptq_utils.py

@@ -141,12 +141,14 @@ def update_kv_cfg_for_mla(model: torch.nn.Module, kv_quant_cfg: list) -> list:
     `k_bmm_quantizer` and `v_bmm_quantizer`. This function copies the
     config from `*[kv]_bmm_quantizer` to also cover `*kv_c_bmm_quantizer`.
     """
-    try:
-        from vllm.attention.layer import MLAAttention
-    except ImportError:
-        return kv_quant_cfg
-
-    if not any(isinstance(m, MLAAttention) for m in model.modules()):
+    # Resolved via the quant plugin, which handles the MLAAttention module
+    # moving across vLLM releases (vllm.attention.layer no longer exists in
+    # vLLM >= 0.26).
+    from modelopt.torch.quantization.plugins.vllm import VllmMLAAttention
+
+    if VllmMLAAttention is None or not any(
+        isinstance(m, VllmMLAAttention) for m in model.modules()
+    ):
         return kv_quant_cfg
 
     kv_entry = next(

examples/vllm_serve/vllm_reload_utils.py

@@ -143,6 +143,14 @@ def _convert_key_for_vllm(key: str, value: Any) -> tuple[str, str | None, Any]:
         )
         return ("group", group_key, value)
 
+    # MLA latent quantizers: self_attn.kv_c/k_pe_bmm_quantizer ->
+    # self_attn.mla_attn.mla_attn.* (DeepSeek-style wrapper nesting in vLLM).
+    # vLLM-native keys that already carry the mla_attn prefix copy as-is below.
+    mla_bmm_match = re.search(r"(.*\.self_attn)\.((?:kv_c|k_pe)_bmm_quantizer.*)$", key)
+    if mla_bmm_match:
+        new_key = mla_bmm_match.group(1) + ".mla_attn.mla_attn." + mla_bmm_match.group(2)
+        return ("copy", new_key, value)
+
     # Transform bmm_quantizer keys: self_attn.q/k/v_bmm_quantizer -> self_attn.attn.q/k/v_bmm_quantizer
     bmm_match = re.search(r"(.*\.self_attn)\.([qkv]_bmm_quantizer.*)$", key) or re.search(
         r"(.*\.mixer)\.([qkv]_bmm_quantizer.*)$", key

modelopt/torch/kernels/quantization/attention/bmm2_qdq.py

@@ -13,12 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""NVFP4 operand helpers for the attention ``P @ V`` matmul (BMM2).
+"""NVFP4 operand helpers for the attention BMM matmuls.
 
-P and V share the low-level ``nvfp4_scalar_qdq`` primitive, but retain thin
-operand-specific wrappers because their layouts and amax reductions differ.
-P is nonnegative with layout ``[M, K]``; V is signed with layout ``[K, N]``.
-Both use block-16 scaling along the BMM2 contraction axis.
+P, V, and the signed A-side share the low-level ``nvfp4_scalar_qdq``
+primitive, but retain thin operand-specific wrappers because their layouts
+and amax reductions differ. P is nonnegative with layout ``[M, K]``; V is
+signed with layout ``[K, N]``; the signed A-side (Q of BMM1) is ``[M, K]``.
+All use block-16 scaling along the BMM contraction axis.
 """
 
 import math
@@ -70,6 +71,21 @@ def _v_qdq_nvfp4(v, global_scale, BLOCK_N: tl.constexpr, BLOCK_D: tl.constexpr):
     return tl.reshape(nvfp4_scalar_qdq(grouped, block_amax, global_scale, 16), (BLOCK_N, BLOCK_D))
 
 
+@triton.jit
+def _a_qdq_nvfp4(x, global_scale, BLOCK_M: tl.constexpr, BLOCK_K: tl.constexpr):
+    """Fake-quantize a signed A-side operand ``[M, K]`` in block-16 groups along K.
+
+    The BMM1 Q-operand counterpart of :func:`_p_qdq_nvfp4`: same ``[M, K]``
+    layout and contraction-axis blocking, but signed, so the block amax uses
+    ``abs``. Zero-padded lanes (masked loads) form all-zero blocks that
+    ``nvfp4_scalar_qdq`` guards to zero.
+    """
+    tl.static_assert(BLOCK_K % 16 == 0, "BLOCK_K must be divisible by 16 for NVFP4")
+    grouped = tl.reshape(x, (BLOCK_M, BLOCK_K // 16, 16))
+    block_amax = tl.expand_dims(tl.max(tl.abs(grouped), axis=2), 2)
+    return tl.reshape(nvfp4_scalar_qdq(grouped, block_amax, global_scale, 16), (BLOCK_M, BLOCK_K))
+
+
 @triton.jit
 def _fake_quant_v_onwrite_kernel(
     V_cache,

modelopt/torch/kernels/quantization/attention/mla/__init__.py

@@ -0,0 +1,21 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MLA attention kernels with fused fake quantization (prefill and decode)."""
+
+from .mla_decode import mla_attention_decode
+from .mla_prefill import mla_prefill_attention
+
+__all__ = ["mla_attention_decode", "mla_prefill_attention"]