From 88a2e7ea98c089bb0136a44bfaaa94371cae65a5 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Thu, 8 Oct 2026 12:57:43 -0700 Subject: [PATCH 1/4] Add vLLM GDN/KDA state-only fake quantization Signed-off-by: Kai Xu --- CHANGELOG.rst | 3 + examples/llm_qat/linear_attention/README.md | 6 + examples/vllm_serve/README.md | 93 ++++++ examples/vllm_serve/fakequant_worker.py | 26 +- .../linear_attention_state_int8.yaml | 15 + examples/vllm_serve/vllm_reload_utils.py | 45 ++- .../nn/modules/tensor_quantizer.py | 7 +- .../torch/quantization/plugins/__init__.py | 1 + .../plugins/vllm_linear_attention.py | 176 ++++++++++ .../test_vllm_linear_attention.py | 315 ++++++++++++++++++ .../quantization/test_tensor_quantizer_cpu.py | 13 + 11 files changed, 687 insertions(+), 13 deletions(-) create mode 100644 examples/vllm_serve/linear_attention_state_int8.yaml create mode 100644 modelopt/torch/quantization/plugins/vllm_linear_attention.py create mode 100644 tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index fd97d8637df..9d1d981e109 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,6 +16,9 @@ Changelog - Added experimental NVFP4 recurrent-state fake quantization using TensorQuantizer block16 E4M3 scales for linear-attention QAT. - Add optional ReplaySSM state QAT for GDN/KDA using the compatible INT8/Hadamard serving fork. The replay recipe controls checkpoint refresh frequency independently of chunked prefill. + +- Add experimental FP8/INT8 recurrent-state fake quantization before native vLLM GDN/KDA prefill and decode; use the state-only recipe with the documented vLLM 0.15.x eager runtime. + - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md index 1c2aef05421..ff37485fa05 100644 --- a/examples/llm_qat/linear_attention/README.md +++ b/examples/llm_qat/linear_attention/README.md @@ -321,3 +321,9 @@ GEMM quantization and approximate inverse are outside this example. See [ReplaySSM state QAT](REPLAYSSM.md) for the optional INT8/Hadamard fork, its checkpoint window, and the proposed two-precision extension. + +## State-only serving with vLLM + +The [vLLM example](../../vllm_serve/README.md#linear-attention-state-quantization) +uses the same state QDQ helper for plain FP8/INT8 policies. Match the runtime, +quantizer grouping, and scheduler prefill boundaries to the training setup. diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 98bfe45c3ce..51a75aef760 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -9,6 +9,10 @@ compact NVFP4 attention worker documented below requires vLLM 0.15.0 or newer. The fakequant launcher does not support vLLM 0.9.0. Use one of the tested releases above. +For GDN/KDA recurrent-state QDQ before native prefill and decode, see the +[linear-attention example](#linear-attention-state-quantization), which +documents the initial vLLM 0.15.x runtime requirements and state-only INT8 recipe. + ## Prepare environment Run the commands below from the ModelOpt repository root (`/workspace/Model-Optimizer` @@ -463,6 +467,95 @@ Supported configurations are regular decoder self-attention with FlashInfer or F 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. +## Linear-attention state quantization + +The existing `FakeQuantWorker` adds a ModelOpt `TensorQuantizer` immediately before +native vLLM prefill and decode for `Qwen3NextGatedDeltaNet` and +`KimiDeltaAttention`. Native attention kernels, projections, convolution, gates, +outputs, and cache management remain in use. + +### Runtime and launch + +The initial adapter targets vLLM 0.15.x V1 with key-first +`[slot, head, Dk, Dv]` FP32 recurrent state. Runtime checks reject other layouts +and versions. Use a vLLM 0.15.x environment for this adapter; the newer +versions supported by the general fakequant example use different model wrappers. + +```bash +PYTHONPATH=.:examples/vllm_serve \ +RECIPE_PATH=examples/vllm_serve/linear_attention_state_int8.yaml \ +bash examples/llm_qat/linear_attention/with_vllm_defaults.sh \ +python examples/vllm_serve/vllm_serve_fakequant.py /path/to/model \ + --tensor-parallel-size 2 --enforce-eager --no-async-scheduling \ + --no-enable-prefix-caching --mamba-cache-dtype float32 +``` + +The example enables signed symmetric dynamic INT8 state QDQ. Import the `configs/ptq/units/gdn_state_fp8_dynamic` quantizer unit for GDN FP8 +E4M3, or set `num_bits: [4, 3]` on the corresponding state quantizers. Each invocation +receives `[active_sequence, local_head, Dk, Dv]`; `axis: [0, 1]` retains one +scale per sequence and local head **within each value-column tile**. The shared +training/serving QDQ helper splits the last axis using `state_block_v` (default +64), then reduces over `[Dk, tile_width]`. State tensors remain FP32. + +For per-key 32-value INT8 blocks, use the existing recipe +`modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml`. +Its `block_sizes: {-1: 32}` is handled directly by TensorQuantizer and overrides +legacy tile grouping. The same recipe can be used for state QAT. + +Dynamic scales need no calibration dataset; `algorithm: null` skips dataset +loading. Weight/activation calibration can use the worker's existing recipe +and calibration loop. Static state calibration is unsupported in this adapter. + +Use `MODELOPT_STATE_PATH=/path/to/modelopt_state.pt` instead of a recipe to +restore quantizer configuration. Reload maps recurrent-state quantizer names +through the model's HF-to-vLLM mapper. Checkpoint weights must already match +the vLLM model architecture. + +### Quantization boundaries + +- Prefill: vLLM gathers initial states and zeros fresh requests. The wrapper + applies `TensorQuantizer` to this tensor and calls the original prefill kernel. +- Decode: the wrapper gathers active native cache slots, quantizes them, writes + them back, and calls the original decode kernel. Other slots are untouched. +- Native kernels compute outputs and update recurrent state normally. No + additional rounding is applied to the kernel's final-state write. +- The next call quantizes that state before reading it. A fresh zero state is + unchanged; prompt-to-decode rounding happens before the first decode call. +- Scheduler-level chunked prefill creates one QDQ boundary per invocation. + Internal kernel chunks do not create extra state-quantization boundaries. + Results can therefore depend on scheduler prompt-chunk sizes. +- Requests, cache-slot reuse, and preemption remain managed by vLLM. The plugin + allocates temporary gathered states, with no additional persistent state cache. +- TP ranks quantize their local heads independently, without scale all-reduce. + +The adapter accepts plain `precision: vllm` policies, including the legacy +`vllm_0_15` spelling. The serving engine supplies the phase boundaries, so no +training phase context is required. W quantization and the native ReplaySSM/ +Hadamard profile require separate serving integration and are rejected. + +Training and serving share state-QDQ formats and grouping. Numerical alignment +also requires the same vLLM runtime, arithmetic settings, and prefill scheduling: +a training prefix must reproduce each serving prompt-continuation boundary. +The QAT handoff rounds the prefix state before its first decode read; this +adapter rounds that same state before the first native decode call. QAT stores the +rounded next-state checkpoint after each update; serving stores the working +state and rounds it at the next read. These feed the same rounded state to +the next recurrence. Raw terminal cache buffers therefore need not match the +QAT checkpoint until the same QDQ is applied. + +The supported scope is eager synchronous execution with TP and PP=DP=CP=1. +Speculative decoding, prefix caching, state transfer, and CUDA graphs require +separate integration. This is floating-point numerical emulation; model-quality, +capacity, and performance claims require separate measurements. + +### Later prefill GEMM support + +Prefill operand QDQ will be a separate change after an optimized fused kernel +is available. It will reuse the eight numerical sites from the training prefill +implementation and preserve this native-cache wrapper. The PyTorch materialized +backend is not exposed in vLLM by this state-only adapter. State QDQ cadence and +prefill operand QDQ remain independent numerical policies. + ## Known Problems 1. **MCore reload does not use `MODELOPT_STATE_PATH`**; use `QUANT_FILE_PATH` and make sure `QUANT_CFG` matches the quantization recipe used for the original MCore model (otherwise quantizer keys/config won’t align). diff --git a/examples/vllm_serve/fakequant_worker.py b/examples/vllm_serve/fakequant_worker.py index 6bc82fbc22f..e3bbcadc7ba 100644 --- a/examples/vllm_serve/fakequant_worker.py +++ b/examples/vllm_serve/fakequant_worker.py @@ -38,6 +38,7 @@ disable_compilation, post_restore_vllm_parallel_linears, ) +from modelopt.torch.quantization.plugins.vllm_linear_attention import bind_vllm_linear_attention from modelopt.torch.utils import safe_load from modelopt.torch.utils.dataset_utils import get_dataset_dataloader @@ -114,17 +115,22 @@ def _fakequant_run_prolog_worker(self, mlflow_tracker: FakeQuantMlflowTracker) - print("Will load quant, so only do a single sample calibration") quant_config["calib_size"] = 1 - calib_dataloader = get_dataset_dataloader( - dataset_name=quant_config["dataset"], - tokenizer=tokenizer, - batch_size=quant_config["calib_batch_size"], - num_samples=quant_config["calib_size"], - device=self.device, - ) + quant_cfg = get_quant_config(quant_config, model) + calibrate_loop = None + if quant_cfg.get("algorithm", "max") is not None: + calib_dataloader = get_dataset_dataloader( + dataset_name=quant_config["dataset"], + tokenizer=tokenizer, + batch_size=quant_config["calib_batch_size"], + num_samples=quant_config["calib_size"], + device=self.device, + ) + worker_loop = calibrate_fun(calib_dataloader, self) - calibrate_loop = calibrate_fun(calib_dataloader, self) + def calibrate_loop(converted): + bind_vllm_linear_attention(converted, self.model_runner) + worker_loop(converted) - quant_cfg = get_quant_config(quant_config, model) # Before calibration, which is the run this artifact is most wanted for if it dies. mlflow_tracker.log_quant_config(quant_cfg) @@ -134,6 +140,7 @@ def _fakequant_run_prolog_worker(self, mlflow_tracker: FakeQuantMlflowTracker) - quantizer_file_path = quant_config["quant_file_path"] if quantizer_file_path: + bind_vllm_linear_attention(model, self.model_runner) self.model_runner._dummy_run(1) current_state_dict = load_state_dict_from_path(self, quantizer_file_path, model) model.load_state_dict(current_state_dict) @@ -148,6 +155,7 @@ def _fakequant_run_prolog_worker(self, mlflow_tracker: FakeQuantMlflowTracker) - if isinstance(module, TensorQuantizer): module.to(self.device) + bind_vllm_linear_attention(model, self.model_runner) if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: mtq.print_quant_summary(model) mlflow_tracker.log_quant_summary(model) diff --git a/examples/vllm_serve/linear_attention_state_int8.yaml b/examples/vllm_serve/linear_attention_state_int8.yaml new file mode 100644 index 00000000000..e0784ac64ed --- /dev/null +++ b/examples/vllm_serve/linear_attention_state_int8.yaml @@ -0,0 +1,15 @@ +imports: + base_disable_all: configs/ptq/units/base_disable_all + state_int8: configs/ptq/units/linear_attention_state_int8_dynamic + +metadata: + recipe_type: ptq + description: GDN/KDA INT8 state-tile QDQ before native vLLM prefill and decode. +quantize: + quant_cfg: + - $import: base_disable_all + - $import: state_int8 + linear_attention: + - module_name: '*' + cfg: {backend: serving, precision: vllm} + algorithm: diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index 1bfc7f95fd4..7973d4e71bf 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -241,6 +241,19 @@ def _merge_values_require_identical(merged_key: str, key_value_pairs: list[tuple return first_value +def _map_linear_attention_names(values, map_fun): + """Map module paths using weight-path probes without dropping numerical policies.""" + if map_fun is None: + return values + probes = {(name + ".weight" if name else "weight"): value for name, value in values.items()} + mapped = map_fun(probes) + if len(mapped) != len(values) or any( + name != "weight" and not name.endswith(".weight") for name in mapped + ): + raise ValueError("HF-to-vLLM mapping must preserve every linear-attention module") + return {name.removesuffix("weight").removesuffix("."): value for name, value in mapped.items()} + + def convert_dict_to_vllm( state_dict: dict[str, Any], max_or_concat: bool = True, @@ -254,6 +267,15 @@ def convert_dict_to_vllm( max_or_concat: Whether to merge grouped values by taking max/concatenate or require identical map_fun: Function to map the state dict to vLLM format """ + state_dict = dict(state_dict) + linear_quantizers = {} + for key in list(state_dict): + match = re.match(r"^(.*)\.((?:gdn|kda)_(?:state|w)_quantizer(?:\..*)?)$", key) + if match: + module, suffix = match.groups() + mapped = _map_linear_attention_names({module: state_dict.pop(key)}, map_fun) + name, value = next(iter(mapped.items())) + linear_quantizers[f"{name}.{suffix}"] = value # If map_fun is provided, pre-transform quantizer key module-path prefixes so that # HF→vLLM model renames (e.g. backbone.layers → model.layers) are applied before # key grouping (q/k/v → qkv, experts.N.up_proj → experts.w13, etc.). @@ -287,7 +309,7 @@ def convert_dict_to_vllm( _, value = key_value_pairs[0] vllm_state_dict[merged_key] = value if map_fun is None: - return vllm_state_dict + return {**vllm_state_dict, **linear_quantizers} # Quantizer module-path keys (e.g. "layers.0.mlp.gate_proj.input_quantizer") must NOT # go through map_fun (hf_to_vllm_mapper.apply_dict), which maps weight tensor paths and # drops any key it doesn't recognise — including all quantizer keys. Split them out, @@ -295,7 +317,7 @@ def convert_dict_to_vllm( quantizer_keys = {k: v for k, v in vllm_state_dict.items() if "_quantizer" in k} non_quantizer_keys = {k: v for k, v in vllm_state_dict.items() if "_quantizer" not in k} mapped = map_fun(non_quantizer_keys) if non_quantizer_keys else {} - return {**mapped, **quantizer_keys} + return {**mapped, **quantizer_keys, **linear_quantizers} def convert_modelopt_state_to_vllm( @@ -320,6 +342,16 @@ def convert_modelopt_state_to_vllm( modelopt_state_dict = modelopt_state.pop("modelopt_state_dict", []) for idx, current_mode in enumerate(modelopt_state_dict): current_mode_metadata = current_mode[1].pop("metadata", {}) + if "linear_attention" in current_mode_metadata: + current_mode_metadata["linear_attention"] = _map_linear_attention_names( + current_mode_metadata["linear_attention"], map_fun + ) + # Use exact saved module names instead of training-framework selector patterns. + if "quant_cfg" in current_mode[1]["config"]: + current_mode[1]["config"]["linear_attention"] = [ + {"module_name": name, "cfg": policy} + for name, policy in current_mode_metadata["linear_attention"].items() + ] current_mode_quant_state = current_mode_metadata.pop("quantizer_state", {}) if current_mode_quant_state: current_mode_metadata["quantizer_state"] = convert_dict_to_vllm( @@ -367,6 +399,15 @@ def filter_modelopt_state_quantizer_state_for_model( metadata = mode_entry[1].get("metadata", {}) if "quantizer_state" in metadata: saved = metadata["quantizer_state"] + missing_linear = [ + name + for name, state in saved.items() + if re.search(r"(?:gdn|kda)_(?:state|w)_quantizer$", name) + and not state.get("_disabled", False) + and name not in model_keys + ] + if missing_linear: + raise ValueError(f"Saved linear-attention quantizers are missing: {missing_linear}") # Keep keys that exist in the model. Remove disabled quantizers UNLESS they # have registered buffers (e.g. _pre_quant_scale from AWQ/smoothquant on a diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index a698514d53a..5e3b677ea2f 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -999,9 +999,12 @@ def set_quant_params(axis, block_reshape_size, padding, slices, amax_shape=None) if amax_shape: self._amax_shape_for_export = amax_shape - # Reshape size have already been set - if hasattr(self, "_block_reshape_size"): + # Static scales require fixed groups; dynamic scales allow changing input shapes. + if hasattr(self, "_block_reshape_size") and not self._dynamic: return + for attribute in ("_padding", "_slices"): + if hasattr(self, attribute): + delattr(self, attribute) reshape_size, quantize_axis, paddings, slices = [], [], [], [] diff --git a/modelopt/torch/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py index fc56e230838..247dacdd909 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -72,6 +72,7 @@ with import_plugin("vllm"): from .vllm import * from .vllm_indexer import * + from .vllm_linear_attention import * with import_plugin("trl"): from .trl import * diff --git a/modelopt/torch/quantization/plugins/vllm_linear_attention.py b/modelopt/torch/quantization/plugins/vllm_linear_attention.py new file mode 100644 index 00000000000..49131c9a481 --- /dev/null +++ b/modelopt/torch/quantization/plugins/vllm_linear_attention.py @@ -0,0 +1,176 @@ +# 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. + +"""Recurrent-state QDQ at native vLLM prefill and decode call boundaries.""" + +from functools import partial +from types import FunctionType + +import torch +import vllm +from packaging.version import Version +from vllm.forward_context import get_forward_context + +from ..linear_attention.utils import _state_qdq, state_quantizer_config +from ..nn import QuantModuleRegistry +from .custom import CUSTOM_MODEL_PLUGINS +from .linear_attention import _LinearAttentionQuantMixin + +__all__ = ["bind_vllm_linear_attention"] + + +def bind_vllm_linear_attention(model, model_runner): + """Validate state-only adapters before calibration or generation.""" + layers = [m for m in model.modules() if isinstance(m, _QuantVllmLinearAttention)] + for layer in layers: + layer.validate_linear_attention() + if any(layer.linear_attention_is_enabled for layer in layers): + _validate_runtime(model_runner) + return layers + + +def _validate_runtime(runner): + version = Version(vllm.__version__) + if version.release[:2] != (0, 15): + raise NotImplementedError("Linear-attention fakequant currently requires vLLM 0.15.x V1") + from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator + + shape = MambaStateShapeCalculator.gated_delta_net_state_shape(1, 2, 2, 32, 64, 4)[-1] + if shape[-2:] != (32, 64): + raise NotImplementedError( + "Linear-attention fakequant requires the key-first vLLM state ABI" + ) + config = runner.vllm_config + if not config.model_config.enforce_eager or config.scheduler_config.async_scheduling: + raise ValueError("State fakequant requires --enforce-eager --no-async-scheduling") + if config.speculative_config is not None or config.cache_config.enable_prefix_caching: + raise ValueError("Disable speculative decoding and prefix caching for state fakequant") + if config.cache_config.mamba_cache_mode != "none": + raise ValueError("State fakequant requires mamba_cache_mode='none'") + if config.kv_transfer_config is not None or config.ec_transfer_config is not None: + raise ValueError("State fakequant does not support state transfer") + parallel = config.parallel_config + if any( + getattr(parallel, name, 1) != 1 + for name in ( + "pipeline_parallel_size", + "data_parallel_size", + "prefill_context_parallel_size", + "decode_context_parallel_size", + ) + ): + raise ValueError("State fakequant supports TP with PP=1, DP=1, and CP=1") + + +class _QuantVllmLinearAttention(_LinearAttentionQuantMixin): + linear_attention_quantizer_names = ("gdn_state_quantizer", "gdn_w_quantizer") + + def validate_linear_attention(self): + # vLLM owns phase boundaries; the training-only validation needs a context. + if self._linear_attn_state.is_enabled: + state_quantizer_config(self._linear_attn_state) + if ( + any( + getattr(self, name).is_enabled for name in self.linear_attention_quantizer_names[1:] + ) + or self.linear_attention_config.precision == "replayssm" + ): + raise ValueError( + "vLLM state fakequant supports plain state QDQ; W and ReplaySSM are unsupported" + ) + + def _state_qdq(self, state): + return _state_qdq( + state, + block_v=self.linear_attention_config.state_block_v, + state_format="int8" if self._linear_attn_state.num_bits == 8 else "fp8_e4m3", + state_quantizer=self._linear_attn_state, + ) + + def _quantized_state_call(self, native, *args, **kwargs): + state = kwargs.get("initial_state") + if state is not None and self._linear_attn_state.is_enabled: + if state.dtype != torch.float32: + raise ValueError("State fakequant requires an FP32 recurrent cache") + indices = kwargs.get("ssm_state_indices") + if indices is None: + # Native prefill has already initialized fresh requests to zero. + kwargs["initial_state"] = self._state_qdq(state) + else: + # Eager, non-speculative decode has one cache slot per active sequence. + indices = indices[: kwargs["cu_seqlens"].numel() - 1].long() + state.index_copy_(0, indices, self._state_qdq(state.index_select(0, indices))) + return native(*args, **kwargs) + + def _forward_with_state_qdq(self, original, kernel_names, *args, **kwargs): + if not self.linear_attention_is_enabled: + return original(*args, **kwargs) + context = get_forward_context() + if context.attn_metadata is None: + return original(*args, **kwargs) + if context.attn_metadata[self.prefix].spec_sequence_masks is not None: + raise ValueError("Speculative metadata is unsupported for state fakequant") + function = original.__func__ + replacements = { + name: partial(self._quantized_state_call, function.__globals__[name]) + for name in kernel_names + } + # Bind wrappers for this invocation only; every wrapper calls the original kernel. + forward = FunctionType( + function.__code__, + {**function.__globals__, **replacements}, + function.__name__, + function.__defaults__, + function.__closure__, + ) + forward.__kwdefaults__ = function.__kwdefaults__ + return forward(self, *args, **kwargs) + + +class _QuantVllmGDN(_QuantVllmLinearAttention): + def _forward_core(self, *args, **kwargs): + return self._forward_with_state_qdq( + super()._forward_core, + ("chunk_gated_delta_rule", "fused_recurrent_gated_delta_rule"), + *args, + **kwargs, + ) + + +class _QuantVllmKDA(_QuantVllmLinearAttention): + linear_attention_quantizer_names = ("kda_state_quantizer",) + + def _forward(self, *args, **kwargs): + return self._forward_with_state_qdq( + super()._forward, + ("chunk_kda", "fused_recurrent_kda"), + *args, + **kwargs, + ) + + +def _register_vllm_linear_attention(model): + adapters = { + ("vllm.model_executor.models.qwen3_next", "Qwen3NextGatedDeltaNet"): _QuantVllmGDN, + ("vllm.model_executor.layers.kda", "KimiDeltaAttention"): _QuantVllmKDA, + } + for module in model.modules(): + cls = type(module) + adapter = adapters.get((cls.__module__, cls.__name__)) + if adapter is not None and cls not in QuantModuleRegistry: + QuantModuleRegistry.register({cls: f"vllm_{cls.__name__}"})(adapter) + + +CUSTOM_MODEL_PLUGINS.add(_register_vllm_linear_attention) diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py b/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py new file mode 100644 index 00000000000..b3f3193522e --- /dev/null +++ b/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py @@ -0,0 +1,315 @@ +# 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. + +"""Native vLLM state-boundary QDQ and worker reload on tiny offline models.""" + +import copy +import gc +import importlib.util +import math +import os +import shutil +from pathlib import Path + +import pytest +import torch +from packaging.version import Version +from transformers import Qwen3NextConfig +from vllm import LLM, SamplingParams +from vllm import __version__ as vllm_version +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig + +import modelopt.torch.opt as mto +from modelopt.torch.quantization.linear_attention import LinearAttentionConfig +from modelopt.torch.quantization.plugins.vllm_linear_attention import _QuantVllmLinearAttention + +pytestmark = pytest.mark.skipif( + Version(vllm_version).release[:2] != (0, 15), reason="Requires the pinned vLLM 0.15.x ABI" +) +ROOT = Path(__file__).parents[4] + + +def _example_module(name): + spec = importlib.util.spec_from_file_location( + name + "_test", ROOT / "examples/vllm_serve" / (name + ".py") + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _set_state_format_worker(worker, *, format="int8"): + for layer in worker.model_runner.model.modules(): + if not isinstance(layer, _QuantVllmLinearAttention): + continue + quantizer = layer._linear_attn_state + quantizer.num_bits = (4, 3) if format == "fp8" else 8 + quantizer.block_sizes = {-1: 32} if format == "block32" else None + if format != "block32": + quantizer.axis = (0, 1) + # Smaller than both tiny models' value dimensions, exposing tile grouping. + layer.linear_attention_config.state_block_v = 16 + layer.validate_linear_attention() + + +def _disable_worker(worker): + for layer in worker.model_runner.model.modules(): + if not isinstance(layer, _QuantVllmLinearAttention): + continue + layer._linear_attn_state.disable() + layer.linear_attention_config = LinearAttentionConfig() + + +def _audit_worker(worker, *, checkpoint=None): + model = worker.model_runner.model + if hasattr(model, "unwrap"): + model = model.unwrap() + layers = [m for m in model.modules() if isinstance(m, _QuantVllmLinearAttention)] + assert layers + if checkpoint is not None and torch.distributed.get_rank() == 0: + torch.save(mto.modelopt_state(model), checkpoint) + reports = [] + for layer in layers: + assert not hasattr(layer, "_linear_attention_cache") + if not hasattr(layer, "_state_audit"): + layer._state_audit = {"prefill": 0, "decode": 0, "changed": 0, "heads": 0} + original = layer._quantized_state_call + + def checked_call(native, *args, _layer=layer, _original=original, **kwargs): + state = kwargs["initial_state"] + indices = kwargs.get("ssm_state_indices") + if indices is not None: + indices = indices[: kwargs["cu_seqlens"].numel() - 1].long() + expected_state = state.clone() + selected = state if indices is None else state.index_select(0, indices) + reference_quantizer = copy.deepcopy(_layer._linear_attn_state) + if reference_quantizer.block_sizes is not None: + rounded = reference_quantizer(selected.clone()) + else: + rounded = torch.cat( + [ + reference_quantizer(part.clone()) + for part in selected.split( + _layer.linear_attention_config.state_block_v, dim=-1 + ) + ], + dim=-1, + ) + if indices is None: + expected_state = rounded + else: + expected_state.index_copy_(0, indices.long(), rounded) + # Native KDA prefill writes its output into the value-input buffer. + control_args = tuple(x.clone() if isinstance(x, torch.Tensor) else x for x in args) + control_kwargs = { + name: x.clone() if isinstance(x, torch.Tensor) else x + for name, x in kwargs.items() + } + control_kwargs["initial_state"] = expected_state.clone() + expected = native(*control_args, **control_kwargs) + native_calls = [] + + def checked_native(*native_args, **native_kwargs): + incoming = native_kwargs["initial_state"] + # Full-cache equality also verifies untouched inactive slots. + torch.testing.assert_close(incoming, expected_state, atol=0, rtol=0) + native_calls.append(True) + return native(*native_args, **native_kwargs) + + actual = _original(checked_native, *args, **kwargs) + if indices is not None: + torch.testing.assert_close( + state, control_kwargs["initial_state"], atol=0, rtol=0 + ) + assert native_calls == [True] + for result, control in zip(actual, expected): + if control is not None: + torch.testing.assert_close(result, control, atol=0, rtol=0) + _layer._state_audit["prefill" if indices is None else "decode"] += 1 + _layer._state_audit["changed"] += int(not torch.equal(selected, rounded)) + _layer._state_audit["heads"] = state.shape[1] + return actual + + layer._quantized_state_call = checked_call + reports.append(dict(layer._state_audit)) + return reports + + +def _tiny_model(path, kind): + common = { + "torch_dtype": "bfloat16", + "hidden_size": 128, + "intermediate_size": 256, + "num_attention_heads": 4, + "head_dim": 32, + "vocab_size": 512, + "max_position_embeddings": 128, + } + if kind == "gdn": + config = Qwen3NextConfig( + **common, + num_hidden_layers=2, + num_key_value_heads=2, + linear_num_value_heads=4, + linear_num_key_heads=2, + linear_key_head_dim=32, + linear_value_head_dim=64, + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=64, + shared_expert_intermediate_size=64, + mlp_only_layers=[], + full_attention_interval=2, + architectures=["Qwen3NextForCausalLM"], + ) + else: + config = KimiLinearConfig( + **common, + num_hidden_layers=1, + architectures=["KimiLinearForCausalLM"], + linear_attn_config={ + "head_dim": 32, + "num_heads": 4, + "short_conv_kernel_size": 4, + "kda_layers": [1], + "full_attn_layers": [], + }, + ) + config.save_pretrained(path) + shutil.copytree(ROOT / "tests/_test_utils/torch/tokenizer", path, dirs_exist_ok=True) + + +def _generate(llm): + outputs = llm.generate( + [{"prompt_token_ids": [3] * 73}, {"prompt_token_ids": [4]}], + SamplingParams(temperature=0, max_tokens=3, ignore_eos=True, logprobs=1), + use_tqdm=False, + ) + values = [(o.outputs[0].token_ids, o.outputs[0].cumulative_logprob) for o in outputs] + assert all(len(tokens) == 3 and math.isfinite(score) for tokens, score in values) + return values + + +@pytest.mark.parametrize("kind", ["gdn", "kda"]) +@pytest.mark.timeout(360) +def test_fakequant_worker_generation_reference_and_reload(tmp_path, monkeypatch, kind): + if not torch.cuda.is_available(): + pytest.skip("Requires a CUDA GPU") + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1") + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + paths = [str(ROOT), str(ROOT / "examples/vllm_serve"), str(Path(__file__).parent)] + # Spawn restores the parent interpreter search path after reading PYTHONPATH. + for path in reversed(paths): + monkeypatch.syspath_prepend(path) + monkeypatch.setenv("PYTHONPATH", os.pathsep.join([*paths, os.environ.get("PYTHONPATH", "")])) + monkeypatch.setenv( + "RECIPE_PATH", str(ROOT / "examples/vllm_serve/linear_attention_state_int8.yaml") + ) + monkeypatch.delenv("MODELOPT_STATE_PATH", raising=False) + model_path, checkpoint = tmp_path / "model", tmp_path / "state.pt" + _tiny_model(model_path, kind) + kwargs = { + "model": str(model_path), + "load_format": "dummy", + "dtype": "bfloat16", + "max_model_len": 128, + "max_num_seqs": 4, + "max_num_batched_tokens": 64, + "enforce_eager": True, + "enable_prefix_caching": False, + "enable_chunked_prefill": True, + "async_scheduling": False, + "mamba_cache_dtype": "float32", + "kv_cache_memory_bytes": 128 * 1024**2, + "worker_cls": "fakequant_worker.FakeQuantWorker", + "disable_custom_all_reduce": True, + "tensor_parallel_size": 1, + "gpu_memory_utilization": 0.1, + "seed": 17, + } + llm = LLM(**kwargs) + try: + llm.collective_rpc(_set_state_format_worker) + llm.collective_rpc(_audit_worker, kwargs={"checkpoint": str(checkpoint)}) + actual = _generate(llm) + assert _generate(llm) == actual # Fresh requests can reuse the same native state slots. + reports = llm.collective_rpc(_audit_worker) + assert len(reports) == 1 + for rank in reports: + assert all( + r["prefill"] > 0 and r["decode"] > 0 and r["changed"] > 0 and r["heads"] == 4 + for r in rank + ) + for format in ("fp8", "block32"): + llm.collective_rpc(_set_state_format_worker, kwargs={"format": format}) + _generate(llm) + llm.collective_rpc(_disable_worker) + disabled = _generate(llm) + finally: + llm.llm_engine.engine_core.shutdown() + del llm + gc.collect() + + monkeypatch.delenv("RECIPE_PATH") + monkeypatch.setenv("MODELOPT_STATE_PATH", str(checkpoint)) + llm = LLM(**kwargs) + try: + assert _generate(llm) == actual + finally: + llm.llm_engine.engine_core.shutdown() + del llm + gc.collect() + + monkeypatch.delenv("MODELOPT_STATE_PATH") + llm = LLM(**kwargs) + try: + assert _generate(llm) == disabled + finally: + llm.llm_engine.engine_core.shutdown() + del llm + gc.collect() + + +def test_saved_linear_attention_policy_and_quantizer_names_follow_mapper(): + module = _example_module("vllm_reload_utils") + + metadata = { + "linear_attention": {"backbone.layers.0.attn": {"backend": "serving", "state_block_v": 16}}, + "quantizer_state": {"backbone.layers.0.attn.kda_state_quantizer": {"_disabled": False}}, + } + state = { + "modelopt_state_dict": [("quantize", {"config": {"quant_cfg": []}, "metadata": metadata})] + } + + def mapper(values): + return { + k.replace("backbone.", "model.").replace(".attn.", ".self_attn."): v + for k, v in values.items() + } + + state["modelopt_state_dict"].append( + ("quantize_algo", {"config": {}, "metadata": copy.deepcopy(metadata)}) + ) + result = module.convert_modelopt_state_to_vllm(state, mapper) + converted = result["modelopt_state_dict"][0][1] + assert list(converted["metadata"]["linear_attention"]) == ["model.layers.0.self_attn"] + assert list(converted["metadata"]["quantizer_state"]) == [ + "model.layers.0.self_attn.kda_state_quantizer" + ] + assert converted["config"]["linear_attention"][0] == { + "module_name": "model.layers.0.self_attn", + "cfg": {"backend": "serving", "state_block_v": 16}, + } + assert "linear_attention" not in result["modelopt_state_dict"][1][1]["config"] diff --git a/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py b/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py index d1512547643..bf8238cb85c 100644 --- a/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py @@ -119,3 +119,16 @@ def test_grouped_quantizer_preserves_nested_sequential_state_dict_layout(): ) assert list(grouped.state_dict()) == ["0.0._amax", "0.1._amax", "1._amax"] + + +def test_dynamic_block_quantization_accepts_changing_shapes(): + cfg = QuantizerAttributeConfig( + num_bits=8, type="dynamic", block_sizes={-1: 32}, pass_through_bwd=True + ) + quantizer = TensorQuantizer(cfg) + for shape in ((2, 3, 17), (1, 2, 4, 64), (3, 2, 33)): + value = torch.randn(shape, requires_grad=True) + output = quantizer(value) + torch.testing.assert_close(output, TensorQuantizer(cfg)(value), rtol=0, atol=0) + output.sum().backward() + torch.testing.assert_close(value.grad, torch.ones_like(value), rtol=0, atol=0) From 022c6fbf8fd8534faa5fe7152766311216ab4f5b Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Thu, 8 Oct 2026 13:38:24 -0700 Subject: [PATCH 2/4] Trim vLLM state quantization tests Signed-off-by: Kai Xu --- .../test_vllm_linear_attention.py | 207 ++++++------------ 1 file changed, 68 insertions(+), 139 deletions(-) diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py b/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py index b3f3193522e..6f4a58e0e0e 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Native vLLM state-boundary QDQ and worker reload on tiny offline models.""" +"""Native vLLM state-boundary QDQ and policy restoration on tiny offline models.""" import copy import gc @@ -32,12 +32,11 @@ from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig import modelopt.torch.opt as mto +from modelopt.torch.opt.conversion import ModeloptStateManager +from modelopt.torch.quantization.conversion import restore_quantizer_state from modelopt.torch.quantization.linear_attention import LinearAttentionConfig from modelopt.torch.quantization.plugins.vllm_linear_attention import _QuantVllmLinearAttention -pytestmark = pytest.mark.skipif( - Version(vllm_version).release[:2] != (0, 15), reason="Requires the pinned vLLM 0.15.x ABI" -) ROOT = Path(__file__).parents[4] @@ -50,41 +49,25 @@ def _example_module(name): return module -def _set_state_format_worker(worker, *, format="int8"): - for layer in worker.model_runner.model.modules(): - if not isinstance(layer, _QuantVllmLinearAttention): - continue - quantizer = layer._linear_attn_state - quantizer.num_bits = (4, 3) if format == "fp8" else 8 - quantizer.block_sizes = {-1: 32} if format == "block32" else None - if format != "block32": - quantizer.axis = (0, 1) - # Smaller than both tiny models' value dimensions, exposing tile grouping. - layer.linear_attention_config.state_block_v = 16 - layer.validate_linear_attention() - - -def _disable_worker(worker): - for layer in worker.model_runner.model.modules(): - if not isinstance(layer, _QuantVllmLinearAttention): - continue - layer._linear_attn_state.disable() - layer.linear_attention_config = LinearAttentionConfig() - - -def _audit_worker(worker, *, checkpoint=None): +def _state_worker(worker, *, action="audit"): model = worker.model_runner.model if hasattr(model, "unwrap"): model = model.unwrap() layers = [m for m in model.modules() if isinstance(m, _QuantVllmLinearAttention)] assert layers - if checkpoint is not None and torch.distributed.get_rank() == 0: - torch.save(mto.modelopt_state(model), checkpoint) - reports = [] + if action == "restore": + state = worker._state_checkpoint + manager = ModeloptStateManager() + manager.load_state_dict(state["modelopt_state_dict"], state["modelopt_version"]) + for _, config, metadata in manager.modes_with_states(): + restore_quantizer_state(model, config, metadata) for layer in layers: - assert not hasattr(layer, "_linear_attention_cache") - if not hasattr(layer, "_state_audit"): - layer._state_audit = {"prefill": 0, "decode": 0, "changed": 0, "heads": 0} + if action == "disable": + layer._linear_attn_state.disable() + layer.linear_attention_config = LinearAttentionConfig() + elif action == "setup": + layer.linear_attention_config.state_block_v = 16 + layer._state_calls = {"prefill": 0, "decode": 0, "changed": 0} original = layer._quantized_state_call def checked_call(native, *args, _layer=layer, _original=original, **kwargs): @@ -92,59 +75,28 @@ def checked_call(native, *args, _layer=layer, _original=original, **kwargs): indices = kwargs.get("ssm_state_indices") if indices is not None: indices = indices[: kwargs["cu_seqlens"].numel() - 1].long() - expected_state = state.clone() selected = state if indices is None else state.index_select(0, indices) - reference_quantizer = copy.deepcopy(_layer._linear_attn_state) - if reference_quantizer.block_sizes is not None: - rounded = reference_quantizer(selected.clone()) - else: - rounded = torch.cat( - [ - reference_quantizer(part.clone()) - for part in selected.split( - _layer.linear_attention_config.state_block_v, dim=-1 - ) - ], - dim=-1, - ) + quantizer = copy.deepcopy(_layer._linear_attn_state) + rounded = torch.cat([quantizer(x) for x in selected.split(16, -1)], -1) + expected = state.clone() if indices is None: - expected_state = rounded + expected = rounded else: - expected_state.index_copy_(0, indices.long(), rounded) - # Native KDA prefill writes its output into the value-input buffer. - control_args = tuple(x.clone() if isinstance(x, torch.Tensor) else x for x in args) - control_kwargs = { - name: x.clone() if isinstance(x, torch.Tensor) else x - for name, x in kwargs.items() - } - control_kwargs["initial_state"] = expected_state.clone() - expected = native(*control_args, **control_kwargs) - native_calls = [] + expected.index_copy_(0, indices, rounded) - def checked_native(*native_args, **native_kwargs): - incoming = native_kwargs["initial_state"] - # Full-cache equality also verifies untouched inactive slots. - torch.testing.assert_close(incoming, expected_state, atol=0, rtol=0) - native_calls.append(True) - return native(*native_args, **native_kwargs) + def checked_native(*a, **kw): + # Check the full cache, including inactive slots, before delegating. + torch.testing.assert_close(kw["initial_state"], expected, atol=0, rtol=0) + return native(*a, **kw) - actual = _original(checked_native, *args, **kwargs) - if indices is not None: - torch.testing.assert_close( - state, control_kwargs["initial_state"], atol=0, rtol=0 - ) - assert native_calls == [True] - for result, control in zip(actual, expected): - if control is not None: - torch.testing.assert_close(result, control, atol=0, rtol=0) - _layer._state_audit["prefill" if indices is None else "decode"] += 1 - _layer._state_audit["changed"] += int(not torch.equal(selected, rounded)) - _layer._state_audit["heads"] = state.shape[1] - return actual + _layer._state_calls["prefill" if indices is None else "decode"] += 1 + _layer._state_calls["changed"] += int(not torch.equal(selected, rounded)) + return _original(checked_native, *args, **kwargs) layer._quantized_state_call = checked_call - reports.append(dict(layer._state_audit)) - return reports + if action == "setup": + worker._state_checkpoint = copy.deepcopy(mto.modelopt_state(model)) + return [dict(layer._state_calls) for layer in layers] def _tiny_model(path, kind): @@ -202,15 +154,15 @@ def _generate(llm): return values -@pytest.mark.parametrize("kind", ["gdn", "kda"]) -@pytest.mark.timeout(360) -def test_fakequant_worker_generation_reference_and_reload(tmp_path, monkeypatch, kind): +@pytest.fixture(params=["gdn", "kda"]) +def compiled_worker(request, tmp_path, monkeypatch): + """Compile one tiny native model outside the functional test's timeout.""" + if Version(vllm_version).release[:2] != (0, 15): + pytest.skip("The state adapter requires vLLM 0.15.x") if not torch.cuda.is_available(): pytest.skip("Requires a CUDA GPU") - monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1") monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") paths = [str(ROOT), str(ROOT / "examples/vllm_serve"), str(Path(__file__).parent)] - # Spawn restores the parent interpreter search path after reading PYTHONPATH. for path in reversed(paths): monkeypatch.syspath_prepend(path) monkeypatch.setenv("PYTHONPATH", os.pathsep.join([*paths, os.environ.get("PYTHONPATH", "")])) @@ -218,68 +170,45 @@ def test_fakequant_worker_generation_reference_and_reload(tmp_path, monkeypatch, "RECIPE_PATH", str(ROOT / "examples/vllm_serve/linear_attention_state_int8.yaml") ) monkeypatch.delenv("MODELOPT_STATE_PATH", raising=False) - model_path, checkpoint = tmp_path / "model", tmp_path / "state.pt" - _tiny_model(model_path, kind) - kwargs = { - "model": str(model_path), - "load_format": "dummy", - "dtype": "bfloat16", - "max_model_len": 128, - "max_num_seqs": 4, - "max_num_batched_tokens": 64, - "enforce_eager": True, - "enable_prefix_caching": False, - "enable_chunked_prefill": True, - "async_scheduling": False, - "mamba_cache_dtype": "float32", - "kv_cache_memory_bytes": 128 * 1024**2, - "worker_cls": "fakequant_worker.FakeQuantWorker", - "disable_custom_all_reduce": True, - "tensor_parallel_size": 1, - "gpu_memory_utilization": 0.1, - "seed": 17, - } - llm = LLM(**kwargs) + model_path = tmp_path / "model" + _tiny_model(model_path, request.param) + llm = LLM( + model=str(model_path), + load_format="dummy", + dtype="bfloat16", + max_model_len=128, + max_num_seqs=4, + max_num_batched_tokens=64, + enforce_eager=True, + enable_prefix_caching=False, + enable_chunked_prefill=True, + async_scheduling=False, + mamba_cache_dtype="float32", + kv_cache_memory_bytes=128 * 1024**2, + gpu_memory_utilization=0.1, + worker_cls="fakequant_worker.FakeQuantWorker", + seed=17, + ) try: - llm.collective_rpc(_set_state_format_worker) - llm.collective_rpc(_audit_worker, kwargs={"checkpoint": str(checkpoint)}) - actual = _generate(llm) - assert _generate(llm) == actual # Fresh requests can reuse the same native state slots. - reports = llm.collective_rpc(_audit_worker) - assert len(reports) == 1 - for rank in reports: - assert all( - r["prefill"] > 0 and r["decode"] > 0 and r["changed"] > 0 and r["heads"] == 4 - for r in rank - ) - for format in ("fp8", "block32"): - llm.collective_rpc(_set_state_format_worker, kwargs={"format": format}) - _generate(llm) - llm.collective_rpc(_disable_worker) - disabled = _generate(llm) + llm.collective_rpc(_state_worker, kwargs={"action": "setup"}) + _generate(llm) + yield llm finally: llm.llm_engine.engine_core.shutdown() del llm gc.collect() - monkeypatch.delenv("RECIPE_PATH") - monkeypatch.setenv("MODELOPT_STATE_PATH", str(checkpoint)) - llm = LLM(**kwargs) - try: - assert _generate(llm) == actual - finally: - llm.llm_engine.engine_core.shutdown() - del llm - gc.collect() - monkeypatch.delenv("MODELOPT_STATE_PATH") - llm = LLM(**kwargs) - try: - assert _generate(llm) == disabled - finally: - llm.llm_engine.engine_core.shutdown() - del llm - gc.collect() +def test_state_qdq_and_restore(compiled_worker): + llm = compiled_worker + expected = _generate(llm) + reports = llm.collective_rpc(_state_worker) + assert all(all(count > 0 for count in layer.values()) for rank in reports for layer in rank) + llm.collective_rpc(_state_worker, kwargs={"action": "disable"}) + _generate(llm) + assert llm.collective_rpc(_state_worker) == reports + llm.collective_rpc(_state_worker, kwargs={"action": "restore"}) + assert _generate(llm) == expected def test_saved_linear_attention_policy_and_quantizer_names_follow_mapper(): From 6455795ca5071af68b8ab91feea7708d48158b86 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 13 Oct 2026 12:22:42 -0700 Subject: [PATCH 3/4] Simplify vLLM linear attention state QDQ for 0.30 Signed-off-by: Kai Xu --- CHANGELOG.rst | 3 +- examples/vllm_serve/README.md | 129 +++++++----------- .../linear_attention_state_int8.yaml | 7 +- modelopt/torch/quantization/model_quant.py | 9 +- .../plugins/vllm_linear_attention.py | 114 ++++++---------- .../test_vllm_linear_attention.py | 60 ++++---- .../torch/quantization/plugins/test_gdn.py | 1 + 7 files changed, 133 insertions(+), 190 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9d1d981e109..59cf9d40da3 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,8 +16,7 @@ Changelog - Added experimental NVFP4 recurrent-state fake quantization using TensorQuantizer block16 E4M3 scales for linear-attention QAT. - Add optional ReplaySSM state QAT for GDN/KDA using the compatible INT8/Hadamard serving fork. The replay recipe controls checkpoint refresh frequency independently of chunked prefill. - -- Add experimental FP8/INT8 recurrent-state fake quantization before native vLLM GDN/KDA prefill and decode; use the state-only recipe with the documented vLLM 0.15.x eager runtime. +- Add experimental FP8/INT8 recurrent-state fake quantization before native vLLM GDN/KDA prefill and decode; use the state-only recipe with the documented vLLM 0.30.x eager runtime. - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 51a75aef760..7d9b695fd82 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -11,7 +11,7 @@ The fakequant launcher does not support vLLM 0.9.0. Use one of the tested releas For GDN/KDA recurrent-state QDQ before native prefill and decode, see the [linear-attention example](#linear-attention-state-quantization), which -documents the initial vLLM 0.15.x runtime requirements and state-only INT8 recipe. +documents the vLLM 0.30.x runtime requirements and state-only INT8 recipe. ## Prepare environment @@ -469,92 +469,57 @@ Unsupported features are sliding window, ALiBi, softcap, sinks, FP8 KV cache, cr ## Linear-attention state quantization -The existing `FakeQuantWorker` adds a ModelOpt `TensorQuantizer` immediately before -native vLLM prefill and decode for `Qwen3NextGatedDeltaNet` and -`KimiDeltaAttention`. Native attention kernels, projections, convolution, gates, -outputs, and cache management remain in use. - -### Runtime and launch - -The initial adapter targets vLLM 0.15.x V1 with key-first -`[slot, head, Dk, Dv]` FP32 recurrent state. Runtime checks reject other layouts -and versions. Use a vLLM 0.15.x environment for this adapter; the newer -versions supported by the general fakequant example use different model wrappers. +`FakeQuantWorker` applies ModelOpt `TensorQuantizer` to recurrent state before +native vLLM 0.30.x GDN/KDA layers run. It supports Qwen GDN and Kimi-Linear/Kimi-K3 +KDA. The state-only recipe uses dynamic INT8 with one scale per key row and 32 +value channels, matching the training recipe. The cache remains FP32; kernels, +gates, projections and outputs use vLLM's native implementation. ```bash -PYTHONPATH=.:examples/vllm_serve \ -RECIPE_PATH=examples/vllm_serve/linear_attention_state_int8.yaml \ -bash examples/llm_qat/linear_attention/with_vllm_defaults.sh \ +VLLM_GDN_DECODE_KERNEL=triton VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE=1 \ python examples/vllm_serve/vllm_serve_fakequant.py /path/to/model \ - --tensor-parallel-size 2 --enforce-eager --no-async-scheduling \ - --no-enable-prefix-caching --mamba-cache-dtype float32 + --modelopt-recipe-path examples/vllm_serve/linear_attention_state_int8.yaml \ + --dtype bfloat16 --tensor-parallel-size 1 --enforce-eager --no-async-scheduling \ + --no-enable-prefix-caching --mamba-cache-mode none --mamba-cache-dtype float32 \ + --additional-config '{"gdn_prefill_backend":"triton","kda_prefill_backend":"triton","kda_decode_backend":"triton"}' ``` -The example enables signed symmetric dynamic INT8 state QDQ. Import the `configs/ptq/units/gdn_state_fp8_dynamic` quantizer unit for GDN FP8 -E4M3, or set `num_bits: [4, 3]` on the corresponding state quantizers. Each invocation -receives `[active_sequence, local_head, Dk, Dv]`; `axis: [0, 1]` retains one -scale per sequence and local head **within each value-column tile**. The shared -training/serving QDQ helper splits the last axis using `state_block_v` (default -64), then reduces over `[Dk, tile_width]`. State tensors remain FP32. - -For per-key 32-value INT8 blocks, use the existing recipe -`modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml`. -Its `block_sizes: {-1: 32}` is handled directly by TensorQuantizer and overrides -legacy tile grouping. The same recipe can be used for state QAT. - -Dynamic scales need no calibration dataset; `algorithm: null` skips dataset -loading. Weight/activation calibration can use the worker's existing recipe -and calibration loop. Static state calibration is unsupported in this adapter. - -Use `MODELOPT_STATE_PATH=/path/to/modelopt_state.pt` instead of a recipe to -restore quantizer configuration. Reload maps recurrent-state quantizer names -through the model's HF-to-vLLM mapper. Checkpoint weights must already match -the vLLM model architecture. - -### Quantization boundaries - -- Prefill: vLLM gathers initial states and zeros fresh requests. The wrapper - applies `TensorQuantizer` to this tensor and calls the original prefill kernel. -- Decode: the wrapper gathers active native cache slots, quantizes them, writes - them back, and calls the original decode kernel. Other slots are untouched. -- Native kernels compute outputs and update recurrent state normally. No - additional rounding is applied to the kernel's final-state write. -- The next call quantizes that state before reading it. A fresh zero state is - unchanged; prompt-to-decode rounding happens before the first decode call. -- Scheduler-level chunked prefill creates one QDQ boundary per invocation. - Internal kernel chunks do not create extra state-quantization boundaries. - Results can therefore depend on scheduler prompt-chunk sizes. -- Requests, cache-slot reuse, and preemption remain managed by vLLM. The plugin - allocates temporary gathered states, with no additional persistent state cache. -- TP ranks quantize their local heads independently, without scale all-reduce. - -The adapter accepts plain `precision: vllm` policies, including the legacy -`vllm_0_15` spelling. The serving engine supplies the phase boundaries, so no -training phase context is required. W quantization and the native ReplaySSM/ -Hadamard profile require separate serving integration and are rejected. - -Training and serving share state-QDQ formats and grouping. Numerical alignment -also requires the same vLLM runtime, arithmetic settings, and prefill scheduling: -a training prefix must reproduce each serving prompt-continuation boundary. -The QAT handoff rounds the prefix state before its first decode read; this -adapter rounds that same state before the first native decode call. QAT stores the -rounded next-state checkpoint after each update; serving stores the working -state and rounds it at the next read. These feed the same rounded state to -the next recurrence. Raw terminal cache buffers therefore need not match the -QAT checkpoint until the same QDQ is applied. - -The supported scope is eager synchronous execution with TP and PP=DP=CP=1. -Speculative decoding, prefix caching, state transfer, and CUDA graphs require -separate integration. This is floating-point numerical emulation; model-quality, -capacity, and performance claims require separate measurements. - -### Later prefill GEMM support - -Prefill operand QDQ will be a separate change after an optimized fused kernel -is available. It will reuse the eight numerical sites from the training prefill -implementation and preserve this native-cache wrapper. The PyTorch materialized -backend is not exposed in vLLM by this state-only adapter. State QDQ cadence and -prefill operand QDQ remain independent numerical policies. +`algorithm: null` needs no calibration dataset. For FP8 E4M3, change the state +quantizers' `num_bits` to `[4, 3]`. Both formats use the same `TensorQuantizer` +configuration and grouping as training. Weight/activation recipes can use the +worker's existing calibration loop. + +### State boundaries and training alignment + +- At layer entry, gather active initialized cache slots, transpose vLLM's + `[slot, head, V, K]` state to `[slot, head, K, V]`, apply QDQ, and write it back. + Fresh prefill and inactive slots are untouched. +- Native prefill/decode reads the rounded state and writes its result normally. + Prompt-to-decode rounding happens before the first decode call; subsequent + rounding happens before each state read. +- Scheduler-level chunked prefill adds one QDQ boundary per continuation call. + Internal kernel chunks do not add QDQ boundaries. Match these boundaries when + comparing with a training prefix. + +Training rounds the state after an update; serving rounds it before the next +read. These produce the same input to the next recurrence. Terminal cache buffers +need not match until the same QDQ is applied. Matching the full training profile +also requires Triton backends as above, BF16 projections, unbounded Kimi gates, +and GDN Triton packed decode (both environment settings above). Mixed +prefill/decode batches remain native vLLM execution and are outside the training +profile's arithmetic guarantee. A saved training policy does not select the +server's kernels. + +Use `--modelopt-state-path` with an HF fakequant export to reload its quantizers +and policy through the HF-to-vLLM name mapper. For the Megatron QAT/QAD example's +HF weight export, pass this state recipe explicitly; its Megatron metadata +sidecar is for Megatron reload, not the HF mapper. + +This adapter requires eager synchronous execution, an FP32 value-first recurrent +cache, and PP=DP=CP=1 (TP heads are quantized independently). Speculative decoding, +prefix caching, state transfer, W quantization and ReplaySSM are unsupported. +Prefill GEMM QDQ remains a separate change requiring optimized kernels. This is +floating-point emulation, with no cache compression or serving speedup claim. ## Known Problems diff --git a/examples/vllm_serve/linear_attention_state_int8.yaml b/examples/vllm_serve/linear_attention_state_int8.yaml index e0784ac64ed..401d6dcd4cb 100644 --- a/examples/vllm_serve/linear_attention_state_int8.yaml +++ b/examples/vllm_serve/linear_attention_state_int8.yaml @@ -1,15 +1,12 @@ imports: base_disable_all: configs/ptq/units/base_disable_all - state_int8: configs/ptq/units/linear_attention_state_int8_dynamic + state_int8: configs/ptq/units/linear_attention_state_int8_block32_dynamic metadata: recipe_type: ptq - description: GDN/KDA INT8 state-tile QDQ before native vLLM prefill and decode. + description: GDN/KDA INT8 block32 state QDQ before native vLLM prefill and decode. quantize: quant_cfg: - $import: base_disable_all - $import: state_int8 - linear_attention: - - module_name: '*' - cfg: {backend: serving, precision: vllm} algorithm: diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index f928acc2ad6..d2c3adc6e44 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -235,12 +235,17 @@ def _check_weight_quantization_took_effect(model: nn.Module, config: QuantizeCon # what those modules are, recipe units listing the supported models). _PLUGIN_QUANTIZERS = { "gdn_state_quantizer": ( - ("*GatedDeltaNet", "*GatedDeltaNetAttention"), + ("*GatedDeltaNet", "*Qwen*GatedDeltaNetAttention"), "GDN layer", "linear_attention_state_int8*", ), "kda_state_quantizer": ( - ("*KimiDeltaAttention", "*KimiK3DeltaAttention", "*KimiLinearDeltaAttention"), + ( + "*KimiDeltaAttention", + "*KimiK3DeltaAttention", + "*KimiLinearDeltaAttention", + "*KimiGatedDeltaNetAttention", + ), "KDA layer", "linear_attention_state_int8*", ), diff --git a/modelopt/torch/quantization/plugins/vllm_linear_attention.py b/modelopt/torch/quantization/plugins/vllm_linear_attention.py index 49131c9a481..257d23a879b 100644 --- a/modelopt/torch/quantization/plugins/vllm_linear_attention.py +++ b/modelopt/torch/quantization/plugins/vllm_linear_attention.py @@ -15,9 +15,6 @@ """Recurrent-state QDQ at native vLLM prefill and decode call boundaries.""" -from functools import partial -from types import FunctionType - import torch import vllm from packaging.version import Version @@ -27,6 +24,7 @@ from ..nn import QuantModuleRegistry from .custom import CUSTOM_MODEL_PLUGINS from .linear_attention import _LinearAttentionQuantMixin +from .vllm import create_parallel_state __all__ = ["bind_vllm_linear_attention"] @@ -43,14 +41,14 @@ def bind_vllm_linear_attention(model, model_runner): def _validate_runtime(runner): version = Version(vllm.__version__) - if version.release[:2] != (0, 15): - raise NotImplementedError("Linear-attention fakequant currently requires vLLM 0.15.x V1") + if version.release[:2] != (0, 30): + raise NotImplementedError("Linear-attention fakequant currently requires vLLM 0.30.x") from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator shape = MambaStateShapeCalculator.gated_delta_net_state_shape(1, 2, 2, 32, 64, 4)[-1] - if shape[-2:] != (32, 64): + if shape[-2:] != (64, 32): raise NotImplementedError( - "Linear-attention fakequant requires the key-first vLLM state ABI" + "Linear-attention fakequant requires the value-first vLLM state ABI" ) config = runner.vllm_config if not config.model_config.enforce_eager or config.scheduler_config.async_scheduling: @@ -77,8 +75,12 @@ def _validate_runtime(runner): class _QuantVllmLinearAttention(_LinearAttentionQuantMixin): linear_attention_quantizer_names = ("gdn_state_quantizer", "gdn_w_quantizer") + def _setup(self): + super()._setup() + self.parallel_state = create_parallel_state() + def validate_linear_attention(self): - # vLLM owns phase boundaries; the training-only validation needs a context. + # vLLM owns execution and phase boundaries; validate only the state-QDQ policy. if self._linear_attn_state.is_enabled: state_quantizer_config(self._linear_attn_state) if ( @@ -91,80 +93,52 @@ def validate_linear_attention(self): "vLLM state fakequant supports plain state QDQ; W and ReplaySSM are unsupported" ) - def _state_qdq(self, state): - return _state_qdq( - state, - block_v=self.linear_attention_config.state_block_v, - state_format="int8" if self._linear_attn_state.num_bits == 8 else "fp8_e4m3", - state_quantizer=self._linear_attn_state, - ) + def warn_linear_attention_profile(self): + """Native vLLM selects the kernels; training profiles do not override them.""" - def _quantized_state_call(self, native, *args, **kwargs): - state = kwargs.get("initial_state") - if state is not None and self._linear_attn_state.is_enabled: - if state.dtype != torch.float32: - raise ValueError("State fakequant requires an FP32 recurrent cache") - indices = kwargs.get("ssm_state_indices") - if indices is None: - # Native prefill has already initialized fresh requests to zero. - kwargs["initial_state"] = self._state_qdq(state) - else: - # Eager, non-speculative decode has one cache slot per active sequence. - indices = indices[: kwargs["cu_seqlens"].numel() - 1].long() - state.index_copy_(0, indices, self._state_qdq(state.index_select(0, indices))) - return native(*args, **kwargs) - - def _forward_with_state_qdq(self, original, kernel_names, *args, **kwargs): - if not self.linear_attention_is_enabled: - return original(*args, **kwargs) - context = get_forward_context() - if context.attn_metadata is None: - return original(*args, **kwargs) - if context.attn_metadata[self.prefix].spec_sequence_masks is not None: + def _quantize_cache(self, metadata): + if metadata.spec_sequence_masks is not None: raise ValueError("Speculative metadata is unsupported for state fakequant") - function = original.__func__ - replacements = { - name: partial(self._quantized_state_call, function.__globals__[name]) - for name in kernel_names - } - # Bind wrappers for this invocation only; every wrapper calls the original kernel. - forward = FunctionType( - function.__code__, - {**function.__globals__, **replacements}, - function.__name__, - function.__defaults__, - function.__closure__, + count = metadata.num_decodes + metadata.num_prefills + indices = metadata.non_spec_state_indices_tensor[:count].long() + if metadata.has_initial_state is not None: + # Fresh prefill starts from zero in vLLM; do not touch its recycled cache slot. + indices = indices[metadata.has_initial_state[:count]] + if indices.numel() == 0: + return + state = self.kv_cache[1] + if state.dtype != torch.float32: + raise ValueError("State fakequant requires an FP32 recurrent cache") + # vLLM stores [slot, head, V, K]; ModelOpt groups along the value dimension. + selected = state.index_select(0, indices).transpose(-1, -2).contiguous() + rounded = _state_qdq( + selected, self._linear_attn_state, block_v=self.linear_attention_config.state_block_v ) - forward.__kwdefaults__ = function.__kwdefaults__ - return forward(self, *args, **kwargs) - + state.index_copy_(0, indices, rounded.transpose(-1, -2)) -class _QuantVllmGDN(_QuantVllmLinearAttention): - def _forward_core(self, *args, **kwargs): - return self._forward_with_state_qdq( - super()._forward_core, - ("chunk_gated_delta_rule", "fused_recurrent_gated_delta_rule"), - *args, - **kwargs, - ) + def forward(self, *args, **kwargs): + if self._linear_attn_state.is_enabled: + metadata = get_forward_context().attn_metadata + if metadata is not None and (layer_metadata := metadata.get(self.prefix)) is not None: + self._quantize_cache(layer_metadata) + return super().forward(*args, **kwargs) class _QuantVllmKDA(_QuantVllmLinearAttention): linear_attention_quantizer_names = ("kda_state_quantizer",) - def _forward(self, *args, **kwargs): - return self._forward_with_state_qdq( - super()._forward, - ("chunk_kda", "fused_recurrent_kda"), - *args, - **kwargs, - ) - def _register_vllm_linear_attention(model): adapters = { - ("vllm.model_executor.models.qwen3_next", "Qwen3NextGatedDeltaNet"): _QuantVllmGDN, - ("vllm.model_executor.layers.kda", "KimiDeltaAttention"): _QuantVllmKDA, + ("vllm.models.kimi_k3.nvidia.kda", "KimiK3DeltaAttention"): _QuantVllmKDA, + ( + "vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn", + "QwenGatedDeltaNetAttention", + ): _QuantVllmLinearAttention, + ( + "vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn", + "KimiGatedDeltaNetAttention", + ): _QuantVllmKDA, } for module in model.modules(): cls = type(module) diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py b/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py index 6f4a58e0e0e..f01ba487102 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py @@ -66,34 +66,28 @@ def _state_worker(worker, *, action="audit"): layer._linear_attn_state.disable() layer.linear_attention_config = LinearAttentionConfig() elif action == "setup": - layer.linear_attention_config.state_block_v = 16 layer._state_calls = {"prefill": 0, "decode": 0, "changed": 0} - original = layer._quantized_state_call - - def checked_call(native, *args, _layer=layer, _original=original, **kwargs): - state = kwargs["initial_state"] - indices = kwargs.get("ssm_state_indices") - if indices is not None: - indices = indices[: kwargs["cu_seqlens"].numel() - 1].long() - selected = state if indices is None else state.index_select(0, indices) - quantizer = copy.deepcopy(_layer._linear_attn_state) - rounded = torch.cat([quantizer(x) for x in selected.split(16, -1)], -1) + original = layer._quantize_cache + + def checked_call(metadata, _layer=layer, _original=original): + state = _layer.kv_cache[1] + count = metadata.num_decodes + metadata.num_prefills + indices = metadata.non_spec_state_indices_tensor[:count].long() + if metadata.has_initial_state is not None: + indices = indices[metadata.has_initial_state[:count]] expected = state.clone() - if indices is None: - expected = rounded - else: - expected.index_copy_(0, indices, rounded) - - def checked_native(*a, **kw): - # Check the full cache, including inactive slots, before delegating. - torch.testing.assert_close(kw["initial_state"], expected, atol=0, rtol=0) - return native(*a, **kw) - - _layer._state_calls["prefill" if indices is None else "decode"] += 1 - _layer._state_calls["changed"] += int(not torch.equal(selected, rounded)) - return _original(checked_native, *args, **kwargs) - - layer._quantized_state_call = checked_call + if indices.numel(): + selected = state[indices].transpose(-1, -2).contiguous() + quantizer = copy.deepcopy(_layer._linear_attn_state) + rounded = quantizer(selected) + expected[indices] = rounded.transpose(-1, -2) + _layer._state_calls["prefill" if metadata.num_prefills else "decode"] += 1 + _layer._state_calls["changed"] += int(not torch.equal(selected, rounded)) + _original(metadata) + # Assert at native layer entry, including fresh and inactive cache slots. + torch.testing.assert_close(state, expected, atol=0, rtol=0, equal_nan=True) + + layer._quantize_cache = checked_call if action == "setup": worker._state_checkpoint = copy.deepcopy(mto.modelopt_state(model)) return [dict(layer._state_calls) for layer in layers] @@ -117,7 +111,7 @@ def _tiny_model(path, kind): linear_num_value_heads=4, linear_num_key_heads=2, linear_key_head_dim=32, - linear_value_head_dim=64, + linear_value_head_dim=32, num_experts=4, num_experts_per_tok=2, moe_intermediate_size=64, @@ -157,11 +151,13 @@ def _generate(llm): @pytest.fixture(params=["gdn", "kda"]) def compiled_worker(request, tmp_path, monkeypatch): """Compile one tiny native model outside the functional test's timeout.""" - if Version(vllm_version).release[:2] != (0, 15): - pytest.skip("The state adapter requires vLLM 0.15.x") + if Version(vllm_version).release[:2] != (0, 30): + pytest.skip("The state adapter requires vLLM 0.30.x") if not torch.cuda.is_available(): pytest.skip("Requires a CUDA GPU") monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + monkeypatch.setenv("VLLM_GDN_DECODE_KERNEL", "triton") + monkeypatch.setenv("VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE", "1") paths = [str(ROOT), str(ROOT / "examples/vllm_serve"), str(Path(__file__).parent)] for path in reversed(paths): monkeypatch.syspath_prepend(path) @@ -183,7 +179,13 @@ def compiled_worker(request, tmp_path, monkeypatch): enable_prefix_caching=False, enable_chunked_prefill=True, async_scheduling=False, + additional_config={ + "gdn_prefill_backend": "triton", + "kda_prefill_backend": "triton", + "kda_decode_backend": "triton", + }, mamba_cache_dtype="float32", + mamba_cache_mode="none", kv_cache_memory_bytes=128 * 1024**2, gpu_memory_utilization=0.1, worker_cls="fakequant_worker.FakeQuantWorker", diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index 0368f8747a9..a733630a8b4 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -306,6 +306,7 @@ def test_quant_cfg_refinement_and_policy_precedence(): ("KimiDeltaAttention", "kda_state_quantizer"), ("KimiK3DeltaAttention", "kda_state_quantizer"), ("KimiLinearDeltaAttention", "kda_state_quantizer"), + ("KimiGatedDeltaNetAttention", "kda_state_quantizer"), ], ) def test_state_recipe_rejects_unconverted_layers(class_name, quantizer_name): From 1ecee6f821ae5a2ba8651050ca73c95362e37217 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 14 Oct 2026 10:45:02 -0700 Subject: [PATCH 4/4] Document and validate vLLM state-only PTQ and reload Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 36 ++++++++++++++++--- examples/vllm_serve/vllm_reload_utils.py | 16 +++++---- .../test_vllm_linear_attention.py | 36 ++++++++++++++++--- 3 files changed, 72 insertions(+), 16 deletions(-) diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 7d9b695fd82..8375b2994cf 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -475,19 +475,45 @@ KDA. The state-only recipe uses dynamic INT8 with one scale per key row and 32 value channels, matching the training recipe. The cache remains FP32; kernels, gates, projections and outputs use vLLM's native implementation. +### PTQ on a pretrained model + +Point the existing fakequant launcher at the original pretrained checkpoint. This +state-only recipe leaves weights unchanged and requires no QAT/QAD or optimizer +steps. Dynamic scales are computed from the live recurrent state, so +`algorithm: null` needs no calibration dataset. + ```bash VLLM_GDN_DECODE_KERNEL=triton VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE=1 \ -python examples/vllm_serve/vllm_serve_fakequant.py /path/to/model \ +python examples/vllm_serve/vllm_serve_fakequant.py /path/to/pretrained-model \ --modelopt-recipe-path examples/vllm_serve/linear_attention_state_int8.yaml \ --dtype bfloat16 --tensor-parallel-size 1 --enforce-eager --no-async-scheduling \ --no-enable-prefix-caching --mamba-cache-mode none --mamba-cache-dtype float32 \ --additional-config '{"gdn_prefill_backend":"triton","kda_prefill_backend":"triton","kda_decode_backend":"triton"}' ``` -`algorithm: null` needs no calibration dataset. For FP8 E4M3, change the state -quantizers' `num_bits` to `[4, 3]`. Both formats use the same `TensorQuantizer` -configuration and grouping as training. Weight/activation recipes can use the -worker's existing calibration loop. +Restart with the same pretrained checkpoint and recipe to reproduce the state +quantization setup; dynamic state scales do not need to be saved. Request caches +are transient and are not checkpointed. To compare against unquantized state, +run the same command without `--modelopt-recipe-path`, keeping all other settings +and prompts fixed. Weight/activation PTQ can use the worker's existing +calibration loop but is separate from this state-only example. + +For FP8 E4M3, replace the state INT8 import with these quantizer entries: + +```yaml + - quantizer_name: '*gdn_state_quantizer' + cfg: &state_fp8 + num_bits: [4, 3] + axis: [0, 1] + type: dynamic + pass_through_bwd: true + - quantizer_name: '*kda_state_quantizer' + cfg: *state_fp8 +``` + +Keep the disable-all entry and `algorithm: null`. FP8 uses one scale per head's +`[K, state_block_v]` tile (64 value channels by default); it does not accept the +INT8 `block_sizes` setting. Use the same format and grouping in training. ### State boundaries and training alignment diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index 7973d4e71bf..8c902b3cdb9 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -372,7 +372,8 @@ def filter_modelopt_state_quantizer_state_for_model( Align quantizer_state in modelopt_state metadata with the model. - Removes keys not in the model (handles TP sharding - each rank has a subset). - - Removes keys only when the quantizer is disabled (in the model). + - Preserves saved linear-attention state regardless of conversion defaults. + - Removes other keys when the quantizer is disabled (in the model). - Adds keys for quantizers in the model but not in metadata (e.g. disabled/excluded). Modifies modelopt_state in place. Call after convert_to_quantized_model so the model has quantizers. @@ -399,12 +400,13 @@ def filter_modelopt_state_quantizer_state_for_model( metadata = mode_entry[1].get("metadata", {}) if "quantizer_state" in metadata: saved = metadata["quantizer_state"] + linear_keys = { + name for name in saved if re.search(r"(?:gdn|kda)_(?:state|w)_quantizer$", name) + } missing_linear = [ name - for name, state in saved.items() - if re.search(r"(?:gdn|kda)_(?:state|w)_quantizer$", name) - and not state.get("_disabled", False) - and name not in model_keys + for name in linear_keys + if not saved[name].get("_disabled", False) and name not in model_keys ] if missing_linear: raise ValueError(f"Saved linear-attention quantizers are missing: {missing_linear}") @@ -419,7 +421,9 @@ def _has_buffers(state: dict) -> bool: filtered = { k: v for k, v in saved.items() - if k in model_keys and (k not in disabled_keys or _has_buffers(v)) + # Saved state wins over defaults when training selectors no longer match. + if k in model_keys + and (k in linear_keys or k not in disabled_keys or _has_buffers(v)) } # Add state for quantizers in model but not in metadata (e.g. disabled/excluded) for k in model_keys - filtered.keys(): diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py b/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py index f01ba487102..532be24373f 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py @@ -35,7 +35,9 @@ from modelopt.torch.opt.conversion import ModeloptStateManager from modelopt.torch.quantization.conversion import restore_quantizer_state from modelopt.torch.quantization.linear_attention import LinearAttentionConfig +from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.quantization.plugins.vllm_linear_attention import _QuantVllmLinearAttention +from modelopt.torch.utils import safe_load ROOT = Path(__file__).parents[4] @@ -49,18 +51,33 @@ def _example_module(name): return module -def _state_worker(worker, *, action="audit"): +def _state_worker(worker, *, action="audit", checkpoint_path=None): model = worker.model_runner.model if hasattr(model, "unwrap"): model = model.unwrap() layers = [m for m in model.modules() if isinstance(m, _QuantVllmLinearAttention)] assert layers if action == "restore": - state = worker._state_checkpoint + reload = _example_module("vllm_reload_utils") + state = safe_load(worker._state_checkpoint, map_location="cpu") + for _, entry in state["modelopt_state_dict"]: + for field in ("quantizer_state", "linear_attention"): + if field in entry["metadata"]: + entry["metadata"][field] = { + "hf." + name: value for name, value in entry["metadata"][field].items() + } + state = reload.convert_modelopt_state_to_vllm( + state, + lambda values: {name.removeprefix("hf."): value for name, value in values.items()}, + ) + reload.filter_modelopt_state_quantizer_state_for_model(state, model) manager = ModeloptStateManager() manager.load_state_dict(state["modelopt_state_dict"], state["modelopt_version"]) for _, config, metadata in manager.modes_with_states(): restore_quantizer_state(model, config, metadata) + assert all(layer._linear_attn_state.is_enabled for layer in layers) + for name, weight in model.named_parameters(): + torch.testing.assert_close(weight.cpu(), worker._state_weights[name], atol=0, rtol=0) for layer in layers: if action == "disable": layer._linear_attn_state.disable() @@ -89,7 +106,13 @@ def checked_call(metadata, _layer=layer, _original=original): layer._quantize_cache = checked_call if action == "setup": - worker._state_checkpoint = copy.deepcopy(mto.modelopt_state(model)) + enabled = [m for m in model.modules() if isinstance(m, TensorQuantizer) and m.is_enabled] + assert {id(m) for m in enabled} == {id(layer._linear_attn_state) for layer in layers} + worker._state_weights = { + name: p.detach().cpu().clone() for name, p in model.named_parameters() + } + worker._state_checkpoint = checkpoint_path + torch.save(mto.modelopt_state(model), checkpoint_path) return [dict(layer._state_calls) for layer in layers] @@ -192,7 +215,10 @@ def compiled_worker(request, tmp_path, monkeypatch): seed=17, ) try: - llm.collective_rpc(_state_worker, kwargs={"action": "setup"}) + llm.collective_rpc( + _state_worker, + kwargs={"action": "setup", "checkpoint_path": str(tmp_path / "state.pt")}, + ) _generate(llm) yield llm finally: @@ -201,7 +227,7 @@ def compiled_worker(request, tmp_path, monkeypatch): gc.collect() -def test_state_qdq_and_restore(compiled_worker): +def test_state_ptq_and_restore(compiled_worker): llm = compiled_worker expected = _generate(llm) reports = llm.collective_rpc(_state_worker)