[3/6] Fused Triton GDN/KDA decode QAT

[kaix-nv]

Linear-attention PR stack — 6 PRs

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

All six PRs form native GitHub stack #2563 in the order shown above. #2541 applies TensorQuantizer before native vLLM prefill/decode calls. A separate vLLM prefill-GEMM PR waits for an optimized fused kernel. #2506 and #2509 are superseded and closed.

What does this PR do?

Type of change: New feature.

Add an optional fused Triton implementation of the GDN/KDA decode-QAT recurrence introduced by #2519. It fuses token/replay updates with FP8 or INT8 state QDQ and supplies a checkpointed custom backward, differentiable continuation carry, and INT8 Hadamard integration. Replay supports encode-once factors.

The kernel uses shared FP8 QDQ primitives, including software E4M3 conversion before SM89. Torch remains the default backend; the KDA example selects Triton. GPU tests compare outputs, reconstructed states, and gradients with the Torch implementation and exercise real KDA layer training/checkpoint restoration.

This is a review split of the existing implementation. The combined tree is identical to the previous #2519 head 30e0659ad478085c79ce2f54b6096d6066b1f814. It does not change the vLLM state-only integration.

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/gdn_state_int8_dynamic").quantize
recipe.linear_attention[0].cfg.decode.implementation = "triton"
mtq.quantize(model, recipe)
with linear_attention_training_phase(model, [64]):
    loss = model(input_ids=ids, labels=labels, use_cache=False).loss
    loss.backward()

The fused implementation supports first-order gradients, key dimensions up to 128, and value blocks 16/32/64/128. Use Torch for higher-order differentiation.

Testing

Compilation-fixture update: this branch is restacked on #2497's separate follow-up commit 95766de709ac. The changed test modules passed at #2497 (26 passed, 20 hardware skips in each cold/warm run) and at the #2507 stack tip (64 passed, 20 hardware skips on two RTX A6000 GPUs); intermediate PRs were not separately rerun. Functional calls created no new tracked kernel binaries and kept the default 120-second cap. All six source trees match the validated trees, and commit hooks passed. Native FP8-state/Hopper cases remain hardware-gated.

Earlier scope-specific validation follows.

103 passed, 11 skipped on RTX A6000 with Torch 2.9.1+cu128, Triton 3.5.1, fla-core and flash-linear-attention 0.5.1. Ran the complete decode-kernel suite, TestScaledE4M3, and real KDA training/checkpoint tests. Nine skips require native FP8 conversion on SM89+; two skip unsupported CPU dtypes. Pre-commit and git diff --check passed. Existing tolerances were unchanged.

The Torch-only parent independently passed 102 tests, including eight KDA GPU integration tests.

The restacked #2507 tip passed 180 CPU tests. The new Triton tip and all three descendant source trees are byte-for-byte identical to their pre-split counterparts.

Before your PR is "Ready for review"

• Is this change backward compatible?: Yes; extends [2/6] Torch GDN/KDA decode QAT with INT8 recurrent state #2519 with an optional backend. The example explicitly selects Triton; the config default remains Torch.
• Copied code/new dependency guidance: Existing provenance retained; no additional dependencies beyond [2/6] Torch GDN/KDA decode QAT with INT8 recurrent state #2519.
• Necessary tests added?: Yes; fused kernel and FP8 parity coverage, plus KDA integration tests selecting the fused backend.
• Changelog updated?: The decode-QAT feature entry is in [2/6] Torch GDN/KDA decode QAT with INT8 recurrent state #2519; this PR separates its implementation.
• Claude approval?: Pending; draft only.

Additional Information

Depends on #2519. The previously identified low-level Triton replay-anchor gradient issue remains open. This split does not fix it and the change remains a draft. Native FP8 parity on SM89+, distributed training, model-quality recovery, and serving performance require separate qualification.

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

[codecov]

Codecov Report

❌ Patch coverage is 1.60000% with 246 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (kaix/linear-attention-decode-first@7b5caf1). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
...ch/kernels/quantization/linear_attention/decode.py	0.00%	195 Missing ⚠️
...lopt/torch/quantization/linear_attention/decode.py	5.88%	32 Missing ⚠️
...opt/torch/kernels/quantization/common/fp8_quant.py	0.00%	19 Missing ⚠️

Additional details and impacted files

@@                          Coverage Diff                          @@
##             kaix/linear-attention-decode-first    #2562   +/-   ##
=====================================================================
  Coverage                                      ?   77.54%           
=====================================================================
  Files                                         ?      619           
  Lines                                         ?    68376           
  Branches                                      ?        0           
=====================================================================
  Hits                                          ?    53025           
  Misses                                        ?    15351           
  Partials                                      ?        0           

Flag	Coverage Δ
unit	57.60% <1.60%> (?)

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.