Add experimental GDN/KDA decode and replay QAT

CHANGELOG.rst

@@ -13,6 +13,8 @@ Changelog
 
 *Quantization*
 
+- Add experimental GDN/KDA decode-aware QAT with explicit prefix lengths, per-token state fake quantization, rounded log retention, and encoded-update replay. Select a saved decode policy and keep its phase context active through checkpointed backward.
+
 - Added opt-in differentiable Neumann solve approximation for GDN/KDA QAT, with an explicit saved degree and a CUDA FP32 candidate. Exact solve remains the default.
 
 - Add experimental KDA prefill QAT for FLA 0.5.1 layers with stable per-channel decay and independent operand/state fake quantization. Enable the explicit ``matmul`` backend; numerical policies and quantizers persist through ModelOpt save/restore.

docs/linear_attention_decode.md

@@ -0,0 +1,61 @@
+# Experimental GDN/KDA decode-aware QAT
+
+This extension models token-state writes, rounded KDA log retention, and encoded-update replay during training. Select the existing materialized prefill backend, an explicit decode policy, and a prefix length per sequence. The default configuration remains unchanged.
+
+```python
+import torch
+import modelopt.torch.quantization as mtq
+from modelopt.torch.quantization.linear_attention import linear_attention_training_phase
+
+mtq.quantize(model, {
+    "quant_cfg": [
+        {"quantizer_name": "*", "enable": False},
+        {"quantizer_name": "*kda_state_quantizer", "cfg": {
+            "num_bits": (4, 3), "type": "dynamic", "axis": (0, 1),
+        }},
+    ],
+    "algorithm": None,
+    "linear_attention": [{"module_name": "*", "cfg": {
+        "backend": "matmul",
+        "state": {"block_v": 64},
+        "decode": {
+            "mode": "replay", "implementation": "triton", "readout": "stored",
+            "replay": {"window": 8, "factor_qdq": True, "encoding": "once"},
+        },
+    }}],
+})
+# One sequence with 64 prefix tokens; all remaining tokens use the decode policy.
+# Keep phase metadata active during activation-checkpoint recomputation.
+with linear_attention_training_phase(model, [64]):
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        loss = model(input_ids=ids, labels=labels, use_cache=False).loss
+    loss.backward()
+```
+
+For GDN, use `*gdn_state_quantizer`. Policies persist through ModelOpt save/restore; per-batch phase lengths are runtime metadata and must be supplied for each workload. The context restores previous lengths on exit, supports nesting, and never patches process-global functions. Concurrent forwards on the same model instance with different phase contexts are unsupported.
+
+## Numerical contract
+
+For activated input tensors, a token computes `D = exp(g)`, `R = D * S`, `u = beta * (v - k^T R)`, and `W = R + k u^T`. GDN uses one log retention per head; KDA uses one per key channel. The output reads either the working state `W` or the stored state, according to `readout`. All operations use FP32 working values (FP64 is supported by the Torch reference).
+
+- **Token mode:** fake-quantize `W` at every state write when the module state quantizer is enabled. The next token consumes that stored value.
+- **Decay approximation:** an optional positive `decay_log_step` rounds log retention to the nearest grid point before exponentiation. Its derivative is an identity straight-through estimator (STE); no hidden clamping is applied. Prefix gates retain the prefill policy.
+- **Replay mode:** maintain an anchor and ordered encoded rank-one updates. Encode each key per head and each update per value block, with optional FP8 E4M3 QDQ. Compute a new residual from the reconstructed current trajectory and the encoded key. At `window` updates, encode the current state as a new anchor and clear the buffer. Anchor QDQ follows the state quantizer toggle; factor QDQ has its own toggle.
+- **Initial state:** `quantize_initial` applies once at the first nonempty decode handoff. Continuation and empty calls do not create extra writes. Prefix state QDQ is independently selected by `prefill_state_qdq`.
+
+State scale grouping is one dynamic scale for each `[Dk, block_v]` head tile. Update scales span one value block; key scales span the complete key vector. Zero tiles use scale one. Codec scales are detached, and rounded values have identity STE gradients. The software codec uses E4M3 nearest-even rounding, including subnormals. Key reductions use a fixed pairwise FP32 tree; this avoids changing recurrence rounding merely by changing the reduction schedule.
+
+The public numerical APIs `recurrent_decode_reference` and `recurrent_decode` accept one aligned `[T,H,D]` sequence and return outputs plus a `LinearAttentionCarry`. It retains anchor values/scales, buffered key/update values/scales, log retentions, position, cursor, and a policy signature. Returned tensors preserve gradients through prefix handoff, refreshes, and continuation. Restart with a new carry to reset; incompatible policy or state shape is rejected.
+
+The Torch reference can explicitly re-encode stored entries using `encoding="reencode"`. It re-encodes the already stored values and does not regenerate factors from a teacher trajectory. The fused implementation supports encode-once replay. Carrying its reconstructed state incrementally applies the same ordered transitions as replaying the fixed entries; it does not measure serving-time reconstruction cost.
+
+## Training implementation and limits
+
+The CUDA Triton implementation supports FP32 working inputs, key dimensions up to 128, scalar/channel gates, and value blocks 16/32/64/128 with masked tails. It saves a state every eight tokens and recomputes intermediate states in backward. This internal checkpoint interval does not change write cadence. Its custom backward supports first-order gradients; use the Torch reference for higher-order differentiation.
+
+Direct prefill APIs accept `prefill_lengths`, grouped query/key heads, both state layouts, packed sequences, empty entries, and full/zero prefixes. Positive prefixes are evaluated together so prefill quantizer scale domains remain shared. FLA KDA layer integration requires 0.5.1 and `use_cache=False`; serving cache objects are rejected. The numerical carry API provides explicit continuation, but this change does not implement a serving cache or export format. Megatron GDN tensor parallelism is covered by dedicated tests; context parallelism remains rejected.
+
+All encoded values remain floating tensors. This is training-time fake quantization and approximation, with no compressed storage or inference speed claim. Use `examples/llm_qat/linear_attention/benchmark_decode.py` for matched-policy training overhead and `decode_study_plan.json` for the fixed model-quality protocol. The exact inverse remains selected in these studies; the earlier rejected Neumann candidate is not enabled.
+
+See [the qualification study](linear_attention_decode_study.md) for the fixed
+model/data protocol, quality comparisons, and measured training overhead.

