From ae9d2a76cc5f5a416397f3f53008dd00b2848a64 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 22 Sep 2026 14:09:32 -0700 Subject: [PATCH] Add experimental GDN and KDA decode and replay QAT Signed-off-by: Kai Xu --- CHANGELOG.rst | 2 + docs/linear_attention_decode.md | 61 +++ docs/linear_attention_decode_study.md | 113 +++++ examples/llm_qat/linear_attention/README.md | 25 + .../linear_attention/benchmark_decode.py | 194 ++++++++ .../linear_attention/compare_quality.py | 2 + .../configs/kda_decode_decay_1024.json | 35 ++ .../configs/kda_decode_decay_256.json | 35 ++ .../configs/kda_decode_decay_4096.json | 35 ++ .../configs/kda_decode_exact.json | 20 + .../configs/kda_decode_replay_fp8.json | 40 ++ .../configs/kda_decode_state_fp8.json | 34 ++ .../linear_attention/decode_study_plan.json | 38 ++ .../llm_qat/linear_attention/quality_study.py | 68 ++- .../quantization/linear_attention/decode.py | 429 ++++++++++++++++++ .../quantization/linear_attention/__init__.py | 2 + .../quantization/linear_attention/config.py | 36 +- .../quantization/linear_attention/decode.py | 352 ++++++++++++++ .../linear_attention/decode_prefill.py | 180 ++++++++ .../quantization/linear_attention/kda.py | 22 + .../quantization/linear_attention/prefill.py | 22 + .../quantization/plugins/gated_delta_net.py | 1 + modelopt/torch/quantization/plugins/kda.py | 16 + .../quantization/plugins/linear_attention.py | 2 + .../linear_attention/test_decode.py | 305 +++++++++++++ .../torch/quantization/plugins/test_kda.py | 90 +++- .../plugins/test_megatron_gated_delta_net.py | 48 +- .../test_linear_attention_decode.py | 266 +++++++++++ .../test_linear_attention_study.py | 12 +- 29 files changed, 2461 insertions(+), 24 deletions(-) create mode 100644 docs/linear_attention_decode.md create mode 100644 docs/linear_attention_decode_study.md create mode 100644 examples/llm_qat/linear_attention/benchmark_decode.py create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_decay_1024.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_decay_256.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_decay_4096.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_exact.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_replay_fp8.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_state_fp8.json create mode 100644 examples/llm_qat/linear_attention/decode_study_plan.json create mode 100644 modelopt/torch/kernels/quantization/linear_attention/decode.py create mode 100644 modelopt/torch/quantization/linear_attention/decode.py create mode 100644 modelopt/torch/quantization/linear_attention/decode_prefill.py create mode 100644 tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py create mode 100644 tests/unit/torch/quantization/test_linear_attention_decode.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d743fa08125..3ca7f96c45e 100755 --- a/CHANGELOG.rst +++ b/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. diff --git a/docs/linear_attention_decode.md b/docs/linear_attention_decode.md new file mode 100644 index 00000000000..24d5bdcf1d9 --- /dev/null +++ b/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. diff --git a/docs/linear_attention_decode_study.md b/docs/linear_attention_decode_study.md new file mode 100644 index 00000000000..a5f2c53209a --- /dev/null +++ b/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. diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md index af549e079d1..2b44ae4dc75 100644 --- a/examples/llm_qat/linear_attention/README.md +++ b/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). diff --git a/examples/llm_qat/linear_attention/benchmark_decode.py b/examples/llm_qat/linear_attention/benchmark_decode.py new file mode 100644 index 00000000000..a234a15f9de --- /dev/null +++ b/examples/llm_qat/linear_attention/benchmark_decode.py @@ -0,0 +1,194 @@ +# 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. + +"""Compare complete prefix/suffix QAT cost under matched token/replay policies.""" + +import argparse +import hashlib +import json +import random +import statistics +import time +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn.functional as F +from fla.ops.gated_delta_rule import chunk_gated_delta_rule +from fla.ops.kda import chunk_kda + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + LinearAttentionMatmulSites, + matmul_gdn, + matmul_kda, +) +from modelopt.torch.quantization.nn import TensorQuantizer + +if TYPE_CHECKING: + from collections.abc import Callable + + +def main(): + """Check outputs/gradients before interleaved, warmup-excluded measurements.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--attention", choices=["gdn", "kda"], default="kda") + parser.add_argument("--length", type=int, default=257) + parser.add_argument("--prefill", type=int, default=64) + parser.add_argument("--heads", type=int, default=4) + parser.add_argument("--dim", type=int, default=128) + parser.add_argument("--repeats", type=int, default=20) + options = parser.parse_args() + if options.repeats < 2 or not 0 <= options.prefill < options.length: + parser.error("Require >=2 repeats and a nonempty suffix") + torch.manual_seed(2026) + torch.backends.cuda.matmul.allow_tf32 = False + shape = (1, options.length, options.heads, options.dim) + q, k = [ + F.normalize(torch.randn(shape, device="cuda"), dim=-1).requires_grad_() for _ in range(2) + ] + v = torch.randn(shape, device="cuda", requires_grad=True) + g = ( + -torch.rand(shape if options.attention == "kda" else shape[:-1], device="cuda") * 0.03 + ).requires_grad_() + beta = (torch.rand(shape[:-1], device="cuda") * 0.4).requires_grad_() + initial = ( + torch.randn(1, options.heads, options.dim, options.dim, device="cuda") * 0.1 + ).requires_grad_() + args = q, k, v, g, beta + sites = LinearAttentionMatmulSites() + w = TensorQuantizer(QuantizerAttributeConfig(enable=False)) + function = matmul_kda if options.attention == "kda" else matmul_gdn + variants: dict[str, Callable[[], tuple[torch.Tensor, torch.Tensor]]] = {} + for mode in ["exact", "token", "decay", "replay"]: + for impl in ["torch", "triton"]: + decode: dict[str, Any] = {"implementation": impl} + if mode == "replay": + decode.update(mode="replay", replay={"window": 8}) + if mode == "decay": + decode["decay_log_step"] = 1 / 256 + policy = LinearAttentionConfig(backend="matmul", decode=decode) + + def run(policy=policy, mode=mode): + return function( + *args, + initial_state=initial, + output_final_state=True, + sites=sites, + w_quantizer=w, + policy=policy, + state_qdq=mode != "exact", + prefill_lengths=[options.prefill], + ) + + variants[f"{mode}_{impl}"] = run + native = chunk_kda if options.attention == "kda" else chunk_gated_delta_rule + variants["fla_bf16"] = lambda: native( + q.bfloat16(), + k.bfloat16(), + v.bfloat16(), + g, + beta, + initial_state=initial, + output_final_state=True, + ) + probes = [torch.randn_like(v), torch.randn_like(initial)] + correctness = {} + for mode in ["exact", "token", "decay", "replay"]: + results = [] + for impl in ["torch", "triton"]: + output = variants[f"{mode}_{impl}"]() + grads = torch.autograd.grad( + sum((x * p).sum() for x, p in zip(output, probes)), (*args, initial) + ) + results.append(tuple(x.detach() for x in (*output, *grads))) + errors = [float((a - b).norm() / a.norm().clamp_min(1e-8)) for a, b in zip(*results)] + if max(errors) > 1e-4 or not all( + torch.isfinite(x).all() for result in results for x in result + ): + raise RuntimeError(f"{mode}: correctness gate failed: {errors}") + correctness[mode] = errors + + def measure(fn): + torch.cuda.synchronize() + allocated = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + a, b, c = [torch.cuda.Event(enable_timing=True) for _ in range(3)] + start = time.perf_counter() + a.record() + output, final = fn() + b.record() + torch.autograd.grad(output.square().mean() + final.square().mean(), (*args, initial)) + c.record() + torch.cuda.synchronize() + return { + "forward_ms": a.elapsed_time(b), + "backward_ms": b.elapsed_time(c), + "wall_ms": 1000 * (time.perf_counter() - start), + "peak_extra_bytes": torch.cuda.max_memory_allocated() - allocated, + } + + for fn in variants.values(): + for _ in range(3): + measure(fn) + samples = {name: [] for name in variants} + rng = random.Random(2026) + for _ in range(options.repeats): + order = list(variants) + rng.shuffle(order) + for name in order: + samples[name].append(measure(variants[name])) + summary = { + name: {key: statistics.median(s[key] for s in values) for key in values[0]} + for name, values in samples.items() + } + for mode in ["exact", "token", "decay", "replay"]: + ratios = [ + a["wall_ms"] / b["wall_ms"] + for a, b in zip(samples[f"{mode}_torch"], samples[f"{mode}_triton"]) + ] + bootstrap = sorted( + statistics.median(rng.choices(ratios, k=len(ratios))) for _ in range(5000) + ) + summary[f"{mode}_triton"]["paired_speedup_vs_reference"] = { + "median": statistics.median(ratios), + "interval_95": [bootstrap[124], bootstrap[4874]], + } + sources = [ + *list(Path("modelopt/torch/quantization/linear_attention").glob("*.py")), + Path("modelopt/torch/kernels/quantization/linear_attention/decode.py"), + Path(__file__), + ] + result = { + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "shape": shape, + "prefill": options.prefill, + "attention": options.attention, + "dtype": "float32; FLA baseline casts Q/K/V to bfloat16", + "repeats": options.repeats, + "correctness_relative_norms": correctness, + "source_sha256": {str(p): hashlib.sha256(p.read_bytes()).hexdigest() for p in sources}, + "summary": summary, + "samples": samples, + } + options.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(summary, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/llm_qat/linear_attention/compare_quality.py b/examples/llm_qat/linear_attention/compare_quality.py index 481a92a4231..24b3fe953e4 100644 --- a/examples/llm_qat/linear_attention/compare_quality.py +++ b/examples/llm_qat/linear_attention/compare_quality.py @@ -34,6 +34,8 @@ "train_tokens_sha256", "eval_tokens_sha256", "sequence_length", + "prefill_tokens", + "loss_scope", "seed", "train_order", "training", diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_decay_1024.json b/examples/llm_qat/linear_attention/configs/kda_decode_decay_1024.json new file mode 100644 index 00000000000..3e49626b796 --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_decay_1024.json @@ -0,0 +1,35 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": [ + 4, + 3 + ], + "type": "dynamic", + "axis": [ + 0, + 1 + ] + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "decay_log_step": 0.0009765625 + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_decay_256.json b/examples/llm_qat/linear_attention/configs/kda_decode_decay_256.json new file mode 100644 index 00000000000..0b3b972c25d --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_decay_256.json @@ -0,0 +1,35 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": [ + 4, + 3 + ], + "type": "dynamic", + "axis": [ + 0, + 1 + ] + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "decay_log_step": 0.00390625 + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_decay_4096.json b/examples/llm_qat/linear_attention/configs/kda_decode_decay_4096.json new file mode 100644 index 00000000000..3b488c92bbb --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_decay_4096.json @@ -0,0 +1,35 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": [ + 4, + 3 + ], + "type": "dynamic", + "axis": [ + 0, + 1 + ] + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "decay_log_step": 0.000244140625 + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_exact.json b/examples/llm_qat/linear_attention/configs/kda_decode_exact.json new file mode 100644 index 00000000000..4402b29f1bc --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_exact.json @@ -0,0 +1,20 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton" + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_replay_fp8.json b/examples/llm_qat/linear_attention/configs/kda_decode_replay_fp8.json new file mode 100644 index 00000000000..e7a86e20ca3 --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_replay_fp8.json @@ -0,0 +1,40 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": [ + 4, + 3 + ], + "type": "dynamic", + "axis": [ + 0, + 1 + ] + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "mode": "replay", + "replay": { + "window": 8, + "factor_qdq": true, + "encoding": "once" + } + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_state_fp8.json b/examples/llm_qat/linear_attention/configs/kda_decode_state_fp8.json new file mode 100644 index 00000000000..27ab6b2fbdd --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_state_fp8.json @@ -0,0 +1,34 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": [ + 4, + 3 + ], + "type": "dynamic", + "axis": [ + 0, + 1 + ] + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton" + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/decode_study_plan.json b/examples/llm_qat/linear_attention/decode_study_plan.json new file mode 100644 index 00000000000..ef9692169fc --- /dev/null +++ b/examples/llm_qat/linear_attention/decode_study_plan.json @@ -0,0 +1,38 @@ +{ + "model": "arcee-ai/AFM-4.5B-Base-KDA-Only", + "model_revision": "01ad2e06ee4f1214193c17b69e09105a9b257e80", + "dataset": "Salesforce/wikitext/wikitext-2-raw-v1", + "dataset_revision": "b08601e04326c79dfdd32d625aee71d232d685c3", + "seed": 2026, + "length": 128, + "prefill_tokens": 64, + "loss_scope": "suffix", + "validation": { + "blocks": 32, + "train_steps": 0, + "controls": [ + "exact", + "state_fp8" + ], + "candidates": [ + "decay_256", + "decay_1024", + "decay_4096", + "replay_fp8" + ], + "rule": "Select coarsest decay grid whose paired NLL upper bound is <=0.02 versus state_fp8 and exact; qualify replay independently versus exact. State FP8 must pass versus exact. Do not access test until selection is frozen." + }, + "held_out": { + "eval_split": "test", + "blocks": 64, + "train_steps": 32, + "learning_rate": 1e-05, + "gradient_clip": 1.0, + "modes": "exact, state_fp8 if eligible, selected decay if eligible, replay if eligible", + "rule": "Run matched attention-only QAT pilots. Report before/after paired NLL deltas, intervals, finite gradients and updates; a passing pilot does not establish broad quality recovery." + }, + "fallback": "Keep exact default; explicitly report failed gates. Do not tune on test.", + "integration": { + "failed_policy_smoke": "One attention-only training step and four validation blocks for any family excluded by screening; report solely as gradient/update integration, never recovery or held-out qualification." + } +} diff --git a/examples/llm_qat/linear_attention/quality_study.py b/examples/llm_qat/linear_attention/quality_study.py index 779fdc4e087..8531ab9d228 100644 --- a/examples/llm_qat/linear_attention/quality_study.py +++ b/examples/llm_qat/linear_attention/quality_study.py @@ -22,6 +22,7 @@ import math import random import time +from contextlib import nullcontext from pathlib import Path import pyarrow.parquet as pq @@ -30,6 +31,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase def _token_blocks(path, tokenizer, length, count): @@ -43,13 +45,32 @@ def _token_blocks(path, tokenizer, length, count): return tensor, digest -def _evaluate(model, blocks): +def _phase(model, prefill_tokens): + return ( + nullcontext() + if prefill_tokens is None + else linear_attention_training_phase(model, [prefill_tokens]) + ) + + +def _labels(ids, prefill_tokens): + labels = ids.clone() + if prefill_tokens is not None: + labels[:, :prefill_tokens] = -100 + return labels + + +def _evaluate(model, blocks, prefill_tokens=None): model.eval() losses = [] - with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + with ( + _phase(model, prefill_tokens), + 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 + loss = model(input_ids=ids, labels=_labels(ids, prefill_tokens), use_cache=False).loss if not torch.isfinite(loss): raise RuntimeError("Non-finite evaluation loss") losses.append(float(loss.detach())) @@ -58,7 +79,7 @@ def _evaluate(model, blocks): "mean_nll": mean, "perplexity": math.exp(mean), "block_nll": losses, - "predicted_tokens": len(blocks) * (blocks.shape[1] - 1), + "predicted_tokens": len(blocks) * (blocks.shape[1] - max(1, prefill_tokens or 0)), } @@ -78,6 +99,9 @@ def main(): 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( + "--prefill-tokens", type=int, help="Explicit prefix; score/train suffix labels only" + ) parser.add_argument("--seed", type=int, default=2026) parser.add_argument("--learning-rate", type=float, default=1e-5) options = parser.parse_args() @@ -85,6 +109,8 @@ def main(): parser.error("length must be >=64, train-steps nonnegative, and eval-blocks positive") if options.train_data.resolve() == options.eval_data.resolve(): parser.error("Training and evaluation data must be distinct splits") + if options.prefill_tokens is not None and not 0 <= options.prefill_tokens <= options.length: + parser.error("prefill-tokens must be between zero and length") torch.manual_seed(options.seed) random.seed(options.seed) torch.set_float32_matmul_precision("highest") @@ -122,7 +148,7 @@ def main(): 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) + before = _evaluate(model, evaluation, options.prefill_tokens) 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() @@ -139,11 +165,15 @@ def main(): 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() + # Activation checkpoint recomputation needs the same explicit phase metadata. + with _phase(model, options.prefill_tokens): + with torch.autocast("cuda", dtype=torch.bfloat16): + loss = model( + input_ids=ids, labels=_labels(ids, options.prefill_tokens), 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) @@ -172,15 +202,24 @@ def main(): 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) if options.train_steps else before + after = _evaluate(model, evaluation, options.prefill_tokens) if options.train_steps else before 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", "solve.py"] + for name in [ + "config.py", + "kda.py", + "prefill.py", + "matmul.py", + "solve.py", + "decode.py", + "decode_prefill.py", + ] ] sources.extend( [ root / "modelopt/torch/kernels/quantization/linear_attention/neumann.py", + root / "modelopt/torch/kernels/quantization/linear_attention/decode.py", root / "modelopt/torch/quantization/plugins/kda.py", root / "modelopt/torch/quantization/plugins/linear_attention.py", Path(__file__).resolve(), @@ -202,6 +241,8 @@ def main(): "train_tokens_sha256": train_hash, "eval_tokens_sha256": eval_hash, "sequence_length": options.length, + "prefill_tokens": options.prefill_tokens, + "loss_scope": "full" if options.prefill_tokens is None else "suffix", "seed": options.seed, "train_order": order, "training": "KDA attention parameters only; FP32 AdamW, BF16 autocast, gradient checkpointing", @@ -216,7 +257,8 @@ def main(): "training_losses": losses, "gradient_norms": grad_norms, "step_seconds": timings, - "train_predicted_tokens": options.train_steps * options.length, + "train_predicted_tokens": options.train_steps + * (options.length + 1 - max(1, options.prefill_tokens or 0)), "peak_train_bytes": peak, "query_weight_changed": query_changed, "gpu": torch.cuda.get_device_name(), diff --git a/modelopt/torch/kernels/quantization/linear_attention/decode.py b/modelopt/torch/kernels/quantization/linear_attention/decode.py new file mode 100644 index 00000000000..9732ff7efda --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/decode.py @@ -0,0 +1,429 @@ +# 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. + +"""Fused FP32 token/replay training recurrence with state checkpoint recomputation.""" + +import torch +import triton +import triton.language as tl +from torch.autograd.function import once_differentiable +from triton.language.extra.cuda import libdevice + +__all__ = ["fused_recurrence"] + + +@triton.jit +def _fp8_qdq(value): + amax = tl.max(tl.abs(value)) + scale = tl.where(amax > 0, amax * (1.0 / 448.0), 1.0) + normalized = tl.div_rn(value, scale) + normalized = tl.where( + normalized < -448.0, -448.0, tl.where(normalized > 448.0, 448.0, normalized) + ) + bits = tl.abs(normalized).to(tl.int32, bitcast=True) + exponent = ((bits >> 23) & 255) - 127 + step = ((tl.maximum(exponent, -6) - 3 + 127) << 23).to(tl.float32, bitcast=True) + rounded = libdevice.nearbyint(normalized / step) * step + return rounded * scale, scale + + +@triton.jit +def _sum_keys(value): + BK: tl.constexpr = value.shape[0] + rows = tl.arange(0, BK) + for level in tl.static_range(0, tl.constexpr(triton.next_power_of_2(BK).bit_length() - 1)): + offset = BK >> (level + 1) + value = value + tl.gather( + value, tl.broadcast_to((rows ^ offset)[:, None], value.shape), axis=0 + ) + return tl.sum(tl.where(rows[:, None] == 0, value, 0.0), axis=0) + + +@triton.jit +def _advance( + state, + key, + value, + gate, + beta, + write_state, + CHANNEL_GATE: tl.constexpr, + STATE_QDQ: tl.constexpr, + FACTOR_QDQ: tl.constexpr, +): + decay = libdevice.exp(gate) + if CHANNEL_GATE: + decayed = state * decay[:, None] + else: + decayed = state * decay + residual = value - _sum_keys(key[:, None] * decayed) + update = beta * residual + update_scale = 1.0 + if FACTOR_QDQ: + update, update_scale = _fp8_qdq(update) + working = decayed + key[:, None] * update[None, :] + stored = working + state_scale = 1.0 + if STATE_QDQ: + if write_state: + stored, state_scale = _fp8_qdq(working) + return stored, working, decayed, residual, update, update_scale, state_scale + + +@triton.jit +def _decode_fwd( + Q, + K, + V, + G, + Beta, + Initial, + Out, + Final, + Updates, + Anchor, + AnchorScale, + UpdateScale, + Checkpoints, + T: tl.constexpr, + H: tl.constexpr, + DK: tl.constexpr, + DV: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NV: tl.constexpr, + CHANNEL_GATE: tl.constexpr, + STATE_QDQ: tl.constexpr, + REPLAY: tl.constexpr, + FACTOR_QDQ: tl.constexpr, + WINDOW: tl.constexpr, + CURSOR: tl.constexpr, + READ_STORED: tl.constexpr, + SCALE: tl.constexpr, + INTERVAL: tl.constexpr, +): + h = tl.program_id(0) + tile = tl.program_id(1) + keys = tl.arange(0, BK) + values = tile * BV + tl.arange(0, BV) + offsets = h * DK * DV + keys[:, None] * DV + values[None, :] + mask = (keys[:, None] < DK) & (values[None, :] < DV) + state = tl.load(Initial + offsets, mask, 0) + tl.store(Anchor + offsets, state, mask) + tl.store(AnchorScale + h * NV + tile, 1.0) + for t in range(T): + if t % INTERVAL == 0: + tl.store(Checkpoints + (t // INTERVAL) * H * DK * DV + offsets, state, mask) + query = tl.load(Q + (t * H + h) * DK + keys, keys < DK, 0) + key = tl.load(K + (t * H + h) * DK + keys, keys < DK, 0) + value = tl.load(V + (t * H + h) * DV + values, values < DV, 0) + beta = tl.load(Beta + t * H + h) + if CHANNEL_GATE: + gate = tl.load(G + (t * H + h) * DK + keys, keys < DK, 0) + else: + gate = tl.load(G + t * H + h) + write_state = not REPLAY or (CURSOR + t + 1) % WINDOW == 0 + state, working, decayed, residual, update, update_scale, state_scale = _advance( + state, + key, + value, + gate, + beta, + write_state, + CHANNEL_GATE, + STATE_QDQ, + FACTOR_QDQ, + ) + read = state if READ_STORED else working + output = _sum_keys(query[:, None] * read) * SCALE + tl.store(Out + (t * H + h) * DV + values, output, values < DV) + tl.store(Updates + (t * H + h) * DV + values, update, values < DV) + tl.store(UpdateScale + (t * H + h) * NV + tile, update_scale) + if write_state: + tl.store(Anchor + offsets, state, mask) + tl.store(AnchorScale + h * NV + tile, state_scale) + tl.store(Final + offsets, state, mask) + + +@triton.jit +def _decode_bwd( + Q, + K, + V, + G, + Beta, + Checkpoints, + DOut, + DFinal, + DUpdates, + DAnchor, + DQ, + DKOut, + DVOut, + DG, + DBeta, + DInitial, + T: tl.constexpr, + H: tl.constexpr, + DK: tl.constexpr, + DV: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NV: tl.constexpr, + CHANNEL_GATE: tl.constexpr, + STATE_QDQ: tl.constexpr, + REPLAY: tl.constexpr, + FACTOR_QDQ: tl.constexpr, + WINDOW: tl.constexpr, + CURSOR: tl.constexpr, + READ_STORED: tl.constexpr, + SCALE: tl.constexpr, + INTERVAL: tl.constexpr, +): + h = tl.program_id(0) + tile = tl.program_id(1) + keys = tl.arange(0, BK) + values = tile * BV + tl.arange(0, BV) + offsets = h * DK * DV + keys[:, None] * DV + values[None, :] + mask = (keys[:, None] < DK) & (values[None, :] < DV) + adjoint = tl.load(DFinal + offsets, mask, 0) + LAST_REFRESH: tl.constexpr = ((T + CURSOR) // WINDOW) * WINDOW - CURSOR - 1 + for t in range(T - 1, -1, -1): + start = (t // INTERVAL) * INTERVAL + state = tl.load(Checkpoints + (t // INTERVAL) * H * DK * DV + offsets, mask, 0) + for j in range(start, t): + key = tl.load(K + (j * H + h) * DK + keys, keys < DK, 0) + value = tl.load(V + (j * H + h) * DV + values, values < DV, 0) + beta = tl.load(Beta + j * H + h) + if CHANNEL_GATE: + gate = tl.load(G + (j * H + h) * DK + keys, keys < DK, 0) + else: + gate = tl.load(G + j * H + h) + write_state = not REPLAY or (CURSOR + j + 1) % WINDOW == 0 + state, _, _, _, _, _, _ = _advance( + state, + key, + value, + gate, + beta, + write_state, + CHANNEL_GATE, + STATE_QDQ, + FACTOR_QDQ, + ) + query = tl.load(Q + (t * H + h) * DK + keys, keys < DK, 0) + key = tl.load(K + (t * H + h) * DK + keys, keys < DK, 0) + value = tl.load(V + (t * H + h) * DV + values, values < DV, 0) + beta = tl.load(Beta + t * H + h) + if CHANNEL_GATE: + gate = tl.load(G + (t * H + h) * DK + keys, keys < DK, 0) + else: + gate = tl.load(G + t * H + h) + write_state = not REPLAY or (CURSOR + t + 1) % WINDOW == 0 + stored, working, decayed, residual, update, _, _ = _advance( + state, + key, + value, + gate, + beta, + write_state, + CHANNEL_GATE, + STATE_QDQ, + FACTOR_QDQ, + ) + if REPLAY: + if t == LAST_REFRESH: + adjoint += tl.load(DAnchor + offsets, mask, 0) + elif t == T - 1: + adjoint += tl.load(DAnchor + offsets, mask, 0) + doutput = tl.load(DOut + (t * H + h) * DV + values, values < DV, 0) + read = stored if READ_STORED else working + dquery = tl.sum(read * doutput[None, :], 1) * SCALE + dstate = adjoint + query[:, None] * (doutput[None, :] * SCALE) + dupdate = tl.sum(key[:, None] * dstate, 0) + dupdate += tl.load(DUpdates + (t * H + h) * DV + values, values < DV, 0) + dbeta = tl.sum(dupdate * residual, 0) + dvalue = beta * dupdate + dkey = tl.sum(update[None, :] * dstate - decayed * dvalue[None, :], 1) + ddecayed = dstate - key[:, None] * dvalue[None, :] + dgate = tl.sum(ddecayed * decayed, 1) + decay = libdevice.exp(gate) + if CHANNEL_GATE: + adjoint = ddecayed * decay[:, None] + tl.store(DG + ((tile * T + t) * H + h) * DK + keys, dgate, keys < DK) + else: + adjoint = ddecayed * decay + tl.store(DG + (tile * T + t) * H + h, tl.sum(dgate, 0)) + tl.store(DQ + ((tile * T + t) * H + h) * DK + keys, dquery, keys < DK) + tl.store(DKOut + ((tile * T + t) * H + h) * DK + keys, dkey, keys < DK) + tl.store(DVOut + (t * H + h) * DV + values, dvalue, values < DV) + tl.store(DBeta + (tile * T + t) * H + h, dbeta) + tl.store(DInitial + offsets, adjoint, mask) + + +class _FusedRecurrence(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q, + k, + v, + g, + beta, + initial, + state_qdq, + block_v, + replay, + factor_qdq, + window, + cursor, + read_stored, + scale, + interval, + ): + length, heads, keys = q.shape + values = v.shape[-1] + tiles = triton.cdiv(values, block_v) + checkpoints = torch.empty( + (triton.cdiv(length, interval), heads, keys, values), device=q.device, dtype=q.dtype + ) + output, updates = torch.empty_like(v), torch.empty_like(v) + final, anchor = torch.empty_like(initial), torch.empty_like(initial) + anchor_scales = torch.empty((heads, tiles), device=q.device, dtype=q.dtype) + update_scales = torch.empty((length, heads, tiles), device=q.device, dtype=q.dtype) + kwargs = { + "T": length, + "H": heads, + "DK": keys, + "DV": values, + "BK": triton.next_power_of_2(keys), + "BV": block_v, + "NV": tiles, + "CHANNEL_GATE": g.ndim == 3, + "STATE_QDQ": state_qdq, + "REPLAY": replay, + "FACTOR_QDQ": factor_qdq, + "WINDOW": window, + "CURSOR": cursor, + "READ_STORED": read_stored, + "SCALE": scale, + "INTERVAL": interval, + "num_warps": 4 if triton.next_power_of_2(keys) * block_v <= 4096 else 8, + "enable_fp_fusion": False, + } + _decode_fwd[(heads, tiles)]( + q, + k, + v, + g, + beta, + initial, + output, + final, + updates, + anchor, + anchor_scales, + update_scales, + checkpoints, + **kwargs, + ) + ctx.save_for_backward(q, k, v, g, beta, initial, checkpoints) + ctx.kwargs = kwargs + ctx.mark_non_differentiable(anchor_scales, update_scales) + return output, final, updates, anchor, anchor_scales, update_scales + + @staticmethod + @once_differentiable + def backward(ctx, doutput, dfinal, dupdates, danchor, _danchor_scales, _dupdate_scales): + q, k, v, g, beta, initial, checkpoints = ctx.saved_tensors + tiles = ctx.kwargs["NV"] + dq = torch.empty((tiles, *q.shape), device=q.device, dtype=q.dtype) + dk = torch.empty_like(dq) + dv = torch.empty_like(v) + dg = torch.empty((tiles, *g.shape), device=q.device, dtype=q.dtype) + dbeta = torch.empty((tiles, *beta.shape), device=q.device, dtype=q.dtype) + dinitial = torch.empty_like(initial) + doutput = torch.zeros_like(v) if doutput is None else doutput.contiguous() + dfinal = torch.zeros_like(initial) if dfinal is None else dfinal.contiguous() + dupdates = torch.zeros_like(v) if dupdates is None else dupdates.contiguous() + danchor = torch.zeros_like(initial) if danchor is None else danchor.contiguous() + _decode_bwd[(q.shape[1], tiles)]( + q, + k, + v, + g, + beta, + checkpoints, + doutput, + dfinal, + dupdates, + danchor, + dq, + dk, + dv, + dg, + dbeta, + dinitial, + **ctx.kwargs, + ) + return dq.sum(0), dk.sum(0), dv, dg.sum(0), dbeta.sum(0), dinitial, *([None] * 9) + + +def fused_recurrence( + q, + k, + v, + g, + beta, + initial, + *, + state_qdq=False, + block_v=16, + replay=False, + factor_qdq=False, + window=8, + cursor=0, + read_stored=True, + scale=None, + checkpoint_interval=8, +): + """Run a nonempty activated sequence; encoding keys/gates and carry metadata is external.""" + if not all( + x.is_cuda and x.dtype == torch.float32 and x.device == q.device + for x in (q, k, v, g, beta, initial) + ): + raise ValueError("Fused recurrence requires CUDA FP32 inputs on one device") + if q.ndim != 3 or q.shape != k.shape or v.shape[:2] != q.shape[:2] or beta.shape != q.shape[:2]: + raise ValueError("Expected aligned [T,H,D] q/k/v and [T,H] beta") + length, heads, keys = q.shape + if length == 0 or not 1 <= keys <= 128 or initial.shape != (heads, keys, v.shape[-1]): + raise ValueError("Require nonempty sequence, Dk<=128, and matching initial state") + if g.shape not in (q.shape, beta.shape) or block_v not in (16, 32, 64, 128): + raise ValueError("Invalid gate shape or state value-block grouping") + if not 1 <= window <= 64 or not 0 <= cursor < window or not 1 <= checkpoint_interval <= 32: + raise ValueError("Invalid replay cursor/window or recomputation interval") + scale = keys**-0.5 if scale is None else scale + return _FusedRecurrence.apply( + *(x.contiguous() for x in (q, k, v, g, beta, initial)), + state_qdq, + block_v, + replay, + factor_qdq, + window, + cursor, + read_stored, + scale, + checkpoint_interval, + ) diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py index 4007ce8dd6a..69b186ff196 100644 --- a/modelopt/torch/quantization/linear_attention/__init__.py +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -16,6 +16,8 @@ """Numerical policies and differentiable references for linear attention.""" from .config import * +from .decode import * +from .decode_prefill import * from .kda import * from .matmul import * from .prefill import * diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index bf2c22dec02..d3f73552d20 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -23,8 +23,10 @@ __all__ = [ "LinearAttentionConfig", + "LinearAttentionDecodeConfig", "LinearAttentionMatmulConfig", "LinearAttentionPolicyEntry", + "LinearAttentionReplayConfig", "LinearAttentionSolveConfig", ] @@ -89,6 +91,34 @@ def _validate_method(self): return self +class LinearAttentionReplayConfig(ModeloptBaseConfig): + """Anchor refresh and encoded rank-one update policy.""" + + window: int = Field(default=8, ge=1, le=64, strict=True) + factor_qdq: bool = ModeloptField(default=True) + encoding: Literal["once", "reencode"] = ModeloptField(default="once") + + +class LinearAttentionDecodeConfig(ModeloptBaseConfig): + """Explicit suffix recurrence; workload supplies per-sequence prefix lengths.""" + + mode: Literal["token", "replay"] = ModeloptField(default="token") + implementation: Literal["torch", "triton"] = ModeloptField(default="torch") + readout: Literal["working", "stored"] = ModeloptField(default="stored") + quantize_initial: bool = ModeloptField(default=True) + prefill_state_qdq: bool = ModeloptField(default=False) + decay_log_step: float | None = Field(default=None, gt=0, allow_inf_nan=False) + replay: LinearAttentionReplayConfig | None = ModeloptField(default=None) + + @model_validator(mode="after") + def _validate_replay(self): + if (self.mode == "replay") != (self.replay is not None): + raise ValueError("replay settings must be supplied exactly when mode='replay'") + if self.implementation == "triton" and self.replay and self.replay.encoding != "once": + raise ValueError("The Triton replay candidate implements encode-once only") + return self + + class LinearAttentionConfig(ModeloptBaseConfig): """GDN/KDA chunk-64 policy; unsupported numerical modes fail config validation. @@ -102,13 +132,17 @@ class LinearAttentionConfig(ModeloptBaseConfig): chunk_size: Literal[64] = ModeloptField(default=64) state: _StateConfig = ModeloptField(default=_StateConfig()) solve: LinearAttentionSolveConfig = ModeloptField(default=LinearAttentionSolveConfig()) + decode: LinearAttentionDecodeConfig | None = ModeloptField(default=None) matmul: dict[_PrefillSite, LinearAttentionMatmulConfig] = ModeloptField(default={}) elementwise: dict[_ElementwiseSite, _ArithmeticDtype] = ModeloptField(default={}) @model_validator(mode="after") def _validate_arithmetic_backend(self): if self.backend == "fla" and ( - self.matmul or self.elementwise or self.solve.method != "exact" + self.matmul + or self.elementwise + or self.solve.method != "exact" + or self.decode is not None ): raise ValueError("Prefill arithmetic policies require backend='matmul'") return self diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py new file mode 100644 index 00000000000..b037abaccfc --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -0,0 +1,352 @@ +# 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. + +"""Explicit token-state and encoded-update replay references for QAT.""" + +from dataclasses import dataclass + +import torch + +from .config import LinearAttentionDecodeConfig + +__all__ = [ + "EncodedLinearAttentionTensor", + "LinearAttentionCarry", + "ReplayEntry", + "recurrent_decode", + "recurrent_decode_reference", +] + + +@dataclass +class EncodedLinearAttentionTensor: + """Fake decoded values plus detached scale metadata; no compressed storage claim.""" + + values: torch.Tensor + scales: torch.Tensor | None + format: str + block_v: int | None + + +@dataclass +class ReplayEntry: + """An already computed, encoded rank-one update and its log retention.""" + + key: EncodedLinearAttentionTensor + update: EncodedLinearAttentionTensor + log_retention: torch.Tensor + + +@dataclass +class LinearAttentionCarry: + """Explicit anchor/update state, position, and codec contract across calls.""" + + anchor: EncodedLinearAttentionTensor + entries: tuple[ReplayEntry, ...] + position: int + started: bool + signature: str + + @property + def cursor(self): + """Number of encoded updates since the last anchor refresh.""" + return len(self.entries) + + def reconstruct(self): + """Replay stored entries in order, preserving all anchor/update gradients.""" + state = self.anchor.values + for entry in self.entries: + gate = entry.log_retention + decay = gate.exp().unsqueeze(-1) + if gate.ndim == 1: + decay = decay.unsqueeze(-1) + state = state * decay + entry.key.values.unsqueeze(-1) * entry.update.values.unsqueeze( + -2 + ) + return state + + +class _IdentityGradient(torch.autograd.Function): + @staticmethod + def forward(ctx, value, rounded): + return rounded + + @staticmethod + def backward(ctx, gradient): + return gradient, None + + +def _encode(value, enabled, block_v, *, state=False): + if not enabled: + return EncodedLinearAttentionTensor(value, None, "identity", None) + rounded, scales = [], [] + with torch.no_grad(): + for part in value.float().split(block_v, dim=-1): + axes = (-2, -1) if state else (-1,) + amax = part.abs().amax(dim=axes, keepdim=True) + scale = torch.where(amax > 0, amax * (1.0 / 448.0), torch.ones_like(amax)) + rounded.append((part / scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() * scale) + scales.append(scale[..., 0, 0] if state else scale[..., 0]) + decoded = torch.cat(rounded, dim=-1).to(value.dtype) + metadata = torch.stack(scales, dim=-1) + return EncodedLinearAttentionTensor( + _IdentityGradient.apply(value, decoded), metadata, "fp8_e4m3", block_v + ) + + +def _signature(config, state_qdq, block_v): + return ( + config.model_dump_json(exclude={"implementation", "prefill_state_qdq"}) + + f"/{state_qdq}/{block_v}" + ) + + +def _sum_keys(value): + keys = value.shape[-2] + padded = 1 << (keys - 1).bit_length() + if keys != padded: + value = torch.nn.functional.pad(value, (0, 0, 0, padded - keys)) + while value.shape[-2] > 1: + half = value.shape[-2] // 2 + value = value[..., :half, :] + value[..., half:, :] + return value[..., 0, :] + + +def _round_log_gate(gate, step): + if step is None: + return gate + rounded = (gate / step).round() * step + return _IdentityGradient.apply(gate, rounded.detach()) + + +def _prepare_carry(q, k, v, g, beta, config, state_qdq, block_v, initial_state, carry, position): + if q.ndim != 3 or k.shape != q.shape or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k/v must have aligned [T,H,D] shapes") + if beta.shape != q.shape[:2] or g.shape not in (beta.shape, q.shape): + raise ValueError("beta must be [T,H]; g must be [T,H] or [T,H,Dk]") + if block_v not in (16, 32, 64, 128): + raise ValueError("block_v must be 16, 32, 64, or 128") + signature = _signature(config, state_qdq, block_v) + if carry is not None and initial_state is not None: + raise ValueError("Supply either carry or initial_state") + shape = (q.shape[1], q.shape[2], v.shape[2]) + if carry is None: + initial_state = q.new_zeros(shape) if initial_state is None else initial_state + carry = LinearAttentionCarry( + _encode(initial_state, False, block_v, state=True), (), position, False, signature + ) + if carry.signature != signature or carry.anchor.values.shape != shape: + raise ValueError("Carry policy or state shape does not match this recurrence") + if config.mode == "token" and carry.entries: + raise ValueError("Token carry cannot contain replay entries") + if config.replay is not None and carry.cursor >= config.replay.window: + raise ValueError("Replay cursor must be below the refresh window") + if len(q) and not carry.started: + anchor = _encode( + carry.anchor.values, state_qdq and config.quantize_initial, block_v, state=True + ) + carry = LinearAttentionCarry(anchor, (), carry.position, True, signature) + return carry, signature + + +def recurrent_decode_reference( + q, + k, + v, + g, + beta, + *, + config: LinearAttentionDecodeConfig, + state_qdq=False, + block_v=64, + initial_state=None, + carry=None, + position=0, + scale=None, +): + """Run one preactivated sequence [T,H,D] with explicit token/replay write events. + + Keys and value heads must already be aligned. Scalar GDN or per-key-channel + KDA log gates are accepted. Outputs and all returned carry values retain their + graphs. An empty call performs no write or initial-state quantization. + """ + carry, signature = _prepare_carry( + q, k, v, g, beta, config, state_qdq, block_v, initial_state, carry, position + ) + if len(q) == 0: + # Keep empty input gradients defined without introducing a state write. + zero = (q.sum() + k.sum() + v.sum() + g.sum() + beta.sum()) * 0 + output = v + zero + return output, carry + scale = q.shape[-1] ** -0.5 if scale is None else scale + outputs = [] + with torch.autocast(device_type=q.device.type, enabled=False): + for t in range(len(q)): + entries = carry.entries + if config.replay is not None and config.replay.encoding == "reencode": + entries = tuple( + ReplayEntry( + _encode(e.key.values, config.replay.factor_qdq, k.shape[-1]), + _encode(e.update.values, config.replay.factor_qdq, block_v), + e.log_retention, + ) + for e in entries + ) + carry = LinearAttentionCarry(carry.anchor, entries, carry.position, True, signature) + state = carry.reconstruct() + gate = _round_log_gate(g[t], config.decay_log_step) + decay = gate.exp().unsqueeze(-1) + if gate.ndim == 1: + decay = decay.unsqueeze(-1) + key = _encode(k[t], config.replay is not None and config.replay.factor_qdq, k.shape[-1]) + decayed = state * decay + residual = v[t] - _sum_keys(key.values.unsqueeze(-1) * decayed) + update = _encode( + beta[t].unsqueeze(-1) * residual, + config.replay is not None and config.replay.factor_qdq, + block_v, + ) + working = decayed + key.values.unsqueeze(-1) * update.values.unsqueeze(-2) + if config.mode == "token": + anchor = _encode(working, state_qdq, block_v, state=True) + next_carry = LinearAttentionCarry(anchor, (), carry.position + 1, True, signature) + stored = anchor.values + else: + assert config.replay is not None + entries = (*entries, ReplayEntry(key, update, gate)) + if len(entries) == config.replay.window: + anchor = _encode(working, state_qdq, block_v, state=True) + next_carry = LinearAttentionCarry( + anchor, (), carry.position + 1, True, signature + ) + stored = anchor.values + else: + next_carry = LinearAttentionCarry( + carry.anchor, entries, carry.position + 1, True, signature + ) + stored = working + read = working if config.readout == "working" else stored + outputs.append(_sum_keys(q[t].unsqueeze(-1) * read) * scale) + carry = next_carry + return torch.stack(outputs), carry + + +def recurrent_decode( + q, + k, + v, + g, + beta, + *, + config: LinearAttentionDecodeConfig, + state_qdq=False, + block_v=64, + initial_state=None, + carry=None, + position=0, + scale=None, + checkpoint_interval=8, +): + """Run the selected training implementation with explicit differentiable carry. + + The fused encode-once path carries the same reconstructed state incrementally + and exposes the anchor and encoded updates needed for continuation gradients. + """ + if config.implementation == "torch": + return recurrent_decode_reference( + q, + k, + v, + g, + beta, + config=config, + state_qdq=state_qdq, + block_v=block_v, + initial_state=initial_state, + carry=carry, + position=position, + scale=scale, + ) + carry, signature = _prepare_carry( + q, k, v, g, beta, config, state_qdq, block_v, initial_state, carry, position + ) + if len(q) == 0: + return v + (q.sum() + k.sum() + v.sum() + g.sum() + beta.sum()) * 0, carry + if config.replay is not None and config.replay.encoding != "once": + raise ValueError("The fused implementation supports encode-once replay only") + # Keep the CPU reference importable without the optional CUDA/Triton backend. + from modelopt.torch.kernels.quantization.linear_attention.decode import fused_recurrence + + replay = config.replay is not None + factor_qdq = config.replay is not None and config.replay.factor_qdq + key = _encode(k, factor_qdq, k.shape[-1]) + gate = _round_log_gate(g, config.decay_log_step) + window = config.replay.window if config.replay is not None else 1 + result = fused_recurrence( + q, + key.values, + v, + gate, + beta, + carry.reconstruct(), + state_qdq=state_qdq, + block_v=block_v, + replay=replay, + factor_qdq=factor_qdq, + window=window, + cursor=carry.cursor, + read_stored=config.readout == "stored", + scale=scale, + checkpoint_interval=checkpoint_interval, + ) + output, final, updates, last_anchor, anchor_scales, update_scales = result + state_format = "fp8_e4m3" if state_qdq else "identity" + if not replay: + anchor = EncodedLinearAttentionTensor( + final, + anchor_scales if state_qdq else None, + state_format, + block_v if state_qdq else None, + ) + return output, LinearAttentionCarry(anchor, (), carry.position + len(q), True, signature) + if carry.cursor + len(q) >= window: + anchor = EncodedLinearAttentionTensor( + last_anchor, + anchor_scales if state_qdq else None, + state_format, + block_v if state_qdq else None, + ) + start = ((carry.cursor + len(q)) // window) * window - carry.cursor + entries = () + else: + anchor, entries, start = carry.anchor, carry.entries, 0 + additions = [] + for t in range(start, len(q)): + encoded_key = EncodedLinearAttentionTensor( + key.values[t], + key.scales[t] if key.scales is not None else None, + key.format, + key.block_v, + ) + encoded_update = EncodedLinearAttentionTensor( + updates[t], + update_scales[t] if factor_qdq else None, + "fp8_e4m3" if factor_qdq else "identity", + block_v if factor_qdq else None, + ) + additions.append(ReplayEntry(encoded_key, encoded_update, gate[t])) + return output, LinearAttentionCarry( + anchor, (*entries, *additions), carry.position + len(q), True, signature + ) diff --git a/modelopt/torch/quantization/linear_attention/decode_prefill.py b/modelopt/torch/quantization/linear_attention/decode_prefill.py new file mode 100644 index 00000000000..51a751a44a5 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/decode_prefill.py @@ -0,0 +1,180 @@ +# 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 explicit per-sequence prefill/decode phase handoff.""" + +from contextlib import contextmanager +from itertools import pairwise + +import torch + +from .decode import recurrent_decode + +__all__ = ["linear_attention_training_phase"] + + +def _lengths(values): + if isinstance(values, torch.Tensor): + values = values.tolist() + values = tuple(values) + if any(type(value) is not int or value < 0 for value in values): + raise ValueError("Prefill lengths must be nonnegative integers") + return values + + +@contextmanager +def linear_attention_training_phase(model, prefill_lengths): + """Supply explicit sequence phases through forward and checkpointed backward. + + Runtime phase metadata is local to the converted layers and restored on exit. + Keep this context active through backward when activation checkpointing is used. + """ + # The plugin imports this numerical package during quantization initialization. + from ..plugins.linear_attention import _LinearAttentionQuantMixin + + lengths = _lengths(prefill_lengths) + layers = [ + m + for m in model.modules() + if isinstance(m, _LinearAttentionQuantMixin) + and m.linear_attention_config.decode is not None + ] + if not layers: + raise ValueError("The model has no converted decode-aware linear-attention layers") + previous = [getattr(m, "_linear_attention_prefill_lengths", None) for m in layers] + try: + for module in layers: + module._linear_attention_prefill_lengths = lengths + yield model + finally: + for module, original in zip(layers, previous): + module._linear_attention_prefill_lengths = original + + +def _decode_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, + output_dtype, + prefill_lengths, +): + # Prefill dispatch imports this wrapper; defer its shared core to avoid a cycle. + from .prefill import _matmul_prefill + + if prefill_lengths is None: + raise ValueError("Decode-aware training requires explicit per-sequence prefill lengths") + if q.ndim != 4 or k.shape != q.shape or v.ndim != 4 or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k and v must have compatible [B,T,H,D] shapes") + batch, length, key_heads, keys = q.shape + heads, values = v.shape[2:] + if batch < 1 or key_heads < 1 or heads % key_heads or beta.shape != (batch, length, heads): + raise ValueError("Invalid batch/head dimensions or beta shape") + if g.shape not in (beta.shape, (*beta.shape, keys)): + raise ValueError("Invalid GDN/KDA log-retention shape") + q, k = (x.repeat_interleave(heads // key_heads, dim=2) for x in (q, k)) + boundaries = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens + if boundaries is None: + sequences = [(b, 0, length) for b in range(batch)] + else: + bounds = _lengths(boundaries) + if ( + batch != 1 + or len(bounds) < 2 + or bounds[0] != 0 + or bounds[-1] != length + or any(a > b for a, b in pairwise(bounds)) + ): + raise ValueError( + "Packed boundaries must partition a batch of one, allowing empty entries" + ) + sequences = [(0, a, b) for a, b in pairwise(bounds)] + prefixes = _lengths(prefill_lengths) + if len(prefixes) != len(sequences) or any( + p > end - start for p, (_, start, end) in zip(prefixes, sequences) + ): + raise ValueError("Supply one valid prefill length per sequence") + if initial_state is None: + states = q.new_zeros(len(sequences), heads, keys, values) + else: + states = initial_state.transpose(-1, -2) if state_v_first else initial_state + states = states.to(q.dtype) + if states.shape != (len(sequences), heads, keys, values): + raise ValueError("Initial state shape does not match sequence/head dimensions") + active = [n for n, prefix in enumerate(prefixes) if prefix] + prefix_outputs, prefix_states = {}, {} + if active: + packed = [ + torch.cat( + [ + tensor[sequences[n][0], sequences[n][1] : sequences[n][1] + prefixes[n]] + for n in active + ] + ).unsqueeze(0) + for tensor in (q, k, v, g, beta) + ] + offsets = [0] + for n in active: + offsets.append(offsets[-1] + prefixes[n]) + output, final = _matmul_prefill( + *packed, + sites=sites, + policy=policy, + w_quantizer=w_quantizer, + state_qdq=state_qdq and policy.decode.prefill_state_qdq, + scale=scale, + initial_state=states[active], + output_final_state=True, + cu_seqlens=None, + cu_seqlens_cpu=torch.tensor(offsets), + state_v_first=False, + chunk_size=policy.chunk_size, + output_dtype=q.dtype, + ) + for index, n in enumerate(active): + prefix_outputs[n] = output[0, offsets[index] : offsets[index + 1]] + prefix_states[n] = final[index] + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + split = start + prefixes[n] + suffix, carry = recurrent_decode( + *(x[b, split:end] for x in (q, k, v, g, beta)), + config=policy.decode, + state_qdq=state_qdq, + block_v=policy.state.block_v, + initial_state=prefix_states.get(n, states[n]), + position=prefixes[n], + scale=scale, + ) + prefix = prefix_outputs.get(n, q.new_empty(0, heads, values)) + outputs.append(torch.cat((prefix, suffix))) + finals.append(carry.reconstruct()) + output = torch.stack(outputs) if boundaries is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + if state_v_first: + final = final.transpose(-1, -2) + return output.to(output_dtype), final if output_final_state else None diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py index 5cd733a558c..fc6923a108e 100644 --- a/modelopt/torch/quantization/linear_attention/kda.py +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -18,6 +18,7 @@ import torch import torch.nn.functional as F +from .decode_prefill import _decode_prefill from .matmul import _validate_operand_quantizer from .prefill import _matmul_prefill @@ -53,6 +54,7 @@ def matmul_kda( cp_context=None, disable_recompute=False, return_intermediate_states=False, + prefill_lengths=None, ): """Run KDA chunk prefill with actual-operand QDQ and differentiable state carry. @@ -90,6 +92,26 @@ def matmul_kda( 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) + if policy.decode is not None: + return _decode_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, + output_dtype=output_dtype, + prefill_lengths=prefill_lengths, + ) return _matmul_prefill( q, k, diff --git a/modelopt/torch/quantization/linear_attention/prefill.py b/modelopt/torch/quantization/linear_attention/prefill.py index bbe88dcc802..dff8243a13b 100644 --- a/modelopt/torch/quantization/linear_attention/prefill.py +++ b/modelopt/torch/quantization/linear_attention/prefill.py @@ -19,6 +19,7 @@ import torch.nn.functional as F from .config import LinearAttentionConfig +from .decode_prefill import _decode_prefill from .matmul import LinearAttentionMatmulSites, _validate_operand_quantizer from .reference import _prepare, state_fp8_qdq_reference from .solve import triangular_inverse @@ -51,6 +52,7 @@ def matmul_gdn( state_v_first=False, chunk_size=64, cp_context=None, + prefill_lengths=None, ): """Execute all eight GDN prefill sites with autograd through their QDQ operators. @@ -78,6 +80,26 @@ 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") + if policy.decode is not None: + return _decode_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, + output_dtype=output_dtype, + prefill_lengths=prefill_lengths, + ) return _matmul_prefill( q, k, diff --git a/modelopt/torch/quantization/plugins/gated_delta_net.py b/modelopt/torch/quantization/plugins/gated_delta_net.py index 856d6307ff6..228a1c450ca 100644 --- a/modelopt/torch/quantization/plugins/gated_delta_net.py +++ b/modelopt/torch/quantization/plugins/gated_delta_net.py @@ -95,6 +95,7 @@ def _state_quantized_chunk_gated_delta_rule( w_quantizer=self.gdn_w_quantizer, state_qdq=quantize_state, chunk_size=chunk_size, + prefill_lengths=self._linear_attention_prefill_lengths, **kwargs, ) return _state_qdq_chunk_gated_delta_rule()( diff --git a/modelopt/torch/quantization/plugins/kda.py b/modelopt/torch/quantization/plugins/kda.py index f972c501cc0..d355cb5955f 100644 --- a/modelopt/torch/quantization/plugins/kda.py +++ b/modelopt/torch/quantization/plugins/kda.py @@ -16,6 +16,8 @@ """QAT integration for the optional FLA KimiDeltaAttention layer.""" import importlib.metadata +import inspect +from functools import lru_cache from types import FunctionType from ..linear_attention.kda import matmul_kda @@ -26,6 +28,11 @@ __all__ = [] +@lru_cache(maxsize=8) +def _forward_signature(function): + return inspect.signature(function) + + class _QuantKimiDeltaAttention(_LinearAttentionQuantMixin): linear_attention_quantizer_names = ("kda_state_quantizer", "kda_w_quantizer") @@ -39,6 +46,12 @@ def forward(self, *args, **kwargs): if not self.linear_attention_is_enabled: return super().forward(*args, **kwargs) original = super().forward.__func__ + if self.linear_attention_config.decode is not None: + arguments = _forward_signature(original).bind(self, *args, **kwargs).arguments + if arguments.get("use_cache", False) or arguments.get("past_key_values") is not None: + raise NotImplementedError( + "Decode-aware QAT requires use_cache=False; use explicit numerical carry for continuation" + ) # Bind this invocation so copied modules cannot retain another module's quantizers. namespace = { **original.__globals__, @@ -62,10 +75,13 @@ def _quantized_chunk(self, *args, **kwargs): 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, + prefill_lengths=self._linear_attention_prefill_lengths, **kwargs, ) def _unsupported_recurrent(self, *args, **kwargs): + if self.linear_attention_config.decode is not None: + return self._quantized_chunk(*args, **kwargs) raise NotImplementedError( "KDA prefill QAT requires the chunk path; recurrent decode is not yet supported" ) diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py index 33a11c0c244..8aa893acd24 100644 --- a/modelopt/torch/quantization/plugins/linear_attention.py +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -32,6 +32,7 @@ def _setup(self): setattr(self, name, TensorQuantizer(QuantizerAttributeConfig(enable=False))) self.linear_attn_sites = LinearAttentionMatmulSites() self.linear_attention_config = LinearAttentionConfig() + self._linear_attention_prefill_lengths = None @property def _linear_attn_state(self): @@ -50,6 +51,7 @@ def linear_attention_is_enabled(self): or self.linear_attn_sites.is_enabled or bool(self.linear_attention_config.matmul or self.linear_attention_config.elementwise) or self.linear_attention_config.solve.method != "exact" + or self.linear_attention_config.decode is not None ) def validate_linear_attention(self): diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py new file mode 100644 index 00000000000..ab3138d9079 --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py @@ -0,0 +1,305 @@ +# 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 +import torch.nn.functional as F +import triton +import triton.language as tl + +from modelopt.torch.kernels.quantization.linear_attention.decode import _fp8_qdq, fused_recurrence +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + LinearAttentionMatmulSites, + matmul_gdn, + matmul_kda, +) +from modelopt.torch.quantization.linear_attention.config import LinearAttentionDecodeConfig +from modelopt.torch.quantization.linear_attention.decode import ( + _encode, + recurrent_decode, + recurrent_decode_reference, +) +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("state_qdq", [False, True]) +@pytest.mark.parametrize("replay", [False, True]) +@pytest.mark.parametrize("read_stored", [False, True]) +def test_fused_recurrence_outputs_state_and_gradients(kda, state_qdq, replay, read_stored): + torch.manual_seed(193) + shape = (37, 2, 16) + q, k = [ + F.normalize(torch.randn(shape, device="cuda"), dim=-1).requires_grad_() for _ in range(2) + ] + v = torch.randn(37, 2, 11, device="cuda", requires_grad=True) + g = (-torch.rand(shape if kda else shape[:-1], device="cuda") * 0.03).requires_grad_() + beta = (torch.rand(shape[:-1], device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(2, 16, 11, device="cuda") * 0.1).requires_grad_() + cfg = LinearAttentionDecodeConfig( + mode="replay" if replay else "token", + quantize_initial=False, + readout="stored" if read_stored else "working", + replay={"window": 5} if replay else None, + ) + encoded_key = _encode(k, replay, 16).values + actual = fused_recurrence( + q, + encoded_key, + v, + g, + beta, + initial, + state_qdq=state_qdq, + block_v=16, + replay=replay, + factor_qdq=replay, + window=5, + read_stored=read_stored, + ) + expected, carry = recurrent_decode_reference( + q, k, v, g, beta, config=cfg, initial_state=initial, state_qdq=state_qdq, block_v=16 + ) + final = carry.reconstruct() + targets = (q, k, v, g, beta, initial) + probes = [torch.randn_like(x) for x in actual[:2]] + gradients = [ + torch.autograd.grad(sum((a * p).sum() for a, p in zip(result, probes)), targets) + for result in (actual[:2], (expected, final)) + ] + for a, e in zip((*actual[:2], *gradients[0]), (expected, final, *gradients[1])): + torch.testing.assert_close(a, e, rtol=1e-3, atol=8e-6) + + +@pytest.mark.parametrize("mode", ["token", "replay"]) +@pytest.mark.parametrize("interval", [3, 8]) +def test_fused_carry_handoff_and_auxiliary_gradients(mode, interval): + torch.manual_seed(217) + q, k = [ + F.normalize(torch.randn(23, 2, 16, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(23, 2, 19, device="cuda", requires_grad=True) + g = (-torch.rand_like(q) * 0.05).requires_grad_() + beta = (torch.rand(23, 2, device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(2, 16, 19, device="cuda") * 0.2).requires_grad_() + args = (q, k, v, g, beta) + reference = LinearAttentionDecodeConfig( + mode=mode, decay_log_step=0.01, replay={"window": 5} if mode == "replay" else None + ) + candidate = reference.model_copy(update={"implementation": "triton"}) + expected, expected_carry = recurrent_decode_reference( + *args, config=reference, state_qdq=True, block_v=16, initial_state=initial + ) + carry = None + outputs = [] + start = 0 + for end in [0, 3, 11, 12, 23]: + output, carry = recurrent_decode( + *(x[start:end] for x in args), + config=candidate, + state_qdq=True, + block_v=16, + initial_state=initial if carry is None else None, + carry=carry, + checkpoint_interval=interval, + ) + outputs.append(output) + start = end + actual = torch.cat(outputs) + actual_state, expected_state = carry.reconstruct(), expected_carry.reconstruct() + assert carry.cursor == expected_carry.cursor + assert carry.position == 23 and carry.started + torch.testing.assert_close( + carry.anchor.scales, expected_carry.anchor.scales, rtol=5e-5, atol=2e-8 + ) + for a, e in zip(carry.entries, expected_carry.entries): + torch.testing.assert_close(a.update.scales, e.update.scales, rtol=5e-5, atol=2e-8) + torch.testing.assert_close(a.key.values, e.key.values, rtol=0, atol=0) + probes = [torch.randn_like(actual), torch.randn_like(actual_state)] + gradients = [ + torch.autograd.grad(sum((a * p).sum() for a, p in zip(result, probes)), (*args, initial)) + for result in ((actual, actual_state), (expected, expected_state)) + ] + for a, e in zip( + (actual, actual_state, *gradients[0]), (expected, expected_state, *gradients[1]) + ): + torch.testing.assert_close(a, e, rtol=1e-3, atol=1e-5) + + +@triton.jit +def _codec_kernel(X, Y, S, N: tl.constexpr, BLOCK: tl.constexpr): # noqa: N803 + offsets = tl.arange(0, BLOCK) + values = tl.load(X + offsets, offsets < N, 0) + decoded, scale = _fp8_qdq(values) + tl.store(Y + offsets, decoded, offsets < N) + tl.store(S, scale) + + +@pytest.mark.parametrize("scale", [0.0, 2.0**-12, 1.0, 2.0**12, 0.00314159, 731.173]) +def test_software_codec_matches_float8_boundaries(scale): + # Exhaust all finite positive E4M3 values and their rounding ties, plus signs. + representable = torch.arange(127, dtype=torch.uint8).view(torch.float8_e4m3fn).float() + ties = (representable[:-1] + representable[1:]) / 2 + probe = torch.cat( + ( + representable, + ties, + torch.nextafter(ties, torch.full_like(ties, float("inf"))), + torch.nextafter(ties, torch.zeros_like(ties)), + ) + ) + probe = torch.cat((probe, -probe)) * scale + values = probe.cuda() + actual = torch.empty_like(values) + actual_scale = torch.empty((), device="cuda") + _codec_kernel[(1,)]( + values, actual, actual_scale, len(values), triton.next_power_of_2(len(values)) + ) + expected = _encode(values, True, len(values)) + torch.testing.assert_close(actual, expected.values, rtol=0, atol=0) + torch.testing.assert_close(actual_scale, expected.scales[0], rtol=0, atol=0) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("mode", ["token", "replay"]) +@pytest.mark.parametrize("horizon", ["near_unit", "mixed"]) +def test_wide_long_trajectory(kda, mode, horizon): + torch.manual_seed(329) + q, k = [ + F.normalize(torch.randn(257, 2, 128, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(257, 2, 70, device="cuda", requires_grad=True) + shape = q.shape if kda else q.shape[:-1] + g = -torch.rand(shape, device="cuda") * (1e-4 if horizon == "near_unit" else 3) + if horizon == "mixed": + g[::7] = -30 + g.requires_grad_() + beta = (torch.rand(257, 2, device="cuda") * 0.6).requires_grad_() + initial = (torch.randn(2, 128, 70, device="cuda") * 0.1).requires_grad_() + cfg = LinearAttentionDecodeConfig(mode=mode, replay={"window": 8} if mode == "replay" else None) + args = q, k, v, g, beta + expected, ec = recurrent_decode_reference( + *args, config=cfg, state_qdq=True, block_v=64, initial_state=initial + ) + actual, ac = recurrent_decode( + *args, + config=cfg.model_copy(update={"implementation": "triton"}), + state_qdq=True, + block_v=64, + initial_state=initial, + ) + probes = [torch.randn_like(actual), torch.randn_like(ac.reconstruct())] + gradients = [ + torch.autograd.grad(sum((a * p).sum() for a, p in zip(result, probes)), (*args, initial)) + for result in ((actual, ac.reconstruct()), (expected, ec.reconstruct())) + ] + # Fixed reduction and scale arithmetic must also agree over long trajectories. + for a, e in zip( + (actual, ac.reconstruct(), *gradients[0]), (expected, ec.reconstruct(), *gradients[1]) + ): + relative = (a - e).norm() / e.norm().clamp_min(1e-8) + assert relative < 1e-4, relative.item() + assert torch.isfinite(a).all() + + +@pytest.mark.parametrize("block_v", [32, 128]) +@pytest.mark.parametrize("kda", [False, True]) +def test_full_key_width_and_value_block_tail(block_v, kda): + torch.manual_seed(418) + q, k = [ + F.normalize(torch.randn(9, 2, 128, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(9, 2, 133, device="cuda", requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1], device="cuda") * 0.1).requires_grad_() + beta = torch.rand(9, 2, device="cuda", requires_grad=True) + initial = torch.randn(2, 128, 133, device="cuda", requires_grad=True) + args = q, k, v, g, beta + cfg = LinearAttentionDecodeConfig(mode="replay", replay={"window": 3}) + results = [] + for implementation in ("torch", "triton"): + output, carry = recurrent_decode( + *args, + config=cfg.model_copy(update={"implementation": implementation}), + state_qdq=True, + block_v=block_v, + initial_state=initial, + ) + final = carry.reconstruct() + gradients = torch.autograd.grad( + output.square().sum() + final.square().sum(), (*args, initial) + ) + results.append((output, final, *gradients)) + for actual, expected in zip(*results): + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=2e-5) + + +@pytest.mark.parametrize("kda", [False, True]) +def test_packed_quantized_prefix_and_replay_composition(kda): + torch.manual_seed(591) + q, k = [ + F.normalize(torch.randn(1, 73, 1, 16, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(1, 73, 2, 19, device="cuda", requires_grad=True) + g = (-torch.rand((1, 73, 2, 16) if kda else (1, 73, 2), device="cuda") * 0.03).requires_grad_() + beta = (torch.rand(1, 73, 2, device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(3, 2, 19, 16, device="cuda") * 0.1).requires_grad_() + sites = LinearAttentionMatmulSites() + attributes = {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)} + for quantizer in sites.modules(): + if isinstance(quantizer, TensorQuantizer): + quantizer.set_from_attribute_config(attributes) + quantizer.enable() + w = TensorQuantizer(QuantizerAttributeConfig(**attributes)) + outputs = [] + function = matmul_kda if kda else matmul_gdn + args = q, k, v, g, beta + for implementation in ("torch", "triton"): + policy = LinearAttentionConfig( + backend="matmul", + state={"block_v": 16}, + solve={"method": "neumann", "degree": 3}, + elementwise={"value_residual": "bfloat16"}, + decode={ + "mode": "replay", + "implementation": implementation, + "prefill_state_qdq": True, + "decay_log_step": 1 / 256, + "replay": {"window": 5}, + }, + ) + output, final = function( + *args, + initial_state=initial, + state_v_first=True, + cu_seqlens=torch.tensor([0, 0, 9, 73]), + prefill_lengths=[0, 3, 33], + output_final_state=True, + sites=sites, + w_quantizer=w, + policy=policy, + state_qdq=True, + ) + gradients = torch.autograd.grad( + output.square().sum() + final.square().sum(), (*args, initial) + ) + outputs.append((output, final, *gradients)) + for actual, expected in zip(*outputs): + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=2e-5) diff --git a/tests/gpu/torch/quantization/plugins/test_kda.py b/tests/gpu/torch/quantization/plugins/test_kda.py index 36de6566d02..d0396c8469a 100644 --- a/tests/gpu/torch/quantization/plugins/test_kda.py +++ b/tests/gpu/torch/quantization/plugins/test_kda.py @@ -14,13 +14,18 @@ # limitations under the License. import copy +from contextlib import nullcontext import pytest import torch +from torch.utils.checkpoint import checkpoint import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq -from modelopt.torch.quantization.linear_attention import LinearAttentionConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + linear_attention_training_phase, +) from modelopt.torch.quantization.nn import TensorQuantizer KimiDeltaAttention = pytest.importorskip("fla.layers.kda").KimiDeltaAttention @@ -35,12 +40,29 @@ def _layer(): def _forward(model, hidden): - with torch.autocast("cuda", dtype=torch.bfloat16): + policy = getattr(model, "linear_attention_config", None) + phase = ( + linear_attention_training_phase(model, [min(31, hidden.shape[1])]) + if policy is not None and policy.decode is not None + else nullcontext() + ) + with phase, torch.autocast("cuda", dtype=torch.bfloat16): return model(hidden)[0] @pytest.mark.parametrize( - "mode", ["state", "w", "prefill_fp8", "prefill_nvfp4", "arithmetic", "solve"] + "mode", + [ + "state", + "w", + "prefill_fp8", + "prefill_nvfp4", + "arithmetic", + "solve", + "decode_token", + "decode_replay", + "decode_decay", + ], ) @pytest.mark.timeout(180) def test_fla_layer_qat_restore_and_optimizer(tmp_path, mode): @@ -77,6 +99,20 @@ def test_fla_layer_qat_restore_and_optimizer(tmp_path, mode): "degree": 3, "implementation": "triton", } + if mode.startswith("decode"): + cfg["quant_cfg"].append( + {"quantizer_name": "*kda_state_quantizer", "cfg": {**attributes, "axis": (0, 1)}} + ) + policy_cfg = cfg["linear_attention"][0]["cfg"] + policy_cfg["state"] = {"block_v": 16} + policy_cfg["decode"] = { + "mode": "replay" if mode == "decode_replay" else "token", + "implementation": "triton", + } + if mode == "decode_replay": + policy_cfg["decode"]["replay"] = {"window": 5} + if mode == "decode_decay": + policy_cfg["decode"]["decay_log_step"] = 1 / 256 mtq.quantize(model, cfg) output = _forward(model, hidden) assert torch.isfinite(output).all() @@ -113,5 +149,49 @@ def test_fla_layer_qat_restore_and_optimizer(tmp_path, mode): 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]) + if mode.startswith("decode"): + with ( + linear_attention_training_phase(restored, [0]), + torch.autocast("cuda", dtype=torch.bfloat16), + ): + assert torch.isfinite(restored(hidden[:, :1])[0]).all() + with pytest.raises(NotImplementedError, match="use_cache=False"): + restored(hidden[:, :1], use_cache=True) + with ( + pytest.raises(ValueError, match="explicit"), + torch.autocast("cuda", dtype=torch.bfloat16), + ): + restored(hidden) + else: + with pytest.raises(NotImplementedError, match="chunk path"): + _forward(restored, hidden[:, :1]) + + +def test_decode_phase_spans_activation_checkpoint_backward(): + model = _layer() + mtq.quantize( + model, + { + "quant_cfg": [{"quantizer_name": "*", "enable": False}], + "algorithm": None, + "linear_attention": [ + { + "module_name": "*", + "cfg": {"backend": "matmul", "decode": {"implementation": "triton"}}, + } + ], + }, + ) + hidden = torch.randn(1, 73, 32, device="cuda", requires_grad=True) + with linear_attention_training_phase(model, [31]): + with torch.autocast("cuda", dtype=torch.bfloat16): + actual = checkpoint(lambda x: model(x)[0], hidden, use_reentrant=False) + actual.square().mean().backward() + assert model._linear_attention_prefill_lengths is None + assert torch.isfinite(hidden.grad).all() + assert torch.count_nonzero(model.q_proj.weight.grad) > 0 + with linear_attention_training_phase(model, [0]): + with linear_attention_training_phase(model, [73]): + assert model._linear_attention_prefill_lengths == (73,) + assert model._linear_attention_prefill_lengths == (0,) + assert model._linear_attention_prefill_lengths is None diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index 9739a2da391..21dcebaf826 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import nullcontext + import pytest import torch from _test_utils.torch.megatron.models import get_mcore_gpt_model @@ -23,7 +25,10 @@ ) import modelopt.torch.quantization as mtq -from modelopt.torch.quantization.linear_attention import LinearAttentionConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + linear_attention_training_phase, +) from modelopt.torch.quantization.nn import TensorQuantizer pytest.importorskip("fla") # Megatron-Core GatedDeltaNet and the state QDQ kernel need fla @@ -90,7 +95,14 @@ def _config_for_mode(mode): sites = ("state", "w") if mode == "both" else ((mode,) if mode in ("state", "w") else ()) if mode in ("prefill_fp8", "prefill_nvfp4"): sites = ("w",) + if mode.startswith("decode"): + sites = ("state",) cfg = _gdn_config(sites) + if mode.startswith("decode"): + policy = cfg["linear_attention"][0]["cfg"] + policy.update(backend="matmul", decode={"implementation": "triton"}) + if mode == "decode_replay": + policy["decode"].update(mode="replay", replay={"window": 8}) if mode.startswith("prefill"): policy = cfg["linear_attention"][0]["cfg"] policy["backend"] = "matmul" @@ -119,7 +131,16 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED ) model = _make_model(size) - forward = get_forward(model) + original_forward = get_forward(model) + + def forward(m): + enabled = any( + getattr(getattr(layer, "linear_attention_config", None), "decode", None) is not None + for layer in m.modules() + ) + with linear_attention_training_phase(m, [64, 64]) if enabled else nullcontext(): + return original_forward(m) + outputs = [] handles = [ module.register_forward_hook( @@ -200,7 +221,17 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): @pytest.mark.parametrize("tp_size", [1, 2]) @pytest.mark.parametrize( "mode", - ["state", "w", "both", "prefill_fp8", "prefill_nvfp4", "prefill_arithmetic", "prefill_solve"], + [ + "state", + "w", + "both", + "prefill_fp8", + "prefill_nvfp4", + "prefill_arithmetic", + "prefill_solve", + "decode_token", + "decode_replay", + ], ) def test_gdn_qat_and_sharded_restore(request, tmp_path, tp_size, mode): """Train through state/W QDQ after a Megatron distributed-checkpoint round trip.""" @@ -218,7 +249,16 @@ def _test_gdn_context_parallel_helper(rank, size, mode): @pytest.mark.parametrize( - "mode", ["state", "w", "prefill_fp8", "prefill_arithmetic", "prefill_solve"] + "mode", + [ + "state", + "w", + "prefill_fp8", + "prefill_arithmetic", + "prefill_solve", + "decode_token", + "decode_replay", + ], ) def test_gdn_context_parallel_rejected(dist_workers_size_2, mode): """Reject unqualified Megatron CP even when it does not pass an FLA CP context.""" diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py new file mode 100644 index 00000000000..f1daed9856b --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -0,0 +1,266 @@ +# 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 itertools + +import pytest +import torch +import torch.nn.functional as F + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + LinearAttentionMatmulSites, + matmul_gdn, + matmul_kda, +) +from modelopt.torch.quantization.linear_attention.config import LinearAttentionDecodeConfig +from modelopt.torch.quantization.linear_attention.decode import recurrent_decode_reference +from modelopt.torch.quantization.linear_attention.reference import recurrent_delta_rule_reference +from modelopt.torch.quantization.nn import TensorQuantizer + + +def _inputs(kda=True, length=73): + torch.manual_seed(193) + q, k = [ + F.normalize(torch.randn(length, 2, 16, dtype=torch.float64), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(length, 2, 11, dtype=torch.float64, requires_grad=True) + g = ( + -torch.rand((length, 2, 16) if kda else (length, 2), dtype=torch.float64) * 0.03 + ).requires_grad_() + beta = (torch.rand(length, 2, dtype=torch.float64) * 0.4).requires_grad_() + state = (torch.randn(2, 16, 11, dtype=torch.float64) * 0.1).requires_grad_() + return (q, k, v, g, beta), state + + +def _values_and_grads(output, state, args, initial): + loss = output.square().sum() + state.square().sum() + return output, state, *torch.autograd.grad(loss, (*args, initial), retain_graph=True) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("mode", ["token", "replay"]) +def test_unquantized_decode_recurrence_and_gradients(kda, mode): + args, state = _inputs(kda) + cfg = LinearAttentionDecodeConfig( + mode=mode, replay={"window": 8, "factor_qdq": False} if mode == "replay" else None + ) + output, carry = recurrent_decode_reference(*args, config=cfg, initial_state=state) + exact, final = recurrent_delta_rule_reference( + *(x.unsqueeze(0) for x in args), initial_state=state.unsqueeze(0) + ) + for a, b in zip( + _values_and_grads(output, carry.reconstruct(), args, state), + _values_and_grads(exact[0], final[0], args, state), + ): + torch.testing.assert_close(a, b, rtol=2e-11, atol=1e-12) + assert carry.position == 73 + assert carry.cursor == (1 if mode == "replay" else 0) + + +@pytest.mark.parametrize("mode", ["token", "replay", "reencode"]) +@pytest.mark.parametrize("readout", ["working", "stored"]) +def test_quantized_carry_continuation_metadata_and_gradients(mode, readout): + args, state = _inputs() + cfg = LinearAttentionDecodeConfig( + mode="token" if mode == "token" else "replay", + readout=readout, + replay={"window": 8, "encoding": "reencode" if mode == "reencode" else "once"} + if mode != "token" + else None, + ) + expected, final = recurrent_decode_reference( + *args, config=cfg, initial_state=state, state_qdq=True, block_v=16 + ) + chunks = [] + carry = None + lo = 0 + for hi in [0, 5, 16, 16, 37, 73]: + result, carry = recurrent_decode_reference( + *(x[lo:hi] for x in args), + config=cfg, + carry=carry, + initial_state=state if carry is None else None, + state_qdq=True, + block_v=16, + ) + chunks.append(result) + lo = hi + actual = torch.cat(chunks) + for a, b in zip( + _values_and_grads(actual, carry.reconstruct(), args, state), + _values_and_grads(expected, final.reconstruct(), args, state), + ): + torch.testing.assert_close(a, b, rtol=0, atol=0) + assert carry.position == 73 + assert carry.anchor.format == "fp8_e4m3" + assert carry.anchor.scales.shape == (2, 1) + assert not carry.anchor.scales.requires_grad + for entry in carry.entries: + assert entry.key.scales.shape == (2, 1) + assert entry.update.values.requires_grad + assert not entry.update.scales.requires_grad + + +def test_initial_write_and_readout_order_are_explicit(): + args, state = _inputs(length=1) + outputs = [] + for readout in ["working", "stored"]: + output, carry = recurrent_decode_reference( + *args, + config=LinearAttentionDecodeConfig(readout=readout), + initial_state=state, + state_qdq=True, + block_v=16, + ) + outputs.append(output) + assert not torch.equal(*outputs) + empty, untouched = recurrent_decode_reference( + *(x[:0] for x in args), + config=LinearAttentionDecodeConfig(), + initial_state=state, + state_qdq=True, + ) + assert empty.shape == (0, 2, 11) + assert untouched.anchor.values is state + assert not untouched.started + + +def test_grid_gate_ste_keeps_gate_gradients_and_changes_trajectory(): + args, state = _inputs() + output, carry = recurrent_decode_reference( + *args, config=LinearAttentionDecodeConfig(decay_log_step=0.02), initial_state=state + ) + grad = torch.autograd.grad(output.square().sum() + carry.reconstruct().square().sum(), args[3])[ + 0 + ] + assert torch.isfinite(grad).all() and torch.count_nonzero(grad) > 0 + exact, _ = recurrent_decode_reference( + *args, config=LinearAttentionDecodeConfig(), initial_state=state + ) + assert not torch.equal(output, exact) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("prefixes", [[0, 0, 0, 0], [0, 3, 65, 0], [0, 5, 68, 0]]) +def test_explicit_packed_phase_with_empty_entries(kda, prefixes): + args, one_state = _inputs(kda) + args = tuple(x.unsqueeze(0) for x in args) + initial = torch.stack([one_state.detach().clone() for _ in range(4)]).requires_grad_() + boundaries = [0, 0, 5, 73, 73] + function = matmul_kda if kda else matmul_gdn + actual = function( + *args, + sites=LinearAttentionMatmulSites(), + policy=LinearAttentionConfig(backend="matmul", decode={}), + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(enable=False)), + prefill_lengths=prefixes, + cu_seqlens=torch.tensor(boundaries), + initial_state=initial, + output_final_state=True, + ) + outputs, states = [], [] + for n, (lo, hi) in enumerate(itertools.pairwise(boundaries)): + if lo == hi: + outputs.append(args[2][0, lo:hi]) + states.append(initial[n]) + continue + o, s = recurrent_delta_rule_reference( + *(x[:, lo:hi] for x in args), initial_state=initial[n : n + 1] + ) + outputs.append(o[0]) + states.append(s[0]) + expected = torch.cat(outputs).unsqueeze(0), torch.stack(states) + for a, e in zip(actual, expected): + torch.testing.assert_close(a, e, rtol=2e-11, atol=1e-12) + gradients = [ + torch.autograd.grad( + sum(x.square().sum() for x in result), (*args, initial), retain_graph=True + ) + for result in (actual, expected) + ] + for a, e in zip(*gradients): + torch.testing.assert_close(a, e, rtol=2e-10, atol=1e-11) + + +@pytest.mark.parametrize("mode", ["token", "replay"]) +def test_carry_rejects_changed_policy_and_shape(mode): + args, state = _inputs(length=4) + cfg = LinearAttentionDecodeConfig(mode=mode, replay={"window": 8} if mode == "replay" else None) + _, carry = recurrent_decode_reference(*args, config=cfg, initial_state=state, state_qdq=True) + for changed in [ + cfg.model_copy(update={"readout": "working"}), + cfg.model_copy(update={"decay_log_step": 0.01}), + ]: + with pytest.raises(ValueError, match="policy or state shape"): + recurrent_decode_reference(*args, config=changed, carry=carry, state_qdq=True) + with pytest.raises(ValueError, match="policy or state shape"): + recurrent_decode_reference(*args, config=cfg, carry=carry, state_qdq=True, block_v=16) + with pytest.raises(ValueError, match="either carry or initial"): + recurrent_decode_reference( + *args, config=cfg, carry=carry, initial_state=state, state_qdq=True + ) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("all_empty", [False, True]) +def test_grouped_heads_transposed_state_and_empty_batch(kda, all_empty): + args, state = _inputs(kda, length=0 if all_empty else 9) + q, k, v, g, beta = [x.unsqueeze(0).repeat(2, *([1] * x.ndim)) for x in args] + q, k = q[:, :, :1], k[:, :, :1] + initial = state.unsqueeze(0).repeat(2, 1, 1, 1).transpose(-1, -2).detach().requires_grad_() + function = matmul_kda if kda else matmul_gdn + policy = LinearAttentionConfig( + backend="matmul", decode={"mode": "replay", "replay": {"window": 3}} + ) + shared = { + "sites": LinearAttentionMatmulSites(), + "policy": policy, + "w_quantizer": TensorQuantizer(QuantizerAttributeConfig(enable=False)), + "state_qdq": True, + "output_final_state": True, + } + actual = function( + q, + k, + v, + g, + beta, + initial_state=initial, + state_v_first=True, + prefill_lengths=[0, 0] if all_empty else [2, 5], + **shared, + ) + expected = function( + q.repeat_interleave(2, dim=2), + k.repeat_interleave(2, dim=2), + v, + g, + beta, + initial_state=initial.transpose(-1, -2), + prefill_lengths=[0, 0] if all_empty else [2, 5], + **shared, + ) + torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0) + torch.testing.assert_close(actual[1].transpose(-1, -2), expected[1], rtol=0, atol=0) + gradients = [ + torch.autograd.grad(result[0].sum() + result[1].sum(), initial, retain_graph=True)[0] + for result in (actual, expected) + ] + torch.testing.assert_close(*gradients, rtol=0, atol=0) + if all_empty: + torch.testing.assert_close(actual[1], initial, rtol=0, atol=0) diff --git a/tests/unit/torch/quantization/test_linear_attention_study.py b/tests/unit/torch/quantization/test_linear_attention_study.py index fabdafa8755..49028e2b475 100644 --- a/tests/unit/torch/quantization/test_linear_attention_study.py +++ b/tests/unit/torch/quantization/test_linear_attention_study.py @@ -39,6 +39,8 @@ def _receipt(): "train_tokens_sha256": "train_tokens", "eval_tokens_sha256": "test_tokens", "sequence_length": 128, + "prefill_tokens": None, + "loss_scope": "full", "seed": 2026, "train_order": [1, 0], "training": "attention_only", @@ -69,7 +71,15 @@ def test_paired_quality_bound_and_failure(): @pytest.mark.parametrize( - "field", ["eval_tokens_sha256", "source_sha256", "train_order", "learning_rate"] + "field", + [ + "eval_tokens_sha256", + "source_sha256", + "train_order", + "learning_rate", + "prefill_tokens", + "loss_scope", + ], ) def test_quality_refuses_unmatched_studies(field): control, candidate = _receipt(), _receipt()