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

[kaix-nv]

Linear-attention series — 6 PRs

Order	PR	Depends on
1/6	#2497 GDN state/W QAT foundation	main
2/6	#2519 GDN/KDA state QAT + INT8	#2497
3/6	#2657 Megatron Bridge linear attention QAT/QAD example	#2519
4/6	#2541 vLLM GDN/KDA state-only fake quantization	#2657
5/6	#2503 GDN/KDA prefill GEMM quantization	#2541
6/6	#2507 Experimental GDN/KDA approximate inverse	#2503

The five open PRs form one native GitHub stack in the order shown. #2497 has landed, so #2519 targets main. #2541 now targets #2657. Rebase each remaining descendant after its immediate parent merges.

#2541 applies TensorQuantizer before native vLLM prefill/decode calls. Serving-time prefill-GEMM quantization remains deferred until an optimized fused kernel is available.

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.

This follow-up also introduces the shared decode benchmark and quality-comparison tooling, configs, and comparison tests deferred from #2519. Its study workflow extends the minimal QAT example with fixed-data evaluation and measurement receipts.

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

Compilation-fixture update: this branch is restacked on #2497's separate follow-up commit 95766de709ac. The changed test modules passed at #2497 (26 passed, 20 hardware skips in each cold/warm run) and at the #2507 stack tip (64 passed, 20 hardware skips on two RTX A6000 GPUs); intermediate PRs were not separately rerun. Functional calls created no new tracked kernel binaries and kept the default 120-second cap. All six source trees match the validated trees, and commit hooks passed. Native FP8-state/Hopper cases remain hardware-gated.

Earlier scope-specific validation follows.

For the earlier README-only restack, all runtime code, tests, and study scripts were byte-for-byte identical to its preceding head; the README inherits the state-quantization enablement guidance. Focused README pre-commit hooks, diff checks, and signed-commit verification passed. Model training, distributed integration, and quality measurements were not rerun.

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. Those earlier results did not include Megatron/full-FLA-layer or model-quality qualification. 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

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (kaix/linear-attention-vllm@2b9e3a3). Learn more about missing BASE report.

Additional details and impacted files

@@                      Coverage Diff                      @@
##             kaix/linear-attention-vllm    #2503   +/-   ##
=============================================================
  Coverage                              ?   77.50%           
=============================================================
  Files                                 ?      621           
  Lines                                 ?    68539           
  Branches                              ?        0           
=============================================================
  Hits                                  ?    53122           
  Misses                                ?    15417           
  Partials                              ?        0           

Flag	Coverage Δ
unit	57.60% <100.00%> (?)

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-28 06:12 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[ifed-ucsd]

Does this actually test against the triton kernel? or is the reference a pure pytorch implementation?