From 21f697fa7d049092a1818612165e5c6ca7f2618f Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Thu, 2 Jul 2026 22:39:38 -0700 Subject: [PATCH 01/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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 117725f8ad74836d3d5b94c87ea89d0f3458032e Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 7 Jul 2026 17:58:47 -0700 Subject: [PATCH 20/20] Add skip-softmax to dedicated decode attention Signed-off-by: Kai Xu --- .../common/attention/decode_attention.py | 132 ++++++++++------ .../attention_sparsity/plugins/vllm.py | 13 +- .../common/attention/test_decode_attention.py | 142 ++++++++++++++++++ .../test_sparse_attn_worker.py | 64 ++++++-- 4 files changed, 284 insertions(+), 67 deletions(-) diff --git a/modelopt/torch/kernels/common/attention/decode_attention.py b/modelopt/torch/kernels/common/attention/decode_attention.py index d12f66f16eb..bb4bd8ed873 100644 --- a/modelopt/torch/kernels/common/attention/decode_attention.py +++ b/modelopt/torch/kernels/common/attention/decode_attention.py @@ -14,10 +14,10 @@ # limitations under the License. -"""Split-K decode attention for the ModelOpt paged NVFP4 serving path. +"""Paged decode attention for the ModelOpt 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. +Normal decode uses fixed-size split-K tiles. Skip-softmax instead scans 32-key +tiles sequentially so its running maximum and tile decisions match prefill. """ import math @@ -37,7 +37,9 @@ __all__ = ["attention_decode"] _BLOCK_N = 128 +_SKIP_BLOCK_N = 32 _DEFAULT_KV_SPLITS = 32 +_SKIP_KV_SPLITS = 1 _MAX_KV_SPLITS = 32 _QDQ_MODES = {None, "nvfp4"} @@ -78,6 +80,8 @@ def _decode_split_kernel( P_QDQ: tl.constexpr, V_QDQ: tl.constexpr, V_CACHE_QUANTIZED: tl.constexpr, + APPLY_SKIP: tl.constexpr, + skip_threshold_log2, ): """Compute one partial softmax for one request, query head, and KV split.""" batch_idx = tl.program_id(0) @@ -131,53 +135,58 @@ def _decode_split_kernel( 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 + skip_tile = False + if APPLY_SKIP: + skip_tile = tile_max < running_max + skip_threshold_log2 - if P_QDQ: - p = tl.reshape( - _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,), - ) + if not skip_tile: + 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 - 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) - 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) - else: - v = v_qdq + if P_QDQ: + p = tl.reshape( + _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,), + ) + + 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) + 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) + else: + v = v_qdq - acc += tl.sum(p[:, None] * v, axis=0) - running_max = new_max + 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) @@ -260,6 +269,7 @@ def attention_decode( softmax_scale: float | None = None, page_size: int = 16, num_kv_splits: int = _DEFAULT_KV_SPLITS, + skip_softmax_threshold: float | None = None, p_qdq: str | None = None, p_qdq_amax: float = 1.0, v_qdq: str | None = None, @@ -273,15 +283,35 @@ 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. + Without skip-softmax, P QDQ follows the split-local 32-split, 128-key-tile + schedule by default; explicit split counts retain that same tile size. A + finite positive ``skip_softmax_threshold`` instead uses one sequential split + with 32-key tiles and skips tiles whose contribution is negligible versus + the running maximum. Finite non-positive values disable skipping. In skip + mode, ``num_kv_splits`` must be the legacy default (which is overridden) or + ``1``; other explicit values are rejected. """ + if skip_softmax_threshold is not None and not math.isfinite(skip_softmax_threshold): + raise ValueError(f"skip_softmax_threshold must be finite, got {skip_softmax_threshold}") 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}") + if skip_softmax_threshold is not None and skip_softmax_threshold > 0.0: + apply_skip = True + skip_threshold_log2 = math.log2(skip_softmax_threshold) + else: + apply_skip = False + skip_threshold_log2 = 0.0 + if apply_skip and num_kv_splits not in {_DEFAULT_KV_SPLITS, _SKIP_KV_SPLITS}: + raise ValueError( + "skip-softmax requires num_kv_splits=1; omit num_kv_splits to use its owned schedule" + ) + if apply_skip: + num_kv_splits = _SKIP_KV_SPLITS + block_n = _SKIP_BLOCK_N if apply_skip else _BLOCK_N batch, num_q_heads, head_dim = q.shape num_kv_heads = k_cache.shape[2] if num_q_heads % num_kv_heads: @@ -333,7 +363,7 @@ def attention_decode( v_qdq_scale, kv_group_num=num_q_heads // num_kv_heads, BLOCK_D=block_d, - BLOCK_N=_BLOCK_N, + BLOCK_N=block_n, HEAD_DIM=head_dim, PAGE_SIZE=page_size, max_blocks_per_seq=block_table.shape[1], @@ -341,6 +371,8 @@ def attention_decode( P_QDQ=p_qdq == "nvfp4", V_QDQ=v_qdq == "nvfp4", V_CACHE_QUANTIZED=v_cache_quantized, + APPLY_SKIP=apply_skip, + skip_threshold_log2=skip_threshold_log2, num_warps=4, num_stages=2, ) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index e15c215b522..64877832cd1 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -46,6 +46,8 @@ 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 +_DEDICATED_DECODE_SPARSE_KWARGS = frozenset({"skip_softmax_threshold"}) + def _target_sparse_ratio_for_phase(target_sparse_ratio, phase: str) -> float: """Return target sparsity for a phase, defaulting old checkpoint metadata.""" @@ -317,12 +319,10 @@ def _forward_modelopt( 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: + # The earlier active-work gate makes empty sparse kwargs valid for Q/K or forced decode; + # unknown sparse keys remain on the shared kernel. + use_dedicated_decode = is_decode_only and sparse_kw.keys() <= _DEDICATED_DECODE_SPARSE_KWARGS + if use_dedicated_decode: triton_out = triton_decode_attention( q[:batch], key_cache, @@ -331,6 +331,7 @@ def _forward_modelopt( seq_lens, softmax_scale=impl.scale, page_size=page_size, + skip_softmax_threshold=sparse_kw.get("skip_softmax_threshold"), p_qdq=p_qdq, p_qdq_amax=p_qdq_amax, v_qdq=v_qdq, 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..c88a41a6962 100644 --- a/tests/gpu/torch/kernels/common/attention/test_decode_attention.py +++ b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py @@ -24,6 +24,7 @@ if TRITON_KERNEL_AVAILABLE: from modelopt.torch.kernels.common.attention.decode_attention import attention_decode + from modelopt.torch.kernels.common.attention.triton_fa import attention as triton_fa_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) @@ -109,6 +110,147 @@ def test_split_k_varlen_gqa_matches_dense(num_kv_splits): ) +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") +@pytest.mark.parametrize( + "skip_softmax_threshold", + [ + pytest.param(float("nan"), id="nan"), + pytest.param(float("inf"), id="positive-inf"), + pytest.param(float("-inf"), id="negative-inf"), + ], +) +def test_skip_softmax_rejects_non_finite_threshold(skip_softmax_threshold): + q = torch.zeros(1, 1, 16, device="cuda", dtype=torch.float16) + k = torch.zeros(1, 1, 1, 16, device="cuda", dtype=torch.float16) + v = torch.zeros_like(k) + k_cache, v_cache, block_table = _paged_cache(k, v, (1,)) + seq_lens = torch.ones(1, device="cuda", dtype=torch.int32) + + with pytest.raises(ValueError, match="finite"): + attention_decode( + q, + k_cache, + v_cache, + block_table, + seq_lens, + skip_softmax_threshold=skip_softmax_threshold, + ) + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") +@pytest.mark.parametrize( + "qdq_kwargs", + [ + pytest.param({}, id="dense"), + pytest.param( + {"p_qdq": "nvfp4", "v_qdq": "nvfp4"}, + marks=requires_native_e4m3, + id="nvfp4-p-v", + ), + ], +) +def test_skip_softmax_matches_32_token_paged_reference(qdq_kwargs): + seq_len, head_dim = 96, 64 + scale = head_dim**-0.5 + skip_softmax_threshold = 0.1 + q = torch.zeros(1, 1, head_dim, device="cuda", dtype=torch.float16) + q[..., 0] = 1.0 + k = torch.zeros(1, 1, seq_len, head_dim, device="cuda", dtype=torch.float16) + k[:, :, 32:, 0] = -4.0 / scale + v = torch.full_like(k, 4.0) + v[:, :, :32] = -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) + decode_kwargs = { + "softmax_scale": scale, + **qdq_kwargs, + } + + no_skip_output = attention_decode( + q, + k_cache, + v_cache, + block_table, + seq_lens, + **decode_kwargs, + ) + + # The reference forward configs use BLOCK_N=32, so only the first tile contributes. + k_flat = k[0].transpose(0, 1).contiguous() + v_flat = v[0].transpose(0, 1).contiguous() + query_starts = torch.zeros(1, device="cuda", dtype=torch.int32) + query_lens = torch.ones(1, device="cuda", dtype=torch.int32) + reference = triton_fa_attention( + q, + k_flat, + v_flat, + query_starts, + query_lens, + 1, + is_causal=False, + softmax_scale=scale, + b_start_loc_k=query_starts, + b_seq_len_k=seq_lens, + max_input_len_k=seq_len, + skip_softmax_threshold=skip_softmax_threshold, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + **qdq_kwargs, + ) + # Skip decode owns the sequential one-split, BLOCK_N=32 schedule. + skip_output = attention_decode( + q, + k_cache, + v_cache, + block_table, + seq_lens, + skip_softmax_threshold=skip_softmax_threshold, + **decode_kwargs, + ) + + torch.testing.assert_close(skip_output, reference, rtol=5e-3, atol=5e-3) + assert (skip_output.float() - no_skip_output.float()).abs().max() > 0.05 + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") +def test_skip_softmax_cuda_graph_replay_uses_live_seq_len(): + captured_seq_len, replay_seq_len, head_dim = 64, 96, 64 + scale = head_dim**-0.5 + q = torch.zeros(1, 1, head_dim, device="cuda", dtype=torch.float16) + q[..., 0] = 1.0 + k = torch.zeros(1, 1, replay_seq_len, head_dim, device="cuda", dtype=torch.float16) + k[:, :, 32:64, 0] = -4.0 / scale + v = torch.ones_like(k) + v[:, :, :32] = -1.0 + v[:, :, 32:64] = 4.0 + k_cache, v_cache, block_table = _paged_cache(k, v, (replay_seq_len,)) + seq_lens = torch.tensor([captured_seq_len], device="cuda", dtype=torch.int32) + decode_kwargs = {"softmax_scale": scale, "skip_softmax_threshold": 0.1} + + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + for _ in range(3): + attention_decode(q, k_cache, v_cache, block_table, seq_lens, **decode_kwargs) + torch.cuda.current_stream().wait_stream(warmup_stream) + + captured_length_output = attention_decode( + q, k_cache, v_cache, block_table, seq_lens, **decode_kwargs + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = attention_decode(q, k_cache, v_cache, block_table, seq_lens, **decode_kwargs) + + seq_lens.fill_(replay_seq_len) + graph.replay() + replay_output = graph_output.clone() + eager_output = attention_decode(q, k_cache, v_cache, block_table, seq_lens, **decode_kwargs) + + torch.testing.assert_close(replay_output, eager_output, rtol=5e-3, atol=5e-3) + assert (eager_output.float() - captured_length_output.float()).abs().max() > 0.5 + + @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") @requires_native_e4m3 @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 72e50753600..e466a1b552f 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 @@ -524,7 +524,9 @@ def fake_attention(query, **kwargs): assert calls["finalize"] == [] -def test_flashinfer_forced_kernel_mixed_batch_never_uses_native_fallback(monkeypatch): +def test_flashinfer_forced_kernel_mixed_batch_uses_dedicated_decode_and_shared_prefill( + monkeypatch, +): class Layer: force_kernel_reads = 0 @@ -540,18 +542,18 @@ def _modelopt_force_kernel(self): 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} + calls = {"native": 0, "dedicated_decode": 0, "shared_decode": 0, "shared_prefill": 0} def native_fallback(*args, **kwargs): calls["native"] += 1 raise AssertionError("native FlashInfer fallback") def fake_decode(query, *args, **kwargs): - calls["decode"] += 1 + calls["dedicated_decode"] += 1 return torch.zeros_like(query) def fake_attention(query, **kwargs): - phase = "decode" if kwargs["max_input_len"] == 1 else "prefill" + phase = "shared_decode" if kwargs["max_input_len"] == 1 else "shared_prefill" calls[phase] += 1 return torch.zeros_like(query) @@ -561,7 +563,12 @@ 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 calls == { + "native": 0, + "dedicated_decode": 1, + "shared_decode": 0, + "shared_prefill": 1, + } assert layer.force_kernel_reads == 1 @@ -1024,8 +1031,8 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): assert calls["query"].dtype == torch.float32 -def test_quantized_skip_softmax_decode_stays_on_shared_kernel(monkeypatch): - """Split-local maxima must not change calibrated skip-softmax semantics.""" +def test_quantized_skip_softmax_decode_uses_dedicated_kernel(monkeypatch): + """Dedicated decode must receive skip-softmax and P/V QDQ configuration.""" impl = _clone_sparse_impl(_make_old_impl()) impl.quant_kw = { "p_qdq": "nvfp4", @@ -1040,10 +1047,45 @@ def test_quantized_skip_softmax_decode_stays_on_shared_kernel(monkeypatch): captured = {} monkeypatch.setattr(vllm_plugin, "fake_quant_v_onwrite", lambda *args, **kwargs: None) + + def fake_decode(query, *args, **kwargs): + captured.update(kwargs) + return torch.zeros_like(query) + + monkeypatch.setattr(vllm_plugin, "triton_decode_attention", fake_decode) + monkeypatch.setattr( + vllm_plugin, + "triton_attention", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("shared kernel")), + ) + output = torch.empty_like(q) + assert impl.forward(SimpleNamespace(), 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["p_qdq_amax"] == pytest.approx(1.0) + assert captured["v_qdq"] == "nvfp4" + assert captured["v_qdq_amax"] == pytest.approx(6.0 * 448.0) + + +def test_decode_with_unrecognized_sparse_kw_uses_shared_kernel(monkeypatch): + """Decode must preserve future sparse options by routing them through the shared kernel.""" + impl = _clone_sparse_impl(_make_old_impl()) + impl.quant_kw = { + "p_qdq": "nvfp4", + "p_qdq_amax": 1.0, + "v_qdq": None, + "v_qdq_amax": None, + } + impl.sparse_kw = {"future_sparse_option": True} + 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 = _flash_attention_metadata(1, 16) + captured = {} + monkeypatch.setattr( vllm_plugin, "triton_decode_attention", - lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("split-K kernel")), + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("dedicated kernel")), ) def fake_attention(query, **kwargs): @@ -1052,10 +1094,10 @@ def fake_attention(query, **kwargs): 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 impl.forward(SimpleNamespace(), q, q, q, kv_cache, metadata, output=output) is output + assert captured["future_sparse_option"] is True assert captured["p_qdq"] == "nvfp4" - assert captured["v_qdq"] == "nvfp4" + assert captured["p_qdq_amax"] == pytest.approx(1.0) def test_resolve_calibrated_skip_softmax_threshold_for_decode():