[1/5] GDN state/W QAT foundation

[kaix-nv]

Linear-attention PR stack — 5 drafts

Order	PR	Depends on
1/5	#2497 GDN state/W QAT foundation	main
2/5	#2519 GDN/KDA decode QAT + INT8	#2497
3/5	#2541 vLLM GDN/KDA state-only fake quantization	#2519
4/5	#2503 GDN/KDA prefill GEMM quantization	#2541
5/5	#2507 Experimental GDN/KDA approximate inverse	#2503

All five drafts form a linear GitHub stack in the review order shown above. #2541 applies TensorQuantizer before native vLLM prefill/decode calls.

A separate vLLM prefill-GEMM PR will wait for an optimized fused kernel. #2506 and #2509 are superseded and closed.

What does this PR do?

Type of change: new feature

GatedDeltaNet training keeps recurrent states inside a chunked kernel, so projection quantizers cannot emulate rounding at state boundaries. This draft adds dynamic FP8 E4M3 fake QDQ to the recurrent state and WY-transformed W activations, with identity straight-through gradients for QAT/QAD.

Both sites use the standard quant_cfg interface and start disabled. State QDQ uses 64-token chunks and one scale per full-key by 64-value-column tile; W grouping is applied by TensorQuantizer. Quantizer settings use normal ModelOpt checkpoint state. There is no QuantizeConfig.linear_attention field in this PR; #2519 introduces execution policies for decode and ReplaySSM, and later PRs extend them for prefill and approximate inverse. Configurations or checkpoints from earlier experimental drafts that use those execution policies require #2519.

The Megatron adapter supports the direct-forward and older split-forward call layouts, restores the original kernel when disabled, and removes temporary quantizer attributes on export. Independent recurrent/chunk numerical references live under tests/_test_utils/torch/quantization/; shared runtime capability checks live in linear_attention/utils.py.

The fused path requires fla-core==0.5.1 and chunk size 64. State FP8 emulation requires SM89 or newer. The Hopper path has additional dtype/TileLang restrictions enforced before launch. This PR simulates numerical error; it does not add compressed state storage or faster inference.

Usage

import modelopt.torch.quantization as mtq