docs/linear_attention_decode_study.md

@@ -0,0 +1,113 @@
+# Decode-aware QAT qualification study
+
+## Fixed protocol
+
+The study uses `arcee-ai/AFM-4.5B-Base-KDA-Only` at revision
+`01ad2e06ee4f1214193c17b69e09105a9b257e80` and Wikitext-2 raw at revision
+`b08601e04326c79dfdd32d625aee71d232d685c3`. The actual checkpoint has 36 KDA
+layers, hidden size 2560, 20 heads, and key/value dimension 128. Parameters,
+source, dataset files, and token blocks are bound by stored hashes.
+
+`decode_study_plan.json` fixes selection before held-out test access. Every
+block contains 129 input tokens: 64 explicit prefix tokens followed by a decode
+suffix. The loss scores labels at positions 64 through 128 (65 predictions per
+block); it does not average prefix loss into decode quality. Controls and
+candidates use the exact prefill solve, FP32 working arithmetic, BF16 outer
+autocast, and seed 2026.
+
+Validation uses 32 blocks (2,080 scored tokens) without training. State FP8 and
+replay must each pass an NLL upper-bound margin of 0.02 against exact. Decay
+selection chooses the coarsest passing grid among 1/256, 1/1024, and 1/4096,
+requiring the same margin against both state FP8 and exact. Eligible modes then
+receive matched 32-step attention-only QAT pilots and 64 held-out test blocks.
+These descriptive paired-block bootstrap bounds are not broad quality guarantees.
+
+## Validation selection
+
+| Policy | Suffix perplexity | Mean NLL increase vs exact | 95% interval |
+| --- | ---: | ---: | --- |
+| Exact | 9.019945 | 0 | — |
+| Token state FP8 | 9.023610 | 0.000406 | [-0.002017, 0.002718] |
+| State FP8 + log grid 1/256 | 9.029220 | 0.001028 | [-0.001118, 0.002986] |
+| State FP8 + log grid 1/1024 | 9.035839 | 0.001761 | [-0.000115, 0.003723] |
+| State FP8 + log grid 1/4096 | 9.024412 | 0.000495 | [-0.001268, 0.002299] |
+| Replay FP8, window 8 | 9.019499 | -0.000049 | [-0.001867, 0.001630] |
+
+All families passed the validation gate. The predeclared coarsest-grid rule
+selected **1/256**; its paired comparison against state FP8 also passed. The
+held-out set was frozen to exact, token state FP8, state FP8 with grid 1/256,
+and replay FP8 with window 8, before accessing test data.
+
+## Held-out training pilot
+
+All four selected modes completed 32 attention-only updates on 32 training
+blocks (2,080 scored training tokens) using FP32 AdamW, learning rate 1e-5,
+zero weight decay, gradient clipping at 1.0, and non-reentrant activation
+checkpointing. All 994,043,088 trainable attention parameters had finite gradients
+on every update, and the query projection changed in each run. Evaluation used
+64 held-out blocks (4,160 scored tokens), with identical token hashes and
+training order across modes.
+
+| Policy | PPL before | PPL after | Post-training NLL delta vs exact | 95% interval |
+| --- | ---: | ---: | ---: | --- |
+| Exact | 12.899382 | 12.641259 | 0 | — |
+| Token state FP8 | 12.904970 | 12.656691 | 0.001220 | [-0.000290, 0.002703] |
+| State FP8 + grid 1/256 | 12.915614 | 12.653026 | 0.000930 | [-0.000802, 0.002673] |
+| Replay FP8, window 8 | 12.909424 | 12.646554 | 0.000419 | [-0.000566, 0.001410] |
+
+All three approximations met the 0.02 held-out pilot margin after training; the
+largest upper bound was 0.002703. Their intervals include zero, so this study does
+not establish a quality advantage for any approximation. All modes, including
+the exact control, improved after fine-tuning. These small short-context pilots
+verify model-level training behavior and bounded measured differences; they do
+not demonstrate recovery of a substantial accuracy gap or broad downstream quality.
+
+## Complete training cost
+
+The synthetic H100 workload is `[B,T,H,D] = [1,257,4,128]`, with 64 prefix tokens,
+FP32 working inputs, TF32 disabled, three warmups, and 20 interleaved samples.
+The table includes the complete prefix/suffix forward and backward, including
+returned final-state gradients. Speedups compare each fused implementation with
+its matching direct Torch numerical reference, with descriptive paired-median
+bootstrap intervals. Memory is peak allocated memory above the live input baseline.
+
+| Policy | Torch wall ms | Fused wall ms | Speedup [95% interval] | Torch / fused extra MiB |
+| --- | ---: | ---: | --- | ---: |
+| GDN exact | 188.830 | 4.758 | 39.64 [39.52, 40.19] | 103.75 / 14.68 |
+| GDN token | 224.532 | 7.239 | 31.02 [30.73, 31.06] | 103.75 / 14.68 |
+| GDN decay | 231.847 | 7.321 | 31.70 [31.57, 31.82] | 103.75 / 14.69 |
+| GDN replay | 325.854 | 6.206 | 52.47 [51.12, 52.73] | 272.08 / 14.94 |
+| KDA exact | 193.183 | 11.110 | 17.40 [17.33, 17.53] | 152.35 / 63.88 |
+| KDA token | 228.732 | 11.265 | 20.32 [20.25, 20.43] | 152.35 / 63.88 |
+| KDA decay | 235.972 | 11.310 | 20.87 [20.77, 20.98] | 152.35 / 64.26 |
+| KDA replay | 327.361 | 11.569 | 28.34 [28.16, 28.39] | 321.54 / 64.51 |
+
+The exact FLA reference, including conversion of Q/K/V to BF16, took 1.806 ms
+for GDN and 2.347 ms for KDA. Fused numerical emulation remains slower than that
+reference. These measurements establish a reduction in QAT emulation overhead,
+not native low-precision execution, inference latency, or compressed-cache cost.
+Every policy first passed checks of outputs, final state, and all input gradients
+against the matching Torch reference; source hashes accompany the raw samples.
+
+## Validation boundary
+
+The H100 regression suite passed 148 cases, with nine expected skips for FP32
+Hopper/TileLang modes that are rejected before launch. It includes GDN TP=1/2
+sharded save/restore and optimizer updates, KDA layer training/restore, explicit
+phase handling, and custom backward tests. All eight long-trajectory cases also
+passed the tightened `1e-4` relative-norm gate on H100 and SM86.
+
+Local checks passed 139 unit cases plus two re-encoding continuation cases,
+44 kernel/layer cases, and four additional full-width/value-tail cases covering
+32- and 128-column codec blocks (also passed on H100), plus two composed-policy
+cases (also passed on H100). The combined H100 total is **154 passed, nine
+expected skips**; the local GPU total is 50 passed. The full Sphinx build passed with FLA imports
+disabled. Measurements pin retained source snapshots. Final cleanup moved import
+boundaries and added explanatory comments; non-import AST equality was verified
+and focused regression/import checks were rerun. The retained benchmark harness
+differs only in the order of two imported typing names, with AST equality
+verified after normalizing that order.
+
+Model results apply to this checkpoint, short context, small corpus slice, and
+seed. Serving cache integration, real compression, long-context quality, and
+multi-model recovery remain unqualified.

