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

[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 configurable GDN and KDA prefill operand fake quantization after the decode
and INT8 state infrastructure in #2519. This PR is stacked on the state-only
vLLM plugin in #2541. 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.

GDN and KDA share the dispatch between pure prefill and explicit prefill/decode execution. Operand modules follow ModelOpt temporary-attribute cleanup. The prefill benchmark reuses the decode benchmark measurement helper, with one configuration pass per quantizer.

Execution configuration is introduced by #2519 and extended here for prefill arithmetic. This PR replaces the decode prefix implementation with the shared batched prefill core, removes the superseded prefix helper, and keeps independent numerical oracles under tests.

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

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.

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

• 145 CPU tests passed: GDN lifecycle, references, decode/INT8/Hadamard, GDN prefill, and KDA prefill.
• 117 GPU tests passed: GDN/KDA prefill plus decode, including outputs, final states, gradients, FP8/NVFP4 operands, arithmetic composition, and continuation.
• Pre-commit, diff checks, and commit-signature verification passed.

The tests preserve independent numerical oracles under the test package and cover checkpoint compatibility, default-disabled handles, packed tails, grouped heads, and autocast. Megatron/full-FLA-layer and model-quality qualification were not rerun during this restack. No serving-speed or quality-recovery claim is made.

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. #2541 supplies
state-only vLLM integration; serving-time prefill-GEMM quantization remains
deferred until an optimized fused kernel is available.

[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.28%. Comparing base (087123e) to head (580c3dc).

Additional details and impacted files

@@                      Coverage Diff                       @@
##           kaix/linear-attention-vllm    #2503      +/-   ##
==============================================================
+ Coverage                       70.20%   70.28%   +0.08%     
==============================================================
  Files                             619      621       +2     
  Lines                           68475    68649     +174     
==============================================================
+ Hits                            48070    48253     +183     
+ Misses                          20405    20396       -9     

Flag	Coverage Δ
unit	57.69% <100.00%> (+0.12%)	⬆️

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
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2503/
Built to branch gh-pages at 2026-09-24 06:38 UTC. Preview will be ready when the GitHub Pages deployment is complete.