[5/8] ReplaySSM state QAT and two-precision GDN

[kaix-nv]

Linear-attention series — 8 PRs

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

The seven open PRs form native GitHub stack #2806 in this order. #2497 has landed, so #2519 targets main. #2805 contains NVFP4 and is based on #2657; #2788 and #2541 are rebased in sequence. #2503 and #2507 retain their existing commits; their rebases remain deferred until their immediate parents merge.

#2788 owns ReplaySSM training, including the implemented two-precision GDN profile. #2541 owns ordinary state-QDQ serving. The experimental ReplaySSM serving follow-up is published separately on kaix/linear-attention-replayssm-serving; it is not included in this PR or the numbered stack yet. 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 ReplaySSM state QAT independently of the ordinary GDN/KDA state-QDQ path. Preserve the existing private-fork INT8/Hadamard comparison profile and implement a two-precision GDN profile with an FP32 recurrence anchor and dynamic FP8 E4M3 read copy.

The two-precision suffix imports upstream emulation.GDNReplay at 73a9592e9d8a2268620e8883a82f5cec7442ed36: one native forward per token, with a window-wise adjoint during backward. Each 16-token flush recomputes recurrence updates from raw inputs and the FP32 anchor. Read-side quantization errors do not overwrite the recurrence anchor. ModelOpt owns the quantizer configuration and checkpoint lifecycle; no upstream kernels are copied.

Initial scope: GDN, normalized Q/K, K=128, FP32 recurrence, FP8 or disabled read quantization, and first-order gradients. KDA two-precision support, lower-precision anchors, NVFP4 reads and stochastic rounding remain follow-ups. Quantized values use floating storage; this is numerical emulation, with no cache-compression or speed claim.

Usage

Expose only the pinned upstream checkout's emulation package on PYTHONPATH, alongside the example's vLLM 0.30/FLA environment. Do not expose the checkout's bundled vLLM package.

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

cfg = load_recipe("general/ptq/gdn_replayssm_fp8_read").quantize.model_dump()
model = mtq.quantize(model, cfg)
with linear_attention_training_phase(model, [64], sequence_lengths=[97]):
    loss = suffix_training_loss(model, batch)
    loss.backward()

The caller supplies an initialized supported model and suffix-only QAT or QAD loss. The recipe disables gdn_state_quantizer and enables dynamic E4M3 gdn_read_quantizer, grouped by 128 key channels per value row. See examples/llm_qat/linear_attention/REPLAYSSM.md for setup and the forward/backward schedule. The existing INT8/Hadamard profile still requires its separate compatible fork environment.

Testing

Published training head: 0a4625eb04. The following are local/cluster results; GitHub CI for this head is pending.

• 40 focused CPU tests passed for policy, checkpoint lifecycle and recipe documentation.
• Two focused GPU tests passed on RTX A6000 and DFW H100. H100 jobs 20242052 and 20242382 completed successfully.
• Synthetic single-layer GDN: 64 prefix tokens, 33 suffix tokens and 3 padding tokens; three SGD steps each for QAT and QAD. Weight updates were nonzero, suffix gradients reached the prefix, and padding gradients were zero.
• ModelOpt save/restore preserved outputs exactly. Reloaded inference matched direct native prefill plus ReplaySSM decode. Activation recomputation preserved outputs and gradients.
• QAT cross-entropy: 2.770704746 to 2.770206928. QAD KL: 1.9964632e-7 to 1.9879703e-7. These tiny synthetic runs validate workflow execution, not pretrained accuracy recovery.
• K=256 failed the predefined numerical gate and is not enabled. No tolerance was relaxed.
• Changed-file pre-commit and git diff --check passed. The private-fork INT8/Hadamard GPU cases were retained but not requalified in this update.

Full-model Megatron/Bridge QAT/QAD, HF export-to-serving qualification and performance evaluation remain outstanding for this new profile. Engine-owned ReplaySSM cache integration is in the separate serving branch, not in this PR.

Before your PR is "Ready for review"

• Backward compatible: additive opt-in profile; existing ordinary state-QDQ behavior is unchanged.
• Copied code or dependencies: no upstream kernels copied; the optional standalone ReplaySSM dependency is imported lazily at the documented pin.
• Necessary tests: focused CPU lifecycle tests and real GPU native-forward/gradient tests.
• Changelog: updated.
• Claude approval: pending; keep draft.

Additional Information

Base: #2805 (kaix/linear-attention-state-nvfp4). Ordinary vLLM state-QDQ serving remains in #2541. ReplaySSM serving uses a separate worker and cache lifecycle, published on kaix/linear-attention-replayssm-serving at f0ddc3221e.

[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

❌ Patch coverage is 22.32416% with 254 lines in your changes missing coverage. Please review.
✅ Project coverage is 68.23%. Comparing base (09bdcee) to head (0a4625e).
⚠️ Report is 1 commits behind head on kaix/linear-attention-state-nvfp4.

Files with missing lines	Patch %	Lines
...ch/quantization/linear_attention/_two_precision.py	20.20%	79 Missing ⚠️
...ls/quantization/linear_attention/serving/replay.py	0.00%	74 Missing ⚠️
...lopt/torch/quantization/linear_attention/replay.py	28.88%	64 Missing ⚠️
...pt/torch/quantization/linear_attention/training.py	8.00%	23 Missing ⚠️
...lopt/torch/quantization/linear_attention/decode.py	14.28%	6 Missing ⚠️
...s/quantization/linear_attention/serving/forward.py	0.00%	3 Missing ⚠️
...elopt/torch/quantization/linear_attention/utils.py	40.00%	3 Missing ⚠️
...odelopt/torch/quantization/linear_attention/gdn.py	0.00%	1 Missing ⚠️
...odelopt/torch/quantization/linear_attention/kda.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@                          Coverage Diff                          @@
##           kaix/linear-attention-state-nvfp4    #2788      +/-   ##
=====================================================================
- Coverage                              68.99%   68.23%   -0.77%     
=====================================================================
  Files                                    656      662       +6     
  Lines                                  72133    74071    +1938     
=====================================================================
+ Hits                                   49769    50540     +771     
- Misses                                 22364    23531    +1167     

Flag	Coverage Δ
unit	59.55% <23.70%> (-0.09%)	⬇️

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-2788/
Built to branch gh-pages at 2026-10-15 03:21 UTC. Preview will be ready when the GitHub Pages deployment is complete.