Skip to content

Commit 6a10687

Browse files
sychen52claude
andauthored
Add vLLM NVFP4 MLA KV-cache fake quant (#2620)
### What does this PR do? Type of change: New feature NVFP4 fake quant of the MLA latent KV cache (GLM-5.3-Flash, DeepSeek-V3) in vLLM fake-quant serving: - Add the `configs/ptq/units/kv_nvfp4_mla` unit: NVFP4 on `*kv_c_bmm_quantizer` (the cached latent) with the global scale fixed to 1. The existing KV units (`*[kv]_bmm_quantizer`) don't match it. The RoPE key part (`k_pe`, e.g. in DeepSeek-V3) stays unquantized. - Add an example recipe to `examples/vllm_serve/README.md`. ### Usage ```bash # recipe: examples/vllm_serve/README.md RECIPE_PATH=kv_nvfp4_mla_only.yaml python vllm_serve_fakequant.py zai-org/GLM-5.3-Flash -tp 4 --moe-backend auto ``` ### Testing B200, vLLM 0.30: - `tests/gpu_vllm/torch/quantization`: 51 passed, including the 2 new tests. The 1 error, `test_tiny_deepseek_mla_quantize`, already fails on vLLM 0.30. - End to end with the README recipe on tiny random-weight models (run on an earlier revision of this branch; the unit and recipe are unchanged): - GLM-5.3-Flash (BF16 and FP8; FP8 also at TP2) and DeepSeek-V3: every latent row written to the cache outside CUDA-graph replay is on the NVFP4 grid. - GLM-5.3-Flash at TP1: CUDA-graph logprobs match eager. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ ### Additional Information Related: #2552, #2539. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added support for serving pre-quantized checkpoints when recipes leave their existing quantized layers unchanged. Recipes that attempt to fake-quantize those layers are not supported. * Added NVFP4 fake quantization for MLA latent KV caches and support for MLA models with FP8 KV caches. * Added serving guidance for MoE checkpoints and MLA KV-cache configurations, including the vLLM 0.16+ preset. * **Bug Fixes** * NVFP4 and cast-mode KV-cache serving no longer requires `--enforce-eager`. * The example serving launcher disables vLLM’s compile cache by default, unless it is already configured. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shiyang Chen <shiychen@nvidia.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
1 parent a57c5eb commit 6a10687

5 files changed

Lines changed: 165 additions & 0 deletions

File tree

‎CHANGELOG.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ Changelog
3030
- Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``.
3131
- 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.
3232
- 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.
33+
- 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.
3334
- 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.
3435
- ``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.
3536
- 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: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,48 @@ 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; RoPE models such as DeepSeek-V3 also cache a small RoPE key. Their fake
105+
quantizers are `kv_c_bmm_quantizer` and `k_pe_bmm_quantizer` on vLLM's `MLAAttention`.
106+
`KV_QUANT_CFG` presets (e.g. `NVFP4_KV_CFG`) are extended to both automatically, with the same
107+
format for both. The KV-cache units in a recipe (`*[kv]_bmm_quantizer`) do not match them, so a
108+
recipe imports the `configs/ptq/units/kv_nvfp4_mla` unit instead, which uses NVFP4 for the latent
109+
and FP8 for the RoPE key. For example, `kv_nvfp4_mla_only.yaml` quantizes the MLA KV cache alone
110+
(weights and activations stay unquantized):
111+
112+
```yaml
113+
# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe
114+
imports:
115+
base_disable_all: configs/ptq/units/base_disable_all
116+
kv_nvfp4_mla: configs/ptq/units/kv_nvfp4_mla
117+
118+
metadata:
119+
description: Fake quantization of the MLA KV cache only (NVFP4 latent, FP8 RoPE key).
120+
quantize:
121+
algorithm: max
122+
quant_cfg:
123+
- $import: base_disable_all
124+
- $import: kv_nvfp4_mla
125+
```
126+
127+
```bash
128+
RECIPE_PATH=kv_nvfp4_mla_only.yaml python vllm_serve_fakequant.py <model_path> -tp 8 \
129+
--host 0.0.0.0 --port 8000
130+
```
131+
132+
This recipe also runs on the FP8 GLM-5.3-Flash release, since it leaves the FP8 layers alone (add
133+
`--moe-backend auto`, see above).
134+
135+
Notes:
136+
137+
- The latent and RoPE key are fake-quantized before vLLM writes them to the cache, so attention
138+
over the tokens of the current step sees the quantized values too.
139+
- Serve with a BF16 KV cache: an FP8 cache would quantize the fake-quantized latent a second time.
140+
`--kv-cache-dtype auto` is BF16 unless the checkpoint declares a quantized KV cache (e.g. a
141+
ModelOpt export with an FP8 KV cache); then pass `--kv-cache-dtype bfloat16`.
142+
101143
## Tracking a serve with MLflow
102144

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

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

Lines changed: 82 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 (
@@ -841,6 +843,86 @@ def test_update_kv_cfg_for_mla_skips_non_mla_and_warns_on_affine():
841843
}
842844

843845

846+
class _NativeMLAAttention(torch.nn.Module):
847+
def forward(self, query, kv_c, k_pe, *args, **kwargs):
848+
return query, kv_c, k_pe
849+
850+
851+
@pytest.mark.skipif(VllmMLAAttention is None, reason="this vLLM has no MLAAttention")
852+
@pytest.mark.parametrize("rope_dim", [0, 64], ids=("nope", "rope")) # GLM-5.3-Flash, DeepSeek-V3
853+
def test_kv_nvfp4_mla_unit_quantizes_the_mla_kv_cache(monkeypatch, rope_dim):
854+
"""Like vLLM's ``nvfp4_ds_mla`` cache: an NVFP4 latent and an unscaled FP8 RoPE key ``k_pe``.
855+
An empty NoPE ``k_pe`` passes through."""
856+
monkeypatch.setattr(
857+
vllm_plugin,
858+
"create_parallel_state",
859+
lambda: vllm_plugin.ParallelState(data_parallel_group=None),
860+
)
861+
862+
class _TestQuantVLLMMLAAttention(vllm_plugin._QuantVLLMMLAAttention, _NativeMLAAttention):
863+
pass
864+
865+
mla = _new_attention(_TestQuantVLLMMLAAttention)
866+
mla._setup()
867+
set_quantizer_by_cfg(
868+
mla,
869+
[{"quantizer_name": "*", "enable": False}, *load_config("configs/ptq/units/kv_nvfp4_mla")],
870+
)
871+
872+
assert mla.kv_c_bmm_quantizer.is_enabled
873+
assert mla.kv_c_bmm_quantizer.amax == 6.0 * 448.0
874+
assert mla.k_pe_bmm_quantizer.is_enabled
875+
assert not mla.q_bmm_quantizer.is_enabled
876+
877+
mla.to("cuda")
878+
query = torch.randn(5, 4, 512, device="cuda", dtype=torch.bfloat16)
879+
kv_c = torch.randn(5, 512, device="cuda", dtype=torch.bfloat16)
880+
k_pe = torch.randn(5, 1, rope_dim, device="cuda", dtype=torch.bfloat16)
881+
out_query, out_kv_c, out_k_pe = mla(query, kv_c, k_pe)
882+
assert out_query is query
883+
assert not torch.equal(out_kv_c, kv_c)
884+
# NVFP4 with a fixed global scale is idempotent: the output is already on the grid.
885+
assert torch.equal(mla.kv_c_bmm_quantizer(out_kv_c), out_kv_c)
886+
if rope_dim:
887+
assert not torch.equal(out_k_pe, k_pe)
888+
assert torch.equal(out_k_pe, k_pe.to(torch.float8_e4m3fn).to(torch.bfloat16))
889+
else:
890+
assert out_k_pe is k_pe
891+
892+
893+
@pytest.mark.parametrize(
894+
("name", "shape"),
895+
[("kv_c_bmm_quantizer", (8, 512)), ("k_pe_bmm_quantizer", (8, 1, 64))],
896+
ids=("kv_c", "k_pe"),
897+
)
898+
def test_kv_nvfp4_mla_quantizer_replays_in_cuda_graph(name, shape):
899+
"""The latent's constant amax is created on the CPU; on the GPU (FakeQuantWorker moves every
900+
quantizer there before CUDA graph capture) each fake quant of the unit captures and replays
901+
like eager."""
902+
layer = torch.nn.Module()
903+
setattr(layer, name, TensorQuantizer())
904+
set_quantizer_by_cfg(layer, load_config("configs/ptq/units/kv_nvfp4_mla"))
905+
quantizer = getattr(layer, name)
906+
if name == "kv_c_bmm_quantizer":
907+
assert quantizer._amax.device.type == "cpu"
908+
quantizer.to("cuda")
909+
910+
static_in = torch.randn(*shape, device="cuda", dtype=torch.bfloat16)
911+
stream = torch.cuda.Stream()
912+
stream.wait_stream(torch.cuda.current_stream())
913+
with torch.cuda.stream(stream):
914+
quantizer(static_in) # compile the kernel before capture
915+
torch.cuda.current_stream().wait_stream(stream)
916+
graph = torch.cuda.CUDAGraph()
917+
with torch.cuda.graph(graph):
918+
static_out = quantizer(static_in)
919+
920+
new_in = torch.randn_like(static_in)
921+
static_in.copy_(new_in)
922+
graph.replay()
923+
assert torch.equal(static_out, quantizer(new_in))
924+
925+
844926
def _quantize_and_summarize(self):
845927
"""Run on the worker via ``LLM.collective_rpc``.
846928

0 commit comments

Comments
 (0)