[3/6] Megatron Bridge linear attention QAT and QAD example

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

Add a Megatron Bridge example for recurrent-state QAT and QAD. The training entry point selects a native chunked prefix and recurrent suffix, applies loss to the suffix, and keeps the phase context active through backward. Bridge owns optimization, distributed scheduling, and checkpoints. A frozen unquantized teacher enables QAD.

The example contains the training script, dependency files, launcher, two minimal training tests, and concise usage/alignment documentation. Raw experiment records and development links are kept outside the release PR. State formats, execution policies, and kernels are supplied by the parent PR.

Usage

bash examples/llm_qat/linear_attention/with_vllm_defaults.sh \
  torchrun --standalone --nproc-per-node=1 examples/llm_qat/linear_attention/train.py \
  --model /path/to/local-model \
  --train-data /path/to/tokenized/train_text_document \
  --output /path/to/megatron-qat-checkpoint \
  --recipe general/ptq/linear_attention_state_int8_block32_dynamic \
  --train-steps 1 --length 128 --prefill-tokens 64

Add --teacher-model /path/to/unquantized-model for QAD. The ordinary INT8 recipe uses public vLLM kernels. INT8 + Hadamard and replay require a compatible native ReplaySSM fork. A KDA trainer additionally requires a compatible Bridge provider. TP/PP/EP options follow Bridge; context parallelism stays at one, and each local pipeline chunk must contain linear attention.

Testing

Rebased onto #2519 at 03f6076dd8; this PR's head is 57317404d0.

• Scoped pre-commit hooks and git diff --check passed. git range-diff confirms all nine example commits replayed without patch changes. The parent-relative diff is 8 files, 987 additions, with no runtime/kernel or workflow changes.
• GPU workflow revalidation was attempted on one RTX A6000 but is blocked by the local Megatron dependency setup. The previous Bridge/Core build lacks the _forward_compute hook now required by [2/6] Torch GDN/KDA decode QAT with INT8 recurrent state #2519. Core ac100f773f9d provides the required hooks, but collection fails with the local nvidia_resiliency_ext package; compatible published wheels require a newer system runtime. No passing QAT/QAD workflow result is claimed for this rebased head.
• On the previous parent 140735b962, pytest tests/examples/megatron_bridge/test_linear_attention.py passed 2 tests in 57.84 seconds. Those tests cover QAT/QAD student updates, a frozen QAD teacher, checkpoint save/resume, HF export/load, and Megatron reimport with restored state-quantization metadata and matching logits/gradients.

Previous workflow checks exercised the real GDN Bridge training entry point with tiny random weights and mock data, including Hadamard token mode and replay QAT/QAD. KDA checks exercised Megatron layers rather than a Bridge KDA trainer. These do not establish pretrained-model quality recovery, full serving-engine equivalence, or performance.

Before your PR is "Ready for review"

• Is this change backward compatible?: Additive example. The native state-QAT runtime and checkpoint policy are defined by the parent PR.
• If you copied code or added a dependency, did you follow CONTRIBUTING.md?: The optional public vLLM dependency and its license are listed separately. INT8/Hadamard and replay require the documented native fork.
• Did you write any new necessary tests?: One QAT and one QAD training-step test with shared compilation setup. This cleanup adds no tests.
• Did you update Changelog?: N/A; the runtime feature entry belongs to the parent PR.
• Did you get Claude approval on this PR?: Pending review of the updated example.

Additional Information

Step 3/6. Merge #2519 first, then this example, followed by #2541. This example trains state quantization; prefill GEMM quantization and approximate inverse remain follow-up work. Rebased onto #2519 at 03f6076dd8; the review diff contains only this example (8 files, 987 added lines).

[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

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.11%. Comparing base (03f6076) to head (5731740).

Additional details and impacted files

@@                         Coverage Diff                         @@
##           kaix/linear-attention-decode-first    #2657   +/-   ##
===================================================================
  Coverage                               78.11%   78.11%           
===================================================================
  Files                                     659      659           
  Lines                                   72634    72634           
===================================================================
  Hits                                    56740    56740           
  Misses                                  15894    15894           

Flag	Coverage Δ
unit	59.50% <ø> (+<0.01%)	⬆️

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.