[5/7] vLLM GDN/KDA state-only fake quantization

CHANGELOG.rst

@@ -14,6 +14,8 @@ Changelog
 *Quantization*
 
 - 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.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.
 - 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.

examples/llm_qat/linear_attention/README.md

@@ -292,3 +292,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.

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 vLLM 0.30.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,86 @@ 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
+
+`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.
+
+### 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/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"}'
+```
+
+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
+
+- 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
 
 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).

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)

examples/vllm_serve/linear_attention_state_int8.yaml

@@ -0,0 +1,12 @@
+imports:
+  base_disable_all: configs/ptq/units/base_disable_all
+  state_int8: configs/ptq/units/linear_attention_state_int8_block32_dynamic
+
+metadata:
+  recipe_type: ptq
+  description: GDN/KDA INT8 block32 state QDQ before native vLLM prefill and decode.
+quantize:
+  quant_cfg:
+    - $import: base_disable_all
+    - $import: state_int8
+  algorithm:

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,15 +309,15 @@ 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,
     # apply map_fun only to non-quantizer keys, then merge back.
     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(
@@ -340,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.
@@ -367,6 +400,16 @@ 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 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}")
 
             # 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
@@ -378,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():

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*",
     ),

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 = [], [], [], []
 

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 *