diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2686c8828ac..2f1f856650f 100755 --- a/CHANGELOG.rst +++ b/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 `_ 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. diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 08f145a25a4..e02d4e5f01d 100644 --- a/examples/vllm_serve/README.md +++ b/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=,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 -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 `, or set MLflow's own `MLFLOW_TRACKING_URI`, to record what diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index 4fa5997027d..f98c6f01051 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/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 | diff --git a/modelopt_recipes/configs/ptq/units/kv_nvfp4_mla.yaml b/modelopt_recipes/configs/ptq/units/kv_nvfp4_mla.yaml new file mode 100644 index 00000000000..3445db9da88 --- /dev/null +++ b/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 diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py index 9c951363e01..ecf24f9b6cd 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/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``.