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

[kaix-nv]

Linear-attention PR stack — 5 PRs

Order	PR	Depends on
1/5	#2497 GDN state/W QAT foundation	main
2/5	#2519 GDN/KDA decode QAT + INT8	#2497
3/5	#2541 vLLM GDN/KDA state-only fake quantization	#2519
4/5	#2503 GDN/KDA prefill GEMM quantization	#2541
5/5	#2507 Experimental GDN/KDA approximate inverse	#2503

All five PRs form a linear GitHub stack in the review order shown above. #2541 applies TensorQuantizer before native vLLM prefill/decode calls.

A separate vLLM prefill-GEMM PR will wait for an optimized fused kernel. #2506 and #2509 are superseded and closed.

What does this PR do?

Type of change: New feature.

Add GDN/KDA decode-aware QAT without depending on configurable prefill operands
or approximate inverse. An explicit prefix length selects exact chunk prefill
followed by token-state writes, optional KDA log-retention rounding, or
encoded-update replay. Both Torch and Triton paths propagate gradients through
initial states, replay anchors, and continuation.

Recurrent states and replay anchors support FP8 E4M3 and signed symmetric dynamic
INT8. INT8 uses per-head [Dk, block_v] scales, zero point zero, [-127,127]
codes, nearest-even rounding, and identity STE. The existing GDN chunk kernel
also accepts INT8; replay key/update factors keep their separate FP8 setting.
FP8 remains the default. Quantizer format and execution policy survive save/restore.

This is the first delivery after the shared GDN foundation (#2497). Prefill
operand emulation and approximate inverse are follow-on patches, in that order.

The optional int8_hadamard32 state codec quantizes recurrent state in a Hadamard value basis. Token-state writes and replay-window refresh share one state-encoding path while preserving their separate rounding cadence and readout choices. The decode benchmark uses the same measurement helper as the later prefill and inverse examples.

This PR introduces QuantizeConfig.linear_attention, the module-selection rules, and execution-policy checkpoint metadata. #2497 supplies only standard quant_cfg state/W handles with fixed execution behavior. Decode/replay settings live here; later PRs extend the same configuration for prefill arithmetic and solve approximation. Test oracles remain under tests/_test_utils/torch/quantization/; runtime prefix and quantization helpers stay in the library.

The QAT entry point is examples/llm_qat/linear_attention/train.py. Usage and numerical contracts live with the example; historical study reports remain beside the scripts in the PRs that introduce them. The training example writes metrics without saving a trained checkpoint, and its source manifest hashes the current implementation files.

Usage

import modelopt.torch.quantization as mtq
from modelopt.torch.quantization.linear_attention import linear_attention_training_phase

model = mtq.quantize(model, {
    "quant_cfg": [
        {"quantizer_name": "*", "enable": False},
        {"quantizer_name": "*kda_state_quantizer", "cfg": {
            "num_bits": 8, "unsigned": False, "narrow_range": True,
            "type": "dynamic", "axis": (0, 1), "pass_through_bwd": True,
        }},
    ],
    "algorithm": None,
    "linear_attention": [{"module_name": "*", "cfg": {
        "backend": "matmul", "state": {"block_v": 64},
        "decode": {"mode": "replay", "implementation": "triton",
                   "replay": {"window": 8}},
    }}],
})
with linear_attention_training_phase(model, [64]):
    loss = model(input_ids=ids, labels=labels, use_cache=False).loss
    loss.backward()

Use *gdn_state_quantizer for Megatron GDN. Keep the phase context active through
activation-checkpoint backward. The prefix solve is exact in this delivery.

Testing

For this documentation/example amendment, pre-commit, Markdown links and anchors, command/Python syntax, source-manifest readability, and stale-path checks passed. Runtime kernels were not changed; model training, distributed integration, and quality studies were not rerun for this amendment.

The preceding focused review also passed 124 CPU/GPU decode, INT8, Hadamard, and study-comparison tests on RTX A6000. The separately identified low-level replay-anchor gradient issue remains open and is outside this documentation amendment.

Prior runtime validation on RTX A6000/SM86, Torch 2.9.1+cu128, Triton 3.5.1, and fla-core 0.5.1:

• 193 passed, 13 skipped across GDN conversion/checkpoint, independent references, decode/INT8/Hadamard/study CPU tests, fused GDN state/W kernels, and decode GPU tests.
• A checkpoint produced by the simplified [1/5] GDN state/W QAT foundation #2497 with state/W quantizers enabled restored successfully with default execution settings.
• Twelve skips are hardware-specific (native state FP8/Hopper); one module skip is due to the absent fla.layers.kda package.
• Pre-commit, diff checks, and commit-signature verification passed.

These tests cover outputs, final states, input/state gradients, packed and grouped-head layouts, continuation, readouts, state tile widths, and saved policy settings. Megatron-Core is absent locally, so its integration tests were not rerun. Earlier full-FLA/Hopper and model-quality measurements remain scoped to their original revisions. State-only vLLM validation belongs to #2541.

Before your PR is "Ready for review"

• Is this change backward compatible?: Yes; existing FP8 recipes/checkpoints retain their behavior.
• Copied code/new dependency guidance: Existing FLA-derived kernel provenance and licenses retained; FLA layer testing uses the already pinned 0.5.1 package.
• Necessary tests added?: Yes.
• Changelog updated?: Yes.
• Claude approval?: Pending; draft only.

Additional Information

All encoded values remain floating tensors. INT8 fake QDQ does not claim integer
cache storage, serving acceleration, or model-quality recovery. State-only vLLM cache and worker integration is covered separately by #2541,
with QDQ at native invocation boundaries rather than the training replay cadence.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

❌ Patch coverage is 56.33001% with 307 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.28%. Comparing base (8ea824b) to head (cbb8874).

Files with missing lines	Patch %	Lines
...ch/kernels/quantization/linear_attention/decode.py	0.00%	183 Missing ⚠️
...quantization/linear_attention/fla_chunk_delta_h.py	0.00%	30 Missing ⚠️
...lopt/torch/quantization/linear_attention/decode.py	82.71%	28 Missing ⚠️
modelopt/torch/quantization/plugins/kda.py	51.11%	22 Missing ⚠️
...ch/quantization/linear_attention/decode_prefill.py	83.09%	12 Missing ⚠️
...t/torch/quantization/linear_attention/reference.py	88.60%	9 Missing ⚠️
...orch/kernels/quantization/linear_attention/int8.py	0.00%	8 Missing ⚠️
...odelopt/torch/quantization/linear_attention/kda.py	73.07%	7 Missing ⚠️
...opt/torch/quantization/linear_attention/prefill.py	80.95%	4 Missing ⚠️
...ion/linear_attention/fla_chunk_gated_delta_rule.py	0.00%	2 Missing ⚠️
... and 2 more

Additional details and impacted files

@@                       Coverage Diff                        @@
##           kaix/linear-attention-qat-m1    #2519      +/-   ##
================================================================
- Coverage                         70.50%   70.28%   -0.23%     
================================================================
  Files                               610      618       +8     
  Lines                             68036    68339     +303     
================================================================
+ Hits                              47969    48029      +60     
- Misses                            20067    20310     +243     

Flag	Coverage Δ
unit	57.63% <56.33%> (-0.06%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.