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

[kaix-nv]

Linear-attention PR stack — 5 drafts

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	#2519
5/5	#2507 Experimental GDN/KDA approximate inverse	#2503

Numbering gives review order; the dependency column records the Git bases. #2541 and #2503 can proceed independently after #2519. #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 configurable GDN and KDA prefill operand fake quantization after the decode
and INT8 state infrastructure in #2519. Each of the eight logical matmul sites
quantizes its actual transformed operands through ModelOpt TensorQuantizer,
with independent FP8/NVFP4 settings, accumulator rounding schedules, and named
elementwise rounding points. State carry and all rounding sites remain differentiable.
KDA uses causal per-channel gate differences to avoid overflowing inverse-decay factors.

This combines the prefill functionality previously split between this PR and
#2506. The triangular solve is exact; approximate inverse belongs to #2507.
Existing FP8/INT8 state formats, decode policies, explicit phase handoff, and
ModelOpt save/restore remain available. Operand scales and state-write scales
are independent. Working arithmetic remains FP32 inside BF16/FP16 autocast.

Usage

import torch
import modelopt.torch.quantization as mtq

torch.set_float32_matmul_precision("highest")
fp8 = {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)}
model = mtq.quantize(model, {
    "quant_cfg": [
        {"quantizer_name": "*", "enable": False},
        {"quantizer_name": "*linear_attn_sites.*", "cfg": fp8},
        {"quantizer_name": "*gdn_w_quantizer", "cfg": fp8},
        {"quantizer_name": "*kda_w_quantizer", "cfg": fp8},
    ],
    "algorithm": None,
    "linear_attention": [{"module_name": "*", "cfg": {"backend": "matmul"}}],
})

The state-read LHS uses the existing GDN/KDA W handle. See
the GDN guide
and the KDA guide
for NVFP4, individual sites, scale domains, arithmetic policies, and framework limits.

Testing

• Restacked source tree exactly matches the validated prefill snapshot
  a5b2bfa9c864782ad194b66e1a6ef87cd2331e47; only commit ancestry changes.
• Fresh CPU suite and all pre-commit hooks: 134 CPU tests passed; all hooks passed.
• Earlier local validation passed four CUDA packed-prefix/replay composition
  cases covering GDN/KDA and FP8/INT8 states. GPU tests were not rerun during
  restacking; the fresh validation above is CPU and pre-commit.
• Coverage includes outputs, final states, input/state gradients, grouped heads,
  packed tails, autocast, independent QDQ/rounding oracles, config validation,
  checkpoint restore, and compatibility with earlier disabled handles.

Local environment: RTX A6000/SM86, Torch 2.9.1+cu128, Triton 3.5.1,
fla-core==0.5.1; layer tests use flash-linear-attention==0.5.1.
Historical H100/quality measurements remain scoped to their documented revisions.
Distributed Megatron and INT8 model-quality requalification for this stack remain pending.

Before your PR is "Ready for review"

• Backward compatible: Yes; new operand handles start disabled and old state formats remain valid.
• Copied code/dependency guidance: Existing provenance retained; FLA remains an optional pinned dependency.
• Necessary tests: Yes.
• Changelog: Updated.
• Claude approval: Pending; keep draft.

Additional Information

This materialized backend emulates training numerics. It does not provide native
low-precision MMA, compressed states, or a serving speedup. vLLM cache and
multi-GPU calibration integration remain separate required decode work.

[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

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 70.37%. Comparing base (cbb8874) to head (bd1f350).
⚠️ Report is 1 commits behind head on kaix/linear-attention-decode-first.

Additional details and impacted files

@@                          Coverage Diff                           @@
##           kaix/linear-attention-decode-first    #2503      +/-   ##
======================================================================
+ Coverage                               70.28%   70.37%   +0.09%     
======================================================================
  Files                                     618      620       +2     
  Lines                                   68339    68513     +174     
======================================================================
+ Hits                                    48029    48218     +189     
+ Misses                                  20310    20295      -15     

Flag	Coverage Δ
unit	57.76% <100.00%> (+0.13%)	⬆️

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.