Add vLLM NVFP4 MLA KV-cache fake quant

CHANGELOG.rst

@@ -30,6 +30,7 @@ Changelog
 - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``.
 - Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training through the standard ``quant_cfg`` interface. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer.
 - Add fake quantization of the sparse-attention indexer key cache and query for DeepSeek-V4 (vLLM and Megatron-Core) and GLM-5.3-Flash (vLLM) through the new ``indexer_k_quantizer`` and ``indexer_q_quantizer``. Enable them by importing the ``configs/ptq/units/indexer_k_nvfp4`` and ``configs/ptq/units/indexer_q_nvfp4`` units (NVFP4 with the global scale fixed to 1) into a recipe.
+- Add the ``configs/ptq/units/kv_nvfp4_mla`` recipe unit for fake quantization of the MLA KV cache (DeepSeek-V3, GLM-5.3-Flash, ...) in vLLM fake-quant serving: NVFP4 for the latent and FP8 for the RoPE key. See `examples/vllm_serve/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/vllm_serve>`_ for an example recipe.
 - vLLM fake-quant serving now runs on pre-quantized checkpoints such as FP8 when the recipe leaves those layers unquantized (for example a KV-cache-only recipe), and on MLA models with an FP8 KV cache; both previously failed during quantization.
 - ``KV_QUANT_CFG`` presets in vLLM fake-quant serving now quantize the MLA KV cache; on vLLM 0.16 and later they silently quantized nothing for MLA models.
 - vLLM fake-quant serving of NVFP4 and cast-mode KV-cache configs no longer needs ``--enforce-eager``. ``examples/vllm_serve/vllm_serve_fakequant.py`` now sets ``VLLM_DISABLE_COMPILE_CACHE=1`` by default, because a cached torch.compile graph of the same model without the fake quant would otherwise be reused.

examples/vllm_serve/README.md

