Skip to content

Commit 81b0637

Browse files
committed
Add the kv_nvfp4_mla unit for NVFP4 MLA KV-cache fake quant in vLLM
NVFP4 on MLAAttention's kv_c_bmm_quantizer (the cached latent) with the global scale fixed to 1, plus a README example recipe. Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
1 parent 5008545 commit 81b0637

5 files changed

Lines changed: 143 additions & 0 deletions

File tree

‎CHANGELOG.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ Changelog
2828
- Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/onnx_ptq/bevformer>`_ for details.
2929
- Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``.
3030
- 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.
31+
- Add the ``configs/ptq/units/kv_nvfp4_mla`` recipe unit for NVFP4 fake quantization of the MLA latent KV cache (DeepSeek-V3, GLM-5.3-Flash, ...) in vLLM fake-quant serving. See `examples/vllm_serve/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/vllm_serve>`_ for an example recipe.
3132
- 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.
3233
- ``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.
3334
- 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‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,46 @@ Step 4 (Optional): using lm_eval to run evaluation
9898
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
9999
```
100100

101+
## Fake-quantize the MLA KV cache
102+
103+
MLA models such as DeepSeek-V3 and GLM-5.3-Flash cache one latent vector per token instead of
104+
separate keys and values. Its fake quantizer is `kv_c_bmm_quantizer` on vLLM's `MLAAttention`.
105+
`KV_QUANT_CFG` presets (e.g. `NVFP4_KV_CFG`) are extended to it automatically, but the KV-cache
106+
units in a recipe (`*[kv]_bmm_quantizer`) do not match it, so a recipe imports the
107+
`configs/ptq/units/kv_nvfp4_mla` unit instead. For example, `kv_nvfp4_mla_only.yaml` quantizes the
108+
latent KV cache alone (weights and activations stay unquantized):
109+
110+
```yaml
111+
# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe
112+
imports:
113+
base_disable_all: configs/ptq/units/base_disable_all
114+
kv_nvfp4_mla: configs/ptq/units/kv_nvfp4_mla
115+
116+
metadata:
117+
description: NVFP4 fake quantization of the MLA latent KV cache only.
118+
quantize:
119+
algorithm: max
120+
quant_cfg:
121+
- $import: base_disable_all
122+
- $import: kv_nvfp4_mla
123+
```
124+
125+
```bash
126+
RECIPE_PATH=kv_nvfp4_mla_only.yaml python vllm_serve_fakequant.py <model_path> -tp 8 \
127+
--host 0.0.0.0 --port 8000
128+
```
129+
130+
This recipe also runs on the FP8 GLM-5.3-Flash release, since it leaves the FP8 layers alone (add
131+
`--moe-backend auto`, see above).
132+
133+
Notes:
134+
135+
- The latent is fake-quantized before vLLM writes it to the cache, so attention over the tokens
136+
of the current step sees the quantized values too.
137+
- Serve with a BF16 KV cache: an FP8 cache would quantize the fake-quantized latent a second time.
138+
`--kv-cache-dtype auto` is BF16 unless the checkpoint declares a quantized KV cache (e.g. a
139+
ModelOpt export with an FP8 KV cache); then pass `--kv-cache-dtype bfloat16`.
140+
101141
## Tracking a serve with MLflow
102142

103143
Pass `--mlflow <tracking-uri>`, or set MLflow's own `MLFLOW_TRACKING_URI`, to record what

‎modelopt_recipes/configs/ptq/units/README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ recipes (under `general/` or `models/`) or presets (under `presets/`).
2525
| `kv_nvfp4_affine.yaml` | NVFP4 affine KV cache quantizer entries; supported on Blackwell+ GPUs |
2626
| `kv_nvfp4_cast.yaml` | NVFP4 KV cache with constant amax (skips KV calibration); supported on Blackwell+ GPUs |
2727
| `kv_nvfp4_rotate.yaml` | NVFP4 rotated KV cache quantizer entries; supported on Blackwell+ GPUs |
28+
| `kv_nvfp4_mla.yaml` | NVFP4 fake quantization of the vLLM MLA latent KV cache (`*kv_c_bmm_quantizer`), global scale fixed to 1; compute capability 8.9+ GPUs |
2829
| `mamba_moe_disabled_quantizers.yaml` | Shared Mamba-MoE quantizer exclusions |
2930
| `w8a8_fp8_fp8.yaml` | FP8 weight + activation quantizer entries (W8A8); supported on Hopper+ GPUs |
3031
| `w4a4_nvfp4_nvfp4.yaml` | NVFP4 weight + activation quantizer entries (W4A4); supported on Blackwell+ GPUs |
Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
# QuantizerCfgList snippet that enables NVFP4 (E2M1, block 16, E4M3 scales) fake quantization of the
17+
# MLA latent KV cache in vLLM: `kv_c_bmm_quantizer` on MLAAttention, the normalized latent that vLLM
18+
# caches and attention reads as both K and V. The KV-cache units (`*[kv]_bmm_quantizer`) do not
19+
# match it. The RoPE key part (`k_pe_bmm_quantizer`) stays unquantized; NoPE MLA models such as
20+
# GLM-5.3-Flash have none.
21+
# The global scale is fixed to 1 (amax = 6 * 448) instead of calibrated, so the per-block E4M3 scales
22+
# carry the whole range and no scale has to be calibrated or exported.
23+
# The NVFP4 fake-quant kernel needs compute capability 8.9+ (Ada, Hopper or newer).
24+
25+
# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig
26+
imports:
27+
nvfp4: configs/numerics/nvfp4
28+
---
29+
- quantizer_name: '*kv_c_bmm_quantizer'
30+
cfg:
31+
$import: nvfp4
32+
constant_amax: 2688.0

