Add KDA prefill numerical emulation for QAT

[kaix-nv]

Superseded: KDA prefill is now included in #2503, after the shared GDN/KDA decode infrastructure in #2519. This closed PR and its branch retain the original review and qualification history.

---

What does this PR do?

Type of change: New feature, new example, new tests, documentation.

Adds experimental training-time numerical emulation to FLA 0.5.1 KimiDeltaAttention. ModelOpt conversion creates disabled-by-default KDA state/W and eight prefill operand-site controls. Enabling a policy routes the existing layer through differentiable FP32 materialized prefill with dynamic FP8 or NVFP4 QDQ; disabling all controls preserves the original FLA forward. This is stacked on #2503.

KDA uses per-key-channel retention. Causal pairwise exponentials are computed from prefix differences, avoiding an overflowing inverse-decay factorization. The adapter retains the model's projections, short convolution, gate formulas, and normalization, and preserves gradients and saved numerical policies. Invocation-local dispatch also works for copied modules without changing global FLA functions.

A pinned public checkpoint/WikiText runner records matched pre/post-training NLL, token hashes, source hashes, optimizer updates, and time/memory. The exact solve remains the default. Approximate inverse and decode/replay training are subsequent milestones.

Usage

import modelopt.torch.quantization as mtq

model = mtq.quantize(model, {
    "quant_cfg": [
        {"quantizer_name": "*", "enable": False},
        {"quantizer_name": "*linear_attn_sites.*",
         "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)}},
        {"quantizer_name": "*kda_w_quantizer",
         "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)}},
    ],
    "algorithm": None,
    "linear_attention": [{
        "module_name": "model.layers.*.self_attn",
        "cfg": {"backend": "matmul"},
    }],
})

Install the example's pinned FLA dependencies only when using this adapter. Core ModelOpt imports do not require FLA.

Testing

Revision: e3987d01bb2b54188bb18165b478972a101ac54e (signed). Remote staged source hashes matched the implementation.

• 150 CPU tests passed: existing policy/conversion/GDN tests plus independent KDA recurrence/operand tests.
• 64 H100 tests passed: 34 GDN kernels, 9 KDA kernels, 5 KDA layer integrations, and 16 Megatron TP=1/2/checkpoint/optimizer/CP-rejection cases. The shared autocast guard is covered. Local SM86 KDA kernel/layer tests also passed.
• Pre-commit hooks passed; complete Sphinx build passed with warnings as errors and FLA imports deliberately blocked.
• Three full-checkpoint training smoke runs passed on H100: pinned arcee-ai/AFM-4.5B-Base-KDA-Only at 01ad2e06ee4f1214193c17b69e09105a9b257e80, WikiText-2 raw revision b08601e04326c79dfdd32d625aee71d232d685c3. All checkpoint keys loaded, all trainable attention gradients were finite, and an optimizer update changed query weights. Runs used identical token hashes/seed, one 128-token training step and four validation blocks (512 predicted tokens), attention-only FP32 AdamW under BF16 autocast with activation checkpointing.

Policy	Validation PPL before	Validation PPL after
Fused FLA	10.1556	10.1665
Exact materialized	10.1603	10.1537
FP8 prefill operands	10.1257	10.1105

These are one-step integration results, not evidence of quality recovery. The model payload has 36 KDA layers; configuration and weight index were checked rather than relying on the model-card layer count. Original FP32 gate parameters are preserved during loading.

H100 environment: NeMo 26.08, Torch 2.13.0a0+8145d630e8.nv26.06, CUDA 13.3, Triton 3.7, FLA/core 0.5.1, Transformers 4.57.1, TileLang 0.1.8, TVM-FFI 0.1.9, Megatron-Core 0.19.1, Transformer Engine 2.18.0. Numerical qualification excludes TF32.

Measured KDA training overhead

One H100, shape [1,1024,4,64], BF16 Q/K/V, FP32 working math, TF32 off. Three warmups and 20 interleaved trials; median times and paired wall ratios. Extra memory is peak tensor allocation above the pre-forward baseline. Compilation is excluded; samples and source hashes retained.

Policy	Forward ms	Backward ms	Wall ms	Wall / FLA	Extra MiB
Exact fused FLA	1.026	1.421	2.490	1.00	15.25
Exact matmul	8.716	13.864	22.580	9.17	416.03
FP8 operands	37.476	15.628	53.184	21.68	433.56
NVFP4 operands	33.298	15.491	48.906	19.86	424.16
FP16 accumulator, reduction block 16	29.519	43.503	73.243	29.58	416.03
FP8 operands + chunk state	40.346	15.634	56.139	22.95	433.56

The per-channel interaction tensors account for substantial materialization cost. This is an explicit numerical research path, with no speed claim.

python -m pytest -q -x \
  tests/gpu/torch/kernels/quantization/linear_attention/test_matmul_gdn.py \
  tests/gpu/torch/kernels/quantization/linear_attention/test_matmul_kda.py \
  tests/gpu/torch/quantization/plugins/test_kda.py \
  tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py
python examples/llm_qat/linear_attention/benchmark_prefill.py \
  --attention kda --batch 1 --length 1024 --heads 4 --dim 64 --repeats 20 \
  --output prefill-kda-overhead.json

The model-study command and pinned dependencies are in examples/llm_qat/linear_attention/README.md; run once with no quant config, once with kda_exact_matmul.json, and once with kda_prefill_fp8.json.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ New controls default off; existing GDN contracts retained.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ✅ Uses existing optional FLA APIs; example dependencies are permissively licensed and pinned where needed. No new mandatory core dependency.
• Did you write any new necessary tests?: ✅ Independent FP64 recurrence, operand/STE oracles, native FLA gradients, and model-layer restore/optimizer tests.
• Did you update Changelog?: ✅
• Did you get Claude approval on this PR?: ❌ Pending; keep draft.

Additional Information

The materialized path is a numerical research implementation. It does not claim native low-precision execution, serving acceleration, or model-quality recovery. Qualification uses FLA's general gate path; the optional safe_gate TensorCore rounding is not emulated. CP, recurrent inference, and FLA intermediate/recompute flags are rejected. FLA selects recurrent inference for evaluation sequences of at most 64 tokens; these QAT prefill studies use longer sequences.

[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 86.31579% with 26 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.51%. Comparing base (86c4600) to head (e3987d0).

Files with missing lines	Patch %	Lines
modelopt/torch/quantization/plugins/kda.py	54.28%	16 Missing ⚠️
...odelopt/torch/quantization/linear_attention/kda.py	75.00%	7 Missing ⚠️
...t/torch/quantization/linear_attention/reference.py	93.61%	3 Missing ⚠️

Additional details and impacted files

@@                       Coverage Diff                        @@
##           kaix/linear-attention-qat-m2    #2506      +/-   ##
================================================================
+ Coverage                         70.49%   70.51%   +0.01%     
================================================================
  Files                               612      615       +3     
  Lines                             67875    68025     +150     
================================================================
+ Hits                              47850    47969     +119     
- Misses                            20025    20056      +31     

Flag	Coverage Δ
unit	57.81% <86.31%> (+0.04%)	⬆️

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.

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-09-23 01:07 UTC