From e38b2f9e3a01b6c45dcb34b7b3092e15d681b359 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sun, 27 Sep 2026 20:55:54 -0700 Subject: [PATCH 01/27] Add GDN/KDA Torch decode QAT with INT8 recurrent state Signed-off-by: Kai Xu --- CHANGELOG.rst | 1 + examples/llm_qat/README.md | 1 + examples/llm_qat/linear_attention/README.md | 304 ++++++++++++++++ .../configs/kda_decode_state_int8.json | 35 ++ .../llm_qat/linear_attention/requirements.txt | 4 + examples/llm_qat/linear_attention/train.py | 109 ++++++ .../linear_attention/fla_chunk_delta_h.py | 68 ++-- .../fla_chunk_gated_delta_rule.py | 12 +- .../quantization/linear_attention/int8.py | 28 ++ modelopt/torch/quantization/config.py | 7 + modelopt/torch/quantization/conversion.py | 55 ++- .../quantization/linear_attention/__init__.py | 8 +- .../linear_attention/_chunk_prefill.py | 146 ++++++++ .../quantization/linear_attention/config.py | 108 ++++++ .../quantization/linear_attention/decode.py | 341 ++++++++++++++++++ .../linear_attention/decode_prefill.py | 176 +++++++++ .../quantization/linear_attention/kda.py | 107 ++++++ .../quantization/linear_attention/prefill.py | 94 +++++ .../quantization/linear_attention/utils.py | 109 +++++- modelopt/torch/quantization/model_quant.py | 4 +- .../torch/quantization/plugins/__init__.py | 3 +- .../plugins/{gated_delta_net.py => gdn.py} | 65 ++-- modelopt/torch/quantization/plugins/kda.py | 103 ++++++ .../quantization/plugins/linear_attention.py | 85 +++++ .../torch/quantization/plugins/megatron.py | 6 +- modelopt/torch/quantization/tensor_quant.py | 3 +- modelopt_recipes/configs/ptq/units/README.md | 1 + .../units/default_disabled_quantizers.yaml | 4 + .../ptq/units/gdn_state_fp8_dynamic.yaml | 7 +- .../ptq/units/gdn_state_int8_dynamic.yaml | 26 ++ .../general/ptq/gdn_state_int8_dynamic.yaml | 38 ++ modelopt_recipes/ptq.md | 1 + .../linear_attention_reference.py | 183 +++++++++- .../test_fla_chunk_gated_delta_rule.py | 16 +- .../torch/quantization/plugins/test_kda.py | 189 ++++++++++ .../quantization/test_tensor_quant_cuda.py | 22 +- .../plugins/test_megatron_gated_delta_net.py | 100 ++++- .../{test_gated_delta_net.py => test_gdn.py} | 92 ++++- .../test_linear_attention_decode.py | 296 +++++++++++++++ .../test_linear_attention_hadamard.py | 227 ++++++++++++ .../test_linear_attention_int8.py | 99 +++++ 41 files changed, 3127 insertions(+), 156 deletions(-) create mode 100644 examples/llm_qat/linear_attention/README.md create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json create mode 100644 examples/llm_qat/linear_attention/requirements.txt create mode 100644 examples/llm_qat/linear_attention/train.py create mode 100644 modelopt/torch/kernels/quantization/linear_attention/int8.py create mode 100644 modelopt/torch/quantization/linear_attention/_chunk_prefill.py create mode 100644 modelopt/torch/quantization/linear_attention/config.py create mode 100644 modelopt/torch/quantization/linear_attention/decode.py create mode 100644 modelopt/torch/quantization/linear_attention/decode_prefill.py create mode 100644 modelopt/torch/quantization/linear_attention/kda.py create mode 100644 modelopt/torch/quantization/linear_attention/prefill.py rename modelopt/torch/quantization/plugins/{gated_delta_net.py => gdn.py} (70%) create mode 100644 modelopt/torch/quantization/plugins/kda.py create mode 100644 modelopt/torch/quantization/plugins/linear_attention.py create mode 100644 modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml create mode 100644 modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml create mode 100644 tests/gpu/torch/quantization/plugins/test_kda.py rename tests/unit/torch/quantization/plugins/{test_gated_delta_net.py => test_gdn.py} (74%) create mode 100644 tests/unit/torch/quantization/test_linear_attention_decode.py create mode 100644 tests/unit/torch/quantization/test_linear_attention_hadamard.py create mode 100644 tests/unit/torch/quantization/test_linear_attention_int8.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index df983a8804e..c7bcc2789c6 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,6 +16,7 @@ Changelog - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. +- 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. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before a separate KV-cache AutoQuantize stage, with independent resumable checkpoints for the weight and diff --git a/examples/llm_qat/README.md b/examples/llm_qat/README.md index cd88ae30498..a32c25368d7 100644 --- a/examples/llm_qat/README.md +++ b/examples/llm_qat/README.md @@ -13,6 +13,7 @@ For background on QAT and QAD and help choosing between Hugging Face, Megatron B | Arguments | Full CLI/YAML argument reference | \[[Link](ARGUMENTS.md)\] | | 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 | Trainer APIs, 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..7c44058512a --- /dev/null +++ b/examples/llm_qat/linear_attention/README.md @@ -0,0 +1,304 @@ +# 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, state codec, decay approximation, and replay settings. The decode +implementation uses Torch operations and autograd. 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. Both complete recipes use the Torch 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 implementation uses autograd through initial states, chunk handoff, +token writes, and replay refreshes. The decode recurrence runs token by token in +Python, so long training suffixes can be slow. + +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..8a29cda8139 --- /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": "torch", + "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/linear_attention/fla_chunk_delta_h.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py index 875cb39a7d0..84da262fd9f 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -1,5 +1,5 @@ # 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 fake quantization +# 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. # @@ -42,15 +42,18 @@ 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): - """[ModelOpt] Dynamic FP8 E4M3 scale per full [K, BV] tile of one sequence and head.""" +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))) @@ -58,7 +61,18 @@ def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr): 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))) - return tl.where(b_amax > 0, b_amax / 448.0, 1.0) + 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] @@ -215,23 +229,23 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( 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 an FP8 cache is quantized before the first chunk uses it. - if STATE_QDQ == 1: + # [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) + 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) + 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) + 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) - b_h1 = fp8_scalar_qdq(b_h1, b_scale) + 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 = fp8_scalar_qdq(b_h2, b_scale) + b_h2 = _state_scalar_qdq(b_h2, b_scale, STATE_QDQ) if K > 128: - b_h3 = fp8_scalar_qdq(b_h3, b_scale) + b_h3 = _state_scalar_qdq(b_h3, b_scale, STATE_QDQ) if K > 192: - b_h4 = fp8_scalar_qdq(b_h4, b_scale) + b_h4 = _state_scalar_qdq(b_h4, b_scale, STATE_QDQ) # main recurrence for i_t in range(NT): @@ -385,24 +399,24 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h4 = tl.dot(b_k, b_v, b_h4) - # [ModelOpt] Dynamic per-tile FP8 QDQ of the next-chunk or stored final state. - # Recompute amax over [K, BV]; fp8_scalar_qdq applies that tile's scalar scale. - if STATE_QDQ == 1: + # [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) + 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) + 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) + 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) - b_h1 = fp8_scalar_qdq(b_h1, b_scale) + 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 = fp8_scalar_qdq(b_h2, b_scale) + b_h2 = _state_scalar_qdq(b_h2, b_scale, STATE_QDQ) if K > 128: - b_h3 = fp8_scalar_qdq(b_h3, b_scale) + b_h3 = _state_scalar_qdq(b_h3, b_scale, STATE_QDQ) if K > 192: - b_h4 = fp8_scalar_qdq(b_h4, b_scale) + b_h4 = _state_scalar_qdq(b_h4, b_scale, STATE_QDQ) if STORE_FINAL_STATE: if STATE_V_FIRST: @@ -885,8 +899,8 @@ def state_qdq_tile_v(V: int, state_qdq: int, state_qdq_block_v: int | None) -> i """ if state_qdq == STATE_QDQ_OFF: return 64 if check_shared_mem("ada") else 32 - if state_qdq != STATE_QDQ_FP8_DYNAMIC: - raise ValueError(f"Unsupported state_qdq mode {state_qdq}; expected 0 or 1.") + 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( 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 index a5609cab102..f34b8a3384b 100644 --- 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 @@ -59,6 +59,7 @@ 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, @@ -636,7 +637,8 @@ def chunk_gated_delta_rule( 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 with a dynamic scale per [K, state_qdq_block_v] tile of each head. + # 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 @@ -647,11 +649,13 @@ def chunk_gated_delta_rule( 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): - raise ValueError(f"`state_qdq` must be 0 or 1, got {state_qdq}.") + 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 and (not q.is_cuda or torch.cuda.get_device_capability(q.device) < (8, 9)): + 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`.") 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 6a538e16d5b..3d2d64e82e0 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.""" @@ -1575,6 +1577,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 78a6e03a16b..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,7 +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", [])) - _validate_linear_attention_quantizers(model) + _apply_linear_attention_policy(model, config) metadata = {} update_quantize_metadata(model, config, metadata) @@ -137,6 +139,7 @@ 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 @@ -144,13 +147,14 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: return model quantizer_state_dict = dict(metadata["quantizer_state"]) - # Older checkpoints predate these disabled handles; preserve their baseline path. - for name, module in _linear_attention_modules(model).items(): - for handle in ("gdn_state_quantizer", "gdn_w_quantizer"): - 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() + 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() @@ -203,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 @@ -212,18 +222,39 @@ def update_quantize_metadata( def _linear_attention_modules(model): # Optional framework plugins import conversion; defer this import to avoid that cycle. - from .plugins.gated_delta_net import GatedDeltaNetStateQuantMixin + from .plugins.linear_attention import _LinearAttentionQuantMixin return { get_unwrapped_name(name, model): module for name, module in model.named_modules() - if isinstance(module, GatedDeltaNetStateQuantMixin) + if isinstance(module, _LinearAttentionQuantMixin) } -def _validate_linear_attention_quantizers(model): - for module in _linear_attention_modules(model).values(): +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]: diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py index 477eeb89847..305bb2cd81f 100644 --- a/modelopt/torch/quantization/linear_attention/__init__.py +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -13,4 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Helpers for linear-attention quantization.""" +"""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..1f3b610fe75 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -0,0 +1,108 @@ +# 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"] = 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.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..2a981c597ec --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -0,0 +1,341 @@ +# 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, +): + """Run the Torch training implementation with explicit differentiable carry.""" + 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, + ) 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 index 4f1e18a174d..68353bd28df 100644 --- a/modelopt/torch/quantization/linear_attention/utils.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -15,18 +15,35 @@ """Shared helpers for linear-attention quantization.""" -from ..nn import TensorQuantizer +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) -> None: - """Check the dynamic E4M3 contract and the custom backward's identity STE.""" +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 == (4, 3) + 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 @@ -37,7 +54,89 @@ def validate_gdn_quantizer(quantizer: TensorQuantizer, *, name: str) -> None: and not quantizer._use_constant_amax ): raise ValueError( - f"{name} supports only dynamic E4M3 fake quantization with " + 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 328b040e18b..c488b7f8fab 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -33,7 +33,7 @@ from modelopt.torch.opt.utils import forward_with_reshard from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.conversion import ( - _validate_linear_attention_quantizers, + _apply_linear_attention_policy, preserve_quantizer_attributes_context, set_quantizer_attributes_partial, set_quantizer_by_cfg, @@ -403,7 +403,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) - _validate_linear_attention_quantizers(model) + _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) _check_plugin_quantization_took_effect(model, quantize_config) diff --git a/modelopt/torch/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py index 7234ca81f76..c072f91d954 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -39,7 +39,8 @@ from .attention import * from .custom import * -from .gated_delta_net import * +from .gdn import * +from .kda import * with import_plugin("diffusers"): from .diffusion.diffusers import * diff --git a/modelopt/torch/quantization/plugins/gated_delta_net.py b/modelopt/torch/quantization/plugins/gdn.py similarity index 70% rename from modelopt/torch/quantization/plugins/gated_delta_net.py rename to modelopt/torch/quantization/plugins/gdn.py index df70683cc2f..8d0fe616390 100644 --- a/modelopt/torch/quantization/plugins/gated_delta_net.py +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -17,9 +17,10 @@ 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, ModelOpt runs an adapted copy of the kernel +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 at the end of every chunk, with a scale computed inside the kernel from the state itself. +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, @@ -32,9 +33,8 @@ import torch -from ..config import QuantizerAttributeConfig -from ..linear_attention.utils import validate_gdn_quantizer -from ..nn import QuantModule, TensorQuantizer +from ..linear_attention.prefill import matmul_gdn +from .linear_attention import _LinearAttentionQuantMixin __all__ = ["GatedDeltaNetStateQuantMixin"] @@ -63,39 +63,22 @@ def _state_qdq_chunk_gated_delta_rule() -> GatedDeltaRuleFn: return chunk_gated_delta_rule -class GatedDeltaNetStateQuantMixin(QuantModule): +class GatedDeltaNetStateQuantMixin(_LinearAttentionQuantMixin): """Adds ``gdn_state_quantizer`` and ``gdn_w_quantizer`` to a GatedDeltaNet module. Subclasses route the module's chunked gated-delta-rule call through :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. Both sites currently require dynamic - E4M3 and identity STE. State QDQ uses fixed 64-column value tiles and 64-token chunks. + 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. """ - def _setup(self): - for name in ("gdn_state_quantizer", "gdn_w_quantizer"): - self._register_temp_attribute( - name, TensorQuantizer(QuantizerAttributeConfig(enable=False)) - ) - - def validate_linear_attention(self) -> None: - """Reject numerical settings that the fused training path cannot implement.""" - for name in ("gdn_state_quantizer", "gdn_w_quantizer"): - quantizer = getattr(self, name) - if quantizer.is_enabled: - validate_gdn_quantizer(quantizer, name=name) - # The state handle configures fused QDQ; W's grouping is executed by TensorQuantizer. - if self.gdn_state_quantizer.is_enabled and self.gdn_state_quantizer.axis != (0, 1): - raise ValueError( - "gdn_state_quantizer supports only axis=(0, 1) with 64-column value tiling" - ) - - def modelopt_post_restore(self, prefix: str = ""): - """Validate restored quantizers against the fused training kernel requirements.""" - super().modelopt_post_restore(prefix) - self.validate_linear_attention() + @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 @@ -104,7 +87,7 @@ def _state_quantized_chunk_gated_delta_rule( 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 (quantize_state or quantize_w): + 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) @@ -115,14 +98,26 @@ def _state_quantized_chunk_gated_delta_rule( "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", 64) - if chunk_size != 64: + 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=int(quantize_state), - state_qdq_block_v=64, + 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 5e422235e14..7a06b4dfe68 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -70,7 +70,7 @@ from ..utils import sync_moe_expert_amax from ..utils.layerwise_calib import LayerActivationCollector from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear -from .gated_delta_net import GatedDeltaNetStateQuantMixin +from .gdn import GatedDeltaNetStateQuantMixin from .transformer_engine import _QuantTEGroupedLinear, _QuantTELayerNormLinear, _QuantTELinear try: @@ -1159,9 +1159,7 @@ def forward_pre_attn_and_core_attn(self, *args, **kwargs): def validate_linear_attention(self): super().validate_linear_attention() - if self.config.context_parallel_size > 1 and ( - self.gdn_state_quantizer.is_enabled or self.gdn_w_quantizer.is_enabled - ): + if self.config.context_parallel_size > 1 and self.linear_attention_is_enabled: raise NotImplementedError("GDN QAT does not support Megatron context parallelism.") 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 f44fbc620ca..6fb77c2163b 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -34,6 +34,7 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `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 | | `indexer_k_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer key cache (`*indexer_k_quantizer`), global scale fixed to 1; Blackwell+ GPUs | | `indexer_q_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer query (`*indexer_q_quantizer`), global scale fixed to 1; Blackwell+ GPUs | diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index cac4f6d80aa..e409bff7aa6 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -21,6 +21,10 @@ enable: false - quantizer_name: '*gdn_w_quantizer' enable: false + - quantizer_name: '*kda_state_quantizer' + enable: false + - quantizer_name: '*kda_w_quantizer' + enable: false - quantizer_name: '*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 index ed5768ee599..4d4b13278ba 100644 --- a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -14,10 +14,11 @@ # limitations under the License. # Dynamic per-tile FP8 E4M3 fake quantization of the GatedDeltaNet recurrent state. -# Recompute amax over each [K, 64] tile inside the kernel, independently per sequence and head +# 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; -# the kernel uses one scale per full-K by 64-value-column tile (two per head for V = 128). -# This is the only configuration the module supports. Applied at every kernel chunk boundary +# 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``. 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/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 4c3751fa394..189d773299f 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -64,6 +64,7 @@ supported combinations. | `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | | `iq2_s` | IQ2_S W2A16 (2.56 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | | `q8_0` | Q8_0 W8A16 (8.5 bpw), eligible linears | 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/tests/_test_utils/torch/quantization/linear_attention_reference.py b/tests/_test_utils/torch/quantization/linear_attention_reference.py index 65ea183a318..57d7448bcb1 100644 --- a/tests/_test_utils/torch/quantization/linear_attention_reference.py +++ b/tests/_test_utils/torch/quantization/linear_attention_reference.py @@ -25,23 +25,43 @@ import torch -__all__ = ["chunk_gdn_reference", "recurrent_delta_rule_reference", "state_fp8_qdq_reference"] +__all__ = [ + "chunk_gdn_reference", + "chunk_kda_reference", + "recurrent_delta_rule_reference", + "state_fp8_qdq_reference", + "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. The zero tile uses scale one; partial value tiles are valid. + 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) - scale = torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)) - rounded.append((tile / scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() * scale) + 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() @@ -140,34 +160,61 @@ def chunk_gdn_reference( 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 = g[b, lo:hi].transpose(0, 1).cumsum(-1) + 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 = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() - gram = kc @ kc.transpose(-1, -2) - lower = (bc * gram * decay).tril(-1) + 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) - rhs = torch.cat((bc * vc, bc * kc * gc.exp().unsqueeze(-1)), 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) + 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) @@ -179,21 +226,27 @@ def chunk_gdn_reference( if n != previous_n: state = states[n] if state_qdq: - state = state_fp8_qdq_reference(state, state_qdq_block_v) + 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 = u - wc @ state # state_read - local_scores = ((qc * scale) @ kc.transpose(-1, -2) * decay).tril() - output = (qc * (scale * gc.exp()).unsqueeze(-1)) @ state - output = output + local_scores @ updated_values # local_readout + 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 * (gc[..., -1:] - gc).exp().unsqueeze(-1) - state = state * gc[..., -1].exp()[:, None, None] - state = state + weighted_keys.transpose(-1, -2) @ updated_values # state_update + 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_fp8_qdq_reference(state, state_qdq_block_v) + state = state_qdq_reference(state, state_qdq_block_v, state_format) if len(finals) <= n: finals.append(state) else: @@ -201,3 +254,93 @@ def chunk_gdn_reference( output = torch.cat(outputs).reshape(*v.shape) final = torch.stack(finals) return output, final.transpose(-1, -2) if state_v_first else final + + +def chunk_kda_reference( + q, + k, + v, + g, + beta, + *, + chunk_size=64, + scale=None, + initial_state=None, + cu_seqlens=None, + state_v_first=False, + state_qdq=False, + state_qdq_block_v=64, + 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_fla_chunk_gated_delta_rule.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py index defe05aa5d5..53776affca1 100644 --- 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 @@ -61,15 +61,15 @@ def compare(actual, expected, tolerance): ) -@pytest.fixture(scope="module", params=["disabled", "w", "state-w"]) +@pytest.fixture(scope="module", params=["disabled", "w", "state-w", "state-int8"]) def compiled_gdn_case(request): """Compile only the selected BF16 forward/backward path, outside the test-call timer.""" - state_qdq = request.param == "state-w" - if state_qdq and torch.cuda.get_device_capability() < (8, 9): + state_qdq = {"state-w": 1, "state-int8": 2}.get(request.param, 0) + if state_qdq == 1 and torch.cuda.get_device_capability() < (8, 9): pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") quantizer = ( TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), axis=(0, 1, 2), type="dynamic")) - if request.param != "disabled" + if request.param in ("w", "state-w") else None ) args, state = make_inputs() @@ -83,7 +83,13 @@ def test_gdn_forward_and_backward(compiled_gdn_case): args, state, kwargs = compiled_gdn_case reference_args = [x.detach().float().requires_grad_() for x in args] reference_state = state.detach().clone().requires_grad_() - expected = values_and_grads(chunk_gdn_reference, reference_args, reference_state, **kwargs) + expected = values_and_grads( + chunk_gdn_reference, + reference_args, + reference_state, + state_format="int8" if kwargs["state_qdq"] == 2 else "fp8_e4m3", + **kwargs, + ) actual = values_and_grads( chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs ) 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..ed67787de54 --- /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": "torch", + } + 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": "torch", + "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 cc3262e1917..1376b8eedd7 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 index 0756f2d0d40..ea5f96865eb 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import nullcontext + import pytest import torch from _test_utils.torch.megatron.models import get_mcore_gpt_model @@ -23,11 +25,16 @@ ) 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.gated_delta_net import _state_qdq_chunk_gated_delta_rule +from modelopt.torch.quantization.plugins.gdn import _state_qdq_chunk_gated_delta_rule from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet try: @@ -80,15 +87,47 @@ def _gdn_config(sites): for site in sites ], "algorithm": None, + "linear_attention": [ + {"module_name": "decoder.layers.*.self_attention", "cfg": {"state": {"block_v": 16}}} + ], } -def _test_gdn_qat_helper(rank, size, checkpoint_path, sites): +def _config_for_mode(mode): + # Ampere exercises W QDQ; native FP8 devices also exercise chunk state QDQ. + sites = ("state", "w") if torch.cuda.get_device_capability() >= (8, 9) else ("w",) + if mode == "decode_replay_int8": + sites = ("state",) + cfg = _gdn_config(sites) + if mode == "decode_replay_int8": + cfg["quant_cfg"][1]["cfg"].update(num_bits=8, unsigned=False, narrow_range=True) + cfg["linear_attention"][0]["cfg"].update( + backend="matmul", decode={"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 = get_forward(model) + forward = _gdn_forward(model) + outputs = [] handles = [ module.register_forward_hook( @@ -102,12 +141,18 @@ def _test_gdn_qat_helper(rank, size, checkpoint_path, sites): gdn_ref = outputs.copy() outputs.clear() - mtq.quantize(model, _gdn_config(sites)) + 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) @@ -120,20 +165,34 @@ def _test_gdn_qat_helper(rank, size, checkpoint_path, sites): for handle in handles: handle.remove() - for site in sites: - mtq.disable_quantizer(model, f"*gdn_{site}_quantizer") + 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 site in sites: - mtq.enable_quantizer(model, f"*gdn_{site}_quantizer") + 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 in restored_gdn: + 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() @@ -144,15 +203,18 @@ def _test_gdn_qat_helper(rank, size, checkpoint_path, sites): assert not torch.equal(restored_gdn[0].in_proj.weight, before) -def _compile_gdn_qat_kernels(rank, size, sites): +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 = get_forward(model) + forward = _gdn_forward(model) with torch.no_grad(): forward(model) - mtq.quantize(model, _gdn_config(sites)) + 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() @@ -160,16 +222,14 @@ def _compile_gdn_qat_kernels(rank, size, sites): torch.cuda.synchronize() -@pytest.fixture -def compiled_gdn_workers(dist_workers_size_1): +@pytest.fixture(params=["chunk", "decode_replay_int8"]) +def compiled_gdn_workers(request, dist_workers_size_1): """Warm one small QAT model in the same worker, outside the test-call budget.""" - # Ampere exercises W QDQ; native FP8 devices also exercise state QDQ. - sites = ("state", "w") if torch.cuda.get_device_capability() >= (8, 9) else ("w",) - dist_workers_size_1.run(_compile_gdn_qat_kernels, sites) - return dist_workers_size_1, sites + dist_workers_size_1.run(_compile_gdn_qat_kernels, request.param) + return dist_workers_size_1, request.param def test_gdn_qat_and_sharded_restore(compiled_gdn_workers, tmp_path): """Train through QDQ after a Megatron distributed-checkpoint round trip.""" - workers, sites = compiled_gdn_workers - workers.run(_test_gdn_qat_helper, tmp_path, sites) + workers, mode = compiled_gdn_workers + workers.run(_test_gdn_qat_helper, mode, tmp_path) diff --git a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py b/tests/unit/torch/quantization/plugins/test_gdn.py similarity index 74% rename from tests/unit/torch/quantization/plugins/test_gated_delta_net.py rename to tests/unit/torch/quantization/plugins/test_gdn.py index 3793edcb189..aed7f8caf42 100644 --- a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -22,9 +22,19 @@ 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 gated_delta_net -from modelopt.torch.quantization.plugins.gated_delta_net import GatedDeltaNetStateQuantMixin +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"} @@ -36,9 +46,7 @@ def chunk_gated_delta_rule(q, k, v, g, beta, **kwargs): @pytest.fixture(autouse=True) def mock_fla_kernel(monkeypatch): - monkeypatch.setattr( - gated_delta_net, "_fla_chunk_gated_delta_rule", lambda: chunk_gated_delta_rule - ) + monkeypatch.setattr(gdn, "_fla_chunk_gated_delta_rule", lambda: chunk_gated_delta_rule) class TinyGatedDeltaNet(nn.Module): @@ -106,7 +114,12 @@ 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"): + for name in ( + "gdn_state_quantizer", + "gdn_w_quantizer", + "linear_attention_config", + "_linear_attention_prefill_lengths", + ): assert not hasattr(model, name) @@ -124,27 +137,29 @@ def test_disabled_state_quantizer_calls_original_kernel(): 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, partial_kernel): +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( - gated_delta_net, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel - ) + 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) - mtq.quantize(model, quant_cfg(), lambda m: m(x)) + 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": 1, + "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 {}), @@ -167,9 +182,7 @@ def fake_state_qdq_kernel(*args, **kwargs): calls.append(kwargs) return chunk_gated_delta_rule(*args) - monkeypatch.setattr( - gated_delta_net, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel - ) + 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)) @@ -214,7 +227,13 @@ 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 @@ -222,6 +241,8 @@ def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): 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) @@ -241,7 +262,12 @@ def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): 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 @@ -256,7 +282,9 @@ def test_restore_legacy_gdn_without_new_quantizer_handles(): 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() @@ -274,3 +302,37 @@ def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): ) 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"} + ) From c396e5c07e6e932038a4748bdb5ba3ea364accbb Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 2 Oct 2026 19:55:34 -0700 Subject: [PATCH 02/27] Use Megatron training adapters for GDN and KDA Signed-off-by: Kai Xu --- CHANGELOG.rst | 2 +- examples/llm_qat/linear_attention/README.md | 98 ++++--- ...state_int8.json => decode_state_int8.json} | 13 + .../llm_qat/linear_attention/requirements.txt | 2 - examples/llm_qat/linear_attention/train.py | 135 +++++---- .../quantization/linear_attention/kda.py | 2 +- .../torch/quantization/plugins/__init__.py | 1 - modelopt/torch/quantization/plugins/kda.py | 80 ++---- .../torch/quantization/plugins/megatron.py | 110 +++++--- .../linear_attention_reference.py | 151 +--------- .../torch/quantization/plugins/test_kda.py | 189 ------------- .../quantization/test_tensor_quant_cuda.py | 22 +- .../quantization/plugins/test_megatron_kda.py | 167 +++++++++++ .../torch/quantization/plugins/test_gdn.py | 39 +-- .../test_linear_attention_decode.py | 266 ++++-------------- .../test_linear_attention_hadamard.py | 160 ++--------- .../test_linear_attention_int8.py | 99 ------- 17 files changed, 505 insertions(+), 1031 deletions(-) rename examples/llm_qat/linear_attention/configs/{kda_decode_state_int8.json => decode_state_int8.json} (71%) delete mode 100644 tests/gpu/torch/quantization/plugins/test_kda.py create mode 100644 tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py delete mode 100644 tests/unit/torch/quantization/test_linear_attention_int8.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index c7bcc2789c6..c65b55dfd96 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,7 +16,7 @@ Changelog - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. -- 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 experimental Megatron-Core 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. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before a separate KV-cache AutoQuantize stage, with independent resumable checkpoints for the weight and diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md index 7c44058512a..b4e950dde88 100644 --- a/examples/llm_qat/linear_attention/README.md +++ b/examples/llm_qat/linear_attention/README.md @@ -1,39 +1,62 @@ # 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 +This example fine-tunes Megatron-Core GDN or KDA attention parameters with +recurrent-state fake quantization, then saves a Megatron checkpoint with ModelOpt +state. Training adapters support Megatron `GatedDeltaNet` and +`KimiDeltaAttention`; FLA supplies kernels, not model-layer adapters. +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. +Use the [Megatron Bridge environment](../../megatron_bridge/README.md#pre-requisites) +and install the example dependencies. Run commands from the repository root. + +KDA requires a Megatron revision exporting +`megatron.core.ssm.gated_delta_net.KimiDeltaAttention`. It is available on +[Megatron's dev branch](https://github.com/NVIDIA/Megatron-LM/tree/ac100f773f9d8fac918ce3e64698be4ed84fa428); +Megatron-Core 0.19.2 does not include it. Older versions can still use GDN. +Retain `fla-core==0.5.1` for the kernel dependency. ModelOpt registers no FLA +model-layer adapters; Megatron or Bridge may still install the full +`flash-linear-attention` package as their dependency. Use mutually compatible +Bridge/Core revisions: Bridge's pinned Core revision may differ from the KDA +`dev` revision. Native Megatron KDA training does not require Bridge. ```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 \ +torchrun --standalone --nproc-per-node=1 examples/llm_qat/linear_attention/train.py \ + --model /path/to/local-model \ --train-data /path/to/train.parquet \ - --output /path/to/qat-checkpoint \ + --output /path/to/megatron-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. +Use a local HF model/tokenizer snapshot that your installed Megatron Bridge can +convert to a Megatron model containing GDN or KDA layers, a Parquet file with a +`text` column, and one CUDA GPU. Bridge's architecture/checkpoint conversion +support is a separate requirement from ModelOpt's layer adapter; an arbitrary +FLA model checkpoint cannot be loaded through this example. For a model already +built in Megatron, use the training-loop integration below. If the local model +requires custom Python code, explicitly pass `--trust-remote-code`. + +The shared [INT8 configuration](configs/decode_state_int8.json) enables GDN and +KDA state quantizers. It 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. Only linear-attention parameters are trained, in FP32 +under BF16 autocast; other parameters are frozen in BF16. Megatron receives +shifted next-token labels; loss is averaged over suffix positions. + +The output is a Megatron model checkpoint containing ModelOpt quantizers and +execution policies, without optimizer state. Reload it through +[the Megatron Bridge workflow](../../megatron_bridge/README.md), preserving +ModelOpt state. Export to a supported serving model separately. The current vLLM +state-only plugin accepts its own boundary-QDQ recipe; it does not implement +the Hadamard or replay training policies. + +This one-GPU example demonstrates the integration. Distributed training needs a +Megatron training schedule, gradient synchronization, and a compatible exporter. +Context parallelism is currently unsupported by these ModelOpt adapters. +This example does not measure model-quality recovery or performance. ## Enable state quantization @@ -67,9 +90,9 @@ must be supplied again for each workload. ### Load a state recipe -State quantizers start disabled. The state recipe enables `*kda_state_quantizer` +State quantizers start disabled. The shared state recipe enables `*gdn_state_quantizer` and `*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 +`*gdn_state_quantizer` or `*kda_state_quantizer` to select just one layer type. 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. @@ -82,7 +105,7 @@ import json from pathlib import Path recipe = json.loads( - Path("examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json").read_text() + Path("examples/llm_qat/linear_attention/configs/decode_state_int8.json").read_text() ) policy = recipe["linear_attention"][0]["cfg"] decode = policy["decode"] @@ -212,9 +235,9 @@ It requires `prefill_state_qdq=False`; use the tile codec to quantize prefix sta 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 +### Megatron GDN with chunk-only FLA kernels -For GDN's existing chunked FLA path, enable the GDN state quantizer and omit the +For Megatron GDN's existing chunked FLA kernel path, enable the GDN state quantizer and omit the `linear_attention` execution-policy entry: ```python @@ -255,16 +278,23 @@ 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. +# ids and shifted next-token labels have shape [2, 128]. +# model is a Megatron model. +positions = torch.arange(128, device=ids.device).expand_as(ids) 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 + losses = model( + input_ids=ids, position_ids=positions, attention_mask=None, labels=labels + ) + mask = torch.arange(128, device=ids.device)[None, :] >= torch.tensor( + [64, 96], device=ids.device + )[:, None] + loss = (losses * mask).sum() / mask.sum() 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 +The shared recipe selects both GDN and KDA state quantizers. State quantization is dynamic, so these recipes use `algorithm=None` without a calibration pass. Replay factors have their own FP8 toggle. @@ -298,7 +328,7 @@ token writes, and replay refreshes. The decode recurrence runs token by token in Python, so long training suffixes can be slow. 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 +has no prefill GEMM QDQ or approximate inverse. Use the Megatron training forward +without a serving inference context. 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/decode_state_int8.json similarity index 71% rename from examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json rename to examples/llm_qat/linear_attention/configs/decode_state_int8.json index 8a29cda8139..aad890b22b0 100644 --- a/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json +++ b/examples/llm_qat/linear_attention/configs/decode_state_int8.json @@ -4,6 +4,19 @@ "quantizer_name": "*", "enable": false }, + { + "quantizer_name": "*gdn_state_quantizer", + "cfg": { + "num_bits": 8, + "type": "dynamic", + "axis": [ + 0, + 1 + ], + "unsigned": false, + "narrow_range": true + } + }, { "quantizer_name": "*kda_state_quantizer", "cfg": { diff --git a/examples/llm_qat/linear_attention/requirements.txt b/examples/llm_qat/linear_attention/requirements.txt index 29636d78e52..9be29cb9a2a 100644 --- a/examples/llm_qat/linear_attention/requirements.txt +++ b/examples/llm_qat/linear_attention/requirements.txt @@ -1,4 +1,2 @@ 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 index 7662c1ca593..19fbe41f52f 100644 --- a/examples/llm_qat/linear_attention/train.py +++ b/examples/llm_qat/linear_attention/train.py @@ -13,24 +13,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Fine-tune KDA attention parameters with recurrent-state fake quantization.""" +"""Fine-tune Megatron GDN/KDA attention with recurrent-state fake quantization.""" import argparse import json +import os from pathlib import Path import pyarrow.parquet as pq import torch -from fla.layers.kda import KimiDeltaAttention -from transformers import AutoModelForCausalLM, AutoTokenizer +from megatron.core import parallel_state +from megatron.core.ssm import gated_delta_net -import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.quantization.linear_attention import linear_attention_training_phase +from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_hf def main(): - """Quantize, train on local text, and save a ModelOpt-aware Hugging Face checkpoint.""" + """Load a local model through Megatron Bridge and save a Megatron QAT checkpoint.""" parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--model", type=Path, required=True) parser.add_argument( @@ -39,7 +40,7 @@ def main(): parser.add_argument( "--quant-config", type=Path, - default=Path(__file__).with_name("configs") / "kda_decode_state_int8.json", + default=Path(__file__).with_name("configs") / "decode_state_int8.json", ) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--trust-remote-code", action="store_true") @@ -53,56 +54,78 @@ def main(): 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 int(os.environ.get("WORLD_SIZE", "1")) != 1: + parser.error( + "This example uses one GPU; use a Megatron training schedule for multiple GPUs" + ) + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", "0"))) + torch.distributed.init_process_group("nccl") + try: + bridge, _, models, model, tokenizer = load_mbridge_model_from_hf( + hf_model_name_or_path=str(options.model), + trust_remote_code=options.trust_remote_code, + provider_overrides={ + "tensor_model_parallel_size": 1, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + "expert_model_parallel_size": 1, + "expert_tensor_parallel_size": 1, + "sequence_parallel": False, + "seq_length": options.length, + "gradient_accumulation_fusion": False, + "bf16": False, + "params_dtype": torch.float32, + "pipeline_dtype": torch.float32, + }, + ) + torch.manual_seed(options.seed) + layer_types = (gated_delta_net.GatedDeltaNet,) + if hasattr(gated_delta_net, "KimiDeltaAttention"): + layer_types += (gated_delta_net.KimiDeltaAttention,) + layers = [module for module in model.modules() if isinstance(module, layer_types)] + if not layers: + raise ValueError("The model must contain Megatron GatedDeltaNet or KimiDeltaAttention") + # 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) + mtq.quantize(model, json.loads(options.quant_config.read_text())) + 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 + 1) + 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 + 1) + 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, labels = block[:-1].unsqueeze(0).cuda(), block[1:].unsqueeze(0).cuda() + positions = torch.arange(options.length, device=ids.device).unsqueeze(0) + optimizer.zero_grad(set_to_none=True) + # Megatron takes shifted labels and returns per-token losses, unlike HF models. + with linear_attention_training_phase(model, [options.prefill_tokens]): + with torch.autocast("cuda", dtype=torch.bfloat16): + losses = model( + input_ids=ids, position_ids=positions, attention_mask=None, labels=labels + ) + loss = losses[:, options.prefill_tokens :].float().mean() + 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) + bridge.save_megatron_model( + models, + str(options.output), + hf_tokenizer_path=str(options.model), + hf_tokenizer_kwargs={"trust_remote_code": options.trust_remote_code}, + ) + finally: + parallel_state.destroy_model_parallel() + torch.distributed.destroy_process_group() if __name__ == "__main__": diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py index 446dd6f7d0d..fbde42dbc29 100644 --- a/modelopt/torch/quantization/linear_attention/kda.py +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -55,7 +55,7 @@ def matmul_kda( ): """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. + Gates follow FLA's kernel 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") diff --git a/modelopt/torch/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py index c072f91d954..fc56e230838 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -40,7 +40,6 @@ 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/kda.py b/modelopt/torch/quantization/plugins/kda.py index fbc4b1cf0ef..92f6e88047c 100644 --- a/modelopt/torch/quantization/plugins/kda.py +++ b/modelopt/torch/quantization/plugins/kda.py @@ -13,62 +13,40 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""QAT integration for the optional FLA KimiDeltaAttention layer.""" +"""Kernel routing for Megatron Kimi Delta Attention quantization.""" -import importlib.metadata -import inspect -from functools import lru_cache -from types import FunctionType +from functools import partial from ..linear_attention.kda import matmul_kda -from ..nn import QuantModuleRegistry -from .custom import CUSTOM_MODEL_PLUGINS from .linear_attention import _LinearAttentionQuantMixin -__all__ = [] +__all__ = ["KimiDeltaAttentionStateQuantMixin"] -@lru_cache(maxsize=8) -def _forward_signature(function): - return inspect.signature(function) +class KimiDeltaAttentionStateQuantMixin(_LinearAttentionQuantMixin): + """Adds state quantizers and decode-aware kernel routing to Kimi Delta Attention.""" - -class _QuantKimiDeltaAttention(_LinearAttentionQuantMixin): linear_attention_quantizer_names = ("kda_state_quantizer", "kda_w_quantizer") def validate_linear_attention(self): + """Require the materialized backend for KDA numerical emulation.""" 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): + def _state_quantized_chunk_kda(self, kernel, *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 kernel(*args, **kwargs) + # FLA kernels are an optional dependency; the training layer belongs to Megatron. + from fla.ops.kda import chunk_kda + + while isinstance(kernel, partial): + args = (*kernel.args, *args) + kwargs = {**kernel.keywords, **kwargs} + kernel = kernel.func + if kernel is not chunk_kda: + raise NotImplementedError("KDA quantization requires FLA's chunk_kda callable") return matmul_kda( *args, policy=self.linear_attention_config, @@ -77,27 +55,3 @@ def _quantized_chunk(self, *args, **kwargs): 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/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 7a06b4dfe68..c36b08944a4 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -71,6 +71,8 @@ from ..utils.layerwise_calib import LayerActivationCollector from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear from .gdn import GatedDeltaNetStateQuantMixin +from .kda import KimiDeltaAttentionStateQuantMixin +from .linear_attention import _LinearAttentionQuantMixin from .transformer_engine import _QuantTEGroupedLinear, _QuantTELayerNormLinear, _QuantTELinear try: @@ -87,6 +89,13 @@ except ImportError: HAS_GDN = False +try: + from megatron.core.ssm.gated_delta_net import KimiDeltaAttention + + HAS_KDA = True +except ImportError: + HAS_KDA = False + __all__ = [] @@ -1106,16 +1115,61 @@ def set_extra_state(self, state): quant_module_set_extra_state(self, state) +class _MegatronLinearAttentionMixin(_LinearAttentionQuantMixin): + """Wrap the kernel call in both direct and recomputed Megatron forwards.""" + + 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_linear_attention_kernel(self): + kernel = self.gated_delta_rule + self.gated_delta_rule = partial(self._linear_attention_kernel, kernel) + try: + yield + finally: + self.gated_delta_rule = kernel + + def forward(self, *args, **kwargs): + # Newer Megatron versions recompute the core independently during backward. + if hasattr(super(), "_forward_compute") or hasattr( + super(), "forward_pre_attn_and_core_attn" + ): + return super().forward(*args, **kwargs) + with self._quantized_linear_attention_kernel(): + return super().forward(*args, **kwargs) + + def _forward_compute(self, *args, **kwargs): + with self._quantized_linear_attention_kernel(): + return super()._forward_compute(*args, **kwargs) + + def forward_pre_attn_and_core_attn(self, *args, **kwargs): + with self._quantized_linear_attention_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/KDA QAT does not support Megatron context parallelism.") + + if HAS_GDN: @QuantModuleRegistry.register({GatedDeltaNet: "megatron_GatedDeltaNet"}) - class _QuantGatedDeltaNet(GatedDeltaNetStateQuantMixin): - """GatedDeltaNet with fake quantization of the recurrent state at kernel chunk boundaries. + class _QuantGatedDeltaNet(_MegatronLinearAttentionMixin, GatedDeltaNetStateQuantMixin): + """Megatron GDN with chunk or decode-aware state fake quantization.""" - 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. - """ + _linear_attention_kernel = ( + GatedDeltaNetStateQuantMixin._state_quantized_chunk_gated_delta_rule + ) # Class-level overrides so torch routes the quantizer state through ``_extra_state`` # (GatedDeltaNet has none); see _QuantDSAttention. @@ -1125,42 +1179,16 @@ def get_extra_state(self): def set_extra_state(self, state): quant_module_set_extra_state(self, state) - 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.") +if HAS_KDA: + + @QuantModuleRegistry.register({KimiDeltaAttention: "megatron_KimiDeltaAttention"}) + class _QuantKimiDeltaAttention( + _MegatronLinearAttentionMixin, KimiDeltaAttentionStateQuantMixin + ): + """Megatron KDA with decode-aware state fake quantization.""" + + _linear_attention_kernel = KimiDeltaAttentionStateQuantMixin._state_quantized_chunk_kda def _is_supported_megatron_model(model: torch.nn.Module) -> bool: diff --git a/tests/_test_utils/torch/quantization/linear_attention_reference.py b/tests/_test_utils/torch/quantization/linear_attention_reference.py index 57d7448bcb1..ad43a2226e1 100644 --- a/tests/_test_utils/torch/quantization/linear_attention_reference.py +++ b/tests/_test_utils/torch/quantization/linear_attention_reference.py @@ -27,7 +27,6 @@ __all__ = [ "chunk_gdn_reference", - "chunk_kda_reference", "recurrent_delta_rule_reference", "state_fp8_qdq_reference", "state_qdq_reference", @@ -162,59 +161,33 @@ def chunk_gdn_reference( 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)) + 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 = 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) + decay = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() + gram = kc @ kc.transpose(-1, -2) + lower = (bc * gram * 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) + rhs = torch.cat((bc * vc, bc * kc * gc.exp().unsqueeze(-1)), 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) chunks.append((n, qc, kc, gc, decay, u)) all_w.append(w.transpose(0, 1)) w = torch.cat(all_w).reshape(q.shape) @@ -231,20 +204,14 @@ def chunk_gdn_reference( 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)) + updated_values = u - wc @ state # state_read + local_scores = ((qc * scale) @ kc.transpose(-1, -2) * decay).tril() + output = (qc * (scale * gc.exp()).unsqueeze(-1)) @ state + output = output + local_scores @ updated_values # local_readout 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) - ) + 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 # state_update if state_qdq: state = state_qdq_reference(state, state_qdq_block_v, state_format) if len(finals) <= n: @@ -254,93 +221,3 @@ def chunk_gdn_reference( output = torch.cat(outputs).reshape(*v.shape) final = torch.stack(finals) return output, final.transpose(-1, -2) if state_v_first else final - - -def chunk_kda_reference( - q, - k, - v, - g, - beta, - *, - chunk_size=64, - scale=None, - initial_state=None, - cu_seqlens=None, - state_v_first=False, - state_qdq=False, - state_qdq_block_v=64, - 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/quantization/plugins/test_kda.py b/tests/gpu/torch/quantization/plugins/test_kda.py deleted file mode 100644 index ed67787de54..00000000000 --- a/tests/gpu/torch/quantization/plugins/test_kda.py +++ /dev/null @@ -1,189 +0,0 @@ -# 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": "torch", - } - 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": "torch", - "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 1376b8eedd7..bf156882b15 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -120,24 +120,16 @@ 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): + @pytest.fixture + def fp8_extension(self): + return get_cuda_ext_fp8() + + def test_eager_matches_cuda_scale_rounding(self, fp8_extension): # 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)] - ) + actual = fp8_extension.fake_e4m3fy_with_axis(values, amax, 0) + expected = tensor_quant.fp8_eager(values, amax[:, None]) torch.testing.assert_close(actual, expected, rtol=0, atol=0) @pytest.mark.parametrize("device", ["cuda", "cpu"]) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py new file mode 100644 index 00000000000..b93f3e81967 --- /dev/null +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py @@ -0,0 +1,167 @@ +# 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 +from contextlib import nullcontext + +import pytest +import torch +from _test_utils.torch.megatron.utils import ( + initialize_for_megatron, + load_distributed_checkpoint, + save_distributed_checkpoint, +) +from megatron.core import parallel_state +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer import TransformerConfig + +import modelopt.torch.quantization as mtq +from modelopt.torch.opt.plugins.mcore_dist_checkpointing import ( + restore_sharded_modelopt_state, + save_sharded_modelopt_state, +) +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + linear_attention_training_phase, +) + +KimiDeltaAttention = pytest.importorskip("megatron.core.ssm.gated_delta_net.kda").KimiDeltaAttention + + +def _layer(): + config = TransformerConfig( + num_layers=1, + hidden_size=64, + num_attention_heads=2, + linear_num_key_heads=2, + linear_num_value_heads=2, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_conv_kernel_dim=4, + experimental_attention_variant="kda", + is_hybrid_model=True, + normalization="RMSNorm", + activation_func=torch.nn.functional.silu, + params_dtype=torch.float32, + gradient_accumulation_fusion=False, + recompute_granularity="selective", + recompute_modules=["gdn"], + ) + spec = hybrid_stack_spec.submodules.kda_layer.submodules.self_attention + model = ( + KimiDeltaAttention( + config, + submodules=spec.submodules, + layer_number=1, + pg_collection=ProcessGroupCollection( + tp=parallel_state.get_tensor_model_parallel_group(), + cp=parallel_state.get_context_parallel_group(), + ), + ) + .cuda() + .train() + ) + with torch.no_grad(): + model.A_log.fill_(-4) + model.dt_bias.zero_() + return model + + +def _case(cfg): + initialize_for_megatron(tensor_model_parallel_size=1, pipeline_model_parallel_size=1, seed=73) + model = _layer() + hidden = torch.randn(73, 1, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + def forward(layer): + policy = getattr(layer, "linear_attention_config", None) + phase = ( + linear_attention_training_phase(layer, [31]) + if policy is not None and policy.decode is not None + else nullcontext() + ) + with phase, torch.autocast("cuda", dtype=torch.bfloat16): + return layer(hidden, attention_mask=None)[0] + + with torch.no_grad(): + baseline = forward(model) + mtq.quantize(model, cfg) + return model, hidden, forward, baseline + + +def _compile_kda(rank, size, cfg): + model, hidden, forward, _ = _case(cfg) + with linear_attention_training_phase(model, [31]): + forward(model).float().square().mean().backward() + torch.cuda.synchronize() + + +def _test_kda(rank, size, cfg, checkpoint_path): + model, hidden, forward, baseline = _case(cfg) + assert model.kda_state_quantizer.is_enabled + assert model.kda_state_quantizer.num_bits == 8 + assert model.linear_attention_config.decode.state_codec == "int8_hadamard32" + kernel = model.gated_delta_rule + with torch.no_grad(): + assert not torch.equal(forward(model), baseline) + assert model.gated_delta_rule is kernel + + policy = model.linear_attention_config + model.kda_state_quantizer.disable() + model.linear_attention_config = LinearAttentionConfig() + with torch.no_grad(): + torch.testing.assert_close(forward(model), baseline, rtol=0, atol=0) + model.kda_state_quantizer.enable() + model.linear_attention_config = policy + + with torch.no_grad(): + expected = forward(model) + save_distributed_checkpoint(checkpoint_path, model) + save_sharded_modelopt_state([model], checkpoint_path) + restored = _layer() + restore_sharded_modelopt_state([restored], checkpoint_path) + load_distributed_checkpoint(checkpoint_path, restored) + # Megatron recomputes the core during backward, after forward's kernel wrapper exits. + with linear_attention_training_phase(restored, [31]): + actual = forward(restored) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + actual.float().square().mean().backward() + assert restored._linear_attention_prefill_lengths is None + assert restored.linear_attention_config == policy + assert restored.gated_delta_rule is kernel + assert torch.isfinite(hidden.grad).all() + for parameter in restored.parameters(): + assert parameter.grad is not None + assert torch.isfinite(parameter.grad).all() + before = restored.in_proj.weight.detach().clone() + torch.optim.SGD(restored.parameters(), lr=0.1).step() + assert not torch.equal(before, restored.in_proj.weight) + + +@pytest.fixture(scope="module") +def compiled_kda_workers(dist_workers_size_1, project_root_path): + """Warm one KDA shape outside the functional test timer.""" + path = project_root_path / "examples/llm_qat/linear_attention/configs/decode_state_int8.json" + cfg = json.loads(path.read_text()) + cfg["linear_attention"][0]["cfg"]["decode"].update( + mode="replay", replay={"window": 5}, decay_log_step=1 / 256 + ) + dist_workers_size_1.run(_compile_kda, cfg) + return dist_workers_size_1, cfg + + +def test_kda_qat_and_sharded_restore(compiled_kda_workers, tmp_path): + workers, cfg = compiled_kda_workers + workers.run(_test_kda, cfg, tmp_path) diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index aed7f8caf42..3f24f2cd64e 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -22,20 +22,11 @@ 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"} @@ -137,9 +128,8 @@ def test_disabled_state_quantizer_calls_original_kernel(): 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): +def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, partial_kernel): calls = [] def fake_state_qdq_kernel(*args, **kwargs): @@ -151,15 +141,12 @@ def fake_state_qdq_kernel(*args, **kwargs): 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)) + mtq.quantize(model, quant_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": 1, "state_qdq_block_v": 64, "w_quantizer": None, **({"output_final_state": True} if partial_kernel else {}), @@ -304,26 +291,6 @@ def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): 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( diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py index e0ced168024..b9e81c3359b 100644 --- a/tests/unit/torch/quantization/test_linear_attention_decode.py +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -13,8 +13,6 @@ # 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 @@ -23,50 +21,40 @@ from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.linear_attention import ( LinearAttentionConfig, + LinearAttentionDecodeConfig, matmul_gdn, matmul_kda, + recurrent_decode_reference, ) -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.linear_attention.decode import _encode 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): +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_state_qdq_matches_tensor_quantizer(state_format): 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))) + # Exact INT8 scales exercise signed endpoints, half ties, and a zero tile. + value = torch.zeros(2, 3, 5, 19) + value[..., 0, :9] = torch.tensor([-127, -126.5, -1.5, -0.5, 0, 0.5, 1.5, 126.5, 127]) * 2 + value[0, :, 0, 16:] = torch.tensor([-254, 1, 254]) value.requires_grad_() - quantizer = TensorQuantizer( - QuantizerAttributeConfig( - num_bits=(4, 3), type="dynamic", axis=(0, 1), pass_through_bwd=True - ) - ) + cfg = {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1)} + if state_format == "int8": + cfg.update(num_bits=8, unsigned=False, narrow_range=True) + quantizer = TensorQuantizer(QuantizerAttributeConfig(**cfg)) expected = torch.cat( - [ - quantizer(tile.flatten(-2) if state else tile).reshape_as(tile) - for tile in value.split(16, -1) - ], - -1, + [quantizer(tile.flatten(-2)).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)) + encoded = _encode(value, True, 16, state=True, state_format=state_format) + torch.testing.assert_close(encoded.values, expected, rtol=0, atol=0) 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) + (gradient,) = torch.autograd.grad((encoded.values * probe).sum(), value) + torch.testing.assert_close(gradient, probe, rtol=0, atol=0) + assert encoded.scales.shape == (2, 3, 2) assert not encoded.scales.requires_grad -def _inputs(kda=True, length=73): +def _inputs(kda=True, length=11): torch.manual_seed(193) q, k = [ F.normalize(torch.randn(length, 2, 16, dtype=torch.float64), dim=-1).requires_grad_() @@ -86,92 +74,49 @@ def _values_and_grads(output, state, args, initial): 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): +@pytest.mark.parametrize(("kda", "mode"), [(False, "token"), (True, "replay")]) +def test_prefill_decode_matches_recurrence_and_gradients(kda, mode): args, state = _inputs(kda) - cfg = LinearAttentionDecodeConfig( - mode=mode, replay={"window": 8, "factor_qdq": False} if mode == "replay" else None + policy = LinearAttentionConfig( + backend="matmul", + decode={ + "mode": mode, + "replay": {"window": 3, "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( + function = matmul_kda if kda else matmul_gdn + output, final = function( + *(x.unsqueeze(0) for x in args), + policy=policy, + prefill_lengths=[3], + initial_state=state.unsqueeze(0), + output_final_state=True, + ) + exact, expected_state = 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), + for actual, expected in zip( + _values_and_grads(output[0], final[0], args, state), + _values_and_grads(exact[0], expected_state[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) + torch.testing.assert_close(actual, expected, rtol=2e-10, atol=1e-11) -@pytest.mark.parametrize("mode", ["token", "replay", "reencode"]) -@pytest.mark.parametrize("readout", ["working", "stored"]) -def test_quantized_carry_continuation_metadata_and_gradients(mode, readout): +def test_quantized_replay_carry_continuation_and_gradients(): 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), + cfg = LinearAttentionDecodeConfig(mode="replay", replay={"window": 3}) + kwargs = {"config": cfg, "state_qdq": True, "block_v": 16} + expected, final = recurrent_decode_reference(*args, initial_state=state, **kwargs) + first, carry = recurrent_decode_reference(*(x[:5] for x in args), initial_state=state, **kwargs) + second, carry = recurrent_decode_reference(*(x[5:] for x in args), carry=carry, **kwargs) + for actual, expected in zip( + _values_and_grads(torch.cat((first, second)), 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 + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert carry.position == 11 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 + assert carry.cursor == 2 def test_grid_gate_ste_keeps_gate_gradients_and_changes_trajectory(): @@ -187,110 +132,3 @@ def test_grid_gate_ste_keeps_gate_gradients_and_changes_trajectory(): *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 index 5f6b40f4c5c..c9020e4a67e 100644 --- a/tests/unit/torch/quantization/test_linear_attention_hadamard.py +++ b/tests/unit/torch/quantization/test_linear_attention_hadamard.py @@ -13,22 +13,13 @@ # 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): @@ -47,30 +38,28 @@ def _qdq(value): return (value - value.detach()) + rounded -def _inputs(kda, length=11): +def _inputs(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 + g = -torch.rand(k.shape, 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): +def _oracle(args, initial): q, k, v, g, beta = args state = _rotate(initial) - if quantize: - state = _qdq(state) + 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) + state = _qdq(working) if (t + 1) % 3 == 0 else working + outputs.append(_rotate(torch.einsum("hk,hkv->hv", q[t], state)) / q.shape[-1] ** 0.5) return torch.stack(outputs), _rotate(state) @@ -87,141 +76,28 @@ def _check_with_grads(actual, expected, args, initial, tolerance=2e-10): 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) +def test_hadamard_replay_matches_dense_oracle_and_split_carry(): + args, initial = _inputs() cfg = LinearAttentionDecodeConfig( - mode=mode, - readout=readout, + mode="replay", state_codec="int8_hadamard32", - replay={"window": 3, "factor_qdq": False} if mode == "replay" else None, + replay={"window": 3, "factor_qdq": False}, ) 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 - ) + _check_with_grads((output, carry.reconstruct()), _oracle(args, initial), 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 + first, split_carry = recurrent_decode_reference( + *(x[:2] for x in args), initial_state=initial, **kwargs + ) + second, split_carry = recurrent_decode_reference( + *(x[2:] for x in args), carry=split_carry, **kwargs + ) _check_with_grads( - (torch.cat(pieces), split_carry.reconstruct()), + (torch.cat((first, second)), 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 deleted file mode 100644 index b3de8c94d47..00000000000 --- a/tests/unit/torch/quantization/test_linear_attention_int8.py +++ /dev/null @@ -1,99 +0,0 @@ -# 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"} - ) From 4e583c1b60ad86b6e0834db5dc60ca2b662c743d Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 2 Oct 2026 23:53:24 -0700 Subject: [PATCH 03/27] Use Megatron Bridge workflows for linear attention QAT and QAD Signed-off-by: Kai Xu --- examples/llm_qat/linear_attention/README.md | 56 ++-- .../llm_qat/linear_attention/requirements.txt | 1 - examples/llm_qat/linear_attention/train.py | 263 +++++++++++++----- .../linear_attention/fla_chunk_delta_h.py | 3 + .../megatron_bridge/test_linear_attention.py | 188 +++++++++++++ 5 files changed, 416 insertions(+), 95 deletions(-) create mode 100644 tests/examples/megatron_bridge/test_linear_attention.py diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md index b4e950dde88..8eeb6b21706 100644 --- a/examples/llm_qat/linear_attention/README.md +++ b/examples/llm_qat/linear_attention/README.md @@ -1,9 +1,10 @@ -# Quantization-Aware Training for Linear Attention +# Quantization-Aware Training and Distillation for Linear Attention This example fine-tunes Megatron-Core GDN or KDA attention parameters with -recurrent-state fake quantization, then saves a Megatron checkpoint with ModelOpt -state. Training adapters support Megatron `GatedDeltaNet` and -`KimiDeltaAttention`; FLA supplies kernels, not model-layer adapters. +recurrent-state fake quantization using Megatron Bridge's training loop (QAT), or +its distillation loop with an unquantized teacher (QAD). Bridge handles optimization, +gradient accumulation, logging, and checkpoints. Training adapters support Megatron +`GatedDeltaNet` and `KimiDeltaAttention`; FLA supplies kernels, not model-layer adapters. Runtime support includes token writes, ReplaySSM, KDA decay approximation, and FP8 or INT8 state QDQ. The INT8 recipes enable Hadamard rotation by default. @@ -26,15 +27,24 @@ Bridge/Core revisions: Bridge's pinned Core revision may differ from the KDA pip install -r examples/llm_qat/linear_attention/requirements.txt torchrun --standalone --nproc-per-node=1 examples/llm_qat/linear_attention/train.py \ --model /path/to/local-model \ - --train-data /path/to/train.parquet \ + --train-data /path/to/tokenized/train_text_document \ --output /path/to/megatron-qat-checkpoint \ --train-steps 1 --length 128 --prefill-tokens 64 ``` +Add `--teacher-model /path/to/unquantized-model` to run QAD. The teacher and +student must share the tokenizer vocabulary and output vocabulary dimensions. +The teacher stays frozen and unquantized. QAT uses next-token cross-entropy; +QAD uses Bridge's logits distillation loss with the language-model loss disabled. +Both losses are masked to the suffix after `--prefill-tokens`. + Use a local HF model/tokenizer snapshot that your installed Megatron Bridge can -convert to a Megatron model containing GDN or KDA layers, a Parquet file with a -`text` column, and one CUDA GPU. Bridge's architecture/checkpoint conversion -support is a separate requirement from ModelOpt's layer adapter; an arbitrary +convert to a Megatron model containing GDN or KDA layers, a tokenized Megatron +`.bin`/`.idx` dataset pair, and one CUDA GPU. `--train-data` is the shared filename +prefix without either extension; use the model's tokenizer when +[preparing the data](../../megatron_bridge/README.md#data-preparation). +Bridge's architecture/checkpoint conversion support is a separate requirement +from ModelOpt's layer adapter; an arbitrary FLA model checkpoint cannot be loaded through this example. For a model already built in Megatron, use the training-loop integration below. If the local model requires custom Python code, explicitly pass `--trust-remote-code`. @@ -42,21 +52,27 @@ requires custom Python code, explicitly pass `--trust-remote-code`. The shared [INT8 configuration](configs/decode_state_int8.json) enables GDN and KDA state quantizers. It 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. Only linear-attention parameters are trained, in FP32 -under BF16 autocast; other parameters are frozen in BF16. Megatron receives -shifted next-token labels; loss is averaged over suffix positions. - -The output is a Megatron model checkpoint containing ModelOpt quantizers and -execution policies, without optimizer state. Reload it through -[the Megatron Bridge workflow](../../megatron_bridge/README.md), preserving -ModelOpt state. Export to a supported serving model separately. The current vLLM +must be divisible by 32. Only linear-attention parameters are trained, using +Bridge's BF16 mixed precision and optimizer; other model parameters are frozen. +Every dense sequence has the same fixed prefill boundary. The example retains +the phase context throughout training so backward recomputation sees that boundary. +Packed sequences and variable per-sequence boundaries require their own batch integration. + +Bridge saves model weights, optimizer/scheduler state, and ModelOpt state under +`/checkpoints`. Reusing `--output` resumes from its latest checkpoint; +increase `--train-steps` to the desired total step count and retain the same model, +teacher, dataset, and prefill boundary. A resumed checkpoint supplies the saved +quantization policy. `--global-batch-size` controls accumulation with microbatch size one. +Export to a supported serving model separately. The current vLLM state-only plugin accepts its own boundary-QDQ recipe; it does not implement the Hadamard or replay training policies. -This one-GPU example demonstrates the integration. Distributed training needs a -Megatron training schedule, gradient synchronization, and a compatible exporter. -Context parallelism is currently unsupported by these ModelOpt adapters. -This example does not measure model-quality recovery or performance. +This one-GPU example demonstrates integration with Bridge's existing training +workflows; distributed execution needs separate qualification. Context parallelism +is currently unsupported by these ModelOpt adapters. The Hugging Face +`QATTrainer`/`QADTrainer` example in [the parent directory](../README.md#using-qattrainer-and-qadtrainer) +targets a different training framework. This example does not measure model-quality +recovery or performance. ## Enable state quantization diff --git a/examples/llm_qat/linear_attention/requirements.txt b/examples/llm_qat/linear_attention/requirements.txt index 9be29cb9a2a..915edf246cd 100644 --- a/examples/llm_qat/linear_attention/requirements.txt +++ b/examples/llm_qat/linear_attention/requirements.txt @@ -1,2 +1 @@ fla-core==0.5.1 -pyarrow>=15 diff --git a/examples/llm_qat/linear_attention/train.py b/examples/llm_qat/linear_attention/train.py index 19fbe41f52f..28558a4a0e1 100644 --- a/examples/llm_qat/linear_attention/train.py +++ b/examples/llm_qat/linear_attention/train.py @@ -13,29 +13,135 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Fine-tune Megatron GDN/KDA attention with recurrent-state fake quantization.""" +"""Run Megatron Bridge QAT or QAD with GDN/KDA recurrent-state fake quantization.""" import argparse import json import os +from contextlib import ExitStack from pathlib import Path -import pyarrow.parquet as pq -import torch -from megatron.core import parallel_state +from megatron.bridge import AutoBridge +from megatron.bridge.models.distillation_provider import convert_to_distillation_provider +from megatron.bridge.training.config import ( + CheckpointConfig, + ConfigContainer, + DistributedDataParallelConfig, + GPTDatasetConfig, + LoggerConfig, + OptimizerConfig, + RNGConfig, + SchedulerConfig, + TokenizerConfig, + TrainingConfig, + ValidationConfig, +) +from megatron.bridge.training.distill import distill +from megatron.bridge.training.gpt_step import forward_step_modelopt +from megatron.bridge.training.post_training.checkpointing import ( + has_modelopt_state, + load_modelopt_state, +) +from megatron.bridge.training.post_training.distillation import ModelOptDistillConfig +from megatron.bridge.training.pretrain import pretrain +from megatron.bridge.training.state import GlobalState +from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size +from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding from megatron.core.ssm import gated_delta_net +from megatron.core.utils import unwrap_model +from transformers import AutoTokenizer import modelopt.torch.quantization as mtq from modelopt.torch.quantization.linear_attention import linear_attention_training_phase -from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_hf +from modelopt.torch.quantization.utils import is_quantized + + +def run_training(config, quant_config, prefill_tokens, teacher_provider=None): + """Delegate optimization and checkpointing to Bridge, with a fixed phase per sequence.""" + if not 0 <= prefill_tokens < config.dataset.seq_length: + raise ValueError("prefill-tokens must leave at least one suffix label") + resume = config.checkpoint.load and has_modelopt_state(config.checkpoint.load) + with ExitStack() as phases: + + def prepare_student(models): + # Restore the saved policy before quantization and before the QAD conversion hook. + if resume: + load_modelopt_state(models, config.checkpoint.load) + student = unwrap_model(models[0]) + layer_types = (gated_delta_net.GatedDeltaNet,) + if hasattr(gated_delta_net, "KimiDeltaAttention"): + layer_types += (gated_delta_net.KimiDeltaAttention,) + layers = [m for m in student.modules() if isinstance(m, layer_types)] + if not layers: + raise ValueError( + "The model must contain Megatron GatedDeltaNet or KimiDeltaAttention" + ) + student.requires_grad_(False) + for layer in layers: + layer.requires_grad_(True) + if config.model.recompute_granularity == "full": + + def require_input_grad(module, args, output): + return output.requires_grad_(True) if module.training else output + + # Reentrant checkpointing needs a differentiable input even with frozen embeddings. + for module in student.modules(): + if isinstance(module, LanguageModelEmbedding): + handle = module.register_forward_hook(require_input_grad) + phases.callback(handle.remove) + if not is_quantized(student): + mtq.quantize(student, quant_config) + # Every dense microbatch uses the same boundary; retain it through recomputed backward. + phases.enter_context( + linear_attention_training_phase( + student, [prefill_tokens] * config.train.micro_batch_size + ) + ) + return models + + def forward_step(state: GlobalState, data_iterator, model, return_schedule_plan=False): + batch = dict(next(data_iterator)) + batch["loss_mask"] = batch["loss_mask"].clone() + batch["loss_mask"][..., :prefill_tokens] = 0 + # Bridge applies this same mask to the language-model and distillation losses. + return forward_step_modelopt(state, iter([batch]), model, return_schedule_plan) + + config.model.register_pre_wrap_hook(prepare_student) + if teacher_provider is not None: + config.model = convert_to_distillation_provider( + config.model, teacher_provider, ModelOptDistillConfig(skip_lm_loss=True) + ) + distill(config, forward_step) + else: + pretrain(config, forward_step) + + +def model_provider(path, options, *, load_weights=True): + """Build a single-GPU Megatron provider using the model's Bridge converter.""" + bridge = AutoBridge.from_hf_pretrained(str(path), trust_remote_code=options.trust_remote_code) + provider = bridge.to_megatron_provider(load_weights=load_weights) + provider.tensor_model_parallel_size = 1 + provider.pipeline_model_parallel_size = 1 + provider.context_parallel_size = 1 + provider.expert_model_parallel_size = 1 + provider.expert_tensor_parallel_size = 1 + provider.sequence_parallel = False + provider.gradient_accumulation_fusion = False + provider.calculate_per_token_loss = True + provider.seq_length = options.length + return provider def main(): - """Load a local model through Megatron Bridge and save a Megatron QAT checkpoint.""" + """Train a local model with QAT, or add a frozen teacher for QAD.""" parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--teacher-model", type=Path, help="Unquantized teacher; enables QAD") parser.add_argument( - "--train-data", type=Path, required=True, help="Parquet file with a text column" + "--train-data", + type=Path, + required=True, + help="Megatron tokenized dataset prefix (.bin/.idx)", ) parser.add_argument( "--quant-config", @@ -49,83 +155,92 @@ def main(): 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) + parser.add_argument("--global-batch-size", type=int, default=1) 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") + if options.global_batch_size < 1: + parser.error("global-batch-size must be positive") if int(os.environ.get("WORLD_SIZE", "1")) != 1: parser.error( "This example uses one GPU; use a Megatron training schedule for multiple GPUs" ) - torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", "0"))) - torch.distributed.init_process_group("nccl") - try: - bridge, _, models, model, tokenizer = load_mbridge_model_from_hf( - hf_model_name_or_path=str(options.model), - trust_remote_code=options.trust_remote_code, - provider_overrides={ - "tensor_model_parallel_size": 1, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 1, - "expert_model_parallel_size": 1, - "expert_tensor_parallel_size": 1, - "sequence_parallel": False, - "seq_length": options.length, - "gradient_accumulation_fusion": False, - "bf16": False, - "params_dtype": torch.float32, - "pipeline_dtype": torch.float32, - }, - ) - torch.manual_seed(options.seed) - layer_types = (gated_delta_net.GatedDeltaNet,) - if hasattr(gated_delta_net, "KimiDeltaAttention"): - layer_types += (gated_delta_net.KimiDeltaAttention,) - layers = [module for module in model.modules() if isinstance(module, layer_types)] - if not layers: - raise ValueError("The model must contain Megatron GatedDeltaNet or KimiDeltaAttention") - # 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) - mtq.quantize(model, json.loads(options.quant_config.read_text())) - 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 + 1) - 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 + 1) - 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, labels = block[:-1].unsqueeze(0).cuda(), block[1:].unsqueeze(0).cuda() - positions = torch.arange(options.length, device=ids.device).unsqueeze(0) - optimizer.zero_grad(set_to_none=True) - # Megatron takes shifted labels and returns per-token losses, unlike HF models. - with linear_attention_training_phase(model, [options.prefill_tokens]): - with torch.autocast("cuda", dtype=torch.bfloat16): - losses = model( - input_ids=ids, position_ids=positions, attention_mask=None, labels=labels - ) - loss = losses[:, options.prefill_tokens :].float().mean() - 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) - bridge.save_megatron_model( - models, - str(options.output), - hf_tokenizer_path=str(options.model), + if any(not Path(f"{options.train_data}.{suffix}").is_file() for suffix in ("bin", "idx")): + parser.error("train-data must name an existing Megatron .bin/.idx dataset prefix") + + checkpoint_dir = str(options.output / "checkpoints") + resume = has_modelopt_state(checkpoint_dir) + student = model_provider(options.model, options, load_weights=not resume) + teacher = None + if options.teacher_model is not None: + tokenizer_kwargs = {"trust_remote_code": options.trust_remote_code} + student_vocab = AutoTokenizer.from_pretrained(options.model, **tokenizer_kwargs).get_vocab() + teacher_vocab = AutoTokenizer.from_pretrained( + options.teacher_model, **tokenizer_kwargs + ).get_vocab() + if student_vocab != teacher_vocab: + parser.error("QAD requires student and teacher to use the same tokenizer vocabulary") + teacher = model_provider(options.teacher_model, options) + vocab_sizes = [ + calculate_padded_vocab_size( + p.vocab_size, p.make_vocab_size_divisible_by, p.tensor_model_parallel_size + ) + for p in (student, teacher) + ] + if vocab_sizes[0] != vocab_sizes[1]: + parser.error("QAD requires matching student and teacher output vocabulary dimensions") + + config = ConfigContainer( + model=student, + train=TrainingConfig( + train_iters=options.train_steps, + global_batch_size=options.global_batch_size, + micro_batch_size=1, + ), + validation=ValidationConfig(eval_iters=0, eval_interval=options.train_steps), + optimizer=OptimizerConfig( + optimizer="adam", lr=options.learning_rate, min_lr=0, weight_decay=0, clip_grad=1.0 + ), + scheduler=SchedulerConfig( + lr_decay_style="constant", + lr_warmup_iters=0, + start_weight_decay=0, + end_weight_decay=0, + use_checkpoint_opt_param_scheduler=True, + ), + ddp=DistributedDataParallelConfig(average_in_collective=False), + dataset=GPTDatasetConfig( + seq_length=options.length, + blend=([str(options.train_data)], None), + split="100,0,0", + random_seed=options.seed, + reset_position_ids=False, + reset_attention_mask=False, + eod_mask_loss=False, + dataloader_type="single", + num_workers=0, + ), + tokenizer=TokenizerConfig( + tokenizer_type="HuggingFaceTokenizer", + tokenizer_model=str(options.model), hf_tokenizer_kwargs={"trust_remote_code": options.trust_remote_code}, - ) - finally: - parallel_state.destroy_model_parallel() - torch.distributed.destroy_process_group() + ), + checkpoint=CheckpointConfig( + save=checkpoint_dir, + load=checkpoint_dir, + save_interval=options.train_steps, + async_save=False, + ckpt_format="torch_dist", + ), + logger=LoggerConfig(log_interval=1), + rng=RNGConfig(seed=options.seed), + mixed_precision="bf16_mixed", + ) + run_training( + config, json.loads(options.quant_config.read_text()), options.prefill_tokens, teacher + ) if __name__ == "__main__": 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 index 84da262fd9f..e316fbf5150 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -54,6 +54,7 @@ @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.""" + # The b_h blocks partition K; their joint amax gives one scale for the full [K, BV] tile. b_amax = tl.max(tl.abs(b_h1)) if K > 64: b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h2))) @@ -61,6 +62,8 @@ def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr, STATE_QDQ: tl.cons 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))) + # STATE_QDQ=2 uses symmetric INT8 (max 127); STATE_QDQ=1 uses FP8 E4M3 (max 448). + # Use scale=1 for zero tiles to avoid division by zero during QDQ. if STATE_QDQ == 2: return tl.where(b_amax > 0, b_amax * (1.0 / 127.0), 1.0) else: diff --git a/tests/examples/megatron_bridge/test_linear_attention.py b/tests/examples/megatron_bridge/test_linear_attention.py new file mode 100644 index 00000000000..6ccbf60bf02 --- /dev/null +++ b/tests/examples/megatron_bridge/test_linear_attention.py @@ -0,0 +1,188 @@ +# 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 runpy + +import pytest +import torch +from _test_utils.examples.megatron_example_runner import reset_megatron_global_state +from _test_utils.examples.run_command import MODELOPT_ROOT +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.training.config import ( + CheckpointConfig, + ConfigContainer, + DistributedDataParallelConfig, + LoggerConfig, + MockGPTDatasetConfig, + OptimizerConfig, + RNGConfig, + SchedulerConfig, + TokenizerConfig, + TrainingConfig, + ValidationConfig, +) +from megatron.bridge.training.post_training.checkpointing import has_modelopt_state +from megatron.core.utils import unwrap_model + +from modelopt.torch.quantization.utils import is_quantized + +run_training = runpy.run_path(str(MODELOPT_ROOT / "examples/llm_qat/linear_attention/train.py"))[ + "run_training" +] + + +def _train(tmp_path, qad, recipe, train_steps=1): + captured = {} + + def provider(name): + model = HybridModelProvider( + num_layers=1, + hidden_size=64, + ffn_hidden_size=128, + num_attention_heads=2, + hybrid_layer_pattern="G", + vocab_size=128, + seq_length=16, + linear_num_key_heads=2, + linear_num_value_heads=2, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_conv_kernel_dim=4, + linear_attention_freq=1, + experimental_attention_variant="gated_delta_net", + is_hybrid_model=True, + activation_func=torch.nn.functional.silu, + calculate_per_token_loss=True, + gradient_accumulation_fusion=False, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, + cross_entropy_loss_fusion=False, + ) + + def capture(models): + module = unwrap_model(models[0]) + captured[name] = module + captured[name + "_before"] = { + key: value.detach().clone() for key, value in module.named_parameters() + } + if name == "student": + + def check_mask(module, args, kwargs): + mask = kwargs["loss_mask"] + assert not mask[:, :8].any() + assert mask[:, 8:].all() + + module.register_forward_pre_hook(check_mask, with_kwargs=True) + return models + + model.register_post_wrap_hook(capture) + return model + + checkpoint = str(tmp_path / "checkpoints") + config = ConfigContainer( + model=provider("student"), + train=TrainingConfig(train_iters=train_steps, global_batch_size=2, micro_batch_size=1), + validation=ValidationConfig(eval_iters=0, eval_interval=1), + optimizer=OptimizerConfig(optimizer="adam", lr=1e-2, min_lr=0, weight_decay=0), + scheduler=SchedulerConfig( + lr_decay_style="constant", + lr_warmup_iters=0, + start_weight_decay=0, + end_weight_decay=0, + use_checkpoint_opt_param_scheduler=True, + ), + ddp=DistributedDataParallelConfig(average_in_collective=False), + dataset=MockGPTDatasetConfig( + seq_length=16, + random_seed=123, + reset_position_ids=False, + reset_attention_mask=False, + eod_mask_loss=False, + dataloader_type="single", + num_workers=0, + ), + tokenizer=TokenizerConfig(tokenizer_type="NullTokenizer", vocab_size=128), + checkpoint=CheckpointConfig( + save=checkpoint, load=checkpoint, save_interval=1, async_save=False + ), + logger=LoggerConfig(log_interval=1), + rng=RNGConfig(seed=123), + mixed_precision="bf16_mixed", + ) + run_training(config, recipe, 8, provider("teacher") if qad else None) + return captured, checkpoint + + +@pytest.fixture(scope="module") +def compiled_state_training(tmp_path_factory, project_root_path): + """Compile one small GDN shape before timing the QAT/QAD integration tests.""" + recipe = json.loads( + ( + project_root_path / "examples/llm_qat/linear_attention/configs/decode_state_int8.json" + ).read_text() + ) + try: + _train(tmp_path_factory.mktemp("compile_state_training"), True, recipe) + finally: + reset_megatron_global_state() + return recipe + + +@pytest.mark.parametrize("qad", [False, True], ids=["qat", "qad"]) +def test_state_training(compiled_state_training, tmp_path, qad): + captured, checkpoint = _train(tmp_path, qad, compiled_state_training) + student = captured["student"] + assert has_modelopt_state(checkpoint) + assert is_quantized(student) + layers = [m for m in student.modules() if hasattr(m, "gdn_state_quantizer")] + assert len(layers) == 1 + assert layers[0].gdn_state_quantizer.is_enabled + assert layers[0].linear_attention_config.decode.state_codec == "int8_hadamard32" + changed = [] + for name, parameter in student.named_parameters(): + assert torch.isfinite(parameter).all() + if not torch.equal(parameter, captured["student_before"][name]): + assert parameter.requires_grad + changed.append(name) + assert changed + assert all( + module._linear_attention_prefill_lengths is None + for module in student.modules() + if hasattr(module, "_linear_attention_prefill_lengths") + ) + if qad: + teacher = captured["teacher"] + assert not is_quantized(teacher) + for name, parameter in teacher.named_parameters(): + assert not parameter.requires_grad + assert torch.equal(parameter, captured["teacher_before"][name]) + before_resume = { + name: parameter.detach().clone() + for name, parameter in student.named_parameters() + if parameter.requires_grad + } + reset_megatron_global_state() + # An empty recipe verifies that resume restores the saved quantization policy. + resumed, _ = _train(tmp_path, True, {}, train_steps=2) + assert has_modelopt_state(str(tmp_path / "checkpoints/iter_0000002")) + assert is_quantized(resumed["student"]) + assert not is_quantized(resumed["teacher"]) + assert any( + not torch.equal(parameter, before_resume[name]) + for name, parameter in resumed["student"].named_parameters() + if name in before_resume + ) From 82c95285d662e89eeaa144ae9b7d88435323d33c Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 3 Oct 2026 10:47:23 -0700 Subject: [PATCH 04/27] Enable distributed Megatron Bridge state QAT and QAD Signed-off-by: Kai Xu --- examples/llm_qat/linear_attention/README.md | 37 +++++-- examples/llm_qat/linear_attention/train.py | 96 +++++++++++-------- .../megatron_bridge/test_linear_attention.py | 46 +++++++-- 3 files changed, 121 insertions(+), 58 deletions(-) diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md index 8eeb6b21706..c06f20fc6a5 100644 --- a/examples/llm_qat/linear_attention/README.md +++ b/examples/llm_qat/linear_attention/README.md @@ -38,9 +38,32 @@ The teacher stays frozen and unquantized. QAT uses next-token cross-entropy; QAD uses Bridge's logits distillation loss with the language-model loss disabled. Both losses are masked to the suffix after `--prefill-tokens`. +For multiple GPUs, use the same `--tp_size`, `--pp_size`, and `--ep_size` options +as [the Megatron Bridge example](../../megatron_bridge/distill.py). Bridge builds +the process groups, schedules forward/backward, and manages the distributed optimizer. +For example, run QAD with two-way tensor parallelism: + +```bash +torchrun --standalone --nproc-per-node=2 examples/llm_qat/linear_attention/train.py \ + --model /path/to/local-model \ + --teacher-model /path/to/unquantized-model \ + --train-data /path/to/tokenized/train_text_document \ + --output /path/to/megatron-qad-checkpoint \ + --tp_size 2 --pp_size 1 --ep_size 1 \ + --global-batch-size 2 --train-steps 1 --length 128 --prefill-tokens 64 +``` + +Student and teacher use the same topology. Sequence parallelism is enabled with +TP greater than one. With TP/PP/EP all set to one, additional ranks use data +parallelism; the model weights remain replicated on each GPU. The global batch +size must be divisible by the data-parallel size (microbatch size is one). +Choose a topology supported by the model's Bridge/Core implementation. This example +currently requires linear-attention layers on every pipeline stage. Context +parallelism stays fixed at one; GDN/KDA state QAT does not yet support CP. + Use a local HF model/tokenizer snapshot that your installed Megatron Bridge can convert to a Megatron model containing GDN or KDA layers, a tokenized Megatron -`.bin`/`.idx` dataset pair, and one CUDA GPU. `--train-data` is the shared filename +`.bin`/`.idx` dataset pair, and CUDA GPUs. `--train-data` is the shared filename prefix without either extension; use the model's tokenizer when [preparing the data](../../megatron_bridge/README.md#data-preparation). Bridge's architecture/checkpoint conversion support is a separate requirement @@ -54,22 +77,22 @@ KDA state quantizers. It leaves the 64-token prefix state unquantized and applie INT8 QDQ in a 32-value Hadamard basis during the decode suffix. Value dimensions must be divisible by 32. Only linear-attention parameters are trained, using Bridge's BF16 mixed precision and optimizer; other model parameters are frozen. -Every dense sequence has the same fixed prefill boundary. The example retains -the phase context throughout training so backward recomputation sees that boundary. +Every dense sequence has the same fixed prefill boundary. The example attaches +the phase context to local linear-attention layers on each rank and retains it +throughout training so backward recomputation sees that boundary. Packed sequences and variable per-sequence boundaries require their own batch integration. Bridge saves model weights, optimizer/scheduler state, and ModelOpt state under `/checkpoints`. Reusing `--output` resumes from its latest checkpoint; increase `--train-steps` to the desired total step count and retain the same model, -teacher, dataset, and prefill boundary. A resumed checkpoint supplies the saved +teacher, topology, dataset, and prefill boundary. A resumed checkpoint supplies the saved quantization policy. `--global-batch-size` controls accumulation with microbatch size one. Export to a supported serving model separately. The current vLLM state-only plugin accepts its own boundary-QDQ recipe; it does not implement the Hadamard or replay training policies. -This one-GPU example demonstrates integration with Bridge's existing training -workflows; distributed execution needs separate qualification. Context parallelism -is currently unsupported by these ModelOpt adapters. The Hugging Face +This example uses Bridge's existing training workflows with state quantization. +The Hugging Face `QATTrainer`/`QADTrainer` example in [the parent directory](../README.md#using-qattrainer-and-qadtrainer) targets a different training framework. This example does not measure model-quality recovery or performance. diff --git a/examples/llm_qat/linear_attention/train.py b/examples/llm_qat/linear_attention/train.py index 28558a4a0e1..5fd681ed656 100644 --- a/examples/llm_qat/linear_attention/train.py +++ b/examples/llm_qat/linear_attention/train.py @@ -17,10 +17,10 @@ import argparse import json -import os from contextlib import ExitStack from pathlib import Path +import torch from megatron.bridge import AutoBridge from megatron.bridge.models.distillation_provider import convert_to_distillation_provider from megatron.bridge.training.config import ( @@ -67,44 +67,49 @@ def prepare_student(models): # Restore the saved policy before quantization and before the QAD conversion hook. if resume: load_modelopt_state(models, config.checkpoint.load) - student = unwrap_model(models[0]) layer_types = (gated_delta_net.GatedDeltaNet,) if hasattr(gated_delta_net, "KimiDeltaAttention"): layer_types += (gated_delta_net.KimiDeltaAttention,) - layers = [m for m in student.modules() if isinstance(m, layer_types)] - if not layers: - raise ValueError( - "The model must contain Megatron GatedDeltaNet or KimiDeltaAttention" - ) - student.requires_grad_(False) - for layer in layers: - layer.requires_grad_(True) - if config.model.recompute_granularity == "full": - - def require_input_grad(module, args, output): - return output.requires_grad_(True) if module.training else output - - # Reentrant checkpointing needs a differentiable input even with frozen embeddings. - for module in student.modules(): - if isinstance(module, LanguageModelEmbedding): - handle = module.register_forward_hook(require_input_grad) - phases.callback(handle.remove) - if not is_quantized(student): - mtq.quantize(student, quant_config) - # Every dense microbatch uses the same boundary; retain it through recomputed backward. - phases.enter_context( - linear_attention_training_phase( - student, [prefill_tokens] * config.train.micro_batch_size + for student in unwrap_model(models): + layers = [m for m in student.modules() if isinstance(m, layer_types)] + if not layers: + raise ValueError( + "Each local model chunk must contain Megatron GatedDeltaNet or " + "KimiDeltaAttention; choose a pipeline layout with linear attention on every stage" + ) + student.requires_grad_(False) + for layer in layers: + layer.requires_grad_(True) + if config.model.recompute_granularity == "full": + + def require_input_grad(module, args, output): + return output.requires_grad_(True) if module.training else output + + # Reentrant checkpointing needs a differentiable input with frozen embeddings. + for module in student.modules(): + if isinstance(module, LanguageModelEmbedding): + handle = module.register_forward_hook(require_input_grad) + phases.callback(handle.remove) + if not is_quantized(student): + mtq.quantize(student, quant_config) + phases.enter_context( + linear_attention_training_phase( + student, [prefill_tokens] * config.train.micro_batch_size + ) ) - ) return models - def forward_step(state: GlobalState, data_iterator, model, return_schedule_plan=False): + def masked_batch(data_iterator): batch = dict(next(data_iterator)) batch["loss_mask"] = batch["loss_mask"].clone() batch["loss_mask"][..., :prefill_tokens] = 0 - # Bridge applies this same mask to the language-model and distillation losses. - return forward_step_modelopt(state, iter([batch]), model, return_schedule_plan) + yield batch + + def forward_step(state: GlobalState, data_iterator, model, return_schedule_plan=False): + # Consume lazily: middle pipeline stages may not need any batch tensors. + return forward_step_modelopt( + state, masked_batch(data_iterator), model, return_schedule_plan + ) config.model.register_pre_wrap_hook(prepare_student) if teacher_provider is not None: @@ -117,15 +122,16 @@ def forward_step(state: GlobalState, data_iterator, model, return_schedule_plan= def model_provider(path, options, *, load_weights=True): - """Build a single-GPU Megatron provider using the model's Bridge converter.""" + """Build a Megatron provider with the same topology options as the Bridge example.""" bridge = AutoBridge.from_hf_pretrained(str(path), trust_remote_code=options.trust_remote_code) provider = bridge.to_megatron_provider(load_weights=load_weights) - provider.tensor_model_parallel_size = 1 - provider.pipeline_model_parallel_size = 1 + provider.tensor_model_parallel_size = options.tp_size + provider.pipeline_model_parallel_size = options.pp_size + provider.pipeline_dtype = torch.bfloat16 provider.context_parallel_size = 1 - provider.expert_model_parallel_size = 1 + provider.expert_model_parallel_size = options.ep_size provider.expert_tensor_parallel_size = 1 - provider.sequence_parallel = False + provider.sequence_parallel = options.tp_size > 1 provider.gradient_accumulation_fusion = False provider.calculate_per_token_loss = True provider.seq_length = options.length @@ -156,6 +162,9 @@ def main(): parser.add_argument("--seed", type=int, default=2026) parser.add_argument("--learning-rate", type=float, default=1e-5) parser.add_argument("--global-batch-size", type=int, default=1) + parser.add_argument("--tp_size", type=int, default=1, help="Tensor parallel size") + parser.add_argument("--pp_size", type=int, default=1, help="Pipeline parallel size") + parser.add_argument("--ep_size", type=int, default=1, help="Expert parallel size") 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") @@ -163,10 +172,8 @@ def main(): parser.error("prefill-tokens must leave at least one suffix label") if options.global_batch_size < 1: parser.error("global-batch-size must be positive") - if int(os.environ.get("WORLD_SIZE", "1")) != 1: - parser.error( - "This example uses one GPU; use a Megatron training schedule for multiple GPUs" - ) + if min(options.tp_size, options.pp_size, options.ep_size) < 1: + parser.error("Parallel sizes must be positive") if any(not Path(f"{options.train_data}.{suffix}").is_file() for suffix in ("bin", "idx")): parser.error("train-data must name an existing Megatron .bin/.idx dataset prefix") @@ -201,7 +208,12 @@ def main(): ), validation=ValidationConfig(eval_iters=0, eval_interval=options.train_steps), optimizer=OptimizerConfig( - optimizer="adam", lr=options.learning_rate, min_lr=0, weight_decay=0, clip_grad=1.0 + optimizer="adam", + lr=options.learning_rate, + min_lr=0, + weight_decay=0, + clip_grad=1.0, + use_distributed_optimizer=True, ), scheduler=SchedulerConfig( lr_decay_style="constant", @@ -210,7 +222,9 @@ def main(): end_weight_decay=0, use_checkpoint_opt_param_scheduler=True, ), - ddp=DistributedDataParallelConfig(average_in_collective=False), + ddp=DistributedDataParallelConfig( + average_in_collective=False, use_distributed_optimizer=True + ), dataset=GPTDatasetConfig( seq_length=options.length, blend=([str(options.train_data)], None), diff --git a/tests/examples/megatron_bridge/test_linear_attention.py b/tests/examples/megatron_bridge/test_linear_attention.py index 6ccbf60bf02..c156b57a53c 100644 --- a/tests/examples/megatron_bridge/test_linear_attention.py +++ b/tests/examples/megatron_bridge/test_linear_attention.py @@ -20,6 +20,7 @@ import torch from _test_utils.examples.megatron_example_runner import reset_megatron_global_state from _test_utils.examples.run_command import MODELOPT_ROOT +from _test_utils.torch.distributed.utils import DistributedWorkerPool from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider from megatron.bridge.training.config import ( CheckpointConfig, @@ -44,7 +45,7 @@ ] -def _train(tmp_path, qad, recipe, train_steps=1): +def _train(tmp_path, qad, recipe, train_steps=1, tp_size=1): captured = {} def provider(name): @@ -55,6 +56,8 @@ def provider(name): num_attention_heads=2, hybrid_layer_pattern="G", vocab_size=128, + tensor_model_parallel_size=tp_size, + sequence_parallel=tp_size > 1, seq_length=16, linear_num_key_heads=2, linear_num_value_heads=2, @@ -97,7 +100,9 @@ def check_mask(module, args, kwargs): model=provider("student"), train=TrainingConfig(train_iters=train_steps, global_batch_size=2, micro_batch_size=1), validation=ValidationConfig(eval_iters=0, eval_interval=1), - optimizer=OptimizerConfig(optimizer="adam", lr=1e-2, min_lr=0, weight_decay=0), + optimizer=OptimizerConfig( + optimizer="adam", lr=1e-2, min_lr=0, weight_decay=0, use_distributed_optimizer=True + ), scheduler=SchedulerConfig( lr_decay_style="constant", lr_warmup_iters=0, @@ -105,7 +110,9 @@ def check_mask(module, args, kwargs): end_weight_decay=0, use_checkpoint_opt_param_scheduler=True, ), - ddp=DistributedDataParallelConfig(average_in_collective=False), + ddp=DistributedDataParallelConfig( + average_in_collective=False, use_distributed_optimizer=True + ), dataset=MockGPTDatasetConfig( seq_length=16, random_seed=123, @@ -127,24 +134,43 @@ def check_mask(module, args, kwargs): return captured, checkpoint +def _reset_worker(rank, world_size): + reset_megatron_global_state() + + +def _warmup(rank, world_size, path, recipe): + _train(path / "dp", False, recipe) + reset_megatron_global_state() + _train(path / "tp", True, recipe, tp_size=world_size) + + @pytest.fixture(scope="module") -def compiled_state_training(tmp_path_factory, project_root_path): - """Compile one small GDN shape before timing the QAT/QAD integration tests.""" +def compiled_state_training(tmp_path_factory, project_root_path, num_gpus): + """Warm up at most two ranks before timing the DP QAT / TP QAD tests.""" + if not num_gpus: + pytest.skip("Requires CUDA") recipe = json.loads( ( project_root_path / "examples/llm_qat/linear_attention/configs/decode_state_int8.json" ).read_text() ) + workers = DistributedWorkerPool(min(num_gpus, 2), teardown_fn=_reset_worker) try: - _train(tmp_path_factory.mktemp("compile_state_training"), True, recipe) + workers.run(_warmup, tmp_path_factory.mktemp("compile_state_training"), recipe) + yield workers, recipe finally: - reset_megatron_global_state() - return recipe + workers.shutdown() @pytest.mark.parametrize("qad", [False, True], ids=["qat", "qad"]) def test_state_training(compiled_state_training, tmp_path, qad): - captured, checkpoint = _train(tmp_path, qad, compiled_state_training) + workers, recipe = compiled_state_training + workers.run(_check_training, tmp_path, qad, recipe) + + +def _check_training(rank, world_size, tmp_path, qad, recipe): + tp_size = world_size if qad else 1 + captured, checkpoint = _train(tmp_path, qad, recipe, tp_size=tp_size) student = captured["student"] assert has_modelopt_state(checkpoint) assert is_quantized(student) @@ -177,7 +203,7 @@ def test_state_training(compiled_state_training, tmp_path, qad): } reset_megatron_global_state() # An empty recipe verifies that resume restores the saved quantization policy. - resumed, _ = _train(tmp_path, True, {}, train_steps=2) + resumed, _ = _train(tmp_path, True, {}, train_steps=2, tp_size=tp_size) assert has_modelopt_state(str(tmp_path / "checkpoints/iter_0000002")) assert is_quantized(resumed["student"]) assert not is_quantized(resumed["teacher"]) From 739c1370e45525b23548e8b890e61c1a5ac44ede Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 3 Oct 2026 14:56:01 -0700 Subject: [PATCH 05/27] Keep linear attention training tests minimal Signed-off-by: Kai Xu --- .../megatron_bridge/test_linear_attention.py | 105 ++++-------------- 1 file changed, 20 insertions(+), 85 deletions(-) diff --git a/tests/examples/megatron_bridge/test_linear_attention.py b/tests/examples/megatron_bridge/test_linear_attention.py index c156b57a53c..2ab89aad2b2 100644 --- a/tests/examples/megatron_bridge/test_linear_attention.py +++ b/tests/examples/megatron_bridge/test_linear_attention.py @@ -20,7 +20,6 @@ import torch from _test_utils.examples.megatron_example_runner import reset_megatron_global_state from _test_utils.examples.run_command import MODELOPT_ROOT -from _test_utils.torch.distributed.utils import DistributedWorkerPool from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider from megatron.bridge.training.config import ( CheckpointConfig, @@ -35,7 +34,6 @@ TrainingConfig, ValidationConfig, ) -from megatron.bridge.training.post_training.checkpointing import has_modelopt_state from megatron.core.utils import unwrap_model from modelopt.torch.quantization.utils import is_quantized @@ -45,7 +43,7 @@ ] -def _train(tmp_path, qad, recipe, train_steps=1, tp_size=1): +def _train(qad, recipe): captured = {} def provider(name): @@ -56,8 +54,6 @@ def provider(name): num_attention_heads=2, hybrid_layer_pattern="G", vocab_size=128, - tensor_model_parallel_size=tp_size, - sequence_parallel=tp_size > 1, seq_length=16, linear_num_key_heads=2, linear_num_value_heads=2, @@ -79,26 +75,18 @@ def provider(name): def capture(models): module = unwrap_model(models[0]) captured[name] = module - captured[name + "_before"] = { - key: value.detach().clone() for key, value in module.named_parameters() - } if name == "student": - - def check_mask(module, args, kwargs): - mask = kwargs["loss_mask"] - assert not mask[:, :8].any() - assert mask[:, 8:].all() - - module.register_forward_pre_hook(check_mask, with_kwargs=True) + captured["before"] = ( + module.decoder.layers[0].self_attention.out_proj.weight.detach().clone() + ) return models model.register_post_wrap_hook(capture) return model - checkpoint = str(tmp_path / "checkpoints") config = ConfigContainer( model=provider("student"), - train=TrainingConfig(train_iters=train_steps, global_batch_size=2, micro_batch_size=1), + train=TrainingConfig(train_iters=1, global_batch_size=1, micro_batch_size=1), validation=ValidationConfig(eval_iters=0, eval_interval=1), optimizer=OptimizerConfig( optimizer="adam", lr=1e-2, min_lr=0, weight_decay=0, use_distributed_optimizer=True @@ -108,7 +96,6 @@ def check_mask(module, args, kwargs): lr_warmup_iters=0, start_weight_decay=0, end_weight_decay=0, - use_checkpoint_opt_param_scheduler=True, ), ddp=DistributedDataParallelConfig( average_in_collective=False, use_distributed_optimizer=True @@ -123,92 +110,40 @@ def check_mask(module, args, kwargs): num_workers=0, ), tokenizer=TokenizerConfig(tokenizer_type="NullTokenizer", vocab_size=128), - checkpoint=CheckpointConfig( - save=checkpoint, load=checkpoint, save_interval=1, async_save=False - ), + checkpoint=CheckpointConfig(async_save=False), logger=LoggerConfig(log_interval=1), rng=RNGConfig(seed=123), mixed_precision="bf16_mixed", ) run_training(config, recipe, 8, provider("teacher") if qad else None) - return captured, checkpoint - - -def _reset_worker(rank, world_size): - reset_megatron_global_state() - - -def _warmup(rank, world_size, path, recipe): - _train(path / "dp", False, recipe) - reset_megatron_global_state() - _train(path / "tp", True, recipe, tp_size=world_size) + return captured @pytest.fixture(scope="module") -def compiled_state_training(tmp_path_factory, project_root_path, num_gpus): - """Warm up at most two ranks before timing the DP QAT / TP QAD tests.""" - if not num_gpus: +def compiled_state_training(project_root_path): + """Compile one tiny GDN shape before timing the single-GPU training checks.""" + if not torch.cuda.is_available(): pytest.skip("Requires CUDA") recipe = json.loads( ( project_root_path / "examples/llm_qat/linear_attention/configs/decode_state_int8.json" ).read_text() ) - workers = DistributedWorkerPool(min(num_gpus, 2), teardown_fn=_reset_worker) try: - workers.run(_warmup, tmp_path_factory.mktemp("compile_state_training"), recipe) - yield workers, recipe + _train(True, recipe) finally: - workers.shutdown() + reset_megatron_global_state() + return recipe @pytest.mark.parametrize("qad", [False, True], ids=["qat", "qad"]) -def test_state_training(compiled_state_training, tmp_path, qad): - workers, recipe = compiled_state_training - workers.run(_check_training, tmp_path, qad, recipe) - - -def _check_training(rank, world_size, tmp_path, qad, recipe): - tp_size = world_size if qad else 1 - captured, checkpoint = _train(tmp_path, qad, recipe, tp_size=tp_size) - student = captured["student"] - assert has_modelopt_state(checkpoint) - assert is_quantized(student) - layers = [m for m in student.modules() if hasattr(m, "gdn_state_quantizer")] - assert len(layers) == 1 - assert layers[0].gdn_state_quantizer.is_enabled - assert layers[0].linear_attention_config.decode.state_codec == "int8_hadamard32" - changed = [] - for name, parameter in student.named_parameters(): - assert torch.isfinite(parameter).all() - if not torch.equal(parameter, captured["student_before"][name]): - assert parameter.requires_grad - changed.append(name) - assert changed - assert all( - module._linear_attention_prefill_lengths is None - for module in student.modules() - if hasattr(module, "_linear_attention_prefill_lengths") - ) +def test_state_training(compiled_state_training, qad): + captured = _train(qad, compiled_state_training) + attention = captured["student"].decoder.layers[0].self_attention + assert attention.gdn_state_quantizer.is_enabled + assert torch.isfinite(attention.out_proj.weight).all() + assert not torch.equal(attention.out_proj.weight, captured["before"]) if qad: teacher = captured["teacher"] assert not is_quantized(teacher) - for name, parameter in teacher.named_parameters(): - assert not parameter.requires_grad - assert torch.equal(parameter, captured["teacher_before"][name]) - before_resume = { - name: parameter.detach().clone() - for name, parameter in student.named_parameters() - if parameter.requires_grad - } - reset_megatron_global_state() - # An empty recipe verifies that resume restores the saved quantization policy. - resumed, _ = _train(tmp_path, True, {}, train_steps=2, tp_size=tp_size) - assert has_modelopt_state(str(tmp_path / "checkpoints/iter_0000002")) - assert is_quantized(resumed["student"]) - assert not is_quantized(resumed["teacher"]) - assert any( - not torch.equal(parameter, before_resume[name]) - for name, parameter in resumed["student"].named_parameters() - if name in before_resume - ) + assert not any(parameter.requires_grad for parameter in teacher.parameters()) From 8142af282af33b23003239270d29e4d77d76b544 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 3 Oct 2026 21:06:39 -0700 Subject: [PATCH 06/27] Make linear attention state QDQ format dispatch explicit Signed-off-by: Kai Xu --- .../quantization/linear_attention/__init__.py | 8 ++++- .../linear_attention/fla_chunk_delta_h.py | 30 +++++++++---------- .../fla_chunk_gated_delta_rule.py | 15 ++++------ modelopt/torch/quantization/plugins/gdn.py | 6 ++-- .../quantization/plugins/linear_attention.py | 9 ++++-- .../torch/quantization/plugins/test_gdn.py | 16 +++++++--- 6 files changed, 50 insertions(+), 34 deletions(-) diff --git a/modelopt/torch/kernels/quantization/linear_attention/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/__init__.py index 3f003f403f8..30a2ee9517a 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/__init__.py +++ b/modelopt/torch/kernels/quantization/linear_attention/__init__.py @@ -18,7 +18,13 @@ ``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 +recurrent state carried between chunks to FP8 or INT8 (``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. """ + +# Kernel format IDs, independent of scale grouping. The current kernels use dynamic scales. +STATE_QDQ_OFF = 0 +STATE_QDQ_FP8_E4M3 = 1 +STATE_QDQ_INT8 = 2 +STATE_QDQ_FORMATS = {"fp8_e4m3": STATE_QDQ_FP8_E4M3, "int8": STATE_QDQ_INT8} 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 index e316fbf5150..815d48a7583 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -40,14 +40,11 @@ check_shared_mem, ) +from modelopt.torch.kernels.quantization import linear_attention as state_formats 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 @@ -62,20 +59,23 @@ def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr, STATE_QDQ: tl.cons 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))) - # STATE_QDQ=2 uses symmetric INT8 (max 127); STATE_QDQ=1 uses FP8 E4M3 (max 448). # Use scale=1 for zero tiles to avoid division by zero during QDQ. - if STATE_QDQ == 2: + if STATE_QDQ == state_formats.STATE_QDQ_INT8: return tl.where(b_amax > 0, b_amax * (1.0 / 127.0), 1.0) - else: + elif STATE_QDQ == state_formats.STATE_QDQ_FP8_E4M3: return tl.where(b_amax > 0, b_amax / 448.0, 1.0) + else: + tl.static_assert(False, "Unsupported state QDQ format") @triton.jit def _state_scalar_qdq(value, scale, STATE_QDQ: tl.constexpr): - if STATE_QDQ == 2: + if STATE_QDQ == state_formats.STATE_QDQ_INT8: return int8_scalar_qdq(value, scale) - else: + elif STATE_QDQ == state_formats.STATE_QDQ_FP8_E4M3: return fp8_scalar_qdq(value, scale) + else: + tl.static_assert(False, "Unsupported state QDQ format") NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] @@ -233,7 +233,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( 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 STATE_QDQ != state_formats.STATE_QDQ_OFF: 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: @@ -404,7 +404,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( # [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 STATE_QDQ != state_formats.STATE_QDQ_OFF: 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: @@ -838,7 +838,7 @@ def chunk_gated_delta_rule_fwd_h( 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: int = state_formats.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] @@ -900,10 +900,10 @@ def state_qdq_tile_v(V: int, state_qdq: int, state_qdq_block_v: int | None) -> i 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: + if state_qdq == state_formats.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.") + if state_qdq not in state_formats.STATE_QDQ_FORMATS.values(): + raise ValueError(f"Unsupported state_qdq format {state_qdq}.") 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( 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 index f34b8a3384b..0cf33a2c947 100644 --- 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 @@ -57,13 +57,8 @@ 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, -) +from . import STATE_QDQ_FORMATS, STATE_QDQ_FP8_E4M3, STATE_QDQ_OFF +from .fla_chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h def chunk_gated_delta_rule_fwd( @@ -649,11 +644,11 @@ def chunk_gated_delta_rule( 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 state_qdq != STATE_QDQ_OFF and state_qdq not in STATE_QDQ_FORMATS.values(): + raise ValueError(f"Unsupported state_qdq format {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 ( + if state_qdq == STATE_QDQ_FP8_E4M3 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.") diff --git a/modelopt/torch/quantization/plugins/gdn.py b/modelopt/torch/quantization/plugins/gdn.py index 8d0fe616390..049b16acd66 100644 --- a/modelopt/torch/quantization/plugins/gdn.py +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -33,6 +33,8 @@ import torch +from modelopt.torch.kernels.quantization.linear_attention import STATE_QDQ_FORMATS, STATE_QDQ_OFF + from ..linear_attention.prefill import matmul_gdn from .linear_attention import _LinearAttentionQuantMixin @@ -114,9 +116,9 @@ def _state_quantized_chunk_gated_delta_rule( return _state_qdq_chunk_gated_delta_rule()( *args, chunk_size=chunk_size, - state_qdq=(2 if self._linear_attn_state_format == "int8" else 1) + state_qdq=STATE_QDQ_FORMATS[self._linear_attn_state_format] if quantize_state - else 0, + else STATE_QDQ_OFF, 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/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py index 5559c9f6606..036bec66f10 100644 --- a/modelopt/torch/quantization/plugins/linear_attention.py +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -22,6 +22,8 @@ __all__ = [] +_STATE_FORMATS: dict[int | tuple[int, int], str] = {(4, 3): "fp8_e4m3", 8: "int8"} + class _LinearAttentionQuantMixin(QuantModule): linear_attention_quantizer_names = ("gdn_state_quantizer", "gdn_w_quantizer") @@ -40,7 +42,10 @@ def _linear_attn_state(self): @property def _linear_attn_state_format(self): - return "int8" if self._linear_attn_state.num_bits == 8 else "fp8_e4m3" + # Disabled quantizers impose no format requirement on the numerical backend. + if not self._linear_attn_state.is_enabled: + return "fp8_e4m3" + return _STATE_FORMATS[self._linear_attn_state.num_bits] @property def _linear_attn_w(self): @@ -61,7 +66,7 @@ def validate_linear_attention(self): validate_gdn_quantizer( self._linear_attn_state, name=self.linear_attention_quantizer_names[0], - num_bits=((4, 3), 8), + num_bits=tuple(_STATE_FORMATS), ) if self._linear_attn_state.axis != (0, 1): raise ValueError( diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index 3f24f2cd64e..a24d1cc4123 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -84,6 +84,7 @@ def quant_cfg(state=True, w=False): {"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, "axis": (0, 1), "type": "dynamic"}, # unsupported format {"num_bits": (4, 3), "type": "dynamic", "block_sizes": {-1: 16}}, # blockwise ], ) @@ -128,8 +129,13 @@ def test_disabled_state_quantizer_calls_original_kernel(): assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" -@pytest.mark.parametrize("partial_kernel", [False, True]) -def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, partial_kernel): +@pytest.mark.parametrize( + ("partial_kernel", "num_bits", "state_qdq"), + [(False, (4, 3), 1), (True, (4, 3), 1), (False, 8, 2)], +) +def test_enabled_state_quantizer_uses_state_qdq_kernel( + monkeypatch, partial_kernel, num_bits, state_qdq +): calls = [] def fake_state_qdq_kernel(*args, **kwargs): @@ -141,12 +147,14 @@ def fake_state_qdq_kernel(*args, **kwargs): if partial_kernel: model.gated_delta_rule = partial(chunk_gated_delta_rule, output_final_state=True) x = torch.randn(2, 8, 3, 4) - mtq.quantize(model, quant_cfg(), lambda m: m(x)) + cfg = deepcopy(quant_cfg()) + cfg["quant_cfg"][-1]["cfg"].update(num_bits=num_bits, unsigned=False, narrow_range=True) + mtq.quantize(model, cfg, lambda m: m(x)) model(x) assert calls and calls[-1] == { "chunk_size": 64, - "state_qdq": 1, + "state_qdq": state_qdq, "state_qdq_block_v": 64, "w_quantizer": None, **({"output_final_state": True} if partial_kernel else {}), From 17ccf0e97d48062f3cb8cf4b08b71285075623ee Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 3 Oct 2026 22:07:06 -0700 Subject: [PATCH 07/27] Support TensorQuantizer blockwise INT8 linear attention state QDQ Use registered state quantizers for format and last-axis grouping, with INT8 groups of 16, 32, or 64 independent of execution tile width. Preserve existing tile and Hadamard codecs and document the blockwise configuration. Keep minimal output/gradient, configuration, and checkpoint coverage; omit redundant mocked block routing and quantizer call-count checks. Signed-off-by: Kai Xu --- examples/llm_qat/linear_attention/README.md | 34 ++++++++++ .../linear_attention/fla_chunk_delta_h.py | 58 +++++++++++----- .../fla_chunk_gated_delta_rule.py | 28 +++++++- .../quantization/linear_attention/int8.py | 17 +++++ .../linear_attention/_chunk_prefill.py | 10 +-- .../quantization/linear_attention/config.py | 2 + .../quantization/linear_attention/decode.py | 68 ++++++++++++++++--- .../linear_attention/decode_prefill.py | 8 +++ .../quantization/linear_attention/kda.py | 2 + .../quantization/linear_attention/prefill.py | 2 + .../quantization/linear_attention/utils.py | 40 ++++++++++- modelopt/torch/quantization/plugins/gdn.py | 12 +--- modelopt/torch/quantization/plugins/kda.py | 3 +- .../quantization/plugins/linear_attention.py | 18 ++--- .../linear_attention_reference.py | 14 +++- .../test_fla_chunk_gated_delta_rule.py | 12 +++- .../torch/quantization/plugins/test_gdn.py | 25 ++++--- .../test_linear_attention_decode.py | 44 ++++++++++++ 18 files changed, 330 insertions(+), 67 deletions(-) diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md index c06f20fc6a5..f4db6058454 100644 --- a/examples/llm_qat/linear_attention/README.md +++ b/examples/llm_qat/linear_attention/README.md @@ -299,6 +299,40 @@ context and leaves W/projection quantizers disabled. KDA uses the materialized backend and can use the all-prefix context shown above. +### Per-row blockwise INT8 through TensorQuantizer + +Use the standard quantizer settings to select one dynamic scale per key row and +16, 32, or 64 consecutive value channels. For example, replace the state quantizer +configuration in `gdn_recipe` above with: + +```python +gdn_recipe["quant_cfg"][1]["cfg"] = { + "num_bits": 8, + "type": "dynamic", + "block_sizes": {-1: 32}, + "unsigned": False, + "narrow_range": True, + "pass_through_bwd": True, +} +``` + +The logical state is `[N, H, K, V]`; `-1` denotes V even when a kernel stores the +state transposed. Omit `axis` when using `block_sizes`. The outer `type="dynamic"` +uses ordinary INT8 with dynamic amax; a nested `block_sizes["type"]="dynamic"` +selects a different specialized ModelOpt path and is not supported here. + +The GDN fused kernel applies the quantizer's grouping inside each execution tile. +It can enlarge `state.block_v` to fit a complete quantization group; changing the +execution tile does not change the scale groups. The Torch GDN/KDA path calls the +registered `TensorQuantizer` at the same configured state-write boundaries. +For KDA, use `*kda_state_quantizer` and the materialized backend. For token decode +or replay, explicitly select `state_codec="tile"`; the existing Hadamard recipe +keeps its fixed codec and cannot be combined with this `block_sizes` setting. + +Configurations without `block_sizes` retain their existing `[K, block_v]` scale +grouping, or their fixed Hadamard codec. The default INT8 + Hadamard recipe is +unchanged. This remains floating-storage fake QDQ for QAT/QAD. + ## Integrate with a training loop Apply the configured `recipe` above with `mtq.quantize`, then supply one prefix length per sequence. 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 index 815d48a7583..ee6a8cd1947 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -43,7 +43,7 @@ from modelopt.torch.kernels.quantization import linear_attention as state_formats from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq -from .int8 import int8_scalar_qdq +from .int8 import int8_block_qdq, int8_scalar_qdq STATE_QDQ_MAX_BLOCK_V = 128 @@ -69,10 +69,15 @@ def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr, STATE_QDQ: tl.cons @triton.jit -def _state_scalar_qdq(value, scale, STATE_QDQ: tl.constexpr): +def _state_scalar_qdq( + value, scale, STATE_QDQ: tl.constexpr, GROUP_SIZE: tl.constexpr, STATE_V_FIRST: tl.constexpr +): if STATE_QDQ == state_formats.STATE_QDQ_INT8: + if GROUP_SIZE: + return int8_block_qdq(value, GROUP_SIZE, STATE_V_FIRST) return int8_scalar_qdq(value, scale) elif STATE_QDQ == state_formats.STATE_QDQ_FP8_E4M3: + tl.static_assert(GROUP_SIZE == 0, "Blockwise state QDQ currently supports only INT8") return fp8_scalar_qdq(value, scale) else: tl.static_assert(False, "Unsupported state QDQ format") @@ -111,7 +116,7 @@ def _state_scalar_qdq(value, scale, STATE_QDQ: tl.constexpr): 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"], + key=["H", "HV", "K", "V", "BT", "BV", "STATE_V_FIRST", "STATE_QDQ", "STATE_QDQ_GROUP_SIZE"], ) @triton.jit(do_not_specialize=["T"]) def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( @@ -141,6 +146,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( STATE_V_FIRST: tl.constexpr, IS_VARLEN: tl.constexpr, STATE_QDQ: tl.constexpr, + STATE_QDQ_GROUP_SIZE: tl.constexpr, ): pid = tl.program_id(0) NV = tl.cdiv(V, BV) @@ -234,7 +240,9 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( 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 != state_formats.STATE_QDQ_OFF: - if K > 192: + if STATE_QDQ_GROUP_SIZE: + b_scale = 1.0 + elif 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) @@ -242,13 +250,19 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( 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) + b_h1 = _state_scalar_qdq(b_h1, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST) if K > 64: - b_h2 = _state_scalar_qdq(b_h2, b_scale, STATE_QDQ) + b_h2 = _state_scalar_qdq( + b_h2, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST + ) if K > 128: - b_h3 = _state_scalar_qdq(b_h3, b_scale, STATE_QDQ) + b_h3 = _state_scalar_qdq( + b_h3, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST + ) if K > 192: - b_h4 = _state_scalar_qdq(b_h4, b_scale, STATE_QDQ) + b_h4 = _state_scalar_qdq( + b_h4, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST + ) # main recurrence for i_t in range(NT): @@ -403,9 +417,11 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( 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. + # Legacy quantizers share a [K, BV] scale; block quantizers scale per key/value group. if STATE_QDQ != state_formats.STATE_QDQ_OFF: - if K > 192: + if STATE_QDQ_GROUP_SIZE: + b_scale = 1.0 + elif 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) @@ -413,13 +429,19 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( 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) + b_h1 = _state_scalar_qdq(b_h1, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST) if K > 64: - b_h2 = _state_scalar_qdq(b_h2, b_scale, STATE_QDQ) + b_h2 = _state_scalar_qdq( + b_h2, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST + ) if K > 128: - b_h3 = _state_scalar_qdq(b_h3, b_scale, STATE_QDQ) + b_h3 = _state_scalar_qdq( + b_h3, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST + ) if K > 192: - b_h4 = _state_scalar_qdq(b_h4, b_scale, STATE_QDQ) + b_h4 = _state_scalar_qdq( + b_h4, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST + ) if STORE_FINAL_STATE: if STATE_V_FIRST: @@ -840,10 +862,11 @@ def chunk_gated_delta_rule_fwd_h( chunk_offsets: torch.LongTensor | None = None, state_qdq: int = state_formats.STATE_QDQ_OFF, state_qdq_block_v: int | None = None, + state_qdq_group_size: int = 0, ) -> 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) + BV = max(state_qdq_tile_v(V, state_qdq, state_qdq_block_v), state_qdq_group_size) if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) @@ -887,6 +910,7 @@ def chunk_gated_delta_rule_fwd_h( BV=BV, STATE_V_FIRST=state_v_first, STATE_QDQ=state_qdq, + STATE_QDQ_GROUP_SIZE=state_qdq_group_size, ) return h, v_new, final_state @@ -894,8 +918,8 @@ def chunk_gated_delta_rule_fwd_h( 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 + Without state quantization this is fla's largest tile. For legacy quantizers the tile is + also the quantization granularity: one 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. 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 index 0cf33a2c947..703b36b1303 100644 --- 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 @@ -54,7 +54,10 @@ input_guard, ) -from modelopt.torch.quantization.linear_attention.utils import validate_gdn_quantizer +from modelopt.torch.quantization.linear_attention.utils import ( + state_quantizer_config, + validate_gdn_quantizer, +) from modelopt.torch.quantization.nn import TensorQuantizer from . import STATE_QDQ_FORMATS, STATE_QDQ_FP8_E4M3, STATE_QDQ_OFF @@ -80,6 +83,7 @@ def chunk_gated_delta_rule_fwd( chunk_size: int = 64, state_qdq: int = STATE_QDQ_OFF, state_qdq_block_v: int | None = None, + state_qdq_group_size: int = 0, w_quantizer: TensorQuantizer | None = None, ): g_input = g if use_gate_in_kernel else None @@ -143,6 +147,7 @@ def chunk_gated_delta_rule_fwd( chunk_size=chunk_size, state_qdq=state_qdq, state_qdq_block_v=state_qdq_block_v, + state_qdq_group_size=state_qdq_group_size, ) if cp_context is not None: @@ -185,6 +190,7 @@ def chunk_gated_delta_rule_bwd( chunk_size: int = 64, state_qdq: int = STATE_QDQ_OFF, state_qdq_block_v: int | None = None, + state_qdq_group_size: int = 0, quantized_w: torch.Tensor | None = None, ): w, u = recompute_w_u_fwd( @@ -219,6 +225,7 @@ def chunk_gated_delta_rule_bwd( chunk_size=chunk_size, state_qdq=state_qdq, state_qdq_block_v=state_qdq_block_v, + state_qdq_group_size=state_qdq_group_size, ) dv = chunk_bwd_dv_local( q=q, @@ -331,6 +338,7 @@ def forward( chunk_size: int = 64, state_qdq: int = STATE_QDQ_OFF, state_qdq_block_v: int | None = None, + state_qdq_group_size: int = 0, w_quantizer: TensorQuantizer | None = None, ): q_rstd, k_rstd = None, None @@ -365,6 +373,7 @@ def forward( chunk_size=chunk_size, state_qdq=state_qdq, state_qdq_block_v=state_qdq_block_v, + state_qdq_group_size=state_qdq_group_size, w_quantizer=w_quantizer, ) ctx.save_for_backward( @@ -395,6 +404,7 @@ def forward( ctx.use_gate_in_kernel = use_gate_in_kernel ctx.state_qdq = state_qdq ctx.state_qdq_block_v = state_qdq_block_v + ctx.state_qdq_group_size = state_qdq_group_size return o.to(q.dtype), final_state @staticmethod @@ -445,6 +455,7 @@ def backward( chunk_size=ctx.chunk_size, state_qdq=ctx.state_qdq, state_qdq_block_v=ctx.state_qdq_block_v, + state_qdq_group_size=ctx.state_qdq_group_size, quantized_w=quantized_w, ) if ctx.use_qk_l2norm_in_kernel: @@ -476,6 +487,7 @@ def backward( None, None, None, + None, ) @@ -549,6 +561,9 @@ def chunk_gated_delta_rule( 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``. + state_quantizer (Optional[TensorQuantizer]): + [ModelOpt] State QDQ settings, including optional INT8 last-axis ``block_sizes``. + Blocks always refer to the logical V dimension, regardless of storage layout. cu_seqlens (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. @@ -636,6 +651,16 @@ def chunk_gated_delta_rule( # [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) + state_quantizer = kwargs.pop("state_quantizer", None) + state_qdq_group_size = 0 + if state_quantizer is not None: + if state_qdq != STATE_QDQ_OFF: + raise ValueError("Specify state_quantizer or state_qdq, not both") + if state_quantizer.is_enabled: + state_format, group_size = state_quantizer_config(state_quantizer) + if state_quantizer._if_quant: + state_qdq = STATE_QDQ_FORMATS[state_format] + state_qdq_group_size = group_size # 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) @@ -712,6 +737,7 @@ def chunk_gated_delta_rule( chunk_size, state_qdq, state_qdq_block_v, + state_qdq_group_size, w_quantizer, ) return o, final_state diff --git a/modelopt/torch/kernels/quantization/linear_attention/int8.py b/modelopt/torch/kernels/quantization/linear_attention/int8.py index d9e16a29dc8..1fa5c621af9 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/int8.py +++ b/modelopt/torch/kernels/quantization/linear_attention/int8.py @@ -26,3 +26,20 @@ 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 + + +@triton.jit +def int8_block_qdq(value, GROUP_SIZE: tl.constexpr, STATE_V_FIRST: tl.constexpr): + """Match TensorQuantizer's dynamic INT8 QDQ per key row and value group.""" + if STATE_V_FIRST: + value = tl.trans(value) + groups = tl.reshape(value, (value.shape[0], value.shape[1] // GROUP_SIZE, GROUP_SIZE)) + amax = tl.max(tl.abs(groups), axis=2, keep_dims=True) + # Match the CUDA TensorQuantizer's scale direction, ties-to-even, and tiny-group handling. + tiny = amax < 2.0**-24 + scale = tl.div_rn(127.0, tl.where(tiny, 1.0, amax)) + codes = tl.clamp(libdevice.nearbyint(groups * scale), -127.0, 127.0) + rounded = tl.reshape(tl.where(tiny, 0.0, tl.div_rn(codes, scale)), value.shape) + if STATE_V_FIRST: + rounded = tl.trans(rounded) + return rounded diff --git a/modelopt/torch/quantization/linear_attention/_chunk_prefill.py b/modelopt/torch/quantization/linear_attention/_chunk_prefill.py index fddff0303ae..c4dfc48dbd1 100644 --- a/modelopt/torch/quantization/linear_attention/_chunk_prefill.py +++ b/modelopt/torch/quantization/linear_attention/_chunk_prefill.py @@ -37,6 +37,7 @@ def chunk_gdn( state_qdq: bool = False, state_qdq_block_v: int = 64, state_format: str = "fp8_e4m3", + state_quantizer=None, ) -> 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: @@ -47,7 +48,7 @@ def chunk_gdn( for n, (b, start, end) in enumerate(sequences): state = states[n] if state_qdq: - state = _state_qdq(state, state_qdq_block_v, state_format) + state = _state_qdq(state, state_qdq_block_v, state_format, state_quantizer) pieces = [] for lo in range(start, end, chunk_size): hi = min(lo + chunk_size, end) @@ -71,7 +72,7 @@ def chunk_gdn( 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) + state = _state_qdq(state, state_qdq_block_v, state_format, state_quantizer) outputs.append(torch.cat(pieces)) finals.append(state) output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) @@ -94,6 +95,7 @@ def chunk_kda( state_qdq=False, state_qdq_block_v=64, state_format="fp8_e4m3", + state_quantizer=None, ): """Compute a KDA prefix with optional QDQ on the initial state and chunk writes.""" if g.ndim != 4 or chunk_size <= 0: @@ -104,7 +106,7 @@ def chunk_kda( for n, (b, start, end) in enumerate(sequences): state = states[n] if state_qdq: - state = _state_qdq(state, state_qdq_block_v, state_format) + state = _state_qdq(state, state_qdq_block_v, state_format, state_quantizer) pieces = [] for lo in range(start, end, chunk_size): hi = min(lo + chunk_size, end) @@ -138,7 +140,7 @@ def chunk_kda( 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) + state = _state_qdq(state, state_qdq_block_v, state_format, state_quantizer) outputs.append(torch.cat(pieces)) finals.append(state) output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index 1f3b610fe75..4d3cc108fd5 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -78,6 +78,8 @@ class LinearAttentionConfig(ModeloptBaseConfig): ``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. + With INT8 ``TensorQuantizer.block_sizes``, scales instead follow per-key value + groups; ``state.block_v`` controls only execution tiling (at least one group). """ schema_version: Literal[1] = ModeloptField(default=1) diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py index 2a981c597ec..e15dd8a14ea 100644 --- a/modelopt/torch/quantization/linear_attention/decode.py +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -20,7 +20,7 @@ import torch from .config import LinearAttentionDecodeConfig -from .utils import _fp8_quantize +from .utils import _fp8_quantize, state_quantizer_config __all__ = [ "EncodedLinearAttentionTensor", @@ -33,7 +33,7 @@ @dataclass class EncodedLinearAttentionTensor: - """Fake decoded values plus detached scale metadata; no compressed storage claim.""" + """Fake decoded values plus optional scale metadata; no compressed storage claim.""" values: torch.Tensor scales: torch.Tensor | None @@ -102,9 +102,23 @@ def backward(ctx, gradient): return gradient, None -def _encode(value, enabled, block_v, *, state=False, state_format="fp8_e4m3", state_codec="tile"): +def _encode( + value, + enabled, + block_v, + *, + state=False, + state_format="fp8_e4m3", + state_codec="tile", + state_quantizer=None, +): if not enabled: return EncodedLinearAttentionTensor(value, None, "identity", None) + if state and state_quantizer is not None and state_quantizer.block_sizes is not None: + # TensorQuantizer owns dynamic scales; the floating carry needs only its QDQ output. + return EncodedLinearAttentionTensor( + state_quantizer(value), None, state_format, state_quantizer.block_sizes[-1] + ) if state and state_codec == "int8_hadamard32": with torch.no_grad(): groups = value.float().reshape(*value.shape[:-1], -1, 32) @@ -138,11 +152,14 @@ def _encode(value, enabled, block_v, *, state=False, state_format="fp8_e4m3", st return EncodedLinearAttentionTensor(decoded, torch.stack(scales, dim=-1), state_format, block_v) -def _signature(config, state_qdq, block_v, state_format): - return ( +def _signature(config, state_qdq, block_v, state_format, state_quantizer): + signature = ( config.model_dump_json(exclude={"implementation", "prefill_state_qdq"}) + f"/{state_qdq}/{block_v}/{state_format}" ) + if state_quantizer is not None and state_quantizer.block_sizes is not None: + signature += f"/group={state_quantizer.block_sizes[-1]}" + return signature def _sum_keys(value): @@ -164,7 +181,19 @@ def _round_log_gate(gate, step): def _prepare_carry( - q, k, v, g, beta, config, state_qdq, block_v, initial_state, carry, position, state_format + q, + k, + v, + g, + beta, + config, + state_qdq, + block_v, + initial_state, + carry, + position, + state_format, + state_quantizer, ): 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") @@ -176,11 +205,13 @@ def _prepare_carry( raise ValueError("State format must be fp8_e4m3 or int8") hadamard = config.state_codec == "int8_hadamard32" if hadamard: + if state_quantizer is not None and state_quantizer.block_sizes is not None: + raise ValueError("TensorQuantizer block_sizes requires state_codec='tile'") 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) + signature = _signature(config, state_qdq, block_v, state_format, state_quantizer) 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]) @@ -208,6 +239,7 @@ def _prepare_carry( state=True, state_format=state_format, state_codec=config.state_codec, + state_quantizer=state_quantizer, ) carry = LinearAttentionCarry( anchor, (), carry.position, True, signature, "hadamard32" if hadamard else "identity" @@ -225,6 +257,7 @@ def recurrent_decode_reference( config: LinearAttentionDecodeConfig, state_qdq=False, state_format="fp8_e4m3", + state_quantizer=None, block_v=64, initial_state=None, carry=None, @@ -237,8 +270,24 @@ def recurrent_decode_reference( KDA log gates are accepted. Outputs and all returned carry values retain their graphs. An empty call performs no write or initial-state quantization. """ + if state_quantizer is not None: + state_qdq = state_quantizer.is_enabled and state_quantizer._if_quant + if state_quantizer.is_enabled: + state_format, _ = state_quantizer_config(state_quantizer) carry, signature = _prepare_carry( - q, k, v, g, beta, config, state_qdq, block_v, initial_state, carry, position, state_format + q, + k, + v, + g, + beta, + config, + state_qdq, + block_v, + initial_state, + carry, + position, + state_format, + state_quantizer, ) if len(q) == 0: # Keep empty input gradients defined without introducing a state write. @@ -289,6 +338,7 @@ def recurrent_decode_reference( state=True, state_format=state_format, state_codec=config.state_codec, + state_quantizer=state_quantizer, ) entries = () stored = anchor.values @@ -317,6 +367,7 @@ def recurrent_decode( config: LinearAttentionDecodeConfig, state_qdq=False, state_format="fp8_e4m3", + state_quantizer=None, block_v=64, initial_state=None, carry=None, @@ -333,6 +384,7 @@ def recurrent_decode( config=config, state_qdq=state_qdq, state_format=state_format, + state_quantizer=state_quantizer, block_v=block_v, initial_state=initial_state, carry=carry, diff --git a/modelopt/torch/quantization/linear_attention/decode_prefill.py b/modelopt/torch/quantization/linear_attention/decode_prefill.py index 0a0f69b4ae7..28dbfa4466b 100644 --- a/modelopt/torch/quantization/linear_attention/decode_prefill.py +++ b/modelopt/torch/quantization/linear_attention/decode_prefill.py @@ -22,6 +22,7 @@ from ._chunk_prefill import chunk_gdn, chunk_kda from .decode import recurrent_decode +from .utils import state_quantizer_config __all__ = ["linear_attention_training_phase"] @@ -74,6 +75,7 @@ def _decode_prefill( policy, state_qdq, state_format, + state_quantizer, scale, initial_state, output_final_state, @@ -83,6 +85,10 @@ def _decode_prefill( output_dtype, prefill_lengths, ): + if state_quantizer is not None: + state_qdq = state_quantizer.is_enabled and state_quantizer._if_quant + if state_quantizer.is_enabled: + state_format, _ = state_quantizer_config(state_quantizer) 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]: @@ -143,6 +149,7 @@ def _decode_prefill( *packed, state_qdq=state_qdq and policy.decode.prefill_state_qdq, state_format=state_format, + state_quantizer=state_quantizer, scale=scale, initial_state=states[active], cu_seqlens=torch.tensor(offsets), @@ -161,6 +168,7 @@ def _decode_prefill( config=policy.decode, state_qdq=state_qdq, state_format=state_format, + state_quantizer=state_quantizer, block_v=policy.state.block_v, initial_state=prefix_states.get(n, states[n]), position=prefixes[n], diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py index fbde42dbc29..d7cc0229669 100644 --- a/modelopt/torch/quantization/linear_attention/kda.py +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -33,6 +33,7 @@ def matmul_kda( policy, state_qdq=False, state_format="fp8_e4m3", + state_quantizer=None, scale=None, initial_state=None, output_final_state=False, @@ -95,6 +96,7 @@ def matmul_kda( policy=policy, state_qdq=state_qdq, state_format=state_format, + state_quantizer=state_quantizer, scale=scale, initial_state=initial_state, output_final_state=output_final_state, diff --git a/modelopt/torch/quantization/linear_attention/prefill.py b/modelopt/torch/quantization/linear_attention/prefill.py index 4a602f8f5bc..9bd514a3ab6 100644 --- a/modelopt/torch/quantization/linear_attention/prefill.py +++ b/modelopt/torch/quantization/linear_attention/prefill.py @@ -34,6 +34,7 @@ def matmul_gdn( policy: LinearAttentionConfig, state_qdq=False, state_format="fp8_e4m3", + state_quantizer=None, scale=None, initial_state=None, output_final_state=False, @@ -82,6 +83,7 @@ def matmul_gdn( policy=policy, state_qdq=state_qdq, state_format=state_format, + state_quantizer=state_quantizer, scale=scale, initial_state=initial_state, output_final_state=output_final_state, diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py index 68353bd28df..d51bb76101b 100644 --- a/modelopt/torch/quantization/linear_attention/utils.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -27,12 +27,15 @@ __all__ = [] +_STATE_FORMATS: dict[int | tuple[int, int], str] = {(4, 3): "fp8_e4m3", 8: "int8"} + def validate_gdn_quantizer( quantizer: TensorQuantizer, *, name: str, num_bits: tuple[int | tuple[int, int], ...] = ((4, 3),), + block_sizes: tuple[int, ...] = (), ) -> None: """Check supported formats and the custom backward's identity STE.""" # Numerical helpers load while QuantizeConfig initializes; defer the quantizer import. @@ -44,7 +47,13 @@ def validate_gdn_quantizer( 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.block_sizes is None + or ( + quantizer.num_bits == 8 + and quantizer.block_sizes in ({-1: size} for size in block_sizes) + ) + ) and quantizer.fake_quant and quantizer._pass_through_bwd and not quantizer.rotate_is_enabled @@ -55,11 +64,29 @@ def validate_gdn_quantizer( ): 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 " + f"INT8 block_sizes in {block_sizes} (or no blocks), " + "pass_through_bwd=True, no rotation, pre-scaling, bias, constant " "amax, or custom backend. Other gradient rules and formats are not implemented." ) +def state_quantizer_config( + quantizer: TensorQuantizer, *, name="state_quantizer" +) -> tuple[str, int]: + """Validate the state quantizer and derive its format and last-axis group size. + + A zero group size retains the legacy policy's full-key state tiles. + """ + validate_gdn_quantizer( + quantizer, name=name, num_bits=tuple(_STATE_FORMATS), block_sizes=(16, 32, 64) + ) + if quantizer.block_sizes is not None: + return _STATE_FORMATS[quantizer.num_bits], quantizer.block_sizes[-1] + if quantizer.axis != (0, 1): + raise ValueError(f"{name} supports only axis=(0, 1) with state.block_v tiling") + return _STATE_FORMATS[quantizer.num_bits], 0 + + 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]") @@ -114,8 +141,15 @@ def _fp8_quantize(value: torch.Tensor, axis): return quantized, torch.div(448.0, safe_amax).reciprocal() -def _state_qdq(state: torch.Tensor, block_v: int = 64, state_format: str = "fp8_e4m3"): +def _state_qdq( + state: torch.Tensor, + block_v: int = 64, + state_format: str = "fp8_e4m3", + state_quantizer: TensorQuantizer | None = None, +): """Dynamic state-tile QDQ with detached scales and identity STE.""" + if state_quantizer is not None and state_quantizer.block_sizes is not None: + return state_quantizer(state) 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): diff --git a/modelopt/torch/quantization/plugins/gdn.py b/modelopt/torch/quantization/plugins/gdn.py index 049b16acd66..c8f01773eee 100644 --- a/modelopt/torch/quantization/plugins/gdn.py +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -33,8 +33,6 @@ import torch -from modelopt.torch.kernels.quantization.linear_attention import STATE_QDQ_FORMATS, STATE_QDQ_OFF - from ..linear_attention.prefill import matmul_gdn from .linear_attention import _LinearAttentionQuantMixin @@ -79,7 +77,7 @@ class GatedDeltaNetStateQuantMixin(_LinearAttentionQuantMixin): @property def gdn_state_qdq_block_v(self) -> int: - """Value-column scale grouping from the saved execution policy.""" + """Execution tile width; also sets grouping for legacy tile quantizers.""" return self.linear_attention_config.state.block_v def _state_quantized_chunk_gated_delta_rule( @@ -87,7 +85,6 @@ def _state_quantized_chunk_gated_delta_rule( ) -> 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) @@ -107,8 +104,7 @@ def _state_quantized_chunk_gated_delta_rule( return matmul_gdn( *args, policy=self.linear_attention_config, - state_qdq=quantize_state, - state_format=self._linear_attn_state_format, + state_quantizer=self.gdn_state_quantizer, chunk_size=chunk_size, prefill_lengths=self._linear_attention_prefill_lengths, **kwargs, @@ -116,9 +112,7 @@ def _state_quantized_chunk_gated_delta_rule( return _state_qdq_chunk_gated_delta_rule()( *args, chunk_size=chunk_size, - state_qdq=STATE_QDQ_FORMATS[self._linear_attn_state_format] - if quantize_state - else STATE_QDQ_OFF, + state_quantizer=self.gdn_state_quantizer, 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 index 92f6e88047c..d98aeff4698 100644 --- a/modelopt/torch/quantization/plugins/kda.py +++ b/modelopt/torch/quantization/plugins/kda.py @@ -50,8 +50,7 @@ def _state_quantized_chunk_kda(self, kernel, *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, + state_quantizer=self.kda_state_quantizer, prefill_lengths=self._linear_attention_prefill_lengths, **kwargs, ) diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py index 036bec66f10..1139255fc7a 100644 --- a/modelopt/torch/quantization/plugins/linear_attention.py +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -17,13 +17,11 @@ from ..config import QuantizerAttributeConfig from ..linear_attention.config import LinearAttentionConfig -from ..linear_attention.utils import validate_gdn_quantizer +from ..linear_attention.utils import state_quantizer_config, validate_gdn_quantizer from ..nn import QuantModule, TensorQuantizer __all__ = [] -_STATE_FORMATS: dict[int | tuple[int, int], str] = {(4, 3): "fp8_e4m3", 8: "int8"} - class _LinearAttentionQuantMixin(QuantModule): linear_attention_quantizer_names = ("gdn_state_quantizer", "gdn_w_quantizer") @@ -45,7 +43,7 @@ def _linear_attn_state_format(self): # Disabled quantizers impose no format requirement on the numerical backend. if not self._linear_attn_state.is_enabled: return "fp8_e4m3" - return _STATE_FORMATS[self._linear_attn_state.num_bits] + return state_quantizer_config(self._linear_attn_state)[0] @property def _linear_attn_w(self): @@ -63,20 +61,16 @@ def linear_attention_is_enabled(self): def validate_linear_attention(self): """Validate quantizer contracts shared by GDN and KDA.""" if self._linear_attn_state.is_enabled: - validate_gdn_quantizer( + state_format, group_size = state_quantizer_config( self._linear_attn_state, name=self.linear_attention_quantizer_names[0], - num_bits=tuple(_STATE_FORMATS), ) - 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": + if state_format != "int8": raise ValueError("int8_hadamard32 requires INT8 state quantization") + if group_size: + raise ValueError("TensorQuantizer block_sizes requires state_codec='tile'") if self._linear_attn_w.is_enabled: validate_gdn_quantizer( self._linear_attn_w, name=self.linear_attention_quantizer_names[1] diff --git a/tests/_test_utils/torch/quantization/linear_attention_reference.py b/tests/_test_utils/torch/quantization/linear_attention_reference.py index ad43a2226e1..1e31c44838f 100644 --- a/tests/_test_utils/torch/quantization/linear_attention_reference.py +++ b/tests/_test_utils/torch/quantization/linear_attention_reference.py @@ -161,6 +161,7 @@ def chunk_gdn_reference( state_qdq_block_v: int = 64, state_format: str = "fp8_e4m3", w_quantizer: Callable[[torch.Tensor], torch.Tensor] | None = None, + state_quantizer=None, ) -> tuple[torch.Tensor, torch.Tensor]: """Exact GDN chunk algebra with optional state/W fake quantization. @@ -168,6 +169,15 @@ def chunk_gdn_reference( 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. """ + if state_quantizer is not None: + state_qdq = state_quantizer.is_enabled and state_quantizer._if_quant + state_format = "int8" if state_quantizer.num_bits == 8 else "fp8_e4m3" + + def quantize_state(state): + if state_quantizer is not None and state_quantizer.block_sizes is not None: + return state_quantizer(state) + return state_qdq_reference(state, state_qdq_block_v, state_format) + 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) @@ -199,7 +209,7 @@ def chunk_gdn_reference( if n != previous_n: state = states[n] if state_qdq: - state = state_qdq_reference(state, state_qdq_block_v, state_format) + state = quantize_state(state) previous_n = n length = qc.shape[1] wc = w[offset : offset + length].transpose(0, 1) @@ -213,7 +223,7 @@ def chunk_gdn_reference( state = state * gc[..., -1].exp()[:, None, None] state = state + weighted_keys.transpose(-1, -2) @ updated_values # state_update if state_qdq: - state = state_qdq_reference(state, state_qdq_block_v, state_format) + state = quantize_state(state) if len(finals) <= n: finals.append(state) else: 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 index 53776affca1..0732c97eeb3 100644 --- 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 @@ -61,7 +61,9 @@ def compare(actual, expected, tolerance): ) -@pytest.fixture(scope="module", params=["disabled", "w", "state-w", "state-int8"]) +@pytest.fixture( + scope="module", params=["disabled", "w", "state-w", "state-int8", "state-int8-block"] +) def compiled_gdn_case(request): """Compile only the selected BF16 forward/backward path, outside the test-call timer.""" state_qdq = {"state-w": 1, "state-int8": 2}.get(request.param, 0) @@ -74,6 +76,14 @@ def compiled_gdn_case(request): ) args, state = make_inputs() kwargs = {"state_qdq": state_qdq, "w_quantizer": quantizer} + if request.param == "state-int8-block": + kwargs["state_quantizer"] = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=8, type="dynamic", block_sizes={-1: 16}, narrow_range=True + ) + ) + # Compile the reference quantizer's CUDA extension outside the test-call timer too. + kwargs["state_quantizer"](state[0]) values_and_grads(chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs) torch.cuda.synchronize() return args, state, kwargs diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index a24d1cc4123..816195adf32 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -85,6 +85,7 @@ def quant_cfg(state=True, w=False): {"num_bits": (4, 3), "type": "dynamic"}, # per tensor {"num_bits": 8, "axis": (0, 1), "type": "dynamic"}, # int8 {"num_bits": 4, "axis": (0, 1), "type": "dynamic"}, # unsupported format + {"num_bits": 8, "narrow_range": True, "type": "dynamic", "block_sizes": {-1: 8}}, {"num_bits": (4, 3), "type": "dynamic", "block_sizes": {-1: 16}}, # blockwise ], ) @@ -130,12 +131,10 @@ def test_disabled_state_quantizer_calls_original_kernel(): @pytest.mark.parametrize( - ("partial_kernel", "num_bits", "state_qdq"), - [(False, (4, 3), 1), (True, (4, 3), 1), (False, 8, 2)], + ("partial_kernel", "num_bits"), + [(False, (4, 3)), (True, (4, 3)), (False, 8)], ) -def test_enabled_state_quantizer_uses_state_qdq_kernel( - monkeypatch, partial_kernel, num_bits, state_qdq -): +def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, partial_kernel, num_bits): calls = [] def fake_state_qdq_kernel(*args, **kwargs): @@ -154,7 +153,7 @@ def fake_state_qdq_kernel(*args, **kwargs): model(x) assert calls and calls[-1] == { "chunk_size": 64, - "state_qdq": state_qdq, + "state_quantizer": model.gdn_state_quantizer, "state_qdq_block_v": 64, "w_quantizer": None, **({"output_final_state": True} if partial_kernel else {}), @@ -184,7 +183,7 @@ def fake_state_qdq_kernel(*args, **kwargs): 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]["state_quantizer"] is model.gdn_state_quantizer 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, @@ -220,8 +219,12 @@ def test_w_grouping_is_preserved_during_conversion(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 = deepcopy(quant_cfg(state=state, w=w)) cfg["algorithm"] = None + if state and not w: + cfg["quant_cfg"][-1]["cfg"].update( + num_bits=8, axis=None, block_sizes={-1: 32}, unsigned=False, narrow_range=True + ) cfg["linear_attention"] = [ {"module_name": "*", "cfg": {"state": {"block_v": 128}}}, {"module_name": "0", "cfg": {"state": {"block_v": 32}}}, @@ -229,6 +232,9 @@ def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): mtq.quantize(model, cfg) assert model[0].gdn_state_qdq_block_v == 32 model[0].linear_attention_config.state.block_v = 16 + if state and not w: + sample = torch.randn(2, 4, 19) + expected = model[0].gdn_state_quantizer(sample) if w: # Save the actual quantizer settings, including edits after conversion. model[0].gdn_w_quantizer.axis = None @@ -244,8 +250,11 @@ def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): assert quantizer.is_enabled == enabled assert quantizer.axis == original.axis assert quantizer.num_bits == original.num_bits + assert quantizer.block_sizes == original.block_sizes assert quantizer._dynamic == original._dynamic assert not hasattr(restored[1], name) + if state and not w: + torch.testing.assert_close(restored[0].gdn_state_quantizer(sample), expected) def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py index b9e81c3359b..7e37b6d3548 100644 --- a/tests/unit/torch/quantization/test_linear_attention_decode.py +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -74,6 +74,50 @@ def _values_and_grads(output, state, args, initial): return output, state, *torch.autograd.grad(loss, (*args, initial), retain_graph=True) +@pytest.mark.parametrize(("kda", "block_size"), [(False, 32), (True, 64)]) +def test_block_state_quantizer_prefill_decode_and_gradients(kda, block_size): + args, initial = _inputs(kda, length=5) + # Partial value groups exercise TensorQuantizer padding as well as per-key scales. + quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=8, type="dynamic", block_sizes={-1: block_size}, narrow_range=True + ) + ) + policy = LinearAttentionConfig( + backend="matmul", decode={"prefill_state_qdq": True}, state={"block_v": 16} + ) + function = matmul_kda if kda else matmul_gdn + output, final = function( + *(x.unsqueeze(0) for x in args), + policy=policy, + state_quantizer=quantizer, + prefill_lengths=[3], + initial_state=initial.unsqueeze(0), + output_final_state=True, + ) + # Sequential oracle: the prefix reads unrounded working states; decode reads stored states. + state = quantizer(initial) + expected = [] + q, k, v, g, beta = args + for t in range(len(q)): + if t == 3: + state = quantizer(state) + decay = g[t].exp().unsqueeze(-1) if kda else g[t].exp()[:, None, None] + decayed = state * decay + residual = v[t] - (k[t].unsqueeze(-1) * decayed).sum(-2) + state = decayed + k[t].unsqueeze(-1) * (beta[t].unsqueeze(-1) * residual).unsqueeze(-2) + if t >= 3: + state = quantizer(state) + expected.append((q[t].unsqueeze(-1) * state).sum(-2) / q.shape[-1] ** 0.5) + if t == 2: + state = quantizer(state) + for actual, expected in zip( + _values_and_grads(output[0], final[0], args, initial), + _values_and_grads(torch.stack(expected), state, args, initial), + ): + torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-7) + + @pytest.mark.parametrize(("kda", "mode"), [(False, "token"), (True, "replay")]) def test_prefill_decode_matches_recurrence_and_gradients(kda, mode): args, state = _inputs(kda) From 9cf28aff79d8ef4bab2de6a7c84f31986307961a Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 3 Oct 2026 22:38:38 -0700 Subject: [PATCH 08/27] Simplify linear attention prefill and decode plumbing Make recurrent_decode the single Torch implementation. Process each prepared prefix directly and remove duplicate packing, shape preparation, and the obsolete helper without changing supported numerical policies. Validated 40 focused CPU tests, the existing Bridge QAT/QAD GPU smoke tests, and 12 before/after output, state, and gradient comparisons. Signed-off-by: Kai Xu --- .../linear_attention/_chunk_prefill.py | 156 ++++++++---------- .../quantization/linear_attention/decode.py | 41 +---- .../linear_attention/decode_prefill.py | 49 ++---- .../quantization/linear_attention/utils.py | 40 ----- .../test_linear_attention_decode.py | 14 +- .../test_linear_attention_hadamard.py | 12 +- 6 files changed, 94 insertions(+), 218 deletions(-) diff --git a/modelopt/torch/quantization/linear_attention/_chunk_prefill.py b/modelopt/torch/quantization/linear_attention/_chunk_prefill.py index c4dfc48dbd1..b38aab5e482 100644 --- a/modelopt/torch/quantization/linear_attention/_chunk_prefill.py +++ b/modelopt/torch/quantization/linear_attention/_chunk_prefill.py @@ -17,7 +17,7 @@ import torch -from .utils import _prepare, _state_qdq +from .utils import _state_qdq __all__ = [] @@ -31,53 +31,42 @@ def chunk_gdn( *, 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, + initial_state: torch.Tensor, state_qdq: bool = False, state_qdq_block_v: int = 64, state_format: str = "fp8_e4m3", state_quantizer=None, ) -> 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) + """Compute one prepared, nonempty GDN prefix [T,H,D] with a key-first state.""" 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] + state = initial_state + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format, state_quantizer) + pieces = [] + for lo in range(0, len(q), chunk_size): + hi = min(lo + chunk_size, len(q)) + qc, kc, vc = (x[lo:hi].transpose(0, 1) for x in (q, k, v)) + gc = g[lo:hi].transpose(0, 1).cumsum(-1) + bc = beta[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, state_quantizer) - 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, state_quantizer) - 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 + return torch.cat(pieces), state def chunk_kda( @@ -89,60 +78,49 @@ def chunk_kda( *, chunk_size=64, scale=None, - initial_state=None, - cu_seqlens=None, - state_v_first=False, + initial_state, state_qdq=False, state_qdq_block_v=64, state_format="fp8_e4m3", state_quantizer=None, ): - """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) + """Compute one prepared, nonempty KDA prefix [T,H,D] with a key-first state.""" 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] + state = initial_state + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format, state_quantizer) + pieces = [] + for lo in range(0, len(q), chunk_size): + hi = min(lo + chunk_size, len(q)) + qc, kc, vc = (x[lo:hi].transpose(0, 1) for x in (q, k, v)) + prefix = g[lo:hi].transpose(0, 1).cumsum(-2) + bc = beta[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, state_quantizer) - 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, state_quantizer) - 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 + return torch.cat(pieces), state diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py index e15dd8a14ea..45c681c5ba4 100644 --- a/modelopt/torch/quantization/linear_attention/decode.py +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Explicit token-state and encoded-update replay references for QAT.""" +"""Differentiable token-state and encoded-update replay for QAT.""" from dataclasses import dataclass @@ -27,7 +27,6 @@ "LinearAttentionCarry", "ReplayEntry", "recurrent_decode", - "recurrent_decode_reference", ] @@ -247,7 +246,7 @@ def _prepare_carry( return carry, signature -def recurrent_decode_reference( +def recurrent_decode( q, k, v, @@ -355,39 +354,3 @@ def recurrent_decode_reference( 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", - state_quantizer=None, - block_v=64, - initial_state=None, - carry=None, - position=0, - scale=None, -): - """Run the Torch training implementation with explicit differentiable carry.""" - return recurrent_decode_reference( - q, - k, - v, - g, - beta, - config=config, - state_qdq=state_qdq, - state_format=state_format, - state_quantizer=state_quantizer, - block_v=block_v, - initial_state=initial_state, - carry=carry, - position=position, - scale=scale, - ) diff --git a/modelopt/torch/quantization/linear_attention/decode_prefill.py b/modelopt/torch/quantization/linear_attention/decode_prefill.py index 28dbfa4466b..8544599e1e9 100644 --- a/modelopt/torch/quantization/linear_attention/decode_prefill.py +++ b/modelopt/torch/quantization/linear_attention/decode_prefill.py @@ -128,41 +128,23 @@ def _decode_prefill( 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, - state_quantizer=state_quantizer, - 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] + prefix_fn = chunk_kda if g.ndim == 4 else chunk_gdn outputs, finals = [], [] for n, (b, start, end) in enumerate(sequences): split = start + prefixes[n] + prefix, state = q.new_empty(0, heads, values), states[n] + if prefixes[n]: + with torch.autocast(device_type=q.device.type, enabled=False): + prefix, state = prefix_fn( + *(x[b, start:split] for x in (q, k, v, g, beta)), + state_qdq=state_qdq and policy.decode.prefill_state_qdq, + state_format=state_format, + state_quantizer=state_quantizer, + scale=scale, + initial_state=state, + chunk_size=policy.chunk_size, + state_qdq_block_v=policy.state.block_v, + ) suffix, carry = recurrent_decode( *(x[b, split:end] for x in (q, k, v, g, beta)), config=policy.decode, @@ -170,11 +152,10 @@ def _decode_prefill( state_format=state_format, state_quantizer=state_quantizer, block_v=policy.state.block_v, - initial_state=prefix_states.get(n, states[n]), + initial_state=state, 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) diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py index d51bb76101b..af41f4ed146 100644 --- a/modelopt/torch/quantization/linear_attention/utils.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -17,7 +17,6 @@ from __future__ import annotations -from itertools import pairwise from typing import TYPE_CHECKING import torch @@ -87,45 +86,6 @@ def state_quantizer_config( return _STATE_FORMATS[quantizer.num_bits], 0 -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. diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py index 7e37b6d3548..953e344dfcd 100644 --- a/tests/unit/torch/quantization/test_linear_attention_decode.py +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -24,7 +24,7 @@ LinearAttentionDecodeConfig, matmul_gdn, matmul_kda, - recurrent_decode_reference, + recurrent_decode, ) from modelopt.torch.quantization.linear_attention.decode import _encode from modelopt.torch.quantization.nn import TensorQuantizer @@ -150,9 +150,9 @@ def test_quantized_replay_carry_continuation_and_gradients(): args, state = _inputs() cfg = LinearAttentionDecodeConfig(mode="replay", replay={"window": 3}) kwargs = {"config": cfg, "state_qdq": True, "block_v": 16} - expected, final = recurrent_decode_reference(*args, initial_state=state, **kwargs) - first, carry = recurrent_decode_reference(*(x[:5] for x in args), initial_state=state, **kwargs) - second, carry = recurrent_decode_reference(*(x[5:] for x in args), carry=carry, **kwargs) + expected, final = recurrent_decode(*args, initial_state=state, **kwargs) + first, carry = recurrent_decode(*(x[:5] for x in args), initial_state=state, **kwargs) + second, carry = recurrent_decode(*(x[5:] for x in args), carry=carry, **kwargs) for actual, expected in zip( _values_and_grads(torch.cat((first, second)), carry.reconstruct(), args, state), _values_and_grads(expected, final.reconstruct(), args, state), @@ -165,14 +165,12 @@ def test_quantized_replay_carry_continuation_and_gradients(): def test_grid_gate_ste_keeps_gate_gradients_and_changes_trajectory(): args, state = _inputs() - output, carry = recurrent_decode_reference( + output, carry = recurrent_decode( *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 - ) + exact, _ = recurrent_decode(*args, config=LinearAttentionDecodeConfig(), initial_state=state) assert not torch.equal(output, exact) diff --git a/tests/unit/torch/quantization/test_linear_attention_hadamard.py b/tests/unit/torch/quantization/test_linear_attention_hadamard.py index c9020e4a67e..905db03e9c5 100644 --- a/tests/unit/torch/quantization/test_linear_attention_hadamard.py +++ b/tests/unit/torch/quantization/test_linear_attention_hadamard.py @@ -18,7 +18,7 @@ from modelopt.torch.quantization.linear_attention import ( LinearAttentionDecodeConfig, - recurrent_decode_reference, + recurrent_decode, ) @@ -84,16 +84,12 @@ def test_hadamard_replay_matches_dense_oracle_and_split_carry(): replay={"window": 3, "factor_qdq": False}, ) kwargs = {"config": cfg, "state_format": "int8", "state_qdq": True} - output, carry = recurrent_decode_reference(*args, initial_state=initial, **kwargs) + output, carry = recurrent_decode(*args, initial_state=initial, **kwargs) _check_with_grads((output, carry.reconstruct()), _oracle(args, initial), args, initial) assert carry.value_basis == "hadamard32" and carry.anchor.block_v == 32 assert carry.anchor.scales.shape == (2, 4, 2) - first, split_carry = recurrent_decode_reference( - *(x[:2] for x in args), initial_state=initial, **kwargs - ) - second, split_carry = recurrent_decode_reference( - *(x[2:] for x in args), carry=split_carry, **kwargs - ) + first, split_carry = recurrent_decode(*(x[:2] for x in args), initial_state=initial, **kwargs) + second, split_carry = recurrent_decode(*(x[2:] for x in args), carry=split_carry, **kwargs) _check_with_grads( (torch.cat((first, second)), split_carry.reconstruct()), (output, carry.reconstruct()), From ee689559de95a571d088a7bd450dfb4d22deea40 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sun, 4 Oct 2026 10:19:25 -0700 Subject: [PATCH 09/27] Remove redundant linear attention quantization scaffolding Signed-off-by: Kai Xu --- .../configs/decode_state_int8.json | 1 - modelopt/torch/quantization/conversion.py | 4 ++ .../quantization/linear_attention/config.py | 30 ++++++++++++-- .../quantization/linear_attention/decode.py | 24 ++---------- .../quantization/linear_attention/utils.py | 39 ++++++++++--------- modelopt/torch/quantization/plugins/gdn.py | 11 ++++++ modelopt/torch/quantization/plugins/kda.py | 11 +++++- .../quantization/plugins/linear_attention.py | 24 ++---------- .../torch/quantization/plugins/megatron.py | 4 +- .../units/default_disabled_quantizers.yaml | 2 - .../linear_attention_reference.py | 10 ----- .../plugins/test_megatron_gated_delta_net.py | 2 +- .../quantization/plugins/test_megatron_kda.py | 6 +++ .../torch/quantization/plugins/test_gdn.py | 7 ++++ .../test_linear_attention_reference.py | 6 +-- 15 files changed, 99 insertions(+), 82 deletions(-) diff --git a/examples/llm_qat/linear_attention/configs/decode_state_int8.json b/examples/llm_qat/linear_attention/configs/decode_state_int8.json index aad890b22b0..736e212b338 100644 --- a/examples/llm_qat/linear_attention/configs/decode_state_int8.json +++ b/examples/llm_qat/linear_attention/configs/decode_state_int8.json @@ -38,7 +38,6 @@ "cfg": { "backend": "matmul", "decode": { - "implementation": "torch", "state_codec": "int8_hadamard32", "prefill_state_qdq": false } diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index db194a0f465..df24856ceca 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -146,7 +146,11 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: # QuantModule.set_extra_state(). return model + # Optional plugins import conversion during initialization, so defer this import. + from .plugins.kda import _discard_legacy_w_quantizer_state + quantizer_state_dict = dict(metadata["quantizer_state"]) + _discard_legacy_w_quantizer_state(quantizer_state_dict) 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(): diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index 4d3cc108fd5..2168c870bfb 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -30,12 +30,21 @@ 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) + + @model_validator(mode="before") + @classmethod + def _load_legacy_defaults(cls, values): + if isinstance(values, dict): + values = dict(values) + for key, default in (("mode", "chunk"), ("quantize_initial", True)): + if values.get(key) == default: + values.pop(key) + return values class _SolveConfig(ModeloptBaseConfig): + # Retained solely to deserialize old pickled policies and validate their fixed setting. method: Literal["exact"] = ModeloptField(default="exact") @@ -51,7 +60,6 @@ class LinearAttentionDecodeConfig(ModeloptBaseConfig): """Explicit suffix recurrence; workload supplies per-sequence prefix lengths.""" mode: Literal["token", "replay"] = ModeloptField(default="token") - implementation: Literal["torch"] = ModeloptField(default="torch") readout: Literal["working", "stored"] = ModeloptField(default="stored") quantize_initial: bool = ModeloptField(default=True) prefill_state_qdq: bool = ModeloptField(default=False) @@ -59,6 +67,13 @@ class LinearAttentionDecodeConfig(ModeloptBaseConfig): decay_log_step: float | None = Field(default=None, gt=0, allow_inf_nan=False) replay: LinearAttentionReplayConfig | None = ModeloptField(default=None) + @model_validator(mode="before") + @classmethod + def _load_legacy_defaults(cls, values): + if isinstance(values, dict) and values.get("implementation") == "torch": + values = {key: value for key, value in values.items() if key != "implementation"} + return values + @model_validator(mode="after") def _validate_replay(self): if (self.mode == "replay") != (self.replay is not None): @@ -86,9 +101,16 @@ class LinearAttentionConfig(ModeloptBaseConfig): 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="before") + @classmethod + def _load_legacy_defaults(cls, values): + if isinstance(values, dict) and "solve" in values: + _SolveConfig.model_validate(values["solve"]) + values = {key: value for key, value in values.items() if key != "solve"} + return values + @model_validator(mode="after") def _validate_decode_backend(self): if (self.backend == "matmul") != (self.decode is not None): diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py index 45c681c5ba4..eb4b6fa3779 100644 --- a/modelopt/torch/quantization/linear_attention/decode.py +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -20,7 +20,7 @@ import torch from .config import LinearAttentionDecodeConfig -from .utils import _fp8_quantize, state_quantizer_config +from .utils import _tile_qdq, state_quantizer_config __all__ = [ "EncodedLinearAttentionTensor", @@ -129,31 +129,15 @@ def _encode( 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) + decoded, scales = _tile_qdq(value, block_v, state_format, state=state) if state_format == "int8": decoded = _IdentityGradient.apply(value, decoded) - return EncodedLinearAttentionTensor(decoded, torch.stack(scales, dim=-1), state_format, block_v) + return EncodedLinearAttentionTensor(decoded, scales, state_format, block_v) def _signature(config, state_qdq, block_v, state_format, state_quantizer): signature = ( - config.model_dump_json(exclude={"implementation", "prefill_state_qdq"}) + config.model_dump_json(exclude={"prefill_state_qdq"}) + f"/{state_qdq}/{block_v}/{state_format}" ) if state_quantizer is not None and state_quantizer.block_sizes is not None: diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py index af41f4ed146..2c650877cf5 100644 --- a/modelopt/torch/quantization/linear_attention/utils.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -114,23 +114,26 @@ def _state_qdq( 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") + quantized, _ = _tile_qdq(state, block_v, state_format, state=True) 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 quantized return state + (quantized - state).detach() + + +def _tile_qdq(value, block_v, state_format, *, state=False): + """Return tile-rounded values and detached scales; INT8 callers supply their STE.""" + rounded, scales = [], [] + for part in value.split(block_v, dim=-1): + tensor = part.flatten(-2) if state else part + if state_format == "fp8_e4m3": + axis = tuple(range(tensor.ndim - 1)) or None + decoded, scale = _fp8_quantize(tensor, axis) + else: + with torch.no_grad(): + tensor = tensor.float() + amax = tensor.abs().amax(dim=-1, keepdim=True) + scale = torch.where(amax > 0, amax / 127.0, torch.ones_like(amax)) + decoded = (tensor / scale).round().clamp(-127, 127) * scale + rounded.append(decoded.reshape_as(part)) + scales.append(scale.squeeze(-1)) + return torch.cat(rounded, dim=-1).to(value.dtype), torch.stack(scales, dim=-1) diff --git a/modelopt/torch/quantization/plugins/gdn.py b/modelopt/torch/quantization/plugins/gdn.py index c8f01773eee..f82848b1d0e 100644 --- a/modelopt/torch/quantization/plugins/gdn.py +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -34,6 +34,7 @@ import torch from ..linear_attention.prefill import matmul_gdn +from ..linear_attention.utils import validate_gdn_quantizer from .linear_attention import _LinearAttentionQuantMixin __all__ = ["GatedDeltaNetStateQuantMixin"] @@ -75,6 +76,16 @@ class GatedDeltaNetStateQuantMixin(_LinearAttentionQuantMixin): Both sites require identity STE. The execution policy is saved in ModelOpt metadata. """ + linear_attention_quantizer_names = ("gdn_state_quantizer", "gdn_w_quantizer") + + def validate_linear_attention(self): + """Validate state and GDN WY operand quantization.""" + super().validate_linear_attention() + if self.gdn_w_quantizer.is_enabled: + validate_gdn_quantizer(self.gdn_w_quantizer, name="gdn_w_quantizer") + if self.linear_attention_config.decode is not None: + raise ValueError("Decode's exact prefix does not support WY operand QDQ") + @property def gdn_state_qdq_block_v(self) -> int: """Execution tile width; also sets grouping for legacy tile quantizers.""" diff --git a/modelopt/torch/quantization/plugins/kda.py b/modelopt/torch/quantization/plugins/kda.py index d98aeff4698..c459e7fccce 100644 --- a/modelopt/torch/quantization/plugins/kda.py +++ b/modelopt/torch/quantization/plugins/kda.py @@ -23,10 +23,19 @@ __all__ = ["KimiDeltaAttentionStateQuantMixin"] +def _discard_legacy_w_quantizer_state(quantizer_state): + # Early KDA checkpoints included a disabled placeholder; WY QDQ was never supported. + for name in list(quantizer_state): + if name.rsplit(".", 1)[-1] == "kda_w_quantizer": + if not quantizer_state[name].get("_disabled", False): + raise ValueError("KDA checkpoints with enabled kda_w_quantizer are unsupported") + del quantizer_state[name] + + class KimiDeltaAttentionStateQuantMixin(_LinearAttentionQuantMixin): """Adds state quantizers and decode-aware kernel routing to Kimi Delta Attention.""" - linear_attention_quantizer_names = ("kda_state_quantizer", "kda_w_quantizer") + linear_attention_quantizer_names = ("kda_state_quantizer",) def validate_linear_attention(self): """Require the materialized backend for KDA numerical emulation.""" diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py index 1139255fc7a..1c6b823a4a0 100644 --- a/modelopt/torch/quantization/plugins/linear_attention.py +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -17,14 +17,14 @@ from ..config import QuantizerAttributeConfig from ..linear_attention.config import LinearAttentionConfig -from ..linear_attention.utils import state_quantizer_config, validate_gdn_quantizer +from ..linear_attention.utils import state_quantizer_config from ..nn import QuantModule, TensorQuantizer __all__ = [] class _LinearAttentionQuantMixin(QuantModule): - linear_attention_quantizer_names = ("gdn_state_quantizer", "gdn_w_quantizer") + linear_attention_quantizer_names: tuple[str, ...] = () def _setup(self): for name in self.linear_attention_quantizer_names: @@ -38,23 +38,11 @@ def _setup(self): def _linear_attn_state(self): return getattr(self, self.linear_attention_quantizer_names[0]) - @property - def _linear_attn_state_format(self): - # Disabled quantizers impose no format requirement on the numerical backend. - if not self._linear_attn_state.is_enabled: - return "fp8_e4m3" - return state_quantizer_config(self._linear_attn_state)[0] - - @property - def _linear_attn_w(self): - return getattr(self, self.linear_attention_quantizer_names[1]) - @property def linear_attention_is_enabled(self): """Whether an operand, state, or arithmetic policy changes the computation.""" return ( - self._linear_attn_state.is_enabled - or self._linear_attn_w.is_enabled + any(getattr(self, name).is_enabled for name in self.linear_attention_quantizer_names) or self.linear_attention_config.decode is not None ) @@ -71,12 +59,6 @@ def validate_linear_attention(self): raise ValueError("int8_hadamard32 requires INT8 state quantization") if group_size: raise ValueError("TensorQuantizer block_sizes requires state_codec='tile'") - 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.""" diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index c36b08944a4..3fded293ab0 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -71,7 +71,7 @@ from ..utils.layerwise_calib import LayerActivationCollector from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear from .gdn import GatedDeltaNetStateQuantMixin -from .kda import KimiDeltaAttentionStateQuantMixin +from .kda import KimiDeltaAttentionStateQuantMixin, _discard_legacy_w_quantizer_state from .linear_attention import _LinearAttentionQuantMixin from .transformer_engine import _QuantTEGroupedLinear, _QuantTELayerNormLinear, _QuantTELinear @@ -300,6 +300,8 @@ def quant_module_set_extra_state(self, state: Any): quantizer_state = state.get("modelopt_quantizer_state", None) if quantizer_state is not None: + quantizer_state = dict(quantizer_state) + _discard_legacy_w_quantizer_state(quantizer_state) for name, module in self.named_modules(): if isinstance(module, TensorQuantizer): quantizer_substate = quantizer_state.get(name) diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index e409bff7aa6..8618024222f 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -23,8 +23,6 @@ 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/tests/_test_utils/torch/quantization/linear_attention_reference.py b/tests/_test_utils/torch/quantization/linear_attention_reference.py index 1e31c44838f..d0cfc70989b 100644 --- a/tests/_test_utils/torch/quantization/linear_attention_reference.py +++ b/tests/_test_utils/torch/quantization/linear_attention_reference.py @@ -28,20 +28,10 @@ __all__ = [ "chunk_gdn_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"): diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index ea5f96865eb..123fa1f68c5 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -191,7 +191,7 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): 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_quantizer.num_bits == original.gdn_state_quantizer.num_bits 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() diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py index b93f3e81967..6b243289bbf 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py @@ -33,10 +33,12 @@ restore_sharded_modelopt_state, save_sharded_modelopt_state, ) +from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.linear_attention import ( LinearAttentionConfig, linear_attention_training_phase, ) +from modelopt.torch.quantization.nn import TensorQuantizer KimiDeltaAttention = pytest.importorskip("megatron.core.ssm.gated_delta_net.kda").KimiDeltaAttention @@ -128,8 +130,11 @@ def _test_kda(rank, size, cfg, checkpoint_path): with torch.no_grad(): expected = forward(model) + # Old checkpoints carried a disabled WY placeholder; it must not reappear on restore. + model.kda_w_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) save_distributed_checkpoint(checkpoint_path, model) save_sharded_modelopt_state([model], checkpoint_path) + del model.kda_w_quantizer restored = _layer() restore_sharded_modelopt_state([restored], checkpoint_path) load_distributed_checkpoint(checkpoint_path, restored) @@ -138,6 +143,7 @@ def _test_kda(rank, size, cfg, checkpoint_path): actual = forward(restored) torch.testing.assert_close(actual, expected, rtol=0, atol=0) actual.float().square().mean().backward() + assert not hasattr(restored, "kda_w_quantizer") assert restored._linear_attention_prefill_lengths is None assert restored.linear_attention_config == policy assert restored.gated_delta_rule is kernel diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index 816195adf32..5baf6244d14 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -240,6 +240,13 @@ def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): model[0].gdn_w_quantizer.axis = None path = tmp_path / "gdn.pth" mto.save(model, path) + # Load policies saved before fixed, non-configurable fields were removed. + checkpoint = torch.load(path) + for _, mode_state in checkpoint["modelopt_state"]["modelopt_state_dict"]: + for policy in mode_state["metadata"]["linear_attention"].values(): + policy["state"].update(mode="chunk", quantize_initial=True) + policy["solve"] = {"method": "exact"} + torch.save(checkpoint, path) restored = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) mto.restore(restored, path) assert restored[0].linear_attention_config == model[0].linear_attention_config diff --git a/tests/unit/torch/quantization/test_linear_attention_reference.py b/tests/unit/torch/quantization/test_linear_attention_reference.py index d71818d0163..6b3669c47d8 100644 --- a/tests/unit/torch/quantization/test_linear_attention_reference.py +++ b/tests/unit/torch/quantization/test_linear_attention_reference.py @@ -19,7 +19,7 @@ from _test_utils.torch.quantization.linear_attention_reference import ( chunk_gdn_reference, recurrent_delta_rule_reference, - state_fp8_qdq_reference, + state_qdq_reference, ) @@ -98,11 +98,11 @@ def test_state_qdq_granularity_zero_tail_and_identity_ste(): with torch.no_grad(): state[0].zero_() state[1, ..., :16] *= 10 - quantized = state_fp8_qdq_reference(state, block_v=16) + quantized = state_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)) + assert not torch.equal(quantized, state_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) From 09c085aeafba537176447b1f44444759ce7996fa Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sun, 4 Oct 2026 17:30:27 -0700 Subject: [PATCH 10/27] Separate linear attention training example from core support Signed-off-by: Kai Xu --- examples/llm_qat/README.md | 1 - examples/llm_qat/linear_attention/README.md | 407 ------------------ .../configs/decode_state_int8.json | 47 -- .../llm_qat/linear_attention/requirements.txt | 1 - examples/llm_qat/linear_attention/train.py | 261 ----------- modelopt_recipes/configs/ptq/units/README.md | 2 +- ... linear_attention_state_int8_dynamic.yaml} | 12 +- ... linear_attention_state_int8_dynamic.yaml} | 8 +- modelopt_recipes/ptq.md | 2 +- .../megatron_bridge/test_linear_attention.py | 149 ------- .../quantization/plugins/test_megatron_kda.py | 7 +- 11 files changed, 19 insertions(+), 878 deletions(-) delete mode 100644 examples/llm_qat/linear_attention/README.md delete mode 100644 examples/llm_qat/linear_attention/configs/decode_state_int8.json delete mode 100644 examples/llm_qat/linear_attention/requirements.txt delete mode 100644 examples/llm_qat/linear_attention/train.py rename modelopt_recipes/configs/ptq/units/{gdn_state_int8_dynamic.yaml => linear_attention_state_int8_dynamic.yaml} (73%) rename modelopt_recipes/general/ptq/{gdn_state_int8_dynamic.yaml => linear_attention_state_int8_dynamic.yaml} (83%) delete mode 100644 tests/examples/megatron_bridge/test_linear_attention.py diff --git a/examples/llm_qat/README.md b/examples/llm_qat/README.md index a32c25368d7..cd88ae30498 100644 --- a/examples/llm_qat/README.md +++ b/examples/llm_qat/README.md @@ -13,7 +13,6 @@ For background on QAT and QAD and help choosing between Hugging Face, Megatron B | Arguments | Full CLI/YAML argument reference | \[[Link](ARGUMENTS.md)\] | | 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 | Trainer APIs, 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 deleted file mode 100644 index f4db6058454..00000000000 --- a/examples/llm_qat/linear_attention/README.md +++ /dev/null @@ -1,407 +0,0 @@ -# Quantization-Aware Training and Distillation for Linear Attention - -This example fine-tunes Megatron-Core GDN or KDA attention parameters with -recurrent-state fake quantization using Megatron Bridge's training loop (QAT), or -its distillation loop with an unquantized teacher (QAD). Bridge handles optimization, -gradient accumulation, logging, and checkpoints. Training adapters support Megatron -`GatedDeltaNet` and `KimiDeltaAttention`; FLA supplies kernels, not model-layer adapters. -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 - -Use the [Megatron Bridge environment](../../megatron_bridge/README.md#pre-requisites) -and install the example dependencies. Run commands from the repository root. - -KDA requires a Megatron revision exporting -`megatron.core.ssm.gated_delta_net.KimiDeltaAttention`. It is available on -[Megatron's dev branch](https://github.com/NVIDIA/Megatron-LM/tree/ac100f773f9d8fac918ce3e64698be4ed84fa428); -Megatron-Core 0.19.2 does not include it. Older versions can still use GDN. -Retain `fla-core==0.5.1` for the kernel dependency. ModelOpt registers no FLA -model-layer adapters; Megatron or Bridge may still install the full -`flash-linear-attention` package as their dependency. Use mutually compatible -Bridge/Core revisions: Bridge's pinned Core revision may differ from the KDA -`dev` revision. Native Megatron KDA training does not require Bridge. - -```bash -pip install -r examples/llm_qat/linear_attention/requirements.txt -torchrun --standalone --nproc-per-node=1 examples/llm_qat/linear_attention/train.py \ - --model /path/to/local-model \ - --train-data /path/to/tokenized/train_text_document \ - --output /path/to/megatron-qat-checkpoint \ - --train-steps 1 --length 128 --prefill-tokens 64 -``` - -Add `--teacher-model /path/to/unquantized-model` to run QAD. The teacher and -student must share the tokenizer vocabulary and output vocabulary dimensions. -The teacher stays frozen and unquantized. QAT uses next-token cross-entropy; -QAD uses Bridge's logits distillation loss with the language-model loss disabled. -Both losses are masked to the suffix after `--prefill-tokens`. - -For multiple GPUs, use the same `--tp_size`, `--pp_size`, and `--ep_size` options -as [the Megatron Bridge example](../../megatron_bridge/distill.py). Bridge builds -the process groups, schedules forward/backward, and manages the distributed optimizer. -For example, run QAD with two-way tensor parallelism: - -```bash -torchrun --standalone --nproc-per-node=2 examples/llm_qat/linear_attention/train.py \ - --model /path/to/local-model \ - --teacher-model /path/to/unquantized-model \ - --train-data /path/to/tokenized/train_text_document \ - --output /path/to/megatron-qad-checkpoint \ - --tp_size 2 --pp_size 1 --ep_size 1 \ - --global-batch-size 2 --train-steps 1 --length 128 --prefill-tokens 64 -``` - -Student and teacher use the same topology. Sequence parallelism is enabled with -TP greater than one. With TP/PP/EP all set to one, additional ranks use data -parallelism; the model weights remain replicated on each GPU. The global batch -size must be divisible by the data-parallel size (microbatch size is one). -Choose a topology supported by the model's Bridge/Core implementation. This example -currently requires linear-attention layers on every pipeline stage. Context -parallelism stays fixed at one; GDN/KDA state QAT does not yet support CP. - -Use a local HF model/tokenizer snapshot that your installed Megatron Bridge can -convert to a Megatron model containing GDN or KDA layers, a tokenized Megatron -`.bin`/`.idx` dataset pair, and CUDA GPUs. `--train-data` is the shared filename -prefix without either extension; use the model's tokenizer when -[preparing the data](../../megatron_bridge/README.md#data-preparation). -Bridge's architecture/checkpoint conversion support is a separate requirement -from ModelOpt's layer adapter; an arbitrary -FLA model checkpoint cannot be loaded through this example. For a model already -built in Megatron, use the training-loop integration below. If the local model -requires custom Python code, explicitly pass `--trust-remote-code`. - -The shared [INT8 configuration](configs/decode_state_int8.json) enables GDN and -KDA state quantizers. It 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. Only linear-attention parameters are trained, using -Bridge's BF16 mixed precision and optimizer; other model parameters are frozen. -Every dense sequence has the same fixed prefill boundary. The example attaches -the phase context to local linear-attention layers on each rank and retains it -throughout training so backward recomputation sees that boundary. -Packed sequences and variable per-sequence boundaries require their own batch integration. - -Bridge saves model weights, optimizer/scheduler state, and ModelOpt state under -`/checkpoints`. Reusing `--output` resumes from its latest checkpoint; -increase `--train-steps` to the desired total step count and retain the same model, -teacher, topology, dataset, and prefill boundary. A resumed checkpoint supplies the saved -quantization policy. `--global-batch-size` controls accumulation with microbatch size one. -Export to a supported serving model separately. The current vLLM -state-only plugin accepts its own boundary-QDQ recipe; it does not implement -the Hadamard or replay training policies. - -This example uses Bridge's existing training workflows with state quantization. -The Hugging Face -`QATTrainer`/`QADTrainer` example in [the parent directory](../README.md#using-qattrainer-and-qadtrainer) -targets a different training framework. This 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, state codec, decay approximation, and replay settings. The decode -implementation uses Torch operations and autograd. 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 shared state recipe enables `*gdn_state_quantizer` and `*kda_state_quantizer` -with signed narrow-range INT8, dynamic scales, and `axis=(0, 1)`. Use -`*gdn_state_quantizer` or `*kda_state_quantizer` to select just one layer type. 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/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. Both complete recipes use the Torch 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. - -### Megatron GDN with chunk-only FLA kernels - -For Megatron GDN's existing chunked FLA kernel 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. - -### Per-row blockwise INT8 through TensorQuantizer - -Use the standard quantizer settings to select one dynamic scale per key row and -16, 32, or 64 consecutive value channels. For example, replace the state quantizer -configuration in `gdn_recipe` above with: - -```python -gdn_recipe["quant_cfg"][1]["cfg"] = { - "num_bits": 8, - "type": "dynamic", - "block_sizes": {-1: 32}, - "unsigned": False, - "narrow_range": True, - "pass_through_bwd": True, -} -``` - -The logical state is `[N, H, K, V]`; `-1` denotes V even when a kernel stores the -state transposed. Omit `axis` when using `block_sizes`. The outer `type="dynamic"` -uses ordinary INT8 with dynamic amax; a nested `block_sizes["type"]="dynamic"` -selects a different specialized ModelOpt path and is not supported here. - -The GDN fused kernel applies the quantizer's grouping inside each execution tile. -It can enlarge `state.block_v` to fit a complete quantization group; changing the -execution tile does not change the scale groups. The Torch GDN/KDA path calls the -registered `TensorQuantizer` at the same configured state-write boundaries. -For KDA, use `*kda_state_quantizer` and the materialized backend. For token decode -or replay, explicitly select `state_codec="tile"`; the existing Hadamard recipe -keeps its fixed codec and cannot be combined with this `block_sizes` setting. - -Configurations without `block_sizes` retain their existing `[K, block_v]` scale -grouping, or their fixed Hadamard codec. The default INT8 + Hadamard recipe is -unchanged. This remains floating-storage fake QDQ for QAT/QAD. - -## 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 shifted next-token labels have shape [2, 128]. -# model is a Megatron model. -positions = torch.arange(128, device=ids.device).expand_as(ids) -with linear_attention_training_phase(model, [64, 96]): - with torch.autocast("cuda", dtype=torch.bfloat16): - losses = model( - input_ids=ids, position_ids=positions, attention_mask=None, labels=labels - ) - mask = torch.arange(128, device=ids.device)[None, :] >= torch.tensor( - [64, 96], device=ids.device - )[:, None] - loss = (losses * mask).sum() / mask.sum() - loss.backward() -optimizer.step() -``` - -The shared recipe selects both GDN and KDA state quantizers. 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 implementation uses autograd through initial states, chunk handoff, -token writes, and replay refreshes. The decode recurrence runs token by token in -Python, so long training suffixes can be slow. - -Prefill prefixes use exact chunk algebra with optional state QDQ. This example -has no prefill GEMM QDQ or approximate inverse. Use the Megatron training forward -without a serving inference context. 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/decode_state_int8.json b/examples/llm_qat/linear_attention/configs/decode_state_int8.json deleted file mode 100644 index 736e212b338..00000000000 --- a/examples/llm_qat/linear_attention/configs/decode_state_int8.json +++ /dev/null @@ -1,47 +0,0 @@ -{ - "quant_cfg": [ - { - "quantizer_name": "*", - "enable": false - }, - { - "quantizer_name": "*gdn_state_quantizer", - "cfg": { - "num_bits": 8, - "type": "dynamic", - "axis": [ - 0, - 1 - ], - "unsigned": false, - "narrow_range": true - } - }, - { - "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": { - "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 deleted file mode 100644 index 915edf246cd..00000000000 --- a/examples/llm_qat/linear_attention/requirements.txt +++ /dev/null @@ -1 +0,0 @@ -fla-core==0.5.1 diff --git a/examples/llm_qat/linear_attention/train.py b/examples/llm_qat/linear_attention/train.py deleted file mode 100644 index 5fd681ed656..00000000000 --- a/examples/llm_qat/linear_attention/train.py +++ /dev/null @@ -1,261 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Run Megatron Bridge QAT or QAD with GDN/KDA recurrent-state fake quantization.""" - -import argparse -import json -from contextlib import ExitStack -from pathlib import Path - -import torch -from megatron.bridge import AutoBridge -from megatron.bridge.models.distillation_provider import convert_to_distillation_provider -from megatron.bridge.training.config import ( - CheckpointConfig, - ConfigContainer, - DistributedDataParallelConfig, - GPTDatasetConfig, - LoggerConfig, - OptimizerConfig, - RNGConfig, - SchedulerConfig, - TokenizerConfig, - TrainingConfig, - ValidationConfig, -) -from megatron.bridge.training.distill import distill -from megatron.bridge.training.gpt_step import forward_step_modelopt -from megatron.bridge.training.post_training.checkpointing import ( - has_modelopt_state, - load_modelopt_state, -) -from megatron.bridge.training.post_training.distillation import ModelOptDistillConfig -from megatron.bridge.training.pretrain import pretrain -from megatron.bridge.training.state import GlobalState -from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size -from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding -from megatron.core.ssm import gated_delta_net -from megatron.core.utils import unwrap_model -from transformers import AutoTokenizer - -import modelopt.torch.quantization as mtq -from modelopt.torch.quantization.linear_attention import linear_attention_training_phase -from modelopt.torch.quantization.utils import is_quantized - - -def run_training(config, quant_config, prefill_tokens, teacher_provider=None): - """Delegate optimization and checkpointing to Bridge, with a fixed phase per sequence.""" - if not 0 <= prefill_tokens < config.dataset.seq_length: - raise ValueError("prefill-tokens must leave at least one suffix label") - resume = config.checkpoint.load and has_modelopt_state(config.checkpoint.load) - with ExitStack() as phases: - - def prepare_student(models): - # Restore the saved policy before quantization and before the QAD conversion hook. - if resume: - load_modelopt_state(models, config.checkpoint.load) - layer_types = (gated_delta_net.GatedDeltaNet,) - if hasattr(gated_delta_net, "KimiDeltaAttention"): - layer_types += (gated_delta_net.KimiDeltaAttention,) - for student in unwrap_model(models): - layers = [m for m in student.modules() if isinstance(m, layer_types)] - if not layers: - raise ValueError( - "Each local model chunk must contain Megatron GatedDeltaNet or " - "KimiDeltaAttention; choose a pipeline layout with linear attention on every stage" - ) - student.requires_grad_(False) - for layer in layers: - layer.requires_grad_(True) - if config.model.recompute_granularity == "full": - - def require_input_grad(module, args, output): - return output.requires_grad_(True) if module.training else output - - # Reentrant checkpointing needs a differentiable input with frozen embeddings. - for module in student.modules(): - if isinstance(module, LanguageModelEmbedding): - handle = module.register_forward_hook(require_input_grad) - phases.callback(handle.remove) - if not is_quantized(student): - mtq.quantize(student, quant_config) - phases.enter_context( - linear_attention_training_phase( - student, [prefill_tokens] * config.train.micro_batch_size - ) - ) - return models - - def masked_batch(data_iterator): - batch = dict(next(data_iterator)) - batch["loss_mask"] = batch["loss_mask"].clone() - batch["loss_mask"][..., :prefill_tokens] = 0 - yield batch - - def forward_step(state: GlobalState, data_iterator, model, return_schedule_plan=False): - # Consume lazily: middle pipeline stages may not need any batch tensors. - return forward_step_modelopt( - state, masked_batch(data_iterator), model, return_schedule_plan - ) - - config.model.register_pre_wrap_hook(prepare_student) - if teacher_provider is not None: - config.model = convert_to_distillation_provider( - config.model, teacher_provider, ModelOptDistillConfig(skip_lm_loss=True) - ) - distill(config, forward_step) - else: - pretrain(config, forward_step) - - -def model_provider(path, options, *, load_weights=True): - """Build a Megatron provider with the same topology options as the Bridge example.""" - bridge = AutoBridge.from_hf_pretrained(str(path), trust_remote_code=options.trust_remote_code) - provider = bridge.to_megatron_provider(load_weights=load_weights) - provider.tensor_model_parallel_size = options.tp_size - provider.pipeline_model_parallel_size = options.pp_size - provider.pipeline_dtype = torch.bfloat16 - provider.context_parallel_size = 1 - provider.expert_model_parallel_size = options.ep_size - provider.expert_tensor_parallel_size = 1 - provider.sequence_parallel = options.tp_size > 1 - provider.gradient_accumulation_fusion = False - provider.calculate_per_token_loss = True - provider.seq_length = options.length - return provider - - -def main(): - """Train a local model with QAT, or add a frozen teacher for QAD.""" - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--model", type=Path, required=True) - parser.add_argument("--teacher-model", type=Path, help="Unquantized teacher; enables QAD") - parser.add_argument( - "--train-data", - type=Path, - required=True, - help="Megatron tokenized dataset prefix (.bin/.idx)", - ) - parser.add_argument( - "--quant-config", - type=Path, - default=Path(__file__).with_name("configs") / "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) - parser.add_argument("--global-batch-size", type=int, default=1) - parser.add_argument("--tp_size", type=int, default=1, help="Tensor parallel size") - parser.add_argument("--pp_size", type=int, default=1, help="Pipeline parallel size") - parser.add_argument("--ep_size", type=int, default=1, help="Expert parallel size") - 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") - if options.global_batch_size < 1: - parser.error("global-batch-size must be positive") - if min(options.tp_size, options.pp_size, options.ep_size) < 1: - parser.error("Parallel sizes must be positive") - if any(not Path(f"{options.train_data}.{suffix}").is_file() for suffix in ("bin", "idx")): - parser.error("train-data must name an existing Megatron .bin/.idx dataset prefix") - - checkpoint_dir = str(options.output / "checkpoints") - resume = has_modelopt_state(checkpoint_dir) - student = model_provider(options.model, options, load_weights=not resume) - teacher = None - if options.teacher_model is not None: - tokenizer_kwargs = {"trust_remote_code": options.trust_remote_code} - student_vocab = AutoTokenizer.from_pretrained(options.model, **tokenizer_kwargs).get_vocab() - teacher_vocab = AutoTokenizer.from_pretrained( - options.teacher_model, **tokenizer_kwargs - ).get_vocab() - if student_vocab != teacher_vocab: - parser.error("QAD requires student and teacher to use the same tokenizer vocabulary") - teacher = model_provider(options.teacher_model, options) - vocab_sizes = [ - calculate_padded_vocab_size( - p.vocab_size, p.make_vocab_size_divisible_by, p.tensor_model_parallel_size - ) - for p in (student, teacher) - ] - if vocab_sizes[0] != vocab_sizes[1]: - parser.error("QAD requires matching student and teacher output vocabulary dimensions") - - config = ConfigContainer( - model=student, - train=TrainingConfig( - train_iters=options.train_steps, - global_batch_size=options.global_batch_size, - micro_batch_size=1, - ), - validation=ValidationConfig(eval_iters=0, eval_interval=options.train_steps), - optimizer=OptimizerConfig( - optimizer="adam", - lr=options.learning_rate, - min_lr=0, - weight_decay=0, - clip_grad=1.0, - use_distributed_optimizer=True, - ), - scheduler=SchedulerConfig( - lr_decay_style="constant", - lr_warmup_iters=0, - start_weight_decay=0, - end_weight_decay=0, - use_checkpoint_opt_param_scheduler=True, - ), - ddp=DistributedDataParallelConfig( - average_in_collective=False, use_distributed_optimizer=True - ), - dataset=GPTDatasetConfig( - seq_length=options.length, - blend=([str(options.train_data)], None), - split="100,0,0", - random_seed=options.seed, - reset_position_ids=False, - reset_attention_mask=False, - eod_mask_loss=False, - dataloader_type="single", - num_workers=0, - ), - tokenizer=TokenizerConfig( - tokenizer_type="HuggingFaceTokenizer", - tokenizer_model=str(options.model), - hf_tokenizer_kwargs={"trust_remote_code": options.trust_remote_code}, - ), - checkpoint=CheckpointConfig( - save=checkpoint_dir, - load=checkpoint_dir, - save_interval=options.train_steps, - async_save=False, - ckpt_format="torch_dist", - ), - logger=LoggerConfig(log_interval=1), - rng=RNGConfig(seed=options.seed), - mixed_precision="bf16_mixed", - ) - run_training( - config, json.loads(options.quant_config.read_text()), options.prefill_tokens, teacher - ) - - -if __name__ == "__main__": - main() diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index 6fb77c2163b..62a31a58c18 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -34,7 +34,7 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `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 | +| `linear_attention_state_int8_dynamic.yaml` | GDN/KDA INT8 state quantizer entries; the [complete recipe](../../../general/ptq/linear_attention_state_int8_dynamic.yaml) enables Hadamard during decode | | `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 | | `indexer_k_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer key cache (`*indexer_k_quantizer`), global scale fixed to 1; Blackwell+ GPUs | | `indexer_q_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer query (`*indexer_q_quantizer`), global scale fixed to 1; Blackwell+ GPUs | diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/linear_attention_state_int8_dynamic.yaml similarity index 73% rename from modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml rename to modelopt_recipes/configs/ptq/units/linear_attention_state_int8_dynamic.yaml index e86390cfef5..eceadf5d5d0 100644 --- a/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml +++ b/modelopt_recipes/configs/ptq/units/linear_attention_state_int8_dynamic.yaml @@ -13,8 +13,8 @@ # 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. +# Quantizer-only INT8 unit for GDN and KDA recurrent states. +# The complete general/ptq/linear_attention_state_int8_dynamic recipe enables Hadamard. # modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig - quantizer_name: '*gdn_state_quantizer' cfg: @@ -24,3 +24,11 @@ axis: [0, 1] type: dynamic pass_through_bwd: true + - quantizer_name: '*kda_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/general/ptq/gdn_state_int8_dynamic.yaml b/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml similarity index 83% rename from modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml rename to modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml index e960282121b..3cc24515e25 100644 --- a/modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml +++ b/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml @@ -13,21 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe imports: base_disable_all: configs/ptq/units/base_disable_all - gdn_state_int8_dynamic: configs/ptq/units/gdn_state_int8_dynamic + state_int8: configs/ptq/units/linear_attention_state_int8_dynamic metadata: - recipe_type: ptq description: >- - Dynamic INT8 GDN state QDQ with 32-value Hadamard rotation during decode. + Dynamic INT8 GDN/KDA 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 + - $import: state_int8 linear_attention: - module_name: '*' cfg: diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 189d773299f..1ec8753f698 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -64,7 +64,7 @@ supported combinations. | `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | | `iq2_s` | IQ2_S W2A16 (2.56 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | | `q8_0` | Q8_0 W8A16 (8.5 bpw), eligible linears | 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) | +| `linear_attention_state_int8_dynamic` | GDN/KDA decode state INT8 + Hadamard; weights unchanged | none | none (dynamic scales; requires a prefix/decode phase context) | diff --git a/tests/examples/megatron_bridge/test_linear_attention.py b/tests/examples/megatron_bridge/test_linear_attention.py deleted file mode 100644 index 2ab89aad2b2..00000000000 --- a/tests/examples/megatron_bridge/test_linear_attention.py +++ /dev/null @@ -1,149 +0,0 @@ -# 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 runpy - -import pytest -import torch -from _test_utils.examples.megatron_example_runner import reset_megatron_global_state -from _test_utils.examples.run_command import MODELOPT_ROOT -from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider -from megatron.bridge.training.config import ( - CheckpointConfig, - ConfigContainer, - DistributedDataParallelConfig, - LoggerConfig, - MockGPTDatasetConfig, - OptimizerConfig, - RNGConfig, - SchedulerConfig, - TokenizerConfig, - TrainingConfig, - ValidationConfig, -) -from megatron.core.utils import unwrap_model - -from modelopt.torch.quantization.utils import is_quantized - -run_training = runpy.run_path(str(MODELOPT_ROOT / "examples/llm_qat/linear_attention/train.py"))[ - "run_training" -] - - -def _train(qad, recipe): - captured = {} - - def provider(name): - model = HybridModelProvider( - num_layers=1, - hidden_size=64, - ffn_hidden_size=128, - num_attention_heads=2, - hybrid_layer_pattern="G", - vocab_size=128, - seq_length=16, - linear_num_key_heads=2, - linear_num_value_heads=2, - linear_key_head_dim=32, - linear_value_head_dim=32, - linear_conv_kernel_dim=4, - linear_attention_freq=1, - experimental_attention_variant="gated_delta_net", - is_hybrid_model=True, - activation_func=torch.nn.functional.silu, - calculate_per_token_loss=True, - gradient_accumulation_fusion=False, - recompute_granularity="full", - recompute_method="uniform", - recompute_num_layers=1, - cross_entropy_loss_fusion=False, - ) - - def capture(models): - module = unwrap_model(models[0]) - captured[name] = module - if name == "student": - captured["before"] = ( - module.decoder.layers[0].self_attention.out_proj.weight.detach().clone() - ) - return models - - model.register_post_wrap_hook(capture) - return model - - config = ConfigContainer( - model=provider("student"), - train=TrainingConfig(train_iters=1, global_batch_size=1, micro_batch_size=1), - validation=ValidationConfig(eval_iters=0, eval_interval=1), - optimizer=OptimizerConfig( - optimizer="adam", lr=1e-2, min_lr=0, weight_decay=0, use_distributed_optimizer=True - ), - scheduler=SchedulerConfig( - lr_decay_style="constant", - lr_warmup_iters=0, - start_weight_decay=0, - end_weight_decay=0, - ), - ddp=DistributedDataParallelConfig( - average_in_collective=False, use_distributed_optimizer=True - ), - dataset=MockGPTDatasetConfig( - seq_length=16, - random_seed=123, - reset_position_ids=False, - reset_attention_mask=False, - eod_mask_loss=False, - dataloader_type="single", - num_workers=0, - ), - tokenizer=TokenizerConfig(tokenizer_type="NullTokenizer", vocab_size=128), - checkpoint=CheckpointConfig(async_save=False), - logger=LoggerConfig(log_interval=1), - rng=RNGConfig(seed=123), - mixed_precision="bf16_mixed", - ) - run_training(config, recipe, 8, provider("teacher") if qad else None) - return captured - - -@pytest.fixture(scope="module") -def compiled_state_training(project_root_path): - """Compile one tiny GDN shape before timing the single-GPU training checks.""" - if not torch.cuda.is_available(): - pytest.skip("Requires CUDA") - recipe = json.loads( - ( - project_root_path / "examples/llm_qat/linear_attention/configs/decode_state_int8.json" - ).read_text() - ) - try: - _train(True, recipe) - finally: - reset_megatron_global_state() - return recipe - - -@pytest.mark.parametrize("qad", [False, True], ids=["qat", "qad"]) -def test_state_training(compiled_state_training, qad): - captured = _train(qad, compiled_state_training) - attention = captured["student"].decoder.layers[0].self_attention - assert attention.gdn_state_quantizer.is_enabled - assert torch.isfinite(attention.out_proj.weight).all() - assert not torch.equal(attention.out_proj.weight, captured["before"]) - if qad: - teacher = captured["teacher"] - assert not is_quantized(teacher) - assert not any(parameter.requires_grad for parameter in teacher.parameters()) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py index 6b243289bbf..b1cd73a6da4 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json from contextlib import nullcontext import pytest @@ -29,6 +28,7 @@ from megatron.core.transformer import TransformerConfig import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe from modelopt.torch.opt.plugins.mcore_dist_checkpointing import ( restore_sharded_modelopt_state, save_sharded_modelopt_state, @@ -157,10 +157,9 @@ def _test_kda(rank, size, cfg, checkpoint_path): @pytest.fixture(scope="module") -def compiled_kda_workers(dist_workers_size_1, project_root_path): +def compiled_kda_workers(dist_workers_size_1): """Warm one KDA shape outside the functional test timer.""" - path = project_root_path / "examples/llm_qat/linear_attention/configs/decode_state_int8.json" - cfg = json.loads(path.read_text()) + cfg = load_recipe("general/ptq/linear_attention_state_int8_dynamic").quantize.model_dump() cfg["linear_attention"][0]["cfg"]["decode"].update( mode="replay", replay={"window": 5}, decay_log_step=1 / 256 ) From f67e37fc3111c869be38f5f5c014923ec5ccb375 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 6 Oct 2026 10:25:26 -0700 Subject: [PATCH 11/27] Align GDN and KDA state QAT with native vLLM arithmetic Signed-off-by: Kai Xu --- .pre-commit-config.yaml | 1 + CHANGELOG.rst | 2 +- LICENSE | 1 + .../int8.py => common/int8_quant.py} | 27 ++-- .../linear_attention/fla_chunk_delta_h.py | 64 ++++---- .../linear_attention/serving/__init__.py | 24 +++ .../linear_attention/serving/chunk_delta_h.py | 62 ++++++++ .../linear_attention/serving/forward.py | 120 +++++++++++++++ modelopt/torch/quantization/conversion.py | 7 + .../quantization/linear_attention/__init__.py | 2 +- .../linear_attention/_vllm_autograd.py | 145 ++++++++++++++++++ .../quantization/linear_attention/config.py | 48 +++++- .../quantization/linear_attention/decode.py | 109 ++++++++----- .../quantization/linear_attention/kda.py | 76 +++++---- .../quantization/linear_attention/prefill.py | 69 +++++---- .../{decode_prefill.py => training.py} | 97 ++++++++++-- .../quantization/linear_attention/utils.py | 79 +++++++--- .../nn/modules/tensor_quantizer.py | 7 +- modelopt/torch/quantization/plugins/gdn.py | 2 + modelopt/torch/quantization/plugins/kda.py | 2 + .../quantization/plugins/linear_attention.py | 42 ++++- .../torch/quantization/plugins/megatron.py | 58 ++++++- modelopt_recipes/configs/ptq/units/README.md | 21 +++ .../linear_attention_replay_fp8_dynamic.yaml | 28 ++++ ..._attention_state_int8_block32_dynamic.yaml | 33 ++++ ..._attention_state_int8_block32_dynamic.yaml | 42 +++++ modelopt_recipes/ptq.md | 6 +- .../linear_attention_reference.py | 10 +- .../test_fla_chunk_gated_delta_rule.py | 19 +++ .../quantization/plugins/test_megatron_kda.py | 11 +- .../test_linear_attention_training.py | 88 +++++++++++ .../torch/quantization/plugins/test_gdn.py | 76 +++++++-- .../test_linear_attention_decode.py | 92 ++++------- .../test_linear_attention_hadamard.py | 2 +- .../quantization/test_tensor_quantizer_cpu.py | 13 ++ 35 files changed, 1203 insertions(+), 282 deletions(-) rename modelopt/torch/kernels/quantization/{linear_attention/int8.py => common/int8_quant.py} (57%) create mode 100644 modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py create mode 100644 modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py create mode 100644 modelopt/torch/kernels/quantization/linear_attention/serving/forward.py create mode 100644 modelopt/torch/quantization/linear_attention/_vllm_autograd.py rename modelopt/torch/quantization/linear_attention/{decode_prefill.py => training.py} (58%) create mode 100644 modelopt_recipes/configs/ptq/units/linear_attention_replay_fp8_dynamic.yaml create mode 100644 modelopt_recipes/configs/ptq/units/linear_attention_state_int8_block32_dynamic.yaml create mode 100644 modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml create mode 100644 tests/gpu_vllm/torch/quantization/test_linear_attention_training.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f60fad5cc15..eeddf83e257 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -112,6 +112,7 @@ repos: modelopt/torch/_deploy/utils/onnx_utils.py| modelopt/torch/export/transformer_engine.py| modelopt/torch/fastgen/plugins/qwen_image_pdd.py| + modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.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| diff --git a/CHANGELOG.rst b/CHANGELOG.rst index c65b55dfd96..7e039f49507 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,7 +16,7 @@ Changelog - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. -- Add experimental Megatron-Core 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 experimental Megatron-Core 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; the optional ``vllm_0_15`` precision profile requires ``vllm==0.15.1`` for serving-aligned forward values. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before a separate KV-cache AutoQuantize stage, with independent resumable checkpoints for the weight and diff --git a/LICENSE b/LICENSE index abc6a0cb761..0e99a82684b 100644 --- a/LICENSE +++ b/LICENSE @@ -216,6 +216,7 @@ Portions of this repository were adapted from code originally authored by the following copyright holders, licensed under the Apache License, Version 2.0 (see full license text above): + Copyright contributors to the vLLM project Copyright 2021 The HuggingFace Inc. team Copyright 2022 The HuggingFace Team Copyright 2022, Lefebvre Dalloz Services diff --git a/modelopt/torch/kernels/quantization/linear_attention/int8.py b/modelopt/torch/kernels/quantization/common/int8_quant.py similarity index 57% rename from modelopt/torch/kernels/quantization/linear_attention/int8.py rename to modelopt/torch/kernels/quantization/common/int8_quant.py index 1fa5c621af9..bcd27d7e089 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/int8.py +++ b/modelopt/torch/kernels/quantization/common/int8_quant.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Signed narrow-range INT8 QDQ for recurrent-state tiles.""" +"""Composable signed narrow-range INT8 QDQ matching CUDA TensorQuantizer.""" import triton import triton.language as tl @@ -23,23 +23,18 @@ @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 +def int8_scalar_qdq(value, amax): + """QDQ with a scalar or broadcastable amax, using CUDA TensorQuantizer arithmetic.""" + # Match gemm/tensor_quant_gpu.cu: tiny groups become zero; ties round to even. + tiny = amax < 2.0**-24 + scale = tl.div_rn(127.0, tl.where(tiny, 1.0, amax)) + codes = tl.clamp(libdevice.nearbyint(value * scale), -127.0, 127.0) + return tl.where(tiny, 0.0, tl.div_rn(codes, scale)) @triton.jit -def int8_block_qdq(value, GROUP_SIZE: tl.constexpr, STATE_V_FIRST: tl.constexpr): - """Match TensorQuantizer's dynamic INT8 QDQ per key row and value group.""" - if STATE_V_FIRST: - value = tl.trans(value) +def int8_block_qdq(value, GROUP_SIZE: tl.constexpr): + """Dynamic INT8 QDQ per row and last-axis group of a two-dimensional tile.""" groups = tl.reshape(value, (value.shape[0], value.shape[1] // GROUP_SIZE, GROUP_SIZE)) amax = tl.max(tl.abs(groups), axis=2, keep_dims=True) - # Match the CUDA TensorQuantizer's scale direction, ties-to-even, and tiny-group handling. - tiny = amax < 2.0**-24 - scale = tl.div_rn(127.0, tl.where(tiny, 1.0, amax)) - codes = tl.clamp(libdevice.nearbyint(groups * scale), -127.0, 127.0) - rounded = tl.reshape(tl.where(tiny, 0.0, tl.div_rn(codes, scale)), value.shape) - if STATE_V_FIRST: - rounded = tl.trans(rounded) - return rounded + return tl.reshape(int8_scalar_qdq(groups, amax), value.shape) 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 index ee6a8cd1947..fb5ee76ec2e 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -42,15 +42,14 @@ from modelopt.torch.kernels.quantization import linear_attention as state_formats from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq - -from .int8 import int8_block_qdq, int8_scalar_qdq +from modelopt.torch.kernels.quantization.common.int8_quant import int8_block_qdq, int8_scalar_qdq 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.""" +def _state_amax(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr): + """[ModelOpt] Joint amax of the full [K, BV] tile for one sequence and head.""" # The b_h blocks partition K; their joint amax gives one scale for the full [K, BV] tile. b_amax = tl.max(tl.abs(b_h1)) if K > 64: @@ -59,25 +58,26 @@ def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr, STATE_QDQ: tl.cons 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))) - # Use scale=1 for zero tiles to avoid division by zero during QDQ. - if STATE_QDQ == state_formats.STATE_QDQ_INT8: - return tl.where(b_amax > 0, b_amax * (1.0 / 127.0), 1.0) - elif STATE_QDQ == state_formats.STATE_QDQ_FP8_E4M3: - return tl.where(b_amax > 0, b_amax / 448.0, 1.0) - else: - tl.static_assert(False, "Unsupported state QDQ format") + return b_amax @triton.jit def _state_scalar_qdq( - value, scale, STATE_QDQ: tl.constexpr, GROUP_SIZE: tl.constexpr, STATE_V_FIRST: tl.constexpr + value, amax, STATE_QDQ: tl.constexpr, GROUP_SIZE: tl.constexpr, STATE_V_FIRST: tl.constexpr ): if STATE_QDQ == state_formats.STATE_QDQ_INT8: if GROUP_SIZE: - return int8_block_qdq(value, GROUP_SIZE, STATE_V_FIRST) - return int8_scalar_qdq(value, scale) + # Quantization groups run along V regardless of the kernel's state layout. + if STATE_V_FIRST: + value = tl.trans(value) + value = int8_block_qdq(value, GROUP_SIZE) + if STATE_V_FIRST: + value = tl.trans(value) + return value + return int8_scalar_qdq(value, amax) elif STATE_QDQ == state_formats.STATE_QDQ_FP8_E4M3: tl.static_assert(GROUP_SIZE == 0, "Blockwise state QDQ currently supports only INT8") + scale = tl.where(amax > 0, amax / 448.0, 1.0) return fp8_scalar_qdq(value, scale) else: tl.static_assert(False, "Unsupported state QDQ format") @@ -241,27 +241,27 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( # [ModelOpt] A state read from a quantized cache is quantized before the first chunk uses it. if STATE_QDQ != state_formats.STATE_QDQ_OFF: if STATE_QDQ_GROUP_SIZE: - b_scale = 1.0 + b_amax = 0.0 elif K > 192: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K, STATE_QDQ=STATE_QDQ) + b_amax = _state_amax(b_h1, b_h2, b_h3, b_h4, K=K) elif K > 128: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K, STATE_QDQ=STATE_QDQ) + b_amax = _state_amax(b_h1, b_h2, b_h3, b_h3, K=K) elif K > 64: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K, STATE_QDQ=STATE_QDQ) + b_amax = _state_amax(b_h1, b_h2, b_h2, b_h2, K=K) 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, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST) + b_amax = _state_amax(b_h1, b_h1, b_h1, b_h1, K=K) + b_h1 = _state_scalar_qdq(b_h1, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST) if K > 64: b_h2 = _state_scalar_qdq( - b_h2, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST + b_h2, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST ) if K > 128: b_h3 = _state_scalar_qdq( - b_h3, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST + b_h3, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST ) if K > 192: b_h4 = _state_scalar_qdq( - b_h4, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST + b_h4, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST ) # main recurrence @@ -420,27 +420,27 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( # Legacy quantizers share a [K, BV] scale; block quantizers scale per key/value group. if STATE_QDQ != state_formats.STATE_QDQ_OFF: if STATE_QDQ_GROUP_SIZE: - b_scale = 1.0 + b_amax = 0.0 elif K > 192: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K, STATE_QDQ=STATE_QDQ) + b_amax = _state_amax(b_h1, b_h2, b_h3, b_h4, K=K) elif K > 128: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K, STATE_QDQ=STATE_QDQ) + b_amax = _state_amax(b_h1, b_h2, b_h3, b_h3, K=K) elif K > 64: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K, STATE_QDQ=STATE_QDQ) + b_amax = _state_amax(b_h1, b_h2, b_h2, b_h2, K=K) 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, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST) + b_amax = _state_amax(b_h1, b_h1, b_h1, b_h1, K=K) + b_h1 = _state_scalar_qdq(b_h1, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST) if K > 64: b_h2 = _state_scalar_qdq( - b_h2, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST + b_h2, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST ) if K > 128: b_h3 = _state_scalar_qdq( - b_h3, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST + b_h3, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST ) if K > 192: b_h4 = _state_scalar_qdq( - b_h4, b_scale, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST + b_h4, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST ) if STORE_FINAL_STATE: diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py new file mode 100644 index 00000000000..740c75f3bd3 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/__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. + +"""Optional vLLM forward kernels for the ``vllm_0_15`` precision profile.""" + +try: + import vllm +except ImportError as error: + raise ImportError( + "decode.precision='vllm_0_15' requires the optional vLLM dependency; " + "install the dependencies in the linear-attention QAT example's requirements-vllm.txt." + ) from error diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py b/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py new file mode 100644 index 00000000000..95ca3882dbc --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py @@ -0,0 +1,62 @@ +# Adapted from: https://github.com/vllm-project/vllm/blob/1892993bc18e243e2c05841314c5e9c06a80c70d/vllm/model_executor/layers/fla/ops/chunk_delta_h.py +# Modifications: import the kernel; retain FP32 intermediates for the training adjoint. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Original FLA code: Copyright (c) 2023-2025, Songlin Yang, Yu Zhang. +# FLA is licensed under the MIT license reproduced in the root LICENSE. + +# 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. + +"""Save FP32 training intermediates using vLLM's unchanged chunk-state kernel.""" + +import torch +import triton +from vllm.model_executor.layers.fla.ops.chunk_delta_h import ( + chunk_gated_delta_rule_fwd_kernel_h_blockdim64, +) +from vllm.model_executor.layers.fla.ops.index import prepare_chunk_offsets + + +def chunk_state(k, w, u, g, gk, initial_state, cu_seqlens): + """Return chunk-start states, residual values, and final state for one sequence.""" + _, length, heads, key_dim = k.shape + value_dim = u.shape[-1] + # The Torch adjoint needs unrounded values. Output kernels receive BF16 casts. + h = k.new_empty(1, triton.cdiv(length, 64), heads, key_dim, value_dim, dtype=torch.float32) + updated = torch.empty_like(u, dtype=torch.float32) + final = torch.empty_like(initial_state, dtype=torch.float32) + chunk_gated_delta_rule_fwd_kernel_h_blockdim64[ + lambda meta: (triton.cdiv(value_dim, meta["BV"]), heads) + ]( + k=k, + v=u, + w=w, + v_new=updated, + g=g, + gk=gk, + h=h, + h0=initial_state, + ht=final, + cu_seqlens=cu_seqlens, + chunk_offsets=prepare_chunk_offsets(cu_seqlens, 64), + T=length, + H=heads, + Hg=heads, + K=key_dim, + V=value_dim, + BT=64, + ) + return h, updated, final diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py new file mode 100644 index 00000000000..3d670566bde --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py @@ -0,0 +1,120 @@ +# 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. + +"""vLLM forward primitives with FP32 saved values; autograd lives in quantization.""" + +import torch +from vllm.model_executor.layers.fla.ops.chunk_o import chunk_fwd_o +from vllm.model_executor.layers.fla.ops.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from vllm.model_executor.layers.fla.ops.cumsum import chunk_local_cumsum +from vllm.model_executor.layers.fla.ops.fused_recurrent import fused_recurrent_gated_delta_rule +from vllm.model_executor.layers.fla.ops.kda import ( + chunk_gla_fwd_o_gk, + chunk_kda_scaled_dot_kkt_fwd, + fused_recurrent_kda, +) +from vllm.model_executor.layers.fla.ops.kda import fused_kda_gate as fused_kda_gate +from vllm.model_executor.layers.fla.ops.kda import recompute_w_u_fwd as kda_wu +from vllm.model_executor.layers.fla.ops.l2norm import l2norm_fwd +from vllm.model_executor.layers.fla.ops.solve_tril import solve_tril +from vllm.model_executor.layers.fla.ops.wy_fast import recompute_w_u_fwd as gdn_wu + +from .chunk_delta_h import chunk_state + + +def prefill(q, k, v, g, beta, state, scale, normalize=False): + """Return one sequence's output, final state, and rounded forward intermediates.""" + q, k, v, g, beta = [x.unsqueeze(0).contiguous() for x in (q, k, v, g, beta)] + cu = torch.tensor([0, q.shape[1]], device=q.device, dtype=torch.int32) + if normalize: + q, k = l2norm_fwd(q), l2norm_fwd(k) + channel = g.ndim == 4 + gc = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu) + if channel: + lower, scores = chunk_kda_scaled_dot_kkt_fwd(q, k, gc, beta, scale=scale, cu_seqlens=cu) + else: + lower = chunk_scaled_dot_kkt_fwd( + k=k, beta=beta, g=gc, cu_seqlens=cu, output_dtype=torch.float32 + ) + scores = None + inverse = solve_tril(A=lower, cu_seqlens=cu, output_dtype=k.dtype) + if channel: + w, u, _, kg = kda_wu(k=k, v=v, beta=beta, A=inverse, gk=gc, cu_seqlens=cu) + else: + w, u = gdn_wu(k=k, v=v, beta=beta, A=inverse, g_cumsum=gc, cu_seqlens=cu) + kg = k + assert kg is not None + h, updated, final = chunk_state( + k=kg, + w=w, + u=u, + g=None if channel else gc, + gk=gc if channel else None, + initial_state=state.unsqueeze(0).contiguous(), + cu_seqlens=cu, + ) + assert updated is not None and final is not None + if channel: + out = chunk_gla_fwd_o_gk( + q=q, + v=updated.to(v.dtype), + g=gc, + A=scores, + h=h.to(k.dtype), + scale=scale, + o=torch.empty_like(v), + cu_seqlens=cu, + chunk_size=64, + ) + else: + out = chunk_fwd_o( + q=q, k=k, v=updated.to(v.dtype), h=h.to(k.dtype), g=gc, scale=scale, cu_seqlens=cu + ) + intermediates = { + "q": q[0], + "k": k[0], + "g": gc[0], + "lower": lower[0], + "inverse": inverse[0], + "w": w[0], + "u": u[0], + "h": h[0], + "updated": updated[0], + "kg": kg[0], + "scores": None if scores is None else scores[0], + } + return out[0], final[0], intermediates + + +def step(q, k, v, g, beta, state, scale, normalize=False): + """Run one native token update without modifying the incoming state.""" + recurrent = fused_recurrent_kda if g.ndim == 2 else fused_recurrent_gated_delta_rule + out, final = recurrent( + *[x[None, None].contiguous() for x in (q, k, v, g, beta)], + initial_state=state.unsqueeze(0).contiguous(), + inplace_final_state=False, + scale=scale, + use_qk_l2norm_in_kernel=normalize, + cu_seqlens=torch.tensor([0, 1], device=q.device, dtype=torch.int32), + ssm_state_indices=torch.tensor([0], device=q.device, dtype=torch.int32), + ) + return out[0, 0], final[0] + + +def fused_gdn_gating(*args, **kwargs): + """Load the optional vLLM model only when Megatron needs GDN gate preparation.""" + from vllm.model_executor.models.qwen3_next import fused_gdn_gating as native_gate + + return native_gate(*args, **kwargs) diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index df24856ceca..a961b9d1e09 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -151,6 +151,12 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: quantizer_state_dict = dict(metadata["quantizer_state"]) _discard_legacy_w_quantizer_state(quantizer_state_dict) + for name, module in _linear_attention_modules(model).items(): + policy = metadata.get("linear_attention", {}).get(name, {}) + if policy.get("schema_version", 1) < 2: + for handle in module.replay_quantizer_names: + key = f"{name}.{handle}" if name else handle + quantizer_state_dict.setdefault(key, getattr(module, handle).get_modelopt_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(): @@ -254,6 +260,7 @@ def _apply_linear_attention_policy(model, config, saved_policies=None): for name, policy in saved_policies.items(): modules[name].linear_attention_config = LinearAttentionConfig(**policy) for name, module in modules.items(): + module._migrate_legacy_replay_quantizers() module.validate_linear_attention() if getattr(config, "linear_attention", []): print_rank_0( diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py index 305bb2cd81f..883c701bf4c 100644 --- a/modelopt/torch/quantization/linear_attention/__init__.py +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -17,6 +17,6 @@ from .config import * from .decode import * -from .decode_prefill import * from .kda import * from .prefill import * +from .training import * diff --git a/modelopt/torch/quantization/linear_attention/_vllm_autograd.py b/modelopt/torch/quantization/linear_attention/_vllm_autograd.py new file mode 100644 index 00000000000..dfd67ea925a --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/_vllm_autograd.py @@ -0,0 +1,145 @@ +# 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 adapter for pinned serving arithmetic. + +The kernels supply forward values. Autograd differentiates the corresponding +operations evaluated at those values; operand casts and QDQ use identity STE. +The optional vLLM dependency supplies kernels; training owns its state and needs no server. +""" + +import torch + +from .utils import forward_value + + +def rounded(value): + return forward_value(value, value.to(torch.bfloat16)) + + +def normalized(value): + return value / (value.square().sum(-1, keepdim=True) + 1e-6).sqrt() + + +def prefix(q, k, v, g, beta, state, scale, beta_dtype, normalize=False): + # Keep the optional vLLM dependency isolated to the native precision profile. + from ...kernels.quantization.linear_attention.serving.forward import prefill + + with torch.no_grad(): + out, final, saved = prefill( + q.to(torch.bfloat16), + k.to(torch.bfloat16), + v.to(torch.bfloat16), + g.float(), + beta.to(beta_dtype), + state.float(), + scale, + normalize, + ) + if not torch.is_grad_enabled() or not any(x.requires_grad for x in (q, k, v, g, beta, state)): + return out.float(), final + q, k = [ + forward_value(normalized(x) if normalize else x, saved[n]) for n, x in (("q", q), ("k", k)) + ] + outputs = [] + channel = g.ndim == 3 + for chunk, lo in enumerate(range(0, len(q), 64)): + hi = min(lo + 64, len(q)) + qc, kc, vc = [x[lo:hi].transpose(0, 1) for x in (q, k, v)] + bc = beta[lo:hi].transpose(0, 1).unsqueeze(-1) + gc = forward_value(g[lo:hi].transpose(0, 1).cumsum(1), saved["g"][lo:hi].transpose(0, 1)) + state = forward_value(state, saved["h"][chunk]) + hs = rounded(state) + count = hi - lo + + def matrix(name, value): + return forward_value(value, saved[name][lo:hi].transpose(0, 1)[..., :count]) + + if channel: + lower_rows, score_rows = [], [] + for row in range(count): + right = ( + kc[:, : row + 1] * (gc[:, row : row + 1] - gc[:, : row + 1]).exp() + ).transpose(-1, -2) + lower_rows.append( + torch.nn.functional.pad( + (bc[:, row : row + 1] * kc[:, row : row + 1]) @ right, (0, count - row - 1) + ) + ) + score_rows.append( + torch.nn.functional.pad( + (qc[:, row : row + 1] * scale) @ right, (0, count - row - 1) + ) + ) + lower = matrix("lower", torch.cat(lower_rows, dim=1).tril(-1)) + scores = matrix("scores", torch.cat(score_rows, dim=1)) + gate = gc.exp() + else: + causal = torch.ones(count, count, device=q.device, dtype=torch.bool).tril() + decay = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() + lower = matrix("lower", (bc * (kc @ kc.transpose(-1, -2)) * decay).tril(-1)) + scores = ((qc @ kc.transpose(-1, -2)) * decay).tril() + gate = gc.exp().unsqueeze(-1) + eye = torch.eye(count, device=q.device, dtype=q.dtype).expand_as(lower) + inverse = matrix( + "inverse", + torch.linalg.solve_triangular(eye + lower, eye, upper=False, unitriangular=True), + ) + u = forward_value(inverse @ rounded(bc * vc), saved["u"][lo:hi].transpose(0, 1)) + kb = rounded(bc * kc) if beta_dtype == torch.bfloat16 else bc * kc + w = forward_value(inverse @ rounded(kb * gate), saved["w"][lo:hi].transpose(0, 1)) + updated = forward_value(u - w @ hs, saved["updated"][lo:hi].transpose(0, 1)) + if channel: + output = rounded(rounded(qc * scale) * gate) @ hs + rounded(scores) @ rounded(updated) + kg = forward_value(kc * (gc[:, -1:] - gc).exp(), saved["kg"][lo:hi].transpose(0, 1)) + state = state * gate[:, -1, :, None] + kg.transpose(-1, -2) @ rounded(updated) + else: + output = ((qc @ hs) * gate + rounded(scores) @ rounded(updated)) * scale + weighted = rounded(updated * (gc[:, -1:] - gc).exp().unsqueeze(-1)) + state = state * gate[:, -1, :, None] + kc.transpose(-1, -2) @ weighted + outputs.append(forward_value(output.transpose(0, 1), out[lo:hi])) + return torch.cat(outputs), forward_value(state, final) + + +def step(q, k, v, gate, beta, state, scale, normalize=False): + # Keep the optional vLLM dependency isolated to the native precision profile. + from ...kernels.quantization.linear_attention.serving.forward import step as native_step + + with torch.no_grad(): + out, final = native_step( + q.to(torch.bfloat16), + k.to(torch.bfloat16), + v.to(torch.bfloat16), + gate.float(), + beta.float(), + state.float(), + scale, + normalize, + ) + if not torch.is_grad_enabled() or not any( + x.requires_grad for x in (q, k, v, gate, beta, state) + ): + return out.float(), final + if normalize: + q, k = normalized(q), normalized(k) + decay = gate.exp().unsqueeze(-1) + if gate.ndim == 1: + decay = decay.unsqueeze(-1) + decayed = state * decay + residual = v - (decayed * k.unsqueeze(-1)).sum(-2) + update = residual * beta.unsqueeze(-1) + working = forward_value(decayed + k.unsqueeze(-1) * update.unsqueeze(-2), final) + output = ((q * scale).unsqueeze(-1) * working).sum(-2) + return forward_value(output, out), working diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index 2168c870bfb..010d7fca426 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -17,7 +17,7 @@ from typing import Literal -from pydantic import Field, model_validator +from pydantic import Field, PrivateAttr, model_validator from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField @@ -49,20 +49,41 @@ class _SolveConfig(ModeloptBaseConfig): class LinearAttentionReplayConfig(ModeloptBaseConfig): - """Anchor refresh and encoded rank-one update policy.""" + """Anchor refresh and update encoding schedule; quantizers are configured by quant_cfg.""" window: int = Field(default=8, ge=1, le=64, strict=True) - factor_qdq: bool = ModeloptField(default=True) encoding: Literal["once", "reencode"] = ModeloptField(default="once") + _legacy_factor_qdq: bool | None = PrivateAttr(default=None) + + @model_validator(mode="wrap") + @classmethod + def _load_legacy_factor_qdq(cls, values, handler): + legacy = None + if isinstance(values, dict) and "factor_qdq" in values: + values = dict(values) + legacy = values.pop("factor_qdq") + if not isinstance(legacy, bool): + raise ValueError("Legacy factor_qdq must be a boolean") + result = handler(values) + if legacy is not None: + result._legacy_factor_qdq = legacy + return result class LinearAttentionDecodeConfig(ModeloptBaseConfig): - """Explicit suffix recurrence; workload supplies per-sequence prefix lengths.""" + """Explicit suffix recurrence; workload supplies per-sequence prefix lengths. + + ``precision='vllm_0_15'`` uses the pinned serving forward arithmetic for + the chunked prefix and token suffix, with a differentiable Torch adjoint. + ``full`` retains the FP32/FP64 Torch reference. Quantization is configured + independently by TensorQuantizer; this profile does not enable QDQ. + """ mode: Literal["token", "replay"] = ModeloptField(default="token") readout: Literal["working", "stored"] = ModeloptField(default="stored") quantize_initial: bool = ModeloptField(default=True) prefill_state_qdq: bool = ModeloptField(default=False) + precision: Literal["full", "vllm_0_15"] = ModeloptField(default="full") 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) @@ -82,6 +103,18 @@ def _validate_replay(self): raise ValueError( "Hadamard state QDQ starts at decode handoff; disable prefill_state_qdq" ) + if self.precision == "vllm_0_15" and ( + self.mode != "token" + or self.readout != "working" + or self.state_codec != "tile" + or not self.quantize_initial + or self.prefill_state_qdq + or self.decay_log_step is not None + ): + raise ValueError( + "vllm_0_15 requires token mode, working readout, tile codec, initial-state QDQ, " + "and no additional prefill or gate rounding" + ) return self @@ -97,7 +130,7 @@ class LinearAttentionConfig(ModeloptBaseConfig): groups; ``state.block_v`` controls only execution tiling (at least one group). """ - schema_version: Literal[1] = ModeloptField(default=1) + schema_version: Literal[1, 2] = ModeloptField(default=2) backend: Literal["fla", "matmul"] = ModeloptField(default="fla") chunk_size: Literal[64] = ModeloptField(default=64) state: _StateConfig = ModeloptField(default=_StateConfig()) @@ -106,6 +139,11 @@ class LinearAttentionConfig(ModeloptBaseConfig): @model_validator(mode="before") @classmethod def _load_legacy_defaults(cls, values): + if isinstance(values, dict) and values.get("schema_version") == 1: + decode = values.get("decode") + if isinstance(decode, dict) and isinstance(decode.get("replay"), dict): + replay = {"factor_qdq": True, **decode["replay"]} + values = {**values, "decode": {**decode, "replay": replay}} if isinstance(values, dict) and "solve" in values: _SolveConfig.model_validate(values["solve"]) values = {key: value for key, value in values.items() if key != "solve"} diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py index eb4b6fa3779..2e995c47242 100644 --- a/modelopt/torch/quantization/linear_attention/decode.py +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -20,7 +20,7 @@ import torch from .config import LinearAttentionDecodeConfig -from .utils import _tile_qdq, state_quantizer_config +from .utils import _resolve_state_quantizer, _tile_qdq, forward_value __all__ = [ "EncodedLinearAttentionTensor", @@ -91,16 +91,6 @@ def _hadamard32(value): 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, @@ -126,22 +116,38 @@ def _encode( 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 - ) - decoded, scales = _tile_qdq(value, block_v, state_format, state=state) - if state_format == "int8": - decoded = _IdentityGradient.apply(value, decoded) + return EncodedLinearAttentionTensor(forward_value(value, decoded), metadata, "int8", 32) + decoded, scales = _tile_qdq(value, block_v, state_format, state_quantizer=state_quantizer) return EncodedLinearAttentionTensor(decoded, scales, state_format, block_v) -def _signature(config, state_qdq, block_v, state_format, state_quantizer): +def _encode_factor(value, quantizer): + if quantizer is None: + return EncodedLinearAttentionTensor(value, None, "identity", None) + decoded = quantizer(value) + enabled = quantizer.is_enabled and quantizer._if_quant + format = str(quantizer.num_bits) if enabled else "identity" + return EncodedLinearAttentionTensor(decoded, None, format, None) + + +def _signature(config, state_qdq, block_v, state_format, state_quantizer, replay_quantizers): signature = ( config.model_dump_json(exclude={"prefill_state_qdq"}) + f"/{state_qdq}/{block_v}/{state_format}" ) if state_quantizer is not None and state_quantizer.block_sizes is not None: signature += f"/group={state_quantizer.block_sizes[-1]}" + for quantizer in replay_quantizers: + if quantizer is None: + signature += "/factor=None" + else: + grouping = ( + quantizer.block_sizes if quantizer.block_sizes is not None else quantizer.axis + ) + signature += ( + f"/factor={quantizer.is_enabled}/{quantizer._if_quant}/{quantizer.num_bits}" + f"/{grouping}/{quantizer.backend}/{quantizer.backend_extra_args}" + ) return signature @@ -160,7 +166,7 @@ def _round_log_gate(gate, step): if step is None: return gate rounded = (gate / step).round() * step - return _IdentityGradient.apply(gate, rounded.detach()) + return forward_value(gate, rounded) def _prepare_carry( @@ -177,6 +183,7 @@ def _prepare_carry( position, state_format, state_quantizer, + replay_quantizers, ): 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") @@ -194,7 +201,9 @@ def _prepare_carry( 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, state_quantizer) + signature = _signature( + config, state_qdq, block_v, state_format, state_quantizer, replay_quantizers + ) 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]) @@ -241,22 +250,38 @@ def recurrent_decode( state_qdq=False, state_format="fp8_e4m3", state_quantizer=None, + replay_key_quantizer=None, + replay_update_quantizer=None, block_v=64, initial_state=None, carry=None, position=0, scale=None, + use_qk_l2norm_in_kernel=False, ): """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. + Replay factors use the supplied TensorQuantizers; omitted quantizers leave factors unchanged. """ - if state_quantizer is not None: - state_qdq = state_quantizer.is_enabled and state_quantizer._if_quant - if state_quantizer.is_enabled: - state_format, _ = state_quantizer_config(state_quantizer) + state_quantizer, state_qdq, state_format = _resolve_state_quantizer( + state_quantizer, state_qdq, state_format + ) + serving = config.precision == "vllm_0_15" + if serving and q.device.type != "cuda": + raise ValueError("vllm_0_15 requires CUDA") + replay_quantizers = (replay_key_quantizer, replay_update_quantizer) + if config.replay is not None and config.replay._legacy_factor_qdq: + if any(quantizer is None for quantizer in replay_quantizers): + raise ValueError("Replace factor_qdq with explicit replay key/update TensorQuantizers") + for quantizer in replay_quantizers: + if quantizer is not None and quantizer.is_enabled: + if config.replay is None or not quantizer.fake_quant: + raise ValueError( + "Replay factor quantizers require replay mode and fake quantization" + ) carry, signature = _prepare_carry( q, k, @@ -271,6 +296,7 @@ def recurrent_decode( position, state_format, state_quantizer, + replay_quantizers, ) if len(q) == 0: # Keep empty input gradients defined without introducing a state write. @@ -287,8 +313,8 @@ def recurrent_decode( 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), + _encode_factor(e.key.values, replay_key_quantizer), + _encode_factor(e.update.values, replay_update_quantizer), e.log_retention, ) for e in entries @@ -298,18 +324,21 @@ def recurrent_decode( ) 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 serving: + from ._vllm_autograd import step as serving_step + + native_output, working = serving_step( + q[t], k[t], v[t], gate, beta[t], state, scale, use_qk_l2norm_in_kernel + ) + else: + decay = gate.exp().unsqueeze(-1) + if gate.ndim == 1: + decay = decay.unsqueeze(-1) + key = _encode_factor(k[t], replay_key_quantizer) + decayed = state * decay + residual = v[t] - _sum_keys(key.values.unsqueeze(-1) * decayed) + update = _encode_factor(beta[t].unsqueeze(-1) * residual, replay_update_quantizer) + 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 @@ -332,7 +361,9 @@ def recurrent_decode( 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) + outputs.append( + native_output if serving else _sum_keys(q[t].unsqueeze(-1) * read) * scale + ) carry = next_carry output = torch.stack(outputs) if config.state_codec == "int8_hadamard32": diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py index d7cc0229669..54b7b8efce2 100644 --- a/modelopt/torch/quantization/linear_attention/kda.py +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -18,7 +18,8 @@ import torch import torch.nn.functional as F -from .decode_prefill import _decode_prefill +from .training import _prefill_decode_forward, _prepare_prefill_inputs +from .utils import forward_value __all__ = ["matmul_kda"] @@ -34,6 +35,8 @@ def matmul_kda( state_qdq=False, state_format="fp8_e4m3", state_quantizer=None, + replay_key_quantizer=None, + replay_update_quantizer=None, scale=None, initial_state=None, output_final_state=False, @@ -54,12 +57,12 @@ def matmul_kda( return_intermediate_states=False, prefill_lengths=None, ): - """Normalize KDA inputs and run exact prefix plus configured suffix recurrence. + """Normalize KDA inputs and run a chunked prefix plus configured suffix recurrence. Gates follow FLA's kernel activation formula and per-key log retention. + The optional serving precision profile supplies matching native forward values + and a Torch adjoint through the rounded state trajectory. """ - 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" @@ -71,39 +74,54 @@ def matmul_kda( 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)) + beta_dtype = beta.dtype + q, k, v, g, beta, serving = _prepare_prefill_inputs( + q, k, v, g, beta, policy=policy, chunk_size=chunk_size, normalize=use_qk_l2norm_in_kernel + ) + dtype = q.dtype 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: + raw_gate = g 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 serving: + if lower_bound is not None or safe_gate: + raise ValueError("vllm_0_15 uses the native softplus KDA gate") + # Import the optional vLLM backend only for the native precision profile. + from ...kernels.quantization.linear_attention.serving.forward import fused_kda_gate + + with torch.no_grad(): + native_gate = fused_kda_gate( + raw_gate.flatten(-2).contiguous(), A_log, raw_gate.shape[-1], g_bias=dt_bias + ) + g = forward_value(g, native_gate) 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, - state_quantizer=state_quantizer, - 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") + return _prefill_decode_forward( + q, + k, + v, + g, + beta, + policy=policy, + state_qdq=state_qdq, + state_format=state_format, + state_quantizer=state_quantizer, + replay_key_quantizer=replay_key_quantizer, + replay_update_quantizer=replay_update_quantizer, + 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, + beta_dtype=beta_dtype, + prefill_lengths=prefill_lengths, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel if serving else False, + ) diff --git a/modelopt/torch/quantization/linear_attention/prefill.py b/modelopt/torch/quantization/linear_attention/prefill.py index 9bd514a3ab6..2f67670c77b 100644 --- a/modelopt/torch/quantization/linear_attention/prefill.py +++ b/modelopt/torch/quantization/linear_attention/prefill.py @@ -15,11 +15,10 @@ """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 +from .training import _prefill_decode_forward, _prepare_prefill_inputs __all__ = ["matmul_gdn"] @@ -35,6 +34,8 @@ def matmul_gdn( state_qdq=False, state_format="fp8_e4m3", state_quantizer=None, + replay_key_quantizer=None, + replay_update_quantizer=None, scale=None, initial_state=None, output_final_state=False, @@ -51,20 +52,24 @@ def matmul_gdn( cp_context=None, prefill_lengths=None, ): - """Normalize GDN inputs and run exact prefix plus configured suffix recurrence. + """Normalize GDN inputs and run a chunked 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. + The optional serving precision profile supplies matching native forward values + and a Torch adjoint through the rounded state trajectory. """ - 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)) + beta_dtype = beta.dtype + q, k, v, g, beta, serving = _prepare_prefill_inputs( + q, k, v, g, beta, policy=policy, chunk_size=chunk_size, normalize=use_qk_l2norm_in_kernel + ) + dtype = q.dtype + if serving and (use_gate_in_kernel or use_beta_sigmoid_in_kernel): + raise ValueError( + "vllm_0_15 GDN expects prepared log gates and beta from the Megatron adapter" + ) 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") @@ -73,24 +78,26 @@ def matmul_gdn( beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) if g.ndim != 3: raise ValueError("matmul_gdn requires scalar GDN log gates") - if policy.decode is not None: - return _decode_prefill( - q, - k, - v, - g, - beta, - policy=policy, - state_qdq=state_qdq, - state_format=state_format, - state_quantizer=state_quantizer, - 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") + return _prefill_decode_forward( + q, + k, + v, + g, + beta, + policy=policy, + state_qdq=state_qdq, + state_format=state_format, + state_quantizer=state_quantizer, + replay_key_quantizer=replay_key_quantizer, + replay_update_quantizer=replay_update_quantizer, + 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, + beta_dtype=beta_dtype, + prefill_lengths=prefill_lengths, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel if serving else False, + ) diff --git a/modelopt/torch/quantization/linear_attention/decode_prefill.py b/modelopt/torch/quantization/linear_attention/training.py similarity index 58% rename from modelopt/torch/quantization/linear_attention/decode_prefill.py rename to modelopt/torch/quantization/linear_attention/training.py index 8544599e1e9..e5319410368 100644 --- a/modelopt/torch/quantization/linear_attention/decode_prefill.py +++ b/modelopt/torch/quantization/linear_attention/training.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Differentiable explicit per-sequence prefill/decode phase handoff.""" +"""Training forwards with a chunked prefill prefix and recurrent decode suffix.""" from contextlib import contextmanager from itertools import pairwise @@ -22,7 +22,7 @@ from ._chunk_prefill import chunk_gdn, chunk_kda from .decode import recurrent_decode -from .utils import state_quantizer_config +from .utils import _resolve_state_quantizer, _state_qdq __all__ = ["linear_attention_training_phase"] @@ -65,7 +65,23 @@ def linear_attention_training_phase(model, prefill_lengths): module._linear_attention_prefill_lengths = original -def _decode_prefill( +def _prepare_prefill_inputs(q, k, v, g, beta, *, policy, chunk_size, normalize): + """Validate the shared policy and prepare GDN/KDA working dtypes and Q/K normalization.""" + if policy.backend != "matmul" or chunk_size != policy.chunk_size: + raise ValueError("Matmul prefill requires backend='matmul' and its configured chunk size") + if policy.decode is None: + raise ValueError("An explicit decode policy is required") + serving = policy.decode.precision == "vllm_0_15" + if serving and (q.device.type != "cuda" or any(x.dtype != torch.bfloat16 for x in (q, k, v))): + raise ValueError("precision='vllm_0_15' requires CUDA BF16 Q/K/V inputs") + dtype = torch.float64 if q.dtype == torch.float64 else torch.float32 + q, k, v, g, beta = (x.to(dtype) for x in (q, k, v, g, beta)) + if normalize and not serving: + q, k = (x * (x.square().sum(-1, keepdim=True) + 1e-6).rsqrt() for x in (q, k)) + return q, k, v, g, beta, serving + + +def _prefill_decode_forward( q, k, v, @@ -76,6 +92,8 @@ def _decode_prefill( state_qdq, state_format, state_quantizer, + replay_key_quantizer, + replay_update_quantizer, scale, initial_state, output_final_state, @@ -83,12 +101,22 @@ def _decode_prefill( cu_seqlens_cpu, state_v_first, output_dtype, + beta_dtype, prefill_lengths, + use_qk_l2norm_in_kernel=False, ): - if state_quantizer is not None: - state_qdq = state_quantizer.is_enabled and state_quantizer._if_quant - if state_quantizer.is_enabled: - state_format, _ = state_quantizer_config(state_quantizer) + """Run both prefill and decode phases in one differentiable training forward. + + Each sequence's chunked prefix produces the state for its token or ReplaySSM + suffix. Their outputs are joined in token order for the training loss. + + Args: + prefill_lengths: Prefix token count per sequence. For 128 tokens, a value + of 64 selects 64 chunked prefill tokens followed by 64 recurrent tokens. + """ + state_quantizer, state_qdq, state_format = _resolve_state_quantizer( + state_quantizer, state_qdq, state_format + ) 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]: @@ -121,6 +149,15 @@ def _decode_prefill( p > end - start for p, (_, start, end) in zip(prefixes, sequences) ): raise ValueError("Supply one valid prefill length per sequence") + serving_precision = policy.decode.precision == "vllm_0_15" + if serving_precision and ( + keys > 256 + or (g.ndim == 4 and values != keys) + or (use_qk_l2norm_in_kernel and keys & (keys - 1)) + ): + raise ValueError( + "vllm_0_15 requires K <= 256, KDA V=K, and power-of-two K for normalization" + ) if initial_state is None: states = q.new_zeros(len(sequences), heads, keys, values) else: @@ -135,25 +172,51 @@ def _decode_prefill( prefix, state = q.new_empty(0, heads, values), states[n] if prefixes[n]: with torch.autocast(device_type=q.device.type, enabled=False): - prefix, state = prefix_fn( - *(x[b, start:split] for x in (q, k, v, g, beta)), - state_qdq=state_qdq and policy.decode.prefill_state_qdq, - state_format=state_format, - state_quantizer=state_quantizer, - scale=scale, - initial_state=state, - chunk_size=policy.chunk_size, - state_qdq_block_v=policy.state.block_v, - ) + if serving_precision: + from ._vllm_autograd import prefix as serving_prefix + + # A continuation prefill consumes a stored cache just as native serving + # does. A fresh zero-state prefix has no incoming cache to quantize. + if initial_state is not None and state_qdq: + state = _state_qdq( + state, policy.state.block_v, state_format, state_quantizer + ) + prefix, state = serving_prefix( + q=q[b, start:split], + k=k[b, start:split], + v=v[b, start:split], + g=g[b, start:split], + beta=beta[b, start:split], + state=state, + scale=keys**-0.5 if scale is None else scale, + beta_dtype=beta_dtype, + normalize=use_qk_l2norm_in_kernel, + ) + else: + prefix, state = prefix_fn( + *(x[b, start:split] for x in (q, k, v, g, beta)), + state_qdq=state_qdq and policy.decode.prefill_state_qdq, + state_format=state_format, + state_quantizer=state_quantizer, + scale=scale, + initial_state=state, + chunk_size=policy.chunk_size, + state_qdq_block_v=policy.state.block_v, + ) + # Keep the prefix state attached so suffix losses backpropagate through prefill. + # recurrent_decode applies configured state QDQ at the handoff and suffix writes. suffix, carry = recurrent_decode( *(x[b, split:end] for x in (q, k, v, g, beta)), config=policy.decode, + replay_key_quantizer=replay_key_quantizer, + replay_update_quantizer=replay_update_quantizer, state_qdq=state_qdq, state_format=state_format, state_quantizer=state_quantizer, block_v=policy.state.block_v, initial_state=state, position=prefixes[n], + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, scale=scale, ) outputs.append(torch.cat((prefix, suffix))) diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py index 2c650877cf5..bed74203474 100644 --- a/modelopt/torch/quantization/linear_attention/utils.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -29,6 +29,21 @@ _STATE_FORMATS: dict[int | tuple[int, int], str] = {(4, 3): "fp8_e4m3", 8: "int8"} +class _ForwardValue(torch.autograd.Function): + @staticmethod + def forward(ctx, value, rounded): + return rounded.to(value.dtype) + + @staticmethod + def backward(ctx, grad): + return grad, None + + +def forward_value(value, rounded): + # Unlike value + (rounded - value).detach(), this cannot lose low bits by cancellation. + return _ForwardValue.apply(value, rounded) + + def validate_gdn_quantizer( quantizer: TensorQuantizer, *, @@ -86,19 +101,34 @@ def state_quantizer_config( return _STATE_FORMATS[quantizer.num_bits], 0 -def _fp8_quantize(value: torch.Tensor, axis): - """Apply ModelOpt FP8 QDQ and return its detached dequantization scales.""" +def _make_state_quantizer(state_format): + """Build a standalone-call quantizer; converted modules supply their registered instance.""" # 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) + if state_format not in _STATE_FORMATS.values(): + raise ValueError("State format must be fp8_e4m3 or int8") + return TensorQuantizer( + QuantizerAttributeConfig( + num_bits=(4, 3) if state_format == "fp8_e4m3" else 8, + type="dynamic", + axis=(0, 1), + narrow_range=True, + 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 _resolve_state_quantizer(state_quantizer, state_qdq, state_format): + """Resolve legacy flags and registered quantizer settings at either training entry point.""" + if state_quantizer is None and state_qdq: + state_quantizer = _make_state_quantizer(state_format) + if state_quantizer is not None: + state_qdq = state_quantizer.is_enabled and state_quantizer._if_quant + if state_quantizer.is_enabled: + state_format, _ = state_quantizer_config(state_quantizer) + return state_quantizer, state_qdq, state_format def _state_qdq( @@ -114,26 +144,27 @@ def _state_qdq( 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") - quantized, _ = _tile_qdq(state, block_v, state_format, state=True) - if state_format == "fp8_e4m3": - return quantized - return state + (quantized - state).detach() + quantized, _ = _tile_qdq(state, block_v, state_format, state_quantizer=state_quantizer) + return quantized -def _tile_qdq(value, block_v, state_format, *, state=False): - """Return tile-rounded values and detached scales; INT8 callers supply their STE.""" +def _tile_qdq(value, block_v, state_format, *, state_quantizer=None): + """Return TensorQuantizer tile QDQ with identity STE and detached scales.""" + quantizer = ( + state_quantizer if state_quantizer is not None else _make_state_quantizer(state_format) + ) rounded, scales = [], [] for part in value.split(block_v, dim=-1): - tensor = part.flatten(-2) if state else part - if state_format == "fp8_e4m3": - axis = tuple(range(tensor.ndim - 1)) or None - decoded, scale = _fp8_quantize(tensor, axis) + # Canonical [N, H, K*BV] shape preserves per-head tile scales for both state ranks. + tensor = part.flatten(-2) + inputs = tensor.reshape(1, -1, tensor.shape[-1]) + decoded = quantizer(inputs) + amax = quantizer._get_amax(inputs).float() + if quantizer.num_bits == 8: + scale = amax / quantizer.maxbound else: - with torch.no_grad(): - tensor = tensor.float() - amax = tensor.abs().amax(dim=-1, keepdim=True) - scale = torch.where(amax > 0, amax / 127.0, torch.ones_like(amax)) - decoded = (tensor / scale).round().clamp(-127, 127) * scale + safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) + scale = torch.div(quantizer.maxbound, safe_amax).reciprocal() rounded.append(decoded.reshape_as(part)) - scales.append(scale.squeeze(-1)) + scales.append(scale.reshape(tensor.shape[:-1])) return torch.cat(rounded, dim=-1).to(value.dtype), torch.stack(scales, dim=-1) diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index a698514d53a..5e3b677ea2f 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -999,9 +999,12 @@ def set_quant_params(axis, block_reshape_size, padding, slices, amax_shape=None) if amax_shape: self._amax_shape_for_export = amax_shape - # Reshape size have already been set - if hasattr(self, "_block_reshape_size"): + # Static scales require fixed groups; dynamic scales allow changing input shapes. + if hasattr(self, "_block_reshape_size") and not self._dynamic: return + for attribute in ("_padding", "_slices"): + if hasattr(self, attribute): + delattr(self, attribute) reshape_size, quantize_axis, paddings, slices = [], [], [], [] diff --git a/modelopt/torch/quantization/plugins/gdn.py b/modelopt/torch/quantization/plugins/gdn.py index f82848b1d0e..1f00ef8c20a 100644 --- a/modelopt/torch/quantization/plugins/gdn.py +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -116,6 +116,8 @@ def _state_quantized_chunk_gated_delta_rule( *args, policy=self.linear_attention_config, state_quantizer=self.gdn_state_quantizer, + replay_key_quantizer=self.replay_key_quantizer, + replay_update_quantizer=self.replay_update_quantizer, chunk_size=chunk_size, prefill_lengths=self._linear_attention_prefill_lengths, **kwargs, diff --git a/modelopt/torch/quantization/plugins/kda.py b/modelopt/torch/quantization/plugins/kda.py index c459e7fccce..52218dbffe8 100644 --- a/modelopt/torch/quantization/plugins/kda.py +++ b/modelopt/torch/quantization/plugins/kda.py @@ -60,6 +60,8 @@ def _state_quantized_chunk_kda(self, kernel, *args, **kwargs): *args, policy=self.linear_attention_config, state_quantizer=self.kda_state_quantizer, + replay_key_quantizer=self.replay_key_quantizer, + replay_update_quantizer=self.replay_update_quantizer, prefill_lengths=self._linear_attention_prefill_lengths, **kwargs, ) diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py index 1c6b823a4a0..22f75cfe142 100644 --- a/modelopt/torch/quantization/plugins/linear_attention.py +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -25,9 +25,10 @@ class _LinearAttentionQuantMixin(QuantModule): linear_attention_quantizer_names: tuple[str, ...] = () + replay_quantizer_names = ("replay_key_quantizer", "replay_update_quantizer") def _setup(self): - for name in self.linear_attention_quantizer_names: + for name in (*self.linear_attention_quantizer_names, *self.replay_quantizer_names): self._register_temp_attribute( name, TensorQuantizer(QuantizerAttributeConfig(enable=False)) ) @@ -42,18 +43,53 @@ def _linear_attn_state(self): def linear_attention_is_enabled(self): """Whether an operand, state, or arithmetic policy changes the computation.""" return ( - any(getattr(self, name).is_enabled for name in self.linear_attention_quantizer_names) + any( + getattr(self, name).is_enabled + for name in (*self.linear_attention_quantizer_names, *self.replay_quantizer_names) + ) or self.linear_attention_config.decode is not None ) + def _migrate_legacy_replay_quantizers(self): + decode = self.linear_attention_config.decode + replay = decode.replay if decode is not None else None + legacy = getattr(replay, "_legacy_factor_qdq", None) + if ( + replay is not None + and legacy is None + and self.linear_attention_config.schema_version == 1 + ): + legacy = replay.__dict__.pop("factor_qdq", True) + if replay is not None and legacy is not None: + for name in self.replay_quantizer_names: + grouping = ( + {"axis": (0,)} + if name == "replay_key_quantizer" + else {"block_sizes": {-1: self.linear_attention_config.state.block_v}} + ) + getattr(self, name).set_from_attribute_config( + QuantizerAttributeConfig( + num_bits=(4, 3), type="dynamic", enable=legacy, **grouping + ) + ) + replay._legacy_factor_qdq = None + self.linear_attention_config.schema_version = 2 + def validate_linear_attention(self): """Validate quantizer contracts shared by GDN and KDA.""" + decode = self.linear_attention_config.decode + for name in self.replay_quantizer_names: + quantizer = getattr(self, name) + if quantizer.is_enabled: + if decode is None or decode.mode != "replay": + raise ValueError(f"{name} requires a replay execution policy") + if not quantizer.fake_quant: + raise ValueError(f"{name} requires fake quantization for training") if self._linear_attn_state.is_enabled: state_format, group_size = state_quantizer_config( self._linear_attn_state, name=self.linear_attention_quantizer_names[0], ) - decode = self.linear_attention_config.decode if decode is not None and decode.state_codec == "int8_hadamard32": if state_format != "int8": raise ValueError("int8_hadamard32 requires INT8 state quantization") diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 3fded293ab0..3791da54e1e 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -1131,10 +1131,66 @@ def _setup(self): mcore_parallel.get_tensor_model_parallel_group(), ) + @property + def _serving_arithmetic(self): + decode = self.linear_attention_config.decode + return ( + self.linear_attention_is_enabled + and decode is not None + and decode.precision == "vllm_0_15" + ) + + def _prepare_input_for_gated_delta_rule(self, *args, **kwargs): + # Preserve raw BF16 Q/K: prefill stores normalized BF16 operands, whereas + # decode normalizes inside its FP32 update. One shared pre-normalization + # would irreversibly change the suffix state trajectory. + normalize = self.use_qk_l2norm + if self._serving_arithmetic: + self.use_qk_l2norm = False + try: + return super()._prepare_input_for_gated_delta_rule(*args, **kwargs) + finally: + self.use_qk_l2norm = normalize + + def _compute_gates(self, a_log, dt_bias, batch, seq_len, *gate_feats): + gate, inputs = super()._compute_gates(a_log, dt_bias, batch, seq_len, *gate_feats) + if self._serving_arithmetic and hasattr(self, "gdn_state_quantizer"): + # Import the optional vLLM backend only for the native precision profile. + from ...kernels.quantization.linear_attention.serving.forward import fused_gdn_gating + from ..linear_attention.utils import forward_value + + raw_beta, raw_gate = gate_feats + with torch.no_grad(): + native_gate, native_beta = fused_gdn_gating( + a_log, + raw_gate.reshape(-1, raw_gate.shape[-1]).contiguous(), + raw_beta.reshape(-1, raw_beta.shape[-1]).contiguous(), + dt_bias, + ) + gate = forward_value(gate, native_gate.reshape_as(gate)) + inputs["beta"] = forward_value( + inputs["beta"], native_beta.reshape_as(inputs["beta"]) + ).to(raw_beta.dtype) + return gate, inputs + @contextmanager def _quantized_linear_attention_kernel(self): kernel = self.gated_delta_rule - self.gated_delta_rule = partial(self._linear_attention_kernel, kernel) + if self._serving_arithmetic: + if not hasattr(super(), "_prepare_input_for_gated_delta_rule") or getattr( + self, "gdn_pre_gated_delta_rule_fusion", False + ): + raise NotImplementedError( + "Serving arithmetic requires Megatron's unfused input-preparation hook" + ) + + def run_kernel(*args, **kwargs): + kwargs["use_qk_l2norm_in_kernel"] = self.use_qk_l2norm + return self._linear_attention_kernel(kernel, *args, **kwargs) + + self.gated_delta_rule = run_kernel + else: + self.gated_delta_rule = partial(self._linear_attention_kernel, kernel) try: yield finally: diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index 62a31a58c18..a867006aadd 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -35,8 +35,29 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `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+ | | `linear_attention_state_int8_dynamic.yaml` | GDN/KDA INT8 state quantizer entries; the [complete recipe](../../../general/ptq/linear_attention_state_int8_dynamic.yaml) enables Hadamard during decode | +| `linear_attention_state_int8_block32_dynamic.yaml` | Standard dynamic INT8 with one scale per key row and 32 value channels; the [complete training recipe](../../../general/ptq/linear_attention_state_int8_block32_dynamic.yaml) uses token QDQ and working-state readout | +| `linear_attention_replay_fp8_dynamic.yaml` | GDN/KDA ReplaySSM key/update FP8 quantizers; requires a replay execution policy | | `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 | | `indexer_k_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer key cache (`*indexer_k_quantizer`), global scale fixed to 1; Blackwell+ GPUs | | `indexer_q_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer query (`*indexer_q_quantizer`), global scale fixed to 1; Blackwell+ GPUs | | `moe_dispatch_nvfp4.yaml` | NVFP4 fake quantization of the MoE expert-parallel dispatch (`*dispatch_quantizer`, vLLM); Blackwell+ GPUs | | `moe_combine_nvfp4.yaml` | NVFP4 fake quantization of the MoE expert-parallel combine (`*combine_quantizer`, vLLM); Blackwell+ GPUs | + +Replay factor formats, grouping, and enable/disable controls use `quant_cfg` through +`*replay_key_quantizer` and `*replay_update_quantizer`. The replay policy controls +scheduling only; for example: + +```yaml +linear_attention: + - module_name: "*" + cfg: + backend: matmul + decode: + mode: replay + replay: {window: 8, encoding: once} +``` + +Both factor quantizers start disabled. Import `linear_attention_replay_fp8_dynamic.yaml` +to enable their dynamic FP8 recipe. Previous checkpoints using `replay.factor_qdq` +are migrated to these quantizers when restored. State quantizer recipes are independent; +this unit does not enable state quantization. diff --git a/modelopt_recipes/configs/ptq/units/linear_attention_replay_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/linear_attention_replay_fp8_dynamic.yaml new file mode 100644 index 00000000000..eaa81f957c6 --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/linear_attention_replay_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. + +# Replay factors for GDN/KDA; requires a linear_attention decode policy with mode: replay. +# Keys use one scale per head; updates use one scale per head and 64 value channels. +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*replay_key_quantizer' + cfg: + num_bits: e4m3 + axis: [0] + type: dynamic + - quantizer_name: '*replay_update_quantizer' + cfg: + num_bits: e4m3 + block_sizes: {-1: 64} + type: dynamic diff --git a/modelopt_recipes/configs/ptq/units/linear_attention_state_int8_block32_dynamic.yaml b/modelopt_recipes/configs/ptq/units/linear_attention_state_int8_block32_dynamic.yaml new file mode 100644 index 00000000000..00284f84ce3 --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/linear_attention_state_int8_block32_dynamic.yaml @@ -0,0 +1,33 @@ +# 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 INT8 with one scale per key row and 32 value channels for training and serving. +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*gdn_state_quantizer' + cfg: + num_bits: 8 + unsigned: false + narrow_range: true + block_sizes: {-1: 32} + type: dynamic + pass_through_bwd: true + - quantizer_name: '*kda_state_quantizer' + cfg: + num_bits: 8 + unsigned: false + narrow_range: true + block_sizes: {-1: 32} + type: dynamic + pass_through_bwd: true diff --git a/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml new file mode 100644 index 00000000000..ac1d078bc01 --- /dev/null +++ b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml @@ -0,0 +1,42 @@ +# 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. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + base_disable_all: configs/ptq/units/base_disable_all + state_int8: configs/ptq/units/linear_attention_state_int8_block32_dynamic + +metadata: + description: >- + Dynamic INT8 GDN/KDA state QDQ with per-key 32-value scales and working-state readout. + Emulates vLLM 0.15 BF16 prefill arithmetic without additional state QDQ. + Supply per-sequence prefix lengths through + linear_attention_training_phase. No calibration is required. +quantize: + algorithm: + quant_cfg: + - $import: base_disable_all + - $import: state_int8 + linear_attention: + - module_name: '*' + cfg: + backend: matmul + decode: + mode: token + state_codec: tile + readout: working + quantize_initial: true + prefill_state_qdq: false + precision: vllm_0_15 diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 1ec8753f698..882d16cb15f 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 32 general/ptq/ recipes (click to expand) +All 33 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -65,9 +65,13 @@ supported combinations. | `iq2_s` | IQ2_S W2A16 (2.56 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | | `q8_0` | Q8_0 W8A16 (8.5 bpw), eligible linears | none | none (no calibration) | | `linear_attention_state_int8_dynamic` | GDN/KDA decode state INT8 + Hadamard; weights unchanged | none | none (dynamic scales; requires a prefix/decode phase context) | +| `linear_attention_state_int8_block32_dynamic` | GDN/KDA state INT8, 32 value channels per key row; weights unchanged | none | none (dynamic scales; requires vLLM and a prefix/decode phase context) |
+The block32 linear-attention recipe selects `decode.precision="vllm_0_15"` +and working-state readout. It applies fake QDQ with floating-point state storage. + --- ### Model-body schemes diff --git a/tests/_test_utils/torch/quantization/linear_attention_reference.py b/tests/_test_utils/torch/quantization/linear_attention_reference.py index d0cfc70989b..3acf3fc0e5c 100644 --- a/tests/_test_utils/torch/quantization/linear_attention_reference.py +++ b/tests/_test_utils/torch/quantization/linear_attention_reference.py @@ -43,14 +43,18 @@ def state_qdq_reference(state: torch.Tensor, block_v: int = 64, state_format: st 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() + # CUDA TensorQuantizer uses a quantization multiplier and zeros tiny groups. + tiny = amax < 2**-24 + quant_scale = 127.0 / torch.where(tiny, torch.ones_like(amax), amax) + codes = (tile * quant_scale).round().clamp(-127, 127) + decoded = torch.where(tiny, 0.0, codes / quant_scale) 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) + decoded = codes * scale + rounded.append(decoded) quantized = torch.cat(rounded, dim=-1).to(state.dtype) return state + (quantized - state).detach() 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 index 0732c97eeb3..3077229a1f0 100644 --- 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 @@ -76,6 +76,8 @@ def compiled_gdn_case(request): ) args, state = make_inputs() kwargs = {"state_qdq": state_qdq, "w_quantizer": quantizer} + if state_qdq == 2: + TensorQuantizer(QuantizerAttributeConfig(num_bits=8, type="dynamic"))(state) if request.param == "state-int8-block": kwargs["state_quantizer"] = TensorQuantizer( QuantizerAttributeConfig( @@ -105,3 +107,20 @@ def test_gdn_forward_and_backward(compiled_gdn_case): ) compare(actual[0], expected[0], 0.03) compare(actual[1], expected[1], 0.05) + + if kwargs["state_qdq"] == 2 or "state_quantizer" in kwargs: + quantizer = kwargs.get("state_quantizer") or TensorQuantizer( + QuantizerAttributeConfig(num_bits=8, type="dynamic") + ) + # Zero updates isolate QDQ, reusing the compiled shape for both tile and block INT8. + zero_args = [torch.zeros_like(x) for x in args] + pattern = torch.linspace(-1, 1, state.numel(), device=state.device).reshape_as(state) + for amax in (2**-25, 2**-24, 2.0): + initial = pattern * amax + expected_state = initial + for _ in range(3): # Initial handoff and two chunk boundaries. + expected_state = quantizer(expected_state[0]).unsqueeze(0) + _, actual_state = chunk_gated_delta_rule( + *zero_args, initial_state=initial, output_final_state=True, **kwargs + ) + torch.testing.assert_close(actual_state, expected_state, rtol=0, atol=0) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py index b1cd73a6da4..57e36b226bb 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py @@ -29,6 +29,7 @@ import modelopt.torch.quantization as mtq from modelopt.recipe import load_recipe +from modelopt.torch.opt.config_loader import load_config from modelopt.torch.opt.plugins.mcore_dist_checkpointing import ( restore_sharded_modelopt_state, save_sharded_modelopt_state, @@ -121,12 +122,13 @@ def _test_kda(rank, size, cfg, checkpoint_path): assert model.gated_delta_rule is kernel policy = model.linear_attention_config - model.kda_state_quantizer.disable() + mtq.disable_quantizer(model, "*") model.linear_attention_config = LinearAttentionConfig() with torch.no_grad(): torch.testing.assert_close(forward(model), baseline, rtol=0, atol=0) - model.kda_state_quantizer.enable() model.linear_attention_config = policy + mtq.enable_quantizer(model, "*kda_state_quantizer") + mtq.enable_quantizer(model, "*replay_*_quantizer") with torch.no_grad(): expected = forward(model) @@ -146,6 +148,7 @@ def _test_kda(rank, size, cfg, checkpoint_path): assert not hasattr(restored, "kda_w_quantizer") assert restored._linear_attention_prefill_lengths is None assert restored.linear_attention_config == policy + assert restored.replay_key_quantizer.is_enabled and restored.replay_update_quantizer.is_enabled assert restored.gated_delta_rule is kernel assert torch.isfinite(hidden.grad).all() for parameter in restored.parameters(): @@ -163,6 +166,10 @@ def compiled_kda_workers(dist_workers_size_1): cfg["linear_attention"][0]["cfg"]["decode"].update( mode="replay", replay={"window": 5}, decay_log_step=1 / 256 ) + cfg["quant_cfg"].extend( + entry.model_dump(exclude_unset=True) + for entry in load_config("configs/ptq/units/linear_attention_replay_fp8_dynamic") + ) dist_workers_size_1.run(_compile_kda, cfg) return dist_workers_size_1, cfg diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py new file mode 100644 index 00000000000..deae03a2f4b --- /dev/null +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from functools import partial + +import pytest +import torch + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + matmul_gdn, + matmul_kda, +) +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.fixture(scope="module", params=[False, True]) +def compiled_serving_case(request): + """Compile one shared BF16 shape per model, outside the test-call timer.""" + pytest.importorskip("vllm.model_executor.layers.fla.ops.kda", exc_type=ModuleNotFoundError) + kda = request.param + torch.manual_seed(73) + args = [torch.randn(1, 73, 1, 32, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + args += [ + -torch.rand((1, 73, 1, 32) if kda else (1, 73, 1), device="cuda") * 0.03, + torch.rand(1, 73, 1, device="cuda") * 0.4, + ] + args = [x.requires_grad_() for x in args] + policy = LinearAttentionConfig( + backend="matmul", decode={"precision": "vllm_0_15", "readout": "working"} + ) + quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=8, + type="dynamic", + block_sizes={-1: 32}, + narrow_range=True, + pass_through_bwd=True, + ) + ).cuda() + forward = partial( + matmul_kda if kda else matmul_gdn, + *args, + policy=policy, + state_quantizer=quantizer, + prefill_lengths=[65], + use_qk_l2norm_in_kernel=True, + output_final_state=True, + ) + output, state = forward() + torch.autograd.grad(output.float().sum() + state.sum(), args) + torch.cuda.synchronize() + return args, quantizer, forward + + +def test_serving_state_qdq_and_handoff_gradient(compiled_serving_case): + args, quantizer, forward = compiled_serving_case + calls = [] + handle = quantizer.register_forward_hook(lambda *_: calls.append(True)) + output, state = forward() + handle.remove() + # One handoff QDQ, then one QDQ per suffix update; none inside the fresh prefill. + assert len(calls) == 9 + with torch.no_grad(): + expected, expected_state = forward() + torch.testing.assert_close(output, expected, rtol=0, atol=0) + torch.testing.assert_close(state, expected_state, rtol=0, atol=0) + quantizer.disable() + plain, _ = forward() + assert not torch.equal(output[:, 65:], plain[:, 65:]) + gradients = torch.autograd.grad( + output[:, 65:].float().square().sum() + state.square().sum(), args + ) + assert all(torch.isfinite(x).all() for x in gradients) + assert gradients[1][:, :65].abs().sum() > 0 diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index 5baf6244d14..60ce199473a 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -23,6 +23,7 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.quantization.config import QuantizeConfig +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins import gdn from modelopt.torch.quantization.plugins.gdn import GatedDeltaNetStateQuantMixin @@ -84,8 +85,6 @@ def quant_cfg(state=True, w=False): {"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, "axis": (0, 1), "type": "dynamic"}, # unsupported format - {"num_bits": 8, "narrow_range": True, "type": "dynamic", "block_sizes": {-1: 8}}, {"num_bits": (4, 3), "type": "dynamic", "block_sizes": {-1: 16}}, # blockwise ], ) @@ -110,6 +109,8 @@ def test_dynamic_export_removes_linear_attention_attributes(): for name in ( "gdn_state_quantizer", "gdn_w_quantizer", + "replay_key_quantizer", + "replay_update_quantizer", "linear_attention_config", "_linear_attention_prefill_lengths", ): @@ -130,11 +131,8 @@ def test_disabled_state_quantizer_calls_original_kernel(): assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" -@pytest.mark.parametrize( - ("partial_kernel", "num_bits"), - [(False, (4, 3)), (True, (4, 3)), (False, 8)], -) -def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, partial_kernel, num_bits): +@pytest.mark.parametrize("partial_kernel", [False, True]) +def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, partial_kernel): calls = [] def fake_state_qdq_kernel(*args, **kwargs): @@ -146,9 +144,7 @@ def fake_state_qdq_kernel(*args, **kwargs): if partial_kernel: model.gated_delta_rule = partial(chunk_gated_delta_rule, output_final_state=True) x = torch.randn(2, 8, 3, 4) - cfg = deepcopy(quant_cfg()) - cfg["quant_cfg"][-1]["cfg"].update(num_bits=num_bits, unsigned=False, narrow_range=True) - mtq.quantize(model, cfg, lambda m: m(x)) + mtq.quantize(model, quant_cfg(), lambda m: m(x)) model(x) assert calls and calls[-1] == { @@ -313,6 +309,66 @@ def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): ) assert not model.gdn_state_quantizer.is_enabled assert not model.gdn_w_quantizer.is_enabled + assert not model.replay_key_quantizer.is_enabled + assert not model.replay_update_quantizer.is_enabled + + +def test_replay_quantizers_use_standard_controls_and_restore(): + cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*replay_key_quantizer", + "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0,)}, + }, + { + "quantizer_name": "*replay_update_quantizer", + "cfg": {"num_bits": (4, 3), "type": "dynamic", "block_sizes": {-1: 16}}, + }, + ], + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "state": {"block_v": 16}, + "decode": {"mode": "replay", "replay": {"window": 3}}, + }, + } + ], + "algorithm": None, + } + torch.manual_seed(53) + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + x = torch.randn(1, 7, 1, 4) * 0.1 + with linear_attention_training_phase(model, [2]): + quantized = model(x) + saved = deepcopy(mto.modelopt_state(model)) + weights = deepcopy(model.state_dict()) + mtq.disable_quantizer(model, "*") + with linear_attention_training_phase(model, [2]): + plain = model(x) + assert not torch.equal(plain, quantized) + # Current checkpoints and pre-handle replay checkpoints reproduce the same computation. + for legacy in (None, False, True): + checkpoint = deepcopy(saved) + if legacy is not None: + for _, mode_state in checkpoint["modelopt_state_dict"]: + policy = mode_state["metadata"]["linear_attention"][""] + policy["schema_version"] = 1 + if not legacy: + policy["decode"]["replay"]["factor_qdq"] = False + for name in model.replay_quantizer_names: + mode_state["metadata"]["quantizer_state"].pop(name) + restored = mto.restore_from_modelopt_state(TinyGatedDeltaNet(), checkpoint) + restored.load_state_dict(weights) + assert restored.replay_key_quantizer.is_enabled == (legacy is not False) + assert restored.replay_update_quantizer.is_enabled == (legacy is not False) + assert "factor_qdq" not in restored.linear_attention_config.decode.replay.model_dump() + with linear_attention_training_phase(restored, [2]): + torch.testing.assert_close( + restored(x), plain if legacy is False else quantized, rtol=0, atol=0 + ) def test_policy_rejects_unmatched_and_unimplemented_modes(): diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py index 953e344dfcd..09cfb56e36f 100644 --- a/tests/unit/torch/quantization/test_linear_attention_decode.py +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -16,8 +16,8 @@ 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.recipe import load_recipe from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.linear_attention import ( LinearAttentionConfig, @@ -33,11 +33,8 @@ @pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) def test_state_qdq_matches_tensor_quantizer(state_format): torch.manual_seed(762) - # Exact INT8 scales exercise signed endpoints, half ties, and a zero tile. - value = torch.zeros(2, 3, 5, 19) - value[..., 0, :9] = torch.tensor([-127, -126.5, -1.5, -0.5, 0, 0.5, 1.5, 126.5, 127]) * 2 - value[0, :, 0, 16:] = torch.tensor([-254, 1, 254]) - value.requires_grad_() + # A partial value tile checks grouping and padding against TensorQuantizer. + value = torch.randn(2, 3, 5, 19, requires_grad=True) cfg = {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1)} if state_format == "int8": cfg.update(num_bits=8, unsigned=False, narrow_range=True) @@ -45,7 +42,9 @@ def test_state_qdq_matches_tensor_quantizer(state_format): expected = torch.cat( [quantizer(tile.flatten(-2)).reshape_as(tile) for tile in value.split(16, -1)], -1 ) - encoded = _encode(value, True, 16, state=True, state_format=state_format) + encoded = _encode( + value, True, 16, state=True, state_format=state_format, state_quantizer=quantizer + ) torch.testing.assert_close(encoded.values, expected, rtol=0, atol=0) probe = torch.randn_like(value) (gradient,) = torch.autograd.grad((encoded.values * probe).sum(), value) @@ -78,14 +77,18 @@ def _values_and_grads(output, state, args, initial): def test_block_state_quantizer_prefill_decode_and_gradients(kda, block_size): args, initial = _inputs(kda, length=5) # Partial value groups exercise TensorQuantizer padding as well as per-key scales. - quantizer = TensorQuantizer( - QuantizerAttributeConfig( - num_bits=8, type="dynamic", block_sizes={-1: block_size}, narrow_range=True - ) - ) - policy = LinearAttentionConfig( - backend="matmul", decode={"prefill_state_qdq": True}, state={"block_v": 16} - ) + recipe = load_recipe( + "general/ptq/linear_attention_state_int8_block32_dynamic" + ).quantize.model_dump() + cfg = recipe["quant_cfg"][2 if kda else 1]["cfg"] + cfg["block_sizes"] = {-1: block_size} + quantizer = TensorQuantizer(QuantizerAttributeConfig(**cfg)) + policy = LinearAttentionConfig(**recipe["linear_attention"][0]["cfg"]) + policy.decode.precision = "full" # Double-precision recurrence/gradient oracle. + policy.state.block_v = 16 + if kda: + policy.decode.prefill_state_qdq = True + policy.decode.readout = "stored" function = matmul_kda if kda else matmul_gdn output, final = function( *(x.unsqueeze(0) for x in args), @@ -95,8 +98,8 @@ def test_block_state_quantizer_prefill_decode_and_gradients(kda, block_size): initial_state=initial.unsqueeze(0), output_final_state=True, ) - # Sequential oracle: the prefix reads unrounded working states; decode reads stored states. - state = quantizer(initial) + # GDN exercises the serving-aligned recipe; KDA retains prefix-QDQ/stored-read coverage. + state = quantizer(initial) if policy.decode.prefill_state_qdq else initial expected = [] q, k, v, g, beta = args for t in range(len(q)): @@ -106,10 +109,11 @@ def test_block_state_quantizer_prefill_decode_and_gradients(kda, block_size): decayed = state * decay residual = v[t] - (k[t].unsqueeze(-1) * decayed).sum(-2) state = decayed + k[t].unsqueeze(-1) * (beta[t].unsqueeze(-1) * residual).unsqueeze(-2) - if t >= 3: - state = quantizer(state) - expected.append((q[t].unsqueeze(-1) * state).sum(-2) / q.shape[-1] ** 0.5) - if t == 2: + stored = quantizer(state) if t >= 3 else state + read = stored if t >= 3 and policy.decode.readout == "stored" else state + expected.append((q[t].unsqueeze(-1) * read).sum(-2) / q.shape[-1] ** 0.5) + state = stored + if t == 2 and policy.decode.prefill_state_qdq: state = quantizer(state) for actual, expected in zip( _values_and_grads(output[0], final[0], args, initial), @@ -118,49 +122,9 @@ def test_block_state_quantizer_prefill_decode_and_gradients(kda, block_size): torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-7) -@pytest.mark.parametrize(("kda", "mode"), [(False, "token"), (True, "replay")]) -def test_prefill_decode_matches_recurrence_and_gradients(kda, mode): - args, state = _inputs(kda) - policy = LinearAttentionConfig( - backend="matmul", - decode={ - "mode": mode, - "replay": {"window": 3, "factor_qdq": False} if mode == "replay" else None, - }, - ) - function = matmul_kda if kda else matmul_gdn - output, final = function( - *(x.unsqueeze(0) for x in args), - policy=policy, - prefill_lengths=[3], - initial_state=state.unsqueeze(0), - output_final_state=True, - ) - exact, expected_state = recurrent_delta_rule_reference( - *(x.unsqueeze(0) for x in args), initial_state=state.unsqueeze(0) - ) - for actual, expected in zip( - _values_and_grads(output[0], final[0], args, state), - _values_and_grads(exact[0], expected_state[0], args, state), - ): - torch.testing.assert_close(actual, expected, rtol=2e-10, atol=1e-11) - - -def test_quantized_replay_carry_continuation_and_gradients(): - args, state = _inputs() - cfg = LinearAttentionDecodeConfig(mode="replay", replay={"window": 3}) - kwargs = {"config": cfg, "state_qdq": True, "block_v": 16} - expected, final = recurrent_decode(*args, initial_state=state, **kwargs) - first, carry = recurrent_decode(*(x[:5] for x in args), initial_state=state, **kwargs) - second, carry = recurrent_decode(*(x[5:] for x in args), carry=carry, **kwargs) - for actual, expected in zip( - _values_and_grads(torch.cat((first, second)), carry.reconstruct(), args, state), - _values_and_grads(expected, final.reconstruct(), args, state), - ): - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - assert carry.position == 11 - assert carry.anchor.format == "fp8_e4m3" - assert carry.cursor == 2 +def test_serving_precision_requires_native_schedule(): + with pytest.raises(ValueError, match="working readout"): + LinearAttentionDecodeConfig(precision="vllm_0_15") def test_grid_gate_ste_keeps_gate_gradients_and_changes_trajectory(): diff --git a/tests/unit/torch/quantization/test_linear_attention_hadamard.py b/tests/unit/torch/quantization/test_linear_attention_hadamard.py index 905db03e9c5..71a4cb23fa4 100644 --- a/tests/unit/torch/quantization/test_linear_attention_hadamard.py +++ b/tests/unit/torch/quantization/test_linear_attention_hadamard.py @@ -81,7 +81,7 @@ def test_hadamard_replay_matches_dense_oracle_and_split_carry(): cfg = LinearAttentionDecodeConfig( mode="replay", state_codec="int8_hadamard32", - replay={"window": 3, "factor_qdq": False}, + replay={"window": 3}, ) kwargs = {"config": cfg, "state_format": "int8", "state_qdq": True} output, carry = recurrent_decode(*args, initial_state=initial, **kwargs) diff --git a/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py b/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py index d1512547643..bf8238cb85c 100644 --- a/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py @@ -119,3 +119,16 @@ def test_grouped_quantizer_preserves_nested_sequential_state_dict_layout(): ) assert list(grouped.state_dict()) == ["0.0._amax", "0.1._amax", "1._amax"] + + +def test_dynamic_block_quantization_accepts_changing_shapes(): + cfg = QuantizerAttributeConfig( + num_bits=8, type="dynamic", block_sizes={-1: 32}, pass_through_bwd=True + ) + quantizer = TensorQuantizer(cfg) + for shape in ((2, 3, 17), (1, 2, 4, 64), (3, 2, 33)): + value = torch.randn(shape, requires_grad=True) + output = quantizer(value) + torch.testing.assert_close(output, TensorQuantizer(cfg)(value), rtol=0, atol=0) + output.sum().backward() + torch.testing.assert_close(value.grad, torch.ones_like(value), rtol=0, atol=0) From 4e8fdf69be2d8f4eeec23e26ebcea99b00a7e957 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 6 Oct 2026 16:10:39 -0700 Subject: [PATCH 12/27] Align state QAT with native Hadamard and ReplaySSM arithmetic Signed-off-by: Kai Xu --- .pre-commit-config.yaml | 1 - CHANGELOG.rst | 4 +- .../quantization/linear_attention/__init__.py | 16 +- .../linear_attention/fla_chunk_delta_h.py | 1019 ----------------- .../fla_chunk_gated_delta_rule.py | 87 +- .../linear_attention/serving/replay.py | 130 +++ .../quantization/linear_attention/config.py | 59 +- .../quantization/linear_attention/decode.py | 148 ++- .../quantization/linear_attention/prefill.py | 4 +- .../quantization/linear_attention/training.py | 43 +- modelopt/torch/quantization/plugins/gdn.py | 40 +- modelopt/torch/quantization/plugins/kda.py | 4 +- .../torch/quantization/plugins/megatron.py | 12 +- modelopt_recipes/configs/ptq/units/README.md | 6 +- ..._attention_state_int8_block32_dynamic.yaml | 2 +- .../linear_attention_state_int8_dynamic.yaml | 5 +- .../test_fla_chunk_gated_delta_rule.py | 41 +- .../plugins/test_megatron_gated_delta_net.py | 7 +- .../quantization/plugins/test_megatron_kda.py | 4 +- .../test_linear_attention_replay.py | 148 +++ .../test_linear_attention_training.py | 2 +- .../torch/quantization/plugins/test_gdn.py | 54 +- 22 files changed, 582 insertions(+), 1254 deletions(-) delete mode 100644 modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py create mode 100644 modelopt/torch/kernels/quantization/linear_attention/serving/replay.py create mode 100644 tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index eeddf83e257..43f4049f807 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -113,7 +113,6 @@ repos: modelopt/torch/export/transformer_engine.py| modelopt/torch/fastgen/plugins/qwen_image_pdd.py| modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.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| diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7e039f49507..d20f81453ec 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,7 +16,7 @@ Changelog - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. -- Add experimental Megatron-Core 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; the optional ``vllm_0_15`` precision profile requires ``vllm==0.15.1`` for serving-aligned forward values. +- Add experimental serving-aligned GDN/KDA state QAT with explicit prefill/decode boundaries; migrate chunk-only state recipes to a serving policy and supply prefix lengths through the training phase context. Plain state QDQ requires ``vllm==0.15.1``; INT8 Hadamard and ReplaySSM require the compatible quantized-ReplaySSM serving fork. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before a separate KV-cache AutoQuantize stage, with independent resumable checkpoints for the weight and @@ -34,7 +34,7 @@ Changelog - 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. +- Add experimental dynamic FP8 fake quantization of GatedDeltaNet WY activations through the standard ``quant_cfg`` interface. W QAT uses ``fla-core==0.5.1`` with chunk size 64; state QAT uses the serving-aligned workflow. - Add fake quantization of the sparse-attention indexer key cache and query for DeepSeek-V4 (vLLM and Megatron-Core) and GLM-5.3-Flash (vLLM) through the new ``indexer_k_quantizer`` and ``indexer_q_quantizer``. Enable them by importing the ``configs/ptq/units/indexer_k_nvfp4`` and ``configs/ptq/units/indexer_q_nvfp4`` units (NVFP4 with the global scale fixed to 1) into a recipe. - Add the ``configs/ptq/units/kv_nvfp4_mla`` recipe unit for fake quantization of the MLA KV cache (DeepSeek-V3, GLM-5.3-Flash, ...) in vLLM fake-quant serving: NVFP4 for the latent and FP8 for the RoPE key. See `examples/vllm_serve/README.md `_ for an example recipe. - Add fake quantization of MoE expert-parallel communication to vLLM fake-quant serving through the new ``dispatch_quantizer`` and ``combine_quantizer`` on vLLM's routed experts, for MoE layers built on vLLM's fused-MoE layer and any all-to-all backend. Enable them by importing the ``configs/ptq/units/moe_dispatch_nvfp4`` and ``configs/ptq/units/moe_combine_nvfp4`` units into a recipe; see ``examples/vllm_serve/README.md``. diff --git a/modelopt/torch/kernels/quantization/linear_attention/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/__init__.py index 30a2ee9517a..fb6476f0950 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/__init__.py +++ b/modelopt/torch/kernels/quantization/linear_attention/__init__.py @@ -13,18 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Linear-attention kernels for quantization. +"""Linear-attention quantization adapters for optional FLA and serving kernels. -``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 or INT8 (``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. +FLA supplies the state kernels used by W QAT. Serving adapters import forward +kernels only when selected; this package initializer has no kernel dependencies. """ - -# Kernel format IDs, independent of scale grouping. The current kernels use dynamic scales. -STATE_QDQ_OFF = 0 -STATE_QDQ_FP8_E4M3 = 1 -STATE_QDQ_INT8 = 2 -STATE_QDQ_FORMATS = {"fp8_e4m3": STATE_QDQ_FP8_E4M3, "int8": STATE_QDQ_INT8} 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 deleted file mode 100644 index fb5ee76ec2e..00000000000 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py +++ /dev/null @@ -1,1019 +0,0 @@ -# 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 import linear_attention as state_formats -from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq -from modelopt.torch.kernels.quantization.common.int8_quant import int8_block_qdq, int8_scalar_qdq - -STATE_QDQ_MAX_BLOCK_V = 128 - - -@triton.jit -def _state_amax(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr): - """[ModelOpt] Joint amax of the full [K, BV] tile for one sequence and head.""" - # The b_h blocks partition K; their joint amax gives one scale for the full [K, BV] tile. - 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))) - return b_amax - - -@triton.jit -def _state_scalar_qdq( - value, amax, STATE_QDQ: tl.constexpr, GROUP_SIZE: tl.constexpr, STATE_V_FIRST: tl.constexpr -): - if STATE_QDQ == state_formats.STATE_QDQ_INT8: - if GROUP_SIZE: - # Quantization groups run along V regardless of the kernel's state layout. - if STATE_V_FIRST: - value = tl.trans(value) - value = int8_block_qdq(value, GROUP_SIZE) - if STATE_V_FIRST: - value = tl.trans(value) - return value - return int8_scalar_qdq(value, amax) - elif STATE_QDQ == state_formats.STATE_QDQ_FP8_E4M3: - tl.static_assert(GROUP_SIZE == 0, "Blockwise state QDQ currently supports only INT8") - scale = tl.where(amax > 0, amax / 448.0, 1.0) - return fp8_scalar_qdq(value, scale) - else: - tl.static_assert(False, "Unsupported state QDQ format") - - -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", "STATE_QDQ_GROUP_SIZE"], -) -@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, - STATE_QDQ_GROUP_SIZE: 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 != state_formats.STATE_QDQ_OFF: - if STATE_QDQ_GROUP_SIZE: - b_amax = 0.0 - elif K > 192: - b_amax = _state_amax(b_h1, b_h2, b_h3, b_h4, K=K) - elif K > 128: - b_amax = _state_amax(b_h1, b_h2, b_h3, b_h3, K=K) - elif K > 64: - b_amax = _state_amax(b_h1, b_h2, b_h2, b_h2, K=K) - else: - b_amax = _state_amax(b_h1, b_h1, b_h1, b_h1, K=K) - b_h1 = _state_scalar_qdq(b_h1, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST) - if K > 64: - b_h2 = _state_scalar_qdq( - b_h2, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST - ) - if K > 128: - b_h3 = _state_scalar_qdq( - b_h3, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST - ) - if K > 192: - b_h4 = _state_scalar_qdq( - b_h4, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST - ) - - # 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. - # Legacy quantizers share a [K, BV] scale; block quantizers scale per key/value group. - if STATE_QDQ != state_formats.STATE_QDQ_OFF: - if STATE_QDQ_GROUP_SIZE: - b_amax = 0.0 - elif K > 192: - b_amax = _state_amax(b_h1, b_h2, b_h3, b_h4, K=K) - elif K > 128: - b_amax = _state_amax(b_h1, b_h2, b_h3, b_h3, K=K) - elif K > 64: - b_amax = _state_amax(b_h1, b_h2, b_h2, b_h2, K=K) - else: - b_amax = _state_amax(b_h1, b_h1, b_h1, b_h1, K=K) - b_h1 = _state_scalar_qdq(b_h1, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST) - if K > 64: - b_h2 = _state_scalar_qdq( - b_h2, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST - ) - if K > 128: - b_h3 = _state_scalar_qdq( - b_h3, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST - ) - if K > 192: - b_h4 = _state_scalar_qdq( - b_h4, b_amax, STATE_QDQ, STATE_QDQ_GROUP_SIZE, STATE_V_FIRST - ) - - 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_formats.STATE_QDQ_OFF, - state_qdq_block_v: int | None = None, - state_qdq_group_size: int = 0, -) -> 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 = max(state_qdq_tile_v(V, state_qdq, state_qdq_block_v), state_qdq_group_size) - - 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, - STATE_QDQ_GROUP_SIZE=state_qdq_group_size, - ) - 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. For legacy quantizers the tile is - also the quantization granularity: one 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_formats.STATE_QDQ_OFF: - return 64 if check_shared_mem("ada") else 32 - if state_qdq not in state_formats.STATE_QDQ_FORMATS.values(): - raise ValueError(f"Unsupported state_qdq format {state_qdq}.") - 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 index 703b36b1303..dde6b75e3da 100644 --- 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 @@ -1,7 +1,6 @@ # 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. +# Modifications: optional W fake QDQ, saved rounded W for backward recomputation, +# and explicit validation. State forward/backward kernels are imported from FLA. # # Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li # @@ -31,6 +30,10 @@ import fla import torch from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.common.chunk_delta_h import ( + chunk_gated_delta_rule_bwd_dhu, + chunk_gated_delta_rule_fwd_h, +) 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 @@ -54,15 +57,9 @@ input_guard, ) -from modelopt.torch.quantization.linear_attention.utils import ( - state_quantizer_config, - validate_gdn_quantizer, -) +from modelopt.torch.quantization.linear_attention.utils import validate_gdn_quantizer from modelopt.torch.quantization.nn import TensorQuantizer -from . import STATE_QDQ_FORMATS, STATE_QDQ_FP8_E4M3, STATE_QDQ_OFF -from .fla_chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h - def chunk_gated_delta_rule_fwd( q: torch.Tensor, @@ -81,9 +78,6 @@ def chunk_gated_delta_rule_fwd( 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, - state_qdq_group_size: int = 0, w_quantizer: TensorQuantizer | None = None, ): g_input = g if use_gate_in_kernel else None @@ -145,9 +139,6 @@ def chunk_gated_delta_rule_fwd( 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, - state_qdq_group_size=state_qdq_group_size, ) if cp_context is not None: @@ -188,9 +179,6 @@ def chunk_gated_delta_rule_bwd( 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, - state_qdq_group_size: int = 0, quantized_w: torch.Tensor | None = None, ): w, u = recompute_w_u_fwd( @@ -210,8 +198,6 @@ def chunk_gated_delta_rule_bwd( 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, @@ -223,9 +209,6 @@ def chunk_gated_delta_rule_bwd( 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, - state_qdq_group_size=state_qdq_group_size, ) dv = chunk_bwd_dv_local( q=q, @@ -336,9 +319,6 @@ def forward( 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, - state_qdq_group_size: int = 0, w_quantizer: TensorQuantizer | None = None, ): q_rstd, k_rstd = None, None @@ -371,9 +351,6 @@ def forward( A_log=A_log, dt_bias=dt_bias, chunk_size=chunk_size, - state_qdq=state_qdq, - state_qdq_block_v=state_qdq_block_v, - state_qdq_group_size=state_qdq_group_size, w_quantizer=w_quantizer, ) ctx.save_for_backward( @@ -402,9 +379,6 @@ def forward( 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 - ctx.state_qdq_group_size = state_qdq_group_size return o.to(q.dtype), final_state @staticmethod @@ -453,9 +427,6 @@ def backward( 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, - state_qdq_group_size=ctx.state_qdq_group_size, quantized_w=quantized_w, ) if ctx.use_qk_l2norm_in_kernel: @@ -485,14 +456,11 @@ def backward( 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. +# call that asks for W quantization. @torch.compiler.disable def chunk_gated_delta_rule( q: torch.Tensor, @@ -561,9 +529,8 @@ def chunk_gated_delta_rule( 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``. - state_quantizer (Optional[TensorQuantizer]): - [ModelOpt] State QDQ settings, including optional INT8 last-axis ``block_sizes``. - Blocks always refer to the logical V dimension, regardless of storage layout. + w_quantizer (Optional[TensorQuantizer]): + [ModelOpt] Dynamic FP8 QDQ of the WY operand W, with identity STE. cu_seqlens (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. @@ -646,39 +613,24 @@ def chunk_gated_delta_rule( 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) state_quantizer = kwargs.pop("state_quantizer", None) - state_qdq_group_size = 0 - if state_quantizer is not None: - if state_qdq != STATE_QDQ_OFF: - raise ValueError("Specify state_quantizer or state_qdq, not both") - if state_quantizer.is_enabled: - state_format, group_size = state_quantizer_config(state_quantizer) - if state_quantizer._if_quant: - state_qdq = STATE_QDQ_FORMATS[state_format] - state_qdq_group_size = group_size + kwargs.pop("state_qdq_block_v", None) + if kwargs.pop("state_qdq", 0) or (state_quantizer is not None and state_quantizer.is_enabled): + raise ValueError( + "Chunk-only state QAT is retired; use the serving-aligned training workflow" + ) # 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`` of shape [B, T, HV, K] before it multiplies the unquantized state. 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 != STATE_QDQ_OFF and state_qdq not in STATE_QDQ_FORMATS.values(): - raise ValueError(f"Unsupported state_qdq format {state_qdq}.") if w_quantizer is not None: validate_gdn_quantizer(w_quantizer, name="gdn_w_quantizer") - if state_qdq == STATE_QDQ_FP8_E4M3 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: + raise ValueError("W quantization does not support context parallelism") if cp_context is not None: assert initial_state is None, "Initial state is not supported for CP" @@ -735,9 +687,6 @@ def chunk_gated_delta_rule( allow_neg_eigval, cp_context, chunk_size, - state_qdq, - state_qdq_block_v, - state_qdq_group_size, w_quantizer, ) return o, final_state diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/replay.py b/modelopt/torch/kernels/quantization/linear_attention/serving/replay.py new file mode 100644 index 00000000000..17cd816a758 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/replay.py @@ -0,0 +1,130 @@ +# 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. + +"""Functional training launches of the optional quantized-ReplaySSM serving kernels.""" + +import torch +from vllm.model_executor.layers.fla.ops.fused_recurrent_replayssm import ( + _apply_h_value_axis, + fused_recurrent_gated_delta_rule_replayssm, + prefill_write_checkpoint, +) + + +def _storage(value, enabled): + heads, keys, values = value.shape + state = torch.zeros( + 2, heads, values, keys, device=value.device, dtype=torch.int8 if enabled else torch.float32 + ) + scales = ( + torch.zeros(2, heads, values // 32, keys, device=value.device, dtype=torch.float16) + if enabled + else None + ) + indices = torch.ones(1, device=value.device, dtype=torch.int32) + return state, scales, indices + + +def _decoded(state, scales): + values = state[1].transpose(-1, -2).float() + if scales is None: + return values, None + metadata = scales[1].transpose(-1, -2).contiguous() + return (values.unflatten(-1, (-1, 32)) * metadata.float().unsqueeze(-1)).flatten(-2), metadata + + +def checkpoint(value, enabled): + """Return the serving checkpoint's decoded Hadamard-basis values and FP16 scales.""" + state, scales, indices = _storage(value, enabled) + original = value.transpose(-1, -2).unsqueeze(0).contiguous().float() + if enabled: + prefill_write_checkpoint( + original, state, None, None, scales, None, None, indices, 8, hadamard=True + ) + else: + state[1] = _apply_h_value_axis(original)[0] + return _decoded(state, scales) + + +def original_basis(value): + """Use serving's FP32 checkpoint transform when returning a dense state.""" + return _apply_h_value_axis(value.transpose(-1, -2).unsqueeze(0).contiguous())[0].transpose( + -1, -2 + ) + + +def step(q, k, v, gate, beta, carry, window, enabled, scale, normalize, gate_inputs): + """Launch one token on private cache buffers; never mutate tensors saved for backward.""" + anchor = carry.anchor + state, scales, indices = _storage(anchor.values, enabled) + if enabled: + metadata = anchor.scales + codes = (anchor.values.unflatten(-1, (-1, 32)) / metadata.float().unsqueeze(-1)).round() + state[1] = codes.flatten(-2).transpose(-1, -2).to(torch.int8) + scales[1] = metadata.transpose(-1, -2) + else: + state[1] = anchor.values.transpose(-1, -2) + heads, keys, values = anchor.values.shape + d = torch.zeros(2, heads, window, values, device=q.device, dtype=torch.bfloat16) + kc = torch.zeros(2, heads, window, keys, device=q.device, dtype=torch.bfloat16) + channel = gate.ndim == 2 + gc = torch.zeros( + (2, heads, window, keys) if channel else (2, heads, window), + device=q.device, + dtype=torch.float32, + ) + for i, entry in enumerate(carry.entries): + d[1, :, i] = entry.update.values + kc[1, :, i] = entry.key.values + gc[1, :, i] = entry.log_retention + mixed = torch.cat([x.to(torch.bfloat16).flatten() for x in (q, k, v)])[None] + if channel: + a, b, a_log, bias = gate.flatten()[None], beta[None], None, None + else: + raw_gate, raw_beta, a_log, bias = gate_inputs + a, b = raw_gate[None].contiguous(), raw_beta[None].contiguous() + output = torch.empty(1, heads, values, device=q.device, dtype=torch.bfloat16) + cursor = carry.cursor + fused_recurrent_gated_delta_rule_replayssm( + mixed, + a, + b, + a_log, + bias, + scale, + state, + d, + kc, + gc, + output, + indices, + torch.full((1,), cursor, device=q.device, dtype=torch.int32), + use_qk_l2norm_in_kernel=normalize, + quant_state_bits=8 if enabled else 0, + state_scale=scales, + hadamard_value_basis=True, + block_v=64, + nk=2 if keys == 32 else None, + _is_kda=channel, + ) + decoded, metadata = _decoded(state, scales) + return ( + output[0].float(), + decoded, + metadata, + kc[1, :, cursor].float(), + d[1, :, cursor].float(), + gc[1, :, cursor], + ) diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index 010d7fca426..ae86af71400 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -75,15 +75,16 @@ class LinearAttentionDecodeConfig(ModeloptBaseConfig): ``precision='vllm_0_15'`` uses the pinned serving forward arithmetic for the chunked prefix and token suffix, with a differentiable Torch adjoint. - ``full`` retains the FP32/FP64 Torch reference. Quantization is configured - independently by TensorQuantizer; this profile does not enable QDQ. + ``replayssm`` imports INT8/Hadamard checkpoint and ring kernels from the + quantized-ReplaySSM serving fork. ``full`` retains the FP32/FP64 Torch reference. + TensorQuantizer configures quantization independently; this profile does not enable QDQ. """ mode: Literal["token", "replay"] = ModeloptField(default="token") readout: Literal["working", "stored"] = ModeloptField(default="stored") quantize_initial: bool = ModeloptField(default=True) prefill_state_qdq: bool = ModeloptField(default=False) - precision: Literal["full", "vllm_0_15"] = ModeloptField(default="full") + precision: Literal["full", "vllm_0_15", "replayssm"] = ModeloptField(default="full") 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) @@ -115,23 +116,32 @@ def _validate_replay(self): "vllm_0_15 requires token mode, working readout, tile codec, initial-state QDQ, " "and no additional prefill or gate rounding" ) + if self.precision == "replayssm" and ( + self.readout != "working" + or self.state_codec != "int8_hadamard32" + or not self.quantize_initial + or self.prefill_state_qdq + or self.decay_log_step is not None + or (self.replay is not None and self.replay.encoding != "once") + ): + raise ValueError( + "replayssm requires working readout, int8_hadamard32, initial-state QDQ, " + "encode-once replay, and no additional prefill or gate rounding" + ) 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. - With INT8 ``TensorQuantizer.block_sizes``, scales instead follow per-key value - groups; ``state.block_v`` controls only execution tiling (at least one group). + """Execution policy shared by GDN/KDA state QAT and retained GDN W QAT. + + ``serving`` uses native forward arithmetic with a differentiable adjoint; + ``reference`` retains explicit mathematical experiments. ``fla`` is the + disabled/default or W-only path. State formats and grouping come from + TensorQuantizer; ``state.block_v`` controls legacy tile grouping. """ schema_version: Literal[1, 2] = ModeloptField(default=2) - backend: Literal["fla", "matmul"] = ModeloptField(default="fla") + backend: Literal["fla", "serving", "reference"] = ModeloptField(default="fla") chunk_size: Literal[64] = ModeloptField(default=64) state: _StateConfig = ModeloptField(default=_StateConfig()) decode: LinearAttentionDecodeConfig | None = ModeloptField(default=None) @@ -139,6 +149,19 @@ class LinearAttentionConfig(ModeloptBaseConfig): @model_validator(mode="before") @classmethod def _load_legacy_defaults(cls, values): + if isinstance(values, dict) and values.get("backend") == "matmul": + decode = values.get("decode") or {} + precision = ( + decode.get("precision", "full") if isinstance(decode, dict) else decode.precision + ) + values = {**values, "backend": "reference" if precision == "full" else "serving"} + if isinstance(values, dict) and values.get("backend") == "serving": + decode = values.get("decode") + if isinstance(decode, dict): + values = { + **values, + "decode": {"precision": "vllm_0_15", "readout": "working", **decode}, + } if isinstance(values, dict) and values.get("schema_version") == 1: decode = values.get("decode") if isinstance(decode, dict) and isinstance(decode.get("replay"), dict): @@ -151,8 +174,14 @@ def _load_legacy_defaults(cls, values): @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.backend != "fla") != (self.decode is not None): + raise ValueError("State QAT requires an explicit decode policy") + if self.decode is not None and (self.backend == "serving") != ( + self.decode.precision != "full" + ): + raise ValueError( + "Serving state QAT requires a native precision profile; full uses reference" + ) 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") diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py index 2e995c47242..74694ada60e 100644 --- a/modelopt/torch/quantization/linear_attention/decode.py +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -59,6 +59,7 @@ class LinearAttentionCarry: started: bool signature: str value_basis: str = "identity" + precision: str = "full" @property def cursor(self): @@ -67,17 +68,28 @@ def cursor(self): 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 self.precision == "replayssm": + state = _replay_state(self) + else: + 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) + if self.precision == "replayssm": + # The optional serving fork owns the checkpoint basis transform. + from ...kernels.quantization.linear_attention.serving.replay import ( + original_basis as restore_basis, + ) + + state = restore_basis(state) + else: + state = _hadamard32(state) return state @@ -91,6 +103,29 @@ def _hadamard32(value): return value * (32**-0.5) +def _replay_state(carry, current_gate=None): + """Differentiate serving's weighted BF16 ring reconstruction, including KDA decay.""" + state = carry.anchor.values + if not carry.entries: + if current_gate is not None: + state = state * current_gate.exp().unsqueeze(-1) + return state + gates = torch.stack([entry.log_retention for entry in carry.entries], dim=1) + keys = torch.stack([entry.key.values for entry in carry.entries], dim=1) + updates = torch.stack([entry.update.values for entry in carry.entries], dim=1) + total = gates.sum(1) + if current_gate is not None: + total = total + current_gate + weights = (total.unsqueeze(1) - gates.cumsum(1)).exp() + if gates.ndim == 3: + weighted = keys * weights + weighted = forward_value(weighted, weighted.to(torch.bfloat16)) + return state * total.exp().unsqueeze(-1) + weighted.transpose(-1, -2) @ updates + weighted = updates * weights.unsqueeze(-1) + weighted = forward_value(weighted, weighted.to(torch.bfloat16)) + return state * total.exp()[:, None, None] + keys.transpose(-1, -2) @ weighted + + def _encode( value, enabled, @@ -233,8 +268,22 @@ def _prepare_carry( state_codec=config.state_codec, state_quantizer=state_quantizer, ) + if config.precision == "replayssm": + # The optional serving fork owns checkpoint rounding and Hadamard arithmetic. + from ...kernels.quantization.linear_attention.serving.replay import checkpoint + + with torch.no_grad(): + decoded, metadata = checkpoint(carry.anchor.values, state_qdq) + anchor.values = forward_value(initial, decoded) + anchor.scales = metadata carry = LinearAttentionCarry( - anchor, (), carry.position, True, signature, "hadamard32" if hadamard else "identity" + anchor, + (), + carry.position, + True, + signature, + "hadamard32" if hadamard else "identity", + config.precision, ) return carry, signature @@ -258,6 +307,7 @@ def recurrent_decode( position=0, scale=None, use_qk_l2norm_in_kernel=False, + replay_gate_inputs=None, ): """Run one preactivated sequence [T,H,D] with explicit token/replay write events. @@ -270,14 +320,29 @@ def recurrent_decode( state_quantizer, state_qdq, state_format ) serving = config.precision == "vllm_0_15" - if serving and q.device.type != "cuda": - raise ValueError("vllm_0_15 requires CUDA") + native_replay = config.precision == "replayssm" + if (serving or native_replay) and q.device.type != "cuda": + raise ValueError("Serving arithmetic requires CUDA") + if ( + native_replay + and len(q) + and ( + q.shape[-1] < 32 + or q.shape[-1] & (q.shape[-1] - 1) + or (g.ndim == 2 and replay_gate_inputs is None) + ) + ): + raise ValueError("ReplaySSM requires power-of-two K >= 32 and raw GDN gate inputs") replay_quantizers = (replay_key_quantizer, replay_update_quantizer) if config.replay is not None and config.replay._legacy_factor_qdq: if any(quantizer is None for quantizer in replay_quantizers): raise ValueError("Replace factor_qdq with explicit replay key/update TensorQuantizers") for quantizer in replay_quantizers: if quantizer is not None and quantizer.is_enabled: + if native_replay: + raise ValueError( + "Native ReplaySSM stores BF16 factors; additional factor QDQ is unsupported" + ) if config.replay is None or not quantizer.fake_quant: raise ValueError( "Replay factor quantizers require replay mode and fake quantization" @@ -303,10 +368,12 @@ def recurrent_decode( zero = (q.sum() + k.sum() + v.sum() + g.sum() + beta.sum()) * 0 output = v + zero return output, carry + original_v = v if config.state_codec == "int8_hadamard32": v = _hadamard32(v) scale = q.shape[-1] ** -0.5 if scale is None else scale outputs = [] + native_outputs = [] with torch.autocast(device_type=q.device.type, enabled=False): for t in range(len(q)): entries = carry.entries @@ -324,6 +391,30 @@ def recurrent_decode( ) state = carry.reconstruct(original_basis=False) gate = _round_log_gate(g[t], config.decay_log_step) + if native_replay: + from ...kernels.quantization.linear_attention.serving.replay import step + + with torch.no_grad(): + native = step( + q[t], + k[t], + original_v[t], + gate, + beta[t], + carry, + config.replay.window if config.replay is not None else 1, + state_qdq, + scale, + use_qk_l2norm_in_kernel, + ( + replay_gate_inputs[0][t], + replay_gate_inputs[1][t], + *replay_gate_inputs[2:], + ) + if replay_gate_inputs is not None + else None, + ) + native_outputs.append(native[0]) if serving: from ._vllm_autograd import step as serving_step @@ -334,12 +425,23 @@ def recurrent_decode( decay = gate.exp().unsqueeze(-1) if gate.ndim == 1: decay = decay.unsqueeze(-1) - key = _encode_factor(k[t], replay_key_quantizer) + current_key = k[t] + if native_replay and use_qk_l2norm_in_kernel: + current_key = ( + current_key / (current_key.square().sum(-1, keepdim=True) + 1e-6).sqrt() + ) + key = _encode_factor(current_key, replay_key_quantizer) decayed = state * decay + if native_replay and gate.ndim == 2: + decayed = _replay_state(carry, gate) residual = v[t] - _sum_keys(key.values.unsqueeze(-1) * decayed) update = _encode_factor(beta[t].unsqueeze(-1) * residual, replay_update_quantizer) working = decayed + key.values.unsqueeze(-1) * update.values.unsqueeze(-2) if config.replay is not None: + if native_replay and carry.cursor + 1 < config.replay.window: + key.values = forward_value(key.values, native[3]) + update.values = forward_value(update.values, native[4]) + gate = forward_value(gate, native[5]) entries = (*entries, ReplayEntry(key, update, gate)) refresh = config.replay is None or len(entries) == config.replay.window if refresh: @@ -352,20 +454,34 @@ def recurrent_decode( state_codec=config.state_codec, state_quantizer=state_quantizer, ) + if native_replay: + anchor.values = forward_value(working, native[1]) + anchor.scales = native[2] entries = () stored = anchor.values else: anchor = carry.anchor stored = working next_carry = LinearAttentionCarry( - anchor, entries, carry.position + 1, True, signature, carry.value_basis + anchor, + entries, + carry.position + 1, + True, + signature, + carry.value_basis, + config.precision, ) read = working if config.readout == "working" else stored + query = q[t] + if native_replay and use_qk_l2norm_in_kernel: + query = query / (query.square().sum(-1, keepdim=True) + 1e-6).sqrt() outputs.append( - native_output if serving else _sum_keys(q[t].unsqueeze(-1) * read) * scale + native_output if serving else _sum_keys(query.unsqueeze(-1) * read) * scale ) carry = next_carry output = torch.stack(outputs) if config.state_codec == "int8_hadamard32": output = _hadamard32(output) + if native_replay: + output = forward_value(output, torch.stack(native_outputs)) return output, carry diff --git a/modelopt/torch/quantization/linear_attention/prefill.py b/modelopt/torch/quantization/linear_attention/prefill.py index 2f67670c77b..7bbb42deb09 100644 --- a/modelopt/torch/quantization/linear_attention/prefill.py +++ b/modelopt/torch/quantization/linear_attention/prefill.py @@ -51,6 +51,7 @@ def matmul_gdn( chunk_size=64, cp_context=None, prefill_lengths=None, + replay_gate_inputs=None, ): """Normalize GDN inputs and run a chunked prefix plus configured suffix recurrence. @@ -68,7 +69,7 @@ def matmul_gdn( dtype = q.dtype if serving and (use_gate_in_kernel or use_beta_sigmoid_in_kernel): raise ValueError( - "vllm_0_15 GDN expects prepared log gates and beta from the Megatron adapter" + "Serving GDN expects prepared log gates and beta from the Megatron adapter" ) if use_gate_in_kernel: if A_log is None or dt_bias is None: @@ -99,5 +100,6 @@ def matmul_gdn( output_dtype=output_dtype, beta_dtype=beta_dtype, prefill_lengths=prefill_lengths, + replay_gate_inputs=replay_gate_inputs, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel if serving else False, ) diff --git a/modelopt/torch/quantization/linear_attention/training.py b/modelopt/torch/quantization/linear_attention/training.py index e5319410368..fa9987e6296 100644 --- a/modelopt/torch/quantization/linear_attention/training.py +++ b/modelopt/torch/quantization/linear_attention/training.py @@ -22,7 +22,7 @@ from ._chunk_prefill import chunk_gdn, chunk_kda from .decode import recurrent_decode -from .utils import _resolve_state_quantizer, _state_qdq +from .utils import _resolve_state_quantizer, _state_qdq, forward_value __all__ = ["linear_attention_training_phase"] @@ -67,13 +67,15 @@ def linear_attention_training_phase(model, prefill_lengths): def _prepare_prefill_inputs(q, k, v, g, beta, *, policy, chunk_size, normalize): """Validate the shared policy and prepare GDN/KDA working dtypes and Q/K normalization.""" - if policy.backend != "matmul" or chunk_size != policy.chunk_size: - raise ValueError("Matmul prefill requires backend='matmul' and its configured chunk size") + if policy.backend not in ("serving", "reference") or chunk_size != policy.chunk_size: + raise ValueError( + "State training requires backend='serving' or 'reference' and its configured chunk size" + ) if policy.decode is None: raise ValueError("An explicit decode policy is required") - serving = policy.decode.precision == "vllm_0_15" + serving = policy.decode.precision != "full" if serving and (q.device.type != "cuda" or any(x.dtype != torch.bfloat16 for x in (q, k, v))): - raise ValueError("precision='vllm_0_15' requires CUDA BF16 Q/K/V inputs") + raise ValueError("Serving arithmetic requires CUDA BF16 Q/K/V inputs") dtype = torch.float64 if q.dtype == torch.float64 else torch.float32 q, k, v, g, beta = (x.to(dtype) for x in (q, k, v, g, beta)) if normalize and not serving: @@ -104,6 +106,7 @@ def _prefill_decode_forward( beta_dtype, prefill_lengths, use_qk_l2norm_in_kernel=False, + replay_gate_inputs=None, ): """Run both prefill and decode phases in one differentiable training forward. @@ -149,14 +152,14 @@ def _prefill_decode_forward( p > end - start for p, (_, start, end) in zip(prefixes, sequences) ): raise ValueError("Supply one valid prefill length per sequence") - serving_precision = policy.decode.precision == "vllm_0_15" + serving_precision = policy.decode.precision != "full" if serving_precision and ( keys > 256 or (g.ndim == 4 and values != keys) or (use_qk_l2norm_in_kernel and keys & (keys - 1)) ): raise ValueError( - "vllm_0_15 requires K <= 256, KDA V=K, and power-of-two K for normalization" + "Serving arithmetic requires K <= 256, KDA V=K, and power-of-two K for normalization" ) if initial_state is None: states = q.new_zeros(len(sequences), heads, keys, values) @@ -178,9 +181,20 @@ def _prefill_decode_forward( # A continuation prefill consumes a stored cache just as native serving # does. A fresh zero-state prefix has no incoming cache to quantize. if initial_state is not None and state_qdq: - state = _state_qdq( - state, policy.state.block_v, state_format, state_quantizer - ) + if policy.decode.precision == "replayssm": + from ...kernels.quantization.linear_attention.serving.replay import ( + checkpoint, + original_basis, + ) + + with torch.no_grad(): + decoded, _ = checkpoint(state, True) + decoded = original_basis(decoded) + state = forward_value(state, decoded) + else: + state = _state_qdq( + state, policy.state.block_v, state_format, state_quantizer + ) prefix, state = serving_prefix( q=q[b, start:split], k=k[b, start:split], @@ -216,6 +230,15 @@ def _prefill_decode_forward( block_v=policy.state.block_v, initial_state=state, position=prefixes[n], + replay_gate_inputs=( + ( + replay_gate_inputs[0][b, split:end], + replay_gate_inputs[1][b, split:end], + *replay_gate_inputs[2:], + ) + if replay_gate_inputs is not None + else None + ), use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, scale=scale, ) diff --git a/modelopt/torch/quantization/plugins/gdn.py b/modelopt/torch/quantization/plugins/gdn.py index 1f00ef8c20a..93c561fc2fd 100644 --- a/modelopt/torch/quantization/plugins/gdn.py +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -13,18 +13,11 @@ # 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. +"""GatedDeltaNet state and WY activation fake quantization. + +State QAT uses an explicit prefix/suffix policy with native serving arithmetic. +W-only QAT applies TensorQuantizer to the WY-transformed keys and reuses FLA's +state kernels. Both paths retain gradients through quantization with identity STE. """ from collections.abc import Callable @@ -49,7 +42,7 @@ def _fla_chunk_gated_delta_rule() -> GatedDeltaRuleFn: return chunk_gated_delta_rule -def _state_qdq_chunk_gated_delta_rule() -> GatedDeltaRuleFn: +def _w_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: @@ -59,7 +52,7 @@ def _state_qdq_chunk_gated_delta_rule() -> GatedDeltaRuleFn: 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}." + f"device; importing the W-quantizing kernel failed with {e!r}." ) from e return chunk_gated_delta_rule @@ -71,9 +64,10 @@ class GatedDeltaNetStateQuantMixin(_LinearAttentionQuantMixin): :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 + quantizer carries the 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. + Both sites require identity STE; state QAT uses an explicit serving policy. + The execution policy is saved in ModelOpt metadata. """ linear_attention_quantizer_names = ("gdn_state_quantizer", "gdn_w_quantizer") @@ -94,11 +88,17 @@ def gdn_state_qdq_block_v(self) -> int: 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.""" + """Route enabled state or W quantization to its training implementation.""" self.validate_linear_attention() quantize_w = self.gdn_w_quantizer.is_enabled if not self.linear_attention_is_enabled: return gated_delta_rule(*args, **kwargs) + if self.linear_attention_config.backend == "fla" and self.gdn_state_quantizer.is_enabled: + raise ValueError( + "Chunk-only state QAT is retired. Select backend='serving' with an explicit " + "decode policy and supply prefill lengths through linear_attention_training_phase. " + "Existing state checkpoints are not silently migrated to a different QDQ schedule." + ) while isinstance(gated_delta_rule, partial): args = (*gated_delta_rule.args, *args) kwargs = {**gated_delta_rule.keywords, **kwargs} @@ -111,7 +111,7 @@ def _state_quantized_chunk_gated_delta_rule( 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": + if self.linear_attention_config.backend != "fla": return matmul_gdn( *args, policy=self.linear_attention_config, @@ -122,11 +122,9 @@ def _state_quantized_chunk_gated_delta_rule( prefill_lengths=self._linear_attention_prefill_lengths, **kwargs, ) - return _state_qdq_chunk_gated_delta_rule()( + return _w_qdq_chunk_gated_delta_rule()( *args, chunk_size=chunk_size, - state_quantizer=self.gdn_state_quantizer, - 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 index 52218dbffe8..820ee3abad5 100644 --- a/modelopt/torch/quantization/plugins/kda.py +++ b/modelopt/torch/quantization/plugins/kda.py @@ -40,8 +40,8 @@ class KimiDeltaAttentionStateQuantMixin(_LinearAttentionQuantMixin): def validate_linear_attention(self): """Require the materialized backend for KDA numerical emulation.""" 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'") + if self.linear_attention_is_enabled and self.linear_attention_config.backend == "fla": + raise ValueError("KDA numerical emulation requires backend='serving'") def _state_quantized_chunk_kda(self, kernel, *args, **kwargs): self.validate_linear_attention() diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 3791da54e1e..1b00e52beca 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -1122,6 +1122,7 @@ class _MegatronLinearAttentionMixin(_LinearAttentionQuantMixin): def _setup(self): super()._setup() + self._register_temp_attribute("_linear_attention_replay_gate_inputs", None) try: data_parallel_group = get_data_parallel_group(with_context_parallel=True) except AssertionError: @@ -1135,9 +1136,7 @@ def _setup(self): def _serving_arithmetic(self): decode = self.linear_attention_config.decode return ( - self.linear_attention_is_enabled - and decode is not None - and decode.precision == "vllm_0_15" + self.linear_attention_is_enabled and decode is not None and decode.precision != "full" ) def _prepare_input_for_gated_delta_rule(self, *args, **kwargs): @@ -1160,6 +1159,8 @@ def _compute_gates(self, a_log, dt_bias, batch, seq_len, *gate_feats): from ..linear_attention.utils import forward_value raw_beta, raw_gate = gate_feats + if self.linear_attention_config.decode.precision == "replayssm": + self._linear_attention_replay_gate_inputs = (raw_gate, raw_beta, a_log, dt_bias) with torch.no_grad(): native_gate, native_beta = fused_gdn_gating( a_log, @@ -1176,6 +1177,7 @@ def _compute_gates(self, a_log, dt_bias, batch, seq_len, *gate_feats): @contextmanager def _quantized_linear_attention_kernel(self): kernel = self.gated_delta_rule + previous_gates = self._linear_attention_replay_gate_inputs if self._serving_arithmetic: if not hasattr(super(), "_prepare_input_for_gated_delta_rule") or getattr( self, "gdn_pre_gated_delta_rule_fusion", False @@ -1186,15 +1188,19 @@ def _quantized_linear_attention_kernel(self): def run_kernel(*args, **kwargs): kwargs["use_qk_l2norm_in_kernel"] = self.use_qk_l2norm + if self._linear_attention_replay_gate_inputs is not None: + kwargs["replay_gate_inputs"] = self._linear_attention_replay_gate_inputs return self._linear_attention_kernel(kernel, *args, **kwargs) self.gated_delta_rule = run_kernel else: self.gated_delta_rule = partial(self._linear_attention_kernel, kernel) + self._linear_attention_replay_gate_inputs = None try: yield finally: self.gated_delta_rule = kernel + self._linear_attention_replay_gate_inputs = previous_gates def forward(self, *args, **kwargs): # Newer Megatron versions recompute the core independently during backward. diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index a867006aadd..f40d76da3a3 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -33,10 +33,10 @@ 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_fp8_dynamic.yaml` | FP8 E4M3 dynamic fake quantization of the GatedDeltaNet recurrent state (per sequence, head and 64-column tile) at serving cache boundaries; requires an explicit serving execution policy | | `linear_attention_state_int8_dynamic.yaml` | GDN/KDA INT8 state quantizer entries; the [complete recipe](../../../general/ptq/linear_attention_state_int8_dynamic.yaml) enables Hadamard during decode | | `linear_attention_state_int8_block32_dynamic.yaml` | Standard dynamic INT8 with one scale per key row and 32 value channels; the [complete training recipe](../../../general/ptq/linear_attention_state_int8_block32_dynamic.yaml) uses token QDQ and working-state readout | -| `linear_attention_replay_fp8_dynamic.yaml` | GDN/KDA ReplaySSM key/update FP8 quantizers; requires a replay execution policy | +| `linear_attention_replay_fp8_dynamic.yaml` | GDN/KDA ReplaySSM key/update FP8 quantizers; for the explicit reference backend; native ReplaySSM uses BF16 ring entries | | `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 | | `indexer_k_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer key cache (`*indexer_k_quantizer`), global scale fixed to 1; Blackwell+ GPUs | | `indexer_q_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer query (`*indexer_q_quantizer`), global scale fixed to 1; Blackwell+ GPUs | @@ -51,7 +51,7 @@ scheduling only; for example: linear_attention: - module_name: "*" cfg: - backend: matmul + backend: reference decode: mode: replay replay: {window: 8, encoding: once} diff --git a/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml index ac1d078bc01..272d3ce2825 100644 --- a/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml +++ b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml @@ -32,7 +32,7 @@ quantize: linear_attention: - module_name: '*' cfg: - backend: matmul + backend: serving decode: mode: token state_codec: tile diff --git a/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml b/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml index 3cc24515e25..0c07838b8ad 100644 --- a/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml +++ b/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml @@ -21,6 +21,7 @@ imports: metadata: description: >- Dynamic INT8 GDN/KDA state QDQ with 32-value Hadamard rotation during decode. + Imports the quantized-ReplaySSM serving kernels and uses working-state readout. Prefill state remains unquantized; supply per-sequence prefix lengths through linear_attention_training_phase. No calibration is required. quantize: @@ -31,8 +32,10 @@ quantize: linear_attention: - module_name: '*' cfg: - backend: matmul + backend: serving decode: mode: token + precision: replayssm + readout: working state_codec: int8_hadamard32 prefill_state_qdq: false 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 index 3077229a1f0..a3fa7585f56 100644 --- 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 @@ -61,31 +61,16 @@ def compare(actual, expected, tolerance): ) -@pytest.fixture( - scope="module", params=["disabled", "w", "state-w", "state-int8", "state-int8-block"] -) +@pytest.fixture(scope="module", params=[False, True]) def compiled_gdn_case(request): - """Compile only the selected BF16 forward/backward path, outside the test-call timer.""" - state_qdq = {"state-w": 1, "state-int8": 2}.get(request.param, 0) - if state_qdq == 1 and torch.cuda.get_device_capability() < (8, 9): - pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") + """Compile the retained W-QAT path outside the test-call timer.""" quantizer = ( TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), axis=(0, 1, 2), type="dynamic")) - if request.param in ("w", "state-w") + if request.param else None ) args, state = make_inputs() - kwargs = {"state_qdq": state_qdq, "w_quantizer": quantizer} - if state_qdq == 2: - TensorQuantizer(QuantizerAttributeConfig(num_bits=8, type="dynamic"))(state) - if request.param == "state-int8-block": - kwargs["state_quantizer"] = TensorQuantizer( - QuantizerAttributeConfig( - num_bits=8, type="dynamic", block_sizes={-1: 16}, narrow_range=True - ) - ) - # Compile the reference quantizer's CUDA extension outside the test-call timer too. - kwargs["state_quantizer"](state[0]) + kwargs = {"w_quantizer": quantizer} values_and_grads(chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs) torch.cuda.synchronize() return args, state, kwargs @@ -99,7 +84,6 @@ def test_gdn_forward_and_backward(compiled_gdn_case): chunk_gdn_reference, reference_args, reference_state, - state_format="int8" if kwargs["state_qdq"] == 2 else "fp8_e4m3", **kwargs, ) actual = values_and_grads( @@ -107,20 +91,3 @@ def test_gdn_forward_and_backward(compiled_gdn_case): ) compare(actual[0], expected[0], 0.03) compare(actual[1], expected[1], 0.05) - - if kwargs["state_qdq"] == 2 or "state_quantizer" in kwargs: - quantizer = kwargs.get("state_quantizer") or TensorQuantizer( - QuantizerAttributeConfig(num_bits=8, type="dynamic") - ) - # Zero updates isolate QDQ, reusing the compiled shape for both tile and block INT8. - zero_args = [torch.zeros_like(x) for x in args] - pattern = torch.linspace(-1, 1, state.numel(), device=state.device).reshape_as(state) - for amax in (2**-25, 2**-24, 2.0): - initial = pattern * amax - expected_state = initial - for _ in range(3): # Initial handoff and two chunk boundaries. - expected_state = quantizer(expected_state[0]).unsqueeze(0) - _, actual_state = chunk_gated_delta_rule( - *zero_args, initial_state=initial, output_final_state=True, **kwargs - ) - torch.testing.assert_close(actual_state, expected_state, rtol=0, atol=0) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index 123fa1f68c5..b9e97aa85ba 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -34,11 +34,11 @@ 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.gdn import _w_qdq_chunk_gated_delta_rule from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet try: - _state_qdq_chunk_gated_delta_rule() + _w_qdq_chunk_gated_delta_rule() except RuntimeError as e: pytest.skip(str(e), allow_module_level=True) @@ -94,8 +94,7 @@ def _gdn_config(sites): def _config_for_mode(mode): - # Ampere exercises W QDQ; native FP8 devices also exercise chunk state QDQ. - sites = ("state", "w") if torch.cuda.get_device_capability() >= (8, 9) else ("w",) + sites = ("w",) if mode == "decode_replay_int8": sites = ("state",) cfg = _gdn_config(sites) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py index 57e36b226bb..51b6c3f2790 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py @@ -163,8 +163,10 @@ def _test_kda(rank, size, cfg, checkpoint_path): def compiled_kda_workers(dist_workers_size_1): """Warm one KDA shape outside the functional test timer.""" cfg = load_recipe("general/ptq/linear_attention_state_int8_dynamic").quantize.model_dump() + # Exercise mathematical replay/checkpointing without the optional serving fork. + cfg["linear_attention"][0]["cfg"]["backend"] = "reference" cfg["linear_attention"][0]["cfg"]["decode"].update( - mode="replay", replay={"window": 5}, decay_log_step=1 / 256 + precision="full", mode="replay", replay={"window": 5}, decay_log_step=1 / 256 ) cfg["quant_cfg"].extend( entry.model_dump(exclude_unset=True) diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py new file mode 100644 index 00000000000..f8084a88122 --- /dev/null +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py @@ -0,0 +1,148 @@ +# 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. + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Minimal native-cache and backward checks for the optional ReplaySSM serving fork.""" + +import pytest +import torch + +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionDecodeConfig, + recurrent_decode, +) + + +@pytest.fixture(scope="module", params=[False, True], ids=["gdn-token", "kda-replay"]) +def compiled_replay(request): + native = pytest.importorskip( + "vllm.model_executor.layers.fla.ops.fused_recurrent_replayssm", + exc_type=ModuleNotFoundError, + ) + channel = request.param + window = 4 if channel else 1 + torch.manual_seed(71) + q, k = [ + torch.nn.functional.normalize(torch.randn(5, 1, 64, device="cuda"), dim=-1).bfloat16() + for _ in range(2) + ] + v = torch.randn_like(q) + raw = torch.randn(5, 1, device="cuda", dtype=torch.bfloat16) + raw_beta = torch.randn_like(raw) + rate = torch.full((1,), -3.0, device="cuda") + bias = torch.zeros_like(rate) + g = ( + -torch.rand_like(q, dtype=torch.float32) * 0.03 + if channel + else -rate.exp() * torch.nn.functional.softplus(raw.float()) + ) + beta = raw_beta.float().sigmoid() + args = tuple(x.detach().requires_grad_() for x in (q, k, v, g, beta)) + initial = (torch.randn(1, 64, 64, device="cuda") * 0.1).requires_grad_() + config = LinearAttentionDecodeConfig( + precision="replayssm", + state_codec="int8_hadamard32", + readout="working", + mode="replay" if channel else "token", + replay={"window": window} if channel else None, + ) + kwargs = { + "config": config, + "state_qdq": True, + "state_format": "int8", + "replay_gate_inputs": None if channel else (raw, raw_beta, rate, bias), + } + + def forward(): + return recurrent_decode(*args, initial_state=initial, **kwargs) + + out, carry = forward() + torch.autograd.grad(out.sum() + carry.reconstruct().sum(), (*args, initial)) + torch.cuda.synchronize() + return native, args, initial, kwargs, forward + + +def test_replay_matches_persistent_serving_cache(compiled_replay): + native, args, initial, kwargs, forward = compiled_replay + output, carry = forward() + channel = args[3].ndim == 3 + window = kwargs["config"].replay.window if channel else 1 + state = torch.zeros(2, 1, 64, 64, device="cuda", dtype=torch.int8) + scales = torch.zeros(2, 1, 2, 64, device="cuda", dtype=torch.float16) + updates = torch.zeros(2, 1, window, 64, device="cuda", dtype=torch.bfloat16) + keys = torch.zeros_like(updates) + gates = torch.zeros((2, 1, window, 64) if channel else (2, 1, window), device="cuda") + indices = torch.ones(1, device="cuda", dtype=torch.int32) + with torch.no_grad(): + native.prefill_write_checkpoint( + initial.transpose(-1, -2)[None].contiguous(), + state, + None, + None, + scales, + None, + None, + indices, + 8, + hadamard=True, + ) + for t in range(len(args[0])): + q, k, v, gate, beta = (x[t] for x in args) + if channel: + a, b, rate, bias = gate.flatten()[None], beta[None], None, None + else: + raw, raw_beta, rate, bias = kwargs["replay_gate_inputs"] + a, b = raw[t : t + 1], raw_beta[t : t + 1] + expected = torch.empty(1, 1, 64, device="cuda", dtype=torch.bfloat16) + native.fused_recurrent_gated_delta_rule_replayssm( + torch.cat([x.flatten() for x in (q, k, v)])[None], + a, + b, + rate, + bias, + 64**-0.5, + state, + updates, + keys, + gates, + expected, + indices, + torch.tensor([t % window], device="cuda", dtype=torch.int32), + quant_state_bits=8, + state_scale=scales, + hadamard_value_basis=True, + block_v=64, + _is_kda=channel, + ) + torch.testing.assert_close(output[t].bfloat16(), expected[0], rtol=0, atol=0) + decoded = state[1].transpose(-1, -2).float().unflatten(-1, (-1, 32)) + metadata = scales[1].transpose(-1, -2) + torch.testing.assert_close( + carry.anchor.values, (decoded * metadata[..., None]).flatten(-2), rtol=0, atol=0 + ) + torch.testing.assert_close(carry.anchor.scales, metadata, rtol=0, atol=0) + for i, entry in enumerate(carry.entries): + torch.testing.assert_close( + entry.key.values.float(), keys[1, :, i].float(), rtol=0, atol=0 + ) + torch.testing.assert_close( + entry.update.values.float(), updates[1, :, i].float(), rtol=0, atol=0 + ) + grads = torch.autograd.grad( + output.square().mean() + carry.reconstruct().square().mean(), (*args, initial) + ) + assert all(torch.isfinite(grad).all() and grad.abs().sum() > 0 for grad in grads) diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py index deae03a2f4b..623be1bf48b 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py @@ -40,7 +40,7 @@ def compiled_serving_case(request): ] args = [x.requires_grad_() for x in args] policy = LinearAttentionConfig( - backend="matmul", decode={"precision": "vllm_0_15", "readout": "working"} + backend="serving", decode={"precision": "vllm_0_15", "readout": "working"} ) quantizer = TensorQuantizer( QuantizerAttributeConfig( diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index 60ce199473a..0d2d498e895 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -14,7 +14,6 @@ # limitations under the License. from copy import deepcopy -from functools import partial import pytest import torch @@ -23,7 +22,10 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.quantization.config import QuantizeConfig -from modelopt.torch.quantization.linear_attention import linear_attention_training_phase +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + linear_attention_training_phase, +) from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins import gdn from modelopt.torch.quantization.plugins.gdn import GatedDeltaNetStateQuantMixin @@ -131,55 +133,39 @@ def test_disabled_state_quantizer_calls_original_kernel(): assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" -@pytest.mark.parametrize("partial_kernel", [False, True]) -def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, partial_kernel): +def test_state_qat_requires_explicit_serving_policy(monkeypatch): + model = mtq.quantize(TinyGatedDeltaNet(), {**quant_cfg(), "algorithm": None}) + x = torch.randn(2, 8, 3, 4) + with pytest.raises(ValueError, match="Chunk-only state QAT is retired"): + model(x) calls = [] - def fake_state_qdq_kernel(*args, **kwargs): + def forward(*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) - mtq.quantize(model, quant_cfg(), lambda m: m(x)) - - model(x) - assert calls and calls[-1] == { - "chunk_size": 64, - "state_quantizer": model.gdn_state_quantizer, - "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"): + monkeypatch.setattr(gdn, "matmul_gdn", forward) + model.linear_attention_config = LinearAttentionConfig(backend="serving", decode={}) + with linear_attention_training_phase(model, [4, 4]): model(x) + assert calls[-1]["prefill_lengths"] == (4, 4) -@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.""" +def test_w_quantizer_is_passed_to_the_kernel(monkeypatch): + """Pass the configured W TensorQuantizer to the FLA wrapper.""" calls = [] - def fake_state_qdq_kernel(*args, **kwargs): + def fake_w_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) + monkeypatch.setattr(gdn, "_w_qdq_chunk_gated_delta_rule", lambda: fake_w_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 + mtq.quantize(model, quant_cfg(state=False, w=True), lambda m: m(x)) + assert model.gdn_w_quantizer.is_enabled and not model.gdn_state_quantizer.is_enabled model(x) - assert calls[-1]["state_quantizer"] is model.gdn_state_quantizer 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, From 309010ab8b3499b811bca13c3496768cf9f1b57c Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 6 Oct 2026 22:50:07 -0700 Subject: [PATCH 13/27] Simplify serving-aligned linear-attention state QAT Unify prefix, decode, and replay execution settings in LinearAttentionConfig with migration for supported nested configs and saved policy objects. Keep formats and grouping in TensorQuantizer. Retire W-only and reference training paths, move mathematical references under tests, and remove the copied FLA W-QAT implementation and unused recipes. Validation: 44 focused CPU tests, 6 native and Megatron GPU tests, legacy checkpoint loading, recipe validation, and pre-commit hooks passed. Signed-off-by: Kai Xu --- .pre-commit-config.yaml | 1 - CHANGELOG.rst | 4 +- .../quantization/linear_attention/__init__.py | 6 +- .../fla_chunk_gated_delta_rule.py | 695 ------------------ .../linear_attention/serving/__init__.py | 2 +- modelopt/torch/quantization/conversion.py | 1 - .../linear_attention/_chunk_prefill.py | 126 ---- .../quantization/linear_attention/config.py | 269 ++++--- .../quantization/linear_attention/decode.py | 116 +-- .../quantization/linear_attention/kda.py | 33 +- .../quantization/linear_attention/prefill.py | 21 +- .../quantization/linear_attention/training.py | 103 +-- .../quantization/linear_attention/utils.py | 2 +- modelopt/torch/quantization/plugins/gdn.py | 77 +- .../quantization/plugins/linear_attention.py | 37 +- .../torch/quantization/plugins/megatron.py | 7 +- modelopt_recipes/configs/ptq/units/README.md | 22 +- .../configs/ptq/units/gdn_w_fp8_dynamic.yaml | 28 - .../linear_attention_replay_fp8_dynamic.yaml | 28 - ..._attention_state_int8_block32_dynamic.yaml | 8 +- .../linear_attention_state_int8_dynamic.yaml | 7 +- modelopt_recipes/ptq.md | 2 +- pyproject.toml | 9 - .../linear_attention_reference.py | 171 +++-- .../test_fla_chunk_gated_delta_rule.py | 93 --- .../plugins/test_megatron_gated_delta_net.py | 84 +-- .../quantization/plugins/test_megatron_kda.py | 24 +- .../test_linear_attention_replay.py | 15 +- .../test_linear_attention_training.py | 4 +- .../torch/quantization/plugins/test_gdn.py | 185 ++--- .../test_linear_attention_decode.py | 135 ++-- .../test_linear_attention_hadamard.py | 78 +- .../test_linear_attention_reference.py | 13 +- 33 files changed, 549 insertions(+), 1857 deletions(-) delete mode 100644 modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py delete mode 100644 modelopt/torch/quantization/linear_attention/_chunk_prefill.py delete mode 100644 modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml delete mode 100644 modelopt_recipes/configs/ptq/units/linear_attention_replay_fp8_dynamic.yaml delete mode 100644 tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 43f4049f807..a5a3b97cd8c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -113,7 +113,6 @@ repos: modelopt/torch/export/transformer_engine.py| modelopt/torch/fastgen/plugins/qwen_image_pdd.py| modelopt/torch/kernels/quantization/linear_attention/serving/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 d20f81453ec..c139abb7377 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -34,7 +34,6 @@ Changelog - 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 WY activations through the standard ``quant_cfg`` interface. W QAT uses ``fla-core==0.5.1`` with chunk size 64; state QAT uses the serving-aligned workflow. - Add fake quantization of the sparse-attention indexer key cache and query for DeepSeek-V4 (vLLM and Megatron-Core) and GLM-5.3-Flash (vLLM) through the new ``indexer_k_quantizer`` and ``indexer_q_quantizer``. Enable them by importing the ``configs/ptq/units/indexer_k_nvfp4`` and ``configs/ptq/units/indexer_q_nvfp4`` units (NVFP4 with the global scale fixed to 1) into a recipe. - Add the ``configs/ptq/units/kv_nvfp4_mla`` recipe unit for fake quantization of the MLA KV cache (DeepSeek-V3, GLM-5.3-Flash, ...) in vLLM fake-quant serving: NVFP4 for the latent and FP8 for the RoPE key. See `examples/vllm_serve/README.md `_ for an example recipe. - Add fake quantization of MoE expert-parallel communication to vLLM fake-quant serving through the new ``dispatch_quantizer`` and ``combine_quantizer`` on vLLM's routed experts, for MoE layers built on vLLM's fused-MoE layer and any all-to-all backend. Enable them by importing the ``configs/ptq/units/moe_dispatch_nvfp4`` and ``configs/ptq/units/moe_combine_nvfp4`` units into a recipe; see ``examples/vllm_serve/README.md``. @@ -69,6 +68,9 @@ Changelog **Backward Breaking Changes** - Bump minimum transformers version to ``5.5`` instead of ``4.57``; transformers 4.x is no longer supported. Upgrade with ``pip install -U "nvidia-modelopt[hf]"``. +- Experimental linear-attention reference training is removed, including ``precision="full"`` and replay-factor QDQ; select ``backend="serving"`` with a native precision profile. Mathematical reference implementations are available only in tests. +- Experimental GDN W quantization and its ``gdn_w_fp8_dynamic`` recipe unit are removed; keep ``gdn_w_quantizer`` disabled when using state QAT or restoring existing checkpoints. + - ``modelopt.torch.distill.plugins.megatron.TopKLogitsKLLoss`` (``logit_kl_topk`` in ``DistillationConfig``) is renamed to ``TopLogitsKLLoss``, keeping the old name as a deprecated alias. It now normalizes both distributions over the full vocabulary instead of re-normalizing over the Top-K entries, and always appends a "ghost" token holding the probability mass outside the Top-K to both student and teacher (matching Megatron-LM's offline cached-logits KD loss). Loss values change for existing ``logit_kl_topk`` runs. - ``LogitsAndIntermediatesLossBalancer`` (Megatron distillation plugin) no longer rescales the distillation loss to the magnitude of the LM loss when the LM loss is included. The total is now the fixed convex combination ``(1 - alpha) * lm_loss + alpha * kd_loss`` with ``DistillationConfig.kd_loss_alpha`` (in [0, 1]), matching Megatron-LM's offline cached-logits KD. The default ``kd_loss_alpha=1.0`` skips the LM loss, as before. ``DistillationConfig.skip_lm_loss`` and ``kd_loss_scale`` are removed and raise ``ValueError`` if passed; set ``kd_loss_alpha`` instead. In ``examples/megatron_bridge/distill.py``, ``--kd_loss_alpha`` replaces ``--no_skip_lm_loss`` and ``--kd_loss_scale``. - The ``examples/vllm_serve`` fakequant launcher no longer supports vLLM 0.9.0. Upgrade to a version listed as tested in the example README. diff --git a/modelopt/torch/kernels/quantization/linear_attention/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/__init__.py index fb6476f0950..617bb6dcbd8 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/__init__.py +++ b/modelopt/torch/kernels/quantization/linear_attention/__init__.py @@ -13,8 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Linear-attention quantization adapters for optional FLA and serving kernels. +"""Linear-attention quantization adapters for optional serving kernels. -FLA supplies the state kernels used by W QAT. Serving adapters import forward -kernels only when selected; this package initializer has no kernel dependencies. +Serving adapters import forward kernels only when selected; this package +initializer has no kernel dependencies. """ 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 deleted file mode 100644 index dde6b75e3da..00000000000 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py +++ /dev/null @@ -1,695 +0,0 @@ -# Adapted from: https://github.com/fla-org/flash-linear-attention/blob/516143e31fce/fla/ops/gated_delta_rule/chunk.py -# Modifications: optional W fake QDQ, saved rounded W for backward recomputation, -# and explicit validation. State forward/backward kernels are imported from FLA. -# -# 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_delta_h import ( - chunk_gated_delta_rule_bwd_dhu, - chunk_gated_delta_rule_fwd_h, -) -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 - - -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, - 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, - ) - - 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, - 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) - - 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, - ) - 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, - 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, - 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 - 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, - 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, - ) - - -# [ModelOpt] Not registered with fla's backend dispatch: another backend must not take over a -# call that asks for W 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``. - w_quantizer (Optional[TensorQuantizer]): - [ModelOpt] Dynamic FP8 QDQ of the WY operand W, with identity STE. - 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.") - - state_quantizer = kwargs.pop("state_quantizer", None) - kwargs.pop("state_qdq_block_v", None) - if kwargs.pop("state_qdq", 0) or (state_quantizer is not None and state_quantizer.is_enabled): - raise ValueError( - "Chunk-only state QAT is retired; use the serving-aligned training workflow" - ) - # w_quantizer: dynamic FP8 TensorQuantizer applied to the WY tensor - # ``w`` of shape [B, T, HV, K] before it multiplies the unquantized state. - 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 w_quantizer is not None: - validate_gdn_quantizer(w_quantizer, name="gdn_w_quantizer") - if cp_context is not None: - raise ValueError("W quantization does not support context parallelism") - - 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, - w_quantizer, - ) - return o, final_state - - -chunk_gdn = chunk_gated_delta_rule diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py index 740c75f3bd3..6ee02050242 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py @@ -19,6 +19,6 @@ import vllm except ImportError as error: raise ImportError( - "decode.precision='vllm_0_15' requires the optional vLLM dependency; " + "precision='vllm_0_15' requires the optional vLLM dependency; " "install the dependencies in the linear-attention QAT example's requirements-vllm.txt." ) from error diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index a961b9d1e09..84d7256c76b 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -260,7 +260,6 @@ def _apply_linear_attention_policy(model, config, saved_policies=None): for name, policy in saved_policies.items(): modules[name].linear_attention_config = LinearAttentionConfig(**policy) for name, module in modules.items(): - module._migrate_legacy_replay_quantizers() module.validate_linear_attention() if getattr(config, "linear_attention", []): print_rank_0( diff --git a/modelopt/torch/quantization/linear_attention/_chunk_prefill.py b/modelopt/torch/quantization/linear_attention/_chunk_prefill.py deleted file mode 100644 index b38aab5e482..00000000000 --- a/modelopt/torch/quantization/linear_attention/_chunk_prefill.py +++ /dev/null @@ -1,126 +0,0 @@ -# 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 _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, - state_qdq: bool = False, - state_qdq_block_v: int = 64, - state_format: str = "fp8_e4m3", - state_quantizer=None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Compute one prepared, nonempty GDN prefix [T,H,D] with a key-first state.""" - scale = q.shape[-1] ** -0.5 if scale is None else scale - state = initial_state - if state_qdq: - state = _state_qdq(state, state_qdq_block_v, state_format, state_quantizer) - pieces = [] - for lo in range(0, len(q), chunk_size): - hi = min(lo + chunk_size, len(q)) - qc, kc, vc = (x[lo:hi].transpose(0, 1) for x in (q, k, v)) - gc = g[lo:hi].transpose(0, 1).cumsum(-1) - bc = beta[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, state_quantizer) - return torch.cat(pieces), state - - -def chunk_kda( - q, - k, - v, - g, - beta, - *, - chunk_size=64, - scale=None, - initial_state, - state_qdq=False, - state_qdq_block_v=64, - state_format="fp8_e4m3", - state_quantizer=None, -): - """Compute one prepared, nonempty KDA prefix [T,H,D] with a key-first state.""" - scale = q.shape[-1] ** -0.5 if scale is None else scale - state = initial_state - if state_qdq: - state = _state_qdq(state, state_qdq_block_v, state_format, state_quantizer) - pieces = [] - for lo in range(0, len(q), chunk_size): - hi = min(lo + chunk_size, len(q)) - qc, kc, vc = (x[lo:hi].transpose(0, 1) for x in (q, k, v)) - prefix = g[lo:hi].transpose(0, 1).cumsum(-2) - bc = beta[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, state_quantizer) - return torch.cat(pieces), state diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index ae86af71400..88134b77779 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -13,185 +13,178 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Saved execution policy for GDN training-time numerical emulation.""" +"""Saved execution policy for GDN/KDA state QAT.""" from typing import Literal -from pydantic import Field, PrivateAttr, model_validator +from pydantic import Field, model_validator from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField -__all__ = [ - "LinearAttentionConfig", - "LinearAttentionDecodeConfig", - "LinearAttentionPolicyEntry", - "LinearAttentionReplayConfig", -] +__all__ = ["LinearAttentionConfig", "LinearAttentionPolicyEntry"] +# Deserialization-only names for pre-unification pickles. ModeloptBaseConfig +# registers each original class name with torch.serialization's safe globals. class _StateConfig(ModeloptBaseConfig): - block_v: Literal[16, 32, 64, 128] = ModeloptField(default=64) - - @model_validator(mode="before") - @classmethod - def _load_legacy_defaults(cls, values): - if isinstance(values, dict): - values = dict(values) - for key, default in (("mode", "chunk"), ("quantize_initial", True)): - if values.get(key) == default: - values.pop(key) - return values + """Legacy state settings; migrated by LinearAttentionConfig on load.""" class _SolveConfig(ModeloptBaseConfig): - # Retained solely to deserialize old pickled policies and validate their fixed setting. - method: Literal["exact"] = ModeloptField(default="exact") + """Legacy solve settings; migrated by LinearAttentionConfig on load.""" class LinearAttentionReplayConfig(ModeloptBaseConfig): - """Anchor refresh and update encoding schedule; quantizers are configured by quant_cfg.""" - - window: int = Field(default=8, ge=1, le=64, strict=True) - encoding: Literal["once", "reencode"] = ModeloptField(default="once") - _legacy_factor_qdq: bool | None = PrivateAttr(default=None) - - @model_validator(mode="wrap") - @classmethod - def _load_legacy_factor_qdq(cls, values, handler): - legacy = None - if isinstance(values, dict) and "factor_qdq" in values: - values = dict(values) - legacy = values.pop("factor_qdq") - if not isinstance(legacy, bool): - raise ValueError("Legacy factor_qdq must be a boolean") - result = handler(values) - if legacy is not None: - result._legacy_factor_qdq = legacy - return result + """Legacy replay settings; migrated by LinearAttentionConfig on load.""" class LinearAttentionDecodeConfig(ModeloptBaseConfig): - """Explicit suffix recurrence; workload supplies per-sequence prefix lengths. + """Legacy decode settings; migrated by LinearAttentionConfig on load.""" - ``precision='vllm_0_15'`` uses the pinned serving forward arithmetic for - the chunked prefix and token suffix, with a differentiable Torch adjoint. - ``replayssm`` imports INT8/Hadamard checkpoint and ring kernels from the - quantized-ReplaySSM serving fork. ``full`` retains the FP32/FP64 Torch reference. - TensorQuantizer configures quantization independently; this profile does not enable QDQ. - """ - mode: Literal["token", "replay"] = ModeloptField(default="token") - readout: Literal["working", "stored"] = ModeloptField(default="stored") - quantize_initial: bool = ModeloptField(default=True) - prefill_state_qdq: bool = ModeloptField(default=False) - precision: Literal["full", "vllm_0_15", "replayssm"] = ModeloptField(default="full") - 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) +def _legacy_fields(value): + if isinstance(value, ModeloptBaseConfig): + # model_dump would omit the fields of these deserialization-only classes. + fields = dict(vars(value)) + factor_qdq = (value.__pydantic_private__ or {}).get("_legacy_factor_qdq") + if factor_qdq is not None: + fields["factor_qdq"] = factor_qdq + return fields + if not isinstance(value, dict): + raise ValueError("Legacy linear-attention settings must be dictionaries") + return dict(value) - @model_validator(mode="before") - @classmethod - def _load_legacy_defaults(cls, values): - if isinstance(values, dict) and values.get("implementation") == "torch": - values = {key: value for key, value in values.items() if key != "implementation"} - return values - @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.state_codec == "int8_hadamard32" and self.prefill_state_qdq: - raise ValueError( - "Hadamard state QDQ starts at decode handoff; disable prefill_state_qdq" - ) - if self.precision == "vllm_0_15" and ( - self.mode != "token" - or self.readout != "working" - or self.state_codec != "tile" - or not self.quantize_initial - or self.prefill_state_qdq - or self.decay_log_step is not None - ): - raise ValueError( - "vllm_0_15 requires token mode, working readout, tile codec, initial-state QDQ, " - "and no additional prefill or gate rounding" - ) - if self.precision == "replayssm" and ( - self.readout != "working" - or self.state_codec != "int8_hadamard32" - or not self.quantize_initial - or self.prefill_state_qdq - or self.decay_log_step is not None - or (self.replay is not None and self.replay.encoding != "once") - ): - raise ValueError( - "replayssm requires working readout, int8_hadamard32, initial-state QDQ, " - "encode-once replay, and no additional prefill or gate rounding" - ) - return self +def _consume_fixed_settings(values, **defaults): + """Remove fixed legacy settings without silently changing their numerics.""" + for key, default in defaults.items(): + if values.pop(key, default) != default: + raise ValueError(f"Native state QAT requires {key}={default!r}") class LinearAttentionConfig(ModeloptBaseConfig): - """Execution policy shared by GDN/KDA state QAT and retained GDN W QAT. + """One execution policy shared by the chunked prefix and recurrent suffix. ``serving`` uses native forward arithmetic with a differentiable adjoint; - ``reference`` retains explicit mathematical experiments. ``fla`` is the - disabled/default or W-only path. State formats and grouping come from - TensorQuantizer; ``state.block_v`` controls legacy tile grouping. + ``fla`` is the disabled/default path. ``vllm_0_15`` uses public vLLM arithmetic; + ``replayssm`` uses the serving fork's INT8/Hadamard checkpoint and ring kernels. + ``replay_window=1`` refreshes the state every token; larger windows enable replay. + Both profiles leave prefill unquantized and quantize the handoff state. + + TensorQuantizer independently enables QDQ and owns its format and grouping. + ``state_block_v`` controls grouping only for legacy per-tile quantizers. """ - schema_version: Literal[1, 2] = ModeloptField(default=2) - backend: Literal["fla", "serving", "reference"] = ModeloptField(default="fla") - chunk_size: Literal[64] = ModeloptField(default=64) - state: _StateConfig = ModeloptField(default=_StateConfig()) - decode: LinearAttentionDecodeConfig | None = ModeloptField(default=None) + schema_version: Literal[3] = ModeloptField(default=3) + backend: Literal["fla", "serving"] = ModeloptField(default="fla") + precision: Literal["vllm_0_15", "replayssm"] = ModeloptField(default="vllm_0_15") + replay_window: int = Field(default=1, ge=1, le=64, strict=True) + state_block_v: Literal[16, 32, 64, 128] = ModeloptField(default=64) + + @property + def state_codec(self): + """The native profile fixes the codec; it is not an independent setting.""" + return "int8_hadamard32" if self.precision == "replayssm" else "tile" @model_validator(mode="before") @classmethod - def _load_legacy_defaults(cls, values): - if isinstance(values, dict) and values.get("backend") == "matmul": - decode = values.get("decode") or {} - precision = ( - decode.get("precision", "full") if isinstance(decode, dict) else decode.precision - ) - values = {**values, "backend": "reference" if precision == "full" else "serving"} - if isinstance(values, dict) and values.get("backend") == "serving": - decode = values.get("decode") - if isinstance(decode, dict): - values = { - **values, - "decode": {"precision": "vllm_0_15", "readout": "working", **decode}, - } - if isinstance(values, dict) and values.get("schema_version") == 1: - decode = values.get("decode") - if isinstance(decode, dict) and isinstance(decode.get("replay"), dict): - replay = {"factor_qdq": True, **decode["replay"]} - values = {**values, "decode": {**decode, "replay": replay}} - if isinstance(values, dict) and "solve" in values: - _SolveConfig.model_validate(values["solve"]) - values = {key: value for key, value in values.items() if key != "solve"} + def _load_legacy_settings(cls, values): + if not isinstance(values, dict): + return values + values = dict(values) + version = values.get("schema_version", 3) + if version in (1, 2): + values["schema_version"] = 3 + if values.get("backend") == "reference": + raise ValueError("Reference training is retired; select backend='serving'") + _consume_fixed_settings(values, chunk_size=64) + if "solve" in values: + solve = _legacy_fields(values.pop("solve")) + _consume_fixed_settings(solve, method="exact") + if solve: + raise ValueError(f"Unsupported legacy solve settings: {solve}") + if "state" in values: + state = _legacy_fields(values.pop("state")) + _consume_fixed_settings(state, mode="chunk", quantize_initial=True) + if "state_block_v" in values: + raise ValueError("Supply state_block_v without legacy state settings") + values["state_block_v"] = state.pop("block_v", 64) + if state: + raise ValueError(f"Unsupported legacy state settings: {state}") + if "decode" in values: + decode = values.pop("decode") + if (values.get("backend", "fla") != "fla") != (decode is not None): + raise ValueError("Legacy state QAT requires an explicit decode policy") + if decode is not None: + if "precision" in values or "replay_window" in values: + raise ValueError( + "Supply precision/replay_window without legacy decode settings" + ) + decode = _legacy_fields(decode) + precision = decode.pop( + "precision", "full" if values.get("backend") == "matmul" else "vllm_0_15" + ) + if precision == "full": + raise ValueError( + "Reference training is retired; select a native serving precision" + ) + values["precision"] = precision + replay = decode.pop("replay", None) + mode = decode.pop("mode", "token") + if mode not in ("token", "replay") or (mode == "replay") != (replay is not None): + raise ValueError("Legacy replay settings require mode='replay'") + _consume_fixed_settings( + decode, + readout="working", + quantize_initial=True, + prefill_state_qdq=False, + decay_log_step=None, + implementation="torch", + state_codec="int8_hadamard32" if precision == "replayssm" else "tile", + ) + if decode: + raise ValueError(f"Unsupported legacy decode settings: {decode}") + if replay is not None: + replay = _legacy_fields(replay) + factor_qdq = replay.pop("factor_qdq", version == 1) + if factor_qdq is not False: + raise ValueError( + "Replay factor QDQ is retired; native ReplaySSM stores BF16 factors" + ) + values["replay_window"] = replay.pop("window", 8) + _consume_fixed_settings(replay, encoding="once") + if replay: + raise ValueError(f"Unsupported legacy replay settings: {replay}") + if precision != "replayssm": + raise ValueError("Replay requires precision='replayssm'") + if values.get("backend") == "matmul": + if "precision" not in values: + raise ValueError("Reference training is retired; select a native serving precision") + values["backend"] = "serving" return values @model_validator(mode="after") - def _validate_decode_backend(self): - if (self.backend != "fla") != (self.decode is not None): - raise ValueError("State QAT requires an explicit decode policy") - if self.decode is not None and (self.backend == "serving") != ( - self.decode.precision != "full" - ): - raise ValueError( - "Serving state QAT requires a native precision profile; full uses reference" - ) - 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") + def _validate_profile(self): + if self.precision == "replayssm": + if self.backend != "serving": + raise ValueError("replayssm requires backend='serving'") + if self.state_block_v < 32: + raise ValueError("int8_hadamard32 requires state_block_v >= 32") + elif self.replay_window != 1: + raise ValueError("Replay requires precision='replayssm'") return self + def __setstate__(self, state): + # Config objects can be pickled inside ModelOpt checkpoints. Normalize + # their old nested fields before any runtime consumer sees the policy. + migrated = type(self).model_validate(state["__dict__"]) + super().__setstate__(migrated.__getstate__()) + 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. + Rules apply in order: the last match wins, without merging fields. A rule must match at least one supported linear-attention module. """ diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py index 74694ada60e..d6fd0702756 100644 --- a/modelopt/torch/quantization/linear_attention/decode.py +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -19,7 +19,7 @@ import torch -from .config import LinearAttentionDecodeConfig +from .config import LinearAttentionConfig from .utils import _resolve_state_quantizer, _tile_qdq, forward_value __all__ = [ @@ -59,7 +59,7 @@ class LinearAttentionCarry: started: bool signature: str value_basis: str = "identity" - precision: str = "full" + precision: str = "vllm_0_15" @property def cursor(self): @@ -68,28 +68,14 @@ def cursor(self): def reconstruct(self, *, original_basis=True): """Replay entries, returning the original value basis unless explicitly disabled.""" - if self.precision == "replayssm": - state = _replay_state(self) - else: - 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) + state = _replay_state(self) if self.precision == "replayssm" else self.anchor.values if original_basis and self.value_basis == "hadamard32": - if self.precision == "replayssm": - # The optional serving fork owns the checkpoint basis transform. - from ...kernels.quantization.linear_attention.serving.replay import ( - original_basis as restore_basis, - ) + # The optional serving fork owns the checkpoint basis transform. + from ...kernels.quantization.linear_attention.serving.replay import ( + original_basis as restore_basis, + ) - state = restore_basis(state) - else: - state = _hadamard32(state) + state = restore_basis(state) return state @@ -156,20 +142,8 @@ def _encode( return EncodedLinearAttentionTensor(decoded, scales, state_format, block_v) -def _encode_factor(value, quantizer): - if quantizer is None: - return EncodedLinearAttentionTensor(value, None, "identity", None) - decoded = quantizer(value) - enabled = quantizer.is_enabled and quantizer._if_quant - format = str(quantizer.num_bits) if enabled else "identity" - return EncodedLinearAttentionTensor(decoded, None, format, None) - - def _signature(config, state_qdq, block_v, state_format, state_quantizer, replay_quantizers): - signature = ( - config.model_dump_json(exclude={"prefill_state_qdq"}) - + f"/{state_qdq}/{block_v}/{state_format}" - ) + signature = config.model_dump_json() + f"/{state_qdq}/{block_v}/{state_format}" if state_quantizer is not None and state_quantizer.block_sizes is not None: signature += f"/group={state_quantizer.block_sizes[-1]}" for quantizer in replay_quantizers: @@ -197,13 +171,6 @@ def _sum_keys(value): return value[..., 0, :] -def _round_log_gate(gate, step): - if step is None: - return gate - rounded = (gate / step).round() * step - return forward_value(gate, rounded) - - def _prepare_carry( q, k, @@ -253,15 +220,13 @@ def _prepare_carry( ) 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: + if 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, + state_qdq, block_v, state=True, state_format=state_format, @@ -295,13 +260,12 @@ def recurrent_decode( g, beta, *, - config: LinearAttentionDecodeConfig, + config: LinearAttentionConfig, state_qdq=False, state_format="fp8_e4m3", state_quantizer=None, replay_key_quantizer=None, replay_update_quantizer=None, - block_v=64, initial_state=None, carry=None, position=0, @@ -314,14 +278,18 @@ def recurrent_decode( 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. - Replay factors use the supplied TensorQuantizers; omitted quantizers leave factors unchanged. + ``config`` is the same LinearAttentionConfig used by the chunked prefix. + Native ReplaySSM stores BF16 factors; legacy factor quantizers must remain disabled. """ state_quantizer, state_qdq, state_format = _resolve_state_quantizer( state_quantizer, state_qdq, state_format ) + if config.backend != "serving": + raise ValueError("State QAT requires backend='serving'") + block_v = config.state_block_v serving = config.precision == "vllm_0_15" native_replay = config.precision == "replayssm" - if (serving or native_replay) and q.device.type != "cuda": + if q.device.type != "cuda": raise ValueError("Serving arithmetic requires CUDA") if ( native_replay @@ -334,19 +302,8 @@ def recurrent_decode( ): raise ValueError("ReplaySSM requires power-of-two K >= 32 and raw GDN gate inputs") replay_quantizers = (replay_key_quantizer, replay_update_quantizer) - if config.replay is not None and config.replay._legacy_factor_qdq: - if any(quantizer is None for quantizer in replay_quantizers): - raise ValueError("Replace factor_qdq with explicit replay key/update TensorQuantizers") - for quantizer in replay_quantizers: - if quantizer is not None and quantizer.is_enabled: - if native_replay: - raise ValueError( - "Native ReplaySSM stores BF16 factors; additional factor QDQ is unsupported" - ) - if config.replay is None or not quantizer.fake_quant: - raise ValueError( - "Replay factor quantizers require replay mode and fake quantization" - ) + if any(quantizer is not None and quantizer.is_enabled for quantizer in replay_quantizers): + raise ValueError("Replay factor QDQ is retired; keep replay key/update quantizers disabled") carry, signature = _prepare_carry( q, k, @@ -377,20 +334,8 @@ def recurrent_decode( 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_factor(e.key.values, replay_key_quantizer), - _encode_factor(e.update.values, replay_update_quantizer), - 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) + gate = g[t] if native_replay: from ...kernels.quantization.linear_attention.serving.replay import step @@ -402,7 +347,7 @@ def recurrent_decode( gate, beta[t], carry, - config.replay.window if config.replay is not None else 1, + config.replay_window, state_qdq, scale, use_qk_l2norm_in_kernel, @@ -430,20 +375,22 @@ def recurrent_decode( current_key = ( current_key / (current_key.square().sum(-1, keepdim=True) + 1e-6).sqrt() ) - key = _encode_factor(current_key, replay_key_quantizer) + key = EncodedLinearAttentionTensor(current_key, None, "identity", None) decayed = state * decay if native_replay and gate.ndim == 2: decayed = _replay_state(carry, gate) residual = v[t] - _sum_keys(key.values.unsqueeze(-1) * decayed) - update = _encode_factor(beta[t].unsqueeze(-1) * residual, replay_update_quantizer) + update = EncodedLinearAttentionTensor( + beta[t].unsqueeze(-1) * residual, None, "identity", None + ) working = decayed + key.values.unsqueeze(-1) * update.values.unsqueeze(-2) - if config.replay is not None: - if native_replay and carry.cursor + 1 < config.replay.window: + if config.replay_window > 1: + if native_replay and carry.cursor + 1 < config.replay_window: key.values = forward_value(key.values, native[3]) update.values = forward_value(update.values, native[4]) gate = forward_value(gate, native[5]) entries = (*entries, ReplayEntry(key, update, gate)) - refresh = config.replay is None or len(entries) == config.replay.window + refresh = config.replay_window == 1 or len(entries) == config.replay_window if refresh: anchor = _encode( working, @@ -458,10 +405,8 @@ def recurrent_decode( anchor.values = forward_value(working, native[1]) anchor.scales = native[2] entries = () - stored = anchor.values else: anchor = carry.anchor - stored = working next_carry = LinearAttentionCarry( anchor, entries, @@ -471,12 +416,11 @@ def recurrent_decode( carry.value_basis, config.precision, ) - read = working if config.readout == "working" else stored query = q[t] if native_replay and use_qk_l2norm_in_kernel: query = query / (query.square().sum(-1, keepdim=True) + 1e-6).sqrt() outputs.append( - native_output if serving else _sum_keys(query.unsqueeze(-1) * read) * scale + native_output if serving else _sum_keys(query.unsqueeze(-1) * working) * scale ) carry = next_carry output = torch.stack(outputs) diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py index 54b7b8efce2..d9d39a3d8b9 100644 --- a/modelopt/torch/quantization/linear_attention/kda.py +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -60,7 +60,7 @@ def matmul_kda( """Normalize KDA inputs and run a chunked prefix plus configured suffix recurrence. Gates follow FLA's kernel activation formula and per-key log retention. - The optional serving precision profile supplies matching native forward values + The serving precision profile supplies matching native forward values and a Torch adjoint through the rounded state trajectory. """ if cp_context is not None or disable_recompute or return_intermediate_states: @@ -69,14 +69,12 @@ def matmul_kda( ) 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") + if lower_bound is not None or safe_gate: + raise ValueError("Serving arithmetic uses the native softplus KDA gate") output_dtype = q.dtype beta_dtype = beta.dtype - q, k, v, g, beta, serving = _prepare_prefill_inputs( - q, k, v, g, beta, policy=policy, chunk_size=chunk_size, normalize=use_qk_l2norm_in_kernel + q, k, v, g, beta = _prepare_prefill_inputs( + q, k, v, g, beta, policy=policy, chunk_size=chunk_size ) dtype = q.dtype if g.ndim != 4: @@ -88,18 +86,15 @@ def matmul_kda( 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 serving: - if lower_bound is not None or safe_gate: - raise ValueError("vllm_0_15 uses the native softplus KDA gate") - # Import the optional vLLM backend only for the native precision profile. - from ...kernels.quantization.linear_attention.serving.forward import fused_kda_gate + g = -rate * F.softplus(g) + # Import the optional vLLM backend only for the native precision profile. + from ...kernels.quantization.linear_attention.serving.forward import fused_kda_gate - with torch.no_grad(): - native_gate = fused_kda_gate( - raw_gate.flatten(-2).contiguous(), A_log, raw_gate.shape[-1], g_bias=dt_bias - ) - g = forward_value(g, native_gate) + with torch.no_grad(): + native_gate = fused_kda_gate( + raw_gate.flatten(-2).contiguous(), A_log, raw_gate.shape[-1], g_bias=dt_bias + ) + g = forward_value(g, native_gate) if use_beta_sigmoid_in_kernel: beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) return _prefill_decode_forward( @@ -123,5 +118,5 @@ def matmul_kda( output_dtype=output_dtype, beta_dtype=beta_dtype, prefill_lengths=prefill_lengths, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel if serving else False, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ) diff --git a/modelopt/torch/quantization/linear_attention/prefill.py b/modelopt/torch/quantization/linear_attention/prefill.py index 7bbb42deb09..18573af06a8 100644 --- a/modelopt/torch/quantization/linear_attention/prefill.py +++ b/modelopt/torch/quantization/linear_attention/prefill.py @@ -15,8 +15,6 @@ """Batched differentiable GDN prefill with materialized numerical boundaries.""" -import torch.nn.functional as F - from .config import LinearAttentionConfig from .training import _prefill_decode_forward, _prepare_prefill_inputs @@ -55,28 +53,21 @@ def matmul_gdn( ): """Normalize GDN inputs and run a chunked prefix plus configured suffix recurrence. - Uses FP32 working arithmetic (FP64 for double inputs), with floating QDQ state. - The optional serving precision profile supplies matching native forward values + Uses FP32 working arithmetic with floating QDQ state. + The serving precision profile supplies matching native forward values and a Torch adjoint through the rounded state trajectory. """ if cp_context is not None: raise NotImplementedError("GDN matmul emulation does not support context parallelism") output_dtype = q.dtype beta_dtype = beta.dtype - q, k, v, g, beta, serving = _prepare_prefill_inputs( - q, k, v, g, beta, policy=policy, chunk_size=chunk_size, normalize=use_qk_l2norm_in_kernel + q, k, v, g, beta = _prepare_prefill_inputs( + q, k, v, g, beta, policy=policy, chunk_size=chunk_size ) - dtype = q.dtype - if serving and (use_gate_in_kernel or use_beta_sigmoid_in_kernel): + if use_gate_in_kernel or use_beta_sigmoid_in_kernel: raise ValueError( "Serving GDN expects prepared log gates and beta from the Megatron adapter" ) - 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") return _prefill_decode_forward( @@ -101,5 +92,5 @@ def matmul_gdn( beta_dtype=beta_dtype, prefill_lengths=prefill_lengths, replay_gate_inputs=replay_gate_inputs, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel if serving else False, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ) diff --git a/modelopt/torch/quantization/linear_attention/training.py b/modelopt/torch/quantization/linear_attention/training.py index fa9987e6296..96c74f41fa5 100644 --- a/modelopt/torch/quantization/linear_attention/training.py +++ b/modelopt/torch/quantization/linear_attention/training.py @@ -20,7 +20,6 @@ import torch -from ._chunk_prefill import chunk_gdn, chunk_kda from .decode import recurrent_decode from .utils import _resolve_state_quantizer, _state_qdq, forward_value @@ -51,7 +50,7 @@ def linear_attention_training_phase(model, prefill_lengths): m for m in model.modules() if isinstance(m, _LinearAttentionQuantMixin) - and m.linear_attention_config.decode is not None + and m.linear_attention_config.backend == "serving" ] if not layers: raise ValueError("The model has no converted decode-aware linear-attention layers") @@ -65,22 +64,13 @@ def linear_attention_training_phase(model, prefill_lengths): module._linear_attention_prefill_lengths = original -def _prepare_prefill_inputs(q, k, v, g, beta, *, policy, chunk_size, normalize): - """Validate the shared policy and prepare GDN/KDA working dtypes and Q/K normalization.""" - if policy.backend not in ("serving", "reference") or chunk_size != policy.chunk_size: - raise ValueError( - "State training requires backend='serving' or 'reference' and its configured chunk size" - ) - if policy.decode is None: - raise ValueError("An explicit decode policy is required") - serving = policy.decode.precision != "full" - if serving and (q.device.type != "cuda" or any(x.dtype != torch.bfloat16 for x in (q, k, v))): +def _prepare_prefill_inputs(q, k, v, g, beta, *, policy, chunk_size): + """Validate the native policy and prepare FP32 values for the training adjoint.""" + if policy.backend != "serving" or chunk_size != 64: + raise ValueError("State training requires backend='serving' and chunk_size=64") + if q.device.type != "cuda" or any(x.dtype != torch.bfloat16 for x in (q, k, v)): raise ValueError("Serving arithmetic requires CUDA BF16 Q/K/V inputs") - dtype = torch.float64 if q.dtype == torch.float64 else torch.float32 - q, k, v, g, beta = (x.to(dtype) for x in (q, k, v, g, beta)) - if normalize and not serving: - q, k = (x * (x.square().sum(-1, keepdim=True) + 1e-6).rsqrt() for x in (q, k)) - return q, k, v, g, beta, serving + return tuple(x.float() for x in (q, k, v, g, beta)) def _prefill_decode_forward( @@ -152,8 +142,7 @@ def _prefill_decode_forward( p > end - start for p, (_, start, end) in zip(prefixes, sequences) ): raise ValueError("Supply one valid prefill length per sequence") - serving_precision = policy.decode.precision != "full" - if serving_precision and ( + if ( keys > 256 or (g.ndim == 4 and values != keys) or (use_qk_l2norm_in_kernel and keys & (keys - 1)) @@ -168,66 +157,52 @@ def _prefill_decode_forward( states = states.to(q.dtype) if states.shape != (len(sequences), heads, keys, values): raise ValueError("Initial state shape does not match sequence/head dimensions") - prefix_fn = chunk_kda if g.ndim == 4 else chunk_gdn outputs, finals = [], [] for n, (b, start, end) in enumerate(sequences): split = start + prefixes[n] prefix, state = q.new_empty(0, heads, values), states[n] if prefixes[n]: with torch.autocast(device_type=q.device.type, enabled=False): - if serving_precision: - from ._vllm_autograd import prefix as serving_prefix - - # A continuation prefill consumes a stored cache just as native serving - # does. A fresh zero-state prefix has no incoming cache to quantize. - if initial_state is not None and state_qdq: - if policy.decode.precision == "replayssm": - from ...kernels.quantization.linear_attention.serving.replay import ( - checkpoint, - original_basis, - ) - - with torch.no_grad(): - decoded, _ = checkpoint(state, True) - decoded = original_basis(decoded) - state = forward_value(state, decoded) - else: - state = _state_qdq( - state, policy.state.block_v, state_format, state_quantizer - ) - prefix, state = serving_prefix( - q=q[b, start:split], - k=k[b, start:split], - v=v[b, start:split], - g=g[b, start:split], - beta=beta[b, start:split], - state=state, - scale=keys**-0.5 if scale is None else scale, - beta_dtype=beta_dtype, - normalize=use_qk_l2norm_in_kernel, - ) - else: - prefix, state = prefix_fn( - *(x[b, start:split] for x in (q, k, v, g, beta)), - state_qdq=state_qdq and policy.decode.prefill_state_qdq, - state_format=state_format, - state_quantizer=state_quantizer, - scale=scale, - initial_state=state, - chunk_size=policy.chunk_size, - state_qdq_block_v=policy.state.block_v, - ) + from ._vllm_autograd import prefix as serving_prefix + + # A continuation prefill consumes a stored cache just as native serving + # does. A fresh zero-state prefix has no incoming cache to quantize. + if initial_state is not None and state_qdq: + if policy.precision == "replayssm": + from ...kernels.quantization.linear_attention.serving.replay import ( + checkpoint, + original_basis, + ) + + with torch.no_grad(): + decoded, _ = checkpoint(state, True) + decoded = original_basis(decoded) + state = forward_value(state, decoded) + else: + state = _state_qdq( + state, policy.state_block_v, state_format, state_quantizer + ) + prefix, state = serving_prefix( + q=q[b, start:split], + k=k[b, start:split], + v=v[b, start:split], + g=g[b, start:split], + beta=beta[b, start:split], + state=state, + scale=keys**-0.5 if scale is None else scale, + beta_dtype=beta_dtype, + normalize=use_qk_l2norm_in_kernel, + ) # Keep the prefix state attached so suffix losses backpropagate through prefill. # recurrent_decode applies configured state QDQ at the handoff and suffix writes. suffix, carry = recurrent_decode( *(x[b, split:end] for x in (q, k, v, g, beta)), - config=policy.decode, + config=policy, replay_key_quantizer=replay_key_quantizer, replay_update_quantizer=replay_update_quantizer, state_qdq=state_qdq, state_format=state_format, state_quantizer=state_quantizer, - block_v=policy.state.block_v, initial_state=state, position=prefixes[n], replay_gate_inputs=( diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py index bed74203474..afb40b73c5c 100644 --- a/modelopt/torch/quantization/linear_attention/utils.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -97,7 +97,7 @@ def state_quantizer_config( if quantizer.block_sizes is not None: return _STATE_FORMATS[quantizer.num_bits], quantizer.block_sizes[-1] if quantizer.axis != (0, 1): - raise ValueError(f"{name} supports only axis=(0, 1) with state.block_v tiling") + raise ValueError(f"{name} supports only axis=(0, 1) with state_block_v tiling") return _STATE_FORMATS[quantizer.num_bits], 0 diff --git a/modelopt/torch/quantization/plugins/gdn.py b/modelopt/torch/quantization/plugins/gdn.py index 93c561fc2fd..87df30ca7e0 100644 --- a/modelopt/torch/quantization/plugins/gdn.py +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -13,11 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""GatedDeltaNet state and WY activation fake quantization. +"""GatedDeltaNet recurrent-state fake quantization. State QAT uses an explicit prefix/suffix policy with native serving arithmetic. -W-only QAT applies TensorQuantizer to the WY-transformed keys and reuses FLA's -state kernels. Both paths retain gradients through quantization with identity STE. """ from collections.abc import Callable @@ -27,7 +25,6 @@ import torch from ..linear_attention.prefill import matmul_gdn -from ..linear_attention.utils import validate_gdn_quantizer from .linear_attention import _LinearAttentionQuantMixin __all__ = ["GatedDeltaNetStateQuantMixin"] @@ -42,61 +39,44 @@ def _fla_chunk_gated_delta_rule() -> GatedDeltaRuleFn: return chunk_gated_delta_rule -def _w_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 W-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. + """Adds recurrent-state fake quantization 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 QDQ configuration. State supports dynamic E4M3 or signed - narrow-range INT8; W supports dynamic E4M3. - Both sites require identity STE; state QAT uses an explicit serving policy. - The execution policy is saved in ModelOpt metadata. + :meth:`_state_quantized_chunk_gated_delta_rule`. Enable ``*gdn_state_quantizer`` + through ``quant_cfg`` and select an explicit serving policy. State supports dynamic + E4M3 or signed narrow-range INT8 with identity STE. The execution policy is saved in + ModelOpt metadata. ``gdn_w_quantizer`` is a disabled legacy checkpoint handle; + enabling WY operand quantization is no longer supported. """ linear_attention_quantizer_names = ("gdn_state_quantizer", "gdn_w_quantizer") def validate_linear_attention(self): - """Validate state and GDN WY operand quantization.""" - super().validate_linear_attention() + """Reject retired W QAT and validate the state quantizer.""" if self.gdn_w_quantizer.is_enabled: - validate_gdn_quantizer(self.gdn_w_quantizer, name="gdn_w_quantizer") - if self.linear_attention_config.decode is not None: - raise ValueError("Decode's exact prefix does not support WY operand QDQ") + raise ValueError( + "GDN W quantization is no longer supported. Disable gdn_w_quantizer; " + "use gdn_state_quantizer with an explicit serving policy for state QAT." + ) + super().validate_linear_attention() @property def gdn_state_qdq_block_v(self) -> int: """Execution tile width; also sets grouping for legacy tile quantizers.""" - return self.linear_attention_config.state.block_v + 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]: - """Route enabled state or W quantization to its training implementation.""" + """Route enabled state quantization to its training implementation.""" self.validate_linear_attention() - quantize_w = self.gdn_w_quantizer.is_enabled if not self.linear_attention_is_enabled: return gated_delta_rule(*args, **kwargs) if self.linear_attention_config.backend == "fla" and self.gdn_state_quantizer.is_enabled: raise ValueError( - "Chunk-only state QAT is retired. Select backend='serving' with an explicit " - "decode policy and supply prefill lengths through linear_attention_training_phase. " + "Chunk-only state QAT is retired. Select backend='serving' and supply prefill " + "lengths through linear_attention_training_phase. " "Existing state checkpoints are not silently migrated to a different QDQ schedule." ) while isinstance(gated_delta_rule, partial): @@ -108,23 +88,16 @@ def _state_quantized_chunk_gated_delta_rule( "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: + chunk_size = kwargs.pop("chunk_size", 64) + if chunk_size != 64: raise ValueError("GDN fake quantization supports only chunk_size=64") - if self.linear_attention_config.backend != "fla": - return matmul_gdn( - *args, - policy=self.linear_attention_config, - state_quantizer=self.gdn_state_quantizer, - replay_key_quantizer=self.replay_key_quantizer, - replay_update_quantizer=self.replay_update_quantizer, - chunk_size=chunk_size, - prefill_lengths=self._linear_attention_prefill_lengths, - **kwargs, - ) - return _w_qdq_chunk_gated_delta_rule()( + return matmul_gdn( *args, + policy=self.linear_attention_config, + state_quantizer=self.gdn_state_quantizer, + replay_key_quantizer=self.replay_key_quantizer, + replay_update_quantizer=self.replay_update_quantizer, chunk_size=chunk_size, - w_quantizer=self.gdn_w_quantizer if quantize_w else None, + prefill_lengths=self._linear_attention_prefill_lengths, **kwargs, ) diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py index 22f75cfe142..a445f1e1c91 100644 --- a/modelopt/torch/quantization/plugins/linear_attention.py +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -47,50 +47,23 @@ def linear_attention_is_enabled(self): getattr(self, name).is_enabled for name in (*self.linear_attention_quantizer_names, *self.replay_quantizer_names) ) - or self.linear_attention_config.decode is not None + or self.linear_attention_config.backend == "serving" ) - def _migrate_legacy_replay_quantizers(self): - decode = self.linear_attention_config.decode - replay = decode.replay if decode is not None else None - legacy = getattr(replay, "_legacy_factor_qdq", None) - if ( - replay is not None - and legacy is None - and self.linear_attention_config.schema_version == 1 - ): - legacy = replay.__dict__.pop("factor_qdq", True) - if replay is not None and legacy is not None: - for name in self.replay_quantizer_names: - grouping = ( - {"axis": (0,)} - if name == "replay_key_quantizer" - else {"block_sizes": {-1: self.linear_attention_config.state.block_v}} - ) - getattr(self, name).set_from_attribute_config( - QuantizerAttributeConfig( - num_bits=(4, 3), type="dynamic", enable=legacy, **grouping - ) - ) - replay._legacy_factor_qdq = None - self.linear_attention_config.schema_version = 2 - def validate_linear_attention(self): """Validate quantizer contracts shared by GDN and KDA.""" - decode = self.linear_attention_config.decode for name in self.replay_quantizer_names: quantizer = getattr(self, name) if quantizer.is_enabled: - if decode is None or decode.mode != "replay": - raise ValueError(f"{name} requires a replay execution policy") - if not quantizer.fake_quant: - raise ValueError(f"{name} requires fake quantization for training") + raise ValueError( + f"{name} is no longer supported; native ReplaySSM stores BF16 factors" + ) if self._linear_attn_state.is_enabled: state_format, group_size = state_quantizer_config( self._linear_attn_state, name=self.linear_attention_quantizer_names[0], ) - if decode is not None and decode.state_codec == "int8_hadamard32": + if self.linear_attention_config.state_codec == "int8_hadamard32": if state_format != "int8": raise ValueError("int8_hadamard32 requires INT8 state quantization") if group_size: diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 1b00e52beca..5283c193640 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -1134,10 +1134,7 @@ def _setup(self): @property def _serving_arithmetic(self): - decode = self.linear_attention_config.decode - return ( - self.linear_attention_is_enabled and decode is not None and decode.precision != "full" - ) + return self.linear_attention_config.backend == "serving" def _prepare_input_for_gated_delta_rule(self, *args, **kwargs): # Preserve raw BF16 Q/K: prefill stores normalized BF16 operands, whereas @@ -1159,7 +1156,7 @@ def _compute_gates(self, a_log, dt_bias, batch, seq_len, *gate_feats): from ..linear_attention.utils import forward_value raw_beta, raw_gate = gate_feats - if self.linear_attention_config.decode.precision == "replayssm": + if self.linear_attention_config.precision == "replayssm": self._linear_attention_replay_gate_inputs = (raw_gate, raw_beta, a_log, dt_bias) with torch.no_grad(): native_gate, native_beta = fused_gdn_gating( diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index f40d76da3a3..1e685a93f9d 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -36,28 +36,10 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `gdn_state_fp8_dynamic.yaml` | FP8 E4M3 dynamic fake quantization of the GatedDeltaNet recurrent state (per sequence, head and 64-column tile) at serving cache boundaries; requires an explicit serving execution policy | | `linear_attention_state_int8_dynamic.yaml` | GDN/KDA INT8 state quantizer entries; the [complete recipe](../../../general/ptq/linear_attention_state_int8_dynamic.yaml) enables Hadamard during decode | | `linear_attention_state_int8_block32_dynamic.yaml` | Standard dynamic INT8 with one scale per key row and 32 value channels; the [complete training recipe](../../../general/ptq/linear_attention_state_int8_block32_dynamic.yaml) uses token QDQ and working-state readout | -| `linear_attention_replay_fp8_dynamic.yaml` | GDN/KDA ReplaySSM key/update FP8 quantizers; for the explicit reference backend; native ReplaySSM uses BF16 ring entries | -| `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 | | `indexer_k_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer key cache (`*indexer_k_quantizer`), global scale fixed to 1; Blackwell+ GPUs | | `indexer_q_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer query (`*indexer_q_quantizer`), global scale fixed to 1; Blackwell+ GPUs | | `moe_dispatch_nvfp4.yaml` | NVFP4 fake quantization of the MoE expert-parallel dispatch (`*dispatch_quantizer`, vLLM); Blackwell+ GPUs | | `moe_combine_nvfp4.yaml` | NVFP4 fake quantization of the MoE expert-parallel combine (`*combine_quantizer`, vLLM); Blackwell+ GPUs | -Replay factor formats, grouping, and enable/disable controls use `quant_cfg` through -`*replay_key_quantizer` and `*replay_update_quantizer`. The replay policy controls -scheduling only; for example: - -```yaml -linear_attention: - - module_name: "*" - cfg: - backend: reference - decode: - mode: replay - replay: {window: 8, encoding: once} -``` - -Both factor quantizers start disabled. Import `linear_attention_replay_fp8_dynamic.yaml` -to enable their dynamic FP8 recipe. Previous checkpoints using `replay.factor_qdq` -are migrated to these quantizers when restored. State quantizer recipes are independent; -this unit does not enable state quantization. +Native ReplaySSM uses BF16 key/update vectors. Legacy `replay_key_quantizer` and +`replay_update_quantizer` handles must remain disabled; factor QDQ is unsupported. diff --git a/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml deleted file mode 100644 index 2c73fdd6eaa..00000000000 --- a/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml +++ /dev/null @@ -1,28 +0,0 @@ -# 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/configs/ptq/units/linear_attention_replay_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/linear_attention_replay_fp8_dynamic.yaml deleted file mode 100644 index eaa81f957c6..00000000000 --- a/modelopt_recipes/configs/ptq/units/linear_attention_replay_fp8_dynamic.yaml +++ /dev/null @@ -1,28 +0,0 @@ -# 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. - -# Replay factors for GDN/KDA; requires a linear_attention decode policy with mode: replay. -# Keys use one scale per head; updates use one scale per head and 64 value channels. -# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig - - quantizer_name: '*replay_key_quantizer' - cfg: - num_bits: e4m3 - axis: [0] - type: dynamic - - quantizer_name: '*replay_update_quantizer' - cfg: - num_bits: e4m3 - block_sizes: {-1: 64} - type: dynamic diff --git a/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml index 272d3ce2825..8bff5421d9e 100644 --- a/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml +++ b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml @@ -33,10 +33,4 @@ quantize: - module_name: '*' cfg: backend: serving - decode: - mode: token - state_codec: tile - readout: working - quantize_initial: true - prefill_state_qdq: false - precision: vllm_0_15 + precision: vllm_0_15 diff --git a/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml b/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml index 0c07838b8ad..3c38ccfe568 100644 --- a/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml +++ b/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml @@ -33,9 +33,4 @@ quantize: - module_name: '*' cfg: backend: serving - decode: - mode: token - precision: replayssm - readout: working - state_codec: int8_hadamard32 - prefill_state_qdq: false + precision: replayssm diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 882d16cb15f..e2aa898890f 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -69,7 +69,7 @@ supported combinations. -The block32 linear-attention recipe selects `decode.precision="vllm_0_15"` +The block32 linear-attention recipe selects `precision="vllm_0_15"` and working-state readout. It applies fake QDQ with floating-point state storage. --- diff --git a/pyproject.toml b/pyproject.toml index cc95e6bca7c..b792592022b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -329,15 +329,6 @@ 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 index 3acf3fc0e5c..1373b5ea46b 100644 --- a/tests/_test_utils/torch/quantization/linear_attention_reference.py +++ b/tests/_test_utils/torch/quantization/linear_attention_reference.py @@ -20,13 +20,12 @@ 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_delta_rule_reference", "recurrent_delta_rule_reference", "state_qdq_reference", ] @@ -139,7 +138,7 @@ def recurrent_delta_rule_reference( return output, final.transpose(-1, -2) if state_v_first else final -def chunk_gdn_reference( +def chunk_gdn( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -148,80 +147,114 @@ def chunk_gdn_reference( *, 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, + initial_state: torch.Tensor, state_qdq: bool = False, state_qdq_block_v: int = 64, state_format: str = "fp8_e4m3", - w_quantizer: Callable[[torch.Tensor], torch.Tensor] | None = None, - state_quantizer=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. - """ - if state_quantizer is not None: - state_qdq = state_quantizer.is_enabled and state_quantizer._if_quant - state_format = "int8" if state_quantizer.num_bits == 8 else "fp8_e4m3" - - def quantize_state(state): - if state_quantizer is not None and state_quantizer.block_sizes is not None: - return state_quantizer(state) - return state_qdq_reference(state, state_qdq_block_v, state_format) - - 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) + """Compute one prepared, nonempty GDN prefix [T,H,D] with a key-first state.""" scale = q.shape[-1] ** -0.5 if scale is None else scale - 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 = 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() - gram = kc @ kc.transpose(-1, -2) - lower = (bc * gram * decay).tril(-1) - matrix = lower + torch.eye(hi - lo, device=q.device, dtype=q.dtype) - rhs = torch.cat((bc * vc, bc * kc * gc.exp().unsqueeze(-1)), 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) - 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 = quantize_state(state) - previous_n = n - length = qc.shape[1] - wc = w[offset : offset + length].transpose(0, 1) - offset += length - updated_values = u - wc @ state # state_read - local_scores = ((qc * scale) @ kc.transpose(-1, -2) * decay).tril() + state = initial_state + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + pieces = [] + for lo in range(0, len(q), chunk_size): + hi = min(lo + chunk_size, len(q)) + qc, kc, vc = (x[lo:hi].transpose(0, 1) for x in (q, k, v)) + gc = g[lo:hi].transpose(0, 1).cumsum(-1) + bc = beta[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 - output = output + local_scores @ updated_values # local_readout - outputs.append(output.transpose(0, 1)) + 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 # state_update + state = state + weighted_keys.transpose(-1, -2) @ updated_values if state_qdq: - state = quantize_state(state) - if len(finals) <= n: - finals.append(state) - else: - finals[n] = state - output = torch.cat(outputs).reshape(*v.shape) + state = state_qdq_reference(state, state_qdq_block_v, state_format) + return torch.cat(pieces), state + + +def chunk_kda( + q, + k, + v, + g, + beta, + *, + chunk_size=64, + scale=None, + initial_state, + state_qdq=False, + state_qdq_block_v=64, + state_format="fp8_e4m3", +): + """Compute one prepared, nonempty KDA prefix [T,H,D] with a key-first state.""" + scale = q.shape[-1] ** -0.5 if scale is None else scale + state = initial_state + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + pieces = [] + for lo in range(0, len(q), chunk_size): + hi = min(lo + chunk_size, len(q)) + qc, kc, vc = (x[lo:hi].transpose(0, 1) for x in (q, k, v)) + prefix = g[lo:hi].transpose(0, 1).cumsum(-2) + bc = beta[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_reference(state, state_qdq_block_v, state_format) + return torch.cat(pieces), state + + +def chunk_delta_rule_reference( + q, k, v, g, beta, *, initial_state=None, cu_seqlens=None, state_v_first=False, chunk_size=64 +): + """Evaluate GDN/KDA chunk algebra independently for each packed or batched sequence.""" + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + chunk = chunk_kda if g.ndim == 4 else chunk_gdn + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + output, state = chunk( + *(x[b, start:end] for x in (q, k, v, g, beta)), + initial_state=states[n], + chunk_size=chunk_size, + ) + outputs.append(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 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 deleted file mode 100644 index a3fa7585f56..00000000000 --- a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py +++ /dev/null @@ -1,93 +0,0 @@ -# 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. - -"""Minimal forward/backward checks for GDN training fake quantization.""" - -import pytest -import torch -import torch.nn.functional as F -from _test_utils.torch.quantization.linear_attention_reference import chunk_gdn_reference - -from modelopt.torch.quantization.config import QuantizerAttributeConfig -from modelopt.torch.quantization.nn import TensorQuantizer - -pytest.importorskip("fla.ops.gated_delta_rule") - -from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( - chunk_gated_delta_rule, -) - - -def make_inputs(): - torch.manual_seed(123) - # Two chunks exercise recurrence with one shared BF16 shape for every mode. - shape = (1, 128, 1, 32) - q, k = [F.normalize(torch.randn(shape, device="cuda"), dim=-1) for _ in range(2)] - v = torch.randn(shape, device="cuda") - g = -torch.rand(shape[:3], device="cuda") * 0.1 - beta = torch.rand_like(g) - args = [x.to(torch.bfloat16).requires_grad_() for x in (q, k, v, g, beta)] - state = (torch.randn(1, 1, 32, 32, device="cuda") * 0.1).requires_grad_() - return args, state - - -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 - - -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}" - ) - - -@pytest.fixture(scope="module", params=[False, True]) -def compiled_gdn_case(request): - """Compile the retained W-QAT path outside the test-call timer.""" - quantizer = ( - TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), axis=(0, 1, 2), type="dynamic")) - if request.param - else None - ) - args, state = make_inputs() - kwargs = {"w_quantizer": quantizer} - values_and_grads(chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs) - torch.cuda.synchronize() - return args, state, kwargs - - -def test_gdn_forward_and_backward(compiled_gdn_case): - args, state, kwargs = compiled_gdn_case - reference_args = [x.detach().float().requires_grad_() for x in args] - reference_state = state.detach().clone().requires_grad_() - expected = values_and_grads( - chunk_gdn_reference, - reference_args, - reference_state, - **kwargs, - ) - actual = values_and_grads( - chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs - ) - compare(actual[0], expected[0], 0.03) - compare(actual[1], expected[1], 0.05) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index b9e97aa85ba..72e0229db24 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -31,17 +31,12 @@ ) from modelopt.torch.quantization.nn import TensorQuantizer -pytest.importorskip("fla") # Megatron-Core GatedDeltaNet and the state QDQ kernel need fla +pytest.importorskip("fla") # Megatron-Core GatedDeltaNet needs FLA for its baseline kernels. GatedDeltaNet = pytest.importorskip("megatron.core.ssm.gated_delta_net").GatedDeltaNet +pytest.importorskip("vllm.model_executor.layers.fla.ops.kda", exc_type=ModuleNotFoundError) -from modelopt.torch.quantization.plugins.gdn import _w_qdq_chunk_gated_delta_rule from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet -try: - _w_qdq_chunk_gated_delta_rule() -except RuntimeError as e: - pytest.skip(str(e), allow_module_level=True) - SEED = 1234 @@ -63,7 +58,7 @@ def _make_model(tp_size): .cuda() .eval() ) - # Retain enough history for chunk-boundary state rounding to affect the next chunk. + # Retain enough history for recurrent-state rounding to affect later tokens. with torch.no_grad(): for module in model.modules(): if isinstance(module, GatedDeltaNet): @@ -72,46 +67,41 @@ def _make_model(tp_size): return model -def _gdn_config(sites): +def _gdn_config(): return { - "quant_cfg": [{"quantizer_name": "*", "enable": False}] - + [ + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, { - "quantizer_name": f"*gdn_{site}_quantizer", + "quantizer_name": "*gdn_state_quantizer", "cfg": { - "num_bits": (4, 3), - "axis": (0, 1) if site == "state" else (0, 1, 2), + "num_bits": 8, + "block_sizes": {-1: 32}, "type": "dynamic", + "unsigned": False, + "narrow_range": True, }, - } - for site in sites + }, ], "algorithm": None, "linear_attention": [ - {"module_name": "decoder.layers.*.self_attention", "cfg": {"state": {"block_v": 16}}} + { + "module_name": "decoder.layers.*.self_attention", + "cfg": { + "backend": "serving", + "state_block_v": 16, + "precision": "vllm_0_15", + }, + } ], } -def _config_for_mode(mode): - sites = ("w",) - if mode == "decode_replay_int8": - sites = ("state",) - cfg = _gdn_config(sites) - if mode == "decode_replay_int8": - cfg["quant_cfg"][1]["cfg"].update(num_bits=8, unsigned=False, narrow_range=True) - cfg["linear_attention"][0]["cfg"].update( - backend="matmul", decode={"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 + getattr(getattr(layer, "linear_attention_config", None), "backend", None) == "serving" for layer in m.modules() ) with linear_attention_training_phase(m, [64, 64]) if enabled else nullcontext(): @@ -120,7 +110,7 @@ def forward(m): return forward -def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): +def _test_gdn_qat_helper(rank, size, checkpoint_path): initialize_for_megatron( tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED ) @@ -140,17 +130,16 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): gdn_ref = outputs.copy() outputs.clear() - cfg, sites = _config_for_mode(mode) - mtq.quantize(model, cfg) + mtq.quantize(model, _gdn_config()) 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) + assert module.gdn_state_quantizer.is_enabled + assert not module.gdn_w_quantizer.is_enabled # 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}} + **{**module.linear_attention_config.model_dump(), "state_block_v": 32} ) with torch.no_grad(): @@ -188,8 +177,8 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): 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.gdn_state_quantizer.is_enabled + assert not module.gdn_w_quantizer.is_enabled assert module.gdn_state_quantizer.num_bits == original.gdn_state_quantizer.num_bits assert module.gdn_state_qdq_block_v == 32 assert module.in_proj.weight.grad is not None @@ -202,7 +191,7 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): assert not torch.equal(restored_gdn[0].in_proj.weight, before) -def _compile_gdn_qat_kernels(rank, size, mode): +def _compile_gdn_qat_kernels(rank, size): initialize_for_megatron( tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED ) @@ -210,9 +199,9 @@ def _compile_gdn_qat_kernels(rank, size, mode): forward = _gdn_forward(model) with torch.no_grad(): forward(model) - cfg, _ = _config_for_mode(mode) + cfg = _gdn_config() # The functional case switches state tiles before running inference. - cfg["linear_attention"][0]["cfg"]["state"]["block_v"] = 32 + cfg["linear_attention"][0]["cfg"]["state_block_v"] = 32 mtq.quantize(model, cfg) with torch.no_grad(): forward(model) @@ -221,14 +210,13 @@ def _compile_gdn_qat_kernels(rank, size, mode): torch.cuda.synchronize() -@pytest.fixture(params=["chunk", "decode_replay_int8"]) -def compiled_gdn_workers(request, dist_workers_size_1): +@pytest.fixture +def compiled_gdn_workers(dist_workers_size_1): """Warm one small QAT model in the same worker, outside the test-call budget.""" - dist_workers_size_1.run(_compile_gdn_qat_kernels, request.param) - return dist_workers_size_1, request.param + dist_workers_size_1.run(_compile_gdn_qat_kernels) + return dist_workers_size_1 def test_gdn_qat_and_sharded_restore(compiled_gdn_workers, tmp_path): """Train through QDQ after a Megatron distributed-checkpoint round trip.""" - workers, mode = compiled_gdn_workers - workers.run(_test_gdn_qat_helper, mode, tmp_path) + compiled_gdn_workers.run(_test_gdn_qat_helper, tmp_path) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py index 51b6c3f2790..15dc812f5ec 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py @@ -29,7 +29,6 @@ import modelopt.torch.quantization as mtq from modelopt.recipe import load_recipe -from modelopt.torch.opt.config_loader import load_config from modelopt.torch.opt.plugins.mcore_dist_checkpointing import ( restore_sharded_modelopt_state, save_sharded_modelopt_state, @@ -42,6 +41,7 @@ from modelopt.torch.quantization.nn import TensorQuantizer KimiDeltaAttention = pytest.importorskip("megatron.core.ssm.gated_delta_net.kda").KimiDeltaAttention +pytest.importorskip("vllm.model_executor.layers.fla.ops.kda", exc_type=ModuleNotFoundError) def _layer(): @@ -92,7 +92,7 @@ def forward(layer): policy = getattr(layer, "linear_attention_config", None) phase = ( linear_attention_training_phase(layer, [31]) - if policy is not None and policy.decode is not None + if policy is not None and policy.backend == "serving" else nullcontext() ) with phase, torch.autocast("cuda", dtype=torch.bfloat16): @@ -115,7 +115,8 @@ def _test_kda(rank, size, cfg, checkpoint_path): model, hidden, forward, baseline = _case(cfg) assert model.kda_state_quantizer.is_enabled assert model.kda_state_quantizer.num_bits == 8 - assert model.linear_attention_config.decode.state_codec == "int8_hadamard32" + assert model.linear_attention_config.precision == "vllm_0_15" + assert model.kda_state_quantizer.block_sizes == {-1: 32} kernel = model.gated_delta_rule with torch.no_grad(): assert not torch.equal(forward(model), baseline) @@ -128,7 +129,6 @@ def _test_kda(rank, size, cfg, checkpoint_path): torch.testing.assert_close(forward(model), baseline, rtol=0, atol=0) model.linear_attention_config = policy mtq.enable_quantizer(model, "*kda_state_quantizer") - mtq.enable_quantizer(model, "*replay_*_quantizer") with torch.no_grad(): expected = forward(model) @@ -148,7 +148,8 @@ def _test_kda(rank, size, cfg, checkpoint_path): assert not hasattr(restored, "kda_w_quantizer") assert restored._linear_attention_prefill_lengths is None assert restored.linear_attention_config == policy - assert restored.replay_key_quantizer.is_enabled and restored.replay_update_quantizer.is_enabled + assert not restored.replay_key_quantizer.is_enabled + assert not restored.replay_update_quantizer.is_enabled assert restored.gated_delta_rule is kernel assert torch.isfinite(hidden.grad).all() for parameter in restored.parameters(): @@ -162,16 +163,9 @@ def _test_kda(rank, size, cfg, checkpoint_path): @pytest.fixture(scope="module") def compiled_kda_workers(dist_workers_size_1): """Warm one KDA shape outside the functional test timer.""" - cfg = load_recipe("general/ptq/linear_attention_state_int8_dynamic").quantize.model_dump() - # Exercise mathematical replay/checkpointing without the optional serving fork. - cfg["linear_attention"][0]["cfg"]["backend"] = "reference" - cfg["linear_attention"][0]["cfg"]["decode"].update( - precision="full", mode="replay", replay={"window": 5}, decay_log_step=1 / 256 - ) - cfg["quant_cfg"].extend( - entry.model_dump(exclude_unset=True) - for entry in load_config("configs/ptq/units/linear_attention_replay_fp8_dynamic") - ) + cfg = load_recipe( + "general/ptq/linear_attention_state_int8_block32_dynamic" + ).quantize.model_dump() dist_workers_size_1.run(_compile_kda, cfg) return dist_workers_size_1, cfg diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py index f8084a88122..da1335342f6 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py @@ -21,10 +21,7 @@ import pytest import torch -from modelopt.torch.quantization.linear_attention import ( - LinearAttentionDecodeConfig, - recurrent_decode, -) +from modelopt.torch.quantization.linear_attention import LinearAttentionConfig, recurrent_decode @pytest.fixture(scope="module", params=[False, True], ids=["gdn-token", "kda-replay"]) @@ -53,12 +50,10 @@ def compiled_replay(request): beta = raw_beta.float().sigmoid() args = tuple(x.detach().requires_grad_() for x in (q, k, v, g, beta)) initial = (torch.randn(1, 64, 64, device="cuda") * 0.1).requires_grad_() - config = LinearAttentionDecodeConfig( + config = LinearAttentionConfig( + backend="serving", precision="replayssm", - state_codec="int8_hadamard32", - readout="working", - mode="replay" if channel else "token", - replay={"window": window} if channel else None, + replay_window=window, ) kwargs = { "config": config, @@ -80,7 +75,7 @@ def test_replay_matches_persistent_serving_cache(compiled_replay): native, args, initial, kwargs, forward = compiled_replay output, carry = forward() channel = args[3].ndim == 3 - window = kwargs["config"].replay.window if channel else 1 + window = kwargs["config"].replay_window state = torch.zeros(2, 1, 64, 64, device="cuda", dtype=torch.int8) scales = torch.zeros(2, 1, 2, 64, device="cuda", dtype=torch.float16) updates = torch.zeros(2, 1, window, 64, device="cuda", dtype=torch.bfloat16) diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py index 623be1bf48b..616551b4973 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py @@ -39,9 +39,7 @@ def compiled_serving_case(request): torch.rand(1, 73, 1, device="cuda") * 0.4, ] args = [x.requires_grad_() for x in args] - policy = LinearAttentionConfig( - backend="serving", decode={"precision": "vllm_0_15", "readout": "working"} - ) + policy = LinearAttentionConfig(backend="serving", precision="vllm_0_15") quantizer = TensorQuantizer( QuantizerAttributeConfig( num_bits=8, diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index 0d2d498e895..9bfcc33143e 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -69,16 +69,14 @@ def forward(self, x): 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"} +def quant_cfg(): + return { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*gdn_state_quantizer", "cfg": GDN_STATE_FP8_DYNAMIC}, + ], + "algorithm": "max", + } @pytest.mark.parametrize( @@ -145,95 +143,77 @@ def forward(*args, **kwargs): return chunk_gated_delta_rule(*args) monkeypatch.setattr(gdn, "matmul_gdn", forward) - model.linear_attention_config = LinearAttentionConfig(backend="serving", decode={}) + model.linear_attention_config = LinearAttentionConfig(backend="serving") with linear_attention_training_phase(model, [4, 4]): model(x) assert calls[-1]["prefill_lengths"] == (4, 4) -def test_w_quantizer_is_passed_to_the_kernel(monkeypatch): - """Pass the configured W TensorQuantizer to the FLA wrapper.""" - calls = [] - - def fake_w_qdq_kernel(*args, **kwargs): - calls.append(kwargs) - return chunk_gated_delta_rule(*args) - - monkeypatch.setattr(gdn, "_w_qdq_chunk_gated_delta_rule", lambda: fake_w_qdq_kernel) - model = TinyGatedDeltaNet() - x = torch.randn(2, 8, 3, 4) - mtq.quantize(model, quant_cfg(state=False, w=True), lambda m: m(x)) - assert model.gdn_w_quantizer.is_enabled and not model.gdn_state_quantizer.is_enabled - - model(x) - 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("phase", ["convert", "forward", "restore"]) +def test_enabled_legacy_w_quantizer_is_rejected(phase): + cfg = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None} + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + with pytest.raises(ValueError, match="GDN W quantization is no longer supported"): + if phase == "convert": + cfg["quant_cfg"].append({"quantizer_name": "*gdn_w_quantizer", "enable": True}) + mtq.quantize(TinyGatedDeltaNet(), cfg) + else: + model.gdn_w_quantizer.enable() + if phase == "forward": + model(torch.randn(2, 8, 3, 4)) + else: + mto.restore_from_modelopt_state(TinyGatedDeltaNet(), mto.modelopt_state(model)) -@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) +def test_unsupported_quantizer_fails_during_conversion(overrides): + cfg = deepcopy(quant_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): +def test_quantizer_roundtrip_and_hybrid_selection(tmp_path): model = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) - cfg = deepcopy(quant_cfg(state=state, w=w)) + cfg = deepcopy(quant_cfg()) cfg["algorithm"] = None - if state and not w: - cfg["quant_cfg"][-1]["cfg"].update( - num_bits=8, axis=None, block_sizes={-1: 32}, unsigned=False, narrow_range=True - ) + cfg["quant_cfg"][-1]["cfg"].update( + num_bits=8, axis=None, block_sizes={-1: 32}, unsigned=False, narrow_range=True + ) cfg["linear_attention"] = [ - {"module_name": "*", "cfg": {"state": {"block_v": 128}}}, - {"module_name": "0", "cfg": {"state": {"block_v": 32}}}, + {"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 state and not w: - sample = torch.randn(2, 4, 19) - expected = model[0].gdn_state_quantizer(sample) - if w: - # Save the actual quantizer settings, including edits after conversion. - model[0].gdn_w_quantizer.axis = None + model[0].linear_attention_config.state_block_v = 16 + sample = torch.randn(2, 4, 19) + expected = model[0].gdn_state_quantizer(sample) path = tmp_path / "gdn.pth" mto.save(model, path) # Load policies saved before fixed, non-configurable fields were removed. checkpoint = torch.load(path) for _, mode_state in checkpoint["modelopt_state"]["modelopt_state_dict"]: for policy in mode_state["metadata"]["linear_attention"].values(): - policy["state"].update(mode="chunk", quantize_initial=True) + policy["schema_version"] = 2 + policy["state"] = { + "block_v": policy.pop("state_block_v"), + "mode": "chunk", + "quantize_initial": True, + } + policy["decode"] = None + policy.pop("precision") + policy.pop("replay_window") policy["solve"] = {"method": "exact"} torch.save(checkpoint, 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)): + for name, enabled in (("gdn_state_quantizer", True), ("gdn_w_quantizer", False)): original = getattr(model[0], name) quantizer = getattr(restored[0], name) assert quantizer.is_enabled == enabled @@ -242,8 +222,7 @@ def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): assert quantizer.block_sizes == original.block_sizes assert quantizer._dynamic == original._dynamic assert not hasattr(restored[1], name) - if state and not w: - torch.testing.assert_close(restored[0].gdn_state_quantizer(sample), expected) + torch.testing.assert_close(restored[0].gdn_state_quantizer(sample), expected) def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): @@ -253,16 +232,16 @@ def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): assert model.gdn_state_quantizer.is_enabled assert not model.gdn_w_quantizer.is_enabled - cfg = quant_cfg(state=False, w=True) + cfg = deepcopy(quant_cfg()) cfg["algorithm"] = None - cfg["linear_attention"] = [{"module_name": "", "cfg": {"state": {"block_v": 32}}}] + 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 + assert model.gdn_state_quantizer.is_enabled + assert not model.gdn_w_quantizer.is_enabled cfg = deepcopy(cfg) cfg["quant_cfg"][-1]["cfg"]["pass_through_bwd"] = False @@ -299,62 +278,16 @@ def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): assert not model.replay_update_quantizer.is_enabled -def test_replay_quantizers_use_standard_controls_and_restore(): - cfg = { - "quant_cfg": [ - {"quantizer_name": "*", "enable": False}, - { - "quantizer_name": "*replay_key_quantizer", - "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0,)}, - }, - { - "quantizer_name": "*replay_update_quantizer", - "cfg": {"num_bits": (4, 3), "type": "dynamic", "block_sizes": {-1: 16}}, - }, - ], - "linear_attention": [ - { - "module_name": "*", - "cfg": { - "backend": "matmul", - "state": {"block_v": 16}, - "decode": {"mode": "replay", "replay": {"window": 3}}, - }, - } - ], - "algorithm": None, - } - torch.manual_seed(53) +@pytest.mark.parametrize("name", ["replay_key_quantizer", "replay_update_quantizer"]) +def test_enabled_legacy_replay_quantizer_is_rejected(name): + cfg = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None} model = mtq.quantize(TinyGatedDeltaNet(), cfg) - x = torch.randn(1, 7, 1, 4) * 0.1 - with linear_attention_training_phase(model, [2]): - quantized = model(x) - saved = deepcopy(mto.modelopt_state(model)) - weights = deepcopy(model.state_dict()) - mtq.disable_quantizer(model, "*") - with linear_attention_training_phase(model, [2]): - plain = model(x) - assert not torch.equal(plain, quantized) - # Current checkpoints and pre-handle replay checkpoints reproduce the same computation. - for legacy in (None, False, True): - checkpoint = deepcopy(saved) - if legacy is not None: - for _, mode_state in checkpoint["modelopt_state_dict"]: - policy = mode_state["metadata"]["linear_attention"][""] - policy["schema_version"] = 1 - if not legacy: - policy["decode"]["replay"]["factor_qdq"] = False - for name in model.replay_quantizer_names: - mode_state["metadata"]["quantizer_state"].pop(name) - restored = mto.restore_from_modelopt_state(TinyGatedDeltaNet(), checkpoint) - restored.load_state_dict(weights) - assert restored.replay_key_quantizer.is_enabled == (legacy is not False) - assert restored.replay_update_quantizer.is_enabled == (legacy is not False) - assert "factor_qdq" not in restored.linear_attention_config.decode.replay.model_dump() - with linear_attention_training_phase(restored, [2]): - torch.testing.assert_close( - restored(x), plain if legacy is False else quantized, rtol=0, atol=0 - ) + getattr(model, name).enable() + with pytest.raises(ValueError, match=f"{name} is no longer supported"): + mto.restore_from_modelopt_state(TinyGatedDeltaNet(), mto.modelopt_state(model)) + cfg["quant_cfg"].append({"quantizer_name": f"*{name}", "enable": True}) + with pytest.raises(ValueError, match=f"{name} is no longer supported"): + mtq.quantize(TinyGatedDeltaNet(), cfg) def test_policy_rejects_unmatched_and_unimplemented_modes(): diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py index 09cfb56e36f..5a5b54220fd 100644 --- a/tests/unit/torch/quantization/test_linear_attention_decode.py +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -15,17 +15,9 @@ import pytest import torch -import torch.nn.functional as F -from modelopt.recipe import load_recipe from modelopt.torch.quantization.config import QuantizerAttributeConfig -from modelopt.torch.quantization.linear_attention import ( - LinearAttentionConfig, - LinearAttentionDecodeConfig, - matmul_gdn, - matmul_kda, - recurrent_decode, -) +from modelopt.torch.quantization.linear_attention import LinearAttentionConfig from modelopt.torch.quantization.linear_attention.decode import _encode from modelopt.torch.quantization.nn import TensorQuantizer @@ -53,88 +45,57 @@ def test_state_qdq_matches_tensor_quantizer(state_format): assert not encoded.scales.requires_grad -def _inputs(kda=True, length=11): - 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( + "config", + [ + {"backend": "reference", "decode": {}}, + {"backend": "matmul", "decode": {}}, + {"backend": "serving", "decode": {"precision": "full"}}, + ], +) +def test_reference_training_is_rejected(config): + with pytest.raises(ValueError, match="Reference training is retired"): + LinearAttentionConfig(**config) -@pytest.mark.parametrize(("kda", "block_size"), [(False, 32), (True, 64)]) -def test_block_state_quantizer_prefill_decode_and_gradients(kda, block_size): - args, initial = _inputs(kda, length=5) - # Partial value groups exercise TensorQuantizer padding as well as per-key scales. - recipe = load_recipe( - "general/ptq/linear_attention_state_int8_block32_dynamic" - ).quantize.model_dump() - cfg = recipe["quant_cfg"][2 if kda else 1]["cfg"] - cfg["block_sizes"] = {-1: block_size} - quantizer = TensorQuantizer(QuantizerAttributeConfig(**cfg)) - policy = LinearAttentionConfig(**recipe["linear_attention"][0]["cfg"]) - policy.decode.precision = "full" # Double-precision recurrence/gradient oracle. - policy.state.block_v = 16 - if kda: - policy.decode.prefill_state_qdq = True - policy.decode.readout = "stored" - function = matmul_kda if kda else matmul_gdn - output, final = function( - *(x.unsqueeze(0) for x in args), - policy=policy, - state_quantizer=quantizer, - prefill_lengths=[3], - initial_state=initial.unsqueeze(0), - output_final_state=True, +@pytest.mark.parametrize(("precision", "window"), [("vllm_0_15", 1), ("replayssm", 4)]) +def test_native_legacy_policy_preserves_precision(precision, window, tmp_path): + decode = {"precision": precision, "readout": "working"} + if precision == "replayssm": + decode.update(state_codec="int8_hadamard32", mode="replay", replay={"window": window}) + policy = LinearAttentionConfig(schema_version=2, backend="matmul", decode=decode) + assert policy == LinearAttentionConfig( + backend="serving", precision=precision, replay_window=window ) - # GDN exercises the serving-aligned recipe; KDA retains prefix-QDQ/stored-read coverage. - state = quantizer(initial) if policy.decode.prefill_state_qdq else initial - expected = [] - q, k, v, g, beta = args - for t in range(len(q)): - if t == 3: - state = quantizer(state) - decay = g[t].exp().unsqueeze(-1) if kda else g[t].exp()[:, None, None] - decayed = state * decay - residual = v[t] - (k[t].unsqueeze(-1) * decayed).sum(-2) - state = decayed + k[t].unsqueeze(-1) * (beta[t].unsqueeze(-1) * residual).unsqueeze(-2) - stored = quantizer(state) if t >= 3 else state - read = stored if t >= 3 and policy.decode.readout == "stored" else state - expected.append((q[t].unsqueeze(-1) * read).sum(-2) / q.shape[-1] ** 0.5) - state = stored - if t == 2 and policy.decode.prefill_state_qdq: - state = quantizer(state) - for actual, expected in zip( - _values_and_grads(output[0], final[0], args, initial), - _values_and_grads(torch.stack(expected), state, args, initial), - ): - torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-7) + assert "decode" not in policy.model_dump() + assert "state_codec" not in policy.model_dump() + # Reproduce the field layout saved by the old class, bypassing new validation. + legacy = LinearAttentionConfig() + legacy.__dict__.clear() + legacy.__dict__.update(schema_version=2, backend="matmul", decode=decode) + path = tmp_path / "legacy-policy.pt" + torch.save(legacy, path) + assert torch.load(path, weights_only=True) == policy -def test_serving_precision_requires_native_schedule(): - with pytest.raises(ValueError, match="working readout"): - LinearAttentionDecodeConfig(precision="vllm_0_15") +@pytest.mark.parametrize( + "settings", [{"readout": "stored"}, {"prefill_state_qdq": True}, {"decay_log_step": 0.02}] +) +def test_serving_precision_requires_native_schedule(settings): + with pytest.raises(ValueError, match="Native state QAT requires"): + LinearAttentionConfig(backend="serving", decode=settings) -def test_grid_gate_ste_keeps_gate_gradients_and_changes_trajectory(): - args, state = _inputs() - output, carry = recurrent_decode( - *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(*args, config=LinearAttentionDecodeConfig(), initial_state=state) - assert not torch.equal(output, exact) +@pytest.mark.parametrize( + "settings", + [ + {"replay_window": 4}, + {"precision": "replayssm", "replay_window": 0}, + {"precision": "replayssm", "replay_window": 65}, + {"precision": "replayssm", "state_block_v": 16}, + {"decode": {}, "precision": "replayssm"}, + ], +) +def test_unified_policy_rejects_incompatible_settings(settings): + with pytest.raises(ValueError): + LinearAttentionConfig(backend="serving", **settings) diff --git a/tests/unit/torch/quantization/test_linear_attention_hadamard.py b/tests/unit/torch/quantization/test_linear_attention_hadamard.py index 71a4cb23fa4..5f51d6344c3 100644 --- a/tests/unit/torch/quantization/test_linear_attention_hadamard.py +++ b/tests/unit/torch/quantization/test_linear_attention_hadamard.py @@ -14,12 +14,8 @@ # limitations under the License. import torch -import torch.nn.functional as F -from modelopt.torch.quantization.linear_attention import ( - LinearAttentionDecodeConfig, - recurrent_decode, -) +from modelopt.torch.quantization.linear_attention.decode import _encode, _hadamard32 def _rotate(value): @@ -38,62 +34,22 @@ def _qdq(value): return (value - value.detach()) + rounded -def _inputs(length=11): +def test_hadamard_codec_matches_dense_oracle_and_identity_ste(): 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, 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): - q, k, v, g, beta = args - state = _rotate(initial) - 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 (t + 1) % 3 == 0 else working - outputs.append(_rotate(torch.einsum("hk,hkv->hv", q[t], state)) / 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_replay_matches_dense_oracle_and_split_carry(): - args, initial = _inputs() - cfg = LinearAttentionDecodeConfig( - mode="replay", + state = torch.randn(2, 4, 64, dtype=torch.float64, requires_grad=True) + encoded = _encode( + _hadamard32(state), + True, + 64, + state=True, + state_format="int8", state_codec="int8_hadamard32", - replay={"window": 3}, - ) - kwargs = {"config": cfg, "state_format": "int8", "state_qdq": True} - output, carry = recurrent_decode(*args, initial_state=initial, **kwargs) - _check_with_grads((output, carry.reconstruct()), _oracle(args, initial), args, initial) - assert carry.value_basis == "hadamard32" and carry.anchor.block_v == 32 - assert carry.anchor.scales.shape == (2, 4, 2) - first, split_carry = recurrent_decode(*(x[:2] for x in args), initial_state=initial, **kwargs) - second, split_carry = recurrent_decode(*(x[2:] for x in args), carry=split_carry, **kwargs) - _check_with_grads( - (torch.cat((first, second)), split_carry.reconstruct()), - (output, carry.reconstruct()), - args, - initial, - 0, ) + actual = _hadamard32(encoded.values) + expected = _rotate(_qdq(_rotate(state))) + torch.testing.assert_close(actual, expected, rtol=1e-10, atol=1e-10) + assert encoded.scales.shape == (2, 4, 2) + assert encoded.scales.dtype == torch.float16 + probe = torch.randn_like(state) + (gradient,) = torch.autograd.grad((actual * probe).sum(), state) + torch.testing.assert_close(gradient, probe, rtol=1e-10, atol=1e-10) diff --git a/tests/unit/torch/quantization/test_linear_attention_reference.py b/tests/unit/torch/quantization/test_linear_attention_reference.py index 6b3669c47d8..1744fc07d9c 100644 --- a/tests/unit/torch/quantization/test_linear_attention_reference.py +++ b/tests/unit/torch/quantization/test_linear_attention_reference.py @@ -17,7 +17,7 @@ import torch import torch.nn.functional as F from _test_utils.torch.quantization.linear_attention_reference import ( - chunk_gdn_reference, + chunk_delta_rule_reference, recurrent_delta_rule_reference, state_qdq_reference, ) @@ -34,10 +34,11 @@ def inputs(length=7, kda=False): 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() +@pytest.mark.parametrize( + ("kda", "packed", "state_v_first"), [(False, True, True), (True, False, False)] +) +def test_chunk_recurrence_output_state_and_gradients(kda, packed, state_v_first): + args = inputs(kda=kda) state = torch.randn(2 if packed else 1, 2, 3, 4, dtype=torch.float64) if state_v_first: state = state.transpose(-1, -2).contiguous() @@ -47,7 +48,7 @@ def test_chunk_recurrence_output_state_and_gradients(packed, state_v_first): "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) + actual = chunk_delta_rule_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) From 0d036200460627e278cf6b7b37f839080957fbe8 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 7 Oct 2026 10:40:47 -0700 Subject: [PATCH 14/27] Simplify state QAT interfaces and fix BF16 replay Normalize recurrent working values to FP32 and require an explicit serving policy for both GDN and KDA state quantization. Remove unpublished config migrations, unused replay-factor plumbing, and the unreferenced Triton INT8 helper while retaining native INT8/Hadamard replay. Align GDN/KDA API and runtime-state names, document the state handoff, and keep shared TensorQuantizer/FP8 corrections outside this PR. Update recipe guidance and retain minimal real-path tests. Validation: 28 focused CPU tests, 5 native GPU cases, 2 Megatron QAT/sharded-restore tests, scoped pre-commit hooks, and git diff --check passed. Signed-off-by: Kai Xu --- CHANGELOG.rst | 1 - .../kernels/quantization/common/int8_quant.py | 40 ---- modelopt/torch/quantization/conversion.py | 10 - .../quantization/linear_attention/__init__.py | 2 +- .../quantization/linear_attention/config.py | 125 +---------- .../quantization/linear_attention/decode.py | 202 +++++++++++------- .../linear_attention/{prefill.py => gdn.py} | 26 +-- .../quantization/linear_attention/kda.py | 26 +-- .../quantization/linear_attention/training.py | 8 +- .../nn/modules/tensor_quantizer.py | 7 +- modelopt/torch/quantization/plugins/gdn.py | 12 +- modelopt/torch/quantization/plugins/kda.py | 21 +- .../quantization/plugins/linear_attention.py | 21 +- .../torch/quantization/plugins/megatron.py | 6 +- modelopt/torch/quantization/tensor_quant.py | 3 +- modelopt_recipes/configs/ptq/units/README.md | 3 +- .../ptq/units/gdn_state_fp8_dynamic.yaml | 16 +- .../linear_attention_state_int8_dynamic.yaml | 2 +- .../quantization/test_tensor_quant_cuda.py | 14 +- .../quantization/plugins/test_megatron_kda.py | 7 - .../test_linear_attention_replay.py | 12 +- .../test_linear_attention_training.py | 6 +- .../torch/quantization/plugins/test_gdn.py | 70 +++--- .../test_linear_attention_decode.py | 46 +--- .../test_linear_attention_hadamard.py | 30 +-- .../quantization/test_tensor_quantizer_cpu.py | 13 -- 26 files changed, 220 insertions(+), 509 deletions(-) delete mode 100644 modelopt/torch/kernels/quantization/common/int8_quant.py rename modelopt/torch/quantization/linear_attention/{prefill.py => gdn.py} (76%) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index c139abb7377..a28a3597900 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -68,7 +68,6 @@ Changelog **Backward Breaking Changes** - Bump minimum transformers version to ``5.5`` instead of ``4.57``; transformers 4.x is no longer supported. Upgrade with ``pip install -U "nvidia-modelopt[hf]"``. -- Experimental linear-attention reference training is removed, including ``precision="full"`` and replay-factor QDQ; select ``backend="serving"`` with a native precision profile. Mathematical reference implementations are available only in tests. - Experimental GDN W quantization and its ``gdn_w_fp8_dynamic`` recipe unit are removed; keep ``gdn_w_quantizer`` disabled when using state QAT or restoring existing checkpoints. - ``modelopt.torch.distill.plugins.megatron.TopKLogitsKLLoss`` (``logit_kl_topk`` in ``DistillationConfig``) is renamed to ``TopLogitsKLLoss``, keeping the old name as a deprecated alias. It now normalizes both distributions over the full vocabulary instead of re-normalizing over the Top-K entries, and always appends a "ghost" token holding the probability mass outside the Top-K to both student and teacher (matching Megatron-LM's offline cached-logits KD loss). Loss values change for existing ``logit_kl_topk`` runs. diff --git a/modelopt/torch/kernels/quantization/common/int8_quant.py b/modelopt/torch/kernels/quantization/common/int8_quant.py deleted file mode 100644 index bcd27d7e089..00000000000 --- a/modelopt/torch/kernels/quantization/common/int8_quant.py +++ /dev/null @@ -1,40 +0,0 @@ -# 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. - -"""Composable signed narrow-range INT8 QDQ matching CUDA TensorQuantizer.""" - -import triton -import triton.language as tl -from triton.language.extra.cuda import libdevice - -__all__ = [] - - -@triton.jit -def int8_scalar_qdq(value, amax): - """QDQ with a scalar or broadcastable amax, using CUDA TensorQuantizer arithmetic.""" - # Match gemm/tensor_quant_gpu.cu: tiny groups become zero; ties round to even. - tiny = amax < 2.0**-24 - scale = tl.div_rn(127.0, tl.where(tiny, 1.0, amax)) - codes = tl.clamp(libdevice.nearbyint(value * scale), -127.0, 127.0) - return tl.where(tiny, 0.0, tl.div_rn(codes, scale)) - - -@triton.jit -def int8_block_qdq(value, GROUP_SIZE: tl.constexpr): - """Dynamic INT8 QDQ per row and last-axis group of a two-dimensional tile.""" - groups = tl.reshape(value, (value.shape[0], value.shape[1] // GROUP_SIZE, GROUP_SIZE)) - amax = tl.max(tl.abs(groups), axis=2, keep_dims=True) - return tl.reshape(int8_scalar_qdq(groups, amax), value.shape) diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 84d7256c76b..db194a0f465 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -146,17 +146,7 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: # QuantModule.set_extra_state(). return model - # Optional plugins import conversion during initialization, so defer this import. - from .plugins.kda import _discard_legacy_w_quantizer_state - quantizer_state_dict = dict(metadata["quantizer_state"]) - _discard_legacy_w_quantizer_state(quantizer_state_dict) - for name, module in _linear_attention_modules(model).items(): - policy = metadata.get("linear_attention", {}).get(name, {}) - if policy.get("schema_version", 1) < 2: - for handle in module.replay_quantizer_names: - key = f"{name}.{handle}" if name else handle - quantizer_state_dict.setdefault(key, getattr(module, handle).get_modelopt_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(): diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py index 883c701bf4c..42bda38f4e4 100644 --- a/modelopt/torch/quantization/linear_attention/__init__.py +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -17,6 +17,6 @@ from .config import * from .decode import * +from .gdn import * from .kda import * -from .prefill import * from .training import * diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index 88134b77779..af8dd9dd1a9 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -24,44 +24,6 @@ __all__ = ["LinearAttentionConfig", "LinearAttentionPolicyEntry"] -# Deserialization-only names for pre-unification pickles. ModeloptBaseConfig -# registers each original class name with torch.serialization's safe globals. -class _StateConfig(ModeloptBaseConfig): - """Legacy state settings; migrated by LinearAttentionConfig on load.""" - - -class _SolveConfig(ModeloptBaseConfig): - """Legacy solve settings; migrated by LinearAttentionConfig on load.""" - - -class LinearAttentionReplayConfig(ModeloptBaseConfig): - """Legacy replay settings; migrated by LinearAttentionConfig on load.""" - - -class LinearAttentionDecodeConfig(ModeloptBaseConfig): - """Legacy decode settings; migrated by LinearAttentionConfig on load.""" - - -def _legacy_fields(value): - if isinstance(value, ModeloptBaseConfig): - # model_dump would omit the fields of these deserialization-only classes. - fields = dict(vars(value)) - factor_qdq = (value.__pydantic_private__ or {}).get("_legacy_factor_qdq") - if factor_qdq is not None: - fields["factor_qdq"] = factor_qdq - return fields - if not isinstance(value, dict): - raise ValueError("Legacy linear-attention settings must be dictionaries") - return dict(value) - - -def _consume_fixed_settings(values, **defaults): - """Remove fixed legacy settings without silently changing their numerics.""" - for key, default in defaults.items(): - if values.pop(key, default) != default: - raise ValueError(f"Native state QAT requires {key}={default!r}") - - class LinearAttentionConfig(ModeloptBaseConfig): """One execution policy shared by the chunked prefix and recurrent suffix. @@ -69,13 +31,13 @@ class LinearAttentionConfig(ModeloptBaseConfig): ``fla`` is the disabled/default path. ``vllm_0_15`` uses public vLLM arithmetic; ``replayssm`` uses the serving fork's INT8/Hadamard checkpoint and ring kernels. ``replay_window=1`` refreshes the state every token; larger windows enable replay. - Both profiles leave prefill unquantized and quantize the handoff state. + Fresh prefill has no internal state QDQ; incoming continuation state and + the handoff to a nonempty suffix use the selected checkpoint encoding. TensorQuantizer independently enables QDQ and owns its format and grouping. ``state_block_v`` controls grouping only for legacy per-tile quantizers. """ - schema_version: Literal[3] = ModeloptField(default=3) backend: Literal["fla", "serving"] = ModeloptField(default="fla") precision: Literal["vllm_0_15", "replayssm"] = ModeloptField(default="vllm_0_15") replay_window: int = Field(default=1, ge=1, le=64, strict=True) @@ -86,83 +48,6 @@ def state_codec(self): """The native profile fixes the codec; it is not an independent setting.""" return "int8_hadamard32" if self.precision == "replayssm" else "tile" - @model_validator(mode="before") - @classmethod - def _load_legacy_settings(cls, values): - if not isinstance(values, dict): - return values - values = dict(values) - version = values.get("schema_version", 3) - if version in (1, 2): - values["schema_version"] = 3 - if values.get("backend") == "reference": - raise ValueError("Reference training is retired; select backend='serving'") - _consume_fixed_settings(values, chunk_size=64) - if "solve" in values: - solve = _legacy_fields(values.pop("solve")) - _consume_fixed_settings(solve, method="exact") - if solve: - raise ValueError(f"Unsupported legacy solve settings: {solve}") - if "state" in values: - state = _legacy_fields(values.pop("state")) - _consume_fixed_settings(state, mode="chunk", quantize_initial=True) - if "state_block_v" in values: - raise ValueError("Supply state_block_v without legacy state settings") - values["state_block_v"] = state.pop("block_v", 64) - if state: - raise ValueError(f"Unsupported legacy state settings: {state}") - if "decode" in values: - decode = values.pop("decode") - if (values.get("backend", "fla") != "fla") != (decode is not None): - raise ValueError("Legacy state QAT requires an explicit decode policy") - if decode is not None: - if "precision" in values or "replay_window" in values: - raise ValueError( - "Supply precision/replay_window without legacy decode settings" - ) - decode = _legacy_fields(decode) - precision = decode.pop( - "precision", "full" if values.get("backend") == "matmul" else "vllm_0_15" - ) - if precision == "full": - raise ValueError( - "Reference training is retired; select a native serving precision" - ) - values["precision"] = precision - replay = decode.pop("replay", None) - mode = decode.pop("mode", "token") - if mode not in ("token", "replay") or (mode == "replay") != (replay is not None): - raise ValueError("Legacy replay settings require mode='replay'") - _consume_fixed_settings( - decode, - readout="working", - quantize_initial=True, - prefill_state_qdq=False, - decay_log_step=None, - implementation="torch", - state_codec="int8_hadamard32" if precision == "replayssm" else "tile", - ) - if decode: - raise ValueError(f"Unsupported legacy decode settings: {decode}") - if replay is not None: - replay = _legacy_fields(replay) - factor_qdq = replay.pop("factor_qdq", version == 1) - if factor_qdq is not False: - raise ValueError( - "Replay factor QDQ is retired; native ReplaySSM stores BF16 factors" - ) - values["replay_window"] = replay.pop("window", 8) - _consume_fixed_settings(replay, encoding="once") - if replay: - raise ValueError(f"Unsupported legacy replay settings: {replay}") - if precision != "replayssm": - raise ValueError("Replay requires precision='replayssm'") - if values.get("backend") == "matmul": - if "precision" not in values: - raise ValueError("Reference training is retired; select a native serving precision") - values["backend"] = "serving" - return values - @model_validator(mode="after") def _validate_profile(self): if self.precision == "replayssm": @@ -174,12 +59,6 @@ def _validate_profile(self): raise ValueError("Replay requires precision='replayssm'") return self - def __setstate__(self, state): - # Config objects can be pickled inside ModelOpt checkpoints. Normalize - # their old nested fields before any runtime consumer sees the policy. - migrated = type(self).model_validate(state["__dict__"]) - super().__setstate__(migrated.__getstate__()) - class LinearAttentionPolicyEntry(ModeloptBaseConfig): """Assign a complete policy to supported modules matching ``module_name``. diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py index d6fd0702756..9efa231f9eb 100644 --- a/modelopt/torch/quantization/linear_attention/decode.py +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -25,6 +25,7 @@ __all__ = [ "EncodedLinearAttentionTensor", "LinearAttentionCarry", + "LinearAttentionState", "ReplayEntry", "recurrent_decode", ] @@ -32,7 +33,11 @@ @dataclass class EncodedLinearAttentionTensor: - """Fake decoded values plus optional scale metadata; no compressed storage claim.""" + """Floating QDQ values with autograd history and optional quantization scales. + + ``format`` and ``block_v`` describe the emulated encoding; ``values`` stores + decoded values, not packed INT8/FP8 codes. Scale metadata is non-differentiable. + """ values: torch.Tensor scales: torch.Tensor | None @@ -42,7 +47,12 @@ class EncodedLinearAttentionTensor: @dataclass class ReplayEntry: - """An already computed, encoded rank-one update and its log retention.""" + """One saved rank-one update since the last state checkpoint. + + ``key`` is [H,K] and ``update`` is the beta-scaled residual [H,V]. ReplaySSM + stores their BF16 forward values once. ``log_retention`` is [H] for GDN or + [H,K] for KDA and records this token's decay for subsequent reconstruction. + """ key: EncodedLinearAttentionTensor update: EncodedLinearAttentionTensor @@ -50,16 +60,22 @@ class ReplayEntry: @dataclass -class LinearAttentionCarry: - """Explicit anchor/update state, position, and codec contract across calls.""" +class LinearAttentionState: + """Runtime state passed between recurrent tokens and successive decode calls. - anchor: EncodedLinearAttentionTensor - entries: tuple[ReplayEntry, ...] - position: int - started: bool - signature: str - value_basis: str = "identity" - precision: str = "vllm_0_15" + With ``replay_window=1``, every token refreshes ``anchor`` and ``entries`` + stays empty. With a larger window, the anchor stays fixed while updates + accumulate; reaching the window refreshes the anchor and clears the entries. + Tensor values retain their autograd history for QAT/QAD across call boundaries. + """ + + anchor: EncodedLinearAttentionTensor # Checkpoint [H,K,V]; QDQ is applied when enabled. + entries: tuple[ReplayEntry, ...] # Updates since the checkpoint, in token order. + position: int # Next token position, including the supplied prefix offset. + started: bool # Initial-state encoding has run; resuming must not repeat it. + signature: str # Execution/quantizer settings used to check compatible reuse. + value_basis: str = "identity" # Anchor and updates may use the Hadamard value basis. + precision: str = "vllm_0_15" # Selects the native profile's reconstruction arithmetic. @property def cursor(self): @@ -67,7 +83,13 @@ def cursor(self): return len(self.entries) def reconstruct(self, *, original_basis=True): - """Replay entries, returning the original value basis unless explicitly disabled.""" + """Return the current dense [H,K,V] state without refreshing or quantizing it. + + Apply pending updates and decay to the anchor, preserving gradients. + Decode uses ``original_basis=False`` to continue in the stored basis; + the default converts Hadamard state back for the caller's final-state output. + Keep this object, rather than only this tensor, to resume a replay window. + """ state = _replay_state(self) if self.precision == "replayssm" else self.anchor.values if original_basis and self.value_basis == "hadamard32": # The optional serving fork owns the checkpoint basis transform. @@ -79,6 +101,10 @@ def reconstruct(self, *, original_basis=True): return state +# Compatibility name for existing imports and pickled runtime state objects. +LinearAttentionCarry = LinearAttentionState + + def _hadamard32(value): """Apply the orthonormal Sylvester transform to contiguous 32-value groups.""" shape = value.shape @@ -89,8 +115,12 @@ def _hadamard32(value): return value * (32**-0.5) -def _replay_state(carry, current_gate=None): - """Differentiate serving's weighted BF16 ring reconstruction, including KDA decay.""" +def _replay_state(carry: LinearAttentionState, current_gate=None): + """Decay the anchor and add saved updates weighted by their subsequent gates. + + KDA weights keys per channel; GDN weights updates per head. ``current_gate`` + includes the next KDA token's decay before BF16 weighting, without adding its update. + """ state = carry.anchor.values if not carry.entries: if current_gate is not None: @@ -117,46 +147,26 @@ def _encode( enabled, block_v, *, - state=False, state_format="fp8_e4m3", - state_codec="tile", state_quantizer=None, ): + """Apply TensorQuantizer state QDQ with straight-through gradients.""" if not enabled: return EncodedLinearAttentionTensor(value, None, "identity", None) - if state and state_quantizer is not None and state_quantizer.block_sizes is not None: + if state_quantizer is not None and state_quantizer.block_sizes is not None: # TensorQuantizer owns dynamic scales; the floating carry needs only its QDQ output. return EncodedLinearAttentionTensor( state_quantizer(value), None, state_format, state_quantizer.block_sizes[-1] ) - 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(forward_value(value, decoded), metadata, "int8", 32) decoded, scales = _tile_qdq(value, block_v, state_format, state_quantizer=state_quantizer) return EncodedLinearAttentionTensor(decoded, scales, state_format, block_v) -def _signature(config, state_qdq, block_v, state_format, state_quantizer, replay_quantizers): +def _signature(config, state_qdq, block_v, state_format, state_quantizer): + """Record the execution and quantizer settings required when resuming a state.""" signature = config.model_dump_json() + f"/{state_qdq}/{block_v}/{state_format}" if state_quantizer is not None and state_quantizer.block_sizes is not None: signature += f"/group={state_quantizer.block_sizes[-1]}" - for quantizer in replay_quantizers: - if quantizer is None: - signature += "/factor=None" - else: - grouping = ( - quantizer.block_sizes if quantizer.block_sizes is not None else quantizer.axis - ) - signature += ( - f"/factor={quantizer.is_enabled}/{quantizer._if_quant}/{quantizer.num_bits}" - f"/{grouping}/{quantizer.backend}/{quantizer.backend_extra_args}" - ) return signature @@ -185,8 +195,13 @@ def _prepare_carry( position, state_format, state_quantizer, - replay_quantizers, ): + """Initialize from a dense state or validate and resume an existing runtime state. + + The first nonempty call applies the handoff transform/QDQ once. A resumed + state preserves its anchor, replay cursor, and token position; an empty call + leaves initial encoding deferred until a token is actually processed. + """ 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): @@ -203,16 +218,14 @@ def _prepare_carry( 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, state_quantizer, replay_quantizers - ) + signature = _signature(config, state_qdq, block_v, state_format, state_quantizer) 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), + carry = LinearAttentionState( + _encode(initial_state, False, block_v, state_format=state_format), (), position, False, @@ -224,24 +237,27 @@ def _prepare_carry( 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, - block_v, - state=True, - state_format=state_format, - state_codec=config.state_codec, - state_quantizer=state_quantizer, - ) if config.precision == "replayssm": # The optional serving fork owns checkpoint rounding and Hadamard arithmetic. from ...kernels.quantization.linear_attention.serving.replay import checkpoint with torch.no_grad(): decoded, metadata = checkpoint(carry.anchor.values, state_qdq) - anchor.values = forward_value(initial, decoded) - anchor.scales = metadata - carry = LinearAttentionCarry( + anchor = EncodedLinearAttentionTensor( + forward_value(initial, decoded), + metadata, + "int8" if state_qdq else "identity", + 32 if state_qdq else None, + ) + else: + anchor = _encode( + initial, + state_qdq, + block_v, + state_format=state_format, + state_quantizer=state_quantizer, + ) + carry = LinearAttentionState( anchor, (), carry.position, @@ -264,22 +280,34 @@ def recurrent_decode( state_qdq=False, state_format="fp8_e4m3", state_quantizer=None, - replay_key_quantizer=None, - replay_update_quantizer=None, initial_state=None, - carry=None, + carry: LinearAttentionState | None = None, position=0, scale=None, use_qk_l2norm_in_kernel=False, replay_gate_inputs=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. - ``config`` is the same LinearAttentionConfig used by the chunked prefix. - Native ReplaySSM stores BF16 factors; legacy factor quantizers must remain disabled. +) -> tuple[torch.Tensor, LinearAttentionState]: + """Run a recurrent suffix and return its outputs and resumable runtime state. + + Q/K are [T,H,K], V is [T,H,V], and beta is [T,H]; heads must be aligned. + Log gates are [T,H] for GDN or [T,H,K] for KDA. The first token consumes the + handoff state after enabled QDQ. Each token reads its output from the working + state before checkpoint rounding; replay_window determines when that state + is saved. Inputs and the initial state are promoted to FP32 working values; + native forward values use a differentiable adjoint for training. + + Args: + config: The same LinearAttentionConfig used by the chunked prefix. + initial_state: Dense [H,K,V] state in the original value basis, normally + from prefill. Defaults to zeros; mutually exclusive with ``carry``. + carry: LinearAttentionState returned by an earlier call. Pass it directly + to preserve the replay window and autograd history across calls. + position: Starting token offset when creating a state, usually the prefix + length. An existing carry retains its own position. + + Returns: + Outputs [T,H,V] and the updated LinearAttentionState, both retaining their + graphs. An empty call performs no state write or initial-state QDQ. """ state_quantizer, state_qdq, state_format = _resolve_state_quantizer( state_quantizer, state_qdq, state_format @@ -291,6 +319,10 @@ def recurrent_decode( native_replay = config.precision == "replayssm" if q.device.type != "cuda": raise ValueError("Serving arithmetic requires CUDA") + # BF16 inputs need FP32 working values for both replay reconstruction and its adjoint. + q, k, v, g, beta = (x.float() for x in (q, k, v, g, beta)) + if initial_state is not None: + initial_state = initial_state.float() if ( native_replay and len(q) @@ -301,9 +333,6 @@ def recurrent_decode( ) ): raise ValueError("ReplaySSM requires power-of-two K >= 32 and raw GDN gate inputs") - replay_quantizers = (replay_key_quantizer, replay_update_quantizer) - if any(quantizer is not None and quantizer.is_enabled for quantizer in replay_quantizers): - raise ValueError("Replay factor QDQ is retired; keep replay key/update quantizers disabled") carry, signature = _prepare_carry( q, k, @@ -318,7 +347,6 @@ def recurrent_decode( position, state_format, state_quantizer, - replay_quantizers, ) if len(q) == 0: # Keep empty input gradients defined without introducing a state write. @@ -390,24 +418,31 @@ def recurrent_decode( update.values = forward_value(update.values, native[4]) gate = forward_value(gate, native[5]) entries = (*entries, ReplayEntry(key, update, gate)) + # Window 1 checkpoints every token; replay checkpoints only at a full window. + # Between checkpoints, keep the anchor and save the native BF16 update values. refresh = config.replay_window == 1 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, - state_quantizer=state_quantizer, - ) if native_replay: - anchor.values = forward_value(working, native[1]) - anchor.scales = native[2] + # Native step already quantized the checkpoint; attach its values to the adjoint. + anchor = EncodedLinearAttentionTensor( + forward_value(working, native[1]), + native[2], + "int8" if state_qdq else "identity", + 32 if state_qdq else None, + ) + else: + anchor = _encode( + working, + state_qdq, + block_v, + state_format=state_format, + state_quantizer=state_quantizer, + ) entries = () else: anchor = carry.anchor - next_carry = LinearAttentionCarry( + # A new container keeps tensors needed by earlier tokens' backward passes intact. + next_carry = LinearAttentionState( anchor, entries, carry.position + 1, @@ -419,6 +454,7 @@ def recurrent_decode( query = q[t] if native_replay and use_qk_l2norm_in_kernel: query = query / (query.square().sum(-1, keepdim=True) + 1e-6).sqrt() + # Read this token from the working state, before the next-state checkpoint QDQ. outputs.append( native_output if serving else _sum_keys(query.unsqueeze(-1) * working) * scale ) diff --git a/modelopt/torch/quantization/linear_attention/prefill.py b/modelopt/torch/quantization/linear_attention/gdn.py similarity index 76% rename from modelopt/torch/quantization/linear_attention/prefill.py rename to modelopt/torch/quantization/linear_attention/gdn.py index 18573af06a8..7b73724dec7 100644 --- a/modelopt/torch/quantization/linear_attention/prefill.py +++ b/modelopt/torch/quantization/linear_attention/gdn.py @@ -13,15 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Batched differentiable GDN prefill with materialized numerical boundaries.""" +"""GDN adapter for serving-aligned recurrent-state QAT.""" from .config import LinearAttentionConfig from .training import _prefill_decode_forward, _prepare_prefill_inputs -__all__ = ["matmul_gdn"] +__all__ = ["gdn_state_qat", "matmul_gdn"] -def matmul_gdn( +def gdn_state_qat( q, k, v, @@ -32,8 +32,6 @@ def matmul_gdn( state_qdq=False, state_format="fp8_e4m3", state_quantizer=None, - replay_key_quantizer=None, - replay_update_quantizer=None, scale=None, initial_state=None, output_final_state=False, @@ -51,14 +49,14 @@ def matmul_gdn( prefill_lengths=None, replay_gate_inputs=None, ): - """Normalize GDN inputs and run a chunked prefix plus configured suffix recurrence. + """Adapt Megatron's FLA-style GDN call to serving-aligned state QAT. - Uses FP32 working arithmetic with floating QDQ state. - The serving precision profile supplies matching native forward values - and a Torch adjoint through the rounded state trajectory. + Validate prepared scalar log gates and promote inputs to FP32 working values. + The shared training forward runs the chunked prefix and recurrent suffix with + native forward values, configured state QDQ, and a differentiable Torch adjoint. """ if cp_context is not None: - raise NotImplementedError("GDN matmul emulation does not support context parallelism") + raise NotImplementedError("GDN state QAT does not support context parallelism") output_dtype = q.dtype beta_dtype = beta.dtype q, k, v, g, beta = _prepare_prefill_inputs( @@ -69,7 +67,7 @@ def matmul_gdn( "Serving GDN expects prepared log gates and beta from the Megatron adapter" ) if g.ndim != 3: - raise ValueError("matmul_gdn requires scalar GDN log gates") + raise ValueError("gdn_state_qat requires scalar GDN log gates") return _prefill_decode_forward( q, k, @@ -80,8 +78,6 @@ def matmul_gdn( state_qdq=state_qdq, state_format=state_format, state_quantizer=state_quantizer, - replay_key_quantizer=replay_key_quantizer, - replay_update_quantizer=replay_update_quantizer, scale=scale, initial_state=initial_state, output_final_state=output_final_state, @@ -94,3 +90,7 @@ def matmul_gdn( replay_gate_inputs=replay_gate_inputs, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ) + + +# Compatibility name for callers using the original adapter API. +matmul_gdn = gdn_state_qat diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py index d9d39a3d8b9..f2238c2c7ad 100644 --- a/modelopt/torch/quantization/linear_attention/kda.py +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Differentiable KDA prefill with stable per-channel decay interactions.""" +"""KDA adapter for serving-aligned recurrent-state QAT.""" import torch import torch.nn.functional as F @@ -21,10 +21,10 @@ from .training import _prefill_decode_forward, _prepare_prefill_inputs from .utils import forward_value -__all__ = ["matmul_kda"] +__all__ = ["kda_state_qat", "matmul_kda"] -def matmul_kda( +def kda_state_qat( q, k, v, @@ -35,8 +35,6 @@ def matmul_kda( state_qdq=False, state_format="fp8_e4m3", state_quantizer=None, - replay_key_quantizer=None, - replay_update_quantizer=None, scale=None, initial_state=None, output_final_state=False, @@ -57,15 +55,15 @@ def matmul_kda( return_intermediate_states=False, prefill_lengths=None, ): - """Normalize KDA inputs and run a chunked prefix plus configured suffix recurrence. + """Adapt Megatron's FLA-style KDA call to serving-aligned state QAT. - Gates follow FLA's kernel activation formula and per-key log retention. - The serving precision profile supplies matching native forward values - and a Torch adjoint through the rounded state trajectory. + Prepare per-key-channel gates and optional beta activation using FLA's formula. + The shared training forward runs the chunked prefix and recurrent suffix with + native forward values, configured state QDQ, and a differentiable Torch adjoint. """ 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" + "KDA state QAT 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") @@ -78,7 +76,7 @@ def matmul_kda( ) dtype = q.dtype if g.ndim != 4: - raise ValueError("matmul_kda requires per-key-channel KDA log gates") + raise ValueError("kda_state_qat requires per-key-channel KDA log gates") if use_gate_in_kernel: raw_gate = g if A_log is None: @@ -107,8 +105,6 @@ def matmul_kda( state_qdq=state_qdq, state_format=state_format, state_quantizer=state_quantizer, - replay_key_quantizer=replay_key_quantizer, - replay_update_quantizer=replay_update_quantizer, scale=scale, initial_state=initial_state, output_final_state=output_final_state, @@ -120,3 +116,7 @@ def matmul_kda( prefill_lengths=prefill_lengths, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ) + + +# Compatibility name for callers using the original adapter API. +matmul_kda = kda_state_qat diff --git a/modelopt/torch/quantization/linear_attention/training.py b/modelopt/torch/quantization/linear_attention/training.py index 96c74f41fa5..af1541e8187 100644 --- a/modelopt/torch/quantization/linear_attention/training.py +++ b/modelopt/torch/quantization/linear_attention/training.py @@ -84,8 +84,6 @@ def _prefill_decode_forward( state_qdq, state_format, state_quantizer, - replay_key_quantizer, - replay_update_quantizer, scale, initial_state, output_final_state, @@ -101,7 +99,9 @@ def _prefill_decode_forward( """Run both prefill and decode phases in one differentiable training forward. Each sequence's chunked prefix produces the state for its token or ReplaySSM - suffix. Their outputs are joined in token order for the training loss. + suffix. ``recurrent_decode`` wraps that dense state in LinearAttentionState + and preserves its gradient connection to prefill. Outputs are joined in token + order; the final runtime state is reconstructed to the caller's dense layout. Args: prefill_lengths: Prefix token count per sequence. For 128 tokens, a value @@ -198,8 +198,6 @@ def _prefill_decode_forward( suffix, carry = recurrent_decode( *(x[b, split:end] for x in (q, k, v, g, beta)), config=policy, - replay_key_quantizer=replay_key_quantizer, - replay_update_quantizer=replay_update_quantizer, state_qdq=state_qdq, state_format=state_format, state_quantizer=state_quantizer, diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 5e3b677ea2f..a698514d53a 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -999,12 +999,9 @@ def set_quant_params(axis, block_reshape_size, padding, slices, amax_shape=None) if amax_shape: self._amax_shape_for_export = amax_shape - # Static scales require fixed groups; dynamic scales allow changing input shapes. - if hasattr(self, "_block_reshape_size") and not self._dynamic: + # Reshape size have already been set + if hasattr(self, "_block_reshape_size"): return - for attribute in ("_padding", "_slices"): - if hasattr(self, attribute): - delattr(self, attribute) reshape_size, quantize_axis, paddings, slices = [], [], [], [] diff --git a/modelopt/torch/quantization/plugins/gdn.py b/modelopt/torch/quantization/plugins/gdn.py index 87df30ca7e0..d1cdc04b7c5 100644 --- a/modelopt/torch/quantization/plugins/gdn.py +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -24,7 +24,7 @@ import torch -from ..linear_attention.prefill import matmul_gdn +from ..linear_attention.gdn import gdn_state_qat from .linear_attention import _LinearAttentionQuantMixin __all__ = ["GatedDeltaNetStateQuantMixin"] @@ -73,12 +73,6 @@ def _state_quantized_chunk_gated_delta_rule( self.validate_linear_attention() if not self.linear_attention_is_enabled: return gated_delta_rule(*args, **kwargs) - if self.linear_attention_config.backend == "fla" and self.gdn_state_quantizer.is_enabled: - raise ValueError( - "Chunk-only state QAT is retired. Select backend='serving' and supply prefill " - "lengths through linear_attention_training_phase. " - "Existing state checkpoints are not silently migrated to a different QDQ schedule." - ) while isinstance(gated_delta_rule, partial): args = (*gated_delta_rule.args, *args) kwargs = {**gated_delta_rule.keywords, **kwargs} @@ -91,12 +85,10 @@ def _state_quantized_chunk_gated_delta_rule( chunk_size = kwargs.pop("chunk_size", 64) if chunk_size != 64: raise ValueError("GDN fake quantization supports only chunk_size=64") - return matmul_gdn( + return gdn_state_qat( *args, policy=self.linear_attention_config, state_quantizer=self.gdn_state_quantizer, - replay_key_quantizer=self.replay_key_quantizer, - replay_update_quantizer=self.replay_update_quantizer, chunk_size=chunk_size, prefill_lengths=self._linear_attention_prefill_lengths, **kwargs, diff --git a/modelopt/torch/quantization/plugins/kda.py b/modelopt/torch/quantization/plugins/kda.py index 820ee3abad5..04766a5d8e9 100644 --- a/modelopt/torch/quantization/plugins/kda.py +++ b/modelopt/torch/quantization/plugins/kda.py @@ -17,32 +17,17 @@ from functools import partial -from ..linear_attention.kda import matmul_kda +from ..linear_attention.kda import kda_state_qat from .linear_attention import _LinearAttentionQuantMixin __all__ = ["KimiDeltaAttentionStateQuantMixin"] -def _discard_legacy_w_quantizer_state(quantizer_state): - # Early KDA checkpoints included a disabled placeholder; WY QDQ was never supported. - for name in list(quantizer_state): - if name.rsplit(".", 1)[-1] == "kda_w_quantizer": - if not quantizer_state[name].get("_disabled", False): - raise ValueError("KDA checkpoints with enabled kda_w_quantizer are unsupported") - del quantizer_state[name] - - class KimiDeltaAttentionStateQuantMixin(_LinearAttentionQuantMixin): """Adds state quantizers and decode-aware kernel routing to Kimi Delta Attention.""" linear_attention_quantizer_names = ("kda_state_quantizer",) - def validate_linear_attention(self): - """Require the materialized backend for KDA numerical emulation.""" - super().validate_linear_attention() - if self.linear_attention_is_enabled and self.linear_attention_config.backend == "fla": - raise ValueError("KDA numerical emulation requires backend='serving'") - def _state_quantized_chunk_kda(self, kernel, *args, **kwargs): self.validate_linear_attention() if not self.linear_attention_is_enabled: @@ -56,12 +41,10 @@ def _state_quantized_chunk_kda(self, kernel, *args, **kwargs): kernel = kernel.func if kernel is not chunk_kda: raise NotImplementedError("KDA quantization requires FLA's chunk_kda callable") - return matmul_kda( + return kda_state_qat( *args, policy=self.linear_attention_config, state_quantizer=self.kda_state_quantizer, - replay_key_quantizer=self.replay_key_quantizer, - replay_update_quantizer=self.replay_update_quantizer, prefill_lengths=self._linear_attention_prefill_lengths, **kwargs, ) diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py index a445f1e1c91..90b969435c2 100644 --- a/modelopt/torch/quantization/plugins/linear_attention.py +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -25,10 +25,9 @@ class _LinearAttentionQuantMixin(QuantModule): linear_attention_quantizer_names: tuple[str, ...] = () - replay_quantizer_names = ("replay_key_quantizer", "replay_update_quantizer") def _setup(self): - for name in (*self.linear_attention_quantizer_names, *self.replay_quantizer_names): + for name in self.linear_attention_quantizer_names: self._register_temp_attribute( name, TensorQuantizer(QuantizerAttributeConfig(enable=False)) ) @@ -41,28 +40,24 @@ def _linear_attn_state(self): @property def linear_attention_is_enabled(self): - """Whether an operand, state, or arithmetic policy changes the computation.""" + """Whether state quantization or the serving arithmetic policy is enabled.""" return ( - any( - getattr(self, name).is_enabled - for name in (*self.linear_attention_quantizer_names, *self.replay_quantizer_names) - ) + any(getattr(self, name).is_enabled for name in self.linear_attention_quantizer_names) or self.linear_attention_config.backend == "serving" ) def validate_linear_attention(self): """Validate quantizer contracts shared by GDN and KDA.""" - for name in self.replay_quantizer_names: - quantizer = getattr(self, name) - if quantizer.is_enabled: - raise ValueError( - f"{name} is no longer supported; native ReplaySSM stores BF16 factors" - ) if self._linear_attn_state.is_enabled: state_format, group_size = state_quantizer_config( self._linear_attn_state, name=self.linear_attention_quantizer_names[0], ) + if self.linear_attention_config.backend != "serving": + raise ValueError( + "GDN/KDA state QAT requires backend='serving' and prefill lengths " + "through linear_attention_training_phase" + ) if self.linear_attention_config.state_codec == "int8_hadamard32": if state_format != "int8": raise ValueError("int8_hadamard32 requires INT8 state quantization") diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 5283c193640..033e553a109 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -71,7 +71,7 @@ from ..utils.layerwise_calib import LayerActivationCollector from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear from .gdn import GatedDeltaNetStateQuantMixin -from .kda import KimiDeltaAttentionStateQuantMixin, _discard_legacy_w_quantizer_state +from .kda import KimiDeltaAttentionStateQuantMixin from .linear_attention import _LinearAttentionQuantMixin from .transformer_engine import _QuantTEGroupedLinear, _QuantTELayerNormLinear, _QuantTELinear @@ -300,8 +300,6 @@ def quant_module_set_extra_state(self, state: Any): quantizer_state = state.get("modelopt_quantizer_state", None) if quantizer_state is not None: - quantizer_state = dict(quantizer_state) - _discard_legacy_w_quantizer_state(quantizer_state) for name, module in self.named_modules(): if isinstance(module, TensorQuantizer): quantizer_substate = quantizer_state.get(name) @@ -1226,7 +1224,7 @@ def validate_linear_attention(self): @QuantModuleRegistry.register({GatedDeltaNet: "megatron_GatedDeltaNet"}) class _QuantGatedDeltaNet(_MegatronLinearAttentionMixin, GatedDeltaNetStateQuantMixin): - """Megatron GDN with chunk or decode-aware state fake quantization.""" + """Megatron GDN with serving-aligned state fake quantization.""" _linear_attention_kernel = ( GatedDeltaNetStateQuantMixin._state_quantized_chunk_gated_delta_rule diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 67513218137..20e083491aa 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -50,8 +50,7 @@ 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) - # Use direct division, matching the CUDA extension rather than reciprocal multiplication. - scale = torch.div(448.0, safe_amax) + scale = 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 1e685a93f9d..d02eff834ca 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -41,5 +41,4 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `moe_dispatch_nvfp4.yaml` | NVFP4 fake quantization of the MoE expert-parallel dispatch (`*dispatch_quantizer`, vLLM); Blackwell+ GPUs | | `moe_combine_nvfp4.yaml` | NVFP4 fake quantization of the MoE expert-parallel combine (`*combine_quantizer`, vLLM); Blackwell+ GPUs | -Native ReplaySSM uses BF16 key/update vectors. Legacy `replay_key_quantizer` and -`replay_update_quantizer` handles must remain disabled; factor QDQ is unsupported. +Native ReplaySSM uses BF16 key/update vectors; they are not independently quantized. diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml index 4d4b13278ba..ba02ff508d2 100644 --- a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -14,13 +14,15 @@ # 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``. +# TensorQuantizer computes amax over each [K, BV] tile of the [N, H, K, V] state, +# independently per sequence/head (axis: [0, 1]), without calibration. +# Set BV through linear_attention[].cfg.state_block_v (default: 64). +# This unit enables only the quantizer. Select backend='serving' in the execution +# policy and provide prefill lengths through linear_attention_training_phase. +# With precision='vllm_0_15', state QDQ applies at the handoff to a nonempty decode +# suffix and at each suffix token's cache write; fresh prefill has no internal QDQ. +# The serving profile requires public vLLM's GDN kernels. +# See modelopt.torch.quantization.plugins.gdn. # modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig - quantizer_name: '*gdn_state_quantizer' diff --git a/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml b/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml index 3c38ccfe568..b21cc4c2653 100644 --- a/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml +++ b/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml @@ -22,7 +22,7 @@ metadata: description: >- Dynamic INT8 GDN/KDA state QDQ with 32-value Hadamard rotation during decode. Imports the quantized-ReplaySSM serving kernels and uses working-state readout. - Prefill state remains unquantized; supply per-sequence prefix lengths through + Fresh prefill has no internal state QDQ; supply per-sequence prefix lengths through linear_attention_training_phase. No calibration is required. quantize: algorithm: diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index bf156882b15..cc3262e1917 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_fp8, get_cuda_ext_mx +from modelopt.torch.quantization.extensions import get_cuda_ext, get_cuda_ext_mx from modelopt.torch.quantization.tensor_quant import mx_format_map if triton_kernel.IS_AVAILABLE: @@ -120,18 +120,6 @@ def test_cuda_ext_tiny_amax(self): class TestScaledE4M3: - @pytest.fixture - def fp8_extension(self): - return get_cuda_ext_fp8() - - def test_eager_matches_cuda_scale_rounding(self, fp8_extension): - # 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] - actual = fp8_extension.fake_e4m3fy_with_axis(values, amax, 0) - expected = tensor_quant.fp8_eager(values, amax[:, None]) - 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_kda.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py index 15dc812f5ec..471b09cc81d 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py @@ -33,12 +33,10 @@ restore_sharded_modelopt_state, save_sharded_modelopt_state, ) -from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.linear_attention import ( LinearAttentionConfig, linear_attention_training_phase, ) -from modelopt.torch.quantization.nn import TensorQuantizer KimiDeltaAttention = pytest.importorskip("megatron.core.ssm.gated_delta_net.kda").KimiDeltaAttention pytest.importorskip("vllm.model_executor.layers.fla.ops.kda", exc_type=ModuleNotFoundError) @@ -132,11 +130,8 @@ def _test_kda(rank, size, cfg, checkpoint_path): with torch.no_grad(): expected = forward(model) - # Old checkpoints carried a disabled WY placeholder; it must not reappear on restore. - model.kda_w_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) save_distributed_checkpoint(checkpoint_path, model) save_sharded_modelopt_state([model], checkpoint_path) - del model.kda_w_quantizer restored = _layer() restore_sharded_modelopt_state([restored], checkpoint_path) load_distributed_checkpoint(checkpoint_path, restored) @@ -148,8 +143,6 @@ def _test_kda(rank, size, cfg, checkpoint_path): assert not hasattr(restored, "kda_w_quantizer") assert restored._linear_attention_prefill_lengths is None assert restored.linear_attention_config == policy - assert not restored.replay_key_quantizer.is_enabled - assert not restored.replay_update_quantizer.is_enabled assert restored.gated_delta_rule is kernel assert torch.isfinite(hidden.grad).all() for parameter in restored.parameters(): diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py index da1335342f6..d6e696b2b40 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py @@ -13,9 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Minimal native-cache and backward checks for the optional ReplaySSM serving fork.""" import pytest @@ -24,14 +21,17 @@ from modelopt.torch.quantization.linear_attention import LinearAttentionConfig, recurrent_decode -@pytest.fixture(scope="module", params=[False, True], ids=["gdn-token", "kda-replay"]) +@pytest.fixture( + scope="module", + params=[(False, 1), (False, 4), (True, 4)], + ids=["gdn-token", "gdn-replay", "kda-replay"], +) def compiled_replay(request): native = pytest.importorskip( "vllm.model_executor.layers.fla.ops.fused_recurrent_replayssm", exc_type=ModuleNotFoundError, ) - channel = request.param - window = 4 if channel else 1 + channel, window = request.param torch.manual_seed(71) q, k = [ torch.nn.functional.normalize(torch.randn(5, 1, 64, device="cuda"), dim=-1).bfloat16() diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py index 616551b4973..67f1e280276 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py @@ -21,8 +21,8 @@ from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.linear_attention import ( LinearAttentionConfig, - matmul_gdn, - matmul_kda, + gdn_state_qat, + kda_state_qat, ) from modelopt.torch.quantization.nn import TensorQuantizer @@ -50,7 +50,7 @@ def compiled_serving_case(request): ) ).cuda() forward = partial( - matmul_kda if kda else matmul_gdn, + kda_state_qat if kda else gdn_state_qat, *args, policy=policy, state_quantizer=quantizer, diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index 9bfcc33143e..56ef1aefe7b 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -22,13 +22,11 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.quantization.config import QuantizeConfig -from modelopt.torch.quantization.linear_attention import ( - LinearAttentionConfig, - linear_attention_training_phase, -) +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins import gdn from modelopt.torch.quantization.plugins.gdn import GatedDeltaNetStateQuantMixin +from modelopt.torch.quantization.plugins.kda import KimiDeltaAttentionStateQuantMixin GDN_STATE_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"} @@ -75,6 +73,7 @@ def quant_cfg(): {"quantizer_name": "*", "enable": False}, {"quantizer_name": "*gdn_state_quantizer", "cfg": GDN_STATE_FP8_DYNAMIC}, ], + "linear_attention": [{"module_name": "*", "cfg": {"backend": "serving"}}], "algorithm": "max", } @@ -109,8 +108,6 @@ def test_dynamic_export_removes_linear_attention_attributes(): for name in ( "gdn_state_quantizer", "gdn_w_quantizer", - "replay_key_quantizer", - "replay_update_quantizer", "linear_attention_config", "_linear_attention_prefill_lengths", ): @@ -131,19 +128,29 @@ def test_disabled_state_quantizer_calls_original_kernel(): assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" -def test_state_qat_requires_explicit_serving_policy(monkeypatch): - model = mtq.quantize(TinyGatedDeltaNet(), {**quant_cfg(), "algorithm": None}) +@pytest.mark.parametrize("mixin", [GatedDeltaNetStateQuantMixin, KimiDeltaAttentionStateQuantMixin]) +def test_state_qat_requires_explicit_serving_policy(mixin): + model = mixin.convert(TinyGatedDeltaNet()) + model._linear_attn_state.set_from_attribute_config(GDN_STATE_FP8_DYNAMIC) + model._linear_attn_state.enable() + with pytest.raises(ValueError, match="requires backend='serving'"): + model.validate_linear_attention() + + +def test_training_phase_routes_prefill_lengths(monkeypatch): + cfg = {**quant_cfg(), "algorithm": None} + missing_policy = {key: value for key, value in cfg.items() if key != "linear_attention"} + with pytest.raises(ValueError, match="requires backend='serving'"): + mtq.quantize(TinyGatedDeltaNet(), missing_policy) + model = mtq.quantize(TinyGatedDeltaNet(), cfg) x = torch.randn(2, 8, 3, 4) - with pytest.raises(ValueError, match="Chunk-only state QAT is retired"): - model(x) calls = [] def forward(*args, **kwargs): calls.append(kwargs) return chunk_gated_delta_rule(*args) - monkeypatch.setattr(gdn, "matmul_gdn", forward) - model.linear_attention_config = LinearAttentionConfig(backend="serving") + monkeypatch.setattr(gdn, "gdn_state_qat", forward) with linear_attention_training_phase(model, [4, 4]): model(x) assert calls[-1]["prefill_lengths"] == (4, 4) @@ -184,8 +191,8 @@ def test_quantizer_roundtrip_and_hybrid_selection(tmp_path): num_bits=8, axis=None, block_sizes={-1: 32}, unsigned=False, narrow_range=True ) cfg["linear_attention"] = [ - {"module_name": "*", "cfg": {"state_block_v": 128}}, - {"module_name": "0", "cfg": {"state_block_v": 32}}, + {"module_name": "*", "cfg": {"backend": "serving", "state_block_v": 128}}, + {"module_name": "0", "cfg": {"backend": "serving", "state_block_v": 32}}, ] mtq.quantize(model, cfg) assert model[0].gdn_state_qdq_block_v == 32 @@ -194,21 +201,6 @@ def test_quantizer_roundtrip_and_hybrid_selection(tmp_path): expected = model[0].gdn_state_quantizer(sample) path = tmp_path / "gdn.pth" mto.save(model, path) - # Load policies saved before fixed, non-configurable fields were removed. - checkpoint = torch.load(path) - for _, mode_state in checkpoint["modelopt_state"]["modelopt_state_dict"]: - for policy in mode_state["metadata"]["linear_attention"].values(): - policy["schema_version"] = 2 - policy["state"] = { - "block_v": policy.pop("state_block_v"), - "mode": "chunk", - "quantize_initial": True, - } - policy["decode"] = None - policy.pop("precision") - policy.pop("replay_window") - policy["solve"] = {"method": "exact"} - torch.save(checkpoint, path) restored = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) mto.restore(restored, path) assert restored[0].linear_attention_config == model[0].linear_attention_config @@ -234,10 +226,12 @@ def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): cfg = deepcopy(quant_cfg()) cfg["algorithm"] = None - cfg["linear_attention"] = [{"module_name": "", "cfg": {"state_block_v": 32}}] + cfg["linear_attention"] = [ + {"module_name": "", "cfg": {"backend": "serving", "state_block_v": 32}} + ] mtq.quantize(model, cfg) assert model.gdn_state_qdq_block_v == 32 - cfg["linear_attention"].append({"module_name": "", "cfg": {}}) + cfg["linear_attention"].append({"module_name": "", "cfg": {"backend": "serving"}}) mtq.quantize(model, cfg) assert model.gdn_state_qdq_block_v == 64 assert model.gdn_state_quantizer.is_enabled @@ -274,20 +268,6 @@ def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): ) assert not model.gdn_state_quantizer.is_enabled assert not model.gdn_w_quantizer.is_enabled - assert not model.replay_key_quantizer.is_enabled - assert not model.replay_update_quantizer.is_enabled - - -@pytest.mark.parametrize("name", ["replay_key_quantizer", "replay_update_quantizer"]) -def test_enabled_legacy_replay_quantizer_is_rejected(name): - cfg = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None} - model = mtq.quantize(TinyGatedDeltaNet(), cfg) - getattr(model, name).enable() - with pytest.raises(ValueError, match=f"{name} is no longer supported"): - mto.restore_from_modelopt_state(TinyGatedDeltaNet(), mto.modelopt_state(model)) - cfg["quant_cfg"].append({"quantizer_name": f"*{name}", "enable": True}) - with pytest.raises(ValueError, match=f"{name} is no longer supported"): - mtq.quantize(TinyGatedDeltaNet(), cfg) def test_policy_rejects_unmatched_and_unimplemented_modes(): diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py index 5a5b54220fd..3c8058df60c 100644 --- a/tests/unit/torch/quantization/test_linear_attention_decode.py +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -34,9 +34,7 @@ def test_state_qdq_matches_tensor_quantizer(state_format): expected = torch.cat( [quantizer(tile.flatten(-2)).reshape_as(tile) for tile in value.split(16, -1)], -1 ) - encoded = _encode( - value, True, 16, state=True, state_format=state_format, state_quantizer=quantizer - ) + encoded = _encode(value, True, 16, state_format=state_format, state_quantizer=quantizer) torch.testing.assert_close(encoded.values, expected, rtol=0, atol=0) probe = torch.randn_like(value) (gradient,) = torch.autograd.grad((encoded.values * probe).sum(), value) @@ -45,47 +43,6 @@ def test_state_qdq_matches_tensor_quantizer(state_format): assert not encoded.scales.requires_grad -@pytest.mark.parametrize( - "config", - [ - {"backend": "reference", "decode": {}}, - {"backend": "matmul", "decode": {}}, - {"backend": "serving", "decode": {"precision": "full"}}, - ], -) -def test_reference_training_is_rejected(config): - with pytest.raises(ValueError, match="Reference training is retired"): - LinearAttentionConfig(**config) - - -@pytest.mark.parametrize(("precision", "window"), [("vllm_0_15", 1), ("replayssm", 4)]) -def test_native_legacy_policy_preserves_precision(precision, window, tmp_path): - decode = {"precision": precision, "readout": "working"} - if precision == "replayssm": - decode.update(state_codec="int8_hadamard32", mode="replay", replay={"window": window}) - policy = LinearAttentionConfig(schema_version=2, backend="matmul", decode=decode) - assert policy == LinearAttentionConfig( - backend="serving", precision=precision, replay_window=window - ) - assert "decode" not in policy.model_dump() - assert "state_codec" not in policy.model_dump() - # Reproduce the field layout saved by the old class, bypassing new validation. - legacy = LinearAttentionConfig() - legacy.__dict__.clear() - legacy.__dict__.update(schema_version=2, backend="matmul", decode=decode) - path = tmp_path / "legacy-policy.pt" - torch.save(legacy, path) - assert torch.load(path, weights_only=True) == policy - - -@pytest.mark.parametrize( - "settings", [{"readout": "stored"}, {"prefill_state_qdq": True}, {"decay_log_step": 0.02}] -) -def test_serving_precision_requires_native_schedule(settings): - with pytest.raises(ValueError, match="Native state QAT requires"): - LinearAttentionConfig(backend="serving", decode=settings) - - @pytest.mark.parametrize( "settings", [ @@ -93,7 +50,6 @@ def test_serving_precision_requires_native_schedule(settings): {"precision": "replayssm", "replay_window": 0}, {"precision": "replayssm", "replay_window": 65}, {"precision": "replayssm", "state_block_v": 16}, - {"decode": {}, "precision": "replayssm"}, ], ) def test_unified_policy_rejects_incompatible_settings(settings): diff --git a/tests/unit/torch/quantization/test_linear_attention_hadamard.py b/tests/unit/torch/quantization/test_linear_attention_hadamard.py index 5f51d6344c3..41acfc74aa3 100644 --- a/tests/unit/torch/quantization/test_linear_attention_hadamard.py +++ b/tests/unit/torch/quantization/test_linear_attention_hadamard.py @@ -15,7 +15,7 @@ import torch -from modelopt.torch.quantization.linear_attention.decode import _encode, _hadamard32 +from modelopt.torch.quantization.linear_attention.decode import _hadamard32 def _rotate(value): @@ -24,32 +24,12 @@ def _rotate(value): 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 test_hadamard_codec_matches_dense_oracle_and_identity_ste(): +def test_hadamard_matches_dense_oracle_and_backward(): torch.manual_seed(321) state = torch.randn(2, 4, 64, dtype=torch.float64, requires_grad=True) - encoded = _encode( - _hadamard32(state), - True, - 64, - state=True, - state_format="int8", - state_codec="int8_hadamard32", - ) - actual = _hadamard32(encoded.values) - expected = _rotate(_qdq(_rotate(state))) + actual = _hadamard32(state) + expected = _rotate(state) torch.testing.assert_close(actual, expected, rtol=1e-10, atol=1e-10) - assert encoded.scales.shape == (2, 4, 2) - assert encoded.scales.dtype == torch.float16 probe = torch.randn_like(state) (gradient,) = torch.autograd.grad((actual * probe).sum(), state) - torch.testing.assert_close(gradient, probe, rtol=1e-10, atol=1e-10) + torch.testing.assert_close(gradient, _rotate(probe), rtol=1e-10, atol=1e-10) diff --git a/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py b/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py index bf8238cb85c..d1512547643 100644 --- a/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py @@ -119,16 +119,3 @@ def test_grouped_quantizer_preserves_nested_sequential_state_dict_layout(): ) assert list(grouped.state_dict()) == ["0.0._amax", "0.1._amax", "1._amax"] - - -def test_dynamic_block_quantization_accepts_changing_shapes(): - cfg = QuantizerAttributeConfig( - num_bits=8, type="dynamic", block_sizes={-1: 32}, pass_through_bwd=True - ) - quantizer = TensorQuantizer(cfg) - for shape in ((2, 3, 17), (1, 2, 4, 64), (3, 2, 33)): - value = torch.randn(shape, requires_grad=True) - output = quantizer(value) - torch.testing.assert_close(output, TensorQuantizer(cfg)(value), rtol=0, atol=0) - output.sum().backward() - torch.testing.assert_close(value.grad, torch.ones_like(value), rtol=0, atol=0) From 406343da5b3f636af2f1686e7f60877cef436d62 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 7 Oct 2026 12:42:56 -0700 Subject: [PATCH 15/27] Refactor linear attention policy and checkpoint lifecycle Signed-off-by: Kai Xu --- .../opt/plugins/mcore_dist_checkpointing.py | 1 + modelopt/torch/quantization/conversion.py | 78 ++++++----------- modelopt/torch/quantization/model_quant.py | 8 +- .../quantization/plugins/linear_attention.py | 61 ++++++++++++++ .../torch/quantization/plugins/megatron.py | 84 +++++++++++-------- modelopt_recipes/ptq.md | 2 +- .../plugins/test_megatron_gated_delta_net.py | 7 ++ .../quantization/plugins/test_megatron_kda.py | 6 ++ .../torch/quantization/plugins/test_gdn.py | 35 +++++++- 9 files changed, 191 insertions(+), 91 deletions(-) diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py index a4796162f6d..f1bd15154fb 100644 --- a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -60,6 +60,7 @@ def remove_per_module_state( _ = metadata.pop("subnet_config", None) _ = metadata.pop("real_quantizer_state", None) _ = metadata.pop("q_tensor_state", None) + _ = metadata.pop("linear_attention", None) else: config["metadata"] = {} diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index db194a0f465..405f9b8876b 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -28,7 +28,6 @@ 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, @@ -37,7 +36,6 @@ _QuantizeExportConfig, normalize_quant_cfg_list, ) -from .linear_attention.config import LinearAttentionConfig from .nn import ( QuantModule, QuantModuleRegistry, @@ -64,12 +62,16 @@ def convert_to_quantized_model(model: ModelLikeModule, config: QuantizeConfig) -> ConvertReturnType: """Convert the model to a quantized one as per `config`.""" + # Defer shared plugin imports to avoid the plugin/conversion initialization cycle. + from .plugins.linear_attention import _apply_linear_attention_policy, _validate_linear_attention + # initialize the true module if necessary model = model.init_modellike() if isinstance(model, ModelLikeModule) else model replace_quant_module(model, version=ModeloptStateManager(model).state_version) set_quantizer_by_cfg(model, config.get("quant_cfg", [])) _apply_linear_attention_policy(model, config) + _validate_linear_attention(model) metadata = {} update_quantize_metadata(model, config, metadata) @@ -97,8 +99,15 @@ def restore_quantized_model( model: ModelLikeModule, config: QuantizeConfig, metadata: MetadataDict ) -> nn.Module: """Insert quantizers to the model and restore the quantizer states from the given state dict.""" + # Defer the shared plugin import to avoid the plugin/conversion initialization cycle. + from .plugins.linear_attention import _apply_linear_attention_policy + # initialize the true module if necessary - convert_to_quantized_model(model, config) + model = model.init_modellike() if isinstance(model, ModelLikeModule) else model + replace_quant_module(model, version=ModeloptStateManager(model).state_version) + set_quantizer_by_cfg(model, config.get("quant_cfg", [])) + # Rebuild from the recipe without validating against partially restored state. + _apply_linear_attention_policy(model, config) return restore_quantizer_state(model, config, metadata) @@ -139,7 +148,13 @@ 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")) + # Defer shared plugin imports to avoid the plugin/conversion initialization cycle. + from .plugins.linear_attention import ( + _restore_legacy_linear_attention_quantizers, + _restore_linear_attention_policy, + ) + + _restore_linear_attention_policy(model, 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 @@ -148,13 +163,7 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: 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() + _restore_legacy_linear_attention_quantizers(model, quantizer_state_dict) unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys() extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys() @@ -207,12 +216,14 @@ 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() - } + # Defer the shared plugin import to avoid the plugin/conversion initialization cycle. + from .plugins.linear_attention import _linear_attention_state + + policies = _linear_attention_state(model) if policies: metadata["linear_attention"] = policies + else: + metadata.pop("linear_attention", None) metadata["quantizer_state"] = quantizer_state(model) if shared_state_metadata := SharedWeightGlobalAmaxState.metadata(model): metadata["shared_quant_states"] = shared_state_metadata @@ -220,43 +231,6 @@ 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/model_quant.py b/modelopt/torch/quantization/model_quant.py index c488b7f8fab..c1430fd39fc 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -33,7 +33,6 @@ 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, @@ -401,9 +400,16 @@ def forward_loop(model) -> None: if not is_quantized(model): model = apply_mode(model, mode=[("quantize", dict(config))], registry=QuantizeModeRegistry) else: + # Defer shared plugin imports to avoid the plugin/conversion initialization cycle. + from .plugins.linear_attention import ( + _apply_linear_attention_policy, + _validate_linear_attention, + ) + # 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) + _validate_linear_attention(model) # Fail before calibration rather than after exporting an unquantized checkpoint. _check_weight_quantization_took_effect(model, quantize_config) _check_plugin_quantization_took_effect(model, quantize_config) diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py index 90b969435c2..baaddeb1fd0 100644 --- a/modelopt/torch/quantization/plugins/linear_attention.py +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -15,6 +15,11 @@ """Shared module policy and checkpoint support for linear-attention QAT.""" +import fnmatch + +from modelopt.torch.opt.conversion import ApplyModeError +from modelopt.torch.utils import get_unwrapped_name + from ..config import QuantizerAttributeConfig from ..linear_attention.config import LinearAttentionConfig from ..linear_attention.utils import state_quantizer_config @@ -23,6 +28,62 @@ __all__ = [] +def _linear_attention_modules(model): + 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): + """Assign complete policies in rule order, before compatibility validation.""" + modules = _linear_attention_modules(model) + # getattr also handles pickled configs that predate the policy field. + 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) + + +def _validate_linear_attention(model): + """Validate configured modules after both policy and quantizers are assigned.""" + for module in _linear_attention_modules(model).values(): + module.validate_linear_attention() + + +def _linear_attention_state(model): + return { + name: module.linear_attention_config.model_dump() + for name, module in _linear_attention_modules(model).items() + } + + +def _restore_linear_attention_policy(model, saved_policies): + """Restore saved policies; modelopt_post_restore validates the completed state.""" + if saved_policies is None: + return + modules = _linear_attention_modules(model) + 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) + + +def _restore_legacy_linear_attention_quantizers(model, quantizer_state): + """Keep newly introduced, disabled handles loadable from older checkpoints.""" + 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 and not quantizer.is_enabled: + quantizer_state[key] = quantizer.get_modelopt_state() + + class _LinearAttentionQuantMixin(QuantModule): linear_attention_quantizer_names: tuple[str, ...] = () diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 033e553a109..63c6acc029e 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -58,6 +58,7 @@ from ..algorithms import AutoQuantizeGradientSearcher from ..conversion import maybe_promote_nvfp4_static_quantizer +from ..linear_attention.config import LinearAttentionConfig from ..nn import ( GroupedQuantizer, QuantModule, @@ -219,6 +220,8 @@ def quant_module_get_extra_state(self) -> dict: quantizer_state[name] = module.get_modelopt_state() extra_state["modelopt_quantizer_state"] = quantizer_state + if isinstance(self, _LinearAttentionQuantMixin): + extra_state["modelopt_linear_attention_state"] = self.linear_attention_config.model_dump() # Handle real_quantizer_state and q_tensor_state extra_state.update(real_quant_module_get_extra_state(self)) @@ -297,6 +300,12 @@ def quant_module_set_extra_state(self, state: Any): if state is None or not self.allow_post_restore: return + if isinstance(self, _LinearAttentionQuantMixin) and "modelopt_linear_attention_state" in state: + # Restore the policy before quantizer restoration invokes modelopt_post_restore. + self.linear_attention_config = LinearAttentionConfig( + **state["modelopt_linear_attention_state"] + ) + quantizer_state = state.get("modelopt_quantizer_state", None) if quantizer_state is not None: @@ -1118,6 +1127,14 @@ def set_extra_state(self, state): class _MegatronLinearAttentionMixin(_LinearAttentionQuantMixin): """Wrap the kernel call in both direct and recomputed Megatron forwards.""" + # PyTorch requires class-level overrides to save/load ``_extra_state``; + # instance-level ModelOpt callbacks alone are not sufficient for GDN or KDA. + def get_extra_state(self): + return quant_module_get_extra_state(self) + + def set_extra_state(self, state): + quant_module_set_extra_state(self, state) + def _setup(self): super()._setup() self._register_temp_attribute("_linear_attention_replay_gate_inputs", None) @@ -1146,40 +1163,11 @@ def _prepare_input_for_gated_delta_rule(self, *args, **kwargs): finally: self.use_qk_l2norm = normalize - def _compute_gates(self, a_log, dt_bias, batch, seq_len, *gate_feats): - gate, inputs = super()._compute_gates(a_log, dt_bias, batch, seq_len, *gate_feats) - if self._serving_arithmetic and hasattr(self, "gdn_state_quantizer"): - # Import the optional vLLM backend only for the native precision profile. - from ...kernels.quantization.linear_attention.serving.forward import fused_gdn_gating - from ..linear_attention.utils import forward_value - - raw_beta, raw_gate = gate_feats - if self.linear_attention_config.precision == "replayssm": - self._linear_attention_replay_gate_inputs = (raw_gate, raw_beta, a_log, dt_bias) - with torch.no_grad(): - native_gate, native_beta = fused_gdn_gating( - a_log, - raw_gate.reshape(-1, raw_gate.shape[-1]).contiguous(), - raw_beta.reshape(-1, raw_beta.shape[-1]).contiguous(), - dt_bias, - ) - gate = forward_value(gate, native_gate.reshape_as(gate)) - inputs["beta"] = forward_value( - inputs["beta"], native_beta.reshape_as(inputs["beta"]) - ).to(raw_beta.dtype) - return gate, inputs - @contextmanager def _quantized_linear_attention_kernel(self): kernel = self.gated_delta_rule previous_gates = self._linear_attention_replay_gate_inputs if self._serving_arithmetic: - if not hasattr(super(), "_prepare_input_for_gated_delta_rule") or getattr( - self, "gdn_pre_gated_delta_rule_fusion", False - ): - raise NotImplementedError( - "Serving arithmetic requires Megatron's unfused input-preparation hook" - ) def run_kernel(*args, **kwargs): kwargs["use_qk_l2norm_in_kernel"] = self.use_qk_l2norm @@ -1218,6 +1206,13 @@ 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/KDA QAT does not support Megatron context parallelism.") + if self._serving_arithmetic and ( + not hasattr(super(), "_prepare_input_for_gated_delta_rule") + or getattr(self, "gdn_pre_gated_delta_rule_fusion", False) + ): + raise NotImplementedError( + "Serving arithmetic requires Megatron's unfused input-preparation hook" + ) if HAS_GDN: @@ -1230,13 +1225,30 @@ class _QuantGatedDeltaNet(_MegatronLinearAttentionMixin, GatedDeltaNetStateQuant GatedDeltaNetStateQuantMixin._state_quantized_chunk_gated_delta_rule ) - # Class-level overrides so torch routes the quantizer state through ``_extra_state`` - # (GatedDeltaNet has none); see _QuantDSAttention. - def get_extra_state(self): - return quant_module_get_extra_state(self) - - def set_extra_state(self, state): - quant_module_set_extra_state(self, state) + def _compute_gates(self, a_log, dt_bias, batch, seq_len, *gate_feats): + gate, inputs = super()._compute_gates(a_log, dt_bias, batch, seq_len, *gate_feats) + if self._serving_arithmetic: + # Import the optional vLLM backend only for the native precision profile. + from ...kernels.quantization.linear_attention.serving.forward import ( + fused_gdn_gating, + ) + from ..linear_attention.utils import forward_value + + raw_beta, raw_gate = gate_feats + if self.linear_attention_config.precision == "replayssm": + self._linear_attention_replay_gate_inputs = (raw_gate, raw_beta, a_log, dt_bias) + with torch.no_grad(): + native_gate, native_beta = fused_gdn_gating( + a_log, + raw_gate.reshape(-1, raw_gate.shape[-1]).contiguous(), + raw_beta.reshape(-1, raw_beta.shape[-1]).contiguous(), + dt_bias, + ) + gate = forward_value(gate, native_gate.reshape_as(gate)) + inputs["beta"] = forward_value( + inputs["beta"], native_beta.reshape_as(inputs["beta"]) + ).to(raw_beta.dtype) + return gate, inputs if HAS_KDA: diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index e2aa898890f..13d1a925844 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 33 general/ptq/ recipes (click to expand) +All 34 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index 72e0229db24..524cc7618fe 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -190,6 +190,13 @@ def _test_gdn_qat_helper(rank, size, checkpoint_path): optimizer.step() assert not torch.equal(restored_gdn[0].in_proj.weight, before) + fused_model = _make_model(size) + for module in fused_model.modules(): + if isinstance(module, GatedDeltaNet): + module.gdn_pre_gated_delta_rule_fusion = True + with pytest.raises(NotImplementedError, match="unfused input-preparation hook"): + mtq.quantize(fused_model, _gdn_config()) + def _compile_gdn_qat_kernels(rank, size): initialize_for_megatron( diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py index 471b09cc81d..12cecf6a30b 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py @@ -120,6 +120,8 @@ def _test_kda(rank, size, cfg, checkpoint_path): assert not torch.equal(forward(model), baseline) assert model.gated_delta_rule is kernel + # The sharded checkpoint must restore this per-module change, not recipe defaults. + model.linear_attention_config.state_block_v = 32 policy = model.linear_attention_config mtq.disable_quantizer(model, "*") model.linear_attention_config = LinearAttentionConfig() @@ -130,11 +132,15 @@ def _test_kda(rank, size, cfg, checkpoint_path): with torch.no_grad(): expected = forward(model) + # A change made after conversion must survive beyond the saved recipe defaults. + mtq.disable_quantizer(model, "*kda_state_quantizer") save_distributed_checkpoint(checkpoint_path, model) save_sharded_modelopt_state([model], checkpoint_path) restored = _layer() restore_sharded_modelopt_state([restored], checkpoint_path) load_distributed_checkpoint(checkpoint_path, restored) + assert not restored.kda_state_quantizer.is_enabled + mtq.enable_quantizer(restored, "*kda_state_quantizer") # Megatron recomputes the core during backward, after forward's kernel wrapper exits. with linear_attention_training_phase(restored, [31]): actual = forward(restored) diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index 56ef1aefe7b..16ea8f5762a 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -22,7 +22,10 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.quantization.config import QuantizeConfig -from modelopt.torch.quantization.linear_attention import linear_attention_training_phase +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + linear_attention_training_phase, +) from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins import gdn from modelopt.torch.quantization.plugins.gdn import GatedDeltaNetStateQuantMixin @@ -217,6 +220,36 @@ def test_quantizer_roundtrip_and_hybrid_selection(tmp_path): torch.testing.assert_close(restored[0].gdn_state_quantizer(sample), expected) +@pytest.mark.parametrize("reverse", [False, True], ids=["fp8-to-replay", "replay-to-fp8"]) +def test_restore_changed_policy_and_quantizer_together(tmp_path, reverse): + states = [ + (GDN_STATE_FP8_DYNAMIC, LinearAttentionConfig(backend="serving")), + ( + {"num_bits": 8, "axis": (0, 1), "type": "dynamic", "narrow_range": True}, + LinearAttentionConfig(backend="serving", precision="replayssm"), + ), + ] + if reverse: + states.reverse() + (original_quantizer, original_policy), (saved_quantizer, saved_policy) = states + cfg = quant_cfg() + cfg["algorithm"] = None + cfg["quant_cfg"][-1]["cfg"] = original_quantizer + cfg["linear_attention"][0]["cfg"] = original_policy.model_dump() + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + model.gdn_state_quantizer.set_from_attribute_config(saved_quantizer) + model.linear_attention_config = saved_policy + model.validate_linear_attention() + + path = tmp_path / "changed-policy.pth" + mto.save(model, path) + restored = mto.restore(TinyGatedDeltaNet(), path) + assert restored.linear_attention_config == saved_policy + assert restored.gdn_state_quantizer.num_bits == saved_quantizer["num_bits"] + assert restored.gdn_state_quantizer.is_enabled + restored.validate_linear_attention() + + def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): cfg = quant_cfg() cfg["algorithm"] = None From 68759da7ab9a51b1627c7bce1fd640549dd0fae4 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 7 Oct 2026 19:21:39 -0700 Subject: [PATCH 16/27] Support native vLLM kernel revisions in linear attention QAT Adapt native imports, state layouts, recurrent indexing, and KDA gate arithmetic while keeping ModelOpt state QDQ in canonical layout. Use the vllm profile name and retain vllm_0_15 as a legacy alias. Keep lazy gate discovery outside compiled graphs. Validate policy matches across distributed stages and correct the optional dependency message. Verified 94 CPU tests, native GDN/KDA GPU checks on vLLM 0.15.1, 0.20.0, and 0.30.0, ReplaySSM cases, Megatron backward/checkpoint restore, and pre-commit hooks. Full serving-engine and model-quality qualification remain separate. Signed-off-by: Kai Xu --- CHANGELOG.rst | 2 +- .../linear_attention/serving/__init__.py | 7 +- .../linear_attention/serving/_compat.py | 40 +++++++++++ .../linear_attention/serving/chunk_delta_h.py | 19 +++-- .../linear_attention/serving/forward.py | 72 +++++++++++++------ .../linear_attention/_vllm_autograd.py | 2 +- .../quantization/linear_attention/config.py | 7 +- .../quantization/linear_attention/decode.py | 4 +- .../quantization/linear_attention/training.py | 5 +- .../quantization/plugins/linear_attention.py | 10 ++- .../ptq/units/gdn_state_fp8_dynamic.yaml | 2 +- ..._attention_state_int8_block32_dynamic.yaml | 2 +- modelopt_recipes/ptq.md | 2 +- .../plugins/test_megatron_gated_delta_net.py | 2 +- .../quantization/plugins/test_megatron_kda.py | 4 +- .../test_linear_attention_training.py | 56 +++++++++++++-- .../torch/quantization/plugins/test_gdn.py | 33 ++++++++- 17 files changed, 216 insertions(+), 53 deletions(-) create mode 100644 modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a28a3597900..6440b085116 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,7 +16,7 @@ Changelog - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. -- Add experimental serving-aligned GDN/KDA state QAT with explicit prefill/decode boundaries; migrate chunk-only state recipes to a serving policy and supply prefix lengths through the training phase context. Plain state QDQ requires ``vllm==0.15.1``; INT8 Hadamard and ReplaySSM require the compatible quantized-ReplaySSM serving fork. +- Add experimental serving-aligned GDN/KDA state QAT with explicit prefill/decode boundaries; migrate chunk-only state recipes to a serving policy and supply prefix lengths through the training phase context. Plain state QDQ uses the installed vLLM's native kernels via ``precision="vllm"`` (the legacy ``vllm_0_15`` spelling remains accepted); INT8 Hadamard and ReplaySSM require the compatible quantized-ReplaySSM serving fork. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before a separate KV-cache AutoQuantize stage, with independent resumable checkpoints for the weight and diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py index 6ee02050242..1156582d24e 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py @@ -13,12 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Optional vLLM forward kernels for the ``vllm_0_15`` precision profile.""" +"""Optional vLLM forward kernels for linear-attention serving precision profiles.""" try: import vllm except ImportError as error: raise ImportError( - "precision='vllm_0_15' requires the optional vLLM dependency; " - "install the dependencies in the linear-attention QAT example's requirements-vllm.txt." + "Linear-attention serving profiles require the optional vLLM dependency: " + "public vLLM for 'vllm' (legacy spelling 'vllm_0_15'), or the compatible quantized-ReplaySSM " + "vLLM fork for 'replayssm'." ) from error diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py b/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py new file mode 100644 index 00000000000..81ba6dd1d88 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py @@ -0,0 +1,40 @@ +# 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. + +"""Keep native vLLM imports and state layout changes at the kernel boundary.""" + +from importlib import import_module +from importlib.util import find_spec + +import vllm + +__all__ = [] + +# vLLM 0.16 transposed the recurrent cache; ModelOpt keeps [H,K,V] for QDQ/autograd. +STATE_V_FIRST = vllm.__version_tuple__[:2] >= (0, 16) +_FLA_OPS = ( + "vllm.third_party.flash_linear_attention.ops" + if find_spec("vllm.third_party.flash_linear_attention") is not None + else "vllm.model_executor.layers.fla.ops" +) + + +def fla_module(name): + return import_module(f"{_FLA_OPS}.{name}") + + +def state_layout(state): + """Convert between ModelOpt and native layout; transposing is its own inverse.""" + return state.transpose(-1, -2).contiguous() if STATE_V_FIRST else state.contiguous() diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py b/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py index 95ca3882dbc..b8759c9541c 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py @@ -22,20 +22,28 @@ """Save FP32 training intermediates using vLLM's unchanged chunk-state kernel.""" +from inspect import signature + import torch import triton -from vllm.model_executor.layers.fla.ops.chunk_delta_h import ( - chunk_gated_delta_rule_fwd_kernel_h_blockdim64, + +from ._compat import STATE_V_FIRST, fla_module + +_chunk = fla_module("chunk_delta_h") +chunk_gated_delta_rule_fwd_kernel_h_blockdim64 = ( + _chunk.chunk_gated_delta_rule_fwd_kernel_h_blockdim64 ) -from vllm.model_executor.layers.fla.ops.index import prepare_chunk_offsets +prepare_chunk_offsets = fla_module("index").prepare_chunk_offsets +_HAS_EXP2 = "use_exp2" in signature(_chunk.chunk_gated_delta_rule_fwd_h).parameters -def chunk_state(k, w, u, g, gk, initial_state, cu_seqlens): +def chunk_state(k, w, u, g, gk, initial_state, cu_seqlens, use_exp2=False): """Return chunk-start states, residual values, and final state for one sequence.""" _, length, heads, key_dim = k.shape value_dim = u.shape[-1] # The Torch adjoint needs unrounded values. Output kernels receive BF16 casts. - h = k.new_empty(1, triton.cdiv(length, 64), heads, key_dim, value_dim, dtype=torch.float32) + state_shape = (value_dim, key_dim) if STATE_V_FIRST else (key_dim, value_dim) + h = k.new_empty(1, triton.cdiv(length, 64), heads, *state_shape, dtype=torch.float32) updated = torch.empty_like(u, dtype=torch.float32) final = torch.empty_like(initial_state, dtype=torch.float32) chunk_gated_delta_rule_fwd_kernel_h_blockdim64[ @@ -58,5 +66,6 @@ def chunk_state(k, w, u, g, gk, initial_state, cu_seqlens): K=key_dim, V=value_dim, BT=64, + **({"USE_EXP2": use_exp2} if _HAS_EXP2 else {}), ) return h, updated, final diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py index 3d670566bde..1d0866ca21f 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py @@ -15,24 +15,29 @@ """vLLM forward primitives with FP32 saved values; autograd lives in quantization.""" +from functools import cache +from importlib import import_module + import torch -from vllm.model_executor.layers.fla.ops.chunk_o import chunk_fwd_o -from vllm.model_executor.layers.fla.ops.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd -from vllm.model_executor.layers.fla.ops.cumsum import chunk_local_cumsum -from vllm.model_executor.layers.fla.ops.fused_recurrent import fused_recurrent_gated_delta_rule -from vllm.model_executor.layers.fla.ops.kda import ( - chunk_gla_fwd_o_gk, - chunk_kda_scaled_dot_kkt_fwd, - fused_recurrent_kda, -) -from vllm.model_executor.layers.fla.ops.kda import fused_kda_gate as fused_kda_gate -from vllm.model_executor.layers.fla.ops.kda import recompute_w_u_fwd as kda_wu -from vllm.model_executor.layers.fla.ops.l2norm import l2norm_fwd -from vllm.model_executor.layers.fla.ops.solve_tril import solve_tril -from vllm.model_executor.layers.fla.ops.wy_fast import recompute_w_u_fwd as gdn_wu +from ._compat import fla_module, state_layout from .chunk_delta_h import chunk_state +chunk_fwd_o = fla_module("chunk_o").chunk_fwd_o +chunk_scaled_dot_kkt_fwd = fla_module("chunk_scaled_dot_kkt").chunk_scaled_dot_kkt_fwd +chunk_local_cumsum = fla_module("cumsum").chunk_local_cumsum +fused_recurrent_gated_delta_rule = fla_module("fused_recurrent").fused_recurrent_gated_delta_rule +_kda = fla_module("kda") +chunk_gla_fwd_o_gk = _kda.chunk_gla_fwd_o_gk +chunk_kda_scaled_dot_kkt_fwd = _kda.chunk_kda_scaled_dot_kkt_fwd +fused_recurrent_kda = _kda.fused_recurrent_kda +fused_kda_gate = _kda.fused_kda_gate +kda_wu = _kda.recompute_w_u_fwd +l2norm_fwd = fla_module("l2norm").l2norm_fwd +solve_tril = fla_module("solve_tril").solve_tril +gdn_wu = fla_module("wy_fast").recompute_w_u_fwd +_KDA_GATE_SCALE = getattr(_kda, "RCP_LN2", 1.0) + def prefill(q, k, v, g, beta, state, scale, normalize=False): """Return one sequence's output, final state, and rounded forward intermediates.""" @@ -42,6 +47,10 @@ def prefill(q, k, v, g, beta, state, scale, normalize=False): q, k = l2norm_fwd(q), l2norm_fwd(k) channel = g.ndim == 4 gc = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu) + # Newer native KDA kernels evaluate exp2 of base-2 cumulative log gates. + natural_gc = gc + if channel: + gc = gc * _KDA_GATE_SCALE if channel: lower, scores = chunk_kda_scaled_dot_kkt_fwd(q, k, gc, beta, scale=scale, cu_seqlens=cu) else: @@ -62,8 +71,9 @@ def prefill(q, k, v, g, beta, state, scale, normalize=False): u=u, g=None if channel else gc, gk=gc if channel else None, - initial_state=state.unsqueeze(0).contiguous(), + initial_state=state_layout(state.unsqueeze(0)), cu_seqlens=cu, + use_exp2=channel and _KDA_GATE_SCALE != 1.0, ) assert updated is not None and final is not None if channel: @@ -85,17 +95,17 @@ def prefill(q, k, v, g, beta, state, scale, normalize=False): intermediates = { "q": q[0], "k": k[0], - "g": gc[0], + "g": natural_gc[0], "lower": lower[0], "inverse": inverse[0], "w": w[0], "u": u[0], - "h": h[0], + "h": state_layout(h[0]), "updated": updated[0], "kg": kg[0], "scores": None if scores is None else scores[0], } - return out[0], final[0], intermediates + return out[0], state_layout(final[0]), intermediates def step(q, k, v, g, beta, state, scale, normalize=False): @@ -103,18 +113,34 @@ def step(q, k, v, g, beta, state, scale, normalize=False): recurrent = fused_recurrent_kda if g.ndim == 2 else fused_recurrent_gated_delta_rule out, final = recurrent( *[x[None, None].contiguous() for x in (q, k, v, g, beta)], - initial_state=state.unsqueeze(0).contiguous(), + initial_state=state_layout(state.unsqueeze(0)), inplace_final_state=False, scale=scale, use_qk_l2norm_in_kernel=normalize, cu_seqlens=torch.tensor([0, 1], device=q.device, dtype=torch.int32), - ssm_state_indices=torch.tensor([0], device=q.device, dtype=torch.int32), + # Private dense training state needs no paged-cache index (0 is now reserved). + ssm_state_indices=None, ) - return out[0, 0], final[0] + return out[0, 0], state_layout(final[0]) def fused_gdn_gating(*args, **kwargs): """Load the optional vLLM model only when Megatron needs GDN gate preparation.""" - from vllm.model_executor.models.qwen3_next import fused_gdn_gating as native_gate + return _gdn_gate()(*args, **kwargs) + - return native_gate(*args, **kwargs) +# Megatron compiles gate preparation; module discovery must execute outside that graph. +@torch.compiler.disable +@cache +def _gdn_gate(): + for path in ( + "vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn", + "vllm.model_executor.layers.mamba.gdn_linear_attn", + "vllm.model_executor.models.qwen3_next", + ): + try: + return import_module(path).fused_gdn_gating + except ModuleNotFoundError as error: # noqa: PERF203 - cached, one-time import discovery + if error.name is None or not (path == error.name or path.startswith(error.name + ".")): + raise + raise ImportError("The installed vLLM does not provide fused_gdn_gating") diff --git a/modelopt/torch/quantization/linear_attention/_vllm_autograd.py b/modelopt/torch/quantization/linear_attention/_vllm_autograd.py index dfd67ea925a..e25971171b5 100644 --- a/modelopt/torch/quantization/linear_attention/_vllm_autograd.py +++ b/modelopt/torch/quantization/linear_attention/_vllm_autograd.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Differentiable adapter for pinned serving arithmetic. +"""Differentiable adapter for the installed serving runtime's arithmetic. The kernels supply forward values. Autograd differentiates the corresponding operations evaluated at those values; operand casts and QDQ use identity STE. diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index af8dd9dd1a9..a616a94dba3 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -28,7 +28,8 @@ class LinearAttentionConfig(ModeloptBaseConfig): """One execution policy shared by the chunked prefix and recurrent suffix. ``serving`` uses native forward arithmetic with a differentiable adjoint; - ``fla`` is the disabled/default path. ``vllm_0_15`` uses public vLLM arithmetic; + ``fla`` is the disabled/default path. ``vllm`` uses the installed vLLM's arithmetic; + ``vllm_0_15`` remains an accepted legacy spelling for existing checkpoints. ``replayssm`` uses the serving fork's INT8/Hadamard checkpoint and ring kernels. ``replay_window=1`` refreshes the state every token; larger windows enable replay. Fresh prefill has no internal state QDQ; incoming continuation state and @@ -39,7 +40,7 @@ class LinearAttentionConfig(ModeloptBaseConfig): """ backend: Literal["fla", "serving"] = ModeloptField(default="fla") - precision: Literal["vllm_0_15", "replayssm"] = ModeloptField(default="vllm_0_15") + precision: Literal["vllm", "vllm_0_15", "replayssm"] = ModeloptField(default="vllm") replay_window: int = Field(default=1, ge=1, le=64, strict=True) state_block_v: Literal[16, 32, 64, 128] = ModeloptField(default=64) @@ -64,7 +65,7 @@ class LinearAttentionPolicyEntry(ModeloptBaseConfig): """Assign a complete policy to supported modules matching ``module_name``. Rules apply in order: the last match wins, without merging fields. - A rule must match at least one supported linear-attention module. + A rule must match at least one supported module across distributed stages. """ module_name: str = Field(...) diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py index 9efa231f9eb..622a825b900 100644 --- a/modelopt/torch/quantization/linear_attention/decode.py +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -75,7 +75,7 @@ class LinearAttentionState: started: bool # Initial-state encoding has run; resuming must not repeat it. signature: str # Execution/quantizer settings used to check compatible reuse. value_basis: str = "identity" # Anchor and updates may use the Hadamard value basis. - precision: str = "vllm_0_15" # Selects the native profile's reconstruction arithmetic. + precision: str = "vllm" # Selects the native profile's reconstruction arithmetic. @property def cursor(self): @@ -315,7 +315,7 @@ def recurrent_decode( if config.backend != "serving": raise ValueError("State QAT requires backend='serving'") block_v = config.state_block_v - serving = config.precision == "vllm_0_15" + serving = config.precision in ("vllm", "vllm_0_15") native_replay = config.precision == "replayssm" if q.device.type != "cuda": raise ValueError("Serving arithmetic requires CUDA") diff --git a/modelopt/torch/quantization/linear_attention/training.py b/modelopt/torch/quantization/linear_attention/training.py index af1541e8187..ed8401d5107 100644 --- a/modelopt/torch/quantization/linear_attention/training.py +++ b/modelopt/torch/quantization/linear_attention/training.py @@ -20,6 +20,8 @@ import torch +import modelopt.torch.utils.distributed as dist + from .decode import recurrent_decode from .utils import _resolve_state_quantizer, _state_qdq, forward_value @@ -52,7 +54,8 @@ def linear_attention_training_phase(model, prefill_lengths): if isinstance(m, _LinearAttentionQuantMixin) and m.linear_attention_config.backend == "serving" ] - if not layers: + # A pipeline stage may have no selected layers; conversion validates global matches. + if not layers and dist.size() == 1: 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: diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py index baaddeb1fd0..d5945b57bab 100644 --- a/modelopt/torch/quantization/plugins/linear_attention.py +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -17,6 +17,9 @@ import fnmatch +import torch + +import modelopt.torch.utils.distributed as dist from modelopt.torch.opt.conversion import ApplyModeError from modelopt.torch.utils import get_unwrapped_name @@ -42,7 +45,12 @@ def _apply_linear_attention_policy(model, config): # getattr also handles pickled configs that predate the policy field. for entry in getattr(config, "linear_attention", []): matches = [name for name in modules if fnmatch.fnmatch(name, entry.module_name)] - if not matches: + matched_on_rank = [bool(matches)] + if dist.size() > 1: + # Shared recipes can select layers owned by another pipeline stage. + matched_on_rank = [False] * dist.size() + torch.distributed.all_gather_object(matched_on_rank, bool(matches)) + if not any(matched_on_rank): raise ValueError( f"linear_attention rule {entry.module_name!r} matches no supported modules" ) diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml index ba02ff508d2..ac36caab0f4 100644 --- a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -19,7 +19,7 @@ # Set BV through linear_attention[].cfg.state_block_v (default: 64). # This unit enables only the quantizer. Select backend='serving' in the execution # policy and provide prefill lengths through linear_attention_training_phase. -# With precision='vllm_0_15', state QDQ applies at the handoff to a nonempty decode +# With precision='vllm', state QDQ applies at the handoff to a nonempty decode # suffix and at each suffix token's cache write; fresh prefill has no internal QDQ. # The serving profile requires public vLLM's GDN kernels. # See modelopt.torch.quantization.plugins.gdn. diff --git a/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml index 8bff5421d9e..540f48a978a 100644 --- a/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml +++ b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml @@ -33,4 +33,4 @@ quantize: - module_name: '*' cfg: backend: serving - precision: vllm_0_15 + precision: vllm diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 13d1a925844..25a7c9df0c9 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -69,7 +69,7 @@ supported combinations.
-The block32 linear-attention recipe selects `precision="vllm_0_15"` +The block32 linear-attention recipe selects `precision="vllm"` and working-state readout. It applies fake QDQ with floating-point state storage. --- diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index 524cc7618fe..2092e8b9da0 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -33,7 +33,7 @@ pytest.importorskip("fla") # Megatron-Core GatedDeltaNet needs FLA for its baseline kernels. GatedDeltaNet = pytest.importorskip("megatron.core.ssm.gated_delta_net").GatedDeltaNet -pytest.importorskip("vllm.model_executor.layers.fla.ops.kda", exc_type=ModuleNotFoundError) +pytest.importorskip("vllm") from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py index 12cecf6a30b..79c4b89e56c 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py @@ -39,7 +39,7 @@ ) KimiDeltaAttention = pytest.importorskip("megatron.core.ssm.gated_delta_net.kda").KimiDeltaAttention -pytest.importorskip("vllm.model_executor.layers.fla.ops.kda", exc_type=ModuleNotFoundError) +pytest.importorskip("vllm") def _layer(): @@ -113,7 +113,7 @@ def _test_kda(rank, size, cfg, checkpoint_path): model, hidden, forward, baseline = _case(cfg) assert model.kda_state_quantizer.is_enabled assert model.kda_state_quantizer.num_bits == 8 - assert model.linear_attention_config.precision == "vllm_0_15" + assert model.linear_attention_config.precision == "vllm" assert model.kda_state_quantizer.block_sizes == {-1: 32} kernel = model.gated_delta_rule with torch.no_grad(): diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py index 67f1e280276..785a7b530d4 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py @@ -26,20 +26,64 @@ ) from modelopt.torch.quantization.nn import TensorQuantizer +vllm = pytest.importorskip("vllm") + +from modelopt.torch.kernels.quantization.linear_attention.serving._compat import fla_module + + +@torch.no_grad() +def _native_forward(args, quantizer, kda): + """Execute native prefill/decode directly, with QDQ in ModelOpt's [K,V] basis.""" + prefill = fla_module("kda").chunk_kda if kda else fla_module("chunk").chunk_gated_delta_rule + step = ( + fla_module("kda").fused_recurrent_kda + if kda + else fla_module("fused_recurrent").fused_recurrent_gated_delta_rule + ) + value_first = vllm.__version_tuple__[:2] >= (0, 16) + + def layout(state): + return state.transpose(-1, -2).contiguous() if value_first else state + + state = layout(args[0].new_zeros(1, 1, args[0].shape[-1], args[2].shape[-1]).float()) + prefix, state = prefill( + *[x[:, :65].detach().clone() for x in args], + initial_state=state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=torch.tensor([0, 65], device="cuda", dtype=torch.int32), + ) + outputs = [prefix] + state = layout(quantizer(layout(state)[0])[None]) + for token in range(65, 73): + output, state = step( + *[x[:, token : token + 1].detach().clone() for x in args], + initial_state=state, + inplace_final_state=False, + use_qk_l2norm_in_kernel=True, + cu_seqlens=torch.tensor([0, 1], device="cuda", dtype=torch.int32), + ) + outputs.append(output) + state = layout(quantizer(layout(state)[0])[None]) + return torch.cat(outputs, dim=1), layout(state) + @pytest.fixture(scope="module", params=[False, True]) def compiled_serving_case(request): """Compile one shared BF16 shape per model, outside the test-call timer.""" - pytest.importorskip("vllm.model_executor.layers.fla.ops.kda", exc_type=ModuleNotFoundError) kda = request.param torch.manual_seed(73) - args = [torch.randn(1, 73, 1, 32, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + # A rectangular GDN state detects accidental key/value-axis swaps. + args = [ + torch.randn(1, 73, 1, dim, device="cuda", dtype=torch.bfloat16) + for dim in (32, 32, 32 if kda else 64) + ] args += [ -torch.rand((1, 73, 1, 32) if kda else (1, 73, 1), device="cuda") * 0.03, torch.rand(1, 73, 1, device="cuda") * 0.4, ] args = [x.requires_grad_() for x in args] - policy = LinearAttentionConfig(backend="serving", precision="vllm_0_15") + policy = LinearAttentionConfig(backend="serving", precision="vllm") quantizer = TensorQuantizer( QuantizerAttributeConfig( num_bits=8, @@ -60,12 +104,13 @@ def compiled_serving_case(request): ) output, state = forward() torch.autograd.grad(output.float().sum() + state.sum(), args) + expected = _native_forward(args, quantizer, kda) torch.cuda.synchronize() - return args, quantizer, forward + return args, quantizer, forward, expected def test_serving_state_qdq_and_handoff_gradient(compiled_serving_case): - args, quantizer, forward = compiled_serving_case + args, quantizer, forward, (expected, expected_state) = compiled_serving_case calls = [] handle = quantizer.register_forward_hook(lambda *_: calls.append(True)) output, state = forward() @@ -73,7 +118,6 @@ def test_serving_state_qdq_and_handoff_gradient(compiled_serving_case): # One handoff QDQ, then one QDQ per suffix update; none inside the fresh prefill. assert len(calls) == 9 with torch.no_grad(): - expected, expected_state = forward() torch.testing.assert_close(output, expected, rtol=0, atol=0) torch.testing.assert_close(state, expected_state, rtol=0, atol=0) quantizer.disable() diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index 16ea8f5762a..c701b7c1e7b 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -18,6 +18,7 @@ import pytest import torch import torch.nn as nn +from _test_utils.torch.distributed.utils import spawn_multiprocess_job import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq @@ -223,7 +224,10 @@ def test_quantizer_roundtrip_and_hybrid_selection(tmp_path): @pytest.mark.parametrize("reverse", [False, True], ids=["fp8-to-replay", "replay-to-fp8"]) def test_restore_changed_policy_and_quantizer_together(tmp_path, reverse): states = [ - (GDN_STATE_FP8_DYNAMIC, LinearAttentionConfig(backend="serving")), + ( + GDN_STATE_FP8_DYNAMIC, + LinearAttentionConfig(backend="serving", precision="vllm_0_15"), + ), ( {"num_bits": 8, "axis": (0, 1), "type": "dynamic", "narrow_range": True}, LinearAttentionConfig(backend="serving", precision="replayssm"), @@ -315,3 +319,30 @@ def test_policy_rejects_unmatched_and_unimplemented_modes(): ): with pytest.raises(ValueError): QuantizeConfig(linear_attention=[{"module_name": "*", "cfg": policy}]) + + +def _test_policy_matches_across_stages(rank, size): + config = { + "quant_cfg": [{"quantizer_name": "*", "enable": False}], + "linear_attention": [{"module_name": "layer", "cfg": {"backend": "serving"}}], + "algorithm": None, + } + try: + # Only rank 1 owns a selected layer; both stages use the same recipe and phase. + layer = TinyGatedDeltaNet() if rank else nn.Linear(4, 4) + model = mtq.quantize(nn.ModuleDict({"layer": layer}), config) + with linear_attention_training_phase(model, [3]): + if rank: + assert model["layer"]._linear_attention_prefill_lengths == (3,) + if rank: + assert model["layer"]._linear_attention_prefill_lengths is None + + config["linear_attention"][0]["module_name"] = "missing" + with pytest.raises(ValueError, match="matches no supported"): + mtq.quantize(model, config) + finally: + torch.distributed.destroy_process_group() + + +def test_policy_matches_across_stages(skip_on_windows): + spawn_multiprocess_job(2, _test_policy_matches_across_stages, backend="gloo") From 5cbfe54780d751fbd99a16cb695e862e7501b25f Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 7 Oct 2026 21:22:16 -0700 Subject: [PATCH 17/27] Defer native vLLM inspection for documentation imports Cache state-layout and kernel-signature checks at first use so Sphinx can import serving adapters with mocked optional dependencies. Preserve runtime dispatch and kernel arithmetic. Validation: focused recursive Sphinx HTML build using repository configuration and warnings as errors; native metadata and CPU state-layout checks on vLLM 0.15.1, 0.20.0, and 0.30.0; scoped pre-commit and git diff --check. Signed-off-by: Kai Xu --- .../linear_attention/serving/_compat.py | 12 +++++++++--- .../linear_attention/serving/chunk_delta_h.py | 14 ++++++++++---- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py b/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py index 81ba6dd1d88..04153111bdb 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py @@ -15,6 +15,7 @@ """Keep native vLLM imports and state layout changes at the kernel boundary.""" +from functools import cache from importlib import import_module from importlib.util import find_spec @@ -22,8 +23,6 @@ __all__ = [] -# vLLM 0.16 transposed the recurrent cache; ModelOpt keeps [H,K,V] for QDQ/autograd. -STATE_V_FIRST = vllm.__version_tuple__[:2] >= (0, 16) _FLA_OPS = ( "vllm.third_party.flash_linear_attention.ops" if find_spec("vllm.third_party.flash_linear_attention") is not None @@ -31,10 +30,17 @@ ) +@cache +def state_v_first(): + """Inspect the native layout at first use, after optional-dependency docs imports.""" + # vLLM 0.16 transposed the recurrent cache; ModelOpt keeps [H,K,V] for QDQ/autograd. + return vllm.__version_tuple__[:2] >= (0, 16) + + def fla_module(name): return import_module(f"{_FLA_OPS}.{name}") def state_layout(state): """Convert between ModelOpt and native layout; transposing is its own inverse.""" - return state.transpose(-1, -2).contiguous() if STATE_V_FIRST else state.contiguous() + return state.transpose(-1, -2).contiguous() if state_v_first() else state.contiguous() diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py b/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py index b8759c9541c..08e43dc6489 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py @@ -22,19 +22,25 @@ """Save FP32 training intermediates using vLLM's unchanged chunk-state kernel.""" +from functools import cache from inspect import signature import torch import triton -from ._compat import STATE_V_FIRST, fla_module +from ._compat import fla_module, state_v_first _chunk = fla_module("chunk_delta_h") chunk_gated_delta_rule_fwd_kernel_h_blockdim64 = ( _chunk.chunk_gated_delta_rule_fwd_kernel_h_blockdim64 ) prepare_chunk_offsets = fla_module("index").prepare_chunk_offsets -_HAS_EXP2 = "use_exp2" in signature(_chunk.chunk_gated_delta_rule_fwd_h).parameters + + +@cache +def _has_exp2(): + """Inspect the native signature only when a kernel runs, not during docs imports.""" + return "use_exp2" in signature(_chunk.chunk_gated_delta_rule_fwd_h).parameters def chunk_state(k, w, u, g, gk, initial_state, cu_seqlens, use_exp2=False): @@ -42,7 +48,7 @@ def chunk_state(k, w, u, g, gk, initial_state, cu_seqlens, use_exp2=False): _, length, heads, key_dim = k.shape value_dim = u.shape[-1] # The Torch adjoint needs unrounded values. Output kernels receive BF16 casts. - state_shape = (value_dim, key_dim) if STATE_V_FIRST else (key_dim, value_dim) + state_shape = (value_dim, key_dim) if state_v_first() else (key_dim, value_dim) h = k.new_empty(1, triton.cdiv(length, 64), heads, *state_shape, dtype=torch.float32) updated = torch.empty_like(u, dtype=torch.float32) final = torch.empty_like(initial_state, dtype=torch.float32) @@ -66,6 +72,6 @@ def chunk_state(k, w, u, g, gk, initial_state, cu_seqlens, use_exp2=False): K=key_dim, V=value_dim, BT=64, - **({"USE_EXP2": use_exp2} if _HAS_EXP2 else {}), + **({"USE_EXP2": use_exp2} if _has_exp2() else {}), ) return h, updated, final From b3c7d6707dec3f909fe6a8607fe501ad5e166c42 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 9 Oct 2026 15:47:17 -0700 Subject: [PATCH 18/27] Fix linear-attention state QAT phase handling and serving profiles Consolidate the unpublished PR #2519 review fixes. Preserve valid packed lengths through selective recompute, expose the public phase API, and keep the Bridge workflow integration separate. Retain the existing FLA kernels for their follow-up PR. Signed-off-by: Kai Xu --- .pre-commit-config.yaml | 3 +- CHANGELOG.rst | 4 +- LICENSE | 1 - examples/megatron_bridge/quantize.py | 25 +- .../quantization/linear_attention/__init__.py | 3 +- .../linear_attention/fla_chunk_delta_h.py | 978 ++++++++++++++++++ .../fla_chunk_gated_delta_rule.py | 721 +++++++++++++ .../linear_attention/serving/__init__.py | 2 +- .../linear_attention/serving/_compat.py | 178 +++- .../linear_attention/serving/chunk_delta_h.py | 77 -- .../linear_attention/serving/forward.py | 281 +++-- modelopt/torch/quantization/conversion.py | 7 +- .../linear_attention/_vllm_autograd.py | 103 +- .../quantization/linear_attention/config.py | 21 +- .../quantization/linear_attention/decode.py | 107 +- .../quantization/linear_attention/gdn.py | 21 +- .../quantization/linear_attention/kda.py | 27 +- .../quantization/linear_attention/training.py | 204 +++- .../quantization/linear_attention/utils.py | 38 +- modelopt/torch/quantization/plugins/gdn.py | 11 +- modelopt/torch/quantization/plugins/kda.py | 4 +- .../quantization/plugins/linear_attention.py | 64 +- .../torch/quantization/plugins/megatron.py | 180 +++- ..._attention_state_int8_block32_dynamic.yaml | 5 +- ...on_state_int8_block32_dynamic_kimi_k3.yaml | 38 + modelopt_recipes/ptq.md | 66 +- pyproject.toml | 9 + .../linear_attention_reference.py | 191 ++-- .../megatron_bridge/test_quantize_export.py | 26 + .../test_fla_chunk_gated_delta_rule.py | 91 ++ .../plugins/test_megatron_gated_delta_net.py | 220 ++-- .../quantization/plugins/test_megatron_kda.py | 108 +- .../test_linear_attention_replay.py | 11 +- .../test_linear_attention_training.py | 367 +++++-- .../torch/quantization/plugins/test_gdn.py | 343 +++--- .../test_linear_attention_decode.py | 16 +- .../test_linear_attention_reference.py | 19 +- 37 files changed, 3562 insertions(+), 1008 deletions(-) create mode 100644 modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py create mode 100644 modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py delete mode 100644 modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py create mode 100644 modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic_kimi_k3.yaml create mode 100644 tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a5a3b97cd8c..f60fad5cc15 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -112,7 +112,8 @@ repos: modelopt/torch/_deploy/utils/onnx_utils.py| modelopt/torch/export/transformer_engine.py| modelopt/torch/fastgen/plugins/qwen_image_pdd.py| - modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.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 6440b085116..e6b458b5bea 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,7 +16,7 @@ Changelog - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. -- Add experimental serving-aligned GDN/KDA state QAT with explicit prefill/decode boundaries; migrate chunk-only state recipes to a serving policy and supply prefix lengths through the training phase context. Plain state QDQ uses the installed vLLM's native kernels via ``precision="vllm"`` (the legacy ``vllm_0_15`` spelling remains accepted); INT8 Hadamard and ReplaySSM require the compatible quantized-ReplaySSM serving fork. +- Add experimental GDN/KDA recurrent-state QAT with native vLLM Triton forwards and explicit prefill/decode phases through ``linear_attention_training_phase``. INT8 Hadamard and ReplaySSM require the compatible quantized-ReplaySSM serving fork. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before a separate KV-cache AutoQuantize stage, with independent resumable checkpoints for the weight and @@ -68,8 +68,6 @@ Changelog **Backward Breaking Changes** - Bump minimum transformers version to ``5.5`` instead of ``4.57``; transformers 4.x is no longer supported. Upgrade with ``pip install -U "nvidia-modelopt[hf]"``. -- Experimental GDN W quantization and its ``gdn_w_fp8_dynamic`` recipe unit are removed; keep ``gdn_w_quantizer`` disabled when using state QAT or restoring existing checkpoints. - - ``modelopt.torch.distill.plugins.megatron.TopKLogitsKLLoss`` (``logit_kl_topk`` in ``DistillationConfig``) is renamed to ``TopLogitsKLLoss``, keeping the old name as a deprecated alias. It now normalizes both distributions over the full vocabulary instead of re-normalizing over the Top-K entries, and always appends a "ghost" token holding the probability mass outside the Top-K to both student and teacher (matching Megatron-LM's offline cached-logits KD loss). Loss values change for existing ``logit_kl_topk`` runs. - ``LogitsAndIntermediatesLossBalancer`` (Megatron distillation plugin) no longer rescales the distillation loss to the magnitude of the LM loss when the LM loss is included. The total is now the fixed convex combination ``(1 - alpha) * lm_loss + alpha * kd_loss`` with ``DistillationConfig.kd_loss_alpha`` (in [0, 1]), matching Megatron-LM's offline cached-logits KD. The default ``kd_loss_alpha=1.0`` skips the LM loss, as before. ``DistillationConfig.skip_lm_loss`` and ``kd_loss_scale`` are removed and raise ``ValueError`` if passed; set ``kd_loss_alpha`` instead. In ``examples/megatron_bridge/distill.py``, ``--kd_loss_alpha`` replaces ``--no_skip_lm_loss`` and ``--kd_loss_scale``. - The ``examples/vllm_serve`` fakequant launcher no longer supports vLLM 0.9.0. Upgrade to a version listed as tested in the example README. diff --git a/LICENSE b/LICENSE index 0e99a82684b..abc6a0cb761 100644 --- a/LICENSE +++ b/LICENSE @@ -216,7 +216,6 @@ Portions of this repository were adapted from code originally authored by the following copyright holders, licensed under the Apache License, Version 2.0 (see full license text above): - Copyright contributors to the vLLM project Copyright 2021 The HuggingFace Inc. team Copyright 2022 The HuggingFace Team Copyright 2022, Lefebvre Dalloz Services diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index 0cc0999f202..a13723ee2cf 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -76,6 +76,7 @@ QUANT_CFG_CHOICES, RecipeSupersededAction, ) +from modelopt.torch.quantization.linear_attention import get_linear_attention_layers from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0 from modelopt.torch.utils.dataset_utils import get_supported_datasets from modelopt.torch.utils.mlflow import Tool, masked_args, resolved_recipe_texts @@ -307,6 +308,13 @@ def get_quant_config(args: argparse.Namespace) -> dict: def main(args: argparse.Namespace): + mtq_config = get_quant_config(args) + if mtq_config.get("linear_attention") and mtq.need_calibration(mtq_config): + raise ValueError( + "This example supports linear_attention recipes only without calibration. " + "For calibrated projection quantization combined with state quantization, " + "use mtq.quantize with a calibration loop inside linear_attention_training_phase." + ) trust_remote_code = is_safe_repo( trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_name_or_path ) @@ -362,8 +370,6 @@ def main(args: argparse.Namespace): ) print_rank_0(f"Using calibration dataset: {args.calib_dataset_name}") - mtq_config = get_quant_config(args) - # Quantize only the language model: disable quantizers on every top-level submodule that is not # the language model (vision tower + projector). Skip aliases of language-model submodules (e.g. # Qwen's ``self.decoder = language_model.decoder``) so the LM's own layers stay enabled. @@ -474,7 +480,20 @@ def forward_loop(_model=None): warn_rank_0( "Skipping the post-quantization generation sanity check because --compress is set." ) - if not args.skip_generate and not args.compress: + state_quantized = torch.tensor( + bool(get_linear_attention_layers(language_model, enabled_only=True)), + dtype=torch.int32, + device="cuda", + ) + if torch.distributed.is_initialized(): + torch.distributed.all_reduce(state_quantized, op=torch.distributed.ReduceOp.MAX) + state_quantized = bool(state_quantized.item()) + if state_quantized and not args.skip_generate: + warn_rank_0( + "Skipping generation sanity check for state quantization; " + "use an explicit prefix/suffix evaluation to measure state QDQ." + ) + if not args.skip_generate and not args.compress and not state_quantized: print_rank_0("\nTesting quantized model with custom prompts...") # Sanity-check text generation on the quantized language model. language_model.eval() diff --git a/modelopt/torch/kernels/quantization/linear_attention/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/__init__.py index 617bb6dcbd8..2d1f4587cfb 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/__init__.py +++ b/modelopt/torch/kernels/quantization/linear_attention/__init__.py @@ -16,5 +16,6 @@ """Linear-attention quantization adapters for optional serving kernels. Serving adapters import forward kernels only when selected; this package -initializer has no kernel dependencies. +initializer has no kernel dependencies. The existing FLA chunk kernels remain +available for direct calls with fla-core==0.5.1. """ 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..875cb39a7d0 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -0,0 +1,978 @@ +# 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 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 + +# ``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_MAX_BLOCK_V = 128 + + +@triton.jit +def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr): + """[ModelOpt] Dynamic FP8 E4M3 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))) + return tl.where(b_amax > 0, b_amax / 448.0, 1.0) + + +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 an FP8 cache is quantized before the first chunk uses it. + if STATE_QDQ == 1: + if K > 192: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) + elif K > 128: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K) + elif K > 64: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K) + else: + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K) + b_h1 = fp8_scalar_qdq(b_h1, b_scale) + if K > 64: + b_h2 = fp8_scalar_qdq(b_h2, b_scale) + if K > 128: + b_h3 = fp8_scalar_qdq(b_h3, b_scale) + if K > 192: + b_h4 = fp8_scalar_qdq(b_h4, b_scale) + + # 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 QDQ of the next-chunk or stored final state. + # Recompute amax over [K, BV]; fp8_scalar_qdq applies that tile's scalar scale. + if STATE_QDQ == 1: + if K > 192: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) + elif K > 128: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K) + elif K > 64: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K) + else: + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K) + b_h1 = fp8_scalar_qdq(b_h1, b_scale) + if K > 64: + b_h2 = fp8_scalar_qdq(b_h2, b_scale) + if K > 128: + b_h3 = fp8_scalar_qdq(b_h3, b_scale) + if K > 192: + b_h4 = fp8_scalar_qdq(b_h4, b_scale) + + 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 != STATE_QDQ_FP8_DYNAMIC: + raise ValueError(f"Unsupported state_qdq mode {state_qdq}; expected 0 or 1.") + 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..a5609cab102 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py @@ -0,0 +1,721 @@ +# 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_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 with 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): + raise ValueError(f"`state_qdq` must be 0 or 1, got {state_qdq}.") + if w_quantizer is not None: + validate_gdn_quantizer(w_quantizer, name="gdn_w_quantizer") + if state_qdq 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/serving/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py index 1156582d24e..e808e5de96d 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py @@ -20,6 +20,6 @@ except ImportError as error: raise ImportError( "Linear-attention serving profiles require the optional vLLM dependency: " - "public vLLM for 'vllm' (legacy spelling 'vllm_0_15'), or the compatible quantized-ReplaySSM " + "public vLLM for 'vllm', or the compatible quantized-ReplaySSM " "vLLM fork for 'replayssm'." ) from error diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py b/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py index 04153111bdb..418ca86e0e9 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py @@ -18,8 +18,9 @@ from functools import cache from importlib import import_module from importlib.util import find_spec +from inspect import signature -import vllm +import torch __all__ = [] @@ -32,15 +33,178 @@ @cache def state_v_first(): - """Inspect the native layout at first use, after optional-dependency docs imports.""" - # vLLM 0.16 transposed the recurrent cache; ModelOpt keeps [H,K,V] for QDQ/autograd. - return vllm.__version_tuple__[:2] >= (0, 16) + """Read the serving cache contract using unequal dimensions, independent of version tags.""" + from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator + + _, shape = MambaStateShapeCalculator.gated_delta_net_state_shape( + tp_world_size=1, + num_k_heads=1, + num_v_heads=1, + head_k_dim=32, + head_v_dim=64, + conv_kernel_size=4, + ) + if shape == (1, 64, 32): + return True + if shape == (1, 32, 64): + return False + raise NotImplementedError(f"Unsupported vLLM recurrent-state layout: {shape}") def fla_module(name): return import_module(f"{_FLA_OPS}.{name}") -def state_layout(state): - """Convert between ModelOpt and native layout; transposing is its own inverse.""" - return state.transpose(-1, -2).contiguous() if state_v_first() else state.contiguous() +@cache +def optional_module(path): + """Allow missing APIs, but never hide a broken dependency inside an installed module.""" + try: + return import_module(path) + except ModuleNotFoundError as error: + if error.name is not None and (path == error.name or path.startswith(error.name + ".")): + return None + raise + + +@cache +def kda_module(kimi_k3=False): + """Load only the explicitly selected KDA implementation, without backend fallback.""" + if not kimi_k3: + return fla_module("kda") + try: + return import_module("vllm.models.kimi_k3.nvidia.ops.third_party.kda") + except ImportError as error: + raise ImportError( + "The selected vLLM KDA Triton backend requires its model package dependencies. " + "Install the dependencies for this vLLM build; ModelOpt will not switch arithmetic " + "to a different KDA backend after an import failure." + ) from error + + +@cache +def gdn_module(name): + return optional_module(f"{_FLA_OPS}.{name}") + + +@cache +def kda_prefill_beta(kimi_k3=False): + """Resolve the raw-gate chunk callable's beta contract, independent of decode APIs.""" + kernel = getattr(kda_module(kimi_k3), "chunk_kda_with_fused_gate", None) + if kernel is None: + return None + parameters = signature(kernel).parameters + for name in ("raw_beta", "beta"): + if name in parameters: + return name + raise NotImplementedError("Unsupported vLLM chunk_kda_with_fused_gate beta signature") + + +def kda_chunk(q, k, v, g, beta, *, gate_inputs=None, kimi_k3=False, **kwargs): + """Use the same native chunk dispatch for execution and cache-contract checks.""" + native = kda_module(kimi_k3) + if gate_inputs is not None and (beta_name := kda_prefill_beta(kimi_k3)) is not None: + raw_g, raw_beta, rate, bias = gate_inputs + return native.chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_g, + **{beta_name: raw_beta if beta_name == "raw_beta" else beta}, + A_log=rate, + g_bias=bias, + **kwargs, + ) + if kimi_k3: + raise ValueError("vllm_kimi_k3 requires native fused-gate prefill and raw gate_inputs") + return native.chunk_kda(q, k, v, g, beta, **kwargs) + + +@cache +def _kda_state_v_first(device, kimi_k3, keys, raw_gates=False): + """Check the selected KDA paths at the actual head size, once per device.""" + with torch.no_grad(): + q = torch.zeros(1, 1, 1, keys, device=device, dtype=torch.bfloat16) + q[..., 0] = 1 + zero = torch.zeros_like(q) + # Distinct row/column readouts made from exactly representable BF16 values, + # including head sizes such as 192 where arange(K*K) would round. + pattern = torch.zeros(keys, keys, device=device, dtype=torch.float32) + markers = torch.arange(1, keys + 1, device=device, dtype=torch.float32) + pattern[0, :] = -markers + pattern[:, 0] = markers + state = pattern.expand(2, 1, keys, keys).contiguous() + output, _ = kda_module(kimi_k3).fused_recurrent_kda_fwd( + q=q, + k=zero, + v=zero, + g=zero.float(), + beta=torch.zeros(1, 1, 1, device=device), + scale=1.0, + initial_state=state, + inplace_final_state=True, + cu_seqlens=torch.arange(2, device=device, dtype=torch.int32), + ssm_state_indices=torch.ones(1, device=device, dtype=torch.int32), + use_qk_l2norm_in_kernel=False, + ) + row = output.flatten().float() + if torch.equal(row, pattern[:, 0]): + value_first = True + elif torch.equal(row, pattern[0]): + value_first = False + else: + raise NotImplementedError("The installed vLLM KDA cache layout could not be validated") + # Some native revisions changed cache storage without updating chunk readout. + # Reject that backend rather than training against a silently transposed prefix. + chunk_args = { + "q": q, + "k": zero, + "v": zero.clone(), + "scale": 1.0, + "initial_state": state[1:].clone(), + "output_final_state": True, + "use_qk_l2norm_in_kernel": kimi_k3, + "cu_seqlens": torch.arange(2, device=device, dtype=torch.int32), + } + gate_inputs = None + if raw_gates or kimi_k3: + raw_g = torch.full_like(q, -80) + raw_beta = torch.full((1, 1, 1), -80.0, device=device) + rate, bias = torch.zeros(1, device=device), torch.zeros(keys, device=device) + gate_inputs = (raw_g, raw_beta, rate, bias) + if kimi_k3: + packed, _ = kda_module(True).fused_recurrent_kda_packed_decode( + mixed_qkv=torch.cat((q.flatten(), zero.flatten(), zero.flatten()))[None], + raw_g=raw_g, + raw_beta=raw_beta, + A_log=rate, + dt_bias=bias, + lower_bound=None, + initial_state=state.clone(), + state_indices=torch.ones(1, device=device, dtype=torch.int32), + scale=1.0, + ) + if not torch.equal(packed.flatten().float(), row): + raise NotImplementedError("The installed vLLM KDA packed cache layout disagrees") + chunk_output, _ = kda_chunk( + g=zero.float(), + beta=torch.zeros(1, 1, 1, device=device), + gate_inputs=gate_inputs, + kimi_k3=kimi_k3, + **chunk_args, + ) + if not torch.equal(chunk_output.flatten().float(), row): + raise NotImplementedError( + "The installed vLLM KDA chunk/decode cache readout disagrees; " + "use a vLLM build with the KDA chunk-state layout fix" + ) + return value_first + + +def state_layout(state, *, kda=False, kimi_k3=False, raw_gates=False): + """Convert between ModelOpt and the selected model's native state layout.""" + value_first = ( + _kda_state_v_first(state.device, kimi_k3, state.shape[-1], raw_gates) + if kda + else state_v_first() + ) + return state.transpose(-1, -2).contiguous() if value_first else state.contiguous() diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py b/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py deleted file mode 100644 index 08e43dc6489..00000000000 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/chunk_delta_h.py +++ /dev/null @@ -1,77 +0,0 @@ -# Adapted from: https://github.com/vllm-project/vllm/blob/1892993bc18e243e2c05841314c5e9c06a80c70d/vllm/model_executor/layers/fla/ops/chunk_delta_h.py -# Modifications: import the kernel; retain FP32 intermediates for the training adjoint. -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# SPDX-License-Identifier: Apache-2.0 -# Original FLA code: Copyright (c) 2023-2025, Songlin Yang, Yu Zhang. -# FLA is licensed under the MIT license reproduced in the root LICENSE. - -# 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. - -"""Save FP32 training intermediates using vLLM's unchanged chunk-state kernel.""" - -from functools import cache -from inspect import signature - -import torch -import triton - -from ._compat import fla_module, state_v_first - -_chunk = fla_module("chunk_delta_h") -chunk_gated_delta_rule_fwd_kernel_h_blockdim64 = ( - _chunk.chunk_gated_delta_rule_fwd_kernel_h_blockdim64 -) -prepare_chunk_offsets = fla_module("index").prepare_chunk_offsets - - -@cache -def _has_exp2(): - """Inspect the native signature only when a kernel runs, not during docs imports.""" - return "use_exp2" in signature(_chunk.chunk_gated_delta_rule_fwd_h).parameters - - -def chunk_state(k, w, u, g, gk, initial_state, cu_seqlens, use_exp2=False): - """Return chunk-start states, residual values, and final state for one sequence.""" - _, length, heads, key_dim = k.shape - value_dim = u.shape[-1] - # The Torch adjoint needs unrounded values. Output kernels receive BF16 casts. - state_shape = (value_dim, key_dim) if state_v_first() else (key_dim, value_dim) - h = k.new_empty(1, triton.cdiv(length, 64), heads, *state_shape, dtype=torch.float32) - updated = torch.empty_like(u, dtype=torch.float32) - final = torch.empty_like(initial_state, dtype=torch.float32) - chunk_gated_delta_rule_fwd_kernel_h_blockdim64[ - lambda meta: (triton.cdiv(value_dim, meta["BV"]), heads) - ]( - k=k, - v=u, - w=w, - v_new=updated, - g=g, - gk=gk, - h=h, - h0=initial_state, - ht=final, - cu_seqlens=cu_seqlens, - chunk_offsets=prepare_chunk_offsets(cu_seqlens, 64), - T=length, - H=heads, - Hg=heads, - K=key_dim, - V=value_dim, - BT=64, - **({"USE_EXP2": use_exp2} if _has_exp2() else {}), - ) - return h, updated, final diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py index 1d0866ca21f..7c54c755f17 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py @@ -13,124 +13,180 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""vLLM forward primitives with FP32 saved values; autograd lives in quantization.""" +"""Native vLLM Triton forwards; training owns private, differentiable state.""" -from functools import cache -from importlib import import_module +from functools import cache, lru_cache import torch -from ._compat import fla_module, state_layout -from .chunk_delta_h import chunk_state - -chunk_fwd_o = fla_module("chunk_o").chunk_fwd_o -chunk_scaled_dot_kkt_fwd = fla_module("chunk_scaled_dot_kkt").chunk_scaled_dot_kkt_fwd -chunk_local_cumsum = fla_module("cumsum").chunk_local_cumsum -fused_recurrent_gated_delta_rule = fla_module("fused_recurrent").fused_recurrent_gated_delta_rule -_kda = fla_module("kda") -chunk_gla_fwd_o_gk = _kda.chunk_gla_fwd_o_gk -chunk_kda_scaled_dot_kkt_fwd = _kda.chunk_kda_scaled_dot_kkt_fwd -fused_recurrent_kda = _kda.fused_recurrent_kda -fused_kda_gate = _kda.fused_kda_gate -kda_wu = _kda.recompute_w_u_fwd -l2norm_fwd = fla_module("l2norm").l2norm_fwd -solve_tril = fla_module("solve_tril").solve_tril -gdn_wu = fla_module("wy_fast").recompute_w_u_fwd -_KDA_GATE_SCALE = getattr(_kda, "RCP_LN2", 1.0) - - -def prefill(q, k, v, g, beta, state, scale, normalize=False): - """Return one sequence's output, final state, and rounded forward intermediates.""" +from ._compat import fla_module, gdn_module, kda_chunk, kda_module, optional_module, state_layout + + +def _gdn_inputs(q, k, v, gate_inputs, normalize): + raw_g, raw_beta, a_log, bias = gate_inputs + prep = gdn_module("fused_gdn_prefill_post_conv") + if prep is not None: + return prep.fused_post_conv_prep( + torch.cat([x.flatten(1) for x in (q, k, v)], dim=-1), + raw_g.flatten(0, -2), + raw_beta.flatten(0, -2), + a_log, + bias, + q.shape[-2], + q.shape[-1], + v.shape[-1], + apply_l2norm=normalize, + ), False + g, beta = _gdn_gate()( + a_log, raw_g.reshape(-1, raw_g.shape[-1]), raw_beta.reshape(-1, raw_beta.shape[-1]), bias + ) + return (q, k, v, g[0], beta[0]), normalize + + +def prefill(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, kimi_k3=False): + """Run the complete native chunk entry point, including raw-gate preparation.""" + if kimi_k3 and gate_inputs is None: + raise ValueError("vllm_kimi_k3 requires raw gate_inputs") + channel = g.ndim == 3 + if not channel and gate_inputs is not None: + (q, k, v, g, beta), normalize = _gdn_inputs(q, k, v, gate_inputs, normalize) q, k, v, g, beta = [x.unsqueeze(0).contiguous() for x in (q, k, v, g, beta)] - cu = torch.tensor([0, q.shape[1]], device=q.device, dtype=torch.int32) - if normalize: - q, k = l2norm_fwd(q), l2norm_fwd(k) - channel = g.ndim == 4 - gc = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu) - # Newer native KDA kernels evaluate exp2 of base-2 cumulative log gates. - natural_gc = gc - if channel: - gc = gc * _KDA_GATE_SCALE - if channel: - lower, scores = chunk_kda_scaled_dot_kkt_fwd(q, k, gc, beta, scale=scale, cu_seqlens=cu) - else: - lower = chunk_scaled_dot_kkt_fwd( - k=k, beta=beta, g=gc, cu_seqlens=cu, output_dtype=torch.float32 - ) - scores = None - inverse = solve_tril(A=lower, cu_seqlens=cu, output_dtype=k.dtype) + # Some native chunk implementations reuse V as their output buffer. + v = v.clone() + kwargs = { + "scale": scale, + "initial_state": state_layout( + state.unsqueeze(0), kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None + ), + "output_final_state": True, + "use_qk_l2norm_in_kernel": normalize, + "cu_seqlens": _prefill_metadata(q.device, q.shape[1]), + } if channel: - w, u, _, kg = kda_wu(k=k, v=v, beta=beta, A=inverse, gk=gc, cu_seqlens=cu) + gates = None + if gate_inputs is not None: + raw_g, raw_beta, rate, bias = gate_inputs + gates = (raw_g[None].contiguous(), raw_beta[None].contiguous(), rate, bias) + out, final = kda_chunk(q, k, v, g, beta, gate_inputs=gates, kimi_k3=kimi_k3, **kwargs) else: - w, u = gdn_wu(k=k, v=v, beta=beta, A=inverse, g_cumsum=gc, cu_seqlens=cu) - kg = k - assert kg is not None - h, updated, final = chunk_state( - k=kg, - w=w, - u=u, - g=None if channel else gc, - gk=gc if channel else None, - initial_state=state_layout(state.unsqueeze(0)), - cu_seqlens=cu, - use_exp2=channel and _KDA_GATE_SCALE != 1.0, + out, final = fla_module("chunk").chunk_gated_delta_rule(q, k, v, g, beta, **kwargs) + return out[0], state_layout( + final[0], kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None + ) + + +def step(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, kimi_k3=False): + """Run native fused decode on a private cache; never mutate the incoming state.""" + if kimi_k3 and gate_inputs is None: + raise ValueError("vllm_kimi_k3 requires raw gate_inputs") + channel = g.ndim == 2 + cu, indices = _decode_metadata(q.device) + native_state = state_layout( + state.unsqueeze(0), kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None ) - assert updated is not None and final is not None + if gate_inputs is not None: + raw_g, raw_beta, a_log, bias = gate_inputs + if channel and kimi_k3: + if not normalize: + raise ValueError("vLLM packed KDA decode requires Q/K normalization") + # Slot zero is reserved by vLLM's paged-cache kernels. + native_state = torch.cat((torch.zeros_like(native_state), native_state), dim=0) + out, final = kda_module(kimi_k3).fused_recurrent_kda_packed_decode( + mixed_qkv=torch.cat([x.flatten() for x in (q, k, v)])[None], + raw_g=raw_g[None, None].contiguous(), + raw_beta=raw_beta[None, None].contiguous(), + A_log=a_log.reshape(-1).contiguous(), + dt_bias=(torch.zeros_like(raw_g).flatten() if bias is None else bias.contiguous()), + lower_bound=None, + initial_state=native_state, + state_indices=indices, + scale=scale, + ) + return out[0, 0], state_layout( + final[1], kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None + ) + packed = ( + None + if channel + else getattr( + fla_module("fused_recurrent"), + "fused_recurrent_gated_delta_rule_packed_decode", + None, + ) + ) + if packed is not None: + # Match vLLM's pure, non-speculative decode path on private cache slot one. + native_state = torch.cat((torch.zeros_like(native_state), native_state), dim=0) + out, final = packed( + mixed_qkv=torch.cat([x.flatten() for x in (q, k, v)])[None], + a=raw_g[None].contiguous(), + b=raw_beta[None].contiguous(), + A_log=a_log.reshape(-1).contiguous(), + dt_bias=bias.contiguous(), + scale=scale, + initial_state=native_state, + out=torch.empty_like(v)[None, None], + ssm_state_indices=indices, + use_qk_l2norm_in_kernel=normalize, + ) + return out[0, 0], state_layout( + final[1], kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None + ) + fused = None if channel else gdn_module("fused_sigmoid_gating") + if fused is not None: + out, final = fused.fused_sigmoid_gating_delta_rule_update( + A_log=a_log, + a=raw_g[None].contiguous(), + b=raw_beta[None].contiguous(), + dt_bias=bias, + q=q[None, None], + k=k[None, None], + v=v[None, None], + initial_state=native_state, + inplace_final_state=False, + scale=scale, + use_qk_l2norm_in_kernel=normalize, + cu_seqlens=cu, + ) + return out[0, 0], state_layout( + final[0], kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None + ) + if not channel: + g, beta = _gdn_gate()( + a_log, raw_g[None].contiguous(), raw_beta[None].contiguous(), bias + ) + g, beta = g[0, 0], beta[0, 0] + # Older vLLM models and standalone calls consume already activated gates. if channel: - out = chunk_gla_fwd_o_gk( - q=q, - v=updated.to(v.dtype), - g=gc, - A=scores, - h=h.to(k.dtype), + native_state = torch.cat((torch.zeros_like(native_state), native_state), dim=0) + out, final = kda_module(kimi_k3).fused_recurrent_kda_fwd( + *[x[None, None].contiguous() for x in (q, k, v, g, beta)], + initial_state=native_state, + inplace_final_state=True, scale=scale, - o=torch.empty_like(v), + use_qk_l2norm_in_kernel=normalize, cu_seqlens=cu, - chunk_size=64, + ssm_state_indices=indices, ) - else: - out = chunk_fwd_o( - q=q, k=k, v=updated.to(v.dtype), h=h.to(k.dtype), g=gc, scale=scale, cu_seqlens=cu + return out[0, 0], state_layout( + final[1], kda=True, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None ) - intermediates = { - "q": q[0], - "k": k[0], - "g": natural_gc[0], - "lower": lower[0], - "inverse": inverse[0], - "w": w[0], - "u": u[0], - "h": state_layout(h[0]), - "updated": updated[0], - "kg": kg[0], - "scores": None if scores is None else scores[0], - } - return out[0], state_layout(final[0]), intermediates - - -def step(q, k, v, g, beta, state, scale, normalize=False): - """Run one native token update without modifying the incoming state.""" - recurrent = fused_recurrent_kda if g.ndim == 2 else fused_recurrent_gated_delta_rule + recurrent = fla_module("fused_recurrent").fused_recurrent_gated_delta_rule out, final = recurrent( *[x[None, None].contiguous() for x in (q, k, v, g, beta)], - initial_state=state_layout(state.unsqueeze(0)), + initial_state=native_state, inplace_final_state=False, scale=scale, use_qk_l2norm_in_kernel=normalize, - cu_seqlens=torch.tensor([0, 1], device=q.device, dtype=torch.int32), - # Private dense training state needs no paged-cache index (0 is now reserved). + cu_seqlens=cu, ssm_state_indices=None, ) - return out[0, 0], state_layout(final[0]) - - -def fused_gdn_gating(*args, **kwargs): - """Load the optional vLLM model only when Megatron needs GDN gate preparation.""" - return _gdn_gate()(*args, **kwargs) + return out[0, 0], state_layout( + final[0], kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None + ) -# Megatron compiles gate preparation; module discovery must execute outside that graph. -@torch.compiler.disable +# Old runtimes keep gate preparation beside the model rather than in the FLA package. @cache def _gdn_gate(): for path in ( @@ -138,9 +194,30 @@ def _gdn_gate(): "vllm.model_executor.layers.mamba.gdn_linear_attn", "vllm.model_executor.models.qwen3_next", ): - try: - return import_module(path).fused_gdn_gating - except ModuleNotFoundError as error: # noqa: PERF203 - cached, one-time import discovery - if error.name is None or not (path == error.name or path.startswith(error.name + ".")): - raise - raise ImportError("The installed vLLM does not provide fused_gdn_gating") + module = optional_module(path) + if module is not None and hasattr(module, "fused_gdn_gating"): + return module.fused_gdn_gating + raise NotImplementedError("The installed vLLM does not expose supported GDN gate preparation") + + +@lru_cache(maxsize=128) +def _prefill_metadata(device, length): + # Native kernels only read these boundaries; reuse them across equal-length layers. + return torch.arange(2, device=device, dtype=torch.int32) * length + + +@cache +def _decode_metadata(device): + # Immutable single-token boundaries and cache slot; construct without a CPU-to-GPU copy. + return torch.arange(2, device=device, dtype=torch.int32), torch.ones( + 1, device=device, dtype=torch.int32 + ) + + +def validate_profile(gate_inputs, normalize, *, kimi_k3=False): + """Validate the selected native serving contract before launching prefill.""" + if kimi_k3: + if gate_inputs is None: + raise ValueError("vllm_kimi_k3 requires raw gate_inputs") + if not normalize: + raise ValueError("vLLM packed KDA decode requires Q/K normalization") diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 405f9b8876b..c77ab1b3623 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -106,8 +106,9 @@ def restore_quantized_model( model = model.init_modellike() if isinstance(model, ModelLikeModule) else model replace_quant_module(model, version=ModeloptStateManager(model).state_version) set_quantizer_by_cfg(model, config.get("quant_cfg", [])) - # Rebuild from the recipe without validating against partially restored state. - _apply_linear_attention_policy(model, config) + # Current checkpoints save resolved local policies; only older ones need matching. + if "linear_attention" not in metadata: + _apply_linear_attention_policy(model, config) return restore_quantizer_state(model, config, metadata) @@ -220,7 +221,7 @@ def update_quantize_metadata( from .plugins.linear_attention import _linear_attention_state policies = _linear_attention_state(model) - if policies: + if policies or getattr(config, "linear_attention", []): metadata["linear_attention"] = policies else: metadata.pop("linear_attention", None) diff --git a/modelopt/torch/quantization/linear_attention/_vllm_autograd.py b/modelopt/torch/quantization/linear_attention/_vllm_autograd.py index e25971171b5..c5165daaae1 100644 --- a/modelopt/torch/quantization/linear_attention/_vllm_autograd.py +++ b/modelopt/torch/quantization/linear_attention/_vllm_autograd.py @@ -15,8 +15,8 @@ """Differentiable adapter for the installed serving runtime's arithmetic. -The kernels supply forward values. Autograd differentiates the corresponding -operations evaluated at those values; operand casts and QDQ use identity STE. +Native kernels supply forward values. Prefill uses FLA's recurrence backward; +decode uses a Torch adjoint. QDQ uses identity STE. The optional vLLM dependency supplies kernels; training owns its state and needs no server. """ @@ -25,20 +25,18 @@ from .utils import forward_value -def rounded(value): - return forward_value(value, value.to(torch.bfloat16)) - - def normalized(value): return value / (value.square().sum(-1, keepdim=True) + 1e-6).sqrt() -def prefix(q, k, v, g, beta, state, scale, beta_dtype, normalize=False): - # Keep the optional vLLM dependency isolated to the native precision profile. +def prefix( + q, k, v, g, beta, state, scale, beta_dtype, normalize=False, gate_inputs=None, precision="vllm" +): + # vLLM supplies the forward; the optional FLA dependency supplies the recurrence adjoint. from ...kernels.quantization.linear_attention.serving.forward import prefill with torch.no_grad(): - out, final, saved = prefill( + out, final = prefill( q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16), @@ -47,73 +45,30 @@ def prefix(q, k, v, g, beta, state, scale, beta_dtype, normalize=False): state.float(), scale, normalize, + gate_inputs, + kimi_k3=precision == "vllm_kimi_k3", ) if not torch.is_grad_enabled() or not any(x.requires_grad for x in (q, k, v, g, beta, state)): return out.float(), final - q, k = [ - forward_value(normalized(x) if normalize else x, saved[n]) for n, x in (("q", q), ("k", k)) - ] - outputs = [] - channel = g.ndim == 3 - for chunk, lo in enumerate(range(0, len(q), 64)): - hi = min(lo + 64, len(q)) - qc, kc, vc = [x[lo:hi].transpose(0, 1) for x in (q, k, v)] - bc = beta[lo:hi].transpose(0, 1).unsqueeze(-1) - gc = forward_value(g[lo:hi].transpose(0, 1).cumsum(1), saved["g"][lo:hi].transpose(0, 1)) - state = forward_value(state, saved["h"][chunk]) - hs = rounded(state) - count = hi - lo - - def matrix(name, value): - return forward_value(value, saved[name][lo:hi].transpose(0, 1)[..., :count]) - - if channel: - lower_rows, score_rows = [], [] - for row in range(count): - right = ( - kc[:, : row + 1] * (gc[:, row : row + 1] - gc[:, : row + 1]).exp() - ).transpose(-1, -2) - lower_rows.append( - torch.nn.functional.pad( - (bc[:, row : row + 1] * kc[:, row : row + 1]) @ right, (0, count - row - 1) - ) - ) - score_rows.append( - torch.nn.functional.pad( - (qc[:, row : row + 1] * scale) @ right, (0, count - row - 1) - ) - ) - lower = matrix("lower", torch.cat(lower_rows, dim=1).tril(-1)) - scores = matrix("scores", torch.cat(score_rows, dim=1)) - gate = gc.exp() - else: - causal = torch.ones(count, count, device=q.device, dtype=torch.bool).tril() - decay = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() - lower = matrix("lower", (bc * (kc @ kc.transpose(-1, -2)) * decay).tril(-1)) - scores = ((qc @ kc.transpose(-1, -2)) * decay).tril() - gate = gc.exp().unsqueeze(-1) - eye = torch.eye(count, device=q.device, dtype=q.dtype).expand_as(lower) - inverse = matrix( - "inverse", - torch.linalg.solve_triangular(eye + lower, eye, upper=False, unitriangular=True), - ) - u = forward_value(inverse @ rounded(bc * vc), saved["u"][lo:hi].transpose(0, 1)) - kb = rounded(bc * kc) if beta_dtype == torch.bfloat16 else bc * kc - w = forward_value(inverse @ rounded(kb * gate), saved["w"][lo:hi].transpose(0, 1)) - updated = forward_value(u - w @ hs, saved["updated"][lo:hi].transpose(0, 1)) - if channel: - output = rounded(rounded(qc * scale) * gate) @ hs + rounded(scores) @ rounded(updated) - kg = forward_value(kc * (gc[:, -1:] - gc).exp(), saved["kg"][lo:hi].transpose(0, 1)) - state = state * gate[:, -1, :, None] + kg.transpose(-1, -2) @ rounded(updated) - else: - output = ((qc @ hs) * gate + rounded(scores) @ rounded(updated)) * scale - weighted = rounded(updated * (gc[:, -1:] - gc).exp().unsqueeze(-1)) - state = state * gate[:, -1, :, None] + kc.transpose(-1, -2) @ weighted - outputs.append(forward_value(output.transpose(0, 1), out[lo:hi])) - return torch.cat(outputs), forward_value(state, final) - - -def step(q, k, v, gate, beta, state, scale, normalize=False): + # FLA differentiates the same recurrence. Its intermediate rounding is a backward + # surrogate, not a claim that vLLM inference kernels have an exact backward. + if g.ndim == 3: + from fla.ops.kda import chunk_kda as chunk + else: + from fla.ops.gated_delta_rule import chunk_gated_delta_rule as chunk + adjoint, adjoint_state = chunk( + *[x.to(torch.bfloat16)[None] for x in (q, k, v)], + g[None], + beta[None], + initial_state=state[None], + output_final_state=True, + scale=scale, + use_qk_l2norm_in_kernel=normalize, + ) + return forward_value(adjoint[0].float(), out), forward_value(adjoint_state[0], final) + + +def step(q, k, v, gate, beta, state, scale, normalize=False, gate_inputs=None, precision="vllm"): # Keep the optional vLLM dependency isolated to the native precision profile. from ...kernels.quantization.linear_attention.serving.forward import step as native_step @@ -127,6 +82,8 @@ def step(q, k, v, gate, beta, state, scale, normalize=False): state.float(), scale, normalize, + gate_inputs, + kimi_k3=precision == "vllm_kimi_k3", ) if not torch.is_grad_enabled() or not any( x.requires_grad for x in (q, k, v, gate, beta, state) diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index a616a94dba3..f1536a3510c 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -28,8 +28,20 @@ class LinearAttentionConfig(ModeloptBaseConfig): """One execution policy shared by the chunked prefix and recurrent suffix. ``serving`` uses native forward arithmetic with a differentiable adjoint; - ``fla`` is the disabled/default path. ``vllm`` uses the installed vLLM's arithmetic; - ``vllm_0_15`` remains an accepted legacy spelling for existing checkpoints. + ``fla`` is the disabled/default path. ``vllm`` targets the installed vLLM's + Triton prefill and single-token, non-speculative decode with BF16 operands and + FP32 recurrent storage. Deployment must explicitly select Triton prefill + (``gdn_prefill_backend="triton"``) and enable packed recurrent decode where + available (``VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE=1``). Mixed prefill/decode + batches and automatic FlashInfer/CuTe dispatch are outside this profile. + KDA uses the standalone FLA/Triton API, with no model-level serving guarantee. + ``vllm_kimi_k3`` explicitly selects + vLLM's Kimi-K3 Triton kernels (also used by Kimi-Linear in vLLM 0.30), + raw projections and that model package's dependencies. + These targets are saved in ModelOpt/Megatron policy metadata, not HF export; + they do not configure or validate a remote server. Kimi serving must set both + ``kda_prefill_backend`` and ``kda_decode_backend`` to ``triton``; + the supported gate is unbounded softplus. ``replayssm`` uses the serving fork's INT8/Hadamard checkpoint and ring kernels. ``replay_window=1`` refreshes the state every token; larger windows enable replay. Fresh prefill has no internal state QDQ; incoming continuation state and @@ -40,7 +52,7 @@ class LinearAttentionConfig(ModeloptBaseConfig): """ backend: Literal["fla", "serving"] = ModeloptField(default="fla") - precision: Literal["vllm", "vllm_0_15", "replayssm"] = ModeloptField(default="vllm") + precision: Literal["vllm", "vllm_kimi_k3", "replayssm"] = ModeloptField(default="vllm") replay_window: int = Field(default=1, ge=1, le=64, strict=True) state_block_v: Literal[16, 32, 64, 128] = ModeloptField(default=64) @@ -65,7 +77,8 @@ class LinearAttentionPolicyEntry(ModeloptBaseConfig): """Assign a complete policy to supported modules matching ``module_name``. Rules apply in order: the last match wins, without merging fields. - A rule must match at least one supported module across distributed stages. + Unmatched rules are ignored locally, as with quant_cfg, so a shared recipe + can be applied independently to pipeline stages and virtual-pipeline chunks. """ module_name: str = Field(...) diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py index 622a825b900..cac33968d4b 100644 --- a/modelopt/torch/quantization/linear_attention/decode.py +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -20,11 +20,15 @@ import torch from .config import LinearAttentionConfig -from .utils import _resolve_state_quantizer, _tile_qdq, forward_value +from .utils import ( + _resolve_state_quantizer, + _state_qdq_with_scales, + _validate_gate_inputs, + forward_value, +) __all__ = [ "EncodedLinearAttentionTensor", - "LinearAttentionCarry", "LinearAttentionState", "ReplayEntry", "recurrent_decode", @@ -101,10 +105,6 @@ def reconstruct(self, *, original_basis=True): return state -# Compatibility name for existing imports and pickled runtime state objects. -LinearAttentionCarry = LinearAttentionState - - def _hadamard32(value): """Apply the orthonormal Sylvester transform to contiguous 32-value groups.""" shape = value.shape @@ -153,12 +153,9 @@ def _encode( """Apply TensorQuantizer state QDQ with straight-through gradients.""" if not enabled: return EncodedLinearAttentionTensor(value, None, "identity", None) + decoded, scales = _state_qdq_with_scales(value, block_v, state_format, state_quantizer) if state_quantizer is not None and state_quantizer.block_sizes is not None: - # TensorQuantizer owns dynamic scales; the floating carry needs only its QDQ output. - return EncodedLinearAttentionTensor( - state_quantizer(value), None, state_format, state_quantizer.block_sizes[-1] - ) - decoded, scales = _tile_qdq(value, block_v, state_format, state_quantizer=state_quantizer) + block_v = state_quantizer.block_sizes[-1] return EncodedLinearAttentionTensor(decoded, scales, state_format, block_v) @@ -285,7 +282,7 @@ def recurrent_decode( position=0, scale=None, use_qk_l2norm_in_kernel=False, - replay_gate_inputs=None, + gate_inputs=None, ) -> tuple[torch.Tensor, LinearAttentionState]: """Run a recurrent suffix and return its outputs and resumable runtime state. @@ -312,10 +309,60 @@ def recurrent_decode( state_quantizer, state_qdq, state_format = _resolve_state_quantizer( state_quantizer, state_qdq, state_format ) + _validate_gate_inputs(g, beta, gate_inputs) + if config.precision == "vllm_kimi_k3" and g.ndim != 3: + raise ValueError("vllm_kimi_k3 is a KDA-only precision profile") + if config.precision != "replayssm": + # Standalone decode validates before any cache writes. + from ...kernels.quantization.linear_attention.serving.forward import validate_profile + + validate_profile( + gate_inputs, + use_qk_l2norm_in_kernel, + kimi_k3=config.precision == "vllm_kimi_k3", + ) + return _recurrent_decode( + q, + k, + v, + g, + beta, + config=config, + state_qdq=state_qdq, + state_format=state_format, + state_quantizer=state_quantizer, + initial_state=initial_state, + carry=carry, + position=position, + scale=scale, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + gate_inputs=gate_inputs, + ) + + +def _recurrent_decode( + q, + k, + v, + g, + beta, + *, + config, + state_qdq, + state_format, + state_quantizer, + initial_state=None, + carry=None, + position=0, + scale=None, + use_qk_l2norm_in_kernel=False, + gate_inputs=None, +): + """Execute a suffix after its caller validates the shared quantizer once.""" if config.backend != "serving": raise ValueError("State QAT requires backend='serving'") block_v = config.state_block_v - serving = config.precision in ("vllm", "vllm_0_15") + serving = config.precision != "replayssm" native_replay = config.precision == "replayssm" if q.device.type != "cuda": raise ValueError("Serving arithmetic requires CUDA") @@ -329,7 +376,7 @@ def recurrent_decode( and ( q.shape[-1] < 32 or q.shape[-1] & (q.shape[-1] - 1) - or (g.ndim == 2 and replay_gate_inputs is None) + or (g.ndim == 2 and gate_inputs is None) ) ): raise ValueError("ReplaySSM requires power-of-two K >= 32 and raw GDN gate inputs") @@ -362,7 +409,6 @@ def recurrent_decode( with torch.autocast(device_type=q.device.type, enabled=False): for t in range(len(q)): entries = carry.entries - state = carry.reconstruct(original_basis=False) gate = g[t] if native_replay: from ...kernels.quantization.linear_attention.serving.replay import step @@ -380,11 +426,11 @@ def recurrent_decode( scale, use_qk_l2norm_in_kernel, ( - replay_gate_inputs[0][t], - replay_gate_inputs[1][t], - *replay_gate_inputs[2:], + gate_inputs[0][t], + gate_inputs[1][t], + *gate_inputs[2:], ) - if replay_gate_inputs is not None + if gate_inputs is not None else None, ) native_outputs.append(native[0]) @@ -392,7 +438,20 @@ def recurrent_decode( from ._vllm_autograd import step as serving_step native_output, working = serving_step( - q[t], k[t], v[t], gate, beta[t], state, scale, use_qk_l2norm_in_kernel + q[t], + k[t], + v[t], + gate, + beta[t], + carry.anchor.values, + scale, + use_qk_l2norm_in_kernel, + precision=config.precision, + gate_inputs=( + (gate_inputs[0][t], gate_inputs[1][t], *gate_inputs[2:]) + if gate_inputs is not None + else None + ), ) else: decay = gate.exp().unsqueeze(-1) @@ -404,9 +463,11 @@ def recurrent_decode( current_key / (current_key.square().sum(-1, keepdim=True) + 1e-6).sqrt() ) key = EncodedLinearAttentionTensor(current_key, None, "identity", None) - decayed = state * decay - if native_replay and gate.ndim == 2: - decayed = _replay_state(carry, gate) + decayed = ( + _replay_state(carry, gate) + if gate.ndim == 2 + else carry.reconstruct(original_basis=False) * decay + ) residual = v[t] - _sum_keys(key.values.unsqueeze(-1) * decayed) update = EncodedLinearAttentionTensor( beta[t].unsqueeze(-1) * residual, None, "identity", None diff --git a/modelopt/torch/quantization/linear_attention/gdn.py b/modelopt/torch/quantization/linear_attention/gdn.py index 7b73724dec7..09f89efe0ef 100644 --- a/modelopt/torch/quantization/linear_attention/gdn.py +++ b/modelopt/torch/quantization/linear_attention/gdn.py @@ -17,8 +17,9 @@ from .config import LinearAttentionConfig from .training import _prefill_decode_forward, _prepare_prefill_inputs +from .utils import _validate_gate_inputs -__all__ = ["gdn_state_qat", "matmul_gdn"] +__all__ = ["gdn_state_qat"] def gdn_state_qat( @@ -47,21 +48,28 @@ def gdn_state_qat( chunk_size=64, cp_context=None, prefill_lengths=None, - replay_gate_inputs=None, + sequence_lengths=None, + gate_inputs=None, ): """Adapt Megatron's FLA-style GDN call to serving-aligned state QAT. Validate prepared scalar log gates and promote inputs to FP32 working values. The shared training forward runs the chunked prefix and recurrent suffix with - native forward values, configured state QDQ, and a differentiable Torch adjoint. + native forward values and configured state QDQ. FLA supplies the prefix backward; + a Torch recurrence supplies the suffix adjoint. """ + if policy.precision == "vllm_kimi_k3": + raise ValueError("vllm_kimi_k3 is a KDA-only precision profile") if cp_context is not None: raise NotImplementedError("GDN state QAT does not support context parallelism") + _validate_gate_inputs(g, beta, gate_inputs) output_dtype = q.dtype beta_dtype = beta.dtype q, k, v, g, beta = _prepare_prefill_inputs( q, k, v, g, beta, policy=policy, chunk_size=chunk_size ) + if allow_neg_eigval: + raise ValueError("Serving GDN does not support allow_neg_eigval") if use_gate_in_kernel or use_beta_sigmoid_in_kernel: raise ValueError( "Serving GDN expects prepared log gates and beta from the Megatron adapter" @@ -87,10 +95,7 @@ def gdn_state_qat( output_dtype=output_dtype, beta_dtype=beta_dtype, prefill_lengths=prefill_lengths, - replay_gate_inputs=replay_gate_inputs, + sequence_lengths=sequence_lengths, + gate_inputs=gate_inputs, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ) - - -# Compatibility name for callers using the original adapter API. -matmul_gdn = gdn_state_qat diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py index f2238c2c7ad..f9653decf12 100644 --- a/modelopt/torch/quantization/linear_attention/kda.py +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -19,9 +19,9 @@ import torch.nn.functional as F from .training import _prefill_decode_forward, _prepare_prefill_inputs -from .utils import forward_value +from .utils import _validate_gate_inputs, forward_value -__all__ = ["kda_state_qat", "matmul_kda"] +__all__ = ["kda_state_qat"] def kda_state_qat( @@ -54,21 +54,25 @@ def kda_state_qat( disable_recompute=False, return_intermediate_states=False, prefill_lengths=None, + sequence_lengths=None, + gate_inputs=None, ): """Adapt Megatron's FLA-style KDA call to serving-aligned state QAT. Prepare per-key-channel gates and optional beta activation using FLA's formula. The shared training forward runs the chunked prefix and recurrent suffix with - native forward values, configured state QDQ, and a differentiable Torch adjoint. + native forward values and configured state QDQ. FLA supplies the prefix backward; + a Torch recurrence supplies the suffix adjoint. """ if cp_context is not None or disable_recompute or return_intermediate_states: raise NotImplementedError( "KDA state QAT 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 allow_neg_eigval and (not use_beta_sigmoid_in_kernel or gate_inputs is not None): + raise ValueError("allow_neg_eigval requires sigmoid beta without raw gate_inputs") if lower_bound is not None or safe_gate: raise ValueError("Serving arithmetic uses the native softplus KDA gate") + _validate_gate_inputs(g, beta, gate_inputs) output_dtype = q.dtype beta_dtype = beta.dtype q, k, v, g, beta = _prepare_prefill_inputs( @@ -86,15 +90,18 @@ def kda_state_qat( rate = A_log.to(dtype).exp().reshape(g.shape[-2], 1) g = -rate * F.softplus(g) # Import the optional vLLM backend only for the native precision profile. - from ...kernels.quantization.linear_attention.serving.forward import fused_kda_gate + from ...kernels.quantization.linear_attention.serving._compat import fla_module with torch.no_grad(): - native_gate = fused_kda_gate( + native_gate = fla_module("kda").fused_kda_gate( raw_gate.flatten(-2).contiguous(), A_log, raw_gate.shape[-1], g_bias=dt_bias ) g = forward_value(g, native_gate) if use_beta_sigmoid_in_kernel: + if use_gate_in_kernel and gate_inputs is None and not allow_neg_eigval: + gate_inputs = (raw_gate, beta, A_log, dt_bias) beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) + beta_dtype = torch.float32 return _prefill_decode_forward( q, k, @@ -114,9 +121,7 @@ def kda_state_qat( output_dtype=output_dtype, beta_dtype=beta_dtype, prefill_lengths=prefill_lengths, + sequence_lengths=sequence_lengths, + gate_inputs=gate_inputs, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ) - - -# Compatibility name for callers using the original adapter API. -matmul_kda = kda_state_qat diff --git a/modelopt/torch/quantization/linear_attention/training.py b/modelopt/torch/quantization/linear_attention/training.py index ed8401d5107..9d9ab9c1d3e 100644 --- a/modelopt/torch/quantization/linear_attention/training.py +++ b/modelopt/torch/quantization/linear_attention/training.py @@ -17,54 +17,125 @@ from contextlib import contextmanager from itertools import pairwise +from numbers import Integral import torch -import modelopt.torch.utils.distributed as dist - -from .decode import recurrent_decode +from .decode import _recurrent_decode from .utils import _resolve_state_quantizer, _state_qdq, forward_value -__all__ = ["linear_attention_training_phase"] +__all__ = ["get_linear_attention_layers", "linear_attention_training_phase"] -def _lengths(values): +def _lengths(values, name): 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 + if any( + isinstance(value, bool) or not isinstance(value, Integral) or value < 0 for value in values + ): + raise ValueError(f"{name} must be nonnegative integers") + return tuple(int(value) for value in values) -@contextmanager -def linear_attention_training_phase(model, prefill_lengths): - """Supply explicit sequence phases through forward and checkpointed backward. +class _PackedBoundaries: + """Own immutable host boundaries and verify each layer's current packed tensor.""" + + def __init__(self, values): + self.values = None if values is None else _lengths(values, "Packed boundaries") - Runtime phase metadata is local to the converted layers and restored on exit. - Keep this context active through backward when activation checkpointing is used. + def resolve(self, actual): + if actual is None: + if self.values is not None: + raise ValueError("Phase cu_seqlens requires packed boundaries from the layer") + return None + # Tensor versions do not track writes through .data, NumPy or DLPack aliases. + values = _lengths(actual, "Packed boundaries") + if self.values is not None and values != self.values: + raise ValueError("Phase cu_seqlens does not match the layer's packed boundaries") + self.values = values + return values + + +def get_linear_attention_layers(model, *, enabled_only=False, unphased_only=False): + """Return serving-policy layers from a module or iterable of model chunks. + + Disabled state quantizers are included by default so a baseline can use the + same serving arithmetic and phase as QAT. Use ``enabled_only`` for setup checks + and ``unphased_only`` to select layers needing adapter-provided defaults. """ # 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.backend == "serving" + roots = (model,) if isinstance(model, torch.nn.Module) else model + return tuple( + dict.fromkeys( + module + for root in roots + for module in root.modules() + if isinstance(module, _LinearAttentionQuantMixin) + and module.linear_attention_config.backend == "serving" + and (not enabled_only or module._linear_attn_state.is_enabled) + and (not unphased_only or module._linear_attention_prefill_lengths is None) + ) + ) + + +def _prefix_lengths(values, valid_lengths): + prefixes = _lengths(values, "Prefill lengths") + if valid_lengths is not None and ( + len(prefixes) != len(valid_lengths) + or any(prefix > end for prefix, end in zip(prefixes, valid_lengths)) + ): + raise ValueError("Supply one prefill length <= sequence length per sequence") + return prefixes + + +@contextmanager +def linear_attention_training_phase( + model, prefill_lengths, *, sequence_lengths=None, cu_seqlens=None, preserve_existing=False +): + """Supply per-sequence phases to a module or iterable of model chunks/layers. + + Enabled state quantization requires a phase, including calibration/eval. Disabled + quantizers use the original module outside a phase and serving arithmetic inside. + Megatron selective recompute captures the phase for backward after context exit. + + ``sequence_lengths`` excludes right padding in dense rows or padded THD segments. + ``cu_seqlens`` describes physical packed storage, including alignment padding. + Megatron derives omitted valid lengths from unpadded ``cu_seqlens_q`` metadata; + padded-only metadata requires explicit ``sequence_lengths``. + Padding produces zero attention outputs and never updates the recurrent state. + ``preserve_existing=True`` installs defaults only on layers without a phase. + Empty pipeline stages participate as no-ops. + """ + layers = get_linear_attention_layers(model) + if preserve_existing: + layers = tuple(m for m in layers if m._linear_attention_prefill_lengths is None) + valid_lengths = ( + None if sequence_lengths is None else _lengths(sequence_lengths, "Sequence lengths") + ) + lengths = _prefix_lengths(prefill_lengths, valid_lengths) + boundaries = _PackedBoundaries(cu_seqlens) + previous = [ + ( + m._linear_attention_prefill_lengths, + m._linear_attention_sequence_lengths, + m._linear_attention_cu_seqlens, + ) + for m in layers ] - # A pipeline stage may have no selected layers; conversion validates global matches. - if not layers and dist.size() == 1: - 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 + module._linear_attention_sequence_lengths = valid_lengths + module._linear_attention_cu_seqlens = boundaries yield model finally: - for module, original in zip(layers, previous): + for module, (original, original_lengths, original_boundaries) in zip(layers, previous): module._linear_attention_prefill_lengths = original + module._linear_attention_sequence_lengths = original_lengths + module._linear_attention_cu_seqlens = original_boundaries def _prepare_prefill_inputs(q, k, v, g, beta, *, policy, chunk_size): @@ -96,8 +167,9 @@ def _prefill_decode_forward( output_dtype, beta_dtype, prefill_lengths, + sequence_lengths=None, use_qk_l2norm_in_kernel=False, - replay_gate_inputs=None, + gate_inputs=None, ): """Run both prefill and decode phases in one differentiable training forward. @@ -109,12 +181,17 @@ def _prefill_decode_forward( Args: prefill_lengths: Prefix token count per sequence. For 128 tokens, a value of 64 selects 64 chunked prefill tokens followed by 64 recurrent tokens. + sequence_lengths: Valid token count per sequence, excluding right padding. + Defaults to the dense row length or each packed storage segment length. """ state_quantizer, state_qdq, state_format = _resolve_state_quantizer( state_quantizer, state_qdq, state_format ) - if prefill_lengths is None: - raise ValueError("Decode-aware training requires explicit per-sequence prefill lengths") + if state_qdq and prefill_lengths is None: + raise ValueError( + "State quantization requires explicit prefill lengths through " + "linear_attention_training_phase" + ) 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 @@ -124,11 +201,16 @@ def _prefill_decode_forward( 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 isinstance(cu_seqlens_cpu, _PackedBoundaries): + boundaries = cu_seqlens_cpu.resolve(cu_seqlens) + elif cu_seqlens_cpu is not None: + boundaries = _PackedBoundaries(cu_seqlens_cpu).resolve(cu_seqlens) + else: + boundaries = cu_seqlens if boundaries is None: sequences = [(b, 0, length) for b in range(batch)] else: - bounds = _lengths(boundaries) + bounds = _lengths(boundaries, "Packed boundaries") if ( batch != 1 or len(bounds) < 2 @@ -140,11 +222,18 @@ def _prefill_decode_forward( "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) + valid_lengths = ( + tuple(end - start for _, start, end in sequences) + if sequence_lengths is None + else _lengths(sequence_lengths, "Sequence lengths") + ) + if len(valid_lengths) != len(sequences) or any( + valid > end - start for valid, (_, start, end) in zip(valid_lengths, sequences) ): - raise ValueError("Supply one valid prefill length per sequence") + raise ValueError("Supply one sequence length within its storage segment per sequence") + prefixes = _prefix_lengths( + valid_lengths if prefill_lengths is None else prefill_lengths, valid_lengths + ) if ( keys > 256 or (g.ndim == 4 and values != keys) @@ -153,6 +242,15 @@ def _prefill_decode_forward( raise ValueError( "Serving arithmetic requires K <= 256, KDA V=K, and power-of-two K for normalization" ) + if policy.precision != "replayssm": + # Validate the optional native backend before launching any prefix kernels. + from ...kernels.quantization.linear_attention.serving.forward import validate_profile + + validate_profile( + gate_inputs, + use_qk_l2norm_in_kernel, + kimi_k3=policy.precision == "vllm_kimi_k3", + ) if initial_state is None: states = q.new_zeros(len(sequences), heads, keys, values) else: @@ -161,7 +259,8 @@ def _prefill_decode_forward( if states.shape != (len(sequences), heads, keys, values): raise ValueError("Initial state shape does not match sequence/head dimensions") outputs, finals = [], [] - for n, (b, start, end) in enumerate(sequences): + for n, (b, start, storage_end) in enumerate(sequences): + end = start + valid_lengths[n] split = start + prefixes[n] prefix, state = q.new_empty(0, heads, values), states[n] if prefixes[n]: @@ -194,11 +293,21 @@ def _prefill_decode_forward( state=state, scale=keys**-0.5 if scale is None else scale, beta_dtype=beta_dtype, + precision=policy.precision, normalize=use_qk_l2norm_in_kernel, + gate_inputs=( + ( + gate_inputs[0][b, start:split], + gate_inputs[1][b, start:split], + *gate_inputs[2:], + ) + if gate_inputs is not None + else None + ), ) # Keep the prefix state attached so suffix losses backpropagate through prefill. # recurrent_decode applies configured state QDQ at the handoff and suffix writes. - suffix, carry = recurrent_decode( + suffix, carry = _recurrent_decode( *(x[b, split:end] for x in (q, k, v, g, beta)), config=policy, state_qdq=state_qdq, @@ -206,22 +315,27 @@ def _prefill_decode_forward( state_quantizer=state_quantizer, initial_state=state, position=prefixes[n], - replay_gate_inputs=( + gate_inputs=( ( - replay_gate_inputs[0][b, split:end], - replay_gate_inputs[1][b, split:end], - *replay_gate_inputs[2:], + gate_inputs[0][b, split:end], + gate_inputs[1][b, split:end], + *gate_inputs[2:], ) - if replay_gate_inputs is not None + if gate_inputs is not None else None ), use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, scale=scale, ) - outputs.append(torch.cat((prefix, suffix))) - finals.append(carry.reconstruct()) + # Keep storage offsets for downstream layers while excluding padding from the graph. + pieces = (prefix, suffix) + if storage_end > end: + pieces += (v.new_zeros(storage_end - end, heads, values),) + outputs.append(torch.cat(pieces)) + if output_final_state: + 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 = torch.stack(finals) if output_final_state else None + if final is not None and state_v_first: final = final.transpose(-1, -2) - return output.to(output_dtype), final if output_final_state else None + return output.to(output_dtype), final diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py index afb40b73c5c..133649a0e04 100644 --- a/modelopt/torch/quantization/linear_attention/utils.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -44,6 +44,32 @@ def forward_value(value, rounded): return _ForwardValue.apply(value, rounded) +def _validate_gate_inputs(g, beta, gate_inputs): + """Check raw projection shapes and types before dispatching native kernels.""" + if gate_inputs is None: + return + if not isinstance(gate_inputs, (tuple, list)) or len(gate_inputs) != 4: + raise ValueError("gate_inputs must contain raw gate, raw beta, A_log, and bias") + raw_g, raw_beta, rate, bias = gate_inputs + for value in (raw_g, raw_beta, rate, bias): + if value is not None and ( + not isinstance(value, torch.Tensor) + or value.device != g.device + or value.dtype not in (torch.bfloat16, torch.float32) + ): + raise ValueError("Raw gate inputs must be BF16/FP32 tensors on the gate device") + if raw_g is None or raw_beta is None or raw_g.shape != g.shape or raw_beta.shape != beta.shape: + raise ValueError("Raw gate and beta shapes must match their prepared operands") + heads = beta.shape[-1] + channels = g.shape[-1] if g.ndim == beta.ndim + 1 else 1 + if rate is None or rate.numel() != heads or rate.ndim not in (1, 2): + raise ValueError("Raw gate A_log must contain one rate per value head") + if (bias is None and g.ndim == beta.ndim) or ( + bias is not None and (bias.ndim != 1 or bias.numel() != heads * channels) + ): + raise ValueError("Raw gate bias must contain one value per gate channel") + + def validate_gdn_quantizer( quantizer: TensorQuantizer, *, @@ -137,15 +163,19 @@ def _state_qdq( state_format: str = "fp8_e4m3", state_quantizer: TensorQuantizer | None = None, ): - """Dynamic state-tile QDQ with detached scales and identity STE.""" + """Apply the same state encoding used by recurrent writes.""" + return _state_qdq_with_scales(state, block_v, state_format, state_quantizer)[0] + + +def _state_qdq_with_scales(state, block_v, state_format, state_quantizer=None): + """Shared continuation/decode QDQ; blockwise scales remain owned by TensorQuantizer.""" if state_quantizer is not None and state_quantizer.block_sizes is not None: - return state_quantizer(state) + return state_quantizer(state), None 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") - quantized, _ = _tile_qdq(state, block_v, state_format, state_quantizer=state_quantizer) - return quantized + return _tile_qdq(state, block_v, state_format, state_quantizer=state_quantizer) def _tile_qdq(value, block_v, state_format, *, state_quantizer=None): diff --git a/modelopt/torch/quantization/plugins/gdn.py b/modelopt/torch/quantization/plugins/gdn.py index d1cdc04b7c5..91dfc33c4dd 100644 --- a/modelopt/torch/quantization/plugins/gdn.py +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -59,6 +59,8 @@ def validate_linear_attention(self): "GDN W quantization is no longer supported. Disable gdn_w_quantizer; " "use gdn_state_quantizer with an explicit serving policy for state QAT." ) + if self.linear_attention_config.precision == "vllm_kimi_k3": + raise ValueError("vllm_kimi_k3 is a KDA-only precision profile") super().validate_linear_attention() @property @@ -70,7 +72,11 @@ def _state_quantized_chunk_gated_delta_rule( self, gated_delta_rule: GatedDeltaRuleFn, *args: Any, **kwargs: Any ) -> tuple[torch.Tensor, torch.Tensor | None]: """Route enabled state quantization to its training implementation.""" - self.validate_linear_attention() + if ( + self.gdn_w_quantizer.is_enabled + or self.linear_attention_config.precision == "vllm_kimi_k3" + ): + self.validate_linear_attention() if not self.linear_attention_is_enabled: return gated_delta_rule(*args, **kwargs) while isinstance(gated_delta_rule, partial): @@ -85,11 +91,14 @@ def _state_quantized_chunk_gated_delta_rule( chunk_size = kwargs.pop("chunk_size", 64) if chunk_size != 64: raise ValueError("GDN fake quantization supports only chunk_size=64") + if self._linear_attention_cu_seqlens is not None: + kwargs["cu_seqlens_cpu"] = self._linear_attention_cu_seqlens return gdn_state_qat( *args, policy=self.linear_attention_config, state_quantizer=self.gdn_state_quantizer, chunk_size=chunk_size, prefill_lengths=self._linear_attention_prefill_lengths, + sequence_lengths=self._linear_attention_sequence_lengths, **kwargs, ) diff --git a/modelopt/torch/quantization/plugins/kda.py b/modelopt/torch/quantization/plugins/kda.py index 04766a5d8e9..9c9acce694b 100644 --- a/modelopt/torch/quantization/plugins/kda.py +++ b/modelopt/torch/quantization/plugins/kda.py @@ -29,7 +29,6 @@ class KimiDeltaAttentionStateQuantMixin(_LinearAttentionQuantMixin): linear_attention_quantizer_names = ("kda_state_quantizer",) def _state_quantized_chunk_kda(self, kernel, *args, **kwargs): - self.validate_linear_attention() if not self.linear_attention_is_enabled: return kernel(*args, **kwargs) # FLA kernels are an optional dependency; the training layer belongs to Megatron. @@ -41,10 +40,13 @@ def _state_quantized_chunk_kda(self, kernel, *args, **kwargs): kernel = kernel.func if kernel is not chunk_kda: raise NotImplementedError("KDA quantization requires FLA's chunk_kda callable") + if self._linear_attention_cu_seqlens is not None: + kwargs["cu_seqlens_cpu"] = self._linear_attention_cu_seqlens return kda_state_qat( *args, policy=self.linear_attention_config, state_quantizer=self.kda_state_quantizer, prefill_lengths=self._linear_attention_prefill_lengths, + sequence_lengths=self._linear_attention_sequence_lengths, **kwargs, ) diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py index d5945b57bab..4e651acbfd6 100644 --- a/modelopt/torch/quantization/plugins/linear_attention.py +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -16,10 +16,8 @@ """Shared module policy and checkpoint support for linear-attention QAT.""" import fnmatch +import warnings -import torch - -import modelopt.torch.utils.distributed as dist from modelopt.torch.opt.conversion import ApplyModeError from modelopt.torch.utils import get_unwrapped_name @@ -40,28 +38,22 @@ def _linear_attention_modules(model): def _apply_linear_attention_policy(model, config): - """Assign complete policies in rule order, before compatibility validation.""" + """Apply rules locally, like quant_cfg; a PP/VPP chunk may have no matches.""" modules = _linear_attention_modules(model) # getattr also handles pickled configs that predate the policy field. for entry in getattr(config, "linear_attention", []): matches = [name for name in modules if fnmatch.fnmatch(name, entry.module_name)] - matched_on_rank = [bool(matches)] - if dist.size() > 1: - # Shared recipes can select layers owned by another pipeline stage. - matched_on_rank = [False] * dist.size() - torch.distributed.all_gather_object(matched_on_rank, bool(matches)) - if not any(matched_on_rank): - 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) def _validate_linear_attention(model): """Validate configured modules after both policy and quantizers are assigned.""" - for module in _linear_attention_modules(model).values(): + modules = _linear_attention_modules(model).values() + for module in modules: module.validate_linear_attention() + module.warn_linear_attention_profile() + module._validate_linear_attention_config() def _linear_attention_state(model): @@ -102,6 +94,8 @@ def _setup(self): ) self._register_temp_attribute("linear_attention_config", LinearAttentionConfig()) self._register_temp_attribute("_linear_attention_prefill_lengths", None) + self._register_temp_attribute("_linear_attention_sequence_lengths", None) + self._register_temp_attribute("_linear_attention_cu_seqlens", None) @property def _linear_attn_state(self): @@ -109,10 +103,12 @@ def _linear_attn_state(self): @property def linear_attention_is_enabled(self): - """Whether state quantization or the serving arithmetic policy is enabled.""" - return ( - any(getattr(self, name).is_enabled for name in self.linear_attention_quantizer_names) - or self.linear_attention_config.backend == "serving" + """Retain an unquantized serving baseline only inside an explicit phase.""" + return any( + getattr(self, name).is_enabled for name in self.linear_attention_quantizer_names + ) or ( + self.linear_attention_config.backend == "serving" + and self._linear_attention_prefill_lengths is not None ) def validate_linear_attention(self): @@ -124,8 +120,8 @@ def validate_linear_attention(self): ) if self.linear_attention_config.backend != "serving": raise ValueError( - "GDN/KDA state QAT requires backend='serving' and prefill lengths " - "through linear_attention_training_phase" + "GDN/KDA state QAT requires backend='serving'; use " + "linear_attention_training_phase to select a recurrent training suffix" ) if self.linear_attention_config.state_codec == "int8_hadamard32": if state_format != "int8": @@ -133,7 +129,35 @@ def validate_linear_attention(self): if group_size: raise ValueError("TensorQuantizer block_sizes requires state_codec='tile'") + def warn_linear_attention_profile(self): + """Warn at conversion/restore on one TP/DP representative of each owning stage.""" + if ( + not self._linear_attn_state.is_enabled + or "kda_state_quantizer" not in self.linear_attention_quantizer_names + or self.linear_attention_config.precision != "vllm" + ): + return + parallel = getattr(self, "parallel_state", None) + if parallel is not None and any( + group.rank() > 0 + for group in (parallel.tensor_parallel_group, parallel.data_parallel_group) + ): + return + warnings.warn( + "KDA precision='vllm' selects standalone FLA/Triton kernels, not a model " + "serving profile. For Kimi-Linear/Kimi-K3 on vLLM 0.30, select " + "precision='vllm_kimi_k3' and configure the server's matching Triton " + "backends and unbounded gates." + ) + + def _validate_linear_attention_config(self): + """Check framework settings at conversion/forward, without restricting offline restore.""" + + def validate_linear_attention_execution(self): + """Check framework execution restrictions at forward, not conversion or restore.""" + def modelopt_post_restore(self, prefix=""): """Validate the restored numerical policy and quantizers.""" super().modelopt_post_restore(prefix) self.validate_linear_attention() + self.warn_linear_attention_profile() diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 63c6acc029e..7e8b5df29e7 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -21,7 +21,9 @@ import textwrap import types from contextlib import contextmanager +from copy import copy from functools import cache, partial +from itertools import pairwise from typing import Any import megatron.core.parallel_state as mcore_parallel @@ -29,6 +31,7 @@ import megatron.core.transformer.mlp as megatron_mlp import megatron.core.transformer.moe.experts as megatron_moe import torch +from megatron.core import tensor_parallel from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding from megatron.core.extensions.transformer_engine import ( TEColumnParallelGroupedLinear, @@ -59,6 +62,7 @@ from ..algorithms import AutoQuantizeGradientSearcher from ..conversion import maybe_promote_nvfp4_static_quantizer from ..linear_attention.config import LinearAttentionConfig +from ..linear_attention.training import _lengths, linear_attention_training_phase from ..nn import ( GroupedQuantizer, QuantModule, @@ -1137,7 +1141,8 @@ def set_extra_state(self, state): def _setup(self): super()._setup() - self._register_temp_attribute("_linear_attention_replay_gate_inputs", None) + self._register_temp_attribute("_linear_attention_gate_inputs", None) + self._register_temp_attribute("_linear_attention_recompute", False) try: data_parallel_group = get_data_parallel_group(with_context_parallel=True) except AssertionError: @@ -1149,7 +1154,9 @@ def _setup(self): @property def _serving_arithmetic(self): - return self.linear_attention_config.backend == "serving" + return ( + self.linear_attention_config.backend == "serving" and self.linear_attention_is_enabled + ) def _prepare_input_for_gated_delta_rule(self, *args, **kwargs): # Preserve raw BF16 Q/K: prefill stores normalized BF16 operands, whereas @@ -1166,46 +1173,133 @@ def _prepare_input_for_gated_delta_rule(self, *args, **kwargs): @contextmanager def _quantized_linear_attention_kernel(self): kernel = self.gated_delta_rule - previous_gates = self._linear_attention_replay_gate_inputs + previous_gates = self._linear_attention_gate_inputs if self._serving_arithmetic: def run_kernel(*args, **kwargs): kwargs["use_qk_l2norm_in_kernel"] = self.use_qk_l2norm - if self._linear_attention_replay_gate_inputs is not None: - kwargs["replay_gate_inputs"] = self._linear_attention_replay_gate_inputs + if self._linear_attention_gate_inputs is not None: + kwargs["gate_inputs"] = self._linear_attention_gate_inputs return self._linear_attention_kernel(kernel, *args, **kwargs) self.gated_delta_rule = run_kernel else: self.gated_delta_rule = partial(self._linear_attention_kernel, kernel) - self._linear_attention_replay_gate_inputs = None + self._linear_attention_gate_inputs = None try: yield finally: self.gated_delta_rule = kernel - self._linear_attention_replay_gate_inputs = previous_gates + self._linear_attention_gate_inputs = previous_gates def forward(self, *args, **kwargs): - # Newer Megatron versions recompute the core independently during backward. - if hasattr(super(), "_forward_compute") or hasattr( - super(), "forward_pre_attn_and_core_attn" - ): - return super().forward(*args, **kwargs) - with self._quantized_linear_attention_kernel(): - return super().forward(*args, **kwargs) + self.validate_linear_attention_execution() + recompute = getattr(self, "recompute_gdn", False) + capture = self._serving_arithmetic and self.training and recompute + if capture and not hasattr(super(), "_forward_compute"): + raise NotImplementedError("Selective state QAT recompute requires Megatron's core hook") + previous = self._linear_attention_recompute + self._linear_attention_recompute = capture + if capture: + # The adapter checkpoints the same core, with microbatch-local metadata. + self.recompute_gdn = False + try: + if hasattr(super(), "_forward_compute") or hasattr( + super(), "forward_pre_attn_and_core_attn" + ): + return super().forward(*args, **kwargs) + with self._quantized_linear_attention_kernel(): + return super().forward(*args, **kwargs) + finally: + self._linear_attention_recompute = previous + if capture: + self.recompute_gdn = recompute + + def _forward_compute(self, hidden_states, *args, **kwargs): + compute = super()._forward_compute + if not self._serving_arithmetic: + with self._quantized_linear_attention_kernel(): + return compute(hidden_states, *args, **kwargs) + lengths = self._linear_attention_prefill_lengths + sequence_lengths = self._linear_attention_sequence_lengths + if lengths is None: + raise ValueError("State quantization requires explicit prefill lengths") + arguments = inspect.signature(compute).bind(hidden_states, *args, **kwargs) + boundaries = self._linear_attention_cu_seqlens.resolve( + arguments.arguments.get("cu_seqlens_q") + ) + packed = arguments.arguments.get("packed_seq_params") + if packed is not None and packed.qkv_format == "thd": + actual = packed.cu_seqlens_q + if actual is None: + if sequence_lengths is None: + raise ValueError( + "Packed state QAT requires valid sequence_lengths or cu_seqlens_q" + ) + else: + real_bounds = _lengths(actual, "Packed boundaries") + valid = tuple(end - start for start, end in pairwise(real_bounds)) + if sequence_lengths is None: + sequence_lengths = valid + elif len(sequence_lengths) != len(valid) or any( + length > end for length, end in zip(sequence_lengths, valid) + ): + raise ValueError("Sequence lengths exceed the packed valid token counts") + if self._linear_attention_recompute: + # Packed metadata may refer to buffers reused by a later pipeline microbatch. + for name, value in list(arguments.arguments.items()): + if name == "cu_seqlens_q" and value is not None: + arguments.arguments[name] = value.clone() + elif name == "packed_seq_params" and value is not None: + snapshot = copy(value) + for key, tensor in vars(value).items(): + if isinstance(tensor, torch.Tensor): + setattr(snapshot, key, tensor.clone()) + arguments.arguments[name] = snapshot + + # Megatron checkpoint preserves RNG, but does not capture PyTorch autocast. + autocast_enabled = torch.is_autocast_enabled("cuda") + autocast_dtype = torch.get_autocast_dtype("cuda") + + def run(hidden): + with ( + torch.autocast("cuda", enabled=autocast_enabled, dtype=autocast_dtype), + linear_attention_training_phase( + self, lengths, sequence_lengths=sequence_lengths, cu_seqlens=boundaries + ), + self._quantized_linear_attention_kernel(), + ): + return compute(hidden, *arguments.args[1:], **arguments.kwargs) - def _forward_compute(self, *args, **kwargs): - with self._quantized_linear_attention_kernel(): - return super()._forward_compute(*args, **kwargs) + if self._linear_attention_recompute: + return tensor_parallel.checkpoint(run, False, hidden_states) + return run(hidden_states) def forward_pre_attn_and_core_attn(self, *args, **kwargs): with self._quantized_linear_attention_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: + def validate_linear_attention_execution(self): + super().validate_linear_attention_execution() + self._validate_linear_attention_config() + active = self.linear_attention_is_enabled + serving = self.linear_attention_config.backend == "serving" and active + if self.config.context_parallel_size > 1 and active: raise NotImplementedError("GDN/KDA QAT does not support Megatron context parallelism.") + if serving and self.training and self.config.recompute_granularity == "full": + raise NotImplementedError( + "State QAT supports selective GDN recompute, not full-layer recompute" + ) + if ( + serving + and getattr(self, "recompute_gdn", False) + and not hasattr(super(), "_forward_compute") + ): + raise NotImplementedError("Selective state QAT recompute requires Megatron's core hook") + + def _validate_linear_attention_config(self): + if self._serving_arithmetic and self.config.deterministic_mode: + raise NotImplementedError("Serving state QAT requires non-deterministic kernels") if self._serving_arithmetic and ( not hasattr(super(), "_prepare_input_for_gated_delta_rule") or getattr(self, "gdn_pre_gated_delta_rule_fusion", False) @@ -1228,26 +1322,14 @@ class _QuantGatedDeltaNet(_MegatronLinearAttentionMixin, GatedDeltaNetStateQuant def _compute_gates(self, a_log, dt_bias, batch, seq_len, *gate_feats): gate, inputs = super()._compute_gates(a_log, dt_bias, batch, seq_len, *gate_feats) if self._serving_arithmetic: - # Import the optional vLLM backend only for the native precision profile. - from ...kernels.quantization.linear_attention.serving.forward import ( - fused_gdn_gating, - ) - from ..linear_attention.utils import forward_value - raw_beta, raw_gate = gate_feats - if self.linear_attention_config.precision == "replayssm": - self._linear_attention_replay_gate_inputs = (raw_gate, raw_beta, a_log, dt_bias) - with torch.no_grad(): - native_gate, native_beta = fused_gdn_gating( - a_log, - raw_gate.reshape(-1, raw_gate.shape[-1]).contiguous(), - raw_beta.reshape(-1, raw_beta.shape[-1]).contiguous(), - dt_bias, - ) - gate = forward_value(gate, native_gate.reshape_as(gate)) - inputs["beta"] = forward_value( - inputs["beta"], native_beta.reshape_as(inputs["beta"]) - ).to(raw_beta.dtype) + # Native prefill/decode own gate activation and its phase-specific rounding. + self._linear_attention_gate_inputs = ( + raw_gate.reshape(batch, seq_len, -1), + raw_beta.reshape(batch, seq_len, -1), + a_log, + dt_bias, + ) return gate, inputs @@ -1261,6 +1343,26 @@ class _QuantKimiDeltaAttention( _linear_attention_kernel = KimiDeltaAttentionStateQuantMixin._state_quantized_chunk_kda + def _validate_linear_attention_config(self): + super()._validate_linear_attention_config() + if self._serving_arithmetic and ( + getattr(self.config, "kda_safe_gate", False) + or getattr(self.config, "kda_lower_bound", None) is not None + ): + raise NotImplementedError("Serving KDA state QAT requires unbounded softplus gates") + + def _compute_gates(self, a_log, dt_bias, batch, seq_len, *gate_feats): + gate, inputs = super()._compute_gates(a_log, dt_bias, batch, seq_len, *gate_feats) + if self._serving_arithmetic: + raw_gate, raw_beta = gate_feats + self._linear_attention_gate_inputs = ( + raw_gate.reshape_as(gate), + raw_beta.reshape_as(inputs["beta"]), + a_log, + dt_bias, + ) + return gate, inputs + def _is_supported_megatron_model(model: torch.nn.Module) -> bool: return isinstance(model, MegatronModule) diff --git a/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml index 540f48a978a..f5539e01c05 100644 --- a/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml +++ b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml @@ -21,7 +21,10 @@ imports: metadata: description: >- Dynamic INT8 GDN/KDA state QDQ with per-key 32-value scales and working-state readout. - Emulates vLLM 0.15 BF16 prefill arithmetic without additional state QDQ. + For Kimi-K3/Kimi-Linear, use the explicit kimi_k3 recipe instead. + Uses the installed vLLM GDN or standalone FLA/KDA Triton backend with BF16 operands + and FP32 state storage. + Fresh prefill has no internal state QDQ; handoff and recurrent writes use the configured QDQ. Supply per-sequence prefix lengths through linear_attention_training_phase. No calibration is required. quantize: diff --git a/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic_kimi_k3.yaml b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic_kimi_k3.yaml new file mode 100644 index 00000000000..7801afbee37 --- /dev/null +++ b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic_kimi_k3.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. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + base_disable_all: configs/ptq/units/base_disable_all + state_int8: configs/ptq/units/linear_attention_state_int8_block32_dynamic + +metadata: + description: >- + Dynamic INT8 Kimi-K3/Kimi-Linear KDA state QDQ with per-key 32-value scales and working-state readout. + Requires raw projections and unbounded softplus gates. In vLLM, set both + kda_prefill_backend and kda_decode_backend to triton; mixed batches are not covered. + Fresh prefill has no internal state QDQ; handoff and recurrent writes use the configured QDQ. + Supply per-sequence prefix lengths through + linear_attention_training_phase. No calibration is required. +quantize: + algorithm: + quant_cfg: + - $import: base_disable_all + - $import: state_int8 + linear_attention: + - module_name: '*' + cfg: + backend: serving + precision: vllm_kimi_k3 diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 25a7c9df0c9..bd591e69a6c 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 34 general/ptq/ recipes (click to expand) +All 35 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -66,11 +66,75 @@ supported combinations. | `q8_0` | Q8_0 W8A16 (8.5 bpw), eligible linears | none | none (no calibration) | | `linear_attention_state_int8_dynamic` | GDN/KDA decode state INT8 + Hadamard; weights unchanged | none | none (dynamic scales; requires a prefix/decode phase context) | | `linear_attention_state_int8_block32_dynamic` | GDN/KDA state INT8, 32 value channels per key row; weights unchanged | none | none (dynamic scales; requires vLLM and a prefix/decode phase context) | +| `linear_attention_state_int8_block32_dynamic_kimi_k3` | Kimi-K3/Kimi-Linear KDA state INT8, 32 value channels per key row | none | none (dynamic scales; explicit Triton Kimi profile) |
The block32 linear-attention recipe selects `precision="vllm"` and working-state readout. It applies fake QDQ with floating-point state storage. +For GDN, select `gdn_prefill_backend="triton"` and +`VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE=1` where packed decode is available. +This profile covers pure, single-token, non-speculative decode. Mixed prefill/decode +batches can select different native kernels, so arbitrary continuous batching is +not numerically qualified by these tests. + +For KDA, `precision="vllm"` selects the standalone FLA/Triton API, without a +model-level serving parity claim; conversion emits a diagnostic for this choice. Use the +`linear_attention_state_int8_block32_dynamic_kimi_k3` recipe for Kimi-K3 and +Kimi-Linear on vLLM 0.30, which share model kernels (`precision="vllm_kimi_k3"`). Set vLLM's `additional_config` to +`{"kda_prefill_backend": "triton", "kda_decode_backend": "triton"}`. The profile +requires raw projections, Q/K normalization, unbounded softplus gates, and the +Kimi-K3 package dependencies (`gate_lower_bound=None` in the serving HF config). +It never falls back to prepared-only arithmetic. +KDA validates the selected cache paths at the actual head size once per device; +builds with inconsistent state readout are rejected before training. + +State QAT training requires explicit per-sequence prefix lengths. Megatron selective +GDN recompute captures each microbatch's phase and packed metadata, allowing backward +after its context exits. Full-layer recompute and CP remain unsupported while state +QAT is active; disabled quantizers outside a phase retain the original module path. +Framework training, evaluation and generation loops must establish this context; +automatic Bridge and evaluation-helper integration is a separate change. Dynamic +state quantizers need no calibration. `examples/megatron_bridge/quantize.py` rejects +linear-attention recipes that require calibration before loading a model. To combine +calibrated projection quantization with state QDQ, provide `mtq.quantize` with a +calibration loop inside `linear_attention_training_phase`; a full-prefill phase +calibrates projections but does not measure decode state-QDQ quality. +To compare state QDQ, convert both models +with the same serving policy, disable the baseline's state quantizers, and use +identical prefix and valid-token lengths: + +```python +for model in (baseline, candidate): + mtq.quantize(model, state_recipe) +mtq.disable_quantizer(baseline, "*state_quantizer") +for model in (baseline, candidate): + with linear_attention_training_phase( + model, prefill_lengths, sequence_lengths=valid_lengths + ): + logits = model(tokens, position_ids, attention_mask) + # Score only valid continuation tokens, identically for both models. +``` + +`sequence_lengths` excludes right padding in dense rows or padded THD segments. +The recurrent suffix ends at that valid length; padding produces zero attention +outputs and does not update the cache. Omit valid lengths only when all storage +tokens are real. For Megatron THD input, omitted lengths are derived from the +unpadded `packed_seq_params.cu_seqlens_q`. If only padded boundaries are available, +pass `sequence_lengths` explicitly, even when every storage token is real. +Full-prefill scoring alone does not exercise decode state QDQ. + +The selected precision profile is saved in ModelOpt/Megatron checkpoints. +HF export currently does not carry this serving contract, and neither checkpoint +format configures or validates a remote vLLM process. Apply the backend and gate +settings above to the actual server; a matching trainer environment variable is no proof. +Server configuration and mixed-batch dispatch enforcement belong to the vLLM plugin. +The tests qualify the selected native kernels, not arbitrary server dispatch. + +Packed boundaries must match the layer's actual (including padding) boundaries. +Validation reads their values rather than trusting tensor versions, so it can synchronize +CUDA metadata. Prefix backward still runs a separate FLA recurrence as an adjoint; +this implementation is not an optimized fused training kernel. --- diff --git a/pyproject.toml b/pyproject.toml index b792592022b..cc95e6bca7c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -329,6 +329,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 index 1373b5ea46b..65ea183a318 100644 --- a/tests/_test_utils/torch/quantization/linear_attention_reference.py +++ b/tests/_test_utils/torch/quantization/linear_attention_reference.py @@ -20,40 +20,28 @@ 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_delta_rule_reference", - "recurrent_delta_rule_reference", - "state_qdq_reference", -] +__all__ = ["chunk_gdn_reference", "recurrent_delta_rule_reference", "state_fp8_qdq_reference"] -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") +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. The zero tile uses scale one; partial value tiles are valid. + """ 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": - # CUDA TensorQuantizer uses a quantization multiplier and zeros tiny groups. - tiny = amax < 2**-24 - quant_scale = 127.0 / torch.where(tiny, torch.ones_like(amax), amax) - codes = (tile * quant_scale).round().clamp(-127, 127) - decoded = torch.where(tiny, 0.0, codes / quant_scale) - 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() - decoded = codes * scale - rounded.append(decoded) + scale = torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)) + rounded.append((tile / scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() * scale) quantized = torch.cat(rounded, dim=-1).to(state.dtype) return state + (quantized - state).detach() @@ -138,7 +126,7 @@ def recurrent_delta_rule_reference( return output, final.transpose(-1, -2) if state_v_first else final -def chunk_gdn( +def chunk_gdn_reference( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -147,114 +135,69 @@ def chunk_gdn( *, chunk_size: int = 64, scale: float | None = None, - initial_state: torch.Tensor, + 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, ) -> tuple[torch.Tensor, torch.Tensor]: - """Compute one prepared, nonempty GDN prefix [T,H,D] with a key-first state.""" + """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. + """ + 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 - state = initial_state - if state_qdq: - state = state_qdq_reference(state, state_qdq_block_v, state_format) - pieces = [] - for lo in range(0, len(q), chunk_size): - hi = min(lo + chunk_size, len(q)) - qc, kc, vc = (x[lo:hi].transpose(0, 1) for x in (q, k, v)) - gc = g[lo:hi].transpose(0, 1).cumsum(-1) - bc = beta[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() + 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 = 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() + gram = kc @ kc.transpose(-1, -2) + lower = (bc * gram * decay).tril(-1) + matrix = lower + torch.eye(hi - lo, device=q.device, dtype=q.dtype) + rhs = torch.cat((bc * vc, bc * kc * gc.exp().unsqueeze(-1)), 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) + 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_fp8_qdq_reference(state, state_qdq_block_v) + previous_n = n + length = qc.shape[1] + wc = w[offset : offset + length].transpose(0, 1) + offset += length + updated_values = u - wc @ state # state_read + 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)) + output = output + local_scores @ updated_values # local_readout + outputs.append(output.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 + state = state + weighted_keys.transpose(-1, -2) @ updated_values # state_update if state_qdq: - state = state_qdq_reference(state, state_qdq_block_v, state_format) - return torch.cat(pieces), state - - -def chunk_kda( - q, - k, - v, - g, - beta, - *, - chunk_size=64, - scale=None, - initial_state, - state_qdq=False, - state_qdq_block_v=64, - state_format="fp8_e4m3", -): - """Compute one prepared, nonempty KDA prefix [T,H,D] with a key-first state.""" - scale = q.shape[-1] ** -0.5 if scale is None else scale - state = initial_state - if state_qdq: - state = state_qdq_reference(state, state_qdq_block_v, state_format) - pieces = [] - for lo in range(0, len(q), chunk_size): - hi = min(lo + chunk_size, len(q)) - qc, kc, vc = (x[lo:hi].transpose(0, 1) for x in (q, k, v)) - prefix = g[lo:hi].transpose(0, 1).cumsum(-2) - bc = beta[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_reference(state, state_qdq_block_v, state_format) - return torch.cat(pieces), state - - -def chunk_delta_rule_reference( - q, k, v, g, beta, *, initial_state=None, cu_seqlens=None, state_v_first=False, chunk_size=64 -): - """Evaluate GDN/KDA chunk algebra independently for each packed or batched sequence.""" - q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) - chunk = chunk_kda if g.ndim == 4 else chunk_gdn - outputs, finals = [], [] - for n, (b, start, end) in enumerate(sequences): - output, state = chunk( - *(x[b, start:end] for x in (q, k, v, g, beta)), - initial_state=states[n], - chunk_size=chunk_size, - ) - outputs.append(output) - finals.append(state) - output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + state = state_fp8_qdq_reference(state, state_qdq_block_v) + 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 diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py index 9d46482f328..ab234b14c12 100644 --- a/tests/examples/megatron_bridge/test_quantize_export.py +++ b/tests/examples/megatron_bridge/test_quantize_export.py @@ -17,6 +17,7 @@ from pathlib import Path import pytest +import yaml from _test_utils.examples.run_command import extend_cmd_parts, run_example_command from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches from _test_utils.torch.megatron.modelopt_state import assert_has_modelopt_state @@ -26,6 +27,8 @@ create_tiny_qwen3vl_dir, ) +from modelopt.recipe import load_recipe + # Per-architecture export *mappings* are covered in-process by # tests/gpu_megatron/torch/export/test_unified_export_megatron.py; these cases cover the script # wiring (CLI, recipe, checkpoint hand-off) that only running quantize.py + the exporter exercises. @@ -40,6 +43,29 @@ } +def test_calibrated_state_recipe_rejected_before_model_load(tmp_path): + recipe = { + "metadata": {"recipe_type": "ptq"}, + "quantize": load_recipe( + "general/ptq/linear_attention_state_int8_block32_dynamic" + ).quantize.model_dump(), + } + recipe["quantize"]["algorithm"] = "max" + recipe["quantize"]["quant_cfg"].append( + {"quantizer_name": "*input_quantizer", "cfg": {"num_bits": [4, 3], "axis": None}} + ) + recipe_path = tmp_path / "calibrated_state.yaml" + recipe_path.write_text(yaml.safe_dump(recipe)) + cmd = extend_cmd_parts( + ["torchrun", "--nproc_per_node=1", "quantize.py"], + hf_model_name_or_path=tmp_path / "no_model", + export_megatron_path=tmp_path / "output", + recipe=recipe_path, + ) + with pytest.raises(ValueError, match="calibration loop inside linear_attention_training_phase"): + run_example_command(cmd, example_path="megatron_bridge") + + @pytest.mark.parametrize( ("create_model", "model_kwargs", "recipe", "quantize_parallelism", "export_parallelism"), [ 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..defe05aa5d5 --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py @@ -0,0 +1,91 @@ +# 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. + +"""Minimal forward/backward checks for GDN training fake quantization.""" + +import pytest +import torch +import torch.nn.functional as F +from _test_utils.torch.quantization.linear_attention_reference import chunk_gdn_reference + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer + +pytest.importorskip("fla.ops.gated_delta_rule") + +from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( + chunk_gated_delta_rule, +) + + +def make_inputs(): + torch.manual_seed(123) + # Two chunks exercise recurrence with one shared BF16 shape for every mode. + shape = (1, 128, 1, 32) + q, k = [F.normalize(torch.randn(shape, device="cuda"), dim=-1) for _ in range(2)] + v = torch.randn(shape, device="cuda") + g = -torch.rand(shape[:3], device="cuda") * 0.1 + beta = torch.rand_like(g) + args = [x.to(torch.bfloat16).requires_grad_() for x in (q, k, v, g, beta)] + state = (torch.randn(1, 1, 32, 32, device="cuda") * 0.1).requires_grad_() + return args, state + + +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 + + +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}" + ) + + +@pytest.fixture(scope="module", params=["disabled", "w", "state-w"]) +def compiled_gdn_case(request): + """Compile only the selected BF16 forward/backward path, outside the test-call timer.""" + state_qdq = request.param == "state-w" + if state_qdq and torch.cuda.get_device_capability() < (8, 9): + pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") + quantizer = ( + TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), axis=(0, 1, 2), type="dynamic")) + if request.param != "disabled" + else None + ) + args, state = make_inputs() + kwargs = {"state_qdq": state_qdq, "w_quantizer": quantizer} + values_and_grads(chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs) + torch.cuda.synchronize() + return args, state, kwargs + + +def test_gdn_forward_and_backward(compiled_gdn_case): + args, state, kwargs = compiled_gdn_case + reference_args = [x.detach().float().requires_grad_() for x in args] + reference_state = state.detach().clone().requires_grad_() + expected = values_and_grads(chunk_gdn_reference, reference_args, reference_state, **kwargs) + actual = values_and_grads( + chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs + ) + compare(actual[0], expected[0], 0.03) + compare(actual[1], expected[1], 0.05) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index 2092e8b9da0..be822e18c35 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -23,13 +23,12 @@ initialize_for_megatron, sharded_state_dict_test_helper, ) +from megatron.core.packed_seq_params import PackedSeqParams +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 +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase pytest.importorskip("fla") # Megatron-Core GatedDeltaNet needs FLA for its baseline kernels. GatedDeltaNet = pytest.importorskip("megatron.core.ssm.gated_delta_net").GatedDeltaNet @@ -64,47 +63,23 @@ def _make_model(tp_size): if isinstance(module, GatedDeltaNet): module.A_log.fill_(-4) module.dt_bias.zero_() + # Shared hybrid-model KDA settings must not reject a GDN layer. + module.config.kda_safe_gate = True + module.config.kda_lower_bound = -5.0 return model def _gdn_config(): - return { - "quant_cfg": [ - {"quantizer_name": "*", "enable": False}, - { - "quantizer_name": "*gdn_state_quantizer", - "cfg": { - "num_bits": 8, - "block_sizes": {-1: 32}, - "type": "dynamic", - "unsigned": False, - "narrow_range": True, - }, - }, - ], - "algorithm": None, - "linear_attention": [ - { - "module_name": "decoder.layers.*.self_attention", - "cfg": { - "backend": "serving", - "state_block_v": 16, - "precision": "vllm_0_15", - }, - } - ], - } + return load_recipe( + "general/ptq/linear_attention_state_int8_block32_dynamic" + ).quantize.model_dump() def _gdn_forward(model): original_forward = get_forward(model) - def forward(m): - enabled = any( - getattr(getattr(layer, "linear_attention_config", None), "backend", None) == "serving" - for layer in m.modules() - ) - with linear_attention_training_phase(m, [64, 64]) if enabled else nullcontext(): + def forward(m, *, decode_aware=True): + with linear_attention_training_phase(m, [64, 64]) if decode_aware else nullcontext(): return original_forward(m) return forward @@ -116,86 +91,107 @@ def _test_gdn_qat_helper(rank, size, checkpoint_path): ) 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() - + baseline = forward(model, decode_aware=False) mtq.quantize(model, _gdn_config()) - 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 - assert not module.gdn_w_quantizer.is_enabled - # 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} - ) - + module = next(m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)) + calls = [] + hook = module.gdn_state_quantizer.register_forward_hook(lambda *_: calls.append(True)) + try: + assert torch.isfinite(forward(model)).all() + assert calls + finally: + hook.remove() + mtq.disable_quantizer(model, "*") 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 + torch.testing.assert_close(forward(model, decode_aware=False), baseline, rtol=0, atol=0) + mtq.enable_quantizer(model, "*gdn_state_quantizer") + module.linear_attention_config.state_block_v = 32 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 - assert not module.gdn_w_quantizer.is_enabled - assert module.gdn_state_quantizer.num_bits == original.gdn_state_quantizer.num_bits - 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) - - fused_model = _make_model(size) - for module in fused_model.modules(): - if isinstance(module, GatedDeltaNet): - module.gdn_pre_gated_delta_rule_fusion = True - with pytest.raises(NotImplementedError, match="unfused input-preparation hook"): - mtq.quantize(fused_model, _gdn_config()) + restored_gdn = next(m for m in restored.modules() if isinstance(m, _QuantGatedDeltaNet)) + assert restored_gdn.linear_attention_config == module.linear_attention_config + assert restored_gdn.gdn_state_quantizer.is_enabled + assert torch.isfinite(restored_gdn.in_proj.weight.grad).all() + before = restored_gdn.in_proj.weight.detach().clone() + torch.optim.SGD(restored.parameters(), lr=1e-3).step() + assert not torch.equal(restored_gdn.in_proj.weight, before) + _check_packed_recompute(restored_gdn) + + for name, value, restore in ( + ("context_parallel_size", 2, True), + ("recompute_granularity", "full", True), + ("recompute_granularity", "full", False), + ): + receiver = _make_model(size) + layer = next(m for m in receiver.modules() if isinstance(m, GatedDeltaNet)) + original = getattr(layer.config, name) + try: + setattr(layer.config, name, value) + receiver.train() + if restore: + mto.restore_from_modelopt_state(receiver, mto.modelopt_state(model)) + else: + mtq.quantize(receiver, _gdn_config()) + assert layer.training + for enabled in (False, True): + if enabled: + layer.gdn_state_quantizer.enable() + else: + layer.gdn_state_quantizer.disable() + layer.modelopt_post_restore() + with ( + pytest.raises(NotImplementedError, match=r"context parallelism|full-layer"), + linear_attention_training_phase(layer, [0]), + ): + layer(torch.zeros(1, 1, 64, device="cuda"), attention_mask=None) + finally: + setattr(layer.config, name, original) + + +def _check_packed_recompute(module): + hidden = torch.randn(128, 1, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) + packed = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=torch.tensor([0, 60, 118], device="cuda", dtype=torch.int32), + cu_seqlens_kv=torch.tensor([0, 60, 118], device="cuda", dtype=torch.int32), + cu_seqlens_q_padded=torch.tensor([0, 64, 128], device="cuda", dtype=torch.int32), + cu_seqlens_kv_padded=torch.tensor([0, 64, 128], device="cuda", dtype=torch.int32), + max_seqlen_q=64, + max_seqlen_kv=64, + total_tokens=128, + ) + original = module.recompute_gdn + was_training = module.training + module.train() + try: + results = [] + for recompute, lengths in ((False, [60, 58]), (False, None), (True, None)): + module.recompute_gdn = recompute + module.zero_grad(set_to_none=True) + hidden.grad = None + with ( + linear_attention_training_phase( + module, + [56, 54], + sequence_lengths=lengths, + cu_seqlens=packed.cu_seqlens_q_padded, + ), + torch.autocast("cuda", dtype=torch.bfloat16), + ): + output = module(hidden, attention_mask=None, packed_seq_params=packed)[0] + # Backward runs after the phase exits, as in pipeline schedules. + output.float().square().sum().backward() + results.append( + (output.detach(), hidden.grad.clone(), module.in_proj.weight.grad.clone()) + ) + for actual in results[1:]: + for value, expected in zip(actual, results[0]): + torch.testing.assert_close(value, expected, rtol=0, atol=0) + finally: + module.recompute_gdn = original + module.train(was_training) def _compile_gdn_qat_kernels(rank, size): @@ -207,13 +203,13 @@ def _compile_gdn_qat_kernels(rank, size): with torch.no_grad(): forward(model) cfg = _gdn_config() - # 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() + _check_packed_recompute(next(m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet))) torch.cuda.synchronize() diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py index 79c4b89e56c..8157fd2ccd9 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py @@ -17,11 +17,7 @@ import pytest import torch -from _test_utils.torch.megatron.utils import ( - initialize_for_megatron, - load_distributed_checkpoint, - save_distributed_checkpoint, -) +from _test_utils.torch.megatron.utils import initialize_for_megatron from megatron.core import parallel_state from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec from megatron.core.process_groups_config import ProcessGroupCollection @@ -29,14 +25,7 @@ import modelopt.torch.quantization as mtq from modelopt.recipe import load_recipe -from modelopt.torch.opt.plugins.mcore_dist_checkpointing import ( - restore_sharded_modelopt_state, - save_sharded_modelopt_state, -) -from modelopt.torch.quantization.linear_attention import ( - LinearAttentionConfig, - linear_attention_training_phase, -) +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase KimiDeltaAttention = pytest.importorskip("megatron.core.ssm.gated_delta_net.kda").KimiDeltaAttention pytest.importorskip("vllm") @@ -86,11 +75,13 @@ def _case(cfg): model = _layer() hidden = torch.randn(73, 1, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) - def forward(layer): + def forward(layer, *, decode_aware=True, prefix=31, length=None): policy = getattr(layer, "linear_attention_config", None) phase = ( - linear_attention_training_phase(layer, [31]) - if policy is not None and policy.backend == "serving" + linear_attention_training_phase( + layer, [prefix], sequence_lengths=None if length is None else [length] + ) + if policy is not None and policy.backend == "serving" and decode_aware else nullcontext() ) with phase, torch.autocast("cuda", dtype=torch.bfloat16): @@ -98,19 +89,32 @@ def forward(layer): with torch.no_grad(): baseline = forward(model) - mtq.quantize(model, cfg) + + def calibrate(m): + with linear_attention_training_phase(m, [73]): + return forward(m, decode_aware=False) + + mtq.quantize(model, {**cfg, "algorithm": "max"}, calibrate) return model, hidden, forward, baseline def _compile_kda(rank, size, cfg): model, hidden, forward, _ = _case(cfg) - with linear_attention_training_phase(model, [31]): - forward(model).float().square().mean().backward() + for prefix in (31, 32): + forward(model, prefix=prefix, length=prefix + 4).float().square().mean().backward() torch.cuda.synchronize() -def _test_kda(rank, size, cfg, checkpoint_path): +def _test_kda(rank, size, cfg): model, hidden, forward, baseline = _case(cfg) + for name, value in (("kda_safe_gate", True), ("kda_lower_bound", -5.0)): + original = getattr(model.config, name) + try: + setattr(model.config, name, value) + with pytest.raises(NotImplementedError, match="unbounded softplus gates"): + mtq.quantize(model, cfg) + finally: + setattr(model.config, name, original) assert model.kda_state_quantizer.is_enabled assert model.kda_state_quantizer.num_bits == 8 assert model.linear_attention_config.precision == "vllm" @@ -120,43 +124,35 @@ def _test_kda(rank, size, cfg, checkpoint_path): assert not torch.equal(forward(model), baseline) assert model.gated_delta_rule is kernel - # The sharded checkpoint must restore this per-module change, not recipe defaults. - model.linear_attention_config.state_block_v = 32 - policy = model.linear_attention_config + with pytest.raises(ValueError, match="explicit prefill lengths"): + forward(model, decode_aware=False) + # Two in-flight microbatches must keep their own phases after both contexts exit. + expected_grads = [] + model.recompute_gdn = False + for prefix in (31, 32): + model.zero_grad(set_to_none=True) + hidden.grad = None + forward(model, prefix=prefix, length=prefix + 4).float().square().mean().backward() + expected_grads.append((hidden.grad.clone(), model.in_proj.weight.grad.clone())) + model.recompute_gdn = True + pending = [forward(model, prefix=prefix, length=prefix + 4) for prefix in (31, 32)] + for output, (hidden_grad, weight_grad) in zip(pending, expected_grads): + model.zero_grad(set_to_none=True) + hidden.grad = None + output.float().square().mean().backward() + torch.testing.assert_close(hidden.grad, hidden_grad, rtol=0, atol=0) + torch.testing.assert_close(model.in_proj.weight.grad, weight_grad, rtol=0, atol=0) + model.eval() + forward(model).float().sum().backward() # No checkpoint hook may reject eval gradients. + model.train() + mtq.disable_quantizer(model, "*") - model.linear_attention_config = LinearAttentionConfig() with torch.no_grad(): - torch.testing.assert_close(forward(model), baseline, rtol=0, atol=0) - model.linear_attention_config = policy + torch.testing.assert_close(forward(model, decode_aware=False), baseline, rtol=0, atol=0) mtq.enable_quantizer(model, "*kda_state_quantizer") - - with torch.no_grad(): - expected = forward(model) - # A change made after conversion must survive beyond the saved recipe defaults. - mtq.disable_quantizer(model, "*kda_state_quantizer") - save_distributed_checkpoint(checkpoint_path, model) - save_sharded_modelopt_state([model], checkpoint_path) - restored = _layer() - restore_sharded_modelopt_state([restored], checkpoint_path) - load_distributed_checkpoint(checkpoint_path, restored) - assert not restored.kda_state_quantizer.is_enabled - mtq.enable_quantizer(restored, "*kda_state_quantizer") - # Megatron recomputes the core during backward, after forward's kernel wrapper exits. - with linear_attention_training_phase(restored, [31]): - actual = forward(restored) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - actual.float().square().mean().backward() - assert not hasattr(restored, "kda_w_quantizer") - assert restored._linear_attention_prefill_lengths is None - assert restored.linear_attention_config == policy - assert restored.gated_delta_rule is kernel - assert torch.isfinite(hidden.grad).all() - for parameter in restored.parameters(): - assert parameter.grad is not None - assert torch.isfinite(parameter.grad).all() - before = restored.in_proj.weight.detach().clone() - torch.optim.SGD(restored.parameters(), lr=0.1).step() - assert not torch.equal(before, restored.in_proj.weight) + before = model.in_proj.weight.detach().clone() + torch.optim.SGD(model.parameters(), lr=0.1).step() + assert not torch.equal(before, model.in_proj.weight) @pytest.fixture(scope="module") @@ -169,6 +165,6 @@ def compiled_kda_workers(dist_workers_size_1): return dist_workers_size_1, cfg -def test_kda_qat_and_sharded_restore(compiled_kda_workers, tmp_path): +def test_kda_qat_recompute(compiled_kda_workers): workers, cfg = compiled_kda_workers - workers.run(_test_kda, cfg, tmp_path) + workers.run(_test_kda, cfg) diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py index d6e696b2b40..860b9d37383 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py @@ -59,7 +59,7 @@ def compiled_replay(request): "config": config, "state_qdq": True, "state_format": "int8", - "replay_gate_inputs": None if channel else (raw, raw_beta, rate, bias), + "gate_inputs": None if channel else (raw, raw_beta, rate, bias), } def forward(): @@ -100,7 +100,7 @@ def test_replay_matches_persistent_serving_cache(compiled_replay): if channel: a, b, rate, bias = gate.flatten()[None], beta[None], None, None else: - raw, raw_beta, rate, bias = kwargs["replay_gate_inputs"] + raw, raw_beta, rate, bias = kwargs["gate_inputs"] a, b = raw[t : t + 1], raw_beta[t : t + 1] expected = torch.empty(1, 1, 64, device="cuda", dtype=torch.bfloat16) native.fused_recurrent_gated_delta_rule_replayssm( @@ -130,13 +130,6 @@ def test_replay_matches_persistent_serving_cache(compiled_replay): carry.anchor.values, (decoded * metadata[..., None]).flatten(-2), rtol=0, atol=0 ) torch.testing.assert_close(carry.anchor.scales, metadata, rtol=0, atol=0) - for i, entry in enumerate(carry.entries): - torch.testing.assert_close( - entry.key.values.float(), keys[1, :, i].float(), rtol=0, atol=0 - ) - torch.testing.assert_close( - entry.update.values.float(), updates[1, :, i].float(), rtol=0, atol=0 - ) grads = torch.autograd.grad( output.square().mean() + carry.reconstruct().square().mean(), (*args, initial) ) diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py index 785a7b530d4..b99981f9c6d 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py @@ -13,10 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from functools import partial +from inspect import signature import pytest import torch +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 ( @@ -28,62 +29,144 @@ vllm = pytest.importorskip("vllm") -from modelopt.torch.kernels.quantization.linear_attention.serving._compat import fla_module +from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator + +from modelopt.torch.kernels.quantization.linear_attention.serving._compat import ( + _kda_state_v_first, + fla_module, + gdn_module, + kda_module, + optional_module, +) @torch.no_grad() -def _native_forward(args, quantizer, kda): - """Execute native prefill/decode directly, with QDQ in ModelOpt's [K,V] basis.""" - prefill = fla_module("kda").chunk_kda if kda else fla_module("chunk").chunk_gated_delta_rule - step = ( - fla_module("kda").fused_recurrent_kda - if kda - else fla_module("fused_recurrent").fused_recurrent_gated_delta_rule +def _native_forward(args, quantizer, kda, kimi_k3): + """Run the CI vLLM 0.20/0.30 Triton APIs with a persistent quantized cache.""" + q, k, v, raw_g, raw_beta, rate, bias = [x.detach().clone() for x in args] + _, shape = MambaStateShapeCalculator.gated_delta_net_state_shape( + 1, 1, 1, q.shape[-1], v.shape[-1], 4 ) - value_first = vllm.__version_tuple__[:2] >= (0, 16) + # CI KDA kernels specify [H,V,K]; rectangular GDN advertises its own layout. + value_first = kda or shape == (1, v.shape[-1], q.shape[-1]) def layout(state): - return state.transpose(-1, -2).contiguous() if value_first else state - - state = layout(args[0].new_zeros(1, 1, args[0].shape[-1], args[2].shape[-1]).float()) - prefix, state = prefill( - *[x[:, :65].detach().clone() for x in args], - initial_state=state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - cu_seqlens=torch.tensor([0, 65], device="cuda", dtype=torch.int32), - ) + return (state.transpose(-1, -2) if value_first else state).contiguous() + + cache = torch.zeros(2, *shape, device=q.device) + prefix_args = [x[:, :65].contiguous() for x in (q, k, v)] + kwargs = { + "initial_state": cache[1:], + "output_final_state": True, + "use_qk_l2norm_in_kernel": True, + "cu_seqlens": torch.tensor([0, 65], device=q.device, dtype=torch.int32), + } + if kda: + native = kda_module(kimi_k3) + g = fla_module("kda").fused_kda_gate(raw_g.flatten(-2), rate, q.shape[-1], g_bias=bias) + beta = raw_beta.float().sigmoid() + raw_chunk = getattr(native, "chunk_kda_with_fused_gate", None) + if raw_chunk is not None: + beta_name = "raw_beta" if "raw_beta" in signature(raw_chunk).parameters else "beta" + prefix, state = raw_chunk( + *prefix_args, + raw_g=raw_g[:, :65].contiguous(), + **{beta_name: (raw_beta if beta_name == "raw_beta" else beta)[:, :65].contiguous()}, + A_log=rate, + g_bias=bias, + **kwargs, + ) + else: # vLLM 0.20 standalone KDA takes activated gates. + prefix, state = native.chunk_kda(*prefix_args, g[:, :65], beta[:, :65], **kwargs) + else: + prepared = gdn_module("fused_gdn_prefill_post_conv").fused_post_conv_prep( + torch.cat([x[0].flatten(1) for x in prefix_args], dim=-1), + raw_g[0, :65], + raw_beta[0, :65], + rate, + bias, + 1, + q.shape[-1], + v.shape[-1], + ) + kwargs["use_qk_l2norm_in_kernel"] = False + prefix, state = fla_module("chunk").chunk_gated_delta_rule( + *[x[None] for x in prepared], **kwargs + ) + cache[1] = layout(quantizer(layout(state)[0])) outputs = [prefix] - state = layout(quantizer(layout(state)[0])[None]) + indices = torch.ones(1, device=q.device, dtype=torch.int32) + cu = torch.tensor([0, 1], device=q.device, dtype=torch.int32) for token in range(65, 73): - output, state = step( - *[x[:, token : token + 1].detach().clone() for x in args], - initial_state=state, - inplace_final_state=False, - use_qk_l2norm_in_kernel=True, - cu_seqlens=torch.tensor([0, 1], device="cuda", dtype=torch.int32), - ) + operands = [x[:, token : token + 1].contiguous() for x in (q, k, v)] + if kimi_k3: + output, _ = native.fused_recurrent_kda_packed_decode( + torch.cat([x.flatten() for x in operands])[None], + raw_g[:, token : token + 1].contiguous(), + raw_beta[:, token : token + 1].contiguous(), + rate, + bias, + None, + cache, + indices, + ) + state = cache[1:] + elif kda: + output, state = native.fused_recurrent_kda( + *operands, + g[:, token : token + 1], + beta[:, token : token + 1], + initial_state=cache[1:], + inplace_final_state=False, + cu_seqlens=cu, + use_qk_l2norm_in_kernel=True, + ) + else: + output, _ = fla_module( + "fused_recurrent" + ).fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=torch.cat([x.flatten() for x in operands])[None], + a=raw_g[0, token : token + 1], + b=raw_beta[0, token : token + 1], + A_log=rate, + dt_bias=bias, + scale=q.shape[-1] ** -0.5, + initial_state=cache, + out=torch.empty_like(operands[-1]), + ssm_state_indices=indices, + use_qk_l2norm_in_kernel=True, + ) + state = cache[1:] outputs.append(output) - state = layout(quantizer(layout(state)[0])[None]) - return torch.cat(outputs, dim=1), layout(state) + cache[1] = layout(quantizer(layout(state)[0])) + return torch.cat(outputs, dim=1), layout(cache[1:]) -@pytest.fixture(scope="module", params=[False, True]) +@pytest.fixture(scope="module", params=["gdn", "kda", "kimi_k3"]) def compiled_serving_case(request): """Compile one shared BF16 shape per model, outside the test-call timer.""" - kda = request.param + kda = request.param != "gdn" + kimi_k3 = request.param == "kimi_k3" + if kimi_k3 and optional_module("vllm.models.kimi_k3.nvidia.ops.third_party.kda") is None: + pytest.skip("This vLLM build does not provide the Kimi-K3 backend") torch.manual_seed(73) - # A rectangular GDN state detects accidental key/value-axis swaps. + keys = 128 if kda else 32 args = [ torch.randn(1, 73, 1, dim, device="cuda", dtype=torch.bfloat16) - for dim in (32, 32, 32 if kda else 64) + # Rectangular GDN state catches transposes; equal padded tile widths also + # satisfy the fused post-conv kernel's tile-width constraint. + for dim in (keys, keys, keys if kda else 24) ] args += [ - -torch.rand((1, 73, 1, 32) if kda else (1, 73, 1), device="cuda") * 0.03, - torch.rand(1, 73, 1, device="cuda") * 0.4, + torch.randn((1, 73, 1, keys) if kda else (1, 73, 1), device="cuda", dtype=torch.bfloat16), + torch.randn(1, 73, 1, device="cuda", dtype=torch.bfloat16), + torch.full((1,), -3.0, device="cuda"), + torch.zeros(keys if kda else 1, device="cuda"), ] args = [x.requires_grad_() for x in args] - policy = LinearAttentionConfig(backend="serving", precision="vllm") + policy = LinearAttentionConfig( + backend="serving", precision="vllm_kimi_k3" if kimi_k3 else "vllm" + ) quantizer = TensorQuantizer( QuantizerAttributeConfig( num_bits=8, @@ -93,38 +176,202 @@ def compiled_serving_case(request): pass_through_bwd=True, ) ).cuda() - forward = partial( - kda_state_qat if kda else gdn_state_qat, - *args, - policy=policy, - state_quantizer=quantizer, - prefill_lengths=[65], - use_qk_l2norm_in_kernel=True, - output_final_state=True, - ) - output, state = forward() + + def forward(inputs=args, **overrides): + q, k, v, raw_g, raw_beta, rate, bias = inputs + kwargs = { + "policy": policy, + "state_quantizer": quantizer, + "prefill_lengths": [65], + "use_qk_l2norm_in_kernel": True, + "output_final_state": True, + "gate_inputs": (raw_g, raw_beta, rate, bias), + } + if kda: + kwargs.update( + A_log=rate, + dt_bias=bias, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + ) + kwargs.update(overrides) + if kda: + return kda_state_qat(q, k, v, raw_g, raw_beta, **kwargs) + gate = -rate.exp() * torch.nn.functional.softplus(raw_g.float() + bias) + return gdn_state_qat(q, k, v, gate, raw_beta.float().sigmoid(), **kwargs) + + try: + output, state = forward() + except NotImplementedError as error: + if vllm.__version__ == "0.20.0" and "KDA chunk/decode cache readout disagrees" in str( + error + ): + pytest.xfail(str(error)) + raise torch.autograd.grad(output.float().sum() + state.sum(), args) - expected = _native_forward(args, quantizer, kda) + expected = _native_forward(args, quantizer, kda, kimi_k3) torch.cuda.synchronize() - return args, quantizer, forward, expected + return args, quantizer, forward, expected, kimi_k3 def test_serving_state_qdq_and_handoff_gradient(compiled_serving_case): - args, quantizer, forward, (expected, expected_state) = compiled_serving_case + args, quantizer, forward, (expected, expected_state), _ = compiled_serving_case calls = [] handle = quantizer.register_forward_hook(lambda *_: calls.append(True)) - output, state = forward() - handle.remove() - # One handoff QDQ, then one QDQ per suffix update; none inside the fresh prefill. - assert len(calls) == 9 - with torch.no_grad(): - torch.testing.assert_close(output, expected, rtol=0, atol=0) - torch.testing.assert_close(state, expected_state, rtol=0, atol=0) - quantizer.disable() - plain, _ = forward() - assert not torch.equal(output[:, 65:], plain[:, 65:]) + try: + output, state = forward() + finally: + handle.remove() + assert len(calls) == 9 # Handoff plus eight suffix writes. + torch.testing.assert_close(output, expected, rtol=0, atol=0) + torch.testing.assert_close(state, expected_state, rtol=0, atol=0) gradients = torch.autograd.grad( output[:, 65:].float().square().sum() + state.square().sum(), args ) assert all(torch.isfinite(x).all() for x in gradients) assert gradients[1][:, :65].abs().sum() > 0 + with pytest.raises(ValueError, match="explicit prefill lengths"): + forward(prefill_lengths=None) + + with torch.no_grad(): + quantizer.disable() + try: + plain, _ = forward() + assert not torch.equal(output[:, 65:], plain[:, 65:]) + # A nonzero nonsymmetric state checks layout against an independent oracle. + initial = torch.randn(1, 1, args[0].shape[-1], args[2].shape[-1], device="cuda") + for prefix in (73, 65): + actual, final = forward(prefill_lengths=[prefix], initial_state=initial) + q, k, v, raw_g, raw_beta, rate, bias = [x.float() for x in args] + q, k = [x / (x.square().sum(-1, keepdim=True) + 1e-6).sqrt() for x in (q, k)] + q[:, :prefix] = q[:, :prefix].bfloat16().float() + k[:, :prefix] = k[:, :prefix].bfloat16().float() + gate = -rate.exp().reshape( + (1, 1, -1, 1) if raw_g.ndim == 4 else (1, 1, -1) + ) * torch.nn.functional.softplus(raw_g + bias.reshape(raw_g.shape[2:])) + expected, expected_state = recurrent_delta_rule_reference( + q, k, v, gate, raw_beta.sigmoid(), initial_state=initial + ) + torch.testing.assert_close(actual.float(), expected, rtol=0.02, atol=0.01) + torch.testing.assert_close(final, expected_state, rtol=0.02, atol=0.01) + finally: + quantizer.enable() + + +@pytest.mark.parametrize( + ("overrides", "message", "kimi_only"), + [ + ( + {"gate_inputs": None, "use_gate_in_kernel": False, "use_beta_sigmoid_in_kernel": False}, + "raw gate_inputs", + True, + ), + ({"use_qk_l2norm_in_kernel": False}, "normalization", True), + ({"allow_neg_eigval": True}, "allow_neg_eigval", False), + ], +) +def test_rejects_silent_serving_mismatches(compiled_serving_case, overrides, message, kimi_only): + args, _, forward, _, kimi_k3 = compiled_serving_case + if args[3].ndim != 4 or (kimi_only and not kimi_k3): + pytest.skip("Guard applies only to the selected KDA profile") + with pytest.raises(ValueError, match=message): + forward(**overrides) + + +def test_kda_non_power_of_two_layout(monkeypatch): + # Probe values must be BF16-exact for K=192, including fused-gate chunk dispatch. + try: + for raw_gates in (False, True): + assert isinstance( + _kda_state_v_first(torch.device("cuda:0"), False, 192, raw_gates), bool + ) + except NotImplementedError as error: + if vllm.__version__ == "0.20.0" and "KDA chunk/decode cache readout disagrees" in str( + error + ): + pytest.xfail(str(error)) + raise + native = kda_module() + recurrent = native.fused_recurrent_kda_fwd + + def shifted_readout(*args, **kwargs): + output, state = recurrent(*args, **kwargs) + return output.roll(1, -1), state + + _kda_state_v_first.cache_clear() + try: + with monkeypatch.context() as patch: + patch.setattr(native, "fused_recurrent_kda_fwd", shifted_readout) + with pytest.raises(NotImplementedError, match="cache layout could not be validated"): + _kda_state_v_first(torch.device("cuda:0"), False, 192) + finally: + _kda_state_v_first.cache_clear() + + +@pytest.mark.parametrize("packed", [False, True], ids=["dense", "packed"]) +def test_padding_does_not_extend_decode(compiled_serving_case, packed): + args, quantizer, forward, _, _ = compiled_serving_case + valid = 69 # The existing compiled shape uses a 65-token prefix. + trimmed = [ + (x[:, :valid] if i < 5 else x).detach().clone().requires_grad_() for i, x in enumerate(args) + ] + initial = args[0].new_zeros( + 2, args[2].shape[2], args[0].shape[-1], args[2].shape[-1], dtype=torch.float32 + ) + initial[1].fill_(1.0) # Empty rows preserve an incoming nonzero state without QDQ. + expected, expected_state = forward(inputs=trimmed, initial_state=initial[:1]) + expected_grads = torch.autograd.grad( + expected.square().sum() + expected_state.square().sum(), trimmed + ) + padded = [] + for i, x in enumerate(args): + if i < 5: + x = torch.cat((x, torch.zeros_like(x)), dim=0) + if packed: + x = x.flatten(0, 1).unsqueeze(0) + padded.append(x.detach().clone().requires_grad_()) + calls = [] + hook = quantizer.register_forward_hook(lambda *_: calls.append(True)) + try: + output, state = forward( + inputs=padded, + initial_state=initial, + prefill_lengths=[65, 0], + sequence_lengths=[valid, 0], + cu_seqlens=torch.tensor([0, 73, 146], device="cuda", dtype=torch.int32) + if packed + else None, + ) + finally: + hook.remove() + assert len(calls) == 6 # Incoming prefix state, handoff, four writes; none for the empty row. + output = output.reshape(2, 73, *output.shape[-2:]) + torch.testing.assert_close(output[:1, :valid], expected, rtol=0, atol=0) + torch.testing.assert_close(state[:1], expected_state, rtol=0, atol=0) + assert not output[0, valid:].count_nonzero() and not output[1].count_nonzero() + torch.testing.assert_close(state[1], initial[1], rtol=0, atol=0) + grads = torch.autograd.grad(output.square().sum() + state.square().sum(), padded) + for i, (grad, expected_grad) in enumerate(zip(grads, expected_grads)): + if i < 5: + grad = grad.reshape(2, 73, *grad.shape[2:]) + assert not grad[0, valid:].count_nonzero() and not grad[1].count_nonzero() + grad = grad[:1, :valid] + torch.testing.assert_close(grad, expected_grad) + + if not packed: + with torch.no_grad(): + prefill, prefill_state = forward( + inputs=trimmed, prefill_lengths=[valid], initial_state=initial[:1] + ) + padded_prefill, padded_state = forward( + inputs=padded, + prefill_lengths=[valid, 0], + sequence_lengths=[valid, 0], + initial_state=initial, + ) + torch.testing.assert_close(padded_prefill[:1, :valid], prefill, rtol=0, atol=0) + torch.testing.assert_close(padded_state[:1], prefill_state, rtol=0, atol=0) + torch.testing.assert_close(padded_state[1], initial[1], rtol=0, atol=0) + assert ( + not padded_prefill[0, valid:].count_nonzero() and not padded_prefill[1].count_nonzero() + ) diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index c701b7c1e7b..e243589bedb 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -13,24 +13,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -from copy import deepcopy - +import numpy as np import pytest import torch import torch.nn as nn -from _test_utils.torch.distributed.utils import spawn_multiprocess_job import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq -from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.linear_attention import ( LinearAttentionConfig, + get_linear_attention_layers, linear_attention_training_phase, ) from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins import gdn from modelopt.torch.quantization.plugins.gdn import GatedDeltaNetStateQuantMixin from modelopt.torch.quantization.plugins.kda import KimiDeltaAttentionStateQuantMixin +from modelopt.torch.quantization.plugins.linear_attention import _validate_linear_attention GDN_STATE_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"} @@ -82,72 +81,28 @@ def quant_cfg(): } -@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(): +def test_disabled_quantization_preserves_baseline(): 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("mixin", [GatedDeltaNetStateQuantMixin, KimiDeltaAttentionStateQuantMixin]) -def test_state_qat_requires_explicit_serving_policy(mixin): - model = mixin.convert(TinyGatedDeltaNet()) - model._linear_attn_state.set_from_attribute_config(GDN_STATE_FP8_DYNAMIC) - model._linear_attn_state.enable() - with pytest.raises(ValueError, match="requires backend='serving'"): - model.validate_linear_attention() - - -def test_training_phase_routes_prefill_lengths(monkeypatch): + cfg = quant_cfg() + cfg["quant_cfg"] = [{"quantizer_name": "*", "enable": False}] + mtq.quantize(model, cfg, lambda m: m(x)) + torch.testing.assert_close(model(x), expected, rtol=0, atol=0) + assert model.gated_delta_rule is chunk_gated_delta_rule + assert not model.linear_attention_is_enabled + + +def test_phase_routes_lengths_and_preserves_outer_context(monkeypatch): + monkeypatch.setattr( + _QuantTinyGatedDeltaNet, + "validate_linear_attention_execution", + lambda _: pytest.fail("Conversion must not validate runtime execution settings"), + ) cfg = {**quant_cfg(), "algorithm": None} - missing_policy = {key: value for key, value in cfg.items() if key != "linear_attention"} with pytest.raises(ValueError, match="requires backend='serving'"): - mtq.quantize(TinyGatedDeltaNet(), missing_policy) + mtq.quantize(TinyGatedDeltaNet(), {k: v for k, v in cfg.items() if k != "linear_attention"}) model = mtq.quantize(TinyGatedDeltaNet(), cfg) - x = torch.randn(2, 8, 3, 4) calls = [] def forward(*args, **kwargs): @@ -155,129 +110,101 @@ def forward(*args, **kwargs): return chunk_gated_delta_rule(*args) monkeypatch.setattr(gdn, "gdn_state_qat", forward) - with linear_attention_training_phase(model, [4, 4]): - model(x) + with ( + linear_attention_training_phase( + model, np.array([4, 4]), sequence_lengths=[7, 6], cu_seqlens=[0, 8, 16] + ), + linear_attention_training_phase(model, [1, 1], preserve_existing=True), + ): + model(torch.randn(2, 8, 3, 4)) assert calls[-1]["prefill_lengths"] == (4, 4) + assert calls[-1]["sequence_lengths"] == (7, 6) + boundaries = calls[-1]["cu_seqlens_cpu"] + actual = torch.tensor([0, 8, 16]) + assert boundaries.resolve(actual) == (0, 8, 16) + actual.data[1] = 7 # Alias writes do not bump the tensor version. + with pytest.raises(ValueError, match="does not match"): + boundaries.resolve(actual) + assert model._linear_attention_prefill_lengths is None + assert model._linear_attention_sequence_lengths is None + mtq.disable_quantizer(model, "*") + assert get_linear_attention_layers(model) == (model,) + assert not get_linear_attention_layers(model, enabled_only=True) + with linear_attention_training_phase(model, [4, 4]): + assert not get_linear_attention_layers(model, unphased_only=True) + assert model.linear_attention_is_enabled + assert not model.linear_attention_is_enabled -@pytest.mark.parametrize("phase", ["convert", "forward", "restore"]) -def test_enabled_legacy_w_quantizer_is_rejected(phase): - cfg = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None} - model = mtq.quantize(TinyGatedDeltaNet(), cfg) - with pytest.raises(ValueError, match="GDN W quantization is no longer supported"): - if phase == "convert": - cfg["quant_cfg"].append({"quantizer_name": "*gdn_w_quantizer", "enable": True}) - mtq.quantize(TinyGatedDeltaNet(), cfg) - else: - model.gdn_w_quantizer.enable() - if phase == "forward": - model(torch.randn(2, 8, 3, 4)) - else: - mto.restore_from_modelopt_state(TinyGatedDeltaNet(), mto.modelopt_state(model)) - - -@pytest.mark.parametrize( - "overrides", - [{"pass_through_bwd": False}, {"type": "static"}, {"fake_quant": False}, {"rotate": True}], -) -def test_unsupported_quantizer_fails_during_conversion(overrides): - cfg = deepcopy(quant_cfg()) - cfg["quant_cfg"][-1]["cfg"].update(overrides) - with pytest.raises(ValueError, match="supports only"): - mtq.quantize(TinyGatedDeltaNet(), cfg) - - -def test_quantizer_roundtrip_and_hybrid_selection(tmp_path): +def test_restore_uses_saved_quantizer_and_policy(tmp_path, monkeypatch): model = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) - cfg = deepcopy(quant_cfg()) - cfg["algorithm"] = None - cfg["quant_cfg"][-1]["cfg"].update( - num_bits=8, axis=None, block_sizes={-1: 32}, unsigned=False, narrow_range=True + mtq.quantize(model, {**quant_cfg(), "algorithm": None}) + # Change both after conversion: restore must not validate one against stale recipe defaults. + model[0].gdn_state_quantizer.set_from_attribute_config( + {"num_bits": 8, "axis": (0, 1), "type": "dynamic", "narrow_range": True} + ) + model[0].linear_attention_config = LinearAttentionConfig( + backend="serving", precision="replayssm", replay_window=4 ) - cfg["linear_attention"] = [ - {"module_name": "*", "cfg": {"backend": "serving", "state_block_v": 128}}, - {"module_name": "0", "cfg": {"backend": "serving", "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 sample = torch.randn(2, 4, 19) expected = model[0].gdn_state_quantizer(sample) path = tmp_path / "gdn.pth" mto.save(model, path) - restored = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) - mto.restore(restored, path) + + def no_rematching(*args): + pytest.fail("Restore must use saved resolved policies") + + monkeypatch.setattr( + "modelopt.torch.quantization.plugins.linear_attention._apply_linear_attention_policy", + no_rematching, + ) + restored = mto.restore(nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)), 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", True), ("gdn_w_quantizer", False)): - 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.block_sizes == original.block_sizes - assert quantizer._dynamic == original._dynamic - assert not hasattr(restored[1], name) + assert restored[0].gdn_state_quantizer.is_enabled + assert not restored[0].gdn_w_quantizer.is_enabled + assert not hasattr(restored[1], "gdn_state_quantizer") torch.testing.assert_close(restored[0].gdn_state_quantizer(sample), expected) -@pytest.mark.parametrize("reverse", [False, True], ids=["fp8-to-replay", "replay-to-fp8"]) -def test_restore_changed_policy_and_quantizer_together(tmp_path, reverse): - states = [ - ( - GDN_STATE_FP8_DYNAMIC, - LinearAttentionConfig(backend="serving", precision="vllm_0_15"), - ), - ( - {"num_bits": 8, "axis": (0, 1), "type": "dynamic", "narrow_range": True}, - LinearAttentionConfig(backend="serving", precision="replayssm"), - ), - ] - if reverse: - states.reverse() - (original_quantizer, original_policy), (saved_quantizer, saved_policy) = states - cfg = quant_cfg() - cfg["algorithm"] = None - cfg["quant_cfg"][-1]["cfg"] = original_quantizer - cfg["linear_attention"][0]["cfg"] = original_policy.model_dump() - model = mtq.quantize(TinyGatedDeltaNet(), cfg) - model.gdn_state_quantizer.set_from_attribute_config(saved_quantizer) - model.linear_attention_config = saved_policy - model.validate_linear_attention() - - path = tmp_path / "changed-policy.pth" - mto.save(model, path) - restored = mto.restore(TinyGatedDeltaNet(), path) - assert restored.linear_attention_config == saved_policy - assert restored.gdn_state_quantizer.num_bits == saved_quantizer["num_bits"] - assert restored.gdn_state_quantizer.is_enabled - restored.validate_linear_attention() +@pytest.mark.parametrize("mixin", [GatedDeltaNetStateQuantMixin, KimiDeltaAttentionStateQuantMixin]) +def test_state_qat_requires_serving_policy(mixin): + model = mixin.convert(TinyGatedDeltaNet()) + model._linear_attn_state.set_from_attribute_config(GDN_STATE_FP8_DYNAMIC) + model._linear_attn_state.enable() + with pytest.raises(ValueError, match="requires backend='serving'"): + model.validate_linear_attention() + model.linear_attention_config = LinearAttentionConfig(backend="serving") + if mixin is KimiDeltaAttentionStateQuantMixin: + with pytest.warns(UserWarning, match="standalone FLA/Triton"): + _validate_linear_attention(model) + with pytest.warns(UserWarning, match="standalone FLA/Triton"): + model.modelopt_post_restore() -def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): +@pytest.mark.parametrize("overrides", [{"pass_through_bwd": False}, {"type": "static"}]) +def test_unsupported_quantizer_fails_at_conversion(overrides): 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"][-1]["cfg"] = {**GDN_STATE_FP8_DYNAMIC, **overrides} + with pytest.raises(ValueError, match="supports only"): + mtq.quantize(TinyGatedDeltaNet(), cfg) - cfg = deepcopy(quant_cfg()) - cfg["algorithm"] = None - cfg["linear_attention"] = [ - {"module_name": "", "cfg": {"backend": "serving", "state_block_v": 32}} - ] - mtq.quantize(model, cfg) - assert model.gdn_state_qdq_block_v == 32 - cfg["linear_attention"].append({"module_name": "", "cfg": {"backend": "serving"}}) - mtq.quantize(model, cfg) - assert model.gdn_state_qdq_block_v == 64 - assert model.gdn_state_quantizer.is_enabled - assert not 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) +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"prefill_lengths": [-1]}, "Prefill lengths"), + ({"prefill_lengths": [0], "sequence_lengths": [True]}, "Sequence lengths"), + ({"prefill_lengths": [2], "sequence_lengths": [1]}, "prefill length <= sequence length"), + ({"prefill_lengths": [0, 0], "sequence_lengths": [1]}, "per sequence"), + ], +) +def test_invalid_phase_lengths(kwargs, message): + with ( + pytest.raises(ValueError, match=message), + linear_attention_training_phase(nn.Identity(), **kwargs), + ): + pass def test_restore_legacy_gdn_without_new_quantizer_handles(): @@ -307,42 +234,52 @@ def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): assert not model.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}]) - - -def _test_policy_matches_across_stages(rank, size): - config = { +def test_empty_stage_conversion_and_offline_restore(monkeypatch): + cfg = { "quant_cfg": [{"quantizer_name": "*", "enable": False}], "linear_attention": [{"module_name": "layer", "cfg": {"backend": "serving"}}], "algorithm": None, } - try: - # Only rank 1 owns a selected layer; both stages use the same recipe and phase. - layer = TinyGatedDeltaNet() if rank else nn.Linear(4, 4) - model = mtq.quantize(nn.ModuleDict({"layer": layer}), config) - with linear_attention_training_phase(model, [3]): - if rank: - assert model["layer"]._linear_attention_prefill_lengths == (3,) - if rank: - assert model["layer"]._linear_attention_prefill_lengths is None - - config["linear_attention"][0]["module_name"] = "missing" - with pytest.raises(ValueError, match="matches no supported"): - mtq.quantize(model, config) - finally: - torch.distributed.destroy_process_group() - - -def test_policy_matches_across_stages(skip_on_windows): - spawn_multiprocess_job(2, _test_policy_matches_across_stages, backend="gloo") + model = mtq.quantize(nn.ModuleDict({"layer": nn.Linear(4, 4)}), cfg) + with linear_attention_training_phase(model, [3]): + pass + state = mto.modelopt_state(model) + + # Saved resolved policies must restore without evaluating the original global recipe. + def no_rematching(*args): + pytest.fail("Restore must use saved resolved policies") + + monkeypatch.setattr( + "modelopt.torch.quantization.plugins.linear_attention._apply_linear_attention_policy", + no_rematching, + ) + restored = nn.ModuleDict({"layer": nn.Linear(4, 4)}) + mto.restore_from_modelopt_state(restored, state) + restored.load_state_dict(model.state_dict()) + x = torch.randn(2, 4) + torch.testing.assert_close(restored["layer"](x), model["layer"](x)) + + +def test_enabled_legacy_w_quantizer_is_rejected_on_restore(): + cfg = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None} + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + model.gdn_w_quantizer.enable() + with pytest.raises(ValueError, match="GDN W quantization is no longer supported"): + mto.restore_from_modelopt_state(TinyGatedDeltaNet(), mto.modelopt_state(model)) + + +def test_quant_cfg_refinement_and_policy_precedence(): + cfg = {**quant_cfg(), "algorithm": None} + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + cfg["linear_attention"] = [ + {"module_name": "*", "cfg": {"backend": "serving", "state_block_v": 128}}, + {"module_name": "", "cfg": {"backend": "serving", "state_block_v": 32}}, + ] + assert mtq.quantize(model, cfg) is model + assert model.gdn_state_qdq_block_v == 32 + cfg["linear_attention"].reverse() + mtq.quantize(model, cfg) + assert model.gdn_state_qdq_block_v == 128 + cfg["quant_cfg"][-1]["cfg"] = {**GDN_STATE_FP8_DYNAMIC, "pass_through_bwd": False} + with pytest.raises(ValueError, match="supports only"): + mtq.quantize(model, cfg) diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py index 3c8058df60c..aefe454398e 100644 --- a/tests/unit/torch/quantization/test_linear_attention_decode.py +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -16,9 +16,9 @@ import pytest import torch -from modelopt.torch.quantization.config import QuantizerAttributeConfig -from modelopt.torch.quantization.linear_attention import LinearAttentionConfig +from modelopt.torch.quantization.config import QuantizeConfig, QuantizerAttributeConfig from modelopt.torch.quantization.linear_attention.decode import _encode +from modelopt.torch.quantization.linear_attention.utils import _state_qdq from modelopt.torch.quantization.nn import TensorQuantizer @@ -36,6 +36,9 @@ def test_state_qdq_matches_tensor_quantizer(state_format): ) encoded = _encode(value, True, 16, state_format=state_format, state_quantizer=quantizer) torch.testing.assert_close(encoded.values, expected, rtol=0, atol=0) + torch.testing.assert_close( + _state_qdq(value, 16, state_format, quantizer), expected, rtol=0, atol=0 + ) probe = torch.randn_like(value) (gradient,) = torch.autograd.grad((encoded.values * probe).sum(), value) torch.testing.assert_close(gradient, probe, rtol=0, atol=0) @@ -46,12 +49,17 @@ def test_state_qdq_matches_tensor_quantizer(state_format): @pytest.mark.parametrize( "settings", [ + {"chunk_size": 32}, + {"solve": {"method": "neumann"}}, + {"state": {"mode": "token"}}, {"replay_window": 4}, {"precision": "replayssm", "replay_window": 0}, {"precision": "replayssm", "replay_window": 65}, {"precision": "replayssm", "state_block_v": 16}, ], ) -def test_unified_policy_rejects_incompatible_settings(settings): +def test_linear_attention_policy_rejects_invalid_settings(settings): with pytest.raises(ValueError): - LinearAttentionConfig(backend="serving", **settings) + QuantizeConfig( + linear_attention=[{"module_name": "*", "cfg": {"backend": "serving", **settings}}] + ) diff --git a/tests/unit/torch/quantization/test_linear_attention_reference.py b/tests/unit/torch/quantization/test_linear_attention_reference.py index 1744fc07d9c..d71818d0163 100644 --- a/tests/unit/torch/quantization/test_linear_attention_reference.py +++ b/tests/unit/torch/quantization/test_linear_attention_reference.py @@ -17,9 +17,9 @@ import torch import torch.nn.functional as F from _test_utils.torch.quantization.linear_attention_reference import ( - chunk_delta_rule_reference, + chunk_gdn_reference, recurrent_delta_rule_reference, - state_qdq_reference, + state_fp8_qdq_reference, ) @@ -34,11 +34,10 @@ def inputs(length=7, kda=False): return [x.requires_grad_() for x in (q, k, v, g, beta)] -@pytest.mark.parametrize( - ("kda", "packed", "state_v_first"), [(False, True, True), (True, False, False)] -) -def test_chunk_recurrence_output_state_and_gradients(kda, packed, state_v_first): - args = inputs(kda=kda) +@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() @@ -48,7 +47,7 @@ def test_chunk_recurrence_output_state_and_gradients(kda, packed, state_v_first) "state_v_first": state_v_first, "cu_seqlens": torch.tensor([0, 2, 7]) if packed else None, } - actual = chunk_delta_rule_reference(*args, chunk_size=3, **kwargs) + 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) @@ -99,11 +98,11 @@ def test_state_qdq_granularity_zero_tail_and_identity_ste(): with torch.no_grad(): state[0].zero_() state[1, ..., :16] *= 10 - quantized = state_qdq_reference(state, block_v=16) + 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_qdq_reference(state, block_v=64)) + 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) From dd016399ba80c64d2590145760e83141b7c09b04 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 9 Oct 2026 18:47:55 -0700 Subject: [PATCH 19/27] Fix FLA dependency setup for vLLM GPU tests Signed-off-by: Kai Xu --- noxfile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/noxfile.py b/noxfile.py index afa96010ac4..e30de083bd7 100644 --- a/noxfile.py +++ b/noxfile.py @@ -151,7 +151,7 @@ def gpu_trtllm(session): # Pin must stay in sync with examples/vllm_serve/Dockerfile. @nox.session(venv_backend="none") def gpu_vllm(session): - session.run("python3", "-m", "pip", "install", "-e", ".[hf,puzzletron,dev-test]") + session.run("python3", "-m", "pip", "install", "-e", ".[hf,puzzletron,dev-test,dev-fla]") session.run("python3", "-m", "pytest", "tests/gpu_vllm", *_cov_args()) From 140735b9628d4c23cf1150999575af4ed40650a7 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 9 Oct 2026 19:52:31 -0700 Subject: [PATCH 20/27] Fix linear attention CI dependency and Megatron compatibility Signed-off-by: Kai Xu --- noxfile.py | 5 ++++- .../quantization/plugins/test_megatron_gated_delta_net.py | 3 +++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/noxfile.py b/noxfile.py index e30de083bd7..4caa443ef9c 100644 --- a/noxfile.py +++ b/noxfile.py @@ -151,7 +151,10 @@ def gpu_trtllm(session): # Pin must stay in sync with examples/vllm_serve/Dockerfile. @nox.session(venv_backend="none") def gpu_vllm(session): - session.run("python3", "-m", "pip", "install", "-e", ".[hf,puzzletron,dev-test,dev-fla]") + # Keep vLLM's TileLang/TVM pins; dev-fla pins older versions for other test images. + session.run( + "python3", "-m", "pip", "install", "-e", ".[hf,puzzletron,dev-test]", "fla-core==0.5.1" + ) session.run("python3", "-m", "pytest", "tests/gpu_vllm", *_cov_args()) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index be822e18c35..c9ab51c77a0 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -151,6 +151,9 @@ def _test_gdn_qat_helper(rank, size, checkpoint_path): def _check_packed_recompute(module): + # Older Megatron releases lack selective GDN recompute and its native core hook. + if not hasattr(module, "recompute_gdn") or not hasattr(GatedDeltaNet, "_forward_compute"): + return hidden = torch.randn(128, 1, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) packed = PackedSeqParams( qkv_format="thd", From fee00a108b02580065b7b81392b46bb320d3f61d Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 10 Oct 2026 15:58:57 -0700 Subject: [PATCH 21/27] Guard linear attention GPU tests on older Triton and reduce KDA compilation Signed-off-by: Kai Xu --- .../torch/quantization/test_linear_attention_training.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py index b99981f9c6d..06337bb13f3 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py @@ -17,7 +17,9 @@ import pytest import torch +import triton from _test_utils.torch.quantization.linear_attention_reference import recurrent_delta_rule_reference +from packaging.version import Version from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.linear_attention import ( @@ -145,12 +147,15 @@ def layout(state): @pytest.fixture(scope="module", params=["gdn", "kda", "kimi_k3"]) def compiled_serving_case(request): """Compile one shared BF16 shape per model, outside the test-call timer.""" + # Skip before any kernels: an FLA backward fault poisons later tests' CUDA context. + if torch.cuda.get_device_capability()[0] >= 12 and Version(triton.__version__) < Version("3.7"): + pytest.skip("FLA backward causes misaligned CUDA access on sm_120+ with Triton < 3.7") kda = request.param != "gdn" kimi_k3 = request.param == "kimi_k3" if kimi_k3 and optional_module("vllm.models.kimi_k3.nvidia.ops.third_party.kda") is None: pytest.skip("This vLLM build does not provide the Kimi-K3 backend") torch.manual_seed(73) - keys = 128 if kda else 32 + keys = 128 if kimi_k3 else 32 args = [ torch.randn(1, 73, 1, dim, device="cuda", dtype=torch.bfloat16) # Rectangular GDN state catches transposes; equal padded tile widths also From 52b75c6e595dab4f00765c081f39ae6827bc1931 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 10 Oct 2026 19:14:45 -0700 Subject: [PATCH 22/27] Simplify linear attention state QAT and serving compatibility Signed-off-by: Kai Xu --- CHANGELOG.rst | 2 +- .../linear_attention/serving/__init__.py | 2 +- .../linear_attention/serving/_compat.py | 210 ------- .../linear_attention/serving/forward.py | 254 ++++----- .../linear_attention/serving/replay.py | 44 +- modelopt/torch/quantization/conversion.py | 3 +- .../quantization/linear_attention/__init__.py | 1 - .../linear_attention/_vllm_autograd.py | 4 +- .../quantization/linear_attention/config.py | 6 +- .../quantization/linear_attention/decode.py | 535 ++++-------------- .../quantization/linear_attention/gdn.py | 7 +- .../quantization/linear_attention/kda.py | 12 +- .../quantization/linear_attention/training.py | 55 +- .../quantization/linear_attention/utils.py | 96 +--- modelopt/torch/quantization/model_quant.py | 16 +- modelopt/torch/quantization/plugins/gdn.py | 5 - .../quantization/plugins/linear_attention.py | 8 +- .../torch/quantization/plugins/megatron.py | 49 +- modelopt_recipes/ptq.md | 32 +- .../plugins/test_megatron_gated_delta_net.py | 53 +- .../quantization/plugins/test_megatron_kda.py | 4 +- .../test_linear_attention_replay.py | 82 ++- .../test_linear_attention_training.py | 156 ++--- .../torch/quantization/plugins/test_gdn.py | 48 +- .../test_linear_attention_decode.py | 48 +- 25 files changed, 653 insertions(+), 1079 deletions(-) delete mode 100644 modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index e6b458b5bea..7dcbe340bff 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,7 +16,7 @@ Changelog - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. -- Add experimental GDN/KDA recurrent-state QAT with native vLLM Triton forwards and explicit prefill/decode phases through ``linear_attention_training_phase``. INT8 Hadamard and ReplaySSM require the compatible quantized-ReplaySSM serving fork. +- Add experimental GDN/KDA recurrent-state QAT with native vLLM 0.30.0 Triton forwards and explicit prefill/decode phases through ``linear_attention_training_phase``, including dynamic blockwise INT8 state quantization. Megatron users need a Core build with native gate and packed-forward hooks; 0.19.2 does not support this workflow. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before a separate KV-cache AutoQuantize stage, with independent resumable checkpoints for the weight and diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py index e808e5de96d..2f4c03234b5 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py @@ -20,6 +20,6 @@ except ImportError as error: raise ImportError( "Linear-attention serving profiles require the optional vLLM dependency: " - "public vLLM for 'vllm', or the compatible quantized-ReplaySSM " + "vllm==0.30.0 for 'vllm', or the compatible quantized-ReplaySSM " "vLLM fork for 'replayssm'." ) from error diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py b/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py deleted file mode 100644 index 418ca86e0e9..00000000000 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/_compat.py +++ /dev/null @@ -1,210 +0,0 @@ -# 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. - -"""Keep native vLLM imports and state layout changes at the kernel boundary.""" - -from functools import cache -from importlib import import_module -from importlib.util import find_spec -from inspect import signature - -import torch - -__all__ = [] - -_FLA_OPS = ( - "vllm.third_party.flash_linear_attention.ops" - if find_spec("vllm.third_party.flash_linear_attention") is not None - else "vllm.model_executor.layers.fla.ops" -) - - -@cache -def state_v_first(): - """Read the serving cache contract using unequal dimensions, independent of version tags.""" - from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator - - _, shape = MambaStateShapeCalculator.gated_delta_net_state_shape( - tp_world_size=1, - num_k_heads=1, - num_v_heads=1, - head_k_dim=32, - head_v_dim=64, - conv_kernel_size=4, - ) - if shape == (1, 64, 32): - return True - if shape == (1, 32, 64): - return False - raise NotImplementedError(f"Unsupported vLLM recurrent-state layout: {shape}") - - -def fla_module(name): - return import_module(f"{_FLA_OPS}.{name}") - - -@cache -def optional_module(path): - """Allow missing APIs, but never hide a broken dependency inside an installed module.""" - try: - return import_module(path) - except ModuleNotFoundError as error: - if error.name is not None and (path == error.name or path.startswith(error.name + ".")): - return None - raise - - -@cache -def kda_module(kimi_k3=False): - """Load only the explicitly selected KDA implementation, without backend fallback.""" - if not kimi_k3: - return fla_module("kda") - try: - return import_module("vllm.models.kimi_k3.nvidia.ops.third_party.kda") - except ImportError as error: - raise ImportError( - "The selected vLLM KDA Triton backend requires its model package dependencies. " - "Install the dependencies for this vLLM build; ModelOpt will not switch arithmetic " - "to a different KDA backend after an import failure." - ) from error - - -@cache -def gdn_module(name): - return optional_module(f"{_FLA_OPS}.{name}") - - -@cache -def kda_prefill_beta(kimi_k3=False): - """Resolve the raw-gate chunk callable's beta contract, independent of decode APIs.""" - kernel = getattr(kda_module(kimi_k3), "chunk_kda_with_fused_gate", None) - if kernel is None: - return None - parameters = signature(kernel).parameters - for name in ("raw_beta", "beta"): - if name in parameters: - return name - raise NotImplementedError("Unsupported vLLM chunk_kda_with_fused_gate beta signature") - - -def kda_chunk(q, k, v, g, beta, *, gate_inputs=None, kimi_k3=False, **kwargs): - """Use the same native chunk dispatch for execution and cache-contract checks.""" - native = kda_module(kimi_k3) - if gate_inputs is not None and (beta_name := kda_prefill_beta(kimi_k3)) is not None: - raw_g, raw_beta, rate, bias = gate_inputs - return native.chunk_kda_with_fused_gate( - q=q, - k=k, - v=v, - raw_g=raw_g, - **{beta_name: raw_beta if beta_name == "raw_beta" else beta}, - A_log=rate, - g_bias=bias, - **kwargs, - ) - if kimi_k3: - raise ValueError("vllm_kimi_k3 requires native fused-gate prefill and raw gate_inputs") - return native.chunk_kda(q, k, v, g, beta, **kwargs) - - -@cache -def _kda_state_v_first(device, kimi_k3, keys, raw_gates=False): - """Check the selected KDA paths at the actual head size, once per device.""" - with torch.no_grad(): - q = torch.zeros(1, 1, 1, keys, device=device, dtype=torch.bfloat16) - q[..., 0] = 1 - zero = torch.zeros_like(q) - # Distinct row/column readouts made from exactly representable BF16 values, - # including head sizes such as 192 where arange(K*K) would round. - pattern = torch.zeros(keys, keys, device=device, dtype=torch.float32) - markers = torch.arange(1, keys + 1, device=device, dtype=torch.float32) - pattern[0, :] = -markers - pattern[:, 0] = markers - state = pattern.expand(2, 1, keys, keys).contiguous() - output, _ = kda_module(kimi_k3).fused_recurrent_kda_fwd( - q=q, - k=zero, - v=zero, - g=zero.float(), - beta=torch.zeros(1, 1, 1, device=device), - scale=1.0, - initial_state=state, - inplace_final_state=True, - cu_seqlens=torch.arange(2, device=device, dtype=torch.int32), - ssm_state_indices=torch.ones(1, device=device, dtype=torch.int32), - use_qk_l2norm_in_kernel=False, - ) - row = output.flatten().float() - if torch.equal(row, pattern[:, 0]): - value_first = True - elif torch.equal(row, pattern[0]): - value_first = False - else: - raise NotImplementedError("The installed vLLM KDA cache layout could not be validated") - # Some native revisions changed cache storage without updating chunk readout. - # Reject that backend rather than training against a silently transposed prefix. - chunk_args = { - "q": q, - "k": zero, - "v": zero.clone(), - "scale": 1.0, - "initial_state": state[1:].clone(), - "output_final_state": True, - "use_qk_l2norm_in_kernel": kimi_k3, - "cu_seqlens": torch.arange(2, device=device, dtype=torch.int32), - } - gate_inputs = None - if raw_gates or kimi_k3: - raw_g = torch.full_like(q, -80) - raw_beta = torch.full((1, 1, 1), -80.0, device=device) - rate, bias = torch.zeros(1, device=device), torch.zeros(keys, device=device) - gate_inputs = (raw_g, raw_beta, rate, bias) - if kimi_k3: - packed, _ = kda_module(True).fused_recurrent_kda_packed_decode( - mixed_qkv=torch.cat((q.flatten(), zero.flatten(), zero.flatten()))[None], - raw_g=raw_g, - raw_beta=raw_beta, - A_log=rate, - dt_bias=bias, - lower_bound=None, - initial_state=state.clone(), - state_indices=torch.ones(1, device=device, dtype=torch.int32), - scale=1.0, - ) - if not torch.equal(packed.flatten().float(), row): - raise NotImplementedError("The installed vLLM KDA packed cache layout disagrees") - chunk_output, _ = kda_chunk( - g=zero.float(), - beta=torch.zeros(1, 1, 1, device=device), - gate_inputs=gate_inputs, - kimi_k3=kimi_k3, - **chunk_args, - ) - if not torch.equal(chunk_output.flatten().float(), row): - raise NotImplementedError( - "The installed vLLM KDA chunk/decode cache readout disagrees; " - "use a vLLM build with the KDA chunk-state layout fix" - ) - return value_first - - -def state_layout(state, *, kda=False, kimi_k3=False, raw_gates=False): - """Convert between ModelOpt and the selected model's native state layout.""" - value_first = ( - _kda_state_v_first(state.device, kimi_k3, state.shape[-1], raw_gates) - if kda - else state_v_first() - ) - return state.transpose(-1, -2).contiguous() if value_first else state.contiguous() diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py index 7c54c755f17..84913b224a7 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py @@ -13,153 +13,114 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Native vLLM Triton forwards; training owns private, differentiable state.""" +"""Native vLLM 0.30 Triton forwards; training owns private, differentiable state.""" from functools import cache, lru_cache import torch -from ._compat import fla_module, gdn_module, kda_chunk, kda_module, optional_module, state_layout + +@cache +def _kernels(precision): + # Defer optional serving/model dependencies until their explicit profile is selected. + if precision == "replayssm": + from .replay import _prefill_kernels + + return _prefill_kernels() + else: + import vllm + + if vllm.__version__ != "0.30.0": + raise ImportError("Linear-attention vLLM profiles require vllm==0.30.0") + from vllm.third_party.flash_linear_attention import ops as gdn_ops + + if precision == "vllm_kimi_k3": + from vllm.models.kimi_k3.nvidia.ops.third_party import kda as kda_ops + else: + from vllm.third_party.flash_linear_attention.ops import kda as kda_ops + return gdn_ops, kda_ops -def _gdn_inputs(q, k, v, gate_inputs, normalize): - raw_g, raw_beta, a_log, bias = gate_inputs - prep = gdn_module("fused_gdn_prefill_post_conv") - if prep is not None: - return prep.fused_post_conv_prep( +def _state_layout(state): + # ModelOpt uses [..., K, V]; both supported native backends store [..., V, K]. + return state.transpose(-1, -2).contiguous() + + +def prefill(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, precision="vllm"): + """Run native chunk prefill, including the selected profile's raw-gate preparation.""" + gdn, kda = _kernels(precision) + kimi_k3 = precision == "vllm_kimi_k3" + channel = g.ndim == 3 + if not channel and gate_inputs is not None: + raw_g, raw_beta, rate, bias = gate_inputs + q, k, v, g, beta = gdn.fused_post_conv_prep( torch.cat([x.flatten(1) for x in (q, k, v)], dim=-1), raw_g.flatten(0, -2), raw_beta.flatten(0, -2), - a_log, + rate, bias, q.shape[-2], q.shape[-1], v.shape[-1], apply_l2norm=normalize, - ), False - g, beta = _gdn_gate()( - a_log, raw_g.reshape(-1, raw_g.shape[-1]), raw_beta.reshape(-1, raw_beta.shape[-1]), bias - ) - return (q, k, v, g[0], beta[0]), normalize - - -def prefill(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, kimi_k3=False): - """Run the complete native chunk entry point, including raw-gate preparation.""" - if kimi_k3 and gate_inputs is None: - raise ValueError("vllm_kimi_k3 requires raw gate_inputs") - channel = g.ndim == 3 - if not channel and gate_inputs is not None: - (q, k, v, g, beta), normalize = _gdn_inputs(q, k, v, gate_inputs, normalize) + ) + normalize = False q, k, v, g, beta = [x.unsqueeze(0).contiguous() for x in (q, k, v, g, beta)] - # Some native chunk implementations reuse V as their output buffer. + # Native chunk implementations may reuse V as their output buffer. v = v.clone() kwargs = { "scale": scale, - "initial_state": state_layout( - state.unsqueeze(0), kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None - ), + "initial_state": _state_layout(state.unsqueeze(0)), "output_final_state": True, "use_qk_l2norm_in_kernel": normalize, "cu_seqlens": _prefill_metadata(q.device, q.shape[1]), } - if channel: - gates = None - if gate_inputs is not None: - raw_g, raw_beta, rate, bias = gate_inputs - gates = (raw_g[None].contiguous(), raw_beta[None].contiguous(), rate, bias) - out, final = kda_chunk(q, k, v, g, beta, gate_inputs=gates, kimi_k3=kimi_k3, **kwargs) + if channel and gate_inputs is not None: + raw_g, raw_beta, rate, bias = gate_inputs + # Kimi fuses beta activation; standalone KDA takes already activated beta. + beta_arg = {"raw_beta": raw_beta[None].contiguous()} if kimi_k3 else {"beta": beta} + out, final = kda.chunk_kda_with_fused_gate( + q, + k, + v, + raw_g=raw_g[None].contiguous(), + A_log=rate, + g_bias=bias, + **beta_arg, + **kwargs, + ) + elif channel: + out, final = kda.chunk_kda(q, k, v, g, beta, **kwargs) else: - out, final = fla_module("chunk").chunk_gated_delta_rule(q, k, v, g, beta, **kwargs) - return out[0], state_layout( - final[0], kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None - ) + out, final = gdn.chunk_gated_delta_rule(q, k, v, g, beta, **kwargs) + return out[0], _state_layout(final[0]) -def step(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, kimi_k3=False): - """Run native fused decode on a private cache; never mutate the incoming state.""" - if kimi_k3 and gate_inputs is None: - raise ValueError("vllm_kimi_k3 requires raw gate_inputs") +def step(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, precision="vllm"): + """Run native fused decode on a private cache, preserving the incoming autograd state.""" + gdn, kda = _kernels(precision) + kimi_k3 = precision == "vllm_kimi_k3" channel = g.ndim == 2 cu, indices = _decode_metadata(q.device) - native_state = state_layout( - state.unsqueeze(0), kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None - ) - if gate_inputs is not None: - raw_g, raw_beta, a_log, bias = gate_inputs - if channel and kimi_k3: - if not normalize: - raise ValueError("vLLM packed KDA decode requires Q/K normalization") - # Slot zero is reserved by vLLM's paged-cache kernels. - native_state = torch.cat((torch.zeros_like(native_state), native_state), dim=0) - out, final = kda_module(kimi_k3).fused_recurrent_kda_packed_decode( - mixed_qkv=torch.cat([x.flatten() for x in (q, k, v)])[None], - raw_g=raw_g[None, None].contiguous(), - raw_beta=raw_beta[None, None].contiguous(), - A_log=a_log.reshape(-1).contiguous(), - dt_bias=(torch.zeros_like(raw_g).flatten() if bias is None else bias.contiguous()), - lower_bound=None, - initial_state=native_state, - state_indices=indices, - scale=scale, - ) - return out[0, 0], state_layout( - final[1], kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None - ) - packed = ( - None - if channel - else getattr( - fla_module("fused_recurrent"), - "fused_recurrent_gated_delta_rule_packed_decode", - None, - ) - ) - if packed is not None: - # Match vLLM's pure, non-speculative decode path on private cache slot one. - native_state = torch.cat((torch.zeros_like(native_state), native_state), dim=0) - out, final = packed( - mixed_qkv=torch.cat([x.flatten() for x in (q, k, v)])[None], - a=raw_g[None].contiguous(), - b=raw_beta[None].contiguous(), - A_log=a_log.reshape(-1).contiguous(), - dt_bias=bias.contiguous(), - scale=scale, - initial_state=native_state, - out=torch.empty_like(v)[None, None], - ssm_state_indices=indices, - use_qk_l2norm_in_kernel=normalize, - ) - return out[0, 0], state_layout( - final[1], kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None - ) - fused = None if channel else gdn_module("fused_sigmoid_gating") - if fused is not None: - out, final = fused.fused_sigmoid_gating_delta_rule_update( - A_log=a_log, - a=raw_g[None].contiguous(), - b=raw_beta[None].contiguous(), - dt_bias=bias, - q=q[None, None], - k=k[None, None], - v=v[None, None], - initial_state=native_state, - inplace_final_state=False, - scale=scale, - use_qk_l2norm_in_kernel=normalize, - cu_seqlens=cu, - ) - return out[0, 0], state_layout( - final[0], kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None - ) - if not channel: - g, beta = _gdn_gate()( - a_log, raw_g[None].contiguous(), raw_beta[None].contiguous(), bias - ) - g, beta = g[0, 0], beta[0, 0] - # Older vLLM models and standalone calls consume already activated gates. - if channel: + native_state = _state_layout(state.unsqueeze(0)) + if channel or gate_inputs is not None: + # Slot zero is reserved by vLLM's paged-cache kernels. native_state = torch.cat((torch.zeros_like(native_state), native_state), dim=0) - out, final = kda_module(kimi_k3).fused_recurrent_kda_fwd( + if kimi_k3: + raw_g, raw_beta, rate, bias = gate_inputs + out, final = kda.fused_recurrent_kda_packed_decode( + mixed_qkv=torch.cat([x.flatten() for x in (q, k, v)])[None], + raw_g=raw_g[None, None].contiguous(), + raw_beta=raw_beta[None, None].contiguous(), + A_log=rate.reshape(-1).contiguous(), + dt_bias=(torch.zeros_like(raw_g).flatten() if bias is None else bias.contiguous()), + lower_bound=None, + initial_state=native_state, + state_indices=indices, + scale=scale, + ) + elif channel: + out, final = kda.fused_recurrent_kda_fwd( *[x[None, None].contiguous() for x in (q, k, v, g, beta)], initial_state=native_state, inplace_final_state=True, @@ -168,36 +129,30 @@ def step(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, k cu_seqlens=cu, ssm_state_indices=indices, ) - return out[0, 0], state_layout( - final[1], kda=True, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None + elif gate_inputs is not None: + raw_g, raw_beta, rate, bias = gate_inputs + out, final = gdn.fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=torch.cat([x.flatten() for x in (q, k, v)])[None], + a=raw_g[None].contiguous(), + b=raw_beta[None].contiguous(), + A_log=rate.reshape(-1).contiguous(), + dt_bias=bias.contiguous(), + scale=scale, + initial_state=native_state, + out=torch.empty_like(v)[None, None], + ssm_state_indices=indices, + use_qk_l2norm_in_kernel=normalize, ) - recurrent = fla_module("fused_recurrent").fused_recurrent_gated_delta_rule - out, final = recurrent( - *[x[None, None].contiguous() for x in (q, k, v, g, beta)], - initial_state=native_state, - inplace_final_state=False, - scale=scale, - use_qk_l2norm_in_kernel=normalize, - cu_seqlens=cu, - ssm_state_indices=None, - ) - return out[0, 0], state_layout( - final[0], kda=channel, kimi_k3=kimi_k3, raw_gates=gate_inputs is not None - ) - - -# Old runtimes keep gate preparation beside the model rather than in the FLA package. -@cache -def _gdn_gate(): - for path in ( - "vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn", - "vllm.model_executor.layers.mamba.gdn_linear_attn", - "vllm.model_executor.models.qwen3_next", - ): - module = optional_module(path) - if module is not None and hasattr(module, "fused_gdn_gating"): - return module.fused_gdn_gating - raise NotImplementedError("The installed vLLM does not expose supported GDN gate preparation") + else: + out, final = gdn.fused_recurrent_gated_delta_rule( + *[x[None, None].contiguous() for x in (q, k, v, g, beta)], + initial_state=native_state, + inplace_final_state=False, + scale=scale, + use_qk_l2norm_in_kernel=normalize, + cu_seqlens=cu, + ) + return out[0, 0], _state_layout(final[-1]) @lru_cache(maxsize=128) @@ -214,9 +169,10 @@ def _decode_metadata(device): ) -def validate_profile(gate_inputs, normalize, *, kimi_k3=False): +def validate_profile(gate_inputs, normalize, *, precision="vllm"): """Validate the selected native serving contract before launching prefill.""" - if kimi_k3: + _kernels(precision) + if precision == "vllm_kimi_k3": if gate_inputs is None: raise ValueError("vllm_kimi_k3 requires raw gate_inputs") if not normalize: diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/replay.py b/modelopt/torch/kernels/quantization/linear_attention/serving/replay.py index 17cd816a758..67f43db24a9 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/replay.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/replay.py @@ -16,11 +16,17 @@ """Functional training launches of the optional quantized-ReplaySSM serving kernels.""" import torch -from vllm.model_executor.layers.fla.ops.fused_recurrent_replayssm import ( - _apply_h_value_axis, - fused_recurrent_gated_delta_rule_replayssm, - prefill_write_checkpoint, -) + +# This optional fork has an explicit API separate from the public vLLM 0.30 profile. +from vllm.model_executor.layers.fla.ops import fused_recurrent_replayssm as _native + + +def _prefill_kernels(): + # Keep this fork's import contract isolated from the public vLLM profile. + from vllm.model_executor.layers.fla import ops + from vllm.model_executor.layers.fla.ops import kda + + return ops, kda def _storage(value, enabled): @@ -50,33 +56,33 @@ def checkpoint(value, enabled): state, scales, indices = _storage(value, enabled) original = value.transpose(-1, -2).unsqueeze(0).contiguous().float() if enabled: - prefill_write_checkpoint( + _native.prefill_write_checkpoint( original, state, None, None, scales, None, None, indices, 8, hadamard=True ) else: - state[1] = _apply_h_value_axis(original)[0] + state[1] = _native._apply_h_value_axis(original)[0] return _decoded(state, scales) def original_basis(value): """Use serving's FP32 checkpoint transform when returning a dense state.""" - return _apply_h_value_axis(value.transpose(-1, -2).unsqueeze(0).contiguous())[0].transpose( - -1, -2 - ) + return _native._apply_h_value_axis(value.transpose(-1, -2).unsqueeze(0).contiguous())[ + 0 + ].transpose(-1, -2) def step(q, k, v, gate, beta, carry, window, enabled, scale, normalize, gate_inputs): """Launch one token on private cache buffers; never mutate tensors saved for backward.""" anchor = carry.anchor - state, scales, indices = _storage(anchor.values, enabled) + state, scales, indices = _storage(anchor, enabled) if enabled: - metadata = anchor.scales - codes = (anchor.values.unflatten(-1, (-1, 32)) / metadata.float().unsqueeze(-1)).round() + metadata = carry.scales + codes = (anchor.unflatten(-1, (-1, 32)) / metadata.float().unsqueeze(-1)).round() state[1] = codes.flatten(-2).transpose(-1, -2).to(torch.int8) scales[1] = metadata.transpose(-1, -2) else: - state[1] = anchor.values.transpose(-1, -2) - heads, keys, values = anchor.values.shape + state[1] = anchor.transpose(-1, -2) + heads, keys, values = anchor.shape d = torch.zeros(2, heads, window, values, device=q.device, dtype=torch.bfloat16) kc = torch.zeros(2, heads, window, keys, device=q.device, dtype=torch.bfloat16) channel = gate.ndim == 2 @@ -86,8 +92,8 @@ def step(q, k, v, gate, beta, carry, window, enabled, scale, normalize, gate_inp dtype=torch.float32, ) for i, entry in enumerate(carry.entries): - d[1, :, i] = entry.update.values - kc[1, :, i] = entry.key.values + d[1, :, i] = entry.update + kc[1, :, i] = entry.key gc[1, :, i] = entry.log_retention mixed = torch.cat([x.to(torch.bfloat16).flatten() for x in (q, k, v)])[None] if channel: @@ -96,8 +102,8 @@ def step(q, k, v, gate, beta, carry, window, enabled, scale, normalize, gate_inp raw_gate, raw_beta, a_log, bias = gate_inputs a, b = raw_gate[None].contiguous(), raw_beta[None].contiguous() output = torch.empty(1, heads, values, device=q.device, dtype=torch.bfloat16) - cursor = carry.cursor - fused_recurrent_gated_delta_rule_replayssm( + cursor = len(carry.entries) + _native.fused_recurrent_gated_delta_rule_replayssm( mixed, a, b, diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index c77ab1b3623..9ca1781e224 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -163,8 +163,7 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: return model quantizer_state_dict = dict(metadata["quantizer_state"]) - if "linear_attention" not in metadata: - _restore_legacy_linear_attention_quantizers(model, quantizer_state_dict) + _restore_legacy_linear_attention_quantizers(model, quantizer_state_dict) unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys() extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys() diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py index 42bda38f4e4..56e50a5174b 100644 --- a/modelopt/torch/quantization/linear_attention/__init__.py +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -16,7 +16,6 @@ """Numerical policies and differentiable kernels for linear attention.""" from .config import * -from .decode import * from .gdn import * from .kda import * from .training import * diff --git a/modelopt/torch/quantization/linear_attention/_vllm_autograd.py b/modelopt/torch/quantization/linear_attention/_vllm_autograd.py index c5165daaae1..c32c213ff51 100644 --- a/modelopt/torch/quantization/linear_attention/_vllm_autograd.py +++ b/modelopt/torch/quantization/linear_attention/_vllm_autograd.py @@ -46,7 +46,7 @@ def prefix( scale, normalize, gate_inputs, - kimi_k3=precision == "vllm_kimi_k3", + precision=precision, ) if not torch.is_grad_enabled() or not any(x.requires_grad for x in (q, k, v, g, beta, state)): return out.float(), final @@ -83,7 +83,7 @@ def step(q, k, v, gate, beta, state, scale, normalize=False, gate_inputs=None, p scale, normalize, gate_inputs, - kimi_k3=precision == "vllm_kimi_k3", + precision=precision, ) if not torch.is_grad_enabled() or not any( x.requires_grad for x in (q, k, v, gate, beta, state) diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index f1536a3510c..1ee036ee809 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -28,11 +28,11 @@ class LinearAttentionConfig(ModeloptBaseConfig): """One execution policy shared by the chunked prefix and recurrent suffix. ``serving`` uses native forward arithmetic with a differentiable adjoint; - ``fla`` is the disabled/default path. ``vllm`` targets the installed vLLM's + ``fla`` is the disabled/default path. ``vllm`` targets vLLM 0.30.0's Triton prefill and single-token, non-speculative decode with BF16 operands and FP32 recurrent storage. Deployment must explicitly select Triton prefill - (``gdn_prefill_backend="triton"``) and enable packed recurrent decode where - available (``VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE=1``). Mixed prefill/decode + (``gdn_prefill_backend="triton"``) and enable packed recurrent decode + (``VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE=1``). Mixed prefill/decode batches and automatic FlashInfer/CuTe dispatch are outside this profile. KDA uses the standalone FLA/Triton API, with no model-level serving guarantee. ``vllm_kimi_k3`` explicitly selects diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py index cac33968d4b..7cc94d11170 100644 --- a/modelopt/torch/quantization/linear_attention/decode.py +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -19,38 +19,14 @@ import torch -from .config import LinearAttentionConfig -from .utils import ( - _resolve_state_quantizer, - _state_qdq_with_scales, - _validate_gate_inputs, - forward_value, -) +from ._vllm_autograd import normalized +from .utils import _state_qdq, forward_value -__all__ = [ - "EncodedLinearAttentionTensor", - "LinearAttentionState", - "ReplayEntry", - "recurrent_decode", -] +__all__ = [] @dataclass -class EncodedLinearAttentionTensor: - """Floating QDQ values with autograd history and optional quantization scales. - - ``format`` and ``block_v`` describe the emulated encoding; ``values`` stores - decoded values, not packed INT8/FP8 codes. Scale metadata is non-differentiable. - """ - - values: torch.Tensor - scales: torch.Tensor | None - format: str - block_v: int | None - - -@dataclass -class ReplayEntry: +class _ReplayEntry: """One saved rank-one update since the last state checkpoint. ``key`` is [H,K] and ``update`` is the beta-scaled residual [H,V]. ReplaySSM @@ -58,51 +34,29 @@ class ReplayEntry: [H,K] for KDA and records this token's decay for subsequent reconstruction. """ - key: EncodedLinearAttentionTensor - update: EncodedLinearAttentionTensor + key: torch.Tensor + update: torch.Tensor log_retention: torch.Tensor @dataclass -class LinearAttentionState: - """Runtime state passed between recurrent tokens and successive decode calls. +class _ReplayState: + """Checkpoint and pending updates within one training suffix. - With ``replay_window=1``, every token refreshes ``anchor`` and ``entries`` - stays empty. With a larger window, the anchor stays fixed while updates - accumulate; reaching the window refreshes the anchor and clears the entries. - Tensor values retain their autograd history for QAT/QAD across call boundaries. + Refreshing the replay window replaces the checkpoint and clears the entries. + Tensors retain their autograd history; this state never crosses forward calls. """ - anchor: EncodedLinearAttentionTensor # Checkpoint [H,K,V]; QDQ is applied when enabled. - entries: tuple[ReplayEntry, ...] # Updates since the checkpoint, in token order. - position: int # Next token position, including the supplied prefix offset. - started: bool # Initial-state encoding has run; resuming must not repeat it. - signature: str # Execution/quantizer settings used to check compatible reuse. - value_basis: str = "identity" # Anchor and updates may use the Hadamard value basis. - precision: str = "vllm" # Selects the native profile's reconstruction arithmetic. - - @property - def cursor(self): - """Number of encoded updates since the last anchor refresh.""" - return len(self.entries) + anchor: torch.Tensor # Decoded checkpoint [H,K,V] in the Hadamard value basis. + scales: torch.Tensor | None # Native FP16 scales needed to rebuild an INT8 cache. + entries: tuple[_ReplayEntry, ...] = () - def reconstruct(self, *, original_basis=True): - """Return the current dense [H,K,V] state without refreshing or quantizing it. + def reconstruct(self): + """Return the dense state in the original value basis without another QDQ.""" + # The optional serving fork owns the checkpoint basis transform. + from ...kernels.quantization.linear_attention.serving.replay import original_basis - Apply pending updates and decay to the anchor, preserving gradients. - Decode uses ``original_basis=False`` to continue in the stored basis; - the default converts Hadamard state back for the caller's final-state output. - Keep this object, rather than only this tensor, to resume a replay window. - """ - state = _replay_state(self) if self.precision == "replayssm" else self.anchor.values - if original_basis and self.value_basis == "hadamard32": - # The optional serving fork owns the checkpoint basis transform. - from ...kernels.quantization.linear_attention.serving.replay import ( - original_basis as restore_basis, - ) - - state = restore_basis(state) - return state + return original_basis(_replay_state(self)) def _hadamard32(value): @@ -115,20 +69,20 @@ def _hadamard32(value): return value * (32**-0.5) -def _replay_state(carry: LinearAttentionState, current_gate=None): +def _replay_state(carry: _ReplayState, current_gate=None): """Decay the anchor and add saved updates weighted by their subsequent gates. KDA weights keys per channel; GDN weights updates per head. ``current_gate`` includes the next KDA token's decay before BF16 weighting, without adding its update. """ - state = carry.anchor.values + state = carry.anchor if not carry.entries: if current_gate is not None: state = state * current_gate.exp().unsqueeze(-1) return state gates = torch.stack([entry.log_retention for entry in carry.entries], dim=1) - keys = torch.stack([entry.key.values for entry in carry.entries], dim=1) - updates = torch.stack([entry.update.values for entry in carry.entries], dim=1) + keys = torch.stack([entry.key for entry in carry.entries], dim=1) + updates = torch.stack([entry.update for entry in carry.entries], dim=1) total = gates.sum(1) if current_gate is not None: total = total + current_gate @@ -142,31 +96,6 @@ def _replay_state(carry: LinearAttentionState, current_gate=None): return state * total.exp()[:, None, None] + keys.transpose(-1, -2) @ weighted -def _encode( - value, - enabled, - block_v, - *, - state_format="fp8_e4m3", - state_quantizer=None, -): - """Apply TensorQuantizer state QDQ with straight-through gradients.""" - if not enabled: - return EncodedLinearAttentionTensor(value, None, "identity", None) - decoded, scales = _state_qdq_with_scales(value, block_v, state_format, state_quantizer) - if state_quantizer is not None and state_quantizer.block_sizes is not None: - block_v = state_quantizer.block_sizes[-1] - return EncodedLinearAttentionTensor(decoded, scales, state_format, block_v) - - -def _signature(config, state_qdq, block_v, state_format, state_quantizer): - """Record the execution and quantizer settings required when resuming a state.""" - signature = config.model_dump_json() + f"/{state_qdq}/{block_v}/{state_format}" - if state_quantizer is not None and state_quantizer.block_sizes is not None: - signature += f"/group={state_quantizer.block_sizes[-1]}" - return signature - - def _sum_keys(value): keys = value.shape[-2] padded = 1 << (keys - 1).bit_length() @@ -178,168 +107,6 @@ def _sum_keys(value): return value[..., 0, :] -def _prepare_carry( - q, - k, - v, - g, - beta, - config, - state_qdq, - block_v, - initial_state, - carry, - position, - state_format, - state_quantizer, -): - """Initialize from a dense state or validate and resume an existing runtime state. - - The first nonempty call applies the handoff transform/QDQ once. A resumed - state preserves its anchor, replay cursor, and token position; an empty call - leaves initial encoding deferred until a token is actually processed. - """ - 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_quantizer is not None and state_quantizer.block_sizes is not None: - raise ValueError("TensorQuantizer block_sizes requires state_codec='tile'") - 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, state_quantizer) - 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 = LinearAttentionState( - _encode(initial_state, False, block_v, 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 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 - if config.precision == "replayssm": - # The optional serving fork owns checkpoint rounding and Hadamard arithmetic. - from ...kernels.quantization.linear_attention.serving.replay import checkpoint - - with torch.no_grad(): - decoded, metadata = checkpoint(carry.anchor.values, state_qdq) - anchor = EncodedLinearAttentionTensor( - forward_value(initial, decoded), - metadata, - "int8" if state_qdq else "identity", - 32 if state_qdq else None, - ) - else: - anchor = _encode( - initial, - state_qdq, - block_v, - state_format=state_format, - state_quantizer=state_quantizer, - ) - carry = LinearAttentionState( - anchor, - (), - carry.position, - True, - signature, - "hadamard32" if hadamard else "identity", - config.precision, - ) - return carry, signature - - -def recurrent_decode( - q, - k, - v, - g, - beta, - *, - config: LinearAttentionConfig, - state_qdq=False, - state_format="fp8_e4m3", - state_quantizer=None, - initial_state=None, - carry: LinearAttentionState | None = None, - position=0, - scale=None, - use_qk_l2norm_in_kernel=False, - gate_inputs=None, -) -> tuple[torch.Tensor, LinearAttentionState]: - """Run a recurrent suffix and return its outputs and resumable runtime state. - - Q/K are [T,H,K], V is [T,H,V], and beta is [T,H]; heads must be aligned. - Log gates are [T,H] for GDN or [T,H,K] for KDA. The first token consumes the - handoff state after enabled QDQ. Each token reads its output from the working - state before checkpoint rounding; replay_window determines when that state - is saved. Inputs and the initial state are promoted to FP32 working values; - native forward values use a differentiable adjoint for training. - - Args: - config: The same LinearAttentionConfig used by the chunked prefix. - initial_state: Dense [H,K,V] state in the original value basis, normally - from prefill. Defaults to zeros; mutually exclusive with ``carry``. - carry: LinearAttentionState returned by an earlier call. Pass it directly - to preserve the replay window and autograd history across calls. - position: Starting token offset when creating a state, usually the prefix - length. An existing carry retains its own position. - - Returns: - Outputs [T,H,V] and the updated LinearAttentionState, both retaining their - graphs. An empty call performs no state write or initial-state QDQ. - """ - state_quantizer, state_qdq, state_format = _resolve_state_quantizer( - state_quantizer, state_qdq, state_format - ) - _validate_gate_inputs(g, beta, gate_inputs) - if config.precision == "vllm_kimi_k3" and g.ndim != 3: - raise ValueError("vllm_kimi_k3 is a KDA-only precision profile") - if config.precision != "replayssm": - # Standalone decode validates before any cache writes. - from ...kernels.quantization.linear_attention.serving.forward import validate_profile - - validate_profile( - gate_inputs, - use_qk_l2norm_in_kernel, - kimi_k3=config.precision == "vllm_kimi_k3", - ) - return _recurrent_decode( - q, - k, - v, - g, - beta, - config=config, - state_qdq=state_qdq, - state_format=state_format, - state_quantizer=state_quantizer, - initial_state=initial_state, - carry=carry, - position=position, - scale=scale, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - gate_inputs=gate_inputs, - ) - - def _recurrent_decode( q, k, @@ -348,181 +115,125 @@ def _recurrent_decode( beta, *, config, - state_qdq, - state_format, state_quantizer, - initial_state=None, - carry=None, - position=0, + initial_state, + output_final_state=False, scale=None, use_qk_l2norm_in_kernel=False, gate_inputs=None, ): - """Execute a suffix after its caller validates the shared quantizer once.""" - if config.backend != "serving": - raise ValueError("State QAT requires backend='serving'") - block_v = config.state_block_v - serving = config.precision != "replayssm" - native_replay = config.precision == "replayssm" - if q.device.type != "cuda": - raise ValueError("Serving arithmetic requires CUDA") - # BF16 inputs need FP32 working values for both replay reconstruction and its adjoint. - q, k, v, g, beta = (x.float() for x in (q, k, v, g, beta)) - if initial_state is not None: - initial_state = initial_state.float() - if ( - native_replay - and len(q) - and ( - q.shape[-1] < 32 - or q.shape[-1] & (q.shape[-1] - 1) - or (g.ndim == 2 and gate_inputs is None) - ) - ): - raise ValueError("ReplaySSM requires power-of-two K >= 32 and raw GDN gate inputs") - carry, signature = _prepare_carry( - q, - k, - v, - g, - beta, - config, - state_qdq, - block_v, - initial_state, - carry, - position, - state_format, - state_quantizer, - ) + """Run one training suffix and optionally return its dense final state. + + The caller validates shapes, policy and quantizer, and supplies FP32 inputs and prefill state. + Empty suffixes preserve that state without a handoff QDQ or native launch. + """ 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 - original_v = v - if config.state_codec == "int8_hadamard32": - v = _hadamard32(v) + return v + zero, initial_state if output_final_state else None scale = q.shape[-1] ** -0.5 if scale is None else scale - outputs = [] - native_outputs = [] + if config.precision == "replayssm": + output, carry = _replay_decode( + q, + k, + v, + g, + beta, + initial_state, + config.replay_window, + state_quantizer is not None, + scale, + use_qk_l2norm_in_kernel, + gate_inputs, + ) + return output, carry.reconstruct() if output_final_state else None + # Import the optional native backend only for the ordinary state-QDQ path. + from ._vllm_autograd import step + + state = _state_qdq(initial_state, state_quantizer, config.state_block_v) with torch.autocast(device_type=q.device.type, enabled=False): + outputs = [] + for t in range(len(q)): + # The output reads the working state before QDQ writes the next token's cache. + output, working = step( + q[t], + k[t], + v[t], + g[t], + beta[t], + state, + scale, + use_qk_l2norm_in_kernel, + precision=config.precision, + gate_inputs=(gate_inputs[0][t], gate_inputs[1][t], *gate_inputs[2:]) + if gate_inputs is not None + else None, + ) + outputs.append(output) + # A new tensor preserves the state saved by earlier tokens' backward passes. + state = _state_qdq(working, state_quantizer, config.state_block_v) + return torch.stack(outputs), state if output_final_state else None + + +def _replay_decode(q, k, v, g, beta, initial_state, window, enabled, scale, normalize, gate_inputs): + """Attach native ReplaySSM values to the differentiable checkpoint/update recurrence.""" + # This optional fork owns checkpoint rounding and the Hadamard value basis. + from ...kernels.quantization.linear_attention.serving.replay import checkpoint, step + + initial = _hadamard32(initial_state) + with torch.no_grad(): + decoded, scales = checkpoint(initial_state, enabled) + carry = _ReplayState(forward_value(initial, decoded), scales) + original_v, v = v, _hadamard32(v) + with torch.autocast(device_type=q.device.type, enabled=False): + outputs, native_outputs = [], [] for t in range(len(q)): - entries = carry.entries gate = g[t] - if native_replay: - from ...kernels.quantization.linear_attention.serving.replay import step - - with torch.no_grad(): - native = step( - q[t], - k[t], - original_v[t], - gate, - beta[t], - carry, - config.replay_window, - state_qdq, - scale, - use_qk_l2norm_in_kernel, - ( - gate_inputs[0][t], - gate_inputs[1][t], - *gate_inputs[2:], - ) - if gate_inputs is not None - else None, - ) - native_outputs.append(native[0]) - if serving: - from ._vllm_autograd import step as serving_step - - native_output, working = serving_step( + with torch.no_grad(): + native = step( q[t], k[t], - v[t], + original_v[t], gate, beta[t], - carry.anchor.values, + carry, + window, + enabled, scale, - use_qk_l2norm_in_kernel, - precision=config.precision, - gate_inputs=( - (gate_inputs[0][t], gate_inputs[1][t], *gate_inputs[2:]) - if gate_inputs is not None - else None - ), - ) - else: - decay = gate.exp().unsqueeze(-1) - if gate.ndim == 1: - decay = decay.unsqueeze(-1) - current_key = k[t] - if native_replay and use_qk_l2norm_in_kernel: - current_key = ( - current_key / (current_key.square().sum(-1, keepdim=True) + 1e-6).sqrt() - ) - key = EncodedLinearAttentionTensor(current_key, None, "identity", None) - decayed = ( - _replay_state(carry, gate) - if gate.ndim == 2 - else carry.reconstruct(original_basis=False) * decay - ) - residual = v[t] - _sum_keys(key.values.unsqueeze(-1) * decayed) - update = EncodedLinearAttentionTensor( - beta[t].unsqueeze(-1) * residual, None, "identity", None + normalize, + (gate_inputs[0][t], gate_inputs[1][t], *gate_inputs[2:]) + if gate_inputs is not None + else None, ) - working = decayed + key.values.unsqueeze(-1) * update.values.unsqueeze(-2) - if config.replay_window > 1: - if native_replay and carry.cursor + 1 < config.replay_window: - key.values = forward_value(key.values, native[3]) - update.values = forward_value(update.values, native[4]) - gate = forward_value(gate, native[5]) - entries = (*entries, ReplayEntry(key, update, gate)) - # Window 1 checkpoints every token; replay checkpoints only at a full window. - # Between checkpoints, keep the anchor and save the native BF16 update values. - refresh = config.replay_window == 1 or len(entries) == config.replay_window - if refresh: - if native_replay: - # Native step already quantized the checkpoint; attach its values to the adjoint. - anchor = EncodedLinearAttentionTensor( - forward_value(working, native[1]), - native[2], - "int8" if state_qdq else "identity", - 32 if state_qdq else None, - ) - else: - anchor = _encode( - working, - state_qdq, - block_v, - state_format=state_format, - state_quantizer=state_quantizer, - ) + native_outputs.append(native[0]) + decay = gate.exp().unsqueeze(-1) + if gate.ndim == 1: + decay = decay.unsqueeze(-1) + key = k[t] + if normalize: + key = normalized(key) + decayed = _replay_state(carry, gate) if gate.ndim == 2 else _replay_state(carry) * decay + residual = v[t] - _sum_keys(key.unsqueeze(-1) * decayed) + update = beta[t].unsqueeze(-1) * residual + working = decayed + key.unsqueeze(-1) * update.unsqueeze(-2) + if len(carry.entries) + 1 == window: + # Native step already quantized the checkpoint; preserve its scales and adjoint. + anchor, scales = forward_value(working, native[1]), native[2] entries = () else: - anchor = carry.anchor - # A new container keeps tensors needed by earlier tokens' backward passes intact. - next_carry = LinearAttentionState( - anchor, - entries, - carry.position + 1, - True, - signature, - carry.value_basis, - config.precision, - ) + anchor, scales = carry.anchor, carry.scales + entries = ( + *carry.entries, + _ReplayEntry( + forward_value(key, native[3]), + forward_value(update, native[4]), + forward_value(gate, native[5]), + ), + ) + carry = _ReplayState(anchor, scales, entries) query = q[t] - if native_replay and use_qk_l2norm_in_kernel: - query = query / (query.square().sum(-1, keepdim=True) + 1e-6).sqrt() - # Read this token from the working state, before the next-state checkpoint QDQ. - outputs.append( - native_output if serving else _sum_keys(query.unsqueeze(-1) * working) * scale - ) - carry = next_carry - output = torch.stack(outputs) - if config.state_codec == "int8_hadamard32": - output = _hadamard32(output) - if native_replay: - output = forward_value(output, torch.stack(native_outputs)) - return output, carry + if normalize: + query = normalized(query) + # Read the working state before checkpoint rounding, including when the window fills. + outputs.append(_sum_keys(query.unsqueeze(-1) * working) * scale) + return forward_value(_hadamard32(torch.stack(outputs)), torch.stack(native_outputs)), carry diff --git a/modelopt/torch/quantization/linear_attention/gdn.py b/modelopt/torch/quantization/linear_attention/gdn.py index 09f89efe0ef..bcaf7951230 100644 --- a/modelopt/torch/quantization/linear_attention/gdn.py +++ b/modelopt/torch/quantization/linear_attention/gdn.py @@ -30,8 +30,6 @@ def gdn_state_qat( beta, *, policy: LinearAttentionConfig, - state_qdq=False, - state_format="fp8_e4m3", state_quantizer=None, scale=None, initial_state=None, @@ -56,7 +54,8 @@ def gdn_state_qat( Validate prepared scalar log gates and promote inputs to FP32 working values. The shared training forward runs the chunked prefix and recurrent suffix with native forward values and configured state QDQ. FLA supplies the prefix backward; - a Torch recurrence supplies the suffix adjoint. + a Torch recurrence supplies the suffix adjoint. ``state_quantizer`` owns state + QDQ; passing None or a disabled quantizer retains the unquantized native path. """ if policy.precision == "vllm_kimi_k3": raise ValueError("vllm_kimi_k3 is a KDA-only precision profile") @@ -83,8 +82,6 @@ def gdn_state_qat( g, beta, policy=policy, - state_qdq=state_qdq, - state_format=state_format, state_quantizer=state_quantizer, scale=scale, initial_state=initial_state, diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py index f9653decf12..f6bdb7565ae 100644 --- a/modelopt/torch/quantization/linear_attention/kda.py +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -32,8 +32,6 @@ def kda_state_qat( beta, *, policy, - state_qdq=False, - state_format="fp8_e4m3", state_quantizer=None, scale=None, initial_state=None, @@ -62,7 +60,8 @@ def kda_state_qat( Prepare per-key-channel gates and optional beta activation using FLA's formula. The shared training forward runs the chunked prefix and recurrent suffix with native forward values and configured state QDQ. FLA supplies the prefix backward; - a Torch recurrence supplies the suffix adjoint. + a Torch recurrence supplies the suffix adjoint. ``state_quantizer`` owns state + QDQ; passing None or a disabled quantizer retains the unquantized native path. """ if cp_context is not None or disable_recompute or return_intermediate_states: raise NotImplementedError( @@ -90,10 +89,11 @@ def kda_state_qat( rate = A_log.to(dtype).exp().reshape(g.shape[-2], 1) g = -rate * F.softplus(g) # Import the optional vLLM backend only for the native precision profile. - from ...kernels.quantization.linear_attention.serving._compat import fla_module + from ...kernels.quantization.linear_attention.serving.forward import _kernels with torch.no_grad(): - native_gate = fla_module("kda").fused_kda_gate( + _, native_kda = _kernels("replayssm" if policy.precision == "replayssm" else "vllm") + native_gate = native_kda.fused_kda_gate( raw_gate.flatten(-2).contiguous(), A_log, raw_gate.shape[-1], g_bias=dt_bias ) g = forward_value(g, native_gate) @@ -109,8 +109,6 @@ def kda_state_qat( g, beta, policy=policy, - state_qdq=state_qdq, - state_format=state_format, state_quantizer=state_quantizer, scale=scale, initial_state=initial_state, diff --git a/modelopt/torch/quantization/linear_attention/training.py b/modelopt/torch/quantization/linear_attention/training.py index 9d9ab9c1d3e..d40a9fae5c0 100644 --- a/modelopt/torch/quantization/linear_attention/training.py +++ b/modelopt/torch/quantization/linear_attention/training.py @@ -16,10 +16,13 @@ """Training forwards with a chunked prefill prefix and recurrent decode suffix.""" from contextlib import contextmanager +from functools import cache +from importlib.metadata import version from itertools import pairwise from numbers import Integral import torch +from packaging.version import Version from .decode import _recurrent_decode from .utils import _resolve_state_quantizer, _state_qdq, forward_value @@ -138,12 +141,26 @@ def linear_attention_training_phase( module._linear_attention_cu_seqlens = original_boundaries +@cache +def _validate_training_device(device): + if torch.cuda.get_device_capability(device)[0] >= 12 and Version(version("triton")) < Version( + "3.7" + ): + raise NotImplementedError( + "Linear-attention state QAT is unsupported on sm_120+ with Triton < 3.7 " + "due to a CUDA misaligned-address fault. Use a vLLM build with Triton >= 3.7." + ) + + def _prepare_prefill_inputs(q, k, v, g, beta, *, policy, chunk_size): """Validate the native policy and prepare FP32 values for the training adjoint.""" if policy.backend != "serving" or chunk_size != 64: raise ValueError("State training requires backend='serving' and chunk_size=64") if q.device.type != "cuda" or any(x.dtype != torch.bfloat16 for x in (q, k, v)): raise ValueError("Serving arithmetic requires CUDA BF16 Q/K/V inputs") + if policy.state_codec == "int8_hadamard32" and v.shape[-1] % 32: + raise ValueError("int8_hadamard32 requires Dv divisible by 32") + _validate_training_device(q.device) return tuple(x.float() for x in (q, k, v, g, beta)) @@ -155,8 +172,6 @@ def _prefill_decode_forward( beta, *, policy, - state_qdq, - state_format, state_quantizer, scale, initial_state, @@ -174,9 +189,8 @@ def _prefill_decode_forward( """Run both prefill and decode phases in one differentiable training forward. Each sequence's chunked prefix produces the state for its token or ReplaySSM - suffix. ``recurrent_decode`` wraps that dense state in LinearAttentionState - and preserves its gradient connection to prefill. Outputs are joined in token - order; the final runtime state is reconstructed to the caller's dense layout. + suffix. The suffix preserves its gradient connection to prefill and returns + a dense final state only when requested. Outputs are joined in token order. Args: prefill_lengths: Prefix token count per sequence. For 128 tokens, a value @@ -184,10 +198,8 @@ def _prefill_decode_forward( sequence_lengths: Valid token count per sequence, excluding right padding. Defaults to the dense row length or each packed storage segment length. """ - state_quantizer, state_qdq, state_format = _resolve_state_quantizer( - state_quantizer, state_qdq, state_format - ) - if state_qdq and prefill_lengths is None: + state_quantizer = _resolve_state_quantizer(state_quantizer, state_codec=policy.state_codec) + if state_quantizer is not None and prefill_lengths is None: raise ValueError( "State quantization requires explicit prefill lengths through " "linear_attention_training_phase" @@ -242,14 +254,19 @@ def _prefill_decode_forward( raise ValueError( "Serving arithmetic requires K <= 256, KDA V=K, and power-of-two K for normalization" ) - if policy.precision != "replayssm": + if policy.precision == "replayssm": + if any(prefix < length for prefix, length in zip(prefixes, valid_lengths)) and ( + keys < 32 or keys & (keys - 1) or (g.ndim == 3 and gate_inputs is None) + ): + raise ValueError("ReplaySSM requires power-of-two K >= 32 and raw GDN gate inputs") + else: # Validate the optional native backend before launching any prefix kernels. from ...kernels.quantization.linear_attention.serving.forward import validate_profile validate_profile( gate_inputs, use_qk_l2norm_in_kernel, - kimi_k3=policy.precision == "vllm_kimi_k3", + precision=policy.precision, ) if initial_state is None: states = q.new_zeros(len(sequences), heads, keys, values) @@ -269,7 +286,7 @@ def _prefill_decode_forward( # A continuation prefill consumes a stored cache just as native serving # does. A fresh zero-state prefix has no incoming cache to quantize. - if initial_state is not None and state_qdq: + if initial_state is not None and state_quantizer is not None: if policy.precision == "replayssm": from ...kernels.quantization.linear_attention.serving.replay import ( checkpoint, @@ -281,9 +298,7 @@ def _prefill_decode_forward( decoded = original_basis(decoded) state = forward_value(state, decoded) else: - state = _state_qdq( - state, policy.state_block_v, state_format, state_quantizer - ) + state = _state_qdq(state, state_quantizer, policy.state_block_v) prefix, state = serving_prefix( q=q[b, start:split], k=k[b, start:split], @@ -306,15 +321,13 @@ def _prefill_decode_forward( ), ) # Keep the prefix state attached so suffix losses backpropagate through prefill. - # recurrent_decode applies configured state QDQ at the handoff and suffix writes. - suffix, carry = _recurrent_decode( + # The suffix applies configured state QDQ at the handoff and recurrent writes. + suffix, final_state = _recurrent_decode( *(x[b, split:end] for x in (q, k, v, g, beta)), config=policy, - state_qdq=state_qdq, - state_format=state_format, state_quantizer=state_quantizer, initial_state=state, - position=prefixes[n], + output_final_state=output_final_state, gate_inputs=( ( gate_inputs[0][b, split:end], @@ -333,7 +346,7 @@ def _prefill_decode_forward( pieces += (v.new_zeros(storage_end - end, heads, values),) outputs.append(torch.cat(pieces)) if output_final_state: - finals.append(carry.reconstruct()) + finals.append(final_state) output = torch.stack(outputs) if boundaries is None else torch.cat(outputs).unsqueeze(0) final = torch.stack(finals) if output_final_state else None if final is not None and state_v_first: diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py index 133649a0e04..9f44d2a44d7 100644 --- a/modelopt/torch/quantization/linear_attention/utils.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -83,6 +83,12 @@ def validate_gdn_quantizer( if not isinstance(quantizer, TensorQuantizer): raise ValueError(f"{name} requires a single TensorQuantizer") + if quantizer.block_sizes is not None and quantizer.block_sizes.get("type") == "dynamic": + raise ValueError( + f"{name} does not support block_sizes['type']='dynamic'. For INT8 state QDQ, " + "use top-level type='dynamic' with block_sizes={-1: N}; " + "quantized block scales are not supported." + ) if not ( quantizer._dynamic and quantizer.num_bits in num_bits @@ -111,7 +117,7 @@ def validate_gdn_quantizer( def state_quantizer_config( - quantizer: TensorQuantizer, *, name="state_quantizer" + quantizer: TensorQuantizer, *, name="state_quantizer", state_codec="tile" ) -> tuple[str, int]: """Validate the state quantizer and derive its format and last-axis group size. @@ -120,6 +126,11 @@ def state_quantizer_config( validate_gdn_quantizer( quantizer, name=name, num_bits=tuple(_STATE_FORMATS), block_sizes=(16, 32, 64) ) + if state_codec == "int8_hadamard32": + if quantizer.num_bits != 8: + raise ValueError("int8_hadamard32 requires INT8 state quantization") + if quantizer.block_sizes is not None: + raise ValueError("TensorQuantizer block_sizes requires state_codec='tile'") if quantizer.block_sizes is not None: return _STATE_FORMATS[quantizer.num_bits], quantizer.block_sizes[-1] if quantizer.axis != (0, 1): @@ -127,74 +138,25 @@ def state_quantizer_config( return _STATE_FORMATS[quantizer.num_bits], 0 -def _make_state_quantizer(state_format): - """Build a standalone-call quantizer; converted modules supply their registered instance.""" - # QuantizeConfig imports this module before the quantizer classes are initialized. - from ..config import QuantizerAttributeConfig - from ..nn import TensorQuantizer - - if state_format not in _STATE_FORMATS.values(): - raise ValueError("State format must be fp8_e4m3 or int8") - return TensorQuantizer( - QuantizerAttributeConfig( - num_bits=(4, 3) if state_format == "fp8_e4m3" else 8, - type="dynamic", - axis=(0, 1), - narrow_range=True, - pass_through_bwd=True, - ) - ) +def _resolve_state_quantizer(state_quantizer, *, state_codec="tile"): + """Validate an enabled quantizer once; return None when state QDQ is inactive.""" + if state_quantizer is not None and state_quantizer.is_enabled: + state_quantizer_config(state_quantizer, state_codec=state_codec) + if state_quantizer._if_quant: + return state_quantizer + return None -def _resolve_state_quantizer(state_quantizer, state_qdq, state_format): - """Resolve legacy flags and registered quantizer settings at either training entry point.""" - if state_quantizer is None and state_qdq: - state_quantizer = _make_state_quantizer(state_format) - if state_quantizer is not None: - state_qdq = state_quantizer.is_enabled and state_quantizer._if_quant - if state_quantizer.is_enabled: - state_format, _ = state_quantizer_config(state_quantizer) - return state_quantizer, state_qdq, state_format - - -def _state_qdq( - state: torch.Tensor, - block_v: int = 64, - state_format: str = "fp8_e4m3", - state_quantizer: TensorQuantizer | None = None, -): - """Apply the same state encoding used by recurrent writes.""" - return _state_qdq_with_scales(state, block_v, state_format, state_quantizer)[0] - - -def _state_qdq_with_scales(state, block_v, state_format, state_quantizer=None): - """Shared continuation/decode QDQ; blockwise scales remain owned by TensorQuantizer.""" - if state_quantizer is not None and state_quantizer.block_sizes is not None: - return state_quantizer(state), None - 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") - return _tile_qdq(state, block_v, state_format, state_quantizer=state_quantizer) - - -def _tile_qdq(value, block_v, state_format, *, state_quantizer=None): - """Return TensorQuantizer tile QDQ with identity STE and detached scales.""" - quantizer = ( - state_quantizer if state_quantizer is not None else _make_state_quantizer(state_format) - ) - rounded, scales = [], [] - for part in value.split(block_v, dim=-1): +def _state_qdq(state: torch.Tensor, state_quantizer: TensorQuantizer | None, block_v: int = 64): + """Apply state QDQ; TensorQuantizer owns the format, scales, and backward rule.""" + if state_quantizer is None: + return state + if state_quantizer.block_sizes is not None: + return state_quantizer(state) + rounded = [] + for part in state.split(block_v, dim=-1): # Canonical [N, H, K*BV] shape preserves per-head tile scales for both state ranks. tensor = part.flatten(-2) inputs = tensor.reshape(1, -1, tensor.shape[-1]) - decoded = quantizer(inputs) - amax = quantizer._get_amax(inputs).float() - if quantizer.num_bits == 8: - scale = amax / quantizer.maxbound - else: - safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) - scale = torch.div(quantizer.maxbound, safe_amax).reciprocal() - rounded.append(decoded.reshape_as(part)) - scales.append(scale.reshape(tensor.shape[:-1])) - return torch.cat(rounded, dim=-1).to(value.dtype), torch.stack(scales, dim=-1) + rounded.append(state_quantizer(inputs).reshape_as(part)) + return torch.cat(rounded, dim=-1).to(state.dtype) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index c1430fd39fc..f928acc2ad6 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -234,6 +234,16 @@ def _check_weight_quantization_took_effect(model: nn.Module, config: QuantizeCon # Quantizers that only framework plugins add: (class-name patterns of the modules they belong on, # what those modules are, recipe units listing the supported models). _PLUGIN_QUANTIZERS = { + "gdn_state_quantizer": ( + ("*GatedDeltaNet", "*GatedDeltaNetAttention"), + "GDN layer", + "linear_attention_state_int8*", + ), + "kda_state_quantizer": ( + ("*KimiDeltaAttention", "*KimiK3DeltaAttention", "*KimiLinearDeltaAttention"), + "KDA layer", + "linear_attention_state_int8*", + ), "indexer_q_quantizer": (("*Indexer",), "sparse-attention indexer", "indexer_*_nvfp4"), "indexer_k_quantizer": (("*Indexer",), "sparse-attention indexer", "indexer_*_nvfp4"), "dispatch_quantizer": (("*Experts*", "*FusedMoE"), "MoE experts module", "moe_*_nvfp4"), @@ -244,9 +254,9 @@ def _check_weight_quantization_took_effect(model: nn.Module, config: QuantizeCon def _check_plugin_quantization_took_effect(model: nn.Module, config: QuantizeConfig) -> None: """Raise when a config enables a plugin-only quantizer that quantizes nothing. - The quantizers in ``_PLUGIN_QUANTIZERS`` only exist on modules (sparse-attention indexers, MoE - experts) that a framework plugin (vLLM, Megatron-Core) converted. When a plugin does not - recognize the installed framework's module class, or the patterns match none of the modules, + The quantizers in ``_PLUGIN_QUANTIZERS`` only exist on modules (linear-attention layers, + sparse-attention indexers, MoE experts) that a framework plugin (vLLM, Megatron-Core) converted. + When a plugin does not recognize the installed framework's module class, or the patterns match none of the modules, the tensors silently stay unquantized. Intent is read like in :func:`_check_weight_quantization_took_effect`, but only from patterns naming the quantizer, so catch-all patterns do not count. A model without such a module (another architecture) is diff --git a/modelopt/torch/quantization/plugins/gdn.py b/modelopt/torch/quantization/plugins/gdn.py index 91dfc33c4dd..9eea217fa50 100644 --- a/modelopt/torch/quantization/plugins/gdn.py +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -63,11 +63,6 @@ def validate_linear_attention(self): raise ValueError("vllm_kimi_k3 is a KDA-only precision profile") super().validate_linear_attention() - @property - def gdn_state_qdq_block_v(self) -> int: - """Execution tile width; also sets grouping for legacy tile quantizers.""" - 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]: diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py index 4e651acbfd6..7eea8b4bfbc 100644 --- a/modelopt/torch/quantization/plugins/linear_attention.py +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -114,20 +114,16 @@ def linear_attention_is_enabled(self): def validate_linear_attention(self): """Validate quantizer contracts shared by GDN and KDA.""" if self._linear_attn_state.is_enabled: - state_format, group_size = state_quantizer_config( + state_quantizer_config( self._linear_attn_state, name=self.linear_attention_quantizer_names[0], + state_codec=self.linear_attention_config.state_codec, ) if self.linear_attention_config.backend != "serving": raise ValueError( "GDN/KDA state QAT requires backend='serving'; use " "linear_attention_training_phase to select a recurrent training suffix" ) - if self.linear_attention_config.state_codec == "int8_hadamard32": - if state_format != "int8": - raise ValueError("int8_hadamard32 requires INT8 state quantization") - if group_size: - raise ValueError("TensorQuantizer block_sizes requires state_codec='tile'") def warn_linear_attention_profile(self): """Warn at conversion/restore on one TP/DP representative of each owning stage.""" diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 7e8b5df29e7..6aeab2c5179 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -22,7 +22,7 @@ import types from contextlib import contextmanager from copy import copy -from functools import cache, partial +from functools import cache from itertools import pairwise from typing import Any @@ -1174,17 +1174,17 @@ def _prepare_input_for_gated_delta_rule(self, *args, **kwargs): def _quantized_linear_attention_kernel(self): kernel = self.gated_delta_rule previous_gates = self._linear_attention_gate_inputs - if self._serving_arithmetic: - def run_kernel(*args, **kwargs): - kwargs["use_qk_l2norm_in_kernel"] = self.use_qk_l2norm - if self._linear_attention_gate_inputs is not None: - kwargs["gate_inputs"] = self._linear_attention_gate_inputs - return self._linear_attention_kernel(kernel, *args, **kwargs) + def run_kernel(*args, **kwargs): + kwargs["use_qk_l2norm_in_kernel"] = self.use_qk_l2norm + if self._linear_attention_gate_inputs is None: + raise NotImplementedError( + "Serving state QAT requires raw gates captured by Megatron's _compute_gates hook" + ) + kwargs["gate_inputs"] = self._linear_attention_gate_inputs + return self._linear_attention_kernel(kernel, *args, **kwargs) - self.gated_delta_rule = run_kernel - else: - self.gated_delta_rule = partial(self._linear_attention_kernel, kernel) + self.gated_delta_rule = run_kernel self._linear_attention_gate_inputs = None try: yield @@ -1196,20 +1196,13 @@ def forward(self, *args, **kwargs): self.validate_linear_attention_execution() recompute = getattr(self, "recompute_gdn", False) capture = self._serving_arithmetic and self.training and recompute - if capture and not hasattr(super(), "_forward_compute"): - raise NotImplementedError("Selective state QAT recompute requires Megatron's core hook") previous = self._linear_attention_recompute self._linear_attention_recompute = capture if capture: # The adapter checkpoints the same core, with microbatch-local metadata. self.recompute_gdn = False try: - if hasattr(super(), "_forward_compute") or hasattr( - super(), "forward_pre_attn_and_core_attn" - ): - return super().forward(*args, **kwargs) - with self._quantized_linear_attention_kernel(): - return super().forward(*args, **kwargs) + return super().forward(*args, **kwargs) finally: self._linear_attention_recompute = previous if capture: @@ -1218,8 +1211,7 @@ def forward(self, *args, **kwargs): def _forward_compute(self, hidden_states, *args, **kwargs): compute = super()._forward_compute if not self._serving_arithmetic: - with self._quantized_linear_attention_kernel(): - return compute(hidden_states, *args, **kwargs) + return compute(hidden_states, *args, **kwargs) lengths = self._linear_attention_prefill_lengths sequence_lengths = self._linear_attention_sequence_lengths if lengths is None: @@ -1275,27 +1267,26 @@ def run(hidden): return tensor_parallel.checkpoint(run, False, hidden_states) return run(hidden_states) - def forward_pre_attn_and_core_attn(self, *args, **kwargs): - with self._quantized_linear_attention_kernel(): - return super().forward_pre_attn_and_core_attn(*args, **kwargs) - def validate_linear_attention_execution(self): super().validate_linear_attention_execution() self._validate_linear_attention_config() active = self.linear_attention_is_enabled serving = self.linear_attention_config.backend == "serving" and active + if active and not serving: + self.validate_linear_attention() if self.config.context_parallel_size > 1 and active: raise NotImplementedError("GDN/KDA QAT does not support Megatron context parallelism.") if serving and self.training and self.config.recompute_granularity == "full": raise NotImplementedError( "State QAT supports selective GDN recompute, not full-layer recompute" ) - if ( - serving - and getattr(self, "recompute_gdn", False) - and not hasattr(super(), "_forward_compute") + if serving and ( + not hasattr(super(), "_compute_gates") or not hasattr(super(), "_forward_compute") ): - raise NotImplementedError("Selective state QAT recompute requires Megatron's core hook") + raise NotImplementedError( + "Serving state QAT requires Megatron-Core's _compute_gates and _forward_compute " + "hooks to preserve raw gates and packed sequence lengths. Upgrade Megatron-Core." + ) def _validate_linear_attention_config(self): if self._serving_arithmetic and self.config.deterministic_mode: diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index bd591e69a6c..3088b53d776 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -64,16 +64,23 @@ supported combinations. | `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | | `iq2_s` | IQ2_S W2A16 (2.56 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | | `q8_0` | Q8_0 W8A16 (8.5 bpw), eligible linears | none | none (no calibration) | -| `linear_attention_state_int8_dynamic` | GDN/KDA decode state INT8 + Hadamard; weights unchanged | none | none (dynamic scales; requires a prefix/decode phase context) | +| `linear_attention_state_int8_dynamic` | GDN/KDA decode state INT8 + Hadamard; requires a private ReplaySSM fork | none | none (dynamic scales; requires a prefix/decode phase context) | | `linear_attention_state_int8_block32_dynamic` | GDN/KDA state INT8, 32 value channels per key row; weights unchanged | none | none (dynamic scales; requires vLLM and a prefix/decode phase context) | | `linear_attention_state_int8_block32_dynamic_kimi_k3` | Kimi-K3/Kimi-Linear KDA state INT8, 32 value channels per key row | none | none (dynamic scales; explicit Triton Kimi profile) | +The public linear-attention profiles require `vllm==0.30.0` and `fla-core==0.5.1` +(the latter supplies the prefix backward). Install these in the training environment: + +```bash +pip install vllm==0.30.0 fla-core==0.5.1 +``` + The block32 linear-attention recipe selects `precision="vllm"` and working-state readout. It applies fake QDQ with floating-point state storage. For GDN, select `gdn_prefill_backend="triton"` and -`VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE=1` where packed decode is available. +`VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE=1` for packed decode. This profile covers pure, single-token, non-speculative decode. Mixed prefill/decode batches can select different native kernels, so arbitrary continuous batching is not numerically qualified by these tests. @@ -86,8 +93,25 @@ Kimi-Linear on vLLM 0.30, which share model kernels (`precision="vllm_kimi_k3"`) requires raw projections, Q/K normalization, unbounded softplus gates, and the Kimi-K3 package dependencies (`gate_lower_bound=None` in the serving HF config). It never falls back to prepared-only arithmetic. -KDA validates the selected cache paths at the actual head size once per device; -builds with inconsistent state readout are rejected before training. +The adapters use the fixed vLLM 0.30.0 value-first cache layout. Native-forward +comparisons and independent recurrence checks validate it in GPU tests. + +The `linear_attention_state_int8_dynamic` recipe uses an optional private +quantized-ReplaySSM fork; public vLLM releases do not provide those kernels. +Its imports and kernel contract are isolated from the public 0.30.0 profiles; +it requires the fork's `vllm.model_executor.layers.fla.ops` API. Without access +to that fork, use a block32 recipe above. ReplaySSM native tests +skip when the fork is absent, so the public CI images do not qualify this profile. + +The public profiles do not support vLLM 0.20 or arbitrary development builds. +The tested GDN/KDA profiles pass bit-exact native forward comparisons on vLLM 0.30.0. +The existing sm_120+/Triton <3.7 training guard also applies to custom ReplaySSM +environments; changing the public profile does not remove that hardware restriction. + +Megatron state QAT requires native `_compute_gates` and `_forward_compute` hooks +for raw gates and packed padding; Megatron-Core 0.19.2 lacks these hooks and is +rejected before an active QAT forward. Loading checkpoints and running the original +module with state quantizers disabled outside a phase remain supported. State QAT training requires explicit per-sequence prefix lengths. Megatron selective GDN recompute captures each microbatch's phase and packed metadata, allowing backward diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index c9ab51c77a0..1b41ab01fb8 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -32,11 +32,16 @@ pytest.importorskip("fla") # Megatron-Core GatedDeltaNet needs FLA for its baseline kernels. GatedDeltaNet = pytest.importorskip("megatron.core.ssm.gated_delta_net").GatedDeltaNet -pytest.importorskip("vllm") +vllm = pytest.importorskip("vllm") +if vllm.__version__ != "0.30.0": + pytest.skip("State QAT profiles target vLLM 0.30.0", allow_module_level=True) from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet SEED = 1234 +HAS_SERVING_HOOKS = all( + hasattr(GatedDeltaNet, name) for name in ("_compute_gates", "_forward_compute") +) def _make_model(tp_size): @@ -95,6 +100,12 @@ def _test_gdn_qat_helper(rank, size, checkpoint_path): baseline = forward(model, decode_aware=False) mtq.quantize(model, _gdn_config()) module = next(m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)) + module.linear_attention_config.backend = "fla" + try: + with pytest.raises(ValueError, match="requires backend='serving'"): + forward(model, decode_aware=False) + finally: + module.linear_attention_config.backend = "serving" calls = [] hook = module.gdn_state_quantizer.register_forward_hook(lambda *_: calls.append(True)) try: @@ -205,6 +216,8 @@ def _compile_gdn_qat_kernels(rank, size): forward = _gdn_forward(model) with torch.no_grad(): forward(model) + if not HAS_SERVING_HOOKS: + return # The compatibility test only needs the original Megatron kernel. cfg = _gdn_config() cfg["linear_attention"][0]["cfg"]["state_block_v"] = 32 mtq.quantize(model, cfg) @@ -216,7 +229,7 @@ def _compile_gdn_qat_kernels(rank, size): torch.cuda.synchronize() -@pytest.fixture +@pytest.fixture(scope="module") def compiled_gdn_workers(dist_workers_size_1): """Warm one small QAT model in the same worker, outside the test-call budget.""" dist_workers_size_1.run(_compile_gdn_qat_kernels) @@ -225,4 +238,40 @@ def compiled_gdn_workers(dist_workers_size_1): def test_gdn_qat_and_sharded_restore(compiled_gdn_workers, tmp_path): """Train through QDQ after a Megatron distributed-checkpoint round trip.""" + if not HAS_SERVING_HOOKS: + pytest.skip("Serving QAT requires Megatron-Core gate and packed-forward hooks") compiled_gdn_workers.run(_test_gdn_qat_helper, tmp_path) + + +def _test_missing_serving_hooks(rank, size): + 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(): + baseline = forward(model, decode_aware=False) + mtq.quantize(model, _gdn_config()) + layer = next(m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)) + with ( + linear_attention_training_phase(layer, [64, 64]), + layer._quantized_linear_attention_kernel(), + pytest.raises(NotImplementedError, match="raw gates captured"), + ): + layer.gated_delta_rule() + for name in ("_compute_gates", "_forward_compute"): + with pytest.MonkeyPatch.context() as patch: + for cls in GatedDeltaNet.__mro__: + if name in vars(cls): + patch.delattr(cls, name) + restored = _make_model(size) + mto.restore_from_modelopt_state(restored, mto.modelopt_state(model)) + with pytest.raises(NotImplementedError, match="Upgrade Megatron-Core"): + forward(restored) + mtq.disable_quantizer(model, "*") + with torch.no_grad(): + torch.testing.assert_close(forward(model, decode_aware=False), baseline, rtol=0, atol=0) + + +def test_gdn_missing_serving_hooks(compiled_gdn_workers): + compiled_gdn_workers.run(_test_missing_serving_hooks) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py index 8157fd2ccd9..779a7c036f3 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py @@ -28,7 +28,9 @@ from modelopt.torch.quantization.linear_attention import linear_attention_training_phase KimiDeltaAttention = pytest.importorskip("megatron.core.ssm.gated_delta_net.kda").KimiDeltaAttention -pytest.importorskip("vllm") +vllm = pytest.importorskip("vllm") +if vllm.__version__ != "0.30.0": + pytest.skip("State QAT profiles target vLLM 0.30.0", allow_module_level=True) def _layer(): diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py index 860b9d37383..dd094805e0c 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py @@ -18,7 +18,25 @@ import pytest import torch -from modelopt.torch.quantization.linear_attention import LinearAttentionConfig, recurrent_decode +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import LinearAttentionConfig, gdn_state_qat +from modelopt.torch.quantization.linear_attention.decode import _replay_decode +from modelopt.torch.quantization.nn import TensorQuantizer + +pytest.importorskip("vllm") + + +@pytest.fixture(scope="module") +def state_quantizer(): + return TensorQuantizer( + QuantizerAttributeConfig( + num_bits=8, + type="dynamic", + axis=(0, 1), + narrow_range=True, + pass_through_bwd=True, + ) + ) @pytest.fixture( @@ -26,11 +44,15 @@ params=[(False, 1), (False, 4), (True, 4)], ids=["gdn-token", "gdn-replay", "kda-replay"], ) -def compiled_replay(request): - native = pytest.importorskip( - "vllm.model_executor.layers.fla.ops.fused_recurrent_replayssm", - exc_type=ModuleNotFoundError, - ) +def compiled_replay(request, state_quantizer): + try: + from vllm.model_executor.layers.fla.ops import fused_recurrent_replayssm as native + except ModuleNotFoundError as error: + if error.name and "vllm.model_executor.layers.fla.ops.fused_recurrent_replayssm".startswith( + error.name + ): + pytest.skip("Requires the optional quantized-ReplaySSM serving fork") + raise channel, window = request.param torch.manual_seed(71) q, k = [ @@ -57,13 +79,20 @@ def compiled_replay(request): ) kwargs = { "config": config, - "state_qdq": True, - "state_format": "int8", + "state_quantizer": state_quantizer, "gate_inputs": None if channel else (raw, raw_beta, rate, bias), } def forward(): - return recurrent_decode(*args, initial_state=initial, **kwargs) + return _replay_decode( + *(x.float() for x in args), + initial, + window, + state_quantizer.is_enabled, + 64**-0.5, + False, + kwargs["gate_inputs"], + ) out, carry = forward() torch.autograd.grad(out.sum() + carry.reconstruct().sum(), (*args, initial)) @@ -71,6 +100,37 @@ def forward(): return native, args, initial, kwargs, forward +@pytest.mark.parametrize( + ("keys", "values", "overrides", "message"), + [ + (32, 48, {}, "Dv divisible by 32"), + (32, 32, {"num_bits": (4, 3)}, "requires INT8"), + (32, 32, {"block_sizes": {-1: 32}, "axis": None}, "requires state_codec='tile'"), + (16, 32, {}, "power-of-two K >= 32"), + (32, 32, {}, "raw GDN gate inputs"), + ], +) +def test_replay_rejects_invalid_inputs_before_backend(keys, values, overrides, message): + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=8, type="dynamic", axis=(0, 1), narrow_range=True) + ) + quantizer.set_from_attribute_config(overrides) + q = torch.ones(1, 2, 1, keys, device="cuda", dtype=torch.bfloat16) + v = q.new_ones(1, 2, 1, values) + with pytest.raises(ValueError, match=message): + gdn_state_qat( + q, + q, + v, + q[..., 0], + q[..., 0], + policy=LinearAttentionConfig(backend="serving", precision="replayssm"), + state_quantizer=quantizer, + prefill_lengths=[1], + initial_state=q.new_ones(1, 1, keys, values), + ) + + def test_replay_matches_persistent_serving_cache(compiled_replay): native, args, initial, kwargs, forward = compiled_replay output, carry = forward() @@ -127,9 +187,9 @@ def test_replay_matches_persistent_serving_cache(compiled_replay): decoded = state[1].transpose(-1, -2).float().unflatten(-1, (-1, 32)) metadata = scales[1].transpose(-1, -2) torch.testing.assert_close( - carry.anchor.values, (decoded * metadata[..., None]).flatten(-2), rtol=0, atol=0 + carry.anchor, (decoded * metadata[..., None]).flatten(-2), rtol=0, atol=0 ) - torch.testing.assert_close(carry.anchor.scales, metadata, rtol=0, atol=0) + torch.testing.assert_close(carry.scales, metadata, rtol=0, atol=0) grads = torch.autograd.grad( output.square().mean() + carry.reconstruct().square().mean(), (*args, initial) ) diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py index 06337bb13f3..83d55033c25 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py @@ -13,13 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from inspect import signature - import pytest import torch -import triton from _test_utils.torch.quantization.linear_attention_reference import recurrent_delta_rule_reference -from packaging.version import Version from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.linear_attention import ( @@ -31,29 +27,25 @@ vllm = pytest.importorskip("vllm") -from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator +if vllm.__version__ != "0.30.0": + pytest.skip("State QAT profiles target vLLM 0.30.0", allow_module_level=True) -from modelopt.torch.kernels.quantization.linear_attention.serving._compat import ( - _kda_state_v_first, - fla_module, - gdn_module, - kda_module, - optional_module, -) +from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator +from vllm.third_party.flash_linear_attention import ops as gdn_ops +from vllm.third_party.flash_linear_attention.ops import kda as kda_ops @torch.no_grad() def _native_forward(args, quantizer, kda, kimi_k3): - """Run the CI vLLM 0.20/0.30 Triton APIs with a persistent quantized cache.""" + """Run the vLLM 0.30 Triton APIs with a persistent quantized cache.""" q, k, v, raw_g, raw_beta, rate, bias = [x.detach().clone() for x in args] _, shape = MambaStateShapeCalculator.gated_delta_net_state_shape( 1, 1, 1, q.shape[-1], v.shape[-1], 4 ) - # CI KDA kernels specify [H,V,K]; rectangular GDN advertises its own layout. - value_first = kda or shape == (1, v.shape[-1], q.shape[-1]) + assert shape == (1, v.shape[-1], q.shape[-1]) def layout(state): - return (state.transpose(-1, -2) if value_first else state).contiguous() + return state.transpose(-1, -2).contiguous() cache = torch.zeros(2, *shape, device=q.device) prefix_args = [x[:, :65].contiguous() for x in (q, k, v)] @@ -64,24 +56,27 @@ def layout(state): "cu_seqlens": torch.tensor([0, 65], device=q.device, dtype=torch.int32), } if kda: - native = kda_module(kimi_k3) - g = fla_module("kda").fused_kda_gate(raw_g.flatten(-2), rate, q.shape[-1], g_bias=bias) + native = kda_ops + if kimi_k3: + from vllm.models.kimi_k3.nvidia.ops.third_party import kda as native + g = kda_ops.fused_kda_gate(raw_g.flatten(-2), rate, q.shape[-1], g_bias=bias) beta = raw_beta.float().sigmoid() - raw_chunk = getattr(native, "chunk_kda_with_fused_gate", None) - if raw_chunk is not None: - beta_name = "raw_beta" if "raw_beta" in signature(raw_chunk).parameters else "beta" - prefix, state = raw_chunk( - *prefix_args, - raw_g=raw_g[:, :65].contiguous(), - **{beta_name: (raw_beta if beta_name == "raw_beta" else beta)[:, :65].contiguous()}, - A_log=rate, - g_bias=bias, - **kwargs, - ) - else: # vLLM 0.20 standalone KDA takes activated gates. - prefix, state = native.chunk_kda(*prefix_args, g[:, :65], beta[:, :65], **kwargs) + # vLLM 0.30's two explicit KDA profiles have different beta contracts. + beta_arg = ( + {"raw_beta": raw_beta[:, :65].contiguous()} + if kimi_k3 + else {"beta": beta[:, :65].contiguous()} + ) + prefix, state = native.chunk_kda_with_fused_gate( + *prefix_args, + raw_g=raw_g[:, :65].contiguous(), + A_log=rate, + g_bias=bias, + **beta_arg, + **kwargs, + ) else: - prepared = gdn_module("fused_gdn_prefill_post_conv").fused_post_conv_prep( + prepared = gdn_ops.fused_post_conv_prep( torch.cat([x[0].flatten(1) for x in prefix_args], dim=-1), raw_g[0, :65], raw_beta[0, :65], @@ -92,14 +87,12 @@ def layout(state): v.shape[-1], ) kwargs["use_qk_l2norm_in_kernel"] = False - prefix, state = fla_module("chunk").chunk_gated_delta_rule( - *[x[None] for x in prepared], **kwargs - ) + prefix, state = gdn_ops.chunk_gated_delta_rule(*[x[None] for x in prepared], **kwargs) cache[1] = layout(quantizer(layout(state)[0])) outputs = [prefix] indices = torch.ones(1, device=q.device, dtype=torch.int32) cu = torch.tensor([0, 1], device=q.device, dtype=torch.int32) - for token in range(65, 73): + for token in range(65, 67): operands = [x[:, token : token + 1].contiguous() for x in (q, k, v)] if kimi_k3: output, _ = native.fused_recurrent_kda_packed_decode( @@ -124,9 +117,7 @@ def layout(state): use_qk_l2norm_in_kernel=True, ) else: - output, _ = fla_module( - "fused_recurrent" - ).fused_recurrent_gated_delta_rule_packed_decode( + output, _ = gdn_ops.fused_recurrent_gated_delta_rule_packed_decode( mixed_qkv=torch.cat([x.flatten() for x in operands])[None], a=raw_g[0, token : token + 1], b=raw_beta[0, token : token + 1], @@ -147,24 +138,19 @@ def layout(state): @pytest.fixture(scope="module", params=["gdn", "kda", "kimi_k3"]) def compiled_serving_case(request): """Compile one shared BF16 shape per model, outside the test-call timer.""" - # Skip before any kernels: an FLA backward fault poisons later tests' CUDA context. - if torch.cuda.get_device_capability()[0] >= 12 and Version(triton.__version__) < Version("3.7"): - pytest.skip("FLA backward causes misaligned CUDA access on sm_120+ with Triton < 3.7") kda = request.param != "gdn" kimi_k3 = request.param == "kimi_k3" - if kimi_k3 and optional_module("vllm.models.kimi_k3.nvidia.ops.third_party.kda") is None: - pytest.skip("This vLLM build does not provide the Kimi-K3 backend") torch.manual_seed(73) - keys = 128 if kimi_k3 else 32 + keys = 32 # Share the smallest supported KDA shape across both native profiles. args = [ - torch.randn(1, 73, 1, dim, device="cuda", dtype=torch.bfloat16) + torch.randn(1, 67, 1, dim, device="cuda", dtype=torch.bfloat16) # Rectangular GDN state catches transposes; equal padded tile widths also # satisfy the fused post-conv kernel's tile-width constraint. for dim in (keys, keys, keys if kda else 24) ] args += [ - torch.randn((1, 73, 1, keys) if kda else (1, 73, 1), device="cuda", dtype=torch.bfloat16), - torch.randn(1, 73, 1, device="cuda", dtype=torch.bfloat16), + torch.randn((1, 67, 1, keys) if kda else (1, 67, 1), device="cuda", dtype=torch.bfloat16), + torch.randn(1, 67, 1, device="cuda", dtype=torch.bfloat16), torch.full((1,), -3.0, device="cuda"), torch.zeros(keys if kda else 1, device="cuda"), ] @@ -205,14 +191,7 @@ def forward(inputs=args, **overrides): gate = -rate.exp() * torch.nn.functional.softplus(raw_g.float() + bias) return gdn_state_qat(q, k, v, gate, raw_beta.float().sigmoid(), **kwargs) - try: - output, state = forward() - except NotImplementedError as error: - if vllm.__version__ == "0.20.0" and "KDA chunk/decode cache readout disagrees" in str( - error - ): - pytest.xfail(str(error)) - raise + output, state = forward() torch.autograd.grad(output.float().sum() + state.sum(), args) expected = _native_forward(args, quantizer, kda, kimi_k3) torch.cuda.synchronize() @@ -227,7 +206,7 @@ def test_serving_state_qdq_and_handoff_gradient(compiled_serving_case): output, state = forward() finally: handle.remove() - assert len(calls) == 9 # Handoff plus eight suffix writes. + assert len(calls) == 3 # Handoff plus two recurrent writes. torch.testing.assert_close(output, expected, rtol=0, atol=0) torch.testing.assert_close(state, expected_state, rtol=0, atol=0) gradients = torch.autograd.grad( @@ -245,7 +224,7 @@ def test_serving_state_qdq_and_handoff_gradient(compiled_serving_case): assert not torch.equal(output[:, 65:], plain[:, 65:]) # A nonzero nonsymmetric state checks layout against an independent oracle. initial = torch.randn(1, 1, args[0].shape[-1], args[2].shape[-1], device="cuda") - for prefix in (73, 65): + for prefix in (67, 65): actual, final = forward(prefill_lengths=[prefix], initial_state=initial) q, k, v, raw_g, raw_beta, rate, bias = [x.float() for x in args] q, k = [x / (x.square().sum(-1, keepdim=True) + 1e-6).sqrt() for x in (q, k)] @@ -263,60 +242,33 @@ def test_serving_state_qdq_and_handoff_gradient(compiled_serving_case): quantizer.enable() +@pytest.mark.parametrize("compiled_serving_case", ["kimi_k3"], indirect=True, scope="module") @pytest.mark.parametrize( - ("overrides", "message", "kimi_only"), + ("overrides", "message"), [ ( {"gate_inputs": None, "use_gate_in_kernel": False, "use_beta_sigmoid_in_kernel": False}, "raw gate_inputs", - True, ), - ({"use_qk_l2norm_in_kernel": False}, "normalization", True), - ({"allow_neg_eigval": True}, "allow_neg_eigval", False), + ({"use_qk_l2norm_in_kernel": False}, "normalization"), ], ) -def test_rejects_silent_serving_mismatches(compiled_serving_case, overrides, message, kimi_only): - args, _, forward, _, kimi_k3 = compiled_serving_case - if args[3].ndim != 4 or (kimi_only and not kimi_k3): - pytest.skip("Guard applies only to the selected KDA profile") +def test_rejects_silent_serving_mismatches(compiled_serving_case, overrides, message): + _, _, forward, _, _ = compiled_serving_case with pytest.raises(ValueError, match=message): forward(**overrides) -def test_kda_non_power_of_two_layout(monkeypatch): - # Probe values must be BF16-exact for K=192, including fused-gate chunk dispatch. - try: - for raw_gates in (False, True): - assert isinstance( - _kda_state_v_first(torch.device("cuda:0"), False, 192, raw_gates), bool - ) - except NotImplementedError as error: - if vllm.__version__ == "0.20.0" and "KDA chunk/decode cache readout disagrees" in str( - error - ): - pytest.xfail(str(error)) - raise - native = kda_module() - recurrent = native.fused_recurrent_kda_fwd - - def shifted_readout(*args, **kwargs): - output, state = recurrent(*args, **kwargs) - return output.roll(1, -1), state - - _kda_state_v_first.cache_clear() - try: - with monkeypatch.context() as patch: - patch.setattr(native, "fused_recurrent_kda_fwd", shifted_readout) - with pytest.raises(NotImplementedError, match="cache layout could not be validated"): - _kda_state_v_first(torch.device("cuda:0"), False, 192) - finally: - _kda_state_v_first.cache_clear() - - -@pytest.mark.parametrize("packed", [False, True], ids=["dense", "packed"]) +@pytest.mark.parametrize( + ("compiled_serving_case", "packed"), + [("gdn", False), ("kda", True)], + indirect=["compiled_serving_case"], + scope="module", + ids=["gdn-dense", "kda-packed"], +) def test_padding_does_not_extend_decode(compiled_serving_case, packed): args, quantizer, forward, _, _ = compiled_serving_case - valid = 69 # The existing compiled shape uses a 65-token prefix. + valid = 66 # One recurrent write and one padding token after the 65-token prefix. trimmed = [ (x[:, :valid] if i < 5 else x).detach().clone().requires_grad_() for i, x in enumerate(args) ] @@ -343,14 +295,14 @@ def test_padding_does_not_extend_decode(compiled_serving_case, packed): initial_state=initial, prefill_lengths=[65, 0], sequence_lengths=[valid, 0], - cu_seqlens=torch.tensor([0, 73, 146], device="cuda", dtype=torch.int32) + cu_seqlens=torch.tensor([0, 67, 134], device="cuda", dtype=torch.int32) if packed else None, ) finally: hook.remove() - assert len(calls) == 6 # Incoming prefix state, handoff, four writes; none for the empty row. - output = output.reshape(2, 73, *output.shape[-2:]) + assert len(calls) == 3 # Incoming prefix state, handoff, one write; none for the empty row. + output = output.reshape(2, 67, *output.shape[-2:]) torch.testing.assert_close(output[:1, :valid], expected, rtol=0, atol=0) torch.testing.assert_close(state[:1], expected_state, rtol=0, atol=0) assert not output[0, valid:].count_nonzero() and not output[1].count_nonzero() @@ -358,7 +310,7 @@ def test_padding_does_not_extend_decode(compiled_serving_case, packed): grads = torch.autograd.grad(output.square().sum() + state.square().sum(), padded) for i, (grad, expected_grad) in enumerate(zip(grads, expected_grads)): if i < 5: - grad = grad.reshape(2, 73, *grad.shape[2:]) + grad = grad.reshape(2, 67, *grad.shape[2:]) assert not grad[0, valid:].count_nonzero() and not grad[1].count_nonzero() grad = grad[:1, :valid] torch.testing.assert_close(grad, expected_grad) diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index e243589bedb..aab9a3d284d 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -20,6 +20,7 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe from modelopt.torch.quantization.linear_attention import ( LinearAttentionConfig, get_linear_attention_layers, @@ -181,12 +182,22 @@ def test_state_qat_requires_serving_policy(mixin): model.modelopt_post_restore() -@pytest.mark.parametrize("overrides", [{"pass_through_bwd": False}, {"type": "static"}]) -def test_unsupported_quantizer_fails_at_conversion(overrides): +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"pass_through_bwd": False}, "supports only"), + ({"type": "static"}, "supports only"), + ( + {"num_bits": 8, "axis": None, "block_sizes": {-1: 32, "type": "dynamic"}}, + "use top-level type='dynamic'", + ), + ], +) +def test_unsupported_quantizer_fails_at_conversion(overrides, message): cfg = quant_cfg() cfg["algorithm"] = None cfg["quant_cfg"][-1]["cfg"] = {**GDN_STATE_FP8_DYNAMIC, **overrides} - with pytest.raises(ValueError, match="supports only"): + with pytest.raises(ValueError, match=message): mtq.quantize(TinyGatedDeltaNet(), cfg) @@ -207,14 +218,18 @@ def test_invalid_phase_lengths(kwargs, message): pass -def test_restore_legacy_gdn_without_new_quantizer_handles(): +@pytest.mark.parametrize("has_policy", [False, True]) +def test_restore_legacy_gdn_without_new_quantizer_handles(has_policy): 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) + if not has_policy: + metadata.pop("linear_attention", None) + else: + assert "linear_attention" in metadata for name in ("gdn_state_quantizer", "gdn_w_quantizer"): metadata["quantizer_state"].pop(name, None) restored = TinyGatedDeltaNet() @@ -276,10 +291,29 @@ def test_quant_cfg_refinement_and_policy_precedence(): {"module_name": "", "cfg": {"backend": "serving", "state_block_v": 32}}, ] assert mtq.quantize(model, cfg) is model - assert model.gdn_state_qdq_block_v == 32 + assert model.linear_attention_config.state_block_v == 32 cfg["linear_attention"].reverse() mtq.quantize(model, cfg) - assert model.gdn_state_qdq_block_v == 128 + assert model.linear_attention_config.state_block_v == 128 cfg["quant_cfg"][-1]["cfg"] = {**GDN_STATE_FP8_DYNAMIC, "pass_through_bwd": False} with pytest.raises(ValueError, match="supports only"): mtq.quantize(model, cfg) + + +@pytest.mark.parametrize( + ("class_name", "quantizer_name"), + [ + ("GatedDeltaNet", "gdn_state_quantizer"), + ("Qwen3NextGatedDeltaNetAttention", "gdn_state_quantizer"), + ("KimiDeltaAttention", "kda_state_quantizer"), + ("KimiK3DeltaAttention", "kda_state_quantizer"), + ("KimiLinearDeltaAttention", "kda_state_quantizer"), + ], +) +def test_state_recipe_rejects_unconverted_layers(class_name, quantizer_name): + # An unknown framework class has no registered state-quantization adapter. + model = type(class_name, (nn.Module,), {})() + model.proj = nn.Linear(4, 4) + config = load_recipe("general/ptq/linear_attention_state_int8_block32_dynamic").quantize + with pytest.raises(RuntimeError, match=quantizer_name): + mtq.quantize(model, config) diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py index aefe454398e..8cd5a994821 100644 --- a/tests/unit/torch/quantization/test_linear_attention_decode.py +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -17,7 +17,8 @@ import torch from modelopt.torch.quantization.config import QuantizeConfig, QuantizerAttributeConfig -from modelopt.torch.quantization.linear_attention.decode import _encode +from modelopt.torch.quantization.linear_attention import LinearAttentionConfig, kda_state_qat +from modelopt.torch.quantization.linear_attention import training as state_training from modelopt.torch.quantization.linear_attention.utils import _state_qdq from modelopt.torch.quantization.nn import TensorQuantizer @@ -34,16 +35,11 @@ def test_state_qdq_matches_tensor_quantizer(state_format): expected = torch.cat( [quantizer(tile.flatten(-2)).reshape_as(tile) for tile in value.split(16, -1)], -1 ) - encoded = _encode(value, True, 16, state_format=state_format, state_quantizer=quantizer) - torch.testing.assert_close(encoded.values, expected, rtol=0, atol=0) - torch.testing.assert_close( - _state_qdq(value, 16, state_format, quantizer), expected, rtol=0, atol=0 - ) + actual = _state_qdq(value, quantizer, 16) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) probe = torch.randn_like(value) - (gradient,) = torch.autograd.grad((encoded.values * probe).sum(), value) + (gradient,) = torch.autograd.grad((actual * probe).sum(), value) torch.testing.assert_close(gradient, probe, rtol=0, atol=0) - assert encoded.scales.shape == (2, 3, 2) - assert not encoded.scales.requires_grad @pytest.mark.parametrize( @@ -63,3 +59,37 @@ def test_linear_attention_policy_rejects_invalid_settings(settings): QuantizeConfig( linear_attention=[{"module_name": "*", "cfg": {"backend": "serving", **settings}}] ) + + +@pytest.mark.parametrize("raw_gates", [False, True]) +def test_kda_rejects_incompatible_negative_eigenvalue_gates(raw_gates): + q = torch.zeros(1, 1, 1, 32) + beta = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="allow_neg_eigval"): + kda_state_qat( + q, + q, + q, + q, + beta, + policy=LinearAttentionConfig(backend="serving"), + allow_neg_eigval=True, + use_beta_sigmoid_in_kernel=raw_gates, + gate_inputs=(q, beta, torch.zeros(1), torch.zeros(32)) if raw_gates else None, + ) + + +@pytest.mark.parametrize(("major", "triton_version"), [(12, "3.6.0"), (12, "3.7.0"), (8, "3.6.0")]) +def test_state_training_device_guard(monkeypatch, major, triton_version): + # CI cannot safely run the faulting kernels; simulate only the absent hardware/version. + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (major, 0)) + monkeypatch.setattr(state_training, "version", lambda name: triton_version) + state_training._validate_training_device.cache_clear() + try: + if major >= 12 and triton_version == "3.6.0": + with pytest.raises(NotImplementedError, match=r"sm_120.*Triton < 3\.7"): + state_training._validate_training_device(torch.device("cuda:0")) + else: + state_training._validate_training_device(torch.device("cuda:0")) + finally: + state_training._validate_training_device.cache_clear() From 3bbaf8634b88a6ecb09a343014b3dc98d98cc1f4 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 12 Oct 2026 10:41:40 -0700 Subject: [PATCH 23/27] Fix linear attention serving contracts and GPU tests Signed-off-by: Kai Xu --- .../linear_attention/serving/forward.py | 56 +++++++------ .../quantization/linear_attention/gdn.py | 4 + .../torch/quantization/plugins/megatron.py | 33 ++++++-- modelopt_recipes/ptq.md | 5 +- .../plugins/test_megatron_gated_delta_net.py | 49 ++++++++---- .../test_linear_attention_replay.py | 10 ++- .../test_linear_attention_training.py | 79 ++++++++++++++++--- 7 files changed, 174 insertions(+), 62 deletions(-) diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py index 84913b224a7..6088458c391 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py @@ -27,17 +27,16 @@ def _kernels(precision): from .replay import _prefill_kernels return _prefill_kernels() - else: - import vllm + import vllm - if vllm.__version__ != "0.30.0": - raise ImportError("Linear-attention vLLM profiles require vllm==0.30.0") - from vllm.third_party.flash_linear_attention import ops as gdn_ops + if vllm.__version__ != "0.30.0": + raise ImportError("Linear-attention vLLM profiles require vllm==0.30.0") + from vllm.third_party.flash_linear_attention import ops as gdn_ops - if precision == "vllm_kimi_k3": - from vllm.models.kimi_k3.nvidia.ops.third_party import kda as kda_ops - else: - from vllm.third_party.flash_linear_attention.ops import kda as kda_ops + if precision == "vllm_kimi_k3": + from vllm.models.kimi_k3.nvidia.ops.third_party import kda as kda_ops + else: + from vllm.third_party.flash_linear_attention.ops import kda as kda_ops return gdn_ops, kda_ops @@ -48,6 +47,7 @@ def _state_layout(state): def prefill(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, precision="vllm"): """Run native chunk prefill, including the selected profile's raw-gate preparation.""" + # Match BF16 projection outputs; A_log and bias retain their parameter precision. gdn, kda = _kernels(precision) kimi_k3 = precision == "vllm_kimi_k3" channel = g.ndim == 3 @@ -55,10 +55,10 @@ def prefill(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, * raw_g, raw_beta, rate, bias = gate_inputs q, k, v, g, beta = gdn.fused_post_conv_prep( torch.cat([x.flatten(1) for x in (q, k, v)], dim=-1), - raw_g.flatten(0, -2), - raw_beta.flatten(0, -2), - rate, - bias, + raw_g.flatten(0, -2).bfloat16().contiguous(), + raw_beta.flatten(0, -2).bfloat16().contiguous(), + rate.reshape(-1).contiguous(), + bias.contiguous(), q.shape[-2], q.shape[-1], v.shape[-1], @@ -78,14 +78,16 @@ def prefill(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, * if channel and gate_inputs is not None: raw_g, raw_beta, rate, bias = gate_inputs # Kimi fuses beta activation; standalone KDA takes already activated beta. - beta_arg = {"raw_beta": raw_beta[None].contiguous()} if kimi_k3 else {"beta": beta} + beta_arg = ( + {"raw_beta": raw_beta[None].bfloat16().contiguous()} if kimi_k3 else {"beta": beta} + ) out, final = kda.chunk_kda_with_fused_gate( q, k, v, - raw_g=raw_g[None].contiguous(), - A_log=rate, - g_bias=bias, + raw_g=raw_g[None].bfloat16().contiguous(), + A_log=rate.reshape(-1).contiguous(), + g_bias=bias.contiguous() if bias is not None else None, **beta_arg, **kwargs, ) @@ -102,16 +104,19 @@ def step(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, p kimi_k3 = precision == "vllm_kimi_k3" channel = g.ndim == 2 cu, indices = _decode_metadata(q.device) - native_state = _state_layout(state.unsqueeze(0)) if channel or gate_inputs is not None: - # Slot zero is reserved by vLLM's paged-cache kernels. - native_state = torch.cat((torch.zeros_like(native_state), native_state), dim=0) + # Only slot one is read/written; slot zero is reserved and needs no initialization. + native_state = state.new_empty(2, state.shape[0], state.shape[2], state.shape[1]) + native_state[1].copy_(state.transpose(-1, -2)) + else: + native_state = _state_layout(state.unsqueeze(0)) if kimi_k3: + # The Kimi kernel always normalizes Q/K; validate_profile enforces normalize=True. raw_g, raw_beta, rate, bias = gate_inputs out, final = kda.fused_recurrent_kda_packed_decode( mixed_qkv=torch.cat([x.flatten() for x in (q, k, v)])[None], - raw_g=raw_g[None, None].contiguous(), - raw_beta=raw_beta[None, None].contiguous(), + raw_g=raw_g[None, None].bfloat16().contiguous(), + raw_beta=raw_beta[None, None].bfloat16().contiguous(), A_log=rate.reshape(-1).contiguous(), dt_bias=(torch.zeros_like(raw_g).flatten() if bias is None else bias.contiguous()), lower_bound=None, @@ -133,13 +138,13 @@ def step(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, p raw_g, raw_beta, rate, bias = gate_inputs out, final = gdn.fused_recurrent_gated_delta_rule_packed_decode( mixed_qkv=torch.cat([x.flatten() for x in (q, k, v)])[None], - a=raw_g[None].contiguous(), - b=raw_beta[None].contiguous(), + a=raw_g[None].bfloat16().contiguous(), + b=raw_beta[None].bfloat16().contiguous(), A_log=rate.reshape(-1).contiguous(), dt_bias=bias.contiguous(), scale=scale, initial_state=native_state, - out=torch.empty_like(v)[None, None], + out=torch.empty_like(v, memory_format=torch.contiguous_format)[None, None], ssm_state_indices=indices, use_qk_l2norm_in_kernel=normalize, ) @@ -152,6 +157,7 @@ def step(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, p use_qk_l2norm_in_kernel=normalize, cu_seqlens=cu, ) + # Indexed kernels update slot 1; the non-in-place fallback returns one state. return out[0, 0], _state_layout(final[-1]) diff --git a/modelopt/torch/quantization/linear_attention/gdn.py b/modelopt/torch/quantization/linear_attention/gdn.py index bcaf7951230..6fed19cdf13 100644 --- a/modelopt/torch/quantization/linear_attention/gdn.py +++ b/modelopt/torch/quantization/linear_attention/gdn.py @@ -56,9 +56,13 @@ def gdn_state_qat( native forward values and configured state QDQ. FLA supplies the prefix backward; a Torch recurrence supplies the suffix adjoint. ``state_quantizer`` owns state QDQ; passing None or a disabled quantizer retains the unquantized native path. + ``precision="vllm"`` requires ``gate_inputs=(raw_g, raw_beta, A_log, bias)`` + for native gate preparation; the prepared ``g`` and ``beta`` supply the adjoint. """ if policy.precision == "vllm_kimi_k3": raise ValueError("vllm_kimi_k3 is a KDA-only precision profile") + if policy.precision == "vllm" and gate_inputs is None: + raise ValueError("vllm GDN requires raw gate_inputs=(raw_g, raw_beta, A_log, bias)") if cp_context is not None: raise NotImplementedError("GDN state QAT does not support context parallelism") _validate_gate_inputs(g, beta, gate_inputs) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 6aeab2c5179..a5abb7a3b7e 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -1209,7 +1209,18 @@ def forward(self, *args, **kwargs): self.recompute_gdn = recompute def _forward_compute(self, hidden_states, *args, **kwargs): - compute = super()._forward_compute + return self._run_linear_attention_compute( + super()._forward_compute, hidden_states, *args, **kwargs + ) + + def forward_pre_attn_and_core_attn(self, hidden_states, *args, **kwargs): + # Core's two-stage schedules can call this entry point without forward(). + self.validate_linear_attention_execution() + return self._run_linear_attention_compute( + super().forward_pre_attn_and_core_attn, hidden_states, *args, **kwargs + ) + + def _run_linear_attention_compute(self, compute, hidden_states, *args, **kwargs): if not self._serving_arithmetic: return compute(hidden_states, *args, **kwargs) lengths = self._linear_attention_prefill_lengths @@ -1217,11 +1228,13 @@ def _forward_compute(self, hidden_states, *args, **kwargs): if lengths is None: raise ValueError("State quantization requires explicit prefill lengths") arguments = inspect.signature(compute).bind(hidden_states, *args, **kwargs) - boundaries = self._linear_attention_cu_seqlens.resolve( - arguments.arguments.get("cu_seqlens_q") - ) packed = arguments.arguments.get("packed_seq_params") + cu_seqlens = arguments.arguments.get("cu_seqlens_q") if packed is not None and packed.qkv_format == "thd": + if cu_seqlens is None: + cu_seqlens = packed.cu_seqlens_q_padded + if cu_seqlens is None: + cu_seqlens = packed.cu_seqlens_q actual = packed.cu_seqlens_q if actual is None: if sequence_lengths is None: @@ -1237,6 +1250,7 @@ def _forward_compute(self, hidden_states, *args, **kwargs): length > end for length, end in zip(sequence_lengths, valid) ): raise ValueError("Sequence lengths exceed the packed valid token counts") + boundaries = self._linear_attention_cu_seqlens.resolve(cu_seqlens) if self._linear_attention_recompute: # Packed metadata may refer to buffers reused by a later pipeline microbatch. for name, value in list(arguments.arguments.items()): @@ -1281,11 +1295,16 @@ def validate_linear_attention_execution(self): "State QAT supports selective GDN recompute, not full-layer recompute" ) if serving and ( - not hasattr(super(), "_compute_gates") or not hasattr(super(), "_forward_compute") + not hasattr(super(), "_compute_gates") + or not ( + hasattr(super(), "_forward_compute") + or hasattr(super(), "forward_pre_attn_and_core_attn") + ) ): raise NotImplementedError( - "Serving state QAT requires Megatron-Core's _compute_gates and _forward_compute " - "hooks to preserve raw gates and packed sequence lengths. Upgrade Megatron-Core." + "Serving state QAT requires Megatron-Core's _compute_gates and either " + "_forward_compute or forward_pre_attn_and_core_attn to preserve raw gates and " + "packed sequence lengths. Upgrade Megatron-Core." ) def _validate_linear_attention_config(self): diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 3088b53d776..0fe785c28f7 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -108,8 +108,9 @@ The tested GDN/KDA profiles pass bit-exact native forward comparisons on vLLM 0. The existing sm_120+/Triton <3.7 training guard also applies to custom ReplaySSM environments; changing the public profile does not remove that hardware restriction. -Megatron state QAT requires native `_compute_gates` and `_forward_compute` hooks -for raw gates and packed padding; Megatron-Core 0.19.2 lacks these hooks and is +Megatron state QAT requires native `_compute_gates` and either +`forward_pre_attn_and_core_attn` or `_forward_compute` for raw gates and packed +padding; Megatron-Core 0.19.2 lacks these hooks and is rejected before an active QAT forward. Loading checkpoints and running the original module with state quantizers disabled outside a phase remain supported. diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index 1b41ab01fb8..0ecfadcb8d8 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -39,8 +39,8 @@ from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet SEED = 1234 -HAS_SERVING_HOOKS = all( - hasattr(GatedDeltaNet, name) for name in ("_compute_gates", "_forward_compute") +HAS_SERVING_HOOKS = hasattr(GatedDeltaNet, "_compute_gates") and any( + hasattr(GatedDeltaNet, name) for name in ("_forward_compute", "forward_pre_attn_and_core_attn") ) @@ -128,7 +128,7 @@ def _test_gdn_qat_helper(rank, size, checkpoint_path): before = restored_gdn.in_proj.weight.detach().clone() torch.optim.SGD(restored.parameters(), lr=1e-3).step() assert not torch.equal(restored_gdn.in_proj.weight, before) - _check_packed_recompute(restored_gdn) + _check_packed_forward(restored_gdn) for name, value, restore in ( ("context_parallel_size", 2, True), @@ -161,10 +161,7 @@ def _test_gdn_qat_helper(rank, size, checkpoint_path): setattr(layer.config, name, original) -def _check_packed_recompute(module): - # Older Megatron releases lack selective GDN recompute and its native core hook. - if not hasattr(module, "recompute_gdn") or not hasattr(GatedDeltaNet, "_forward_compute"): - return +def _check_packed_forward(module): hidden = torch.randn(128, 1, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) packed = PackedSeqParams( qkv_format="thd", @@ -176,13 +173,19 @@ def _check_packed_recompute(module): max_seqlen_kv=64, total_tokens=128, ) - original = module.recompute_gdn + original = getattr(module, "recompute_gdn", None) was_training = module.training module.train() try: results = [] - for recompute, lengths in ((False, [60, 58]), (False, None), (True, None)): - module.recompute_gdn = recompute + cases = [(False, [60, 58], False), (False, None, False)] + if original is not None and hasattr(GatedDeltaNet, "_forward_compute"): + cases.append((True, None, False)) + if hasattr(GatedDeltaNet, "forward_pre_attn_and_core_attn"): + cases.append((False, None, True)) + for recompute, lengths, two_stage in cases: + if original is not None: + module.recompute_gdn = recompute module.zero_grad(set_to_none=True) hidden.grad = None with ( @@ -194,9 +197,19 @@ def _check_packed_recompute(module): ), torch.autocast("cuda", dtype=torch.bfloat16), ): - output = module(hidden, attention_mask=None, packed_seq_params=packed)[0] + if two_stage: + norm = module.forward_pre_attn_and_core_attn( + hidden, attention_mask=None, packed_seq_params=packed + ) + output = module.forward_post_core_attn(norm)[0] + else: + output = module(hidden, attention_mask=None, packed_seq_params=packed)[0] + # Exclude padding from the loss: trained output norms may have a bias. # Backward runs after the phase exits, as in pipeline schedules. - output.float().square().sum().backward() + valid_output = torch.cat((output[:60], output[64:122])) + valid_output.float().square().sum().backward() + assert not hidden.grad[60:64].count_nonzero() + assert not hidden.grad[122:].count_nonzero() results.append( (output.detach(), hidden.grad.clone(), module.in_proj.weight.grad.clone()) ) @@ -204,7 +217,8 @@ def _check_packed_recompute(module): for value, expected in zip(actual, results[0]): torch.testing.assert_close(value, expected, rtol=0, atol=0) finally: - module.recompute_gdn = original + if original is not None: + module.recompute_gdn = original module.train(was_training) @@ -225,7 +239,7 @@ def _compile_gdn_qat_kernels(rank, size): forward(model) model.train() forward(model).sum().backward() - _check_packed_recompute(next(m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet))) + _check_packed_forward(next(m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet))) torch.cuda.synchronize() @@ -259,11 +273,12 @@ def _test_missing_serving_hooks(rank, size): pytest.raises(NotImplementedError, match="raw gates captured"), ): layer.gated_delta_rule() - for name in ("_compute_gates", "_forward_compute"): + for names in (("_compute_gates",), ("_forward_compute", "forward_pre_attn_and_core_attn")): with pytest.MonkeyPatch.context() as patch: for cls in GatedDeltaNet.__mro__: - if name in vars(cls): - patch.delattr(cls, name) + for name in names: + if name in vars(cls): + patch.delattr(cls, name) restored = _make_model(size) mto.restore_from_modelopt_state(restored, mto.modelopt_state(model)) with pytest.raises(NotImplementedError, match="Upgrade Megatron-Core"): diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py index dd094805e0c..86c9aedfbe4 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py @@ -15,16 +15,24 @@ """Minimal native-cache and backward checks for the optional ReplaySSM serving fork.""" +from importlib import import_module + import pytest import torch from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.linear_attention import LinearAttentionConfig, gdn_state_qat from modelopt.torch.quantization.linear_attention.decode import _replay_decode +from modelopt.torch.quantization.linear_attention.training import _validate_training_device from modelopt.torch.quantization.nn import TensorQuantizer pytest.importorskip("vllm") +try: + _validate_training_device(torch.device("cuda")) +except NotImplementedError as error: + pytest.skip(str(error), allow_module_level=True) + @pytest.fixture(scope="module") def state_quantizer(): @@ -46,7 +54,7 @@ def state_quantizer(): ) def compiled_replay(request, state_quantizer): try: - from vllm.model_executor.layers.fla.ops import fused_recurrent_replayssm as native + native = import_module("vllm.model_executor.layers.fla.ops.fused_recurrent_replayssm") except ModuleNotFoundError as error: if error.name and "vllm.model_executor.layers.fla.ops.fused_recurrent_replayssm".startswith( error.name diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py index 83d55033c25..d29101b3dd3 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py @@ -141,13 +141,8 @@ def compiled_serving_case(request): kda = request.param != "gdn" kimi_k3 = request.param == "kimi_k3" torch.manual_seed(73) - keys = 32 # Share the smallest supported KDA shape across both native profiles. - args = [ - torch.randn(1, 67, 1, dim, device="cuda", dtype=torch.bfloat16) - # Rectangular GDN state catches transposes; equal padded tile widths also - # satisfy the fused post-conv kernel's tile-width constraint. - for dim in (keys, keys, keys if kda else 24) - ] + keys = 32 # Share one head shape across GDN/KDA to reuse native and FLA compilation. + args = [torch.randn(1, 67, 1, keys, device="cuda", dtype=torch.bfloat16) for _ in range(3)] args += [ torch.randn((1, 67, 1, keys) if kda else (1, 67, 1), device="cuda", dtype=torch.bfloat16), torch.randn(1, 67, 1, device="cuda", dtype=torch.bfloat16), @@ -242,16 +237,65 @@ def test_serving_state_qdq_and_handoff_gradient(compiled_serving_case): quantizer.enable() -@pytest.mark.parametrize("compiled_serving_case", ["kimi_k3"], indirect=True, scope="module") +@pytest.fixture(scope="module") +@torch.no_grad() +def compiled_gdn_layout_case(): + """Compile a two-head case so strided head/value axes are observable.""" + torch.manual_seed(732) + q, k, v = [torch.randn(1, 2, 2, 32, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + raw_g, raw_beta = [torch.randn(1, 2, 2, device="cuda", dtype=torch.bfloat16) for _ in range(2)] + gates = (raw_g, raw_beta, torch.full((2,), -3.0, device="cuda"), torch.randn(2, device="cuda")) + initial = torch.randn(1, 2, 32, 32, device="cuda") + + def forward(values, gate_inputs): + raw_g, raw_beta, rate, bias = gate_inputs + return gdn_state_qat( + q, + k, + values, + -rate.exp() * torch.nn.functional.softplus(raw_g.float() + bias), + raw_beta.float().sigmoid(), + policy=LinearAttentionConfig(backend="serving"), + initial_state=initial, + prefill_lengths=[1], + output_final_state=True, + use_qk_l2norm_in_kernel=True, + gate_inputs=gate_inputs, + ) + + expected = forward(v, gates) + torch.cuda.synchronize() + return v, gates, initial, forward, expected + + +@pytest.mark.parametrize("strided", ["values", "gates"]) +@torch.no_grad() +def test_gdn_strided_inputs_match_contiguous(compiled_gdn_layout_case, strided): + v, gates, initial, forward, expected = compiled_gdn_layout_case + original_state = initial.clone() + if strided == "values": + v = v.transpose(-1, -2).contiguous().transpose(-1, -2) + else: + gates = tuple(torch.stack((x, torch.zeros_like(x)), dim=-1)[..., 0] for x in gates) + actual = forward(v, gates) + for value, reference in zip(actual, expected): + torch.testing.assert_close(value, reference, rtol=0, atol=0) + torch.testing.assert_close(initial, original_state, rtol=0, atol=0) + + @pytest.mark.parametrize( - ("overrides", "message"), + ("compiled_serving_case", "overrides", "message"), [ + ("gdn", {"gate_inputs": None}, "raw gate_inputs"), ( + "kimi_k3", {"gate_inputs": None, "use_gate_in_kernel": False, "use_beta_sigmoid_in_kernel": False}, "raw gate_inputs", ), - ({"use_qk_l2norm_in_kernel": False}, "normalization"), + ("kimi_k3", {"use_qk_l2norm_in_kernel": False}, "normalization"), ], + indirect=["compiled_serving_case"], + scope="module", ) def test_rejects_silent_serving_mismatches(compiled_serving_case, overrides, message): _, _, forward, _, _ = compiled_serving_case @@ -259,6 +303,21 @@ def test_rejects_silent_serving_mismatches(compiled_serving_case, overrides, mes forward(**overrides) +@pytest.mark.parametrize("compiled_serving_case", ["gdn", "kimi_k3"], indirect=True, scope="module") +@torch.no_grad() +def test_raw_projections_match_bf16_serving(compiled_serving_case): + args, quantizer, forward, _, kimi_k3 = compiled_serving_case + inputs = [x.detach() for x in args] + inputs[3:5] = [x.float() + 0.0007 for x in inputs[3:5]] + # Rates and biases keep their FP32 precision, independent of projection rounding. + inputs[5:] = [x + 0.0123 for x in inputs[5:]] + rounded = [*inputs[:3], *[x.bfloat16() for x in inputs[3:5]], *inputs[5:]] + expected = _native_forward(rounded, quantizer, kimi_k3, kimi_k3) + actual = forward(inputs=inputs) + for value, reference in zip(actual, expected): + torch.testing.assert_close(value, reference, rtol=0, atol=0) + + @pytest.mark.parametrize( ("compiled_serving_case", "packed"), [("gdn", False), ("kda", True)], From 5c9797aff38d79920fc6e83e0ac30b534505b45d Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 13 Oct 2026 23:14:01 -0700 Subject: [PATCH 24/27] Move ReplaySSM state QAT into a separate follow-up Signed-off-by: Kai Xu --- .../linear_attention/serving/__init__.py | 4 +- .../linear_attention/serving/forward.py | 4 - .../linear_attention/serving/replay.py | 136 ------------ .../quantization/linear_attention/config.py | 25 +-- .../quantization/linear_attention/decode.py | 168 +-------------- .../quantization/linear_attention/kda.py | 2 +- .../quantization/linear_attention/training.py | 45 ++-- .../quantization/linear_attention/utils.py | 11 +- .../quantization/plugins/linear_attention.py | 1 - modelopt_recipes/configs/ptq/units/README.md | 3 - .../linear_attention_state_int8_dynamic.yaml | 34 --- .../linear_attention_state_int8_dynamic.yaml | 36 ---- modelopt_recipes/ptq.md | 12 +- .../test_linear_attention_replay.py | 204 ------------------ .../torch/quantization/plugins/test_gdn.py | 4 +- .../test_linear_attention_decode.py | 4 - .../test_linear_attention_hadamard.py | 35 --- 17 files changed, 25 insertions(+), 703 deletions(-) delete mode 100644 modelopt/torch/kernels/quantization/linear_attention/serving/replay.py delete mode 100644 modelopt_recipes/configs/ptq/units/linear_attention_state_int8_dynamic.yaml delete mode 100644 modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml delete mode 100644 tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py delete mode 100644 tests/unit/torch/quantization/test_linear_attention_hadamard.py diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py index 2f4c03234b5..de4e3c23c18 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py @@ -19,7 +19,5 @@ import vllm except ImportError as error: raise ImportError( - "Linear-attention serving profiles require the optional vLLM dependency: " - "vllm==0.30.0 for 'vllm', or the compatible quantized-ReplaySSM " - "vLLM fork for 'replayssm'." + "Linear-attention serving profiles require the optional vllm==0.30.0 dependency." ) from error diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py index 6088458c391..f0916d224e9 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py @@ -23,10 +23,6 @@ @cache def _kernels(precision): # Defer optional serving/model dependencies until their explicit profile is selected. - if precision == "replayssm": - from .replay import _prefill_kernels - - return _prefill_kernels() import vllm if vllm.__version__ != "0.30.0": diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/replay.py b/modelopt/torch/kernels/quantization/linear_attention/serving/replay.py deleted file mode 100644 index 67f43db24a9..00000000000 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/replay.py +++ /dev/null @@ -1,136 +0,0 @@ -# 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. - -"""Functional training launches of the optional quantized-ReplaySSM serving kernels.""" - -import torch - -# This optional fork has an explicit API separate from the public vLLM 0.30 profile. -from vllm.model_executor.layers.fla.ops import fused_recurrent_replayssm as _native - - -def _prefill_kernels(): - # Keep this fork's import contract isolated from the public vLLM profile. - from vllm.model_executor.layers.fla import ops - from vllm.model_executor.layers.fla.ops import kda - - return ops, kda - - -def _storage(value, enabled): - heads, keys, values = value.shape - state = torch.zeros( - 2, heads, values, keys, device=value.device, dtype=torch.int8 if enabled else torch.float32 - ) - scales = ( - torch.zeros(2, heads, values // 32, keys, device=value.device, dtype=torch.float16) - if enabled - else None - ) - indices = torch.ones(1, device=value.device, dtype=torch.int32) - return state, scales, indices - - -def _decoded(state, scales): - values = state[1].transpose(-1, -2).float() - if scales is None: - return values, None - metadata = scales[1].transpose(-1, -2).contiguous() - return (values.unflatten(-1, (-1, 32)) * metadata.float().unsqueeze(-1)).flatten(-2), metadata - - -def checkpoint(value, enabled): - """Return the serving checkpoint's decoded Hadamard-basis values and FP16 scales.""" - state, scales, indices = _storage(value, enabled) - original = value.transpose(-1, -2).unsqueeze(0).contiguous().float() - if enabled: - _native.prefill_write_checkpoint( - original, state, None, None, scales, None, None, indices, 8, hadamard=True - ) - else: - state[1] = _native._apply_h_value_axis(original)[0] - return _decoded(state, scales) - - -def original_basis(value): - """Use serving's FP32 checkpoint transform when returning a dense state.""" - return _native._apply_h_value_axis(value.transpose(-1, -2).unsqueeze(0).contiguous())[ - 0 - ].transpose(-1, -2) - - -def step(q, k, v, gate, beta, carry, window, enabled, scale, normalize, gate_inputs): - """Launch one token on private cache buffers; never mutate tensors saved for backward.""" - anchor = carry.anchor - state, scales, indices = _storage(anchor, enabled) - if enabled: - metadata = carry.scales - codes = (anchor.unflatten(-1, (-1, 32)) / metadata.float().unsqueeze(-1)).round() - state[1] = codes.flatten(-2).transpose(-1, -2).to(torch.int8) - scales[1] = metadata.transpose(-1, -2) - else: - state[1] = anchor.transpose(-1, -2) - heads, keys, values = anchor.shape - d = torch.zeros(2, heads, window, values, device=q.device, dtype=torch.bfloat16) - kc = torch.zeros(2, heads, window, keys, device=q.device, dtype=torch.bfloat16) - channel = gate.ndim == 2 - gc = torch.zeros( - (2, heads, window, keys) if channel else (2, heads, window), - device=q.device, - dtype=torch.float32, - ) - for i, entry in enumerate(carry.entries): - d[1, :, i] = entry.update - kc[1, :, i] = entry.key - gc[1, :, i] = entry.log_retention - mixed = torch.cat([x.to(torch.bfloat16).flatten() for x in (q, k, v)])[None] - if channel: - a, b, a_log, bias = gate.flatten()[None], beta[None], None, None - else: - raw_gate, raw_beta, a_log, bias = gate_inputs - a, b = raw_gate[None].contiguous(), raw_beta[None].contiguous() - output = torch.empty(1, heads, values, device=q.device, dtype=torch.bfloat16) - cursor = len(carry.entries) - _native.fused_recurrent_gated_delta_rule_replayssm( - mixed, - a, - b, - a_log, - bias, - scale, - state, - d, - kc, - gc, - output, - indices, - torch.full((1,), cursor, device=q.device, dtype=torch.int32), - use_qk_l2norm_in_kernel=normalize, - quant_state_bits=8 if enabled else 0, - state_scale=scales, - hadamard_value_basis=True, - block_v=64, - nk=2 if keys == 32 else None, - _is_kda=channel, - ) - decoded, metadata = _decoded(state, scales) - return ( - output[0].float(), - decoded, - metadata, - kc[1, :, cursor].float(), - d[1, :, cursor].float(), - gc[1, :, cursor], - ) diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index 1ee036ee809..09c07c44485 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -17,7 +17,7 @@ from typing import Literal -from pydantic import Field, model_validator +from pydantic import Field from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField @@ -42,36 +42,17 @@ class LinearAttentionConfig(ModeloptBaseConfig): they do not configure or validate a remote server. Kimi serving must set both ``kda_prefill_backend`` and ``kda_decode_backend`` to ``triton``; the supported gate is unbounded softplus. - ``replayssm`` uses the serving fork's INT8/Hadamard checkpoint and ring kernels. - ``replay_window=1`` refreshes the state every token; larger windows enable replay. Fresh prefill has no internal state QDQ; incoming continuation state and the handoff to a nonempty suffix use the selected checkpoint encoding. - TensorQuantizer independently enables QDQ and owns its format and grouping. + TensorQuantizer enables QDQ and owns its format and grouping. ``state_block_v`` controls grouping only for legacy per-tile quantizers. """ backend: Literal["fla", "serving"] = ModeloptField(default="fla") - precision: Literal["vllm", "vllm_kimi_k3", "replayssm"] = ModeloptField(default="vllm") - replay_window: int = Field(default=1, ge=1, le=64, strict=True) + precision: Literal["vllm", "vllm_kimi_k3"] = ModeloptField(default="vllm") state_block_v: Literal[16, 32, 64, 128] = ModeloptField(default=64) - @property - def state_codec(self): - """The native profile fixes the codec; it is not an independent setting.""" - return "int8_hadamard32" if self.precision == "replayssm" else "tile" - - @model_validator(mode="after") - def _validate_profile(self): - if self.precision == "replayssm": - if self.backend != "serving": - raise ValueError("replayssm requires backend='serving'") - if self.state_block_v < 32: - raise ValueError("int8_hadamard32 requires state_block_v >= 32") - elif self.replay_window != 1: - raise ValueError("Replay requires precision='replayssm'") - return self - class LinearAttentionPolicyEntry(ModeloptBaseConfig): """Assign a complete policy to supported modules matching ``module_name``. diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py index 7cc94d11170..52a611ad62d 100644 --- a/modelopt/torch/quantization/linear_attention/decode.py +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -13,100 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Differentiable token-state and encoded-update replay for QAT.""" - -from dataclasses import dataclass +"""Differentiable per-token state quantization for a training suffix.""" import torch -from ._vllm_autograd import normalized -from .utils import _state_qdq, forward_value +from .utils import _state_qdq __all__ = [] -@dataclass -class _ReplayEntry: - """One saved rank-one update since the last state checkpoint. - - ``key`` is [H,K] and ``update`` is the beta-scaled residual [H,V]. ReplaySSM - stores their BF16 forward values once. ``log_retention`` is [H] for GDN or - [H,K] for KDA and records this token's decay for subsequent reconstruction. - """ - - key: torch.Tensor - update: torch.Tensor - log_retention: torch.Tensor - - -@dataclass -class _ReplayState: - """Checkpoint and pending updates within one training suffix. - - Refreshing the replay window replaces the checkpoint and clears the entries. - Tensors retain their autograd history; this state never crosses forward calls. - """ - - anchor: torch.Tensor # Decoded checkpoint [H,K,V] in the Hadamard value basis. - scales: torch.Tensor | None # Native FP16 scales needed to rebuild an INT8 cache. - entries: tuple[_ReplayEntry, ...] = () - - def reconstruct(self): - """Return the dense state in the original value basis without another QDQ.""" - # The optional serving fork owns the checkpoint basis transform. - from ...kernels.quantization.linear_attention.serving.replay import original_basis - - return original_basis(_replay_state(self)) - - -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) - - -def _replay_state(carry: _ReplayState, current_gate=None): - """Decay the anchor and add saved updates weighted by their subsequent gates. - - KDA weights keys per channel; GDN weights updates per head. ``current_gate`` - includes the next KDA token's decay before BF16 weighting, without adding its update. - """ - state = carry.anchor - if not carry.entries: - if current_gate is not None: - state = state * current_gate.exp().unsqueeze(-1) - return state - gates = torch.stack([entry.log_retention for entry in carry.entries], dim=1) - keys = torch.stack([entry.key for entry in carry.entries], dim=1) - updates = torch.stack([entry.update for entry in carry.entries], dim=1) - total = gates.sum(1) - if current_gate is not None: - total = total + current_gate - weights = (total.unsqueeze(1) - gates.cumsum(1)).exp() - if gates.ndim == 3: - weighted = keys * weights - weighted = forward_value(weighted, weighted.to(torch.bfloat16)) - return state * total.exp().unsqueeze(-1) + weighted.transpose(-1, -2) @ updates - weighted = updates * weights.unsqueeze(-1) - weighted = forward_value(weighted, weighted.to(torch.bfloat16)) - return state * total.exp()[:, None, None] + keys.transpose(-1, -2) @ weighted - - -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 _recurrent_decode( q, k, @@ -132,21 +47,6 @@ def _recurrent_decode( zero = (q.sum() + k.sum() + v.sum() + g.sum() + beta.sum()) * 0 return v + zero, initial_state if output_final_state else None scale = q.shape[-1] ** -0.5 if scale is None else scale - if config.precision == "replayssm": - output, carry = _replay_decode( - q, - k, - v, - g, - beta, - initial_state, - config.replay_window, - state_quantizer is not None, - scale, - use_qk_l2norm_in_kernel, - gate_inputs, - ) - return output, carry.reconstruct() if output_final_state else None # Import the optional native backend only for the ordinary state-QDQ path. from ._vllm_autograd import step @@ -173,67 +73,3 @@ def _recurrent_decode( # A new tensor preserves the state saved by earlier tokens' backward passes. state = _state_qdq(working, state_quantizer, config.state_block_v) return torch.stack(outputs), state if output_final_state else None - - -def _replay_decode(q, k, v, g, beta, initial_state, window, enabled, scale, normalize, gate_inputs): - """Attach native ReplaySSM values to the differentiable checkpoint/update recurrence.""" - # This optional fork owns checkpoint rounding and the Hadamard value basis. - from ...kernels.quantization.linear_attention.serving.replay import checkpoint, step - - initial = _hadamard32(initial_state) - with torch.no_grad(): - decoded, scales = checkpoint(initial_state, enabled) - carry = _ReplayState(forward_value(initial, decoded), scales) - original_v, v = v, _hadamard32(v) - with torch.autocast(device_type=q.device.type, enabled=False): - outputs, native_outputs = [], [] - for t in range(len(q)): - gate = g[t] - with torch.no_grad(): - native = step( - q[t], - k[t], - original_v[t], - gate, - beta[t], - carry, - window, - enabled, - scale, - normalize, - (gate_inputs[0][t], gate_inputs[1][t], *gate_inputs[2:]) - if gate_inputs is not None - else None, - ) - native_outputs.append(native[0]) - decay = gate.exp().unsqueeze(-1) - if gate.ndim == 1: - decay = decay.unsqueeze(-1) - key = k[t] - if normalize: - key = normalized(key) - decayed = _replay_state(carry, gate) if gate.ndim == 2 else _replay_state(carry) * decay - residual = v[t] - _sum_keys(key.unsqueeze(-1) * decayed) - update = beta[t].unsqueeze(-1) * residual - working = decayed + key.unsqueeze(-1) * update.unsqueeze(-2) - if len(carry.entries) + 1 == window: - # Native step already quantized the checkpoint; preserve its scales and adjoint. - anchor, scales = forward_value(working, native[1]), native[2] - entries = () - else: - anchor, scales = carry.anchor, carry.scales - entries = ( - *carry.entries, - _ReplayEntry( - forward_value(key, native[3]), - forward_value(update, native[4]), - forward_value(gate, native[5]), - ), - ) - carry = _ReplayState(anchor, scales, entries) - query = q[t] - if normalize: - query = normalized(query) - # Read the working state before checkpoint rounding, including when the window fills. - outputs.append(_sum_keys(query.unsqueeze(-1) * working) * scale) - return forward_value(_hadamard32(torch.stack(outputs)), torch.stack(native_outputs)), carry diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py index f6bdb7565ae..d16ebd69bfa 100644 --- a/modelopt/torch/quantization/linear_attention/kda.py +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -92,7 +92,7 @@ def kda_state_qat( from ...kernels.quantization.linear_attention.serving.forward import _kernels with torch.no_grad(): - _, native_kda = _kernels("replayssm" if policy.precision == "replayssm" else "vllm") + _, native_kda = _kernels("vllm") native_gate = native_kda.fused_kda_gate( raw_gate.flatten(-2).contiguous(), A_log, raw_gate.shape[-1], g_bias=dt_bias ) diff --git a/modelopt/torch/quantization/linear_attention/training.py b/modelopt/torch/quantization/linear_attention/training.py index d40a9fae5c0..011ecea2753 100644 --- a/modelopt/torch/quantization/linear_attention/training.py +++ b/modelopt/torch/quantization/linear_attention/training.py @@ -25,7 +25,7 @@ from packaging.version import Version from .decode import _recurrent_decode -from .utils import _resolve_state_quantizer, _state_qdq, forward_value +from .utils import _resolve_state_quantizer, _state_qdq __all__ = ["get_linear_attention_layers", "linear_attention_training_phase"] @@ -158,8 +158,6 @@ def _prepare_prefill_inputs(q, k, v, g, beta, *, policy, chunk_size): raise ValueError("State training requires backend='serving' and chunk_size=64") if q.device.type != "cuda" or any(x.dtype != torch.bfloat16 for x in (q, k, v)): raise ValueError("Serving arithmetic requires CUDA BF16 Q/K/V inputs") - if policy.state_codec == "int8_hadamard32" and v.shape[-1] % 32: - raise ValueError("int8_hadamard32 requires Dv divisible by 32") _validate_training_device(q.device) return tuple(x.float() for x in (q, k, v, g, beta)) @@ -188,9 +186,9 @@ def _prefill_decode_forward( ): """Run both prefill and decode phases in one differentiable training forward. - Each sequence's chunked prefix produces the state for its token or ReplaySSM - suffix. The suffix preserves its gradient connection to prefill and returns - a dense final state only when requested. Outputs are joined in token order. + Each sequence's chunked prefix produces the state for its recurrent suffix. + The suffix preserves its gradient connection to prefill and returns a dense + final state only when requested. Outputs are joined in token order. Args: prefill_lengths: Prefix token count per sequence. For 128 tokens, a value @@ -198,7 +196,7 @@ def _prefill_decode_forward( sequence_lengths: Valid token count per sequence, excluding right padding. Defaults to the dense row length or each packed storage segment length. """ - state_quantizer = _resolve_state_quantizer(state_quantizer, state_codec=policy.state_codec) + state_quantizer = _resolve_state_quantizer(state_quantizer) if state_quantizer is not None and prefill_lengths is None: raise ValueError( "State quantization requires explicit prefill lengths through " @@ -254,20 +252,14 @@ def _prefill_decode_forward( raise ValueError( "Serving arithmetic requires K <= 256, KDA V=K, and power-of-two K for normalization" ) - if policy.precision == "replayssm": - if any(prefix < length for prefix, length in zip(prefixes, valid_lengths)) and ( - keys < 32 or keys & (keys - 1) or (g.ndim == 3 and gate_inputs is None) - ): - raise ValueError("ReplaySSM requires power-of-two K >= 32 and raw GDN gate inputs") - else: - # Validate the optional native backend before launching any prefix kernels. - from ...kernels.quantization.linear_attention.serving.forward import validate_profile + # Validate the optional native backend before launching any prefix kernels. + from ...kernels.quantization.linear_attention.serving.forward import validate_profile - validate_profile( - gate_inputs, - use_qk_l2norm_in_kernel, - precision=policy.precision, - ) + validate_profile( + gate_inputs, + use_qk_l2norm_in_kernel, + precision=policy.precision, + ) if initial_state is None: states = q.new_zeros(len(sequences), heads, keys, values) else: @@ -287,18 +279,7 @@ def _prefill_decode_forward( # A continuation prefill consumes a stored cache just as native serving # does. A fresh zero-state prefix has no incoming cache to quantize. if initial_state is not None and state_quantizer is not None: - if policy.precision == "replayssm": - from ...kernels.quantization.linear_attention.serving.replay import ( - checkpoint, - original_basis, - ) - - with torch.no_grad(): - decoded, _ = checkpoint(state, True) - decoded = original_basis(decoded) - state = forward_value(state, decoded) - else: - state = _state_qdq(state, state_quantizer, policy.state_block_v) + state = _state_qdq(state, state_quantizer, policy.state_block_v) prefix, state = serving_prefix( q=q[b, start:split], k=k[b, start:split], diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py index 9f44d2a44d7..f539919d472 100644 --- a/modelopt/torch/quantization/linear_attention/utils.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -117,7 +117,7 @@ def validate_gdn_quantizer( def state_quantizer_config( - quantizer: TensorQuantizer, *, name="state_quantizer", state_codec="tile" + quantizer: TensorQuantizer, *, name="state_quantizer" ) -> tuple[str, int]: """Validate the state quantizer and derive its format and last-axis group size. @@ -126,11 +126,6 @@ def state_quantizer_config( validate_gdn_quantizer( quantizer, name=name, num_bits=tuple(_STATE_FORMATS), block_sizes=(16, 32, 64) ) - if state_codec == "int8_hadamard32": - if quantizer.num_bits != 8: - raise ValueError("int8_hadamard32 requires INT8 state quantization") - if quantizer.block_sizes is not None: - raise ValueError("TensorQuantizer block_sizes requires state_codec='tile'") if quantizer.block_sizes is not None: return _STATE_FORMATS[quantizer.num_bits], quantizer.block_sizes[-1] if quantizer.axis != (0, 1): @@ -138,10 +133,10 @@ def state_quantizer_config( return _STATE_FORMATS[quantizer.num_bits], 0 -def _resolve_state_quantizer(state_quantizer, *, state_codec="tile"): +def _resolve_state_quantizer(state_quantizer): """Validate an enabled quantizer once; return None when state QDQ is inactive.""" if state_quantizer is not None and state_quantizer.is_enabled: - state_quantizer_config(state_quantizer, state_codec=state_codec) + state_quantizer_config(state_quantizer) if state_quantizer._if_quant: return state_quantizer return None diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py index 7eea8b4bfbc..726384b2fbc 100644 --- a/modelopt/torch/quantization/plugins/linear_attention.py +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -117,7 +117,6 @@ def validate_linear_attention(self): state_quantizer_config( self._linear_attn_state, name=self.linear_attention_quantizer_names[0], - state_codec=self.linear_attention_config.state_codec, ) if self.linear_attention_config.backend != "serving": raise ValueError( diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index d02eff834ca..e997841ab47 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -34,11 +34,8 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `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 serving cache boundaries; requires an explicit serving execution policy | -| `linear_attention_state_int8_dynamic.yaml` | GDN/KDA INT8 state quantizer entries; the [complete recipe](../../../general/ptq/linear_attention_state_int8_dynamic.yaml) enables Hadamard during decode | | `linear_attention_state_int8_block32_dynamic.yaml` | Standard dynamic INT8 with one scale per key row and 32 value channels; the [complete training recipe](../../../general/ptq/linear_attention_state_int8_block32_dynamic.yaml) uses token QDQ and working-state readout | | `indexer_k_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer key cache (`*indexer_k_quantizer`), global scale fixed to 1; Blackwell+ GPUs | | `indexer_q_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer query (`*indexer_q_quantizer`), global scale fixed to 1; Blackwell+ GPUs | | `moe_dispatch_nvfp4.yaml` | NVFP4 fake quantization of the MoE expert-parallel dispatch (`*dispatch_quantizer`, vLLM); Blackwell+ GPUs | | `moe_combine_nvfp4.yaml` | NVFP4 fake quantization of the MoE expert-parallel combine (`*combine_quantizer`, vLLM); Blackwell+ GPUs | - -Native ReplaySSM uses BF16 key/update vectors; they are not independently quantized. diff --git a/modelopt_recipes/configs/ptq/units/linear_attention_state_int8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/linear_attention_state_int8_dynamic.yaml deleted file mode 100644 index eceadf5d5d0..00000000000 --- a/modelopt_recipes/configs/ptq/units/linear_attention_state_int8_dynamic.yaml +++ /dev/null @@ -1,34 +0,0 @@ -# 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 for GDN and KDA recurrent states. -# The complete general/ptq/linear_attention_state_int8_dynamic recipe enables Hadamard. -# 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 - - quantizer_name: '*kda_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/general/ptq/linear_attention_state_int8_dynamic.yaml b/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml deleted file mode 100644 index b21cc4c2653..00000000000 --- a/modelopt_recipes/general/ptq/linear_attention_state_int8_dynamic.yaml +++ /dev/null @@ -1,36 +0,0 @@ -# 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. - -# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe -imports: - base_disable_all: configs/ptq/units/base_disable_all - state_int8: configs/ptq/units/linear_attention_state_int8_dynamic - -metadata: - description: >- - Dynamic INT8 GDN/KDA state QDQ with 32-value Hadamard rotation during decode. - Imports the quantized-ReplaySSM serving kernels and uses working-state readout. - Fresh prefill has no internal state QDQ; supply per-sequence prefix lengths through - linear_attention_training_phase. No calibration is required. -quantize: - algorithm: - quant_cfg: - - $import: base_disable_all - - $import: state_int8 - linear_attention: - - module_name: '*' - cfg: - backend: serving - precision: replayssm diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 0fe785c28f7..845c1643bc1 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 35 general/ptq/ recipes (click to expand) +All 34 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -64,7 +64,6 @@ supported combinations. | `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | | `iq2_s` | IQ2_S W2A16 (2.56 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | | `q8_0` | Q8_0 W8A16 (8.5 bpw), eligible linears | none | none (no calibration) | -| `linear_attention_state_int8_dynamic` | GDN/KDA decode state INT8 + Hadamard; requires a private ReplaySSM fork | none | none (dynamic scales; requires a prefix/decode phase context) | | `linear_attention_state_int8_block32_dynamic` | GDN/KDA state INT8, 32 value channels per key row; weights unchanged | none | none (dynamic scales; requires vLLM and a prefix/decode phase context) | | `linear_attention_state_int8_block32_dynamic_kimi_k3` | Kimi-K3/Kimi-Linear KDA state INT8, 32 value channels per key row | none | none (dynamic scales; explicit Triton Kimi profile) | @@ -96,17 +95,8 @@ It never falls back to prepared-only arithmetic. The adapters use the fixed vLLM 0.30.0 value-first cache layout. Native-forward comparisons and independent recurrence checks validate it in GPU tests. -The `linear_attention_state_int8_dynamic` recipe uses an optional private -quantized-ReplaySSM fork; public vLLM releases do not provide those kernels. -Its imports and kernel contract are isolated from the public 0.30.0 profiles; -it requires the fork's `vllm.model_executor.layers.fla.ops` API. Without access -to that fork, use a block32 recipe above. ReplaySSM native tests -skip when the fork is absent, so the public CI images do not qualify this profile. - The public profiles do not support vLLM 0.20 or arbitrary development builds. The tested GDN/KDA profiles pass bit-exact native forward comparisons on vLLM 0.30.0. -The existing sm_120+/Triton <3.7 training guard also applies to custom ReplaySSM -environments; changing the public profile does not remove that hardware restriction. Megatron state QAT requires native `_compute_gates` and either `forward_pre_attn_and_core_attn` or `_forward_compute` for raw gates and packed diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py deleted file mode 100644 index 86c9aedfbe4..00000000000 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_replay.py +++ /dev/null @@ -1,204 +0,0 @@ -# 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. - -"""Minimal native-cache and backward checks for the optional ReplaySSM serving fork.""" - -from importlib import import_module - -import pytest -import torch - -from modelopt.torch.quantization.config import QuantizerAttributeConfig -from modelopt.torch.quantization.linear_attention import LinearAttentionConfig, gdn_state_qat -from modelopt.torch.quantization.linear_attention.decode import _replay_decode -from modelopt.torch.quantization.linear_attention.training import _validate_training_device -from modelopt.torch.quantization.nn import TensorQuantizer - -pytest.importorskip("vllm") - -try: - _validate_training_device(torch.device("cuda")) -except NotImplementedError as error: - pytest.skip(str(error), allow_module_level=True) - - -@pytest.fixture(scope="module") -def state_quantizer(): - return TensorQuantizer( - QuantizerAttributeConfig( - num_bits=8, - type="dynamic", - axis=(0, 1), - narrow_range=True, - pass_through_bwd=True, - ) - ) - - -@pytest.fixture( - scope="module", - params=[(False, 1), (False, 4), (True, 4)], - ids=["gdn-token", "gdn-replay", "kda-replay"], -) -def compiled_replay(request, state_quantizer): - try: - native = import_module("vllm.model_executor.layers.fla.ops.fused_recurrent_replayssm") - except ModuleNotFoundError as error: - if error.name and "vllm.model_executor.layers.fla.ops.fused_recurrent_replayssm".startswith( - error.name - ): - pytest.skip("Requires the optional quantized-ReplaySSM serving fork") - raise - channel, window = request.param - torch.manual_seed(71) - q, k = [ - torch.nn.functional.normalize(torch.randn(5, 1, 64, device="cuda"), dim=-1).bfloat16() - for _ in range(2) - ] - v = torch.randn_like(q) - raw = torch.randn(5, 1, device="cuda", dtype=torch.bfloat16) - raw_beta = torch.randn_like(raw) - rate = torch.full((1,), -3.0, device="cuda") - bias = torch.zeros_like(rate) - g = ( - -torch.rand_like(q, dtype=torch.float32) * 0.03 - if channel - else -rate.exp() * torch.nn.functional.softplus(raw.float()) - ) - beta = raw_beta.float().sigmoid() - args = tuple(x.detach().requires_grad_() for x in (q, k, v, g, beta)) - initial = (torch.randn(1, 64, 64, device="cuda") * 0.1).requires_grad_() - config = LinearAttentionConfig( - backend="serving", - precision="replayssm", - replay_window=window, - ) - kwargs = { - "config": config, - "state_quantizer": state_quantizer, - "gate_inputs": None if channel else (raw, raw_beta, rate, bias), - } - - def forward(): - return _replay_decode( - *(x.float() for x in args), - initial, - window, - state_quantizer.is_enabled, - 64**-0.5, - False, - kwargs["gate_inputs"], - ) - - out, carry = forward() - torch.autograd.grad(out.sum() + carry.reconstruct().sum(), (*args, initial)) - torch.cuda.synchronize() - return native, args, initial, kwargs, forward - - -@pytest.mark.parametrize( - ("keys", "values", "overrides", "message"), - [ - (32, 48, {}, "Dv divisible by 32"), - (32, 32, {"num_bits": (4, 3)}, "requires INT8"), - (32, 32, {"block_sizes": {-1: 32}, "axis": None}, "requires state_codec='tile'"), - (16, 32, {}, "power-of-two K >= 32"), - (32, 32, {}, "raw GDN gate inputs"), - ], -) -def test_replay_rejects_invalid_inputs_before_backend(keys, values, overrides, message): - quantizer = TensorQuantizer( - QuantizerAttributeConfig(num_bits=8, type="dynamic", axis=(0, 1), narrow_range=True) - ) - quantizer.set_from_attribute_config(overrides) - q = torch.ones(1, 2, 1, keys, device="cuda", dtype=torch.bfloat16) - v = q.new_ones(1, 2, 1, values) - with pytest.raises(ValueError, match=message): - gdn_state_qat( - q, - q, - v, - q[..., 0], - q[..., 0], - policy=LinearAttentionConfig(backend="serving", precision="replayssm"), - state_quantizer=quantizer, - prefill_lengths=[1], - initial_state=q.new_ones(1, 1, keys, values), - ) - - -def test_replay_matches_persistent_serving_cache(compiled_replay): - native, args, initial, kwargs, forward = compiled_replay - output, carry = forward() - channel = args[3].ndim == 3 - window = kwargs["config"].replay_window - state = torch.zeros(2, 1, 64, 64, device="cuda", dtype=torch.int8) - scales = torch.zeros(2, 1, 2, 64, device="cuda", dtype=torch.float16) - updates = torch.zeros(2, 1, window, 64, device="cuda", dtype=torch.bfloat16) - keys = torch.zeros_like(updates) - gates = torch.zeros((2, 1, window, 64) if channel else (2, 1, window), device="cuda") - indices = torch.ones(1, device="cuda", dtype=torch.int32) - with torch.no_grad(): - native.prefill_write_checkpoint( - initial.transpose(-1, -2)[None].contiguous(), - state, - None, - None, - scales, - None, - None, - indices, - 8, - hadamard=True, - ) - for t in range(len(args[0])): - q, k, v, gate, beta = (x[t] for x in args) - if channel: - a, b, rate, bias = gate.flatten()[None], beta[None], None, None - else: - raw, raw_beta, rate, bias = kwargs["gate_inputs"] - a, b = raw[t : t + 1], raw_beta[t : t + 1] - expected = torch.empty(1, 1, 64, device="cuda", dtype=torch.bfloat16) - native.fused_recurrent_gated_delta_rule_replayssm( - torch.cat([x.flatten() for x in (q, k, v)])[None], - a, - b, - rate, - bias, - 64**-0.5, - state, - updates, - keys, - gates, - expected, - indices, - torch.tensor([t % window], device="cuda", dtype=torch.int32), - quant_state_bits=8, - state_scale=scales, - hadamard_value_basis=True, - block_v=64, - _is_kda=channel, - ) - torch.testing.assert_close(output[t].bfloat16(), expected[0], rtol=0, atol=0) - decoded = state[1].transpose(-1, -2).float().unflatten(-1, (-1, 32)) - metadata = scales[1].transpose(-1, -2) - torch.testing.assert_close( - carry.anchor, (decoded * metadata[..., None]).flatten(-2), rtol=0, atol=0 - ) - torch.testing.assert_close(carry.scales, metadata, rtol=0, atol=0) - grads = torch.autograd.grad( - output.square().mean() + carry.reconstruct().square().mean(), (*args, initial) - ) - assert all(torch.isfinite(grad).all() and grad.abs().sum() > 0 for grad in grads) diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index aab9a3d284d..b7af95761c4 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -144,9 +144,7 @@ def test_restore_uses_saved_quantizer_and_policy(tmp_path, monkeypatch): model[0].gdn_state_quantizer.set_from_attribute_config( {"num_bits": 8, "axis": (0, 1), "type": "dynamic", "narrow_range": True} ) - model[0].linear_attention_config = LinearAttentionConfig( - backend="serving", precision="replayssm", replay_window=4 - ) + model[0].linear_attention_config = LinearAttentionConfig(backend="serving", state_block_v=32) sample = torch.randn(2, 4, 19) expected = model[0].gdn_state_quantizer(sample) path = tmp_path / "gdn.pth" diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py index 8cd5a994821..b8efe8a597c 100644 --- a/tests/unit/torch/quantization/test_linear_attention_decode.py +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -48,10 +48,6 @@ def test_state_qdq_matches_tensor_quantizer(state_format): {"chunk_size": 32}, {"solve": {"method": "neumann"}}, {"state": {"mode": "token"}}, - {"replay_window": 4}, - {"precision": "replayssm", "replay_window": 0}, - {"precision": "replayssm", "replay_window": 65}, - {"precision": "replayssm", "state_block_v": 16}, ], ) def test_linear_attention_policy_rejects_invalid_settings(settings): diff --git a/tests/unit/torch/quantization/test_linear_attention_hadamard.py b/tests/unit/torch/quantization/test_linear_attention_hadamard.py deleted file mode 100644 index 41acfc74aa3..00000000000 --- a/tests/unit/torch/quantization/test_linear_attention_hadamard.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 torch - -from modelopt.torch.quantization.linear_attention.decode import _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 test_hadamard_matches_dense_oracle_and_backward(): - torch.manual_seed(321) - state = torch.randn(2, 4, 64, dtype=torch.float64, requires_grad=True) - actual = _hadamard32(state) - expected = _rotate(state) - torch.testing.assert_close(actual, expected, rtol=1e-10, atol=1e-10) - probe = torch.randn_like(state) - (gradient,) = torch.autograd.grad((actual * probe).sum(), state) - torch.testing.assert_close(gradient, _rotate(probe), rtol=1e-10, atol=1e-10) From bcb2263f95e83befa0a53c8354023fab0d43913b Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 14 Oct 2026 12:29:10 -0700 Subject: [PATCH 25/27] Use one native state-QAT forward and reuse CI kernel shapes Signed-off-by: Kai Xu --- .../linear_attention/_vllm_autograd.py | 180 +++++++++++++----- .../test_linear_attention_training.py | 31 ++- .../torch/quantization/test_vllm_autograd.py | 144 ++++++++++++++ 3 files changed, 287 insertions(+), 68 deletions(-) create mode 100644 tests/gpu_vllm/torch/quantization/test_vllm_autograd.py diff --git a/modelopt/torch/quantization/linear_attention/_vllm_autograd.py b/modelopt/torch/quantization/linear_attention/_vllm_autograd.py index c32c213ff51..67169cf5847 100644 --- a/modelopt/torch/quantization/linear_attention/_vllm_autograd.py +++ b/modelopt/torch/quantization/linear_attention/_vllm_autograd.py @@ -15,31 +15,28 @@ """Differentiable adapter for the installed serving runtime's arithmetic. -Native kernels supply forward values. Prefill uses FLA's recurrence backward; -decode uses a Torch adjoint. QDQ uses identity STE. +Native kernels supply forward values. Prefill prepares only the intermediates for FLA backward; +decode uses an explicit recurrence adjoint. QDQ uses identity STE. The optional vLLM dependency supplies kernels; training owns its state and needs no server. """ import torch +from torch.autograd.function import once_differentiable -from .utils import forward_value +class _Prefill(torch.autograd.Function): + """Run native prefill once and reconstruct FLA's backward inputs on demand.""" -def normalized(value): - return value / (value.square().sum(-1, keepdim=True) + 1e-6).sqrt() + @staticmethod + def forward(ctx, q, k, v, g, beta, state, scale, beta_dtype, normalize, gate_inputs, precision): + # Keep optional inference kernels isolated to the serving profile. + from ...kernels.quantization.linear_attention.serving.forward import prefill - -def prefix( - q, k, v, g, beta, state, scale, beta_dtype, normalize=False, gate_inputs=None, precision="vllm" -): - # vLLM supplies the forward; the optional FLA dependency supplies the recurrence adjoint. - from ...kernels.quantization.linear_attention.serving.forward import prefill - - with torch.no_grad(): + q, k, v = (x.to(torch.bfloat16) for x in (q, k, v)) out, final = prefill( - q.to(torch.bfloat16), - k.to(torch.bfloat16), - v.to(torch.bfloat16), + q, + k, + v, g.float(), beta.to(beta_dtype), state.float(), @@ -48,32 +45,79 @@ def prefix( gate_inputs, precision=precision, ) - if not torch.is_grad_enabled() or not any(x.requires_grad for x in (q, k, v, g, beta, state)): + ctx.save_for_backward(q, k, v, g, beta, state) + ctx.scale, ctx.normalize = scale, normalize + ctx.set_materialize_grads(False) return out.float(), final - # FLA differentiates the same recurrence. Its intermediate rounding is a backward - # surrogate, not a claim that vLLM inference kernels have an exact backward. - if g.ndim == 3: - from fla.ops.kda import chunk_kda as chunk - else: - from fla.ops.gated_delta_rule import chunk_gated_delta_rule as chunk - adjoint, adjoint_state = chunk( - *[x.to(torch.bfloat16)[None] for x in (q, k, v)], - g[None], - beta[None], - initial_state=state[None], - output_final_state=True, - scale=scale, - use_qk_l2norm_in_kernel=normalize, + + @staticmethod + @once_differentiable + def backward(ctx, grad_output, grad_final): + # FLA is optional and needed only for prefix gradients, never native inference. + from fla.modules.l2norm import l2norm_bwd, l2norm_fwd + from fla.ops.utils import chunk_local_cumsum + from fla.ops.utils.constant import RCP_LN2 + + q, k, v, g, beta, state = (x[None].contiguous() for x in ctx.saved_tensors) + if ctx.normalize: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + cumulative = chunk_local_cumsum(g, chunk_size=64, scale=RCP_LN2) + kwargs = { + "q": q, + "k": k, + "v": v, + "g": cumulative, + "beta": beta, + "scale": ctx.scale, + "initial_state": state, + "do": torch.zeros_like(v) + if grad_output is None + else grad_output[None].to(v).contiguous(), + "dht": None if grad_final is None else grad_final[None].contiguous(), + } + # Rebuild only intra-chunk factors. The upstream backward already reconstructs + # chunk states; a second full prefill would also compute unused outputs/states. + if g.ndim == 4: + from fla.ops.kda.chunk_bwd import chunk_kda_bwd + from fla.ops.kda.chunk_intra import chunk_kda_fwd_intra + + aqk, akk = chunk_kda_fwd_intra(q, k, v, cumulative, beta, ctx.scale)[-2:] + dq, dk, dv, db, dg, dh, *_ = chunk_kda_bwd(**kwargs, Aqk=aqk, Akk=akk) + else: + from fla.ops.gated_delta_rule.chunk import chunk_gated_delta_rule_bwd + from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra + + a = chunk_gated_delta_rule_fwd_intra(k, v, cumulative, beta)[-1] + dq, dk, dv, db, dg, dh, *_ = chunk_gated_delta_rule_bwd(**kwargs, A=a) + if ctx.normalize: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + # Match FLA's BF16 operand-gradient rounding before upstream FP32 casts. + grads = tuple( + dx[0].to(x.dtype) for dx, x in zip((dq, dk, dv, dg, db, dh), (q, k, v, g, beta, state)) + ) + return (*grads, None, None, None, None, None) + + +def prefix( + q, k, v, g, beta, state, scale, beta_dtype, normalize=False, gate_inputs=None, precision="vllm" +): + """Return native prefix output/state with an upstream FLA recurrence backward.""" + return _Prefill.apply( + q, k, v, g, beta, state, scale, beta_dtype, normalize, gate_inputs, precision ) - return forward_value(adjoint[0].float(), out), forward_value(adjoint_state[0], final) -def step(q, k, v, gate, beta, state, scale, normalize=False, gate_inputs=None, precision="vllm"): - # Keep the optional vLLM dependency isolated to the native precision profile. - from ...kernels.quantization.linear_attention.serving.forward import step as native_step +class _DecodeStep(torch.autograd.Function): + """Run native decode once and differentiate the recurrence only during backward.""" + + @staticmethod + def forward(ctx, q, k, v, gate, beta, state, scale, normalize, gate_inputs, precision): + # Autograd Function.forward runs without recording the inference kernel's operations. + from ...kernels.quantization.linear_attention.serving.forward import step as native_step - with torch.no_grad(): - out, final = native_step( + output, final = native_step( q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16), @@ -85,18 +129,50 @@ def step(q, k, v, gate, beta, state, scale, normalize=False, gate_inputs=None, p gate_inputs, precision=precision, ) - if not torch.is_grad_enabled() or not any( - x.requires_grad for x in (q, k, v, gate, beta, state) - ): - return out.float(), final - if normalize: - q, k = normalized(q), normalized(k) - decay = gate.exp().unsqueeze(-1) - if gate.ndim == 1: - decay = decay.unsqueeze(-1) - decayed = state * decay - residual = v - (decayed * k.unsqueeze(-1)).sum(-2) - update = residual * beta.unsqueeze(-1) - working = forward_value(decayed + k.unsqueeze(-1) * update.unsqueeze(-2), final) - output = ((q * scale).unsqueeze(-1) * working).sum(-2) - return forward_value(output, out), working + ctx.save_for_backward(q, k, v, gate, beta, state, final) + ctx.scale, ctx.normalize = scale, normalize + ctx.set_materialize_grads(False) + return output.float(), final + + @staticmethod + @once_differentiable + def backward(ctx, grad_output, grad_final): + if grad_output is None and grad_final is None: + return (None,) * 10 + q, k, v, gate, beta, state, final = ctx.saved_tensors + if ctx.normalize: + q_norm = (q.square().sum(-1, keepdim=True) + 1e-6).sqrt() + k_norm = (k.square().sum(-1, keepdim=True) + 1e-6).sqrt() + q, k = q / q_norm, k / k_norm + decay = gate.exp().unsqueeze(-1) + if gate.ndim == 1: + decay = decay.unsqueeze(-1) + decayed = state * decay + residual = v - (decayed * k.unsqueeze(-1)).sum(-2) + + # Output reads the native working state before the external QDQ cache write. + grad_q, grad_working = None, grad_final + if grad_output is not None: + grad_q = (final * grad_output.unsqueeze(-2)).sum(-1) * ctx.scale + readout_grad = (q * ctx.scale).unsqueeze(-1) * grad_output.unsqueeze(-2) + grad_working = readout_grad if grad_final is None else grad_final + readout_grad + grad_update = (grad_working * k.unsqueeze(-1)).sum(-2) + grad_v = grad_update * beta.unsqueeze(-1) + grad_beta = (grad_update * residual).sum(-1) + grad_k = (grad_working * (residual * beta.unsqueeze(-1)).unsqueeze(-2)).sum(-1) + grad_k = grad_k - (decayed * grad_v.unsqueeze(-2)).sum(-1) + grad_decayed = grad_working - k.unsqueeze(-1) * grad_v.unsqueeze(-2) + grad_state = grad_decayed * decay + grad_gate = (grad_decayed * decayed).sum(-1) + if gate.ndim == 1: + grad_gate = grad_gate.sum(-1) + if ctx.normalize: + if grad_q is not None: + grad_q = (grad_q - q * (grad_q * q).sum(-1, keepdim=True)) / q_norm + grad_k = (grad_k - k * (grad_k * k).sum(-1, keepdim=True)) / k_norm + return grad_q, grad_k, grad_v, grad_gate, grad_beta, grad_state, None, None, None, None + + +def step(q, k, v, gate, beta, state, scale, normalize=False, gate_inputs=None, precision="vllm"): + """Return native output/state with a recurrence backward; the caller owns state QDQ.""" + return _DecodeStep.apply(q, k, v, gate, beta, state, scale, normalize, gate_inputs, precision) diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py index d29101b3dd3..5054ffddb7e 100644 --- a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py @@ -219,12 +219,14 @@ def test_serving_state_qdq_and_handoff_gradient(compiled_serving_case): assert not torch.equal(output[:, 65:], plain[:, 65:]) # A nonzero nonsymmetric state checks layout against an independent oracle. initial = torch.randn(1, 1, args[0].shape[-1], args[2].shape[-1], device="cuda") - for prefix in (67, 65): - actual, final = forward(prefill_lengths=[prefix], initial_state=initial) - q, k, v, raw_g, raw_beta, rate, bias = [x.float() for x in args] + # Reuse the compiled 65-token prefix for full-prefill and suffix cases. + for tokens in (65, 67): + inputs = [x[:, :tokens] if i < 5 else x for i, x in enumerate(args)] + actual, final = forward(inputs=inputs, initial_state=initial) + q, k, v, raw_g, raw_beta, rate, bias = [x.float() for x in inputs] q, k = [x / (x.square().sum(-1, keepdim=True) + 1e-6).sqrt() for x in (q, k)] - q[:, :prefix] = q[:, :prefix].bfloat16().float() - k[:, :prefix] = k[:, :prefix].bfloat16().float() + q[:, :65] = q[:, :65].bfloat16().float() + k[:, :65] = k[:, :65].bfloat16().float() gate = -rate.exp().reshape( (1, 1, -1, 1) if raw_g.ndim == 4 else (1, 1, -1) ) * torch.nn.functional.softplus(raw_g + bias.reshape(raw_g.shape[2:])) @@ -240,7 +242,7 @@ def test_serving_state_qdq_and_handoff_gradient(compiled_serving_case): @pytest.fixture(scope="module") @torch.no_grad() def compiled_gdn_layout_case(): - """Compile a two-head case so strided head/value axes are observable.""" + """Compile two-head decode so strided head/value axes are observable.""" torch.manual_seed(732) q, k, v = [torch.randn(1, 2, 2, 32, device="cuda", dtype=torch.bfloat16) for _ in range(3)] raw_g, raw_beta = [torch.randn(1, 2, 2, device="cuda", dtype=torch.bfloat16) for _ in range(2)] @@ -257,7 +259,7 @@ def forward(values, gate_inputs): raw_beta.float().sigmoid(), policy=LinearAttentionConfig(backend="serving"), initial_state=initial, - prefill_lengths=[1], + prefill_lengths=[0], output_final_state=True, use_qk_l2norm_in_kernel=True, gate_inputs=gate_inputs, @@ -376,18 +378,15 @@ def test_padding_does_not_extend_decode(compiled_serving_case, packed): if not packed: with torch.no_grad(): - prefill, prefill_state = forward( - inputs=trimmed, prefill_lengths=[valid], initial_state=initial[:1] - ) + prefix_inputs = [x[:, :65] if i < 5 else x for i, x in enumerate(trimmed)] + prefill, prefill_state = forward(inputs=prefix_inputs, initial_state=initial[:1]) padded_prefill, padded_state = forward( inputs=padded, - prefill_lengths=[valid, 0], - sequence_lengths=[valid, 0], + prefill_lengths=[65, 0], + sequence_lengths=[65, 0], initial_state=initial, ) - torch.testing.assert_close(padded_prefill[:1, :valid], prefill, rtol=0, atol=0) + torch.testing.assert_close(padded_prefill[:1, :65], prefill, rtol=0, atol=0) torch.testing.assert_close(padded_state[:1], prefill_state, rtol=0, atol=0) torch.testing.assert_close(padded_state[1], initial[1], rtol=0, atol=0) - assert ( - not padded_prefill[0, valid:].count_nonzero() and not padded_prefill[1].count_nonzero() - ) + assert not padded_prefill[0, 65:].count_nonzero() and not padded_prefill[1].count_nonzero() diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_autograd.py b/tests/gpu_vllm/torch/quantization/test_vllm_autograd.py new file mode 100644 index 00000000000..0a67fe7da97 --- /dev/null +++ b/tests/gpu_vllm/torch/quantization/test_vllm_autograd.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +from _test_utils.torch.quantization.linear_attention_reference import recurrent_delta_rule_reference + +from modelopt.torch.quantization.linear_attention._vllm_autograd import prefix, step +from modelopt.torch.quantization.linear_attention.utils import forward_value + +vllm = pytest.importorskip("vllm") +if vllm.__version__ != "0.30.0": + pytest.skip("State QAT profiles target vLLM 0.30.0", allow_module_level=True) + + +@pytest.fixture(scope="module", params=[(False, True), (True, True), (False, False)]) +def compiled_decode(request): + """Compile the existing one-head decode shape outside the test timer.""" + channel, normalize = request.param + torch.manual_seed(97) + q, k, v = [torch.randn(1, 32, device="cuda").bfloat16().float() for _ in range(3)] + raw_gate = torch.randn((1, 32) if channel else (1,), device="cuda").bfloat16().float() + raw_beta = torch.randn(1, device="cuda").bfloat16().float() + rate, bias = torch.full((1,), -3.0, device="cuda"), torch.zeros_like(raw_gate) + initial = torch.randn(1, 32, 32, device="cuda") + inputs = [x.requires_grad_() for x in (q, k, v, raw_gate, raw_beta, rate, bias, initial)] + gate = -rate.exp() * torch.nn.functional.softplus(raw_gate + bias) + operands = q, k, v, gate, raw_beta.sigmoid(), initial + kwargs = {"scale": 0.3, "normalize": normalize, "gate_inputs": inputs[3:7]} + with torch.no_grad(): + native = step(*operands, **kwargs) + torch.cuda.synchronize() + return inputs, operands, kwargs, native + + +@pytest.mark.parametrize("loss_on", ["both", "output", "state"]) +def test_decode_backward_matches_recurrence(compiled_decode, loss_on): + inputs, operands, kwargs, native = compiled_decode + output, final = step(*operands, **kwargs) + for actual, expected in zip((output, final), native): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + q, k, v, gate, beta, initial = operands + if kwargs["normalize"]: + q, k = [x / (x.square().sum(-1, keepdim=True) + 1e-6).sqrt() for x in (q, k)] + _, reference_state = recurrent_delta_rule_reference( + *[x[None, None] for x in (q, k, v, gate, beta)], + initial_state=initial[None], + scale=kwargs["scale"], + ) + # The former Torch adjoint evaluated the readout at the native working state. + working = forward_value(reference_state[0], native[1]) + reference_output = torch.einsum("hk,hkv->hv", q * kwargs["scale"], working) + probes = [torch.randn_like(x) for x in (output, final)] + active = (loss_on != "state", loss_on != "output") + losses = [ + sum((value * probe).sum() for value, probe, use in zip(values, probes, active) if use) + for values in ((output, final), (reference_output, working)) + ] + gradients = [ + torch.autograd.grad(loss, inputs, allow_unused=True, retain_graph=True) for loss in losses + ] + for actual, expected in zip(*gradients): + if actual is None or expected is None: + assert actual is expected + else: + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=2e-6) + + +@pytest.fixture(scope="module", params=[False, True]) +def compiled_prefix(request): + """Reuse the integration tests' 65-token prefix shape for both recurrence types.""" + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + from fla.ops.kda import chunk_kda + + channel = request.param + torch.manual_seed(98) + inputs = [torch.randn(65, 1, 32, device="cuda").bfloat16().float() for _ in range(3)] + inputs += [ + torch.randn((65, 1, 32) if channel else (65, 1), device="cuda").bfloat16().float(), + torch.randn(65, 1, device="cuda").bfloat16().float(), + torch.full((1,), -3.0, device="cuda"), + torch.zeros(32 if channel else 1, device="cuda"), + torch.randn(1, 32, 32, device="cuda"), + ] + inputs = [x.requires_grad_() for x in inputs] + q, k, v, raw_gate, raw_beta, rate, bias, state = inputs + gate = -rate.exp() * torch.nn.functional.softplus(raw_gate + bias) + beta = raw_beta.sigmoid() + kwargs = { + "scale": 0.3, + "beta_dtype": torch.float32, + "normalize": True, + "gate_inputs": (raw_gate, raw_beta, rate, bias), + } + operands = q, k, v, gate, beta, state + actual = prefix(*operands, **kwargs) + torch.autograd.grad(sum(x.sum() for x in actual), inputs, retain_graph=True) + chunk = chunk_kda if channel else chunk_gated_delta_rule + expected = chunk( + *[x[None].bfloat16() for x in (q, k, v)], + gate[None], + beta[None], + initial_state=state[None], + output_final_state=True, + scale=0.3, + use_qk_l2norm_in_kernel=True, + ) + torch.cuda.synchronize() + return inputs, operands, kwargs, tuple(x[0].float() for x in expected) + + +@pytest.mark.parametrize("loss_on", ["both", "output", "state"]) +def test_prefix_backward_matches_fla(compiled_prefix, loss_on): + inputs, operands, kwargs, expected = compiled_prefix + actual = prefix(*operands, **kwargs) + with torch.no_grad(): + native = prefix(*operands, **kwargs) + for value, reference in zip(actual, native): + torch.testing.assert_close(value, reference, rtol=0, atol=0) + probes = [torch.randn_like(x) for x in actual] + active = (loss_on != "state", loss_on != "output") + gradients = [ + torch.autograd.grad( + sum((x * p).sum() for x, p, use in zip(values, probes, active) if use), + inputs, + retain_graph=True, + ) + for values in (actual, expected) + ] + for value, reference in zip(*gradients): + torch.testing.assert_close(value, reference, rtol=1e-5, atol=2e-6) From 8cbec76f1134b67f064ee48fb54bc12c2a8f2809 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 14 Oct 2026 12:53:15 -0700 Subject: [PATCH 26/27] Expose FP8 state block sizes through TensorQuantizer Signed-off-by: Kai Xu --- .../linear_attention/serving/__init__.py | 6 ++++- .../quantization/linear_attention/utils.py | 9 +++----- modelopt_recipes/ptq.md | 23 +++++++++++++++++-- .../test_linear_attention_decode.py | 18 ++++++++++++++- 4 files changed, 46 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py index de4e3c23c18..eb9d4dfd539 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py @@ -13,7 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Optional vLLM forward kernels for linear-attention serving precision profiles.""" +"""Training adapter for native vLLM linear-attention forward kernels. + +ModelOpt owns the training state, QDQ and backward. This package imports the +selected serving arithmetic; it does not modify the vLLM engine or scheduler. +""" try: import vllm diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py index f539919d472..1f6f753fe4b 100644 --- a/modelopt/torch/quantization/linear_attention/utils.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -95,10 +95,7 @@ def validate_gdn_quantizer( and (quantizer.num_bits != 8 or (not quantizer.unsigned and quantizer.narrow_range)) and ( quantizer.block_sizes is None - or ( - quantizer.num_bits == 8 - and quantizer.block_sizes in ({-1: size} for size in block_sizes) - ) + or (quantizer.block_sizes in ({-1: size} for size in block_sizes)) ) and quantizer.fake_quant and quantizer._pass_through_bwd @@ -110,7 +107,7 @@ def validate_gdn_quantizer( ): raise ValueError( f"{name} supports only dynamic fake quantization with num_bits in {num_bits}, " - f"INT8 block_sizes in {block_sizes} (or no blocks), " + f"block_sizes in {block_sizes} (or no blocks), " "pass_through_bwd=True, no rotation, pre-scaling, bias, constant " "amax, or custom backend. Other gradient rules and formats are not implemented." ) @@ -124,7 +121,7 @@ def state_quantizer_config( A zero group size retains the legacy policy's full-key state tiles. """ validate_gdn_quantizer( - quantizer, name=name, num_bits=tuple(_STATE_FORMATS), block_sizes=(16, 32, 64) + quantizer, name=name, num_bits=tuple(_STATE_FORMATS), block_sizes=(16, 32, 64, 128) ) if quantizer.block_sizes is not None: return _STATE_FORMATS[quantizer.num_bits], quantizer.block_sizes[-1] diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 845c1643bc1..a06c80b75ee 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -148,8 +148,27 @@ The tests qualify the selected native kernels, not arbitrary server dispatch. Packed boundaries must match the layer's actual (including padding) boundaries. Validation reads their values rather than trusting tensor versions, so it can synchronize -CUDA metadata. Prefix backward still runs a separate FLA recurrence as an adjoint; -this implementation is not an optimized fused training kernel. +CUDA metadata. Forward runs one native vLLM prefix and one native decode per +suffix token. Prefix backward reconstructs intermediates and calls upstream FLA +backward; decode uses an explicit recurrence backward. QDQ uses an identity STE. +The `linear_attention/serving` package is a training adapter for these kernels; +it does not integrate with the vLLM engine or scheduler. + +State FP8 block size uses the existing TensorQuantizer settings in `quant_cfg`: + +```python +state_fp8 = { + "num_bits": (4, 3), + "type": "dynamic", + "block_sizes": {-1: 32}, # 16, 32, 64 or 128 value channels per key row + "pass_through_bwd": True, +} +cfg["quant_cfg"].append({"quantizer_name": "*state_quantizer", "cfg": state_fp8}) +``` + +These blocks each have their own dynamic scale. The legacy `axis=(0, 1)` setting +instead shares a scale across a full-key tile; `state_block_v` only controls that +legacy tile's width. NVFP4 state support is a separate extension. --- diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py index b8efe8a597c..f43ad30f2a4 100644 --- a/tests/unit/torch/quantization/test_linear_attention_decode.py +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -19,7 +19,7 @@ from modelopt.torch.quantization.config import QuantizeConfig, QuantizerAttributeConfig from modelopt.torch.quantization.linear_attention import LinearAttentionConfig, kda_state_qat from modelopt.torch.quantization.linear_attention import training as state_training -from modelopt.torch.quantization.linear_attention.utils import _state_qdq +from modelopt.torch.quantization.linear_attention.utils import _state_qdq, state_quantizer_config from modelopt.torch.quantization.nn import TensorQuantizer @@ -42,6 +42,22 @@ def test_state_qdq_matches_tensor_quantizer(state_format): torch.testing.assert_close(gradient, probe, rtol=0, atol=0) +@pytest.mark.parametrize("block_size", [16, 32, 64, 128]) +def test_fp8_state_blocks_match_tensor_quantizer(block_size): + value = torch.randn(1, 2, 3, 133, requires_grad=True) + quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=(4, 3), type="dynamic", block_sizes={-1: block_size}, pass_through_bwd=True + ) + ) + assert state_quantizer_config(quantizer) == ("fp8_e4m3", block_size) + actual = _state_qdq(value, quantizer) + torch.testing.assert_close(actual, quantizer(value), rtol=0, atol=0) + probe = torch.randn_like(value) + (gradient,) = torch.autograd.grad((actual * probe).sum(), value) + torch.testing.assert_close(gradient, probe, rtol=0, atol=0) + + @pytest.mark.parametrize( "settings", [ From 662fd69622dfb3a52d714c84d32ec1bed539b4c8 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 14 Oct 2026 13:10:24 -0700 Subject: [PATCH 27/27] Reuse the integration scale for prefix backward compilation Signed-off-by: Kai Xu --- tests/gpu_vllm/torch/quantization/test_vllm_autograd.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_autograd.py b/tests/gpu_vllm/torch/quantization/test_vllm_autograd.py index 0a67fe7da97..9474f5d49b9 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_autograd.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_autograd.py @@ -100,7 +100,7 @@ def compiled_prefix(request): gate = -rate.exp() * torch.nn.functional.softplus(raw_gate + bias) beta = raw_beta.sigmoid() kwargs = { - "scale": 0.3, + "scale": 32**-0.5, "beta_dtype": torch.float32, "normalize": True, "gate_inputs": (raw_gate, raw_beta, rate, bias), @@ -115,7 +115,7 @@ def compiled_prefix(request): beta[None], initial_state=state[None], output_final_state=True, - scale=0.3, + scale=kwargs["scale"], use_qk_l2norm_in_kernel=True, ) torch.cuda.synchronize()