diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index c38e462912a..f0c5a38bbab 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.github/workflows/gpu_tests.yml @@ -44,7 +44,8 @@ jobs: matrix: include: - example: gpu - timeout: 60 + # Includes dependency builds and cold compilation of the FLA forward/backward tests. + timeout: 75 # Pinned to 26.05: benchmark.py uses trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH, # which newer TensorRT (26.06) removed. Bump once the source is updated for TensorRT 10. container_image: nvcr.io/nvidia/pytorch:26.05-py3 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7654acab550..e677fb2dff2 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -108,6 +108,8 @@ repos: modelopt/onnx/quantization/ort_patching.py| modelopt/torch/_deploy/utils/onnx_utils.py| modelopt/torch/export/transformer_engine.py| + modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py| + modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py| modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_pruned_to_mxfp4.py| modelopt/torch/quantization/export_onnx.py| modelopt/torch/quantization/plugins/attention.py| diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 22169f974c0..147e9a88632 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,10 +13,19 @@ Changelog *Quantization* +- Add experimental FP8/INT8 recurrent-state fake quantization before native vLLM GDN/KDA prefill and decode. Use the state-only recipe with the documented vLLM 0.15.x eager runtime. + +- 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 KDA prefill QAT for FLA 0.5.1 layers with stable per-channel decay and independent operand/state fake quantization. Enable the explicit ``matmul`` backend; numerical policies and quantizers persist through ModelOpt save/restore. + +- Add experimental per-operand FP8/NVFP4 GDN prefill QAT with independently configurable accumulator and elementwise rounding. Select the materialized ``matmul`` backend to enable these sites; the existing fused FLA state/W path remains the default. + - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration. - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md `_ for details. - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``. +- Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training through the standard ``quant_cfg`` interface. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer. *Megatron Framework (M-LM / M-Bridge)* diff --git a/LICENSE b/LICENSE index 57c1104ae67..e08eaefc016 100644 --- a/LICENSE +++ b/LICENSE @@ -251,6 +251,7 @@ the following copyright holders, licensed under the MIT License: Copyright (c) 2025 sgl-project Copyright (c) 2026 The DeepSpec Authors Copyright (c) 2023-2026 The ggml authors + Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/docs/source/_templates/autosummary/module.rst b/docs/source/_templates/autosummary/module.rst index 36eb087fbec..c693958c63d 100644 --- a/docs/source/_templates/autosummary/module.rst +++ b/docs/source/_templates/autosummary/module.rst @@ -11,7 +11,7 @@ :recursive: {% for item in modules %} {% set full_item = fullname + '.' + item.split('.')[-1] %} -{% if ('.plugins.' not in full_item or full_item == 'modelopt.torch.opt.plugins.huggingface') and full_item != 'modelopt.torch.quantization.backends.fp8_per_tensor_gemm' %} +{% if ('.plugins.' not in full_item or full_item == 'modelopt.torch.opt.plugins.huggingface') and full_item != 'modelopt.torch.quantization.backends.fp8_per_tensor_gemm' and not full_item.startswith('modelopt.torch.kernels.quantization.linear_attention.fla_') %} {{ full_item }} {% endif %} {%- endfor %} diff --git a/examples/llm_qat/README.md b/examples/llm_qat/README.md index b4900ebea1f..aa099d88ebd 100644 --- a/examples/llm_qat/README.md +++ b/examples/llm_qat/README.md @@ -14,6 +14,7 @@ For background on how QAT enables low-precision accuracy recovery, see the [QAT/ | Background | How QAT/QAD work and when to use each | \[[Link](#background)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | | Support Matrix | Supported models, quantization formats, and backends | \[[Link](#support-matrix)\] | | | QLoRA | Model training with reduced GPU memory | \[[Link](#qlora-real-quantization)\] | | +| Linear Attention | GDN/KDA recurrent-state QAT and ReplaySSM example | \[[Link](linear_attention/README.md)\] | | | Advanced Topics | FSDP2 config, YAML options | \[[Link](#advanced-topics)\] | | | Results | Accuracy benchmarks | \[[Link](#results)\] | | | Resources | Extra links and references | \[[Link](#resources)\] | | diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md new file mode 100644 index 00000000000..9b6110deca8 --- /dev/null +++ b/examples/llm_qat/linear_attention/README.md @@ -0,0 +1,721 @@ +# Quantization-Aware Training for Linear Attention + +This example fine-tunes linear-attention models with differentiable recurrent-state +fake quantization. It supports GDN/KDA token-state writes, KDA decay approximation, +and ReplaySSM. The default INT8 token-state recipe enables Hadamard rotation. +The included model-level training script uses FLA `KimiDeltaAttention` layers. + +## Prerequisites + +Install ModelOpt from this checkout using the [QAT setup instructions](../README.md#quick-start), +then install the example dependencies. Run the commands below from the repository root. + +```bash +pip install -r examples/llm_qat/linear_attention/requirements.txt +``` + +The example requires a CUDA GPU, `fla-core==0.5.1`, and +`flash-linear-attention==0.5.1`. Stage a local model/tokenizer snapshot and separate +training and evaluation Parquet files with a `text` column. The pinned KDA model +and WikiText revisions are recorded in [decode_study_plan.json](decode_study_plan.json); +the script defaults to those revision identifiers, so the local inputs must match. + +## Run a QAT example + +The following command runs one training step with INT8 recurrent-state QDQ +in a 32-value Hadamard basis. Value dimensions must be divisible by 32. +Each sequence has 64 exact prefix tokens; the remaining tokens use the decode +policy. This recipe enables state QDQ only in the decode suffix; prefix state +QDQ is disabled by default. Loss is computed on suffix labels. + +```bash +python examples/llm_qat/linear_attention/train.py \ + --model /path/to/model-snapshot \ + --train-data /path/to/train.parquet \ + --eval-data /path/to/validation.parquet \ + --quant-config examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json \ + --prefill-tokens 64 \ + --length 128 \ + --train-steps 1 \ + --eval-blocks 4 \ + --output tmp/linear_attention/state_int8.json +``` + +Add `--trust-remote-code` if the selected checkpoint requires its custom model code. +The script trains attention parameters with FP32 AdamW, BF16 autocast, and activation +checkpointing. Other model parameters remain frozen. It writes pre/post-training +NLL, gradient norms, timings, and input/source hashes to JSON; it does not save a +trained checkpoint. This short run checks training integration. Use the study plan +for matched quality comparisons before making accuracy-recovery claims. + +### Select a recipe + +Pass one of the following files to `--quant-config`: + +| Recipe | Behavior | +| --- | --- | +| [kda_decode_exact.json](configs/kda_decode_exact.json) | Exact recurrence control with state QDQ disabled. | +| [kda_decode_state_int8.json](configs/kda_decode_state_int8.json) | INT8 + Hadamard state QDQ at handoff and after every token. | +| [kda_decode_state_fp8.json](configs/kda_decode_state_fp8.json) | FP8 E4M3 state QDQ after every token. | +| [kda_decode_replay_int8.json](configs/kda_decode_replay_int8.json) | Unrotated INT8 replay-anchor comparison with separate FP8 factor QDQ. | +| [kda_decode_replay_fp8.json](configs/kda_decode_replay_fp8.json) | Replay with FP8 anchors and factors. | +| [kda_decode_replay_int8_hadamard.json](configs/kda_decode_replay_int8_hadamard.json) | Replay with rotated INT8 anchors and factor QDQ disabled. | +| [kda_decode_decay_256.json](configs/kda_decode_decay_256.json), [1024](configs/kda_decode_decay_1024.json), [4096](configs/kda_decode_decay_4096.json) | FP8 state QDQ with different log-retention rounding grids. | + +For an exact control, repeat the training command with `kda_decode_exact.json` and +`--output tmp/linear_attention/exact.json`, keeping model, data, seed, phase lengths, +and training settings identical. Compare the resulting records with: + +```bash +python examples/llm_qat/linear_attention/compare_quality.py \ + --control tmp/linear_attention/exact.json \ + --candidate tmp/linear_attention/state_int8.json \ + --output tmp/linear_attention/comparison.json +``` + +The comparison reports paired block-NLL differences and bootstrap bounds. +INT8 and Hadamard model-quality qualification remains pending. + +## Enable state quantization + +### What the execution policy controls + +An **execution policy** is the layer's `LinearAttentionConfig`: the settings that +determine how the recurrence runs and when it applies quantization. Training +uses three separate inputs: + +| Input | Purpose | Supplied through | +| --- | --- | --- | +| `TensorQuantizer` settings | Enable a quantization site and choose its numerical format, scales, and gradient behavior | Recipe `quant_cfg`, such as dynamic INT8 with identity STE | +| Execution policy | Choose token or replay mode, Hadamard rotation, handoff quantization, and the execution backend | Recipe `linear_attention` entries | +| Batch metadata | Specify where each sequence switches from prefill to decode | `linear_attention_training_phase(model, prefill_lengths)` | + +The same INT8 quantizer settings can round the state after every token or round +a replay anchor every eight tokens. Those schedules produce different recurrent +states, so the policy must specify which computation training should emulate. + +`LinearAttentionConfig` holds the overall backend, chunk size, and state settings. +Its optional `decode` field contains a `LinearAttentionDecodeConfig` for token or +replay mode, Torch or Triton implementation, state codec, decay approximation, +and replay settings. This is one nested configuration: supplying a `decode` +dictionary in the recipe constructs the nested config automatically. + +Each `linear_attention` entry uses `module_name` to select attention layers and +`cfg` to specify their policy. During `mtq.quantize`, ModelOpt stores that policy +as each matched layer's `linear_attention_config`. Quantizer settings and the +policy persist through ModelOpt save/restore; batch-specific prefill lengths +must be supplied again for each workload. + +### Load a state recipe + +State quantizers start disabled. The state recipe enables `*kda_state_quantizer` +with signed narrow-range INT8, dynamic scales, and `axis=(0, 1)`. Use +`*gdn_state_quantizer` for GDN. Select the corresponding FP8 recipe for E4M3. +If changing the INT8 recipe manually, also set `decode.state_codec="tile"`. +Setting `prefill_state_qdq=True` alone does not enable a quantizer. + +Start with the token-state recipe and select one of the schedules below **before** +calling `mtq.quantize`: + +```python +import json +from pathlib import Path + +recipe = json.loads( + Path("examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json").read_text() +) +policy = recipe["linear_attention"][0]["cfg"] +decode = policy["decode"] +``` + +For GDN, the complete recipe composes the INT8 quantizer unit with the Hadamard +execution policy: + +```python +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe + +gdn_recipe = load_recipe("general/ptq/gdn_state_int8_dynamic").quantize +mtq.quantize(gdn_model, gdn_recipe) +``` + +Use the phase context shown below for each forward/backward. Importing only +`configs/ptq/units/gdn_state_int8_dynamic` configures the quantizer without +selecting Hadamard. The complete GDN recipe uses the Torch implementation; +set `gdn_recipe.linear_attention[0].cfg.decode.implementation="triton"` before +conversion to select the fused decode implementation. + +### Why training needs prefill/decode boundaries + +One training forward can simulate prompt prefill followed by recurrent decode. +The default 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 default INT8 token-state recipe uses +`"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`. +The FP8 token-state recipe also leaves prefix state QDQ disabled. +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; the supplied replay recipes enable +that separate factor quantization unless stated otherwise. + +`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. + +To use a modified recipe with `train.py`, save it as JSON and pass that file to +`--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. + +For pure-prefill GDN/KDA training with the materialized backend, use +`backend="matmul"` with no `decode` policy and enable the state quantizer. +That path also rounds initial and chunk-final state and needs no phase context; +`decode.prefill_state_qdq` applies only when a decode policy is present. + +## 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. + +## Numerical contract + +For activated input tensors, a token computes `D = exp(g)`, `R = D * S`, `u = beta * (v - k^T R)`, and `W = R + k u^T`. GDN uses one log retention per head; KDA uses one per key channel. The output reads either the working state `W` or the stored state, according to `readout`. All operations use FP32 working values (FP64 is supported by the Torch reference). + +- **Token mode:** fake-quantize `W` at every state write when the module state quantizer is enabled. The next token consumes that stored value. +- **Decay approximation:** an optional positive `decay_log_step` rounds log retention to the nearest grid point before exponentiation. Its derivative is an identity straight-through estimator (STE); no hidden clamping is applied. Prefix gates retain the prefill policy. +- **Replay mode:** maintain an anchor and ordered encoded rank-one updates. Encode each key per head and each update per value block, with optional FP8 E4M3 QDQ. Compute a new residual from the reconstructed current trajectory and the encoded key. At `window` updates, encode the current state as a new anchor and clear the buffer. Anchor QDQ follows the state quantizer toggle; factor QDQ has its own toggle. +- **Initial state:** `quantize_initial` applies once at the first nonempty decode handoff. Continuation and empty calls do not create extra writes. Prefix state QDQ is independently selected by `prefill_state_qdq`. + +With `state_codec="tile"`, state scale grouping is one dynamic scale for each `[Dk, block_v]` head tile. Update scales span one value block; key scales span the complete key vector. Zero tiles use scale one. Codec scales are detached, and rounded values have identity STE gradients. State codecs support E4M3 nearest-even rounding (including subnormals) and signed +symmetric INT8. INT8 uses FP32 `amax / 127` scales, zero point zero, codes in +`[-127, 127]`, nearest-even rounding, and saturation. Replay key/update factors +retain their independent E4M3 or identity format. Key reductions use a fixed pairwise FP32 tree; this avoids changing recurrence rounding merely by changing the reduction schedule. + +The public numerical APIs `recurrent_decode_reference` and `recurrent_decode` accept `state_format="fp8_e4m3"` (default) or `"int8"` and one aligned `[T,H,D]` sequence and return outputs plus a `LinearAttentionCarry`. It retains anchor values/scales, buffered key/update values/scales, log retentions, position, cursor, and a policy signature. Returned tensors preserve gradients through prefix handoff, refreshes, and continuation. Restart with a new carry to reset; incompatible policy, state format, or state shape is rejected. + +The Torch reference can explicitly re-encode stored entries using `encoding="reencode"`. It re-encodes the already stored values and does not regenerate factors from a teacher trajectory. The fused implementation supports encode-once replay. Carrying its reconstructed state incrementally applies the same ordered transitions as replaying the fixed entries; it does not measure serving-time reconstruction cost. + +## INT8 state with Hadamard rotation + +Set `decode.state_codec="int8_hadamard32"` with an INT8 state quantizer to enable +32-point orthonormal Sylvester Hadamard transforms along the **value** dimension. +This codec supports both token writes and replay anchors. The default INT8 +token-state recipe and complete GDN INT8 recipe select it. The explicit `"tile"` +codec retains the unrotated FP8/INT8 behavior above. + +```python +# Execution policy; pair with the INT8 quantizer attributes in the example above. +policy = { + "backend": "matmul", + "state": {"block_v": 64}, + "decode": { + "mode": "replay", + "implementation": "triton", + "state_codec": "int8_hadamard32", + "readout": "working", + "replay": {"window": 8, "factor_qdq": False, "encoding": "once"}, + }, +} +``` + +For key-first state `S`, the internal state is `S @ H`, with a block-diagonal +Hadamard matrix `H`. Inputs `v` and replay update vectors use the same basis; +queries, keys, and scalar/channel decay gates keep their original coordinates. +Outputs are transformed back. The first nonempty decode call transforms the +incoming state once; continuation retains the rotated anchor and updates without +another rotation or initial-state quantization. `carry.value_basis` records the +basis. `carry.reconstruct()` returns the **original** basis; passing +`original_basis=False` exposes the internal basis. + +The state codec uses one scale per **key channel and 32 contiguous values**, even +when the execution tile `state.block_v` is 64 or 128. For each group, compute +`scale=max(amax/127, 6e-8)` in FP32, divide by that scale, round half ties away from +zero, and saturate codes to `[-127,127]`. Store the scale in FP16 and reconstruct +with that FP16 scale. Scale metadata has shape `[H,Dk,Dv/32]`. Rounding has +identity STE and scales are detached. Basis transforms use their actual linear +derivatives. Replay factor QDQ remains independently configurable; the example +disables it to isolate state/anchor quantization. + +The Hadamard basis and INT8 scale/rounding conventions follow +[quantized-replayssm at 29c35508](https://github.com/mxinO/quantized-replayssm/blob/29c355086070d2c7973550d3670171af53433533/vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py). +Model-quality evaluation of this ModelOpt implementation remains pending. + +`Dv` must be divisible by 32, and `state.block_v` must be 32, 64, or 128. Prefix +computation remains in the original basis with the exact solve; this codec starts +at decode handoff and rejects `prefill_state_qdq=True`. An all-prefix or empty +sequence does not create a decode quantization event. The policy is saved and +restored with the model. A complete example is +[configs/kda_decode_replay_int8_hadamard.json](configs/kda_decode_replay_int8_hadamard.json). + +## Training implementation and limits + +The CUDA Triton implementation supports FP32 working inputs, key dimensions up to 128, scalar/channel gates, and value blocks 16/32/64/128 with masked tails. It saves a state every eight tokens and recomputes intermediate states in backward. This internal checkpoint interval does not change write cadence. Its custom backward supports first-order gradients; use the Torch reference for higher-order differentiation. + +Direct prefill APIs accept `prefill_lengths`, grouped query/key heads, both state layouts, packed sequences, empty entries, and full/zero prefixes. Positive prefixes are evaluated together so prefill quantizer scale domains remain shared. FLA KDA layer integration requires 0.5.1 and `use_cache=False`; serving cache objects are rejected. The numerical carry API provides explicit continuation, but this change does not implement a serving cache or export format. Megatron GDN tensor parallelism is covered by dedicated tests; context parallelism remains rejected. + +All encoded values remain floating tensors. This is training-time fake quantization and approximation, with no compressed storage or inference speed claim. Use [benchmark_decode.py](benchmark_decode.py) for matched-policy training overhead and `decode_study_plan.json` for the fixed model-quality protocol. The exact inverse remains selected in these studies; the earlier rejected Neumann candidate is not enabled. + +See [the qualification study](decode_study.md) for the fixed +model/data protocol, quality comparisons, and measured training overhead. + +## GDN prefill arithmetic for QAT + +GDN training can expose all eight logical prefill matmuls through an explicit +`backend="matmul"` policy. Each operand uses ModelOpt `TensorQuantizer`; standard +PyTorch autograd propagates its identity straight-through gradient. The computation +batches chunk-local work and carries state between chunks without detaching it. +This backend materializes operands and the triangular inverse. Its training cost +must be measured separately from the fused FLA state/W path. + +### Enable operand quantization + +```python +import modelopt.torch.quantization as mtq + +model = mtq.quantize(model, { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*linear_attn_sites.*", + "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)}}, + {"quantizer_name": "*gdn_w_quantizer", + "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)}}, + ], + "algorithm": None, + "linear_attention": [{ + "module_name": "decoder.layers.*.self_attention", + "cfg": {"backend": "matmul"}, + }], +}) +``` + +Each site has `lhs_quantizer` and `rhs_quantizer`. The sole exception is +`state_read.lhs_quantizer`: use the existing `gdn_w_quantizer` handle. It is not +registered twice. Standard projection recipes keep all these handles disabled. + +Both operands are presented as four-dimensional row vectors: `[chunk_or_sequence, +head, row, reduction]`. The right operand is transposed after quantization. +E4M3 uses one dynamic scale per row. The exact transformed operands are: + +| Site | Left operand | Right operand before transpose | +| --- | --- | --- | +| `key_interaction` | beta-scaled keys | keys | +| `wy_value` | triangular inverse | transposed beta-scaled values | +| `wy_key` | triangular inverse | transposed beta- and prefix-decay-scaled keys | +| `state_read` | WY keys (`gdn_w_quantizer`) | transposed incoming state | +| `state_update` | transposed end-decayed keys | transposed corrected values | +| `output_state` | scaled, prefix-decayed queries | transposed incoming state | +| `output_score` | scaled queries | keys | +| `output_value` | causal, decayed local scores | transposed corrected values | + +Stored-state rounding remains separate from temporary state-operand quantization. +Enable `gdn_state_quantizer` separately to round the initial state and chunk writes. + +For NVFP4 operands, replace the operand configuration with: + +```python +{ + "num_bits": (2, 1), + "type": "dynamic", + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, +} +``` + +This uses E2M1 values, E4M3 block scales, and a dynamically computed tensor amax +for each quantizer invocation. Chunk-local sites batch all sequence chunks in one +invocation. State-dependent sites invoke once per chunk over the sequence batch. +These scale domains are part of the numerical contract. Chunk padding is zero; +state carries for completed packed sequences do not incur additional state writes. +NVFP4 support here covers operands, including WY keys, and does not change the +FP8 stored-state codec. Static scales and other quantizer transformations are rejected. + +### Select arithmetic experiments independently + +The optional `matmul` and `elementwise` policy fields operate even when all +quantizers are disabled: + +```python +policy = { + "backend": "matmul", + "matmul": { + "output_value": {"accumulator_dtype": "float16", "reduction_block": 16}, + }, + "elementwise": {"value_residual": "bfloat16"}, +} +``` + +A matmul schedule computes partial products in the working dtype, traverses the +reduction dimension left to right in `reduction_block` elements, and rounds the +accumulator after each addition. Both schedule fields are required together. +`float32`, `float16`, and `bfloat16` are supported rounding formats. This schedule +is a numerical experiment; it does not assert instruction-level hardware behavior. + +Elementwise sites are `gate_prefix` (after cumulative summation), `gate_exp` +(after each exponential), `value_residual`, `state_decay`, `state_add`, and +`output_add`. Rounding applies only at the named boundary. These operations and +operand QDQ use identity STE, preserving baseline backward arithmetic. + +Working arithmetic is FP32 for FP16/BF16/FP32 input and FP64 for double input, +including inside an outer training-autocast context. +Matmuls respect PyTorch's global precision setting. Set +`torch.set_float32_matmul_precision("highest")` to exclude TF32 from comparisons; +the kernel tests and benchmark do this explicitly. +Outputs are cast to the query dtype; carried state keeps the working dtype. The +triangular solve remains exact. No gate clamps or automatic inverse approximation +are enabled. Gate activation and Q/K normalization follow FLA's formulas. + +### Persistence and verification + +The policy and enabled handles survive ModelOpt save/restore. Checkpoints from +before these new handles existed restore with the new handles disabled. Context +parallelism is rejected; qualification is limited to unchanged TP topology and +PP=1. All-disabled policies preserve the original framework kernel. + +The numerical tests compare outputs, final states, and every differentiable input +gradient, including raw gate parameters, and exercise packed tails and state layout. +FP8 single-site cases use elementwise comparisons. Composed FP32 state/operand QDQ +uses the established GDN relative-L2 bounds (1% outputs/states, 2% gradients): +batched and unbatched GEMM roundoff can cross a later rounding threshold. NVFP4 +checks the codec independently at each actual operand shape and compares the full +computation and gradients with the chunk oracle using the declared scale domains. +Activation-checkpoint recomputation must reproduce the same gradients. + +Measure forward/backward time and peak extra tensor memory with: + +```bash +PYTHONPATH=. python examples/llm_qat/linear_attention/benchmark_prefill.py \ + --batch 1 --length 1024 --heads 4 --dim 64 --repeats 20 \ + --output prefill-overhead.json +``` + +The script warms each variant, interleaves trial order, and reports paired wall-time +ratios against the exact FLA path. Its JSON records the GPU, package versions, shape, +source hashes, and all samples. These are training-emulation costs; no native +low-precision inference speedup or model-quality recovery is implied. + +## KDA prefill arithmetic for QAT + +ModelOpt recognizes FLA 0.5.1 `KimiDeltaAttention` layers during conversion and adds +`kda_state_quantizer`, `kda_w_quantizer`, and the eight `linear_attn_sites` operand +pairs. These handles start disabled. The existing projection recipes leave them +disabled; enable them explicitly with `backend="matmul"`. + +The adapter reuses the layer's projections, convolution, normalization, and gate +parameters. Its numerical path uses the materialized differentiable prefill +backend. It does not patch process-wide FLA functions. When all numerical policies +are off, the original FLA forward runs unchanged. + +### KDA site map and decay stability + +For a chunk, let `G[i,d]` be the cumulative natural-log retention for key channel +`d`. KDA's interaction terms contain `exp(G[i,d] - G[j,d])`. The implementation +computes that causal difference directly. It never constructs `exp(-G[j,d])`; +strong forgetting can underflow a retention to zero without introducing the +infinite inverse factor of an unstable factorization. Future positions are +masked before the exponential and zeroed before operand quantization. + +The shared prefill sites have the following KDA operands. Right operands below +are the row vectors presented to the quantizer before transposition. + +| Site | Left operand | Right operand | +| --- | --- | --- | +| `key_interaction` | `beta[i] * k[i]` | `k[j] * exp(G[i] - G[j])`, `j <= i` | +| `wy_value` | exact triangular inverse | transposed beta-scaled values | +| `wy_key` | exact triangular inverse | transposed beta-scaled, prefix-decayed keys | +| `state_read` | WY keys, via `kda_w_quantizer` | transposed incoming state | +| `state_update` | transposed end-decayed keys | transposed corrected values | +| `output_state` | scaled, prefix-decayed queries | transposed incoming state | +| `output_score` | scaled `q[i]` | `k[j] * exp(G[i] - G[j])`, `j <= i` | +| `output_value` | causal local scores | transposed corrected values | + +For the two interaction sites, query rows are grouped in fixed tiles of eight +and folded into the batch dimension. Quantizer inputs remain four-dimensional: +`[chunk_or_sequence_times_query_rows, head, row, reduction]`. FP8 scales are +per row. NVFP4 uses block-16 E2M1 values and E4M3 scales with one dynamic tensor +amax per invocation; the query tile of eight is therefore part of its scale +contract. Other sites use the [GDN prefill scale domains](#gdn-prefill-arithmetic-for-qat). + +Chunk size is 64. Working arithmetic is FP32, including inside an outer BF16 +training-autocast context; double inputs use FP64. Output returns to the query +dtype. Select `torch.set_float32_matmul_precision("highest")` to exclude TF32. +Operand and arithmetic rounding use identity STE; state carry is differentiable. +Exact triangular solve is the only supported solve in this delivery. + +Raw gate activation follows the loaded FLA model: without `lower_bound`, it is +`-exp(A_log) * softplus(raw_gate + dt_bias)`; with a lower bound, it is +`lower_bound * sigmoid(exp(A_log) * (raw_gate + dt_bias))`. The latter is the +model's bounded gate formula, not a new clamp. No additional decay approximation +is enabled by this integration. The materialized path does not emulate the +optional FLA `safe_gate` TensorCore rounding. Numerical parity checks use FLA's +general computation (`safe_gate=False`) for both gate formulas. + +### Enable KDA operand QDQ + +```python +import modelopt.torch.quantization as mtq + +model = mtq.quantize(model, { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*linear_attn_sites.*", + "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)}}, + {"quantizer_name": "*kda_w_quantizer", + "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)}}, + ], + "algorithm": None, + "linear_attention": [{ + "module_name": "model.layers.*.self_attn", + "cfg": {"backend": "matmul"}, + }], +}) +``` + +To emulate stored-state writes, enable `*kda_state_quantizer` with dynamic E4M3 +or signed narrow-range INT8 and `axis=(0,1)`. The policy's `state.block_v` specifies +the value-column tile; initial state and every chunk write are rounded. This is distinct from temporary +state-read operand QDQ. State-V-first layout, grouped value heads, packed tails, +and gradients through the initial state are supported. + +`flash-linear-attention==0.5.1` and `fla-core==0.5.1` are required for the layer +adapter. Importing ModelOpt itself does not require either package. Qualification +covers a single-GPU FLA model. CP, FLA-specific intermediate/recompute flags, and +quantized recurrent inference are rejected. In evaluation mode FLA selects its +recurrent path for lengths at most 64; use longer prefill sequences for this +integration. For token-state training and decode/replay policies, see +[decode-aware QAT](#numerical-contract). + +### Reproducible model study + +The [study example](#prefill-model-study) uses a public, +pinned Arcee KDA checkpoint and WikiText-2 splits. It reports pre/post-training +held-out NLL and perplexity with per-block results, hashes, seeds, and fixed token +budgets. A successful training step verifies integration; it is not evidence of +model-quality recovery. The materialized backend's time and memory must be +measured separately from fused FLA. + +## Prefill model study + +The initial model is [arcee-ai/AFM-4.5B-Base-KDA-Only](https://huggingface.co/arcee-ai/AFM-4.5B-Base-KDA-Only) +at revision `01ad2e06ee4f1214193c17b69e09105a9b257e80`. Its implementation uses +FLA's `KimiDeltaAttention` directly. The dataset is +[Salesforce/wikitext](https://huggingface.co/datasets/Salesforce/wikitext), configuration +`wikitext-2-raw-v1`, revision `b08601e04326c79dfdd32d625aee71d232d685c3`. + +Stage those exact revisions with model +weights, tokenizer, audited model code, and the three Parquet splits. Supply local +paths to the runner; it does not fetch weights or data. The checkpoint's custom +model code requires an explicit `--trust-remote-code` opt-in. Review its pinned +Python files before enabling that flag. + +```bash +PYTHONPATH=. python examples/llm_qat/linear_attention/train.py \ + --model /path/to/arcee-kda \ + --train-data /path/to/wikitext-2-raw-v1/train-00000-of-00001.parquet \ + --eval-data /path/to/wikitext-2-raw-v1/validation-00000-of-00001.parquet \ + --eval-split validation --trust-remote-code \ + --quant-config examples/llm_qat/linear_attention/configs/kda_prefill_fp8.json \ + --length 128 --train-steps 1 --eval-blocks 4 --seed 2026 \ + --output fp8-smoke.json +``` + +The one-step command is an integration smoke test. For matched studies, run the +same workload from the same original checkpoint with each numerical configuration: + +| Trial | Configuration | +| --- | --- | +| Exact fused FLA control | Omit `--quant-config` | +| Exact materialized control | `configs/kda_exact_matmul.json` | +| FP8 prefill operands | `configs/kda_prefill_fp8.json` | +| NVFP4 prefill operands | `configs/kda_prefill_nvfp4.json` | + +The exact materialized control explicitly rounds the already-FP32 `output_add` +to FP32 to select the emulation path. It introduces no lower-precision rounding. +First establish agreement with the fused control before interpreting a candidate's +quality. Candidate `before` measurements isolate its untrained numerical change; +compare trained candidates against an equally trained exact control. Choose +settings using validation data, then run the fixed selected settings on the test +split. Never choose approximation degrees or training hyperparameters on test data. + +### Protocol + +- Join Parquet text rows with two newlines and tokenize without added special + tokens. Use the first requested contiguous blocks, each containing `length+1` + tokens and predicting `length` tokens. Blocks do not overlap; cache/state reset + for each block. This protocol is for paired comparisons, not a claim of the + checkpoint's published WikiText perplexity. +- Shuffle the training blocks with the recorded seed. Use the same seed, block + count, token hashes, length, and learning rate across all trials. +- Train KDA attention parameters only. Keep them and AdamW moments in FP32, use + BF16 autocast and activation checkpointing, and freeze other model parameters + in BF16. Preserve the checkpoint's FP32 gate parameters when loading. AdamW has + zero weight decay, default learning rate `1e-5`, and gradient clipping at 1. +- Record per-block NLL, predicted token counts, source/data/token hashes, package + versions, training losses, finite-gradient checks, optimizer updates, step times, + and peak allocated training memory. The output contains no generated answers or + source document text. + +A short matched study can screen numerical sensitivity. It does not establish +broad downstream-task quality, long-context equivalence, or native serving speed. + +## Measure prefix and decode training cost + +```bash +python examples/llm_qat/linear_attention/benchmark_decode.py \ + --attention kda --length 257 --prefill 64 --dim 128 --output decode-cost.json +``` + +The benchmark checks outputs, final states, and input gradients against the +matching Torch numerical policy before measuring forward/backward time and peak +memory. See the [recorded decode study](decode_study.md) for historical results. +The comparison measures training emulation, not inference or compressed-cache cost. + +## State-only serving with vLLM + +See the [vLLM example](../../vllm_serve/README.md#linear-attention-state-quantization) +for incoming-state QDQ before native prefill/decode. Its invocation boundaries +differ from the token-write and ReplaySSM policies used during training. diff --git a/examples/llm_qat/linear_attention/benchmark_decode.py b/examples/llm_qat/linear_attention/benchmark_decode.py new file mode 100644 index 00000000000..db401a3a787 --- /dev/null +++ b/examples/llm_qat/linear_attention/benchmark_decode.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compare complete prefix/suffix QAT cost under matched token/replay policies.""" + +import argparse +import hashlib +import json +import random +import statistics +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn.functional as F +from benchmark_utils import benchmark_variants +from fla.ops.gated_delta_rule import chunk_gated_delta_rule +from fla.ops.kda import chunk_kda + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + LinearAttentionMatmulSites, + matmul_gdn, + matmul_kda, +) +from modelopt.torch.quantization.nn import TensorQuantizer + +if TYPE_CHECKING: + from collections.abc import Callable + + +def main(): + """Check outputs/gradients before interleaved, warmup-excluded measurements.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--attention", choices=["gdn", "kda"], default="kda") + parser.add_argument("--length", type=int, default=257) + parser.add_argument("--prefill", type=int, default=64) + parser.add_argument("--heads", type=int, default=4) + parser.add_argument("--dim", type=int, default=128) + parser.add_argument("--repeats", type=int, default=20) + parser.add_argument("--state-format", choices=["fp8_e4m3", "int8"], default="fp8_e4m3") + options = parser.parse_args() + if options.repeats < 2 or not 0 <= options.prefill < options.length: + parser.error("Require >=2 repeats and a nonempty suffix") + torch.manual_seed(2026) + torch.backends.cuda.matmul.allow_tf32 = False + shape = (1, options.length, options.heads, options.dim) + q, k = [ + F.normalize(torch.randn(shape, device="cuda"), dim=-1).requires_grad_() for _ in range(2) + ] + v = torch.randn(shape, device="cuda", requires_grad=True) + g = ( + -torch.rand(shape if options.attention == "kda" else shape[:-1], device="cuda") * 0.03 + ).requires_grad_() + beta = (torch.rand(shape[:-1], device="cuda") * 0.4).requires_grad_() + initial = ( + torch.randn(1, options.heads, options.dim, options.dim, device="cuda") * 0.1 + ).requires_grad_() + args = q, k, v, g, beta + sites = LinearAttentionMatmulSites() + w = TensorQuantizer(QuantizerAttributeConfig(enable=False)) + function = matmul_kda if options.attention == "kda" else matmul_gdn + variants: dict[str, Callable[[], tuple[torch.Tensor, torch.Tensor]]] = {} + for mode in ["exact", "token", "decay", "replay"]: + for impl in ["torch", "triton"]: + decode: dict[str, Any] = {"implementation": impl} + if mode == "replay": + decode.update(mode="replay", replay={"window": 8}) + if mode == "decay": + decode["decay_log_step"] = 1 / 256 + policy = LinearAttentionConfig(backend="matmul", decode=decode) + + def run(policy=policy, mode=mode): + return function( + *args, + initial_state=initial, + output_final_state=True, + sites=sites, + w_quantizer=w, + policy=policy, + state_qdq=mode != "exact", + state_format=options.state_format, + prefill_lengths=[options.prefill], + ) + + variants[f"{mode}_{impl}"] = run + native = chunk_kda if options.attention == "kda" else chunk_gated_delta_rule + variants["fla_bf16"] = lambda: native( + q.bfloat16(), + k.bfloat16(), + v.bfloat16(), + g, + beta, + initial_state=initial, + output_final_state=True, + ) + probes = [torch.randn_like(v), torch.randn_like(initial)] + correctness = {} + for mode in ["exact", "token", "decay", "replay"]: + results = [] + for impl in ["torch", "triton"]: + output = variants[f"{mode}_{impl}"]() + grads = torch.autograd.grad( + sum((x * p).sum() for x, p in zip(output, probes)), (*args, initial) + ) + results.append(tuple(x.detach() for x in (*output, *grads))) + errors = [float((a - b).norm() / a.norm().clamp_min(1e-8)) for a, b in zip(*results)] + if max(errors) > 1e-4 or not all( + torch.isfinite(x).all() for result in results for x in result + ): + raise RuntimeError(f"{mode}: correctness gate failed: {errors}") + correctness[mode] = errors + + rng = random.Random(2026) + samples, summary = benchmark_variants( + variants, + lambda result: result[0].square().mean() + result[1].square().mean(), + (*args, initial), + options.repeats, + rng, + ) + for mode in ["exact", "token", "decay", "replay"]: + ratios = [ + a["wall_ms"] / b["wall_ms"] + for a, b in zip(samples[f"{mode}_torch"], samples[f"{mode}_triton"]) + ] + bootstrap = sorted( + statistics.median(rng.choices(ratios, k=len(ratios))) for _ in range(5000) + ) + summary[f"{mode}_triton"]["paired_speedup_vs_reference"] = { + "median": statistics.median(ratios), + "interval_95": [bootstrap[124], bootstrap[4874]], + } + sources = [ + *list(Path("modelopt/torch/quantization/linear_attention").glob("*.py")), + Path("modelopt/torch/kernels/quantization/linear_attention/decode.py"), + Path(__file__), + Path(__file__).with_name("benchmark_utils.py"), + ] + result = { + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "shape": shape, + "prefill": options.prefill, + "state_format": options.state_format, + "attention": options.attention, + "dtype": "float32; FLA baseline casts Q/K/V to bfloat16", + "repeats": options.repeats, + "correctness_relative_norms": correctness, + "source_sha256": {str(p): hashlib.sha256(p.read_bytes()).hexdigest() for p in sources}, + "summary": summary, + "samples": samples, + } + options.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(summary, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/llm_qat/linear_attention/benchmark_prefill.py b/examples/llm_qat/linear_attention/benchmark_prefill.py new file mode 100644 index 00000000000..8f5e03dd479 --- /dev/null +++ b/examples/llm_qat/linear_attention/benchmark_prefill.py @@ -0,0 +1,164 @@ +# 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. + +"""Measure GDN/KDA prefill numerical-emulation overhead against the exact FLA path.""" + +import argparse +import hashlib +import importlib.metadata +import json +import random +import statistics +import subprocess +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as F +from benchmark_utils import benchmark_variants +from fla.ops.kda import chunk_kda + +from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( + chunk_gated_delta_rule, +) +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + matmul_gdn, + matmul_kda, +) +from modelopt.torch.quantization.linear_attention.matmul import LinearAttentionMatmulSites +from modelopt.torch.quantization.nn import TensorQuantizer + + +def main(): + """Run interleaved fixed-workload forward/backward trials and write a JSON receipt.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--attention", choices=["gdn", "kda"], default="gdn") + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--length", type=int, default=1024) + parser.add_argument("--heads", type=int, default=4) + parser.add_argument("--dim", type=int, default=64) + parser.add_argument("--repeats", type=int, default=20) + parser.add_argument("--revision", default=None) + options = parser.parse_args() + torch.manual_seed(2026) + torch.backends.cuda.matmul.allow_tf32 = False + shape = (options.batch, options.length, options.heads, options.dim) + q, k = [ + F.normalize(torch.randn(shape, device="cuda"), dim=-1).bfloat16().requires_grad_() + for _ in range(2) + ] + v = torch.randn(shape, device="cuda", dtype=torch.bfloat16, requires_grad=True) + gate_shape = shape if options.attention == "kda" else shape[:-1] + g = (-torch.rand(gate_shape, device="cuda") * 0.03).requires_grad_() + beta = torch.rand(shape[:-1], device="cuda", requires_grad=True) + state = ( + torch.randn(options.batch, options.heads, options.dim, options.dim, device="cuda") * 0.1 + ).requires_grad_() + inputs = (q, k, v, g, beta) + baseline = chunk_kda if options.attention == "kda" else chunk_gated_delta_rule + materialized = matmul_kda if options.attention == "kda" else matmul_gdn + variants = { + "exact_fla": lambda: baseline(*inputs, initial_state=state, output_final_state=True) + } + for mode in ( + "exact_matmul", + "fp8_operands", + "nvfp4_operands", + "fp16_accumulation", + "fp8_with_state", + ): + sites = LinearAttentionMatmulSites() + w = TensorQuantizer(QuantizerAttributeConfig(enable=False)) + policy = LinearAttentionConfig(backend="matmul") + if mode in ("fp8_operands", "fp8_with_state", "nvfp4_operands"): + attributes: dict[str, Any] = {"num_bits": (4, 3), "axis": (0, 1, 2), "type": "dynamic"} + if mode == "nvfp4_operands": + attributes = { + "num_bits": (2, 1), + "type": "dynamic", + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + } + for quantizer in sites.modules(): + if isinstance(quantizer, TensorQuantizer): + quantizer.set_from_attribute_config(attributes) + quantizer.enable() + w.set_from_attribute_config(attributes) + w.enable() + if mode == "fp16_accumulation": + policy = LinearAttentionConfig( + backend="matmul", + matmul={ + name: {"accumulator_dtype": "float16", "reduction_block": 16} for name in sites + }, + ) + + def run(sites=sites, w=w, policy=policy, mode=mode): + return materialized( + *inputs, + initial_state=state, + output_final_state=True, + sites=sites, + w_quantizer=w, + policy=policy, + state_qdq=mode == "fp8_with_state", + ) + + variants[mode] = run + + samples, summary = benchmark_variants( + variants, + lambda result: result[0].float().square().mean() + result[1].square().mean(), + (*inputs, state), + options.repeats, + random.Random(2026), + ) + for name, values in samples.items(): + ratios = [ + sample["wall_ms"] / control["wall_ms"] + for sample, control in zip(values, samples["exact_fla"]) + ] + summary[name]["median_paired_wall_ratio"] = statistics.median(ratios) + revision = options.revision + if revision is None: + revision = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip() + sources = [ + Path("modelopt/torch/quantization/linear_attention") / name + for name in ("prefill.py", "kda.py", "matmul.py", "reference.py", "config.py") + ] + sources.append(Path(__file__).with_name("benchmark_utils.py")) + result = { + "attention": options.attention, + "git_head": revision, + "source_sha256": { + str(path): hashlib.sha256(path.read_bytes()).hexdigest() for path in sources + }, + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "fla_core": importlib.metadata.version("fla-core"), + "shape": shape, + "dtype": "bfloat16", + "repeats": options.repeats, + "summary": summary, + "samples": samples, + } + options.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(summary, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/examples/llm_qat/linear_attention/benchmark_utils.py b/examples/llm_qat/linear_attention/benchmark_utils.py new file mode 100644 index 00000000000..e1b1510b13b --- /dev/null +++ b/examples/llm_qat/linear_attention/benchmark_utils.py @@ -0,0 +1,59 @@ +# 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 measurements for the linear-attention training examples.""" + +import statistics +import time + +import torch + + +def benchmark_variants(variants, loss_fn, inputs, repeats, rng): + """Measure interleaved forward/backward samples after three warmup rounds.""" + + def measure(fn): + torch.cuda.synchronize() + allocated = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + start, mid, end = [torch.cuda.Event(enable_timing=True) for _ in range(3)] + wall_start = time.perf_counter() + start.record() + output = fn() + mid.record() + torch.autograd.grad(loss_fn(output), inputs) + end.record() + torch.cuda.synchronize() + return { + "forward_ms": start.elapsed_time(mid), + "backward_ms": mid.elapsed_time(end), + "wall_ms": 1000 * (time.perf_counter() - wall_start), + "peak_extra_bytes": torch.cuda.max_memory_allocated() - allocated, + } + + for fn in variants.values(): + for _ in range(3): + measure(fn) + samples = {name: [] for name in variants} + for _ in range(repeats): + order = list(variants) + rng.shuffle(order) + for name in order: + samples[name].append(measure(variants[name])) + summary = { + name: {key: statistics.median(row[key] for row in rows) for key in rows[0]} + for name, rows in samples.items() + } + return samples, summary diff --git a/examples/llm_qat/linear_attention/compare_quality.py b/examples/llm_qat/linear_attention/compare_quality.py new file mode 100644 index 00000000000..24b3fe953e4 --- /dev/null +++ b/examples/llm_qat/linear_attention/compare_quality.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compare matched numerical-policy studies using paired block NLL differences.""" + +import argparse +import json +import math +import random +import statistics +from pathlib import Path + +_MATCHED_FIELDS = ( + "model", + "model_revision", + "dataset", + "dataset_config", + "dataset_revision", + "eval_split", + "train_file_sha256", + "eval_file_sha256", + "train_tokens_sha256", + "eval_tokens_sha256", + "sequence_length", + "prefill_tokens", + "loss_scope", + "seed", + "train_order", + "training", + "trainable_parameters", + "kda_layers", + "learning_rate", + "weight_decay", + "gradient_clip", + "train_predicted_tokens", + "source_sha256", + "packages", + "torch", + "gpu", +) + + +def compare(control, candidate, *, margin=0.02, samples=5000, seed=2026): + """Reject unmatched inputs and report descriptive paired block-bootstrap bounds.""" + if not math.isfinite(margin) or samples < 2: + raise ValueError("margin must be finite and samples must be at least two") + mismatches = [key for key in _MATCHED_FIELDS if control[key] != candidate[key]] + if mismatches: + raise ValueError(f"Unmatched study fields: {mismatches}") + result = {} + for phase in ("before", "after"): + left, right = control[phase], candidate[phase] + a, b = left["block_nll"], right["block_nll"] + if len(a) != len(b) or len(a) < 2 or left["predicted_tokens"] != right["predicted_tokens"]: + raise ValueError("At least two matched evaluation blocks are required") + if not all(math.isfinite(x) for x in (*a, *b)): + raise ValueError("NLL values must be finite") + delta = [y - x for x, y in zip(a, b)] + rng = random.Random(seed) + boot = sorted(statistics.mean(rng.choices(delta, k=len(delta))) for _ in range(samples)) + lower, upper = boot[int(0.025 * (samples - 1))], boot[int(0.975 * (samples - 1))] + mean = statistics.mean(delta) + result[phase] = { + "mean_nll_delta": mean, + "perplexity_ratio": math.exp(mean), + "block_nll_delta": delta, + "bootstrap_interval_95": [lower, upper], + "margin": margin, + "upper_bound_within_margin": upper <= margin, + "blocks": len(delta), + "predicted_tokens": left["predicted_tokens"], + } + return { + "method": "paired block percentile bootstrap; descriptive pilot interval", + "samples": samples, + "seed": seed, + "comparison": result, + } + + +def main(): + """Write a reproducible matched comparison receipt.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--control", type=Path, required=True) + parser.add_argument("--candidate", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--margin", type=float, default=0.02) + options = parser.parse_args() + result = compare( + json.loads(options.control.read_text()), + json.loads(options.candidate.read_text()), + margin=options.margin, + ) + result["control"] = str(options.control) + result["candidate"] = str(options.candidate) + options.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result["comparison"], indent=2)) + + +if __name__ == "__main__": + main() diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_decay_1024.json b/examples/llm_qat/linear_attention/configs/kda_decode_decay_1024.json new file mode 100644 index 00000000000..3e49626b796 --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_decay_1024.json @@ -0,0 +1,35 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": [ + 4, + 3 + ], + "type": "dynamic", + "axis": [ + 0, + 1 + ] + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "decay_log_step": 0.0009765625 + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_decay_256.json b/examples/llm_qat/linear_attention/configs/kda_decode_decay_256.json new file mode 100644 index 00000000000..0b3b972c25d --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_decay_256.json @@ -0,0 +1,35 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": [ + 4, + 3 + ], + "type": "dynamic", + "axis": [ + 0, + 1 + ] + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "decay_log_step": 0.00390625 + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_decay_4096.json b/examples/llm_qat/linear_attention/configs/kda_decode_decay_4096.json new file mode 100644 index 00000000000..3b488c92bbb --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_decay_4096.json @@ -0,0 +1,35 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": [ + 4, + 3 + ], + "type": "dynamic", + "axis": [ + 0, + 1 + ] + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "decay_log_step": 0.000244140625 + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_exact.json b/examples/llm_qat/linear_attention/configs/kda_decode_exact.json new file mode 100644 index 00000000000..4402b29f1bc --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_exact.json @@ -0,0 +1,20 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton" + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_replay_fp8.json b/examples/llm_qat/linear_attention/configs/kda_decode_replay_fp8.json new file mode 100644 index 00000000000..e7a86e20ca3 --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_replay_fp8.json @@ -0,0 +1,40 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": [ + 4, + 3 + ], + "type": "dynamic", + "axis": [ + 0, + 1 + ] + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "mode": "replay", + "replay": { + "window": 8, + "factor_qdq": true, + "encoding": "once" + } + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_replay_int8.json b/examples/llm_qat/linear_attention/configs/kda_decode_replay_int8.json new file mode 100644 index 00000000000..0ddd98e9ad7 --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_replay_int8.json @@ -0,0 +1,39 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": 8, + "type": "dynamic", + "axis": [ + 0, + 1 + ], + "unsigned": false, + "narrow_range": true + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "mode": "replay", + "replay": { + "window": 8, + "factor_qdq": true, + "encoding": "once" + } + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_replay_int8_hadamard.json b/examples/llm_qat/linear_attention/configs/kda_decode_replay_int8_hadamard.json new file mode 100644 index 00000000000..2a6f0acb88b --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_replay_int8_hadamard.json @@ -0,0 +1,44 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": 8, + "type": "dynamic", + "axis": [ + 0, + 1 + ], + "unsigned": false, + "narrow_range": true + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "mode": "replay", + "replay": { + "window": 8, + "factor_qdq": false, + "encoding": "once" + }, + "state_codec": "int8_hadamard32", + "readout": "working" + }, + "state": { + "block_v": 64 + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_state_fp8.json b/examples/llm_qat/linear_attention/configs/kda_decode_state_fp8.json new file mode 100644 index 00000000000..27ab6b2fbdd --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_state_fp8.json @@ -0,0 +1,34 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": [ + 4, + 3 + ], + "type": "dynamic", + "axis": [ + 0, + 1 + ] + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton" + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json b/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json new file mode 100644 index 00000000000..f6cc5fef650 --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json @@ -0,0 +1,35 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": 8, + "type": "dynamic", + "axis": [ + 0, + 1 + ], + "unsigned": false, + "narrow_range": true + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "state_codec": "int8_hadamard32", + "prefill_state_qdq": false + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_exact_matmul.json b/examples/llm_qat/linear_attention/configs/kda_exact_matmul.json new file mode 100644 index 00000000000..dc0ddcbb311 --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_exact_matmul.json @@ -0,0 +1,20 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "model.layers.*.self_attn", + "cfg": { + "backend": "matmul", + "elementwise": { + "output_add": "float32" + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_prefill_fp8.json b/examples/llm_qat/linear_attention/configs/kda_prefill_fp8.json new file mode 100644 index 00000000000..ba6b824c434 --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_prefill_fp8.json @@ -0,0 +1,47 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*linear_attn_sites.*", + "cfg": { + "num_bits": [ + 4, + 3 + ], + "type": "dynamic", + "axis": [ + 0, + 1, + 2 + ] + } + }, + { + "quantizer_name": "*kda_w_quantizer", + "cfg": { + "num_bits": [ + 4, + 3 + ], + "type": "dynamic", + "axis": [ + 0, + 1, + 2 + ] + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "model.layers.*.self_attn", + "cfg": { + "backend": "matmul" + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/configs/kda_prefill_nvfp4.json b/examples/llm_qat/linear_attention/configs/kda_prefill_nvfp4.json new file mode 100644 index 00000000000..13e4b66dd0a --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_prefill_nvfp4.json @@ -0,0 +1,53 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*linear_attn_sites.*", + "cfg": { + "num_bits": [ + 2, + 1 + ], + "type": "dynamic", + "block_sizes": { + "-1": 16, + "type": "dynamic", + "scale_bits": [ + 4, + 3 + ] + } + } + }, + { + "quantizer_name": "*kda_w_quantizer", + "cfg": { + "num_bits": [ + 2, + 1 + ], + "type": "dynamic", + "block_sizes": { + "-1": 16, + "type": "dynamic", + "scale_bits": [ + 4, + 3 + ] + } + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "model.layers.*.self_attn", + "cfg": { + "backend": "matmul" + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/decode_study.md b/examples/llm_qat/linear_attention/decode_study.md new file mode 100644 index 00000000000..4c92c57a321 --- /dev/null +++ b/examples/llm_qat/linear_attention/decode_study.md @@ -0,0 +1,116 @@ +# Decode-aware QAT qualification study + +These measurements describe the recorded implementation and workloads below; +they do not qualify later revisions or INT8/Hadamard policies. + +## Fixed protocol + +The study uses `arcee-ai/AFM-4.5B-Base-KDA-Only` at revision +`01ad2e06ee4f1214193c17b69e09105a9b257e80` and Wikitext-2 raw at revision +`b08601e04326c79dfdd32d625aee71d232d685c3`. The actual checkpoint has 36 KDA +layers, hidden size 2560, 20 heads, and key/value dimension 128. Parameters, +source, dataset files, and token blocks are bound by stored hashes. + +[decode_study_plan.json](decode_study_plan.json) fixes selection before held-out test access. Every +block contains 129 input tokens: 64 explicit prefix tokens followed by a decode +suffix. The loss scores labels at positions 64 through 128 (65 predictions per +block); it does not average prefix loss into decode quality. Controls and +candidates use the exact prefill solve, FP32 working arithmetic, BF16 outer +autocast, and seed 2026. + +Validation uses 32 blocks (2,080 scored tokens) without training. State FP8 and +replay must each pass an NLL upper-bound margin of 0.02 against exact. Decay +selection chooses the coarsest passing grid among 1/256, 1/1024, and 1/4096, +requiring the same margin against both state FP8 and exact. Eligible modes then +receive matched 32-step attention-only QAT pilots and 64 held-out test blocks. +These descriptive paired-block bootstrap bounds are not broad quality guarantees. + +## Validation selection + +| Policy | Suffix perplexity | Mean NLL increase vs exact | 95% interval | +| --- | ---: | ---: | --- | +| Exact | 9.019945 | 0 | — | +| Token state FP8 | 9.023610 | 0.000406 | [-0.002017, 0.002718] | +| State FP8 + log grid 1/256 | 9.029220 | 0.001028 | [-0.001118, 0.002986] | +| State FP8 + log grid 1/1024 | 9.035839 | 0.001761 | [-0.000115, 0.003723] | +| State FP8 + log grid 1/4096 | 9.024412 | 0.000495 | [-0.001268, 0.002299] | +| Replay FP8, window 8 | 9.019499 | -0.000049 | [-0.001867, 0.001630] | + +All families passed the validation gate. The predeclared coarsest-grid rule +selected **1/256**; its paired comparison against state FP8 also passed. The +held-out set was frozen to exact, token state FP8, state FP8 with grid 1/256, +and replay FP8 with window 8, before accessing test data. + +## Held-out training pilot + +All four selected modes completed 32 attention-only updates on 32 training +blocks (2,080 scored training tokens) using FP32 AdamW, learning rate 1e-5, +zero weight decay, gradient clipping at 1.0, and non-reentrant activation +checkpointing. All 994,043,088 trainable attention parameters had finite gradients +on every update, and the query projection changed in each run. Evaluation used +64 held-out blocks (4,160 scored tokens), with identical token hashes and +training order across modes. + +| Policy | PPL before | PPL after | Post-training NLL delta vs exact | 95% interval | +| --- | ---: | ---: | ---: | --- | +| Exact | 12.899382 | 12.641259 | 0 | — | +| Token state FP8 | 12.904970 | 12.656691 | 0.001220 | [-0.000290, 0.002703] | +| State FP8 + grid 1/256 | 12.915614 | 12.653026 | 0.000930 | [-0.000802, 0.002673] | +| Replay FP8, window 8 | 12.909424 | 12.646554 | 0.000419 | [-0.000566, 0.001410] | + +All three approximations met the 0.02 held-out pilot margin after training; the +largest upper bound was 0.002703. Their intervals include zero, so this study does +not establish a quality advantage for any approximation. All modes, including +the exact control, improved after fine-tuning. These small short-context pilots +verify model-level training behavior and bounded measured differences; they do +not demonstrate recovery of a substantial accuracy gap or broad downstream quality. + +## Complete training cost + +The synthetic H100 workload is `[B,T,H,D] = [1,257,4,128]`, with 64 prefix tokens, +FP32 working inputs, TF32 disabled, three warmups, and 20 interleaved samples. +The table includes the complete prefix/suffix forward and backward, including +returned final-state gradients. Speedups compare each fused implementation with +its matching direct Torch numerical reference, with descriptive paired-median +bootstrap intervals. Memory is peak allocated memory above the live input baseline. + +| Policy | Torch wall ms | Fused wall ms | Speedup [95% interval] | Torch / fused extra MiB | +| --- | ---: | ---: | --- | ---: | +| GDN exact | 188.830 | 4.758 | 39.64 [39.52, 40.19] | 103.75 / 14.68 | +| GDN token | 224.532 | 7.239 | 31.02 [30.73, 31.06] | 103.75 / 14.68 | +| GDN decay | 231.847 | 7.321 | 31.70 [31.57, 31.82] | 103.75 / 14.69 | +| GDN replay | 325.854 | 6.206 | 52.47 [51.12, 52.73] | 272.08 / 14.94 | +| KDA exact | 193.183 | 11.110 | 17.40 [17.33, 17.53] | 152.35 / 63.88 | +| KDA token | 228.732 | 11.265 | 20.32 [20.25, 20.43] | 152.35 / 63.88 | +| KDA decay | 235.972 | 11.310 | 20.87 [20.77, 20.98] | 152.35 / 64.26 | +| KDA replay | 327.361 | 11.569 | 28.34 [28.16, 28.39] | 321.54 / 64.51 | + +The exact FLA reference, including conversion of Q/K/V to BF16, took 1.806 ms +for GDN and 2.347 ms for KDA. Fused numerical emulation remains slower than that +reference. These measurements establish a reduction in QAT emulation overhead, +not native low-precision execution, inference latency, or compressed-cache cost. +Every policy first passed checks of outputs, final state, and all input gradients +against the matching Torch reference; source hashes accompany the raw samples. + +## Validation boundary + +The H100 regression suite passed 148 cases, with nine expected skips for FP32 +Hopper/TileLang modes that are rejected before launch. It includes GDN TP=1/2 +sharded save/restore and optimizer updates, KDA layer training/restore, explicit +phase handling, and custom backward tests. All eight long-trajectory cases also +passed the tightened `1e-4` relative-norm gate on H100 and SM86. + +Local checks passed 139 unit cases plus two re-encoding continuation cases, +44 kernel/layer cases, and four additional full-width/value-tail cases covering +32- and 128-column codec blocks (also passed on H100), plus two composed-policy +cases (also passed on H100). The combined H100 total is **154 passed, nine +expected skips**; the local GPU total is 50 passed. The full Sphinx build passed with FLA imports +disabled. Measurements pin retained source snapshots. Final cleanup moved import +boundaries and added explanatory comments; non-import AST equality was verified +and focused regression/import checks were rerun. The retained benchmark harness +differs only in the order of two imported typing names, with AST equality +verified after normalizing that order. + +Model results apply to this checkpoint, short context, small corpus slice, and +seed. Serving cache integration, real compression, long-context quality, and +multi-model recovery remain unqualified. diff --git a/examples/llm_qat/linear_attention/decode_study_plan.json b/examples/llm_qat/linear_attention/decode_study_plan.json new file mode 100644 index 00000000000..17ba96d7173 --- /dev/null +++ b/examples/llm_qat/linear_attention/decode_study_plan.json @@ -0,0 +1,40 @@ +{ + "model": "arcee-ai/AFM-4.5B-Base-KDA-Only", + "model_revision": "01ad2e06ee4f1214193c17b69e09105a9b257e80", + "dataset": "Salesforce/wikitext/wikitext-2-raw-v1", + "dataset_revision": "b08601e04326c79dfdd32d625aee71d232d685c3", + "seed": 2026, + "length": 128, + "prefill_tokens": 64, + "loss_scope": "suffix", + "validation": { + "blocks": 32, + "train_steps": 0, + "controls": [ + "exact", + "state_fp8" + ], + "candidates": [ + "decay_256", + "decay_1024", + "decay_4096", + "replay_fp8", + "state_int8", + "replay_int8" + ], + "rule": "Select coarsest decay grid whose paired NLL upper bound is <=0.02 versus state_fp8 and exact; qualify replay independently versus exact. State FP8 must pass versus exact. Do not access test until selection is frozen. Qualify INT8 state and INT8-anchor replay separately against exact and their matching FP8 controls with the same 0.02 paired NLL upper bound." + }, + "held_out": { + "eval_split": "test", + "blocks": 64, + "train_steps": 32, + "learning_rate": 1e-05, + "gradient_clip": 1.0, + "modes": "exact, state_fp8 if eligible, selected decay if eligible, replay if eligible, INT8 state/replay if individually eligible", + "rule": "Run matched attention-only QAT pilots. Report before/after paired NLL deltas, intervals, finite gradients and updates; a passing pilot does not establish broad quality recovery." + }, + "fallback": "Keep exact default; explicitly report failed gates. Do not tune on test.", + "integration": { + "failed_policy_smoke": "One attention-only training step and four validation blocks for any family excluded by screening; report solely as gradient/update integration, never recovery or held-out qualification." + } +} diff --git a/examples/llm_qat/linear_attention/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..618cc74de50 --- /dev/null +++ b/examples/llm_qat/linear_attention/train.py @@ -0,0 +1,287 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run a reproducible linear-attention QAT trial on pinned local model and data files.""" + +import argparse +import hashlib +import importlib.metadata +import json +import math +import random +import time +from contextlib import nullcontext +from pathlib import Path + +import pyarrow.parquet as pq +import torch +from fla.layers.kda import KimiDeltaAttention +from transformers import AutoModelForCausalLM, AutoTokenizer + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase + + +def _token_blocks(path, tokenizer, length, count): + text = "\n\n".join(pq.read_table(path, columns=["text"])["text"].to_pylist()) + ids = tokenizer.encode(text, add_special_tokens=False) + needed = count * (length + 1) + if len(ids) < needed: + raise ValueError(f"{path} has {len(ids)} tokens; {needed} are required") + tensor = torch.tensor(ids[:needed], dtype=torch.int64).reshape(count, length + 1) + digest = hashlib.sha256(tensor.numpy().tobytes()).hexdigest() + return tensor, digest + + +def _phase(model, prefill_tokens): + return ( + nullcontext() + if prefill_tokens is None + else linear_attention_training_phase(model, [prefill_tokens]) + ) + + +def _labels(ids, prefill_tokens): + labels = ids.clone() + if prefill_tokens is not None: + labels[:, :prefill_tokens] = -100 + return labels + + +def _evaluate(model, blocks, prefill_tokens=None): + model.eval() + losses = [] + with ( + _phase(model, prefill_tokens), + torch.no_grad(), + torch.autocast("cuda", dtype=torch.bfloat16), + ): + for block in blocks: + ids = block.unsqueeze(0).cuda() + loss = model(input_ids=ids, labels=_labels(ids, prefill_tokens), use_cache=False).loss + if not torch.isfinite(loss): + raise RuntimeError("Non-finite evaluation loss") + losses.append(float(loss.detach())) + mean = sum(losses) / len(losses) + return { + "mean_nll": mean, + "perplexity": math.exp(mean), + "block_nll": losses, + "predicted_tokens": len(blocks) * (blocks.shape[1] - max(1, prefill_tokens or 0)), + } + + +def main(): + """Compare pre/post-training held-out NLL using an explicit numerical policy.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--model-id", default="arcee-ai/AFM-4.5B-Base-KDA-Only") + parser.add_argument("--model-revision", default="01ad2e06ee4f1214193c17b69e09105a9b257e80") + parser.add_argument("--train-data", type=Path, required=True) + parser.add_argument("--eval-data", type=Path, required=True) + parser.add_argument("--eval-split", choices=["validation", "test"], default="validation") + parser.add_argument("--dataset-revision", default="b08601e04326c79dfdd32d625aee71d232d685c3") + parser.add_argument("--quant-config", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--train-steps", type=int, default=1) + parser.add_argument("--eval-blocks", type=int, default=4) + parser.add_argument("--length", type=int, default=128) + parser.add_argument( + "--prefill-tokens", type=int, help="Explicit prefix; score/train suffix labels only" + ) + parser.add_argument("--seed", type=int, default=2026) + parser.add_argument("--learning-rate", type=float, default=1e-5) + options = parser.parse_args() + if options.length < 64 or options.train_steps < 0 or options.eval_blocks < 1: + parser.error("length must be >=64, train-steps nonnegative, and eval-blocks positive") + if options.train_data.resolve() == options.eval_data.resolve(): + parser.error("Training and evaluation data must be distinct splits") + if options.prefill_tokens is not None and not 0 <= options.prefill_tokens <= options.length: + parser.error("prefill-tokens must be between zero and length") + torch.manual_seed(options.seed) + random.seed(options.seed) + torch.set_float32_matmul_precision("highest") + tokenizer = AutoTokenizer.from_pretrained( + options.model, local_files_only=True, trust_remote_code=options.trust_remote_code + ) + train, train_hash = _token_blocks( + options.train_data, tokenizer, options.length, options.train_steps + ) + evaluation, eval_hash = _token_blocks( + options.eval_data, tokenizer, options.length, options.eval_blocks + ) + model, loading = AutoModelForCausalLM.from_pretrained( + options.model, + local_files_only=True, + trust_remote_code=options.trust_remote_code, + torch_dtype=torch.float32, + output_loading_info=True, + ) + if loading["missing_keys"] or loading["unexpected_keys"] or loading.get("mismatched_keys"): + raise RuntimeError(f"Checkpoint/model mismatch: {loading}") + model.requires_grad_(False) + layers = [m for m in model.modules() if isinstance(m, KimiDeltaAttention)] + if not layers: + raise RuntimeError("The selected model contains no FLA KimiDeltaAttention layers") + # Keep trainable attention parameters and optimizer moments in FP32 under BF16 autocast. + for layer in layers: + layer.requires_grad_(True) + for parameter in model.parameters(): + if not parameter.requires_grad: + parameter.data = parameter.data.to(torch.bfloat16) + model.cuda() + quant_cfg = None + if options.quant_config: + quant_cfg = json.loads(options.quant_config.read_text()) + mtq.quantize(model, quant_cfg) + model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + before = _evaluate(model, evaluation, options.prefill_tokens) + trainable = [p for p in model.parameters() if p.requires_grad] + optimizer = torch.optim.AdamW(trainable, lr=options.learning_rate, weight_decay=0) + model.train() + losses = [] + grad_norms = [] + timings = [] + torch.cuda.reset_peak_memory_stats() + probe = layers[0].q_proj.weight + original = probe.detach().clone() + order = list(range(len(train))) + random.Random(options.seed).shuffle(order) + for index in order: + optimizer.zero_grad(set_to_none=True) + ids = train[index].unsqueeze(0).cuda() + torch.cuda.synchronize() + start = time.perf_counter() + # Activation checkpoint recomputation needs the same explicit phase metadata. + with _phase(model, options.prefill_tokens): + with torch.autocast("cuda", dtype=torch.bfloat16): + loss = model( + input_ids=ids, labels=_labels(ids, options.prefill_tokens), use_cache=False + ).loss + if not torch.isfinite(loss): + raise RuntimeError("Non-finite training loss") + loss.backward() + gradients = [p.grad for p in trainable if p.grad is not None] + if ( + len(gradients) != len(trainable) + or not torch.stack([torch.isfinite(g).all() for g in gradients]).all() + ): + raise RuntimeError("Missing or non-finite training gradients") + norm = torch.nn.utils.clip_grad_norm_(trainable, 1.0, error_if_nonfinite=True) + optimizer.step() + torch.cuda.synchronize() + losses.append(float(loss.detach())) + grad_norms.append(float(norm)) + timings.append(time.perf_counter() - start) + print( + json.dumps( + { + "step": len(losses), + "loss": losses[-1], + "gradient_norm": grad_norms[-1], + "seconds": timings[-1], + } + ), + flush=True, + ) + query_changed = not torch.equal(original, probe) + if options.train_steps and not query_changed: + raise RuntimeError("Training did not update the query projection") + peak = torch.cuda.max_memory_allocated() + optimizer.zero_grad(set_to_none=True) + after = _evaluate(model, evaluation, options.prefill_tokens) if options.train_steps else before + root = Path(__file__).resolve().parents[3] + sources = [ + root / "modelopt/torch/quantization/linear_attention" / name + for name in [ + "config.py", + "kda.py", + "prefill.py", + "matmul.py", + "solve.py", + "decode.py", + "decode_prefill.py", + "utils.py", + ] + ] + sources.extend( + [ + root / "modelopt/torch/kernels/quantization/linear_attention/decode.py", + root / "modelopt/torch/kernels/quantization/linear_attention/int8.py", + root / "modelopt/torch/quantization/plugins/kda.py", + root / "modelopt/torch/quantization/plugins/linear_attention.py", + Path(__file__).resolve(), + ] + ) + result = { + "source_sha256": { + str(path.relative_to(root)): hashlib.sha256(path.read_bytes()).hexdigest() + for path in sources + }, + "model": options.model_id, + "model_revision": options.model_revision, + "dataset": "Salesforce/wikitext", + "dataset_config": "wikitext-2-raw-v1", + "dataset_revision": options.dataset_revision, + "eval_split": options.eval_split, + "train_file_sha256": hashlib.sha256(options.train_data.read_bytes()).hexdigest(), + "eval_file_sha256": hashlib.sha256(options.eval_data.read_bytes()).hexdigest(), + "train_tokens_sha256": train_hash, + "eval_tokens_sha256": eval_hash, + "sequence_length": options.length, + "prefill_tokens": options.prefill_tokens, + "loss_scope": "full" if options.prefill_tokens is None else "suffix", + "seed": options.seed, + "train_order": order, + "training": "KDA attention parameters only; FP32 AdamW, BF16 autocast, gradient checkpointing", + "trainable_parameters": sum(p.numel() for p in trainable), + "kda_layers": len(layers), + "learning_rate": options.learning_rate, + "weight_decay": 0, + "gradient_clip": 1.0, + "quant_config": quant_cfg, + "before": before, + "after": after, + "training_losses": losses, + "gradient_norms": grad_norms, + "step_seconds": timings, + "train_predicted_tokens": options.train_steps + * (options.length + 1 - max(1, options.prefill_tokens or 0)), + "peak_train_bytes": peak, + "query_weight_changed": query_changed, + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "packages": { + name: importlib.metadata.version(name) + for name in ["transformers", "flash-linear-attention", "fla-core", "triton"] + }, + } + options.output.parent.mkdir(parents=True, exist_ok=True) + options.output.write_text(json.dumps(result, indent=2) + "\n") + print( + json.dumps( + { + "before": before["perplexity"], + "after": after["perplexity"], + "output": str(options.output), + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 9c4eec609b7..e12e429f17d 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -7,6 +7,10 @@ Compared with realquant, fakequant is 2-5x slower, but doesn't require dedicated The general fakequant example is tested with vLLM 0.9.0, 0.19.1, 0.26.0, and 0.28.0. The compact NVFP4 attention worker documented below requires vLLM 0.15.0 or newer. +For GDN/KDA recurrent-state QDQ before native prefill and decode, see the +[linear-attention example](#linear-attention-state-quantization), which +documents the initial vLLM 0.15.x runtime requirements and state-only INT8 recipe. + ## Prepare environment Use the Dockerfile to build an environment with vLLM 0.28.0: @@ -279,6 +283,79 @@ Supported configurations are regular decoder self-attention with FlashInfer or F Unsupported features are sliding window, ALiBi, softcap, sinks, FP8 KV cache, cross/encoder/MLA attention, KV sharing or transfer, prefix caching, speculative decoding, DBO/ubatching, and `FULL` mixed/prefill CUDA graphs. +## Linear-attention state quantization + +The existing `FakeQuantWorker` adds a ModelOpt `TensorQuantizer` immediately before +native vLLM prefill and decode for `Qwen3NextGatedDeltaNet` and +`KimiDeltaAttention`. Native attention kernels, projections, convolution, gates, +outputs, and cache management remain in use. + +### Runtime and launch + +The initial adapter targets vLLM 0.15.x V1 with key-first +`[slot, head, Dk, Dv]` FP32 recurrent state. Runtime checks reject other layouts +and versions. Qualification uses source revision `930288170` with Torch +2.9.1/CUDA 12.8 on RTX A6000 GPUs. + +```bash +PYTHONPATH=.:examples/vllm_serve \ +RECIPE_PATH=examples/vllm_serve/linear_attention_state_int8.yaml \ +python examples/vllm_serve/vllm_serve_fakequant.py /path/to/model \ + --tensor-parallel-size 2 --enforce-eager --no-async-scheduling \ + --no-enable-prefix-caching --mamba-cache-dtype float32 +``` + +The example enables signed symmetric dynamic INT8 state QDQ. Set +`num_bits: [4, 3]` in each state quantizer to select FP8 E4M3. Each invocation +receives `[active_sequence, local_head, Dk, Dv]`; `axis: [0, 1]` retains one +scale per sequence and local head, reducing over the whole state matrix. +State tensors remain FP32 after dequantization. + +Dynamic scales need no calibration dataset; `algorithm: null` skips dataset +loading. Weight/activation calibration can use the worker's existing recipe +and calibration loop. Static state calibration is unsupported in this adapter. + +Use `MODELOPT_STATE_PATH=/path/to/modelopt_state.pt` instead of a recipe to +restore quantizer configuration. Reload maps recurrent-state quantizer names +through the model's HF-to-vLLM mapper. Checkpoint weights must already match +the vLLM model architecture. + +### Quantization boundaries + +- Prefill: vLLM gathers initial states and zeros fresh requests. The wrapper + applies `TensorQuantizer` to this tensor and calls the original prefill kernel. +- Decode: the wrapper gathers active native cache slots, quantizes them, writes + them back, and calls the original decode kernel. Other slots are untouched. +- Native kernels compute outputs and update recurrent state normally. No + additional rounding is applied to the kernel's final-state write. +- The next call quantizes that state before reading it. A fresh zero state is + unchanged; prompt-to-decode rounding happens before the first decode call. +- Scheduler-level chunked prefill creates one QDQ boundary per invocation. + Internal kernel chunks do not create extra state-quantization boundaries. + Results can therefore depend on scheduler prompt-chunk sizes. +- Requests, cache-slot reuse, and preemption remain managed by vLLM. The plugin + allocates temporary gathered states, with no additional persistent state cache. +- TP ranks quantize their local heads independently, without scale all-reduce. + +This state-only adapter accepts the default `linear_attention` execution policy +and state quantizer configuration. Enabled W/operand quantizers and nondefault +training, replay, decay, or solve policies are rejected rather than reinterpreted. +The previous experimental replay-serving recipe is replaced by +`linear_attention_state_int8.yaml`; its saved decode policies are unsupported. + +The qualified scope is eager synchronous execution with TP=1/2 and PP=DP=CP=1. +Speculative decoding, prefix caching, state transfer, and CUDA graphs require +separate integration. This is floating-point numerical emulation; model-quality, +capacity, and performance claims require separate measurements. + +### Later prefill GEMM support + +Prefill operand QDQ will be a separate change after an optimized fused kernel +is available. It will reuse the eight numerical sites from the training prefill +implementation and preserve this native-cache wrapper. The PyTorch materialized +backend is not exposed in vLLM by this state-only adapter. State QDQ cadence and +prefill operand QDQ remain independent numerical policies. + ## Known Problems 1. **MCore reload does not use `MODELOPT_STATE_PATH`**; use `QUANT_FILE_PATH` and make sure `QUANT_CFG` matches the quantization recipe used for the original MCore model (otherwise quantizer keys/config won’t align). diff --git a/examples/vllm_serve/fakequant_worker.py b/examples/vllm_serve/fakequant_worker.py index a7d924c06d3..c78abe49985 100644 --- a/examples/vllm_serve/fakequant_worker.py +++ b/examples/vllm_serve/fakequant_worker.py @@ -37,6 +37,7 @@ disable_compilation, post_restore_vllm_parallel_linears, ) +from modelopt.torch.quantization.plugins.vllm_linear_attention import bind_vllm_linear_attention from modelopt.torch.utils import safe_load from modelopt.torch.utils.dataset_utils import get_dataset_dataloader @@ -114,17 +115,22 @@ def _fakequant_run_prolog_worker(self, mlflow_tracker: FakeQuantMlflowTracker) - print("Will load quant, so only do a single sample calibration") quant_config["calib_size"] = 1 - calib_dataloader = get_dataset_dataloader( - dataset_name=quant_config["dataset"], - tokenizer=tokenizer, - batch_size=quant_config["calib_batch_size"], - num_samples=quant_config["calib_size"], - device=self.device, - ) + quant_cfg = get_quant_config(quant_config, model) + calibrate_loop = None + if quant_cfg.get("algorithm", "max") is not None: + calib_dataloader = get_dataset_dataloader( + dataset_name=quant_config["dataset"], + tokenizer=tokenizer, + batch_size=quant_config["calib_batch_size"], + num_samples=quant_config["calib_size"], + device=self.device, + ) + worker_loop = calibrate_fun(calib_dataloader, self) - calibrate_loop = calibrate_fun(calib_dataloader, self) + def calibrate_loop(converted): + bind_vllm_linear_attention(converted, self.model_runner) + worker_loop(converted) - quant_cfg = get_quant_config(quant_config, model) # Before calibration, which is the run this artifact is most wanted for if it dies. mlflow_tracker.log_quant_config(quant_cfg) @@ -134,6 +140,7 @@ def _fakequant_run_prolog_worker(self, mlflow_tracker: FakeQuantMlflowTracker) - quantizer_file_path = quant_config["quant_file_path"] if quantizer_file_path: + bind_vllm_linear_attention(model, self.model_runner) self.model_runner._dummy_run(1) current_state_dict = load_state_dict_from_path(self, quantizer_file_path, model) model.load_state_dict(current_state_dict) @@ -142,6 +149,7 @@ def _fakequant_run_prolog_worker(self, mlflow_tracker: FakeQuantMlflowTracker) - if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1: torch.distributed.barrier() + bind_vllm_linear_attention(model, self.model_runner) if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: mtq.print_quant_summary(model) mlflow_tracker.log_quant_summary(model) diff --git a/examples/vllm_serve/linear_attention_state_int8.yaml b/examples/vllm_serve/linear_attention_state_int8.yaml new file mode 100644 index 00000000000..5ca88bfd959 --- /dev/null +++ b/examples/vllm_serve/linear_attention_state_int8.yaml @@ -0,0 +1,12 @@ +metadata: + recipe_type: ptq + description: GDN/KDA INT8 recurrent-state QDQ before native vLLM prefill and decode. +quantize: + 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: diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index 1bfc7f95fd4..7973d4e71bf 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -241,6 +241,19 @@ def _merge_values_require_identical(merged_key: str, key_value_pairs: list[tuple return first_value +def _map_linear_attention_names(values, map_fun): + """Map module paths using weight-path probes without dropping numerical policies.""" + if map_fun is None: + return values + probes = {(name + ".weight" if name else "weight"): value for name, value in values.items()} + mapped = map_fun(probes) + if len(mapped) != len(values) or any( + name != "weight" and not name.endswith(".weight") for name in mapped + ): + raise ValueError("HF-to-vLLM mapping must preserve every linear-attention module") + return {name.removesuffix("weight").removesuffix("."): value for name, value in mapped.items()} + + def convert_dict_to_vllm( state_dict: dict[str, Any], max_or_concat: bool = True, @@ -254,6 +267,15 @@ def convert_dict_to_vllm( max_or_concat: Whether to merge grouped values by taking max/concatenate or require identical map_fun: Function to map the state dict to vLLM format """ + state_dict = dict(state_dict) + linear_quantizers = {} + for key in list(state_dict): + match = re.match(r"^(.*)\.((?:gdn|kda)_(?:state|w)_quantizer(?:\..*)?)$", key) + if match: + module, suffix = match.groups() + mapped = _map_linear_attention_names({module: state_dict.pop(key)}, map_fun) + name, value = next(iter(mapped.items())) + linear_quantizers[f"{name}.{suffix}"] = value # If map_fun is provided, pre-transform quantizer key module-path prefixes so that # HF→vLLM model renames (e.g. backbone.layers → model.layers) are applied before # key grouping (q/k/v → qkv, experts.N.up_proj → experts.w13, etc.). @@ -287,7 +309,7 @@ def convert_dict_to_vllm( _, value = key_value_pairs[0] vllm_state_dict[merged_key] = value if map_fun is None: - return vllm_state_dict + return {**vllm_state_dict, **linear_quantizers} # Quantizer module-path keys (e.g. "layers.0.mlp.gate_proj.input_quantizer") must NOT # go through map_fun (hf_to_vllm_mapper.apply_dict), which maps weight tensor paths and # drops any key it doesn't recognise — including all quantizer keys. Split them out, @@ -295,7 +317,7 @@ def convert_dict_to_vllm( quantizer_keys = {k: v for k, v in vllm_state_dict.items() if "_quantizer" in k} non_quantizer_keys = {k: v for k, v in vllm_state_dict.items() if "_quantizer" not in k} mapped = map_fun(non_quantizer_keys) if non_quantizer_keys else {} - return {**mapped, **quantizer_keys} + return {**mapped, **quantizer_keys, **linear_quantizers} def convert_modelopt_state_to_vllm( @@ -320,6 +342,16 @@ def convert_modelopt_state_to_vllm( modelopt_state_dict = modelopt_state.pop("modelopt_state_dict", []) for idx, current_mode in enumerate(modelopt_state_dict): current_mode_metadata = current_mode[1].pop("metadata", {}) + if "linear_attention" in current_mode_metadata: + current_mode_metadata["linear_attention"] = _map_linear_attention_names( + current_mode_metadata["linear_attention"], map_fun + ) + # Use exact saved module names instead of training-framework selector patterns. + if "quant_cfg" in current_mode[1]["config"]: + current_mode[1]["config"]["linear_attention"] = [ + {"module_name": name, "cfg": policy} + for name, policy in current_mode_metadata["linear_attention"].items() + ] current_mode_quant_state = current_mode_metadata.pop("quantizer_state", {}) if current_mode_quant_state: current_mode_metadata["quantizer_state"] = convert_dict_to_vllm( @@ -367,6 +399,15 @@ def filter_modelopt_state_quantizer_state_for_model( metadata = mode_entry[1].get("metadata", {}) if "quantizer_state" in metadata: saved = metadata["quantizer_state"] + missing_linear = [ + name + for name, state in saved.items() + if re.search(r"(?:gdn|kda)_(?:state|w)_quantizer$", name) + and not state.get("_disabled", False) + and name not in model_keys + ] + if missing_linear: + raise ValueError(f"Saved linear-attention quantizers are missing: {missing_linear}") # Keep keys that exist in the model. Remove disabled quantizers UNLESS they # have registered buffers (e.g. _pre_quant_scale from AWQ/smoothquant on a diff --git a/modelopt/torch/kernels/quantization/common/fp8_quant.py b/modelopt/torch/kernels/quantization/common/fp8_quant.py index 2b76499b968..30c2d148629 100644 --- a/modelopt/torch/kernels/quantization/common/fp8_quant.py +++ b/modelopt/torch/kernels/quantization/common/fp8_quant.py @@ -16,11 +16,36 @@ """Composable Triton JIT functions for FP8 (E4M3) fake quantization. Counterpart of ``nvfp4_quant.py`` for per-tensor FP8. Used by the unified -flash-attention kernel's softmax-P qdq (``common/attention/triton_fa.py``). +flash-attention kernel's softmax-P qdq and linear-attention state/factor qdq. """ import triton import triton.language as tl +from triton.language.extra.cuda import libdevice + +__all__ = ["fp8_dynamic_qdq", "fp8_scalar_qdq"] + + +@triton.jit +def _fp8_cast(value, NATIVE_FP8: tl.constexpr): + value = tl.clamp(value, -448.0, 448.0) + if NATIVE_FP8: + return value.to(tl.float8e4nv).to(tl.float32) + # GPUs before SM89 need a software E4M3 conversion, including subnormals and ties. + bits = tl.abs(value).to(tl.int32, bitcast=True) + exponent = ((bits >> 23) & 255) - 127 + step = ((tl.maximum(exponent, -6) - 3 + 127) << 23).to(tl.float32, bitcast=True) + return libdevice.nearbyint(value / step) * step + + +@triton.jit +def fp8_dynamic_qdq(value, NATIVE_FP8: tl.constexpr = True): + """Dynamic per-tile FP8 QDQ with TensorQuantizer scaling and returned decode scale.""" + amax = tl.max(tl.abs(value)) + safe_amax = tl.where(amax <= 2.0**-24, 1.0, amax) + quant_scale = tl.div_rn(448.0, safe_amax) + decode_scale = tl.div_rn(1.0, quant_scale) + return _fp8_cast(value * quant_scale, NATIVE_FP8) * decode_scale, decode_scale @triton.jit @@ -38,6 +63,4 @@ def fp8_scalar_qdq(x, scale): Returns: Fake-quantized tensor of the same shape as ``x``, in float32. """ - FP8_E4M3_MAX: tl.constexpr = 448.0 - x_scaled = tl.clamp(x / scale, -FP8_E4M3_MAX, FP8_E4M3_MAX) - return x_scaled.to(tl.float8e4nv).to(tl.float32) * scale + return _fp8_cast(x / scale, True) * scale diff --git a/modelopt/torch/kernels/quantization/linear_attention/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/__init__.py new file mode 100644 index 00000000000..3f003f403f8 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/__init__.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Linear-attention kernels for quantization. + +``fla_chunk_delta_h.py`` and ``fla_chunk_gated_delta_rule.py`` are adapted copies of the chunked +GatedDeltaNet kernels of `flash-linear-attention `_ +(``fla.ops.common.chunk_delta_h`` and ``fla.ops.gated_delta_rule.chunk``) that can fake-quantize the +recurrent state carried between chunks to FP8 (``state_qdq``). They still import the surrounding +fla operators, so ``fla-core==0.5.1`` and Triton must be installed to use +them. This package initializer does not import the kernels, so importing it needs neither. +""" diff --git a/modelopt/torch/kernels/quantization/linear_attention/decode.py b/modelopt/torch/kernels/quantization/linear_attention/decode.py new file mode 100644 index 00000000000..a156f824b12 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/decode.py @@ -0,0 +1,468 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fused FP32 token/replay training recurrence with state checkpoint recomputation.""" + +import torch +import triton +import triton.language as tl +from torch.autograd.function import once_differentiable +from triton.language.extra.cuda import libdevice + +from ..common.fp8_quant import fp8_dynamic_qdq +from .int8 import int8_scalar_qdq + +__all__ = ["fused_recurrence"] + + +@triton.jit +def _sum_keys(value): + BK: tl.constexpr = value.shape[0] + rows = tl.arange(0, BK) + for level in tl.static_range(0, tl.constexpr(triton.next_power_of_2(BK).bit_length() - 1)): + offset = BK >> (level + 1) + value = value + tl.gather( + value, tl.broadcast_to((rows ^ offset)[:, None], value.shape), axis=0 + ) + return tl.sum(tl.where(rows[:, None] == 0, value, 0.0), axis=0) + + +@triton.jit +def _advance( + state, + key, + value, + gate, + beta, + write_state, + CHANNEL_GATE: tl.constexpr, + STATE_QDQ: tl.constexpr, + FACTOR_QDQ: tl.constexpr, + NATIVE_FP8: tl.constexpr, +): + decay = libdevice.exp(gate) + if CHANNEL_GATE: + decayed = state * decay[:, None] + else: + decayed = state * decay + residual = value - _sum_keys(key[:, None] * decayed) + update = beta * residual + update_scale = 1.0 + if FACTOR_QDQ: + update, update_scale = fp8_dynamic_qdq(update, NATIVE_FP8) + working = decayed + key[:, None] * update[None, :] + stored = working + if STATE_QDQ == 3: + state_scale = tl.full((working.shape[0], working.shape[1] // 32), 1.0, tl.float16) + else: + state_scale = 1.0 + if STATE_QDQ: + if write_state: + if STATE_QDQ == 3: + groups = tl.reshape(working, (working.shape[0], working.shape[1] // 32, 32)) + scale = tl.maximum(tl.max(tl.abs(groups), 2) * (1.0 / 127.0), 6e-8) + normalized = tl.div_rn(groups, scale[:, :, None]) + codes = tl.floor(tl.abs(normalized) + 0.5) + codes = tl.where(normalized < 0, -codes, codes) + state_scale = scale.to(tl.float16) + stored = tl.reshape( + tl.minimum(tl.maximum(codes, -127.0), 127.0) + * state_scale.to(tl.float32)[:, :, None], + working.shape, + ) + elif STATE_QDQ == 2: + amax = tl.max(tl.abs(working)) + # Match Torch scalar division: FP32 multiply by the rounded reciprocal. + state_scale = tl.where(amax > 0, amax * (1.0 / 127.0), 1.0) + stored = int8_scalar_qdq(working, state_scale) + else: + stored, state_scale = fp8_dynamic_qdq(working, NATIVE_FP8) + return stored, working, decayed, residual, update, update_scale, state_scale + + +@triton.jit +def _decode_fwd( + Q, + K, + V, + G, + Beta, + Initial, + Out, + Final, + Updates, + Anchor, + AnchorScale, + UpdateScale, + Checkpoints, + T: tl.constexpr, + H: tl.constexpr, + DK: tl.constexpr, + DV: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NV: tl.constexpr, + CHANNEL_GATE: tl.constexpr, + STATE_QDQ: tl.constexpr, + REPLAY: tl.constexpr, + FACTOR_QDQ: tl.constexpr, + WINDOW: tl.constexpr, + CURSOR: tl.constexpr, + READ_STORED: tl.constexpr, + SCALE: tl.constexpr, + INTERVAL: tl.constexpr, + NATIVE_FP8: tl.constexpr, +): + h = tl.program_id(0) + tile = tl.program_id(1) + keys = tl.arange(0, BK) + values = tile * BV + tl.arange(0, BV) + offsets = h * DK * DV + keys[:, None] * DV + values[None, :] + mask = (keys[:, None] < DK) & (values[None, :] < DV) + state = tl.load(Initial + offsets, mask, 0) + tl.store(Anchor + offsets, state, mask) + if STATE_QDQ == 3: + groups = tile * (BV // 32) + tl.arange(0, BV // 32) + scale_offsets = (h * DK + keys[:, None]) * (DV // 32) + groups[None, :] + scale_mask = (keys[:, None] < DK) & (groups[None, :] < DV // 32) + tl.store(AnchorScale + scale_offsets, 1.0, scale_mask) + else: + tl.store(AnchorScale + h * NV + tile, 1.0) + for t in range(T): + if t % INTERVAL == 0: + tl.store(Checkpoints + (t // INTERVAL) * H * DK * DV + offsets, state, mask) + query = tl.load(Q + (t * H + h) * DK + keys, keys < DK, 0) + key = tl.load(K + (t * H + h) * DK + keys, keys < DK, 0) + value = tl.load(V + (t * H + h) * DV + values, values < DV, 0) + beta = tl.load(Beta + t * H + h) + if CHANNEL_GATE: + gate = tl.load(G + (t * H + h) * DK + keys, keys < DK, 0) + else: + gate = tl.load(G + t * H + h) + write_state = not REPLAY or (CURSOR + t + 1) % WINDOW == 0 + state, working, decayed, residual, update, update_scale, state_scale = _advance( + state, + key, + value, + gate, + beta, + write_state, + CHANNEL_GATE, + STATE_QDQ, + FACTOR_QDQ, + NATIVE_FP8, + ) + read = state if READ_STORED else working + output = _sum_keys(query[:, None] * read) * SCALE + tl.store(Out + (t * H + h) * DV + values, output, values < DV) + tl.store(Updates + (t * H + h) * DV + values, update, values < DV) + tl.store(UpdateScale + (t * H + h) * NV + tile, update_scale) + if write_state: + tl.store(Anchor + offsets, state, mask) + if STATE_QDQ == 3: + tl.store(AnchorScale + scale_offsets, state_scale, scale_mask) + else: + tl.store(AnchorScale + h * NV + tile, state_scale) + tl.store(Final + offsets, state, mask) + + +@triton.jit +def _decode_bwd( + Q, + K, + V, + G, + Beta, + Checkpoints, + DOut, + DFinal, + DUpdates, + DAnchor, + DQ, + DKOut, + DVOut, + DG, + DBeta, + DInitial, + T: tl.constexpr, + H: tl.constexpr, + DK: tl.constexpr, + DV: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NV: tl.constexpr, + CHANNEL_GATE: tl.constexpr, + STATE_QDQ: tl.constexpr, + REPLAY: tl.constexpr, + FACTOR_QDQ: tl.constexpr, + WINDOW: tl.constexpr, + CURSOR: tl.constexpr, + READ_STORED: tl.constexpr, + SCALE: tl.constexpr, + INTERVAL: tl.constexpr, + NATIVE_FP8: tl.constexpr, +): + h = tl.program_id(0) + tile = tl.program_id(1) + keys = tl.arange(0, BK) + values = tile * BV + tl.arange(0, BV) + offsets = h * DK * DV + keys[:, None] * DV + values[None, :] + mask = (keys[:, None] < DK) & (values[None, :] < DV) + adjoint = tl.load(DFinal + offsets, mask, 0) + LAST_REFRESH: tl.constexpr = ((T + CURSOR) // WINDOW) * WINDOW - CURSOR - 1 + for t in range(T - 1, -1, -1): + start = (t // INTERVAL) * INTERVAL + state = tl.load(Checkpoints + (t // INTERVAL) * H * DK * DV + offsets, mask, 0) + for j in range(start, t): + key = tl.load(K + (j * H + h) * DK + keys, keys < DK, 0) + value = tl.load(V + (j * H + h) * DV + values, values < DV, 0) + beta = tl.load(Beta + j * H + h) + if CHANNEL_GATE: + gate = tl.load(G + (j * H + h) * DK + keys, keys < DK, 0) + else: + gate = tl.load(G + j * H + h) + write_state = not REPLAY or (CURSOR + j + 1) % WINDOW == 0 + state, _, _, _, _, _, _ = _advance( + state, + key, + value, + gate, + beta, + write_state, + CHANNEL_GATE, + STATE_QDQ, + FACTOR_QDQ, + NATIVE_FP8, + ) + query = tl.load(Q + (t * H + h) * DK + keys, keys < DK, 0) + key = tl.load(K + (t * H + h) * DK + keys, keys < DK, 0) + value = tl.load(V + (t * H + h) * DV + values, values < DV, 0) + beta = tl.load(Beta + t * H + h) + if CHANNEL_GATE: + gate = tl.load(G + (t * H + h) * DK + keys, keys < DK, 0) + else: + gate = tl.load(G + t * H + h) + write_state = not REPLAY or (CURSOR + t + 1) % WINDOW == 0 + stored, working, decayed, residual, update, _, _ = _advance( + state, + key, + value, + gate, + beta, + write_state, + CHANNEL_GATE, + STATE_QDQ, + FACTOR_QDQ, + NATIVE_FP8, + ) + if REPLAY: + if t == LAST_REFRESH: + adjoint += tl.load(DAnchor + offsets, mask, 0) + elif t == T - 1: + adjoint += tl.load(DAnchor + offsets, mask, 0) + doutput = tl.load(DOut + (t * H + h) * DV + values, values < DV, 0) + read = stored if READ_STORED else working + dquery = tl.sum(read * doutput[None, :], 1) * SCALE + dstate = adjoint + query[:, None] * (doutput[None, :] * SCALE) + dupdate = tl.sum(key[:, None] * dstate, 0) + dupdate += tl.load(DUpdates + (t * H + h) * DV + values, values < DV, 0) + dbeta = tl.sum(dupdate * residual, 0) + dvalue = beta * dupdate + dkey = tl.sum(update[None, :] * dstate - decayed * dvalue[None, :], 1) + ddecayed = dstate - key[:, None] * dvalue[None, :] + dgate = tl.sum(ddecayed * decayed, 1) + decay = libdevice.exp(gate) + if CHANNEL_GATE: + adjoint = ddecayed * decay[:, None] + tl.store(DG + ((tile * T + t) * H + h) * DK + keys, dgate, keys < DK) + else: + adjoint = ddecayed * decay + tl.store(DG + (tile * T + t) * H + h, tl.sum(dgate, 0)) + tl.store(DQ + ((tile * T + t) * H + h) * DK + keys, dquery, keys < DK) + tl.store(DKOut + ((tile * T + t) * H + h) * DK + keys, dkey, keys < DK) + tl.store(DVOut + (t * H + h) * DV + values, dvalue, values < DV) + tl.store(DBeta + (tile * T + t) * H + h, dbeta) + tl.store(DInitial + offsets, adjoint, mask) + + +class _FusedRecurrence(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q, + k, + v, + g, + beta, + initial, + state_qdq, + block_v, + replay, + factor_qdq, + window, + cursor, + read_stored, + scale, + interval, + ): + length, heads, keys = q.shape + values = v.shape[-1] + tiles = triton.cdiv(values, block_v) + checkpoints = torch.empty( + (triton.cdiv(length, interval), heads, keys, values), device=q.device, dtype=q.dtype + ) + output, updates = torch.empty_like(v), torch.empty_like(v) + final, anchor = torch.empty_like(initial), torch.empty_like(initial) + scale_shape = (heads, keys, values // 32) if state_qdq == 3 else (heads, tiles) + anchor_scales = torch.empty( + scale_shape, device=q.device, dtype=torch.float16 if state_qdq == 3 else q.dtype + ) + update_scales = torch.empty((length, heads, tiles), device=q.device, dtype=q.dtype) + kwargs = { + "T": length, + "H": heads, + "DK": keys, + "DV": values, + "BK": triton.next_power_of_2(keys), + "BV": block_v, + "NV": tiles, + "CHANNEL_GATE": g.ndim == 3, + "STATE_QDQ": state_qdq, + "REPLAY": replay, + "FACTOR_QDQ": factor_qdq, + "WINDOW": window, + "CURSOR": cursor, + "READ_STORED": read_stored, + "SCALE": scale, + "INTERVAL": interval, + "NATIVE_FP8": torch.cuda.get_device_capability(q.device) >= (8, 9), + "num_warps": 4 if triton.next_power_of_2(keys) * block_v <= 4096 else 8, + "enable_fp_fusion": False, + } + _decode_fwd[(heads, tiles)]( + q, + k, + v, + g, + beta, + initial, + output, + final, + updates, + anchor, + anchor_scales, + update_scales, + checkpoints, + **kwargs, + ) + ctx.save_for_backward(q, k, v, g, beta, initial, checkpoints) + ctx.kwargs = kwargs + ctx.mark_non_differentiable(anchor_scales, update_scales) + return output, final, updates, anchor, anchor_scales, update_scales + + @staticmethod + @once_differentiable + def backward(ctx, doutput, dfinal, dupdates, danchor, _danchor_scales, _dupdate_scales): + q, k, v, g, beta, initial, checkpoints = ctx.saved_tensors + tiles = ctx.kwargs["NV"] + dq = torch.empty((tiles, *q.shape), device=q.device, dtype=q.dtype) + dk = torch.empty_like(dq) + dv = torch.empty_like(v) + dg = torch.empty((tiles, *g.shape), device=q.device, dtype=q.dtype) + dbeta = torch.empty((tiles, *beta.shape), device=q.device, dtype=q.dtype) + dinitial = torch.empty_like(initial) + doutput = torch.zeros_like(v) if doutput is None else doutput.contiguous() + dfinal = torch.zeros_like(initial) if dfinal is None else dfinal.contiguous() + dupdates = torch.zeros_like(v) if dupdates is None else dupdates.contiguous() + danchor = torch.zeros_like(initial) if danchor is None else danchor.contiguous() + _decode_bwd[(q.shape[1], tiles)]( + q, + k, + v, + g, + beta, + checkpoints, + doutput, + dfinal, + dupdates, + danchor, + dq, + dk, + dv, + dg, + dbeta, + dinitial, + **ctx.kwargs, + ) + return dq.sum(0), dk.sum(0), dv, dg.sum(0), dbeta.sum(0), dinitial, *([None] * 9) + + +def fused_recurrence( + q, + k, + v, + g, + beta, + initial, + *, + state_qdq=False, + state_format="fp8_e4m3", + state_codec="tile", + block_v=16, + replay=False, + factor_qdq=False, + window=8, + cursor=0, + read_stored=True, + scale=None, + checkpoint_interval=8, +): + """Run a nonempty activated sequence; encoding keys/gates and carry metadata is external.""" + if not all( + x.is_cuda and x.dtype == torch.float32 and x.device == q.device + for x in (q, k, v, g, beta, initial) + ): + raise ValueError("Fused recurrence requires CUDA FP32 inputs on one device") + if q.ndim != 3 or q.shape != k.shape or v.shape[:2] != q.shape[:2] or beta.shape != q.shape[:2]: + raise ValueError("Expected aligned [T,H,D] q/k/v and [T,H] beta") + length, heads, keys = q.shape + if length == 0 or not 1 <= keys <= 128 or initial.shape != (heads, keys, v.shape[-1]): + raise ValueError("Require nonempty sequence, Dk<=128, and matching initial state") + if g.shape not in (q.shape, beta.shape) or block_v not in (16, 32, 64, 128): + raise ValueError("Invalid gate shape or state value-block grouping") + if not 1 <= window <= 64 or not 0 <= cursor < window or not 1 <= checkpoint_interval <= 32: + raise ValueError("Invalid replay cursor/window or recomputation interval") + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") + if state_codec not in ("tile", "int8_hadamard32"): + raise ValueError("State codec must be tile or int8_hadamard32") + if state_codec == "int8_hadamard32": + if (state_qdq and state_format != "int8") or v.shape[-1] % 32 or block_v < 32: + raise ValueError("Hadamard codec requires INT8, Dv divisible by 32 and block_v >= 32") + state_mode = 3 if state_codec == "int8_hadamard32" else (2 if state_format == "int8" else 1) + state_qdq = state_mode if state_qdq else 0 + scale = keys**-0.5 if scale is None else scale + return _FusedRecurrence.apply( + *(x.contiguous() for x in (q, k, v, g, beta, initial)), + state_qdq, + block_v, + replay, + factor_qdq, + window, + cursor, + read_stored, + scale, + checkpoint_interval, + ) diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py new file mode 100644 index 00000000000..84da262fd9f --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -0,0 +1,992 @@ +# Adapted from: https://github.com/fla-org/flash-linear-attention/blob/516143e31fce/fla/ops/common/chunk_delta_h.py +# Adapted with modifications (marked [ModelOpt]): optional in-kernel FP8 E4M3 / INT8 fake quantization +# of the carried chunk state (STATE_QDQ), BV as an explicit launch argument, no fla backend +# dispatch or config-cache autotune. +# +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import triton +import triton.language as tl +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.cache import fla_cache_autotune +from fla.ops.utils.op import exp2 +from fla.utils import ( + IS_INTEL, + IS_NVIDIA_BLACKWELL, + IS_NVIDIA_HOPPER, + autotune_cache_kwargs, + check_shared_mem, +) + +from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq + +from .int8 import int8_scalar_qdq + +# ``STATE_QDQ`` modes of the forward state kernel. +STATE_QDQ_OFF = 0 +STATE_QDQ_FP8_DYNAMIC = 1 # FP8 E4M3, one dynamic scale per program tile ([K, BV] of one head) +STATE_QDQ_INT8_DYNAMIC = 2 +STATE_QDQ_MAX_BLOCK_V = 128 + + +@triton.jit +def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr, STATE_QDQ: tl.constexpr): + """[ModelOpt] Dynamic scale per full [K, BV] tile of one sequence and head.""" + b_amax = tl.max(tl.abs(b_h1)) + if K > 64: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h2))) + if K > 128: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h3))) + if K > 192: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h4))) + if STATE_QDQ == 2: + return tl.where(b_amax > 0, b_amax * (1.0 / 127.0), 1.0) + else: + return tl.where(b_amax > 0, b_amax / 448.0, 1.0) + + +@triton.jit +def _state_scalar_qdq(value, scale, STATE_QDQ: tl.constexpr): + if STATE_QDQ == 2: + return int8_scalar_qdq(value, scale) + else: + return fp8_scalar_qdq(value, scale) + + +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] + +# TODO: Triton mainline fixes a Blackwell tl.dot recurrence race. +# Keep this kernel on num_warps=2 for Blackwell until Triton 3.8 is released +# and we re-validate the wider config space. +# Intel needs more warps than NVIDIA here: 8 warps is ~1.5x faster than the best +# config reachable under the [2, 4] cap. +if IS_NVIDIA_BLACKWELL: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2] +elif IS_INTEL: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2, 4, 8, 16] +else: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2, 4] + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +# ``BV`` is an explicit argument rather than an autotuned config: with ``STATE_QDQ`` it sets the +# quantization granularity, so it must not vary with the autotuner's choice. +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in GATED_DELTA_RULE_FWD_H_NUM_WARPS + for num_stages in ([2, 3, 4] if check_shared_mem("ampere") else [2, 1]) + ], + key=["H", "HV", "K", "V", "BT", "BV", "STATE_V_FIRST", "STATE_QDQ"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( + k, + v, + w, + v_new, + g, + gk, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + SAVE_NEW_VALUE: tl.constexpr, + STATE_V_FIRST: tl.constexpr, + IS_VARLEN: tl.constexpr, + STATE_QDQ: tl.constexpr, +): + pid = tl.program_id(0) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) + i_n, i_h = i_nh // HV, i_nh % HV + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + if STATE_V_FIRST: + b_h1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([BV, 64], dtype=tl.float32) + else: + b_h1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + h += (boh * HV + i_h).to(tl.int64) * K * V + v += (bos * HV + i_h).to(tl.int64) * V + k += (bos * H + i_h // (HV // H)).to(tl.int64) * K + w += (bos * HV + i_h).to(tl.int64) * K + if SAVE_NEW_VALUE: + v_new += (bos * HV + i_h).to(tl.int64) * V + + if USE_INITIAL_STATE: + h0 = h0 + i_nh * K * V + if STORE_FINAL_STATE: + ht = ht + i_nh * K * V + + # load initial state + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K + if USE_INITIAL_STATE: + if STATE_V_FIRST: + p_h0_1 = h0 + o_v[:, None] * K + o_k1[None, :] + m_h0_1 = m_v[:, None] & m_k1[None, :] + else: + p_h0_1 = h0 + o_k1[:, None] * V + o_v[None, :] + m_h0_1 = m_k1[:, None] & m_v[None, :] + b_h1 += tl.load(p_h0_1, mask=m_h0_1, other=0.0).to(tl.float32) + if K > 64: + if STATE_V_FIRST: + p_h0_2 = h0 + o_v[:, None] * K + o_k2[None, :] + m_h0_2 = m_v[:, None] & m_k2[None, :] + else: + p_h0_2 = h0 + o_k2[:, None] * V + o_v[None, :] + m_h0_2 = m_k2[:, None] & m_v[None, :] + b_h2 += tl.load(p_h0_2, mask=m_h0_2, other=0.0).to(tl.float32) + if K > 128: + if STATE_V_FIRST: + p_h0_3 = h0 + o_v[:, None] * K + o_k3[None, :] + m_h0_3 = m_v[:, None] & m_k3[None, :] + else: + p_h0_3 = h0 + o_k3[:, None] * V + o_v[None, :] + m_h0_3 = m_k3[:, None] & m_v[None, :] + b_h3 += tl.load(p_h0_3, mask=m_h0_3, other=0.0).to(tl.float32) + if K > 192: + if STATE_V_FIRST: + p_h0_4 = h0 + o_v[:, None] * K + o_k4[None, :] + m_h0_4 = m_v[:, None] & m_k4[None, :] + else: + p_h0_4 = h0 + o_k4[:, None] * V + o_v[None, :] + m_h0_4 = m_k4[:, None] & m_v[None, :] + b_h4 += tl.load(p_h0_4, mask=m_h0_4, other=0.0).to(tl.float32) + # [ModelOpt] A state read from a quantized cache is quantized before the first chunk uses it. + if STATE_QDQ != 0: + if K > 192: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K, STATE_QDQ=STATE_QDQ) + elif K > 128: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K, STATE_QDQ=STATE_QDQ) + elif K > 64: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K, STATE_QDQ=STATE_QDQ) + else: + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K, STATE_QDQ=STATE_QDQ) + b_h1 = _state_scalar_qdq(b_h1, b_scale, STATE_QDQ) + if K > 64: + b_h2 = _state_scalar_qdq(b_h2, b_scale, STATE_QDQ) + if K > 128: + b_h3 = _state_scalar_qdq(b_h3, b_scale, STATE_QDQ) + if K > 192: + b_h4 = _state_scalar_qdq(b_h4, b_scale, STATE_QDQ) + + # main recurrence + for i_t in range(NT): + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + if STATE_V_FIRST: + p_h1 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k1[None, :] + m_h1 = m_v[:, None] & m_k1[None, :] + else: + p_h1 = h + i_t_int64 * HV * K * V + o_k1[:, None] * V + o_v[None, :] + m_h1 = m_k1[:, None] & m_v[None, :] + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), mask=m_h1) + if K > 64: + if STATE_V_FIRST: + p_h2 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k2[None, :] + m_h2 = m_v[:, None] & m_k2[None, :] + else: + p_h2 = h + i_t_int64 * HV * K * V + o_k2[:, None] * V + o_v[None, :] + m_h2 = m_k2[:, None] & m_v[None, :] + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), mask=m_h2) + if K > 128: + if STATE_V_FIRST: + p_h3 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k3[None, :] + m_h3 = m_v[:, None] & m_k3[None, :] + else: + p_h3 = h + i_t_int64 * HV * K * V + o_k3[:, None] * V + o_v[None, :] + m_h3 = m_k3[:, None] & m_v[None, :] + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), mask=m_h3) + if K > 192: + if STATE_V_FIRST: + p_h4 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k4[None, :] + m_h4 = m_v[:, None] & m_k4[None, :] + else: + p_h4 = h + i_t_int64 * HV * K * V + o_k4[:, None] * V + o_v[None, :] + m_h4 = m_k4[:, None] & m_v[None, :] + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), mask=m_h4) + + p_w = w + o_t[:, None] * (HV * K) + o_k1[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k1[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) + else: + b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) + if K > 64: + p_w = w + o_t[:, None] * (HV * K) + o_k2[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k2[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h2.to(b_w.dtype), b_v) + if K > 128: + p_w = w + o_t[:, None] * (HV * K) + o_k3[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k3[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h3.to(b_w.dtype), b_v) + if K > 192: + p_w = w + o_t[:, None] * (HV * K) + o_k4[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k4[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h4.to(b_w.dtype), b_v) + p_v = v + o_t[:, None] * (HV * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & m_v[None, :], other=0.0) - b_v + + if SAVE_NEW_VALUE: + p_v = v_new + o_t[:, None] * (HV * V) + o_v[None, :] + tl.store(p_v, b_v.to(p_v.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) + + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + b_g_last = tl.load(g + (bos * HV + last_idx * HV + i_h).to(tl.int64)).to(tl.float32) + p_g = g + (bos * HV + i_h).to(tl.int64) + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] + b_g_last = exp2(b_g_last) + b_h1 *= b_g_last + if K > 64: + b_h2 *= b_g_last + if K > 128: + b_h3 *= b_g_last + if K > 192: + b_h4 *= b_g_last + + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k1, mask=(o_k1 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h1 *= exp2(b_gk_last1)[None, :] + else: + b_h1 *= exp2(b_gk_last1)[:, None] + if K > 64: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k2, mask=(o_k2 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h2 *= exp2(b_gk_last2)[None, :] + else: + b_h2 *= exp2(b_gk_last2)[:, None] + if K > 128: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k3, mask=(o_k3 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h3 *= exp2(b_gk_last3)[None, :] + else: + b_h3 *= exp2(b_gk_last3)[:, None] + if K > 192: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k4, mask=(o_k4 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h4 *= exp2(b_gk_last4)[None, :] + else: + b_h4 *= exp2(b_gk_last4)[:, None] + b_v = b_v.to(k.dtype.element_ty) + + p_k = k + o_k1[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k1[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h1 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h1 = tl.dot(b_k, b_v, b_h1) + if K > 64: + p_k = k + o_k2[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k2[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h2 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h2 = tl.dot(b_k, b_v, b_h2) + if K > 128: + p_k = k + o_k3[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k3[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h3 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h3 = tl.dot(b_k, b_v, b_h3) + if K > 192: + p_k = k + o_k4[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k4[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h4 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h4 = tl.dot(b_k, b_v, b_h4) + + # [ModelOpt] Dynamic per-tile FP8/INT8 QDQ of the next-chunk or stored final state. + # Recompute amax over [K, BV]; _state_scalar_qdq applies that tile's scalar scale. + if STATE_QDQ != 0: + if K > 192: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K, STATE_QDQ=STATE_QDQ) + elif K > 128: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K, STATE_QDQ=STATE_QDQ) + elif K > 64: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K, STATE_QDQ=STATE_QDQ) + else: + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K, STATE_QDQ=STATE_QDQ) + b_h1 = _state_scalar_qdq(b_h1, b_scale, STATE_QDQ) + if K > 64: + b_h2 = _state_scalar_qdq(b_h2, b_scale, STATE_QDQ) + if K > 128: + b_h3 = _state_scalar_qdq(b_h3, b_scale, STATE_QDQ) + if K > 192: + b_h4 = _state_scalar_qdq(b_h4, b_scale, STATE_QDQ) + + if STORE_FINAL_STATE: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k1[None, :] + m_ht = m_v[:, None] & m_k1[None, :] + else: + p_ht = ht + o_k1[:, None] * V + o_v[None, :] + m_ht = m_k1[:, None] & m_v[None, :] + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 64: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k2[None, :] + m_ht = m_v[:, None] & m_k2[None, :] + else: + p_ht = ht + o_k2[:, None] * V + o_v[None, :] + m_ht = m_k2[:, None] & m_v[None, :] + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 128: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k3[None, :] + m_ht = m_v[:, None] & m_k3[None, :] + else: + p_ht = ht + o_k3[:, None] * V + o_v[None, :] + m_ht = m_k3[:, None] & m_v[None, :] + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 192: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k4[None, :] + m_ht = m_v[:, None] & m_k4[None, :] + else: + p_ht = ht + o_k4[:, None] * V + o_v[None, :] + m_ht = m_k4[:, None] & m_v[None, :] + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), mask=m_ht) + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "USE_INITIAL_STATE": lambda args: args["dh0"] is not None, + "USE_FINAL_STATE_GRADIENT": lambda args: args["dht"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@fla_cache_autotune( + configs=[ + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in ([2, 3, 4] if check_shared_mem("ampere") else [1]) + for BV in ([32, 64] if check_shared_mem("ada") else [32]) + ], + key=["H", "HV", "K", "V", "BT", "BV", "USE_G", "STATE_V_FIRST"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( + q, + k, + w, + g, + gk, + dht, + dh0, + do, + dh, + dv, + dv2, + cu_seqlens, + chunk_offsets, + scale, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + STATE_V_FIRST: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + pid = tl.program_id(0) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) + i_n, i_h = i_nh // HV, i_nh % HV + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + if STATE_V_FIRST: + b_dh1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([BV, 64], dtype=tl.float32) + else: + b_dh1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + q += (bos * H + i_h // (HV // H)).to(tl.int64) * K + k += (bos * H + i_h // (HV // H)).to(tl.int64) * K + w += (bos * HV + i_h).to(tl.int64) * K + do += (bos * HV + i_h).to(tl.int64) * V + dv += (bos * HV + i_h).to(tl.int64) * V + dv2 += (bos * HV + i_h).to(tl.int64) * V + dh += (boh * HV + i_h).to(tl.int64) * K * V + if USE_GK: + gk += (bos * HV + i_h).to(tl.int64) * K + + if USE_INITIAL_STATE: + dh0 += i_nh * K * V + if USE_FINAL_STATE_GRADIENT: + dht += i_nh * K * V + + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K + if USE_FINAL_STATE_GRADIENT: + if STATE_V_FIRST: + p_dht1 = dht + o_v[:, None] * K + o_k1[None, :] + m_dht1 = m_v[:, None] & m_k1[None, :] + else: + p_dht1 = dht + o_k1[:, None] * V + o_v[None, :] + m_dht1 = m_k1[:, None] & m_v[None, :] + b_dh1 += tl.load(p_dht1, mask=m_dht1, other=0.0) + if K > 64: + if STATE_V_FIRST: + p_dht2 = dht + o_v[:, None] * K + o_k2[None, :] + m_dht2 = m_v[:, None] & m_k2[None, :] + else: + p_dht2 = dht + o_k2[:, None] * V + o_v[None, :] + m_dht2 = m_k2[:, None] & m_v[None, :] + b_dh2 += tl.load(p_dht2, mask=m_dht2, other=0.0) + if K > 128: + if STATE_V_FIRST: + p_dht3 = dht + o_v[:, None] * K + o_k3[None, :] + m_dht3 = m_v[:, None] & m_k3[None, :] + else: + p_dht3 = dht + o_k3[:, None] * V + o_v[None, :] + m_dht3 = m_k3[:, None] & m_v[None, :] + b_dh3 += tl.load(p_dht3, mask=m_dht3, other=0.0) + if K > 192: + if STATE_V_FIRST: + p_dht4 = dht + o_v[:, None] * K + o_k4[None, :] + m_dht4 = m_v[:, None] & m_k4[None, :] + else: + p_dht4 = dht + o_k4[:, None] * V + o_v[None, :] + m_dht4 = m_k4[:, None] & m_v[None, :] + b_dh4 += tl.load(p_dht4, mask=m_dht4, other=0.0) + + for i_t in range(NT - 1, -1, -1): + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + if STATE_V_FIRST: + p_dh1 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k1[None, :] + m_dh1 = m_v[:, None] & m_k1[None, :] + else: + p_dh1 = dh + i_t_int64 * HV * K * V + o_k1[:, None] * V + o_v[None, :] + m_dh1 = m_k1[:, None] & m_v[None, :] + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), mask=m_dh1) + if K > 64: + if STATE_V_FIRST: + p_dh2 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k2[None, :] + m_dh2 = m_v[:, None] & m_k2[None, :] + else: + p_dh2 = dh + i_t_int64 * HV * K * V + o_k2[:, None] * V + o_v[None, :] + m_dh2 = m_k2[:, None] & m_v[None, :] + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), mask=m_dh2) + if K > 128: + if STATE_V_FIRST: + p_dh3 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k3[None, :] + m_dh3 = m_v[:, None] & m_k3[None, :] + else: + p_dh3 = dh + i_t_int64 * HV * K * V + o_k3[:, None] * V + o_v[None, :] + m_dh3 = m_k3[:, None] & m_v[None, :] + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), mask=m_dh3) + if K > 192: + if STATE_V_FIRST: + p_dh4 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k4[None, :] + m_dh4 = m_v[:, None] & m_k4[None, :] + else: + p_dh4 = dh + i_t_int64 * HV * K * V + o_k4[:, None] * V + o_v[None, :] + m_dh4 = m_k4[:, None] & m_v[None, :] + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), mask=m_dh4) + + last_idx = min((i_t_int64 + 1) * BT, T) - 1 + if USE_G: + bg_last = tl.load(g + (bos + last_idx) * HV + i_h).to(tl.float32) + p_g = g + bos * HV + i_h + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + bg_last_exp = exp2(bg_last) + b_g_exp = exp2(b_g) + p_dv = dv + o_t[:, None] * (HV * V) + o_v[None, :] + p_dv2 = dv2 + o_t[:, None] * (HV * V) + o_v[None, :] + p_do = do + o_t[:, None] * (HV * V) + o_v[None, :] + + b_do = tl.load(p_do, mask=m_t[:, None] & m_v[None, :], other=0.0) + + # Update dv + p_k = k + o_t[:, None] * (H * K) + o_k1[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k1[None, :], other=0.0) + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load(gk + last_idx * HV * K + o_k1, mask=(o_k1 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh1).to(b_k.dtype)) + else: + b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) + + if K > 64: + p_k = k + o_t[:, None] * (H * K) + o_k2[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k2[None, :], other=0.0) + if USE_GK: + b_gk_last2 = tl.load(gk + last_idx * HV * K + o_k2, mask=(o_k2 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh2).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh2.to(b_k.dtype), b_dv) + + if K > 128: + p_k = k + o_t[:, None] * (H * K) + o_k3[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k3[None, :], other=0.0) + if USE_GK: + b_gk_last3 = tl.load(gk + last_idx * HV * K + o_k3, mask=(o_k3 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh3).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh3.to(b_k.dtype), b_dv) + + if K > 192: + p_k = k + o_t[:, None] * (H * K) + o_k4[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k4[None, :], other=0.0) + if USE_GK: + b_gk_last4 = tl.load(gk + last_idx * HV * K + o_k4, mask=(o_k4 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh4).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh4.to(b_k.dtype), b_dv) + + if USE_G: + b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] + b_dv += tl.load(p_dv, mask=m_t[:, None] & m_v[None, :], other=0.0) + + tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) + # Update dh + p_w = w + o_k1[:, None] + o_t[None, :] * (HV * K) + p_q = q + o_k1[:, None] + o_t[None, :] * (H * K) + b_w = tl.load(p_w, mask=m_k1[:, None] & m_t[None, :], other=0.0) + b_q = tl.load(p_q, mask=m_k1[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh1 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh1 *= exp2(b_gk_last1)[None, :] + else: + b_dh1 *= exp2(b_gk_last1[:, None]) + if STATE_V_FIRST: + b_dh1 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 64: + p_q = q + o_k2[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k2[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k2[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k2[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh2 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh2 *= exp2(b_gk_last2)[None, :] + else: + b_dh2 *= exp2(b_gk_last2[:, None]) + if STATE_V_FIRST: + b_dh2 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 128: + p_q = q + o_k3[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k3[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k3[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k3[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh3 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh3 *= exp2(b_gk_last3)[None, :] + else: + b_dh3 *= exp2(b_gk_last3[:, None]) + if STATE_V_FIRST: + b_dh3 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 192: + p_q = q + o_k4[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k4[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k4[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k4[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh4 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh4 *= exp2(b_gk_last4)[None, :] + else: + b_dh4 *= exp2(b_gk_last4[:, None]) + if STATE_V_FIRST: + b_dh4 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + + if USE_INITIAL_STATE: + if STATE_V_FIRST: + p_dh0 = dh0 + o_v[:, None] * K + o_k1[None, :] + m_dh0 = m_v[:, None] & m_k1[None, :] + else: + p_dh0 = dh0 + o_k1[:, None] * V + o_v[None, :] + m_dh0 = m_k1[:, None] & m_v[None, :] + tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), mask=m_dh0) + if K > 64: + if STATE_V_FIRST: + p_dh1 = dh0 + o_v[:, None] * K + o_k2[None, :] + m_dh1 = m_v[:, None] & m_k2[None, :] + else: + p_dh1 = dh0 + o_k2[:, None] * V + o_v[None, :] + m_dh1 = m_k2[:, None] & m_v[None, :] + tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), mask=m_dh1) + if K > 128: + if STATE_V_FIRST: + p_dh2 = dh0 + o_v[:, None] * K + o_k3[None, :] + m_dh2 = m_v[:, None] & m_k3[None, :] + else: + p_dh2 = dh0 + o_k3[:, None] * V + o_v[None, :] + m_dh2 = m_k3[:, None] & m_v[None, :] + tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), mask=m_dh2) + if K > 192: + if STATE_V_FIRST: + p_dh3 = dh0 + o_v[:, None] * K + o_k4[None, :] + m_dh3 = m_v[:, None] & m_k4[None, :] + else: + p_dh3 = dh0 + o_k4[:, None] * V + o_v[None, :] + m_dh3 = m_k4[:, None] & m_v[None, :] + tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), mask=m_dh3) + + +def chunk_gated_delta_rule_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + B, T, H, K, V, HV = *k.shape, u.shape[-1], u.shape[2] + BT = chunk_size + BV = state_qdq_tile_v(V, state_qdq, state_qdq_block_v) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + if chunk_offsets is None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + assert K <= 256, "current kernel does not support head dimension larger than 256." + + if state_v_first: + h = k.new_empty(B, NT, HV, V, K) + final_state = k.new_zeros(N, HV, V, K, dtype=torch.float32) if output_final_state else None + else: + h = k.new_empty(B, NT, HV, K, V) + final_state = k.new_zeros(N, HV, K, V, dtype=torch.float32) if output_final_state else None + + v_new = torch.empty_like(u) if save_new_value else None + + grid = (triton.cdiv(V, BV) * N * HV,) + chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( + k=k, + v=u, + w=w, + v_new=v_new, + g=g, + gk=gk, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BV=BV, + STATE_V_FIRST=state_v_first, + STATE_QDQ=state_qdq, + ) + return h, v_new, final_state + + +def state_qdq_tile_v(V: int, state_qdq: int, state_qdq_block_v: int | None) -> int: + """Return the V tile width ``BV`` of the forward state kernel. + + Without state quantization this is fla's largest tile. With it, the tile is also the + quantization granularity: one dynamic scale per ``[K, BV]`` block of a head's state. The + default is fla's 64-column tile, i.e. one scale per sequence and head for ``V <= 64`` and two + for the usual ``V == 128``. ``state_qdq_block_v=128`` gives one scale per 128-wide head but + exceeds the register budget where the kernel is limited to two warps (Blackwell) and spills. + """ + if state_qdq == STATE_QDQ_OFF: + return 64 if check_shared_mem("ada") else 32 + if state_qdq not in (STATE_QDQ_FP8_DYNAMIC, STATE_QDQ_INT8_DYNAMIC): + raise ValueError(f"Unsupported state_qdq mode {state_qdq}; expected 0, 1, or 2.") + BV = min(triton.next_power_of_2(V), 64) if state_qdq_block_v is None else state_qdq_block_v + if BV < 16 or BV > STATE_QDQ_MAX_BLOCK_V or BV & (BV - 1): + raise ValueError( + f"state_qdq_block_v must be a power of two in [16, {STATE_QDQ_MAX_BLOCK_V}], got {BV}." + ) + return BV + + +def chunk_gated_delta_rule_bwd_dhu( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + do: torch.Tensor, + dv: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + h0: torch.Tensor | None = None, + dht: torch.Tensor | None = None, + scale: float | None = None, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, + use_graph: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V, HV = *q.shape, do.shape[-1], do.shape[2] + # N: the actual number of sequences in the batch with either equal or variable lengths + BT = chunk_size + assert K <= 256, "current kernel does not support head dimension being larger than 256." + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + if chunk_offsets is None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + + if use_graph: + # [ModelOpt] Imported here: fla.ops.utils.graph is absent from released fla versions. + from fla.ops.utils.graph import get_static_buffer + + if state_v_first: + dh = get_static_buffer("dhu_dh_vf", (B, NT, HV, V, K), q.dtype, q.device) + else: + dh = get_static_buffer("dhu_dh", (B, NT, HV, K, V), q.dtype, q.device) + dh0 = ( + get_static_buffer("dhu_dh0", tuple(h0.shape), torch.float32, h0.device) + if h0 is not None + else None + ) + dv2 = get_static_buffer("dhu_dv2", tuple(dv.shape), dv.dtype, dv.device) + else: + if state_v_first: + dh = q.new_empty(B, NT, HV, V, K) + else: + dh = q.new_empty(B, NT, HV, K, V) + dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None + dv2 = torch.empty_like(dv) + + def grid(meta): + return (triton.cdiv(V, meta["BV"]) * N * HV,) + + chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64[grid]( + q=q, + k=k, + w=w, + g=g, + gk=gk, + dht=dht, + dh0=dh0, + do=do, + dh=dh, + dv=dv, + dv2=dv2, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + scale=scale, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + STATE_V_FIRST=state_v_first, + ) + return dh, dh0, dv2 diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py new file mode 100644 index 00000000000..f34b8a3384b --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py @@ -0,0 +1,725 @@ +# Adapted from: https://github.com/fla-org/flash-linear-attention/blob/516143e31fce/fla/ops/gated_delta_rule/chunk.py +# Adapted with modifications (marked [ModelOpt]): threads state_qdq / state_qdq_block_v through +# the autograd function, applies an optional w_quantizer to the WY tensor w, and imports the +# state kernels from the vendored sibling module. +# +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import warnings + +import fla +import torch +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o +from fla.ops.common.gate import fused_beta_sigmoid, fused_beta_sigmoid_bwd +from fla.ops.cp import FLACPContext +from fla.ops.cp.chunk_delta_h import ( + chunk_gated_delta_rule_bwd_dhu_pre_process, + chunk_gated_delta_rule_fwd_h_pre_process, + compress_h0, + expand_h0, +) +from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra +from fla.ops.gated_delta_rule.gate import gdn_gate_bwd, gdn_gate_chunk_cumsum +from fla.ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd +from fla.ops.utils import chunk_local_cumsum +from fla.ops.utils.constant import RCP_LN2 +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import ( + IS_NVIDIA_HOPPER, + TRITON_ABOVE_3_4_0, + autocast_custom_bwd, + autocast_custom_fwd, + input_guard, +) + +from modelopt.torch.quantization.linear_attention.utils import validate_gdn_quantizer +from modelopt.torch.quantization.nn import TensorQuantizer + +from .fla_chunk_delta_h import ( + STATE_QDQ_FP8_DYNAMIC, + STATE_QDQ_INT8_DYNAMIC, + STATE_QDQ_OFF, + chunk_gated_delta_rule_bwd_dhu, + chunk_gated_delta_rule_fwd_h, +) + + +def chunk_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + w_quantizer: TensorQuantizer | None = None, +): + g_input = g if use_gate_in_kernel else None + if use_gate_in_kernel: + g = gdn_gate_chunk_cumsum( + g=g, + A_log=A_log, + chunk_size=chunk_size, + scale=RCP_LN2, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + else: + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + scale=RCP_LN2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # obtain WY representation. u is actually the new v. + # fused kkt + solve_tril + recompute_w_u + w, u, A = chunk_gated_delta_rule_fwd_intra( + k=k, + v=v, + g=g, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + # [ModelOpt] w is an activation (the WY form of the chunk's keys) that lands in memory here, + # so it is fake-quantized once per forward instead of tile by tile inside the kernel. + if w_quantizer is not None: + w = w_quantizer(w) + + if cp_context is not None: + initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=k, + w=w, + u=u, + g=g, + cu_seqlens=cu_seqlens, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + ) + + if cp_context is not None: + initial_state = compress_h0(initial_state, context=cp_context) + + o = chunk_fwd_o( + q=q, + k=k, + v=v_new, + h=h, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + return g, o, A, final_state, initial_state, g_input, w if w_quantizer is not None else None + + +def chunk_gated_delta_rule_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + g_input: torch.Tensor | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + quantized_w: torch.Tensor | None = None, +): + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # [ModelOpt] Reuse forward QDQ exactly; the validated policy uses identity STE. + # This also avoids calling observers or dynamic scale computation in backward. + if quantized_w is not None: + w = quantized_w + + if cp_context is not None: + initial_state = expand_h0(initial_state, context=cp_context) + + # [ModelOpt] The backward recomputes the forward's chunk states, so it sees the same + # fake-quantized states; the state gradient itself passes straight through the QDQ. + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + ) + dv = chunk_bwd_dv_local( + q=q, + k=k, + g=g, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + if cp_context is not None: + # initial_state is None in the CP mode + # We only need to compute dht of current rank and pass it to the backward kernel + dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=q, + k=k, + w=w, + do=do, + dv=dv, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + dht=dht, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=q, + k=k, + w=w, + g=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dq, dk, dw, dg = chunk_bwd_dqkwg( + q=q, + k=k, + v=v_new, + w=w, + g=g, + h=h, + dv=dv, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dk2, dv, db, dg2 = prepare_wy_repr_bwd( + k=k, + v=v, + beta=beta, + g=g, + A=A, + dw=dw, + du=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + dk.add_(dk2) + dg.add_(dg2) + dg = chunk_local_cumsum( + dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices + ) + dA_log, ddt_bias = None, None + if use_gate_in_kernel: + dg, dA_log, ddt_bias = gdn_gate_bwd(g=g_input, A_log=A_log, dt_bias=dt_bias, dyg=dg) + return dq, dk, dv, db, dg, dh0, dA_log, ddt_bias + + +class ChunkGatedDeltaRuleFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + cp_context: FLACPContext | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + w_quantizer: TensorQuantizer | None = None, + ): + q_rstd, k_rstd = None, None + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + + beta_raw = beta + if use_beta_sigmoid_in_kernel: + beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu + ) + g, o, A, final_state, initial_state, g_input, quantized_w = chunk_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cp_context=cp_context, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + use_gate_in_kernel=use_gate_in_kernel, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + w_quantizer=w_quantizer, + ) + ctx.save_for_backward( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + quantized_w, + ) + ctx.scale = scale + ctx.chunk_size = chunk_size + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.use_beta_sigmoid_in_kernel = use_beta_sigmoid_in_kernel + ctx.allow_neg_eigval = allow_neg_eigval + ctx.cp_context = cp_context + ctx.state_v_first = state_v_first + ctx.use_gate_in_kernel = use_gate_in_kernel + ctx.state_qdq = state_qdq + ctx.state_qdq_block_v = state_qdq_block_v + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + ( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + quantized_w, + ) = ctx.saved_tensors + dq, dk, dv, db, dg, dh0, dA_log, ddt_bias = chunk_gated_delta_rule_bwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A=A, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + cp_context=ctx.cp_context, + chunk_indices=chunk_indices, + state_v_first=ctx.state_v_first, + use_gate_in_kernel=ctx.use_gate_in_kernel, + g_input=g_input, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=ctx.chunk_size, + state_qdq=ctx.state_qdq, + state_qdq_block_v=ctx.state_qdq_block_v, + quantized_w=quantized_w, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + if ctx.use_beta_sigmoid_in_kernel: + db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0) + return ( + dq.to(q), + dk.to(k), + dv.to(v), + dg.to(g), + db.to(beta_raw), + None, + dh0, + None, + None, + None, + None, + None, + None, + None, + dA_log, + ddt_bias, + None, + None, + None, + None, + None, + None, + None, + ) + + +# [ModelOpt] Not registered with fla's backend dispatch: another backend must not take over a +# call that asks for state quantization. +@torch.compiler.disable +def chunk_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + **kwargs, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`. + g (torch.Tensor): + (forget) gating tensor of shape `[B, T, HV]`. + When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay). + When `use_gate_in_kernel=True`, `g` is the raw input before gate activation; + the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, HV, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2norm to the q/k tensor internally. Default: `False`. + use_gate_in_kernel (bool): + Whether to compute the log-space GDN decay internally. + When `True`, the passed `g` is the raw input, and `A_log` must be provided. + The kernel fuses gate activation + chunk cumsum in a single pass. + Default: `False`. + A_log (Optional[torch.Tensor]): + Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`. + dt_bias (Optional[torch.Tensor]): + Bias added to `g` before activation, of shape `[HV]`. + Only used when `use_gate_in_kernel=True`. + use_beta_sigmoid_in_kernel (bool): + Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel. + - If `True`, the passed `beta` acts as the raw beta logits. + - If `False`, `beta` is expected to already be in post-sigmoid space. + Default: `False`. + allow_neg_eigval (bool): + Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`. + Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case + the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`. + state_v_first (Optional[bool]): + Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + chunk_indices (Optional[torch.LongTensor]): + Pre-computed chunk indices for variable-length inputs. + If provided, they are used directly instead of being computed from `cu_seqlens`. Default: `None`. + cp_context (Optional[FLACPContext]): + Context parallel context for distributed training across multiple devices. + When provided, `initial_state` and `output_final_state` are not supported, + and `cu_seqlens` will be overridden by the context. Default: `None`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda') + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if fla.__version__ != "0.5.1": + raise RuntimeError(f"ModelOpt GDN requires fla-core==0.5.1, got {fla.__version__}.") + if "transpose_state_layout" in kwargs: + if state_v_first: + raise ValueError( + "Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`." + ) + warnings.warn( + "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", + DeprecationWarning, + stacklevel=2, + ) + state_v_first = kwargs.pop("transpose_state_layout") + + # Validate head dimensions + if q.shape[2] != k.shape[2]: + raise ValueError( + f"q and k must have the same number of heads, " + f"but got q.shape[2]={q.shape[2]} and k.shape[2]={k.shape[2]}" + ) + H, HV = q.shape[2], v.shape[2] + if HV % H != 0: + raise ValueError( + f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by " + f"num_heads (H={H}), but got HV % H = {HV % H}" + ) + + if "head_first" in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) + + chunk_size = kwargs.pop("chunk_size", 64) + if chunk_size != 64: + raise ValueError("ModelOpt GDN supports only chunk_size=64; FLA WY backward assumes 64.") + + # [ModelOpt] state_qdq: 0 keeps fla's numerics; 1 fake-quantizes the state carried between + # chunks to FP8 E4M3; 2 uses signed narrow-range INT8. Both use a dynamic scale per + # [K, state_qdq_block_v] tile of each head. + state_qdq = kwargs.pop("state_qdq", STATE_QDQ_OFF) + state_qdq_block_v = kwargs.pop("state_qdq_block_v", None) + # w_quantizer: dynamic FP8 TensorQuantizer applied to the WY tensor + # ``w`` of shape [B, T, HV, K] before it multiplies the state, emulating an FP8 x FP8 matmul. + w_quantizer = kwargs.pop("w_quantizer", None) + use_gate_in_kernel = kwargs.pop("use_gate_in_kernel", False) + A_log = kwargs.pop("A_log", None) + dt_bias = kwargs.pop("dt_bias", None) + if kwargs: + raise TypeError(f"Unexpected keyword arguments: {', '.join(sorted(kwargs))}") + if state_qdq not in (STATE_QDQ_OFF, STATE_QDQ_FP8_DYNAMIC, STATE_QDQ_INT8_DYNAMIC): + raise ValueError(f"`state_qdq` must be 0, 1, or 2, got {state_qdq}.") + if w_quantizer is not None: + validate_gdn_quantizer(w_quantizer, name="gdn_w_quantizer") + if state_qdq == STATE_QDQ_FP8_DYNAMIC and ( + not q.is_cuda or torch.cuda.get_device_capability(q.device) < (8, 9) + ): + raise RuntimeError("GDN state QDQ requires native E4M3 conversion on CUDA SM89 or newer.") + if (state_qdq != STATE_QDQ_OFF or w_quantizer is not None) and cp_context is not None: + raise ValueError("State or w quantization is not supported together with `cp_context`.") + + if cp_context is not None: + assert initial_state is None, "Initial state is not supported for CP" + assert output_final_state is False, "Output final state is not supported for CP" + assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" + cu_seqlens = cp_context.cu_seqlens + if cp_context.cu_seqlens_cpu is not None: + cu_seqlens_cpu = cp_context.cu_seqlens_cpu + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if use_gate_in_kernel: + assert A_log is not None, "A_log must be provided when use_gate_in_kernel=True." + if allow_neg_eigval and not use_beta_sigmoid_in_kernel: + raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") + + if scale is None: + scale = k.shape[-1] ** -0.5 + # [ModelOpt] Hopper's TileLang backward needs BF16 and equal head counts. Expand outside + # custom autograd so repeat_interleave reduces q/k gradients back to the original heads. + if IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0: + if any(x.dtype != torch.bfloat16 for x in (q, k, v)): + raise ValueError("Hopper with Triton >= 3.4 requires BF16 q/k/v for GDN training.") + if H != HV: + q = q.repeat_interleave(HV // H, dim=2) + k = k.repeat_interleave(HV // H, dim=2) + o, final_state = ChunkGatedDeltaRuleFunction.apply( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + state_v_first, + cu_seqlens, + cu_seqlens_cpu, + chunk_indices, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + A_log, + dt_bias, + use_beta_sigmoid_in_kernel, + allow_neg_eigval, + cp_context, + chunk_size, + state_qdq, + state_qdq_block_v, + w_quantizer, + ) + return o, final_state + + +chunk_gdn = chunk_gated_delta_rule diff --git a/modelopt/torch/kernels/quantization/linear_attention/int8.py b/modelopt/torch/kernels/quantization/linear_attention/int8.py new file mode 100644 index 00000000000..d9e16a29dc8 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/int8.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Signed narrow-range INT8 QDQ for recurrent-state tiles.""" + +import triton +import triton.language as tl +from triton.language.extra.cuda import libdevice + +__all__ = [] + + +@triton.jit +def int8_scalar_qdq(value, scale): + codes = libdevice.nearbyint(tl.div_rn(value, scale)) + return tl.minimum(tl.maximum(codes, -127.0), 127.0) * scale diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index f8847aa1fc1..e8db5c8446a 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -161,6 +161,8 @@ from modelopt.torch.opt.config_loader import load_config from modelopt.torch.utils.network import ConstructorLike +from .linear_attention.config import LinearAttentionPolicyEntry + class QuantizerCfgEntry(ModeloptBaseConfig): """A single entry in a ``quant_cfg`` list.""" @@ -1567,6 +1569,11 @@ def _dict_to_entry(key: str, value) -> list[dict[str, Any]]: class QuantizeConfig(ModeloptBaseConfig): """Default configuration for ``quantize`` mode.""" + linear_attention: list[LinearAttentionPolicyEntry] = ModeloptField( + default=[], + title="Execution policies for supported linear-attention training modules", + ) + quant_cfg: QuantizeQuantCfgType = ModeloptField( default=[{"quantizer_name": "*", "cfg": {"num_bits": 8, "axis": None}}], title="Quantization configuration", diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 7fbd7d10dea..f36af6e77cd 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -28,6 +28,7 @@ from modelopt.torch.opt.dynamic import _DMRegistryCls from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict from modelopt.torch.utils import get_unwrapped_name +from modelopt.torch.utils.logging import print_rank_0 from .config import ( QuantizeConfig, @@ -36,6 +37,7 @@ _QuantizeExportConfig, normalize_quant_cfg_list, ) +from .linear_attention.config import LinearAttentionConfig from .nn import ( QuantModule, QuantModuleRegistry, @@ -67,6 +69,7 @@ def convert_to_quantized_model(model: ModelLikeModule, config: QuantizeConfig) - replace_quant_module(model, version=ModeloptStateManager(model).state_version) set_quantizer_by_cfg(model, config.get("quant_cfg", [])) + _apply_linear_attention_policy(model, config) metadata = {} update_quantize_metadata(model, config, metadata) @@ -136,13 +139,28 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: details regarding how MCore sharded checkpoint is restored, see modelopt.torch.opt.plugins.mcore_dist_checkpointing.restore_sharded_modelopt_state. """ + _apply_linear_attention_policy(model, config, metadata.get("linear_attention")) if "quantizer_state" not in metadata: # MCore sharded checkpoint (`torch-dist`) has its quantizer_state stored as the # extra_state of `QuantModule`. The quantizer_state is resumed with # QuantModule.set_extra_state(). return model - quantizer_state_dict = metadata["quantizer_state"] + quantizer_state_dict = dict(metadata["quantizer_state"]) + # M1 checkpoints predate prefill operand handles; newly introduced handles stay disabled. + for name, module in _linear_attention_modules(model).items(): + for site_name, quantizer in module.linear_attn_sites.named_modules(): + if isinstance(quantizer, TensorQuantizer) and not quantizer.is_enabled: + key = ".".join(filter(None, (name, "linear_attn_sites", site_name))) + quantizer_state_dict.setdefault(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() @@ -195,6 +213,12 @@ def update_quantize_metadata( model: nn.Module, config: QuantizeConfig, metadata: MetadataDict ) -> None: """Update the quantizer state in the metadata dict.""" + policies = { + name: module.linear_attention_config.model_dump() + for name, module in _linear_attention_modules(model).items() + } + if policies: + metadata["linear_attention"] = policies metadata["quantizer_state"] = quantizer_state(model) if shared_state_metadata := SharedWeightGlobalAmaxState.metadata(model): metadata["shared_quant_states"] = shared_state_metadata @@ -202,6 +226,43 @@ def update_quantize_metadata( metadata.pop("shared_quant_states", None) +def _linear_attention_modules(model): + # Optional framework plugins import conversion; defer this import to avoid that cycle. + from .plugins.linear_attention import _LinearAttentionQuantMixin + + return { + get_unwrapped_name(name, model): module + for name, module in model.named_modules() + if isinstance(module, _LinearAttentionQuantMixin) + } + + +def _apply_linear_attention_policy(model, config, saved_policies=None): + modules = _linear_attention_modules(model) + # getattr also handles pickled pre-extension config objects with no field in __dict__. + for entry in getattr(config, "linear_attention", []): + matches = [name for name in modules if fnmatch.fnmatch(name, entry.module_name)] + if not matches: + raise ValueError( + f"linear_attention rule {entry.module_name!r} matches no supported modules" + ) + for name in matches: + modules[name].linear_attention_config = entry.cfg.model_copy(deep=True) + if saved_policies is not None: + if saved_policies.keys() != modules.keys(): + raise ApplyModeError( + "Saved linear_attention policies do not match the restored modules" + ) + for name, policy in saved_policies.items(): + modules[name].linear_attention_config = LinearAttentionConfig(**policy) + for name, module in modules.items(): + module.validate_linear_attention() + if getattr(config, "linear_attention", []): + print_rank_0( + f"Linear attention {name!r}: {module.linear_attention_config.model_dump()}" + ) + + def quantizer_state(model: nn.Module) -> dict[str, Any]: """Returns the quantizer state dict describing the quantizer states in the model.""" return { diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py new file mode 100644 index 00000000000..fe41be40476 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Numerical policies and differentiable kernels for linear attention.""" + +from .config import * +from .decode import * +from .decode_prefill import * +from .kda import * +from .matmul import * +from .prefill import * +from .solve import * diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py new file mode 100644 index 00000000000..af41fd148d9 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -0,0 +1,158 @@ +# 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", + "LinearAttentionMatmulConfig", + "LinearAttentionPolicyEntry", + "LinearAttentionReplayConfig", + "LinearAttentionSolveConfig", +] + +_PrefillSite = Literal[ + "key_interaction", + "wy_value", + "wy_key", + "state_read", + "state_update", + "output_state", + "output_score", + "output_value", +] +_ElementwiseSite = Literal[ + "gate_prefix", + "gate_exp", + "value_residual", + "state_decay", + "state_add", + "output_add", +] +_ArithmeticDtype = Literal["float32", "float16", "bfloat16"] + + +class LinearAttentionMatmulConfig(ModeloptBaseConfig): + """Round an accumulator after each left-to-right reduction block. + + Partial products use the baseline working dtype. This specifies an emulation + schedule, not the internal accumulation order of a hardware MMA instruction. + """ + + accumulator_dtype: _ArithmeticDtype | None = ModeloptField(default=None) + reduction_block: int | None = Field(default=None, ge=1) + + @model_validator(mode="after") + def _require_complete_schedule(self): + if (self.accumulator_dtype is None) != (self.reduction_block is None): + raise ValueError("accumulator_dtype and reduction_block must be specified together") + return self + + +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 LinearAttentionSolveConfig(ModeloptBaseConfig): + """Exact solve; inverse approximation is a later delivery.""" + + method: Literal["exact"] = ModeloptField(default="exact") + + +class LinearAttentionReplayConfig(ModeloptBaseConfig): + """Anchor refresh and encoded rank-one update policy.""" + + window: int = Field(default=8, ge=1, le=64, strict=True) + factor_qdq: bool = ModeloptField(default=True) + encoding: Literal["once", "reencode"] = ModeloptField(default="once") + + +class LinearAttentionDecodeConfig(ModeloptBaseConfig): + """Explicit suffix recurrence; workload supplies per-sequence prefix lengths.""" + + mode: Literal["token", "replay"] = ModeloptField(default="token") + implementation: Literal["torch", "triton"] = ModeloptField(default="torch") + readout: Literal["working", "stored"] = ModeloptField(default="stored") + quantize_initial: bool = ModeloptField(default=True) + prefill_state_qdq: bool = ModeloptField(default=False) + state_codec: Literal["tile", "int8_hadamard32"] = ModeloptField(default="tile") + decay_log_step: float | None = Field(default=None, gt=0, allow_inf_nan=False) + replay: LinearAttentionReplayConfig | None = ModeloptField(default=None) + + @model_validator(mode="after") + def _validate_replay(self): + if (self.mode == "replay") != (self.replay is not None): + raise ValueError("replay settings must be supplied exactly when mode='replay'") + if self.implementation == "triton" and self.replay and self.replay.encoding != "once": + raise ValueError("The Triton replay candidate implements encode-once only") + if self.state_codec == "int8_hadamard32" and self.prefill_state_qdq: + raise ValueError( + "Hadamard state QDQ starts at decode handoff; disable prefill_state_qdq" + ) + return self + + +class LinearAttentionConfig(ModeloptBaseConfig): + """GDN/KDA chunk-64 policy; unsupported numerical modes fail config validation. + + ``state.block_v`` defines one dynamic scale per ``[Dk, block_v]`` tile of each + sequence/head. The initial state and every chunk's final state are rounded when + the module's 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: LinearAttentionSolveConfig = ModeloptField(default=LinearAttentionSolveConfig()) + decode: LinearAttentionDecodeConfig | None = ModeloptField(default=None) + matmul: dict[_PrefillSite, LinearAttentionMatmulConfig] = ModeloptField(default={}) + elementwise: dict[_ElementwiseSite, _ArithmeticDtype] = ModeloptField(default={}) + + @model_validator(mode="after") + def _validate_arithmetic_backend(self): + if self.backend == "fla" and ( + self.matmul + or self.elementwise + or self.solve.method != "exact" + or self.decode is not None + ): + raise ValueError("Prefill arithmetic policies require backend='matmul'") + if self.decode is not None and self.decode.state_codec == "int8_hadamard32": + if self.state.block_v < 32: + raise ValueError("int8_hadamard32 requires block_v >= 32") + return self + + +class LinearAttentionPolicyEntry(ModeloptBaseConfig): + """Assign a complete policy to supported modules matching ``module_name``. + + Rules apply in order: the last match wins, without merging nested fields. + A rule must match at least one supported linear-attention module. + """ + + module_name: str = Field(...) + cfg: LinearAttentionConfig = ModeloptField(default=LinearAttentionConfig()) diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py new file mode 100644 index 00000000000..80905958bb9 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -0,0 +1,427 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Explicit token-state and encoded-update replay references for QAT.""" + +from dataclasses import dataclass + +import torch + +from .config import LinearAttentionDecodeConfig +from .utils import _fp8_quantize + +__all__ = [ + "EncodedLinearAttentionTensor", + "LinearAttentionCarry", + "ReplayEntry", + "recurrent_decode", + "recurrent_decode_reference", +] + + +@dataclass +class EncodedLinearAttentionTensor: + """Fake decoded values plus detached scale metadata; no compressed storage claim.""" + + values: torch.Tensor + scales: torch.Tensor | None + format: str + block_v: int | None + + +@dataclass +class ReplayEntry: + """An already computed, encoded rank-one update and its log retention.""" + + key: EncodedLinearAttentionTensor + update: EncodedLinearAttentionTensor + log_retention: torch.Tensor + + +@dataclass +class LinearAttentionCarry: + """Explicit anchor/update state, position, and codec contract across calls.""" + + anchor: EncodedLinearAttentionTensor + entries: tuple[ReplayEntry, ...] + position: int + started: bool + signature: str + value_basis: str = "identity" + + @property + def cursor(self): + """Number of encoded updates since the last anchor refresh.""" + return len(self.entries) + + def reconstruct(self, *, original_basis=True): + """Replay entries, returning the original value basis unless explicitly disabled.""" + state = self.anchor.values + for entry in self.entries: + gate = entry.log_retention + decay = gate.exp().unsqueeze(-1) + if gate.ndim == 1: + decay = decay.unsqueeze(-1) + state = state * decay + entry.key.values.unsqueeze(-1) * entry.update.values.unsqueeze( + -2 + ) + if original_basis and self.value_basis == "hadamard32": + state = _hadamard32(state) + return state + + +def _hadamard32(value): + """Apply the orthonormal Sylvester transform to contiguous 32-value groups.""" + shape = value.shape + for width in (1, 2, 4, 8, 16): + pairs = value.reshape(*shape[:-1], -1, 2, width) + left, right = pairs.unbind(-2) + value = torch.stack((left + right, left - right), dim=-2).reshape(shape) + return value * (32**-0.5) + + +class _IdentityGradient(torch.autograd.Function): + @staticmethod + def forward(ctx, value, rounded): + return rounded + + @staticmethod + def backward(ctx, gradient): + return gradient, None + + +def _encode(value, enabled, block_v, *, state=False, state_format="fp8_e4m3", state_codec="tile"): + if not enabled: + return EncodedLinearAttentionTensor(value, None, "identity", None) + if state and state_codec == "int8_hadamard32": + with torch.no_grad(): + groups = value.float().reshape(*value.shape[:-1], -1, 32) + scales = (groups.abs().amax(-1) / 127).clamp_min(6e-8) + normalized = groups / scales.unsqueeze(-1) + codes = normalized.sign() * (normalized.abs() + 0.5).floor() + metadata = scales.to(torch.float16) + decoded = (codes.clamp(-127, 127) * metadata.float().unsqueeze(-1)).reshape_as(value) + return EncodedLinearAttentionTensor( + _IdentityGradient.apply(value, decoded.to(value.dtype)), metadata, "int8", 32 + ) + rounded, scales = [], [] + for part in value.split(block_v, dim=-1): + if state_format == "fp8_e4m3": + tensor = part.flatten(-2) if state else part + axis = tuple(range(tensor.ndim - 1)) or None + decoded, scale = _fp8_quantize(tensor, axis) + rounded.append(decoded.reshape_as(part)) + scales.append(scale.squeeze(-1)) + else: + with torch.no_grad(): + axes = (-2, -1) if state else (-1,) + part = part.float() + amax = part.abs().amax(dim=axes, keepdim=True) + scale = torch.where(amax > 0, amax / 127.0, torch.ones_like(amax)) + rounded.append((part / scale).round().clamp(-127, 127) * scale) + scales.append(scale[..., 0, 0] if state else scale[..., 0]) + decoded = torch.cat(rounded, dim=-1).to(value.dtype) + if state_format == "int8": + decoded = _IdentityGradient.apply(value, decoded) + return EncodedLinearAttentionTensor(decoded, torch.stack(scales, dim=-1), state_format, block_v) + + +def _signature(config, state_qdq, block_v, state_format): + return ( + config.model_dump_json(exclude={"implementation", "prefill_state_qdq"}) + + f"/{state_qdq}/{block_v}/{state_format}" + ) + + +def _sum_keys(value): + keys = value.shape[-2] + padded = 1 << (keys - 1).bit_length() + if keys != padded: + value = torch.nn.functional.pad(value, (0, 0, 0, padded - keys)) + while value.shape[-2] > 1: + half = value.shape[-2] // 2 + value = value[..., :half, :] + value[..., half:, :] + return value[..., 0, :] + + +def _round_log_gate(gate, step): + if step is None: + return gate + rounded = (gate / step).round() * step + return _IdentityGradient.apply(gate, rounded.detach()) + + +def _prepare_carry( + q, k, v, g, beta, config, state_qdq, block_v, initial_state, carry, position, state_format +): + if q.ndim != 3 or k.shape != q.shape or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k/v must have aligned [T,H,D] shapes") + if beta.shape != q.shape[:2] or g.shape not in (beta.shape, q.shape): + raise ValueError("beta must be [T,H]; g must be [T,H] or [T,H,Dk]") + if block_v not in (16, 32, 64, 128): + raise ValueError("block_v must be 16, 32, 64, or 128") + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") + hadamard = config.state_codec == "int8_hadamard32" + if hadamard: + if state_qdq and state_format != "int8": + raise ValueError("int8_hadamard32 requires INT8 state quantization") + if v.shape[-1] % 32 or block_v < 32: + raise ValueError("int8_hadamard32 requires Dv divisible by 32 and block_v >= 32") + signature = _signature(config, state_qdq, block_v, state_format) + if carry is not None and initial_state is not None: + raise ValueError("Supply either carry or initial_state") + shape = (q.shape[1], q.shape[2], v.shape[2]) + if carry is None: + initial_state = q.new_zeros(shape) if initial_state is None else initial_state + carry = LinearAttentionCarry( + _encode(initial_state, False, block_v, state=True, state_format=state_format), + (), + position, + False, + signature, + ) + if carry.signature != signature or carry.anchor.values.shape != shape: + raise ValueError("Carry policy or state shape does not match this recurrence") + if config.mode == "token" and carry.entries: + raise ValueError("Token carry cannot contain replay entries") + if config.replay is not None and carry.cursor >= config.replay.window: + raise ValueError("Replay cursor must be below the refresh window") + if len(q) and not carry.started: + initial = _hadamard32(carry.anchor.values) if hadamard else carry.anchor.values + anchor = _encode( + initial, + state_qdq and config.quantize_initial, + block_v, + state=True, + state_format=state_format, + state_codec=config.state_codec, + ) + carry = LinearAttentionCarry( + anchor, (), carry.position, True, signature, "hadamard32" if hadamard else "identity" + ) + return carry, signature + + +def recurrent_decode_reference( + q, + k, + v, + g, + beta, + *, + config: LinearAttentionDecodeConfig, + state_qdq=False, + state_format="fp8_e4m3", + block_v=64, + initial_state=None, + carry=None, + position=0, + scale=None, +): + """Run one preactivated sequence [T,H,D] with explicit token/replay write events. + + Keys and value heads must already be aligned. Scalar GDN or per-key-channel + KDA log gates are accepted. Outputs and all returned carry values retain their + graphs. An empty call performs no write or initial-state quantization. + """ + carry, signature = _prepare_carry( + q, k, v, g, beta, config, state_qdq, block_v, initial_state, carry, position, state_format + ) + if len(q) == 0: + # Keep empty input gradients defined without introducing a state write. + zero = (q.sum() + k.sum() + v.sum() + g.sum() + beta.sum()) * 0 + output = v + zero + return output, carry + if config.state_codec == "int8_hadamard32": + v = _hadamard32(v) + scale = q.shape[-1] ** -0.5 if scale is None else scale + outputs = [] + with torch.autocast(device_type=q.device.type, enabled=False): + for t in range(len(q)): + entries = carry.entries + if config.replay is not None and config.replay.encoding == "reencode": + entries = tuple( + ReplayEntry( + _encode(e.key.values, config.replay.factor_qdq, k.shape[-1]), + _encode(e.update.values, config.replay.factor_qdq, block_v), + e.log_retention, + ) + for e in entries + ) + carry = LinearAttentionCarry( + carry.anchor, entries, carry.position, True, signature, carry.value_basis + ) + state = carry.reconstruct(original_basis=False) + gate = _round_log_gate(g[t], config.decay_log_step) + decay = gate.exp().unsqueeze(-1) + if gate.ndim == 1: + decay = decay.unsqueeze(-1) + key = _encode(k[t], config.replay is not None and config.replay.factor_qdq, k.shape[-1]) + decayed = state * decay + residual = v[t] - _sum_keys(key.values.unsqueeze(-1) * decayed) + update = _encode( + beta[t].unsqueeze(-1) * residual, + config.replay is not None and config.replay.factor_qdq, + block_v, + ) + working = decayed + key.values.unsqueeze(-1) * update.values.unsqueeze(-2) + if config.replay is not None: + entries = (*entries, ReplayEntry(key, update, gate)) + refresh = config.replay is None or len(entries) == config.replay.window + if refresh: + anchor = _encode( + working, + state_qdq, + block_v, + state=True, + state_format=state_format, + state_codec=config.state_codec, + ) + entries = () + stored = anchor.values + else: + anchor = carry.anchor + stored = working + next_carry = LinearAttentionCarry( + anchor, entries, carry.position + 1, True, signature, carry.value_basis + ) + read = working if config.readout == "working" else stored + outputs.append(_sum_keys(q[t].unsqueeze(-1) * read) * scale) + carry = next_carry + output = torch.stack(outputs) + if config.state_codec == "int8_hadamard32": + output = _hadamard32(output) + return output, carry + + +def recurrent_decode( + q, + k, + v, + g, + beta, + *, + config: LinearAttentionDecodeConfig, + state_qdq=False, + state_format="fp8_e4m3", + block_v=64, + initial_state=None, + carry=None, + position=0, + scale=None, + checkpoint_interval=8, +): + """Run the selected training implementation with explicit differentiable carry. + + The fused encode-once path carries the same reconstructed state incrementally + and exposes the anchor and encoded updates needed for continuation gradients. + """ + if config.implementation == "torch": + return recurrent_decode_reference( + q, + k, + v, + g, + beta, + config=config, + state_qdq=state_qdq, + state_format=state_format, + block_v=block_v, + initial_state=initial_state, + carry=carry, + position=position, + scale=scale, + ) + carry, signature = _prepare_carry( + q, k, v, g, beta, config, state_qdq, block_v, initial_state, carry, position, state_format + ) + if len(q) == 0: + return v + (q.sum() + k.sum() + v.sum() + g.sum() + beta.sum()) * 0, carry + if config.replay is not None and config.replay.encoding != "once": + raise ValueError("The fused implementation supports encode-once replay only") + # Keep the CPU reference importable without the optional CUDA/Triton backend. + from modelopt.torch.kernels.quantization.linear_attention.decode import fused_recurrence + + if config.state_codec == "int8_hadamard32": + v = _hadamard32(v) + replay = config.replay is not None + factor_qdq = config.replay is not None and config.replay.factor_qdq + key = _encode(k, factor_qdq, k.shape[-1]) + gate = _round_log_gate(g, config.decay_log_step) + window = config.replay.window if config.replay is not None else 1 + result = fused_recurrence( + q, + key.values, + v, + gate, + beta, + carry.reconstruct(original_basis=False), + state_qdq=state_qdq, + state_format=state_format, + state_codec=config.state_codec, + block_v=block_v, + replay=replay, + factor_qdq=factor_qdq, + window=window, + cursor=carry.cursor, + read_stored=config.readout == "stored", + scale=scale, + checkpoint_interval=checkpoint_interval, + ) + output, final, updates, last_anchor, anchor_scales, update_scales = result + if config.state_codec == "int8_hadamard32": + output = _hadamard32(output) + state_format = state_format if state_qdq else "identity" + state_block_v = 32 if config.state_codec == "int8_hadamard32" else block_v + if not replay: + anchor = EncodedLinearAttentionTensor( + final, + anchor_scales if state_qdq else None, + state_format, + state_block_v if state_qdq else None, + ) + return output, LinearAttentionCarry( + anchor, (), carry.position + len(q), True, signature, carry.value_basis + ) + if carry.cursor + len(q) >= window: + anchor = EncodedLinearAttentionTensor( + last_anchor, + anchor_scales if state_qdq else None, + state_format, + state_block_v if state_qdq else None, + ) + start = ((carry.cursor + len(q)) // window) * window - carry.cursor + entries = () + else: + anchor, entries, start = carry.anchor, carry.entries, 0 + additions = [] + for t in range(start, len(q)): + encoded_key = EncodedLinearAttentionTensor( + key.values[t], + key.scales[t] if key.scales is not None else None, + key.format, + key.block_v, + ) + encoded_update = EncodedLinearAttentionTensor( + updates[t], + update_scales[t] if factor_qdq else None, + "fp8_e4m3" if factor_qdq else "identity", + block_v if factor_qdq else None, + ) + additions.append(ReplayEntry(encoded_key, encoded_update, gate[t])) + return output, LinearAttentionCarry( + anchor, (*entries, *additions), carry.position + len(q), True, signature, carry.value_basis + ) diff --git a/modelopt/torch/quantization/linear_attention/decode_prefill.py b/modelopt/torch/quantization/linear_attention/decode_prefill.py new file mode 100644 index 00000000000..3ad194c017f --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/decode_prefill.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differentiable explicit per-sequence prefill/decode phase handoff.""" + +from contextlib import contextmanager +from itertools import pairwise + +import torch + +from .decode import recurrent_decode + +__all__ = ["linear_attention_training_phase"] + + +def _lengths(values): + if isinstance(values, torch.Tensor): + values = values.tolist() + values = tuple(values) + if any(type(value) is not int or value < 0 for value in values): + raise ValueError("Prefill lengths must be nonnegative integers") + return values + + +@contextmanager +def linear_attention_training_phase(model, prefill_lengths): + """Supply explicit sequence phases through forward and checkpointed backward. + + Runtime phase metadata is local to the converted layers and restored on exit. + Keep this context active through backward when activation checkpointing is used. + """ + # The plugin imports this numerical package during quantization initialization. + from ..plugins.linear_attention import _LinearAttentionQuantMixin + + lengths = _lengths(prefill_lengths) + layers = [ + m + for m in model.modules() + if isinstance(m, _LinearAttentionQuantMixin) + and m.linear_attention_config.decode is not None + ] + if not layers: + raise ValueError("The model has no converted decode-aware linear-attention layers") + previous = [getattr(m, "_linear_attention_prefill_lengths", None) for m in layers] + try: + for module in layers: + module._linear_attention_prefill_lengths = lengths + yield model + finally: + for module, original in zip(layers, previous): + module._linear_attention_prefill_lengths = original + + +def _decode_prefill( + q, + k, + v, + g, + beta, + *, + sites, + policy, + w_quantizer, + state_qdq, + state_format, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + state_v_first, + output_dtype, + prefill_lengths, +): + # Prefill dispatch imports this wrapper; defer its shared core to avoid a cycle. + from .prefill import _matmul_prefill + + if prefill_lengths is None: + raise ValueError("Decode-aware training requires explicit per-sequence prefill lengths") + if q.ndim != 4 or k.shape != q.shape or v.ndim != 4 or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k and v must have compatible [B,T,H,D] shapes") + batch, length, key_heads, keys = q.shape + heads, values = v.shape[2:] + if batch < 1 or key_heads < 1 or heads % key_heads or beta.shape != (batch, length, heads): + raise ValueError("Invalid batch/head dimensions or beta shape") + if g.shape not in (beta.shape, (*beta.shape, keys)): + raise ValueError("Invalid GDN/KDA log-retention shape") + q, k = (x.repeat_interleave(heads // key_heads, dim=2) for x in (q, k)) + boundaries = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens + if boundaries is None: + sequences = [(b, 0, length) for b in range(batch)] + else: + bounds = _lengths(boundaries) + if ( + batch != 1 + or len(bounds) < 2 + or bounds[0] != 0 + or bounds[-1] != length + or any(a > b for a, b in pairwise(bounds)) + ): + raise ValueError( + "Packed boundaries must partition a batch of one, allowing empty entries" + ) + sequences = [(0, a, b) for a, b in pairwise(bounds)] + prefixes = _lengths(prefill_lengths) + if len(prefixes) != len(sequences) or any( + p > end - start for p, (_, start, end) in zip(prefixes, sequences) + ): + raise ValueError("Supply one valid prefill length per sequence") + if initial_state is None: + states = q.new_zeros(len(sequences), heads, keys, values) + else: + states = initial_state.transpose(-1, -2) if state_v_first else initial_state + states = states.to(q.dtype) + if states.shape != (len(sequences), heads, keys, values): + raise ValueError("Initial state shape does not match sequence/head dimensions") + active = [n for n, prefix in enumerate(prefixes) if prefix] + prefix_outputs, prefix_states = {}, {} + if active: + packed = [ + torch.cat( + [ + tensor[sequences[n][0], sequences[n][1] : sequences[n][1] + prefixes[n]] + for n in active + ] + ).unsqueeze(0) + for tensor in (q, k, v, g, beta) + ] + offsets = [0] + for n in active: + offsets.append(offsets[-1] + prefixes[n]) + output, final = _matmul_prefill( + *packed, + sites=sites, + policy=policy, + w_quantizer=w_quantizer, + state_qdq=state_qdq and policy.decode.prefill_state_qdq, + state_format=state_format, + scale=scale, + initial_state=states[active], + output_final_state=True, + cu_seqlens=None, + cu_seqlens_cpu=torch.tensor(offsets), + state_v_first=False, + chunk_size=policy.chunk_size, + output_dtype=q.dtype, + ) + for index, n in enumerate(active): + prefix_outputs[n] = output[0, offsets[index] : offsets[index + 1]] + prefix_states[n] = final[index] + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + split = start + prefixes[n] + suffix, carry = recurrent_decode( + *(x[b, split:end] for x in (q, k, v, g, beta)), + config=policy.decode, + state_qdq=state_qdq, + 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..1e6e4a4c398 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differentiable KDA prefill with stable per-channel decay interactions.""" + +import torch +import torch.nn.functional as F + +from .matmul import _validate_operand_quantizer +from .prefill import _dispatch_prefill + +__all__ = ["matmul_kda"] + + +def matmul_kda( + q, + k, + v, + g, + beta, + *, + sites, + policy, + w_quantizer, + 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, +): + """Run KDA chunk prefill with actual-operand QDQ and differentiable state carry. + + Gates are per-key-channel natural log retentions. Causal interactions form + exp(prefix[i] - prefix[j]) directly, avoiding an overflowing inverse decay. + Raw-gate activation, when requested, follows the selected FLA gate formula. + """ + if policy.backend != "matmul" or chunk_size != policy.chunk_size: + raise ValueError("matmul_kda requires backend='matmul' and its configured chunk size") + if cp_context is not None or disable_recompute or return_intermediate_states: + raise NotImplementedError( + "KDA matmul does not support CP or FLA recompute/intermediate flags" + ) + if allow_neg_eigval and not use_beta_sigmoid_in_kernel: + raise ValueError("allow_neg_eigval requires use_beta_sigmoid_in_kernel") + if safe_gate and use_gate_in_kernel and (lower_bound is None or not -5 <= lower_bound < 0): + raise ValueError("safe_gate requires a lower_bound in [-5, 0)") + if lower_bound is not None and lower_bound >= 0: + raise ValueError("lower_bound must be negative") + sites.validate() + _validate_operand_quantizer(w_quantizer, "kda_w_quantizer") + output_dtype = q.dtype + dtype = torch.float64 if output_dtype == torch.float64 else torch.float32 + q, k, v, g, beta = (x.to(dtype) for x in (q, k, v, g, beta)) + if g.ndim != 4: + raise ValueError("matmul_kda requires per-key-channel KDA log gates") + if use_qk_l2norm_in_kernel: + q, k = (x * (x.square().sum(-1, keepdim=True) + 1e-6).rsqrt() for x in (q, k)) + if use_gate_in_kernel: + if A_log is None: + raise ValueError("Fused KDA gate requires A_log") + if dt_bias is not None: + g = g + dt_bias.to(dtype).reshape(g.shape[-2:]) + rate = A_log.to(dtype).exp().reshape(g.shape[-2], 1) + g = -rate * F.softplus(g) if lower_bound is None else lower_bound * (rate * g).sigmoid() + if use_beta_sigmoid_in_kernel: + beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) + return _dispatch_prefill( + q, + k, + v, + g, + beta, + sites=sites, + policy=policy, + w_quantizer=w_quantizer, + 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, + chunk_size=chunk_size, + output_dtype=output_dtype, + prefill_lengths=prefill_lengths, + ) diff --git a/modelopt/torch/quantization/linear_attention/matmul.py b/modelopt/torch/quantization/linear_attention/matmul.py new file mode 100644 index 00000000000..e8cadfb6d20 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/matmul.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Operand QDQ and explicit arithmetic schedules for linear-attention matmuls.""" + +from typing import get_args + +import torch +from torch import nn + +from .config import LinearAttentionMatmulConfig, _PrefillSite + +__all__ = ["LinearAttentionMatmulSites"] + + +class _MatmulSite(nn.Module): + def __init__(self, *, legacy_w=False): + # QuantizeConfig imports this package before TensorQuantizer is defined. + from modelopt.torch.quantization.config import QuantizerAttributeConfig + from modelopt.torch.quantization.nn import TensorQuantizer + + super().__init__() + if not legacy_w: + self.lhs_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) + self.rhs_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) + + def forward(self, lhs, rhs, policy, *, lhs_quantizer=None): + # The emulation working dtype must not inherit an outer training autocast. + with torch.autocast(device_type=lhs.device.type, enabled=False): + # Operands are [sequence-or-chunk, head, row, reduction] before transpose. + quantizer = lhs_quantizer if lhs_quantizer is not None else self.lhs_quantizer + lhs, rhs = quantizer(lhs), self.rhs_quantizer(rhs) + if policy.accumulator_dtype is None: + return lhs @ rhs.transpose(-1, -2) + acc = lhs.new_zeros(*lhs.shape[:-1], rhs.shape[-2]) + for lo in range(0, lhs.shape[-1], policy.reduction_block): + hi = lo + policy.reduction_block + acc = acc + lhs[..., lo:hi] @ rhs[..., lo:hi].transpose(-1, -2) + rounded = acc.to(getattr(torch, policy.accumulator_dtype)).to(lhs.dtype) + acc = acc + (rounded - acc).detach() + return acc + + +class LinearAttentionMatmulSites(nn.ModuleDict): + """The eight prefill matmuls, with independent QDQ on their actual operands. + + Inputs to each handle are four-dimensional, with reduction on the last axis. + The state-read LHS uses the parent's ``gdn_w_quantizer`` or ``kda_w_quantizer``; + no second W quantizer is registered here. New handles start disabled. + """ + + def __init__(self): + """Create disabled operand handles for the eight prefill sites.""" + super().__init__( + {name: _MatmulSite(legacy_w=name == "state_read") for name in get_args(_PrefillSite)} + ) + + @property + def is_enabled(self): + """Whether any of the additional operand handles is enabled.""" + return any(q.is_enabled for site in self.values() for q in site.children()) + + def matmul(self, name, lhs, rhs, policy, *, w_quantizer=None): + """Multiply LHS by transposed RHS using the named numerical site.""" + return self[name]( + lhs, + rhs, + policy.matmul.get(name, LinearAttentionMatmulConfig()), + lhs_quantizer=w_quantizer if name == "state_read" else None, + ) + + def validate(self): + """Reject formats whose scale grouping has not been qualified at these sites.""" + for name, site in self.items(): + for handle, quantizer in site.named_children(): + _validate_operand_quantizer(quantizer, f"linear_attn_sites.{name}.{handle}") + + +def _validate_operand_quantizer(quantizer, name): + # Deferred because QuantizeConfig imports this package before TensorQuantizer exists. + from modelopt.torch.quantization.nn import TensorQuantizer + + if not isinstance(quantizer, TensorQuantizer): + raise ValueError(f"{name} requires a single TensorQuantizer") + if not quantizer.is_enabled: + return + fp8 = ( + quantizer.num_bits == (4, 3) + and quantizer.axis == (0, 1, 2) + and quantizer.block_sizes is None + ) + nvfp4 = ( + quantizer.num_bits == (2, 1) + and quantizer.axis is None + and quantizer.block_sizes == {-1: 16, "type": "dynamic", "scale_bits": (4, 3)} + ) + if not ( + (fp8 or nvfp4) + and quantizer._dynamic + and quantizer.fake_quant + and quantizer._pass_through_bwd + and not quantizer.rotate_is_enabled + and quantizer.pre_quant_scale is None + and quantizer.backend is None + and not quantizer._bias + and not quantizer._use_constant_amax + ): + raise ValueError( + f"{name} requires dynamic per-row E4M3 or block-16 NVFP4 " + "with dynamic tensor amax, identity STE, and no extra transformations" + ) diff --git a/modelopt/torch/quantization/linear_attention/prefill.py b/modelopt/torch/quantization/linear_attention/prefill.py new file mode 100644 index 00000000000..ffc54f12ba9 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/prefill.py @@ -0,0 +1,249 @@ +# 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 +from .matmul import LinearAttentionMatmulSites, _validate_operand_quantizer +from .solve import triangular_inverse +from .utils import _prepare, _state_qdq + +__all__ = ["matmul_gdn"] + + +def matmul_gdn( + q, + k, + v, + g, + beta, + *, + sites: LinearAttentionMatmulSites, + policy: LinearAttentionConfig, + w_quantizer, + 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, +): + """Execute all eight GDN prefill sites with autograd through their QDQ operators. + + This path batches chunk-local matmuls and carries states across chunks without + detaching them. It uses FP32 working arithmetic (FP64 for double inputs) and + casts outputs back to Q's dtype. It materializes the inverse and intermediates; + it is a numerical-emulation backend, not the fused FLA implementation. + """ + 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") + sites.validate() + _validate_operand_quantizer(w_quantizer, "gdn_w_quantizer") + output_dtype = q.dtype + dtype = torch.float64 if output_dtype == torch.float64 else torch.float32 + q, k, v, g, beta = (x.to(dtype) for x in (q, k, v, g, beta)) + if 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") + return _dispatch_prefill( + q, + k, + v, + g, + beta, + sites=sites, + policy=policy, + w_quantizer=w_quantizer, + 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, + chunk_size=chunk_size, + output_dtype=output_dtype, + prefill_lengths=prefill_lengths, + ) + + +def _dispatch_prefill(*args, policy, chunk_size, prefill_lengths, **kwargs): + if policy.decode is not None: + return _decode_prefill(*args, policy=policy, prefill_lengths=prefill_lengths, **kwargs) + return _matmul_prefill(*args, policy=policy, chunk_size=chunk_size, **kwargs) + + +def _matmul_prefill( + q, + k, + v, + g, + beta, + *, + sites, + policy, + w_quantizer, + state_qdq, + state_format, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + state_v_first, + chunk_size, + output_dtype, +): + dtype = q.dtype + channel_gate = g.ndim == 4 + initial_state = initial_state.to(dtype) if initial_state is not None else None + boundaries = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens + q, k, state, sequences = _prepare( + q, + k, + v, + g, + beta, + initial_state, + boundaries, + state_v_first, + ) + scale = q.shape[-1] ** -0.5 if scale is None else scale + count = (max(end - start for _, start, end in sequences) + chunk_size - 1) // chunk_size + padded_length = count * chunk_size + heads = v.shape[2] + + def chunks(x): + pieces = [] + for b, start, end in sequences: + part = x[b, start:end] + padding = part.new_zeros(padded_length - len(part), *part.shape[1:]) + pieces.append(torch.cat((part, padding))) + packed = torch.stack(pieces) + # [sequence, chunk, head, token, optional feature] -> batch matmuls over chunks. + packed = packed.reshape(len(sequences), count, chunk_size, heads, *x.shape[3:]) + return packed.transpose(2, 3).reshape(-1, heads, chunk_size, *x.shape[3:]) + + def arithmetic(name, x): + target = policy.elementwise.get(name) + if target is None: + return x + rounded = x.to(getattr(torch, target)).to(dtype) + return x + (rounded - x).detach() + + def mm(name, lhs, rhs): + return sites.matmul(name, lhs, rhs, policy, w_quantizer=w_quantizer) + + qc, kc, vc = (chunks(x) for x in (q, k, v)) + gc = arithmetic("gate_prefix", chunks(g).cumsum(-2 if channel_gate else -1)) + bc = chunks(beta).unsqueeze(-1) + gate = arithmetic("gate_exp", gc.exp()) + if channel_gate: + lower = _kda_interaction("key_interaction", bc * kc, kc, gc, mm, arithmetic).tril(-1) + scores = _kda_interaction("output_score", qc * scale, kc, gc, mm, arithmetic).tril() + weighted_keys = kc * arithmetic("gate_exp", (gc[..., -1:, :] - gc).exp()) + final_decay = arithmetic("gate_exp", gc[..., -1, :].exp()).unsqueeze(-1) + else: + causal = torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device).tril() + relative = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0) + decay = arithmetic("gate_exp", relative.exp()) + gate = gate.unsqueeze(-1) + lower = (mm("key_interaction", bc * kc, kc) * decay).tril(-1) + scores = (mm("output_score", qc * scale, kc) * decay).tril() + weighted_keys = kc * arithmetic("gate_exp", (gc[..., -1:] - gc).exp()).unsqueeze(-1) + final_decay = arithmetic("gate_exp", gc[..., -1].exp())[..., None, None] + inverse = triangular_inverse(lower, policy.solve) + u = mm("wy_value", inverse, (bc * vc).transpose(-1, -2)) + w = mm("wy_key", inverse, (bc * kc * gate).transpose(-1, -2)) + + def unbatch(x): + return x.reshape(len(sequences), count, *x.shape[1:]) + + qc, gate, final_decay, u, w, scores, weighted_keys = map( + unbatch, (qc, gate, final_decay, u, w, scores, weighted_keys) + ) + if state_qdq: + state = _state_qdq(state, policy.state.block_v, state_format) + outputs = [] + for c in range(count): + # The legacy W handle is invoked once per state read and is never rerun in backward. + updated = arithmetic( + "value_residual", u[:, c] - mm("state_read", w[:, c], state.transpose(-1, -2)) + ) + memory = mm("output_state", qc[:, c] * scale * gate[:, c], state.transpose(-1, -2)) + local = mm("output_value", scores[:, c], updated.transpose(-1, -2)) + outputs.append(arithmetic("output_add", memory + local)) + decayed = arithmetic("state_decay", state * final_decay[:, c]) + update = mm( + "state_update", weighted_keys[:, c].transpose(-1, -2), updated.transpose(-1, -2) + ) + next_state = arithmetic("state_add", decayed + update) + if state_qdq: + next_state = _state_qdq(next_state, policy.state.block_v, state_format) + # Padded chunks must not introduce extra state writes for shorter packed sequences. + active = torch.tensor( + [c * chunk_size < end - start for _, start, end in sequences], device=q.device + ) + state = torch.where(active[:, None, None, None], next_state, state) + output = torch.cat(outputs, dim=2).transpose(1, 2) + pieces = [output[n, : end - start] for n, (_, start, end) in enumerate(sequences)] + output = torch.stack(pieces) if boundaries is None else torch.cat(pieces).unsqueeze(0) + final = state.transpose(-1, -2) if state_v_first else state + return output.to(output_dtype), final if output_final_state else None + + +def _kda_interaction(name, lhs, keys, prefix, mm, arithmetic): + # Causal pairwise decays avoid exp(-prefix), which overflows for strong forgetting. + batch, heads, length, width = keys.shape + result = [] + columns = torch.arange(length, device=keys.device) + for lo in range(0, length, 8): + hi = min(lo + 8, length) + rows = torch.arange(lo, hi, device=keys.device) + causal = (rows[:, None] >= columns[None, :])[None, None, ..., None] + relative = prefix[:, :, lo:hi, None] - prefix[:, :, None] + relative = relative.masked_fill(~causal, 0) + decay = arithmetic("gate_exp", relative.exp()).masked_fill(~causal, 0) + rhs = (keys[:, :, None] * decay).permute(0, 2, 1, 3, 4) + right = rhs.reshape(batch * (hi - lo), heads, length, width) + left = lhs[:, :, lo:hi].permute(0, 2, 1, 3).reshape(batch * (hi - lo), heads, 1, width) + product = mm(name, left, right).reshape(batch, hi - lo, heads, length) + result.append(product.permute(0, 2, 1, 3)) + return torch.cat(result, dim=-2) diff --git a/modelopt/torch/quantization/linear_attention/solve.py b/modelopt/torch/quantization/linear_attention/solve.py new file mode 100644 index 00000000000..de4a1ffe17c --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/solve.py @@ -0,0 +1,35 @@ +# 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 triangular inverse policies for differentiable prefill emulation.""" + +import torch + +from .config import LinearAttentionSolveConfig + +__all__ = ["triangular_inverse"] + + +def triangular_inverse(lower: torch.Tensor, policy: LinearAttentionSolveConfig) -> torch.Tensor: + """Materialize the exact unit-lower triangular inverse for WY operand QDQ.""" + if lower.ndim < 2 or lower.shape[-1] != lower.shape[-2]: + raise ValueError("lower must contain square matrices") + with torch.autocast(device_type=lower.device.type, enabled=False): + identity = torch.eye(lower.shape[-1], device=lower.device, dtype=lower.dtype).expand_as( + lower + ) + return torch.linalg.solve_triangular( + identity + lower, identity, upper=False, unitriangular=True + ) diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py new file mode 100644 index 00000000000..68353bd28df --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared helpers for linear-attention quantization.""" + +from __future__ import annotations + +from itertools import pairwise +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from ..nn import TensorQuantizer + +__all__ = [] + + +def validate_gdn_quantizer( + quantizer: TensorQuantizer, + *, + name: str, + num_bits: tuple[int | tuple[int, int], ...] = ((4, 3),), +) -> None: + """Check supported formats and the custom backward's identity STE.""" + # Numerical helpers load while QuantizeConfig initializes; defer the quantizer import. + from ..nn import TensorQuantizer + + if not isinstance(quantizer, TensorQuantizer): + raise ValueError(f"{name} requires a single TensorQuantizer") + if not ( + quantizer._dynamic + and quantizer.num_bits in num_bits + and (quantizer.num_bits != 8 or (not quantizer.unsigned and quantizer.narrow_range)) + and quantizer.block_sizes is None + and quantizer.fake_quant + and quantizer._pass_through_bwd + and not quantizer.rotate_is_enabled + and quantizer.pre_quant_scale is None + and quantizer.backend is None + and not quantizer._bias + and not quantizer._use_constant_amax + ): + raise ValueError( + f"{name} supports only dynamic fake quantization with num_bits in {num_bits}, " + "pass_through_bwd=True, no block_sizes, rotation, pre-scaling, bias, constant " + "amax, or custom backend. Other gradient rules and formats are not implemented." + ) + + +def _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first): + if q.ndim != 4 or k.shape != q.shape or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k must have shape [B,T,Hk,Dk] and v shape [B,T,Hv,Dv]") + batch, length, heads, keys = q.shape + value_heads, values = v.shape[2:] + if length == 0 or value_heads % heads: + raise ValueError("nonempty sequences and Hv divisible by Hk are required") + if beta.shape != (batch, length, value_heads) or g.shape not in ( + beta.shape, + (*beta.shape, keys), + ): + raise ValueError("beta must be [B,T,Hv]; g must be [B,T,Hv] or [B,T,Hv,Dk]") + if cu_seqlens is None: + sequences = [(b, 0, length) for b in range(batch)] + else: + boundaries = cu_seqlens.tolist() + if ( + batch != 1 + or len(boundaries) < 2 + or boundaries[0] != 0 + or boundaries[-1] != length + or any(a >= b for a, b in pairwise(boundaries)) + ): + raise ValueError("cu_seqlens must partition a packed batch of size one") + sequences = [(0, a, b) for a, b in pairwise(boundaries)] + if initial_state is None: + state = q.new_zeros(len(sequences), value_heads, keys, values) + else: + state = initial_state.transpose(-1, -2) if state_v_first else initial_state + if state.shape != (len(sequences), value_heads, keys, values): + raise ValueError("initial_state shape does not match sequence/head dimensions") + return ( + q.repeat_interleave(value_heads // heads, dim=2), + k.repeat_interleave(value_heads // heads, dim=2), + state, + sequences, + ) + + +def _fp8_quantize(value: torch.Tensor, axis): + """Apply ModelOpt FP8 QDQ and return its detached dequantization scales.""" + # QuantizeConfig imports this module before the quantizer classes are initialized. + from ..config import QuantizerAttributeConfig + from ..nn import TensorQuantizer + + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=(4, 3), type="dynamic", axis=axis, pass_through_bwd=True) + ) + quantized = quantizer(value) + amax = quantizer._get_amax(value).float() + safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) + return quantized, torch.div(448.0, safe_amax).reciprocal() + + +def _state_qdq(state: torch.Tensor, block_v: int = 64, state_format: str = "fp8_e4m3"): + """Dynamic state-tile QDQ with detached scales and identity STE.""" + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") + if block_v not in (16, 32, 64, 128): + raise ValueError("block_v must be 16, 32, 64, or 128") + if state_format == "fp8_e4m3": + axis = tuple(range(state.ndim - 2)) or None + return torch.cat( + [ + _fp8_quantize(tile.flatten(-2), axis)[0].reshape_as(tile) + for tile in state.split(block_v, dim=-1) + ], + dim=-1, + ) + with torch.no_grad(): + rounded = [] + for tile in state.float().split(block_v, dim=-1): + amax = tile.abs().amax(dim=(-2, -1), keepdim=True) + limit = 127.0 + scale = torch.where(amax > 0, amax / limit, torch.ones_like(amax)) + normalized = (tile / scale).clamp(-limit, limit) + codes = normalized.round() + rounded.append(codes * scale) + quantized = torch.cat(rounded, dim=-1).to(state.dtype) + return state + (quantized - state).detach() diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index d7ef04f750d..5aac82cacfe 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -32,6 +32,7 @@ from modelopt.torch.opt.utils import forward_with_reshard from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.conversion import ( + _apply_linear_attention_policy, preserve_quantizer_attributes_context, set_quantizer_attributes_partial, set_quantizer_by_cfg, @@ -329,6 +330,7 @@ def forward_loop(model) -> None: else: # Already quantized, so lets apply the quant_cfg from the config set_quantizer_by_cfg(model, quantize_config.quant_cfg) + _apply_linear_attention_policy(model, quantize_config) # Fail before calibration rather than after exporting an unquantized checkpoint. _check_weight_quantization_took_effect(model, quantize_config) return calibrate(model, config.get("algorithm"), forward_loop=forward_loop) diff --git a/modelopt/torch/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py index 22b4bc2e3cc..91b95257cce 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -39,6 +39,8 @@ from .attention import * from .custom import * +from .gdn import * +from .kda import * with import_plugin("diffusers"): from .diffusion.diffusers import * @@ -69,6 +71,7 @@ with import_plugin("vllm"): from .vllm import * + from .vllm_linear_attention import * with import_plugin("trl"): from .trl import * diff --git a/modelopt/torch/quantization/plugins/gdn.py b/modelopt/torch/quantization/plugins/gdn.py new file mode 100644 index 00000000000..dd9ac102440 --- /dev/null +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fake quantization of the GatedDeltaNet (GDN) recurrent state. + +The chunked gated-delta-rule kernel keeps each head's ``[K, V]`` recurrent state in fp32 inside +one Triton launch and carries it from chunk to chunk. To emulate a deployment that stores that +state in FP8 or INT8, ModelOpt runs an adapted copy of the kernel +(:mod:`modelopt.torch.kernels.quantization.linear_attention`) that fake-quantizes the state to +E4M3 or signed narrow-range INT8 at the end of every chunk, with a scale computed inside the +kernel from the state itself. +The backward pass recomputes the same quantized states and passes the state gradient straight +through the quantization, so QAT and QAD train against the quantized recurrence. A second +quantizer covers ``w``, the WY-transformed keys that multiply the state; ``w`` is a regular tensor, +so it is fake-quantized by the ``TensorQuantizer`` itself before the kernel reads it. +""" + +from collections.abc import Callable +from functools import partial +from typing import Any + +import torch + +from ..linear_attention.prefill import matmul_gdn +from .linear_attention import _LinearAttentionQuantMixin + +__all__ = ["GatedDeltaNetStateQuantMixin"] + +GatedDeltaRuleFn = Callable[..., tuple[torch.Tensor, torch.Tensor | None]] + + +def _fla_chunk_gated_delta_rule() -> GatedDeltaRuleFn: + # FLA is an optional, heavy dependency needed only when quantization is enabled. + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + + return chunk_gated_delta_rule + + +def _state_qdq_chunk_gated_delta_rule() -> GatedDeltaRuleFn: + # Imported on first use: flash-linear-attention is a heavy optional dependency that only the + # enabled quantizer needs, and importing it warns on machines without a GPU. + try: + from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( + chunk_gated_delta_rule, + ) + except ImportError as e: + raise RuntimeError( + "GDN fake quantization needs Triton and fla-core==0.5.1 on a CUDA " + f"device; importing the state-quantizing kernel failed with {e!r}." + ) from e + return chunk_gated_delta_rule + + +class GatedDeltaNetStateQuantMixin(_LinearAttentionQuantMixin): + """Adds ``gdn_state_quantizer`` and ``gdn_w_quantizer`` to a GatedDeltaNet module. + + Subclasses route the module's chunked gated-delta-rule call through + :meth:`_state_quantized_chunk_gated_delta_rule`. Both quantizers start disabled; enable them + with ``quant_cfg`` entries on ``*gdn_state_quantizer`` / ``*gdn_w_quantizer`` such as the + ``configs/ptq/units/gdn_state_fp8_dynamic`` and ``gdn_w_fp8_dynamic`` recipe units. The state + quantizer carries the fused QDQ configuration. State supports dynamic E4M3 or signed + narrow-range INT8; W supports dynamic E4M3. + Both sites require identity STE. The execution policy is saved in ModelOpt metadata. + """ + + @property + def gdn_state_qdq_block_v(self) -> int: + """Value-column scale grouping from the saved execution policy.""" + return self.linear_attention_config.state.block_v + + def _state_quantized_chunk_gated_delta_rule( + self, gated_delta_rule: GatedDeltaRuleFn, *args: Any, **kwargs: Any + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Call ``gated_delta_rule`` or, if a quantizer is on, the vendored quantizing copy.""" + self.validate_linear_attention() + quantize_state = self.gdn_state_quantizer.is_enabled and self.gdn_state_quantizer._if_quant + quantize_w = self.gdn_w_quantizer.is_enabled + if not self.linear_attention_is_enabled: + return gated_delta_rule(*args, **kwargs) + while isinstance(gated_delta_rule, partial): + args = (*gated_delta_rule.args, *args) + kwargs = {**gated_delta_rule.keywords, **kwargs} + gated_delta_rule = gated_delta_rule.func + if gated_delta_rule is not _fla_chunk_gated_delta_rule(): + raise NotImplementedError( + "GatedDeltaNet quantization supports only FLA's chunk_gated_delta_rule " + f"callable or a functools.partial of it; got {gated_delta_rule!r}." + ) + chunk_size = kwargs.pop("chunk_size", self.linear_attention_config.chunk_size) + if chunk_size != self.linear_attention_config.chunk_size: + raise ValueError("GDN fake quantization supports only chunk_size=64") + if self.linear_attention_config.backend == "matmul": + return matmul_gdn( + *args, + sites=self.linear_attn_sites, + policy=self.linear_attention_config, + w_quantizer=self.gdn_w_quantizer, + state_qdq=quantize_state, + state_format=self._linear_attn_state_format, + chunk_size=chunk_size, + prefill_lengths=self._linear_attention_prefill_lengths, + **kwargs, + ) + return _state_qdq_chunk_gated_delta_rule()( + *args, + chunk_size=chunk_size, + state_qdq=(2 if self._linear_attn_state_format == "int8" else 1) + if quantize_state + else 0, + state_qdq_block_v=self.gdn_state_qdq_block_v, + w_quantizer=self.gdn_w_quantizer if quantize_w else None, + **kwargs, + ) diff --git a/modelopt/torch/quantization/plugins/kda.py b/modelopt/torch/quantization/plugins/kda.py new file mode 100644 index 00000000000..9c0ead2fe9f --- /dev/null +++ b/modelopt/torch/quantization/plugins/kda.py @@ -0,0 +1,105 @@ +# 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, + sites=self.linear_attn_sites, + policy=self.linear_attention_config, + w_quantizer=self.kda_w_quantizer, + state_qdq=self.kda_state_quantizer.is_enabled and self.kda_state_quantizer._if_quant, + 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..36d5b1c08c0 --- /dev/null +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared module policy and checkpoint support for linear-attention QAT.""" + +from ..config import QuantizerAttributeConfig +from ..linear_attention.config import LinearAttentionConfig +from ..linear_attention.matmul import LinearAttentionMatmulSites, _validate_operand_quantizer +from ..linear_attention.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) + self._register_temp_attribute("linear_attn_sites", LinearAttentionMatmulSites()) + + @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_attn_sites.is_enabled + or bool(self.linear_attention_config.matmul or self.linear_attention_config.elementwise) + or self.linear_attention_config.solve.method != "exact" + or self.linear_attention_config.decode is not None + ) + + def validate_linear_attention(self): + """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_attention_config.backend == "matmul": + _validate_operand_quantizer( + self._linear_attn_w, self.linear_attention_quantizer_names[1] + ) + elif self._linear_attn_w.is_enabled: + validate_gdn_quantizer( + self._linear_attn_w, name=self.linear_attention_quantizer_names[1] + ) + self.linear_attn_sites.validate() + if self.linear_attn_sites.is_enabled and self.linear_attention_config.backend != "matmul": + raise ValueError("Additional linear-attention operand sites require backend='matmul'") + + def modelopt_post_restore(self, prefix=""): + """Validate the restored numerical policy and quantizers.""" + super().modelopt_post_restore(prefix) + self.validate_linear_attention() diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index ec1958649dd..49195b44d29 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -21,7 +21,7 @@ import textwrap import types from contextlib import contextmanager -from functools import cache +from functools import cache, partial from typing import Any import megatron.core.parallel_state as mcore_parallel @@ -61,6 +61,7 @@ from ..utils import sync_moe_expert_amax from ..utils.layerwise_calib import LayerActivationCollector from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear +from .gdn import GatedDeltaNetStateQuantMixin try: from megatron.core.extensions.transformer_engine import ( @@ -79,6 +80,13 @@ except ImportError: HAS_TE = False +try: + from megatron.core.ssm.gated_delta_net import GatedDeltaNet + + HAS_GDN = True +except ImportError: + HAS_GDN = False + __all__ = [] @@ -1069,6 +1077,55 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): return make_sharded_tensors_for_checkpoint(state_dict, prefix, {}, sharded_offsets) +if HAS_GDN: + + @QuantModuleRegistry.register({GatedDeltaNet: "megatron_GatedDeltaNet"}) + class _QuantGatedDeltaNet(GatedDeltaNetStateQuantMixin): + """GatedDeltaNet with fake quantization of the recurrent state at kernel chunk boundaries. + + Routes ``self.gated_delta_rule`` through state/W QDQ from both Megatron's direct + forward and older split-forward layouts. The older dynamic-batching inference + paths (``ssm_prefill`` / ``ssm_decode``) are left untouched. + """ + + def _setup(self): + super()._setup() + try: + data_parallel_group = get_data_parallel_group(with_context_parallel=True) + except AssertionError: + data_parallel_group = get_data_parallel_group() + self.parallel_state = ParallelState( + data_parallel_group, + mcore_parallel.get_tensor_model_parallel_group(), + ) + + @contextmanager + def _quantized_gdn_kernel(self): + gated_delta_rule = self.gated_delta_rule + self.gated_delta_rule = partial( + self._state_quantized_chunk_gated_delta_rule, gated_delta_rule + ) + try: + yield + finally: + self.gated_delta_rule = gated_delta_rule + + def forward(self, *args, **kwargs): + if hasattr(GatedDeltaNet, "forward_pre_attn_and_core_attn"): + return super().forward(*args, **kwargs) + with self._quantized_gdn_kernel(): + return super().forward(*args, **kwargs) + + def forward_pre_attn_and_core_attn(self, *args, **kwargs): + with self._quantized_gdn_kernel(): + return super().forward_pre_attn_and_core_attn(*args, **kwargs) + + def validate_linear_attention(self): + super().validate_linear_attention() + if self.config.context_parallel_size > 1 and self.linear_attention_is_enabled: + raise NotImplementedError("GDN QAT does not support Megatron context parallelism.") + + def _is_supported_megatron_model(model: torch.nn.Module) -> bool: return isinstance(model, MegatronModule) diff --git a/modelopt/torch/quantization/plugins/vllm_linear_attention.py b/modelopt/torch/quantization/plugins/vllm_linear_attention.py new file mode 100644 index 00000000000..af98837dd95 --- /dev/null +++ b/modelopt/torch/quantization/plugins/vllm_linear_attention.py @@ -0,0 +1,165 @@ +# 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. + +"""Recurrent-state QDQ at native vLLM prefill and decode call boundaries.""" + +from functools import partial +from types import FunctionType + +import torch +import vllm +from packaging.version import Version +from vllm.forward_context import get_forward_context + +from ..linear_attention.config import LinearAttentionConfig +from ..nn import QuantModuleRegistry +from .custom import CUSTOM_MODEL_PLUGINS +from .linear_attention import _LinearAttentionQuantMixin + +__all__ = ["bind_vllm_linear_attention"] + + +def bind_vllm_linear_attention(model, model_runner): + """Validate state-only adapters before calibration or generation.""" + layers = [m for m in model.modules() if isinstance(m, _QuantVllmLinearAttention)] + for layer in layers: + layer.validate_linear_attention() + if any(layer.linear_attention_is_enabled for layer in layers): + _validate_runtime(model_runner) + return layers + + +def _validate_runtime(runner): + version = Version(vllm.__version__) + if version.release[:2] != (0, 15): + raise NotImplementedError("Linear-attention fakequant currently requires vLLM 0.15.x V1") + from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator + + shape = MambaStateShapeCalculator.gated_delta_net_state_shape(1, 2, 2, 32, 64, 4)[-1] + if shape[-2:] != (32, 64): + raise NotImplementedError( + "Linear-attention fakequant requires the key-first vLLM state ABI" + ) + config = runner.vllm_config + if not config.model_config.enforce_eager or config.scheduler_config.async_scheduling: + raise ValueError("State fakequant requires --enforce-eager --no-async-scheduling") + if config.speculative_config is not None or config.cache_config.enable_prefix_caching: + raise ValueError("Disable speculative decoding and prefix caching for state fakequant") + if config.cache_config.mamba_cache_mode != "none": + raise ValueError("State fakequant requires mamba_cache_mode='none'") + if config.kv_transfer_config is not None or config.ec_transfer_config is not None: + raise ValueError("State fakequant does not support state transfer") + parallel = config.parallel_config + if any( + getattr(parallel, name, 1) != 1 + for name in ( + "pipeline_parallel_size", + "data_parallel_size", + "prefill_context_parallel_size", + "decode_context_parallel_size", + ) + ): + raise ValueError("State fakequant supports TP with PP=1, DP=1, and CP=1") + + +class _QuantVllmLinearAttention(_LinearAttentionQuantMixin): + def validate_linear_attention(self): + super().validate_linear_attention() + if ( + self._linear_attn_w.is_enabled + or self.linear_attention_config != LinearAttentionConfig() + ): + raise ValueError( + "vLLM supports only state TensorQuantizer at native call boundaries; " + "prefill operand, arithmetic, and decode policies are unsupported" + ) + + def _quantized_state_call(self, native, *args, **kwargs): + state = kwargs.get("initial_state") + if state is not None and self._linear_attn_state.is_enabled: + if state.dtype != torch.float32: + raise ValueError("State fakequant requires an FP32 recurrent cache") + indices = kwargs.get("ssm_state_indices") + if indices is None: + # Native prefill has already initialized fresh requests to zero. + kwargs["initial_state"] = self._linear_attn_state(state) + else: + # Eager, non-speculative decode has one cache slot per active sequence. + indices = indices[: kwargs["cu_seqlens"].numel() - 1].long() + state.index_copy_( + 0, indices, self._linear_attn_state(state.index_select(0, indices)) + ) + return native(*args, **kwargs) + + def _forward_with_state_qdq(self, original, kernel_names, *args, **kwargs): + if not self.linear_attention_is_enabled: + return original(*args, **kwargs) + context = get_forward_context() + if context.attn_metadata is None: + return original(*args, **kwargs) + if context.attn_metadata[self.prefix].spec_sequence_masks is not None: + raise ValueError("Speculative metadata is unsupported for state fakequant") + function = original.__func__ + replacements = { + name: partial(self._quantized_state_call, function.__globals__[name]) + for name in kernel_names + } + # Bind wrappers for this invocation only; every wrapper calls the original kernel. + forward = FunctionType( + function.__code__, + {**function.__globals__, **replacements}, + function.__name__, + function.__defaults__, + function.__closure__, + ) + forward.__kwdefaults__ = function.__kwdefaults__ + return forward(self, *args, **kwargs) + + +class _QuantVllmGDN(_QuantVllmLinearAttention): + def _forward_core(self, *args, **kwargs): + return self._forward_with_state_qdq( + super()._forward_core, + ("chunk_gated_delta_rule", "fused_recurrent_gated_delta_rule"), + *args, + **kwargs, + ) + + +class _QuantVllmKDA(_QuantVllmLinearAttention): + linear_attention_quantizer_names = ("kda_state_quantizer", "kda_w_quantizer") + + def _forward(self, *args, **kwargs): + return self._forward_with_state_qdq( + super()._forward, + ("chunk_kda", "fused_recurrent_kda"), + *args, + **kwargs, + ) + + +def _register_vllm_linear_attention(model): + adapters = { + ("vllm.model_executor.models.qwen3_next", "Qwen3NextGatedDeltaNet"): _QuantVllmGDN, + ("vllm.model_executor.layers.kda", "KimiDeltaAttention"): _QuantVllmKDA, + } + for module in model.modules(): + cls = type(module) + adapter = adapters.get((cls.__module__, cls.__name__)) + if adapter is not None and cls not in QuantModuleRegistry: + QuantModuleRegistry.register({cls: f"vllm_{cls.__name__}"})(adapter) + + +CUSTOM_MODEL_PLUGINS.add(_register_vllm_linear_attention) diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 20e083491aa..67513218137 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -50,7 +50,8 @@ def _fp8_eager(x, amax=None): epsilon = 1.0 / (1 << 24) zero_amax_mask = amax <= epsilon safe_amax = torch.where(zero_amax_mask, torch.ones_like(amax), amax) - scale = 448.0 / safe_amax + # Use direct division, matching the CUDA extension rather than reciprocal multiplication. + scale = torch.div(448.0, safe_amax) scale_inv = 1 / scale x = (x.to(torch.float32) * scale).clamp(min=-448.0, max=448.0) x = x.to(torch.float8_e4m3fn) diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index db37b9222ca..f00932169ca 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -32,3 +32,8 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `experts_nvfp4.yaml` | NVFP4 W4A4 on `*.experts.*` weight/input quantizers | | `mixer_mlp_nvfp4.yaml` | NVFP4 W4A4 on dense `*.mixer.{up,down}_proj` weight/input quantizers | | `attention_qkv_fp8.yaml` | FP8 E4M3 on attention q/k/v bmm and softmax quantizers | +| `gdn_state_fp8_dynamic.yaml` | FP8 E4M3 dynamic fake quantization of the GatedDeltaNet recurrent state (per sequence, head and 64-column tile) at every kernel chunk boundary; needs `fla-core==0.5.1`, Triton, and SM89+ | +| `gdn_state_int8_dynamic.yaml` | INT8 quantizer unit; use the [complete recipe](../../../general/ptq/gdn_state_int8_dynamic.yaml) to enable INT8 state QDQ and Hadamard together | +| `gdn_w_fp8_dynamic.yaml` | FP8 E4M3 dynamic (per token and head) fake quantization of the WY tensor `w` that multiplies the GatedDeltaNet state; requires identity STE, `fla-core==0.5.1`, and Triton | +| `gdn_prefill_fp8_dynamic.yaml` | Dynamic per-row E4M3 operand fake QDQ for all GDN prefill matmuls; requires the explicit `linear_attention` `matmul` backend | +| `gdn_prefill_nvfp4_dynamic.yaml` | Dynamic block-16 NVFP4 operand fake QDQ for all GDN prefill matmuls; requires the explicit `linear_attention` `matmul` backend | diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index 2fee6279973..3559145b48d 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -16,6 +16,17 @@ # QuantizerCfgList snippet for standard module patterns that should remain unquantized. # modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + # Linear-attention numerical emulation is opt-in, independently of projection quantization. + - quantizer_name: '*gdn_state_quantizer' + enable: false + - quantizer_name: '*gdn_w_quantizer' + enable: false + - quantizer_name: '*kda_state_quantizer' + enable: false + - quantizer_name: '*kda_w_quantizer' + enable: false + - quantizer_name: '*linear_attn_sites.*' + enable: false - quantizer_name: '*block_sparse_moe.gate*' enable: false - quantizer_name: '*linear_attn.conv1d*' diff --git a/modelopt_recipes/configs/ptq/units/gdn_prefill_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_prefill_fp8_dynamic.yaml new file mode 100644 index 00000000000..c9ab34d8ffa --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/gdn_prefill_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. + +# Dynamic per-row E4M3 operand QDQ at all eight GDN prefill matmuls. +# Requires linear_attention backend="matmul" on matching GDN modules. +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*linear_attn_sites.*' + cfg: + num_bits: [4, 3] + type: dynamic + axis: [0, 1, 2] + - quantizer_name: '*gdn_w_quantizer' + cfg: + num_bits: [4, 3] + type: dynamic + axis: [0, 1, 2] diff --git a/modelopt_recipes/configs/ptq/units/gdn_prefill_nvfp4_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_prefill_nvfp4_dynamic.yaml new file mode 100644 index 00000000000..a274affcb46 --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/gdn_prefill_nvfp4_dynamic.yaml @@ -0,0 +1,34 @@ +# 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 block-16 NVFP4 operand QDQ at all eight GDN prefill matmuls. +# Requires linear_attention backend="matmul" on matching GDN modules. +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*linear_attn_sites.*' + cfg: + num_bits: [2, 1] + type: dynamic + block_sizes: + -1: 16 + type: dynamic + scale_bits: [4, 3] + - quantizer_name: '*gdn_w_quantizer' + cfg: + num_bits: [2, 1] + type: dynamic + block_sizes: + -1: 16 + type: dynamic + scale_bits: [4, 3] diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml new file mode 100644 index 00000000000..4d4b13278ba --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Dynamic per-tile FP8 E4M3 fake quantization of the GatedDeltaNet recurrent state. +# Recompute amax over each [K, BV] tile inside the kernel, independently per sequence and head +# of the [N, H, K, V] state (no calibration). ``axis: [0, 1]`` specifies sequence/head grouping; +# BV is ``gdn_state_qdq_block_v`` (default: the kernel's 64-column tile, +# two scales per head for V = 128; set 128 on the quantized module for one scale per head). This +# selects the FP8 state format. Applied at every kernel chunk boundary +# (64 tokens) during calibration, QAT and QAD; needs fla-core==0.5.1 and Triton. +# See ``modelopt.torch.quantization.plugins.gated_delta_net``. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*gdn_state_quantizer' + cfg: + num_bits: e4m3 + axis: [0, 1] + type: dynamic diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml new file mode 100644 index 00000000000..e86390cfef5 --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Quantizer-only INT8 unit; general/ptq/gdn_state_int8_dynamic pairs it with Hadamard. +# The execution policy selects state scale groups and chunk/token/replay boundaries. +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*gdn_state_quantizer' + cfg: + num_bits: 8 + unsigned: false + narrow_range: true + axis: [0, 1] + type: dynamic + pass_through_bwd: true diff --git a/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml new file mode 100644 index 00000000000..2c73fdd6eaa --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# FP8 E4M3 fake quantization of ``w``, the WY-transformed keys that multiply the GatedDeltaNet +# recurrent state inside the chunked kernel, with a dynamic scale per token and head +# (``w`` is ``[B, T, H, K]``; ``axis: [0, 1, 2]`` reduces over K). Pairs with +# ``gdn_state_fp8_dynamic`` to emulate an FP8 x FP8 state matmul. Unlike the state quantizer this +# one runs on the tensor itself. The initial integration supports dynamic E4M3 with identity +# STE only. Needs fla-core==0.5.1 and Triton. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*gdn_w_quantizer' + cfg: + num_bits: e4m3 + axis: [0, 1, 2] + type: dynamic diff --git a/modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml b/modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml new file mode 100644 index 00000000000..e960282121b --- /dev/null +++ b/modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + gdn_state_int8_dynamic: configs/ptq/units/gdn_state_int8_dynamic + +metadata: + recipe_type: ptq + description: >- + Dynamic INT8 GDN state QDQ with 32-value Hadamard rotation during decode. + Prefill state remains unquantized; supply per-sequence prefix lengths through + linear_attention_training_phase. No calibration is required. +quantize: + algorithm: + quant_cfg: + - $import: base_disable_all + - $import: gdn_state_int8_dynamic + linear_attention: + - module_name: '*' + cfg: + backend: matmul + decode: + mode: token + state_codec: int8_hadamard32 + prefill_state_qdq: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 2ad373ffcab..53da4aac5fc 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 26 general/ptq/ recipes (click to expand) +All 27 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -58,6 +58,7 @@ supported combinations. | `int4_blockwise_weight_only` | INT4 W4A16, block 128, weights only | none | max | | `nvfp4_mlp_weight_only` | NVFP4 W4A16 (block 32), MLP + MoE weights only | none | max | | `mxfp4_mlp_weight_only` | MXFP4 W4A16, MLP + MoE weights only | none | none (no calibration) | +| `gdn_state_int8_dynamic` | GDN decode state INT8 + Hadamard; weights unchanged | none | none (dynamic scales; requires a prefix/decode phase context) |
diff --git a/noxfile.py b/noxfile.py index 135465ed445..cd9e96ecffa 100644 --- a/noxfile.py +++ b/noxfile.py @@ -130,6 +130,15 @@ def gpu(session): "mamba_ssm", "causal-conv1d", ) + session.run( + "python", + "-m", + "pip", + "install", + "fla-core==0.5.1", + "tilelang==0.1.8", + "apache-tvm-ffi==0.1.9", + ) session.run("python", "-m", "pytest", "tests/gpu", *_cov_args()) @@ -141,6 +150,15 @@ def gpu_megatron(session): # Pre-installed nvidia-modelopt shadows the editable install session.run("pip", "uninstall", "-y", "nvidia-modelopt") session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]") + session.run( + "python", + "-m", + "pip", + "install", + "fla-core==0.5.1", + "tilelang==0.1.8", + "apache-tvm-ffi==0.1.9", + ) session.run("python", "-m", "pytest", "tests/gpu_megatron", *_cov_args()) diff --git a/pyproject.toml b/pyproject.toml index 0886e0380b3..47c8fbdf5ab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -248,6 +248,7 @@ extend-ignore = [ "tests/*" = ["B017", "D", "E402", "PT012"] "plugins/modelopt/skills/*/tests/test_*.py" = ["D", "E402"] # Skill test scripts: docstring (D) + sys.path import-order (E402) exemptions "*/_[a-zA-Z]*" = ["D"] # Private packages (_abc/*.py) or modules (_xyz.py) +"modelopt/torch/kernels/quantization/linear_attention/fla_*.py" = ["D", "E501"] # Vendored flash-linear-attention kernels "*.ipynb" = ["D", "E501"] # Ignore missing docstrings or line length for Jupyter notebooks "modelopt/torch/kernels/*" = ["N803", "N806", "E731"] # triton style "modelopt/torch/puzzletron/*" = [ @@ -322,6 +323,15 @@ disable_error_code = ["attr-defined"] module = ["examples.diffusers.fastgen.preprocess.*"] ignore_errors = true +# Vendored from fla-org/flash-linear-attention (MIT); kept faithful to upstream rather than +# annotated to modelopt's strict mypy (Triton loop indices, ``*tensor.shape`` unpacking). +[[tool.mypy.overrides]] +module = [ + "modelopt.torch.kernels.quantization.linear_attention.fla_chunk_delta_h", + "modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule", +] +ignore_errors = true + [tool.bandit] exclude_dirs = [".github/", "examples/", "noxfile.py", "tests/"] # Do not change `skips`. It should be consistent with NVIDIA's Wheel-CI-CD bandit.yml config. diff --git a/tests/_test_utils/torch/linear_attention.py b/tests/_test_utils/torch/linear_attention.py new file mode 100644 index 00000000000..303ba1fc344 --- /dev/null +++ b/tests/_test_utils/torch/linear_attention.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. + +"""Independent numerical oracles for materialized linear-attention sites.""" + +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.linear_attention import matmul_gdn +from modelopt.torch.quantization.linear_attention.matmul import LinearAttentionMatmulSites +from modelopt.torch.quantization.nn import TensorQuantizer + +FP8 = {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)} + + +def inputs(device="cpu", *, packed=False, dtype=torch.float32, state_v_first=False): + torch.manual_seed(113) + shape = (1, 73, 1, 16) + q, k = [F.normalize(torch.randn(shape, device=device, dtype=dtype), dim=-1) for _ in range(2)] + v = torch.randn(1, 73, 2, 16, device=device, dtype=dtype) + g = -torch.rand(1, 73, 2, device=device, dtype=dtype) * 0.05 + beta = torch.rand_like(g) * 0.4 + state = torch.randn(2 if packed else 1, 2, 16, 16, device=device, dtype=dtype) * 0.1 + if state_v_first: + state = state.transpose(-1, -2).contiguous() + return [x.requires_grad_() for x in (q, k, v, g, beta)], state.requires_grad_() + + +def fp8_oracle(x): + with torch.no_grad(): + amax = x.float().abs().amax(-1, keepdim=True) + scale = torch.div(448.0, torch.where(amax > 2**-24, amax, torch.ones_like(amax))) + value = (x.float() * scale).clamp(-448, 448).to( + torch.float8_e4m3fn + ).float() * scale.reciprocal() + return x + (value.to(x.dtype) - x).detach() + + +def nvfp4_oracle(x, tensor_amax=None): + """Independent block-16 E2M1 rounding with E4M3 scales and a tensor amax.""" + with torch.no_grad(): + width = x.shape[-1] + padded = F.pad(x.float(), (0, (-width) % 16)) + blocks = padded.reshape(*padded.shape[:-1], -1, 16) + local_amax = blocks.abs().amax(-1, keepdim=True) + global_amax = blocks.abs().amax() if tensor_amax is None else tensor_amax + safe_global = torch.where(global_amax > 0, global_amax, torch.ones_like(global_amax)) + two_level = (6 / safe_global).double() * 448 + scaled = ((local_amax / 6).double() * two_level).float().clamp(0, 448) + unscale64 = scaled.to(torch.float8_e4m3fn).double() / two_level + unscale64 = torch.where(local_amax > 0, unscale64, torch.ones_like(unscale64)) + scale = unscale64.reciprocal().float() + normalized = blocks * scale + # Even codes precede odd codes so argmin implements round-to-nearest-even ties. + grid = x.new_tensor([0.0, 1.0, 2.0, 4.0, 0.5, 1.5, 3.0, 6.0], dtype=torch.float32) + index = (normalized.abs().unsqueeze(-1) - grid).abs().argmin(-1) + rounded = grid[index] * normalized.sign() * unscale64.float() + value = rounded.reshape(padded.shape)[..., :width].to(x.dtype) + return x + (value - x).detach() + + +def reference_matmul(enabled, policy): + def mm(name, lhs, rhs): + if f"{name}.lhs_quantizer" in enabled: + lhs = fp8_oracle(lhs) + if f"{name}.rhs_quantizer" in enabled: + rhs = fp8_oracle(rhs.transpose(-1, -2)).transpose(-1, -2) + cfg = policy.matmul.get(name) + if cfg is None or cfg.accumulator_dtype is None: + return torch.einsum("hij,hjk->hik", lhs, rhs) + result = lhs.new_zeros((*lhs.shape[:-1], rhs.shape[-1])) + for start in range(0, lhs.shape[-1], cfg.reduction_block): + end = start + cfg.reduction_block + result = result + torch.einsum( + "hij,hjk->hik", lhs[..., start:end], rhs[..., start:end, :] + ) + rounded = result.to(getattr(torch, cfg.accumulator_dtype)).to(lhs.dtype) + result = result + (rounded - result).detach() + return result + + return mm + + +def values_and_gradients(result, args, state): + probes = [ + torch.linspace(-0.7, 0.9, x.numel(), device=x.device, dtype=x.dtype).reshape(x.shape) + for x in result + ] + grad = torch.autograd.grad( + sum((x * p).sum() for x, p in zip(result, probes)), (*args, state), retain_graph=True + ) + return (*result, *grad) + + +def check_prefill_case( + policy, + enabled=(), + *, + device="cpu", + packed=False, + state_v_first=False, + state_qdq=False, + w_qdq=False, + relative_tolerances=None, +): + args, state = inputs(device, packed=packed, state_v_first=state_v_first) + sites = LinearAttentionMatmulSites() + for name in enabled: + sites.get_submodule(name).set_from_attribute_config(FP8) + sites.get_submodule(name).enable() + w = TensorQuantizer(QuantizerAttributeConfig(**FP8, enable=w_qdq)) + kwargs = { + "initial_state": state, + "state_v_first": state_v_first, + "cu_seqlens": torch.tensor([0, 5, 73], device=device) if packed else None, + } + actual = matmul_gdn( + *args, + sites=sites, + policy=policy, + w_quantizer=w, + state_qdq=state_qdq, + output_final_state=True, + **kwargs, + ) + + def arithmetic(name, x): + target = policy.elementwise.get(name) + if target is None: + return x + return x + (x.to(getattr(torch, target)).to(x.dtype) - x).detach() + + expected = chunk_gdn_reference( + *args, + matmul=reference_matmul(enabled, policy), + arithmetic=arithmetic, + w_quantizer=fp8_oracle if w_qdq else None, + state_qdq=state_qdq, + state_qdq_block_v=policy.state.block_v, + **kwargs, + ) + for i, (a, e) in enumerate( + zip(values_and_gradients(actual, args, state), values_and_gradients(expected, args, state)) + ): + if relative_tolerances is None: + torch.testing.assert_close(a, e, rtol=5e-4, atol=3e-6) + else: + tolerance = relative_tolerances[i >= 2] + error = (a - e).norm() / e.norm().clamp_min(1e-6) + assert error <= tolerance, (i, error.item(), tolerance) + if enabled or w_qdq or state_qdq or policy.matmul or policy.elementwise: + baseline = chunk_gdn_reference(*args, **kwargs) + assert any(not torch.equal(a, b) for a, b in zip(actual, baseline)) diff --git a/tests/_test_utils/torch/quantization/linear_attention_reference.py b/tests/_test_utils/torch/quantization/linear_attention_reference.py new file mode 100644 index 00000000000..57d7448bcb1 --- /dev/null +++ b/tests/_test_utils/torch/quantization/linear_attention_reference.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Small differentiable oracles, independent of FLA and Triton. + +These references prioritize explicit arithmetic and gradients over speed. Gates are +natural logarithms; inputs are already normalized/activated. Packed sequence boundaries +are read on the CPU. Accumulation follows the input dtype (use float64 for algebra tests). +""" + +from collections.abc import Callable +from itertools import pairwise + +import torch + +__all__ = [ + "chunk_gdn_reference", + "chunk_kda_reference", + "recurrent_delta_rule_reference", + "state_fp8_qdq_reference", + "state_qdq_reference", +] + + +def state_fp8_qdq_reference(state: torch.Tensor, block_v: int = 64) -> torch.Tensor: + """Dynamic E4M3 QDQ with identity STE and one scale per ``[Dk, block_v]`` tile. + + ``state`` is in key-first layout ``[..., Dk, Dv]``. Scales and rounding do not + contribute derivatives. FP8 uses TensorQuantizer small-range handling; partial tiles are valid. + """ + return state_qdq_reference(state, block_v) + + +def state_qdq_reference(state: torch.Tensor, block_v: int = 64, state_format: str = "fp8_e4m3"): + """Dynamic state-tile QDQ with detached scales and identity STE.""" + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") + if block_v not in (16, 32, 64, 128): + raise ValueError("block_v must be 16, 32, 64, or 128") + with torch.no_grad(): + rounded = [] + for tile in state.float().split(block_v, dim=-1): + amax = tile.abs().amax(dim=(-2, -1), keepdim=True) + if state_format == "int8": + scale = torch.where(amax > 0, amax / 127.0, torch.ones_like(amax)) + codes = (tile / scale).clamp(-127, 127).round() + else: + safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) + quant_scale = torch.div(448.0, safe_amax) + scale = quant_scale.reciprocal() + codes = (tile * quant_scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() + rounded.append(codes * scale) + quantized = torch.cat(rounded, dim=-1).to(state.dtype) + return state + (quantized - state).detach() + + +def _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first): + if q.ndim != 4 or k.shape != q.shape or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k must have shape [B,T,Hk,Dk] and v shape [B,T,Hv,Dv]") + batch, length, heads, keys = q.shape + value_heads, values = v.shape[2:] + if length == 0 or value_heads % heads: + raise ValueError("nonempty sequences and Hv divisible by Hk are required") + if beta.shape != (batch, length, value_heads) or g.shape not in ( + beta.shape, + (*beta.shape, keys), + ): + raise ValueError("beta must be [B,T,Hv]; g must be [B,T,Hv] or [B,T,Hv,Dk]") + if cu_seqlens is None: + sequences = [(b, 0, length) for b in range(batch)] + else: + boundaries = cu_seqlens.tolist() + if ( + batch != 1 + or len(boundaries) < 2 + or boundaries[0] != 0 + or boundaries[-1] != length + or any(a >= b for a, b in pairwise(boundaries)) + ): + raise ValueError("cu_seqlens must partition a packed batch of size one") + sequences = [(0, a, b) for a, b in pairwise(boundaries)] + if initial_state is None: + state = q.new_zeros(len(sequences), value_heads, keys, values) + else: + state = initial_state.transpose(-1, -2) if state_v_first else initial_state + if state.shape != (len(sequences), value_heads, keys, values): + raise ValueError("initial_state shape does not match sequence/head dimensions") + return ( + q.repeat_interleave(value_heads // heads, dim=2), + k.repeat_interleave(value_heads // heads, dim=2), + state, + sequences, + ) + + +def recurrent_delta_rule_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + state_v_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """GDN (scalar decay) or KDA (per-key decay) exact token recurrence. + + Inputs have shapes ``q,k:[B,T,Hk,Dk]``, ``v:[B,T,Hv,Dv]``, + ``beta:[B,T,Hv]``, and ``g:[B,T,Hv]`` (GDN) or ``[B,T,Hv,Dk]`` (KDA). + State is ``[N,Hv,Dk,Dv]``, or ``[N,Hv,Dv,Dk]`` with ``state_v_first``. + Returns output and final state, preserving gradients through the initial state. + """ + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + sequence_output = [] + for t in range(start, end): + decay = g[b, t].exp() + if g.ndim == 3: + decay = decay.unsqueeze(-1) + decayed = state * decay.unsqueeze(-1) + residual = v[b, t] - (k[b, t].unsqueeze(-1) * decayed).sum(-2) + update = beta[b, t].unsqueeze(-1) * residual + state = decayed + k[b, t].unsqueeze(-1) * update.unsqueeze(-2) + sequence_output.append((q[b, t].unsqueeze(-1) * state).sum(-2) * scale) + outputs.append(torch.stack(sequence_output)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final + + +def chunk_gdn_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + chunk_size: int = 64, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + state_v_first: bool = False, + state_qdq: bool = False, + state_qdq_block_v: int = 64, + state_format: str = "fp8_e4m3", + w_quantizer: Callable[[torch.Tensor], torch.Tensor] | None = None, + matmul: Callable[[str, torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + arithmetic: Callable[[str, torch.Tensor], torch.Tensor] | None = None, + inverse_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Exact GDN chunk algebra with optional state/W fake quantization. + + The solve is a unit-lower triangular solve. ``w_quantizer`` sees the complete + materialized ``[B,T,Hv,Dk]`` WY operand once, with its own autograd semantics. + State QDQ occurs on the initial state and each chunk's final state, after readout. + Optional callbacks expose the eight named matmuls and elementwise boundaries. + Matmul callbacks receive conventional ``lhs @ rhs`` shapes, without padding. + """ + if g.ndim != 3 or chunk_size <= 0: + raise ValueError("chunk_gdn_reference requires scalar GDN gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + mm = matmul if matmul is not None else lambda name, lhs, rhs: lhs @ rhs + arithmetic = arithmetic if arithmetic is not None else lambda name, x: x + chunks, all_w = [], [] + for n, (b, start, end) in enumerate(sequences): + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + gc = arithmetic("gate_prefix", g[b, lo:hi].transpose(0, 1).cumsum(-1)) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + # Mask before exp: upper-triangle positive differences can overflow for long decay. + causal = torch.ones(hi - lo, hi - lo, device=q.device, dtype=torch.bool).tril() + decay = arithmetic( + "gate_exp", (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() + ) + if matmul is None: + lower = (bc * (kc @ kc.transpose(-1, -2)) * decay).tril(-1) + else: + lower = (mm("key_interaction", bc * kc, kc.transpose(-1, -2)) * decay).tril(-1) + matrix = lower + torch.eye(hi - lo, device=q.device, dtype=q.dtype) + gate = arithmetic("gate_exp", gc.exp()).unsqueeze(-1) + if matmul is None and inverse_fn is None: + rhs = torch.cat((bc * vc, bc * kc * gate), dim=-1) + solved = torch.linalg.solve_triangular(matrix, rhs, upper=False, unitriangular=True) + u, w = solved.split((v.shape[-1], k.shape[-1]), dim=-1) + else: + inverse = ( + inverse_fn(lower) + if inverse_fn is not None + else torch.linalg.solve_triangular( + matrix, + torch.eye(hi - lo, device=q.device, dtype=q.dtype).expand_as(matrix), + upper=False, + unitriangular=True, + ) + ) + u = mm("wy_value", inverse, bc * vc) + w = mm("wy_key", inverse, bc * kc * gate) + chunks.append((n, qc, kc, gc, decay, u)) + all_w.append(w.transpose(0, 1)) + w = torch.cat(all_w).reshape(q.shape) + if w_quantizer is not None: + w = w_quantizer(w) + w = w.reshape(-1, *w.shape[2:]) + outputs, finals, offset, previous_n = [], [], 0, -1 + for n, qc, kc, gc, decay, u in chunks: + if n != previous_n: + state = states[n] + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + previous_n = n + length = qc.shape[1] + wc = w[offset : offset + length].transpose(0, 1) + offset += length + updated_values = arithmetic("value_residual", u - mm("state_read", wc, state)) + local_scores = (mm("output_score", qc * scale, kc.transpose(-1, -2)) * decay).tril() + output = mm( + "output_state", qc * (scale * arithmetic("gate_exp", gc.exp())).unsqueeze(-1), state + ) + output = arithmetic("output_add", output + mm("output_value", local_scores, updated_values)) + outputs.append(output.transpose(0, 1)) + weighted_keys = kc * arithmetic("gate_exp", (gc[..., -1:] - gc).exp()).unsqueeze(-1) + state = arithmetic( + "state_decay", state * arithmetic("gate_exp", gc[..., -1].exp())[:, None, None] + ) + state = arithmetic( + "state_add", state + mm("state_update", weighted_keys.transpose(-1, -2), updated_values) + ) + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + if len(finals) <= n: + finals.append(state) + else: + finals[n] = state + output = torch.cat(outputs).reshape(*v.shape) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final + + +def chunk_kda_reference( + q, + k, + v, + g, + beta, + *, + chunk_size=64, + scale=None, + initial_state=None, + cu_seqlens=None, + state_v_first=False, + state_qdq=False, + state_qdq_block_v=64, + state_format="fp8_e4m3", + w_quantizer=None, + matmul=None, + inverse_fn=None, +): + """KDA chunk oracle with causal per-channel decays and optional operand QDQ. + + The eight-site callback receives conventional three-dimensional matmul + operands. Interaction sites are invoked one query row at a time with decay + already applied to the right-hand keys. This avoids inverse-prefix factors. + """ + if g.ndim != 4 or chunk_size <= 0: + raise ValueError("chunk_kda_reference requires per-key gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + mm = matmul if matmul is not None else lambda name, lhs, rhs: lhs @ rhs + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + pieces = [] + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + prefix = g[b, lo:hi].transpose(0, 1).cumsum(-2) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + gate = prefix.exp() + lower_rows, score_rows = [], [] + for row in range(hi - lo): + decayed_keys = ( + kc[..., : row + 1, :] + * (prefix[..., row : row + 1, :] - prefix[..., : row + 1, :]).exp() + ) + right = decayed_keys.transpose(-1, -2) + left = kc[..., row : row + 1, :] * bc[..., row : row + 1, :] + interaction = mm("key_interaction", left, right) + score = mm("output_score", qc[..., row : row + 1, :] * scale, right) + # Zero padding follows the computation and never contributes a scale statistic. + padding = hi - lo - row - 1 + lower_rows.append(torch.nn.functional.pad(interaction, (0, padding))) + score_rows.append(torch.nn.functional.pad(score, (0, padding))) + lower = torch.cat(lower_rows, dim=-2).tril(-1) + scores = torch.cat(score_rows, dim=-2) + identity = torch.eye(hi - lo, device=q.device, dtype=q.dtype).expand_as(lower) + inverse = ( + inverse_fn(lower) + if inverse_fn is not None + else torch.linalg.solve_triangular( + identity + lower, identity, upper=False, unitriangular=True + ) + ) + u = mm("wy_value", inverse, bc * vc) + w = mm("wy_key", inverse, bc * kc * gate) + if w_quantizer is not None: + w = w_quantizer(w) + updated = u - mm("state_read", w, state) + pieces.append( + ( + mm("output_state", qc * scale * gate, state) + + mm("output_value", scores, updated) + ).transpose(0, 1) + ) + weighted_keys = kc * (prefix[..., -1:, :] - prefix).exp() + state = state * gate[..., -1, :, None] + mm( + "state_update", weighted_keys.transpose(-1, -2), updated + ) + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + outputs.append(torch.cat(pieces)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py new file mode 100644 index 00000000000..17cf0112640 --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py @@ -0,0 +1,474 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_dynamic_qdq +from modelopt.torch.kernels.quantization.linear_attention.decode import fused_recurrence +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + LinearAttentionMatmulSites, + matmul_gdn, + matmul_kda, +) +from modelopt.torch.quantization.linear_attention.config import LinearAttentionDecodeConfig +from modelopt.torch.quantization.linear_attention.decode import ( + _encode, + recurrent_decode, + recurrent_decode_reference, +) +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("state_format", [None, "fp8_e4m3", "int8"]) +@pytest.mark.parametrize("replay", [False, True]) +@pytest.mark.parametrize("read_stored", [False, True]) +def test_fused_recurrence_outputs_state_and_gradients(kda, state_format, replay, read_stored): + state_qdq = state_format is not None + state_format = state_format or "fp8_e4m3" + torch.manual_seed(193) + shape = (37, 2, 16) + q, k = [ + F.normalize(torch.randn(shape, device="cuda"), dim=-1).requires_grad_() for _ in range(2) + ] + v = torch.randn(37, 2, 11, device="cuda", requires_grad=True) + g = (-torch.rand(shape if kda else shape[:-1], device="cuda") * 0.03).requires_grad_() + beta = (torch.rand(shape[:-1], device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(2, 16, 11, device="cuda") * 0.1).requires_grad_() + cfg = LinearAttentionDecodeConfig( + mode="replay" if replay else "token", + quantize_initial=False, + readout="stored" if read_stored else "working", + replay={"window": 5} if replay else None, + ) + encoded_key = _encode(k, replay, 16).values + actual = fused_recurrence( + q, + encoded_key, + v, + g, + beta, + initial, + state_qdq=state_qdq, + state_format=state_format, + block_v=16, + replay=replay, + factor_qdq=replay, + window=5, + read_stored=read_stored, + ) + expected, carry = recurrent_decode_reference( + q, + k, + v, + g, + beta, + config=cfg, + initial_state=initial, + state_qdq=state_qdq, + state_format=state_format, + block_v=16, + ) + final = carry.reconstruct() + targets = (q, k, v, g, beta, initial) + probes = [torch.randn_like(x) for x in actual[:2]] + gradients = [ + torch.autograd.grad(sum((a * p).sum() for a, p in zip(result, probes)), targets) + for result in (actual[:2], (expected, final)) + ] + for a, e in zip((*actual[:2], *gradients[0]), (expected, final, *gradients[1])): + torch.testing.assert_close(a, e, rtol=1e-3, atol=8e-6) + + +@pytest.mark.parametrize("mode", ["token", "replay"]) +@pytest.mark.parametrize("interval", [3, 8]) +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_fused_carry_handoff_and_auxiliary_gradients(mode, interval, state_format): + torch.manual_seed(217) + q, k = [ + F.normalize(torch.randn(23, 2, 16, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(23, 2, 19, device="cuda", requires_grad=True) + g = (-torch.rand_like(q) * 0.05).requires_grad_() + beta = (torch.rand(23, 2, device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(2, 16, 19, device="cuda") * 0.2).requires_grad_() + args = (q, k, v, g, beta) + reference = LinearAttentionDecodeConfig( + mode=mode, decay_log_step=0.01, replay={"window": 5} if mode == "replay" else None + ) + candidate = reference.model_copy(update={"implementation": "triton"}) + expected, expected_carry = recurrent_decode_reference( + *args, + config=reference, + state_qdq=True, + state_format=state_format, + block_v=16, + initial_state=initial, + ) + carry = None + outputs = [] + start = 0 + for end in [0, 3, 11, 12, 23]: + output, carry = recurrent_decode( + *(x[start:end] for x in args), + config=candidate, + state_qdq=True, + state_format=state_format, + block_v=16, + initial_state=initial if carry is None else None, + carry=carry, + checkpoint_interval=interval, + ) + outputs.append(output) + start = end + actual = torch.cat(outputs) + actual_state, expected_state = carry.reconstruct(), expected_carry.reconstruct() + assert carry.cursor == expected_carry.cursor + assert carry.position == 23 and carry.started + torch.testing.assert_close( + carry.anchor.scales, expected_carry.anchor.scales, rtol=5e-5, atol=2e-8 + ) + for a, e in zip(carry.entries, expected_carry.entries): + torch.testing.assert_close(a.update.scales, e.update.scales, rtol=5e-5, atol=2e-8) + torch.testing.assert_close(a.key.values, e.key.values, rtol=0, atol=0) + probes = [torch.randn_like(actual), torch.randn_like(actual_state)] + gradients = [ + torch.autograd.grad(sum((a * p).sum() for a, p in zip(result, probes)), (*args, initial)) + for result in ((actual, actual_state), (expected, expected_state)) + ] + for a, e in zip( + (actual, actual_state, *gradients[0]), (expected, expected_state, *gradients[1]) + ): + torch.testing.assert_close(a, e, rtol=1e-3, atol=1e-5) + + +@triton.jit +def _codec_kernel(X, Y, S, N: tl.constexpr, BLOCK: tl.constexpr, NATIVE_FP8: tl.constexpr): # noqa: N803 + offsets = tl.arange(0, BLOCK) + values = tl.load(X + offsets, offsets < N, 0) + decoded, scale = fp8_dynamic_qdq(values, NATIVE_FP8) + tl.store(Y + offsets, decoded, offsets < N) + tl.store(S, scale) + + +@pytest.mark.parametrize("native_fp8", [False, True]) +@pytest.mark.parametrize( + "scale", + [0.0, 1e-8 / 448, 2.0**-24 / 448, 2.0**-23 / 448, 2.0**-12, 1.0, 2.0**12, 0.00314159, 731.173], +) +def test_fp8_codec_matches_tensor_quantizer_boundaries(scale, native_fp8): + if native_fp8 and torch.cuda.get_device_capability() < (8, 9): + pytest.skip("Native FP8 conversion requires SM89 or later") + # Exhaust all finite positive E4M3 values and their rounding ties, plus signs. + representable = torch.arange(127, dtype=torch.uint8).view(torch.float8_e4m3fn).float() + ties = (representable[:-1] + representable[1:]) / 2 + probe = torch.cat( + ( + representable, + ties, + torch.nextafter(ties, torch.full_like(ties, float("inf"))), + torch.nextafter(ties, torch.zeros_like(ties)), + ) + ) + probe = torch.cat((probe, -probe)) * scale + values = probe.cuda() + actual = torch.empty_like(values) + actual_scale = torch.empty((), device="cuda") + _codec_kernel[(1,)]( + values, + actual, + actual_scale, + len(values), + triton.next_power_of_2(len(values)), + native_fp8, + ) + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), type="dynamic")) + torch.testing.assert_close(actual, quantizer(values), rtol=0, atol=0) + amax = values.abs().amax() + safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) + torch.testing.assert_close( + actual_scale, torch.div(448.0, safe_amax).reciprocal(), rtol=0, atol=0 + ) + + +def test_fused_fp8_state_matches_tensor_quantizer_with_partial_tiles(): + torch.manual_seed(762) + amplitudes = torch.tensor([0.0, 1e-8, 2**-24, 2**-23, 1.0, 731.173], device="cuda") + state = (torch.rand(6, 3, 19, device="cuda") * 2 - 1) * amplitudes[:, None, None] + state[:, 0, 0] = amplitudes + state.requires_grad_() + q = torch.zeros(1, 6, 3, device="cuda") + v = torch.zeros(1, 6, 19, device="cuda") + gate = torch.zeros(1, 6, device="cuda") + _, actual, *_ = fused_recurrence(q, q, v, gate, gate, state, state_qdq=True, block_v=16) + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=(4, 3), type="dynamic", axis=0, pass_through_bwd=True) + ) + expected = torch.cat( + [quantizer(tile.flatten(-2)).reshape_as(tile) for tile in state.split(16, -1)], -1 + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + probe = torch.randn_like(state) + gradients = [ + torch.autograd.grad((value * probe).sum(), state)[0] for value in (actual, expected) + ] + torch.testing.assert_close(*gradients, rtol=0, atol=0) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("mode", ["token", "replay"]) +@pytest.mark.parametrize("horizon", ["near_unit", "mixed"]) +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_wide_long_trajectory(kda, mode, horizon, state_format): + torch.manual_seed(329) + q, k = [ + F.normalize(torch.randn(257, 2, 128, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(257, 2, 70, device="cuda", requires_grad=True) + shape = q.shape if kda else q.shape[:-1] + g = -torch.rand(shape, device="cuda") * (1e-4 if horizon == "near_unit" else 3) + if horizon == "mixed": + g[::7] = -30 + g.requires_grad_() + beta = (torch.rand(257, 2, device="cuda") * 0.6).requires_grad_() + initial = (torch.randn(2, 128, 70, device="cuda") * 0.1).requires_grad_() + cfg = LinearAttentionDecodeConfig(mode=mode, replay={"window": 8} if mode == "replay" else None) + args = q, k, v, g, beta + expected, ec = recurrent_decode_reference( + *args, + config=cfg, + state_qdq=True, + state_format=state_format, + block_v=64, + initial_state=initial, + ) + actual, ac = recurrent_decode( + *args, + config=cfg.model_copy(update={"implementation": "triton"}), + state_qdq=True, + state_format=state_format, + block_v=64, + initial_state=initial, + ) + probes = [torch.randn_like(actual), torch.randn_like(ac.reconstruct())] + gradients = [ + torch.autograd.grad(sum((a * p).sum() for a, p in zip(result, probes)), (*args, initial)) + for result in ((actual, ac.reconstruct()), (expected, ec.reconstruct())) + ] + # Fixed reduction and scale arithmetic must also agree over long trajectories. + for a, e in zip( + (actual, ac.reconstruct(), *gradients[0]), (expected, ec.reconstruct(), *gradients[1]) + ): + relative = (a - e).norm() / e.norm().clamp_min(1e-8) + assert relative < 1e-4, relative.item() + assert torch.isfinite(a).all() + + +@pytest.mark.parametrize("block_v", [32, 128]) +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_full_key_width_and_value_block_tail(block_v, kda, state_format): + torch.manual_seed(418) + q, k = [ + F.normalize(torch.randn(9, 2, 128, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(9, 2, 133, device="cuda", requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1], device="cuda") * 0.1).requires_grad_() + beta = torch.rand(9, 2, device="cuda", requires_grad=True) + initial = torch.randn(2, 128, 133, device="cuda", requires_grad=True) + args = q, k, v, g, beta + cfg = LinearAttentionDecodeConfig(mode="replay", replay={"window": 3}) + results = [] + for implementation in ("torch", "triton"): + output, carry = recurrent_decode( + *args, + config=cfg.model_copy(update={"implementation": implementation}), + state_qdq=True, + state_format=state_format, + block_v=block_v, + initial_state=initial, + ) + final = carry.reconstruct() + gradients = torch.autograd.grad( + output.square().sum() + final.square().sum(), (*args, initial) + ) + results.append((output, final, *gradients)) + for actual, expected in zip(*results): + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=2e-5) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize( + ("mode", "block_v", "readout", "factor_qdq"), + [ + ("token", 32, "stored", False), + ("replay", 64, "working", False), + ("replay", 128, "stored", True), + ], +) +@pytest.mark.timeout(180) +def test_hadamard_fused_split_carry_and_gradients(kda, mode, block_v, readout, factor_qdq): + torch.manual_seed(917) + q, k = [ + F.normalize(torch.randn(37, 2, 17, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(37, 2, 96, device="cuda", requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1], device="cuda") * 0.03).requires_grad_() + beta = (torch.rand(37, 2, device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(2, 17, 96, device="cuda") * 0.1).requires_grad_() + args = (q, k, v, g, beta) + cfg = LinearAttentionDecodeConfig( + mode=mode, + readout=readout, + state_codec="int8_hadamard32", + decay_log_step=1 / 256, + replay={"window": 8, "factor_qdq": factor_qdq} if mode == "replay" else None, + ) + shared = {"state_qdq": True, "state_format": "int8", "block_v": block_v} + expected, ref_carry = recurrent_decode_reference( + *args, config=cfg, initial_state=initial, **shared + ) + carry, start, pieces = None, 0, [] + for end in (0, 2, 8, 8, 19, 37): + output, carry = recurrent_decode( + *(x[start:end] for x in args), + config=cfg.model_copy(update={"implementation": "triton"}), + initial_state=initial if carry is None else None, + carry=carry, + **shared, + ) + pieces.append(output) + start = end + actual = torch.cat(pieces) + assert carry.value_basis == "hadamard32" and carry.cursor == ref_carry.cursor + assert carry.anchor.scales.shape == (2, 17, 3) and carry.anchor.scales.dtype == torch.float16 + torch.testing.assert_close(carry.anchor.scales, ref_carry.anchor.scales, rtol=0, atol=0) + actual_tensors = ( + actual, + carry.reconstruct(), + carry.anchor.values, + *(e.update.values for e in carry.entries), + ) + expected_tensors = ( + expected, + ref_carry.reconstruct(), + ref_carry.anchor.values, + *(e.update.values for e in ref_carry.entries), + ) + probes = [torch.randn_like(x) for x in expected_tensors] + gradients = [ + torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probes)), (*args, initial)) + for result in (actual_tensors, expected_tensors) + ] + for a, e in zip((*actual_tensors, *gradients[0]), (*expected_tensors, *gradients[1])): + torch.testing.assert_close(a, e, rtol=1e-3, atol=1e-5) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.timeout(180) +def test_hadamard_long_trajectory(kda): + torch.manual_seed(991) + q, k = [ + F.normalize(torch.randn(257, 1, 128, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(257, 1, 64, device="cuda", requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1], device="cuda") * 1e-4).requires_grad_() + beta = (torch.rand(257, 1, device="cuda") * 0.1).requires_grad_() + initial = (torch.randn(1, 128, 64, device="cuda") * 0.1).requires_grad_() + args = (q, k, v, g, beta) + results = [] + for implementation in ("torch", "triton"): + cfg = LinearAttentionDecodeConfig( + mode="replay", + implementation=implementation, + state_codec="int8_hadamard32", + replay={"window": 8, "factor_qdq": False}, + ) + out, carry = recurrent_decode( + *args, config=cfg, initial_state=initial, state_qdq=True, state_format="int8" + ) + final = carry.reconstruct() + grads = torch.autograd.grad(out.square().sum() + final.square().sum(), (*args, initial)) + results.append((out, final, *grads)) + for a, e in zip(*results): + assert torch.isfinite(a).all() and torch.isfinite(e).all() + assert (a - e).norm() / e.norm().clamp_min(1e-8) < 1e-4 + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_packed_quantized_prefix_and_replay_composition(kda, state_format): + torch.manual_seed(591) + q, k = [ + F.normalize(torch.randn(1, 73, 1, 16, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(1, 73, 2, 19, device="cuda", requires_grad=True) + g = (-torch.rand((1, 73, 2, 16) if kda else (1, 73, 2), device="cuda") * 0.03).requires_grad_() + beta = (torch.rand(1, 73, 2, device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(3, 2, 19, 16, device="cuda") * 0.1).requires_grad_() + sites = LinearAttentionMatmulSites() + attributes = {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)} + for quantizer in sites.modules(): + if isinstance(quantizer, TensorQuantizer): + quantizer.set_from_attribute_config(attributes) + quantizer.enable() + w = TensorQuantizer(QuantizerAttributeConfig(**attributes)) + outputs = [] + function = matmul_kda if kda else matmul_gdn + args = q, k, v, g, beta + for implementation in ("torch", "triton"): + policy = LinearAttentionConfig( + backend="matmul", + state={"block_v": 16}, + elementwise={"value_residual": "bfloat16"}, + decode={ + "mode": "replay", + "implementation": implementation, + "prefill_state_qdq": True, + "decay_log_step": 1 / 256, + "replay": {"window": 5}, + }, + ) + output, final = function( + *args, + initial_state=initial, + state_v_first=True, + cu_seqlens=torch.tensor([0, 0, 9, 73]), + prefill_lengths=[0, 3, 33], + output_final_state=True, + sites=sites, + w_quantizer=w, + policy=policy, + state_qdq=True, + state_format=state_format, + ) + gradients = torch.autograd.grad( + output.square().sum() + final.square().sum(), (*args, initial) + ) + outputs.append((output, final, *gradients)) + for actual, expected in zip(*outputs): + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=2e-5) diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py new file mode 100644 index 00000000000..84b344ba224 --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py @@ -0,0 +1,403 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Numerical forward/backward coverage for the GDN training emulation contract.""" + +from functools import partial + +import pytest +import torch +import torch.nn.functional as F +from _test_utils.torch.quantization.linear_attention_reference import ( + chunk_gdn_reference, + state_fp8_qdq_reference, +) +from torch.utils.checkpoint import checkpoint + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.plugins.gdn import GatedDeltaNetStateQuantMixin + +fla = pytest.importorskip("fla.ops.gated_delta_rule") +from fla.utils import IS_NVIDIA_HOPPER, TRITON_ABOVE_3_4_0 + +from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( + chunk_gated_delta_rule, +) + +DTYPES = (torch.float32, torch.bfloat16) +QDQ_MODES = ((False, False), (False, True), (True, False), (True, True), (2, False), (2, True)) +STATE_LAYOUTS = ((False, False), (True, True)) +W_AXES = (None, (0, 1), (0, 1, 2)) + + +def require_state_qdq(): + if torch.cuda.get_device_capability() < (8, 9): + pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") + + +def make_inputs(dtype=None, packed=False, state_v_first=False, grouped=True): + hopper_backward = IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0 + if dtype is None: + dtype = torch.bfloat16 if hopper_backward else torch.float32 + if dtype == torch.float32 and hopper_backward: + pytest.skip("Hopper/TileLang training supports BF16; FP32 is rejected before launch") + torch.manual_seed(123) + # Two full chunks plus a tail; grouped-value heads and a partial state scale tile. + batch, length, heads, value_heads, keys, values = 1, 145, 1, 2, 32, 48 + heads = heads if grouped else value_heads + q, k = [ + F.normalize(torch.randn(batch, length, heads, keys, device="cuda"), dim=-1) + for _ in range(2) + ] + v = torch.randn(batch, length, value_heads, values, device="cuda") + g = -torch.rand(batch, length, value_heads, device="cuda") * 0.1 + beta = torch.rand_like(g) + state = torch.randn(2 if packed else batch, value_heads, keys, values, device="cuda") * 0.1 + if state_v_first: + state = state.transpose(-1, -2).contiguous() + args = [x.to(dtype).requires_grad_() for x in (q, k, v, g, beta)] + return args, state.requires_grad_() + + +def w_quantizer(axis=(0, 1, 2)): + return TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), axis=axis, type="dynamic")) + + +def compare(actual, expected, tolerance): + for a, e in zip(actual, expected): + assert torch.isfinite(a).all() + error = (a.float() - e.float()).norm() + bound = tolerance * e.float().norm().clamp_min(1e-6) + assert error <= bound, ( + f"relative L2 error {(error / e.float().norm()).item():.5g} > {tolerance}" + ) + + +def compare_forward(actual, expected, dtype, state_qdq): + tolerance = 0.03 if dtype == torch.bfloat16 else 0.01 + compare(actual[:1], expected[:1], tolerance) + # FP8 bin crossings amplify small GEMM differences in the stored final state. + state_tolerance = max(tolerance, 0.02) if state_qdq else tolerance + compare(actual[1:], expected[1:], state_tolerance) + + +def values_and_grads(fn, args, state, **kwargs): + result = fn(*args, initial_state=state, **kwargs) + torch.manual_seed(15) + probes = [torch.randn(x.shape, device=x.device, dtype=torch.float32) for x in result] + grads = torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probes)), (*args, state)) + return result, grads + + +@pytest.fixture(scope="module", autouse=True) +def _compile_gdn_kernels(): + """Keep compilation/autotuning in setup, including when selecting a single test.""" + supports_state_qdq = torch.cuda.get_device_capability() >= (8, 9) + for dtype in DTYPES: + if dtype == torch.float32 and IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0: + continue + args, state = make_inputs(dtype, grouped=False) + for kernel in (fla.chunk_gated_delta_rule, chunk_gated_delta_rule): + values_and_grads(kernel, args, state, output_final_state=True) + for packed, state_v_first in STATE_LAYOUTS: + args, state = make_inputs(dtype, packed=packed, state_v_first=state_v_first) + for state_qdq, quantize_w in QDQ_MODES: + if state_qdq == 1 and not supports_state_qdq: + continue + values_and_grads( + chunk_gated_delta_rule, + args, + state, + output_final_state=True, + state_qdq=state_qdq, + state_qdq_block_v=32, + w_quantizer=w_quantizer() if quantize_w else None, + state_v_first=state_v_first, + cu_seqlens=torch.tensor([0, 67, 145], device="cuda", dtype=torch.int32) + if packed + else None, + ) + + args, state = make_inputs() + quantizer = w_quantizer() + for state_qdq in (0, 2): + values_and_grads( + chunk_gated_delta_rule, + args, + state, + output_final_state=True, + w_quantizer=quantizer, + state_qdq=state_qdq, + ) + for axis in W_AXES: + output, _ = chunk_gated_delta_rule(*args, w_quantizer=w_quantizer(axis)) + torch.autograd.grad(output.square().sum(), args) + a_log = torch.zeros(2, device="cuda", requires_grad=True) + bias = torch.zeros_like(a_log, requires_grad=True) + output, final = chunk_gated_delta_rule( + *args, + initial_state=state, + output_final_state=True, + w_quantizer=quantizer, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + A_log=a_log, + dt_bias=bias, + use_beta_sigmoid_in_kernel=True, + allow_neg_eigval=True, + ) + torch.autograd.grad(output.square().sum() + final.square().sum(), (*args, state, a_log, bias)) + for state_qdq in (1, 2): + if state_qdq == 1 and not supports_state_qdq: + continue + for block_v in (16, 64, 128): + values_and_grads( + chunk_gated_delta_rule, + args, + state, + output_final_state=True, + state_qdq=state_qdq, + state_qdq_block_v=block_v, + ) + if supports_state_qdq: + args, state = make_inputs(packed=True, state_v_first=True) + for block_v in (16, 32, 64, 128): + _, final = chunk_gated_delta_rule( + *args, + initial_state=state, + output_final_state=True, + state_v_first=True, + cu_seqlens=torch.tensor([0, 67, 145], device="cuda", dtype=torch.int32), + state_qdq=True, + state_qdq_block_v=block_v, + ) + torch.autograd.grad(final.sum(), state) + torch.cuda.synchronize() + + +@pytest.mark.parametrize("dtype", DTYPES) +def test_disabled_matches_upstream_forward_and_backward(dtype): + # Upstream Hopper/TileLang backward cannot handle grouped value heads. + args, state = make_inputs(dtype, grouped=False) + expected = values_and_grads(fla.chunk_gated_delta_rule, args, state, output_final_state=True) + actual = values_and_grads(chunk_gated_delta_rule, args, state, output_final_state=True) + for a, e in zip(actual, expected): + compare(a, e, 0.005 if dtype == torch.bfloat16 else 0.001) + + +@pytest.mark.parametrize(("state_qdq", "quantize_w"), QDQ_MODES) +@pytest.mark.parametrize(("packed", "state_v_first"), STATE_LAYOUTS) +@pytest.mark.parametrize("dtype", DTYPES) +def test_numerical_surrogate_reference(state_qdq, quantize_w, packed, state_v_first, dtype): + if state_qdq == 1: + require_state_qdq() + args, state = make_inputs(dtype, packed=packed, state_v_first=state_v_first) + quantizer = w_quantizer() if quantize_w else None + kwargs = { + "state_qdq": state_qdq, + "state_qdq_block_v": 32, + "w_quantizer": quantizer, + "state_v_first": state_v_first, + "cu_seqlens": torch.tensor([0, 67, 145], device="cuda", dtype=torch.int32) + if packed + else None, + } + reference_args = [x.detach().float().requires_grad_() for x in args] + reference_state = state.detach().float().requires_grad_() + expected = values_and_grads( + chunk_gdn_reference, + reference_args, + reference_state, + state_format="int8" if state_qdq == 2 else "fp8_e4m3", + **kwargs, + ) + actual = values_and_grads( + chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs + ) + compare_forward(actual[0], expected[0], dtype, state_qdq) + compare(actual[1], expected[1], 0.05 if dtype == torch.bfloat16 else 0.02) + + +@pytest.mark.parametrize("state_qdq", [0, 2]) +def test_w_qdq_saved_once_and_activation_checkpoint_parity(state_qdq): + args, state = make_inputs() + quantizer = w_quantizer() + calls = [] + handle = quantizer.register_forward_hook(lambda *args: calls.append(1)) + + def fn(*x): + return chunk_gated_delta_rule( + *x[:5], + initial_state=x[5], + output_final_state=True, + w_quantizer=quantizer, + state_qdq=state_qdq, + ) + + output = fn(*args, state) + grads = torch.autograd.grad(sum(x.square().sum() for x in output), (*args, state)) + assert len(calls) == 1, "Backward must not rerun the quantizer or its observers" + recomputed = checkpoint(fn, *args, state, use_reentrant=False) + checkpoint_grads = torch.autograd.grad( + sum(x.square().sum() for x in recomputed), (*args, state) + ) + compare(recomputed, output, 1e-6) + compare(checkpoint_grads, grads, 1e-6) + handle.remove() + + +@pytest.mark.parametrize("axis", W_AXES) +def test_zero_initial_state_and_output_only_training_loss(axis): + args, _ = make_inputs() + quantizer = w_quantizer(axis) + expected, _ = chunk_gdn_reference(*(x.float() for x in args), w_quantizer=quantizer) + actual, final = chunk_gated_delta_rule(*args, w_quantizer=quantizer) + assert final is None + compare((actual,), (expected,), 0.03 if args[0].dtype == torch.bfloat16 else 0.01) + grads = [torch.autograd.grad(x.square().sum(), args) for x in (actual, expected)] + compare(*grads, 0.05 if args[0].dtype == torch.bfloat16 else 0.02) + + +@pytest.mark.parametrize("block_v", [16, 64, 128]) +@pytest.mark.parametrize("state_qdq", [1, 2]) +def test_state_scale_tile_forward_and_backward(block_v, state_qdq): + if state_qdq == 1: + require_state_qdq() + args, state = make_inputs() + kwargs = {"state_qdq": state_qdq, "state_qdq_block_v": block_v} + reference_args = [x.detach().float().requires_grad_() for x in args] + expected = values_and_grads( + chunk_gdn_reference, + reference_args, + state, + state_format="int8" if state_qdq == 2 else "fp8_e4m3", + **kwargs, + ) + actual = values_and_grads( + chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs + ) + compare_forward(actual[0], expected[0], args[0].dtype, state_qdq=True) + compare(actual[1], expected[1], 0.05 if args[0].dtype == torch.bfloat16 else 0.02) + + +@pytest.mark.parametrize("block_v", [16, 32, 64, 128]) +def test_state_qdq_without_recurrent_updates(block_v): + require_state_qdq() + args, state = make_inputs(packed=True, state_v_first=True) + q, k, v, g, beta = args + _, final = chunk_gated_delta_rule( + q, + torch.zeros_like(k), + torch.zeros_like(v), + torch.zeros_like(g), + beta, + initial_state=state, + output_final_state=True, + state_v_first=True, + cu_seqlens=torch.tensor([0, 67, 145], device="cuda", dtype=torch.int32), + state_qdq=True, + state_qdq_block_v=block_v, + ) + expected = state_fp8_qdq_reference(state.transpose(-1, -2), block_v).transpose(-1, -2) + torch.testing.assert_close(final, expected, rtol=1e-6, atol=1e-7) + (grad,) = torch.autograd.grad(final.sum(), state) + torch.testing.assert_close(grad, torch.ones_like(state), rtol=0, atol=0) + + +def test_fused_gate_normalization_and_beta_gradients(): + args, state = make_inputs() + q, k, v, raw_g, raw_beta = args + a_log = torch.randn(2, device="cuda", requires_grad=True) + bias = torch.randn(2, device="cuda", requires_grad=True) + quantizer = w_quantizer() + expected = chunk_gdn_reference( + F.normalize(q.float(), dim=-1), + F.normalize(k.float(), dim=-1), + v.float(), + -a_log.exp() * F.softplus(raw_g.float() + bias), + 2 * raw_beta.float().sigmoid(), + initial_state=state, + w_quantizer=quantizer, + ) + actual = chunk_gated_delta_rule( + *args, + initial_state=state, + output_final_state=True, + w_quantizer=quantizer, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + A_log=a_log, + dt_bias=bias, + use_beta_sigmoid_in_kernel=True, + allow_neg_eigval=True, + ) + compare(actual, expected, 0.03 if q.dtype == torch.bfloat16 else 0.01) + grads = [ + torch.autograd.grad(sum(x.square().sum() for x in result), (*args, state, a_log, bias)) + for result in (actual, expected) + ] + compare(*grads, 0.05 if q.dtype == torch.bfloat16 else 0.02) + + +def test_plugin_dispatch_preserves_fla_partial(): + args, state = make_inputs() + module = GatedDeltaNetStateQuantMixin.convert(torch.nn.Module()) + module.gdn_w_quantizer = w_quantizer() + expected = values_and_grads( + chunk_gated_delta_rule, + args, + state, + output_final_state=True, + w_quantizer=module.gdn_w_quantizer, + ) + + def wrapped(q, *remaining, **kwargs): + kernel = partial(fla.chunk_gated_delta_rule, q, output_final_state=True) + return module._state_quantized_chunk_gated_delta_rule(kernel, *remaining, **kwargs) + + actual = values_and_grads(wrapped, args, state) + for a, e in zip(actual, expected): + compare(a, e, 1e-6) + + +@pytest.mark.parametrize("option", ["use_gate_in_kernal", "unsupported_option"]) +def test_unknown_keyword_rejected_before_launch(option): + args, _ = make_inputs() + with pytest.raises(TypeError, match=f"Unexpected keyword arguments: {option}"): + chunk_gated_delta_rule(*args, **{option: True}) + + +@pytest.mark.parametrize("chunk_size", [16, 32, 128]) +def test_unsupported_chunk_rejected_before_launch(chunk_size): + args, state = make_inputs() + with pytest.raises(ValueError, match="only chunk_size=64"): + chunk_gated_delta_rule(*args, chunk_size=chunk_size) + + +def test_unsupported_w_gradient_rejected_before_launch(): + args, _ = make_inputs() + quantizer = w_quantizer() + quantizer.set_from_attribute_config({"pass_through_bwd": False}) + with pytest.raises(ValueError, match="pass_through_bwd=True"): + chunk_gated_delta_rule(*args, w_quantizer=quantizer) + + +def test_hopper_fp32_rejected_before_launch(): + if not (IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0): + pytest.skip("Hopper/TileLang capability guard") + args, _ = make_inputs() + with pytest.raises(ValueError, match="requires BF16 q/k/v"): + chunk_gated_delta_rule(*(x.float() for x in args)) diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_matmul_gdn.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_matmul_gdn.py new file mode 100644 index 00000000000..3879eae6342 --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_matmul_gdn.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +from _test_utils.torch.linear_attention import ( + FP8, + check_prefill_case, + inputs, + nvfp4_oracle, + values_and_gradients, +) +from _test_utils.torch.quantization.linear_attention_reference import chunk_gdn_reference +from torch.utils.checkpoint import checkpoint + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import LinearAttentionConfig, matmul_gdn +from modelopt.torch.quantization.linear_attention.matmul import LinearAttentionMatmulSites +from modelopt.torch.quantization.nn import TensorQuantizer + +HANDLES = [ + name + for name, module in LinearAttentionMatmulSites().named_modules() + if name.endswith("quantizer") +] +NVFP4 = { + "num_bits": (2, 1), + "type": "dynamic", + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, +} + + +@pytest.fixture(autouse=True) +def full_precision_matmul(): + """Keep operand-format comparisons independent of the container TF32 default.""" + previous = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") + yield + torch.set_float32_matmul_precision(previous) + + +@pytest.mark.parametrize("handle", HANDLES) +def test_cuda_fp8_each_operand(handle): + check_prefill_case(LinearAttentionConfig(backend="matmul"), [handle], device="cuda") + + +def test_cuda_composed_arithmetic_state_and_w(): + check_prefill_case( + LinearAttentionConfig( + backend="matmul", + state={"block_v": 16}, + matmul={"output_value": {"accumulator_dtype": "float16", "reduction_block": 16}}, + elementwise={"value_residual": "bfloat16"}, + ), + # The established FP32 state-QDQ bounds allow GEMM roundoff crossing a QDQ tie. + HANDLES, + device="cuda", + packed=True, + state_v_first=True, + state_qdq=True, + w_qdq=True, + relative_tolerances=(0.01, 0.02), + ) + + +@pytest.mark.parametrize("handle", [*HANDLES, "state_read.lhs_quantizer"]) +def test_cuda_nvfp4_operand_and_surrogate(handle): + """Independent codec, GEMM, and operand VJPs at each actual site shape.""" + args, state = inputs("cuda") + sites = LinearAttentionMatmulSites() + w = TensorQuantizer(QuantizerAttributeConfig(enable=False)) + target = w if handle == "state_read.lhs_quantizer" else sites.get_submodule(handle) + target.set_from_attribute_config(NVFP4) + target.enable() + captures = [] + observer = target.register_forward_hook( + lambda module, args, result: captures.append((args[0], result)) + ) + result = matmul_gdn( + *args, + sites=sites, + policy=LinearAttentionConfig(backend="matmul"), + w_quantizer=w, + initial_state=state, + output_final_state=True, + ) + observer.remove() + assert captures + for value, actual in captures: + expected = nvfp4_oracle(value) + torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-6) + probe = torch.randn_like(actual) + actual_grad = torch.autograd.grad((actual * probe).sum(), value, retain_graph=True)[0] + torch.testing.assert_close(actual_grad, probe, rtol=0, atol=0) + # Match the declared global scale domain: local stages batch all chunks, recurrent + # stages invoke the quantizer once per chunk over the sequence batch. + tensor_amaxes = [value.detach().abs().amax() for value, _ in captures] + target_site, side = handle.split(".") + call_index = 0 + + def reference_mm(name, lhs, rhs): + nonlocal call_index + if name == target_site: + index = call_index if len(tensor_amaxes) > 1 else 0 + call_index += 1 + if side == "lhs_quantizer": + lhs = nvfp4_oracle(lhs, tensor_amaxes[index]) + else: + rhs = nvfp4_oracle(rhs.transpose(-1, -2), tensor_amaxes[index]).transpose(-1, -2) + return torch.einsum("hij,hjk->hik", lhs, rhs) + + expected = chunk_gdn_reference(*args, initial_state=state, matmul=reference_mm) + for a, e in zip( + values_and_gradients(result, args, state), values_and_gradients(expected, args, state) + ): + torch.testing.assert_close(a, e, rtol=5e-4, atol=3e-6) + + +def test_cuda_checkpoint_preserves_prefill_gradients(): + args, state = inputs("cuda") + sites = LinearAttentionMatmulSites() + for name in HANDLES: + sites.get_submodule(name).set_from_attribute_config(FP8) + sites.get_submodule(name).enable() + w = TensorQuantizer(QuantizerAttributeConfig(**FP8)) + policy = LinearAttentionConfig(backend="matmul") + + def run(*x): + return matmul_gdn( + *x[:-1], + sites=sites, + policy=policy, + w_quantizer=w, + initial_state=x[-1], + output_final_state=True, + ) + + plain = values_and_gradients(run(*args, state), args, state) + recomputed = values_and_gradients( + checkpoint(run, *args, state, use_reentrant=False), args, state + ) + with torch.autocast("cuda", dtype=torch.bfloat16): + mixed_result = run(*args, state) + mixed = values_and_gradients(mixed_result, args, state) + for result in (recomputed, mixed): + for a, b in zip(plain, result): + torch.testing.assert_close(a, b, rtol=0, atol=0) + + +def test_cuda_matmul_baseline_and_fused_gate_gradients(): + # FLA is optional outside this GPU qualification test. + pytest.importorskip("fla") + from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( + chunk_gated_delta_rule, + ) + + args, state = inputs("cuda", dtype=torch.bfloat16) + state = state.float().detach().requires_grad_() + a_log = torch.full((2,), -3.0, device="cuda", requires_grad=True) + bias = torch.zeros(2, device="cuda", requires_grad=True) + kwargs = { + "initial_state": state, + "output_final_state": True, + "use_qk_l2norm_in_kernel": True, + "use_gate_in_kernel": True, + "use_beta_sigmoid_in_kernel": True, + "A_log": a_log, + "dt_bias": bias, + } + actual = matmul_gdn( + *args, + sites=LinearAttentionMatmulSites(), + policy=LinearAttentionConfig(backend="matmul"), + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(enable=False)), + **kwargs, + ) + expected = chunk_gated_delta_rule(*args, **kwargs) + probes = [torch.randn_like(x) for x in actual] + gradients = [ + torch.autograd.grad( + sum((x * p).sum() for x, p in zip(result, probes)), (*args, state, a_log, bias) + ) + for result in (actual, expected) + ] + for i, (a, e) in enumerate(zip((*actual, *gradients[0]), (*expected, *gradients[1]))): + error = (a.float() - e.float()).norm() / e.float().norm().clamp_min(1e-6) + assert error <= (0.03 if i < 2 else 0.05), (i, error.item()) diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_matmul_kda.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_matmul_kda.py new file mode 100644 index 00000000000..5bc80b44aab --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_matmul_kda.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +from _test_utils.torch.linear_attention import ( + FP8, + fp8_oracle, + inputs, + nvfp4_oracle, + reference_matmul, + values_and_gradients, +) +from _test_utils.torch.quantization.linear_attention_reference import chunk_kda_reference + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + LinearAttentionMatmulSites, + matmul_kda, +) +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.fixture(autouse=True) +def full_precision_matmul(): + previous = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") + yield + torch.set_float32_matmul_precision(previous) + + +def _inputs(*, packed=False, dtype=torch.float32): + args, state = inputs("cuda", packed=packed, dtype=dtype) + args[3] = ( + -torch.rand(*args[3].shape, args[0].shape[-1], device="cuda") * 0.05 + ).requires_grad_() + return args, state.float().detach().requires_grad_() + + +@pytest.mark.timeout(300) +@pytest.mark.parametrize("lower_bound", [None, -5.0]) +@pytest.mark.parametrize("packed", [False, True]) +def test_fla_baseline_and_fused_gate_gradients(lower_bound, packed): + pytest.importorskip("fla") + # FLA is optional except for comparison with its actual CUDA implementation. + from fla.ops.kda import chunk_kda + + args, state = _inputs(packed=packed, dtype=torch.bfloat16) + a_log = torch.full((2,), -3.0, device="cuda", requires_grad=True) + bias = torch.zeros(32, device="cuda", requires_grad=True) + kwargs = { + "initial_state": state, + "output_final_state": True, + "use_qk_l2norm_in_kernel": True, + "use_gate_in_kernel": True, + "use_beta_sigmoid_in_kernel": True, + "allow_neg_eigval": True, + "A_log": a_log, + "dt_bias": bias, + "lower_bound": lower_bound, + "safe_gate": lower_bound is not None, + "cu_seqlens": torch.tensor([0, 5, 73], device="cuda", dtype=torch.int32) + if packed + else None, + } + actual = matmul_kda( + *args, + sites=LinearAttentionMatmulSites(), + policy=LinearAttentionConfig(backend="matmul"), + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(enable=False)), + **kwargs, + ) + # Compare the general FLA algebra, not its optional safe_gate TensorCore rounding. + expected = chunk_kda(*args, **{**kwargs, "safe_gate": False}) + probes = [torch.randn_like(x) for x in actual] + gradients = [ + torch.autograd.grad( + sum((x * p).sum() for x, p in zip(result, probes)), (*args, state, a_log, bias) + ) + for result in (actual, expected) + ] + for i, (a, e) in enumerate(zip((*actual, *gradients[0]), (*expected, *gradients[1]))): + error = (a.float() - e.float()).norm() / e.float().norm().clamp_min(1e-6) + assert error <= (0.03 if i < 2 else 0.05), (i, error.item()) + + +def test_composed_fp8_with_chunk_state_and_w(): + args, state = _inputs(packed=True) + sites = LinearAttentionMatmulSites() + handles = [name for name, _ in sites.named_modules() if name.endswith("quantizer")] + for name in handles: + sites.get_submodule(name).set_from_attribute_config(FP8) + sites.get_submodule(name).enable() + policy = LinearAttentionConfig(backend="matmul", state={"block_v": 16}) + boundaries = torch.tensor([0, 5, 73], device="cuda") + actual = matmul_kda( + *args, + sites=sites, + policy=policy, + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(**FP8)), + state_qdq=True, + initial_state=state, + output_final_state=True, + cu_seqlens=boundaries, + ) + expected = chunk_kda_reference( + *args, + initial_state=state, + cu_seqlens=boundaries, + state_qdq=True, + state_qdq_block_v=16, + w_quantizer=fp8_oracle, + matmul=reference_matmul(handles, policy), + ) + for i, (a, e) in enumerate( + zip(values_and_gradients(actual, args, state), values_and_gradients(expected, args, state)) + ): + error = (a - e).norm() / e.norm().clamp_min(1e-6) + assert error <= (0.01 if i < 2 else 0.02), (i, error.item()) + + +@pytest.mark.parametrize("site", ["key_interaction", "output_score"]) +@pytest.mark.parametrize("side", ["lhs_quantizer", "rhs_quantizer"]) +def test_nvfp4_channel_decay_interaction(site, side): + args, state = _inputs() + sites = LinearAttentionMatmulSites() + target = sites.get_submodule(f"{site}.{side}") + target.set_from_attribute_config( + { + "num_bits": (2, 1), + "type": "dynamic", + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + } + ) + target.enable() + captured = [] + hook = target.register_forward_hook(lambda m, a, o: captured.append((a[0], o))) + actual = matmul_kda( + *args, + sites=sites, + policy=LinearAttentionConfig(backend="matmul"), + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(enable=False)), + initial_state=state, + output_final_state=True, + ) + hook.remove() + assert len(captured) == 8 + amaxes = [] + for x, y in captured: + torch.testing.assert_close(y, nvfp4_oracle(x), rtol=2e-5, atol=2e-6) + probe = torch.randn_like(y) + torch.testing.assert_close( + torch.autograd.grad((y * probe).sum(), x, retain_graph=True)[0], probe, rtol=0, atol=0 + ) + amaxes.append(x.detach().abs().amax()) + index = 0 + + def mm(name, lhs, rhs): + nonlocal index + if name == site: + amax = amaxes[(index % 64) // 8] + index += 1 + if side == "lhs_quantizer": + lhs = nvfp4_oracle(lhs, amax) + else: + rhs = nvfp4_oracle(rhs.transpose(-1, -2), amax).transpose(-1, -2) + return torch.einsum("hij,hjk->hik", lhs, rhs) + + expected = chunk_kda_reference(*args, initial_state=state, matmul=mm) + for a, e in zip( + values_and_gradients(actual, args, state), values_and_gradients(expected, args, state) + ): + torch.testing.assert_close(a, e, rtol=5e-4, atol=3e-6) diff --git a/tests/gpu/torch/quantization/plugins/test_kda.py b/tests/gpu/torch/quantization/plugins/test_kda.py new file mode 100644 index 00000000000..291e3b57dc8 --- /dev/null +++ b/tests/gpu/torch/quantization/plugins/test_kda.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +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", + [ + "state", + "state_int8", + "w", + "prefill_fp8", + "prefill_nvfp4", + "arithmetic", + "decode_token", + "decode_replay", + "decode_token_int8", + "decode_replay_int8", + "decode_decay", + "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)} + if mode == "prefill_nvfp4": + attributes = { + "num_bits": (2, 1), + "type": "dynamic", + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + } + cfg = { + "quant_cfg": [{"quantizer_name": "*", "enable": False}], + "algorithm": None, + "linear_attention": [{"module_name": "*", "cfg": {"backend": "matmul"}}], + } + if mode in ("w", "prefill_fp8", "prefill_nvfp4"): + cfg["quant_cfg"].append({"quantizer_name": "*kda_w_quantizer", "cfg": attributes}) + if mode.startswith("prefill"): + cfg["quant_cfg"].append({"quantizer_name": "*linear_attn_sites.*", "cfg": attributes}) + state_attributes = {**attributes, "axis": (0, 1)} + if mode.endswith("int8"): + state_attributes.update(num_bits=8, unsigned=False, narrow_range=True) + if mode.startswith("state"): + cfg["quant_cfg"].append({"quantizer_name": "*kda_state_quantizer", "cfg": state_attributes}) + if mode == "arithmetic": + cfg["linear_attention"][0]["cfg"]["elementwise"] = {"value_residual": "bfloat16"} + if mode.startswith("decode"): + cfg["quant_cfg"].append({"quantizer_name": "*kda_state_quantizer", "cfg": state_attributes}) + policy_cfg = cfg["linear_attention"][0]["cfg"] + policy_cfg["state"] = {"block_v": 16} + policy_cfg["decode"] = { + "mode": "replay" if mode.startswith("decode_replay") else "token", + "implementation": "triton", + } + if mode.startswith("decode_replay"): + policy_cfg["decode"]["replay"] = {"window": 5} + if mode == "decode_decay": + policy_cfg["decode"]["decay_log_step"] = 1 / 256 + if "hadamard" in mode: + policy_cfg["state"]["block_v"] = 32 + policy_cfg["decode"]["state_codec"] = "int8_hadamard32" + mtq.quantize(model, cfg) + output = _forward(model, hidden) + assert torch.isfinite(output).all() + enabled = [q for q in model.modules() if isinstance(q, TensorQuantizer) and q.is_enabled] + policy = model.linear_attention_config + for q in enabled: + q.disable() + model.linear_attention_config = LinearAttentionConfig() + with torch.no_grad(): + torch.testing.assert_close(_forward(model, hidden), baseline, rtol=0, atol=0) + for q in enabled: + q.enable() + model.linear_attention_config = policy + checkpoint = tmp_path / "model.pth" + mto.save(model, checkpoint) + restored = mto.restore(_layer(head_dim), checkpoint) + actual = _forward(restored, hidden) + torch.testing.assert_close(actual, output, rtol=0, atol=0) + assert restored.linear_attention_config == policy + actual.float().square().mean().backward() + assert restored.q_proj.weight.grad is not None + for parameter in restored.parameters(): + assert parameter.grad is not None + assert torch.isfinite(parameter.grad).all() + if mode == "prefill_fp8": + copied = copy.deepcopy(restored) + calls = [] + hook = copied.kda_w_quantizer.register_forward_hook(lambda *args: calls.append(True)) + with torch.no_grad(): + torch.testing.assert_close(_forward(copied, hidden), actual, rtol=0, atol=0) + hook.remove() + assert calls, "a copied layer must use its own quantizer handles" + before = restored.q_proj.weight.detach().clone() + torch.optim.SGD(restored.parameters(), lr=0.1).step() + assert not torch.equal(before, restored.q_proj.weight) + restored.eval() + if mode.startswith("decode"): + with ( + linear_attention_training_phase(restored, [0]), + torch.autocast("cuda", dtype=torch.bfloat16), + ): + assert torch.isfinite(restored(hidden[:, :1])[0]).all() + with pytest.raises(NotImplementedError, match="use_cache=False"): + restored(hidden[:, :1], use_cache=True) + with ( + pytest.raises(ValueError, match="explicit"), + torch.autocast("cuda", dtype=torch.bfloat16), + ): + restored(hidden) + else: + with pytest.raises(NotImplementedError, match="chunk path"): + _forward(restored, hidden[:, :1]) + + +@pytest.mark.parametrize("hadamard", [False, True]) +def test_decode_phase_spans_activation_checkpoint_backward(hadamard): + model = _layer(32 if hadamard else 16) + quant_cfg = [{"quantizer_name": "*", "enable": False}] + if hadamard: + quant_cfg.append( + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": 8, + "type": "dynamic", + "axis": (0, 1), + "unsigned": False, + "narrow_range": True, + }, + } + ) + mtq.quantize( + model, + { + "quant_cfg": quant_cfg, + "algorithm": None, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "triton", + "state_codec": "int8_hadamard32" if hadamard else "tile", + }, + }, + } + ], + }, + ) + hidden = torch.randn(1, 73, 32, device="cuda", requires_grad=True) + with linear_attention_training_phase(model, [31]): + with torch.autocast("cuda", dtype=torch.bfloat16): + actual = checkpoint(lambda x: model(x)[0], hidden, use_reentrant=False) + actual.square().mean().backward() + assert model._linear_attention_prefill_lengths is None + assert torch.isfinite(hidden.grad).all() + assert torch.count_nonzero(model.q_proj.weight.grad) > 0 + with linear_attention_training_phase(model, [0]): + with linear_attention_training_phase(model, [73]): + assert model._linear_attention_prefill_lengths == (73,) + assert model._linear_attention_prefill_lengths == (0,) + assert model._linear_attention_prefill_lengths is None diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index cf802abaf40..2039217a8e8 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -25,7 +25,7 @@ import modelopt.torch.kernels.quantization.gemm as triton_kernel import modelopt.torch.quantization.utils as quant_utils from modelopt.torch.quantization import tensor_quant -from modelopt.torch.quantization.extensions import get_cuda_ext, get_cuda_ext_mx +from modelopt.torch.quantization.extensions import get_cuda_ext, get_cuda_ext_fp8, get_cuda_ext_mx from modelopt.torch.quantization.tensor_quant import mx_format_map if triton_kernel.IS_AVAILABLE: @@ -120,6 +120,26 @@ def test_cuda_ext_tiny_amax(self): class TestScaledE4M3: + @pytest.mark.parametrize("per_channel", [False, True]) + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_eager_matches_cuda_scale_rounding(self, per_channel, dtype): + # These ranges distinguish direct division from reciprocal multiplication. + amax = torch.tensor([0.0, 1e-8, 2**-24, 1.04207686e-7, 0.91567558], device="cuda") + values = torch.linspace(-1, 1, 129, device="cuda")[None, :] * amax[:, None] + values = values.to(dtype) + extension = get_cuda_ext_fp8() + if per_channel: + actual = extension.fake_e4m3fy_with_axis(values, amax, 0) + expected = tensor_quant.fp8_eager(values, amax[:, None]) + else: + actual = torch.stack( + [extension.fake_e4m3fy(row, limit) for row, limit in zip(values, amax)] + ) + expected = torch.stack( + [tensor_quant.fp8_eager(row, limit) for row, limit in zip(values, amax)] + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + @pytest.mark.parametrize("device", ["cuda", "cpu"]) def test_e4m3_no_scale(self, device): x = torch.randn(4, 4, device=device, dtype=torch.float32) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py new file mode 100644 index 00000000000..e825477d7e8 --- /dev/null +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from contextlib import nullcontext + +import pytest +import torch +from _test_utils.torch.megatron.models import get_mcore_gpt_model +from _test_utils.torch.megatron.utils import ( + get_forward, + initialize_for_megatron, + sharded_state_dict_test_helper, +) + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + linear_attention_training_phase, +) +from modelopt.torch.quantization.nn import TensorQuantizer + +pytest.importorskip("fla") # Megatron-Core GatedDeltaNet and the state QDQ kernel need fla +GatedDeltaNet = pytest.importorskip("megatron.core.ssm.gated_delta_net").GatedDeltaNet + +from modelopt.torch.quantization.plugins.gdn import _state_qdq_chunk_gated_delta_rule +from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet + +try: + _state_qdq_chunk_gated_delta_rule() +except RuntimeError as e: + pytest.skip(str(e), allow_module_level=True) + +SEED = 1234 + + +def _make_model(tp_size, **config_kwargs): + model = ( + get_mcore_gpt_model( + tensor_model_parallel_size=tp_size, + num_layers=2, + hidden_size=64, + num_attention_heads=4, + vocab_size=32, + max_sequence_length=128, + experimental_attention_variant="gated_delta_net", + linear_value_head_dim=64, + **config_kwargs, + ) + .cuda() + .eval() + ) + # Retain enough history for chunk-boundary state rounding to affect the next chunk. + with torch.no_grad(): + for module in model.modules(): + if isinstance(module, GatedDeltaNet): + module.A_log.fill_(-4) + module.dt_bias.zero_() + return model + + +def _gdn_config(sites): + return { + "quant_cfg": [{"quantizer_name": "*", "enable": False}] + + [ + { + "quantizer_name": f"*gdn_{site}_quantizer", + "cfg": { + "num_bits": (4, 3), + "axis": (0, 1) if site == "state" else (0, 1, 2), + "type": "dynamic", + }, + } + for site in sites + ], + "algorithm": None, + "linear_attention": [ + {"module_name": "decoder.layers.*.self_attention", "cfg": {"state": {"block_v": 16}}} + ], + } + + +def _config_for_mode(mode): + sites = ("state", "w") if mode == "both" else ((mode,) if mode in ("state", "w") else ()) + if mode in ("prefill_fp8", "prefill_nvfp4"): + sites = ("w",) + if mode.startswith("decode"): + sites = ("state",) + if mode == "state_int8": + sites = ("state",) + cfg = _gdn_config(sites) + if mode.endswith("int8"): + cfg["quant_cfg"][1]["cfg"].update(num_bits=8, unsigned=False, narrow_range=True) + if mode.startswith("decode"): + policy = cfg["linear_attention"][0]["cfg"] + policy.update(backend="matmul", decode={"implementation": "triton"}) + if mode.startswith("decode_replay"): + policy["decode"].update(mode="replay", replay={"window": 8}) + if mode.startswith("prefill"): + policy = cfg["linear_attention"][0]["cfg"] + policy["backend"] = "matmul" + if mode == "prefill_arithmetic": + policy["matmul"] = { + "output_value": {"accumulator_dtype": "float16", "reduction_block": 16} + } + policy["elementwise"] = {"value_residual": "bfloat16"} + else: + attributes = {"num_bits": (4, 3), "axis": (0, 1, 2), "type": "dynamic"} + if mode == "prefill_nvfp4": + attributes = { + "num_bits": (2, 1), + "type": "dynamic", + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + } + cfg["quant_cfg"].append({"quantizer_name": "*linear_attn_sites.*", "cfg": attributes}) + cfg["quant_cfg"].append({"quantizer_name": "*gdn_w_quantizer", "cfg": attributes}) + return cfg, sites + + +def _gdn_forward(model): + original_forward = get_forward(model) + + def forward(m): + enabled = any( + getattr(getattr(layer, "linear_attention_config", None), "decode", None) is not None + for layer in m.modules() + ) + with linear_attention_training_phase(m, [64, 64]) if enabled else nullcontext(): + return original_forward(m) + + return forward + + +def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): + initialize_for_megatron( + tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED + ) + model = _make_model(size) + forward = _gdn_forward(model) + + outputs = [] + handles = [ + module.register_forward_hook( + lambda module, args, result: outputs.append(result[0].detach().clone()) + ) + for module in model.modules() + if isinstance(module, GatedDeltaNet) + ] + with torch.no_grad(): + loss_ref = forward(model) + gdn_ref = outputs.copy() + outputs.clear() + + cfg, sites = _config_for_mode(mode) + mtq.quantize(model, cfg) + gdn_modules = [m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)] + assert gdn_modules, "no GatedDeltaNet layer was wrapped" + for module in gdn_modules: + assert module.gdn_state_qdq_block_v == 16 + assert module.gdn_state_quantizer.is_enabled == ("state" in sites) + assert module.gdn_w_quantizer.is_enabled == ("w" in sites) + # Checkpointing must save the resolved policy, including edits after conversion. + module.linear_attention_config = LinearAttentionConfig( + **{**module.linear_attention_config.model_dump(), "state": {"block_v": 32}} + ) + + with torch.no_grad(): + loss_quant = forward(model) + assert torch.isfinite(loss_quant).all() + # BF16 residual additions can hide a small GDN change at the final loss. + assert len(outputs) == len(gdn_ref) == len(gdn_modules) + assert any(not torch.equal(actual, ref) for actual, ref in zip(outputs, gdn_ref)), ( + "enabled QDQ must change the GDN branch output" + ) + for handle in handles: + handle.remove() + + enabled = [ + q + for m in gdn_modules + for q in m.modules() + if isinstance(q, TensorQuantizer) and q.is_enabled + ] + policies = [m.linear_attention_config for m in gdn_modules] + for q in enabled: + q.disable() + for module in gdn_modules: + module.linear_attention_config = LinearAttentionConfig() + with torch.no_grad(): + torch.testing.assert_close(forward(model), loss_ref, rtol=1e-4, atol=1e-4) + for q in enabled: + q.enable() + for module, policy in zip(gdn_modules, policies): + module.linear_attention_config = policy + + restored = _make_model(size) + sharded_state_dict_test_helper(checkpoint_path, model, restored, forward) + restored_gdn = [m for m in restored.modules() if isinstance(m, _QuantGatedDeltaNet)] + assert len(restored_gdn) == len(gdn_modules) + for module, original in zip(restored_gdn, gdn_modules): + assert module.linear_attention_config == original.linear_attention_config + assert module.linear_attn_sites.is_enabled == original.linear_attn_sites.is_enabled + assert module.gdn_state_quantizer.is_enabled == ("state" in sites) + assert module.gdn_w_quantizer.is_enabled == ("w" in sites) + assert module._linear_attn_state_format == original._linear_attn_state_format + assert module.gdn_state_qdq_block_v == 32 + assert module.in_proj.weight.grad is not None + assert torch.isfinite(module.in_proj.weight.grad).all() + + # The checkpoint helper already ran backward through the restored model. + optimizer = torch.optim.SGD(restored.parameters(), lr=1e-3) + before = restored_gdn[0].in_proj.weight.detach().clone() + optimizer.step() + assert not torch.equal(restored_gdn[0].in_proj.weight, before) + + +def _compile_gdn_qat_kernels(rank, size, mode): + initialize_for_megatron( + tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED + ) + model = _make_model(size) + forward = _gdn_forward(model) + with torch.no_grad(): + forward(model) + cfg, _ = _config_for_mode(mode) + # The functional case switches state tiles before running inference. + cfg["linear_attention"][0]["cfg"]["state"]["block_v"] = 32 + mtq.quantize(model, cfg) + with torch.no_grad(): + forward(model) + model.train() + forward(model).sum().backward() + torch.cuda.synchronize() + + +@pytest.fixture +def compiled_gdn_workers(request, tp_size, mode): + """Warm the selected QAT path in the same workers, outside the test-call budget.""" + if mode in ("state", "both") and torch.cuda.get_device_capability() < (8, 9): + pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") + workers = request.getfixturevalue(f"dist_workers_size_{tp_size}") + workers.run(_compile_gdn_qat_kernels, mode) + return workers + + +@pytest.mark.parametrize("tp_size", [1, 2]) +@pytest.mark.parametrize( + "mode", + [ + "state", + "w", + "both", + "prefill_fp8", + "prefill_nvfp4", + "prefill_arithmetic", + "decode_token", + "decode_replay", + "state_int8", + "decode_token_int8", + "decode_replay_int8", + ], +) +def test_gdn_qat_and_sharded_restore(compiled_gdn_workers, tmp_path, mode): + """Train through state/W QDQ after a Megatron distributed-checkpoint round trip.""" + compiled_gdn_workers.run(_test_gdn_qat_helper, mode, tmp_path) + + +def _test_gdn_context_parallel_helper(rank, size, mode): + initialize_for_megatron(context_parallel_size=size, seed=SEED) + model = _make_model(1, context_parallel_size=size) + with pytest.raises(NotImplementedError, match="context parallelism"): + mtq.quantize(model, _config_for_mode(mode)[0]) + + +@pytest.mark.parametrize( + "mode", + [ + "state", + "w", + "prefill_fp8", + "prefill_arithmetic", + "decode_token", + "decode_replay", + "state_int8", + "decode_token_int8", + "decode_replay_int8", + ], +) +def test_gdn_context_parallel_rejected(dist_workers_size_2, mode): + """Reject unqualified Megatron CP even when it does not pass an FLA CP context.""" + dist_workers_size_2.run(_test_gdn_context_parallel_helper, mode) diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py b/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py new file mode 100644 index 00000000000..76f0d5b1a7c --- /dev/null +++ b/tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py @@ -0,0 +1,315 @@ +# 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. + +"""Native vLLM state-boundary QDQ, worker reload, and TP on tiny offline models.""" + +import copy +import gc +import importlib.util +import math +import os +import shutil +from pathlib import Path + +import pytest +import torch +from packaging.version import Version +from transformers import Qwen3NextConfig +from vllm import LLM, SamplingParams +from vllm import __version__ as vllm_version +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import LinearAttentionConfig +from modelopt.torch.quantization.plugins.vllm_linear_attention import _QuantVllmLinearAttention + +pytestmark = pytest.mark.skipif( + Version(vllm_version).release[:2] != (0, 15), reason="Requires the pinned vLLM 0.15.x ABI" +) +ROOT = Path(__file__).parents[4] + + +def _example_module(name): + spec = importlib.util.spec_from_file_location( + name + "_test", ROOT / "examples/vllm_serve" / (name + ".py") + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _calibrate_worker(worker): + model = worker.model_runner.model + if hasattr(model, "unwrap"): + model = model.unwrap() + layer = next(m for m in model.modules() if isinstance(m, _QuantVllmLinearAttention)) + projection = layer.q_proj if hasattr(layer, "q_proj") else layer.in_proj_qkvz + quantizer = projection.input_quantizer + quantizer.enable() + try: + loop = _example_module("vllm_ptq_utils").calibrate_fun( + [{"input_ids": torch.tensor([[3] * 5, [4] * 5])}], worker + ) + mtq.calibrate(model, algorithm="max", forward_loop=loop) + assert torch.isfinite(quantizer.amax).all() and quantizer.amax.max() > 0 + finally: + quantizer.disable() + return True + + +def _set_state_format_worker(worker, *, fp8): + for layer in worker.model_runner.model.modules(): + if not isinstance(layer, _QuantVllmLinearAttention): + continue + layer._linear_attn_state.num_bits = (4, 3) if fp8 else 8 + layer.validate_linear_attention() + + +def _disable_worker(worker): + for layer in worker.model_runner.model.modules(): + if not isinstance(layer, _QuantVllmLinearAttention): + continue + layer._linear_attn_state.disable() + layer._linear_attn_w.disable() + layer.linear_attention_config = LinearAttentionConfig() + + +def _audit_worker(worker, *, checkpoint=None): + model = worker.model_runner.model + if hasattr(model, "unwrap"): + model = model.unwrap() + layers = [m for m in model.modules() if isinstance(m, _QuantVllmLinearAttention)] + assert layers + if checkpoint is not None and torch.distributed.get_rank() == 0: + torch.save(mto.modelopt_state(model), checkpoint) + reports = [] + for layer in layers: + assert not hasattr(layer, "_linear_attention_cache") + if not hasattr(layer, "_state_audit"): + layer._state_audit = {"prefill": 0, "decode": 0, "changed": 0, "heads": 0} + original = layer._quantized_state_call + + def checked_call(native, *args, _layer=layer, _original=original, **kwargs): + state = kwargs["initial_state"] + indices = kwargs.get("ssm_state_indices") + expected_state = state.clone() + selected = state if indices is None else state.index_select(0, indices.long()) + reference_quantizer = copy.deepcopy(_layer._linear_attn_state) + rounded = reference_quantizer(selected.clone()) + if indices is None: + expected_state = rounded + else: + expected_state.index_copy_(0, indices.long(), rounded) + # Native KDA prefill writes its output into the value-input buffer. + control_args = tuple(x.clone() if isinstance(x, torch.Tensor) else x for x in args) + control_kwargs = { + name: x.clone() if isinstance(x, torch.Tensor) else x + for name, x in kwargs.items() + } + control_kwargs["initial_state"] = expected_state + expected = native(*control_args, **control_kwargs) + native_calls = [] + + def checked_native(*native_args, **native_kwargs): + incoming = native_kwargs["initial_state"] + incoming = ( + incoming if indices is None else incoming.index_select(0, indices.long()) + ) + torch.testing.assert_close(incoming, rounded, atol=0, rtol=0) + native_calls.append(True) + return native(*native_args, **native_kwargs) + + actual = _original(checked_native, *args, **kwargs) + assert native_calls == [True] + for result, control in zip(actual, expected): + if control is not None: + torch.testing.assert_close(result, control, atol=0, rtol=0) + _layer._state_audit["prefill" if indices is None else "decode"] += 1 + _layer._state_audit["changed"] += int(not torch.equal(selected, rounded)) + _layer._state_audit["heads"] = state.shape[1] + return actual + + layer._quantized_state_call = checked_call + reports.append(dict(layer._state_audit)) + return reports + + +def _tiny_model(path, kind): + common = { + "torch_dtype": "bfloat16", + "hidden_size": 128, + "intermediate_size": 256, + "num_attention_heads": 4, + "head_dim": 32, + "vocab_size": 512, + "max_position_embeddings": 128, + } + if kind == "gdn": + config = Qwen3NextConfig( + **common, + num_hidden_layers=2, + num_key_value_heads=2, + linear_num_value_heads=4, + linear_num_key_heads=2, + linear_key_head_dim=32, + linear_value_head_dim=64, + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=64, + shared_expert_intermediate_size=64, + mlp_only_layers=[], + full_attention_interval=2, + architectures=["Qwen3NextForCausalLM"], + ) + else: + config = KimiLinearConfig( + **common, + num_hidden_layers=1, + architectures=["KimiLinearForCausalLM"], + linear_attn_config={ + "head_dim": 32, + "num_heads": 4, + "short_conv_kernel_size": 4, + "kda_layers": [1], + "full_attn_layers": [], + }, + ) + config.save_pretrained(path) + shutil.copytree(ROOT / "tests/_test_utils/torch/tokenizer", path, dirs_exist_ok=True) + + +def _generate(llm): + outputs = llm.generate( + [{"prompt_token_ids": [3] * 73}, {"prompt_token_ids": [4]}], + SamplingParams(temperature=0, max_tokens=12, ignore_eos=True, logprobs=1), + use_tqdm=False, + ) + values = [(o.outputs[0].token_ids, o.outputs[0].cumulative_logprob) for o in outputs] + assert all(len(tokens) == 12 and math.isfinite(score) for tokens, score in values) + return values + + +@pytest.mark.parametrize("kind", ["gdn", "kda"]) +@pytest.mark.parametrize("tp", [1, 2]) +@pytest.mark.timeout(360) +def test_fakequant_worker_generation_reference_reload_and_tp(tmp_path, monkeypatch, kind, tp): + if torch.cuda.device_count() < tp: + pytest.skip(f"Requires {tp} GPUs") + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1") + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + paths = [str(ROOT), str(ROOT / "examples/vllm_serve"), str(Path(__file__).parent)] + # Spawn restores the parent interpreter search path after reading PYTHONPATH. + for path in reversed(paths): + monkeypatch.syspath_prepend(path) + monkeypatch.setenv("PYTHONPATH", os.pathsep.join([*paths, os.environ.get("PYTHONPATH", "")])) + monkeypatch.setenv( + "RECIPE_PATH", str(ROOT / "examples/vllm_serve/linear_attention_state_int8.yaml") + ) + monkeypatch.delenv("MODELOPT_STATE_PATH", raising=False) + model_path, checkpoint = tmp_path / "model", tmp_path / "state.pt" + _tiny_model(model_path, kind) + kwargs = { + "model": str(model_path), + "load_format": "dummy", + "dtype": "bfloat16", + "max_model_len": 128, + "max_num_seqs": 4, + "max_num_batched_tokens": 64, + "enforce_eager": True, + "enable_prefix_caching": False, + "enable_chunked_prefill": True, + "async_scheduling": False, + "mamba_cache_dtype": "float32", + "kv_cache_memory_bytes": 128 * 1024**2, + "worker_cls": "fakequant_worker.FakeQuantWorker", + "disable_custom_all_reduce": True, + "tensor_parallel_size": tp, + "seed": 17, + } + llm = LLM(**kwargs) + try: + llm.collective_rpc(_audit_worker, kwargs={"checkpoint": str(checkpoint)}) + actual = _generate(llm) + assert _generate(llm) == actual # Fresh requests can reuse the same native state slots. + reports = llm.collective_rpc(_audit_worker) + assert len(reports) == tp + for rank in reports: + assert all( + r["prefill"] > 0 and r["decode"] > 0 and r["changed"] > 0 and r["heads"] == 4 // tp + for r in rank + ) + assert _generate(llm) == actual + llm.collective_rpc(_set_state_format_worker, kwargs={"fp8": True}) + fp8 = _generate(llm) + assert _generate(llm) == fp8 + llm.collective_rpc(_set_state_format_worker, kwargs={"fp8": False}) + assert _generate(llm) == actual + assert all(llm.collective_rpc(_calibrate_worker)) + llm.collective_rpc(_disable_worker) + disabled = _generate(llm) + finally: + llm.llm_engine.engine_core.shutdown() + del llm + gc.collect() + + monkeypatch.delenv("RECIPE_PATH") + monkeypatch.setenv("MODELOPT_STATE_PATH", str(checkpoint)) + llm = LLM(**kwargs) + try: + assert _generate(llm) == actual + finally: + llm.llm_engine.engine_core.shutdown() + del llm + gc.collect() + + monkeypatch.delenv("MODELOPT_STATE_PATH") + llm = LLM(**kwargs) + try: + assert _generate(llm) == disabled + finally: + llm.llm_engine.engine_core.shutdown() + del llm + gc.collect() + + +def test_saved_linear_attention_policy_and_quantizer_names_follow_mapper(): + module = _example_module("vllm_reload_utils") + + metadata = { + "linear_attention": {"backbone.layers.0.attn": {"backend": "fla"}}, + "quantizer_state": {"backbone.layers.0.attn.kda_state_quantizer": {"_disabled": False}}, + } + state = { + "modelopt_state_dict": [("quantize", {"config": {"quant_cfg": []}, "metadata": metadata})] + } + + def mapper(values): + return { + k.replace("backbone.", "model.").replace(".attn.", ".self_attn."): v + for k, v in values.items() + } + + state["modelopt_state_dict"].append( + ("quantize_algo", {"config": {}, "metadata": copy.deepcopy(metadata)}) + ) + result = module.convert_modelopt_state_to_vllm(state, mapper) + converted = result["modelopt_state_dict"][0][1] + assert list(converted["metadata"]["linear_attention"]) == ["model.layers.0.self_attn"] + assert list(converted["metadata"]["quantizer_state"]) == [ + "model.layers.0.self_attn.kda_state_quantizer" + ] + assert converted["config"]["linear_attention"][0]["module_name"] == "model.layers.0.self_attn" + assert "linear_attention" not in result["modelopt_state_dict"][1][1]["config"] diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py new file mode 100644 index 00000000000..cda0f0b263b --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -0,0 +1,399 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from copy import deepcopy +from functools import partial + +import pytest +import torch +import torch.nn as nn + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.config import QuantizeConfig +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.plugins import gdn +from modelopt.torch.quantization.plugins.gdn import GatedDeltaNetStateQuantMixin + +GDN_STATE_INT8_DYNAMIC = { + "num_bits": 8, + "unsigned": False, + "narrow_range": True, + "type": "dynamic", + "axis": (0, 1), +} + +GDN_STATE_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"} + + +def chunk_gated_delta_rule(q, k, v, g, beta, **kwargs): + """CPU stand-in for the optional FLA kernel.""" + return q + k + v, None + + +@pytest.fixture(autouse=True) +def mock_fla_kernel(monkeypatch): + monkeypatch.setattr(gdn, "_fla_chunk_gated_delta_rule", lambda: chunk_gated_delta_rule) + + +class TinyGatedDeltaNet(nn.Module): + """A module that, like Megatron-Core's GatedDeltaNet, calls ``self.gated_delta_rule``.""" + + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4) + self.gated_delta_rule = chunk_gated_delta_rule + + def forward(self, x): + out, _ = self.gated_delta_rule(x, x, x, x[..., 0], x[..., 0]) + return self.proj(out) + + +@QuantModuleRegistry.register({TinyGatedDeltaNet: "TinyGatedDeltaNet"}) +class _QuantTinyGatedDeltaNet(GatedDeltaNetStateQuantMixin): + def forward(self, x): + gated_delta_rule = self.gated_delta_rule + self.gated_delta_rule = lambda *a, **kw: self._state_quantized_chunk_gated_delta_rule( + gated_delta_rule, *a, **kw + ) + try: + return super().forward(x) + finally: + self.gated_delta_rule = gated_delta_rule + + +GDN_W_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1, 2), "type": "dynamic"} + + +def quant_cfg(state=True, w=False): + entries = [{"quantizer_name": "*", "enable": False}] + if state: + entries.append({"quantizer_name": "*gdn_state_quantizer", "cfg": GDN_STATE_FP8_DYNAMIC}) + if w: + entries.append({"quantizer_name": "*gdn_w_quantizer", "cfg": GDN_W_FP8_DYNAMIC}) + return {"quant_cfg": entries, "algorithm": "max"} + + +@pytest.mark.parametrize( + "attributes", + [ + {"num_bits": (4, 3), "axis": (0, 1)}, # static + {"num_bits": (4, 3), "type": "dynamic"}, # per tensor + {"num_bits": 8, "axis": (0, 1), "type": "dynamic"}, # int8 + {"num_bits": (4, 3), "type": "dynamic", "block_sizes": {-1: 16}}, # blockwise + ], +) +def test_validate_state_quantizer_rejects_unsupported(attributes): + with pytest.raises(ValueError, match="supports only"): + mtq.quantize( + TinyGatedDeltaNet(), + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*gdn_state_quantizer", "cfg": attributes}, + ], + "algorithm": None, + }, + ) + + +def test_dynamic_export_removes_linear_attention_attributes(): + model = _QuantTinyGatedDeltaNet.convert(TinyGatedDeltaNet()) + model.export() + assert type(model) is TinyGatedDeltaNet + for name in ( + "gdn_state_quantizer", + "gdn_w_quantizer", + "linear_attention_config", + "_linear_attention_prefill_lengths", + "linear_attn_sites", + ): + assert not hasattr(model, name) + + +def test_disabled_state_quantizer_calls_original_kernel(): + model = TinyGatedDeltaNet() + x = torch.randn(2, 8, 3, 4) + expected = model(x) + + disable_all = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": "max"} + mtq.quantize(model, disable_all, lambda m: m(x)) + + assert isinstance(model, _QuantTinyGatedDeltaNet) + assert not model.gdn_state_quantizer.is_enabled and not model.gdn_w_quantizer.is_enabled + assert torch.equal(model(x), expected) + assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" + + +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +@pytest.mark.parametrize("partial_kernel", [False, True]) +def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, state_format, partial_kernel): + calls = [] + + def fake_state_qdq_kernel(*args, **kwargs): + calls.append(kwargs) + return chunk_gated_delta_rule(*args) + + monkeypatch.setattr(gdn, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel) + model = TinyGatedDeltaNet() + if partial_kernel: + model.gated_delta_rule = partial(chunk_gated_delta_rule, output_final_state=True) + x = torch.randn(2, 8, 3, 4) + cfg = quant_cfg() + if state_format == "int8": + cfg["quant_cfg"][1]["cfg"] = GDN_STATE_INT8_DYNAMIC + mtq.quantize(model, cfg, lambda m: m(x)) + + model(x) + assert calls and calls[-1] == { + "chunk_size": 64, + "state_qdq": 2 if state_format == "int8" else 1, + "state_qdq_block_v": 64, + "w_quantizer": None, + **({"output_final_state": True} if partial_kernel else {}), + } + + # An unrelated callable must not pass validation merely by copying the FLA name. + model.gated_delta_rule = lambda *a, **kw: chunk_gated_delta_rule(*a, **kw) + model.gated_delta_rule.__name__ = "chunk_gated_delta_rule" + with pytest.raises(NotImplementedError, match="supports only FLA"): + model(x) + + +@pytest.mark.parametrize("state", [False, True]) +def test_w_quantizer_is_passed_to_the_kernel(monkeypatch, state): + """``*gdn_w_quantizer`` in the config hands the module's TensorQuantizer to the kernel, with + or without the state quantizer.""" + calls = [] + + def fake_state_qdq_kernel(*args, **kwargs): + calls.append(kwargs) + return chunk_gated_delta_rule(*args) + + monkeypatch.setattr(gdn, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel) + model = TinyGatedDeltaNet() + x = torch.randn(2, 8, 3, 4) + mtq.quantize(model, quant_cfg(state=state, w=True), lambda m: m(x)) + assert model.gdn_w_quantizer.is_enabled and model.gdn_state_quantizer.is_enabled == state + + model(x) + assert calls[-1]["state_qdq"] == int(state) + assert calls[-1]["w_quantizer"] is model.gdn_w_quantizer + + # The w quantizer really quantizes: 256 random values per token collapse onto the E4M3 grid, + # which has at most 127 distinct magnitudes per (row-specific) scale. + w = torch.randn(1, 1, 1, 256) + quantized = model.gdn_w_quantizer(w) + assert not torch.equal(quantized, w) + assert torch.unique(quantized.abs()).numel() <= 127 < torch.unique(w.abs()).numel() + + +@pytest.mark.parametrize("site", ["state", "w"]) +@pytest.mark.parametrize( + "overrides", + [{"pass_through_bwd": False}, {"type": "static"}, {"fake_quant": False}, {"rotate": True}], +) +def test_unsupported_quantizer_fails_during_conversion(site, overrides): + cfg = quant_cfg(state=site == "state", w=site == "w") + cfg = deepcopy(cfg) + cfg["quant_cfg"][-1]["cfg"].update(overrides) + with pytest.raises(ValueError, match="supports only"): + mtq.quantize(TinyGatedDeltaNet(), cfg) + + +@pytest.mark.parametrize("axis", [None, (0, 1), (0, 1, 2)]) +def test_w_grouping_is_preserved_during_conversion(axis): + cfg = deepcopy(quant_cfg(state=False, w=True)) + cfg["algorithm"] = None + cfg["quant_cfg"][-1]["cfg"]["axis"] = axis + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + assert model.gdn_w_quantizer.axis == axis + + +@pytest.mark.parametrize(("state", "w"), [(True, False), (False, True), (True, True)]) +def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): + model = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) + cfg = quant_cfg(state=state, w=w) + cfg["algorithm"] = None + cfg["linear_attention"] = [ + {"module_name": "*", "cfg": {"state": {"block_v": 128}}}, + {"module_name": "0", "cfg": {"state": {"block_v": 32}}}, + ] + mtq.quantize(model, cfg) + assert model[0].gdn_state_qdq_block_v == 32 + model[0].linear_attention_config.state.block_v = 16 + if w: + # Save the actual quantizer settings, including edits after conversion. + model[0].gdn_w_quantizer.axis = None + path = tmp_path / "gdn.pth" + mto.save(model, path) + restored = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) + mto.restore(restored, path) + assert restored[0].linear_attention_config == model[0].linear_attention_config + assert restored[0].gdn_state_qdq_block_v == 16 + for name, enabled in (("gdn_state_quantizer", state), ("gdn_w_quantizer", w)): + original = getattr(model[0], name) + quantizer = getattr(restored[0], name) + assert quantizer.is_enabled == enabled + assert quantizer.axis == original.axis + assert quantizer.num_bits == original.num_bits + assert quantizer._dynamic == original._dynamic + assert not hasattr(restored[1], name) + + +def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): + cfg = quant_cfg() + cfg["algorithm"] = None + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + assert model.gdn_state_quantizer.is_enabled + assert not model.gdn_w_quantizer.is_enabled + + cfg = quant_cfg(state=False, w=True) + cfg["algorithm"] = None + cfg["linear_attention"] = [{"module_name": "", "cfg": {"state": {"block_v": 32}}}] + mtq.quantize(model, cfg) + assert model.gdn_state_qdq_block_v == 32 + cfg["linear_attention"].append({"module_name": "", "cfg": {}}) + mtq.quantize(model, cfg) + assert model.gdn_state_qdq_block_v == 64 + assert not model.gdn_state_quantizer.is_enabled + assert model.gdn_w_quantizer.is_enabled + + cfg = deepcopy(cfg) + cfg["quant_cfg"][-1]["cfg"]["pass_through_bwd"] = False + with pytest.raises(ValueError, match="supports only"): + mtq.quantize(model, cfg) + + +def test_restore_legacy_gdn_without_new_quantizer_handles(): + config = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None} + model = mtq.quantize(TinyGatedDeltaNet(), config) + state = mto.modelopt_state(model) + for _, mode_state in state["modelopt_state_dict"]: + mode_state["config"].pop("linear_attention", None) + metadata = mode_state["metadata"] + metadata.pop("linear_attention", None) + for name in ("gdn_state_quantizer", "gdn_w_quantizer"): + metadata["quantizer_state"].pop(name, None) + restored = TinyGatedDeltaNet() + mto.restore_from_modelopt_state(restored, state) + restored.load_state_dict(model.state_dict()) + x = torch.randn(2, 8, 3, 4) + torch.testing.assert_close(restored(x), model(x)) + assert not restored.gdn_state_quantizer.is_enabled + assert not restored.gdn_w_quantizer.is_enabled + + +def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): + model = mtq.quantize( + TinyGatedDeltaNet(), mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(2, 8, 3, 4)) + ) + assert not model.gdn_state_quantizer.is_enabled + assert not model.gdn_w_quantizer.is_enabled + + +@pytest.mark.parametrize("arithmetic_only", [False, True]) +def test_prefill_policy_executes_and_survives_restore(tmp_path, arithmetic_only): + cfg = { + "quant_cfg": [{"quantizer_name": "*", "enable": False}], + "algorithm": None, + "linear_attention": [{"module_name": "", "cfg": {"backend": "matmul"}}], + } + if arithmetic_only: + cfg["linear_attention"][0]["cfg"]["elementwise"] = {"output_add": "float16"} + else: + cfg["quant_cfg"].append( + { + "quantizer_name": "linear_attn_sites.output_score.lhs_quantizer", + "cfg": GDN_W_FP8_DYNAMIC, + } + ) + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + x = torch.randn(2, 8, 3, 4) * 0.1 + expected = model(x) + expected.square().sum().backward() + assert torch.isfinite(model.proj.weight.grad).all() + path = tmp_path / "prefill.pth" + mto.save(model, path) + restored = mto.restore(TinyGatedDeltaNet(), path) + torch.testing.assert_close(restored(x), expected) + assert restored.linear_attention_config == model.linear_attention_config + assert restored.linear_attention_is_enabled + + +def test_restore_m1_checkpoint_without_prefill_handles(): + model = mtq.quantize(TinyGatedDeltaNet(), quant_cfg(state=False, w=True)) + state = mto.modelopt_state(model) + for _, mode_state in state["modelopt_state_dict"]: + quantizers = mode_state["metadata"]["quantizer_state"] + for key in list(quantizers): + if key.startswith("linear_attn_sites."): + del quantizers[key] + restored = mto.restore_from_modelopt_state(TinyGatedDeltaNet(), state) + assert restored.gdn_w_quantizer.is_enabled + assert not restored.linear_attn_sites.is_enabled + assert restored.linear_attention_config.backend == "fla" + + +def test_prefill_sites_require_explicit_backend(): + with pytest.raises(ValueError, match="require backend='matmul'"): + mtq.quantize( + TinyGatedDeltaNet(), + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "linear_attn_sites.output_score.lhs_quantizer", + "cfg": GDN_W_FP8_DYNAMIC, + }, + ], + "algorithm": None, + }, + ) + + +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_kda_prefill.py b/tests/unit/torch/quantization/test_kda_prefill.py new file mode 100644 index 00000000000..5288f42dc68 --- /dev/null +++ b/tests/unit/torch/quantization/test_kda_prefill.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +from _test_utils.torch.linear_attention import FP8, inputs, reference_matmul, values_and_gradients +from _test_utils.torch.quantization.linear_attention_reference import ( + chunk_kda_reference, + recurrent_delta_rule_reference, +) + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + LinearAttentionMatmulSites, + matmul_kda, +) +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.mark.parametrize(("packed", "state_v_first"), [(False, False), (True, True)]) +@pytest.mark.parametrize("strong_decay", [False, True]) +def test_kda_exact_outputs_state_and_gradients(packed, state_v_first, strong_decay): + args, state = inputs(packed=packed, state_v_first=state_v_first, dtype=torch.float64) + args[2] = args[2][..., :11].detach().requires_grad_() + state = (state[..., :11, :] if state_v_first else state[..., :11]).detach().requires_grad_() + gates = -torch.rand(*args[3].shape, args[0].shape[-1], dtype=torch.float64) + if strong_decay: + gates = gates * torch.logspace(-8, 3, gates.shape[-1], dtype=torch.float64) + else: + gates = gates * 0.1 + args[3] = gates.requires_grad_() + kwargs = { + "initial_state": state, + "state_v_first": state_v_first, + "cu_seqlens": torch.tensor([0, 5, 73]) if packed else None, + } + actual = matmul_kda( + *args, + sites=LinearAttentionMatmulSites(), + policy=LinearAttentionConfig(backend="matmul"), + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(enable=False)), + output_final_state=True, + **kwargs, + ) + recurrent = recurrent_delta_rule_reference(*args, **kwargs) + chunk = chunk_kda_reference(*args, **kwargs) + for result in (actual, chunk): + for a, e in zip( + values_and_gradients(result, args, state), values_and_gradients(recurrent, args, state) + ): + assert torch.isfinite(a).all() + torch.testing.assert_close(a, e, rtol=2e-9, atol=1e-10) + + +HANDLES = [ + name for name, _ in LinearAttentionMatmulSites().named_modules() if name.endswith("quantizer") +] + + +@pytest.mark.parametrize("handle", HANDLES) +def test_kda_fp8_site_matches_chunk_oracle(handle): + args, state = inputs() + args[3] = (-torch.rand(*args[3].shape, args[0].shape[-1]) * 0.1).requires_grad_() + sites = LinearAttentionMatmulSites() + sites.get_submodule(handle).set_from_attribute_config(FP8) + sites.get_submodule(handle).enable() + policy = LinearAttentionConfig(backend="matmul") + actual = matmul_kda( + *args, + sites=sites, + policy=policy, + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(enable=False)), + initial_state=state, + output_final_state=True, + ) + expected = chunk_kda_reference( + *args, initial_state=state, matmul=reference_matmul([handle], policy) + ) + for a, e in zip( + values_and_gradients(actual, args, state), values_and_gradients(expected, args, state) + ): + torch.testing.assert_close(a, e, rtol=5e-4, atol=3e-6) 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..393b555faca --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -0,0 +1,301 @@ +# 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, + LinearAttentionMatmulSites, + matmul_gdn, + matmul_kda, +) +from modelopt.torch.quantization.linear_attention.config import LinearAttentionDecodeConfig +from modelopt.torch.quantization.linear_attention.decode import _encode, recurrent_decode_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, + sites=LinearAttentionMatmulSites(), + policy=LinearAttentionConfig(backend="matmul", decode={}), + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(enable=False)), + prefill_lengths=prefixes, + cu_seqlens=torch.tensor(boundaries), + initial_state=initial, + output_final_state=True, + ) + outputs, states = [], [] + for n, (lo, hi) in enumerate(itertools.pairwise(boundaries)): + if lo == hi: + outputs.append(args[2][0, lo:hi]) + states.append(initial[n]) + continue + o, s = recurrent_delta_rule_reference( + *(x[:, lo:hi] for x in args), initial_state=initial[n : n + 1] + ) + outputs.append(o[0]) + states.append(s[0]) + expected = torch.cat(outputs).unsqueeze(0), torch.stack(states) + for a, e in zip(actual, expected): + torch.testing.assert_close(a, e, rtol=2e-11, atol=1e-12) + gradients = [ + torch.autograd.grad( + sum(x.square().sum() for x in result), (*args, initial), retain_graph=True + ) + for result in (actual, expected) + ] + for a, e in zip(*gradients): + torch.testing.assert_close(a, e, rtol=2e-10, atol=1e-11) + + +@pytest.mark.parametrize("mode", ["token", "replay"]) +def test_carry_rejects_changed_policy_and_shape(mode): + args, state = _inputs(length=4) + cfg = LinearAttentionDecodeConfig(mode=mode, replay={"window": 8} if mode == "replay" else None) + _, carry = recurrent_decode_reference(*args, config=cfg, initial_state=state, state_qdq=True) + for changed in [ + cfg.model_copy(update={"readout": "working"}), + cfg.model_copy(update={"decay_log_step": 0.01}), + ]: + with pytest.raises(ValueError, match="policy or state shape"): + recurrent_decode_reference(*args, config=changed, carry=carry, state_qdq=True) + with pytest.raises(ValueError, match="policy or state shape"): + recurrent_decode_reference(*args, config=cfg, carry=carry, state_qdq=True, block_v=16) + with pytest.raises(ValueError, match="either carry or initial"): + recurrent_decode_reference( + *args, config=cfg, carry=carry, initial_state=state, state_qdq=True + ) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("all_empty", [False, True]) +def test_grouped_heads_transposed_state_and_empty_batch(kda, all_empty): + args, state = _inputs(kda, length=0 if all_empty else 9) + q, k, v, g, beta = [x.unsqueeze(0).repeat(2, *([1] * x.ndim)) for x in args] + q, k = q[:, :, :1], k[:, :, :1] + initial = state.unsqueeze(0).repeat(2, 1, 1, 1).transpose(-1, -2).detach().requires_grad_() + function = matmul_kda if kda else matmul_gdn + policy = LinearAttentionConfig( + backend="matmul", decode={"mode": "replay", "replay": {"window": 3}} + ) + shared = { + "sites": LinearAttentionMatmulSites(), + "policy": policy, + "w_quantizer": TensorQuantizer(QuantizerAttributeConfig(enable=False)), + "state_qdq": True, + "output_final_state": True, + } + actual = function( + q, + k, + v, + g, + beta, + initial_state=initial, + state_v_first=True, + prefill_lengths=[0, 0] if all_empty else [2, 5], + **shared, + ) + expected = function( + q.repeat_interleave(2, dim=2), + k.repeat_interleave(2, dim=2), + v, + g, + beta, + initial_state=initial.transpose(-1, -2), + prefill_lengths=[0, 0] if all_empty else [2, 5], + **shared, + ) + torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0) + torch.testing.assert_close(actual[1].transpose(-1, -2), expected[1], rtol=0, atol=0) + gradients = [ + torch.autograd.grad(result[0].sum() + result[1].sum(), initial, retain_graph=True)[0] + for result in (actual, expected) + ] + torch.testing.assert_close(*gradients, rtol=0, atol=0) + if all_empty: + torch.testing.assert_close(actual[1], initial, rtol=0, atol=0) diff --git a/tests/unit/torch/quantization/test_linear_attention_hadamard.py b/tests/unit/torch/quantization/test_linear_attention_hadamard.py new file mode 100644 index 00000000000..77208748aa0 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_hadamard.py @@ -0,0 +1,231 @@ +# 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, QuantizerAttributeConfig +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 +from modelopt.torch.quantization.linear_attention.matmul import LinearAttentionMatmulSites +from modelopt.torch.quantization.nn import TensorQuantizer + + +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), + sites=LinearAttentionMatmulSites(), + policy=policy, + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(enable=False)), + prefill_lengths=[prefix], + state_qdq=True, + state_format="int8", + initial_state=initial.transpose(-1, -2).unsqueeze(0), + state_v_first=True, + output_final_state=True, + ) + prefix_out, state = recurrent_decode_reference( + *(x[:prefix] for x in args), initial_state=initial, config=LinearAttentionDecodeConfig() + ) + suffix, carry = recurrent_decode_reference( + *(x[prefix:] for x in args), + initial_state=state.reconstruct(), + config=policy.decode, + state_qdq=True, + state_format="int8", + ) + expected = (torch.cat((prefix_out, suffix)), carry.reconstruct()) + _check_with_grads((actual[0][0], actual[1][0].transpose(-1, -2)), expected, args, initial) + + +def test_hadamard_rejects_unsupported_codec_combinations(): + args, initial = _inputs(False) + cfg = LinearAttentionDecodeConfig(state_codec="int8_hadamard32") + for kwargs, message in [ + ({"state_qdq": True}, "requires INT8"), + ({"state_format": "int8", "block_v": 16}, "block_v >= 32"), + ]: + with pytest.raises(ValueError, match=message): + recurrent_decode_reference(*args, config=cfg, initial_state=initial, **kwargs) + with pytest.raises(ValueError, match="Dv divisible"): + recurrent_decode_reference(args[0], args[1], args[2][..., :33], *args[3:], config=cfg) + with pytest.raises(ValueError, match="decode handoff"): + LinearAttentionDecodeConfig(state_codec="int8_hadamard32", prefill_state_qdq=True) + with pytest.raises(ValueError, match="block_v >= 32"): + LinearAttentionConfig(backend="matmul", state={"block_v": 16}, decode=cfg) diff --git a/tests/unit/torch/quantization/test_linear_attention_int8.py b/tests/unit/torch/quantization/test_linear_attention_int8.py new file mode 100644 index 00000000000..b3de8c94d47 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_int8.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import pytest +import torch +import torch.nn.functional as F +from _test_utils.torch.quantization.linear_attention_reference import state_qdq_reference + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionDecodeConfig, + recurrent_decode_reference, +) +from modelopt.torch.quantization.linear_attention.decode import _encode +from modelopt.torch.quantization.nn import TensorQuantizer + +INT8 = {"num_bits": 8, "unsigned": False, "narrow_range": True, "type": "dynamic", "axis": (0, 1)} + + +@pytest.mark.parametrize("block_v", [16, 32, 64, 128]) +def test_int8_state_codec_matches_modelopt_and_identity_gradient(block_v): + # Exact scale two exercises both signs, endpoints and nearest-even half ties. + state = torch.zeros(2, 3, 5, block_v + 3) + state[..., 0, :9] = torch.tensor([-127, -126.5, -1.5, -0.5, 0, 0.5, 1.5, 126.5, 127]) * 2 + state[0, :, 0, block_v:] = torch.tensor([-254, 1, 254]) + state.requires_grad_() + quantizer = TensorQuantizer(QuantizerAttributeConfig(**INT8)) + expected = torch.cat( + [quantizer(tile.flatten(-2)).reshape_as(tile) for tile in state.split(block_v, -1)], -1 + ) + encoded = _encode(state, True, block_v, state=True, state_format="int8") + reference = state_qdq_reference(state, block_v, "int8") + torch.testing.assert_close(encoded.values, expected, rtol=0, atol=0) + torch.testing.assert_close(reference, expected, rtol=0, atol=0) + assert encoded.format == "int8" + assert torch.equal(encoded.scales[1, :, -1], torch.ones(3)) + assert not encoded.scales.requires_grad + probe = torch.randn_like(state) + for result in (encoded.values, reference): + (gradient,) = torch.autograd.grad((result * probe).sum(), state, retain_graph=True) + torch.testing.assert_close(gradient, probe, rtol=0, atol=0) + + +def _inputs(kda): + torch.manual_seed(913) + q, k = [F.normalize(torch.randn(23, 2, 16), dim=-1).requires_grad_() for _ in range(2)] + v = torch.randn(23, 2, 19, requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1]) * 0.03).requires_grad_() + beta = (torch.rand(23, 2) * 0.4).requires_grad_() + state = (torch.randn(2, 16, 19) * 0.1).requires_grad_() + return (q, k, v, g, beta), state + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("mode", ["token", "replay"]) +def test_int8_decode_continuation_and_format_rejection(kda, mode): + args, initial = _inputs(kda) + cfg = LinearAttentionDecodeConfig(mode=mode, replay={"window": 5} if mode == "replay" else None) + kwargs = {"config": cfg, "state_qdq": True, "state_format": "int8", "block_v": 16} + expected, final = recurrent_decode_reference(*args, initial_state=initial, **kwargs) + pieces, carry, start = [], None, 0 + for end in (0, 3, 10, 10, 23): + output, carry = recurrent_decode_reference( + *(x[start:end] for x in args), + carry=carry, + initial_state=initial if carry is None else None, + **kwargs, + ) + pieces.append(output) + start = end + actual = torch.cat(pieces) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(carry.reconstruct(), final.reconstruct(), rtol=0, atol=0) + probes = (torch.randn_like(expected), torch.randn_like(initial)) + grads = [ + torch.autograd.grad((o * probes[0]).sum() + (s * probes[1]).sum(), (*args, initial)) + for o, s in ((actual, carry.reconstruct()), (expected, final.reconstruct())) + ] + for actual_grad, expected_grad in zip(*grads): + torch.testing.assert_close(actual_grad, expected_grad, rtol=0, atol=0) + assert carry.anchor.format == "int8" + assert all(e.key.format == e.update.format == "fp8_e4m3" for e in carry.entries) + with pytest.raises(ValueError, match="Carry policy"): + recurrent_decode_reference( + *(x[:0] for x in args), carry=carry, **{**kwargs, "state_format": "fp8_e4m3"} + ) diff --git a/tests/unit/torch/quantization/test_linear_attention_prefill.py b/tests/unit/torch/quantization/test_linear_attention_prefill.py new file mode 100644 index 00000000000..9261f7352a6 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_prefill.py @@ -0,0 +1,126 @@ +# 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 typing import get_args + +import pytest +import torch +from _test_utils.torch.linear_attention import check_prefill_case, inputs, values_and_gradients +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 +from modelopt.torch.quantization.linear_attention.config import _ElementwiseSite, _PrefillSite +from modelopt.torch.quantization.linear_attention.matmul import LinearAttentionMatmulSites +from modelopt.torch.quantization.nn import TensorQuantizer + +SITES = get_args(_PrefillSite) +HANDLES = [ + name + for name, module in LinearAttentionMatmulSites().named_modules() + if name.endswith("quantizer") +] + + +@pytest.mark.parametrize("handle", HANDLES) +def test_prefill_operand_outputs_state_and_all_input_gradients(handle): + check_prefill_case(LinearAttentionConfig(backend="matmul"), [handle]) + + +@pytest.mark.parametrize("site", SITES) +def test_blocked_accumulator_surrogate(site): + policy = LinearAttentionConfig( + backend="matmul", matmul={site: {"accumulator_dtype": "float16", "reduction_block": 4}} + ) + check_prefill_case(policy) + + +@pytest.mark.parametrize("site", get_args(_ElementwiseSite)) +def test_elementwise_rounding_surrogate(site): + check_prefill_case(LinearAttentionConfig(backend="matmul", elementwise={site: "float16"})) + + +@pytest.mark.parametrize(("packed", "state_v_first"), [(False, False), (True, True)]) +def test_composed_prefill_with_state_w_and_layouts(packed, state_v_first): + check_prefill_case( + LinearAttentionConfig( + backend="matmul", state={"block_v": 16}, elementwise={"value_residual": "bfloat16"} + ), + HANDLES, + packed=packed, + state_v_first=state_v_first, + state_qdq=True, + w_qdq=True, + ) + + +@pytest.mark.parametrize( + "config", + [ + {"matmul": {"output_score": {"accumulator_dtype": "float16", "reduction_block": 16}}}, + {"backend": "matmul", "matmul": {"unknown": {}}}, + {"backend": "matmul", "matmul": {"output_score": {"accumulator_dtype": "float16"}}}, + {"backend": "matmul", "matmul": {"output_score": {"reduction_block": 8}}}, + {"backend": "matmul", "elementwise": {"output_add": "fp8"}}, + ], +) +def test_reject_unimplemented_or_incomplete_arithmetic(config): + with pytest.raises(ValueError): + LinearAttentionConfig(**config) + + +def test_host_packing_metadata_with_unequal_key_value_dimensions(): + args, state = inputs(packed=True, dtype=torch.float64) + args[2] = args[2][..., :11].detach().requires_grad_() + state = state[..., :11].detach().requires_grad_() + boundaries = torch.tensor([0, 5, 73]) + actual = matmul_gdn( + *args, + initial_state=state, + cu_seqlens_cpu=boundaries, + sites=LinearAttentionMatmulSites(), + policy=LinearAttentionConfig(backend="matmul"), + w_quantizer=TensorQuantizer(QuantizerAttributeConfig(enable=False)), + output_final_state=True, + ) + expected = recurrent_delta_rule_reference(*args, initial_state=state, cu_seqlens=boundaries) + for a, e in zip( + values_and_gradients(actual, args, state), values_and_gradients(expected, args, state) + ): + torch.testing.assert_close(a, e, rtol=1e-9, atol=1e-10) + + +def test_outer_autocast_preserves_working_precision_and_gradients(): + args, state = inputs() + sites = LinearAttentionMatmulSites() + for name in HANDLES: + sites.get_submodule(name).set_from_attribute_config( + {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)} + ) + sites.get_submodule(name).enable() + kwargs = { + "sites": sites, + "policy": LinearAttentionConfig(backend="matmul"), + "w_quantizer": TensorQuantizer(QuantizerAttributeConfig(enable=False)), + "initial_state": state, + "output_final_state": True, + } + expected = matmul_gdn(*args, **kwargs) + with torch.autocast("cpu", dtype=torch.bfloat16): + actual = matmul_gdn(*args, **kwargs) + for a, e in zip( + values_and_gradients(actual, args, state), values_and_gradients(expected, args, state) + ): + torch.testing.assert_close(a, e, rtol=0, atol=0) diff --git a/tests/unit/torch/quantization/test_linear_attention_reference.py b/tests/unit/torch/quantization/test_linear_attention_reference.py new file mode 100644 index 00000000000..d71818d0163 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_reference.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +import torch.nn.functional as F +from _test_utils.torch.quantization.linear_attention_reference import ( + chunk_gdn_reference, + recurrent_delta_rule_reference, + state_fp8_qdq_reference, +) + + +def inputs(length=7, kda=False): + torch.manual_seed(81) + q, k = [ + F.normalize(torch.randn(1, length, 1, 3, dtype=torch.float64), dim=-1) for _ in range(2) + ] + v = torch.randn(1, length, 2, 4, dtype=torch.float64) + g = -torch.rand(1, length, 2, *([3] if kda else []), dtype=torch.float64) + beta = torch.rand(1, length, 2, dtype=torch.float64) + return [x.requires_grad_() for x in (q, k, v, g, beta)] + + +@pytest.mark.parametrize("packed", [False, True]) +@pytest.mark.parametrize("state_v_first", [False, True]) +def test_chunk_recurrence_output_state_and_gradients(packed, state_v_first): + args = inputs() + state = torch.randn(2 if packed else 1, 2, 3, 4, dtype=torch.float64) + if state_v_first: + state = state.transpose(-1, -2).contiguous() + state.requires_grad_() + kwargs = { + "initial_state": state, + "state_v_first": state_v_first, + "cu_seqlens": torch.tensor([0, 2, 7]) if packed else None, + } + actual = chunk_gdn_reference(*args, chunk_size=3, **kwargs) + expected = recurrent_delta_rule_reference(*args, **kwargs) + for a, e in zip(actual, expected): + torch.testing.assert_close(a, e, rtol=1e-10, atol=1e-10) + probe = [torch.randn_like(x) for x in actual] + grads = [ + torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probe)), (*args, state)) + for result in (actual, expected) + ] + for a, e in zip(*grads): + torch.testing.assert_close(a, e, rtol=1e-9, atol=1e-9) + + +@pytest.mark.parametrize("kda", [False, True]) +def test_recurrence_smooth_gradcheck(kda): + args = inputs(length=2, kda=kda) + state = torch.randn(1, 2, 3, 4, dtype=torch.float64, requires_grad=True) + assert torch.autograd.gradcheck( + lambda *x: recurrent_delta_rule_reference(*x[:-1], initial_state=x[-1]), + (*args, state), + fast_mode=True, + ) + + +def test_kda_one_step_and_scalar_decay_equivalence(): + q, k, v, g, beta = inputs(length=1, kda=True) + state = torch.randn(1, 2, 3, 4, dtype=torch.float64) + out, final = recurrent_delta_rule_reference(q, k, v, g, beta, initial_state=state) + keys = k[:, 0].expand(-1, 2, -1) + decayed = state * g[:, 0].exp().unsqueeze(-1) + residual = v[:, 0] - torch.einsum("bhk,bhkv->bhv", keys, decayed) + expected = decayed + torch.einsum("bhk,bhv->bhkv", keys, beta[:, 0, :, None] * residual) + torch.testing.assert_close(final, expected) + torch.testing.assert_close( + out[:, 0], torch.einsum("bhk,bhkv->bhv", q[:, 0].expand(-1, 2, -1), expected) / 3**0.5 + ) + scalar = g[..., 0] + gdn = recurrent_delta_rule_reference(q, k, v, scalar, beta, initial_state=state) + kda = recurrent_delta_rule_reference( + q, k, v, scalar.unsqueeze(-1).expand_as(g), beta, initial_state=state + ) + for a, e in zip(kda, gdn): + torch.testing.assert_close(a, e) + + +def test_state_qdq_granularity_zero_tail_and_identity_ste(): + torch.manual_seed(14) + state = torch.randn(2, 2, 3, 37, requires_grad=True) + with torch.no_grad(): + state[0].zero_() + state[1, ..., :16] *= 10 + quantized = state_fp8_qdq_reference(state, block_v=16) + assert torch.isfinite(quantized).all() + assert torch.equal(quantized[0], state[0]) + assert not torch.equal(quantized, state) + assert not torch.equal(quantized, state_fp8_qdq_reference(state, block_v=64)) + probe = torch.randn_like(state) + (grad,) = torch.autograd.grad((quantized * probe).sum(), state) + torch.testing.assert_close(grad, probe, rtol=0, atol=0) + + +def test_packed_sequences_reset_the_state(): + args = inputs(length=7, kda=True) + initial = torch.randn(2, 2, 3, 4, dtype=torch.float64) + packed = recurrent_delta_rule_reference( + *args, initial_state=initial, cu_seqlens=torch.tensor([0, 2, 7]) + ) + separate = [ + recurrent_delta_rule_reference( + *(x[:, lo:hi] for x in args), initial_state=initial[n : n + 1] + ) + for n, (lo, hi) in enumerate(((0, 2), (2, 7))) + ] + torch.testing.assert_close(packed[0], torch.cat([x[0] for x in separate], dim=1)) + torch.testing.assert_close(packed[1], torch.cat([x[1] for x in separate])) diff --git a/tests/unit/torch/quantization/test_linear_attention_study.py b/tests/unit/torch/quantization/test_linear_attention_study.py new file mode 100644 index 00000000000..49028e2b475 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_study.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import runpy +from pathlib import Path + +import pytest + +_STUDY = runpy.run_path( + str( + Path(__file__).resolve().parents[4] / "examples/llm_qat/linear_attention/compare_quality.py" + ) +) + + +def _receipt(): + return { + "model": "fixture", + "model_revision": "pinned", + "dataset": "fixture", + "dataset_config": "raw", + "dataset_revision": "pinned", + "eval_split": "test", + "train_file_sha256": "train", + "eval_file_sha256": "test", + "train_tokens_sha256": "train_tokens", + "eval_tokens_sha256": "test_tokens", + "sequence_length": 128, + "prefill_tokens": None, + "loss_scope": "full", + "seed": 2026, + "train_order": [1, 0], + "training": "attention_only", + "trainable_parameters": 100, + "kda_layers": 2, + "learning_rate": 1e-5, + "weight_decay": 0, + "gradient_clip": 1, + "train_predicted_tokens": 256, + "source_sha256": {"solve.py": "pinned"}, + "packages": {"fla-core": "0.5.1"}, + "torch": "pinned", + "gpu": "same_device", + "before": {"block_nll": [1.0, 2.0, 3.0, 4.0], "predicted_tokens": 512}, + "after": {"block_nll": [0.9, 1.9, 2.9, 3.9], "predicted_tokens": 512}, + } + + +def test_paired_quality_bound_and_failure(): + control = _receipt() + candidate = copy.deepcopy(control) + candidate["before"]["block_nll"] = [x + 0.01 for x in control["before"]["block_nll"]] + candidate["after"]["block_nll"] = [x + 0.03 for x in control["after"]["block_nll"]] + result = _STUDY["compare"](control, candidate, samples=100)["comparison"] + assert result["before"]["mean_nll_delta"] == pytest.approx(0.01) + assert result["before"]["upper_bound_within_margin"] + assert not result["after"]["upper_bound_within_margin"] + + +@pytest.mark.parametrize( + "field", + [ + "eval_tokens_sha256", + "source_sha256", + "train_order", + "learning_rate", + "prefill_tokens", + "loss_scope", + ], +) +def test_quality_refuses_unmatched_studies(field): + control, candidate = _receipt(), _receipt() + candidate[field] = "different" + with pytest.raises(ValueError, match=field): + _STUDY["compare"](control, candidate) + + +def test_quality_refuses_nonfinite_values(): + control, candidate = _receipt(), _receipt() + candidate["after"]["block_nll"][0] = float("nan") + with pytest.raises(ValueError, match="finite"): + _STUDY["compare"](control, candidate, samples=10)