[kaix-nv]

Linear-attention series — 10 PRs

Order	PR	Depends on
1/10	#2497 GDN state/W QAT foundation	main
2/10	#2519 GDN/KDA state QAT + INT8/FP8	#2497
3/10	#2657 Megatron Bridge QAT/QAD example	#2519
4/10	#2805 NVFP4 recurrent-state QAT	#2657
5/10	#2829 TensorQuantizer stochastic rounding	#2805
6/10	#2788 ReplaySSM state QAT and two-precision GDN	#2829
7/10	#2541 vLLM state-only fake quantization	#2788
8/10	#2878 ReplaySSM state quantization serving	#2541
9/10	#2503 GDN/KDA prefill GEMM quantization	#2878
10/10	#2507 Experimental approximate inverse	#2503

The nine open PRs form native GitHub stack #2879 in this order. #2497 has landed, so #2519 targets main.

#2829 supplies the shared rounding policy before ReplaySSM training and serving. #2788 owns the ReplaySSM training foundation; #2541 owns ordinary state-QDQ serving; #2878 adds the experimental two-precision ReplaySSM serving backend and native stochastic NVFP4 read integration. The serving branch inherits rounding from #2829 without a temporary backport.

#2503 and #2507 retain their existing code and commit heads. Their functional ports and rebases remain deferred until their immediate parents land; their bases and table positions express the intended merge order. Serving-time prefill-GEMM QDQ remains deferred until an optimized fused kernel is available.

What does this PR do?

Type of change: New feature.

Add optional dynamic NVFP4 recurrent-state fake quantization through ModelOpt's existing TensorQuantizer. This follows #2657 and precedes #2788; the state-QAT runtime remains in #2519. State values use E2M1, 16-value last-axis blocks, E4M3 block scales, a dynamic global scale and identity STE. The new recipe leaves weights and activations unchanged.

The implementation extends format validation and adds a recipe; it introduces no quantization kernel. State storage remains floating point. This PR does not qualify NVFP4 in a vLLM engine or implement ReplaySSM's planned two-precision design.

Usage

import modelopt.torch.quantization as mtq
from modelopt.recipe import load_recipe
from modelopt.torch.quantization.linear_attention import linear_attention_training_phase

recipe = load_recipe("general/ptq/linear_attention_state_nvfp4_dynamic")
model = mtq.quantize(model, recipe.quantize)
with linear_attention_training_phase(model, prefill_lengths=[64], sequence_lengths=[128]):
    loss = training_step(model, batch)
loss.backward()

Use a supported Megatron-Core GDN/KDA model. For Kimi-Linear/Kimi-K3, explicitly select the vllm_kimi_k3 policy and its serving restrictions. Quantizer configuration is num_bits=(2,1), block_sizes={-1:16, "type":"dynamic", "scale_bits":(4,3)}, top-level type="dynamic", and pass_through_bwd=True.

Testing

• H100 GPU: native TensorQuantizer NVFP4 state QDQ matches the state path, changes finite inputs, and propagates the identity gradient; passed in DFW job 20217113.
• CPU recipe conversion and ModelOpt save/restore preserve the NVFP4 quantizer metadata; 31 focused CPU tests passed on this slice.
• After restacking, 52 focused CPU tests passed across state formats, policies/checkpoints, recipe documentation, replay numerics and example batch metadata.
• Changed-file pre-commit hooks and git diff --check passed.

The GPU job also passed the parent example's QAT/QAD checkpoint round trips and prompt-derived QAD test (four tests total, 90.50s). Those workflow tests use INT8, not NVFP4; NVFP4 training quality and engine qualification remain unmeasured.

Before your PR is "Ready for review"

• Backward compatibility: additive format and recipe; existing INT8/FP8 settings remain supported.
• Dependencies/copying: no new dependency or copied kernel; reuses TensorQuantizer's NVFP4 implementation and numerical recipe.
• Necessary tests: CPU recipe/restore and one focused CUDA QDQ/gradient test.
• Changelog: one user-facing feature entry.
• Review: draft; human/Claude review pending.

Additional Information

Step 4/10. About 101 added non-test lines over #2657. NVFP4 was separated from #2519 in response to review.

[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

• Autofix · 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.
✅ Project coverage is 68.13%. Comparing base (acdc5ab) to head (fa97dd3).

Additional details and impacted files

@@                        Coverage Diff                         @@
##           kaix/linear-attention-qat-example    #2805   +/-   ##
==================================================================
  Coverage                              68.12%   68.13%           
==================================================================
  Files                                    665      665           
  Lines                                  75687    75689    +2     
==================================================================
+ Hits                                   51559    51567    +8     
+ Misses                                 24128    24122    -6     

Flag	Coverage Δ
unit	59.70% <100.00%> (+<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-2805/
Built to branch gh-pages at 2026-10-14 21:45 UTC. Preview will be ready when the GitHub Pages deployment is complete.