[2/6] Torch GDN/KDA decode QAT with INT8 recurrent state

[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 recurrent-state QAT for Megatron-Core GDN and KDA. Each sequence uses a chunked prefill prefix followed by a recurrent suffix. TensorQuantizer applies fake QDQ at the handoff and subsequent state writes, so training includes the recurrent quantization effects encountered during decode.

An explicit linear-attention policy selects the native forward arithmetic and state/checkpoint schedule. Training uses native forward values with differentiable FLA/Torch surrogate gradients and identity STE through QDQ; this does not claim an exact backward for inference kernels.

• precision="vllm" selects the fixed vLLM 0.30.0 Triton APIs. Public vLLM profiles require vllm==0.30.0; older-version compatibility probing is removed. GDN qualification covers pure single-token decode with matching Triton prefill and packed-decode settings. Generic KDA selects standalone FLA/Triton APIs without model-level serving parity.
• precision="vllm_kimi_k3" selects the Kimi-K3/Kimi-Linear Triton profile. Serving must use Triton prefill/decode, Q/K normalization and gate_lower_bound=None.
• precision="replayssm" requires an optional private quantized-ReplaySSM fork, unavailable in public vLLM releases. Without access to it, use a block32 recipe. The fork profile uses Hadamard checkpoints with replay_window=1 by default.

The public training phase supplies per-sequence prefix and valid-token lengths, including packed padding. Active state QAT requires native Megatron _compute_gates and _forward_compute hooks for raw gates and packed sequence lengths; Core 0.19.2 lacks them and is rejected before forward execution. Selective GDN recompute captures the phase for backward on supported Core builds. Active state QAT rejects context parallelism and full-layer recompute. Disabled state quantizers outside a phase preserve the original module path.

This change supports floating-storage state fake quantization. It does not add a vLLM server or native compressed cache. #2541 owns serving integration; #2657 owns the training example. Prefill GEMM quantization remains a follow-up. Descendant rebases remain deferred until their immediate parent merges.

Megatron-Core QAT/QAD workflow

Megatron-Core owns the model, distributed forward/backward schedule and optimizer. ModelOpt replaces the GDN/KDA attention-core call and inserts differentiable state fake quantization. vLLM supplies imported forward kernels; training does not require a running vLLM server. FLA supplies the prefix gradient computation.

1. Load and configure the student. Build the Megatron-Core model and load pretrained weights, then call mtq.quantize(student, recipe.quantize) on each local model chunk. quant_cfg configures the gdn_state_quantizer / kda_state_quantizer; linear_attention selects the execution profile. The shipped dynamic INT8 state-only recipe needs no calibration. A quantized student checkpoint can instead restore its saved ModelOpt configuration.
2. Choose a prefix and suffix for each microbatch. Enter linear_attention_training_phase around the student's forward, providing prefix lengths and valid sequence lengths. For a 128-token sequence with prefill_lengths=[64], the first 64 tokens use chunked prefill and the last 64 use recurrent updates. Both parts consume the training sequence's tokens; this is teacher-forced training, not autoregressive sampling. Exclude padding from valid lengths and apply the loss mask to the suffix. The phase controls attention execution; the caller controls loss masking. Packed storage boundaries and valid lengths remain separate.
3. Run the quantized student forward. Megatron captures raw gates and routes its attention-core call through gdn_state_qat / kda_state_qat into _prefill_decode_forward. The prefix produces a state without internal state QDQ. For a nonempty ordinary suffix, QDQ is applied at the handoff; each token then reads the stored state, computes its output from the updated working state, and applies QDQ before storing state for the next token. State remains in floating-point storage. The optional ReplaySSM profile follows its configured checkpoint/replay schedule instead.
4. Compute the loss and update student weights. QAT uses the language-model loss. QAD uses a frozen unquantized teacher on the same tokens and a distillation loss on student logits, optionally combined with the language-model loss. Both use the same student state-QDQ path. Native kernels provide forward values; FLA provides prefix surrogate gradients, Torch provides recurrent surrogate gradients, and QDQ uses an identity straight-through estimator. The handoff state stays attached, so suffix losses also train the prefix. Megatron's schedule performs backward and the optimizer step; supported selective recompute captures the microbatch's phase for backward after context exit.
5. Validate, save and resume. Evaluate with the same prefix/suffix split and mask. The comparison baseline must also be converted with the same serving policy and phase, with state quantizers disabled, to isolate quantization effects. Megatron checkpoints retain model weights and ModelOpt quantizer/policy metadata through the existing extra-state machinery; supply fresh per-microbatch phase metadata after restore. Full-prefill-only evaluation does not measure decode state-quantization quality.

Ordinary state-QDQ path:

flowchart LR
    A[Training tokens] --> B[Student chunked prefix]
    B --> C[State QDQ at handoff]
    C --> D[Student recurrent suffix with state QDQ writes]
    D --> L[Loss on valid suffix tokens]
    T[Labels: QAT] --> L
    A --> E[Frozen teacher: QAD]
    E --> L
    L --> F[Megatron backward and student optimizer step]

Loading

The call path is plugins/megatron.py → plugins/gdn.py / plugins/kda.py → linear_attention/gdn.py / kda.py → training.py → _vllm_autograd.py / decode.py → imported serving kernels. This PR provides the core APIs and Megatron adapter. #2657 adds the Bridge training entry point and QAT/QAD checkpoint export/load workflow; #2541 adds the vLLM serving integration. They are separate from this PR's local kernel/adapter validation.

Usage

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/linear_attention_state_int8_block32_dynamic").quantize
model = mtq.quantize(model, cfg)
with linear_attention_training_phase(model, prefill_lengths=[64], sequence_lengths=[128]):
    loss = training_step(model, batch)
loss.backward()

This expects an initialized Megatron model and a loss on the decode suffix. Training needs compatible vLLM and FLA packages. Validation/calibration loops must also establish the appropriate phase. A full-prefill pass does not measure decode state-QDQ quality. Compare against a converted model with the same serving policy and disabled state quantizers.

Testing

Current head: 52b75c6e59. GitHub CI for this amended head is pending.

Local validation on RTX A6000 during the cleanup:

• Focused CPU policy/checkpoint/state-QDQ and recipe checks: 47 passed.
• vLLM 0.30 GPU suite: 12 passed, 3 skipped. Coverage includes native bit-exact output/state comparisons, handoff gradients, mismatch rejection, GDN/KDA padding gradients and prefill-only behavior. The skips are optional private ReplaySSM cases.
• Modern Megatron-Core (ac100f773f9d) with vLLM 0.30: 3 passed, covering GDN/KDA state QAT, recompute, GDN sharded restore and missing-hook/raw-gate rejection. The NeMo 26.08 CI image cannot exercise the active serving-QAT adapter because it lacks the required Core hooks/KDA module.
• Optional ReplaySSM fork: 8 tests passed with the actual fork modules loaded locally into the vLLM 0.30 namespace. This qualifies the kernel adapter, not a complete installed fork or serving engine.
• After the final removal of redundant FP32 casts: 7 vLLM GPU training tests passed, plus 32 exact ReplaySSM output/state/scale/entry/gradient comparisons and 2 direct GDN/KDA QAT output/state/gradient comparisons against the pre-cleanup implementation.
• Changed-file pre-commit hooks and git diff --check: passed.

Tests use K=32 to reduce compilation, retain the native serving comparisons, and no longer include the K=192 regression. The GPU workflow timeout remains unchanged at 15 minutes. Local runs used warm caches and do not establish cold CI compilation time.

Public profiles now target vLLM 0.30.0 only; vLLM 0.20 state-training tests skip. The runtime hardware guard remains for custom/fork environments using sm_120+ with Triton < 3.7, where the observed CUDA misaligned-address fault has not been isolated to a specific kernel.

Bit-exact qualification covers the tested vLLM 0.30 kernel profiles and shapes only. This does not establish full-engine scheduling parity, model-quality recovery or training speed. The private ReplaySSM profile remains outside public CI coverage.

Before your PR is "Ready for review"

• Backward compatibility: checkpoints missing newly added disabled state-quantizer handles remain loadable, with or without saved linear-attention policy metadata. Enabled W quantization is rejected with a clear error. State QAT requires an explicit serving policy and phase.
• Dependencies/copying: native vLLM/ReplaySSM kernels are imported. FLA supplies training kernels; the vLLM GPU session installs fla-core==0.5.1 without replacing vLLM's TileLang/TVM pins.
• Tests: focused CPU policy/QDQ/checkpoint tests, native serving comparisons, and Megatron training/checkpoint restore. Compilation remains in setup fixtures. Core builds without the native gate/packed-forward hooks run a restore/rejection/disabled-baseline check and skip the unsupported full QAT test.
• Changelog: updated for state QAT.
• Review: human CODEOWNER approval remains required.

Additional Information

Step 2/6. Known limits include per-layer boundary synchronization, an additional FLA prefix forward for gradients, per-token native cache allocations, and explicit valid lengths for padded-only packed metadata. The source size is intentionally deferred for this review.

Summary by CodeRabbit

• New Features
  • Added experimental recurrent-state quantization for GDN and KDA, with vLLM serving support for prefill and decode.
  • Added dynamic INT8 state recipes, including block-32 and ReplaySSM/Hadamard options, plus a Kimi-K3 serving profile.
  • Added configurable linear-attention policies and training-phase handling for sequence lengths and state handoff.
• Documentation
  • Updated recipe guidance with serving requirements, supported profiles, and quantization behavior.
  • Clarified optional serving-kernel requirements and compatibility limitations.

[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]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: d9fff262-effd-42d4-b57f-51349d5837b3

---

---

---

---

---

---

---

---

---

📥 Commits

Reviewing files that changed from the base of the PR and between 2c90acd and fec406f.

---

---

---

---

---

---

---

---

---

📒 Files selected for processing (28)

• CHANGELOG.rst
• examples/megatron_bridge/quantize.py
• modelopt/torch/kernels/quantization/linear_attention/__init__.py
• modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py
• modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py
• modelopt/torch/kernels/quantization/linear_attention/serving/forward.py
• modelopt/torch/quantization/conversion.py
• modelopt/torch/quantization/linear_attention/_vllm_autograd.py
• modelopt/torch/quantization/linear_attention/config.py
• modelopt/torch/quantization/linear_attention/decode.py
• modelopt/torch/quantization/linear_attention/gdn.py
• modelopt/torch/quantization/linear_attention/kda.py
• modelopt/torch/quantization/linear_attention/training.py
• modelopt/torch/quantization/linear_attention/utils.py
• modelopt/torch/quantization/plugins/gdn.py
• modelopt/torch/quantization/plugins/kda.py
• modelopt/torch/quantization/plugins/linear_attention.py
• modelopt/torch/quantization/plugins/megatron.py
• modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml
• modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic_kimi_k3.yaml
• modelopt_recipes/ptq.md
• tests/examples/megatron_bridge/test_quantize_export.py
• tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py
• tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py
• tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py
• tests/gpu_vllm/torch/quantization/test_linear_attention_training.py
• tests/unit/torch/quantization/plugins/test_gdn.py
• tests/unit/torch/quantization/test_linear_attention_decode.py

---

---

---

---

---

---

---

---

---

🚧 Files skipped from review as they are similar to previous changes (3)

• modelopt/torch/kernels/quantization/linear_attention/serving/init.py
• CHANGELOG.rst
• modelopt/torch/kernels/quantization/linear_attention/init.py

---

---

---

---

---

---

---

---

---

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

---

---

---

---

---

---

---

---

---

📝 Walkthrough

📝 Walkthrough

📝 Walkthrough

📝 Walkthrough

📝 Walkthrough

📝 Walkthrough

📝 Walkthrough

📝 Walkthrough

📝 Walkthrough

Walkthrough

The change adds serving-aligned GDN and KDA recurrent-state QAT, vLLM forwards, differentiable training and decode paths, policy and checkpoint support, Megatron integration, and dynamic INT8 PTQ recipes. It also updates related tests, documentation, and changelog entries.

Changes

Serving-aligned linear-attention state QAT

Layer / File(s)	Summary
Policy configuration and persistence modelopt/torch/quantization/config.py, modelopt/torch/quantization/linear_attention/config.py, modelopt/torch/quantization/plugins/linear_attention.py, modelopt/torch/quantization/conversion.py, modelopt/torch/quantization/model_quant.py, modelopt/torch/quantization/plugins/__init__.py, modelopt/torch/quantization/linear_attention/__init__.py, modelopt/torch/opt/plugins/mcore_dist_checkpointing.py	QuantizeConfig accepts linear-attention policies. Conversion and restoration apply and validate policies, and checkpoint metadata stores or removes resolved policy state.
Serving kernels and differentiable adapters modelopt/torch/kernels/quantization/linear_attention/serving/*, modelopt/torch/kernels/quantization/linear_attention/__init__.py, modelopt/torch/quantization/linear_attention/_vllm_autograd.py	vLLM compatibility helpers select kernels and state layouts. Serving forwards implement prefill and decode; the differentiable adapters use native forward results with gradient paths for training. ReplaySSM helpers manage checkpoints and recurrent-kernel calls.
Recurrent state and decode modelopt/torch/quantization/linear_attention/decode.py, modelopt/torch/quantization/linear_attention/utils.py, tests/unit/torch/quantization/test_linear_attention_decode.py, tests/unit/torch/quantization/test_linear_attention_hadamard.py, tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py	Recurrent decoding adds resumable carries, replay entries, state encoding, and checkpoint refresh. State QDQ supports FP8 and INT8 formats. Unit and GPU tests cover quantization, configuration checks, Hadamard behavior, ReplaySSM output, carry state, and gradients.
Training phases and GDN/KDA integration modelopt/torch/quantization/linear_attention/{training,gdn,kda}.py, modelopt/torch/quantization/plugins/{gdn,kda,megatron}.py, tests/unit/torch/quantization/plugins/test_gdn.py, tests/gpu_megatron/torch/quantization/plugins/*, tests/gpu_vllm/torch/quantization/test_linear_attention_training.py	Training phases provide per-sequence lengths and packed-boundary metadata. GDN and KDA adapters route supported calls through the shared prefill/decode path. Megatron integration handles serving inputs, selective recomputation, and execution validation. Tests cover policy restoration, supported modes, outputs, state, and gradients.
PTQ recipes and execution constraints modelopt_recipes/configs/ptq/units/*, modelopt_recipes/general/ptq/linear_attention_state_int8*.yaml, modelopt_recipes/ptq.md, examples/megatron_bridge/quantize.py, tests/examples/megatron_bridge/test_quantize_export.py, CHANGELOG.rst	Recipes add dynamic INT8 state configurations, including block-32 and Kimi-K3 profiles. The Megatron example rejects calibration recipes that include linear-attention policies and skips generation when state quantization is enabled. Documentation and the changelog describe the profiles and requirements.

Priority: ➖ Normal

Estimated code review effort: 5 (Critical) | ~120 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant TrainingPhase as linear_attention_training_phase
  participant Megatron as Megatron linear-attention adapter
  participant QAT as gdn_state_qat or kda_state_qat
  participant Forward as _prefill_decode_forward
  TrainingPhase->>Megatron: set prefill and sequence lengths
  Megatron->>QAT: route supported GDN or KDA call
  QAT->>Forward: pass prepared inputs and state options
  Forward-->>QAT: return output and final state

Loading

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

Merge Risk: ⚪ Minimal · up to fec40

This change adds GDN/KDA recurrent-state QAT with serving-aligned forwards and INT8 state recipes. The concerns raised earlier about checkpoint restore and configuration handling have been addressed, and no outstanding defect is established at this head. The change appears ready to merge, subject to the runtime qualification limits the PR already documents.

Pre-merge checks | 5 | 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 168 functions across 34 files. (4 skipped…	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		PASS. The authoritative PR diff adds no listed security anti-pattern in modelopt or examples Python code. It adds no torch.load(..., weights_only=False), numpy.load/np.load(..., allow_pickle=True), ha…
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly identifies the main change: Torch GDN/KDA decode QAT with INT8 recurrent-state support.

---

Full details: Docstring Coverage

Explanation

Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 168 functions across 34 files. (4 skipped: 4 unsupported.)

---

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 3

📝 Generate docstrings 💡

• Commit to this branch
• Create a new PR

---

---

---

---

---

---

---

---

---

⚔️ Resolve merge conflicts 💡

• Resolve merge conflict in branch kaix/linear-attention-decode-first

---

---

---

---

---

---

---

---

---

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

---

---

---

---

---

---

---

---

---

---

🛠️ Fix failing CI checks 💡

• Commit to this branch
• Create a new PR

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

---

• 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 44.18023% with 446 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.32%. Comparing base (4c1f813) to head (52b75c6).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
...lopt/torch/quantization/linear_attention/decode.py	25.47%	79 Missing ⚠️
modelopt/torch/quantization/plugins/megatron.py	46.80%	75 Missing ⚠️
...pt/torch/quantization/linear_attention/training.py	45.96%	67 Missing ⚠️
...s/quantization/linear_attention/serving/forward.py	0.00%	65 Missing ⚠️
...ls/quantization/linear_attention/serving/replay.py	0.00%	55 Missing ⚠️
...ch/quantization/linear_attention/_vllm_autograd.py	15.62%	27 Missing ⚠️
...odelopt/torch/quantization/linear_attention/kda.py	24.24%	25 Missing ⚠️
...elopt/torch/quantization/linear_attention/utils.py	59.25%	22 Missing ⚠️
modelopt/torch/quantization/plugins/kda.py	38.88%	11 Missing ⚠️
...odelopt/torch/quantization/linear_attention/gdn.py	33.33%	10 Missing ⚠️
... and 3 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2519      +/-   ##
==========================================
+ Coverage   71.76%   76.32%   +4.55%     
==========================================
  Files         644      658      +14     
  Lines       71477    72435     +958     
==========================================
+ Hits        51298    55287    +3989     
+ Misses      20179    17148    -3031     

Flag	Coverage Δ
examples-diffusers	20.24% <21.02%> (-0.01%)	⬇️
examples-gpt-oss	13.51% <17.89%> (+0.03%)	⬆️
examples-hf_ptq	23.30% <19.52%> (-0.09%)	⬇️
examples-llm_distill	13.59% <17.89%> (+0.03%)	⬆️
examples-llm_eval	17.29% <19.52%> (+0.01%)	⬆️
examples-llm_qat	17.60% <21.02%> (+0.01%)	⬆️
examples-llm_sparsity	15.83% <17.89%> (+0.01%)	⬆️
examples-megatron_bridge	26.51% <32.66%> (-0.11%)	⬇️
examples-specdec_bench	13.29% <17.89%> (+0.04%)	⬆️
examples-speculative_decoding	17.58% <19.52%> (-0.07%)	⬇️
examples-torch_onnx	21.60% <19.52%> (-0.04%)	⬇️
examples-torch_trt	15.28% <19.52%> (+0.03%)	⬆️
examples-vllm_serve	13.96% <17.89%> (+0.03%)	⬆️
gpu	57.03% <32.16%> (+23.38%)	⬆️
regression	15.12% <17.89%> (+0.01%)	⬆️
unit	59.61% <36.17%> (-0.26%)	⬇️

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.