diff --git a/CHANGELOG.rst b/CHANGELOG.rst index df983a8804e..7dcbe340bff 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,6 +16,7 @@ Changelog - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. +- Add experimental GDN/KDA recurrent-state QAT with native vLLM 0.30.0 Triton forwards and explicit prefill/decode phases through ``linear_attention_training_phase``, including dynamic blockwise INT8 state quantization. Megatron users need a Core build with native gate and packed-forward hooks; 0.19.2 does not support this workflow. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before a separate KV-cache AutoQuantize stage, with independent resumable checkpoints for the weight and @@ -33,7 +34,6 @@ Changelog - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration. - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md `_ for details. - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``. -- Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training through the standard ``quant_cfg`` interface. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer. - Add fake quantization of the sparse-attention indexer key cache and query for DeepSeek-V4 (vLLM and Megatron-Core) and GLM-5.3-Flash (vLLM) through the new ``indexer_k_quantizer`` and ``indexer_q_quantizer``. Enable them by importing the ``configs/ptq/units/indexer_k_nvfp4`` and ``configs/ptq/units/indexer_q_nvfp4`` units (NVFP4 with the global scale fixed to 1) into a recipe. - Add the ``configs/ptq/units/kv_nvfp4_mla`` recipe unit for fake quantization of the MLA KV cache (DeepSeek-V3, GLM-5.3-Flash, ...) in vLLM fake-quant serving: NVFP4 for the latent and FP8 for the RoPE key. See `examples/vllm_serve/README.md `_ for an example recipe. - Add fake quantization of MoE expert-parallel communication to vLLM fake-quant serving through the new ``dispatch_quantizer`` and ``combine_quantizer`` on vLLM's routed experts, for MoE layers built on vLLM's fused-MoE layer and any all-to-all backend. Enable them by importing the ``configs/ptq/units/moe_dispatch_nvfp4`` and ``configs/ptq/units/moe_combine_nvfp4`` units into a recipe; see ``examples/vllm_serve/README.md``. diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index 0cc0999f202..a13723ee2cf 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -76,6 +76,7 @@ QUANT_CFG_CHOICES, RecipeSupersededAction, ) +from modelopt.torch.quantization.linear_attention import get_linear_attention_layers from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0 from modelopt.torch.utils.dataset_utils import get_supported_datasets from modelopt.torch.utils.mlflow import Tool, masked_args, resolved_recipe_texts @@ -307,6 +308,13 @@ def get_quant_config(args: argparse.Namespace) -> dict: def main(args: argparse.Namespace): + mtq_config = get_quant_config(args) + if mtq_config.get("linear_attention") and mtq.need_calibration(mtq_config): + raise ValueError( + "This example supports linear_attention recipes only without calibration. " + "For calibrated projection quantization combined with state quantization, " + "use mtq.quantize with a calibration loop inside linear_attention_training_phase." + ) trust_remote_code = is_safe_repo( trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_name_or_path ) @@ -362,8 +370,6 @@ def main(args: argparse.Namespace): ) print_rank_0(f"Using calibration dataset: {args.calib_dataset_name}") - mtq_config = get_quant_config(args) - # Quantize only the language model: disable quantizers on every top-level submodule that is not # the language model (vision tower + projector). Skip aliases of language-model submodules (e.g. # Qwen's ``self.decoder = language_model.decoder``) so the LM's own layers stay enabled. @@ -474,7 +480,20 @@ def forward_loop(_model=None): warn_rank_0( "Skipping the post-quantization generation sanity check because --compress is set." ) - if not args.skip_generate and not args.compress: + state_quantized = torch.tensor( + bool(get_linear_attention_layers(language_model, enabled_only=True)), + dtype=torch.int32, + device="cuda", + ) + if torch.distributed.is_initialized(): + torch.distributed.all_reduce(state_quantized, op=torch.distributed.ReduceOp.MAX) + state_quantized = bool(state_quantized.item()) + if state_quantized and not args.skip_generate: + warn_rank_0( + "Skipping generation sanity check for state quantization; " + "use an explicit prefix/suffix evaluation to measure state QDQ." + ) + if not args.skip_generate and not args.compress and not state_quantized: print_rank_0("\nTesting quantized model with custom prompts...") # Sanity-check text generation on the quantized language model. language_model.eval() diff --git a/modelopt/torch/kernels/quantization/linear_attention/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/__init__.py index 3f003f403f8..2d1f4587cfb 100644 --- a/modelopt/torch/kernels/quantization/linear_attention/__init__.py +++ b/modelopt/torch/kernels/quantization/linear_attention/__init__.py @@ -13,12 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Linear-attention kernels for quantization. +"""Linear-attention quantization adapters for optional serving kernels. -``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. +Serving adapters import forward kernels only when selected; this package +initializer has no kernel dependencies. The existing FLA chunk kernels remain +available for direct calls with fla-core==0.5.1. """ diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py new file mode 100644 index 00000000000..eb9d4dfd539 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/__init__.py @@ -0,0 +1,27 @@ +# 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. + +"""Training adapter for native vLLM linear-attention forward kernels. + +ModelOpt owns the training state, QDQ and backward. This package imports the +selected serving arithmetic; it does not modify the vLLM engine or scheduler. +""" + +try: + import vllm +except ImportError as error: + raise ImportError( + "Linear-attention serving profiles require the optional vllm==0.30.0 dependency." + ) from error diff --git a/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py new file mode 100644 index 00000000000..f0916d224e9 --- /dev/null +++ b/modelopt/torch/kernels/quantization/linear_attention/serving/forward.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Native vLLM 0.30 Triton forwards; training owns private, differentiable state.""" + +from functools import cache, lru_cache + +import torch + + +@cache +def _kernels(precision): + # Defer optional serving/model dependencies until their explicit profile is selected. + import vllm + + if vllm.__version__ != "0.30.0": + raise ImportError("Linear-attention vLLM profiles require vllm==0.30.0") + from vllm.third_party.flash_linear_attention import ops as gdn_ops + + if precision == "vllm_kimi_k3": + from vllm.models.kimi_k3.nvidia.ops.third_party import kda as kda_ops + else: + from vllm.third_party.flash_linear_attention.ops import kda as kda_ops + return gdn_ops, kda_ops + + +def _state_layout(state): + # ModelOpt uses [..., K, V]; both supported native backends store [..., V, K]. + return state.transpose(-1, -2).contiguous() + + +def prefill(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, precision="vllm"): + """Run native chunk prefill, including the selected profile's raw-gate preparation.""" + # Match BF16 projection outputs; A_log and bias retain their parameter precision. + gdn, kda = _kernels(precision) + kimi_k3 = precision == "vllm_kimi_k3" + channel = g.ndim == 3 + if not channel and gate_inputs is not None: + raw_g, raw_beta, rate, bias = gate_inputs + q, k, v, g, beta = gdn.fused_post_conv_prep( + torch.cat([x.flatten(1) for x in (q, k, v)], dim=-1), + raw_g.flatten(0, -2).bfloat16().contiguous(), + raw_beta.flatten(0, -2).bfloat16().contiguous(), + rate.reshape(-1).contiguous(), + bias.contiguous(), + q.shape[-2], + q.shape[-1], + v.shape[-1], + apply_l2norm=normalize, + ) + normalize = False + q, k, v, g, beta = [x.unsqueeze(0).contiguous() for x in (q, k, v, g, beta)] + # Native chunk implementations may reuse V as their output buffer. + v = v.clone() + kwargs = { + "scale": scale, + "initial_state": _state_layout(state.unsqueeze(0)), + "output_final_state": True, + "use_qk_l2norm_in_kernel": normalize, + "cu_seqlens": _prefill_metadata(q.device, q.shape[1]), + } + if channel and gate_inputs is not None: + raw_g, raw_beta, rate, bias = gate_inputs + # Kimi fuses beta activation; standalone KDA takes already activated beta. + beta_arg = ( + {"raw_beta": raw_beta[None].bfloat16().contiguous()} if kimi_k3 else {"beta": beta} + ) + out, final = kda.chunk_kda_with_fused_gate( + q, + k, + v, + raw_g=raw_g[None].bfloat16().contiguous(), + A_log=rate.reshape(-1).contiguous(), + g_bias=bias.contiguous() if bias is not None else None, + **beta_arg, + **kwargs, + ) + elif channel: + out, final = kda.chunk_kda(q, k, v, g, beta, **kwargs) + else: + out, final = gdn.chunk_gated_delta_rule(q, k, v, g, beta, **kwargs) + return out[0], _state_layout(final[0]) + + +def step(q, k, v, g, beta, state, scale, normalize=False, gate_inputs=None, *, precision="vllm"): + """Run native fused decode on a private cache, preserving the incoming autograd state.""" + gdn, kda = _kernels(precision) + kimi_k3 = precision == "vllm_kimi_k3" + channel = g.ndim == 2 + cu, indices = _decode_metadata(q.device) + if channel or gate_inputs is not None: + # Only slot one is read/written; slot zero is reserved and needs no initialization. + native_state = state.new_empty(2, state.shape[0], state.shape[2], state.shape[1]) + native_state[1].copy_(state.transpose(-1, -2)) + else: + native_state = _state_layout(state.unsqueeze(0)) + if kimi_k3: + # The Kimi kernel always normalizes Q/K; validate_profile enforces normalize=True. + raw_g, raw_beta, rate, bias = gate_inputs + out, final = kda.fused_recurrent_kda_packed_decode( + mixed_qkv=torch.cat([x.flatten() for x in (q, k, v)])[None], + raw_g=raw_g[None, None].bfloat16().contiguous(), + raw_beta=raw_beta[None, None].bfloat16().contiguous(), + A_log=rate.reshape(-1).contiguous(), + dt_bias=(torch.zeros_like(raw_g).flatten() if bias is None else bias.contiguous()), + lower_bound=None, + initial_state=native_state, + state_indices=indices, + scale=scale, + ) + elif channel: + out, final = kda.fused_recurrent_kda_fwd( + *[x[None, None].contiguous() for x in (q, k, v, g, beta)], + initial_state=native_state, + inplace_final_state=True, + scale=scale, + use_qk_l2norm_in_kernel=normalize, + cu_seqlens=cu, + ssm_state_indices=indices, + ) + elif gate_inputs is not None: + raw_g, raw_beta, rate, bias = gate_inputs + out, final = gdn.fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=torch.cat([x.flatten() for x in (q, k, v)])[None], + a=raw_g[None].bfloat16().contiguous(), + b=raw_beta[None].bfloat16().contiguous(), + A_log=rate.reshape(-1).contiguous(), + dt_bias=bias.contiguous(), + scale=scale, + initial_state=native_state, + out=torch.empty_like(v, memory_format=torch.contiguous_format)[None, None], + ssm_state_indices=indices, + use_qk_l2norm_in_kernel=normalize, + ) + else: + out, final = gdn.fused_recurrent_gated_delta_rule( + *[x[None, None].contiguous() for x in (q, k, v, g, beta)], + initial_state=native_state, + inplace_final_state=False, + scale=scale, + use_qk_l2norm_in_kernel=normalize, + cu_seqlens=cu, + ) + # Indexed kernels update slot 1; the non-in-place fallback returns one state. + return out[0, 0], _state_layout(final[-1]) + + +@lru_cache(maxsize=128) +def _prefill_metadata(device, length): + # Native kernels only read these boundaries; reuse them across equal-length layers. + return torch.arange(2, device=device, dtype=torch.int32) * length + + +@cache +def _decode_metadata(device): + # Immutable single-token boundaries and cache slot; construct without a CPU-to-GPU copy. + return torch.arange(2, device=device, dtype=torch.int32), torch.ones( + 1, device=device, dtype=torch.int32 + ) + + +def validate_profile(gate_inputs, normalize, *, precision="vllm"): + """Validate the selected native serving contract before launching prefill.""" + _kernels(precision) + if precision == "vllm_kimi_k3": + if gate_inputs is None: + raise ValueError("vllm_kimi_k3 requires raw gate_inputs") + if not normalize: + raise ValueError("vLLM packed KDA decode requires Q/K normalization") diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py index a4796162f6d..f1bd15154fb 100644 --- a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -60,6 +60,7 @@ def remove_per_module_state( _ = metadata.pop("subnet_config", None) _ = metadata.pop("real_quantizer_state", None) _ = metadata.pop("q_tensor_state", None) + _ = metadata.pop("linear_attention", None) else: config["metadata"] = {} diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 6a538e16d5b..3d2d64e82e0 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.""" @@ -1575,6 +1577,11 @@ def _dict_to_entry(key: str, value) -> list[dict[str, Any]]: class QuantizeConfig(ModeloptBaseConfig): """Default configuration for ``quantize`` mode.""" + linear_attention: list[LinearAttentionPolicyEntry] = ModeloptField( + default=[], + title="Execution policies for supported linear-attention training modules", + ) + quant_cfg: QuantizeQuantCfgType = ModeloptField( default=[{"quantizer_name": "*", "cfg": {"num_bits": 8, "axis": None}}], title="Quantization configuration", diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 78a6e03a16b..9ca1781e224 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -62,12 +62,16 @@ def convert_to_quantized_model(model: ModelLikeModule, config: QuantizeConfig) -> ConvertReturnType: """Convert the model to a quantized one as per `config`.""" + # Defer shared plugin imports to avoid the plugin/conversion initialization cycle. + from .plugins.linear_attention import _apply_linear_attention_policy, _validate_linear_attention + # initialize the true module if necessary model = model.init_modellike() if isinstance(model, ModelLikeModule) else model replace_quant_module(model, version=ModeloptStateManager(model).state_version) set_quantizer_by_cfg(model, config.get("quant_cfg", [])) - _validate_linear_attention_quantizers(model) + _apply_linear_attention_policy(model, config) + _validate_linear_attention(model) metadata = {} update_quantize_metadata(model, config, metadata) @@ -95,8 +99,16 @@ def restore_quantized_model( model: ModelLikeModule, config: QuantizeConfig, metadata: MetadataDict ) -> nn.Module: """Insert quantizers to the model and restore the quantizer states from the given state dict.""" + # Defer the shared plugin import to avoid the plugin/conversion initialization cycle. + from .plugins.linear_attention import _apply_linear_attention_policy + # initialize the true module if necessary - convert_to_quantized_model(model, config) + model = model.init_modellike() if isinstance(model, ModelLikeModule) else model + replace_quant_module(model, version=ModeloptStateManager(model).state_version) + set_quantizer_by_cfg(model, config.get("quant_cfg", [])) + # Current checkpoints save resolved local policies; only older ones need matching. + if "linear_attention" not in metadata: + _apply_linear_attention_policy(model, config) return restore_quantizer_state(model, config, metadata) @@ -137,6 +149,13 @@ 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. """ + # Defer shared plugin imports to avoid the plugin/conversion initialization cycle. + from .plugins.linear_attention import ( + _restore_legacy_linear_attention_quantizers, + _restore_linear_attention_policy, + ) + + _restore_linear_attention_policy(model, metadata.get("linear_attention")) if "quantizer_state" not in metadata: # MCore sharded checkpoint (`torch-dist`) has its quantizer_state stored as the # extra_state of `QuantModule`. The quantizer_state is resumed with @@ -144,13 +163,7 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: return model quantizer_state_dict = dict(metadata["quantizer_state"]) - # Older checkpoints predate these disabled handles; preserve their baseline path. - for name, module in _linear_attention_modules(model).items(): - for handle in ("gdn_state_quantizer", "gdn_w_quantizer"): - key = f"{name}.{handle}" if name else handle - quantizer = getattr(module, handle) - if key not in quantizer_state_dict and not quantizer.is_enabled: - quantizer_state_dict[key] = quantizer.get_modelopt_state() + _restore_legacy_linear_attention_quantizers(model, quantizer_state_dict) unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys() extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys() @@ -203,6 +216,14 @@ def update_quantize_metadata( model: nn.Module, config: QuantizeConfig, metadata: MetadataDict ) -> None: """Update the quantizer state in the metadata dict.""" + # Defer the shared plugin import to avoid the plugin/conversion initialization cycle. + from .plugins.linear_attention import _linear_attention_state + + policies = _linear_attention_state(model) + if policies or getattr(config, "linear_attention", []): + metadata["linear_attention"] = policies + else: + metadata.pop("linear_attention", None) metadata["quantizer_state"] = quantizer_state(model) if shared_state_metadata := SharedWeightGlobalAmaxState.metadata(model): metadata["shared_quant_states"] = shared_state_metadata @@ -210,22 +231,6 @@ def update_quantize_metadata( metadata.pop("shared_quant_states", None) -def _linear_attention_modules(model): - # Optional framework plugins import conversion; defer this import to avoid that cycle. - from .plugins.gated_delta_net import GatedDeltaNetStateQuantMixin - - return { - get_unwrapped_name(name, model): module - for name, module in model.named_modules() - if isinstance(module, GatedDeltaNetStateQuantMixin) - } - - -def _validate_linear_attention_quantizers(model): - for module in _linear_attention_modules(model).values(): - module.validate_linear_attention() - - def quantizer_state(model: nn.Module) -> dict[str, Any]: """Returns the quantizer state dict describing the quantizer states in the model.""" return { diff --git a/modelopt/torch/quantization/linear_attention/__init__.py b/modelopt/torch/quantization/linear_attention/__init__.py index 477eeb89847..56e50a5174b 100644 --- a/modelopt/torch/quantization/linear_attention/__init__.py +++ b/modelopt/torch/quantization/linear_attention/__init__.py @@ -13,4 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Helpers for linear-attention quantization.""" +"""Numerical policies and differentiable kernels for linear attention.""" + +from .config import * +from .gdn import * +from .kda import * +from .training import * diff --git a/modelopt/torch/quantization/linear_attention/_vllm_autograd.py b/modelopt/torch/quantization/linear_attention/_vllm_autograd.py new file mode 100644 index 00000000000..67169cf5847 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/_vllm_autograd.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differentiable adapter for the installed serving runtime's arithmetic. + +Native kernels supply forward values. Prefill prepares only the intermediates for FLA backward; +decode uses an explicit recurrence adjoint. QDQ uses identity STE. +The optional vLLM dependency supplies kernels; training owns its state and needs no server. +""" + +import torch +from torch.autograd.function import once_differentiable + + +class _Prefill(torch.autograd.Function): + """Run native prefill once and reconstruct FLA's backward inputs on demand.""" + + @staticmethod + def forward(ctx, q, k, v, g, beta, state, scale, beta_dtype, normalize, gate_inputs, precision): + # Keep optional inference kernels isolated to the serving profile. + from ...kernels.quantization.linear_attention.serving.forward import prefill + + q, k, v = (x.to(torch.bfloat16) for x in (q, k, v)) + out, final = prefill( + q, + k, + v, + g.float(), + beta.to(beta_dtype), + state.float(), + scale, + normalize, + gate_inputs, + precision=precision, + ) + ctx.save_for_backward(q, k, v, g, beta, state) + ctx.scale, ctx.normalize = scale, normalize + ctx.set_materialize_grads(False) + return out.float(), final + + @staticmethod + @once_differentiable + def backward(ctx, grad_output, grad_final): + # FLA is optional and needed only for prefix gradients, never native inference. + from fla.modules.l2norm import l2norm_bwd, l2norm_fwd + from fla.ops.utils import chunk_local_cumsum + from fla.ops.utils.constant import RCP_LN2 + + q, k, v, g, beta, state = (x[None].contiguous() for x in ctx.saved_tensors) + if ctx.normalize: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + cumulative = chunk_local_cumsum(g, chunk_size=64, scale=RCP_LN2) + kwargs = { + "q": q, + "k": k, + "v": v, + "g": cumulative, + "beta": beta, + "scale": ctx.scale, + "initial_state": state, + "do": torch.zeros_like(v) + if grad_output is None + else grad_output[None].to(v).contiguous(), + "dht": None if grad_final is None else grad_final[None].contiguous(), + } + # Rebuild only intra-chunk factors. The upstream backward already reconstructs + # chunk states; a second full prefill would also compute unused outputs/states. + if g.ndim == 4: + from fla.ops.kda.chunk_bwd import chunk_kda_bwd + from fla.ops.kda.chunk_intra import chunk_kda_fwd_intra + + aqk, akk = chunk_kda_fwd_intra(q, k, v, cumulative, beta, ctx.scale)[-2:] + dq, dk, dv, db, dg, dh, *_ = chunk_kda_bwd(**kwargs, Aqk=aqk, Akk=akk) + else: + from fla.ops.gated_delta_rule.chunk import chunk_gated_delta_rule_bwd + from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra + + a = chunk_gated_delta_rule_fwd_intra(k, v, cumulative, beta)[-1] + dq, dk, dv, db, dg, dh, *_ = chunk_gated_delta_rule_bwd(**kwargs, A=a) + if ctx.normalize: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + # Match FLA's BF16 operand-gradient rounding before upstream FP32 casts. + grads = tuple( + dx[0].to(x.dtype) for dx, x in zip((dq, dk, dv, dg, db, dh), (q, k, v, g, beta, state)) + ) + return (*grads, None, None, None, None, None) + + +def prefix( + q, k, v, g, beta, state, scale, beta_dtype, normalize=False, gate_inputs=None, precision="vllm" +): + """Return native prefix output/state with an upstream FLA recurrence backward.""" + return _Prefill.apply( + q, k, v, g, beta, state, scale, beta_dtype, normalize, gate_inputs, precision + ) + + +class _DecodeStep(torch.autograd.Function): + """Run native decode once and differentiate the recurrence only during backward.""" + + @staticmethod + def forward(ctx, q, k, v, gate, beta, state, scale, normalize, gate_inputs, precision): + # Autograd Function.forward runs without recording the inference kernel's operations. + from ...kernels.quantization.linear_attention.serving.forward import step as native_step + + output, final = native_step( + q.to(torch.bfloat16), + k.to(torch.bfloat16), + v.to(torch.bfloat16), + gate.float(), + beta.float(), + state.float(), + scale, + normalize, + gate_inputs, + precision=precision, + ) + ctx.save_for_backward(q, k, v, gate, beta, state, final) + ctx.scale, ctx.normalize = scale, normalize + ctx.set_materialize_grads(False) + return output.float(), final + + @staticmethod + @once_differentiable + def backward(ctx, grad_output, grad_final): + if grad_output is None and grad_final is None: + return (None,) * 10 + q, k, v, gate, beta, state, final = ctx.saved_tensors + if ctx.normalize: + q_norm = (q.square().sum(-1, keepdim=True) + 1e-6).sqrt() + k_norm = (k.square().sum(-1, keepdim=True) + 1e-6).sqrt() + q, k = q / q_norm, k / k_norm + decay = gate.exp().unsqueeze(-1) + if gate.ndim == 1: + decay = decay.unsqueeze(-1) + decayed = state * decay + residual = v - (decayed * k.unsqueeze(-1)).sum(-2) + + # Output reads the native working state before the external QDQ cache write. + grad_q, grad_working = None, grad_final + if grad_output is not None: + grad_q = (final * grad_output.unsqueeze(-2)).sum(-1) * ctx.scale + readout_grad = (q * ctx.scale).unsqueeze(-1) * grad_output.unsqueeze(-2) + grad_working = readout_grad if grad_final is None else grad_final + readout_grad + grad_update = (grad_working * k.unsqueeze(-1)).sum(-2) + grad_v = grad_update * beta.unsqueeze(-1) + grad_beta = (grad_update * residual).sum(-1) + grad_k = (grad_working * (residual * beta.unsqueeze(-1)).unsqueeze(-2)).sum(-1) + grad_k = grad_k - (decayed * grad_v.unsqueeze(-2)).sum(-1) + grad_decayed = grad_working - k.unsqueeze(-1) * grad_v.unsqueeze(-2) + grad_state = grad_decayed * decay + grad_gate = (grad_decayed * decayed).sum(-1) + if gate.ndim == 1: + grad_gate = grad_gate.sum(-1) + if ctx.normalize: + if grad_q is not None: + grad_q = (grad_q - q * (grad_q * q).sum(-1, keepdim=True)) / q_norm + grad_k = (grad_k - k * (grad_k * k).sum(-1, keepdim=True)) / k_norm + return grad_q, grad_k, grad_v, grad_gate, grad_beta, grad_state, None, None, None, None + + +def step(q, k, v, gate, beta, state, scale, normalize=False, gate_inputs=None, precision="vllm"): + """Return native output/state with a recurrence backward; the caller owns state QDQ.""" + return _DecodeStep.apply(q, k, v, gate, beta, state, scale, normalize, gate_inputs, precision) diff --git a/modelopt/torch/quantization/linear_attention/config.py b/modelopt/torch/quantization/linear_attention/config.py new file mode 100644 index 00000000000..09c07c44485 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/config.py @@ -0,0 +1,66 @@ +# 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/KDA state QAT.""" + +from typing import Literal + +from pydantic import Field + +from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField + +__all__ = ["LinearAttentionConfig", "LinearAttentionPolicyEntry"] + + +class LinearAttentionConfig(ModeloptBaseConfig): + """One execution policy shared by the chunked prefix and recurrent suffix. + + ``serving`` uses native forward arithmetic with a differentiable adjoint; + ``fla`` is the disabled/default path. ``vllm`` targets vLLM 0.30.0's + Triton prefill and single-token, non-speculative decode with BF16 operands and + FP32 recurrent storage. Deployment must explicitly select Triton prefill + (``gdn_prefill_backend="triton"``) and enable packed recurrent decode + (``VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE=1``). Mixed prefill/decode + batches and automatic FlashInfer/CuTe dispatch are outside this profile. + KDA uses the standalone FLA/Triton API, with no model-level serving guarantee. + ``vllm_kimi_k3`` explicitly selects + vLLM's Kimi-K3 Triton kernels (also used by Kimi-Linear in vLLM 0.30), + raw projections and that model package's dependencies. + These targets are saved in ModelOpt/Megatron policy metadata, not HF export; + they do not configure or validate a remote server. Kimi serving must set both + ``kda_prefill_backend`` and ``kda_decode_backend`` to ``triton``; + the supported gate is unbounded softplus. + Fresh prefill has no internal state QDQ; incoming continuation state and + the handoff to a nonempty suffix use the selected checkpoint encoding. + + TensorQuantizer enables QDQ and owns its format and grouping. + ``state_block_v`` controls grouping only for legacy per-tile quantizers. + """ + + backend: Literal["fla", "serving"] = ModeloptField(default="fla") + precision: Literal["vllm", "vllm_kimi_k3"] = ModeloptField(default="vllm") + state_block_v: Literal[16, 32, 64, 128] = ModeloptField(default=64) + + +class LinearAttentionPolicyEntry(ModeloptBaseConfig): + """Assign a complete policy to supported modules matching ``module_name``. + + Rules apply in order: the last match wins, without merging fields. + Unmatched rules are ignored locally, as with quant_cfg, so a shared recipe + can be applied independently to pipeline stages and virtual-pipeline chunks. + """ + + module_name: str = Field(...) + cfg: LinearAttentionConfig = ModeloptField(default=LinearAttentionConfig()) diff --git a/modelopt/torch/quantization/linear_attention/decode.py b/modelopt/torch/quantization/linear_attention/decode.py new file mode 100644 index 00000000000..52a611ad62d --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/decode.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differentiable per-token state quantization for a training suffix.""" + +import torch + +from .utils import _state_qdq + +__all__ = [] + + +def _recurrent_decode( + q, + k, + v, + g, + beta, + *, + config, + state_quantizer, + initial_state, + output_final_state=False, + scale=None, + use_qk_l2norm_in_kernel=False, + gate_inputs=None, +): + """Run one training suffix and optionally return its dense final state. + + The caller validates shapes, policy and quantizer, and supplies FP32 inputs and prefill state. + Empty suffixes preserve that state without a handoff QDQ or native launch. + """ + if len(q) == 0: + # Keep empty input gradients defined without introducing a state write. + zero = (q.sum() + k.sum() + v.sum() + g.sum() + beta.sum()) * 0 + return v + zero, initial_state if output_final_state else None + scale = q.shape[-1] ** -0.5 if scale is None else scale + # Import the optional native backend only for the ordinary state-QDQ path. + from ._vllm_autograd import step + + state = _state_qdq(initial_state, state_quantizer, config.state_block_v) + with torch.autocast(device_type=q.device.type, enabled=False): + outputs = [] + for t in range(len(q)): + # The output reads the working state before QDQ writes the next token's cache. + output, working = step( + q[t], + k[t], + v[t], + g[t], + beta[t], + state, + scale, + use_qk_l2norm_in_kernel, + precision=config.precision, + gate_inputs=(gate_inputs[0][t], gate_inputs[1][t], *gate_inputs[2:]) + if gate_inputs is not None + else None, + ) + outputs.append(output) + # A new tensor preserves the state saved by earlier tokens' backward passes. + state = _state_qdq(working, state_quantizer, config.state_block_v) + return torch.stack(outputs), state if output_final_state else None diff --git a/modelopt/torch/quantization/linear_attention/gdn.py b/modelopt/torch/quantization/linear_attention/gdn.py new file mode 100644 index 00000000000..6fed19cdf13 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/gdn.py @@ -0,0 +1,102 @@ +# 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. + +"""GDN adapter for serving-aligned recurrent-state QAT.""" + +from .config import LinearAttentionConfig +from .training import _prefill_decode_forward, _prepare_prefill_inputs +from .utils import _validate_gate_inputs + +__all__ = ["gdn_state_qat"] + + +def gdn_state_qat( + q, + k, + v, + g, + beta, + *, + policy: LinearAttentionConfig, + state_quantizer=None, + scale=None, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=False, + use_beta_sigmoid_in_kernel=False, + allow_neg_eigval=False, + A_log=None, # noqa: N803 - match the FLA kernel signature + dt_bias=None, + cu_seqlens=None, + cu_seqlens_cpu=None, + state_v_first=False, + chunk_size=64, + cp_context=None, + prefill_lengths=None, + sequence_lengths=None, + gate_inputs=None, +): + """Adapt Megatron's FLA-style GDN call to serving-aligned state QAT. + + Validate prepared scalar log gates and promote inputs to FP32 working values. + The shared training forward runs the chunked prefix and recurrent suffix with + native forward values and configured state QDQ. FLA supplies the prefix backward; + a Torch recurrence supplies the suffix adjoint. ``state_quantizer`` owns state + QDQ; passing None or a disabled quantizer retains the unquantized native path. + ``precision="vllm"`` requires ``gate_inputs=(raw_g, raw_beta, A_log, bias)`` + for native gate preparation; the prepared ``g`` and ``beta`` supply the adjoint. + """ + if policy.precision == "vllm_kimi_k3": + raise ValueError("vllm_kimi_k3 is a KDA-only precision profile") + if policy.precision == "vllm" and gate_inputs is None: + raise ValueError("vllm GDN requires raw gate_inputs=(raw_g, raw_beta, A_log, bias)") + if cp_context is not None: + raise NotImplementedError("GDN state QAT does not support context parallelism") + _validate_gate_inputs(g, beta, gate_inputs) + output_dtype = q.dtype + beta_dtype = beta.dtype + q, k, v, g, beta = _prepare_prefill_inputs( + q, k, v, g, beta, policy=policy, chunk_size=chunk_size + ) + if allow_neg_eigval: + raise ValueError("Serving GDN does not support allow_neg_eigval") + if use_gate_in_kernel or use_beta_sigmoid_in_kernel: + raise ValueError( + "Serving GDN expects prepared log gates and beta from the Megatron adapter" + ) + if g.ndim != 3: + raise ValueError("gdn_state_qat requires scalar GDN log gates") + return _prefill_decode_forward( + q, + k, + v, + g, + beta, + policy=policy, + state_quantizer=state_quantizer, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + state_v_first=state_v_first, + output_dtype=output_dtype, + beta_dtype=beta_dtype, + prefill_lengths=prefill_lengths, + sequence_lengths=sequence_lengths, + gate_inputs=gate_inputs, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) diff --git a/modelopt/torch/quantization/linear_attention/kda.py b/modelopt/torch/quantization/linear_attention/kda.py new file mode 100644 index 00000000000..d16ebd69bfa --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/kda.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. + +"""KDA adapter for serving-aligned recurrent-state QAT.""" + +import torch +import torch.nn.functional as F + +from .training import _prefill_decode_forward, _prepare_prefill_inputs +from .utils import _validate_gate_inputs, forward_value + +__all__ = ["kda_state_qat"] + + +def kda_state_qat( + q, + k, + v, + g, + beta, + *, + policy, + state_quantizer=None, + scale=None, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=False, + use_beta_sigmoid_in_kernel=False, + allow_neg_eigval=False, + A_log=None, # noqa: N803 - match the FLA kernel signature + dt_bias=None, + safe_gate=False, + lower_bound=None, + cu_seqlens=None, + cu_seqlens_cpu=None, + state_v_first=False, + chunk_size=64, + cp_context=None, + disable_recompute=False, + return_intermediate_states=False, + prefill_lengths=None, + sequence_lengths=None, + gate_inputs=None, +): + """Adapt Megatron's FLA-style KDA call to serving-aligned state QAT. + + Prepare per-key-channel gates and optional beta activation using FLA's formula. + The shared training forward runs the chunked prefix and recurrent suffix with + native forward values and configured state QDQ. FLA supplies the prefix backward; + a Torch recurrence supplies the suffix adjoint. ``state_quantizer`` owns state + QDQ; passing None or a disabled quantizer retains the unquantized native path. + """ + if cp_context is not None or disable_recompute or return_intermediate_states: + raise NotImplementedError( + "KDA state QAT does not support CP or FLA recompute/intermediate flags" + ) + if allow_neg_eigval and (not use_beta_sigmoid_in_kernel or gate_inputs is not None): + raise ValueError("allow_neg_eigval requires sigmoid beta without raw gate_inputs") + if lower_bound is not None or safe_gate: + raise ValueError("Serving arithmetic uses the native softplus KDA gate") + _validate_gate_inputs(g, beta, gate_inputs) + output_dtype = q.dtype + beta_dtype = beta.dtype + q, k, v, g, beta = _prepare_prefill_inputs( + q, k, v, g, beta, policy=policy, chunk_size=chunk_size + ) + dtype = q.dtype + if g.ndim != 4: + raise ValueError("kda_state_qat requires per-key-channel KDA log gates") + if use_gate_in_kernel: + raw_gate = g + if A_log is None: + raise ValueError("Fused KDA gate requires A_log") + if dt_bias is not None: + g = g + dt_bias.to(dtype).reshape(g.shape[-2:]) + rate = A_log.to(dtype).exp().reshape(g.shape[-2], 1) + g = -rate * F.softplus(g) + # Import the optional vLLM backend only for the native precision profile. + from ...kernels.quantization.linear_attention.serving.forward import _kernels + + with torch.no_grad(): + _, native_kda = _kernels("vllm") + native_gate = native_kda.fused_kda_gate( + raw_gate.flatten(-2).contiguous(), A_log, raw_gate.shape[-1], g_bias=dt_bias + ) + g = forward_value(g, native_gate) + if use_beta_sigmoid_in_kernel: + if use_gate_in_kernel and gate_inputs is None and not allow_neg_eigval: + gate_inputs = (raw_gate, beta, A_log, dt_bias) + beta = beta.sigmoid() * (2.0 if allow_neg_eigval else 1.0) + beta_dtype = torch.float32 + return _prefill_decode_forward( + q, + k, + v, + g, + beta, + policy=policy, + state_quantizer=state_quantizer, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + state_v_first=state_v_first, + output_dtype=output_dtype, + beta_dtype=beta_dtype, + prefill_lengths=prefill_lengths, + sequence_lengths=sequence_lengths, + gate_inputs=gate_inputs, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) diff --git a/modelopt/torch/quantization/linear_attention/training.py b/modelopt/torch/quantization/linear_attention/training.py new file mode 100644 index 00000000000..011ecea2753 --- /dev/null +++ b/modelopt/torch/quantization/linear_attention/training.py @@ -0,0 +1,335 @@ +# 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. + +"""Training forwards with a chunked prefill prefix and recurrent decode suffix.""" + +from contextlib import contextmanager +from functools import cache +from importlib.metadata import version +from itertools import pairwise +from numbers import Integral + +import torch +from packaging.version import Version + +from .decode import _recurrent_decode +from .utils import _resolve_state_quantizer, _state_qdq + +__all__ = ["get_linear_attention_layers", "linear_attention_training_phase"] + + +def _lengths(values, name): + if isinstance(values, torch.Tensor): + values = values.tolist() + values = tuple(values) + if any( + isinstance(value, bool) or not isinstance(value, Integral) or value < 0 for value in values + ): + raise ValueError(f"{name} must be nonnegative integers") + return tuple(int(value) for value in values) + + +class _PackedBoundaries: + """Own immutable host boundaries and verify each layer's current packed tensor.""" + + def __init__(self, values): + self.values = None if values is None else _lengths(values, "Packed boundaries") + + def resolve(self, actual): + if actual is None: + if self.values is not None: + raise ValueError("Phase cu_seqlens requires packed boundaries from the layer") + return None + # Tensor versions do not track writes through .data, NumPy or DLPack aliases. + values = _lengths(actual, "Packed boundaries") + if self.values is not None and values != self.values: + raise ValueError("Phase cu_seqlens does not match the layer's packed boundaries") + self.values = values + return values + + +def get_linear_attention_layers(model, *, enabled_only=False, unphased_only=False): + """Return serving-policy layers from a module or iterable of model chunks. + + Disabled state quantizers are included by default so a baseline can use the + same serving arithmetic and phase as QAT. Use ``enabled_only`` for setup checks + and ``unphased_only`` to select layers needing adapter-provided defaults. + """ + # The plugin imports this numerical package during quantization initialization. + from ..plugins.linear_attention import _LinearAttentionQuantMixin + + roots = (model,) if isinstance(model, torch.nn.Module) else model + return tuple( + dict.fromkeys( + module + for root in roots + for module in root.modules() + if isinstance(module, _LinearAttentionQuantMixin) + and module.linear_attention_config.backend == "serving" + and (not enabled_only or module._linear_attn_state.is_enabled) + and (not unphased_only or module._linear_attention_prefill_lengths is None) + ) + ) + + +def _prefix_lengths(values, valid_lengths): + prefixes = _lengths(values, "Prefill lengths") + if valid_lengths is not None and ( + len(prefixes) != len(valid_lengths) + or any(prefix > end for prefix, end in zip(prefixes, valid_lengths)) + ): + raise ValueError("Supply one prefill length <= sequence length per sequence") + return prefixes + + +@contextmanager +def linear_attention_training_phase( + model, prefill_lengths, *, sequence_lengths=None, cu_seqlens=None, preserve_existing=False +): + """Supply per-sequence phases to a module or iterable of model chunks/layers. + + Enabled state quantization requires a phase, including calibration/eval. Disabled + quantizers use the original module outside a phase and serving arithmetic inside. + Megatron selective recompute captures the phase for backward after context exit. + + ``sequence_lengths`` excludes right padding in dense rows or padded THD segments. + ``cu_seqlens`` describes physical packed storage, including alignment padding. + Megatron derives omitted valid lengths from unpadded ``cu_seqlens_q`` metadata; + padded-only metadata requires explicit ``sequence_lengths``. + Padding produces zero attention outputs and never updates the recurrent state. + ``preserve_existing=True`` installs defaults only on layers without a phase. + Empty pipeline stages participate as no-ops. + """ + layers = get_linear_attention_layers(model) + if preserve_existing: + layers = tuple(m for m in layers if m._linear_attention_prefill_lengths is None) + valid_lengths = ( + None if sequence_lengths is None else _lengths(sequence_lengths, "Sequence lengths") + ) + lengths = _prefix_lengths(prefill_lengths, valid_lengths) + boundaries = _PackedBoundaries(cu_seqlens) + previous = [ + ( + m._linear_attention_prefill_lengths, + m._linear_attention_sequence_lengths, + m._linear_attention_cu_seqlens, + ) + for m in layers + ] + try: + for module in layers: + module._linear_attention_prefill_lengths = lengths + module._linear_attention_sequence_lengths = valid_lengths + module._linear_attention_cu_seqlens = boundaries + yield model + finally: + for module, (original, original_lengths, original_boundaries) in zip(layers, previous): + module._linear_attention_prefill_lengths = original + module._linear_attention_sequence_lengths = original_lengths + module._linear_attention_cu_seqlens = original_boundaries + + +@cache +def _validate_training_device(device): + if torch.cuda.get_device_capability(device)[0] >= 12 and Version(version("triton")) < Version( + "3.7" + ): + raise NotImplementedError( + "Linear-attention state QAT is unsupported on sm_120+ with Triton < 3.7 " + "due to a CUDA misaligned-address fault. Use a vLLM build with Triton >= 3.7." + ) + + +def _prepare_prefill_inputs(q, k, v, g, beta, *, policy, chunk_size): + """Validate the native policy and prepare FP32 values for the training adjoint.""" + if policy.backend != "serving" or chunk_size != 64: + raise ValueError("State training requires backend='serving' and chunk_size=64") + if q.device.type != "cuda" or any(x.dtype != torch.bfloat16 for x in (q, k, v)): + raise ValueError("Serving arithmetic requires CUDA BF16 Q/K/V inputs") + _validate_training_device(q.device) + return tuple(x.float() for x in (q, k, v, g, beta)) + + +def _prefill_decode_forward( + q, + k, + v, + g, + beta, + *, + policy, + state_quantizer, + scale, + initial_state, + output_final_state, + cu_seqlens, + cu_seqlens_cpu, + state_v_first, + output_dtype, + beta_dtype, + prefill_lengths, + sequence_lengths=None, + use_qk_l2norm_in_kernel=False, + gate_inputs=None, +): + """Run both prefill and decode phases in one differentiable training forward. + + Each sequence's chunked prefix produces the state for its recurrent suffix. + The suffix preserves its gradient connection to prefill and returns a dense + final state only when requested. Outputs are joined in token order. + + Args: + prefill_lengths: Prefix token count per sequence. For 128 tokens, a value + of 64 selects 64 chunked prefill tokens followed by 64 recurrent tokens. + sequence_lengths: Valid token count per sequence, excluding right padding. + Defaults to the dense row length or each packed storage segment length. + """ + state_quantizer = _resolve_state_quantizer(state_quantizer) + if state_quantizer is not None and prefill_lengths is None: + raise ValueError( + "State quantization requires explicit prefill lengths through " + "linear_attention_training_phase" + ) + if q.ndim != 4 or k.shape != q.shape or v.ndim != 4 or v.shape[:2] != q.shape[:2]: + raise ValueError("q/k and v must have compatible [B,T,H,D] shapes") + batch, length, key_heads, keys = q.shape + heads, values = v.shape[2:] + if batch < 1 or key_heads < 1 or heads % key_heads or beta.shape != (batch, length, heads): + raise ValueError("Invalid batch/head dimensions or beta shape") + if g.shape not in (beta.shape, (*beta.shape, keys)): + raise ValueError("Invalid GDN/KDA log-retention shape") + q, k = (x.repeat_interleave(heads // key_heads, dim=2) for x in (q, k)) + if isinstance(cu_seqlens_cpu, _PackedBoundaries): + boundaries = cu_seqlens_cpu.resolve(cu_seqlens) + elif cu_seqlens_cpu is not None: + boundaries = _PackedBoundaries(cu_seqlens_cpu).resolve(cu_seqlens) + else: + boundaries = cu_seqlens + if boundaries is None: + sequences = [(b, 0, length) for b in range(batch)] + else: + bounds = _lengths(boundaries, "Packed boundaries") + if ( + batch != 1 + or len(bounds) < 2 + or bounds[0] != 0 + or bounds[-1] != length + or any(a > b for a, b in pairwise(bounds)) + ): + raise ValueError( + "Packed boundaries must partition a batch of one, allowing empty entries" + ) + sequences = [(0, a, b) for a, b in pairwise(bounds)] + valid_lengths = ( + tuple(end - start for _, start, end in sequences) + if sequence_lengths is None + else _lengths(sequence_lengths, "Sequence lengths") + ) + if len(valid_lengths) != len(sequences) or any( + valid > end - start for valid, (_, start, end) in zip(valid_lengths, sequences) + ): + raise ValueError("Supply one sequence length within its storage segment per sequence") + prefixes = _prefix_lengths( + valid_lengths if prefill_lengths is None else prefill_lengths, valid_lengths + ) + if ( + keys > 256 + or (g.ndim == 4 and values != keys) + or (use_qk_l2norm_in_kernel and keys & (keys - 1)) + ): + raise ValueError( + "Serving arithmetic requires K <= 256, KDA V=K, and power-of-two K for normalization" + ) + # Validate the optional native backend before launching any prefix kernels. + from ...kernels.quantization.linear_attention.serving.forward import validate_profile + + validate_profile( + gate_inputs, + use_qk_l2norm_in_kernel, + precision=policy.precision, + ) + if initial_state is None: + states = q.new_zeros(len(sequences), heads, keys, values) + else: + states = initial_state.transpose(-1, -2) if state_v_first else initial_state + states = states.to(q.dtype) + if states.shape != (len(sequences), heads, keys, values): + raise ValueError("Initial state shape does not match sequence/head dimensions") + outputs, finals = [], [] + for n, (b, start, storage_end) in enumerate(sequences): + end = start + valid_lengths[n] + split = start + prefixes[n] + prefix, state = q.new_empty(0, heads, values), states[n] + if prefixes[n]: + with torch.autocast(device_type=q.device.type, enabled=False): + from ._vllm_autograd import prefix as serving_prefix + + # A continuation prefill consumes a stored cache just as native serving + # does. A fresh zero-state prefix has no incoming cache to quantize. + if initial_state is not None and state_quantizer is not None: + state = _state_qdq(state, state_quantizer, policy.state_block_v) + prefix, state = serving_prefix( + q=q[b, start:split], + k=k[b, start:split], + v=v[b, start:split], + g=g[b, start:split], + beta=beta[b, start:split], + state=state, + scale=keys**-0.5 if scale is None else scale, + beta_dtype=beta_dtype, + precision=policy.precision, + normalize=use_qk_l2norm_in_kernel, + gate_inputs=( + ( + gate_inputs[0][b, start:split], + gate_inputs[1][b, start:split], + *gate_inputs[2:], + ) + if gate_inputs is not None + else None + ), + ) + # Keep the prefix state attached so suffix losses backpropagate through prefill. + # The suffix applies configured state QDQ at the handoff and recurrent writes. + suffix, final_state = _recurrent_decode( + *(x[b, split:end] for x in (q, k, v, g, beta)), + config=policy, + state_quantizer=state_quantizer, + initial_state=state, + output_final_state=output_final_state, + gate_inputs=( + ( + gate_inputs[0][b, split:end], + gate_inputs[1][b, split:end], + *gate_inputs[2:], + ) + if gate_inputs is not None + else None + ), + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + scale=scale, + ) + # Keep storage offsets for downstream layers while excluding padding from the graph. + pieces = (prefix, suffix) + if storage_end > end: + pieces += (v.new_zeros(storage_end - end, heads, values),) + outputs.append(torch.cat(pieces)) + if output_final_state: + finals.append(final_state) + output = torch.stack(outputs) if boundaries is None else torch.cat(outputs).unsqueeze(0) + final = torch.stack(finals) if output_final_state else None + if final is not None and state_v_first: + final = final.transpose(-1, -2) + return output.to(output_dtype), final diff --git a/modelopt/torch/quantization/linear_attention/utils.py b/modelopt/torch/quantization/linear_attention/utils.py index 4f1e18a174d..1f6f753fe4b 100644 --- a/modelopt/torch/quantization/linear_attention/utils.py +++ b/modelopt/torch/quantization/linear_attention/utils.py @@ -15,19 +15,88 @@ """Shared helpers for linear-attention quantization.""" -from ..nn import TensorQuantizer +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from ..nn import TensorQuantizer __all__ = [] +_STATE_FORMATS: dict[int | tuple[int, int], str] = {(4, 3): "fp8_e4m3", 8: "int8"} + + +class _ForwardValue(torch.autograd.Function): + @staticmethod + def forward(ctx, value, rounded): + return rounded.to(value.dtype) + + @staticmethod + def backward(ctx, grad): + return grad, None + + +def forward_value(value, rounded): + # Unlike value + (rounded - value).detach(), this cannot lose low bits by cancellation. + return _ForwardValue.apply(value, rounded) + + +def _validate_gate_inputs(g, beta, gate_inputs): + """Check raw projection shapes and types before dispatching native kernels.""" + if gate_inputs is None: + return + if not isinstance(gate_inputs, (tuple, list)) or len(gate_inputs) != 4: + raise ValueError("gate_inputs must contain raw gate, raw beta, A_log, and bias") + raw_g, raw_beta, rate, bias = gate_inputs + for value in (raw_g, raw_beta, rate, bias): + if value is not None and ( + not isinstance(value, torch.Tensor) + or value.device != g.device + or value.dtype not in (torch.bfloat16, torch.float32) + ): + raise ValueError("Raw gate inputs must be BF16/FP32 tensors on the gate device") + if raw_g is None or raw_beta is None or raw_g.shape != g.shape or raw_beta.shape != beta.shape: + raise ValueError("Raw gate and beta shapes must match their prepared operands") + heads = beta.shape[-1] + channels = g.shape[-1] if g.ndim == beta.ndim + 1 else 1 + if rate is None or rate.numel() != heads or rate.ndim not in (1, 2): + raise ValueError("Raw gate A_log must contain one rate per value head") + if (bias is None and g.ndim == beta.ndim) or ( + bias is not None and (bias.ndim != 1 or bias.numel() != heads * channels) + ): + raise ValueError("Raw gate bias must contain one value per gate channel") + + +def validate_gdn_quantizer( + quantizer: TensorQuantizer, + *, + name: str, + num_bits: tuple[int | tuple[int, int], ...] = ((4, 3),), + block_sizes: tuple[int, ...] = (), +) -> None: + """Check supported formats and the custom backward's identity STE.""" + # Numerical helpers load while QuantizeConfig initializes; defer the quantizer import. + from ..nn import TensorQuantizer -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 quantizer.block_sizes is not None and quantizer.block_sizes.get("type") == "dynamic": + raise ValueError( + f"{name} does not support block_sizes['type']='dynamic'. For INT8 state QDQ, " + "use top-level type='dynamic' with block_sizes={-1: N}; " + "quantized block scales are not supported." + ) if not ( quantizer._dynamic - and quantizer.num_bits == (4, 3) - and quantizer.block_sizes is None + and quantizer.num_bits in num_bits + and (quantizer.num_bits != 8 or (not quantizer.unsigned and quantizer.narrow_range)) + and ( + quantizer.block_sizes is None + or (quantizer.block_sizes in ({-1: size} for size in block_sizes)) + ) and quantizer.fake_quant and quantizer._pass_through_bwd and not quantizer.rotate_is_enabled @@ -37,7 +106,49 @@ def validate_gdn_quantizer(quantizer: TensorQuantizer, *, name: str) -> None: and not quantizer._use_constant_amax ): raise ValueError( - f"{name} supports only dynamic E4M3 fake quantization with " - "pass_through_bwd=True, no block_sizes, rotation, pre-scaling, bias, constant " + f"{name} supports only dynamic fake quantization with num_bits in {num_bits}, " + f"block_sizes in {block_sizes} (or no blocks), " + "pass_through_bwd=True, no rotation, pre-scaling, bias, constant " "amax, or custom backend. Other gradient rules and formats are not implemented." ) + + +def state_quantizer_config( + quantizer: TensorQuantizer, *, name="state_quantizer" +) -> tuple[str, int]: + """Validate the state quantizer and derive its format and last-axis group size. + + A zero group size retains the legacy policy's full-key state tiles. + """ + validate_gdn_quantizer( + quantizer, name=name, num_bits=tuple(_STATE_FORMATS), block_sizes=(16, 32, 64, 128) + ) + if quantizer.block_sizes is not None: + return _STATE_FORMATS[quantizer.num_bits], quantizer.block_sizes[-1] + if quantizer.axis != (0, 1): + raise ValueError(f"{name} supports only axis=(0, 1) with state_block_v tiling") + return _STATE_FORMATS[quantizer.num_bits], 0 + + +def _resolve_state_quantizer(state_quantizer): + """Validate an enabled quantizer once; return None when state QDQ is inactive.""" + if state_quantizer is not None and state_quantizer.is_enabled: + state_quantizer_config(state_quantizer) + if state_quantizer._if_quant: + return state_quantizer + return None + + +def _state_qdq(state: torch.Tensor, state_quantizer: TensorQuantizer | None, block_v: int = 64): + """Apply state QDQ; TensorQuantizer owns the format, scales, and backward rule.""" + if state_quantizer is None: + return state + if state_quantizer.block_sizes is not None: + return state_quantizer(state) + rounded = [] + for part in state.split(block_v, dim=-1): + # Canonical [N, H, K*BV] shape preserves per-head tile scales for both state ranks. + tensor = part.flatten(-2) + inputs = tensor.reshape(1, -1, tensor.shape[-1]) + rounded.append(state_quantizer(inputs).reshape_as(part)) + return torch.cat(rounded, dim=-1).to(state.dtype) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 328b040e18b..f928acc2ad6 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -33,7 +33,6 @@ from modelopt.torch.opt.utils import forward_with_reshard from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.conversion import ( - _validate_linear_attention_quantizers, preserve_quantizer_attributes_context, set_quantizer_attributes_partial, set_quantizer_by_cfg, @@ -235,6 +234,16 @@ def _check_weight_quantization_took_effect(model: nn.Module, config: QuantizeCon # Quantizers that only framework plugins add: (class-name patterns of the modules they belong on, # what those modules are, recipe units listing the supported models). _PLUGIN_QUANTIZERS = { + "gdn_state_quantizer": ( + ("*GatedDeltaNet", "*GatedDeltaNetAttention"), + "GDN layer", + "linear_attention_state_int8*", + ), + "kda_state_quantizer": ( + ("*KimiDeltaAttention", "*KimiK3DeltaAttention", "*KimiLinearDeltaAttention"), + "KDA layer", + "linear_attention_state_int8*", + ), "indexer_q_quantizer": (("*Indexer",), "sparse-attention indexer", "indexer_*_nvfp4"), "indexer_k_quantizer": (("*Indexer",), "sparse-attention indexer", "indexer_*_nvfp4"), "dispatch_quantizer": (("*Experts*", "*FusedMoE"), "MoE experts module", "moe_*_nvfp4"), @@ -245,9 +254,9 @@ def _check_weight_quantization_took_effect(model: nn.Module, config: QuantizeCon def _check_plugin_quantization_took_effect(model: nn.Module, config: QuantizeConfig) -> None: """Raise when a config enables a plugin-only quantizer that quantizes nothing. - The quantizers in ``_PLUGIN_QUANTIZERS`` only exist on modules (sparse-attention indexers, MoE - experts) that a framework plugin (vLLM, Megatron-Core) converted. When a plugin does not - recognize the installed framework's module class, or the patterns match none of the modules, + The quantizers in ``_PLUGIN_QUANTIZERS`` only exist on modules (linear-attention layers, + sparse-attention indexers, MoE experts) that a framework plugin (vLLM, Megatron-Core) converted. + When a plugin does not recognize the installed framework's module class, or the patterns match none of the modules, the tensors silently stay unquantized. Intent is read like in :func:`_check_weight_quantization_took_effect`, but only from patterns naming the quantizer, so catch-all patterns do not count. A model without such a module (another architecture) is @@ -401,9 +410,16 @@ def forward_loop(model) -> None: if not is_quantized(model): model = apply_mode(model, mode=[("quantize", dict(config))], registry=QuantizeModeRegistry) else: + # Defer shared plugin imports to avoid the plugin/conversion initialization cycle. + from .plugins.linear_attention import ( + _apply_linear_attention_policy, + _validate_linear_attention, + ) + # Already quantized, so lets apply the quant_cfg from the config set_quantizer_by_cfg(model, quantize_config.quant_cfg) - _validate_linear_attention_quantizers(model) + _apply_linear_attention_policy(model, quantize_config) + _validate_linear_attention(model) # Fail before calibration rather than after exporting an unquantized checkpoint. _check_weight_quantization_took_effect(model, quantize_config) _check_plugin_quantization_took_effect(model, quantize_config) diff --git a/modelopt/torch/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py index 7234ca81f76..fc56e230838 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -39,7 +39,7 @@ from .attention import * from .custom import * -from .gated_delta_net import * +from .gdn 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 deleted file mode 100644 index df70683cc2f..00000000000 --- a/modelopt/torch/quantization/plugins/gated_delta_net.py +++ /dev/null @@ -1,128 +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. - -"""Fake quantization of the GatedDeltaNet (GDN) recurrent state. - -The chunked gated-delta-rule kernel keeps each head's ``[K, V]`` recurrent state in fp32 inside -one Triton launch and carries it from chunk to chunk. To emulate a deployment that stores that -state in FP8, ModelOpt runs an adapted copy of the kernel -(:mod:`modelopt.torch.kernels.quantization.linear_attention`) that fake-quantizes the state to -E4M3 at the end of every chunk, with a scale computed inside the kernel from the state itself. -The backward pass recomputes the same quantized states and passes the state gradient straight -through the quantization, so QAT and QAD train against the quantized recurrence. A second -quantizer covers ``w``, the WY-transformed keys that multiply the state; ``w`` is a regular tensor, -so it is fake-quantized by the ``TensorQuantizer`` itself before the kernel reads it. -""" - -from collections.abc import Callable -from functools import partial -from typing import Any - -import torch - -from ..config import QuantizerAttributeConfig -from ..linear_attention.utils import validate_gdn_quantizer -from ..nn import QuantModule, TensorQuantizer - -__all__ = ["GatedDeltaNetStateQuantMixin"] - -GatedDeltaRuleFn = Callable[..., tuple[torch.Tensor, torch.Tensor | None]] - - -def _fla_chunk_gated_delta_rule() -> GatedDeltaRuleFn: - # FLA is an optional, heavy dependency needed only when quantization is enabled. - from fla.ops.gated_delta_rule import chunk_gated_delta_rule - - return chunk_gated_delta_rule - - -def _state_qdq_chunk_gated_delta_rule() -> GatedDeltaRuleFn: - # Imported on first use: flash-linear-attention is a heavy optional dependency that only the - # enabled quantizer needs, and importing it warns on machines without a GPU. - try: - from modelopt.torch.kernels.quantization.linear_attention.fla_chunk_gated_delta_rule import ( - chunk_gated_delta_rule, - ) - except ImportError as e: - raise RuntimeError( - "GDN fake quantization needs Triton and fla-core==0.5.1 on a CUDA " - f"device; importing the state-quantizing kernel failed with {e!r}." - ) from e - return chunk_gated_delta_rule - - -class GatedDeltaNetStateQuantMixin(QuantModule): - """Adds ``gdn_state_quantizer`` and ``gdn_w_quantizer`` to a GatedDeltaNet module. - - Subclasses route the module's chunked gated-delta-rule call through - :meth:`_state_quantized_chunk_gated_delta_rule`. Both quantizers start disabled; enable them - with ``quant_cfg`` entries on ``*gdn_state_quantizer`` / ``*gdn_w_quantizer`` such as the - ``configs/ptq/units/gdn_state_fp8_dynamic`` and ``gdn_w_fp8_dynamic`` recipe units. The state - quantizer carries the fused QDQ configuration. Both sites currently require dynamic - E4M3 and identity STE. State QDQ uses fixed 64-column value tiles and 64-token chunks. - """ - - def _setup(self): - for name in ("gdn_state_quantizer", "gdn_w_quantizer"): - self._register_temp_attribute( - name, TensorQuantizer(QuantizerAttributeConfig(enable=False)) - ) - - def validate_linear_attention(self) -> None: - """Reject numerical settings that the fused training path cannot implement.""" - for name in ("gdn_state_quantizer", "gdn_w_quantizer"): - quantizer = getattr(self, name) - if quantizer.is_enabled: - validate_gdn_quantizer(quantizer, name=name) - # The state handle configures fused QDQ; W's grouping is executed by TensorQuantizer. - if self.gdn_state_quantizer.is_enabled and self.gdn_state_quantizer.axis != (0, 1): - raise ValueError( - "gdn_state_quantizer supports only axis=(0, 1) with 64-column value tiling" - ) - - def modelopt_post_restore(self, prefix: str = ""): - """Validate restored quantizers against the fused training kernel requirements.""" - super().modelopt_post_restore(prefix) - self.validate_linear_attention() - - def _state_quantized_chunk_gated_delta_rule( - self, gated_delta_rule: GatedDeltaRuleFn, *args: Any, **kwargs: Any - ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Call ``gated_delta_rule`` or, if a quantizer is on, the vendored quantizing copy.""" - self.validate_linear_attention() - quantize_state = self.gdn_state_quantizer.is_enabled and self.gdn_state_quantizer._if_quant - quantize_w = self.gdn_w_quantizer.is_enabled - if not (quantize_state or quantize_w): - return gated_delta_rule(*args, **kwargs) - while isinstance(gated_delta_rule, partial): - args = (*gated_delta_rule.args, *args) - kwargs = {**gated_delta_rule.keywords, **kwargs} - gated_delta_rule = gated_delta_rule.func - if gated_delta_rule is not _fla_chunk_gated_delta_rule(): - raise NotImplementedError( - "GatedDeltaNet quantization supports only FLA's chunk_gated_delta_rule " - f"callable or a functools.partial of it; got {gated_delta_rule!r}." - ) - chunk_size = kwargs.pop("chunk_size", 64) - if chunk_size != 64: - raise ValueError("GDN fake quantization supports only chunk_size=64") - return _state_qdq_chunk_gated_delta_rule()( - *args, - chunk_size=chunk_size, - state_qdq=int(quantize_state), - state_qdq_block_v=64, - w_quantizer=self.gdn_w_quantizer if quantize_w else None, - **kwargs, - ) diff --git a/modelopt/torch/quantization/plugins/gdn.py b/modelopt/torch/quantization/plugins/gdn.py new file mode 100644 index 00000000000..9eea217fa50 --- /dev/null +++ b/modelopt/torch/quantization/plugins/gdn.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GatedDeltaNet recurrent-state fake quantization. + +State QAT uses an explicit prefix/suffix policy with native serving arithmetic. +""" + +from collections.abc import Callable +from functools import partial +from typing import Any + +import torch + +from ..linear_attention.gdn import gdn_state_qat +from .linear_attention import _LinearAttentionQuantMixin + +__all__ = ["GatedDeltaNetStateQuantMixin"] + +GatedDeltaRuleFn = Callable[..., tuple[torch.Tensor, torch.Tensor | None]] + + +def _fla_chunk_gated_delta_rule() -> GatedDeltaRuleFn: + # FLA is an optional, heavy dependency needed only when quantization is enabled. + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + + return chunk_gated_delta_rule + + +class GatedDeltaNetStateQuantMixin(_LinearAttentionQuantMixin): + """Adds recurrent-state fake quantization to a GatedDeltaNet module. + + Subclasses route the module's chunked gated-delta-rule call through + :meth:`_state_quantized_chunk_gated_delta_rule`. Enable ``*gdn_state_quantizer`` + through ``quant_cfg`` and select an explicit serving policy. State supports dynamic + E4M3 or signed narrow-range INT8 with identity STE. The execution policy is saved in + ModelOpt metadata. ``gdn_w_quantizer`` is a disabled legacy checkpoint handle; + enabling WY operand quantization is no longer supported. + """ + + linear_attention_quantizer_names = ("gdn_state_quantizer", "gdn_w_quantizer") + + def validate_linear_attention(self): + """Reject retired W QAT and validate the state quantizer.""" + if self.gdn_w_quantizer.is_enabled: + raise ValueError( + "GDN W quantization is no longer supported. Disable gdn_w_quantizer; " + "use gdn_state_quantizer with an explicit serving policy for state QAT." + ) + if self.linear_attention_config.precision == "vllm_kimi_k3": + raise ValueError("vllm_kimi_k3 is a KDA-only precision profile") + super().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]: + """Route enabled state quantization to its training implementation.""" + if ( + self.gdn_w_quantizer.is_enabled + or self.linear_attention_config.precision == "vllm_kimi_k3" + ): + self.validate_linear_attention() + if not self.linear_attention_is_enabled: + return gated_delta_rule(*args, **kwargs) + while isinstance(gated_delta_rule, partial): + args = (*gated_delta_rule.args, *args) + kwargs = {**gated_delta_rule.keywords, **kwargs} + gated_delta_rule = gated_delta_rule.func + if gated_delta_rule is not _fla_chunk_gated_delta_rule(): + raise NotImplementedError( + "GatedDeltaNet quantization supports only FLA's chunk_gated_delta_rule " + f"callable or a functools.partial of it; got {gated_delta_rule!r}." + ) + chunk_size = kwargs.pop("chunk_size", 64) + if chunk_size != 64: + raise ValueError("GDN fake quantization supports only chunk_size=64") + if self._linear_attention_cu_seqlens is not None: + kwargs["cu_seqlens_cpu"] = self._linear_attention_cu_seqlens + return gdn_state_qat( + *args, + policy=self.linear_attention_config, + state_quantizer=self.gdn_state_quantizer, + chunk_size=chunk_size, + prefill_lengths=self._linear_attention_prefill_lengths, + sequence_lengths=self._linear_attention_sequence_lengths, + **kwargs, + ) diff --git a/modelopt/torch/quantization/plugins/kda.py b/modelopt/torch/quantization/plugins/kda.py new file mode 100644 index 00000000000..9c9acce694b --- /dev/null +++ b/modelopt/torch/quantization/plugins/kda.py @@ -0,0 +1,52 @@ +# 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. + +"""Kernel routing for Megatron Kimi Delta Attention quantization.""" + +from functools import partial + +from ..linear_attention.kda import kda_state_qat +from .linear_attention import _LinearAttentionQuantMixin + +__all__ = ["KimiDeltaAttentionStateQuantMixin"] + + +class KimiDeltaAttentionStateQuantMixin(_LinearAttentionQuantMixin): + """Adds state quantizers and decode-aware kernel routing to Kimi Delta Attention.""" + + linear_attention_quantizer_names = ("kda_state_quantizer",) + + def _state_quantized_chunk_kda(self, kernel, *args, **kwargs): + if not self.linear_attention_is_enabled: + return kernel(*args, **kwargs) + # FLA kernels are an optional dependency; the training layer belongs to Megatron. + from fla.ops.kda import chunk_kda + + while isinstance(kernel, partial): + args = (*kernel.args, *args) + kwargs = {**kernel.keywords, **kwargs} + kernel = kernel.func + if kernel is not chunk_kda: + raise NotImplementedError("KDA quantization requires FLA's chunk_kda callable") + if self._linear_attention_cu_seqlens is not None: + kwargs["cu_seqlens_cpu"] = self._linear_attention_cu_seqlens + return kda_state_qat( + *args, + policy=self.linear_attention_config, + state_quantizer=self.kda_state_quantizer, + prefill_lengths=self._linear_attention_prefill_lengths, + sequence_lengths=self._linear_attention_sequence_lengths, + **kwargs, + ) diff --git a/modelopt/torch/quantization/plugins/linear_attention.py b/modelopt/torch/quantization/plugins/linear_attention.py new file mode 100644 index 00000000000..726384b2fbc --- /dev/null +++ b/modelopt/torch/quantization/plugins/linear_attention.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared module policy and checkpoint support for linear-attention QAT.""" + +import fnmatch +import warnings + +from modelopt.torch.opt.conversion import ApplyModeError +from modelopt.torch.utils import get_unwrapped_name + +from ..config import QuantizerAttributeConfig +from ..linear_attention.config import LinearAttentionConfig +from ..linear_attention.utils import state_quantizer_config +from ..nn import QuantModule, TensorQuantizer + +__all__ = [] + + +def _linear_attention_modules(model): + return { + get_unwrapped_name(name, model): module + for name, module in model.named_modules() + if isinstance(module, _LinearAttentionQuantMixin) + } + + +def _apply_linear_attention_policy(model, config): + """Apply rules locally, like quant_cfg; a PP/VPP chunk may have no matches.""" + modules = _linear_attention_modules(model) + # getattr also handles pickled configs that predate the policy field. + for entry in getattr(config, "linear_attention", []): + matches = [name for name in modules if fnmatch.fnmatch(name, entry.module_name)] + for name in matches: + modules[name].linear_attention_config = entry.cfg.model_copy(deep=True) + + +def _validate_linear_attention(model): + """Validate configured modules after both policy and quantizers are assigned.""" + modules = _linear_attention_modules(model).values() + for module in modules: + module.validate_linear_attention() + module.warn_linear_attention_profile() + module._validate_linear_attention_config() + + +def _linear_attention_state(model): + return { + name: module.linear_attention_config.model_dump() + for name, module in _linear_attention_modules(model).items() + } + + +def _restore_linear_attention_policy(model, saved_policies): + """Restore saved policies; modelopt_post_restore validates the completed state.""" + if saved_policies is None: + return + modules = _linear_attention_modules(model) + 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) + + +def _restore_legacy_linear_attention_quantizers(model, quantizer_state): + """Keep newly introduced, disabled handles loadable from older checkpoints.""" + for name, module in _linear_attention_modules(model).items(): + for handle in module.linear_attention_quantizer_names: + key = f"{name}.{handle}" if name else handle + quantizer = getattr(module, handle) + if key not in quantizer_state and not quantizer.is_enabled: + quantizer_state[key] = quantizer.get_modelopt_state() + + +class _LinearAttentionQuantMixin(QuantModule): + linear_attention_quantizer_names: tuple[str, ...] = () + + def _setup(self): + for name in self.linear_attention_quantizer_names: + self._register_temp_attribute( + name, TensorQuantizer(QuantizerAttributeConfig(enable=False)) + ) + self._register_temp_attribute("linear_attention_config", LinearAttentionConfig()) + self._register_temp_attribute("_linear_attention_prefill_lengths", None) + self._register_temp_attribute("_linear_attention_sequence_lengths", None) + self._register_temp_attribute("_linear_attention_cu_seqlens", None) + + @property + def _linear_attn_state(self): + return getattr(self, self.linear_attention_quantizer_names[0]) + + @property + def linear_attention_is_enabled(self): + """Retain an unquantized serving baseline only inside an explicit phase.""" + return any( + getattr(self, name).is_enabled for name in self.linear_attention_quantizer_names + ) or ( + self.linear_attention_config.backend == "serving" + and self._linear_attention_prefill_lengths is not None + ) + + def validate_linear_attention(self): + """Validate quantizer contracts shared by GDN and KDA.""" + if self._linear_attn_state.is_enabled: + state_quantizer_config( + self._linear_attn_state, + name=self.linear_attention_quantizer_names[0], + ) + if self.linear_attention_config.backend != "serving": + raise ValueError( + "GDN/KDA state QAT requires backend='serving'; use " + "linear_attention_training_phase to select a recurrent training suffix" + ) + + def warn_linear_attention_profile(self): + """Warn at conversion/restore on one TP/DP representative of each owning stage.""" + if ( + not self._linear_attn_state.is_enabled + or "kda_state_quantizer" not in self.linear_attention_quantizer_names + or self.linear_attention_config.precision != "vllm" + ): + return + parallel = getattr(self, "parallel_state", None) + if parallel is not None and any( + group.rank() > 0 + for group in (parallel.tensor_parallel_group, parallel.data_parallel_group) + ): + return + warnings.warn( + "KDA precision='vllm' selects standalone FLA/Triton kernels, not a model " + "serving profile. For Kimi-Linear/Kimi-K3 on vLLM 0.30, select " + "precision='vllm_kimi_k3' and configure the server's matching Triton " + "backends and unbounded gates." + ) + + def _validate_linear_attention_config(self): + """Check framework settings at conversion/forward, without restricting offline restore.""" + + def validate_linear_attention_execution(self): + """Check framework execution restrictions at forward, not conversion or restore.""" + + def modelopt_post_restore(self, prefix=""): + """Validate the restored numerical policy and quantizers.""" + super().modelopt_post_restore(prefix) + self.validate_linear_attention() + self.warn_linear_attention_profile() diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 5e422235e14..a5abb7a3b7e 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -21,7 +21,9 @@ import textwrap import types from contextlib import contextmanager -from functools import cache, partial +from copy import copy +from functools import cache +from itertools import pairwise from typing import Any import megatron.core.parallel_state as mcore_parallel @@ -29,6 +31,7 @@ import megatron.core.transformer.mlp as megatron_mlp import megatron.core.transformer.moe.experts as megatron_moe import torch +from megatron.core import tensor_parallel from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding from megatron.core.extensions.transformer_engine import ( TEColumnParallelGroupedLinear, @@ -58,6 +61,8 @@ from ..algorithms import AutoQuantizeGradientSearcher from ..conversion import maybe_promote_nvfp4_static_quantizer +from ..linear_attention.config import LinearAttentionConfig +from ..linear_attention.training import _lengths, linear_attention_training_phase from ..nn import ( GroupedQuantizer, QuantModule, @@ -70,7 +75,9 @@ from ..utils import sync_moe_expert_amax from ..utils.layerwise_calib import LayerActivationCollector from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear -from .gated_delta_net import GatedDeltaNetStateQuantMixin +from .gdn import GatedDeltaNetStateQuantMixin +from .kda import KimiDeltaAttentionStateQuantMixin +from .linear_attention import _LinearAttentionQuantMixin from .transformer_engine import _QuantTEGroupedLinear, _QuantTELayerNormLinear, _QuantTELinear try: @@ -87,6 +94,13 @@ except ImportError: HAS_GDN = False +try: + from megatron.core.ssm.gated_delta_net import KimiDeltaAttention + + HAS_KDA = True +except ImportError: + HAS_KDA = False + __all__ = [] @@ -210,6 +224,8 @@ def quant_module_get_extra_state(self) -> dict: quantizer_state[name] = module.get_modelopt_state() extra_state["modelopt_quantizer_state"] = quantizer_state + if isinstance(self, _LinearAttentionQuantMixin): + extra_state["modelopt_linear_attention_state"] = self.linear_attention_config.model_dump() # Handle real_quantizer_state and q_tensor_state extra_state.update(real_quant_module_get_extra_state(self)) @@ -288,6 +304,12 @@ def quant_module_set_extra_state(self, state: Any): if state is None or not self.allow_post_restore: return + if isinstance(self, _LinearAttentionQuantMixin) and "modelopt_linear_attention_state" in state: + # Restore the policy before quantizer restoration invokes modelopt_post_restore. + self.linear_attention_config = LinearAttentionConfig( + **state["modelopt_linear_attention_state"] + ) + quantizer_state = state.get("modelopt_quantizer_state", None) if quantizer_state is not None: @@ -1106,63 +1128,250 @@ def set_extra_state(self, state): quant_module_set_extra_state(self, state) -if HAS_GDN: +class _MegatronLinearAttentionMixin(_LinearAttentionQuantMixin): + """Wrap the kernel call in both direct and recomputed Megatron forwards.""" - @QuantModuleRegistry.register({GatedDeltaNet: "megatron_GatedDeltaNet"}) - class _QuantGatedDeltaNet(GatedDeltaNetStateQuantMixin): - """GatedDeltaNet with fake quantization of the recurrent state at kernel chunk boundaries. + # PyTorch requires class-level overrides to save/load ``_extra_state``; + # instance-level ModelOpt callbacks alone are not sufficient for GDN or KDA. + def get_extra_state(self): + return quant_module_get_extra_state(self) - 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 set_extra_state(self, state): + quant_module_set_extra_state(self, state) - # Class-level overrides so torch routes the quantizer state through ``_extra_state`` - # (GatedDeltaNet has none); see _QuantDSAttention. - def get_extra_state(self): - return quant_module_get_extra_state(self) + def _setup(self): + super()._setup() + self._register_temp_attribute("_linear_attention_gate_inputs", None) + self._register_temp_attribute("_linear_attention_recompute", False) + 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 set_extra_state(self, state): - quant_module_set_extra_state(self, state) + @property + def _serving_arithmetic(self): + return ( + self.linear_attention_config.backend == "serving" and self.linear_attention_is_enabled + ) - 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 _prepare_input_for_gated_delta_rule(self, *args, **kwargs): + # Preserve raw BF16 Q/K: prefill stores normalized BF16 operands, whereas + # decode normalizes inside its FP32 update. One shared pre-normalization + # would irreversibly change the suffix state trajectory. + normalize = self.use_qk_l2norm + if self._serving_arithmetic: + self.use_qk_l2norm = False + try: + return super()._prepare_input_for_gated_delta_rule(*args, **kwargs) + finally: + self.use_qk_l2norm = normalize + + @contextmanager + def _quantized_linear_attention_kernel(self): + kernel = self.gated_delta_rule + previous_gates = self._linear_attention_gate_inputs + + def run_kernel(*args, **kwargs): + kwargs["use_qk_l2norm_in_kernel"] = self.use_qk_l2norm + if self._linear_attention_gate_inputs is None: + raise NotImplementedError( + "Serving state QAT requires raw gates captured by Megatron's _compute_gates hook" + ) + kwargs["gate_inputs"] = self._linear_attention_gate_inputs + return self._linear_attention_kernel(kernel, *args, **kwargs) + + self.gated_delta_rule = run_kernel + self._linear_attention_gate_inputs = None + try: + yield + finally: + self.gated_delta_rule = kernel + self._linear_attention_gate_inputs = previous_gates + + def forward(self, *args, **kwargs): + self.validate_linear_attention_execution() + recompute = getattr(self, "recompute_gdn", False) + capture = self._serving_arithmetic and self.training and recompute + previous = self._linear_attention_recompute + self._linear_attention_recompute = capture + if capture: + # The adapter checkpoints the same core, with microbatch-local metadata. + self.recompute_gdn = False + try: + return super().forward(*args, **kwargs) + finally: + self._linear_attention_recompute = previous + if capture: + self.recompute_gdn = recompute + + def _forward_compute(self, hidden_states, *args, **kwargs): + return self._run_linear_attention_compute( + super()._forward_compute, hidden_states, *args, **kwargs + ) + + def forward_pre_attn_and_core_attn(self, hidden_states, *args, **kwargs): + # Core's two-stage schedules can call this entry point without forward(). + self.validate_linear_attention_execution() + return self._run_linear_attention_compute( + super().forward_pre_attn_and_core_attn, hidden_states, *args, **kwargs + ) + + def _run_linear_attention_compute(self, compute, hidden_states, *args, **kwargs): + if not self._serving_arithmetic: + return compute(hidden_states, *args, **kwargs) + lengths = self._linear_attention_prefill_lengths + sequence_lengths = self._linear_attention_sequence_lengths + if lengths is None: + raise ValueError("State quantization requires explicit prefill lengths") + arguments = inspect.signature(compute).bind(hidden_states, *args, **kwargs) + packed = arguments.arguments.get("packed_seq_params") + cu_seqlens = arguments.arguments.get("cu_seqlens_q") + if packed is not None and packed.qkv_format == "thd": + if cu_seqlens is None: + cu_seqlens = packed.cu_seqlens_q_padded + if cu_seqlens is None: + cu_seqlens = packed.cu_seqlens_q + actual = packed.cu_seqlens_q + if actual is None: + if sequence_lengths is None: + raise ValueError( + "Packed state QAT requires valid sequence_lengths or cu_seqlens_q" + ) + else: + real_bounds = _lengths(actual, "Packed boundaries") + valid = tuple(end - start for start, end in pairwise(real_bounds)) + if sequence_lengths is None: + sequence_lengths = valid + elif len(sequence_lengths) != len(valid) or any( + length > end for length, end in zip(sequence_lengths, valid) + ): + raise ValueError("Sequence lengths exceed the packed valid token counts") + boundaries = self._linear_attention_cu_seqlens.resolve(cu_seqlens) + if self._linear_attention_recompute: + # Packed metadata may refer to buffers reused by a later pipeline microbatch. + for name, value in list(arguments.arguments.items()): + if name == "cu_seqlens_q" and value is not None: + arguments.arguments[name] = value.clone() + elif name == "packed_seq_params" and value is not None: + snapshot = copy(value) + for key, tensor in vars(value).items(): + if isinstance(tensor, torch.Tensor): + setattr(snapshot, key, tensor.clone()) + arguments.arguments[name] = snapshot + + # Megatron checkpoint preserves RNG, but does not capture PyTorch autocast. + autocast_enabled = torch.is_autocast_enabled("cuda") + autocast_dtype = torch.get_autocast_dtype("cuda") + + def run(hidden): + with ( + torch.autocast("cuda", enabled=autocast_enabled, dtype=autocast_dtype), + linear_attention_training_phase( + self, lengths, sequence_lengths=sequence_lengths, cu_seqlens=boundaries + ), + self._quantized_linear_attention_kernel(), + ): + return compute(hidden, *arguments.args[1:], **arguments.kwargs) + + if self._linear_attention_recompute: + return tensor_parallel.checkpoint(run, False, hidden_states) + return run(hidden_states) + + def validate_linear_attention_execution(self): + super().validate_linear_attention_execution() + self._validate_linear_attention_config() + active = self.linear_attention_is_enabled + serving = self.linear_attention_config.backend == "serving" and active + if active and not serving: + self.validate_linear_attention() + if self.config.context_parallel_size > 1 and active: + raise NotImplementedError("GDN/KDA QAT does not support Megatron context parallelism.") + if serving and self.training and self.config.recompute_granularity == "full": + raise NotImplementedError( + "State QAT supports selective GDN recompute, not full-layer recompute" + ) + if serving and ( + not hasattr(super(), "_compute_gates") + or not ( + hasattr(super(), "_forward_compute") + or hasattr(super(), "forward_pre_attn_and_core_attn") + ) + ): + raise NotImplementedError( + "Serving state QAT requires Megatron-Core's _compute_gates and either " + "_forward_compute or forward_pre_attn_and_core_attn to preserve raw gates and " + "packed sequence lengths. Upgrade Megatron-Core." ) - @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 + def _validate_linear_attention_config(self): + if self._serving_arithmetic and self.config.deterministic_mode: + raise NotImplementedError("Serving state QAT requires non-deterministic kernels") + if self._serving_arithmetic and ( + not hasattr(super(), "_prepare_input_for_gated_delta_rule") + or getattr(self, "gdn_pre_gated_delta_rule_fusion", False) + ): + raise NotImplementedError( + "Serving arithmetic requires Megatron's unfused input-preparation hook" ) - try: - yield - finally: - self.gated_delta_rule = gated_delta_rule - - def forward(self, *args, **kwargs): - if hasattr(GatedDeltaNet, "forward_pre_attn_and_core_attn"): - return super().forward(*args, **kwargs) - with self._quantized_gdn_kernel(): - return super().forward(*args, **kwargs) - - def forward_pre_attn_and_core_attn(self, *args, **kwargs): - with self._quantized_gdn_kernel(): - return super().forward_pre_attn_and_core_attn(*args, **kwargs) - - def validate_linear_attention(self): - super().validate_linear_attention() - if self.config.context_parallel_size > 1 and ( - self.gdn_state_quantizer.is_enabled or self.gdn_w_quantizer.is_enabled + + +if HAS_GDN: + + @QuantModuleRegistry.register({GatedDeltaNet: "megatron_GatedDeltaNet"}) + class _QuantGatedDeltaNet(_MegatronLinearAttentionMixin, GatedDeltaNetStateQuantMixin): + """Megatron GDN with serving-aligned state fake quantization.""" + + _linear_attention_kernel = ( + GatedDeltaNetStateQuantMixin._state_quantized_chunk_gated_delta_rule + ) + + def _compute_gates(self, a_log, dt_bias, batch, seq_len, *gate_feats): + gate, inputs = super()._compute_gates(a_log, dt_bias, batch, seq_len, *gate_feats) + if self._serving_arithmetic: + raw_beta, raw_gate = gate_feats + # Native prefill/decode own gate activation and its phase-specific rounding. + self._linear_attention_gate_inputs = ( + raw_gate.reshape(batch, seq_len, -1), + raw_beta.reshape(batch, seq_len, -1), + a_log, + dt_bias, + ) + return gate, inputs + + +if HAS_KDA: + + @QuantModuleRegistry.register({KimiDeltaAttention: "megatron_KimiDeltaAttention"}) + class _QuantKimiDeltaAttention( + _MegatronLinearAttentionMixin, KimiDeltaAttentionStateQuantMixin + ): + """Megatron KDA with decode-aware state fake quantization.""" + + _linear_attention_kernel = KimiDeltaAttentionStateQuantMixin._state_quantized_chunk_kda + + def _validate_linear_attention_config(self): + super()._validate_linear_attention_config() + if self._serving_arithmetic and ( + getattr(self.config, "kda_safe_gate", False) + or getattr(self.config, "kda_lower_bound", None) is not None ): - raise NotImplementedError("GDN QAT does not support Megatron context parallelism.") + raise NotImplementedError("Serving KDA state QAT requires unbounded softplus gates") + + def _compute_gates(self, a_log, dt_bias, batch, seq_len, *gate_feats): + gate, inputs = super()._compute_gates(a_log, dt_bias, batch, seq_len, *gate_feats) + if self._serving_arithmetic: + raw_gate, raw_beta = gate_feats + self._linear_attention_gate_inputs = ( + raw_gate.reshape_as(gate), + raw_beta.reshape_as(inputs["beta"]), + a_log, + dt_bias, + ) + return gate, inputs def _is_supported_megatron_model(model: torch.nn.Module) -> bool: diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index f44fbc620ca..e997841ab47 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -33,8 +33,8 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `experts_nvfp4.yaml` | NVFP4 W4A4 on `*.experts.*` weight/input quantizers | | `mixer_mlp_nvfp4.yaml` | NVFP4 W4A4 on dense `*.mixer.{up,down}_proj` weight/input quantizers | | `attention_qkv_fp8.yaml` | FP8 E4M3 on attention q/k/v bmm and softmax quantizers | -| `gdn_state_fp8_dynamic.yaml` | FP8 E4M3 dynamic fake quantization of the GatedDeltaNet recurrent state (per sequence, head and 64-column tile) at every kernel chunk boundary; needs `fla-core==0.5.1`, Triton, and SM89+ | -| `gdn_w_fp8_dynamic.yaml` | FP8 E4M3 dynamic (per token and head) fake quantization of the WY tensor `w` that multiplies the GatedDeltaNet state; requires identity STE, `fla-core==0.5.1`, and Triton | +| `gdn_state_fp8_dynamic.yaml` | FP8 E4M3 dynamic fake quantization of the GatedDeltaNet recurrent state (per sequence, head and 64-column tile) at serving cache boundaries; requires an explicit serving execution policy | +| `linear_attention_state_int8_block32_dynamic.yaml` | Standard dynamic INT8 with one scale per key row and 32 value channels; the [complete training recipe](../../../general/ptq/linear_attention_state_int8_block32_dynamic.yaml) uses token QDQ and working-state readout | | `indexer_k_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer key cache (`*indexer_k_quantizer`), global scale fixed to 1; Blackwell+ GPUs | | `indexer_q_nvfp4.yaml` | NVFP4 fake quantization of the sparse-attention indexer query (`*indexer_q_quantizer`), global scale fixed to 1; Blackwell+ GPUs | | `moe_dispatch_nvfp4.yaml` | NVFP4 fake quantization of the MoE expert-parallel dispatch (`*dispatch_quantizer`, vLLM); Blackwell+ GPUs | diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index cac4f6d80aa..8618024222f 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -21,6 +21,8 @@ enable: false - quantizer_name: '*gdn_w_quantizer' enable: false + - quantizer_name: '*kda_state_quantizer' + enable: false - quantizer_name: '*block_sparse_moe.gate*' enable: false - quantizer_name: '*linear_attn.conv1d*' diff --git a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml index ed5768ee599..ac36caab0f4 100644 --- a/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml +++ b/modelopt_recipes/configs/ptq/units/gdn_state_fp8_dynamic.yaml @@ -14,12 +14,15 @@ # limitations under the License. # Dynamic per-tile FP8 E4M3 fake quantization of the GatedDeltaNet recurrent state. -# Recompute amax over each [K, 64] tile inside the kernel, independently per sequence and head -# of the [N, H, K, V] state (no calibration). ``axis: [0, 1]`` specifies sequence/head grouping; -# the kernel uses one scale per full-K by 64-value-column tile (two per head for V = 128). -# This is the only configuration the module supports. Applied at every kernel chunk boundary -# (64 tokens) during calibration, QAT and QAD; needs fla-core==0.5.1 and Triton. -# See ``modelopt.torch.quantization.plugins.gated_delta_net``. +# TensorQuantizer computes amax over each [K, BV] tile of the [N, H, K, V] state, +# independently per sequence/head (axis: [0, 1]), without calibration. +# Set BV through linear_attention[].cfg.state_block_v (default: 64). +# This unit enables only the quantizer. Select backend='serving' in the execution +# policy and provide prefill lengths through linear_attention_training_phase. +# With precision='vllm', state QDQ applies at the handoff to a nonempty decode +# suffix and at each suffix token's cache write; fresh prefill has no internal QDQ. +# The serving profile requires public vLLM's GDN kernels. +# See modelopt.torch.quantization.plugins.gdn. # modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig - quantizer_name: '*gdn_state_quantizer' diff --git a/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml b/modelopt_recipes/configs/ptq/units/linear_attention_state_int8_block32_dynamic.yaml similarity index 57% rename from modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml rename to modelopt_recipes/configs/ptq/units/linear_attention_state_int8_block32_dynamic.yaml index 2c73fdd6eaa..00284f84ce3 100644 --- a/modelopt_recipes/configs/ptq/units/gdn_w_fp8_dynamic.yaml +++ b/modelopt_recipes/configs/ptq/units/linear_attention_state_int8_block32_dynamic.yaml @@ -13,16 +13,21 @@ # 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. - +# Dynamic INT8 with one scale per key row and 32 value channels for training and serving. # modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig - - quantizer_name: '*gdn_w_quantizer' + - quantizer_name: '*gdn_state_quantizer' + cfg: + num_bits: 8 + unsigned: false + narrow_range: true + block_sizes: {-1: 32} + type: dynamic + pass_through_bwd: true + - quantizer_name: '*kda_state_quantizer' cfg: - num_bits: e4m3 - axis: [0, 1, 2] + num_bits: 8 + unsigned: false + narrow_range: true + block_sizes: {-1: 32} type: dynamic + pass_through_bwd: true diff --git a/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml new file mode 100644 index 00000000000..f5539e01c05 --- /dev/null +++ b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic.yaml @@ -0,0 +1,39 @@ +# 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. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + base_disable_all: configs/ptq/units/base_disable_all + state_int8: configs/ptq/units/linear_attention_state_int8_block32_dynamic + +metadata: + description: >- + Dynamic INT8 GDN/KDA state QDQ with per-key 32-value scales and working-state readout. + For Kimi-K3/Kimi-Linear, use the explicit kimi_k3 recipe instead. + Uses the installed vLLM GDN or standalone FLA/KDA Triton backend with BF16 operands + and FP32 state storage. + Fresh prefill has no internal state QDQ; handoff and recurrent writes use the configured QDQ. + Supply per-sequence prefix lengths through + linear_attention_training_phase. No calibration is required. +quantize: + algorithm: + quant_cfg: + - $import: base_disable_all + - $import: state_int8 + linear_attention: + - module_name: '*' + cfg: + backend: serving + precision: vllm diff --git a/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic_kimi_k3.yaml b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic_kimi_k3.yaml new file mode 100644 index 00000000000..7801afbee37 --- /dev/null +++ b/modelopt_recipes/general/ptq/linear_attention_state_int8_block32_dynamic_kimi_k3.yaml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + base_disable_all: configs/ptq/units/base_disable_all + state_int8: configs/ptq/units/linear_attention_state_int8_block32_dynamic + +metadata: + description: >- + Dynamic INT8 Kimi-K3/Kimi-Linear KDA state QDQ with per-key 32-value scales and working-state readout. + Requires raw projections and unbounded softplus gates. In vLLM, set both + kda_prefill_backend and kda_decode_backend to triton; mixed batches are not covered. + Fresh prefill has no internal state QDQ; handoff and recurrent writes use the configured QDQ. + Supply per-sequence prefix lengths through + linear_attention_training_phase. No calibration is required. +quantize: + algorithm: + quant_cfg: + - $import: base_disable_all + - $import: state_int8 + linear_attention: + - module_name: '*' + cfg: + backend: serving + precision: vllm_kimi_k3 diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 4c3751fa394..a06c80b75ee 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 32 general/ptq/ recipes (click to expand) +All 34 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -64,9 +64,112 @@ supported combinations. | `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | | `iq2_s` | IQ2_S W2A16 (2.56 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | | `q8_0` | Q8_0 W8A16 (8.5 bpw), eligible linears | none | none (no calibration) | +| `linear_attention_state_int8_block32_dynamic` | GDN/KDA state INT8, 32 value channels per key row; weights unchanged | none | none (dynamic scales; requires vLLM and a prefix/decode phase context) | +| `linear_attention_state_int8_block32_dynamic_kimi_k3` | Kimi-K3/Kimi-Linear KDA state INT8, 32 value channels per key row | none | none (dynamic scales; explicit Triton Kimi profile) |
+The public linear-attention profiles require `vllm==0.30.0` and `fla-core==0.5.1` +(the latter supplies the prefix backward). Install these in the training environment: + +```bash +pip install vllm==0.30.0 fla-core==0.5.1 +``` + +The block32 linear-attention recipe selects `precision="vllm"` +and working-state readout. It applies fake QDQ with floating-point state storage. +For GDN, select `gdn_prefill_backend="triton"` and +`VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE=1` for packed decode. +This profile covers pure, single-token, non-speculative decode. Mixed prefill/decode +batches can select different native kernels, so arbitrary continuous batching is +not numerically qualified by these tests. + +For KDA, `precision="vllm"` selects the standalone FLA/Triton API, without a +model-level serving parity claim; conversion emits a diagnostic for this choice. Use the +`linear_attention_state_int8_block32_dynamic_kimi_k3` recipe for Kimi-K3 and +Kimi-Linear on vLLM 0.30, which share model kernels (`precision="vllm_kimi_k3"`). Set vLLM's `additional_config` to +`{"kda_prefill_backend": "triton", "kda_decode_backend": "triton"}`. The profile +requires raw projections, Q/K normalization, unbounded softplus gates, and the +Kimi-K3 package dependencies (`gate_lower_bound=None` in the serving HF config). +It never falls back to prepared-only arithmetic. +The adapters use the fixed vLLM 0.30.0 value-first cache layout. Native-forward +comparisons and independent recurrence checks validate it in GPU tests. + +The public profiles do not support vLLM 0.20 or arbitrary development builds. +The tested GDN/KDA profiles pass bit-exact native forward comparisons on vLLM 0.30.0. + +Megatron state QAT requires native `_compute_gates` and either +`forward_pre_attn_and_core_attn` or `_forward_compute` for raw gates and packed +padding; Megatron-Core 0.19.2 lacks these hooks and is +rejected before an active QAT forward. Loading checkpoints and running the original +module with state quantizers disabled outside a phase remain supported. + +State QAT training requires explicit per-sequence prefix lengths. Megatron selective +GDN recompute captures each microbatch's phase and packed metadata, allowing backward +after its context exits. Full-layer recompute and CP remain unsupported while state +QAT is active; disabled quantizers outside a phase retain the original module path. +Framework training, evaluation and generation loops must establish this context; +automatic Bridge and evaluation-helper integration is a separate change. Dynamic +state quantizers need no calibration. `examples/megatron_bridge/quantize.py` rejects +linear-attention recipes that require calibration before loading a model. To combine +calibrated projection quantization with state QDQ, provide `mtq.quantize` with a +calibration loop inside `linear_attention_training_phase`; a full-prefill phase +calibrates projections but does not measure decode state-QDQ quality. +To compare state QDQ, convert both models +with the same serving policy, disable the baseline's state quantizers, and use +identical prefix and valid-token lengths: + +```python +for model in (baseline, candidate): + mtq.quantize(model, state_recipe) +mtq.disable_quantizer(baseline, "*state_quantizer") +for model in (baseline, candidate): + with linear_attention_training_phase( + model, prefill_lengths, sequence_lengths=valid_lengths + ): + logits = model(tokens, position_ids, attention_mask) + # Score only valid continuation tokens, identically for both models. +``` + +`sequence_lengths` excludes right padding in dense rows or padded THD segments. +The recurrent suffix ends at that valid length; padding produces zero attention +outputs and does not update the cache. Omit valid lengths only when all storage +tokens are real. For Megatron THD input, omitted lengths are derived from the +unpadded `packed_seq_params.cu_seqlens_q`. If only padded boundaries are available, +pass `sequence_lengths` explicitly, even when every storage token is real. +Full-prefill scoring alone does not exercise decode state QDQ. + +The selected precision profile is saved in ModelOpt/Megatron checkpoints. +HF export currently does not carry this serving contract, and neither checkpoint +format configures or validates a remote vLLM process. Apply the backend and gate +settings above to the actual server; a matching trainer environment variable is no proof. +Server configuration and mixed-batch dispatch enforcement belong to the vLLM plugin. +The tests qualify the selected native kernels, not arbitrary server dispatch. + +Packed boundaries must match the layer's actual (including padding) boundaries. +Validation reads their values rather than trusting tensor versions, so it can synchronize +CUDA metadata. Forward runs one native vLLM prefix and one native decode per +suffix token. Prefix backward reconstructs intermediates and calls upstream FLA +backward; decode uses an explicit recurrence backward. QDQ uses an identity STE. +The `linear_attention/serving` package is a training adapter for these kernels; +it does not integrate with the vLLM engine or scheduler. + +State FP8 block size uses the existing TensorQuantizer settings in `quant_cfg`: + +```python +state_fp8 = { + "num_bits": (4, 3), + "type": "dynamic", + "block_sizes": {-1: 32}, # 16, 32, 64 or 128 value channels per key row + "pass_through_bwd": True, +} +cfg["quant_cfg"].append({"quantizer_name": "*state_quantizer", "cfg": state_fp8}) +``` + +These blocks each have their own dynamic scale. The legacy `axis=(0, 1)` setting +instead shares a scale across a full-key tile; `state_block_v` only controls that +legacy tile's width. NVFP4 state support is a separate extension. + --- ### Model-body schemes diff --git a/noxfile.py b/noxfile.py index afa96010ac4..4caa443ef9c 100644 --- a/noxfile.py +++ b/noxfile.py @@ -151,7 +151,10 @@ def gpu_trtllm(session): # Pin must stay in sync with examples/vllm_serve/Dockerfile. @nox.session(venv_backend="none") def gpu_vllm(session): - session.run("python3", "-m", "pip", "install", "-e", ".[hf,puzzletron,dev-test]") + # Keep vLLM's TileLang/TVM pins; dev-fla pins older versions for other test images. + session.run( + "python3", "-m", "pip", "install", "-e", ".[hf,puzzletron,dev-test]", "fla-core==0.5.1" + ) session.run("python3", "-m", "pytest", "tests/gpu_vllm", *_cov_args()) diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py index 9d46482f328..ab234b14c12 100644 --- a/tests/examples/megatron_bridge/test_quantize_export.py +++ b/tests/examples/megatron_bridge/test_quantize_export.py @@ -17,6 +17,7 @@ from pathlib import Path import pytest +import yaml from _test_utils.examples.run_command import extend_cmd_parts, run_example_command from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches from _test_utils.torch.megatron.modelopt_state import assert_has_modelopt_state @@ -26,6 +27,8 @@ create_tiny_qwen3vl_dir, ) +from modelopt.recipe import load_recipe + # Per-architecture export *mappings* are covered in-process by # tests/gpu_megatron/torch/export/test_unified_export_megatron.py; these cases cover the script # wiring (CLI, recipe, checkpoint hand-off) that only running quantize.py + the exporter exercises. @@ -40,6 +43,29 @@ } +def test_calibrated_state_recipe_rejected_before_model_load(tmp_path): + recipe = { + "metadata": {"recipe_type": "ptq"}, + "quantize": load_recipe( + "general/ptq/linear_attention_state_int8_block32_dynamic" + ).quantize.model_dump(), + } + recipe["quantize"]["algorithm"] = "max" + recipe["quantize"]["quant_cfg"].append( + {"quantizer_name": "*input_quantizer", "cfg": {"num_bits": [4, 3], "axis": None}} + ) + recipe_path = tmp_path / "calibrated_state.yaml" + recipe_path.write_text(yaml.safe_dump(recipe)) + cmd = extend_cmd_parts( + ["torchrun", "--nproc_per_node=1", "quantize.py"], + hf_model_name_or_path=tmp_path / "no_model", + export_megatron_path=tmp_path / "output", + recipe=recipe_path, + ) + with pytest.raises(ValueError, match="calibration loop inside linear_attention_training_phase"): + run_example_command(cmd, example_path="megatron_bridge") + + @pytest.mark.parametrize( ("create_model", "model_kwargs", "recipe", "quantize_parallelism", "export_parallelism"), [ 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 0756f2d0d40..0ecfadcb8d8 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import nullcontext + import pytest import torch from _test_utils.torch.megatron.models import get_mcore_gpt_model @@ -21,21 +23,25 @@ initialize_for_megatron, sharded_state_dict_test_helper, ) +from megatron.core.packed_seq_params import PackedSeqParams +import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase -pytest.importorskip("fla") # Megatron-Core GatedDeltaNet and the state QDQ kernel need fla +pytest.importorskip("fla") # Megatron-Core GatedDeltaNet needs FLA for its baseline kernels. GatedDeltaNet = pytest.importorskip("megatron.core.ssm.gated_delta_net").GatedDeltaNet +vllm = pytest.importorskip("vllm") +if vllm.__version__ != "0.30.0": + pytest.skip("State QAT profiles target vLLM 0.30.0", allow_module_level=True) -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 +HAS_SERVING_HOOKS = hasattr(GatedDeltaNet, "_compute_gates") and any( + hasattr(GatedDeltaNet, name) for name in ("_forward_compute", "forward_pre_attn_and_core_attn") +) def _make_model(tp_size): @@ -56,120 +62,231 @@ def _make_model(tp_size): .cuda() .eval() ) - # Retain enough history for chunk-boundary state rounding to affect the next chunk. + # Retain enough history for recurrent-state rounding to affect later tokens. with torch.no_grad(): for module in model.modules(): if isinstance(module, GatedDeltaNet): module.A_log.fill_(-4) module.dt_bias.zero_() + # Shared hybrid-model KDA settings must not reject a GDN layer. + module.config.kda_safe_gate = True + module.config.kda_lower_bound = -5.0 return model -def _gdn_config(sites): - return { - "quant_cfg": [{"quantizer_name": "*", "enable": False}] - + [ - { - "quantizer_name": f"*gdn_{site}_quantizer", - "cfg": { - "num_bits": (4, 3), - "axis": (0, 1) if site == "state" else (0, 1, 2), - "type": "dynamic", - }, - } - for site in sites - ], - "algorithm": None, - } - - -def _test_gdn_qat_helper(rank, size, checkpoint_path, sites): +def _gdn_config(): + return load_recipe( + "general/ptq/linear_attention_state_int8_block32_dynamic" + ).quantize.model_dump() + + +def _gdn_forward(model): + original_forward = get_forward(model) + + def forward(m, *, decode_aware=True): + with linear_attention_training_phase(m, [64, 64]) if decode_aware else nullcontext(): + return original_forward(m) + + return forward + + +def _test_gdn_qat_helper(rank, size, checkpoint_path): initialize_for_megatron( tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED ) model = _make_model(size) - forward = get_forward(model) - 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) - ] + forward = _gdn_forward(model) with torch.no_grad(): - loss_ref = forward(model) - gdn_ref = outputs.copy() - outputs.clear() + baseline = forward(model, decode_aware=False) + mtq.quantize(model, _gdn_config()) + module = next(m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)) + module.linear_attention_config.backend = "fla" + try: + with pytest.raises(ValueError, match="requires backend='serving'"): + forward(model, decode_aware=False) + finally: + module.linear_attention_config.backend = "serving" + calls = [] + hook = module.gdn_state_quantizer.register_forward_hook(lambda *_: calls.append(True)) + try: + assert torch.isfinite(forward(model)).all() + assert calls + finally: + hook.remove() + mtq.disable_quantizer(model, "*") + with torch.no_grad(): + torch.testing.assert_close(forward(model, decode_aware=False), baseline, rtol=0, atol=0) + mtq.enable_quantizer(model, "*gdn_state_quantizer") + module.linear_attention_config.state_block_v = 32 - mtq.quantize(model, _gdn_config(sites)) - gdn_modules = [m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)] - assert gdn_modules, "no GatedDeltaNet layer was wrapped" - for module in gdn_modules: - assert module.gdn_state_quantizer.is_enabled == ("state" in sites) - assert module.gdn_w_quantizer.is_enabled == ("w" in sites) + restored = _make_model(size) + sharded_state_dict_test_helper(checkpoint_path, model, restored, forward) + restored_gdn = next(m for m in restored.modules() if isinstance(m, _QuantGatedDeltaNet)) + assert restored_gdn.linear_attention_config == module.linear_attention_config + assert restored_gdn.gdn_state_quantizer.is_enabled + assert torch.isfinite(restored_gdn.in_proj.weight.grad).all() + before = restored_gdn.in_proj.weight.detach().clone() + torch.optim.SGD(restored.parameters(), lr=1e-3).step() + assert not torch.equal(restored_gdn.in_proj.weight, before) + _check_packed_forward(restored_gdn) - with torch.no_grad(): - loss_quant = forward(model) - assert torch.isfinite(loss_quant).all() - # BF16 residual additions can hide a small GDN change at the final loss. - assert len(outputs) == len(gdn_ref) == len(gdn_modules) - assert any(not torch.equal(actual, ref) for actual, ref in zip(outputs, gdn_ref)), ( - "enabled QDQ must change the GDN branch output" + for name, value, restore in ( + ("context_parallel_size", 2, True), + ("recompute_granularity", "full", True), + ("recompute_granularity", "full", False), + ): + receiver = _make_model(size) + layer = next(m for m in receiver.modules() if isinstance(m, GatedDeltaNet)) + original = getattr(layer.config, name) + try: + setattr(layer.config, name, value) + receiver.train() + if restore: + mto.restore_from_modelopt_state(receiver, mto.modelopt_state(model)) + else: + mtq.quantize(receiver, _gdn_config()) + assert layer.training + for enabled in (False, True): + if enabled: + layer.gdn_state_quantizer.enable() + else: + layer.gdn_state_quantizer.disable() + layer.modelopt_post_restore() + with ( + pytest.raises(NotImplementedError, match=r"context parallelism|full-layer"), + linear_attention_training_phase(layer, [0]), + ): + layer(torch.zeros(1, 1, 64, device="cuda"), attention_mask=None) + finally: + setattr(layer.config, name, original) + + +def _check_packed_forward(module): + hidden = torch.randn(128, 1, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) + packed = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=torch.tensor([0, 60, 118], device="cuda", dtype=torch.int32), + cu_seqlens_kv=torch.tensor([0, 60, 118], device="cuda", dtype=torch.int32), + cu_seqlens_q_padded=torch.tensor([0, 64, 128], device="cuda", dtype=torch.int32), + cu_seqlens_kv_padded=torch.tensor([0, 64, 128], device="cuda", dtype=torch.int32), + max_seqlen_q=64, + max_seqlen_kv=64, + total_tokens=128, ) - for handle in handles: - handle.remove() + original = getattr(module, "recompute_gdn", None) + was_training = module.training + module.train() + try: + results = [] + cases = [(False, [60, 58], False), (False, None, False)] + if original is not None and hasattr(GatedDeltaNet, "_forward_compute"): + cases.append((True, None, False)) + if hasattr(GatedDeltaNet, "forward_pre_attn_and_core_attn"): + cases.append((False, None, True)) + for recompute, lengths, two_stage in cases: + if original is not None: + module.recompute_gdn = recompute + module.zero_grad(set_to_none=True) + hidden.grad = None + with ( + linear_attention_training_phase( + module, + [56, 54], + sequence_lengths=lengths, + cu_seqlens=packed.cu_seqlens_q_padded, + ), + torch.autocast("cuda", dtype=torch.bfloat16), + ): + if two_stage: + norm = module.forward_pre_attn_and_core_attn( + hidden, attention_mask=None, packed_seq_params=packed + ) + output = module.forward_post_core_attn(norm)[0] + else: + output = module(hidden, attention_mask=None, packed_seq_params=packed)[0] + # Exclude padding from the loss: trained output norms may have a bias. + # Backward runs after the phase exits, as in pipeline schedules. + valid_output = torch.cat((output[:60], output[64:122])) + valid_output.float().square().sum().backward() + assert not hidden.grad[60:64].count_nonzero() + assert not hidden.grad[122:].count_nonzero() + results.append( + (output.detach(), hidden.grad.clone(), module.in_proj.weight.grad.clone()) + ) + for actual in results[1:]: + for value, expected in zip(actual, results[0]): + torch.testing.assert_close(value, expected, rtol=0, atol=0) + finally: + if original is not None: + module.recompute_gdn = original + module.train(was_training) - for site in sites: - mtq.disable_quantizer(model, f"*gdn_{site}_quantizer") - with torch.no_grad(): - torch.testing.assert_close(forward(model), loss_ref, rtol=1e-4, atol=1e-4) - for site in sites: - mtq.enable_quantizer(model, f"*gdn_{site}_quantizer") - restored = _make_model(size) - sharded_state_dict_test_helper(checkpoint_path, model, restored, forward) - restored_gdn = [m for m in restored.modules() if isinstance(m, _QuantGatedDeltaNet)] - assert len(restored_gdn) == len(gdn_modules) - for module in restored_gdn: - assert module.gdn_state_quantizer.is_enabled == ("state" in sites) - assert module.gdn_w_quantizer.is_enabled == ("w" in sites) - assert module.in_proj.weight.grad is not None - assert torch.isfinite(module.in_proj.weight.grad).all() - - # The checkpoint helper already ran backward through the restored model. - optimizer = torch.optim.SGD(restored.parameters(), lr=1e-3) - before = restored_gdn[0].in_proj.weight.detach().clone() - optimizer.step() - assert not torch.equal(restored_gdn[0].in_proj.weight, before) - - -def _compile_gdn_qat_kernels(rank, size, sites): +def _compile_gdn_qat_kernels(rank, size): initialize_for_megatron( tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED ) model = _make_model(size) - forward = get_forward(model) + forward = _gdn_forward(model) with torch.no_grad(): forward(model) - mtq.quantize(model, _gdn_config(sites)) + if not HAS_SERVING_HOOKS: + return # The compatibility test only needs the original Megatron kernel. + cfg = _gdn_config() + cfg["linear_attention"][0]["cfg"]["state_block_v"] = 32 + mtq.quantize(model, cfg) with torch.no_grad(): forward(model) model.train() forward(model).sum().backward() + _check_packed_forward(next(m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet))) torch.cuda.synchronize() -@pytest.fixture +@pytest.fixture(scope="module") 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 + dist_workers_size_1.run(_compile_gdn_qat_kernels) + return dist_workers_size_1 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) + if not HAS_SERVING_HOOKS: + pytest.skip("Serving QAT requires Megatron-Core gate and packed-forward hooks") + compiled_gdn_workers.run(_test_gdn_qat_helper, tmp_path) + + +def _test_missing_serving_hooks(rank, size): + initialize_for_megatron( + tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED + ) + model = _make_model(size) + forward = _gdn_forward(model) + with torch.no_grad(): + baseline = forward(model, decode_aware=False) + mtq.quantize(model, _gdn_config()) + layer = next(m for m in model.modules() if isinstance(m, _QuantGatedDeltaNet)) + with ( + linear_attention_training_phase(layer, [64, 64]), + layer._quantized_linear_attention_kernel(), + pytest.raises(NotImplementedError, match="raw gates captured"), + ): + layer.gated_delta_rule() + for names in (("_compute_gates",), ("_forward_compute", "forward_pre_attn_and_core_attn")): + with pytest.MonkeyPatch.context() as patch: + for cls in GatedDeltaNet.__mro__: + for name in names: + if name in vars(cls): + patch.delattr(cls, name) + restored = _make_model(size) + mto.restore_from_modelopt_state(restored, mto.modelopt_state(model)) + with pytest.raises(NotImplementedError, match="Upgrade Megatron-Core"): + forward(restored) + mtq.disable_quantizer(model, "*") + with torch.no_grad(): + torch.testing.assert_close(forward(model, decode_aware=False), baseline, rtol=0, atol=0) + + +def test_gdn_missing_serving_hooks(compiled_gdn_workers): + compiled_gdn_workers.run(_test_missing_serving_hooks) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py new file mode 100644 index 00000000000..779a7c036f3 --- /dev/null +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_kda.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from contextlib import nullcontext + +import pytest +import torch +from _test_utils.torch.megatron.utils import initialize_for_megatron +from megatron.core import parallel_state +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer import TransformerConfig + +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.linear_attention import linear_attention_training_phase + +KimiDeltaAttention = pytest.importorskip("megatron.core.ssm.gated_delta_net.kda").KimiDeltaAttention +vllm = pytest.importorskip("vllm") +if vllm.__version__ != "0.30.0": + pytest.skip("State QAT profiles target vLLM 0.30.0", allow_module_level=True) + + +def _layer(): + config = TransformerConfig( + num_layers=1, + hidden_size=64, + num_attention_heads=2, + linear_num_key_heads=2, + linear_num_value_heads=2, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_conv_kernel_dim=4, + experimental_attention_variant="kda", + is_hybrid_model=True, + normalization="RMSNorm", + activation_func=torch.nn.functional.silu, + params_dtype=torch.float32, + gradient_accumulation_fusion=False, + recompute_granularity="selective", + recompute_modules=["gdn"], + ) + spec = hybrid_stack_spec.submodules.kda_layer.submodules.self_attention + model = ( + KimiDeltaAttention( + config, + submodules=spec.submodules, + layer_number=1, + pg_collection=ProcessGroupCollection( + tp=parallel_state.get_tensor_model_parallel_group(), + cp=parallel_state.get_context_parallel_group(), + ), + ) + .cuda() + .train() + ) + with torch.no_grad(): + model.A_log.fill_(-4) + model.dt_bias.zero_() + return model + + +def _case(cfg): + initialize_for_megatron(tensor_model_parallel_size=1, pipeline_model_parallel_size=1, seed=73) + model = _layer() + hidden = torch.randn(73, 1, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + def forward(layer, *, decode_aware=True, prefix=31, length=None): + policy = getattr(layer, "linear_attention_config", None) + phase = ( + linear_attention_training_phase( + layer, [prefix], sequence_lengths=None if length is None else [length] + ) + if policy is not None and policy.backend == "serving" and decode_aware + else nullcontext() + ) + with phase, torch.autocast("cuda", dtype=torch.bfloat16): + return layer(hidden, attention_mask=None)[0] + + with torch.no_grad(): + baseline = forward(model) + + def calibrate(m): + with linear_attention_training_phase(m, [73]): + return forward(m, decode_aware=False) + + mtq.quantize(model, {**cfg, "algorithm": "max"}, calibrate) + return model, hidden, forward, baseline + + +def _compile_kda(rank, size, cfg): + model, hidden, forward, _ = _case(cfg) + for prefix in (31, 32): + forward(model, prefix=prefix, length=prefix + 4).float().square().mean().backward() + torch.cuda.synchronize() + + +def _test_kda(rank, size, cfg): + model, hidden, forward, baseline = _case(cfg) + for name, value in (("kda_safe_gate", True), ("kda_lower_bound", -5.0)): + original = getattr(model.config, name) + try: + setattr(model.config, name, value) + with pytest.raises(NotImplementedError, match="unbounded softplus gates"): + mtq.quantize(model, cfg) + finally: + setattr(model.config, name, original) + assert model.kda_state_quantizer.is_enabled + assert model.kda_state_quantizer.num_bits == 8 + assert model.linear_attention_config.precision == "vllm" + assert model.kda_state_quantizer.block_sizes == {-1: 32} + kernel = model.gated_delta_rule + with torch.no_grad(): + assert not torch.equal(forward(model), baseline) + assert model.gated_delta_rule is kernel + + with pytest.raises(ValueError, match="explicit prefill lengths"): + forward(model, decode_aware=False) + # Two in-flight microbatches must keep their own phases after both contexts exit. + expected_grads = [] + model.recompute_gdn = False + for prefix in (31, 32): + model.zero_grad(set_to_none=True) + hidden.grad = None + forward(model, prefix=prefix, length=prefix + 4).float().square().mean().backward() + expected_grads.append((hidden.grad.clone(), model.in_proj.weight.grad.clone())) + model.recompute_gdn = True + pending = [forward(model, prefix=prefix, length=prefix + 4) for prefix in (31, 32)] + for output, (hidden_grad, weight_grad) in zip(pending, expected_grads): + model.zero_grad(set_to_none=True) + hidden.grad = None + output.float().square().mean().backward() + torch.testing.assert_close(hidden.grad, hidden_grad, rtol=0, atol=0) + torch.testing.assert_close(model.in_proj.weight.grad, weight_grad, rtol=0, atol=0) + model.eval() + forward(model).float().sum().backward() # No checkpoint hook may reject eval gradients. + model.train() + + mtq.disable_quantizer(model, "*") + with torch.no_grad(): + torch.testing.assert_close(forward(model, decode_aware=False), baseline, rtol=0, atol=0) + mtq.enable_quantizer(model, "*kda_state_quantizer") + before = model.in_proj.weight.detach().clone() + torch.optim.SGD(model.parameters(), lr=0.1).step() + assert not torch.equal(before, model.in_proj.weight) + + +@pytest.fixture(scope="module") +def compiled_kda_workers(dist_workers_size_1): + """Warm one KDA shape outside the functional test timer.""" + cfg = load_recipe( + "general/ptq/linear_attention_state_int8_block32_dynamic" + ).quantize.model_dump() + dist_workers_size_1.run(_compile_kda, cfg) + return dist_workers_size_1, cfg + + +def test_kda_qat_recompute(compiled_kda_workers): + workers, cfg = compiled_kda_workers + workers.run(_test_kda, cfg) diff --git a/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py new file mode 100644 index 00000000000..5054ffddb7e --- /dev/null +++ b/tests/gpu_vllm/torch/quantization/test_linear_attention_training.py @@ -0,0 +1,392 @@ +# 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.quantization.linear_attention_reference import recurrent_delta_rule_reference + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + gdn_state_qat, + kda_state_qat, +) +from modelopt.torch.quantization.nn import TensorQuantizer + +vllm = pytest.importorskip("vllm") + +if vllm.__version__ != "0.30.0": + pytest.skip("State QAT profiles target vLLM 0.30.0", allow_module_level=True) + +from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator +from vllm.third_party.flash_linear_attention import ops as gdn_ops +from vllm.third_party.flash_linear_attention.ops import kda as kda_ops + + +@torch.no_grad() +def _native_forward(args, quantizer, kda, kimi_k3): + """Run the vLLM 0.30 Triton APIs with a persistent quantized cache.""" + q, k, v, raw_g, raw_beta, rate, bias = [x.detach().clone() for x in args] + _, shape = MambaStateShapeCalculator.gated_delta_net_state_shape( + 1, 1, 1, q.shape[-1], v.shape[-1], 4 + ) + assert shape == (1, v.shape[-1], q.shape[-1]) + + def layout(state): + return state.transpose(-1, -2).contiguous() + + cache = torch.zeros(2, *shape, device=q.device) + prefix_args = [x[:, :65].contiguous() for x in (q, k, v)] + kwargs = { + "initial_state": cache[1:], + "output_final_state": True, + "use_qk_l2norm_in_kernel": True, + "cu_seqlens": torch.tensor([0, 65], device=q.device, dtype=torch.int32), + } + if kda: + native = kda_ops + if kimi_k3: + from vllm.models.kimi_k3.nvidia.ops.third_party import kda as native + g = kda_ops.fused_kda_gate(raw_g.flatten(-2), rate, q.shape[-1], g_bias=bias) + beta = raw_beta.float().sigmoid() + # vLLM 0.30's two explicit KDA profiles have different beta contracts. + beta_arg = ( + {"raw_beta": raw_beta[:, :65].contiguous()} + if kimi_k3 + else {"beta": beta[:, :65].contiguous()} + ) + prefix, state = native.chunk_kda_with_fused_gate( + *prefix_args, + raw_g=raw_g[:, :65].contiguous(), + A_log=rate, + g_bias=bias, + **beta_arg, + **kwargs, + ) + else: + prepared = gdn_ops.fused_post_conv_prep( + torch.cat([x[0].flatten(1) for x in prefix_args], dim=-1), + raw_g[0, :65], + raw_beta[0, :65], + rate, + bias, + 1, + q.shape[-1], + v.shape[-1], + ) + kwargs["use_qk_l2norm_in_kernel"] = False + prefix, state = gdn_ops.chunk_gated_delta_rule(*[x[None] for x in prepared], **kwargs) + cache[1] = layout(quantizer(layout(state)[0])) + outputs = [prefix] + indices = torch.ones(1, device=q.device, dtype=torch.int32) + cu = torch.tensor([0, 1], device=q.device, dtype=torch.int32) + for token in range(65, 67): + operands = [x[:, token : token + 1].contiguous() for x in (q, k, v)] + if kimi_k3: + output, _ = native.fused_recurrent_kda_packed_decode( + torch.cat([x.flatten() for x in operands])[None], + raw_g[:, token : token + 1].contiguous(), + raw_beta[:, token : token + 1].contiguous(), + rate, + bias, + None, + cache, + indices, + ) + state = cache[1:] + elif kda: + output, state = native.fused_recurrent_kda( + *operands, + g[:, token : token + 1], + beta[:, token : token + 1], + initial_state=cache[1:], + inplace_final_state=False, + cu_seqlens=cu, + use_qk_l2norm_in_kernel=True, + ) + else: + output, _ = gdn_ops.fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=torch.cat([x.flatten() for x in operands])[None], + a=raw_g[0, token : token + 1], + b=raw_beta[0, token : token + 1], + A_log=rate, + dt_bias=bias, + scale=q.shape[-1] ** -0.5, + initial_state=cache, + out=torch.empty_like(operands[-1]), + ssm_state_indices=indices, + use_qk_l2norm_in_kernel=True, + ) + state = cache[1:] + outputs.append(output) + cache[1] = layout(quantizer(layout(state)[0])) + return torch.cat(outputs, dim=1), layout(cache[1:]) + + +@pytest.fixture(scope="module", params=["gdn", "kda", "kimi_k3"]) +def compiled_serving_case(request): + """Compile one shared BF16 shape per model, outside the test-call timer.""" + kda = request.param != "gdn" + kimi_k3 = request.param == "kimi_k3" + torch.manual_seed(73) + keys = 32 # Share one head shape across GDN/KDA to reuse native and FLA compilation. + args = [torch.randn(1, 67, 1, keys, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + args += [ + torch.randn((1, 67, 1, keys) if kda else (1, 67, 1), device="cuda", dtype=torch.bfloat16), + torch.randn(1, 67, 1, device="cuda", dtype=torch.bfloat16), + torch.full((1,), -3.0, device="cuda"), + torch.zeros(keys if kda else 1, device="cuda"), + ] + args = [x.requires_grad_() for x in args] + policy = LinearAttentionConfig( + backend="serving", precision="vllm_kimi_k3" if kimi_k3 else "vllm" + ) + quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=8, + type="dynamic", + block_sizes={-1: 32}, + narrow_range=True, + pass_through_bwd=True, + ) + ).cuda() + + def forward(inputs=args, **overrides): + q, k, v, raw_g, raw_beta, rate, bias = inputs + kwargs = { + "policy": policy, + "state_quantizer": quantizer, + "prefill_lengths": [65], + "use_qk_l2norm_in_kernel": True, + "output_final_state": True, + "gate_inputs": (raw_g, raw_beta, rate, bias), + } + if kda: + kwargs.update( + A_log=rate, + dt_bias=bias, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + ) + kwargs.update(overrides) + if kda: + return kda_state_qat(q, k, v, raw_g, raw_beta, **kwargs) + gate = -rate.exp() * torch.nn.functional.softplus(raw_g.float() + bias) + return gdn_state_qat(q, k, v, gate, raw_beta.float().sigmoid(), **kwargs) + + output, state = forward() + torch.autograd.grad(output.float().sum() + state.sum(), args) + expected = _native_forward(args, quantizer, kda, kimi_k3) + torch.cuda.synchronize() + return args, quantizer, forward, expected, kimi_k3 + + +def test_serving_state_qdq_and_handoff_gradient(compiled_serving_case): + args, quantizer, forward, (expected, expected_state), _ = compiled_serving_case + calls = [] + handle = quantizer.register_forward_hook(lambda *_: calls.append(True)) + try: + output, state = forward() + finally: + handle.remove() + assert len(calls) == 3 # Handoff plus two recurrent writes. + torch.testing.assert_close(output, expected, rtol=0, atol=0) + torch.testing.assert_close(state, expected_state, rtol=0, atol=0) + gradients = torch.autograd.grad( + output[:, 65:].float().square().sum() + state.square().sum(), args + ) + assert all(torch.isfinite(x).all() for x in gradients) + assert gradients[1][:, :65].abs().sum() > 0 + with pytest.raises(ValueError, match="explicit prefill lengths"): + forward(prefill_lengths=None) + + with torch.no_grad(): + quantizer.disable() + try: + plain, _ = forward() + assert not torch.equal(output[:, 65:], plain[:, 65:]) + # A nonzero nonsymmetric state checks layout against an independent oracle. + initial = torch.randn(1, 1, args[0].shape[-1], args[2].shape[-1], device="cuda") + # Reuse the compiled 65-token prefix for full-prefill and suffix cases. + for tokens in (65, 67): + inputs = [x[:, :tokens] if i < 5 else x for i, x in enumerate(args)] + actual, final = forward(inputs=inputs, initial_state=initial) + q, k, v, raw_g, raw_beta, rate, bias = [x.float() for x in inputs] + q, k = [x / (x.square().sum(-1, keepdim=True) + 1e-6).sqrt() for x in (q, k)] + q[:, :65] = q[:, :65].bfloat16().float() + k[:, :65] = k[:, :65].bfloat16().float() + gate = -rate.exp().reshape( + (1, 1, -1, 1) if raw_g.ndim == 4 else (1, 1, -1) + ) * torch.nn.functional.softplus(raw_g + bias.reshape(raw_g.shape[2:])) + expected, expected_state = recurrent_delta_rule_reference( + q, k, v, gate, raw_beta.sigmoid(), initial_state=initial + ) + torch.testing.assert_close(actual.float(), expected, rtol=0.02, atol=0.01) + torch.testing.assert_close(final, expected_state, rtol=0.02, atol=0.01) + finally: + quantizer.enable() + + +@pytest.fixture(scope="module") +@torch.no_grad() +def compiled_gdn_layout_case(): + """Compile two-head decode so strided head/value axes are observable.""" + torch.manual_seed(732) + q, k, v = [torch.randn(1, 2, 2, 32, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + raw_g, raw_beta = [torch.randn(1, 2, 2, device="cuda", dtype=torch.bfloat16) for _ in range(2)] + gates = (raw_g, raw_beta, torch.full((2,), -3.0, device="cuda"), torch.randn(2, device="cuda")) + initial = torch.randn(1, 2, 32, 32, device="cuda") + + def forward(values, gate_inputs): + raw_g, raw_beta, rate, bias = gate_inputs + return gdn_state_qat( + q, + k, + values, + -rate.exp() * torch.nn.functional.softplus(raw_g.float() + bias), + raw_beta.float().sigmoid(), + policy=LinearAttentionConfig(backend="serving"), + initial_state=initial, + prefill_lengths=[0], + output_final_state=True, + use_qk_l2norm_in_kernel=True, + gate_inputs=gate_inputs, + ) + + expected = forward(v, gates) + torch.cuda.synchronize() + return v, gates, initial, forward, expected + + +@pytest.mark.parametrize("strided", ["values", "gates"]) +@torch.no_grad() +def test_gdn_strided_inputs_match_contiguous(compiled_gdn_layout_case, strided): + v, gates, initial, forward, expected = compiled_gdn_layout_case + original_state = initial.clone() + if strided == "values": + v = v.transpose(-1, -2).contiguous().transpose(-1, -2) + else: + gates = tuple(torch.stack((x, torch.zeros_like(x)), dim=-1)[..., 0] for x in gates) + actual = forward(v, gates) + for value, reference in zip(actual, expected): + torch.testing.assert_close(value, reference, rtol=0, atol=0) + torch.testing.assert_close(initial, original_state, rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("compiled_serving_case", "overrides", "message"), + [ + ("gdn", {"gate_inputs": None}, "raw gate_inputs"), + ( + "kimi_k3", + {"gate_inputs": None, "use_gate_in_kernel": False, "use_beta_sigmoid_in_kernel": False}, + "raw gate_inputs", + ), + ("kimi_k3", {"use_qk_l2norm_in_kernel": False}, "normalization"), + ], + indirect=["compiled_serving_case"], + scope="module", +) +def test_rejects_silent_serving_mismatches(compiled_serving_case, overrides, message): + _, _, forward, _, _ = compiled_serving_case + with pytest.raises(ValueError, match=message): + forward(**overrides) + + +@pytest.mark.parametrize("compiled_serving_case", ["gdn", "kimi_k3"], indirect=True, scope="module") +@torch.no_grad() +def test_raw_projections_match_bf16_serving(compiled_serving_case): + args, quantizer, forward, _, kimi_k3 = compiled_serving_case + inputs = [x.detach() for x in args] + inputs[3:5] = [x.float() + 0.0007 for x in inputs[3:5]] + # Rates and biases keep their FP32 precision, independent of projection rounding. + inputs[5:] = [x + 0.0123 for x in inputs[5:]] + rounded = [*inputs[:3], *[x.bfloat16() for x in inputs[3:5]], *inputs[5:]] + expected = _native_forward(rounded, quantizer, kimi_k3, kimi_k3) + actual = forward(inputs=inputs) + for value, reference in zip(actual, expected): + torch.testing.assert_close(value, reference, rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("compiled_serving_case", "packed"), + [("gdn", False), ("kda", True)], + indirect=["compiled_serving_case"], + scope="module", + ids=["gdn-dense", "kda-packed"], +) +def test_padding_does_not_extend_decode(compiled_serving_case, packed): + args, quantizer, forward, _, _ = compiled_serving_case + valid = 66 # One recurrent write and one padding token after the 65-token prefix. + trimmed = [ + (x[:, :valid] if i < 5 else x).detach().clone().requires_grad_() for i, x in enumerate(args) + ] + initial = args[0].new_zeros( + 2, args[2].shape[2], args[0].shape[-1], args[2].shape[-1], dtype=torch.float32 + ) + initial[1].fill_(1.0) # Empty rows preserve an incoming nonzero state without QDQ. + expected, expected_state = forward(inputs=trimmed, initial_state=initial[:1]) + expected_grads = torch.autograd.grad( + expected.square().sum() + expected_state.square().sum(), trimmed + ) + padded = [] + for i, x in enumerate(args): + if i < 5: + x = torch.cat((x, torch.zeros_like(x)), dim=0) + if packed: + x = x.flatten(0, 1).unsqueeze(0) + padded.append(x.detach().clone().requires_grad_()) + calls = [] + hook = quantizer.register_forward_hook(lambda *_: calls.append(True)) + try: + output, state = forward( + inputs=padded, + initial_state=initial, + prefill_lengths=[65, 0], + sequence_lengths=[valid, 0], + cu_seqlens=torch.tensor([0, 67, 134], device="cuda", dtype=torch.int32) + if packed + else None, + ) + finally: + hook.remove() + assert len(calls) == 3 # Incoming prefix state, handoff, one write; none for the empty row. + output = output.reshape(2, 67, *output.shape[-2:]) + torch.testing.assert_close(output[:1, :valid], expected, rtol=0, atol=0) + torch.testing.assert_close(state[:1], expected_state, rtol=0, atol=0) + assert not output[0, valid:].count_nonzero() and not output[1].count_nonzero() + torch.testing.assert_close(state[1], initial[1], rtol=0, atol=0) + grads = torch.autograd.grad(output.square().sum() + state.square().sum(), padded) + for i, (grad, expected_grad) in enumerate(zip(grads, expected_grads)): + if i < 5: + grad = grad.reshape(2, 67, *grad.shape[2:]) + assert not grad[0, valid:].count_nonzero() and not grad[1].count_nonzero() + grad = grad[:1, :valid] + torch.testing.assert_close(grad, expected_grad) + + if not packed: + with torch.no_grad(): + prefix_inputs = [x[:, :65] if i < 5 else x for i, x in enumerate(trimmed)] + prefill, prefill_state = forward(inputs=prefix_inputs, initial_state=initial[:1]) + padded_prefill, padded_state = forward( + inputs=padded, + prefill_lengths=[65, 0], + sequence_lengths=[65, 0], + initial_state=initial, + ) + torch.testing.assert_close(padded_prefill[:1, :65], prefill, rtol=0, atol=0) + torch.testing.assert_close(padded_state[:1], prefill_state, rtol=0, atol=0) + torch.testing.assert_close(padded_state[1], initial[1], rtol=0, atol=0) + assert not padded_prefill[0, 65:].count_nonzero() and not padded_prefill[1].count_nonzero() diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_autograd.py b/tests/gpu_vllm/torch/quantization/test_vllm_autograd.py new file mode 100644 index 00000000000..9474f5d49b9 --- /dev/null +++ b/tests/gpu_vllm/torch/quantization/test_vllm_autograd.py @@ -0,0 +1,144 @@ +# 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.quantization.linear_attention_reference import recurrent_delta_rule_reference + +from modelopt.torch.quantization.linear_attention._vllm_autograd import prefix, step +from modelopt.torch.quantization.linear_attention.utils import forward_value + +vllm = pytest.importorskip("vllm") +if vllm.__version__ != "0.30.0": + pytest.skip("State QAT profiles target vLLM 0.30.0", allow_module_level=True) + + +@pytest.fixture(scope="module", params=[(False, True), (True, True), (False, False)]) +def compiled_decode(request): + """Compile the existing one-head decode shape outside the test timer.""" + channel, normalize = request.param + torch.manual_seed(97) + q, k, v = [torch.randn(1, 32, device="cuda").bfloat16().float() for _ in range(3)] + raw_gate = torch.randn((1, 32) if channel else (1,), device="cuda").bfloat16().float() + raw_beta = torch.randn(1, device="cuda").bfloat16().float() + rate, bias = torch.full((1,), -3.0, device="cuda"), torch.zeros_like(raw_gate) + initial = torch.randn(1, 32, 32, device="cuda") + inputs = [x.requires_grad_() for x in (q, k, v, raw_gate, raw_beta, rate, bias, initial)] + gate = -rate.exp() * torch.nn.functional.softplus(raw_gate + bias) + operands = q, k, v, gate, raw_beta.sigmoid(), initial + kwargs = {"scale": 0.3, "normalize": normalize, "gate_inputs": inputs[3:7]} + with torch.no_grad(): + native = step(*operands, **kwargs) + torch.cuda.synchronize() + return inputs, operands, kwargs, native + + +@pytest.mark.parametrize("loss_on", ["both", "output", "state"]) +def test_decode_backward_matches_recurrence(compiled_decode, loss_on): + inputs, operands, kwargs, native = compiled_decode + output, final = step(*operands, **kwargs) + for actual, expected in zip((output, final), native): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + q, k, v, gate, beta, initial = operands + if kwargs["normalize"]: + q, k = [x / (x.square().sum(-1, keepdim=True) + 1e-6).sqrt() for x in (q, k)] + _, reference_state = recurrent_delta_rule_reference( + *[x[None, None] for x in (q, k, v, gate, beta)], + initial_state=initial[None], + scale=kwargs["scale"], + ) + # The former Torch adjoint evaluated the readout at the native working state. + working = forward_value(reference_state[0], native[1]) + reference_output = torch.einsum("hk,hkv->hv", q * kwargs["scale"], working) + probes = [torch.randn_like(x) for x in (output, final)] + active = (loss_on != "state", loss_on != "output") + losses = [ + sum((value * probe).sum() for value, probe, use in zip(values, probes, active) if use) + for values in ((output, final), (reference_output, working)) + ] + gradients = [ + torch.autograd.grad(loss, inputs, allow_unused=True, retain_graph=True) for loss in losses + ] + for actual, expected in zip(*gradients): + if actual is None or expected is None: + assert actual is expected + else: + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=2e-6) + + +@pytest.fixture(scope="module", params=[False, True]) +def compiled_prefix(request): + """Reuse the integration tests' 65-token prefix shape for both recurrence types.""" + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + from fla.ops.kda import chunk_kda + + channel = request.param + torch.manual_seed(98) + inputs = [torch.randn(65, 1, 32, device="cuda").bfloat16().float() for _ in range(3)] + inputs += [ + torch.randn((65, 1, 32) if channel else (65, 1), device="cuda").bfloat16().float(), + torch.randn(65, 1, device="cuda").bfloat16().float(), + torch.full((1,), -3.0, device="cuda"), + torch.zeros(32 if channel else 1, device="cuda"), + torch.randn(1, 32, 32, device="cuda"), + ] + inputs = [x.requires_grad_() for x in inputs] + q, k, v, raw_gate, raw_beta, rate, bias, state = inputs + gate = -rate.exp() * torch.nn.functional.softplus(raw_gate + bias) + beta = raw_beta.sigmoid() + kwargs = { + "scale": 32**-0.5, + "beta_dtype": torch.float32, + "normalize": True, + "gate_inputs": (raw_gate, raw_beta, rate, bias), + } + operands = q, k, v, gate, beta, state + actual = prefix(*operands, **kwargs) + torch.autograd.grad(sum(x.sum() for x in actual), inputs, retain_graph=True) + chunk = chunk_kda if channel else chunk_gated_delta_rule + expected = chunk( + *[x[None].bfloat16() for x in (q, k, v)], + gate[None], + beta[None], + initial_state=state[None], + output_final_state=True, + scale=kwargs["scale"], + use_qk_l2norm_in_kernel=True, + ) + torch.cuda.synchronize() + return inputs, operands, kwargs, tuple(x[0].float() for x in expected) + + +@pytest.mark.parametrize("loss_on", ["both", "output", "state"]) +def test_prefix_backward_matches_fla(compiled_prefix, loss_on): + inputs, operands, kwargs, expected = compiled_prefix + actual = prefix(*operands, **kwargs) + with torch.no_grad(): + native = prefix(*operands, **kwargs) + for value, reference in zip(actual, native): + torch.testing.assert_close(value, reference, rtol=0, atol=0) + probes = [torch.randn_like(x) for x in actual] + active = (loss_on != "state", loss_on != "output") + gradients = [ + torch.autograd.grad( + sum((x * p).sum() for x, p, use in zip(values, probes, active) if use), + inputs, + retain_graph=True, + ) + for values in (actual, expected) + ] + for value, reference in zip(*gradients): + torch.testing.assert_close(value, reference, rtol=1e-5, atol=2e-6) diff --git a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py b/tests/unit/torch/quantization/plugins/test_gated_delta_net.py deleted file mode 100644 index 3793edcb189..00000000000 --- a/tests/unit/torch/quantization/plugins/test_gated_delta_net.py +++ /dev/null @@ -1,276 +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. - -from copy import deepcopy -from functools import partial - -import pytest -import torch -import torch.nn as nn - -import modelopt.torch.opt as mto -import modelopt.torch.quantization as mtq -from modelopt.torch.quantization.nn import QuantModuleRegistry -from modelopt.torch.quantization.plugins import gated_delta_net -from modelopt.torch.quantization.plugins.gated_delta_net import GatedDeltaNetStateQuantMixin - -GDN_STATE_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"} - - -def chunk_gated_delta_rule(q, k, v, g, beta, **kwargs): - """CPU stand-in for the optional FLA kernel.""" - return q + k + v, None - - -@pytest.fixture(autouse=True) -def mock_fla_kernel(monkeypatch): - monkeypatch.setattr( - gated_delta_net, "_fla_chunk_gated_delta_rule", lambda: chunk_gated_delta_rule - ) - - -class TinyGatedDeltaNet(nn.Module): - """A module that, like Megatron-Core's GatedDeltaNet, calls ``self.gated_delta_rule``.""" - - def __init__(self): - super().__init__() - self.proj = nn.Linear(4, 4) - self.gated_delta_rule = chunk_gated_delta_rule - - def forward(self, x): - out, _ = self.gated_delta_rule(x, x, x, x[..., 0], x[..., 0]) - return self.proj(out) - - -@QuantModuleRegistry.register({TinyGatedDeltaNet: "TinyGatedDeltaNet"}) -class _QuantTinyGatedDeltaNet(GatedDeltaNetStateQuantMixin): - def forward(self, x): - gated_delta_rule = self.gated_delta_rule - self.gated_delta_rule = lambda *a, **kw: self._state_quantized_chunk_gated_delta_rule( - gated_delta_rule, *a, **kw - ) - try: - return super().forward(x) - finally: - self.gated_delta_rule = gated_delta_rule - - -GDN_W_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1, 2), "type": "dynamic"} - - -def quant_cfg(state=True, w=False): - entries = [{"quantizer_name": "*", "enable": False}] - if state: - entries.append({"quantizer_name": "*gdn_state_quantizer", "cfg": GDN_STATE_FP8_DYNAMIC}) - if w: - entries.append({"quantizer_name": "*gdn_w_quantizer", "cfg": GDN_W_FP8_DYNAMIC}) - return {"quant_cfg": entries, "algorithm": "max"} - - -@pytest.mark.parametrize( - "attributes", - [ - {"num_bits": (4, 3), "axis": (0, 1)}, # static - {"num_bits": (4, 3), "type": "dynamic"}, # per tensor - {"num_bits": 8, "axis": (0, 1), "type": "dynamic"}, # int8 - {"num_bits": (4, 3), "type": "dynamic", "block_sizes": {-1: 16}}, # blockwise - ], -) -def test_validate_state_quantizer_rejects_unsupported(attributes): - with pytest.raises(ValueError, match="supports only"): - mtq.quantize( - TinyGatedDeltaNet(), - { - "quant_cfg": [ - {"quantizer_name": "*", "enable": False}, - {"quantizer_name": "*gdn_state_quantizer", "cfg": attributes}, - ], - "algorithm": None, - }, - ) - - -def test_dynamic_export_removes_linear_attention_attributes(): - model = _QuantTinyGatedDeltaNet.convert(TinyGatedDeltaNet()) - model.export() - assert type(model) is TinyGatedDeltaNet - for name in ("gdn_state_quantizer", "gdn_w_quantizer"): - assert not hasattr(model, name) - - -def test_disabled_state_quantizer_calls_original_kernel(): - model = TinyGatedDeltaNet() - x = torch.randn(2, 8, 3, 4) - expected = model(x) - - disable_all = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": "max"} - mtq.quantize(model, disable_all, lambda m: m(x)) - - assert isinstance(model, _QuantTinyGatedDeltaNet) - assert not model.gdn_state_quantizer.is_enabled and not model.gdn_w_quantizer.is_enabled - assert torch.equal(model(x), expected) - assert model.gated_delta_rule is chunk_gated_delta_rule, "the kernel swap must be undone" - - -@pytest.mark.parametrize("partial_kernel", [False, True]) -def test_enabled_state_quantizer_uses_state_qdq_kernel(monkeypatch, partial_kernel): - calls = [] - - def fake_state_qdq_kernel(*args, **kwargs): - calls.append(kwargs) - return chunk_gated_delta_rule(*args) - - monkeypatch.setattr( - gated_delta_net, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel - ) - model = TinyGatedDeltaNet() - if partial_kernel: - model.gated_delta_rule = partial(chunk_gated_delta_rule, output_final_state=True) - x = torch.randn(2, 8, 3, 4) - mtq.quantize(model, quant_cfg(), lambda m: m(x)) - - model(x) - assert calls and calls[-1] == { - "chunk_size": 64, - "state_qdq": 1, - "state_qdq_block_v": 64, - "w_quantizer": None, - **({"output_final_state": True} if partial_kernel else {}), - } - - # An unrelated callable must not pass validation merely by copying the FLA name. - model.gated_delta_rule = lambda *a, **kw: chunk_gated_delta_rule(*a, **kw) - model.gated_delta_rule.__name__ = "chunk_gated_delta_rule" - with pytest.raises(NotImplementedError, match="supports only FLA"): - model(x) - - -@pytest.mark.parametrize("state", [False, True]) -def test_w_quantizer_is_passed_to_the_kernel(monkeypatch, state): - """``*gdn_w_quantizer`` in the config hands the module's TensorQuantizer to the kernel, with - or without the state quantizer.""" - calls = [] - - def fake_state_qdq_kernel(*args, **kwargs): - calls.append(kwargs) - return chunk_gated_delta_rule(*args) - - monkeypatch.setattr( - gated_delta_net, "_state_qdq_chunk_gated_delta_rule", lambda: fake_state_qdq_kernel - ) - model = TinyGatedDeltaNet() - x = torch.randn(2, 8, 3, 4) - mtq.quantize(model, quant_cfg(state=state, w=True), lambda m: m(x)) - assert model.gdn_w_quantizer.is_enabled and model.gdn_state_quantizer.is_enabled == state - - model(x) - assert calls[-1]["state_qdq"] == int(state) - assert calls[-1]["w_quantizer"] is model.gdn_w_quantizer - - # The w quantizer really quantizes: 256 random values per token collapse onto the E4M3 grid, - # which has at most 127 distinct magnitudes per (row-specific) scale. - w = torch.randn(1, 1, 1, 256) - quantized = model.gdn_w_quantizer(w) - assert not torch.equal(quantized, w) - assert torch.unique(quantized.abs()).numel() <= 127 < torch.unique(w.abs()).numel() - - -@pytest.mark.parametrize("site", ["state", "w"]) -@pytest.mark.parametrize( - "overrides", - [{"pass_through_bwd": False}, {"type": "static"}, {"fake_quant": False}, {"rotate": True}], -) -def test_unsupported_quantizer_fails_during_conversion(site, overrides): - cfg = quant_cfg(state=site == "state", w=site == "w") - cfg = deepcopy(cfg) - cfg["quant_cfg"][-1]["cfg"].update(overrides) - with pytest.raises(ValueError, match="supports only"): - mtq.quantize(TinyGatedDeltaNet(), cfg) - - -@pytest.mark.parametrize("axis", [None, (0, 1), (0, 1, 2)]) -def test_w_grouping_is_preserved_during_conversion(axis): - cfg = deepcopy(quant_cfg(state=False, w=True)) - cfg["algorithm"] = None - cfg["quant_cfg"][-1]["cfg"]["axis"] = axis - model = mtq.quantize(TinyGatedDeltaNet(), cfg) - assert model.gdn_w_quantizer.axis == axis - - -@pytest.mark.parametrize(("state", "w"), [(True, False), (False, True), (True, True)]) -def test_quantizer_roundtrip_and_hybrid_selection(tmp_path, state, w): - model = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) - cfg = quant_cfg(state=state, w=w) - cfg["algorithm"] = None - mtq.quantize(model, cfg) - if w: - # Save the actual quantizer settings, including edits after conversion. - model[0].gdn_w_quantizer.axis = None - path = tmp_path / "gdn.pth" - mto.save(model, path) - restored = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) - mto.restore(restored, path) - for name, enabled in (("gdn_state_quantizer", state), ("gdn_w_quantizer", w)): - original = getattr(model[0], name) - quantizer = getattr(restored[0], name) - assert quantizer.is_enabled == enabled - assert quantizer.axis == original.axis - assert quantizer.num_bits == original.num_bits - assert quantizer._dynamic == original._dynamic - assert not hasattr(restored[1], name) - - -def test_quant_cfg_refinement_updates_and_validates_existing_quantized_module(): - cfg = quant_cfg() - cfg["algorithm"] = None - model = mtq.quantize(TinyGatedDeltaNet(), cfg) - assert model.gdn_state_quantizer.is_enabled - assert not model.gdn_w_quantizer.is_enabled - - cfg = quant_cfg(state=False, w=True) - cfg["algorithm"] = None - mtq.quantize(model, cfg) - assert not model.gdn_state_quantizer.is_enabled - assert model.gdn_w_quantizer.is_enabled - - cfg = deepcopy(cfg) - cfg["quant_cfg"][-1]["cfg"]["pass_through_bwd"] = False - with pytest.raises(ValueError, match="supports only"): - mtq.quantize(model, cfg) - - -def test_restore_legacy_gdn_without_new_quantizer_handles(): - config = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None} - model = mtq.quantize(TinyGatedDeltaNet(), config) - state = mto.modelopt_state(model) - for _, mode_state in state["modelopt_state_dict"]: - metadata = mode_state["metadata"] - for name in ("gdn_state_quantizer", "gdn_w_quantizer"): - metadata["quantizer_state"].pop(name, None) - restored = TinyGatedDeltaNet() - mto.restore_from_modelopt_state(restored, state) - restored.load_state_dict(model.state_dict()) - x = torch.randn(2, 8, 3, 4) - torch.testing.assert_close(restored(x), model(x)) - assert not restored.gdn_state_quantizer.is_enabled - assert not restored.gdn_w_quantizer.is_enabled - - -def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): - model = mtq.quantize( - TinyGatedDeltaNet(), mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(2, 8, 3, 4)) - ) - assert not model.gdn_state_quantizer.is_enabled - assert not model.gdn_w_quantizer.is_enabled diff --git a/tests/unit/torch/quantization/plugins/test_gdn.py b/tests/unit/torch/quantization/plugins/test_gdn.py new file mode 100644 index 00000000000..b7af95761c4 --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_gdn.py @@ -0,0 +1,317 @@ +# 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 numpy as np +import pytest +import torch +import torch.nn as nn + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.linear_attention import ( + LinearAttentionConfig, + get_linear_attention_layers, + linear_attention_training_phase, +) +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.plugins import gdn +from modelopt.torch.quantization.plugins.gdn import GatedDeltaNetStateQuantMixin +from modelopt.torch.quantization.plugins.kda import KimiDeltaAttentionStateQuantMixin +from modelopt.torch.quantization.plugins.linear_attention import _validate_linear_attention + +GDN_STATE_FP8_DYNAMIC = {"num_bits": (4, 3), "axis": (0, 1), "type": "dynamic"} + + +def chunk_gated_delta_rule(q, k, v, g, beta, **kwargs): + """CPU stand-in for the optional FLA kernel.""" + return q + k + v, None + + +@pytest.fixture(autouse=True) +def mock_fla_kernel(monkeypatch): + monkeypatch.setattr(gdn, "_fla_chunk_gated_delta_rule", lambda: chunk_gated_delta_rule) + + +class TinyGatedDeltaNet(nn.Module): + """A module that, like Megatron-Core's GatedDeltaNet, calls ``self.gated_delta_rule``.""" + + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4) + self.gated_delta_rule = chunk_gated_delta_rule + + def forward(self, x): + out, _ = self.gated_delta_rule(x, x, x, x[..., 0], x[..., 0]) + return self.proj(out) + + +@QuantModuleRegistry.register({TinyGatedDeltaNet: "TinyGatedDeltaNet"}) +class _QuantTinyGatedDeltaNet(GatedDeltaNetStateQuantMixin): + def forward(self, x): + gated_delta_rule = self.gated_delta_rule + self.gated_delta_rule = lambda *a, **kw: self._state_quantized_chunk_gated_delta_rule( + gated_delta_rule, *a, **kw + ) + try: + return super().forward(x) + finally: + self.gated_delta_rule = gated_delta_rule + + +def quant_cfg(): + return { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*gdn_state_quantizer", "cfg": GDN_STATE_FP8_DYNAMIC}, + ], + "linear_attention": [{"module_name": "*", "cfg": {"backend": "serving"}}], + "algorithm": "max", + } + + +def test_disabled_quantization_preserves_baseline(): + model = TinyGatedDeltaNet() + x = torch.randn(2, 8, 3, 4) + expected = model(x) + cfg = quant_cfg() + cfg["quant_cfg"] = [{"quantizer_name": "*", "enable": False}] + mtq.quantize(model, cfg, lambda m: m(x)) + torch.testing.assert_close(model(x), expected, rtol=0, atol=0) + assert model.gated_delta_rule is chunk_gated_delta_rule + assert not model.linear_attention_is_enabled + + +def test_phase_routes_lengths_and_preserves_outer_context(monkeypatch): + monkeypatch.setattr( + _QuantTinyGatedDeltaNet, + "validate_linear_attention_execution", + lambda _: pytest.fail("Conversion must not validate runtime execution settings"), + ) + cfg = {**quant_cfg(), "algorithm": None} + with pytest.raises(ValueError, match="requires backend='serving'"): + mtq.quantize(TinyGatedDeltaNet(), {k: v for k, v in cfg.items() if k != "linear_attention"}) + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + calls = [] + + def forward(*args, **kwargs): + calls.append(kwargs) + return chunk_gated_delta_rule(*args) + + monkeypatch.setattr(gdn, "gdn_state_qat", forward) + with ( + linear_attention_training_phase( + model, np.array([4, 4]), sequence_lengths=[7, 6], cu_seqlens=[0, 8, 16] + ), + linear_attention_training_phase(model, [1, 1], preserve_existing=True), + ): + model(torch.randn(2, 8, 3, 4)) + assert calls[-1]["prefill_lengths"] == (4, 4) + assert calls[-1]["sequence_lengths"] == (7, 6) + boundaries = calls[-1]["cu_seqlens_cpu"] + actual = torch.tensor([0, 8, 16]) + assert boundaries.resolve(actual) == (0, 8, 16) + actual.data[1] = 7 # Alias writes do not bump the tensor version. + with pytest.raises(ValueError, match="does not match"): + boundaries.resolve(actual) + assert model._linear_attention_prefill_lengths is None + assert model._linear_attention_sequence_lengths is None + mtq.disable_quantizer(model, "*") + assert get_linear_attention_layers(model) == (model,) + assert not get_linear_attention_layers(model, enabled_only=True) + with linear_attention_training_phase(model, [4, 4]): + assert not get_linear_attention_layers(model, unphased_only=True) + assert model.linear_attention_is_enabled + assert not model.linear_attention_is_enabled + + +def test_restore_uses_saved_quantizer_and_policy(tmp_path, monkeypatch): + model = nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)) + mtq.quantize(model, {**quant_cfg(), "algorithm": None}) + # Change both after conversion: restore must not validate one against stale recipe defaults. + model[0].gdn_state_quantizer.set_from_attribute_config( + {"num_bits": 8, "axis": (0, 1), "type": "dynamic", "narrow_range": True} + ) + model[0].linear_attention_config = LinearAttentionConfig(backend="serving", state_block_v=32) + sample = torch.randn(2, 4, 19) + expected = model[0].gdn_state_quantizer(sample) + path = tmp_path / "gdn.pth" + mto.save(model, path) + + def no_rematching(*args): + pytest.fail("Restore must use saved resolved policies") + + monkeypatch.setattr( + "modelopt.torch.quantization.plugins.linear_attention._apply_linear_attention_policy", + no_rematching, + ) + restored = mto.restore(nn.Sequential(TinyGatedDeltaNet(), nn.Linear(4, 4)), path) + assert restored[0].linear_attention_config == model[0].linear_attention_config + assert restored[0].gdn_state_quantizer.is_enabled + assert not restored[0].gdn_w_quantizer.is_enabled + assert not hasattr(restored[1], "gdn_state_quantizer") + torch.testing.assert_close(restored[0].gdn_state_quantizer(sample), expected) + + +@pytest.mark.parametrize("mixin", [GatedDeltaNetStateQuantMixin, KimiDeltaAttentionStateQuantMixin]) +def test_state_qat_requires_serving_policy(mixin): + model = mixin.convert(TinyGatedDeltaNet()) + model._linear_attn_state.set_from_attribute_config(GDN_STATE_FP8_DYNAMIC) + model._linear_attn_state.enable() + with pytest.raises(ValueError, match="requires backend='serving'"): + model.validate_linear_attention() + model.linear_attention_config = LinearAttentionConfig(backend="serving") + if mixin is KimiDeltaAttentionStateQuantMixin: + with pytest.warns(UserWarning, match="standalone FLA/Triton"): + _validate_linear_attention(model) + with pytest.warns(UserWarning, match="standalone FLA/Triton"): + model.modelopt_post_restore() + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"pass_through_bwd": False}, "supports only"), + ({"type": "static"}, "supports only"), + ( + {"num_bits": 8, "axis": None, "block_sizes": {-1: 32, "type": "dynamic"}}, + "use top-level type='dynamic'", + ), + ], +) +def test_unsupported_quantizer_fails_at_conversion(overrides, message): + cfg = quant_cfg() + cfg["algorithm"] = None + cfg["quant_cfg"][-1]["cfg"] = {**GDN_STATE_FP8_DYNAMIC, **overrides} + with pytest.raises(ValueError, match=message): + mtq.quantize(TinyGatedDeltaNet(), cfg) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"prefill_lengths": [-1]}, "Prefill lengths"), + ({"prefill_lengths": [0], "sequence_lengths": [True]}, "Sequence lengths"), + ({"prefill_lengths": [2], "sequence_lengths": [1]}, "prefill length <= sequence length"), + ({"prefill_lengths": [0, 0], "sequence_lengths": [1]}, "per sequence"), + ], +) +def test_invalid_phase_lengths(kwargs, message): + with ( + pytest.raises(ValueError, match=message), + linear_attention_training_phase(nn.Identity(), **kwargs), + ): + pass + + +@pytest.mark.parametrize("has_policy", [False, True]) +def test_restore_legacy_gdn_without_new_quantizer_handles(has_policy): + 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"] + if not has_policy: + metadata.pop("linear_attention", None) + else: + assert "linear_attention" in metadata + for name in ("gdn_state_quantizer", "gdn_w_quantizer"): + metadata["quantizer_state"].pop(name, None) + restored = TinyGatedDeltaNet() + mto.restore_from_modelopt_state(restored, state) + restored.load_state_dict(model.state_dict()) + x = torch.randn(2, 8, 3, 4) + torch.testing.assert_close(restored(x), model(x)) + assert not restored.gdn_state_quantizer.is_enabled + assert not restored.gdn_w_quantizer.is_enabled + + +def test_standard_projection_recipe_leaves_gdn_emulation_disabled(): + model = mtq.quantize( + TinyGatedDeltaNet(), mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(2, 8, 3, 4)) + ) + assert not model.gdn_state_quantizer.is_enabled + assert not model.gdn_w_quantizer.is_enabled + + +def test_empty_stage_conversion_and_offline_restore(monkeypatch): + cfg = { + "quant_cfg": [{"quantizer_name": "*", "enable": False}], + "linear_attention": [{"module_name": "layer", "cfg": {"backend": "serving"}}], + "algorithm": None, + } + model = mtq.quantize(nn.ModuleDict({"layer": nn.Linear(4, 4)}), cfg) + with linear_attention_training_phase(model, [3]): + pass + state = mto.modelopt_state(model) + + # Saved resolved policies must restore without evaluating the original global recipe. + def no_rematching(*args): + pytest.fail("Restore must use saved resolved policies") + + monkeypatch.setattr( + "modelopt.torch.quantization.plugins.linear_attention._apply_linear_attention_policy", + no_rematching, + ) + restored = nn.ModuleDict({"layer": 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["layer"](x), model["layer"](x)) + + +def test_enabled_legacy_w_quantizer_is_rejected_on_restore(): + cfg = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None} + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + model.gdn_w_quantizer.enable() + with pytest.raises(ValueError, match="GDN W quantization is no longer supported"): + mto.restore_from_modelopt_state(TinyGatedDeltaNet(), mto.modelopt_state(model)) + + +def test_quant_cfg_refinement_and_policy_precedence(): + cfg = {**quant_cfg(), "algorithm": None} + model = mtq.quantize(TinyGatedDeltaNet(), cfg) + cfg["linear_attention"] = [ + {"module_name": "*", "cfg": {"backend": "serving", "state_block_v": 128}}, + {"module_name": "", "cfg": {"backend": "serving", "state_block_v": 32}}, + ] + assert mtq.quantize(model, cfg) is model + assert model.linear_attention_config.state_block_v == 32 + cfg["linear_attention"].reverse() + mtq.quantize(model, cfg) + assert model.linear_attention_config.state_block_v == 128 + cfg["quant_cfg"][-1]["cfg"] = {**GDN_STATE_FP8_DYNAMIC, "pass_through_bwd": False} + with pytest.raises(ValueError, match="supports only"): + mtq.quantize(model, cfg) + + +@pytest.mark.parametrize( + ("class_name", "quantizer_name"), + [ + ("GatedDeltaNet", "gdn_state_quantizer"), + ("Qwen3NextGatedDeltaNetAttention", "gdn_state_quantizer"), + ("KimiDeltaAttention", "kda_state_quantizer"), + ("KimiK3DeltaAttention", "kda_state_quantizer"), + ("KimiLinearDeltaAttention", "kda_state_quantizer"), + ], +) +def test_state_recipe_rejects_unconverted_layers(class_name, quantizer_name): + # An unknown framework class has no registered state-quantization adapter. + model = type(class_name, (nn.Module,), {})() + model.proj = nn.Linear(4, 4) + config = load_recipe("general/ptq/linear_attention_state_int8_block32_dynamic").quantize + with pytest.raises(RuntimeError, match=quantizer_name): + mtq.quantize(model, config) diff --git a/tests/unit/torch/quantization/test_linear_attention_decode.py b/tests/unit/torch/quantization/test_linear_attention_decode.py new file mode 100644 index 00000000000..f43ad30f2a4 --- /dev/null +++ b/tests/unit/torch/quantization/test_linear_attention_decode.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from modelopt.torch.quantization.config import QuantizeConfig, QuantizerAttributeConfig +from modelopt.torch.quantization.linear_attention import LinearAttentionConfig, kda_state_qat +from modelopt.torch.quantization.linear_attention import training as state_training +from modelopt.torch.quantization.linear_attention.utils import _state_qdq, state_quantizer_config +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.mark.parametrize("state_format", ["fp8_e4m3", "int8"]) +def test_state_qdq_matches_tensor_quantizer(state_format): + torch.manual_seed(762) + # A partial value tile checks grouping and padding against TensorQuantizer. + value = torch.randn(2, 3, 5, 19, requires_grad=True) + cfg = {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1)} + if state_format == "int8": + cfg.update(num_bits=8, unsigned=False, narrow_range=True) + quantizer = TensorQuantizer(QuantizerAttributeConfig(**cfg)) + expected = torch.cat( + [quantizer(tile.flatten(-2)).reshape_as(tile) for tile in value.split(16, -1)], -1 + ) + actual = _state_qdq(value, quantizer, 16) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + probe = torch.randn_like(value) + (gradient,) = torch.autograd.grad((actual * probe).sum(), value) + torch.testing.assert_close(gradient, probe, rtol=0, atol=0) + + +@pytest.mark.parametrize("block_size", [16, 32, 64, 128]) +def test_fp8_state_blocks_match_tensor_quantizer(block_size): + value = torch.randn(1, 2, 3, 133, requires_grad=True) + quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=(4, 3), type="dynamic", block_sizes={-1: block_size}, pass_through_bwd=True + ) + ) + assert state_quantizer_config(quantizer) == ("fp8_e4m3", block_size) + actual = _state_qdq(value, quantizer) + torch.testing.assert_close(actual, quantizer(value), rtol=0, atol=0) + probe = torch.randn_like(value) + (gradient,) = torch.autograd.grad((actual * probe).sum(), value) + torch.testing.assert_close(gradient, probe, rtol=0, atol=0) + + +@pytest.mark.parametrize( + "settings", + [ + {"chunk_size": 32}, + {"solve": {"method": "neumann"}}, + {"state": {"mode": "token"}}, + ], +) +def test_linear_attention_policy_rejects_invalid_settings(settings): + with pytest.raises(ValueError): + QuantizeConfig( + linear_attention=[{"module_name": "*", "cfg": {"backend": "serving", **settings}}] + ) + + +@pytest.mark.parametrize("raw_gates", [False, True]) +def test_kda_rejects_incompatible_negative_eigenvalue_gates(raw_gates): + q = torch.zeros(1, 1, 1, 32) + beta = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="allow_neg_eigval"): + kda_state_qat( + q, + q, + q, + q, + beta, + policy=LinearAttentionConfig(backend="serving"), + allow_neg_eigval=True, + use_beta_sigmoid_in_kernel=raw_gates, + gate_inputs=(q, beta, torch.zeros(1), torch.zeros(32)) if raw_gates else None, + ) + + +@pytest.mark.parametrize(("major", "triton_version"), [(12, "3.6.0"), (12, "3.7.0"), (8, "3.6.0")]) +def test_state_training_device_guard(monkeypatch, major, triton_version): + # CI cannot safely run the faulting kernels; simulate only the absent hardware/version. + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (major, 0)) + monkeypatch.setattr(state_training, "version", lambda name: triton_version) + state_training._validate_training_device.cache_clear() + try: + if major >= 12 and triton_version == "3.6.0": + with pytest.raises(NotImplementedError, match=r"sm_120.*Triton < 3\.7"): + state_training._validate_training_device(torch.device("cuda:0")) + else: + state_training._validate_training_device(torch.device("cuda:0")) + finally: + state_training._validate_training_device.cache_clear()