From 21f697fa7d049092a1818612165e5c6ca7f2618f Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Thu, 2 Jul 2026 22:39:38 -0700 Subject: [PATCH 01/28] Add compact vLLM attention quant carrier Signed-off-by: Kai Xu --- modelopt/torch/quantization/plugins/vllm.py | 52 +++++-- .../quantization/test_vllm_dynamic_modules.py | 138 ++++++++++++++++++ 2 files changed, 174 insertions(+), 16 deletions(-) diff --git a/modelopt/torch/quantization/plugins/vllm.py b/modelopt/torch/quantization/plugins/vllm.py index 95ca3240b73..a81b71b57b3 100644 --- a/modelopt/torch/quantization/plugins/vllm.py +++ b/modelopt/torch/quantization/plugins/vllm.py @@ -23,13 +23,6 @@ from itertools import chain import torch - -# Try multiple import paths for vLLM compatibility across versions -if importlib.util.find_spec("vllm.attention"): - import vllm.attention as vllm_attention # vllm < 0.16.0 -else: - import vllm.model_executor.layers.attention as vllm_attention # vllm >= 0.16.0 - import vllm.model_executor.layers.fused_moe.layer as vllm_fused_moe_layer import vllm.model_executor.layers.linear as vllm_linear from vllm.distributed.parallel_state import get_dp_group, get_ep_group, get_tp_group @@ -38,6 +31,25 @@ from ..nn import QuantLinearConvBase, QuantModule, QuantModuleRegistry, TensorQuantizer from .custom import CUSTOM_MODEL_PLUGINS + +def _import_attention_module(): + """Import a vLLM module that exports the concrete ``Attention`` class.""" + for module_name in ( + "vllm.attention.layer", + "vllm.model_executor.layers.attention", + "vllm.attention", + ): + try: + module = importlib.import_module(module_name) + except ImportError: + continue + if hasattr(module, "Attention"): + return module + raise ImportError("No supported vLLM Attention module was found") + + +vllm_attention = _import_attention_module() + # Try multiple import paths for vLLM compatibility across versions vllm_shared_fused_moe_layer = None for module_path in [ @@ -68,14 +80,6 @@ except ImportError: EncoderOnlyAttention = None -try: - _has_attention_layer = importlib.util.find_spec("vllm.attention.layer") is not None -except (ModuleNotFoundError, ValueError): - _has_attention_layer = False - -if _has_attention_layer: - import vllm.attention.layer as vllm_attention - try: VllmMLAAttention = vllm_attention.MLAAttention except (AttributeError, ImportError): @@ -188,6 +192,21 @@ def vllm_replace_quant_module_hook(model: torch.nn.Module) -> None: CUSTOM_MODEL_PLUGINS.add(vllm_replace_quant_module_hook) +def _set_vllm_attention_kv_default_amax(module, device: torch.device) -> None: + """Set a global-scale-one amax on uncalibrated block-16 NVFP4 K/V quantizers.""" + for name in ("k_bmm_quantizer", "v_bmm_quantizer"): + quantizer = getattr(module, name, None) + if ( + not isinstance(quantizer, TensorQuantizer) + or not quantizer.is_enabled + or not quantizer.is_nvfp4_dynamic + or (quantizer.block_sizes or {}).get(-1) != 16 + or hasattr(quantizer, "_amax") + ): + continue + quantizer.amax = torch.tensor(6.0 * 448.0, device=device, dtype=torch.float32) + + def _vllm_attention_modelopt_post_restore(self) -> None: """Move Attention module to its correct device after ModelOpt state restore.""" device, dtype = _get_device_dtype(self) @@ -506,7 +525,8 @@ def _setup(self): def forward(self, query, key, value, *args, **kwargs): query = self.q_bmm_quantizer(query) key = self.k_bmm_quantizer(key) - value = self.v_bmm_quantizer(value) + if not getattr(self, "_value_quant_in_kernel", False): + value = self.v_bmm_quantizer(value) return super().forward(query, key, value, *args, **kwargs) def modelopt_post_restore(self, prefix: str = "") -> None: diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py index 2136f959754..6576c700faf 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -29,8 +29,11 @@ from __future__ import annotations import gc +from types import SimpleNamespace +from unittest.mock import Mock import pytest +import torch from _test_utils.torch.transformers_models import ( create_tiny_deepseek_v3_dir, create_tiny_llama_dir, @@ -40,16 +43,151 @@ from vllm.distributed import cleanup_dist_env_and_memory import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.plugins import vllm as vllm_plugin from modelopt.torch.quantization.plugins.vllm import ( _ATTENTION_TYPES, VllmMLAAttention, _QuantFusedMoEBase, + _QuantVLLMAttention, _VLLMParallelLinear, disable_compilation, ) +class _NativeAttention(torch.nn.Module): + def forward(self, query, key, value, *args, **kwargs): + return query, key, value + + +class _TestQuantVLLMAttention(_QuantVLLMAttention, _NativeAttention): + pass + + +def _new_attention(cls): + attention = object.__new__(cls) + torch.nn.Module.__init__(attention) + return attention + + +def _nvfp4_quantizer(*, block_size=16, enabled=True): + quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: block_size, "type": "dynamic", "scale_bits": (4, 3)}, + enable=enabled, + ) + ) + return quantizer + + +def test_import_attention_module_resolves_installed_concrete_attention(): + attention_module = vllm_plugin._import_attention_module() + assert isinstance(attention_module.Attention, type) + assert vllm_plugin.vllm_attention.Attention is attention_module.Attention + + +@pytest.mark.parametrize( + "wrapper_name", + ["_QuantVLLMAttention", "_QuantVLLMCrossAttention", "_QuantVLLMEncoderOnlyAttention"], +) +def test_attention_setup_keeps_qkv_only_checkpoint_surface(monkeypatch, wrapper_name): + monkeypatch.setattr( + vllm_plugin, + "create_parallel_state", + lambda: vllm_plugin.ParallelState(data_parallel_group=None), + ) + wrapper = getattr(vllm_plugin, wrapper_name) + test_wrapper = type(f"Test{wrapper.__name__}", (wrapper, _NativeAttention), {}) + attention = _new_attention(test_wrapper) + + attention._setup() + + quantizer_names = ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer") + assert set(dict(attention.named_children())) == set(quantizer_names) + for name in quantizer_names: + getattr(attention, name).amax = torch.tensor(1.0) + assert set(attention.state_dict()) == {f"{name}._amax" for name in quantizer_names} + assert not hasattr(attention, "_value_quant_in_kernel") + + attention.k_bmm_quantizer = _nvfp4_quantizer() + attention.v_bmm_quantizer = _nvfp4_quantizer() + attention.device, attention.dtype = torch.device("cpu"), torch.float32 + attention.modelopt_post_restore() + assert not hasattr(attention.k_bmm_quantizer, "_amax") + assert not hasattr(attention.v_bmm_quantizer, "_amax") + + +def test_quant_vllm_attention_forward_skips_only_in_kernel_value_quantization(): + attention = _new_attention(_TestQuantVLLMAttention) + attention.q_bmm_quantizer = Mock(side_effect=lambda inputs: inputs + 1) + attention.k_bmm_quantizer = Mock(side_effect=lambda inputs: inputs + 2) + attention.v_bmm_quantizer = Mock(side_effect=lambda inputs: inputs + 3) + query = torch.tensor(10) + key = torch.tensor(20) + value = torch.tensor(30) + + assert not hasattr(attention, "_value_quant_in_kernel") + quantized = attention(query, key, value) + attention._value_quant_in_kernel = True + in_kernel = attention(query, key, value) + + assert quantized[:3] == (torch.tensor(11), torch.tensor(22), torch.tensor(33)) + assert in_kernel[:3] == (torch.tensor(11), torch.tensor(22), value) + assert tuple(q.call_count for q in (attention.q_bmm_quantizer, attention.k_bmm_quantizer)) == ( + 2, + 2, + ) + assert attention.v_bmm_quantizer.call_count == 1 + + +@pytest.mark.parametrize( + "wrapper_name", ["_QuantVLLMCrossAttention", "_QuantVLLMEncoderOnlyAttention"] +) +def test_cross_encoder_attention_wrappers_keep_qkv_only_behavior(wrapper_name): + wrapper = getattr(vllm_plugin, wrapper_name) + test_wrapper = type(f"Test{wrapper.__name__}", (wrapper, _NativeAttention), {}) + attention = _new_attention(test_wrapper) + attention.q_bmm_quantizer = Mock(side_effect=lambda inputs: inputs + 1) + attention.k_bmm_quantizer = Mock(side_effect=lambda inputs: inputs + 2) + attention.v_bmm_quantizer = Mock(side_effect=lambda inputs: inputs + 3) + + output = attention(torch.tensor(10), torch.tensor(20), torch.tensor(30)) + + assert output == (torch.tensor(11), torch.tensor(22), torch.tensor(33)) + assert attention.v_bmm_quantizer.call_count == 1 + + +def test_attention_kv_defaults_set_only_uncalibrated_dynamic_block16_quantizers(): + calibrated_amax = 7.25 + layer = SimpleNamespace( + q_bmm_quantizer=_nvfp4_quantizer(), + k_bmm_quantizer=_nvfp4_quantizer(), + v_bmm_quantizer=_nvfp4_quantizer(), + p_bmm_quantizer=_nvfp4_quantizer(), + ) + layer.v_bmm_quantizer.amax = calibrated_amax + + vllm_plugin._set_vllm_attention_kv_default_amax(layer, torch.device("cpu")) + + assert layer.k_bmm_quantizer._amax.item() == 6.0 * 448.0 + assert layer.v_bmm_quantizer._amax.item() == calibrated_amax + assert not hasattr(layer.q_bmm_quantizer, "_amax") + assert not hasattr(layer.p_bmm_quantizer, "_amax") + + +def test_attention_kv_defaults_ignore_unsupported_quantizers(): + for quantizer in ( + TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3))), + _nvfp4_quantizer(block_size=32), + _nvfp4_quantizer(enabled=False), + ): + layer = SimpleNamespace(k_bmm_quantizer=quantizer, v_bmm_quantizer=quantizer) + vllm_plugin._set_vllm_attention_kv_default_amax(layer, torch.device("cpu")) + assert not hasattr(quantizer, "_amax") + + def _quantize_and_summarize(self): """Run on the worker via ``LLM.collective_rpc``. From 7436ba64821acbddfbdf267efbbc98b2c8898cf6 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Thu, 2 Jul 2026 23:08:48 -0700 Subject: [PATCH 02/28] Add compact NVFP4 V attention path Signed-off-by: Kai Xu --- .../kernels/common/attention/triton_fa.py | 85 ++++++++-- .../kernels/quantization/attention/v_qdq.py | 129 ++++++++++++++++ .../quantization/common/nvfp4_quant.py | 3 +- .../quantization/gemm/fp4_kernel_hopper.py | 1 - .../common/attention/test_triton_fa_p_qdq.py | 30 +++- .../common/attention/test_triton_fa_paged.py | 146 ++++++++++++++++++ 6 files changed, 376 insertions(+), 18 deletions(-) create mode 100644 modelopt/torch/kernels/quantization/attention/v_qdq.py diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 751ce052a5c..7e51ad074aa 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -41,8 +41,9 @@ _apply_sparse_nm_to_qk_tile: Any = None _is_dense_region: Any = None _skip_softmax_decision: Any = None -_p_qdq_fp8: Any = None +_qdq_fp8: Any = None _p_qdq_nvfp4: Any = None +_v_qdq_nvfp4: Any = None def _load_sparsity_helpers() -> None: @@ -63,18 +64,20 @@ def _load_sparsity_helpers() -> None: _skip_softmax_decision = _skip -def _load_p_qdq_helpers() -> None: - global _p_qdq_fp8, _p_qdq_nvfp4 - if _p_qdq_fp8 is None: - from modelopt.torch.kernels.quantization.attention.p_qdq import _p_qdq_nvfp4 as _nvfp4 +def _load_qdq_helpers() -> None: + global _qdq_fp8, _p_qdq_nvfp4, _v_qdq_nvfp4 + if _qdq_fp8 is None: + from modelopt.torch.kernels.quantization.attention.p_qdq import _p_qdq_nvfp4 as _p_nvfp4 + from modelopt.torch.kernels.quantization.attention.v_qdq import _v_qdq_nvfp4 as _v_nvfp4 from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq as _fp8 - _p_qdq_fp8 = _fp8 - _p_qdq_nvfp4 = _nvfp4 + _qdq_fp8 = _fp8 + _p_qdq_nvfp4 = _p_nvfp4 + _v_qdq_nvfp4 = _v_nvfp4 -# Maps the public p_qdq option to the kernel's P_QDQ constexpr. -_P_QDQ_MODES = {None: 0, "fp8": 1, "nvfp4": 2} +# Maps public P/V QDQ options to kernel constexpr values. +_QDQ_MODES = {None: 0, "fp8": 1, "nvfp4": 2} LOG2E: float = 1.44269504088896 @@ -263,6 +266,9 @@ def _attn_fwd( SKIP_THRESHOLD_LOG2: tl.constexpr = 0.0, # log2(lambda) in the kernel's scaled log2 score space P_QDQ: tl.constexpr = 0, # Fake quant-dequant of softmax P: 0=off, 1=FP8 E4M3, 2=NVFP4 p_qdq_scale=1.0, # Per-tensor scale for softmax qdq (runtime scalar; amax/448 or amax/(6*448)) + V_QDQ: tl.constexpr = 0, # Fake quant-dequant of V: 0=off, 1=FP8 E4M3, 2=NVFP4 + v_qdq_scale=1.0, + V_CACHE_QUANTIZED: tl.constexpr = False, # complete block-16 groups are already QDQ Sparsity_total=None, # Optional int64 scalar for counting total tiles (atomic) Sparsity_skipped=None, # Optional int64 scalar for counting skipped tiles (atomic) MEASURE_SPARSITY: tl.constexpr = False, # When True, count total/skipped tiles via atomic adds @@ -323,6 +329,7 @@ def _attn_fwd( if not IS_CAUSAL else tl.minimum(causal_offset + (tile_q + 1) * BLOCK_M, seq_len_kv) ) + v_quantized_boundary = (seq_len_kv // 16) * 16 # --- Main loop: iterate over KV tiles --- for kv_start in range(0, kv_bound, BLOCK_N): @@ -404,7 +411,7 @@ def _attn_fwd( # row_sum keeps the unquantized p: the softmax denominator stays in # fp32 and only the quantized P is fed to BMM2. if P_QDQ == 1: - p = _p_qdq_fp8(p, p_qdq_scale) + p = _qdq_fp8(p, p_qdq_scale) elif P_QDQ == 2: p = _p_qdq_nvfp4(p, p_qdq_scale, BLOCK_M, BLOCK_N) @@ -435,6 +442,19 @@ def _attn_fwd( mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :], other=0.0, ) + if V_QDQ != 0 and ( + (not V_CACHE_QUANTIZED) or (kv_start + BLOCK_N > v_quantized_boundary) + ): + if V_QDQ == 1: + v_qdq = _qdq_fp8(v.to(tl.float32), v_qdq_scale) + else: + v_qdq = _v_qdq_nvfp4(v.to(tl.float32), v_qdq_scale, BLOCK_N, BLOCK_D) + v_qdq = v_qdq.to(v.dtype) + if V_CACHE_QUANTIZED: + use_qdq = (kv_start + kv_pos) >= v_quantized_boundary + v = tl.where(use_qdq[:, None], v_qdq, v) + else: + v = v_qdq acc = tl.dot(p.to(v.dtype), v, acc) row_max = m_new # else: tile skipped — no softmax, no V load, no BMM2 for this tile @@ -833,6 +853,9 @@ def forward( measure_sparsity, p_qdq_mode, p_qdq_scale, + v_qdq_mode, + v_qdq_scale, + v_cache_quantized, k_cache, v_cache, block_table, @@ -932,6 +955,9 @@ def forward( "SKIP_THRESHOLD_LOG2": skip_threshold_log2, "P_QDQ": p_qdq_mode, "p_qdq_scale": p_qdq_scale, + "V_QDQ": v_qdq_mode, + "v_qdq_scale": v_qdq_scale, + "V_CACHE_QUANTIZED": v_cache_quantized, "Sparsity_total": sparsity_total, "Sparsity_skipped": sparsity_skipped, "MEASURE_SPARSITY": do_measure, @@ -1137,6 +1163,9 @@ def backward(ctx, grad_output): None, # measure_sparsity None, # p_qdq_mode None, # p_qdq_scale + None, # v_qdq_mode + None, # v_qdq_scale + None, # v_cache_quantized None, # k_cache None, # v_cache None, # block_table @@ -1165,6 +1194,9 @@ def attention( measure_sparsity: bool = False, p_qdq: str | None = None, p_qdq_amax: float = 1.0, + v_qdq: str | None = None, + v_qdq_amax: float | None = None, + v_cache_quantized: bool = False, k_cache: torch.Tensor | None = None, v_cache: torch.Tensor | None = None, block_table: torch.Tensor | None = None, @@ -1221,6 +1253,14 @@ def attention( and the global scale ``amax / (6 * 448)`` for NVFP4. A runtime scalar — user-set or calibrated values do not recompile the kernel. Values above amax saturate. + v_qdq: Fake quant-dequant of V before ``P @ V``. ``"fp8"`` uses a + per-tensor E4M3 scale; ``"nvfp4"`` uses signed E2M1 values with + one E4M3 scale per 16 keys. ``None`` disables V QDQ. + v_qdq_amax: Optional per-tensor V amax. ``None`` uses global scale 1; + otherwise converts to ``amax / 448`` (FP8) or ``amax / (6 * 448)`` + (NVFP4). + v_cache_quantized: Complete block-16 groups in the paged V cache are + already QDQ; only the pristine partial group is QDQ on read. k_cache: Paged K cache [num_blocks, page_size, num_kv_heads, head_dim]. When provided, K/V are read from paged cache via block_table instead of from contiguous k/v tensors. @@ -1244,12 +1284,12 @@ def attention( # permanently excluded from the cache key and later edits to them would # silently reuse stale compiled kernels from the on-disk cache. _load_sparsity_helpers() - _load_p_qdq_helpers() - if p_qdq not in _P_QDQ_MODES: + _load_qdq_helpers() + if p_qdq not in _QDQ_MODES: raise ValueError( - f"p_qdq must be one of {sorted(k for k in _P_QDQ_MODES if k)} or None, got {p_qdq!r}" + f"p_qdq must be one of {sorted(k for k in _QDQ_MODES if k)} or None, got {p_qdq!r}" ) - p_qdq_mode = _P_QDQ_MODES[p_qdq] + p_qdq_mode = _QDQ_MODES[p_qdq] # Convert the per-tensor amax to the kernel's scale convention # (``q = cast(p / scale) * scale``): FP8 uses ``amax / 448``; NVFP4 uses the # global scale ``amax / (6 * 448)``. amax=1 (the default, the theoretical @@ -1259,6 +1299,20 @@ def attention( if not (math.isfinite(p_qdq_amax) and p_qdq_amax > 0): raise ValueError(f"p_qdq_amax must be a finite positive value, got {p_qdq_amax}") p_qdq_scale = p_qdq_amax / 448.0 if p_qdq == "fp8" else p_qdq_amax / (6.0 * 448.0) + if v_qdq not in _QDQ_MODES: + raise ValueError( + f"v_qdq must be one of {sorted(k for k in _QDQ_MODES if k)} or None, got {v_qdq!r}" + ) + v_qdq_mode = _QDQ_MODES[v_qdq] + v_qdq_scale = 1.0 + if v_qdq_mode and v_qdq_amax is not None: + if not (math.isfinite(v_qdq_amax) and v_qdq_amax > 0): + raise ValueError(f"v_qdq_amax must be a finite positive value, got {v_qdq_amax}") + v_qdq_scale = v_qdq_amax / 448.0 if v_qdq == "fp8" else v_qdq_amax / (6.0 * 448.0) + if v_cache_quantized and v_qdq != "nvfp4": + raise ValueError("v_cache_quantized requires v_qdq='nvfp4'") + if v_cache_quantized and any(x is None for x in (k_cache, v_cache, block_table)): + raise ValueError("v_cache_quantized requires a paged KV cache") sm_scale = 1.0 / (q.shape[2] ** 0.5) if softmax_scale is None else softmax_scale return _Attention.apply( q, @@ -1280,6 +1334,9 @@ def attention( measure_sparsity, p_qdq_mode, p_qdq_scale, + v_qdq_mode, + v_qdq_scale, + v_cache_quantized, k_cache, v_cache, block_table, diff --git a/modelopt/torch/kernels/quantization/attention/v_qdq.py b/modelopt/torch/kernels/quantization/attention/v_qdq.py new file mode 100644 index 00000000000..876a189787e --- /dev/null +++ b/modelopt/torch/kernels/quantization/attention/v_qdq.py @@ -0,0 +1,129 @@ +# 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. + +"""Value-operand NVFP4 helpers for flash attention.""" + +import math + +import torch +import triton +import triton.language as tl + +from modelopt.torch.kernels.quantization.common.nvfp4_quant import nvfp4_scalar_qdq + +__all__ = ["fake_quant_v_onwrite"] + +_BLOCK_N = 16 + + +@triton.jit +def _v_qdq_nvfp4(v, global_scale, BLOCK_N: tl.constexpr, BLOCK_D: tl.constexpr): + """Fake-quantize signed V in block-16 groups along its key axis.""" + tl.static_assert(BLOCK_N % 16 == 0, "BLOCK_N must be divisible by 16 for NVFP4") + grouped = tl.reshape(v, (BLOCK_N // 16, 16, BLOCK_D)) + block_amax = tl.expand_dims(tl.max(tl.abs(grouped), axis=1), 1) + return tl.reshape(nvfp4_scalar_qdq(grouped, block_amax, global_scale, 16), (BLOCK_N, BLOCK_D)) + + +@triton.jit +def _fake_quant_v_onwrite_kernel( + V_cache, + Block_table, + V_lo, + V_hi, + stride_vc_block, + stride_vc_pos, + stride_vc_head, + v_qdq_scale, + BLOCK_N: tl.constexpr, + BLOCK_D: tl.constexpr, + HEAD_DIM: tl.constexpr, + PAGE_SIZE: tl.constexpr, + max_blocks_per_seq, +): + """Finalize one block-16 V group for one request and KV head.""" + batch_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + v_lo = tl.load(V_lo + batch_idx) + v_hi = tl.load(V_hi + batch_idx) + kv_start = (v_lo // BLOCK_N + tl.program_id(2)) * BLOCK_N + kv_abs = kv_start + tl.arange(0, BLOCK_N) + dim_pos = tl.arange(0, BLOCK_D) + mask = (kv_abs >= v_lo) & (kv_abs < v_hi) + page_local = kv_abs // PAGE_SIZE + page_offset = kv_abs % PAGE_SIZE + page = tl.load(Block_table + batch_idx * max_blocks_per_seq + page_local, mask=mask, other=0) + ptrs = ( + page[:, None].to(tl.int64) * stride_vc_block + + page_offset[:, None] * stride_vc_pos + + kv_head_idx * stride_vc_head + + dim_pos[None, :] + ) + load_mask = mask[:, None] & (dim_pos[None, :] < HEAD_DIM) + v = tl.load(V_cache + ptrs, mask=load_mask, other=0.0).to(tl.float32) + v = _v_qdq_nvfp4(v, v_qdq_scale, BLOCK_N, BLOCK_D) + tl.store(V_cache + ptrs, v.to(V_cache.dtype.element_ty), mask=load_mask) + + +def fake_quant_v_onwrite( + v_cache: torch.Tensor, + block_table: torch.Tensor, + v_lo: torch.Tensor, + v_hi: torch.Tensor, + *, + page_size: int = 16, + v_qdq_scale: float = 1.0, + decode: bool = False, +) -> None: + """NVFP4-finalize complete block-16 groups in ``[v_lo, v_hi)`` in place. + + Decode uses one fixed group per request without reading device metadata on the + host. Eager prefill sizes the grid to cover every newly completed group. + ``v_lo`` and ``v_hi`` must describe aligned, completed block-16 boundaries; + their device values are not host-validated. + """ + if page_size != v_cache.shape[1]: + raise ValueError(f"page_size {page_size} must match v_cache.shape[1] {v_cache.shape[1]}") + if not (math.isfinite(v_qdq_scale) and v_qdq_scale > 0): + raise ValueError(f"v_qdq_scale must be finite and positive, got {v_qdq_scale}") + + batch, max_blocks = block_table.shape + num_kv_heads, head_dim = v_cache.shape[2:] + if decode: + num_groups = 1 + else: + first_group = v_lo // _BLOCK_N + span = int((v_hi - first_group * _BLOCK_N).max().item()) + if span <= 0: + return + num_groups = triton.cdiv(span, _BLOCK_N) + + with torch.cuda.device(v_cache.device): + _fake_quant_v_onwrite_kernel[(batch, num_kv_heads, num_groups)]( + v_cache, + block_table, + v_lo, + v_hi, + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + v_qdq_scale, + BLOCK_N=_BLOCK_N, + BLOCK_D=triton.next_power_of_2(head_dim), + HEAD_DIM=head_dim, + PAGE_SIZE=page_size, + max_blocks_per_seq=max_blocks, + num_warps=4, + ) diff --git a/modelopt/torch/kernels/quantization/common/nvfp4_quant.py b/modelopt/torch/kernels/quantization/common/nvfp4_quant.py index 4aea5a88688..2f0db87f5a4 100644 --- a/modelopt/torch/kernels/quantization/common/nvfp4_quant.py +++ b/modelopt/torch/kernels/quantization/common/nvfp4_quant.py @@ -118,8 +118,9 @@ def fp8_quantize_scale(block_amax, global_scale): FP8-quantized per-block scale(s), same shape as ``block_amax``. """ FP8_E4M3_MAX: tl.constexpr = 448.0 + FP8_E4M3_MIN: tl.constexpr = 2**-9 scale_in_fp8_range = block_amax / (6.0 * global_scale) - scale_clamped = tl.minimum(scale_in_fp8_range, FP8_E4M3_MAX) + scale_clamped = tl.minimum(tl.maximum(scale_in_fp8_range, FP8_E4M3_MIN), FP8_E4M3_MAX) return scale_clamped.to(tl.float8e4nv).to(tl.float32) * global_scale diff --git a/modelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.py b/modelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.py index 742c58a3969..a3686a4fa65 100644 --- a/modelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.py +++ b/modelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.py @@ -81,7 +81,6 @@ def fp4_fake_quant_kernel( block_max = tl.max(x_abs, axis=2, keep_dims=True) block_max_quant = fp8_quantize_scale(block_max, global_scale_safe) - block_max_quant = tl.where(block_max_quant >= 1e-5, block_max_quant, 1.0) block_max_quant_broadcast = tl.broadcast_to( block_max_quant, (TILE_M, NUM_FP4_BLOCKS, BLOCK_SIZE) diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py index f0663dc6b44..5a12a10eba2 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -27,6 +27,11 @@ from modelopt.torch.kernels.common.attention import attention from modelopt.torch.kernels.common.attention.triton_fa import LOG2E +NATIVE_E4M3_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (8, 9) +requires_native_e4m3 = pytest.mark.skipif( + not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute capability >= 8.9" +) + # The kernel runs with a single pinned config under pytest (see _FWD_CONFIGS): # BLOCK_M=128, BLOCK_N=64. The tile-looped reference below relies on it. BLOCK_N = 64 @@ -67,8 +72,8 @@ def _qdq_nvfp4(p, global_scale=1.0): shape = p.shape g = p.reshape(*shape[:-1], shape[-1] // 16, 16) block_amax = g.amax(dim=-1, keepdim=True) # p >= 0, so max == amax - scale = fp8_eager(block_amax / 6.0, torch.tensor(FP8_E4M3_MAX * global_scale, device=p.device)) - scale = torch.where(scale == 0.0, torch.ones_like(scale), scale) + raw_scale = (block_amax / 6.0).clamp(2**-9 * global_scale, FP8_E4M3_MAX * global_scale) + scale = fp8_eager(raw_scale, torch.tensor(FP8_E4M3_MAX * global_scale, device=p.device)) q = _fp4_round(g / scale) * scale return q.reshape(shape) @@ -135,6 +140,7 @@ class TestSoftmaxQdqForward: """Forward correctness of FP8/NVFP4 softmax quant-dequant.""" @pytest.mark.parametrize("mode", ["fp8", "nvfp4"]) + @requires_native_e4m3 def test_prefill_matches_tile_reference(self, mode): """Kernel qdq output matches the tile-looped torch reference.""" seq_len, num_heads, num_kv_heads, head_dim = 128, 4, 2, 64 @@ -153,6 +159,7 @@ def test_prefill_matches_tile_reference(self, mode): assert not torch.equal(o, o_dense) @pytest.mark.parametrize("mode", ["fp8", "nvfp4"]) + @requires_native_e4m3 def test_varlen_partial_tiles(self, mode): """Variable-length batch with partial KV tiles (seq % BLOCK_N != 0).""" seq_lens = [96, 80] @@ -171,6 +178,7 @@ def test_varlen_partial_tiles(self, mode): torch.testing.assert_close(o[s : s + n].float(), ref, rtol=5e-3, atol=5e-3) @pytest.mark.parametrize(("mode", "tol"), [("fp8", 5e-2), ("nvfp4", 0.25)]) + @requires_native_e4m3 def test_qdq_close_to_dense(self, mode, tol): """Quantization is an approximation: output stays near dense attention.""" seq_len, num_heads, num_kv_heads, head_dim = 128, 4, 2, 64 @@ -185,6 +193,7 @@ def test_qdq_close_to_dense(self, mode, tol): torch.testing.assert_close(o, ref, rtol=tol, atol=tol) @pytest.mark.parametrize("mode", ["fp8", "nvfp4"]) + @requires_native_e4m3 def test_decode(self, mode): """Decode (seq_q=1 vs KV cache) matches the non-causal tile reference.""" batch, num_heads, num_kv_heads, head_dim = 2, 4, 2, 64 @@ -232,6 +241,7 @@ def test_decode(self, mode): ("nvfp4", 0.7), ], ) + @requires_native_e4m3 def test_custom_amax_matches_tile_reference(self, mode, amax): """User-supplied p_qdq_amax changes the grid and matches the reference.""" seq_len, num_heads, num_kv_heads, head_dim = 128, 4, 2, 64 @@ -265,6 +275,7 @@ def test_invalid_amax_raises(self): with pytest.raises(ValueError, match="p_qdq_amax"): attention(q, k, v, locs, lens, 8, p_qdq="fp8", p_qdq_amax=0.0) + @requires_native_e4m3 def test_composes_with_skip_softmax(self): """p_qdq composes with the skip-softmax feature in one launch.""" seq_len, num_heads, num_kv_heads, head_dim = 256, 4, 2, 64 @@ -294,8 +305,23 @@ def test_invalid_mode_raises(self): with pytest.raises(ValueError, match="p_qdq"): attention(q, k, v, locs, lens, 8, p_qdq="int8") + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"v_qdq": "int8"}, "v_qdq"), + ({"v_qdq": "nvfp4", "v_qdq_amax": 0.0}, "v_qdq_amax"), + ({"v_qdq": "nvfp4", "v_cache_quantized": True}, "v_cache_quantized"), + ], + ) + def test_invalid_v_qdq_configuration(self, kwargs, match): + q, k, v = make_qkv(8, 2, 2, 32, dtype=torch.float16) + locs, lens = make_varlen_meta([8]) + with pytest.raises(ValueError, match=match): + attention(q, k, v, locs, lens, 8, **kwargs) + @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +@requires_native_e4m3 class TestSoftmaxQdqBackward: """Backward uses the straight-through estimator (no qdq re-applied).""" diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py index 0a51e48a1c7..f4a65366eb4 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py @@ -20,9 +20,16 @@ from conftest import make_qkv, make_varlen_meta from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor if TRITON_KERNEL_AVAILABLE: from modelopt.torch.kernels.common.attention import attention + from modelopt.torch.kernels.quantization.attention.v_qdq import fake_quant_v_onwrite + +NATIVE_E4M3_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (8, 9) +requires_native_e4m3 = pytest.mark.skipif( + not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute capability >= 8.9" +) # --------------------------------------------------------------------------- @@ -113,6 +120,23 @@ def _suffix_causal_reference(q, k, v, q_lens, kv_lens, num_heads, num_kv_heads, class TestPagedKV: """Paged KV cache mode tests — verify paged output matches contiguous.""" + @pytest.mark.parametrize( + ("page_size", "v_qdq_scale", "match"), + [(8, 1.0, "page_size"), (16, 0.0, "v_qdq_scale"), (16, float("inf"), "v_qdq_scale")], + ) + def test_v_onwrite_validates_page_size_and_scale(self, page_size, v_qdq_scale, match): + cache = torch.empty(1, 16, 1, 1, device="cuda") + zeros = torch.zeros(1, device="cuda", dtype=torch.int32) + with pytest.raises(ValueError, match=match): + fake_quant_v_onwrite( + cache, + zeros.view(1, 1), + zeros, + zeros, + page_size=page_size, + v_qdq_scale=v_qdq_scale, + ) + def test_paged_matches_contiguous(self): """Paged mode produces same output as contiguous mode with identical data.""" batch = 2 @@ -495,3 +519,125 @@ def test_paged_chunked_prefill_matches_suffix_causal_reference(self): ref = _suffix_causal_reference(q, k, v, q_lens, kv_lens, num_heads, num_kv_heads, scale) torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-2) + + @requires_native_e4m3 + def test_v_cache_finalizes_complete_groups_once(self): + """Eager prefill and fixed-grid decode bake groups once and leave tails raw.""" + seq_len, head_dim, page_size = 49, 32, 16 + k = torch.zeros(seq_len, 1, head_dim, device="cuda", dtype=torch.float16) + v = torch.full_like(k, 0.017578125) # once: 0.015625; twice: 0.01171875 + locs, lens = make_varlen_meta([seq_len]) + _, raw, block_table = _scatter_to_paged_cache(k, v, locs, lens, 1, head_dim, page_size) + baked = raw.clone() + fake_quant_v_onwrite( + baked, + block_table, + torch.zeros(1, device="cuda", dtype=torch.int32), + torch.tensor([32], device="cuda", dtype=torch.int32), + page_size=page_size, + ) + after_prefill = baked.clone() + fake_quant_v_onwrite( + baked, + block_table, + torch.tensor([32], device="cuda", dtype=torch.int32), + torch.tensor([48], device="cuda", dtype=torch.int32), + page_size=page_size, + decode=True, + ) + torch.testing.assert_close(baked[:2], after_prefill[:2], rtol=0, atol=0) + assert torch.all(baked[:3] == 0.015625) + torch.testing.assert_close(baked[3, 0], raw[3, 0], rtol=0, atol=0) + + tiny = torch.full((1, 16, 1, 16), 2e-6, device="cuda") + fake_quant_v_onwrite( + tiny, + torch.zeros(1, 1, device="cuda", dtype=torch.int32), + torch.zeros(1, device="cuda", dtype=torch.int32), + torch.tensor([16], device="cuda", dtype=torch.int32), + v_qdq_scale=1.0 / (6.0 * 448.0), + ) + assert torch.isfinite(tiny).all() and torch.all(tiny != 0) + + @requires_native_e4m3 + def test_v_cache_matches_independent_signed_key_axis_oracle(self): + page_size, num_kv_heads, head_dim = 8, 2, 4 + logical = ( + torch.arange(16 * num_kv_heads * head_dim, device="cuda", dtype=torch.float16) + .reshape(16, num_kv_heads, head_dim) + .sub_(61) + .div_(13) + ) + cache = torch.full( + (3, page_size, num_kv_heads, head_dim), + 99.0, + device="cuda", + dtype=logical.dtype, + ) + block_table = torch.tensor([[2, 0]], device="cuda", dtype=torch.int32) + cache[2], cache[0] = logical[:page_size], logical[page_size:] + unused_page = cache[1].clone() + + fake_quant_v_onwrite( + cache, + block_table, + torch.zeros(1, device="cuda", dtype=torch.int32), + torch.tensor([16], device="cuda", dtype=torch.int32), + page_size=page_size, + ) + key_last = logical.permute(1, 2, 0).contiguous() + q, scale, double_scale = NVFP4QTensor.quantize( + key_last, + 16, + weights_scaling_factor_2=torch.tensor(1.0, device="cuda"), + try_tensorrt=False, + ) + expected = q.dequantize( + dtype=logical.dtype, + scale=scale, + double_scale=double_scale, + block_sizes={-1: 16}, + ).permute(2, 0, 1) + actual = torch.cat((cache[2], cache[0])) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(cache[1], unused_page, rtol=0, atol=0) + + @pytest.mark.parametrize("q_len", [1, 7]) + @requires_native_e4m3 + def test_baked_prefix_and_raw_tail_match_full_onread(self, q_len): + """The same paged Triton path handles pure decode and chunked prefill.""" + seq_len, num_heads, head_dim, page_size = 17, 2, 32, 16 + q = torch.zeros(q_len, num_heads, head_dim, device="cuda", dtype=torch.float16) + k = torch.zeros(seq_len, 1, head_dim, device="cuda", dtype=torch.float16) + v = torch.full_like(k, 0.017578125) + q_locs, q_lens = make_varlen_meta([q_len]) + locs, lens = make_varlen_meta([seq_len]) + k_cache, raw, block_table = _scatter_to_paged_cache( + k, v, locs, lens, 1, head_dim, page_size + ) + common = { + "is_causal": False, + "b_seq_len_k": lens, + "max_input_len_k": seq_len, + "k_cache": k_cache, + "block_table": block_table, + "page_size": page_size, + "v_qdq": "nvfp4", + } + out_onread = attention(q, k, v, q_locs, q_lens, q_len, v_cache=raw, **common) + baked = raw.clone() + fake_quant_v_onwrite( + baked, block_table, locs, torch.tensor([16], device="cuda", dtype=torch.int32) + ) + out_baked = attention( + q, + k, + v, + q_locs, + q_lens, + q_len, + v_cache=baked, + v_cache_quantized=True, + **common, + ) + torch.testing.assert_close(out_baked, out_onread, rtol=1e-4, atol=1e-5) From d392f96c0bde85c0d5335a427afceaa9889874ac Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Thu, 2 Jul 2026 23:43:48 -0700 Subject: [PATCH 03/28] Add compact NVFP4 vLLM attention worker Signed-off-by: Kai Xu --- .../vllm_serve/quant_sparse_attn_worker.py | 192 ++++++++++++++++++ .../attention_sparsity/plugins/vllm.py | 83 ++++++-- .../test_quant_sparse_attn_worker.py | 165 +++++++++++++++ .../test_sparse_attn_worker.py | 96 ++++++++- 4 files changed, 517 insertions(+), 19 deletions(-) create mode 100644 examples/vllm_serve/quant_sparse_attn_worker.py create mode 100644 tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py diff --git a/examples/vllm_serve/quant_sparse_attn_worker.py b/examples/vllm_serve/quant_sparse_attn_worker.py new file mode 100644 index 00000000000..a0bfb086d68 --- /dev/null +++ b/examples/vllm_serve/quant_sparse_attn_worker.py @@ -0,0 +1,192 @@ +# 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. + +"""Fixed NVFP4 Q/K/P/V worker for ModelOpt sparse attention on vLLM.""" + +import torch +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.attention.backend import AttentionType +from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl +from vllm.v1.worker.gpu_worker import Worker as BaseWorker + +from modelopt.torch.quantization.conversion import set_quantizer_by_cfg +from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer +from modelopt.torch.quantization.plugins.vllm import ( + _ATTENTION_TYPES, + _get_device_dtype, + _set_vllm_attention_kv_default_amax, + disable_compilation, + vllm_attention, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_config import ( + load_from_checkpoint_metadata, + match_sparse_config, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( + _build_sparse_kw, + _clone_sparse_impl, + _p_qdq_from_layer, + _v_qdq_from_layer, +) + +__all__ = ["QuantSparseAttnWorker"] + +VLLMAttention = vllm_attention.Attention +_NVFP4_CFG = { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, +} +_BMM_CFG = [ + {"quantizer_name": "*_bmm_quantizer", "enable": False}, + *( + {"quantizer_name": f"*{name}_bmm_quantizer", "cfg": _NVFP4_CFG, "enable": True} + for name in ("q", "k", "p", "v") + ), +] + + +def _unwrapped_model(worker): + model = worker.model_runner.model + return model.unwrap() if hasattr(model, "unwrap") else model + + +def _global_errors(worker) -> list[str]: + config = worker.model_runner.vllm_config + parallel = config.parallel_config + cache = config.cache_config + model_config = config.model_config + errors = [] + if getattr(parallel, "decode_context_parallel_size", 1) != 1: + errors.append("decode_context_parallel_size must be 1") + if getattr(parallel, "enable_dbo", False) or getattr(parallel, "use_ubatching", False): + errors.append("DBO/ubatching is unsupported") + if getattr(cache, "enable_prefix_caching", False): + errors.append("prefix caching is unsupported") + if getattr(config, "kv_transfer_config", None) is not None: + errors.append("KV transfer is unsupported") + if getattr(config, "speculative_config", None) is not None: + errors.append("speculative decoding is unsupported") + if config.compilation_config.cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL: + errors.append("FULL mixed-batch cudagraph mode is unsupported") + if getattr(model_config, "dtype", None) not in (torch.float16, torch.bfloat16): + errors.append("resolved model/KV-cache dtype must be fp16 or bf16") + cache_dtype = getattr(cache, "cache_dtype", "auto") + if str(cache_dtype) not in {"auto", "bfloat16", "float16", "torch.bfloat16", "torch.float16"}: + errors.append(f"resolved KV-cache dtype {cache_dtype!r} must be fp16 or bf16") + return errors + + +def _layer_errors(module) -> list[str]: + impl = getattr(module, "impl", None) + errors = [] + if type(module) is not VLLMAttention: + errors.append(f"layout {type(module).__name__} is not regular decoder self-attention") + if getattr(module, "attn_type", None) != AttentionType.DECODER: + errors.append("attn_type must be DECODER") + if not isinstance(impl, FlashAttentionImpl): + errors.append(f"backend {type(impl).__name__} is not FlashAttentionImpl") + head_size = getattr(module, "head_size", None) + if not isinstance(head_size, int) or head_size % 16: + errors.append(f"head_size={head_size!r} must be a multiple of 16") + head_size_v = getattr(module, "head_size_v", head_size) + if head_size_v != head_size: + errors.append(f"head_size_v={head_size_v!r} must equal head_size={head_size!r}") + if getattr(module, "sliding_window", None) is not None: + errors.append("sliding_window is unsupported") + if getattr(module, "kv_sharing_target_layer_name", None) is not None: + errors.append("cross-layer KV sharing is unsupported") + if str(getattr(module, "kv_cache_dtype", "")).startswith("fp8"): + errors.append("FP8 KV cache is unsupported") + if getattr(impl, "alibi_slopes", None) is not None: + errors.append("ALiBi is unsupported") + if getattr(impl, "logits_soft_cap", None): + errors.append("logits soft cap is unsupported") + if getattr(impl, "sinks", None) is not None or getattr(impl, "has_sinks", False): + errors.append("attention sinks are unsupported") + return errors + + +def _validated_attention_plans(worker): + """Validate every attention layout without mutation, then return regular-layer tuples.""" + model = _unwrapped_model(worker) + detected = load_from_checkpoint_metadata(worker.model_runner.model_config.hf_config) + sparse_cfg = detected[0] if detected is not None else None + errors = _global_errors(worker) + plans = [] + attention_count = 0 + for name, module in model.named_modules(): + if not isinstance(module, _ATTENTION_TYPES): + continue + attention_count += 1 + reasons = _layer_errors(module) + device, dtype = _get_device_dtype(module) + if device is None or dtype is None: + reasons.append("device/dtype could not be resolved") + elif dtype not in (torch.float16, torch.bfloat16): + reasons.append(f"resolved dtype {dtype} must be fp16 or bf16") + if reasons: + errors.extend(f"{name or ''}: {reason}" for reason in reasons) + continue + layer_cfg = match_sparse_config(name, sparse_cfg) if sparse_cfg is not None else None + sparse_kw = ( + _build_sparse_kw(layer_cfg) + if layer_cfg is not None and layer_cfg.get("enable", True) + else {} + ) + plans.append((name, module, sparse_kw, device, dtype)) + if attention_count == 0: + errors.append("no regular attention layers were found") + if errors: + raise NotImplementedError( + "Unsupported ModelOpt attention plan:\n - " + "\n - ".join(errors) + ) + return plans + + +def _install_quant_sparse_attn(worker) -> None: + plans = _validated_attention_plans(worker) + for _name, module, sparse_kw, device, dtype in plans: + module.device, module.dtype = device, dtype + QuantModuleRegistry.convert(module) + module.p_bmm_quantizer = TensorQuantizer() + set_quantizer_by_cfg(module, _BMM_CFG) + _set_vllm_attention_kv_default_amax(module, device) + new_impl = _clone_sparse_impl(module.impl) + new_impl.sparse_kw = sparse_kw + p_qdq, p_qdq_amax = _p_qdq_from_layer(module) + v_qdq, v_qdq_amax = _v_qdq_from_layer(module) + new_impl.quant_kw = { + "p_qdq": p_qdq, + "p_qdq_amax": p_qdq_amax, + "v_qdq": v_qdq, + "v_qdq_amax": v_qdq_amax, + } + module.impl = new_impl + module._value_quant_in_kernel = True + worker.model_runner.cascade_attn_enabled = False + print(f"[ModelOpt] Installed NVFP4 quant+sparse attention on {len(plans)} layers") + + +class QuantSparseAttnWorker(BaseWorker): + """Install the fixed attention-only recipe before vLLM warmup and graph capture.""" + + @torch.inference_mode() + def determine_available_memory(self) -> int: + with disable_compilation(_unwrapped_model(self)): + return BaseWorker.determine_available_memory(self) + + def compile_or_warm_up_model(self) -> float: + _install_quant_sparse_attn(self) + return BaseWorker.compile_or_warm_up_model(self) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 734755e3bba..6a1aada4cb7 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -38,11 +38,12 @@ ) from modelopt.torch.kernels.common.attention.triton_fa import attention as triton_attention +from modelopt.torch.kernels.quantization.attention.v_qdq import fake_quant_v_onwrite def _target_sparse_ratio_for_phase(target_sparse_ratio, phase: str) -> float: """Return target sparsity for a phase, defaulting old checkpoint metadata.""" - if isinstance(target_sparse_ratio, (float, int)): + if isinstance(target_sparse_ratio, float | int): return float(target_sparse_ratio) if isinstance(target_sparse_ratio, dict): return float(target_sparse_ratio.get(phase, 0.5)) @@ -113,6 +114,34 @@ def _build_sparse_kw(layer_cfg: dict) -> dict: return sparse_kw +def _bmm_qdq_from_layer(layer, attr: str, default_amax: float | None): + """Map an enabled BMM2 quantizer to the kernel's NVFP4 mode and scalar amax.""" + quantizer = getattr(layer, attr, None) + if quantizer is None or not getattr(quantizer, "is_enabled", False): + return None, default_amax + if ( + not getattr(quantizer, "is_nvfp4_dynamic", False) + or (quantizer.block_sizes or {}).get(-1) != 16 + ): + raise NotImplementedError( + f"{attr} is enabled with an unsupported format; only dynamic block-16 NVFP4 is supported" + ) + amax = getattr(quantizer, "_amax", None) + if amax is None: + return "nvfp4", default_amax + if getattr(amax, "numel", lambda: 1)() != 1: + raise NotImplementedError(f"{attr} requires a scalar amax, got shape {tuple(amax.shape)}") + return "nvfp4", float(amax) + + +def _p_qdq_from_layer(layer) -> tuple[str | None, float]: + return _bmm_qdq_from_layer(layer, "p_bmm_quantizer", 1.0) + + +def _v_qdq_from_layer(layer) -> tuple[str | None, float | None]: + return _bmm_qdq_from_layer(layer, "v_bmm_quantizer", None) + + class ModelOptSparseAttentionImpl(FlashAttentionImpl): """Attention implementation that uses the ModelOpt Triton kernel. @@ -166,10 +195,21 @@ def forward( # Profiling run return output.fill_(0) + quant_kw = getattr(self, "quant_kw", None) + if quant_kw is None: + p_qdq, p_qdq_amax = _p_qdq_from_layer(layer) + v_qdq, v_qdq_amax = _v_qdq_from_layer(layer) + else: + p_qdq, p_qdq_amax = quant_kw["p_qdq"], quant_kw["p_qdq_amax"] + v_qdq, v_qdq_amax = quant_kw["v_qdq"], quant_kw["v_qdq_amax"] if getattr(attn_metadata, "use_cascade", False): # vLLM cascade metadata splits the request into shared-prefix and # suffix pieces. The ModelOpt paged kernel consumes plain per-request # KV lengths, so delegate cascade launches back to vLLM's impl. + if p_qdq or v_qdq or getattr(self, "sparse_kw", None): + raise NotImplementedError( + "vLLM cascade attention is incompatible with an active ModelOpt attention transform" + ) return self._forward_vllm_flash_attn( layer, query, @@ -209,23 +249,7 @@ def forward( # N:M sparse softmax is prefill-only. for name in ("sparsity_n", "sparsity_m", "dense_sink_tokens", "dense_recent_tokens"): sparse_kw.pop(name, None) - if set(sparse_kw) <= {"skip_softmax_threshold"}: - # The current ModelOpt paged kernel is only validated for - # sparse prefill in vLLM. Decode-only skip-softmax would route - # through the dense Triton path for every non-skipped tile, so - # keep decode on vLLM FlashAttention until that path is covered. - return self._forward_vllm_flash_attn( - layer, - query, - key, - value, - kv_cache, - attn_metadata, - output, - output_scale, - output_block_scale, - ) - if not sparse_kw: + if not sparse_kw and p_qdq is None and v_qdq is None: # Dynamic calibration can disable sparse work for a launch, e.g. # short-prefill thresholds outside the valid lambda range. Avoid # swapping in the ModelOpt dense kernel when no sparse feature is active. @@ -241,6 +265,24 @@ def forward( output_block_scale, ) + v_cache_quantized = v_qdq == "nvfp4" + if v_cache_quantized: + v_qdq_scale = 1.0 if v_qdq_amax is None else v_qdq_amax / (6.0 * 448.0) + if not (math.isfinite(v_qdq_scale) and v_qdq_scale > 0): + raise ValueError( + f"v_bmm_quantizer amax must be finite and positive, got {v_qdq_amax}" + ) + prev = seq_lens - b_seq_len + fake_quant_v_onwrite( + value_cache, + attn_metadata.block_table, + (prev // 16) * 16, + (seq_lens // 16) * 16, + page_size=page_size, + v_qdq_scale=v_qdq_scale, + decode=is_decode_only, + ) + # Prepare metadata for our kernel q = query[:num_actual_tokens].contiguous() # Dummy K/V for paged mode: not used by the kernel (KV are read from @@ -271,6 +313,11 @@ def forward( v_cache=value_cache, # [num_blocks, page_size, num_kv_heads, head_dim] block_table=attn_metadata.block_table, # [batch, max_blocks] page_size=page_size, # tokens per page in the KV cache + p_qdq=p_qdq, + p_qdq_amax=p_qdq_amax, + v_qdq=v_qdq, + v_qdq_amax=v_qdq_amax, + v_cache_quantized=v_cache_quantized, **sparse_kw, ) diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py new file mode 100644 index 00000000000..52e086db52b --- /dev/null +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Focused tests for the fixed-recipe quant+sparse vLLM worker.""" + +import importlib.util +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.attention.backend import AttentionType +from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl +from vllm.v1.worker.gpu_worker import Worker as BaseWorker + +from modelopt.torch.quantization.plugins import vllm as quant_plugin +from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ModelOptSparseAttentionImpl + +_WORKER_PATH = Path(__file__).parents[5] / "examples/vllm_serve/quant_sparse_attn_worker.py" +_SPEC = importlib.util.spec_from_file_location("quant_sparse_attn_worker", _WORKER_PATH) +worker_module = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(worker_module) + + +def _attention(attn_type=AttentionType.DECODER): + module = object.__new__(quant_plugin.vllm_attention.Attention) + nn.Module.__init__(module) + module.dummy = nn.Parameter(torch.zeros(1, dtype=torch.float16)) + module.impl = object.__new__(FlashAttentionImpl) + module.impl.__dict__.update(alibi_slopes=None, logits_soft_cap=None, sinks=None) + module.__dict__.update( + attn_type=attn_type, + head_size=64, + head_size_v=64, + sliding_window=None, + kv_cache_dtype="auto", + kv_sharing_target_layer_name=None, + ) + return module + + +def _worker(model): + runner = SimpleNamespace( + model=model, model_config=SimpleNamespace(hf_config=None), cascade_attn_enabled=True + ) + return SimpleNamespace(model_runner=runner) + + +def _patch_conversion(monkeypatch): + monkeypatch.setattr( + quant_plugin, + "create_parallel_state", + lambda: quant_plugin.ParallelState(data_parallel_group=None), + ) + monkeypatch.setattr(worker_module, "load_from_checkpoint_metadata", lambda _: None) + monkeypatch.setattr(worker_module, "_global_errors", lambda _: []) + + +def test_install_converts_only_attention_and_configures_fixed_recipe(monkeypatch): + _patch_conversion(monkeypatch) + attention = _attention() + linear = nn.Linear(4, 4) + model = nn.ModuleDict({"attn": attention, "linear": linear}) + state = _worker(model) + + worker_module._install_quant_sparse_attn(state) + + converted = model["attn"] + assert isinstance(converted, quant_plugin._QuantVLLMAttention) + for name in ("q_bmm_quantizer", "k_bmm_quantizer", "p_bmm_quantizer", "v_bmm_quantizer"): + quantizer = getattr(converted, name) + assert quantizer.is_enabled and quantizer.is_nvfp4_dynamic + assert quantizer.block_sizes[-1] == 16 + assert converted.k_bmm_quantizer._amax == 6.0 * 448.0 + assert converted.v_bmm_quantizer._amax == 6.0 * 448.0 + assert converted._value_quant_in_kernel is True + assert isinstance(converted.impl, ModelOptSparseAttentionImpl) + assert converted.impl.quant_kw == { + "p_qdq": "nvfp4", + "p_qdq_amax": 1.0, + "v_qdq": "nvfp4", + "v_qdq_amax": 6.0 * 448.0, + } + assert model["linear"] is linear + assert state.model_runner.cascade_attn_enabled is False + + +def test_validation_of_all_layouts_precedes_mutation(monkeypatch): + _patch_conversion(monkeypatch) + good, bad = _attention(), _attention(AttentionType.ENCODER) + bad.head_size_v = 32 + bad.dummy = nn.Parameter(torch.zeros(1, dtype=torch.float32)) + model = nn.ModuleDict({"good": good, "bad": bad}) + original_impl = good.impl + + with pytest.raises(NotImplementedError) as exc: + worker_module._install_quant_sparse_attn(_worker(model)) + + assert all(needle in str(exc.value) for needle in ("bad: attn_type", "head_size_v", "float32")) + assert model["good"] is good and good.impl is original_impl + assert not isinstance(good, quant_plugin._QuantVLLMAttention) + + +def test_worker_installs_before_base_warmup(monkeypatch): + events = [] + monkeypatch.setattr( + worker_module, "_install_quant_sparse_attn", lambda _: events.append("install") + ) + monkeypatch.setattr( + BaseWorker, "compile_or_warm_up_model", lambda _: events.append("base") or 1.25 + ) + instance = object.__new__(worker_module.QuantSparseAttnWorker) + assert instance.compile_or_warm_up_model() == 1.25 + assert events == ["install", "base"] + + +def test_memory_profile_disables_compilation(monkeypatch): + events, model = [], object() + + @contextmanager + def disabled(actual): + events.append(("enter", actual)) + yield + events.append(("exit", actual)) + + monkeypatch.setattr(worker_module, "disable_compilation", disabled, raising=False) + monkeypatch.setattr( + BaseWorker, "determine_available_memory", lambda _: events.append(("profile", model)) or 7 + ) + instance = object.__new__(worker_module.QuantSparseAttnWorker) + instance.model_runner = SimpleNamespace(model=SimpleNamespace(unwrap=lambda: model)) + assert instance.determine_available_memory() == 7 + assert events == [("enter", model), ("profile", model), ("exit", model)] + + +@pytest.mark.parametrize( + ("mode", "rejected"), + [(CUDAGraphMode.FULL, True), (CUDAGraphMode.FULL_AND_PIECEWISE, False)], +) +def test_full_mixed_cudagraph_validation(mode, rejected): + config = SimpleNamespace( + parallel_config=SimpleNamespace(), + cache_config=SimpleNamespace(), + model_config=SimpleNamespace(dtype=torch.float16), + compilation_config=SimpleNamespace(cudagraph_mode=mode), + ) + errors = worker_module._global_errors( + SimpleNamespace(model_runner=SimpleNamespace(vllm_config=config)) + ) + assert any("mixed" in error for error in errors) is rejected diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index e9f8ee73d95..218c06942e9 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -17,6 +17,7 @@ import math from contextlib import nullcontext +from types import SimpleNamespace import pytest import torch @@ -117,7 +118,6 @@ def fake_forward( @pytest.mark.parametrize( ("sparse_kw", "max_query_len", "max_seq_len"), [ - ({"skip_softmax_threshold": 0.001}, 1, 128), ( { "threshold_scale_factor": { @@ -250,6 +250,100 @@ def fake_attention(q, **kwargs): assert "target_sparse_ratio" not in captured +def test_nvfp4_bmm_mapping_and_unsupported_format_failure(): + """Enabled P/V quantizers map only dynamic block-16 NVFP4.""" + assert vllm_plugin._p_qdq_from_layer(SimpleNamespace()) == (None, 1.0) + assert vllm_plugin._v_qdq_from_layer(SimpleNamespace()) == (None, None) + + good = SimpleNamespace( + is_enabled=True, + is_nvfp4_dynamic=True, + block_sizes={-1: 16}, + _amax=torch.tensor(6.0 * 448.0), + ) + layer = SimpleNamespace(p_bmm_quantizer=good, v_bmm_quantizer=good) + assert vllm_plugin._p_qdq_from_layer(layer) == ("nvfp4", 6.0 * 448.0) + assert vllm_plugin._v_qdq_from_layer(layer) == ("nvfp4", 6.0 * 448.0) + + good.block_sizes = {-1: 32} + with pytest.raises(NotImplementedError, match="p_bmm_quantizer"): + vllm_plugin._p_qdq_from_layer(layer) + with pytest.raises(NotImplementedError, match="v_bmm_quantizer"): + vllm_plugin._v_qdq_from_layer(layer) + + +def test_quantized_decode_finalizes_v_then_calls_shared_paged_kernel(monkeypatch): + """Pure decode uses exact aligned V ranges and the shared Triton FA path.""" + impl = _clone_sparse_impl(_make_old_impl()) + + class UnreadableAmax: + def numel(self): + return 1 + + def __float__(self): + raise AssertionError("forward read live quantizer amax") + + quantizer = SimpleNamespace( + is_enabled=True, + is_nvfp4_dynamic=True, + block_sizes={-1: 16}, + _amax=UnreadableAmax(), + ) + layer = SimpleNamespace(p_bmm_quantizer=quantizer, v_bmm_quantizer=quantizer) + impl.quant_kw = { + "p_qdq": "nvfp4", + "p_qdq_amax": 1.0, + "v_qdq": "nvfp4", + "v_qdq_amax": 6.0 * 448.0, + } + q = torch.zeros(2, impl.num_heads, impl.head_size, dtype=torch.float16) + kv_cache = torch.zeros(2, 4, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + metadata = SimpleNamespace( + num_actual_tokens=2, + max_query_len=1, + max_seq_len=34, + query_start_loc=torch.tensor([0, 1, 2], dtype=torch.int32), + seq_lens=torch.tensor([16, 34], dtype=torch.int32), + block_table=torch.zeros(2, 3, dtype=torch.int32), + ) + calls = {} + + def fake_finalize(value_cache, block_table, v_lo, v_hi, **kwargs): + calls["finalize"] = (v_lo.clone(), v_hi.clone(), kwargs) + + def fake_attention(query, **kwargs): + calls["attention"] = kwargs + return torch.ones_like(query) + + monkeypatch.setattr(vllm_plugin, "fake_quant_v_onwrite", fake_finalize) + monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) + monkeypatch.setattr( + FlashAttentionImpl, + "forward", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("native fallback")), + ) + + output = torch.empty_like(q) + assert impl.forward(layer, q, q, q, kv_cache, metadata, output=output) is output + v_lo, v_hi, finalizer_kw = calls["finalize"] + torch.testing.assert_close(v_lo, torch.tensor([0, 32], dtype=torch.int32)) + torch.testing.assert_close(v_hi, torch.tensor([16, 32], dtype=torch.int32)) + assert finalizer_kw == {"page_size": 16, "v_qdq_scale": 1.0, "decode": True} + assert calls["attention"]["p_qdq"] == "nvfp4" + assert calls["attention"]["v_qdq"] == "nvfp4" + assert calls["attention"]["v_cache_quantized"] is True + + +def test_active_cascade_fails_loud(): + """Cascade must not silently drop an active ModelOpt transform.""" + impl = _clone_sparse_impl(_make_old_impl()) + impl.sparse_kw = {"skip_softmax_threshold": 0.001} + q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16) + cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + with pytest.raises(NotImplementedError, match="cascade"): + impl.forward(None, q, q, q, cache, SimpleNamespace(use_cascade=True), output=q) + + def test_resolve_calibrated_skip_softmax_threshold_for_decode(): """Calibration conversion is phase-aware even when decode later delegates.""" sparse_kw = { From 08369ba94ecbe4c0ca471bb64de545f46bffc60e Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Thu, 2 Jul 2026 23:54:33 -0700 Subject: [PATCH 04/28] Document compact NVFP4 vLLM attention serving Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 22 +++++++++++++++++-- examples/vllm_serve/sparse_attn_worker.py | 4 ++-- examples/vllm_serve/vllm_serve_sparse_attn.py | 10 ++++----- 3 files changed, 27 insertions(+), 9 deletions(-) diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index a8f7d2ead18..352327b118e 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -114,12 +114,30 @@ Workflow: python vllm_serve_sparse_attn.py --enforce-eager -tp 8 --host 0.0.0.0 --port 8000 ``` -If the checkpoint has no `sparse_attention_config`, the worker logs a message and passes through — vLLM runs unchanged. Quant-only flows are handled by `vllm_serve_fakequant.py`; combined sparse + quant will land in a follow-up PR. +If the checkpoint has no `sparse_attention_config`, the worker logs a message and passes through — vLLM runs unchanged. Whole-model fakequant flows remain handled by `vllm_serve_fakequant.py`; the compact attention-only path is below. Limitations: - vLLM V1 chunked prefill and prefix-cache suffix attention are supported by offsetting query positions into the longer KV span. -- CUDA graph capture is not validated yet — use `--enforce-eager`. +- `SparseAttnWorker` CUDA graph capture is not validated yet — use `--enforce-eager`. + +### Compact NVFP4 attention worker + +Use the same launcher with the compact worker and FlashAttention selected explicitly: + +```bash +python vllm_serve_sparse_attn.py -tp 8 \ + --attention-backend FLASH_ATTN \ + --worker-cls quant_sparse_attn_worker.QuantSparseAttnWorker +``` + +This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant recipe to Q/K/P/V. Q is dynamic; uncalibrated K/V use global scale 1.0, P uses amax 1.0, and scalar calibrated K/V amax values are preserved. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored. + +K is QDQ before its cache write, while V is written pristine. Complete 16-token V groups are finalized once in cache; an incomplete tail remains pristine and is QDQ on read. P@V therefore sees uniform fakequant values without re-quantizing the tail. + +Supported configurations are regular decoder self-attention with FlashAttention, fp16/bf16 model and KV cache, equal Q/K/V head dimensions that are multiples of 16, and DCP 1. The default `FULL_AND_PIECEWISE` CUDA graph mode and full decode graphs are supported. + +Unsupported features are sliding window, ALiBi, softcap, sinks, FP8 KV cache, cross/encoder/MLA attention, KV sharing or transfer, prefix caching, speculative decoding, DBO/ubatching, and `FULL` mixed/prefill CUDA graphs. ## Known Problems diff --git a/examples/vllm_serve/sparse_attn_worker.py b/examples/vllm_serve/sparse_attn_worker.py index 1057baa870e..1cd27f4756d 100644 --- a/examples/vllm_serve/sparse_attn_worker.py +++ b/examples/vllm_serve/sparse_attn_worker.py @@ -26,8 +26,8 @@ checkpoint has no such block, the worker logs a message and passes through unchanged. -Quantization combined with sparse attention is not handled by this worker -and will land in a follow-up PR once the combined path is tested. +This worker remains sparse-only; the compact quant+sparse path uses +``quant_sparse_attn_worker.QuantSparseAttnWorker``. Usage: python vllm_serve_sparse_attn.py diff --git a/examples/vllm_serve/vllm_serve_sparse_attn.py b/examples/vllm_serve/vllm_serve_sparse_attn.py index e65ae3e44fb..1f077fe0f37 100644 --- a/examples/vllm_serve/vllm_serve_sparse_attn.py +++ b/examples/vllm_serve/vllm_serve_sparse_attn.py @@ -15,14 +15,14 @@ """Launch vLLM with sparse attention. -Configuration is read exclusively from ``/config.json``'s -``sparse_attention_config`` block, written during calibration by +The default ``SparseAttnWorker`` reads configuration exclusively from +``/config.json``'s ``sparse_attention_config`` block, written by ``examples/llm_sparsity/attention_sparsity/hf_sa.py``. If the checkpoint has -no such block, the worker logs a message and the server runs as standard +no such block, that worker logs a message and the server runs as standard vLLM. -Combined sparse attention + quantization is not handled by this launcher; it -will be added in a follow-up PR once the combined path is tested. +The launcher defaults to ``sparse_attn_worker.SparseAttnWorker``. Pass +``--worker-cls quant_sparse_attn_worker.QuantSparseAttnWorker`` for quant+sparse. Usage: python vllm_serve_sparse_attn.py From 081e7e3097c4051b400fd7b6e84e189e050f399e Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 3 Jul 2026 00:21:30 -0700 Subject: [PATCH 05/28] Harden compact NVFP4 attention numerics Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 3 ++- .../vllm_serve/quant_sparse_attn_worker.py | 1 + .../kernels/common/attention/triton_fa.py | 4 ++++ modelopt/torch/quantization/plugins/vllm.py | 5 +++-- .../attention_sparsity/plugins/vllm.py | 4 ++++ .../common/attention/test_triton_fa_p_qdq.py | 17 +++++++++++++- .../common/attention/test_triton_fa_paged.py | 11 +++++++++- .../quantization/test_vllm_dynamic_modules.py | 19 +++++++++------- .../test_quant_sparse_attn_worker.py | 1 + .../test_sparse_attn_worker.py | 22 ++++++++++++++++--- 10 files changed, 71 insertions(+), 16 deletions(-) diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 352327b118e..3a2f1c6ccd5 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -128,10 +128,11 @@ Use the same launcher with the compact worker and FlashAttention selected explic ```bash python vllm_serve_sparse_attn.py -tp 8 \ --attention-backend FLASH_ATTN \ + --no-enable-prefix-caching \ --worker-cls quant_sparse_attn_worker.QuantSparseAttnWorker ``` -This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant recipe to Q/K/P/V. Q is dynamic; uncalibrated K/V use global scale 1.0, P uses amax 1.0, and scalar calibrated K/V amax values are preserved. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored. +This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant recipe to Q/K/P/V. Q is dynamic, K/V use global scale 1.0, and P uses amax 1.0; calibrated attention amax restore is not part of this fixed path. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored. K is QDQ before its cache write, while V is written pristine. Complete 16-token V groups are finalized once in cache; an incomplete tail remains pristine and is QDQ on read. P@V therefore sees uniform fakequant values without re-quantizing the tail. diff --git a/examples/vllm_serve/quant_sparse_attn_worker.py b/examples/vllm_serve/quant_sparse_attn_worker.py index a0bfb086d68..fe30593bed3 100644 --- a/examples/vllm_serve/quant_sparse_attn_worker.py +++ b/examples/vllm_serve/quant_sparse_attn_worker.py @@ -174,6 +174,7 @@ def _install_quant_sparse_attn(worker) -> None: "v_qdq_amax": v_qdq_amax, } module.impl = new_impl + module._query_quant_in_kernel = True module._value_quant_in_kernel = True worker.model_runner.cascade_attn_enabled = False print(f"[ModelOpt] Installed NVFP4 quant+sparse attention on {len(plans)} layers") diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 7e51ad074aa..7ab93cd181d 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -141,6 +141,7 @@ def _load_paged_k_tile( mask=kv_valid, other=0, ) + page_global = page_global.to(tl.int64) # Load K values: K_cache[page_global, offset_in_page, kv_head_idx, dim] # K^T layout [BLOCK_D, BLOCK_N] for Q @ K^T matmul @@ -184,6 +185,7 @@ def _load_paged_v_tile( mask=kv_valid, other=0, ) + page_global = page_global.to(tl.int64) # V layout [BLOCK_N, BLOCK_D] v_ptrs = ( @@ -1304,6 +1306,8 @@ def attention( f"v_qdq must be one of {sorted(k for k in _QDQ_MODES if k)} or None, got {v_qdq!r}" ) v_qdq_mode = _QDQ_MODES[v_qdq] + if v_qdq_mode and any(t.requires_grad for t in (q, k, v)): + raise NotImplementedError("v_qdq is inference-only and does not support autograd") v_qdq_scale = 1.0 if v_qdq_mode and v_qdq_amax is not None: if not (math.isfinite(v_qdq_amax) and v_qdq_amax > 0): diff --git a/modelopt/torch/quantization/plugins/vllm.py b/modelopt/torch/quantization/plugins/vllm.py index a81b71b57b3..b37e10125d8 100644 --- a/modelopt/torch/quantization/plugins/vllm.py +++ b/modelopt/torch/quantization/plugins/vllm.py @@ -166,7 +166,7 @@ def _get_device_dtype(module: torch.nn.Module) -> tuple: # kv_cache is a list of tensors (v0) or a single tensor (v1). kv = getattr(module, "kv_cache", None) if kv is not None: - t0 = kv[0] if isinstance(kv, (list, tuple)) and len(kv) > 0 else kv + t0 = kv[0] if isinstance(kv, list | tuple) and len(kv) > 0 else kv if isinstance(t0, torch.Tensor) and t0.numel() > 0: spec = getattr(module, "kv_cache_dtype", t0.dtype) out_dtype = ( @@ -523,7 +523,8 @@ def _setup(self): self.parallel_state = create_parallel_state() def forward(self, query, key, value, *args, **kwargs): - query = self.q_bmm_quantizer(query) + if not getattr(self, "_query_quant_in_kernel", False): + query = self.q_bmm_quantizer(query) key = self.k_bmm_quantizer(key) if not getattr(self, "_value_quant_in_kernel", False): value = self.v_bmm_quantizer(value) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 6a1aada4cb7..f289a68104e 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -285,6 +285,10 @@ def forward( # Prepare metadata for our kernel q = query[:num_actual_tokens].contiguous() + if getattr(layer, "_query_quant_in_kernel", False): + valid_q = torch.arange(q.shape[0], device=q.device) < cu_seqlens_q[-1] + q = q.masked_fill(~valid_q[:, None, None], 0) + q = layer.q_bmm_quantizer(q) # Dummy K/V for paged mode: not used by the kernel (KV are read from # k_cache/v_cache via block_table), but shape[1] must be num_kv_heads # so the kernel computes the correct GQA ratio (num_q_heads // num_kv_heads). diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py index 5a12a10eba2..ff99df759e3 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -15,6 +15,8 @@ """GPU tests for the softmax quant-dequant (P_QDQ) feature of the Triton FA kernel.""" +from unittest.mock import Mock + import pytest import torch from conftest import make_qkv, make_varlen_meta, sdpa_reference @@ -24,7 +26,7 @@ from modelopt.torch.quantization.tensor_quant import fp8_eager if TRITON_KERNEL_AVAILABLE: - from modelopt.torch.kernels.common.attention import attention + from modelopt.torch.kernels.common.attention import attention, triton_fa from modelopt.torch.kernels.common.attention.triton_fa import LOG2E NATIVE_E4M3_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (8, 9) @@ -32,6 +34,19 @@ not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute capability >= 8.9" ) + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need triton") +def test_v_qdq_rejects_autograd_before_kernel_launch(monkeypatch): + apply = Mock(side_effect=AssertionError("_Attention.apply reached")) + monkeypatch.setattr(triton_fa._Attention, "apply", apply) + q, k, v = (torch.zeros(1, 1, 16, requires_grad=True) for _ in range(3)) + locs = torch.zeros(1, dtype=torch.int32) + lens = torch.ones(1, dtype=torch.int32) + with pytest.raises(NotImplementedError, match=r"v_qdq.*autograd"): + attention(q, k, v, locs, lens, 1, v_qdq="nvfp4") + apply.assert_not_called() + + # The kernel runs with a single pinned config under pytest (see _FWD_CONFIGS): # BLOCK_M=128, BLOCK_N=64. The tile-looped reference below relies on it. BLOCK_N = 64 diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py index f4a65366eb4..bb33059e65a 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py @@ -15,6 +15,8 @@ """GPU tests for paged KV cache mode of the Triton flash attention kernel.""" +import inspect + import pytest import torch from conftest import make_qkv, make_varlen_meta @@ -23,7 +25,7 @@ from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor if TRITON_KERNEL_AVAILABLE: - from modelopt.torch.kernels.common.attention import attention + from modelopt.torch.kernels.common.attention import attention, triton_fa from modelopt.torch.kernels.quantization.attention.v_qdq import fake_quant_v_onwrite NATIVE_E4M3_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (8, 9) @@ -32,6 +34,13 @@ ) +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need triton") +@pytest.mark.parametrize("loader", ["_load_paged_k_tile", "_load_paged_v_tile"]) +def test_paged_loaders_widen_page_ids_before_pointer_math(loader): + source = inspect.getsource(getattr(triton_fa, loader).fn) + assert "page_global = page_global.to(tl.int64)" in source + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py index 6576c700faf..41fee80eda2 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -109,6 +109,7 @@ def test_attention_setup_keeps_qkv_only_checkpoint_surface(monkeypatch, wrapper_ for name in quantizer_names: getattr(attention, name).amax = torch.tensor(1.0) assert set(attention.state_dict()) == {f"{name}._amax" for name in quantizer_names} + assert not hasattr(attention, "_query_quant_in_kernel") assert not hasattr(attention, "_value_quant_in_kernel") attention.k_bmm_quantizer = _nvfp4_quantizer() @@ -119,7 +120,7 @@ def test_attention_setup_keeps_qkv_only_checkpoint_surface(monkeypatch, wrapper_ assert not hasattr(attention.v_bmm_quantizer, "_amax") -def test_quant_vllm_attention_forward_skips_only_in_kernel_value_quantization(): +def test_quant_vllm_attention_forward_skips_only_in_kernel_qv_quantization(): attention = _new_attention(_TestQuantVLLMAttention) attention.q_bmm_quantizer = Mock(side_effect=lambda inputs: inputs + 1) attention.k_bmm_quantizer = Mock(side_effect=lambda inputs: inputs + 2) @@ -128,18 +129,20 @@ def test_quant_vllm_attention_forward_skips_only_in_kernel_value_quantization(): key = torch.tensor(20) value = torch.tensor(30) + assert not hasattr(attention, "_query_quant_in_kernel") assert not hasattr(attention, "_value_quant_in_kernel") quantized = attention(query, key, value) + attention._query_quant_in_kernel = True + query_in_kernel = attention(query, key, value) attention._value_quant_in_kernel = True - in_kernel = attention(query, key, value) + qv_in_kernel = attention(query, key, value) assert quantized[:3] == (torch.tensor(11), torch.tensor(22), torch.tensor(33)) - assert in_kernel[:3] == (torch.tensor(11), torch.tensor(22), value) - assert tuple(q.call_count for q in (attention.q_bmm_quantizer, attention.k_bmm_quantizer)) == ( - 2, - 2, - ) - assert attention.v_bmm_quantizer.call_count == 1 + assert query_in_kernel[:3] == (query, torch.tensor(22), torch.tensor(33)) + assert qv_in_kernel[:3] == (query, torch.tensor(22), value) + assert attention.q_bmm_quantizer.call_count == 1 + assert attention.k_bmm_quantizer.call_count == 3 + assert attention.v_bmm_quantizer.call_count == 2 @pytest.mark.parametrize( diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py index 52e086db52b..2dd074681af 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py @@ -88,6 +88,7 @@ def test_install_converts_only_attention_and_configures_fixed_recipe(monkeypatch assert quantizer.block_sizes[-1] == 16 assert converted.k_bmm_quantizer._amax == 6.0 * 448.0 assert converted.v_bmm_quantizer._amax == 6.0 * 448.0 + assert converted._query_quant_in_kernel is True assert converted._value_quant_in_kernel is True assert isinstance(converted.impl, ModelOptSparseAttentionImpl) assert converted.impl.quant_kw == { diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 218c06942e9..9bc57366ac2 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -289,17 +289,29 @@ def __float__(self): block_sizes={-1: 16}, _amax=UnreadableAmax(), ) - layer = SimpleNamespace(p_bmm_quantizer=quantizer, v_bmm_quantizer=quantizer) + q_inputs = [] + + def quantize_q(query): + q_inputs.append(query.clone()) + return query + 1 + + layer = SimpleNamespace( + p_bmm_quantizer=quantizer, + q_bmm_quantizer=quantize_q, + v_bmm_quantizer=quantizer, + _query_quant_in_kernel=True, + ) impl.quant_kw = { "p_qdq": "nvfp4", "p_qdq_amax": 1.0, "v_qdq": "nvfp4", "v_qdq_amax": 6.0 * 448.0, } - q = torch.zeros(2, impl.num_heads, impl.head_size, dtype=torch.float16) + q = torch.full((4, impl.num_heads, impl.head_size), 2.0, dtype=torch.float16) + q[2:] = 10_000 kv_cache = torch.zeros(2, 4, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) metadata = SimpleNamespace( - num_actual_tokens=2, + num_actual_tokens=q.shape[0], max_query_len=1, max_seq_len=34, query_start_loc=torch.tensor([0, 1, 2], dtype=torch.int32), @@ -312,6 +324,7 @@ def fake_finalize(value_cache, block_table, v_lo, v_hi, **kwargs): calls["finalize"] = (v_lo.clone(), v_hi.clone(), kwargs) def fake_attention(query, **kwargs): + calls["query"] = query.clone() calls["attention"] = kwargs return torch.ones_like(query) @@ -332,6 +345,9 @@ def fake_attention(query, **kwargs): assert calls["attention"]["p_qdq"] == "nvfp4" assert calls["attention"]["v_qdq"] == "nvfp4" assert calls["attention"]["v_cache_quantized"] is True + assert len(q_inputs) == 1 and q_inputs[0].shape[0] == q.shape[0] + torch.testing.assert_close(q_inputs[0][:2], q[:2]) + assert torch.all(q_inputs[0][2:] == 0) def test_active_cascade_fails_loud(): From 94c60cc4f45546bf57f3b0232d9622068f7080bc Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 3 Jul 2026 11:28:51 -0700 Subject: [PATCH 06/28] Add split-K NVFP4 decode attention Adapt the per-head split-K and FP32 combine structure from internal ModelOpt commit 6c08d08476. Reuse the compact branch paged loaders and NVFP4 P/V QDQ helpers, while preserving the Option-3 pristine V tail. Signed-off-by: Kai Xu --- .../common/attention/decode_attention.py | 350 ++++++++++++++++++ .../attention_sparsity/plugins/vllm.py | 26 ++ .../common/attention/test_decode_attention.py | 132 +++++++ .../test_sparse_attn_worker.py | 68 +++- 4 files changed, 568 insertions(+), 8 deletions(-) create mode 100644 modelopt/torch/kernels/common/attention/decode_attention.py create mode 100644 tests/gpu/torch/kernels/common/attention/test_decode_attention.py diff --git a/modelopt/torch/kernels/common/attention/decode_attention.py b/modelopt/torch/kernels/common/attention/decode_attention.py new file mode 100644 index 00000000000..4f3a8c741fa --- /dev/null +++ b/modelopt/torch/kernels/common/attention/decode_attention.py @@ -0,0 +1,350 @@ +# 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. + + +"""Split-K decode attention for the ModelOpt paged NVFP4 serving path.""" + +import math + +import torch +import triton +import triton.language as tl + +from modelopt.torch.kernels.common.attention.triton_fa import ( + LOG2E, + _load_paged_k_tile, + _load_paged_v_tile, +) +from modelopt.torch.kernels.quantization.attention.p_qdq import _p_qdq_nvfp4 +from modelopt.torch.kernels.quantization.attention.v_qdq import _v_qdq_nvfp4 + +__all__ = ["attention_decode"] + +_BLOCK_N = 128 +_DEFAULT_KV_SPLITS = 32 +_MAX_KV_SPLITS = 32 +_QDQ_MODES = {None, "nvfp4"} + + +@triton.jit +def _decode_split_kernel( + Q, + B_seq_len_k, + M_partial, + L_partial, + Acc_partial, + K_cache, + V_cache, + Block_table, + qk_scale, + stride_qb, + stride_qh, + stride_mb, + stride_mh, + stride_ab, + stride_ah, + stride_as, + stride_kc_block, + stride_kc_pos, + stride_kc_head, + stride_vc_block, + stride_vc_pos, + stride_vc_head, + p_qdq_scale, + v_qdq_scale, + kv_group_num: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, + PAGE_SIZE: tl.constexpr, + max_blocks_per_seq, + NUM_KV_SPLITS: tl.constexpr, + P_QDQ: tl.constexpr, + V_QDQ: tl.constexpr, + V_CACHE_QUANTIZED: tl.constexpr, +): + """Compute one partial softmax for one request, query head, and KV split.""" + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + split_idx = tl.program_id(2) + kv_head_idx = head_idx // kv_group_num + seq_len_kv = tl.load(B_seq_len_k + batch_idx) + v_quantized_boundary = (seq_len_kv // 16) * 16 + + num_tiles = tl.cdiv(seq_len_kv, BLOCK_N) + tiles_per_split = tl.cdiv(num_tiles, NUM_KV_SPLITS) + tile_lo = split_idx * tiles_per_split + tile_hi = tl.minimum(tile_lo + tiles_per_split, num_tiles) + kv_lo = tile_lo * BLOCK_N + kv_hi = tile_hi * BLOCK_N + + dim_pos = tl.arange(0, BLOCK_D) + d_mask = dim_pos < HEAD_DIM + kv_pos = tl.arange(0, BLOCK_N) + q = tl.load( + Q + batch_idx * stride_qb + head_idx * stride_qh + dim_pos, + mask=d_mask, + other=0.0, + ).to(tl.float32) + + running_max = -float("inf") + running_sum = 0.0 + acc = tl.zeros([BLOCK_D], dtype=tl.float32) + + for kv_start in range(kv_lo, kv_hi, BLOCK_N): + kv_start = tl.multiple_of(kv_start, BLOCK_N) + kv_valid = kv_start + kv_pos < seq_len_kv + k = _load_paged_k_tile( + K_cache, + Block_table, + batch_idx, + kv_head_idx, + kv_start, + kv_pos, + dim_pos, + seq_len_kv, + stride_kc_block, + stride_kc_pos, + stride_kc_head, + PAGE_SIZE, + BLOCK_N, + BLOCK_D, + HEAD_DIM, + max_blocks_per_seq, + ).to(tl.float32) + scores = tl.sum(q[:, None] * k, axis=0) * qk_scale + scores = tl.where(kv_valid, scores, -float("inf")) + tile_max = tl.max(scores, axis=0) + new_max = tl.maximum(running_max, tile_max) + p = tl.math.exp2(scores - new_max) + p = tl.where(kv_valid, p, 0.0) + correction = tl.math.exp2(running_max - new_max) + running_sum = running_sum * correction + tl.sum(p, axis=0) + acc *= correction + + if P_QDQ: + p = tl.reshape( + _p_qdq_nvfp4(tl.reshape(p, (1, BLOCK_N)), p_qdq_scale, 1, BLOCK_N), + (BLOCK_N,), + ) + + v = _load_paged_v_tile( + V_cache, + Block_table, + batch_idx, + kv_head_idx, + kv_start, + kv_pos, + dim_pos, + seq_len_kv, + stride_vc_block, + stride_vc_pos, + stride_vc_head, + PAGE_SIZE, + BLOCK_N, + BLOCK_D, + HEAD_DIM, + max_blocks_per_seq, + ).to(tl.float32) + if V_QDQ and ((not V_CACHE_QUANTIZED) or kv_start + BLOCK_N > v_quantized_boundary): + v_qdq = _v_qdq_nvfp4(v, v_qdq_scale, BLOCK_N, BLOCK_D) + if V_CACHE_QUANTIZED: + use_qdq = kv_start + kv_pos >= v_quantized_boundary + v = tl.where(use_qdq[:, None], v_qdq, v) + else: + v = v_qdq + + acc += tl.sum(p[:, None] * v, axis=0) + running_max = new_max + + partial_offset = batch_idx * stride_mb + head_idx * stride_mh + split_idx + tl.store(M_partial + partial_offset, running_max) + tl.store(L_partial + partial_offset, running_sum) + acc_offset = batch_idx * stride_ab + head_idx * stride_ah + split_idx * stride_as + dim_pos + tl.store(Acc_partial + acc_offset, acc, mask=d_mask) + + +@triton.jit +def _decode_combine_kernel( + M_partial, + L_partial, + Acc_partial, + Out, + stride_mb, + stride_mh, + stride_ab, + stride_ah, + stride_as, + stride_ob, + stride_oh, + BLOCK_D: tl.constexpr, + HEAD_DIM: tl.constexpr, + NUM_KV_SPLITS: tl.constexpr, +): + """Merge split-local online-softmax states.""" + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + dim_pos = tl.arange(0, BLOCK_D) + d_mask = dim_pos < HEAD_DIM + base_ml = batch_idx * stride_mb + head_idx * stride_mh + base_acc = batch_idx * stride_ab + head_idx * stride_ah + + running_max = -float("inf") + running_sum = 0.0 + acc = tl.zeros([BLOCK_D], dtype=tl.float32) + for split_idx in range(NUM_KV_SPLITS): + split_sum = tl.load(L_partial + base_ml + split_idx) + if split_sum > 0.0: + split_max = tl.load(M_partial + base_ml + split_idx) + split_acc = tl.load( + Acc_partial + base_acc + split_idx * stride_as + dim_pos, + mask=d_mask, + other=0.0, + ) + new_max = tl.maximum(running_max, split_max) + correction = tl.math.exp2(running_max - new_max) + split_correction = tl.math.exp2(split_max - new_max) + acc = acc * correction + split_acc * split_correction + running_sum = running_sum * correction + split_sum * split_correction + running_max = new_max + + output = acc / tl.maximum(running_sum, 1e-6) + tl.store( + Out + batch_idx * stride_ob + head_idx * stride_oh + dim_pos, + output, + mask=d_mask, + ) + + +def _qdq_scale(mode: str | None, amax: float | None, operand: str) -> float: + if mode not in _QDQ_MODES: + raise ValueError(f"{operand}_qdq must be 'nvfp4' or None, got {mode!r}") + if mode is None: + return 1.0 + if amax is None: + return 1.0 + if not (math.isfinite(amax) and amax > 0.0): + raise ValueError(f"{operand}_qdq_amax must be finite and positive, got {amax}") + return amax / (6.0 * 448.0) + + +def attention_decode( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + b_seq_len_k: torch.Tensor, + *, + softmax_scale: float | None = None, + page_size: int = 16, + num_kv_splits: int = _DEFAULT_KV_SPLITS, + p_qdq: str | None = None, + p_qdq_amax: float = 1.0, + v_qdq: str | None = None, + v_qdq_amax: float | None = None, + v_cache_quantized: bool = False, +) -> torch.Tensor: + """Decode one query token per request over a paged KV cache. + + Q and K are expected to be fake-quantized before this call. P is quantized + per split, while complete block-16 V groups may be finalized in the cache; + only the pristine partial group is then quantized on read. + """ + if q.ndim != 3: + raise ValueError(f"q must have shape [batch, heads, head_dim], got {tuple(q.shape)}") + if page_size != k_cache.shape[1] or page_size != v_cache.shape[1]: + raise ValueError("page_size must match both paged KV cache tensors") + if not 1 <= num_kv_splits <= _MAX_KV_SPLITS: + raise ValueError(f"num_kv_splits must be in [1, {_MAX_KV_SPLITS}], got {num_kv_splits}") + batch, num_q_heads, head_dim = q.shape + num_kv_heads = k_cache.shape[2] + if num_q_heads % num_kv_heads: + raise ValueError("num_q_heads must be divisible by num_kv_heads") + if b_seq_len_k.shape != (batch,) or block_table.shape[0] != batch: + raise ValueError("decode metadata batch dimension must match q") + + p_qdq_scale = _qdq_scale(p_qdq, p_qdq_amax, "p") + v_qdq_scale = _qdq_scale(v_qdq, v_qdq_amax, "v") + if v_cache_quantized and v_qdq != "nvfp4": + raise ValueError("v_cache_quantized requires v_qdq='nvfp4'") + q = q.contiguous() + block_d = triton.next_power_of_2(head_dim) + if (p_qdq == "nvfp4" or v_qdq == "nvfp4") and head_dim % 16: + raise ValueError("NVFP4 decode requires dimensions divisible by 16") + qk_scale = (head_dim**-0.5 if softmax_scale is None else softmax_scale) * LOG2E + m_partial = torch.empty(batch, num_q_heads, num_kv_splits, dtype=torch.float32, device=q.device) + l_partial = torch.empty_like(m_partial) + acc_partial = torch.empty( + batch, num_q_heads, num_kv_splits, block_d, dtype=torch.float32, device=q.device + ) + output = torch.empty_like(q) + + with torch.cuda.device(q.device): + _decode_split_kernel[(batch, num_q_heads, num_kv_splits)]( + q, + b_seq_len_k, + m_partial, + l_partial, + acc_partial, + k_cache, + v_cache, + block_table, + qk_scale, + q.stride(0), + q.stride(1), + m_partial.stride(0), + m_partial.stride(1), + acc_partial.stride(0), + acc_partial.stride(1), + acc_partial.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + p_qdq_scale, + v_qdq_scale, + kv_group_num=num_q_heads // num_kv_heads, + BLOCK_D=block_d, + BLOCK_N=_BLOCK_N, + HEAD_DIM=head_dim, + PAGE_SIZE=page_size, + max_blocks_per_seq=block_table.shape[1], + NUM_KV_SPLITS=num_kv_splits, + P_QDQ=p_qdq == "nvfp4", + V_QDQ=v_qdq == "nvfp4", + V_CACHE_QUANTIZED=v_cache_quantized, + num_warps=4, + num_stages=2, + ) + _decode_combine_kernel[(batch, num_q_heads)]( + m_partial, + l_partial, + acc_partial, + output, + m_partial.stride(0), + m_partial.stride(1), + acc_partial.stride(0), + acc_partial.stride(1), + acc_partial.stride(2), + output.stride(0), + output.stride(1), + BLOCK_D=block_d, + HEAD_DIM=head_dim, + NUM_KV_SPLITS=num_kv_splits, + num_warps=4, + ) + return output diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index f289a68104e..86dba87ba75 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -37,6 +37,9 @@ FlashAttentionMetadata, ) +from modelopt.torch.kernels.common.attention.decode_attention import ( + attention_decode as triton_decode_attention, +) from modelopt.torch.kernels.common.attention.triton_fa import attention as triton_attention from modelopt.torch.kernels.quantization.attention.v_qdq import fake_quant_v_onwrite @@ -289,6 +292,29 @@ def forward( valid_q = torch.arange(q.shape[0], device=q.device) < cu_seqlens_q[-1] q = q.masked_fill(~valid_q[:, None, None], 0) q = layer.q_bmm_quantizer(q) + use_split_k_decode = ( + is_decode_only + and "skip_softmax_threshold" not in sparse_kw + and (p_qdq == "nvfp4" or v_qdq == "nvfp4") + ) + if use_split_k_decode: + triton_out = triton_decode_attention( + q[:batch], + key_cache, + value_cache, + attn_metadata.block_table, + seq_lens, + softmax_scale=self.scale, + page_size=page_size, + p_qdq=p_qdq, + p_qdq_amax=p_qdq_amax, + v_qdq=v_qdq, + v_qdq_amax=v_qdq_amax, + v_cache_quantized=v_cache_quantized, + ) + output[:batch] = triton_out + return output + # Dummy K/V for paged mode: not used by the kernel (KV are read from # k_cache/v_cache via block_table), but shape[1] must be num_kv_heads # so the kernel computes the correct GQA ratio (num_q_heads // num_kv_heads). diff --git a/tests/gpu/torch/kernels/common/attention/test_decode_attention.py b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py new file mode 100644 index 00000000000..06aa0f9af64 --- /dev/null +++ b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py @@ -0,0 +1,132 @@ +# 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. + +"""GPU tests for the minimal paged split-K decode kernel.""" + +import pytest +import torch + +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE + +if TRITON_KERNEL_AVAILABLE: + from modelopt.torch.kernels.common.attention.decode_attention import attention_decode + from modelopt.torch.kernels.quantization.attention.v_qdq import fake_quant_v_onwrite + +NATIVE_E4M3_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (8, 9) +requires_native_e4m3 = pytest.mark.skipif( + not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute capability >= 8.9" +) + + +def _paged_cache(k, v, seq_lens, page_size=16): + batch, num_kv_heads, _, head_dim = k.shape + blocks = [(int(length) + page_size - 1) // page_size for length in seq_lens] + k_cache = torch.zeros( + sum(blocks), page_size, num_kv_heads, head_dim, device=k.device, dtype=k.dtype + ) + v_cache = torch.zeros_like(k_cache) + block_table = torch.zeros(batch, max(blocks), device=k.device, dtype=torch.int32) + physical = 0 + for batch_idx, length in enumerate(seq_lens): + for logical in range(blocks[batch_idx]): + block_table[batch_idx, logical] = physical + start = logical * page_size + stop = min(start + page_size, int(length)) + k_cache[physical, : stop - start] = k[batch_idx, :, start:stop].transpose(0, 1) + v_cache[physical, : stop - start] = v[batch_idx, :, start:stop].transpose(0, 1) + physical += 1 + return k_cache, v_cache, block_table + + +def _dense_decode(q, k, v, seq_lens, scale): + output = torch.empty_like(q) + group_size = q.shape[1] // k.shape[1] + for batch_idx, length in enumerate(seq_lens): + for head_idx in range(q.shape[1]): + kv_head = head_idx // group_size + scores = ( + torch.matmul(k[batch_idx, kv_head, :length].float(), q[batch_idx, head_idx].float()) + * scale + ) + output[batch_idx, head_idx] = torch.matmul( + torch.softmax(scores, dim=0), v[batch_idx, kv_head, :length].float() + ).to(q.dtype) + return output + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") +@pytest.mark.parametrize("num_kv_splits", [1, 8, 32]) +def test_split_k_varlen_gqa_matches_dense(num_kv_splits): + torch.manual_seed(13) + batch, num_q_heads, num_kv_heads, seq_len, head_dim = 2, 8, 2, 511, 64 + seq_lens = (130, seq_len) + q = torch.randn(batch, num_q_heads, head_dim, device="cuda", dtype=torch.float16) + k = torch.randn(batch, num_kv_heads, seq_len, head_dim, device="cuda", dtype=torch.float16) + v = torch.randn_like(k) + k_cache, v_cache, block_table = _paged_cache(k, v, seq_lens) + seq_lens_device = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + scale = head_dim**-0.5 + + output = attention_decode( + q, + k_cache, + v_cache, + block_table, + seq_lens_device, + softmax_scale=scale, + num_kv_splits=num_kv_splits, + ) + + torch.testing.assert_close( + output, _dense_decode(q, k, v, seq_lens, scale), rtol=5e-3, atol=5e-3 + ) + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") +@requires_native_e4m3 +def test_baked_v_prefix_and_pristine_tail_match_full_onread(): + seq_len, num_q_heads, num_kv_heads, head_dim = 17, 4, 1, 64 + q = torch.zeros(1, num_q_heads, head_dim, device="cuda", dtype=torch.float16) + k = torch.zeros(1, num_kv_heads, seq_len, head_dim, device="cuda", dtype=torch.float16) + v = torch.full_like(k, 0.017578125) + seq_lens = torch.tensor([seq_len], device="cuda", dtype=torch.int32) + k_cache, raw_v_cache, block_table = _paged_cache(k, v, (seq_len,)) + common = { + "p_qdq": "nvfp4", + "v_qdq": "nvfp4", + "num_kv_splits": 1, + } + + onread = attention_decode(q, k_cache, raw_v_cache, block_table, seq_lens, **common) + baked_v_cache = raw_v_cache.clone() + fake_quant_v_onwrite( + baked_v_cache, + block_table, + torch.zeros(1, device="cuda", dtype=torch.int32), + torch.tensor([16], device="cuda", dtype=torch.int32), + ) + baked_v_before = baked_v_cache.clone() + baked = attention_decode( + q, + k_cache, + baked_v_cache, + block_table, + seq_lens, + v_cache_quantized=True, + **common, + ) + + torch.testing.assert_close(baked, onread, rtol=0, atol=0) + torch.testing.assert_close(baked_v_cache, baked_v_before, rtol=0, atol=0) diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 9bc57366ac2..7730e5a80de 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -272,8 +272,8 @@ def test_nvfp4_bmm_mapping_and_unsupported_format_failure(): vllm_plugin._v_qdq_from_layer(layer) -def test_quantized_decode_finalizes_v_then_calls_shared_paged_kernel(monkeypatch): - """Pure decode uses exact aligned V ranges and the shared Triton FA path.""" +def test_quantized_decode_finalizes_v_then_calls_split_k_kernel(monkeypatch): + """Pure decode finalizes V before dispatching the valid query rows to split-K.""" impl = _clone_sparse_impl(_make_old_impl()) class UnreadableAmax: @@ -323,13 +323,18 @@ def quantize_q(query): def fake_finalize(value_cache, block_table, v_lo, v_hi, **kwargs): calls["finalize"] = (v_lo.clone(), v_hi.clone(), kwargs) - def fake_attention(query, **kwargs): + def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): calls["query"] = query.clone() - calls["attention"] = kwargs + calls["decode"] = (key_cache, value_cache, block_table, seq_lens, kwargs) return torch.ones_like(query) monkeypatch.setattr(vllm_plugin, "fake_quant_v_onwrite", fake_finalize) - monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) + monkeypatch.setattr(vllm_plugin, "triton_decode_attention", fake_decode) + monkeypatch.setattr( + vllm_plugin, + "triton_attention", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("shared kernel")), + ) monkeypatch.setattr( FlashAttentionImpl, "forward", @@ -342,14 +347,61 @@ def fake_attention(query, **kwargs): torch.testing.assert_close(v_lo, torch.tensor([0, 32], dtype=torch.int32)) torch.testing.assert_close(v_hi, torch.tensor([16, 32], dtype=torch.int32)) assert finalizer_kw == {"page_size": 16, "v_qdq_scale": 1.0, "decode": True} - assert calls["attention"]["p_qdq"] == "nvfp4" - assert calls["attention"]["v_qdq"] == "nvfp4" - assert calls["attention"]["v_cache_quantized"] is True + key_cache, value_cache, block_table, seq_lens, decode_kw = calls["decode"] + assert key_cache.data_ptr() == kv_cache[0].data_ptr() + assert value_cache.data_ptr() == kv_cache[1].data_ptr() + assert block_table is metadata.block_table + assert seq_lens is metadata.seq_lens + assert calls["query"].shape[0] == metadata.seq_lens.shape[0] + assert decode_kw["p_qdq"] == "nvfp4" + assert decode_kw["v_qdq"] == "nvfp4" + assert decode_kw["v_cache_quantized"] is True assert len(q_inputs) == 1 and q_inputs[0].shape[0] == q.shape[0] torch.testing.assert_close(q_inputs[0][:2], q[:2]) assert torch.all(q_inputs[0][2:] == 0) +def test_quantized_skip_softmax_decode_stays_on_shared_kernel(monkeypatch): + """Split-local maxima must not change calibrated skip-softmax semantics.""" + impl = _clone_sparse_impl(_make_old_impl()) + impl.quant_kw = { + "p_qdq": "nvfp4", + "p_qdq_amax": 1.0, + "v_qdq": "nvfp4", + "v_qdq_amax": 6.0 * 448.0, + } + impl.sparse_kw = {"skip_softmax_threshold": 0.001} + q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16) + kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + metadata = SimpleNamespace( + num_actual_tokens=1, + max_query_len=1, + max_seq_len=16, + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + seq_lens=torch.tensor([16], dtype=torch.int32), + block_table=torch.zeros(1, 1, dtype=torch.int32), + ) + captured = {} + + monkeypatch.setattr(vllm_plugin, "fake_quant_v_onwrite", lambda *args, **kwargs: None) + monkeypatch.setattr( + vllm_plugin, + "triton_decode_attention", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("split-K kernel")), + ) + + def fake_attention(query, **kwargs): + captured.update(kwargs) + return torch.zeros_like(query) + + monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) + output = torch.empty_like(q) + assert impl.forward(None, q, q, q, kv_cache, metadata, output=output) is output + assert captured["skip_softmax_threshold"] == pytest.approx(0.001) + assert captured["p_qdq"] == "nvfp4" + assert captured["v_qdq"] == "nvfp4" + + def test_active_cascade_fails_loud(): """Cascade must not silently drop an active ModelOpt transform.""" impl = _clone_sparse_impl(_make_old_impl()) From 50e1e0e75f035f938b51b51f235c1d32e1a92af4 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 3 Jul 2026 14:27:20 -0700 Subject: [PATCH 07/28] Match NVFP4 attention fake quant to native numerics Signed-off-by: Kai Xu --- .../common/attention/decode_attention.py | 15 ++-- .../kernels/common/attention/triton_fa.py | 18 ++++- .../quantization/common/nvfp4_quant.py | 15 ++-- .../attention_sparsity/plugins/vllm.py | 2 +- .../common/attention/test_decode_attention.py | 50 +++++++++++++- .../common/attention/test_triton_fa_p_qdq.py | 68 +++++++++++++++++-- .../common/attention/test_triton_fa_paged.py | 49 ++++++++++++- .../quantization/test_tensor_quant_cuda.py | 42 ++++++++++++ .../test_quant_sparse_attn_worker.py | 1 + .../test_sparse_attn_worker.py | 4 +- 10 files changed, 239 insertions(+), 25 deletions(-) diff --git a/modelopt/torch/kernels/common/attention/decode_attention.py b/modelopt/torch/kernels/common/attention/decode_attention.py index 4f3a8c741fa..377b6450161 100644 --- a/modelopt/torch/kernels/common/attention/decode_attention.py +++ b/modelopt/torch/kernels/common/attention/decode_attention.py @@ -136,7 +136,12 @@ def _decode_split_kernel( if P_QDQ: p = tl.reshape( - _p_qdq_nvfp4(tl.reshape(p, (1, BLOCK_N)), p_qdq_scale, 1, BLOCK_N), + _p_qdq_nvfp4( + tl.reshape(p.to(V_cache.dtype.element_ty).to(tl.float32), (1, BLOCK_N)), + p_qdq_scale, + 1, + BLOCK_N, + ), (BLOCK_N,), ) @@ -258,9 +263,11 @@ def attention_decode( ) -> torch.Tensor: """Decode one query token per request over a paged KV cache. - Q and K are expected to be fake-quantized before this call. P is quantized - per split, while complete block-16 V groups may be finalized in the cache; - only the pristine partial group is then quantized on read. + Q and K are expected to be fake-quantized before this call. Dynamic NVFP4 + Q should use an FP32 QDQ carrier; K may remain BF16 when its global scale is + one. P is rounded to the model/cache dtype before native-style quantization, + then its QDQ result remains FP32. Complete block-16 V groups may be finalized + in the cache; only the pristine partial group is then quantized on read. """ if q.ndim != 3: raise ValueError(f"q must have shape [batch, heads, head_dim], got {tuple(q.shape)}") diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 7ab93cd181d..8755757cf55 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -260,6 +260,7 @@ def _attn_fwd( IS_CAUSAL: tl.constexpr, # Whether to apply causal mask HEAD_DIM: tl.constexpr, # Actual head dimension (for d_mask) STORE_LSE: tl.constexpr, # Whether to save LSE for backward pass + Q_IS_FP32: tl.constexpr, # Dynamic NVFP4 QDQ carrier uses FP32 SPARSITY_N: tl.constexpr = 0, # N:M sparsity — keep top-N of every M elements (0 = disabled) SPARSITY_M: tl.constexpr = 4, # N:M sparsity — group size (4 or 8) DENSE_SINK_TOKENS: tl.constexpr = 0, # Leading KV tokens kept dense (attention sinks) @@ -366,7 +367,10 @@ def _attn_fwd( ) # scores = Q @ K^T * scale [BLOCK_M, BLOCK_N] - scores = tl.dot(q, k) * qk_scale + if Q_IS_FP32: + scores = tl.dot(q, k.to(tl.float32), input_precision="ieee") * qk_scale + else: + scores = tl.dot(q, k) * qk_scale scores = _apply_mask(scores, q_pos, kv_pos, seq_len_q, seq_len_kv, kv_start, IS_CAUSAL) # --- Optional N:M sparse softmax --- @@ -415,6 +419,12 @@ def _attn_fwd( if P_QDQ == 1: p = _qdq_fp8(p, p_qdq_scale) elif P_QDQ == 2: + # Native packing consumes the model dtype, but its QDQ value + # remains FP32 for the scaled MMA accumulation. + if IS_PAGED: + p = p.to(V_cache.dtype.element_ty).to(tl.float32) + else: + p = p.to(V.dtype.element_ty).to(tl.float32) p = _p_qdq_nvfp4(p, p_qdq_scale, BLOCK_M, BLOCK_N) # Load V and accumulate @@ -457,7 +467,10 @@ def _attn_fwd( v = tl.where(use_qdq[:, None], v_qdq, v) else: v = v_qdq - acc = tl.dot(p.to(v.dtype), v, acc) + if P_QDQ == 2: + acc = tl.dot(p, v.to(tl.float32), acc, input_precision="ieee") + else: + acc = tl.dot(p.to(v.dtype), v, acc) row_max = m_new # else: tile skipped — no softmax, no V load, no BMM2 for this tile @@ -949,6 +962,7 @@ def forward( "IS_CAUSAL": is_causal, "HEAD_DIM": HEAD_DIM, "STORE_LSE": True, + "Q_IS_FP32": q.dtype == torch.float32 and (p_qdq_mode == 2 or v_qdq_mode == 2), "SPARSITY_N": sparsity_n, "SPARSITY_M": sparsity_m, "DENSE_SINK_TOKENS": dense_sink_tokens, diff --git a/modelopt/torch/kernels/quantization/common/nvfp4_quant.py b/modelopt/torch/kernels/quantization/common/nvfp4_quant.py index 2f0db87f5a4..85b73fe6388 100644 --- a/modelopt/torch/kernels/quantization/common/nvfp4_quant.py +++ b/modelopt/torch/kernels/quantization/common/nvfp4_quant.py @@ -73,6 +73,7 @@ def nvfp4_scalar_quant( Quantizes each element independently: divide by scale, round to nearest FP4 (E2M1) value via ``fp4_round_magnitude``, multiply by scale. + A zero scale produces an all-zero block. All ops are element-wise, so any shape works with a broadcastable ``scale`` (e.g. a [M, B, 16] tile with [M, B, 1] per-block scales, as @@ -87,9 +88,9 @@ def nvfp4_scalar_quant( x_quant: [N] float32, fake-quantized values. """ x_abs = tl.abs(x) - # Guard against degenerate scale (matching CUDA kernel behavior) + zero_scale = scale == 0.0 scale_safe = tl.where( - (scale == 0.0) | libdevice.isnan(scale) | (tl.abs(scale) == float("inf")), + zero_scale | libdevice.isnan(scale) | (tl.abs(scale) == float("inf")), 1.0, scale, ) @@ -97,7 +98,7 @@ def nvfp4_scalar_quant( q_val = fp4_round_magnitude(abs_scaled) x_rescaled = q_val * scale_safe x_quant = tl.where(x >= 0, x_rescaled, -x_rescaled) - return x_quant + return tl.where(zero_scale, 0.0, x_quant) @triton.jit @@ -105,7 +106,8 @@ def fp8_quantize_scale(block_amax, global_scale): """FP8 E4M3 fake-quantize the per-block NVFP4 scale. Computes ``scale = block_amax / 6.0``, then round-trips it through - FP8 E4M3 using ``global_scale`` for the second-level scaling. + FP8 E4M3 using ``global_scale`` for the second-level scaling. Values above + the E4M3 maximum saturate; values below its rounding range may become zero. Works with any tensor shape (scalar, 1-D, or higher) since all ops are element-wise. @@ -118,10 +120,9 @@ def fp8_quantize_scale(block_amax, global_scale): FP8-quantized per-block scale(s), same shape as ``block_amax``. """ FP8_E4M3_MAX: tl.constexpr = 448.0 - FP8_E4M3_MIN: tl.constexpr = 2**-9 scale_in_fp8_range = block_amax / (6.0 * global_scale) - scale_clamped = tl.minimum(tl.maximum(scale_in_fp8_range, FP8_E4M3_MIN), FP8_E4M3_MAX) - return scale_clamped.to(tl.float8e4nv).to(tl.float32) * global_scale + scale_saturated = tl.minimum(scale_in_fp8_range, FP8_E4M3_MAX) + return scale_saturated.to(tl.float8e4nv).to(tl.float32) * global_scale @triton.jit diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 86dba87ba75..bc5d0792a78 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -291,7 +291,7 @@ def forward( if getattr(layer, "_query_quant_in_kernel", False): valid_q = torch.arange(q.shape[0], device=q.device) < cu_seqlens_q[-1] q = q.masked_fill(~valid_q[:, None, None], 0) - q = layer.q_bmm_quantizer(q) + q = layer.q_bmm_quantizer(q.float()) use_split_k_decode = ( is_decode_only and "skip_softmax_threshold" not in sparse_kw diff --git a/tests/gpu/torch/kernels/common/attention/test_decode_attention.py b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py index 06aa0f9af64..cee3e03b16f 100644 --- a/tests/gpu/torch/kernels/common/attention/test_decode_attention.py +++ b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py @@ -15,6 +15,8 @@ """GPU tests for the minimal paged split-K decode kernel.""" +import math + import pytest import torch @@ -66,6 +68,19 @@ def _dense_decode(q, k, v, seq_lens, scale): return output +def _nvfp4_qdq_reference(x, global_scale=1.0 / (6.0 * 448.0)): + blocks = x.reshape(-1, 16) + block_amax = blocks.abs().amax(dim=-1, keepdim=True) + scales = (block_amax / (6.0 * global_scale)).clamp(max=448.0) + scales = scales.to(torch.float8_e4m3fn).float() * global_scale + scale_safe = torch.where(scales == 0, 1.0, scales) + levels = x.new_tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + scaled = blocks.abs() / scale_safe + quantized = levels[(scaled[..., None] - levels).abs().argmin(dim=-1)] * scale_safe + quantized = torch.copysign(quantized, blocks) + return torch.where(scales == 0, 0.0, quantized).reshape_as(x) + + @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") @pytest.mark.parametrize("num_kv_splits", [1, 8, 32]) def test_split_k_varlen_gqa_matches_dense(num_kv_splits): @@ -98,7 +113,7 @@ def test_split_k_varlen_gqa_matches_dense(num_kv_splits): @requires_native_e4m3 def test_baked_v_prefix_and_pristine_tail_match_full_onread(): seq_len, num_q_heads, num_kv_heads, head_dim = 17, 4, 1, 64 - q = torch.zeros(1, num_q_heads, head_dim, device="cuda", dtype=torch.float16) + q = torch.zeros(1, num_q_heads, head_dim, device="cuda", dtype=torch.float32) k = torch.zeros(1, num_kv_heads, seq_len, head_dim, device="cuda", dtype=torch.float16) v = torch.full_like(k, 0.017578125) seq_lens = torch.tensor([seq_len], device="cuda", dtype=torch.int32) @@ -130,3 +145,36 @@ def test_baked_v_prefix_and_pristine_tail_match_full_onread(): torch.testing.assert_close(baked, onread, rtol=0, atol=0) torch.testing.assert_close(baked_v_cache, baked_v_before, rtol=0, atol=0) + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") +@requires_native_e4m3 +def test_p_quantizes_model_dtype_input_but_accumulates_fp32(): + seq_len, head_dim = 128, 16 + scale = head_dim**-0.5 + q = torch.zeros(1, 1, head_dim, device="cuda", dtype=torch.float32) + boundary_p = 0.04163 + q[..., 0] = -math.log2(boundary_p) / (scale * 1.4426950408889634) + k = torch.zeros(1, 1, seq_len, head_dim, device="cuda", dtype=torch.bfloat16) + k[:, :, 1:, 0] = -1.0 + v = torch.zeros_like(k) + v[:, :, 1:16, 0] = 1.0 + k_cache, v_cache, block_table = _paged_cache(k, v, (seq_len,)) + seq_lens = torch.tensor([seq_len], device="cuda", dtype=torch.int32) + + output = attention_decode( + q, + k_cache, + v_cache, + block_table, + seq_lens, + softmax_scale=scale, + num_kv_splits=1, + p_qdq="nvfp4", + ) + scores = torch.matmul(k[0, 0].float(), q[0, 0]) * (scale * 1.4426950408889634) + p = torch.exp2(scores - scores.max()) + p_qdq = _nvfp4_qdq_reference(p.to(torch.bfloat16).float()) + reference = (p_qdq[:, None] * v[0, 0].float()).sum(0) / p.sum() + + torch.testing.assert_close(output[0, 0], reference, rtol=5e-5, atol=5e-5) diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py index ff99df759e3..f144e7c4a94 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -15,6 +15,7 @@ """GPU tests for the softmax quant-dequant (P_QDQ) feature of the Triton FA kernel.""" +import math from unittest.mock import Mock import pytest @@ -87,9 +88,11 @@ def _qdq_nvfp4(p, global_scale=1.0): shape = p.shape g = p.reshape(*shape[:-1], shape[-1] // 16, 16) block_amax = g.amax(dim=-1, keepdim=True) # p >= 0, so max == amax - raw_scale = (block_amax / 6.0).clamp(2**-9 * global_scale, FP8_E4M3_MAX * global_scale) - scale = fp8_eager(raw_scale, torch.tensor(FP8_E4M3_MAX * global_scale, device=p.device)) - q = _fp4_round(g / scale) * scale + scale_in_fp8_range = (block_amax / (6.0 * global_scale)).clamp(max=FP8_E4M3_MAX) + scale = scale_in_fp8_range.to(torch.float8_e4m3fn).float() * global_scale + scale_safe = torch.where(scale == 0, 1.0, scale) + q = _fp4_round(g / scale_safe) * scale_safe + q = torch.where(scale == 0, 0.0, q) return q.reshape(shape) @@ -106,7 +109,9 @@ def qdq_attention_reference(q, k, v, scale, mode, is_causal=True, amax=1.0): Single sequence: q [s, h, d], k/v [s_kv, h_kv, d] (fp16). Walks KV tiles of BLOCK_N exactly like the kernel, keeps the softmax denominator unquantized, applies qdq to the unnormalized p of each tile, and mirrors - the kernel's ``p.to(v.dtype)`` cast before the P @ V dot. + the kernel's operand carrier before the P @ V dot. Native NVFP4 rounds P + to the model dtype before packing, then keeps the QDQ result in FP32. FP8 + retains the legacy post-QDQ cast to V's dtype. Returns [s, h, d] float32. ``amax`` mirrors the kernel's ``p_qdq_amax`` and is converted to the same @@ -141,9 +146,11 @@ def qdq_attention_reference(q, k, v, scale, mode, is_causal=True, amax=1.0): corr = torch.exp2(row_max - m_new) row_sum = row_sum * corr + l_new acc = acc * corr[..., None] + if mode == "nvfp4": + p = p.to(v.dtype).float() p = _apply_qdq(p, mode, qdq_scale) - # Kernel casts p to v.dtype for the BMM2 dot - p = p.to(v.dtype).float() + if mode == "fp8": + p = p.to(v.dtype).float() acc = acc + torch.einsum("hqk,khd->hqd", p, vv[start : start + BLOCK_N].float()) row_max = m_new out = acc / row_sum[..., None] @@ -173,6 +180,50 @@ def test_prefill_matches_tile_reference(self, mode): o_dense = attention(q, k, v, locs, lens, seq_len, softmax_scale=scale) assert not torch.equal(o, o_dense) + @requires_native_e4m3 + def test_nvfp4_uses_native_p_input_and_fp32_carriers(self): + """Dynamic-Q and P QDQ stay FP32 until their native MMA boundaries.""" + num_heads = num_kv_heads = 1 + head_dim = 16 + seq_len_kv = BLOCK_N + scale = 1.0 / math.sqrt(head_dim) + + q = torch.zeros(1, num_heads, head_dim, device="cuda", dtype=torch.float32) + boundary_p = 0.1248 + q[..., 0] = -math.log2(boundary_p) / (scale * LOG2E) + k = torch.zeros(seq_len_kv, num_kv_heads, head_dim, device="cuda", dtype=torch.bfloat16) + # The first block contains the tile max and P at the 1/24 decision + # boundary. Rounding P to BF16 before packing moves it to the adjacent + # E2M1 value; rounding the QDQ result afterward does not model this. + k[1:, 0, 0] = -1.0 + v = torch.zeros_like(k) + v[1:16, 0, 0] = 1.0 + q_locs = torch.zeros(1, device="cuda", dtype=torch.int32) + q_lens = torch.ones(1, device="cuda", dtype=torch.int32) + kv_locs = torch.zeros(1, device="cuda", dtype=torch.int32) + kv_lens = torch.full((1,), seq_len_kv, device="cuda", dtype=torch.int32) + + out = attention( + q, + k, + v, + q_locs, + q_lens, + 1, + is_causal=False, + softmax_scale=scale, + b_start_loc_k=kv_locs, + b_seq_len_k=kv_lens, + max_input_len_k=seq_len_kv, + p_qdq="nvfp4", + ) + reference = qdq_attention_reference(q, k, v, scale, "nvfp4", is_causal=False) + + # The Triton exp2 approximation shifts the unquantized denominator + # slightly at this decision-boundary case; the old FP32 pack input is + # still separated by more than two orders of magnitude. + torch.testing.assert_close(out, reference, rtol=2e-3, atol=5e-4) + @pytest.mark.parametrize("mode", ["fp8", "nvfp4"]) @requires_native_e4m3 def test_varlen_partial_tiles(self, mode): @@ -190,7 +241,10 @@ def test_varlen_partial_tiles(self, mode): for b, n in enumerate(seq_lens): s = int(locs[b].item()) ref = qdq_attention_reference(q[s : s + n], k[s : s + n], v[s : s + n], scale, mode) - torch.testing.assert_close(o[s : s + n].float(), ref, rtol=5e-3, atol=5e-3) + # BF16/FP16 pre-pack rounding makes NVFP4 code selection more + # sensitive to Triton-vs-PyTorch exp2 differences near thresholds. + atol = 2e-2 if mode == "nvfp4" else 5e-3 + torch.testing.assert_close(o[s : s + n].float(), ref, rtol=5e-3, atol=atol) @pytest.mark.parametrize(("mode", "tol"), [("fp8", 5e-2), ("nvfp4", 0.25)]) @requires_native_e4m3 diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py index bb33059e65a..202137c4c07 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py @@ -16,6 +16,7 @@ """GPU tests for paged KV cache mode of the Triton flash attention kernel.""" import inspect +import math import pytest import torch @@ -566,7 +567,8 @@ def test_v_cache_finalizes_complete_groups_once(self): torch.tensor([16], device="cuda", dtype=torch.int32), v_qdq_scale=1.0 / (6.0 * 448.0), ) - assert torch.isfinite(tiny).all() and torch.all(tiny != 0) + # Native E4M3 block scales below half the minimum subnormal round to zero. + assert torch.isfinite(tiny).all() and torch.count_nonzero(tiny) == 0 @requires_native_e4m3 def test_v_cache_matches_independent_signed_key_axis_oracle(self): @@ -616,7 +618,7 @@ def test_v_cache_matches_independent_signed_key_axis_oracle(self): def test_baked_prefix_and_raw_tail_match_full_onread(self, q_len): """The same paged Triton path handles pure decode and chunked prefill.""" seq_len, num_heads, head_dim, page_size = 17, 2, 32, 16 - q = torch.zeros(q_len, num_heads, head_dim, device="cuda", dtype=torch.float16) + q = torch.zeros(q_len, num_heads, head_dim, device="cuda", dtype=torch.float32) k = torch.zeros(seq_len, 1, head_dim, device="cuda", dtype=torch.float16) v = torch.full_like(k, 0.017578125) q_locs, q_lens = make_varlen_meta([q_len]) @@ -631,6 +633,7 @@ def test_baked_prefix_and_raw_tail_match_full_onread(self, q_len): "k_cache": k_cache, "block_table": block_table, "page_size": page_size, + "p_qdq": "nvfp4", "v_qdq": "nvfp4", } out_onread = attention(q, k, v, q_locs, q_lens, q_len, v_cache=raw, **common) @@ -650,3 +653,45 @@ def test_baked_prefix_and_raw_tail_match_full_onread(self, q_len): **common, ) torch.testing.assert_close(out_baked, out_onread, rtol=1e-4, atol=1e-5) + + @requires_native_e4m3 + def test_p_qdq_uses_cache_dtype_with_fp32_dummy_tensors(self): + seq_len, head_dim, page_size = 128, 16, 16 + scale = head_dim**-0.5 + boundary_p = 0.1248 + q = torch.zeros(1, 1, head_dim, device="cuda", dtype=torch.float32) + q[..., 0] = -math.log2(boundary_p) / (scale * triton_fa.LOG2E) + k = torch.zeros(seq_len, 1, head_dim, device="cuda", dtype=torch.bfloat16) + k[1:, 0, 0] = -1.0 + v = torch.zeros_like(k) + v[1:16, 0, 0] = 1.0 + q_locs, q_lens = make_varlen_meta([1]) + kv_locs, kv_lens = make_varlen_meta([seq_len]) + k_cache, v_cache, block_table = _scatter_to_paged_cache( + k, v, kv_locs, kv_lens, 1, head_dim, page_size + ) + common = { + "is_causal": False, + "softmax_scale": scale, + "b_start_loc_k": kv_locs, + "b_seq_len_k": kv_lens, + "max_input_len_k": seq_len, + "p_qdq": "nvfp4", + } + contiguous = attention(q, k, v, q_locs, q_lens, 1, **common) + dummy = torch.empty(0, 1, head_dim, device="cuda", dtype=torch.float32) + paged = attention( + q, + dummy, + dummy, + q_locs, + q_lens, + 1, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + page_size=page_size, + **common, + ) + + torch.testing.assert_close(paged, contiguous, rtol=2e-3, atol=5e-4) diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index d2503669ac9..cf802abaf40 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -15,6 +15,8 @@ """Tests of tensor quantization function and module""" +import os + import pytest import torch from _test_utils.torch.quantization.quant_utils import quant @@ -26,6 +28,24 @@ from modelopt.torch.quantization.extensions import get_cuda_ext, get_cuda_ext_mx from modelopt.torch.quantization.tensor_quant import mx_format_map +if triton_kernel.IS_AVAILABLE: + import triton + import triton.language as tl + + from modelopt.torch.kernels.quantization.common.nvfp4_quant import fp8_quantize_scale + + @triton.jit + def _fp8_quantize_scale_test_kernel(block_amax_ptr, output_ptr, global_scale_ptr): + block_amax = tl.load(block_amax_ptr).to(tl.float32) + global_scale = tl.load(global_scale_ptr).to(tl.float32) + output = fp8_quantize_scale(block_amax, global_scale) + tl.store(output_ptr, output) + + +NATIVE_E4M3_AVAILABLE = triton_kernel.IS_AVAILABLE and ( + os.environ.get("TRITON_INTERPRET") == "1" or torch.cuda.get_device_capability() >= (8, 9) +) + class TestFakeTensorQuantCuda(FakeTensorQuantTester): device = "cuda" @@ -160,6 +180,28 @@ def test_zero_amax_per_channel_is_finite(self): class Testfp4: + @pytest.mark.skipif(not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute >= 8.9") + def test_native_block_scale_underflows_to_zero(self): + raw_scale = 2**-11 # Below half the smallest E4M3 subnormal (2**-9). + block_amax = torch.tensor(6.0 * raw_scale, device="cuda") + output = torch.empty(1, device="cuda") + global_scale = torch.tensor(1.0, device="cuda") + + _fp8_quantize_scale_test_kernel[(1,)](block_amax, output, global_scale) + + assert torch.equal(output, torch.zeros_like(output)) + + @pytest.mark.skipif(not triton_kernel.IS_AVAILABLE, reason="triton kernel is not available") + def test_zero_block_scale_zeroes_block(self): + x = torch.ones((1, 16), device="cuda") + output = triton_kernel.static_blockwise_fp4_fake_quant( + x, + amax=torch.zeros(1, device="cuda"), + quantize_block_scales=False, + ) + + assert torch.equal(output, torch.zeros_like(output)) + @pytest.mark.skipif(get_cuda_ext_mx() is None, reason="cuda_ext_mx is not available") @pytest.mark.parametrize( "set_torch_dtype", [torch.float, torch.float16, torch.bfloat16], indirect=True diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py index 2dd074681af..7783ad2fae8 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py @@ -86,6 +86,7 @@ def test_install_converts_only_attention_and_configures_fixed_recipe(monkeypatch quantizer = getattr(converted, name) assert quantizer.is_enabled and quantizer.is_nvfp4_dynamic assert quantizer.block_sizes[-1] == 16 + assert not hasattr(converted.q_bmm_quantizer, "_amax") assert converted.k_bmm_quantizer._amax == 6.0 * 448.0 assert converted.v_bmm_quantizer._amax == 6.0 * 448.0 assert converted._query_quant_in_kernel is True diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 7730e5a80de..2b427adedf2 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -292,6 +292,7 @@ def __float__(self): q_inputs = [] def quantize_q(query): + assert query.dtype == torch.float32 q_inputs.append(query.clone()) return query + 1 @@ -357,8 +358,9 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): assert decode_kw["v_qdq"] == "nvfp4" assert decode_kw["v_cache_quantized"] is True assert len(q_inputs) == 1 and q_inputs[0].shape[0] == q.shape[0] - torch.testing.assert_close(q_inputs[0][:2], q[:2]) + torch.testing.assert_close(q_inputs[0][:2], q[:2].float()) assert torch.all(q_inputs[0][2:] == 0) + assert calls["query"].dtype == torch.float32 def test_quantized_skip_softmax_decode_stays_on_shared_kernel(monkeypatch): From 46744e1bafdd2f45fd229ebf1b3bcd1d8d530fff Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 3 Jul 2026 21:58:43 -0700 Subject: [PATCH 08/28] Fix NVFP4 attention tests Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 9 +- .../vllm_serve/quant_sparse_attn_worker.py | 59 ++++++++----- .../common/attention/decode_attention.py | 9 +- .../kernels/quantization/attention/v_qdq.py | 17 ++-- .../attention_sparsity/plugins/vllm.py | 1 + .../common/attention/test_decode_attention.py | 87 +++++++++++++++++++ .../common/attention/test_triton_fa_paged.py | 20 ++++- .../test_quant_sparse_attn_worker.py | 21 ++++- .../test_sparse_attn_worker.py | 42 ++++++++- 9 files changed, 224 insertions(+), 41 deletions(-) diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 3a2f1c6ccd5..b8316bab15c 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -4,7 +4,8 @@ This is a simple example to demonstrate calibrating and serving ModelOpt fakequa Compared with realquant, fakequant is 2-5x slower, but doesn't require dedicated kernel support and facilitates research. -This example is tested with vllm 0.9.0 and 0.19.1 +The general fakequant example is tested with vLLM 0.9.0 and 0.19.1. The compact +NVFP4 attention worker documented below requires vLLM 0.14.0 or newer. ## Prepare environment @@ -123,6 +124,8 @@ Limitations: ### Compact NVFP4 attention worker +This worker requires vLLM 0.14.0 or newer and fails at import time on older releases. + Use the same launcher with the compact worker and FlashAttention selected explicitly: ```bash @@ -134,6 +137,10 @@ python vllm_serve_sparse_attn.py -tp 8 \ This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant recipe to Q/K/P/V. Q is dynamic, K/V use global scale 1.0, and P uses amax 1.0; calibrated attention amax restore is not part of this fixed path. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored. +Decode uses a fixed 32-split, 128-key-tile schedule. P QDQ consumes split-local, +unnormalized online-softmax probabilities, so changing that schedule can change +quantized results; split count is part of the numerical contract. + K is QDQ before its cache write, while V is written pristine. Complete 16-token V groups are finalized once in cache; an incomplete tail remains pristine and is QDQ on read. P@V therefore sees uniform fakequant values without re-quantizing the tail. Supported configurations are regular decoder self-attention with FlashAttention, fp16/bf16 model and KV cache, equal Q/K/V head dimensions that are multiples of 16, and DCP 1. The default `FULL_AND_PIECEWISE` CUDA graph mode and full decode graphs are supported. diff --git a/examples/vllm_serve/quant_sparse_attn_worker.py b/examples/vllm_serve/quant_sparse_attn_worker.py index fe30593bed3..ba0f7087a2f 100644 --- a/examples/vllm_serve/quant_sparse_attn_worker.py +++ b/examples/vllm_serve/quant_sparse_attn_worker.py @@ -16,30 +16,41 @@ """Fixed NVFP4 Q/K/P/V worker for ModelOpt sparse attention on vLLM.""" import torch -from vllm.config.compilation import CUDAGraphMode -from vllm.v1.attention.backend import AttentionType -from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl -from vllm.v1.worker.gpu_worker import Worker as BaseWorker - -from modelopt.torch.quantization.conversion import set_quantizer_by_cfg -from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer -from modelopt.torch.quantization.plugins.vllm import ( - _ATTENTION_TYPES, - _get_device_dtype, - _set_vllm_attention_kv_default_amax, - disable_compilation, - vllm_attention, -) -from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_config import ( - load_from_checkpoint_metadata, - match_sparse_config, -) -from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( - _build_sparse_kw, - _clone_sparse_impl, - _p_qdq_from_layer, - _v_qdq_from_layer, -) +import vllm +from packaging import version + +try: + from vllm.config.compilation import CUDAGraphMode + from vllm.v1.attention.backend import AttentionType + from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl + from vllm.v1.worker.gpu_worker import Worker as BaseWorker + + from modelopt.torch.quantization.conversion import set_quantizer_by_cfg + from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer + from modelopt.torch.quantization.plugins.vllm import ( + _ATTENTION_TYPES, + _get_device_dtype, + _set_vllm_attention_kv_default_amax, + disable_compilation, + vllm_attention, + ) + from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_config import ( + load_from_checkpoint_metadata, + match_sparse_config, + ) + from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( + _build_sparse_kw, + _clone_sparse_impl, + _p_qdq_from_layer, + _v_qdq_from_layer, + ) +except ImportError as err: + if version.parse(vllm.__version__) < version.parse("0.14.0"): + raise RuntimeError("The compact NVFP4 attention worker requires vLLM >= 0.14.0") from err + raise + +if version.parse(vllm.__version__) < version.parse("0.14.0"): + raise RuntimeError("The compact NVFP4 attention worker requires vLLM >= 0.14.0") __all__ = ["QuantSparseAttnWorker"] diff --git a/modelopt/torch/kernels/common/attention/decode_attention.py b/modelopt/torch/kernels/common/attention/decode_attention.py index 377b6450161..d12f66f16eb 100644 --- a/modelopt/torch/kernels/common/attention/decode_attention.py +++ b/modelopt/torch/kernels/common/attention/decode_attention.py @@ -14,7 +14,11 @@ # limitations under the License. -"""Split-K decode attention for the ModelOpt paged NVFP4 serving path.""" +"""Split-K decode attention for the ModelOpt paged NVFP4 serving path. + +P QDQ operates on split-local, unnormalized online-softmax probabilities. Its +numerics therefore include the fixed split count as part of the kernel schedule. +""" import math @@ -165,6 +169,7 @@ def _decode_split_kernel( ).to(tl.float32) if V_QDQ and ((not V_CACHE_QUANTIZED) or kv_start + BLOCK_N > v_quantized_boundary): v_qdq = _v_qdq_nvfp4(v, v_qdq_scale, BLOCK_N, BLOCK_D) + v_qdq = v_qdq.to(V_cache.dtype.element_ty).to(tl.float32) if V_CACHE_QUANTIZED: use_qdq = kv_start + kv_pos >= v_quantized_boundary v = tl.where(use_qdq[:, None], v_qdq, v) @@ -268,6 +273,8 @@ def attention_decode( one. P is rounded to the model/cache dtype before native-style quantization, then its QDQ result remains FP32. Complete block-16 V groups may be finalized in the cache; only the pristine partial group is then quantized on read. + P QDQ intentionally follows the split-local online-softmax schedule; changing + ``num_kv_splits`` can therefore change quantized results. """ if q.ndim != 3: raise ValueError(f"q must have shape [batch, heads, head_dim], got {tuple(q.shape)}") diff --git a/modelopt/torch/kernels/quantization/attention/v_qdq.py b/modelopt/torch/kernels/quantization/attention/v_qdq.py index 876a189787e..edbd0e0f2ec 100644 --- a/modelopt/torch/kernels/quantization/attention/v_qdq.py +++ b/modelopt/torch/kernels/quantization/attention/v_qdq.py @@ -83,14 +83,16 @@ def fake_quant_v_onwrite( v_lo: torch.Tensor, v_hi: torch.Tensor, *, + max_new_tokens: int, page_size: int = 16, v_qdq_scale: float = 1.0, decode: bool = False, ) -> None: """NVFP4-finalize complete block-16 groups in ``[v_lo, v_hi)`` in place. - Decode uses one fixed group per request without reading device metadata on the - host. Eager prefill sizes the grid to cover every newly completed group. + ``max_new_tokens`` is host metadata used to size the masked launch grid. + Decode uses one fixed group per request. Eager prefill covers every group + that the largest query chunk can complete without reading device metadata. ``v_lo`` and ``v_hi`` must describe aligned, completed block-16 boundaries; their device values are not host-validated. """ @@ -98,17 +100,12 @@ def fake_quant_v_onwrite( raise ValueError(f"page_size {page_size} must match v_cache.shape[1] {v_cache.shape[1]}") if not (math.isfinite(v_qdq_scale) and v_qdq_scale > 0): raise ValueError(f"v_qdq_scale must be finite and positive, got {v_qdq_scale}") + if max_new_tokens < 1: + raise ValueError(f"max_new_tokens must be positive, got {max_new_tokens}") batch, max_blocks = block_table.shape num_kv_heads, head_dim = v_cache.shape[2:] - if decode: - num_groups = 1 - else: - first_group = v_lo // _BLOCK_N - span = int((v_hi - first_group * _BLOCK_N).max().item()) - if span <= 0: - return - num_groups = triton.cdiv(span, _BLOCK_N) + num_groups = 1 if decode else triton.cdiv(max_new_tokens, _BLOCK_N) with torch.cuda.device(v_cache.device): _fake_quant_v_onwrite_kernel[(batch, num_kv_heads, num_groups)]( diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index bc5d0792a78..f778fbac000 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -281,6 +281,7 @@ def forward( attn_metadata.block_table, (prev // 16) * 16, (seq_lens // 16) * 16, + max_new_tokens=attn_metadata.max_query_len, page_size=page_size, v_qdq_scale=v_qdq_scale, decode=is_decode_only, diff --git a/tests/gpu/torch/kernels/common/attention/test_decode_attention.py b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py index cee3e03b16f..e50ba6ca360 100644 --- a/tests/gpu/torch/kernels/common/attention/test_decode_attention.py +++ b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py @@ -131,6 +131,7 @@ def test_baked_v_prefix_and_pristine_tail_match_full_onread(): block_table, torch.zeros(1, device="cuda", dtype=torch.int32), torch.tensor([16], device="cuda", dtype=torch.int32), + max_new_tokens=16, ) baked_v_before = baked_v_cache.clone() baked = attention_decode( @@ -147,6 +148,45 @@ def test_baked_v_prefix_and_pristine_tail_match_full_onread(): torch.testing.assert_close(baked_v_cache, baked_v_before, rtol=0, atol=0) +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") +@requires_native_e4m3 +def test_baked_v_prefix_uses_same_cache_carrier_as_tail(): + """A non-default V scale must not change values at the bake boundary.""" + seq_len, head_dim = 17, 16 + q = torch.zeros(1, 1, head_dim, device="cuda", dtype=torch.float32) + k = torch.zeros(1, 1, seq_len, head_dim, device="cuda", dtype=torch.bfloat16) + v = torch.full_like(k, 0.019) + seq_lens = torch.tensor([seq_len], device="cuda", dtype=torch.int32) + k_cache, raw_v_cache, block_table = _paged_cache(k, v, (seq_len,)) + common = { + "num_kv_splits": 1, + "v_qdq": "nvfp4", + "v_qdq_amax": 1.0, + } + + onread = attention_decode(q, k_cache, raw_v_cache, block_table, seq_lens, **common) + baked_v_cache = raw_v_cache.clone() + fake_quant_v_onwrite( + baked_v_cache, + block_table, + torch.zeros(1, device="cuda", dtype=torch.int32), + torch.tensor([16], device="cuda", dtype=torch.int32), + max_new_tokens=16, + v_qdq_scale=1.0 / (6.0 * 448.0), + ) + baked = attention_decode( + q, + k_cache, + baked_v_cache, + block_table, + seq_lens, + v_cache_quantized=True, + **common, + ) + + torch.testing.assert_close(baked, onread, rtol=0, atol=0) + + @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") @requires_native_e4m3 def test_p_quantizes_model_dtype_input_but_accumulates_fp32(): @@ -178,3 +218,50 @@ def test_p_quantizes_model_dtype_input_but_accumulates_fp32(): reference = (p_qdq[:, None] * v[0, 0].float()).sum(0) / p.sum() torch.testing.assert_close(output[0, 0], reference, rtol=5e-5, atol=5e-5) + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") +@requires_native_e4m3 +def test_p_qdq_matches_fixed_split_local_oracle(): + """The production 32-split schedule quantizes split-local unnormalized P.""" + seq_len, head_dim, num_splits = 4096, 16, 32 + scale = head_dim**-0.5 + q = torch.zeros(1, 1, head_dim, device="cuda", dtype=torch.float32) + q[..., 0] = 1.0 + k = torch.zeros(1, 1, seq_len, head_dim, device="cuda", dtype=torch.bfloat16) + k[..., 0] = torch.linspace(-8.0, 8.0, seq_len, device="cuda", dtype=torch.bfloat16) + torch.manual_seed(19) + v = torch.randn_like(k) + k_cache, v_cache, block_table = _paged_cache(k, v, (seq_len,)) + seq_lens = torch.tensor([seq_len], device="cuda", dtype=torch.int32) + + output = attention_decode( + q, + k_cache, + v_cache, + block_table, + seq_lens, + softmax_scale=scale, + num_kv_splits=num_splits, + p_qdq="nvfp4", + ) + + scores = torch.matmul(k[0, 0].float(), q[0, 0]) * (scale * 1.4426950408889634) + split_scores = scores.reshape(num_splits, -1) + split_max = split_scores.amax(dim=1) + p = torch.exp2(split_scores - split_max[:, None]) + p_qdq = _nvfp4_qdq_reference(p.to(torch.bfloat16).float()) + split_acc = torch.einsum("sk,skd->sd", p_qdq, v[0, 0].float().reshape(num_splits, -1, head_dim)) + running_max = torch.tensor(-float("inf"), device="cuda") + running_sum = torch.tensor(0.0, device="cuda") + acc = torch.zeros(head_dim, device="cuda") + for split_idx in range(num_splits): + new_max = torch.maximum(running_max, split_max[split_idx]) + correction = torch.exp2(running_max - new_max) + split_correction = torch.exp2(split_max[split_idx] - new_max) + acc = acc * correction + split_acc[split_idx] * split_correction + running_sum = running_sum * correction + p[split_idx].sum() * split_correction + running_max = new_max + reference = acc / running_sum + + torch.testing.assert_close(output[0, 0], reference, rtol=5e-5, atol=5e-5) diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py index 202137c4c07..ca0deb8b7fd 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py @@ -35,6 +35,15 @@ ) +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") +def test_v_onwrite_prefill_grid_uses_host_metadata(): + signature = inspect.signature(fake_quant_v_onwrite) + source = inspect.getsource(fake_quant_v_onwrite) + + assert "max_new_tokens" in signature.parameters + assert ".item()" not in source + + @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need triton") @pytest.mark.parametrize("loader", ["_load_paged_k_tile", "_load_paged_v_tile"]) def test_paged_loaders_widen_page_ids_before_pointer_math(loader): @@ -143,6 +152,7 @@ def test_v_onwrite_validates_page_size_and_scale(self, page_size, v_qdq_scale, m zeros.view(1, 1), zeros, zeros, + max_new_tokens=1, page_size=page_size, v_qdq_scale=v_qdq_scale, ) @@ -544,6 +554,7 @@ def test_v_cache_finalizes_complete_groups_once(self): block_table, torch.zeros(1, device="cuda", dtype=torch.int32), torch.tensor([32], device="cuda", dtype=torch.int32), + max_new_tokens=32, page_size=page_size, ) after_prefill = baked.clone() @@ -552,6 +563,7 @@ def test_v_cache_finalizes_complete_groups_once(self): block_table, torch.tensor([32], device="cuda", dtype=torch.int32), torch.tensor([48], device="cuda", dtype=torch.int32), + max_new_tokens=1, page_size=page_size, decode=True, ) @@ -565,6 +577,7 @@ def test_v_cache_finalizes_complete_groups_once(self): torch.zeros(1, 1, device="cuda", dtype=torch.int32), torch.zeros(1, device="cuda", dtype=torch.int32), torch.tensor([16], device="cuda", dtype=torch.int32), + max_new_tokens=16, v_qdq_scale=1.0 / (6.0 * 448.0), ) # Native E4M3 block scales below half the minimum subnormal round to zero. @@ -594,6 +607,7 @@ def test_v_cache_matches_independent_signed_key_axis_oracle(self): block_table, torch.zeros(1, device="cuda", dtype=torch.int32), torch.tensor([16], device="cuda", dtype=torch.int32), + max_new_tokens=16, page_size=page_size, ) key_last = logical.permute(1, 2, 0).contiguous() @@ -639,7 +653,11 @@ def test_baked_prefix_and_raw_tail_match_full_onread(self, q_len): out_onread = attention(q, k, v, q_locs, q_lens, q_len, v_cache=raw, **common) baked = raw.clone() fake_quant_v_onwrite( - baked, block_table, locs, torch.tensor([16], device="cuda", dtype=torch.int32) + baked, + block_table, + locs, + torch.tensor([16], device="cuda", dtype=torch.int32), + max_new_tokens=16, ) out_baked = attention( q, diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py index 7783ad2fae8..344cb472d1b 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py @@ -22,6 +22,7 @@ import pytest import torch +import vllm from torch import nn from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.backend import AttentionType @@ -32,9 +33,23 @@ from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ModelOptSparseAttentionImpl _WORKER_PATH = Path(__file__).parents[5] / "examples/vllm_serve/quant_sparse_attn_worker.py" -_SPEC = importlib.util.spec_from_file_location("quant_sparse_attn_worker", _WORKER_PATH) -worker_module = importlib.util.module_from_spec(_SPEC) -_SPEC.loader.exec_module(worker_module) + + +def _load_worker_module(): + spec = importlib.util.spec_from_file_location("quant_sparse_attn_worker", _WORKER_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +worker_module = _load_worker_module() + + +def test_worker_rejects_vllm_before_v1_attention_api(monkeypatch): + monkeypatch.setattr(vllm, "__version__", "0.9.0") + + with pytest.raises(RuntimeError, match=r"vLLM >= 0\.14\.0"): + _load_worker_module() def _attention(attn_type=AttentionType.DECODER): diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 2b427adedf2..ebcda487e34 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -347,7 +347,12 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): v_lo, v_hi, finalizer_kw = calls["finalize"] torch.testing.assert_close(v_lo, torch.tensor([0, 32], dtype=torch.int32)) torch.testing.assert_close(v_hi, torch.tensor([16, 32], dtype=torch.int32)) - assert finalizer_kw == {"page_size": 16, "v_qdq_scale": 1.0, "decode": True} + assert finalizer_kw == { + "max_new_tokens": 1, + "page_size": 16, + "v_qdq_scale": 1.0, + "decode": True, + } key_cache, value_cache, block_table, seq_lens, decode_kw = calls["decode"] assert key_cache.data_ptr() == kv_cache[0].data_ptr() assert value_cache.data_ptr() == kv_cache[1].data_ptr() @@ -363,6 +368,41 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): assert calls["query"].dtype == torch.float32 +def test_quantized_prefill_passes_host_v_group_bound(monkeypatch): + impl = _clone_sparse_impl(_make_old_impl()) + impl.quant_kw = { + "p_qdq": None, + "p_qdq_amax": 1.0, + "v_qdq": "nvfp4", + "v_qdq_amax": 6.0 * 448.0, + } + q_len, kv_len = 17, 31 + q = torch.zeros(q_len, impl.num_heads, impl.head_size, dtype=torch.float16) + kv_cache = torch.zeros(2, 2, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + metadata = SimpleNamespace( + num_actual_tokens=q_len, + max_query_len=q_len, + max_seq_len=kv_len, + query_start_loc=torch.tensor([0, q_len], dtype=torch.int32), + seq_lens=torch.tensor([kv_len], dtype=torch.int32), + block_table=torch.zeros(1, 2, dtype=torch.int32), + ) + calls = {} + + def fake_finalize(*args, **kwargs): + calls["finalize"] = kwargs + + monkeypatch.setattr(vllm_plugin, "fake_quant_v_onwrite", fake_finalize) + monkeypatch.setattr( + vllm_plugin, "triton_attention", lambda query, **kwargs: torch.zeros_like(query) + ) + + impl.forward(None, q, q, q, kv_cache, metadata, output=torch.empty_like(q)) + + assert calls["finalize"]["max_new_tokens"] == q_len + assert calls["finalize"]["decode"] is False + + def test_quantized_skip_softmax_decode_stays_on_shared_kernel(monkeypatch): """Split-local maxima must not change calibrated skip-softmax semantics.""" impl = _clone_sparse_impl(_make_old_impl()) From b57f8b8684583471e7ef0b5e44dde227353917c8 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 4 Jul 2026 15:08:31 -0700 Subject: [PATCH 09/28] Support FlashInfer vLLM attention backend Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 11 +- .../vllm_serve/quant_sparse_attn_worker.py | 70 ++- examples/vllm_serve/sparse_attn_worker.py | 46 +- .../kernels/common/attention/triton_fa.py | 5 +- .../attention_sparsity/plugins/vllm.py | 586 +++++++++++++----- .../common/attention/test_triton_fa_paged.py | 12 - .../quantization/test_vllm_dynamic_modules.py | 17 - .../test_quant_sparse_attn_worker.py | 54 +- .../test_sparse_attn_worker.py | 579 ++++++++++++----- .../common/attention/test_triton_fa.py | 42 +- 10 files changed, 1059 insertions(+), 363 deletions(-) diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index b8316bab15c..3a6cefc59f5 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -102,9 +102,9 @@ QUANT_CFG= QUANT_FILE_PATH= python vllm_serve_fa ## Serve a model with sparse attention in vLLM -Apply ModelOpt sparse attention at serve time. The launcher replaces vLLM's `FlashAttentionImpl` with `ModelOptSparseAttentionImpl` (Triton kernel with paged KV cache support) on every attention layer right after model load. +Apply ModelOpt sparse attention at serve time. Right after model load, the launcher replaces each native attention implementation with its matching ModelOpt adapter: `ModelOptSparseAttentionImpl` for FlashAttention or `ModelOptSparseFlashInferImpl` for FlashInfer. Both adapters use the same Triton kernel with paged KV cache support. -The configuration is read from the checkpoint's `config.json` `sparse_attention_config` block, written by ModelOpt's HF export. The launcher restores calibrated skip-softmax metadata and N:M sparse-softmax metadata (`sparsity_n`, `sparsity_m`, `dense_sink_tokens`, `dense_recent_tokens`). Checkpoints exported with both metadata entries use ModelOpt Triton for sparse prefill launches; decode-only launches and launches without active sparse work delegate back to vLLM FlashAttention. +The configuration is read from the checkpoint's `config.json` `sparse_attention_config` block, written by ModelOpt's HF export. The launcher restores calibrated skip-softmax metadata and N:M sparse-softmax metadata (`sparsity_n`, `sparsity_m`, `dense_sink_tokens`, `dense_recent_tokens`). Checkpoints exported with both metadata entries use ModelOpt Triton for sparse prefill launches; launches without active sparse work delegate back to the native backend selected by vLLM. Workflow: @@ -126,15 +126,16 @@ Limitations: This worker requires vLLM 0.14.0 or newer and fails at import time on older releases. -Use the same launcher with the compact worker and FlashAttention selected explicitly: +Use the same launcher with the compact worker. By default, vLLM selects the backend for the model and platform; NemotronH on Blackwell selects FlashInfer: ```bash python vllm_serve_sparse_attn.py -tp 8 \ - --attention-backend FLASH_ATTN \ --no-enable-prefix-caching \ --worker-cls quant_sparse_attn_worker.QuantSparseAttnWorker ``` +The worker supports both FlashInfer and FlashAttention and prints the installed adapter counts. Pass `--attention-backend FLASHINFER` or `--attention-backend FLASH_ATTN` only when an explicit override is needed. + This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant recipe to Q/K/P/V. Q is dynamic, K/V use global scale 1.0, and P uses amax 1.0; calibrated attention amax restore is not part of this fixed path. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored. Decode uses a fixed 32-split, 128-key-tile schedule. P QDQ consumes split-local, @@ -143,7 +144,7 @@ quantized results; split count is part of the numerical contract. K is QDQ before its cache write, while V is written pristine. Complete 16-token V groups are finalized once in cache; an incomplete tail remains pristine and is QDQ on read. P@V therefore sees uniform fakequant values without re-quantizing the tail. -Supported configurations are regular decoder self-attention with FlashAttention, fp16/bf16 model and KV cache, equal Q/K/V head dimensions that are multiples of 16, and DCP 1. The default `FULL_AND_PIECEWISE` CUDA graph mode and full decode graphs are supported. +Supported configurations are regular decoder self-attention with FlashInfer or FlashAttention, fp16/bf16 model and KV cache, equal Q/K/V head dimensions that are multiples of 16, and DCP 1. The FlashInfer adapter preserves both NHD and HND cache strides and separates mixed decode/prefill launches so each phase keeps its own kernel contract. The default `FULL_AND_PIECEWISE` mode remains enabled for fixed N:M and attention-only NVFP4; checkpoints with calibrated decode `threshold_scale_factor` must use a non-`FULL` decode graph mode such as `--enforce-eager` because the live sequence length is not replayed as a Python scalar. Unsupported features are sliding window, ALiBi, softcap, sinks, FP8 KV cache, cross/encoder/MLA attention, KV sharing or transfer, prefix caching, speculative decoding, DBO/ubatching, and `FULL` mixed/prefill CUDA graphs. diff --git a/examples/vllm_serve/quant_sparse_attn_worker.py b/examples/vllm_serve/quant_sparse_attn_worker.py index ba0f7087a2f..cd9f9b4d22c 100644 --- a/examples/vllm_serve/quant_sparse_attn_worker.py +++ b/examples/vllm_serve/quant_sparse_attn_worker.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Fixed NVFP4 Q/K/P/V worker for ModelOpt sparse attention on vLLM.""" +"""Fixed NVFP4 Q/K/P/V worker for FlashAttention and FlashInfer in vLLM.""" import torch import vllm @@ -22,7 +22,6 @@ try: from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.backend import AttentionType - from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl from vllm.v1.worker.gpu_worker import Worker as BaseWorker from modelopt.torch.quantization.conversion import set_quantizer_by_cfg @@ -43,6 +42,7 @@ _clone_sparse_impl, _p_qdq_from_layer, _v_qdq_from_layer, + select_sparse_impl_cls, ) except ImportError as err: if version.parse(vllm.__version__) < version.parse("0.14.0"): @@ -106,8 +106,6 @@ def _layer_errors(module) -> list[str]: errors.append(f"layout {type(module).__name__} is not regular decoder self-attention") if getattr(module, "attn_type", None) != AttentionType.DECODER: errors.append("attn_type must be DECODER") - if not isinstance(impl, FlashAttentionImpl): - errors.append(f"backend {type(impl).__name__} is not FlashAttentionImpl") head_size = getattr(module, "head_size", None) if not isinstance(head_size, int) or head_size % 16: errors.append(f"head_size={head_size!r} must be a multiple of 16") @@ -129,12 +127,29 @@ def _layer_errors(module) -> list[str]: return errors +def _sparse_graph_error(sparse_kw: dict, mode: CUDAGraphMode) -> str | None: + """Reject decode calibration whose live length would be frozen by a full graph.""" + params = sparse_kw.get("threshold_scale_factor") + if ( + mode.decode_mode() == CUDAGraphMode.FULL + and isinstance(params, dict) + and isinstance(params.get("decode"), dict) + ): + return "calibrated decode skip-softmax requires a non-FULL CUDA graph mode" + return None + + def _validated_attention_plans(worker): """Validate every attention layout without mutation, then return regular-layer tuples.""" model = _unwrapped_model(worker) - detected = load_from_checkpoint_metadata(worker.model_runner.model_config.hf_config) + model_config = worker.model_runner.model_config + model_dtype = getattr(model_config, "dtype", None) + detected = load_from_checkpoint_metadata(model_config.hf_config) sparse_cfg = detected[0] if detected is not None else None errors = _global_errors(worker) + vllm_config = getattr(worker.model_runner, "vllm_config", None) + compilation_config = getattr(vllm_config, "compilation_config", None) + cudagraph_mode = getattr(compilation_config, "cudagraph_mode", CUDAGraphMode.NONE) plans = [] attention_count = 0 for name, module in model.named_modules(): @@ -142,7 +157,19 @@ def _validated_attention_plans(worker): continue attention_count += 1 reasons = _layer_errors(module) + new_impl_cls = select_sparse_impl_cls(module.impl) + if new_impl_cls is None: + reasons.append( + f"backend {type(module.impl).__name__} is not supported; " + "expected FlashAttentionImpl or FlashInferImpl" + ) device, dtype = _get_device_dtype(module) + # The install runs in load_model (before KV-cache allocation), so the + # buffer scan in _get_device_dtype can pick up fp32 scale buffers + # (_k_scale/_v_scale/...) on the attention module. Prefer the + # authoritative model compute dtype when it is fp16/bf16. + if model_dtype in (torch.float16, torch.bfloat16): + dtype = model_dtype if device is None or dtype is None: reasons.append("device/dtype could not be resolved") elif dtype not in (torch.float16, torch.bfloat16): @@ -156,7 +183,11 @@ def _validated_attention_plans(worker): if layer_cfg is not None and layer_cfg.get("enable", True) else {} ) - plans.append((name, module, sparse_kw, device, dtype)) + graph_error = _sparse_graph_error(sparse_kw, cudagraph_mode) + if graph_error is not None: + errors.append(f"{name or ''}: {graph_error}") + continue + plans.append((name, module, new_impl_cls, sparse_kw, device, dtype)) if attention_count == 0: errors.append("no regular attention layers were found") if errors: @@ -168,13 +199,14 @@ def _validated_attention_plans(worker): def _install_quant_sparse_attn(worker) -> None: plans = _validated_attention_plans(worker) - for _name, module, sparse_kw, device, dtype in plans: + installed = {} + for _name, module, new_impl_cls, sparse_kw, device, dtype in plans: module.device, module.dtype = device, dtype QuantModuleRegistry.convert(module) module.p_bmm_quantizer = TensorQuantizer() set_quantizer_by_cfg(module, _BMM_CFG) _set_vllm_attention_kv_default_amax(module, device) - new_impl = _clone_sparse_impl(module.impl) + new_impl = _clone_sparse_impl(module.impl, new_impl_cls) new_impl.sparse_kw = sparse_kw p_qdq, p_qdq_amax = _p_qdq_from_layer(module) v_qdq, v_qdq_amax = _v_qdq_from_layer(module) @@ -187,18 +219,28 @@ def _install_quant_sparse_attn(worker) -> None: module.impl = new_impl module._query_quant_in_kernel = True module._value_quant_in_kernel = True + impl_name = type(new_impl).__name__ + installed[impl_name] = installed.get(impl_name, 0) + 1 worker.model_runner.cascade_attn_enabled = False - print(f"[ModelOpt] Installed NVFP4 quant+sparse attention on {len(plans)} layers") + print(f"[ModelOpt] Installed NVFP4 quant+sparse attention on {len(plans)} layers: {installed}") class QuantSparseAttnWorker(BaseWorker): - """Install the fixed attention-only recipe before vLLM warmup and graph capture.""" + """Install the fixed NVFP4 attention recipe right after model load. + + The impl swap must happen before the first forward pass -- including the + ``determine_available_memory`` profiling run -- so that profiling exercises + the ModelOpt Triton path that real serving uses. Installing here also patches + FlashInfer metadata before cache and metadata-builder initialization. This + mirrors ``SparseAttnWorker``, which installs in ``load_model`` for the same + reason. + """ + + def load_model(self, *args, **kwargs) -> None: + super().load_model(*args, **kwargs) + _install_quant_sparse_attn(self) @torch.inference_mode() def determine_available_memory(self) -> int: with disable_compilation(_unwrapped_model(self)): return BaseWorker.determine_available_memory(self) - - def compile_or_warm_up_model(self) -> float: - _install_quant_sparse_attn(self) - return BaseWorker.compile_or_warm_up_model(self) diff --git a/examples/vllm_serve/sparse_attn_worker.py b/examples/vllm_serve/sparse_attn_worker.py index 1cd27f4756d..4b8c5e84579 100644 --- a/examples/vllm_serve/sparse_attn_worker.py +++ b/examples/vllm_serve/sparse_attn_worker.py @@ -15,11 +15,10 @@ """Custom vLLM worker for sparse attention. -``SparseAttnWorker``: Replaces ``FlashAttentionImpl`` with -``ModelOptSparseAttentionImpl`` on each Attention module after model loading. -The sparse impl uses the ModelOpt Triton kernel for sparse prefill launches. -Decode-only launches and launches without active sparse work delegate back to -vLLM FlashAttention. +``SparseAttnWorker`` replaces the native FlashAttention or FlashInfer impl with +the matching ModelOpt adapter on each Attention module after model loading. +The sparse impl uses the ModelOpt Triton kernel for sparse prefill launches and +delegates inactive launches to the selected native backend. Configuration flows exclusively through the loaded checkpoint's ``sparse_attention_config`` block (written by ModelOpt's HF export). If the @@ -54,11 +53,12 @@ from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( _build_sparse_kw, _clone_sparse_impl, + select_sparse_impl_cls, ) def _replace_attention_impl(worker): - """Replace FlashAttentionImpl with ModelOptSparseAttentionImpl on all Attention layers. + """Install the backend-matched ModelOpt sparse impl on attention layers. The sole configuration source is the checkpoint's ``sparse_attention_config`` metadata. No-op if the checkpoint has no such block. @@ -80,7 +80,8 @@ def _replace_attention_impl(worker): if hasattr(model, "unwrap"): model = model.unwrap() - patched = 0 + plans = [] + errors = [] for name, module in model.named_modules(): if not isinstance(module, VLLMAttention): continue @@ -94,11 +95,31 @@ def _replace_attention_impl(worker): # Keep vLLM's original impl when the exported layer config does not # enable any sparse feature. continue - new_impl = _clone_sparse_impl(module.impl) + new_impl_cls = select_sparse_impl_cls(module.impl) + if new_impl_cls is None: + errors.append(f"{name or ''}: unsupported backend {type(module.impl).__name__}") + continue + try: + new_impl = _clone_sparse_impl(module.impl, new_impl_cls) + except (NotImplementedError, TypeError) as err: + errors.append(f"{name or ''}: {err}") + continue + plans.append((module, new_impl, sparse_kw)) + + if errors: + raise NotImplementedError( + "Unsupported ModelOpt sparse attention plan:\n - " + "\n - ".join(errors) + ) + + installed = {} + for module, new_impl, sparse_kw in plans: new_impl.sparse_kw = sparse_kw module.impl = new_impl - patched += 1 - print(f"[ModelOpt] Sparse attention: replaced impl on {patched} attention layers") + impl_name = type(new_impl).__name__ + installed[impl_name] = installed.get(impl_name, 0) + 1 + print( + f"[ModelOpt] Sparse attention: replaced impl on {len(plans)} attention layers: {installed}" + ) # --------------------------------------------------------------------------- @@ -109,9 +130,8 @@ def _replace_attention_impl(worker): class SparseAttnWorker(BaseWorker): """vLLM worker that uses the ModelOpt sparse attention backend. - Replaces FlashAttentionImpl with ModelOptSparseAttentionImpl on each - Attention module right after model loading — before any forward pass - (including determine_available_memory profiling). + Replaces FlashAttention or FlashInfer with its matching ModelOpt adapter on + each Attention module right after model loading, before any forward pass. """ def load_model(self, *args, **kwargs) -> None: diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 8755757cf55..93d0faa738a 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -956,7 +956,10 @@ def forward( lse.stride(1), ) fwd_kwargs = { - "N_CTX": max_input_len, + # N_CTX is an autotune key only. Bucket variable prefill lengths so + # each power-of-two regime reuses one tuned configuration; the grid + # below still uses the exact max_input_len. + "N_CTX": triton.next_power_of_2(max(1, max_input_len)), "kv_group_num": kv_group_num, "BLOCK_D": BLOCK_D, "IS_CAUSAL": is_causal, diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index f778fbac000..6ff552ff7aa 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -15,18 +15,20 @@ """ModelOpt sparse attention backend for vLLM. -Registers a custom vLLM attention backend that uses the ModelOpt Triton kernel -with paged KV cache support. Integration approach: +Installs backend-matched vLLM attention implementations that use the ModelOpt +Triton kernel with paged KV cache support. Integration approach: - No module replacement — the Attention module stays intact with all its state -- Only ``impl`` is swapped from FlashAttentionImpl to ModelOptSparseAttentionImpl -- KV cache update is handled by vLLM (inherited ``do_kv_cache_update``) -- ``forward()`` calls ModelOpt Triton only when a validated sparse path is active +- Only ``impl`` is swapped to the matching FlashAttention or FlashInfer adapter +- KV cache update follows the selected backend's native version-specific contract +- ``forward()`` calls ModelOpt Triton only when a validated transform is active Vllm-free config helpers (``match_sparse_config`` / ``load_from_checkpoint_metadata``) live in ``plugins/sparse_attn_config.py`` and are unit-testable without vLLM. """ +import functools +import inspect import math import warnings @@ -145,14 +147,159 @@ def _v_qdq_from_layer(layer) -> tuple[str | None, float | None]: return _bmm_qdq_from_layer(layer, "v_bmm_quantizer", None) -class ModelOptSparseAttentionImpl(FlashAttentionImpl): - """Attention implementation that uses the ModelOpt Triton kernel. +def _quant_kw_from_impl(impl, layer): + """Resolve the compact P/V QDQ contract once for one attention launch.""" + quant_kw = getattr(impl, "quant_kw", None) + if quant_kw is None: + p_qdq, p_qdq_amax = _p_qdq_from_layer(layer) + v_qdq, v_qdq_amax = _v_qdq_from_layer(layer) + else: + p_qdq, p_qdq_amax = quant_kw["p_qdq"], quant_kw["p_qdq_amax"] + v_qdq, v_qdq_amax = quant_kw["v_qdq"], quant_kw["v_qdq_amax"] + return p_qdq, p_qdq_amax, v_qdq, v_qdq_amax + + +def _has_active_quant_transform(layer, p_qdq, v_qdq) -> bool: + """Return whether native fallback would omit any Q/K/P/V transform.""" + k_quantizer = getattr(layer, "k_bmm_quantizer", None) + return bool( + p_qdq + or v_qdq + or getattr(layer, "_query_quant_in_kernel", False) + or getattr(k_quantizer, "is_enabled", False) + ) + + +def _has_active_transform(impl, layer, p_qdq, v_qdq) -> bool: + """Return whether native fallback would omit any ModelOpt transform.""" + return bool( + getattr(impl, "sparse_kw", None) or _has_active_quant_transform(layer, p_qdq, v_qdq) + ) - Inherits from FlashAttentionImpl to reuse: - - __init__ (all configuration) - - do_kv_cache_update (KV cache writing) - Only overrides forward() to replace sparse prefill attention computation. - """ + +def _forward_modelopt( + impl, + *, + layer, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + num_actual_tokens: int, + max_query_len: int, + max_seq_len: int, + is_causal: bool, + output: torch.Tensor, + p_qdq: str | None, + p_qdq_amax: float, + v_qdq: str | None, + v_qdq_amax: float | None, + quant_transform_active: bool, + dense_fallback, + prepare_modelopt=None, +) -> torch.Tensor: + """Run the compact ModelOpt path over a backend-normalized paged cache.""" + batch = seq_lens.shape[0] + b_start_loc = cu_seqlens_q[:batch] + b_seq_len = cu_seqlens_q[1 : batch + 1] - cu_seqlens_q[:batch] + is_decode_only = max_query_len <= 1 + page_size = key_cache.shape[1] + + sparse_kw = dict(getattr(impl, "sparse_kw", {})) + _resolve_skip_softmax_calibration( + sparse_kw, + is_prefill=not is_decode_only, + max_seq_len=max_seq_len, + ) + if is_decode_only: + # N:M sparse softmax is prefill-only. + for name in ("sparsity_n", "sparsity_m", "dense_sink_tokens", "dense_recent_tokens"): + sparse_kw.pop(name, None) + if not sparse_kw and not quant_transform_active: + # Dynamic calibration can disable sparse work for a launch. Preserve the + # backend's native dense path when no ModelOpt transform remains active. + return dense_fallback() + if prepare_modelopt is not None: + prepare_modelopt() + + v_cache_quantized = v_qdq == "nvfp4" + if v_cache_quantized: + v_qdq_scale = 1.0 if v_qdq_amax is None else v_qdq_amax / (6.0 * 448.0) + if not (math.isfinite(v_qdq_scale) and v_qdq_scale > 0): + raise ValueError(f"v_bmm_quantizer amax must be finite and positive, got {v_qdq_amax}") + prev = seq_lens - b_seq_len + fake_quant_v_onwrite( + value_cache, + block_table, + (prev // 16) * 16, + (seq_lens // 16) * 16, + max_new_tokens=max_query_len, + page_size=page_size, + v_qdq_scale=v_qdq_scale, + decode=is_decode_only, + ) + + q = query[:num_actual_tokens].contiguous() + if getattr(layer, "_query_quant_in_kernel", False): + valid_q = torch.arange(q.shape[0], device=q.device) < cu_seqlens_q[-1] + q = q.masked_fill(~valid_q[:, None, None], 0) + q = layer.q_bmm_quantizer(q.float()) + use_split_k_decode = ( + is_decode_only + and "skip_softmax_threshold" not in sparse_kw + and (p_qdq == "nvfp4" or v_qdq == "nvfp4") + ) + if use_split_k_decode: + triton_out = triton_decode_attention( + q[:batch], + key_cache, + value_cache, + block_table, + seq_lens, + softmax_scale=impl.scale, + page_size=page_size, + p_qdq=p_qdq, + p_qdq_amax=p_qdq_amax, + v_qdq=v_qdq, + v_qdq_amax=v_qdq_amax, + v_cache_quantized=v_cache_quantized, + ) + output[:batch] = triton_out + return output + + # Paged mode reads K/V through the cache. The dummy shape provides the GQA ratio. + k_dummy = torch.empty(0, impl.num_kv_heads, impl.head_size, device=q.device, dtype=q.dtype) + triton_out = triton_attention( + q, + k=k_dummy, + v=k_dummy, + b_start_loc=b_start_loc, + b_seq_len=b_seq_len, + max_input_len=max_query_len, + is_causal=is_causal, + softmax_scale=impl.scale, + b_start_loc_k=None, + b_seq_len_k=seq_lens, + max_input_len_k=max_seq_len, + k_cache=key_cache, + v_cache=value_cache, + block_table=block_table, + page_size=page_size, + p_qdq=p_qdq, + p_qdq_amax=p_qdq_amax, + v_qdq=v_qdq, + v_qdq_amax=v_qdq_amax, + v_cache_quantized=v_cache_quantized, + **sparse_kw, + ) + output[:num_actual_tokens] = triton_out + return output + + +class ModelOptSparseAttentionImpl(FlashAttentionImpl): + """FlashAttention adapter for the compact ModelOpt Triton path.""" def _forward_vllm_flash_attn( self, @@ -195,21 +342,13 @@ def forward( assert output is not None, "Output tensor must be provided." if attn_metadata is None: - # Profiling run return output.fill_(0) - quant_kw = getattr(self, "quant_kw", None) - if quant_kw is None: - p_qdq, p_qdq_amax = _p_qdq_from_layer(layer) - v_qdq, v_qdq_amax = _v_qdq_from_layer(layer) - else: - p_qdq, p_qdq_amax = quant_kw["p_qdq"], quant_kw["p_qdq_amax"] - v_qdq, v_qdq_amax = quant_kw["v_qdq"], quant_kw["v_qdq_amax"] + p_qdq, p_qdq_amax, v_qdq, v_qdq_amax = _quant_kw_from_impl(self, layer) + quant_transform_active = _has_active_quant_transform(layer, p_qdq, v_qdq) + active = _has_active_transform(self, layer, p_qdq, v_qdq) if getattr(attn_metadata, "use_cascade", False): - # vLLM cascade metadata splits the request into shared-prefix and - # suffix pieces. The ModelOpt paged kernel consumes plain per-request - # KV lengths, so delegate cascade launches back to vLLM's impl. - if p_qdq or v_qdq or getattr(self, "sparse_kw", None): + if active: raise NotImplementedError( "vLLM cascade attention is incompatible with an active ModelOpt attention transform" ) @@ -224,39 +363,31 @@ def forward( output_scale, output_block_scale, ) + if active and (output_scale is not None or output_block_scale is not None): + raise NotImplementedError("Fused attention output quantization is unsupported") - num_actual_tokens = attn_metadata.num_actual_tokens - cu_seqlens_q = attn_metadata.query_start_loc - seq_lens = attn_metadata.seq_lens - batch = seq_lens.shape[0] - b_start_loc = cu_seqlens_q[:batch] - b_seq_len = cu_seqlens_q[1 : batch + 1] - cu_seqlens_q[:batch] - - # Standard decode schedules one query token per request. Chunked - # prefill and mixed prefill/decode launches use the prefill path. - is_decode_only = attn_metadata.max_query_len <= 1 - is_causal = getattr(attn_metadata, "causal", not is_decode_only) - - # Unpack paged KV cache: [2, num_blocks, page_size, num_kv_heads, head_dim] key_cache, value_cache = kv_cache.unbind(0) - page_size = key_cache.shape[1] - - # Per-layer sparse kwargs (set by _replace_attention_impl in the worker) - sparse_kw = dict(getattr(self, "sparse_kw", {})) - _resolve_skip_softmax_calibration( - sparse_kw, - is_prefill=not is_decode_only, + is_decode_only = attn_metadata.max_query_len <= 1 + return _forward_modelopt( + self, + layer=layer, + query=query, + key_cache=key_cache, + value_cache=value_cache, + block_table=attn_metadata.block_table, + seq_lens=attn_metadata.seq_lens, + cu_seqlens_q=attn_metadata.query_start_loc, + num_actual_tokens=attn_metadata.num_actual_tokens, + max_query_len=attn_metadata.max_query_len, max_seq_len=attn_metadata.max_seq_len, - ) - if is_decode_only: - # N:M sparse softmax is prefill-only. - for name in ("sparsity_n", "sparsity_m", "dense_sink_tokens", "dense_recent_tokens"): - sparse_kw.pop(name, None) - if not sparse_kw and p_qdq is None and v_qdq is None: - # Dynamic calibration can disable sparse work for a launch, e.g. - # short-prefill thresholds outside the valid lambda range. Avoid - # swapping in the ModelOpt dense kernel when no sparse feature is active. - return self._forward_vllm_flash_attn( + is_causal=getattr(attn_metadata, "causal", not is_decode_only), + output=output, + p_qdq=p_qdq, + p_qdq_amax=p_qdq_amax, + v_qdq=v_qdq, + v_qdq_amax=v_qdq_amax, + quant_transform_active=quant_transform_active, + dense_fallback=lambda: self._forward_vllm_flash_attn( layer, query, key, @@ -266,94 +397,8 @@ def forward( output, output_scale, output_block_scale, - ) - - v_cache_quantized = v_qdq == "nvfp4" - if v_cache_quantized: - v_qdq_scale = 1.0 if v_qdq_amax is None else v_qdq_amax / (6.0 * 448.0) - if not (math.isfinite(v_qdq_scale) and v_qdq_scale > 0): - raise ValueError( - f"v_bmm_quantizer amax must be finite and positive, got {v_qdq_amax}" - ) - prev = seq_lens - b_seq_len - fake_quant_v_onwrite( - value_cache, - attn_metadata.block_table, - (prev // 16) * 16, - (seq_lens // 16) * 16, - max_new_tokens=attn_metadata.max_query_len, - page_size=page_size, - v_qdq_scale=v_qdq_scale, - decode=is_decode_only, - ) - - # Prepare metadata for our kernel - q = query[:num_actual_tokens].contiguous() - if getattr(layer, "_query_quant_in_kernel", False): - valid_q = torch.arange(q.shape[0], device=q.device) < cu_seqlens_q[-1] - q = q.masked_fill(~valid_q[:, None, None], 0) - q = layer.q_bmm_quantizer(q.float()) - use_split_k_decode = ( - is_decode_only - and "skip_softmax_threshold" not in sparse_kw - and (p_qdq == "nvfp4" or v_qdq == "nvfp4") + ), ) - if use_split_k_decode: - triton_out = triton_decode_attention( - q[:batch], - key_cache, - value_cache, - attn_metadata.block_table, - seq_lens, - softmax_scale=self.scale, - page_size=page_size, - p_qdq=p_qdq, - p_qdq_amax=p_qdq_amax, - v_qdq=v_qdq, - v_qdq_amax=v_qdq_amax, - v_cache_quantized=v_cache_quantized, - ) - output[:batch] = triton_out - return output - - # Dummy K/V for paged mode: not used by the kernel (KV are read from - # k_cache/v_cache via block_table), but shape[1] must be num_kv_heads - # so the kernel computes the correct GQA ratio (num_q_heads // num_kv_heads). - k_dummy = torch.empty(0, self.num_kv_heads, self.head_size, device=q.device, dtype=q.dtype) - - # Call ModelOpt Triton kernel with paged KV. - # b_seq_len is the query length (e.g., 6 for prefill, 1 for decode). - # b_seq_len_k is the total KV length including cache (e.g., 6 for first - # prefill, 7/8/... for subsequent decode steps). - triton_out = triton_attention( - q, - k=k_dummy, - v=k_dummy, - # Query metadata - b_start_loc=b_start_loc, - b_seq_len=b_seq_len, - max_input_len=attn_metadata.max_query_len, - is_causal=is_causal, - softmax_scale=self.scale, - # KV metadata - b_start_loc_k=None, # paged mode: KV offsets not needed - b_seq_len_k=seq_lens, # total KV length per sequence - max_input_len_k=attn_metadata.max_seq_len, - # Paged KV cache - k_cache=key_cache, # [num_blocks, page_size, num_kv_heads, head_dim] - v_cache=value_cache, # [num_blocks, page_size, num_kv_heads, head_dim] - block_table=attn_metadata.block_table, # [batch, max_blocks] - page_size=page_size, # tokens per page in the KV cache - p_qdq=p_qdq, - p_qdq_amax=p_qdq_amax, - v_qdq=v_qdq, - v_qdq_amax=v_qdq_amax, - v_cache_quantized=v_cache_quantized, - **sparse_kw, - ) - - output[:num_actual_tokens] = triton_out - return output class ModelOptSparseAttentionBackend(FlashAttentionBackend): @@ -373,13 +418,260 @@ def get_impl_cls() -> type: return ModelOptSparseAttentionImpl -def _clone_sparse_impl(old_impl): +_FLASHINFER_PATCHED = False +_FLASHINFER_IMPL_CLS: type | None = None +_FLASHINFER_METADATA_FIELDS = { + "_modelopt_block_table": "block_table_tensor", + "_modelopt_seq_lens": "seq_lens", + "_modelopt_query_start_loc": "query_start_loc", + "_modelopt_num_actual_tokens": "num_actual_tokens", + "_modelopt_max_query_len": "max_query_len", + "_modelopt_max_seq_len": "max_seq_len", + "_modelopt_causal": "causal", +} + + +def patch_flashinfer_metadata_builder() -> bool: + """Attach the common paged metadata needed by the ModelOpt kernels.""" + global _FLASHINFER_PATCHED + if _FLASHINFER_PATCHED: + return True + try: + from vllm.v1.attention.backends.flashinfer import FlashInferMetadataBuilder + except ImportError: + return False + + orig_build = FlashInferMetadataBuilder.build + if getattr(orig_build, "_modelopt_sparse_metadata_patch", False): + _FLASHINFER_PATCHED = True + return True + build_sig = inspect.signature(orig_build) + + @functools.wraps(orig_build) + def build(*args, **kwargs): + metadata = orig_build(*args, **kwargs) + common = build_sig.bind(*args, **kwargs).arguments["common_attn_metadata"] + for target, source in _FLASHINFER_METADATA_FIELDS.items(): + setattr(metadata, target, getattr(common, source)) + return metadata + + setattr(build, "_modelopt_sparse_metadata_patch", True) + FlashInferMetadataBuilder.build = build + _FLASHINFER_PATCHED = True + return True + + +def _flashinfer_cache_write(layer, key, value, kv_cache, attn_metadata, impl) -> None: + """Issue FlashInfer's native paged K/V cache write.""" + torch.ops._C_cache_ops.reshape_and_cache_flash( + key, + value, + kv_cache[:, 0], + kv_cache[:, 1], + attn_metadata.slot_mapping, + impl.kv_cache_dtype, + layer._k_scale, + layer._v_scale, + ) + + +def _maybe_update_flashinfer_cache(layer, key, value, kv_cache, attn_metadata, impl) -> None: + """Write K/V when the selected vLLM release performs updates in forward.""" + from vllm.v1.attention.backends.flashinfer import FlashInferBackend + + if not getattr(FlashInferBackend, "forward_includes_kv_cache_update", True): + return + if getattr(impl, "kv_sharing_target_layer_name", None) is not None: + return + _flashinfer_cache_write(layer, key, value, kv_cache, attn_metadata, impl) + + +def get_flashinfer_sparse_impl_cls() -> type: + """Return the lazy FlashInfer adapter without requiring it for FA users.""" + global _FLASHINFER_IMPL_CLS + if _FLASHINFER_IMPL_CLS is not None: + return _FLASHINFER_IMPL_CLS + + from vllm.v1.attention.backends.flashinfer import FlashInferImpl + + class ModelOptSparseFlashInferImpl(FlashInferImpl): + """FlashInfer adapter for the compact ModelOpt Triton path.""" + + def forward( + self, + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output=None, + output_scale=None, + output_block_scale=None, + ): + assert output is not None, "Output tensor must be provided." + if attn_metadata is None: + return output.fill_(0) + + dense_output = None + cache_prepared = False + + def dense_fallback(): + nonlocal cache_prepared, dense_output + if dense_output is None: + dense_output = FlashInferImpl.forward( + self, + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output, + output_scale, + output_block_scale, + ) + cache_prepared = True + return dense_output + + def prepare_modelopt(): + nonlocal cache_prepared + if not cache_prepared: + _maybe_update_flashinfer_cache(layer, key, value, kv_cache, attn_metadata, self) + cache_prepared = True + + p_qdq, p_qdq_amax, v_qdq, v_qdq_amax = _quant_kw_from_impl(self, layer) + quant_transform_active = _has_active_quant_transform(layer, p_qdq, v_qdq) + active = _has_active_transform(self, layer, p_qdq, v_qdq) + if getattr(attn_metadata, "use_cascade", False): + if active: + raise NotImplementedError( + "vLLM cascade attention is incompatible with an active " + "ModelOpt attention transform" + ) + return dense_fallback() + + missing = [ + name for name in _FLASHINFER_METADATA_FIELDS if not hasattr(attn_metadata, name) + ] + if missing: + if active: + raise NotImplementedError( + "FlashInfer metadata is missing the ModelOpt attention transform " + f"fields: {', '.join(missing)}" + ) + return dense_fallback() + if active and (output_scale is not None or output_block_scale is not None): + raise NotImplementedError("Fused attention output quantization is unsupported") + if kv_cache.ndim != 5 or kv_cache.shape[1] != 2: + raise ValueError( + "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim]" + ) + + key_cache = kv_cache[:, 0] + value_cache = kv_cache[:, 1] + max_query_len = attn_metadata._modelopt_max_query_len + is_decode_only = max_query_len <= 1 + common_kw = { + "layer": layer, + "key_cache": key_cache, + "value_cache": value_cache, + "max_seq_len": attn_metadata._modelopt_max_seq_len, + "is_causal": getattr(attn_metadata, "_modelopt_causal", not is_decode_only), + "p_qdq": p_qdq, + "p_qdq_amax": p_qdq_amax, + "v_qdq": v_qdq, + "v_qdq_amax": v_qdq_amax, + "quant_transform_active": quant_transform_active, + "dense_fallback": dense_fallback, + "prepare_modelopt": prepare_modelopt, + } + num_decodes = getattr(attn_metadata, "num_decodes", 0) + num_prefills = getattr(attn_metadata, "num_prefills", 0) + if not (num_decodes and num_prefills): + return _forward_modelopt( + self, + query=query, + block_table=attn_metadata._modelopt_block_table, + seq_lens=attn_metadata._modelopt_seq_lens, + cu_seqlens_q=attn_metadata._modelopt_query_start_loc, + num_actual_tokens=attn_metadata._modelopt_num_actual_tokens, + max_query_len=max_query_len, + output=output, + **common_kw, + ) + + num_decode_tokens = attn_metadata.num_decode_tokens + num_prefill_tokens = attn_metadata.num_prefill_tokens + if num_decode_tokens % num_decodes: + raise NotImplementedError("Non-uniform mixed FlashInfer decode is unsupported") + if num_decode_tokens + num_prefill_tokens != attn_metadata._modelopt_num_actual_tokens: + raise ValueError("FlashInfer mixed-batch token counts do not match common metadata") + + # Sparse-only launches may have an inactive phase (for example N:M + # is prefill-only). Compute the native result once, then overwrite + # each active phase with its ModelOpt result. + if not quant_transform_active: + dense_fallback() + + block_table = attn_metadata._modelopt_block_table + seq_lens = attn_metadata._modelopt_seq_lens + cu_seqlens_q = attn_metadata._modelopt_query_start_loc + _forward_modelopt( + self, + query=query[:num_decode_tokens], + block_table=block_table[:num_decodes], + seq_lens=seq_lens[:num_decodes], + cu_seqlens_q=cu_seqlens_q[: num_decodes + 1], + num_actual_tokens=num_decode_tokens, + max_query_len=num_decode_tokens // num_decodes, + output=output[:num_decode_tokens], + **common_kw, + ) + prefill_start = num_decode_tokens + prefill_cu_seqlens_q = cu_seqlens_q[num_decodes:] - cu_seqlens_q[num_decodes] + _forward_modelopt( + self, + query=query[prefill_start : prefill_start + num_prefill_tokens], + block_table=block_table[num_decodes : num_decodes + num_prefills], + seq_lens=seq_lens[num_decodes : num_decodes + num_prefills], + cu_seqlens_q=prefill_cu_seqlens_q, + num_actual_tokens=num_prefill_tokens, + max_query_len=max_query_len, + output=output[prefill_start : prefill_start + num_prefill_tokens], + **common_kw, + ) + return output + + _FLASHINFER_IMPL_CLS = ModelOptSparseFlashInferImpl + return _FLASHINFER_IMPL_CLS + + +def select_sparse_impl_cls(impl) -> type | None: + """Return the ModelOpt adapter matching a native vLLM implementation.""" + if isinstance(impl, ModelOptSparseAttentionImpl): + return None + if _FLASHINFER_IMPL_CLS is not None and isinstance(impl, _FLASHINFER_IMPL_CLS): + return None + if isinstance(impl, FlashAttentionImpl): + return ModelOptSparseAttentionImpl + try: + from vllm.v1.attention.backends.flashinfer import FlashInferImpl + except ImportError: + return None + if isinstance(impl, FlashInferImpl) and patch_flashinfer_metadata_builder(): + return get_flashinfer_sparse_impl_cls() + return None + + +def _clone_sparse_impl(old_impl, new_cls=None): """Create a sparse impl while preserving vLLM's initialized runtime state.""" + if new_cls is None: + new_cls = select_sparse_impl_cls(old_impl) + if new_cls is None: + raise TypeError(f"Unsupported vLLM attention implementation: {type(old_impl).__name__}") if getattr(old_impl, "sinks", None) is not None: - # vLLM passes sinks to FlashAttention as s_aux; our Triton path does not support sinks yet. - raise NotImplementedError( - "ModelOptSparseAttentionImpl does not support vLLM FlashAttention sinks yet." - ) + raise NotImplementedError(f"{new_cls.__name__} does not support attention sinks yet.") try: old_state = vars(old_impl) @@ -388,6 +680,6 @@ def _clone_sparse_impl(old_impl): "Cannot clone vLLM attention impl state: old impl does not expose __dict__." ) from err - new_impl = ModelOptSparseAttentionImpl.__new__(ModelOptSparseAttentionImpl) + new_impl = object.__new__(new_cls) new_impl.__dict__.update(old_state) return new_impl diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py index ca0deb8b7fd..c7f68c27cdf 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py @@ -571,18 +571,6 @@ def test_v_cache_finalizes_complete_groups_once(self): assert torch.all(baked[:3] == 0.015625) torch.testing.assert_close(baked[3, 0], raw[3, 0], rtol=0, atol=0) - tiny = torch.full((1, 16, 1, 16), 2e-6, device="cuda") - fake_quant_v_onwrite( - tiny, - torch.zeros(1, 1, device="cuda", dtype=torch.int32), - torch.zeros(1, device="cuda", dtype=torch.int32), - torch.tensor([16], device="cuda", dtype=torch.int32), - max_new_tokens=16, - v_qdq_scale=1.0 / (6.0 * 448.0), - ) - # Native E4M3 block scales below half the minimum subnormal round to zero. - assert torch.isfinite(tiny).all() and torch.count_nonzero(tiny) == 0 - @requires_native_e4m3 def test_v_cache_matches_independent_signed_key_axis_oracle(self): page_size, num_kv_heads, head_dim = 8, 2, 4 diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py index 41fee80eda2..dcc447f4f82 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -145,23 +145,6 @@ def test_quant_vllm_attention_forward_skips_only_in_kernel_qv_quantization(): assert attention.v_bmm_quantizer.call_count == 2 -@pytest.mark.parametrize( - "wrapper_name", ["_QuantVLLMCrossAttention", "_QuantVLLMEncoderOnlyAttention"] -) -def test_cross_encoder_attention_wrappers_keep_qkv_only_behavior(wrapper_name): - wrapper = getattr(vllm_plugin, wrapper_name) - test_wrapper = type(f"Test{wrapper.__name__}", (wrapper, _NativeAttention), {}) - attention = _new_attention(test_wrapper) - attention.q_bmm_quantizer = Mock(side_effect=lambda inputs: inputs + 1) - attention.k_bmm_quantizer = Mock(side_effect=lambda inputs: inputs + 2) - attention.v_bmm_quantizer = Mock(side_effect=lambda inputs: inputs + 3) - - output = attention(torch.tensor(10), torch.tensor(20), torch.tensor(30)) - - assert output == (torch.tensor(11), torch.tensor(22), torch.tensor(33)) - assert attention.v_bmm_quantizer.call_count == 1 - - def test_attention_kv_defaults_set_only_uncalibrated_dynamic_block16_quantizers(): calibrated_amax = 7.25 layer = SimpleNamespace( diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py index 344cb472d1b..68ab3b6bce0 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py @@ -27,10 +27,14 @@ from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.backend import AttentionType from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl +from vllm.v1.attention.backends.flashinfer import FlashInferImpl from vllm.v1.worker.gpu_worker import Worker as BaseWorker from modelopt.torch.quantization.plugins import vllm as quant_plugin -from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ModelOptSparseAttentionImpl +from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( + ModelOptSparseAttentionImpl, + get_flashinfer_sparse_impl_cls, +) _WORKER_PATH = Path(__file__).parents[5] / "examples/vllm_serve/quant_sparse_attn_worker.py" @@ -52,11 +56,11 @@ def test_worker_rejects_vllm_before_v1_attention_api(monkeypatch): _load_worker_module() -def _attention(attn_type=AttentionType.DECODER): +def _attention(attn_type=AttentionType.DECODER, impl_cls=FlashAttentionImpl): module = object.__new__(quant_plugin.vllm_attention.Attention) nn.Module.__init__(module) module.dummy = nn.Parameter(torch.zeros(1, dtype=torch.float16)) - module.impl = object.__new__(FlashAttentionImpl) + module.impl = object.__new__(impl_cls) module.impl.__dict__.update(alibi_slopes=None, logits_soft_cap=None, sinks=None) module.__dict__.update( attn_type=attn_type, @@ -86,9 +90,10 @@ def _patch_conversion(monkeypatch): monkeypatch.setattr(worker_module, "_global_errors", lambda _: []) -def test_install_converts_only_attention_and_configures_fixed_recipe(monkeypatch): +@pytest.mark.parametrize("impl_cls", [FlashAttentionImpl, FlashInferImpl]) +def test_install_converts_only_attention_and_configures_fixed_recipe(monkeypatch, impl_cls): _patch_conversion(monkeypatch) - attention = _attention() + attention = _attention(impl_cls=impl_cls) linear = nn.Linear(4, 4) model = nn.ModuleDict({"attn": attention, "linear": linear}) state = _worker(model) @@ -106,7 +111,13 @@ def test_install_converts_only_attention_and_configures_fixed_recipe(monkeypatch assert converted.v_bmm_quantizer._amax == 6.0 * 448.0 assert converted._query_quant_in_kernel is True assert converted._value_quant_in_kernel is True - assert isinstance(converted.impl, ModelOptSparseAttentionImpl) + expected_impl_cls = ( + ModelOptSparseAttentionImpl + if impl_cls is FlashAttentionImpl + else get_flashinfer_sparse_impl_cls() + ) + assert isinstance(converted.impl, impl_cls) + assert type(converted.impl) is expected_impl_cls assert converted.impl.quant_kw == { "p_qdq": "nvfp4", "p_qdq_amax": 1.0, @@ -133,17 +144,15 @@ def test_validation_of_all_layouts_precedes_mutation(monkeypatch): assert not isinstance(good, quant_plugin._QuantVLLMAttention) -def test_worker_installs_before_base_warmup(monkeypatch): +def test_worker_installs_after_base_load(monkeypatch): events = [] monkeypatch.setattr( worker_module, "_install_quant_sparse_attn", lambda _: events.append("install") ) - monkeypatch.setattr( - BaseWorker, "compile_or_warm_up_model", lambda _: events.append("base") or 1.25 - ) + monkeypatch.setattr(BaseWorker, "load_model", lambda *_args, **_kwargs: events.append("base")) instance = object.__new__(worker_module.QuantSparseAttnWorker) - assert instance.compile_or_warm_up_model() == 1.25 - assert events == ["install", "base"] + assert instance.load_model() is None + assert events == ["base", "install"] def test_memory_profile_disables_compilation(monkeypatch): @@ -180,3 +189,24 @@ def test_full_mixed_cudagraph_validation(mode, rejected): SimpleNamespace(model_runner=SimpleNamespace(vllm_config=config)) ) assert any("mixed" in error for error in errors) is rejected + + +def test_calibrated_decode_skip_softmax_rejects_full_decode_graphs(): + sparse_kw = { + "threshold_scale_factor": { + "prefill": {"a": 1.0, "b": 2.0}, + "decode": {"a": 0.1, "b": 1.0}, + } + } + + assert "calibrated decode skip-softmax" in worker_module._sparse_graph_error( + sparse_kw, CUDAGraphMode.FULL_AND_PIECEWISE + ) + assert worker_module._sparse_graph_error(sparse_kw, CUDAGraphMode.PIECEWISE) is None + assert ( + worker_module._sparse_graph_error( + {"threshold_scale_factor": {"prefill": {"a": 1.0, "b": 2.0}}}, + CUDAGraphMode.FULL_AND_PIECEWISE, + ) + is None + ) diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index ebcda487e34..1144e21b982 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -15,19 +15,31 @@ """Tests for sparse attention vLLM worker compatibility helpers.""" +import importlib.util import math from contextlib import nullcontext +from itertools import accumulate +from pathlib import Path from types import SimpleNamespace import pytest import torch +from torch import nn from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl +from vllm.v1.attention.backends.flashinfer import ( + FlashInferBackend, + FlashInferImpl, + FlashInferMetadataBuilder, +) from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as vllm_plugin from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( ModelOptSparseAttentionImpl, _build_sparse_kw, _clone_sparse_impl, + get_flashinfer_sparse_impl_cls, + patch_flashinfer_metadata_builder, + select_sparse_impl_cls, ) @@ -67,6 +79,420 @@ def test_clone_sparse_impl_rejects_non_none_sinks(): _clone_sparse_impl(old_impl) +def _make_old_flashinfer_impl(): + """Create a bare FlashInfer impl without requiring a live vLLM config.""" + impl = object.__new__(FlashInferImpl) + impl.__dict__.update( + num_heads=2, + num_kv_heads=2, + head_size=64, + scale=0.125, + sinks=None, + alibi_slopes=None, + logits_soft_cap=None, + kv_cache_dtype="auto", + ) + return impl + + +def _make_flashinfer_impl(*, sparse=False, quantized=False): + impl = _clone_sparse_impl(_make_old_flashinfer_impl(), get_flashinfer_sparse_impl_cls()) + impl.sparse_kw = {"sparsity_n": 2, "sparsity_m": 4} if sparse else {} + impl.quant_kw = { + "p_qdq": "nvfp4" if quantized else None, + "p_qdq_amax": 1.0, + "v_qdq": "nvfp4" if quantized else None, + "v_qdq_amax": 6.0 * 448.0 if quantized else None, + } + return impl + + +def test_flashinfer_metadata_builder_patch_stashes_common_metadata(monkeypatch): + """The wrapper must bind the current vLLM signature without shifting arguments.""" + calls = [] + + def fake_build(self, common_prefix_len, common_attn_metadata, fast_build=False): + calls.append((common_prefix_len, common_attn_metadata, fast_build)) + return SimpleNamespace() + + monkeypatch.setattr(FlashInferMetadataBuilder, "build", fake_build) + monkeypatch.setattr(vllm_plugin, "_FLASHINFER_PATCHED", False) + common = SimpleNamespace( + block_table_tensor=object(), + seq_lens=object(), + query_start_loc=object(), + num_actual_tokens=7, + max_query_len=3, + max_seq_len=11, + causal=False, + ) + + assert patch_flashinfer_metadata_builder() is True + patched_build = FlashInferMetadataBuilder.build + metadata = patched_build(object(), 5, common, fast_build=True) + assert calls == [(5, common, True)] + assert metadata._modelopt_block_table is common.block_table_tensor + assert metadata._modelopt_seq_lens is common.seq_lens + assert metadata._modelopt_query_start_loc is common.query_start_loc + assert metadata._modelopt_num_actual_tokens == 7 + assert metadata._modelopt_max_query_len == 3 + assert metadata._modelopt_max_seq_len == 11 + assert metadata._modelopt_causal is False + + assert patch_flashinfer_metadata_builder() is True + assert FlashInferMetadataBuilder.build is patched_build + + +def test_select_and_clone_flashinfer_impl_preserves_runtime_state(monkeypatch): + monkeypatch.setattr(vllm_plugin, "_FLASHINFER_PATCHED", False) + monkeypatch.setattr(vllm_plugin, "_FLASHINFER_IMPL_CLS", None) + old_impl = _make_old_flashinfer_impl() + old_impl.future_attr = object() + + new_cls = select_sparse_impl_cls(old_impl) + new_impl = _clone_sparse_impl(old_impl, new_cls) + + assert new_cls is get_flashinfer_sparse_impl_cls() + assert isinstance(new_impl, FlashInferImpl) + assert type(new_impl).__name__ == "ModelOptSparseFlashInferImpl" + assert new_impl.future_attr is old_impl.future_attr + assert new_impl.__dict__.items() >= old_impl.__dict__.items() + if hasattr(FlashInferImpl, "do_kv_cache_update"): + assert type(new_impl).do_kv_cache_update is FlashInferImpl.do_kv_cache_update + assert select_sparse_impl_cls(new_impl) is None + + +def _flashinfer_metadata(*, query_lens=(1, 1), seq_lens=(16, 34), max_query_len=None, mixed=False): + query_start_loc = list(accumulate(query_lens, initial=0)) + max_query_len = max_query_len or max(query_lens) + metadata = SimpleNamespace( + use_cascade=False, + _modelopt_block_table=torch.zeros(len(seq_lens), 3, dtype=torch.int32), + _modelopt_seq_lens=torch.tensor(seq_lens, dtype=torch.int32), + _modelopt_query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32), + _modelopt_num_actual_tokens=sum(query_lens), + _modelopt_max_query_len=max_query_len, + _modelopt_max_seq_len=max(seq_lens), + _modelopt_causal=max_query_len > 1, + slot_mapping=torch.arange(sum(query_lens), dtype=torch.int64), + ) + if mixed: + metadata.num_decodes = 1 + metadata.num_prefills = len(query_lens) - 1 + metadata.num_decode_tokens = query_lens[0] + metadata.num_prefill_tokens = sum(query_lens[1:]) + return metadata + + +def _flash_attention_metadata(q_len, kv_len): + return SimpleNamespace( + num_actual_tokens=q_len, + max_query_len=q_len, + max_seq_len=kv_len, + query_start_loc=torch.tensor([0, q_len], dtype=torch.int32), + seq_lens=torch.tensor([kv_len], dtype=torch.int32), + block_table=torch.zeros(1, 1, dtype=torch.int32), + ) + + +@pytest.mark.parametrize("layout", ["NHD", "HND"]) +def test_flashinfer_quantized_decode_preserves_cache_layout(monkeypatch, layout): + """Both FI layouts expose one logical cache shape with different strides.""" + impl = _make_flashinfer_impl(quantized=True) + page_size, num_heads, head_dim = 16, 2, 64 + if layout == "NHD": + kv_cache = torch.zeros(4, 2, page_size, num_heads, head_dim, dtype=torch.float16) + else: + physical = torch.zeros(4, 2, num_heads, page_size, head_dim, dtype=torch.float16) + kv_cache = physical.permute(0, 1, 3, 2, 4) + metadata = _flashinfer_metadata() + query = torch.zeros(4, num_heads, head_dim, dtype=torch.float16) + layer = SimpleNamespace( + _query_quant_in_kernel=True, + q_bmm_quantizer=lambda value: value, + ) + calls = {} + + def fake_finalize(value_cache, block_table, v_lo, v_hi, **kwargs): + calls["finalize"] = (value_cache, block_table, kwargs) + + def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): + calls["decode"] = (key_cache, value_cache, block_table, seq_lens, kwargs) + return torch.ones_like(query) + + monkeypatch.setattr(vllm_plugin, "fake_quant_v_onwrite", fake_finalize) + monkeypatch.setattr(vllm_plugin, "triton_decode_attention", fake_decode) + + output = torch.empty_like(query) + assert impl.forward(layer, query, query, query, kv_cache, metadata, output=output) is output + + key_cache, value_cache, block_table, seq_lens, decode_kw = calls["decode"] + assert key_cache.shape == (4, page_size, num_heads, head_dim) + assert value_cache.shape == key_cache.shape + assert key_cache.stride() == kv_cache[:, 0].stride() + assert value_cache.stride() == kv_cache[:, 1].stride() + assert key_cache.data_ptr() == kv_cache[:, 0].data_ptr() + assert value_cache.data_ptr() == kv_cache[:, 1].data_ptr() + assert block_table is metadata._modelopt_block_table + assert seq_lens is metadata._modelopt_seq_lens + assert decode_kw["page_size"] == page_size + assert decode_kw["p_qdq"] == "nvfp4" + assert decode_kw["v_qdq"] == "nvfp4" + finalized_cache, finalized_table, finalize_kw = calls["finalize"] + assert finalized_cache.data_ptr() == value_cache.data_ptr() + assert finalized_table is block_table + assert finalize_kw["page_size"] == page_size + + +def test_flashinfer_sparse_prefill_uses_shared_triton_kernel(monkeypatch): + """FlashInfer must route 2:4 prefill through the current ModelOpt kernel.""" + impl = _make_flashinfer_impl(sparse=True) + page_size, num_heads, head_dim = 16, 2, 64 + physical = torch.zeros(4, 2, num_heads, page_size, head_dim, dtype=torch.float16) + kv_cache = physical.permute(0, 1, 3, 2, 4) + metadata = _flashinfer_metadata(query_lens=(2, 2), max_query_len=4) + query = torch.zeros(4, num_heads, head_dim, dtype=torch.float16) + captured = {} + + def fake_attention(query, **kwargs): + captured.update(kwargs) + return torch.zeros_like(query) + + monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) + output = torch.empty_like(query) + + assert ( + impl.forward(SimpleNamespace(), query, query, query, kv_cache, metadata, output=output) + is output + ) + assert captured["sparsity_n"] == 2 + assert captured["sparsity_m"] == 4 + assert captured["page_size"] == page_size + assert captured["k_cache"].stride() == kv_cache[:, 0].stride() + assert captured["v_cache"].stride() == kv_cache[:, 1].stride() + assert captured["block_table"] is metadata._modelopt_block_table + assert captured["b_seq_len_k"] is metadata._modelopt_seq_lens + + +@pytest.mark.parametrize(("updates_in_forward", "expected_writes"), [(True, 1), (False, 0)]) +def test_flashinfer_legacy_forward_writes_kv_cache( + monkeypatch, updates_in_forward, expected_writes +): + """vLLM releases where FlashInfer updates in forward must retain that write.""" + impl = _make_flashinfer_impl(sparse=True) + impl.kv_sharing_target_layer_name = None + query = torch.zeros(4, 2, 64, dtype=torch.float16) + kv_cache = torch.zeros(4, 2, 16, 2, 64, dtype=torch.float16) + metadata = _flashinfer_metadata(query_lens=(2, 2), max_query_len=4) + writes = [] + + monkeypatch.setattr(FlashInferBackend, "forward_includes_kv_cache_update", updates_in_forward) + monkeypatch.setattr(vllm_plugin, "_flashinfer_cache_write", lambda *args: writes.append(args)) + monkeypatch.setattr( + vllm_plugin, "triton_attention", lambda query, **kwargs: torch.zeros_like(query) + ) + + layer = SimpleNamespace() + impl.forward(layer, query, query, query, kv_cache, metadata, output=torch.empty_like(query)) + + assert len(writes) == expected_writes + if expected_writes: + assert writes == [(layer, query, query, kv_cache, metadata, impl)] + + +def test_flashinfer_q_only_transform_does_not_fallback(monkeypatch): + """Withheld Q QDQ must run even when sparse/P/V transforms are inactive.""" + impl = _make_flashinfer_impl() + query = torch.zeros(4, 2, 64, dtype=torch.float16) + kv_cache = torch.zeros(4, 2, 16, 2, 64, dtype=torch.float16) + metadata = _flashinfer_metadata(query_lens=(2, 2), max_query_len=4) + captured = {} + layer = SimpleNamespace( + _query_quant_in_kernel=True, + q_bmm_quantizer=lambda value: value + 1, + ) + + monkeypatch.setattr( + FlashInferImpl, + "forward", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("native fallback")), + ) + + def fake_attention(query, **kwargs): + captured["query"] = query + return torch.zeros_like(query) + + monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) + + impl.forward(layer, query, query, query, kv_cache, metadata, output=torch.empty_like(query)) + + assert torch.all(captured["query"] == 1) + + +def test_flashinfer_legacy_inactive_launch_writes_only_in_native_fallback(monkeypatch): + impl = _make_flashinfer_impl(sparse=True) + impl.kv_sharing_target_layer_name = None + query = torch.zeros(1, 2, 64, dtype=torch.float16) + kv_cache = torch.zeros(1, 2, 16, 2, 64, dtype=torch.float16) + metadata = _flashinfer_metadata(query_lens=(1,), seq_lens=(16,)) + manual_writes = [] + native_calls = [] + + monkeypatch.setattr(FlashInferBackend, "forward_includes_kv_cache_update", True) + monkeypatch.setattr( + vllm_plugin, "_flashinfer_cache_write", lambda *args: manual_writes.append(args) + ) + monkeypatch.setattr( + FlashInferImpl, + "forward", + lambda *args, **kwargs: native_calls.append(True) or kwargs.get("output", args[7]), + ) + + impl.forward( + SimpleNamespace(), query, query, query, kv_cache, metadata, output=torch.empty_like(query) + ) + + assert manual_writes == [] + assert native_calls == [True] + + +@pytest.mark.parametrize("quantized", [False, True], ids=["sparse-only", "quantized"]) +def test_flashinfer_mixed_batch_splits_decode_and_prefill(monkeypatch, quantized): + prefill_tokens = 17 + impl = _make_flashinfer_impl(sparse=True, quantized=quantized) + query = torch.zeros(1 + prefill_tokens, 2, 64, dtype=torch.float16) + kv_cache = torch.zeros(4, 2, 16, 2, 64, dtype=torch.float16) + metadata = _flashinfer_metadata(query_lens=(1, prefill_tokens), mixed=True) + layer = SimpleNamespace( + _query_quant_in_kernel=quantized, + q_bmm_quantizer=lambda value: value, + ) + calls = {"native": [], "decode": [], "prefill": [], "finalize": []} + + def fake_decode(query, *args, **kwargs): + calls["decode"].append((query, kwargs)) + return torch.zeros_like(query) + + def fake_attention(query, **kwargs): + calls["prefill"].append((query, kwargs)) + return torch.zeros_like(query) + + monkeypatch.setattr( + FlashInferImpl, + "forward", + lambda *args, **kwargs: calls["native"].append(True) or args[7], + ) + monkeypatch.setattr( + vllm_plugin, + "fake_quant_v_onwrite", + lambda *args, **kwargs: calls["finalize"].append(kwargs), + ) + monkeypatch.setattr(vllm_plugin, "triton_decode_attention", fake_decode) + monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) + + impl.forward(layer, query, query, query, kv_cache, metadata, output=torch.empty_like(query)) + + assert len(calls["prefill"]) == 1 + prefill_query, prefill_kw = calls["prefill"][0] + assert prefill_query.shape[0] == prefill_tokens + assert prefill_kw["sparsity_n"] == 2 + assert prefill_kw["sparsity_m"] == 4 + torch.testing.assert_close( + prefill_kw["b_seq_len"], torch.tensor([prefill_tokens], dtype=torch.int32) + ) + + if quantized: + assert calls["native"] == [] + assert len(calls["decode"]) == 1 + decode_query, decode_kw = calls["decode"][0] + assert decode_query.shape[0] == 1 + assert decode_kw["p_qdq"] == "nvfp4" + assert "sparsity_n" not in decode_kw + assert prefill_kw["p_qdq"] == "nvfp4" + assert [(kw["max_new_tokens"], kw["decode"]) for kw in calls["finalize"]] == [ + (1, True), + (prefill_tokens, False), + ] + else: + assert calls["native"] == [True] + assert calls["decode"] == [] + assert calls["finalize"] == [] + + +def test_sparse_worker_rejects_unsupported_backend_before_mutation(monkeypatch): + worker_path = Path(__file__).parents[5] / "examples/vllm_serve/sparse_attn_worker.py" + spec = importlib.util.spec_from_file_location("sparse_attn_worker_test", worker_path) + worker_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(worker_module) + good = object.__new__(worker_module.VLLMAttention) + nn.Module.__init__(good) + good.impl = object.__new__(FlashAttentionImpl) + good.impl.sinks = None + bad = object.__new__(worker_module.VLLMAttention) + nn.Module.__init__(bad) + bad.impl = SimpleNamespace() + original_good_impl = good.impl + model = nn.ModuleDict({"good": good, "bad": bad}) + state = SimpleNamespace( + model_runner=SimpleNamespace( + model=model, + model_config=SimpleNamespace(hf_config=SimpleNamespace()), + ) + ) + monkeypatch.setattr( + worker_module, + "load_from_checkpoint_metadata", + lambda _: ({}, "test"), + ) + monkeypatch.setattr( + worker_module, + "match_sparse_config", + lambda *_: {"enable": True, "sparsity_n": 2, "sparsity_m": 4}, + ) + + with pytest.raises(NotImplementedError, match=r"bad.*SimpleNamespace"): + worker_module._replace_attention_impl(state) + + assert good.impl is original_good_impl + assert bad.impl.__class__ is SimpleNamespace + + +@pytest.mark.parametrize( + ("failure", "active"), + [("cascade", True), ("cascade", False), ("metadata", True), ("metadata", False)], +) +def test_flashinfer_transform_safety_for_unsupported_metadata(monkeypatch, failure, active): + impl = _make_flashinfer_impl(sparse=active) + metadata = ( + SimpleNamespace(use_cascade=True) + if failure == "cascade" + else SimpleNamespace(use_cascade=False) + ) + native_calls = [] + query = torch.zeros(1, 2, 64, dtype=torch.float16) + output = torch.empty_like(query) + + def fake_forward(self, *args, **kwargs): + native_calls.append((self, args[5])) + output = args[6] + output.fill_(9) + return output + + monkeypatch.setattr(FlashInferImpl, "forward", fake_forward) + + if active: + with pytest.raises(NotImplementedError, match="ModelOpt attention transform"): + impl.forward(None, query, query, query, torch.empty(0), metadata, output=output) + assert native_calls == [] + else: + assert ( + impl.forward(None, query, query, query, torch.empty(0), metadata, output=output) + is output + ) + assert native_calls == [(impl, metadata)] + assert torch.all(output == 9) + + def test_forward_delegates_cascade_metadata_to_vllm(monkeypatch): """Cascade/prefix-cache metadata should use vLLM's native implementation.""" impl = _clone_sparse_impl(_make_old_impl()) @@ -129,6 +555,16 @@ def fake_forward( 4, 4, ), + ( + { + "sparsity_n": 2, + "sparsity_m": 4, + "dense_sink_tokens": 4, + "dense_recent_tokens": 128, + }, + 1, + 16, + ), ], ) def test_forward_delegates_launches_without_effective_sparse_work( @@ -142,18 +578,7 @@ def test_forward_delegates_launches_without_effective_sparse_work( 2, 1, max_seq_len, impl.num_kv_heads, impl.head_size, dtype=torch.float16 ) output = torch.empty_like(q) - attn_metadata = type( - "AttnMetadata", - (), - { - "num_actual_tokens": max_query_len, - "max_query_len": max_query_len, - "max_seq_len": max_seq_len, - "query_start_loc": torch.tensor([0, max_query_len], dtype=torch.int32), - "seq_lens": torch.tensor([max_seq_len], dtype=torch.int32), - "block_table": torch.zeros(1, 1, dtype=torch.int32), - }, - )() + attn_metadata = _flash_attention_metadata(max_query_len, max_seq_len) called = {} def fake_attention(*args, **kwargs): @@ -215,18 +640,7 @@ def test_forward_resolves_calibrated_skip_softmax_threshold(monkeypatch): } q = torch.zeros(max_query_len, impl.num_heads, impl.head_size, dtype=torch.float16) kv_cache = torch.zeros(2, 1, seq_len, impl.num_kv_heads, impl.head_size, dtype=torch.float16) - attn_metadata = type( - "AttnMetadata", - (), - { - "num_actual_tokens": max_query_len, - "max_query_len": max_query_len, - "max_seq_len": seq_len, - "query_start_loc": torch.tensor([0, max_query_len], dtype=torch.int32), - "seq_lens": torch.tensor([seq_len], dtype=torch.int32), - "block_table": torch.zeros(1, 1, dtype=torch.int32), - }, - )() + attn_metadata = _flash_attention_metadata(max_query_len, seq_len) captured = {} def fake_attention(q, **kwargs): @@ -368,41 +782,6 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): assert calls["query"].dtype == torch.float32 -def test_quantized_prefill_passes_host_v_group_bound(monkeypatch): - impl = _clone_sparse_impl(_make_old_impl()) - impl.quant_kw = { - "p_qdq": None, - "p_qdq_amax": 1.0, - "v_qdq": "nvfp4", - "v_qdq_amax": 6.0 * 448.0, - } - q_len, kv_len = 17, 31 - q = torch.zeros(q_len, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 2, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) - metadata = SimpleNamespace( - num_actual_tokens=q_len, - max_query_len=q_len, - max_seq_len=kv_len, - query_start_loc=torch.tensor([0, q_len], dtype=torch.int32), - seq_lens=torch.tensor([kv_len], dtype=torch.int32), - block_table=torch.zeros(1, 2, dtype=torch.int32), - ) - calls = {} - - def fake_finalize(*args, **kwargs): - calls["finalize"] = kwargs - - monkeypatch.setattr(vllm_plugin, "fake_quant_v_onwrite", fake_finalize) - monkeypatch.setattr( - vllm_plugin, "triton_attention", lambda query, **kwargs: torch.zeros_like(query) - ) - - impl.forward(None, q, q, q, kv_cache, metadata, output=torch.empty_like(q)) - - assert calls["finalize"]["max_new_tokens"] == q_len - assert calls["finalize"]["decode"] is False - - def test_quantized_skip_softmax_decode_stays_on_shared_kernel(monkeypatch): """Split-local maxima must not change calibrated skip-softmax semantics.""" impl = _clone_sparse_impl(_make_old_impl()) @@ -415,14 +794,7 @@ def test_quantized_skip_softmax_decode_stays_on_shared_kernel(monkeypatch): impl.sparse_kw = {"skip_softmax_threshold": 0.001} q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16) kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) - metadata = SimpleNamespace( - num_actual_tokens=1, - max_query_len=1, - max_seq_len=16, - query_start_loc=torch.tensor([0, 1], dtype=torch.int32), - seq_lens=torch.tensor([16], dtype=torch.int32), - block_table=torch.zeros(1, 1, dtype=torch.int32), - ) + metadata = _flash_attention_metadata(1, 16) captured = {} monkeypatch.setattr(vllm_plugin, "fake_quant_v_onwrite", lambda *args, **kwargs: None) @@ -516,66 +888,6 @@ def test_build_sparse_kw_restores_checkpoint_sparse_metadata(): } -def test_forward_delegates_sparse_nm_only_decode_to_vllm(monkeypatch): - """N:M sparse softmax is prefill-only, so N:M-only decode uses vLLM.""" - impl = _clone_sparse_impl(_make_old_impl()) - impl.sparse_kw = { - "sparsity_n": 2, - "sparsity_m": 4, - "dense_sink_tokens": 4, - "dense_recent_tokens": 128, - } - q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) - attn_metadata = type( - "AttnMetadata", - (), - { - "num_actual_tokens": 1, - "max_query_len": 1, - "max_seq_len": 16, - "query_start_loc": torch.tensor([0, 1], dtype=torch.int32), - "seq_lens": torch.tensor([16], dtype=torch.int32), - "block_table": torch.zeros(1, 1, dtype=torch.int32), - }, - )() - - def fake_attention(q, **kwargs): - raise AssertionError("N:M-only decode should not call ModelOpt Triton") - - def fake_forward( - self, - layer, - query, - key, - value, - kv_cache_arg, - attn_metadata_arg, - output_arg=None, - output_scale=None, - output_block_scale=None, - ): - output_arg.fill_(7) - return output_arg - - monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) - monkeypatch.setattr(FlashAttentionImpl, "forward", fake_forward) - - output = torch.empty_like(q) - result = impl.forward( - layer=None, - query=q, - key=q, - value=q, - kv_cache=kv_cache, - attn_metadata=attn_metadata, - output=output, - ) - - assert result is output - assert torch.all(result == 7) - - def test_forward_allows_chunked_prefill_metadata(monkeypatch): """vLLM V1 can pass suffix-Q/chunked-prefill metadata; the kernel handles it.""" impl = _clone_sparse_impl(_make_old_impl()) @@ -584,18 +896,7 @@ def test_forward_allows_chunked_prefill_metadata(monkeypatch): kv_len = 10 q = torch.zeros(q_len, impl.num_heads, impl.head_size, dtype=torch.float16) kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) - attn_metadata = type( - "AttnMetadata", - (), - { - "num_actual_tokens": q_len, - "max_query_len": q_len, - "max_seq_len": kv_len, - "query_start_loc": torch.tensor([0, q_len], dtype=torch.int32), - "seq_lens": torch.tensor([kv_len], dtype=torch.int32), - "block_table": torch.zeros(1, 1, dtype=torch.int32), - }, - )() + attn_metadata = _flash_attention_metadata(q_len, kv_len) captured = {} def fake_attention(q, **kwargs): diff --git a/tests/unit/torch/kernels/common/attention/test_triton_fa.py b/tests/unit/torch/kernels/common/attention/test_triton_fa.py index 6969ae0a0e3..0c4aff9c948 100644 --- a/tests/unit/torch/kernels/common/attention/test_triton_fa.py +++ b/tests/unit/torch/kernels/common/attention/test_triton_fa.py @@ -13,17 +13,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""CPU smoke tests for the Triton flash attention module. +"""CPU tests for the Triton flash attention module. The ``@triton.jit`` kernels and the ``attention`` / ``attention_calibrate`` Python wrappers require a GPU and are fully exercised in ``tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa*.py``. -This file only verifies that the module is importable on CPU-only CI runners, -so upstream code paths that conditionally import it don't break. +These tests verify CPU-safe wrapper behavior without executing a Triton kernel. """ +from contextlib import nullcontext + import pytest +import torch def test_triton_fa_importable_on_cpu(): @@ -38,3 +40,37 @@ def test_triton_fa_importable_on_cpu(): assert "attention" in triton_fa.__all__ assert callable(calibrate.attention_calibrate) + + +def test_forward_buckets_autotune_key_without_bucketing_grid(monkeypatch): + """Reuse autotune results by length regime without launching extra query tiles.""" + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + class CapturingKernel: + def __getitem__(self, grid): + self.grid = grid + + def launch(*args, **kwargs): + self.kwargs = kwargs + + return launch + + kernel = CapturingKernel() + monkeypatch.setattr(triton_fa, "_attn_fwd", kernel) + monkeypatch.setattr(triton_fa.torch.cuda, "device", lambda _device: nullcontext()) + monkeypatch.setattr(triton_fa, "_load_sparsity_helpers", lambda: None) + monkeypatch.setattr(triton_fa, "_load_qdq_helpers", lambda: None) + + seq_len = 129 + q = torch.empty(seq_len, 2, 16) + k = torch.empty(seq_len, 1, 16) + v = torch.empty_like(k) + starts = torch.tensor([0], dtype=torch.int32) + lengths = torch.tensor([seq_len], dtype=torch.int32) + + triton_fa.attention(q, k, v, starts, lengths, seq_len) + + assert kernel.kwargs["N_CTX"] == 256 + assert kernel.grid({"BLOCK_M": 64}) == (1, 2, 3) From 633e341169c23f5b3f2397189d5118591e9cff0d Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sun, 5 Jul 2026 17:52:13 -0700 Subject: [PATCH 10/28] Document minimal MNI attention autotune design Signed-off-by: Kai Xu --- ...5-minimal-mni-attention-autotune-design.md | 101 ++++++++++++++++++ 1 file changed, 101 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-05-minimal-mni-attention-autotune-design.md diff --git a/docs/superpowers/specs/2026-07-05-minimal-mni-attention-autotune-design.md b/docs/superpowers/specs/2026-07-05-minimal-mni-attention-autotune-design.md new file mode 100644 index 00000000000..4732cc808d1 --- /dev/null +++ b/docs/superpowers/specs/2026-07-05-minimal-mni-attention-autotune-design.md @@ -0,0 +1,101 @@ +# Minimal MNI-Aligned Attention Autotune + +## Status + +Approved in conversation on 2026-07-05. + +## Context + +The compact attention branch currently builds a union of dense and P-QDQ Triton +configurations, uses a long semantic cache key, and prunes configurations according +to P-QDQ and sparse-mask flags. This makes the launch policy difficult to review. + +The authoritative MNI reference uses one prefill schedule for dense and every +Q/K/P/V QDQ combination. Its KV tile is fixed at 32, independent of quantization +flags. B200 measurements of the ModelOpt kernel with that fixed KV tile selected +query tiles 16, 64, and 128 at sequence lengths 512, 2048, and 8192 respectively; +all selected two stages and four warps. + +## Goals + +- Use one small, identical eligible configuration set for all normal forward modes. +- Fix the KV tile to 32 to match the MNI prefill numerical schedule. +- Retain query-tile autotuning for short and long prefill efficiency. +- Remove the union configuration construction and mode-aware pruning function. +- Keep sparsity-counter measurement as a direct single launch because autotune + trials would mutate its atomic counters. + +## Non-Goals + +- Reproduce MNI decode tiling or replace the compact split-K decode kernel. +- Align Q/K/P/V carrier dtypes, dot precision, global-scale lifecycle, or every + other MNI arithmetic detail. +- Redesign skip-softmax or N:M dense-region semantics. +- Claim that KV tile 32 is the unrestricted throughput winner. It is the MNI + compatibility contract; unrestricted B200 tuning selected a different tile. + +## Design + +Normal forward launches use one Triton autotuner with three configurations: + +```text +BLOCK_M = {16, 64, 128} +BLOCK_N = 32 +num_stages = 2 +num_warps = 4 +``` + +The autotune key contains only the shape regime and QDQ properties that can change +the best query tile: + +```text +N_CTX, HEAD_DIM, Q_IS_FP32, P_QDQ, V_QDQ +``` + +K QDQ does not need a kernel key because it is materialized before this kernel and +does not change the kernel's dtype or temporary storage. Under pytest, the eligible +set is reduced to the single `BLOCK_M=16, BLOCK_N=32` configuration so tests do not +run benchmark warmups. + +The following implementation elements are removed: + +- `_FWD_BLOCK_M_CHOICES` +- `_FWD_BLOCK_N_CHOICES` +- `_P_QDQ_BLOCK_M_CHOICES` +- `_FWD_AUTOTUNE_KEYS` +- `_prune_fwd_configs` +- Tests coupled to Triton's early-pruning internals + +The direct measurement launch remains outside the autotuner. Its existing stable +measurement geometry remains unchanged in this refactor; changing skip-softmax +calibration geometry is separate work. + +## Behavioral Consequences + +- Dense, Q-only, K-only, P-only, V-only, and combined Q/K/P/V paths see the same + eligible physical tile set. +- P-QDQ output no longer depends on a performance-selected KV tile because every + eligible configuration uses KV tile 32. +- The selected query tile may differ by sequence-length bucket and QDQ combination, + but query tiling does not change P-QDQ's per-row block-16 quantization groups. +- Normal P-QDQ plus 2:4 or skip-softmax uses the same small configuration set instead + of a special fixed autotune profile. The BM16 candidate remains available as the + low-resource fallback. + +## Testing + +- Assert the production configuration set has exactly the three declared tile + shapes and fixed launch parameters. +- Assert the autotune key is the reduced five-field key. +- Assert all non-measurement QDQ combinations route through `_attn_fwd`. +- Retain the direct measurement-launch tests. +- Update the P-QDQ oracle to use KV tile 32. +- Run focused CPU routing tests, P-QDQ GPU tests, and a real dense Triton launch. +- Run repository pre-commit hooks on every touched file. + +## Known Risk + +MNI's `BLOCK_M` fuses GQA heads and query tokens, while ModelOpt's `BLOCK_M` counts +query tokens for one head. Matching the numeric value does not make the physical +grids identical. This design aligns the P-QDQ-sensitive KV tile and simplifies the +policy without claiming bitwise MNI kernel equivalence. From 8dbb2ca4f100e48102d497ce93ab922e1f8ec01d Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 6 Jul 2026 11:59:07 -0700 Subject: [PATCH 11/28] Document unified vLLM attention worker design Signed-off-by: Kai Xu --- ...06-unified-vllm-attention-worker-design.md | 218 ++++++++++++++++++ 1 file changed, 218 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-06-unified-vllm-attention-worker-design.md diff --git a/docs/superpowers/specs/2026-07-06-unified-vllm-attention-worker-design.md b/docs/superpowers/specs/2026-07-06-unified-vllm-attention-worker-design.md new file mode 100644 index 00000000000..6479bd2727c --- /dev/null +++ b/docs/superpowers/specs/2026-07-06-unified-vllm-attention-worker-design.md @@ -0,0 +1,218 @@ +# Unified vLLM Attention Worker Design + +## Goal + +Remove duplicated sparse-attention planning and installation code from the vLLM +example workers while preserving two explicit serving policies: + +- `SparseAttnWorker` applies checkpoint-driven sparse attention only. +- `QuantSparseAttnWorker` applies the fixed block-16 NVFP4 Q/K/P/V recipe to + every supported attention layer and optionally adds checkpoint-driven sparse + attention. + +The refactor must reduce production code, retain atomic validation, and leave +attention numerics and kernel dispatch unchanged. + +## Current Problem + +`sparse_attn_worker.py` and `quant_sparse_attn_worker.py` independently perform +the same orchestration: + +1. Unwrap the loaded vLLM model. +2. Read sparse-attention checkpoint metadata. +3. Traverse attention modules. +4. Match per-layer sparse configuration. +5. Select and clone the backend-matched ModelOpt implementation. +6. Validate the complete plan before mutating the model. +7. Install the replacement immediately after `BaseWorker.load_model`. + +The quant worker adds strict runtime and layout validation, Q/K/P/V quantizer +configuration, quantization arguments, cascade disablement, and compilation- +disabled memory profiling. Those are policy extensions of the same attention +installation lifecycle, not a separate serving architecture. + +The launcher is not part of this duplication. It owns CLI parsing, worker-module +import setup, and API-server startup, while the worker runs inside each spawned +vLLM GPU process where the model exists. + +## Architecture + +### One Source Module + +`examples/vllm_serve/sparse_attn_worker.py` becomes the single attention-worker +module and declares both public classes in `__all__`: + +```python +class _ModelOptAttentionWorker(BaseWorker): + quantize_attention = False + + def load_model(self, *args, **kwargs): + super().load_model(*args, **kwargs) + _install_attention(self, quantize=self.quantize_attention) + + +class SparseAttnWorker(_ModelOptAttentionWorker): + pass + + +class QuantSparseAttnWorker(_ModelOptAttentionWorker): + quantize_attention = True +``` + +`QuantSparseAttnWorker` retains its `determine_available_memory` override. The +override remains quant-only because sparse-only serving does not need the +compilation workaround and must not inherit its custom-all-reduce interaction. + +`examples/vllm_serve/quant_sparse_attn_worker.py` is deleted. The new quant +worker path is: + +```text +sparse_attn_worker.QuantSparseAttnWorker +``` + +The existing `sparse_attn_worker.SparseAttnWorker` path remains unchanged. + +### Shared Planning Pipeline + +The common pipeline produces immutable plan records before any module mutation. +Each record contains the layer name, attention module, backend-matched +replacement implementation, sparse keyword arguments, and quant-specific +device/dtype fields when quantization is enabled. + +The pipeline performs these shared steps once: + +1. Resolve the loaded model and optional sparse checkpoint configuration. +2. Enumerate vLLM attention modules. +3. Match and normalize each layer's sparse configuration. +4. Select and clone the FlashAttention or FlashInfer adapter. +5. Accumulate all errors and reject the complete plan before mutation. + +Policy-specific validation is invoked from this pipeline rather than mixed into +the shared traversal. + +### Sparse-Only Policy + +With `quantize=False`: + +- Missing `sparse_attention_config` logs the existing message and leaves the + model unchanged. +- Only layers with enabled, nonempty sparse configuration enter the plan. +- Unsupported backends are rejected only for layers requesting a sparse + transform. +- No quantizers, quantization flags, quantization arguments, or profiling + overrides are installed. +- Existing native fallback behavior for inactive sparse launches is preserved. + +### Quant-Plus-Sparse Policy + +With `quantize=True`: + +- vLLM 0.14 or newer remains required. +- Every regular decoder self-attention layer must enter the plan, even when the + checkpoint has no sparse metadata. +- Existing global and per-layer quant validation remains mandatory. +- Each layer receives the fixed dynamic block-16 NVFP4 Q/K/P/V recipe. +- Optional sparse metadata augments the same backend-matched implementation. +- P/V quantization arguments, Q/V kernel ownership flags, K/V default scales, + and cascade disablement remain unchanged. +- The compilation-disabled memory-profile override remains active. + +Quant-only vLLM and ModelOpt imports are resolved only when quant mode is used. +Importing or running `SparseAttnWorker` must not acquire the quant worker's +vLLM-version floor. The quant version/API check therefore runs when the quant +policy is selected, not merely when the shared module is imported. + +## Startup Data Flow + +The vLLM lifecycle remains: + +1. The launcher selects a worker class by dotted path. +2. Each GPU process creates that worker and calls `load_model`. +3. The base worker constructs the native attention modules and implementations. +4. `_ModelOptAttentionWorker.load_model` builds and validates the complete + ModelOpt plan, then installs it. +5. vLLM discovers KV-cache specs and runs memory profiling using the installed + implementation. +6. vLLM initializes cache/backend metadata and performs warmup and CUDA-graph + capture. + +Installation must not move to `compile_or_warm_up_model`; that hook is too late +for attention profiling, FlashInfer metadata setup, and cache initialization. + +## Compatibility + +- Preserve the `sparse_attn_worker.SparseAttnWorker` public path and observable + sparse-only behavior. +- Change the unreleased compact-worker path from + `quant_sparse_attn_worker.QuantSparseAttnWorker` to + `sparse_attn_worker.QuantSparseAttnWorker` in tracked examples, documentation, + and tests. +- Do not add an environment variable or CLI mode flag. The two class names are + the explicit policy selectors. +- Preserve sparse-only import behavior on vLLM releases older than the quant + worker's 0.14 minimum. +- Keep `vllm_serve_sparse_attn.py` as a separate launcher and retain its + sparse-only default. +- Do not merge this path with `FakeQuantWorker`; whole-model fakequant uses a + later lifecycle hook because calibration requires initialized KV-cache state. +- Preserve the current `MODELOPT_ATTN_QUANT_OFF` isolation behavior if it is + present when implementation begins. Removing that benchmark-only behavior is + a separate cleanup decision, not part of worker consolidation. + +## Error Handling + +- Both policies validate the entire selected layer set before the first model + mutation. +- Sparse-only mode continues to no-op when no sparse metadata is present. +- Quant mode fails if no regular attention layers are found or if any selected + layer/runtime feature violates the existing quant contract. +- Backend clone errors and unsupported attention implementations are reported + with layer names in one aggregated `NotImplementedError`. +- A quant-only import/version failure is raised only when + `QuantSparseAttnWorker` is selected. + +## Testing + +Retain the focused sparse and quant worker tests as separate test modules while +pointing both at `sparse_attn_worker.py`. + +Required coverage: + +- Both public classes install after the base `load_model` call. +- Sparse mode remains a no-op without checkpoint metadata. +- Sparse mode mutates only layers with active sparse configuration. +- Quant mode converts every supported regular attention layer and preserves the + exact Q/K/P/V NVFP4 recipe for FlashAttention and FlashInfer. +- Both modes reject a multi-layer plan atomically before mutation. +- Quant-only runtime/layout validation does not reject sparse-only serving. +- Sparse-only import remains valid when quant-only vLLM APIs are unavailable. +- Selecting the quant class below vLLM 0.14 produces the existing clear error. +- Quant memory profiling runs under `disable_compilation`; sparse-only profiling + does not gain that override. +- Full and calibrated-decode CUDA-graph validation remains unchanged. + +Run the focused worker and plugin suites, formatting checks, and an import smoke +for both dotted class paths. A single vLLM serving smoke with the quant class is +required after the refactor because the worker path changes. + +## Scope And Expected Reduction + +The implementation should remove duplicated model traversal, sparse metadata +handling, adapter planning/cloning, load lifecycle, module header, and imports. +The combined production source must be materially smaller than the current 404 +lines across the two workers; a net reduction of approximately 40-70 lines is +expected without weakening validation or deleting regression coverage. + +If preserving older sparse-only imports requires conditional imports, keep that +logic local to the quant policy. Do not create another helper module solely to +move lines between files. + +## Non-Goals + +- Changing attention kernels, tile sizes, quantization scales, or numerical + behavior. +- Replacing the two explicit public worker policies with automatic inference. +- Merging the worker with `vllm_serve_sparse_attn.py`. +- Merging attention-only serving with whole-model `FakeQuantWorker` calibration. +- Adding new sparse algorithms, attention backends, or vLLM feature support. +- Removing focused tests solely to reduce the pull request's line count. From 26e628f0a4e46864e9cb768fb3a7e37768e20deb Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 6 Jul 2026 12:28:15 -0700 Subject: [PATCH 12/28] Unify vLLM attention worker entry points Signed-off-by: Kai Xu --- examples/vllm_serve/sparse_attn_worker.py | 97 +++++++-- .../test_sparse_attn_worker.py | 186 +++++++++++++++++- 2 files changed, 257 insertions(+), 26 deletions(-) diff --git a/examples/vllm_serve/sparse_attn_worker.py b/examples/vllm_serve/sparse_attn_worker.py index 4b8c5e84579..14798f7579c 100644 --- a/examples/vllm_serve/sparse_attn_worker.py +++ b/examples/vllm_serve/sparse_attn_worker.py @@ -13,26 +13,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Custom vLLM worker for sparse attention. +"""Custom vLLM workers for ModelOpt attention policies. ``SparseAttnWorker`` replaces the native FlashAttention or FlashInfer impl with -the matching ModelOpt adapter on each Attention module after model loading. -The sparse impl uses the ModelOpt Triton kernel for sparse prefill launches and -delegates inactive launches to the selected native backend. +the matching checkpoint-driven ModelOpt sparse adapter. ``QuantSparseAttnWorker`` +installs fixed NVFP4 attention plus optional checkpoint-driven sparsity. Both +policies install their attention implementation after model loading. Configuration flows exclusively through the loaded checkpoint's ``sparse_attention_config`` block (written by ModelOpt's HF export). If the -checkpoint has no such block, the worker logs a message and passes through -unchanged. - -This worker remains sparse-only; the compact quant+sparse path uses -``quant_sparse_attn_worker.QuantSparseAttnWorker``. +checkpoint has no such block, the sparse-only policy logs a message and passes +through unchanged. Usage: python vllm_serve_sparse_attn.py """ import importlib +from functools import cache +from types import SimpleNamespace try: _has_legacy_attention_layer = importlib.util.find_spec("vllm.attention.layer") is not None @@ -56,6 +55,45 @@ select_sparse_impl_cls, ) +__all__ = ["SparseAttnWorker", "QuantSparseAttnWorker"] # noqa: RUF022 + + +def _unwrapped_model(worker): + model = worker.model_runner.model + return model.unwrap() if hasattr(model, "unwrap") else model + + +@cache +def _load_quant_api(vllm_version: str): + # Keep sparse-only module loading independent of quant-specific vLLM APIs. + import torch + from packaging import version + + if version.parse(vllm_version) < version.parse("0.14.0"): + raise RuntimeError("The compact NVFP4 attention worker requires vLLM >= 0.14.0") + + from vllm.config import compilation + from vllm.v1.attention import backend + + from modelopt.torch.quantization import conversion + from modelopt.torch.quantization import nn as quant_nn + from modelopt.torch.quantization.plugins import vllm as quant_plugin + + return SimpleNamespace( + torch=torch, + compilation=compilation, + backend=backend, + conversion=conversion, + nn=quant_nn, + plugin=quant_plugin, + ) + + +def _quant_api(): + import vllm + + return _load_quant_api(vllm.__version__) + def _replace_attention_impl(worker): """Install the backend-matched ModelOpt sparse impl on attention layers. @@ -76,9 +114,7 @@ def _replace_attention_impl(worker): cfg, preset_name = detected print(f"[ModelOpt] Sparse attention config: algo -> {preset_name}") - model = worker.model_runner.model - if hasattr(model, "unwrap"): - model = model.unwrap() + model = _unwrapped_model(worker) plans = [] errors = [] @@ -122,19 +158,40 @@ def _replace_attention_impl(worker): ) +def _install_attention(worker, *, quantize: bool) -> None: + if quantize: + _quant_api() + # Keep the compatibility delegation lazy until both policies share one planner. + from quant_sparse_attn_worker import _install_quant_sparse_attn + + _install_quant_sparse_attn(worker) + return + _replace_attention_impl(worker) + + # --------------------------------------------------------------------------- # Workers # --------------------------------------------------------------------------- -class SparseAttnWorker(BaseWorker): - """vLLM worker that uses the ModelOpt sparse attention backend. - - Replaces FlashAttention or FlashInfer with its matching ModelOpt adapter on - each Attention module right after model loading, before any forward pass. - """ +class _ModelOptAttentionWorker(BaseWorker): + quantize_attention = False def load_model(self, *args, **kwargs) -> None: - """Load model, then replace attention impl with sparse variant.""" super().load_model(*args, **kwargs) - _replace_attention_impl(self) + _install_attention(self, quantize=self.quantize_attention) + + +class SparseAttnWorker(_ModelOptAttentionWorker): + """Install checkpoint-driven ModelOpt sparse attention after model load.""" + + +class QuantSparseAttnWorker(_ModelOptAttentionWorker): + """Install fixed NVFP4 attention plus optional checkpoint sparsity.""" + + quantize_attention = True + + def determine_available_memory(self) -> int: + api = _quant_api() + with api.torch.inference_mode(), api.plugin.disable_compilation(_unwrapped_model(self)): + return BaseWorker.determine_available_memory(self) diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 1144e21b982..7695f81c6cd 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -15,15 +15,18 @@ """Tests for sparse attention vLLM worker compatibility helpers.""" +import builtins import importlib.util import math -from contextlib import nullcontext +import sys +from contextlib import contextmanager, nullcontext from itertools import accumulate from pathlib import Path -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace import pytest import torch +import vllm from torch import nn from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl from vllm.v1.attention.backends.flashinfer import ( @@ -31,6 +34,7 @@ FlashInferImpl, FlashInferMetadataBuilder, ) +from vllm.v1.worker.gpu_worker import Worker as BaseWorker from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as vllm_plugin from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( @@ -42,6 +46,179 @@ select_sparse_impl_cls, ) +_WORKER_PATH = Path(__file__).parents[5] / "examples/vllm_serve/sparse_attn_worker.py" + + +def _load_worker_module(name="sparse_attn_worker_test"): + spec = importlib.util.spec_from_file_location(name, _WORKER_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _legacy_quant_module(installer): + module = ModuleType("quant_sparse_attn_worker") + setattr(module, "_install_quant_sparse_attn", installer) + return module + + +def test_shared_worker_import_does_not_resolve_quant_only_apis(monkeypatch): + forbidden = { + "vllm.config.compilation", + "vllm.v1.attention.backend", + "modelopt.torch.quantization.conversion", + "modelopt.torch.quantization.nn", + "modelopt.torch.quantization.plugins.vllm", + } + real_import = builtins.__import__ + + def guarded_import(name, *args, **kwargs): + fromlist = kwargs.get("fromlist", args[2] if len(args) > 2 else ()) + requested = {name, *(f"{name}.{item}" for item in fromlist or ())} + if blocked := requested & forbidden: + raise AssertionError( + f"quant-only import during sparse module load: {sorted(blocked)[0]}" + ) + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(vllm, "__version__", "0.9.0") + monkeypatch.setattr(builtins, "__import__", guarded_import) + worker_module = _load_worker_module("sparse_attn_worker_import_test") + + assert worker_module.__all__ == ["SparseAttnWorker", "QuantSparseAttnWorker"] + + +@pytest.mark.parametrize( + ("class_name", "quantize"), + [("SparseAttnWorker", False), ("QuantSparseAttnWorker", True)], +) +def test_public_workers_install_after_base_load(monkeypatch, class_name, quantize): + worker_module = _load_worker_module(f"worker_order_{class_name}") + events = [] + monkeypatch.setattr(BaseWorker, "load_model", lambda *_a, **_k: events.append("base")) + monkeypatch.setattr( + worker_module, + "_install_attention", + lambda _worker, *, quantize: events.append(("install", quantize)), + ) + + instance = object.__new__(getattr(worker_module, class_name)) + assert instance.load_model() is None + assert events == ["base", ("install", quantize)] + + +def test_shared_worker_quant_install_rejects_old_vllm_before_delegation(monkeypatch): + worker_module = _load_worker_module("quant_install_version_test") + worker = object() + delegated = [] + monkeypatch.setitem( + sys.modules, "quant_sparse_attn_worker", _legacy_quant_module(delegated.append) + ) + monkeypatch.setattr(vllm, "__version__", "0.9.0") + worker_module._load_quant_api.cache_clear() + + try: + with pytest.raises(RuntimeError, match=r"vLLM >= 0\.14\.0"): + worker_module._install_attention(worker, quantize=True) + finally: + worker_module._load_quant_api.cache_clear() + + assert delegated == [] + + +def test_shared_worker_quant_install_validates_before_delegation(monkeypatch): + worker_module = _load_worker_module("quant_install_order_test") + worker = object() + events = [] + delegated = [] + real_import = builtins.__import__ + + def install(actual_worker): + events.append("delegate") + delegated.append(actual_worker) + + def recording_import(name, *args, **kwargs): + if name == "quant_sparse_attn_worker": + events.append("import") + return real_import(name, *args, **kwargs) + + monkeypatch.setitem(sys.modules, "quant_sparse_attn_worker", _legacy_quant_module(install)) + monkeypatch.setattr(builtins, "__import__", recording_import) + monkeypatch.setattr(worker_module, "_quant_api", lambda: events.append("validate")) + + worker_module._install_attention(worker, quantize=True) + + assert events == ["validate", "import", "delegate"] + assert len(delegated) == 1 and delegated[0] is worker + + +def test_shared_worker_sparse_install_does_not_import_legacy_quant_worker(monkeypatch): + worker_module = _load_worker_module("sparse_install_test") + worker = object() + replaced = [] + real_import = builtins.__import__ + + def guarded_import(name, *args, **kwargs): + if name == "quant_sparse_attn_worker": + raise AssertionError("legacy quant worker imported for sparse attention") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", guarded_import) + monkeypatch.setattr(worker_module, "_replace_attention_impl", replaced.append) + + worker_module._install_attention(worker, quantize=False) + + assert len(replaced) == 1 and replaced[0] is worker + + +def test_sparse_worker_keeps_base_memory_profile(): + worker_module = _load_worker_module("sparse_profile_test") + + assert ( + worker_module.SparseAttnWorker.determine_available_memory + is BaseWorker.determine_available_memory + ) + + +def test_shared_worker_quant_memory_profile_uses_both_contexts(monkeypatch): + worker_module = _load_worker_module("quant_profile_test") + events = [] + model = object() + + @contextmanager + def recorded_context(name, value=None): + events.append(("enter", name, value)) + try: + yield + finally: + events.append(("exit", name, value)) + + api = SimpleNamespace( + torch=SimpleNamespace(inference_mode=lambda: recorded_context("inference")), + plugin=SimpleNamespace( + disable_compilation=lambda actual_model: recorded_context("compilation", actual_model) + ), + ) + monkeypatch.setattr(worker_module, "_quant_api", lambda: api) + + instance = object.__new__(worker_module.QuantSparseAttnWorker) + instance.model_runner = SimpleNamespace(model=SimpleNamespace(unwrap=lambda: model)) + + def profile(actual_worker): + events.append(("profile", actual_worker)) + return 73 + + monkeypatch.setattr(BaseWorker, "determine_available_memory", profile) + + assert instance.determine_available_memory() == 73 + assert events == [ + ("enter", "inference", None), + ("enter", "compilation", model), + ("profile", instance), + ("exit", "compilation", model), + ("exit", "inference", None), + ] + def _make_old_impl(): """Create a vLLM FlashAttention impl with initialized runtime state.""" @@ -420,10 +597,7 @@ def fake_attention(query, **kwargs): def test_sparse_worker_rejects_unsupported_backend_before_mutation(monkeypatch): - worker_path = Path(__file__).parents[5] / "examples/vllm_serve/sparse_attn_worker.py" - spec = importlib.util.spec_from_file_location("sparse_attn_worker_test", worker_path) - worker_module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(worker_module) + worker_module = _load_worker_module() good = object.__new__(worker_module.VLLMAttention) nn.Module.__init__(good) good.impl = object.__new__(FlashAttentionImpl) From a5980c183229c1aa04b1a90173b74ccfc00beb63 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 6 Jul 2026 12:58:54 -0700 Subject: [PATCH 13/28] Consolidate quant and sparse attention workers Signed-off-by: Kai Xu --- .../vllm_serve/quant_sparse_attn_worker.py | 246 --------------- examples/vllm_serve/sparse_attn_worker.py | 296 ++++++++++++++---- .../test_quant_sparse_attn_worker.py | 98 ++++-- .../test_sparse_attn_worker.py | 206 +++++------- 4 files changed, 379 insertions(+), 467 deletions(-) delete mode 100644 examples/vllm_serve/quant_sparse_attn_worker.py diff --git a/examples/vllm_serve/quant_sparse_attn_worker.py b/examples/vllm_serve/quant_sparse_attn_worker.py deleted file mode 100644 index cd9f9b4d22c..00000000000 --- a/examples/vllm_serve/quant_sparse_attn_worker.py +++ /dev/null @@ -1,246 +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. - -"""Fixed NVFP4 Q/K/P/V worker for FlashAttention and FlashInfer in vLLM.""" - -import torch -import vllm -from packaging import version - -try: - from vllm.config.compilation import CUDAGraphMode - from vllm.v1.attention.backend import AttentionType - from vllm.v1.worker.gpu_worker import Worker as BaseWorker - - from modelopt.torch.quantization.conversion import set_quantizer_by_cfg - from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer - from modelopt.torch.quantization.plugins.vllm import ( - _ATTENTION_TYPES, - _get_device_dtype, - _set_vllm_attention_kv_default_amax, - disable_compilation, - vllm_attention, - ) - from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_config import ( - load_from_checkpoint_metadata, - match_sparse_config, - ) - from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( - _build_sparse_kw, - _clone_sparse_impl, - _p_qdq_from_layer, - _v_qdq_from_layer, - select_sparse_impl_cls, - ) -except ImportError as err: - if version.parse(vllm.__version__) < version.parse("0.14.0"): - raise RuntimeError("The compact NVFP4 attention worker requires vLLM >= 0.14.0") from err - raise - -if version.parse(vllm.__version__) < version.parse("0.14.0"): - raise RuntimeError("The compact NVFP4 attention worker requires vLLM >= 0.14.0") - -__all__ = ["QuantSparseAttnWorker"] - -VLLMAttention = vllm_attention.Attention -_NVFP4_CFG = { - "num_bits": (2, 1), - "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, -} -_BMM_CFG = [ - {"quantizer_name": "*_bmm_quantizer", "enable": False}, - *( - {"quantizer_name": f"*{name}_bmm_quantizer", "cfg": _NVFP4_CFG, "enable": True} - for name in ("q", "k", "p", "v") - ), -] - - -def _unwrapped_model(worker): - model = worker.model_runner.model - return model.unwrap() if hasattr(model, "unwrap") else model - - -def _global_errors(worker) -> list[str]: - config = worker.model_runner.vllm_config - parallel = config.parallel_config - cache = config.cache_config - model_config = config.model_config - errors = [] - if getattr(parallel, "decode_context_parallel_size", 1) != 1: - errors.append("decode_context_parallel_size must be 1") - if getattr(parallel, "enable_dbo", False) or getattr(parallel, "use_ubatching", False): - errors.append("DBO/ubatching is unsupported") - if getattr(cache, "enable_prefix_caching", False): - errors.append("prefix caching is unsupported") - if getattr(config, "kv_transfer_config", None) is not None: - errors.append("KV transfer is unsupported") - if getattr(config, "speculative_config", None) is not None: - errors.append("speculative decoding is unsupported") - if config.compilation_config.cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL: - errors.append("FULL mixed-batch cudagraph mode is unsupported") - if getattr(model_config, "dtype", None) not in (torch.float16, torch.bfloat16): - errors.append("resolved model/KV-cache dtype must be fp16 or bf16") - cache_dtype = getattr(cache, "cache_dtype", "auto") - if str(cache_dtype) not in {"auto", "bfloat16", "float16", "torch.bfloat16", "torch.float16"}: - errors.append(f"resolved KV-cache dtype {cache_dtype!r} must be fp16 or bf16") - return errors - - -def _layer_errors(module) -> list[str]: - impl = getattr(module, "impl", None) - errors = [] - if type(module) is not VLLMAttention: - errors.append(f"layout {type(module).__name__} is not regular decoder self-attention") - if getattr(module, "attn_type", None) != AttentionType.DECODER: - errors.append("attn_type must be DECODER") - head_size = getattr(module, "head_size", None) - if not isinstance(head_size, int) or head_size % 16: - errors.append(f"head_size={head_size!r} must be a multiple of 16") - head_size_v = getattr(module, "head_size_v", head_size) - if head_size_v != head_size: - errors.append(f"head_size_v={head_size_v!r} must equal head_size={head_size!r}") - if getattr(module, "sliding_window", None) is not None: - errors.append("sliding_window is unsupported") - if getattr(module, "kv_sharing_target_layer_name", None) is not None: - errors.append("cross-layer KV sharing is unsupported") - if str(getattr(module, "kv_cache_dtype", "")).startswith("fp8"): - errors.append("FP8 KV cache is unsupported") - if getattr(impl, "alibi_slopes", None) is not None: - errors.append("ALiBi is unsupported") - if getattr(impl, "logits_soft_cap", None): - errors.append("logits soft cap is unsupported") - if getattr(impl, "sinks", None) is not None or getattr(impl, "has_sinks", False): - errors.append("attention sinks are unsupported") - return errors - - -def _sparse_graph_error(sparse_kw: dict, mode: CUDAGraphMode) -> str | None: - """Reject decode calibration whose live length would be frozen by a full graph.""" - params = sparse_kw.get("threshold_scale_factor") - if ( - mode.decode_mode() == CUDAGraphMode.FULL - and isinstance(params, dict) - and isinstance(params.get("decode"), dict) - ): - return "calibrated decode skip-softmax requires a non-FULL CUDA graph mode" - return None - - -def _validated_attention_plans(worker): - """Validate every attention layout without mutation, then return regular-layer tuples.""" - model = _unwrapped_model(worker) - model_config = worker.model_runner.model_config - model_dtype = getattr(model_config, "dtype", None) - detected = load_from_checkpoint_metadata(model_config.hf_config) - sparse_cfg = detected[0] if detected is not None else None - errors = _global_errors(worker) - vllm_config = getattr(worker.model_runner, "vllm_config", None) - compilation_config = getattr(vllm_config, "compilation_config", None) - cudagraph_mode = getattr(compilation_config, "cudagraph_mode", CUDAGraphMode.NONE) - plans = [] - attention_count = 0 - for name, module in model.named_modules(): - if not isinstance(module, _ATTENTION_TYPES): - continue - attention_count += 1 - reasons = _layer_errors(module) - new_impl_cls = select_sparse_impl_cls(module.impl) - if new_impl_cls is None: - reasons.append( - f"backend {type(module.impl).__name__} is not supported; " - "expected FlashAttentionImpl or FlashInferImpl" - ) - device, dtype = _get_device_dtype(module) - # The install runs in load_model (before KV-cache allocation), so the - # buffer scan in _get_device_dtype can pick up fp32 scale buffers - # (_k_scale/_v_scale/...) on the attention module. Prefer the - # authoritative model compute dtype when it is fp16/bf16. - if model_dtype in (torch.float16, torch.bfloat16): - dtype = model_dtype - if device is None or dtype is None: - reasons.append("device/dtype could not be resolved") - elif dtype not in (torch.float16, torch.bfloat16): - reasons.append(f"resolved dtype {dtype} must be fp16 or bf16") - if reasons: - errors.extend(f"{name or ''}: {reason}" for reason in reasons) - continue - layer_cfg = match_sparse_config(name, sparse_cfg) if sparse_cfg is not None else None - sparse_kw = ( - _build_sparse_kw(layer_cfg) - if layer_cfg is not None and layer_cfg.get("enable", True) - else {} - ) - graph_error = _sparse_graph_error(sparse_kw, cudagraph_mode) - if graph_error is not None: - errors.append(f"{name or ''}: {graph_error}") - continue - plans.append((name, module, new_impl_cls, sparse_kw, device, dtype)) - if attention_count == 0: - errors.append("no regular attention layers were found") - if errors: - raise NotImplementedError( - "Unsupported ModelOpt attention plan:\n - " + "\n - ".join(errors) - ) - return plans - - -def _install_quant_sparse_attn(worker) -> None: - plans = _validated_attention_plans(worker) - installed = {} - for _name, module, new_impl_cls, sparse_kw, device, dtype in plans: - module.device, module.dtype = device, dtype - QuantModuleRegistry.convert(module) - module.p_bmm_quantizer = TensorQuantizer() - set_quantizer_by_cfg(module, _BMM_CFG) - _set_vllm_attention_kv_default_amax(module, device) - new_impl = _clone_sparse_impl(module.impl, new_impl_cls) - new_impl.sparse_kw = sparse_kw - p_qdq, p_qdq_amax = _p_qdq_from_layer(module) - v_qdq, v_qdq_amax = _v_qdq_from_layer(module) - new_impl.quant_kw = { - "p_qdq": p_qdq, - "p_qdq_amax": p_qdq_amax, - "v_qdq": v_qdq, - "v_qdq_amax": v_qdq_amax, - } - module.impl = new_impl - module._query_quant_in_kernel = True - module._value_quant_in_kernel = True - impl_name = type(new_impl).__name__ - installed[impl_name] = installed.get(impl_name, 0) + 1 - worker.model_runner.cascade_attn_enabled = False - print(f"[ModelOpt] Installed NVFP4 quant+sparse attention on {len(plans)} layers: {installed}") - - -class QuantSparseAttnWorker(BaseWorker): - """Install the fixed NVFP4 attention recipe right after model load. - - The impl swap must happen before the first forward pass -- including the - ``determine_available_memory`` profiling run -- so that profiling exercises - the ModelOpt Triton path that real serving uses. Installing here also patches - FlashInfer metadata before cache and metadata-builder initialization. This - mirrors ``SparseAttnWorker``, which installs in ``load_model`` for the same - reason. - """ - - def load_model(self, *args, **kwargs) -> None: - super().load_model(*args, **kwargs) - _install_quant_sparse_attn(self) - - @torch.inference_mode() - def determine_available_memory(self) -> int: - with disable_compilation(_unwrapped_model(self)): - return BaseWorker.determine_available_memory(self) diff --git a/examples/vllm_serve/sparse_attn_worker.py b/examples/vllm_serve/sparse_attn_worker.py index 14798f7579c..4da178d56de 100644 --- a/examples/vllm_serve/sparse_attn_worker.py +++ b/examples/vllm_serve/sparse_attn_worker.py @@ -13,25 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Custom vLLM workers for ModelOpt attention policies. - -``SparseAttnWorker`` replaces the native FlashAttention or FlashInfer impl with -the matching checkpoint-driven ModelOpt sparse adapter. ``QuantSparseAttnWorker`` -installs fixed NVFP4 attention plus optional checkpoint-driven sparsity. Both -policies install their attention implementation after model loading. - -Configuration flows exclusively through the loaded checkpoint's -``sparse_attention_config`` block (written by ModelOpt's HF export). If the -checkpoint has no such block, the sparse-only policy logs a message and passes -through unchanged. - -Usage: - python vllm_serve_sparse_attn.py -""" +"""Custom vLLM workers for checkpoint-driven sparse and fixed-NVFP4 attention.""" import importlib +import os from functools import cache from types import SimpleNamespace +from typing import NamedTuple try: _has_legacy_attention_layer = importlib.util.find_spec("vllm.attention.layer") is not None @@ -52,11 +40,34 @@ from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( _build_sparse_kw, _clone_sparse_impl, + _p_qdq_from_layer, + _v_qdq_from_layer, select_sparse_impl_cls, ) __all__ = ["SparseAttnWorker", "QuantSparseAttnWorker"] # noqa: RUF022 +_NVFP4_CFG = { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, +} +_BMM_CFG = [ + {"quantizer_name": "*_bmm_quantizer", "enable": False}, + *( + {"quantizer_name": f"*{name}_bmm_quantizer", "cfg": _NVFP4_CFG, "enable": True} + for name in ("q", "k", "p", "v") + ), +] + + +class _AttentionPlan(NamedTuple): + name: str + module: object + new_impl: object + sparse_kw: dict + device: object | None + dtype: object | None + def _unwrapped_model(worker): model = worker.model_runner.model @@ -95,83 +106,236 @@ def _quant_api(): return _load_quant_api(vllm.__version__) -def _replace_attention_impl(worker): - """Install the backend-matched ModelOpt sparse impl on attention layers. +def _cudagraph_mode(worker, api): + config = getattr(worker.model_runner, "vllm_config", None) + compilation = getattr(config, "compilation_config", None) + mode = getattr(compilation, "cudagraph_mode", None) + return mode if mode is not None else api.compilation.CUDAGraphMode.NONE - The sole configuration source is the checkpoint's ``sparse_attention_config`` - metadata. No-op if the checkpoint has no such block. - """ - hf_config = getattr(worker.model_runner.model_config, "hf_config", None) - detected = load_from_checkpoint_metadata(hf_config) - if detected is None: + +def _global_errors(worker, api=None) -> list[str]: + api = api or _quant_api() + config = worker.model_runner.vllm_config + parallel, cache, model_config = config.parallel_config, config.cache_config, config.model_config + errors = [] + if getattr(parallel, "decode_context_parallel_size", 1) != 1: + errors.append("decode_context_parallel_size must be 1") + if getattr(parallel, "enable_dbo", False) or getattr(parallel, "use_ubatching", False): + errors.append("DBO/ubatching is unsupported") + if getattr(cache, "enable_prefix_caching", False): + errors.append("prefix caching is unsupported") + if getattr(config, "kv_transfer_config", None) is not None: + errors.append("KV transfer is unsupported") + if getattr(config, "speculative_config", None) is not None: + errors.append("speculative decoding is unsupported") + if _cudagraph_mode(worker, api).mixed_mode() == api.compilation.CUDAGraphMode.FULL: + errors.append("FULL mixed-batch cudagraph mode is unsupported") + if getattr(model_config, "dtype", None) not in (api.torch.float16, api.torch.bfloat16): + errors.append("resolved model/KV-cache dtype must be fp16 or bf16") + cache_dtype = getattr(cache, "cache_dtype", "auto") + if str(cache_dtype) not in {"auto", "bfloat16", "float16", "torch.bfloat16", "torch.float16"}: + errors.append(f"resolved KV-cache dtype {cache_dtype!r} must be fp16 or bf16") + return errors + + +def _quant_layer_errors(module, api) -> list[str]: + impl = getattr(module, "impl", None) + errors = [] + if type(module) is not api.plugin.vllm_attention.Attention: + errors.append(f"layout {type(module).__name__} is not regular decoder self-attention") + if getattr(module, "attn_type", None) != api.backend.AttentionType.DECODER: + errors.append("attn_type must be DECODER") + head_size = getattr(module, "head_size", None) + if not isinstance(head_size, int) or head_size % 16: + errors.append(f"head_size={head_size!r} must be a multiple of 16") + head_size_v = getattr(module, "head_size_v", head_size) + if head_size_v != head_size: + errors.append(f"head_size_v={head_size_v!r} must equal head_size={head_size!r}") + if getattr(module, "sliding_window", None) is not None: + errors.append("sliding_window is unsupported") + if getattr(module, "kv_sharing_target_layer_name", None) is not None: + errors.append("cross-layer KV sharing is unsupported") + if str(getattr(module, "kv_cache_dtype", "")).startswith("fp8"): + errors.append("FP8 KV cache is unsupported") + if getattr(impl, "alibi_slopes", None) is not None: + errors.append("ALiBi is unsupported") + if getattr(impl, "logits_soft_cap", None): + errors.append("logits soft cap is unsupported") + if getattr(impl, "sinks", None) is not None or getattr(impl, "has_sinks", False): + errors.append("attention sinks are unsupported") + return errors + + +def _validated_device_dtype(module, model_config, api): + device, dtype = api.plugin._get_device_dtype(module) + model_dtype = getattr(model_config, "dtype", None) + if model_dtype in (api.torch.float16, api.torch.bfloat16): + dtype = model_dtype + if device is None or dtype is None: + return device, dtype, "device/dtype could not be resolved" + if dtype not in (api.torch.float16, api.torch.bfloat16): + return device, dtype, f"resolved dtype {dtype} must be fp16 or bf16" + return device, dtype, None + + +def _sparse_graph_error(sparse_kw: dict, mode, api=None) -> str | None: + """Reject decode calibration whose live length would be frozen by a full graph.""" + api = api or _quant_api() + params = sparse_kw.get("threshold_scale_factor") + if ( + mode.decode_mode() == api.compilation.CUDAGraphMode.FULL + and isinstance(params, dict) + and isinstance(params.get("decode"), dict) + ): + return "calibrated decode skip-softmax requires a non-FULL CUDA graph mode" + return None + + +def _validated_attention_plans(worker, *, quantize: bool): + """Validate and clone every selected attention adapter before mutating layers.""" + api = _quant_api() if quantize else None + model = _unwrapped_model(worker) + model_config = worker.model_runner.model_config + detected = load_from_checkpoint_metadata(getattr(model_config, "hf_config", None)) + if not quantize and detected is None: print( "[ModelOpt] No sparse_attention_config found in the checkpoint; " "skipping sparse attention. Run examples/llm_sparsity/" "attention_sparsity/hf_sa.py to calibrate and export a checkpoint " "with the config embedded." ) - return - cfg, preset_name = detected - print(f"[ModelOpt] Sparse attention config: algo -> {preset_name}") - - model = _unwrapped_model(worker) + return None + sparse_cfg = detected[0] if detected is not None else None + if quantize: + assert api is not None + errors = _global_errors(worker, api) + mode = _cudagraph_mode(worker, api) + attention_types = api.plugin._ATTENTION_TYPES + else: + assert detected is not None + print(f"[ModelOpt] Sparse attention config: algo -> {detected[1]}") + errors, mode, attention_types = [], None, (VLLMAttention,) plans = [] - errors = [] + attention_count = 0 for name, module in model.named_modules(): - if not isinstance(module, VLLMAttention): - continue - - layer_cfg = match_sparse_config(name, cfg) - if layer_cfg is None or not layer_cfg.get("enable", True): - continue - - sparse_kw = _build_sparse_kw(layer_cfg) - if not sparse_kw: - # Keep vLLM's original impl when the exported layer config does not - # enable any sparse feature. - continue - new_impl_cls = select_sparse_impl_cls(module.impl) - if new_impl_cls is None: - errors.append(f"{name or ''}: unsupported backend {type(module.impl).__name__}") + if not isinstance(module, attention_types): continue + if quantize: + assert api is not None + attention_count += 1 + reasons = _quant_layer_errors(module, api) + device, dtype, dtype_error = _validated_device_dtype(module, model_config, api) + if dtype_error: + reasons.append(dtype_error) + layer_cfg = match_sparse_config(name, sparse_cfg) if sparse_cfg is not None else None + sparse_kw = ( + _build_sparse_kw(layer_cfg) + if layer_cfg is not None and layer_cfg.get("enable", True) + else {} + ) + if graph_error := _sparse_graph_error(sparse_kw, mode, api): + reasons.append(graph_error) + else: + assert sparse_cfg is not None + layer_cfg = match_sparse_config(name, sparse_cfg) + if layer_cfg is None or not layer_cfg.get("enable", True): + continue + sparse_kw = _build_sparse_kw(layer_cfg) + if not sparse_kw: + continue + reasons, device, dtype = [], None, None + + new_impl = None try: - new_impl = _clone_sparse_impl(module.impl, new_impl_cls) + new_impl_cls = select_sparse_impl_cls(module.impl) + if new_impl_cls is None: + backend = type(module.impl).__name__ + message = ( + f"backend {backend} is not supported; expected FlashAttentionImpl or FlashInferImpl" + if quantize + else f"unsupported backend {backend}" + ) + reasons.append(message) + else: + new_impl = _clone_sparse_impl(module.impl, new_impl_cls) except (NotImplementedError, TypeError) as err: - errors.append(f"{name or ''}: {err}") - continue - plans.append((module, new_impl, sparse_kw)) - + reasons.append(str(err)) + layer_name = name or "" + if reasons: + errors.extend(f"{layer_name}: {reason}" for reason in reasons) + else: + plans.append(_AttentionPlan(name, module, new_impl, sparse_kw, device, dtype)) + + if quantize and attention_count == 0: + errors.append("no regular attention layers were found") if errors: + policy = "attention" if quantize else "sparse attention" raise NotImplementedError( - "Unsupported ModelOpt sparse attention plan:\n - " + "\n - ".join(errors) + f"Unsupported ModelOpt {policy} plan:\n - " + "\n - ".join(errors) ) + return tuple(plans) + +def _install_sparse_plans(plans) -> None: installed = {} - for module, new_impl, sparse_kw in plans: - new_impl.sparse_kw = sparse_kw - module.impl = new_impl - impl_name = type(new_impl).__name__ + for plan in plans: + plan.new_impl.sparse_kw = plan.sparse_kw + plan.module.impl = plan.new_impl + impl_name = type(plan.new_impl).__name__ installed[impl_name] = installed.get(impl_name, 0) + 1 print( f"[ModelOpt] Sparse attention: replaced impl on {len(plans)} attention layers: {installed}" ) -def _install_attention(worker, *, quantize: bool) -> None: - if quantize: - _quant_api() - # Keep the compatibility delegation lazy until both policies share one planner. - from quant_sparse_attn_worker import _install_quant_sparse_attn - - _install_quant_sparse_attn(worker) - return - _replace_attention_impl(worker) +def _install_quant_plans(worker, plans) -> None: + api = _quant_api() + quant_off = os.environ.get("MODELOPT_ATTN_QUANT_OFF") == "1" + installed = {} + for plan in plans: + module = plan.module + module.device, module.dtype = plan.device, plan.dtype + api.nn.QuantModuleRegistry.convert(module) + module.p_bmm_quantizer = api.nn.TensorQuantizer() + api.conversion.set_quantizer_by_cfg(module, _BMM_CFG) + if quant_off: + # Isolation knob: keep the ModelOpt kernel fixed while disabling all + # Q/K/P/V NVFP4 transforms, so quant-on minus quant-off isolates the + # fakequant loss without triggering the native dense fallback. + for name in ("q", "k", "p", "v"): + quantizer = getattr(module, f"{name}_bmm_quantizer", None) + if quantizer is not None: + quantizer.disable() + api.plugin._set_vllm_attention_kv_default_amax(module, plan.device) + plan.new_impl.sparse_kw = plan.sparse_kw + p_qdq, p_qdq_amax = _p_qdq_from_layer(module) + v_qdq, v_qdq_amax = _v_qdq_from_layer(module) + plan.new_impl.quant_kw = { + "p_qdq": p_qdq, + "p_qdq_amax": p_qdq_amax, + "v_qdq": v_qdq, + "v_qdq_amax": v_qdq_amax, + } + module.impl = plan.new_impl + module._query_quant_in_kernel = not quant_off + module._value_quant_in_kernel = not quant_off + if quant_off: + module._modelopt_force_kernel = True + impl_name = type(plan.new_impl).__name__ + installed[impl_name] = installed.get(impl_name, 0) + 1 + worker.model_runner.cascade_attn_enabled = False + print(f"[ModelOpt] Installed NVFP4 quant+sparse attention on {len(plans)} layers: {installed}") -# --------------------------------------------------------------------------- -# Workers -# --------------------------------------------------------------------------- +def _install_attention(worker, *, quantize: bool) -> None: + plans = _validated_attention_plans(worker, quantize=quantize) + if plans is None: + return + if quantize: + _install_quant_plans(worker, plans) + else: + _install_sparse_plans(plans) class _ModelOptAttentionWorker(BaseWorker): diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py index 68ab3b6bce0..25c7a7623c3 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py @@ -36,11 +36,13 @@ get_flashinfer_sparse_impl_cls, ) -_WORKER_PATH = Path(__file__).parents[5] / "examples/vllm_serve/quant_sparse_attn_worker.py" +_WORKER_PATH = Path(__file__).parents[5] / "examples/vllm_serve/sparse_attn_worker.py" def _load_worker_module(): - spec = importlib.util.spec_from_file_location("quant_sparse_attn_worker", _WORKER_PATH) + spec = importlib.util.spec_from_file_location( + "shared_attention_worker_quant_test", _WORKER_PATH + ) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module @@ -49,11 +51,15 @@ def _load_worker_module(): worker_module = _load_worker_module() -def test_worker_rejects_vllm_before_v1_attention_api(monkeypatch): +def test_quant_policy_rejects_old_vllm(monkeypatch): monkeypatch.setattr(vllm, "__version__", "0.9.0") + worker_module._load_quant_api.cache_clear() - with pytest.raises(RuntimeError, match=r"vLLM >= 0\.14\.0"): - _load_worker_module() + try: + with pytest.raises(RuntimeError, match=r"vLLM >= 0\.14\.0"): + worker_module._install_attention(SimpleNamespace(), quantize=True) + finally: + worker_module._load_quant_api.cache_clear() def _attention(attn_type=AttentionType.DECODER, impl_cls=FlashAttentionImpl): @@ -87,7 +93,7 @@ def _patch_conversion(monkeypatch): lambda: quant_plugin.ParallelState(data_parallel_group=None), ) monkeypatch.setattr(worker_module, "load_from_checkpoint_metadata", lambda _: None) - monkeypatch.setattr(worker_module, "_global_errors", lambda _: []) + monkeypatch.setattr(worker_module, "_global_errors", lambda _worker, _api=None: []) @pytest.mark.parametrize("impl_cls", [FlashAttentionImpl, FlashInferImpl]) @@ -98,7 +104,7 @@ def test_install_converts_only_attention_and_configures_fixed_recipe(monkeypatch model = nn.ModuleDict({"attn": attention, "linear": linear}) state = _worker(model) - worker_module._install_quant_sparse_attn(state) + worker_module._install_attention(state, quantize=True) converted = model["attn"] assert isinstance(converted, quant_plugin._QuantVLLMAttention) @@ -128,6 +134,29 @@ def test_install_converts_only_attention_and_configures_fixed_recipe(monkeypatch assert state.model_runner.cascade_attn_enabled is False +def test_quant_off_disables_recipe_but_keeps_modelopt_adapter(monkeypatch): + _patch_conversion(monkeypatch) + monkeypatch.setenv("MODELOPT_ATTN_QUANT_OFF", "1") + attention = _attention() + state = _worker(nn.ModuleDict({"attn": attention})) + + worker_module._install_attention(state, quantize=True) + + converted = state.model_runner.model["attn"] + for name in ("q_bmm_quantizer", "k_bmm_quantizer", "p_bmm_quantizer", "v_bmm_quantizer"): + assert not getattr(converted, name).is_enabled + assert converted._query_quant_in_kernel is False + assert converted._value_quant_in_kernel is False + assert converted._modelopt_force_kernel is True + assert type(converted.impl) is ModelOptSparseAttentionImpl + assert converted.impl.quant_kw == { + "p_qdq": None, + "p_qdq_amax": 1.0, + "v_qdq": None, + "v_qdq_amax": None, + } + + def test_validation_of_all_layouts_precedes_mutation(monkeypatch): _patch_conversion(monkeypatch) good, bad = _attention(), _attention(AttentionType.ENCODER) @@ -137,41 +166,50 @@ def test_validation_of_all_layouts_precedes_mutation(monkeypatch): original_impl = good.impl with pytest.raises(NotImplementedError) as exc: - worker_module._install_quant_sparse_attn(_worker(model)) + worker_module._install_attention(_worker(model), quantize=True) assert all(needle in str(exc.value) for needle in ("bad: attn_type", "head_size_v", "float32")) assert model["good"] is good and good.impl is original_impl assert not isinstance(good, quant_plugin._QuantVLLMAttention) -def test_worker_installs_after_base_load(monkeypatch): +def test_quant_memory_profile_uses_inference_mode_and_disables_compilation(monkeypatch): events = [] - monkeypatch.setattr( - worker_module, "_install_quant_sparse_attn", lambda _: events.append("install") - ) - monkeypatch.setattr(BaseWorker, "load_model", lambda *_args, **_kwargs: events.append("base")) - instance = object.__new__(worker_module.QuantSparseAttnWorker) - assert instance.load_model() is None - assert events == ["base", "install"] - - -def test_memory_profile_disables_compilation(monkeypatch): - events, model = [], object() + model = object() @contextmanager - def disabled(actual): - events.append(("enter", actual)) - yield - events.append(("exit", actual)) - - monkeypatch.setattr(worker_module, "disable_compilation", disabled, raising=False) - monkeypatch.setattr( - BaseWorker, "determine_available_memory", lambda _: events.append(("profile", model)) or 7 + def recorded_context(name, value=None): + events.append(("enter", name, value)) + try: + yield + finally: + events.append(("exit", name, value)) + + api = SimpleNamespace( + torch=SimpleNamespace(inference_mode=lambda: recorded_context("inference")), + plugin=SimpleNamespace( + disable_compilation=lambda actual_model: recorded_context("compilation", actual_model) + ), ) + monkeypatch.setattr(worker_module, "_quant_api", lambda: api) + instance = object.__new__(worker_module.QuantSparseAttnWorker) instance.model_runner = SimpleNamespace(model=SimpleNamespace(unwrap=lambda: model)) - assert instance.determine_available_memory() == 7 - assert events == [("enter", model), ("profile", model), ("exit", model)] + + def profile(actual_worker): + events.append(("profile", actual_worker)) + return 73 + + monkeypatch.setattr(BaseWorker, "determine_available_memory", profile) + + assert instance.determine_available_memory() == 73 + assert events == [ + ("enter", "inference", None), + ("enter", "compilation", model), + ("profile", instance), + ("exit", "compilation", model), + ("exit", "inference", None), + ] @pytest.mark.parametrize( diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 7695f81c6cd..c48a0e2d028 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -18,11 +18,10 @@ import builtins import importlib.util import math -import sys -from contextlib import contextmanager, nullcontext +from contextlib import nullcontext from itertools import accumulate from pathlib import Path -from types import ModuleType, SimpleNamespace +from types import SimpleNamespace import pytest import torch @@ -56,12 +55,23 @@ def _load_worker_module(name="sparse_attn_worker_test"): return module -def _legacy_quant_module(installer): - module = ModuleType("quant_sparse_attn_worker") - setattr(module, "_install_quant_sparse_attn", installer) +def _bare_attention(worker_module, impl_cls=FlashAttentionImpl): + module = object.__new__(worker_module.VLLMAttention) + nn.Module.__init__(module) + module.impl = object.__new__(impl_cls) + module.impl.sinks = None return module +def _sparse_worker_state(model): + return SimpleNamespace( + model_runner=SimpleNamespace( + model=model, + model_config=SimpleNamespace(hf_config=SimpleNamespace()), + ) + ) + + def test_shared_worker_import_does_not_resolve_quant_only_apis(monkeypatch): forbidden = { "vllm.config.compilation", @@ -107,70 +117,6 @@ def test_public_workers_install_after_base_load(monkeypatch, class_name, quantiz assert events == ["base", ("install", quantize)] -def test_shared_worker_quant_install_rejects_old_vllm_before_delegation(monkeypatch): - worker_module = _load_worker_module("quant_install_version_test") - worker = object() - delegated = [] - monkeypatch.setitem( - sys.modules, "quant_sparse_attn_worker", _legacy_quant_module(delegated.append) - ) - monkeypatch.setattr(vllm, "__version__", "0.9.0") - worker_module._load_quant_api.cache_clear() - - try: - with pytest.raises(RuntimeError, match=r"vLLM >= 0\.14\.0"): - worker_module._install_attention(worker, quantize=True) - finally: - worker_module._load_quant_api.cache_clear() - - assert delegated == [] - - -def test_shared_worker_quant_install_validates_before_delegation(monkeypatch): - worker_module = _load_worker_module("quant_install_order_test") - worker = object() - events = [] - delegated = [] - real_import = builtins.__import__ - - def install(actual_worker): - events.append("delegate") - delegated.append(actual_worker) - - def recording_import(name, *args, **kwargs): - if name == "quant_sparse_attn_worker": - events.append("import") - return real_import(name, *args, **kwargs) - - monkeypatch.setitem(sys.modules, "quant_sparse_attn_worker", _legacy_quant_module(install)) - monkeypatch.setattr(builtins, "__import__", recording_import) - monkeypatch.setattr(worker_module, "_quant_api", lambda: events.append("validate")) - - worker_module._install_attention(worker, quantize=True) - - assert events == ["validate", "import", "delegate"] - assert len(delegated) == 1 and delegated[0] is worker - - -def test_shared_worker_sparse_install_does_not_import_legacy_quant_worker(monkeypatch): - worker_module = _load_worker_module("sparse_install_test") - worker = object() - replaced = [] - real_import = builtins.__import__ - - def guarded_import(name, *args, **kwargs): - if name == "quant_sparse_attn_worker": - raise AssertionError("legacy quant worker imported for sparse attention") - return real_import(name, *args, **kwargs) - - monkeypatch.setattr(builtins, "__import__", guarded_import) - monkeypatch.setattr(worker_module, "_replace_attention_impl", replaced.append) - - worker_module._install_attention(worker, quantize=False) - - assert len(replaced) == 1 and replaced[0] is worker - - def test_sparse_worker_keeps_base_memory_profile(): worker_module = _load_worker_module("sparse_profile_test") @@ -180,46 +126,6 @@ def test_sparse_worker_keeps_base_memory_profile(): ) -def test_shared_worker_quant_memory_profile_uses_both_contexts(monkeypatch): - worker_module = _load_worker_module("quant_profile_test") - events = [] - model = object() - - @contextmanager - def recorded_context(name, value=None): - events.append(("enter", name, value)) - try: - yield - finally: - events.append(("exit", name, value)) - - api = SimpleNamespace( - torch=SimpleNamespace(inference_mode=lambda: recorded_context("inference")), - plugin=SimpleNamespace( - disable_compilation=lambda actual_model: recorded_context("compilation", actual_model) - ), - ) - monkeypatch.setattr(worker_module, "_quant_api", lambda: api) - - instance = object.__new__(worker_module.QuantSparseAttnWorker) - instance.model_runner = SimpleNamespace(model=SimpleNamespace(unwrap=lambda: model)) - - def profile(actual_worker): - events.append(("profile", actual_worker)) - return 73 - - monkeypatch.setattr(BaseWorker, "determine_available_memory", profile) - - assert instance.determine_available_memory() == 73 - assert events == [ - ("enter", "inference", None), - ("enter", "compilation", model), - ("profile", instance), - ("exit", "compilation", model), - ("exit", "inference", None), - ] - - def _make_old_impl(): """Create a vLLM FlashAttention impl with initialized runtime state.""" return FlashAttentionImpl( @@ -596,23 +502,70 @@ def fake_attention(query, **kwargs): assert calls["finalize"] == [] +def test_sparse_worker_is_noop_without_checkpoint_metadata(monkeypatch, capsys): + worker_module = _load_worker_module() + attention = _bare_attention(worker_module) + original_impl = attention.impl + state = _sparse_worker_state(nn.ModuleDict({"attn": attention})) + monkeypatch.setattr(worker_module, "load_from_checkpoint_metadata", lambda _: None) + + worker_module._install_attention(state, quantize=False) + + output = capsys.readouterr().out + assert attention.impl is original_impl + assert "No sparse_attention_config found in the checkpoint" in output + assert "replaced impl on 0" not in output + + +def test_sparse_worker_installs_only_active_layer(monkeypatch): + worker_module = _load_worker_module() + active = _bare_attention(worker_module) + disabled = _bare_attention(worker_module, FlashInferImpl) + empty = _bare_attention(worker_module) + original_disabled_impl = disabled.impl + original_empty_impl = empty.impl + model = nn.ModuleDict({"active": active, "disabled": disabled, "empty": empty}) + state = _sparse_worker_state(model) + configs = { + "active": {"enable": True, "sparsity_n": 2, "sparsity_m": 4}, + "disabled": {"enable": False, "sparsity_n": 2, "sparsity_m": 4}, + "empty": {"enable": True}, + } + monkeypatch.setattr(worker_module, "load_from_checkpoint_metadata", lambda _: ({}, "test")) + monkeypatch.setattr(worker_module, "match_sparse_config", lambda name, _cfg: configs[name]) + + worker_module._install_attention(state, quantize=False) + + assert type(active.impl) is ModelOptSparseAttentionImpl + assert active.impl.sparse_kw["sparsity_n"] == 2 + assert disabled.impl is original_disabled_impl + assert empty.impl is original_empty_impl + + +def test_sparse_worker_does_not_load_quant_api(monkeypatch): + worker_module = _load_worker_module() + attention = _bare_attention(worker_module) + state = _sparse_worker_state(nn.ModuleDict({"attn": attention})) + monkeypatch.setattr(worker_module, "load_from_checkpoint_metadata", lambda _: ({}, "test")) + monkeypatch.setattr( + worker_module, + "match_sparse_config", + lambda *_: {"enable": True, "sparsity_n": 2, "sparsity_m": 4}, + ) + monkeypatch.setattr(worker_module, "_quant_api", lambda: pytest.fail("loaded quant API")) + + worker_module._install_attention(state, quantize=False) + + assert type(attention.impl) is ModelOptSparseAttentionImpl + + def test_sparse_worker_rejects_unsupported_backend_before_mutation(monkeypatch): worker_module = _load_worker_module() - good = object.__new__(worker_module.VLLMAttention) - nn.Module.__init__(good) - good.impl = object.__new__(FlashAttentionImpl) - good.impl.sinks = None - bad = object.__new__(worker_module.VLLMAttention) - nn.Module.__init__(bad) + good = _bare_attention(worker_module) + bad = _bare_attention(worker_module) bad.impl = SimpleNamespace() original_good_impl = good.impl - model = nn.ModuleDict({"good": good, "bad": bad}) - state = SimpleNamespace( - model_runner=SimpleNamespace( - model=model, - model_config=SimpleNamespace(hf_config=SimpleNamespace()), - ) - ) + state = _sparse_worker_state(nn.ModuleDict({"good": good, "bad": bad})) monkeypatch.setattr( worker_module, "load_from_checkpoint_metadata", @@ -624,9 +577,12 @@ def test_sparse_worker_rejects_unsupported_backend_before_mutation(monkeypatch): lambda *_: {"enable": True, "sparsity_n": 2, "sparsity_m": 4}, ) - with pytest.raises(NotImplementedError, match=r"bad.*SimpleNamespace"): - worker_module._replace_attention_impl(state) + with pytest.raises(NotImplementedError) as exc: + worker_module._install_attention(state, quantize=False) + assert str(exc.value) == ( + "Unsupported ModelOpt sparse attention plan:\n - bad: unsupported backend SimpleNamespace" + ) assert good.impl is original_good_impl assert bad.impl.__class__ is SimpleNamespace From 29a967f1c9864674e8d85a1a928144cf73f9f171 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 6 Jul 2026 13:13:38 -0700 Subject: [PATCH 14/28] Honor forced ModelOpt attention kernel Signed-off-by: Kai Xu --- .../attention_sparsity/plugins/vllm.py | 15 +++++++-- .../test_sparse_attn_worker.py | 31 +++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 6ff552ff7aa..02b4d42dd44 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -173,7 +173,9 @@ def _has_active_quant_transform(layer, p_qdq, v_qdq) -> bool: def _has_active_transform(impl, layer, p_qdq, v_qdq) -> bool: """Return whether native fallback would omit any ModelOpt transform.""" return bool( - getattr(impl, "sparse_kw", None) or _has_active_quant_transform(layer, p_qdq, v_qdq) + getattr(layer, "_modelopt_force_kernel", False) + or getattr(impl, "sparse_kw", None) + or _has_active_quant_transform(layer, p_qdq, v_qdq) ) @@ -217,9 +219,16 @@ def _forward_modelopt( # N:M sparse softmax is prefill-only. for name in ("sparsity_n", "sparsity_m", "dense_sink_tokens", "dense_recent_tokens"): sparse_kw.pop(name, None) - if not sparse_kw and not quant_transform_active: + if ( + not sparse_kw + and not quant_transform_active + and not getattr(layer, "_modelopt_force_kernel", False) + ): # Dynamic calibration can disable sparse work for a launch. Preserve the # backend's native dense path when no ModelOpt transform remains active. + # ``_modelopt_force_kernel`` (isolation knob) keeps the bf16 kernel path + # even with all quant disabled, so the on/off PPL diff holds the kernel + # constant. return dense_fallback() if prepare_modelopt is not None: prepare_modelopt() @@ -611,7 +620,7 @@ def prepare_modelopt(): # Sparse-only launches may have an inactive phase (for example N:M # is prefill-only). Compute the native result once, then overwrite # each active phase with its ModelOpt result. - if not quant_transform_active: + if not quant_transform_active and not getattr(layer, "_modelopt_force_kernel", False): dense_fallback() block_table = attn_metadata._modelopt_block_table diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index c48a0e2d028..6337c83f3fa 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -502,6 +502,37 @@ def fake_attention(query, **kwargs): assert calls["finalize"] == [] +def test_flashinfer_forced_kernel_mixed_batch_never_uses_native_fallback(monkeypatch): + prefill_tokens = 17 + impl = _make_flashinfer_impl(sparse=False, quantized=False) + query = torch.zeros(1 + prefill_tokens, 2, 64, dtype=torch.float16) + kv_cache = torch.zeros(4, 2, 16, 2, 64, dtype=torch.float16) + metadata = _flashinfer_metadata(query_lens=(1, prefill_tokens), mixed=True) + layer = SimpleNamespace(_modelopt_force_kernel=True) + calls = {"native": 0, "decode": 0, "prefill": 0} + + def native_fallback(*args, **kwargs): + calls["native"] += 1 + raise AssertionError("native FlashInfer fallback") + + def fake_decode(query, *args, **kwargs): + calls["decode"] += 1 + return torch.zeros_like(query) + + def fake_attention(query, **kwargs): + phase = "decode" if kwargs["max_input_len"] == 1 else "prefill" + calls[phase] += 1 + return torch.zeros_like(query) + + monkeypatch.setattr(FlashInferImpl, "forward", native_fallback) + monkeypatch.setattr(vllm_plugin, "triton_decode_attention", fake_decode) + monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) + + impl.forward(layer, query, query, query, kv_cache, metadata, output=torch.empty_like(query)) + + assert calls == {"native": 0, "decode": 1, "prefill": 1} + + def test_sparse_worker_is_noop_without_checkpoint_metadata(monkeypatch, capsys): worker_module = _load_worker_module() attention = _bare_attention(worker_module) From bccd3e6abe8ddd43fa411290b021127ec6a49921 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 6 Jul 2026 13:18:26 -0700 Subject: [PATCH 15/28] Update unified attention worker examples Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 6 ++++-- examples/vllm_serve/vllm_serve_sparse_attn.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 3a6cefc59f5..b816218c1bc 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -117,6 +117,8 @@ Workflow: If the checkpoint has no `sparse_attention_config`, the worker logs a message and passes through — vLLM runs unchanged. Whole-model fakequant flows remain handled by `vllm_serve_fakequant.py`; the compact attention-only path is below. +Both explicit serving policies live in `sparse_attn_worker.py`: `SparseAttnWorker` is checkpoint-driven sparse-only, while `QuantSparseAttnWorker` uses fixed NVFP4 Q/K/P/V plus optional checkpoint sparsity. The launcher keeps `SparseAttnWorker` as its default. + Limitations: - vLLM V1 chunked prefill and prefix-cache suffix attention are supported by offsetting query positions into the longer KV span. @@ -124,14 +126,14 @@ Limitations: ### Compact NVFP4 attention worker -This worker requires vLLM 0.14.0 or newer and fails at import time on older releases. +vLLM 0.14.0 or newer is checked when `QuantSparseAttnWorker` is selected. Importing or using `SparseAttnWorker` does not resolve quant-only APIs. Use the same launcher with the compact worker. By default, vLLM selects the backend for the model and platform; NemotronH on Blackwell selects FlashInfer: ```bash python vllm_serve_sparse_attn.py -tp 8 \ --no-enable-prefix-caching \ - --worker-cls quant_sparse_attn_worker.QuantSparseAttnWorker + --worker-cls sparse_attn_worker.QuantSparseAttnWorker ``` The worker supports both FlashInfer and FlashAttention and prints the installed adapter counts. Pass `--attention-backend FLASHINFER` or `--attention-backend FLASH_ATTN` only when an explicit override is needed. diff --git a/examples/vllm_serve/vllm_serve_sparse_attn.py b/examples/vllm_serve/vllm_serve_sparse_attn.py index 1f077fe0f37..777aff8fc7d 100644 --- a/examples/vllm_serve/vllm_serve_sparse_attn.py +++ b/examples/vllm_serve/vllm_serve_sparse_attn.py @@ -22,7 +22,7 @@ vLLM. The launcher defaults to ``sparse_attn_worker.SparseAttnWorker``. Pass -``--worker-cls quant_sparse_attn_worker.QuantSparseAttnWorker`` for quant+sparse. +``--worker-cls sparse_attn_worker.QuantSparseAttnWorker`` for quant+sparse. Usage: python vllm_serve_sparse_attn.py From da15426d10f358b6f3d014e8b2bd452d9c0a4bdb Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 6 Jul 2026 16:12:14 -0700 Subject: [PATCH 16/28] Simplify attention autotune schedule Signed-off-by: Kai Xu --- .../kernels/common/attention/triton_fa.py | 116 ++++++++----- .../attention/skip_softmax_helpers.py | 32 ---- .../common/attention/test_triton_fa_p_qdq.py | 160 ++++++++++++++++-- .../common/attention/test_triton_fa.py | 148 ++++++++++++++-- 4 files changed, 356 insertions(+), 100 deletions(-) diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 93d0faa738a..898d1550127 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -29,17 +29,16 @@ import triton import triton.language as tl -# Helpers for optional N:M sparsity and sink/window-aware dense regions live -# in the sparsity package. The baseline forward kernel below calls them -# conditionally under constexpr guards, so the unified single-kernel design -# stays intact while keeping feature-specific logic in its own subpackage. +# Helpers for optional N:M sparsity and skip-softmax live in the sparsity +# package. The baseline forward kernel below calls them conditionally under +# constexpr guards, so the unified single-kernel design stays intact while +# keeping feature-specific logic in its own subpackage. # # Lazy import: Triton resolves @triton.jit names at kernel compile time (first # call), not at definition time, so populating the module globals before the # first ``attention()`` call is sufficient. Deferring avoids a circular import # (common.attention/__init__.py ↔ sparsity.attention/__init__.py via this file). _apply_sparse_nm_to_qk_tile: Any = None -_is_dense_region: Any = None _skip_softmax_decision: Any = None _qdq_fp8: Any = None _p_qdq_nvfp4: Any = None @@ -47,20 +46,16 @@ def _load_sparsity_helpers() -> None: - global _apply_sparse_nm_to_qk_tile, _is_dense_region, _skip_softmax_decision + global _apply_sparse_nm_to_qk_tile, _skip_softmax_decision if _apply_sparse_nm_to_qk_tile is None: from modelopt.torch.kernels.sparsity.attention.skip_softmax_helpers import ( _apply_sparse_nm_to_qk_tile as _nm, ) - from modelopt.torch.kernels.sparsity.attention.skip_softmax_helpers import ( - _is_dense_region as _dense, - ) from modelopt.torch.kernels.sparsity.attention.skip_softmax_helpers import ( _skip_softmax_decision as _skip, ) _apply_sparse_nm_to_qk_tile = _nm - _is_dense_region = _dense _skip_softmax_decision = _skip @@ -86,21 +81,15 @@ def _load_qdq_helpers() -> None: # Autotune configs for forward kernel # --------------------------------------------------------------------------- _FWD_CONFIGS = [ - triton.Config({"BLOCK_M": bm, "BLOCK_N": bn}, num_stages=s, num_warps=w) - for bm in [64, 128] - for bn in [32, 64, 128] - for s in [1, 2, 3] - for w in [4, 8] + triton.Config({"BLOCK_M": block_m, "BLOCK_N": 32}, num_stages=2, num_warps=4) + for block_m in (16, 64, 128) ] -# Use a single config in testing for reproducibility -if "PYTEST_VERSION" in __import__("os").environ: - _FWD_CONFIGS = [triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_stages=1, num_warps=4)] - _MEASURE_BLOCK_M = 128 +_P_QDQ_MEASURE_BLOCK_M = 16 # 128 so the kernel sparsity-measurement block matches the PyTorch -# flash_skip_softmax calibration block (br = bc = 128) and the Triton -# calibration kernel; otherwise the two measure at different granularities. +# calibration/reference granularity. This is deliberately independent of the +# autotuned compute tile. _MEASURE_BLOCK_N = 128 _MEASURE_NUM_STAGES = 1 _MEASURE_NUM_WARPS = 4 @@ -226,10 +215,41 @@ def _apply_mask( return scores +@triton.jit +def _apply_sparse_nm_with_dense_tokens( + scores, + kv_start, + q_pos, + kv_pos, + seq_len_q, + seq_len_kv, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + SPARSITY_N: tl.constexpr, + SPARSITY_M: tl.constexpr, + DENSE_SINK_TOKENS: tl.constexpr, + DENSE_RECENT_TOKENS: tl.constexpr, +): + """Apply N:M sparsity outside token-exact sink and recent regions.""" + sparse_scores = _apply_sparse_nm_to_qk_tile(scores, BLOCK_M, BLOCK_N, SPARSITY_N, SPARSITY_M) + q_abs_pos = q_pos[:, None] + seq_len_kv - seq_len_q + kv_abs_pos = kv_start + kv_pos[None, :] + token_distance = q_abs_pos - kv_abs_pos + dense_tokens = ( + (seq_len_q <= 1) + | (kv_abs_pos < DENSE_SINK_TOKENS) + | ((token_distance >= 0) & (token_distance < DENSE_RECENT_TOKENS)) + ) + return tl.where(dense_tokens, scores, sparse_scores) + + # --------------------------------------------------------------------------- # Forward kernel # --------------------------------------------------------------------------- -@triton.autotune(configs=_FWD_CONFIGS, key=["N_CTX", "HEAD_DIM"]) +@triton.autotune( + configs=(_FWD_CONFIGS[:1] if "PYTEST_VERSION" in __import__("os").environ else _FWD_CONFIGS), + key=["N_CTX", "HEAD_DIM", "Q_IS_FP32", "P_QDQ", "V_QDQ"], +) @triton.jit def _attn_fwd( Q, # [total_q, num_q_heads, head_dim] query tensor @@ -375,18 +395,20 @@ def _attn_fwd( # --- Optional N:M sparse softmax --- if SPARSITY_N > 0: - if not _is_dense_region( + scores = _apply_sparse_nm_with_dense_tokens( + scores, kv_start, - tile_q, + q_pos, + kv_pos, seq_len_q, seq_len_kv, BLOCK_M, + BLOCK_N, + SPARSITY_N, + SPARSITY_M, DENSE_SINK_TOKENS, DENSE_RECENT_TOKENS, - ): - scores = _apply_sparse_nm_to_qk_tile( - scores, BLOCK_M, BLOCK_N, SPARSITY_N, SPARSITY_M - ) + ) # Optional skip-softmax decision — the decision logic (and optional # atomic counter updates) lives in sparsity/attention; this kernel @@ -650,24 +672,26 @@ def _attn_bwd_dq( # Re-apply N:M sparse softmax to match forward pass if SPARSITY_N > 0: - if not _is_dense_region( + scores = _apply_sparse_nm_with_dense_tokens( + scores, kv_start, - tile_q, + q_pos, + kv_pos, seq_len_q, seq_len_kv, BLOCK_M, + BLOCK_N, + SPARSITY_N, + SPARSITY_M, DENSE_SINK_TOKENS, DENSE_RECENT_TOKENS, - ): - scores = _apply_sparse_nm_to_qk_tile( - scores, BLOCK_M, BLOCK_N, SPARSITY_N, SPARSITY_M - ) + ) p = tl.math.exp2(scores - lse[:, None]) # Skip-softmax backward: zero out P for rows with negligible contribution. # Per-row using final LSE because forward/backward tile sizes may differ - # (forward autotunes BLOCK_N; backward uses a fixed size), so per-tile + # (forward autotunes BLOCK_M; backward uses a different fixed tile), so per-tile # skip masks from forward wouldn't align. LSE >= any intermediate running # max, so this conservatively zeros out at least what forward skipped. if APPLY_SKIP_SOFTMAX: @@ -804,24 +828,26 @@ def _attn_bwd_dkdv( # Re-apply N:M sparse softmax to match forward pass if SPARSITY_N > 0: - if not _is_dense_region( + scores = _apply_sparse_nm_with_dense_tokens( + scores, kv_start, - qi, + q_pos, + kv_pos, seq_len_q, seq_len_kv, BLOCK_M, + BLOCK_N, + SPARSITY_N, + SPARSITY_M, DENSE_SINK_TOKENS, DENSE_RECENT_TOKENS, - ): - scores = _apply_sparse_nm_to_qk_tile( - scores, BLOCK_M, BLOCK_N, SPARSITY_N, SPARSITY_M - ) + ) p = tl.math.exp2(scores - lse[:, None]) # Skip-softmax backward: zero out P for rows with negligible contribution. # Per-row using final LSE because forward/backward tile sizes may differ - # (forward autotunes BLOCK_N; backward uses a fixed size), so per-tile + # (forward autotunes BLOCK_M; backward uses a different fixed tile), so per-tile # skip masks from forward wouldn't align. LSE >= any intermediate running # max, so this conservatively zeros out at least what forward skipped. if APPLY_SKIP_SOFTMAX: @@ -1010,7 +1036,7 @@ def grid(META): _attn_fwd.fn[grid]( *fwd_args, **fwd_kwargs, - BLOCK_M=_MEASURE_BLOCK_M, + BLOCK_M=_P_QDQ_MEASURE_BLOCK_M if p_qdq_mode else _MEASURE_BLOCK_M, BLOCK_N=_MEASURE_BLOCK_N, num_warps=_MEASURE_NUM_WARPS, num_stages=_MEASURE_NUM_STAGES, @@ -1262,8 +1288,8 @@ def attention( BLOCK_N is a multiple of 16). The softmax denominator stays unquantized. The backward pass uses the straight-through estimator: gradients are computed from the unquantized P, matching QAT - references that keep the backward dots in high precision. - Set to ``None`` to disable. + references that keep the backward dots in high precision. Set to + ``None`` to disable. p_qdq_amax: Per-tensor amax for the softmax-P quant-dequant. The kernel's unnormalized P lies in [0, 1] (the max-subtraction caps every entry at ``exp2(0) = 1``), so 1 is the theoretical upper diff --git a/modelopt/torch/kernels/sparsity/attention/skip_softmax_helpers.py b/modelopt/torch/kernels/sparsity/attention/skip_softmax_helpers.py index 044e54b2e8e..013a4bff0b1 100644 --- a/modelopt/torch/kernels/sparsity/attention/skip_softmax_helpers.py +++ b/modelopt/torch/kernels/sparsity/attention/skip_softmax_helpers.py @@ -188,35 +188,3 @@ def _skip_softmax_decision( tl.atomic_add(Sparsity_skipped, 1) # count skipped tiles return skip_tile - - -# --------------------------------------------------------------------------- -# Sink/window dense-region check -# --------------------------------------------------------------------------- -@triton.jit -def _is_dense_region( - kv_start, - tile_q, - seq_len_q, - seq_len_kv, - BLOCK_M: tl.constexpr, - DENSE_SINK_TOKENS: tl.constexpr, - DENSE_RECENT_TOKENS: tl.constexpr, -): - """Check if a KV tile falls in a dense region (sink tokens or local window). - - Uses absolute token positions so the result is BLOCK_N-independent, - ensuring forward and backward (which may use different BLOCK_N) agree. - - Returns: - True if the tile should be kept dense (skip N:M sparsification). - """ - # N:M sparse softmax is a prefill optimization. Keep decode rows dense, - # including decode rows that are scheduled in a mixed prefill/decode launch. - is_decode = seq_len_q <= 1 - is_sink = kv_start < DENSE_SINK_TOKENS - causal_offset = seq_len_kv - seq_len_q - q_abs_pos = tile_q * BLOCK_M + causal_offset - token_distance = q_abs_pos - kv_start - is_local = (token_distance >= 0) and (token_distance < DENSE_RECENT_TOKENS) - return is_decode or is_sink or is_local diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py index f144e7c4a94..40a0c5fdbe9 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""GPU tests for the softmax quant-dequant (P_QDQ) feature of the Triton FA kernel.""" +"""GPU tests for P-QDQ and tile-sensitive behavior of the Triton FA kernel.""" import math from unittest.mock import Mock @@ -48,12 +48,143 @@ def test_v_qdq_rejects_autograd_before_kernel_launch(monkeypatch): apply.assert_not_called() -# The kernel runs with a single pinned config under pytest (see _FWD_CONFIGS): -# BLOCK_M=128, BLOCK_N=64. The tile-looped reference below relies on it. -BLOCK_N = 64 +# P-QDQ is tile-local, so this is the normal autotuned forward path's numerical +# contract rather than an implementation detail inferred from whichever +# configuration autotune selects. Direct measurement uses separate geometry. +P_QDQ_BLOCK_N = 32 FP8_E4M3_MAX = 448.0 +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need triton") +def test_p_qdq_block_n_matches_forward_configs(): + assert {config.kwargs["BLOCK_N"] for config in triton_fa._FWD_CONFIGS} == {P_QDQ_BLOCK_N} + + +class _FixedBlockMForward: + """Launch the raw forward kernel with one selected query tile.""" + + def __init__(self, autotuner, block_m): + self.fn = autotuner.fn + self._block_m = block_m + + def __getitem__(self, grid): + resolved_grid = grid({"BLOCK_M": self._block_m}) if callable(grid) else grid + + def launch(*args, **kwargs): + return self.fn[resolved_grid]( + *args, + **kwargs, + BLOCK_M=self._block_m, + BLOCK_N=P_QDQ_BLOCK_N, + num_stages=2, + num_warps=4, + ) + + return launch + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need triton") +@pytest.mark.parametrize( + ("seq_len_q", "seq_len_kv", "is_causal"), + [(256, 256, True), (197, 233, True), (256, 256, False)], + ids=["causal", "chunked-prefill", "non-causal"], +) +def test_sparse_dense_window_is_block_m_invariant(monkeypatch, seq_len_q, seq_len_kv, is_causal): + """The dense recent-token policy must not depend on the compute query tile.""" + num_heads, head_dim = 2, 64 + torch.manual_seed(20260706) + q = torch.randn(seq_len_q, num_heads, head_dim, device="cuda", dtype=torch.float16) + k = torch.randn(seq_len_kv, num_heads, head_dim, device="cuda", dtype=torch.float16) + v = torch.randn_like(k) + q_locs, q_lens = make_varlen_meta([seq_len_q]) + kv_locs, kv_lens = make_varlen_meta([seq_len_kv]) + + autotuner = triton_fa._attn_fwd + outputs = {} + for block_m in (16, 64, 128): + monkeypatch.setattr(triton_fa, "_attn_fwd", _FixedBlockMForward(autotuner, block_m)) + outputs[block_m] = attention( + q, + k, + v, + q_locs, + q_lens, + seq_len_q, + is_causal=is_causal, + b_start_loc_k=kv_locs, + b_seq_len_k=kv_lens, + max_input_len_k=seq_len_kv, + sparsity_n=2, + sparsity_m=4, + dense_recent_tokens=64, + ) + + for block_m in (64, 128): + torch.testing.assert_close(outputs[16], outputs[block_m], rtol=2e-3, atol=2e-4) + + +def _sparse_attention_reference(q, k, v, scale, dense_recent_tokens): + """Differentiable 2:4 attention with token-exact dense recent positions.""" + seq_len_q, seq_len_kv = q.shape[0], k.shape[0] + scores = torch.einsum("qhd,khd->hqk", q, k) * scale + q_abs = torch.arange(seq_len_q, device=q.device) + seq_len_kv - seq_len_q + kv_abs = torch.arange(seq_len_kv, device=q.device) + causal = q_abs[:, None] >= kv_abs[None, :] + scores = scores.masked_fill(~causal[None, :, :], float("-inf")) + + grouped = scores.reshape(scores.shape[0], seq_len_q, seq_len_kv // 4, 4) + kept = torch.zeros_like(grouped, dtype=torch.bool) + kept.scatter_(-1, grouped.topk(2, dim=-1).indices, True) + sparse_scores = grouped.masked_fill(~kept, float("-inf")).reshape_as(scores) + + distance = q_abs[:, None] - kv_abs[None, :] + dense = (distance >= 0) & (distance < dense_recent_tokens) + scores = torch.where(dense[None, :, :], scores, sparse_scores) + probabilities = torch.softmax(scores, dim=-1) + return torch.einsum("hqk,khd->qhd", probabilities, v) + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need triton") +def test_sparse_dense_window_forward_and_backward_match_token_reference(): + """Forward and both backward phases use the same tile-independent dense policy.""" + seq_len, num_heads, head_dim = 96, 1, 32 + scale = 1.0 / math.sqrt(head_dim) + dense_recent_tokens = 17 + torch.manual_seed(31415) + q = torch.randn(seq_len, num_heads, head_dim, device="cuda", dtype=torch.float32) + k = torch.randn_like(q) + v = torch.randn_like(q) + weights = torch.randn_like(q) + locs, lens = make_varlen_meta([seq_len]) + + q_actual, k_actual, v_actual = (tensor.clone().requires_grad_(True) for tensor in (q, k, v)) + actual = attention( + q_actual, + k_actual, + v_actual, + locs, + lens, + seq_len, + softmax_scale=scale, + sparsity_n=2, + sparsity_m=4, + dense_recent_tokens=dense_recent_tokens, + ) + (actual * weights).sum().backward() + + q_ref, k_ref, v_ref = (tensor.clone().requires_grad_(True) for tensor in (q, k, v)) + reference = _sparse_attention_reference(q_ref, k_ref, v_ref, scale, dense_recent_tokens) + (reference * weights).sum().backward() + + torch.testing.assert_close(actual, reference, rtol=1e-2, atol=1e-2) + for actual_grad, reference_grad in ( + (q_actual.grad, q_ref.grad), + (k_actual.grad, k_ref.grad), + (v_actual.grad, v_ref.grad), + ): + torch.testing.assert_close(actual_grad, reference_grad, rtol=2e-2, atol=2e-2) + + def _qdq_fp8(p, scale=1.0): """Per-tensor FP8 E4M3 qdq, mirroring the kernel's fp8_scalar_qdq. @@ -103,11 +234,11 @@ def _apply_qdq(p, mode, qdq_scale=1.0): return _qdq_nvfp4(p, qdq_scale) -def qdq_attention_reference(q, k, v, scale, mode, is_causal=True, amax=1.0): +def qdq_attention_reference(q, k, v, scale, mode, is_causal=True, amax=1.0, block_n=P_QDQ_BLOCK_N): """Tile-looped online-softmax reference replicating kernel P_QDQ semantics. Single sequence: q [s, h, d], k/v [s_kv, h_kv, d] (fp16). Walks KV tiles - of BLOCK_N exactly like the kernel, keeps the softmax denominator + of ``block_n`` exactly like the kernel, keeps the softmax denominator unquantized, applies qdq to the unnormalized p of each tile, and mirrors the kernel's operand carrier before the P @ V dot. Native NVFP4 rounds P to the model dtype before packing, then keeps the QDQ result in FP32. FP8 @@ -138,8 +269,8 @@ def qdq_attention_reference(q, k, v, scale, mode, is_causal=True, amax=1.0): row_max = torch.full((h, s), float("-inf"), device=q.device) row_sum = torch.zeros(h, s, device=q.device) acc = torch.zeros(h, s, d, device=q.device) - for start in range(0, s_kv, BLOCK_N): - tile = t[:, :, start : start + BLOCK_N] + for start in range(0, s_kv, block_n): + tile = t[:, :, start : start + block_n] m_new = torch.maximum(row_max, tile.amax(dim=-1)) p = torch.exp2(tile - m_new[..., None]) l_new = p.sum(dim=-1) @@ -151,7 +282,7 @@ def qdq_attention_reference(q, k, v, scale, mode, is_causal=True, amax=1.0): p = _apply_qdq(p, mode, qdq_scale) if mode == "fp8": p = p.to(v.dtype).float() - acc = acc + torch.einsum("hqk,khd->hqd", p, vv[start : start + BLOCK_N].float()) + acc = acc + torch.einsum("hqk,khd->hqd", p, vv[start : start + block_n].float()) row_max = m_new out = acc / row_sum[..., None] return out.permute(1, 0, 2) @@ -165,7 +296,7 @@ class TestSoftmaxQdqForward: @requires_native_e4m3 def test_prefill_matches_tile_reference(self, mode): """Kernel qdq output matches the tile-looped torch reference.""" - seq_len, num_heads, num_kv_heads, head_dim = 128, 4, 2, 64 + seq_len, num_heads, num_kv_heads, head_dim = 256, 4, 2, 64 scale = 1.0 / (head_dim**0.5) torch.manual_seed(7) @@ -174,7 +305,10 @@ def test_prefill_matches_tile_reference(self, mode): o = attention(q, k, v, locs, lens, seq_len, softmax_scale=scale, p_qdq=mode) ref = qdq_attention_reference(q, k, v, scale, mode) - torch.testing.assert_close(o.float(), ref, rtol=5e-3, atol=5e-3) + atol = 2e-2 if mode == "nvfp4" else 5e-3 + # Isolated Triton-vs-PyTorch exp2 differences can cross an NVFP4 quantization + # boundary, so use the same NVFP4 tolerance as the partial-tile coverage. + torch.testing.assert_close(o.float(), ref, rtol=5e-3, atol=atol) # The feature must actually change the output vs dense attention. o_dense = attention(q, k, v, locs, lens, seq_len, softmax_scale=scale) @@ -185,7 +319,7 @@ def test_nvfp4_uses_native_p_input_and_fp32_carriers(self): """Dynamic-Q and P QDQ stay FP32 until their native MMA boundaries.""" num_heads = num_kv_heads = 1 head_dim = 16 - seq_len_kv = BLOCK_N + seq_len_kv = P_QDQ_BLOCK_N scale = 1.0 / math.sqrt(head_dim) q = torch.zeros(1, num_heads, head_dim, device="cuda", dtype=torch.float32) @@ -227,7 +361,7 @@ def test_nvfp4_uses_native_p_input_and_fp32_carriers(self): @pytest.mark.parametrize("mode", ["fp8", "nvfp4"]) @requires_native_e4m3 def test_varlen_partial_tiles(self, mode): - """Variable-length batch with partial KV tiles (seq % BLOCK_N != 0).""" + """Variable-length batch with partial KV tiles (seq % P_QDQ_BLOCK_N != 0).""" seq_lens = [96, 80] total = sum(seq_lens) num_heads, num_kv_heads, head_dim = 4, 2, 64 diff --git a/tests/unit/torch/kernels/common/attention/test_triton_fa.py b/tests/unit/torch/kernels/common/attention/test_triton_fa.py index 0c4aff9c948..941d11c441c 100644 --- a/tests/unit/torch/kernels/common/attention/test_triton_fa.py +++ b/tests/unit/torch/kernels/common/attention/test_triton_fa.py @@ -22,12 +22,41 @@ These tests verify CPU-safe wrapper behavior without executing a Triton kernel. """ +import os from contextlib import nullcontext import pytest import torch +class _CapturingKernel: + def __init__(self): + self.launch_count = 0 + + def __getitem__(self, grid): + self.grid = grid + + def launch(*args, **kwargs): + self.launch_count += 1 + self.kwargs = kwargs + + return launch + + +class _ForbiddenKernel: + def __getitem__(self, grid): + raise AssertionError("unexpected kernel launch") + + +def _config_signature(config): + return ( + config.kwargs["BLOCK_M"], + config.kwargs["BLOCK_N"], + config.num_stages, + config.num_warps, + ) + + def test_triton_fa_importable_on_cpu(): """Module imports cleanly without CUDA; exports the public API names.""" try: @@ -48,16 +77,7 @@ def test_forward_buckets_autotune_key_without_bucketing_grid(monkeypatch): from modelopt.torch.kernels.common.attention import triton_fa - class CapturingKernel: - def __getitem__(self, grid): - self.grid = grid - - def launch(*args, **kwargs): - self.kwargs = kwargs - - return launch - - kernel = CapturingKernel() + kernel = _CapturingKernel() monkeypatch.setattr(triton_fa, "_attn_fwd", kernel) monkeypatch.setattr(triton_fa.torch.cuda, "device", lambda _device: nullcontext()) monkeypatch.setattr(triton_fa, "_load_sparsity_helpers", lambda: None) @@ -74,3 +94,111 @@ def launch(*args, **kwargs): assert kernel.kwargs["N_CTX"] == 256 assert kernel.grid({"BLOCK_M": 64}) == (1, 2, 3) + + +def test_forward_uses_minimal_shared_autotune_configs(): + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + expected_configs = [ + (16, 32, 2, 4), + (64, 32, 2, 4), + (128, 32, 2, 4), + ] + assert [_config_signature(config) for config in triton_fa._FWD_CONFIGS] == expected_configs + + expected_active_configs = ( + expected_configs[:1] if "PYTEST_VERSION" in os.environ else expected_configs + ) + assert [ + _config_signature(config) for config in triton_fa._attn_fwd.configs + ] == expected_active_configs + assert triton_fa._attn_fwd.keys == ["N_CTX", "HEAD_DIM", "Q_IS_FP32", "P_QDQ", "V_QDQ"] + assert "_attn_fwd_p_qdq" not in vars(triton_fa) + assert "_prune_fwd_configs" not in vars(triton_fa) + + +@pytest.mark.parametrize( + ("attention_kwargs", "expected_p_qdq", "expected_v_qdq"), + [ + ({}, 0, 0), + ({"p_qdq": "fp8"}, 1, 0), + ({"p_qdq": "nvfp4"}, 2, 0), + ({"v_qdq": "fp8"}, 0, 1), + ({"v_qdq": "nvfp4"}, 0, 2), + ({"p_qdq": "nvfp4", "v_qdq": "nvfp4"}, 2, 2), + ({"p_qdq": "nvfp4", "sparsity_n": 2, "sparsity_m": 4}, 2, 0), + ({"p_qdq": "nvfp4", "skip_softmax_threshold": 0.1}, 2, 0), + ], +) +def test_forward_routes_every_mode_to_single_autotuner( + monkeypatch, attention_kwargs, expected_p_qdq, expected_v_qdq +): + """Every non-measurement launch uses the unified autotuner.""" + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + kernel = _CapturingKernel() + kernel.fn = _ForbiddenKernel() + monkeypatch.setattr(triton_fa, "_attn_fwd", kernel) + monkeypatch.setattr(triton_fa, "_attn_fwd_p_qdq", _ForbiddenKernel(), raising=False) + monkeypatch.setattr(triton_fa.torch.cuda, "device", lambda _device: nullcontext()) + monkeypatch.setattr(triton_fa, "_load_sparsity_helpers", lambda: None) + monkeypatch.setattr(triton_fa, "_load_qdq_helpers", lambda: None) + + seq_len = 129 + q = torch.empty(seq_len, 2, 16) + k = torch.empty(seq_len, 1, 16) + v = torch.empty_like(k) + starts = torch.tensor([0], dtype=torch.int32) + lengths = torch.tensor([seq_len], dtype=torch.int32) + + triton_fa.attention(q, k, v, starts, lengths, seq_len, **attention_kwargs) + + assert kernel.kwargs["P_QDQ"] == expected_p_qdq + assert kernel.kwargs["V_QDQ"] == expected_v_qdq + + +@pytest.mark.parametrize( + ("attention_kwargs", "expected_block_m"), + [ + ({"skip_softmax_threshold": 0.1, "measure_sparsity": True}, 128), + ( + { + "p_qdq": "nvfp4", + "skip_softmax_threshold": 0.1, + "measure_sparsity": True, + }, + 16, + ), + ], +) +def test_forward_measurement_uses_one_fixed_launch(monkeypatch, attention_kwargs, expected_block_m): + """Counter measurement bypasses autotuning to avoid repeated atomic updates.""" + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + kernel = _ForbiddenKernel() + kernel.fn = _CapturingKernel() + monkeypatch.setattr(triton_fa, "_attn_fwd", kernel) + monkeypatch.setattr(triton_fa.torch.cuda, "device", lambda _device: nullcontext()) + monkeypatch.setattr(triton_fa, "_load_sparsity_helpers", lambda: None) + monkeypatch.setattr(triton_fa, "_load_qdq_helpers", lambda: None) + + seq_len = 129 + q = torch.empty(seq_len, 2, 16) + k = torch.empty(seq_len, 1, 16) + v = torch.empty_like(k) + starts = torch.tensor([0], dtype=torch.int32) + lengths = torch.tensor([seq_len], dtype=torch.int32) + + triton_fa.attention(q, k, v, starts, lengths, seq_len, **attention_kwargs) + + assert kernel.fn.launch_count == 1 + assert kernel.fn.kwargs["BLOCK_M"] == expected_block_m + assert kernel.fn.kwargs["BLOCK_N"] == 128 + assert kernel.fn.kwargs["num_stages"] == 1 + assert kernel.fn.kwargs["num_warps"] == 4 From 93952da35774eb5b8e1a1d35aa6ed038638c931c Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 7 Jul 2026 10:40:41 -0700 Subject: [PATCH 17/28] Remove internal attention design specs Signed-off-by: Kai Xu --- ...5-minimal-mni-attention-autotune-design.md | 101 -------- ...06-unified-vllm-attention-worker-design.md | 218 ------------------ 2 files changed, 319 deletions(-) delete mode 100644 docs/superpowers/specs/2026-07-05-minimal-mni-attention-autotune-design.md delete mode 100644 docs/superpowers/specs/2026-07-06-unified-vllm-attention-worker-design.md diff --git a/docs/superpowers/specs/2026-07-05-minimal-mni-attention-autotune-design.md b/docs/superpowers/specs/2026-07-05-minimal-mni-attention-autotune-design.md deleted file mode 100644 index 4732cc808d1..00000000000 --- a/docs/superpowers/specs/2026-07-05-minimal-mni-attention-autotune-design.md +++ /dev/null @@ -1,101 +0,0 @@ -# Minimal MNI-Aligned Attention Autotune - -## Status - -Approved in conversation on 2026-07-05. - -## Context - -The compact attention branch currently builds a union of dense and P-QDQ Triton -configurations, uses a long semantic cache key, and prunes configurations according -to P-QDQ and sparse-mask flags. This makes the launch policy difficult to review. - -The authoritative MNI reference uses one prefill schedule for dense and every -Q/K/P/V QDQ combination. Its KV tile is fixed at 32, independent of quantization -flags. B200 measurements of the ModelOpt kernel with that fixed KV tile selected -query tiles 16, 64, and 128 at sequence lengths 512, 2048, and 8192 respectively; -all selected two stages and four warps. - -## Goals - -- Use one small, identical eligible configuration set for all normal forward modes. -- Fix the KV tile to 32 to match the MNI prefill numerical schedule. -- Retain query-tile autotuning for short and long prefill efficiency. -- Remove the union configuration construction and mode-aware pruning function. -- Keep sparsity-counter measurement as a direct single launch because autotune - trials would mutate its atomic counters. - -## Non-Goals - -- Reproduce MNI decode tiling or replace the compact split-K decode kernel. -- Align Q/K/P/V carrier dtypes, dot precision, global-scale lifecycle, or every - other MNI arithmetic detail. -- Redesign skip-softmax or N:M dense-region semantics. -- Claim that KV tile 32 is the unrestricted throughput winner. It is the MNI - compatibility contract; unrestricted B200 tuning selected a different tile. - -## Design - -Normal forward launches use one Triton autotuner with three configurations: - -```text -BLOCK_M = {16, 64, 128} -BLOCK_N = 32 -num_stages = 2 -num_warps = 4 -``` - -The autotune key contains only the shape regime and QDQ properties that can change -the best query tile: - -```text -N_CTX, HEAD_DIM, Q_IS_FP32, P_QDQ, V_QDQ -``` - -K QDQ does not need a kernel key because it is materialized before this kernel and -does not change the kernel's dtype or temporary storage. Under pytest, the eligible -set is reduced to the single `BLOCK_M=16, BLOCK_N=32` configuration so tests do not -run benchmark warmups. - -The following implementation elements are removed: - -- `_FWD_BLOCK_M_CHOICES` -- `_FWD_BLOCK_N_CHOICES` -- `_P_QDQ_BLOCK_M_CHOICES` -- `_FWD_AUTOTUNE_KEYS` -- `_prune_fwd_configs` -- Tests coupled to Triton's early-pruning internals - -The direct measurement launch remains outside the autotuner. Its existing stable -measurement geometry remains unchanged in this refactor; changing skip-softmax -calibration geometry is separate work. - -## Behavioral Consequences - -- Dense, Q-only, K-only, P-only, V-only, and combined Q/K/P/V paths see the same - eligible physical tile set. -- P-QDQ output no longer depends on a performance-selected KV tile because every - eligible configuration uses KV tile 32. -- The selected query tile may differ by sequence-length bucket and QDQ combination, - but query tiling does not change P-QDQ's per-row block-16 quantization groups. -- Normal P-QDQ plus 2:4 or skip-softmax uses the same small configuration set instead - of a special fixed autotune profile. The BM16 candidate remains available as the - low-resource fallback. - -## Testing - -- Assert the production configuration set has exactly the three declared tile - shapes and fixed launch parameters. -- Assert the autotune key is the reduced five-field key. -- Assert all non-measurement QDQ combinations route through `_attn_fwd`. -- Retain the direct measurement-launch tests. -- Update the P-QDQ oracle to use KV tile 32. -- Run focused CPU routing tests, P-QDQ GPU tests, and a real dense Triton launch. -- Run repository pre-commit hooks on every touched file. - -## Known Risk - -MNI's `BLOCK_M` fuses GQA heads and query tokens, while ModelOpt's `BLOCK_M` counts -query tokens for one head. Matching the numeric value does not make the physical -grids identical. This design aligns the P-QDQ-sensitive KV tile and simplifies the -policy without claiming bitwise MNI kernel equivalence. diff --git a/docs/superpowers/specs/2026-07-06-unified-vllm-attention-worker-design.md b/docs/superpowers/specs/2026-07-06-unified-vllm-attention-worker-design.md deleted file mode 100644 index 6479bd2727c..00000000000 --- a/docs/superpowers/specs/2026-07-06-unified-vllm-attention-worker-design.md +++ /dev/null @@ -1,218 +0,0 @@ -# Unified vLLM Attention Worker Design - -## Goal - -Remove duplicated sparse-attention planning and installation code from the vLLM -example workers while preserving two explicit serving policies: - -- `SparseAttnWorker` applies checkpoint-driven sparse attention only. -- `QuantSparseAttnWorker` applies the fixed block-16 NVFP4 Q/K/P/V recipe to - every supported attention layer and optionally adds checkpoint-driven sparse - attention. - -The refactor must reduce production code, retain atomic validation, and leave -attention numerics and kernel dispatch unchanged. - -## Current Problem - -`sparse_attn_worker.py` and `quant_sparse_attn_worker.py` independently perform -the same orchestration: - -1. Unwrap the loaded vLLM model. -2. Read sparse-attention checkpoint metadata. -3. Traverse attention modules. -4. Match per-layer sparse configuration. -5. Select and clone the backend-matched ModelOpt implementation. -6. Validate the complete plan before mutating the model. -7. Install the replacement immediately after `BaseWorker.load_model`. - -The quant worker adds strict runtime and layout validation, Q/K/P/V quantizer -configuration, quantization arguments, cascade disablement, and compilation- -disabled memory profiling. Those are policy extensions of the same attention -installation lifecycle, not a separate serving architecture. - -The launcher is not part of this duplication. It owns CLI parsing, worker-module -import setup, and API-server startup, while the worker runs inside each spawned -vLLM GPU process where the model exists. - -## Architecture - -### One Source Module - -`examples/vllm_serve/sparse_attn_worker.py` becomes the single attention-worker -module and declares both public classes in `__all__`: - -```python -class _ModelOptAttentionWorker(BaseWorker): - quantize_attention = False - - def load_model(self, *args, **kwargs): - super().load_model(*args, **kwargs) - _install_attention(self, quantize=self.quantize_attention) - - -class SparseAttnWorker(_ModelOptAttentionWorker): - pass - - -class QuantSparseAttnWorker(_ModelOptAttentionWorker): - quantize_attention = True -``` - -`QuantSparseAttnWorker` retains its `determine_available_memory` override. The -override remains quant-only because sparse-only serving does not need the -compilation workaround and must not inherit its custom-all-reduce interaction. - -`examples/vllm_serve/quant_sparse_attn_worker.py` is deleted. The new quant -worker path is: - -```text -sparse_attn_worker.QuantSparseAttnWorker -``` - -The existing `sparse_attn_worker.SparseAttnWorker` path remains unchanged. - -### Shared Planning Pipeline - -The common pipeline produces immutable plan records before any module mutation. -Each record contains the layer name, attention module, backend-matched -replacement implementation, sparse keyword arguments, and quant-specific -device/dtype fields when quantization is enabled. - -The pipeline performs these shared steps once: - -1. Resolve the loaded model and optional sparse checkpoint configuration. -2. Enumerate vLLM attention modules. -3. Match and normalize each layer's sparse configuration. -4. Select and clone the FlashAttention or FlashInfer adapter. -5. Accumulate all errors and reject the complete plan before mutation. - -Policy-specific validation is invoked from this pipeline rather than mixed into -the shared traversal. - -### Sparse-Only Policy - -With `quantize=False`: - -- Missing `sparse_attention_config` logs the existing message and leaves the - model unchanged. -- Only layers with enabled, nonempty sparse configuration enter the plan. -- Unsupported backends are rejected only for layers requesting a sparse - transform. -- No quantizers, quantization flags, quantization arguments, or profiling - overrides are installed. -- Existing native fallback behavior for inactive sparse launches is preserved. - -### Quant-Plus-Sparse Policy - -With `quantize=True`: - -- vLLM 0.14 or newer remains required. -- Every regular decoder self-attention layer must enter the plan, even when the - checkpoint has no sparse metadata. -- Existing global and per-layer quant validation remains mandatory. -- Each layer receives the fixed dynamic block-16 NVFP4 Q/K/P/V recipe. -- Optional sparse metadata augments the same backend-matched implementation. -- P/V quantization arguments, Q/V kernel ownership flags, K/V default scales, - and cascade disablement remain unchanged. -- The compilation-disabled memory-profile override remains active. - -Quant-only vLLM and ModelOpt imports are resolved only when quant mode is used. -Importing or running `SparseAttnWorker` must not acquire the quant worker's -vLLM-version floor. The quant version/API check therefore runs when the quant -policy is selected, not merely when the shared module is imported. - -## Startup Data Flow - -The vLLM lifecycle remains: - -1. The launcher selects a worker class by dotted path. -2. Each GPU process creates that worker and calls `load_model`. -3. The base worker constructs the native attention modules and implementations. -4. `_ModelOptAttentionWorker.load_model` builds and validates the complete - ModelOpt plan, then installs it. -5. vLLM discovers KV-cache specs and runs memory profiling using the installed - implementation. -6. vLLM initializes cache/backend metadata and performs warmup and CUDA-graph - capture. - -Installation must not move to `compile_or_warm_up_model`; that hook is too late -for attention profiling, FlashInfer metadata setup, and cache initialization. - -## Compatibility - -- Preserve the `sparse_attn_worker.SparseAttnWorker` public path and observable - sparse-only behavior. -- Change the unreleased compact-worker path from - `quant_sparse_attn_worker.QuantSparseAttnWorker` to - `sparse_attn_worker.QuantSparseAttnWorker` in tracked examples, documentation, - and tests. -- Do not add an environment variable or CLI mode flag. The two class names are - the explicit policy selectors. -- Preserve sparse-only import behavior on vLLM releases older than the quant - worker's 0.14 minimum. -- Keep `vllm_serve_sparse_attn.py` as a separate launcher and retain its - sparse-only default. -- Do not merge this path with `FakeQuantWorker`; whole-model fakequant uses a - later lifecycle hook because calibration requires initialized KV-cache state. -- Preserve the current `MODELOPT_ATTN_QUANT_OFF` isolation behavior if it is - present when implementation begins. Removing that benchmark-only behavior is - a separate cleanup decision, not part of worker consolidation. - -## Error Handling - -- Both policies validate the entire selected layer set before the first model - mutation. -- Sparse-only mode continues to no-op when no sparse metadata is present. -- Quant mode fails if no regular attention layers are found or if any selected - layer/runtime feature violates the existing quant contract. -- Backend clone errors and unsupported attention implementations are reported - with layer names in one aggregated `NotImplementedError`. -- A quant-only import/version failure is raised only when - `QuantSparseAttnWorker` is selected. - -## Testing - -Retain the focused sparse and quant worker tests as separate test modules while -pointing both at `sparse_attn_worker.py`. - -Required coverage: - -- Both public classes install after the base `load_model` call. -- Sparse mode remains a no-op without checkpoint metadata. -- Sparse mode mutates only layers with active sparse configuration. -- Quant mode converts every supported regular attention layer and preserves the - exact Q/K/P/V NVFP4 recipe for FlashAttention and FlashInfer. -- Both modes reject a multi-layer plan atomically before mutation. -- Quant-only runtime/layout validation does not reject sparse-only serving. -- Sparse-only import remains valid when quant-only vLLM APIs are unavailable. -- Selecting the quant class below vLLM 0.14 produces the existing clear error. -- Quant memory profiling runs under `disable_compilation`; sparse-only profiling - does not gain that override. -- Full and calibrated-decode CUDA-graph validation remains unchanged. - -Run the focused worker and plugin suites, formatting checks, and an import smoke -for both dotted class paths. A single vLLM serving smoke with the quant class is -required after the refactor because the worker path changes. - -## Scope And Expected Reduction - -The implementation should remove duplicated model traversal, sparse metadata -handling, adapter planning/cloning, load lifecycle, module header, and imports. -The combined production source must be materially smaller than the current 404 -lines across the two workers; a net reduction of approximately 40-70 lines is -expected without weakening validation or deleting regression coverage. - -If preserving older sparse-only imports requires conditional imports, keep that -logic local to the quant policy. Do not create another helper module solely to -move lines between files. - -## Non-Goals - -- Changing attention kernels, tile sizes, quantization scales, or numerical - behavior. -- Replacing the two explicit public worker policies with automatic inference. -- Merging the worker with `vllm_serve_sparse_attn.py`. -- Merging attention-only serving with whole-model `FakeQuantWorker` calibration. -- Adding new sparse algorithms, attention backends, or vLLM feature support. -- Removing focused tests solely to reduce the pull request's line count. From 6649c287fcb25eb5c30d25c31042d4ac76f1c79b Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 7 Jul 2026 13:40:04 -0700 Subject: [PATCH 18/28] Simplify sparse and quant attention plan validation Signed-off-by: Kai Xu --- examples/vllm_serve/sparse_attn_worker.py | 201 +++++++++--------- .../test_quant_sparse_attn_worker.py | 11 +- .../test_sparse_attn_worker.py | 3 +- 3 files changed, 108 insertions(+), 107 deletions(-) diff --git a/examples/vllm_serve/sparse_attn_worker.py b/examples/vllm_serve/sparse_attn_worker.py index 4da178d56de..35b75d476ed 100644 --- a/examples/vllm_serve/sparse_attn_worker.py +++ b/examples/vllm_serve/sparse_attn_worker.py @@ -17,6 +17,7 @@ import importlib import os +from collections import Counter from functools import cache from types import SimpleNamespace from typing import NamedTuple @@ -61,7 +62,6 @@ class _AttentionPlan(NamedTuple): - name: str module: object new_impl: object sparse_kw: dict @@ -74,6 +74,15 @@ def _unwrapped_model(worker): return model.unwrap() if hasattr(model, "unwrap") else model +def _sparse_kwargs(name: str, sparse_cfg: dict | None) -> dict: + if sparse_cfg is None: + return {} + layer_cfg = match_sparse_config(name, sparse_cfg) + if layer_cfg is None or not layer_cfg.get("enable", True): + return {} + return _build_sparse_kw(layer_cfg) + + @cache def _load_quant_api(vllm_version: str): # Keep sparse-only module loading independent of quant-specific vLLM APIs. @@ -113,16 +122,16 @@ def _cudagraph_mode(worker, api): return mode if mode is not None else api.compilation.CUDAGraphMode.NONE -def _global_errors(worker, api=None) -> list[str]: - api = api or _quant_api() +def _global_errors(worker, api) -> list[str]: config = worker.model_runner.vllm_config - parallel, cache, model_config = config.parallel_config, config.cache_config, config.model_config + parallel = config.parallel_config + cache_config, model_config = config.cache_config, config.model_config errors = [] if getattr(parallel, "decode_context_parallel_size", 1) != 1: errors.append("decode_context_parallel_size must be 1") if getattr(parallel, "enable_dbo", False) or getattr(parallel, "use_ubatching", False): errors.append("DBO/ubatching is unsupported") - if getattr(cache, "enable_prefix_caching", False): + if getattr(cache_config, "enable_prefix_caching", False): errors.append("prefix caching is unsupported") if getattr(config, "kv_transfer_config", None) is not None: errors.append("KV transfer is unsupported") @@ -132,7 +141,7 @@ def _global_errors(worker, api=None) -> list[str]: errors.append("FULL mixed-batch cudagraph mode is unsupported") if getattr(model_config, "dtype", None) not in (api.torch.float16, api.torch.bfloat16): errors.append("resolved model/KV-cache dtype must be fp16 or bf16") - cache_dtype = getattr(cache, "cache_dtype", "auto") + cache_dtype = getattr(cache_config, "cache_dtype", "auto") if str(cache_dtype) not in {"auto", "bfloat16", "float16", "torch.bfloat16", "torch.float16"}: errors.append(f"resolved KV-cache dtype {cache_dtype!r} must be fp16 or bf16") return errors @@ -166,21 +175,8 @@ def _quant_layer_errors(module, api) -> list[str]: return errors -def _validated_device_dtype(module, model_config, api): - device, dtype = api.plugin._get_device_dtype(module) - model_dtype = getattr(model_config, "dtype", None) - if model_dtype in (api.torch.float16, api.torch.bfloat16): - dtype = model_dtype - if device is None or dtype is None: - return device, dtype, "device/dtype could not be resolved" - if dtype not in (api.torch.float16, api.torch.bfloat16): - return device, dtype, f"resolved dtype {dtype} must be fp16 or bf16" - return device, dtype, None - - -def _sparse_graph_error(sparse_kw: dict, mode, api=None) -> str | None: +def _sparse_graph_error(sparse_kw: dict, mode, api) -> str | None: """Reject decode calibration whose live length would be frozen by a full graph.""" - api = api or _quant_api() params = sparse_kw.get("threshold_scale_factor") if ( mode.decode_mode() == api.compilation.CUDAGraphMode.FULL @@ -191,99 +187,103 @@ def _sparse_graph_error(sparse_kw: dict, mode, api=None) -> str | None: return None -def _validated_attention_plans(worker, *, quantize: bool): - """Validate and clone every selected attention adapter before mutating layers.""" - api = _quant_api() if quantize else None +def _select_new_impl(module): + """Clone the module's attention impl into its sparse-capable subclass; return (impl, error).""" + try: + cls = select_sparse_impl_cls(module.impl) + except (NotImplementedError, TypeError) as err: + return None, str(err) + if cls is None: + return None, ( + f"backend {type(module.impl).__name__} is not supported; " + "expected FlashAttentionImpl or FlashInferImpl" + ) + return _clone_sparse_impl(module.impl, cls), None + + +def _raise_unsupported(errors: list[str], policy: str) -> None: + if errors: + raise NotImplementedError( + f"Unsupported ModelOpt {policy} plan:\n - " + "\n - ".join(errors) + ) + + +def _sparse_plans(worker): + """Plans for checkpoint-driven sparse attention; skips layers without a sparse config.""" model = _unwrapped_model(worker) - model_config = worker.model_runner.model_config - detected = load_from_checkpoint_metadata(getattr(model_config, "hf_config", None)) - if not quantize and detected is None: + detected = load_from_checkpoint_metadata( + getattr(worker.model_runner.model_config, "hf_config", None) + ) + if detected is None: print( "[ModelOpt] No sparse_attention_config found in the checkpoint; " - "skipping sparse attention. Run examples/llm_sparsity/" - "attention_sparsity/hf_sa.py to calibrate and export a checkpoint " - "with the config embedded." + "skipping sparse attention. Run examples/llm_sparsity/attention_sparsity/" + "hf_sa.py to calibrate and export a checkpoint with the config embedded." ) return None + sparse_cfg, sparse_algo = detected + print(f"[ModelOpt] Sparse attention config: algo -> {sparse_algo}") + plans, errors = [], [] + for name, module in model.named_modules(): + if not isinstance(module, VLLMAttention): + continue + sparse_kw = _sparse_kwargs(name, sparse_cfg) + if not sparse_kw: + continue + new_impl, error = _select_new_impl(module) + if error: + errors.append(f"{name or ''}: {error}") + else: + plans.append(_AttentionPlan(module, new_impl, sparse_kw, None, None)) + _raise_unsupported(errors, "sparse attention") + return tuple(plans) + +def _quant_plans(worker): + """Plans for fixed-NVFP4 attention on every decoder self-attention layer (+ optional sparsity).""" + api = _quant_api() + model = _unwrapped_model(worker) + model_config = worker.model_runner.model_config + detected = load_from_checkpoint_metadata(getattr(model_config, "hf_config", None)) sparse_cfg = detected[0] if detected is not None else None - if quantize: - assert api is not None - errors = _global_errors(worker, api) - mode = _cudagraph_mode(worker, api) - attention_types = api.plugin._ATTENTION_TYPES - else: - assert detected is not None - print(f"[ModelOpt] Sparse attention config: algo -> {detected[1]}") - errors, mode, attention_types = [], None, (VLLMAttention,) - plans = [] - attention_count = 0 + errors = _global_errors(worker, api) + mode = _cudagraph_mode(worker, api) + plans, attention_count = [], 0 for name, module in model.named_modules(): - if not isinstance(module, attention_types): + if not isinstance(module, api.plugin._ATTENTION_TYPES): continue - if quantize: - assert api is not None - attention_count += 1 - reasons = _quant_layer_errors(module, api) - device, dtype, dtype_error = _validated_device_dtype(module, model_config, api) - if dtype_error: - reasons.append(dtype_error) - layer_cfg = match_sparse_config(name, sparse_cfg) if sparse_cfg is not None else None - sparse_kw = ( - _build_sparse_kw(layer_cfg) - if layer_cfg is not None and layer_cfg.get("enable", True) - else {} - ) - if graph_error := _sparse_graph_error(sparse_kw, mode, api): - reasons.append(graph_error) - else: - assert sparse_cfg is not None - layer_cfg = match_sparse_config(name, sparse_cfg) - if layer_cfg is None or not layer_cfg.get("enable", True): - continue - sparse_kw = _build_sparse_kw(layer_cfg) - if not sparse_kw: - continue - reasons, device, dtype = [], None, None - - new_impl = None - try: - new_impl_cls = select_sparse_impl_cls(module.impl) - if new_impl_cls is None: - backend = type(module.impl).__name__ - message = ( - f"backend {backend} is not supported; expected FlashAttentionImpl or FlashInferImpl" - if quantize - else f"unsupported backend {backend}" - ) - reasons.append(message) - else: - new_impl = _clone_sparse_impl(module.impl, new_impl_cls) - except (NotImplementedError, TypeError) as err: - reasons.append(str(err)) - layer_name = name or "" + attention_count += 1 + reasons = _quant_layer_errors(module, api) + # Prefer the model compute dtype (fp16/bf16); _get_device_dtype's buffer scan + # can otherwise report fp32 from the attention module's scale buffers. + device, dtype = api.plugin._get_device_dtype(module) + if getattr(model_config, "dtype", None) in (api.torch.float16, api.torch.bfloat16): + dtype = model_config.dtype + if device is None or dtype is None: + reasons.append("device/dtype could not be resolved") + elif dtype not in (api.torch.float16, api.torch.bfloat16): + reasons.append(f"resolved dtype {dtype} must be fp16 or bf16") + sparse_kw = _sparse_kwargs(name, sparse_cfg) + if graph_error := _sparse_graph_error(sparse_kw, mode, api): + reasons.append(graph_error) + new_impl, error = _select_new_impl(module) + if error: + reasons.append(error) if reasons: - errors.extend(f"{layer_name}: {reason}" for reason in reasons) + errors.extend(f"{name or ''}: {reason}" for reason in reasons) else: - plans.append(_AttentionPlan(name, module, new_impl, sparse_kw, device, dtype)) - - if quantize and attention_count == 0: + plans.append(_AttentionPlan(module, new_impl, sparse_kw, device, dtype)) + if attention_count == 0: errors.append("no regular attention layers were found") - if errors: - policy = "attention" if quantize else "sparse attention" - raise NotImplementedError( - f"Unsupported ModelOpt {policy} plan:\n - " + "\n - ".join(errors) - ) + _raise_unsupported(errors, "attention") return tuple(plans) def _install_sparse_plans(plans) -> None: - installed = {} for plan in plans: plan.new_impl.sparse_kw = plan.sparse_kw plan.module.impl = plan.new_impl - impl_name = type(plan.new_impl).__name__ - installed[impl_name] = installed.get(impl_name, 0) + 1 + installed = dict(Counter(type(plan.new_impl).__name__ for plan in plans)) print( f"[ModelOpt] Sparse attention: replaced impl on {len(plans)} attention layers: {installed}" ) @@ -292,7 +292,6 @@ def _install_sparse_plans(plans) -> None: def _install_quant_plans(worker, plans) -> None: api = _quant_api() quant_off = os.environ.get("MODELOPT_ATTN_QUANT_OFF") == "1" - installed = {} for plan in plans: module = plan.module module.device, module.dtype = plan.device, plan.dtype @@ -322,20 +321,18 @@ def _install_quant_plans(worker, plans) -> None: module._value_quant_in_kernel = not quant_off if quant_off: module._modelopt_force_kernel = True - impl_name = type(plan.new_impl).__name__ - installed[impl_name] = installed.get(impl_name, 0) + 1 worker.model_runner.cascade_attn_enabled = False + installed = dict(Counter(type(plan.new_impl).__name__ for plan in plans)) print(f"[ModelOpt] Installed NVFP4 quant+sparse attention on {len(plans)} layers: {installed}") def _install_attention(worker, *, quantize: bool) -> None: - plans = _validated_attention_plans(worker, quantize=quantize) - if plans is None: - return if quantize: - _install_quant_plans(worker, plans) + _install_quant_plans(worker, _quant_plans(worker)) else: - _install_sparse_plans(plans) + plans = _sparse_plans(worker) + if plans is not None: + _install_sparse_plans(plans) class _ModelOptAttentionWorker(BaseWorker): diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py index 25c7a7623c3..d706ec11508 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py @@ -93,7 +93,7 @@ def _patch_conversion(monkeypatch): lambda: quant_plugin.ParallelState(data_parallel_group=None), ) monkeypatch.setattr(worker_module, "load_from_checkpoint_metadata", lambda _: None) - monkeypatch.setattr(worker_module, "_global_errors", lambda _worker, _api=None: []) + monkeypatch.setattr(worker_module, "_global_errors", lambda _worker, _api: []) @pytest.mark.parametrize("impl_cls", [FlashAttentionImpl, FlashInferImpl]) @@ -217,6 +217,7 @@ def profile(actual_worker): [(CUDAGraphMode.FULL, True), (CUDAGraphMode.FULL_AND_PIECEWISE, False)], ) def test_full_mixed_cudagraph_validation(mode, rejected): + api = worker_module._quant_api() config = SimpleNamespace( parallel_config=SimpleNamespace(), cache_config=SimpleNamespace(), @@ -224,12 +225,13 @@ def test_full_mixed_cudagraph_validation(mode, rejected): compilation_config=SimpleNamespace(cudagraph_mode=mode), ) errors = worker_module._global_errors( - SimpleNamespace(model_runner=SimpleNamespace(vllm_config=config)) + SimpleNamespace(model_runner=SimpleNamespace(vllm_config=config)), api ) assert any("mixed" in error for error in errors) is rejected def test_calibrated_decode_skip_softmax_rejects_full_decode_graphs(): + api = worker_module._quant_api() sparse_kw = { "threshold_scale_factor": { "prefill": {"a": 1.0, "b": 2.0}, @@ -238,13 +240,14 @@ def test_calibrated_decode_skip_softmax_rejects_full_decode_graphs(): } assert "calibrated decode skip-softmax" in worker_module._sparse_graph_error( - sparse_kw, CUDAGraphMode.FULL_AND_PIECEWISE + sparse_kw, CUDAGraphMode.FULL_AND_PIECEWISE, api ) - assert worker_module._sparse_graph_error(sparse_kw, CUDAGraphMode.PIECEWISE) is None + assert worker_module._sparse_graph_error(sparse_kw, CUDAGraphMode.PIECEWISE, api) is None assert ( worker_module._sparse_graph_error( {"threshold_scale_factor": {"prefill": {"a": 1.0, "b": 2.0}}}, CUDAGraphMode.FULL_AND_PIECEWISE, + api, ) is None ) diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 6337c83f3fa..d940f64c725 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -612,7 +612,8 @@ def test_sparse_worker_rejects_unsupported_backend_before_mutation(monkeypatch): worker_module._install_attention(state, quantize=False) assert str(exc.value) == ( - "Unsupported ModelOpt sparse attention plan:\n - bad: unsupported backend SimpleNamespace" + "Unsupported ModelOpt sparse attention plan:\n - bad: backend SimpleNamespace " + "is not supported; expected FlashAttentionImpl or FlashInferImpl" ) assert good.impl is original_good_impl assert bad.impl.__class__ is SimpleNamespace From 5ef7dea2ffdf7202408b3b786a6f9af2b4b5d092 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 7 Jul 2026 15:21:01 -0700 Subject: [PATCH 19/28] Simplify sparse attention vLLM plugin forward paths Signed-off-by: Kai Xu --- .../attention_sparsity/plugins/vllm.py | 421 +++++++++++------- .../common/attention/test_decode_attention.py | 71 ++- .../common/attention/test_triton_fa_p_qdq.py | 30 +- .../common/attention/test_triton_fa_paged.py | 165 +------ .../quantization/test_vllm_dynamic_modules.py | 16 +- .../test_sparse_attn_worker.py | 230 ++++++---- .../common/attention/test_triton_fa.py | 36 +- 7 files changed, 465 insertions(+), 504 deletions(-) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 02b4d42dd44..e15c215b522 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -31,6 +31,7 @@ import inspect import math import warnings +from dataclasses import dataclass import torch from vllm.v1.attention.backends.flash_attn import ( @@ -159,7 +160,7 @@ def _quant_kw_from_impl(impl, layer): return p_qdq, p_qdq_amax, v_qdq, v_qdq_amax -def _has_active_quant_transform(layer, p_qdq, v_qdq) -> bool: +def _any_quant_active(layer, p_qdq, v_qdq) -> bool: """Return whether native fallback would omit any Q/K/P/V transform.""" k_quantizer = getattr(layer, "k_bmm_quantizer", None) return bool( @@ -170,15 +171,75 @@ def _has_active_quant_transform(layer, p_qdq, v_qdq) -> bool: ) -def _has_active_transform(impl, layer, p_qdq, v_qdq) -> bool: - """Return whether native fallback would omit any ModelOpt transform.""" - return bool( - getattr(layer, "_modelopt_force_kernel", False) - or getattr(impl, "sparse_kw", None) - or _has_active_quant_transform(layer, p_qdq, v_qdq) +def _should_run_modelopt_kernel(sparse_kw, quant_active: bool, force_kernel: bool) -> bool: + """Return whether a launch has effective work for the ModelOpt kernel.""" + return bool(sparse_kw or quant_active or force_kernel) + + +def _mixed_batch_needs_dense_base(quant_active: bool, force_kernel: bool) -> bool: + """Return whether sparse-only mixed phases need one native output seed.""" + return not quant_active and not force_kernel + + +@dataclass(frozen=True, slots=True) +class _ResolvedForward: + p_qdq: str | None + p_qdq_amax: float + v_qdq: str | None + v_qdq_amax: float | None + quant_active: bool + force_kernel: bool + + +def _resolve_forward( + impl, + layer, + attn_metadata, + output_scale, + output_block_scale, + *, + require_flashinfer_metadata: bool = False, +) -> _ResolvedForward | None: + """Resolve shared transform state or request the backend's native path.""" + p_qdq, p_qdq_amax, v_qdq, v_qdq_amax = _quant_kw_from_impl(impl, layer) + quant_active = _any_quant_active(layer, p_qdq, v_qdq) + force_kernel = getattr(layer, "_modelopt_force_kernel", False) + transform_active = _should_run_modelopt_kernel( + getattr(impl, "sparse_kw", None), quant_active, force_kernel + ) + + if getattr(attn_metadata, "use_cascade", False): + if transform_active: + raise NotImplementedError( + "vLLM cascade attention is incompatible with an active ModelOpt attention transform" + ) + return None + + if require_flashinfer_metadata: + missing = [name for name in _FLASHINFER_METADATA_FIELDS if not hasattr(attn_metadata, name)] + if missing: + if transform_active: + raise NotImplementedError( + "FlashInfer metadata is missing the ModelOpt attention transform " + f"fields: {', '.join(missing)}" + ) + return None + + if transform_active and (output_scale is not None or output_block_scale is not None): + raise NotImplementedError("Fused attention output quantization is unsupported") + + return _ResolvedForward( + p_qdq=p_qdq, + p_qdq_amax=p_qdq_amax, + v_qdq=v_qdq, + v_qdq_amax=v_qdq_amax, + quant_active=quant_active, + force_kernel=force_kernel, ) +# Resolution guards raw configured transforms; dispatch rechecks effective +# sparse work after calibration and decode-only pruning. def _forward_modelopt( impl, *, @@ -198,7 +259,8 @@ def _forward_modelopt( p_qdq_amax: float, v_qdq: str | None, v_qdq_amax: float | None, - quant_transform_active: bool, + quant_active: bool, + force_kernel: bool, dense_fallback, prepare_modelopt=None, ) -> torch.Tensor: @@ -219,10 +281,10 @@ def _forward_modelopt( # N:M sparse softmax is prefill-only. for name in ("sparsity_n", "sparsity_m", "dense_sink_tokens", "dense_recent_tokens"): sparse_kw.pop(name, None) - if ( - not sparse_kw - and not quant_transform_active - and not getattr(layer, "_modelopt_force_kernel", False) + if not _should_run_modelopt_kernel( + sparse_kw, + quant_active, + force_kernel, ): # Dynamic calibration can disable sparse work for a launch. Preserve the # backend's native dense path when no ModelOpt transform remains active. @@ -353,14 +415,7 @@ def forward( if attn_metadata is None: return output.fill_(0) - p_qdq, p_qdq_amax, v_qdq, v_qdq_amax = _quant_kw_from_impl(self, layer) - quant_transform_active = _has_active_quant_transform(layer, p_qdq, v_qdq) - active = _has_active_transform(self, layer, p_qdq, v_qdq) - if getattr(attn_metadata, "use_cascade", False): - if active: - raise NotImplementedError( - "vLLM cascade attention is incompatible with an active ModelOpt attention transform" - ) + def native_forward(): return self._forward_vllm_flash_attn( layer, query, @@ -372,8 +427,16 @@ def forward( output_scale, output_block_scale, ) - if active and (output_scale is not None or output_block_scale is not None): - raise NotImplementedError("Fused attention output quantization is unsupported") + + resolved = _resolve_forward( + self, + layer, + attn_metadata, + output_scale, + output_block_scale, + ) + if resolved is None: + return native_forward() key_cache, value_cache = kv_cache.unbind(0) is_decode_only = attn_metadata.max_query_len <= 1 @@ -391,22 +454,13 @@ def forward( max_seq_len=attn_metadata.max_seq_len, is_causal=getattr(attn_metadata, "causal", not is_decode_only), output=output, - p_qdq=p_qdq, - p_qdq_amax=p_qdq_amax, - v_qdq=v_qdq, - v_qdq_amax=v_qdq_amax, - quant_transform_active=quant_transform_active, - dense_fallback=lambda: self._forward_vllm_flash_attn( - layer, - query, - key, - value, - kv_cache, - attn_metadata, - output, - output_scale, - output_block_scale, - ), + p_qdq=resolved.p_qdq, + p_qdq_amax=resolved.p_qdq_amax, + v_qdq=resolved.v_qdq, + v_qdq_amax=resolved.v_qdq_amax, + quant_active=resolved.quant_active, + force_kernel=resolved.force_kernel, + dense_fallback=native_forward, ) @@ -440,6 +494,13 @@ def get_impl_cls() -> type: } +def _reset_flashinfer_state_for_tests() -> None: + """Clear lazy state without unwrapping the process-wide builder patch.""" + global _FLASHINFER_PATCHED, _FLASHINFER_IMPL_CLS + _FLASHINFER_PATCHED = False + _FLASHINFER_IMPL_CLS = None + + def patch_flashinfer_metadata_builder() -> bool: """Attach the common paged metadata needed by the ModelOpt kernels.""" global _FLASHINFER_PATCHED @@ -454,6 +515,8 @@ def patch_flashinfer_metadata_builder() -> bool: if getattr(orig_build, "_modelopt_sparse_metadata_patch", False): _FLASHINFER_PATCHED = True return True + # vLLM compatibility contract: build has a named ``common_attn_metadata`` + # argument and returns a mutable metadata object that accepts attached fields. build_sig = inspect.signature(orig_build) @functools.wraps(orig_build) @@ -495,6 +558,146 @@ def _maybe_update_flashinfer_cache(layer, key, value, kv_cache, attn_metadata, i _flashinfer_cache_write(layer, key, value, kv_cache, attn_metadata, impl) +def _flashinfer_forward( + impl, + native_forward, + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output=None, + output_scale=None, + output_block_scale=None, +): + """Run the FlashInfer adapter with module-scope, directly testable logic.""" + assert output is not None, "Output tensor must be provided." + if attn_metadata is None: + return output.fill_(0) + + dense_output = None + cache_prepared = False + + def dense_fallback(): + nonlocal cache_prepared, dense_output + if dense_output is None: + dense_output = native_forward( + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output, + output_scale, + output_block_scale, + ) + cache_prepared = True + return dense_output + + def prepare_modelopt(): + nonlocal cache_prepared + if not cache_prepared: + _maybe_update_flashinfer_cache(layer, key, value, kv_cache, attn_metadata, impl) + cache_prepared = True + + resolved = _resolve_forward( + impl, + layer, + attn_metadata, + output_scale, + output_block_scale, + require_flashinfer_metadata=True, + ) + if resolved is None: + return dense_fallback() + + if kv_cache.ndim != 5 or kv_cache.shape[1] != 2: + raise ValueError( + "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim]" + ) + + key_cache = kv_cache[:, 0] + value_cache = kv_cache[:, 1] + max_query_len = attn_metadata._modelopt_max_query_len + is_decode_only = max_query_len <= 1 + common_kw = { + "layer": layer, + "key_cache": key_cache, + "value_cache": value_cache, + "max_seq_len": attn_metadata._modelopt_max_seq_len, + "is_causal": getattr(attn_metadata, "_modelopt_causal", not is_decode_only), + "p_qdq": resolved.p_qdq, + "p_qdq_amax": resolved.p_qdq_amax, + "v_qdq": resolved.v_qdq, + "v_qdq_amax": resolved.v_qdq_amax, + "quant_active": resolved.quant_active, + "force_kernel": resolved.force_kernel, + "dense_fallback": dense_fallback, + "prepare_modelopt": prepare_modelopt, + } + num_decodes = getattr(attn_metadata, "num_decodes", 0) + num_prefills = getattr(attn_metadata, "num_prefills", 0) + if not (num_decodes and num_prefills): + return _forward_modelopt( + impl, + query=query, + block_table=attn_metadata._modelopt_block_table, + seq_lens=attn_metadata._modelopt_seq_lens, + cu_seqlens_q=attn_metadata._modelopt_query_start_loc, + num_actual_tokens=attn_metadata._modelopt_num_actual_tokens, + max_query_len=max_query_len, + output=output, + **common_kw, + ) + + num_decode_tokens = attn_metadata.num_decode_tokens + num_prefill_tokens = attn_metadata.num_prefill_tokens + if num_decode_tokens % num_decodes: + raise NotImplementedError("Non-uniform mixed FlashInfer decode is unsupported") + if num_decode_tokens + num_prefill_tokens != attn_metadata._modelopt_num_actual_tokens: + raise ValueError("FlashInfer mixed-batch token counts do not match common metadata") + + # Sparse-only launches may have an inactive phase (for example N:M + # is prefill-only). Compute the native result once, then overwrite + # each active phase with its ModelOpt result. + if _mixed_batch_needs_dense_base( + resolved.quant_active, + resolved.force_kernel, + ): + dense_fallback() + + block_table = attn_metadata._modelopt_block_table + seq_lens = attn_metadata._modelopt_seq_lens + cu_seqlens_q = attn_metadata._modelopt_query_start_loc + _forward_modelopt( + impl, + query=query[:num_decode_tokens], + block_table=block_table[:num_decodes], + seq_lens=seq_lens[:num_decodes], + cu_seqlens_q=cu_seqlens_q[: num_decodes + 1], + num_actual_tokens=num_decode_tokens, + max_query_len=num_decode_tokens // num_decodes, + output=output[:num_decode_tokens], + **common_kw, + ) + prefill_start = num_decode_tokens + prefill_cu_seqlens_q = cu_seqlens_q[num_decodes:] - cu_seqlens_q[num_decodes] + _forward_modelopt( + impl, + query=query[prefill_start : prefill_start + num_prefill_tokens], + block_table=block_table[num_decodes : num_decodes + num_prefills], + seq_lens=seq_lens[num_decodes : num_decodes + num_prefills], + cu_seqlens_q=prefill_cu_seqlens_q, + num_actual_tokens=num_prefill_tokens, + max_query_len=max_query_len, + output=output[prefill_start : prefill_start + num_prefill_tokens], + **common_kw, + ) + return output + + def get_flashinfer_sparse_impl_cls() -> type: """Return the lazy FlashInfer adapter without requiring it for FA users.""" global _FLASHINFER_IMPL_CLS @@ -518,139 +721,19 @@ def forward( output_scale=None, output_block_scale=None, ): - assert output is not None, "Output tensor must be provided." - if attn_metadata is None: - return output.fill_(0) - - dense_output = None - cache_prepared = False - - def dense_fallback(): - nonlocal cache_prepared, dense_output - if dense_output is None: - dense_output = FlashInferImpl.forward( - self, - layer, - query, - key, - value, - kv_cache, - attn_metadata, - output, - output_scale, - output_block_scale, - ) - cache_prepared = True - return dense_output - - def prepare_modelopt(): - nonlocal cache_prepared - if not cache_prepared: - _maybe_update_flashinfer_cache(layer, key, value, kv_cache, attn_metadata, self) - cache_prepared = True - - p_qdq, p_qdq_amax, v_qdq, v_qdq_amax = _quant_kw_from_impl(self, layer) - quant_transform_active = _has_active_quant_transform(layer, p_qdq, v_qdq) - active = _has_active_transform(self, layer, p_qdq, v_qdq) - if getattr(attn_metadata, "use_cascade", False): - if active: - raise NotImplementedError( - "vLLM cascade attention is incompatible with an active " - "ModelOpt attention transform" - ) - return dense_fallback() - - missing = [ - name for name in _FLASHINFER_METADATA_FIELDS if not hasattr(attn_metadata, name) - ] - if missing: - if active: - raise NotImplementedError( - "FlashInfer metadata is missing the ModelOpt attention transform " - f"fields: {', '.join(missing)}" - ) - return dense_fallback() - if active and (output_scale is not None or output_block_scale is not None): - raise NotImplementedError("Fused attention output quantization is unsupported") - if kv_cache.ndim != 5 or kv_cache.shape[1] != 2: - raise ValueError( - "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim]" - ) - - key_cache = kv_cache[:, 0] - value_cache = kv_cache[:, 1] - max_query_len = attn_metadata._modelopt_max_query_len - is_decode_only = max_query_len <= 1 - common_kw = { - "layer": layer, - "key_cache": key_cache, - "value_cache": value_cache, - "max_seq_len": attn_metadata._modelopt_max_seq_len, - "is_causal": getattr(attn_metadata, "_modelopt_causal", not is_decode_only), - "p_qdq": p_qdq, - "p_qdq_amax": p_qdq_amax, - "v_qdq": v_qdq, - "v_qdq_amax": v_qdq_amax, - "quant_transform_active": quant_transform_active, - "dense_fallback": dense_fallback, - "prepare_modelopt": prepare_modelopt, - } - num_decodes = getattr(attn_metadata, "num_decodes", 0) - num_prefills = getattr(attn_metadata, "num_prefills", 0) - if not (num_decodes and num_prefills): - return _forward_modelopt( - self, - query=query, - block_table=attn_metadata._modelopt_block_table, - seq_lens=attn_metadata._modelopt_seq_lens, - cu_seqlens_q=attn_metadata._modelopt_query_start_loc, - num_actual_tokens=attn_metadata._modelopt_num_actual_tokens, - max_query_len=max_query_len, - output=output, - **common_kw, - ) - - num_decode_tokens = attn_metadata.num_decode_tokens - num_prefill_tokens = attn_metadata.num_prefill_tokens - if num_decode_tokens % num_decodes: - raise NotImplementedError("Non-uniform mixed FlashInfer decode is unsupported") - if num_decode_tokens + num_prefill_tokens != attn_metadata._modelopt_num_actual_tokens: - raise ValueError("FlashInfer mixed-batch token counts do not match common metadata") - - # Sparse-only launches may have an inactive phase (for example N:M - # is prefill-only). Compute the native result once, then overwrite - # each active phase with its ModelOpt result. - if not quant_transform_active and not getattr(layer, "_modelopt_force_kernel", False): - dense_fallback() - - block_table = attn_metadata._modelopt_block_table - seq_lens = attn_metadata._modelopt_seq_lens - cu_seqlens_q = attn_metadata._modelopt_query_start_loc - _forward_modelopt( + return _flashinfer_forward( self, - query=query[:num_decode_tokens], - block_table=block_table[:num_decodes], - seq_lens=seq_lens[:num_decodes], - cu_seqlens_q=cu_seqlens_q[: num_decodes + 1], - num_actual_tokens=num_decode_tokens, - max_query_len=num_decode_tokens // num_decodes, - output=output[:num_decode_tokens], - **common_kw, - ) - prefill_start = num_decode_tokens - prefill_cu_seqlens_q = cu_seqlens_q[num_decodes:] - cu_seqlens_q[num_decodes] - _forward_modelopt( - self, - query=query[prefill_start : prefill_start + num_prefill_tokens], - block_table=block_table[num_decodes : num_decodes + num_prefills], - seq_lens=seq_lens[num_decodes : num_decodes + num_prefills], - cu_seqlens_q=prefill_cu_seqlens_q, - num_actual_tokens=num_prefill_tokens, - max_query_len=max_query_len, - output=output[prefill_start : prefill_start + num_prefill_tokens], - **common_kw, + super().forward, + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output, + output_scale, + output_block_scale, ) - return output _FLASHINFER_IMPL_CLS = ModelOptSparseFlashInferImpl return _FLASHINFER_IMPL_CLS diff --git a/tests/gpu/torch/kernels/common/attention/test_decode_attention.py b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py index e50ba6ca360..490fcac9fd1 100644 --- a/tests/gpu/torch/kernels/common/attention/test_decode_attention.py +++ b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py @@ -82,7 +82,7 @@ def _nvfp4_qdq_reference(x, global_scale=1.0 / (6.0 * 448.0)): @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") -@pytest.mark.parametrize("num_kv_splits", [1, 8, 32]) +@pytest.mark.parametrize("num_kv_splits", [1, 32]) def test_split_k_varlen_gqa_matches_dense(num_kv_splits): torch.manual_seed(13) batch, num_q_heads, num_kv_heads, seq_len, head_dim = 2, 8, 2, 511, 64 @@ -111,17 +111,36 @@ def test_split_k_varlen_gqa_matches_dense(num_kv_splits): @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") @requires_native_e4m3 -def test_baked_v_prefix_and_pristine_tail_match_full_onread(): - seq_len, num_q_heads, num_kv_heads, head_dim = 17, 4, 1, 64 +@pytest.mark.parametrize( + ("cache_dtype", "num_q_heads", "head_dim", "value", "v_qdq_amax", "v_qdq_scale"), + [ + pytest.param(torch.float16, 4, 64, 0.017578125, None, 1.0, id="fp16-default-gqa"), + pytest.param( + torch.bfloat16, + 1, + 16, + 0.019, + 1.0, + 1.0 / (6.0 * 448.0), + id="bf16-custom-amax-carrier", + ), + ], +) +def test_baked_v_prefix_and_pristine_tail_match_full_onread( + cache_dtype, num_q_heads, head_dim, value, v_qdq_amax, v_qdq_scale +): + """Baked prefixes and pristine tails share the cache carrier at either V scale.""" + seq_len, num_kv_heads = 17, 1 q = torch.zeros(1, num_q_heads, head_dim, device="cuda", dtype=torch.float32) - k = torch.zeros(1, num_kv_heads, seq_len, head_dim, device="cuda", dtype=torch.float16) - v = torch.full_like(k, 0.017578125) + k = torch.zeros(1, num_kv_heads, seq_len, head_dim, device="cuda", dtype=cache_dtype) + v = torch.full_like(k, value) seq_lens = torch.tensor([seq_len], device="cuda", dtype=torch.int32) k_cache, raw_v_cache, block_table = _paged_cache(k, v, (seq_len,)) common = { "p_qdq": "nvfp4", - "v_qdq": "nvfp4", "num_kv_splits": 1, + "v_qdq": "nvfp4", + "v_qdq_amax": v_qdq_amax, } onread = attention_decode(q, k_cache, raw_v_cache, block_table, seq_lens, **common) @@ -132,6 +151,7 @@ def test_baked_v_prefix_and_pristine_tail_match_full_onread(): torch.zeros(1, device="cuda", dtype=torch.int32), torch.tensor([16], device="cuda", dtype=torch.int32), max_new_tokens=16, + v_qdq_scale=v_qdq_scale, ) baked_v_before = baked_v_cache.clone() baked = attention_decode( @@ -148,45 +168,6 @@ def test_baked_v_prefix_and_pristine_tail_match_full_onread(): torch.testing.assert_close(baked_v_cache, baked_v_before, rtol=0, atol=0) -@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") -@requires_native_e4m3 -def test_baked_v_prefix_uses_same_cache_carrier_as_tail(): - """A non-default V scale must not change values at the bake boundary.""" - seq_len, head_dim = 17, 16 - q = torch.zeros(1, 1, head_dim, device="cuda", dtype=torch.float32) - k = torch.zeros(1, 1, seq_len, head_dim, device="cuda", dtype=torch.bfloat16) - v = torch.full_like(k, 0.019) - seq_lens = torch.tensor([seq_len], device="cuda", dtype=torch.int32) - k_cache, raw_v_cache, block_table = _paged_cache(k, v, (seq_len,)) - common = { - "num_kv_splits": 1, - "v_qdq": "nvfp4", - "v_qdq_amax": 1.0, - } - - onread = attention_decode(q, k_cache, raw_v_cache, block_table, seq_lens, **common) - baked_v_cache = raw_v_cache.clone() - fake_quant_v_onwrite( - baked_v_cache, - block_table, - torch.zeros(1, device="cuda", dtype=torch.int32), - torch.tensor([16], device="cuda", dtype=torch.int32), - max_new_tokens=16, - v_qdq_scale=1.0 / (6.0 * 448.0), - ) - baked = attention_decode( - q, - k_cache, - baked_v_cache, - block_table, - seq_lens, - v_cache_quantized=True, - **common, - ) - - torch.testing.assert_close(baked, onread, rtol=0, atol=0) - - @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") @requires_native_e4m3 def test_p_quantizes_model_dtype_input_but_accumulates_fp32(): diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py index 40a0c5fdbe9..8ce41718d9a 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -55,11 +55,6 @@ def test_v_qdq_rejects_autograd_before_kernel_launch(monkeypatch): FP8_E4M3_MAX = 448.0 -@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need triton") -def test_p_qdq_block_n_matches_forward_configs(): - assert {config.kwargs["BLOCK_N"] for config in triton_fa._FWD_CONFIGS} == {P_QDQ_BLOCK_N} - - class _FixedBlockMForward: """Launch the raw forward kernel with one selected query tile.""" @@ -84,13 +79,9 @@ def launch(*args, **kwargs): @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need triton") -@pytest.mark.parametrize( - ("seq_len_q", "seq_len_kv", "is_causal"), - [(256, 256, True), (197, 233, True), (256, 256, False)], - ids=["causal", "chunked-prefill", "non-causal"], -) -def test_sparse_dense_window_is_block_m_invariant(monkeypatch, seq_len_q, seq_len_kv, is_causal): +def test_sparse_dense_window_is_block_m_invariant(monkeypatch): """The dense recent-token policy must not depend on the compute query tile.""" + seq_len_q, seq_len_kv = 197, 233 num_heads, head_dim = 2, 64 torch.manual_seed(20260706) q = torch.randn(seq_len_q, num_heads, head_dim, device="cuda", dtype=torch.float16) @@ -110,7 +101,7 @@ def test_sparse_dense_window_is_block_m_invariant(monkeypatch, seq_len_q, seq_le q_locs, q_lens, seq_len_q, - is_causal=is_causal, + is_causal=True, b_start_loc_k=kv_locs, b_seq_len_k=kv_lens, max_input_len_k=seq_len_kv, @@ -380,21 +371,6 @@ def test_varlen_partial_tiles(self, mode): atol = 2e-2 if mode == "nvfp4" else 5e-3 torch.testing.assert_close(o[s : s + n].float(), ref, rtol=5e-3, atol=atol) - @pytest.mark.parametrize(("mode", "tol"), [("fp8", 5e-2), ("nvfp4", 0.25)]) - @requires_native_e4m3 - def test_qdq_close_to_dense(self, mode, tol): - """Quantization is an approximation: output stays near dense attention.""" - seq_len, num_heads, num_kv_heads, head_dim = 128, 4, 2, 64 - scale = 1.0 / (head_dim**0.5) - - torch.manual_seed(13) - q, k, v = make_qkv(seq_len, num_heads, num_kv_heads, head_dim, dtype=torch.float16) - locs, lens = make_varlen_meta([seq_len]) - - o = attention(q, k, v, locs, lens, seq_len, softmax_scale=scale, p_qdq=mode) - ref = sdpa_reference(q, k, v, locs, lens) - torch.testing.assert_close(o, ref, rtol=tol, atol=tol) - @pytest.mark.parametrize("mode", ["fp8", "nvfp4"]) @requires_native_e4m3 def test_decode(self, mode): diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py index c7f68c27cdf..bdfbeefe47c 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py @@ -35,15 +35,6 @@ ) -@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") -def test_v_onwrite_prefill_grid_uses_host_metadata(): - signature = inspect.signature(fake_quant_v_onwrite) - source = inspect.getsource(fake_quant_v_onwrite) - - assert "max_new_tokens" in signature.parameters - assert ".item()" not in source - - @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need triton") @pytest.mark.parametrize("loader", ["_load_paged_k_tile", "_load_paged_v_tile"]) def test_paged_loaders_widen_page_ids_before_pointer_math(loader): @@ -141,7 +132,7 @@ class TestPagedKV: @pytest.mark.parametrize( ("page_size", "v_qdq_scale", "match"), - [(8, 1.0, "page_size"), (16, 0.0, "v_qdq_scale"), (16, float("inf"), "v_qdq_scale")], + [(8, 1.0, "page_size"), (16, 0.0, "v_qdq_scale")], ) def test_v_onwrite_validates_page_size_and_scale(self, page_size, v_qdq_scale, match): cache = torch.empty(1, 16, 1, 1, device="cuda") @@ -157,21 +148,34 @@ def test_v_onwrite_validates_page_size_and_scale(self, page_size, v_qdq_scale, m v_qdq_scale=v_qdq_scale, ) - def test_paged_matches_contiguous(self): - """Paged mode produces same output as contiguous mode with identical data.""" + @pytest.mark.parametrize( + ("seq_len", "seed", "attention_kwargs"), + [ + pytest.param(128, 42, {}, id="dense"), + pytest.param( + 256, + 99, + {"sparsity_n": 2, "sparsity_m": 4, "dense_recent_tokens": 64}, + id="sparse-2-4", + ), + ], + ) + def test_paged_matches_contiguous(self, seq_len, seed, attention_kwargs): + """Paged dense and sparse prefill match their contiguous counterparts.""" batch = 2 - seq_len = 128 num_heads, num_kv_heads, head_dim = 4, 2, 64 page_size = 16 scale = 1.0 / (head_dim**0.5) total = batch * seq_len - torch.manual_seed(42) + torch.manual_seed(seed) q, k, v = make_qkv(total, num_heads, num_kv_heads, head_dim) locs, lens = make_varlen_meta([seq_len] * batch) # Contiguous reference - out_contig = attention(q, k, v, locs, lens, seq_len, softmax_scale=scale) + out_contig = attention( + q, k, v, locs, lens, seq_len, softmax_scale=scale, **attention_kwargs + ) # Build paged cache from the same K/V locs_k, lens_k = locs, lens @@ -195,46 +199,11 @@ def test_paged_matches_contiguous(self): v_cache=v_cache, block_table=block_table, page_size=page_size, + **attention_kwargs, ) torch.testing.assert_close(out_paged, out_contig, rtol=1e-2, atol=1e-2) - def test_paged_no_nan(self): - """Paged mode output is finite.""" - batch = 2 - seq_len = 256 - num_heads, num_kv_heads, head_dim = 4, 2, 64 - page_size = 16 - scale = 1.0 / (head_dim**0.5) - total = batch * seq_len - - torch.manual_seed(55) - q, k, v = make_qkv(total, num_heads, num_kv_heads, head_dim) - locs, lens = make_varlen_meta([seq_len] * batch) - - k_cache, v_cache, block_table = _scatter_to_paged_cache( - k, v, locs, lens, num_kv_heads, head_dim, page_size - ) - - out = attention( - q, - k, - v, - locs, - lens, - seq_len, - softmax_scale=scale, - b_seq_len_k=lens, - max_input_len_k=seq_len, - k_cache=k_cache, - v_cache=v_cache, - block_table=block_table, - page_size=page_size, - ) - - assert not torch.isnan(out).any(), "NaN in paged output" - assert not torch.isinf(out).any(), "Inf in paged output" - def test_paged_variable_length(self): """Paged mode works with variable-length sequences.""" seq_lens = [64, 128] @@ -310,93 +279,6 @@ def test_paged_different_page_sizes(self, page_size): torch.testing.assert_close(out_paged, out_contig, rtol=1e-2, atol=1e-2) - def test_paged_with_sparsity(self): - """Paged mode works with N:M sparsity enabled.""" - batch = 2 - seq_len = 256 - num_heads, num_kv_heads, head_dim = 4, 2, 64 - page_size = 16 - scale = 1.0 / (head_dim**0.5) - total = batch * seq_len - - torch.manual_seed(99) - q, k, v = make_qkv(total, num_heads, num_kv_heads, head_dim) - locs, lens = make_varlen_meta([seq_len] * batch) - - k_cache, v_cache, block_table = _scatter_to_paged_cache( - k, v, locs, lens, num_kv_heads, head_dim, page_size - ) - - out_paged_sparse = attention( - q, - k, - v, - locs, - lens, - seq_len, - softmax_scale=scale, - b_seq_len_k=lens, - max_input_len_k=seq_len, - k_cache=k_cache, - v_cache=v_cache, - block_table=block_table, - page_size=page_size, - sparsity_n=2, - sparsity_m=4, - ) - - assert not torch.isnan(out_paged_sparse).any(), "NaN in paged + sparse output" - assert not torch.isinf(out_paged_sparse).any(), "Inf in paged + sparse output" - assert out_paged_sparse.shape == q.shape - - def test_paged_decode(self): - """Paged mode works for decode (single Q token, long KV context).""" - batch = 2 - seq_lens_k = [64, 128] - num_heads, num_kv_heads, head_dim = 4, 2, 64 - page_size = 16 - scale = 1.0 / (head_dim**0.5) - total_kv = sum(seq_lens_k) - - torch.manual_seed(33) - q_flat = torch.randn(batch, num_heads, head_dim, device="cuda", dtype=torch.float16) - k_flat = torch.randn(total_kv, num_kv_heads, head_dim, device="cuda", dtype=torch.float16) - v_flat = torch.randn(total_kv, num_kv_heads, head_dim, device="cuda", dtype=torch.float16) - - b_start_loc_q = torch.arange(batch, device="cuda", dtype=torch.int32) - b_seq_len_q = torch.ones(batch, device="cuda", dtype=torch.int32) - cumsum = [0] - for sl in seq_lens_k: - cumsum.append(cumsum[-1] + sl) - b_start_loc_k = torch.tensor(cumsum[:-1], device="cuda", dtype=torch.int32) - b_seq_len_k = torch.tensor(seq_lens_k, device="cuda", dtype=torch.int32) - - # Build paged cache - k_cache, v_cache, block_table = _scatter_to_paged_cache( - k_flat, v_flat, b_start_loc_k, b_seq_len_k, num_kv_heads, head_dim, page_size - ) - - out = attention( - q_flat, - k_flat, - v_flat, - b_start_loc_q, - b_seq_len_q, - 1, - is_causal=False, - softmax_scale=scale, - b_start_loc_k=b_start_loc_k, - b_seq_len_k=b_seq_len_k, - max_input_len_k=max(seq_lens_k), - k_cache=k_cache, - v_cache=v_cache, - block_table=block_table, - page_size=page_size, - ) - - assert out.shape == q_flat.shape - assert not torch.isnan(out).any(), "NaN in paged decode output" - def test_paged_decode_ignores_sparse_nm(self): """N:M sparse softmax is prefill-only; paged decode remains dense.""" batch = 2 @@ -615,11 +497,10 @@ def test_v_cache_matches_independent_signed_key_axis_oracle(self): torch.testing.assert_close(actual, expected, rtol=0, atol=0) torch.testing.assert_close(cache[1], unused_page, rtol=0, atol=0) - @pytest.mark.parametrize("q_len", [1, 7]) @requires_native_e4m3 - def test_baked_prefix_and_raw_tail_match_full_onread(self, q_len): - """The same paged Triton path handles pure decode and chunked prefill.""" - seq_len, num_heads, head_dim, page_size = 17, 2, 32, 16 + def test_baked_prefix_and_raw_tail_match_full_onread(self): + """The paged Triton chunked-prefill path handles a baked prefix and raw tail.""" + q_len, seq_len, num_heads, head_dim, page_size = 7, 17, 2, 32, 16 q = torch.zeros(q_len, num_heads, head_dim, device="cuda", dtype=torch.float32) k = torch.zeros(seq_len, 1, head_dim, device="cuda", dtype=torch.float16) v = torch.full_like(k, 0.017578125) diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py index dcc447f4f82..a8011805daf 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -82,25 +82,13 @@ def _nvfp4_quantizer(*, block_size=16, enabled=True): return quantizer -def test_import_attention_module_resolves_installed_concrete_attention(): - attention_module = vllm_plugin._import_attention_module() - assert isinstance(attention_module.Attention, type) - assert vllm_plugin.vllm_attention.Attention is attention_module.Attention - - -@pytest.mark.parametrize( - "wrapper_name", - ["_QuantVLLMAttention", "_QuantVLLMCrossAttention", "_QuantVLLMEncoderOnlyAttention"], -) -def test_attention_setup_keeps_qkv_only_checkpoint_surface(monkeypatch, wrapper_name): +def test_attention_setup_keeps_qkv_only_checkpoint_surface(monkeypatch): monkeypatch.setattr( vllm_plugin, "create_parallel_state", lambda: vllm_plugin.ParallelState(data_parallel_group=None), ) - wrapper = getattr(vllm_plugin, wrapper_name) - test_wrapper = type(f"Test{wrapper.__name__}", (wrapper, _NativeAttention), {}) - attention = _new_attention(test_wrapper) + attention = _new_attention(_TestQuantVLLMAttention) attention._setup() diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index d940f64c725..72e50753600 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -27,10 +27,13 @@ import torch import vllm from torch import nn +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends import flashinfer as flashinfer_backend from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl from vllm.v1.attention.backends.flashinfer import ( FlashInferBackend, FlashInferImpl, + FlashInferMetadata, FlashInferMetadataBuilder, ) from vllm.v1.worker.gpu_worker import Worker as BaseWorker @@ -94,8 +97,19 @@ def guarded_import(name, *args, **kwargs): monkeypatch.setattr(vllm, "__version__", "0.9.0") monkeypatch.setattr(builtins, "__import__", guarded_import) worker_module = _load_worker_module("sparse_attn_worker_import_test") + attention = _bare_attention(worker_module) + state = _sparse_worker_state(nn.ModuleDict({"attn": attention})) + monkeypatch.setattr(worker_module, "load_from_checkpoint_metadata", lambda _: ({}, "test")) + monkeypatch.setattr( + worker_module, + "match_sparse_config", + lambda *_: {"enable": True, "sparsity_n": 2, "sparsity_m": 4}, + ) + + worker_module._install_attention(state, quantize=False) assert worker_module.__all__ == ["SparseAttnWorker", "QuantSparseAttnWorker"] + assert type(attention.impl) is ModelOptSparseAttentionImpl @pytest.mark.parametrize( @@ -117,15 +131,6 @@ def test_public_workers_install_after_base_load(monkeypatch, class_name, quantiz assert events == ["base", ("install", quantize)] -def test_sparse_worker_keeps_base_memory_profile(): - worker_module = _load_worker_module("sparse_profile_test") - - assert ( - worker_module.SparseAttnWorker.determine_available_memory - is BaseWorker.determine_available_memory - ) - - def _make_old_impl(): """Create a vLLM FlashAttention impl with initialized runtime state.""" return FlashAttentionImpl( @@ -190,47 +195,66 @@ def _make_flashinfer_impl(*, sparse=False, quantized=False): return impl -def test_flashinfer_metadata_builder_patch_stashes_common_metadata(monkeypatch): - """The wrapper must bind the current vLLM signature without shifting arguments.""" - calls = [] +@pytest.fixture +def isolated_flashinfer_builder_patch(): + saved_build = FlashInferMetadataBuilder.build + vllm_plugin._reset_flashinfer_state_for_tests() + try: + yield + finally: + FlashInferMetadataBuilder.build = saved_build + vllm_plugin._reset_flashinfer_state_for_tests() - def fake_build(self, common_prefix_len, common_attn_metadata, fast_build=False): - calls.append((common_prefix_len, common_attn_metadata, fast_build)) - return SimpleNamespace() - monkeypatch.setattr(FlashInferMetadataBuilder, "build", fake_build) - monkeypatch.setattr(vllm_plugin, "_FLASHINFER_PATCHED", False) - common = SimpleNamespace( - block_table_tensor=object(), - seq_lens=object(), - query_start_loc=object(), - num_actual_tokens=7, - max_query_len=3, - max_seq_len=11, +def test_flashinfer_metadata_builder_patch_stashes_common_metadata( + monkeypatch, isolated_flashinfer_builder_patch +): + """The real builder result must retain the common metadata contract.""" + monkeypatch.setattr(flashinfer_backend, "use_trtllm_attention", lambda *_args, **_kwargs: True) + assert patch_flashinfer_metadata_builder() is True + builder = object.__new__(FlashInferMetadataBuilder) + builder.__dict__.update( + reorder_batch_threshold=1, + page_size=16, + num_qo_heads=2, + num_kv_heads=2, + dcp_world_size=1, + cache_dtype=torch.float16, + q_data_type=torch.float16, + attention_config=SimpleNamespace(use_trtllm_attention=True), + has_sinks=False, + use_trtllm_decode_attention=True, + use_dcp=False, + ) + common = CommonAttentionMetadata( + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32), + seq_lens=torch.tensor([16], dtype=torch.int32), + num_reqs=1, + num_actual_tokens=1, + max_query_len=1, + max_seq_len=16, + block_table_tensor=torch.tensor([[0]], dtype=torch.int32), + slot_mapping=torch.tensor([0], dtype=torch.int64), causal=False, ) - assert patch_flashinfer_metadata_builder() is True - patched_build = FlashInferMetadataBuilder.build - metadata = patched_build(object(), 5, common, fast_build=True) - assert calls == [(5, common, True)] - assert metadata._modelopt_block_table is common.block_table_tensor - assert metadata._modelopt_seq_lens is common.seq_lens - assert metadata._modelopt_query_start_loc is common.query_start_loc - assert metadata._modelopt_num_actual_tokens == 7 - assert metadata._modelopt_max_query_len == 3 - assert metadata._modelopt_max_seq_len == 11 - assert metadata._modelopt_causal is False + metadata = builder.build(0, common, fast_build=False) - assert patch_flashinfer_metadata_builder() is True - assert FlashInferMetadataBuilder.build is patched_build + assert isinstance(metadata, FlashInferMetadata) + for target, source in vllm_plugin._FLASHINFER_METADATA_FIELDS.items(): + actual = getattr(metadata, target) + expected = getattr(common, source) + if isinstance(expected, torch.Tensor): + assert actual is expected + else: + assert actual == expected -def test_select_and_clone_flashinfer_impl_preserves_runtime_state(monkeypatch): - monkeypatch.setattr(vllm_plugin, "_FLASHINFER_PATCHED", False) - monkeypatch.setattr(vllm_plugin, "_FLASHINFER_IMPL_CLS", None) +def test_select_and_clone_flashinfer_impl_preserves_runtime_state( + isolated_flashinfer_builder_patch, +): old_impl = _make_old_flashinfer_impl() - old_impl.future_attr = object() new_cls = select_sparse_impl_cls(old_impl) new_impl = _clone_sparse_impl(old_impl, new_cls) @@ -238,8 +262,6 @@ def test_select_and_clone_flashinfer_impl_preserves_runtime_state(monkeypatch): assert new_cls is get_flashinfer_sparse_impl_cls() assert isinstance(new_impl, FlashInferImpl) assert type(new_impl).__name__ == "ModelOptSparseFlashInferImpl" - assert new_impl.future_attr is old_impl.future_attr - assert new_impl.__dict__.items() >= old_impl.__dict__.items() if hasattr(FlashInferImpl, "do_kv_cache_update"): assert type(new_impl).do_kv_cache_update is FlashInferImpl.do_kv_cache_update assert select_sparse_impl_cls(new_impl) is None @@ -503,12 +525,21 @@ def fake_attention(query, **kwargs): def test_flashinfer_forced_kernel_mixed_batch_never_uses_native_fallback(monkeypatch): + class Layer: + force_kernel_reads = 0 + + @property + def _modelopt_force_kernel(self): + self.force_kernel_reads += 1 + assert self.force_kernel_reads == 1, "force-kernel state was resolved more than once" + return True + prefill_tokens = 17 impl = _make_flashinfer_impl(sparse=False, quantized=False) query = torch.zeros(1 + prefill_tokens, 2, 64, dtype=torch.float16) kv_cache = torch.zeros(4, 2, 16, 2, 64, dtype=torch.float16) metadata = _flashinfer_metadata(query_lens=(1, prefill_tokens), mixed=True) - layer = SimpleNamespace(_modelopt_force_kernel=True) + layer = Layer() calls = {"native": 0, "decode": 0, "prefill": 0} def native_fallback(*args, **kwargs): @@ -531,6 +562,36 @@ def fake_attention(query, **kwargs): impl.forward(layer, query, query, query, kv_cache, metadata, output=torch.empty_like(query)) assert calls == {"native": 0, "decode": 1, "prefill": 1} + assert layer.force_kernel_reads == 1 + + +def test_flashinfer_invalid_mixed_metadata_has_no_side_effects(monkeypatch): + impl = _make_flashinfer_impl(sparse=True) + metadata = _flashinfer_metadata(query_lens=(1, 17), mixed=True) + metadata._modelopt_num_actual_tokens += 1 + query = torch.zeros(18, 2, 64, dtype=torch.float16) + kv_cache = torch.zeros(4, 2, 16, 2, 64, dtype=torch.float16) + + def unexpected_side_effect(*args, **kwargs): + pytest.fail("invalid metadata triggered attention side effects") + + monkeypatch.setattr(FlashInferImpl, "forward", unexpected_side_effect) + for name in ("_flashinfer_cache_write", "triton_attention"): + monkeypatch.setattr(vllm_plugin, name, unexpected_side_effect) + + with pytest.raises( + ValueError, + match="FlashInfer mixed-batch token counts do not match common metadata", + ): + impl.forward( + SimpleNamespace(), + query, + query, + query, + kv_cache, + metadata, + output=torch.empty_like(query), + ) def test_sparse_worker_is_noop_without_checkpoint_metadata(monkeypatch, capsys): @@ -573,23 +634,6 @@ def test_sparse_worker_installs_only_active_layer(monkeypatch): assert empty.impl is original_empty_impl -def test_sparse_worker_does_not_load_quant_api(monkeypatch): - worker_module = _load_worker_module() - attention = _bare_attention(worker_module) - state = _sparse_worker_state(nn.ModuleDict({"attn": attention})) - monkeypatch.setattr(worker_module, "load_from_checkpoint_metadata", lambda _: ({}, "test")) - monkeypatch.setattr( - worker_module, - "match_sparse_config", - lambda *_: {"enable": True, "sparsity_n": 2, "sparsity_m": 4}, - ) - monkeypatch.setattr(worker_module, "_quant_api", lambda: pytest.fail("loaded quant API")) - - worker_module._install_attention(state, quantize=False) - - assert type(attention.impl) is ModelOptSparseAttentionImpl - - def test_sparse_worker_rejects_unsupported_backend_before_mutation(monkeypatch): worker_module = _load_worker_module() good = _bare_attention(worker_module) @@ -620,15 +664,32 @@ def test_sparse_worker_rejects_unsupported_backend_before_mutation(monkeypatch): @pytest.mark.parametrize( - ("failure", "active"), - [("cascade", True), ("cascade", False), ("metadata", True), ("metadata", False)], + ("failure", "active", "expected"), + [ + ( + "cascade", + True, + "vLLM cascade attention is incompatible with an active ModelOpt attention transform", + ), + ("cascade", False, None), + ( + "metadata", + True, + "FlashInfer metadata is missing the ModelOpt attention transform fields: " + + ", ".join(vllm_plugin._FLASHINFER_METADATA_FIELDS), + ), + ("metadata", False, None), + ("output", True, "Fused attention output quantization is unsupported"), + ], ) -def test_flashinfer_transform_safety_for_unsupported_metadata(monkeypatch, failure, active): +def test_flashinfer_transform_safety_for_unsupported_metadata( + monkeypatch, failure, active, expected +): impl = _make_flashinfer_impl(sparse=active) metadata = ( - SimpleNamespace(use_cascade=True) - if failure == "cascade" - else SimpleNamespace(use_cascade=False) + _flashinfer_metadata() + if failure == "output" + else SimpleNamespace(use_cascade=failure == "cascade") ) native_calls = [] query = torch.zeros(1, 2, 64, dtype=torch.float16) @@ -643,12 +704,31 @@ def fake_forward(self, *args, **kwargs): monkeypatch.setattr(FlashInferImpl, "forward", fake_forward) if active: - with pytest.raises(NotImplementedError, match="ModelOpt attention transform"): - impl.forward(None, query, query, query, torch.empty(0), metadata, output=output) + with pytest.raises(NotImplementedError) as exc: + impl.forward( + None, + query, + query, + query, + torch.empty(0), + metadata, + output=output, + output_scale=torch.ones(1), + ) + assert str(exc.value) == expected assert native_calls == [] else: assert ( - impl.forward(None, query, query, query, torch.empty(0), metadata, output=output) + impl.forward( + None, + query, + query, + query, + torch.empty(0), + metadata, + output=output, + output_scale=torch.ones(1), + ) is output ) assert native_calls == [(impl, metadata)] @@ -978,16 +1058,6 @@ def fake_attention(query, **kwargs): assert captured["v_qdq"] == "nvfp4" -def test_active_cascade_fails_loud(): - """Cascade must not silently drop an active ModelOpt transform.""" - impl = _clone_sparse_impl(_make_old_impl()) - impl.sparse_kw = {"skip_softmax_threshold": 0.001} - q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16) - cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) - with pytest.raises(NotImplementedError, match="cascade"): - impl.forward(None, q, q, q, cache, SimpleNamespace(use_cascade=True), output=q) - - def test_resolve_calibrated_skip_softmax_threshold_for_decode(): """Calibration conversion is phase-aware even when decode later delegates.""" sparse_kw = { diff --git a/tests/unit/torch/kernels/common/attention/test_triton_fa.py b/tests/unit/torch/kernels/common/attention/test_triton_fa.py index 941d11c441c..bb170eaa8d9 100644 --- a/tests/unit/torch/kernels/common/attention/test_triton_fa.py +++ b/tests/unit/torch/kernels/common/attention/test_triton_fa.py @@ -48,15 +48,6 @@ def __getitem__(self, grid): raise AssertionError("unexpected kernel launch") -def _config_signature(config): - return ( - config.kwargs["BLOCK_M"], - config.kwargs["BLOCK_N"], - config.num_stages, - config.num_warps, - ) - - def test_triton_fa_importable_on_cpu(): """Module imports cleanly without CUDA; exports the public API names.""" try: @@ -101,32 +92,23 @@ def test_forward_uses_minimal_shared_autotune_configs(): from modelopt.torch.kernels.common.attention import triton_fa - expected_configs = [ - (16, 32, 2, 4), - (64, 32, 2, 4), - (128, 32, 2, 4), + configs = triton_fa._FWD_CONFIGS + assert [(config.kwargs["BLOCK_M"], config.kwargs["BLOCK_N"]) for config in configs] == [ + (16, 32), + (64, 32), + (128, 32), ] - assert [_config_signature(config) for config in triton_fa._FWD_CONFIGS] == expected_configs - - expected_active_configs = ( - expected_configs[:1] if "PYTEST_VERSION" in os.environ else expected_configs - ) - assert [ - _config_signature(config) for config in triton_fa._attn_fwd.configs - ] == expected_active_configs + + expected_active_configs = configs[:1] if "PYTEST_VERSION" in os.environ else configs + assert triton_fa._attn_fwd.configs == expected_active_configs assert triton_fa._attn_fwd.keys == ["N_CTX", "HEAD_DIM", "Q_IS_FP32", "P_QDQ", "V_QDQ"] - assert "_attn_fwd_p_qdq" not in vars(triton_fa) - assert "_prune_fwd_configs" not in vars(triton_fa) @pytest.mark.parametrize( ("attention_kwargs", "expected_p_qdq", "expected_v_qdq"), [ ({}, 0, 0), - ({"p_qdq": "fp8"}, 1, 0), - ({"p_qdq": "nvfp4"}, 2, 0), - ({"v_qdq": "fp8"}, 0, 1), - ({"v_qdq": "nvfp4"}, 0, 2), + ({"p_qdq": "fp8", "v_qdq": "fp8"}, 1, 1), ({"p_qdq": "nvfp4", "v_qdq": "nvfp4"}, 2, 2), ({"p_qdq": "nvfp4", "sparsity_n": 2, "sparsity_m": 4}, 2, 0), ({"p_qdq": "nvfp4", "skip_softmax_threshold": 0.1}, 2, 0), From 37d1b178ab79b2dd5153e116a75f940cc9b02be9 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 7 Jul 2026 20:10:42 -0700 Subject: [PATCH 20/28] Remove redundant NVFP4 attention paths Signed-off-by: Kai Xu --- examples/vllm_serve/sparse_attn_worker.py | 16 +---- .../kernels/common/attention/triton_fa.py | 37 +++++------ .../kernels/quantization/attention/v_qdq.py | 11 ++-- .../quantization/gemm/fp4_kernel_hopper.py | 1 + .../attention_sparsity/plugins/vllm.py | 34 ++-------- .../common/attention/test_triton_fa_p_qdq.py | 1 + .../common/attention/test_triton_fa_paged.py | 55 ---------------- .../test_quant_sparse_attn_worker.py | 23 ------- .../test_sparse_attn_worker.py | 66 ++----------------- .../common/attention/test_triton_fa.py | 6 +- 10 files changed, 38 insertions(+), 212 deletions(-) diff --git a/examples/vllm_serve/sparse_attn_worker.py b/examples/vllm_serve/sparse_attn_worker.py index 35b75d476ed..40a04585265 100644 --- a/examples/vllm_serve/sparse_attn_worker.py +++ b/examples/vllm_serve/sparse_attn_worker.py @@ -16,7 +16,6 @@ """Custom vLLM workers for checkpoint-driven sparse and fixed-NVFP4 attention.""" import importlib -import os from collections import Counter from functools import cache from types import SimpleNamespace @@ -291,21 +290,12 @@ def _install_sparse_plans(plans) -> None: def _install_quant_plans(worker, plans) -> None: api = _quant_api() - quant_off = os.environ.get("MODELOPT_ATTN_QUANT_OFF") == "1" for plan in plans: module = plan.module module.device, module.dtype = plan.device, plan.dtype api.nn.QuantModuleRegistry.convert(module) module.p_bmm_quantizer = api.nn.TensorQuantizer() api.conversion.set_quantizer_by_cfg(module, _BMM_CFG) - if quant_off: - # Isolation knob: keep the ModelOpt kernel fixed while disabling all - # Q/K/P/V NVFP4 transforms, so quant-on minus quant-off isolates the - # fakequant loss without triggering the native dense fallback. - for name in ("q", "k", "p", "v"): - quantizer = getattr(module, f"{name}_bmm_quantizer", None) - if quantizer is not None: - quantizer.disable() api.plugin._set_vllm_attention_kv_default_amax(module, plan.device) plan.new_impl.sparse_kw = plan.sparse_kw p_qdq, p_qdq_amax = _p_qdq_from_layer(module) @@ -317,10 +307,8 @@ def _install_quant_plans(worker, plans) -> None: "v_qdq_amax": v_qdq_amax, } module.impl = plan.new_impl - module._query_quant_in_kernel = not quant_off - module._value_quant_in_kernel = not quant_off - if quant_off: - module._modelopt_force_kernel = True + module._query_quant_in_kernel = True + module._value_quant_in_kernel = True worker.model_runner.cascade_attn_enabled = False installed = dict(Counter(type(plan.new_impl).__name__ for plan in plans)) print(f"[ModelOpt] Installed NVFP4 quant+sparse attention on {len(plans)} layers: {installed}") diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 898d1550127..1d2d5c1145d 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -71,8 +71,9 @@ def _load_qdq_helpers() -> None: _v_qdq_nvfp4 = _v_nvfp4 -# Maps public P/V QDQ options to kernel constexpr values. -_QDQ_MODES = {None: 0, "fp8": 1, "nvfp4": 2} +# Maps public QDQ options to kernel constexpr values. +_P_QDQ_MODES = {None: 0, "fp8": 1, "nvfp4": 2} +_V_QDQ_MODES = {None: 0, "nvfp4": 2} LOG2E: float = 1.44269504088896 @@ -289,7 +290,7 @@ def _attn_fwd( SKIP_THRESHOLD_LOG2: tl.constexpr = 0.0, # log2(lambda) in the kernel's scaled log2 score space P_QDQ: tl.constexpr = 0, # Fake quant-dequant of softmax P: 0=off, 1=FP8 E4M3, 2=NVFP4 p_qdq_scale=1.0, # Per-tensor scale for softmax qdq (runtime scalar; amax/448 or amax/(6*448)) - V_QDQ: tl.constexpr = 0, # Fake quant-dequant of V: 0=off, 1=FP8 E4M3, 2=NVFP4 + V_QDQ: tl.constexpr = 0, # Fake quant-dequant of V: 0=off, 2=NVFP4 v_qdq_scale=1.0, V_CACHE_QUANTIZED: tl.constexpr = False, # complete block-16 groups are already QDQ Sparsity_total=None, # Optional int64 scalar for counting total tiles (atomic) @@ -476,13 +477,10 @@ def _attn_fwd( mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :], other=0.0, ) - if V_QDQ != 0 and ( + if V_QDQ == 2 and ( (not V_CACHE_QUANTIZED) or (kv_start + BLOCK_N > v_quantized_boundary) ): - if V_QDQ == 1: - v_qdq = _qdq_fp8(v.to(tl.float32), v_qdq_scale) - else: - v_qdq = _v_qdq_nvfp4(v.to(tl.float32), v_qdq_scale, BLOCK_N, BLOCK_D) + v_qdq = _v_qdq_nvfp4(v.to(tl.float32), v_qdq_scale, BLOCK_N, BLOCK_D) v_qdq = v_qdq.to(v.dtype) if V_CACHE_QUANTIZED: use_qdq = (kv_start + kv_pos) >= v_quantized_boundary @@ -1298,12 +1296,11 @@ def attention( and the global scale ``amax / (6 * 448)`` for NVFP4. A runtime scalar — user-set or calibrated values do not recompile the kernel. Values above amax saturate. - v_qdq: Fake quant-dequant of V before ``P @ V``. ``"fp8"`` uses a - per-tensor E4M3 scale; ``"nvfp4"`` uses signed E2M1 values with - one E4M3 scale per 16 keys. ``None`` disables V QDQ. + v_qdq: Fake quant-dequant of V before ``P @ V``. ``"nvfp4"`` uses + signed E2M1 values with one E4M3 scale per 16 keys. ``None`` + disables V QDQ. v_qdq_amax: Optional per-tensor V amax. ``None`` uses global scale 1; - otherwise converts to ``amax / 448`` (FP8) or ``amax / (6 * 448)`` - (NVFP4). + otherwise converts to the NVFP4 global scale ``amax / (6 * 448)``. v_cache_quantized: Complete block-16 groups in the paged V cache are already QDQ; only the pristine partial group is QDQ on read. k_cache: Paged K cache [num_blocks, page_size, num_kv_heads, head_dim]. @@ -1330,11 +1327,11 @@ def attention( # silently reuse stale compiled kernels from the on-disk cache. _load_sparsity_helpers() _load_qdq_helpers() - if p_qdq not in _QDQ_MODES: + if p_qdq not in _P_QDQ_MODES: raise ValueError( - f"p_qdq must be one of {sorted(k for k in _QDQ_MODES if k)} or None, got {p_qdq!r}" + f"p_qdq must be one of {sorted(k for k in _P_QDQ_MODES if k)} or None, got {p_qdq!r}" ) - p_qdq_mode = _QDQ_MODES[p_qdq] + p_qdq_mode = _P_QDQ_MODES[p_qdq] # Convert the per-tensor amax to the kernel's scale convention # (``q = cast(p / scale) * scale``): FP8 uses ``amax / 448``; NVFP4 uses the # global scale ``amax / (6 * 448)``. amax=1 (the default, the theoretical @@ -1344,18 +1341,18 @@ def attention( if not (math.isfinite(p_qdq_amax) and p_qdq_amax > 0): raise ValueError(f"p_qdq_amax must be a finite positive value, got {p_qdq_amax}") p_qdq_scale = p_qdq_amax / 448.0 if p_qdq == "fp8" else p_qdq_amax / (6.0 * 448.0) - if v_qdq not in _QDQ_MODES: + if v_qdq not in _V_QDQ_MODES: raise ValueError( - f"v_qdq must be one of {sorted(k for k in _QDQ_MODES if k)} or None, got {v_qdq!r}" + f"v_qdq must be one of {sorted(k for k in _V_QDQ_MODES if k)} or None, got {v_qdq!r}" ) - v_qdq_mode = _QDQ_MODES[v_qdq] + v_qdq_mode = _V_QDQ_MODES[v_qdq] if v_qdq_mode and any(t.requires_grad for t in (q, k, v)): raise NotImplementedError("v_qdq is inference-only and does not support autograd") v_qdq_scale = 1.0 if v_qdq_mode and v_qdq_amax is not None: if not (math.isfinite(v_qdq_amax) and v_qdq_amax > 0): raise ValueError(f"v_qdq_amax must be a finite positive value, got {v_qdq_amax}") - v_qdq_scale = v_qdq_amax / 448.0 if v_qdq == "fp8" else v_qdq_amax / (6.0 * 448.0) + v_qdq_scale = v_qdq_amax / (6.0 * 448.0) if v_cache_quantized and v_qdq != "nvfp4": raise ValueError("v_cache_quantized requires v_qdq='nvfp4'") if v_cache_quantized and any(x is None for x in (k_cache, v_cache, block_table)): diff --git a/modelopt/torch/kernels/quantization/attention/v_qdq.py b/modelopt/torch/kernels/quantization/attention/v_qdq.py index edbd0e0f2ec..8772fa3c605 100644 --- a/modelopt/torch/kernels/quantization/attention/v_qdq.py +++ b/modelopt/torch/kernels/quantization/attention/v_qdq.py @@ -86,15 +86,14 @@ def fake_quant_v_onwrite( max_new_tokens: int, page_size: int = 16, v_qdq_scale: float = 1.0, - decode: bool = False, ) -> None: """NVFP4-finalize complete block-16 groups in ``[v_lo, v_hi)`` in place. ``max_new_tokens`` is host metadata used to size the masked launch grid. - Decode uses one fixed group per request. Eager prefill covers every group - that the largest query chunk can complete without reading device metadata. - ``v_lo`` and ``v_hi`` must describe aligned, completed block-16 boundaries; - their device values are not host-validated. + The grid covers every group that the largest query chunk can complete + without reading device metadata. ``v_lo`` and ``v_hi`` must describe + aligned, completed block-16 boundaries; their device values are not + host-validated. """ if page_size != v_cache.shape[1]: raise ValueError(f"page_size {page_size} must match v_cache.shape[1] {v_cache.shape[1]}") @@ -105,7 +104,7 @@ def fake_quant_v_onwrite( batch, max_blocks = block_table.shape num_kv_heads, head_dim = v_cache.shape[2:] - num_groups = 1 if decode else triton.cdiv(max_new_tokens, _BLOCK_N) + num_groups = triton.cdiv(max_new_tokens, _BLOCK_N) with torch.cuda.device(v_cache.device): _fake_quant_v_onwrite_kernel[(batch, num_kv_heads, num_groups)]( diff --git a/modelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.py b/modelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.py index a3686a4fa65..742c58a3969 100644 --- a/modelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.py +++ b/modelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.py @@ -81,6 +81,7 @@ def fp4_fake_quant_kernel( block_max = tl.max(x_abs, axis=2, keep_dims=True) block_max_quant = fp8_quantize_scale(block_max, global_scale_safe) + block_max_quant = tl.where(block_max_quant >= 1e-5, block_max_quant, 1.0) block_max_quant_broadcast = tl.broadcast_to( block_max_quant, (TILE_M, NUM_FP4_BLOCKS, BLOCK_SIZE) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index e15c215b522..65167e2d1bc 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -171,14 +171,9 @@ def _any_quant_active(layer, p_qdq, v_qdq) -> bool: ) -def _should_run_modelopt_kernel(sparse_kw, quant_active: bool, force_kernel: bool) -> bool: +def _should_run_modelopt_kernel(sparse_kw, quant_active: bool) -> bool: """Return whether a launch has effective work for the ModelOpt kernel.""" - return bool(sparse_kw or quant_active or force_kernel) - - -def _mixed_batch_needs_dense_base(quant_active: bool, force_kernel: bool) -> bool: - """Return whether sparse-only mixed phases need one native output seed.""" - return not quant_active and not force_kernel + return bool(sparse_kw or quant_active) @dataclass(frozen=True, slots=True) @@ -188,7 +183,6 @@ class _ResolvedForward: v_qdq: str | None v_qdq_amax: float | None quant_active: bool - force_kernel: bool def _resolve_forward( @@ -203,10 +197,7 @@ def _resolve_forward( """Resolve shared transform state or request the backend's native path.""" p_qdq, p_qdq_amax, v_qdq, v_qdq_amax = _quant_kw_from_impl(impl, layer) quant_active = _any_quant_active(layer, p_qdq, v_qdq) - force_kernel = getattr(layer, "_modelopt_force_kernel", False) - transform_active = _should_run_modelopt_kernel( - getattr(impl, "sparse_kw", None), quant_active, force_kernel - ) + transform_active = _should_run_modelopt_kernel(getattr(impl, "sparse_kw", None), quant_active) if getattr(attn_metadata, "use_cascade", False): if transform_active: @@ -234,7 +225,6 @@ def _resolve_forward( v_qdq=v_qdq, v_qdq_amax=v_qdq_amax, quant_active=quant_active, - force_kernel=force_kernel, ) @@ -260,7 +250,6 @@ def _forward_modelopt( v_qdq: str | None, v_qdq_amax: float | None, quant_active: bool, - force_kernel: bool, dense_fallback, prepare_modelopt=None, ) -> torch.Tensor: @@ -281,16 +270,9 @@ def _forward_modelopt( # N:M sparse softmax is prefill-only. for name in ("sparsity_n", "sparsity_m", "dense_sink_tokens", "dense_recent_tokens"): sparse_kw.pop(name, None) - if not _should_run_modelopt_kernel( - sparse_kw, - quant_active, - force_kernel, - ): + if not _should_run_modelopt_kernel(sparse_kw, quant_active): # Dynamic calibration can disable sparse work for a launch. Preserve the # backend's native dense path when no ModelOpt transform remains active. - # ``_modelopt_force_kernel`` (isolation knob) keeps the bf16 kernel path - # even with all quant disabled, so the on/off PPL diff holds the kernel - # constant. return dense_fallback() if prepare_modelopt is not None: prepare_modelopt() @@ -309,7 +291,6 @@ def _forward_modelopt( max_new_tokens=max_query_len, page_size=page_size, v_qdq_scale=v_qdq_scale, - decode=is_decode_only, ) q = query[:num_actual_tokens].contiguous() @@ -459,7 +440,6 @@ def native_forward(): v_qdq=resolved.v_qdq, v_qdq_amax=resolved.v_qdq_amax, quant_active=resolved.quant_active, - force_kernel=resolved.force_kernel, dense_fallback=native_forward, ) @@ -633,7 +613,6 @@ def prepare_modelopt(): "v_qdq": resolved.v_qdq, "v_qdq_amax": resolved.v_qdq_amax, "quant_active": resolved.quant_active, - "force_kernel": resolved.force_kernel, "dense_fallback": dense_fallback, "prepare_modelopt": prepare_modelopt, } @@ -662,10 +641,7 @@ def prepare_modelopt(): # Sparse-only launches may have an inactive phase (for example N:M # is prefill-only). Compute the native result once, then overwrite # each active phase with its ModelOpt result. - if _mixed_batch_needs_dense_base( - resolved.quant_active, - resolved.force_kernel, - ): + if not resolved.quant_active: dense_fallback() block_table = attn_metadata._modelopt_block_table diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py index 8ce41718d9a..60523d44457 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -488,6 +488,7 @@ def test_invalid_mode_raises(self): ("kwargs", "match"), [ ({"v_qdq": "int8"}, "v_qdq"), + ({"v_qdq": "fp8"}, "v_qdq"), ({"v_qdq": "nvfp4", "v_qdq_amax": 0.0}, "v_qdq_amax"), ({"v_qdq": "nvfp4", "v_cache_quantized": True}, "v_cache_quantized"), ], diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py index bdfbeefe47c..44eaf452863 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py @@ -279,60 +279,6 @@ def test_paged_different_page_sizes(self, page_size): torch.testing.assert_close(out_paged, out_contig, rtol=1e-2, atol=1e-2) - def test_paged_decode_ignores_sparse_nm(self): - """N:M sparse softmax is prefill-only; paged decode remains dense.""" - batch = 2 - seq_lens_k = [64, 128] - num_heads, num_kv_heads, head_dim = 4, 2, 64 - page_size = 16 - scale = 1.0 / (head_dim**0.5) - total_kv = sum(seq_lens_k) - - torch.manual_seed(34) - q_flat = torch.randn(batch, num_heads, head_dim, device="cuda", dtype=torch.float16) - k_flat = torch.randn(total_kv, num_kv_heads, head_dim, device="cuda", dtype=torch.float16) - v_flat = torch.randn(total_kv, num_kv_heads, head_dim, device="cuda", dtype=torch.float16) - - b_start_loc_q = torch.arange(batch, device="cuda", dtype=torch.int32) - b_seq_len_q = torch.ones(batch, device="cuda", dtype=torch.int32) - cumsum = [0] - for sl in seq_lens_k: - cumsum.append(cumsum[-1] + sl) - b_start_loc_k = torch.tensor(cumsum[:-1], device="cuda", dtype=torch.int32) - b_seq_len_k = torch.tensor(seq_lens_k, device="cuda", dtype=torch.int32) - k_cache, v_cache, block_table = _scatter_to_paged_cache( - k_flat, v_flat, b_start_loc_k, b_seq_len_k, num_kv_heads, head_dim, page_size - ) - - common_kwargs = { - "is_causal": False, - "softmax_scale": scale, - "b_start_loc_k": b_start_loc_k, - "b_seq_len_k": b_seq_len_k, - "max_input_len_k": max(seq_lens_k), - "k_cache": k_cache, - "v_cache": v_cache, - "block_table": block_table, - "page_size": page_size, - } - out_dense = attention( - q_flat, k_flat, v_flat, b_start_loc_q, b_seq_len_q, 1, **common_kwargs - ) - out_sparse = attention( - q_flat, - k_flat, - v_flat, - b_start_loc_q, - b_seq_len_q, - 1, - **common_kwargs, - sparsity_n=2, - sparsity_m=4, - dense_recent_tokens=64, - ) - - torch.testing.assert_close(out_sparse, out_dense, rtol=1e-2, atol=1e-2) - def test_paged_mixed_prefill_decode_sparse_nm_keeps_decode_dense(self): """Decode rows stay dense even when batched with sparse prefill rows.""" q_lens = [64, 1] @@ -447,7 +393,6 @@ def test_v_cache_finalizes_complete_groups_once(self): torch.tensor([48], device="cuda", dtype=torch.int32), max_new_tokens=1, page_size=page_size, - decode=True, ) torch.testing.assert_close(baked[:2], after_prefill[:2], rtol=0, atol=0) assert torch.all(baked[:3] == 0.015625) diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py index d706ec11508..4839be76614 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py @@ -134,29 +134,6 @@ def test_install_converts_only_attention_and_configures_fixed_recipe(monkeypatch assert state.model_runner.cascade_attn_enabled is False -def test_quant_off_disables_recipe_but_keeps_modelopt_adapter(monkeypatch): - _patch_conversion(monkeypatch) - monkeypatch.setenv("MODELOPT_ATTN_QUANT_OFF", "1") - attention = _attention() - state = _worker(nn.ModuleDict({"attn": attention})) - - worker_module._install_attention(state, quantize=True) - - converted = state.model_runner.model["attn"] - for name in ("q_bmm_quantizer", "k_bmm_quantizer", "p_bmm_quantizer", "v_bmm_quantizer"): - assert not getattr(converted, name).is_enabled - assert converted._query_quant_in_kernel is False - assert converted._value_quant_in_kernel is False - assert converted._modelopt_force_kernel is True - assert type(converted.impl) is ModelOptSparseAttentionImpl - assert converted.impl.quant_kw == { - "p_qdq": None, - "p_qdq_amax": 1.0, - "v_qdq": None, - "v_qdq_amax": None, - } - - def test_validation_of_all_layouts_precedes_mutation(monkeypatch): _patch_conversion(monkeypatch) good, bad = _attention(), _attention(AttentionType.ENCODER) diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 72e50753600..e463b413787 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -514,57 +514,13 @@ def fake_attention(query, **kwargs): assert decode_kw["p_qdq"] == "nvfp4" assert "sparsity_n" not in decode_kw assert prefill_kw["p_qdq"] == "nvfp4" - assert [(kw["max_new_tokens"], kw["decode"]) for kw in calls["finalize"]] == [ - (1, True), - (prefill_tokens, False), - ] + assert [kw["max_new_tokens"] for kw in calls["finalize"]] == [1, prefill_tokens] else: assert calls["native"] == [True] assert calls["decode"] == [] assert calls["finalize"] == [] -def test_flashinfer_forced_kernel_mixed_batch_never_uses_native_fallback(monkeypatch): - class Layer: - force_kernel_reads = 0 - - @property - def _modelopt_force_kernel(self): - self.force_kernel_reads += 1 - assert self.force_kernel_reads == 1, "force-kernel state was resolved more than once" - return True - - prefill_tokens = 17 - impl = _make_flashinfer_impl(sparse=False, quantized=False) - query = torch.zeros(1 + prefill_tokens, 2, 64, dtype=torch.float16) - kv_cache = torch.zeros(4, 2, 16, 2, 64, dtype=torch.float16) - metadata = _flashinfer_metadata(query_lens=(1, prefill_tokens), mixed=True) - layer = Layer() - calls = {"native": 0, "decode": 0, "prefill": 0} - - def native_fallback(*args, **kwargs): - calls["native"] += 1 - raise AssertionError("native FlashInfer fallback") - - def fake_decode(query, *args, **kwargs): - calls["decode"] += 1 - return torch.zeros_like(query) - - def fake_attention(query, **kwargs): - phase = "decode" if kwargs["max_input_len"] == 1 else "prefill" - calls[phase] += 1 - return torch.zeros_like(query) - - monkeypatch.setattr(FlashInferImpl, "forward", native_fallback) - monkeypatch.setattr(vllm_plugin, "triton_decode_attention", fake_decode) - monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) - - impl.forward(layer, query, query, query, kv_cache, metadata, output=torch.empty_like(query)) - - assert calls == {"native": 0, "decode": 1, "prefill": 1} - assert layer.force_kernel_reads == 1 - - def test_flashinfer_invalid_mixed_metadata_has_no_side_effects(monkeypatch): impl = _make_flashinfer_impl(sparse=True) metadata = _flashinfer_metadata(query_lens=(1, 17), mixed=True) @@ -671,7 +627,6 @@ def test_sparse_worker_rejects_unsupported_backend_before_mutation(monkeypatch): True, "vLLM cascade attention is incompatible with an active ModelOpt attention transform", ), - ("cascade", False, None), ( "metadata", True, @@ -906,22 +861,14 @@ def fake_attention(q, **kwargs): assert "target_sparse_ratio" not in captured -def test_nvfp4_bmm_mapping_and_unsupported_format_failure(): - """Enabled P/V quantizers map only dynamic block-16 NVFP4.""" - assert vllm_plugin._p_qdq_from_layer(SimpleNamespace()) == (None, 1.0) - assert vllm_plugin._v_qdq_from_layer(SimpleNamespace()) == (None, None) - - good = SimpleNamespace( +def test_unsupported_nvfp4_bmm_block_size_raises(): + """P/V mappings reject NVFP4 quantizers that are not block-16.""" + quantizer = SimpleNamespace( is_enabled=True, is_nvfp4_dynamic=True, - block_sizes={-1: 16}, - _amax=torch.tensor(6.0 * 448.0), + block_sizes={-1: 32}, ) - layer = SimpleNamespace(p_bmm_quantizer=good, v_bmm_quantizer=good) - assert vllm_plugin._p_qdq_from_layer(layer) == ("nvfp4", 6.0 * 448.0) - assert vllm_plugin._v_qdq_from_layer(layer) == ("nvfp4", 6.0 * 448.0) - - good.block_sizes = {-1: 32} + layer = SimpleNamespace(p_bmm_quantizer=quantizer, v_bmm_quantizer=quantizer) with pytest.raises(NotImplementedError, match="p_bmm_quantizer"): vllm_plugin._p_qdq_from_layer(layer) with pytest.raises(NotImplementedError, match="v_bmm_quantizer"): @@ -1007,7 +954,6 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): "max_new_tokens": 1, "page_size": 16, "v_qdq_scale": 1.0, - "decode": True, } key_cache, value_cache, block_table, seq_lens, decode_kw = calls["decode"] assert key_cache.data_ptr() == kv_cache[0].data_ptr() diff --git a/tests/unit/torch/kernels/common/attention/test_triton_fa.py b/tests/unit/torch/kernels/common/attention/test_triton_fa.py index bb170eaa8d9..62395ff5a7a 100644 --- a/tests/unit/torch/kernels/common/attention/test_triton_fa.py +++ b/tests/unit/torch/kernels/common/attention/test_triton_fa.py @@ -22,7 +22,6 @@ These tests verify CPU-safe wrapper behavior without executing a Triton kernel. """ -import os from contextlib import nullcontext import pytest @@ -99,8 +98,6 @@ def test_forward_uses_minimal_shared_autotune_configs(): (128, 32), ] - expected_active_configs = configs[:1] if "PYTEST_VERSION" in os.environ else configs - assert triton_fa._attn_fwd.configs == expected_active_configs assert triton_fa._attn_fwd.keys == ["N_CTX", "HEAD_DIM", "Q_IS_FP32", "P_QDQ", "V_QDQ"] @@ -108,10 +105,9 @@ def test_forward_uses_minimal_shared_autotune_configs(): ("attention_kwargs", "expected_p_qdq", "expected_v_qdq"), [ ({}, 0, 0), - ({"p_qdq": "fp8", "v_qdq": "fp8"}, 1, 1), + ({"p_qdq": "fp8"}, 1, 0), ({"p_qdq": "nvfp4", "v_qdq": "nvfp4"}, 2, 2), ({"p_qdq": "nvfp4", "sparsity_n": 2, "sparsity_m": 4}, 2, 0), - ({"p_qdq": "nvfp4", "skip_softmax_threshold": 0.1}, 2, 0), ], ) def test_forward_routes_every_mode_to_single_autotuner( From 64f3340485e2d653f5a65541b26cb75fee4fb0e6 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 8 Jul 2026 22:52:13 -0700 Subject: [PATCH 21/28] Consolidate attention BMM2 quantization helpers Signed-off-by: Kai Xu --- .../common/attention/decode_attention.py | 3 +- .../kernels/common/attention/triton_fa.py | 4 +- .../quantization/attention/__init__.py | 9 +-- .../attention/{v_qdq.py => bmm2_qdq.py} | 35 +++++++++- .../kernels/quantization/attention/p_qdq.py | 65 ------------------- .../quantization/common/nvfp4_quant.py | 2 +- .../attention_sparsity/plugins/vllm.py | 2 +- .../common/attention/test_decode_attention.py | 2 +- .../common/attention/test_triton_fa_paged.py | 2 +- 9 files changed, 44 insertions(+), 80 deletions(-) rename modelopt/torch/kernels/quantization/attention/{v_qdq.py => bmm2_qdq.py} (75%) delete mode 100644 modelopt/torch/kernels/quantization/attention/p_qdq.py diff --git a/modelopt/torch/kernels/common/attention/decode_attention.py b/modelopt/torch/kernels/common/attention/decode_attention.py index d12f66f16eb..1fad8d7cc57 100644 --- a/modelopt/torch/kernels/common/attention/decode_attention.py +++ b/modelopt/torch/kernels/common/attention/decode_attention.py @@ -31,8 +31,7 @@ _load_paged_k_tile, _load_paged_v_tile, ) -from modelopt.torch.kernels.quantization.attention.p_qdq import _p_qdq_nvfp4 -from modelopt.torch.kernels.quantization.attention.v_qdq import _v_qdq_nvfp4 +from modelopt.torch.kernels.quantization.attention.bmm2_qdq import _p_qdq_nvfp4, _v_qdq_nvfp4 __all__ = ["attention_decode"] diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 1d2d5c1145d..404a5dc5ca4 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -62,8 +62,8 @@ def _load_sparsity_helpers() -> None: def _load_qdq_helpers() -> None: global _qdq_fp8, _p_qdq_nvfp4, _v_qdq_nvfp4 if _qdq_fp8 is None: - from modelopt.torch.kernels.quantization.attention.p_qdq import _p_qdq_nvfp4 as _p_nvfp4 - from modelopt.torch.kernels.quantization.attention.v_qdq import _v_qdq_nvfp4 as _v_nvfp4 + from modelopt.torch.kernels.quantization.attention.bmm2_qdq import _p_qdq_nvfp4 as _p_nvfp4 + from modelopt.torch.kernels.quantization.attention.bmm2_qdq import _v_qdq_nvfp4 as _v_nvfp4 from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq as _fp8 _qdq_fp8 = _fp8 diff --git a/modelopt/torch/kernels/quantization/attention/__init__.py b/modelopt/torch/kernels/quantization/attention/__init__.py index 4232082dc60..1d859f07c33 100644 --- a/modelopt/torch/kernels/quantization/attention/__init__.py +++ b/modelopt/torch/kernels/quantization/attention/__init__.py @@ -15,10 +15,7 @@ """Quantization-specific attention kernel pieces. -``p_qdq.py`` holds the softmax-P (``p_bmm_quantizer``) quant-dequant -``@triton.jit`` helpers invoked by the unified flash-attention kernel in -``common/attention/triton_fa.py`` under its ``P_QDQ`` constexpr guard. -Only NVFP4 needs a P-specific helper (tiling and block-amax policy on top of -``quantization/common/nvfp4_quant.py``); the FP8 mode uses -``quantization/common/fp8_quant.fp8_scalar_qdq`` directly. +``bmm2_qdq.py`` holds the operand-specific NVFP4 helpers for the attention +``P @ V`` matmul and the V-cache finalization kernel. This package initializer +does not import that module, so importing the package alone does not require Triton. """ diff --git a/modelopt/torch/kernels/quantization/attention/v_qdq.py b/modelopt/torch/kernels/quantization/attention/bmm2_qdq.py similarity index 75% rename from modelopt/torch/kernels/quantization/attention/v_qdq.py rename to modelopt/torch/kernels/quantization/attention/bmm2_qdq.py index 8772fa3c605..b33900546eb 100644 --- a/modelopt/torch/kernels/quantization/attention/v_qdq.py +++ b/modelopt/torch/kernels/quantization/attention/bmm2_qdq.py @@ -13,7 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Value-operand NVFP4 helpers for flash attention.""" +"""NVFP4 operand helpers for the attention ``P @ V`` matmul (BMM2). + +P and V share the low-level ``nvfp4_scalar_qdq`` primitive, but retain thin +operand-specific wrappers because their layouts and amax reductions differ. +P is nonnegative with layout ``[M, K]``; V is signed with layout ``[K, N]``. +Both use block-16 scaling along the BMM2 contraction axis. +""" import math @@ -28,6 +34,33 @@ _BLOCK_N = 16 +@triton.jit +def _p_qdq_nvfp4( + p, + global_scale, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """NVFP4 fake quant-dequant of softmax probabilities. + + Two-level scaling per the NVFP4 recipe: E2M1 elements with one FP8 E4M3 + scale per 16 contiguous elements along the key dimension (the contraction + axis of ``P @ V``), and a per-tensor ``global_scale`` (runtime scalar, + ``amax / (6 * 448)``; ``attention()`` derives it from ``p_qdq_amax``, + which defaults to 1, the theoretical upper bound of P's amax). + + ``p >= 0``, so the block amax is a plain max (no ``abs``), and + ``nvfp4_scalar_qdq`` guards the degenerate all-zero blocks of fully + masked or padded positions. + """ + tl.static_assert(BLOCK_N % 16 == 0, "BLOCK_N must be divisible by 16 for NVFP4") + + grouped = tl.reshape(p, (BLOCK_M, BLOCK_N // 16, 16)) + block_amax = tl.expand_dims(tl.max(grouped, axis=2), 2) # p >= 0, so max == amax + q = nvfp4_scalar_qdq(grouped, block_amax, global_scale, 16) + return tl.reshape(q, (BLOCK_M, BLOCK_N)) + + @triton.jit def _v_qdq_nvfp4(v, global_scale, BLOCK_N: tl.constexpr, BLOCK_D: tl.constexpr): """Fake-quantize signed V in block-16 groups along its key axis.""" diff --git a/modelopt/torch/kernels/quantization/attention/p_qdq.py b/modelopt/torch/kernels/quantization/attention/p_qdq.py deleted file mode 100644 index 2603f6364af..00000000000 --- a/modelopt/torch/kernels/quantization/attention/p_qdq.py +++ /dev/null @@ -1,65 +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. - -"""Softmax-P quant-dequant helpers for the unified flash attention kernel. - -These ``@triton.jit`` helpers fake-quantize the softmax probabilities ``P`` -before the ``P @ V`` matmul (BMM2) — the in-kernel counterpart of the -``p_bmm_quantizer`` config. They are called conditionally from the baseline -flash-attention kernel in ``common/attention/triton_fa.py`` under the -``P_QDQ`` constexpr guard, following the same composition pattern as -the sparsity helpers in ``sparsity/attention/skip_softmax_helpers.py``. - -Only NVFP4 needs a P-specific helper (tiling policy and block amaxes); the -per-tensor FP8 mode uses ``quantization/common/fp8_quant.fp8_scalar_qdq`` -directly. What is P-specific here: the kernel's online-softmax ``p`` is -unnormalized and bounded (``0 <= p <= 1``, since the max-subtraction caps -every entry at ``exp2(0) = 1``), so 1 is the theoretical upper bound of its -amax; block amaxes need no ``abs``; and the NVFP4 scale blocks of 16 run -along the key dimension — the contraction axis of ``P @ V``. The caller -(``attention()``) converts the amax to the ``global_scale`` below. -""" - -import triton -import triton.language as tl - -from modelopt.torch.kernels.quantization.common.nvfp4_quant import nvfp4_scalar_qdq - - -@triton.jit -def _p_qdq_nvfp4( - p, - global_scale, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, -): - """NVFP4 fake quant-dequant of softmax probabilities. - - Two-level scaling per the NVFP4 recipe: E2M1 elements with one FP8 E4M3 - scale per 16 contiguous elements along the key dimension (the contraction - axis of ``P @ V``), and a per-tensor ``global_scale`` (runtime scalar, - ``amax / (6 * 448)``; ``attention()`` derives it from ``p_qdq_amax``, - which defaults to 1, the theoretical upper bound of P's amax). - - ``p >= 0``, so the block amax is a plain max (no ``abs``), and - ``nvfp4_scalar_qdq`` guards the degenerate all-zero blocks of fully - masked or padded positions. - """ - tl.static_assert(BLOCK_N % 16 == 0, "BLOCK_N must be divisible by 16 for NVFP4") - - grouped = tl.reshape(p, (BLOCK_M, BLOCK_N // 16, 16)) - block_amax = tl.expand_dims(tl.max(grouped, axis=2), 2) # p >= 0, so max == amax - q = nvfp4_scalar_qdq(grouped, block_amax, global_scale, 16) - return tl.reshape(q, (BLOCK_M, BLOCK_N)) diff --git a/modelopt/torch/kernels/quantization/common/nvfp4_quant.py b/modelopt/torch/kernels/quantization/common/nvfp4_quant.py index 85b73fe6388..d74a346df90 100644 --- a/modelopt/torch/kernels/quantization/common/nvfp4_quant.py +++ b/modelopt/torch/kernels/quantization/common/nvfp4_quant.py @@ -19,7 +19,7 @@ - ``../gemm/fp4_kernel.py`` (standalone blockwise fake quant) - ``../gemm/fp4_kernel_hopper.py`` (Hopper block-pointer variant) - ``../gemm/gptq_fused_kernel.py`` (fused GPTQ scalar path) - - ``../attention/p_qdq.py`` (softmax-P qdq in the flash-attention kernel) + - ``../attention/bmm2_qdq.py`` (P/V operand qdq in the attention BMM2 kernel) FP4 (E2M1) representable magnitudes: {0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0} """ diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 65167e2d1bc..0e25dbdadd3 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -44,7 +44,7 @@ attention_decode as triton_decode_attention, ) from modelopt.torch.kernels.common.attention.triton_fa import attention as triton_attention -from modelopt.torch.kernels.quantization.attention.v_qdq import fake_quant_v_onwrite +from modelopt.torch.kernels.quantization.attention.bmm2_qdq import fake_quant_v_onwrite def _target_sparse_ratio_for_phase(target_sparse_ratio, phase: str) -> float: diff --git a/tests/gpu/torch/kernels/common/attention/test_decode_attention.py b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py index 490fcac9fd1..a018074abcb 100644 --- a/tests/gpu/torch/kernels/common/attention/test_decode_attention.py +++ b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py @@ -24,7 +24,7 @@ if TRITON_KERNEL_AVAILABLE: from modelopt.torch.kernels.common.attention.decode_attention import attention_decode - from modelopt.torch.kernels.quantization.attention.v_qdq import fake_quant_v_onwrite + from modelopt.torch.kernels.quantization.attention.bmm2_qdq import fake_quant_v_onwrite NATIVE_E4M3_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (8, 9) requires_native_e4m3 = pytest.mark.skipif( diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py index 44eaf452863..a1513b497c3 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py @@ -27,7 +27,7 @@ if TRITON_KERNEL_AVAILABLE: from modelopt.torch.kernels.common.attention import attention, triton_fa - from modelopt.torch.kernels.quantization.attention.v_qdq import fake_quant_v_onwrite + from modelopt.torch.kernels.quantization.attention.bmm2_qdq import fake_quant_v_onwrite NATIVE_E4M3_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (8, 9) requires_native_e4m3 = pytest.mark.skipif( From 74cde538e389e13c973fc2371f86759b034671aa Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 8 Jul 2026 15:08:55 -0700 Subject: [PATCH 22/28] Document mixed-FP16 softmax design Signed-off-by: Kai Xu --- .../2026-07-08-mixed-fp16-softmax-design.md | 88 +++++++++++++++++++ 1 file changed, 88 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-08-mixed-fp16-softmax-design.md diff --git a/docs/superpowers/specs/2026-07-08-mixed-fp16-softmax-design.md b/docs/superpowers/specs/2026-07-08-mixed-fp16-softmax-design.md new file mode 100644 index 00000000000..e641f4797b1 --- /dev/null +++ b/docs/superpowers/specs/2026-07-08-mixed-fp16-softmax-design.md @@ -0,0 +1,88 @@ +# Mixed-FP16 Softmax Design + +## Goal + +Add the reference mixed-FP16 online-softmax path to the compact NVFP4 attention worker while preserving one public `QuantSparseAttnWorker` and the current fixed block-16 NVFP4 Q/K/P/V recipe. + +The worker must select either the existing FP32 softmax or mixed-FP16 softmax before model installation and CUDA-graph capture. This is inference-only support for regular decoder self-attention through the existing FlashAttention and FlashInfer adapters. + +## Configuration Boundary + +Q/K/P/V formats remain `TensorQuantizer` configuration because they describe tensor QDQ operations. Mixed-FP16 softmax is a kernel compute policy: it changes the two online-softmax exponentials but is not a tensor QDQ format. + +Use one compact-only environment selector: + +```text +MODELOPT_ATTN_SOFTMAX_MODE=fp32 +MODELOPT_ATTN_SOFTMAX_MODE=mixed_fp16 +``` + +The default is `fp32`. Any other value fails before an attention layer is mutated. The launcher must forward the selector to Ray workers so every tensor-parallel worker installs the same static policy. + +The worker resolves the value once and stores the normalized string in the existing immutable launch-policy snapshot: + +```python +impl.quant_kw["softmax_mode"] = "fp32" | "mixed_fp16" +``` + +No additional worker class, softmax `TensorQuantizer`, generic per-point precision matrix, or integration-branch compatibility alias is added. + +## Numerical Contract + +`fp32` preserves current behavior. + +`mixed_fp16` implements the reference path: + +- Convert the input of `exp2(scores - new_max)` from FP32 to FP16 with round-to-nearest-even, execute native `ex2.approx.f16`, and return its FP16-valued result as FP32. +- Apply the same native FP16 operation to the online correction `exp2(old_max - new_max)`. +- Sum probabilities into the denominator in FP32 without another rounding step. +- Keep the denominator state, weighted-value accumulator, and matrix accumulators in FP32. +- Apply existing NVFP4 P QDQ after the mixed-FP16 probability computation and after the unquantized FP32 denominator sum. +- Keep split-K reconciliation in FP32. Mixed FP16 applies only inside each split's online-softmax loop. +- Leave Q/K/P/V QDQ, V-cache finalization, sparse masking, and skip-softmax decisions unchanged. + +The mode is a Triton compile-time constant. It introduces no device scalar, host synchronization, or graph-time mutable state. + +## Data Flow + +```text +MODELOPT_ATTN_SOFTMAX_MODE + -> worker validation and normalization + -> impl.quant_kw["softmax_mode"] + -> _ResolvedForward + -> shared FlashAttention/FlashInfer _forward_modelopt path + -> triton_fa.attention / decode_attention.attention_decode + -> MIXED_FP16 constexpr +``` + +The mode counts as an active ModelOpt transform. A non-FP32 mode must therefore never delegate to the backend's native dense path, including FlashInfer mixed prefill/decode batches. + +## API And Error Handling + +The public prefill and decode kernel wrappers accept `softmax_mode="fp32"` and validate against the two supported values. Internal Triton kernels receive only a boolean `MIXED_FP16` constexpr. + +Autograd with `mixed_fp16` raises `NotImplementedError`; backward recomputation is outside this inference-focused change. Existing autograd behavior for `fp32` remains unchanged. + +Worker configuration is read before plan installation. Invalid configuration fails atomically, without converting modules, replacing implementations, or partially installing a policy. + +## Testing + +Keep checked-in coverage focused on behavioral contracts: + +- Worker selection defaults to `fp32`, accepts `mixed_fp16`, rejects invalid modes atomically, and forwards the selector to Ray workers. +- FlashAttention and FlashInfer routes propagate the mode for prefill, decode, and FlashInfer mixed batches without native fallback. +- Prefill matches a mixed-FP16 reference while composing with NVFP4 Q/K/P/V and 2:4 sparsity. +- Fixed 32-split decode matches a split-local mixed-FP16 reference with FP32 split reconciliation. +- The native FP16 `ex2` helper covers representative finite inputs and masked negative infinity. +- Mixed-FP16 autograd is rejected before launch. +- Default `fp32` routing and numerics remain unchanged. + +CUDA-graph support is claimed only after a capture/replay smoke confirms that the static mode is preserved. + +## Non-Goals + +- Porting the integration branch's `MODELOPT_ATTN_FP16_SOFTMAX` alias or generic `MODELOPT_ATTN_SOFTMAX_QUANT` matrix. +- Supporting independently configurable DIFF, EXP2, DELTA, ALPHA, ACC, reciprocal, or output formats. +- Changing the fixed NVFP4 recipe, split count, tile sizes, sparse behavior, or supported vLLM attention envelope. +- Implementing mixed-FP16 backward kernels. +- Adding the design document to the final pull request. From 389e7b0522623f340e2db37cd7d4906331d33b3c Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 8 Jul 2026 15:41:54 -0700 Subject: [PATCH 23/28] Add mixed-FP16 softmax primitive Signed-off-by: Kai Xu --- .../attention/softmax_fakequant.py | 43 +++++++++++++++++++ .../common/attention/test_triton_fa_p_qdq.py | 40 +++++++++++++++++ .../common/attention/test_triton_fa.py | 11 +++++ 3 files changed, 94 insertions(+) create mode 100644 modelopt/torch/kernels/quantization/attention/softmax_fakequant.py diff --git a/modelopt/torch/kernels/quantization/attention/softmax_fakequant.py b/modelopt/torch/kernels/quantization/attention/softmax_fakequant.py new file mode 100644 index 00000000000..c78649afd7a --- /dev/null +++ b/modelopt/torch/kernels/quantization/attention/softmax_fakequant.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Mixed-FP16 softmax helpers for Triton attention kernels.""" + +import triton +import triton.language as tl + +__all__ = ["SOFTMAX_MODES", "ex2_fp16", "resolve_softmax_mode"] + +SOFTMAX_MODES = frozenset(("fp32", "mixed_fp16")) + + +def resolve_softmax_mode(mode: str) -> bool: + """Validate ``mode`` and report whether mixed-FP16 softmax is enabled.""" + if mode not in SOFTMAX_MODES: + raise ValueError(f"softmax_mode must be one of {sorted(SOFTMAX_MODES)}, got {mode!r}") + return mode == "mixed_fp16" + + +@triton.jit +def ex2_fp16(x): + """Evaluate base-2 exponentiation in native FP16 on sm_75 or newer and return FP32.""" + return tl.inline_asm_elementwise( + "{ .reg .b16 h; cvt.rn.f16.f32 h, $1; ex2.approx.f16 h, h; cvt.f32.f16 $0, h; }", + "=r,r", + args=[x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py index 60523d44457..4c4236265e2 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -27,15 +27,55 @@ from modelopt.torch.quantization.tensor_quant import fp8_eager if TRITON_KERNEL_AVAILABLE: + import triton + import triton.language as tl + from modelopt.torch.kernels.common.attention import attention, triton_fa from modelopt.torch.kernels.common.attention.triton_fa import LOG2E + from modelopt.torch.kernels.quantization.attention.softmax_fakequant import ex2_fp16 + + @triton.jit + def _ex2_fp16_test_kernel(x, out, n, BLOCK: tl.constexpr): # noqa: N803 + offsets = tl.arange(0, BLOCK) + mask = offsets < n + tl.store(out + offsets, ex2_fp16(tl.load(x + offsets, mask=mask)), mask=mask) + +NATIVE_FP16_EX2_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (7, 5) +requires_native_fp16_ex2 = pytest.mark.skipif( + not NATIVE_FP16_EX2_AVAILABLE, + reason="Native FP16 ex2 requires compute capability >= 7.5", +) NATIVE_E4M3_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (8, 9) requires_native_e4m3 = pytest.mark.skipif( not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute capability >= 8.9" ) +@requires_native_fp16_ex2 +def test_ex2_fp16_matches_fp16_reference(): + x = torch.tensor( + [float("-inf"), -12.0, -3.25, -1.0, -0.125, -0.1, 0.0], + dtype=torch.float32, + device="cuda", + ) + actual = torch.empty_like(x) + + _ex2_fp16_test_kernel[(1,)](x, actual, x.numel(), BLOCK=8) + + expected = torch.exp2(x.to(torch.float16)).to(torch.float32) + # PTX specifies a maximum relative error of 2**-9.9 for ex2.approx.f16. + torch.testing.assert_close(actual, expected, rtol=2**-9, atol=0.0) + actual_finite = actual[torch.isfinite(actual)] + torch.testing.assert_close( + actual_finite, + actual_finite.to(torch.float16).float(), + rtol=0.0, + atol=0.0, + ) + assert actual[0].item() == 0.0 + + @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need triton") def test_v_qdq_rejects_autograd_before_kernel_launch(monkeypatch): apply = Mock(side_effect=AssertionError("_Attention.apply reached")) diff --git a/tests/unit/torch/kernels/common/attention/test_triton_fa.py b/tests/unit/torch/kernels/common/attention/test_triton_fa.py index 62395ff5a7a..9c9f1599443 100644 --- a/tests/unit/torch/kernels/common/attention/test_triton_fa.py +++ b/tests/unit/torch/kernels/common/attention/test_triton_fa.py @@ -61,6 +61,17 @@ def test_triton_fa_importable_on_cpu(): assert callable(calibrate.attention_calibrate) +def test_softmax_mode_validation(): + pytest.importorskip("triton") + + from modelopt.torch.kernels.quantization.attention.softmax_fakequant import resolve_softmax_mode + + assert resolve_softmax_mode("fp32") is False + assert resolve_softmax_mode("mixed_fp16") is True + with pytest.raises(ValueError, match="softmax_mode"): + resolve_softmax_mode("fp16") + + def test_forward_buckets_autotune_key_without_bucketing_grid(monkeypatch): """Reuse autotune results by length regime without launching extra query tiles.""" pytest.importorskip("triton") From da3cdb6870a5f136283f57865cba446ec52e05e0 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 8 Jul 2026 16:04:11 -0700 Subject: [PATCH 24/28] Add mixed-FP16 attention prefill Signed-off-by: Kai Xu --- .../kernels/common/attention/triton_fa.py | 30 +++- .../common/attention/test_triton_fa_p_qdq.py | 149 +++++++++++++++++- .../common/attention/test_triton_fa.py | 20 ++- 3 files changed, 184 insertions(+), 15 deletions(-) diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 404a5dc5ca4..a784e17bfc8 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -29,6 +29,11 @@ import triton import triton.language as tl +from modelopt.torch.kernels.quantization.attention.softmax_fakequant import ( + ex2_fp16, + resolve_softmax_mode, +) + # Helpers for optional N:M sparsity and skip-softmax live in the sparsity # package. The baseline forward kernel below calls them conditionally under # constexpr guards, so the unified single-kernel design stays intact while @@ -282,6 +287,7 @@ def _attn_fwd( HEAD_DIM: tl.constexpr, # Actual head dimension (for d_mask) STORE_LSE: tl.constexpr, # Whether to save LSE for backward pass Q_IS_FP32: tl.constexpr, # Dynamic NVFP4 QDQ carrier uses FP32 + MIXED_FP16: tl.constexpr = False, # FP16 exp2 with FP32 reductions and accumulators SPARSITY_N: tl.constexpr = 0, # N:M sparsity — keep top-N of every M elements (0 = disabled) SPARSITY_M: tl.constexpr = 4, # N:M sparsity — group size (4 or 8) DENSE_SINK_TOKENS: tl.constexpr = 0, # Leading KV tokens kept dense (attention sinks) @@ -430,9 +436,13 @@ def _attn_fwd( if not skip_tile: # --- Online softmax update --- m_new = tl.maximum(row_max, tl.max(scores, 1)) - p = tl.math.exp2(scores - m_new[:, None]) + if MIXED_FP16: + p = ex2_fp16(scores - m_new[:, None]) + correction = ex2_fp16(row_max - m_new) + else: + p = tl.math.exp2(scores - m_new[:, None]) + correction = tl.math.exp2(row_max - m_new) l_new = tl.sum(p, 1) - correction = tl.math.exp2(row_max - m_new) row_sum = row_sum * correction + l_new acc = acc * correction[:, None] @@ -881,6 +891,7 @@ def forward( max_input_len, is_causal, sm_scale, + mixed_fp16, b_start_loc_k, b_seq_len_k, max_input_len_k, @@ -990,6 +1001,7 @@ def forward( "HEAD_DIM": HEAD_DIM, "STORE_LSE": True, "Q_IS_FP32": q.dtype == torch.float32 and (p_qdq_mode == 2 or v_qdq_mode == 2), + "MIXED_FP16": mixed_fp16, "SPARSITY_N": sparsity_n, "SPARSITY_M": sparsity_m, "DENSE_SINK_TOKENS": dense_sink_tokens, @@ -1195,6 +1207,7 @@ def backward(ctx, grad_output): None, # max_input_len None, # is_causal None, # sm_scale + None, # mixed_fp16 None, # b_start_loc_k None, # b_seq_len_k None, # max_input_len_k @@ -1235,6 +1248,7 @@ def attention( dense_recent_tokens: int = 64, skip_softmax_threshold: float | None = None, measure_sparsity: bool = False, + softmax_mode: str = "fp32", p_qdq: str | None = None, p_qdq_amax: float = 1.0, v_qdq: str | None = None, @@ -1277,6 +1291,12 @@ def attention( and skipped tiles via atomic counters. The counts are stored as ``_sparsity_total`` and ``_sparsity_skipped`` attributes on the returned output tensor. + softmax_mode: Exponentiation mode for the forward online-softmax update. + ``"fp32"`` evaluates probability and running-maximum correction + exponentials in FP32. ``"mixed_fp16"`` evaluates those exponentials + in native FP16 and converts their results back to FP32; row maxima, + denominator sums, and output accumulation remain FP32. Mixed mode is + inference-only because backward continues to use FP32 exponentials. p_qdq: Fake quant-dequant of the softmax probabilities ``P`` before the ``P @ V`` matmul (BMM2), emulating quantized attention. ``"fp8"`` round-trips P through FP8 E4M3 with a static per-tensor @@ -1327,6 +1347,11 @@ def attention( # silently reuse stale compiled kernels from the on-disk cache. _load_sparsity_helpers() _load_qdq_helpers() + mixed_fp16 = resolve_softmax_mode(softmax_mode) + if mixed_fp16 and any(t.requires_grad for t in (q, k, v)): + raise NotImplementedError( + "mixed_fp16 softmax is inference-only and does not support autograd" + ) if p_qdq not in _P_QDQ_MODES: raise ValueError( f"p_qdq must be one of {sorted(k for k in _P_QDQ_MODES if k)} or None, got {p_qdq!r}" @@ -1367,6 +1392,7 @@ def attention( max_input_len, is_causal, sm_scale, + mixed_fp16, b_start_loc_k, b_seq_len_k, max_input_len_k, diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py index 4c4236265e2..04d4227f63a 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -77,14 +77,21 @@ def test_ex2_fp16_matches_fp16_reference(): @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need triton") -def test_v_qdq_rejects_autograd_before_kernel_launch(monkeypatch): +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"v_qdq": "nvfp4"}, r"v_qdq.*autograd"), + ({"softmax_mode": "mixed_fp16"}, r"mixed_fp16.*autograd"), + ], +) +def test_inference_only_modes_reject_autograd_before_kernel_launch(monkeypatch, kwargs, match): apply = Mock(side_effect=AssertionError("_Attention.apply reached")) monkeypatch.setattr(triton_fa._Attention, "apply", apply) q, k, v = (torch.zeros(1, 1, 16, requires_grad=True) for _ in range(3)) locs = torch.zeros(1, dtype=torch.int32) lens = torch.ones(1, dtype=torch.int32) - with pytest.raises(NotImplementedError, match=r"v_qdq.*autograd"): - attention(q, k, v, locs, lens, 1, v_qdq="nvfp4") + with pytest.raises(NotImplementedError, match=match): + attention(q, k, v, locs, lens, 1, **kwargs) apply.assert_not_called() @@ -258,6 +265,22 @@ def _qdq_nvfp4(p, global_scale=1.0): return q.reshape(shape) +def _qdq_v_nvfp4(v, global_scale=1.0): + """Signed NVFP4 V qdq in block-16 groups along the key axis.""" + carrier_dtype = v.dtype + shape = v.shape + assert shape[0] % 16 == 0 + g = v.float().reshape(shape[0] // 16, 16, *shape[1:]) + block_amax = g.abs().amax(dim=1, keepdim=True) + scale_in_fp8_range = (block_amax / (6.0 * global_scale)).clamp(max=FP8_E4M3_MAX) + scale = scale_in_fp8_range.to(torch.float8_e4m3fn).float() * global_scale + scale_safe = torch.where(scale == 0, 1.0, scale) + q = _fp4_round(g.abs() / scale_safe) * scale_safe + q = torch.where(g < 0, -q, q) + q = torch.where(scale == 0, 0.0, q) + return q.reshape(shape).to(carrier_dtype).float() + + def _apply_qdq(p, mode, qdq_scale=1.0): if mode == "fp8": return _qdq_fp8(p, qdq_scale) @@ -265,7 +288,45 @@ def _apply_qdq(p, mode, qdq_scale=1.0): return _qdq_nvfp4(p, qdq_scale) -def qdq_attention_reference(q, k, v, scale, mode, is_causal=True, amax=1.0, block_n=P_QDQ_BLOCK_N): +def _softmax_exp2(x, softmax_mode): + if softmax_mode == "mixed_fp16": + return torch.exp2(x.to(torch.float16)).to(torch.float32) + assert softmax_mode == "fp32" + return torch.exp2(x) + + +def _apply_top_n_of_m_tile(scores, n, m): + """Apply the kernel's per-tile top-N-of-M masking along the key axis.""" + assert m in (4, 8) and scores.shape[-1] % m == 0 + grouped = scores.reshape(*scores.shape[:-1], scores.shape[-1] // m, m) + if m == 4: + # The Triton comparison network breaks ties in favor of lower indices. + order = torch.argsort(grouped, dim=-1, descending=True, stable=True) + kept = torch.zeros_like(grouped, dtype=torch.bool) + kept.scatter_(-1, order[..., :n], True) + else: + # The Triton M=8 path keeps every value tied at the Nth-largest threshold. + threshold = grouped.sort(dim=-1).values[..., m - n] + kept = grouped >= threshold[..., None] + return grouped.masked_fill(~kept, float("-inf")).reshape_as(scores) + + +def qdq_attention_reference( + q, + k, + v, + scale, + mode, + is_causal=True, + amax=1.0, + block_n=P_QDQ_BLOCK_N, + softmax_mode="fp32", + v_qdq=None, + sparsity_n=0, + sparsity_m=4, + dense_sink_tokens=0, + dense_recent_tokens=64, +): """Tile-looped online-softmax reference replicating kernel P_QDQ semantics. Single sequence: q [s, h, d], k/v [s_kv, h_kv, d] (fp16). Walks KV tiles @@ -302,10 +363,21 @@ def qdq_attention_reference(q, k, v, scale, mode, is_causal=True, amax=1.0, bloc acc = torch.zeros(h, s, d, device=q.device) for start in range(0, s_kv, block_n): tile = t[:, :, start : start + block_n] + if sparsity_n: + sparse_tile = _apply_top_n_of_m_tile(tile, sparsity_n, sparsity_m) + q_abs = torch.arange(s, device=q.device) + s_kv - s + kv_abs = torch.arange(start, start + tile.shape[-1], device=q.device) + distance = q_abs[:, None] - kv_abs[None, :] + dense = (kv_abs[None, :] < dense_sink_tokens) | ( + (distance >= 0) & (distance < dense_recent_tokens) + ) + if s <= 1: + dense = torch.ones_like(dense) + tile = torch.where(dense[None], tile, sparse_tile) m_new = torch.maximum(row_max, tile.amax(dim=-1)) - p = torch.exp2(tile - m_new[..., None]) + p = _softmax_exp2(tile - m_new[..., None], softmax_mode) l_new = p.sum(dim=-1) - corr = torch.exp2(row_max - m_new) + corr = _softmax_exp2(row_max - m_new, softmax_mode) row_sum = row_sum * corr + l_new acc = acc * corr[..., None] if mode == "nvfp4": @@ -313,7 +385,13 @@ def qdq_attention_reference(q, k, v, scale, mode, is_causal=True, amax=1.0, bloc p = _apply_qdq(p, mode, qdq_scale) if mode == "fp8": p = p.to(v.dtype).float() - acc = acc + torch.einsum("hqk,khd->hqd", p, vv[start : start + block_n].float()) + v_tile = vv[start : start + block_n] + if v_qdq == "nvfp4": + # The kernel QDQs in FP32, then rounds back to the cache/model carrier dtype. + v_tile = _qdq_v_nvfp4(v_tile) + else: + assert v_qdq is None + acc = acc + torch.einsum("hqk,khd->hqd", p, v_tile.float()) row_max = m_new out = acc / row_sum[..., None] return out.permute(1, 0, 2) @@ -345,6 +423,63 @@ def test_prefill_matches_tile_reference(self, mode): o_dense = attention(q, k, v, locs, lens, seq_len, softmax_scale=scale) assert not torch.equal(o, o_dense) + @requires_native_fp16_ex2 + @requires_native_e4m3 + def test_mixed_fp16_prefill_matches_tile_reference(self): + """Mixed softmax composes with P/V NVFP4 and tile-local 2:4 sparsity.""" + seq_len, num_heads, num_kv_heads, head_dim = 128, 2, 1, 64 + scale = 1.0 / math.sqrt(head_dim) + + torch.manual_seed(29) + q, k, v = make_qkv(seq_len, num_heads, num_kv_heads, head_dim, dtype=torch.float16) + locs, lens = make_varlen_meta([seq_len]) + + actual = attention( + q, + k, + v, + locs, + lens, + seq_len, + softmax_scale=scale, + softmax_mode="mixed_fp16", + p_qdq="nvfp4", + v_qdq="nvfp4", + sparsity_n=2, + sparsity_m=4, + dense_recent_tokens=0, + ) + reference = qdq_attention_reference( + q, + k, + v, + scale, + "nvfp4", + softmax_mode="mixed_fp16", + v_qdq="nvfp4", + sparsity_n=2, + sparsity_m=4, + dense_recent_tokens=0, + ) + fp32 = attention( + q, + k, + v, + locs, + lens, + seq_len, + softmax_scale=scale, + softmax_mode="fp32", + p_qdq="nvfp4", + v_qdq="nvfp4", + sparsity_n=2, + sparsity_m=4, + dense_recent_tokens=0, + ) + + torch.testing.assert_close(actual.float(), reference, rtol=5e-3, atol=2e-2) + assert not torch.equal(actual, fp32) + @requires_native_e4m3 def test_nvfp4_uses_native_p_input_and_fp32_carriers(self): """Dynamic-Q and P QDQ stay FP32 until their native MMA boundaries.""" diff --git a/tests/unit/torch/kernels/common/attention/test_triton_fa.py b/tests/unit/torch/kernels/common/attention/test_triton_fa.py index 9c9f1599443..d78ae5295b2 100644 --- a/tests/unit/torch/kernels/common/attention/test_triton_fa.py +++ b/tests/unit/torch/kernels/common/attention/test_triton_fa.py @@ -113,16 +113,23 @@ def test_forward_uses_minimal_shared_autotune_configs(): @pytest.mark.parametrize( - ("attention_kwargs", "expected_p_qdq", "expected_v_qdq"), + ("attention_kwargs", "expected_p_qdq", "expected_v_qdq", "expected_mixed_fp16"), [ - ({}, 0, 0), - ({"p_qdq": "fp8"}, 1, 0), - ({"p_qdq": "nvfp4", "v_qdq": "nvfp4"}, 2, 2), - ({"p_qdq": "nvfp4", "sparsity_n": 2, "sparsity_m": 4}, 2, 0), + ({}, 0, 0, False), + ({"softmax_mode": "mixed_fp16"}, 0, 0, True), + ({"p_qdq": "fp8"}, 1, 0, False), + ({"p_qdq": "nvfp4", "v_qdq": "nvfp4"}, 2, 2, False), + ( + {"softmax_mode": "mixed_fp16", "p_qdq": "nvfp4", "v_qdq": "nvfp4"}, + 2, + 2, + True, + ), + ({"p_qdq": "nvfp4", "sparsity_n": 2, "sparsity_m": 4}, 2, 0, False), ], ) def test_forward_routes_every_mode_to_single_autotuner( - monkeypatch, attention_kwargs, expected_p_qdq, expected_v_qdq + monkeypatch, attention_kwargs, expected_p_qdq, expected_v_qdq, expected_mixed_fp16 ): """Every non-measurement launch uses the unified autotuner.""" pytest.importorskip("triton") @@ -148,6 +155,7 @@ def test_forward_routes_every_mode_to_single_autotuner( assert kernel.kwargs["P_QDQ"] == expected_p_qdq assert kernel.kwargs["V_QDQ"] == expected_v_qdq + assert kernel.kwargs["MIXED_FP16"] is expected_mixed_fp16 @pytest.mark.parametrize( From 17bc8c95788e1cb89a9abef88495a74d0a482c1e Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 8 Jul 2026 16:20:27 -0700 Subject: [PATCH 25/28] Add mixed-FP16 split-K decode Signed-off-by: Kai Xu --- .../common/attention/decode_attention.py | 21 +++- .../common/attention/test_decode_attention.py | 106 +++++++++++++----- 2 files changed, 96 insertions(+), 31 deletions(-) diff --git a/modelopt/torch/kernels/common/attention/decode_attention.py b/modelopt/torch/kernels/common/attention/decode_attention.py index 1fad8d7cc57..a77eb0ca44a 100644 --- a/modelopt/torch/kernels/common/attention/decode_attention.py +++ b/modelopt/torch/kernels/common/attention/decode_attention.py @@ -32,6 +32,10 @@ _load_paged_v_tile, ) from modelopt.torch.kernels.quantization.attention.bmm2_qdq import _p_qdq_nvfp4, _v_qdq_nvfp4 +from modelopt.torch.kernels.quantization.attention.softmax_fakequant import ( + ex2_fp16, + resolve_softmax_mode, +) __all__ = ["attention_decode"] @@ -74,6 +78,7 @@ def _decode_split_kernel( PAGE_SIZE: tl.constexpr, max_blocks_per_seq, NUM_KV_SPLITS: tl.constexpr, + MIXED_FP16: tl.constexpr, P_QDQ: tl.constexpr, V_QDQ: tl.constexpr, V_CACHE_QUANTIZED: tl.constexpr, @@ -131,9 +136,13 @@ def _decode_split_kernel( scores = tl.where(kv_valid, scores, -float("inf")) tile_max = tl.max(scores, axis=0) new_max = tl.maximum(running_max, tile_max) - p = tl.math.exp2(scores - new_max) + if MIXED_FP16: + p = ex2_fp16(scores - new_max) + correction = ex2_fp16(running_max - new_max) + else: + p = tl.math.exp2(scores - new_max) + correction = tl.math.exp2(running_max - new_max) p = tl.where(kv_valid, p, 0.0) - correction = tl.math.exp2(running_max - new_max) running_sum = running_sum * correction + tl.sum(p, axis=0) acc *= correction @@ -257,6 +266,7 @@ def attention_decode( b_seq_len_k: torch.Tensor, *, softmax_scale: float | None = None, + softmax_mode: str = "fp32", page_size: int = 16, num_kv_splits: int = _DEFAULT_KV_SPLITS, p_qdq: str | None = None, @@ -274,7 +284,13 @@ def attention_decode( in the cache; only the pristine partial group is then quantized on read. P QDQ intentionally follows the split-local online-softmax schedule; changing ``num_kv_splits`` can therefore change quantized results. + + ``softmax_mode="mixed_fp16"`` evaluates split-local probabilities and + running-max corrections with native FP16 exponentiation while retaining + FP32 denominator state and accumulators. Split-state combination remains + FP32. The default ``"fp32"`` preserves the existing behavior. """ + mixed_fp16 = resolve_softmax_mode(softmax_mode) if q.ndim != 3: raise ValueError(f"q must have shape [batch, heads, head_dim], got {tuple(q.shape)}") if page_size != k_cache.shape[1] or page_size != v_cache.shape[1]: @@ -337,6 +353,7 @@ def attention_decode( PAGE_SIZE=page_size, max_blocks_per_seq=block_table.shape[1], NUM_KV_SPLITS=num_kv_splits, + MIXED_FP16=mixed_fp16, P_QDQ=p_qdq == "nvfp4", V_QDQ=v_qdq == "nvfp4", V_CACHE_QUANTIZED=v_cache_quantized, diff --git a/tests/gpu/torch/kernels/common/attention/test_decode_attention.py b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py index a018074abcb..d59a618c243 100644 --- a/tests/gpu/torch/kernels/common/attention/test_decode_attention.py +++ b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py @@ -81,6 +81,27 @@ def _nvfp4_qdq_reference(x, global_scale=1.0 / (6.0 * 448.0)): return torch.where(scales == 0, 0.0, quantized).reshape_as(x) +def _softmax_exp2(x, softmax_mode): + if softmax_mode == "mixed_fp16": + return torch.exp2(x.half()).float() + assert softmax_mode == "fp32" + return torch.exp2(x) + + +def _combine_split_states(split_max, split_sum, split_acc, softmax_mode): + running_max = torch.tensor(-float("inf"), device=split_max.device) + running_sum = torch.tensor(0.0, device=split_max.device) + acc = torch.zeros_like(split_acc[0]) + for split_idx in range(split_max.shape[0]): + new_max = torch.maximum(running_max, split_max[split_idx]) + correction = _softmax_exp2(running_max - new_max, softmax_mode) + split_correction = _softmax_exp2(split_max[split_idx] - new_max, softmax_mode) + acc = acc * correction + split_acc[split_idx] * split_correction + running_sum = running_sum * correction + split_sum[split_idx] * split_correction + running_max = new_max + return acc / running_sum + + @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") @pytest.mark.parametrize("num_kv_splits", [1, 32]) def test_split_k_varlen_gqa_matches_dense(num_kv_splits): @@ -203,46 +224,73 @@ def test_p_quantizes_model_dtype_input_but_accumulates_fp32(): @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") @requires_native_e4m3 -def test_p_qdq_matches_fixed_split_local_oracle(): +@pytest.mark.parametrize("softmax_mode", ["fp32", "mixed_fp16"]) +def test_p_qdq_matches_fixed_split_local_oracle(softmax_mode): """The production 32-split schedule quantizes split-local unnormalized P.""" - seq_len, head_dim, num_splits = 4096, 16, 32 + seq_len, head_dim, num_splits = 8192, 16, 32 + block_n = 128 scale = head_dim**-0.5 q = torch.zeros(1, 1, head_dim, device="cuda", dtype=torch.float32) - q[..., 0] = 1.0 + q[..., 0] = 1.0 / (scale * 1.4426950408889634) k = torch.zeros(1, 1, seq_len, head_dim, device="cuda", dtype=torch.bfloat16) - k[..., 0] = torch.linspace(-8.0, 8.0, seq_len, device="cuda", dtype=torch.bfloat16) - torch.manual_seed(19) - v = torch.randn_like(k) + token_idx = torch.arange(seq_len, device="cuda") + token_split = token_idx // (2 * block_n) + tile_phase = (token_idx // block_n) % 2 + score_pattern = -0.25 * ((token_idx % block_n) % 3) + 0.25 * tile_phase + 0.25 * token_split + k[0, 0, :, 0] = score_pattern.to(torch.bfloat16) + v = torch.zeros_like(k) + # Channel 0 isolates the split-local correction; channel 1 nearly cancels only + # under the intended FP32 combine, making combine precision observable. + v[0, 0, :, 0] = torch.where(tile_phase == 0, 256.0, -215.0).to(torch.bfloat16) + v[0, 0, :, 1] = torch.where(token_split == num_splits - 1, -5.25, 1.0).to(torch.bfloat16) k_cache, v_cache, block_table = _paged_cache(k, v, (seq_len,)) seq_lens = torch.tensor([seq_len], device="cuda", dtype=torch.int32) + common = { + "softmax_scale": scale, + "num_kv_splits": num_splits, + "p_qdq": "nvfp4", + } output = attention_decode( - q, - k_cache, - v_cache, - block_table, - seq_lens, - softmax_scale=scale, - num_kv_splits=num_splits, - p_qdq="nvfp4", + q, k_cache, v_cache, block_table, seq_lens, softmax_mode=softmax_mode, **common ) + if softmax_mode == "mixed_fp16": + fp32_output = attention_decode( + q, k_cache, v_cache, block_table, seq_lens, softmax_mode="fp32", **common + ) + assert not torch.equal(output.to(k.dtype), fp32_output.to(k.dtype)) scores = torch.matmul(k[0, 0].float(), q[0, 0]) * (scale * 1.4426950408889634) - split_scores = scores.reshape(num_splits, -1) - split_max = split_scores.amax(dim=1) - p = torch.exp2(split_scores - split_max[:, None]) - p_qdq = _nvfp4_qdq_reference(p.to(torch.bfloat16).float()) - split_acc = torch.einsum("sk,skd->sd", p_qdq, v[0, 0].float().reshape(num_splits, -1, head_dim)) - running_max = torch.tensor(-float("inf"), device="cuda") - running_sum = torch.tensor(0.0, device="cuda") - acc = torch.zeros(head_dim, device="cuda") + split_scores = scores.reshape(num_splits, 2, block_n) + split_values = v[0, 0].float().reshape(num_splits, 2, block_n, head_dim) + split_max = [] + split_sum = [] + split_acc = [] for split_idx in range(num_splits): - new_max = torch.maximum(running_max, split_max[split_idx]) - correction = torch.exp2(running_max - new_max) - split_correction = torch.exp2(split_max[split_idx] - new_max) - acc = acc * correction + split_acc[split_idx] * split_correction - running_sum = running_sum * correction + p[split_idx].sum() * split_correction - running_max = new_max - reference = acc / running_sum + running_max = torch.tensor(-float("inf"), device="cuda") + running_sum = torch.tensor(0.0, device="cuda") + local_acc = torch.zeros(head_dim, device="cuda") + for tile_idx in range(2): + tile_scores = split_scores[split_idx, tile_idx] + new_max = torch.maximum(running_max, tile_scores.amax()) + p = _softmax_exp2(tile_scores - new_max, softmax_mode) + correction = _softmax_exp2(running_max - new_max, softmax_mode) + running_sum = running_sum * correction + p.sum() + local_acc *= correction + p_qdq = _nvfp4_qdq_reference(p.to(torch.bfloat16).float()) + local_acc += (p_qdq[:, None] * split_values[split_idx, tile_idx]).sum(dim=0) + running_max = new_max + split_max.append(running_max) + split_sum.append(running_sum) + split_acc.append(local_acc) + split_max = torch.stack(split_max) + split_sum = torch.stack(split_sum) + split_acc = torch.stack(split_acc) + reference = _combine_split_states(split_max, split_sum, split_acc, "fp32") + if softmax_mode == "mixed_fp16": + mixed_combine_reference = _combine_split_states( + split_max, split_sum, split_acc, "mixed_fp16" + ) + assert reference.to(k.dtype)[1] != mixed_combine_reference.to(k.dtype)[1] torch.testing.assert_close(output[0, 0], reference, rtol=5e-5, atol=5e-5) From 544e10a4ad39959d67ad2f9a493c0fd9f3cdae2f Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 8 Jul 2026 16:30:54 -0700 Subject: [PATCH 26/28] Propagate mixed-FP16 attention policy Signed-off-by: Kai Xu --- .../attention_sparsity/plugins/vllm.py | 22 ++++++--- .../test_sparse_attn_worker.py | 45 ++++++++++++++----- 2 files changed, 49 insertions(+), 18 deletions(-) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 0e25dbdadd3..0b283cf032b 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -149,23 +149,26 @@ def _v_qdq_from_layer(layer) -> tuple[str | None, float | None]: def _quant_kw_from_impl(impl, layer): - """Resolve the compact P/V QDQ contract once for one attention launch.""" + """Resolve the compact attention quantization contract for one launch.""" quant_kw = getattr(impl, "quant_kw", None) if quant_kw is None: p_qdq, p_qdq_amax = _p_qdq_from_layer(layer) v_qdq, v_qdq_amax = _v_qdq_from_layer(layer) + softmax_mode = "fp32" else: p_qdq, p_qdq_amax = quant_kw["p_qdq"], quant_kw["p_qdq_amax"] v_qdq, v_qdq_amax = quant_kw["v_qdq"], quant_kw["v_qdq_amax"] - return p_qdq, p_qdq_amax, v_qdq, v_qdq_amax + softmax_mode = quant_kw.get("softmax_mode", "fp32") + return p_qdq, p_qdq_amax, v_qdq, v_qdq_amax, softmax_mode -def _any_quant_active(layer, p_qdq, v_qdq) -> bool: - """Return whether native fallback would omit any Q/K/P/V transform.""" +def _any_quant_active(layer, p_qdq, v_qdq, softmax_mode) -> bool: + """Return whether native fallback would omit any quantized operation.""" k_quantizer = getattr(layer, "k_bmm_quantizer", None) return bool( p_qdq or v_qdq + or softmax_mode != "fp32" or getattr(layer, "_query_quant_in_kernel", False) or getattr(k_quantizer, "is_enabled", False) ) @@ -182,6 +185,7 @@ class _ResolvedForward: p_qdq_amax: float v_qdq: str | None v_qdq_amax: float | None + softmax_mode: str quant_active: bool @@ -195,8 +199,8 @@ def _resolve_forward( require_flashinfer_metadata: bool = False, ) -> _ResolvedForward | None: """Resolve shared transform state or request the backend's native path.""" - p_qdq, p_qdq_amax, v_qdq, v_qdq_amax = _quant_kw_from_impl(impl, layer) - quant_active = _any_quant_active(layer, p_qdq, v_qdq) + p_qdq, p_qdq_amax, v_qdq, v_qdq_amax, softmax_mode = _quant_kw_from_impl(impl, layer) + quant_active = _any_quant_active(layer, p_qdq, v_qdq, softmax_mode) transform_active = _should_run_modelopt_kernel(getattr(impl, "sparse_kw", None), quant_active) if getattr(attn_metadata, "use_cascade", False): @@ -224,6 +228,7 @@ def _resolve_forward( p_qdq_amax=p_qdq_amax, v_qdq=v_qdq, v_qdq_amax=v_qdq_amax, + softmax_mode=softmax_mode, quant_active=quant_active, ) @@ -249,6 +254,7 @@ def _forward_modelopt( p_qdq_amax: float, v_qdq: str | None, v_qdq_amax: float | None, + softmax_mode: str, quant_active: bool, dense_fallback, prepare_modelopt=None, @@ -317,6 +323,7 @@ def _forward_modelopt( v_qdq=v_qdq, v_qdq_amax=v_qdq_amax, v_cache_quantized=v_cache_quantized, + softmax_mode=softmax_mode, ) output[:batch] = triton_out return output @@ -344,6 +351,7 @@ def _forward_modelopt( v_qdq=v_qdq, v_qdq_amax=v_qdq_amax, v_cache_quantized=v_cache_quantized, + softmax_mode=softmax_mode, **sparse_kw, ) output[:num_actual_tokens] = triton_out @@ -439,6 +447,7 @@ def native_forward(): p_qdq_amax=resolved.p_qdq_amax, v_qdq=resolved.v_qdq, v_qdq_amax=resolved.v_qdq_amax, + softmax_mode=resolved.softmax_mode, quant_active=resolved.quant_active, dense_fallback=native_forward, ) @@ -612,6 +621,7 @@ def prepare_modelopt(): "p_qdq_amax": resolved.p_qdq_amax, "v_qdq": resolved.v_qdq, "v_qdq_amax": resolved.v_qdq_amax, + "softmax_mode": resolved.softmax_mode, "quant_active": resolved.quant_active, "dense_fallback": dense_fallback, "prepare_modelopt": prepare_modelopt, diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index e463b413787..37d91cf5447 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -183,7 +183,7 @@ def _make_old_flashinfer_impl(): return impl -def _make_flashinfer_impl(*, sparse=False, quantized=False): +def _make_flashinfer_impl(*, sparse=False, quantized=False, softmax_mode="fp32"): impl = _clone_sparse_impl(_make_old_flashinfer_impl(), get_flashinfer_sparse_impl_cls()) impl.sparse_kw = {"sparsity_n": 2, "sparsity_m": 4} if sparse else {} impl.quant_kw = { @@ -191,6 +191,7 @@ def _make_flashinfer_impl(*, sparse=False, quantized=False): "p_qdq_amax": 1.0, "v_qdq": "nvfp4" if quantized else None, "v_qdq_amax": 6.0 * 448.0 if quantized else None, + "softmax_mode": softmax_mode, } return impl @@ -405,17 +406,21 @@ def test_flashinfer_legacy_forward_writes_kv_cache( assert writes == [(layer, query, query, kv_cache, metadata, impl)] -def test_flashinfer_q_only_transform_does_not_fallback(monkeypatch): - """Withheld Q QDQ must run even when sparse/P/V transforms are inactive.""" - impl = _make_flashinfer_impl() +@pytest.mark.parametrize( + ("query_quantized", "softmax_mode"), + [(True, "fp32"), (False, "mixed_fp16")], + ids=["q-only", "mixed-softmax-only"], +) +def test_flashinfer_q_only_transform_does_not_fallback(monkeypatch, query_quantized, softmax_mode): + """Q QDQ and mixed softmax must run when other transforms are inactive.""" + impl = _make_flashinfer_impl(softmax_mode=softmax_mode) query = torch.zeros(4, 2, 64, dtype=torch.float16) kv_cache = torch.zeros(4, 2, 16, 2, 64, dtype=torch.float16) metadata = _flashinfer_metadata(query_lens=(2, 2), max_query_len=4) captured = {} - layer = SimpleNamespace( - _query_quant_in_kernel=True, - q_bmm_quantizer=lambda value: value + 1, - ) + layer = SimpleNamespace(_query_quant_in_kernel=query_quantized) + if query_quantized: + layer.q_bmm_quantizer = lambda value: value + 1 monkeypatch.setattr( FlashInferImpl, @@ -425,13 +430,16 @@ def test_flashinfer_q_only_transform_does_not_fallback(monkeypatch): def fake_attention(query, **kwargs): captured["query"] = query + captured.update(kwargs) return torch.zeros_like(query) monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) impl.forward(layer, query, query, query, kv_cache, metadata, output=torch.empty_like(query)) - assert torch.all(captured["query"] == 1) + expected_query = 1 if query_quantized else 0 + assert torch.all(captured["query"] == expected_query) + assert captured["softmax_mode"] == softmax_mode def test_flashinfer_legacy_inactive_launch_writes_only_in_native_fallback(monkeypatch): @@ -461,10 +469,18 @@ def test_flashinfer_legacy_inactive_launch_writes_only_in_native_fallback(monkey assert native_calls == [True] -@pytest.mark.parametrize("quantized", [False, True], ids=["sparse-only", "quantized"]) -def test_flashinfer_mixed_batch_splits_decode_and_prefill(monkeypatch, quantized): +@pytest.mark.parametrize( + ("quantized", "softmax_mode"), + [(False, "fp32"), (True, "mixed_fp16")], + ids=["sparse-only", "quantized-mixed-fp16"], +) +def test_flashinfer_mixed_batch_splits_decode_and_prefill(monkeypatch, quantized, softmax_mode): prefill_tokens = 17 - impl = _make_flashinfer_impl(sparse=True, quantized=quantized) + impl = _make_flashinfer_impl( + sparse=True, + quantized=quantized, + softmax_mode=softmax_mode, + ) query = torch.zeros(1 + prefill_tokens, 2, 64, dtype=torch.float16) kv_cache = torch.zeros(4, 2, 16, 2, 64, dtype=torch.float16) metadata = _flashinfer_metadata(query_lens=(1, prefill_tokens), mixed=True) @@ -502,6 +518,7 @@ def fake_attention(query, **kwargs): assert prefill_query.shape[0] == prefill_tokens assert prefill_kw["sparsity_n"] == 2 assert prefill_kw["sparsity_m"] == 4 + assert prefill_kw["softmax_mode"] == softmax_mode torch.testing.assert_close( prefill_kw["b_seq_len"], torch.tensor([prefill_tokens], dtype=torch.int32) ) @@ -512,6 +529,7 @@ def fake_attention(query, **kwargs): decode_query, decode_kw = calls["decode"][0] assert decode_query.shape[0] == 1 assert decode_kw["p_qdq"] == "nvfp4" + assert decode_kw["softmax_mode"] == "mixed_fp16" assert "sparsity_n" not in decode_kw assert prefill_kw["p_qdq"] == "nvfp4" assert [kw["max_new_tokens"] for kw in calls["finalize"]] == [1, prefill_tokens] @@ -910,6 +928,7 @@ def quantize_q(query): "p_qdq_amax": 1.0, "v_qdq": "nvfp4", "v_qdq_amax": 6.0 * 448.0, + "softmax_mode": "mixed_fp16", } q = torch.full((4, impl.num_heads, impl.head_size), 2.0, dtype=torch.float16) q[2:] = 10_000 @@ -963,6 +982,7 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): assert calls["query"].shape[0] == metadata.seq_lens.shape[0] assert decode_kw["p_qdq"] == "nvfp4" assert decode_kw["v_qdq"] == "nvfp4" + assert decode_kw["softmax_mode"] == "mixed_fp16" assert decode_kw["v_cache_quantized"] is True assert len(q_inputs) == 1 and q_inputs[0].shape[0] == q.shape[0] torch.testing.assert_close(q_inputs[0][:2], q[:2].float()) @@ -1002,6 +1022,7 @@ def fake_attention(query, **kwargs): assert captured["skip_softmax_threshold"] == pytest.approx(0.001) assert captured["p_qdq"] == "nvfp4" assert captured["v_qdq"] == "nvfp4" + assert captured["softmax_mode"] == "fp32" def test_resolve_calibrated_skip_softmax_threshold_for_decode(): From bf5d7e9e40cd7460a523215e59ae92b803471b57 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 8 Jul 2026 16:36:38 -0700 Subject: [PATCH 27/28] Select compact attention softmax mode Signed-off-by: Kai Xu --- examples/vllm_serve/sparse_attn_worker.py | 19 ++++- examples/vllm_serve/vllm_serve_sparse_attn.py | 27 ++++++ .../test_quant_sparse_attn_worker.py | 83 ++++++++++++++++++- 3 files changed, 126 insertions(+), 3 deletions(-) diff --git a/examples/vllm_serve/sparse_attn_worker.py b/examples/vllm_serve/sparse_attn_worker.py index 40a04585265..396c80c03e6 100644 --- a/examples/vllm_serve/sparse_attn_worker.py +++ b/examples/vllm_serve/sparse_attn_worker.py @@ -16,6 +16,7 @@ """Custom vLLM workers for checkpoint-driven sparse and fixed-NVFP4 attention.""" import importlib +import os from collections import Counter from functools import cache from types import SimpleNamespace @@ -33,6 +34,7 @@ from vllm.v1.worker.gpu_worker import Worker as BaseWorker +from modelopt.torch.kernels.quantization.attention.softmax_fakequant import SOFTMAX_MODES from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_config import ( load_from_checkpoint_metadata, match_sparse_config, @@ -47,6 +49,7 @@ __all__ = ["SparseAttnWorker", "QuantSparseAttnWorker"] # noqa: RUF022 +MODELOPT_ATTN_SOFTMAX_MODE = "MODELOPT_ATTN_SOFTMAX_MODE" _NVFP4_CFG = { "num_bits": (2, 1), "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, @@ -114,6 +117,16 @@ def _quant_api(): return _load_quant_api(vllm.__version__) +def _softmax_mode_from_env() -> str: + raw_mode = os.environ.get(MODELOPT_ATTN_SOFTMAX_MODE, "fp32") + mode = raw_mode.strip().lower() + if mode not in SOFTMAX_MODES: + raise ValueError( + f"{MODELOPT_ATTN_SOFTMAX_MODE} must be one of {sorted(SOFTMAX_MODES)}, got {raw_mode!r}" + ) + return mode + + def _cudagraph_mode(worker, api): config = getattr(worker.model_runner, "vllm_config", None) compilation = getattr(config, "compilation_config", None) @@ -288,7 +301,7 @@ def _install_sparse_plans(plans) -> None: ) -def _install_quant_plans(worker, plans) -> None: +def _install_quant_plans(worker, plans, *, softmax_mode: str) -> None: api = _quant_api() for plan in plans: module = plan.module @@ -305,6 +318,7 @@ def _install_quant_plans(worker, plans) -> None: "p_qdq_amax": p_qdq_amax, "v_qdq": v_qdq, "v_qdq_amax": v_qdq_amax, + "softmax_mode": softmax_mode, } module.impl = plan.new_impl module._query_quant_in_kernel = True @@ -316,7 +330,8 @@ def _install_quant_plans(worker, plans) -> None: def _install_attention(worker, *, quantize: bool) -> None: if quantize: - _install_quant_plans(worker, _quant_plans(worker)) + softmax_mode = _softmax_mode_from_env() + _install_quant_plans(worker, _quant_plans(worker), softmax_mode=softmax_mode) else: plans = _sparse_plans(worker) if plans is not None: diff --git a/examples/vllm_serve/vllm_serve_sparse_attn.py b/examples/vllm_serve/vllm_serve_sparse_attn.py index 777aff8fc7d..6c6559e9990 100644 --- a/examples/vllm_serve/vllm_serve_sparse_attn.py +++ b/examples/vllm_serve/vllm_serve_sparse_attn.py @@ -28,6 +28,7 @@ python vllm_serve_sparse_attn.py """ +import importlib import os import sys from pathlib import Path @@ -45,6 +46,32 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser +_MODELOPT_ATTN_SOFTMAX_MODE = "MODELOPT_ATTN_SOFTMAX_MODE" +_RAY_EXECUTOR_MODULES = ( + "vllm.executor.ray_distributed_executor", + "vllm.v1.executor.ray_distributed_executor", +) + + +def _propagate_env_var_to_ray_workers(env_var: str) -> None: + for module_name in _RAY_EXECUTOR_MODULES: + try: + executor = importlib.import_module(module_name).RayDistributedExecutor + executor.ADDITIONAL_ENV_VARS.update({env_var}) + except (ImportError, AttributeError): + continue + return + + extra_env_var = "VLLM_RAY_EXTRA_ENV_VARS_TO_COPY" + merged_env_vars = { + name.strip() for name in os.environ.get(extra_env_var, "").split(",") if name.strip() + } | {env_var} + os.environ[extra_env_var] = ",".join(sorted(merged_env_vars)) + + +_propagate_env_var_to_ray_workers(_MODELOPT_ATTN_SOFTMAX_MODE) + + def main(): """Launch vLLM with sparse attention worker.""" parser = FlexibleArgumentParser(description="vLLM model server with sparse attention") diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py index 4839be76614..9c81d637bfe 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py @@ -16,6 +16,7 @@ """Focused tests for the fixed-recipe quant+sparse vLLM worker.""" import importlib.util +import os from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace @@ -37,6 +38,7 @@ ) _WORKER_PATH = Path(__file__).parents[5] / "examples/vllm_serve/sparse_attn_worker.py" +_LAUNCHER_PATH = Path(__file__).parents[5] / "examples/vllm_serve/vllm_serve_sparse_attn.py" def _load_worker_module(): @@ -51,6 +53,13 @@ def _load_worker_module(): worker_module = _load_worker_module() +def _load_launcher_module(name="shared_attention_launcher_quant_test"): + spec = importlib.util.spec_from_file_location(name, _LAUNCHER_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + def test_quant_policy_rejects_old_vllm(monkeypatch): monkeypatch.setattr(vllm, "__version__", "0.9.0") worker_module._load_quant_api.cache_clear() @@ -97,8 +106,18 @@ def _patch_conversion(monkeypatch): @pytest.mark.parametrize("impl_cls", [FlashAttentionImpl, FlashInferImpl]) -def test_install_converts_only_attention_and_configures_fixed_recipe(monkeypatch, impl_cls): +@pytest.mark.parametrize( + ("softmax_env", "expected_softmax_mode"), + [(None, "fp32"), (" MiXeD_FP16 ", "mixed_fp16")], +) +def test_install_converts_only_attention_and_configures_fixed_recipe( + monkeypatch, impl_cls, softmax_env, expected_softmax_mode +): _patch_conversion(monkeypatch) + if softmax_env is None: + monkeypatch.delenv("MODELOPT_ATTN_SOFTMAX_MODE", raising=False) + else: + monkeypatch.setenv("MODELOPT_ATTN_SOFTMAX_MODE", softmax_env) attention = _attention(impl_cls=impl_cls) linear = nn.Linear(4, 4) model = nn.ModuleDict({"attn": attention, "linear": linear}) @@ -129,11 +148,73 @@ def test_install_converts_only_attention_and_configures_fixed_recipe(monkeypatch "p_qdq_amax": 1.0, "v_qdq": "nvfp4", "v_qdq_amax": 6.0 * 448.0, + "softmax_mode": expected_softmax_mode, } assert model["linear"] is linear assert state.model_runner.cascade_attn_enabled is False +def test_invalid_softmax_mode_is_rejected_before_quant_mutation(monkeypatch): + _patch_conversion(monkeypatch) + monkeypatch.setenv("MODELOPT_ATTN_SOFTMAX_MODE", " fp64 ") + attention = _attention() + model = nn.ModuleDict({"attn": attention}) + state = _worker(model) + original_impl = attention.impl + + with pytest.raises(ValueError) as exc: + worker_module._install_attention(state, quantize=True) + + assert "MODELOPT_ATTN_SOFTMAX_MODE" in str(exc.value) + assert " fp64 " in str(exc.value) + assert model["attn"] is attention + assert type(model["attn"]) is quant_plugin.vllm_attention.Attention + assert attention.impl is original_impl + assert state.model_runner.cascade_attn_enabled is True + + +@pytest.mark.parametrize( + "available_module", + [ + pytest.param("vllm.executor.ray_distributed_executor", id="legacy"), + pytest.param("vllm.v1.executor.ray_distributed_executor", id="v1"), + pytest.param(None, id="env-fallback"), + ], +) +def test_launcher_propagates_softmax_mode_to_ray_workers(monkeypatch, available_module): + ray_modules = ( + "vllm.executor.ray_distributed_executor", + "vllm.v1.executor.ray_distributed_executor", + ) + softmax_env = "MODELOPT_ATTN_SOFTMAX_MODE" + extra_env_var = "VLLM_RAY_EXTRA_ENV_VARS_TO_COPY" + existing_env = "EXISTING_ENV" + monkeypatch.setenv(extra_env_var, existing_env) + executor = SimpleNamespace(ADDITIONAL_ENV_VARS={"EXISTING_CLASS_VAR"}) + real_import_module = importlib.import_module + attempted_ray_modules = [] + + def import_module(name, package=None): + if name not in ray_modules: + return real_import_module(name, package) + attempted_ray_modules.append(name) + if name == available_module: + return SimpleNamespace(RayDistributedExecutor=executor) + raise ModuleNotFoundError(name) + + monkeypatch.setattr(importlib, "import_module", import_module) + + _load_launcher_module(f"shared_attention_launcher_quant_test_{available_module}") + + expected_attempts = ray_modules[:1] if available_module == ray_modules[0] else ray_modules + assert tuple(attempted_ray_modules) == expected_attempts + if available_module is not None: + assert {"EXISTING_CLASS_VAR", softmax_env} == executor.ADDITIONAL_ENV_VARS + assert os.environ[extra_env_var] == existing_env + else: + assert set(os.environ[extra_env_var].split(",")) == {existing_env, softmax_env} + + def test_validation_of_all_layouts_precedes_mutation(monkeypatch): _patch_conversion(monkeypatch) good, bad = _attention(), _attention(AttentionType.ENCODER) From a3f54acbb267e285cfe39be1ad9b486660df1a72 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 8 Jul 2026 16:48:09 -0700 Subject: [PATCH 28/28] Document mixed-FP16 attention serving Signed-off-by: Kai Xu --- .../2026-07-08-mixed-fp16-softmax-design.md | 88 ------------------- examples/vllm_serve/README.md | 16 ++-- 2 files changed, 10 insertions(+), 94 deletions(-) delete mode 100644 docs/superpowers/specs/2026-07-08-mixed-fp16-softmax-design.md diff --git a/docs/superpowers/specs/2026-07-08-mixed-fp16-softmax-design.md b/docs/superpowers/specs/2026-07-08-mixed-fp16-softmax-design.md deleted file mode 100644 index e641f4797b1..00000000000 --- a/docs/superpowers/specs/2026-07-08-mixed-fp16-softmax-design.md +++ /dev/null @@ -1,88 +0,0 @@ -# Mixed-FP16 Softmax Design - -## Goal - -Add the reference mixed-FP16 online-softmax path to the compact NVFP4 attention worker while preserving one public `QuantSparseAttnWorker` and the current fixed block-16 NVFP4 Q/K/P/V recipe. - -The worker must select either the existing FP32 softmax or mixed-FP16 softmax before model installation and CUDA-graph capture. This is inference-only support for regular decoder self-attention through the existing FlashAttention and FlashInfer adapters. - -## Configuration Boundary - -Q/K/P/V formats remain `TensorQuantizer` configuration because they describe tensor QDQ operations. Mixed-FP16 softmax is a kernel compute policy: it changes the two online-softmax exponentials but is not a tensor QDQ format. - -Use one compact-only environment selector: - -```text -MODELOPT_ATTN_SOFTMAX_MODE=fp32 -MODELOPT_ATTN_SOFTMAX_MODE=mixed_fp16 -``` - -The default is `fp32`. Any other value fails before an attention layer is mutated. The launcher must forward the selector to Ray workers so every tensor-parallel worker installs the same static policy. - -The worker resolves the value once and stores the normalized string in the existing immutable launch-policy snapshot: - -```python -impl.quant_kw["softmax_mode"] = "fp32" | "mixed_fp16" -``` - -No additional worker class, softmax `TensorQuantizer`, generic per-point precision matrix, or integration-branch compatibility alias is added. - -## Numerical Contract - -`fp32` preserves current behavior. - -`mixed_fp16` implements the reference path: - -- Convert the input of `exp2(scores - new_max)` from FP32 to FP16 with round-to-nearest-even, execute native `ex2.approx.f16`, and return its FP16-valued result as FP32. -- Apply the same native FP16 operation to the online correction `exp2(old_max - new_max)`. -- Sum probabilities into the denominator in FP32 without another rounding step. -- Keep the denominator state, weighted-value accumulator, and matrix accumulators in FP32. -- Apply existing NVFP4 P QDQ after the mixed-FP16 probability computation and after the unquantized FP32 denominator sum. -- Keep split-K reconciliation in FP32. Mixed FP16 applies only inside each split's online-softmax loop. -- Leave Q/K/P/V QDQ, V-cache finalization, sparse masking, and skip-softmax decisions unchanged. - -The mode is a Triton compile-time constant. It introduces no device scalar, host synchronization, or graph-time mutable state. - -## Data Flow - -```text -MODELOPT_ATTN_SOFTMAX_MODE - -> worker validation and normalization - -> impl.quant_kw["softmax_mode"] - -> _ResolvedForward - -> shared FlashAttention/FlashInfer _forward_modelopt path - -> triton_fa.attention / decode_attention.attention_decode - -> MIXED_FP16 constexpr -``` - -The mode counts as an active ModelOpt transform. A non-FP32 mode must therefore never delegate to the backend's native dense path, including FlashInfer mixed prefill/decode batches. - -## API And Error Handling - -The public prefill and decode kernel wrappers accept `softmax_mode="fp32"` and validate against the two supported values. Internal Triton kernels receive only a boolean `MIXED_FP16` constexpr. - -Autograd with `mixed_fp16` raises `NotImplementedError`; backward recomputation is outside this inference-focused change. Existing autograd behavior for `fp32` remains unchanged. - -Worker configuration is read before plan installation. Invalid configuration fails atomically, without converting modules, replacing implementations, or partially installing a policy. - -## Testing - -Keep checked-in coverage focused on behavioral contracts: - -- Worker selection defaults to `fp32`, accepts `mixed_fp16`, rejects invalid modes atomically, and forwards the selector to Ray workers. -- FlashAttention and FlashInfer routes propagate the mode for prefill, decode, and FlashInfer mixed batches without native fallback. -- Prefill matches a mixed-FP16 reference while composing with NVFP4 Q/K/P/V and 2:4 sparsity. -- Fixed 32-split decode matches a split-local mixed-FP16 reference with FP32 split reconciliation. -- The native FP16 `ex2` helper covers representative finite inputs and masked negative infinity. -- Mixed-FP16 autograd is rejected before launch. -- Default `fp32` routing and numerics remain unchanged. - -CUDA-graph support is claimed only after a capture/replay smoke confirms that the static mode is preserved. - -## Non-Goals - -- Porting the integration branch's `MODELOPT_ATTN_FP16_SOFTMAX` alias or generic `MODELOPT_ATTN_SOFTMAX_QUANT` matrix. -- Supporting independently configurable DIFF, EXP2, DELTA, ALPHA, ACC, reciprocal, or output formats. -- Changing the fixed NVFP4 recipe, split count, tile sizes, sparse behavior, or supported vLLM attention envelope. -- Implementing mixed-FP16 backward kernels. -- Adding the design document to the final pull request. diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index b816218c1bc..d71c60535b2 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -128,25 +128,29 @@ Limitations: vLLM 0.14.0 or newer is checked when `QuantSparseAttnWorker` is selected. Importing or using `SparseAttnWorker` does not resolve quant-only APIs. -Use the same launcher with the compact worker. By default, vLLM selects the backend for the model and platform; NemotronH on Blackwell selects FlashInfer: +Use the same launcher with the compact worker. By default, vLLM selects the backend for the model and platform; NemotronH on Blackwell selects FlashInfer. Select the mixed-FP16 softmax policy explicitly and use eager execution because CUDA graph capture has not been validated for this path: ```bash +MODELOPT_ATTN_SOFTMAX_MODE=mixed_fp16 \ python vllm_serve_sparse_attn.py -tp 8 \ + --enforce-eager \ --no-enable-prefix-caching \ --worker-cls sparse_attn_worker.QuantSparseAttnWorker ``` The worker supports both FlashInfer and FlashAttention and prints the installed adapter counts. Pass `--attention-backend FLASHINFER` or `--attention-backend FLASH_ATTN` only when an explicit override is needed. -This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant recipe to Q/K/P/V. Q is dynamic, K/V use global scale 1.0, and P uses amax 1.0; calibrated attention amax restore is not part of this fixed path. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored. +`MODELOPT_ATTN_SOFTMAX_MODE` accepts `fp32` (the default) or `mixed_fp16`. Values are stripped and lowercased; an invalid value fails before adapters are installed. The launcher forwards the variable to Ray workers. -Decode uses a fixed 32-split, 128-key-tile schedule. P QDQ consumes split-local, -unnormalized online-softmax probabilities, so changing that schedule can change -quantized results; split count is part of the numerical contract. +Mixed mode uses native FP16 `exp2` only for tile probabilities and the online-maximum correction, then converts those results back to FP32. Denominator reductions and running state, weighted-value accumulators, and split reconciliation remain FP32. P NVFP4 QDQ happens after the unquantized denominator update. Native FP16 `exp2` requires SM75 or newer, while this worker's NVFP4 Q/K/P/V path uses native E4M3 and requires SM89 or newer. The mode is inference-only; backward and autograd are unsupported. + +This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant recipe to Q/K/P/V. Q is dynamic, K/V use global scale 1.0, and P uses amax 1.0; calibrated attention amax restore is not part of this fixed path. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored, so NVFP4 and mixed softmax can compose with N:M sparsity. N:M sparsity applies only to prefill. + +Decode-only launches use a fixed 32-split, 128-key-tile schedule when NVFP4 P or V is active and decode skip-softmax is not active. P QDQ consumes split-local, unnormalized online-softmax probabilities, so changing that schedule can change quantized results; split count is part of the numerical contract. Split reconciliation remains FP32. K is QDQ before its cache write, while V is written pristine. Complete 16-token V groups are finalized once in cache; an incomplete tail remains pristine and is QDQ on read. P@V therefore sees uniform fakequant values without re-quantizing the tail. -Supported configurations are regular decoder self-attention with FlashInfer or FlashAttention, fp16/bf16 model and KV cache, equal Q/K/V head dimensions that are multiples of 16, and DCP 1. The FlashInfer adapter preserves both NHD and HND cache strides and separates mixed decode/prefill launches so each phase keeps its own kernel contract. The default `FULL_AND_PIECEWISE` mode remains enabled for fixed N:M and attention-only NVFP4; checkpoints with calibrated decode `threshold_scale_factor` must use a non-`FULL` decode graph mode such as `--enforce-eager` because the live sequence length is not replayed as a Python scalar. +Supported configurations are regular decoder self-attention with FlashInfer or FlashAttention, fp16/bf16 model and KV cache, equal Q/K/V head dimensions that are multiples of 16, and DCP 1. The FlashInfer adapter preserves both NHD and HND cache strides and separates mixed decode/prefill launches so each phase keeps its own kernel contract. Checkpoints with calibrated decode `threshold_scale_factor` must use a non-`FULL` decode graph mode such as `--enforce-eager` because the live sequence length is not replayed as a Python scalar. Unsupported features are sliding window, ALiBi, softcap, sinks, FP8 KV cache, cross/encoder/MLA attention, KV sharing or transfer, prefix caching, speculative decoding, DBO/ubatching, and `FULL` mixed/prefill CUDA graphs.