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

[kaix-nv]

Stack / merge order: #2497 foundation → #2519 decode + INT8 → #2503 GDN/KDA prefill → #2507 approximate inverse.

All entries remain drafts. #2506 (KDA prefill) and #2509 (original decode) 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.

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

• Restacked source tree exactly matches the validated combined snapshot
  7c00dd086d0a882a5bd88337f4796f421ec63426; only commit ancestry changes.
• Fresh CPU suite and all pre-commit hooks: 151 CPU tests passed; all hooks passed.
• Earlier local integration runs passed four CUDA composed prefix/decode cases
  and thirteen FLA KDA layer cases, including INT8 and experimental solve
  integration. These GPU runs predate the restacked commits and were not rerun
  during restacking; the fresh validation above is CPU and pre-commit.
• Tests cover actual-polynomial gradients, config/degree validation, residual
  identities, composed operand/state QDQ, save/restore, and packed sequence tails.

Historical model screening used arcee-ai/AFM-4.5B-Base-KDA-Only and a fixed
WikiText-2 validation set. Degrees 3, 7, 15, and 31 all failed the declared NLL
margin. Degree 63 is not a numerically safe fallback. Synthetic H100 measurements
also established no speed advantage. The
qualification report
records exact revisions, data, intervals, and limitations. These historical
measurements were not rerun during restacking; no new quality claim is made.

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,
Megatron distributed requalification, and vLLM integration remain pending work
elsewhere in the stack. 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

---

_{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.37%. Comparing base (bd1f350) to head (ff21e39).

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.37%   70.37%   -0.01%     
================================================================
  Files                               620      621       +1     
  Lines                             68513    68598      +85     
================================================================
+ Hits                              48218    48273      +55     
- Misses                            20295    20325      +30     

Flag	Coverage Δ
unit	57.77% <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.