Add experimental GDN/KDA decode and replay QAT

[kaix-nv]

Superseded by #2519: Decode is delivered before prefill and inverse, with signed INT8 state and replay-anchor support. Prefill composition follows in #2503 and inverse in #2507. This closed PR and its branch retain the original qualification history.

---

What does this PR do?

Type of change: New feature, new example, new tests, documentation.

Adds experimental differentiable GDN/KDA decode policies for training: per-token FP8 state writes, optional log-retention grid rounding, and anchor-plus-encoded-update replay. Workloads supply explicit prefix lengths, preserving gradients across prefill/decode handoff, anchor refresh, and continuation. Policies persist through ModelOpt save/restore.

A Torch reference defines write/readout order and codec metadata. A fused FP32 Triton forward/backward implements token and encode-once replay with internal state checkpoints. A fixed key-reduction tree and matching dynamic FP8 scale arithmetic prevent tiny arithmetic differences from flipping rounding ties and accumulating along near-unit-decay trajectories. Replay factors are computed from the current reconstructed trajectory, never a teacher state.

The default execution path is unchanged. Serving caches, compressed storage, native low-precision speedups, and higher-order fused differentiation are outside this draft. The study keeps the exact solve; it does not enable the Neumann candidate rejected in #2507.

Usage

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

mtq.quantize(model, {
    "quant_cfg": [
        {"quantizer_name": "*", "enable": False},
        {"quantizer_name": "*kda_state_quantizer", "cfg": {
            "num_bits": (4, 3), "type": "dynamic", "axis": (0, 1),
        }},
    ],
    "algorithm": None,
    "linear_attention": [{"module_name": "*", "cfg": {
        "backend": "matmul",
        "decode": {"mode": "replay", "implementation": "triton",
                   "replay": {"window": 8, "factor_qdq": True}},
    }}],
})
with linear_attention_training_phase(model, [64]):
    with torch.autocast("cuda", dtype=torch.bfloat16):
        loss = model(input_ids=ids, labels=labels, use_cache=False).loss
    loss.backward()  # Context also covers activation-checkpoint recomputation.

Use *gdn_state_quantizer for GDN. Token mode omits the replay settings. See docs/linear_attention_decode.md for carry/readout contracts and unsupported cases.

Testing

• Local CPU: 139 numerical/configuration/comparison tests plus two replay-reencoding continuation cases passed.
• Local GPU: 50 kernel/layer/composition tests passed; eight long-trajectory cases also passed a tightened relative-error limit of 1e-4.
• Two H100s: 154 passed (148 regression + 4 value-block tails + 2 composition), 9 expected Hopper/TileLang FP32 skips. Includes GDN TP=1/2 sharded save/restore and optimizer updates, KDA layer restore/training, prefix/empty/carry handling, and reference/fused gradients. All eight tighter long-trajectory cases also passed on H100.
• Full Sphinx build with FLA imports disabled passed; pre-commit, mypy, and diff checks passed.

Pinned Arcee KDA / Wikitext-2 pilot: 32 validation blocks selected log grid 1/256 before test access. Exact, token-state FP8, state FP8 + selected grid, and replay each completed 32 attention-only updates with finite gradients for all trainable parameters and verified weight changes. On 64 held-out suffix-scored blocks, post-training perplexities were 12.6413 / 12.6567 / 12.6530 / 12.6466. All approximations met the 0.02 NLL pilot margin; the largest 95% upper bound was 0.002703. This is a short-context, single-model pilot, not broad quality recovery.

For complete [1,257,4,128] prefix/suffix training on H100 (prefix 64, three warmups, 20 interleaved samples), KDA token/decay/replay took 11.265 / 11.310 / 11.569 ms versus 228.732 / 235.972 / 327.361 ms for matching Torch references. Exact FLA with BF16 Q/K/V took 2.347 ms. Numerical outputs/state matched the references and maximum gradient relative-norm error was below 3.4e-7. GDN results and paired intervals are in docs/linear_attention_decode_study.md; raw receipts bind retained source snapshots. Final import/comment cleanup preserves the non-import AST and was regression-tested; the benchmark harness differs only in typing import-name ordering.

Additional tests cover full-width state/value tails at codec blocks 32/128 and composition of prefix FP8, Neumann solve, elementwise rounding, and replay across packed/empty sequences and state layouts. These synthetic composition tests do not promote the Neumann algorithm for model quality.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ Opt-in saved policy; defaults unchanged.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A No new dependency or copied third-party implementation.
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅
• Did you get Claude approval on this PR?: ❌ Not requested; remains a draft.

Additional Information

Stacked on #2507 (M4), following #2506 (KDA prefill), #2503 (GDN materialized prefill), and #2497 (GDN state/W QDQ). This draft models numerical behavior for QAT and has no inference performance or compressed-storage claim.

[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 49.54545% with 222 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.38%. Comparing base (34e549e) to head (ae9d2a7).

Files with missing lines	Patch %	Lines
...ch/kernels/quantization/linear_attention/decode.py	0.00%	176 Missing ⚠️
...lopt/torch/quantization/linear_attention/decode.py	82.27%	28 Missing ⚠️
...ch/quantization/linear_attention/decode_prefill.py	82.60%	12 Missing ⚠️
modelopt/torch/quantization/plugins/kda.py	40.00%	6 Missing ⚠️

Additional details and impacted files

@@                       Coverage Diff                        @@
##           kaix/linear-attention-qat-m4    #2509      +/-   ##
================================================================
- Coverage                         70.52%   70.38%   -0.14%     
================================================================
  Files                               617      620       +3     
  Lines                             68119    68559     +440     
================================================================
+ Hits                              48039    48257     +218     
- Misses                            20080    20302     +222     

Flag	Coverage Δ
unit	57.78% <49.54%> (-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.

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-09-23 01:07 UTC