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

[kaix-nv]

Stack / merge order: #2497 foundation → #2519 decode + INT8 → #2503 GDN/KDA prefill → #2507 approximate inverse.

All entries remain drafts. #2506 (KDA prefill) and #2509 (original decode) 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.

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

On RTX A6000/SM86, Torch 2.9.1+cu128, Triton 3.5.1, FLA 0.5.1:

• CPU reference, policy, INT8, checkpoint, and quality-comparison tests: 73 passed.
• Decode and GDN chunk GPU suites: 91 passed, 12 expected hardware skips.
• FLA KDA layer tests: 6 passed, including INT8 restore, backward, optimizer steps,
  and activation checkpointing.
• INT8 benchmark script functional smoke: passed; no performance conclusion.
• All pre-commit hooks passed, including Ruff, mypy, and Markdown checks.
• Sequential application of decode, prefill, and inverse patches reproduced all
  three recorded Git trees exactly.

The INT8 long-trajectory tests caught and fixed FP32 scale evaluation drift:
Triton now uses the same rounded reciprocal multiplication as Torch scalar
division. Tolerances were not relaxed. Tests cover 257-token trajectories,
nonzero initial states, grouped heads, packed tails, both readouts, split carry,
all state scale widths, and first-order input/state gradients.

Megatron tests were added for INT8 and decode at TP=1/2, but local collection
was blocked by the absent megatron package. Native FP8/Hopper requalification,
INT8 model-quality comparisons, and vLLM serving/cache integration are pending.
Keep this PR in draft until its remaining qualification is complete.

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. vLLM cache and
multi-GPU calibration integration remain required follow-up work in the decode delivery.

[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.