diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0078076d0bf..92ec264e274 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 706bc4a5baf..cec6dc16497 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -27,6 +27,7 @@ Changelog - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration. - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md `_ for details. - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``. +- Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training through the standard ``quant_cfg`` interface. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer. - Add fake quantization of the sparse-attention indexer key cache and query for DeepSeek-V4 (vLLM and Megatron-Core) and GLM-5.3-Flash (vLLM) through the new ``indexer_k_quantizer`` and ``indexer_q_quantizer``. Enable them by importing the ``configs/ptq/units/indexer_k_nvfp4`` and ``configs/ptq/units/indexer_q_nvfp4`` units (NVFP4 with the global scale fixed to 1) into a recipe. *Speculative Decoding* diff --git a/LICENSE b/LICENSE index 57c1104ae67..e08eaefc016 100644 --- a/LICENSE +++ b/LICENSE @@ -251,6 +251,7 @@ the following copyright holders, licensed under the MIT License: Copyright (c) 2025 sgl-project Copyright (c) 2026 The DeepSpec Authors Copyright (c) 2023-2026 The ggml authors + Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/docs/source/_templates/autosummary/module.rst b/docs/source/_templates/autosummary/module.rst index 36eb087fbec..c693958c63d 100644 --- a/docs/source/_templates/autosummary/module.rst +++ b/docs/source/_templates/autosummary/module.rst @@ -11,7 +11,7 @@ :recursive: {% for item in modules %} {% set full_item = fullname + '.' + item.split('.')[-1] %} -{% if ('.plugins.' not in full_item or full_item == 'modelopt.torch.opt.plugins.huggingface') and full_item != 'modelopt.torch.quantization.backends.fp8_per_tensor_gemm' %} +{% if ('.plugins.' not in full_item or full_item == 'modelopt.torch.opt.plugins.huggingface') and full_item != 'modelopt.torch.quantization.backends.fp8_per_tensor_gemm' and not full_item.startswith('modelopt.torch.kernels.quantization.linear_attention.fla_') %} {{ full_item }} {% endif %} {%- endfor %} diff --git a/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..875cb39a7d0 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -0,0 +1,978 @@ +# Adapted from: https://github.com/fla-org/flash-linear-attention/blob/516143e31fce/fla/ops/common/chunk_delta_h.py +# Adapted with modifications (marked [ModelOpt]): optional in-kernel FP8 E4M3 fake quantization +# of the carried chunk state (STATE_QDQ), BV as an explicit launch argument, no fla backend +# dispatch or config-cache autotune. +# +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import triton +import triton.language as tl +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.cache import fla_cache_autotune +from fla.ops.utils.op import exp2 +from fla.utils import ( + IS_INTEL, + IS_NVIDIA_BLACKWELL, + IS_NVIDIA_HOPPER, + autotune_cache_kwargs, + check_shared_mem, +) + +from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq + +# ``STATE_QDQ`` modes of the forward state kernel. +STATE_QDQ_OFF = 0 +STATE_QDQ_FP8_DYNAMIC = 1 # FP8 E4M3, one dynamic scale per program tile ([K, BV] of one head) +STATE_QDQ_MAX_BLOCK_V = 128 + + +@triton.jit +def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr): + """[ModelOpt] Dynamic FP8 E4M3 scale per full [K, BV] tile of one sequence and head.""" + b_amax = tl.max(tl.abs(b_h1)) + if K > 64: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h2))) + if K > 128: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h3))) + if K > 192: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h4))) + return tl.where(b_amax > 0, b_amax / 448.0, 1.0) + + +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] + +# TODO: Triton mainline fixes a Blackwell tl.dot recurrence race. +# Keep this kernel on num_warps=2 for Blackwell until Triton 3.8 is released +# and we re-validate the wider config space. +# Intel needs more warps than NVIDIA here: 8 warps is ~1.5x faster than the best +# config reachable under the [2, 4] cap. +if IS_NVIDIA_BLACKWELL: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2] +elif IS_INTEL: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2, 4, 8, 16] +else: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2, 4] + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +# ``BV`` is an explicit argument rather than an autotuned config: with ``STATE_QDQ`` it sets the +# quantization granularity, so it must not vary with the autotuner's choice. +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in GATED_DELTA_RULE_FWD_H_NUM_WARPS + for num_stages in ([2, 3, 4] if check_shared_mem("ampere") else [2, 1]) + ], + key=["H", "HV", "K", "V", "BT", "BV", "STATE_V_FIRST", "STATE_QDQ"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( + k, + v, + w, + v_new, + g, + gk, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + SAVE_NEW_VALUE: tl.constexpr, + STATE_V_FIRST: tl.constexpr, + IS_VARLEN: tl.constexpr, + STATE_QDQ: tl.constexpr, +): + pid = tl.program_id(0) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) + i_n, i_h = i_nh // HV, i_nh % HV + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + if STATE_V_FIRST: + b_h1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([BV, 64], dtype=tl.float32) + else: + b_h1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + h += (boh * HV + i_h).to(tl.int64) * K * V + v += (bos * HV + i_h).to(tl.int64) * V + k += (bos * H + i_h // (HV // H)).to(tl.int64) * K + w += (bos * HV + i_h).to(tl.int64) * K + if SAVE_NEW_VALUE: + v_new += (bos * HV + i_h).to(tl.int64) * V + + if USE_INITIAL_STATE: + h0 = h0 + i_nh * K * V + if STORE_FINAL_STATE: + ht = ht + i_nh * K * V + + # load initial state + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K + if USE_INITIAL_STATE: + if STATE_V_FIRST: + p_h0_1 = h0 + o_v[:, None] * K + o_k1[None, :] + m_h0_1 = m_v[:, None] & m_k1[None, :] + else: + p_h0_1 = h0 + o_k1[:, None] * V + o_v[None, :] + m_h0_1 = m_k1[:, None] & m_v[None, :] + b_h1 += tl.load(p_h0_1, mask=m_h0_1, other=0.0).to(tl.float32) + if K > 64: + if STATE_V_FIRST: + p_h0_2 = h0 + o_v[:, None] * K + o_k2[None, :] + m_h0_2 = m_v[:, None] & m_k2[None, :] + else: + p_h0_2 = h0 + o_k2[:, None] * V + o_v[None, :] + m_h0_2 = m_k2[:, None] & m_v[None, :] + b_h2 += tl.load(p_h0_2, mask=m_h0_2, other=0.0).to(tl.float32) + if K > 128: + if STATE_V_FIRST: + p_h0_3 = h0 + o_v[:, None] * K + o_k3[None, :] + m_h0_3 = m_v[:, None] & m_k3[None, :] + else: + p_h0_3 = h0 + o_k3[:, None] * V + o_v[None, :] + m_h0_3 = m_k3[:, None] & m_v[None, :] + b_h3 += tl.load(p_h0_3, mask=m_h0_3, other=0.0).to(tl.float32) + if K > 192: + if STATE_V_FIRST: + p_h0_4 = h0 + o_v[:, None] * K + o_k4[None, :] + m_h0_4 = m_v[:, None] & m_k4[None, :] + else: + p_h0_4 = h0 + o_k4[:, None] * V + o_v[None, :] + m_h0_4 = m_k4[:, None] & m_v[None, :] + b_h4 += tl.load(p_h0_4, mask=m_h0_4, other=0.0).to(tl.float32) + # [ModelOpt] A state read from an FP8 cache is quantized before the first chunk uses it. + if STATE_QDQ == 1: + if K > 192: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) + elif K > 128: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K) + elif K > 64: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K) + else: + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K) + b_h1 = fp8_scalar_qdq(b_h1, b_scale) + if K > 64: + b_h2 = fp8_scalar_qdq(b_h2, b_scale) + if K > 128: + b_h3 = fp8_scalar_qdq(b_h3, b_scale) + if K > 192: + b_h4 = fp8_scalar_qdq(b_h4, b_scale) + + # main recurrence + for i_t in range(NT): + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + if STATE_V_FIRST: + p_h1 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k1[None, :] + m_h1 = m_v[:, None] & m_k1[None, :] + else: + p_h1 = h + i_t_int64 * HV * K * V + o_k1[:, None] * V + o_v[None, :] + m_h1 = m_k1[:, None] & m_v[None, :] + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), mask=m_h1) + if K > 64: + if STATE_V_FIRST: + p_h2 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k2[None, :] + m_h2 = m_v[:, None] & m_k2[None, :] + else: + p_h2 = h + i_t_int64 * HV * K * V + o_k2[:, None] * V + o_v[None, :] + m_h2 = m_k2[:, None] & m_v[None, :] + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), mask=m_h2) + if K > 128: + if STATE_V_FIRST: + p_h3 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k3[None, :] + m_h3 = m_v[:, None] & m_k3[None, :] + else: + p_h3 = h + i_t_int64 * HV * K * V + o_k3[:, None] * V + o_v[None, :] + m_h3 = m_k3[:, None] & m_v[None, :] + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), mask=m_h3) + if K > 192: + if STATE_V_FIRST: + p_h4 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k4[None, :] + m_h4 = m_v[:, None] & m_k4[None, :] + else: + p_h4 = h + i_t_int64 * HV * K * V + o_k4[:, None] * V + o_v[None, :] + m_h4 = m_k4[:, None] & m_v[None, :] + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), mask=m_h4) + + p_w = w + o_t[:, None] * (HV * K) + o_k1[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k1[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) + else: + b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) + if K > 64: + p_w = w + o_t[:, None] * (HV * K) + o_k2[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k2[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h2.to(b_w.dtype), b_v) + if K > 128: + p_w = w + o_t[:, None] * (HV * K) + o_k3[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k3[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h3.to(b_w.dtype), b_v) + if K > 192: + p_w = w + o_t[:, None] * (HV * K) + o_k4[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k4[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h4.to(b_w.dtype), b_v) + p_v = v + o_t[:, None] * (HV * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & m_v[None, :], other=0.0) - b_v + + if SAVE_NEW_VALUE: + p_v = v_new + o_t[:, None] * (HV * V) + o_v[None, :] + tl.store(p_v, b_v.to(p_v.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) + + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + b_g_last = tl.load(g + (bos * HV + last_idx * HV + i_h).to(tl.int64)).to(tl.float32) + p_g = g + (bos * HV + i_h).to(tl.int64) + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] + b_g_last = exp2(b_g_last) + b_h1 *= b_g_last + if K > 64: + b_h2 *= b_g_last + if K > 128: + b_h3 *= b_g_last + if K > 192: + b_h4 *= b_g_last + + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k1, mask=(o_k1 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h1 *= exp2(b_gk_last1)[None, :] + else: + b_h1 *= exp2(b_gk_last1)[:, None] + if K > 64: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k2, mask=(o_k2 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h2 *= exp2(b_gk_last2)[None, :] + else: + b_h2 *= exp2(b_gk_last2)[:, None] + if K > 128: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k3, mask=(o_k3 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h3 *= exp2(b_gk_last3)[None, :] + else: + b_h3 *= exp2(b_gk_last3)[:, None] + if K > 192: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k4, mask=(o_k4 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h4 *= exp2(b_gk_last4)[None, :] + else: + b_h4 *= exp2(b_gk_last4)[:, None] + b_v = b_v.to(k.dtype.element_ty) + + p_k = k + o_k1[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k1[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h1 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h1 = tl.dot(b_k, b_v, b_h1) + if K > 64: + p_k = k + o_k2[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k2[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h2 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h2 = tl.dot(b_k, b_v, b_h2) + if K > 128: + p_k = k + o_k3[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k3[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h3 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h3 = tl.dot(b_k, b_v, b_h3) + if K > 192: + p_k = k + o_k4[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k4[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h4 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h4 = tl.dot(b_k, b_v, b_h4) + + # [ModelOpt] Dynamic per-tile FP8 QDQ of the next-chunk or stored final state. + # Recompute amax over [K, BV]; fp8_scalar_qdq applies that tile's scalar scale. + if STATE_QDQ == 1: + if K > 192: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) + elif K > 128: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K) + elif K > 64: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K) + else: + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K) + b_h1 = fp8_scalar_qdq(b_h1, b_scale) + if K > 64: + b_h2 = fp8_scalar_qdq(b_h2, b_scale) + if K > 128: + b_h3 = fp8_scalar_qdq(b_h3, b_scale) + if K > 192: + b_h4 = fp8_scalar_qdq(b_h4, b_scale) + + if STORE_FINAL_STATE: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k1[None, :] + m_ht = m_v[:, None] & m_k1[None, :] + else: + p_ht = ht + o_k1[:, None] * V + o_v[None, :] + m_ht = m_k1[:, None] & m_v[None, :] + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 64: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k2[None, :] + m_ht = m_v[:, None] & m_k2[None, :] + else: + p_ht = ht + o_k2[:, None] * V + o_v[None, :] + m_ht = m_k2[:, None] & m_v[None, :] + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 128: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k3[None, :] + m_ht = m_v[:, None] & m_k3[None, :] + else: + p_ht = ht + o_k3[:, None] * V + o_v[None, :] + m_ht = m_k3[:, None] & m_v[None, :] + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 192: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k4[None, :] + m_ht = m_v[:, None] & m_k4[None, :] + else: + p_ht = ht + o_k4[:, None] * V + o_v[None, :] + m_ht = m_k4[:, None] & m_v[None, :] + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), mask=m_ht) + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "USE_INITIAL_STATE": lambda args: args["dh0"] is not None, + "USE_FINAL_STATE_GRADIENT": lambda args: args["dht"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@fla_cache_autotune( + configs=[ + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in ([2, 3, 4] if check_shared_mem("ampere") else [1]) + for BV in ([32, 64] if check_shared_mem("ada") else [32]) + ], + key=["H", "HV", "K", "V", "BT", "BV", "USE_G", "STATE_V_FIRST"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( + q, + k, + w, + g, + gk, + dht, + dh0, + do, + dh, + dv, + dv2, + cu_seqlens, + chunk_offsets, + scale, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + STATE_V_FIRST: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + pid = tl.program_id(0) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) + i_n, i_h = i_nh // HV, i_nh % HV + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + if STATE_V_FIRST: + b_dh1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([BV, 64], dtype=tl.float32) + else: + b_dh1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + q += (bos * H + i_h // (HV // H)).to(tl.int64) * K + k += (bos * H + i_h // (HV // H)).to(tl.int64) * K + w += (bos * HV + i_h).to(tl.int64) * K + do += (bos * HV + i_h).to(tl.int64) * V + dv += (bos * HV + i_h).to(tl.int64) * V + dv2 += (bos * HV + i_h).to(tl.int64) * V + dh += (boh * HV + i_h).to(tl.int64) * K * V + if USE_GK: + gk += (bos * HV + i_h).to(tl.int64) * K + + if USE_INITIAL_STATE: + dh0 += i_nh * K * V + if USE_FINAL_STATE_GRADIENT: + dht += i_nh * K * V + + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K + if USE_FINAL_STATE_GRADIENT: + if STATE_V_FIRST: + p_dht1 = dht + o_v[:, None] * K + o_k1[None, :] + m_dht1 = m_v[:, None] & m_k1[None, :] + else: + p_dht1 = dht + o_k1[:, None] * V + o_v[None, :] + m_dht1 = m_k1[:, None] & m_v[None, :] + b_dh1 += tl.load(p_dht1, mask=m_dht1, other=0.0) + if K > 64: + if STATE_V_FIRST: + p_dht2 = dht + o_v[:, None] * K + o_k2[None, :] + m_dht2 = m_v[:, None] & m_k2[None, :] + else: + p_dht2 = dht + o_k2[:, None] * V + o_v[None, :] + m_dht2 = m_k2[:, None] & m_v[None, :] + b_dh2 += tl.load(p_dht2, mask=m_dht2, other=0.0) + if K > 128: + if STATE_V_FIRST: + p_dht3 = dht + o_v[:, None] * K + o_k3[None, :] + m_dht3 = m_v[:, None] & m_k3[None, :] + else: + p_dht3 = dht + o_k3[:, None] * V + o_v[None, :] + m_dht3 = m_k3[:, None] & m_v[None, :] + b_dh3 += tl.load(p_dht3, mask=m_dht3, other=0.0) + if K > 192: + if STATE_V_FIRST: + p_dht4 = dht + o_v[:, None] * K + o_k4[None, :] + m_dht4 = m_v[:, None] & m_k4[None, :] + else: + p_dht4 = dht + o_k4[:, None] * V + o_v[None, :] + m_dht4 = m_k4[:, None] & m_v[None, :] + b_dh4 += tl.load(p_dht4, mask=m_dht4, other=0.0) + + for i_t in range(NT - 1, -1, -1): + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + if STATE_V_FIRST: + p_dh1 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k1[None, :] + m_dh1 = m_v[:, None] & m_k1[None, :] + else: + p_dh1 = dh + i_t_int64 * HV * K * V + o_k1[:, None] * V + o_v[None, :] + m_dh1 = m_k1[:, None] & m_v[None, :] + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), mask=m_dh1) + if K > 64: + if STATE_V_FIRST: + p_dh2 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k2[None, :] + m_dh2 = m_v[:, None] & m_k2[None, :] + else: + p_dh2 = dh + i_t_int64 * HV * K * V + o_k2[:, None] * V + o_v[None, :] + m_dh2 = m_k2[:, None] & m_v[None, :] + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), mask=m_dh2) + if K > 128: + if STATE_V_FIRST: + p_dh3 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k3[None, :] + m_dh3 = m_v[:, None] & m_k3[None, :] + else: + p_dh3 = dh + i_t_int64 * HV * K * V + o_k3[:, None] * V + o_v[None, :] + m_dh3 = m_k3[:, None] & m_v[None, :] + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), mask=m_dh3) + if K > 192: + if STATE_V_FIRST: + p_dh4 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k4[None, :] + m_dh4 = m_v[:, None] & m_k4[None, :] + else: + p_dh4 = dh + i_t_int64 * HV * K * V + o_k4[:, None] * V + o_v[None, :] + m_dh4 = m_k4[:, None] & m_v[None, :] + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), mask=m_dh4) + + last_idx = min((i_t_int64 + 1) * BT, T) - 1 + if USE_G: + bg_last = tl.load(g + (bos + last_idx) * HV + i_h).to(tl.float32) + p_g = g + bos * HV + i_h + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + bg_last_exp = exp2(bg_last) + b_g_exp = exp2(b_g) + p_dv = dv + o_t[:, None] * (HV * V) + o_v[None, :] + p_dv2 = dv2 + o_t[:, None] * (HV * V) + o_v[None, :] + p_do = do + o_t[:, None] * (HV * V) + o_v[None, :] + + b_do = tl.load(p_do, mask=m_t[:, None] & m_v[None, :], other=0.0) + + # Update dv + p_k = k + o_t[:, None] * (H * K) + o_k1[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k1[None, :], other=0.0) + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load(gk + last_idx * HV * K + o_k1, mask=(o_k1 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh1).to(b_k.dtype)) + else: + b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) + + if K > 64: + p_k = k + o_t[:, None] * (H * K) + o_k2[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k2[None, :], other=0.0) + if USE_GK: + b_gk_last2 = tl.load(gk + last_idx * HV * K + o_k2, mask=(o_k2 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh2).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh2.to(b_k.dtype), b_dv) + + if K > 128: + p_k = k + o_t[:, None] * (H * K) + o_k3[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k3[None, :], other=0.0) + if USE_GK: + b_gk_last3 = tl.load(gk + last_idx * HV * K + o_k3, mask=(o_k3 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh3).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh3.to(b_k.dtype), b_dv) + + if K > 192: + p_k = k + o_t[:, None] * (H * K) + o_k4[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k4[None, :], other=0.0) + if USE_GK: + b_gk_last4 = tl.load(gk + last_idx * HV * K + o_k4, mask=(o_k4 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh4).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh4.to(b_k.dtype), b_dv) + + if USE_G: + b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] + b_dv += tl.load(p_dv, mask=m_t[:, None] & m_v[None, :], other=0.0) + + tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) + # Update dh + p_w = w + o_k1[:, None] + o_t[None, :] * (HV * K) + p_q = q + o_k1[:, None] + o_t[None, :] * (H * K) + b_w = tl.load(p_w, mask=m_k1[:, None] & m_t[None, :], other=0.0) + b_q = tl.load(p_q, mask=m_k1[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh1 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh1 *= exp2(b_gk_last1)[None, :] + else: + b_dh1 *= exp2(b_gk_last1[:, None]) + if STATE_V_FIRST: + b_dh1 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 64: + p_q = q + o_k2[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k2[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k2[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k2[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh2 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh2 *= exp2(b_gk_last2)[None, :] + else: + b_dh2 *= exp2(b_gk_last2[:, None]) + if STATE_V_FIRST: + b_dh2 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 128: + p_q = q + o_k3[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k3[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k3[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k3[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh3 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh3 *= exp2(b_gk_last3)[None, :] + else: + b_dh3 *= exp2(b_gk_last3[:, None]) + if STATE_V_FIRST: + b_dh3 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 192: + p_q = q + o_k4[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k4[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k4[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k4[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh4 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh4 *= exp2(b_gk_last4)[None, :] + else: + b_dh4 *= exp2(b_gk_last4[:, None]) + if STATE_V_FIRST: + b_dh4 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + + if USE_INITIAL_STATE: + if STATE_V_FIRST: + p_dh0 = dh0 + o_v[:, None] * K + o_k1[None, :] + m_dh0 = m_v[:, None] & m_k1[None, :] + else: + p_dh0 = dh0 + o_k1[:, None] * V + o_v[None, :] + m_dh0 = m_k1[:, None] & m_v[None, :] + tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), mask=m_dh0) + if K > 64: + if STATE_V_FIRST: + p_dh1 = dh0 + o_v[:, None] * K + o_k2[None, :] + m_dh1 = m_v[:, None] & m_k2[None, :] + else: + p_dh1 = dh0 + o_k2[:, None] * V + o_v[None, :] + m_dh1 = m_k2[:, None] & m_v[None, :] + tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), mask=m_dh1) + if K > 128: + if STATE_V_FIRST: + p_dh2 = dh0 + o_v[:, None] * K + o_k3[None, :] + m_dh2 = m_v[:, None] & m_k3[None, :] + else: + p_dh2 = dh0 + o_k3[:, None] * V + o_v[None, :] + m_dh2 = m_k3[:, None] & m_v[None, :] + tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), mask=m_dh2) + if K > 192: + if STATE_V_FIRST: + p_dh3 = dh0 + o_v[:, None] * K + o_k4[None, :] + m_dh3 = m_v[:, None] & m_k4[None, :] + else: + p_dh3 = dh0 + o_k4[:, None] * V + o_v[None, :] + m_dh3 = m_k4[:, None] & m_v[None, :] + tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), mask=m_dh3) + + +def chunk_gated_delta_rule_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + B, T, H, K, V, HV = *k.shape, u.shape[-1], u.shape[2] + BT = chunk_size + BV = state_qdq_tile_v(V, state_qdq, state_qdq_block_v) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + if chunk_offsets is None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + assert K <= 256, "current kernel does not support head dimension larger than 256." + + if state_v_first: + h = k.new_empty(B, NT, HV, V, K) + final_state = k.new_zeros(N, HV, V, K, dtype=torch.float32) if output_final_state else None + else: + h = k.new_empty(B, NT, HV, K, V) + final_state = k.new_zeros(N, HV, K, V, dtype=torch.float32) if output_final_state else None + + v_new = torch.empty_like(u) if save_new_value else None + + grid = (triton.cdiv(V, BV) * N * HV,) + chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( + k=k, + v=u, + w=w, + v_new=v_new, + g=g, + gk=gk, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BV=BV, + STATE_V_FIRST=state_v_first, + STATE_QDQ=state_qdq, + ) + return h, v_new, final_state + + +def state_qdq_tile_v(V: int, state_qdq: int, state_qdq_block_v: int | None) -> int: + """Return the V tile width ``BV`` of the forward state kernel. + + Without state quantization this is fla's largest tile. With it, the tile is also the + quantization granularity: one dynamic scale per ``[K, BV]`` block of a head's state. The + default is fla's 64-column tile, i.e. one scale per sequence and head for ``V <= 64`` and two + for the usual ``V == 128``. ``state_qdq_block_v=128`` gives one scale per 128-wide head but + exceeds the register budget where the kernel is limited to two warps (Blackwell) and spills. + """ + if state_qdq == STATE_QDQ_OFF: + return 64 if check_shared_mem("ada") else 32 + if state_qdq != STATE_QDQ_FP8_DYNAMIC: + raise ValueError(f"Unsupported state_qdq mode {state_qdq}; expected 0 or 1.") + BV = min(triton.next_power_of_2(V), 64) if state_qdq_block_v is None else state_qdq_block_v + if BV < 16 or BV > STATE_QDQ_MAX_BLOCK_V or BV & (BV - 1): + raise ValueError( + f"state_qdq_block_v must be a power of two in [16, {STATE_QDQ_MAX_BLOCK_V}], got {BV}." + ) + return BV + + +def chunk_gated_delta_rule_bwd_dhu( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + do: torch.Tensor, + dv: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + h0: torch.Tensor | None = None, + dht: torch.Tensor | None = None, + scale: float | None = None, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, + use_graph: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V, HV = *q.shape, do.shape[-1], do.shape[2] + # N: the actual number of sequences in the batch with either equal or variable lengths + BT = chunk_size + assert K <= 256, "current kernel does not support head dimension being larger than 256." + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + if chunk_offsets is None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + + if use_graph: + # [ModelOpt] Imported here: fla.ops.utils.graph is absent from released fla versions. + from fla.ops.utils.graph import get_static_buffer + + if state_v_first: + dh = get_static_buffer("dhu_dh_vf", (B, NT, HV, V, K), q.dtype, q.device) + else: + dh = get_static_buffer("dhu_dh", (B, NT, HV, K, V), q.dtype, q.device) + dh0 = ( + get_static_buffer("dhu_dh0", tuple(h0.shape), torch.float32, h0.device) + if h0 is not None + else None + ) + dv2 = get_static_buffer("dhu_dv2", tuple(dv.shape), dv.dtype, dv.device) + else: + if state_v_first: + dh = q.new_empty(B, NT, HV, V, K) + else: + dh = q.new_empty(B, NT, HV, K, V) + dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None + dv2 = torch.empty_like(dv) + + def grid(meta): + return (triton.cdiv(V, meta["BV"]) * N * HV,) + + chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64[grid]( + q=q, + k=k, + w=w, + g=g, + gk=gk, + dht=dht, + dh0=dh0, + do=do, + dh=dh, + dv=dv, + dv2=dv2, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + scale=scale, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + STATE_V_FIRST=state_v_first, + ) + return dh, dh0, dv2 diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py new file mode 100644 index 00000000000..a5609cab102 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py @@ -0,0 +1,721 @@ +# Adapted from: https://github.com/fla-org/flash-linear-attention/blob/516143e31fce/fla/ops/gated_delta_rule/chunk.py +# Adapted with modifications (marked [ModelOpt]): threads state_qdq / state_qdq_block_v through +# the autograd function, applies an optional w_quantizer to the WY tensor w, and imports the +# state kernels from the vendored sibling module. +# +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import warnings + +import fla +import torch +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o +from fla.ops.common.gate import fused_beta_sigmoid, fused_beta_sigmoid_bwd +from fla.ops.cp import FLACPContext +from fla.ops.cp.chunk_delta_h import ( + chunk_gated_delta_rule_bwd_dhu_pre_process, + chunk_gated_delta_rule_fwd_h_pre_process, + compress_h0, + expand_h0, +) +from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra +from fla.ops.gated_delta_rule.gate import gdn_gate_bwd, gdn_gate_chunk_cumsum +from fla.ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd +from fla.ops.utils import chunk_local_cumsum +from fla.ops.utils.constant import RCP_LN2 +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import ( + IS_NVIDIA_HOPPER, + TRITON_ABOVE_3_4_0, + autocast_custom_bwd, + autocast_custom_fwd, + input_guard, +) + +from modelopt.torch.quantization.linear_attention.utils import validate_gdn_quantizer +from modelopt.torch.quantization.nn import TensorQuantizer + +from .fla_chunk_delta_h import ( + STATE_QDQ_FP8_DYNAMIC, + STATE_QDQ_OFF, + chunk_gated_delta_rule_bwd_dhu, + chunk_gated_delta_rule_fwd_h, +) + + +def chunk_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + w_quantizer: TensorQuantizer | None = None, +): + g_input = g if use_gate_in_kernel else None + if use_gate_in_kernel: + g = gdn_gate_chunk_cumsum( + g=g, + A_log=A_log, + chunk_size=chunk_size, + scale=RCP_LN2, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + else: + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + scale=RCP_LN2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # obtain WY representation. u is actually the new v. + # fused kkt + solve_tril + recompute_w_u + w, u, A = chunk_gated_delta_rule_fwd_intra( + k=k, + v=v, + g=g, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + # [ModelOpt] w is an activation (the WY form of the chunk's keys) that lands in memory here, + # so it is fake-quantized once per forward instead of tile by tile inside the kernel. + if w_quantizer is not None: + w = w_quantizer(w) + + if cp_context is not None: + initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=k, + w=w, + u=u, + g=g, + cu_seqlens=cu_seqlens, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + ) + + if cp_context is not None: + initial_state = compress_h0(initial_state, context=cp_context) + + o = chunk_fwd_o( + q=q, + k=k, + v=v_new, + h=h, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + return g, o, A, final_state, initial_state, g_input, w if w_quantizer is not None else None + + +def chunk_gated_delta_rule_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + g_input: torch.Tensor | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + quantized_w: torch.Tensor | None = None, +): + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # [ModelOpt] Reuse forward QDQ exactly; the validated policy uses identity STE. + # This also avoids calling observers or dynamic scale computation in backward. + if quantized_w is not None: + w = quantized_w + + if cp_context is not None: + initial_state = expand_h0(initial_state, context=cp_context) + + # [ModelOpt] The backward recomputes the forward's chunk states, so it sees the same + # fake-quantized states; the state gradient itself passes straight through the QDQ. + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + ) + dv = chunk_bwd_dv_local( + q=q, + k=k, + g=g, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + if cp_context is not None: + # initial_state is None in the CP mode + # We only need to compute dht of current rank and pass it to the backward kernel + dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=q, + k=k, + w=w, + do=do, + dv=dv, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + dht=dht, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=q, + k=k, + w=w, + g=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dq, dk, dw, dg = chunk_bwd_dqkwg( + q=q, + k=k, + v=v_new, + w=w, + g=g, + h=h, + dv=dv, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dk2, dv, db, dg2 = prepare_wy_repr_bwd( + k=k, + v=v, + beta=beta, + g=g, + A=A, + dw=dw, + du=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + dk.add_(dk2) + dg.add_(dg2) + dg = chunk_local_cumsum( + dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices + ) + dA_log, ddt_bias = None, None + if use_gate_in_kernel: + dg, dA_log, ddt_bias = gdn_gate_bwd(g=g_input, A_log=A_log, dt_bias=dt_bias, dyg=dg) + return dq, dk, dv, db, dg, dh0, dA_log, ddt_bias + + +class ChunkGatedDeltaRuleFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + cp_context: FLACPContext | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + w_quantizer: TensorQuantizer | None = None, + ): + q_rstd, k_rstd = None, None + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + + beta_raw = beta + if use_beta_sigmoid_in_kernel: + beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu + ) + g, o, A, final_state, initial_state, g_input, quantized_w = chunk_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cp_context=cp_context, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + use_gate_in_kernel=use_gate_in_kernel, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + w_quantizer=w_quantizer, + ) + ctx.save_for_backward( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + quantized_w, + ) + ctx.scale = scale + ctx.chunk_size = chunk_size + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.use_beta_sigmoid_in_kernel = use_beta_sigmoid_in_kernel + ctx.allow_neg_eigval = allow_neg_eigval + ctx.cp_context = cp_context + ctx.state_v_first = state_v_first + ctx.use_gate_in_kernel = use_gate_in_kernel + ctx.state_qdq = state_qdq + ctx.state_qdq_block_v = state_qdq_block_v + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + ( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + quantized_w, + ) = ctx.saved_tensors + dq, dk, dv, db, dg, dh0, dA_log, ddt_bias = chunk_gated_delta_rule_bwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A=A, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + cp_context=ctx.cp_context, + chunk_indices=chunk_indices, + state_v_first=ctx.state_v_first, + use_gate_in_kernel=ctx.use_gate_in_kernel, + g_input=g_input, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=ctx.chunk_size, + state_qdq=ctx.state_qdq, + state_qdq_block_v=ctx.state_qdq_block_v, + quantized_w=quantized_w, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + if ctx.use_beta_sigmoid_in_kernel: + db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0) + return ( + dq.to(q), + dk.to(k), + dv.to(v), + dg.to(g), + db.to(beta_raw), + None, + dh0, + None, + None, + None, + None, + None, + None, + None, + dA_log, + ddt_bias, + None, + None, + None, + None, + None, + None, + None, + ) + + +# [ModelOpt] Not registered with fla's backend dispatch: another backend must not take over a +# call that asks for state quantization. +@torch.compiler.disable +def chunk_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + **kwargs, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`. + g (torch.Tensor): + (forget) gating tensor of shape `[B, T, HV]`. + When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay). + When `use_gate_in_kernel=True`, `g` is the raw input before gate activation; + the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, HV, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2norm to the q/k tensor internally. Default: `False`. + use_gate_in_kernel (bool): + Whether to compute the log-space GDN decay internally. + When `True`, the passed `g` is the raw input, and `A_log` must be provided. + The kernel fuses gate activation + chunk cumsum in a single pass. + Default: `False`. + A_log (Optional[torch.Tensor]): + Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`. + dt_bias (Optional[torch.Tensor]): + Bias added to `g` before activation, of shape `[HV]`. + Only used when `use_gate_in_kernel=True`. + use_beta_sigmoid_in_kernel (bool): + Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel. + - If `True`, the passed `beta` acts as the raw beta logits. + - If `False`, `beta` is expected to already be in post-sigmoid space. + Default: `False`. + allow_neg_eigval (bool): + Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`. + Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case + the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`. + state_v_first (Optional[bool]): + Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + chunk_indices (Optional[torch.LongTensor]): + Pre-computed chunk indices for variable-length inputs. + If provided, they are used directly instead of being computed from `cu_seqlens`. Default: `None`. + cp_context (Optional[FLACPContext]): + Context parallel context for distributed training across multiple devices. + When provided, `initial_state` and `output_final_state` are not supported, + and `cu_seqlens` will be overridden by the context. Default: `None`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda') + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if fla.__version__ != "0.5.1": + raise RuntimeError(f"ModelOpt GDN requires fla-core==0.5.1, got {fla.__version__}.") + if "transpose_state_layout" in kwargs: + if state_v_first: + raise ValueError( + "Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`." + ) + warnings.warn( + "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", + DeprecationWarning, + stacklevel=2, + ) + state_v_first = kwargs.pop("transpose_state_layout") + + # Validate head dimensions + if q.shape[2] != k.shape[2]: + raise ValueError( + f"q and k must have the same number of heads, " + f"but got q.shape[2]={q.shape[2]} and k.shape[2]={k.shape[2]}" + ) + H, HV = q.shape[2], v.shape[2] + if HV % H != 0: + raise ValueError( + f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by " + f"num_heads (H={H}), but got HV % H = {HV % H}" + ) + + if "head_first" in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) + + chunk_size = kwargs.pop("chunk_size", 64) + if chunk_size != 64: + raise ValueError("ModelOpt GDN supports only chunk_size=64; FLA WY backward assumes 64.") + + # [ModelOpt] state_qdq: 0 keeps fla's numerics; 1 fake-quantizes the state carried between + # chunks to FP8 E4M3 with a dynamic scale per [K, state_qdq_block_v] tile of each head. + state_qdq = kwargs.pop("state_qdq", STATE_QDQ_OFF) + state_qdq_block_v = kwargs.pop("state_qdq_block_v", None) + # w_quantizer: dynamic FP8 TensorQuantizer applied to the WY tensor + # ``w`` of shape [B, T, HV, K] before it multiplies the state, emulating an FP8 x FP8 matmul. + w_quantizer = kwargs.pop("w_quantizer", None) + use_gate_in_kernel = kwargs.pop("use_gate_in_kernel", False) + A_log = kwargs.pop("A_log", None) + dt_bias = kwargs.pop("dt_bias", None) + if kwargs: + raise TypeError(f"Unexpected keyword arguments: {', '.join(sorted(kwargs))}") + if state_qdq not in (STATE_QDQ_OFF, STATE_QDQ_FP8_DYNAMIC): + raise ValueError(f"`state_qdq` must be 0 or 1, got {state_qdq}.") + if w_quantizer is not None: + validate_gdn_quantizer(w_quantizer, name="gdn_w_quantizer") + if state_qdq and (not q.is_cuda or torch.cuda.get_device_capability(q.device) < (8, 9)): + raise RuntimeError("GDN state QDQ requires native E4M3 conversion on CUDA SM89 or newer.") + if (state_qdq != STATE_QDQ_OFF or w_quantizer is not None) and cp_context is not None: + raise ValueError("State or w quantization is not supported together with `cp_context`.") + + if cp_context is not None: + assert initial_state is None, "Initial state is not supported for CP" + assert output_final_state is False, "Output final state is not supported for CP" + assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" + cu_seqlens = cp_context.cu_seqlens + if cp_context.cu_seqlens_cpu is not None: + cu_seqlens_cpu = cp_context.cu_seqlens_cpu + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if use_gate_in_kernel: + assert A_log is not None, "A_log must be provided when use_gate_in_kernel=True." + if allow_neg_eigval and not use_beta_sigmoid_in_kernel: + raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") + + if scale is None: + scale = k.shape[-1] ** -0.5 + # [ModelOpt] Hopper's TileLang backward needs BF16 and equal head counts. Expand outside + # custom autograd so repeat_interleave reduces q/k gradients back to the original heads. + if IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0: + if any(x.dtype != torch.bfloat16 for x in (q, k, v)): + raise ValueError("Hopper with Triton >= 3.4 requires BF16 q/k/v for GDN training.") + if H != HV: + q = q.repeat_interleave(HV // H, dim=2) + k = k.repeat_interleave(HV // H, dim=2) + o, final_state = ChunkGatedDeltaRuleFunction.apply( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + state_v_first, + cu_seqlens, + cu_seqlens_cpu, + chunk_indices, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + A_log, + dt_bias, + use_beta_sigmoid_in_kernel, + allow_neg_eigval, + cp_context, + chunk_size, + state_qdq, + state_qdq_block_v, + w_quantizer, + ) + return o, final_state + + +chunk_gdn = chunk_gated_delta_rule diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 7fbd7d10dea..78a6e03a16b 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -67,6 +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", [])) + _validate_linear_attention_quantizers(model) metadata = {} update_quantize_metadata(model, config, metadata) @@ -142,7 +143,14 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: # QuantModule.set_extra_state(). return model - quantizer_state_dict = metadata["quantizer_state"] + 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() unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys() extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys() @@ -202,6 +210,22 @@ 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 _validate_linear_attention_quantizers(model): + for module in _linear_attention_modules(model).values(): + module.validate_linear_attention() + + 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..477eeb89847 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -0,0 +1,16 @@ +# 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. + +"""Helpers for linear-attention quantization.""" diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py new file mode 100644 index 00000000000..4f1e18a174d --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared helpers for linear-attention quantization.""" + +from ..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.""" + 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.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 " + "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 5f1afad1dea..5b53526aba6 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -33,6 +33,7 @@ from modelopt.torch.opt.utils import forward_with_reshard from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.conversion import ( + _validate_linear_attention_quantizers, preserve_quantizer_attributes_context, set_quantizer_attributes_partial, set_quantizer_by_cfg, @@ -391,6 +392,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) # Fail before calibration rather than after exporting an unquantized checkpoint. _check_weight_quantization_took_effect(model, quantize_config) _check_indexer_quantization_took_effect(model, quantize_config) diff --git a/modelopt/torch/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py index ac5edc087f5..7234ca81f76 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..df70683cc2f --- /dev/null +++ b/modelopt/torch/quantization/plugins/gated_delta_net.py @@ -0,0 +1,128 @@ +# 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 functools import partial +from typing import Any + +import torch + +from ..config import QuantizerAttributeConfig +from ..linear_attention.utils import validate_gdn_quantizer +from ..nn import QuantModule, TensorQuantizer + +__all__ = ["GatedDeltaNetStateQuantMixin"] + +GatedDeltaRuleFn = Callable[..., tuple[torch.Tensor, torch.Tensor | None]] + + +def _fla_chunk_gated_delta_rule() -> GatedDeltaRuleFn: + # FLA is an optional, heavy dependency needed only when quantization is enabled. + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + + return chunk_gated_delta_rule + + +def _state_qdq_chunk_gated_delta_rule() -> GatedDeltaRuleFn: + # Imported on first use: flash-linear-attention is a heavy optional dependency that only the + # enabled quantizer needs, and importing it warns on machines without a GPU. + try: + from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( + chunk_gated_delta_rule, + ) + except ImportError as e: + raise RuntimeError( + "GDN fake quantization needs Triton and fla-core==0.5.1 on a CUDA " + f"device; importing the state-quantizing kernel failed with {e!r}." + ) from e + return chunk_gated_delta_rule + + +class GatedDeltaNetStateQuantMixin(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. State QDQ uses fixed 64-column value tiles and 64-token chunks. + """ + + 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() + + 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) + while isinstance(gated_delta_rule, partial): + args = (*gated_delta_rule.args, *args) + kwargs = {**gated_delta_rule.keywords, **kwargs} + gated_delta_rule = gated_delta_rule.func + if gated_delta_rule is not _fla_chunk_gated_delta_rule(): + raise NotImplementedError( + "GatedDeltaNet quantization supports only FLA's chunk_gated_delta_rule " + f"callable or a functools.partial of it; got {gated_delta_rule!r}." + ) + chunk_size = kwargs.pop("chunk_size", 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=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 2980d9e5d30..aaacf198a12 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__ = [] @@ -1074,6 +1082,57 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): return make_sharded_tensors_for_checkpoint(state_dict, prefix, {}, sharded_offsets) +if HAS_GDN: + + @QuantModuleRegistry.register({GatedDeltaNet: "megatron_GatedDeltaNet"}) + class _QuantGatedDeltaNet(GatedDeltaNetStateQuantMixin): + """GatedDeltaNet with fake quantization of the recurrent state at kernel chunk boundaries. + + Routes ``self.gated_delta_rule`` through state/W QDQ from both Megatron's direct + forward and older split-forward layouts. The older dynamic-batching inference + paths (``ssm_prefill`` / ``ssm_decode``) are left untouched. + """ + + def _setup(self): + super()._setup() + try: + data_parallel_group = get_data_parallel_group(with_context_parallel=True) + except AssertionError: + data_parallel_group = get_data_parallel_group() + self.parallel_state = ParallelState( + data_parallel_group, + mcore_parallel.get_tensor_model_parallel_group(), + ) + + @contextmanager + def _quantized_gdn_kernel(self): + gated_delta_rule = self.gated_delta_rule + self.gated_delta_rule = partial( + self._state_quantized_chunk_gated_delta_rule, gated_delta_rule + ) + try: + yield + finally: + self.gated_delta_rule = gated_delta_rule + + def forward(self, *args, **kwargs): + if hasattr(GatedDeltaNet, "forward_pre_attn_and_core_attn"): + return super().forward(*args, **kwargs) + with self._quantized_gdn_kernel(): + return super().forward(*args, **kwargs) + + def forward_pre_attn_and_core_attn(self, *args, **kwargs): + with self._quantized_gdn_kernel(): + return super().forward_pre_attn_and_core_attn(*args, **kwargs) + + def validate_linear_attention(self): + super().validate_linear_attention() + if self.config.context_parallel_size > 1 and ( + self.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 796a19e6e04..4fa5997027d 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -32,5 +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 | | `indexer_k_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer key cache (`*indexer_k_quantizer`), global scale fixed to 1; Blackwell+ GPUs | | `indexer_q_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer query (`*indexer_q_quantizer`), global scale fixed to 1; Blackwell+ GPUs | diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index 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..ed5768ee599 --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Dynamic per-tile FP8 E4M3 fake quantization of the GatedDeltaNet recurrent state. +# Recompute amax over each [K, 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``. + +# 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..2211b0d6f9c 100644 --- a/noxfile.py +++ b/noxfile.py @@ -117,7 +117,7 @@ def gpu(session): "--no-build-isolation", "git+https://github.com/Dao-AILab/fast-hadamard-transform.git", ) - session.run("python", "-m", "pip", "install", "-e", ".[all,dev-test]") + session.run("python", "-m", "pip", "install", "-e", ".[all,dev-test,dev-fla]") session.run("python", "-m", "pip", "uninstall", "-y", "cupy-cuda12x") session.run("python", "-m", "pip", "install", "cupy-cuda13x") session.run( @@ -140,7 +140,7 @@ def gpu_megatron(session): session.run("pip", "uninstall", "-y", "tensorrt_llm") # 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", "-e", ".[hf,dev-test,dev-fla]") session.run("python", "-m", "pytest", "tests/gpu_megatron", *_cov_args()) diff --git a/pyproject.toml b/pyproject.toml index 31bac4b523c..1404766235d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -133,6 +133,11 @@ dev-test = [ "torchvision", "torch-geometric", ] +dev-fla = [ + "apache-tvm-ffi==0.1.9", + "fla-core==0.5.1", + "tilelang==0.1.8", +] # Compound extras via self-references mlflow = [ # Dependencies for modelopt.torch.utils.mlflow run tracking "mlflow-skinny>=2.9", # client only; the tracking server's own deps are not needed here @@ -249,6 +254,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/*" = [ @@ -323,6 +329,15 @@ disable_error_code = ["attr-defined"] module = ["examples.diffusers.fastgen.preprocess.*"] ignore_errors = true +# Vendored from fla-org/flash-linear-attention (MIT); kept faithful to upstream rather than +# annotated to modelopt's strict mypy (Triton loop indices, ``*tensor.shape`` unpacking). +[[tool.mypy.overrides]] +module = [ + "modelopt.torch.kernels.quantization.linear_attention.fla_chunk_delta_h", + "modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule", +] +ignore_errors = true + [tool.bandit] exclude_dirs = [".github/", "examples/", "noxfile.py", "tests/"] # Do not change `skips`. It should be consistent with NVIDIA's Wheel-CI-CD bandit.yml config. diff --git a/tests/_test_utils/torch/quantization/linear_attention_reference.py b/tests/_test_utils/torch/quantization/linear_attention_reference.py new file mode 100644 index 00000000000..65ea183a318 --- /dev/null +++ b/tests/_test_utils/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/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py new file mode 100644 index 00000000000..defe05aa5d5 --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Minimal forward/backward checks for GDN training fake quantization.""" + +import pytest +import torch +import torch.nn.functional as F +from _test_utils.torch.quantization.linear_attention_reference import chunk_gdn_reference + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer + +pytest.importorskip("fla.ops.gated_delta_rule") + +from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( + chunk_gated_delta_rule, +) + + +def make_inputs(): + torch.manual_seed(123) + # Two chunks exercise recurrence with one shared BF16 shape for every mode. + shape = (1, 128, 1, 32) + q, k = [F.normalize(torch.randn(shape, device="cuda"), dim=-1) for _ in range(2)] + v = torch.randn(shape, device="cuda") + g = -torch.rand(shape[:3], device="cuda") * 0.1 + beta = torch.rand_like(g) + args = [x.to(torch.bfloat16).requires_grad_() for x in (q, k, v, g, beta)] + state = (torch.randn(1, 1, 32, 32, device="cuda") * 0.1).requires_grad_() + return args, state + + +def values_and_grads(fn, args, state, **kwargs): + result = fn(*args, initial_state=state, **kwargs) + torch.manual_seed(15) + probes = [torch.randn(x.shape, device=x.device, dtype=torch.float32) for x in result] + grads = torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probes)), (*args, state)) + return result, grads + + +def compare(actual, expected, tolerance): + for a, e in zip(actual, expected): + assert torch.isfinite(a).all() + error = (a.float() - e.float()).norm() + bound = tolerance * e.float().norm().clamp_min(1e-6) + assert error <= bound, ( + f"relative L2 error {(error / e.float().norm()).item():.5g} > {tolerance}" + ) + + +@pytest.fixture(scope="module", params=["disabled", "w", "state-w"]) +def compiled_gdn_case(request): + """Compile only the selected BF16 forward/backward path, outside the test-call timer.""" + state_qdq = request.param == "state-w" + if state_qdq and torch.cuda.get_device_capability() < (8, 9): + pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") + quantizer = ( + TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), axis=(0, 1, 2), type="dynamic")) + if request.param != "disabled" + else None + ) + args, state = make_inputs() + kwargs = {"state_qdq": state_qdq, "w_quantizer": quantizer} + values_and_grads(chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs) + torch.cuda.synchronize() + return args, state, kwargs + + +def test_gdn_forward_and_backward(compiled_gdn_case): + args, state, kwargs = compiled_gdn_case + reference_args = [x.detach().float().requires_grad_() for x in args] + reference_state = state.detach().clone().requires_grad_() + expected = values_and_grads(chunk_gdn_reference, reference_args, reference_state, **kwargs) + actual = values_and_grads( + chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs + ) + compare(actual[0], expected[0], 0.03) + compare(actual[1], expected[1], 0.05) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py new file mode 100644 index 00000000000..0756f2d0d40 --- /dev/null +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -0,0 +1,175 @@ +# 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, + sharded_state_dict_test_helper, +) + +import modelopt.torch.quantization as mtq + +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.megatron import _QuantGatedDeltaNet + +try: + _state_qdq_chunk_gated_delta_rule() +except RuntimeError as e: + pytest.skip(str(e), allow_module_level=True) + +SEED = 1234 + + +def _make_model(tp_size): + model = ( + get_mcore_gpt_model( + tensor_model_parallel_size=tp_size, + num_layers=1, + hidden_size=64, + num_attention_heads=4, + vocab_size=32, + max_sequence_length=128, + experimental_attention_variant="gated_delta_net", + linear_num_key_heads=1, + linear_num_value_heads=1, + linear_key_head_dim=32, + linear_value_head_dim=32, + ) + .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, checkpoint_path, sites): + initialize_for_megatron( + tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED + ) + 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() + + 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" + 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() + # 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() + + for site in sites: + mtq.disable_quantizer(model, f"*gdn_{site}_quantizer") + 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") + + 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(restored_gdn[0].in_proj.weight, before) + + +def _compile_gdn_qat_kernels(rank, size, sites): + 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) + 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(dist_workers_size_1): + """Warm one small QAT model in the same worker, outside the test-call budget.""" + # Ampere exercises W QDQ; native FP8 devices also exercise state QDQ. + sites = ("state", "w") if torch.cuda.get_device_capability() >= (8, 9) else ("w",) + dist_workers_size_1.run(_compile_gdn_qat_kernels, sites) + return dist_workers_size_1, sites + + +def test_gdn_qat_and_sharded_restore(compiled_gdn_workers, tmp_path): + """Train through QDQ after a Megatron distributed-checkpoint round trip.""" + workers, sites = compiled_gdn_workers + workers.run(_test_gdn_qat_helper, tmp_path, sites) 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..3793edcb189 --- /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 +from functools import partial + +import pytest +import torch +import torch.nn as nn + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.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): + """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``.""" + + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4) + self.gated_delta_rule = chunk_gated_delta_rule + + def forward(self, x): + out, _ = self.gated_delta_rule(x, x, x, x[..., 0], x[..., 0]) + return self.proj(out) + + +@QuantModuleRegistry.register({TinyGatedDeltaNet: "TinyGatedDeltaNet"}) +class _QuantTinyGatedDeltaNet(GatedDeltaNetStateQuantMixin): + def forward(self, x): + gated_delta_rule = self.gated_delta_rule + self.gated_delta_rule = lambda *a, **kw: self._state_quantized_chunk_gated_delta_rule( + gated_delta_rule, *a, **kw + ) + try: + return super().forward(x) + finally: + self.gated_delta_rule = gated_delta_rule + + +GDN_W_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1, 2), "type": "dynamic"} + + +def quant_cfg(state=True, w=False): + entries = [{"quantizer_name": "*", "enable": False}] + if state: + entries.append({"quantizer_name": "*gdn_state_quantizer", "cfg": GDN_STATE_FP8_DYNAMIC}) + if w: + entries.append({"quantizer_name": "*gdn_w_quantizer", "cfg": GDN_W_FP8_DYNAMIC}) + return {"quant_cfg": entries, "algorithm": "max"} + + +@pytest.mark.parametrize( + "attributes", + [ + {"num_bits": (4, 3), "axis": (0, 1)}, # static + {"num_bits": (4, 3), "type": "dynamic"}, # per tensor + {"num_bits": 8, "axis": (0, 1), "type": "dynamic"}, # int8 + {"num_bits": (4, 3), "type": "dynamic", "block_sizes": {-1: 16}}, # blockwise + ], +) +def test_validate_state_quantizer_rejects_unsupported(attributes): + with pytest.raises(ValueError, match="supports only"): + mtq.quantize( + TinyGatedDeltaNet(), + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*gdn_state_quantizer", "cfg": attributes}, + ], + "algorithm": None, + }, + ) + + +def test_dynamic_export_removes_linear_attention_attributes(): + model = _QuantTinyGatedDeltaNet.convert(TinyGatedDeltaNet()) + model.export() + assert type(model) is TinyGatedDeltaNet + for name in ("gdn_state_quantizer", "gdn_w_quantizer"): + assert not hasattr(model, name) + + +def test_disabled_state_quantizer_calls_original_kernel(): + model = TinyGatedDeltaNet() + x = torch.randn(2, 8, 3, 4) + expected = model(x) + + disable_all = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": "max"} + mtq.quantize(model, disable_all, lambda m: m(x)) + + assert isinstance(model, _QuantTinyGatedDeltaNet) + assert not model.gdn_state_quantizer.is_enabled and not model.gdn_w_quantizer.is_enabled + assert torch.equal(model(x), expected) + assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" + + +@pytest.mark.parametrize("partial_kernel", [False, True]) +def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, 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 + ) + model = TinyGatedDeltaNet() + if partial_kernel: + model.gated_delta_rule = partial(chunk_gated_delta_rule, output_final_state=True) + x = torch.randn(2, 8, 3, 4) + mtq.quantize(model, quant_cfg(), lambda m: m(x)) + + model(x) + assert calls and calls[-1] == { + "chunk_size": 64, + "state_qdq": 1, + "state_qdq_block_v": 64, + "w_quantizer": None, + **({"output_final_state": True} if partial_kernel else {}), + } + + # An unrelated callable must not pass validation merely by copying the FLA name. + model.gated_delta_rule = lambda *a, **kw: chunk_gated_delta_rule(*a, **kw) + model.gated_delta_rule.__name__ = "chunk_gated_delta_rule" + with pytest.raises(NotImplementedError, match="supports only FLA"): + model(x) + + +@pytest.mark.parametrize("state", [False, True]) +def test_w_quantizer_is_passed_to_the_kernel(monkeypatch, state): + """``*gdn_w_quantizer`` in the config hands the module's TensorQuantizer to the kernel, with + or without the state quantizer.""" + calls = [] + + def fake_state_qdq_kernel(*args, **kwargs): + calls.append(kwargs) + return chunk_gated_delta_rule(*args) + + monkeypatch.setattr( + 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) + + +@pytest.mark.parametrize("axis", [None, (0, 1), (0, 1, 2)]) +def test_w_grouping_is_preserved_during_conversion(axis): + cfg = deepcopy(quant_cfg(state=False, w=True)) + cfg["algorithm"] = None + cfg["quant_cfg"][-1]["cfg"]["axis"] = axis + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + assert model.gdn_w_quantizer.axis == axis + + +@pytest.mark.parametrize(("state", "w"), [(True, False), (False, True), (True, True)]) +def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): + model = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) + cfg = quant_cfg(state=state, w=w) + cfg["algorithm"] = None + mtq.quantize(model, cfg) + 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) + for name, enabled in (("gdn_state_quantizer", state), ("gdn_w_quantizer", w)): + original = getattr(model[0], name) + quantizer = getattr(restored[0], name) + assert quantizer.is_enabled == enabled + assert quantizer.axis == original.axis + assert quantizer.num_bits == original.num_bits + assert quantizer._dynamic == original._dynamic + assert not hasattr(restored[1], name) + + +def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): + cfg = quant_cfg() + cfg["algorithm"] = None + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + assert model.gdn_state_quantizer.is_enabled + assert not model.gdn_w_quantizer.is_enabled + + cfg = quant_cfg(state=False, w=True) + cfg["algorithm"] = None + mtq.quantize(model, cfg) + assert not model.gdn_state_quantizer.is_enabled + assert model.gdn_w_quantizer.is_enabled + + cfg = deepcopy(cfg) + cfg["quant_cfg"][-1]["cfg"]["pass_through_bwd"] = False + with pytest.raises(ValueError, match="supports only"): + mtq.quantize(model, cfg) + + +def test_restore_legacy_gdn_without_new_quantizer_handles(): + config = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None} + model = mtq.quantize(TinyGatedDeltaNet(), config) + state = mto.modelopt_state(model) + for _, mode_state in state["modelopt_state_dict"]: + metadata = mode_state["metadata"] + 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..d71818d0163 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_reference.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +import torch.nn.functional as F +from _test_utils.torch.quantization.linear_attention_reference import ( + chunk_gdn_reference, + recurrent_delta_rule_reference, + state_fp8_qdq_reference, +) + + +def inputs(length=7, kda=False): + torch.manual_seed(81) + q, k = [ + F.normalize(torch.randn(1, length, 1, 3, dtype=torch.float64), dim=-1) for _ in range(2) + ] + v = torch.randn(1, length, 2, 4, dtype=torch.float64) + g = -torch.rand(1, length, 2, *([3] if kda else []), dtype=torch.float64) + beta = torch.rand(1, length, 2, dtype=torch.float64) + return [x.requires_grad_() for x in (q, k, v, g, beta)] + + +@pytest.mark.parametrize("packed", [False, True]) +@pytest.mark.parametrize("state_v_first", [False, True]) +def test_chunk_recurrence_output_state_and_gradients(packed, state_v_first): + args = inputs() + state = torch.randn(2 if packed else 1, 2, 3, 4, dtype=torch.float64) + if state_v_first: + state = state.transpose(-1, -2).contiguous() + state.requires_grad_() + kwargs = { + "initial_state": state, + "state_v_first": state_v_first, + "cu_seqlens": torch.tensor([0, 2, 7]) if packed else None, + } + actual = chunk_gdn_reference(*args, chunk_size=3, **kwargs) + expected = recurrent_delta_rule_reference(*args, **kwargs) + for a, e in zip(actual, expected): + torch.testing.assert_close(a, e, rtol=1e-10, atol=1e-10) + probe = [torch.randn_like(x) for x in actual] + grads = [ + torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probe)), (*args, state)) + for result in (actual, expected) + ] + for a, e in zip(*grads): + torch.testing.assert_close(a, e, rtol=1e-9, atol=1e-9) + + +@pytest.mark.parametrize("kda", [False, True]) +def test_recurrence_smooth_gradcheck(kda): + args = inputs(length=2, kda=kda) + state = torch.randn(1, 2, 3, 4, dtype=torch.float64, requires_grad=True) + assert torch.autograd.gradcheck( + lambda *x: recurrent_delta_rule_reference(*x[:-1], initial_state=x[-1]), + (*args, state), + fast_mode=True, + ) + + +def test_kda_one_step_and_scalar_decay_equivalence(): + q, k, v, g, beta = inputs(length=1, kda=True) + state = torch.randn(1, 2, 3, 4, dtype=torch.float64) + out, final = recurrent_delta_rule_reference(q, k, v, g, beta, initial_state=state) + keys = k[:, 0].expand(-1, 2, -1) + decayed = state * g[:, 0].exp().unsqueeze(-1) + residual = v[:, 0] - torch.einsum("bhk,bhkv->bhv", keys, decayed) + expected = decayed + torch.einsum("bhk,bhv->bhkv", keys, beta[:, 0, :, None] * residual) + torch.testing.assert_close(final, expected) + torch.testing.assert_close( + out[:, 0], torch.einsum("bhk,bhkv->bhv", q[:, 0].expand(-1, 2, -1), expected) / 3**0.5 + ) + scalar = g[..., 0] + gdn = recurrent_delta_rule_reference(q, k, v, scalar, beta, initial_state=state) + kda = recurrent_delta_rule_reference( + q, k, v, scalar.unsqueeze(-1).expand_as(g), beta, initial_state=state + ) + for a, e in zip(kda, gdn): + torch.testing.assert_close(a, e) + + +def test_state_qdq_granularity_zero_tail_and_identity_ste(): + torch.manual_seed(14) + state = torch.randn(2, 2, 3, 37, requires_grad=True) + with torch.no_grad(): + state[0].zero_() + state[1, ..., :16] *= 10 + quantized = state_fp8_qdq_reference(state, block_v=16) + assert torch.isfinite(quantized).all() + assert torch.equal(quantized[0], state[0]) + assert not torch.equal(quantized, state) + assert not torch.equal(quantized, state_fp8_qdq_reference(state, block_v=64)) + probe = torch.randn_like(state) + (grad,) = torch.autograd.grad((quantized * probe).sum(), state) + torch.testing.assert_close(grad, probe, rtol=0, atol=0) + + +def test_packed_sequences_reset_the_state(): + args = inputs(length=7, kda=True) + initial = torch.randn(2, 2, 3, 4, dtype=torch.float64) + packed = recurrent_delta_rule_reference( + *args, initial_state=initial, cu_seqlens=torch.tensor([0, 2, 7]) + ) + separate = [ + recurrent_delta_rule_reference( + *(x[:, lo:hi] for x in args), initial_state=initial[n : n + 1] + ) + for n, (lo, hi) in enumerate(((0, 2), (2, 7))) + ] + torch.testing.assert_close(packed[0], torch.cat([x[0] for x in separate], dim=1)) + torch.testing.assert_close(packed[1], torch.cat([x[1] for x in separate]))