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 7b43e6efa5ca914057e44fddc814a83894f1a376 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Thu, 1 Oct 2026 21:58:08 -0700 Subject: [PATCH 5/6] Address GDN dependency and compilation review Centralize the pinned FLA test dependencies in the dev-fla optional extra used by both GPU nox sessions. Replace the exhaustive compilation and numerical matrices with three shared-shape BF16 checks and one small Megatron QAT/checkpoint test; compile only the selected paths in setup fixtures. Consolidates the dependency, extra-naming, and minimal-test follow-ups without adding Triton/TileLang CI cache plumbing. On RTX A6000, cold and warm focused runs each passed 3 tests with 1 SM89+ hardware skip; the warm run took 38.20s total and 2.41s in functional calls. Dependency wiring and pre-commit checks passed. Signed-off-by: Kai Xu --- noxfile.py | 22 +- pyproject.toml | 5 + .../test_fla_chunk_gated_delta_rule.py | 358 ++---------------- .../plugins/test_megatron_gated_delta_net.py | 55 +-- 4 files changed, 64 insertions(+), 376 deletions(-) diff --git a/noxfile.py b/noxfile.py index cd9e96ecffa..2211b0d6f9c 100644 --- a/noxfile.py +++ b/noxfile.py @@ -117,7 +117,7 @@ def gpu(session): "--no-build-isolation", "git+https://github.com/Dao-AILab/fast-hadamard-transform.git", ) - session.run("python", "-m", "pip", "install", "-e", ".[all,dev-test]") + session.run("python", "-m", "pip", "install", "-e", ".[all,dev-test,dev-fla]") session.run("python", "-m", "pip", "uninstall", "-y", "cupy-cuda12x") session.run("python", "-m", "pip", "install", "cupy-cuda13x") session.run( @@ -130,15 +130,6 @@ def gpu(session): "mamba_ssm", "causal-conv1d", ) - session.run( - "python", - "-m", - "pip", - "install", - "fla-core==0.5.1", - "tilelang==0.1.8", - "apache-tvm-ffi==0.1.9", - ) session.run("python", "-m", "pytest", "tests/gpu", *_cov_args()) @@ -149,16 +140,7 @@ def gpu_megatron(session): session.run("pip", "uninstall", "-y", "tensorrt_llm") # Pre-installed nvidia-modelopt shadows the editable install session.run("pip", "uninstall", "-y", "nvidia-modelopt") - session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]") - session.run( - "python", - "-m", - "pip", - "install", - "fla-core==0.5.1", - "tilelang==0.1.8", - "apache-tvm-ffi==0.1.9", - ) + session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test,dev-fla]") session.run("python", "-m", "pytest", "tests/gpu_megatron", *_cov_args()) diff --git a/pyproject.toml b/pyproject.toml index 47c8fbdf5ab..68ff9c32cfa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -132,6 +132,11 @@ dev-test = [ "torchvision", "torch-geometric", ] +dev-fla = [ + "apache-tvm-ffi==0.1.9", + "fla-core==0.5.1", + "tilelang==0.1.8", +] # Compound extras via self-references mlflow = [ # Dependencies for modelopt.torch.utils.mlflow run tracking "mlflow-skinny>=2.9", # client only; the tracking server's own deps are not needed here 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..defe05aa5d5 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 @@ -13,67 +13,42 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Numerical forward/backward coverage for the GDN training emulation contract.""" - -from functools import partial +"""Minimal forward/backward checks for GDN training fake quantization.""" import pytest import torch import torch.nn.functional as F -from _test_utils.torch.quantization.linear_attention_reference import ( - chunk_gdn_reference, - state_fp8_qdq_reference, -) -from torch.utils.checkpoint import checkpoint +from _test_utils.torch.quantization.linear_attention_reference import chunk_gdn_reference from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.nn import TensorQuantizer -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 +pytest.importorskip("fla.ops.gated_delta_rule") from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( chunk_gated_delta_rule, ) -DTYPES = (torch.float32, torch.bfloat16) -QDQ_MODES = ((False, False), (False, True), (True, False), (True, True)) -STATE_LAYOUTS = ((False, False), (True, True)) -W_AXES = (None, (0, 1), (0, 1, 2)) - - -def require_state_qdq(): - if torch.cuda.get_device_capability() < (8, 9): - pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") - -def make_inputs(dtype=None, packed=False, state_v_first=False, grouped=True): - hopper_backward = IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0 - if dtype is None: - dtype = torch.bfloat16 if hopper_backward else torch.float32 - if dtype == torch.float32 and hopper_backward: - pytest.skip("Hopper/TileLang training supports BF16; FP32 is rejected before launch") +def make_inputs(): torch.manual_seed(123) - # Two full chunks plus a tail; grouped-value heads and a partial state scale tile. - batch, length, heads, value_heads, keys, values = 1, 145, 1, 2, 32, 48 - heads = heads if grouped else value_heads - q, k = [ - F.normalize(torch.randn(batch, length, heads, keys, device="cuda"), dim=-1) - for _ in range(2) - ] - v = torch.randn(batch, length, value_heads, values, device="cuda") - g = -torch.rand(batch, length, value_heads, device="cuda") * 0.1 + # Two chunks exercise recurrence with one shared BF16 shape for every mode. + shape = (1, 128, 1, 32) + q, k = [F.normalize(torch.randn(shape, device="cuda"), dim=-1) for _ in range(2)] + v = torch.randn(shape, device="cuda") + g = -torch.rand(shape[:3], device="cuda") * 0.1 beta = torch.rand_like(g) - 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_() + args = [x.to(torch.bfloat16).requires_grad_() for x in (q, k, v, g, beta)] + state = (torch.randn(1, 1, 32, 32, device="cuda") * 0.1).requires_grad_() + return args, state -def w_quantizer(axis=(0, 1, 2)): - return TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), axis=axis, type="dynamic")) +def values_and_grads(fn, args, state, **kwargs): + result = fn(*args, initial_state=state, **kwargs) + torch.manual_seed(15) + probes = [torch.randn(x.shape, device=x.device, dtype=torch.float32) for x in result] + grads = torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probes)), (*args, state)) + return result, grads def compare(actual, expected, tolerance): @@ -86,290 +61,31 @@ 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) - probes = [torch.randn(x.shape, device=x.device, dtype=torch.float32) for x in result] - grads = torch.autograd.grad(sum((x * p).sum() for x, p in zip(result, probes)), (*args, state)) - return result, grads - - -@pytest.fixture(scope="module", autouse=True) -def _compile_gdn_kernels(): - """Keep compilation/autotuning in setup, including when selecting a single test.""" - supports_state_qdq = torch.cuda.get_device_capability() >= (8, 9) - for dtype in DTYPES: - if dtype == torch.float32 and IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0: - continue - args, state = make_inputs(dtype, grouped=False) - for kernel in (fla.chunk_gated_delta_rule, chunk_gated_delta_rule): - values_and_grads(kernel, args, state, output_final_state=True) - for packed, state_v_first in STATE_LAYOUTS: - args, state = make_inputs(dtype, packed=packed, state_v_first=state_v_first) - for state_qdq, quantize_w in QDQ_MODES: - if state_qdq 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, +@pytest.fixture(scope="module", params=["disabled", "w", "state-w"]) +def compiled_gdn_case(request): + """Compile only the selected BF16 forward/backward path, outside the test-call timer.""" + state_qdq = request.param == "state-w" + if state_qdq and torch.cuda.get_device_capability() < (8, 9): + pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") + quantizer = ( + TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), axis=(0, 1, 2), type="dynamic")) + if request.param != "disabled" + else None ) - 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) + args, state = make_inputs() + kwargs = {"state_qdq": state_qdq, "w_quantizer": quantizer} + values_and_grads(chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs) torch.cuda.synchronize() + return args, state, kwargs -@pytest.mark.parametrize("dtype", DTYPES) -def test_disabled_matches_upstream_forward_and_backward(dtype): - # Upstream Hopper/TileLang backward cannot handle grouped value heads. - args, state = make_inputs(dtype, grouped=False) - expected = values_and_grads(fla.chunk_gated_delta_rule, args, state, output_final_state=True) - actual = values_and_grads(chunk_gated_delta_rule, args, state, output_final_state=True) - for a, e in zip(actual, expected): - compare(a, e, 0.005 if dtype == torch.bfloat16 else 0.001) - - -@pytest.mark.parametrize(("state_qdq", "quantize_w"), QDQ_MODES) -@pytest.mark.parametrize(("packed", "state_v_first"), STATE_LAYOUTS) -@pytest.mark.parametrize("dtype", DTYPES) -def test_numerical_surrogate_reference(state_qdq, quantize_w, packed, state_v_first, dtype): - if state_qdq: - 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, - } +def test_gdn_forward_and_backward(compiled_gdn_case): + args, state, kwargs = compiled_gdn_case reference_args = [x.detach().float().requires_grad_() for x in args] - reference_state = state.detach().float().requires_grad_() + reference_state = state.detach().clone().requires_grad_() expected = values_and_grads(chunk_gdn_reference, reference_args, reference_state, **kwargs) actual = values_and_grads( chunk_gated_delta_rule, args, state, output_final_state=True, **kwargs ) - compare_forward(actual[0], expected[0], dtype, state_qdq) - 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() - - -@pytest.mark.parametrize("axis", W_AXES) -def test_zero_initial_state_and_output_only_training_loss(axis): - args, _ = make_inputs() - quantizer = w_quantizer(axis) - expected, _ = chunk_gdn_reference(*(x.float() for x in args), w_quantizer=quantizer) - actual, final = chunk_gated_delta_rule(*args, w_quantizer=quantizer) - assert final is None - compare((actual,), (expected,), 0.03 if args[0].dtype == torch.bfloat16 else 0.01) - grads = [torch.autograd.grad(x.square().sum(), args) for x in (actual, expected)] - compare(*grads, 0.05 if args[0].dtype == torch.bfloat16 else 0.02) - - -@pytest.mark.parametrize("block_v", [16, 64, 128]) -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} - 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_forward(actual[0], expected[0], args[0].dtype, state_qdq=True) - compare(actual[1], expected[1], 0.05 if args[0].dtype == torch.bfloat16 else 0.02) - - -@pytest.mark.parametrize("block_v", [16, 32, 64, 128]) -def test_state_qdq_without_recurrent_updates(block_v): - require_state_qdq() - args, state = make_inputs(packed=True, state_v_first=True) - q, k, v, g, beta = args - _, final = chunk_gated_delta_rule( - q, - torch.zeros_like(k), - torch.zeros_like(v), - torch.zeros_like(g), - beta, - initial_state=state, - output_final_state=True, - state_v_first=True, - cu_seqlens=torch.tensor([0, 67, 145], device="cuda", dtype=torch.int32), - state_qdq=True, - state_qdq_block_v=block_v, - ) - expected = state_fp8_qdq_reference(state.transpose(-1, -2), block_v).transpose(-1, -2) - torch.testing.assert_close(final, expected, rtol=1e-6, atol=1e-7) - (grad,) = torch.autograd.grad(final.sum(), state) - torch.testing.assert_close(grad, torch.ones_like(state), rtol=0, atol=0) - - -def test_fused_gate_normalization_and_beta_gradients(): - args, state = make_inputs() - q, k, v, raw_g, raw_beta = args - a_log = torch.randn(2, device="cuda", requires_grad=True) - bias = torch.randn(2, device="cuda", requires_grad=True) - quantizer = w_quantizer() - expected = chunk_gdn_reference( - F.normalize(q.float(), dim=-1), - F.normalize(k.float(), dim=-1), - v.float(), - -a_log.exp() * F.softplus(raw_g.float() + bias), - 2 * raw_beta.float().sigmoid(), - initial_state=state, - w_quantizer=quantizer, - ) - actual = chunk_gated_delta_rule( - *args, - initial_state=state, - output_final_state=True, - w_quantizer=quantizer, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - A_log=a_log, - dt_bias=bias, - use_beta_sigmoid_in_kernel=True, - allow_neg_eigval=True, - ) - compare(actual, expected, 0.03 if q.dtype == torch.bfloat16 else 0.01) - grads = [ - torch.autograd.grad(sum(x.square().sum() for x in result), (*args, state, a_log, bias)) - for result in (actual, expected) - ] - compare(*grads, 0.05 if q.dtype == torch.bfloat16 else 0.02) - - -def test_plugin_dispatch_preserves_fla_partial(): - args, state = make_inputs() - module = GatedDeltaNetStateQuantMixin.convert(torch.nn.Module()) - module.gdn_w_quantizer = w_quantizer() - expected = values_and_grads( - chunk_gated_delta_rule, - args, - state, - output_final_state=True, - w_quantizer=module.gdn_w_quantizer, - ) - - def wrapped(q, *remaining, **kwargs): - kernel = partial(fla.chunk_gated_delta_rule, q, output_final_state=True) - return module._state_quantized_chunk_gated_delta_rule(kernel, *remaining, **kwargs) - - actual = values_and_grads(wrapped, args, state) - for a, e in zip(actual, expected): - compare(a, e, 1e-6) - - -@pytest.mark.parametrize("option", ["use_gate_in_kernal", "unsupported_option"]) -def test_unknown_keyword_rejected_before_launch(option): - args, _ = make_inputs() - with pytest.raises(TypeError, match=f"Unexpected keyword arguments: {option}"): - chunk_gated_delta_rule(*args, **{option: True}) - - -@pytest.mark.parametrize("chunk_size", [16, 32, 128]) -def test_unsupported_chunk_rejected_before_launch(chunk_size): - args, state = make_inputs() - with pytest.raises(ValueError, match="only chunk_size=64"): - chunk_gated_delta_rule(*args, chunk_size=chunk_size) - - -def test_unsupported_w_gradient_rejected_before_launch(): - args, _ = make_inputs() - quantizer = w_quantizer() - quantizer.set_from_attribute_config({"pass_through_bwd": False}) - with pytest.raises(ValueError, match="pass_through_bwd=True"): - chunk_gated_delta_rule(*args, w_quantizer=quantizer) - - -def test_hopper_fp32_rejected_before_launch(): - if not (IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0): - pytest.skip("Hopper/TileLang capability guard") - args, _ = make_inputs() - with pytest.raises(ValueError, match="requires BF16 q/k/v"): - chunk_gated_delta_rule(*(x.float() for x in args)) + compare(actual[0], expected[0], 0.03) + compare(actual[1], expected[1], 0.05) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index c46959c4d8d..0756f2d0d40 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 @@ -38,18 +38,20 @@ SEED = 1234 -def _make_model(tp_size, **config_kwargs): +def _make_model(tp_size): model = ( get_mcore_gpt_model( tensor_model_parallel_size=tp_size, - num_layers=2, + num_layers=1, 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, + linear_num_key_heads=1, + linear_num_value_heads=1, + linear_key_head_dim=32, + linear_value_head_dim=32, ) .cuda() .eval() @@ -81,7 +83,7 @@ def _gdn_config(sites): } -def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): +def _test_gdn_qat_helper(rank, size, checkpoint_path, sites): initialize_for_megatron( tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED ) @@ -100,7 +102,6 @@ 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)) gdn_modules = [m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)] assert gdn_modules, "no GatedDeltaNet layer was wrapped" @@ -143,7 +144,7 @@ def _test_gdn_qat_helper(rank, size, mode, checkpoint_path): assert not torch.equal(restored_gdn[0].in_proj.weight, before) -def _compile_gdn_qat_kernels(rank, size, mode): +def _compile_gdn_qat_kernels(rank, size, sites): initialize_for_megatron( tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED ) @@ -151,7 +152,6 @@ def _compile_gdn_qat_kernels(rank, size, mode): 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) @@ -161,30 +161,15 @@ 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): - pytest.skip("State QDQ needs native E4M3 conversion (SM89+)") - workers = request.getfixturevalue(f"dist_workers_size_{tp_size}") - workers.run(_compile_gdn_qat_kernels, mode) - return workers - - -@pytest.mark.parametrize("tp_size", [1, 2]) -@pytest.mark.parametrize("mode", ["state", "w", "both"]) -def test_gdn_qat_and_sharded_restore(compiled_gdn_workers, tmp_path, mode): - """Train through state/W QDQ after a Megatron distributed-checkpoint round trip.""" - compiled_gdn_workers.run(_test_gdn_qat_helper, mode, tmp_path) - - -def _test_gdn_context_parallel_helper(rank, size, mode): - initialize_for_megatron(context_parallel_size=size, seed=SEED) - model = _make_model(1, context_parallel_size=size) - with pytest.raises(NotImplementedError, match="context parallelism"): - mtq.quantize(model, _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) +def compiled_gdn_workers(dist_workers_size_1): + """Warm one small QAT model in the same worker, outside the test-call budget.""" + # Ampere exercises W QDQ; native FP8 devices also exercise state QDQ. + sites = ("state", "w") if torch.cuda.get_device_capability() >= (8, 9) else ("w",) + dist_workers_size_1.run(_compile_gdn_qat_kernels, sites) + return dist_workers_size_1, sites + + +def test_gdn_qat_and_sharded_restore(compiled_gdn_workers, tmp_path): + """Train through QDQ after a Megatron distributed-checkpoint round trip.""" + workers, sites = compiled_gdn_workers + workers.run(_test_gdn_qat_helper, tmp_path, sites) From d4895f8353c96162ef329c1209d52c48e408ec1d Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 2 Oct 2026 11:30:10 -0700 Subject: [PATCH 6/6] Restore the general GPU CI timeout to 60 minutes Signed-off-by: Kai Xu --- .github/workflows/gpu_tests.yml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index f0c5a38bbab..c38e462912a 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.github/workflows/gpu_tests.yml @@ -44,8 +44,7 @@ jobs: matrix: include: - example: gpu - # Includes dependency builds and cold compilation of the FLA forward/backward tests. - timeout: 75 + timeout: 60 # 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