diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index c38e462912a..f0c5a38bbab 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.github/workflows/gpu_tests.yml @@ -44,7 +44,8 @@ jobs: matrix: include: - example: gpu - timeout: 60 + # Includes dependency builds and cold compilation of the FLA forward/backward tests. + timeout: 75 # Pinned to 26.05: benchmark.py uses trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH, # which newer TensorRT (26.06) removed. Bump once the source is updated for TensorRT 10. container_image: nvcr.io/nvidia/pytorch:26.05-py3 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7654acab550..e677fb2dff2 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -108,6 +108,8 @@ repos: modelopt/onnx/quantization/ort_patching.py| modelopt/torch/_deploy/utils/onnx_utils.py| modelopt/torch/export/transformer_engine.py| + modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py| + modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py| modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_pruned_to_mxfp4.py| modelopt/torch/quantization/export_onnx.py| modelopt/torch/quantization/plugins/attention.py| diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 22169f974c0..60d905778b0 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,10 +13,13 @@ Changelog *Quantization* +- Add experimental GDN/KDA decode-aware QAT with FP8 or INT8 recurrent states, optional INT8 value-axis Hadamard encoding, KDA decay rounding, and encoded-update replay. Supply explicit prefix lengths through the training phase context; prefix solves remain exact. + - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration. - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md `_ for details. - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``. +- Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training through the standard ``quant_cfg`` interface. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer. *Megatron Framework (M-LM / M-Bridge)* diff --git a/LICENSE b/LICENSE index 57c1104ae67..e08eaefc016 100644 --- a/LICENSE +++ b/LICENSE @@ -251,6 +251,7 @@ the following copyright holders, licensed under the MIT License: Copyright (c) 2025 sgl-project Copyright (c) 2026 The DeepSpec Authors Copyright (c) 2023-2026 The ggml authors + Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/docs/source/_templates/autosummary/module.rst b/docs/source/_templates/autosummary/module.rst index 36eb087fbec..c693958c63d 100644 --- a/docs/source/_templates/autosummary/module.rst +++ b/docs/source/_templates/autosummary/module.rst @@ -11,7 +11,7 @@ :recursive: {% for item in modules %} {% set full_item = fullname + '.' + item.split('.')[-1] %} -{% if ('.plugins.' not in full_item or full_item == 'modelopt.torch.opt.plugins.huggingface') and full_item != 'modelopt.torch.quantization.backends.fp8_per_tensor_gemm' %} +{% if ('.plugins.' not in full_item or full_item == 'modelopt.torch.opt.plugins.huggingface') and full_item != 'modelopt.torch.quantization.backends.fp8_per_tensor_gemm' and not full_item.startswith('modelopt.torch.kernels.quantization.linear_attention.fla_') %} {{ full_item }} {% endif %} {%- endfor %} diff --git a/examples/llm_qat/README.md b/examples/llm_qat/README.md index b4900ebea1f..aa099d88ebd 100644 --- a/examples/llm_qat/README.md +++ b/examples/llm_qat/README.md @@ -14,6 +14,7 @@ For background on how QAT enables low-precision accuracy recovery, see the [QAT/ | Background | How QAT/QAD work and when to use each | \[[Link](#background)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | | Support Matrix | Supported models, quantization formats, and backends | \[[Link](#support-matrix)\] | | | QLoRA | Model training with reduced GPU memory | \[[Link](#qlora-real-quantization)\] | | +| Linear Attention | GDN/KDA recurrent-state QAT and ReplaySSM example | \[[Link](linear_attention/README.md)\] | | | Advanced Topics | FSDP2 config, YAML options | \[[Link](#advanced-topics)\] | | | Results | Accuracy benchmarks | \[[Link](#results)\] | | | Resources | Extra links and references | \[[Link](#resources)\] | | diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md new file mode 100644 index 00000000000..ad0e116c1c3 --- /dev/null +++ b/examples/llm_qat/linear_attention/README.md @@ -0,0 +1,307 @@ +# Quantization-Aware Training for Linear Attention + +This example fine-tunes KDA attention parameters with recurrent-state fake +quantization, then saves a Hugging Face checkpoint with ModelOpt state. GDN/KDA +runtime support includes token writes, ReplaySSM, KDA decay approximation, and +FP8 or INT8 state QDQ. The INT8 recipes enable Hadamard rotation by default. + +## Run the example + +Install ModelOpt using the [QAT setup instructions](../README.md#quick-start), +then install the example dependencies. Run commands from the repository root. + +```bash +pip install -r examples/llm_qat/linear_attention/requirements.txt +python examples/llm_qat/linear_attention/train.py \ + --model /path/to/local-kda-model \ + --train-data /path/to/train.parquet \ + --output /path/to/qat-checkpoint \ + --train-steps 1 --length 128 --prefill-tokens 64 +``` + +Use a local model/tokenizer snapshot containing FLA `KimiDeltaAttention` layers, +a Parquet file with a `text` column, and a CUDA GPU. The example requires +`fla-core==0.5.1` and `flash-linear-attention==0.5.1`. If the local model requires +custom Python code, review it and explicitly pass `--trust-remote-code`. + +The default [INT8 configuration](configs/kda_decode_state_int8.json) leaves the +64-token prefix state unquantized and applies INT8 QDQ in a 32-value Hadamard +basis during the decode suffix. Value dimensions must be divisible by 32. +Loss uses suffix labels. Only KDA attention parameters are trained, in FP32 under +BF16 autocast; other parameters are frozen in BF16. The output contains model +weights, tokenizer files, and ModelOpt quantizer/policy state. Call +`mto.enable_huggingface_checkpointing()` before reloading it with +`AutoModelForCausalLM.from_pretrained` to restore those quantizers and policies. +It is a floating fake-quantized training checkpoint, not a compressed serving model. +This short example does not measure model-quality recovery or performance. + +## Enable state quantization + +### What the execution policy controls + +An **execution policy** is the layer's `LinearAttentionConfig`: the settings that +determine how the recurrence runs and when it applies quantization. Training +uses three separate inputs: + +| Input | Purpose | Supplied through | +| --- | --- | --- | +| `TensorQuantizer` settings | Enable a quantization site and choose its numerical format, scales, and gradient behavior | Recipe `quant_cfg`, such as dynamic INT8 with identity STE | +| Execution policy | Choose token or replay mode, Hadamard rotation, handoff quantization, and the execution backend | Recipe `linear_attention` entries | +| Batch metadata | Specify where each sequence switches from prefill to decode | `linear_attention_training_phase(model, prefill_lengths)` | + +The same INT8 quantizer settings can round the state after every token or round +a replay anchor every eight tokens. Those schedules produce different recurrent +states, so the policy must specify which computation training should emulate. + +`LinearAttentionConfig` holds the overall backend, chunk size, and state settings. +Its optional `decode` field contains a `LinearAttentionDecodeConfig` for token or +replay mode, Torch or Triton implementation, state codec, decay approximation, +and replay settings. This is one nested configuration: supplying a `decode` +dictionary in the recipe constructs the nested config automatically. + +Each `linear_attention` entry uses `module_name` to select attention layers and +`cfg` to specify their policy. During `mtq.quantize`, ModelOpt stores that policy +as each matched layer's `linear_attention_config`. Quantizer settings and the +policy persist through ModelOpt save/restore; batch-specific prefill lengths +must be supplied again for each workload. + +### Load a state recipe + +State quantizers start disabled. The state recipe enables `*kda_state_quantizer` +with signed narrow-range INT8, dynamic scales, and `axis=(0, 1)`. Use +`*gdn_state_quantizer` for GDN. For E4M3, set the quantizer config to +`{"num_bits": [4, 3], "type": "dynamic", "axis": [0, 1]}` and set +`decode.state_codec="tile"`. +Setting `prefill_state_qdq=True` alone does not enable a quantizer. + +Start with the token-state recipe and select one of the schedules below **before** +calling `mtq.quantize`: + +```python +import json +from pathlib import Path + +recipe = json.loads( + Path("examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json").read_text() +) +policy = recipe["linear_attention"][0]["cfg"] +decode = policy["decode"] +``` + +For GDN, the complete recipe composes the INT8 quantizer unit with the Hadamard +execution policy: + +```python +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe + +gdn_recipe = load_recipe("general/ptq/gdn_state_int8_dynamic").quantize +mtq.quantize(gdn_model, gdn_recipe) +``` + +Use the phase context shown below for each forward/backward. Importing only +`configs/ptq/units/gdn_state_int8_dynamic` configures the quantizer without +selecting Hadamard. The complete GDN recipe uses the Torch implementation; +set `gdn_recipe.linear_attention[0].cfg.decode.implementation="triton"` before +conversion to select the fused decode implementation. + +### Why training needs prefill/decode boundaries + +One training forward can simulate prompt prefill followed by recurrent decode. +The INT8 + Hadamard recipes apply different state quantization schedules to those +phases, so the workload must specify where each sequence switches to decode. +The training batch supplies all tokens; this simulates decode arithmetic without +running a text-generation loop. Total sequence length alone does not identify +the prompt portion. + +For a 128-token sequence with `prefill_lengths=[96]`, the default token-state +recipe runs these steps: + +1. **Prefill:** process tokens 0–95 with `prefill_state_qdq=False`, leaving state + unquantized during the prefix. +2. **Handoff:** carry the resulting state into decode, rotate its value dimension + into the Hadamard basis, and apply INT8 QDQ because `quantize_initial=True`. +3. **Decode:** process tokens 96–127 recurrently, applying INT8 QDQ after every + state update. Later tokens consume the rounded state; outputs are transformed + back to the original value basis. + +The boundary is independent of `chunk_size=64`: this 96-token prefix contains a +full 64-token chunk and a partial 32-token chunk. The phase switches after token +95, not after each chunk. + +Two 128-token sequences can use `prefill_lengths=[64, 96]` in the same batch: +their decode suffixes then contain 64 and 32 tokens, respectively. The next batch +can have different lengths while reusing the same quantization recipe. This is +why the execution policy is saved in `linear_attention_config`, while the lengths +are supplied per batch through `linear_attention_training_phase`. They are not +saved as part of the model's quantization policy. The context selects numerical +phases; the caller still supplies training labels and any loss masking. + +The combined prefill/decode path requires explicit lengths, including `[0]` for +decode only or `[T]` for an all-prefix sequence of length `T`. The GDN chunk-only +FLA path described below applies its chunk schedule throughout and needs no +phase context. See the tables below for the corresponding quantization settings. + +### Choose the prefill and decode boundaries + +The table assumes the state quantizer is enabled. The INT8 recipes default to +`"int8_hadamard32"`; prefix state QDQ requires explicitly selecting `"tile"`. +Prefix lengths are supplied separately through `linear_attention_training_phase`. + +| Desired state QDQ | `decode.state_codec` | `decode.mode` | `decode.prefill_state_qdq` | Where rounding occurs | +| --- | --- | --- | --- | --- | +| Token decode only (default) | `"int8_hadamard32"` | `"token"` | `False` | At the first nonempty decode handoff and after every suffix token. | +| Prefill and token decode | `"tile"` | `"token"` | `True` | At prefix initialization, each prefix chunk write, decode handoff, and every suffix token. | +| Replay anchors only | `"int8_hadamard32"` | `"replay"` | `False` | At decode handoff and each replay-window refresh. | +| Prefill and replay anchors | `"tile"` | `"replay"` | `True` | At prefix initialization and chunk writes, then decode handoff and replay-window refreshes. | + +To enable **prefill and token decode**: + +```python +decode.update(mode="token", replay=None, state_codec="tile", prefill_state_qdq=True) +``` + +For **token decode only**, keep the supplied INT8 recipe's defaults: +`state_codec="int8_hadamard32"` and `prefill_state_qdq=False`. +Prefix state remains unquantized until it enters the decode path. + +To enable **ReplaySSM anchor quantization** with an eight-token window: + +```python +decode.update( + mode="replay", + state_codec="int8_hadamard32", + prefill_state_qdq=False, + replay={"window": 8, "factor_qdq": False, "encoding": "once"}, +) +``` + +Set `state_codec="tile"` and `prefill_state_qdq=True` to add prefix state QDQ +to this replay configuration, using unrotated INT8 for both phases. +`factor_qdq=False` above isolates state/anchor quantization. Set it to `True` to +also quantize buffered keys and updates to FP8. + +`decode.quantize_initial=True` is the default: it quantizes the incoming state +once at the first nonempty decode handoff. Set it to `False` to skip that initial +rounding while keeping later token writes or anchor refreshes quantized. This +setting does not disable prefix state QDQ. With both phases enabled, prefix-final +rounding and decode-handoff rounding are separate configured events. + +`chunk_size=64` counts **tokens per prefill chunk**. `state.block_v=64` counts +**value channels per execution tile**. The tile codec shares a scale across all +key channels and this value tile; Hadamard uses one scale per key channel and +32 values. A replay `window=8` counts **suffix tokens between anchor refreshes**. + +### Run only the desired phase + +For a batch containing one sequence of `T` tokens, choose the context lengths as +follows; for larger batches, provide one length per sequence: + +| Workload | Context argument | Required setting | +| --- | --- | --- | +| Prefill followed by decode | `[64]`, with `T > 64` | Select either prefix setting above. | +| Decode only | `[0]` | State quantizer enabled; token or replay policy. | +| Prefill only | `[T]` | `state_codec="tile"`, `prefill_state_qdq=True`; the decode suffix is empty. | + +An empty suffix creates no decode quantization event. The combined interface has +one state quantizer per layer: it does not offer a switch to quantize the prefix +while leaving a **nonempty** decode suffix unquantized. Disable the state quantizer +to disable both state and anchor QDQ; replay factor QDQ has its own toggle. + +The `int8_hadamard32` codec applies to decode token writes or replay anchors only. +It requires `prefill_state_qdq=False`; use the tile codec to quantize prefix state. + +Save a modified recipe as JSON and pass it to `train.py --quant-config`. `--prefill-tokens` sets the phase split; it does not enable state +quantization. The integration loop below applies the configured `recipe` directly. + +### GDN chunk-only FLA training + +For GDN's existing chunked FLA path, enable the GDN state quantizer and omit the +`linear_attention` execution-policy entry: + +```python +gdn_recipe = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*gdn_state_quantizer", + "cfg": recipe["quant_cfg"][1]["cfg"], + }, + ], + "algorithm": None, +} +# Apply mtq.quantize(gdn_model, gdn_recipe), then use normal forward/backward. +``` + +This explicitly selects unrotated tile QDQ. It rounds the initial state and each +64-token chunk's final state, including a partial final chunk. It needs no phase +context and leaves W/projection quantizers +disabled. KDA uses the materialized backend and can use the all-prefix context +shown above. + +## Integrate with a training loop + +Apply the configured `recipe` above with `mtq.quantize`, then supply one prefix length per sequence. +Keep the phase context active through backward so activation-checkpoint +recomputation uses the same prefix/decode split. + +```python +import torch + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase + +mtq.quantize(model, recipe) +optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5) +model.train() +optimizer.zero_grad(set_to_none=True) + +# Two 128-token sequences: 64 and 96 prefill tokens, respectively. +# ids and labels have shape [2, 128] and are on the model's device. +with linear_attention_training_phase(model, [64, 96]): + with torch.autocast("cuda", dtype=torch.bfloat16): + loss = model(input_ids=ids, labels=labels, use_cache=False).loss + loss.backward() +optimizer.step() +``` + +For GDN, select `*gdn_state_quantizer` in the recipe instead of +`*kda_state_quantizer`. State quantization is dynamic, so these recipes use +`algorithm=None` without a calibration pass. Replay factors have their own FP8 +toggle. + +Policies persist through ModelOpt save/restore. Per-batch prefix lengths are +runtime metadata and must be supplied for each workload. The context restores +previous lengths on exit and supports nesting. Concurrent forwards on the same +model instance with different phase contexts are unsupported. + +## Replay, decay, and Hadamard options + +Token mode rounds each recurrent-state write. Replay mode retains an anchor and +ordered key/update factors, rounding the anchor at each `replay.window` refresh. +`factor_qdq` independently enables FP8 QDQ for those factors. `readout="working"` +reads the state before its write quantization; `"stored"` reads it afterward. + +For KDA decay approximation, set `decode["decay_log_step"] = 1 / 256` before +conversion. It rounds suffix log retention before exponentiation with identity +STE gradients. Prefix decay remains exact. + +The default INT8 codec applies a 32-point orthonormal Hadamard transform along +the value axis, quantizes token states or replay anchors in that basis, and +transforms outputs back. Value dimensions must be divisible by 32; `state.block_v` must be 32, 64, +or 128. Scales group one key channel and 32 values, independent of execution tile +width. INT8 codes use half-away-from-zero rounding with FP16 stored scales; +the optional tile codec instead uses nearest-even rounding and FP32 scales. + +## Training boundaries + +The Torch and Triton implementations support first-order QAT gradients through +initial states, chunk handoff, token writes, and replay refreshes. Triton supports +key dimensions up to 128 and value blocks 16/32/64/128, with checkpointed backward. +Use the Torch implementation for higher-order differentiation. + +Prefill prefixes use exact chunk algebra with optional state QDQ. This example +has no prefill GEMM QDQ or approximate inverse. FLA KDA requires `use_cache=False`; +serving cache objects are rejected. ModelOpt saves execution policies, while +per-batch prefix lengths must be supplied again during training. Distributed +decode training and model-quality recovery require separate qualification. diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json b/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json new file mode 100644 index 00000000000..f6cc5fef650 --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json @@ -0,0 +1,35 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": 8, + "type": "dynamic", + "axis": [ + 0, + 1 + ], + "unsigned": false, + "narrow_range": true + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "state_codec": "int8_hadamard32", + "prefill_state_qdq": false + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/requirements.txt b/examples/llm_qat/linear_attention/requirements.txt new file mode 100644 index 00000000000..29636d78e52 --- /dev/null +++ b/examples/llm_qat/linear_attention/requirements.txt @@ -0,0 +1,4 @@ +fla-core==0.5.1 +flash-linear-attention==0.5.1 +pyarrow>=15 +transformers==4.57.1 diff --git a/examples/llm_qat/linear_attention/train.py b/examples/llm_qat/linear_attention/train.py new file mode 100644 index 00000000000..7662c1ca593 --- /dev/null +++ b/examples/llm_qat/linear_attention/train.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fine-tune KDA attention parameters with recurrent-state fake quantization.""" + +import argparse +import json +from pathlib import Path + +import pyarrow.parquet as pq +import torch +from fla.layers.kda import KimiDeltaAttention +from transformers import AutoModelForCausalLM, AutoTokenizer + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase + + +def main(): + """Quantize, train on local text, and save a ModelOpt-aware Hugging Face checkpoint.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument( + "--train-data", type=Path, required=True, help="Parquet file with a text column" + ) + parser.add_argument( + "--quant-config", + type=Path, + default=Path(__file__).with_name("configs") / "kda_decode_state_int8.json", + ) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--train-steps", type=int, default=1) + parser.add_argument("--length", type=int, default=128) + parser.add_argument("--prefill-tokens", type=int, default=64) + parser.add_argument("--seed", type=int, default=2026) + parser.add_argument("--learning-rate", type=float, default=1e-5) + options = parser.parse_args() + if options.train_steps < 1 or options.length < 2: + parser.error("train-steps must be positive and length must be at least two") + if not 0 <= options.prefill_tokens < options.length: + parser.error("prefill-tokens must leave at least one suffix label") + torch.manual_seed(options.seed) + mto.enable_huggingface_checkpointing() + tokenizer = AutoTokenizer.from_pretrained( + options.model, local_files_only=True, trust_remote_code=options.trust_remote_code + ) + text = "\n\n".join(pq.read_table(options.train_data, columns=["text"])["text"].to_pylist()) + tokens = tokenizer.encode(text, add_special_tokens=False) + needed = options.train_steps * options.length + if len(tokens) < needed: + parser.error(f"Training data has {len(tokens)} tokens; {needed} are required") + blocks = torch.tensor(tokens[:needed]).reshape(options.train_steps, options.length) + model = AutoModelForCausalLM.from_pretrained( + options.model, + local_files_only=True, + trust_remote_code=options.trust_remote_code, + torch_dtype=torch.float32, + ) + layers = [module for module in model.modules() if isinstance(module, KimiDeltaAttention)] + if not layers: + raise ValueError("The model must contain FLA KimiDeltaAttention layers") + # Keep attention parameters and optimizer moments in FP32; freeze the rest in BF16. + model.requires_grad_(False) + for layer in layers: + layer.requires_grad_(True) + for parameter in model.parameters(): + if not parameter.requires_grad: + parameter.data = parameter.data.to(torch.bfloat16) + model.cuda() + mtq.quantize(model, json.loads(options.quant_config.read_text())) + model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + 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() + for step, block in enumerate(blocks, 1): + ids = block.unsqueeze(0).cuda() + labels = ids.clone() + labels[:, : options.prefill_tokens] = -100 + optimizer.zero_grad(set_to_none=True) + # Keep phase metadata active during activation-checkpoint recomputation. + with linear_attention_training_phase(model, [options.prefill_tokens]): + with torch.autocast("cuda", dtype=torch.bfloat16): + loss = model(input_ids=ids, labels=labels, use_cache=False).loss + if not torch.isfinite(loss): + raise RuntimeError("Non-finite training loss") + loss.backward() + torch.nn.utils.clip_grad_norm_(trainable, 1.0, error_if_nonfinite=True) + optimizer.step() + print(f"step={step} loss={loss.item():.4f}", flush=True) + model.save_pretrained(options.output) + tokenizer.save_pretrained(options.output) + + +if __name__ == "__main__": + main() diff --git a/modelopt/torch/kernels/quantization/common/fp8_quant.py b/modelopt/torch/kernels/quantization/common/fp8_quant.py index 2b76499b968..30c2d148629 100644 --- a/modelopt/torch/kernels/quantization/common/fp8_quant.py +++ b/modelopt/torch/kernels/quantization/common/fp8_quant.py @@ -16,11 +16,36 @@ """Composable Triton JIT functions for FP8 (E4M3) fake quantization. Counterpart of ``nvfp4_quant.py`` for per-tensor FP8. Used by the unified -flash-attention kernel's softmax-P qdq (``common/attention/triton_fa.py``). +flash-attention kernel's softmax-P qdq and linear-attention state/factor qdq. """ import triton import triton.language as tl +from triton.language.extra.cuda import libdevice + +__all__ = ["fp8_dynamic_qdq", "fp8_scalar_qdq"] + + +@triton.jit +def _fp8_cast(value, NATIVE_FP8: tl.constexpr): + value = tl.clamp(value, -448.0, 448.0) + if NATIVE_FP8: + return value.to(tl.float8e4nv).to(tl.float32) + # GPUs before SM89 need a software E4M3 conversion, including subnormals and ties. + bits = tl.abs(value).to(tl.int32, bitcast=True) + exponent = ((bits >> 23) & 255) - 127 + step = ((tl.maximum(exponent, -6) - 3 + 127) << 23).to(tl.float32, bitcast=True) + return libdevice.nearbyint(value / step) * step + + +@triton.jit +def fp8_dynamic_qdq(value, NATIVE_FP8: tl.constexpr = True): + """Dynamic per-tile FP8 QDQ with TensorQuantizer scaling and returned decode scale.""" + amax = tl.max(tl.abs(value)) + safe_amax = tl.where(amax <= 2.0**-24, 1.0, amax) + quant_scale = tl.div_rn(448.0, safe_amax) + decode_scale = tl.div_rn(1.0, quant_scale) + return _fp8_cast(value * quant_scale, NATIVE_FP8) * decode_scale, decode_scale @triton.jit @@ -38,6 +63,4 @@ def fp8_scalar_qdq(x, scale): Returns: Fake-quantized tensor of the same shape as ``x``, in float32. """ - FP8_E4M3_MAX: tl.constexpr = 448.0 - x_scaled = tl.clamp(x / scale, -FP8_E4M3_MAX, FP8_E4M3_MAX) - return x_scaled.to(tl.float8e4nv).to(tl.float32) * scale + return _fp8_cast(x / scale, True) * scale diff --git a/modelopt/torch/kernels/quantization/linear_attention/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/__init__.py new file mode 100644 index 00000000000..3f003f403f8 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/__init__.py @@ -0,0 +1,24 @@ +# 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. + +"""Linear-attention kernels for quantization. + +``fla_chunk_delta_h.py`` and ``fla_chunk_gated_delta_rule.py`` are adapted copies of the chunked +GatedDeltaNet kernels of `flash-linear-attention `_ +(``fla.ops.common.chunk_delta_h`` and ``fla.ops.gated_delta_rule.chunk``) that can fake-quantize the +recurrent state carried between chunks to FP8 (``state_qdq``). They still import the surrounding +fla operators, so ``fla-core==0.5.1`` and Triton must be installed to use +them. This package initializer does not import the kernels, so importing it needs neither. +""" 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..a156f824b12 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/decode.py @@ -0,0 +1,468 @@ +# 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 + +from ..common.fp8_quant import fp8_dynamic_qdq +from .int8 import int8_scalar_qdq + +__all__ = ["fused_recurrence"] + + +@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, + NATIVE_FP8: 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_dynamic_qdq(update, NATIVE_FP8) + working = decayed + key[:, None] * update[None, :] + stored = working + if STATE_QDQ == 3: + state_scale = tl.full((working.shape[0], working.shape[1] // 32), 1.0, tl.float16) + else: + state_scale = 1.0 + if STATE_QDQ: + if write_state: + if STATE_QDQ == 3: + groups = tl.reshape(working, (working.shape[0], working.shape[1] // 32, 32)) + scale = tl.maximum(tl.max(tl.abs(groups), 2) * (1.0 / 127.0), 6e-8) + normalized = tl.div_rn(groups, scale[:, :, None]) + codes = tl.floor(tl.abs(normalized) + 0.5) + codes = tl.where(normalized < 0, -codes, codes) + state_scale = scale.to(tl.float16) + stored = tl.reshape( + tl.minimum(tl.maximum(codes, -127.0), 127.0) + * state_scale.to(tl.float32)[:, :, None], + working.shape, + ) + elif STATE_QDQ == 2: + amax = tl.max(tl.abs(working)) + # Match Torch scalar division: FP32 multiply by the rounded reciprocal. + state_scale = tl.where(amax > 0, amax * (1.0 / 127.0), 1.0) + stored = int8_scalar_qdq(working, state_scale) + else: + stored, state_scale = fp8_dynamic_qdq(working, NATIVE_FP8) + 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, + NATIVE_FP8: 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) + if STATE_QDQ == 3: + groups = tile * (BV // 32) + tl.arange(0, BV // 32) + scale_offsets = (h * DK + keys[:, None]) * (DV // 32) + groups[None, :] + scale_mask = (keys[:, None] < DK) & (groups[None, :] < DV // 32) + tl.store(AnchorScale + scale_offsets, 1.0, scale_mask) + else: + 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, + NATIVE_FP8, + ) + 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) + if STATE_QDQ == 3: + tl.store(AnchorScale + scale_offsets, state_scale, scale_mask) + else: + 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, + NATIVE_FP8: 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, + NATIVE_FP8, + ) + 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, + NATIVE_FP8, + ) + 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) + scale_shape = (heads, keys, values // 32) if state_qdq == 3 else (heads, tiles) + anchor_scales = torch.empty( + scale_shape, device=q.device, dtype=torch.float16 if state_qdq == 3 else 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, + "NATIVE_FP8": torch.cuda.get_device_capability(q.device) >= (8, 9), + "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, + state_format="fp8_e4m3", + state_codec="tile", + 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") + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") + if state_codec not in ("tile", "int8_hadamard32"): + raise ValueError("State codec must be tile or int8_hadamard32") + if state_codec == "int8_hadamard32": + if (state_qdq and state_format != "int8") or v.shape[-1] % 32 or block_v < 32: + raise ValueError("Hadamard codec requires INT8, Dv divisible by 32 and block_v >= 32") + state_mode = 3 if state_codec == "int8_hadamard32" else (2 if state_format == "int8" else 1) + state_qdq = state_mode if state_qdq else 0 + 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/kernels/quantization/linear_attention/fla_chunk_delta_h.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py new file mode 100644 index 00000000000..84da262fd9f --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -0,0 +1,992 @@ +# Adapted from: https://github.com/fla-org/flash-linear-attention/blob/516143e31fce/fla/ops/common/chunk_delta_h.py +# Adapted with modifications (marked [ModelOpt]): optional in-kernel FP8 E4M3 / INT8 fake quantization +# of the carried chunk state (STATE_QDQ), BV as an explicit launch argument, no fla backend +# dispatch or config-cache autotune. +# +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# 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 torch +import triton +import triton.language as tl +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.cache import fla_cache_autotune +from fla.ops.utils.op import exp2 +from fla.utils import ( + IS_INTEL, + IS_NVIDIA_BLACKWELL, + IS_NVIDIA_HOPPER, + autotune_cache_kwargs, + check_shared_mem, +) + +from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq + +from .int8 import int8_scalar_qdq + +# ``STATE_QDQ`` modes of the forward state kernel. +STATE_QDQ_OFF = 0 +STATE_QDQ_FP8_DYNAMIC = 1 # FP8 E4M3, one dynamic scale per program tile ([K, BV] of one head) +STATE_QDQ_INT8_DYNAMIC = 2 +STATE_QDQ_MAX_BLOCK_V = 128 + + +@triton.jit +def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr, STATE_QDQ: tl.constexpr): + """[ModelOpt] Dynamic scale per full [K, BV] tile of one sequence and head.""" + b_amax = tl.max(tl.abs(b_h1)) + if K > 64: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h2))) + if K > 128: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h3))) + if K > 192: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h4))) + if STATE_QDQ == 2: + return tl.where(b_amax > 0, b_amax * (1.0 / 127.0), 1.0) + else: + return tl.where(b_amax > 0, b_amax / 448.0, 1.0) + + +@triton.jit +def _state_scalar_qdq(value, scale, STATE_QDQ: tl.constexpr): + if STATE_QDQ == 2: + return int8_scalar_qdq(value, scale) + else: + return fp8_scalar_qdq(value, scale) + + +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] + +# TODO: Triton mainline fixes a Blackwell tl.dot recurrence race. +# Keep this kernel on num_warps=2 for Blackwell until Triton 3.8 is released +# and we re-validate the wider config space. +# Intel needs more warps than NVIDIA here: 8 warps is ~1.5x faster than the best +# config reachable under the [2, 4] cap. +if IS_NVIDIA_BLACKWELL: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2] +elif IS_INTEL: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2, 4, 8, 16] +else: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2, 4] + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +# ``BV`` is an explicit argument rather than an autotuned config: with ``STATE_QDQ`` it sets the +# quantization granularity, so it must not vary with the autotuner's choice. +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in GATED_DELTA_RULE_FWD_H_NUM_WARPS + for num_stages in ([2, 3, 4] if check_shared_mem("ampere") else [2, 1]) + ], + key=["H", "HV", "K", "V", "BT", "BV", "STATE_V_FIRST", "STATE_QDQ"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( + k, + v, + w, + v_new, + g, + gk, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + SAVE_NEW_VALUE: tl.constexpr, + STATE_V_FIRST: tl.constexpr, + IS_VARLEN: tl.constexpr, + STATE_QDQ: tl.constexpr, +): + pid = tl.program_id(0) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) + i_n, i_h = i_nh // HV, i_nh % HV + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + if STATE_V_FIRST: + b_h1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([BV, 64], dtype=tl.float32) + else: + b_h1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + h += (boh * HV + i_h).to(tl.int64) * K * V + v += (bos * HV + i_h).to(tl.int64) * V + k += (bos * H + i_h // (HV // H)).to(tl.int64) * K + w += (bos * HV + i_h).to(tl.int64) * K + if SAVE_NEW_VALUE: + v_new += (bos * HV + i_h).to(tl.int64) * V + + if USE_INITIAL_STATE: + h0 = h0 + i_nh * K * V + if STORE_FINAL_STATE: + ht = ht + i_nh * K * V + + # load initial state + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K + if USE_INITIAL_STATE: + if STATE_V_FIRST: + p_h0_1 = h0 + o_v[:, None] * K + o_k1[None, :] + m_h0_1 = m_v[:, None] & m_k1[None, :] + else: + p_h0_1 = h0 + o_k1[:, None] * V + o_v[None, :] + m_h0_1 = m_k1[:, None] & m_v[None, :] + b_h1 += tl.load(p_h0_1, mask=m_h0_1, other=0.0).to(tl.float32) + if K > 64: + if STATE_V_FIRST: + p_h0_2 = h0 + o_v[:, None] * K + o_k2[None, :] + m_h0_2 = m_v[:, None] & m_k2[None, :] + else: + p_h0_2 = h0 + o_k2[:, None] * V + o_v[None, :] + m_h0_2 = m_k2[:, None] & m_v[None, :] + b_h2 += tl.load(p_h0_2, mask=m_h0_2, other=0.0).to(tl.float32) + if K > 128: + if STATE_V_FIRST: + p_h0_3 = h0 + o_v[:, None] * K + o_k3[None, :] + m_h0_3 = m_v[:, None] & m_k3[None, :] + else: + p_h0_3 = h0 + o_k3[:, None] * V + o_v[None, :] + m_h0_3 = m_k3[:, None] & m_v[None, :] + b_h3 += tl.load(p_h0_3, mask=m_h0_3, other=0.0).to(tl.float32) + if K > 192: + if STATE_V_FIRST: + p_h0_4 = h0 + o_v[:, None] * K + o_k4[None, :] + m_h0_4 = m_v[:, None] & m_k4[None, :] + else: + p_h0_4 = h0 + o_k4[:, None] * V + o_v[None, :] + m_h0_4 = m_k4[:, None] & m_v[None, :] + b_h4 += tl.load(p_h0_4, mask=m_h0_4, other=0.0).to(tl.float32) + # [ModelOpt] A state read from a quantized cache is quantized before the first chunk uses it. + if STATE_QDQ != 0: + if K > 192: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K, STATE_QDQ=STATE_QDQ) + elif K > 128: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K, STATE_QDQ=STATE_QDQ) + elif K > 64: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K, STATE_QDQ=STATE_QDQ) + else: + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K, STATE_QDQ=STATE_QDQ) + b_h1 = _state_scalar_qdq(b_h1, b_scale, STATE_QDQ) + if K > 64: + b_h2 = _state_scalar_qdq(b_h2, b_scale, STATE_QDQ) + if K > 128: + b_h3 = _state_scalar_qdq(b_h3, b_scale, STATE_QDQ) + if K > 192: + b_h4 = _state_scalar_qdq(b_h4, b_scale, STATE_QDQ) + + # main recurrence + for i_t in range(NT): + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + if STATE_V_FIRST: + p_h1 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k1[None, :] + m_h1 = m_v[:, None] & m_k1[None, :] + else: + p_h1 = h + i_t_int64 * HV * K * V + o_k1[:, None] * V + o_v[None, :] + m_h1 = m_k1[:, None] & m_v[None, :] + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), mask=m_h1) + if K > 64: + if STATE_V_FIRST: + p_h2 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k2[None, :] + m_h2 = m_v[:, None] & m_k2[None, :] + else: + p_h2 = h + i_t_int64 * HV * K * V + o_k2[:, None] * V + o_v[None, :] + m_h2 = m_k2[:, None] & m_v[None, :] + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), mask=m_h2) + if K > 128: + if STATE_V_FIRST: + p_h3 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k3[None, :] + m_h3 = m_v[:, None] & m_k3[None, :] + else: + p_h3 = h + i_t_int64 * HV * K * V + o_k3[:, None] * V + o_v[None, :] + m_h3 = m_k3[:, None] & m_v[None, :] + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), mask=m_h3) + if K > 192: + if STATE_V_FIRST: + p_h4 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k4[None, :] + m_h4 = m_v[:, None] & m_k4[None, :] + else: + p_h4 = h + i_t_int64 * HV * K * V + o_k4[:, None] * V + o_v[None, :] + m_h4 = m_k4[:, None] & m_v[None, :] + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), mask=m_h4) + + p_w = w + o_t[:, None] * (HV * K) + o_k1[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k1[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) + else: + b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) + if K > 64: + p_w = w + o_t[:, None] * (HV * K) + o_k2[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k2[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h2.to(b_w.dtype), b_v) + if K > 128: + p_w = w + o_t[:, None] * (HV * K) + o_k3[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k3[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h3.to(b_w.dtype), b_v) + if K > 192: + p_w = w + o_t[:, None] * (HV * K) + o_k4[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k4[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h4.to(b_w.dtype), b_v) + p_v = v + o_t[:, None] * (HV * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & m_v[None, :], other=0.0) - b_v + + if SAVE_NEW_VALUE: + p_v = v_new + o_t[:, None] * (HV * V) + o_v[None, :] + tl.store(p_v, b_v.to(p_v.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) + + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + b_g_last = tl.load(g + (bos * HV + last_idx * HV + i_h).to(tl.int64)).to(tl.float32) + p_g = g + (bos * HV + i_h).to(tl.int64) + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] + b_g_last = exp2(b_g_last) + b_h1 *= b_g_last + if K > 64: + b_h2 *= b_g_last + if K > 128: + b_h3 *= b_g_last + if K > 192: + b_h4 *= b_g_last + + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k1, mask=(o_k1 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h1 *= exp2(b_gk_last1)[None, :] + else: + b_h1 *= exp2(b_gk_last1)[:, None] + if K > 64: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k2, mask=(o_k2 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h2 *= exp2(b_gk_last2)[None, :] + else: + b_h2 *= exp2(b_gk_last2)[:, None] + if K > 128: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k3, mask=(o_k3 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h3 *= exp2(b_gk_last3)[None, :] + else: + b_h3 *= exp2(b_gk_last3)[:, None] + if K > 192: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k4, mask=(o_k4 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h4 *= exp2(b_gk_last4)[None, :] + else: + b_h4 *= exp2(b_gk_last4)[:, None] + b_v = b_v.to(k.dtype.element_ty) + + p_k = k + o_k1[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k1[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h1 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h1 = tl.dot(b_k, b_v, b_h1) + if K > 64: + p_k = k + o_k2[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k2[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h2 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h2 = tl.dot(b_k, b_v, b_h2) + if K > 128: + p_k = k + o_k3[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k3[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h3 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h3 = tl.dot(b_k, b_v, b_h3) + if K > 192: + p_k = k + o_k4[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k4[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h4 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h4 = tl.dot(b_k, b_v, b_h4) + + # [ModelOpt] Dynamic per-tile FP8/INT8 QDQ of the next-chunk or stored final state. + # Recompute amax over [K, BV]; _state_scalar_qdq applies that tile's scalar scale. + if STATE_QDQ != 0: + if K > 192: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K, STATE_QDQ=STATE_QDQ) + elif K > 128: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K, STATE_QDQ=STATE_QDQ) + elif K > 64: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K, STATE_QDQ=STATE_QDQ) + else: + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K, STATE_QDQ=STATE_QDQ) + b_h1 = _state_scalar_qdq(b_h1, b_scale, STATE_QDQ) + if K > 64: + b_h2 = _state_scalar_qdq(b_h2, b_scale, STATE_QDQ) + if K > 128: + b_h3 = _state_scalar_qdq(b_h3, b_scale, STATE_QDQ) + if K > 192: + b_h4 = _state_scalar_qdq(b_h4, b_scale, STATE_QDQ) + + if STORE_FINAL_STATE: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k1[None, :] + m_ht = m_v[:, None] & m_k1[None, :] + else: + p_ht = ht + o_k1[:, None] * V + o_v[None, :] + m_ht = m_k1[:, None] & m_v[None, :] + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 64: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k2[None, :] + m_ht = m_v[:, None] & m_k2[None, :] + else: + p_ht = ht + o_k2[:, None] * V + o_v[None, :] + m_ht = m_k2[:, None] & m_v[None, :] + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 128: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k3[None, :] + m_ht = m_v[:, None] & m_k3[None, :] + else: + p_ht = ht + o_k3[:, None] * V + o_v[None, :] + m_ht = m_k3[:, None] & m_v[None, :] + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 192: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k4[None, :] + m_ht = m_v[:, None] & m_k4[None, :] + else: + p_ht = ht + o_k4[:, None] * V + o_v[None, :] + m_ht = m_k4[:, None] & m_v[None, :] + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), mask=m_ht) + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "USE_INITIAL_STATE": lambda args: args["dh0"] is not None, + "USE_FINAL_STATE_GRADIENT": lambda args: args["dht"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@fla_cache_autotune( + configs=[ + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in ([2, 3, 4] if check_shared_mem("ampere") else [1]) + for BV in ([32, 64] if check_shared_mem("ada") else [32]) + ], + key=["H", "HV", "K", "V", "BT", "BV", "USE_G", "STATE_V_FIRST"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( + q, + k, + w, + g, + gk, + dht, + dh0, + do, + dh, + dv, + dv2, + cu_seqlens, + chunk_offsets, + scale, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + STATE_V_FIRST: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + pid = tl.program_id(0) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) + i_n, i_h = i_nh // HV, i_nh % HV + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + if STATE_V_FIRST: + b_dh1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([BV, 64], dtype=tl.float32) + else: + b_dh1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + q += (bos * H + i_h // (HV // H)).to(tl.int64) * K + k += (bos * H + i_h // (HV // H)).to(tl.int64) * K + w += (bos * HV + i_h).to(tl.int64) * K + do += (bos * HV + i_h).to(tl.int64) * V + dv += (bos * HV + i_h).to(tl.int64) * V + dv2 += (bos * HV + i_h).to(tl.int64) * V + dh += (boh * HV + i_h).to(tl.int64) * K * V + if USE_GK: + gk += (bos * HV + i_h).to(tl.int64) * K + + if USE_INITIAL_STATE: + dh0 += i_nh * K * V + if USE_FINAL_STATE_GRADIENT: + dht += i_nh * K * V + + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K + if USE_FINAL_STATE_GRADIENT: + if STATE_V_FIRST: + p_dht1 = dht + o_v[:, None] * K + o_k1[None, :] + m_dht1 = m_v[:, None] & m_k1[None, :] + else: + p_dht1 = dht + o_k1[:, None] * V + o_v[None, :] + m_dht1 = m_k1[:, None] & m_v[None, :] + b_dh1 += tl.load(p_dht1, mask=m_dht1, other=0.0) + if K > 64: + if STATE_V_FIRST: + p_dht2 = dht + o_v[:, None] * K + o_k2[None, :] + m_dht2 = m_v[:, None] & m_k2[None, :] + else: + p_dht2 = dht + o_k2[:, None] * V + o_v[None, :] + m_dht2 = m_k2[:, None] & m_v[None, :] + b_dh2 += tl.load(p_dht2, mask=m_dht2, other=0.0) + if K > 128: + if STATE_V_FIRST: + p_dht3 = dht + o_v[:, None] * K + o_k3[None, :] + m_dht3 = m_v[:, None] & m_k3[None, :] + else: + p_dht3 = dht + o_k3[:, None] * V + o_v[None, :] + m_dht3 = m_k3[:, None] & m_v[None, :] + b_dh3 += tl.load(p_dht3, mask=m_dht3, other=0.0) + if K > 192: + if STATE_V_FIRST: + p_dht4 = dht + o_v[:, None] * K + o_k4[None, :] + m_dht4 = m_v[:, None] & m_k4[None, :] + else: + p_dht4 = dht + o_k4[:, None] * V + o_v[None, :] + m_dht4 = m_k4[:, None] & m_v[None, :] + b_dh4 += tl.load(p_dht4, mask=m_dht4, other=0.0) + + for i_t in range(NT - 1, -1, -1): + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + if STATE_V_FIRST: + p_dh1 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k1[None, :] + m_dh1 = m_v[:, None] & m_k1[None, :] + else: + p_dh1 = dh + i_t_int64 * HV * K * V + o_k1[:, None] * V + o_v[None, :] + m_dh1 = m_k1[:, None] & m_v[None, :] + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), mask=m_dh1) + if K > 64: + if STATE_V_FIRST: + p_dh2 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k2[None, :] + m_dh2 = m_v[:, None] & m_k2[None, :] + else: + p_dh2 = dh + i_t_int64 * HV * K * V + o_k2[:, None] * V + o_v[None, :] + m_dh2 = m_k2[:, None] & m_v[None, :] + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), mask=m_dh2) + if K > 128: + if STATE_V_FIRST: + p_dh3 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k3[None, :] + m_dh3 = m_v[:, None] & m_k3[None, :] + else: + p_dh3 = dh + i_t_int64 * HV * K * V + o_k3[:, None] * V + o_v[None, :] + m_dh3 = m_k3[:, None] & m_v[None, :] + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), mask=m_dh3) + if K > 192: + if STATE_V_FIRST: + p_dh4 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k4[None, :] + m_dh4 = m_v[:, None] & m_k4[None, :] + else: + p_dh4 = dh + i_t_int64 * HV * K * V + o_k4[:, None] * V + o_v[None, :] + m_dh4 = m_k4[:, None] & m_v[None, :] + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), mask=m_dh4) + + last_idx = min((i_t_int64 + 1) * BT, T) - 1 + if USE_G: + bg_last = tl.load(g + (bos + last_idx) * HV + i_h).to(tl.float32) + p_g = g + bos * HV + i_h + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + bg_last_exp = exp2(bg_last) + b_g_exp = exp2(b_g) + p_dv = dv + o_t[:, None] * (HV * V) + o_v[None, :] + p_dv2 = dv2 + o_t[:, None] * (HV * V) + o_v[None, :] + p_do = do + o_t[:, None] * (HV * V) + o_v[None, :] + + b_do = tl.load(p_do, mask=m_t[:, None] & m_v[None, :], other=0.0) + + # Update dv + p_k = k + o_t[:, None] * (H * K) + o_k1[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k1[None, :], other=0.0) + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load(gk + last_idx * HV * K + o_k1, mask=(o_k1 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh1).to(b_k.dtype)) + else: + b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) + + if K > 64: + p_k = k + o_t[:, None] * (H * K) + o_k2[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k2[None, :], other=0.0) + if USE_GK: + b_gk_last2 = tl.load(gk + last_idx * HV * K + o_k2, mask=(o_k2 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh2).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh2.to(b_k.dtype), b_dv) + + if K > 128: + p_k = k + o_t[:, None] * (H * K) + o_k3[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k3[None, :], other=0.0) + if USE_GK: + b_gk_last3 = tl.load(gk + last_idx * HV * K + o_k3, mask=(o_k3 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh3).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh3.to(b_k.dtype), b_dv) + + if K > 192: + p_k = k + o_t[:, None] * (H * K) + o_k4[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k4[None, :], other=0.0) + if USE_GK: + b_gk_last4 = tl.load(gk + last_idx * HV * K + o_k4, mask=(o_k4 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh4).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh4.to(b_k.dtype), b_dv) + + if USE_G: + b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] + b_dv += tl.load(p_dv, mask=m_t[:, None] & m_v[None, :], other=0.0) + + tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) + # Update dh + p_w = w + o_k1[:, None] + o_t[None, :] * (HV * K) + p_q = q + o_k1[:, None] + o_t[None, :] * (H * K) + b_w = tl.load(p_w, mask=m_k1[:, None] & m_t[None, :], other=0.0) + b_q = tl.load(p_q, mask=m_k1[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh1 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh1 *= exp2(b_gk_last1)[None, :] + else: + b_dh1 *= exp2(b_gk_last1[:, None]) + if STATE_V_FIRST: + b_dh1 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 64: + p_q = q + o_k2[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k2[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k2[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k2[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh2 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh2 *= exp2(b_gk_last2)[None, :] + else: + b_dh2 *= exp2(b_gk_last2[:, None]) + if STATE_V_FIRST: + b_dh2 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 128: + p_q = q + o_k3[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k3[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k3[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k3[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh3 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh3 *= exp2(b_gk_last3)[None, :] + else: + b_dh3 *= exp2(b_gk_last3[:, None]) + if STATE_V_FIRST: + b_dh3 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 192: + p_q = q + o_k4[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k4[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k4[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k4[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh4 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh4 *= exp2(b_gk_last4)[None, :] + else: + b_dh4 *= exp2(b_gk_last4[:, None]) + if STATE_V_FIRST: + b_dh4 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + + if USE_INITIAL_STATE: + if STATE_V_FIRST: + p_dh0 = dh0 + o_v[:, None] * K + o_k1[None, :] + m_dh0 = m_v[:, None] & m_k1[None, :] + else: + p_dh0 = dh0 + o_k1[:, None] * V + o_v[None, :] + m_dh0 = m_k1[:, None] & m_v[None, :] + tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), mask=m_dh0) + if K > 64: + if STATE_V_FIRST: + p_dh1 = dh0 + o_v[:, None] * K + o_k2[None, :] + m_dh1 = m_v[:, None] & m_k2[None, :] + else: + p_dh1 = dh0 + o_k2[:, None] * V + o_v[None, :] + m_dh1 = m_k2[:, None] & m_v[None, :] + tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), mask=m_dh1) + if K > 128: + if STATE_V_FIRST: + p_dh2 = dh0 + o_v[:, None] * K + o_k3[None, :] + m_dh2 = m_v[:, None] & m_k3[None, :] + else: + p_dh2 = dh0 + o_k3[:, None] * V + o_v[None, :] + m_dh2 = m_k3[:, None] & m_v[None, :] + tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), mask=m_dh2) + if K > 192: + if STATE_V_FIRST: + p_dh3 = dh0 + o_v[:, None] * K + o_k4[None, :] + m_dh3 = m_v[:, None] & m_k4[None, :] + else: + p_dh3 = dh0 + o_k4[:, None] * V + o_v[None, :] + m_dh3 = m_k4[:, None] & m_v[None, :] + tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), mask=m_dh3) + + +def chunk_gated_delta_rule_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + B, T, H, K, V, HV = *k.shape, u.shape[-1], u.shape[2] + BT = chunk_size + BV = state_qdq_tile_v(V, state_qdq, state_qdq_block_v) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + if chunk_offsets is None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + assert K <= 256, "current kernel does not support head dimension larger than 256." + + if state_v_first: + h = k.new_empty(B, NT, HV, V, K) + final_state = k.new_zeros(N, HV, V, K, dtype=torch.float32) if output_final_state else None + else: + h = k.new_empty(B, NT, HV, K, V) + final_state = k.new_zeros(N, HV, K, V, dtype=torch.float32) if output_final_state else None + + v_new = torch.empty_like(u) if save_new_value else None + + grid = (triton.cdiv(V, BV) * N * HV,) + chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( + k=k, + v=u, + w=w, + v_new=v_new, + g=g, + gk=gk, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BV=BV, + STATE_V_FIRST=state_v_first, + STATE_QDQ=state_qdq, + ) + return h, v_new, final_state + + +def state_qdq_tile_v(V: int, state_qdq: int, state_qdq_block_v: int | None) -> int: + """Return the V tile width ``BV`` of the forward state kernel. + + Without state quantization this is fla's largest tile. With it, the tile is also the + quantization granularity: one dynamic scale per ``[K, BV]`` block of a head's state. The + default is fla's 64-column tile, i.e. one scale per sequence and head for ``V <= 64`` and two + for the usual ``V == 128``. ``state_qdq_block_v=128`` gives one scale per 128-wide head but + exceeds the register budget where the kernel is limited to two warps (Blackwell) and spills. + """ + if state_qdq == STATE_QDQ_OFF: + return 64 if check_shared_mem("ada") else 32 + if state_qdq not in (STATE_QDQ_FP8_DYNAMIC, STATE_QDQ_INT8_DYNAMIC): + raise ValueError(f"Unsupported state_qdq mode {state_qdq}; expected 0, 1, or 2.") + BV = min(triton.next_power_of_2(V), 64) if state_qdq_block_v is None else state_qdq_block_v + if BV < 16 or BV > STATE_QDQ_MAX_BLOCK_V or BV & (BV - 1): + raise ValueError( + f"state_qdq_block_v must be a power of two in [16, {STATE_QDQ_MAX_BLOCK_V}], got {BV}." + ) + return BV + + +def chunk_gated_delta_rule_bwd_dhu( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + do: torch.Tensor, + dv: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + h0: torch.Tensor | None = None, + dht: torch.Tensor | None = None, + scale: float | None = None, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, + use_graph: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V, HV = *q.shape, do.shape[-1], do.shape[2] + # N: the actual number of sequences in the batch with either equal or variable lengths + BT = chunk_size + assert K <= 256, "current kernel does not support head dimension being larger than 256." + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + if chunk_offsets is None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + + if use_graph: + # [ModelOpt] Imported here: fla.ops.utils.graph is absent from released fla versions. + from fla.ops.utils.graph import get_static_buffer + + if state_v_first: + dh = get_static_buffer("dhu_dh_vf", (B, NT, HV, V, K), q.dtype, q.device) + else: + dh = get_static_buffer("dhu_dh", (B, NT, HV, K, V), q.dtype, q.device) + dh0 = ( + get_static_buffer("dhu_dh0", tuple(h0.shape), torch.float32, h0.device) + if h0 is not None + else None + ) + dv2 = get_static_buffer("dhu_dv2", tuple(dv.shape), dv.dtype, dv.device) + else: + if state_v_first: + dh = q.new_empty(B, NT, HV, V, K) + else: + dh = q.new_empty(B, NT, HV, K, V) + dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None + dv2 = torch.empty_like(dv) + + def grid(meta): + return (triton.cdiv(V, meta["BV"]) * N * HV,) + + chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64[grid]( + q=q, + k=k, + w=w, + g=g, + gk=gk, + dht=dht, + dh0=dh0, + do=do, + dh=dh, + dv=dv, + dv2=dv2, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + scale=scale, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + STATE_V_FIRST=state_v_first, + ) + return dh, dh0, dv2 diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py new file mode 100644 index 00000000000..f34b8a3384b --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py @@ -0,0 +1,725 @@ +# Adapted from: https://github.com/fla-org/flash-linear-attention/blob/516143e31fce/fla/ops/gated_delta_rule/chunk.py +# Adapted with modifications (marked [ModelOpt]): threads state_qdq / state_qdq_block_v through +# the autograd function, applies an optional w_quantizer to the WY tensor w, and imports the +# state kernels from the vendored sibling module. +# +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# 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 warnings + +import fla +import torch +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o +from fla.ops.common.gate import fused_beta_sigmoid, fused_beta_sigmoid_bwd +from fla.ops.cp import FLACPContext +from fla.ops.cp.chunk_delta_h import ( + chunk_gated_delta_rule_bwd_dhu_pre_process, + chunk_gated_delta_rule_fwd_h_pre_process, + compress_h0, + expand_h0, +) +from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra +from fla.ops.gated_delta_rule.gate import gdn_gate_bwd, gdn_gate_chunk_cumsum +from fla.ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd +from fla.ops.utils import chunk_local_cumsum +from fla.ops.utils.constant import RCP_LN2 +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import ( + IS_NVIDIA_HOPPER, + TRITON_ABOVE_3_4_0, + autocast_custom_bwd, + autocast_custom_fwd, + input_guard, +) + +from modelopt.torch.quantization.linear_attention.utils import validate_gdn_quantizer +from modelopt.torch.quantization.nn import TensorQuantizer + +from .fla_chunk_delta_h import ( + STATE_QDQ_FP8_DYNAMIC, + STATE_QDQ_INT8_DYNAMIC, + STATE_QDQ_OFF, + chunk_gated_delta_rule_bwd_dhu, + chunk_gated_delta_rule_fwd_h, +) + + +def chunk_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + w_quantizer: TensorQuantizer | None = None, +): + g_input = g if use_gate_in_kernel else None + if use_gate_in_kernel: + g = gdn_gate_chunk_cumsum( + g=g, + A_log=A_log, + chunk_size=chunk_size, + scale=RCP_LN2, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + else: + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + scale=RCP_LN2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # obtain WY representation. u is actually the new v. + # fused kkt + solve_tril + recompute_w_u + w, u, A = chunk_gated_delta_rule_fwd_intra( + k=k, + v=v, + g=g, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + # [ModelOpt] w is an activation (the WY form of the chunk's keys) that lands in memory here, + # so it is fake-quantized once per forward instead of tile by tile inside the kernel. + if w_quantizer is not None: + w = w_quantizer(w) + + if cp_context is not None: + initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=k, + w=w, + u=u, + g=g, + cu_seqlens=cu_seqlens, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + ) + + if cp_context is not None: + initial_state = compress_h0(initial_state, context=cp_context) + + o = chunk_fwd_o( + q=q, + k=k, + v=v_new, + h=h, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + return g, o, A, final_state, initial_state, g_input, w if w_quantizer is not None else None + + +def chunk_gated_delta_rule_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + g_input: torch.Tensor | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + quantized_w: torch.Tensor | None = None, +): + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # [ModelOpt] Reuse forward QDQ exactly; the validated policy uses identity STE. + # This also avoids calling observers or dynamic scale computation in backward. + if quantized_w is not None: + w = quantized_w + + if cp_context is not None: + initial_state = expand_h0(initial_state, context=cp_context) + + # [ModelOpt] The backward recomputes the forward's chunk states, so it sees the same + # fake-quantized states; the state gradient itself passes straight through the QDQ. + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + ) + dv = chunk_bwd_dv_local( + q=q, + k=k, + g=g, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + if cp_context is not None: + # initial_state is None in the CP mode + # We only need to compute dht of current rank and pass it to the backward kernel + dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=q, + k=k, + w=w, + do=do, + dv=dv, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + dht=dht, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=q, + k=k, + w=w, + g=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dq, dk, dw, dg = chunk_bwd_dqkwg( + q=q, + k=k, + v=v_new, + w=w, + g=g, + h=h, + dv=dv, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dk2, dv, db, dg2 = prepare_wy_repr_bwd( + k=k, + v=v, + beta=beta, + g=g, + A=A, + dw=dw, + du=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + dk.add_(dk2) + dg.add_(dg2) + dg = chunk_local_cumsum( + dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices + ) + dA_log, ddt_bias = None, None + if use_gate_in_kernel: + dg, dA_log, ddt_bias = gdn_gate_bwd(g=g_input, A_log=A_log, dt_bias=dt_bias, dyg=dg) + return dq, dk, dv, db, dg, dh0, dA_log, ddt_bias + + +class ChunkGatedDeltaRuleFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + cp_context: FLACPContext | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + w_quantizer: TensorQuantizer | None = None, + ): + q_rstd, k_rstd = None, None + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + + beta_raw = beta + if use_beta_sigmoid_in_kernel: + beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu + ) + g, o, A, final_state, initial_state, g_input, quantized_w = chunk_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cp_context=cp_context, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + use_gate_in_kernel=use_gate_in_kernel, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + w_quantizer=w_quantizer, + ) + ctx.save_for_backward( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + quantized_w, + ) + ctx.scale = scale + ctx.chunk_size = chunk_size + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.use_beta_sigmoid_in_kernel = use_beta_sigmoid_in_kernel + ctx.allow_neg_eigval = allow_neg_eigval + ctx.cp_context = cp_context + ctx.state_v_first = state_v_first + ctx.use_gate_in_kernel = use_gate_in_kernel + ctx.state_qdq = state_qdq + ctx.state_qdq_block_v = state_qdq_block_v + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + ( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + quantized_w, + ) = ctx.saved_tensors + dq, dk, dv, db, dg, dh0, dA_log, ddt_bias = chunk_gated_delta_rule_bwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A=A, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + cp_context=ctx.cp_context, + chunk_indices=chunk_indices, + state_v_first=ctx.state_v_first, + use_gate_in_kernel=ctx.use_gate_in_kernel, + g_input=g_input, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=ctx.chunk_size, + state_qdq=ctx.state_qdq, + state_qdq_block_v=ctx.state_qdq_block_v, + quantized_w=quantized_w, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + if ctx.use_beta_sigmoid_in_kernel: + db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0) + return ( + dq.to(q), + dk.to(k), + dv.to(v), + dg.to(g), + db.to(beta_raw), + None, + dh0, + None, + None, + None, + None, + None, + None, + None, + dA_log, + ddt_bias, + None, + None, + None, + None, + None, + None, + None, + ) + + +# [ModelOpt] Not registered with fla's backend dispatch: another backend must not take over a +# call that asks for state quantization. +@torch.compiler.disable +def chunk_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + **kwargs, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`. + g (torch.Tensor): + (forget) gating tensor of shape `[B, T, HV]`. + When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay). + When `use_gate_in_kernel=True`, `g` is the raw input before gate activation; + the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, HV, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2norm to the q/k tensor internally. Default: `False`. + use_gate_in_kernel (bool): + Whether to compute the log-space GDN decay internally. + When `True`, the passed `g` is the raw input, and `A_log` must be provided. + The kernel fuses gate activation + chunk cumsum in a single pass. + Default: `False`. + A_log (Optional[torch.Tensor]): + Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`. + dt_bias (Optional[torch.Tensor]): + Bias added to `g` before activation, of shape `[HV]`. + Only used when `use_gate_in_kernel=True`. + use_beta_sigmoid_in_kernel (bool): + Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel. + - If `True`, the passed `beta` acts as the raw beta logits. + - If `False`, `beta` is expected to already be in post-sigmoid space. + Default: `False`. + allow_neg_eigval (bool): + Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`. + Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case + the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`. + state_v_first (Optional[bool]): + Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + chunk_indices (Optional[torch.LongTensor]): + Pre-computed chunk indices for variable-length inputs. + If provided, they are used directly instead of being computed from `cu_seqlens`. Default: `None`. + cp_context (Optional[FLACPContext]): + Context parallel context for distributed training across multiple devices. + When provided, `initial_state` and `output_final_state` are not supported, + and `cu_seqlens` will be overridden by the context. Default: `None`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda') + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if fla.__version__ != "0.5.1": + raise RuntimeError(f"ModelOpt GDN requires fla-core==0.5.1, got {fla.__version__}.") + if "transpose_state_layout" in kwargs: + if state_v_first: + raise ValueError( + "Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`." + ) + warnings.warn( + "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", + DeprecationWarning, + stacklevel=2, + ) + state_v_first = kwargs.pop("transpose_state_layout") + + # Validate head dimensions + if q.shape[2] != k.shape[2]: + raise ValueError( + f"q and k must have the same number of heads, " + f"but got q.shape[2]={q.shape[2]} and k.shape[2]={k.shape[2]}" + ) + H, HV = q.shape[2], v.shape[2] + if HV % H != 0: + raise ValueError( + f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by " + f"num_heads (H={H}), but got HV % H = {HV % H}" + ) + + if "head_first" in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) + + chunk_size = kwargs.pop("chunk_size", 64) + if chunk_size != 64: + raise ValueError("ModelOpt GDN supports only chunk_size=64; FLA WY backward assumes 64.") + + # [ModelOpt] state_qdq: 0 keeps fla's numerics; 1 fake-quantizes the state carried between + # chunks to FP8 E4M3; 2 uses signed narrow-range INT8. Both use a dynamic scale per + # [K, state_qdq_block_v] tile of each head. + state_qdq = kwargs.pop("state_qdq", STATE_QDQ_OFF) + state_qdq_block_v = kwargs.pop("state_qdq_block_v", None) + # w_quantizer: dynamic FP8 TensorQuantizer applied to the WY tensor + # ``w`` of shape [B, T, HV, K] before it multiplies the state, emulating an FP8 x FP8 matmul. + w_quantizer = kwargs.pop("w_quantizer", None) + use_gate_in_kernel = kwargs.pop("use_gate_in_kernel", False) + A_log = kwargs.pop("A_log", None) + dt_bias = kwargs.pop("dt_bias", None) + if kwargs: + raise TypeError(f"Unexpected keyword arguments: {', '.join(sorted(kwargs))}") + if state_qdq not in (STATE_QDQ_OFF, STATE_QDQ_FP8_DYNAMIC, STATE_QDQ_INT8_DYNAMIC): + raise ValueError(f"`state_qdq` must be 0, 1, or 2, got {state_qdq}.") + if w_quantizer is not None: + validate_gdn_quantizer(w_quantizer, name="gdn_w_quantizer") + if state_qdq == STATE_QDQ_FP8_DYNAMIC and ( + not q.is_cuda or torch.cuda.get_device_capability(q.device) < (8, 9) + ): + raise RuntimeError("GDN state QDQ requires native E4M3 conversion on CUDA SM89 or newer.") + if (state_qdq != STATE_QDQ_OFF or w_quantizer is not None) and cp_context is not None: + raise ValueError("State or w quantization is not supported together with `cp_context`.") + + if cp_context is not None: + assert initial_state is None, "Initial state is not supported for CP" + assert output_final_state is False, "Output final state is not supported for CP" + assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" + cu_seqlens = cp_context.cu_seqlens + if cp_context.cu_seqlens_cpu is not None: + cu_seqlens_cpu = cp_context.cu_seqlens_cpu + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if use_gate_in_kernel: + assert A_log is not None, "A_log must be provided when use_gate_in_kernel=True." + if allow_neg_eigval and not use_beta_sigmoid_in_kernel: + raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") + + if scale is None: + scale = k.shape[-1] ** -0.5 + # [ModelOpt] Hopper's TileLang backward needs BF16 and equal head counts. Expand outside + # custom autograd so repeat_interleave reduces q/k gradients back to the original heads. + if IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0: + if any(x.dtype != torch.bfloat16 for x in (q, k, v)): + raise ValueError("Hopper with Triton >= 3.4 requires BF16 q/k/v for GDN training.") + if H != HV: + q = q.repeat_interleave(HV // H, dim=2) + k = k.repeat_interleave(HV // H, dim=2) + o, final_state = ChunkGatedDeltaRuleFunction.apply( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + state_v_first, + cu_seqlens, + cu_seqlens_cpu, + chunk_indices, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + A_log, + dt_bias, + use_beta_sigmoid_in_kernel, + allow_neg_eigval, + cp_context, + chunk_size, + state_qdq, + state_qdq_block_v, + w_quantizer, + ) + return o, final_state + + +chunk_gdn = chunk_gated_delta_rule diff --git a/modelopt/torch/kernels/quantization/linear_attention/int8.py b/modelopt/torch/kernels/quantization/linear_attention/int8.py new file mode 100644 index 00000000000..d9e16a29dc8 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/int8.py @@ -0,0 +1,28 @@ +# 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. + +"""Signed narrow-range INT8 QDQ for recurrent-state tiles.""" + +import triton +import triton.language as tl +from triton.language.extra.cuda import libdevice + +__all__ = [] + + +@triton.jit +def int8_scalar_qdq(value, scale): + codes = libdevice.nearbyint(tl.div_rn(value, scale)) + return tl.minimum(tl.maximum(codes, -127.0), 127.0) * scale diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index f8847aa1fc1..e8db5c8446a 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -161,6 +161,8 @@ from modelopt.torch.opt.config_loader import load_config from modelopt.torch.utils.network import ConstructorLike +from .linear_attention.config import LinearAttentionPolicyEntry + class QuantizerCfgEntry(ModeloptBaseConfig): """A single entry in a ``quant_cfg`` list.""" @@ -1567,6 +1569,11 @@ def _dict_to_entry(key: str, value) -> list[dict[str, Any]]: class QuantizeConfig(ModeloptBaseConfig): """Default configuration for ``quantize`` mode.""" + linear_attention: list[LinearAttentionPolicyEntry] = ModeloptField( + default=[], + title="Execution policies for supported linear-attention training modules", + ) + quant_cfg: QuantizeQuantCfgType = ModeloptField( default=[{"quantizer_name": "*", "cfg": {"num_bits": 8, "axis": None}}], title="Quantization configuration", diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 7fbd7d10dea..db194a0f465 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -28,6 +28,7 @@ from modelopt.torch.opt.dynamic import _DMRegistryCls from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict from modelopt.torch.utils import get_unwrapped_name +from modelopt.torch.utils.logging import print_rank_0 from .config import ( QuantizeConfig, @@ -36,6 +37,7 @@ _QuantizeExportConfig, normalize_quant_cfg_list, ) +from .linear_attention.config import LinearAttentionConfig from .nn import ( QuantModule, QuantModuleRegistry, @@ -67,6 +69,7 @@ def convert_to_quantized_model(model: ModelLikeModule, config: QuantizeConfig) - replace_quant_module(model, version=ModeloptStateManager(model).state_version) set_quantizer_by_cfg(model, config.get("quant_cfg", [])) + _apply_linear_attention_policy(model, config) metadata = {} update_quantize_metadata(model, config, metadata) @@ -136,13 +139,22 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: details regarding how MCore sharded checkpoint is restored, see modelopt.torch.opt.plugins.mcore_dist_checkpointing.restore_sharded_modelopt_state. """ + _apply_linear_attention_policy(model, config, metadata.get("linear_attention")) if "quantizer_state" not in metadata: # MCore sharded checkpoint (`torch-dist`) has its quantizer_state stored as the # extra_state of `QuantModule`. The quantizer_state is resumed with # QuantModule.set_extra_state(). return model - quantizer_state_dict = metadata["quantizer_state"] + quantizer_state_dict = dict(metadata["quantizer_state"]) + if "linear_attention" not in metadata: + # Older checkpoints predate these disabled handles; preserve their baseline path. + for name, module in _linear_attention_modules(model).items(): + for handle in module.linear_attention_quantizer_names: + key = f"{name}.{handle}" if name else handle + quantizer = getattr(module, handle) + if key not in quantizer_state_dict and not quantizer.is_enabled: + quantizer_state_dict[key] = quantizer.get_modelopt_state() unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys() extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys() @@ -195,6 +207,12 @@ def update_quantize_metadata( model: nn.Module, config: QuantizeConfig, metadata: MetadataDict ) -> None: """Update the quantizer state in the metadata dict.""" + policies = { + name: module.linear_attention_config.model_dump() + for name, module in _linear_attention_modules(model).items() + } + if policies: + metadata["linear_attention"] = policies metadata["quantizer_state"] = quantizer_state(model) if shared_state_metadata := SharedWeightGlobalAmaxState.metadata(model): metadata["shared_quant_states"] = shared_state_metadata @@ -202,6 +220,43 @@ def update_quantize_metadata( metadata.pop("shared_quant_states", None) +def _linear_attention_modules(model): + # Optional framework plugins import conversion; defer this import to avoid that cycle. + from .plugins.linear_attention import _LinearAttentionQuantMixin + + return { + get_unwrapped_name(name, model): module + for name, module in model.named_modules() + if isinstance(module, _LinearAttentionQuantMixin) + } + + +def _apply_linear_attention_policy(model, config, saved_policies=None): + modules = _linear_attention_modules(model) + # getattr also handles pickled pre-extension config objects with no field in __dict__. + for entry in getattr(config, "linear_attention", []): + matches = [name for name in modules if fnmatch.fnmatch(name, entry.module_name)] + if not matches: + raise ValueError( + f"linear_attention rule {entry.module_name!r} matches no supported modules" + ) + for name in matches: + modules[name].linear_attention_config = entry.cfg.model_copy(deep=True) + if saved_policies is not None: + if saved_policies.keys() != modules.keys(): + raise ApplyModeError( + "Saved linear_attention policies do not match the restored modules" + ) + for name, policy in saved_policies.items(): + modules[name].linear_attention_config = LinearAttentionConfig(**policy) + for name, module in modules.items(): + module.validate_linear_attention() + if getattr(config, "linear_attention", []): + print_rank_0( + f"Linear attention {name!r}: {module.linear_attention_config.model_dump()}" + ) + + def quantizer_state(model: nn.Module) -> dict[str, Any]: """Returns the quantizer state dict describing the quantizer states in the model.""" return { diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py new file mode 100644 index 00000000000..305bb2cd81f --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -0,0 +1,22 @@ +# 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. + +"""Numerical policies and differentiable kernels for linear attention.""" + +from .config import * +from .decode import * +from .decode_prefill import * +from .kda import * +from .prefill import * diff --git a/modelopt/torch/quantization/linear_attention/_chunk_prefill.py b/modelopt/torch/quantization/linear_attention/_chunk_prefill.py new file mode 100644 index 00000000000..fddff0303ae --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/_chunk_prefill.py @@ -0,0 +1,146 @@ +# 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. + +"""Exact chunked prefix computation for decode-aware training.""" + +import torch + +from .utils import _prepare, _state_qdq + +__all__ = [] + + +def chunk_gdn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + chunk_size: int = 64, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + state_v_first: bool = False, + state_qdq: bool = False, + state_qdq_block_v: int = 64, + state_format: str = "fp8_e4m3", +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute a GDN prefix with optional QDQ on the initial state and chunk writes.""" + if g.ndim != 3 or chunk_size <= 0: + raise ValueError("chunk_gdn requires scalar GDN gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format) + pieces = [] + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + gc = g[b, lo:hi].transpose(0, 1).cumsum(-1) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + # Mask before exp: upper-triangle positive differences can overflow for long decay. + causal = torch.ones(hi - lo, hi - lo, device=q.device, dtype=torch.bool).tril() + decay = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() + lower = (bc * (kc @ kc.transpose(-1, -2)) * decay).tril(-1) + matrix = lower + torch.eye(hi - lo, device=q.device, dtype=q.dtype) + gate = gc.exp().unsqueeze(-1) + rhs = torch.cat((bc * vc, bc * kc * gate), dim=-1) + solved = torch.linalg.solve_triangular(matrix, rhs, upper=False, unitriangular=True) + u, w = solved.split((v.shape[-1], k.shape[-1]), dim=-1) + updated_values = u - w @ state + local_scores = ((qc * scale @ kc.transpose(-1, -2)) * decay).tril() + output = (qc * (scale * gc.exp()).unsqueeze(-1)) @ state + pieces.append((output + local_scores @ updated_values).transpose(0, 1)) + weighted_keys = kc * (gc[..., -1:] - gc).exp().unsqueeze(-1) + state = state * gc[..., -1].exp()[:, None, None] + state = state + weighted_keys.transpose(-1, -2) @ updated_values + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format) + outputs.append(torch.cat(pieces)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final + + +def chunk_kda( + q, + k, + v, + g, + beta, + *, + chunk_size=64, + scale=None, + initial_state=None, + cu_seqlens=None, + state_v_first=False, + state_qdq=False, + state_qdq_block_v=64, + state_format="fp8_e4m3", +): + """Compute a KDA prefix with optional QDQ on the initial state and chunk writes.""" + if g.ndim != 4 or chunk_size <= 0: + raise ValueError("chunk_kda requires per-key gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format) + pieces = [] + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + prefix = g[b, lo:hi].transpose(0, 1).cumsum(-2) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + gate = prefix.exp() + lower_rows, score_rows = [], [] + for row in range(hi - lo): + decayed_keys = ( + kc[..., : row + 1, :] + * (prefix[..., row : row + 1, :] - prefix[..., : row + 1, :]).exp() + ) + right = decayed_keys.transpose(-1, -2) + left = kc[..., row : row + 1, :] * bc[..., row : row + 1, :] + interaction = left @ right + score = qc[..., row : row + 1, :] * scale @ right + padding = hi - lo - row - 1 + lower_rows.append(torch.nn.functional.pad(interaction, (0, padding))) + score_rows.append(torch.nn.functional.pad(score, (0, padding))) + lower = torch.cat(lower_rows, dim=-2).tril(-1) + scores = torch.cat(score_rows, dim=-2) + identity = torch.eye(hi - lo, device=q.device, dtype=q.dtype).expand_as(lower) + inverse = torch.linalg.solve_triangular( + identity + lower, identity, upper=False, unitriangular=True + ) + u = inverse @ (bc * vc) + w = inverse @ (bc * kc * gate) + updated = u - w @ state + pieces.append(((qc * scale * gate) @ state + scores @ updated).transpose(0, 1)) + weighted_keys = kc * (prefix[..., -1:, :] - prefix).exp() + state = state * gate[..., -1, :, None] + weighted_keys.transpose(-1, -2) @ updated + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format) + outputs.append(torch.cat(pieces)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py new file mode 100644 index 00000000000..33756c5b459 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -0,0 +1,110 @@ +# 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. + +"""Saved execution policy for GDN training-time numerical emulation.""" + +from typing import Literal + +from pydantic import Field, model_validator + +from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField + +__all__ = [ + "LinearAttentionConfig", + "LinearAttentionDecodeConfig", + "LinearAttentionPolicyEntry", + "LinearAttentionReplayConfig", +] + + +class _StateConfig(ModeloptBaseConfig): + mode: Literal["chunk"] = ModeloptField(default="chunk") + block_v: Literal[16, 32, 64, 128] = ModeloptField(default=64) + quantize_initial: Literal[True] = ModeloptField(default=True) + + +class _SolveConfig(ModeloptBaseConfig): + method: Literal["exact"] = ModeloptField(default="exact") + + +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) + state_codec: Literal["tile", "int8_hadamard32"] = ModeloptField(default="tile") + 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") + if self.state_codec == "int8_hadamard32" and self.prefill_state_qdq: + raise ValueError( + "Hadamard state QDQ starts at decode handoff; disable prefill_state_qdq" + ) + return self + + +class LinearAttentionConfig(ModeloptBaseConfig): + """GDN chunk-64 policy; unsupported numerical modes fail config validation. + + ``state.block_v`` defines one dynamic scale per ``[Dk, block_v]`` tile of each + sequence/head. The initial state and every chunk's final state are rounded when + ``gdn_state_quantizer`` is enabled. Outputs use the incoming rounded state. + Decode's ``int8_hadamard32`` codec instead fixes scales to one key channel and + 32 values; ``state.block_v`` remains the execution tile width. + """ + + schema_version: Literal[1] = ModeloptField(default=1) + backend: Literal["fla", "matmul"] = ModeloptField(default="fla") + chunk_size: Literal[64] = ModeloptField(default=64) + state: _StateConfig = ModeloptField(default=_StateConfig()) + solve: _SolveConfig = ModeloptField(default=_SolveConfig()) + decode: LinearAttentionDecodeConfig | None = ModeloptField(default=None) + + @model_validator(mode="after") + def _validate_decode_backend(self): + if (self.backend == "matmul") != (self.decode is not None): + raise ValueError("The exact-prefix matmul backend requires an explicit decode policy") + if self.decode is not None and self.decode.state_codec == "int8_hadamard32": + if self.state.block_v < 32: + raise ValueError("int8_hadamard32 requires block_v >= 32") + return self + + +class LinearAttentionPolicyEntry(ModeloptBaseConfig): + """Assign a complete policy to supported modules matching ``module_name``. + + Rules apply in order: the last match wins, without merging nested fields. + A rule must match at least one supported linear-attention module. + """ + + module_name: str = Field(...) + cfg: LinearAttentionConfig = ModeloptField(default=LinearAttentionConfig()) diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py new file mode 100644 index 00000000000..80905958bb9 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -0,0 +1,427 @@ +# 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 +from .utils import _fp8_quantize + +__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 + value_basis: str = "identity" + + @property + def cursor(self): + """Number of encoded updates since the last anchor refresh.""" + return len(self.entries) + + def reconstruct(self, *, original_basis=True): + """Replay entries, returning the original value basis unless explicitly disabled.""" + 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 + ) + if original_basis and self.value_basis == "hadamard32": + state = _hadamard32(state) + return state + + +def _hadamard32(value): + """Apply the orthonormal Sylvester transform to contiguous 32-value groups.""" + shape = value.shape + for width in (1, 2, 4, 8, 16): + pairs = value.reshape(*shape[:-1], -1, 2, width) + left, right = pairs.unbind(-2) + value = torch.stack((left + right, left - right), dim=-2).reshape(shape) + return value * (32**-0.5) + + +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, state_format="fp8_e4m3", state_codec="tile"): + if not enabled: + return EncodedLinearAttentionTensor(value, None, "identity", None) + if state and state_codec == "int8_hadamard32": + with torch.no_grad(): + groups = value.float().reshape(*value.shape[:-1], -1, 32) + scales = (groups.abs().amax(-1) / 127).clamp_min(6e-8) + normalized = groups / scales.unsqueeze(-1) + codes = normalized.sign() * (normalized.abs() + 0.5).floor() + metadata = scales.to(torch.float16) + decoded = (codes.clamp(-127, 127) * metadata.float().unsqueeze(-1)).reshape_as(value) + return EncodedLinearAttentionTensor( + _IdentityGradient.apply(value, decoded.to(value.dtype)), metadata, "int8", 32 + ) + rounded, scales = [], [] + for part in value.split(block_v, dim=-1): + if state_format == "fp8_e4m3": + tensor = part.flatten(-2) if state else part + axis = tuple(range(tensor.ndim - 1)) or None + decoded, scale = _fp8_quantize(tensor, axis) + rounded.append(decoded.reshape_as(part)) + scales.append(scale.squeeze(-1)) + else: + with torch.no_grad(): + axes = (-2, -1) if state else (-1,) + part = part.float() + amax = part.abs().amax(dim=axes, keepdim=True) + scale = torch.where(amax > 0, amax / 127.0, torch.ones_like(amax)) + rounded.append((part / scale).round().clamp(-127, 127) * scale) + scales.append(scale[..., 0, 0] if state else scale[..., 0]) + decoded = torch.cat(rounded, dim=-1).to(value.dtype) + if state_format == "int8": + decoded = _IdentityGradient.apply(value, decoded) + return EncodedLinearAttentionTensor(decoded, torch.stack(scales, dim=-1), state_format, block_v) + + +def _signature(config, state_qdq, block_v, state_format): + return ( + config.model_dump_json(exclude={"implementation", "prefill_state_qdq"}) + + f"/{state_qdq}/{block_v}/{state_format}" + ) + + +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, state_format +): + 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") + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") + hadamard = config.state_codec == "int8_hadamard32" + if hadamard: + if state_qdq and state_format != "int8": + raise ValueError("int8_hadamard32 requires INT8 state quantization") + if v.shape[-1] % 32 or block_v < 32: + raise ValueError("int8_hadamard32 requires Dv divisible by 32 and block_v >= 32") + signature = _signature(config, state_qdq, block_v, state_format) + 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, state_format=state_format), + (), + 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: + initial = _hadamard32(carry.anchor.values) if hadamard else carry.anchor.values + anchor = _encode( + initial, + state_qdq and config.quantize_initial, + block_v, + state=True, + state_format=state_format, + state_codec=config.state_codec, + ) + carry = LinearAttentionCarry( + anchor, (), carry.position, True, signature, "hadamard32" if hadamard else "identity" + ) + return carry, signature + + +def recurrent_decode_reference( + q, + k, + v, + g, + beta, + *, + config: LinearAttentionDecodeConfig, + state_qdq=False, + state_format="fp8_e4m3", + 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, state_format + ) + 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 + if config.state_codec == "int8_hadamard32": + v = _hadamard32(v) + 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, carry.value_basis + ) + state = carry.reconstruct(original_basis=False) + 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.replay is not None: + entries = (*entries, ReplayEntry(key, update, gate)) + refresh = config.replay is None or len(entries) == config.replay.window + if refresh: + anchor = _encode( + working, + state_qdq, + block_v, + state=True, + state_format=state_format, + state_codec=config.state_codec, + ) + entries = () + stored = anchor.values + else: + anchor = carry.anchor + stored = working + next_carry = LinearAttentionCarry( + anchor, entries, carry.position + 1, True, signature, carry.value_basis + ) + read = working if config.readout == "working" else stored + outputs.append(_sum_keys(q[t].unsqueeze(-1) * read) * scale) + carry = next_carry + output = torch.stack(outputs) + if config.state_codec == "int8_hadamard32": + output = _hadamard32(output) + return output, carry + + +def recurrent_decode( + q, + k, + v, + g, + beta, + *, + config: LinearAttentionDecodeConfig, + state_qdq=False, + state_format="fp8_e4m3", + 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, + state_format=state_format, + 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, state_format + ) + 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 + + if config.state_codec == "int8_hadamard32": + v = _hadamard32(v) + 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(original_basis=False), + state_qdq=state_qdq, + state_format=state_format, + state_codec=config.state_codec, + 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 + if config.state_codec == "int8_hadamard32": + output = _hadamard32(output) + state_format = state_format if state_qdq else "identity" + state_block_v = 32 if config.state_codec == "int8_hadamard32" else block_v + if not replay: + anchor = EncodedLinearAttentionTensor( + final, + anchor_scales if state_qdq else None, + state_format, + state_block_v if state_qdq else None, + ) + return output, LinearAttentionCarry( + anchor, (), carry.position + len(q), True, signature, carry.value_basis + ) + if carry.cursor + len(q) >= window: + anchor = EncodedLinearAttentionTensor( + last_anchor, + anchor_scales if state_qdq else None, + state_format, + state_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, carry.value_basis + ) 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..0a0f69b4ae7 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/decode_prefill.py @@ -0,0 +1,176 @@ +# 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 ._chunk_prefill import chunk_gdn, chunk_kda +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, + *, + policy, + state_qdq, + state_format, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + state_v_first, + output_dtype, + prefill_lengths, +): + 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]) + prefix_fn = chunk_kda if g.ndim == 4 else chunk_gdn + with torch.autocast(device_type=q.device.type, enabled=False): + output, final = prefix_fn( + *packed, + state_qdq=state_qdq and policy.decode.prefill_state_qdq, + state_format=state_format, + scale=scale, + initial_state=states[active], + cu_seqlens=torch.tensor(offsets), + state_v_first=False, + chunk_size=policy.chunk_size, + state_qdq_block_v=policy.state.block_v, + ) + 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, + state_format=state_format, + 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 new file mode 100644 index 00000000000..446dd6f7d0d --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differentiable KDA prefill with stable per-channel decay interactions.""" + +import torch +import torch.nn.functional as F + +from .decode_prefill import _decode_prefill + +__all__ = ["matmul_kda"] + + +def matmul_kda( + q, + k, + v, + g, + beta, + *, + policy, + state_qdq=False, + state_format="fp8_e4m3", + scale=None, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=False, + use_beta_sigmoid_in_kernel=False, + allow_neg_eigval=False, + A_log=None, # noqa: N803 - match the FLA kernel signature + dt_bias=None, + safe_gate=False, + lower_bound=None, + cu_seqlens=None, + cu_seqlens_cpu=None, + state_v_first=False, + chunk_size=64, + cp_context=None, + disable_recompute=False, + return_intermediate_states=False, + prefill_lengths=None, +): + """Normalize KDA inputs and run exact prefix plus configured suffix recurrence. + + Gates use the loaded FLA model's activation formula and per-key log retention. + """ + if policy.backend != "matmul" or chunk_size != policy.chunk_size: + raise ValueError("matmul_kda requires backend='matmul' and its configured chunk size") + if cp_context is not None or disable_recompute or return_intermediate_states: + raise NotImplementedError( + "KDA matmul does not support CP or FLA recompute/intermediate flags" + ) + if allow_neg_eigval and not use_beta_sigmoid_in_kernel: + raise ValueError("allow_neg_eigval requires use_beta_sigmoid_in_kernel") + if safe_gate and use_gate_in_kernel and (lower_bound is None or not -5 <= lower_bound < 0): + raise ValueError("safe_gate requires a lower_bound in [-5, 0)") + if lower_bound is not None and lower_bound >= 0: + raise ValueError("lower_bound must be negative") + output_dtype = q.dtype + dtype = torch.float64 if output_dtype == torch.float64 else torch.float32 + q, k, v, g, beta = (x.to(dtype) for x in (q, k, v, g, beta)) + if g.ndim != 4: + raise ValueError("matmul_kda requires per-key-channel KDA log gates") + if use_qk_l2norm_in_kernel: + q, k = (x * (x.square().sum(-1, keepdim=True) + 1e-6).rsqrt() for x in (q, k)) + if use_gate_in_kernel: + if A_log is None: + raise ValueError("Fused KDA gate requires A_log") + if dt_bias is not None: + g = g + dt_bias.to(dtype).reshape(g.shape[-2:]) + rate = A_log.to(dtype).exp().reshape(g.shape[-2], 1) + g = -rate * F.softplus(g) if lower_bound is None else lower_bound * (rate * g).sigmoid() + if use_beta_sigmoid_in_kernel: + beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) + if policy.decode is not None: + return _decode_prefill( + q, + k, + v, + g, + beta, + policy=policy, + state_qdq=state_qdq, + state_format=state_format, + 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, + ) + raise ValueError("An explicit decode policy is required") diff --git a/modelopt/torch/quantization/linear_attention/prefill.py b/modelopt/torch/quantization/linear_attention/prefill.py new file mode 100644 index 00000000000..4a602f8f5bc --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/prefill.py @@ -0,0 +1,94 @@ +# 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. + +"""Batched differentiable GDN prefill with materialized numerical boundaries.""" + +import torch +import torch.nn.functional as F + +from .config import LinearAttentionConfig +from .decode_prefill import _decode_prefill + +__all__ = ["matmul_gdn"] + + +def matmul_gdn( + q, + k, + v, + g, + beta, + *, + policy: LinearAttentionConfig, + state_qdq=False, + state_format="fp8_e4m3", + scale=None, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=False, + use_beta_sigmoid_in_kernel=False, + allow_neg_eigval=False, + A_log=None, # noqa: N803 - match the FLA kernel signature + dt_bias=None, + cu_seqlens=None, + cu_seqlens_cpu=None, + state_v_first=False, + chunk_size=64, + cp_context=None, + prefill_lengths=None, +): + """Normalize GDN inputs and run exact prefix plus configured suffix recurrence. + + Uses FP32 working arithmetic (FP64 for double inputs), with floating QDQ state. + Prefixes use exact chunk algebra; no configurable operand rounding is enabled. + """ + if policy.backend != "matmul" or chunk_size != policy.chunk_size: + raise ValueError("matmul_gdn requires backend='matmul' and its configured chunk size") + if cp_context is not None: + raise NotImplementedError("GDN matmul emulation does not support context parallelism") + output_dtype = q.dtype + dtype = torch.float64 if output_dtype == torch.float64 else torch.float32 + q, k, v, g, beta = (x.to(dtype) for x in (q, k, v, g, beta)) + if use_qk_l2norm_in_kernel: + q, k = (x * (x.square().sum(-1, keepdim=True) + 1e-6).rsqrt() for x in (q, k)) + if use_gate_in_kernel: + if A_log is None or dt_bias is None: + raise ValueError("Fused GDN gate requires A_log and dt_bias") + g = -A_log.to(dtype).exp() * F.softplus(g + dt_bias.to(dtype)) + if use_beta_sigmoid_in_kernel: + 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, + policy=policy, + state_qdq=state_qdq, + state_format=state_format, + 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, + ) + raise ValueError("An explicit decode policy is required") diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py new file mode 100644 index 00000000000..68353bd28df --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared helpers for linear-attention quantization.""" + +from __future__ import annotations + +from itertools import pairwise +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from ..nn import TensorQuantizer + +__all__ = [] + + +def validate_gdn_quantizer( + quantizer: TensorQuantizer, + *, + name: str, + num_bits: tuple[int | tuple[int, int], ...] = ((4, 3),), +) -> None: + """Check supported formats and the custom backward's identity STE.""" + # Numerical helpers load while QuantizeConfig initializes; defer the quantizer import. + from ..nn import TensorQuantizer + + if not isinstance(quantizer, TensorQuantizer): + raise ValueError(f"{name} requires a single TensorQuantizer") + if not ( + quantizer._dynamic + and quantizer.num_bits in num_bits + and (quantizer.num_bits != 8 or (not quantizer.unsigned and quantizer.narrow_range)) + and quantizer.block_sizes is None + and quantizer.fake_quant + and quantizer._pass_through_bwd + and not quantizer.rotate_is_enabled + and quantizer.pre_quant_scale is None + and quantizer.backend is None + and not quantizer._bias + and not quantizer._use_constant_amax + ): + raise ValueError( + f"{name} supports only dynamic fake quantization with num_bits in {num_bits}, " + "pass_through_bwd=True, no block_sizes, rotation, pre-scaling, bias, constant " + "amax, or custom backend. Other gradient rules and formats are not implemented." + ) + + +def _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first): + if q.ndim != 4 or k.shape != q.shape or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k must have shape [B,T,Hk,Dk] and v shape [B,T,Hv,Dv]") + batch, length, heads, keys = q.shape + value_heads, values = v.shape[2:] + if length == 0 or value_heads % heads: + raise ValueError("nonempty sequences and Hv divisible by Hk are required") + if beta.shape != (batch, length, value_heads) or g.shape not in ( + beta.shape, + (*beta.shape, keys), + ): + raise ValueError("beta must be [B,T,Hv]; g must be [B,T,Hv] or [B,T,Hv,Dk]") + if cu_seqlens is None: + sequences = [(b, 0, length) for b in range(batch)] + else: + boundaries = cu_seqlens.tolist() + if ( + batch != 1 + or len(boundaries) < 2 + or boundaries[0] != 0 + or boundaries[-1] != length + or any(a >= b for a, b in pairwise(boundaries)) + ): + raise ValueError("cu_seqlens must partition a packed batch of size one") + sequences = [(0, a, b) for a, b in pairwise(boundaries)] + if initial_state is None: + state = q.new_zeros(len(sequences), value_heads, keys, values) + else: + state = initial_state.transpose(-1, -2) if state_v_first else initial_state + if state.shape != (len(sequences), value_heads, keys, values): + raise ValueError("initial_state shape does not match sequence/head dimensions") + return ( + q.repeat_interleave(value_heads // heads, dim=2), + k.repeat_interleave(value_heads // heads, dim=2), + state, + sequences, + ) + + +def _fp8_quantize(value: torch.Tensor, axis): + """Apply ModelOpt FP8 QDQ and return its detached dequantization scales.""" + # QuantizeConfig imports this module before the quantizer classes are initialized. + from ..config import QuantizerAttributeConfig + from ..nn import TensorQuantizer + + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=(4, 3), type="dynamic", axis=axis, pass_through_bwd=True) + ) + quantized = quantizer(value) + amax = quantizer._get_amax(value).float() + safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) + return quantized, torch.div(448.0, safe_amax).reciprocal() + + +def _state_qdq(state: torch.Tensor, block_v: int = 64, state_format: str = "fp8_e4m3"): + """Dynamic state-tile QDQ with detached scales and identity STE.""" + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") + if block_v not in (16, 32, 64, 128): + raise ValueError("block_v must be 16, 32, 64, or 128") + if state_format == "fp8_e4m3": + axis = tuple(range(state.ndim - 2)) or None + return torch.cat( + [ + _fp8_quantize(tile.flatten(-2), axis)[0].reshape_as(tile) + for tile in state.split(block_v, dim=-1) + ], + dim=-1, + ) + with torch.no_grad(): + rounded = [] + for tile in state.float().split(block_v, dim=-1): + amax = tile.abs().amax(dim=(-2, -1), keepdim=True) + limit = 127.0 + scale = torch.where(amax > 0, amax / limit, torch.ones_like(amax)) + normalized = (tile / scale).clamp(-limit, limit) + codes = normalized.round() + rounded.append(codes * scale) + quantized = torch.cat(rounded, dim=-1).to(state.dtype) + return state + (quantized - state).detach() diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index d7ef04f750d..5aac82cacfe 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -32,6 +32,7 @@ from modelopt.torch.opt.utils import forward_with_reshard from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.conversion import ( + _apply_linear_attention_policy, preserve_quantizer_attributes_context, set_quantizer_attributes_partial, set_quantizer_by_cfg, @@ -329,6 +330,7 @@ def forward_loop(model) -> None: else: # Already quantized, so lets apply the quant_cfg from the config set_quantizer_by_cfg(model, quantize_config.quant_cfg) + _apply_linear_attention_policy(model, quantize_config) # Fail before calibration rather than after exporting an unquantized checkpoint. _check_weight_quantization_took_effect(model, quantize_config) return calibrate(model, config.get("algorithm"), forward_loop=forward_loop) diff --git a/modelopt/torch/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py index 22b4bc2e3cc..e1b2ec7a1ae 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -39,6 +39,8 @@ from .attention import * from .custom import * +from .gdn import * +from .kda import * with import_plugin("diffusers"): from .diffusion.diffusers import * diff --git a/modelopt/torch/quantization/plugins/gdn.py b/modelopt/torch/quantization/plugins/gdn.py new file mode 100644 index 00000000000..8d0fe616390 --- /dev/null +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -0,0 +1,123 @@ +# 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. + +"""Fake quantization of the GatedDeltaNet (GDN) recurrent state. + +The chunked gated-delta-rule kernel keeps each head's ``[K, V]`` recurrent state in fp32 inside +one Triton launch and carries it from chunk to chunk. To emulate a deployment that stores that +state in FP8 or INT8, ModelOpt runs an adapted copy of the kernel +(:mod:`modelopt.torch.kernels.quantization.linear_attention`) that fake-quantizes the state to +E4M3 or signed narrow-range INT8 at the end of every chunk, with a scale computed inside the +kernel from the state itself. +The backward pass recomputes the same quantized states and passes the state gradient straight +through the quantization, so QAT and QAD train against the quantized recurrence. A second +quantizer covers ``w``, the WY-transformed keys that multiply the state; ``w`` is a regular tensor, +so it is fake-quantized by the ``TensorQuantizer`` itself before the kernel reads it. +""" + +from collections.abc import Callable +from functools import partial +from typing import Any + +import torch + +from ..linear_attention.prefill import matmul_gdn +from .linear_attention import _LinearAttentionQuantMixin + +__all__ = ["GatedDeltaNetStateQuantMixin"] + +GatedDeltaRuleFn = Callable[..., tuple[torch.Tensor, torch.Tensor | None]] + + +def _fla_chunk_gated_delta_rule() -> GatedDeltaRuleFn: + # FLA is an optional, heavy dependency needed only when quantization is enabled. + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + + return chunk_gated_delta_rule + + +def _state_qdq_chunk_gated_delta_rule() -> GatedDeltaRuleFn: + # Imported on first use: flash-linear-attention is a heavy optional dependency that only the + # enabled quantizer needs, and importing it warns on machines without a GPU. + try: + from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( + chunk_gated_delta_rule, + ) + except ImportError as e: + raise RuntimeError( + "GDN fake quantization needs Triton and fla-core==0.5.1 on a CUDA " + f"device; importing the state-quantizing kernel failed with {e!r}." + ) from e + return chunk_gated_delta_rule + + +class GatedDeltaNetStateQuantMixin(_LinearAttentionQuantMixin): + """Adds ``gdn_state_quantizer`` and ``gdn_w_quantizer`` to a GatedDeltaNet module. + + Subclasses route the module's chunked gated-delta-rule call through + :meth:`_state_quantized_chunk_gated_delta_rule`. Both quantizers start disabled; enable them + with ``quant_cfg`` entries on ``*gdn_state_quantizer`` / ``*gdn_w_quantizer`` such as the + ``configs/ptq/units/gdn_state_fp8_dynamic`` and ``gdn_w_fp8_dynamic`` recipe units. The state + quantizer carries the fused QDQ configuration. State supports dynamic E4M3 or signed + narrow-range INT8; W supports dynamic E4M3. + Both sites require identity STE. The execution policy is saved in ModelOpt metadata. + """ + + @property + def gdn_state_qdq_block_v(self) -> int: + """Value-column scale grouping from the saved execution policy.""" + return self.linear_attention_config.state.block_v + + def _state_quantized_chunk_gated_delta_rule( + self, gated_delta_rule: GatedDeltaRuleFn, *args: Any, **kwargs: Any + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Call ``gated_delta_rule`` or, if a quantizer is on, the vendored quantizing copy.""" + self.validate_linear_attention() + quantize_state = self.gdn_state_quantizer.is_enabled and self.gdn_state_quantizer._if_quant + quantize_w = self.gdn_w_quantizer.is_enabled + if not self.linear_attention_is_enabled: + return gated_delta_rule(*args, **kwargs) + while isinstance(gated_delta_rule, partial): + args = (*gated_delta_rule.args, *args) + kwargs = {**gated_delta_rule.keywords, **kwargs} + gated_delta_rule = gated_delta_rule.func + if gated_delta_rule is not _fla_chunk_gated_delta_rule(): + raise NotImplementedError( + "GatedDeltaNet quantization supports only FLA's chunk_gated_delta_rule " + f"callable or a functools.partial of it; got {gated_delta_rule!r}." + ) + chunk_size = kwargs.pop("chunk_size", self.linear_attention_config.chunk_size) + if chunk_size != self.linear_attention_config.chunk_size: + raise ValueError("GDN fake quantization supports only chunk_size=64") + if self.linear_attention_config.backend == "matmul": + return matmul_gdn( + *args, + policy=self.linear_attention_config, + state_qdq=quantize_state, + state_format=self._linear_attn_state_format, + chunk_size=chunk_size, + prefill_lengths=self._linear_attention_prefill_lengths, + **kwargs, + ) + return _state_qdq_chunk_gated_delta_rule()( + *args, + chunk_size=chunk_size, + state_qdq=(2 if self._linear_attn_state_format == "int8" else 1) + if quantize_state + else 0, + state_qdq_block_v=self.gdn_state_qdq_block_v, + w_quantizer=self.gdn_w_quantizer if quantize_w else None, + **kwargs, + ) diff --git a/modelopt/torch/quantization/plugins/kda.py b/modelopt/torch/quantization/plugins/kda.py new file mode 100644 index 00000000000..fbc4b1cf0ef --- /dev/null +++ b/modelopt/torch/quantization/plugins/kda.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""QAT integration for the optional FLA KimiDeltaAttention layer.""" + +import importlib.metadata +import inspect +from functools import lru_cache +from types import FunctionType + +from ..linear_attention.kda import matmul_kda +from ..nn import QuantModuleRegistry +from .custom import CUSTOM_MODEL_PLUGINS +from .linear_attention import _LinearAttentionQuantMixin + +__all__ = [] + + +@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") + + def validate_linear_attention(self): + super().validate_linear_attention() + if self.linear_attention_is_enabled and self.linear_attention_config.backend != "matmul": + raise ValueError("KDA numerical emulation requires backend='matmul'") + + def forward(self, *args, **kwargs): + self.validate_linear_attention() + if not self.linear_attention_is_enabled: + return super().forward(*args, **kwargs) + original = super().forward.__func__ + 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__, + "chunk_kda": self._quantized_chunk, + "fused_recurrent_kda": self._unsupported_recurrent, + } + forward = FunctionType( + original.__code__, + namespace, + original.__name__, + original.__defaults__, + original.__closure__, + ) + forward.__kwdefaults__ = original.__kwdefaults__ + return forward(self, *args, **kwargs) + + def _quantized_chunk(self, *args, **kwargs): + return matmul_kda( + *args, + policy=self.linear_attention_config, + state_qdq=self.kda_state_quantizer.is_enabled and self.kda_state_quantizer._if_quant, + state_format=self._linear_attn_state_format, + 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" + ) + + +def _register_fla_kda(model): + for module in model.modules(): + cls = type(module) + if cls.__module__ != "fla.layers.kda" or cls.__name__ != "KimiDeltaAttention": + continue + if cls in QuantModuleRegistry: + continue + if importlib.metadata.version("flash-linear-attention") != "0.5.1": + raise RuntimeError("KDA QAT is qualified with flash-linear-attention==0.5.1") + if importlib.metadata.version("fla-core") != "0.5.1": + raise RuntimeError("KDA QAT requires fla-core==0.5.1") + QuantModuleRegistry.register({cls: "fla_KimiDeltaAttention"})(_QuantKimiDeltaAttention) + + +CUSTOM_MODEL_PLUGINS.add(_register_fla_kda) diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py new file mode 100644 index 00000000000..5559c9f6606 --- /dev/null +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared module policy and checkpoint support for linear-attention QAT.""" + +from ..config import QuantizerAttributeConfig +from ..linear_attention.config import LinearAttentionConfig +from ..linear_attention.utils import validate_gdn_quantizer +from ..nn import QuantModule, TensorQuantizer + +__all__ = [] + + +class _LinearAttentionQuantMixin(QuantModule): + linear_attention_quantizer_names = ("gdn_state_quantizer", "gdn_w_quantizer") + + def _setup(self): + for name in self.linear_attention_quantizer_names: + self._register_temp_attribute( + name, TensorQuantizer(QuantizerAttributeConfig(enable=False)) + ) + self._register_temp_attribute("linear_attention_config", LinearAttentionConfig()) + self._register_temp_attribute("_linear_attention_prefill_lengths", None) + + @property + def _linear_attn_state(self): + return getattr(self, self.linear_attention_quantizer_names[0]) + + @property + def _linear_attn_state_format(self): + return "int8" if self._linear_attn_state.num_bits == 8 else "fp8_e4m3" + + @property + def _linear_attn_w(self): + return getattr(self, self.linear_attention_quantizer_names[1]) + + @property + def linear_attention_is_enabled(self): + """Whether an operand, state, or arithmetic policy changes the computation.""" + return ( + self._linear_attn_state.is_enabled + or self._linear_attn_w.is_enabled + or self.linear_attention_config.decode is not None + ) + + def validate_linear_attention(self): + """Validate quantizer contracts shared by GDN and KDA.""" + if self._linear_attn_state.is_enabled: + validate_gdn_quantizer( + self._linear_attn_state, + name=self.linear_attention_quantizer_names[0], + num_bits=((4, 3), 8), + ) + if self._linear_attn_state.axis != (0, 1): + raise ValueError( + f"{self.linear_attention_quantizer_names[0]} supports only axis=(0, 1) " + "with state.block_v tiling" + ) + decode = self.linear_attention_config.decode + if decode is not None and decode.state_codec == "int8_hadamard32": + if self._linear_attn_state_format != "int8": + raise ValueError("int8_hadamard32 requires INT8 state quantization") + if self._linear_attn_w.is_enabled: + validate_gdn_quantizer( + self._linear_attn_w, name=self.linear_attention_quantizer_names[1] + ) + if self.linear_attention_config.decode is not None: + raise ValueError("Decode's exact prefix does not support WY operand QDQ") + + def modelopt_post_restore(self, prefix=""): + """Validate the restored numerical policy and quantizers.""" + super().modelopt_post_restore(prefix) + self.validate_linear_attention() diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index ec1958649dd..49195b44d29 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -21,7 +21,7 @@ import textwrap import types from contextlib import contextmanager -from functools import cache +from functools import cache, partial from typing import Any import megatron.core.parallel_state as mcore_parallel @@ -61,6 +61,7 @@ from ..utils import sync_moe_expert_amax from ..utils.layerwise_calib import LayerActivationCollector from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear +from .gdn import GatedDeltaNetStateQuantMixin try: from megatron.core.extensions.transformer_engine import ( @@ -79,6 +80,13 @@ except ImportError: HAS_TE = False +try: + from megatron.core.ssm.gated_delta_net import GatedDeltaNet + + HAS_GDN = True +except ImportError: + HAS_GDN = False + __all__ = [] @@ -1069,6 +1077,55 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): return make_sharded_tensors_for_checkpoint(state_dict, prefix, {}, sharded_offsets) +if HAS_GDN: + + @QuantModuleRegistry.register({GatedDeltaNet: "megatron_GatedDeltaNet"}) + class _QuantGatedDeltaNet(GatedDeltaNetStateQuantMixin): + """GatedDeltaNet with fake quantization of the recurrent state at kernel chunk boundaries. + + Routes ``self.gated_delta_rule`` through state/W QDQ from both Megatron's direct + forward and older split-forward layouts. The older dynamic-batching inference + paths (``ssm_prefill`` / ``ssm_decode``) are left untouched. + """ + + def _setup(self): + super()._setup() + try: + data_parallel_group = get_data_parallel_group(with_context_parallel=True) + except AssertionError: + data_parallel_group = get_data_parallel_group() + self.parallel_state = ParallelState( + data_parallel_group, + mcore_parallel.get_tensor_model_parallel_group(), + ) + + @contextmanager + def _quantized_gdn_kernel(self): + gated_delta_rule = self.gated_delta_rule + self.gated_delta_rule = partial( + self._state_quantized_chunk_gated_delta_rule, gated_delta_rule + ) + try: + yield + finally: + self.gated_delta_rule = gated_delta_rule + + def forward(self, *args, **kwargs): + if hasattr(GatedDeltaNet, "forward_pre_attn_and_core_attn"): + return super().forward(*args, **kwargs) + with self._quantized_gdn_kernel(): + return super().forward(*args, **kwargs) + + def forward_pre_attn_and_core_attn(self, *args, **kwargs): + with self._quantized_gdn_kernel(): + return super().forward_pre_attn_and_core_attn(*args, **kwargs) + + def validate_linear_attention(self): + super().validate_linear_attention() + if self.config.context_parallel_size > 1 and self.linear_attention_is_enabled: + raise NotImplementedError("GDN QAT does not support Megatron context parallelism.") + + def _is_supported_megatron_model(model: torch.nn.Module) -> bool: return isinstance(model, MegatronModule) diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 20e083491aa..67513218137 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -50,7 +50,8 @@ def _fp8_eager(x, amax=None): epsilon = 1.0 / (1 << 24) zero_amax_mask = amax <= epsilon safe_amax = torch.where(zero_amax_mask, torch.ones_like(amax), amax) - scale = 448.0 / safe_amax + # Use direct division, matching the CUDA extension rather than reciprocal multiplication. + scale = torch.div(448.0, safe_amax) scale_inv = 1 / scale x = (x.to(torch.float32) * scale).clamp(min=-448.0, max=448.0) x = x.to(torch.float8_e4m3fn) diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index db37b9222ca..ea115b6fe88 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -32,3 +32,6 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `experts_nvfp4.yaml` | NVFP4 W4A4 on `*.experts.*` weight/input quantizers | | `mixer_mlp_nvfp4.yaml` | NVFP4 W4A4 on dense `*.mixer.{up,down}_proj` weight/input quantizers | | `attention_qkv_fp8.yaml` | FP8 E4M3 on attention q/k/v bmm and softmax quantizers | +| `gdn_state_fp8_dynamic.yaml` | FP8 E4M3 dynamic fake quantization of the GatedDeltaNet recurrent state (per sequence, head and 64-column tile) at every kernel chunk boundary; needs `fla-core==0.5.1`, Triton, and SM89+ | +| `gdn_state_int8_dynamic.yaml` | INT8 quantizer unit; use the [complete recipe](../../../general/ptq/gdn_state_int8_dynamic.yaml) to enable INT8 state QDQ and Hadamard together | +| `gdn_w_fp8_dynamic.yaml` | FP8 E4M3 dynamic (per token and head) fake quantization of the WY tensor `w` that multiplies the GatedDeltaNet state; requires identity STE, `fla-core==0.5.1`, and Triton | diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index 2fee6279973..e409bff7aa6 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -16,6 +16,15 @@ # QuantizerCfgList snippet for standard module patterns that should remain unquantized. # modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + # Linear-attention numerical emulation is opt-in, independently of projection quantization. + - quantizer_name: '*gdn_state_quantizer' + enable: false + - quantizer_name: '*gdn_w_quantizer' + enable: false + - quantizer_name: '*kda_state_quantizer' + enable: false + - quantizer_name: '*kda_w_quantizer' + enable: false - quantizer_name: '*block_sparse_moe.gate*' enable: false - quantizer_name: '*linear_attn.conv1d*' diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml new file mode 100644 index 00000000000..4d4b13278ba --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -0,0 +1,30 @@ +# 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. + +# Dynamic per-tile FP8 E4M3 fake quantization of the GatedDeltaNet recurrent state. +# Recompute amax over each [K, BV] tile inside the kernel, independently per sequence and head +# of the [N, H, K, V] state (no calibration). ``axis: [0, 1]`` specifies sequence/head grouping; +# BV is ``gdn_state_qdq_block_v`` (default: the kernel's 64-column tile, +# two scales per head for V = 128; set 128 on the quantized module for one scale per head). This +# selects the FP8 state format. Applied at every kernel chunk boundary +# (64 tokens) during calibration, QAT and QAD; needs fla-core==0.5.1 and Triton. +# See ``modelopt.torch.quantization.plugins.gated_delta_net``. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*gdn_state_quantizer' + cfg: + num_bits: e4m3 + axis: [0, 1] + type: dynamic diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml new file mode 100644 index 00000000000..e86390cfef5 --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml @@ -0,0 +1,26 @@ +# 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. + +# Quantizer-only INT8 unit; general/ptq/gdn_state_int8_dynamic pairs it with Hadamard. +# The execution policy selects state scale groups and chunk/token/replay boundaries. +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*gdn_state_quantizer' + cfg: + num_bits: 8 + unsigned: false + narrow_range: true + axis: [0, 1] + type: dynamic + pass_through_bwd: true diff --git a/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml new file mode 100644 index 00000000000..2c73fdd6eaa --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml @@ -0,0 +1,28 @@ +# 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. + +# FP8 E4M3 fake quantization of ``w``, the WY-transformed keys that multiply the GatedDeltaNet +# recurrent state inside the chunked kernel, with a dynamic scale per token and head +# (``w`` is ``[B, T, H, K]``; ``axis: [0, 1, 2]`` reduces over K). Pairs with +# ``gdn_state_fp8_dynamic`` to emulate an FP8 x FP8 state matmul. Unlike the state quantizer this +# one runs on the tensor itself. The initial integration supports dynamic E4M3 with identity +# STE only. Needs fla-core==0.5.1 and Triton. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*gdn_w_quantizer' + cfg: + num_bits: e4m3 + axis: [0, 1, 2] + type: dynamic diff --git a/modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml b/modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml new file mode 100644 index 00000000000..e960282121b --- /dev/null +++ b/modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml @@ -0,0 +1,38 @@ +# 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. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + gdn_state_int8_dynamic: configs/ptq/units/gdn_state_int8_dynamic + +metadata: + recipe_type: ptq + description: >- + Dynamic INT8 GDN state QDQ with 32-value Hadamard rotation during decode. + Prefill state remains unquantized; supply per-sequence prefix lengths through + linear_attention_training_phase. No calibration is required. +quantize: + algorithm: + quant_cfg: + - $import: base_disable_all + - $import: gdn_state_int8_dynamic + linear_attention: + - module_name: '*' + cfg: + backend: matmul + decode: + mode: token + state_codec: int8_hadamard32 + prefill_state_qdq: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 2ad373ffcab..53da4aac5fc 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 26 general/ptq/ recipes (click to expand) +All 27 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -58,6 +58,7 @@ supported combinations. | `int4_blockwise_weight_only` | INT4 W4A16, block 128, weights only | none | max | | `nvfp4_mlp_weight_only` | NVFP4 W4A16 (block 32), MLP + MoE weights only | none | max | | `mxfp4_mlp_weight_only` | MXFP4 W4A16, MLP + MoE weights only | none | none (no calibration) | +| `gdn_state_int8_dynamic` | GDN decode state INT8 + Hadamard; weights unchanged | none | none (dynamic scales; requires a prefix/decode phase context) |
diff --git a/noxfile.py b/noxfile.py index 135465ed445..cd9e96ecffa 100644 --- a/noxfile.py +++ b/noxfile.py @@ -130,6 +130,15 @@ def gpu(session): "mamba_ssm", "causal-conv1d", ) + session.run( + "python", + "-m", + "pip", + "install", + "fla-core==0.5.1", + "tilelang==0.1.8", + "apache-tvm-ffi==0.1.9", + ) session.run("python", "-m", "pytest", "tests/gpu", *_cov_args()) @@ -141,6 +150,15 @@ def gpu_megatron(session): # Pre-installed nvidia-modelopt shadows the editable install session.run("pip", "uninstall", "-y", "nvidia-modelopt") session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]") + session.run( + "python", + "-m", + "pip", + "install", + "fla-core==0.5.1", + "tilelang==0.1.8", + "apache-tvm-ffi==0.1.9", + ) session.run("python", "-m", "pytest", "tests/gpu_megatron", *_cov_args()) diff --git a/pyproject.toml b/pyproject.toml index 0886e0380b3..47c8fbdf5ab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -248,6 +248,7 @@ extend-ignore = [ "tests/*" = ["B017", "D", "E402", "PT012"] "plugins/modelopt/skills/*/tests/test_*.py" = ["D", "E402"] # Skill test scripts: docstring (D) + sys.path import-order (E402) exemptions "*/_[a-zA-Z]*" = ["D"] # Private packages (_abc/*.py) or modules (_xyz.py) +"modelopt/torch/kernels/quantization/linear_attention/fla_*.py" = ["D", "E501"] # Vendored flash-linear-attention kernels "*.ipynb" = ["D", "E501"] # Ignore missing docstrings or line length for Jupyter notebooks "modelopt/torch/kernels/*" = ["N803", "N806", "E731"] # triton style "modelopt/torch/puzzletron/*" = [ @@ -322,6 +323,15 @@ disable_error_code = ["attr-defined"] module = ["examples.diffusers.fastgen.preprocess.*"] ignore_errors = true +# Vendored from fla-org/flash-linear-attention (MIT); kept faithful to upstream rather than +# annotated to modelopt's strict mypy (Triton loop indices, ``*tensor.shape`` unpacking). +[[tool.mypy.overrides]] +module = [ + "modelopt.torch.kernels.quantization.linear_attention.fla_chunk_delta_h", + "modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule", +] +ignore_errors = true + [tool.bandit] exclude_dirs = [".github/", "examples/", "noxfile.py", "tests/"] # Do not change `skips`. It should be consistent with NVIDIA's Wheel-CI-CD bandit.yml config. diff --git a/tests/_test_utils/torch/quantization/linear_attention_reference.py b/tests/_test_utils/torch/quantization/linear_attention_reference.py new file mode 100644 index 00000000000..57d7448bcb1 --- /dev/null +++ b/tests/_test_utils/torch/quantization/linear_attention_reference.py @@ -0,0 +1,346 @@ +# 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. + +"""Small differentiable oracles, independent of FLA and Triton. + +These references prioritize explicit arithmetic and gradients over speed. Gates are +natural logarithms; inputs are already normalized/activated. Packed sequence boundaries +are read on the CPU. Accumulation follows the input dtype (use float64 for algebra tests). +""" + +from collections.abc import Callable +from itertools import pairwise + +import torch + +__all__ = [ + "chunk_gdn_reference", + "chunk_kda_reference", + "recurrent_delta_rule_reference", + "state_fp8_qdq_reference", + "state_qdq_reference", +] + + +def state_fp8_qdq_reference(state: torch.Tensor, block_v: int = 64) -> torch.Tensor: + """Dynamic E4M3 QDQ with identity STE and one scale per ``[Dk, block_v]`` tile. + + ``state`` is in key-first layout ``[..., Dk, Dv]``. Scales and rounding do not + contribute derivatives. FP8 uses TensorQuantizer small-range handling; partial tiles are valid. + """ + return state_qdq_reference(state, block_v) + + +def state_qdq_reference(state: torch.Tensor, block_v: int = 64, state_format: str = "fp8_e4m3"): + """Dynamic state-tile QDQ with detached scales and identity STE.""" + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") + if block_v not in (16, 32, 64, 128): + raise ValueError("block_v must be 16, 32, 64, or 128") + with torch.no_grad(): + rounded = [] + for tile in state.float().split(block_v, dim=-1): + amax = tile.abs().amax(dim=(-2, -1), keepdim=True) + if state_format == "int8": + scale = torch.where(amax > 0, amax / 127.0, torch.ones_like(amax)) + codes = (tile / scale).clamp(-127, 127).round() + else: + safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) + quant_scale = torch.div(448.0, safe_amax) + scale = quant_scale.reciprocal() + codes = (tile * quant_scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() + rounded.append(codes * scale) + quantized = torch.cat(rounded, dim=-1).to(state.dtype) + return state + (quantized - state).detach() + + +def _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first): + if q.ndim != 4 or k.shape != q.shape or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k must have shape [B,T,Hk,Dk] and v shape [B,T,Hv,Dv]") + batch, length, heads, keys = q.shape + value_heads, values = v.shape[2:] + if length == 0 or value_heads % heads: + raise ValueError("nonempty sequences and Hv divisible by Hk are required") + if beta.shape != (batch, length, value_heads) or g.shape not in ( + beta.shape, + (*beta.shape, keys), + ): + raise ValueError("beta must be [B,T,Hv]; g must be [B,T,Hv] or [B,T,Hv,Dk]") + if cu_seqlens is None: + sequences = [(b, 0, length) for b in range(batch)] + else: + boundaries = cu_seqlens.tolist() + if ( + batch != 1 + or len(boundaries) < 2 + or boundaries[0] != 0 + or boundaries[-1] != length + or any(a >= b for a, b in pairwise(boundaries)) + ): + raise ValueError("cu_seqlens must partition a packed batch of size one") + sequences = [(0, a, b) for a, b in pairwise(boundaries)] + if initial_state is None: + state = q.new_zeros(len(sequences), value_heads, keys, values) + else: + state = initial_state.transpose(-1, -2) if state_v_first else initial_state + if state.shape != (len(sequences), value_heads, keys, values): + raise ValueError("initial_state shape does not match sequence/head dimensions") + return ( + q.repeat_interleave(value_heads // heads, dim=2), + k.repeat_interleave(value_heads // heads, dim=2), + state, + sequences, + ) + + +def recurrent_delta_rule_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + state_v_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """GDN (scalar decay) or KDA (per-key decay) exact token recurrence. + + Inputs have shapes ``q,k:[B,T,Hk,Dk]``, ``v:[B,T,Hv,Dv]``, + ``beta:[B,T,Hv]``, and ``g:[B,T,Hv]`` (GDN) or ``[B,T,Hv,Dk]`` (KDA). + State is ``[N,Hv,Dk,Dv]``, or ``[N,Hv,Dv,Dk]`` with ``state_v_first``. + Returns output and final state, preserving gradients through the initial state. + """ + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + sequence_output = [] + for t in range(start, end): + decay = g[b, t].exp() + if g.ndim == 3: + decay = decay.unsqueeze(-1) + decayed = state * decay.unsqueeze(-1) + residual = v[b, t] - (k[b, t].unsqueeze(-1) * decayed).sum(-2) + update = beta[b, t].unsqueeze(-1) * residual + state = decayed + k[b, t].unsqueeze(-1) * update.unsqueeze(-2) + sequence_output.append((q[b, t].unsqueeze(-1) * state).sum(-2) * scale) + outputs.append(torch.stack(sequence_output)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final + + +def chunk_gdn_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + chunk_size: int = 64, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + state_v_first: bool = False, + state_qdq: bool = False, + state_qdq_block_v: int = 64, + state_format: str = "fp8_e4m3", + w_quantizer: Callable[[torch.Tensor], torch.Tensor] | None = None, + matmul: Callable[[str, torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + arithmetic: Callable[[str, torch.Tensor], torch.Tensor] | None = None, + inverse_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Exact GDN chunk algebra with optional state/W fake quantization. + + The solve is a unit-lower triangular solve. ``w_quantizer`` sees the complete + materialized ``[B,T,Hv,Dk]`` WY operand once, with its own autograd semantics. + State QDQ occurs on the initial state and each chunk's final state, after readout. + Optional callbacks expose the eight named matmuls and elementwise boundaries. + Matmul callbacks receive conventional ``lhs @ rhs`` shapes, without padding. + """ + if g.ndim != 3 or chunk_size <= 0: + raise ValueError("chunk_gdn_reference requires scalar GDN gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + mm = matmul if matmul is not None else lambda name, lhs, rhs: lhs @ rhs + arithmetic = arithmetic if arithmetic is not None else lambda name, x: x + chunks, all_w = [], [] + for n, (b, start, end) in enumerate(sequences): + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + gc = arithmetic("gate_prefix", g[b, lo:hi].transpose(0, 1).cumsum(-1)) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + # Mask before exp: upper-triangle positive differences can overflow for long decay. + causal = torch.ones(hi - lo, hi - lo, device=q.device, dtype=torch.bool).tril() + decay = arithmetic( + "gate_exp", (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() + ) + if matmul is None: + lower = (bc * (kc @ kc.transpose(-1, -2)) * decay).tril(-1) + else: + lower = (mm("key_interaction", bc * kc, kc.transpose(-1, -2)) * decay).tril(-1) + matrix = lower + torch.eye(hi - lo, device=q.device, dtype=q.dtype) + gate = arithmetic("gate_exp", gc.exp()).unsqueeze(-1) + if matmul is None and inverse_fn is None: + rhs = torch.cat((bc * vc, bc * kc * gate), dim=-1) + solved = torch.linalg.solve_triangular(matrix, rhs, upper=False, unitriangular=True) + u, w = solved.split((v.shape[-1], k.shape[-1]), dim=-1) + else: + inverse = ( + inverse_fn(lower) + if inverse_fn is not None + else torch.linalg.solve_triangular( + matrix, + torch.eye(hi - lo, device=q.device, dtype=q.dtype).expand_as(matrix), + upper=False, + unitriangular=True, + ) + ) + u = mm("wy_value", inverse, bc * vc) + w = mm("wy_key", inverse, bc * kc * gate) + chunks.append((n, qc, kc, gc, decay, u)) + all_w.append(w.transpose(0, 1)) + w = torch.cat(all_w).reshape(q.shape) + if w_quantizer is not None: + w = w_quantizer(w) + w = w.reshape(-1, *w.shape[2:]) + outputs, finals, offset, previous_n = [], [], 0, -1 + for n, qc, kc, gc, decay, u in chunks: + if n != previous_n: + state = states[n] + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + previous_n = n + length = qc.shape[1] + wc = w[offset : offset + length].transpose(0, 1) + offset += length + updated_values = arithmetic("value_residual", u - mm("state_read", wc, state)) + local_scores = (mm("output_score", qc * scale, kc.transpose(-1, -2)) * decay).tril() + output = mm( + "output_state", qc * (scale * arithmetic("gate_exp", gc.exp())).unsqueeze(-1), state + ) + output = arithmetic("output_add", output + mm("output_value", local_scores, updated_values)) + outputs.append(output.transpose(0, 1)) + weighted_keys = kc * arithmetic("gate_exp", (gc[..., -1:] - gc).exp()).unsqueeze(-1) + state = arithmetic( + "state_decay", state * arithmetic("gate_exp", gc[..., -1].exp())[:, None, None] + ) + state = arithmetic( + "state_add", state + mm("state_update", weighted_keys.transpose(-1, -2), updated_values) + ) + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + if len(finals) <= n: + finals.append(state) + else: + finals[n] = state + output = torch.cat(outputs).reshape(*v.shape) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final + + +def chunk_kda_reference( + q, + k, + v, + g, + beta, + *, + chunk_size=64, + scale=None, + initial_state=None, + cu_seqlens=None, + state_v_first=False, + state_qdq=False, + state_qdq_block_v=64, + state_format="fp8_e4m3", + w_quantizer=None, + matmul=None, + inverse_fn=None, +): + """KDA chunk oracle with causal per-channel decays and optional operand QDQ. + + The eight-site callback receives conventional three-dimensional matmul + operands. Interaction sites are invoked one query row at a time with decay + already applied to the right-hand keys. This avoids inverse-prefix factors. + """ + if g.ndim != 4 or chunk_size <= 0: + raise ValueError("chunk_kda_reference requires per-key gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + mm = matmul if matmul is not None else lambda name, lhs, rhs: lhs @ rhs + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + pieces = [] + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + prefix = g[b, lo:hi].transpose(0, 1).cumsum(-2) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + gate = prefix.exp() + lower_rows, score_rows = [], [] + for row in range(hi - lo): + decayed_keys = ( + kc[..., : row + 1, :] + * (prefix[..., row : row + 1, :] - prefix[..., : row + 1, :]).exp() + ) + right = decayed_keys.transpose(-1, -2) + left = kc[..., row : row + 1, :] * bc[..., row : row + 1, :] + interaction = mm("key_interaction", left, right) + score = mm("output_score", qc[..., row : row + 1, :] * scale, right) + # Zero padding follows the computation and never contributes a scale statistic. + padding = hi - lo - row - 1 + lower_rows.append(torch.nn.functional.pad(interaction, (0, padding))) + score_rows.append(torch.nn.functional.pad(score, (0, padding))) + lower = torch.cat(lower_rows, dim=-2).tril(-1) + scores = torch.cat(score_rows, dim=-2) + identity = torch.eye(hi - lo, device=q.device, dtype=q.dtype).expand_as(lower) + inverse = ( + inverse_fn(lower) + if inverse_fn is not None + else torch.linalg.solve_triangular( + identity + lower, identity, upper=False, unitriangular=True + ) + ) + u = mm("wy_value", inverse, bc * vc) + w = mm("wy_key", inverse, bc * kc * gate) + if w_quantizer is not None: + w = w_quantizer(w) + updated = u - mm("state_read", w, state) + pieces.append( + ( + mm("output_state", qc * scale * gate, state) + + mm("output_value", scores, updated) + ).transpose(0, 1) + ) + weighted_keys = kc * (prefix[..., -1:, :] - prefix).exp() + state = state * gate[..., -1, :, None] + mm( + "state_update", weighted_keys.transpose(-1, -2), updated + ) + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + outputs.append(torch.cat(pieces)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final diff --git a/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..da6fdf6e772 --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py @@ -0,0 +1,412 @@ +# 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.common.fp8_quant import fp8_dynamic_qdq +from modelopt.torch.kernels.quantization.linear_attention.decode import fused_recurrence +from modelopt.torch.quantization.config import QuantizerAttributeConfig +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_format", [None, "fp8_e4m3", "int8"]) +@pytest.mark.parametrize("replay", [False, True]) +@pytest.mark.parametrize("read_stored", [False, True]) +def test_fused_recurrence_outputs_state_and_gradients(kda, state_format, replay, read_stored): + state_qdq = state_format is not None + state_format = state_format or "fp8_e4m3" + 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, + state_format=state_format, + 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, + state_format=state_format, + 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]) +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_fused_carry_handoff_and_auxiliary_gradients(mode, interval, state_format): + 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, + state_format=state_format, + 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, + state_format=state_format, + 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, NATIVE_FP8: tl.constexpr): # noqa: N803 + offsets = tl.arange(0, BLOCK) + values = tl.load(X + offsets, offsets < N, 0) + decoded, scale = fp8_dynamic_qdq(values, NATIVE_FP8) + tl.store(Y + offsets, decoded, offsets < N) + tl.store(S, scale) + + +@pytest.mark.parametrize("native_fp8", [False, True]) +@pytest.mark.parametrize( + "scale", + [0.0, 1e-8 / 448, 2.0**-24 / 448, 2.0**-23 / 448, 2.0**-12, 1.0, 2.0**12, 0.00314159, 731.173], +) +def test_fp8_codec_matches_tensor_quantizer_boundaries(scale, native_fp8): + if native_fp8 and torch.cuda.get_device_capability() < (8, 9): + pytest.skip("Native FP8 conversion requires SM89 or later") + # 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)), + native_fp8, + ) + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), type="dynamic")) + torch.testing.assert_close(actual, quantizer(values), rtol=0, atol=0) + amax = values.abs().amax() + safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) + torch.testing.assert_close( + actual_scale, torch.div(448.0, safe_amax).reciprocal(), rtol=0, atol=0 + ) + + +def test_fused_fp8_state_matches_tensor_quantizer_with_partial_tiles(): + torch.manual_seed(762) + amplitudes = torch.tensor([0.0, 1e-8, 2**-24, 2**-23, 1.0, 731.173], device="cuda") + state = (torch.rand(6, 3, 19, device="cuda") * 2 - 1) * amplitudes[:, None, None] + state[:, 0, 0] = amplitudes + state.requires_grad_() + q = torch.zeros(1, 6, 3, device="cuda") + v = torch.zeros(1, 6, 19, device="cuda") + gate = torch.zeros(1, 6, device="cuda") + _, actual, *_ = fused_recurrence(q, q, v, gate, gate, state, state_qdq=True, block_v=16) + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=(4, 3), type="dynamic", axis=0, pass_through_bwd=True) + ) + expected = torch.cat( + [quantizer(tile.flatten(-2)).reshape_as(tile) for tile in state.split(16, -1)], -1 + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + probe = torch.randn_like(state) + gradients = [ + torch.autograd.grad((value * probe).sum(), state)[0] for value in (actual, expected) + ] + torch.testing.assert_close(*gradients, rtol=0, atol=0) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("mode", ["token", "replay"]) +@pytest.mark.parametrize("horizon", ["near_unit", "mixed"]) +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_wide_long_trajectory(kda, mode, horizon, state_format): + 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, + state_format=state_format, + block_v=64, + initial_state=initial, + ) + actual, ac = recurrent_decode( + *args, + config=cfg.model_copy(update={"implementation": "triton"}), + state_qdq=True, + state_format=state_format, + 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]) +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_full_key_width_and_value_block_tail(block_v, kda, state_format): + 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, + state_format=state_format, + 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]) +@pytest.mark.parametrize( + ("mode", "block_v", "readout", "factor_qdq"), + [ + ("token", 32, "stored", False), + ("replay", 64, "working", False), + ("replay", 128, "stored", True), + ], +) +@pytest.mark.timeout(180) +def test_hadamard_fused_split_carry_and_gradients(kda, mode, block_v, readout, factor_qdq): + torch.manual_seed(917) + q, k = [ + F.normalize(torch.randn(37, 2, 17, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(37, 2, 96, device="cuda", requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1], device="cuda") * 0.03).requires_grad_() + beta = (torch.rand(37, 2, device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(2, 17, 96, device="cuda") * 0.1).requires_grad_() + args = (q, k, v, g, beta) + cfg = LinearAttentionDecodeConfig( + mode=mode, + readout=readout, + state_codec="int8_hadamard32", + decay_log_step=1 / 256, + replay={"window": 8, "factor_qdq": factor_qdq} if mode == "replay" else None, + ) + shared = {"state_qdq": True, "state_format": "int8", "block_v": block_v} + expected, ref_carry = recurrent_decode_reference( + *args, config=cfg, initial_state=initial, **shared + ) + carry, start, pieces = None, 0, [] + for end in (0, 2, 8, 8, 19, 37): + output, carry = recurrent_decode( + *(x[start:end] for x in args), + config=cfg.model_copy(update={"implementation": "triton"}), + initial_state=initial if carry is None else None, + carry=carry, + **shared, + ) + pieces.append(output) + start = end + actual = torch.cat(pieces) + assert carry.value_basis == "hadamard32" and carry.cursor == ref_carry.cursor + assert carry.anchor.scales.shape == (2, 17, 3) and carry.anchor.scales.dtype == torch.float16 + torch.testing.assert_close(carry.anchor.scales, ref_carry.anchor.scales, rtol=0, atol=0) + actual_tensors = ( + actual, + carry.reconstruct(), + carry.anchor.values, + *(e.update.values for e in carry.entries), + ) + expected_tensors = ( + expected, + ref_carry.reconstruct(), + ref_carry.anchor.values, + *(e.update.values for e in ref_carry.entries), + ) + probes = [torch.randn_like(x) for x in expected_tensors] + gradients = [ + torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probes)), (*args, initial)) + for result in (actual_tensors, expected_tensors) + ] + for a, e in zip((*actual_tensors, *gradients[0]), (*expected_tensors, *gradients[1])): + torch.testing.assert_close(a, e, rtol=1e-3, atol=1e-5) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.timeout(180) +def test_hadamard_long_trajectory(kda): + torch.manual_seed(991) + q, k = [ + F.normalize(torch.randn(257, 1, 128, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(257, 1, 64, device="cuda", requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1], device="cuda") * 1e-4).requires_grad_() + beta = (torch.rand(257, 1, device="cuda") * 0.1).requires_grad_() + initial = (torch.randn(1, 128, 64, device="cuda") * 0.1).requires_grad_() + args = (q, k, v, g, beta) + results = [] + for implementation in ("torch", "triton"): + cfg = LinearAttentionDecodeConfig( + mode="replay", + implementation=implementation, + state_codec="int8_hadamard32", + replay={"window": 8, "factor_qdq": False}, + ) + out, carry = recurrent_decode( + *args, config=cfg, initial_state=initial, state_qdq=True, state_format="int8" + ) + final = carry.reconstruct() + grads = torch.autograd.grad(out.square().sum() + final.square().sum(), (*args, initial)) + results.append((out, final, *grads)) + for a, e in zip(*results): + assert torch.isfinite(a).all() and torch.isfinite(e).all() + assert (a - e).norm() / e.norm().clamp_min(1e-8) < 1e-4 diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py new file mode 100644 index 00000000000..84b344ba224 --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py @@ -0,0 +1,403 @@ +# 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. + +"""Numerical forward/backward coverage for the GDN training emulation contract.""" + +from functools import partial + +import pytest +import torch +import torch.nn.functional as F +from _test_utils.torch.quantization.linear_attention_reference import ( + chunk_gdn_reference, + state_fp8_qdq_reference, +) +from torch.utils.checkpoint import checkpoint + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.plugins.gdn import GatedDeltaNetStateQuantMixin + +fla = pytest.importorskip("fla.ops.gated_delta_rule") +from fla.utils import IS_NVIDIA_HOPPER, TRITON_ABOVE_3_4_0 + +from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( + chunk_gated_delta_rule, +) + +DTYPES = (torch.float32, torch.bfloat16) +QDQ_MODES = ((False, False), (False, True), (True, False), (True, True), (2, False), (2, True)) +STATE_LAYOUTS = ((False, False), (True, True)) +W_AXES = (None, (0, 1), (0, 1, 2)) + + +def require_state_qdq(): + if torch.cuda.get_device_capability() < (8, 9): + pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") + + +def make_inputs(dtype=None, packed=False, state_v_first=False, grouped=True): + hopper_backward = IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0 + if dtype is None: + dtype = torch.bfloat16 if hopper_backward else torch.float32 + if dtype == torch.float32 and hopper_backward: + pytest.skip("Hopper/TileLang training supports BF16; FP32 is rejected before launch") + torch.manual_seed(123) + # Two full chunks plus a tail; grouped-value heads and a partial state scale tile. + batch, length, heads, value_heads, keys, values = 1, 145, 1, 2, 32, 48 + heads = heads if grouped else value_heads + q, k = [ + F.normalize(torch.randn(batch, length, heads, keys, device="cuda"), dim=-1) + for _ in range(2) + ] + v = torch.randn(batch, length, value_heads, values, device="cuda") + g = -torch.rand(batch, length, value_heads, device="cuda") * 0.1 + beta = torch.rand_like(g) + state = torch.randn(2 if packed else batch, value_heads, keys, values, device="cuda") * 0.1 + if state_v_first: + state = state.transpose(-1, -2).contiguous() + args = [x.to(dtype).requires_grad_() for x in (q, k, v, g, beta)] + return args, state.requires_grad_() + + +def w_quantizer(axis=(0, 1, 2)): + return TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), axis=axis, type="dynamic")) + + +def compare(actual, expected, tolerance): + for a, e in zip(actual, expected): + assert torch.isfinite(a).all() + error = (a.float() - e.float()).norm() + bound = tolerance * e.float().norm().clamp_min(1e-6) + assert error <= bound, ( + f"relative L2 error {(error / e.float().norm()).item():.5g} > {tolerance}" + ) + + +def compare_forward(actual, expected, dtype, state_qdq): + tolerance = 0.03 if dtype == torch.bfloat16 else 0.01 + compare(actual[:1], expected[:1], tolerance) + # FP8 bin crossings amplify small GEMM differences in the stored final state. + state_tolerance = max(tolerance, 0.02) if state_qdq else tolerance + compare(actual[1:], expected[1:], state_tolerance) + + +def values_and_grads(fn, args, state, **kwargs): + result = fn(*args, initial_state=state, **kwargs) + torch.manual_seed(15) + probes = [torch.randn(x.shape, device=x.device, dtype=torch.float32) for x in result] + grads = torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probes)), (*args, state)) + return result, grads + + +@pytest.fixture(scope="module", autouse=True) +def _compile_gdn_kernels(): + """Keep compilation/autotuning in setup, including when selecting a single test.""" + supports_state_qdq = torch.cuda.get_device_capability() >= (8, 9) + for dtype in DTYPES: + if dtype == torch.float32 and IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0: + continue + args, state = make_inputs(dtype, grouped=False) + for kernel in (fla.chunk_gated_delta_rule, chunk_gated_delta_rule): + values_and_grads(kernel, args, state, output_final_state=True) + for packed, state_v_first in STATE_LAYOUTS: + args, state = make_inputs(dtype, packed=packed, state_v_first=state_v_first) + for state_qdq, quantize_w in QDQ_MODES: + if state_qdq == 1 and not supports_state_qdq: + continue + values_and_grads( + chunk_gated_delta_rule, + args, + state, + output_final_state=True, + state_qdq=state_qdq, + state_qdq_block_v=32, + w_quantizer=w_quantizer() if quantize_w else None, + state_v_first=state_v_first, + cu_seqlens=torch.tensor([0, 67, 145], device="cuda", dtype=torch.int32) + if packed + else None, + ) + + args, state = make_inputs() + quantizer = w_quantizer() + for state_qdq in (0, 2): + values_and_grads( + chunk_gated_delta_rule, + args, + state, + output_final_state=True, + w_quantizer=quantizer, + state_qdq=state_qdq, + ) + for axis in W_AXES: + output, _ = chunk_gated_delta_rule(*args, w_quantizer=w_quantizer(axis)) + torch.autograd.grad(output.square().sum(), args) + a_log = torch.zeros(2, device="cuda", requires_grad=True) + bias = torch.zeros_like(a_log, requires_grad=True) + output, final = chunk_gated_delta_rule( + *args, + initial_state=state, + output_final_state=True, + w_quantizer=quantizer, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + A_log=a_log, + dt_bias=bias, + use_beta_sigmoid_in_kernel=True, + allow_neg_eigval=True, + ) + torch.autograd.grad(output.square().sum() + final.square().sum(), (*args, state, a_log, bias)) + for state_qdq in (1, 2): + if state_qdq == 1 and not supports_state_qdq: + continue + for block_v in (16, 64, 128): + values_and_grads( + chunk_gated_delta_rule, + args, + state, + output_final_state=True, + state_qdq=state_qdq, + state_qdq_block_v=block_v, + ) + if supports_state_qdq: + args, state = make_inputs(packed=True, state_v_first=True) + for block_v in (16, 32, 64, 128): + _, final = chunk_gated_delta_rule( + *args, + initial_state=state, + output_final_state=True, + state_v_first=True, + cu_seqlens=torch.tensor([0, 67, 145], device="cuda", dtype=torch.int32), + state_qdq=True, + state_qdq_block_v=block_v, + ) + torch.autograd.grad(final.sum(), state) + torch.cuda.synchronize() + + +@pytest.mark.parametrize("dtype", DTYPES) +def test_disabled_matches_upstream_forward_and_backward(dtype): + # Upstream Hopper/TileLang backward cannot handle grouped value heads. + args, state = make_inputs(dtype, grouped=False) + expected = values_and_grads(fla.chunk_gated_delta_rule, args, state, output_final_state=True) + actual = values_and_grads(chunk_gated_delta_rule, args, state, output_final_state=True) + for a, e in zip(actual, expected): + compare(a, e, 0.005 if dtype == torch.bfloat16 else 0.001) + + +@pytest.mark.parametrize(("state_qdq", "quantize_w"), QDQ_MODES) +@pytest.mark.parametrize(("packed", "state_v_first"), STATE_LAYOUTS) +@pytest.mark.parametrize("dtype", DTYPES) +def test_numerical_surrogate_reference(state_qdq, quantize_w, packed, state_v_first, dtype): + if state_qdq == 1: + require_state_qdq() + args, state = make_inputs(dtype, packed=packed, state_v_first=state_v_first) + quantizer = w_quantizer() if quantize_w else None + kwargs = { + "state_qdq": state_qdq, + "state_qdq_block_v": 32, + "w_quantizer": quantizer, + "state_v_first": state_v_first, + "cu_seqlens": torch.tensor([0, 67, 145], device="cuda", dtype=torch.int32) + if packed + else None, + } + reference_args = [x.detach().float().requires_grad_() for x in args] + reference_state = state.detach().float().requires_grad_() + expected = values_and_grads( + chunk_gdn_reference, + reference_args, + reference_state, + state_format="int8" if state_qdq == 2 else "fp8_e4m3", + **kwargs, + ) + actual = values_and_grads( + chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs + ) + compare_forward(actual[0], expected[0], dtype, state_qdq) + compare(actual[1], expected[1], 0.05 if dtype == torch.bfloat16 else 0.02) + + +@pytest.mark.parametrize("state_qdq", [0, 2]) +def test_w_qdq_saved_once_and_activation_checkpoint_parity(state_qdq): + args, state = make_inputs() + quantizer = w_quantizer() + calls = [] + handle = quantizer.register_forward_hook(lambda *args: calls.append(1)) + + def fn(*x): + return chunk_gated_delta_rule( + *x[:5], + initial_state=x[5], + output_final_state=True, + w_quantizer=quantizer, + state_qdq=state_qdq, + ) + + output = fn(*args, state) + grads = torch.autograd.grad(sum(x.square().sum() for x in output), (*args, state)) + assert len(calls) == 1, "Backward must not rerun the quantizer or its observers" + recomputed = checkpoint(fn, *args, state, use_reentrant=False) + checkpoint_grads = torch.autograd.grad( + sum(x.square().sum() for x in recomputed), (*args, state) + ) + compare(recomputed, output, 1e-6) + compare(checkpoint_grads, grads, 1e-6) + handle.remove() + + +@pytest.mark.parametrize("axis", W_AXES) +def test_zero_initial_state_and_output_only_training_loss(axis): + args, _ = make_inputs() + quantizer = w_quantizer(axis) + expected, _ = chunk_gdn_reference(*(x.float() for x in args), w_quantizer=quantizer) + actual, final = chunk_gated_delta_rule(*args, w_quantizer=quantizer) + assert final is None + compare((actual,), (expected,), 0.03 if args[0].dtype == torch.bfloat16 else 0.01) + grads = [torch.autograd.grad(x.square().sum(), args) for x in (actual, expected)] + compare(*grads, 0.05 if args[0].dtype == torch.bfloat16 else 0.02) + + +@pytest.mark.parametrize("block_v", [16, 64, 128]) +@pytest.mark.parametrize("state_qdq", [1, 2]) +def test_state_scale_tile_forward_and_backward(block_v, state_qdq): + if state_qdq == 1: + require_state_qdq() + args, state = make_inputs() + kwargs = {"state_qdq": state_qdq, "state_qdq_block_v": block_v} + reference_args = [x.detach().float().requires_grad_() for x in args] + expected = values_and_grads( + chunk_gdn_reference, + reference_args, + state, + state_format="int8" if state_qdq == 2 else "fp8_e4m3", + **kwargs, + ) + actual = values_and_grads( + chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs + ) + compare_forward(actual[0], expected[0], args[0].dtype, state_qdq=True) + compare(actual[1], expected[1], 0.05 if args[0].dtype == torch.bfloat16 else 0.02) + + +@pytest.mark.parametrize("block_v", [16, 32, 64, 128]) +def test_state_qdq_without_recurrent_updates(block_v): + require_state_qdq() + args, state = make_inputs(packed=True, state_v_first=True) + q, k, v, g, beta = args + _, final = chunk_gated_delta_rule( + q, + torch.zeros_like(k), + torch.zeros_like(v), + torch.zeros_like(g), + beta, + initial_state=state, + output_final_state=True, + state_v_first=True, + cu_seqlens=torch.tensor([0, 67, 145], device="cuda", dtype=torch.int32), + state_qdq=True, + state_qdq_block_v=block_v, + ) + expected = state_fp8_qdq_reference(state.transpose(-1, -2), block_v).transpose(-1, -2) + torch.testing.assert_close(final, expected, rtol=1e-6, atol=1e-7) + (grad,) = torch.autograd.grad(final.sum(), state) + torch.testing.assert_close(grad, torch.ones_like(state), rtol=0, atol=0) + + +def test_fused_gate_normalization_and_beta_gradients(): + args, state = make_inputs() + q, k, v, raw_g, raw_beta = args + a_log = torch.randn(2, device="cuda", requires_grad=True) + bias = torch.randn(2, device="cuda", requires_grad=True) + quantizer = w_quantizer() + expected = chunk_gdn_reference( + F.normalize(q.float(), dim=-1), + F.normalize(k.float(), dim=-1), + v.float(), + -a_log.exp() * F.softplus(raw_g.float() + bias), + 2 * raw_beta.float().sigmoid(), + initial_state=state, + w_quantizer=quantizer, + ) + actual = chunk_gated_delta_rule( + *args, + initial_state=state, + output_final_state=True, + w_quantizer=quantizer, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + A_log=a_log, + dt_bias=bias, + use_beta_sigmoid_in_kernel=True, + allow_neg_eigval=True, + ) + compare(actual, expected, 0.03 if q.dtype == torch.bfloat16 else 0.01) + grads = [ + torch.autograd.grad(sum(x.square().sum() for x in result), (*args, state, a_log, bias)) + for result in (actual, expected) + ] + compare(*grads, 0.05 if q.dtype == torch.bfloat16 else 0.02) + + +def test_plugin_dispatch_preserves_fla_partial(): + args, state = make_inputs() + module = GatedDeltaNetStateQuantMixin.convert(torch.nn.Module()) + module.gdn_w_quantizer = w_quantizer() + expected = values_and_grads( + chunk_gated_delta_rule, + args, + state, + output_final_state=True, + w_quantizer=module.gdn_w_quantizer, + ) + + def wrapped(q, *remaining, **kwargs): + kernel = partial(fla.chunk_gated_delta_rule, q, output_final_state=True) + return module._state_quantized_chunk_gated_delta_rule(kernel, *remaining, **kwargs) + + actual = values_and_grads(wrapped, args, state) + for a, e in zip(actual, expected): + compare(a, e, 1e-6) + + +@pytest.mark.parametrize("option", ["use_gate_in_kernal", "unsupported_option"]) +def test_unknown_keyword_rejected_before_launch(option): + args, _ = make_inputs() + with pytest.raises(TypeError, match=f"Unexpected keyword arguments: {option}"): + chunk_gated_delta_rule(*args, **{option: True}) + + +@pytest.mark.parametrize("chunk_size", [16, 32, 128]) +def test_unsupported_chunk_rejected_before_launch(chunk_size): + args, state = make_inputs() + with pytest.raises(ValueError, match="only chunk_size=64"): + chunk_gated_delta_rule(*args, chunk_size=chunk_size) + + +def test_unsupported_w_gradient_rejected_before_launch(): + args, _ = make_inputs() + quantizer = w_quantizer() + quantizer.set_from_attribute_config({"pass_through_bwd": False}) + with pytest.raises(ValueError, match="pass_through_bwd=True"): + chunk_gated_delta_rule(*args, w_quantizer=quantizer) + + +def test_hopper_fp32_rejected_before_launch(): + if not (IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0): + pytest.skip("Hopper/TileLang capability guard") + args, _ = make_inputs() + with pytest.raises(ValueError, match="requires BF16 q/k/v"): + chunk_gated_delta_rule(*(x.float() for x in args)) diff --git a/tests/gpu/torch/quantization/plugins/test_kda.py b/tests/gpu/torch/quantization/plugins/test_kda.py new file mode 100644 index 00000000000..aecf7fe8ef8 --- /dev/null +++ b/tests/gpu/torch/quantization/plugins/test_kda.py @@ -0,0 +1,189 @@ +# 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. + +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, + linear_attention_training_phase, +) +from modelopt.torch.quantization.nn import TensorQuantizer + +KimiDeltaAttention = pytest.importorskip("fla.layers.kda").KimiDeltaAttention + + +def _layer(head_dim=16): + return ( + KimiDeltaAttention(hidden_size=32, head_dim=head_dim, num_heads=2, use_short_conv=True) + .cuda() + .train() + ) + + +def _forward(model, hidden): + 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", + [ + "decode_token", + "decode_replay", + "decode_decay", + "decode_token_int8", + "decode_replay_int8", + "decode_replay_hadamard_int8", + ], +) +@pytest.mark.timeout(180) +def test_fla_layer_qat_restore_and_optimizer(tmp_path, mode): + torch.manual_seed(73) + head_dim = 32 if "hadamard" in mode else 16 + model = _layer(head_dim) + hidden = torch.randn(1, 73, 32, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + baseline = _forward(model, hidden) + attributes = {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)} + cfg = { + "quant_cfg": [{"quantizer_name": "*", "enable": False}], + "algorithm": None, + "linear_attention": [{"module_name": "*", "cfg": {"backend": "matmul"}}], + } + state_attributes = {**attributes, "axis": (0, 1)} + if mode.endswith("int8"): + state_attributes.update(num_bits=8, unsigned=False, narrow_range=True) + if mode.startswith("decode"): + cfg["quant_cfg"].append({"quantizer_name": "*kda_state_quantizer", "cfg": state_attributes}) + policy_cfg = cfg["linear_attention"][0]["cfg"] + policy_cfg["state"] = {"block_v": 16} + policy_cfg["decode"] = { + "mode": "replay" if mode.startswith("decode_replay") else "token", + "implementation": "triton", + } + if mode.startswith("decode_replay"): + policy_cfg["decode"]["replay"] = {"window": 5} + if mode == "decode_decay": + policy_cfg["decode"]["decay_log_step"] = 1 / 256 + if "hadamard" in mode: + policy_cfg["state"]["block_v"] = 32 + policy_cfg["decode"]["state_codec"] = "int8_hadamard32" + mtq.quantize(model, cfg) + output = _forward(model, hidden) + assert torch.isfinite(output).all() + enabled = [q for q in model.modules() if isinstance(q, TensorQuantizer) and q.is_enabled] + policy = model.linear_attention_config + for q in enabled: + q.disable() + model.linear_attention_config = LinearAttentionConfig() + with torch.no_grad(): + torch.testing.assert_close(_forward(model, hidden), baseline, rtol=0, atol=0) + for q in enabled: + q.enable() + model.linear_attention_config = policy + checkpoint = tmp_path / "model.pth" + mto.save(model, checkpoint) + restored = mto.restore(_layer(head_dim), checkpoint) + actual = _forward(restored, hidden) + torch.testing.assert_close(actual, output, rtol=0, atol=0) + assert restored.linear_attention_config == policy + actual.float().square().mean().backward() + assert restored.q_proj.weight.grad is not None + for parameter in restored.parameters(): + assert parameter.grad is not None + assert torch.isfinite(parameter.grad).all() + before = restored.q_proj.weight.detach().clone() + torch.optim.SGD(restored.parameters(), lr=0.1).step() + assert not torch.equal(before, restored.q_proj.weight) + restored.eval() + 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]) + + +@pytest.mark.parametrize("hadamard", [False, True]) +def test_decode_phase_spans_activation_checkpoint_backward(hadamard): + model = _layer(32 if hadamard else 16) + quant_cfg = [{"quantizer_name": "*", "enable": False}] + if hadamard: + quant_cfg.append( + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": 8, + "type": "dynamic", + "axis": (0, 1), + "unsigned": False, + "narrow_range": True, + }, + } + ) + mtq.quantize( + model, + { + "quant_cfg": quant_cfg, + "algorithm": None, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "state_codec": "int8_hadamard32" if hadamard else "tile", + }, + }, + } + ], + }, + ) + 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/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index cf802abaf40..2039217a8e8 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -25,7 +25,7 @@ import modelopt.torch.kernels.quantization.gemm as triton_kernel import modelopt.torch.quantization.utils as quant_utils from modelopt.torch.quantization import tensor_quant -from modelopt.torch.quantization.extensions import get_cuda_ext, get_cuda_ext_mx +from modelopt.torch.quantization.extensions import get_cuda_ext, get_cuda_ext_fp8, get_cuda_ext_mx from modelopt.torch.quantization.tensor_quant import mx_format_map if triton_kernel.IS_AVAILABLE: @@ -120,6 +120,26 @@ def test_cuda_ext_tiny_amax(self): class TestScaledE4M3: + @pytest.mark.parametrize("per_channel", [False, True]) + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_eager_matches_cuda_scale_rounding(self, per_channel, dtype): + # These ranges distinguish direct division from reciprocal multiplication. + amax = torch.tensor([0.0, 1e-8, 2**-24, 1.04207686e-7, 0.91567558], device="cuda") + values = torch.linspace(-1, 1, 129, device="cuda")[None, :] * amax[:, None] + values = values.to(dtype) + extension = get_cuda_ext_fp8() + if per_channel: + actual = extension.fake_e4m3fy_with_axis(values, amax, 0) + expected = tensor_quant.fp8_eager(values, amax[:, None]) + else: + actual = torch.stack( + [extension.fake_e4m3fy(row, limit) for row, limit in zip(values, amax)] + ) + expected = torch.stack( + [tensor_quant.fp8_eager(row, limit) for row, limit in zip(values, amax)] + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + @pytest.mark.parametrize("device", ["cuda", "cpu"]) def test_e4m3_no_scale(self, device): x = torch.randn(4, 4, device=device, dtype=torch.float32) 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 new file mode 100644 index 00000000000..94ad1aaedf1 --- /dev/null +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -0,0 +1,276 @@ +# 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. + +from contextlib import nullcontext + +import pytest +import torch +from _test_utils.torch.megatron.models import get_mcore_gpt_model +from _test_utils.torch.megatron.utils import ( + get_forward, + initialize_for_megatron, + sharded_state_dict_test_helper, +) + +import modelopt.torch.quantization as mtq +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 +GatedDeltaNet = pytest.importorskip("megatron.core.ssm.gated_delta_net").GatedDeltaNet + +from modelopt.torch.quantization.plugins.gdn import _state_qdq_chunk_gated_delta_rule +from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet + +try: + _state_qdq_chunk_gated_delta_rule() +except RuntimeError as e: + pytest.skip(str(e), allow_module_level=True) + +SEED = 1234 + + +def _make_model(tp_size, **config_kwargs): + model = ( + get_mcore_gpt_model( + tensor_model_parallel_size=tp_size, + num_layers=2, + hidden_size=64, + num_attention_heads=4, + vocab_size=32, + max_sequence_length=128, + experimental_attention_variant="gated_delta_net", + linear_value_head_dim=64, + **config_kwargs, + ) + .cuda() + .eval() + ) + # Retain enough history for chunk-boundary state rounding to affect the next chunk. + with torch.no_grad(): + for module in model.modules(): + if isinstance(module, GatedDeltaNet): + module.A_log.fill_(-4) + module.dt_bias.zero_() + return model + + +def _gdn_config(sites): + return { + "quant_cfg": [{"quantizer_name": "*", "enable": False}] + + [ + { + "quantizer_name": f"*gdn_{site}_quantizer", + "cfg": { + "num_bits": (4, 3), + "axis": (0, 1) if site == "state" else (0, 1, 2), + "type": "dynamic", + }, + } + for site in sites + ], + "algorithm": None, + "linear_attention": [ + {"module_name": "decoder.layers.*.self_attention", "cfg": {"state": {"block_v": 16}}} + ], + } + + +def _config_for_mode(mode): + sites = ("state", "w") if mode == "both" else ((mode,) if mode in ("state", "w") else ()) + if mode.startswith("decode"): + sites = ("state",) + if mode == "state_int8": + sites = ("state",) + cfg = _gdn_config(sites) + if mode.endswith("int8"): + cfg["quant_cfg"][1]["cfg"].update(num_bits=8, unsigned=False, narrow_range=True) + if mode.startswith("decode"): + policy = cfg["linear_attention"][0]["cfg"] + policy.update(backend="matmul", decode={"implementation": "triton"}) + if mode.startswith("decode_replay"): + policy["decode"].update(mode="replay", replay={"window": 8}) + return cfg, sites + + +def _gdn_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) + + return forward + + +def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): + initialize_for_megatron( + tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED + ) + model = _make_model(size) + forward = _gdn_forward(model) + + outputs = [] + handles = [ + module.register_forward_hook( + lambda module, args, result: outputs.append(result[0].detach().clone()) + ) + for module in model.modules() + if isinstance(module, GatedDeltaNet) + ] + with torch.no_grad(): + loss_ref = forward(model) + gdn_ref = outputs.copy() + outputs.clear() + + cfg, sites = _config_for_mode(mode) + mtq.quantize(model, cfg) + gdn_modules = [m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)] + assert gdn_modules, "no GatedDeltaNet layer was wrapped" + for module in gdn_modules: + assert module.gdn_state_qdq_block_v == 16 + assert module.gdn_state_quantizer.is_enabled == ("state" in sites) + assert module.gdn_w_quantizer.is_enabled == ("w" in sites) + # Checkpointing must save the resolved policy, including edits after conversion. + module.linear_attention_config = LinearAttentionConfig( + **{**module.linear_attention_config.model_dump(), "state": {"block_v": 32}} + ) + + with torch.no_grad(): + loss_quant = forward(model) + assert torch.isfinite(loss_quant).all() + # BF16 residual additions can hide a small GDN change at the final loss. + assert len(outputs) == len(gdn_ref) == len(gdn_modules) + assert any(not torch.equal(actual, ref) for actual, ref in zip(outputs, gdn_ref)), ( + "enabled QDQ must change the GDN branch output" + ) + for handle in handles: + handle.remove() + + enabled = [ + q + for m in gdn_modules + for q in m.modules() + if isinstance(q, TensorQuantizer) and q.is_enabled + ] + policies = [m.linear_attention_config for m in gdn_modules] + for q in enabled: + q.disable() + for module in gdn_modules: + module.linear_attention_config = LinearAttentionConfig() + with torch.no_grad(): + torch.testing.assert_close(forward(model), loss_ref, rtol=1e-4, atol=1e-4) + for q in enabled: + q.enable() + for module, policy in zip(gdn_modules, policies): + module.linear_attention_config = policy + + restored = _make_model(size) + sharded_state_dict_test_helper(checkpoint_path, model, restored, forward) + restored_gdn = [m for m in restored.modules() if isinstance(m, _QuantGatedDeltaNet)] + assert len(restored_gdn) == len(gdn_modules) + for module, original in zip(restored_gdn, gdn_modules): + assert module.linear_attention_config == original.linear_attention_config + assert module.gdn_state_quantizer.is_enabled == ("state" in sites) + assert module.gdn_w_quantizer.is_enabled == ("w" in sites) + assert module._linear_attn_state_format == original._linear_attn_state_format + assert module.gdn_state_qdq_block_v == 32 + assert module.in_proj.weight.grad is not None + assert torch.isfinite(module.in_proj.weight.grad).all() + + # The checkpoint helper already ran backward through the restored model. + optimizer = torch.optim.SGD(restored.parameters(), lr=1e-3) + before = restored_gdn[0].in_proj.weight.detach().clone() + optimizer.step() + assert not torch.equal(restored_gdn[0].in_proj.weight, before) + + +def _compile_gdn_qat_kernels(rank, size, mode): + initialize_for_megatron( + tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED + ) + model = _make_model(size) + forward = _gdn_forward(model) + with torch.no_grad(): + forward(model) + cfg, _ = _config_for_mode(mode) + # The functional case switches state tiles before running inference. + cfg["linear_attention"][0]["cfg"]["state"]["block_v"] = 32 + mtq.quantize(model, cfg) + with torch.no_grad(): + forward(model) + model.train() + forward(model).sum().backward() + torch.cuda.synchronize() + + +@pytest.fixture +def compiled_gdn_workers(request, tp_size, mode): + """Warm the selected QAT path in the same workers, outside the test-call budget.""" + if mode in ("state", "both") and torch.cuda.get_device_capability() < (8, 9): + pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") + workers = request.getfixturevalue(f"dist_workers_size_{tp_size}") + workers.run(_compile_gdn_qat_kernels, mode) + return workers + + +@pytest.mark.parametrize("tp_size", [1, 2]) +@pytest.mark.parametrize( + "mode", + [ + "state", + "w", + "both", + "decode_token", + "decode_replay", + "state_int8", + "decode_token_int8", + "decode_replay_int8", + ], +) +def test_gdn_qat_and_sharded_restore(compiled_gdn_workers, tmp_path, mode): + """Train through state/W QDQ after a Megatron distributed-checkpoint round trip.""" + compiled_gdn_workers.run(_test_gdn_qat_helper, mode, tmp_path) + + +def _test_gdn_context_parallel_helper(rank, size, mode): + initialize_for_megatron(context_parallel_size=size, seed=SEED) + model = _make_model(1, context_parallel_size=size) + with pytest.raises(NotImplementedError, match="context parallelism"): + mtq.quantize(model, _config_for_mode(mode)[0]) + + +@pytest.mark.parametrize( + "mode", + [ + "state", + "w", + "decode_token", + "decode_replay", + "state_int8", + "decode_token_int8", + "decode_replay_int8", + ], +) +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.""" + dist_workers_size_2.run(_test_gdn_context_parallel_helper, mode) diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py new file mode 100644 index 00000000000..aed7f8caf42 --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -0,0 +1,338 @@ +# 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. + +from copy import deepcopy +from functools import partial + +import pytest +import torch +import torch.nn as nn + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.config import QuantizeConfig +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.plugins import gdn +from modelopt.torch.quantization.plugins.gdn import GatedDeltaNetStateQuantMixin + +GDN_STATE_INT8_DYNAMIC = { + "num_bits": 8, + "unsigned": False, + "narrow_range": True, + "type": "dynamic", + "axis": (0, 1), +} + +GDN_STATE_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"} + + +def chunk_gated_delta_rule(q, k, v, g, beta, **kwargs): + """CPU stand-in for the optional FLA kernel.""" + return q + k + v, None + + +@pytest.fixture(autouse=True) +def mock_fla_kernel(monkeypatch): + monkeypatch.setattr(gdn, "_fla_chunk_gated_delta_rule", lambda: chunk_gated_delta_rule) + + +class TinyGatedDeltaNet(nn.Module): + """A module that, like Megatron-Core's GatedDeltaNet, calls ``self.gated_delta_rule``.""" + + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4) + self.gated_delta_rule = chunk_gated_delta_rule + + def forward(self, x): + out, _ = self.gated_delta_rule(x, x, x, x[..., 0], x[..., 0]) + return self.proj(out) + + +@QuantModuleRegistry.register({TinyGatedDeltaNet: "TinyGatedDeltaNet"}) +class _QuantTinyGatedDeltaNet(GatedDeltaNetStateQuantMixin): + def forward(self, x): + gated_delta_rule = self.gated_delta_rule + self.gated_delta_rule = lambda *a, **kw: self._state_quantized_chunk_gated_delta_rule( + gated_delta_rule, *a, **kw + ) + try: + return super().forward(x) + finally: + self.gated_delta_rule = gated_delta_rule + + +GDN_W_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1, 2), "type": "dynamic"} + + +def quant_cfg(state=True, w=False): + entries = [{"quantizer_name": "*", "enable": False}] + if state: + entries.append({"quantizer_name": "*gdn_state_quantizer", "cfg": GDN_STATE_FP8_DYNAMIC}) + if w: + entries.append({"quantizer_name": "*gdn_w_quantizer", "cfg": GDN_W_FP8_DYNAMIC}) + return {"quant_cfg": entries, "algorithm": "max"} + + +@pytest.mark.parametrize( + "attributes", + [ + {"num_bits": (4, 3), "axis": (0, 1)}, # static + {"num_bits": (4, 3), "type": "dynamic"}, # per tensor + {"num_bits": 8, "axis": (0, 1), "type": "dynamic"}, # int8 + {"num_bits": (4, 3), "type": "dynamic", "block_sizes": {-1: 16}}, # blockwise + ], +) +def test_validate_state_quantizer_rejects_unsupported(attributes): + with pytest.raises(ValueError, match="supports only"): + mtq.quantize( + TinyGatedDeltaNet(), + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*gdn_state_quantizer", "cfg": attributes}, + ], + "algorithm": None, + }, + ) + + +def test_dynamic_export_removes_linear_attention_attributes(): + model = _QuantTinyGatedDeltaNet.convert(TinyGatedDeltaNet()) + model.export() + assert type(model) is TinyGatedDeltaNet + for name in ( + "gdn_state_quantizer", + "gdn_w_quantizer", + "linear_attention_config", + "_linear_attention_prefill_lengths", + ): + assert not hasattr(model, name) + + +def test_disabled_state_quantizer_calls_original_kernel(): + model = TinyGatedDeltaNet() + x = torch.randn(2, 8, 3, 4) + expected = model(x) + + disable_all = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": "max"} + mtq.quantize(model, disable_all, lambda m: m(x)) + + assert isinstance(model, _QuantTinyGatedDeltaNet) + assert not model.gdn_state_quantizer.is_enabled and not model.gdn_w_quantizer.is_enabled + assert torch.equal(model(x), expected) + assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" + + +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +@pytest.mark.parametrize("partial_kernel", [False, True]) +def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, state_format, partial_kernel): + calls = [] + + def fake_state_qdq_kernel(*args, **kwargs): + calls.append(kwargs) + return chunk_gated_delta_rule(*args) + + monkeypatch.setattr(gdn, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel) + model = TinyGatedDeltaNet() + if partial_kernel: + model.gated_delta_rule = partial(chunk_gated_delta_rule, output_final_state=True) + x = torch.randn(2, 8, 3, 4) + cfg = quant_cfg() + if state_format == "int8": + cfg["quant_cfg"][1]["cfg"] = GDN_STATE_INT8_DYNAMIC + mtq.quantize(model, cfg, lambda m: m(x)) + + model(x) + assert calls and calls[-1] == { + "chunk_size": 64, + "state_qdq": 2 if state_format == "int8" else 1, + "state_qdq_block_v": 64, + "w_quantizer": None, + **({"output_final_state": True} if partial_kernel else {}), + } + + # An unrelated callable must not pass validation merely by copying the FLA name. + model.gated_delta_rule = lambda *a, **kw: chunk_gated_delta_rule(*a, **kw) + model.gated_delta_rule.__name__ = "chunk_gated_delta_rule" + with pytest.raises(NotImplementedError, match="supports only FLA"): + model(x) + + +@pytest.mark.parametrize("state", [False, True]) +def test_w_quantizer_is_passed_to_the_kernel(monkeypatch, state): + """``*gdn_w_quantizer`` in the config hands the module's TensorQuantizer to the kernel, with + or without the state quantizer.""" + calls = [] + + def fake_state_qdq_kernel(*args, **kwargs): + calls.append(kwargs) + return chunk_gated_delta_rule(*args) + + monkeypatch.setattr(gdn, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel) + model = TinyGatedDeltaNet() + x = torch.randn(2, 8, 3, 4) + mtq.quantize(model, quant_cfg(state=state, w=True), lambda m: m(x)) + assert model.gdn_w_quantizer.is_enabled and model.gdn_state_quantizer.is_enabled == state + + model(x) + assert calls[-1]["state_qdq"] == int(state) + assert calls[-1]["w_quantizer"] is model.gdn_w_quantizer + + # The w quantizer really quantizes: 256 random values per token collapse onto the E4M3 grid, + # which has at most 127 distinct magnitudes per (row-specific) scale. + w = torch.randn(1, 1, 1, 256) + quantized = model.gdn_w_quantizer(w) + assert not torch.equal(quantized, w) + assert torch.unique(quantized.abs()).numel() <= 127 < torch.unique(w.abs()).numel() + + +@pytest.mark.parametrize("site", ["state", "w"]) +@pytest.mark.parametrize( + "overrides", + [{"pass_through_bwd": False}, {"type": "static"}, {"fake_quant": False}, {"rotate": True}], +) +def test_unsupported_quantizer_fails_during_conversion(site, overrides): + cfg = quant_cfg(state=site == "state", w=site == "w") + cfg = deepcopy(cfg) + cfg["quant_cfg"][-1]["cfg"].update(overrides) + with pytest.raises(ValueError, match="supports only"): + mtq.quantize(TinyGatedDeltaNet(), cfg) + + +@pytest.mark.parametrize("axis", [None, (0, 1), (0, 1, 2)]) +def test_w_grouping_is_preserved_during_conversion(axis): + cfg = deepcopy(quant_cfg(state=False, w=True)) + cfg["algorithm"] = None + cfg["quant_cfg"][-1]["cfg"]["axis"] = axis + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + assert model.gdn_w_quantizer.axis == axis + + +@pytest.mark.parametrize(("state", "w"), [(True, False), (False, True), (True, True)]) +def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): + model = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) + cfg = quant_cfg(state=state, w=w) + cfg["algorithm"] = None + cfg["linear_attention"] = [ + {"module_name": "*", "cfg": {"state": {"block_v": 128}}}, + {"module_name": "0", "cfg": {"state": {"block_v": 32}}}, + ] + mtq.quantize(model, cfg) + assert model[0].gdn_state_qdq_block_v == 32 + model[0].linear_attention_config.state.block_v = 16 + if w: + # Save the actual quantizer settings, including edits after conversion. + model[0].gdn_w_quantizer.axis = None + path = tmp_path / "gdn.pth" + mto.save(model, path) + restored = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) + mto.restore(restored, path) + assert restored[0].linear_attention_config == model[0].linear_attention_config + assert restored[0].gdn_state_qdq_block_v == 16 + for name, enabled in (("gdn_state_quantizer", state), ("gdn_w_quantizer", w)): + original = getattr(model[0], name) + quantizer = getattr(restored[0], name) + assert quantizer.is_enabled == enabled + assert quantizer.axis == original.axis + assert quantizer.num_bits == original.num_bits + assert quantizer._dynamic == original._dynamic + assert not hasattr(restored[1], name) + + +def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): + cfg = quant_cfg() + cfg["algorithm"] = None + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + assert model.gdn_state_quantizer.is_enabled + assert not model.gdn_w_quantizer.is_enabled + + cfg = quant_cfg(state=False, w=True) + cfg["algorithm"] = None + cfg["linear_attention"] = [{"module_name": "", "cfg": {"state": {"block_v": 32}}}] + mtq.quantize(model, cfg) + assert model.gdn_state_qdq_block_v == 32 + cfg["linear_attention"].append({"module_name": "", "cfg": {}}) + mtq.quantize(model, cfg) + assert model.gdn_state_qdq_block_v == 64 + assert not model.gdn_state_quantizer.is_enabled + assert model.gdn_w_quantizer.is_enabled + + cfg = deepcopy(cfg) + cfg["quant_cfg"][-1]["cfg"]["pass_through_bwd"] = False + with pytest.raises(ValueError, match="supports only"): + mtq.quantize(model, cfg) + + +def test_restore_legacy_gdn_without_new_quantizer_handles(): + config = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None} + model = mtq.quantize(TinyGatedDeltaNet(), config) + state = mto.modelopt_state(model) + for _, mode_state in state["modelopt_state_dict"]: + mode_state["config"].pop("linear_attention", None) + metadata = mode_state["metadata"] + metadata.pop("linear_attention", None) + for name in ("gdn_state_quantizer", "gdn_w_quantizer"): + metadata["quantizer_state"].pop(name, None) + restored = TinyGatedDeltaNet() + mto.restore_from_modelopt_state(restored, state) + restored.load_state_dict(model.state_dict()) + x = torch.randn(2, 8, 3, 4) + torch.testing.assert_close(restored(x), model(x)) + assert not restored.gdn_state_quantizer.is_enabled + assert not restored.gdn_w_quantizer.is_enabled + + +def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): + model = mtq.quantize( + TinyGatedDeltaNet(), mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(2, 8, 3, 4)) + ) + assert not model.gdn_state_quantizer.is_enabled + assert not model.gdn_w_quantizer.is_enabled + + +def test_int8_state_conversion_and_checkpoint(tmp_path): + cfg = load_recipe("general/ptq/gdn_state_int8_dynamic").quantize + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + assert model.gdn_state_quantizer.is_enabled + assert model._linear_attn_state_format == "int8" + assert model.linear_attention_config.backend == "matmul" + assert model.linear_attention_config.decode.state_codec == "int8_hadamard32" + assert not model.linear_attention_config.decode.prefill_state_qdq + assert not model.proj.weight_quantizer.is_enabled + path = tmp_path / "int8.pt" + mto.save(model, path) + restored = mto.restore(TinyGatedDeltaNet(), path) + assert restored.linear_attention_config == model.linear_attention_config + assert restored.gdn_state_quantizer.is_enabled + assert restored._linear_attn_state_format == "int8" + assert restored.gdn_state_quantizer.narrow_range + assert not restored.gdn_state_quantizer.unsigned + assert not restored.gdn_w_quantizer.is_enabled + + +def test_policy_rejects_unmatched_and_unimplemented_modes(): + with pytest.raises(ValueError, match="matches no supported"): + mtq.quantize( + nn.Linear(4, 4), {"linear_attention": [{"module_name": "*"}], "algorithm": None} + ) + for policy in ( + {"chunk_size": 32}, + {"solve": {"method": "neumann"}}, + {"state": {"mode": "token"}}, + ): + with pytest.raises(ValueError): + QuantizeConfig(linear_attention=[{"module_name": "*", "cfg": policy}]) 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..e0ced168024 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -0,0 +1,296 @@ +# 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 _test_utils.torch.quantization.linear_attention_reference import recurrent_delta_rule_reference + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + matmul_gdn, + matmul_kda, +) +from modelopt.torch.quantization.linear_attention.config import LinearAttentionDecodeConfig +from modelopt.torch.quantization.linear_attention.decode import _encode, recurrent_decode_reference +from modelopt.torch.quantization.linear_attention.utils import _state_qdq +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.mark.parametrize("state", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_fp8_qdq_matches_tensor_quantizer(state, dtype): + torch.manual_seed(762) + amplitudes = torch.tensor([0.0, 1e-8, 2**-24, 2**-23, 1.0, 731.173], dtype=dtype) + value = torch.rand((2, 6, 5, 19) if state else (2, 6, 19), dtype=dtype) * 2 - 1 + value *= amplitudes.reshape(1, 6, *([1] * (value.ndim - 2))) + value.requires_grad_() + quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=(4, 3), type="dynamic", axis=(0, 1), pass_through_bwd=True + ) + ) + expected = torch.cat( + [ + quantizer(tile.flatten(-2) if state else tile).reshape_as(tile) + for tile in value.split(16, -1) + ], + -1, + ) + encoded = _encode(value, True, 16, state=state) + results = [encoded.values] + if state: + results.append(_state_qdq(value, 16)) + probe = torch.randn_like(value) + (expected_gradient,) = torch.autograd.grad((expected * probe).sum(), value) + for result in results: + torch.testing.assert_close(result, expected, rtol=0, atol=0) + (gradient,) = torch.autograd.grad((result * probe).sum(), value) + torch.testing.assert_close(gradient, expected_gradient, rtol=0, atol=0) + assert encoded.scales.shape == (2, 6, 2) + assert not encoded.scales.requires_grad + + +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, + policy=LinearAttentionConfig(backend="matmul", decode={}), + 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 = { + "policy": policy, + "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_hadamard.py b/tests/unit/torch/quantization/test_linear_attention_hadamard.py new file mode 100644 index 00000000000..5f6b40f4c5c --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_hadamard.py @@ -0,0 +1,227 @@ +# 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 json + +import pytest +import torch +import torch.nn.functional as F + +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.config import QuantizeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + LinearAttentionDecodeConfig, + matmul_gdn, + matmul_kda, + recurrent_decode_reference, +) +from modelopt.torch.quantization.linear_attention.decode import _encode, _hadamard32 + + +def _rotate(value): + # Independent dense Sylvester matrix; production uses a butterfly transform. + matrix = value.new_tensor([[(-1) ** (i & j).bit_count() for j in range(32)] for i in range(32)]) + return (value.reshape(*value.shape[:-1], -1, 32) @ (matrix / 32**0.5)).reshape_as(value) + + +def _qdq(value): + with torch.no_grad(): + groups = value.float().unflatten(-1, (-1, 32)) + scale = (groups.abs().amax(-1, keepdim=True) / 127).clamp_min(6e-8) + normalized = groups / scale + codes = torch.where(normalized >= 0, (normalized + 0.5).floor(), (normalized - 0.5).ceil()) + rounded = (codes.clamp(-127, 127) * scale.half().float()).flatten(-2).to(value.dtype) + return (value - value.detach()) + rounded + + +def _inputs(kda, length=11): + torch.manual_seed(321) + q, k = [F.normalize(torch.randn(length, 2, 4, dtype=torch.float64), dim=-1) for _ in range(2)] + v = torch.randn(length, 2, 64, dtype=torch.float64) + g = -torch.rand(k.shape if kda else k.shape[:-1], dtype=torch.float64) * 0.05 + beta = torch.rand(length, 2, dtype=torch.float64) * 0.4 + initial = torch.randn(2, 4, 64, dtype=torch.float64) * 0.1 + return tuple(x.requires_grad_() for x in (q, k, v, g, beta)), initial.requires_grad_() + + +def _oracle(args, initial, mode, readout, quantize=True): + q, k, v, g, beta = args + state = _rotate(initial) + if quantize: + state = _qdq(state) + outputs = [] + for t in range(len(q)): + decay = g[t].exp().reshape(2, -1, 1) + decayed = state * decay + correction = beta[t, :, None] * (_rotate(v[t]) - torch.einsum("hk,hkv->hv", k[t], decayed)) + working = decayed + k[t, :, :, None] * correction[:, None, :] + state = _qdq(working) if quantize and (mode == "token" or (t + 1) % 3 == 0) else working + read = state if readout == "stored" else working + outputs.append(_rotate(torch.einsum("hk,hkv->hv", q[t], read)) / q.shape[-1] ** 0.5) + return torch.stack(outputs), _rotate(state) + + +def _check_with_grads(actual, expected, args, initial, tolerance=2e-10): + for a, e in zip(actual, expected): + torch.testing.assert_close(a, e, rtol=tolerance, atol=tolerance) + 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=tolerance, atol=tolerance) + + +def test_hadamard_codec_basis_scales_ties_and_identity_gradient(): + raw = torch.zeros(1, 2, 64, dtype=torch.float64) + raw[0, 0, :7] = torch.tensor([254, -254, 1, -1, 3, -3, 0]) + raw[0, 1, 32:] = torch.linspace(-0.123, 0.123, 32) + raw.requires_grad_() + torch.testing.assert_close(_hadamard32(raw), _rotate(raw), rtol=1e-14, atol=1e-14) + torch.testing.assert_close(_hadamard32(_hadamard32(raw)), raw, rtol=1e-14, atol=1e-13) + encoded = _encode(raw, True, 64, state=True, state_format="int8", state_codec="int8_hadamard32") + torch.testing.assert_close(encoded.values, _qdq(raw), rtol=0, atol=0) + torch.testing.assert_close( + encoded.values[0, 0, :7], raw.new_tensor([254, -254, 2, -2, 4, -4, 0]) + ) + assert encoded.scales.shape == (1, 2, 2) + assert encoded.scales.dtype == torch.float16 and not encoded.scales.requires_grad + assert encoded.scales[0, 0, 0] == 2 + assert encoded.scales[0, 0, 1] == torch.tensor(6e-8, dtype=torch.float16) + probe = torch.randn_like(raw) + (gradient,) = torch.autograd.grad((encoded.values * probe).sum(), raw) + torch.testing.assert_close(gradient, probe, rtol=0, atol=0) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize( + ("mode", "readout"), [("token", "stored"), ("replay", "working"), ("replay", "stored")] +) +def test_hadamard_recurrence_matches_dense_oracle_and_split_carry(kda, mode, readout): + args, initial = _inputs(kda) + cfg = LinearAttentionDecodeConfig( + mode=mode, + readout=readout, + state_codec="int8_hadamard32", + replay={"window": 3, "factor_qdq": False} if mode == "replay" else None, + ) + kwargs = {"config": cfg, "state_format": "int8", "state_qdq": True} + output, carry = recurrent_decode_reference(*args, initial_state=initial, **kwargs) + _check_with_grads( + (output, carry.reconstruct()), _oracle(args, initial, mode, readout), args, initial + ) + assert carry.value_basis == "hadamard32" and carry.anchor.block_v == 32 + assert carry.anchor.scales.shape == (2, 4, 2) + pieces, split_carry, start = [], None, 0 + for end in (0, 2, 3, 3, 7, 11): + piece, split_carry = recurrent_decode_reference( + *(x[start:end] for x in args), + carry=split_carry, + initial_state=initial if split_carry is None else None, + **kwargs, + ) + if end == 0: + assert split_carry.reconstruct() is initial and not split_carry.started + pieces.append(piece) + start = end + _check_with_grads( + (torch.cat(pieces), split_carry.reconstruct()), + (output, carry.reconstruct()), + args, + initial, + 0, + ) + with pytest.raises(ValueError, match="Carry policy"): + recurrent_decode_reference( + *args, + carry=carry, + **{**kwargs, "config": cfg.model_copy(update={"state_codec": "tile"})}, + ) + + +@pytest.mark.parametrize("kda", [False, True]) +def test_hadamard_disabled_qdq_is_a_change_of_basis(kda): + args, initial = _inputs(kda) + cfg = LinearAttentionDecodeConfig(mode="replay", replay={"window": 3, "factor_qdq": False}) + results = [] + for codec in ("tile", "int8_hadamard32"): + out, carry = recurrent_decode_reference( + *args, initial_state=initial, config=cfg.model_copy(update={"state_codec": codec}) + ) + results.append((out, carry.reconstruct())) + _check_with_grads(*results, args, initial) + + +@pytest.mark.parametrize(("kda", "prefix"), [(False, 0), (False, 3), (True, 3), (True, 11)]) +def test_hadamard_exact_prefill_handoff_and_state_layout(kda, prefix, project_root_path): + args, initial = _inputs(kda) + function = matmul_kda if kda else matmul_gdn + if kda: + path = ( + project_root_path + / "examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json" + ) + recipe = QuantizeConfig(**json.loads(path.read_text())) + else: + recipe = load_recipe("general/ptq/gdn_state_int8_dynamic").quantize + policy = recipe.linear_attention[0].cfg + assert recipe.algorithm is None + assert recipe.quant_cfg[1].cfg.num_bits == 8 + assert policy.decode.state_codec == "int8_hadamard32" + assert not policy.decode.prefill_state_qdq + # Exercise the shipped policy with the CPU implementation for this algebra test. + policy.decode.implementation = "torch" + actual = function( + *(x.unsqueeze(0) for x in args), + policy=policy, + prefill_lengths=[prefix], + state_qdq=True, + state_format="int8", + initial_state=initial.transpose(-1, -2).unsqueeze(0), + state_v_first=True, + output_final_state=True, + ) + prefix_out, state = recurrent_decode_reference( + *(x[:prefix] for x in args), initial_state=initial, config=LinearAttentionDecodeConfig() + ) + suffix, carry = recurrent_decode_reference( + *(x[prefix:] for x in args), + initial_state=state.reconstruct(), + config=policy.decode, + state_qdq=True, + state_format="int8", + ) + expected = (torch.cat((prefix_out, suffix)), carry.reconstruct()) + _check_with_grads((actual[0][0], actual[1][0].transpose(-1, -2)), expected, args, initial) + + +def test_hadamard_rejects_unsupported_codec_combinations(): + args, initial = _inputs(False) + cfg = LinearAttentionDecodeConfig(state_codec="int8_hadamard32") + for kwargs, message in [ + ({"state_qdq": True}, "requires INT8"), + ({"state_format": "int8", "block_v": 16}, "block_v >= 32"), + ]: + with pytest.raises(ValueError, match=message): + recurrent_decode_reference(*args, config=cfg, initial_state=initial, **kwargs) + with pytest.raises(ValueError, match="Dv divisible"): + recurrent_decode_reference(args[0], args[1], args[2][..., :33], *args[3:], config=cfg) + with pytest.raises(ValueError, match="decode handoff"): + LinearAttentionDecodeConfig(state_codec="int8_hadamard32", prefill_state_qdq=True) + with pytest.raises(ValueError, match="block_v >= 32"): + LinearAttentionConfig(backend="matmul", state={"block_v": 16}, decode=cfg) diff --git a/tests/unit/torch/quantization/test_linear_attention_int8.py b/tests/unit/torch/quantization/test_linear_attention_int8.py new file mode 100644 index 00000000000..b3de8c94d47 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_int8.py @@ -0,0 +1,99 @@ +# 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 +from _test_utils.torch.quantization.linear_attention_reference import state_qdq_reference + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionDecodeConfig, + recurrent_decode_reference, +) +from modelopt.torch.quantization.linear_attention.decode import _encode +from modelopt.torch.quantization.nn import TensorQuantizer + +INT8 = {"num_bits": 8, "unsigned": False, "narrow_range": True, "type": "dynamic", "axis": (0, 1)} + + +@pytest.mark.parametrize("block_v", [16, 32, 64, 128]) +def test_int8_state_codec_matches_modelopt_and_identity_gradient(block_v): + # Exact scale two exercises both signs, endpoints and nearest-even half ties. + state = torch.zeros(2, 3, 5, block_v + 3) + state[..., 0, :9] = torch.tensor([-127, -126.5, -1.5, -0.5, 0, 0.5, 1.5, 126.5, 127]) * 2 + state[0, :, 0, block_v:] = torch.tensor([-254, 1, 254]) + state.requires_grad_() + quantizer = TensorQuantizer(QuantizerAttributeConfig(**INT8)) + expected = torch.cat( + [quantizer(tile.flatten(-2)).reshape_as(tile) for tile in state.split(block_v, -1)], -1 + ) + encoded = _encode(state, True, block_v, state=True, state_format="int8") + reference = state_qdq_reference(state, block_v, "int8") + torch.testing.assert_close(encoded.values, expected, rtol=0, atol=0) + torch.testing.assert_close(reference, expected, rtol=0, atol=0) + assert encoded.format == "int8" + assert torch.equal(encoded.scales[1, :, -1], torch.ones(3)) + assert not encoded.scales.requires_grad + probe = torch.randn_like(state) + for result in (encoded.values, reference): + (gradient,) = torch.autograd.grad((result * probe).sum(), state, retain_graph=True) + torch.testing.assert_close(gradient, probe, rtol=0, atol=0) + + +def _inputs(kda): + torch.manual_seed(913) + q, k = [F.normalize(torch.randn(23, 2, 16), dim=-1).requires_grad_() for _ in range(2)] + v = torch.randn(23, 2, 19, requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1]) * 0.03).requires_grad_() + beta = (torch.rand(23, 2) * 0.4).requires_grad_() + state = (torch.randn(2, 16, 19) * 0.1).requires_grad_() + return (q, k, v, g, beta), state + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("mode", ["token", "replay"]) +def test_int8_decode_continuation_and_format_rejection(kda, mode): + args, initial = _inputs(kda) + cfg = LinearAttentionDecodeConfig(mode=mode, replay={"window": 5} if mode == "replay" else None) + kwargs = {"config": cfg, "state_qdq": True, "state_format": "int8", "block_v": 16} + expected, final = recurrent_decode_reference(*args, initial_state=initial, **kwargs) + pieces, carry, start = [], None, 0 + for end in (0, 3, 10, 10, 23): + output, carry = recurrent_decode_reference( + *(x[start:end] for x in args), + carry=carry, + initial_state=initial if carry is None else None, + **kwargs, + ) + pieces.append(output) + start = end + actual = torch.cat(pieces) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(carry.reconstruct(), final.reconstruct(), rtol=0, atol=0) + probes = (torch.randn_like(expected), torch.randn_like(initial)) + grads = [ + torch.autograd.grad((o * probes[0]).sum() + (s * probes[1]).sum(), (*args, initial)) + for o, s in ((actual, carry.reconstruct()), (expected, final.reconstruct())) + ] + for actual_grad, expected_grad in zip(*grads): + torch.testing.assert_close(actual_grad, expected_grad, rtol=0, atol=0) + assert carry.anchor.format == "int8" + assert all(e.key.format == e.update.format == "fp8_e4m3" for e in carry.entries) + with pytest.raises(ValueError, match="Carry policy"): + recurrent_decode_reference( + *(x[:0] for x in args), carry=carry, **{**kwargs, "state_format": "fp8_e4m3"} + ) diff --git a/tests/unit/torch/quantization/test_linear_attention_reference.py b/tests/unit/torch/quantization/test_linear_attention_reference.py new file mode 100644 index 00000000000..d71818d0163 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_reference.py @@ -0,0 +1,124 @@ +# 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 +from _test_utils.torch.quantization.linear_attention_reference import ( + chunk_gdn_reference, + recurrent_delta_rule_reference, + state_fp8_qdq_reference, +) + + +def inputs(length=7, kda=False): + torch.manual_seed(81) + q, k = [ + F.normalize(torch.randn(1, length, 1, 3, dtype=torch.float64), dim=-1) for _ in range(2) + ] + v = torch.randn(1, length, 2, 4, dtype=torch.float64) + g = -torch.rand(1, length, 2, *([3] if kda else []), dtype=torch.float64) + beta = torch.rand(1, length, 2, dtype=torch.float64) + return [x.requires_grad_() for x in (q, k, v, g, beta)] + + +@pytest.mark.parametrize("packed", [False, True]) +@pytest.mark.parametrize("state_v_first", [False, True]) +def test_chunk_recurrence_output_state_and_gradients(packed, state_v_first): + args = inputs() + state = torch.randn(2 if packed else 1, 2, 3, 4, dtype=torch.float64) + if state_v_first: + state = state.transpose(-1, -2).contiguous() + state.requires_grad_() + kwargs = { + "initial_state": state, + "state_v_first": state_v_first, + "cu_seqlens": torch.tensor([0, 2, 7]) if packed else None, + } + actual = chunk_gdn_reference(*args, chunk_size=3, **kwargs) + expected = recurrent_delta_rule_reference(*args, **kwargs) + for a, e in zip(actual, expected): + torch.testing.assert_close(a, e, rtol=1e-10, atol=1e-10) + probe = [torch.randn_like(x) for x in actual] + grads = [ + torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probe)), (*args, state)) + for result in (actual, expected) + ] + for a, e in zip(*grads): + torch.testing.assert_close(a, e, rtol=1e-9, atol=1e-9) + + +@pytest.mark.parametrize("kda", [False, True]) +def test_recurrence_smooth_gradcheck(kda): + args = inputs(length=2, kda=kda) + state = torch.randn(1, 2, 3, 4, dtype=torch.float64, requires_grad=True) + assert torch.autograd.gradcheck( + lambda *x: recurrent_delta_rule_reference(*x[:-1], initial_state=x[-1]), + (*args, state), + fast_mode=True, + ) + + +def test_kda_one_step_and_scalar_decay_equivalence(): + q, k, v, g, beta = inputs(length=1, kda=True) + state = torch.randn(1, 2, 3, 4, dtype=torch.float64) + out, final = recurrent_delta_rule_reference(q, k, v, g, beta, initial_state=state) + keys = k[:, 0].expand(-1, 2, -1) + decayed = state * g[:, 0].exp().unsqueeze(-1) + residual = v[:, 0] - torch.einsum("bhk,bhkv->bhv", keys, decayed) + expected = decayed + torch.einsum("bhk,bhv->bhkv", keys, beta[:, 0, :, None] * residual) + torch.testing.assert_close(final, expected) + torch.testing.assert_close( + out[:, 0], torch.einsum("bhk,bhkv->bhv", q[:, 0].expand(-1, 2, -1), expected) / 3**0.5 + ) + scalar = g[..., 0] + gdn = recurrent_delta_rule_reference(q, k, v, scalar, beta, initial_state=state) + kda = recurrent_delta_rule_reference( + q, k, v, scalar.unsqueeze(-1).expand_as(g), beta, initial_state=state + ) + for a, e in zip(kda, gdn): + torch.testing.assert_close(a, e) + + +def test_state_qdq_granularity_zero_tail_and_identity_ste(): + torch.manual_seed(14) + state = torch.randn(2, 2, 3, 37, requires_grad=True) + with torch.no_grad(): + state[0].zero_() + state[1, ..., :16] *= 10 + quantized = state_fp8_qdq_reference(state, block_v=16) + assert torch.isfinite(quantized).all() + assert torch.equal(quantized[0], state[0]) + assert not torch.equal(quantized, state) + assert not torch.equal(quantized, state_fp8_qdq_reference(state, block_v=64)) + probe = torch.randn_like(state) + (grad,) = torch.autograd.grad((quantized * probe).sum(), state) + torch.testing.assert_close(grad, probe, rtol=0, atol=0) + + +def test_packed_sequences_reset_the_state(): + args = inputs(length=7, kda=True) + initial = torch.randn(2, 2, 3, 4, dtype=torch.float64) + packed = recurrent_delta_rule_reference( + *args, initial_state=initial, cu_seqlens=torch.tensor([0, 2, 7]) + ) + separate = [ + recurrent_delta_rule_reference( + *(x[:, lo:hi] for x in args), initial_state=initial[n : n + 1] + ) + for n, (lo, hi) in enumerate(((0, 2), (2, 7))) + ] + torch.testing.assert_close(packed[0], torch.cat([x[0] for x in separate], dim=1)) + torch.testing.assert_close(packed[1], torch.cat([x[1] for x in separate]))