@@ -98,6 +98,48 @@ Step 4 (Optional): using lm_eval to run evaluation
 lm_eval --model local-completions --tasks gsm8k --model_args model=<model_name>,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=1,max_retries=3,tokenized_requests=False,batch_size=128,tokenizer_backend=None
 ```
 
+## Fake-quantize the MLA KV cache
+
+MLA models such as DeepSeek-V3 and GLM-5.3-Flash cache one latent vector per token instead of
+separate keys and values; RoPE models such as DeepSeek-V3 also cache a small RoPE key. Their fake
+quantizers are `kv_c_bmm_quantizer` and `k_pe_bmm_quantizer` on vLLM's `MLAAttention`.
+`KV_QUANT_CFG` presets (e.g. `NVFP4_KV_CFG`) are extended to both automatically, with the same
+format for both. The KV-cache units in a recipe (`*[kv]_bmm_quantizer`) do not match them, so a
+recipe imports the `configs/ptq/units/kv_nvfp4_mla` unit instead, which uses NVFP4 for the latent
+and FP8 for the RoPE key. For example, `kv_nvfp4_mla_only.yaml` quantizes the MLA KV cache alone
+(weights and activations stay unquantized):
+
+```yaml
+# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe
+imports:
+  base_disable_all: configs/ptq/units/base_disable_all
+  kv_nvfp4_mla: configs/ptq/units/kv_nvfp4_mla
+
+metadata:
+  description: Fake quantization of the MLA KV cache only (NVFP4 latent, FP8 RoPE key).
+quantize:
+  algorithm: max
+  quant_cfg:
+    - $import: base_disable_all
+    - $import: kv_nvfp4_mla
+```
+
+```bash
+RECIPE_PATH=kv_nvfp4_mla_only.yaml python vllm_serve_fakequant.py <model_path> -tp 8 \
+  --host 0.0.0.0 --port 8000
+```
+
+This recipe also runs on the FP8 GLM-5.3-Flash release, since it leaves the FP8 layers alone (add
+`--moe-backend auto`, see above).
+
+Notes:
+
+- The latent and RoPE key are fake-quantized before vLLM writes them to the cache, so attention
+  over the tokens of the current step sees the quantized values too.
+- Serve with a BF16 KV cache: an FP8 cache would quantize the fake-quantized latent a second time.
+  `--kv-cache-dtype auto` is BF16 unless the checkpoint declares a quantized KV cache (e.g. a
+  ModelOpt export with an FP8 KV cache); then pass `--kv-cache-dtype bfloat16`.
+
 ## Tracking a serve with MLflow
 
 Pass `--mlflow <tracking-uri>`, or set MLflow's own `MLFLOW_TRACKING_URI`, to record what

modelopt_recipes/configs/ptq/units/README.md

@@ -25,6 +25,7 @@ recipes (under `general/` or `models/`) or presets (under `presets/`).
 | `kv_nvfp4_affine.yaml` | NVFP4 affine KV cache quantizer entries; supported on Blackwell+ GPUs |
 | `kv_nvfp4_cast.yaml` | NVFP4 KV cache with constant amax (skips KV calibration); supported on Blackwell+ GPUs |
 | `kv_nvfp4_rotate.yaml` | NVFP4 rotated KV cache quantizer entries; supported on Blackwell+ GPUs |
+| `kv_nvfp4_mla.yaml` | Fake quantization of the vLLM MLA KV cache: NVFP4 latent (`*kv_c_bmm_quantizer`, global scale fixed to 1) and FP8 RoPE key (`*k_pe_bmm_quantizer`, scale 1); compute capability 8.9+ GPUs |
 | `mamba_moe_disabled_quantizers.yaml` | Shared Mamba-MoE quantizer exclusions |
 | `w8a8_fp8_fp8.yaml` | FP8 weight + activation quantizer entries (W8A8); supported on Hopper+ GPUs |
 | `w4a4_nvfp4_nvfp4.yaml` | NVFP4 weight + activation quantizer entries (W4A4); supported on Blackwell+ GPUs |

modelopt_recipes/configs/ptq/units/kv_nvfp4_mla.yaml

@@ -0,0 +1,39 @@
+# 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.
+
+# QuantizerCfgList snippet that fake-quantizes the vLLM MLA KV cache in the formats of vLLM's
+# `nvfp4_ds_mla` cache. It sets two MLAAttention quantizers, which the KV-cache units
+# (`*[kv]_bmm_quantizer`) do not match:
+# - `kv_c_bmm_quantizer`, the normalized latent that vLLM caches and attention reads as both K and
+#   V: NVFP4 (E2M1, block 16, E4M3 scales) with the global scale fixed to 1 (amax = 6 * 448), so the
+#   per-block E4M3 scales carry the whole range.
+# - `k_pe_bmm_quantizer`, the RoPE key: FP8 E4M3 with scale 1 (amax = 448). NoPE MLA models such as
+#   GLM-5.3-Flash have an empty k_pe, which passes through unchanged.
+# No scale is calibrated or exported. The NVFP4 fake-quant kernel needs compute capability 8.9+
+# (Ada, Hopper or newer).
+
+# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig
+imports:
+  nvfp4: configs/numerics/nvfp4
+  fp8: configs/numerics/fp8
+---
+  - quantizer_name: '*kv_c_bmm_quantizer'
+    cfg:
+      $import: nvfp4
+      constant_amax: 2688.0
+  - quantizer_name: '*k_pe_bmm_quantizer'
+    cfg:
+      $import: fp8
+      use_constant_amax: true

tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py

@@ -53,7 +53,9 @@
 from vllm.utils.import_utils import has_deep_gemm
 
 import modelopt.torch.quantization as mtq
+from modelopt.torch.opt.config_loader import load_config
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
+from modelopt.torch.quantization.conversion import set_quantizer_by_cfg
 from modelopt.torch.quantization.nn import TensorQuantizer
 from modelopt.torch.quantization.plugins import vllm as vllm_plugin
 from modelopt.torch.quantization.plugins.vllm import (
@@ -841,6 +843,86 @@ def test_update_kv_cfg_for_mla_skips_non_mla_and_warns_on_affine():
     }
 
 
+class _NativeMLAAttention(torch.nn.Module):
+    def forward(self, query, kv_c, k_pe, *args, **kwargs):
+        return query, kv_c, k_pe
+
+
+@pytest.mark.skipif(VllmMLAAttention is None, reason="this vLLM has no MLAAttention")
+@pytest.mark.parametrize("rope_dim", [0, 64], ids=("nope", "rope"))  # GLM-5.3-Flash, DeepSeek-V3
+def test_kv_nvfp4_mla_unit_quantizes_the_mla_kv_cache(monkeypatch, rope_dim):
+    """Like vLLM's ``nvfp4_ds_mla`` cache: an NVFP4 latent and an unscaled FP8 RoPE key ``k_pe``.
+    An empty NoPE ``k_pe`` passes through."""
+    monkeypatch.setattr(
+        vllm_plugin,
+        "create_parallel_state",
+        lambda: vllm_plugin.ParallelState(data_parallel_group=None),
+    )
+
+    class _TestQuantVLLMMLAAttention(vllm_plugin._QuantVLLMMLAAttention, _NativeMLAAttention):
+        pass
+
+    mla = _new_attention(_TestQuantVLLMMLAAttention)
+    mla._setup()
+    set_quantizer_by_cfg(
+        mla,
+        [{"quantizer_name": "*", "enable": False}, *load_config("configs/ptq/units/kv_nvfp4_mla")],
+    )
+
+    assert mla.kv_c_bmm_quantizer.is_enabled
+    assert mla.kv_c_bmm_quantizer.amax == 6.0 * 448.0
+    assert mla.k_pe_bmm_quantizer.is_enabled
+    assert not mla.q_bmm_quantizer.is_enabled
+
+    mla.to("cuda")
+    query = torch.randn(5, 4, 512, device="cuda", dtype=torch.bfloat16)
+    kv_c = torch.randn(5, 512, device="cuda", dtype=torch.bfloat16)
+    k_pe = torch.randn(5, 1, rope_dim, device="cuda", dtype=torch.bfloat16)
+    out_query, out_kv_c, out_k_pe = mla(query, kv_c, k_pe)
+    assert out_query is query
+    assert not torch.equal(out_kv_c, kv_c)
+    # NVFP4 with a fixed global scale is idempotent: the output is already on the grid.
+    assert torch.equal(mla.kv_c_bmm_quantizer(out_kv_c), out_kv_c)
+    if rope_dim:
+        assert not torch.equal(out_k_pe, k_pe)
+        assert torch.equal(out_k_pe, k_pe.to(torch.float8_e4m3fn).to(torch.bfloat16))
+    else:
+        assert out_k_pe is k_pe
+
+
+@pytest.mark.parametrize(
+    ("name", "shape"),
+    [("kv_c_bmm_quantizer", (8, 512)), ("k_pe_bmm_quantizer", (8, 1, 64))],
+    ids=("kv_c", "k_pe"),
+)
+def test_kv_nvfp4_mla_quantizer_replays_in_cuda_graph(name, shape):
+    """The latent's constant amax is created on the CPU; on the GPU (FakeQuantWorker moves every
+    quantizer there before CUDA graph capture) each fake quant of the unit captures and replays
+    like eager."""
+    layer = torch.nn.Module()
+    setattr(layer, name, TensorQuantizer())
+    set_quantizer_by_cfg(layer, load_config("configs/ptq/units/kv_nvfp4_mla"))
+    quantizer = getattr(layer, name)
+    if name == "kv_c_bmm_quantizer":
+        assert quantizer._amax.device.type == "cpu"
+    quantizer.to("cuda")
+
+    static_in = torch.randn(*shape, device="cuda", dtype=torch.bfloat16)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        quantizer(static_in)  # compile the kernel before capture
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        static_out = quantizer(static_in)
+
+    new_in = torch.randn_like(static_in)
+    static_in.copy_(new_in)
+    graph.replay()
+    assert torch.equal(static_out, quantizer(new_in))
+
+
 def _quantize_and_summarize(self):
     """Run on the worker via ``LLM.collective_rpc``.