[7/7] Experimental GDN/KDA approximate-inverse QAT

[kaix-nv]

Linear-attention series — 7 PRs (1 merged, 6 open)

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

The six open PRs form native GitHub stack #2658 in the order shown; #2497 is retained as the merged foundation in this seven-PR series. #2497 is merged, so #2519 targets main. #2657 contains the training example split from #2519; #2562 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. #2506 and #2509 are superseded and closed.

What does this PR do?

Type of change: New feature.

Add an opt-in approximate inverse after the GDN/KDA prefill infrastructure in
#2503. The saved solve policy selects an explicit Neumann polynomial degree and
either Torch or CUDA FP32 execution. Backward differentiates the actual
polynomial; the implementation never silently changes degree or falls back.
Exact triangular solve remains the default.

The current Neumann candidate failed the pinned KDA model-quality screen.
Keep this PR experimental and in draft. The historical study and failed results
are preserved; a successful kernel or optimizer test is not quality recovery.

The solve benchmark shares CUDA timing, warmup, interleaving, and aggregation with the decode/prefill examples. Each benchmark retains its own loss and correctness checks and records the shared helper source hash.

Execution configuration is introduced by #2519, extended for prefill by #2503, and extended here with the solve policy. Numerical test oracles are imported from the test utility package.

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 modelopt.torch.quantization as mtq

policy = {
    "backend": "matmul",
    "solve": {"method": "neumann", "degree": 3, "implementation": "triton"},
}
model = mtq.quantize(model, {
    "quant_cfg": [{"quantizer_name": "*", "enable": False}],
    "algorithm": None,
    "linear_attention": [{"module_name": "*", "cfg": policy}],
})
# Research example only: this degree was not accepted by the model-quality gate.

Use {"method": "exact"} (the default) for the supported baseline.
See the solve guide.

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:

• 131 CPU tests passed: solve/configuration, GDN lifecycle, decode, and GDN/KDA prefill.
• 59 GPU tests passed: 16 inverse-polynomial/composition cases plus 43 GDN/KDA prefill cases.
• A checkpoint saved by the simplified [1/7] GDN state/W QAT foundation #2497 restored at this stack tip with its state/W settings preserved and new prefill operand quantizers disabled.
• Pre-commit, diff checks, and signed-commit verification passed.

The tests cover gradients of the actual polynomial, degree validation, residual identities, composed QDQ, save/restore, and packed tails. Those earlier results did not include Megatron/full-FLA-layer integration.

Historical screening of arcee-ai/AFM-4.5B-Base-KDA-Only on fixed WikiText-2 validation data failed the declared NLL margin for degrees 3, 7, 15, and 31. Degree 63 is not a safe fallback, and prior synthetic H100 measurements established no speed advantage. The qualification report retains the original revisions and limits. Those measurements were not rerun; the candidate remains experimental and in draft.

Before your PR is "Ready for review"

• Backward compatible: Yes; exact solve remains the default.
• Copied code/dependency guidance: Existing provenance retained; no new restacking dependency.
• Necessary tests: Yes.
• Changelog: Updated.
• Claude approval: Pending; keep draft.

Additional Information

Depends on #2503 and comes last in the stack. INT8 model-quality comparisons
and Megatron distributed requalification remain pending. #2541 supplies
state-only vLLM integration; serving-time prefill-GEMM quantization is deferred
until an optimized fused kernel is available. A better-conditioned inverse
approximation needs its own numerical and model-quality evidence.

[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

❌ Patch coverage is 66.29213% with 30 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.28%. Comparing base (580c3dc) to head (f3ce747).

Files with missing lines	Patch %	Lines
...h/kernels/quantization/linear_attention/neumann.py	34.09%	29 Missing ⚠️
...elopt/torch/quantization/linear_attention/solve.py	97.22%	1 Missing ⚠️

Additional details and impacted files

@@                       Coverage Diff                        @@
##           kaix/linear-attention-qat-m2    #2507      +/-   ##
================================================================
- Coverage                         70.28%   70.28%   -0.01%     
================================================================
  Files                               621      622       +1     
  Lines                             68649    68734      +85     
================================================================
+ Hits                              48253    48308      +55     
- Misses                            20396    20426      +30     

Flag	Coverage Δ
unit	57.70% <66.29%> (+<0.01%)	⬆️

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-2507/
Built to branch gh-pages at 2026-09-24 06:42 UTC. Preview will be ready when the GitHub Pages deployment is complete.