kaix-nv · Sep 22, 2026
@@ -13,6 +13,8 @@ Changelog
 
 *Quantization*
 
+- 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.
+
 - Add experimental per-operand FP8/NVFP4 GDN prefill QAT with independently configurable accumulator and elementwise rounding. Select the materialized ``matmul`` backend to enable these sites; the existing fused FLA state/W path remains the default.
 
 - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts.

@@ -0,0 +1,98 @@
+# KDA prefill arithmetic for QAT
+
+ModelOpt recognizes FLA 0.5.1 `KimiDeltaAttention` layers during conversion and adds
+`kda_state_quantizer`, `kda_w_quantizer`, and the eight `linear_attn_sites` operand
+pairs. These handles start disabled. The existing projection recipes leave them
+disabled; enable them explicitly with `backend="matmul"`.
+
+The adapter reuses the layer's projections, convolution, normalization, and gate
+parameters. Its numerical path uses the materialized differentiable prefill
+backend. It does not patch process-wide FLA functions. When all numerical policies
+are off, the original FLA forward runs unchanged.
+
+## KDA site map and decay stability
+
+For a chunk, let `G[i,d]` be the cumulative natural-log retention for key channel
+`d`. KDA's interaction terms contain `exp(G[i,d] - G[j,d])`. The implementation
+computes that causal difference directly. It never constructs `exp(-G[j,d])`;
+strong forgetting can underflow a retention to zero without introducing the
+infinite inverse factor of an unstable factorization. Future positions are
+masked before the exponential and zeroed before operand quantization.
+
+The shared prefill sites have the following KDA operands. Right operands below
+are the row vectors presented to the quantizer before transposition.
+
+| Site | Left operand | Right operand |
+| --- | --- | --- |
+| `key_interaction` | `beta[i] * k[i]` | `k[j] * exp(G[i] - G[j])`, `j <= i` |
+| `wy_value` | exact triangular inverse | transposed beta-scaled values |
+| `wy_key` | exact triangular inverse | transposed beta-scaled, prefix-decayed keys |
+| `state_read` | WY keys, via `kda_w_quantizer` | transposed incoming state |
+| `state_update` | transposed end-decayed keys | transposed corrected values |
+| `output_state` | scaled, prefix-decayed queries | transposed incoming state |
+| `output_score` | scaled `q[i]` | `k[j] * exp(G[i] - G[j])`, `j <= i` |
+| `output_value` | causal local scores | transposed corrected values |
+
+For the two interaction sites, query rows are grouped in fixed tiles of eight
+and folded into the batch dimension. Quantizer inputs remain four-dimensional:
+`[chunk_or_sequence_times_query_rows, head, row, reduction]`. FP8 scales are
+per row. NVFP4 uses block-16 E2M1 values and E4M3 scales with one dynamic tensor
+amax per invocation; the query tile of eight is therefore part of its scale
+contract. Other sites use the [GDN prefill scale domains](linear_attention_prefill.md).
+
+Chunk size is 64. Working arithmetic is FP32, including inside an outer BF16
+training-autocast context; double inputs use FP64. Output returns to the query
+dtype. Select `torch.set_float32_matmul_precision("highest")` to exclude TF32.
+Operand and arithmetic rounding use identity STE; state carry is differentiable.
+The triangular solve remains exact.
+
+Raw gate activation follows the loaded FLA model: without `lower_bound`, it is
+`-exp(A_log) * softplus(raw_gate + dt_bias)`; with a lower bound, it is
+`lower_bound * sigmoid(exp(A_log) * (raw_gate + dt_bias))`. The latter is the
+model's bounded gate formula, not a new clamp. No additional decay approximation
+is enabled by this integration. The materialized path does not emulate the
+optional FLA `safe_gate` TensorCore rounding. Numerical parity checks use FLA's
+general computation (`safe_gate=False`) for both gate formulas.
+
+## Enable KDA operand QDQ
+
+```python
+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"},
+    }],
+})
+```
+
+To emulate stored-state writes, enable `*kda_state_quantizer` with dynamic E4M3
+and `axis=(0,1)`. The policy's `state.block_v` specifies the value-column tile;
+initial state and every chunk write are rounded. This is distinct from temporary
+state-read operand QDQ. State-V-first layout, grouped value heads, packed tails,
+and gradients through the initial state are supported.
+
+`flash-linear-attention==0.5.1` and `fla-core==0.5.1` are required for the layer
+adapter. Importing ModelOpt itself does not require either package. Qualification
+covers a single-GPU FLA model. CP, FLA-specific intermediate/recompute flags, and
+quantized recurrent inference are rejected. In evaluation mode FLA selects its
+recurrent path for lengths at most 64; use longer prefill sequences for this
+integration. Token-state training and decode/replay policies are separate work.
+
+## Reproducible model study
+
+The [study example](../examples/llm_qat/linear_attention/README.md) uses a public,
+pinned Arcee KDA checkpoint and WikiText-2 splits. It reports pre/post-training
+held-out NLL and perplexity with per-block results, hashes, seeds, and fixed token
+budgets. A successful training step verifies integration; it is not evidence of
+model-quality recovery. The materialized backend's time and memory must be
+measured separately from fused FLA.
@@ -10,7 +10,7 @@ The implementation adapts [PR #2455](https://github.com/NVIDIA/Model-Optimizer/p
 at `13c7e2456f2e9d079c9ef822742eeaa634353802` onto ModelOpt
 `051d6adb204f10cd3e78d0f824f31a5a01d54831`. It is the M0/M1 portion of the linear-attention
 design. [Additional prefill operand sites](linear_attention_prefill.md) have an explicit matmul backend.
-KDA training integration, approximate inverse,
+[KDA prefill training](linear_attention_kda.md) uses the same policy interface. Approximate inverse,
 token-state quantization, decay approximation, and SSM replay remain later milestones.
 
 ## Requirements and scope

@@ -0,0 +1,82 @@
+# Linear-attention numerical studies
+
+These examples study differentiable numerical emulation. The matmul backend
+materializes intermediates and can be substantially slower than fused FLA.
+
+## GDN overhead
+
+```bash
+PYTHONPATH=. python examples/llm_qat/linear_attention/benchmark_prefill.py \
+  --batch 1 --length 1024 --heads 4 --dim 64 --repeats 20 \
+  --output prefill-overhead.json
+```
+
+The JSON contains all interleaved forward/backward timing samples and extra tensor
+memory measurements. Compilation is excluded; CPU submission overhead is included.
+Add `--attention kda` to measure the per-channel KDA implementation.
+
+## KDA checkpoint and data
+
+The initial model is [arcee-ai/AFM-4.5B-Base-KDA-Only](https://huggingface.co/arcee-ai/AFM-4.5B-Base-KDA-Only)
+at revision `01ad2e06ee4f1214193c17b69e09105a9b257e80`. Its implementation uses
+FLA's `KimiDeltaAttention` directly. The dataset is
+[Salesforce/wikitext](https://huggingface.co/datasets/Salesforce/wikitext), configuration
+`wikitext-2-raw-v1`, revision `b08601e04326c79dfdd32d625aee71d232d685c3`.
+
+Install the example dependencies, then stage those exact revisions with model
+weights, tokenizer, audited model code, and the three Parquet splits. Supply local
+paths to the runner; it does not fetch weights or data. The checkpoint's custom
+model code requires an explicit `--trust-remote-code` opt-in. Review its pinned
+Python files before enabling that flag.
+
+```bash
+pip install -r examples/llm_qat/linear_attention/requirements.txt
+
+PYTHONPATH=. python examples/llm_qat/linear_attention/quality_study.py \
+  --model /path/to/arcee-kda \
+  --train-data /path/to/wikitext-2-raw-v1/train-00000-of-00001.parquet \
+  --eval-data /path/to/wikitext-2-raw-v1/validation-00000-of-00001.parquet \
+  --eval-split validation --trust-remote-code \
+  --quant-config examples/llm_qat/linear_attention/configs/kda_prefill_fp8.json \
+  --length 128 --train-steps 1 --eval-blocks 4 --seed 2026 \
+  --output fp8-smoke.json
+```
+
+The one-step command is an integration smoke test. For matched studies, run the
+same workload from the same original checkpoint with each numerical configuration:
+
+| Trial | Configuration |
+| --- | --- |
+| Exact fused FLA control | Omit `--quant-config` |
+| Exact materialized control | `configs/kda_exact_matmul.json` |
+| FP8 prefill operands | `configs/kda_prefill_fp8.json` |
+| NVFP4 prefill operands | `configs/kda_prefill_nvfp4.json` |
+
+The exact materialized control explicitly rounds the already-FP32 `output_add`
+to FP32 to select the emulation path. It introduces no lower-precision rounding.
+First establish agreement with the fused control before interpreting a candidate's
+quality. Candidate `before` measurements isolate its untrained numerical change;
+compare trained candidates against an equally trained exact control. Choose
+settings using validation data, then run the fixed selected settings on the test
+split. Never choose approximation degrees or training hyperparameters on test data.
+
+### Protocol
+
+- Join Parquet text rows with two newlines and tokenize without added special
+  tokens. Use the first requested contiguous blocks, each containing `length+1`
+  tokens and predicting `length` tokens. Blocks do not overlap; cache/state reset
+  for each block. This protocol is for paired comparisons, not a claim of the
+  checkpoint's published WikiText perplexity.
+- Shuffle the training blocks with the recorded seed. Use the same seed, block
+  count, token hashes, length, and learning rate across all trials.
+- Train KDA attention parameters only. Keep them and AdamW moments in FP32, use
+  BF16 autocast and activation checkpointing, and freeze other model parameters
+  in BF16. Preserve the checkpoint's FP32 gate parameters when loading. AdamW has
+  zero weight decay, default learning rate `1e-5`, and gradient clipping at 1.
+- Record per-block NLL, predicted token counts, source/data/token hashes, package
+  versions, training losses, finite-gradient checks, optimizer updates, step times,
+  and peak allocated training memory. The output contains no generated answers or
+  source document text.
+
+A short matched study can screen numerical sensitivity. It does not establish
+broad downstream-task quality, long-context equivalence, or native serving speed.
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Measure GDN prefill numerical-emulation overhead against the exact FLA path."""
+"""Measure GDN/KDA prefill numerical-emulation overhead against the exact FLA path."""
 
 import argparse
 import hashlib
@@ -28,12 +28,17 @@
 
 import torch
 import torch.nn.functional as F
+from fla.ops.kda import chunk_kda
 
 from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import (
     chunk_gated_delta_rule,
 )
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
-from modelopt.torch.quantization.linear_attention import LinearAttentionConfig, matmul_gdn
+from modelopt.torch.quantization.linear_attention import (
+    LinearAttentionConfig,
+    matmul_gdn,
+    matmul_kda,
+)
 from modelopt.torch.quantization.linear_attention.matmul import LinearAttentionMatmulSites
 from modelopt.torch.quantization.nn import TensorQuantizer
 
@@ -42,6 +47,7 @@ def main():
     """Run interleaved fixed-workload forward/backward trials and write a JSON receipt."""
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--attention", choices=["gdn", "kda"], default="gdn")
     parser.add_argument("--batch", type=int, default=1)
     parser.add_argument("--length", type=int, default=1024)
     parser.add_argument("--heads", type=int, default=4)
@@ -57,16 +63,17 @@ def main():
         for _ in range(2)
     ]
     v = torch.randn(shape, device="cuda", dtype=torch.bfloat16, requires_grad=True)
-    g = (-torch.rand(shape[:-1], device="cuda") * 0.03).requires_grad_()
-    beta = torch.rand_like(g, requires_grad=True)
+    gate_shape = shape if options.attention == "kda" else shape[:-1]
+    g = (-torch.rand(gate_shape, device="cuda") * 0.03).requires_grad_()
+    beta = torch.rand(shape[:-1], device="cuda", requires_grad=True)
     state = (
         torch.randn(options.batch, options.heads, options.dim, options.dim, device="cuda") * 0.1
     ).requires_grad_()
     inputs = (q, k, v, g, beta)
+    baseline = chunk_kda if options.attention == "kda" else chunk_gated_delta_rule
+    materialized = matmul_kda if options.attention == "kda" else matmul_gdn
     variants = {
-        "exact_fla": lambda: chunk_gated_delta_rule(
-            *inputs, initial_state=state, output_final_state=True
-        )
+        "exact_fla": lambda: baseline(*inputs, initial_state=state, output_final_state=True)
     }
     for mode in (
         "exact_matmul",
@@ -104,7 +111,7 @@ def main():
             )
 
         def run(sites=sites, w=w, policy=policy, mode=mode):
-            return matmul_gdn(
+            return materialized(
                 *inputs,
                 initial_state=state,
                 output_final_state=True,
@@ -161,9 +168,10 @@ def measure(fn):
         revision = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip()
     sources = [
         Path("modelopt/torch/quantization/linear_attention") / name
-        for name in ("prefill.py", "matmul.py", "reference.py", "config.py")
+        for name in ("prefill.py", "kda.py", "matmul.py", "reference.py", "config.py")
     ]
     result = {
+        "attention": options.attention,
         "git_head": revision,
         "source_sha256": {
             str(path): hashlib.sha256(path.read_bytes()).hexdigest() for path in sources

@@ -0,0 +1,20 @@
+{
+  "quant_cfg": [
+    {
+      "quantizer_name": "*",
+      "enable": false
+    }
+  ],
+  "algorithm": null,
+  "linear_attention": [
+    {
+      "module_name": "model.layers.*.self_attn",
+      "cfg": {
+        "backend": "matmul",
+        "elementwise": {
+          "output_add": "float32"
+        }
+      }
+    }
+  ]
+}
@@ -0,0 +1,47 @@
+{
+  "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": null,
+  "linear_attention": [
+    {
+      "module_name": "model.layers.*.self_attn",
+      "cfg": {
+        "backend": "matmul"
+      }
+    }
+  ]
+}