‎tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,9 @@
5353
from vllm.utils.import_utils import has_deep_gemm
5454

5555
import modelopt.torch.quantization as mtq
56+
from modelopt.torch.opt.config_loader import load_config
5657
from modelopt.torch.quantization.config import QuantizerAttributeConfig
58+
from modelopt.torch.quantization.conversion import set_quantizer_by_cfg
5759
from modelopt.torch.quantization.nn import TensorQuantizer
5860
from modelopt.torch.quantization.plugins import vllm as vllm_plugin
5961
from modelopt.torch.quantization.plugins.vllm import (
@@ -834,6 +836,73 @@ def test_update_kv_cfg_for_mla_skips_non_mla_and_warns_on_affine():
834836
}
835837

836838

839+
class _NativeMLAAttention(torch.nn.Module):
840+
def forward(self, query, kv_c, k_pe, *args, **kwargs):
841+
return query, kv_c, k_pe
842+
843+
844+
@pytest.mark.skipif(VllmMLAAttention is None, reason="this vLLM has no MLAAttention")
845+
def test_kv_nvfp4_mla_unit_quantizes_only_the_latent(monkeypatch):
846+
"""The unit fake-quantizes the cached MLA latent; the empty NoPE ``k_pe`` passes through."""
847+
monkeypatch.setattr(
848+
vllm_plugin,
849+
"create_parallel_state",
850+
lambda: vllm_plugin.ParallelState(data_parallel_group=None),
851+
)
852+
853+
class _TestQuantVLLMMLAAttention(vllm_plugin._QuantVLLMMLAAttention, _NativeMLAAttention):
854+
pass
855+
856+
mla = _new_attention(_TestQuantVLLMMLAAttention)
857+
mla._setup()
858+
set_quantizer_by_cfg(
859+
mla,
860+
[{"quantizer_name": "*", "enable": False}, *load_config("configs/ptq/units/kv_nvfp4_mla")],
861+
)
862+
863+
assert mla.kv_c_bmm_quantizer.is_enabled
864+
assert mla.kv_c_bmm_quantizer.amax == 6.0 * 448.0
865+
assert not mla.q_bmm_quantizer.is_enabled
866+
assert not mla.k_pe_bmm_quantizer.is_enabled
867+
868+
mla.to("cuda")
869+
query = torch.randn(5, 4, 512, device="cuda", dtype=torch.bfloat16)
870+
kv_c = torch.randn(5, 512, device="cuda", dtype=torch.bfloat16)
871+
k_pe = torch.empty(5, 1, 0, device="cuda", dtype=torch.bfloat16) # qk_rope_head_dim=0
872+
out_query, out_kv_c, out_k_pe = mla(query, kv_c, k_pe)
873+
assert out_query is query
874+
assert out_k_pe is k_pe
875+
assert not torch.equal(out_kv_c, kv_c)
876+
# NVFP4 with a fixed global scale is idempotent: the output is already on the grid.
877+
assert torch.equal(mla.kv_c_bmm_quantizer(out_kv_c), out_kv_c)
878+
879+
880+
def test_kv_nvfp4_mla_quantizer_replays_in_cuda_graph():
881+
"""The constant amax is created on the CPU; on the GPU (FakeQuantWorker moves every quantizer
882+
there before CUDA graph capture) the latent fake quant captures and replays like eager."""
883+
layer = torch.nn.Module()
884+
layer.kv_c_bmm_quantizer = TensorQuantizer()
885+
set_quantizer_by_cfg(layer, load_config("configs/ptq/units/kv_nvfp4_mla"))
886+
quantizer = layer.kv_c_bmm_quantizer
887+
assert quantizer._amax.device.type == "cpu"
888+
quantizer.to("cuda")
889+
890+
static_in = torch.randn(8, 512, device="cuda", dtype=torch.bfloat16)
891+
stream = torch.cuda.Stream()
892+
stream.wait_stream(torch.cuda.current_stream())
893+
with torch.cuda.stream(stream):
894+
quantizer(static_in) # compile the kernel before capture
895+
torch.cuda.current_stream().wait_stream(stream)
896+
graph = torch.cuda.CUDAGraph()
897+
with torch.cuda.graph(graph):
898+
static_out = quantizer(static_in)
899+
900+
new_in = torch.randn_like(static_in)
901+
static_in.copy_(new_in)
902+
graph.replay()
903+
assert torch.equal(static_out, quantizer(new_in))
904+
905+
837906
def _quantize_and_summarize(self):
838907
"""Run on the worker via ``LLM.collective_rpc``.
839908

0 commit comments

Comments
 (0)