model = mtq.quantize(model, {
    "quant_cfg": [
        {"quantizer_name": "*", "enable": False},
        {"quantizer_name": "*gdn_state_quantizer",
         "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1)}},
        {"quantizer_name": "*gdn_w_quantizer",
         "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)}},
    ],
    "algorithm": None,
})
# Continue with the framework's normal forward/backward/optimizer steps.

Dynamic scales require no calibration.

Testing

Current amended source: 122 passed, 12 skipped on RTX A6000 (SM86), Python 3.12.8, Torch 2.9.1+cu128, Triton 3.5.1, and fla-core 0.5.1. Coverage includes conversion/refinement, state/W checkpoint round trips, legacy disabled-handle restoration, independent reference numerics, W grouping, fused output/gradient comparisons, existing CPU quantization, and recipe regressions. Eleven skips require native state E4M3 conversion; one is Hopper-specific.

PYTHONPATH=. python -m pytest -q \
  tests/unit/torch/quantization/plugins/test_gated_delta_net.py \
  tests/unit/torch/quantization/test_linear_attention_reference.py \
  tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py \
  tests/unit/torch/quantization/test_quantize_cpu.py \
  tests/unit/recipe/test_presets.py

CI collection fix: the streaming-dataset test now skips when optional httpx is unavailable. Reproduced the missing-dependency collection failure, verified it becomes a skip, and ran all 18 streaming-dataset tests successfully with dependencies installed.

Pre-commit and git diff --check passed. Megatron-Core/Transformer Engine are absent locally, so the updated Megatron test was checked statically but not rerun. Earlier Hopper/Megatron qualification predates this amendment; it is not current-head runtime evidence. Model-quality/QAT-recovery evaluation, pipeline parallelism, and checkpoint resharding remain outside this validation.

Before your PR is "Ready for review"

Contributor and security guidance reviewed. Commits are signed and signed off.

• Is this change backward compatible?: ✅ Disabled-by-default quantizers, standard-recipe exclusions, and legacy-checkpoint coverage; enabled experimental configurations have explicit capability restrictions.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ❌ Internal third-party approval tracking still needs confirmation. Upstream attribution, MIT/Apache headers, LICENSE notice, and license-hook exclusions are included. FLA/TileLang and TVM-FFI license files were reviewed.
• Did you write any new necessary tests?: ✅ Numerical, gradient, conversion/checkpoint, and real framework tests.
• Did you update Changelog?: ✅ Experimental quantization feature entry.
• Did you get Claude approval on this PR?: ❌ Not requested; this remains a draft.

Additional Information

Related: #2455. This is one draft for the first integration slice and does not assume #2455 has merged. Later milestones will extend the numerical boundaries after choosing their approximation contracts.

Summary by CodeRabbit

• New Features

  • Added experimental dynamic FP8 fake quantization for GatedDeltaNet recurrent states and WY activations during training.
  • Added PTQ configuration options for enabling state and WY quantization independently; both are disabled by default.
  • Added support for applying these quantization options to Megatron GatedDeltaNet models.

• Documentation

  • Documented hardware and integration requirements, including SM89+ for state quantization and a chunk size of 64.

• Bug Fixes

  • Improved quantized-model restoration when GatedDeltaNet quantizers are disabled.

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

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

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: 1b714216-45a4-4c38-a128-322e8ee8446b

📥 Commits

Reviewing files that changed from the base of the PR and between ed7e879 and 6686c9b.

📒 Files selected for processing (26)

• .pre-commit-config.yaml
• CHANGELOG.rst
• LICENSE
• docs/source/_templates/autosummary/module.rst
• modelopt/torch/kernels/quantization/linear_attention/__init__.py
• modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py
• modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py
• modelopt/torch/quantization/conversion.py
• modelopt/torch/quantization/linear_attention/__init__.py
• modelopt/torch/quantization/linear_attention/utils.py
• modelopt/torch/quantization/model_quant.py
• modelopt/torch/quantization/plugins/__init__.py
• modelopt/torch/quantization/plugins/gated_delta_net.py
• modelopt/torch/quantization/plugins/megatron.py
• modelopt_recipes/configs/ptq/units/README.md
• modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml
• modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml
• modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml
• noxfile.py
• pyproject.toml
• tests/_test_utils/torch/quantization/linear_attention_reference.py
• tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py
• tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py
• tests/unit/torch/quantization/plugins/test_gated_delta_net.py
• tests/unit/torch/quantization/test_linear_attention_reference.py
• tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py

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

---

📝 Walkthrough

Walkthrough

This change adds experimental dynamic FP8 fake quantization for GatedDeltaNet recurrent states and WY activations. It adds chunked kernels, model and Megatron integration, PTQ configurations, and reference and GPU tests.

Changes

GatedDeltaNet Quantization

Layer / File(s)	Summary
Chunked state kernels modelopt/torch/kernels/quantization/linear_attention/*, .pre-commit-config.yaml, docs/source/_templates/autosummary/module.rst, LICENSE, noxfile.py, pyproject.toml	Adds Triton forward and backward kernels with optional dynamic FP8 state QDQ, state-tile selection, and launch wrappers. Adds package documentation and supporting tooling and GPU test environment configuration.
Fused chunk GatedDeltaNet API modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py	Adds the chunked GatedDeltaNet API and autograd implementation. The path supports optional WY quantization and state QDQ, with validation for supported chunk size, hardware, and context-parallel settings.
Quantizer configuration and model lifecycle modelopt/torch/quantization/linear_attention/*, modelopt/torch/quantization/conversion.py, modelopt/torch/quantization/model_quant.py, modelopt/torch/quantization/plugins/{__init__,gated_delta_net}.py, modelopt_recipes/configs/ptq/units/*, CHANGELOG.rst, tests/unit/torch/quantization/plugins/test_gated_delta_net.py	Adds state and WY quantizers, validates quantizer settings during conversion and restoration, and adds dynamic FP8 PTQ recipes. Unit tests cover configuration, dispatch, and restore behavior.
Megatron GatedDeltaNet wiring modelopt/torch/quantization/plugins/megatron.py, tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py	Registers a quantized Megatron wrapper when GatedDeltaNet is available. Tests cover distributed execution, checkpoint restore, and context-parallel rejection.
Reference and kernel validation tests/_test_utils/torch/quantization/linear_attention_reference.py, tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py, tests/unit/torch/quantization/test_linear_attention_reference.py	Adds differentiable reference functions and tests for outputs and gradients across quantization modes, packed sequences, state layouts, and kernel options.

Optional Test Dependency

Layer / File(s)	Summary
Skip without optional HTTP client tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py	Uses pytest.importorskip so the module skips when httpx is unavailable.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~60 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant GatedDeltaNetStateQuantMixin
  participant chunk_gated_delta_rule
  participant chunk_gated_delta_rule_fwd_h
  participant chunk_gated_delta_rule_bwd_dhu
  GatedDeltaNetStateQuantMixin->>chunk_gated_delta_rule: invoke quantized chunk operation
  chunk_gated_delta_rule->>chunk_gated_delta_rule_fwd_h: launch forward state computation
  chunk_gated_delta_rule->>chunk_gated_delta_rule_bwd_dhu: compute backward state and value gradients

Loading

Merge Risk: ⚪ Minimal · up to 6686c

This change adds experimental FP8 fake quantization for GatedDeltaNet recurrent states and WY activations. It is disabled by default, so existing quantization recipes and restored checkpoints keep their current behavior unless users explicitly enable the new quantizers. No concrete merge-blocking defect remains. Hardware and version requirements are enforced by explicit errors raised before any kernel launches.

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 22.89% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 83 functions across 17 files. (9 skipped:…	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly identifies the main change: foundational GatedDeltaNet state and W quantization-aware training support. It is concise and consistent with the pull request objectives.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	No configured security anti-pattern was introduced. The changed modelopt/examples Python additions contain no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded `tr…

Full details: Docstring Coverage

Explanation

Docstring coverage is 22.89% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 83 functions across 17 files. (9 skipped: 9 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

---

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2497/
Built to branch gh-pages at 2026-09-25 02:01 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 26.44231% with 612 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.58%. Comparing base (051d6ad) to head (6686c9b).
⚠️ Report is 17 commits behind head on main.

Files with missing lines	Patch %	Lines
...quantization/linear_attention/fla_chunk_delta_h.py	8.00%	540 Missing ⚠️
...ion/linear_attention/fla_chunk_gated_delta_rule.py	63.57%	51 Missing ⚠️
modelopt/torch/quantization/plugins/megatron.py	57.14%	21 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2497      +/-   ##
==========================================
- Coverage   71.14%   70.58%   -0.56%     
==========================================
  Files         603      611       +8     
  Lines       66739    71409    +4670     
==========================================
+ Hits        47482    50405    +2923     
- Misses      19257    21004    +1747     

Flag	Coverage Δ
examples-diffusers	21.20% <3.30%> (-0.13%)	⬇️
examples-gpt-oss	13.34% <2.57%> (-0.07%)	⬇️
examples-hf_ptq	22.30% <3.06%> (-0.14%)	⬇️
examples-llm_distill	13.40% <2.57%> (-0.08%)	⬇️
examples-llm_eval	17.26% <3.06%> (-0.08%)	⬇️
examples-llm_qat	17.52% <3.30%> (-0.12%)	⬇️
examples-llm_sparsity	15.82% <2.57%> (-0.12%)	⬇️
examples-megatron_bridge	25.91% <7.84%> (-0.40%)	⬇️
examples-specdec_bench	13.10% <2.57%> (-0.07%)	⬇️
examples-speculative_decoding	17.58% <3.06%> (-0.24%)	⬇️
examples-torch_onnx	21.66% <3.06%> (-0.21%)	⬇️
examples-torch_trt	15.12% <3.06%> (-0.10%)	⬇️
examples-vllm_serve	13.72% <2.57%> (-0.08%)	⬇️
gpu	32.92% <25.98%> (-0.50%)	⬇️
regression	14.93% <2.57%> (-0.17%)	⬇️
unit	58.14% <6.60%> (-0.08%)	⬇️

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.