examples/llm_qat/linear_attention/README.md

@@ -110,3 +110,28 @@ versions, and training settings. It reports per-block deltas and a descriptive
 95% paired block-bootstrap interval. Adjacent text blocks can remain correlated;
 the interval is not a broad downstream-quality guarantee. Compare a combined
 FP8/approximation candidate against its FP8 control to isolate the approximation.
+
+## Explicit prefix and decode training
+
+See [the decode contract](../../../docs/linear_attention_decode.md). The
+`kda_decode_*.json` configurations independently expose token state QDQ, rounded
+log retention, and replay factors/anchors. Pass `--prefill-tokens 64` to
+`quality_study.py` for explicit phase metadata and suffix-only loss; the context
+also spans activation-checkpoint recomputation. `decode_study_plan.json` fixes
+validation selection before held-out test access. These configurations keep the
+exact prefill solve.
+
+Measure complete prefix/suffix forward, backward, and peak allocated memory:
+
+```bash
+python examples/llm_qat/linear_attention/benchmark_decode.py \
+  --attention kda --length 257 --prefill 64 --dim 128 --output decode-cost.json
+```
+
+The benchmark first checks outputs, final states, and all input gradients against
+the same Torch numerical policy, then runs interleaved measurements after warmup.
+Its speedup compares training implementations of the same emulation; it is not
+an inference kernel or compressed-cache benchmark.
+
+The fixed validation selection, held-out pilot results, and measured training
+costs are recorded in [the decode study](../../../docs/linear_attention_decode_study.md).