From 07d39f36a06463c5550d92879d2a563be127d8bd Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 21 Sep 2026 19:07:34 -0700 Subject: [PATCH 1/8] Add GDN state and WY activation QAT emulation Signed-off-by: Kai Xu --- .pre-commit-config.yaml | 2 + CHANGELOG.rst | 1 + LICENSE | 1 + docs/linear_attention_qat.md | 130 +++ docs/source/_templates/autosummary/module.rst | 2 +- .../quantization/linear_attention/__init__.py | 24 + .../linear_attention/fla_chunk_delta_h.py | 979 ++++++++++++++++++ .../fla_chunk_gated_delta_rule.py | 705 +++++++++++++ modelopt/torch/quantization/config.py | 7 + modelopt/torch/quantization/conversion.py | 57 +- .../quantization/linear_attention/__init__.py | 19 + .../quantization/linear_attention/config.py | 60 ++ .../linear_attention/reference.py | 203 ++++ .../linear_attention/validation.py | 45 + modelopt/torch/quantization/model_quant.py | 2 + .../torch/quantization/plugins/__init__.py | 1 + .../quantization/plugins/gated_delta_net.py | 115 ++ .../torch/quantization/plugins/megatron.py | 43 +- modelopt_recipes/configs/ptq/units/README.md | 4 + .../units/default_disabled_quantizers.yaml | 5 + .../ptq/units/gdn_state_fp8_dynamic.yaml | 30 + .../configs/ptq/units/gdn_w_fp8_dynamic.yaml | 28 + noxfile.py | 2 + pyproject.toml | 10 + .../test_fla_chunk_gated_delta_rule.py | 212 ++++ .../plugins/test_megatron_gated_delta_net.py | 101 ++ .../plugins/test_gated_delta_net.py | 276 +++++ .../test_linear_attention_reference.py | 125 +++ 28 files changed, 3186 insertions(+), 3 deletions(-) create mode 100644 docs/linear_attention_qat.md create mode 100644 modelopt/torch/kernels/quantization/linear_attention/__init__.py create mode 100644 modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py create mode 100644 modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py create mode 100644 modelopt/torch/quantization/linear_attention/__init__.py create mode 100644 modelopt/torch/quantization/linear_attention/config.py create mode 100644 modelopt/torch/quantization/linear_attention/reference.py create mode 100644 modelopt/torch/quantization/linear_attention/validation.py create mode 100644 modelopt/torch/quantization/plugins/gated_delta_net.py create mode 100644 modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml create mode 100644 modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml create mode 100644 tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py create mode 100644 tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py create mode 100644 tests/unit/torch/quantization/plugins/test_gated_delta_net.py create mode 100644 tests/unit/torch/quantization/test_linear_attention_reference.py 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..bc6a6ae6e7a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -17,6 +17,7 @@ Changelog - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration. - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md `_ for details. - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``. +- Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training, with saved execution policies and differentiable GDN/KDA references. 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/linear_attention_qat.md b/docs/linear_attention_qat.md new file mode 100644 index 00000000000..0db84b58f22 --- /dev/null +++ b/docs/linear_attention_qat.md @@ -0,0 +1,130 @@ +# Linear-attention numerical emulation for training + +This experimental integration adds dynamic FP8 E4M3 fake quantization to Megatron +GatedDeltaNet's chunked training path. It rounds the recurrent state at chunk boundaries +and the materialized WY activation `W` before the state-read matmul. Both sites use an +identity straight-through estimator (STE). The underlying tensors remain floating point; +this feature does not provide compressed inference state or native FP8 matmul acceleration. + +The implementation adapts [PR #2455](https://github.com/NVIDIA/Model-Optimizer/pull/2455) +at `13c7e2456f2e9d079c9ef822742eeaa634353802` onto ModelOpt +`051d6adb204f10cd3e78d0f824f31a5a01d54831`. It is the M0/M1 portion of the linear-attention +design. Additional prefill operand sites, KDA training integration, approximate inverse, +token-state quantization, decay approximation, and SSM replay remain later milestones. + +## Requirements and scope + +- Install `fla-core==0.5.1` alongside the training framework. FLA and Triton remain optional + for ordinary ModelOpt imports and the PyTorch references. +- The fused GDN implementation accepts **chunk size 64 only**, including backward. +- State QDQ requires NVIDIA SM89 or newer for the native E4M3 conversion instruction. + W-only QDQ can run on the A6000/SM86 environment used for local validation. +- The numerical recipe supports dynamic E4M3 and `pass_through_bwd=True`. Static scales, + clipping-aware backward, real quantization, rotations, and custom format backends are rejected. +- Context parallelism with either numerical site enabled is rejected. Distributed/sharded + checkpoint qualification remains a release gate; ordinary ModelOpt save/restore is covered. +- Megatron's dynamic-batching `ssm_prefill` and `ssm_decode` paths are outside this integration. + +## Configure training + +Apply the recipe after constructing the Megatron model. The module pattern must match +at least one converted GatedDeltaNet layer; dense attention layers in a hybrid model are +ignored by the execution policy. + +```python +import modelopt.torch.quantization as mtq + +config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*gdn_state_quantizer", + "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1)}, + }, + { + "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": "fla", + "chunk_size": 64, + "state": {"mode": "chunk", "block_v": 64, "quantize_initial": True}, + "solve": {"method": "exact"}, + }, + }, + ], +} +model = mtq.quantize(model, config) +``` + +Dynamic scaling requires no calibration loop. Continue with the framework's normal +forward, backward, and optimizer steps. Enable projection quantization through separate +recipe entries if needed. The standard projection presets keep the GDN numerical handles +disabled; append the GDN units after standard exclusions when composing recipes. + +Execution policies use last-match precedence: each match replaces the complete policy, +including defaults for omitted fields. There is no nested merge. The resolved policy is +reported at conversion and is available on each converted module as +`linear_attention_config`. Quantizer enables/formats and resolved policies are saved +through the ModelOpt checkpoint APIs, including changes made after conversion. + +Without an explicit execution policy, enabled GDN recipe units use the defaults above. +When both GDN handles are disabled, the adapter calls the original training kernel. + +## Numerical contract + +`W` is an activation, **not a learned weight**. The W quantizer sees `[B,T,Hv,Dk]` +and reduces over `Dk`, using one scale per token/head. Forward's quantized W is saved +for backward, avoiding a second quantizer invocation and ensuring identical values. +This costs one saved W activation when that site is enabled. + +The state handle specifies a fused operation rather than calling `TensorQuantizer` on a +materialized state. Each sequence/head uses one dynamic scale per `[Dk, block_v]` tile; +`block_v` can be 16, 32, 64, or 128. Partial value tiles are valid. The scale is +`amax / 448`, with scale one for an all-zero tile. Rounding occurs on the initial state +and after each chunk's state update. The chunk's outputs use its incoming rounded state; +they are computed before rounding the outgoing state. Initial/final state gradients +propagate through these rounding events using identity STE. + +The triangular solve remains exact. No approximate-inverse algorithm is selected or +silently enabled by this integration. + +## References and verification + +The CPU-only reference package is `modelopt.torch.quantization.linear_attention`: + +| Function | Contract | +| --- | --- | +| `recurrent_delta_rule_reference` | Exact GDN scalar-decay or KDA per-key-decay recurrence | +| `chunk_gdn_reference` | Exact GDN triangular solve and chunk algebra; optional state/W QDQ | +| `state_fp8_qdq_reference` | Dynamic E4M3 state tiles with identity STE | + +Use float64 for algebra/gradcheck and float32 for GPU comparison. References support +grouped value heads, packed sequences, nonzero initial states, tails, and either state +layout. They favor explicit arithmetic over speed and are not an efficient training backend. + +```bash +PYTHONPATH=. pytest -q tests/unit/torch/quantization/test_linear_attention_reference.py \ + tests/unit/torch/quantization/plugins/test_gated_delta_net.py \ + tests/unit/torch/quantization/test_quantize_cpu.py +PYTHONPATH=. pytest -q \ + tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py +PYTHONPATH=. pytest -q \ + tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +``` + +The GPU suite compares outputs, final states, and gradients for every differentiable +input against the declared surrogate. It also covers packed boundaries, state layout, +fused gate/beta activation, activation checkpointing, and unsupported-policy errors. +The Megatron test includes an optimizer step. The `gpu` and `gpu_megatron` Nox sessions +install the pinned FLA dependency so missing FLA does not silently remove CI coverage. + +Local validation uses Python 3.12.8, Torch 2.9.1, Triton 3.5.1, FLA 0.5.1, and A6000/SM86. +State-QDQ GPU cases require a separate SM89+ run; the local environment also lacks +Megatron-Core. Those gates must pass before M1 is considered fully qualified. No model +quality, QAT recovery, or throughput result is claimed. 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/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/fla_chunk_delta_h.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py new file mode 100644 index 00000000000..ca8b085470d --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -0,0 +1,979 @@ +# Adapted from: https://github.com/fla-org/flash-linear-attention/blob/516143e31fce/fla/ops/common/chunk_delta_h.py +# Adapted with modifications (marked [ModelOpt]): optional in-kernel FP8 E4M3 fake quantization +# of the carried chunk state (STATE_QDQ), BV as an explicit launch argument, no fla backend +# dispatch or config-cache autotune. +# +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import triton +import triton.language as tl +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.cache import fla_cache_autotune +from fla.ops.utils.op import exp2 +from fla.utils import ( + IS_INTEL, + IS_NVIDIA_BLACKWELL, + IS_NVIDIA_HOPPER, + autotune_cache_kwargs, + check_shared_mem, +) + +from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq + +# ``STATE_QDQ`` modes of the forward state kernel. +STATE_QDQ_OFF = 0 +STATE_QDQ_FP8_DYNAMIC = 1 # FP8 E4M3, one dynamic scale per program tile ([K, BV] of one head) +STATE_QDQ_MAX_BLOCK_V = 128 + + +@triton.jit +def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr): + """[ModelOpt] Dynamic FP8 E4M3 scale over the up-to-four K tiles of one program's state.""" + b_amax = tl.max(tl.abs(b_h1)) + if K > 64: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h2))) + if K > 128: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h3))) + if K > 192: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h4))) + return tl.where(b_amax > 0, b_amax / 448.0, 1.0) + + +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] + +# TODO: Triton mainline fixes a Blackwell tl.dot recurrence race. +# Keep this kernel on num_warps=2 for Blackwell until Triton 3.8 is released +# and we re-validate the wider config space. +# Intel needs more warps than NVIDIA here: 8 warps is ~1.5x faster than the best +# config reachable under the [2, 4] cap. +if IS_NVIDIA_BLACKWELL: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2] +elif IS_INTEL: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2, 4, 8, 16] +else: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2, 4] + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +# ``BV`` is an explicit argument rather than an autotuned config: with ``STATE_QDQ`` it sets the +# quantization granularity, so it must not vary with the autotuner's choice. +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in GATED_DELTA_RULE_FWD_H_NUM_WARPS + for num_stages in ([2, 3, 4] if check_shared_mem("ampere") else [2, 1]) + ], + key=["H", "HV", "K", "V", "BT", "BV", "STATE_V_FIRST", "STATE_QDQ"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( + k, + v, + w, + v_new, + g, + gk, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + SAVE_NEW_VALUE: tl.constexpr, + STATE_V_FIRST: tl.constexpr, + IS_VARLEN: tl.constexpr, + STATE_QDQ: tl.constexpr, +): + pid = tl.program_id(0) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) + i_n, i_h = i_nh // HV, i_nh % HV + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + if STATE_V_FIRST: + b_h1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([BV, 64], dtype=tl.float32) + else: + b_h1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + h += (boh * HV + i_h).to(tl.int64) * K * V + v += (bos * HV + i_h).to(tl.int64) * V + k += (bos * H + i_h // (HV // H)).to(tl.int64) * K + w += (bos * HV + i_h).to(tl.int64) * K + if SAVE_NEW_VALUE: + v_new += (bos * HV + i_h).to(tl.int64) * V + + if USE_INITIAL_STATE: + h0 = h0 + i_nh * K * V + if STORE_FINAL_STATE: + ht = ht + i_nh * K * V + + # load initial state + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K + if USE_INITIAL_STATE: + if STATE_V_FIRST: + p_h0_1 = h0 + o_v[:, None] * K + o_k1[None, :] + m_h0_1 = m_v[:, None] & m_k1[None, :] + else: + p_h0_1 = h0 + o_k1[:, None] * V + o_v[None, :] + m_h0_1 = m_k1[:, None] & m_v[None, :] + b_h1 += tl.load(p_h0_1, mask=m_h0_1, other=0.0).to(tl.float32) + if K > 64: + if STATE_V_FIRST: + p_h0_2 = h0 + o_v[:, None] * K + o_k2[None, :] + m_h0_2 = m_v[:, None] & m_k2[None, :] + else: + p_h0_2 = h0 + o_k2[:, None] * V + o_v[None, :] + m_h0_2 = m_k2[:, None] & m_v[None, :] + b_h2 += tl.load(p_h0_2, mask=m_h0_2, other=0.0).to(tl.float32) + if K > 128: + if STATE_V_FIRST: + p_h0_3 = h0 + o_v[:, None] * K + o_k3[None, :] + m_h0_3 = m_v[:, None] & m_k3[None, :] + else: + p_h0_3 = h0 + o_k3[:, None] * V + o_v[None, :] + m_h0_3 = m_k3[:, None] & m_v[None, :] + b_h3 += tl.load(p_h0_3, mask=m_h0_3, other=0.0).to(tl.float32) + if K > 192: + if STATE_V_FIRST: + p_h0_4 = h0 + o_v[:, None] * K + o_k4[None, :] + m_h0_4 = m_v[:, None] & m_k4[None, :] + else: + p_h0_4 = h0 + o_k4[:, None] * V + o_v[None, :] + m_h0_4 = m_k4[:, None] & m_v[None, :] + b_h4 += tl.load(p_h0_4, mask=m_h0_4, other=0.0).to(tl.float32) + # [ModelOpt] A state read from an FP8 cache is quantized before the first chunk uses it. + if STATE_QDQ == 1: + if K > 192: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) + elif K > 128: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K) + elif K > 64: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K) + else: + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K) + b_h1 = fp8_scalar_qdq(b_h1, b_scale) + if K > 64: + b_h2 = fp8_scalar_qdq(b_h2, b_scale) + if K > 128: + b_h3 = fp8_scalar_qdq(b_h3, b_scale) + if K > 192: + b_h4 = fp8_scalar_qdq(b_h4, b_scale) + + # main recurrence + for i_t in range(NT): + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + if STATE_V_FIRST: + p_h1 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k1[None, :] + m_h1 = m_v[:, None] & m_k1[None, :] + else: + p_h1 = h + i_t_int64 * HV * K * V + o_k1[:, None] * V + o_v[None, :] + m_h1 = m_k1[:, None] & m_v[None, :] + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), mask=m_h1) + if K > 64: + if STATE_V_FIRST: + p_h2 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k2[None, :] + m_h2 = m_v[:, None] & m_k2[None, :] + else: + p_h2 = h + i_t_int64 * HV * K * V + o_k2[:, None] * V + o_v[None, :] + m_h2 = m_k2[:, None] & m_v[None, :] + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), mask=m_h2) + if K > 128: + if STATE_V_FIRST: + p_h3 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k3[None, :] + m_h3 = m_v[:, None] & m_k3[None, :] + else: + p_h3 = h + i_t_int64 * HV * K * V + o_k3[:, None] * V + o_v[None, :] + m_h3 = m_k3[:, None] & m_v[None, :] + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), mask=m_h3) + if K > 192: + if STATE_V_FIRST: + p_h4 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k4[None, :] + m_h4 = m_v[:, None] & m_k4[None, :] + else: + p_h4 = h + i_t_int64 * HV * K * V + o_k4[:, None] * V + o_v[None, :] + m_h4 = m_k4[:, None] & m_v[None, :] + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), mask=m_h4) + + p_w = w + o_t[:, None] * (HV * K) + o_k1[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k1[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) + else: + b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) + if K > 64: + p_w = w + o_t[:, None] * (HV * K) + o_k2[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k2[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h2.to(b_w.dtype), b_v) + if K > 128: + p_w = w + o_t[:, None] * (HV * K) + o_k3[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k3[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h3.to(b_w.dtype), b_v) + if K > 192: + p_w = w + o_t[:, None] * (HV * K) + o_k4[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k4[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h4.to(b_w.dtype), b_v) + p_v = v + o_t[:, None] * (HV * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & m_v[None, :], other=0.0) - b_v + + if SAVE_NEW_VALUE: + p_v = v_new + o_t[:, None] * (HV * V) + o_v[None, :] + tl.store(p_v, b_v.to(p_v.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) + + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + b_g_last = tl.load(g + (bos * HV + last_idx * HV + i_h).to(tl.int64)).to(tl.float32) + p_g = g + (bos * HV + i_h).to(tl.int64) + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] + b_g_last = exp2(b_g_last) + b_h1 *= b_g_last + if K > 64: + b_h2 *= b_g_last + if K > 128: + b_h3 *= b_g_last + if K > 192: + b_h4 *= b_g_last + + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k1, mask=(o_k1 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h1 *= exp2(b_gk_last1)[None, :] + else: + b_h1 *= exp2(b_gk_last1)[:, None] + if K > 64: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k2, mask=(o_k2 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h2 *= exp2(b_gk_last2)[None, :] + else: + b_h2 *= exp2(b_gk_last2)[:, None] + if K > 128: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k3, mask=(o_k3 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h3 *= exp2(b_gk_last3)[None, :] + else: + b_h3 *= exp2(b_gk_last3)[:, None] + if K > 192: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k4, mask=(o_k4 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h4 *= exp2(b_gk_last4)[None, :] + else: + b_h4 *= exp2(b_gk_last4)[:, None] + b_v = b_v.to(k.dtype.element_ty) + + p_k = k + o_k1[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k1[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h1 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h1 = tl.dot(b_k, b_v, b_h1) + if K > 64: + p_k = k + o_k2[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k2[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h2 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h2 = tl.dot(b_k, b_v, b_h2) + if K > 128: + p_k = k + o_k3[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k3[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h3 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h3 = tl.dot(b_k, b_v, b_h3) + if K > 192: + p_k = k + o_k4[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k4[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h4 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h4 = tl.dot(b_k, b_v, b_h4) + + # [ModelOpt] Fake-quantize the state carried into the next chunk (and, after the last + # chunk, the stored final state) to FP8 E4M3. The scale is dynamic over this program's + # [K, BV] tile of the head state; BV == V makes it one scale per sequence and head. + if STATE_QDQ == 1: + if K > 192: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) + elif K > 128: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K) + elif K > 64: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K) + else: + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K) + b_h1 = fp8_scalar_qdq(b_h1, b_scale) + if K > 64: + b_h2 = fp8_scalar_qdq(b_h2, b_scale) + if K > 128: + b_h3 = fp8_scalar_qdq(b_h3, b_scale) + if K > 192: + b_h4 = fp8_scalar_qdq(b_h4, b_scale) + + if STORE_FINAL_STATE: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k1[None, :] + m_ht = m_v[:, None] & m_k1[None, :] + else: + p_ht = ht + o_k1[:, None] * V + o_v[None, :] + m_ht = m_k1[:, None] & m_v[None, :] + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 64: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k2[None, :] + m_ht = m_v[:, None] & m_k2[None, :] + else: + p_ht = ht + o_k2[:, None] * V + o_v[None, :] + m_ht = m_k2[:, None] & m_v[None, :] + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 128: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k3[None, :] + m_ht = m_v[:, None] & m_k3[None, :] + else: + p_ht = ht + o_k3[:, None] * V + o_v[None, :] + m_ht = m_k3[:, None] & m_v[None, :] + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 192: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k4[None, :] + m_ht = m_v[:, None] & m_k4[None, :] + else: + p_ht = ht + o_k4[:, None] * V + o_v[None, :] + m_ht = m_k4[:, None] & m_v[None, :] + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), mask=m_ht) + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "USE_INITIAL_STATE": lambda args: args["dh0"] is not None, + "USE_FINAL_STATE_GRADIENT": lambda args: args["dht"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@fla_cache_autotune( + configs=[ + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in ([2, 3, 4] if check_shared_mem("ampere") else [1]) + for BV in ([32, 64] if check_shared_mem("ada") else [32]) + ], + key=["H", "HV", "K", "V", "BT", "BV", "USE_G", "STATE_V_FIRST"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( + q, + k, + w, + g, + gk, + dht, + dh0, + do, + dh, + dv, + dv2, + cu_seqlens, + chunk_offsets, + scale, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + STATE_V_FIRST: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + pid = tl.program_id(0) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) + i_n, i_h = i_nh // HV, i_nh % HV + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + if STATE_V_FIRST: + b_dh1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([BV, 64], dtype=tl.float32) + else: + b_dh1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + q += (bos * H + i_h // (HV // H)).to(tl.int64) * K + k += (bos * H + i_h // (HV // H)).to(tl.int64) * K + w += (bos * HV + i_h).to(tl.int64) * K + do += (bos * HV + i_h).to(tl.int64) * V + dv += (bos * HV + i_h).to(tl.int64) * V + dv2 += (bos * HV + i_h).to(tl.int64) * V + dh += (boh * HV + i_h).to(tl.int64) * K * V + if USE_GK: + gk += (bos * HV + i_h).to(tl.int64) * K + + if USE_INITIAL_STATE: + dh0 += i_nh * K * V + if USE_FINAL_STATE_GRADIENT: + dht += i_nh * K * V + + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K + if USE_FINAL_STATE_GRADIENT: + if STATE_V_FIRST: + p_dht1 = dht + o_v[:, None] * K + o_k1[None, :] + m_dht1 = m_v[:, None] & m_k1[None, :] + else: + p_dht1 = dht + o_k1[:, None] * V + o_v[None, :] + m_dht1 = m_k1[:, None] & m_v[None, :] + b_dh1 += tl.load(p_dht1, mask=m_dht1, other=0.0) + if K > 64: + if STATE_V_FIRST: + p_dht2 = dht + o_v[:, None] * K + o_k2[None, :] + m_dht2 = m_v[:, None] & m_k2[None, :] + else: + p_dht2 = dht + o_k2[:, None] * V + o_v[None, :] + m_dht2 = m_k2[:, None] & m_v[None, :] + b_dh2 += tl.load(p_dht2, mask=m_dht2, other=0.0) + if K > 128: + if STATE_V_FIRST: + p_dht3 = dht + o_v[:, None] * K + o_k3[None, :] + m_dht3 = m_v[:, None] & m_k3[None, :] + else: + p_dht3 = dht + o_k3[:, None] * V + o_v[None, :] + m_dht3 = m_k3[:, None] & m_v[None, :] + b_dh3 += tl.load(p_dht3, mask=m_dht3, other=0.0) + if K > 192: + if STATE_V_FIRST: + p_dht4 = dht + o_v[:, None] * K + o_k4[None, :] + m_dht4 = m_v[:, None] & m_k4[None, :] + else: + p_dht4 = dht + o_k4[:, None] * V + o_v[None, :] + m_dht4 = m_k4[:, None] & m_v[None, :] + b_dh4 += tl.load(p_dht4, mask=m_dht4, other=0.0) + + for i_t in range(NT - 1, -1, -1): + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + if STATE_V_FIRST: + p_dh1 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k1[None, :] + m_dh1 = m_v[:, None] & m_k1[None, :] + else: + p_dh1 = dh + i_t_int64 * HV * K * V + o_k1[:, None] * V + o_v[None, :] + m_dh1 = m_k1[:, None] & m_v[None, :] + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), mask=m_dh1) + if K > 64: + if STATE_V_FIRST: + p_dh2 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k2[None, :] + m_dh2 = m_v[:, None] & m_k2[None, :] + else: + p_dh2 = dh + i_t_int64 * HV * K * V + o_k2[:, None] * V + o_v[None, :] + m_dh2 = m_k2[:, None] & m_v[None, :] + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), mask=m_dh2) + if K > 128: + if STATE_V_FIRST: + p_dh3 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k3[None, :] + m_dh3 = m_v[:, None] & m_k3[None, :] + else: + p_dh3 = dh + i_t_int64 * HV * K * V + o_k3[:, None] * V + o_v[None, :] + m_dh3 = m_k3[:, None] & m_v[None, :] + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), mask=m_dh3) + if K > 192: + if STATE_V_FIRST: + p_dh4 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k4[None, :] + m_dh4 = m_v[:, None] & m_k4[None, :] + else: + p_dh4 = dh + i_t_int64 * HV * K * V + o_k4[:, None] * V + o_v[None, :] + m_dh4 = m_k4[:, None] & m_v[None, :] + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), mask=m_dh4) + + last_idx = min((i_t_int64 + 1) * BT, T) - 1 + if USE_G: + bg_last = tl.load(g + (bos + last_idx) * HV + i_h).to(tl.float32) + p_g = g + bos * HV + i_h + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + bg_last_exp = exp2(bg_last) + b_g_exp = exp2(b_g) + p_dv = dv + o_t[:, None] * (HV * V) + o_v[None, :] + p_dv2 = dv2 + o_t[:, None] * (HV * V) + o_v[None, :] + p_do = do + o_t[:, None] * (HV * V) + o_v[None, :] + + b_do = tl.load(p_do, mask=m_t[:, None] & m_v[None, :], other=0.0) + + # Update dv + p_k = k + o_t[:, None] * (H * K) + o_k1[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k1[None, :], other=0.0) + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load(gk + last_idx * HV * K + o_k1, mask=(o_k1 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh1).to(b_k.dtype)) + else: + b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) + + if K > 64: + p_k = k + o_t[:, None] * (H * K) + o_k2[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k2[None, :], other=0.0) + if USE_GK: + b_gk_last2 = tl.load(gk + last_idx * HV * K + o_k2, mask=(o_k2 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh2).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh2.to(b_k.dtype), b_dv) + + if K > 128: + p_k = k + o_t[:, None] * (H * K) + o_k3[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k3[None, :], other=0.0) + if USE_GK: + b_gk_last3 = tl.load(gk + last_idx * HV * K + o_k3, mask=(o_k3 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh3).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh3.to(b_k.dtype), b_dv) + + if K > 192: + p_k = k + o_t[:, None] * (H * K) + o_k4[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k4[None, :], other=0.0) + if USE_GK: + b_gk_last4 = tl.load(gk + last_idx * HV * K + o_k4, mask=(o_k4 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh4).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh4.to(b_k.dtype), b_dv) + + if USE_G: + b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] + b_dv += tl.load(p_dv, mask=m_t[:, None] & m_v[None, :], other=0.0) + + tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) + # Update dh + p_w = w + o_k1[:, None] + o_t[None, :] * (HV * K) + p_q = q + o_k1[:, None] + o_t[None, :] * (H * K) + b_w = tl.load(p_w, mask=m_k1[:, None] & m_t[None, :], other=0.0) + b_q = tl.load(p_q, mask=m_k1[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh1 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh1 *= exp2(b_gk_last1)[None, :] + else: + b_dh1 *= exp2(b_gk_last1[:, None]) + if STATE_V_FIRST: + b_dh1 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 64: + p_q = q + o_k2[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k2[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k2[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k2[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh2 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh2 *= exp2(b_gk_last2)[None, :] + else: + b_dh2 *= exp2(b_gk_last2[:, None]) + if STATE_V_FIRST: + b_dh2 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 128: + p_q = q + o_k3[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k3[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k3[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k3[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh3 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh3 *= exp2(b_gk_last3)[None, :] + else: + b_dh3 *= exp2(b_gk_last3[:, None]) + if STATE_V_FIRST: + b_dh3 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 192: + p_q = q + o_k4[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k4[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k4[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k4[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh4 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh4 *= exp2(b_gk_last4)[None, :] + else: + b_dh4 *= exp2(b_gk_last4[:, None]) + if STATE_V_FIRST: + b_dh4 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + + if USE_INITIAL_STATE: + if STATE_V_FIRST: + p_dh0 = dh0 + o_v[:, None] * K + o_k1[None, :] + m_dh0 = m_v[:, None] & m_k1[None, :] + else: + p_dh0 = dh0 + o_k1[:, None] * V + o_v[None, :] + m_dh0 = m_k1[:, None] & m_v[None, :] + tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), mask=m_dh0) + if K > 64: + if STATE_V_FIRST: + p_dh1 = dh0 + o_v[:, None] * K + o_k2[None, :] + m_dh1 = m_v[:, None] & m_k2[None, :] + else: + p_dh1 = dh0 + o_k2[:, None] * V + o_v[None, :] + m_dh1 = m_k2[:, None] & m_v[None, :] + tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), mask=m_dh1) + if K > 128: + if STATE_V_FIRST: + p_dh2 = dh0 + o_v[:, None] * K + o_k3[None, :] + m_dh2 = m_v[:, None] & m_k3[None, :] + else: + p_dh2 = dh0 + o_k3[:, None] * V + o_v[None, :] + m_dh2 = m_k3[:, None] & m_v[None, :] + tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), mask=m_dh2) + if K > 192: + if STATE_V_FIRST: + p_dh3 = dh0 + o_v[:, None] * K + o_k4[None, :] + m_dh3 = m_v[:, None] & m_k4[None, :] + else: + p_dh3 = dh0 + o_k4[:, None] * V + o_v[None, :] + m_dh3 = m_k4[:, None] & m_v[None, :] + tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), mask=m_dh3) + + +def chunk_gated_delta_rule_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + B, T, H, K, V, HV = *k.shape, u.shape[-1], u.shape[2] + BT = chunk_size + BV = state_qdq_tile_v(V, state_qdq, state_qdq_block_v) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + if chunk_offsets is None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + assert K <= 256, "current kernel does not support head dimension larger than 256." + + if state_v_first: + h = k.new_empty(B, NT, HV, V, K) + final_state = k.new_zeros(N, HV, V, K, dtype=torch.float32) if output_final_state else None + else: + h = k.new_empty(B, NT, HV, K, V) + final_state = k.new_zeros(N, HV, K, V, dtype=torch.float32) if output_final_state else None + + v_new = torch.empty_like(u) if save_new_value else None + + grid = (triton.cdiv(V, BV) * N * HV,) + chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( + k=k, + v=u, + w=w, + v_new=v_new, + g=g, + gk=gk, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BV=BV, + STATE_V_FIRST=state_v_first, + STATE_QDQ=state_qdq, + ) + return h, v_new, final_state + + +def state_qdq_tile_v(V: int, state_qdq: int, state_qdq_block_v: int | None) -> int: + """Return the V tile width ``BV`` of the forward state kernel. + + Without state quantization this is fla's largest tile. With it, the tile is also the + quantization granularity: one dynamic scale per ``[K, BV]`` block of a head's state. The + default is fla's 64-column tile, i.e. one scale per sequence and head for ``V <= 64`` and two + for the usual ``V == 128``. ``state_qdq_block_v=128`` gives one scale per 128-wide head but + exceeds the register budget where the kernel is limited to two warps (Blackwell) and spills. + """ + if state_qdq == STATE_QDQ_OFF: + return 64 if check_shared_mem("ada") else 32 + if state_qdq != STATE_QDQ_FP8_DYNAMIC: + raise ValueError(f"Unsupported state_qdq mode {state_qdq}; expected 0 or 1.") + BV = min(triton.next_power_of_2(V), 64) if state_qdq_block_v is None else state_qdq_block_v + if BV < 16 or BV > STATE_QDQ_MAX_BLOCK_V or BV & (BV - 1): + raise ValueError( + f"state_qdq_block_v must be a power of two in [16, {STATE_QDQ_MAX_BLOCK_V}], got {BV}." + ) + return BV + + +def chunk_gated_delta_rule_bwd_dhu( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + do: torch.Tensor, + dv: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + h0: torch.Tensor | None = None, + dht: torch.Tensor | None = None, + scale: float | None = None, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, + use_graph: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V, HV = *q.shape, do.shape[-1], do.shape[2] + # N: the actual number of sequences in the batch with either equal or variable lengths + BT = chunk_size + assert K <= 256, "current kernel does not support head dimension being larger than 256." + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + if chunk_offsets is None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + + if use_graph: + # [ModelOpt] Imported here: fla.ops.utils.graph is absent from released fla versions. + from fla.ops.utils.graph import get_static_buffer + + if state_v_first: + dh = get_static_buffer("dhu_dh_vf", (B, NT, HV, V, K), q.dtype, q.device) + else: + dh = get_static_buffer("dhu_dh", (B, NT, HV, K, V), q.dtype, q.device) + dh0 = ( + get_static_buffer("dhu_dh0", tuple(h0.shape), torch.float32, h0.device) + if h0 is not None + else None + ) + dv2 = get_static_buffer("dhu_dv2", tuple(dv.shape), dv.dtype, dv.device) + else: + if state_v_first: + dh = q.new_empty(B, NT, HV, V, K) + else: + dh = q.new_empty(B, NT, HV, K, V) + dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None + dv2 = torch.empty_like(dv) + + def grid(meta): + return (triton.cdiv(V, meta["BV"]) * N * HV,) + + chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64[grid]( + q=q, + k=k, + w=w, + g=g, + gk=gk, + dht=dht, + dh0=dh0, + do=do, + dh=dh, + dv=dv, + dv2=dv2, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + scale=scale, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + STATE_V_FIRST=state_v_first, + ) + return dh, dh0, dv2 diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py new file mode 100644 index 00000000000..aa88b49d71f --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py @@ -0,0 +1,705 @@ +# 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 autocast_custom_bwd, autocast_custom_fwd, input_guard + +from modelopt.torch.quantization.linear_attention.validation import validate_gdn_quantizer +from modelopt.torch.quantization.nn import TensorQuantizer + +from .fla_chunk_delta_h import ( + STATE_QDQ_FP8_DYNAMIC, + STATE_QDQ_OFF, + chunk_gated_delta_rule_bwd_dhu, + chunk_gated_delta_rule_fwd_h, +) + + +def chunk_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + w_quantizer: TensorQuantizer | None = None, +): + g_input = g if use_gate_in_kernel else None + if use_gate_in_kernel: + g = gdn_gate_chunk_cumsum( + g=g, + A_log=A_log, + chunk_size=chunk_size, + scale=RCP_LN2, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + else: + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + scale=RCP_LN2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # obtain WY representation. u is actually the new v. + # fused kkt + solve_tril + recompute_w_u + w, u, A = chunk_gated_delta_rule_fwd_intra( + k=k, + v=v, + g=g, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + # [ModelOpt] w is an activation (the WY form of the chunk's keys) that lands in memory here, + # so it is fake-quantized once per forward instead of tile by tile inside the kernel. + if w_quantizer is not None: + w = w_quantizer(w) + + if cp_context is not None: + initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=k, + w=w, + u=u, + g=g, + cu_seqlens=cu_seqlens, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + ) + + if cp_context is not None: + initial_state = compress_h0(initial_state, context=cp_context) + + o = chunk_fwd_o( + q=q, + k=k, + v=v_new, + h=h, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + return g, o, A, final_state, initial_state, g_input, w if w_quantizer is not None else None + + +def chunk_gated_delta_rule_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + g_input: torch.Tensor | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + quantized_w: torch.Tensor | None = None, +): + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # [ModelOpt] Reuse forward QDQ exactly; the validated policy uses identity STE. + # This also avoids calling observers or dynamic scale computation in backward. + if quantized_w is not None: + w = quantized_w + + if cp_context is not None: + initial_state = expand_h0(initial_state, context=cp_context) + + # [ModelOpt] The backward recomputes the forward's chunk states, so it sees the same + # fake-quantized states; the state gradient itself passes straight through the QDQ. + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + ) + dv = chunk_bwd_dv_local( + q=q, + k=k, + g=g, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + if cp_context is not None: + # initial_state is None in the CP mode + # We only need to compute dht of current rank and pass it to the backward kernel + dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=q, + k=k, + w=w, + do=do, + dv=dv, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + dht=dht, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=q, + k=k, + w=w, + g=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dq, dk, dw, dg = chunk_bwd_dqkwg( + q=q, + k=k, + v=v_new, + w=w, + g=g, + h=h, + dv=dv, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dk2, dv, db, dg2 = prepare_wy_repr_bwd( + k=k, + v=v, + beta=beta, + g=g, + A=A, + dw=dw, + du=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + dk.add_(dk2) + dg.add_(dg2) + dg = chunk_local_cumsum( + dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices + ) + dA_log, ddt_bias = None, None + if use_gate_in_kernel: + dg, dA_log, ddt_bias = gdn_gate_bwd(g=g_input, A_log=A_log, dt_bias=dt_bias, dyg=dg) + return dq, dk, dv, db, dg, dh0, dA_log, ddt_bias + + +class ChunkGatedDeltaRuleFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + cp_context: FLACPContext | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + w_quantizer: TensorQuantizer | None = None, + ): + q_rstd, k_rstd = None, None + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + + beta_raw = beta + if use_beta_sigmoid_in_kernel: + beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu + ) + g, o, A, final_state, initial_state, g_input, quantized_w = chunk_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cp_context=cp_context, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + use_gate_in_kernel=use_gate_in_kernel, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + w_quantizer=w_quantizer, + ) + ctx.save_for_backward( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + quantized_w, + ) + ctx.scale = scale + ctx.chunk_size = chunk_size + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.use_beta_sigmoid_in_kernel = use_beta_sigmoid_in_kernel + ctx.allow_neg_eigval = allow_neg_eigval + ctx.cp_context = cp_context + ctx.state_v_first = state_v_first + ctx.use_gate_in_kernel = use_gate_in_kernel + ctx.state_qdq = state_qdq + ctx.state_qdq_block_v = state_qdq_block_v + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + ( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + quantized_w, + ) = ctx.saved_tensors + dq, dk, dv, db, dg, dh0, dA_log, ddt_bias = chunk_gated_delta_rule_bwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A=A, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + cp_context=ctx.cp_context, + chunk_indices=chunk_indices, + state_v_first=ctx.state_v_first, + use_gate_in_kernel=ctx.use_gate_in_kernel, + g_input=g_input, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=ctx.chunk_size, + state_qdq=ctx.state_qdq, + state_qdq_block_v=ctx.state_qdq_block_v, + quantized_w=quantized_w, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + if ctx.use_beta_sigmoid_in_kernel: + db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0) + return ( + dq.to(q), + dk.to(k), + dv.to(v), + dg.to(g), + db.to(beta_raw), + None, + dh0, + None, + None, + None, + None, + None, + None, + None, + dA_log, + ddt_bias, + None, + None, + None, + None, + None, + None, + None, + ) + + +# [ModelOpt] Not registered with fla's backend dispatch: another backend must not take over a +# call that asks for state quantization. +@torch.compiler.disable +def chunk_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + **kwargs, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`. + g (torch.Tensor): + (forget) gating tensor of shape `[B, T, HV]`. + When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay). + When `use_gate_in_kernel=True`, `g` is the raw input before gate activation; + the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, HV, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2norm to the q/k tensor internally. Default: `False`. + use_gate_in_kernel (bool): + Whether to compute the log-space GDN decay internally. + When `True`, the passed `g` is the raw input, and `A_log` must be provided. + The kernel fuses gate activation + chunk cumsum in a single pass. + Default: `False`. + A_log (Optional[torch.Tensor]): + Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`. + dt_bias (Optional[torch.Tensor]): + Bias added to `g` before activation, of shape `[HV]`. + Only used when `use_gate_in_kernel=True`. + use_beta_sigmoid_in_kernel (bool): + Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel. + - If `True`, the passed `beta` acts as the raw beta logits. + - If `False`, `beta` is expected to already be in post-sigmoid space. + Default: `False`. + allow_neg_eigval (bool): + Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`. + Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case + the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`. + state_v_first (Optional[bool]): + Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + chunk_indices (Optional[torch.LongTensor]): + Pre-computed chunk indices for variable-length inputs. + If provided, they are used directly instead of being computed from `cu_seqlens`. Default: `None`. + cp_context (Optional[FLACPContext]): + Context parallel context for distributed training across multiple devices. + When provided, `initial_state` and `output_final_state` are not supported, + and `cu_seqlens` will be overridden by the context. Default: `None`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda') + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if fla.__version__ != "0.5.1": + raise RuntimeError(f"ModelOpt GDN requires fla-core==0.5.1, got {fla.__version__}.") + if "transpose_state_layout" in kwargs: + if state_v_first: + raise ValueError( + "Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`." + ) + warnings.warn( + "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", + DeprecationWarning, + stacklevel=2, + ) + state_v_first = kwargs.pop("transpose_state_layout") + + # Validate head dimensions + if q.shape[2] != k.shape[2]: + raise ValueError( + f"q and k must have the same number of heads, " + f"but got q.shape[2]={q.shape[2]} and k.shape[2]={k.shape[2]}" + ) + H, HV = q.shape[2], v.shape[2] + if HV % H != 0: + raise ValueError( + f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by " + f"num_heads (H={H}), but got HV % H = {HV % H}" + ) + + if "head_first" in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) + + chunk_size = kwargs.pop("chunk_size", 64) + if chunk_size != 64: + raise ValueError("ModelOpt GDN supports only chunk_size=64; FLA WY backward assumes 64.") + + # [ModelOpt] state_qdq: 0 keeps fla's numerics; 1 fake-quantizes the state carried between + # chunks to FP8 E4M3 with a dynamic scale per [K, state_qdq_block_v] tile of each head. + state_qdq = kwargs.pop("state_qdq", STATE_QDQ_OFF) + state_qdq_block_v = kwargs.pop("state_qdq_block_v", None) + # w_quantizer: dynamic FP8 TensorQuantizer applied to the WY tensor + # ``w`` of shape [B, T, HV, K] before it multiplies the state, emulating an FP8 x FP8 matmul. + w_quantizer = kwargs.pop("w_quantizer", None) + if state_qdq not in (STATE_QDQ_OFF, STATE_QDQ_FP8_DYNAMIC): + raise ValueError(f"`state_qdq` must be 0 or 1, got {state_qdq}.") + if w_quantizer is not None: + validate_gdn_quantizer(w_quantizer, state=False) + if state_qdq and (not q.is_cuda or torch.cuda.get_device_capability(q.device) < (8, 9)): + raise RuntimeError("GDN state QDQ requires native E4M3 conversion on CUDA SM89 or newer.") + if (state_qdq != STATE_QDQ_OFF or w_quantizer is not None) and cp_context is not None: + raise ValueError("State or w quantization is not supported together with `cp_context`.") + + if cp_context is not None: + assert initial_state is None, "Initial state is not supported for CP" + assert output_final_state is False, "Output final state is not supported for CP" + assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" + cu_seqlens = cp_context.cu_seqlens + if cp_context.cu_seqlens_cpu is not None: + cu_seqlens_cpu = cp_context.cu_seqlens_cpu + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + use_gate_in_kernel = kwargs.get("use_gate_in_kernel", False) + A_log = kwargs.get("A_log") + dt_bias = kwargs.get("dt_bias") + 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 + 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/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..0edccdf4649 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -28,6 +28,7 @@ from modelopt.torch.opt.dynamic import _DMRegistryCls from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict from modelopt.torch.utils import get_unwrapped_name +from modelopt.torch.utils.logging import print_rank_0 from .config import ( QuantizeConfig, @@ -36,6 +37,7 @@ _QuantizeExportConfig, normalize_quant_cfg_list, ) +from .linear_attention.config import LinearAttentionConfig from .nn import ( QuantModule, QuantModuleRegistry, @@ -67,6 +69,7 @@ def convert_to_quantized_model(model: ModelLikeModule, config: QuantizeConfig) - replace_quant_module(model, version=ModeloptStateManager(model).state_version) set_quantizer_by_cfg(model, config.get("quant_cfg", [])) + _apply_linear_attention_policy(model, config) metadata = {} update_quantize_metadata(model, config, metadata) @@ -136,13 +139,22 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: details regarding how MCore sharded checkpoint is restored, see modelopt.torch.opt.plugins.mcore_dist_checkpointing.restore_sharded_modelopt_state. """ + _apply_linear_attention_policy(model, config, metadata.get("linear_attention")) if "quantizer_state" not in metadata: # MCore sharded checkpoint (`torch-dist`) has its quantizer_state stored as the # extra_state of `QuantModule`. The quantizer_state is resumed with # QuantModule.set_extra_state(). return model - quantizer_state_dict = metadata["quantizer_state"] + quantizer_state_dict = dict(metadata["quantizer_state"]) + if "linear_attention" not in metadata: + # Older checkpoints predate these disabled handles; preserve their baseline path. + for name, module in _linear_attention_modules(model).items(): + for handle in ("gdn_state_quantizer", "gdn_w_quantizer"): + key = f"{name}.{handle}" if name else handle + quantizer = getattr(module, handle) + if key not in quantizer_state_dict and not quantizer.is_enabled: + quantizer_state_dict[key] = quantizer.get_modelopt_state() unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys() extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys() @@ -195,6 +207,12 @@ def update_quantize_metadata( model: nn.Module, config: QuantizeConfig, metadata: MetadataDict ) -> None: """Update the quantizer state in the metadata dict.""" + policies = { + name: module.linear_attention_config.model_dump() + for name, module in _linear_attention_modules(model).items() + } + if policies: + metadata["linear_attention"] = policies metadata["quantizer_state"] = quantizer_state(model) if shared_state_metadata := SharedWeightGlobalAmaxState.metadata(model): metadata["shared_quant_states"] = shared_state_metadata @@ -202,6 +220,43 @@ def update_quantize_metadata( metadata.pop("shared_quant_states", None) +def _linear_attention_modules(model): + # Optional framework plugins import conversion; defer this import to avoid that cycle. + from .plugins.gated_delta_net import GatedDeltaNetStateQuantMixin + + return { + get_unwrapped_name(name, model): module + for name, module in model.named_modules() + if isinstance(module, GatedDeltaNetStateQuantMixin) + } + + +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..49f5ea3514a --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -0,0 +1,19 @@ +# 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 references for linear attention.""" + +from .config import * +from .reference 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..6405fbbf329 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -0,0 +1,60 @@ +# 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 + +from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField + +__all__ = ["LinearAttentionConfig", "LinearAttentionPolicyEntry"] + + +class _StateConfig(ModeloptBaseConfig): + mode: Literal["chunk"] = ModeloptField(default="chunk") + block_v: Literal[16, 32, 64, 128] = ModeloptField(default=64) + quantize_initial: Literal[True] = ModeloptField(default=True) + + +class _SolveConfig(ModeloptBaseConfig): + method: Literal["exact"] = ModeloptField(default="exact") + + +class LinearAttentionConfig(ModeloptBaseConfig): + """GDN chunk-64 policy; unsupported numerical modes fail config validation. + + ``state.block_v`` defines one dynamic scale per ``[Dk, block_v]`` tile of each + sequence/head. The initial state and every chunk's final state are rounded when + ``gdn_state_quantizer`` is enabled. Outputs use the incoming rounded state. + """ + + schema_version: Literal[1] = ModeloptField(default=1) + backend: Literal["fla"] = ModeloptField(default="fla") + chunk_size: Literal[64] = ModeloptField(default=64) + state: _StateConfig = ModeloptField(default=_StateConfig()) + solve: _SolveConfig = ModeloptField(default=_SolveConfig()) + + +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/reference.py b/modelopt/torch/quantization/linear_attention/reference.py new file mode 100644 index 00000000000..65ea183a318 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/reference.py @@ -0,0 +1,203 @@ +# 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", "recurrent_delta_rule_reference", "state_fp8_qdq_reference"] + + +def state_fp8_qdq_reference(state: torch.Tensor, block_v: int = 64) -> torch.Tensor: + """Dynamic E4M3 QDQ with identity STE and one scale per ``[Dk, block_v]`` tile. + + ``state`` is in key-first layout ``[..., Dk, Dv]``. Scales and rounding do not + contribute derivatives. The zero tile uses scale one; partial value tiles are valid. + """ + if block_v not in (16, 32, 64, 128): + raise ValueError("block_v must be 16, 32, 64, or 128") + with torch.no_grad(): + rounded = [] + for tile in state.float().split(block_v, dim=-1): + amax = tile.abs().amax(dim=(-2, -1), keepdim=True) + scale = torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)) + rounded.append((tile / scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() * scale) + 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, + w_quantizer: 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. + """ + 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 + chunks, all_w = [], [] + for n, (b, start, end) in enumerate(sequences): + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + gc = g[b, lo:hi].transpose(0, 1).cumsum(-1) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + # Mask before exp: upper-triangle positive differences can overflow for long decay. + causal = torch.ones(hi - lo, hi - lo, device=q.device, dtype=torch.bool).tril() + decay = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() + gram = kc @ kc.transpose(-1, -2) + lower = (bc * gram * decay).tril(-1) + matrix = lower + torch.eye(hi - lo, device=q.device, dtype=q.dtype) + rhs = torch.cat((bc * vc, bc * kc * gc.exp().unsqueeze(-1)), dim=-1) + solved = torch.linalg.solve_triangular(matrix, rhs, upper=False, unitriangular=True) + u, w = solved.split((v.shape[-1], k.shape[-1]), dim=-1) + chunks.append((n, qc, kc, gc, decay, u)) + all_w.append(w.transpose(0, 1)) + w = torch.cat(all_w).reshape(q.shape) + if w_quantizer is not None: + w = w_quantizer(w) + w = w.reshape(-1, *w.shape[2:]) + outputs, finals, offset, previous_n = [], [], 0, -1 + for n, qc, kc, gc, decay, u in chunks: + if n != previous_n: + state = states[n] + if state_qdq: + state = state_fp8_qdq_reference(state, state_qdq_block_v) + previous_n = n + length = qc.shape[1] + wc = w[offset : offset + length].transpose(0, 1) + offset += length + updated_values = u - wc @ state # state_read + local_scores = ((qc * scale) @ kc.transpose(-1, -2) * decay).tril() + output = (qc * (scale * gc.exp()).unsqueeze(-1)) @ state + output = output + local_scores @ updated_values # local_readout + outputs.append(output.transpose(0, 1)) + weighted_keys = kc * (gc[..., -1:] - gc).exp().unsqueeze(-1) + state = state * gc[..., -1].exp()[:, None, None] + state = state + weighted_keys.transpose(-1, -2) @ updated_values # state_update + if state_qdq: + state = state_fp8_qdq_reference(state, state_qdq_block_v) + if len(finals) <= n: + finals.append(state) + else: + finals[n] = state + output = torch.cat(outputs).reshape(*v.shape) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final diff --git a/modelopt/torch/quantization/linear_attention/validation.py b/modelopt/torch/quantization/linear_attention/validation.py new file mode 100644 index 00000000000..2bf07525e22 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/validation.py @@ -0,0 +1,45 @@ +# 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. + +"""Capabilities of the initial fused GDN fake-quant path.""" + +from ..nn import TensorQuantizer + +__all__ = [] + + +def validate_gdn_quantizer(quantizer: TensorQuantizer, *, state: bool) -> None: + name = "gdn_state_quantizer" if state else "gdn_w_quantizer" + axis = (0, 1) if state else (0, 1, 2) + if not isinstance(quantizer, TensorQuantizer): + raise ValueError(f"{name} requires a single TensorQuantizer") + if not ( + quantizer._dynamic + and quantizer.num_bits == (4, 3) + and quantizer.axis == axis + 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 E4M3 fake quantization with axis={axis}, " + "pass_through_bwd=True, no block_sizes, rotation, pre-scaling, bias, constant " + "amax, or custom backend. Other gradient rules and formats are not implemented." + ) 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..963452ff509 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -39,6 +39,7 @@ from .attention import * from .custom import * +from .gated_delta_net import * with import_plugin("diffusers"): from .diffusion.diffusers import * diff --git a/modelopt/torch/quantization/plugins/gated_delta_net.py b/modelopt/torch/quantization/plugins/gated_delta_net.py new file mode 100644 index 00000000000..7be046b8800 --- /dev/null +++ b/modelopt/torch/quantization/plugins/gated_delta_net.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. + +"""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, ModelOpt runs an adapted copy of the kernel +(:mod:`modelopt.torch.kernels.quantization.linear_attention`) that fake-quantizes the state to +E4M3 at the end of every chunk, with a scale computed inside the kernel from the state itself. +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 typing import Any + +import torch + +from ..config import QuantizerAttributeConfig +from ..linear_attention.config import LinearAttentionConfig +from ..linear_attention.validation import validate_gdn_quantizer +from ..nn import QuantModule, TensorQuantizer + +__all__ = ["GatedDeltaNetStateQuantMixin"] + +GatedDeltaRuleFn = Callable[..., tuple[torch.Tensor, torch.Tensor | None]] + + +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(QuantModule): + """Adds ``gdn_state_quantizer`` and ``gdn_w_quantizer`` to a GatedDeltaNet module. + + Subclasses route the module's chunked gated-delta-rule call through + :meth:`_state_quantized_chunk_gated_delta_rule`. Both quantizers start disabled; enable them + with ``quant_cfg`` entries on ``*gdn_state_quantizer`` / ``*gdn_w_quantizer`` such as the + ``configs/ptq/units/gdn_state_fp8_dynamic`` and ``gdn_w_fp8_dynamic`` recipe units. The state + quantizer carries the fused QDQ configuration. Both sites currently require dynamic + E4M3 and identity STE. The execution policy is saved in ModelOpt metadata. + """ + + def _setup(self): + self.gdn_state_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) + self.gdn_w_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) + self.linear_attention_config = LinearAttentionConfig() + + @property + def gdn_state_qdq_block_v(self) -> int: + """Value-column scale grouping from the saved execution policy.""" + return self.linear_attention_config.state.block_v + + def validate_linear_attention(self) -> None: + """Reject numerical settings that the fused training path cannot implement.""" + for state, quantizer in ((True, self.gdn_state_quantizer), (False, self.gdn_w_quantizer)): + if quantizer.is_enabled: + validate_gdn_quantizer(quantizer, state=state) + + def modelopt_post_restore(self, prefix: str = ""): + """Validate restored quantizers before using the saved execution policy.""" + super().modelopt_post_restore(prefix) + self.validate_linear_attention() + + def _state_quantized_chunk_gated_delta_rule( + self, gated_delta_rule: GatedDeltaRuleFn, *args: Any, **kwargs: Any + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """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 (quantize_state or quantize_w): + return gated_delta_rule(*args, **kwargs) + if getattr(gated_delta_rule, "__name__", None) != "chunk_gated_delta_rule": + raise NotImplementedError( + "GatedDeltaNet quantizers require the fla chunked kernel; the deterministic torch " + f"kernel ({gated_delta_rule!r}) is not supported." + ) + 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") + return _state_qdq_chunk_gated_delta_rule()( + *args, + chunk_size=chunk_size, + state_qdq=int(quantize_state), + 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/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index ec1958649dd..3a0d4627e38 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 .gated_delta_net 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,39 @@ 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. + + Training and calibration reach the chunked kernel through ``forward_pre_attn_and_core_attn``, + which calls ``self.gated_delta_rule``; that call is routed through ``gdn_state_quantizer``. + The 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(), + ) + + def forward_pre_attn_and_core_attn(self, *args, **kwargs): + gated_delta_rule = self.gated_delta_rule + self.gated_delta_rule = partial( + self._state_quantized_chunk_gated_delta_rule, gated_delta_rule + ) + try: + return super().forward_pre_attn_and_core_attn(*args, **kwargs) + finally: + self.gated_delta_rule = gated_delta_rule + + def _is_supported_megatron_model(model: torch.nn.Module) -> bool: return isinstance(model, MegatronModule) diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index db37b9222ca..4b6ccec395f 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -32,3 +32,7 @@ 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_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 | + +See the [linear-attention training guide](../../../../docs/linear_attention_qat.md) for GDN policy, hardware, checkpoint, and gradient requirements. diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index 2fee6279973..cac4f6d80aa 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -16,6 +16,11 @@ # 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: '*block_sparse_moe.gate*' enable: false - quantizer_name: '*linear_attn.conv1d*' diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml new file mode 100644 index 00000000000..51f9ad1de20 --- /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. + +# FP8 E4M3 fake quantization of the GatedDeltaNet (linear attention) recurrent state with a +# dynamic scale computed inside the kernel from the state itself (no calibration). The state is +# ``[N, H, K, V]`` for ``N`` sequences; ``axis: [0, 1]`` scales each sequence's head state on its +# own, in tiles of ``gdn_state_qdq_block_v`` value columns (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 +# is the only configuration the kernel implements. 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_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/noxfile.py b/noxfile.py index 135465ed445..d061ef31f57 100644 --- a/noxfile.py +++ b/noxfile.py @@ -130,6 +130,7 @@ def gpu(session): "mamba_ssm", "causal-conv1d", ) + session.run("python", "-m", "pip", "install", "fla-core==0.5.1") session.run("python", "-m", "pytest", "tests/gpu", *_cov_args()) @@ -141,6 +142,7 @@ 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") 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/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..c78668febde --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py @@ -0,0 +1,212 @@ +# 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.""" + +import pytest +import torch +import torch.nn.functional as F +from torch.utils.checkpoint import checkpoint + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import chunk_gdn_reference +from modelopt.torch.quantization.nn import TensorQuantizer + +fla = pytest.importorskip("fla.ops.gated_delta_rule") +from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( + chunk_gated_delta_rule, +) + + +def require_state_qdq(): + if torch.cuda.get_device_capability() < (8, 9): + pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") + + +def make_inputs(dtype=torch.float32, packed=False, state_v_first=False): + 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 + 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(): + return TensorQuantizer( + QuantizerAttributeConfig(num_bits=(4, 3), axis=(0, 1, 2), 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 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.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_disabled_matches_upstream_forward_and_backward(dtype): + args, state = make_inputs(dtype) + 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"), [(False, False), (False, True), (True, False), (True, True)] +) +@pytest.mark.parametrize(("packed", "state_v_first"), [(False, False), (True, True)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_numerical_surrogate_reference(state_qdq, quantize_w, packed, state_v_first, dtype): + if state_qdq: + 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, **kwargs) + actual = values_and_grads( + chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs + ) + compare(actual[0], expected[0], 0.03 if dtype == torch.bfloat16 else 0.01) + compare(actual[1], expected[1], 0.05 if dtype == torch.bfloat16 else 0.02) + + +def test_w_qdq_saved_once_and_activation_checkpoint_parity(): + 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 + ) + + 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() + + +def test_zero_initial_state_and_output_only_training_loss(): + args, _ = make_inputs() + quantizer = w_quantizer() + expected, _ = chunk_gdn_reference(*args, w_quantizer=quantizer) + actual, final = chunk_gated_delta_rule(*args, w_quantizer=quantizer) + assert final is None + compare((actual,), (expected,), 0.01) + grads = [torch.autograd.grad(x.square().sum(), args) for x in (actual, expected)] + compare(*grads, 0.02) + + +@pytest.mark.parametrize("block_v", [16, 64, 128]) +def test_state_scale_tile_forward_and_backward(block_v): + require_state_qdq() + args, state = make_inputs() + kwargs = {"state_qdq": True, "state_qdq_block_v": block_v} + expected = values_and_grads(chunk_gdn_reference, args, state, **kwargs) + actual = values_and_grads( + chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs + ) + compare(actual[0], expected[0], 0.01) + compare(actual[1], expected[1], 0.02) + + +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, dim=-1), + F.normalize(k, dim=-1), + v, + -a_log.exp() * F.softplus(raw_g + bias), + 2 * raw_beta.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.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.02) + + +@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) 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..2252eb25a61 --- /dev/null +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -0,0 +1,101 @@ +# 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.megatron.models import get_mcore_gpt_model +from _test_utils.torch.megatron.utils import get_forward, initialize_for_megatron + +import modelopt.torch.quantization as mtq + +pytest.importorskip("fla") # Megatron-Core GatedDeltaNet and the state QDQ kernel need fla +pytest.importorskip("megatron.core.ssm.gated_delta_net") + +from modelopt.torch.quantization.plugins.gated_delta_net 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 +GDN_STATE_QUANT_CFG = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*gdn_state_quantizer", + "cfg": {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"}, + }, + ], + "algorithm": "max", +} + + +def _test_gdn_state_quant_helper(rank, size): + initialize_for_megatron( + tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED + ) + model = get_mcore_gpt_model( + tensor_model_parallel_size=size, + num_layers=2, + hidden_size=64, + num_attention_heads=4, + vocab_size=32, + max_sequence_length=128, + experimental_attention_variant="gated_delta_net", + ).cuda() + model.eval() # no dropout, so the forwards differ only through the state quantizer + forward = get_forward(model) + with torch.no_grad(): + loss_ref = forward(model) + + model = mtq.quantize(model, GDN_STATE_QUANT_CFG, forward) + + gdn_modules = [m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)] + assert gdn_modules, "no GatedDeltaNet layer was wrapped" + assert all(m.gdn_state_quantizer.is_enabled for m in gdn_modules) + assert all(m.gdn_state_qdq_block_v == 64 for m in gdn_modules) # kernel's 64-column tile + + with torch.no_grad(): + loss_quant = forward(model) + assert torch.isfinite(loss_quant).all() + assert not torch.allclose(loss_quant, loss_ref, rtol=1e-4, atol=1e-4), ( + "FP8 state quantization must change the output" + ) + + mtq.disable_quantizer(model, "*gdn_state_quantizer") + with torch.no_grad(): + assert torch.allclose(forward(model), loss_ref, rtol=1e-4, atol=1e-4) + mtq.enable_quantizer(model, "*gdn_state_quantizer") + + # QAD trains through the segmented kernel: gradients must reach the GDN projections. + model.train() + optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) + before = gdn_modules[0].in_proj.weight.detach().clone() + optimizer.zero_grad() + forward(model).sum().backward() + assert all(m.in_proj.weight.grad is not None for m in gdn_modules) + assert all(torch.isfinite(m.in_proj.weight.grad).all() for m in gdn_modules) + optimizer.step() + assert not torch.equal(gdn_modules[0].in_proj.weight, before) + + +def test_gdn_state_quant(dist_workers_size_1): + """GatedDeltaNet layers get a ``gdn_state_quantizer`` whose FP8 quant-dequant of the recurrent + state runs inside the chunked kernel during forward and backward.""" + if torch.cuda.get_device_capability() < (8, 9): + pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") + dist_workers_size_1.run(_test_gdn_state_quant_helper) diff --git a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py b/tests/unit/torch/quantization/plugins/test_gated_delta_net.py new file mode 100644 index 00000000000..2bdfa2e7b09 --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_gated_delta_net.py @@ -0,0 +1,276 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from copy import deepcopy + +import pytest +import torch +import torch.nn as nn + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.config import QuantizeConfig +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.plugins import gated_delta_net +from modelopt.torch.quantization.plugins.gated_delta_net import GatedDeltaNetStateQuantMixin + +GDN_STATE_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"} + + +def chunk_gated_delta_rule(q, k, v, g, beta, **kwargs): + """Stand-in with the fla kernel's name; the real kernel needs a GPU.""" + return q + k + v, None + + +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_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 model.gdn_state_qdq_block_v == 64 + assert torch.equal(model(x), expected) + assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" + + +def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch): + calls = [] + + def fake_state_qdq_kernel(*args, **kwargs): + calls.append(kwargs) + return chunk_gated_delta_rule(*args) + + monkeypatch.setattr( + gated_delta_net, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel + ) + model = TinyGatedDeltaNet() + x = torch.randn(2, 8, 3, 4) + mtq.quantize(model, quant_cfg(), lambda m: m(x)) + + model(x) + assert calls and calls[-1] == { + "chunk_size": 64, + "state_qdq": 1, + "state_qdq_block_v": 64, + "w_quantizer": None, + } + + # The deterministic torch kernel has no quantized counterpart. + model.gated_delta_rule = lambda *a, **kw: chunk_gated_delta_rule(*a, **kw) + with pytest.raises(NotImplementedError, match="deterministic torch kernel"): + 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( + gated_delta_net, "_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) + + +def test_policy_roundtrip_last_match_and_hybrid_selection(tmp_path): + model = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) + cfg = quant_cfg() + 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 + # A post-conversion policy edit must also survive save/load. + model[0].linear_attention_config.state.block_v = 16 + 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_quantizer.is_enabled + assert restored[0].gdn_state_qdq_block_v == 16 + assert not hasattr(restored[1], "linear_attention_config") + + +def test_policy_rejects_unmatched_and_unimplemented_modes(): + with pytest.raises(ValueError, match="matches no supported"): + mtq.quantize( + nn.Linear(4, 4), {"linear_attention": [{"module_name": "*"}], "algorithm": None} + ) + for policy in ( + {"chunk_size": 32}, + {"solve": {"method": "neumann"}}, + {"state": {"mode": "token"}}, + ): + with pytest.raises(ValueError): + QuantizeConfig(linear_attention=[{"module_name": "*", "cfg": policy}]) + + +def test_restore_old_config_without_policy_field(): + model = mtq.quantize( + nn.Linear(4, 4), + {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None}, + ) + state = mto.modelopt_state(model) + for _, mode_state in state["modelopt_state_dict"]: + config = mode_state["config"] + if isinstance(config, dict): + config.pop("linear_attention", None) + else: + config.__dict__.pop("linear_attention", None) + restored = nn.Linear(4, 4) + mto.restore_from_modelopt_state(restored, state) + restored.load_state_dict(model.state_dict()) + x = torch.randn(2, 4) + torch.testing.assert_close(restored(x), model(x)) + + +def test_policy_refinement_updates_existing_quantized_module(): + cfg = quant_cfg() + cfg["algorithm"] = None + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + cfg["linear_attention"] = [{"module_name": "", "cfg": {"state": {"block_v": 32}}}] + mtq.quantize(model, cfg) + assert model.gdn_state_qdq_block_v == 32 + # A later complete rule resets omitted fields instead of merging the previous policy. + cfg["linear_attention"].append({"module_name": "", "cfg": {}}) + mtq.quantize(model, cfg) + assert model.gdn_state_qdq_block_v == 64 + + +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 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..dd2aea8d592 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_reference.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. + +import pytest +import torch +import torch.nn.functional as F + +from modelopt.torch.quantization.linear_attention 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])) From bfbb13ed98292fb37e540788bc72774c5bb030b0 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 21 Sep 2026 22:54:29 -0700 Subject: [PATCH 2/8] Qualify GDN QAT on Hopper and Megatron training Signed-off-by: Kai Xu --- .github/workflows/gpu_tests.yml | 3 +- CHANGELOG.rst | 2 +- docs/linear_attention_qat.md | 130 -------------- .../fla_chunk_gated_delta_rule.py | 28 ++- modelopt/torch/quantization/config.py | 7 - modelopt/torch/quantization/conversion.py | 51 ++---- .../quantization/linear_attention/__init__.py | 5 +- .../quantization/linear_attention/config.py | 60 ------- .../{validation.py => utils.py} | 10 +- modelopt/torch/quantization/model_quant.py | 4 +- .../quantization/plugins/gated_delta_net.py | 53 +++--- .../torch/quantization/plugins/megatron.py | 28 ++- modelopt_recipes/configs/ptq/units/README.md | 2 - .../ptq/units/gdn_state_fp8_dynamic.yaml | 5 +- noxfile.py | 20 ++- .../linear_attention_reference.py | 0 .../test_fla_chunk_gated_delta_rule.py | 131 +++++++++++--- .../plugins/test_megatron_gated_delta_net.py | 165 ++++++++++++------ .../plugins/test_gated_delta_net.py | 122 ++++++------- .../test_linear_attention_reference.py | 3 +- 20 files changed, 404 insertions(+), 425 deletions(-) delete mode 100644 docs/linear_attention_qat.md delete mode 100644 modelopt/torch/quantization/linear_attention/config.py rename modelopt/torch/quantization/linear_attention/{validation.py => utils.py} (83%) rename modelopt/torch/quantization/linear_attention/reference.py => tests/_test_utils/torch/quantization/linear_attention_reference.py (100%) 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/CHANGELOG.rst b/CHANGELOG.rst index bc6a6ae6e7a..347806df32b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -17,7 +17,7 @@ Changelog - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration. - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md `_ for details. - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``. -- Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training, with saved execution policies and differentiable GDN/KDA references. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer. +- Add experimental dynamic FP8 fake quantization of GatedDeltaNet 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/docs/linear_attention_qat.md b/docs/linear_attention_qat.md deleted file mode 100644 index 0db84b58f22..00000000000 --- a/docs/linear_attention_qat.md +++ /dev/null @@ -1,130 +0,0 @@ -# Linear-attention numerical emulation for training - -This experimental integration adds dynamic FP8 E4M3 fake quantization to Megatron -GatedDeltaNet's chunked training path. It rounds the recurrent state at chunk boundaries -and the materialized WY activation `W` before the state-read matmul. Both sites use an -identity straight-through estimator (STE). The underlying tensors remain floating point; -this feature does not provide compressed inference state or native FP8 matmul acceleration. - -The implementation adapts [PR #2455](https://github.com/NVIDIA/Model-Optimizer/pull/2455) -at `13c7e2456f2e9d079c9ef822742eeaa634353802` onto ModelOpt -`051d6adb204f10cd3e78d0f824f31a5a01d54831`. It is the M0/M1 portion of the linear-attention -design. Additional prefill operand sites, KDA training integration, approximate inverse, -token-state quantization, decay approximation, and SSM replay remain later milestones. - -## Requirements and scope - -- Install `fla-core==0.5.1` alongside the training framework. FLA and Triton remain optional - for ordinary ModelOpt imports and the PyTorch references. -- The fused GDN implementation accepts **chunk size 64 only**, including backward. -- State QDQ requires NVIDIA SM89 or newer for the native E4M3 conversion instruction. - W-only QDQ can run on the A6000/SM86 environment used for local validation. -- The numerical recipe supports dynamic E4M3 and `pass_through_bwd=True`. Static scales, - clipping-aware backward, real quantization, rotations, and custom format backends are rejected. -- Context parallelism with either numerical site enabled is rejected. Distributed/sharded - checkpoint qualification remains a release gate; ordinary ModelOpt save/restore is covered. -- Megatron's dynamic-batching `ssm_prefill` and `ssm_decode` paths are outside this integration. - -## Configure training - -Apply the recipe after constructing the Megatron model. The module pattern must match -at least one converted GatedDeltaNet layer; dense attention layers in a hybrid model are -ignored by the execution policy. - -```python -import modelopt.torch.quantization as mtq - -config = { - "quant_cfg": [ - {"quantizer_name": "*", "enable": False}, - { - "quantizer_name": "*gdn_state_quantizer", - "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1)}, - }, - { - "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": "fla", - "chunk_size": 64, - "state": {"mode": "chunk", "block_v": 64, "quantize_initial": True}, - "solve": {"method": "exact"}, - }, - }, - ], -} -model = mtq.quantize(model, config) -``` - -Dynamic scaling requires no calibration loop. Continue with the framework's normal -forward, backward, and optimizer steps. Enable projection quantization through separate -recipe entries if needed. The standard projection presets keep the GDN numerical handles -disabled; append the GDN units after standard exclusions when composing recipes. - -Execution policies use last-match precedence: each match replaces the complete policy, -including defaults for omitted fields. There is no nested merge. The resolved policy is -reported at conversion and is available on each converted module as -`linear_attention_config`. Quantizer enables/formats and resolved policies are saved -through the ModelOpt checkpoint APIs, including changes made after conversion. - -Without an explicit execution policy, enabled GDN recipe units use the defaults above. -When both GDN handles are disabled, the adapter calls the original training kernel. - -## Numerical contract - -`W` is an activation, **not a learned weight**. The W quantizer sees `[B,T,Hv,Dk]` -and reduces over `Dk`, using one scale per token/head. Forward's quantized W is saved -for backward, avoiding a second quantizer invocation and ensuring identical values. -This costs one saved W activation when that site is enabled. - -The state handle specifies a fused operation rather than calling `TensorQuantizer` on a -materialized state. Each sequence/head uses one dynamic scale per `[Dk, block_v]` tile; -`block_v` can be 16, 32, 64, or 128. Partial value tiles are valid. The scale is -`amax / 448`, with scale one for an all-zero tile. Rounding occurs on the initial state -and after each chunk's state update. The chunk's outputs use its incoming rounded state; -they are computed before rounding the outgoing state. Initial/final state gradients -propagate through these rounding events using identity STE. - -The triangular solve remains exact. No approximate-inverse algorithm is selected or -silently enabled by this integration. - -## References and verification - -The CPU-only reference package is `modelopt.torch.quantization.linear_attention`: - -| Function | Contract | -| --- | --- | -| `recurrent_delta_rule_reference` | Exact GDN scalar-decay or KDA per-key-decay recurrence | -| `chunk_gdn_reference` | Exact GDN triangular solve and chunk algebra; optional state/W QDQ | -| `state_fp8_qdq_reference` | Dynamic E4M3 state tiles with identity STE | - -Use float64 for algebra/gradcheck and float32 for GPU comparison. References support -grouped value heads, packed sequences, nonzero initial states, tails, and either state -layout. They favor explicit arithmetic over speed and are not an efficient training backend. - -```bash -PYTHONPATH=. pytest -q tests/unit/torch/quantization/test_linear_attention_reference.py \ - tests/unit/torch/quantization/plugins/test_gated_delta_net.py \ - tests/unit/torch/quantization/test_quantize_cpu.py -PYTHONPATH=. pytest -q \ - tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py -PYTHONPATH=. pytest -q \ - tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py -``` - -The GPU suite compares outputs, final states, and gradients for every differentiable -input against the declared surrogate. It also covers packed boundaries, state layout, -fused gate/beta activation, activation checkpointing, and unsupported-policy errors. -The Megatron test includes an optimizer step. The `gpu` and `gpu_megatron` Nox sessions -install the pinned FLA dependency so missing FLA does not silently remove CI coverage. - -Local validation uses Python 3.12.8, Torch 2.9.1, Triton 3.5.1, FLA 0.5.1, and A6000/SM86. -State-QDQ GPU cases require a separate SM89+ run; the local environment also lacks -Megatron-Core. Those gates must pass before M1 is considered fully qualified. No model -quality, QAT recovery, or throughput result is claimed. diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py index aa88b49d71f..a5609cab102 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py @@ -46,9 +46,15 @@ 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 autocast_custom_bwd, autocast_custom_fwd, input_guard +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.validation import validate_gdn_quantizer +from modelopt.torch.quantization.linear_attention.utils import validate_gdn_quantizer from modelopt.torch.quantization.nn import TensorQuantizer from .fla_chunk_delta_h import ( @@ -636,10 +642,15 @@ def chunk_gated_delta_rule( # w_quantizer: dynamic FP8 TensorQuantizer applied to the WY tensor # ``w`` of shape [B, T, HV, K] before it multiplies the state, emulating an FP8 x FP8 matmul. w_quantizer = kwargs.pop("w_quantizer", None) + use_gate_in_kernel = kwargs.pop("use_gate_in_kernel", False) + A_log = kwargs.pop("A_log", None) + dt_bias = kwargs.pop("dt_bias", None) + if kwargs: + raise TypeError(f"Unexpected keyword arguments: {', '.join(sorted(kwargs))}") if state_qdq not in (STATE_QDQ_OFF, STATE_QDQ_FP8_DYNAMIC): raise ValueError(f"`state_qdq` must be 0 or 1, got {state_qdq}.") if w_quantizer is not None: - validate_gdn_quantizer(w_quantizer, state=False) + validate_gdn_quantizer(w_quantizer, name="gdn_w_quantizer") if state_qdq and (not q.is_cuda or torch.cuda.get_device_capability(q.device) < (8, 9)): raise RuntimeError("GDN state QDQ requires native E4M3 conversion on CUDA SM89 or newer.") if (state_qdq != STATE_QDQ_OFF or w_quantizer is not None) and cp_context is not None: @@ -664,9 +675,6 @@ def chunk_gated_delta_rule( 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]}.", ) - use_gate_in_kernel = kwargs.get("use_gate_in_kernel", False) - A_log = kwargs.get("A_log") - dt_bias = kwargs.get("dt_bias") 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: @@ -674,6 +682,14 @@ def chunk_gated_delta_rule( 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, diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index e8db5c8446a..f8847aa1fc1 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -161,8 +161,6 @@ 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.""" @@ -1569,11 +1567,6 @@ 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 0edccdf4649..78a6e03a16b 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -28,7 +28,6 @@ from modelopt.torch.opt.dynamic import _DMRegistryCls from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict from modelopt.torch.utils import get_unwrapped_name -from modelopt.torch.utils.logging import print_rank_0 from .config import ( QuantizeConfig, @@ -37,7 +36,6 @@ _QuantizeExportConfig, normalize_quant_cfg_list, ) -from .linear_attention.config import LinearAttentionConfig from .nn import ( QuantModule, QuantModuleRegistry, @@ -69,7 +67,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) + _validate_linear_attention_quantizers(model) metadata = {} update_quantize_metadata(model, config, metadata) @@ -139,7 +137,6 @@ 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 @@ -147,14 +144,13 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: return model quantizer_state_dict = dict(metadata["quantizer_state"]) - if "linear_attention" not in metadata: - # Older checkpoints predate these disabled handles; preserve their baseline path. - for name, module in _linear_attention_modules(model).items(): - for handle in ("gdn_state_quantizer", "gdn_w_quantizer"): - key = f"{name}.{handle}" if name else handle - quantizer = getattr(module, handle) - if key not in quantizer_state_dict and not quantizer.is_enabled: - quantizer_state_dict[key] = quantizer.get_modelopt_state() + # Older checkpoints predate these disabled handles; preserve their baseline path. + for name, module in _linear_attention_modules(model).items(): + for handle in ("gdn_state_quantizer", "gdn_w_quantizer"): + key = f"{name}.{handle}" if name else handle + quantizer = getattr(module, handle) + if key not in quantizer_state_dict and not quantizer.is_enabled: + quantizer_state_dict[key] = quantizer.get_modelopt_state() unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys() extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys() @@ -207,12 +203,6 @@ 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 @@ -231,30 +221,9 @@ def _linear_attention_modules(model): } -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(): +def _validate_linear_attention_quantizers(model): + for module in _linear_attention_modules(model).values(): module.validate_linear_attention() - if getattr(config, "linear_attention", []): - print_rank_0( - f"Linear attention {name!r}: {module.linear_attention_config.model_dump()}" - ) def quantizer_state(model: nn.Module) -> dict[str, Any]: diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py index 49f5ea3514a..477eeb89847 100644 --- a/modelopt/torch/quantization/linear_attention/__init__.py +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -13,7 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Numerical policies and differentiable references for linear attention.""" - -from .config import * -from .reference import * +"""Helpers for linear-attention quantization.""" diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py deleted file mode 100644 index 6405fbbf329..00000000000 --- a/modelopt/torch/quantization/linear_attention/config.py +++ /dev/null @@ -1,60 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Saved execution policy for GDN training-time numerical emulation.""" - -from typing import Literal - -from pydantic import Field - -from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField - -__all__ = ["LinearAttentionConfig", "LinearAttentionPolicyEntry"] - - -class _StateConfig(ModeloptBaseConfig): - mode: Literal["chunk"] = ModeloptField(default="chunk") - block_v: Literal[16, 32, 64, 128] = ModeloptField(default=64) - quantize_initial: Literal[True] = ModeloptField(default=True) - - -class _SolveConfig(ModeloptBaseConfig): - method: Literal["exact"] = ModeloptField(default="exact") - - -class LinearAttentionConfig(ModeloptBaseConfig): - """GDN chunk-64 policy; unsupported numerical modes fail config validation. - - ``state.block_v`` defines one dynamic scale per ``[Dk, block_v]`` tile of each - sequence/head. The initial state and every chunk's final state are rounded when - ``gdn_state_quantizer`` is enabled. Outputs use the incoming rounded state. - """ - - schema_version: Literal[1] = ModeloptField(default=1) - backend: Literal["fla"] = ModeloptField(default="fla") - chunk_size: Literal[64] = ModeloptField(default=64) - state: _StateConfig = ModeloptField(default=_StateConfig()) - solve: _SolveConfig = ModeloptField(default=_SolveConfig()) - - -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/validation.py b/modelopt/torch/quantization/linear_attention/utils.py similarity index 83% rename from modelopt/torch/quantization/linear_attention/validation.py rename to modelopt/torch/quantization/linear_attention/utils.py index 2bf07525e22..4f1e18a174d 100644 --- a/modelopt/torch/quantization/linear_attention/validation.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -13,22 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Capabilities of the initial fused GDN fake-quant path.""" +"""Shared helpers for linear-attention quantization.""" from ..nn import TensorQuantizer __all__ = [] -def validate_gdn_quantizer(quantizer: TensorQuantizer, *, state: bool) -> None: - name = "gdn_state_quantizer" if state else "gdn_w_quantizer" - axis = (0, 1) if state else (0, 1, 2) +def validate_gdn_quantizer(quantizer: TensorQuantizer, *, name: str) -> None: + """Check the dynamic E4M3 contract and the custom backward's identity STE.""" if not isinstance(quantizer, TensorQuantizer): raise ValueError(f"{name} requires a single TensorQuantizer") if not ( quantizer._dynamic and quantizer.num_bits == (4, 3) - and quantizer.axis == axis and quantizer.block_sizes is None and quantizer.fake_quant and quantizer._pass_through_bwd @@ -39,7 +37,7 @@ def validate_gdn_quantizer(quantizer: TensorQuantizer, *, state: bool) -> None: and not quantizer._use_constant_amax ): raise ValueError( - f"{name} supports only dynamic E4M3 fake quantization with axis={axis}, " + f"{name} supports only dynamic E4M3 fake quantization with " "pass_through_bwd=True, no block_sizes, rotation, pre-scaling, bias, constant " "amax, or custom backend. Other gradient rules and formats are not implemented." ) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 5aac82cacfe..fbf9bb299fd 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -32,7 +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, + _validate_linear_attention_quantizers, preserve_quantizer_attributes_context, set_quantizer_attributes_partial, set_quantizer_by_cfg, @@ -330,7 +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) + _validate_linear_attention_quantizers(model) # 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/gated_delta_net.py b/modelopt/torch/quantization/plugins/gated_delta_net.py index 7be046b8800..df70683cc2f 100644 --- a/modelopt/torch/quantization/plugins/gated_delta_net.py +++ b/modelopt/torch/quantization/plugins/gated_delta_net.py @@ -27,13 +27,13 @@ """ from collections.abc import Callable +from functools import partial from typing import Any import torch from ..config import QuantizerAttributeConfig -from ..linear_attention.config import LinearAttentionConfig -from ..linear_attention.validation import validate_gdn_quantizer +from ..linear_attention.utils import validate_gdn_quantizer from ..nn import QuantModule, TensorQuantizer __all__ = ["GatedDeltaNetStateQuantMixin"] @@ -41,6 +41,13 @@ 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. @@ -64,27 +71,29 @@ class GatedDeltaNetStateQuantMixin(QuantModule): with ``quant_cfg`` entries on ``*gdn_state_quantizer`` / ``*gdn_w_quantizer`` such as the ``configs/ptq/units/gdn_state_fp8_dynamic`` and ``gdn_w_fp8_dynamic`` recipe units. The state quantizer carries the fused QDQ configuration. Both sites currently require dynamic - E4M3 and identity STE. The execution policy is saved in ModelOpt metadata. + E4M3 and identity STE. State QDQ uses fixed 64-column value tiles and 64-token chunks. """ def _setup(self): - self.gdn_state_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) - self.gdn_w_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) - self.linear_attention_config = LinearAttentionConfig() - - @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 + for name in ("gdn_state_quantizer", "gdn_w_quantizer"): + self._register_temp_attribute( + name, TensorQuantizer(QuantizerAttributeConfig(enable=False)) + ) def validate_linear_attention(self) -> None: """Reject numerical settings that the fused training path cannot implement.""" - for state, quantizer in ((True, self.gdn_state_quantizer), (False, self.gdn_w_quantizer)): + for name in ("gdn_state_quantizer", "gdn_w_quantizer"): + quantizer = getattr(self, name) if quantizer.is_enabled: - validate_gdn_quantizer(quantizer, state=state) + validate_gdn_quantizer(quantizer, name=name) + # The state handle configures fused QDQ; W's grouping is executed by TensorQuantizer. + if self.gdn_state_quantizer.is_enabled and self.gdn_state_quantizer.axis != (0, 1): + raise ValueError( + "gdn_state_quantizer supports only axis=(0, 1) with 64-column value tiling" + ) def modelopt_post_restore(self, prefix: str = ""): - """Validate restored quantizers before using the saved execution policy.""" + """Validate restored quantizers against the fused training kernel requirements.""" super().modelopt_post_restore(prefix) self.validate_linear_attention() @@ -97,19 +106,23 @@ def _state_quantized_chunk_gated_delta_rule( quantize_w = self.gdn_w_quantizer.is_enabled if not (quantize_state or quantize_w): return gated_delta_rule(*args, **kwargs) - if getattr(gated_delta_rule, "__name__", None) != "chunk_gated_delta_rule": + 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 quantizers require the fla chunked kernel; the deterministic torch " - f"kernel ({gated_delta_rule!r}) is not supported." + "GatedDeltaNet quantization supports only FLA's chunk_gated_delta_rule " + f"callable or a functools.partial of it; got {gated_delta_rule!r}." ) - chunk_size = kwargs.pop("chunk_size", self.linear_attention_config.chunk_size) - if chunk_size != self.linear_attention_config.chunk_size: + chunk_size = kwargs.pop("chunk_size", 64) + if chunk_size != 64: raise ValueError("GDN fake quantization supports only chunk_size=64") return _state_qdq_chunk_gated_delta_rule()( *args, chunk_size=chunk_size, state_qdq=int(quantize_state), - state_qdq_block_v=self.gdn_state_qdq_block_v, + state_qdq_block_v=64, w_quantizer=self.gdn_w_quantizer if quantize_w else None, **kwargs, ) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 3a0d4627e38..ec942e0c24c 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -1083,9 +1083,9 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): class _QuantGatedDeltaNet(GatedDeltaNetStateQuantMixin): """GatedDeltaNet with fake quantization of the recurrent state at kernel chunk boundaries. - Training and calibration reach the chunked kernel through ``forward_pre_attn_and_core_attn``, - which calls ``self.gated_delta_rule``; that call is routed through ``gdn_state_quantizer``. - The dynamic-batching inference paths (``ssm_prefill`` / ``ssm_decode``) are left untouched. + 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): @@ -1099,16 +1099,34 @@ def _setup(self): mcore_parallel.get_tensor_model_parallel_group(), ) - def forward_pre_attn_and_core_attn(self, *args, **kwargs): + @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: - return super().forward_pre_attn_and_core_attn(*args, **kwargs) + 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.gdn_state_quantizer.is_enabled or self.gdn_w_quantizer.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_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index 4b6ccec395f..79bc6065c24 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -34,5 +34,3 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `attention_qkv_fp8.yaml` | FP8 E4M3 on attention q/k/v bmm and softmax quantizers | | `gdn_state_fp8_dynamic.yaml` | FP8 E4M3 dynamic fake quantization of the GatedDeltaNet recurrent state (per sequence, head and 64-column tile) at every kernel chunk boundary; needs `fla-core==0.5.1`, Triton, and SM89+ | | `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 | - -See the [linear-attention training guide](../../../../docs/linear_attention_qat.md) for GDN policy, hardware, checkpoint, and gradient requirements. diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml index 51f9ad1de20..114c86d5cad 100644 --- a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -16,9 +16,8 @@ # FP8 E4M3 fake quantization of the GatedDeltaNet (linear attention) recurrent state with a # dynamic scale computed inside the kernel from the state itself (no calibration). The state is # ``[N, H, K, V]`` for ``N`` sequences; ``axis: [0, 1]`` scales each sequence's head state on its -# own, in tiles of ``gdn_state_qdq_block_v`` value columns (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 -# is the only configuration the kernel implements. Applied at every kernel chunk boundary +# own, with one scale per full-K by 64-value-column tile (two scales per head for V = 128). +# This is the only configuration the module supports. Applied at every kernel chunk boundary # (64 tokens) during calibration, QAT and QAD; needs fla-core==0.5.1 and Triton. # See ``modelopt.torch.quantization.plugins.gated_delta_net``. diff --git a/noxfile.py b/noxfile.py index d061ef31f57..cd9e96ecffa 100644 --- a/noxfile.py +++ b/noxfile.py @@ -130,7 +130,15 @@ def gpu(session): "mamba_ssm", "causal-conv1d", ) - session.run("python", "-m", "pip", "install", "fla-core==0.5.1") + 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()) @@ -142,7 +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") + 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/modelopt/torch/quantization/linear_attention/reference.py b/tests/_test_utils/torch/quantization/linear_attention_reference.py similarity index 100% rename from modelopt/torch/quantization/linear_attention/reference.py rename to tests/_test_utils/torch/quantization/linear_attention_reference.py diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py index c78668febde..d470da9f491 100644 --- a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py @@ -15,16 +15,24 @@ """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.linear_attention import chunk_gdn_reference from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.plugins.gated_delta_net 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, ) @@ -35,10 +43,16 @@ def require_state_qdq(): pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") -def make_inputs(dtype=torch.float32, packed=False, state_v_first=False): +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) @@ -53,10 +67,8 @@ def make_inputs(dtype=torch.float32, packed=False, state_v_first=False): return args, state.requires_grad_() -def w_quantizer(): - return TensorQuantizer( - QuantizerAttributeConfig(num_bits=(4, 3), axis=(0, 1, 2), type="dynamic") - ) +def w_quantizer(axis=(0, 1, 2)): + return TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), axis=axis, type="dynamic")) def compare(actual, expected, tolerance): @@ -69,6 +81,14 @@ def compare(actual, expected, 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) @@ -77,9 +97,12 @@ def values_and_grads(fn, args, state, **kwargs): return result, grads +# The first comparison compiles and autotunes both upstream and ModelOpt forward/backward kernels. +@pytest.mark.timeout(300) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_disabled_matches_upstream_forward_and_backward(dtype): - args, state = make_inputs(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): @@ -111,7 +134,7 @@ def test_numerical_surrogate_reference(state_qdq, quantize_w, packed, state_v_fi actual = values_and_grads( chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs ) - compare(actual[0], expected[0], 0.03 if dtype == torch.bfloat16 else 0.01) + compare_forward(actual[0], expected[0], dtype, state_qdq) compare(actual[1], expected[1], 0.05 if dtype == torch.bfloat16 else 0.02) @@ -138,15 +161,16 @@ def fn(*x): handle.remove() -def test_zero_initial_state_and_output_only_training_loss(): +@pytest.mark.parametrize("axis", [None, (0, 1), (0, 1, 2)]) +def test_zero_initial_state_and_output_only_training_loss(axis): args, _ = make_inputs() - quantizer = w_quantizer() - expected, _ = chunk_gdn_reference(*args, w_quantizer=quantizer) + 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.01) + 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.02) + compare(*grads, 0.05 if args[0].dtype == torch.bfloat16 else 0.02) @pytest.mark.parametrize("block_v", [16, 64, 128]) @@ -154,12 +178,37 @@ def test_state_scale_tile_forward_and_backward(block_v): require_state_qdq() args, state = make_inputs() kwargs = {"state_qdq": True, "state_qdq_block_v": block_v} - expected = values_and_grads(chunk_gdn_reference, args, state, **kwargs) + reference_args = [x.detach().float().requires_grad_() for x in args] + expected = values_and_grads(chunk_gdn_reference, reference_args, state, **kwargs) actual = values_and_grads( chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs ) - compare(actual[0], expected[0], 0.01) - compare(actual[1], expected[1], 0.02) + 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(): @@ -169,11 +218,11 @@ def test_fused_gate_normalization_and_beta_gradients(): bias = torch.randn(2, device="cuda", requires_grad=True) quantizer = w_quantizer() expected = chunk_gdn_reference( - F.normalize(q, dim=-1), - F.normalize(k, dim=-1), - v, - -a_log.exp() * F.softplus(raw_g + bias), - 2 * raw_beta.sigmoid(), + 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, ) @@ -189,12 +238,40 @@ def test_fused_gate_normalization_and_beta_gradients(): use_beta_sigmoid_in_kernel=True, allow_neg_eigval=True, ) - compare(actual, expected, 0.01) + 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.02) + 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]) @@ -210,3 +287,11 @@ def test_unsupported_w_gradient_rejected_before_launch(): 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_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index 2252eb25a61..1fe3e7399b8 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -16,12 +16,16 @@ 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 +from _test_utils.torch.megatron.utils import ( + get_forward, + initialize_for_megatron, + sharded_state_dict_test_helper, +) import modelopt.torch.quantization as mtq pytest.importorskip("fla") # Megatron-Core GatedDeltaNet and the state QDQ kernel need fla -pytest.importorskip("megatron.core.ssm.gated_delta_net") +GatedDeltaNet = pytest.importorskip("megatron.core.ssm.gated_delta_net").GatedDeltaNet from modelopt.torch.quantization.plugins.gated_delta_net import _state_qdq_chunk_gated_delta_rule from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet @@ -32,70 +36,133 @@ pytest.skip(str(e), allow_module_level=True) SEED = 1234 -GDN_STATE_QUANT_CFG = { - "quant_cfg": [ - {"quantizer_name": "*", "enable": False}, - { - "quantizer_name": "*gdn_state_quantizer", - "cfg": {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"}, - }, - ], - "algorithm": "max", -} - - -def _test_gdn_state_quant_helper(rank, size): + + +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, + } + + +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 = get_mcore_gpt_model( - tensor_model_parallel_size=size, - num_layers=2, - hidden_size=64, - num_attention_heads=4, - vocab_size=32, - max_sequence_length=128, - experimental_attention_variant="gated_delta_net", - ).cuda() - model.eval() # no dropout, so the forwards differ only through the state quantizer + model = _make_model(size) forward = get_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() - model = mtq.quantize(model, GDN_STATE_QUANT_CFG, forward) - + sites = ("state", "w") if mode == "both" else (mode,) + mtq.quantize(model, _gdn_config(sites)) gdn_modules = [m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)] assert gdn_modules, "no GatedDeltaNet layer was wrapped" - assert all(m.gdn_state_quantizer.is_enabled for m in gdn_modules) - assert all(m.gdn_state_qdq_block_v == 64 for m in gdn_modules) # kernel's 64-column tile + for module in gdn_modules: + assert module.gdn_state_quantizer.is_enabled == ("state" in sites) + assert module.gdn_w_quantizer.is_enabled == ("w" in sites) with torch.no_grad(): loss_quant = forward(model) assert torch.isfinite(loss_quant).all() - assert not torch.allclose(loss_quant, loss_ref, rtol=1e-4, atol=1e-4), ( - "FP8 state quantization must change the output" + # 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() - mtq.disable_quantizer(model, "*gdn_state_quantizer") + for site in sites: + mtq.disable_quantizer(model, f"*gdn_{site}_quantizer") with torch.no_grad(): - assert torch.allclose(forward(model), loss_ref, rtol=1e-4, atol=1e-4) - mtq.enable_quantizer(model, "*gdn_state_quantizer") - - # QAD trains through the segmented kernel: gradients must reach the GDN projections. - model.train() - optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) - before = gdn_modules[0].in_proj.weight.detach().clone() - optimizer.zero_grad() - forward(model).sum().backward() - assert all(m.in_proj.weight.grad is not None for m in gdn_modules) - assert all(torch.isfinite(m.in_proj.weight.grad).all() for m in gdn_modules) + torch.testing.assert_close(forward(model), loss_ref, rtol=1e-4, atol=1e-4) + for site in sites: + mtq.enable_quantizer(model, f"*gdn_{site}_quantizer") + + restored = _make_model(size) + sharded_state_dict_test_helper(checkpoint_path, model, restored, forward) + restored_gdn = [m for m in restored.modules() if isinstance(m, _QuantGatedDeltaNet)] + assert len(restored_gdn) == len(gdn_modules) + for module in restored_gdn: + assert module.gdn_state_quantizer.is_enabled == ("state" in sites) + assert module.gdn_w_quantizer.is_enabled == ("w" in sites) + 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(gdn_modules[0].in_proj.weight, before) + assert not torch.equal(restored_gdn[0].in_proj.weight, before) -def test_gdn_state_quant(dist_workers_size_1): - """GatedDeltaNet layers get a ``gdn_state_quantizer`` whose FP8 quant-dequant of the recurrent - state runs inside the chunked kernel during forward and backward.""" - if torch.cuda.get_device_capability() < (8, 9): +# Cold FLA/TileLang compilation plus a two-rank checkpoint exceeds the lane's 120s default. +@pytest.mark.timeout(300) +@pytest.mark.parametrize("tp_size", [1, 2]) +@pytest.mark.parametrize("mode", ["state", "w", "both"]) +def test_gdn_qat_and_sharded_restore(request, tmp_path, tp_size, mode): + """Train through state/W QDQ after a Megatron distributed-checkpoint round trip.""" + if mode != "w" and torch.cuda.get_device_capability() < (8, 9): pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") - dist_workers_size_1.run(_test_gdn_state_quant_helper) + workers = request.getfixturevalue(f"dist_workers_size_{tp_size}") + 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, _gdn_config((mode,))) + + +@pytest.mark.parametrize("mode", ["state", "w"]) +def test_gdn_context_parallel_rejected(dist_workers_size_2, mode): + """Reject unqualified Megatron CP even when it does not pass an FLA CP context.""" + dist_workers_size_2.run(_test_gdn_context_parallel_helper, mode) diff --git a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py b/tests/unit/torch/quantization/plugins/test_gated_delta_net.py index 2bdfa2e7b09..3793edcb189 100644 --- a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py +++ b/tests/unit/torch/quantization/plugins/test_gated_delta_net.py @@ -14,6 +14,7 @@ # limitations under the License. from copy import deepcopy +from functools import partial import pytest import torch @@ -21,7 +22,6 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq -from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins import gated_delta_net from modelopt.torch.quantization.plugins.gated_delta_net import GatedDeltaNetStateQuantMixin @@ -30,10 +30,17 @@ def chunk_gated_delta_rule(q, k, v, g, beta, **kwargs): - """Stand-in with the fla kernel's name; the real kernel needs a GPU.""" + """CPU stand-in for the optional FLA kernel.""" return q + k + v, None +@pytest.fixture(autouse=True) +def mock_fla_kernel(monkeypatch): + monkeypatch.setattr( + gated_delta_net, "_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``.""" @@ -95,6 +102,14 @@ def test_validate_state_quantizer_rejects_unsupported(attributes): ) +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"): + assert not hasattr(model, name) + + def test_disabled_state_quantizer_calls_original_kernel(): model = TinyGatedDeltaNet() x = torch.randn(2, 8, 3, 4) @@ -105,12 +120,12 @@ def test_disabled_state_quantizer_calls_original_kernel(): assert isinstance(model, _QuantTinyGatedDeltaNet) assert not model.gdn_state_quantizer.is_enabled and not model.gdn_w_quantizer.is_enabled - assert model.gdn_state_qdq_block_v == 64 assert torch.equal(model(x), expected) assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" -def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch): +@pytest.mark.parametrize("partial_kernel", [False, True]) +def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, partial_kernel): calls = [] def fake_state_qdq_kernel(*args, **kwargs): @@ -121,6 +136,8 @@ def fake_state_qdq_kernel(*args, **kwargs): gated_delta_net, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel ) model = TinyGatedDeltaNet() + if partial_kernel: + model.gated_delta_rule = partial(chunk_gated_delta_rule, output_final_state=True) x = torch.randn(2, 8, 3, 4) mtq.quantize(model, quant_cfg(), lambda m: m(x)) @@ -130,11 +147,13 @@ def fake_state_qdq_kernel(*args, **kwargs): "state_qdq": 1, "state_qdq_block_v": 64, "w_quantizer": None, + **({"output_final_state": True} if partial_kernel else {}), } - # The deterministic torch kernel has no quantized counterpart. + # 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) - with pytest.raises(NotImplementedError, match="deterministic torch kernel"): + model.gated_delta_rule.__name__ = "chunk_gated_delta_rule" + with pytest.raises(NotImplementedError, match="supports only FLA"): model(x) @@ -181,72 +200,55 @@ def test_unsupported_quantizer_fails_during_conversion(site, overrides): mtq.quantize(TinyGatedDeltaNet(), cfg) -def test_policy_roundtrip_last_match_and_hybrid_selection(tmp_path): +@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() + 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 - # A post-conversion policy edit must also survive save/load. - 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_quantizer.is_enabled - assert restored[0].gdn_state_qdq_block_v == 16 - assert not hasattr(restored[1], "linear_attention_config") - + 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_policy_rejects_unmatched_and_unimplemented_modes(): - with pytest.raises(ValueError, match="matches no supported"): - mtq.quantize( - nn.Linear(4, 4), {"linear_attention": [{"module_name": "*"}], "algorithm": None} - ) - for policy in ( - {"chunk_size": 32}, - {"solve": {"method": "neumann"}}, - {"state": {"mode": "token"}}, - ): - with pytest.raises(ValueError): - QuantizeConfig(linear_attention=[{"module_name": "*", "cfg": policy}]) - -def test_restore_old_config_without_policy_field(): - model = mtq.quantize( - nn.Linear(4, 4), - {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None}, - ) - state = mto.modelopt_state(model) - for _, mode_state in state["modelopt_state_dict"]: - config = mode_state["config"] - if isinstance(config, dict): - config.pop("linear_attention", None) - else: - config.__dict__.pop("linear_attention", None) - restored = nn.Linear(4, 4) - mto.restore_from_modelopt_state(restored, state) - restored.load_state_dict(model.state_dict()) - x = torch.randn(2, 4) - torch.testing.assert_close(restored(x), model(x)) - - -def test_policy_refinement_updates_existing_quantized_module(): +def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): cfg = quant_cfg() cfg["algorithm"] = None model = mtq.quantize(TinyGatedDeltaNet(), cfg) - cfg["linear_attention"] = [{"module_name": "", "cfg": {"state": {"block_v": 32}}}] - mtq.quantize(model, cfg) - assert model.gdn_state_qdq_block_v == 32 - # A later complete rule resets omitted fields instead of merging the previous policy. - cfg["linear_attention"].append({"module_name": "", "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 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(): @@ -254,9 +256,7 @@ def test_restore_legacy_gdn_without_new_quantizer_handles(): model = mtq.quantize(TinyGatedDeltaNet(), config) state = mto.modelopt_state(model) for _, mode_state in state["modelopt_state_dict"]: - mode_state["config"].pop("linear_attention", None) metadata = mode_state["metadata"] - metadata.pop("linear_attention", None) for name in ("gdn_state_quantizer", "gdn_w_quantizer"): metadata["quantizer_state"].pop(name, None) restored = TinyGatedDeltaNet() diff --git a/tests/unit/torch/quantization/test_linear_attention_reference.py b/tests/unit/torch/quantization/test_linear_attention_reference.py index dd2aea8d592..d71818d0163 100644 --- a/tests/unit/torch/quantization/test_linear_attention_reference.py +++ b/tests/unit/torch/quantization/test_linear_attention_reference.py @@ -16,8 +16,7 @@ import pytest import torch import torch.nn.functional as F - -from modelopt.torch.quantization.linear_attention import ( +from _test_utils.torch.quantization.linear_attention_reference import ( chunk_gdn_reference, recurrent_delta_rule_reference, state_fp8_qdq_reference, From 95766de709ac05ed733be4c01066a7e832533826 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 28 Sep 2026 17:45:36 -0700 Subject: [PATCH 3/8] Move GDN test compilation into setup fixtures Warm standalone and Megatron forward/backward kernels before functional tests. Remove the 300-second overrides so test calls retain the default 120-second limit and report execution separately from compilation. Signed-off-by: Kai Xu --- .../test_fla_chunk_gated_delta_rule.py | 96 +++++++++++++++++-- .../plugins/test_megatron_gated_delta_net.py | 36 +++++-- 2 files changed, 116 insertions(+), 16 deletions(-) diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py index d470da9f491..9a11c454932 100644 --- a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py @@ -37,6 +37,11 @@ chunk_gated_delta_rule, ) +DTYPES = (torch.float32, torch.bfloat16) +QDQ_MODES = ((False, False), (False, True), (True, False), (True, 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): @@ -97,9 +102,84 @@ def values_and_grads(fn, args, state, **kwargs): return result, grads -# The first comparison compiles and autotunes both upstream and ModelOpt forward/backward kernels. -@pytest.mark.timeout(300) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@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 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() + values_and_grads( + chunk_gated_delta_rule, args, state, output_final_state=True, w_quantizer=quantizer + ) + 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)) + if supports_state_qdq: + for block_v in (16, 64, 128): + values_and_grads( + chunk_gated_delta_rule, + args, + state, + output_final_state=True, + state_qdq=True, + state_qdq_block_v=block_v, + ) + 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) @@ -109,11 +189,9 @@ def test_disabled_matches_upstream_forward_and_backward(dtype): compare(a, e, 0.005 if dtype == torch.bfloat16 else 0.001) -@pytest.mark.parametrize( - ("state_qdq", "quantize_w"), [(False, False), (False, True), (True, False), (True, True)] -) -@pytest.mark.parametrize(("packed", "state_v_first"), [(False, False), (True, True)]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@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: require_state_qdq() @@ -161,7 +239,7 @@ def fn(*x): handle.remove() -@pytest.mark.parametrize("axis", [None, (0, 1), (0, 1, 2)]) +@pytest.mark.parametrize("axis", W_AXES) def test_zero_initial_state_and_output_only_training_loss(axis): args, _ = make_inputs() quantizer = w_quantizer(axis) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index 1fe3e7399b8..c46959c4d8d 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -143,16 +143,38 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): assert not torch.equal(restored_gdn[0].in_proj.weight, before) -# Cold FLA/TileLang compilation plus a two-rank checkpoint exceeds the lane's 120s default. -@pytest.mark.timeout(300) -@pytest.mark.parametrize("tp_size", [1, 2]) -@pytest.mark.parametrize("mode", ["state", "w", "both"]) -def test_gdn_qat_and_sharded_restore(request, tmp_path, tp_size, mode): - """Train through state/W QDQ after a Megatron distributed-checkpoint round trip.""" +def _compile_gdn_qat_kernels(rank, size, mode): + initialize_for_megatron( + tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED + ) + model = _make_model(size) + forward = get_forward(model) + with torch.no_grad(): + forward(model) + sites = ("state", "w") if mode == "both" else (mode,) + mtq.quantize(model, _gdn_config(sites)) + 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 != "w" 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(_test_gdn_qat_helper, mode, tmp_path) + workers.run(_compile_gdn_qat_kernels, mode) + return workers + + +@pytest.mark.parametrize("tp_size", [1, 2]) +@pytest.mark.parametrize("mode", ["state", "w", "both"]) +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): From 3896181aef1cecc6c85014b5b3ba37f1fdf0fb0b Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 28 Sep 2026 18:13:31 -0700 Subject: [PATCH 4/8] Clarify dynamic per-tile FP8 state quantization Signed-off-by: Kai Xu --- .../quantization/linear_attention/fla_chunk_delta_h.py | 7 +++---- .../configs/ptq/units/gdn_state_fp8_dynamic.yaml | 8 ++++---- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py index ca8b085470d..875cb39a7d0 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -50,7 +50,7 @@ @triton.jit def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr): - """[ModelOpt] Dynamic FP8 E4M3 scale over the up-to-four K tiles of one program's state.""" + """[ModelOpt] Dynamic FP8 E4M3 scale per full [K, BV] tile of one sequence and head.""" b_amax = tl.max(tl.abs(b_h1)) if K > 64: b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h2))) @@ -385,9 +385,8 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h4 = tl.dot(b_k, b_v, b_h4) - # [ModelOpt] Fake-quantize the state carried into the next chunk (and, after the last - # chunk, the stored final state) to FP8 E4M3. The scale is dynamic over this program's - # [K, BV] tile of the head state; BV == V makes it one scale per sequence and head. + # [ModelOpt] Dynamic per-tile FP8 QDQ of the next-chunk or stored final state. + # Recompute amax over [K, BV]; fp8_scalar_qdq applies that tile's scalar scale. if STATE_QDQ == 1: if K > 192: b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml index 114c86d5cad..ed5768ee599 100644 --- a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -13,10 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -# FP8 E4M3 fake quantization of the GatedDeltaNet (linear attention) recurrent state with a -# dynamic scale computed inside the kernel from the state itself (no calibration). The state is -# ``[N, H, K, V]`` for ``N`` sequences; ``axis: [0, 1]`` scales each sequence's head state on its -# own, with one scale per full-K by 64-value-column tile (two scales per head for V = 128). +# Dynamic per-tile FP8 E4M3 fake quantization of the GatedDeltaNet recurrent state. +# Recompute amax over each [K, 64] tile inside the kernel, independently per sequence and head +# of the [N, H, K, V] state (no calibration). ``axis: [0, 1]`` specifies sequence/head grouping; +# the kernel uses one scale per full-K by 64-value-column tile (two per head for V = 128). # This is the only configuration the module supports. Applied at every kernel chunk boundary # (64 tokens) during calibration, QAT and QAD; needs fla-core==0.5.1 and Triton. # See ``modelopt.torch.quantization.plugins.gated_delta_net``. From d03c73254264cc85943ee7f81e322898e11ae1a1 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sun, 27 Sep 2026 20:55:54 -0700 Subject: [PATCH 5/8] Add GDN/KDA Torch decode QAT with INT8 recurrent state Signed-off-by: Kai Xu --- CHANGELOG.rst | 2 + examples/llm_qat/README.md | 1 + examples/llm_qat/linear_attention/README.md | 304 ++++++++++++++++ .../configs/kda_decode_state_int8.json | 35 ++ .../llm_qat/linear_attention/requirements.txt | 4 + examples/llm_qat/linear_attention/train.py | 109 ++++++ .../linear_attention/fla_chunk_delta_h.py | 68 ++-- .../fla_chunk_gated_delta_rule.py | 12 +- .../quantization/linear_attention/int8.py | 28 ++ modelopt/torch/quantization/config.py | 7 + modelopt/torch/quantization/conversion.py | 55 ++- .../quantization/linear_attention/__init__.py | 8 +- .../linear_attention/_chunk_prefill.py | 146 ++++++++ .../quantization/linear_attention/config.py | 108 ++++++ .../quantization/linear_attention/decode.py | 341 ++++++++++++++++++ .../linear_attention/decode_prefill.py | 176 +++++++++ .../quantization/linear_attention/kda.py | 107 ++++++ .../quantization/linear_attention/prefill.py | 94 +++++ .../quantization/linear_attention/utils.py | 109 +++++- modelopt/torch/quantization/model_quant.py | 4 +- .../torch/quantization/plugins/__init__.py | 3 +- .../plugins/{gated_delta_net.py => gdn.py} | 65 ++-- modelopt/torch/quantization/plugins/kda.py | 103 ++++++ .../quantization/plugins/linear_attention.py | 85 +++++ .../torch/quantization/plugins/megatron.py | 6 +- modelopt/torch/quantization/tensor_quant.py | 3 +- modelopt_recipes/configs/ptq/units/README.md | 1 + .../units/default_disabled_quantizers.yaml | 4 + .../ptq/units/gdn_state_fp8_dynamic.yaml | 7 +- .../ptq/units/gdn_state_int8_dynamic.yaml | 26 ++ .../general/ptq/gdn_state_int8_dynamic.yaml | 38 ++ modelopt_recipes/ptq.md | 3 +- .../linear_attention_reference.py | 183 +++++++++- .../test_fla_chunk_gated_delta_rule.py | 60 ++- .../torch/quantization/plugins/test_kda.py | 189 ++++++++++ .../quantization/test_tensor_quant_cuda.py | 22 +- .../plugins/test_megatron_gated_delta_net.py | 118 +++++- .../{test_gated_delta_net.py => test_gdn.py} | 92 ++++- .../test_linear_attention_decode.py | 296 +++++++++++++++ .../test_linear_attention_hadamard.py | 227 ++++++++++++ .../test_linear_attention_int8.py | 99 +++++ 41 files changed, 3184 insertions(+), 164 deletions(-) create mode 100644 examples/llm_qat/linear_attention/README.md create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json create mode 100644 examples/llm_qat/linear_attention/requirements.txt create mode 100644 examples/llm_qat/linear_attention/train.py create mode 100644 modelopt/torch/kernels/quantization/linear_attention/int8.py create mode 100644 modelopt/torch/quantization/linear_attention/_chunk_prefill.py create mode 100644 modelopt/torch/quantization/linear_attention/config.py create mode 100644 modelopt/torch/quantization/linear_attention/decode.py create mode 100644 modelopt/torch/quantization/linear_attention/decode_prefill.py create mode 100644 modelopt/torch/quantization/linear_attention/kda.py create mode 100644 modelopt/torch/quantization/linear_attention/prefill.py rename modelopt/torch/quantization/plugins/{gated_delta_net.py => gdn.py} (70%) create mode 100644 modelopt/torch/quantization/plugins/kda.py create mode 100644 modelopt/torch/quantization/plugins/linear_attention.py create mode 100644 modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml create mode 100644 modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml create mode 100644 tests/gpu/torch/quantization/plugins/test_kda.py rename tests/unit/torch/quantization/plugins/{test_gated_delta_net.py => test_gdn.py} (74%) create mode 100644 tests/unit/torch/quantization/test_linear_attention_decode.py create mode 100644 tests/unit/torch/quantization/test_linear_attention_hadamard.py create mode 100644 tests/unit/torch/quantization/test_linear_attention_int8.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 347806df32b..60d905778b0 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,6 +13,8 @@ Changelog *Quantization* +- Add experimental GDN/KDA decode-aware QAT with FP8 or INT8 recurrent states, optional INT8 value-axis Hadamard encoding, KDA decay rounding, and encoded-update replay. Supply explicit prefix lengths through the training phase context; prefix solves remain exact. + - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration. - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md `_ for details. 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..7c44058512a --- /dev/null +++ b/examples/llm_qat/linear_attention/README.md @@ -0,0 +1,304 @@ +# Quantization-Aware Training for Linear Attention + +This example fine-tunes KDA attention parameters with recurrent-state fake +quantization, then saves a Hugging Face checkpoint with ModelOpt state. GDN/KDA +runtime support includes token writes, ReplaySSM, KDA decay approximation, and +FP8 or INT8 state QDQ. The INT8 recipes enable Hadamard rotation by default. + +## Run the example + +Install ModelOpt using the [QAT setup instructions](../README.md#quick-start), +then install the example dependencies. Run commands from the repository root. + +```bash +pip install -r examples/llm_qat/linear_attention/requirements.txt +python examples/llm_qat/linear_attention/train.py \ + --model /path/to/local-kda-model \ + --train-data /path/to/train.parquet \ + --output /path/to/qat-checkpoint \ + --train-steps 1 --length 128 --prefill-tokens 64 +``` + +Use a local model/tokenizer snapshot containing FLA `KimiDeltaAttention` layers, +a Parquet file with a `text` column, and a CUDA GPU. The example requires +`fla-core==0.5.1` and `flash-linear-attention==0.5.1`. If the local model requires +custom Python code, review it and explicitly pass `--trust-remote-code`. + +The default [INT8 configuration](configs/kda_decode_state_int8.json) leaves the +64-token prefix state unquantized and applies INT8 QDQ in a 32-value Hadamard +basis during the decode suffix. Value dimensions must be divisible by 32. +Loss uses suffix labels. Only KDA attention parameters are trained, in FP32 under +BF16 autocast; other parameters are frozen in BF16. The output contains model +weights, tokenizer files, and ModelOpt quantizer/policy state. Call +`mto.enable_huggingface_checkpointing()` before reloading it with +`AutoModelForCausalLM.from_pretrained` to restore those quantizers and policies. +It is a floating fake-quantized training checkpoint, not a compressed serving model. +This short example does not measure model-quality recovery or performance. + +## Enable state quantization + +### What the execution policy controls + +An **execution policy** is the layer's `LinearAttentionConfig`: the settings that +determine how the recurrence runs and when it applies quantization. Training +uses three separate inputs: + +| Input | Purpose | Supplied through | +| --- | --- | --- | +| `TensorQuantizer` settings | Enable a quantization site and choose its numerical format, scales, and gradient behavior | Recipe `quant_cfg`, such as dynamic INT8 with identity STE | +| Execution policy | Choose token or replay mode, Hadamard rotation, handoff quantization, and the execution backend | Recipe `linear_attention` entries | +| Batch metadata | Specify where each sequence switches from prefill to decode | `linear_attention_training_phase(model, prefill_lengths)` | + +The same INT8 quantizer settings can round the state after every token or round +a replay anchor every eight tokens. Those schedules produce different recurrent +states, so the policy must specify which computation training should emulate. + +`LinearAttentionConfig` holds the overall backend, chunk size, and state settings. +Its optional `decode` field contains a `LinearAttentionDecodeConfig` for token or +replay mode, state codec, decay approximation, and replay settings. The decode +implementation uses Torch operations and autograd. This is one nested configuration: supplying a `decode` +dictionary in the recipe constructs the nested config automatically. + +Each `linear_attention` entry uses `module_name` to select attention layers and +`cfg` to specify their policy. During `mtq.quantize`, ModelOpt stores that policy +as each matched layer's `linear_attention_config`. Quantizer settings and the +policy persist through ModelOpt save/restore; batch-specific prefill lengths +must be supplied again for each workload. + +### Load a state recipe + +State quantizers start disabled. The state recipe enables `*kda_state_quantizer` +with signed narrow-range INT8, dynamic scales, and `axis=(0, 1)`. Use +`*gdn_state_quantizer` for GDN. For E4M3, set the quantizer config to +`{"num_bits": [4, 3], "type": "dynamic", "axis": [0, 1]}` and set +`decode.state_codec="tile"`. +Setting `prefill_state_qdq=True` alone does not enable a quantizer. + +Start with the token-state recipe and select one of the schedules below **before** +calling `mtq.quantize`: + +```python +import json +from pathlib import Path + +recipe = json.loads( + Path("examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json").read_text() +) +policy = recipe["linear_attention"][0]["cfg"] +decode = policy["decode"] +``` + +For GDN, the complete recipe composes the INT8 quantizer unit with the Hadamard +execution policy: + +```python +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe + +gdn_recipe = load_recipe("general/ptq/gdn_state_int8_dynamic").quantize +mtq.quantize(gdn_model, gdn_recipe) +``` + +Use the phase context shown below for each forward/backward. Importing only +`configs/ptq/units/gdn_state_int8_dynamic` configures the quantizer without +selecting Hadamard. Both complete recipes use the Torch decode implementation. + +### Why training needs prefill/decode boundaries + +One training forward can simulate prompt prefill followed by recurrent decode. +The INT8 + Hadamard recipes apply different state quantization schedules to those +phases, so the workload must specify where each sequence switches to decode. +The training batch supplies all tokens; this simulates decode arithmetic without +running a text-generation loop. Total sequence length alone does not identify +the prompt portion. + +For a 128-token sequence with `prefill_lengths=[96]`, the default token-state +recipe runs these steps: + +1. **Prefill:** process tokens 0–95 with `prefill_state_qdq=False`, leaving state + unquantized during the prefix. +2. **Handoff:** carry the resulting state into decode, rotate its value dimension + into the Hadamard basis, and apply INT8 QDQ because `quantize_initial=True`. +3. **Decode:** process tokens 96–127 recurrently, applying INT8 QDQ after every + state update. Later tokens consume the rounded state; outputs are transformed + back to the original value basis. + +The boundary is independent of `chunk_size=64`: this 96-token prefix contains a +full 64-token chunk and a partial 32-token chunk. The phase switches after token +95, not after each chunk. + +Two 128-token sequences can use `prefill_lengths=[64, 96]` in the same batch: +their decode suffixes then contain 64 and 32 tokens, respectively. The next batch +can have different lengths while reusing the same quantization recipe. This is +why the execution policy is saved in `linear_attention_config`, while the lengths +are supplied per batch through `linear_attention_training_phase`. They are not +saved as part of the model's quantization policy. The context selects numerical +phases; the caller still supplies training labels and any loss masking. + +The combined prefill/decode path requires explicit lengths, including `[0]` for +decode only or `[T]` for an all-prefix sequence of length `T`. The GDN chunk-only +FLA path described below applies its chunk schedule throughout and needs no +phase context. See the tables below for the corresponding quantization settings. + +### Choose the prefill and decode boundaries + +The table assumes the state quantizer is enabled. The INT8 recipes default to +`"int8_hadamard32"`; prefix state QDQ requires explicitly selecting `"tile"`. +Prefix lengths are supplied separately through `linear_attention_training_phase`. + +| Desired state QDQ | `decode.state_codec` | `decode.mode` | `decode.prefill_state_qdq` | Where rounding occurs | +| --- | --- | --- | --- | --- | +| Token decode only (default) | `"int8_hadamard32"` | `"token"` | `False` | At the first nonempty decode handoff and after every suffix token. | +| Prefill and token decode | `"tile"` | `"token"` | `True` | At prefix initialization, each prefix chunk write, decode handoff, and every suffix token. | +| Replay anchors only | `"int8_hadamard32"` | `"replay"` | `False` | At decode handoff and each replay-window refresh. | +| Prefill and replay anchors | `"tile"` | `"replay"` | `True` | At prefix initialization and chunk writes, then decode handoff and replay-window refreshes. | + +To enable **prefill and token decode**: + +```python +decode.update(mode="token", replay=None, state_codec="tile", prefill_state_qdq=True) +``` + +For **token decode only**, keep the supplied INT8 recipe's defaults: +`state_codec="int8_hadamard32"` and `prefill_state_qdq=False`. +Prefix state remains unquantized until it enters the decode path. + +To enable **ReplaySSM anchor quantization** with an eight-token window: + +```python +decode.update( + mode="replay", + state_codec="int8_hadamard32", + prefill_state_qdq=False, + replay={"window": 8, "factor_qdq": False, "encoding": "once"}, +) +``` + +Set `state_codec="tile"` and `prefill_state_qdq=True` to add prefix state QDQ +to this replay configuration, using unrotated INT8 for both phases. +`factor_qdq=False` above isolates state/anchor quantization. Set it to `True` to +also quantize buffered keys and updates to FP8. + +`decode.quantize_initial=True` is the default: it quantizes the incoming state +once at the first nonempty decode handoff. Set it to `False` to skip that initial +rounding while keeping later token writes or anchor refreshes quantized. This +setting does not disable prefix state QDQ. With both phases enabled, prefix-final +rounding and decode-handoff rounding are separate configured events. + +`chunk_size=64` counts **tokens per prefill chunk**. `state.block_v=64` counts +**value channels per execution tile**. The tile codec shares a scale across all +key channels and this value tile; Hadamard uses one scale per key channel and +32 values. A replay `window=8` counts **suffix tokens between anchor refreshes**. + +### Run only the desired phase + +For a batch containing one sequence of `T` tokens, choose the context lengths as +follows; for larger batches, provide one length per sequence: + +| Workload | Context argument | Required setting | +| --- | --- | --- | +| Prefill followed by decode | `[64]`, with `T > 64` | Select either prefix setting above. | +| Decode only | `[0]` | State quantizer enabled; token or replay policy. | +| Prefill only | `[T]` | `state_codec="tile"`, `prefill_state_qdq=True`; the decode suffix is empty. | + +An empty suffix creates no decode quantization event. The combined interface has +one state quantizer per layer: it does not offer a switch to quantize the prefix +while leaving a **nonempty** decode suffix unquantized. Disable the state quantizer +to disable both state and anchor QDQ; replay factor QDQ has its own toggle. + +The `int8_hadamard32` codec applies to decode token writes or replay anchors only. +It requires `prefill_state_qdq=False`; use the tile codec to quantize prefix state. + +Save a modified recipe as JSON and pass it to `train.py --quant-config`. `--prefill-tokens` sets the phase split; it does not enable state +quantization. The integration loop below applies the configured `recipe` directly. + +### GDN chunk-only FLA training + +For GDN's existing chunked FLA path, enable the GDN state quantizer and omit the +`linear_attention` execution-policy entry: + +```python +gdn_recipe = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*gdn_state_quantizer", + "cfg": recipe["quant_cfg"][1]["cfg"], + }, + ], + "algorithm": None, +} +# Apply mtq.quantize(gdn_model, gdn_recipe), then use normal forward/backward. +``` + +This explicitly selects unrotated tile QDQ. It rounds the initial state and each +64-token chunk's final state, including a partial final chunk. It needs no phase +context and leaves W/projection quantizers +disabled. KDA uses the materialized backend and can use the all-prefix context +shown above. + +## Integrate with a training loop + +Apply the configured `recipe` above with `mtq.quantize`, then supply one prefix length per sequence. +Keep the phase context active through backward so activation-checkpoint +recomputation uses the same prefix/decode split. + +```python +import torch + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase + +mtq.quantize(model, recipe) +optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5) +model.train() +optimizer.zero_grad(set_to_none=True) + +# Two 128-token sequences: 64 and 96 prefill tokens, respectively. +# ids and labels have shape [2, 128] and are on the model's device. +with linear_attention_training_phase(model, [64, 96]): + with torch.autocast("cuda", dtype=torch.bfloat16): + loss = model(input_ids=ids, labels=labels, use_cache=False).loss + loss.backward() +optimizer.step() +``` + +For GDN, select `*gdn_state_quantizer` in the recipe instead of +`*kda_state_quantizer`. State quantization is dynamic, so these recipes use +`algorithm=None` without a calibration pass. Replay factors have their own FP8 +toggle. + +Policies persist through ModelOpt save/restore. Per-batch prefix lengths are +runtime metadata and must be supplied for each workload. The context restores +previous lengths on exit and supports nesting. Concurrent forwards on the same +model instance with different phase contexts are unsupported. + +## Replay, decay, and Hadamard options + +Token mode rounds each recurrent-state write. Replay mode retains an anchor and +ordered key/update factors, rounding the anchor at each `replay.window` refresh. +`factor_qdq` independently enables FP8 QDQ for those factors. `readout="working"` +reads the state before its write quantization; `"stored"` reads it afterward. + +For KDA decay approximation, set `decode["decay_log_step"] = 1 / 256` before +conversion. It rounds suffix log retention before exponentiation with identity +STE gradients. Prefix decay remains exact. + +The default INT8 codec applies a 32-point orthonormal Hadamard transform along +the value axis, quantizes token states or replay anchors in that basis, and +transforms outputs back. Value dimensions must be divisible by 32; `state.block_v` must be 32, 64, +or 128. Scales group one key channel and 32 values, independent of execution tile +width. INT8 codes use half-away-from-zero rounding with FP16 stored scales; +the optional tile codec instead uses nearest-even rounding and FP32 scales. + +## Training boundaries + +The Torch implementation uses autograd through initial states, chunk handoff, +token writes, and replay refreshes. The decode recurrence runs token by token in +Python, so long training suffixes can be slow. + +Prefill prefixes use exact chunk algebra with optional state QDQ. This example +has no prefill GEMM QDQ or approximate inverse. FLA KDA requires `use_cache=False`; +serving cache objects are rejected. ModelOpt saves execution policies, while +per-batch prefix lengths must be supplied again during training. Distributed +decode training and model-quality recovery require separate qualification. diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json b/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json new file mode 100644 index 00000000000..8a29cda8139 --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json @@ -0,0 +1,35 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": 8, + "type": "dynamic", + "axis": [ + 0, + 1 + ], + "unsigned": false, + "narrow_range": true + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "torch", + "state_codec": "int8_hadamard32", + "prefill_state_qdq": false + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/requirements.txt b/examples/llm_qat/linear_attention/requirements.txt new file mode 100644 index 00000000000..29636d78e52 --- /dev/null +++ b/examples/llm_qat/linear_attention/requirements.txt @@ -0,0 +1,4 @@ +fla-core==0.5.1 +flash-linear-attention==0.5.1 +pyarrow>=15 +transformers==4.57.1 diff --git a/examples/llm_qat/linear_attention/train.py b/examples/llm_qat/linear_attention/train.py new file mode 100644 index 00000000000..7662c1ca593 --- /dev/null +++ b/examples/llm_qat/linear_attention/train.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fine-tune KDA attention parameters with recurrent-state fake quantization.""" + +import argparse +import json +from pathlib import Path + +import pyarrow.parquet as pq +import torch +from fla.layers.kda import KimiDeltaAttention +from transformers import AutoModelForCausalLM, AutoTokenizer + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase + + +def main(): + """Quantize, train on local text, and save a ModelOpt-aware Hugging Face checkpoint.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument( + "--train-data", type=Path, required=True, help="Parquet file with a text column" + ) + parser.add_argument( + "--quant-config", + type=Path, + default=Path(__file__).with_name("configs") / "kda_decode_state_int8.json", + ) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--train-steps", type=int, default=1) + parser.add_argument("--length", type=int, default=128) + parser.add_argument("--prefill-tokens", type=int, default=64) + parser.add_argument("--seed", type=int, default=2026) + parser.add_argument("--learning-rate", type=float, default=1e-5) + options = parser.parse_args() + if options.train_steps < 1 or options.length < 2: + parser.error("train-steps must be positive and length must be at least two") + if not 0 <= options.prefill_tokens < options.length: + parser.error("prefill-tokens must leave at least one suffix label") + torch.manual_seed(options.seed) + mto.enable_huggingface_checkpointing() + tokenizer = AutoTokenizer.from_pretrained( + options.model, local_files_only=True, trust_remote_code=options.trust_remote_code + ) + text = "\n\n".join(pq.read_table(options.train_data, columns=["text"])["text"].to_pylist()) + tokens = tokenizer.encode(text, add_special_tokens=False) + needed = options.train_steps * options.length + if len(tokens) < needed: + parser.error(f"Training data has {len(tokens)} tokens; {needed} are required") + blocks = torch.tensor(tokens[:needed]).reshape(options.train_steps, options.length) + model = AutoModelForCausalLM.from_pretrained( + options.model, + local_files_only=True, + trust_remote_code=options.trust_remote_code, + torch_dtype=torch.float32, + ) + layers = [module for module in model.modules() if isinstance(module, KimiDeltaAttention)] + if not layers: + raise ValueError("The model must contain FLA KimiDeltaAttention layers") + # Keep attention parameters and optimizer moments in FP32; freeze the rest in BF16. + model.requires_grad_(False) + for layer in layers: + layer.requires_grad_(True) + for parameter in model.parameters(): + if not parameter.requires_grad: + parameter.data = parameter.data.to(torch.bfloat16) + model.cuda() + mtq.quantize(model, json.loads(options.quant_config.read_text())) + model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + trainable = [p for p in model.parameters() if p.requires_grad] + optimizer = torch.optim.AdamW(trainable, lr=options.learning_rate, weight_decay=0) + model.train() + for step, block in enumerate(blocks, 1): + ids = block.unsqueeze(0).cuda() + labels = ids.clone() + labels[:, : options.prefill_tokens] = -100 + optimizer.zero_grad(set_to_none=True) + # Keep phase metadata active during activation-checkpoint recomputation. + with linear_attention_training_phase(model, [options.prefill_tokens]): + with torch.autocast("cuda", dtype=torch.bfloat16): + loss = model(input_ids=ids, labels=labels, use_cache=False).loss + if not torch.isfinite(loss): + raise RuntimeError("Non-finite training loss") + loss.backward() + torch.nn.utils.clip_grad_norm_(trainable, 1.0, error_if_nonfinite=True) + optimizer.step() + print(f"step={step} loss={loss.item():.4f}", flush=True) + model.save_pretrained(options.output) + tokenizer.save_pretrained(options.output) + + +if __name__ == "__main__": + main() diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py index 875cb39a7d0..84da262fd9f 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -1,5 +1,5 @@ # Adapted from: https://github.com/fla-org/flash-linear-attention/blob/516143e31fce/fla/ops/common/chunk_delta_h.py -# Adapted with modifications (marked [ModelOpt]): optional in-kernel FP8 E4M3 fake quantization +# Adapted with modifications (marked [ModelOpt]): optional in-kernel FP8 E4M3 / INT8 fake quantization # of the carried chunk state (STATE_QDQ), BV as an explicit launch argument, no fla backend # dispatch or config-cache autotune. # @@ -42,15 +42,18 @@ from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq +from .int8 import int8_scalar_qdq + # ``STATE_QDQ`` modes of the forward state kernel. STATE_QDQ_OFF = 0 STATE_QDQ_FP8_DYNAMIC = 1 # FP8 E4M3, one dynamic scale per program tile ([K, BV] of one head) +STATE_QDQ_INT8_DYNAMIC = 2 STATE_QDQ_MAX_BLOCK_V = 128 @triton.jit -def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr): - """[ModelOpt] Dynamic FP8 E4M3 scale per full [K, BV] tile of one sequence and head.""" +def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr, STATE_QDQ: tl.constexpr): + """[ModelOpt] Dynamic scale per full [K, BV] tile of one sequence and head.""" b_amax = tl.max(tl.abs(b_h1)) if K > 64: b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h2))) @@ -58,7 +61,18 @@ def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr): b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h3))) if K > 192: b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h4))) - return tl.where(b_amax > 0, b_amax / 448.0, 1.0) + if STATE_QDQ == 2: + return tl.where(b_amax > 0, b_amax * (1.0 / 127.0), 1.0) + else: + return tl.where(b_amax > 0, b_amax / 448.0, 1.0) + + +@triton.jit +def _state_scalar_qdq(value, scale, STATE_QDQ: tl.constexpr): + if STATE_QDQ == 2: + return int8_scalar_qdq(value, scale) + else: + return fp8_scalar_qdq(value, scale) NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] @@ -215,23 +229,23 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( p_h0_4 = h0 + o_k4[:, None] * V + o_v[None, :] m_h0_4 = m_k4[:, None] & m_v[None, :] b_h4 += tl.load(p_h0_4, mask=m_h0_4, other=0.0).to(tl.float32) - # [ModelOpt] A state read from an FP8 cache is quantized before the first chunk uses it. - if STATE_QDQ == 1: + # [ModelOpt] A state read from a quantized cache is quantized before the first chunk uses it. + if STATE_QDQ != 0: if K > 192: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K, STATE_QDQ=STATE_QDQ) elif K > 128: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K) + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K, STATE_QDQ=STATE_QDQ) elif K > 64: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K) + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K, STATE_QDQ=STATE_QDQ) else: - b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K) - b_h1 = fp8_scalar_qdq(b_h1, b_scale) + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K, STATE_QDQ=STATE_QDQ) + b_h1 = _state_scalar_qdq(b_h1, b_scale, STATE_QDQ) if K > 64: - b_h2 = fp8_scalar_qdq(b_h2, b_scale) + b_h2 = _state_scalar_qdq(b_h2, b_scale, STATE_QDQ) if K > 128: - b_h3 = fp8_scalar_qdq(b_h3, b_scale) + b_h3 = _state_scalar_qdq(b_h3, b_scale, STATE_QDQ) if K > 192: - b_h4 = fp8_scalar_qdq(b_h4, b_scale) + b_h4 = _state_scalar_qdq(b_h4, b_scale, STATE_QDQ) # main recurrence for i_t in range(NT): @@ -385,24 +399,24 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h4 = tl.dot(b_k, b_v, b_h4) - # [ModelOpt] Dynamic per-tile FP8 QDQ of the next-chunk or stored final state. - # Recompute amax over [K, BV]; fp8_scalar_qdq applies that tile's scalar scale. - if STATE_QDQ == 1: + # [ModelOpt] Dynamic per-tile FP8/INT8 QDQ of the next-chunk or stored final state. + # Recompute amax over [K, BV]; _state_scalar_qdq applies that tile's scalar scale. + if STATE_QDQ != 0: if K > 192: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K, STATE_QDQ=STATE_QDQ) elif K > 128: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K) + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K, STATE_QDQ=STATE_QDQ) elif K > 64: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K) + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K, STATE_QDQ=STATE_QDQ) else: - b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K) - b_h1 = fp8_scalar_qdq(b_h1, b_scale) + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K, STATE_QDQ=STATE_QDQ) + b_h1 = _state_scalar_qdq(b_h1, b_scale, STATE_QDQ) if K > 64: - b_h2 = fp8_scalar_qdq(b_h2, b_scale) + b_h2 = _state_scalar_qdq(b_h2, b_scale, STATE_QDQ) if K > 128: - b_h3 = fp8_scalar_qdq(b_h3, b_scale) + b_h3 = _state_scalar_qdq(b_h3, b_scale, STATE_QDQ) if K > 192: - b_h4 = fp8_scalar_qdq(b_h4, b_scale) + b_h4 = _state_scalar_qdq(b_h4, b_scale, STATE_QDQ) if STORE_FINAL_STATE: if STATE_V_FIRST: @@ -885,8 +899,8 @@ def state_qdq_tile_v(V: int, state_qdq: int, state_qdq_block_v: int | None) -> i """ if state_qdq == STATE_QDQ_OFF: return 64 if check_shared_mem("ada") else 32 - if state_qdq != STATE_QDQ_FP8_DYNAMIC: - raise ValueError(f"Unsupported state_qdq mode {state_qdq}; expected 0 or 1.") + if state_qdq not in (STATE_QDQ_FP8_DYNAMIC, STATE_QDQ_INT8_DYNAMIC): + raise ValueError(f"Unsupported state_qdq mode {state_qdq}; expected 0, 1, or 2.") BV = min(triton.next_power_of_2(V), 64) if state_qdq_block_v is None else state_qdq_block_v if BV < 16 or BV > STATE_QDQ_MAX_BLOCK_V or BV & (BV - 1): raise ValueError( diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py index a5609cab102..f34b8a3384b 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py @@ -59,6 +59,7 @@ from .fla_chunk_delta_h import ( STATE_QDQ_FP8_DYNAMIC, + STATE_QDQ_INT8_DYNAMIC, STATE_QDQ_OFF, chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h, @@ -636,7 +637,8 @@ def chunk_gated_delta_rule( raise ValueError("ModelOpt GDN supports only chunk_size=64; FLA WY backward assumes 64.") # [ModelOpt] state_qdq: 0 keeps fla's numerics; 1 fake-quantizes the state carried between - # chunks to FP8 E4M3 with a dynamic scale per [K, state_qdq_block_v] tile of each head. + # chunks to FP8 E4M3; 2 uses signed narrow-range INT8. Both use a dynamic scale per + # [K, state_qdq_block_v] tile of each head. state_qdq = kwargs.pop("state_qdq", STATE_QDQ_OFF) state_qdq_block_v = kwargs.pop("state_qdq_block_v", None) # w_quantizer: dynamic FP8 TensorQuantizer applied to the WY tensor @@ -647,11 +649,13 @@ def chunk_gated_delta_rule( dt_bias = kwargs.pop("dt_bias", None) if kwargs: raise TypeError(f"Unexpected keyword arguments: {', '.join(sorted(kwargs))}") - if state_qdq not in (STATE_QDQ_OFF, STATE_QDQ_FP8_DYNAMIC): - raise ValueError(f"`state_qdq` must be 0 or 1, got {state_qdq}.") + if state_qdq not in (STATE_QDQ_OFF, STATE_QDQ_FP8_DYNAMIC, STATE_QDQ_INT8_DYNAMIC): + raise ValueError(f"`state_qdq` must be 0, 1, or 2, got {state_qdq}.") if w_quantizer is not None: validate_gdn_quantizer(w_quantizer, name="gdn_w_quantizer") - if state_qdq and (not q.is_cuda or torch.cuda.get_device_capability(q.device) < (8, 9)): + if state_qdq == STATE_QDQ_FP8_DYNAMIC and ( + not q.is_cuda or torch.cuda.get_device_capability(q.device) < (8, 9) + ): raise RuntimeError("GDN state QDQ requires native E4M3 conversion on CUDA SM89 or newer.") if (state_qdq != STATE_QDQ_OFF or w_quantizer is not None) and cp_context is not None: raise ValueError("State or w quantization is not supported together with `cp_context`.") diff --git a/modelopt/torch/kernels/quantization/linear_attention/int8.py b/modelopt/torch/kernels/quantization/linear_attention/int8.py new file mode 100644 index 00000000000..d9e16a29dc8 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/int8.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Signed narrow-range INT8 QDQ for recurrent-state tiles.""" + +import triton +import triton.language as tl +from triton.language.extra.cuda import libdevice + +__all__ = [] + + +@triton.jit +def int8_scalar_qdq(value, scale): + codes = libdevice.nearbyint(tl.div_rn(value, scale)) + return tl.minimum(tl.maximum(codes, -127.0), 127.0) * scale diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 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 78a6e03a16b..db194a0f465 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -28,6 +28,7 @@ from modelopt.torch.opt.dynamic import _DMRegistryCls from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict from modelopt.torch.utils import get_unwrapped_name +from modelopt.torch.utils.logging import print_rank_0 from .config import ( QuantizeConfig, @@ -36,6 +37,7 @@ _QuantizeExportConfig, normalize_quant_cfg_list, ) +from .linear_attention.config import LinearAttentionConfig from .nn import ( QuantModule, QuantModuleRegistry, @@ -67,7 +69,7 @@ def convert_to_quantized_model(model: ModelLikeModule, config: QuantizeConfig) - replace_quant_module(model, version=ModeloptStateManager(model).state_version) set_quantizer_by_cfg(model, config.get("quant_cfg", [])) - _validate_linear_attention_quantizers(model) + _apply_linear_attention_policy(model, config) metadata = {} update_quantize_metadata(model, config, metadata) @@ -137,6 +139,7 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: details regarding how MCore sharded checkpoint is restored, see modelopt.torch.opt.plugins.mcore_dist_checkpointing.restore_sharded_modelopt_state. """ + _apply_linear_attention_policy(model, config, metadata.get("linear_attention")) if "quantizer_state" not in metadata: # MCore sharded checkpoint (`torch-dist`) has its quantizer_state stored as the # extra_state of `QuantModule`. The quantizer_state is resumed with @@ -144,13 +147,14 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: return model quantizer_state_dict = dict(metadata["quantizer_state"]) - # Older checkpoints predate these disabled handles; preserve their baseline path. - for name, module in _linear_attention_modules(model).items(): - for handle in ("gdn_state_quantizer", "gdn_w_quantizer"): - key = f"{name}.{handle}" if name else handle - quantizer = getattr(module, handle) - if key not in quantizer_state_dict and not quantizer.is_enabled: - quantizer_state_dict[key] = quantizer.get_modelopt_state() + if "linear_attention" not in metadata: + # Older checkpoints predate these disabled handles; preserve their baseline path. + for name, module in _linear_attention_modules(model).items(): + for handle in module.linear_attention_quantizer_names: + key = f"{name}.{handle}" if name else handle + quantizer = getattr(module, handle) + if key not in quantizer_state_dict and not quantizer.is_enabled: + quantizer_state_dict[key] = quantizer.get_modelopt_state() unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys() extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys() @@ -203,6 +207,12 @@ def update_quantize_metadata( model: nn.Module, config: QuantizeConfig, metadata: MetadataDict ) -> None: """Update the quantizer state in the metadata dict.""" + policies = { + name: module.linear_attention_config.model_dump() + for name, module in _linear_attention_modules(model).items() + } + if policies: + metadata["linear_attention"] = policies metadata["quantizer_state"] = quantizer_state(model) if shared_state_metadata := SharedWeightGlobalAmaxState.metadata(model): metadata["shared_quant_states"] = shared_state_metadata @@ -212,18 +222,39 @@ def update_quantize_metadata( def _linear_attention_modules(model): # Optional framework plugins import conversion; defer this import to avoid that cycle. - from .plugins.gated_delta_net import GatedDeltaNetStateQuantMixin + from .plugins.linear_attention import _LinearAttentionQuantMixin return { get_unwrapped_name(name, model): module for name, module in model.named_modules() - if isinstance(module, GatedDeltaNetStateQuantMixin) + if isinstance(module, _LinearAttentionQuantMixin) } -def _validate_linear_attention_quantizers(model): - for module in _linear_attention_modules(model).values(): +def _apply_linear_attention_policy(model, config, saved_policies=None): + modules = _linear_attention_modules(model) + # getattr also handles pickled pre-extension config objects with no field in __dict__. + for entry in getattr(config, "linear_attention", []): + matches = [name for name in modules if fnmatch.fnmatch(name, entry.module_name)] + if not matches: + raise ValueError( + f"linear_attention rule {entry.module_name!r} matches no supported modules" + ) + for name in matches: + modules[name].linear_attention_config = entry.cfg.model_copy(deep=True) + if saved_policies is not None: + if saved_policies.keys() != modules.keys(): + raise ApplyModeError( + "Saved linear_attention policies do not match the restored modules" + ) + for name, policy in saved_policies.items(): + modules[name].linear_attention_config = LinearAttentionConfig(**policy) + for name, module in modules.items(): module.validate_linear_attention() + if getattr(config, "linear_attention", []): + print_rank_0( + f"Linear attention {name!r}: {module.linear_attention_config.model_dump()}" + ) def quantizer_state(model: nn.Module) -> dict[str, Any]: diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py index 477eeb89847..305bb2cd81f 100644 --- a/modelopt/torch/quantization/linear_attention/__init__.py +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -13,4 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Helpers for linear-attention quantization.""" +"""Numerical policies and differentiable kernels for linear attention.""" + +from .config import * +from .decode import * +from .decode_prefill import * +from .kda import * +from .prefill import * diff --git a/modelopt/torch/quantization/linear_attention/_chunk_prefill.py b/modelopt/torch/quantization/linear_attention/_chunk_prefill.py new file mode 100644 index 00000000000..fddff0303ae --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/_chunk_prefill.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Exact chunked prefix computation for decode-aware training.""" + +import torch + +from .utils import _prepare, _state_qdq + +__all__ = [] + + +def chunk_gdn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + chunk_size: int = 64, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + state_v_first: bool = False, + state_qdq: bool = False, + state_qdq_block_v: int = 64, + state_format: str = "fp8_e4m3", +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute a GDN prefix with optional QDQ on the initial state and chunk writes.""" + if g.ndim != 3 or chunk_size <= 0: + raise ValueError("chunk_gdn requires scalar GDN gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format) + pieces = [] + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + gc = g[b, lo:hi].transpose(0, 1).cumsum(-1) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + # Mask before exp: upper-triangle positive differences can overflow for long decay. + causal = torch.ones(hi - lo, hi - lo, device=q.device, dtype=torch.bool).tril() + decay = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() + lower = (bc * (kc @ kc.transpose(-1, -2)) * decay).tril(-1) + matrix = lower + torch.eye(hi - lo, device=q.device, dtype=q.dtype) + gate = gc.exp().unsqueeze(-1) + rhs = torch.cat((bc * vc, bc * kc * gate), dim=-1) + solved = torch.linalg.solve_triangular(matrix, rhs, upper=False, unitriangular=True) + u, w = solved.split((v.shape[-1], k.shape[-1]), dim=-1) + updated_values = u - w @ state + local_scores = ((qc * scale @ kc.transpose(-1, -2)) * decay).tril() + output = (qc * (scale * gc.exp()).unsqueeze(-1)) @ state + pieces.append((output + local_scores @ updated_values).transpose(0, 1)) + weighted_keys = kc * (gc[..., -1:] - gc).exp().unsqueeze(-1) + state = state * gc[..., -1].exp()[:, None, None] + state = state + weighted_keys.transpose(-1, -2) @ updated_values + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format) + outputs.append(torch.cat(pieces)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final + + +def chunk_kda( + q, + k, + v, + g, + beta, + *, + chunk_size=64, + scale=None, + initial_state=None, + cu_seqlens=None, + state_v_first=False, + state_qdq=False, + state_qdq_block_v=64, + state_format="fp8_e4m3", +): + """Compute a KDA prefix with optional QDQ on the initial state and chunk writes.""" + if g.ndim != 4 or chunk_size <= 0: + raise ValueError("chunk_kda requires per-key gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format) + pieces = [] + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + prefix = g[b, lo:hi].transpose(0, 1).cumsum(-2) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + gate = prefix.exp() + lower_rows, score_rows = [], [] + for row in range(hi - lo): + decayed_keys = ( + kc[..., : row + 1, :] + * (prefix[..., row : row + 1, :] - prefix[..., : row + 1, :]).exp() + ) + right = decayed_keys.transpose(-1, -2) + left = kc[..., row : row + 1, :] * bc[..., row : row + 1, :] + interaction = left @ right + score = qc[..., row : row + 1, :] * scale @ right + padding = hi - lo - row - 1 + lower_rows.append(torch.nn.functional.pad(interaction, (0, padding))) + score_rows.append(torch.nn.functional.pad(score, (0, padding))) + lower = torch.cat(lower_rows, dim=-2).tril(-1) + scores = torch.cat(score_rows, dim=-2) + identity = torch.eye(hi - lo, device=q.device, dtype=q.dtype).expand_as(lower) + inverse = torch.linalg.solve_triangular( + identity + lower, identity, upper=False, unitriangular=True + ) + u = inverse @ (bc * vc) + w = inverse @ (bc * kc * gate) + updated = u - w @ state + pieces.append(((qc * scale * gate) @ state + scores @ updated).transpose(0, 1)) + weighted_keys = kc * (prefix[..., -1:, :] - prefix).exp() + state = state * gate[..., -1, :, None] + weighted_keys.transpose(-1, -2) @ updated + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format) + outputs.append(torch.cat(pieces)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py new file mode 100644 index 00000000000..1f3b610fe75 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Saved execution policy for GDN training-time numerical emulation.""" + +from typing import Literal + +from pydantic import Field, model_validator + +from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField + +__all__ = [ + "LinearAttentionConfig", + "LinearAttentionDecodeConfig", + "LinearAttentionPolicyEntry", + "LinearAttentionReplayConfig", +] + + +class _StateConfig(ModeloptBaseConfig): + mode: Literal["chunk"] = ModeloptField(default="chunk") + block_v: Literal[16, 32, 64, 128] = ModeloptField(default=64) + quantize_initial: Literal[True] = ModeloptField(default=True) + + +class _SolveConfig(ModeloptBaseConfig): + method: Literal["exact"] = ModeloptField(default="exact") + + +class LinearAttentionReplayConfig(ModeloptBaseConfig): + """Anchor refresh and encoded rank-one update policy.""" + + window: int = Field(default=8, ge=1, le=64, strict=True) + factor_qdq: bool = ModeloptField(default=True) + encoding: Literal["once", "reencode"] = ModeloptField(default="once") + + +class LinearAttentionDecodeConfig(ModeloptBaseConfig): + """Explicit suffix recurrence; workload supplies per-sequence prefix lengths.""" + + mode: Literal["token", "replay"] = ModeloptField(default="token") + implementation: Literal["torch"] = ModeloptField(default="torch") + readout: Literal["working", "stored"] = ModeloptField(default="stored") + quantize_initial: bool = ModeloptField(default=True) + prefill_state_qdq: bool = ModeloptField(default=False) + state_codec: Literal["tile", "int8_hadamard32"] = ModeloptField(default="tile") + decay_log_step: float | None = Field(default=None, gt=0, allow_inf_nan=False) + replay: LinearAttentionReplayConfig | None = ModeloptField(default=None) + + @model_validator(mode="after") + def _validate_replay(self): + if (self.mode == "replay") != (self.replay is not None): + raise ValueError("replay settings must be supplied exactly when mode='replay'") + if self.state_codec == "int8_hadamard32" and self.prefill_state_qdq: + raise ValueError( + "Hadamard state QDQ starts at decode handoff; disable prefill_state_qdq" + ) + return self + + +class LinearAttentionConfig(ModeloptBaseConfig): + """GDN chunk-64 policy; unsupported numerical modes fail config validation. + + ``state.block_v`` defines one dynamic scale per ``[Dk, block_v]`` tile of each + sequence/head. The initial state and every chunk's final state are rounded when + ``gdn_state_quantizer`` is enabled. Outputs use the incoming rounded state. + Decode's ``int8_hadamard32`` codec instead fixes scales to one key channel and + 32 values; ``state.block_v`` remains the execution tile width. + """ + + schema_version: Literal[1] = ModeloptField(default=1) + backend: Literal["fla", "matmul"] = ModeloptField(default="fla") + chunk_size: Literal[64] = ModeloptField(default=64) + state: _StateConfig = ModeloptField(default=_StateConfig()) + solve: _SolveConfig = ModeloptField(default=_SolveConfig()) + decode: LinearAttentionDecodeConfig | None = ModeloptField(default=None) + + @model_validator(mode="after") + def _validate_decode_backend(self): + if (self.backend == "matmul") != (self.decode is not None): + raise ValueError("The exact-prefix matmul backend requires an explicit decode policy") + if self.decode is not None and self.decode.state_codec == "int8_hadamard32": + if self.state.block_v < 32: + raise ValueError("int8_hadamard32 requires block_v >= 32") + return self + + +class LinearAttentionPolicyEntry(ModeloptBaseConfig): + """Assign a complete policy to supported modules matching ``module_name``. + + Rules apply in order: the last match wins, without merging nested fields. + A rule must match at least one supported linear-attention module. + """ + + module_name: str = Field(...) + cfg: LinearAttentionConfig = ModeloptField(default=LinearAttentionConfig()) diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py new file mode 100644 index 00000000000..2a981c597ec --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -0,0 +1,341 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Explicit token-state and encoded-update replay references for QAT.""" + +from dataclasses import dataclass + +import torch + +from .config import LinearAttentionDecodeConfig +from .utils import _fp8_quantize + +__all__ = [ + "EncodedLinearAttentionTensor", + "LinearAttentionCarry", + "ReplayEntry", + "recurrent_decode", + "recurrent_decode_reference", +] + + +@dataclass +class EncodedLinearAttentionTensor: + """Fake decoded values plus detached scale metadata; no compressed storage claim.""" + + values: torch.Tensor + scales: torch.Tensor | None + format: str + block_v: int | None + + +@dataclass +class ReplayEntry: + """An already computed, encoded rank-one update and its log retention.""" + + key: EncodedLinearAttentionTensor + update: EncodedLinearAttentionTensor + log_retention: torch.Tensor + + +@dataclass +class LinearAttentionCarry: + """Explicit anchor/update state, position, and codec contract across calls.""" + + anchor: EncodedLinearAttentionTensor + entries: tuple[ReplayEntry, ...] + position: int + started: bool + signature: str + value_basis: str = "identity" + + @property + def cursor(self): + """Number of encoded updates since the last anchor refresh.""" + return len(self.entries) + + def reconstruct(self, *, original_basis=True): + """Replay entries, returning the original value basis unless explicitly disabled.""" + state = self.anchor.values + for entry in self.entries: + gate = entry.log_retention + decay = gate.exp().unsqueeze(-1) + if gate.ndim == 1: + decay = decay.unsqueeze(-1) + state = state * decay + entry.key.values.unsqueeze(-1) * entry.update.values.unsqueeze( + -2 + ) + if original_basis and self.value_basis == "hadamard32": + state = _hadamard32(state) + return state + + +def _hadamard32(value): + """Apply the orthonormal Sylvester transform to contiguous 32-value groups.""" + shape = value.shape + for width in (1, 2, 4, 8, 16): + pairs = value.reshape(*shape[:-1], -1, 2, width) + left, right = pairs.unbind(-2) + value = torch.stack((left + right, left - right), dim=-2).reshape(shape) + return value * (32**-0.5) + + +class _IdentityGradient(torch.autograd.Function): + @staticmethod + def forward(ctx, value, rounded): + return rounded + + @staticmethod + def backward(ctx, gradient): + return gradient, None + + +def _encode(value, enabled, block_v, *, state=False, state_format="fp8_e4m3", state_codec="tile"): + if not enabled: + return EncodedLinearAttentionTensor(value, None, "identity", None) + if state and state_codec == "int8_hadamard32": + with torch.no_grad(): + groups = value.float().reshape(*value.shape[:-1], -1, 32) + scales = (groups.abs().amax(-1) / 127).clamp_min(6e-8) + normalized = groups / scales.unsqueeze(-1) + codes = normalized.sign() * (normalized.abs() + 0.5).floor() + metadata = scales.to(torch.float16) + decoded = (codes.clamp(-127, 127) * metadata.float().unsqueeze(-1)).reshape_as(value) + return EncodedLinearAttentionTensor( + _IdentityGradient.apply(value, decoded.to(value.dtype)), metadata, "int8", 32 + ) + rounded, scales = [], [] + for part in value.split(block_v, dim=-1): + if state_format == "fp8_e4m3": + tensor = part.flatten(-2) if state else part + axis = tuple(range(tensor.ndim - 1)) or None + decoded, scale = _fp8_quantize(tensor, axis) + rounded.append(decoded.reshape_as(part)) + scales.append(scale.squeeze(-1)) + else: + with torch.no_grad(): + axes = (-2, -1) if state else (-1,) + part = part.float() + amax = part.abs().amax(dim=axes, keepdim=True) + scale = torch.where(amax > 0, amax / 127.0, torch.ones_like(amax)) + rounded.append((part / scale).round().clamp(-127, 127) * scale) + scales.append(scale[..., 0, 0] if state else scale[..., 0]) + decoded = torch.cat(rounded, dim=-1).to(value.dtype) + if state_format == "int8": + decoded = _IdentityGradient.apply(value, decoded) + return EncodedLinearAttentionTensor(decoded, torch.stack(scales, dim=-1), state_format, block_v) + + +def _signature(config, state_qdq, block_v, state_format): + return ( + config.model_dump_json(exclude={"implementation", "prefill_state_qdq"}) + + f"/{state_qdq}/{block_v}/{state_format}" + ) + + +def _sum_keys(value): + keys = value.shape[-2] + padded = 1 << (keys - 1).bit_length() + if keys != padded: + value = torch.nn.functional.pad(value, (0, 0, 0, padded - keys)) + while value.shape[-2] > 1: + half = value.shape[-2] // 2 + value = value[..., :half, :] + value[..., half:, :] + return value[..., 0, :] + + +def _round_log_gate(gate, step): + if step is None: + return gate + rounded = (gate / step).round() * step + return _IdentityGradient.apply(gate, rounded.detach()) + + +def _prepare_carry( + q, k, v, g, beta, config, state_qdq, block_v, initial_state, carry, position, state_format +): + if q.ndim != 3 or k.shape != q.shape or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k/v must have aligned [T,H,D] shapes") + if beta.shape != q.shape[:2] or g.shape not in (beta.shape, q.shape): + raise ValueError("beta must be [T,H]; g must be [T,H] or [T,H,Dk]") + if block_v not in (16, 32, 64, 128): + raise ValueError("block_v must be 16, 32, 64, or 128") + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") + hadamard = config.state_codec == "int8_hadamard32" + if hadamard: + if state_qdq and state_format != "int8": + raise ValueError("int8_hadamard32 requires INT8 state quantization") + if v.shape[-1] % 32 or block_v < 32: + raise ValueError("int8_hadamard32 requires Dv divisible by 32 and block_v >= 32") + signature = _signature(config, state_qdq, block_v, state_format) + if carry is not None and initial_state is not None: + raise ValueError("Supply either carry or initial_state") + shape = (q.shape[1], q.shape[2], v.shape[2]) + if carry is None: + initial_state = q.new_zeros(shape) if initial_state is None else initial_state + carry = LinearAttentionCarry( + _encode(initial_state, False, block_v, state=True, state_format=state_format), + (), + position, + False, + signature, + ) + if carry.signature != signature or carry.anchor.values.shape != shape: + raise ValueError("Carry policy or state shape does not match this recurrence") + if config.mode == "token" and carry.entries: + raise ValueError("Token carry cannot contain replay entries") + if config.replay is not None and carry.cursor >= config.replay.window: + raise ValueError("Replay cursor must be below the refresh window") + if len(q) and not carry.started: + initial = _hadamard32(carry.anchor.values) if hadamard else carry.anchor.values + anchor = _encode( + initial, + state_qdq and config.quantize_initial, + block_v, + state=True, + state_format=state_format, + state_codec=config.state_codec, + ) + carry = LinearAttentionCarry( + anchor, (), carry.position, True, signature, "hadamard32" if hadamard else "identity" + ) + return carry, signature + + +def recurrent_decode_reference( + q, + k, + v, + g, + beta, + *, + config: LinearAttentionDecodeConfig, + state_qdq=False, + state_format="fp8_e4m3", + block_v=64, + initial_state=None, + carry=None, + position=0, + scale=None, +): + """Run one preactivated sequence [T,H,D] with explicit token/replay write events. + + Keys and value heads must already be aligned. Scalar GDN or per-key-channel + KDA log gates are accepted. Outputs and all returned carry values retain their + graphs. An empty call performs no write or initial-state quantization. + """ + carry, signature = _prepare_carry( + q, k, v, g, beta, config, state_qdq, block_v, initial_state, carry, position, state_format + ) + if len(q) == 0: + # Keep empty input gradients defined without introducing a state write. + zero = (q.sum() + k.sum() + v.sum() + g.sum() + beta.sum()) * 0 + output = v + zero + return output, carry + if config.state_codec == "int8_hadamard32": + v = _hadamard32(v) + scale = q.shape[-1] ** -0.5 if scale is None else scale + outputs = [] + with torch.autocast(device_type=q.device.type, enabled=False): + for t in range(len(q)): + entries = carry.entries + if config.replay is not None and config.replay.encoding == "reencode": + entries = tuple( + ReplayEntry( + _encode(e.key.values, config.replay.factor_qdq, k.shape[-1]), + _encode(e.update.values, config.replay.factor_qdq, block_v), + e.log_retention, + ) + for e in entries + ) + carry = LinearAttentionCarry( + carry.anchor, entries, carry.position, True, signature, carry.value_basis + ) + state = carry.reconstruct(original_basis=False) + gate = _round_log_gate(g[t], config.decay_log_step) + decay = gate.exp().unsqueeze(-1) + if gate.ndim == 1: + decay = decay.unsqueeze(-1) + key = _encode(k[t], config.replay is not None and config.replay.factor_qdq, k.shape[-1]) + decayed = state * decay + residual = v[t] - _sum_keys(key.values.unsqueeze(-1) * decayed) + update = _encode( + beta[t].unsqueeze(-1) * residual, + config.replay is not None and config.replay.factor_qdq, + block_v, + ) + working = decayed + key.values.unsqueeze(-1) * update.values.unsqueeze(-2) + if config.replay is not None: + entries = (*entries, ReplayEntry(key, update, gate)) + refresh = config.replay is None or len(entries) == config.replay.window + if refresh: + anchor = _encode( + working, + state_qdq, + block_v, + state=True, + state_format=state_format, + state_codec=config.state_codec, + ) + entries = () + stored = anchor.values + else: + anchor = carry.anchor + stored = working + next_carry = LinearAttentionCarry( + anchor, entries, carry.position + 1, True, signature, carry.value_basis + ) + read = working if config.readout == "working" else stored + outputs.append(_sum_keys(q[t].unsqueeze(-1) * read) * scale) + carry = next_carry + output = torch.stack(outputs) + if config.state_codec == "int8_hadamard32": + output = _hadamard32(output) + return output, carry + + +def recurrent_decode( + q, + k, + v, + g, + beta, + *, + config: LinearAttentionDecodeConfig, + state_qdq=False, + state_format="fp8_e4m3", + block_v=64, + initial_state=None, + carry=None, + position=0, + scale=None, +): + """Run the Torch training implementation with explicit differentiable carry.""" + return recurrent_decode_reference( + q, + k, + v, + g, + beta, + config=config, + state_qdq=state_qdq, + state_format=state_format, + block_v=block_v, + initial_state=initial_state, + carry=carry, + position=position, + scale=scale, + ) diff --git a/modelopt/torch/quantization/linear_attention/decode_prefill.py b/modelopt/torch/quantization/linear_attention/decode_prefill.py new file mode 100644 index 00000000000..0a0f69b4ae7 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/decode_prefill.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differentiable explicit per-sequence prefill/decode phase handoff.""" + +from contextlib import contextmanager +from itertools import pairwise + +import torch + +from ._chunk_prefill import chunk_gdn, chunk_kda +from .decode import recurrent_decode + +__all__ = ["linear_attention_training_phase"] + + +def _lengths(values): + if isinstance(values, torch.Tensor): + values = values.tolist() + values = tuple(values) + if any(type(value) is not int or value < 0 for value in values): + raise ValueError("Prefill lengths must be nonnegative integers") + return values + + +@contextmanager +def linear_attention_training_phase(model, prefill_lengths): + """Supply explicit sequence phases through forward and checkpointed backward. + + Runtime phase metadata is local to the converted layers and restored on exit. + Keep this context active through backward when activation checkpointing is used. + """ + # The plugin imports this numerical package during quantization initialization. + from ..plugins.linear_attention import _LinearAttentionQuantMixin + + lengths = _lengths(prefill_lengths) + layers = [ + m + for m in model.modules() + if isinstance(m, _LinearAttentionQuantMixin) + and m.linear_attention_config.decode is not None + ] + if not layers: + raise ValueError("The model has no converted decode-aware linear-attention layers") + previous = [getattr(m, "_linear_attention_prefill_lengths", None) for m in layers] + try: + for module in layers: + module._linear_attention_prefill_lengths = lengths + yield model + finally: + for module, original in zip(layers, previous): + module._linear_attention_prefill_lengths = original + + +def _decode_prefill( + q, + k, + v, + g, + beta, + *, + policy, + state_qdq, + state_format, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + state_v_first, + output_dtype, + prefill_lengths, +): + if prefill_lengths is None: + raise ValueError("Decode-aware training requires explicit per-sequence prefill lengths") + if q.ndim != 4 or k.shape != q.shape or v.ndim != 4 or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k and v must have compatible [B,T,H,D] shapes") + batch, length, key_heads, keys = q.shape + heads, values = v.shape[2:] + if batch < 1 or key_heads < 1 or heads % key_heads or beta.shape != (batch, length, heads): + raise ValueError("Invalid batch/head dimensions or beta shape") + if g.shape not in (beta.shape, (*beta.shape, keys)): + raise ValueError("Invalid GDN/KDA log-retention shape") + q, k = (x.repeat_interleave(heads // key_heads, dim=2) for x in (q, k)) + boundaries = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens + if boundaries is None: + sequences = [(b, 0, length) for b in range(batch)] + else: + bounds = _lengths(boundaries) + if ( + batch != 1 + or len(bounds) < 2 + or bounds[0] != 0 + or bounds[-1] != length + or any(a > b for a, b in pairwise(bounds)) + ): + raise ValueError( + "Packed boundaries must partition a batch of one, allowing empty entries" + ) + sequences = [(0, a, b) for a, b in pairwise(bounds)] + prefixes = _lengths(prefill_lengths) + if len(prefixes) != len(sequences) or any( + p > end - start for p, (_, start, end) in zip(prefixes, sequences) + ): + raise ValueError("Supply one valid prefill length per sequence") + if initial_state is None: + states = q.new_zeros(len(sequences), heads, keys, values) + else: + states = initial_state.transpose(-1, -2) if state_v_first else initial_state + states = states.to(q.dtype) + if states.shape != (len(sequences), heads, keys, values): + raise ValueError("Initial state shape does not match sequence/head dimensions") + active = [n for n, prefix in enumerate(prefixes) if prefix] + prefix_outputs, prefix_states = {}, {} + if active: + packed = [ + torch.cat( + [ + tensor[sequences[n][0], sequences[n][1] : sequences[n][1] + prefixes[n]] + for n in active + ] + ).unsqueeze(0) + for tensor in (q, k, v, g, beta) + ] + offsets = [0] + for n in active: + offsets.append(offsets[-1] + prefixes[n]) + prefix_fn = chunk_kda if g.ndim == 4 else chunk_gdn + with torch.autocast(device_type=q.device.type, enabled=False): + output, final = prefix_fn( + *packed, + state_qdq=state_qdq and policy.decode.prefill_state_qdq, + state_format=state_format, + scale=scale, + initial_state=states[active], + cu_seqlens=torch.tensor(offsets), + state_v_first=False, + chunk_size=policy.chunk_size, + state_qdq_block_v=policy.state.block_v, + ) + for index, n in enumerate(active): + prefix_outputs[n] = output[0, offsets[index] : offsets[index + 1]] + prefix_states[n] = final[index] + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + split = start + prefixes[n] + suffix, carry = recurrent_decode( + *(x[b, split:end] for x in (q, k, v, g, beta)), + config=policy.decode, + state_qdq=state_qdq, + state_format=state_format, + block_v=policy.state.block_v, + initial_state=prefix_states.get(n, states[n]), + position=prefixes[n], + scale=scale, + ) + prefix = prefix_outputs.get(n, q.new_empty(0, heads, values)) + outputs.append(torch.cat((prefix, suffix))) + finals.append(carry.reconstruct()) + output = torch.stack(outputs) if boundaries is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + if state_v_first: + final = final.transpose(-1, -2) + return output.to(output_dtype), final if output_final_state else None diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py new file mode 100644 index 00000000000..446dd6f7d0d --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differentiable KDA prefill with stable per-channel decay interactions.""" + +import torch +import torch.nn.functional as F + +from .decode_prefill import _decode_prefill + +__all__ = ["matmul_kda"] + + +def matmul_kda( + q, + k, + v, + g, + beta, + *, + policy, + state_qdq=False, + state_format="fp8_e4m3", + scale=None, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=False, + use_beta_sigmoid_in_kernel=False, + allow_neg_eigval=False, + A_log=None, # noqa: N803 - match the FLA kernel signature + dt_bias=None, + safe_gate=False, + lower_bound=None, + cu_seqlens=None, + cu_seqlens_cpu=None, + state_v_first=False, + chunk_size=64, + cp_context=None, + disable_recompute=False, + return_intermediate_states=False, + prefill_lengths=None, +): + """Normalize KDA inputs and run exact prefix plus configured suffix recurrence. + + Gates use the loaded FLA model's activation formula and per-key log retention. + """ + if policy.backend != "matmul" or chunk_size != policy.chunk_size: + raise ValueError("matmul_kda requires backend='matmul' and its configured chunk size") + if cp_context is not None or disable_recompute or return_intermediate_states: + raise NotImplementedError( + "KDA matmul does not support CP or FLA recompute/intermediate flags" + ) + if allow_neg_eigval and not use_beta_sigmoid_in_kernel: + raise ValueError("allow_neg_eigval requires use_beta_sigmoid_in_kernel") + if safe_gate and use_gate_in_kernel and (lower_bound is None or not -5 <= lower_bound < 0): + raise ValueError("safe_gate requires a lower_bound in [-5, 0)") + if lower_bound is not None and lower_bound >= 0: + raise ValueError("lower_bound must be negative") + output_dtype = q.dtype + dtype = torch.float64 if output_dtype == torch.float64 else torch.float32 + q, k, v, g, beta = (x.to(dtype) for x in (q, k, v, g, beta)) + if g.ndim != 4: + raise ValueError("matmul_kda requires per-key-channel KDA log gates") + if use_qk_l2norm_in_kernel: + q, k = (x * (x.square().sum(-1, keepdim=True) + 1e-6).rsqrt() for x in (q, k)) + if use_gate_in_kernel: + if A_log is None: + raise ValueError("Fused KDA gate requires A_log") + if dt_bias is not None: + g = g + dt_bias.to(dtype).reshape(g.shape[-2:]) + rate = A_log.to(dtype).exp().reshape(g.shape[-2], 1) + g = -rate * F.softplus(g) if lower_bound is None else lower_bound * (rate * g).sigmoid() + if use_beta_sigmoid_in_kernel: + beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) + if policy.decode is not None: + return _decode_prefill( + q, + k, + v, + g, + beta, + policy=policy, + state_qdq=state_qdq, + state_format=state_format, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + state_v_first=state_v_first, + output_dtype=output_dtype, + prefill_lengths=prefill_lengths, + ) + raise ValueError("An explicit decode policy is required") diff --git a/modelopt/torch/quantization/linear_attention/prefill.py b/modelopt/torch/quantization/linear_attention/prefill.py new file mode 100644 index 00000000000..4a602f8f5bc --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/prefill.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Batched differentiable GDN prefill with materialized numerical boundaries.""" + +import torch +import torch.nn.functional as F + +from .config import LinearAttentionConfig +from .decode_prefill import _decode_prefill + +__all__ = ["matmul_gdn"] + + +def matmul_gdn( + q, + k, + v, + g, + beta, + *, + policy: LinearAttentionConfig, + state_qdq=False, + state_format="fp8_e4m3", + scale=None, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=False, + use_beta_sigmoid_in_kernel=False, + allow_neg_eigval=False, + A_log=None, # noqa: N803 - match the FLA kernel signature + dt_bias=None, + cu_seqlens=None, + cu_seqlens_cpu=None, + state_v_first=False, + chunk_size=64, + cp_context=None, + prefill_lengths=None, +): + """Normalize GDN inputs and run exact prefix plus configured suffix recurrence. + + Uses FP32 working arithmetic (FP64 for double inputs), with floating QDQ state. + Prefixes use exact chunk algebra; no configurable operand rounding is enabled. + """ + if policy.backend != "matmul" or chunk_size != policy.chunk_size: + raise ValueError("matmul_gdn requires backend='matmul' and its configured chunk size") + if cp_context is not None: + raise NotImplementedError("GDN matmul emulation does not support context parallelism") + output_dtype = q.dtype + dtype = torch.float64 if output_dtype == torch.float64 else torch.float32 + q, k, v, g, beta = (x.to(dtype) for x in (q, k, v, g, beta)) + if use_qk_l2norm_in_kernel: + q, k = (x * (x.square().sum(-1, keepdim=True) + 1e-6).rsqrt() for x in (q, k)) + if use_gate_in_kernel: + if A_log is None or dt_bias is None: + raise ValueError("Fused GDN gate requires A_log and dt_bias") + g = -A_log.to(dtype).exp() * F.softplus(g + dt_bias.to(dtype)) + if use_beta_sigmoid_in_kernel: + beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) + if g.ndim != 3: + raise ValueError("matmul_gdn requires scalar GDN log gates") + if policy.decode is not None: + return _decode_prefill( + q, + k, + v, + g, + beta, + policy=policy, + state_qdq=state_qdq, + state_format=state_format, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + state_v_first=state_v_first, + output_dtype=output_dtype, + prefill_lengths=prefill_lengths, + ) + raise ValueError("An explicit decode policy is required") diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py index 4f1e18a174d..68353bd28df 100644 --- a/modelopt/torch/quantization/linear_attention/utils.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -15,18 +15,35 @@ """Shared helpers for linear-attention quantization.""" -from ..nn import TensorQuantizer +from __future__ import annotations + +from itertools import pairwise +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from ..nn import TensorQuantizer __all__ = [] -def validate_gdn_quantizer(quantizer: TensorQuantizer, *, name: str) -> None: - """Check the dynamic E4M3 contract and the custom backward's identity STE.""" +def validate_gdn_quantizer( + quantizer: TensorQuantizer, + *, + name: str, + num_bits: tuple[int | tuple[int, int], ...] = ((4, 3),), +) -> None: + """Check supported formats and the custom backward's identity STE.""" + # Numerical helpers load while QuantizeConfig initializes; defer the quantizer import. + from ..nn import TensorQuantizer + if not isinstance(quantizer, TensorQuantizer): raise ValueError(f"{name} requires a single TensorQuantizer") if not ( quantizer._dynamic - and quantizer.num_bits == (4, 3) + and quantizer.num_bits in num_bits + and (quantizer.num_bits != 8 or (not quantizer.unsigned and quantizer.narrow_range)) and quantizer.block_sizes is None and quantizer.fake_quant and quantizer._pass_through_bwd @@ -37,7 +54,89 @@ def validate_gdn_quantizer(quantizer: TensorQuantizer, *, name: str) -> None: and not quantizer._use_constant_amax ): raise ValueError( - f"{name} supports only dynamic E4M3 fake quantization with " + f"{name} supports only dynamic fake quantization with num_bits in {num_bits}, " "pass_through_bwd=True, no block_sizes, rotation, pre-scaling, bias, constant " "amax, or custom backend. Other gradient rules and formats are not implemented." ) + + +def _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first): + if q.ndim != 4 or k.shape != q.shape or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k must have shape [B,T,Hk,Dk] and v shape [B,T,Hv,Dv]") + batch, length, heads, keys = q.shape + value_heads, values = v.shape[2:] + if length == 0 or value_heads % heads: + raise ValueError("nonempty sequences and Hv divisible by Hk are required") + if beta.shape != (batch, length, value_heads) or g.shape not in ( + beta.shape, + (*beta.shape, keys), + ): + raise ValueError("beta must be [B,T,Hv]; g must be [B,T,Hv] or [B,T,Hv,Dk]") + if cu_seqlens is None: + sequences = [(b, 0, length) for b in range(batch)] + else: + boundaries = cu_seqlens.tolist() + if ( + batch != 1 + or len(boundaries) < 2 + or boundaries[0] != 0 + or boundaries[-1] != length + or any(a >= b for a, b in pairwise(boundaries)) + ): + raise ValueError("cu_seqlens must partition a packed batch of size one") + sequences = [(0, a, b) for a, b in pairwise(boundaries)] + if initial_state is None: + state = q.new_zeros(len(sequences), value_heads, keys, values) + else: + state = initial_state.transpose(-1, -2) if state_v_first else initial_state + if state.shape != (len(sequences), value_heads, keys, values): + raise ValueError("initial_state shape does not match sequence/head dimensions") + return ( + q.repeat_interleave(value_heads // heads, dim=2), + k.repeat_interleave(value_heads // heads, dim=2), + state, + sequences, + ) + + +def _fp8_quantize(value: torch.Tensor, axis): + """Apply ModelOpt FP8 QDQ and return its detached dequantization scales.""" + # QuantizeConfig imports this module before the quantizer classes are initialized. + from ..config import QuantizerAttributeConfig + from ..nn import TensorQuantizer + + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=(4, 3), type="dynamic", axis=axis, pass_through_bwd=True) + ) + quantized = quantizer(value) + amax = quantizer._get_amax(value).float() + safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) + return quantized, torch.div(448.0, safe_amax).reciprocal() + + +def _state_qdq(state: torch.Tensor, block_v: int = 64, state_format: str = "fp8_e4m3"): + """Dynamic state-tile QDQ with detached scales and identity STE.""" + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") + if block_v not in (16, 32, 64, 128): + raise ValueError("block_v must be 16, 32, 64, or 128") + if state_format == "fp8_e4m3": + axis = tuple(range(state.ndim - 2)) or None + return torch.cat( + [ + _fp8_quantize(tile.flatten(-2), axis)[0].reshape_as(tile) + for tile in state.split(block_v, dim=-1) + ], + dim=-1, + ) + with torch.no_grad(): + rounded = [] + for tile in state.float().split(block_v, dim=-1): + amax = tile.abs().amax(dim=(-2, -1), keepdim=True) + limit = 127.0 + scale = torch.where(amax > 0, amax / limit, torch.ones_like(amax)) + normalized = (tile / scale).clamp(-limit, limit) + codes = normalized.round() + rounded.append(codes * scale) + quantized = torch.cat(rounded, dim=-1).to(state.dtype) + return state + (quantized - state).detach() diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index fbf9bb299fd..5aac82cacfe 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -32,7 +32,7 @@ from modelopt.torch.opt.utils import forward_with_reshard from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.conversion import ( - _validate_linear_attention_quantizers, + _apply_linear_attention_policy, preserve_quantizer_attributes_context, set_quantizer_attributes_partial, set_quantizer_by_cfg, @@ -330,7 +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) - _validate_linear_attention_quantizers(model) + _apply_linear_attention_policy(model, quantize_config) # Fail before calibration rather than after exporting an unquantized checkpoint. _check_weight_quantization_took_effect(model, quantize_config) 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 963452ff509..e1b2ec7a1ae 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -39,7 +39,8 @@ from .attention import * from .custom import * -from .gated_delta_net import * +from .gdn import * +from .kda import * with import_plugin("diffusers"): from .diffusion.diffusers import * diff --git a/modelopt/torch/quantization/plugins/gated_delta_net.py b/modelopt/torch/quantization/plugins/gdn.py similarity index 70% rename from modelopt/torch/quantization/plugins/gated_delta_net.py rename to modelopt/torch/quantization/plugins/gdn.py index df70683cc2f..8d0fe616390 100644 --- a/modelopt/torch/quantization/plugins/gated_delta_net.py +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -17,9 +17,10 @@ The chunked gated-delta-rule kernel keeps each head's ``[K, V]`` recurrent state in fp32 inside one Triton launch and carries it from chunk to chunk. To emulate a deployment that stores that -state in FP8, ModelOpt runs an adapted copy of the kernel +state in FP8 or INT8, ModelOpt runs an adapted copy of the kernel (:mod:`modelopt.torch.kernels.quantization.linear_attention`) that fake-quantizes the state to -E4M3 at the end of every chunk, with a scale computed inside the kernel from the state itself. +E4M3 or signed narrow-range INT8 at the end of every chunk, with a scale computed inside the +kernel from the state itself. The backward pass recomputes the same quantized states and passes the state gradient straight through the quantization, so QAT and QAD train against the quantized recurrence. A second quantizer covers ``w``, the WY-transformed keys that multiply the state; ``w`` is a regular tensor, @@ -32,9 +33,8 @@ import torch -from ..config import QuantizerAttributeConfig -from ..linear_attention.utils import validate_gdn_quantizer -from ..nn import QuantModule, TensorQuantizer +from ..linear_attention.prefill import matmul_gdn +from .linear_attention import _LinearAttentionQuantMixin __all__ = ["GatedDeltaNetStateQuantMixin"] @@ -63,39 +63,22 @@ def _state_qdq_chunk_gated_delta_rule() -> GatedDeltaRuleFn: return chunk_gated_delta_rule -class GatedDeltaNetStateQuantMixin(QuantModule): +class GatedDeltaNetStateQuantMixin(_LinearAttentionQuantMixin): """Adds ``gdn_state_quantizer`` and ``gdn_w_quantizer`` to a GatedDeltaNet module. Subclasses route the module's chunked gated-delta-rule call through :meth:`_state_quantized_chunk_gated_delta_rule`. Both quantizers start disabled; enable them with ``quant_cfg`` entries on ``*gdn_state_quantizer`` / ``*gdn_w_quantizer`` such as the ``configs/ptq/units/gdn_state_fp8_dynamic`` and ``gdn_w_fp8_dynamic`` recipe units. The state - quantizer carries the fused QDQ configuration. Both sites currently require dynamic - E4M3 and identity STE. State QDQ uses fixed 64-column value tiles and 64-token chunks. + quantizer carries the fused QDQ configuration. State supports dynamic E4M3 or signed + narrow-range INT8; W supports dynamic E4M3. + Both sites require identity STE. The execution policy is saved in ModelOpt metadata. """ - def _setup(self): - for name in ("gdn_state_quantizer", "gdn_w_quantizer"): - self._register_temp_attribute( - name, TensorQuantizer(QuantizerAttributeConfig(enable=False)) - ) - - def validate_linear_attention(self) -> None: - """Reject numerical settings that the fused training path cannot implement.""" - for name in ("gdn_state_quantizer", "gdn_w_quantizer"): - quantizer = getattr(self, name) - if quantizer.is_enabled: - validate_gdn_quantizer(quantizer, name=name) - # The state handle configures fused QDQ; W's grouping is executed by TensorQuantizer. - if self.gdn_state_quantizer.is_enabled and self.gdn_state_quantizer.axis != (0, 1): - raise ValueError( - "gdn_state_quantizer supports only axis=(0, 1) with 64-column value tiling" - ) - - def modelopt_post_restore(self, prefix: str = ""): - """Validate restored quantizers against the fused training kernel requirements.""" - super().modelopt_post_restore(prefix) - self.validate_linear_attention() + @property + def gdn_state_qdq_block_v(self) -> int: + """Value-column scale grouping from the saved execution policy.""" + return self.linear_attention_config.state.block_v def _state_quantized_chunk_gated_delta_rule( self, gated_delta_rule: GatedDeltaRuleFn, *args: Any, **kwargs: Any @@ -104,7 +87,7 @@ def _state_quantized_chunk_gated_delta_rule( self.validate_linear_attention() quantize_state = self.gdn_state_quantizer.is_enabled and self.gdn_state_quantizer._if_quant quantize_w = self.gdn_w_quantizer.is_enabled - if not (quantize_state or quantize_w): + if not self.linear_attention_is_enabled: return gated_delta_rule(*args, **kwargs) while isinstance(gated_delta_rule, partial): args = (*gated_delta_rule.args, *args) @@ -115,14 +98,26 @@ def _state_quantized_chunk_gated_delta_rule( "GatedDeltaNet quantization supports only FLA's chunk_gated_delta_rule " f"callable or a functools.partial of it; got {gated_delta_rule!r}." ) - chunk_size = kwargs.pop("chunk_size", 64) - if chunk_size != 64: + chunk_size = kwargs.pop("chunk_size", self.linear_attention_config.chunk_size) + if chunk_size != self.linear_attention_config.chunk_size: raise ValueError("GDN fake quantization supports only chunk_size=64") + if self.linear_attention_config.backend == "matmul": + return matmul_gdn( + *args, + policy=self.linear_attention_config, + state_qdq=quantize_state, + state_format=self._linear_attn_state_format, + chunk_size=chunk_size, + prefill_lengths=self._linear_attention_prefill_lengths, + **kwargs, + ) return _state_qdq_chunk_gated_delta_rule()( *args, chunk_size=chunk_size, - state_qdq=int(quantize_state), - state_qdq_block_v=64, + state_qdq=(2 if self._linear_attn_state_format == "int8" else 1) + if quantize_state + else 0, + state_qdq_block_v=self.gdn_state_qdq_block_v, w_quantizer=self.gdn_w_quantizer if quantize_w else None, **kwargs, ) diff --git a/modelopt/torch/quantization/plugins/kda.py b/modelopt/torch/quantization/plugins/kda.py new file mode 100644 index 00000000000..fbc4b1cf0ef --- /dev/null +++ b/modelopt/torch/quantization/plugins/kda.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""QAT integration for the optional FLA KimiDeltaAttention layer.""" + +import importlib.metadata +import inspect +from functools import lru_cache +from types import FunctionType + +from ..linear_attention.kda import matmul_kda +from ..nn import QuantModuleRegistry +from .custom import CUSTOM_MODEL_PLUGINS +from .linear_attention import _LinearAttentionQuantMixin + +__all__ = [] + + +@lru_cache(maxsize=8) +def _forward_signature(function): + return inspect.signature(function) + + +class _QuantKimiDeltaAttention(_LinearAttentionQuantMixin): + linear_attention_quantizer_names = ("kda_state_quantizer", "kda_w_quantizer") + + def validate_linear_attention(self): + super().validate_linear_attention() + if self.linear_attention_is_enabled and self.linear_attention_config.backend != "matmul": + raise ValueError("KDA numerical emulation requires backend='matmul'") + + def forward(self, *args, **kwargs): + self.validate_linear_attention() + if not self.linear_attention_is_enabled: + return super().forward(*args, **kwargs) + original = super().forward.__func__ + if self.linear_attention_config.decode is not None: + arguments = _forward_signature(original).bind(self, *args, **kwargs).arguments + if arguments.get("use_cache", False) or arguments.get("past_key_values") is not None: + raise NotImplementedError( + "Decode-aware QAT requires use_cache=False; use explicit numerical carry for continuation" + ) + # Bind this invocation so copied modules cannot retain another module's quantizers. + namespace = { + **original.__globals__, + "chunk_kda": self._quantized_chunk, + "fused_recurrent_kda": self._unsupported_recurrent, + } + forward = FunctionType( + original.__code__, + namespace, + original.__name__, + original.__defaults__, + original.__closure__, + ) + forward.__kwdefaults__ = original.__kwdefaults__ + return forward(self, *args, **kwargs) + + def _quantized_chunk(self, *args, **kwargs): + return matmul_kda( + *args, + policy=self.linear_attention_config, + state_qdq=self.kda_state_quantizer.is_enabled and self.kda_state_quantizer._if_quant, + state_format=self._linear_attn_state_format, + prefill_lengths=self._linear_attention_prefill_lengths, + **kwargs, + ) + + def _unsupported_recurrent(self, *args, **kwargs): + if self.linear_attention_config.decode is not None: + return self._quantized_chunk(*args, **kwargs) + raise NotImplementedError( + "KDA prefill QAT requires the chunk path; recurrent decode is not yet supported" + ) + + +def _register_fla_kda(model): + for module in model.modules(): + cls = type(module) + if cls.__module__ != "fla.layers.kda" or cls.__name__ != "KimiDeltaAttention": + continue + if cls in QuantModuleRegistry: + continue + if importlib.metadata.version("flash-linear-attention") != "0.5.1": + raise RuntimeError("KDA QAT is qualified with flash-linear-attention==0.5.1") + if importlib.metadata.version("fla-core") != "0.5.1": + raise RuntimeError("KDA QAT requires fla-core==0.5.1") + QuantModuleRegistry.register({cls: "fla_KimiDeltaAttention"})(_QuantKimiDeltaAttention) + + +CUSTOM_MODEL_PLUGINS.add(_register_fla_kda) diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py new file mode 100644 index 00000000000..5559c9f6606 --- /dev/null +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared module policy and checkpoint support for linear-attention QAT.""" + +from ..config import QuantizerAttributeConfig +from ..linear_attention.config import LinearAttentionConfig +from ..linear_attention.utils import validate_gdn_quantizer +from ..nn import QuantModule, TensorQuantizer + +__all__ = [] + + +class _LinearAttentionQuantMixin(QuantModule): + linear_attention_quantizer_names = ("gdn_state_quantizer", "gdn_w_quantizer") + + def _setup(self): + for name in self.linear_attention_quantizer_names: + self._register_temp_attribute( + name, TensorQuantizer(QuantizerAttributeConfig(enable=False)) + ) + self._register_temp_attribute("linear_attention_config", LinearAttentionConfig()) + self._register_temp_attribute("_linear_attention_prefill_lengths", None) + + @property + def _linear_attn_state(self): + return getattr(self, self.linear_attention_quantizer_names[0]) + + @property + def _linear_attn_state_format(self): + return "int8" if self._linear_attn_state.num_bits == 8 else "fp8_e4m3" + + @property + def _linear_attn_w(self): + return getattr(self, self.linear_attention_quantizer_names[1]) + + @property + def linear_attention_is_enabled(self): + """Whether an operand, state, or arithmetic policy changes the computation.""" + return ( + self._linear_attn_state.is_enabled + or self._linear_attn_w.is_enabled + or self.linear_attention_config.decode is not None + ) + + def validate_linear_attention(self): + """Validate quantizer contracts shared by GDN and KDA.""" + if self._linear_attn_state.is_enabled: + validate_gdn_quantizer( + self._linear_attn_state, + name=self.linear_attention_quantizer_names[0], + num_bits=((4, 3), 8), + ) + if self._linear_attn_state.axis != (0, 1): + raise ValueError( + f"{self.linear_attention_quantizer_names[0]} supports only axis=(0, 1) " + "with state.block_v tiling" + ) + decode = self.linear_attention_config.decode + if decode is not None and decode.state_codec == "int8_hadamard32": + if self._linear_attn_state_format != "int8": + raise ValueError("int8_hadamard32 requires INT8 state quantization") + if self._linear_attn_w.is_enabled: + validate_gdn_quantizer( + self._linear_attn_w, name=self.linear_attention_quantizer_names[1] + ) + if self.linear_attention_config.decode is not None: + raise ValueError("Decode's exact prefix does not support WY operand QDQ") + + def modelopt_post_restore(self, prefix=""): + """Validate the restored numerical policy and quantizers.""" + super().modelopt_post_restore(prefix) + self.validate_linear_attention() diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index ec942e0c24c..49195b44d29 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -61,7 +61,7 @@ from ..utils import sync_moe_expert_amax from ..utils.layerwise_calib import LayerActivationCollector from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear -from .gated_delta_net import GatedDeltaNetStateQuantMixin +from .gdn import GatedDeltaNetStateQuantMixin try: from megatron.core.extensions.transformer_engine import ( @@ -1122,9 +1122,7 @@ def forward_pre_attn_and_core_attn(self, *args, **kwargs): def validate_linear_attention(self): super().validate_linear_attention() - if self.config.context_parallel_size > 1 and ( - self.gdn_state_quantizer.is_enabled or self.gdn_w_quantizer.is_enabled - ): + if self.config.context_parallel_size > 1 and self.linear_attention_is_enabled: raise NotImplementedError("GDN QAT does not support Megatron context parallelism.") diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 20e083491aa..67513218137 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -50,7 +50,8 @@ def _fp8_eager(x, amax=None): epsilon = 1.0 / (1 << 24) zero_amax_mask = amax <= epsilon safe_amax = torch.where(zero_amax_mask, torch.ones_like(amax), amax) - scale = 448.0 / safe_amax + # Use direct division, matching the CUDA extension rather than reciprocal multiplication. + scale = torch.div(448.0, safe_amax) scale_inv = 1 / scale x = (x.to(torch.float32) * scale).clamp(min=-448.0, max=448.0) x = x.to(torch.float8_e4m3fn) diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index 79bc6065c24..ea115b6fe88 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -33,4 +33,5 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `mixer_mlp_nvfp4.yaml` | NVFP4 W4A4 on dense `*.mixer.{up,down}_proj` weight/input quantizers | | `attention_qkv_fp8.yaml` | FP8 E4M3 on attention q/k/v bmm and softmax quantizers | | `gdn_state_fp8_dynamic.yaml` | FP8 E4M3 dynamic fake quantization of the GatedDeltaNet recurrent state (per sequence, head and 64-column tile) at every kernel chunk boundary; needs `fla-core==0.5.1`, Triton, and SM89+ | +| `gdn_state_int8_dynamic.yaml` | INT8 quantizer unit; use the [complete recipe](../../../general/ptq/gdn_state_int8_dynamic.yaml) to enable INT8 state QDQ and Hadamard together | | `gdn_w_fp8_dynamic.yaml` | FP8 E4M3 dynamic (per token and head) fake quantization of the WY tensor `w` that multiplies the GatedDeltaNet state; requires identity STE, `fla-core==0.5.1`, and Triton | diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index cac4f6d80aa..e409bff7aa6 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -21,6 +21,10 @@ enable: false - quantizer_name: '*gdn_w_quantizer' enable: false + - quantizer_name: '*kda_state_quantizer' + enable: false + - quantizer_name: '*kda_w_quantizer' + enable: false - quantizer_name: '*block_sparse_moe.gate*' enable: false - quantizer_name: '*linear_attn.conv1d*' diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml index ed5768ee599..4d4b13278ba 100644 --- a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -14,10 +14,11 @@ # limitations under the License. # Dynamic per-tile FP8 E4M3 fake quantization of the GatedDeltaNet recurrent state. -# Recompute amax over each [K, 64] tile inside the kernel, independently per sequence and head +# Recompute amax over each [K, BV] tile inside the kernel, independently per sequence and head # of the [N, H, K, V] state (no calibration). ``axis: [0, 1]`` specifies sequence/head grouping; -# the kernel uses one scale per full-K by 64-value-column tile (two per head for V = 128). -# This is the only configuration the module supports. Applied at every kernel chunk boundary +# BV is ``gdn_state_qdq_block_v`` (default: the kernel's 64-column tile, +# two scales per head for V = 128; set 128 on the quantized module for one scale per head). This +# selects the FP8 state format. Applied at every kernel chunk boundary # (64 tokens) during calibration, QAT and QAD; needs fla-core==0.5.1 and Triton. # See ``modelopt.torch.quantization.plugins.gated_delta_net``. diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml new file mode 100644 index 00000000000..e86390cfef5 --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Quantizer-only INT8 unit; general/ptq/gdn_state_int8_dynamic pairs it with Hadamard. +# The execution policy selects state scale groups and chunk/token/replay boundaries. +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*gdn_state_quantizer' + cfg: + num_bits: 8 + unsigned: false + narrow_range: true + axis: [0, 1] + type: dynamic + pass_through_bwd: true diff --git a/modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml b/modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml new file mode 100644 index 00000000000..e960282121b --- /dev/null +++ b/modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + gdn_state_int8_dynamic: configs/ptq/units/gdn_state_int8_dynamic + +metadata: + recipe_type: ptq + description: >- + Dynamic INT8 GDN state QDQ with 32-value Hadamard rotation during decode. + Prefill state remains unquantized; supply per-sequence prefix lengths through + linear_attention_training_phase. No calibration is required. +quantize: + algorithm: + quant_cfg: + - $import: base_disable_all + - $import: gdn_state_int8_dynamic + linear_attention: + - module_name: '*' + cfg: + backend: matmul + decode: + mode: token + state_codec: int8_hadamard32 + prefill_state_qdq: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 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/tests/_test_utils/torch/quantization/linear_attention_reference.py b/tests/_test_utils/torch/quantization/linear_attention_reference.py index 65ea183a318..57d7448bcb1 100644 --- a/tests/_test_utils/torch/quantization/linear_attention_reference.py +++ b/tests/_test_utils/torch/quantization/linear_attention_reference.py @@ -25,23 +25,43 @@ import torch -__all__ = ["chunk_gdn_reference", "recurrent_delta_rule_reference", "state_fp8_qdq_reference"] +__all__ = [ + "chunk_gdn_reference", + "chunk_kda_reference", + "recurrent_delta_rule_reference", + "state_fp8_qdq_reference", + "state_qdq_reference", +] def state_fp8_qdq_reference(state: torch.Tensor, block_v: int = 64) -> torch.Tensor: """Dynamic E4M3 QDQ with identity STE and one scale per ``[Dk, block_v]`` tile. ``state`` is in key-first layout ``[..., Dk, Dv]``. Scales and rounding do not - contribute derivatives. The zero tile uses scale one; partial value tiles are valid. + contribute derivatives. FP8 uses TensorQuantizer small-range handling; partial tiles are valid. """ + return state_qdq_reference(state, block_v) + + +def state_qdq_reference(state: torch.Tensor, block_v: int = 64, state_format: str = "fp8_e4m3"): + """Dynamic state-tile QDQ with detached scales and identity STE.""" + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") if block_v not in (16, 32, 64, 128): raise ValueError("block_v must be 16, 32, 64, or 128") with torch.no_grad(): rounded = [] for tile in state.float().split(block_v, dim=-1): amax = tile.abs().amax(dim=(-2, -1), keepdim=True) - scale = torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)) - rounded.append((tile / scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() * scale) + if state_format == "int8": + scale = torch.where(amax > 0, amax / 127.0, torch.ones_like(amax)) + codes = (tile / scale).clamp(-127, 127).round() + else: + safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) + quant_scale = torch.div(448.0, safe_amax) + scale = quant_scale.reciprocal() + codes = (tile * quant_scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() + rounded.append(codes * scale) quantized = torch.cat(rounded, dim=-1).to(state.dtype) return state + (quantized - state).detach() @@ -140,34 +160,61 @@ def chunk_gdn_reference( state_v_first: bool = False, state_qdq: bool = False, state_qdq_block_v: int = 64, + state_format: str = "fp8_e4m3", w_quantizer: Callable[[torch.Tensor], torch.Tensor] | None = None, + matmul: Callable[[str, torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + arithmetic: Callable[[str, torch.Tensor], torch.Tensor] | None = None, + inverse_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Exact GDN chunk algebra with optional state/W fake quantization. The solve is a unit-lower triangular solve. ``w_quantizer`` sees the complete materialized ``[B,T,Hv,Dk]`` WY operand once, with its own autograd semantics. State QDQ occurs on the initial state and each chunk's final state, after readout. + Optional callbacks expose the eight named matmuls and elementwise boundaries. + Matmul callbacks receive conventional ``lhs @ rhs`` shapes, without padding. """ if g.ndim != 3 or chunk_size <= 0: raise ValueError("chunk_gdn_reference requires scalar GDN gates and positive chunk_size") q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) scale = q.shape[-1] ** -0.5 if scale is None else scale + mm = matmul if matmul is not None else lambda name, lhs, rhs: lhs @ rhs + arithmetic = arithmetic if arithmetic is not None else lambda name, x: x chunks, all_w = [], [] for n, (b, start, end) in enumerate(sequences): for lo in range(start, end, chunk_size): hi = min(lo + chunk_size, end) qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) - gc = g[b, lo:hi].transpose(0, 1).cumsum(-1) + gc = arithmetic("gate_prefix", g[b, lo:hi].transpose(0, 1).cumsum(-1)) bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) # Mask before exp: upper-triangle positive differences can overflow for long decay. causal = torch.ones(hi - lo, hi - lo, device=q.device, dtype=torch.bool).tril() - decay = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() - gram = kc @ kc.transpose(-1, -2) - lower = (bc * gram * decay).tril(-1) + decay = arithmetic( + "gate_exp", (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() + ) + if matmul is None: + lower = (bc * (kc @ kc.transpose(-1, -2)) * decay).tril(-1) + else: + lower = (mm("key_interaction", bc * kc, kc.transpose(-1, -2)) * decay).tril(-1) matrix = lower + torch.eye(hi - lo, device=q.device, dtype=q.dtype) - rhs = torch.cat((bc * vc, bc * kc * gc.exp().unsqueeze(-1)), dim=-1) - solved = torch.linalg.solve_triangular(matrix, rhs, upper=False, unitriangular=True) - u, w = solved.split((v.shape[-1], k.shape[-1]), dim=-1) + gate = arithmetic("gate_exp", gc.exp()).unsqueeze(-1) + if matmul is None and inverse_fn is None: + rhs = torch.cat((bc * vc, bc * kc * gate), dim=-1) + solved = torch.linalg.solve_triangular(matrix, rhs, upper=False, unitriangular=True) + u, w = solved.split((v.shape[-1], k.shape[-1]), dim=-1) + else: + inverse = ( + inverse_fn(lower) + if inverse_fn is not None + else torch.linalg.solve_triangular( + matrix, + torch.eye(hi - lo, device=q.device, dtype=q.dtype).expand_as(matrix), + upper=False, + unitriangular=True, + ) + ) + u = mm("wy_value", inverse, bc * vc) + w = mm("wy_key", inverse, bc * kc * gate) chunks.append((n, qc, kc, gc, decay, u)) all_w.append(w.transpose(0, 1)) w = torch.cat(all_w).reshape(q.shape) @@ -179,21 +226,27 @@ def chunk_gdn_reference( if n != previous_n: state = states[n] if state_qdq: - state = state_fp8_qdq_reference(state, state_qdq_block_v) + state = state_qdq_reference(state, state_qdq_block_v, state_format) previous_n = n length = qc.shape[1] wc = w[offset : offset + length].transpose(0, 1) offset += length - updated_values = u - wc @ state # state_read - local_scores = ((qc * scale) @ kc.transpose(-1, -2) * decay).tril() - output = (qc * (scale * gc.exp()).unsqueeze(-1)) @ state - output = output + local_scores @ updated_values # local_readout + updated_values = arithmetic("value_residual", u - mm("state_read", wc, state)) + local_scores = (mm("output_score", qc * scale, kc.transpose(-1, -2)) * decay).tril() + output = mm( + "output_state", qc * (scale * arithmetic("gate_exp", gc.exp())).unsqueeze(-1), state + ) + output = arithmetic("output_add", output + mm("output_value", local_scores, updated_values)) outputs.append(output.transpose(0, 1)) - weighted_keys = kc * (gc[..., -1:] - gc).exp().unsqueeze(-1) - state = state * gc[..., -1].exp()[:, None, None] - state = state + weighted_keys.transpose(-1, -2) @ updated_values # state_update + weighted_keys = kc * arithmetic("gate_exp", (gc[..., -1:] - gc).exp()).unsqueeze(-1) + state = arithmetic( + "state_decay", state * arithmetic("gate_exp", gc[..., -1].exp())[:, None, None] + ) + state = arithmetic( + "state_add", state + mm("state_update", weighted_keys.transpose(-1, -2), updated_values) + ) if state_qdq: - state = state_fp8_qdq_reference(state, state_qdq_block_v) + state = state_qdq_reference(state, state_qdq_block_v, state_format) if len(finals) <= n: finals.append(state) else: @@ -201,3 +254,93 @@ def chunk_gdn_reference( output = torch.cat(outputs).reshape(*v.shape) final = torch.stack(finals) return output, final.transpose(-1, -2) if state_v_first else final + + +def chunk_kda_reference( + q, + k, + v, + g, + beta, + *, + chunk_size=64, + scale=None, + initial_state=None, + cu_seqlens=None, + state_v_first=False, + state_qdq=False, + state_qdq_block_v=64, + state_format="fp8_e4m3", + w_quantizer=None, + matmul=None, + inverse_fn=None, +): + """KDA chunk oracle with causal per-channel decays and optional operand QDQ. + + The eight-site callback receives conventional three-dimensional matmul + operands. Interaction sites are invoked one query row at a time with decay + already applied to the right-hand keys. This avoids inverse-prefix factors. + """ + if g.ndim != 4 or chunk_size <= 0: + raise ValueError("chunk_kda_reference requires per-key gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + mm = matmul if matmul is not None else lambda name, lhs, rhs: lhs @ rhs + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + pieces = [] + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + prefix = g[b, lo:hi].transpose(0, 1).cumsum(-2) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + gate = prefix.exp() + lower_rows, score_rows = [], [] + for row in range(hi - lo): + decayed_keys = ( + kc[..., : row + 1, :] + * (prefix[..., row : row + 1, :] - prefix[..., : row + 1, :]).exp() + ) + right = decayed_keys.transpose(-1, -2) + left = kc[..., row : row + 1, :] * bc[..., row : row + 1, :] + interaction = mm("key_interaction", left, right) + score = mm("output_score", qc[..., row : row + 1, :] * scale, right) + # Zero padding follows the computation and never contributes a scale statistic. + padding = hi - lo - row - 1 + lower_rows.append(torch.nn.functional.pad(interaction, (0, padding))) + score_rows.append(torch.nn.functional.pad(score, (0, padding))) + lower = torch.cat(lower_rows, dim=-2).tril(-1) + scores = torch.cat(score_rows, dim=-2) + identity = torch.eye(hi - lo, device=q.device, dtype=q.dtype).expand_as(lower) + inverse = ( + inverse_fn(lower) + if inverse_fn is not None + else torch.linalg.solve_triangular( + identity + lower, identity, upper=False, unitriangular=True + ) + ) + u = mm("wy_value", inverse, bc * vc) + w = mm("wy_key", inverse, bc * kc * gate) + if w_quantizer is not None: + w = w_quantizer(w) + updated = u - mm("state_read", w, state) + pieces.append( + ( + mm("output_state", qc * scale * gate, state) + + mm("output_value", scores, updated) + ).transpose(0, 1) + ) + weighted_keys = kc * (prefix[..., -1:, :] - prefix).exp() + state = state * gate[..., -1, :, None] + mm( + "state_update", weighted_keys.transpose(-1, -2), updated + ) + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + outputs.append(torch.cat(pieces)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py index 9a11c454932..84b344ba224 100644 --- a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py @@ -28,7 +28,7 @@ from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.nn import TensorQuantizer -from modelopt.torch.quantization.plugins.gated_delta_net import GatedDeltaNetStateQuantMixin +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 @@ -38,7 +38,7 @@ ) DTYPES = (torch.float32, torch.bfloat16) -QDQ_MODES = ((False, False), (False, True), (True, False), (True, True)) +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)) @@ -115,7 +115,7 @@ def _compile_gdn_kernels(): 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 and not supports_state_qdq: + if state_qdq == 1 and not supports_state_qdq: continue values_and_grads( chunk_gated_delta_rule, @@ -133,9 +133,15 @@ def _compile_gdn_kernels(): args, state = make_inputs() quantizer = w_quantizer() - values_and_grads( - chunk_gated_delta_rule, args, state, output_final_state=True, w_quantizer=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) @@ -154,16 +160,19 @@ def _compile_gdn_kernels(): allow_neg_eigval=True, ) torch.autograd.grad(output.square().sum() + final.square().sum(), (*args, state, a_log, bias)) - if supports_state_qdq: + 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=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( @@ -193,7 +202,7 @@ def test_disabled_matches_upstream_forward_and_backward(dtype): @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: + 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 @@ -208,7 +217,13 @@ def test_numerical_surrogate_reference(state_qdq, quantize_w, packed, state_v_fi } 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, **kwargs) + 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 ) @@ -216,7 +231,8 @@ def test_numerical_surrogate_reference(state_qdq, quantize_w, packed, state_v_fi compare(actual[1], expected[1], 0.05 if dtype == torch.bfloat16 else 0.02) -def test_w_qdq_saved_once_and_activation_checkpoint_parity(): +@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 = [] @@ -224,7 +240,11 @@ def test_w_qdq_saved_once_and_activation_checkpoint_parity(): def fn(*x): return chunk_gated_delta_rule( - *x[:5], initial_state=x[5], output_final_state=True, w_quantizer=quantizer + *x[:5], + initial_state=x[5], + output_final_state=True, + w_quantizer=quantizer, + state_qdq=state_qdq, ) output = fn(*args, state) @@ -252,12 +272,20 @@ def test_zero_initial_state_and_output_only_training_loss(axis): @pytest.mark.parametrize("block_v", [16, 64, 128]) -def test_state_scale_tile_forward_and_backward(block_v): - require_state_qdq() +@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": True, "state_qdq_block_v": block_v} + 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, **kwargs) + 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 ) diff --git a/tests/gpu/torch/quantization/plugins/test_kda.py b/tests/gpu/torch/quantization/plugins/test_kda.py new file mode 100644 index 00000000000..ed67787de54 --- /dev/null +++ b/tests/gpu/torch/quantization/plugins/test_kda.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from contextlib import nullcontext + +import pytest +import torch +from torch.utils.checkpoint import checkpoint + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + linear_attention_training_phase, +) +from modelopt.torch.quantization.nn import TensorQuantizer + +KimiDeltaAttention = pytest.importorskip("fla.layers.kda").KimiDeltaAttention + + +def _layer(head_dim=16): + return ( + KimiDeltaAttention(hidden_size=32, head_dim=head_dim, num_heads=2, use_short_conv=True) + .cuda() + .train() + ) + + +def _forward(model, hidden): + policy = getattr(model, "linear_attention_config", None) + phase = ( + linear_attention_training_phase(model, [min(31, hidden.shape[1])]) + if policy is not None and policy.decode is not None + else nullcontext() + ) + with phase, torch.autocast("cuda", dtype=torch.bfloat16): + return model(hidden)[0] + + +@pytest.mark.parametrize( + "mode", + [ + "decode_token", + "decode_replay", + "decode_decay", + "decode_token_int8", + "decode_replay_int8", + "decode_replay_hadamard_int8", + ], +) +@pytest.mark.timeout(180) +def test_fla_layer_qat_restore_and_optimizer(tmp_path, mode): + torch.manual_seed(73) + head_dim = 32 if "hadamard" in mode else 16 + model = _layer(head_dim) + hidden = torch.randn(1, 73, 32, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + baseline = _forward(model, hidden) + attributes = {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)} + cfg = { + "quant_cfg": [{"quantizer_name": "*", "enable": False}], + "algorithm": None, + "linear_attention": [{"module_name": "*", "cfg": {"backend": "matmul"}}], + } + state_attributes = {**attributes, "axis": (0, 1)} + if mode.endswith("int8"): + state_attributes.update(num_bits=8, unsigned=False, narrow_range=True) + if mode.startswith("decode"): + cfg["quant_cfg"].append({"quantizer_name": "*kda_state_quantizer", "cfg": state_attributes}) + policy_cfg = cfg["linear_attention"][0]["cfg"] + policy_cfg["state"] = {"block_v": 16} + policy_cfg["decode"] = { + "mode": "replay" if mode.startswith("decode_replay") else "token", + "implementation": "torch", + } + if mode.startswith("decode_replay"): + policy_cfg["decode"]["replay"] = {"window": 5} + if mode == "decode_decay": + policy_cfg["decode"]["decay_log_step"] = 1 / 256 + if "hadamard" in mode: + policy_cfg["state"]["block_v"] = 32 + policy_cfg["decode"]["state_codec"] = "int8_hadamard32" + mtq.quantize(model, cfg) + output = _forward(model, hidden) + assert torch.isfinite(output).all() + enabled = [q for q in model.modules() if isinstance(q, TensorQuantizer) and q.is_enabled] + policy = model.linear_attention_config + for q in enabled: + q.disable() + model.linear_attention_config = LinearAttentionConfig() + with torch.no_grad(): + torch.testing.assert_close(_forward(model, hidden), baseline, rtol=0, atol=0) + for q in enabled: + q.enable() + model.linear_attention_config = policy + checkpoint = tmp_path / "model.pth" + mto.save(model, checkpoint) + restored = mto.restore(_layer(head_dim), checkpoint) + actual = _forward(restored, hidden) + torch.testing.assert_close(actual, output, rtol=0, atol=0) + assert restored.linear_attention_config == policy + actual.float().square().mean().backward() + assert restored.q_proj.weight.grad is not None + for parameter in restored.parameters(): + assert parameter.grad is not None + assert torch.isfinite(parameter.grad).all() + before = restored.q_proj.weight.detach().clone() + torch.optim.SGD(restored.parameters(), lr=0.1).step() + assert not torch.equal(before, restored.q_proj.weight) + restored.eval() + if mode.startswith("decode"): + with ( + linear_attention_training_phase(restored, [0]), + torch.autocast("cuda", dtype=torch.bfloat16), + ): + assert torch.isfinite(restored(hidden[:, :1])[0]).all() + with pytest.raises(NotImplementedError, match="use_cache=False"): + restored(hidden[:, :1], use_cache=True) + with ( + pytest.raises(ValueError, match="explicit"), + torch.autocast("cuda", dtype=torch.bfloat16), + ): + restored(hidden) + else: + with pytest.raises(NotImplementedError, match="chunk path"): + _forward(restored, hidden[:, :1]) + + +@pytest.mark.parametrize("hadamard", [False, True]) +def test_decode_phase_spans_activation_checkpoint_backward(hadamard): + model = _layer(32 if hadamard else 16) + quant_cfg = [{"quantizer_name": "*", "enable": False}] + if hadamard: + quant_cfg.append( + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": 8, + "type": "dynamic", + "axis": (0, 1), + "unsigned": False, + "narrow_range": True, + }, + } + ) + mtq.quantize( + model, + { + "quant_cfg": quant_cfg, + "algorithm": None, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "torch", + "state_codec": "int8_hadamard32" if hadamard else "tile", + }, + }, + } + ], + }, + ) + hidden = torch.randn(1, 73, 32, device="cuda", requires_grad=True) + with linear_attention_training_phase(model, [31]): + with torch.autocast("cuda", dtype=torch.bfloat16): + actual = checkpoint(lambda x: model(x)[0], hidden, use_reentrant=False) + actual.square().mean().backward() + assert model._linear_attention_prefill_lengths is None + assert torch.isfinite(hidden.grad).all() + assert torch.count_nonzero(model.q_proj.weight.grad) > 0 + with linear_attention_training_phase(model, [0]): + with linear_attention_training_phase(model, [73]): + assert model._linear_attention_prefill_lengths == (73,) + assert model._linear_attention_prefill_lengths == (0,) + assert model._linear_attention_prefill_lengths is None diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index 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 index c46959c4d8d..94ad1aaedf1 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import nullcontext + import pytest import torch from _test_utils.torch.megatron.models import get_mcore_gpt_model @@ -23,11 +25,16 @@ ) import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + linear_attention_training_phase, +) +from modelopt.torch.quantization.nn import TensorQuantizer pytest.importorskip("fla") # Megatron-Core GatedDeltaNet and the state QDQ kernel need fla GatedDeltaNet = pytest.importorskip("megatron.core.ssm.gated_delta_net").GatedDeltaNet -from modelopt.torch.quantization.plugins.gated_delta_net import _state_qdq_chunk_gated_delta_rule +from modelopt.torch.quantization.plugins.gdn import _state_qdq_chunk_gated_delta_rule from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet try: @@ -78,15 +85,50 @@ def _gdn_config(sites): for site in sites ], "algorithm": None, + "linear_attention": [ + {"module_name": "decoder.layers.*.self_attention", "cfg": {"state": {"block_v": 16}}} + ], } +def _config_for_mode(mode): + sites = ("state", "w") if mode == "both" else ((mode,) if mode in ("state", "w") else ()) + if mode.startswith("decode"): + sites = ("state",) + if mode == "state_int8": + sites = ("state",) + cfg = _gdn_config(sites) + if mode.endswith("int8"): + cfg["quant_cfg"][1]["cfg"].update(num_bits=8, unsigned=False, narrow_range=True) + if mode.startswith("decode"): + policy = cfg["linear_attention"][0]["cfg"] + policy.update(backend="matmul", decode={"implementation": "triton"}) + if mode.startswith("decode_replay"): + policy["decode"].update(mode="replay", replay={"window": 8}) + return cfg, sites + + +def _gdn_forward(model): + original_forward = get_forward(model) + + def forward(m): + enabled = any( + getattr(getattr(layer, "linear_attention_config", None), "decode", None) is not None + for layer in m.modules() + ) + with linear_attention_training_phase(m, [64, 64]) if enabled else nullcontext(): + return original_forward(m) + + return forward + + def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): initialize_for_megatron( tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED ) model = _make_model(size) - forward = get_forward(model) + forward = _gdn_forward(model) + outputs = [] handles = [ module.register_forward_hook( @@ -100,13 +142,18 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): gdn_ref = outputs.copy() outputs.clear() - sites = ("state", "w") if mode == "both" else (mode,) - mtq.quantize(model, _gdn_config(sites)) + cfg, sites = _config_for_mode(mode) + mtq.quantize(model, cfg) gdn_modules = [m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)] assert gdn_modules, "no GatedDeltaNet layer was wrapped" for module in gdn_modules: + assert module.gdn_state_qdq_block_v == 16 assert module.gdn_state_quantizer.is_enabled == ("state" in sites) assert module.gdn_w_quantizer.is_enabled == ("w" in sites) + # Checkpointing must save the resolved policy, including edits after conversion. + module.linear_attention_config = LinearAttentionConfig( + **{**module.linear_attention_config.model_dump(), "state": {"block_v": 32}} + ) with torch.no_grad(): loss_quant = forward(model) @@ -119,20 +166,34 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): for handle in handles: handle.remove() - for site in sites: - mtq.disable_quantizer(model, f"*gdn_{site}_quantizer") + enabled = [ + q + for m in gdn_modules + for q in m.modules() + if isinstance(q, TensorQuantizer) and q.is_enabled + ] + policies = [m.linear_attention_config for m in gdn_modules] + for q in enabled: + q.disable() + for module in gdn_modules: + module.linear_attention_config = LinearAttentionConfig() with torch.no_grad(): torch.testing.assert_close(forward(model), loss_ref, rtol=1e-4, atol=1e-4) - for site in sites: - mtq.enable_quantizer(model, f"*gdn_{site}_quantizer") + for q in enabled: + q.enable() + for module, policy in zip(gdn_modules, policies): + module.linear_attention_config = policy restored = _make_model(size) sharded_state_dict_test_helper(checkpoint_path, model, restored, forward) restored_gdn = [m for m in restored.modules() if isinstance(m, _QuantGatedDeltaNet)] assert len(restored_gdn) == len(gdn_modules) - for module in restored_gdn: + for module, original in zip(restored_gdn, gdn_modules): + assert module.linear_attention_config == original.linear_attention_config assert module.gdn_state_quantizer.is_enabled == ("state" in sites) assert module.gdn_w_quantizer.is_enabled == ("w" in sites) + assert module._linear_attn_state_format == original._linear_attn_state_format + assert module.gdn_state_qdq_block_v == 32 assert module.in_proj.weight.grad is not None assert torch.isfinite(module.in_proj.weight.grad).all() @@ -148,11 +209,13 @@ def _compile_gdn_qat_kernels(rank, size, mode): tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED ) model = _make_model(size) - forward = get_forward(model) + forward = _gdn_forward(model) with torch.no_grad(): forward(model) - sites = ("state", "w") if mode == "both" else (mode,) - mtq.quantize(model, _gdn_config(sites)) + cfg, _ = _config_for_mode(mode) + # The functional case switches state tiles before running inference. + cfg["linear_attention"][0]["cfg"]["state"]["block_v"] = 32 + mtq.quantize(model, cfg) with torch.no_grad(): forward(model) model.train() @@ -163,7 +226,7 @@ def _compile_gdn_qat_kernels(rank, size, mode): @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 != "w" and torch.cuda.get_device_capability() < (8, 9): + 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) @@ -171,7 +234,19 @@ def compiled_gdn_workers(request, tp_size, mode): @pytest.mark.parametrize("tp_size", [1, 2]) -@pytest.mark.parametrize("mode", ["state", "w", "both"]) +@pytest.mark.parametrize( + "mode", + [ + "state", + "w", + "both", + "decode_token", + "decode_replay", + "state_int8", + "decode_token_int8", + "decode_replay_int8", + ], +) def test_gdn_qat_and_sharded_restore(compiled_gdn_workers, tmp_path, mode): """Train through state/W QDQ after a Megatron distributed-checkpoint round trip.""" compiled_gdn_workers.run(_test_gdn_qat_helper, mode, tmp_path) @@ -181,10 +256,21 @@ 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, _gdn_config((mode,))) + mtq.quantize(model, _config_for_mode(mode)[0]) -@pytest.mark.parametrize("mode", ["state", "w"]) +@pytest.mark.parametrize( + "mode", + [ + "state", + "w", + "decode_token", + "decode_replay", + "state_int8", + "decode_token_int8", + "decode_replay_int8", + ], +) def test_gdn_context_parallel_rejected(dist_workers_size_2, mode): """Reject unqualified Megatron CP even when it does not pass an FLA CP context.""" dist_workers_size_2.run(_test_gdn_context_parallel_helper, mode) diff --git a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py b/tests/unit/torch/quantization/plugins/test_gdn.py similarity index 74% rename from tests/unit/torch/quantization/plugins/test_gated_delta_net.py rename to tests/unit/torch/quantization/plugins/test_gdn.py index 3793edcb189..aed7f8caf42 100644 --- a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -22,9 +22,19 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.nn import QuantModuleRegistry -from modelopt.torch.quantization.plugins import gated_delta_net -from modelopt.torch.quantization.plugins.gated_delta_net import GatedDeltaNetStateQuantMixin +from modelopt.torch.quantization.plugins import gdn +from modelopt.torch.quantization.plugins.gdn import GatedDeltaNetStateQuantMixin + +GDN_STATE_INT8_DYNAMIC = { + "num_bits": 8, + "unsigned": False, + "narrow_range": True, + "type": "dynamic", + "axis": (0, 1), +} GDN_STATE_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"} @@ -36,9 +46,7 @@ def chunk_gated_delta_rule(q, k, v, g, beta, **kwargs): @pytest.fixture(autouse=True) def mock_fla_kernel(monkeypatch): - monkeypatch.setattr( - gated_delta_net, "_fla_chunk_gated_delta_rule", lambda: chunk_gated_delta_rule - ) + monkeypatch.setattr(gdn, "_fla_chunk_gated_delta_rule", lambda: chunk_gated_delta_rule) class TinyGatedDeltaNet(nn.Module): @@ -106,7 +114,12 @@ def test_dynamic_export_removes_linear_attention_attributes(): model = _QuantTinyGatedDeltaNet.convert(TinyGatedDeltaNet()) model.export() assert type(model) is TinyGatedDeltaNet - for name in ("gdn_state_quantizer", "gdn_w_quantizer"): + for name in ( + "gdn_state_quantizer", + "gdn_w_quantizer", + "linear_attention_config", + "_linear_attention_prefill_lengths", + ): assert not hasattr(model, name) @@ -124,27 +137,29 @@ def test_disabled_state_quantizer_calls_original_kernel(): assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) @pytest.mark.parametrize("partial_kernel", [False, True]) -def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, partial_kernel): +def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, state_format, partial_kernel): calls = [] def fake_state_qdq_kernel(*args, **kwargs): calls.append(kwargs) return chunk_gated_delta_rule(*args) - monkeypatch.setattr( - gated_delta_net, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel - ) + monkeypatch.setattr(gdn, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel) model = TinyGatedDeltaNet() if partial_kernel: model.gated_delta_rule = partial(chunk_gated_delta_rule, output_final_state=True) x = torch.randn(2, 8, 3, 4) - mtq.quantize(model, quant_cfg(), lambda m: m(x)) + cfg = quant_cfg() + if state_format == "int8": + cfg["quant_cfg"][1]["cfg"] = GDN_STATE_INT8_DYNAMIC + mtq.quantize(model, cfg, lambda m: m(x)) model(x) assert calls and calls[-1] == { "chunk_size": 64, - "state_qdq": 1, + "state_qdq": 2 if state_format == "int8" else 1, "state_qdq_block_v": 64, "w_quantizer": None, **({"output_final_state": True} if partial_kernel else {}), @@ -167,9 +182,7 @@ def fake_state_qdq_kernel(*args, **kwargs): calls.append(kwargs) return chunk_gated_delta_rule(*args) - monkeypatch.setattr( - gated_delta_net, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel - ) + monkeypatch.setattr(gdn, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel) model = TinyGatedDeltaNet() x = torch.randn(2, 8, 3, 4) mtq.quantize(model, quant_cfg(state=state, w=True), lambda m: m(x)) @@ -214,7 +227,13 @@ def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): model = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) cfg = quant_cfg(state=state, w=w) cfg["algorithm"] = None + cfg["linear_attention"] = [ + {"module_name": "*", "cfg": {"state": {"block_v": 128}}}, + {"module_name": "0", "cfg": {"state": {"block_v": 32}}}, + ] mtq.quantize(model, cfg) + assert model[0].gdn_state_qdq_block_v == 32 + model[0].linear_attention_config.state.block_v = 16 if w: # Save the actual quantizer settings, including edits after conversion. model[0].gdn_w_quantizer.axis = None @@ -222,6 +241,8 @@ def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): mto.save(model, path) restored = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) mto.restore(restored, path) + assert restored[0].linear_attention_config == model[0].linear_attention_config + assert restored[0].gdn_state_qdq_block_v == 16 for name, enabled in (("gdn_state_quantizer", state), ("gdn_w_quantizer", w)): original = getattr(model[0], name) quantizer = getattr(restored[0], name) @@ -241,7 +262,12 @@ def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): cfg = quant_cfg(state=False, w=True) cfg["algorithm"] = None + cfg["linear_attention"] = [{"module_name": "", "cfg": {"state": {"block_v": 32}}}] + mtq.quantize(model, cfg) + assert model.gdn_state_qdq_block_v == 32 + cfg["linear_attention"].append({"module_name": "", "cfg": {}}) mtq.quantize(model, cfg) + assert model.gdn_state_qdq_block_v == 64 assert not model.gdn_state_quantizer.is_enabled assert model.gdn_w_quantizer.is_enabled @@ -256,7 +282,9 @@ def test_restore_legacy_gdn_without_new_quantizer_handles(): model = mtq.quantize(TinyGatedDeltaNet(), config) state = mto.modelopt_state(model) for _, mode_state in state["modelopt_state_dict"]: + mode_state["config"].pop("linear_attention", None) metadata = mode_state["metadata"] + metadata.pop("linear_attention", None) for name in ("gdn_state_quantizer", "gdn_w_quantizer"): metadata["quantizer_state"].pop(name, None) restored = TinyGatedDeltaNet() @@ -274,3 +302,37 @@ def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): ) assert not model.gdn_state_quantizer.is_enabled assert not model.gdn_w_quantizer.is_enabled + + +def test_int8_state_conversion_and_checkpoint(tmp_path): + cfg = load_recipe("general/ptq/gdn_state_int8_dynamic").quantize + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + assert model.gdn_state_quantizer.is_enabled + assert model._linear_attn_state_format == "int8" + assert model.linear_attention_config.backend == "matmul" + assert model.linear_attention_config.decode.state_codec == "int8_hadamard32" + assert not model.linear_attention_config.decode.prefill_state_qdq + assert not model.proj.weight_quantizer.is_enabled + path = tmp_path / "int8.pt" + mto.save(model, path) + restored = mto.restore(TinyGatedDeltaNet(), path) + assert restored.linear_attention_config == model.linear_attention_config + assert restored.gdn_state_quantizer.is_enabled + assert restored._linear_attn_state_format == "int8" + assert restored.gdn_state_quantizer.narrow_range + assert not restored.gdn_state_quantizer.unsigned + assert not restored.gdn_w_quantizer.is_enabled + + +def test_policy_rejects_unmatched_and_unimplemented_modes(): + with pytest.raises(ValueError, match="matches no supported"): + mtq.quantize( + nn.Linear(4, 4), {"linear_attention": [{"module_name": "*"}], "algorithm": None} + ) + for policy in ( + {"chunk_size": 32}, + {"solve": {"method": "neumann"}}, + {"state": {"mode": "token"}}, + ): + with pytest.raises(ValueError): + QuantizeConfig(linear_attention=[{"module_name": "*", "cfg": policy}]) diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py new file mode 100644 index 00000000000..e0ced168024 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import itertools + +import pytest +import torch +import torch.nn.functional as F +from _test_utils.torch.quantization.linear_attention_reference import recurrent_delta_rule_reference + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + matmul_gdn, + matmul_kda, +) +from modelopt.torch.quantization.linear_attention.config import LinearAttentionDecodeConfig +from modelopt.torch.quantization.linear_attention.decode import _encode, recurrent_decode_reference +from modelopt.torch.quantization.linear_attention.utils import _state_qdq +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.mark.parametrize("state", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_fp8_qdq_matches_tensor_quantizer(state, dtype): + torch.manual_seed(762) + amplitudes = torch.tensor([0.0, 1e-8, 2**-24, 2**-23, 1.0, 731.173], dtype=dtype) + value = torch.rand((2, 6, 5, 19) if state else (2, 6, 19), dtype=dtype) * 2 - 1 + value *= amplitudes.reshape(1, 6, *([1] * (value.ndim - 2))) + value.requires_grad_() + quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=(4, 3), type="dynamic", axis=(0, 1), pass_through_bwd=True + ) + ) + expected = torch.cat( + [ + quantizer(tile.flatten(-2) if state else tile).reshape_as(tile) + for tile in value.split(16, -1) + ], + -1, + ) + encoded = _encode(value, True, 16, state=state) + results = [encoded.values] + if state: + results.append(_state_qdq(value, 16)) + probe = torch.randn_like(value) + (expected_gradient,) = torch.autograd.grad((expected * probe).sum(), value) + for result in results: + torch.testing.assert_close(result, expected, rtol=0, atol=0) + (gradient,) = torch.autograd.grad((result * probe).sum(), value) + torch.testing.assert_close(gradient, expected_gradient, rtol=0, atol=0) + assert encoded.scales.shape == (2, 6, 2) + assert not encoded.scales.requires_grad + + +def _inputs(kda=True, length=73): + torch.manual_seed(193) + q, k = [ + F.normalize(torch.randn(length, 2, 16, dtype=torch.float64), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(length, 2, 11, dtype=torch.float64, requires_grad=True) + g = ( + -torch.rand((length, 2, 16) if kda else (length, 2), dtype=torch.float64) * 0.03 + ).requires_grad_() + beta = (torch.rand(length, 2, dtype=torch.float64) * 0.4).requires_grad_() + state = (torch.randn(2, 16, 11, dtype=torch.float64) * 0.1).requires_grad_() + return (q, k, v, g, beta), state + + +def _values_and_grads(output, state, args, initial): + loss = output.square().sum() + state.square().sum() + return output, state, *torch.autograd.grad(loss, (*args, initial), retain_graph=True) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("mode", ["token", "replay"]) +def test_unquantized_decode_recurrence_and_gradients(kda, mode): + args, state = _inputs(kda) + cfg = LinearAttentionDecodeConfig( + mode=mode, replay={"window": 8, "factor_qdq": False} if mode == "replay" else None + ) + output, carry = recurrent_decode_reference(*args, config=cfg, initial_state=state) + exact, final = recurrent_delta_rule_reference( + *(x.unsqueeze(0) for x in args), initial_state=state.unsqueeze(0) + ) + for a, b in zip( + _values_and_grads(output, carry.reconstruct(), args, state), + _values_and_grads(exact[0], final[0], args, state), + ): + torch.testing.assert_close(a, b, rtol=2e-11, atol=1e-12) + assert carry.position == 73 + assert carry.cursor == (1 if mode == "replay" else 0) + + +@pytest.mark.parametrize("mode", ["token", "replay", "reencode"]) +@pytest.mark.parametrize("readout", ["working", "stored"]) +def test_quantized_carry_continuation_metadata_and_gradients(mode, readout): + args, state = _inputs() + cfg = LinearAttentionDecodeConfig( + mode="token" if mode == "token" else "replay", + readout=readout, + replay={"window": 8, "encoding": "reencode" if mode == "reencode" else "once"} + if mode != "token" + else None, + ) + expected, final = recurrent_decode_reference( + *args, config=cfg, initial_state=state, state_qdq=True, block_v=16 + ) + chunks = [] + carry = None + lo = 0 + for hi in [0, 5, 16, 16, 37, 73]: + result, carry = recurrent_decode_reference( + *(x[lo:hi] for x in args), + config=cfg, + carry=carry, + initial_state=state if carry is None else None, + state_qdq=True, + block_v=16, + ) + chunks.append(result) + lo = hi + actual = torch.cat(chunks) + for a, b in zip( + _values_and_grads(actual, carry.reconstruct(), args, state), + _values_and_grads(expected, final.reconstruct(), args, state), + ): + torch.testing.assert_close(a, b, rtol=0, atol=0) + assert carry.position == 73 + assert carry.anchor.format == "fp8_e4m3" + assert carry.anchor.scales.shape == (2, 1) + assert not carry.anchor.scales.requires_grad + for entry in carry.entries: + assert entry.key.scales.shape == (2, 1) + assert entry.update.values.requires_grad + assert not entry.update.scales.requires_grad + + +def test_initial_write_and_readout_order_are_explicit(): + args, state = _inputs(length=1) + outputs = [] + for readout in ["working", "stored"]: + output, carry = recurrent_decode_reference( + *args, + config=LinearAttentionDecodeConfig(readout=readout), + initial_state=state, + state_qdq=True, + block_v=16, + ) + outputs.append(output) + assert not torch.equal(*outputs) + empty, untouched = recurrent_decode_reference( + *(x[:0] for x in args), + config=LinearAttentionDecodeConfig(), + initial_state=state, + state_qdq=True, + ) + assert empty.shape == (0, 2, 11) + assert untouched.anchor.values is state + assert not untouched.started + + +def test_grid_gate_ste_keeps_gate_gradients_and_changes_trajectory(): + args, state = _inputs() + output, carry = recurrent_decode_reference( + *args, config=LinearAttentionDecodeConfig(decay_log_step=0.02), initial_state=state + ) + grad = torch.autograd.grad(output.square().sum() + carry.reconstruct().square().sum(), args[3])[ + 0 + ] + assert torch.isfinite(grad).all() and torch.count_nonzero(grad) > 0 + exact, _ = recurrent_decode_reference( + *args, config=LinearAttentionDecodeConfig(), initial_state=state + ) + assert not torch.equal(output, exact) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("prefixes", [[0, 0, 0, 0], [0, 3, 65, 0], [0, 5, 68, 0]]) +def test_explicit_packed_phase_with_empty_entries(kda, prefixes): + args, one_state = _inputs(kda) + args = tuple(x.unsqueeze(0) for x in args) + initial = torch.stack([one_state.detach().clone() for _ in range(4)]).requires_grad_() + boundaries = [0, 0, 5, 73, 73] + function = matmul_kda if kda else matmul_gdn + actual = function( + *args, + policy=LinearAttentionConfig(backend="matmul", decode={}), + prefill_lengths=prefixes, + cu_seqlens=torch.tensor(boundaries), + initial_state=initial, + output_final_state=True, + ) + outputs, states = [], [] + for n, (lo, hi) in enumerate(itertools.pairwise(boundaries)): + if lo == hi: + outputs.append(args[2][0, lo:hi]) + states.append(initial[n]) + continue + o, s = recurrent_delta_rule_reference( + *(x[:, lo:hi] for x in args), initial_state=initial[n : n + 1] + ) + outputs.append(o[0]) + states.append(s[0]) + expected = torch.cat(outputs).unsqueeze(0), torch.stack(states) + for a, e in zip(actual, expected): + torch.testing.assert_close(a, e, rtol=2e-11, atol=1e-12) + gradients = [ + torch.autograd.grad( + sum(x.square().sum() for x in result), (*args, initial), retain_graph=True + ) + for result in (actual, expected) + ] + for a, e in zip(*gradients): + torch.testing.assert_close(a, e, rtol=2e-10, atol=1e-11) + + +@pytest.mark.parametrize("mode", ["token", "replay"]) +def test_carry_rejects_changed_policy_and_shape(mode): + args, state = _inputs(length=4) + cfg = LinearAttentionDecodeConfig(mode=mode, replay={"window": 8} if mode == "replay" else None) + _, carry = recurrent_decode_reference(*args, config=cfg, initial_state=state, state_qdq=True) + for changed in [ + cfg.model_copy(update={"readout": "working"}), + cfg.model_copy(update={"decay_log_step": 0.01}), + ]: + with pytest.raises(ValueError, match="policy or state shape"): + recurrent_decode_reference(*args, config=changed, carry=carry, state_qdq=True) + with pytest.raises(ValueError, match="policy or state shape"): + recurrent_decode_reference(*args, config=cfg, carry=carry, state_qdq=True, block_v=16) + with pytest.raises(ValueError, match="either carry or initial"): + recurrent_decode_reference( + *args, config=cfg, carry=carry, initial_state=state, state_qdq=True + ) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("all_empty", [False, True]) +def test_grouped_heads_transposed_state_and_empty_batch(kda, all_empty): + args, state = _inputs(kda, length=0 if all_empty else 9) + q, k, v, g, beta = [x.unsqueeze(0).repeat(2, *([1] * x.ndim)) for x in args] + q, k = q[:, :, :1], k[:, :, :1] + initial = state.unsqueeze(0).repeat(2, 1, 1, 1).transpose(-1, -2).detach().requires_grad_() + function = matmul_kda if kda else matmul_gdn + policy = LinearAttentionConfig( + backend="matmul", decode={"mode": "replay", "replay": {"window": 3}} + ) + shared = { + "policy": policy, + "state_qdq": True, + "output_final_state": True, + } + actual = function( + q, + k, + v, + g, + beta, + initial_state=initial, + state_v_first=True, + prefill_lengths=[0, 0] if all_empty else [2, 5], + **shared, + ) + expected = function( + q.repeat_interleave(2, dim=2), + k.repeat_interleave(2, dim=2), + v, + g, + beta, + initial_state=initial.transpose(-1, -2), + prefill_lengths=[0, 0] if all_empty else [2, 5], + **shared, + ) + torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0) + torch.testing.assert_close(actual[1].transpose(-1, -2), expected[1], rtol=0, atol=0) + gradients = [ + torch.autograd.grad(result[0].sum() + result[1].sum(), initial, retain_graph=True)[0] + for result in (actual, expected) + ] + torch.testing.assert_close(*gradients, rtol=0, atol=0) + if all_empty: + torch.testing.assert_close(actual[1], initial, rtol=0, atol=0) diff --git a/tests/unit/torch/quantization/test_linear_attention_hadamard.py b/tests/unit/torch/quantization/test_linear_attention_hadamard.py new file mode 100644 index 00000000000..5f6b40f4c5c --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_hadamard.py @@ -0,0 +1,227 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import pytest +import torch +import torch.nn.functional as F + +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.config import QuantizeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + LinearAttentionDecodeConfig, + matmul_gdn, + matmul_kda, + recurrent_decode_reference, +) +from modelopt.torch.quantization.linear_attention.decode import _encode, _hadamard32 + + +def _rotate(value): + # Independent dense Sylvester matrix; production uses a butterfly transform. + matrix = value.new_tensor([[(-1) ** (i & j).bit_count() for j in range(32)] for i in range(32)]) + return (value.reshape(*value.shape[:-1], -1, 32) @ (matrix / 32**0.5)).reshape_as(value) + + +def _qdq(value): + with torch.no_grad(): + groups = value.float().unflatten(-1, (-1, 32)) + scale = (groups.abs().amax(-1, keepdim=True) / 127).clamp_min(6e-8) + normalized = groups / scale + codes = torch.where(normalized >= 0, (normalized + 0.5).floor(), (normalized - 0.5).ceil()) + rounded = (codes.clamp(-127, 127) * scale.half().float()).flatten(-2).to(value.dtype) + return (value - value.detach()) + rounded + + +def _inputs(kda, length=11): + torch.manual_seed(321) + q, k = [F.normalize(torch.randn(length, 2, 4, dtype=torch.float64), dim=-1) for _ in range(2)] + v = torch.randn(length, 2, 64, dtype=torch.float64) + g = -torch.rand(k.shape if kda else k.shape[:-1], dtype=torch.float64) * 0.05 + beta = torch.rand(length, 2, dtype=torch.float64) * 0.4 + initial = torch.randn(2, 4, 64, dtype=torch.float64) * 0.1 + return tuple(x.requires_grad_() for x in (q, k, v, g, beta)), initial.requires_grad_() + + +def _oracle(args, initial, mode, readout, quantize=True): + q, k, v, g, beta = args + state = _rotate(initial) + if quantize: + state = _qdq(state) + outputs = [] + for t in range(len(q)): + decay = g[t].exp().reshape(2, -1, 1) + decayed = state * decay + correction = beta[t, :, None] * (_rotate(v[t]) - torch.einsum("hk,hkv->hv", k[t], decayed)) + working = decayed + k[t, :, :, None] * correction[:, None, :] + state = _qdq(working) if quantize and (mode == "token" or (t + 1) % 3 == 0) else working + read = state if readout == "stored" else working + outputs.append(_rotate(torch.einsum("hk,hkv->hv", q[t], read)) / q.shape[-1] ** 0.5) + return torch.stack(outputs), _rotate(state) + + +def _check_with_grads(actual, expected, args, initial, tolerance=2e-10): + for a, e in zip(actual, expected): + torch.testing.assert_close(a, e, rtol=tolerance, atol=tolerance) + gradients = [ + torch.autograd.grad( + sum(x.square().sum() for x in result), (*args, initial), retain_graph=True + ) + for result in (actual, expected) + ] + for a, e in zip(*gradients): + torch.testing.assert_close(a, e, rtol=tolerance, atol=tolerance) + + +def test_hadamard_codec_basis_scales_ties_and_identity_gradient(): + raw = torch.zeros(1, 2, 64, dtype=torch.float64) + raw[0, 0, :7] = torch.tensor([254, -254, 1, -1, 3, -3, 0]) + raw[0, 1, 32:] = torch.linspace(-0.123, 0.123, 32) + raw.requires_grad_() + torch.testing.assert_close(_hadamard32(raw), _rotate(raw), rtol=1e-14, atol=1e-14) + torch.testing.assert_close(_hadamard32(_hadamard32(raw)), raw, rtol=1e-14, atol=1e-13) + encoded = _encode(raw, True, 64, state=True, state_format="int8", state_codec="int8_hadamard32") + torch.testing.assert_close(encoded.values, _qdq(raw), rtol=0, atol=0) + torch.testing.assert_close( + encoded.values[0, 0, :7], raw.new_tensor([254, -254, 2, -2, 4, -4, 0]) + ) + assert encoded.scales.shape == (1, 2, 2) + assert encoded.scales.dtype == torch.float16 and not encoded.scales.requires_grad + assert encoded.scales[0, 0, 0] == 2 + assert encoded.scales[0, 0, 1] == torch.tensor(6e-8, dtype=torch.float16) + probe = torch.randn_like(raw) + (gradient,) = torch.autograd.grad((encoded.values * probe).sum(), raw) + torch.testing.assert_close(gradient, probe, rtol=0, atol=0) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize( + ("mode", "readout"), [("token", "stored"), ("replay", "working"), ("replay", "stored")] +) +def test_hadamard_recurrence_matches_dense_oracle_and_split_carry(kda, mode, readout): + args, initial = _inputs(kda) + cfg = LinearAttentionDecodeConfig( + mode=mode, + readout=readout, + state_codec="int8_hadamard32", + replay={"window": 3, "factor_qdq": False} if mode == "replay" else None, + ) + kwargs = {"config": cfg, "state_format": "int8", "state_qdq": True} + output, carry = recurrent_decode_reference(*args, initial_state=initial, **kwargs) + _check_with_grads( + (output, carry.reconstruct()), _oracle(args, initial, mode, readout), args, initial + ) + assert carry.value_basis == "hadamard32" and carry.anchor.block_v == 32 + assert carry.anchor.scales.shape == (2, 4, 2) + pieces, split_carry, start = [], None, 0 + for end in (0, 2, 3, 3, 7, 11): + piece, split_carry = recurrent_decode_reference( + *(x[start:end] for x in args), + carry=split_carry, + initial_state=initial if split_carry is None else None, + **kwargs, + ) + if end == 0: + assert split_carry.reconstruct() is initial and not split_carry.started + pieces.append(piece) + start = end + _check_with_grads( + (torch.cat(pieces), split_carry.reconstruct()), + (output, carry.reconstruct()), + args, + initial, + 0, + ) + with pytest.raises(ValueError, match="Carry policy"): + recurrent_decode_reference( + *args, + carry=carry, + **{**kwargs, "config": cfg.model_copy(update={"state_codec": "tile"})}, + ) + + +@pytest.mark.parametrize("kda", [False, True]) +def test_hadamard_disabled_qdq_is_a_change_of_basis(kda): + args, initial = _inputs(kda) + cfg = LinearAttentionDecodeConfig(mode="replay", replay={"window": 3, "factor_qdq": False}) + results = [] + for codec in ("tile", "int8_hadamard32"): + out, carry = recurrent_decode_reference( + *args, initial_state=initial, config=cfg.model_copy(update={"state_codec": codec}) + ) + results.append((out, carry.reconstruct())) + _check_with_grads(*results, args, initial) + + +@pytest.mark.parametrize(("kda", "prefix"), [(False, 0), (False, 3), (True, 3), (True, 11)]) +def test_hadamard_exact_prefill_handoff_and_state_layout(kda, prefix, project_root_path): + args, initial = _inputs(kda) + function = matmul_kda if kda else matmul_gdn + if kda: + path = ( + project_root_path + / "examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json" + ) + recipe = QuantizeConfig(**json.loads(path.read_text())) + else: + recipe = load_recipe("general/ptq/gdn_state_int8_dynamic").quantize + policy = recipe.linear_attention[0].cfg + assert recipe.algorithm is None + assert recipe.quant_cfg[1].cfg.num_bits == 8 + assert policy.decode.state_codec == "int8_hadamard32" + assert not policy.decode.prefill_state_qdq + # Exercise the shipped policy with the CPU implementation for this algebra test. + policy.decode.implementation = "torch" + actual = function( + *(x.unsqueeze(0) for x in args), + policy=policy, + prefill_lengths=[prefix], + state_qdq=True, + state_format="int8", + initial_state=initial.transpose(-1, -2).unsqueeze(0), + state_v_first=True, + output_final_state=True, + ) + prefix_out, state = recurrent_decode_reference( + *(x[:prefix] for x in args), initial_state=initial, config=LinearAttentionDecodeConfig() + ) + suffix, carry = recurrent_decode_reference( + *(x[prefix:] for x in args), + initial_state=state.reconstruct(), + config=policy.decode, + state_qdq=True, + state_format="int8", + ) + expected = (torch.cat((prefix_out, suffix)), carry.reconstruct()) + _check_with_grads((actual[0][0], actual[1][0].transpose(-1, -2)), expected, args, initial) + + +def test_hadamard_rejects_unsupported_codec_combinations(): + args, initial = _inputs(False) + cfg = LinearAttentionDecodeConfig(state_codec="int8_hadamard32") + for kwargs, message in [ + ({"state_qdq": True}, "requires INT8"), + ({"state_format": "int8", "block_v": 16}, "block_v >= 32"), + ]: + with pytest.raises(ValueError, match=message): + recurrent_decode_reference(*args, config=cfg, initial_state=initial, **kwargs) + with pytest.raises(ValueError, match="Dv divisible"): + recurrent_decode_reference(args[0], args[1], args[2][..., :33], *args[3:], config=cfg) + with pytest.raises(ValueError, match="decode handoff"): + LinearAttentionDecodeConfig(state_codec="int8_hadamard32", prefill_state_qdq=True) + with pytest.raises(ValueError, match="block_v >= 32"): + LinearAttentionConfig(backend="matmul", state={"block_v": 16}, decode=cfg) diff --git a/tests/unit/torch/quantization/test_linear_attention_int8.py b/tests/unit/torch/quantization/test_linear_attention_int8.py new file mode 100644 index 00000000000..b3de8c94d47 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_int8.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import pytest +import torch +import torch.nn.functional as F +from _test_utils.torch.quantization.linear_attention_reference import state_qdq_reference + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionDecodeConfig, + recurrent_decode_reference, +) +from modelopt.torch.quantization.linear_attention.decode import _encode +from modelopt.torch.quantization.nn import TensorQuantizer + +INT8 = {"num_bits": 8, "unsigned": False, "narrow_range": True, "type": "dynamic", "axis": (0, 1)} + + +@pytest.mark.parametrize("block_v", [16, 32, 64, 128]) +def test_int8_state_codec_matches_modelopt_and_identity_gradient(block_v): + # Exact scale two exercises both signs, endpoints and nearest-even half ties. + state = torch.zeros(2, 3, 5, block_v + 3) + state[..., 0, :9] = torch.tensor([-127, -126.5, -1.5, -0.5, 0, 0.5, 1.5, 126.5, 127]) * 2 + state[0, :, 0, block_v:] = torch.tensor([-254, 1, 254]) + state.requires_grad_() + quantizer = TensorQuantizer(QuantizerAttributeConfig(**INT8)) + expected = torch.cat( + [quantizer(tile.flatten(-2)).reshape_as(tile) for tile in state.split(block_v, -1)], -1 + ) + encoded = _encode(state, True, block_v, state=True, state_format="int8") + reference = state_qdq_reference(state, block_v, "int8") + torch.testing.assert_close(encoded.values, expected, rtol=0, atol=0) + torch.testing.assert_close(reference, expected, rtol=0, atol=0) + assert encoded.format == "int8" + assert torch.equal(encoded.scales[1, :, -1], torch.ones(3)) + assert not encoded.scales.requires_grad + probe = torch.randn_like(state) + for result in (encoded.values, reference): + (gradient,) = torch.autograd.grad((result * probe).sum(), state, retain_graph=True) + torch.testing.assert_close(gradient, probe, rtol=0, atol=0) + + +def _inputs(kda): + torch.manual_seed(913) + q, k = [F.normalize(torch.randn(23, 2, 16), dim=-1).requires_grad_() for _ in range(2)] + v = torch.randn(23, 2, 19, requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1]) * 0.03).requires_grad_() + beta = (torch.rand(23, 2) * 0.4).requires_grad_() + state = (torch.randn(2, 16, 19) * 0.1).requires_grad_() + return (q, k, v, g, beta), state + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("mode", ["token", "replay"]) +def test_int8_decode_continuation_and_format_rejection(kda, mode): + args, initial = _inputs(kda) + cfg = LinearAttentionDecodeConfig(mode=mode, replay={"window": 5} if mode == "replay" else None) + kwargs = {"config": cfg, "state_qdq": True, "state_format": "int8", "block_v": 16} + expected, final = recurrent_decode_reference(*args, initial_state=initial, **kwargs) + pieces, carry, start = [], None, 0 + for end in (0, 3, 10, 10, 23): + output, carry = recurrent_decode_reference( + *(x[start:end] for x in args), + carry=carry, + initial_state=initial if carry is None else None, + **kwargs, + ) + pieces.append(output) + start = end + actual = torch.cat(pieces) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(carry.reconstruct(), final.reconstruct(), rtol=0, atol=0) + probes = (torch.randn_like(expected), torch.randn_like(initial)) + grads = [ + torch.autograd.grad((o * probes[0]).sum() + (s * probes[1]).sum(), (*args, initial)) + for o, s in ((actual, carry.reconstruct()), (expected, final.reconstruct())) + ] + for actual_grad, expected_grad in zip(*grads): + torch.testing.assert_close(actual_grad, expected_grad, rtol=0, atol=0) + assert carry.anchor.format == "int8" + assert all(e.key.format == e.update.format == "fp8_e4m3" for e in carry.entries) + with pytest.raises(ValueError, match="Carry policy"): + recurrent_decode_reference( + *(x[:0] for x in args), carry=carry, **{**kwargs, "state_format": "fp8_e4m3"} + ) From 756c3304c64457f5d3646c7ba217ac07971a786f Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sun, 27 Sep 2026 20:58:13 -0700 Subject: [PATCH 6/8] Add fused Triton GDN/KDA decode QAT Signed-off-by: Kai Xu --- examples/llm_qat/linear_attention/README.md | 15 +- .../configs/kda_decode_state_int8.json | 2 +- .../kernels/quantization/common/fp8_quant.py | 31 +- .../quantization/linear_attention/decode.py | 468 ++++++++++++++++++ .../quantization/linear_attention/config.py | 4 +- .../quantization/linear_attention/decode.py | 102 +++- .../linear_attention/test_decode.py | 412 +++++++++++++++ .../torch/quantization/plugins/test_kda.py | 4 +- 8 files changed, 1016 insertions(+), 22 deletions(-) create mode 100644 modelopt/torch/kernels/quantization/linear_attention/decode.py create mode 100644 tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md index 7c44058512a..ad0e116c1c3 100644 --- a/examples/llm_qat/linear_attention/README.md +++ b/examples/llm_qat/linear_attention/README.md @@ -55,8 +55,8 @@ states, so the policy must specify which computation training should emulate. `LinearAttentionConfig` holds the overall backend, chunk size, and state settings. Its optional `decode` field contains a `LinearAttentionDecodeConfig` for token or -replay mode, state codec, decay approximation, and replay settings. The decode -implementation uses Torch operations and autograd. This is one nested configuration: supplying a `decode` +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 @@ -101,7 +101,9 @@ mtq.quantize(gdn_model, gdn_recipe) Use the phase context shown below for each forward/backward. Importing only `configs/ptq/units/gdn_state_int8_dynamic` configures the quantizer without -selecting Hadamard. Both complete recipes use the Torch decode implementation. +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 @@ -293,9 +295,10 @@ the optional tile codec instead uses nearest-even rounding and FP32 scales. ## Training boundaries -The Torch implementation uses autograd through initial states, chunk handoff, -token writes, and replay refreshes. The decode recurrence runs token by token in -Python, so long training suffixes can be slow. +The Torch and Triton implementations support first-order QAT gradients through +initial states, chunk handoff, token writes, and replay refreshes. Triton supports +key dimensions up to 128 and value blocks 16/32/64/128, with checkpointed backward. +Use the Torch implementation for higher-order differentiation. Prefill prefixes use exact chunk algebra with optional state QDQ. This example has no prefill GEMM QDQ or approximate inverse. FLA KDA requires `use_cache=False`; 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 index 8a29cda8139..f6cc5fef650 100644 --- a/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json +++ b/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json @@ -25,7 +25,7 @@ "cfg": { "backend": "matmul", "decode": { - "implementation": "torch", + "implementation": "triton", "state_codec": "int8_hadamard32", "prefill_state_qdq": false } 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/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/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index 1f3b610fe75..33756c5b459 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -51,7 +51,7 @@ class LinearAttentionDecodeConfig(ModeloptBaseConfig): """Explicit suffix recurrence; workload supplies per-sequence prefix lengths.""" mode: Literal["token", "replay"] = ModeloptField(default="token") - implementation: Literal["torch"] = ModeloptField(default="torch") + 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) @@ -63,6 +63,8 @@ class LinearAttentionDecodeConfig(ModeloptBaseConfig): 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" diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py index 2a981c597ec..80905958bb9 100644 --- a/modelopt/torch/quantization/linear_attention/decode.py +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -322,20 +322,106 @@ def recurrent_decode( carry=None, position=0, scale=None, + checkpoint_interval=8, ): - """Run the Torch training implementation with explicit differentiable carry.""" - return recurrent_decode_reference( + """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, - k, + key.values, v, - g, + gate, beta, - config=config, + carry.reconstruct(original_basis=False), state_qdq=state_qdq, state_format=state_format, + state_codec=config.state_codec, block_v=block_v, - initial_state=initial_state, - carry=carry, - position=position, + 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/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py new file mode 100644 index 00000000000..da6fdf6e772 --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py @@ -0,0 +1,412 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_dynamic_qdq +from modelopt.torch.kernels.quantization.linear_attention.decode import fused_recurrence +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention.config import LinearAttentionDecodeConfig +from modelopt.torch.quantization.linear_attention.decode import ( + _encode, + recurrent_decode, + recurrent_decode_reference, +) +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("state_format", [None, "fp8_e4m3", "int8"]) +@pytest.mark.parametrize("replay", [False, True]) +@pytest.mark.parametrize("read_stored", [False, True]) +def test_fused_recurrence_outputs_state_and_gradients(kda, state_format, replay, read_stored): + state_qdq = state_format is not None + state_format = state_format or "fp8_e4m3" + torch.manual_seed(193) + shape = (37, 2, 16) + q, k = [ + F.normalize(torch.randn(shape, device="cuda"), dim=-1).requires_grad_() for _ in range(2) + ] + v = torch.randn(37, 2, 11, device="cuda", requires_grad=True) + g = (-torch.rand(shape if kda else shape[:-1], device="cuda") * 0.03).requires_grad_() + beta = (torch.rand(shape[:-1], device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(2, 16, 11, device="cuda") * 0.1).requires_grad_() + cfg = LinearAttentionDecodeConfig( + mode="replay" if replay else "token", + quantize_initial=False, + readout="stored" if read_stored else "working", + replay={"window": 5} if replay else None, + ) + encoded_key = _encode(k, replay, 16).values + actual = fused_recurrence( + q, + encoded_key, + v, + g, + beta, + initial, + state_qdq=state_qdq, + state_format=state_format, + block_v=16, + replay=replay, + factor_qdq=replay, + window=5, + read_stored=read_stored, + ) + expected, carry = recurrent_decode_reference( + q, + k, + v, + g, + beta, + config=cfg, + initial_state=initial, + state_qdq=state_qdq, + state_format=state_format, + block_v=16, + ) + final = carry.reconstruct() + targets = (q, k, v, g, beta, initial) + probes = [torch.randn_like(x) for x in actual[:2]] + gradients = [ + torch.autograd.grad(sum((a * p).sum() for a, p in zip(result, probes)), targets) + for result in (actual[:2], (expected, final)) + ] + for a, e in zip((*actual[:2], *gradients[0]), (expected, final, *gradients[1])): + torch.testing.assert_close(a, e, rtol=1e-3, atol=8e-6) + + +@pytest.mark.parametrize("mode", ["token", "replay"]) +@pytest.mark.parametrize("interval", [3, 8]) +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_fused_carry_handoff_and_auxiliary_gradients(mode, interval, state_format): + torch.manual_seed(217) + q, k = [ + F.normalize(torch.randn(23, 2, 16, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(23, 2, 19, device="cuda", requires_grad=True) + g = (-torch.rand_like(q) * 0.05).requires_grad_() + beta = (torch.rand(23, 2, device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(2, 16, 19, device="cuda") * 0.2).requires_grad_() + args = (q, k, v, g, beta) + reference = LinearAttentionDecodeConfig( + mode=mode, decay_log_step=0.01, replay={"window": 5} if mode == "replay" else None + ) + candidate = reference.model_copy(update={"implementation": "triton"}) + expected, expected_carry = recurrent_decode_reference( + *args, + config=reference, + state_qdq=True, + state_format=state_format, + block_v=16, + initial_state=initial, + ) + carry = None + outputs = [] + start = 0 + for end in [0, 3, 11, 12, 23]: + output, carry = recurrent_decode( + *(x[start:end] for x in args), + config=candidate, + state_qdq=True, + state_format=state_format, + block_v=16, + initial_state=initial if carry is None else None, + carry=carry, + checkpoint_interval=interval, + ) + outputs.append(output) + start = end + actual = torch.cat(outputs) + actual_state, expected_state = carry.reconstruct(), expected_carry.reconstruct() + assert carry.cursor == expected_carry.cursor + assert carry.position == 23 and carry.started + torch.testing.assert_close( + carry.anchor.scales, expected_carry.anchor.scales, rtol=5e-5, atol=2e-8 + ) + for a, e in zip(carry.entries, expected_carry.entries): + torch.testing.assert_close(a.update.scales, e.update.scales, rtol=5e-5, atol=2e-8) + torch.testing.assert_close(a.key.values, e.key.values, rtol=0, atol=0) + probes = [torch.randn_like(actual), torch.randn_like(actual_state)] + gradients = [ + torch.autograd.grad(sum((a * p).sum() for a, p in zip(result, probes)), (*args, initial)) + for result in ((actual, actual_state), (expected, expected_state)) + ] + for a, e in zip( + (actual, actual_state, *gradients[0]), (expected, expected_state, *gradients[1]) + ): + torch.testing.assert_close(a, e, rtol=1e-3, atol=1e-5) + + +@triton.jit +def _codec_kernel(X, Y, S, N: tl.constexpr, BLOCK: tl.constexpr, NATIVE_FP8: tl.constexpr): # noqa: N803 + offsets = tl.arange(0, BLOCK) + values = tl.load(X + offsets, offsets < N, 0) + decoded, scale = fp8_dynamic_qdq(values, NATIVE_FP8) + tl.store(Y + offsets, decoded, offsets < N) + tl.store(S, scale) + + +@pytest.mark.parametrize("native_fp8", [False, True]) +@pytest.mark.parametrize( + "scale", + [0.0, 1e-8 / 448, 2.0**-24 / 448, 2.0**-23 / 448, 2.0**-12, 1.0, 2.0**12, 0.00314159, 731.173], +) +def test_fp8_codec_matches_tensor_quantizer_boundaries(scale, native_fp8): + if native_fp8 and torch.cuda.get_device_capability() < (8, 9): + pytest.skip("Native FP8 conversion requires SM89 or later") + # Exhaust all finite positive E4M3 values and their rounding ties, plus signs. + representable = torch.arange(127, dtype=torch.uint8).view(torch.float8_e4m3fn).float() + ties = (representable[:-1] + representable[1:]) / 2 + probe = torch.cat( + ( + representable, + ties, + torch.nextafter(ties, torch.full_like(ties, float("inf"))), + torch.nextafter(ties, torch.zeros_like(ties)), + ) + ) + probe = torch.cat((probe, -probe)) * scale + values = probe.cuda() + actual = torch.empty_like(values) + actual_scale = torch.empty((), device="cuda") + _codec_kernel[(1,)]( + values, + actual, + actual_scale, + len(values), + triton.next_power_of_2(len(values)), + native_fp8, + ) + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), type="dynamic")) + torch.testing.assert_close(actual, quantizer(values), rtol=0, atol=0) + amax = values.abs().amax() + safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) + torch.testing.assert_close( + actual_scale, torch.div(448.0, safe_amax).reciprocal(), rtol=0, atol=0 + ) + + +def test_fused_fp8_state_matches_tensor_quantizer_with_partial_tiles(): + torch.manual_seed(762) + amplitudes = torch.tensor([0.0, 1e-8, 2**-24, 2**-23, 1.0, 731.173], device="cuda") + state = (torch.rand(6, 3, 19, device="cuda") * 2 - 1) * amplitudes[:, None, None] + state[:, 0, 0] = amplitudes + state.requires_grad_() + q = torch.zeros(1, 6, 3, device="cuda") + v = torch.zeros(1, 6, 19, device="cuda") + gate = torch.zeros(1, 6, device="cuda") + _, actual, *_ = fused_recurrence(q, q, v, gate, gate, state, state_qdq=True, block_v=16) + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=(4, 3), type="dynamic", axis=0, pass_through_bwd=True) + ) + expected = torch.cat( + [quantizer(tile.flatten(-2)).reshape_as(tile) for tile in state.split(16, -1)], -1 + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + probe = torch.randn_like(state) + gradients = [ + torch.autograd.grad((value * probe).sum(), state)[0] for value in (actual, expected) + ] + torch.testing.assert_close(*gradients, rtol=0, atol=0) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("mode", ["token", "replay"]) +@pytest.mark.parametrize("horizon", ["near_unit", "mixed"]) +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_wide_long_trajectory(kda, mode, horizon, state_format): + torch.manual_seed(329) + q, k = [ + F.normalize(torch.randn(257, 2, 128, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(257, 2, 70, device="cuda", requires_grad=True) + shape = q.shape if kda else q.shape[:-1] + g = -torch.rand(shape, device="cuda") * (1e-4 if horizon == "near_unit" else 3) + if horizon == "mixed": + g[::7] = -30 + g.requires_grad_() + beta = (torch.rand(257, 2, device="cuda") * 0.6).requires_grad_() + initial = (torch.randn(2, 128, 70, device="cuda") * 0.1).requires_grad_() + cfg = LinearAttentionDecodeConfig(mode=mode, replay={"window": 8} if mode == "replay" else None) + args = q, k, v, g, beta + expected, ec = recurrent_decode_reference( + *args, + config=cfg, + state_qdq=True, + state_format=state_format, + block_v=64, + initial_state=initial, + ) + actual, ac = recurrent_decode( + *args, + config=cfg.model_copy(update={"implementation": "triton"}), + state_qdq=True, + state_format=state_format, + block_v=64, + initial_state=initial, + ) + probes = [torch.randn_like(actual), torch.randn_like(ac.reconstruct())] + gradients = [ + torch.autograd.grad(sum((a * p).sum() for a, p in zip(result, probes)), (*args, initial)) + for result in ((actual, ac.reconstruct()), (expected, ec.reconstruct())) + ] + # Fixed reduction and scale arithmetic must also agree over long trajectories. + for a, e in zip( + (actual, ac.reconstruct(), *gradients[0]), (expected, ec.reconstruct(), *gradients[1]) + ): + relative = (a - e).norm() / e.norm().clamp_min(1e-8) + assert relative < 1e-4, relative.item() + assert torch.isfinite(a).all() + + +@pytest.mark.parametrize("block_v", [32, 128]) +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_full_key_width_and_value_block_tail(block_v, kda, state_format): + torch.manual_seed(418) + q, k = [ + F.normalize(torch.randn(9, 2, 128, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(9, 2, 133, device="cuda", requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1], device="cuda") * 0.1).requires_grad_() + beta = torch.rand(9, 2, device="cuda", requires_grad=True) + initial = torch.randn(2, 128, 133, device="cuda", requires_grad=True) + args = q, k, v, g, beta + cfg = LinearAttentionDecodeConfig(mode="replay", replay={"window": 3}) + results = [] + for implementation in ("torch", "triton"): + output, carry = recurrent_decode( + *args, + config=cfg.model_copy(update={"implementation": implementation}), + state_qdq=True, + state_format=state_format, + block_v=block_v, + initial_state=initial, + ) + final = carry.reconstruct() + gradients = torch.autograd.grad( + output.square().sum() + final.square().sum(), (*args, initial) + ) + results.append((output, final, *gradients)) + for actual, expected in zip(*results): + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=2e-5) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize( + ("mode", "block_v", "readout", "factor_qdq"), + [ + ("token", 32, "stored", False), + ("replay", 64, "working", False), + ("replay", 128, "stored", True), + ], +) +@pytest.mark.timeout(180) +def test_hadamard_fused_split_carry_and_gradients(kda, mode, block_v, readout, factor_qdq): + torch.manual_seed(917) + q, k = [ + F.normalize(torch.randn(37, 2, 17, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(37, 2, 96, device="cuda", requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1], device="cuda") * 0.03).requires_grad_() + beta = (torch.rand(37, 2, device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(2, 17, 96, device="cuda") * 0.1).requires_grad_() + args = (q, k, v, g, beta) + cfg = LinearAttentionDecodeConfig( + mode=mode, + readout=readout, + state_codec="int8_hadamard32", + decay_log_step=1 / 256, + replay={"window": 8, "factor_qdq": factor_qdq} if mode == "replay" else None, + ) + shared = {"state_qdq": True, "state_format": "int8", "block_v": block_v} + expected, ref_carry = recurrent_decode_reference( + *args, config=cfg, initial_state=initial, **shared + ) + carry, start, pieces = None, 0, [] + for end in (0, 2, 8, 8, 19, 37): + output, carry = recurrent_decode( + *(x[start:end] for x in args), + config=cfg.model_copy(update={"implementation": "triton"}), + initial_state=initial if carry is None else None, + carry=carry, + **shared, + ) + pieces.append(output) + start = end + actual = torch.cat(pieces) + assert carry.value_basis == "hadamard32" and carry.cursor == ref_carry.cursor + assert carry.anchor.scales.shape == (2, 17, 3) and carry.anchor.scales.dtype == torch.float16 + torch.testing.assert_close(carry.anchor.scales, ref_carry.anchor.scales, rtol=0, atol=0) + actual_tensors = ( + actual, + carry.reconstruct(), + carry.anchor.values, + *(e.update.values for e in carry.entries), + ) + expected_tensors = ( + expected, + ref_carry.reconstruct(), + ref_carry.anchor.values, + *(e.update.values for e in ref_carry.entries), + ) + probes = [torch.randn_like(x) for x in expected_tensors] + gradients = [ + torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probes)), (*args, initial)) + for result in (actual_tensors, expected_tensors) + ] + for a, e in zip((*actual_tensors, *gradients[0]), (*expected_tensors, *gradients[1])): + torch.testing.assert_close(a, e, rtol=1e-3, atol=1e-5) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.timeout(180) +def test_hadamard_long_trajectory(kda): + torch.manual_seed(991) + q, k = [ + F.normalize(torch.randn(257, 1, 128, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(257, 1, 64, device="cuda", requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1], device="cuda") * 1e-4).requires_grad_() + beta = (torch.rand(257, 1, device="cuda") * 0.1).requires_grad_() + initial = (torch.randn(1, 128, 64, device="cuda") * 0.1).requires_grad_() + args = (q, k, v, g, beta) + results = [] + for implementation in ("torch", "triton"): + cfg = LinearAttentionDecodeConfig( + mode="replay", + implementation=implementation, + state_codec="int8_hadamard32", + replay={"window": 8, "factor_qdq": False}, + ) + out, carry = recurrent_decode( + *args, config=cfg, initial_state=initial, state_qdq=True, state_format="int8" + ) + final = carry.reconstruct() + grads = torch.autograd.grad(out.square().sum() + final.square().sum(), (*args, initial)) + results.append((out, final, *grads)) + for a, e in zip(*results): + assert torch.isfinite(a).all() and torch.isfinite(e).all() + assert (a - e).norm() / e.norm().clamp_min(1e-8) < 1e-4 diff --git a/tests/gpu/torch/quantization/plugins/test_kda.py b/tests/gpu/torch/quantization/plugins/test_kda.py index ed67787de54..aecf7fe8ef8 100644 --- a/tests/gpu/torch/quantization/plugins/test_kda.py +++ b/tests/gpu/torch/quantization/plugins/test_kda.py @@ -83,7 +83,7 @@ def test_fla_layer_qat_restore_and_optimizer(tmp_path, mode): policy_cfg["state"] = {"block_v": 16} policy_cfg["decode"] = { "mode": "replay" if mode.startswith("decode_replay") else "token", - "implementation": "torch", + "implementation": "triton", } if mode.startswith("decode_replay"): policy_cfg["decode"]["replay"] = {"window": 5} @@ -166,7 +166,7 @@ def test_decode_phase_spans_activation_checkpoint_backward(hadamard): "cfg": { "backend": "matmul", "decode": { - "implementation": "torch", + "implementation": "triton", "state_codec": "int8_hadamard32" if hadamard else "tile", }, }, From bae8fa23a297d62e6a76754cfb2093c56bb87f03 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 23 Sep 2026 13:24:00 -0700 Subject: [PATCH 7/8] Add vLLM GDN/KDA state-only fake quantization Signed-off-by: Kai Xu --- CHANGELOG.rst | 2 + examples/llm_qat/linear_attention/README.md | 6 + examples/vllm_serve/README.md | 77 +++++ examples/vllm_serve/fakequant_worker.py | 26 +- .../linear_attention_state_int8.yaml | 12 + examples/vllm_serve/vllm_reload_utils.py | 45 ++- .../torch/quantization/plugins/__init__.py | 1 + .../plugins/vllm_linear_attention.py | 165 +++++++++ .../test_vllm_linear_attention.py | 315 ++++++++++++++++++ 9 files changed, 638 insertions(+), 11 deletions(-) create mode 100644 examples/vllm_serve/linear_attention_state_int8.yaml create mode 100644 modelopt/torch/quantization/plugins/vllm_linear_attention.py create mode 100644 tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 60d905778b0..0a751f6c2d5 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,6 +13,8 @@ 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 ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md index ad0e116c1c3..43b6c970c39 100644 --- a/examples/llm_qat/linear_attention/README.md +++ b/examples/llm_qat/linear_attention/README.md @@ -305,3 +305,9 @@ has no prefill GEMM QDQ or approximate inverse. FLA KDA requires `use_cache=Fals serving cache objects are rejected. ModelOpt saves execution policies, while per-batch prefix lengths must be supplied again during training. Distributed decode training and model-quality recovery require separate qualification. + +## 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/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/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py index e1b2ec7a1ae..91b95257cce 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -71,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/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/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"] From bd2be2b9cdcaf1155eeac597bb0ce402e4d2a181 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 22 Sep 2026 18:12:46 -0700 Subject: [PATCH 8/8] Add GDN and KDA prefill QAT after decode infrastructure Signed-off-by: Kai Xu --- CHANGELOG.rst | 4 + examples/llm_qat/linear_attention/README.md | 520 ++++++++++++++++-- .../linear_attention/benchmark_decode.py | 173 ++++++ .../linear_attention/benchmark_prefill.py | 164 ++++++ .../linear_attention/benchmark_utils.py | 59 ++ .../linear_attention/compare_quality.py | 113 ++++ .../configs/kda_decode_decay_1024.json | 35 ++ .../configs/kda_decode_decay_256.json | 35 ++ .../configs/kda_decode_decay_4096.json | 35 ++ .../configs/kda_decode_exact.json | 20 + .../configs/kda_decode_replay_fp8.json | 40 ++ .../configs/kda_decode_replay_int8.json | 39 ++ .../kda_decode_replay_int8_hadamard.json | 44 ++ .../configs/kda_decode_state_fp8.json | 34 ++ .../configs/kda_exact_matmul.json | 20 + .../configs/kda_prefill_fp8.json | 47 ++ .../configs/kda_prefill_nvfp4.json | 53 ++ .../llm_qat/linear_attention/decode_study.md | 116 ++++ .../linear_attention/decode_study_plan.json | 40 ++ examples/llm_qat/linear_attention/train.py | 258 +++++++-- modelopt/torch/quantization/conversion.py | 6 + .../quantization/linear_attention/__init__.py | 2 + .../linear_attention/_chunk_prefill.py | 146 ----- .../quantization/linear_attention/config.py | 62 ++- .../linear_attention/decode_prefill.py | 35 +- .../quantization/linear_attention/kda.py | 54 +- .../quantization/linear_attention/matmul.py | 123 +++++ .../quantization/linear_attention/prefill.py | 197 ++++++- .../quantization/linear_attention/solve.py | 35 ++ modelopt/torch/quantization/plugins/gdn.py | 2 + modelopt/torch/quantization/plugins/kda.py | 2 + .../quantization/plugins/linear_attention.py | 16 +- modelopt_recipes/configs/ptq/units/README.md | 2 + .../units/default_disabled_quantizers.yaml | 2 + .../ptq/units/gdn_prefill_fp8_dynamic.yaml | 28 + .../ptq/units/gdn_prefill_nvfp4_dynamic.yaml | 34 ++ tests/_test_utils/torch/linear_attention.py | 167 ++++++ .../linear_attention/test_decode.py | 62 +++ .../linear_attention/test_matmul_gdn.py | 199 +++++++ .../linear_attention/test_matmul_kda.py | 186 +++++++ .../torch/quantization/plugins/test_kda.py | 31 +- .../plugins/test_megatron_gated_delta_net.py | 26 + .../torch/quantization/plugins/test_gdn.py | 61 ++ .../torch/quantization/test_kda_prefill.py | 95 ++++ .../test_linear_attention_decode.py | 5 + .../test_linear_attention_hadamard.py | 6 +- .../test_linear_attention_prefill.py | 126 +++++ .../test_linear_attention_study.py | 95 ++++ 48 files changed, 3342 insertions(+), 312 deletions(-) create mode 100644 examples/llm_qat/linear_attention/benchmark_decode.py create mode 100644 examples/llm_qat/linear_attention/benchmark_prefill.py create mode 100644 examples/llm_qat/linear_attention/benchmark_utils.py create mode 100644 examples/llm_qat/linear_attention/compare_quality.py create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_decay_1024.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_decay_256.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_decay_4096.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_exact.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_replay_fp8.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_replay_int8.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_replay_int8_hadamard.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_state_fp8.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_exact_matmul.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_prefill_fp8.json create mode 100644 examples/llm_qat/linear_attention/configs/kda_prefill_nvfp4.json create mode 100644 examples/llm_qat/linear_attention/decode_study.md create mode 100644 examples/llm_qat/linear_attention/decode_study_plan.json delete mode 100644 modelopt/torch/quantization/linear_attention/_chunk_prefill.py create mode 100644 modelopt/torch/quantization/linear_attention/matmul.py create mode 100644 modelopt/torch/quantization/linear_attention/solve.py create mode 100644 modelopt_recipes/configs/ptq/units/gdn_prefill_fp8_dynamic.yaml create mode 100644 modelopt_recipes/configs/ptq/units/gdn_prefill_nvfp4_dynamic.yaml create mode 100644 tests/_test_utils/torch/linear_attention.py create mode 100644 tests/gpu/torch/kernels/quantization/linear_attention/test_matmul_gdn.py create mode 100644 tests/gpu/torch/kernels/quantization/linear_attention/test_matmul_kda.py create mode 100644 tests/unit/torch/quantization/test_kda_prefill.py create mode 100644 tests/unit/torch/quantization/test_linear_attention_prefill.py create mode 100644 tests/unit/torch/quantization/test_linear_attention_study.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0a751f6c2d5..147e9a88632 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -17,6 +17,10 @@ Changelog - 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. diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md index 43b6c970c39..9b6110deca8 100644 --- a/examples/llm_qat/linear_attention/README.md +++ b/examples/llm_qat/linear_attention/README.md @@ -1,39 +1,80 @@ # Quantization-Aware Training for Linear Attention -This example fine-tunes KDA attention parameters with recurrent-state fake -quantization, then saves a Hugging Face checkpoint with ModelOpt state. GDN/KDA -runtime support includes token writes, ReplaySSM, KDA decay approximation, and -FP8 or INT8 state QDQ. The INT8 recipes enable Hadamard rotation by default. +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. -## Run the example +## Prerequisites -Install ModelOpt using the [QAT setup instructions](../README.md#quick-start), -then install the example dependencies. Run commands from the repository root. +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/local-kda-model \ + --model /path/to/model-snapshot \ --train-data /path/to/train.parquet \ - --output /path/to/qat-checkpoint \ - --train-steps 1 --length 128 --prefill-tokens 64 + --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 ``` -Use a local model/tokenizer snapshot containing FLA `KimiDeltaAttention` layers, -a Parquet file with a `text` column, and a CUDA GPU. The example requires -`fla-core==0.5.1` and `flash-linear-attention==0.5.1`. If the local model requires -custom Python code, review it and explicitly pass `--trust-remote-code`. - -The default [INT8 configuration](configs/kda_decode_state_int8.json) leaves the -64-token prefix state unquantized and applies INT8 QDQ in a 32-value Hadamard -basis during the decode suffix. Value dimensions must be divisible by 32. -Loss uses suffix labels. Only KDA attention parameters are trained, in FP32 under -BF16 autocast; other parameters are frozen in BF16. The output contains model -weights, tokenizer files, and ModelOpt quantizer/policy state. Call -`mto.enable_huggingface_checkpointing()` before reloading it with -`AutoModelForCausalLM.from_pretrained` to restore those quantizers and policies. -It is a floating fake-quantized training checkpoint, not a compressed serving model. -This short example does not measure model-quality recovery or performance. +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 @@ -69,9 +110,8 @@ must be supplied again for each workload. State quantizers start disabled. The state recipe enables `*kda_state_quantizer` with signed narrow-range INT8, dynamic scales, and `axis=(0, 1)`. Use -`*gdn_state_quantizer` for GDN. For E4M3, set the quantizer config to -`{"num_bits": [4, 3], "type": "dynamic", "axis": [0, 1]}` and set -`decode.state_codec="tile"`. +`*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** @@ -108,7 +148,7 @@ conversion to select the fused decode implementation. ### Why training needs prefill/decode boundaries One training forward can simulate prompt prefill followed by recurrent decode. -The INT8 + Hadamard recipes apply different state quantization schedules to those +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 @@ -144,7 +184,7 @@ phase context. See the tables below for the corresponding quantization settings. ### Choose the prefill and decode boundaries -The table assumes the state quantizer is enabled. The INT8 recipes default to +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`. @@ -163,6 +203,7 @@ decode.update(mode="token", replay=None, state_codec="tile", prefill_state_qdq=T 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: @@ -179,7 +220,8 @@ decode.update( 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. +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 @@ -211,7 +253,8 @@ to disable both state and anchor QDQ; replay factor QDQ has its own toggle. The `int8_hadamard32` codec applies to decode token writes or replay anchors only. It requires `prefill_state_qdq=False`; use the tile codec to quantize prefix state. -Save a modified recipe as JSON and pass it to `train.py --quant-config`. `--prefill-tokens` sets the phase split; it does not enable state +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 @@ -239,6 +282,11 @@ 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. @@ -275,36 +323,396 @@ runtime metadata and must be supplied for each workload. The context restores previous lengths on exit and supports nesting. Concurrent forwards on the same model instance with different phase contexts are unsupported. -## Replay, decay, and Hadamard options +## 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. -Token mode rounds each recurrent-state write. Replay mode retains an anchor and -ordered key/update factors, rounding the anchor at each `replay.window` refresh. -`factor_qdq` independently enables FP8 QDQ for those factors. `readout="working"` -reads the state before its write quantization; `"stored"` reads it afterward. +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. -For KDA decay approximation, set `decode["decay_log_step"] = 1 / 256` before -conversion. It rounds suffix log retention before exponentiation with identity -STE gradients. Prefix decay remains exact. +The 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"}, + }], +}) +``` -The default INT8 codec applies a 32-point orthonormal Hadamard transform along -the value axis, quantizes token states or replay anchors in that basis, and -transforms outputs back. Value dimensions must be divisible by 32; `state.block_v` must be 32, 64, -or 128. Scales group one key channel and 32 values, independent of execution tile -width. INT8 codes use half-away-from-zero rounding with FP16 stored scales; -the optional tile codec instead uses nearest-even rounding and FP32 scales. +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. -## Training boundaries +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: -The Torch and Triton implementations support first-order QAT gradients through -initial states, chunk handoff, token writes, and replay refreshes. Triton supports -key dimensions up to 128 and value blocks 16/32/64/128, with checkpointed backward. -Use the Torch implementation for higher-order differentiation. +| 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 +``` -Prefill prefixes use exact chunk algebra with optional state QDQ. This example -has no prefill GEMM QDQ or approximate inverse. FLA KDA requires `use_cache=False`; -serving cache objects are rejected. ModelOpt saves execution policies, while -per-batch prefix lengths must be supplied again during training. Distributed -decode training and model-quality recovery require separate qualification. +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 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_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/train.py b/examples/llm_qat/linear_attention/train.py index 7662c1ca593..618cc74de50 100644 --- a/examples/llm_qat/linear_attention/train.py +++ b/examples/llm_qat/linear_attention/train.py @@ -13,10 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Fine-tune KDA attention parameters with recurrent-state fake quantization.""" +"""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 @@ -24,85 +30,257 @@ from fla.layers.kda import KimiDeltaAttention from transformers import AutoModelForCausalLM, AutoTokenizer -import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.quantization.linear_attention import linear_attention_training_phase +def _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(): - """Quantize, train on local text, and save a ModelOpt-aware Hugging Face checkpoint.""" + """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( - "--train-data", type=Path, required=True, help="Parquet file with a text column" - ) - parser.add_argument( - "--quant-config", - type=Path, - default=Path(__file__).with_name("configs") / "kda_decode_state_int8.json", - ) + parser.add_argument("--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, default=64) + 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.train_steps < 1 or options.length < 2: - parser.error("train-steps must be positive and length must be at least two") - if not 0 <= options.prefill_tokens < options.length: - parser.error("prefill-tokens must leave at least one suffix label") + if options.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) - mto.enable_huggingface_checkpointing() + 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 ) - text = "\n\n".join(pq.read_table(options.train_data, columns=["text"])["text"].to_pylist()) - tokens = tokenizer.encode(text, add_special_tokens=False) - needed = options.train_steps * options.length - if len(tokens) < needed: - parser.error(f"Training data has {len(tokens)} tokens; {needed} are required") - blocks = torch.tensor(tokens[:needed]).reshape(options.train_steps, options.length) - model = AutoModelForCausalLM.from_pretrained( + 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, ) - layers = [module for module in model.modules() if isinstance(module, KimiDeltaAttention)] - if not layers: - raise ValueError("The model must contain FLA KimiDeltaAttention layers") - # Keep attention parameters and optimizer moments in FP32; freeze the rest in BF16. + 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() - mtq.quantize(model, json.loads(options.quant_config.read_text())) + 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() - for step, block in enumerate(blocks, 1): - ids = block.unsqueeze(0).cuda() - labels = ids.clone() - labels[:, : options.prefill_tokens] = -100 + 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) - # Keep phase metadata active during activation-checkpoint recomputation. - with linear_attention_training_phase(model, [options.prefill_tokens]): + 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, use_cache=False).loss + 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() - torch.nn.utils.clip_grad_norm_(trainable, 1.0, error_if_nonfinite=True) + 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() - print(f"step={step} loss={loss.item():.4f}", flush=True) - model.save_pretrained(options.output) - tokenizer.save_pretrained(options.output) + 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__": diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index db194a0f465..f36af6e77cd 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -147,6 +147,12 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: return model 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(): diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py index 305bb2cd81f..fe41be40476 100644 --- a/modelopt/torch/quantization/linear_attention/__init__.py +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -19,4 +19,6 @@ 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/_chunk_prefill.py b/modelopt/torch/quantization/linear_attention/_chunk_prefill.py deleted file mode 100644 index fddff0303ae..00000000000 --- a/modelopt/torch/quantization/linear_attention/_chunk_prefill.py +++ /dev/null @@ -1,146 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Exact chunked prefix computation for decode-aware training.""" - -import torch - -from .utils import _prepare, _state_qdq - -__all__ = [] - - -def chunk_gdn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - *, - chunk_size: int = 64, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - cu_seqlens: torch.Tensor | None = None, - state_v_first: bool = False, - state_qdq: bool = False, - state_qdq_block_v: int = 64, - state_format: str = "fp8_e4m3", -) -> tuple[torch.Tensor, torch.Tensor]: - """Compute a GDN prefix with optional QDQ on the initial state and chunk writes.""" - if g.ndim != 3 or chunk_size <= 0: - raise ValueError("chunk_gdn requires scalar GDN gates and positive chunk_size") - q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) - scale = q.shape[-1] ** -0.5 if scale is None else scale - outputs, finals = [], [] - for n, (b, start, end) in enumerate(sequences): - state = states[n] - if state_qdq: - state = _state_qdq(state, state_qdq_block_v, state_format) - pieces = [] - for lo in range(start, end, chunk_size): - hi = min(lo + chunk_size, end) - qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) - gc = g[b, lo:hi].transpose(0, 1).cumsum(-1) - bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) - # Mask before exp: upper-triangle positive differences can overflow for long decay. - causal = torch.ones(hi - lo, hi - lo, device=q.device, dtype=torch.bool).tril() - decay = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() - lower = (bc * (kc @ kc.transpose(-1, -2)) * decay).tril(-1) - matrix = lower + torch.eye(hi - lo, device=q.device, dtype=q.dtype) - gate = gc.exp().unsqueeze(-1) - rhs = torch.cat((bc * vc, bc * kc * gate), dim=-1) - solved = torch.linalg.solve_triangular(matrix, rhs, upper=False, unitriangular=True) - u, w = solved.split((v.shape[-1], k.shape[-1]), dim=-1) - updated_values = u - w @ state - local_scores = ((qc * scale @ kc.transpose(-1, -2)) * decay).tril() - output = (qc * (scale * gc.exp()).unsqueeze(-1)) @ state - pieces.append((output + local_scores @ updated_values).transpose(0, 1)) - weighted_keys = kc * (gc[..., -1:] - gc).exp().unsqueeze(-1) - state = state * gc[..., -1].exp()[:, None, None] - state = state + weighted_keys.transpose(-1, -2) @ updated_values - if state_qdq: - state = _state_qdq(state, state_qdq_block_v, state_format) - outputs.append(torch.cat(pieces)) - finals.append(state) - output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) - final = torch.stack(finals) - return output, final.transpose(-1, -2) if state_v_first else final - - -def chunk_kda( - q, - k, - v, - g, - beta, - *, - chunk_size=64, - scale=None, - initial_state=None, - cu_seqlens=None, - state_v_first=False, - state_qdq=False, - state_qdq_block_v=64, - state_format="fp8_e4m3", -): - """Compute a KDA prefix with optional QDQ on the initial state and chunk writes.""" - if g.ndim != 4 or chunk_size <= 0: - raise ValueError("chunk_kda requires per-key gates and positive chunk_size") - q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) - scale = q.shape[-1] ** -0.5 if scale is None else scale - outputs, finals = [], [] - for n, (b, start, end) in enumerate(sequences): - state = states[n] - if state_qdq: - state = _state_qdq(state, state_qdq_block_v, state_format) - pieces = [] - for lo in range(start, end, chunk_size): - hi = min(lo + chunk_size, end) - qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) - prefix = g[b, lo:hi].transpose(0, 1).cumsum(-2) - bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) - gate = prefix.exp() - lower_rows, score_rows = [], [] - for row in range(hi - lo): - decayed_keys = ( - kc[..., : row + 1, :] - * (prefix[..., row : row + 1, :] - prefix[..., : row + 1, :]).exp() - ) - right = decayed_keys.transpose(-1, -2) - left = kc[..., row : row + 1, :] * bc[..., row : row + 1, :] - interaction = left @ right - score = qc[..., row : row + 1, :] * scale @ right - padding = hi - lo - row - 1 - lower_rows.append(torch.nn.functional.pad(interaction, (0, padding))) - score_rows.append(torch.nn.functional.pad(score, (0, padding))) - lower = torch.cat(lower_rows, dim=-2).tril(-1) - scores = torch.cat(score_rows, dim=-2) - identity = torch.eye(hi - lo, device=q.device, dtype=q.dtype).expand_as(lower) - inverse = torch.linalg.solve_triangular( - identity + lower, identity, upper=False, unitriangular=True - ) - u = inverse @ (bc * vc) - w = inverse @ (bc * kc * gate) - updated = u - w @ state - pieces.append(((qc * scale * gate) @ state + scores @ updated).transpose(0, 1)) - weighted_keys = kc * (prefix[..., -1:, :] - prefix).exp() - state = state * gate[..., -1, :, None] + weighted_keys.transpose(-1, -2) @ updated - if state_qdq: - state = _state_qdq(state, state_qdq_block_v, state_format) - outputs.append(torch.cat(pieces)) - finals.append(state) - output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) - final = torch.stack(finals) - return output, final.transpose(-1, -2) if state_v_first else final diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index 33756c5b459..af41fd148d9 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -24,10 +24,49 @@ __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") @@ -35,7 +74,9 @@ class _StateConfig(ModeloptBaseConfig): quantize_initial: Literal[True] = ModeloptField(default=True) -class _SolveConfig(ModeloptBaseConfig): +class LinearAttentionSolveConfig(ModeloptBaseConfig): + """Exact solve; inverse approximation is a later delivery.""" + method: Literal["exact"] = ModeloptField(default="exact") @@ -73,11 +114,11 @@ def _validate_replay(self): class LinearAttentionConfig(ModeloptBaseConfig): - """GDN chunk-64 policy; unsupported numerical modes fail config validation. + """GDN/KDA chunk-64 policy; unsupported numerical modes fail config validation. ``state.block_v`` defines one dynamic scale per ``[Dk, block_v]`` tile of each sequence/head. The initial state and every chunk's final state are rounded when - ``gdn_state_quantizer`` is enabled. Outputs use the incoming rounded state. + the module's state quantizer is enabled. Outputs use the incoming rounded state. Decode's ``int8_hadamard32`` codec instead fixes scales to one key channel and 32 values; ``state.block_v`` remains the execution tile width. """ @@ -86,13 +127,20 @@ class LinearAttentionConfig(ModeloptBaseConfig): backend: Literal["fla", "matmul"] = ModeloptField(default="fla") chunk_size: Literal[64] = ModeloptField(default=64) state: _StateConfig = ModeloptField(default=_StateConfig()) - solve: _SolveConfig = ModeloptField(default=_SolveConfig()) + 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_decode_backend(self): - if (self.backend == "matmul") != (self.decode is not None): - raise ValueError("The exact-prefix matmul backend requires an explicit decode policy") + 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") diff --git a/modelopt/torch/quantization/linear_attention/decode_prefill.py b/modelopt/torch/quantization/linear_attention/decode_prefill.py index 0a0f69b4ae7..3ad194c017f 100644 --- a/modelopt/torch/quantization/linear_attention/decode_prefill.py +++ b/modelopt/torch/quantization/linear_attention/decode_prefill.py @@ -20,7 +20,6 @@ import torch -from ._chunk_prefill import chunk_gdn, chunk_kda from .decode import recurrent_decode __all__ = ["linear_attention_training_phase"] @@ -71,7 +70,9 @@ def _decode_prefill( g, beta, *, + sites, policy, + w_quantizer, state_qdq, state_format, scale, @@ -83,6 +84,9 @@ def _decode_prefill( 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]: @@ -137,19 +141,22 @@ def _decode_prefill( offsets = [0] for n in active: offsets.append(offsets[-1] + prefixes[n]) - prefix_fn = chunk_kda if g.ndim == 4 else chunk_gdn - with torch.autocast(device_type=q.device.type, enabled=False): - output, final = prefix_fn( - *packed, - state_qdq=state_qdq and policy.decode.prefill_state_qdq, - state_format=state_format, - scale=scale, - initial_state=states[active], - cu_seqlens=torch.tensor(offsets), - state_v_first=False, - chunk_size=policy.chunk_size, - state_qdq_block_v=policy.state.block_v, - ) + 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] diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py index 446dd6f7d0d..1e6e4a4c398 100644 --- a/modelopt/torch/quantization/linear_attention/kda.py +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -18,7 +18,8 @@ import torch import torch.nn.functional as F -from .decode_prefill import _decode_prefill +from .matmul import _validate_operand_quantizer +from .prefill import _dispatch_prefill __all__ = ["matmul_kda"] @@ -30,7 +31,9 @@ def matmul_kda( g, beta, *, + sites, policy, + w_quantizer, state_qdq=False, state_format="fp8_e4m3", scale=None, @@ -53,9 +56,11 @@ def matmul_kda( return_intermediate_states=False, prefill_lengths=None, ): - """Normalize KDA inputs and run exact prefix plus configured suffix recurrence. + """Run KDA chunk prefill with actual-operand QDQ and differentiable state carry. - Gates use the loaded FLA model's activation formula and per-key log retention. + 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") @@ -69,6 +74,8 @@ def matmul_kda( 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)) @@ -85,23 +92,24 @@ def matmul_kda( g = -rate * F.softplus(g) if lower_bound is None else lower_bound * (rate * g).sigmoid() if use_beta_sigmoid_in_kernel: beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) - if policy.decode is not None: - return _decode_prefill( - q, - k, - v, - g, - beta, - policy=policy, - state_qdq=state_qdq, - state_format=state_format, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - cu_seqlens_cpu=cu_seqlens_cpu, - state_v_first=state_v_first, - output_dtype=output_dtype, - prefill_lengths=prefill_lengths, - ) - raise ValueError("An explicit decode policy is required") + 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 index 4a602f8f5bc..ffc54f12ba9 100644 --- a/modelopt/torch/quantization/linear_attention/prefill.py +++ b/modelopt/torch/quantization/linear_attention/prefill.py @@ -20,6 +20,9 @@ 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"] @@ -31,7 +34,9 @@ def matmul_gdn( g, beta, *, + sites: LinearAttentionMatmulSites, policy: LinearAttentionConfig, + w_quantizer, state_qdq=False, state_format="fp8_e4m3", scale=None, @@ -50,15 +55,19 @@ def matmul_gdn( cp_context=None, prefill_lengths=None, ): - """Normalize GDN inputs and run exact prefix plus configured suffix recurrence. + """Execute all eight GDN prefill sites with autograd through their QDQ operators. - Uses FP32 working arithmetic (FP64 for double inputs), with floating QDQ state. - Prefixes use exact chunk algebra; no configurable operand rounding is enabled. + 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)) @@ -72,23 +81,169 @@ def matmul_gdn( beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) if g.ndim != 3: raise ValueError("matmul_gdn requires scalar GDN log gates") + return _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( - q, - k, - v, - g, - beta, - policy=policy, - state_qdq=state_qdq, - state_format=state_format, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - cu_seqlens_cpu=cu_seqlens_cpu, - state_v_first=state_v_first, - output_dtype=output_dtype, - prefill_lengths=prefill_lengths, + 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 ) - raise ValueError("An explicit decode policy is required") + 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/plugins/gdn.py b/modelopt/torch/quantization/plugins/gdn.py index 8d0fe616390..dd9ac102440 100644 --- a/modelopt/torch/quantization/plugins/gdn.py +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -104,7 +104,9 @@ def _state_quantized_chunk_gated_delta_rule( 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, diff --git a/modelopt/torch/quantization/plugins/kda.py b/modelopt/torch/quantization/plugins/kda.py index fbc4b1cf0ef..9c0ead2fe9f 100644 --- a/modelopt/torch/quantization/plugins/kda.py +++ b/modelopt/torch/quantization/plugins/kda.py @@ -71,7 +71,9 @@ def 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, diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py index 5559c9f6606..36d5b1c08c0 100644 --- a/modelopt/torch/quantization/plugins/linear_attention.py +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -17,6 +17,7 @@ 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 @@ -33,6 +34,7 @@ def _setup(self): ) self._register_temp_attribute("linear_attention_config", LinearAttentionConfig()) self._register_temp_attribute("_linear_attention_prefill_lengths", None) + self._register_temp_attribute("linear_attn_sites", LinearAttentionMatmulSites()) @property def _linear_attn_state(self): @@ -52,6 +54,9 @@ def linear_attention_is_enabled(self): 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 ) @@ -72,12 +77,17 @@ def validate_linear_attention(self): if decode is not None and decode.state_codec == "int8_hadamard32": if self._linear_attn_state_format != "int8": raise ValueError("int8_hadamard32 requires INT8 state quantization") - if self._linear_attn_w.is_enabled: + 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] ) - if self.linear_attention_config.decode is not None: - raise ValueError("Decode's exact prefix does not support WY operand QDQ") + 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.""" diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index ea115b6fe88..f00932169ca 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -35,3 +35,5 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `gdn_state_fp8_dynamic.yaml` | FP8 E4M3 dynamic fake quantization of the GatedDeltaNet recurrent state (per sequence, head and 64-column tile) at 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 e409bff7aa6..3559145b48d 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -25,6 +25,8 @@ 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/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/gpu/torch/kernels/quantization/linear_attention/test_decode.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py index da6fdf6e772..17cf0112640 100644 --- a/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py @@ -22,6 +22,12 @@ 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, @@ -410,3 +416,59 @@ def test_hadamard_long_trajectory(kda): 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_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 index aecf7fe8ef8..291e3b57dc8 100644 --- a/tests/gpu/torch/quantization/plugins/test_kda.py +++ b/tests/gpu/torch/quantization/plugins/test_kda.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy from contextlib import nullcontext import pytest @@ -52,11 +53,17 @@ def _forward(model, hidden): @pytest.mark.parametrize( "mode", [ + "state", + "state_int8", + "w", + "prefill_fp8", + "prefill_nvfp4", + "arithmetic", "decode_token", "decode_replay", - "decode_decay", "decode_token_int8", "decode_replay_int8", + "decode_decay", "decode_replay_hadamard_int8", ], ) @@ -69,14 +76,28 @@ def test_fla_layer_qat_restore_and_optimizer(tmp_path, mode): 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"] @@ -116,6 +137,14 @@ def test_fla_layer_qat_restore_and_optimizer(tmp_path, mode): 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) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index 94ad1aaedf1..e825477d7e8 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -93,6 +93,8 @@ def _gdn_config(sites): 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": @@ -105,6 +107,24 @@ def _config_for_mode(mode): 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 @@ -190,6 +210,7 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): assert len(restored_gdn) == len(gdn_modules) for module, original in zip(restored_gdn, gdn_modules): assert module.linear_attention_config == original.linear_attention_config + assert module.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 @@ -240,6 +261,9 @@ def compiled_gdn_workers(request, tp_size, mode): "state", "w", "both", + "prefill_fp8", + "prefill_nvfp4", + "prefill_arithmetic", "decode_token", "decode_replay", "state_int8", @@ -264,6 +288,8 @@ def _test_gdn_context_parallel_helper(rank, size, mode): [ "state", "w", + "prefill_fp8", + "prefill_arithmetic", "decode_token", "decode_replay", "state_int8", diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py index aed7f8caf42..cda0f0b263b 100644 --- a/tests/unit/torch/quantization/plugins/test_gdn.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -119,6 +119,7 @@ def test_dynamic_export_removes_linear_attention_attributes(): "gdn_w_quantizer", "linear_attention_config", "_linear_attention_prefill_lengths", + "linear_attn_sites", ): assert not hasattr(model, name) @@ -304,6 +305,66 @@ def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): 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) 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 index e0ced168024..393b555faca 100644 --- a/tests/unit/torch/quantization/test_linear_attention_decode.py +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -23,6 +23,7 @@ from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.linear_attention import ( LinearAttentionConfig, + LinearAttentionMatmulSites, matmul_gdn, matmul_kda, ) @@ -199,7 +200,9 @@ def test_explicit_packed_phase_with_empty_entries(kda, prefixes): 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, @@ -260,7 +263,9 @@ def test_grouped_heads_transposed_state_and_empty_batch(kda, all_empty): 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, } diff --git a/tests/unit/torch/quantization/test_linear_attention_hadamard.py b/tests/unit/torch/quantization/test_linear_attention_hadamard.py index 5f6b40f4c5c..77208748aa0 100644 --- a/tests/unit/torch/quantization/test_linear_attention_hadamard.py +++ b/tests/unit/torch/quantization/test_linear_attention_hadamard.py @@ -20,7 +20,7 @@ import torch.nn.functional as F from modelopt.recipe import load_recipe -from modelopt.torch.quantization.config import QuantizeConfig +from modelopt.torch.quantization.config import QuantizeConfig, QuantizerAttributeConfig from modelopt.torch.quantization.linear_attention import ( LinearAttentionConfig, LinearAttentionDecodeConfig, @@ -29,6 +29,8 @@ 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): @@ -188,7 +190,9 @@ def test_hadamard_exact_prefill_handoff_and_state_layout(kda, prefix, project_ro 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", 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_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)