[5/6] GDN/KDA prefill operand and arithmetic QAT

.github/workflows/gpu_tests.yml

@@ -44,7 +44,8 @@ jobs:
       matrix:
         include:
           - example: gpu
-            timeout: 60
+            # Includes dependency builds and cold compilation of the FLA forward/backward tests.
+            timeout: 75
             # Pinned to 26.05: benchmark.py uses trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH,
             # which newer TensorRT (26.06) removed. Bump once the source is updated for TensorRT 10.
             container_image: nvcr.io/nvidia/pytorch:26.05-py3

.pre-commit-config.yaml

@@ -108,6 +108,8 @@ repos:
               modelopt/onnx/quantization/ort_patching.py|
               modelopt/torch/_deploy/utils/onnx_utils.py|
               modelopt/torch/export/transformer_engine.py|
+              modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py|
+              modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py|
               modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_pruned_to_mxfp4.py|
               modelopt/torch/quantization/export_onnx.py|
               modelopt/torch/quantization/plugins/attention.py|

CHANGELOG.rst

@@ -13,10 +13,19 @@ Changelog
 
 *Quantization*
 
+- 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 GDN/KDA decode-aware QAT with FP8 or INT8 recurrent states, optional INT8 value-axis Hadamard encoding, KDA decay rounding, and encoded-update replay. Supply explicit prefix lengths through the training phase context; prefix solves remain exact.
+
+- Add experimental KDA prefill QAT for FLA 0.5.1 layers with stable per-channel decay and independent operand/state fake quantization. Enable the explicit ``matmul`` backend; numerical policies and quantizers persist through ModelOpt save/restore.
+
+- Add experimental per-operand FP8/NVFP4 GDN prefill QAT with independently configurable accumulator and elementwise rounding. Select the materialized ``matmul`` backend to enable these sites; the existing fused FLA state/W path remains the default.
+
 - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts.
 - 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.
 
 *Megatron Framework (M-LM / M-Bridge)*
 

LICENSE

@@ -251,6 +251,7 @@ the following copyright holders, licensed under the MIT License:
   Copyright (c) 2025 sgl-project
   Copyright (c) 2026 The DeepSpec Authors
   Copyright (c) 2023-2026 The ggml authors
+  Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

docs/source/_templates/autosummary/module.rst

@@ -11,7 +11,7 @@
    :recursive:
 {% for item in modules %}
 {% set full_item = fullname + '.' + item.split('.')[-1] %}
-{% if ('.plugins.' not in full_item or full_item == 'modelopt.torch.opt.plugins.huggingface') and full_item != 'modelopt.torch.quantization.backends.fp8_per_tensor_gemm' %}
+{% if ('.plugins.' not in full_item or full_item == 'modelopt.torch.opt.plugins.huggingface') and full_item != 'modelopt.torch.quantization.backends.fp8_per_tensor_gemm' and not full_item.startswith('modelopt.torch.kernels.quantization.linear_attention.fla_') %}
    {{ full_item }}
 {% endif %}
 {%- endfor %}

examples/llm_qat/README.md

@@ -14,6 +14,7 @@ For background on how QAT enables low-precision accuracy recovery, see the [QAT/
 | Background | How QAT/QAD work and when to use each | \[[Link](#background)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] |
 | Support Matrix | Supported models, quantization formats, and backends | \[[Link](#support-matrix)\] | |
 | QLoRA | Model training with reduced GPU memory | \[[Link](#qlora-real-quantization)\] | |
+| Linear Attention | GDN/KDA recurrent-state QAT and ReplaySSM example | \[[Link](linear_attention/README.md)\] | |
 | Advanced Topics | FSDP2 config, YAML options | \[[Link](#advanced-topics)\] | |
 | Results | Accuracy benchmarks | \[[Link](#results)\] | |
 | Resources | Extra links and references | \[[Link](#resources)\] | |