diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 32e711c8767..8b9ee34d828 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -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. diff --git a/docs/linear_attention_kda.md b/docs/linear_attention_kda.md new file mode 100644 index 00000000000..8c82306376f --- /dev/null +++ b/docs/linear_attention_kda.md @@ -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. diff --git a/docs/linear_attention_qat.md b/docs/linear_attention_qat.md index 05d7ed39b0f..f2351f52bd3 100644 --- a/docs/linear_attention_qat.md +++ b/docs/linear_attention_qat.md @@ -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 diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md new file mode 100644 index 00000000000..35757e94e34 --- /dev/null +++ b/examples/llm_qat/linear_attention/README.md @@ -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. diff --git a/examples/llm_qat/linear_attention/benchmark_prefill.py b/examples/llm_qat/linear_attention/benchmark_prefill.py index 8812ce2571d..53c85210527 100644 --- a/examples/llm_qat/linear_attention/benchmark_prefill.py +++ b/examples/llm_qat/linear_attention/benchmark_prefill.py @@ -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 diff --git a/examples/llm_qat/linear_attention/configs/kda_exact_matmul.json b/examples/llm_qat/linear_attention/configs/kda_exact_matmul.json new file mode 100644 index 00000000000..dc0ddcbb311 --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_exact_matmul.json @@ -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" + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_prefill_fp8.json b/examples/llm_qat/linear_attention/configs/kda_prefill_fp8.json new file mode 100644 index 00000000000..ba6b824c434 --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_prefill_fp8.json @@ -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" + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_prefill_nvfp4.json b/examples/llm_qat/linear_attention/configs/kda_prefill_nvfp4.json new file mode 100644 index 00000000000..13e4b66dd0a --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_prefill_nvfp4.json @@ -0,0 +1,53 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*linear_attn_sites.*", + "cfg": { + "num_bits": [ + 2, + 1 + ], + "type": "dynamic", + "block_sizes": { + "-1": 16, + "type": "dynamic", + "scale_bits": [ + 4, + 3 + ] + } + } + }, + { + "quantizer_name": "*kda_w_quantizer", + "cfg": { + "num_bits": [ + 2, + 1 + ], + "type": "dynamic", + "block_sizes": { + "-1": 16, + "type": "dynamic", + "scale_bits": [ + 4, + 3 + ] + } + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "model.layers.*.self_attn", + "cfg": { + "backend": "matmul" + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/quality_study.py b/examples/llm_qat/linear_attention/quality_study.py new file mode 100644 index 00000000000..bbcdf16342d --- /dev/null +++ b/examples/llm_qat/linear_attention/quality_study.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run a reproducible linear-attention QAT trial on pinned local model and data files.""" + +import argparse +import hashlib +import importlib.metadata +import json +import math +import random +import time +from pathlib import Path + +import pyarrow.parquet as pq +import torch +from fla.layers.kda import KimiDeltaAttention +from transformers import AutoModelForCausalLM, AutoTokenizer + +import modelopt.torch.quantization as mtq + + +def _token_blocks(path, tokenizer, length, count): + text = "\n\n".join(pq.read_table(path, columns=["text"])["text"].to_pylist()) + ids = tokenizer.encode(text, add_special_tokens=False) + needed = count * (length + 1) + if len(ids) < needed: + raise ValueError(f"{path} has {len(ids)} tokens; {needed} are required") + tensor = torch.tensor(ids[:needed], dtype=torch.int64).reshape(count, length + 1) + digest = hashlib.sha256(tensor.numpy().tobytes()).hexdigest() + return tensor, digest + + +def _evaluate(model, blocks): + model.eval() + losses = [] + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + for block in blocks: + ids = block.unsqueeze(0).cuda() + loss = model(input_ids=ids, labels=ids, use_cache=False).loss + if not torch.isfinite(loss): + raise RuntimeError("Non-finite evaluation loss") + losses.append(float(loss)) + mean = sum(losses) / len(losses) + return { + "mean_nll": mean, + "perplexity": math.exp(mean), + "block_nll": losses, + "predicted_tokens": len(blocks) * (blocks.shape[1] - 1), + } + + +def main(): + """Compare pre/post-training held-out NLL using an explicit numerical policy.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--model-id", default="arcee-ai/AFM-4.5B-Base-KDA-Only") + parser.add_argument("--model-revision", default="01ad2e06ee4f1214193c17b69e09105a9b257e80") + parser.add_argument("--train-data", type=Path, required=True) + parser.add_argument("--eval-data", type=Path, required=True) + parser.add_argument("--eval-split", choices=["validation", "test"], default="validation") + parser.add_argument("--dataset-revision", default="b08601e04326c79dfdd32d625aee71d232d685c3") + parser.add_argument("--quant-config", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--train-steps", type=int, default=1) + parser.add_argument("--eval-blocks", type=int, default=4) + parser.add_argument("--length", type=int, default=128) + parser.add_argument("--seed", type=int, default=2026) + parser.add_argument("--learning-rate", type=float, default=1e-5) + options = parser.parse_args() + if options.length < 64 or options.train_steps < 1 or options.eval_blocks < 1: + parser.error("length must be >=64 and train-steps/eval-blocks must be positive") + if options.train_data.resolve() == options.eval_data.resolve(): + parser.error("Training and evaluation data must be distinct splits") + torch.manual_seed(options.seed) + random.seed(options.seed) + torch.set_float32_matmul_precision("highest") + tokenizer = AutoTokenizer.from_pretrained( + options.model, local_files_only=True, trust_remote_code=options.trust_remote_code + ) + train, train_hash = _token_blocks( + options.train_data, tokenizer, options.length, options.train_steps + ) + evaluation, eval_hash = _token_blocks( + options.eval_data, tokenizer, options.length, options.eval_blocks + ) + model, loading = AutoModelForCausalLM.from_pretrained( + options.model, + local_files_only=True, + trust_remote_code=options.trust_remote_code, + torch_dtype=torch.float32, + output_loading_info=True, + ) + if loading["missing_keys"] or loading["unexpected_keys"] or loading.get("mismatched_keys"): + raise RuntimeError(f"Checkpoint/model mismatch: {loading}") + model.requires_grad_(False) + layers = [m for m in model.modules() if isinstance(m, KimiDeltaAttention)] + if not layers: + raise RuntimeError("The selected model contains no FLA KimiDeltaAttention layers") + # Keep trainable attention parameters and optimizer moments in FP32 under BF16 autocast. + for layer in layers: + layer.requires_grad_(True) + for parameter in model.parameters(): + if not parameter.requires_grad: + parameter.data = parameter.data.to(torch.bfloat16) + model.cuda() + quant_cfg = None + if options.quant_config: + quant_cfg = json.loads(options.quant_config.read_text()) + mtq.quantize(model, quant_cfg) + model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + before = _evaluate(model, evaluation) + trainable = [p for p in model.parameters() if p.requires_grad] + optimizer = torch.optim.AdamW(trainable, lr=options.learning_rate, weight_decay=0) + model.train() + losses = [] + grad_norms = [] + timings = [] + torch.cuda.reset_peak_memory_stats() + probe = layers[0].q_proj.weight + original = probe.detach().clone() + order = list(range(len(train))) + random.Random(options.seed).shuffle(order) + for index in order: + optimizer.zero_grad(set_to_none=True) + ids = train[index].unsqueeze(0).cuda() + torch.cuda.synchronize() + start = time.perf_counter() + with torch.autocast("cuda", dtype=torch.bfloat16): + loss = model(input_ids=ids, labels=ids, use_cache=False).loss + if not torch.isfinite(loss): + raise RuntimeError("Non-finite training loss") + loss.backward() + gradients = [p.grad for p in trainable if p.grad is not None] + if ( + len(gradients) != len(trainable) + or not torch.stack([torch.isfinite(g).all() for g in gradients]).all() + ): + raise RuntimeError("Missing or non-finite training gradients") + norm = torch.nn.utils.clip_grad_norm_(trainable, 1.0, error_if_nonfinite=True) + optimizer.step() + torch.cuda.synchronize() + losses.append(float(loss)) + grad_norms.append(float(norm)) + timings.append(time.perf_counter() - start) + print( + json.dumps( + { + "step": len(losses), + "loss": losses[-1], + "gradient_norm": grad_norms[-1], + "seconds": timings[-1], + } + ), + flush=True, + ) + if torch.equal(original, probe): + raise RuntimeError("Training did not update the query projection") + peak = torch.cuda.max_memory_allocated() + optimizer.zero_grad(set_to_none=True) + after = _evaluate(model, evaluation) + root = Path(__file__).resolve().parents[3] + sources = [ + root / "modelopt/torch/quantization/linear_attention" / name + for name in ["config.py", "kda.py", "prefill.py", "matmul.py"] + ] + sources.append(Path(__file__).resolve()) + result = { + "source_sha256": { + str(path.relative_to(root)): hashlib.sha256(path.read_bytes()).hexdigest() + for path in sources + }, + "model": options.model_id, + "model_revision": options.model_revision, + "dataset": "Salesforce/wikitext", + "dataset_config": "wikitext-2-raw-v1", + "dataset_revision": options.dataset_revision, + "eval_split": options.eval_split, + "train_file_sha256": hashlib.sha256(options.train_data.read_bytes()).hexdigest(), + "eval_file_sha256": hashlib.sha256(options.eval_data.read_bytes()).hexdigest(), + "train_tokens_sha256": train_hash, + "eval_tokens_sha256": eval_hash, + "sequence_length": options.length, + "seed": options.seed, + "train_order": order, + "training": "KDA attention parameters only; FP32 AdamW, BF16 autocast, gradient checkpointing", + "trainable_parameters": sum(p.numel() for p in trainable), + "kda_layers": len(layers), + "learning_rate": options.learning_rate, + "weight_decay": 0, + "gradient_clip": 1.0, + "quant_config": quant_cfg, + "before": before, + "after": after, + "training_losses": losses, + "gradient_norms": grad_norms, + "step_seconds": timings, + "train_predicted_tokens": options.train_steps * options.length, + "peak_train_bytes": peak, + "query_weight_changed": True, + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "packages": { + name: importlib.metadata.version(name) + for name in ["transformers", "flash-linear-attention", "fla-core", "triton"] + }, + } + options.output.parent.mkdir(parents=True, exist_ok=True) + options.output.write_text(json.dumps(result, indent=2) + "\n") + print( + json.dumps( + { + "before": before["perplexity"], + "after": after["perplexity"], + "output": str(options.output), + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/llm_qat/linear_attention/requirements.txt b/examples/llm_qat/linear_attention/requirements.txt new file mode 100644 index 00000000000..29636d78e52 --- /dev/null +++ b/examples/llm_qat/linear_attention/requirements.txt @@ -0,0 +1,4 @@ +fla-core==0.5.1 +flash-linear-attention==0.5.1 +pyarrow>=15 +transformers==4.57.1 diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 08861eac5c7..f36af6e77cd 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -156,7 +156,7 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: if "linear_attention" not in metadata: # Older checkpoints predate these disabled handles; preserve their baseline path. for name, module in _linear_attention_modules(model).items(): - for handle in ("gdn_state_quantizer", "gdn_w_quantizer"): + for handle in module.linear_attention_quantizer_names: key = f"{name}.{handle}" if name else handle quantizer = getattr(module, handle) if key not in quantizer_state_dict and not quantizer.is_enabled: @@ -228,12 +228,12 @@ def update_quantize_metadata( def _linear_attention_modules(model): # Optional framework plugins import conversion; defer this import to avoid that cycle. - from .plugins.gated_delta_net import GatedDeltaNetStateQuantMixin + from .plugins.linear_attention import _LinearAttentionQuantMixin return { get_unwrapped_name(name, model): module for name, module in model.named_modules() - if isinstance(module, GatedDeltaNetStateQuantMixin) + if isinstance(module, _LinearAttentionQuantMixin) } diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py index 91f61b2849c..6c1b27334b4 100644 --- a/modelopt/torch/quantization/linear_attention/__init__.py +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -16,6 +16,7 @@ """Numerical policies and differentiable references for linear attention.""" from .config import * +from .kda import * from .matmul import * from .prefill import * from .reference import * diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index a3399a3edc8..6c3776e8751 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -72,11 +72,11 @@ class _SolveConfig(ModeloptBaseConfig): class LinearAttentionConfig(ModeloptBaseConfig): - """GDN chunk-64 policy; unsupported numerical modes fail config validation. + """GDN/KDA chunk-64 policy; unsupported numerical modes fail config validation. ``state.block_v`` defines one dynamic scale per ``[Dk, block_v]`` tile of each sequence/head. The initial state and every chunk's final state are rounded when - ``gdn_state_quantizer`` is enabled. Outputs use the incoming rounded state. + the module's state quantizer is enabled. Outputs use the incoming rounded state. """ schema_version: Literal[1] = ModeloptField(default=1) diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py new file mode 100644 index 00000000000..5cd733a558c --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differentiable KDA prefill with stable per-channel decay interactions.""" + +import torch +import torch.nn.functional as F + +from .matmul import _validate_operand_quantizer +from .prefill import _matmul_prefill + +__all__ = ["matmul_kda"] + + +def matmul_kda( + q, + k, + v, + g, + beta, + *, + sites, + policy, + w_quantizer, + state_qdq=False, + scale=None, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=False, + use_beta_sigmoid_in_kernel=False, + allow_neg_eigval=False, + A_log=None, # noqa: N803 - match the FLA kernel signature + dt_bias=None, + safe_gate=False, + lower_bound=None, + cu_seqlens=None, + cu_seqlens_cpu=None, + state_v_first=False, + chunk_size=64, + cp_context=None, + disable_recompute=False, + return_intermediate_states=False, +): + """Run KDA chunk prefill with actual-operand QDQ and differentiable state carry. + + Gates are per-key-channel natural log retentions. Causal interactions form + exp(prefix[i] - prefix[j]) directly, avoiding an overflowing inverse decay. + Raw-gate activation, when requested, follows the selected FLA gate formula. + """ + if policy.backend != "matmul" or chunk_size != policy.chunk_size: + raise ValueError("matmul_kda requires backend='matmul' and its configured chunk size") + if cp_context is not None or disable_recompute or return_intermediate_states: + raise NotImplementedError( + "KDA matmul does not support CP or FLA recompute/intermediate flags" + ) + if allow_neg_eigval and not use_beta_sigmoid_in_kernel: + raise ValueError("allow_neg_eigval requires use_beta_sigmoid_in_kernel") + if safe_gate and use_gate_in_kernel and (lower_bound is None or not -5 <= lower_bound < 0): + raise ValueError("safe_gate requires a lower_bound in [-5, 0)") + if lower_bound is not None and lower_bound >= 0: + raise ValueError("lower_bound must be negative") + sites.validate() + _validate_operand_quantizer(w_quantizer, "kda_w_quantizer") + output_dtype = q.dtype + dtype = torch.float64 if output_dtype == torch.float64 else torch.float32 + q, k, v, g, beta = (x.to(dtype) for x in (q, k, v, g, beta)) + if g.ndim != 4: + raise ValueError("matmul_kda requires per-key-channel KDA log gates") + if use_qk_l2norm_in_kernel: + q, k = (x * (x.square().sum(-1, keepdim=True) + 1e-6).rsqrt() for x in (q, k)) + if use_gate_in_kernel: + if A_log is None: + raise ValueError("Fused KDA gate requires A_log") + if dt_bias is not None: + g = g + dt_bias.to(dtype).reshape(g.shape[-2:]) + rate = A_log.to(dtype).exp().reshape(g.shape[-2], 1) + g = -rate * F.softplus(g) if lower_bound is None else lower_bound * (rate * g).sigmoid() + if use_beta_sigmoid_in_kernel: + beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) + return _matmul_prefill( + q, + k, + v, + g, + beta, + sites=sites, + policy=policy, + w_quantizer=w_quantizer, + state_qdq=state_qdq, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + state_v_first=state_v_first, + chunk_size=chunk_size, + output_dtype=output_dtype, + ) diff --git a/modelopt/torch/quantization/linear_attention/matmul.py b/modelopt/torch/quantization/linear_attention/matmul.py index 49e7bf7694b..e8cadfb6d20 100644 --- a/modelopt/torch/quantization/linear_attention/matmul.py +++ b/modelopt/torch/quantization/linear_attention/matmul.py @@ -57,8 +57,8 @@ class LinearAttentionMatmulSites(nn.ModuleDict): """The eight prefill matmuls, with independent QDQ on their actual operands. Inputs to each handle are four-dimensional, with reduction on the last axis. - The state-read LHS uses the parent's legacy ``gdn_w_quantizer``; no second W - quantizer is registered here. New handles start disabled. + The state-read LHS uses the parent's ``gdn_w_quantizer`` or ``kda_w_quantizer``; + no second W quantizer is registered here. New handles start disabled. """ def __init__(self): diff --git a/modelopt/torch/quantization/linear_attention/prefill.py b/modelopt/torch/quantization/linear_attention/prefill.py index 8217f7189ab..a6e4475edff 100644 --- a/modelopt/torch/quantization/linear_attention/prefill.py +++ b/modelopt/torch/quantization/linear_attention/prefill.py @@ -77,6 +77,49 @@ def matmul_gdn( beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) if g.ndim != 3: raise ValueError("matmul_gdn requires scalar GDN log gates") + return _matmul_prefill( + q, + k, + v, + g, + beta, + sites=sites, + policy=policy, + w_quantizer=w_quantizer, + state_qdq=state_qdq, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + state_v_first=state_v_first, + chunk_size=chunk_size, + output_dtype=output_dtype, + ) + + +def _matmul_prefill( + q, + k, + v, + g, + beta, + *, + sites, + policy, + w_quantizer, + state_qdq, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + state_v_first, + chunk_size, + output_dtype, +): + dtype = q.dtype + channel_gate = g.ndim == 4 initial_state = initial_state.to(dtype) if initial_state is not None else None boundaries = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens q, k, state, sequences = _prepare( @@ -116,27 +159,35 @@ def mm(name, lhs, rhs): return sites.matmul(name, lhs, rhs, policy, w_quantizer=w_quantizer) qc, kc, vc = (chunks(x) for x in (q, k, v)) - gc = arithmetic("gate_prefix", chunks(g).cumsum(-1)) + gc = arithmetic("gate_prefix", chunks(g).cumsum(-2 if channel_gate else -1)) bc = chunks(beta).unsqueeze(-1) - causal = torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device).tril() - relative = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0) - decay = arithmetic("gate_exp", relative.exp()) - gate = arithmetic("gate_exp", gc.exp()).unsqueeze(-1) - lower = (mm("key_interaction", bc * kc, kc) * decay).tril(-1) + gate = arithmetic("gate_exp", gc.exp()) + if channel_gate: + lower = _kda_interaction("key_interaction", bc * kc, kc, gc, mm, arithmetic).tril(-1) + scores = _kda_interaction("output_score", qc * scale, kc, gc, mm, arithmetic).tril() + weighted_keys = kc * arithmetic("gate_exp", (gc[..., -1:, :] - gc).exp()) + final_decay = arithmetic("gate_exp", gc[..., -1, :].exp()).unsqueeze(-1) + else: + causal = torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device).tril() + relative = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0) + decay = arithmetic("gate_exp", relative.exp()) + gate = gate.unsqueeze(-1) + lower = (mm("key_interaction", bc * kc, kc) * decay).tril(-1) + scores = (mm("output_score", qc * scale, kc) * decay).tril() + weighted_keys = kc * arithmetic("gate_exp", (gc[..., -1:] - gc).exp()).unsqueeze(-1) + final_decay = arithmetic("gate_exp", gc[..., -1].exp())[..., None, None] identity = torch.eye(chunk_size, device=q.device, dtype=dtype).expand_as(lower) inverse = torch.linalg.solve_triangular( identity + lower, identity, upper=False, unitriangular=True ) u = mm("wy_value", inverse, (bc * vc).transpose(-1, -2)) w = mm("wy_key", inverse, (bc * kc * gate).transpose(-1, -2)) - scores = (mm("output_score", qc * scale, kc) * decay).tril() - weighted_keys = kc * arithmetic("gate_exp", (gc[..., -1:] - gc).exp()).unsqueeze(-1) def unbatch(x): return x.reshape(len(sequences), count, *x.shape[1:]) - qc, gate, gc, u, w, scores, weighted_keys = map( - unbatch, (qc, gate, gc, u, w, scores, weighted_keys) + qc, gate, final_decay, u, w, scores, weighted_keys = map( + unbatch, (qc, gate, final_decay, u, w, scores, weighted_keys) ) if state_qdq: state = state_fp8_qdq_reference(state, policy.state.block_v) @@ -149,9 +200,7 @@ def unbatch(x): memory = mm("output_state", qc[:, c] * scale * gate[:, c], state.transpose(-1, -2)) local = mm("output_value", scores[:, c], updated.transpose(-1, -2)) outputs.append(arithmetic("output_add", memory + local)) - decayed = arithmetic( - "state_decay", state * arithmetic("gate_exp", gc[:, c, :, -1].exp())[..., None, None] - ) + decayed = arithmetic("state_decay", state * final_decay[:, c]) update = mm( "state_update", weighted_keys[:, c].transpose(-1, -2), updated.transpose(-1, -2) ) @@ -168,3 +217,23 @@ def unbatch(x): output = torch.stack(pieces) if boundaries is None else torch.cat(pieces).unsqueeze(0) final = state.transpose(-1, -2) if state_v_first else state return output.to(output_dtype), final if output_final_state else None + + +def _kda_interaction(name, lhs, keys, prefix, mm, arithmetic): + # Causal pairwise decays avoid exp(-prefix), which overflows for strong forgetting. + batch, heads, length, width = keys.shape + result = [] + columns = torch.arange(length, device=keys.device) + for lo in range(0, length, 8): + hi = min(lo + 8, length) + rows = torch.arange(lo, hi, device=keys.device) + causal = (rows[:, None] >= columns[None, :])[None, None, ..., None] + relative = prefix[:, :, lo:hi, None] - prefix[:, :, None] + relative = relative.masked_fill(~causal, 0) + decay = arithmetic("gate_exp", relative.exp()).masked_fill(~causal, 0) + rhs = (keys[:, :, None] * decay).permute(0, 2, 1, 3, 4) + right = rhs.reshape(batch * (hi - lo), heads, length, width) + left = lhs[:, :, lo:hi].permute(0, 2, 1, 3).reshape(batch * (hi - lo), heads, 1, width) + product = mm(name, left, right).reshape(batch, hi - lo, heads, length) + result.append(product.permute(0, 2, 1, 3)) + return torch.cat(result, dim=-2) diff --git a/modelopt/torch/quantization/linear_attention/reference.py b/modelopt/torch/quantization/linear_attention/reference.py index 4e310079e4f..b94879253ae 100644 --- a/modelopt/torch/quantization/linear_attention/reference.py +++ b/modelopt/torch/quantization/linear_attention/reference.py @@ -25,7 +25,12 @@ import torch -__all__ = ["chunk_gdn_reference", "recurrent_delta_rule_reference", "state_fp8_qdq_reference"] +__all__ = [ + "chunk_gdn_reference", + "chunk_kda_reference", + "recurrent_delta_rule_reference", + "state_fp8_qdq_reference", +] def state_fp8_qdq_reference(state: torch.Tensor, block_v: int = 64) -> torch.Tensor: @@ -228,3 +233,87 @@ def chunk_gdn_reference( output = torch.cat(outputs).reshape(*v.shape) final = torch.stack(finals) return output, final.transpose(-1, -2) if state_v_first else final + + +def chunk_kda_reference( + q, + k, + v, + g, + beta, + *, + chunk_size=64, + scale=None, + initial_state=None, + cu_seqlens=None, + state_v_first=False, + state_qdq=False, + state_qdq_block_v=64, + w_quantizer=None, + matmul=None, +): + """KDA chunk oracle with causal per-channel decays and optional operand QDQ. + + The eight-site callback receives conventional three-dimensional matmul + operands. Interaction sites are invoked one query row at a time with decay + already applied to the right-hand keys. This avoids inverse-prefix factors. + """ + if g.ndim != 4 or chunk_size <= 0: + raise ValueError("chunk_kda_reference requires per-key gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + mm = matmul if matmul is not None else lambda name, lhs, rhs: lhs @ rhs + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + if state_qdq: + state = state_fp8_qdq_reference(state, state_qdq_block_v) + pieces = [] + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + prefix = g[b, lo:hi].transpose(0, 1).cumsum(-2) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + gate = prefix.exp() + lower_rows, score_rows = [], [] + for row in range(hi - lo): + decayed_keys = ( + kc[..., : row + 1, :] + * (prefix[..., row : row + 1, :] - prefix[..., : row + 1, :]).exp() + ) + right = decayed_keys.transpose(-1, -2) + left = kc[..., row : row + 1, :] * bc[..., row : row + 1, :] + interaction = mm("key_interaction", left, right) + score = mm("output_score", qc[..., row : row + 1, :] * scale, right) + # Zero padding follows the computation and never contributes a scale statistic. + padding = hi - lo - row - 1 + lower_rows.append(torch.nn.functional.pad(interaction, (0, padding))) + score_rows.append(torch.nn.functional.pad(score, (0, padding))) + lower = torch.cat(lower_rows, dim=-2).tril(-1) + scores = torch.cat(score_rows, dim=-2) + identity = torch.eye(hi - lo, device=q.device, dtype=q.dtype).expand_as(lower) + inverse = torch.linalg.solve_triangular( + identity + lower, identity, upper=False, unitriangular=True + ) + u = mm("wy_value", inverse, bc * vc) + w = mm("wy_key", inverse, bc * kc * gate) + if w_quantizer is not None: + w = w_quantizer(w) + updated = u - mm("state_read", w, state) + pieces.append( + ( + mm("output_state", qc * scale * gate, state) + + mm("output_value", scores, updated) + ).transpose(0, 1) + ) + weighted_keys = kc * (prefix[..., -1:, :] - prefix).exp() + state = state * gate[..., -1, :, None] + mm( + "state_update", weighted_keys.transpose(-1, -2), updated + ) + if state_qdq: + state = state_fp8_qdq_reference(state, state_qdq_block_v) + outputs.append(torch.cat(pieces)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final diff --git a/modelopt/torch/quantization/linear_attention/validation.py b/modelopt/torch/quantization/linear_attention/validation.py index 2bf07525e22..b3f5828eed6 100644 --- a/modelopt/torch/quantization/linear_attention/validation.py +++ b/modelopt/torch/quantization/linear_attention/validation.py @@ -20,8 +20,10 @@ __all__ = [] -def validate_gdn_quantizer(quantizer: TensorQuantizer, *, state: bool) -> None: - name = "gdn_state_quantizer" if state else "gdn_w_quantizer" +def validate_gdn_quantizer( + quantizer: TensorQuantizer, *, state: bool, name: str | None = None +) -> None: + name = name or ("gdn_state_quantizer" if state else "gdn_w_quantizer") axis = (0, 1) if state else (0, 1, 2) if not isinstance(quantizer, TensorQuantizer): raise ValueError(f"{name} requires a single TensorQuantizer") diff --git a/modelopt/torch/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py index 963452ff509..077b0d5d676 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -40,6 +40,7 @@ from .attention import * from .custom import * from .gated_delta_net import * +from .kda import * with import_plugin("diffusers"): from .diffusion.diffusers import * diff --git a/modelopt/torch/quantization/plugins/gated_delta_net.py b/modelopt/torch/quantization/plugins/gated_delta_net.py index a45d4a7e337..856d6307ff6 100644 --- a/modelopt/torch/quantization/plugins/gated_delta_net.py +++ b/modelopt/torch/quantization/plugins/gated_delta_net.py @@ -31,12 +31,8 @@ import torch -from ..config import QuantizerAttributeConfig -from ..linear_attention.config import LinearAttentionConfig -from ..linear_attention.matmul import LinearAttentionMatmulSites, _validate_operand_quantizer from ..linear_attention.prefill import matmul_gdn -from ..linear_attention.validation import validate_gdn_quantizer -from ..nn import QuantModule, TensorQuantizer +from .linear_attention import _LinearAttentionQuantMixin __all__ = ["GatedDeltaNetStateQuantMixin"] @@ -58,7 +54,7 @@ def _state_qdq_chunk_gated_delta_rule() -> GatedDeltaRuleFn: return chunk_gated_delta_rule -class GatedDeltaNetStateQuantMixin(QuantModule): +class GatedDeltaNetStateQuantMixin(_LinearAttentionQuantMixin): """Adds ``gdn_state_quantizer`` and ``gdn_w_quantizer`` to a GatedDeltaNet module. Subclasses route the module's chunked gated-delta-rule call through @@ -69,44 +65,11 @@ class GatedDeltaNetStateQuantMixin(QuantModule): E4M3 and identity STE. The execution policy is saved in ModelOpt metadata. """ - def _setup(self): - self.gdn_state_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) - self.gdn_w_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) - self.linear_attn_sites = LinearAttentionMatmulSites() - self.linear_attention_config = LinearAttentionConfig() - - @property - def linear_attention_is_enabled(self) -> bool: - """Whether any operand, state, or arithmetic policy changes the computation.""" - return ( - self.gdn_state_quantizer.is_enabled - or self.gdn_w_quantizer.is_enabled - or self.linear_attn_sites.is_enabled - or bool(self.linear_attention_config.matmul or self.linear_attention_config.elementwise) - ) - @property def gdn_state_qdq_block_v(self) -> int: """Value-column scale grouping from the saved execution policy.""" return self.linear_attention_config.state.block_v - def validate_linear_attention(self) -> None: - """Reject numerical settings that the fused training path cannot implement.""" - for state, quantizer in ((True, self.gdn_state_quantizer), (False, self.gdn_w_quantizer)): - if quantizer.is_enabled: - if not state and self.linear_attention_config.backend == "matmul": - _validate_operand_quantizer(quantizer, "gdn_w_quantizer") - else: - validate_gdn_quantizer(quantizer, state=state) - self.linear_attn_sites.validate() - if self.linear_attn_sites.is_enabled and self.linear_attention_config.backend != "matmul": - raise ValueError("Additional GDN operand sites require backend='matmul'") - - def modelopt_post_restore(self, prefix: str = ""): - """Validate restored quantizers before using the saved execution policy.""" - super().modelopt_post_restore(prefix) - self.validate_linear_attention() - def _state_quantized_chunk_gated_delta_rule( self, gated_delta_rule: GatedDeltaRuleFn, *args: Any, **kwargs: Any ) -> tuple[torch.Tensor, torch.Tensor | None]: diff --git a/modelopt/torch/quantization/plugins/kda.py b/modelopt/torch/quantization/plugins/kda.py new file mode 100644 index 00000000000..f972c501cc0 --- /dev/null +++ b/modelopt/torch/quantization/plugins/kda.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""QAT integration for the optional FLA KimiDeltaAttention layer.""" + +import importlib.metadata +from types import FunctionType + +from ..linear_attention.kda import matmul_kda +from ..nn import QuantModuleRegistry +from .custom import CUSTOM_MODEL_PLUGINS +from .linear_attention import _LinearAttentionQuantMixin + +__all__ = [] + + +class _QuantKimiDeltaAttention(_LinearAttentionQuantMixin): + linear_attention_quantizer_names = ("kda_state_quantizer", "kda_w_quantizer") + + def validate_linear_attention(self): + super().validate_linear_attention() + if self.linear_attention_is_enabled and self.linear_attention_config.backend != "matmul": + raise ValueError("KDA numerical emulation requires backend='matmul'") + + def forward(self, *args, **kwargs): + self.validate_linear_attention() + if not self.linear_attention_is_enabled: + return super().forward(*args, **kwargs) + original = super().forward.__func__ + # Bind this invocation so copied modules cannot retain another module's quantizers. + namespace = { + **original.__globals__, + "chunk_kda": self._quantized_chunk, + "fused_recurrent_kda": self._unsupported_recurrent, + } + forward = FunctionType( + original.__code__, + namespace, + original.__name__, + original.__defaults__, + original.__closure__, + ) + forward.__kwdefaults__ = original.__kwdefaults__ + return forward(self, *args, **kwargs) + + def _quantized_chunk(self, *args, **kwargs): + return matmul_kda( + *args, + sites=self.linear_attn_sites, + policy=self.linear_attention_config, + w_quantizer=self.kda_w_quantizer, + state_qdq=self.kda_state_quantizer.is_enabled and self.kda_state_quantizer._if_quant, + **kwargs, + ) + + def _unsupported_recurrent(self, *args, **kwargs): + raise NotImplementedError( + "KDA prefill QAT requires the chunk path; recurrent decode is not yet supported" + ) + + +def _register_fla_kda(model): + for module in model.modules(): + cls = type(module) + if cls.__module__ != "fla.layers.kda" or cls.__name__ != "KimiDeltaAttention": + continue + if cls in QuantModuleRegistry: + continue + if importlib.metadata.version("flash-linear-attention") != "0.5.1": + raise RuntimeError("KDA QAT is qualified with flash-linear-attention==0.5.1") + if importlib.metadata.version("fla-core") != "0.5.1": + raise RuntimeError("KDA QAT requires fla-core==0.5.1") + QuantModuleRegistry.register({cls: "fla_KimiDeltaAttention"})(_QuantKimiDeltaAttention) + + +CUSTOM_MODEL_PLUGINS.add(_register_fla_kda) diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py new file mode 100644 index 00000000000..b1d0f5e3cbd --- /dev/null +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared module policy and checkpoint support for linear-attention QAT.""" + +from ..config import QuantizerAttributeConfig +from ..linear_attention.config import LinearAttentionConfig +from ..linear_attention.matmul import LinearAttentionMatmulSites, _validate_operand_quantizer +from ..linear_attention.validation import validate_gdn_quantizer +from ..nn import QuantModule, TensorQuantizer + +__all__ = [] + + +class _LinearAttentionQuantMixin(QuantModule): + linear_attention_quantizer_names = ("gdn_state_quantizer", "gdn_w_quantizer") + + def _setup(self): + for name in self.linear_attention_quantizer_names: + setattr(self, name, TensorQuantizer(QuantizerAttributeConfig(enable=False))) + self.linear_attn_sites = LinearAttentionMatmulSites() + self.linear_attention_config = LinearAttentionConfig() + + @property + def _linear_attn_state(self): + return getattr(self, self.linear_attention_quantizer_names[0]) + + @property + def _linear_attn_w(self): + return getattr(self, self.linear_attention_quantizer_names[1]) + + @property + def linear_attention_is_enabled(self): + """Whether an operand, state, or arithmetic policy changes the computation.""" + return ( + self._linear_attn_state.is_enabled + or self._linear_attn_w.is_enabled + or self.linear_attn_sites.is_enabled + or bool(self.linear_attention_config.matmul or self.linear_attention_config.elementwise) + ) + + def validate_linear_attention(self): + """Validate quantizer contracts shared by GDN and KDA.""" + if self._linear_attn_state.is_enabled: + validate_gdn_quantizer( + self._linear_attn_state, state=True, name=self.linear_attention_quantizer_names[0] + ) + if self.linear_attention_config.backend == "matmul": + _validate_operand_quantizer( + self._linear_attn_w, self.linear_attention_quantizer_names[1] + ) + elif self._linear_attn_w.is_enabled: + validate_gdn_quantizer( + self._linear_attn_w, state=False, name=self.linear_attention_quantizer_names[1] + ) + self.linear_attn_sites.validate() + if self.linear_attn_sites.is_enabled and self.linear_attention_config.backend != "matmul": + raise ValueError("Additional linear-attention operand sites require backend='matmul'") + + def modelopt_post_restore(self, prefix=""): + """Validate the restored numerical policy and quantizers.""" + super().modelopt_post_restore(prefix) + self.validate_linear_attention() diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index f2b38ca46c7..3559145b48d 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -21,6 +21,10 @@ enable: false - quantizer_name: '*gdn_w_quantizer' enable: false + - quantizer_name: '*kda_state_quantizer' + enable: false + - quantizer_name: '*kda_w_quantizer' + enable: false - quantizer_name: '*linear_attn_sites.*' enable: false - quantizer_name: '*block_sparse_moe.gate*' diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_matmul_kda.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_matmul_kda.py new file mode 100644 index 00000000000..1b5c50becfa --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_matmul_kda.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +from _test_utils.torch.linear_attention import ( + FP8, + fp8_oracle, + inputs, + nvfp4_oracle, + reference_matmul, + values_and_gradients, +) + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + LinearAttentionMatmulSites, + chunk_kda_reference, + matmul_kda, +) +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.fixture(autouse=True) +def full_precision_matmul(): + previous = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") + yield + torch.set_float32_matmul_precision(previous) + + +def _inputs(*, packed=False, dtype=torch.float32): + args, state = inputs("cuda", packed=packed, dtype=dtype) + args[3] = ( + -torch.rand(*args[3].shape, args[0].shape[-1], device="cuda") * 0.05 + ).requires_grad_() + return args, state.float().detach().requires_grad_() + + +@pytest.mark.timeout(300) +@pytest.mark.parametrize("lower_bound", [None, -5.0]) +@pytest.mark.parametrize("packed", [False, True]) +def test_fla_baseline_and_fused_gate_gradients(lower_bound, packed): + pytest.importorskip("fla") + # FLA is optional except for comparison with its actual CUDA implementation. + from fla.ops.kda import chunk_kda + + args, state = _inputs(packed=packed, dtype=torch.bfloat16) + a_log = torch.full((2,), -3.0, device="cuda", requires_grad=True) + bias = torch.zeros(32, device="cuda", requires_grad=True) + kwargs = { + "initial_state": state, + "output_final_state": True, + "use_qk_l2norm_in_kernel": True, + "use_gate_in_kernel": True, + "use_beta_sigmoid_in_kernel": True, + "allow_neg_eigval": True, + "A_log": a_log, + "dt_bias": bias, + "lower_bound": lower_bound, + "safe_gate": lower_bound is not None, + "cu_seqlens": torch.tensor([0, 5, 73], device="cuda", dtype=torch.int32) + if packed + else None, + } + actual = matmul_kda( + *args, + sites=LinearAttentionMatmulSites(), + policy=LinearAttentionConfig(backend="matmul"), + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(enable=False)), + **kwargs, + ) + # Compare the general FLA algebra, not its optional safe_gate TensorCore rounding. + expected = chunk_kda(*args, **{**kwargs, "safe_gate": False}) + probes = [torch.randn_like(x) for x in actual] + gradients = [ + torch.autograd.grad( + sum((x * p).sum() for x, p in zip(result, probes)), (*args, state, a_log, bias) + ) + for result in (actual, expected) + ] + for i, (a, e) in enumerate(zip((*actual, *gradients[0]), (*expected, *gradients[1]))): + error = (a.float() - e.float()).norm() / e.float().norm().clamp_min(1e-6) + assert error <= (0.03 if i < 2 else 0.05), (i, error.item()) + + +def test_composed_fp8_with_chunk_state_and_w(): + args, state = _inputs(packed=True) + sites = LinearAttentionMatmulSites() + handles = [name for name, _ in sites.named_modules() if name.endswith("quantizer")] + for name in handles: + sites.get_submodule(name).set_from_attribute_config(FP8) + sites.get_submodule(name).enable() + policy = LinearAttentionConfig(backend="matmul", state={"block_v": 16}) + boundaries = torch.tensor([0, 5, 73], device="cuda") + actual = matmul_kda( + *args, + sites=sites, + policy=policy, + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(**FP8)), + state_qdq=True, + initial_state=state, + output_final_state=True, + cu_seqlens=boundaries, + ) + expected = chunk_kda_reference( + *args, + initial_state=state, + cu_seqlens=boundaries, + state_qdq=True, + state_qdq_block_v=16, + w_quantizer=fp8_oracle, + matmul=reference_matmul(handles, policy), + ) + for i, (a, e) in enumerate( + zip(values_and_gradients(actual, args, state), values_and_gradients(expected, args, state)) + ): + error = (a - e).norm() / e.norm().clamp_min(1e-6) + assert error <= (0.01 if i < 2 else 0.02), (i, error.item()) + + +@pytest.mark.parametrize("site", ["key_interaction", "output_score"]) +@pytest.mark.parametrize("side", ["lhs_quantizer", "rhs_quantizer"]) +def test_nvfp4_channel_decay_interaction(site, side): + args, state = _inputs() + sites = LinearAttentionMatmulSites() + target = sites.get_submodule(f"{site}.{side}") + target.set_from_attribute_config( + { + "num_bits": (2, 1), + "type": "dynamic", + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + } + ) + target.enable() + captured = [] + hook = target.register_forward_hook(lambda m, a, o: captured.append((a[0], o))) + actual = matmul_kda( + *args, + sites=sites, + policy=LinearAttentionConfig(backend="matmul"), + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(enable=False)), + initial_state=state, + output_final_state=True, + ) + hook.remove() + assert len(captured) == 8 + amaxes = [] + for x, y in captured: + torch.testing.assert_close(y, nvfp4_oracle(x), rtol=2e-5, atol=2e-6) + probe = torch.randn_like(y) + torch.testing.assert_close( + torch.autograd.grad((y * probe).sum(), x, retain_graph=True)[0], probe, rtol=0, atol=0 + ) + amaxes.append(x.detach().abs().amax()) + index = 0 + + def mm(name, lhs, rhs): + nonlocal index + if name == site: + amax = amaxes[(index % 64) // 8] + index += 1 + if side == "lhs_quantizer": + lhs = nvfp4_oracle(lhs, amax) + else: + rhs = nvfp4_oracle(rhs.transpose(-1, -2), amax).transpose(-1, -2) + return torch.einsum("hij,hjk->hik", lhs, rhs) + + expected = chunk_kda_reference(*args, initial_state=state, matmul=mm) + for a, e in zip( + values_and_gradients(actual, args, state), values_and_gradients(expected, args, state) + ): + torch.testing.assert_close(a, e, rtol=5e-4, atol=3e-6) diff --git a/tests/gpu/torch/quantization/plugins/test_kda.py b/tests/gpu/torch/quantization/plugins/test_kda.py new file mode 100644 index 00000000000..1f254dc582e --- /dev/null +++ b/tests/gpu/torch/quantization/plugins/test_kda.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy + +import pytest +import torch + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import LinearAttentionConfig +from modelopt.torch.quantization.nn import TensorQuantizer + +KimiDeltaAttention = pytest.importorskip("fla.layers.kda").KimiDeltaAttention + + +def _layer(): + return ( + KimiDeltaAttention(hidden_size=32, head_dim=16, num_heads=2, use_short_conv=True) + .cuda() + .train() + ) + + +def _forward(model, hidden): + with torch.autocast("cuda", dtype=torch.bfloat16): + return model(hidden)[0] + + +@pytest.mark.parametrize("mode", ["state", "w", "prefill_fp8", "prefill_nvfp4", "arithmetic"]) +@pytest.mark.timeout(180) +def test_fla_layer_qat_restore_and_optimizer(tmp_path, mode): + torch.manual_seed(73) + model = _layer() + hidden = torch.randn(1, 73, 32, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + baseline = _forward(model, hidden) + attributes = {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)} + if mode == "prefill_nvfp4": + attributes = { + "num_bits": (2, 1), + "type": "dynamic", + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + } + cfg = { + "quant_cfg": [{"quantizer_name": "*", "enable": False}], + "algorithm": None, + "linear_attention": [{"module_name": "*", "cfg": {"backend": "matmul"}}], + } + if mode in ("w", "prefill_fp8", "prefill_nvfp4"): + cfg["quant_cfg"].append({"quantizer_name": "*kda_w_quantizer", "cfg": attributes}) + if mode.startswith("prefill"): + cfg["quant_cfg"].append({"quantizer_name": "*linear_attn_sites.*", "cfg": attributes}) + if mode == "state": + cfg["quant_cfg"].append( + {"quantizer_name": "*kda_state_quantizer", "cfg": {**attributes, "axis": (0, 1)}} + ) + if mode == "arithmetic": + cfg["linear_attention"][0]["cfg"]["elementwise"] = {"value_residual": "bfloat16"} + mtq.quantize(model, cfg) + output = _forward(model, hidden) + assert torch.isfinite(output).all() + enabled = [q for q in model.modules() if isinstance(q, TensorQuantizer) and q.is_enabled] + policy = model.linear_attention_config + for q in enabled: + q.disable() + model.linear_attention_config = LinearAttentionConfig() + with torch.no_grad(): + torch.testing.assert_close(_forward(model, hidden), baseline, rtol=0, atol=0) + for q in enabled: + q.enable() + model.linear_attention_config = policy + checkpoint = tmp_path / "model.pth" + mto.save(model, checkpoint) + restored = mto.restore(_layer(), checkpoint) + actual = _forward(restored, hidden) + torch.testing.assert_close(actual, output, rtol=0, atol=0) + assert restored.linear_attention_config == policy + actual.float().square().mean().backward() + assert restored.q_proj.weight.grad is not None + for parameter in restored.parameters(): + assert parameter.grad is not None + assert torch.isfinite(parameter.grad).all() + if mode == "prefill_fp8": + copied = copy.deepcopy(restored) + calls = [] + hook = copied.kda_w_quantizer.register_forward_hook(lambda *args: calls.append(True)) + with torch.no_grad(): + torch.testing.assert_close(_forward(copied, hidden), actual, rtol=0, atol=0) + hook.remove() + assert calls, "a copied layer must use its own quantizer handles" + before = restored.q_proj.weight.detach().clone() + torch.optim.SGD(restored.parameters(), lr=0.1).step() + assert not torch.equal(before, restored.q_proj.weight) + restored.eval() + with pytest.raises(NotImplementedError, match="chunk path"): + _forward(restored, hidden[:, :1]) diff --git a/tests/unit/torch/quantization/test_kda_prefill.py b/tests/unit/torch/quantization/test_kda_prefill.py new file mode 100644 index 00000000000..d8ddecf07c5 --- /dev/null +++ b/tests/unit/torch/quantization/test_kda_prefill.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +from _test_utils.torch.linear_attention import FP8, inputs, reference_matmul, values_and_gradients + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + LinearAttentionMatmulSites, + chunk_kda_reference, + matmul_kda, + recurrent_delta_rule_reference, +) +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.mark.parametrize(("packed", "state_v_first"), [(False, False), (True, True)]) +@pytest.mark.parametrize("strong_decay", [False, True]) +def test_kda_exact_outputs_state_and_gradients(packed, state_v_first, strong_decay): + args, state = inputs(packed=packed, state_v_first=state_v_first, dtype=torch.float64) + args[2] = args[2][..., :11].detach().requires_grad_() + state = (state[..., :11, :] if state_v_first else state[..., :11]).detach().requires_grad_() + gates = -torch.rand(*args[3].shape, args[0].shape[-1], dtype=torch.float64) + if strong_decay: + gates = gates * torch.logspace(-8, 3, gates.shape[-1], dtype=torch.float64) + else: + gates = gates * 0.1 + args[3] = gates.requires_grad_() + kwargs = { + "initial_state": state, + "state_v_first": state_v_first, + "cu_seqlens": torch.tensor([0, 5, 73]) if packed else None, + } + actual = matmul_kda( + *args, + sites=LinearAttentionMatmulSites(), + policy=LinearAttentionConfig(backend="matmul"), + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(enable=False)), + output_final_state=True, + **kwargs, + ) + recurrent = recurrent_delta_rule_reference(*args, **kwargs) + chunk = chunk_kda_reference(*args, **kwargs) + for result in (actual, chunk): + for a, e in zip( + values_and_gradients(result, args, state), values_and_gradients(recurrent, args, state) + ): + assert torch.isfinite(a).all() + torch.testing.assert_close(a, e, rtol=2e-9, atol=1e-10) + + +HANDLES = [ + name for name, _ in LinearAttentionMatmulSites().named_modules() if name.endswith("quantizer") +] + + +@pytest.mark.parametrize("handle", HANDLES) +def test_kda_fp8_site_matches_chunk_oracle(handle): + args, state = inputs() + args[3] = (-torch.rand(*args[3].shape, args[0].shape[-1]) * 0.1).requires_grad_() + sites = LinearAttentionMatmulSites() + sites.get_submodule(handle).set_from_attribute_config(FP8) + sites.get_submodule(handle).enable() + policy = LinearAttentionConfig(backend="matmul") + actual = matmul_kda( + *args, + sites=sites, + policy=policy, + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(enable=False)), + initial_state=state, + output_final_state=True, + ) + expected = chunk_kda_reference( + *args, initial_state=state, matmul=reference_matmul([handle], policy) + ) + for a, e in zip( + values_and_gradients(actual, args, state), values_and_gradients(expected, args, state) + ): + torch.testing.assert_close(a, e, rtol=5e-4, atol=3e-6)