From 07d39f36a06463c5550d92879d2a563be127d8bd Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 21 Sep 2026 19:07:34 -0700 Subject: [PATCH 1/6] Add GDN state and WY activation QAT emulation Signed-off-by: Kai Xu --- .pre-commit-config.yaml | 2 + CHANGELOG.rst | 1 + LICENSE | 1 + docs/linear_attention_qat.md | 130 +++ docs/source/_templates/autosummary/module.rst | 2 +- .../quantization/linear_attention/__init__.py | 24 + .../linear_attention/fla_chunk_delta_h.py | 979 ++++++++++++++++++ .../fla_chunk_gated_delta_rule.py | 705 +++++++++++++ modelopt/torch/quantization/config.py | 7 + modelopt/torch/quantization/conversion.py | 57 +- .../quantization/linear_attention/__init__.py | 19 + .../quantization/linear_attention/config.py | 60 ++ .../linear_attention/reference.py | 203 ++++ .../linear_attention/validation.py | 45 + modelopt/torch/quantization/model_quant.py | 2 + .../torch/quantization/plugins/__init__.py | 1 + .../quantization/plugins/gated_delta_net.py | 115 ++ .../torch/quantization/plugins/megatron.py | 43 +- modelopt_recipes/configs/ptq/units/README.md | 4 + .../units/default_disabled_quantizers.yaml | 5 + .../ptq/units/gdn_state_fp8_dynamic.yaml | 30 + .../configs/ptq/units/gdn_w_fp8_dynamic.yaml | 28 + noxfile.py | 2 + pyproject.toml | 10 + .../test_fla_chunk_gated_delta_rule.py | 212 ++++ .../plugins/test_megatron_gated_delta_net.py | 101 ++ .../plugins/test_gated_delta_net.py | 276 +++++ .../test_linear_attention_reference.py | 125 +++ 28 files changed, 3186 insertions(+), 3 deletions(-) create mode 100644 docs/linear_attention_qat.md create mode 100644 modelopt/torch/kernels/quantization/linear_attention/__init__.py create mode 100644 modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py create mode 100644 modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py create mode 100644 modelopt/torch/quantization/linear_attention/__init__.py create mode 100644 modelopt/torch/quantization/linear_attention/config.py create mode 100644 modelopt/torch/quantization/linear_attention/reference.py create mode 100644 modelopt/torch/quantization/linear_attention/validation.py create mode 100644 modelopt/torch/quantization/plugins/gated_delta_net.py create mode 100644 modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml create mode 100644 modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml create mode 100644 tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py create mode 100644 tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py create mode 100644 tests/unit/torch/quantization/plugins/test_gated_delta_net.py create mode 100644 tests/unit/torch/quantization/test_linear_attention_reference.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7654acab550..e677fb2dff2 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -108,6 +108,8 @@ repos: modelopt/onnx/quantization/ort_patching.py| modelopt/torch/_deploy/utils/onnx_utils.py| modelopt/torch/export/transformer_engine.py| + modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py| + modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py| modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_pruned_to_mxfp4.py| modelopt/torch/quantization/export_onnx.py| modelopt/torch/quantization/plugins/attention.py| diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 22169f974c0..bc6a6ae6e7a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -17,6 +17,7 @@ Changelog - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration. - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md `_ for details. - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``. +- Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training, with saved execution policies and differentiable GDN/KDA references. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer. *Megatron Framework (M-LM / M-Bridge)* diff --git a/LICENSE b/LICENSE index 57c1104ae67..e08eaefc016 100644 --- a/LICENSE +++ b/LICENSE @@ -251,6 +251,7 @@ the following copyright holders, licensed under the MIT License: Copyright (c) 2025 sgl-project Copyright (c) 2026 The DeepSpec Authors Copyright (c) 2023-2026 The ggml authors + Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/docs/linear_attention_qat.md b/docs/linear_attention_qat.md new file mode 100644 index 00000000000..0db84b58f22 --- /dev/null +++ b/docs/linear_attention_qat.md @@ -0,0 +1,130 @@ +# Linear-attention numerical emulation for training + +This experimental integration adds dynamic FP8 E4M3 fake quantization to Megatron +GatedDeltaNet's chunked training path. It rounds the recurrent state at chunk boundaries +and the materialized WY activation `W` before the state-read matmul. Both sites use an +identity straight-through estimator (STE). The underlying tensors remain floating point; +this feature does not provide compressed inference state or native FP8 matmul acceleration. + +The implementation adapts [PR #2455](https://github.com/NVIDIA/Model-Optimizer/pull/2455) +at `13c7e2456f2e9d079c9ef822742eeaa634353802` onto ModelOpt +`051d6adb204f10cd3e78d0f824f31a5a01d54831`. It is the M0/M1 portion of the linear-attention +design. Additional prefill operand sites, KDA training integration, approximate inverse, +token-state quantization, decay approximation, and SSM replay remain later milestones. + +## Requirements and scope + +- Install `fla-core==0.5.1` alongside the training framework. FLA and Triton remain optional + for ordinary ModelOpt imports and the PyTorch references. +- The fused GDN implementation accepts **chunk size 64 only**, including backward. +- State QDQ requires NVIDIA SM89 or newer for the native E4M3 conversion instruction. + W-only QDQ can run on the A6000/SM86 environment used for local validation. +- The numerical recipe supports dynamic E4M3 and `pass_through_bwd=True`. Static scales, + clipping-aware backward, real quantization, rotations, and custom format backends are rejected. +- Context parallelism with either numerical site enabled is rejected. Distributed/sharded + checkpoint qualification remains a release gate; ordinary ModelOpt save/restore is covered. +- Megatron's dynamic-batching `ssm_prefill` and `ssm_decode` paths are outside this integration. + +## Configure training + +Apply the recipe after constructing the Megatron model. The module pattern must match +at least one converted GatedDeltaNet layer; dense attention layers in a hybrid model are +ignored by the execution policy. + +```python +import modelopt.torch.quantization as mtq + +config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*gdn_state_quantizer", + "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1)}, + }, + { + "quantizer_name": "*gdn_w_quantizer", + "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)}, + }, + ], + "algorithm": None, + "linear_attention": [ + { + "module_name": "decoder.layers.*.self_attention", + "cfg": { + "backend": "fla", + "chunk_size": 64, + "state": {"mode": "chunk", "block_v": 64, "quantize_initial": True}, + "solve": {"method": "exact"}, + }, + }, + ], +} +model = mtq.quantize(model, config) +``` + +Dynamic scaling requires no calibration loop. Continue with the framework's normal +forward, backward, and optimizer steps. Enable projection quantization through separate +recipe entries if needed. The standard projection presets keep the GDN numerical handles +disabled; append the GDN units after standard exclusions when composing recipes. + +Execution policies use last-match precedence: each match replaces the complete policy, +including defaults for omitted fields. There is no nested merge. The resolved policy is +reported at conversion and is available on each converted module as +`linear_attention_config`. Quantizer enables/formats and resolved policies are saved +through the ModelOpt checkpoint APIs, including changes made after conversion. + +Without an explicit execution policy, enabled GDN recipe units use the defaults above. +When both GDN handles are disabled, the adapter calls the original training kernel. + +## Numerical contract + +`W` is an activation, **not a learned weight**. The W quantizer sees `[B,T,Hv,Dk]` +and reduces over `Dk`, using one scale per token/head. Forward's quantized W is saved +for backward, avoiding a second quantizer invocation and ensuring identical values. +This costs one saved W activation when that site is enabled. + +The state handle specifies a fused operation rather than calling `TensorQuantizer` on a +materialized state. Each sequence/head uses one dynamic scale per `[Dk, block_v]` tile; +`block_v` can be 16, 32, 64, or 128. Partial value tiles are valid. The scale is +`amax / 448`, with scale one for an all-zero tile. Rounding occurs on the initial state +and after each chunk's state update. The chunk's outputs use its incoming rounded state; +they are computed before rounding the outgoing state. Initial/final state gradients +propagate through these rounding events using identity STE. + +The triangular solve remains exact. No approximate-inverse algorithm is selected or +silently enabled by this integration. + +## References and verification + +The CPU-only reference package is `modelopt.torch.quantization.linear_attention`: + +| Function | Contract | +| --- | --- | +| `recurrent_delta_rule_reference` | Exact GDN scalar-decay or KDA per-key-decay recurrence | +| `chunk_gdn_reference` | Exact GDN triangular solve and chunk algebra; optional state/W QDQ | +| `state_fp8_qdq_reference` | Dynamic E4M3 state tiles with identity STE | + +Use float64 for algebra/gradcheck and float32 for GPU comparison. References support +grouped value heads, packed sequences, nonzero initial states, tails, and either state +layout. They favor explicit arithmetic over speed and are not an efficient training backend. + +```bash +PYTHONPATH=. pytest -q tests/unit/torch/quantization/test_linear_attention_reference.py \ + tests/unit/torch/quantization/plugins/test_gated_delta_net.py \ + tests/unit/torch/quantization/test_quantize_cpu.py +PYTHONPATH=. pytest -q \ + tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py +PYTHONPATH=. pytest -q \ + tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +``` + +The GPU suite compares outputs, final states, and gradients for every differentiable +input against the declared surrogate. It also covers packed boundaries, state layout, +fused gate/beta activation, activation checkpointing, and unsupported-policy errors. +The Megatron test includes an optimizer step. The `gpu` and `gpu_megatron` Nox sessions +install the pinned FLA dependency so missing FLA does not silently remove CI coverage. + +Local validation uses Python 3.12.8, Torch 2.9.1, Triton 3.5.1, FLA 0.5.1, and A6000/SM86. +State-QDQ GPU cases require a separate SM89+ run; the local environment also lacks +Megatron-Core. Those gates must pass before M1 is considered fully qualified. No model +quality, QAT recovery, or throughput result is claimed. diff --git a/docs/source/_templates/autosummary/module.rst b/docs/source/_templates/autosummary/module.rst index 36eb087fbec..c693958c63d 100644 --- a/docs/source/_templates/autosummary/module.rst +++ b/docs/source/_templates/autosummary/module.rst @@ -11,7 +11,7 @@ :recursive: {% for item in modules %} {% set full_item = fullname + '.' + item.split('.')[-1] %} -{% if ('.plugins.' not in full_item or full_item == 'modelopt.torch.opt.plugins.huggingface') and full_item != 'modelopt.torch.quantization.backends.fp8_per_tensor_gemm' %} +{% if ('.plugins.' not in full_item or full_item == 'modelopt.torch.opt.plugins.huggingface') and full_item != 'modelopt.torch.quantization.backends.fp8_per_tensor_gemm' and not full_item.startswith('modelopt.torch.kernels.quantization.linear_attention.fla_') %} {{ full_item }} {% endif %} {%- endfor %} diff --git a/modelopt/torch/kernels/quantization/linear_attention/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/__init__.py new file mode 100644 index 00000000000..3f003f403f8 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/__init__.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Linear-attention kernels for quantization. + +``fla_chunk_delta_h.py`` and ``fla_chunk_gated_delta_rule.py`` are adapted copies of the chunked +GatedDeltaNet kernels of `flash-linear-attention `_ +(``fla.ops.common.chunk_delta_h`` and ``fla.ops.gated_delta_rule.chunk``) that can fake-quantize the +recurrent state carried between chunks to FP8 (``state_qdq``). They still import the surrounding +fla operators, so ``fla-core==0.5.1`` and Triton must be installed to use +them. This package initializer does not import the kernels, so importing it needs neither. +""" diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py new file mode 100644 index 00000000000..ca8b085470d --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -0,0 +1,979 @@ +# Adapted from: https://github.com/fla-org/flash-linear-attention/blob/516143e31fce/fla/ops/common/chunk_delta_h.py +# Adapted with modifications (marked [ModelOpt]): optional in-kernel FP8 E4M3 fake quantization +# of the carried chunk state (STATE_QDQ), BV as an explicit launch argument, no fla backend +# dispatch or config-cache autotune. +# +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import triton +import triton.language as tl +from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from fla.ops.utils.cache import fla_cache_autotune +from fla.ops.utils.op import exp2 +from fla.utils import ( + IS_INTEL, + IS_NVIDIA_BLACKWELL, + IS_NVIDIA_HOPPER, + autotune_cache_kwargs, + check_shared_mem, +) + +from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq + +# ``STATE_QDQ`` modes of the forward state kernel. +STATE_QDQ_OFF = 0 +STATE_QDQ_FP8_DYNAMIC = 1 # FP8 E4M3, one dynamic scale per program tile ([K, BV] of one head) +STATE_QDQ_MAX_BLOCK_V = 128 + + +@triton.jit +def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr): + """[ModelOpt] Dynamic FP8 E4M3 scale over the up-to-four K tiles of one program's state.""" + b_amax = tl.max(tl.abs(b_h1)) + if K > 64: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h2))) + if K > 128: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h3))) + if K > 192: + b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h4))) + return tl.where(b_amax > 0, b_amax / 448.0, 1.0) + + +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] + +# TODO: Triton mainline fixes a Blackwell tl.dot recurrence race. +# Keep this kernel on num_warps=2 for Blackwell until Triton 3.8 is released +# and we re-validate the wider config space. +# Intel needs more warps than NVIDIA here: 8 warps is ~1.5x faster than the best +# config reachable under the [2, 4] cap. +if IS_NVIDIA_BLACKWELL: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2] +elif IS_INTEL: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2, 4, 8, 16] +else: + GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2, 4] + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +# ``BV`` is an explicit argument rather than an autotuned config: with ``STATE_QDQ`` it sets the +# quantization granularity, so it must not vary with the autotuner's choice. +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in GATED_DELTA_RULE_FWD_H_NUM_WARPS + for num_stages in ([2, 3, 4] if check_shared_mem("ampere") else [2, 1]) + ], + key=["H", "HV", "K", "V", "BT", "BV", "STATE_V_FIRST", "STATE_QDQ"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( + k, + v, + w, + v_new, + g, + gk, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + SAVE_NEW_VALUE: tl.constexpr, + STATE_V_FIRST: tl.constexpr, + IS_VARLEN: tl.constexpr, + STATE_QDQ: tl.constexpr, +): + pid = tl.program_id(0) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) + i_n, i_h = i_nh // HV, i_nh % HV + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + if STATE_V_FIRST: + b_h1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([BV, 64], dtype=tl.float32) + else: + b_h1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + h += (boh * HV + i_h).to(tl.int64) * K * V + v += (bos * HV + i_h).to(tl.int64) * V + k += (bos * H + i_h // (HV // H)).to(tl.int64) * K + w += (bos * HV + i_h).to(tl.int64) * K + if SAVE_NEW_VALUE: + v_new += (bos * HV + i_h).to(tl.int64) * V + + if USE_INITIAL_STATE: + h0 = h0 + i_nh * K * V + if STORE_FINAL_STATE: + ht = ht + i_nh * K * V + + # load initial state + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K + if USE_INITIAL_STATE: + if STATE_V_FIRST: + p_h0_1 = h0 + o_v[:, None] * K + o_k1[None, :] + m_h0_1 = m_v[:, None] & m_k1[None, :] + else: + p_h0_1 = h0 + o_k1[:, None] * V + o_v[None, :] + m_h0_1 = m_k1[:, None] & m_v[None, :] + b_h1 += tl.load(p_h0_1, mask=m_h0_1, other=0.0).to(tl.float32) + if K > 64: + if STATE_V_FIRST: + p_h0_2 = h0 + o_v[:, None] * K + o_k2[None, :] + m_h0_2 = m_v[:, None] & m_k2[None, :] + else: + p_h0_2 = h0 + o_k2[:, None] * V + o_v[None, :] + m_h0_2 = m_k2[:, None] & m_v[None, :] + b_h2 += tl.load(p_h0_2, mask=m_h0_2, other=0.0).to(tl.float32) + if K > 128: + if STATE_V_FIRST: + p_h0_3 = h0 + o_v[:, None] * K + o_k3[None, :] + m_h0_3 = m_v[:, None] & m_k3[None, :] + else: + p_h0_3 = h0 + o_k3[:, None] * V + o_v[None, :] + m_h0_3 = m_k3[:, None] & m_v[None, :] + b_h3 += tl.load(p_h0_3, mask=m_h0_3, other=0.0).to(tl.float32) + if K > 192: + if STATE_V_FIRST: + p_h0_4 = h0 + o_v[:, None] * K + o_k4[None, :] + m_h0_4 = m_v[:, None] & m_k4[None, :] + else: + p_h0_4 = h0 + o_k4[:, None] * V + o_v[None, :] + m_h0_4 = m_k4[:, None] & m_v[None, :] + b_h4 += tl.load(p_h0_4, mask=m_h0_4, other=0.0).to(tl.float32) + # [ModelOpt] A state read from an FP8 cache is quantized before the first chunk uses it. + if STATE_QDQ == 1: + if K > 192: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) + elif K > 128: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K) + elif K > 64: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K) + else: + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K) + b_h1 = fp8_scalar_qdq(b_h1, b_scale) + if K > 64: + b_h2 = fp8_scalar_qdq(b_h2, b_scale) + if K > 128: + b_h3 = fp8_scalar_qdq(b_h3, b_scale) + if K > 192: + b_h4 = fp8_scalar_qdq(b_h4, b_scale) + + # main recurrence + for i_t in range(NT): + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + if STATE_V_FIRST: + p_h1 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k1[None, :] + m_h1 = m_v[:, None] & m_k1[None, :] + else: + p_h1 = h + i_t_int64 * HV * K * V + o_k1[:, None] * V + o_v[None, :] + m_h1 = m_k1[:, None] & m_v[None, :] + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), mask=m_h1) + if K > 64: + if STATE_V_FIRST: + p_h2 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k2[None, :] + m_h2 = m_v[:, None] & m_k2[None, :] + else: + p_h2 = h + i_t_int64 * HV * K * V + o_k2[:, None] * V + o_v[None, :] + m_h2 = m_k2[:, None] & m_v[None, :] + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), mask=m_h2) + if K > 128: + if STATE_V_FIRST: + p_h3 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k3[None, :] + m_h3 = m_v[:, None] & m_k3[None, :] + else: + p_h3 = h + i_t_int64 * HV * K * V + o_k3[:, None] * V + o_v[None, :] + m_h3 = m_k3[:, None] & m_v[None, :] + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), mask=m_h3) + if K > 192: + if STATE_V_FIRST: + p_h4 = h + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k4[None, :] + m_h4 = m_v[:, None] & m_k4[None, :] + else: + p_h4 = h + i_t_int64 * HV * K * V + o_k4[:, None] * V + o_v[None, :] + m_h4 = m_k4[:, None] & m_v[None, :] + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), mask=m_h4) + + p_w = w + o_t[:, None] * (HV * K) + o_k1[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k1[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) + else: + b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) + if K > 64: + p_w = w + o_t[:, None] * (HV * K) + o_k2[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k2[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h2.to(b_w.dtype), b_v) + if K > 128: + p_w = w + o_t[:, None] * (HV * K) + o_k3[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k3[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h3.to(b_w.dtype), b_v) + if K > 192: + p_w = w + o_t[:, None] * (HV * K) + o_k4[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k4[None, :], other=0.0) + if STATE_V_FIRST: + b_v = tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype), b_v) + else: + b_v = tl.dot(b_w, b_h4.to(b_w.dtype), b_v) + p_v = v + o_t[:, None] * (HV * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & m_v[None, :], other=0.0) - b_v + + if SAVE_NEW_VALUE: + p_v = v_new + o_t[:, None] * (HV * V) + o_v[None, :] + tl.store(p_v, b_v.to(p_v.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) + + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + b_g_last = tl.load(g + (bos * HV + last_idx * HV + i_h).to(tl.int64)).to(tl.float32) + p_g = g + (bos * HV + i_h).to(tl.int64) + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] + b_g_last = exp2(b_g_last) + b_h1 *= b_g_last + if K > 64: + b_h2 *= b_g_last + if K > 128: + b_h3 *= b_g_last + if K > 192: + b_h4 *= b_g_last + + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k1, mask=(o_k1 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h1 *= exp2(b_gk_last1)[None, :] + else: + b_h1 *= exp2(b_gk_last1)[:, None] + if K > 64: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k2, mask=(o_k2 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h2 *= exp2(b_gk_last2)[None, :] + else: + b_h2 *= exp2(b_gk_last2)[:, None] + if K > 128: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k3, mask=(o_k3 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h3 *= exp2(b_gk_last3)[None, :] + else: + b_h3 *= exp2(b_gk_last3)[:, None] + if K > 192: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load( + gk + (bos + last_idx) * HV * K + i_h * K + o_k4, mask=(o_k4 < K), other=0.0 + ).to(tl.float32) + if STATE_V_FIRST: + b_h4 *= exp2(b_gk_last4)[None, :] + else: + b_h4 *= exp2(b_gk_last4)[:, None] + b_v = b_v.to(k.dtype.element_ty) + + p_k = k + o_k1[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k1[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h1 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h1 = tl.dot(b_k, b_v, b_h1) + if K > 64: + p_k = k + o_k2[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k2[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h2 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h2 = tl.dot(b_k, b_v, b_h2) + if K > 128: + p_k = k + o_k3[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k3[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h3 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h3 = tl.dot(b_k, b_v, b_h3) + if K > 192: + p_k = k + o_k4[:, None] + o_t[None, :] * (H * K) + b_k = tl.load(p_k, mask=m_k4[:, None] & m_t[None, :], other=0.0) + if STATE_V_FIRST: + b_h4 += tl.trans(tl.dot(b_k, b_v)) + else: + b_h4 = tl.dot(b_k, b_v, b_h4) + + # [ModelOpt] Fake-quantize the state carried into the next chunk (and, after the last + # chunk, the stored final state) to FP8 E4M3. The scale is dynamic over this program's + # [K, BV] tile of the head state; BV == V makes it one scale per sequence and head. + if STATE_QDQ == 1: + if K > 192: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) + elif K > 128: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K) + elif K > 64: + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K) + else: + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K) + b_h1 = fp8_scalar_qdq(b_h1, b_scale) + if K > 64: + b_h2 = fp8_scalar_qdq(b_h2, b_scale) + if K > 128: + b_h3 = fp8_scalar_qdq(b_h3, b_scale) + if K > 192: + b_h4 = fp8_scalar_qdq(b_h4, b_scale) + + if STORE_FINAL_STATE: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k1[None, :] + m_ht = m_v[:, None] & m_k1[None, :] + else: + p_ht = ht + o_k1[:, None] * V + o_v[None, :] + m_ht = m_k1[:, None] & m_v[None, :] + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 64: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k2[None, :] + m_ht = m_v[:, None] & m_k2[None, :] + else: + p_ht = ht + o_k2[:, None] * V + o_v[None, :] + m_ht = m_k2[:, None] & m_v[None, :] + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 128: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k3[None, :] + m_ht = m_v[:, None] & m_k3[None, :] + else: + p_ht = ht + o_k3[:, None] * V + o_v[None, :] + m_ht = m_k3[:, None] & m_v[None, :] + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), mask=m_ht) + if K > 192: + if STATE_V_FIRST: + p_ht = ht + o_v[:, None] * K + o_k4[None, :] + m_ht = m_v[:, None] & m_k4[None, :] + else: + p_ht = ht + o_k4[:, None] * V + o_v[None, :] + m_ht = m_k4[:, None] & m_v[None, :] + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), mask=m_ht) + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "USE_INITIAL_STATE": lambda args: args["dh0"] is not None, + "USE_FINAL_STATE_GRADIENT": lambda args: args["dht"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@fla_cache_autotune( + configs=[ + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in ([2, 3, 4] if check_shared_mem("ampere") else [1]) + for BV in ([32, 64] if check_shared_mem("ada") else [32]) + ], + key=["H", "HV", "K", "V", "BT", "BV", "USE_G", "STATE_V_FIRST"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( + q, + k, + w, + g, + gk, + dht, + dh0, + do, + dh, + dv, + dv2, + cu_seqlens, + chunk_offsets, + scale, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, + STATE_V_FIRST: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + pid = tl.program_id(0) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) + i_n, i_h = i_nh // HV, i_nh % HV + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + if STATE_V_FIRST: + b_dh1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([BV, 64], dtype=tl.float32) + else: + b_dh1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_dh2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_dh3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_dh4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + q += (bos * H + i_h // (HV // H)).to(tl.int64) * K + k += (bos * H + i_h // (HV // H)).to(tl.int64) * K + w += (bos * HV + i_h).to(tl.int64) * K + do += (bos * HV + i_h).to(tl.int64) * V + dv += (bos * HV + i_h).to(tl.int64) * V + dv2 += (bos * HV + i_h).to(tl.int64) * V + dh += (boh * HV + i_h).to(tl.int64) * K * V + if USE_GK: + gk += (bos * HV + i_h).to(tl.int64) * K + + if USE_INITIAL_STATE: + dh0 += i_nh * K * V + if USE_FINAL_STATE_GRADIENT: + dht += i_nh * K * V + + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K + if USE_FINAL_STATE_GRADIENT: + if STATE_V_FIRST: + p_dht1 = dht + o_v[:, None] * K + o_k1[None, :] + m_dht1 = m_v[:, None] & m_k1[None, :] + else: + p_dht1 = dht + o_k1[:, None] * V + o_v[None, :] + m_dht1 = m_k1[:, None] & m_v[None, :] + b_dh1 += tl.load(p_dht1, mask=m_dht1, other=0.0) + if K > 64: + if STATE_V_FIRST: + p_dht2 = dht + o_v[:, None] * K + o_k2[None, :] + m_dht2 = m_v[:, None] & m_k2[None, :] + else: + p_dht2 = dht + o_k2[:, None] * V + o_v[None, :] + m_dht2 = m_k2[:, None] & m_v[None, :] + b_dh2 += tl.load(p_dht2, mask=m_dht2, other=0.0) + if K > 128: + if STATE_V_FIRST: + p_dht3 = dht + o_v[:, None] * K + o_k3[None, :] + m_dht3 = m_v[:, None] & m_k3[None, :] + else: + p_dht3 = dht + o_k3[:, None] * V + o_v[None, :] + m_dht3 = m_k3[:, None] & m_v[None, :] + b_dh3 += tl.load(p_dht3, mask=m_dht3, other=0.0) + if K > 192: + if STATE_V_FIRST: + p_dht4 = dht + o_v[:, None] * K + o_k4[None, :] + m_dht4 = m_v[:, None] & m_k4[None, :] + else: + p_dht4 = dht + o_k4[:, None] * V + o_v[None, :] + m_dht4 = m_k4[:, None] & m_v[None, :] + b_dh4 += tl.load(p_dht4, mask=m_dht4, other=0.0) + + for i_t in range(NT - 1, -1, -1): + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + if STATE_V_FIRST: + p_dh1 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k1[None, :] + m_dh1 = m_v[:, None] & m_k1[None, :] + else: + p_dh1 = dh + i_t_int64 * HV * K * V + o_k1[:, None] * V + o_v[None, :] + m_dh1 = m_k1[:, None] & m_v[None, :] + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), mask=m_dh1) + if K > 64: + if STATE_V_FIRST: + p_dh2 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k2[None, :] + m_dh2 = m_v[:, None] & m_k2[None, :] + else: + p_dh2 = dh + i_t_int64 * HV * K * V + o_k2[:, None] * V + o_v[None, :] + m_dh2 = m_k2[:, None] & m_v[None, :] + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), mask=m_dh2) + if K > 128: + if STATE_V_FIRST: + p_dh3 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k3[None, :] + m_dh3 = m_v[:, None] & m_k3[None, :] + else: + p_dh3 = dh + i_t_int64 * HV * K * V + o_k3[:, None] * V + o_v[None, :] + m_dh3 = m_k3[:, None] & m_v[None, :] + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), mask=m_dh3) + if K > 192: + if STATE_V_FIRST: + p_dh4 = dh + i_t_int64 * HV * K * V + o_v[:, None] * K + o_k4[None, :] + m_dh4 = m_v[:, None] & m_k4[None, :] + else: + p_dh4 = dh + i_t_int64 * HV * K * V + o_k4[:, None] * V + o_v[None, :] + m_dh4 = m_k4[:, None] & m_v[None, :] + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), mask=m_dh4) + + last_idx = min((i_t_int64 + 1) * BT, T) - 1 + if USE_G: + bg_last = tl.load(g + (bos + last_idx) * HV + i_h).to(tl.float32) + p_g = g + bos * HV + i_h + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + bg_last_exp = exp2(bg_last) + b_g_exp = exp2(b_g) + p_dv = dv + o_t[:, None] * (HV * V) + o_v[None, :] + p_dv2 = dv2 + o_t[:, None] * (HV * V) + o_v[None, :] + p_do = do + o_t[:, None] * (HV * V) + o_v[None, :] + + b_do = tl.load(p_do, mask=m_t[:, None] & m_v[None, :], other=0.0) + + # Update dv + p_k = k + o_t[:, None] * (H * K) + o_k1[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k1[None, :], other=0.0) + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load(gk + last_idx * HV * K + o_k1, mask=(o_k1 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh1).to(b_k.dtype)) + else: + b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) + + if K > 64: + p_k = k + o_t[:, None] * (H * K) + o_k2[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k2[None, :], other=0.0) + if USE_GK: + b_gk_last2 = tl.load(gk + last_idx * HV * K + o_k2, mask=(o_k2 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh2).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh2.to(b_k.dtype), b_dv) + + if K > 128: + p_k = k + o_t[:, None] * (H * K) + o_k3[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k3[None, :], other=0.0) + if USE_GK: + b_gk_last3 = tl.load(gk + last_idx * HV * K + o_k3, mask=(o_k3 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh3).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh3.to(b_k.dtype), b_dv) + + if K > 192: + p_k = k + o_t[:, None] * (H * K) + o_k4[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k4[None, :], other=0.0) + if USE_GK: + b_gk_last4 = tl.load(gk + last_idx * HV * K + o_k4, mask=(o_k4 < K), other=0.0).to( + tl.float32 + ) + if STATE_V_FIRST: + b_dv = tl.dot(b_k, tl.trans(b_dh4).to(b_k.dtype), b_dv) + else: + b_dv = tl.dot(b_k, b_dh4.to(b_k.dtype), b_dv) + + if USE_G: + b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] + b_dv += tl.load(p_dv, mask=m_t[:, None] & m_v[None, :], other=0.0) + + tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) + # Update dh + p_w = w + o_k1[:, None] + o_t[None, :] * (HV * K) + p_q = q + o_k1[:, None] + o_t[None, :] * (H * K) + b_w = tl.load(p_w, mask=m_k1[:, None] & m_t[None, :], other=0.0) + b_q = tl.load(p_q, mask=m_k1[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh1 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh1 *= exp2(b_gk_last1)[None, :] + else: + b_dh1 *= exp2(b_gk_last1[:, None]) + if STATE_V_FIRST: + b_dh1 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 64: + p_q = q + o_k2[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k2[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k2[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k2[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh2 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh2 *= exp2(b_gk_last2)[None, :] + else: + b_dh2 *= exp2(b_gk_last2[:, None]) + if STATE_V_FIRST: + b_dh2 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 128: + p_q = q + o_k3[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k3[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k3[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k3[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh3 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh3 *= exp2(b_gk_last3)[None, :] + else: + b_dh3 *= exp2(b_gk_last3[:, None]) + if STATE_V_FIRST: + b_dh3 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + if K > 192: + p_q = q + o_k4[:, None] + o_t[None, :] * (H * K) + p_w = w + o_k4[:, None] + o_t[None, :] * (HV * K) + b_q = tl.load(p_q, mask=m_k4[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k4[:, None] & m_t[None, :], other=0.0) + if USE_G: + b_dh4 *= bg_last_exp + b_q = b_q * b_g_exp[None, :] + if USE_GK: + if STATE_V_FIRST: + b_dh4 *= exp2(b_gk_last4)[None, :] + else: + b_dh4 *= exp2(b_gk_last4[:, None]) + if STATE_V_FIRST: + b_dh4 += tl.trans( + tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale + - tl.dot(b_w, b_dv.to(b_w.dtype)) + ) + else: + b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot( + b_w, b_dv.to(b_w.dtype) + ) + + if USE_INITIAL_STATE: + if STATE_V_FIRST: + p_dh0 = dh0 + o_v[:, None] * K + o_k1[None, :] + m_dh0 = m_v[:, None] & m_k1[None, :] + else: + p_dh0 = dh0 + o_k1[:, None] * V + o_v[None, :] + m_dh0 = m_k1[:, None] & m_v[None, :] + tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), mask=m_dh0) + if K > 64: + if STATE_V_FIRST: + p_dh1 = dh0 + o_v[:, None] * K + o_k2[None, :] + m_dh1 = m_v[:, None] & m_k2[None, :] + else: + p_dh1 = dh0 + o_k2[:, None] * V + o_v[None, :] + m_dh1 = m_k2[:, None] & m_v[None, :] + tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), mask=m_dh1) + if K > 128: + if STATE_V_FIRST: + p_dh2 = dh0 + o_v[:, None] * K + o_k3[None, :] + m_dh2 = m_v[:, None] & m_k3[None, :] + else: + p_dh2 = dh0 + o_k3[:, None] * V + o_v[None, :] + m_dh2 = m_k3[:, None] & m_v[None, :] + tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), mask=m_dh2) + if K > 192: + if STATE_V_FIRST: + p_dh3 = dh0 + o_v[:, None] * K + o_k4[None, :] + m_dh3 = m_v[:, None] & m_k4[None, :] + else: + p_dh3 = dh0 + o_k4[:, None] * V + o_v[None, :] + m_dh3 = m_k4[:, None] & m_v[None, :] + tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), mask=m_dh3) + + +def chunk_gated_delta_rule_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + B, T, H, K, V, HV = *k.shape, u.shape[-1], u.shape[2] + BT = chunk_size + BV = state_qdq_tile_v(V, state_qdq, state_qdq_block_v) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + if chunk_offsets is None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + assert K <= 256, "current kernel does not support head dimension larger than 256." + + if state_v_first: + h = k.new_empty(B, NT, HV, V, K) + final_state = k.new_zeros(N, HV, V, K, dtype=torch.float32) if output_final_state else None + else: + h = k.new_empty(B, NT, HV, K, V) + final_state = k.new_zeros(N, HV, K, V, dtype=torch.float32) if output_final_state else None + + v_new = torch.empty_like(u) if save_new_value else None + + grid = (triton.cdiv(V, BV) * N * HV,) + chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( + k=k, + v=u, + w=w, + v_new=v_new, + g=g, + gk=gk, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BV=BV, + STATE_V_FIRST=state_v_first, + STATE_QDQ=state_qdq, + ) + return h, v_new, final_state + + +def state_qdq_tile_v(V: int, state_qdq: int, state_qdq_block_v: int | None) -> int: + """Return the V tile width ``BV`` of the forward state kernel. + + Without state quantization this is fla's largest tile. With it, the tile is also the + quantization granularity: one dynamic scale per ``[K, BV]`` block of a head's state. The + default is fla's 64-column tile, i.e. one scale per sequence and head for ``V <= 64`` and two + for the usual ``V == 128``. ``state_qdq_block_v=128`` gives one scale per 128-wide head but + exceeds the register budget where the kernel is limited to two warps (Blackwell) and spills. + """ + if state_qdq == STATE_QDQ_OFF: + return 64 if check_shared_mem("ada") else 32 + if state_qdq != STATE_QDQ_FP8_DYNAMIC: + raise ValueError(f"Unsupported state_qdq mode {state_qdq}; expected 0 or 1.") + BV = min(triton.next_power_of_2(V), 64) if state_qdq_block_v is None else state_qdq_block_v + if BV < 16 or BV > STATE_QDQ_MAX_BLOCK_V or BV & (BV - 1): + raise ValueError( + f"state_qdq_block_v must be a power of two in [16, {STATE_QDQ_MAX_BLOCK_V}], got {BV}." + ) + return BV + + +def chunk_gated_delta_rule_bwd_dhu( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + do: torch.Tensor, + dv: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + h0: torch.Tensor | None = None, + dht: torch.Tensor | None = None, + scale: float | None = None, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, + use_graph: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V, HV = *q.shape, do.shape[-1], do.shape[2] + # N: the actual number of sequences in the batch with either equal or variable lengths + BT = chunk_size + assert K <= 256, "current kernel does not support head dimension being larger than 256." + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + if chunk_offsets is None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + + if use_graph: + # [ModelOpt] Imported here: fla.ops.utils.graph is absent from released fla versions. + from fla.ops.utils.graph import get_static_buffer + + if state_v_first: + dh = get_static_buffer("dhu_dh_vf", (B, NT, HV, V, K), q.dtype, q.device) + else: + dh = get_static_buffer("dhu_dh", (B, NT, HV, K, V), q.dtype, q.device) + dh0 = ( + get_static_buffer("dhu_dh0", tuple(h0.shape), torch.float32, h0.device) + if h0 is not None + else None + ) + dv2 = get_static_buffer("dhu_dv2", tuple(dv.shape), dv.dtype, dv.device) + else: + if state_v_first: + dh = q.new_empty(B, NT, HV, V, K) + else: + dh = q.new_empty(B, NT, HV, K, V) + dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None + dv2 = torch.empty_like(dv) + + def grid(meta): + return (triton.cdiv(V, meta["BV"]) * N * HV,) + + chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64[grid]( + q=q, + k=k, + w=w, + g=g, + gk=gk, + dht=dht, + dh0=dh0, + do=do, + dh=dh, + dv=dv, + dv2=dv2, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + scale=scale, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + STATE_V_FIRST=state_v_first, + ) + return dh, dh0, dv2 diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py new file mode 100644 index 00000000000..aa88b49d71f --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py @@ -0,0 +1,705 @@ +# Adapted from: https://github.com/fla-org/flash-linear-attention/blob/516143e31fce/fla/ops/gated_delta_rule/chunk.py +# Adapted with modifications (marked [ModelOpt]): threads state_qdq / state_qdq_block_v through +# the autograd function, applies an optional w_quantizer to the WY tensor w, and imports the +# state kernels from the vendored sibling module. +# +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import warnings + +import fla +import torch +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o +from fla.ops.common.gate import fused_beta_sigmoid, fused_beta_sigmoid_bwd +from fla.ops.cp import FLACPContext +from fla.ops.cp.chunk_delta_h import ( + chunk_gated_delta_rule_bwd_dhu_pre_process, + chunk_gated_delta_rule_fwd_h_pre_process, + compress_h0, + expand_h0, +) +from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra +from fla.ops.gated_delta_rule.gate import gdn_gate_bwd, gdn_gate_chunk_cumsum +from fla.ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd +from fla.ops.utils import chunk_local_cumsum +from fla.ops.utils.constant import RCP_LN2 +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + +from modelopt.torch.quantization.linear_attention.validation import validate_gdn_quantizer +from modelopt.torch.quantization.nn import TensorQuantizer + +from .fla_chunk_delta_h import ( + STATE_QDQ_FP8_DYNAMIC, + STATE_QDQ_OFF, + chunk_gated_delta_rule_bwd_dhu, + chunk_gated_delta_rule_fwd_h, +) + + +def chunk_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + w_quantizer: TensorQuantizer | None = None, +): + g_input = g if use_gate_in_kernel else None + if use_gate_in_kernel: + g = gdn_gate_chunk_cumsum( + g=g, + A_log=A_log, + chunk_size=chunk_size, + scale=RCP_LN2, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + else: + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + scale=RCP_LN2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # obtain WY representation. u is actually the new v. + # fused kkt + solve_tril + recompute_w_u + w, u, A = chunk_gated_delta_rule_fwd_intra( + k=k, + v=v, + g=g, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + # [ModelOpt] w is an activation (the WY form of the chunk's keys) that lands in memory here, + # so it is fake-quantized once per forward instead of tile by tile inside the kernel. + if w_quantizer is not None: + w = w_quantizer(w) + + if cp_context is not None: + initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=k, + w=w, + u=u, + g=g, + cu_seqlens=cu_seqlens, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + ) + + if cp_context is not None: + initial_state = compress_h0(initial_state, context=cp_context) + + o = chunk_fwd_o( + q=q, + k=k, + v=v_new, + h=h, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + return g, o, A, final_state, initial_state, g_input, w if w_quantizer is not None else None + + +def chunk_gated_delta_rule_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + g_input: torch.Tensor | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + quantized_w: torch.Tensor | None = None, +): + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # [ModelOpt] Reuse forward QDQ exactly; the validated policy uses identity STE. + # This also avoids calling observers or dynamic scale computation in backward. + if quantized_w is not None: + w = quantized_w + + if cp_context is not None: + initial_state = expand_h0(initial_state, context=cp_context) + + # [ModelOpt] The backward recomputes the forward's chunk states, so it sees the same + # fake-quantized states; the state gradient itself passes straight through the QDQ. + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + ) + dv = chunk_bwd_dv_local( + q=q, + k=k, + g=g, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + if cp_context is not None: + # initial_state is None in the CP mode + # We only need to compute dht of current rank and pass it to the backward kernel + dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=q, + k=k, + w=w, + do=do, + dv=dv, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + dht=dht, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=q, + k=k, + w=w, + g=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dq, dk, dw, dg = chunk_bwd_dqkwg( + q=q, + k=k, + v=v_new, + w=w, + g=g, + h=h, + dv=dv, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dk2, dv, db, dg2 = prepare_wy_repr_bwd( + k=k, + v=v, + beta=beta, + g=g, + A=A, + dw=dw, + du=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + dk.add_(dk2) + dg.add_(dg2) + dg = chunk_local_cumsum( + dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices + ) + dA_log, ddt_bias = None, None + if use_gate_in_kernel: + dg, dA_log, ddt_bias = gdn_gate_bwd(g=g_input, A_log=A_log, dt_bias=dt_bias, dyg=dg) + return dq, dk, dv, db, dg, dh0, dA_log, ddt_bias + + +class ChunkGatedDeltaRuleFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + cp_context: FLACPContext | None = None, + chunk_size: int = 64, + state_qdq: int = STATE_QDQ_OFF, + state_qdq_block_v: int | None = None, + w_quantizer: TensorQuantizer | None = None, + ): + q_rstd, k_rstd = None, None + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + + beta_raw = beta + if use_beta_sigmoid_in_kernel: + beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices( + cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu + ) + g, o, A, final_state, initial_state, g_input, quantized_w = chunk_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cp_context=cp_context, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + use_gate_in_kernel=use_gate_in_kernel, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=chunk_size, + state_qdq=state_qdq, + state_qdq_block_v=state_qdq_block_v, + w_quantizer=w_quantizer, + ) + ctx.save_for_backward( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + quantized_w, + ) + ctx.scale = scale + ctx.chunk_size = chunk_size + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.use_beta_sigmoid_in_kernel = use_beta_sigmoid_in_kernel + ctx.allow_neg_eigval = allow_neg_eigval + ctx.cp_context = cp_context + ctx.state_v_first = state_v_first + ctx.use_gate_in_kernel = use_gate_in_kernel + ctx.state_qdq = state_qdq + ctx.state_qdq_block_v = state_qdq_block_v + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + ( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + quantized_w, + ) = ctx.saved_tensors + dq, dk, dv, db, dg, dh0, dA_log, ddt_bias = chunk_gated_delta_rule_bwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A=A, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + cp_context=ctx.cp_context, + chunk_indices=chunk_indices, + state_v_first=ctx.state_v_first, + use_gate_in_kernel=ctx.use_gate_in_kernel, + g_input=g_input, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=ctx.chunk_size, + state_qdq=ctx.state_qdq, + state_qdq_block_v=ctx.state_qdq_block_v, + quantized_w=quantized_w, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + if ctx.use_beta_sigmoid_in_kernel: + db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0) + return ( + dq.to(q), + dk.to(k), + dv.to(v), + dg.to(g), + db.to(beta_raw), + None, + dh0, + None, + None, + None, + None, + None, + None, + None, + dA_log, + ddt_bias, + None, + None, + None, + None, + None, + None, + None, + ) + + +# [ModelOpt] Not registered with fla's backend dispatch: another backend must not take over a +# call that asks for state quantization. +@torch.compiler.disable +def chunk_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + **kwargs, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`. + g (torch.Tensor): + (forget) gating tensor of shape `[B, T, HV]`. + When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay). + When `use_gate_in_kernel=True`, `g` is the raw input before gate activation; + the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, HV, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2norm to the q/k tensor internally. Default: `False`. + use_gate_in_kernel (bool): + Whether to compute the log-space GDN decay internally. + When `True`, the passed `g` is the raw input, and `A_log` must be provided. + The kernel fuses gate activation + chunk cumsum in a single pass. + Default: `False`. + A_log (Optional[torch.Tensor]): + Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`. + dt_bias (Optional[torch.Tensor]): + Bias added to `g` before activation, of shape `[HV]`. + Only used when `use_gate_in_kernel=True`. + use_beta_sigmoid_in_kernel (bool): + Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel. + - If `True`, the passed `beta` acts as the raw beta logits. + - If `False`, `beta` is expected to already be in post-sigmoid space. + Default: `False`. + allow_neg_eigval (bool): + Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`. + Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case + the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`. + state_v_first (Optional[bool]): + Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + chunk_indices (Optional[torch.LongTensor]): + Pre-computed chunk indices for variable-length inputs. + If provided, they are used directly instead of being computed from `cu_seqlens`. Default: `None`. + cp_context (Optional[FLACPContext]): + Context parallel context for distributed training across multiple devices. + When provided, `initial_state` and `output_final_state` are not supported, + and `cu_seqlens` will be overridden by the context. Default: `None`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda') + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if fla.__version__ != "0.5.1": + raise RuntimeError(f"ModelOpt GDN requires fla-core==0.5.1, got {fla.__version__}.") + if "transpose_state_layout" in kwargs: + if state_v_first: + raise ValueError( + "Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`." + ) + warnings.warn( + "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", + DeprecationWarning, + stacklevel=2, + ) + state_v_first = kwargs.pop("transpose_state_layout") + + # Validate head dimensions + if q.shape[2] != k.shape[2]: + raise ValueError( + f"q and k must have the same number of heads, " + f"but got q.shape[2]={q.shape[2]} and k.shape[2]={k.shape[2]}" + ) + H, HV = q.shape[2], v.shape[2] + if HV % H != 0: + raise ValueError( + f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by " + f"num_heads (H={H}), but got HV % H = {HV % H}" + ) + + if "head_first" in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) + + chunk_size = kwargs.pop("chunk_size", 64) + if chunk_size != 64: + raise ValueError("ModelOpt GDN supports only chunk_size=64; FLA WY backward assumes 64.") + + # [ModelOpt] state_qdq: 0 keeps fla's numerics; 1 fake-quantizes the state carried between + # chunks to FP8 E4M3 with a dynamic scale per [K, state_qdq_block_v] tile of each head. + state_qdq = kwargs.pop("state_qdq", STATE_QDQ_OFF) + state_qdq_block_v = kwargs.pop("state_qdq_block_v", None) + # w_quantizer: dynamic FP8 TensorQuantizer applied to the WY tensor + # ``w`` of shape [B, T, HV, K] before it multiplies the state, emulating an FP8 x FP8 matmul. + w_quantizer = kwargs.pop("w_quantizer", None) + if state_qdq not in (STATE_QDQ_OFF, STATE_QDQ_FP8_DYNAMIC): + raise ValueError(f"`state_qdq` must be 0 or 1, got {state_qdq}.") + if w_quantizer is not None: + validate_gdn_quantizer(w_quantizer, state=False) + if state_qdq and (not q.is_cuda or torch.cuda.get_device_capability(q.device) < (8, 9)): + raise RuntimeError("GDN state QDQ requires native E4M3 conversion on CUDA SM89 or newer.") + if (state_qdq != STATE_QDQ_OFF or w_quantizer is not None) and cp_context is not None: + raise ValueError("State or w quantization is not supported together with `cp_context`.") + + if cp_context is not None: + assert initial_state is None, "Initial state is not supported for CP" + assert output_final_state is False, "Output final state is not supported for CP" + assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" + cu_seqlens = cp_context.cu_seqlens + if cp_context.cu_seqlens_cpu is not None: + cu_seqlens_cpu = cp_context.cu_seqlens_cpu + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + use_gate_in_kernel = kwargs.get("use_gate_in_kernel", False) + A_log = kwargs.get("A_log") + dt_bias = kwargs.get("dt_bias") + if use_gate_in_kernel: + assert A_log is not None, "A_log must be provided when use_gate_in_kernel=True." + if allow_neg_eigval and not use_beta_sigmoid_in_kernel: + raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") + + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = ChunkGatedDeltaRuleFunction.apply( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + state_v_first, + cu_seqlens, + cu_seqlens_cpu, + chunk_indices, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + A_log, + dt_bias, + use_beta_sigmoid_in_kernel, + allow_neg_eigval, + cp_context, + chunk_size, + state_qdq, + state_qdq_block_v, + w_quantizer, + ) + return o, final_state + + +chunk_gdn = chunk_gated_delta_rule diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index f8847aa1fc1..e8db5c8446a 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -161,6 +161,8 @@ from modelopt.torch.opt.config_loader import load_config from modelopt.torch.utils.network import ConstructorLike +from .linear_attention.config import LinearAttentionPolicyEntry + class QuantizerCfgEntry(ModeloptBaseConfig): """A single entry in a ``quant_cfg`` list.""" @@ -1567,6 +1569,11 @@ def _dict_to_entry(key: str, value) -> list[dict[str, Any]]: class QuantizeConfig(ModeloptBaseConfig): """Default configuration for ``quantize`` mode.""" + linear_attention: list[LinearAttentionPolicyEntry] = ModeloptField( + default=[], + title="Execution policies for supported linear-attention training modules", + ) + quant_cfg: QuantizeQuantCfgType = ModeloptField( default=[{"quantizer_name": "*", "cfg": {"num_bits": 8, "axis": None}}], title="Quantization configuration", diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 7fbd7d10dea..0edccdf4649 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -28,6 +28,7 @@ from modelopt.torch.opt.dynamic import _DMRegistryCls from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict from modelopt.torch.utils import get_unwrapped_name +from modelopt.torch.utils.logging import print_rank_0 from .config import ( QuantizeConfig, @@ -36,6 +37,7 @@ _QuantizeExportConfig, normalize_quant_cfg_list, ) +from .linear_attention.config import LinearAttentionConfig from .nn import ( QuantModule, QuantModuleRegistry, @@ -67,6 +69,7 @@ def convert_to_quantized_model(model: ModelLikeModule, config: QuantizeConfig) - replace_quant_module(model, version=ModeloptStateManager(model).state_version) set_quantizer_by_cfg(model, config.get("quant_cfg", [])) + _apply_linear_attention_policy(model, config) metadata = {} update_quantize_metadata(model, config, metadata) @@ -136,13 +139,22 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: details regarding how MCore sharded checkpoint is restored, see modelopt.torch.opt.plugins.mcore_dist_checkpointing.restore_sharded_modelopt_state. """ + _apply_linear_attention_policy(model, config, metadata.get("linear_attention")) if "quantizer_state" not in metadata: # MCore sharded checkpoint (`torch-dist`) has its quantizer_state stored as the # extra_state of `QuantModule`. The quantizer_state is resumed with # QuantModule.set_extra_state(). return model - quantizer_state_dict = metadata["quantizer_state"] + quantizer_state_dict = dict(metadata["quantizer_state"]) + if "linear_attention" not in metadata: + # Older checkpoints predate these disabled handles; preserve their baseline path. + for name, module in _linear_attention_modules(model).items(): + for handle in ("gdn_state_quantizer", "gdn_w_quantizer"): + key = f"{name}.{handle}" if name else handle + quantizer = getattr(module, handle) + if key not in quantizer_state_dict and not quantizer.is_enabled: + quantizer_state_dict[key] = quantizer.get_modelopt_state() unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys() extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys() @@ -195,6 +207,12 @@ def update_quantize_metadata( model: nn.Module, config: QuantizeConfig, metadata: MetadataDict ) -> None: """Update the quantizer state in the metadata dict.""" + policies = { + name: module.linear_attention_config.model_dump() + for name, module in _linear_attention_modules(model).items() + } + if policies: + metadata["linear_attention"] = policies metadata["quantizer_state"] = quantizer_state(model) if shared_state_metadata := SharedWeightGlobalAmaxState.metadata(model): metadata["shared_quant_states"] = shared_state_metadata @@ -202,6 +220,43 @@ def update_quantize_metadata( metadata.pop("shared_quant_states", None) +def _linear_attention_modules(model): + # Optional framework plugins import conversion; defer this import to avoid that cycle. + from .plugins.gated_delta_net import GatedDeltaNetStateQuantMixin + + return { + get_unwrapped_name(name, model): module + for name, module in model.named_modules() + if isinstance(module, GatedDeltaNetStateQuantMixin) + } + + +def _apply_linear_attention_policy(model, config, saved_policies=None): + modules = _linear_attention_modules(model) + # getattr also handles pickled pre-extension config objects with no field in __dict__. + for entry in getattr(config, "linear_attention", []): + matches = [name for name in modules if fnmatch.fnmatch(name, entry.module_name)] + if not matches: + raise ValueError( + f"linear_attention rule {entry.module_name!r} matches no supported modules" + ) + for name in matches: + modules[name].linear_attention_config = entry.cfg.model_copy(deep=True) + if saved_policies is not None: + if saved_policies.keys() != modules.keys(): + raise ApplyModeError( + "Saved linear_attention policies do not match the restored modules" + ) + for name, policy in saved_policies.items(): + modules[name].linear_attention_config = LinearAttentionConfig(**policy) + for name, module in modules.items(): + module.validate_linear_attention() + if getattr(config, "linear_attention", []): + print_rank_0( + f"Linear attention {name!r}: {module.linear_attention_config.model_dump()}" + ) + + def quantizer_state(model: nn.Module) -> dict[str, Any]: """Returns the quantizer state dict describing the quantizer states in the model.""" return { diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py new file mode 100644 index 00000000000..49f5ea3514a --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Numerical policies and differentiable references for linear attention.""" + +from .config import * +from .reference import * diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py new file mode 100644 index 00000000000..6405fbbf329 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Saved execution policy for GDN training-time numerical emulation.""" + +from typing import Literal + +from pydantic import Field + +from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField + +__all__ = ["LinearAttentionConfig", "LinearAttentionPolicyEntry"] + + +class _StateConfig(ModeloptBaseConfig): + mode: Literal["chunk"] = ModeloptField(default="chunk") + block_v: Literal[16, 32, 64, 128] = ModeloptField(default=64) + quantize_initial: Literal[True] = ModeloptField(default=True) + + +class _SolveConfig(ModeloptBaseConfig): + method: Literal["exact"] = ModeloptField(default="exact") + + +class LinearAttentionConfig(ModeloptBaseConfig): + """GDN chunk-64 policy; unsupported numerical modes fail config validation. + + ``state.block_v`` defines one dynamic scale per ``[Dk, block_v]`` tile of each + sequence/head. The initial state and every chunk's final state are rounded when + ``gdn_state_quantizer`` is enabled. Outputs use the incoming rounded state. + """ + + schema_version: Literal[1] = ModeloptField(default=1) + backend: Literal["fla"] = ModeloptField(default="fla") + chunk_size: Literal[64] = ModeloptField(default=64) + state: _StateConfig = ModeloptField(default=_StateConfig()) + solve: _SolveConfig = ModeloptField(default=_SolveConfig()) + + +class LinearAttentionPolicyEntry(ModeloptBaseConfig): + """Assign a complete policy to supported modules matching ``module_name``. + + Rules apply in order: the last match wins, without merging nested fields. + A rule must match at least one supported linear-attention module. + """ + + module_name: str = Field(...) + cfg: LinearAttentionConfig = ModeloptField(default=LinearAttentionConfig()) diff --git a/modelopt/torch/quantization/linear_attention/reference.py b/modelopt/torch/quantization/linear_attention/reference.py new file mode 100644 index 00000000000..65ea183a318 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/reference.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Small differentiable oracles, independent of FLA and Triton. + +These references prioritize explicit arithmetic and gradients over speed. Gates are +natural logarithms; inputs are already normalized/activated. Packed sequence boundaries +are read on the CPU. Accumulation follows the input dtype (use float64 for algebra tests). +""" + +from collections.abc import Callable +from itertools import pairwise + +import torch + +__all__ = ["chunk_gdn_reference", "recurrent_delta_rule_reference", "state_fp8_qdq_reference"] + + +def state_fp8_qdq_reference(state: torch.Tensor, block_v: int = 64) -> torch.Tensor: + """Dynamic E4M3 QDQ with identity STE and one scale per ``[Dk, block_v]`` tile. + + ``state`` is in key-first layout ``[..., Dk, Dv]``. Scales and rounding do not + contribute derivatives. The zero tile uses scale one; partial value tiles are valid. + """ + if block_v not in (16, 32, 64, 128): + raise ValueError("block_v must be 16, 32, 64, or 128") + with torch.no_grad(): + rounded = [] + for tile in state.float().split(block_v, dim=-1): + amax = tile.abs().amax(dim=(-2, -1), keepdim=True) + scale = torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)) + rounded.append((tile / scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() * scale) + quantized = torch.cat(rounded, dim=-1).to(state.dtype) + return state + (quantized - state).detach() + + +def _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first): + if q.ndim != 4 or k.shape != q.shape or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k must have shape [B,T,Hk,Dk] and v shape [B,T,Hv,Dv]") + batch, length, heads, keys = q.shape + value_heads, values = v.shape[2:] + if length == 0 or value_heads % heads: + raise ValueError("nonempty sequences and Hv divisible by Hk are required") + if beta.shape != (batch, length, value_heads) or g.shape not in ( + beta.shape, + (*beta.shape, keys), + ): + raise ValueError("beta must be [B,T,Hv]; g must be [B,T,Hv] or [B,T,Hv,Dk]") + if cu_seqlens is None: + sequences = [(b, 0, length) for b in range(batch)] + else: + boundaries = cu_seqlens.tolist() + if ( + batch != 1 + or len(boundaries) < 2 + or boundaries[0] != 0 + or boundaries[-1] != length + or any(a >= b for a, b in pairwise(boundaries)) + ): + raise ValueError("cu_seqlens must partition a packed batch of size one") + sequences = [(0, a, b) for a, b in pairwise(boundaries)] + if initial_state is None: + state = q.new_zeros(len(sequences), value_heads, keys, values) + else: + state = initial_state.transpose(-1, -2) if state_v_first else initial_state + if state.shape != (len(sequences), value_heads, keys, values): + raise ValueError("initial_state shape does not match sequence/head dimensions") + return ( + q.repeat_interleave(value_heads // heads, dim=2), + k.repeat_interleave(value_heads // heads, dim=2), + state, + sequences, + ) + + +def recurrent_delta_rule_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + state_v_first: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """GDN (scalar decay) or KDA (per-key decay) exact token recurrence. + + Inputs have shapes ``q,k:[B,T,Hk,Dk]``, ``v:[B,T,Hv,Dv]``, + ``beta:[B,T,Hv]``, and ``g:[B,T,Hv]`` (GDN) or ``[B,T,Hv,Dk]`` (KDA). + State is ``[N,Hv,Dk,Dv]``, or ``[N,Hv,Dv,Dk]`` with ``state_v_first``. + Returns output and final state, preserving gradients through the initial state. + """ + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + sequence_output = [] + for t in range(start, end): + decay = g[b, t].exp() + if g.ndim == 3: + decay = decay.unsqueeze(-1) + decayed = state * decay.unsqueeze(-1) + residual = v[b, t] - (k[b, t].unsqueeze(-1) * decayed).sum(-2) + update = beta[b, t].unsqueeze(-1) * residual + state = decayed + k[b, t].unsqueeze(-1) * update.unsqueeze(-2) + sequence_output.append((q[b, t].unsqueeze(-1) * state).sum(-2) * scale) + outputs.append(torch.stack(sequence_output)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final + + +def chunk_gdn_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + chunk_size: int = 64, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + state_v_first: bool = False, + state_qdq: bool = False, + state_qdq_block_v: int = 64, + w_quantizer: Callable[[torch.Tensor], torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Exact GDN chunk algebra with optional state/W fake quantization. + + The solve is a unit-lower triangular solve. ``w_quantizer`` sees the complete + materialized ``[B,T,Hv,Dk]`` WY operand once, with its own autograd semantics. + State QDQ occurs on the initial state and each chunk's final state, after readout. + """ + if g.ndim != 3 or chunk_size <= 0: + raise ValueError("chunk_gdn_reference requires scalar GDN gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + chunks, all_w = [], [] + for n, (b, start, end) in enumerate(sequences): + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + gc = g[b, lo:hi].transpose(0, 1).cumsum(-1) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + # Mask before exp: upper-triangle positive differences can overflow for long decay. + causal = torch.ones(hi - lo, hi - lo, device=q.device, dtype=torch.bool).tril() + decay = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() + gram = kc @ kc.transpose(-1, -2) + lower = (bc * gram * decay).tril(-1) + matrix = lower + torch.eye(hi - lo, device=q.device, dtype=q.dtype) + rhs = torch.cat((bc * vc, bc * kc * gc.exp().unsqueeze(-1)), dim=-1) + solved = torch.linalg.solve_triangular(matrix, rhs, upper=False, unitriangular=True) + u, w = solved.split((v.shape[-1], k.shape[-1]), dim=-1) + chunks.append((n, qc, kc, gc, decay, u)) + all_w.append(w.transpose(0, 1)) + w = torch.cat(all_w).reshape(q.shape) + if w_quantizer is not None: + w = w_quantizer(w) + w = w.reshape(-1, *w.shape[2:]) + outputs, finals, offset, previous_n = [], [], 0, -1 + for n, qc, kc, gc, decay, u in chunks: + if n != previous_n: + state = states[n] + if state_qdq: + state = state_fp8_qdq_reference(state, state_qdq_block_v) + previous_n = n + length = qc.shape[1] + wc = w[offset : offset + length].transpose(0, 1) + offset += length + updated_values = u - wc @ state # state_read + local_scores = ((qc * scale) @ kc.transpose(-1, -2) * decay).tril() + output = (qc * (scale * gc.exp()).unsqueeze(-1)) @ state + output = output + local_scores @ updated_values # local_readout + outputs.append(output.transpose(0, 1)) + weighted_keys = kc * (gc[..., -1:] - gc).exp().unsqueeze(-1) + state = state * gc[..., -1].exp()[:, None, None] + state = state + weighted_keys.transpose(-1, -2) @ updated_values # state_update + if state_qdq: + state = state_fp8_qdq_reference(state, state_qdq_block_v) + if len(finals) <= n: + finals.append(state) + else: + finals[n] = state + output = torch.cat(outputs).reshape(*v.shape) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final diff --git a/modelopt/torch/quantization/linear_attention/validation.py b/modelopt/torch/quantization/linear_attention/validation.py new file mode 100644 index 00000000000..2bf07525e22 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/validation.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Capabilities of the initial fused GDN fake-quant path.""" + +from ..nn import TensorQuantizer + +__all__ = [] + + +def validate_gdn_quantizer(quantizer: TensorQuantizer, *, state: bool) -> None: + name = "gdn_state_quantizer" if state else "gdn_w_quantizer" + axis = (0, 1) if state else (0, 1, 2) + if not isinstance(quantizer, TensorQuantizer): + raise ValueError(f"{name} requires a single TensorQuantizer") + if not ( + quantizer._dynamic + and quantizer.num_bits == (4, 3) + and quantizer.axis == axis + and quantizer.block_sizes is None + and quantizer.fake_quant + and quantizer._pass_through_bwd + and not quantizer.rotate_is_enabled + and quantizer.pre_quant_scale is None + and quantizer.backend is None + and not quantizer._bias + and not quantizer._use_constant_amax + ): + raise ValueError( + f"{name} supports only dynamic E4M3 fake quantization with axis={axis}, " + "pass_through_bwd=True, no block_sizes, rotation, pre-scaling, bias, constant " + "amax, or custom backend. Other gradient rules and formats are not implemented." + ) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index d7ef04f750d..5aac82cacfe 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -32,6 +32,7 @@ from modelopt.torch.opt.utils import forward_with_reshard from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.conversion import ( + _apply_linear_attention_policy, preserve_quantizer_attributes_context, set_quantizer_attributes_partial, set_quantizer_by_cfg, @@ -329,6 +330,7 @@ def forward_loop(model) -> None: else: # Already quantized, so lets apply the quant_cfg from the config set_quantizer_by_cfg(model, quantize_config.quant_cfg) + _apply_linear_attention_policy(model, quantize_config) # Fail before calibration rather than after exporting an unquantized checkpoint. _check_weight_quantization_took_effect(model, quantize_config) return calibrate(model, config.get("algorithm"), forward_loop=forward_loop) diff --git a/modelopt/torch/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py index 22b4bc2e3cc..963452ff509 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -39,6 +39,7 @@ from .attention import * from .custom import * +from .gated_delta_net import * with import_plugin("diffusers"): from .diffusion.diffusers import * diff --git a/modelopt/torch/quantization/plugins/gated_delta_net.py b/modelopt/torch/quantization/plugins/gated_delta_net.py new file mode 100644 index 00000000000..7be046b8800 --- /dev/null +++ b/modelopt/torch/quantization/plugins/gated_delta_net.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fake quantization of the GatedDeltaNet (GDN) recurrent state. + +The chunked gated-delta-rule kernel keeps each head's ``[K, V]`` recurrent state in fp32 inside +one Triton launch and carries it from chunk to chunk. To emulate a deployment that stores that +state in FP8, ModelOpt runs an adapted copy of the kernel +(:mod:`modelopt.torch.kernels.quantization.linear_attention`) that fake-quantizes the state to +E4M3 at the end of every chunk, with a scale computed inside the kernel from the state itself. +The backward pass recomputes the same quantized states and passes the state gradient straight +through the quantization, so QAT and QAD train against the quantized recurrence. A second +quantizer covers ``w``, the WY-transformed keys that multiply the state; ``w`` is a regular tensor, +so it is fake-quantized by the ``TensorQuantizer`` itself before the kernel reads it. +""" + +from collections.abc import Callable +from typing import Any + +import torch + +from ..config import QuantizerAttributeConfig +from ..linear_attention.config import LinearAttentionConfig +from ..linear_attention.validation import validate_gdn_quantizer +from ..nn import QuantModule, TensorQuantizer + +__all__ = ["GatedDeltaNetStateQuantMixin"] + +GatedDeltaRuleFn = Callable[..., tuple[torch.Tensor, torch.Tensor | None]] + + +def _state_qdq_chunk_gated_delta_rule() -> GatedDeltaRuleFn: + # Imported on first use: flash-linear-attention is a heavy optional dependency that only the + # enabled quantizer needs, and importing it warns on machines without a GPU. + try: + from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( + chunk_gated_delta_rule, + ) + except ImportError as e: + raise RuntimeError( + "GDN fake quantization needs Triton and fla-core==0.5.1 on a CUDA " + f"device; importing the state-quantizing kernel failed with {e!r}." + ) from e + return chunk_gated_delta_rule + + +class GatedDeltaNetStateQuantMixin(QuantModule): + """Adds ``gdn_state_quantizer`` and ``gdn_w_quantizer`` to a GatedDeltaNet module. + + Subclasses route the module's chunked gated-delta-rule call through + :meth:`_state_quantized_chunk_gated_delta_rule`. Both quantizers start disabled; enable them + with ``quant_cfg`` entries on ``*gdn_state_quantizer`` / ``*gdn_w_quantizer`` such as the + ``configs/ptq/units/gdn_state_fp8_dynamic`` and ``gdn_w_fp8_dynamic`` recipe units. The state + quantizer carries the fused QDQ configuration. Both sites currently require dynamic + E4M3 and identity STE. The execution policy is saved in ModelOpt metadata. + """ + + def _setup(self): + self.gdn_state_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) + self.gdn_w_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) + self.linear_attention_config = LinearAttentionConfig() + + @property + def gdn_state_qdq_block_v(self) -> int: + """Value-column scale grouping from the saved execution policy.""" + return self.linear_attention_config.state.block_v + + def validate_linear_attention(self) -> None: + """Reject numerical settings that the fused training path cannot implement.""" + for state, quantizer in ((True, self.gdn_state_quantizer), (False, self.gdn_w_quantizer)): + if quantizer.is_enabled: + validate_gdn_quantizer(quantizer, state=state) + + def modelopt_post_restore(self, prefix: str = ""): + """Validate restored quantizers before using the saved execution policy.""" + super().modelopt_post_restore(prefix) + self.validate_linear_attention() + + def _state_quantized_chunk_gated_delta_rule( + self, gated_delta_rule: GatedDeltaRuleFn, *args: Any, **kwargs: Any + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Call ``gated_delta_rule`` or, if a quantizer is on, the vendored quantizing copy.""" + self.validate_linear_attention() + quantize_state = self.gdn_state_quantizer.is_enabled and self.gdn_state_quantizer._if_quant + quantize_w = self.gdn_w_quantizer.is_enabled + if not (quantize_state or quantize_w): + return gated_delta_rule(*args, **kwargs) + if getattr(gated_delta_rule, "__name__", None) != "chunk_gated_delta_rule": + raise NotImplementedError( + "GatedDeltaNet quantizers require the fla chunked kernel; the deterministic torch " + f"kernel ({gated_delta_rule!r}) is not supported." + ) + chunk_size = kwargs.pop("chunk_size", self.linear_attention_config.chunk_size) + if chunk_size != self.linear_attention_config.chunk_size: + raise ValueError("GDN fake quantization supports only chunk_size=64") + return _state_qdq_chunk_gated_delta_rule()( + *args, + chunk_size=chunk_size, + state_qdq=int(quantize_state), + state_qdq_block_v=self.gdn_state_qdq_block_v, + w_quantizer=self.gdn_w_quantizer if quantize_w else None, + **kwargs, + ) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index ec1958649dd..3a0d4627e38 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -21,7 +21,7 @@ import textwrap import types from contextlib import contextmanager -from functools import cache +from functools import cache, partial from typing import Any import megatron.core.parallel_state as mcore_parallel @@ -61,6 +61,7 @@ from ..utils import sync_moe_expert_amax from ..utils.layerwise_calib import LayerActivationCollector from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear +from .gated_delta_net import GatedDeltaNetStateQuantMixin try: from megatron.core.extensions.transformer_engine import ( @@ -79,6 +80,13 @@ except ImportError: HAS_TE = False +try: + from megatron.core.ssm.gated_delta_net import GatedDeltaNet + + HAS_GDN = True +except ImportError: + HAS_GDN = False + __all__ = [] @@ -1069,6 +1077,39 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): return make_sharded_tensors_for_checkpoint(state_dict, prefix, {}, sharded_offsets) +if HAS_GDN: + + @QuantModuleRegistry.register({GatedDeltaNet: "megatron_GatedDeltaNet"}) + class _QuantGatedDeltaNet(GatedDeltaNetStateQuantMixin): + """GatedDeltaNet with fake quantization of the recurrent state at kernel chunk boundaries. + + Training and calibration reach the chunked kernel through ``forward_pre_attn_and_core_attn``, + which calls ``self.gated_delta_rule``; that call is routed through ``gdn_state_quantizer``. + The dynamic-batching inference paths (``ssm_prefill`` / ``ssm_decode``) are left untouched. + """ + + def _setup(self): + super()._setup() + try: + data_parallel_group = get_data_parallel_group(with_context_parallel=True) + except AssertionError: + data_parallel_group = get_data_parallel_group() + self.parallel_state = ParallelState( + data_parallel_group, + mcore_parallel.get_tensor_model_parallel_group(), + ) + + def forward_pre_attn_and_core_attn(self, *args, **kwargs): + gated_delta_rule = self.gated_delta_rule + self.gated_delta_rule = partial( + self._state_quantized_chunk_gated_delta_rule, gated_delta_rule + ) + try: + return super().forward_pre_attn_and_core_attn(*args, **kwargs) + finally: + self.gated_delta_rule = gated_delta_rule + + def _is_supported_megatron_model(model: torch.nn.Module) -> bool: return isinstance(model, MegatronModule) diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index db37b9222ca..4b6ccec395f 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -32,3 +32,7 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `experts_nvfp4.yaml` | NVFP4 W4A4 on `*.experts.*` weight/input quantizers | | `mixer_mlp_nvfp4.yaml` | NVFP4 W4A4 on dense `*.mixer.{up,down}_proj` weight/input quantizers | | `attention_qkv_fp8.yaml` | FP8 E4M3 on attention q/k/v bmm and softmax quantizers | +| `gdn_state_fp8_dynamic.yaml` | FP8 E4M3 dynamic fake quantization of the GatedDeltaNet recurrent state (per sequence, head and 64-column tile) at every kernel chunk boundary; needs `fla-core==0.5.1`, Triton, and SM89+ | +| `gdn_w_fp8_dynamic.yaml` | FP8 E4M3 dynamic (per token and head) fake quantization of the WY tensor `w` that multiplies the GatedDeltaNet state; requires identity STE, `fla-core==0.5.1`, and Triton | + +See the [linear-attention training guide](../../../../docs/linear_attention_qat.md) for GDN policy, hardware, checkpoint, and gradient requirements. diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index 2fee6279973..cac4f6d80aa 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -16,6 +16,11 @@ # QuantizerCfgList snippet for standard module patterns that should remain unquantized. # modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + # Linear-attention numerical emulation is opt-in, independently of projection quantization. + - quantizer_name: '*gdn_state_quantizer' + enable: false + - quantizer_name: '*gdn_w_quantizer' + enable: false - quantizer_name: '*block_sparse_moe.gate*' enable: false - quantizer_name: '*linear_attn.conv1d*' diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml new file mode 100644 index 00000000000..51f9ad1de20 --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# FP8 E4M3 fake quantization of the GatedDeltaNet (linear attention) recurrent state with a +# dynamic scale computed inside the kernel from the state itself (no calibration). The state is +# ``[N, H, K, V]`` for ``N`` sequences; ``axis: [0, 1]`` scales each sequence's head state on its +# own, in tiles of ``gdn_state_qdq_block_v`` value columns (default: the kernel's 64-column tile, +# two scales per head for V = 128; set 128 on the quantized module for one scale per head). This +# is the only configuration the kernel implements. Applied at every kernel chunk boundary +# (64 tokens) during calibration, QAT and QAD; needs fla-core==0.5.1 and Triton. +# See ``modelopt.torch.quantization.plugins.gated_delta_net``. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*gdn_state_quantizer' + cfg: + num_bits: e4m3 + axis: [0, 1] + type: dynamic diff --git a/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml new file mode 100644 index 00000000000..2c73fdd6eaa --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# FP8 E4M3 fake quantization of ``w``, the WY-transformed keys that multiply the GatedDeltaNet +# recurrent state inside the chunked kernel, with a dynamic scale per token and head +# (``w`` is ``[B, T, H, K]``; ``axis: [0, 1, 2]`` reduces over K). Pairs with +# ``gdn_state_fp8_dynamic`` to emulate an FP8 x FP8 state matmul. Unlike the state quantizer this +# one runs on the tensor itself. The initial integration supports dynamic E4M3 with identity +# STE only. Needs fla-core==0.5.1 and Triton. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*gdn_w_quantizer' + cfg: + num_bits: e4m3 + axis: [0, 1, 2] + type: dynamic diff --git a/noxfile.py b/noxfile.py index 135465ed445..d061ef31f57 100644 --- a/noxfile.py +++ b/noxfile.py @@ -130,6 +130,7 @@ def gpu(session): "mamba_ssm", "causal-conv1d", ) + session.run("python", "-m", "pip", "install", "fla-core==0.5.1") session.run("python", "-m", "pytest", "tests/gpu", *_cov_args()) @@ -141,6 +142,7 @@ def gpu_megatron(session): # Pre-installed nvidia-modelopt shadows the editable install session.run("pip", "uninstall", "-y", "nvidia-modelopt") session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]") + session.run("python", "-m", "pip", "install", "fla-core==0.5.1") session.run("python", "-m", "pytest", "tests/gpu_megatron", *_cov_args()) diff --git a/pyproject.toml b/pyproject.toml index 0886e0380b3..47c8fbdf5ab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -248,6 +248,7 @@ extend-ignore = [ "tests/*" = ["B017", "D", "E402", "PT012"] "plugins/modelopt/skills/*/tests/test_*.py" = ["D", "E402"] # Skill test scripts: docstring (D) + sys.path import-order (E402) exemptions "*/_[a-zA-Z]*" = ["D"] # Private packages (_abc/*.py) or modules (_xyz.py) +"modelopt/torch/kernels/quantization/linear_attention/fla_*.py" = ["D", "E501"] # Vendored flash-linear-attention kernels "*.ipynb" = ["D", "E501"] # Ignore missing docstrings or line length for Jupyter notebooks "modelopt/torch/kernels/*" = ["N803", "N806", "E731"] # triton style "modelopt/torch/puzzletron/*" = [ @@ -322,6 +323,15 @@ disable_error_code = ["attr-defined"] module = ["examples.diffusers.fastgen.preprocess.*"] ignore_errors = true +# Vendored from fla-org/flash-linear-attention (MIT); kept faithful to upstream rather than +# annotated to modelopt's strict mypy (Triton loop indices, ``*tensor.shape`` unpacking). +[[tool.mypy.overrides]] +module = [ + "modelopt.torch.kernels.quantization.linear_attention.fla_chunk_delta_h", + "modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule", +] +ignore_errors = true + [tool.bandit] exclude_dirs = [".github/", "examples/", "noxfile.py", "tests/"] # Do not change `skips`. It should be consistent with NVIDIA's Wheel-CI-CD bandit.yml config. diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py new file mode 100644 index 00000000000..c78668febde --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Numerical forward/backward coverage for the GDN training emulation contract.""" + +import pytest +import torch +import torch.nn.functional as F +from torch.utils.checkpoint import checkpoint + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import chunk_gdn_reference +from modelopt.torch.quantization.nn import TensorQuantizer + +fla = pytest.importorskip("fla.ops.gated_delta_rule") +from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( + chunk_gated_delta_rule, +) + + +def require_state_qdq(): + if torch.cuda.get_device_capability() < (8, 9): + pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") + + +def make_inputs(dtype=torch.float32, packed=False, state_v_first=False): + torch.manual_seed(123) + # Two full chunks plus a tail; grouped-value heads and a partial state scale tile. + batch, length, heads, value_heads, keys, values = 1, 145, 1, 2, 32, 48 + q, k = [ + F.normalize(torch.randn(batch, length, heads, keys, device="cuda"), dim=-1) + for _ in range(2) + ] + v = torch.randn(batch, length, value_heads, values, device="cuda") + g = -torch.rand(batch, length, value_heads, device="cuda") * 0.1 + beta = torch.rand_like(g) + state = torch.randn(2 if packed else batch, value_heads, keys, values, device="cuda") * 0.1 + if state_v_first: + state = state.transpose(-1, -2).contiguous() + args = [x.to(dtype).requires_grad_() for x in (q, k, v, g, beta)] + return args, state.requires_grad_() + + +def w_quantizer(): + return TensorQuantizer( + QuantizerAttributeConfig(num_bits=(4, 3), axis=(0, 1, 2), type="dynamic") + ) + + +def compare(actual, expected, tolerance): + for a, e in zip(actual, expected): + assert torch.isfinite(a).all() + error = (a.float() - e.float()).norm() + bound = tolerance * e.float().norm().clamp_min(1e-6) + assert error <= bound, ( + f"relative L2 error {(error / e.float().norm()).item():.5g} > {tolerance}" + ) + + +def values_and_grads(fn, args, state, **kwargs): + result = fn(*args, initial_state=state, **kwargs) + torch.manual_seed(15) + probes = [torch.randn(x.shape, device=x.device, dtype=torch.float32) for x in result] + grads = torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probes)), (*args, state)) + return result, grads + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_disabled_matches_upstream_forward_and_backward(dtype): + args, state = make_inputs(dtype) + expected = values_and_grads(fla.chunk_gated_delta_rule, args, state, output_final_state=True) + actual = values_and_grads(chunk_gated_delta_rule, args, state, output_final_state=True) + for a, e in zip(actual, expected): + compare(a, e, 0.005 if dtype == torch.bfloat16 else 0.001) + + +@pytest.mark.parametrize( + ("state_qdq", "quantize_w"), [(False, False), (False, True), (True, False), (True, True)] +) +@pytest.mark.parametrize(("packed", "state_v_first"), [(False, False), (True, True)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_numerical_surrogate_reference(state_qdq, quantize_w, packed, state_v_first, dtype): + if state_qdq: + require_state_qdq() + args, state = make_inputs(dtype, packed=packed, state_v_first=state_v_first) + quantizer = w_quantizer() if quantize_w else None + kwargs = { + "state_qdq": state_qdq, + "state_qdq_block_v": 32, + "w_quantizer": quantizer, + "state_v_first": state_v_first, + "cu_seqlens": torch.tensor([0, 67, 145], device="cuda", dtype=torch.int32) + if packed + else None, + } + reference_args = [x.detach().float().requires_grad_() for x in args] + reference_state = state.detach().float().requires_grad_() + expected = values_and_grads(chunk_gdn_reference, reference_args, reference_state, **kwargs) + actual = values_and_grads( + chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs + ) + compare(actual[0], expected[0], 0.03 if dtype == torch.bfloat16 else 0.01) + compare(actual[1], expected[1], 0.05 if dtype == torch.bfloat16 else 0.02) + + +def test_w_qdq_saved_once_and_activation_checkpoint_parity(): + args, state = make_inputs() + quantizer = w_quantizer() + calls = [] + handle = quantizer.register_forward_hook(lambda *args: calls.append(1)) + + def fn(*x): + return chunk_gated_delta_rule( + *x[:5], initial_state=x[5], output_final_state=True, w_quantizer=quantizer + ) + + output = fn(*args, state) + grads = torch.autograd.grad(sum(x.square().sum() for x in output), (*args, state)) + assert len(calls) == 1, "Backward must not rerun the quantizer or its observers" + recomputed = checkpoint(fn, *args, state, use_reentrant=False) + checkpoint_grads = torch.autograd.grad( + sum(x.square().sum() for x in recomputed), (*args, state) + ) + compare(recomputed, output, 1e-6) + compare(checkpoint_grads, grads, 1e-6) + handle.remove() + + +def test_zero_initial_state_and_output_only_training_loss(): + args, _ = make_inputs() + quantizer = w_quantizer() + expected, _ = chunk_gdn_reference(*args, w_quantizer=quantizer) + actual, final = chunk_gated_delta_rule(*args, w_quantizer=quantizer) + assert final is None + compare((actual,), (expected,), 0.01) + grads = [torch.autograd.grad(x.square().sum(), args) for x in (actual, expected)] + compare(*grads, 0.02) + + +@pytest.mark.parametrize("block_v", [16, 64, 128]) +def test_state_scale_tile_forward_and_backward(block_v): + require_state_qdq() + args, state = make_inputs() + kwargs = {"state_qdq": True, "state_qdq_block_v": block_v} + expected = values_and_grads(chunk_gdn_reference, args, state, **kwargs) + actual = values_and_grads( + chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs + ) + compare(actual[0], expected[0], 0.01) + compare(actual[1], expected[1], 0.02) + + +def test_fused_gate_normalization_and_beta_gradients(): + args, state = make_inputs() + q, k, v, raw_g, raw_beta = args + a_log = torch.randn(2, device="cuda", requires_grad=True) + bias = torch.randn(2, device="cuda", requires_grad=True) + quantizer = w_quantizer() + expected = chunk_gdn_reference( + F.normalize(q, dim=-1), + F.normalize(k, dim=-1), + v, + -a_log.exp() * F.softplus(raw_g + bias), + 2 * raw_beta.sigmoid(), + initial_state=state, + w_quantizer=quantizer, + ) + actual = chunk_gated_delta_rule( + *args, + initial_state=state, + output_final_state=True, + w_quantizer=quantizer, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + A_log=a_log, + dt_bias=bias, + use_beta_sigmoid_in_kernel=True, + allow_neg_eigval=True, + ) + compare(actual, expected, 0.01) + grads = [ + torch.autograd.grad(sum(x.square().sum() for x in result), (*args, state, a_log, bias)) + for result in (actual, expected) + ] + compare(*grads, 0.02) + + +@pytest.mark.parametrize("chunk_size", [16, 32, 128]) +def test_unsupported_chunk_rejected_before_launch(chunk_size): + args, state = make_inputs() + with pytest.raises(ValueError, match="only chunk_size=64"): + chunk_gated_delta_rule(*args, chunk_size=chunk_size) + + +def test_unsupported_w_gradient_rejected_before_launch(): + args, _ = make_inputs() + quantizer = w_quantizer() + quantizer.set_from_attribute_config({"pass_through_bwd": False}) + with pytest.raises(ValueError, match="pass_through_bwd=True"): + chunk_gated_delta_rule(*args, w_quantizer=quantizer) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py new file mode 100644 index 00000000000..2252eb25a61 --- /dev/null +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +from _test_utils.torch.megatron.models import get_mcore_gpt_model +from _test_utils.torch.megatron.utils import get_forward, initialize_for_megatron + +import modelopt.torch.quantization as mtq + +pytest.importorskip("fla") # Megatron-Core GatedDeltaNet and the state QDQ kernel need fla +pytest.importorskip("megatron.core.ssm.gated_delta_net") + +from modelopt.torch.quantization.plugins.gated_delta_net import _state_qdq_chunk_gated_delta_rule +from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet + +try: + _state_qdq_chunk_gated_delta_rule() +except RuntimeError as e: + pytest.skip(str(e), allow_module_level=True) + +SEED = 1234 +GDN_STATE_QUANT_CFG = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*gdn_state_quantizer", + "cfg": {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"}, + }, + ], + "algorithm": "max", +} + + +def _test_gdn_state_quant_helper(rank, size): + initialize_for_megatron( + tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED + ) + model = get_mcore_gpt_model( + tensor_model_parallel_size=size, + num_layers=2, + hidden_size=64, + num_attention_heads=4, + vocab_size=32, + max_sequence_length=128, + experimental_attention_variant="gated_delta_net", + ).cuda() + model.eval() # no dropout, so the forwards differ only through the state quantizer + forward = get_forward(model) + with torch.no_grad(): + loss_ref = forward(model) + + model = mtq.quantize(model, GDN_STATE_QUANT_CFG, forward) + + gdn_modules = [m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)] + assert gdn_modules, "no GatedDeltaNet layer was wrapped" + assert all(m.gdn_state_quantizer.is_enabled for m in gdn_modules) + assert all(m.gdn_state_qdq_block_v == 64 for m in gdn_modules) # kernel's 64-column tile + + with torch.no_grad(): + loss_quant = forward(model) + assert torch.isfinite(loss_quant).all() + assert not torch.allclose(loss_quant, loss_ref, rtol=1e-4, atol=1e-4), ( + "FP8 state quantization must change the output" + ) + + mtq.disable_quantizer(model, "*gdn_state_quantizer") + with torch.no_grad(): + assert torch.allclose(forward(model), loss_ref, rtol=1e-4, atol=1e-4) + mtq.enable_quantizer(model, "*gdn_state_quantizer") + + # QAD trains through the segmented kernel: gradients must reach the GDN projections. + model.train() + optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) + before = gdn_modules[0].in_proj.weight.detach().clone() + optimizer.zero_grad() + forward(model).sum().backward() + assert all(m.in_proj.weight.grad is not None for m in gdn_modules) + assert all(torch.isfinite(m.in_proj.weight.grad).all() for m in gdn_modules) + optimizer.step() + assert not torch.equal(gdn_modules[0].in_proj.weight, before) + + +def test_gdn_state_quant(dist_workers_size_1): + """GatedDeltaNet layers get a ``gdn_state_quantizer`` whose FP8 quant-dequant of the recurrent + state runs inside the chunked kernel during forward and backward.""" + if torch.cuda.get_device_capability() < (8, 9): + pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") + dist_workers_size_1.run(_test_gdn_state_quant_helper) diff --git a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py b/tests/unit/torch/quantization/plugins/test_gated_delta_net.py new file mode 100644 index 00000000000..2bdfa2e7b09 --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_gated_delta_net.py @@ -0,0 +1,276 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from copy import deepcopy + +import pytest +import torch +import torch.nn as nn + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.config import QuantizeConfig +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.plugins import gated_delta_net +from modelopt.torch.quantization.plugins.gated_delta_net import GatedDeltaNetStateQuantMixin + +GDN_STATE_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"} + + +def chunk_gated_delta_rule(q, k, v, g, beta, **kwargs): + """Stand-in with the fla kernel's name; the real kernel needs a GPU.""" + return q + k + v, None + + +class TinyGatedDeltaNet(nn.Module): + """A module that, like Megatron-Core's GatedDeltaNet, calls ``self.gated_delta_rule``.""" + + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4) + self.gated_delta_rule = chunk_gated_delta_rule + + def forward(self, x): + out, _ = self.gated_delta_rule(x, x, x, x[..., 0], x[..., 0]) + return self.proj(out) + + +@QuantModuleRegistry.register({TinyGatedDeltaNet: "TinyGatedDeltaNet"}) +class _QuantTinyGatedDeltaNet(GatedDeltaNetStateQuantMixin): + def forward(self, x): + gated_delta_rule = self.gated_delta_rule + self.gated_delta_rule = lambda *a, **kw: self._state_quantized_chunk_gated_delta_rule( + gated_delta_rule, *a, **kw + ) + try: + return super().forward(x) + finally: + self.gated_delta_rule = gated_delta_rule + + +GDN_W_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1, 2), "type": "dynamic"} + + +def quant_cfg(state=True, w=False): + entries = [{"quantizer_name": "*", "enable": False}] + if state: + entries.append({"quantizer_name": "*gdn_state_quantizer", "cfg": GDN_STATE_FP8_DYNAMIC}) + if w: + entries.append({"quantizer_name": "*gdn_w_quantizer", "cfg": GDN_W_FP8_DYNAMIC}) + return {"quant_cfg": entries, "algorithm": "max"} + + +@pytest.mark.parametrize( + "attributes", + [ + {"num_bits": (4, 3), "axis": (0, 1)}, # static + {"num_bits": (4, 3), "type": "dynamic"}, # per tensor + {"num_bits": 8, "axis": (0, 1), "type": "dynamic"}, # int8 + {"num_bits": (4, 3), "type": "dynamic", "block_sizes": {-1: 16}}, # blockwise + ], +) +def test_validate_state_quantizer_rejects_unsupported(attributes): + with pytest.raises(ValueError, match="supports only"): + mtq.quantize( + TinyGatedDeltaNet(), + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*gdn_state_quantizer", "cfg": attributes}, + ], + "algorithm": None, + }, + ) + + +def test_disabled_state_quantizer_calls_original_kernel(): + model = TinyGatedDeltaNet() + x = torch.randn(2, 8, 3, 4) + expected = model(x) + + disable_all = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": "max"} + mtq.quantize(model, disable_all, lambda m: m(x)) + + assert isinstance(model, _QuantTinyGatedDeltaNet) + assert not model.gdn_state_quantizer.is_enabled and not model.gdn_w_quantizer.is_enabled + assert model.gdn_state_qdq_block_v == 64 + assert torch.equal(model(x), expected) + assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" + + +def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch): + calls = [] + + def fake_state_qdq_kernel(*args, **kwargs): + calls.append(kwargs) + return chunk_gated_delta_rule(*args) + + monkeypatch.setattr( + gated_delta_net, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel + ) + model = TinyGatedDeltaNet() + x = torch.randn(2, 8, 3, 4) + mtq.quantize(model, quant_cfg(), lambda m: m(x)) + + model(x) + assert calls and calls[-1] == { + "chunk_size": 64, + "state_qdq": 1, + "state_qdq_block_v": 64, + "w_quantizer": None, + } + + # The deterministic torch kernel has no quantized counterpart. + model.gated_delta_rule = lambda *a, **kw: chunk_gated_delta_rule(*a, **kw) + with pytest.raises(NotImplementedError, match="deterministic torch kernel"): + model(x) + + +@pytest.mark.parametrize("state", [False, True]) +def test_w_quantizer_is_passed_to_the_kernel(monkeypatch, state): + """``*gdn_w_quantizer`` in the config hands the module's TensorQuantizer to the kernel, with + or without the state quantizer.""" + calls = [] + + def fake_state_qdq_kernel(*args, **kwargs): + calls.append(kwargs) + return chunk_gated_delta_rule(*args) + + monkeypatch.setattr( + gated_delta_net, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel + ) + model = TinyGatedDeltaNet() + x = torch.randn(2, 8, 3, 4) + mtq.quantize(model, quant_cfg(state=state, w=True), lambda m: m(x)) + assert model.gdn_w_quantizer.is_enabled and model.gdn_state_quantizer.is_enabled == state + + model(x) + assert calls[-1]["state_qdq"] == int(state) + assert calls[-1]["w_quantizer"] is model.gdn_w_quantizer + + # The w quantizer really quantizes: 256 random values per token collapse onto the E4M3 grid, + # which has at most 127 distinct magnitudes per (row-specific) scale. + w = torch.randn(1, 1, 1, 256) + quantized = model.gdn_w_quantizer(w) + assert not torch.equal(quantized, w) + assert torch.unique(quantized.abs()).numel() <= 127 < torch.unique(w.abs()).numel() + + +@pytest.mark.parametrize("site", ["state", "w"]) +@pytest.mark.parametrize( + "overrides", + [{"pass_through_bwd": False}, {"type": "static"}, {"fake_quant": False}, {"rotate": True}], +) +def test_unsupported_quantizer_fails_during_conversion(site, overrides): + cfg = quant_cfg(state=site == "state", w=site == "w") + cfg = deepcopy(cfg) + cfg["quant_cfg"][-1]["cfg"].update(overrides) + with pytest.raises(ValueError, match="supports only"): + mtq.quantize(TinyGatedDeltaNet(), cfg) + + +def test_policy_roundtrip_last_match_and_hybrid_selection(tmp_path): + model = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) + cfg = quant_cfg() + cfg["algorithm"] = None + cfg["linear_attention"] = [ + {"module_name": "*", "cfg": {"state": {"block_v": 128}}}, + {"module_name": "0", "cfg": {"state": {"block_v": 32}}}, + ] + mtq.quantize(model, cfg) + assert model[0].gdn_state_qdq_block_v == 32 + # A post-conversion policy edit must also survive save/load. + model[0].linear_attention_config.state.block_v = 16 + path = tmp_path / "gdn.pth" + mto.save(model, path) + restored = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) + mto.restore(restored, path) + assert restored[0].linear_attention_config == model[0].linear_attention_config + assert restored[0].gdn_state_quantizer.is_enabled + assert restored[0].gdn_state_qdq_block_v == 16 + assert not hasattr(restored[1], "linear_attention_config") + + +def test_policy_rejects_unmatched_and_unimplemented_modes(): + with pytest.raises(ValueError, match="matches no supported"): + mtq.quantize( + nn.Linear(4, 4), {"linear_attention": [{"module_name": "*"}], "algorithm": None} + ) + for policy in ( + {"chunk_size": 32}, + {"solve": {"method": "neumann"}}, + {"state": {"mode": "token"}}, + ): + with pytest.raises(ValueError): + QuantizeConfig(linear_attention=[{"module_name": "*", "cfg": policy}]) + + +def test_restore_old_config_without_policy_field(): + model = mtq.quantize( + nn.Linear(4, 4), + {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None}, + ) + state = mto.modelopt_state(model) + for _, mode_state in state["modelopt_state_dict"]: + config = mode_state["config"] + if isinstance(config, dict): + config.pop("linear_attention", None) + else: + config.__dict__.pop("linear_attention", None) + restored = nn.Linear(4, 4) + mto.restore_from_modelopt_state(restored, state) + restored.load_state_dict(model.state_dict()) + x = torch.randn(2, 4) + torch.testing.assert_close(restored(x), model(x)) + + +def test_policy_refinement_updates_existing_quantized_module(): + cfg = quant_cfg() + cfg["algorithm"] = None + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + cfg["linear_attention"] = [{"module_name": "", "cfg": {"state": {"block_v": 32}}}] + mtq.quantize(model, cfg) + assert model.gdn_state_qdq_block_v == 32 + # A later complete rule resets omitted fields instead of merging the previous policy. + cfg["linear_attention"].append({"module_name": "", "cfg": {}}) + mtq.quantize(model, cfg) + assert model.gdn_state_qdq_block_v == 64 + + +def test_restore_legacy_gdn_without_new_quantizer_handles(): + config = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None} + model = mtq.quantize(TinyGatedDeltaNet(), config) + state = mto.modelopt_state(model) + for _, mode_state in state["modelopt_state_dict"]: + mode_state["config"].pop("linear_attention", None) + metadata = mode_state["metadata"] + metadata.pop("linear_attention", None) + for name in ("gdn_state_quantizer", "gdn_w_quantizer"): + metadata["quantizer_state"].pop(name, None) + restored = TinyGatedDeltaNet() + mto.restore_from_modelopt_state(restored, state) + restored.load_state_dict(model.state_dict()) + x = torch.randn(2, 8, 3, 4) + torch.testing.assert_close(restored(x), model(x)) + assert not restored.gdn_state_quantizer.is_enabled + assert not restored.gdn_w_quantizer.is_enabled + + +def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): + model = mtq.quantize( + TinyGatedDeltaNet(), mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(2, 8, 3, 4)) + ) + assert not model.gdn_state_quantizer.is_enabled + assert not model.gdn_w_quantizer.is_enabled diff --git a/tests/unit/torch/quantization/test_linear_attention_reference.py b/tests/unit/torch/quantization/test_linear_attention_reference.py new file mode 100644 index 00000000000..dd2aea8d592 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_reference.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +import torch.nn.functional as F + +from modelopt.torch.quantization.linear_attention import ( + chunk_gdn_reference, + recurrent_delta_rule_reference, + state_fp8_qdq_reference, +) + + +def inputs(length=7, kda=False): + torch.manual_seed(81) + q, k = [ + F.normalize(torch.randn(1, length, 1, 3, dtype=torch.float64), dim=-1) for _ in range(2) + ] + v = torch.randn(1, length, 2, 4, dtype=torch.float64) + g = -torch.rand(1, length, 2, *([3] if kda else []), dtype=torch.float64) + beta = torch.rand(1, length, 2, dtype=torch.float64) + return [x.requires_grad_() for x in (q, k, v, g, beta)] + + +@pytest.mark.parametrize("packed", [False, True]) +@pytest.mark.parametrize("state_v_first", [False, True]) +def test_chunk_recurrence_output_state_and_gradients(packed, state_v_first): + args = inputs() + state = torch.randn(2 if packed else 1, 2, 3, 4, dtype=torch.float64) + if state_v_first: + state = state.transpose(-1, -2).contiguous() + state.requires_grad_() + kwargs = { + "initial_state": state, + "state_v_first": state_v_first, + "cu_seqlens": torch.tensor([0, 2, 7]) if packed else None, + } + actual = chunk_gdn_reference(*args, chunk_size=3, **kwargs) + expected = recurrent_delta_rule_reference(*args, **kwargs) + for a, e in zip(actual, expected): + torch.testing.assert_close(a, e, rtol=1e-10, atol=1e-10) + probe = [torch.randn_like(x) for x in actual] + grads = [ + torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probe)), (*args, state)) + for result in (actual, expected) + ] + for a, e in zip(*grads): + torch.testing.assert_close(a, e, rtol=1e-9, atol=1e-9) + + +@pytest.mark.parametrize("kda", [False, True]) +def test_recurrence_smooth_gradcheck(kda): + args = inputs(length=2, kda=kda) + state = torch.randn(1, 2, 3, 4, dtype=torch.float64, requires_grad=True) + assert torch.autograd.gradcheck( + lambda *x: recurrent_delta_rule_reference(*x[:-1], initial_state=x[-1]), + (*args, state), + fast_mode=True, + ) + + +def test_kda_one_step_and_scalar_decay_equivalence(): + q, k, v, g, beta = inputs(length=1, kda=True) + state = torch.randn(1, 2, 3, 4, dtype=torch.float64) + out, final = recurrent_delta_rule_reference(q, k, v, g, beta, initial_state=state) + keys = k[:, 0].expand(-1, 2, -1) + decayed = state * g[:, 0].exp().unsqueeze(-1) + residual = v[:, 0] - torch.einsum("bhk,bhkv->bhv", keys, decayed) + expected = decayed + torch.einsum("bhk,bhv->bhkv", keys, beta[:, 0, :, None] * residual) + torch.testing.assert_close(final, expected) + torch.testing.assert_close( + out[:, 0], torch.einsum("bhk,bhkv->bhv", q[:, 0].expand(-1, 2, -1), expected) / 3**0.5 + ) + scalar = g[..., 0] + gdn = recurrent_delta_rule_reference(q, k, v, scalar, beta, initial_state=state) + kda = recurrent_delta_rule_reference( + q, k, v, scalar.unsqueeze(-1).expand_as(g), beta, initial_state=state + ) + for a, e in zip(kda, gdn): + torch.testing.assert_close(a, e) + + +def test_state_qdq_granularity_zero_tail_and_identity_ste(): + torch.manual_seed(14) + state = torch.randn(2, 2, 3, 37, requires_grad=True) + with torch.no_grad(): + state[0].zero_() + state[1, ..., :16] *= 10 + quantized = state_fp8_qdq_reference(state, block_v=16) + assert torch.isfinite(quantized).all() + assert torch.equal(quantized[0], state[0]) + assert not torch.equal(quantized, state) + assert not torch.equal(quantized, state_fp8_qdq_reference(state, block_v=64)) + probe = torch.randn_like(state) + (grad,) = torch.autograd.grad((quantized * probe).sum(), state) + torch.testing.assert_close(grad, probe, rtol=0, atol=0) + + +def test_packed_sequences_reset_the_state(): + args = inputs(length=7, kda=True) + initial = torch.randn(2, 2, 3, 4, dtype=torch.float64) + packed = recurrent_delta_rule_reference( + *args, initial_state=initial, cu_seqlens=torch.tensor([0, 2, 7]) + ) + separate = [ + recurrent_delta_rule_reference( + *(x[:, lo:hi] for x in args), initial_state=initial[n : n + 1] + ) + for n, (lo, hi) in enumerate(((0, 2), (2, 7))) + ] + torch.testing.assert_close(packed[0], torch.cat([x[0] for x in separate], dim=1)) + torch.testing.assert_close(packed[1], torch.cat([x[1] for x in separate])) From bfbb13ed98292fb37e540788bc72774c5bb030b0 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 21 Sep 2026 22:54:29 -0700 Subject: [PATCH 2/6] Qualify GDN QAT on Hopper and Megatron training Signed-off-by: Kai Xu --- .github/workflows/gpu_tests.yml | 3 +- CHANGELOG.rst | 2 +- docs/linear_attention_qat.md | 130 -------------- .../fla_chunk_gated_delta_rule.py | 28 ++- modelopt/torch/quantization/config.py | 7 - modelopt/torch/quantization/conversion.py | 51 ++---- .../quantization/linear_attention/__init__.py | 5 +- .../quantization/linear_attention/config.py | 60 ------- .../{validation.py => utils.py} | 10 +- modelopt/torch/quantization/model_quant.py | 4 +- .../quantization/plugins/gated_delta_net.py | 53 +++--- .../torch/quantization/plugins/megatron.py | 28 ++- modelopt_recipes/configs/ptq/units/README.md | 2 - .../ptq/units/gdn_state_fp8_dynamic.yaml | 5 +- noxfile.py | 20 ++- .../linear_attention_reference.py | 0 .../test_fla_chunk_gated_delta_rule.py | 131 +++++++++++--- .../plugins/test_megatron_gated_delta_net.py | 165 ++++++++++++------ .../plugins/test_gated_delta_net.py | 122 ++++++------- .../test_linear_attention_reference.py | 3 +- 20 files changed, 404 insertions(+), 425 deletions(-) delete mode 100644 docs/linear_attention_qat.md delete mode 100644 modelopt/torch/quantization/linear_attention/config.py rename modelopt/torch/quantization/linear_attention/{validation.py => utils.py} (83%) rename modelopt/torch/quantization/linear_attention/reference.py => tests/_test_utils/torch/quantization/linear_attention_reference.py (100%) diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index c38e462912a..f0c5a38bbab 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.github/workflows/gpu_tests.yml @@ -44,7 +44,8 @@ jobs: matrix: include: - example: gpu - timeout: 60 + # Includes dependency builds and cold compilation of the FLA forward/backward tests. + timeout: 75 # Pinned to 26.05: benchmark.py uses trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH, # which newer TensorRT (26.06) removed. Bump once the source is updated for TensorRT 10. container_image: nvcr.io/nvidia/pytorch:26.05-py3 diff --git a/CHANGELOG.rst b/CHANGELOG.rst index bc6a6ae6e7a..347806df32b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -17,7 +17,7 @@ Changelog - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration. - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md `_ for details. - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``. -- Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training, with saved execution policies and differentiable GDN/KDA references. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer. +- Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training through the standard ``quant_cfg`` interface. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer. *Megatron Framework (M-LM / M-Bridge)* diff --git a/docs/linear_attention_qat.md b/docs/linear_attention_qat.md deleted file mode 100644 index 0db84b58f22..00000000000 --- a/docs/linear_attention_qat.md +++ /dev/null @@ -1,130 +0,0 @@ -# Linear-attention numerical emulation for training - -This experimental integration adds dynamic FP8 E4M3 fake quantization to Megatron -GatedDeltaNet's chunked training path. It rounds the recurrent state at chunk boundaries -and the materialized WY activation `W` before the state-read matmul. Both sites use an -identity straight-through estimator (STE). The underlying tensors remain floating point; -this feature does not provide compressed inference state or native FP8 matmul acceleration. - -The implementation adapts [PR #2455](https://github.com/NVIDIA/Model-Optimizer/pull/2455) -at `13c7e2456f2e9d079c9ef822742eeaa634353802` onto ModelOpt -`051d6adb204f10cd3e78d0f824f31a5a01d54831`. It is the M0/M1 portion of the linear-attention -design. Additional prefill operand sites, KDA training integration, approximate inverse, -token-state quantization, decay approximation, and SSM replay remain later milestones. - -## Requirements and scope - -- Install `fla-core==0.5.1` alongside the training framework. FLA and Triton remain optional - for ordinary ModelOpt imports and the PyTorch references. -- The fused GDN implementation accepts **chunk size 64 only**, including backward. -- State QDQ requires NVIDIA SM89 or newer for the native E4M3 conversion instruction. - W-only QDQ can run on the A6000/SM86 environment used for local validation. -- The numerical recipe supports dynamic E4M3 and `pass_through_bwd=True`. Static scales, - clipping-aware backward, real quantization, rotations, and custom format backends are rejected. -- Context parallelism with either numerical site enabled is rejected. Distributed/sharded - checkpoint qualification remains a release gate; ordinary ModelOpt save/restore is covered. -- Megatron's dynamic-batching `ssm_prefill` and `ssm_decode` paths are outside this integration. - -## Configure training - -Apply the recipe after constructing the Megatron model. The module pattern must match -at least one converted GatedDeltaNet layer; dense attention layers in a hybrid model are -ignored by the execution policy. - -```python -import modelopt.torch.quantization as mtq - -config = { - "quant_cfg": [ - {"quantizer_name": "*", "enable": False}, - { - "quantizer_name": "*gdn_state_quantizer", - "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1)}, - }, - { - "quantizer_name": "*gdn_w_quantizer", - "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)}, - }, - ], - "algorithm": None, - "linear_attention": [ - { - "module_name": "decoder.layers.*.self_attention", - "cfg": { - "backend": "fla", - "chunk_size": 64, - "state": {"mode": "chunk", "block_v": 64, "quantize_initial": True}, - "solve": {"method": "exact"}, - }, - }, - ], -} -model = mtq.quantize(model, config) -``` - -Dynamic scaling requires no calibration loop. Continue with the framework's normal -forward, backward, and optimizer steps. Enable projection quantization through separate -recipe entries if needed. The standard projection presets keep the GDN numerical handles -disabled; append the GDN units after standard exclusions when composing recipes. - -Execution policies use last-match precedence: each match replaces the complete policy, -including defaults for omitted fields. There is no nested merge. The resolved policy is -reported at conversion and is available on each converted module as -`linear_attention_config`. Quantizer enables/formats and resolved policies are saved -through the ModelOpt checkpoint APIs, including changes made after conversion. - -Without an explicit execution policy, enabled GDN recipe units use the defaults above. -When both GDN handles are disabled, the adapter calls the original training kernel. - -## Numerical contract - -`W` is an activation, **not a learned weight**. The W quantizer sees `[B,T,Hv,Dk]` -and reduces over `Dk`, using one scale per token/head. Forward's quantized W is saved -for backward, avoiding a second quantizer invocation and ensuring identical values. -This costs one saved W activation when that site is enabled. - -The state handle specifies a fused operation rather than calling `TensorQuantizer` on a -materialized state. Each sequence/head uses one dynamic scale per `[Dk, block_v]` tile; -`block_v` can be 16, 32, 64, or 128. Partial value tiles are valid. The scale is -`amax / 448`, with scale one for an all-zero tile. Rounding occurs on the initial state -and after each chunk's state update. The chunk's outputs use its incoming rounded state; -they are computed before rounding the outgoing state. Initial/final state gradients -propagate through these rounding events using identity STE. - -The triangular solve remains exact. No approximate-inverse algorithm is selected or -silently enabled by this integration. - -## References and verification - -The CPU-only reference package is `modelopt.torch.quantization.linear_attention`: - -| Function | Contract | -| --- | --- | -| `recurrent_delta_rule_reference` | Exact GDN scalar-decay or KDA per-key-decay recurrence | -| `chunk_gdn_reference` | Exact GDN triangular solve and chunk algebra; optional state/W QDQ | -| `state_fp8_qdq_reference` | Dynamic E4M3 state tiles with identity STE | - -Use float64 for algebra/gradcheck and float32 for GPU comparison. References support -grouped value heads, packed sequences, nonzero initial states, tails, and either state -layout. They favor explicit arithmetic over speed and are not an efficient training backend. - -```bash -PYTHONPATH=. pytest -q tests/unit/torch/quantization/test_linear_attention_reference.py \ - tests/unit/torch/quantization/plugins/test_gated_delta_net.py \ - tests/unit/torch/quantization/test_quantize_cpu.py -PYTHONPATH=. pytest -q \ - tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py -PYTHONPATH=. pytest -q \ - tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py -``` - -The GPU suite compares outputs, final states, and gradients for every differentiable -input against the declared surrogate. It also covers packed boundaries, state layout, -fused gate/beta activation, activation checkpointing, and unsupported-policy errors. -The Megatron test includes an optimizer step. The `gpu` and `gpu_megatron` Nox sessions -install the pinned FLA dependency so missing FLA does not silently remove CI coverage. - -Local validation uses Python 3.12.8, Torch 2.9.1, Triton 3.5.1, FLA 0.5.1, and A6000/SM86. -State-QDQ GPU cases require a separate SM89+ run; the local environment also lacks -Megatron-Core. Those gates must pass before M1 is considered fully qualified. No model -quality, QAT recovery, or throughput result is claimed. diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py index aa88b49d71f..a5609cab102 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py @@ -46,9 +46,15 @@ from fla.ops.utils import chunk_local_cumsum from fla.ops.utils.constant import RCP_LN2 from fla.ops.utils.index import prepare_chunk_indices -from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard +from fla.utils import ( + IS_NVIDIA_HOPPER, + TRITON_ABOVE_3_4_0, + autocast_custom_bwd, + autocast_custom_fwd, + input_guard, +) -from modelopt.torch.quantization.linear_attention.validation import validate_gdn_quantizer +from modelopt.torch.quantization.linear_attention.utils import validate_gdn_quantizer from modelopt.torch.quantization.nn import TensorQuantizer from .fla_chunk_delta_h import ( @@ -636,10 +642,15 @@ def chunk_gated_delta_rule( # w_quantizer: dynamic FP8 TensorQuantizer applied to the WY tensor # ``w`` of shape [B, T, HV, K] before it multiplies the state, emulating an FP8 x FP8 matmul. w_quantizer = kwargs.pop("w_quantizer", None) + use_gate_in_kernel = kwargs.pop("use_gate_in_kernel", False) + A_log = kwargs.pop("A_log", None) + dt_bias = kwargs.pop("dt_bias", None) + if kwargs: + raise TypeError(f"Unexpected keyword arguments: {', '.join(sorted(kwargs))}") if state_qdq not in (STATE_QDQ_OFF, STATE_QDQ_FP8_DYNAMIC): raise ValueError(f"`state_qdq` must be 0 or 1, got {state_qdq}.") if w_quantizer is not None: - validate_gdn_quantizer(w_quantizer, state=False) + validate_gdn_quantizer(w_quantizer, name="gdn_w_quantizer") if state_qdq and (not q.is_cuda or torch.cuda.get_device_capability(q.device) < (8, 9)): raise RuntimeError("GDN state QDQ requires native E4M3 conversion on CUDA SM89 or newer.") if (state_qdq != STATE_QDQ_OFF or w_quantizer is not None) and cp_context is not None: @@ -664,9 +675,6 @@ def chunk_gated_delta_rule( f"The number of initial states is expected to be equal to the number of input sequences, " f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", ) - use_gate_in_kernel = kwargs.get("use_gate_in_kernel", False) - A_log = kwargs.get("A_log") - dt_bias = kwargs.get("dt_bias") if use_gate_in_kernel: assert A_log is not None, "A_log must be provided when use_gate_in_kernel=True." if allow_neg_eigval and not use_beta_sigmoid_in_kernel: @@ -674,6 +682,14 @@ def chunk_gated_delta_rule( if scale is None: scale = k.shape[-1] ** -0.5 + # [ModelOpt] Hopper's TileLang backward needs BF16 and equal head counts. Expand outside + # custom autograd so repeat_interleave reduces q/k gradients back to the original heads. + if IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0: + if any(x.dtype != torch.bfloat16 for x in (q, k, v)): + raise ValueError("Hopper with Triton >= 3.4 requires BF16 q/k/v for GDN training.") + if H != HV: + q = q.repeat_interleave(HV // H, dim=2) + k = k.repeat_interleave(HV // H, dim=2) o, final_state = ChunkGatedDeltaRuleFunction.apply( q, k, diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index e8db5c8446a..f8847aa1fc1 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -161,8 +161,6 @@ from modelopt.torch.opt.config_loader import load_config from modelopt.torch.utils.network import ConstructorLike -from .linear_attention.config import LinearAttentionPolicyEntry - class QuantizerCfgEntry(ModeloptBaseConfig): """A single entry in a ``quant_cfg`` list.""" @@ -1569,11 +1567,6 @@ def _dict_to_entry(key: str, value) -> list[dict[str, Any]]: class QuantizeConfig(ModeloptBaseConfig): """Default configuration for ``quantize`` mode.""" - linear_attention: list[LinearAttentionPolicyEntry] = ModeloptField( - default=[], - title="Execution policies for supported linear-attention training modules", - ) - quant_cfg: QuantizeQuantCfgType = ModeloptField( default=[{"quantizer_name": "*", "cfg": {"num_bits": 8, "axis": None}}], title="Quantization configuration", diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 0edccdf4649..78a6e03a16b 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -28,7 +28,6 @@ from modelopt.torch.opt.dynamic import _DMRegistryCls from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict from modelopt.torch.utils import get_unwrapped_name -from modelopt.torch.utils.logging import print_rank_0 from .config import ( QuantizeConfig, @@ -37,7 +36,6 @@ _QuantizeExportConfig, normalize_quant_cfg_list, ) -from .linear_attention.config import LinearAttentionConfig from .nn import ( QuantModule, QuantModuleRegistry, @@ -69,7 +67,7 @@ def convert_to_quantized_model(model: ModelLikeModule, config: QuantizeConfig) - replace_quant_module(model, version=ModeloptStateManager(model).state_version) set_quantizer_by_cfg(model, config.get("quant_cfg", [])) - _apply_linear_attention_policy(model, config) + _validate_linear_attention_quantizers(model) metadata = {} update_quantize_metadata(model, config, metadata) @@ -139,7 +137,6 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: details regarding how MCore sharded checkpoint is restored, see modelopt.torch.opt.plugins.mcore_dist_checkpointing.restore_sharded_modelopt_state. """ - _apply_linear_attention_policy(model, config, metadata.get("linear_attention")) if "quantizer_state" not in metadata: # MCore sharded checkpoint (`torch-dist`) has its quantizer_state stored as the # extra_state of `QuantModule`. The quantizer_state is resumed with @@ -147,14 +144,13 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: return model quantizer_state_dict = dict(metadata["quantizer_state"]) - if "linear_attention" not in metadata: - # Older checkpoints predate these disabled handles; preserve their baseline path. - for name, module in _linear_attention_modules(model).items(): - for handle in ("gdn_state_quantizer", "gdn_w_quantizer"): - key = f"{name}.{handle}" if name else handle - quantizer = getattr(module, handle) - if key not in quantizer_state_dict and not quantizer.is_enabled: - quantizer_state_dict[key] = quantizer.get_modelopt_state() + # Older checkpoints predate these disabled handles; preserve their baseline path. + for name, module in _linear_attention_modules(model).items(): + for handle in ("gdn_state_quantizer", "gdn_w_quantizer"): + key = f"{name}.{handle}" if name else handle + quantizer = getattr(module, handle) + if key not in quantizer_state_dict and not quantizer.is_enabled: + quantizer_state_dict[key] = quantizer.get_modelopt_state() unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys() extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys() @@ -207,12 +203,6 @@ def update_quantize_metadata( model: nn.Module, config: QuantizeConfig, metadata: MetadataDict ) -> None: """Update the quantizer state in the metadata dict.""" - policies = { - name: module.linear_attention_config.model_dump() - for name, module in _linear_attention_modules(model).items() - } - if policies: - metadata["linear_attention"] = policies metadata["quantizer_state"] = quantizer_state(model) if shared_state_metadata := SharedWeightGlobalAmaxState.metadata(model): metadata["shared_quant_states"] = shared_state_metadata @@ -231,30 +221,9 @@ def _linear_attention_modules(model): } -def _apply_linear_attention_policy(model, config, saved_policies=None): - modules = _linear_attention_modules(model) - # getattr also handles pickled pre-extension config objects with no field in __dict__. - for entry in getattr(config, "linear_attention", []): - matches = [name for name in modules if fnmatch.fnmatch(name, entry.module_name)] - if not matches: - raise ValueError( - f"linear_attention rule {entry.module_name!r} matches no supported modules" - ) - for name in matches: - modules[name].linear_attention_config = entry.cfg.model_copy(deep=True) - if saved_policies is not None: - if saved_policies.keys() != modules.keys(): - raise ApplyModeError( - "Saved linear_attention policies do not match the restored modules" - ) - for name, policy in saved_policies.items(): - modules[name].linear_attention_config = LinearAttentionConfig(**policy) - for name, module in modules.items(): +def _validate_linear_attention_quantizers(model): + for module in _linear_attention_modules(model).values(): module.validate_linear_attention() - if getattr(config, "linear_attention", []): - print_rank_0( - f"Linear attention {name!r}: {module.linear_attention_config.model_dump()}" - ) def quantizer_state(model: nn.Module) -> dict[str, Any]: diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py index 49f5ea3514a..477eeb89847 100644 --- a/modelopt/torch/quantization/linear_attention/__init__.py +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -13,7 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Numerical policies and differentiable references for linear attention.""" - -from .config import * -from .reference import * +"""Helpers for linear-attention quantization.""" diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py deleted file mode 100644 index 6405fbbf329..00000000000 --- a/modelopt/torch/quantization/linear_attention/config.py +++ /dev/null @@ -1,60 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Saved execution policy for GDN training-time numerical emulation.""" - -from typing import Literal - -from pydantic import Field - -from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField - -__all__ = ["LinearAttentionConfig", "LinearAttentionPolicyEntry"] - - -class _StateConfig(ModeloptBaseConfig): - mode: Literal["chunk"] = ModeloptField(default="chunk") - block_v: Literal[16, 32, 64, 128] = ModeloptField(default=64) - quantize_initial: Literal[True] = ModeloptField(default=True) - - -class _SolveConfig(ModeloptBaseConfig): - method: Literal["exact"] = ModeloptField(default="exact") - - -class LinearAttentionConfig(ModeloptBaseConfig): - """GDN chunk-64 policy; unsupported numerical modes fail config validation. - - ``state.block_v`` defines one dynamic scale per ``[Dk, block_v]`` tile of each - sequence/head. The initial state and every chunk's final state are rounded when - ``gdn_state_quantizer`` is enabled. Outputs use the incoming rounded state. - """ - - schema_version: Literal[1] = ModeloptField(default=1) - backend: Literal["fla"] = ModeloptField(default="fla") - chunk_size: Literal[64] = ModeloptField(default=64) - state: _StateConfig = ModeloptField(default=_StateConfig()) - solve: _SolveConfig = ModeloptField(default=_SolveConfig()) - - -class LinearAttentionPolicyEntry(ModeloptBaseConfig): - """Assign a complete policy to supported modules matching ``module_name``. - - Rules apply in order: the last match wins, without merging nested fields. - A rule must match at least one supported linear-attention module. - """ - - module_name: str = Field(...) - cfg: LinearAttentionConfig = ModeloptField(default=LinearAttentionConfig()) diff --git a/modelopt/torch/quantization/linear_attention/validation.py b/modelopt/torch/quantization/linear_attention/utils.py similarity index 83% rename from modelopt/torch/quantization/linear_attention/validation.py rename to modelopt/torch/quantization/linear_attention/utils.py index 2bf07525e22..4f1e18a174d 100644 --- a/modelopt/torch/quantization/linear_attention/validation.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -13,22 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Capabilities of the initial fused GDN fake-quant path.""" +"""Shared helpers for linear-attention quantization.""" from ..nn import TensorQuantizer __all__ = [] -def validate_gdn_quantizer(quantizer: TensorQuantizer, *, state: bool) -> None: - name = "gdn_state_quantizer" if state else "gdn_w_quantizer" - axis = (0, 1) if state else (0, 1, 2) +def validate_gdn_quantizer(quantizer: TensorQuantizer, *, name: str) -> None: + """Check the dynamic E4M3 contract and the custom backward's identity STE.""" if not isinstance(quantizer, TensorQuantizer): raise ValueError(f"{name} requires a single TensorQuantizer") if not ( quantizer._dynamic and quantizer.num_bits == (4, 3) - and quantizer.axis == axis and quantizer.block_sizes is None and quantizer.fake_quant and quantizer._pass_through_bwd @@ -39,7 +37,7 @@ def validate_gdn_quantizer(quantizer: TensorQuantizer, *, state: bool) -> None: and not quantizer._use_constant_amax ): raise ValueError( - f"{name} supports only dynamic E4M3 fake quantization with axis={axis}, " + f"{name} supports only dynamic E4M3 fake quantization with " "pass_through_bwd=True, no block_sizes, rotation, pre-scaling, bias, constant " "amax, or custom backend. Other gradient rules and formats are not implemented." ) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 5aac82cacfe..fbf9bb299fd 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -32,7 +32,7 @@ from modelopt.torch.opt.utils import forward_with_reshard from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.conversion import ( - _apply_linear_attention_policy, + _validate_linear_attention_quantizers, preserve_quantizer_attributes_context, set_quantizer_attributes_partial, set_quantizer_by_cfg, @@ -330,7 +330,7 @@ def forward_loop(model) -> None: else: # Already quantized, so lets apply the quant_cfg from the config set_quantizer_by_cfg(model, quantize_config.quant_cfg) - _apply_linear_attention_policy(model, quantize_config) + _validate_linear_attention_quantizers(model) # Fail before calibration rather than after exporting an unquantized checkpoint. _check_weight_quantization_took_effect(model, quantize_config) return calibrate(model, config.get("algorithm"), forward_loop=forward_loop) diff --git a/modelopt/torch/quantization/plugins/gated_delta_net.py b/modelopt/torch/quantization/plugins/gated_delta_net.py index 7be046b8800..df70683cc2f 100644 --- a/modelopt/torch/quantization/plugins/gated_delta_net.py +++ b/modelopt/torch/quantization/plugins/gated_delta_net.py @@ -27,13 +27,13 @@ """ from collections.abc import Callable +from functools import partial from typing import Any import torch from ..config import QuantizerAttributeConfig -from ..linear_attention.config import LinearAttentionConfig -from ..linear_attention.validation import validate_gdn_quantizer +from ..linear_attention.utils import validate_gdn_quantizer from ..nn import QuantModule, TensorQuantizer __all__ = ["GatedDeltaNetStateQuantMixin"] @@ -41,6 +41,13 @@ GatedDeltaRuleFn = Callable[..., tuple[torch.Tensor, torch.Tensor | None]] +def _fla_chunk_gated_delta_rule() -> GatedDeltaRuleFn: + # FLA is an optional, heavy dependency needed only when quantization is enabled. + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + + return chunk_gated_delta_rule + + def _state_qdq_chunk_gated_delta_rule() -> GatedDeltaRuleFn: # Imported on first use: flash-linear-attention is a heavy optional dependency that only the # enabled quantizer needs, and importing it warns on machines without a GPU. @@ -64,27 +71,29 @@ class GatedDeltaNetStateQuantMixin(QuantModule): with ``quant_cfg`` entries on ``*gdn_state_quantizer`` / ``*gdn_w_quantizer`` such as the ``configs/ptq/units/gdn_state_fp8_dynamic`` and ``gdn_w_fp8_dynamic`` recipe units. The state quantizer carries the fused QDQ configuration. Both sites currently require dynamic - E4M3 and identity STE. The execution policy is saved in ModelOpt metadata. + E4M3 and identity STE. State QDQ uses fixed 64-column value tiles and 64-token chunks. """ def _setup(self): - self.gdn_state_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) - self.gdn_w_quantizer = TensorQuantizer(QuantizerAttributeConfig(enable=False)) - self.linear_attention_config = LinearAttentionConfig() - - @property - def gdn_state_qdq_block_v(self) -> int: - """Value-column scale grouping from the saved execution policy.""" - return self.linear_attention_config.state.block_v + for name in ("gdn_state_quantizer", "gdn_w_quantizer"): + self._register_temp_attribute( + name, TensorQuantizer(QuantizerAttributeConfig(enable=False)) + ) def validate_linear_attention(self) -> None: """Reject numerical settings that the fused training path cannot implement.""" - for state, quantizer in ((True, self.gdn_state_quantizer), (False, self.gdn_w_quantizer)): + for name in ("gdn_state_quantizer", "gdn_w_quantizer"): + quantizer = getattr(self, name) if quantizer.is_enabled: - validate_gdn_quantizer(quantizer, state=state) + validate_gdn_quantizer(quantizer, name=name) + # The state handle configures fused QDQ; W's grouping is executed by TensorQuantizer. + if self.gdn_state_quantizer.is_enabled and self.gdn_state_quantizer.axis != (0, 1): + raise ValueError( + "gdn_state_quantizer supports only axis=(0, 1) with 64-column value tiling" + ) def modelopt_post_restore(self, prefix: str = ""): - """Validate restored quantizers before using the saved execution policy.""" + """Validate restored quantizers against the fused training kernel requirements.""" super().modelopt_post_restore(prefix) self.validate_linear_attention() @@ -97,19 +106,23 @@ def _state_quantized_chunk_gated_delta_rule( quantize_w = self.gdn_w_quantizer.is_enabled if not (quantize_state or quantize_w): return gated_delta_rule(*args, **kwargs) - if getattr(gated_delta_rule, "__name__", None) != "chunk_gated_delta_rule": + while isinstance(gated_delta_rule, partial): + args = (*gated_delta_rule.args, *args) + kwargs = {**gated_delta_rule.keywords, **kwargs} + gated_delta_rule = gated_delta_rule.func + if gated_delta_rule is not _fla_chunk_gated_delta_rule(): raise NotImplementedError( - "GatedDeltaNet quantizers require the fla chunked kernel; the deterministic torch " - f"kernel ({gated_delta_rule!r}) is not supported." + "GatedDeltaNet quantization supports only FLA's chunk_gated_delta_rule " + f"callable or a functools.partial of it; got {gated_delta_rule!r}." ) - chunk_size = kwargs.pop("chunk_size", self.linear_attention_config.chunk_size) - if chunk_size != self.linear_attention_config.chunk_size: + chunk_size = kwargs.pop("chunk_size", 64) + if chunk_size != 64: raise ValueError("GDN fake quantization supports only chunk_size=64") return _state_qdq_chunk_gated_delta_rule()( *args, chunk_size=chunk_size, state_qdq=int(quantize_state), - state_qdq_block_v=self.gdn_state_qdq_block_v, + state_qdq_block_v=64, w_quantizer=self.gdn_w_quantizer if quantize_w else None, **kwargs, ) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 3a0d4627e38..ec942e0c24c 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -1083,9 +1083,9 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): class _QuantGatedDeltaNet(GatedDeltaNetStateQuantMixin): """GatedDeltaNet with fake quantization of the recurrent state at kernel chunk boundaries. - Training and calibration reach the chunked kernel through ``forward_pre_attn_and_core_attn``, - which calls ``self.gated_delta_rule``; that call is routed through ``gdn_state_quantizer``. - The dynamic-batching inference paths (``ssm_prefill`` / ``ssm_decode``) are left untouched. + Routes ``self.gated_delta_rule`` through state/W QDQ from both Megatron's direct + forward and older split-forward layouts. The older dynamic-batching inference + paths (``ssm_prefill`` / ``ssm_decode``) are left untouched. """ def _setup(self): @@ -1099,16 +1099,34 @@ def _setup(self): mcore_parallel.get_tensor_model_parallel_group(), ) - def forward_pre_attn_and_core_attn(self, *args, **kwargs): + @contextmanager + def _quantized_gdn_kernel(self): gated_delta_rule = self.gated_delta_rule self.gated_delta_rule = partial( self._state_quantized_chunk_gated_delta_rule, gated_delta_rule ) try: - return super().forward_pre_attn_and_core_attn(*args, **kwargs) + yield finally: self.gated_delta_rule = gated_delta_rule + def forward(self, *args, **kwargs): + if hasattr(GatedDeltaNet, "forward_pre_attn_and_core_attn"): + return super().forward(*args, **kwargs) + with self._quantized_gdn_kernel(): + return super().forward(*args, **kwargs) + + def forward_pre_attn_and_core_attn(self, *args, **kwargs): + with self._quantized_gdn_kernel(): + return super().forward_pre_attn_and_core_attn(*args, **kwargs) + + def validate_linear_attention(self): + super().validate_linear_attention() + if self.config.context_parallel_size > 1 and ( + self.gdn_state_quantizer.is_enabled or self.gdn_w_quantizer.is_enabled + ): + raise NotImplementedError("GDN QAT does not support Megatron context parallelism.") + def _is_supported_megatron_model(model: torch.nn.Module) -> bool: return isinstance(model, MegatronModule) diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index 4b6ccec395f..79bc6065c24 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -34,5 +34,3 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `attention_qkv_fp8.yaml` | FP8 E4M3 on attention q/k/v bmm and softmax quantizers | | `gdn_state_fp8_dynamic.yaml` | FP8 E4M3 dynamic fake quantization of the GatedDeltaNet recurrent state (per sequence, head and 64-column tile) at every kernel chunk boundary; needs `fla-core==0.5.1`, Triton, and SM89+ | | `gdn_w_fp8_dynamic.yaml` | FP8 E4M3 dynamic (per token and head) fake quantization of the WY tensor `w` that multiplies the GatedDeltaNet state; requires identity STE, `fla-core==0.5.1`, and Triton | - -See the [linear-attention training guide](../../../../docs/linear_attention_qat.md) for GDN policy, hardware, checkpoint, and gradient requirements. diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml index 51f9ad1de20..114c86d5cad 100644 --- a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -16,9 +16,8 @@ # FP8 E4M3 fake quantization of the GatedDeltaNet (linear attention) recurrent state with a # dynamic scale computed inside the kernel from the state itself (no calibration). The state is # ``[N, H, K, V]`` for ``N`` sequences; ``axis: [0, 1]`` scales each sequence's head state on its -# own, in tiles of ``gdn_state_qdq_block_v`` value columns (default: the kernel's 64-column tile, -# two scales per head for V = 128; set 128 on the quantized module for one scale per head). This -# is the only configuration the kernel implements. Applied at every kernel chunk boundary +# own, with one scale per full-K by 64-value-column tile (two scales per head for V = 128). +# This is the only configuration the module supports. Applied at every kernel chunk boundary # (64 tokens) during calibration, QAT and QAD; needs fla-core==0.5.1 and Triton. # See ``modelopt.torch.quantization.plugins.gated_delta_net``. diff --git a/noxfile.py b/noxfile.py index d061ef31f57..cd9e96ecffa 100644 --- a/noxfile.py +++ b/noxfile.py @@ -130,7 +130,15 @@ def gpu(session): "mamba_ssm", "causal-conv1d", ) - session.run("python", "-m", "pip", "install", "fla-core==0.5.1") + session.run( + "python", + "-m", + "pip", + "install", + "fla-core==0.5.1", + "tilelang==0.1.8", + "apache-tvm-ffi==0.1.9", + ) session.run("python", "-m", "pytest", "tests/gpu", *_cov_args()) @@ -142,7 +150,15 @@ def gpu_megatron(session): # Pre-installed nvidia-modelopt shadows the editable install session.run("pip", "uninstall", "-y", "nvidia-modelopt") session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]") - session.run("python", "-m", "pip", "install", "fla-core==0.5.1") + session.run( + "python", + "-m", + "pip", + "install", + "fla-core==0.5.1", + "tilelang==0.1.8", + "apache-tvm-ffi==0.1.9", + ) session.run("python", "-m", "pytest", "tests/gpu_megatron", *_cov_args()) diff --git a/modelopt/torch/quantization/linear_attention/reference.py b/tests/_test_utils/torch/quantization/linear_attention_reference.py similarity index 100% rename from modelopt/torch/quantization/linear_attention/reference.py rename to tests/_test_utils/torch/quantization/linear_attention_reference.py diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py index c78668febde..d470da9f491 100644 --- a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py @@ -15,16 +15,24 @@ """Numerical forward/backward coverage for the GDN training emulation contract.""" +from functools import partial + import pytest import torch import torch.nn.functional as F +from _test_utils.torch.quantization.linear_attention_reference import ( + chunk_gdn_reference, + state_fp8_qdq_reference, +) from torch.utils.checkpoint import checkpoint from modelopt.torch.quantization.config import QuantizerAttributeConfig -from modelopt.torch.quantization.linear_attention import chunk_gdn_reference from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.plugins.gated_delta_net import GatedDeltaNetStateQuantMixin fla = pytest.importorskip("fla.ops.gated_delta_rule") +from fla.utils import IS_NVIDIA_HOPPER, TRITON_ABOVE_3_4_0 + from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( chunk_gated_delta_rule, ) @@ -35,10 +43,16 @@ def require_state_qdq(): pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") -def make_inputs(dtype=torch.float32, packed=False, state_v_first=False): +def make_inputs(dtype=None, packed=False, state_v_first=False, grouped=True): + hopper_backward = IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0 + if dtype is None: + dtype = torch.bfloat16 if hopper_backward else torch.float32 + if dtype == torch.float32 and hopper_backward: + pytest.skip("Hopper/TileLang training supports BF16; FP32 is rejected before launch") torch.manual_seed(123) # Two full chunks plus a tail; grouped-value heads and a partial state scale tile. batch, length, heads, value_heads, keys, values = 1, 145, 1, 2, 32, 48 + heads = heads if grouped else value_heads q, k = [ F.normalize(torch.randn(batch, length, heads, keys, device="cuda"), dim=-1) for _ in range(2) @@ -53,10 +67,8 @@ def make_inputs(dtype=torch.float32, packed=False, state_v_first=False): return args, state.requires_grad_() -def w_quantizer(): - return TensorQuantizer( - QuantizerAttributeConfig(num_bits=(4, 3), axis=(0, 1, 2), type="dynamic") - ) +def w_quantizer(axis=(0, 1, 2)): + return TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), axis=axis, type="dynamic")) def compare(actual, expected, tolerance): @@ -69,6 +81,14 @@ def compare(actual, expected, tolerance): ) +def compare_forward(actual, expected, dtype, state_qdq): + tolerance = 0.03 if dtype == torch.bfloat16 else 0.01 + compare(actual[:1], expected[:1], tolerance) + # FP8 bin crossings amplify small GEMM differences in the stored final state. + state_tolerance = max(tolerance, 0.02) if state_qdq else tolerance + compare(actual[1:], expected[1:], state_tolerance) + + def values_and_grads(fn, args, state, **kwargs): result = fn(*args, initial_state=state, **kwargs) torch.manual_seed(15) @@ -77,9 +97,12 @@ def values_and_grads(fn, args, state, **kwargs): return result, grads +# The first comparison compiles and autotunes both upstream and ModelOpt forward/backward kernels. +@pytest.mark.timeout(300) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_disabled_matches_upstream_forward_and_backward(dtype): - args, state = make_inputs(dtype) + # Upstream Hopper/TileLang backward cannot handle grouped value heads. + args, state = make_inputs(dtype, grouped=False) expected = values_and_grads(fla.chunk_gated_delta_rule, args, state, output_final_state=True) actual = values_and_grads(chunk_gated_delta_rule, args, state, output_final_state=True) for a, e in zip(actual, expected): @@ -111,7 +134,7 @@ def test_numerical_surrogate_reference(state_qdq, quantize_w, packed, state_v_fi actual = values_and_grads( chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs ) - compare(actual[0], expected[0], 0.03 if dtype == torch.bfloat16 else 0.01) + compare_forward(actual[0], expected[0], dtype, state_qdq) compare(actual[1], expected[1], 0.05 if dtype == torch.bfloat16 else 0.02) @@ -138,15 +161,16 @@ def fn(*x): handle.remove() -def test_zero_initial_state_and_output_only_training_loss(): +@pytest.mark.parametrize("axis", [None, (0, 1), (0, 1, 2)]) +def test_zero_initial_state_and_output_only_training_loss(axis): args, _ = make_inputs() - quantizer = w_quantizer() - expected, _ = chunk_gdn_reference(*args, w_quantizer=quantizer) + quantizer = w_quantizer(axis) + expected, _ = chunk_gdn_reference(*(x.float() for x in args), w_quantizer=quantizer) actual, final = chunk_gated_delta_rule(*args, w_quantizer=quantizer) assert final is None - compare((actual,), (expected,), 0.01) + compare((actual,), (expected,), 0.03 if args[0].dtype == torch.bfloat16 else 0.01) grads = [torch.autograd.grad(x.square().sum(), args) for x in (actual, expected)] - compare(*grads, 0.02) + compare(*grads, 0.05 if args[0].dtype == torch.bfloat16 else 0.02) @pytest.mark.parametrize("block_v", [16, 64, 128]) @@ -154,12 +178,37 @@ def test_state_scale_tile_forward_and_backward(block_v): require_state_qdq() args, state = make_inputs() kwargs = {"state_qdq": True, "state_qdq_block_v": block_v} - expected = values_and_grads(chunk_gdn_reference, args, state, **kwargs) + reference_args = [x.detach().float().requires_grad_() for x in args] + expected = values_and_grads(chunk_gdn_reference, reference_args, state, **kwargs) actual = values_and_grads( chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs ) - compare(actual[0], expected[0], 0.01) - compare(actual[1], expected[1], 0.02) + compare_forward(actual[0], expected[0], args[0].dtype, state_qdq=True) + compare(actual[1], expected[1], 0.05 if args[0].dtype == torch.bfloat16 else 0.02) + + +@pytest.mark.parametrize("block_v", [16, 32, 64, 128]) +def test_state_qdq_without_recurrent_updates(block_v): + require_state_qdq() + args, state = make_inputs(packed=True, state_v_first=True) + q, k, v, g, beta = args + _, final = chunk_gated_delta_rule( + q, + torch.zeros_like(k), + torch.zeros_like(v), + torch.zeros_like(g), + beta, + initial_state=state, + output_final_state=True, + state_v_first=True, + cu_seqlens=torch.tensor([0, 67, 145], device="cuda", dtype=torch.int32), + state_qdq=True, + state_qdq_block_v=block_v, + ) + expected = state_fp8_qdq_reference(state.transpose(-1, -2), block_v).transpose(-1, -2) + torch.testing.assert_close(final, expected, rtol=1e-6, atol=1e-7) + (grad,) = torch.autograd.grad(final.sum(), state) + torch.testing.assert_close(grad, torch.ones_like(state), rtol=0, atol=0) def test_fused_gate_normalization_and_beta_gradients(): @@ -169,11 +218,11 @@ def test_fused_gate_normalization_and_beta_gradients(): bias = torch.randn(2, device="cuda", requires_grad=True) quantizer = w_quantizer() expected = chunk_gdn_reference( - F.normalize(q, dim=-1), - F.normalize(k, dim=-1), - v, - -a_log.exp() * F.softplus(raw_g + bias), - 2 * raw_beta.sigmoid(), + F.normalize(q.float(), dim=-1), + F.normalize(k.float(), dim=-1), + v.float(), + -a_log.exp() * F.softplus(raw_g.float() + bias), + 2 * raw_beta.float().sigmoid(), initial_state=state, w_quantizer=quantizer, ) @@ -189,12 +238,40 @@ def test_fused_gate_normalization_and_beta_gradients(): use_beta_sigmoid_in_kernel=True, allow_neg_eigval=True, ) - compare(actual, expected, 0.01) + compare(actual, expected, 0.03 if q.dtype == torch.bfloat16 else 0.01) grads = [ torch.autograd.grad(sum(x.square().sum() for x in result), (*args, state, a_log, bias)) for result in (actual, expected) ] - compare(*grads, 0.02) + compare(*grads, 0.05 if q.dtype == torch.bfloat16 else 0.02) + + +def test_plugin_dispatch_preserves_fla_partial(): + args, state = make_inputs() + module = GatedDeltaNetStateQuantMixin.convert(torch.nn.Module()) + module.gdn_w_quantizer = w_quantizer() + expected = values_and_grads( + chunk_gated_delta_rule, + args, + state, + output_final_state=True, + w_quantizer=module.gdn_w_quantizer, + ) + + def wrapped(q, *remaining, **kwargs): + kernel = partial(fla.chunk_gated_delta_rule, q, output_final_state=True) + return module._state_quantized_chunk_gated_delta_rule(kernel, *remaining, **kwargs) + + actual = values_and_grads(wrapped, args, state) + for a, e in zip(actual, expected): + compare(a, e, 1e-6) + + +@pytest.mark.parametrize("option", ["use_gate_in_kernal", "unsupported_option"]) +def test_unknown_keyword_rejected_before_launch(option): + args, _ = make_inputs() + with pytest.raises(TypeError, match=f"Unexpected keyword arguments: {option}"): + chunk_gated_delta_rule(*args, **{option: True}) @pytest.mark.parametrize("chunk_size", [16, 32, 128]) @@ -210,3 +287,11 @@ def test_unsupported_w_gradient_rejected_before_launch(): quantizer.set_from_attribute_config({"pass_through_bwd": False}) with pytest.raises(ValueError, match="pass_through_bwd=True"): chunk_gated_delta_rule(*args, w_quantizer=quantizer) + + +def test_hopper_fp32_rejected_before_launch(): + if not (IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0): + pytest.skip("Hopper/TileLang capability guard") + args, _ = make_inputs() + with pytest.raises(ValueError, match="requires BF16 q/k/v"): + chunk_gated_delta_rule(*(x.float() for x in args)) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index 2252eb25a61..1fe3e7399b8 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -16,12 +16,16 @@ import pytest import torch from _test_utils.torch.megatron.models import get_mcore_gpt_model -from _test_utils.torch.megatron.utils import get_forward, initialize_for_megatron +from _test_utils.torch.megatron.utils import ( + get_forward, + initialize_for_megatron, + sharded_state_dict_test_helper, +) import modelopt.torch.quantization as mtq pytest.importorskip("fla") # Megatron-Core GatedDeltaNet and the state QDQ kernel need fla -pytest.importorskip("megatron.core.ssm.gated_delta_net") +GatedDeltaNet = pytest.importorskip("megatron.core.ssm.gated_delta_net").GatedDeltaNet from modelopt.torch.quantization.plugins.gated_delta_net import _state_qdq_chunk_gated_delta_rule from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet @@ -32,70 +36,133 @@ pytest.skip(str(e), allow_module_level=True) SEED = 1234 -GDN_STATE_QUANT_CFG = { - "quant_cfg": [ - {"quantizer_name": "*", "enable": False}, - { - "quantizer_name": "*gdn_state_quantizer", - "cfg": {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"}, - }, - ], - "algorithm": "max", -} - - -def _test_gdn_state_quant_helper(rank, size): + + +def _make_model(tp_size, **config_kwargs): + model = ( + get_mcore_gpt_model( + tensor_model_parallel_size=tp_size, + num_layers=2, + hidden_size=64, + num_attention_heads=4, + vocab_size=32, + max_sequence_length=128, + experimental_attention_variant="gated_delta_net", + linear_value_head_dim=64, + **config_kwargs, + ) + .cuda() + .eval() + ) + # Retain enough history for chunk-boundary state rounding to affect the next chunk. + with torch.no_grad(): + for module in model.modules(): + if isinstance(module, GatedDeltaNet): + module.A_log.fill_(-4) + module.dt_bias.zero_() + return model + + +def _gdn_config(sites): + return { + "quant_cfg": [{"quantizer_name": "*", "enable": False}] + + [ + { + "quantizer_name": f"*gdn_{site}_quantizer", + "cfg": { + "num_bits": (4, 3), + "axis": (0, 1) if site == "state" else (0, 1, 2), + "type": "dynamic", + }, + } + for site in sites + ], + "algorithm": None, + } + + +def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): initialize_for_megatron( tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED ) - model = get_mcore_gpt_model( - tensor_model_parallel_size=size, - num_layers=2, - hidden_size=64, - num_attention_heads=4, - vocab_size=32, - max_sequence_length=128, - experimental_attention_variant="gated_delta_net", - ).cuda() - model.eval() # no dropout, so the forwards differ only through the state quantizer + model = _make_model(size) forward = get_forward(model) + outputs = [] + handles = [ + module.register_forward_hook( + lambda module, args, result: outputs.append(result[0].detach().clone()) + ) + for module in model.modules() + if isinstance(module, GatedDeltaNet) + ] with torch.no_grad(): loss_ref = forward(model) + gdn_ref = outputs.copy() + outputs.clear() - model = mtq.quantize(model, GDN_STATE_QUANT_CFG, forward) - + sites = ("state", "w") if mode == "both" else (mode,) + mtq.quantize(model, _gdn_config(sites)) gdn_modules = [m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)] assert gdn_modules, "no GatedDeltaNet layer was wrapped" - assert all(m.gdn_state_quantizer.is_enabled for m in gdn_modules) - assert all(m.gdn_state_qdq_block_v == 64 for m in gdn_modules) # kernel's 64-column tile + for module in gdn_modules: + assert module.gdn_state_quantizer.is_enabled == ("state" in sites) + assert module.gdn_w_quantizer.is_enabled == ("w" in sites) with torch.no_grad(): loss_quant = forward(model) assert torch.isfinite(loss_quant).all() - assert not torch.allclose(loss_quant, loss_ref, rtol=1e-4, atol=1e-4), ( - "FP8 state quantization must change the output" + # BF16 residual additions can hide a small GDN change at the final loss. + assert len(outputs) == len(gdn_ref) == len(gdn_modules) + assert any(not torch.equal(actual, ref) for actual, ref in zip(outputs, gdn_ref)), ( + "enabled QDQ must change the GDN branch output" ) + for handle in handles: + handle.remove() - mtq.disable_quantizer(model, "*gdn_state_quantizer") + for site in sites: + mtq.disable_quantizer(model, f"*gdn_{site}_quantizer") with torch.no_grad(): - assert torch.allclose(forward(model), loss_ref, rtol=1e-4, atol=1e-4) - mtq.enable_quantizer(model, "*gdn_state_quantizer") - - # QAD trains through the segmented kernel: gradients must reach the GDN projections. - model.train() - optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) - before = gdn_modules[0].in_proj.weight.detach().clone() - optimizer.zero_grad() - forward(model).sum().backward() - assert all(m.in_proj.weight.grad is not None for m in gdn_modules) - assert all(torch.isfinite(m.in_proj.weight.grad).all() for m in gdn_modules) + torch.testing.assert_close(forward(model), loss_ref, rtol=1e-4, atol=1e-4) + for site in sites: + mtq.enable_quantizer(model, f"*gdn_{site}_quantizer") + + restored = _make_model(size) + sharded_state_dict_test_helper(checkpoint_path, model, restored, forward) + restored_gdn = [m for m in restored.modules() if isinstance(m, _QuantGatedDeltaNet)] + assert len(restored_gdn) == len(gdn_modules) + for module in restored_gdn: + assert module.gdn_state_quantizer.is_enabled == ("state" in sites) + assert module.gdn_w_quantizer.is_enabled == ("w" in sites) + assert module.in_proj.weight.grad is not None + assert torch.isfinite(module.in_proj.weight.grad).all() + + # The checkpoint helper already ran backward through the restored model. + optimizer = torch.optim.SGD(restored.parameters(), lr=1e-3) + before = restored_gdn[0].in_proj.weight.detach().clone() optimizer.step() - assert not torch.equal(gdn_modules[0].in_proj.weight, before) + assert not torch.equal(restored_gdn[0].in_proj.weight, before) -def test_gdn_state_quant(dist_workers_size_1): - """GatedDeltaNet layers get a ``gdn_state_quantizer`` whose FP8 quant-dequant of the recurrent - state runs inside the chunked kernel during forward and backward.""" - if torch.cuda.get_device_capability() < (8, 9): +# Cold FLA/TileLang compilation plus a two-rank checkpoint exceeds the lane's 120s default. +@pytest.mark.timeout(300) +@pytest.mark.parametrize("tp_size", [1, 2]) +@pytest.mark.parametrize("mode", ["state", "w", "both"]) +def test_gdn_qat_and_sharded_restore(request, tmp_path, tp_size, mode): + """Train through state/W QDQ after a Megatron distributed-checkpoint round trip.""" + if mode != "w" and torch.cuda.get_device_capability() < (8, 9): pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") - dist_workers_size_1.run(_test_gdn_state_quant_helper) + workers = request.getfixturevalue(f"dist_workers_size_{tp_size}") + workers.run(_test_gdn_qat_helper, mode, tmp_path) + + +def _test_gdn_context_parallel_helper(rank, size, mode): + initialize_for_megatron(context_parallel_size=size, seed=SEED) + model = _make_model(1, context_parallel_size=size) + with pytest.raises(NotImplementedError, match="context parallelism"): + mtq.quantize(model, _gdn_config((mode,))) + + +@pytest.mark.parametrize("mode", ["state", "w"]) +def test_gdn_context_parallel_rejected(dist_workers_size_2, mode): + """Reject unqualified Megatron CP even when it does not pass an FLA CP context.""" + dist_workers_size_2.run(_test_gdn_context_parallel_helper, mode) diff --git a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py b/tests/unit/torch/quantization/plugins/test_gated_delta_net.py index 2bdfa2e7b09..3793edcb189 100644 --- a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py +++ b/tests/unit/torch/quantization/plugins/test_gated_delta_net.py @@ -14,6 +14,7 @@ # limitations under the License. from copy import deepcopy +from functools import partial import pytest import torch @@ -21,7 +22,6 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq -from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins import gated_delta_net from modelopt.torch.quantization.plugins.gated_delta_net import GatedDeltaNetStateQuantMixin @@ -30,10 +30,17 @@ def chunk_gated_delta_rule(q, k, v, g, beta, **kwargs): - """Stand-in with the fla kernel's name; the real kernel needs a GPU.""" + """CPU stand-in for the optional FLA kernel.""" return q + k + v, None +@pytest.fixture(autouse=True) +def mock_fla_kernel(monkeypatch): + monkeypatch.setattr( + gated_delta_net, "_fla_chunk_gated_delta_rule", lambda: chunk_gated_delta_rule + ) + + class TinyGatedDeltaNet(nn.Module): """A module that, like Megatron-Core's GatedDeltaNet, calls ``self.gated_delta_rule``.""" @@ -95,6 +102,14 @@ def test_validate_state_quantizer_rejects_unsupported(attributes): ) +def test_dynamic_export_removes_linear_attention_attributes(): + model = _QuantTinyGatedDeltaNet.convert(TinyGatedDeltaNet()) + model.export() + assert type(model) is TinyGatedDeltaNet + for name in ("gdn_state_quantizer", "gdn_w_quantizer"): + assert not hasattr(model, name) + + def test_disabled_state_quantizer_calls_original_kernel(): model = TinyGatedDeltaNet() x = torch.randn(2, 8, 3, 4) @@ -105,12 +120,12 @@ def test_disabled_state_quantizer_calls_original_kernel(): assert isinstance(model, _QuantTinyGatedDeltaNet) assert not model.gdn_state_quantizer.is_enabled and not model.gdn_w_quantizer.is_enabled - assert model.gdn_state_qdq_block_v == 64 assert torch.equal(model(x), expected) assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" -def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch): +@pytest.mark.parametrize("partial_kernel", [False, True]) +def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, partial_kernel): calls = [] def fake_state_qdq_kernel(*args, **kwargs): @@ -121,6 +136,8 @@ def fake_state_qdq_kernel(*args, **kwargs): gated_delta_net, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel ) model = TinyGatedDeltaNet() + if partial_kernel: + model.gated_delta_rule = partial(chunk_gated_delta_rule, output_final_state=True) x = torch.randn(2, 8, 3, 4) mtq.quantize(model, quant_cfg(), lambda m: m(x)) @@ -130,11 +147,13 @@ def fake_state_qdq_kernel(*args, **kwargs): "state_qdq": 1, "state_qdq_block_v": 64, "w_quantizer": None, + **({"output_final_state": True} if partial_kernel else {}), } - # The deterministic torch kernel has no quantized counterpart. + # An unrelated callable must not pass validation merely by copying the FLA name. model.gated_delta_rule = lambda *a, **kw: chunk_gated_delta_rule(*a, **kw) - with pytest.raises(NotImplementedError, match="deterministic torch kernel"): + model.gated_delta_rule.__name__ = "chunk_gated_delta_rule" + with pytest.raises(NotImplementedError, match="supports only FLA"): model(x) @@ -181,72 +200,55 @@ def test_unsupported_quantizer_fails_during_conversion(site, overrides): mtq.quantize(TinyGatedDeltaNet(), cfg) -def test_policy_roundtrip_last_match_and_hybrid_selection(tmp_path): +@pytest.mark.parametrize("axis", [None, (0, 1), (0, 1, 2)]) +def test_w_grouping_is_preserved_during_conversion(axis): + cfg = deepcopy(quant_cfg(state=False, w=True)) + cfg["algorithm"] = None + cfg["quant_cfg"][-1]["cfg"]["axis"] = axis + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + assert model.gdn_w_quantizer.axis == axis + + +@pytest.mark.parametrize(("state", "w"), [(True, False), (False, True), (True, True)]) +def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): model = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) - cfg = quant_cfg() + cfg = quant_cfg(state=state, w=w) cfg["algorithm"] = None - cfg["linear_attention"] = [ - {"module_name": "*", "cfg": {"state": {"block_v": 128}}}, - {"module_name": "0", "cfg": {"state": {"block_v": 32}}}, - ] mtq.quantize(model, cfg) - assert model[0].gdn_state_qdq_block_v == 32 - # A post-conversion policy edit must also survive save/load. - model[0].linear_attention_config.state.block_v = 16 + if w: + # Save the actual quantizer settings, including edits after conversion. + model[0].gdn_w_quantizer.axis = None path = tmp_path / "gdn.pth" mto.save(model, path) restored = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) mto.restore(restored, path) - assert restored[0].linear_attention_config == model[0].linear_attention_config - assert restored[0].gdn_state_quantizer.is_enabled - assert restored[0].gdn_state_qdq_block_v == 16 - assert not hasattr(restored[1], "linear_attention_config") - + for name, enabled in (("gdn_state_quantizer", state), ("gdn_w_quantizer", w)): + original = getattr(model[0], name) + quantizer = getattr(restored[0], name) + assert quantizer.is_enabled == enabled + assert quantizer.axis == original.axis + assert quantizer.num_bits == original.num_bits + assert quantizer._dynamic == original._dynamic + assert not hasattr(restored[1], name) -def test_policy_rejects_unmatched_and_unimplemented_modes(): - with pytest.raises(ValueError, match="matches no supported"): - mtq.quantize( - nn.Linear(4, 4), {"linear_attention": [{"module_name": "*"}], "algorithm": None} - ) - for policy in ( - {"chunk_size": 32}, - {"solve": {"method": "neumann"}}, - {"state": {"mode": "token"}}, - ): - with pytest.raises(ValueError): - QuantizeConfig(linear_attention=[{"module_name": "*", "cfg": policy}]) - -def test_restore_old_config_without_policy_field(): - model = mtq.quantize( - nn.Linear(4, 4), - {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None}, - ) - state = mto.modelopt_state(model) - for _, mode_state in state["modelopt_state_dict"]: - config = mode_state["config"] - if isinstance(config, dict): - config.pop("linear_attention", None) - else: - config.__dict__.pop("linear_attention", None) - restored = nn.Linear(4, 4) - mto.restore_from_modelopt_state(restored, state) - restored.load_state_dict(model.state_dict()) - x = torch.randn(2, 4) - torch.testing.assert_close(restored(x), model(x)) - - -def test_policy_refinement_updates_existing_quantized_module(): +def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): cfg = quant_cfg() cfg["algorithm"] = None model = mtq.quantize(TinyGatedDeltaNet(), cfg) - cfg["linear_attention"] = [{"module_name": "", "cfg": {"state": {"block_v": 32}}}] - mtq.quantize(model, cfg) - assert model.gdn_state_qdq_block_v == 32 - # A later complete rule resets omitted fields instead of merging the previous policy. - cfg["linear_attention"].append({"module_name": "", "cfg": {}}) + assert model.gdn_state_quantizer.is_enabled + assert not model.gdn_w_quantizer.is_enabled + + cfg = quant_cfg(state=False, w=True) + cfg["algorithm"] = None mtq.quantize(model, cfg) - assert model.gdn_state_qdq_block_v == 64 + assert not model.gdn_state_quantizer.is_enabled + assert model.gdn_w_quantizer.is_enabled + + cfg = deepcopy(cfg) + cfg["quant_cfg"][-1]["cfg"]["pass_through_bwd"] = False + with pytest.raises(ValueError, match="supports only"): + mtq.quantize(model, cfg) def test_restore_legacy_gdn_without_new_quantizer_handles(): @@ -254,9 +256,7 @@ def test_restore_legacy_gdn_without_new_quantizer_handles(): model = mtq.quantize(TinyGatedDeltaNet(), config) state = mto.modelopt_state(model) for _, mode_state in state["modelopt_state_dict"]: - mode_state["config"].pop("linear_attention", None) metadata = mode_state["metadata"] - metadata.pop("linear_attention", None) for name in ("gdn_state_quantizer", "gdn_w_quantizer"): metadata["quantizer_state"].pop(name, None) restored = TinyGatedDeltaNet() diff --git a/tests/unit/torch/quantization/test_linear_attention_reference.py b/tests/unit/torch/quantization/test_linear_attention_reference.py index dd2aea8d592..d71818d0163 100644 --- a/tests/unit/torch/quantization/test_linear_attention_reference.py +++ b/tests/unit/torch/quantization/test_linear_attention_reference.py @@ -16,8 +16,7 @@ import pytest import torch import torch.nn.functional as F - -from modelopt.torch.quantization.linear_attention import ( +from _test_utils.torch.quantization.linear_attention_reference import ( chunk_gdn_reference, recurrent_delta_rule_reference, state_fp8_qdq_reference, From 95766de709ac05ed733be4c01066a7e832533826 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 28 Sep 2026 17:45:36 -0700 Subject: [PATCH 3/6] Move GDN test compilation into setup fixtures Warm standalone and Megatron forward/backward kernels before functional tests. Remove the 300-second overrides so test calls retain the default 120-second limit and report execution separately from compilation. Signed-off-by: Kai Xu --- .../test_fla_chunk_gated_delta_rule.py | 96 +++++++++++++++++-- .../plugins/test_megatron_gated_delta_net.py | 36 +++++-- 2 files changed, 116 insertions(+), 16 deletions(-) diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py index d470da9f491..9a11c454932 100644 --- a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py @@ -37,6 +37,11 @@ chunk_gated_delta_rule, ) +DTYPES = (torch.float32, torch.bfloat16) +QDQ_MODES = ((False, False), (False, True), (True, False), (True, True)) +STATE_LAYOUTS = ((False, False), (True, True)) +W_AXES = (None, (0, 1), (0, 1, 2)) + def require_state_qdq(): if torch.cuda.get_device_capability() < (8, 9): @@ -97,9 +102,84 @@ def values_and_grads(fn, args, state, **kwargs): return result, grads -# The first comparison compiles and autotunes both upstream and ModelOpt forward/backward kernels. -@pytest.mark.timeout(300) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.fixture(scope="module", autouse=True) +def _compile_gdn_kernels(): + """Keep compilation/autotuning in setup, including when selecting a single test.""" + supports_state_qdq = torch.cuda.get_device_capability() >= (8, 9) + for dtype in DTYPES: + if dtype == torch.float32 and IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0: + continue + args, state = make_inputs(dtype, grouped=False) + for kernel in (fla.chunk_gated_delta_rule, chunk_gated_delta_rule): + values_and_grads(kernel, args, state, output_final_state=True) + for packed, state_v_first in STATE_LAYOUTS: + args, state = make_inputs(dtype, packed=packed, state_v_first=state_v_first) + for state_qdq, quantize_w in QDQ_MODES: + if state_qdq and not supports_state_qdq: + continue + values_and_grads( + chunk_gated_delta_rule, + args, + state, + output_final_state=True, + state_qdq=state_qdq, + state_qdq_block_v=32, + w_quantizer=w_quantizer() if quantize_w else None, + state_v_first=state_v_first, + cu_seqlens=torch.tensor([0, 67, 145], device="cuda", dtype=torch.int32) + if packed + else None, + ) + + args, state = make_inputs() + quantizer = w_quantizer() + values_and_grads( + chunk_gated_delta_rule, args, state, output_final_state=True, w_quantizer=quantizer + ) + for axis in W_AXES: + output, _ = chunk_gated_delta_rule(*args, w_quantizer=w_quantizer(axis)) + torch.autograd.grad(output.square().sum(), args) + a_log = torch.zeros(2, device="cuda", requires_grad=True) + bias = torch.zeros_like(a_log, requires_grad=True) + output, final = chunk_gated_delta_rule( + *args, + initial_state=state, + output_final_state=True, + w_quantizer=quantizer, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + A_log=a_log, + dt_bias=bias, + use_beta_sigmoid_in_kernel=True, + allow_neg_eigval=True, + ) + torch.autograd.grad(output.square().sum() + final.square().sum(), (*args, state, a_log, bias)) + if supports_state_qdq: + for block_v in (16, 64, 128): + values_and_grads( + chunk_gated_delta_rule, + args, + state, + output_final_state=True, + state_qdq=True, + state_qdq_block_v=block_v, + ) + args, state = make_inputs(packed=True, state_v_first=True) + for block_v in (16, 32, 64, 128): + _, final = chunk_gated_delta_rule( + *args, + initial_state=state, + output_final_state=True, + state_v_first=True, + cu_seqlens=torch.tensor([0, 67, 145], device="cuda", dtype=torch.int32), + state_qdq=True, + state_qdq_block_v=block_v, + ) + torch.autograd.grad(final.sum(), state) + torch.cuda.synchronize() + + +@pytest.mark.parametrize("dtype", DTYPES) def test_disabled_matches_upstream_forward_and_backward(dtype): # Upstream Hopper/TileLang backward cannot handle grouped value heads. args, state = make_inputs(dtype, grouped=False) @@ -109,11 +189,9 @@ def test_disabled_matches_upstream_forward_and_backward(dtype): compare(a, e, 0.005 if dtype == torch.bfloat16 else 0.001) -@pytest.mark.parametrize( - ("state_qdq", "quantize_w"), [(False, False), (False, True), (True, False), (True, True)] -) -@pytest.mark.parametrize(("packed", "state_v_first"), [(False, False), (True, True)]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize(("state_qdq", "quantize_w"), QDQ_MODES) +@pytest.mark.parametrize(("packed", "state_v_first"), STATE_LAYOUTS) +@pytest.mark.parametrize("dtype", DTYPES) def test_numerical_surrogate_reference(state_qdq, quantize_w, packed, state_v_first, dtype): if state_qdq: require_state_qdq() @@ -161,7 +239,7 @@ def fn(*x): handle.remove() -@pytest.mark.parametrize("axis", [None, (0, 1), (0, 1, 2)]) +@pytest.mark.parametrize("axis", W_AXES) def test_zero_initial_state_and_output_only_training_loss(axis): args, _ = make_inputs() quantizer = w_quantizer(axis) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index 1fe3e7399b8..c46959c4d8d 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -143,16 +143,38 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): assert not torch.equal(restored_gdn[0].in_proj.weight, before) -# Cold FLA/TileLang compilation plus a two-rank checkpoint exceeds the lane's 120s default. -@pytest.mark.timeout(300) -@pytest.mark.parametrize("tp_size", [1, 2]) -@pytest.mark.parametrize("mode", ["state", "w", "both"]) -def test_gdn_qat_and_sharded_restore(request, tmp_path, tp_size, mode): - """Train through state/W QDQ after a Megatron distributed-checkpoint round trip.""" +def _compile_gdn_qat_kernels(rank, size, mode): + initialize_for_megatron( + tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED + ) + model = _make_model(size) + forward = get_forward(model) + with torch.no_grad(): + forward(model) + sites = ("state", "w") if mode == "both" else (mode,) + mtq.quantize(model, _gdn_config(sites)) + with torch.no_grad(): + forward(model) + model.train() + forward(model).sum().backward() + torch.cuda.synchronize() + + +@pytest.fixture +def compiled_gdn_workers(request, tp_size, mode): + """Warm the selected QAT path in the same workers, outside the test-call budget.""" if mode != "w" and torch.cuda.get_device_capability() < (8, 9): pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") workers = request.getfixturevalue(f"dist_workers_size_{tp_size}") - workers.run(_test_gdn_qat_helper, mode, tmp_path) + workers.run(_compile_gdn_qat_kernels, mode) + return workers + + +@pytest.mark.parametrize("tp_size", [1, 2]) +@pytest.mark.parametrize("mode", ["state", "w", "both"]) +def test_gdn_qat_and_sharded_restore(compiled_gdn_workers, tmp_path, mode): + """Train through state/W QDQ after a Megatron distributed-checkpoint round trip.""" + compiled_gdn_workers.run(_test_gdn_qat_helper, mode, tmp_path) def _test_gdn_context_parallel_helper(rank, size, mode): From 3896181aef1cecc6c85014b5b3ba37f1fdf0fb0b Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 28 Sep 2026 18:13:31 -0700 Subject: [PATCH 4/6] Clarify dynamic per-tile FP8 state quantization Signed-off-by: Kai Xu --- .../quantization/linear_attention/fla_chunk_delta_h.py | 7 +++---- .../configs/ptq/units/gdn_state_fp8_dynamic.yaml | 8 ++++---- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py index ca8b085470d..875cb39a7d0 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -50,7 +50,7 @@ @triton.jit def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr): - """[ModelOpt] Dynamic FP8 E4M3 scale over the up-to-four K tiles of one program's state.""" + """[ModelOpt] Dynamic FP8 E4M3 scale per full [K, BV] tile of one sequence and head.""" b_amax = tl.max(tl.abs(b_h1)) if K > 64: b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h2))) @@ -385,9 +385,8 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h4 = tl.dot(b_k, b_v, b_h4) - # [ModelOpt] Fake-quantize the state carried into the next chunk (and, after the last - # chunk, the stored final state) to FP8 E4M3. The scale is dynamic over this program's - # [K, BV] tile of the head state; BV == V makes it one scale per sequence and head. + # [ModelOpt] Dynamic per-tile FP8 QDQ of the next-chunk or stored final state. + # Recompute amax over [K, BV]; fp8_scalar_qdq applies that tile's scalar scale. if STATE_QDQ == 1: if K > 192: b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml index 114c86d5cad..ed5768ee599 100644 --- a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -13,10 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -# FP8 E4M3 fake quantization of the GatedDeltaNet (linear attention) recurrent state with a -# dynamic scale computed inside the kernel from the state itself (no calibration). The state is -# ``[N, H, K, V]`` for ``N`` sequences; ``axis: [0, 1]`` scales each sequence's head state on its -# own, with one scale per full-K by 64-value-column tile (two scales per head for V = 128). +# Dynamic per-tile FP8 E4M3 fake quantization of the GatedDeltaNet recurrent state. +# Recompute amax over each [K, 64] tile inside the kernel, independently per sequence and head +# of the [N, H, K, V] state (no calibration). ``axis: [0, 1]`` specifies sequence/head grouping; +# the kernel uses one scale per full-K by 64-value-column tile (two per head for V = 128). # This is the only configuration the module supports. Applied at every kernel chunk boundary # (64 tokens) during calibration, QAT and QAD; needs fla-core==0.5.1 and Triton. # See ``modelopt.torch.quantization.plugins.gated_delta_net``. From d03c73254264cc85943ee7f81e322898e11ae1a1 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sun, 27 Sep 2026 20:55:54 -0700 Subject: [PATCH 5/6] Add GDN/KDA Torch decode QAT with INT8 recurrent state Signed-off-by: Kai Xu --- CHANGELOG.rst | 2 + examples/llm_qat/README.md | 1 + examples/llm_qat/linear_attention/README.md | 304 ++++++++++++++++ .../configs/kda_decode_state_int8.json | 35 ++ .../llm_qat/linear_attention/requirements.txt | 4 + examples/llm_qat/linear_attention/train.py | 109 ++++++ .../linear_attention/fla_chunk_delta_h.py | 68 ++-- .../fla_chunk_gated_delta_rule.py | 12 +- .../quantization/linear_attention/int8.py | 28 ++ modelopt/torch/quantization/config.py | 7 + modelopt/torch/quantization/conversion.py | 55 ++- .../quantization/linear_attention/__init__.py | 8 +- .../linear_attention/_chunk_prefill.py | 146 ++++++++ .../quantization/linear_attention/config.py | 108 ++++++ .../quantization/linear_attention/decode.py | 341 ++++++++++++++++++ .../linear_attention/decode_prefill.py | 176 +++++++++ .../quantization/linear_attention/kda.py | 107 ++++++ .../quantization/linear_attention/prefill.py | 94 +++++ .../quantization/linear_attention/utils.py | 109 +++++- modelopt/torch/quantization/model_quant.py | 4 +- .../torch/quantization/plugins/__init__.py | 3 +- .../plugins/{gated_delta_net.py => gdn.py} | 65 ++-- modelopt/torch/quantization/plugins/kda.py | 103 ++++++ .../quantization/plugins/linear_attention.py | 85 +++++ .../torch/quantization/plugins/megatron.py | 6 +- modelopt/torch/quantization/tensor_quant.py | 3 +- modelopt_recipes/configs/ptq/units/README.md | 1 + .../units/default_disabled_quantizers.yaml | 4 + .../ptq/units/gdn_state_fp8_dynamic.yaml | 7 +- .../ptq/units/gdn_state_int8_dynamic.yaml | 26 ++ .../general/ptq/gdn_state_int8_dynamic.yaml | 38 ++ modelopt_recipes/ptq.md | 3 +- .../linear_attention_reference.py | 183 +++++++++- .../test_fla_chunk_gated_delta_rule.py | 60 ++- .../torch/quantization/plugins/test_kda.py | 189 ++++++++++ .../quantization/test_tensor_quant_cuda.py | 22 +- .../plugins/test_megatron_gated_delta_net.py | 118 +++++- .../{test_gated_delta_net.py => test_gdn.py} | 92 ++++- .../test_linear_attention_decode.py | 296 +++++++++++++++ .../test_linear_attention_hadamard.py | 227 ++++++++++++ .../test_linear_attention_int8.py | 99 +++++ 41 files changed, 3184 insertions(+), 164 deletions(-) create mode 100644 examples/llm_qat/linear_attention/README.md create mode 100644 examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json create mode 100644 examples/llm_qat/linear_attention/requirements.txt create mode 100644 examples/llm_qat/linear_attention/train.py create mode 100644 modelopt/torch/kernels/quantization/linear_attention/int8.py create mode 100644 modelopt/torch/quantization/linear_attention/_chunk_prefill.py create mode 100644 modelopt/torch/quantization/linear_attention/config.py create mode 100644 modelopt/torch/quantization/linear_attention/decode.py create mode 100644 modelopt/torch/quantization/linear_attention/decode_prefill.py create mode 100644 modelopt/torch/quantization/linear_attention/kda.py create mode 100644 modelopt/torch/quantization/linear_attention/prefill.py rename modelopt/torch/quantization/plugins/{gated_delta_net.py => gdn.py} (70%) create mode 100644 modelopt/torch/quantization/plugins/kda.py create mode 100644 modelopt/torch/quantization/plugins/linear_attention.py create mode 100644 modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml create mode 100644 modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml create mode 100644 tests/gpu/torch/quantization/plugins/test_kda.py rename tests/unit/torch/quantization/plugins/{test_gated_delta_net.py => test_gdn.py} (74%) create mode 100644 tests/unit/torch/quantization/test_linear_attention_decode.py create mode 100644 tests/unit/torch/quantization/test_linear_attention_hadamard.py create mode 100644 tests/unit/torch/quantization/test_linear_attention_int8.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 347806df32b..60d905778b0 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,6 +13,8 @@ Changelog *Quantization* +- Add experimental GDN/KDA decode-aware QAT with FP8 or INT8 recurrent states, optional INT8 value-axis Hadamard encoding, KDA decay rounding, and encoded-update replay. Supply explicit prefix lengths through the training phase context; prefix solves remain exact. + - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration. - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md `_ for details. diff --git a/examples/llm_qat/README.md b/examples/llm_qat/README.md index b4900ebea1f..aa099d88ebd 100644 --- a/examples/llm_qat/README.md +++ b/examples/llm_qat/README.md @@ -14,6 +14,7 @@ For background on how QAT enables low-precision accuracy recovery, see the [QAT/ | Background | How QAT/QAD work and when to use each | \[[Link](#background)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | | Support Matrix | Supported models, quantization formats, and backends | \[[Link](#support-matrix)\] | | | QLoRA | Model training with reduced GPU memory | \[[Link](#qlora-real-quantization)\] | | +| Linear Attention | GDN/KDA recurrent-state QAT and ReplaySSM example | \[[Link](linear_attention/README.md)\] | | | Advanced Topics | FSDP2 config, YAML options | \[[Link](#advanced-topics)\] | | | Results | Accuracy benchmarks | \[[Link](#results)\] | | | Resources | Extra links and references | \[[Link](#resources)\] | | diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md new file mode 100644 index 00000000000..7c44058512a --- /dev/null +++ b/examples/llm_qat/linear_attention/README.md @@ -0,0 +1,304 @@ +# Quantization-Aware Training for Linear Attention + +This example fine-tunes KDA attention parameters with recurrent-state fake +quantization, then saves a Hugging Face checkpoint with ModelOpt state. GDN/KDA +runtime support includes token writes, ReplaySSM, KDA decay approximation, and +FP8 or INT8 state QDQ. The INT8 recipes enable Hadamard rotation by default. + +## Run the example + +Install ModelOpt using the [QAT setup instructions](../README.md#quick-start), +then install the example dependencies. Run commands from the repository root. + +```bash +pip install -r examples/llm_qat/linear_attention/requirements.txt +python examples/llm_qat/linear_attention/train.py \ + --model /path/to/local-kda-model \ + --train-data /path/to/train.parquet \ + --output /path/to/qat-checkpoint \ + --train-steps 1 --length 128 --prefill-tokens 64 +``` + +Use a local model/tokenizer snapshot containing FLA `KimiDeltaAttention` layers, +a Parquet file with a `text` column, and a CUDA GPU. The example requires +`fla-core==0.5.1` and `flash-linear-attention==0.5.1`. If the local model requires +custom Python code, review it and explicitly pass `--trust-remote-code`. + +The default [INT8 configuration](configs/kda_decode_state_int8.json) leaves the +64-token prefix state unquantized and applies INT8 QDQ in a 32-value Hadamard +basis during the decode suffix. Value dimensions must be divisible by 32. +Loss uses suffix labels. Only KDA attention parameters are trained, in FP32 under +BF16 autocast; other parameters are frozen in BF16. The output contains model +weights, tokenizer files, and ModelOpt quantizer/policy state. Call +`mto.enable_huggingface_checkpointing()` before reloading it with +`AutoModelForCausalLM.from_pretrained` to restore those quantizers and policies. +It is a floating fake-quantized training checkpoint, not a compressed serving model. +This short example does not measure model-quality recovery or performance. + +## Enable state quantization + +### What the execution policy controls + +An **execution policy** is the layer's `LinearAttentionConfig`: the settings that +determine how the recurrence runs and when it applies quantization. Training +uses three separate inputs: + +| Input | Purpose | Supplied through | +| --- | --- | --- | +| `TensorQuantizer` settings | Enable a quantization site and choose its numerical format, scales, and gradient behavior | Recipe `quant_cfg`, such as dynamic INT8 with identity STE | +| Execution policy | Choose token or replay mode, Hadamard rotation, handoff quantization, and the execution backend | Recipe `linear_attention` entries | +| Batch metadata | Specify where each sequence switches from prefill to decode | `linear_attention_training_phase(model, prefill_lengths)` | + +The same INT8 quantizer settings can round the state after every token or round +a replay anchor every eight tokens. Those schedules produce different recurrent +states, so the policy must specify which computation training should emulate. + +`LinearAttentionConfig` holds the overall backend, chunk size, and state settings. +Its optional `decode` field contains a `LinearAttentionDecodeConfig` for token or +replay mode, state codec, decay approximation, and replay settings. The decode +implementation uses Torch operations and autograd. This is one nested configuration: supplying a `decode` +dictionary in the recipe constructs the nested config automatically. + +Each `linear_attention` entry uses `module_name` to select attention layers and +`cfg` to specify their policy. During `mtq.quantize`, ModelOpt stores that policy +as each matched layer's `linear_attention_config`. Quantizer settings and the +policy persist through ModelOpt save/restore; batch-specific prefill lengths +must be supplied again for each workload. + +### Load a state recipe + +State quantizers start disabled. The state recipe enables `*kda_state_quantizer` +with signed narrow-range INT8, dynamic scales, and `axis=(0, 1)`. Use +`*gdn_state_quantizer` for GDN. For E4M3, set the quantizer config to +`{"num_bits": [4, 3], "type": "dynamic", "axis": [0, 1]}` and set +`decode.state_codec="tile"`. +Setting `prefill_state_qdq=True` alone does not enable a quantizer. + +Start with the token-state recipe and select one of the schedules below **before** +calling `mtq.quantize`: + +```python +import json +from pathlib import Path + +recipe = json.loads( + Path("examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json").read_text() +) +policy = recipe["linear_attention"][0]["cfg"] +decode = policy["decode"] +``` + +For GDN, the complete recipe composes the INT8 quantizer unit with the Hadamard +execution policy: + +```python +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe + +gdn_recipe = load_recipe("general/ptq/gdn_state_int8_dynamic").quantize +mtq.quantize(gdn_model, gdn_recipe) +``` + +Use the phase context shown below for each forward/backward. Importing only +`configs/ptq/units/gdn_state_int8_dynamic` configures the quantizer without +selecting Hadamard. Both complete recipes use the Torch decode implementation. + +### Why training needs prefill/decode boundaries + +One training forward can simulate prompt prefill followed by recurrent decode. +The INT8 + Hadamard recipes apply different state quantization schedules to those +phases, so the workload must specify where each sequence switches to decode. +The training batch supplies all tokens; this simulates decode arithmetic without +running a text-generation loop. Total sequence length alone does not identify +the prompt portion. + +For a 128-token sequence with `prefill_lengths=[96]`, the default token-state +recipe runs these steps: + +1. **Prefill:** process tokens 0–95 with `prefill_state_qdq=False`, leaving state + unquantized during the prefix. +2. **Handoff:** carry the resulting state into decode, rotate its value dimension + into the Hadamard basis, and apply INT8 QDQ because `quantize_initial=True`. +3. **Decode:** process tokens 96–127 recurrently, applying INT8 QDQ after every + state update. Later tokens consume the rounded state; outputs are transformed + back to the original value basis. + +The boundary is independent of `chunk_size=64`: this 96-token prefix contains a +full 64-token chunk and a partial 32-token chunk. The phase switches after token +95, not after each chunk. + +Two 128-token sequences can use `prefill_lengths=[64, 96]` in the same batch: +their decode suffixes then contain 64 and 32 tokens, respectively. The next batch +can have different lengths while reusing the same quantization recipe. This is +why the execution policy is saved in `linear_attention_config`, while the lengths +are supplied per batch through `linear_attention_training_phase`. They are not +saved as part of the model's quantization policy. The context selects numerical +phases; the caller still supplies training labels and any loss masking. + +The combined prefill/decode path requires explicit lengths, including `[0]` for +decode only or `[T]` for an all-prefix sequence of length `T`. The GDN chunk-only +FLA path described below applies its chunk schedule throughout and needs no +phase context. See the tables below for the corresponding quantization settings. + +### Choose the prefill and decode boundaries + +The table assumes the state quantizer is enabled. The INT8 recipes default to +`"int8_hadamard32"`; prefix state QDQ requires explicitly selecting `"tile"`. +Prefix lengths are supplied separately through `linear_attention_training_phase`. + +| Desired state QDQ | `decode.state_codec` | `decode.mode` | `decode.prefill_state_qdq` | Where rounding occurs | +| --- | --- | --- | --- | --- | +| Token decode only (default) | `"int8_hadamard32"` | `"token"` | `False` | At the first nonempty decode handoff and after every suffix token. | +| Prefill and token decode | `"tile"` | `"token"` | `True` | At prefix initialization, each prefix chunk write, decode handoff, and every suffix token. | +| Replay anchors only | `"int8_hadamard32"` | `"replay"` | `False` | At decode handoff and each replay-window refresh. | +| Prefill and replay anchors | `"tile"` | `"replay"` | `True` | At prefix initialization and chunk writes, then decode handoff and replay-window refreshes. | + +To enable **prefill and token decode**: + +```python +decode.update(mode="token", replay=None, state_codec="tile", prefill_state_qdq=True) +``` + +For **token decode only**, keep the supplied INT8 recipe's defaults: +`state_codec="int8_hadamard32"` and `prefill_state_qdq=False`. +Prefix state remains unquantized until it enters the decode path. + +To enable **ReplaySSM anchor quantization** with an eight-token window: + +```python +decode.update( + mode="replay", + state_codec="int8_hadamard32", + prefill_state_qdq=False, + replay={"window": 8, "factor_qdq": False, "encoding": "once"}, +) +``` + +Set `state_codec="tile"` and `prefill_state_qdq=True` to add prefix state QDQ +to this replay configuration, using unrotated INT8 for both phases. +`factor_qdq=False` above isolates state/anchor quantization. Set it to `True` to +also quantize buffered keys and updates to FP8. + +`decode.quantize_initial=True` is the default: it quantizes the incoming state +once at the first nonempty decode handoff. Set it to `False` to skip that initial +rounding while keeping later token writes or anchor refreshes quantized. This +setting does not disable prefix state QDQ. With both phases enabled, prefix-final +rounding and decode-handoff rounding are separate configured events. + +`chunk_size=64` counts **tokens per prefill chunk**. `state.block_v=64` counts +**value channels per execution tile**. The tile codec shares a scale across all +key channels and this value tile; Hadamard uses one scale per key channel and +32 values. A replay `window=8` counts **suffix tokens between anchor refreshes**. + +### Run only the desired phase + +For a batch containing one sequence of `T` tokens, choose the context lengths as +follows; for larger batches, provide one length per sequence: + +| Workload | Context argument | Required setting | +| --- | --- | --- | +| Prefill followed by decode | `[64]`, with `T > 64` | Select either prefix setting above. | +| Decode only | `[0]` | State quantizer enabled; token or replay policy. | +| Prefill only | `[T]` | `state_codec="tile"`, `prefill_state_qdq=True`; the decode suffix is empty. | + +An empty suffix creates no decode quantization event. The combined interface has +one state quantizer per layer: it does not offer a switch to quantize the prefix +while leaving a **nonempty** decode suffix unquantized. Disable the state quantizer +to disable both state and anchor QDQ; replay factor QDQ has its own toggle. + +The `int8_hadamard32` codec applies to decode token writes or replay anchors only. +It requires `prefill_state_qdq=False`; use the tile codec to quantize prefix state. + +Save a modified recipe as JSON and pass it to `train.py --quant-config`. `--prefill-tokens` sets the phase split; it does not enable state +quantization. The integration loop below applies the configured `recipe` directly. + +### GDN chunk-only FLA training + +For GDN's existing chunked FLA path, enable the GDN state quantizer and omit the +`linear_attention` execution-policy entry: + +```python +gdn_recipe = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*gdn_state_quantizer", + "cfg": recipe["quant_cfg"][1]["cfg"], + }, + ], + "algorithm": None, +} +# Apply mtq.quantize(gdn_model, gdn_recipe), then use normal forward/backward. +``` + +This explicitly selects unrotated tile QDQ. It rounds the initial state and each +64-token chunk's final state, including a partial final chunk. It needs no phase +context and leaves W/projection quantizers +disabled. KDA uses the materialized backend and can use the all-prefix context +shown above. + +## Integrate with a training loop + +Apply the configured `recipe` above with `mtq.quantize`, then supply one prefix length per sequence. +Keep the phase context active through backward so activation-checkpoint +recomputation uses the same prefix/decode split. + +```python +import torch + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase + +mtq.quantize(model, recipe) +optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5) +model.train() +optimizer.zero_grad(set_to_none=True) + +# Two 128-token sequences: 64 and 96 prefill tokens, respectively. +# ids and labels have shape [2, 128] and are on the model's device. +with linear_attention_training_phase(model, [64, 96]): + with torch.autocast("cuda", dtype=torch.bfloat16): + loss = model(input_ids=ids, labels=labels, use_cache=False).loss + loss.backward() +optimizer.step() +``` + +For GDN, select `*gdn_state_quantizer` in the recipe instead of +`*kda_state_quantizer`. State quantization is dynamic, so these recipes use +`algorithm=None` without a calibration pass. Replay factors have their own FP8 +toggle. + +Policies persist through ModelOpt save/restore. Per-batch prefix lengths are +runtime metadata and must be supplied for each workload. The context restores +previous lengths on exit and supports nesting. Concurrent forwards on the same +model instance with different phase contexts are unsupported. + +## Replay, decay, and Hadamard options + +Token mode rounds each recurrent-state write. Replay mode retains an anchor and +ordered key/update factors, rounding the anchor at each `replay.window` refresh. +`factor_qdq` independently enables FP8 QDQ for those factors. `readout="working"` +reads the state before its write quantization; `"stored"` reads it afterward. + +For KDA decay approximation, set `decode["decay_log_step"] = 1 / 256` before +conversion. It rounds suffix log retention before exponentiation with identity +STE gradients. Prefix decay remains exact. + +The default INT8 codec applies a 32-point orthonormal Hadamard transform along +the value axis, quantizes token states or replay anchors in that basis, and +transforms outputs back. Value dimensions must be divisible by 32; `state.block_v` must be 32, 64, +or 128. Scales group one key channel and 32 values, independent of execution tile +width. INT8 codes use half-away-from-zero rounding with FP16 stored scales; +the optional tile codec instead uses nearest-even rounding and FP32 scales. + +## Training boundaries + +The Torch implementation uses autograd through initial states, chunk handoff, +token writes, and replay refreshes. The decode recurrence runs token by token in +Python, so long training suffixes can be slow. + +Prefill prefixes use exact chunk algebra with optional state QDQ. This example +has no prefill GEMM QDQ or approximate inverse. FLA KDA requires `use_cache=False`; +serving cache objects are rejected. ModelOpt saves execution policies, while +per-batch prefix lengths must be supplied again during training. Distributed +decode training and model-quality recovery require separate qualification. diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json b/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json new file mode 100644 index 00000000000..8a29cda8139 --- /dev/null +++ b/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json @@ -0,0 +1,35 @@ +{ + "quant_cfg": [ + { + "quantizer_name": "*", + "enable": false + }, + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": 8, + "type": "dynamic", + "axis": [ + 0, + 1 + ], + "unsigned": false, + "narrow_range": true + } + } + ], + "algorithm": null, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "torch", + "state_codec": "int8_hadamard32", + "prefill_state_qdq": false + } + } + } + ] +} diff --git a/examples/llm_qat/linear_attention/requirements.txt b/examples/llm_qat/linear_attention/requirements.txt new file mode 100644 index 00000000000..29636d78e52 --- /dev/null +++ b/examples/llm_qat/linear_attention/requirements.txt @@ -0,0 +1,4 @@ +fla-core==0.5.1 +flash-linear-attention==0.5.1 +pyarrow>=15 +transformers==4.57.1 diff --git a/examples/llm_qat/linear_attention/train.py b/examples/llm_qat/linear_attention/train.py new file mode 100644 index 00000000000..7662c1ca593 --- /dev/null +++ b/examples/llm_qat/linear_attention/train.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fine-tune KDA attention parameters with recurrent-state fake quantization.""" + +import argparse +import json +from pathlib import Path + +import pyarrow.parquet as pq +import torch +from fla.layers.kda import KimiDeltaAttention +from transformers import AutoModelForCausalLM, AutoTokenizer + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase + + +def main(): + """Quantize, train on local text, and save a ModelOpt-aware Hugging Face checkpoint.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument( + "--train-data", type=Path, required=True, help="Parquet file with a text column" + ) + parser.add_argument( + "--quant-config", + type=Path, + default=Path(__file__).with_name("configs") / "kda_decode_state_int8.json", + ) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--train-steps", type=int, default=1) + parser.add_argument("--length", type=int, default=128) + parser.add_argument("--prefill-tokens", type=int, default=64) + parser.add_argument("--seed", type=int, default=2026) + parser.add_argument("--learning-rate", type=float, default=1e-5) + options = parser.parse_args() + if options.train_steps < 1 or options.length < 2: + parser.error("train-steps must be positive and length must be at least two") + if not 0 <= options.prefill_tokens < options.length: + parser.error("prefill-tokens must leave at least one suffix label") + torch.manual_seed(options.seed) + mto.enable_huggingface_checkpointing() + tokenizer = AutoTokenizer.from_pretrained( + options.model, local_files_only=True, trust_remote_code=options.trust_remote_code + ) + text = "\n\n".join(pq.read_table(options.train_data, columns=["text"])["text"].to_pylist()) + tokens = tokenizer.encode(text, add_special_tokens=False) + needed = options.train_steps * options.length + if len(tokens) < needed: + parser.error(f"Training data has {len(tokens)} tokens; {needed} are required") + blocks = torch.tensor(tokens[:needed]).reshape(options.train_steps, options.length) + model = AutoModelForCausalLM.from_pretrained( + options.model, + local_files_only=True, + trust_remote_code=options.trust_remote_code, + torch_dtype=torch.float32, + ) + layers = [module for module in model.modules() if isinstance(module, KimiDeltaAttention)] + if not layers: + raise ValueError("The model must contain FLA KimiDeltaAttention layers") + # Keep attention parameters and optimizer moments in FP32; freeze the rest in BF16. + model.requires_grad_(False) + for layer in layers: + layer.requires_grad_(True) + for parameter in model.parameters(): + if not parameter.requires_grad: + parameter.data = parameter.data.to(torch.bfloat16) + model.cuda() + mtq.quantize(model, json.loads(options.quant_config.read_text())) + model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + trainable = [p for p in model.parameters() if p.requires_grad] + optimizer = torch.optim.AdamW(trainable, lr=options.learning_rate, weight_decay=0) + model.train() + for step, block in enumerate(blocks, 1): + ids = block.unsqueeze(0).cuda() + labels = ids.clone() + labels[:, : options.prefill_tokens] = -100 + optimizer.zero_grad(set_to_none=True) + # Keep phase metadata active during activation-checkpoint recomputation. + with linear_attention_training_phase(model, [options.prefill_tokens]): + with torch.autocast("cuda", dtype=torch.bfloat16): + loss = model(input_ids=ids, labels=labels, use_cache=False).loss + if not torch.isfinite(loss): + raise RuntimeError("Non-finite training loss") + loss.backward() + torch.nn.utils.clip_grad_norm_(trainable, 1.0, error_if_nonfinite=True) + optimizer.step() + print(f"step={step} loss={loss.item():.4f}", flush=True) + model.save_pretrained(options.output) + tokenizer.save_pretrained(options.output) + + +if __name__ == "__main__": + main() diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py index 875cb39a7d0..84da262fd9f 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py @@ -1,5 +1,5 @@ # Adapted from: https://github.com/fla-org/flash-linear-attention/blob/516143e31fce/fla/ops/common/chunk_delta_h.py -# Adapted with modifications (marked [ModelOpt]): optional in-kernel FP8 E4M3 fake quantization +# Adapted with modifications (marked [ModelOpt]): optional in-kernel FP8 E4M3 / INT8 fake quantization # of the carried chunk state (STATE_QDQ), BV as an explicit launch argument, no fla backend # dispatch or config-cache autotune. # @@ -42,15 +42,18 @@ from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq +from .int8 import int8_scalar_qdq + # ``STATE_QDQ`` modes of the forward state kernel. STATE_QDQ_OFF = 0 STATE_QDQ_FP8_DYNAMIC = 1 # FP8 E4M3, one dynamic scale per program tile ([K, BV] of one head) +STATE_QDQ_INT8_DYNAMIC = 2 STATE_QDQ_MAX_BLOCK_V = 128 @triton.jit -def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr): - """[ModelOpt] Dynamic FP8 E4M3 scale per full [K, BV] tile of one sequence and head.""" +def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr, STATE_QDQ: tl.constexpr): + """[ModelOpt] Dynamic scale per full [K, BV] tile of one sequence and head.""" b_amax = tl.max(tl.abs(b_h1)) if K > 64: b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h2))) @@ -58,7 +61,18 @@ def _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K: tl.constexpr): b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h3))) if K > 192: b_amax = tl.maximum(b_amax, tl.max(tl.abs(b_h4))) - return tl.where(b_amax > 0, b_amax / 448.0, 1.0) + if STATE_QDQ == 2: + return tl.where(b_amax > 0, b_amax * (1.0 / 127.0), 1.0) + else: + return tl.where(b_amax > 0, b_amax / 448.0, 1.0) + + +@triton.jit +def _state_scalar_qdq(value, scale, STATE_QDQ: tl.constexpr): + if STATE_QDQ == 2: + return int8_scalar_qdq(value, scale) + else: + return fp8_scalar_qdq(value, scale) NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] @@ -215,23 +229,23 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( p_h0_4 = h0 + o_k4[:, None] * V + o_v[None, :] m_h0_4 = m_k4[:, None] & m_v[None, :] b_h4 += tl.load(p_h0_4, mask=m_h0_4, other=0.0).to(tl.float32) - # [ModelOpt] A state read from an FP8 cache is quantized before the first chunk uses it. - if STATE_QDQ == 1: + # [ModelOpt] A state read from a quantized cache is quantized before the first chunk uses it. + if STATE_QDQ != 0: if K > 192: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K, STATE_QDQ=STATE_QDQ) elif K > 128: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K) + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K, STATE_QDQ=STATE_QDQ) elif K > 64: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K) + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K, STATE_QDQ=STATE_QDQ) else: - b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K) - b_h1 = fp8_scalar_qdq(b_h1, b_scale) + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K, STATE_QDQ=STATE_QDQ) + b_h1 = _state_scalar_qdq(b_h1, b_scale, STATE_QDQ) if K > 64: - b_h2 = fp8_scalar_qdq(b_h2, b_scale) + b_h2 = _state_scalar_qdq(b_h2, b_scale, STATE_QDQ) if K > 128: - b_h3 = fp8_scalar_qdq(b_h3, b_scale) + b_h3 = _state_scalar_qdq(b_h3, b_scale, STATE_QDQ) if K > 192: - b_h4 = fp8_scalar_qdq(b_h4, b_scale) + b_h4 = _state_scalar_qdq(b_h4, b_scale, STATE_QDQ) # main recurrence for i_t in range(NT): @@ -385,24 +399,24 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h4 = tl.dot(b_k, b_v, b_h4) - # [ModelOpt] Dynamic per-tile FP8 QDQ of the next-chunk or stored final state. - # Recompute amax over [K, BV]; fp8_scalar_qdq applies that tile's scalar scale. - if STATE_QDQ == 1: + # [ModelOpt] Dynamic per-tile FP8/INT8 QDQ of the next-chunk or stored final state. + # Recompute amax over [K, BV]; _state_scalar_qdq applies that tile's scalar scale. + if STATE_QDQ != 0: if K > 192: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K) + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h4, K=K, STATE_QDQ=STATE_QDQ) elif K > 128: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K) + b_scale = _state_qdq_scale(b_h1, b_h2, b_h3, b_h3, K=K, STATE_QDQ=STATE_QDQ) elif K > 64: - b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K) + b_scale = _state_qdq_scale(b_h1, b_h2, b_h2, b_h2, K=K, STATE_QDQ=STATE_QDQ) else: - b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K) - b_h1 = fp8_scalar_qdq(b_h1, b_scale) + b_scale = _state_qdq_scale(b_h1, b_h1, b_h1, b_h1, K=K, STATE_QDQ=STATE_QDQ) + b_h1 = _state_scalar_qdq(b_h1, b_scale, STATE_QDQ) if K > 64: - b_h2 = fp8_scalar_qdq(b_h2, b_scale) + b_h2 = _state_scalar_qdq(b_h2, b_scale, STATE_QDQ) if K > 128: - b_h3 = fp8_scalar_qdq(b_h3, b_scale) + b_h3 = _state_scalar_qdq(b_h3, b_scale, STATE_QDQ) if K > 192: - b_h4 = fp8_scalar_qdq(b_h4, b_scale) + b_h4 = _state_scalar_qdq(b_h4, b_scale, STATE_QDQ) if STORE_FINAL_STATE: if STATE_V_FIRST: @@ -885,8 +899,8 @@ def state_qdq_tile_v(V: int, state_qdq: int, state_qdq_block_v: int | None) -> i """ if state_qdq == STATE_QDQ_OFF: return 64 if check_shared_mem("ada") else 32 - if state_qdq != STATE_QDQ_FP8_DYNAMIC: - raise ValueError(f"Unsupported state_qdq mode {state_qdq}; expected 0 or 1.") + if state_qdq not in (STATE_QDQ_FP8_DYNAMIC, STATE_QDQ_INT8_DYNAMIC): + raise ValueError(f"Unsupported state_qdq mode {state_qdq}; expected 0, 1, or 2.") BV = min(triton.next_power_of_2(V), 64) if state_qdq_block_v is None else state_qdq_block_v if BV < 16 or BV > STATE_QDQ_MAX_BLOCK_V or BV & (BV - 1): raise ValueError( diff --git a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py index a5609cab102..f34b8a3384b 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py +++ b/modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py @@ -59,6 +59,7 @@ from .fla_chunk_delta_h import ( STATE_QDQ_FP8_DYNAMIC, + STATE_QDQ_INT8_DYNAMIC, STATE_QDQ_OFF, chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h, @@ -636,7 +637,8 @@ def chunk_gated_delta_rule( raise ValueError("ModelOpt GDN supports only chunk_size=64; FLA WY backward assumes 64.") # [ModelOpt] state_qdq: 0 keeps fla's numerics; 1 fake-quantizes the state carried between - # chunks to FP8 E4M3 with a dynamic scale per [K, state_qdq_block_v] tile of each head. + # chunks to FP8 E4M3; 2 uses signed narrow-range INT8. Both use a dynamic scale per + # [K, state_qdq_block_v] tile of each head. state_qdq = kwargs.pop("state_qdq", STATE_QDQ_OFF) state_qdq_block_v = kwargs.pop("state_qdq_block_v", None) # w_quantizer: dynamic FP8 TensorQuantizer applied to the WY tensor @@ -647,11 +649,13 @@ def chunk_gated_delta_rule( dt_bias = kwargs.pop("dt_bias", None) if kwargs: raise TypeError(f"Unexpected keyword arguments: {', '.join(sorted(kwargs))}") - if state_qdq not in (STATE_QDQ_OFF, STATE_QDQ_FP8_DYNAMIC): - raise ValueError(f"`state_qdq` must be 0 or 1, got {state_qdq}.") + if state_qdq not in (STATE_QDQ_OFF, STATE_QDQ_FP8_DYNAMIC, STATE_QDQ_INT8_DYNAMIC): + raise ValueError(f"`state_qdq` must be 0, 1, or 2, got {state_qdq}.") if w_quantizer is not None: validate_gdn_quantizer(w_quantizer, name="gdn_w_quantizer") - if state_qdq and (not q.is_cuda or torch.cuda.get_device_capability(q.device) < (8, 9)): + if state_qdq == STATE_QDQ_FP8_DYNAMIC and ( + not q.is_cuda or torch.cuda.get_device_capability(q.device) < (8, 9) + ): raise RuntimeError("GDN state QDQ requires native E4M3 conversion on CUDA SM89 or newer.") if (state_qdq != STATE_QDQ_OFF or w_quantizer is not None) and cp_context is not None: raise ValueError("State or w quantization is not supported together with `cp_context`.") diff --git a/modelopt/torch/kernels/quantization/linear_attention/int8.py b/modelopt/torch/kernels/quantization/linear_attention/int8.py new file mode 100644 index 00000000000..d9e16a29dc8 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/int8.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Signed narrow-range INT8 QDQ for recurrent-state tiles.""" + +import triton +import triton.language as tl +from triton.language.extra.cuda import libdevice + +__all__ = [] + + +@triton.jit +def int8_scalar_qdq(value, scale): + codes = libdevice.nearbyint(tl.div_rn(value, scale)) + return tl.minimum(tl.maximum(codes, -127.0), 127.0) * scale diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index f8847aa1fc1..e8db5c8446a 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -161,6 +161,8 @@ from modelopt.torch.opt.config_loader import load_config from modelopt.torch.utils.network import ConstructorLike +from .linear_attention.config import LinearAttentionPolicyEntry + class QuantizerCfgEntry(ModeloptBaseConfig): """A single entry in a ``quant_cfg`` list.""" @@ -1567,6 +1569,11 @@ def _dict_to_entry(key: str, value) -> list[dict[str, Any]]: class QuantizeConfig(ModeloptBaseConfig): """Default configuration for ``quantize`` mode.""" + linear_attention: list[LinearAttentionPolicyEntry] = ModeloptField( + default=[], + title="Execution policies for supported linear-attention training modules", + ) + quant_cfg: QuantizeQuantCfgType = ModeloptField( default=[{"quantizer_name": "*", "cfg": {"num_bits": 8, "axis": None}}], title="Quantization configuration", diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 78a6e03a16b..db194a0f465 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -28,6 +28,7 @@ from modelopt.torch.opt.dynamic import _DMRegistryCls from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict from modelopt.torch.utils import get_unwrapped_name +from modelopt.torch.utils.logging import print_rank_0 from .config import ( QuantizeConfig, @@ -36,6 +37,7 @@ _QuantizeExportConfig, normalize_quant_cfg_list, ) +from .linear_attention.config import LinearAttentionConfig from .nn import ( QuantModule, QuantModuleRegistry, @@ -67,7 +69,7 @@ def convert_to_quantized_model(model: ModelLikeModule, config: QuantizeConfig) - replace_quant_module(model, version=ModeloptStateManager(model).state_version) set_quantizer_by_cfg(model, config.get("quant_cfg", [])) - _validate_linear_attention_quantizers(model) + _apply_linear_attention_policy(model, config) metadata = {} update_quantize_metadata(model, config, metadata) @@ -137,6 +139,7 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: details regarding how MCore sharded checkpoint is restored, see modelopt.torch.opt.plugins.mcore_dist_checkpointing.restore_sharded_modelopt_state. """ + _apply_linear_attention_policy(model, config, metadata.get("linear_attention")) if "quantizer_state" not in metadata: # MCore sharded checkpoint (`torch-dist`) has its quantizer_state stored as the # extra_state of `QuantModule`. The quantizer_state is resumed with @@ -144,13 +147,14 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: return model quantizer_state_dict = dict(metadata["quantizer_state"]) - # Older checkpoints predate these disabled handles; preserve their baseline path. - for name, module in _linear_attention_modules(model).items(): - for handle in ("gdn_state_quantizer", "gdn_w_quantizer"): - key = f"{name}.{handle}" if name else handle - quantizer = getattr(module, handle) - if key not in quantizer_state_dict and not quantizer.is_enabled: - quantizer_state_dict[key] = quantizer.get_modelopt_state() + if "linear_attention" not in metadata: + # Older checkpoints predate these disabled handles; preserve their baseline path. + for name, module in _linear_attention_modules(model).items(): + for handle in module.linear_attention_quantizer_names: + key = f"{name}.{handle}" if name else handle + quantizer = getattr(module, handle) + if key not in quantizer_state_dict and not quantizer.is_enabled: + quantizer_state_dict[key] = quantizer.get_modelopt_state() unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys() extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys() @@ -203,6 +207,12 @@ def update_quantize_metadata( model: nn.Module, config: QuantizeConfig, metadata: MetadataDict ) -> None: """Update the quantizer state in the metadata dict.""" + policies = { + name: module.linear_attention_config.model_dump() + for name, module in _linear_attention_modules(model).items() + } + if policies: + metadata["linear_attention"] = policies metadata["quantizer_state"] = quantizer_state(model) if shared_state_metadata := SharedWeightGlobalAmaxState.metadata(model): metadata["shared_quant_states"] = shared_state_metadata @@ -212,18 +222,39 @@ def update_quantize_metadata( def _linear_attention_modules(model): # Optional framework plugins import conversion; defer this import to avoid that cycle. - from .plugins.gated_delta_net import GatedDeltaNetStateQuantMixin + from .plugins.linear_attention import _LinearAttentionQuantMixin return { get_unwrapped_name(name, model): module for name, module in model.named_modules() - if isinstance(module, GatedDeltaNetStateQuantMixin) + if isinstance(module, _LinearAttentionQuantMixin) } -def _validate_linear_attention_quantizers(model): - for module in _linear_attention_modules(model).values(): +def _apply_linear_attention_policy(model, config, saved_policies=None): + modules = _linear_attention_modules(model) + # getattr also handles pickled pre-extension config objects with no field in __dict__. + for entry in getattr(config, "linear_attention", []): + matches = [name for name in modules if fnmatch.fnmatch(name, entry.module_name)] + if not matches: + raise ValueError( + f"linear_attention rule {entry.module_name!r} matches no supported modules" + ) + for name in matches: + modules[name].linear_attention_config = entry.cfg.model_copy(deep=True) + if saved_policies is not None: + if saved_policies.keys() != modules.keys(): + raise ApplyModeError( + "Saved linear_attention policies do not match the restored modules" + ) + for name, policy in saved_policies.items(): + modules[name].linear_attention_config = LinearAttentionConfig(**policy) + for name, module in modules.items(): module.validate_linear_attention() + if getattr(config, "linear_attention", []): + print_rank_0( + f"Linear attention {name!r}: {module.linear_attention_config.model_dump()}" + ) def quantizer_state(model: nn.Module) -> dict[str, Any]: diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py index 477eeb89847..305bb2cd81f 100644 --- a/modelopt/torch/quantization/linear_attention/__init__.py +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -13,4 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Helpers for linear-attention quantization.""" +"""Numerical policies and differentiable kernels for linear attention.""" + +from .config import * +from .decode import * +from .decode_prefill import * +from .kda import * +from .prefill import * diff --git a/modelopt/torch/quantization/linear_attention/_chunk_prefill.py b/modelopt/torch/quantization/linear_attention/_chunk_prefill.py new file mode 100644 index 00000000000..fddff0303ae --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/_chunk_prefill.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Exact chunked prefix computation for decode-aware training.""" + +import torch + +from .utils import _prepare, _state_qdq + +__all__ = [] + + +def chunk_gdn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + chunk_size: int = 64, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + state_v_first: bool = False, + state_qdq: bool = False, + state_qdq_block_v: int = 64, + state_format: str = "fp8_e4m3", +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute a GDN prefix with optional QDQ on the initial state and chunk writes.""" + if g.ndim != 3 or chunk_size <= 0: + raise ValueError("chunk_gdn requires scalar GDN gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format) + pieces = [] + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + gc = g[b, lo:hi].transpose(0, 1).cumsum(-1) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + # Mask before exp: upper-triangle positive differences can overflow for long decay. + causal = torch.ones(hi - lo, hi - lo, device=q.device, dtype=torch.bool).tril() + decay = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() + lower = (bc * (kc @ kc.transpose(-1, -2)) * decay).tril(-1) + matrix = lower + torch.eye(hi - lo, device=q.device, dtype=q.dtype) + gate = gc.exp().unsqueeze(-1) + rhs = torch.cat((bc * vc, bc * kc * gate), dim=-1) + solved = torch.linalg.solve_triangular(matrix, rhs, upper=False, unitriangular=True) + u, w = solved.split((v.shape[-1], k.shape[-1]), dim=-1) + updated_values = u - w @ state + local_scores = ((qc * scale @ kc.transpose(-1, -2)) * decay).tril() + output = (qc * (scale * gc.exp()).unsqueeze(-1)) @ state + pieces.append((output + local_scores @ updated_values).transpose(0, 1)) + weighted_keys = kc * (gc[..., -1:] - gc).exp().unsqueeze(-1) + state = state * gc[..., -1].exp()[:, None, None] + state = state + weighted_keys.transpose(-1, -2) @ updated_values + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format) + outputs.append(torch.cat(pieces)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final + + +def chunk_kda( + q, + k, + v, + g, + beta, + *, + chunk_size=64, + scale=None, + initial_state=None, + cu_seqlens=None, + state_v_first=False, + state_qdq=False, + state_qdq_block_v=64, + state_format="fp8_e4m3", +): + """Compute a KDA prefix with optional QDQ on the initial state and chunk writes.""" + if g.ndim != 4 or chunk_size <= 0: + raise ValueError("chunk_kda requires per-key gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format) + pieces = [] + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + prefix = g[b, lo:hi].transpose(0, 1).cumsum(-2) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + gate = prefix.exp() + lower_rows, score_rows = [], [] + for row in range(hi - lo): + decayed_keys = ( + kc[..., : row + 1, :] + * (prefix[..., row : row + 1, :] - prefix[..., : row + 1, :]).exp() + ) + right = decayed_keys.transpose(-1, -2) + left = kc[..., row : row + 1, :] * bc[..., row : row + 1, :] + interaction = left @ right + score = qc[..., row : row + 1, :] * scale @ right + padding = hi - lo - row - 1 + lower_rows.append(torch.nn.functional.pad(interaction, (0, padding))) + score_rows.append(torch.nn.functional.pad(score, (0, padding))) + lower = torch.cat(lower_rows, dim=-2).tril(-1) + scores = torch.cat(score_rows, dim=-2) + identity = torch.eye(hi - lo, device=q.device, dtype=q.dtype).expand_as(lower) + inverse = torch.linalg.solve_triangular( + identity + lower, identity, upper=False, unitriangular=True + ) + u = inverse @ (bc * vc) + w = inverse @ (bc * kc * gate) + updated = u - w @ state + pieces.append(((qc * scale * gate) @ state + scores @ updated).transpose(0, 1)) + weighted_keys = kc * (prefix[..., -1:, :] - prefix).exp() + state = state * gate[..., -1, :, None] + weighted_keys.transpose(-1, -2) @ updated + if state_qdq: + state = _state_qdq(state, state_qdq_block_v, state_format) + outputs.append(torch.cat(pieces)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py new file mode 100644 index 00000000000..1f3b610fe75 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Saved execution policy for GDN training-time numerical emulation.""" + +from typing import Literal + +from pydantic import Field, model_validator + +from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField + +__all__ = [ + "LinearAttentionConfig", + "LinearAttentionDecodeConfig", + "LinearAttentionPolicyEntry", + "LinearAttentionReplayConfig", +] + + +class _StateConfig(ModeloptBaseConfig): + mode: Literal["chunk"] = ModeloptField(default="chunk") + block_v: Literal[16, 32, 64, 128] = ModeloptField(default=64) + quantize_initial: Literal[True] = ModeloptField(default=True) + + +class _SolveConfig(ModeloptBaseConfig): + method: Literal["exact"] = ModeloptField(default="exact") + + +class LinearAttentionReplayConfig(ModeloptBaseConfig): + """Anchor refresh and encoded rank-one update policy.""" + + window: int = Field(default=8, ge=1, le=64, strict=True) + factor_qdq: bool = ModeloptField(default=True) + encoding: Literal["once", "reencode"] = ModeloptField(default="once") + + +class LinearAttentionDecodeConfig(ModeloptBaseConfig): + """Explicit suffix recurrence; workload supplies per-sequence prefix lengths.""" + + mode: Literal["token", "replay"] = ModeloptField(default="token") + implementation: Literal["torch"] = ModeloptField(default="torch") + readout: Literal["working", "stored"] = ModeloptField(default="stored") + quantize_initial: bool = ModeloptField(default=True) + prefill_state_qdq: bool = ModeloptField(default=False) + state_codec: Literal["tile", "int8_hadamard32"] = ModeloptField(default="tile") + decay_log_step: float | None = Field(default=None, gt=0, allow_inf_nan=False) + replay: LinearAttentionReplayConfig | None = ModeloptField(default=None) + + @model_validator(mode="after") + def _validate_replay(self): + if (self.mode == "replay") != (self.replay is not None): + raise ValueError("replay settings must be supplied exactly when mode='replay'") + if self.state_codec == "int8_hadamard32" and self.prefill_state_qdq: + raise ValueError( + "Hadamard state QDQ starts at decode handoff; disable prefill_state_qdq" + ) + return self + + +class LinearAttentionConfig(ModeloptBaseConfig): + """GDN chunk-64 policy; unsupported numerical modes fail config validation. + + ``state.block_v`` defines one dynamic scale per ``[Dk, block_v]`` tile of each + sequence/head. The initial state and every chunk's final state are rounded when + ``gdn_state_quantizer`` is enabled. Outputs use the incoming rounded state. + Decode's ``int8_hadamard32`` codec instead fixes scales to one key channel and + 32 values; ``state.block_v`` remains the execution tile width. + """ + + schema_version: Literal[1] = ModeloptField(default=1) + backend: Literal["fla", "matmul"] = ModeloptField(default="fla") + chunk_size: Literal[64] = ModeloptField(default=64) + state: _StateConfig = ModeloptField(default=_StateConfig()) + solve: _SolveConfig = ModeloptField(default=_SolveConfig()) + decode: LinearAttentionDecodeConfig | None = ModeloptField(default=None) + + @model_validator(mode="after") + def _validate_decode_backend(self): + if (self.backend == "matmul") != (self.decode is not None): + raise ValueError("The exact-prefix matmul backend requires an explicit decode policy") + if self.decode is not None and self.decode.state_codec == "int8_hadamard32": + if self.state.block_v < 32: + raise ValueError("int8_hadamard32 requires block_v >= 32") + return self + + +class LinearAttentionPolicyEntry(ModeloptBaseConfig): + """Assign a complete policy to supported modules matching ``module_name``. + + Rules apply in order: the last match wins, without merging nested fields. + A rule must match at least one supported linear-attention module. + """ + + module_name: str = Field(...) + cfg: LinearAttentionConfig = ModeloptField(default=LinearAttentionConfig()) diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py new file mode 100644 index 00000000000..2a981c597ec --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -0,0 +1,341 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Explicit token-state and encoded-update replay references for QAT.""" + +from dataclasses import dataclass + +import torch + +from .config import LinearAttentionDecodeConfig +from .utils import _fp8_quantize + +__all__ = [ + "EncodedLinearAttentionTensor", + "LinearAttentionCarry", + "ReplayEntry", + "recurrent_decode", + "recurrent_decode_reference", +] + + +@dataclass +class EncodedLinearAttentionTensor: + """Fake decoded values plus detached scale metadata; no compressed storage claim.""" + + values: torch.Tensor + scales: torch.Tensor | None + format: str + block_v: int | None + + +@dataclass +class ReplayEntry: + """An already computed, encoded rank-one update and its log retention.""" + + key: EncodedLinearAttentionTensor + update: EncodedLinearAttentionTensor + log_retention: torch.Tensor + + +@dataclass +class LinearAttentionCarry: + """Explicit anchor/update state, position, and codec contract across calls.""" + + anchor: EncodedLinearAttentionTensor + entries: tuple[ReplayEntry, ...] + position: int + started: bool + signature: str + value_basis: str = "identity" + + @property + def cursor(self): + """Number of encoded updates since the last anchor refresh.""" + return len(self.entries) + + def reconstruct(self, *, original_basis=True): + """Replay entries, returning the original value basis unless explicitly disabled.""" + state = self.anchor.values + for entry in self.entries: + gate = entry.log_retention + decay = gate.exp().unsqueeze(-1) + if gate.ndim == 1: + decay = decay.unsqueeze(-1) + state = state * decay + entry.key.values.unsqueeze(-1) * entry.update.values.unsqueeze( + -2 + ) + if original_basis and self.value_basis == "hadamard32": + state = _hadamard32(state) + return state + + +def _hadamard32(value): + """Apply the orthonormal Sylvester transform to contiguous 32-value groups.""" + shape = value.shape + for width in (1, 2, 4, 8, 16): + pairs = value.reshape(*shape[:-1], -1, 2, width) + left, right = pairs.unbind(-2) + value = torch.stack((left + right, left - right), dim=-2).reshape(shape) + return value * (32**-0.5) + + +class _IdentityGradient(torch.autograd.Function): + @staticmethod + def forward(ctx, value, rounded): + return rounded + + @staticmethod + def backward(ctx, gradient): + return gradient, None + + +def _encode(value, enabled, block_v, *, state=False, state_format="fp8_e4m3", state_codec="tile"): + if not enabled: + return EncodedLinearAttentionTensor(value, None, "identity", None) + if state and state_codec == "int8_hadamard32": + with torch.no_grad(): + groups = value.float().reshape(*value.shape[:-1], -1, 32) + scales = (groups.abs().amax(-1) / 127).clamp_min(6e-8) + normalized = groups / scales.unsqueeze(-1) + codes = normalized.sign() * (normalized.abs() + 0.5).floor() + metadata = scales.to(torch.float16) + decoded = (codes.clamp(-127, 127) * metadata.float().unsqueeze(-1)).reshape_as(value) + return EncodedLinearAttentionTensor( + _IdentityGradient.apply(value, decoded.to(value.dtype)), metadata, "int8", 32 + ) + rounded, scales = [], [] + for part in value.split(block_v, dim=-1): + if state_format == "fp8_e4m3": + tensor = part.flatten(-2) if state else part + axis = tuple(range(tensor.ndim - 1)) or None + decoded, scale = _fp8_quantize(tensor, axis) + rounded.append(decoded.reshape_as(part)) + scales.append(scale.squeeze(-1)) + else: + with torch.no_grad(): + axes = (-2, -1) if state else (-1,) + part = part.float() + amax = part.abs().amax(dim=axes, keepdim=True) + scale = torch.where(amax > 0, amax / 127.0, torch.ones_like(amax)) + rounded.append((part / scale).round().clamp(-127, 127) * scale) + scales.append(scale[..., 0, 0] if state else scale[..., 0]) + decoded = torch.cat(rounded, dim=-1).to(value.dtype) + if state_format == "int8": + decoded = _IdentityGradient.apply(value, decoded) + return EncodedLinearAttentionTensor(decoded, torch.stack(scales, dim=-1), state_format, block_v) + + +def _signature(config, state_qdq, block_v, state_format): + return ( + config.model_dump_json(exclude={"implementation", "prefill_state_qdq"}) + + f"/{state_qdq}/{block_v}/{state_format}" + ) + + +def _sum_keys(value): + keys = value.shape[-2] + padded = 1 << (keys - 1).bit_length() + if keys != padded: + value = torch.nn.functional.pad(value, (0, 0, 0, padded - keys)) + while value.shape[-2] > 1: + half = value.shape[-2] // 2 + value = value[..., :half, :] + value[..., half:, :] + return value[..., 0, :] + + +def _round_log_gate(gate, step): + if step is None: + return gate + rounded = (gate / step).round() * step + return _IdentityGradient.apply(gate, rounded.detach()) + + +def _prepare_carry( + q, k, v, g, beta, config, state_qdq, block_v, initial_state, carry, position, state_format +): + if q.ndim != 3 or k.shape != q.shape or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k/v must have aligned [T,H,D] shapes") + if beta.shape != q.shape[:2] or g.shape not in (beta.shape, q.shape): + raise ValueError("beta must be [T,H]; g must be [T,H] or [T,H,Dk]") + if block_v not in (16, 32, 64, 128): + raise ValueError("block_v must be 16, 32, 64, or 128") + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") + hadamard = config.state_codec == "int8_hadamard32" + if hadamard: + if state_qdq and state_format != "int8": + raise ValueError("int8_hadamard32 requires INT8 state quantization") + if v.shape[-1] % 32 or block_v < 32: + raise ValueError("int8_hadamard32 requires Dv divisible by 32 and block_v >= 32") + signature = _signature(config, state_qdq, block_v, state_format) + if carry is not None and initial_state is not None: + raise ValueError("Supply either carry or initial_state") + shape = (q.shape[1], q.shape[2], v.shape[2]) + if carry is None: + initial_state = q.new_zeros(shape) if initial_state is None else initial_state + carry = LinearAttentionCarry( + _encode(initial_state, False, block_v, state=True, state_format=state_format), + (), + position, + False, + signature, + ) + if carry.signature != signature or carry.anchor.values.shape != shape: + raise ValueError("Carry policy or state shape does not match this recurrence") + if config.mode == "token" and carry.entries: + raise ValueError("Token carry cannot contain replay entries") + if config.replay is not None and carry.cursor >= config.replay.window: + raise ValueError("Replay cursor must be below the refresh window") + if len(q) and not carry.started: + initial = _hadamard32(carry.anchor.values) if hadamard else carry.anchor.values + anchor = _encode( + initial, + state_qdq and config.quantize_initial, + block_v, + state=True, + state_format=state_format, + state_codec=config.state_codec, + ) + carry = LinearAttentionCarry( + anchor, (), carry.position, True, signature, "hadamard32" if hadamard else "identity" + ) + return carry, signature + + +def recurrent_decode_reference( + q, + k, + v, + g, + beta, + *, + config: LinearAttentionDecodeConfig, + state_qdq=False, + state_format="fp8_e4m3", + block_v=64, + initial_state=None, + carry=None, + position=0, + scale=None, +): + """Run one preactivated sequence [T,H,D] with explicit token/replay write events. + + Keys and value heads must already be aligned. Scalar GDN or per-key-channel + KDA log gates are accepted. Outputs and all returned carry values retain their + graphs. An empty call performs no write or initial-state quantization. + """ + carry, signature = _prepare_carry( + q, k, v, g, beta, config, state_qdq, block_v, initial_state, carry, position, state_format + ) + if len(q) == 0: + # Keep empty input gradients defined without introducing a state write. + zero = (q.sum() + k.sum() + v.sum() + g.sum() + beta.sum()) * 0 + output = v + zero + return output, carry + if config.state_codec == "int8_hadamard32": + v = _hadamard32(v) + scale = q.shape[-1] ** -0.5 if scale is None else scale + outputs = [] + with torch.autocast(device_type=q.device.type, enabled=False): + for t in range(len(q)): + entries = carry.entries + if config.replay is not None and config.replay.encoding == "reencode": + entries = tuple( + ReplayEntry( + _encode(e.key.values, config.replay.factor_qdq, k.shape[-1]), + _encode(e.update.values, config.replay.factor_qdq, block_v), + e.log_retention, + ) + for e in entries + ) + carry = LinearAttentionCarry( + carry.anchor, entries, carry.position, True, signature, carry.value_basis + ) + state = carry.reconstruct(original_basis=False) + gate = _round_log_gate(g[t], config.decay_log_step) + decay = gate.exp().unsqueeze(-1) + if gate.ndim == 1: + decay = decay.unsqueeze(-1) + key = _encode(k[t], config.replay is not None and config.replay.factor_qdq, k.shape[-1]) + decayed = state * decay + residual = v[t] - _sum_keys(key.values.unsqueeze(-1) * decayed) + update = _encode( + beta[t].unsqueeze(-1) * residual, + config.replay is not None and config.replay.factor_qdq, + block_v, + ) + working = decayed + key.values.unsqueeze(-1) * update.values.unsqueeze(-2) + if config.replay is not None: + entries = (*entries, ReplayEntry(key, update, gate)) + refresh = config.replay is None or len(entries) == config.replay.window + if refresh: + anchor = _encode( + working, + state_qdq, + block_v, + state=True, + state_format=state_format, + state_codec=config.state_codec, + ) + entries = () + stored = anchor.values + else: + anchor = carry.anchor + stored = working + next_carry = LinearAttentionCarry( + anchor, entries, carry.position + 1, True, signature, carry.value_basis + ) + read = working if config.readout == "working" else stored + outputs.append(_sum_keys(q[t].unsqueeze(-1) * read) * scale) + carry = next_carry + output = torch.stack(outputs) + if config.state_codec == "int8_hadamard32": + output = _hadamard32(output) + return output, carry + + +def recurrent_decode( + q, + k, + v, + g, + beta, + *, + config: LinearAttentionDecodeConfig, + state_qdq=False, + state_format="fp8_e4m3", + block_v=64, + initial_state=None, + carry=None, + position=0, + scale=None, +): + """Run the Torch training implementation with explicit differentiable carry.""" + return recurrent_decode_reference( + q, + k, + v, + g, + beta, + config=config, + state_qdq=state_qdq, + state_format=state_format, + block_v=block_v, + initial_state=initial_state, + carry=carry, + position=position, + scale=scale, + ) diff --git a/modelopt/torch/quantization/linear_attention/decode_prefill.py b/modelopt/torch/quantization/linear_attention/decode_prefill.py new file mode 100644 index 00000000000..0a0f69b4ae7 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/decode_prefill.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differentiable explicit per-sequence prefill/decode phase handoff.""" + +from contextlib import contextmanager +from itertools import pairwise + +import torch + +from ._chunk_prefill import chunk_gdn, chunk_kda +from .decode import recurrent_decode + +__all__ = ["linear_attention_training_phase"] + + +def _lengths(values): + if isinstance(values, torch.Tensor): + values = values.tolist() + values = tuple(values) + if any(type(value) is not int or value < 0 for value in values): + raise ValueError("Prefill lengths must be nonnegative integers") + return values + + +@contextmanager +def linear_attention_training_phase(model, prefill_lengths): + """Supply explicit sequence phases through forward and checkpointed backward. + + Runtime phase metadata is local to the converted layers and restored on exit. + Keep this context active through backward when activation checkpointing is used. + """ + # The plugin imports this numerical package during quantization initialization. + from ..plugins.linear_attention import _LinearAttentionQuantMixin + + lengths = _lengths(prefill_lengths) + layers = [ + m + for m in model.modules() + if isinstance(m, _LinearAttentionQuantMixin) + and m.linear_attention_config.decode is not None + ] + if not layers: + raise ValueError("The model has no converted decode-aware linear-attention layers") + previous = [getattr(m, "_linear_attention_prefill_lengths", None) for m in layers] + try: + for module in layers: + module._linear_attention_prefill_lengths = lengths + yield model + finally: + for module, original in zip(layers, previous): + module._linear_attention_prefill_lengths = original + + +def _decode_prefill( + q, + k, + v, + g, + beta, + *, + policy, + state_qdq, + state_format, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + state_v_first, + output_dtype, + prefill_lengths, +): + if prefill_lengths is None: + raise ValueError("Decode-aware training requires explicit per-sequence prefill lengths") + if q.ndim != 4 or k.shape != q.shape or v.ndim != 4 or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k and v must have compatible [B,T,H,D] shapes") + batch, length, key_heads, keys = q.shape + heads, values = v.shape[2:] + if batch < 1 or key_heads < 1 or heads % key_heads or beta.shape != (batch, length, heads): + raise ValueError("Invalid batch/head dimensions or beta shape") + if g.shape not in (beta.shape, (*beta.shape, keys)): + raise ValueError("Invalid GDN/KDA log-retention shape") + q, k = (x.repeat_interleave(heads // key_heads, dim=2) for x in (q, k)) + boundaries = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens + if boundaries is None: + sequences = [(b, 0, length) for b in range(batch)] + else: + bounds = _lengths(boundaries) + if ( + batch != 1 + or len(bounds) < 2 + or bounds[0] != 0 + or bounds[-1] != length + or any(a > b for a, b in pairwise(bounds)) + ): + raise ValueError( + "Packed boundaries must partition a batch of one, allowing empty entries" + ) + sequences = [(0, a, b) for a, b in pairwise(bounds)] + prefixes = _lengths(prefill_lengths) + if len(prefixes) != len(sequences) or any( + p > end - start for p, (_, start, end) in zip(prefixes, sequences) + ): + raise ValueError("Supply one valid prefill length per sequence") + if initial_state is None: + states = q.new_zeros(len(sequences), heads, keys, values) + else: + states = initial_state.transpose(-1, -2) if state_v_first else initial_state + states = states.to(q.dtype) + if states.shape != (len(sequences), heads, keys, values): + raise ValueError("Initial state shape does not match sequence/head dimensions") + active = [n for n, prefix in enumerate(prefixes) if prefix] + prefix_outputs, prefix_states = {}, {} + if active: + packed = [ + torch.cat( + [ + tensor[sequences[n][0], sequences[n][1] : sequences[n][1] + prefixes[n]] + for n in active + ] + ).unsqueeze(0) + for tensor in (q, k, v, g, beta) + ] + offsets = [0] + for n in active: + offsets.append(offsets[-1] + prefixes[n]) + prefix_fn = chunk_kda if g.ndim == 4 else chunk_gdn + with torch.autocast(device_type=q.device.type, enabled=False): + output, final = prefix_fn( + *packed, + state_qdq=state_qdq and policy.decode.prefill_state_qdq, + state_format=state_format, + scale=scale, + initial_state=states[active], + cu_seqlens=torch.tensor(offsets), + state_v_first=False, + chunk_size=policy.chunk_size, + state_qdq_block_v=policy.state.block_v, + ) + for index, n in enumerate(active): + prefix_outputs[n] = output[0, offsets[index] : offsets[index + 1]] + prefix_states[n] = final[index] + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + split = start + prefixes[n] + suffix, carry = recurrent_decode( + *(x[b, split:end] for x in (q, k, v, g, beta)), + config=policy.decode, + state_qdq=state_qdq, + state_format=state_format, + block_v=policy.state.block_v, + initial_state=prefix_states.get(n, states[n]), + position=prefixes[n], + scale=scale, + ) + prefix = prefix_outputs.get(n, q.new_empty(0, heads, values)) + outputs.append(torch.cat((prefix, suffix))) + finals.append(carry.reconstruct()) + output = torch.stack(outputs) if boundaries is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + if state_v_first: + final = final.transpose(-1, -2) + return output.to(output_dtype), final if output_final_state else None diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py new file mode 100644 index 00000000000..446dd6f7d0d --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/kda.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differentiable KDA prefill with stable per-channel decay interactions.""" + +import torch +import torch.nn.functional as F + +from .decode_prefill import _decode_prefill + +__all__ = ["matmul_kda"] + + +def matmul_kda( + q, + k, + v, + g, + beta, + *, + policy, + state_qdq=False, + state_format="fp8_e4m3", + scale=None, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=False, + use_beta_sigmoid_in_kernel=False, + allow_neg_eigval=False, + A_log=None, # noqa: N803 - match the FLA kernel signature + dt_bias=None, + safe_gate=False, + lower_bound=None, + cu_seqlens=None, + cu_seqlens_cpu=None, + state_v_first=False, + chunk_size=64, + cp_context=None, + disable_recompute=False, + return_intermediate_states=False, + prefill_lengths=None, +): + """Normalize KDA inputs and run exact prefix plus configured suffix recurrence. + + Gates use the loaded FLA model's activation formula and per-key log retention. + """ + if policy.backend != "matmul" or chunk_size != policy.chunk_size: + raise ValueError("matmul_kda requires backend='matmul' and its configured chunk size") + if cp_context is not None or disable_recompute or return_intermediate_states: + raise NotImplementedError( + "KDA matmul does not support CP or FLA recompute/intermediate flags" + ) + if allow_neg_eigval and not use_beta_sigmoid_in_kernel: + raise ValueError("allow_neg_eigval requires use_beta_sigmoid_in_kernel") + if safe_gate and use_gate_in_kernel and (lower_bound is None or not -5 <= lower_bound < 0): + raise ValueError("safe_gate requires a lower_bound in [-5, 0)") + if lower_bound is not None and lower_bound >= 0: + raise ValueError("lower_bound must be negative") + output_dtype = q.dtype + dtype = torch.float64 if output_dtype == torch.float64 else torch.float32 + q, k, v, g, beta = (x.to(dtype) for x in (q, k, v, g, beta)) + if g.ndim != 4: + raise ValueError("matmul_kda requires per-key-channel KDA log gates") + if use_qk_l2norm_in_kernel: + q, k = (x * (x.square().sum(-1, keepdim=True) + 1e-6).rsqrt() for x in (q, k)) + if use_gate_in_kernel: + if A_log is None: + raise ValueError("Fused KDA gate requires A_log") + if dt_bias is not None: + g = g + dt_bias.to(dtype).reshape(g.shape[-2:]) + rate = A_log.to(dtype).exp().reshape(g.shape[-2], 1) + g = -rate * F.softplus(g) if lower_bound is None else lower_bound * (rate * g).sigmoid() + if use_beta_sigmoid_in_kernel: + beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) + if policy.decode is not None: + return _decode_prefill( + q, + k, + v, + g, + beta, + policy=policy, + state_qdq=state_qdq, + state_format=state_format, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + state_v_first=state_v_first, + output_dtype=output_dtype, + prefill_lengths=prefill_lengths, + ) + raise ValueError("An explicit decode policy is required") diff --git a/modelopt/torch/quantization/linear_attention/prefill.py b/modelopt/torch/quantization/linear_attention/prefill.py new file mode 100644 index 00000000000..4a602f8f5bc --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/prefill.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Batched differentiable GDN prefill with materialized numerical boundaries.""" + +import torch +import torch.nn.functional as F + +from .config import LinearAttentionConfig +from .decode_prefill import _decode_prefill + +__all__ = ["matmul_gdn"] + + +def matmul_gdn( + q, + k, + v, + g, + beta, + *, + policy: LinearAttentionConfig, + state_qdq=False, + state_format="fp8_e4m3", + scale=None, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=False, + use_beta_sigmoid_in_kernel=False, + allow_neg_eigval=False, + A_log=None, # noqa: N803 - match the FLA kernel signature + dt_bias=None, + cu_seqlens=None, + cu_seqlens_cpu=None, + state_v_first=False, + chunk_size=64, + cp_context=None, + prefill_lengths=None, +): + """Normalize GDN inputs and run exact prefix plus configured suffix recurrence. + + Uses FP32 working arithmetic (FP64 for double inputs), with floating QDQ state. + Prefixes use exact chunk algebra; no configurable operand rounding is enabled. + """ + if policy.backend != "matmul" or chunk_size != policy.chunk_size: + raise ValueError("matmul_gdn requires backend='matmul' and its configured chunk size") + if cp_context is not None: + raise NotImplementedError("GDN matmul emulation does not support context parallelism") + output_dtype = q.dtype + dtype = torch.float64 if output_dtype == torch.float64 else torch.float32 + q, k, v, g, beta = (x.to(dtype) for x in (q, k, v, g, beta)) + if use_qk_l2norm_in_kernel: + q, k = (x * (x.square().sum(-1, keepdim=True) + 1e-6).rsqrt() for x in (q, k)) + if use_gate_in_kernel: + if A_log is None or dt_bias is None: + raise ValueError("Fused GDN gate requires A_log and dt_bias") + g = -A_log.to(dtype).exp() * F.softplus(g + dt_bias.to(dtype)) + if use_beta_sigmoid_in_kernel: + beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) + if g.ndim != 3: + raise ValueError("matmul_gdn requires scalar GDN log gates") + if policy.decode is not None: + return _decode_prefill( + q, + k, + v, + g, + beta, + policy=policy, + state_qdq=state_qdq, + state_format=state_format, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + state_v_first=state_v_first, + output_dtype=output_dtype, + prefill_lengths=prefill_lengths, + ) + raise ValueError("An explicit decode policy is required") diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py index 4f1e18a174d..68353bd28df 100644 --- a/modelopt/torch/quantization/linear_attention/utils.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -15,18 +15,35 @@ """Shared helpers for linear-attention quantization.""" -from ..nn import TensorQuantizer +from __future__ import annotations + +from itertools import pairwise +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from ..nn import TensorQuantizer __all__ = [] -def validate_gdn_quantizer(quantizer: TensorQuantizer, *, name: str) -> None: - """Check the dynamic E4M3 contract and the custom backward's identity STE.""" +def validate_gdn_quantizer( + quantizer: TensorQuantizer, + *, + name: str, + num_bits: tuple[int | tuple[int, int], ...] = ((4, 3),), +) -> None: + """Check supported formats and the custom backward's identity STE.""" + # Numerical helpers load while QuantizeConfig initializes; defer the quantizer import. + from ..nn import TensorQuantizer + if not isinstance(quantizer, TensorQuantizer): raise ValueError(f"{name} requires a single TensorQuantizer") if not ( quantizer._dynamic - and quantizer.num_bits == (4, 3) + and quantizer.num_bits in num_bits + and (quantizer.num_bits != 8 or (not quantizer.unsigned and quantizer.narrow_range)) and quantizer.block_sizes is None and quantizer.fake_quant and quantizer._pass_through_bwd @@ -37,7 +54,89 @@ def validate_gdn_quantizer(quantizer: TensorQuantizer, *, name: str) -> None: and not quantizer._use_constant_amax ): raise ValueError( - f"{name} supports only dynamic E4M3 fake quantization with " + f"{name} supports only dynamic fake quantization with num_bits in {num_bits}, " "pass_through_bwd=True, no block_sizes, rotation, pre-scaling, bias, constant " "amax, or custom backend. Other gradient rules and formats are not implemented." ) + + +def _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first): + if q.ndim != 4 or k.shape != q.shape or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k must have shape [B,T,Hk,Dk] and v shape [B,T,Hv,Dv]") + batch, length, heads, keys = q.shape + value_heads, values = v.shape[2:] + if length == 0 or value_heads % heads: + raise ValueError("nonempty sequences and Hv divisible by Hk are required") + if beta.shape != (batch, length, value_heads) or g.shape not in ( + beta.shape, + (*beta.shape, keys), + ): + raise ValueError("beta must be [B,T,Hv]; g must be [B,T,Hv] or [B,T,Hv,Dk]") + if cu_seqlens is None: + sequences = [(b, 0, length) for b in range(batch)] + else: + boundaries = cu_seqlens.tolist() + if ( + batch != 1 + or len(boundaries) < 2 + or boundaries[0] != 0 + or boundaries[-1] != length + or any(a >= b for a, b in pairwise(boundaries)) + ): + raise ValueError("cu_seqlens must partition a packed batch of size one") + sequences = [(0, a, b) for a, b in pairwise(boundaries)] + if initial_state is None: + state = q.new_zeros(len(sequences), value_heads, keys, values) + else: + state = initial_state.transpose(-1, -2) if state_v_first else initial_state + if state.shape != (len(sequences), value_heads, keys, values): + raise ValueError("initial_state shape does not match sequence/head dimensions") + return ( + q.repeat_interleave(value_heads // heads, dim=2), + k.repeat_interleave(value_heads // heads, dim=2), + state, + sequences, + ) + + +def _fp8_quantize(value: torch.Tensor, axis): + """Apply ModelOpt FP8 QDQ and return its detached dequantization scales.""" + # QuantizeConfig imports this module before the quantizer classes are initialized. + from ..config import QuantizerAttributeConfig + from ..nn import TensorQuantizer + + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=(4, 3), type="dynamic", axis=axis, pass_through_bwd=True) + ) + quantized = quantizer(value) + amax = quantizer._get_amax(value).float() + safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) + return quantized, torch.div(448.0, safe_amax).reciprocal() + + +def _state_qdq(state: torch.Tensor, block_v: int = 64, state_format: str = "fp8_e4m3"): + """Dynamic state-tile QDQ with detached scales and identity STE.""" + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") + if block_v not in (16, 32, 64, 128): + raise ValueError("block_v must be 16, 32, 64, or 128") + if state_format == "fp8_e4m3": + axis = tuple(range(state.ndim - 2)) or None + return torch.cat( + [ + _fp8_quantize(tile.flatten(-2), axis)[0].reshape_as(tile) + for tile in state.split(block_v, dim=-1) + ], + dim=-1, + ) + with torch.no_grad(): + rounded = [] + for tile in state.float().split(block_v, dim=-1): + amax = tile.abs().amax(dim=(-2, -1), keepdim=True) + limit = 127.0 + scale = torch.where(amax > 0, amax / limit, torch.ones_like(amax)) + normalized = (tile / scale).clamp(-limit, limit) + codes = normalized.round() + rounded.append(codes * scale) + quantized = torch.cat(rounded, dim=-1).to(state.dtype) + return state + (quantized - state).detach() diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index fbf9bb299fd..5aac82cacfe 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -32,7 +32,7 @@ from modelopt.torch.opt.utils import forward_with_reshard from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.conversion import ( - _validate_linear_attention_quantizers, + _apply_linear_attention_policy, preserve_quantizer_attributes_context, set_quantizer_attributes_partial, set_quantizer_by_cfg, @@ -330,7 +330,7 @@ def forward_loop(model) -> None: else: # Already quantized, so lets apply the quant_cfg from the config set_quantizer_by_cfg(model, quantize_config.quant_cfg) - _validate_linear_attention_quantizers(model) + _apply_linear_attention_policy(model, quantize_config) # Fail before calibration rather than after exporting an unquantized checkpoint. _check_weight_quantization_took_effect(model, quantize_config) return calibrate(model, config.get("algorithm"), forward_loop=forward_loop) diff --git a/modelopt/torch/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py index 963452ff509..e1b2ec7a1ae 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -39,7 +39,8 @@ from .attention import * from .custom import * -from .gated_delta_net import * +from .gdn import * +from .kda import * with import_plugin("diffusers"): from .diffusion.diffusers import * diff --git a/modelopt/torch/quantization/plugins/gated_delta_net.py b/modelopt/torch/quantization/plugins/gdn.py similarity index 70% rename from modelopt/torch/quantization/plugins/gated_delta_net.py rename to modelopt/torch/quantization/plugins/gdn.py index df70683cc2f..8d0fe616390 100644 --- a/modelopt/torch/quantization/plugins/gated_delta_net.py +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -17,9 +17,10 @@ The chunked gated-delta-rule kernel keeps each head's ``[K, V]`` recurrent state in fp32 inside one Triton launch and carries it from chunk to chunk. To emulate a deployment that stores that -state in FP8, ModelOpt runs an adapted copy of the kernel +state in FP8 or INT8, ModelOpt runs an adapted copy of the kernel (:mod:`modelopt.torch.kernels.quantization.linear_attention`) that fake-quantizes the state to -E4M3 at the end of every chunk, with a scale computed inside the kernel from the state itself. +E4M3 or signed narrow-range INT8 at the end of every chunk, with a scale computed inside the +kernel from the state itself. The backward pass recomputes the same quantized states and passes the state gradient straight through the quantization, so QAT and QAD train against the quantized recurrence. A second quantizer covers ``w``, the WY-transformed keys that multiply the state; ``w`` is a regular tensor, @@ -32,9 +33,8 @@ import torch -from ..config import QuantizerAttributeConfig -from ..linear_attention.utils import validate_gdn_quantizer -from ..nn import QuantModule, TensorQuantizer +from ..linear_attention.prefill import matmul_gdn +from .linear_attention import _LinearAttentionQuantMixin __all__ = ["GatedDeltaNetStateQuantMixin"] @@ -63,39 +63,22 @@ def _state_qdq_chunk_gated_delta_rule() -> GatedDeltaRuleFn: return chunk_gated_delta_rule -class GatedDeltaNetStateQuantMixin(QuantModule): +class GatedDeltaNetStateQuantMixin(_LinearAttentionQuantMixin): """Adds ``gdn_state_quantizer`` and ``gdn_w_quantizer`` to a GatedDeltaNet module. Subclasses route the module's chunked gated-delta-rule call through :meth:`_state_quantized_chunk_gated_delta_rule`. Both quantizers start disabled; enable them with ``quant_cfg`` entries on ``*gdn_state_quantizer`` / ``*gdn_w_quantizer`` such as the ``configs/ptq/units/gdn_state_fp8_dynamic`` and ``gdn_w_fp8_dynamic`` recipe units. The state - quantizer carries the fused QDQ configuration. Both sites currently require dynamic - E4M3 and identity STE. State QDQ uses fixed 64-column value tiles and 64-token chunks. + quantizer carries the fused QDQ configuration. State supports dynamic E4M3 or signed + narrow-range INT8; W supports dynamic E4M3. + Both sites require identity STE. The execution policy is saved in ModelOpt metadata. """ - def _setup(self): - for name in ("gdn_state_quantizer", "gdn_w_quantizer"): - self._register_temp_attribute( - name, TensorQuantizer(QuantizerAttributeConfig(enable=False)) - ) - - def validate_linear_attention(self) -> None: - """Reject numerical settings that the fused training path cannot implement.""" - for name in ("gdn_state_quantizer", "gdn_w_quantizer"): - quantizer = getattr(self, name) - if quantizer.is_enabled: - validate_gdn_quantizer(quantizer, name=name) - # The state handle configures fused QDQ; W's grouping is executed by TensorQuantizer. - if self.gdn_state_quantizer.is_enabled and self.gdn_state_quantizer.axis != (0, 1): - raise ValueError( - "gdn_state_quantizer supports only axis=(0, 1) with 64-column value tiling" - ) - - def modelopt_post_restore(self, prefix: str = ""): - """Validate restored quantizers against the fused training kernel requirements.""" - super().modelopt_post_restore(prefix) - self.validate_linear_attention() + @property + def gdn_state_qdq_block_v(self) -> int: + """Value-column scale grouping from the saved execution policy.""" + return self.linear_attention_config.state.block_v def _state_quantized_chunk_gated_delta_rule( self, gated_delta_rule: GatedDeltaRuleFn, *args: Any, **kwargs: Any @@ -104,7 +87,7 @@ def _state_quantized_chunk_gated_delta_rule( self.validate_linear_attention() quantize_state = self.gdn_state_quantizer.is_enabled and self.gdn_state_quantizer._if_quant quantize_w = self.gdn_w_quantizer.is_enabled - if not (quantize_state or quantize_w): + if not self.linear_attention_is_enabled: return gated_delta_rule(*args, **kwargs) while isinstance(gated_delta_rule, partial): args = (*gated_delta_rule.args, *args) @@ -115,14 +98,26 @@ def _state_quantized_chunk_gated_delta_rule( "GatedDeltaNet quantization supports only FLA's chunk_gated_delta_rule " f"callable or a functools.partial of it; got {gated_delta_rule!r}." ) - chunk_size = kwargs.pop("chunk_size", 64) - if chunk_size != 64: + chunk_size = kwargs.pop("chunk_size", self.linear_attention_config.chunk_size) + if chunk_size != self.linear_attention_config.chunk_size: raise ValueError("GDN fake quantization supports only chunk_size=64") + if self.linear_attention_config.backend == "matmul": + return matmul_gdn( + *args, + policy=self.linear_attention_config, + state_qdq=quantize_state, + state_format=self._linear_attn_state_format, + chunk_size=chunk_size, + prefill_lengths=self._linear_attention_prefill_lengths, + **kwargs, + ) return _state_qdq_chunk_gated_delta_rule()( *args, chunk_size=chunk_size, - state_qdq=int(quantize_state), - state_qdq_block_v=64, + state_qdq=(2 if self._linear_attn_state_format == "int8" else 1) + if quantize_state + else 0, + state_qdq_block_v=self.gdn_state_qdq_block_v, w_quantizer=self.gdn_w_quantizer if quantize_w else None, **kwargs, ) diff --git a/modelopt/torch/quantization/plugins/kda.py b/modelopt/torch/quantization/plugins/kda.py new file mode 100644 index 00000000000..fbc4b1cf0ef --- /dev/null +++ b/modelopt/torch/quantization/plugins/kda.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""QAT integration for the optional FLA KimiDeltaAttention layer.""" + +import importlib.metadata +import inspect +from functools import lru_cache +from types import FunctionType + +from ..linear_attention.kda import matmul_kda +from ..nn import QuantModuleRegistry +from .custom import CUSTOM_MODEL_PLUGINS +from .linear_attention import _LinearAttentionQuantMixin + +__all__ = [] + + +@lru_cache(maxsize=8) +def _forward_signature(function): + return inspect.signature(function) + + +class _QuantKimiDeltaAttention(_LinearAttentionQuantMixin): + linear_attention_quantizer_names = ("kda_state_quantizer", "kda_w_quantizer") + + def validate_linear_attention(self): + super().validate_linear_attention() + if self.linear_attention_is_enabled and self.linear_attention_config.backend != "matmul": + raise ValueError("KDA numerical emulation requires backend='matmul'") + + def forward(self, *args, **kwargs): + self.validate_linear_attention() + if not self.linear_attention_is_enabled: + return super().forward(*args, **kwargs) + original = super().forward.__func__ + if self.linear_attention_config.decode is not None: + arguments = _forward_signature(original).bind(self, *args, **kwargs).arguments + if arguments.get("use_cache", False) or arguments.get("past_key_values") is not None: + raise NotImplementedError( + "Decode-aware QAT requires use_cache=False; use explicit numerical carry for continuation" + ) + # Bind this invocation so copied modules cannot retain another module's quantizers. + namespace = { + **original.__globals__, + "chunk_kda": self._quantized_chunk, + "fused_recurrent_kda": self._unsupported_recurrent, + } + forward = FunctionType( + original.__code__, + namespace, + original.__name__, + original.__defaults__, + original.__closure__, + ) + forward.__kwdefaults__ = original.__kwdefaults__ + return forward(self, *args, **kwargs) + + def _quantized_chunk(self, *args, **kwargs): + return matmul_kda( + *args, + policy=self.linear_attention_config, + state_qdq=self.kda_state_quantizer.is_enabled and self.kda_state_quantizer._if_quant, + state_format=self._linear_attn_state_format, + prefill_lengths=self._linear_attention_prefill_lengths, + **kwargs, + ) + + def _unsupported_recurrent(self, *args, **kwargs): + if self.linear_attention_config.decode is not None: + return self._quantized_chunk(*args, **kwargs) + raise NotImplementedError( + "KDA prefill QAT requires the chunk path; recurrent decode is not yet supported" + ) + + +def _register_fla_kda(model): + for module in model.modules(): + cls = type(module) + if cls.__module__ != "fla.layers.kda" or cls.__name__ != "KimiDeltaAttention": + continue + if cls in QuantModuleRegistry: + continue + if importlib.metadata.version("flash-linear-attention") != "0.5.1": + raise RuntimeError("KDA QAT is qualified with flash-linear-attention==0.5.1") + if importlib.metadata.version("fla-core") != "0.5.1": + raise RuntimeError("KDA QAT requires fla-core==0.5.1") + QuantModuleRegistry.register({cls: "fla_KimiDeltaAttention"})(_QuantKimiDeltaAttention) + + +CUSTOM_MODEL_PLUGINS.add(_register_fla_kda) diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py new file mode 100644 index 00000000000..5559c9f6606 --- /dev/null +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared module policy and checkpoint support for linear-attention QAT.""" + +from ..config import QuantizerAttributeConfig +from ..linear_attention.config import LinearAttentionConfig +from ..linear_attention.utils import validate_gdn_quantizer +from ..nn import QuantModule, TensorQuantizer + +__all__ = [] + + +class _LinearAttentionQuantMixin(QuantModule): + linear_attention_quantizer_names = ("gdn_state_quantizer", "gdn_w_quantizer") + + def _setup(self): + for name in self.linear_attention_quantizer_names: + self._register_temp_attribute( + name, TensorQuantizer(QuantizerAttributeConfig(enable=False)) + ) + self._register_temp_attribute("linear_attention_config", LinearAttentionConfig()) + self._register_temp_attribute("_linear_attention_prefill_lengths", None) + + @property + def _linear_attn_state(self): + return getattr(self, self.linear_attention_quantizer_names[0]) + + @property + def _linear_attn_state_format(self): + return "int8" if self._linear_attn_state.num_bits == 8 else "fp8_e4m3" + + @property + def _linear_attn_w(self): + return getattr(self, self.linear_attention_quantizer_names[1]) + + @property + def linear_attention_is_enabled(self): + """Whether an operand, state, or arithmetic policy changes the computation.""" + return ( + self._linear_attn_state.is_enabled + or self._linear_attn_w.is_enabled + or self.linear_attention_config.decode is not None + ) + + def validate_linear_attention(self): + """Validate quantizer contracts shared by GDN and KDA.""" + if self._linear_attn_state.is_enabled: + validate_gdn_quantizer( + self._linear_attn_state, + name=self.linear_attention_quantizer_names[0], + num_bits=((4, 3), 8), + ) + if self._linear_attn_state.axis != (0, 1): + raise ValueError( + f"{self.linear_attention_quantizer_names[0]} supports only axis=(0, 1) " + "with state.block_v tiling" + ) + decode = self.linear_attention_config.decode + if decode is not None and decode.state_codec == "int8_hadamard32": + if self._linear_attn_state_format != "int8": + raise ValueError("int8_hadamard32 requires INT8 state quantization") + if self._linear_attn_w.is_enabled: + validate_gdn_quantizer( + self._linear_attn_w, name=self.linear_attention_quantizer_names[1] + ) + if self.linear_attention_config.decode is not None: + raise ValueError("Decode's exact prefix does not support WY operand QDQ") + + def modelopt_post_restore(self, prefix=""): + """Validate the restored numerical policy and quantizers.""" + super().modelopt_post_restore(prefix) + self.validate_linear_attention() diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index ec942e0c24c..49195b44d29 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -61,7 +61,7 @@ from ..utils import sync_moe_expert_amax from ..utils.layerwise_calib import LayerActivationCollector from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear -from .gated_delta_net import GatedDeltaNetStateQuantMixin +from .gdn import GatedDeltaNetStateQuantMixin try: from megatron.core.extensions.transformer_engine import ( @@ -1122,9 +1122,7 @@ def forward_pre_attn_and_core_attn(self, *args, **kwargs): def validate_linear_attention(self): super().validate_linear_attention() - if self.config.context_parallel_size > 1 and ( - self.gdn_state_quantizer.is_enabled or self.gdn_w_quantizer.is_enabled - ): + if self.config.context_parallel_size > 1 and self.linear_attention_is_enabled: raise NotImplementedError("GDN QAT does not support Megatron context parallelism.") diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 20e083491aa..67513218137 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -50,7 +50,8 @@ def _fp8_eager(x, amax=None): epsilon = 1.0 / (1 << 24) zero_amax_mask = amax <= epsilon safe_amax = torch.where(zero_amax_mask, torch.ones_like(amax), amax) - scale = 448.0 / safe_amax + # Use direct division, matching the CUDA extension rather than reciprocal multiplication. + scale = torch.div(448.0, safe_amax) scale_inv = 1 / scale x = (x.to(torch.float32) * scale).clamp(min=-448.0, max=448.0) x = x.to(torch.float8_e4m3fn) diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index 79bc6065c24..ea115b6fe88 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -33,4 +33,5 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `mixer_mlp_nvfp4.yaml` | NVFP4 W4A4 on dense `*.mixer.{up,down}_proj` weight/input quantizers | | `attention_qkv_fp8.yaml` | FP8 E4M3 on attention q/k/v bmm and softmax quantizers | | `gdn_state_fp8_dynamic.yaml` | FP8 E4M3 dynamic fake quantization of the GatedDeltaNet recurrent state (per sequence, head and 64-column tile) at every kernel chunk boundary; needs `fla-core==0.5.1`, Triton, and SM89+ | +| `gdn_state_int8_dynamic.yaml` | INT8 quantizer unit; use the [complete recipe](../../../general/ptq/gdn_state_int8_dynamic.yaml) to enable INT8 state QDQ and Hadamard together | | `gdn_w_fp8_dynamic.yaml` | FP8 E4M3 dynamic (per token and head) fake quantization of the WY tensor `w` that multiplies the GatedDeltaNet state; requires identity STE, `fla-core==0.5.1`, and Triton | diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index cac4f6d80aa..e409bff7aa6 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -21,6 +21,10 @@ enable: false - quantizer_name: '*gdn_w_quantizer' enable: false + - quantizer_name: '*kda_state_quantizer' + enable: false + - quantizer_name: '*kda_w_quantizer' + enable: false - quantizer_name: '*block_sparse_moe.gate*' enable: false - quantizer_name: '*linear_attn.conv1d*' diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml index ed5768ee599..4d4b13278ba 100644 --- a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -14,10 +14,11 @@ # limitations under the License. # Dynamic per-tile FP8 E4M3 fake quantization of the GatedDeltaNet recurrent state. -# Recompute amax over each [K, 64] tile inside the kernel, independently per sequence and head +# Recompute amax over each [K, BV] tile inside the kernel, independently per sequence and head # of the [N, H, K, V] state (no calibration). ``axis: [0, 1]`` specifies sequence/head grouping; -# the kernel uses one scale per full-K by 64-value-column tile (two per head for V = 128). -# This is the only configuration the module supports. Applied at every kernel chunk boundary +# BV is ``gdn_state_qdq_block_v`` (default: the kernel's 64-column tile, +# two scales per head for V = 128; set 128 on the quantized module for one scale per head). This +# selects the FP8 state format. Applied at every kernel chunk boundary # (64 tokens) during calibration, QAT and QAD; needs fla-core==0.5.1 and Triton. # See ``modelopt.torch.quantization.plugins.gated_delta_net``. diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml new file mode 100644 index 00000000000..e86390cfef5 --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/gdn_state_int8_dynamic.yaml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Quantizer-only INT8 unit; general/ptq/gdn_state_int8_dynamic pairs it with Hadamard. +# The execution policy selects state scale groups and chunk/token/replay boundaries. +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig + - quantizer_name: '*gdn_state_quantizer' + cfg: + num_bits: 8 + unsigned: false + narrow_range: true + axis: [0, 1] + type: dynamic + pass_through_bwd: true diff --git a/modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml b/modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml new file mode 100644 index 00000000000..e960282121b --- /dev/null +++ b/modelopt_recipes/general/ptq/gdn_state_int8_dynamic.yaml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + gdn_state_int8_dynamic: configs/ptq/units/gdn_state_int8_dynamic + +metadata: + recipe_type: ptq + description: >- + Dynamic INT8 GDN state QDQ with 32-value Hadamard rotation during decode. + Prefill state remains unquantized; supply per-sequence prefix lengths through + linear_attention_training_phase. No calibration is required. +quantize: + algorithm: + quant_cfg: + - $import: base_disable_all + - $import: gdn_state_int8_dynamic + linear_attention: + - module_name: '*' + cfg: + backend: matmul + decode: + mode: token + state_codec: int8_hadamard32 + prefill_state_qdq: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 2ad373ffcab..53da4aac5fc 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 26 general/ptq/ recipes (click to expand) +All 27 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -58,6 +58,7 @@ supported combinations. | `int4_blockwise_weight_only` | INT4 W4A16, block 128, weights only | none | max | | `nvfp4_mlp_weight_only` | NVFP4 W4A16 (block 32), MLP + MoE weights only | none | max | | `mxfp4_mlp_weight_only` | MXFP4 W4A16, MLP + MoE weights only | none | none (no calibration) | +| `gdn_state_int8_dynamic` | GDN decode state INT8 + Hadamard; weights unchanged | none | none (dynamic scales; requires a prefix/decode phase context) |
diff --git a/tests/_test_utils/torch/quantization/linear_attention_reference.py b/tests/_test_utils/torch/quantization/linear_attention_reference.py index 65ea183a318..57d7448bcb1 100644 --- a/tests/_test_utils/torch/quantization/linear_attention_reference.py +++ b/tests/_test_utils/torch/quantization/linear_attention_reference.py @@ -25,23 +25,43 @@ import torch -__all__ = ["chunk_gdn_reference", "recurrent_delta_rule_reference", "state_fp8_qdq_reference"] +__all__ = [ + "chunk_gdn_reference", + "chunk_kda_reference", + "recurrent_delta_rule_reference", + "state_fp8_qdq_reference", + "state_qdq_reference", +] def state_fp8_qdq_reference(state: torch.Tensor, block_v: int = 64) -> torch.Tensor: """Dynamic E4M3 QDQ with identity STE and one scale per ``[Dk, block_v]`` tile. ``state`` is in key-first layout ``[..., Dk, Dv]``. Scales and rounding do not - contribute derivatives. The zero tile uses scale one; partial value tiles are valid. + contribute derivatives. FP8 uses TensorQuantizer small-range handling; partial tiles are valid. """ + return state_qdq_reference(state, block_v) + + +def state_qdq_reference(state: torch.Tensor, block_v: int = 64, state_format: str = "fp8_e4m3"): + """Dynamic state-tile QDQ with detached scales and identity STE.""" + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") if block_v not in (16, 32, 64, 128): raise ValueError("block_v must be 16, 32, 64, or 128") with torch.no_grad(): rounded = [] for tile in state.float().split(block_v, dim=-1): amax = tile.abs().amax(dim=(-2, -1), keepdim=True) - scale = torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)) - rounded.append((tile / scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() * scale) + if state_format == "int8": + scale = torch.where(amax > 0, amax / 127.0, torch.ones_like(amax)) + codes = (tile / scale).clamp(-127, 127).round() + else: + safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) + quant_scale = torch.div(448.0, safe_amax) + scale = quant_scale.reciprocal() + codes = (tile * quant_scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() + rounded.append(codes * scale) quantized = torch.cat(rounded, dim=-1).to(state.dtype) return state + (quantized - state).detach() @@ -140,34 +160,61 @@ def chunk_gdn_reference( state_v_first: bool = False, state_qdq: bool = False, state_qdq_block_v: int = 64, + state_format: str = "fp8_e4m3", w_quantizer: Callable[[torch.Tensor], torch.Tensor] | None = None, + matmul: Callable[[str, torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + arithmetic: Callable[[str, torch.Tensor], torch.Tensor] | None = None, + inverse_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Exact GDN chunk algebra with optional state/W fake quantization. The solve is a unit-lower triangular solve. ``w_quantizer`` sees the complete materialized ``[B,T,Hv,Dk]`` WY operand once, with its own autograd semantics. State QDQ occurs on the initial state and each chunk's final state, after readout. + Optional callbacks expose the eight named matmuls and elementwise boundaries. + Matmul callbacks receive conventional ``lhs @ rhs`` shapes, without padding. """ if g.ndim != 3 or chunk_size <= 0: raise ValueError("chunk_gdn_reference requires scalar GDN gates and positive chunk_size") q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) scale = q.shape[-1] ** -0.5 if scale is None else scale + mm = matmul if matmul is not None else lambda name, lhs, rhs: lhs @ rhs + arithmetic = arithmetic if arithmetic is not None else lambda name, x: x chunks, all_w = [], [] for n, (b, start, end) in enumerate(sequences): for lo in range(start, end, chunk_size): hi = min(lo + chunk_size, end) qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) - gc = g[b, lo:hi].transpose(0, 1).cumsum(-1) + gc = arithmetic("gate_prefix", g[b, lo:hi].transpose(0, 1).cumsum(-1)) bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) # Mask before exp: upper-triangle positive differences can overflow for long decay. causal = torch.ones(hi - lo, hi - lo, device=q.device, dtype=torch.bool).tril() - decay = (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() - gram = kc @ kc.transpose(-1, -2) - lower = (bc * gram * decay).tril(-1) + decay = arithmetic( + "gate_exp", (gc.unsqueeze(-1) - gc.unsqueeze(-2)).masked_fill(~causal, 0).exp() + ) + if matmul is None: + lower = (bc * (kc @ kc.transpose(-1, -2)) * decay).tril(-1) + else: + lower = (mm("key_interaction", bc * kc, kc.transpose(-1, -2)) * decay).tril(-1) matrix = lower + torch.eye(hi - lo, device=q.device, dtype=q.dtype) - rhs = torch.cat((bc * vc, bc * kc * gc.exp().unsqueeze(-1)), dim=-1) - solved = torch.linalg.solve_triangular(matrix, rhs, upper=False, unitriangular=True) - u, w = solved.split((v.shape[-1], k.shape[-1]), dim=-1) + gate = arithmetic("gate_exp", gc.exp()).unsqueeze(-1) + if matmul is None and inverse_fn is None: + rhs = torch.cat((bc * vc, bc * kc * gate), dim=-1) + solved = torch.linalg.solve_triangular(matrix, rhs, upper=False, unitriangular=True) + u, w = solved.split((v.shape[-1], k.shape[-1]), dim=-1) + else: + inverse = ( + inverse_fn(lower) + if inverse_fn is not None + else torch.linalg.solve_triangular( + matrix, + torch.eye(hi - lo, device=q.device, dtype=q.dtype).expand_as(matrix), + upper=False, + unitriangular=True, + ) + ) + u = mm("wy_value", inverse, bc * vc) + w = mm("wy_key", inverse, bc * kc * gate) chunks.append((n, qc, kc, gc, decay, u)) all_w.append(w.transpose(0, 1)) w = torch.cat(all_w).reshape(q.shape) @@ -179,21 +226,27 @@ def chunk_gdn_reference( if n != previous_n: state = states[n] if state_qdq: - state = state_fp8_qdq_reference(state, state_qdq_block_v) + state = state_qdq_reference(state, state_qdq_block_v, state_format) previous_n = n length = qc.shape[1] wc = w[offset : offset + length].transpose(0, 1) offset += length - updated_values = u - wc @ state # state_read - local_scores = ((qc * scale) @ kc.transpose(-1, -2) * decay).tril() - output = (qc * (scale * gc.exp()).unsqueeze(-1)) @ state - output = output + local_scores @ updated_values # local_readout + updated_values = arithmetic("value_residual", u - mm("state_read", wc, state)) + local_scores = (mm("output_score", qc * scale, kc.transpose(-1, -2)) * decay).tril() + output = mm( + "output_state", qc * (scale * arithmetic("gate_exp", gc.exp())).unsqueeze(-1), state + ) + output = arithmetic("output_add", output + mm("output_value", local_scores, updated_values)) outputs.append(output.transpose(0, 1)) - weighted_keys = kc * (gc[..., -1:] - gc).exp().unsqueeze(-1) - state = state * gc[..., -1].exp()[:, None, None] - state = state + weighted_keys.transpose(-1, -2) @ updated_values # state_update + weighted_keys = kc * arithmetic("gate_exp", (gc[..., -1:] - gc).exp()).unsqueeze(-1) + state = arithmetic( + "state_decay", state * arithmetic("gate_exp", gc[..., -1].exp())[:, None, None] + ) + state = arithmetic( + "state_add", state + mm("state_update", weighted_keys.transpose(-1, -2), updated_values) + ) if state_qdq: - state = state_fp8_qdq_reference(state, state_qdq_block_v) + state = state_qdq_reference(state, state_qdq_block_v, state_format) if len(finals) <= n: finals.append(state) else: @@ -201,3 +254,93 @@ def chunk_gdn_reference( output = torch.cat(outputs).reshape(*v.shape) final = torch.stack(finals) return output, final.transpose(-1, -2) if state_v_first else final + + +def chunk_kda_reference( + q, + k, + v, + g, + beta, + *, + chunk_size=64, + scale=None, + initial_state=None, + cu_seqlens=None, + state_v_first=False, + state_qdq=False, + state_qdq_block_v=64, + state_format="fp8_e4m3", + w_quantizer=None, + matmul=None, + inverse_fn=None, +): + """KDA chunk oracle with causal per-channel decays and optional operand QDQ. + + The eight-site callback receives conventional three-dimensional matmul + operands. Interaction sites are invoked one query row at a time with decay + already applied to the right-hand keys. This avoids inverse-prefix factors. + """ + if g.ndim != 4 or chunk_size <= 0: + raise ValueError("chunk_kda_reference requires per-key gates and positive chunk_size") + q, k, states, sequences = _prepare(q, k, v, g, beta, initial_state, cu_seqlens, state_v_first) + scale = q.shape[-1] ** -0.5 if scale is None else scale + mm = matmul if matmul is not None else lambda name, lhs, rhs: lhs @ rhs + outputs, finals = [], [] + for n, (b, start, end) in enumerate(sequences): + state = states[n] + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + pieces = [] + for lo in range(start, end, chunk_size): + hi = min(lo + chunk_size, end) + qc, kc, vc = (x[b, lo:hi].transpose(0, 1) for x in (q, k, v)) + prefix = g[b, lo:hi].transpose(0, 1).cumsum(-2) + bc = beta[b, lo:hi].transpose(0, 1).unsqueeze(-1) + gate = prefix.exp() + lower_rows, score_rows = [], [] + for row in range(hi - lo): + decayed_keys = ( + kc[..., : row + 1, :] + * (prefix[..., row : row + 1, :] - prefix[..., : row + 1, :]).exp() + ) + right = decayed_keys.transpose(-1, -2) + left = kc[..., row : row + 1, :] * bc[..., row : row + 1, :] + interaction = mm("key_interaction", left, right) + score = mm("output_score", qc[..., row : row + 1, :] * scale, right) + # Zero padding follows the computation and never contributes a scale statistic. + padding = hi - lo - row - 1 + lower_rows.append(torch.nn.functional.pad(interaction, (0, padding))) + score_rows.append(torch.nn.functional.pad(score, (0, padding))) + lower = torch.cat(lower_rows, dim=-2).tril(-1) + scores = torch.cat(score_rows, dim=-2) + identity = torch.eye(hi - lo, device=q.device, dtype=q.dtype).expand_as(lower) + inverse = ( + inverse_fn(lower) + if inverse_fn is not None + else torch.linalg.solve_triangular( + identity + lower, identity, upper=False, unitriangular=True + ) + ) + u = mm("wy_value", inverse, bc * vc) + w = mm("wy_key", inverse, bc * kc * gate) + if w_quantizer is not None: + w = w_quantizer(w) + updated = u - mm("state_read", w, state) + pieces.append( + ( + mm("output_state", qc * scale * gate, state) + + mm("output_value", scores, updated) + ).transpose(0, 1) + ) + weighted_keys = kc * (prefix[..., -1:, :] - prefix).exp() + state = state * gate[..., -1, :, None] + mm( + "state_update", weighted_keys.transpose(-1, -2), updated + ) + if state_qdq: + state = state_qdq_reference(state, state_qdq_block_v, state_format) + outputs.append(torch.cat(pieces)) + finals.append(state) + output = torch.stack(outputs) if cu_seqlens is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) + return output, final.transpose(-1, -2) if state_v_first else final diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py index 9a11c454932..84b344ba224 100644 --- a/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py @@ -28,7 +28,7 @@ from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.nn import TensorQuantizer -from modelopt.torch.quantization.plugins.gated_delta_net import GatedDeltaNetStateQuantMixin +from modelopt.torch.quantization.plugins.gdn import GatedDeltaNetStateQuantMixin fla = pytest.importorskip("fla.ops.gated_delta_rule") from fla.utils import IS_NVIDIA_HOPPER, TRITON_ABOVE_3_4_0 @@ -38,7 +38,7 @@ ) DTYPES = (torch.float32, torch.bfloat16) -QDQ_MODES = ((False, False), (False, True), (True, False), (True, True)) +QDQ_MODES = ((False, False), (False, True), (True, False), (True, True), (2, False), (2, True)) STATE_LAYOUTS = ((False, False), (True, True)) W_AXES = (None, (0, 1), (0, 1, 2)) @@ -115,7 +115,7 @@ def _compile_gdn_kernels(): for packed, state_v_first in STATE_LAYOUTS: args, state = make_inputs(dtype, packed=packed, state_v_first=state_v_first) for state_qdq, quantize_w in QDQ_MODES: - if state_qdq and not supports_state_qdq: + if state_qdq == 1 and not supports_state_qdq: continue values_and_grads( chunk_gated_delta_rule, @@ -133,9 +133,15 @@ def _compile_gdn_kernels(): args, state = make_inputs() quantizer = w_quantizer() - values_and_grads( - chunk_gated_delta_rule, args, state, output_final_state=True, w_quantizer=quantizer - ) + for state_qdq in (0, 2): + values_and_grads( + chunk_gated_delta_rule, + args, + state, + output_final_state=True, + w_quantizer=quantizer, + state_qdq=state_qdq, + ) for axis in W_AXES: output, _ = chunk_gated_delta_rule(*args, w_quantizer=w_quantizer(axis)) torch.autograd.grad(output.square().sum(), args) @@ -154,16 +160,19 @@ def _compile_gdn_kernels(): allow_neg_eigval=True, ) torch.autograd.grad(output.square().sum() + final.square().sum(), (*args, state, a_log, bias)) - if supports_state_qdq: + for state_qdq in (1, 2): + if state_qdq == 1 and not supports_state_qdq: + continue for block_v in (16, 64, 128): values_and_grads( chunk_gated_delta_rule, args, state, output_final_state=True, - state_qdq=True, + state_qdq=state_qdq, state_qdq_block_v=block_v, ) + if supports_state_qdq: args, state = make_inputs(packed=True, state_v_first=True) for block_v in (16, 32, 64, 128): _, final = chunk_gated_delta_rule( @@ -193,7 +202,7 @@ def test_disabled_matches_upstream_forward_and_backward(dtype): @pytest.mark.parametrize(("packed", "state_v_first"), STATE_LAYOUTS) @pytest.mark.parametrize("dtype", DTYPES) def test_numerical_surrogate_reference(state_qdq, quantize_w, packed, state_v_first, dtype): - if state_qdq: + if state_qdq == 1: require_state_qdq() args, state = make_inputs(dtype, packed=packed, state_v_first=state_v_first) quantizer = w_quantizer() if quantize_w else None @@ -208,7 +217,13 @@ def test_numerical_surrogate_reference(state_qdq, quantize_w, packed, state_v_fi } reference_args = [x.detach().float().requires_grad_() for x in args] reference_state = state.detach().float().requires_grad_() - expected = values_and_grads(chunk_gdn_reference, reference_args, reference_state, **kwargs) + expected = values_and_grads( + chunk_gdn_reference, + reference_args, + reference_state, + state_format="int8" if state_qdq == 2 else "fp8_e4m3", + **kwargs, + ) actual = values_and_grads( chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs ) @@ -216,7 +231,8 @@ def test_numerical_surrogate_reference(state_qdq, quantize_w, packed, state_v_fi compare(actual[1], expected[1], 0.05 if dtype == torch.bfloat16 else 0.02) -def test_w_qdq_saved_once_and_activation_checkpoint_parity(): +@pytest.mark.parametrize("state_qdq", [0, 2]) +def test_w_qdq_saved_once_and_activation_checkpoint_parity(state_qdq): args, state = make_inputs() quantizer = w_quantizer() calls = [] @@ -224,7 +240,11 @@ def test_w_qdq_saved_once_and_activation_checkpoint_parity(): def fn(*x): return chunk_gated_delta_rule( - *x[:5], initial_state=x[5], output_final_state=True, w_quantizer=quantizer + *x[:5], + initial_state=x[5], + output_final_state=True, + w_quantizer=quantizer, + state_qdq=state_qdq, ) output = fn(*args, state) @@ -252,12 +272,20 @@ def test_zero_initial_state_and_output_only_training_loss(axis): @pytest.mark.parametrize("block_v", [16, 64, 128]) -def test_state_scale_tile_forward_and_backward(block_v): - require_state_qdq() +@pytest.mark.parametrize("state_qdq", [1, 2]) +def test_state_scale_tile_forward_and_backward(block_v, state_qdq): + if state_qdq == 1: + require_state_qdq() args, state = make_inputs() - kwargs = {"state_qdq": True, "state_qdq_block_v": block_v} + kwargs = {"state_qdq": state_qdq, "state_qdq_block_v": block_v} reference_args = [x.detach().float().requires_grad_() for x in args] - expected = values_and_grads(chunk_gdn_reference, reference_args, state, **kwargs) + expected = values_and_grads( + chunk_gdn_reference, + reference_args, + state, + state_format="int8" if state_qdq == 2 else "fp8_e4m3", + **kwargs, + ) actual = values_and_grads( chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs ) diff --git a/tests/gpu/torch/quantization/plugins/test_kda.py b/tests/gpu/torch/quantization/plugins/test_kda.py new file mode 100644 index 00000000000..ed67787de54 --- /dev/null +++ b/tests/gpu/torch/quantization/plugins/test_kda.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from contextlib import nullcontext + +import pytest +import torch +from torch.utils.checkpoint import checkpoint + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + linear_attention_training_phase, +) +from modelopt.torch.quantization.nn import TensorQuantizer + +KimiDeltaAttention = pytest.importorskip("fla.layers.kda").KimiDeltaAttention + + +def _layer(head_dim=16): + return ( + KimiDeltaAttention(hidden_size=32, head_dim=head_dim, num_heads=2, use_short_conv=True) + .cuda() + .train() + ) + + +def _forward(model, hidden): + policy = getattr(model, "linear_attention_config", None) + phase = ( + linear_attention_training_phase(model, [min(31, hidden.shape[1])]) + if policy is not None and policy.decode is not None + else nullcontext() + ) + with phase, torch.autocast("cuda", dtype=torch.bfloat16): + return model(hidden)[0] + + +@pytest.mark.parametrize( + "mode", + [ + "decode_token", + "decode_replay", + "decode_decay", + "decode_token_int8", + "decode_replay_int8", + "decode_replay_hadamard_int8", + ], +) +@pytest.mark.timeout(180) +def test_fla_layer_qat_restore_and_optimizer(tmp_path, mode): + torch.manual_seed(73) + head_dim = 32 if "hadamard" in mode else 16 + model = _layer(head_dim) + hidden = torch.randn(1, 73, 32, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + baseline = _forward(model, hidden) + attributes = {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)} + cfg = { + "quant_cfg": [{"quantizer_name": "*", "enable": False}], + "algorithm": None, + "linear_attention": [{"module_name": "*", "cfg": {"backend": "matmul"}}], + } + state_attributes = {**attributes, "axis": (0, 1)} + if mode.endswith("int8"): + state_attributes.update(num_bits=8, unsigned=False, narrow_range=True) + if mode.startswith("decode"): + cfg["quant_cfg"].append({"quantizer_name": "*kda_state_quantizer", "cfg": state_attributes}) + policy_cfg = cfg["linear_attention"][0]["cfg"] + policy_cfg["state"] = {"block_v": 16} + policy_cfg["decode"] = { + "mode": "replay" if mode.startswith("decode_replay") else "token", + "implementation": "torch", + } + if mode.startswith("decode_replay"): + policy_cfg["decode"]["replay"] = {"window": 5} + if mode == "decode_decay": + policy_cfg["decode"]["decay_log_step"] = 1 / 256 + if "hadamard" in mode: + policy_cfg["state"]["block_v"] = 32 + policy_cfg["decode"]["state_codec"] = "int8_hadamard32" + mtq.quantize(model, cfg) + output = _forward(model, hidden) + assert torch.isfinite(output).all() + enabled = [q for q in model.modules() if isinstance(q, TensorQuantizer) and q.is_enabled] + policy = model.linear_attention_config + for q in enabled: + q.disable() + model.linear_attention_config = LinearAttentionConfig() + with torch.no_grad(): + torch.testing.assert_close(_forward(model, hidden), baseline, rtol=0, atol=0) + for q in enabled: + q.enable() + model.linear_attention_config = policy + checkpoint = tmp_path / "model.pth" + mto.save(model, checkpoint) + restored = mto.restore(_layer(head_dim), checkpoint) + actual = _forward(restored, hidden) + torch.testing.assert_close(actual, output, rtol=0, atol=0) + assert restored.linear_attention_config == policy + actual.float().square().mean().backward() + assert restored.q_proj.weight.grad is not None + for parameter in restored.parameters(): + assert parameter.grad is not None + assert torch.isfinite(parameter.grad).all() + before = restored.q_proj.weight.detach().clone() + torch.optim.SGD(restored.parameters(), lr=0.1).step() + assert not torch.equal(before, restored.q_proj.weight) + restored.eval() + if mode.startswith("decode"): + with ( + linear_attention_training_phase(restored, [0]), + torch.autocast("cuda", dtype=torch.bfloat16), + ): + assert torch.isfinite(restored(hidden[:, :1])[0]).all() + with pytest.raises(NotImplementedError, match="use_cache=False"): + restored(hidden[:, :1], use_cache=True) + with ( + pytest.raises(ValueError, match="explicit"), + torch.autocast("cuda", dtype=torch.bfloat16), + ): + restored(hidden) + else: + with pytest.raises(NotImplementedError, match="chunk path"): + _forward(restored, hidden[:, :1]) + + +@pytest.mark.parametrize("hadamard", [False, True]) +def test_decode_phase_spans_activation_checkpoint_backward(hadamard): + model = _layer(32 if hadamard else 16) + quant_cfg = [{"quantizer_name": "*", "enable": False}] + if hadamard: + quant_cfg.append( + { + "quantizer_name": "*kda_state_quantizer", + "cfg": { + "num_bits": 8, + "type": "dynamic", + "axis": (0, 1), + "unsigned": False, + "narrow_range": True, + }, + } + ) + mtq.quantize( + model, + { + "quant_cfg": quant_cfg, + "algorithm": None, + "linear_attention": [ + { + "module_name": "*", + "cfg": { + "backend": "matmul", + "decode": { + "implementation": "torch", + "state_codec": "int8_hadamard32" if hadamard else "tile", + }, + }, + } + ], + }, + ) + hidden = torch.randn(1, 73, 32, device="cuda", requires_grad=True) + with linear_attention_training_phase(model, [31]): + with torch.autocast("cuda", dtype=torch.bfloat16): + actual = checkpoint(lambda x: model(x)[0], hidden, use_reentrant=False) + actual.square().mean().backward() + assert model._linear_attention_prefill_lengths is None + assert torch.isfinite(hidden.grad).all() + assert torch.count_nonzero(model.q_proj.weight.grad) > 0 + with linear_attention_training_phase(model, [0]): + with linear_attention_training_phase(model, [73]): + assert model._linear_attention_prefill_lengths == (73,) + assert model._linear_attention_prefill_lengths == (0,) + assert model._linear_attention_prefill_lengths is None diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index cf802abaf40..2039217a8e8 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -25,7 +25,7 @@ import modelopt.torch.kernels.quantization.gemm as triton_kernel import modelopt.torch.quantization.utils as quant_utils from modelopt.torch.quantization import tensor_quant -from modelopt.torch.quantization.extensions import get_cuda_ext, get_cuda_ext_mx +from modelopt.torch.quantization.extensions import get_cuda_ext, get_cuda_ext_fp8, get_cuda_ext_mx from modelopt.torch.quantization.tensor_quant import mx_format_map if triton_kernel.IS_AVAILABLE: @@ -120,6 +120,26 @@ def test_cuda_ext_tiny_amax(self): class TestScaledE4M3: + @pytest.mark.parametrize("per_channel", [False, True]) + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_eager_matches_cuda_scale_rounding(self, per_channel, dtype): + # These ranges distinguish direct division from reciprocal multiplication. + amax = torch.tensor([0.0, 1e-8, 2**-24, 1.04207686e-7, 0.91567558], device="cuda") + values = torch.linspace(-1, 1, 129, device="cuda")[None, :] * amax[:, None] + values = values.to(dtype) + extension = get_cuda_ext_fp8() + if per_channel: + actual = extension.fake_e4m3fy_with_axis(values, amax, 0) + expected = tensor_quant.fp8_eager(values, amax[:, None]) + else: + actual = torch.stack( + [extension.fake_e4m3fy(row, limit) for row, limit in zip(values, amax)] + ) + expected = torch.stack( + [tensor_quant.fp8_eager(row, limit) for row, limit in zip(values, amax)] + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + @pytest.mark.parametrize("device", ["cuda", "cpu"]) def test_e4m3_no_scale(self, device): x = torch.randn(4, 4, device=device, dtype=torch.float32) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index c46959c4d8d..94ad1aaedf1 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import nullcontext + import pytest import torch from _test_utils.torch.megatron.models import get_mcore_gpt_model @@ -23,11 +25,16 @@ ) import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + linear_attention_training_phase, +) +from modelopt.torch.quantization.nn import TensorQuantizer pytest.importorskip("fla") # Megatron-Core GatedDeltaNet and the state QDQ kernel need fla GatedDeltaNet = pytest.importorskip("megatron.core.ssm.gated_delta_net").GatedDeltaNet -from modelopt.torch.quantization.plugins.gated_delta_net import _state_qdq_chunk_gated_delta_rule +from modelopt.torch.quantization.plugins.gdn import _state_qdq_chunk_gated_delta_rule from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet try: @@ -78,15 +85,50 @@ def _gdn_config(sites): for site in sites ], "algorithm": None, + "linear_attention": [ + {"module_name": "decoder.layers.*.self_attention", "cfg": {"state": {"block_v": 16}}} + ], } +def _config_for_mode(mode): + sites = ("state", "w") if mode == "both" else ((mode,) if mode in ("state", "w") else ()) + if mode.startswith("decode"): + sites = ("state",) + if mode == "state_int8": + sites = ("state",) + cfg = _gdn_config(sites) + if mode.endswith("int8"): + cfg["quant_cfg"][1]["cfg"].update(num_bits=8, unsigned=False, narrow_range=True) + if mode.startswith("decode"): + policy = cfg["linear_attention"][0]["cfg"] + policy.update(backend="matmul", decode={"implementation": "triton"}) + if mode.startswith("decode_replay"): + policy["decode"].update(mode="replay", replay={"window": 8}) + return cfg, sites + + +def _gdn_forward(model): + original_forward = get_forward(model) + + def forward(m): + enabled = any( + getattr(getattr(layer, "linear_attention_config", None), "decode", None) is not None + for layer in m.modules() + ) + with linear_attention_training_phase(m, [64, 64]) if enabled else nullcontext(): + return original_forward(m) + + return forward + + def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): initialize_for_megatron( tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED ) model = _make_model(size) - forward = get_forward(model) + forward = _gdn_forward(model) + outputs = [] handles = [ module.register_forward_hook( @@ -100,13 +142,18 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): gdn_ref = outputs.copy() outputs.clear() - sites = ("state", "w") if mode == "both" else (mode,) - mtq.quantize(model, _gdn_config(sites)) + cfg, sites = _config_for_mode(mode) + mtq.quantize(model, cfg) gdn_modules = [m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)] assert gdn_modules, "no GatedDeltaNet layer was wrapped" for module in gdn_modules: + assert module.gdn_state_qdq_block_v == 16 assert module.gdn_state_quantizer.is_enabled == ("state" in sites) assert module.gdn_w_quantizer.is_enabled == ("w" in sites) + # Checkpointing must save the resolved policy, including edits after conversion. + module.linear_attention_config = LinearAttentionConfig( + **{**module.linear_attention_config.model_dump(), "state": {"block_v": 32}} + ) with torch.no_grad(): loss_quant = forward(model) @@ -119,20 +166,34 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): for handle in handles: handle.remove() - for site in sites: - mtq.disable_quantizer(model, f"*gdn_{site}_quantizer") + enabled = [ + q + for m in gdn_modules + for q in m.modules() + if isinstance(q, TensorQuantizer) and q.is_enabled + ] + policies = [m.linear_attention_config for m in gdn_modules] + for q in enabled: + q.disable() + for module in gdn_modules: + module.linear_attention_config = LinearAttentionConfig() with torch.no_grad(): torch.testing.assert_close(forward(model), loss_ref, rtol=1e-4, atol=1e-4) - for site in sites: - mtq.enable_quantizer(model, f"*gdn_{site}_quantizer") + for q in enabled: + q.enable() + for module, policy in zip(gdn_modules, policies): + module.linear_attention_config = policy restored = _make_model(size) sharded_state_dict_test_helper(checkpoint_path, model, restored, forward) restored_gdn = [m for m in restored.modules() if isinstance(m, _QuantGatedDeltaNet)] assert len(restored_gdn) == len(gdn_modules) - for module in restored_gdn: + for module, original in zip(restored_gdn, gdn_modules): + assert module.linear_attention_config == original.linear_attention_config assert module.gdn_state_quantizer.is_enabled == ("state" in sites) assert module.gdn_w_quantizer.is_enabled == ("w" in sites) + assert module._linear_attn_state_format == original._linear_attn_state_format + assert module.gdn_state_qdq_block_v == 32 assert module.in_proj.weight.grad is not None assert torch.isfinite(module.in_proj.weight.grad).all() @@ -148,11 +209,13 @@ def _compile_gdn_qat_kernels(rank, size, mode): tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED ) model = _make_model(size) - forward = get_forward(model) + forward = _gdn_forward(model) with torch.no_grad(): forward(model) - sites = ("state", "w") if mode == "both" else (mode,) - mtq.quantize(model, _gdn_config(sites)) + cfg, _ = _config_for_mode(mode) + # The functional case switches state tiles before running inference. + cfg["linear_attention"][0]["cfg"]["state"]["block_v"] = 32 + mtq.quantize(model, cfg) with torch.no_grad(): forward(model) model.train() @@ -163,7 +226,7 @@ def _compile_gdn_qat_kernels(rank, size, mode): @pytest.fixture def compiled_gdn_workers(request, tp_size, mode): """Warm the selected QAT path in the same workers, outside the test-call budget.""" - if mode != "w" and torch.cuda.get_device_capability() < (8, 9): + if mode in ("state", "both") and torch.cuda.get_device_capability() < (8, 9): pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") workers = request.getfixturevalue(f"dist_workers_size_{tp_size}") workers.run(_compile_gdn_qat_kernels, mode) @@ -171,7 +234,19 @@ def compiled_gdn_workers(request, tp_size, mode): @pytest.mark.parametrize("tp_size", [1, 2]) -@pytest.mark.parametrize("mode", ["state", "w", "both"]) +@pytest.mark.parametrize( + "mode", + [ + "state", + "w", + "both", + "decode_token", + "decode_replay", + "state_int8", + "decode_token_int8", + "decode_replay_int8", + ], +) def test_gdn_qat_and_sharded_restore(compiled_gdn_workers, tmp_path, mode): """Train through state/W QDQ after a Megatron distributed-checkpoint round trip.""" compiled_gdn_workers.run(_test_gdn_qat_helper, mode, tmp_path) @@ -181,10 +256,21 @@ def _test_gdn_context_parallel_helper(rank, size, mode): initialize_for_megatron(context_parallel_size=size, seed=SEED) model = _make_model(1, context_parallel_size=size) with pytest.raises(NotImplementedError, match="context parallelism"): - mtq.quantize(model, _gdn_config((mode,))) + mtq.quantize(model, _config_for_mode(mode)[0]) -@pytest.mark.parametrize("mode", ["state", "w"]) +@pytest.mark.parametrize( + "mode", + [ + "state", + "w", + "decode_token", + "decode_replay", + "state_int8", + "decode_token_int8", + "decode_replay_int8", + ], +) def test_gdn_context_parallel_rejected(dist_workers_size_2, mode): """Reject unqualified Megatron CP even when it does not pass an FLA CP context.""" dist_workers_size_2.run(_test_gdn_context_parallel_helper, mode) diff --git a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py b/tests/unit/torch/quantization/plugins/test_gdn.py similarity index 74% rename from tests/unit/torch/quantization/plugins/test_gated_delta_net.py rename to tests/unit/torch/quantization/plugins/test_gdn.py index 3793edcb189..aed7f8caf42 100644 --- a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -22,9 +22,19 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.nn import QuantModuleRegistry -from modelopt.torch.quantization.plugins import gated_delta_net -from modelopt.torch.quantization.plugins.gated_delta_net import GatedDeltaNetStateQuantMixin +from modelopt.torch.quantization.plugins import gdn +from modelopt.torch.quantization.plugins.gdn import GatedDeltaNetStateQuantMixin + +GDN_STATE_INT8_DYNAMIC = { + "num_bits": 8, + "unsigned": False, + "narrow_range": True, + "type": "dynamic", + "axis": (0, 1), +} GDN_STATE_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"} @@ -36,9 +46,7 @@ def chunk_gated_delta_rule(q, k, v, g, beta, **kwargs): @pytest.fixture(autouse=True) def mock_fla_kernel(monkeypatch): - monkeypatch.setattr( - gated_delta_net, "_fla_chunk_gated_delta_rule", lambda: chunk_gated_delta_rule - ) + monkeypatch.setattr(gdn, "_fla_chunk_gated_delta_rule", lambda: chunk_gated_delta_rule) class TinyGatedDeltaNet(nn.Module): @@ -106,7 +114,12 @@ def test_dynamic_export_removes_linear_attention_attributes(): model = _QuantTinyGatedDeltaNet.convert(TinyGatedDeltaNet()) model.export() assert type(model) is TinyGatedDeltaNet - for name in ("gdn_state_quantizer", "gdn_w_quantizer"): + for name in ( + "gdn_state_quantizer", + "gdn_w_quantizer", + "linear_attention_config", + "_linear_attention_prefill_lengths", + ): assert not hasattr(model, name) @@ -124,27 +137,29 @@ def test_disabled_state_quantizer_calls_original_kernel(): assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) @pytest.mark.parametrize("partial_kernel", [False, True]) -def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, partial_kernel): +def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, state_format, partial_kernel): calls = [] def fake_state_qdq_kernel(*args, **kwargs): calls.append(kwargs) return chunk_gated_delta_rule(*args) - monkeypatch.setattr( - gated_delta_net, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel - ) + monkeypatch.setattr(gdn, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel) model = TinyGatedDeltaNet() if partial_kernel: model.gated_delta_rule = partial(chunk_gated_delta_rule, output_final_state=True) x = torch.randn(2, 8, 3, 4) - mtq.quantize(model, quant_cfg(), lambda m: m(x)) + cfg = quant_cfg() + if state_format == "int8": + cfg["quant_cfg"][1]["cfg"] = GDN_STATE_INT8_DYNAMIC + mtq.quantize(model, cfg, lambda m: m(x)) model(x) assert calls and calls[-1] == { "chunk_size": 64, - "state_qdq": 1, + "state_qdq": 2 if state_format == "int8" else 1, "state_qdq_block_v": 64, "w_quantizer": None, **({"output_final_state": True} if partial_kernel else {}), @@ -167,9 +182,7 @@ def fake_state_qdq_kernel(*args, **kwargs): calls.append(kwargs) return chunk_gated_delta_rule(*args) - monkeypatch.setattr( - gated_delta_net, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel - ) + monkeypatch.setattr(gdn, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel) model = TinyGatedDeltaNet() x = torch.randn(2, 8, 3, 4) mtq.quantize(model, quant_cfg(state=state, w=True), lambda m: m(x)) @@ -214,7 +227,13 @@ def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): model = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) cfg = quant_cfg(state=state, w=w) cfg["algorithm"] = None + cfg["linear_attention"] = [ + {"module_name": "*", "cfg": {"state": {"block_v": 128}}}, + {"module_name": "0", "cfg": {"state": {"block_v": 32}}}, + ] mtq.quantize(model, cfg) + assert model[0].gdn_state_qdq_block_v == 32 + model[0].linear_attention_config.state.block_v = 16 if w: # Save the actual quantizer settings, including edits after conversion. model[0].gdn_w_quantizer.axis = None @@ -222,6 +241,8 @@ def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): mto.save(model, path) restored = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) mto.restore(restored, path) + assert restored[0].linear_attention_config == model[0].linear_attention_config + assert restored[0].gdn_state_qdq_block_v == 16 for name, enabled in (("gdn_state_quantizer", state), ("gdn_w_quantizer", w)): original = getattr(model[0], name) quantizer = getattr(restored[0], name) @@ -241,7 +262,12 @@ def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): cfg = quant_cfg(state=False, w=True) cfg["algorithm"] = None + cfg["linear_attention"] = [{"module_name": "", "cfg": {"state": {"block_v": 32}}}] + mtq.quantize(model, cfg) + assert model.gdn_state_qdq_block_v == 32 + cfg["linear_attention"].append({"module_name": "", "cfg": {}}) mtq.quantize(model, cfg) + assert model.gdn_state_qdq_block_v == 64 assert not model.gdn_state_quantizer.is_enabled assert model.gdn_w_quantizer.is_enabled @@ -256,7 +282,9 @@ def test_restore_legacy_gdn_without_new_quantizer_handles(): model = mtq.quantize(TinyGatedDeltaNet(), config) state = mto.modelopt_state(model) for _, mode_state in state["modelopt_state_dict"]: + mode_state["config"].pop("linear_attention", None) metadata = mode_state["metadata"] + metadata.pop("linear_attention", None) for name in ("gdn_state_quantizer", "gdn_w_quantizer"): metadata["quantizer_state"].pop(name, None) restored = TinyGatedDeltaNet() @@ -274,3 +302,37 @@ def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): ) assert not model.gdn_state_quantizer.is_enabled assert not model.gdn_w_quantizer.is_enabled + + +def test_int8_state_conversion_and_checkpoint(tmp_path): + cfg = load_recipe("general/ptq/gdn_state_int8_dynamic").quantize + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + assert model.gdn_state_quantizer.is_enabled + assert model._linear_attn_state_format == "int8" + assert model.linear_attention_config.backend == "matmul" + assert model.linear_attention_config.decode.state_codec == "int8_hadamard32" + assert not model.linear_attention_config.decode.prefill_state_qdq + assert not model.proj.weight_quantizer.is_enabled + path = tmp_path / "int8.pt" + mto.save(model, path) + restored = mto.restore(TinyGatedDeltaNet(), path) + assert restored.linear_attention_config == model.linear_attention_config + assert restored.gdn_state_quantizer.is_enabled + assert restored._linear_attn_state_format == "int8" + assert restored.gdn_state_quantizer.narrow_range + assert not restored.gdn_state_quantizer.unsigned + assert not restored.gdn_w_quantizer.is_enabled + + +def test_policy_rejects_unmatched_and_unimplemented_modes(): + with pytest.raises(ValueError, match="matches no supported"): + mtq.quantize( + nn.Linear(4, 4), {"linear_attention": [{"module_name": "*"}], "algorithm": None} + ) + for policy in ( + {"chunk_size": 32}, + {"solve": {"method": "neumann"}}, + {"state": {"mode": "token"}}, + ): + with pytest.raises(ValueError): + QuantizeConfig(linear_attention=[{"module_name": "*", "cfg": policy}]) diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py new file mode 100644 index 00000000000..e0ced168024 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import itertools + +import pytest +import torch +import torch.nn.functional as F +from _test_utils.torch.quantization.linear_attention_reference import recurrent_delta_rule_reference + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + matmul_gdn, + matmul_kda, +) +from modelopt.torch.quantization.linear_attention.config import LinearAttentionDecodeConfig +from modelopt.torch.quantization.linear_attention.decode import _encode, recurrent_decode_reference +from modelopt.torch.quantization.linear_attention.utils import _state_qdq +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.mark.parametrize("state", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_fp8_qdq_matches_tensor_quantizer(state, dtype): + torch.manual_seed(762) + amplitudes = torch.tensor([0.0, 1e-8, 2**-24, 2**-23, 1.0, 731.173], dtype=dtype) + value = torch.rand((2, 6, 5, 19) if state else (2, 6, 19), dtype=dtype) * 2 - 1 + value *= amplitudes.reshape(1, 6, *([1] * (value.ndim - 2))) + value.requires_grad_() + quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=(4, 3), type="dynamic", axis=(0, 1), pass_through_bwd=True + ) + ) + expected = torch.cat( + [ + quantizer(tile.flatten(-2) if state else tile).reshape_as(tile) + for tile in value.split(16, -1) + ], + -1, + ) + encoded = _encode(value, True, 16, state=state) + results = [encoded.values] + if state: + results.append(_state_qdq(value, 16)) + probe = torch.randn_like(value) + (expected_gradient,) = torch.autograd.grad((expected * probe).sum(), value) + for result in results: + torch.testing.assert_close(result, expected, rtol=0, atol=0) + (gradient,) = torch.autograd.grad((result * probe).sum(), value) + torch.testing.assert_close(gradient, expected_gradient, rtol=0, atol=0) + assert encoded.scales.shape == (2, 6, 2) + assert not encoded.scales.requires_grad + + +def _inputs(kda=True, length=73): + torch.manual_seed(193) + q, k = [ + F.normalize(torch.randn(length, 2, 16, dtype=torch.float64), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(length, 2, 11, dtype=torch.float64, requires_grad=True) + g = ( + -torch.rand((length, 2, 16) if kda else (length, 2), dtype=torch.float64) * 0.03 + ).requires_grad_() + beta = (torch.rand(length, 2, dtype=torch.float64) * 0.4).requires_grad_() + state = (torch.randn(2, 16, 11, dtype=torch.float64) * 0.1).requires_grad_() + return (q, k, v, g, beta), state + + +def _values_and_grads(output, state, args, initial): + loss = output.square().sum() + state.square().sum() + return output, state, *torch.autograd.grad(loss, (*args, initial), retain_graph=True) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("mode", ["token", "replay"]) +def test_unquantized_decode_recurrence_and_gradients(kda, mode): + args, state = _inputs(kda) + cfg = LinearAttentionDecodeConfig( + mode=mode, replay={"window": 8, "factor_qdq": False} if mode == "replay" else None + ) + output, carry = recurrent_decode_reference(*args, config=cfg, initial_state=state) + exact, final = recurrent_delta_rule_reference( + *(x.unsqueeze(0) for x in args), initial_state=state.unsqueeze(0) + ) + for a, b in zip( + _values_and_grads(output, carry.reconstruct(), args, state), + _values_and_grads(exact[0], final[0], args, state), + ): + torch.testing.assert_close(a, b, rtol=2e-11, atol=1e-12) + assert carry.position == 73 + assert carry.cursor == (1 if mode == "replay" else 0) + + +@pytest.mark.parametrize("mode", ["token", "replay", "reencode"]) +@pytest.mark.parametrize("readout", ["working", "stored"]) +def test_quantized_carry_continuation_metadata_and_gradients(mode, readout): + args, state = _inputs() + cfg = LinearAttentionDecodeConfig( + mode="token" if mode == "token" else "replay", + readout=readout, + replay={"window": 8, "encoding": "reencode" if mode == "reencode" else "once"} + if mode != "token" + else None, + ) + expected, final = recurrent_decode_reference( + *args, config=cfg, initial_state=state, state_qdq=True, block_v=16 + ) + chunks = [] + carry = None + lo = 0 + for hi in [0, 5, 16, 16, 37, 73]: + result, carry = recurrent_decode_reference( + *(x[lo:hi] for x in args), + config=cfg, + carry=carry, + initial_state=state if carry is None else None, + state_qdq=True, + block_v=16, + ) + chunks.append(result) + lo = hi + actual = torch.cat(chunks) + for a, b in zip( + _values_and_grads(actual, carry.reconstruct(), args, state), + _values_and_grads(expected, final.reconstruct(), args, state), + ): + torch.testing.assert_close(a, b, rtol=0, atol=0) + assert carry.position == 73 + assert carry.anchor.format == "fp8_e4m3" + assert carry.anchor.scales.shape == (2, 1) + assert not carry.anchor.scales.requires_grad + for entry in carry.entries: + assert entry.key.scales.shape == (2, 1) + assert entry.update.values.requires_grad + assert not entry.update.scales.requires_grad + + +def test_initial_write_and_readout_order_are_explicit(): + args, state = _inputs(length=1) + outputs = [] + for readout in ["working", "stored"]: + output, carry = recurrent_decode_reference( + *args, + config=LinearAttentionDecodeConfig(readout=readout), + initial_state=state, + state_qdq=True, + block_v=16, + ) + outputs.append(output) + assert not torch.equal(*outputs) + empty, untouched = recurrent_decode_reference( + *(x[:0] for x in args), + config=LinearAttentionDecodeConfig(), + initial_state=state, + state_qdq=True, + ) + assert empty.shape == (0, 2, 11) + assert untouched.anchor.values is state + assert not untouched.started + + +def test_grid_gate_ste_keeps_gate_gradients_and_changes_trajectory(): + args, state = _inputs() + output, carry = recurrent_decode_reference( + *args, config=LinearAttentionDecodeConfig(decay_log_step=0.02), initial_state=state + ) + grad = torch.autograd.grad(output.square().sum() + carry.reconstruct().square().sum(), args[3])[ + 0 + ] + assert torch.isfinite(grad).all() and torch.count_nonzero(grad) > 0 + exact, _ = recurrent_decode_reference( + *args, config=LinearAttentionDecodeConfig(), initial_state=state + ) + assert not torch.equal(output, exact) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("prefixes", [[0, 0, 0, 0], [0, 3, 65, 0], [0, 5, 68, 0]]) +def test_explicit_packed_phase_with_empty_entries(kda, prefixes): + args, one_state = _inputs(kda) + args = tuple(x.unsqueeze(0) for x in args) + initial = torch.stack([one_state.detach().clone() for _ in range(4)]).requires_grad_() + boundaries = [0, 0, 5, 73, 73] + function = matmul_kda if kda else matmul_gdn + actual = function( + *args, + policy=LinearAttentionConfig(backend="matmul", decode={}), + prefill_lengths=prefixes, + cu_seqlens=torch.tensor(boundaries), + initial_state=initial, + output_final_state=True, + ) + outputs, states = [], [] + for n, (lo, hi) in enumerate(itertools.pairwise(boundaries)): + if lo == hi: + outputs.append(args[2][0, lo:hi]) + states.append(initial[n]) + continue + o, s = recurrent_delta_rule_reference( + *(x[:, lo:hi] for x in args), initial_state=initial[n : n + 1] + ) + outputs.append(o[0]) + states.append(s[0]) + expected = torch.cat(outputs).unsqueeze(0), torch.stack(states) + for a, e in zip(actual, expected): + torch.testing.assert_close(a, e, rtol=2e-11, atol=1e-12) + gradients = [ + torch.autograd.grad( + sum(x.square().sum() for x in result), (*args, initial), retain_graph=True + ) + for result in (actual, expected) + ] + for a, e in zip(*gradients): + torch.testing.assert_close(a, e, rtol=2e-10, atol=1e-11) + + +@pytest.mark.parametrize("mode", ["token", "replay"]) +def test_carry_rejects_changed_policy_and_shape(mode): + args, state = _inputs(length=4) + cfg = LinearAttentionDecodeConfig(mode=mode, replay={"window": 8} if mode == "replay" else None) + _, carry = recurrent_decode_reference(*args, config=cfg, initial_state=state, state_qdq=True) + for changed in [ + cfg.model_copy(update={"readout": "working"}), + cfg.model_copy(update={"decay_log_step": 0.01}), + ]: + with pytest.raises(ValueError, match="policy or state shape"): + recurrent_decode_reference(*args, config=changed, carry=carry, state_qdq=True) + with pytest.raises(ValueError, match="policy or state shape"): + recurrent_decode_reference(*args, config=cfg, carry=carry, state_qdq=True, block_v=16) + with pytest.raises(ValueError, match="either carry or initial"): + recurrent_decode_reference( + *args, config=cfg, carry=carry, initial_state=state, state_qdq=True + ) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("all_empty", [False, True]) +def test_grouped_heads_transposed_state_and_empty_batch(kda, all_empty): + args, state = _inputs(kda, length=0 if all_empty else 9) + q, k, v, g, beta = [x.unsqueeze(0).repeat(2, *([1] * x.ndim)) for x in args] + q, k = q[:, :, :1], k[:, :, :1] + initial = state.unsqueeze(0).repeat(2, 1, 1, 1).transpose(-1, -2).detach().requires_grad_() + function = matmul_kda if kda else matmul_gdn + policy = LinearAttentionConfig( + backend="matmul", decode={"mode": "replay", "replay": {"window": 3}} + ) + shared = { + "policy": policy, + "state_qdq": True, + "output_final_state": True, + } + actual = function( + q, + k, + v, + g, + beta, + initial_state=initial, + state_v_first=True, + prefill_lengths=[0, 0] if all_empty else [2, 5], + **shared, + ) + expected = function( + q.repeat_interleave(2, dim=2), + k.repeat_interleave(2, dim=2), + v, + g, + beta, + initial_state=initial.transpose(-1, -2), + prefill_lengths=[0, 0] if all_empty else [2, 5], + **shared, + ) + torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0) + torch.testing.assert_close(actual[1].transpose(-1, -2), expected[1], rtol=0, atol=0) + gradients = [ + torch.autograd.grad(result[0].sum() + result[1].sum(), initial, retain_graph=True)[0] + for result in (actual, expected) + ] + torch.testing.assert_close(*gradients, rtol=0, atol=0) + if all_empty: + torch.testing.assert_close(actual[1], initial, rtol=0, atol=0) diff --git a/tests/unit/torch/quantization/test_linear_attention_hadamard.py b/tests/unit/torch/quantization/test_linear_attention_hadamard.py new file mode 100644 index 00000000000..5f6b40f4c5c --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_hadamard.py @@ -0,0 +1,227 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import pytest +import torch +import torch.nn.functional as F + +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.config import QuantizeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + LinearAttentionDecodeConfig, + matmul_gdn, + matmul_kda, + recurrent_decode_reference, +) +from modelopt.torch.quantization.linear_attention.decode import _encode, _hadamard32 + + +def _rotate(value): + # Independent dense Sylvester matrix; production uses a butterfly transform. + matrix = value.new_tensor([[(-1) ** (i & j).bit_count() for j in range(32)] for i in range(32)]) + return (value.reshape(*value.shape[:-1], -1, 32) @ (matrix / 32**0.5)).reshape_as(value) + + +def _qdq(value): + with torch.no_grad(): + groups = value.float().unflatten(-1, (-1, 32)) + scale = (groups.abs().amax(-1, keepdim=True) / 127).clamp_min(6e-8) + normalized = groups / scale + codes = torch.where(normalized >= 0, (normalized + 0.5).floor(), (normalized - 0.5).ceil()) + rounded = (codes.clamp(-127, 127) * scale.half().float()).flatten(-2).to(value.dtype) + return (value - value.detach()) + rounded + + +def _inputs(kda, length=11): + torch.manual_seed(321) + q, k = [F.normalize(torch.randn(length, 2, 4, dtype=torch.float64), dim=-1) for _ in range(2)] + v = torch.randn(length, 2, 64, dtype=torch.float64) + g = -torch.rand(k.shape if kda else k.shape[:-1], dtype=torch.float64) * 0.05 + beta = torch.rand(length, 2, dtype=torch.float64) * 0.4 + initial = torch.randn(2, 4, 64, dtype=torch.float64) * 0.1 + return tuple(x.requires_grad_() for x in (q, k, v, g, beta)), initial.requires_grad_() + + +def _oracle(args, initial, mode, readout, quantize=True): + q, k, v, g, beta = args + state = _rotate(initial) + if quantize: + state = _qdq(state) + outputs = [] + for t in range(len(q)): + decay = g[t].exp().reshape(2, -1, 1) + decayed = state * decay + correction = beta[t, :, None] * (_rotate(v[t]) - torch.einsum("hk,hkv->hv", k[t], decayed)) + working = decayed + k[t, :, :, None] * correction[:, None, :] + state = _qdq(working) if quantize and (mode == "token" or (t + 1) % 3 == 0) else working + read = state if readout == "stored" else working + outputs.append(_rotate(torch.einsum("hk,hkv->hv", q[t], read)) / q.shape[-1] ** 0.5) + return torch.stack(outputs), _rotate(state) + + +def _check_with_grads(actual, expected, args, initial, tolerance=2e-10): + for a, e in zip(actual, expected): + torch.testing.assert_close(a, e, rtol=tolerance, atol=tolerance) + gradients = [ + torch.autograd.grad( + sum(x.square().sum() for x in result), (*args, initial), retain_graph=True + ) + for result in (actual, expected) + ] + for a, e in zip(*gradients): + torch.testing.assert_close(a, e, rtol=tolerance, atol=tolerance) + + +def test_hadamard_codec_basis_scales_ties_and_identity_gradient(): + raw = torch.zeros(1, 2, 64, dtype=torch.float64) + raw[0, 0, :7] = torch.tensor([254, -254, 1, -1, 3, -3, 0]) + raw[0, 1, 32:] = torch.linspace(-0.123, 0.123, 32) + raw.requires_grad_() + torch.testing.assert_close(_hadamard32(raw), _rotate(raw), rtol=1e-14, atol=1e-14) + torch.testing.assert_close(_hadamard32(_hadamard32(raw)), raw, rtol=1e-14, atol=1e-13) + encoded = _encode(raw, True, 64, state=True, state_format="int8", state_codec="int8_hadamard32") + torch.testing.assert_close(encoded.values, _qdq(raw), rtol=0, atol=0) + torch.testing.assert_close( + encoded.values[0, 0, :7], raw.new_tensor([254, -254, 2, -2, 4, -4, 0]) + ) + assert encoded.scales.shape == (1, 2, 2) + assert encoded.scales.dtype == torch.float16 and not encoded.scales.requires_grad + assert encoded.scales[0, 0, 0] == 2 + assert encoded.scales[0, 0, 1] == torch.tensor(6e-8, dtype=torch.float16) + probe = torch.randn_like(raw) + (gradient,) = torch.autograd.grad((encoded.values * probe).sum(), raw) + torch.testing.assert_close(gradient, probe, rtol=0, atol=0) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize( + ("mode", "readout"), [("token", "stored"), ("replay", "working"), ("replay", "stored")] +) +def test_hadamard_recurrence_matches_dense_oracle_and_split_carry(kda, mode, readout): + args, initial = _inputs(kda) + cfg = LinearAttentionDecodeConfig( + mode=mode, + readout=readout, + state_codec="int8_hadamard32", + replay={"window": 3, "factor_qdq": False} if mode == "replay" else None, + ) + kwargs = {"config": cfg, "state_format": "int8", "state_qdq": True} + output, carry = recurrent_decode_reference(*args, initial_state=initial, **kwargs) + _check_with_grads( + (output, carry.reconstruct()), _oracle(args, initial, mode, readout), args, initial + ) + assert carry.value_basis == "hadamard32" and carry.anchor.block_v == 32 + assert carry.anchor.scales.shape == (2, 4, 2) + pieces, split_carry, start = [], None, 0 + for end in (0, 2, 3, 3, 7, 11): + piece, split_carry = recurrent_decode_reference( + *(x[start:end] for x in args), + carry=split_carry, + initial_state=initial if split_carry is None else None, + **kwargs, + ) + if end == 0: + assert split_carry.reconstruct() is initial and not split_carry.started + pieces.append(piece) + start = end + _check_with_grads( + (torch.cat(pieces), split_carry.reconstruct()), + (output, carry.reconstruct()), + args, + initial, + 0, + ) + with pytest.raises(ValueError, match="Carry policy"): + recurrent_decode_reference( + *args, + carry=carry, + **{**kwargs, "config": cfg.model_copy(update={"state_codec": "tile"})}, + ) + + +@pytest.mark.parametrize("kda", [False, True]) +def test_hadamard_disabled_qdq_is_a_change_of_basis(kda): + args, initial = _inputs(kda) + cfg = LinearAttentionDecodeConfig(mode="replay", replay={"window": 3, "factor_qdq": False}) + results = [] + for codec in ("tile", "int8_hadamard32"): + out, carry = recurrent_decode_reference( + *args, initial_state=initial, config=cfg.model_copy(update={"state_codec": codec}) + ) + results.append((out, carry.reconstruct())) + _check_with_grads(*results, args, initial) + + +@pytest.mark.parametrize(("kda", "prefix"), [(False, 0), (False, 3), (True, 3), (True, 11)]) +def test_hadamard_exact_prefill_handoff_and_state_layout(kda, prefix, project_root_path): + args, initial = _inputs(kda) + function = matmul_kda if kda else matmul_gdn + if kda: + path = ( + project_root_path + / "examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json" + ) + recipe = QuantizeConfig(**json.loads(path.read_text())) + else: + recipe = load_recipe("general/ptq/gdn_state_int8_dynamic").quantize + policy = recipe.linear_attention[0].cfg + assert recipe.algorithm is None + assert recipe.quant_cfg[1].cfg.num_bits == 8 + assert policy.decode.state_codec == "int8_hadamard32" + assert not policy.decode.prefill_state_qdq + # Exercise the shipped policy with the CPU implementation for this algebra test. + policy.decode.implementation = "torch" + actual = function( + *(x.unsqueeze(0) for x in args), + policy=policy, + prefill_lengths=[prefix], + state_qdq=True, + state_format="int8", + initial_state=initial.transpose(-1, -2).unsqueeze(0), + state_v_first=True, + output_final_state=True, + ) + prefix_out, state = recurrent_decode_reference( + *(x[:prefix] for x in args), initial_state=initial, config=LinearAttentionDecodeConfig() + ) + suffix, carry = recurrent_decode_reference( + *(x[prefix:] for x in args), + initial_state=state.reconstruct(), + config=policy.decode, + state_qdq=True, + state_format="int8", + ) + expected = (torch.cat((prefix_out, suffix)), carry.reconstruct()) + _check_with_grads((actual[0][0], actual[1][0].transpose(-1, -2)), expected, args, initial) + + +def test_hadamard_rejects_unsupported_codec_combinations(): + args, initial = _inputs(False) + cfg = LinearAttentionDecodeConfig(state_codec="int8_hadamard32") + for kwargs, message in [ + ({"state_qdq": True}, "requires INT8"), + ({"state_format": "int8", "block_v": 16}, "block_v >= 32"), + ]: + with pytest.raises(ValueError, match=message): + recurrent_decode_reference(*args, config=cfg, initial_state=initial, **kwargs) + with pytest.raises(ValueError, match="Dv divisible"): + recurrent_decode_reference(args[0], args[1], args[2][..., :33], *args[3:], config=cfg) + with pytest.raises(ValueError, match="decode handoff"): + LinearAttentionDecodeConfig(state_codec="int8_hadamard32", prefill_state_qdq=True) + with pytest.raises(ValueError, match="block_v >= 32"): + LinearAttentionConfig(backend="matmul", state={"block_v": 16}, decode=cfg) diff --git a/tests/unit/torch/quantization/test_linear_attention_int8.py b/tests/unit/torch/quantization/test_linear_attention_int8.py new file mode 100644 index 00000000000..b3de8c94d47 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_int8.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import pytest +import torch +import torch.nn.functional as F +from _test_utils.torch.quantization.linear_attention_reference import state_qdq_reference + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionDecodeConfig, + recurrent_decode_reference, +) +from modelopt.torch.quantization.linear_attention.decode import _encode +from modelopt.torch.quantization.nn import TensorQuantizer + +INT8 = {"num_bits": 8, "unsigned": False, "narrow_range": True, "type": "dynamic", "axis": (0, 1)} + + +@pytest.mark.parametrize("block_v", [16, 32, 64, 128]) +def test_int8_state_codec_matches_modelopt_and_identity_gradient(block_v): + # Exact scale two exercises both signs, endpoints and nearest-even half ties. + state = torch.zeros(2, 3, 5, block_v + 3) + state[..., 0, :9] = torch.tensor([-127, -126.5, -1.5, -0.5, 0, 0.5, 1.5, 126.5, 127]) * 2 + state[0, :, 0, block_v:] = torch.tensor([-254, 1, 254]) + state.requires_grad_() + quantizer = TensorQuantizer(QuantizerAttributeConfig(**INT8)) + expected = torch.cat( + [quantizer(tile.flatten(-2)).reshape_as(tile) for tile in state.split(block_v, -1)], -1 + ) + encoded = _encode(state, True, block_v, state=True, state_format="int8") + reference = state_qdq_reference(state, block_v, "int8") + torch.testing.assert_close(encoded.values, expected, rtol=0, atol=0) + torch.testing.assert_close(reference, expected, rtol=0, atol=0) + assert encoded.format == "int8" + assert torch.equal(encoded.scales[1, :, -1], torch.ones(3)) + assert not encoded.scales.requires_grad + probe = torch.randn_like(state) + for result in (encoded.values, reference): + (gradient,) = torch.autograd.grad((result * probe).sum(), state, retain_graph=True) + torch.testing.assert_close(gradient, probe, rtol=0, atol=0) + + +def _inputs(kda): + torch.manual_seed(913) + q, k = [F.normalize(torch.randn(23, 2, 16), dim=-1).requires_grad_() for _ in range(2)] + v = torch.randn(23, 2, 19, requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1]) * 0.03).requires_grad_() + beta = (torch.rand(23, 2) * 0.4).requires_grad_() + state = (torch.randn(2, 16, 19) * 0.1).requires_grad_() + return (q, k, v, g, beta), state + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("mode", ["token", "replay"]) +def test_int8_decode_continuation_and_format_rejection(kda, mode): + args, initial = _inputs(kda) + cfg = LinearAttentionDecodeConfig(mode=mode, replay={"window": 5} if mode == "replay" else None) + kwargs = {"config": cfg, "state_qdq": True, "state_format": "int8", "block_v": 16} + expected, final = recurrent_decode_reference(*args, initial_state=initial, **kwargs) + pieces, carry, start = [], None, 0 + for end in (0, 3, 10, 10, 23): + output, carry = recurrent_decode_reference( + *(x[start:end] for x in args), + carry=carry, + initial_state=initial if carry is None else None, + **kwargs, + ) + pieces.append(output) + start = end + actual = torch.cat(pieces) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(carry.reconstruct(), final.reconstruct(), rtol=0, atol=0) + probes = (torch.randn_like(expected), torch.randn_like(initial)) + grads = [ + torch.autograd.grad((o * probes[0]).sum() + (s * probes[1]).sum(), (*args, initial)) + for o, s in ((actual, carry.reconstruct()), (expected, final.reconstruct())) + ] + for actual_grad, expected_grad in zip(*grads): + torch.testing.assert_close(actual_grad, expected_grad, rtol=0, atol=0) + assert carry.anchor.format == "int8" + assert all(e.key.format == e.update.format == "fp8_e4m3" for e in carry.entries) + with pytest.raises(ValueError, match="Carry policy"): + recurrent_decode_reference( + *(x[:0] for x in args), carry=carry, **{**kwargs, "state_format": "fp8_e4m3"} + ) From 756c3304c64457f5d3646c7ba217ac07971a786f Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sun, 27 Sep 2026 20:58:13 -0700 Subject: [PATCH 6/6] Add fused Triton GDN/KDA decode QAT Signed-off-by: Kai Xu --- examples/llm_qat/linear_attention/README.md | 15 +- .../configs/kda_decode_state_int8.json | 2 +- .../kernels/quantization/common/fp8_quant.py | 31 +- .../quantization/linear_attention/decode.py | 468 ++++++++++++++++++ .../quantization/linear_attention/config.py | 4 +- .../quantization/linear_attention/decode.py | 102 +++- .../linear_attention/test_decode.py | 412 +++++++++++++++ .../torch/quantization/plugins/test_kda.py | 4 +- 8 files changed, 1016 insertions(+), 22 deletions(-) create mode 100644 modelopt/torch/kernels/quantization/linear_attention/decode.py create mode 100644 tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py diff --git a/examples/llm_qat/linear_attention/README.md b/examples/llm_qat/linear_attention/README.md index 7c44058512a..ad0e116c1c3 100644 --- a/examples/llm_qat/linear_attention/README.md +++ b/examples/llm_qat/linear_attention/README.md @@ -55,8 +55,8 @@ states, so the policy must specify which computation training should emulate. `LinearAttentionConfig` holds the overall backend, chunk size, and state settings. Its optional `decode` field contains a `LinearAttentionDecodeConfig` for token or -replay mode, state codec, decay approximation, and replay settings. The decode -implementation uses Torch operations and autograd. This is one nested configuration: supplying a `decode` +replay mode, Torch or Triton implementation, state codec, decay approximation, +and replay settings. This is one nested configuration: supplying a `decode` dictionary in the recipe constructs the nested config automatically. Each `linear_attention` entry uses `module_name` to select attention layers and @@ -101,7 +101,9 @@ mtq.quantize(gdn_model, gdn_recipe) Use the phase context shown below for each forward/backward. Importing only `configs/ptq/units/gdn_state_int8_dynamic` configures the quantizer without -selecting Hadamard. Both complete recipes use the Torch decode implementation. +selecting Hadamard. The complete GDN recipe uses the Torch implementation; +set `gdn_recipe.linear_attention[0].cfg.decode.implementation="triton"` before +conversion to select the fused decode implementation. ### Why training needs prefill/decode boundaries @@ -293,9 +295,10 @@ the optional tile codec instead uses nearest-even rounding and FP32 scales. ## Training boundaries -The Torch implementation uses autograd through initial states, chunk handoff, -token writes, and replay refreshes. The decode recurrence runs token by token in -Python, so long training suffixes can be slow. +The Torch and Triton implementations support first-order QAT gradients through +initial states, chunk handoff, token writes, and replay refreshes. Triton supports +key dimensions up to 128 and value blocks 16/32/64/128, with checkpointed backward. +Use the Torch implementation for higher-order differentiation. Prefill prefixes use exact chunk algebra with optional state QDQ. This example has no prefill GEMM QDQ or approximate inverse. FLA KDA requires `use_cache=False`; diff --git a/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json b/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json index 8a29cda8139..f6cc5fef650 100644 --- a/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json +++ b/examples/llm_qat/linear_attention/configs/kda_decode_state_int8.json @@ -25,7 +25,7 @@ "cfg": { "backend": "matmul", "decode": { - "implementation": "torch", + "implementation": "triton", "state_codec": "int8_hadamard32", "prefill_state_qdq": false } diff --git a/modelopt/torch/kernels/quantization/common/fp8_quant.py b/modelopt/torch/kernels/quantization/common/fp8_quant.py index 2b76499b968..30c2d148629 100644 --- a/modelopt/torch/kernels/quantization/common/fp8_quant.py +++ b/modelopt/torch/kernels/quantization/common/fp8_quant.py @@ -16,11 +16,36 @@ """Composable Triton JIT functions for FP8 (E4M3) fake quantization. Counterpart of ``nvfp4_quant.py`` for per-tensor FP8. Used by the unified -flash-attention kernel's softmax-P qdq (``common/attention/triton_fa.py``). +flash-attention kernel's softmax-P qdq and linear-attention state/factor qdq. """ import triton import triton.language as tl +from triton.language.extra.cuda import libdevice + +__all__ = ["fp8_dynamic_qdq", "fp8_scalar_qdq"] + + +@triton.jit +def _fp8_cast(value, NATIVE_FP8: tl.constexpr): + value = tl.clamp(value, -448.0, 448.0) + if NATIVE_FP8: + return value.to(tl.float8e4nv).to(tl.float32) + # GPUs before SM89 need a software E4M3 conversion, including subnormals and ties. + bits = tl.abs(value).to(tl.int32, bitcast=True) + exponent = ((bits >> 23) & 255) - 127 + step = ((tl.maximum(exponent, -6) - 3 + 127) << 23).to(tl.float32, bitcast=True) + return libdevice.nearbyint(value / step) * step + + +@triton.jit +def fp8_dynamic_qdq(value, NATIVE_FP8: tl.constexpr = True): + """Dynamic per-tile FP8 QDQ with TensorQuantizer scaling and returned decode scale.""" + amax = tl.max(tl.abs(value)) + safe_amax = tl.where(amax <= 2.0**-24, 1.0, amax) + quant_scale = tl.div_rn(448.0, safe_amax) + decode_scale = tl.div_rn(1.0, quant_scale) + return _fp8_cast(value * quant_scale, NATIVE_FP8) * decode_scale, decode_scale @triton.jit @@ -38,6 +63,4 @@ def fp8_scalar_qdq(x, scale): Returns: Fake-quantized tensor of the same shape as ``x``, in float32. """ - FP8_E4M3_MAX: tl.constexpr = 448.0 - x_scaled = tl.clamp(x / scale, -FP8_E4M3_MAX, FP8_E4M3_MAX) - return x_scaled.to(tl.float8e4nv).to(tl.float32) * scale + return _fp8_cast(x / scale, True) * scale diff --git a/modelopt/torch/kernels/quantization/linear_attention/decode.py b/modelopt/torch/kernels/quantization/linear_attention/decode.py new file mode 100644 index 00000000000..a156f824b12 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/decode.py @@ -0,0 +1,468 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fused FP32 token/replay training recurrence with state checkpoint recomputation.""" + +import torch +import triton +import triton.language as tl +from torch.autograd.function import once_differentiable +from triton.language.extra.cuda import libdevice + +from ..common.fp8_quant import fp8_dynamic_qdq +from .int8 import int8_scalar_qdq + +__all__ = ["fused_recurrence"] + + +@triton.jit +def _sum_keys(value): + BK: tl.constexpr = value.shape[0] + rows = tl.arange(0, BK) + for level in tl.static_range(0, tl.constexpr(triton.next_power_of_2(BK).bit_length() - 1)): + offset = BK >> (level + 1) + value = value + tl.gather( + value, tl.broadcast_to((rows ^ offset)[:, None], value.shape), axis=0 + ) + return tl.sum(tl.where(rows[:, None] == 0, value, 0.0), axis=0) + + +@triton.jit +def _advance( + state, + key, + value, + gate, + beta, + write_state, + CHANNEL_GATE: tl.constexpr, + STATE_QDQ: tl.constexpr, + FACTOR_QDQ: tl.constexpr, + NATIVE_FP8: tl.constexpr, +): + decay = libdevice.exp(gate) + if CHANNEL_GATE: + decayed = state * decay[:, None] + else: + decayed = state * decay + residual = value - _sum_keys(key[:, None] * decayed) + update = beta * residual + update_scale = 1.0 + if FACTOR_QDQ: + update, update_scale = fp8_dynamic_qdq(update, NATIVE_FP8) + working = decayed + key[:, None] * update[None, :] + stored = working + if STATE_QDQ == 3: + state_scale = tl.full((working.shape[0], working.shape[1] // 32), 1.0, tl.float16) + else: + state_scale = 1.0 + if STATE_QDQ: + if write_state: + if STATE_QDQ == 3: + groups = tl.reshape(working, (working.shape[0], working.shape[1] // 32, 32)) + scale = tl.maximum(tl.max(tl.abs(groups), 2) * (1.0 / 127.0), 6e-8) + normalized = tl.div_rn(groups, scale[:, :, None]) + codes = tl.floor(tl.abs(normalized) + 0.5) + codes = tl.where(normalized < 0, -codes, codes) + state_scale = scale.to(tl.float16) + stored = tl.reshape( + tl.minimum(tl.maximum(codes, -127.0), 127.0) + * state_scale.to(tl.float32)[:, :, None], + working.shape, + ) + elif STATE_QDQ == 2: + amax = tl.max(tl.abs(working)) + # Match Torch scalar division: FP32 multiply by the rounded reciprocal. + state_scale = tl.where(amax > 0, amax * (1.0 / 127.0), 1.0) + stored = int8_scalar_qdq(working, state_scale) + else: + stored, state_scale = fp8_dynamic_qdq(working, NATIVE_FP8) + return stored, working, decayed, residual, update, update_scale, state_scale + + +@triton.jit +def _decode_fwd( + Q, + K, + V, + G, + Beta, + Initial, + Out, + Final, + Updates, + Anchor, + AnchorScale, + UpdateScale, + Checkpoints, + T: tl.constexpr, + H: tl.constexpr, + DK: tl.constexpr, + DV: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NV: tl.constexpr, + CHANNEL_GATE: tl.constexpr, + STATE_QDQ: tl.constexpr, + REPLAY: tl.constexpr, + FACTOR_QDQ: tl.constexpr, + WINDOW: tl.constexpr, + CURSOR: tl.constexpr, + READ_STORED: tl.constexpr, + SCALE: tl.constexpr, + INTERVAL: tl.constexpr, + NATIVE_FP8: tl.constexpr, +): + h = tl.program_id(0) + tile = tl.program_id(1) + keys = tl.arange(0, BK) + values = tile * BV + tl.arange(0, BV) + offsets = h * DK * DV + keys[:, None] * DV + values[None, :] + mask = (keys[:, None] < DK) & (values[None, :] < DV) + state = tl.load(Initial + offsets, mask, 0) + tl.store(Anchor + offsets, state, mask) + if STATE_QDQ == 3: + groups = tile * (BV // 32) + tl.arange(0, BV // 32) + scale_offsets = (h * DK + keys[:, None]) * (DV // 32) + groups[None, :] + scale_mask = (keys[:, None] < DK) & (groups[None, :] < DV // 32) + tl.store(AnchorScale + scale_offsets, 1.0, scale_mask) + else: + tl.store(AnchorScale + h * NV + tile, 1.0) + for t in range(T): + if t % INTERVAL == 0: + tl.store(Checkpoints + (t // INTERVAL) * H * DK * DV + offsets, state, mask) + query = tl.load(Q + (t * H + h) * DK + keys, keys < DK, 0) + key = tl.load(K + (t * H + h) * DK + keys, keys < DK, 0) + value = tl.load(V + (t * H + h) * DV + values, values < DV, 0) + beta = tl.load(Beta + t * H + h) + if CHANNEL_GATE: + gate = tl.load(G + (t * H + h) * DK + keys, keys < DK, 0) + else: + gate = tl.load(G + t * H + h) + write_state = not REPLAY or (CURSOR + t + 1) % WINDOW == 0 + state, working, decayed, residual, update, update_scale, state_scale = _advance( + state, + key, + value, + gate, + beta, + write_state, + CHANNEL_GATE, + STATE_QDQ, + FACTOR_QDQ, + NATIVE_FP8, + ) + read = state if READ_STORED else working + output = _sum_keys(query[:, None] * read) * SCALE + tl.store(Out + (t * H + h) * DV + values, output, values < DV) + tl.store(Updates + (t * H + h) * DV + values, update, values < DV) + tl.store(UpdateScale + (t * H + h) * NV + tile, update_scale) + if write_state: + tl.store(Anchor + offsets, state, mask) + if STATE_QDQ == 3: + tl.store(AnchorScale + scale_offsets, state_scale, scale_mask) + else: + tl.store(AnchorScale + h * NV + tile, state_scale) + tl.store(Final + offsets, state, mask) + + +@triton.jit +def _decode_bwd( + Q, + K, + V, + G, + Beta, + Checkpoints, + DOut, + DFinal, + DUpdates, + DAnchor, + DQ, + DKOut, + DVOut, + DG, + DBeta, + DInitial, + T: tl.constexpr, + H: tl.constexpr, + DK: tl.constexpr, + DV: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NV: tl.constexpr, + CHANNEL_GATE: tl.constexpr, + STATE_QDQ: tl.constexpr, + REPLAY: tl.constexpr, + FACTOR_QDQ: tl.constexpr, + WINDOW: tl.constexpr, + CURSOR: tl.constexpr, + READ_STORED: tl.constexpr, + SCALE: tl.constexpr, + INTERVAL: tl.constexpr, + NATIVE_FP8: tl.constexpr, +): + h = tl.program_id(0) + tile = tl.program_id(1) + keys = tl.arange(0, BK) + values = tile * BV + tl.arange(0, BV) + offsets = h * DK * DV + keys[:, None] * DV + values[None, :] + mask = (keys[:, None] < DK) & (values[None, :] < DV) + adjoint = tl.load(DFinal + offsets, mask, 0) + LAST_REFRESH: tl.constexpr = ((T + CURSOR) // WINDOW) * WINDOW - CURSOR - 1 + for t in range(T - 1, -1, -1): + start = (t // INTERVAL) * INTERVAL + state = tl.load(Checkpoints + (t // INTERVAL) * H * DK * DV + offsets, mask, 0) + for j in range(start, t): + key = tl.load(K + (j * H + h) * DK + keys, keys < DK, 0) + value = tl.load(V + (j * H + h) * DV + values, values < DV, 0) + beta = tl.load(Beta + j * H + h) + if CHANNEL_GATE: + gate = tl.load(G + (j * H + h) * DK + keys, keys < DK, 0) + else: + gate = tl.load(G + j * H + h) + write_state = not REPLAY or (CURSOR + j + 1) % WINDOW == 0 + state, _, _, _, _, _, _ = _advance( + state, + key, + value, + gate, + beta, + write_state, + CHANNEL_GATE, + STATE_QDQ, + FACTOR_QDQ, + NATIVE_FP8, + ) + query = tl.load(Q + (t * H + h) * DK + keys, keys < DK, 0) + key = tl.load(K + (t * H + h) * DK + keys, keys < DK, 0) + value = tl.load(V + (t * H + h) * DV + values, values < DV, 0) + beta = tl.load(Beta + t * H + h) + if CHANNEL_GATE: + gate = tl.load(G + (t * H + h) * DK + keys, keys < DK, 0) + else: + gate = tl.load(G + t * H + h) + write_state = not REPLAY or (CURSOR + t + 1) % WINDOW == 0 + stored, working, decayed, residual, update, _, _ = _advance( + state, + key, + value, + gate, + beta, + write_state, + CHANNEL_GATE, + STATE_QDQ, + FACTOR_QDQ, + NATIVE_FP8, + ) + if REPLAY: + if t == LAST_REFRESH: + adjoint += tl.load(DAnchor + offsets, mask, 0) + elif t == T - 1: + adjoint += tl.load(DAnchor + offsets, mask, 0) + doutput = tl.load(DOut + (t * H + h) * DV + values, values < DV, 0) + read = stored if READ_STORED else working + dquery = tl.sum(read * doutput[None, :], 1) * SCALE + dstate = adjoint + query[:, None] * (doutput[None, :] * SCALE) + dupdate = tl.sum(key[:, None] * dstate, 0) + dupdate += tl.load(DUpdates + (t * H + h) * DV + values, values < DV, 0) + dbeta = tl.sum(dupdate * residual, 0) + dvalue = beta * dupdate + dkey = tl.sum(update[None, :] * dstate - decayed * dvalue[None, :], 1) + ddecayed = dstate - key[:, None] * dvalue[None, :] + dgate = tl.sum(ddecayed * decayed, 1) + decay = libdevice.exp(gate) + if CHANNEL_GATE: + adjoint = ddecayed * decay[:, None] + tl.store(DG + ((tile * T + t) * H + h) * DK + keys, dgate, keys < DK) + else: + adjoint = ddecayed * decay + tl.store(DG + (tile * T + t) * H + h, tl.sum(dgate, 0)) + tl.store(DQ + ((tile * T + t) * H + h) * DK + keys, dquery, keys < DK) + tl.store(DKOut + ((tile * T + t) * H + h) * DK + keys, dkey, keys < DK) + tl.store(DVOut + (t * H + h) * DV + values, dvalue, values < DV) + tl.store(DBeta + (tile * T + t) * H + h, dbeta) + tl.store(DInitial + offsets, adjoint, mask) + + +class _FusedRecurrence(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q, + k, + v, + g, + beta, + initial, + state_qdq, + block_v, + replay, + factor_qdq, + window, + cursor, + read_stored, + scale, + interval, + ): + length, heads, keys = q.shape + values = v.shape[-1] + tiles = triton.cdiv(values, block_v) + checkpoints = torch.empty( + (triton.cdiv(length, interval), heads, keys, values), device=q.device, dtype=q.dtype + ) + output, updates = torch.empty_like(v), torch.empty_like(v) + final, anchor = torch.empty_like(initial), torch.empty_like(initial) + scale_shape = (heads, keys, values // 32) if state_qdq == 3 else (heads, tiles) + anchor_scales = torch.empty( + scale_shape, device=q.device, dtype=torch.float16 if state_qdq == 3 else q.dtype + ) + update_scales = torch.empty((length, heads, tiles), device=q.device, dtype=q.dtype) + kwargs = { + "T": length, + "H": heads, + "DK": keys, + "DV": values, + "BK": triton.next_power_of_2(keys), + "BV": block_v, + "NV": tiles, + "CHANNEL_GATE": g.ndim == 3, + "STATE_QDQ": state_qdq, + "REPLAY": replay, + "FACTOR_QDQ": factor_qdq, + "WINDOW": window, + "CURSOR": cursor, + "READ_STORED": read_stored, + "SCALE": scale, + "INTERVAL": interval, + "NATIVE_FP8": torch.cuda.get_device_capability(q.device) >= (8, 9), + "num_warps": 4 if triton.next_power_of_2(keys) * block_v <= 4096 else 8, + "enable_fp_fusion": False, + } + _decode_fwd[(heads, tiles)]( + q, + k, + v, + g, + beta, + initial, + output, + final, + updates, + anchor, + anchor_scales, + update_scales, + checkpoints, + **kwargs, + ) + ctx.save_for_backward(q, k, v, g, beta, initial, checkpoints) + ctx.kwargs = kwargs + ctx.mark_non_differentiable(anchor_scales, update_scales) + return output, final, updates, anchor, anchor_scales, update_scales + + @staticmethod + @once_differentiable + def backward(ctx, doutput, dfinal, dupdates, danchor, _danchor_scales, _dupdate_scales): + q, k, v, g, beta, initial, checkpoints = ctx.saved_tensors + tiles = ctx.kwargs["NV"] + dq = torch.empty((tiles, *q.shape), device=q.device, dtype=q.dtype) + dk = torch.empty_like(dq) + dv = torch.empty_like(v) + dg = torch.empty((tiles, *g.shape), device=q.device, dtype=q.dtype) + dbeta = torch.empty((tiles, *beta.shape), device=q.device, dtype=q.dtype) + dinitial = torch.empty_like(initial) + doutput = torch.zeros_like(v) if doutput is None else doutput.contiguous() + dfinal = torch.zeros_like(initial) if dfinal is None else dfinal.contiguous() + dupdates = torch.zeros_like(v) if dupdates is None else dupdates.contiguous() + danchor = torch.zeros_like(initial) if danchor is None else danchor.contiguous() + _decode_bwd[(q.shape[1], tiles)]( + q, + k, + v, + g, + beta, + checkpoints, + doutput, + dfinal, + dupdates, + danchor, + dq, + dk, + dv, + dg, + dbeta, + dinitial, + **ctx.kwargs, + ) + return dq.sum(0), dk.sum(0), dv, dg.sum(0), dbeta.sum(0), dinitial, *([None] * 9) + + +def fused_recurrence( + q, + k, + v, + g, + beta, + initial, + *, + state_qdq=False, + state_format="fp8_e4m3", + state_codec="tile", + block_v=16, + replay=False, + factor_qdq=False, + window=8, + cursor=0, + read_stored=True, + scale=None, + checkpoint_interval=8, +): + """Run a nonempty activated sequence; encoding keys/gates and carry metadata is external.""" + if not all( + x.is_cuda and x.dtype == torch.float32 and x.device == q.device + for x in (q, k, v, g, beta, initial) + ): + raise ValueError("Fused recurrence requires CUDA FP32 inputs on one device") + if q.ndim != 3 or q.shape != k.shape or v.shape[:2] != q.shape[:2] or beta.shape != q.shape[:2]: + raise ValueError("Expected aligned [T,H,D] q/k/v and [T,H] beta") + length, heads, keys = q.shape + if length == 0 or not 1 <= keys <= 128 or initial.shape != (heads, keys, v.shape[-1]): + raise ValueError("Require nonempty sequence, Dk<=128, and matching initial state") + if g.shape not in (q.shape, beta.shape) or block_v not in (16, 32, 64, 128): + raise ValueError("Invalid gate shape or state value-block grouping") + if not 1 <= window <= 64 or not 0 <= cursor < window or not 1 <= checkpoint_interval <= 32: + raise ValueError("Invalid replay cursor/window or recomputation interval") + if state_format not in ("fp8_e4m3", "int8"): + raise ValueError("State format must be fp8_e4m3 or int8") + if state_codec not in ("tile", "int8_hadamard32"): + raise ValueError("State codec must be tile or int8_hadamard32") + if state_codec == "int8_hadamard32": + if (state_qdq and state_format != "int8") or v.shape[-1] % 32 or block_v < 32: + raise ValueError("Hadamard codec requires INT8, Dv divisible by 32 and block_v >= 32") + state_mode = 3 if state_codec == "int8_hadamard32" else (2 if state_format == "int8" else 1) + state_qdq = state_mode if state_qdq else 0 + scale = keys**-0.5 if scale is None else scale + return _FusedRecurrence.apply( + *(x.contiguous() for x in (q, k, v, g, beta, initial)), + state_qdq, + block_v, + replay, + factor_qdq, + window, + cursor, + read_stored, + scale, + checkpoint_interval, + ) diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py index 1f3b610fe75..33756c5b459 100644 --- a/modelopt/torch/quantization/linear_attention/config.py +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -51,7 +51,7 @@ class LinearAttentionDecodeConfig(ModeloptBaseConfig): """Explicit suffix recurrence; workload supplies per-sequence prefix lengths.""" mode: Literal["token", "replay"] = ModeloptField(default="token") - implementation: Literal["torch"] = ModeloptField(default="torch") + implementation: Literal["torch", "triton"] = ModeloptField(default="torch") readout: Literal["working", "stored"] = ModeloptField(default="stored") quantize_initial: bool = ModeloptField(default=True) prefill_state_qdq: bool = ModeloptField(default=False) @@ -63,6 +63,8 @@ class LinearAttentionDecodeConfig(ModeloptBaseConfig): def _validate_replay(self): if (self.mode == "replay") != (self.replay is not None): raise ValueError("replay settings must be supplied exactly when mode='replay'") + if self.implementation == "triton" and self.replay and self.replay.encoding != "once": + raise ValueError("The Triton replay candidate implements encode-once only") if self.state_codec == "int8_hadamard32" and self.prefill_state_qdq: raise ValueError( "Hadamard state QDQ starts at decode handoff; disable prefill_state_qdq" diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py index 2a981c597ec..80905958bb9 100644 --- a/modelopt/torch/quantization/linear_attention/decode.py +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -322,20 +322,106 @@ def recurrent_decode( carry=None, position=0, scale=None, + checkpoint_interval=8, ): - """Run the Torch training implementation with explicit differentiable carry.""" - return recurrent_decode_reference( + """Run the selected training implementation with explicit differentiable carry. + + The fused encode-once path carries the same reconstructed state incrementally + and exposes the anchor and encoded updates needed for continuation gradients. + """ + if config.implementation == "torch": + return recurrent_decode_reference( + q, + k, + v, + g, + beta, + config=config, + state_qdq=state_qdq, + state_format=state_format, + block_v=block_v, + initial_state=initial_state, + carry=carry, + position=position, + scale=scale, + ) + carry, signature = _prepare_carry( + q, k, v, g, beta, config, state_qdq, block_v, initial_state, carry, position, state_format + ) + if len(q) == 0: + return v + (q.sum() + k.sum() + v.sum() + g.sum() + beta.sum()) * 0, carry + if config.replay is not None and config.replay.encoding != "once": + raise ValueError("The fused implementation supports encode-once replay only") + # Keep the CPU reference importable without the optional CUDA/Triton backend. + from modelopt.torch.kernels.quantization.linear_attention.decode import fused_recurrence + + if config.state_codec == "int8_hadamard32": + v = _hadamard32(v) + replay = config.replay is not None + factor_qdq = config.replay is not None and config.replay.factor_qdq + key = _encode(k, factor_qdq, k.shape[-1]) + gate = _round_log_gate(g, config.decay_log_step) + window = config.replay.window if config.replay is not None else 1 + result = fused_recurrence( q, - k, + key.values, v, - g, + gate, beta, - config=config, + carry.reconstruct(original_basis=False), state_qdq=state_qdq, state_format=state_format, + state_codec=config.state_codec, block_v=block_v, - initial_state=initial_state, - carry=carry, - position=position, + replay=replay, + factor_qdq=factor_qdq, + window=window, + cursor=carry.cursor, + read_stored=config.readout == "stored", scale=scale, + checkpoint_interval=checkpoint_interval, + ) + output, final, updates, last_anchor, anchor_scales, update_scales = result + if config.state_codec == "int8_hadamard32": + output = _hadamard32(output) + state_format = state_format if state_qdq else "identity" + state_block_v = 32 if config.state_codec == "int8_hadamard32" else block_v + if not replay: + anchor = EncodedLinearAttentionTensor( + final, + anchor_scales if state_qdq else None, + state_format, + state_block_v if state_qdq else None, + ) + return output, LinearAttentionCarry( + anchor, (), carry.position + len(q), True, signature, carry.value_basis + ) + if carry.cursor + len(q) >= window: + anchor = EncodedLinearAttentionTensor( + last_anchor, + anchor_scales if state_qdq else None, + state_format, + state_block_v if state_qdq else None, + ) + start = ((carry.cursor + len(q)) // window) * window - carry.cursor + entries = () + else: + anchor, entries, start = carry.anchor, carry.entries, 0 + additions = [] + for t in range(start, len(q)): + encoded_key = EncodedLinearAttentionTensor( + key.values[t], + key.scales[t] if key.scales is not None else None, + key.format, + key.block_v, + ) + encoded_update = EncodedLinearAttentionTensor( + updates[t], + update_scales[t] if factor_qdq else None, + "fp8_e4m3" if factor_qdq else "identity", + block_v if factor_qdq else None, + ) + additions.append(ReplayEntry(encoded_key, encoded_update, gate[t])) + return output, LinearAttentionCarry( + anchor, (*entries, *additions), carry.position + len(q), True, signature, carry.value_basis ) diff --git a/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py b/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py new file mode 100644 index 00000000000..da6fdf6e772 --- /dev/null +++ b/tests/gpu/torch/kernels/quantization/linear_attention/test_decode.py @@ -0,0 +1,412 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_dynamic_qdq +from modelopt.torch.kernels.quantization.linear_attention.decode import fused_recurrence +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention.config import LinearAttentionDecodeConfig +from modelopt.torch.quantization.linear_attention.decode import ( + _encode, + recurrent_decode, + recurrent_decode_reference, +) +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("state_format", [None, "fp8_e4m3", "int8"]) +@pytest.mark.parametrize("replay", [False, True]) +@pytest.mark.parametrize("read_stored", [False, True]) +def test_fused_recurrence_outputs_state_and_gradients(kda, state_format, replay, read_stored): + state_qdq = state_format is not None + state_format = state_format or "fp8_e4m3" + torch.manual_seed(193) + shape = (37, 2, 16) + q, k = [ + F.normalize(torch.randn(shape, device="cuda"), dim=-1).requires_grad_() for _ in range(2) + ] + v = torch.randn(37, 2, 11, device="cuda", requires_grad=True) + g = (-torch.rand(shape if kda else shape[:-1], device="cuda") * 0.03).requires_grad_() + beta = (torch.rand(shape[:-1], device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(2, 16, 11, device="cuda") * 0.1).requires_grad_() + cfg = LinearAttentionDecodeConfig( + mode="replay" if replay else "token", + quantize_initial=False, + readout="stored" if read_stored else "working", + replay={"window": 5} if replay else None, + ) + encoded_key = _encode(k, replay, 16).values + actual = fused_recurrence( + q, + encoded_key, + v, + g, + beta, + initial, + state_qdq=state_qdq, + state_format=state_format, + block_v=16, + replay=replay, + factor_qdq=replay, + window=5, + read_stored=read_stored, + ) + expected, carry = recurrent_decode_reference( + q, + k, + v, + g, + beta, + config=cfg, + initial_state=initial, + state_qdq=state_qdq, + state_format=state_format, + block_v=16, + ) + final = carry.reconstruct() + targets = (q, k, v, g, beta, initial) + probes = [torch.randn_like(x) for x in actual[:2]] + gradients = [ + torch.autograd.grad(sum((a * p).sum() for a, p in zip(result, probes)), targets) + for result in (actual[:2], (expected, final)) + ] + for a, e in zip((*actual[:2], *gradients[0]), (expected, final, *gradients[1])): + torch.testing.assert_close(a, e, rtol=1e-3, atol=8e-6) + + +@pytest.mark.parametrize("mode", ["token", "replay"]) +@pytest.mark.parametrize("interval", [3, 8]) +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_fused_carry_handoff_and_auxiliary_gradients(mode, interval, state_format): + torch.manual_seed(217) + q, k = [ + F.normalize(torch.randn(23, 2, 16, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(23, 2, 19, device="cuda", requires_grad=True) + g = (-torch.rand_like(q) * 0.05).requires_grad_() + beta = (torch.rand(23, 2, device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(2, 16, 19, device="cuda") * 0.2).requires_grad_() + args = (q, k, v, g, beta) + reference = LinearAttentionDecodeConfig( + mode=mode, decay_log_step=0.01, replay={"window": 5} if mode == "replay" else None + ) + candidate = reference.model_copy(update={"implementation": "triton"}) + expected, expected_carry = recurrent_decode_reference( + *args, + config=reference, + state_qdq=True, + state_format=state_format, + block_v=16, + initial_state=initial, + ) + carry = None + outputs = [] + start = 0 + for end in [0, 3, 11, 12, 23]: + output, carry = recurrent_decode( + *(x[start:end] for x in args), + config=candidate, + state_qdq=True, + state_format=state_format, + block_v=16, + initial_state=initial if carry is None else None, + carry=carry, + checkpoint_interval=interval, + ) + outputs.append(output) + start = end + actual = torch.cat(outputs) + actual_state, expected_state = carry.reconstruct(), expected_carry.reconstruct() + assert carry.cursor == expected_carry.cursor + assert carry.position == 23 and carry.started + torch.testing.assert_close( + carry.anchor.scales, expected_carry.anchor.scales, rtol=5e-5, atol=2e-8 + ) + for a, e in zip(carry.entries, expected_carry.entries): + torch.testing.assert_close(a.update.scales, e.update.scales, rtol=5e-5, atol=2e-8) + torch.testing.assert_close(a.key.values, e.key.values, rtol=0, atol=0) + probes = [torch.randn_like(actual), torch.randn_like(actual_state)] + gradients = [ + torch.autograd.grad(sum((a * p).sum() for a, p in zip(result, probes)), (*args, initial)) + for result in ((actual, actual_state), (expected, expected_state)) + ] + for a, e in zip( + (actual, actual_state, *gradients[0]), (expected, expected_state, *gradients[1]) + ): + torch.testing.assert_close(a, e, rtol=1e-3, atol=1e-5) + + +@triton.jit +def _codec_kernel(X, Y, S, N: tl.constexpr, BLOCK: tl.constexpr, NATIVE_FP8: tl.constexpr): # noqa: N803 + offsets = tl.arange(0, BLOCK) + values = tl.load(X + offsets, offsets < N, 0) + decoded, scale = fp8_dynamic_qdq(values, NATIVE_FP8) + tl.store(Y + offsets, decoded, offsets < N) + tl.store(S, scale) + + +@pytest.mark.parametrize("native_fp8", [False, True]) +@pytest.mark.parametrize( + "scale", + [0.0, 1e-8 / 448, 2.0**-24 / 448, 2.0**-23 / 448, 2.0**-12, 1.0, 2.0**12, 0.00314159, 731.173], +) +def test_fp8_codec_matches_tensor_quantizer_boundaries(scale, native_fp8): + if native_fp8 and torch.cuda.get_device_capability() < (8, 9): + pytest.skip("Native FP8 conversion requires SM89 or later") + # Exhaust all finite positive E4M3 values and their rounding ties, plus signs. + representable = torch.arange(127, dtype=torch.uint8).view(torch.float8_e4m3fn).float() + ties = (representable[:-1] + representable[1:]) / 2 + probe = torch.cat( + ( + representable, + ties, + torch.nextafter(ties, torch.full_like(ties, float("inf"))), + torch.nextafter(ties, torch.zeros_like(ties)), + ) + ) + probe = torch.cat((probe, -probe)) * scale + values = probe.cuda() + actual = torch.empty_like(values) + actual_scale = torch.empty((), device="cuda") + _codec_kernel[(1,)]( + values, + actual, + actual_scale, + len(values), + triton.next_power_of_2(len(values)), + native_fp8, + ) + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), type="dynamic")) + torch.testing.assert_close(actual, quantizer(values), rtol=0, atol=0) + amax = values.abs().amax() + safe_amax = torch.where(amax <= 2**-24, torch.ones_like(amax), amax) + torch.testing.assert_close( + actual_scale, torch.div(448.0, safe_amax).reciprocal(), rtol=0, atol=0 + ) + + +def test_fused_fp8_state_matches_tensor_quantizer_with_partial_tiles(): + torch.manual_seed(762) + amplitudes = torch.tensor([0.0, 1e-8, 2**-24, 2**-23, 1.0, 731.173], device="cuda") + state = (torch.rand(6, 3, 19, device="cuda") * 2 - 1) * amplitudes[:, None, None] + state[:, 0, 0] = amplitudes + state.requires_grad_() + q = torch.zeros(1, 6, 3, device="cuda") + v = torch.zeros(1, 6, 19, device="cuda") + gate = torch.zeros(1, 6, device="cuda") + _, actual, *_ = fused_recurrence(q, q, v, gate, gate, state, state_qdq=True, block_v=16) + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=(4, 3), type="dynamic", axis=0, pass_through_bwd=True) + ) + expected = torch.cat( + [quantizer(tile.flatten(-2)).reshape_as(tile) for tile in state.split(16, -1)], -1 + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + probe = torch.randn_like(state) + gradients = [ + torch.autograd.grad((value * probe).sum(), state)[0] for value in (actual, expected) + ] + torch.testing.assert_close(*gradients, rtol=0, atol=0) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("mode", ["token", "replay"]) +@pytest.mark.parametrize("horizon", ["near_unit", "mixed"]) +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_wide_long_trajectory(kda, mode, horizon, state_format): + torch.manual_seed(329) + q, k = [ + F.normalize(torch.randn(257, 2, 128, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(257, 2, 70, device="cuda", requires_grad=True) + shape = q.shape if kda else q.shape[:-1] + g = -torch.rand(shape, device="cuda") * (1e-4 if horizon == "near_unit" else 3) + if horizon == "mixed": + g[::7] = -30 + g.requires_grad_() + beta = (torch.rand(257, 2, device="cuda") * 0.6).requires_grad_() + initial = (torch.randn(2, 128, 70, device="cuda") * 0.1).requires_grad_() + cfg = LinearAttentionDecodeConfig(mode=mode, replay={"window": 8} if mode == "replay" else None) + args = q, k, v, g, beta + expected, ec = recurrent_decode_reference( + *args, + config=cfg, + state_qdq=True, + state_format=state_format, + block_v=64, + initial_state=initial, + ) + actual, ac = recurrent_decode( + *args, + config=cfg.model_copy(update={"implementation": "triton"}), + state_qdq=True, + state_format=state_format, + block_v=64, + initial_state=initial, + ) + probes = [torch.randn_like(actual), torch.randn_like(ac.reconstruct())] + gradients = [ + torch.autograd.grad(sum((a * p).sum() for a, p in zip(result, probes)), (*args, initial)) + for result in ((actual, ac.reconstruct()), (expected, ec.reconstruct())) + ] + # Fixed reduction and scale arithmetic must also agree over long trajectories. + for a, e in zip( + (actual, ac.reconstruct(), *gradients[0]), (expected, ec.reconstruct(), *gradients[1]) + ): + relative = (a - e).norm() / e.norm().clamp_min(1e-8) + assert relative < 1e-4, relative.item() + assert torch.isfinite(a).all() + + +@pytest.mark.parametrize("block_v", [32, 128]) +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_full_key_width_and_value_block_tail(block_v, kda, state_format): + torch.manual_seed(418) + q, k = [ + F.normalize(torch.randn(9, 2, 128, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(9, 2, 133, device="cuda", requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1], device="cuda") * 0.1).requires_grad_() + beta = torch.rand(9, 2, device="cuda", requires_grad=True) + initial = torch.randn(2, 128, 133, device="cuda", requires_grad=True) + args = q, k, v, g, beta + cfg = LinearAttentionDecodeConfig(mode="replay", replay={"window": 3}) + results = [] + for implementation in ("torch", "triton"): + output, carry = recurrent_decode( + *args, + config=cfg.model_copy(update={"implementation": implementation}), + state_qdq=True, + state_format=state_format, + block_v=block_v, + initial_state=initial, + ) + final = carry.reconstruct() + gradients = torch.autograd.grad( + output.square().sum() + final.square().sum(), (*args, initial) + ) + results.append((output, final, *gradients)) + for actual, expected in zip(*results): + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=2e-5) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize( + ("mode", "block_v", "readout", "factor_qdq"), + [ + ("token", 32, "stored", False), + ("replay", 64, "working", False), + ("replay", 128, "stored", True), + ], +) +@pytest.mark.timeout(180) +def test_hadamard_fused_split_carry_and_gradients(kda, mode, block_v, readout, factor_qdq): + torch.manual_seed(917) + q, k = [ + F.normalize(torch.randn(37, 2, 17, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(37, 2, 96, device="cuda", requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1], device="cuda") * 0.03).requires_grad_() + beta = (torch.rand(37, 2, device="cuda") * 0.4).requires_grad_() + initial = (torch.randn(2, 17, 96, device="cuda") * 0.1).requires_grad_() + args = (q, k, v, g, beta) + cfg = LinearAttentionDecodeConfig( + mode=mode, + readout=readout, + state_codec="int8_hadamard32", + decay_log_step=1 / 256, + replay={"window": 8, "factor_qdq": factor_qdq} if mode == "replay" else None, + ) + shared = {"state_qdq": True, "state_format": "int8", "block_v": block_v} + expected, ref_carry = recurrent_decode_reference( + *args, config=cfg, initial_state=initial, **shared + ) + carry, start, pieces = None, 0, [] + for end in (0, 2, 8, 8, 19, 37): + output, carry = recurrent_decode( + *(x[start:end] for x in args), + config=cfg.model_copy(update={"implementation": "triton"}), + initial_state=initial if carry is None else None, + carry=carry, + **shared, + ) + pieces.append(output) + start = end + actual = torch.cat(pieces) + assert carry.value_basis == "hadamard32" and carry.cursor == ref_carry.cursor + assert carry.anchor.scales.shape == (2, 17, 3) and carry.anchor.scales.dtype == torch.float16 + torch.testing.assert_close(carry.anchor.scales, ref_carry.anchor.scales, rtol=0, atol=0) + actual_tensors = ( + actual, + carry.reconstruct(), + carry.anchor.values, + *(e.update.values for e in carry.entries), + ) + expected_tensors = ( + expected, + ref_carry.reconstruct(), + ref_carry.anchor.values, + *(e.update.values for e in ref_carry.entries), + ) + probes = [torch.randn_like(x) for x in expected_tensors] + gradients = [ + torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probes)), (*args, initial)) + for result in (actual_tensors, expected_tensors) + ] + for a, e in zip((*actual_tensors, *gradients[0]), (*expected_tensors, *gradients[1])): + torch.testing.assert_close(a, e, rtol=1e-3, atol=1e-5) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.timeout(180) +def test_hadamard_long_trajectory(kda): + torch.manual_seed(991) + q, k = [ + F.normalize(torch.randn(257, 1, 128, device="cuda"), dim=-1).requires_grad_() + for _ in range(2) + ] + v = torch.randn(257, 1, 64, device="cuda", requires_grad=True) + g = (-torch.rand(q.shape if kda else q.shape[:-1], device="cuda") * 1e-4).requires_grad_() + beta = (torch.rand(257, 1, device="cuda") * 0.1).requires_grad_() + initial = (torch.randn(1, 128, 64, device="cuda") * 0.1).requires_grad_() + args = (q, k, v, g, beta) + results = [] + for implementation in ("torch", "triton"): + cfg = LinearAttentionDecodeConfig( + mode="replay", + implementation=implementation, + state_codec="int8_hadamard32", + replay={"window": 8, "factor_qdq": False}, + ) + out, carry = recurrent_decode( + *args, config=cfg, initial_state=initial, state_qdq=True, state_format="int8" + ) + final = carry.reconstruct() + grads = torch.autograd.grad(out.square().sum() + final.square().sum(), (*args, initial)) + results.append((out, final, *grads)) + for a, e in zip(*results): + assert torch.isfinite(a).all() and torch.isfinite(e).all() + assert (a - e).norm() / e.norm().clamp_min(1e-8) < 1e-4 diff --git a/tests/gpu/torch/quantization/plugins/test_kda.py b/tests/gpu/torch/quantization/plugins/test_kda.py index ed67787de54..aecf7fe8ef8 100644 --- a/tests/gpu/torch/quantization/plugins/test_kda.py +++ b/tests/gpu/torch/quantization/plugins/test_kda.py @@ -83,7 +83,7 @@ def test_fla_layer_qat_restore_and_optimizer(tmp_path, mode): policy_cfg["state"] = {"block_v": 16} policy_cfg["decode"] = { "mode": "replay" if mode.startswith("decode_replay") else "token", - "implementation": "torch", + "implementation": "triton", } if mode.startswith("decode_replay"): policy_cfg["decode"]["replay"] = {"window": 5} @@ -166,7 +166,7 @@ def test_decode_phase_spans_activation_checkpoint_backward(hadamard): "cfg": { "backend": "matmul", "decode": { - "implementation": "torch", + "implementation": "triton", "state_codec": "int8_hadamard32" if hadamard else "tile", }, },