[2/6] Torch GDN/KDA decode QAT with INT8 recurrent state

CHANGELOG.rst

@@ -16,6 +16,7 @@ Changelog
 - 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.
+- Add experimental GDN/KDA recurrent-state QAT with native vLLM 0.30.0 Triton forwards and explicit prefill/decode phases through ``linear_attention_training_phase``, including dynamic blockwise INT8 state quantization. Megatron users need a Core build with native gate and packed-forward hooks; 0.19.2 does not support this workflow.
 - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body.
 - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before
   a separate KV-cache AutoQuantize stage, with independent resumable checkpoints for the weight and
@@ -33,7 +34,6 @@ Changelog
 - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration.
 - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/onnx_ptq/bevformer>`_ for details.
 - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``.
-- Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training through the standard ``quant_cfg`` interface. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer.
 - Add fake quantization of the sparse-attention indexer key cache and query for DeepSeek-V4 (vLLM and Megatron-Core) and GLM-5.3-Flash (vLLM) through the new ``indexer_k_quantizer`` and ``indexer_q_quantizer``. Enable them by importing the ``configs/ptq/units/indexer_k_nvfp4`` and ``configs/ptq/units/indexer_q_nvfp4`` units (NVFP4 with the global scale fixed to 1) into a recipe.
 - Add the ``configs/ptq/units/kv_nvfp4_mla`` recipe unit for fake quantization of the MLA KV cache (DeepSeek-V3, GLM-5.3-Flash, ...) in vLLM fake-quant serving: NVFP4 for the latent and FP8 for the RoPE key. See `examples/vllm_serve/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/vllm_serve>`_ for an example recipe.
 - Add fake quantization of MoE expert-parallel communication to vLLM fake-quant serving through the new ``dispatch_quantizer`` and ``combine_quantizer`` on vLLM's routed experts, for MoE layers built on vLLM's fused-MoE layer and any all-to-all backend. Enable them by importing the ``configs/ptq/units/moe_dispatch_nvfp4`` and ``configs/ptq/units/moe_combine_nvfp4`` units into a recipe; see ``examples/vllm_serve/README.md``.

examples/megatron_bridge/quantize.py

@@ -76,6 +76,7 @@
     QUANT_CFG_CHOICES,
     RecipeSupersededAction,
 )
+from modelopt.torch.quantization.linear_attention import get_linear_attention_layers
 from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0
 from modelopt.torch.utils.dataset_utils import get_supported_datasets
 from modelopt.torch.utils.mlflow import Tool, masked_args, resolved_recipe_texts
@@ -307,6 +308,13 @@ def get_quant_config(args: argparse.Namespace) -> dict:
 
 
 def main(args: argparse.Namespace):
+    mtq_config = get_quant_config(args)
+    if mtq_config.get("linear_attention") and mtq.need_calibration(mtq_config):
+        raise ValueError(
+            "This example supports linear_attention recipes only without calibration. "
+            "For calibrated projection quantization combined with state quantization, "
+            "use mtq.quantize with a calibration loop inside linear_attention_training_phase."
+        )
     trust_remote_code = is_safe_repo(
         trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_name_or_path
     )
@@ -362,8 +370,6 @@ def main(args: argparse.Namespace):
         )
     print_rank_0(f"Using calibration dataset: {args.calib_dataset_name}")
 
-    mtq_config = get_quant_config(args)
-
     # Quantize only the language model: disable quantizers on every top-level submodule that is not
     # the language model (vision tower + projector). Skip aliases of language-model submodules (e.g.
     # Qwen's ``self.decoder = language_model.decoder``) so the LM's own layers stay enabled.
@@ -474,7 +480,20 @@ def forward_loop(_model=None):
         warn_rank_0(
             "Skipping the post-quantization generation sanity check because --compress is set."
         )
-    if not args.skip_generate and not args.compress:
+    state_quantized = torch.tensor(
+        bool(get_linear_attention_layers(language_model, enabled_only=True)),
+        dtype=torch.int32,
+        device="cuda",
+    )
+    if torch.distributed.is_initialized():
+        torch.distributed.all_reduce(state_quantized, op=torch.distributed.ReduceOp.MAX)
+    state_quantized = bool(state_quantized.item())
+    if state_quantized and not args.skip_generate:
+        warn_rank_0(
+            "Skipping generation sanity check for state quantization; "
+            "use an explicit prefix/suffix evaluation to measure state QDQ."
+        )
+    if not args.skip_generate and not args.compress and not state_quantized:
         print_rank_0("\nTesting quantized model with custom prompts...")
         # Sanity-check text generation on the quantized language model.
         language_model.eval()

modelopt/torch/kernels/quantization/linear_attention/__init__.py

@@ -13,12 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Linear-attention kernels for quantization.
+"""Linear-attention quantization adapters for optional serving kernels.
 
-``fla_chunk_delta_h.py`` and ``fla_chunk_gated_delta_rule.py`` are adapted copies of the chunked
-GatedDeltaNet kernels of `flash-linear-attention <https://github.com/fla-org/flash-linear-attention>`_
-(``fla.ops.common.chunk_delta_h`` and ``fla.ops.gated_delta_rule.chunk``) that can fake-quantize the
-recurrent state carried between chunks to FP8 (``state_qdq``). They still import the surrounding
-fla operators, so ``fla-core==0.5.1`` and Triton must be installed to use
-them. This package initializer does not import the kernels, so importing it needs neither.
+Serving adapters import forward kernels only when selected; this package
+initializer has no kernel dependencies. The existing FLA chunk kernels remain
+available for direct calls with fla-core==0.5.1.
 """

modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py

@@ -0,0 +1,25 @@
+# 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.
+
+"""Optional vLLM forward kernels for linear-attention serving precision profiles."""
+
+try:
+    import vllm
+except ImportError as error:
+    raise ImportError(
+        "Linear-attention serving profiles require the optional vLLM dependency: "
+        "vllm==0.30.0 for 'vllm', or the compatible quantized-ReplaySSM "
+        "vLLM fork for 'replayssm'."
+    ) from error

modelopt/torch/kernels/quantization/linear_attention/serving/forward.py

@@ -0,0 +1,185 @@
+# 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 0.30 Triton forwards; training owns private, differentiable state."""
+
+from functools import cache, lru_cache
+
+import torch
+
+
+@cache
+def _kernels(precision):
+    # Defer optional serving/model dependencies until their explicit profile is selected.
+    if precision == "replayssm":
+        from .replay import _prefill_kernels
+
+        return _prefill_kernels()
+    import vllm
+
+    if vllm.__version__ != "0.30.0":
+        raise ImportError("Linear-attention vLLM profiles require vllm==0.30.0")
+    from vllm.third_party.flash_linear_attention import ops as gdn_ops
+
+    if precision == "vllm_kimi_k3":
+        from vllm.models.kimi_k3.nvidia.ops.third_party import kda as kda_ops
+    else:
+        from vllm.third_party.flash_linear_attention.ops import kda as kda_ops
+    return gdn_ops, kda_ops
+
+
+def _state_layout(state):
+    # ModelOpt uses [..., K, V]; both supported native backends store [..., V, K].
+    return state.transpose(-1, -2).contiguous()
+
+
+def prefill(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, precision="vllm"):
+    """Run native chunk prefill, including the selected profile's raw-gate preparation."""
+    # Match BF16 projection outputs; A_log and bias retain their parameter precision.
+    gdn, kda = _kernels(precision)
+    kimi_k3 = precision == "vllm_kimi_k3"
+    channel = g.ndim == 3
+    if not channel and gate_inputs is not None:
+        raw_g, raw_beta, rate, bias = gate_inputs
+        q, k, v, g, beta = gdn.fused_post_conv_prep(
+            torch.cat([x.flatten(1) for x in (q, k, v)], dim=-1),
+            raw_g.flatten(0, -2).bfloat16().contiguous(),
+            raw_beta.flatten(0, -2).bfloat16().contiguous(),
+            rate.reshape(-1).contiguous(),
+            bias.contiguous(),
+            q.shape[-2],
+            q.shape[-1],
+            v.shape[-1],
+            apply_l2norm=normalize,
+        )
+        normalize = False
+    q, k, v, g, beta = [x.unsqueeze(0).contiguous() for x in (q, k, v, g, beta)]
+    # Native chunk implementations may reuse V as their output buffer.
+    v = v.clone()
+    kwargs = {
+        "scale": scale,
+        "initial_state": _state_layout(state.unsqueeze(0)),
+        "output_final_state": True,
+        "use_qk_l2norm_in_kernel": normalize,
+        "cu_seqlens": _prefill_metadata(q.device, q.shape[1]),
+    }
+    if channel and gate_inputs is not None:
+        raw_g, raw_beta, rate, bias = gate_inputs
+        # Kimi fuses beta activation; standalone KDA takes already activated beta.
+        beta_arg = (
+            {"raw_beta": raw_beta[None].bfloat16().contiguous()} if kimi_k3 else {"beta": beta}
+        )
+        out, final = kda.chunk_kda_with_fused_gate(
+            q,
+            k,
+            v,
+            raw_g=raw_g[None].bfloat16().contiguous(),
+            A_log=rate.reshape(-1).contiguous(),
+            g_bias=bias.contiguous() if bias is not None else None,
+            **beta_arg,
+            **kwargs,
+        )
+    elif channel:
+        out, final = kda.chunk_kda(q, k, v, g, beta, **kwargs)
+    else:
+        out, final = gdn.chunk_gated_delta_rule(q, k, v, g, beta, **kwargs)
+    return out[0], _state_layout(final[0])
+
+
+def step(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, precision="vllm"):
+    """Run native fused decode on a private cache, preserving the incoming autograd state."""
+    gdn, kda = _kernels(precision)
+    kimi_k3 = precision == "vllm_kimi_k3"
+    channel = g.ndim == 2
+    cu, indices = _decode_metadata(q.device)
+    if channel or gate_inputs is not None:
+        # Only slot one is read/written; slot zero is reserved and needs no initialization.
+        native_state = state.new_empty(2, state.shape[0], state.shape[2], state.shape[1])
+        native_state[1].copy_(state.transpose(-1, -2))
+    else:
+        native_state = _state_layout(state.unsqueeze(0))
+    if kimi_k3:
+        # The Kimi kernel always normalizes Q/K; validate_profile enforces normalize=True.
+        raw_g, raw_beta, rate, bias = gate_inputs
+        out, final = kda.fused_recurrent_kda_packed_decode(
+            mixed_qkv=torch.cat([x.flatten() for x in (q, k, v)])[None],
+            raw_g=raw_g[None, None].bfloat16().contiguous(),
+            raw_beta=raw_beta[None, None].bfloat16().contiguous(),
+            A_log=rate.reshape(-1).contiguous(),
+            dt_bias=(torch.zeros_like(raw_g).flatten() if bias is None else bias.contiguous()),
+            lower_bound=None,
+            initial_state=native_state,
+            state_indices=indices,
+            scale=scale,
+        )
+    elif channel:
+        out, final = kda.fused_recurrent_kda_fwd(
+            *[x[None, None].contiguous() for x in (q, k, v, g, beta)],
+            initial_state=native_state,
+            inplace_final_state=True,
+            scale=scale,
+            use_qk_l2norm_in_kernel=normalize,
+            cu_seqlens=cu,
+            ssm_state_indices=indices,
+        )
+    elif gate_inputs is not None:
+        raw_g, raw_beta, rate, bias = gate_inputs
+        out, final = gdn.fused_recurrent_gated_delta_rule_packed_decode(
+            mixed_qkv=torch.cat([x.flatten() for x in (q, k, v)])[None],
+            a=raw_g[None].bfloat16().contiguous(),
+            b=raw_beta[None].bfloat16().contiguous(),
+            A_log=rate.reshape(-1).contiguous(),
+            dt_bias=bias.contiguous(),
+            scale=scale,
+            initial_state=native_state,
+            out=torch.empty_like(v, memory_format=torch.contiguous_format)[None, None],
+            ssm_state_indices=indices,
+            use_qk_l2norm_in_kernel=normalize,
+        )
+    else:
+        out, final = gdn.fused_recurrent_gated_delta_rule(
+            *[x[None, None].contiguous() for x in (q, k, v, g, beta)],
+            initial_state=native_state,
+            inplace_final_state=False,
+            scale=scale,
+            use_qk_l2norm_in_kernel=normalize,
+            cu_seqlens=cu,
+        )
+    # Indexed kernels update slot 1; the non-in-place fallback returns one state.
+    return out[0, 0], _state_layout(final[-1])
+
+
+@lru_cache(maxsize=128)
+def _prefill_metadata(device, length):
+    # Native kernels only read these boundaries; reuse them across equal-length layers.
+    return torch.arange(2, device=device, dtype=torch.int32) * length
+
+
+@cache
+def _decode_metadata(device):
+    # Immutable single-token boundaries and cache slot; construct without a CPU-to-GPU copy.
+    return torch.arange(2, device=device, dtype=torch.int32), torch.ones(
+        1, device=device, dtype=torch.int32
+    )
+
+
+def validate_profile(gate_inputs, normalize, *, precision="vllm"):
+    """Validate the selected native serving contract before launching prefill."""
+    _kernels(precision)
+    if precision == "vllm_kimi_k3":
+        if gate_inputs is None:
+            raise ValueError("vllm_kimi_k3 requires raw gate_inputs")
+        if not normalize:
+            raise ValueError("vLLM packed KDA decode requires Q/K normalization")