[1/4] GDN state/W QAT foundation

[kaix-nv]

What does this PR do?

Type of change: new feature

ModelOpt needs to emulate recurrent-state and WY-activation rounding inside GatedDeltaNet training, where projection quantization does not expose those numerical boundaries. This draft adds dynamic FP8 E4M3 fake QDQ at the state and W sites, with identity straight-through gradients, plus differentiable GDN/KDA references for numerical validation.

This is the M0/M1 implementation slice: GDN state/W QDQ in Megatron's chunked training path. KDA has a recurrent reference only. Additional prefill operand sites, approximate inverse, KDA training/decay approximation, and decode/SSM replay remain subsequent work. The triangular solve stays exact. No compressed state storage, native FP8 matmul speedup, or model-quality recovery is claimed.

The implementation adapts #2455 at 13c7e2456f2e9d079c9ef822742eeaa634353802 with these contracts:

• fla-core==0.5.1, chunk size 64, dynamic E4M3, and identity STE. Unsupported formats, clipping-aware backward, and context parallelism are rejected. Forward's quantized W is saved for backward, avoiding a second quantizer invocation at the cost of one saved W activation.
• A linear_attention execution policy with complete last-match precedence, capability checks, and persistence of resolved state-scale grouping through conversion, calibration, and checkpoints.
• Megatron adapters for both direct-forward and older split-forward layouts. Megatron 0.19.1 calls the kernel directly; wrapping only the old split-forward method silently bypassed QDQ. Integration tests verify that enabling QDQ changes the GDN branch output.
• On Hopper with Triton >=3.4, BF16 Q/K/V with tilelang==0.1.8 and apache-tvm-ffi==0.1.9. FLA disables its gated Triton backward there; its TileLang fallback requires equal head counts. Q/K expansion outside custom autograd preserves grouped-head gradients, with temporary activation-storage overhead. FP32 is rejected before launch.
• Exact recurrent and chunk references, optional-dependency-safe imports, original-kernel dispatch when disabled, and standard-recipe exclusions for the new numerical handles.

Usage

After constructing a supported Megatron GatedDeltaNet model:

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,
    "linear_attention": [{
        "module_name": "decoder.layers.*.self_attention",
        "cfg": {"chunk_size": 64, "state": {"block_v": 64}, "solve": {"method": "exact"}},
    }],
})
# Continue with the framework's normal forward/backward/optimizer steps.

Dynamic scales require no calibration. State QDQ requires SM89 or newer; W-only QDQ is tested on SM86. The selector must match actual GDN module names. See docs/linear_attention_qat.md for the numerical contract and recipe composition.

Testing

Validation uses the branch based on 051d6adb204f10cd3e78d0f824f31a5a01d54831.

• 89 CPU tests passed: exact references/gradcheck, GDN conversion and policy persistence, and existing CPU quantization tests.
• 10 recipe regression tests passed (tests/unit/recipe/test_presets.py).
• SM86 GPU: 17 passed, 12 skipped with Python 3.12.8, Torch 2.9.1, Triton 3.5.1, FLA 0.5.1. Eleven skips need native state E4M3 conversion; one is Hopper-specific.
• H100 GPU: 20 passed, 9 skipped in NeMo 26.08 with Python 3.12.3, Torch 2.13.0a0 (NV 26.6), Triton 3.7.0, FLA 0.5.1, TileLang 0.1.8, TVM-FFI 0.1.9. Tests compare outputs, final states, and all input gradients against float32 references, including all state-scale tiles, grouped heads, packed sequences, fused gates, activation checkpointing, and state/W composition. Nine FP32 cases are outside the supported Hopper contract; explicit early rejection passes.
• Megatron: 8 passed in that H100 environment with Megatron-Core 0.19.1 and Transformer Engine 2.18.0. State-only, W-only, and combined QAT each pass at TP=1 and TP=2: enabled-site output sensitivity, disabled-path parity, distributed checkpoint restore with an edited execution policy, backward, and an optimizer step. Two additional cases verify context-parallel rejection. The adapter also supports the older split-forward layout, but this runtime receipt covers the direct-forward layout in 0.19.1.
• Full Sphinx HTML build with warnings treated as errors passed, with fla deliberately unavailable via an import blocker.
• Pre-commit on changed files and git diff --check: passed.

PYTHONPATH=. pytest -q tests/unit/torch/quantization/test_linear_attention_reference.py \
  tests/unit/torch/quantization/plugins/test_gated_delta_net.py \
  tests/unit/torch/quantization/test_quantize_cpu.py
PYTHONPATH=. pytest -q tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py
PYTHONPATH=. pytest -q tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py

Checkpoint qualification is limited to unchanged TP topology with PP=1. Pipeline parallelism and resharding across topology changes remain unqualified. Dynamic-batching inference and model-quality/QAT-recovery evaluation are outside this draft's validation.

Before your PR is "Ready for review"

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

• Is this change backward compatible?: ✅ Additive policy, disabled-path preservation, 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.

[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.}

[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-22 06:00 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 20.60041% with 767 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.50%. Comparing base (051d6ad) to head (8ea824b).
⚠️ Report is 7 commits behind head on main.

Files with missing lines	Patch %	Lines
...quantization/linear_attention/fla_chunk_delta_h.py	0.00%	587 Missing ⚠️
...ion/linear_attention/fla_chunk_gated_delta_rule.py	0.00%	140 Missing ⚠️
modelopt/torch/quantization/plugins/megatron.py	3.03%	32 Missing ⚠️
...t/torch/quantization/linear_attention/reference.py	96.07%	4 Missing ⚠️
...lopt/torch/quantization/plugins/gated_delta_net.py	92.10%	3 Missing ⚠️
modelopt/torch/quantization/conversion.py	97.05%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2497      +/-   ##
==========================================
- Coverage   71.14%   70.50%   -0.65%     
==========================================
  Files         603      610       +7     
  Lines       66739    68036    +1297     
==========================================
+ Hits        47482    47969     +487     
- Misses      19257    20067     +810     

Flag	Coverage Δ
unit	57.68% <20.60%> (-0.54%)	⬇️

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.