diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index a8f7d2ead18..d71c60535b2 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 @@ -101,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: @@ -114,12 +115,44 @@ 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. + +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. -- 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 + +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. Select the mixed-FP16 softmax policy explicitly and use eager execution because CUDA graph capture has not been validated for this path: + +```bash +MODELOPT_ATTN_SOFTMAX_MODE=mixed_fp16 \ +python vllm_serve_sparse_attn.py -tp 8 \ + --enforce-eager \ + --no-enable-prefix-caching \ + --worker-cls sparse_attn_worker.QuantSparseAttnWorker +``` + +The worker supports both FlashInfer and FlashAttention and prints the installed adapter counts. Pass `--attention-backend FLASHINFER` or `--attention-backend FLASH_ATTN` only when an explicit override is needed. + +`MODELOPT_ATTN_SOFTMAX_MODE` accepts `fp32` (the default) or `mixed_fp16`. Values are stripped and lowercased; an invalid value fails before adapters are installed. The launcher forwards the variable to Ray workers. + +Mixed mode uses native FP16 `exp2` only for tile probabilities and the online-maximum correction, then converts those results back to FP32. Denominator reductions and running state, weighted-value accumulators, and split reconciliation remain FP32. P NVFP4 QDQ happens after the unquantized denominator update. Native FP16 `exp2` requires SM75 or newer, while this worker's NVFP4 Q/K/P/V path uses native E4M3 and requires SM89 or newer. The mode is inference-only; backward and autograd are unsupported. + +This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant recipe to Q/K/P/V. Q is dynamic, K/V use global scale 1.0, and P uses amax 1.0; calibrated attention amax restore is not part of this fixed path. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored, so NVFP4 and mixed softmax can compose with N:M sparsity. N:M sparsity applies only to prefill. + +Decode-only launches use a fixed 32-split, 128-key-tile schedule when NVFP4 P or V is active and decode skip-softmax is not active. P QDQ consumes split-local, unnormalized online-softmax probabilities, so changing that schedule can change quantized results; split count is part of the numerical contract. Split reconciliation remains FP32. + +K is QDQ before its cache write, while V is written pristine. Complete 16-token V groups are finalized once in cache; an incomplete tail remains pristine and is QDQ on read. P@V therefore sees uniform fakequant values without re-quantizing the tail. + +Supported configurations are regular decoder self-attention with FlashInfer or FlashAttention, fp16/bf16 model and KV cache, equal Q/K/V head dimensions that are multiples of 16, and DCP 1. The FlashInfer adapter preserves both NHD and HND cache strides and separates mixed decode/prefill launches so each phase keeps its own kernel contract. 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. ## Known Problems diff --git a/examples/vllm_serve/sparse_attn_worker.py b/examples/vllm_serve/sparse_attn_worker.py index 1057baa870e..396c80c03e6 100644 --- a/examples/vllm_serve/sparse_attn_worker.py +++ b/examples/vllm_serve/sparse_attn_worker.py @@ -13,27 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""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. - -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. - -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. - -Usage: - python vllm_serve_sparse_attn.py -""" +"""Custom vLLM workers for checkpoint-driven sparse and fixed-NVFP4 attention.""" import importlib +import os +from collections import Counter +from functools import cache +from types import SimpleNamespace +from typing import NamedTuple try: _has_legacy_attention_layer = importlib.util.find_spec("vllm.attention.layer") is not None @@ -47,6 +34,7 @@ from vllm.v1.worker.gpu_worker import Worker as BaseWorker +from modelopt.torch.kernels.quantization.attention.softmax_fakequant import SOFTMAX_MODES from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_config import ( load_from_checkpoint_metadata, match_sparse_config, @@ -54,67 +42,320 @@ 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 + +MODELOPT_ATTN_SOFTMAX_MODE = "MODELOPT_ATTN_SOFTMAX_MODE" +_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): + module: object + new_impl: object + sparse_kw: dict + device: object | None + dtype: object | None + + +def _unwrapped_model(worker): + model = worker.model_runner.model + 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. + 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 _softmax_mode_from_env() -> str: + raw_mode = os.environ.get(MODELOPT_ATTN_SOFTMAX_MODE, "fp32") + mode = raw_mode.strip().lower() + if mode not in SOFTMAX_MODES: + raise ValueError( + f"{MODELOPT_ATTN_SOFTMAX_MODE} must be one of {sorted(SOFTMAX_MODES)}, got {raw_mode!r}" + ) + return mode + + +def _cudagraph_mode(worker, api): + config = getattr(worker.model_runner, "vllm_config", None) + compilation = getattr(config, "compilation_config", None) + mode = getattr(compilation, "cudagraph_mode", None) + return mode if mode is not None else api.compilation.CUDAGraphMode.NONE + + +def _global_errors(worker, api) -> list[str]: + config = worker.model_runner.vllm_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_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") + 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_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 -def _replace_attention_impl(worker): - """Replace FlashAttentionImpl with ModelOptSparseAttentionImpl on all Attention layers. - 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) +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 _sparse_graph_error(sparse_kw: dict, mode, api) -> 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() == 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 _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) + 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 - 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() - - patched = 0 + 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 - - layer_cfg = match_sparse_config(name, cfg) - if layer_cfg is None or not layer_cfg.get("enable", True): + 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) - 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. + +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 + 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, api.plugin._ATTENTION_TYPES): continue - new_impl = _clone_sparse_impl(module.impl) - new_impl.sparse_kw = sparse_kw - module.impl = new_impl - patched += 1 - print(f"[ModelOpt] Sparse attention: replaced impl on {patched} attention layers") + 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"{name or ''}: {reason}" for reason in reasons) + else: + plans.append(_AttentionPlan(module, new_impl, sparse_kw, device, dtype)) + if attention_count == 0: + errors.append("no regular attention layers were found") + _raise_unsupported(errors, "attention") + return tuple(plans) + + +def _install_sparse_plans(plans) -> None: + for plan in plans: + plan.new_impl.sparse_kw = plan.sparse_kw + plan.module.impl = plan.new_impl + 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}" + ) + +def _install_quant_plans(worker, plans, *, softmax_mode: str) -> None: + api = _quant_api() + 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) + 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, + "softmax_mode": softmax_mode, + } + module.impl = plan.new_impl + module._query_quant_in_kernel = True + module._value_quant_in_kernel = True + worker.model_runner.cascade_attn_enabled = False + installed = dict(Counter(type(plan.new_impl).__name__ for plan in plans)) + print(f"[ModelOpt] Installed NVFP4 quant+sparse attention on {len(plans)} layers: {installed}") -# --------------------------------------------------------------------------- -# Workers -# --------------------------------------------------------------------------- +def _install_attention(worker, *, quantize: bool) -> None: + if quantize: + softmax_mode = _softmax_mode_from_env() + _install_quant_plans(worker, _quant_plans(worker), softmax_mode=softmax_mode) + else: + plans = _sparse_plans(worker) + if plans is not None: + _install_sparse_plans(plans) -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). - """ +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/examples/vllm_serve/vllm_serve_sparse_attn.py b/examples/vllm_serve/vllm_serve_sparse_attn.py index e65ae3e44fb..6c6559e9990 100644 --- a/examples/vllm_serve/vllm_serve_sparse_attn.py +++ b/examples/vllm_serve/vllm_serve_sparse_attn.py @@ -15,19 +15,20 @@ """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 sparse_attn_worker.QuantSparseAttnWorker`` for quant+sparse. Usage: python vllm_serve_sparse_attn.py """ +import importlib import os import sys from pathlib import Path @@ -45,6 +46,32 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser +_MODELOPT_ATTN_SOFTMAX_MODE = "MODELOPT_ATTN_SOFTMAX_MODE" +_RAY_EXECUTOR_MODULES = ( + "vllm.executor.ray_distributed_executor", + "vllm.v1.executor.ray_distributed_executor", +) + + +def _propagate_env_var_to_ray_workers(env_var: str) -> None: + for module_name in _RAY_EXECUTOR_MODULES: + try: + executor = importlib.import_module(module_name).RayDistributedExecutor + executor.ADDITIONAL_ENV_VARS.update({env_var}) + except (ImportError, AttributeError): + continue + return + + extra_env_var = "VLLM_RAY_EXTRA_ENV_VARS_TO_COPY" + merged_env_vars = { + name.strip() for name in os.environ.get(extra_env_var, "").split(",") if name.strip() + } | {env_var} + os.environ[extra_env_var] = ",".join(sorted(merged_env_vars)) + + +_propagate_env_var_to_ray_workers(_MODELOPT_ATTN_SOFTMAX_MODE) + + def main(): """Launch vLLM with sparse attention worker.""" parser = FlexibleArgumentParser(description="vLLM model server with sparse attention") diff --git a/modelopt/torch/kernels/common/attention/decode_attention.py b/modelopt/torch/kernels/common/attention/decode_attention.py new file mode 100644 index 00000000000..a77eb0ca44a --- /dev/null +++ b/modelopt/torch/kernels/common/attention/decode_attention.py @@ -0,0 +1,380 @@ +# 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. + +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 + +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.bmm2_qdq import _p_qdq_nvfp4, _v_qdq_nvfp4 +from modelopt.torch.kernels.quantization.attention.softmax_fakequant import ( + ex2_fp16, + resolve_softmax_mode, +) + +__all__ = ["attention_decode"] + +_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, + MIXED_FP16: 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) + if MIXED_FP16: + p = ex2_fp16(scores - new_max) + correction = ex2_fp16(running_max - new_max) + else: + p = tl.math.exp2(scores - new_max) + correction = tl.math.exp2(running_max - new_max) + p = tl.where(kv_valid, p, 0.0) + running_sum = running_sum * correction + tl.sum(p, axis=0) + acc *= correction + + 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 + + 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, + softmax_mode: str = "fp32", + 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. 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. + P QDQ intentionally follows the split-local online-softmax schedule; changing + ``num_kv_splits`` can therefore change quantized results. + + ``softmax_mode="mixed_fp16"`` evaluates split-local probabilities and + running-max corrections with native FP16 exponentiation while retaining + FP32 denominator state and accumulators. Split-state combination remains + FP32. The default ``"fp32"`` preserves the existing behavior. + """ + mixed_fp16 = resolve_softmax_mode(softmax_mode) + if q.ndim != 3: + raise ValueError(f"q must have shape [batch, heads, head_dim], got {tuple(q.shape)}") + if page_size != k_cache.shape[1] or page_size != v_cache.shape[1]: + 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, + MIXED_FP16=mixed_fp16, + 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/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 751ce052a5c..a784e17bfc8 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -29,52 +29,56 @@ 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. +from modelopt.torch.kernels.quantization.attention.softmax_fakequant import ( + ex2_fp16, + resolve_softmax_mode, +) + +# Helpers for optional N:M sparsity and skip-softmax live in the sparsity +# package. The baseline forward kernel below calls them conditionally under +# constexpr guards, so the unified single-kernel design stays intact while +# 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 -_p_qdq_fp8: Any = None +_qdq_fp8: Any = None _p_qdq_nvfp4: Any = None +_v_qdq_nvfp4: Any = None 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 -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.bmm2_qdq import _p_qdq_nvfp4 as _p_nvfp4 + from modelopt.torch.kernels.quantization.attention.bmm2_qdq import _v_qdq_nvfp4 as _v_nvfp4 from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq as _fp8 - _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. +# Maps public QDQ options to kernel constexpr values. _P_QDQ_MODES = {None: 0, "fp8": 1, "nvfp4": 2} +_V_QDQ_MODES = {None: 0, "nvfp4": 2} LOG2E: float = 1.44269504088896 @@ -83,21 +87,15 @@ def _load_p_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 @@ -138,6 +136,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 @@ -181,6 +180,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 = ( @@ -221,10 +221,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 @@ -255,6 +286,8 @@ 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 + MIXED_FP16: tl.constexpr = False, # FP16 exp2 with FP32 reductions and accumulators SPARSITY_N: tl.constexpr = 0, # N:M sparsity — keep top-N of every M elements (0 = disabled) SPARSITY_M: tl.constexpr = 4, # N:M sparsity — group size (4 or 8) DENSE_SINK_TOKENS: tl.constexpr = 0, # Leading KV tokens kept dense (attention sinks) @@ -263,6 +296,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, 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 +359,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): @@ -357,23 +394,28 @@ 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 --- 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 @@ -394,9 +436,13 @@ def _attn_fwd( if not skip_tile: # --- Online softmax update --- m_new = tl.maximum(row_max, tl.max(scores, 1)) - p = tl.math.exp2(scores - m_new[:, None]) + if MIXED_FP16: + p = ex2_fp16(scores - m_new[:, None]) + correction = ex2_fp16(row_max - m_new) + else: + p = tl.math.exp2(scores - m_new[:, None]) + correction = tl.math.exp2(row_max - m_new) l_new = tl.sum(p, 1) - correction = tl.math.exp2(row_max - m_new) row_sum = row_sum * correction + l_new acc = acc * correction[:, None] @@ -404,8 +450,14 @@ 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: + # 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 @@ -435,7 +487,20 @@ def _attn_fwd( mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :], other=0.0, ) - acc = tl.dot(p.to(v.dtype), v, acc) + if V_QDQ == 2 and ( + (not V_CACHE_QUANTIZED) or (kv_start + BLOCK_N > v_quantized_boundary) + ): + 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 + 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 @@ -615,24 +680,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: @@ -769,24 +836,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: @@ -822,6 +891,7 @@ def forward( max_input_len, is_causal, sm_scale, + mixed_fp16, b_start_loc_k, b_seq_len_k, max_input_len_k, @@ -833,6 +903,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, @@ -918,12 +991,17 @@ 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, "HEAD_DIM": HEAD_DIM, "STORE_LSE": True, + "Q_IS_FP32": q.dtype == torch.float32 and (p_qdq_mode == 2 or v_qdq_mode == 2), + "MIXED_FP16": mixed_fp16, "SPARSITY_N": sparsity_n, "SPARSITY_M": sparsity_m, "DENSE_SINK_TOKENS": dense_sink_tokens, @@ -932,6 +1010,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, @@ -965,7 +1046,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, @@ -1126,6 +1207,7 @@ def backward(ctx, grad_output): None, # max_input_len None, # is_causal None, # sm_scale + None, # mixed_fp16 None, # b_start_loc_k None, # b_seq_len_k None, # max_input_len_k @@ -1137,6 +1219,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 @@ -1163,8 +1248,12 @@ def attention( dense_recent_tokens: int = 64, skip_softmax_threshold: float | None = None, measure_sparsity: bool = False, + softmax_mode: str = "fp32", p_qdq: str | None = None, p_qdq_amax: float = 1.0, + v_qdq: str | None = None, + 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, @@ -1202,6 +1291,12 @@ def attention( and skipped tiles via atomic counters. The counts are stored as ``_sparsity_total`` and ``_sparsity_skipped`` attributes on the returned output tensor. + softmax_mode: Exponentiation mode for the forward online-softmax update. + ``"fp32"`` evaluates probability and running-maximum correction + exponentials in FP32. ``"mixed_fp16"`` evaluates those exponentials + in native FP16 and converts their results back to FP32; row maxima, + denominator sums, and output accumulation remain FP32. Mixed mode is + inference-only because backward continues to use FP32 exponentials. p_qdq: Fake quant-dequant of the softmax probabilities ``P`` before the ``P @ V`` matmul (BMM2), emulating quantized attention. ``"fp8"`` round-trips P through FP8 E4M3 with a static per-tensor @@ -1211,8 +1306,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 @@ -1221,6 +1316,13 @@ 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``. ``"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 the NVFP4 global scale ``amax / (6 * 448)``. + v_cache_quantized: Complete block-16 groups in the paged V cache are + already QDQ; only the pristine partial group is QDQ on read. k_cache: Paged K cache [num_blocks, page_size, num_kv_heads, head_dim]. When provided, K/V are read from paged cache via block_table instead of from contiguous k/v tensors. @@ -1244,7 +1346,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() + _load_qdq_helpers() + mixed_fp16 = resolve_softmax_mode(softmax_mode) + if mixed_fp16 and any(t.requires_grad for t in (q, k, v)): + raise NotImplementedError( + "mixed_fp16 softmax is inference-only and does not support autograd" + ) if p_qdq not in _P_QDQ_MODES: raise ValueError( f"p_qdq must be one of {sorted(k for k in _P_QDQ_MODES if k)} or None, got {p_qdq!r}" @@ -1259,6 +1366,22 @@ 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 _V_QDQ_MODES: + raise ValueError( + f"v_qdq must be one of {sorted(k for k in _V_QDQ_MODES if k)} or None, got {v_qdq!r}" + ) + v_qdq_mode = _V_QDQ_MODES[v_qdq] + if v_qdq_mode and any(t.requires_grad for t in (q, k, v)): + raise NotImplementedError("v_qdq is inference-only and does not support autograd") + v_qdq_scale = 1.0 + if v_qdq_mode and v_qdq_amax is not None: + if not (math.isfinite(v_qdq_amax) and v_qdq_amax > 0): + raise ValueError(f"v_qdq_amax must be a finite positive value, got {v_qdq_amax}") + v_qdq_scale = v_qdq_amax / (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, @@ -1269,6 +1392,7 @@ def attention( max_input_len, is_causal, sm_scale, + mixed_fp16, b_start_loc_k, b_seq_len_k, max_input_len_k, @@ -1280,6 +1404,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/__init__.py b/modelopt/torch/kernels/quantization/attention/__init__.py index 4232082dc60..1d859f07c33 100644 --- a/modelopt/torch/kernels/quantization/attention/__init__.py +++ b/modelopt/torch/kernels/quantization/attention/__init__.py @@ -15,10 +15,7 @@ """Quantization-specific attention kernel pieces. -``p_qdq.py`` holds the softmax-P (``p_bmm_quantizer``) quant-dequant -``@triton.jit`` helpers invoked by the unified flash-attention kernel in -``common/attention/triton_fa.py`` under its ``P_QDQ`` constexpr guard. -Only NVFP4 needs a P-specific helper (tiling and block-amax policy on top of -``quantization/common/nvfp4_quant.py``); the FP8 mode uses -``quantization/common/fp8_quant.fp8_scalar_qdq`` directly. +``bmm2_qdq.py`` holds the operand-specific NVFP4 helpers for the attention +``P @ V`` matmul and the V-cache finalization kernel. This package initializer +does not import that module, so importing the package alone does not require Triton. """ diff --git a/modelopt/torch/kernels/quantization/attention/bmm2_qdq.py b/modelopt/torch/kernels/quantization/attention/bmm2_qdq.py new file mode 100644 index 00000000000..b33900546eb --- /dev/null +++ b/modelopt/torch/kernels/quantization/attention/bmm2_qdq.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""NVFP4 operand helpers for the attention ``P @ V`` matmul (BMM2). + +P and V share the low-level ``nvfp4_scalar_qdq`` primitive, but retain thin +operand-specific wrappers because their layouts and amax reductions differ. +P is nonnegative with layout ``[M, K]``; V is signed with layout ``[K, N]``. +Both use block-16 scaling along the BMM2 contraction axis. +""" + +import math + +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 _p_qdq_nvfp4( + p, + global_scale, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """NVFP4 fake quant-dequant of softmax probabilities. + + Two-level scaling per the NVFP4 recipe: E2M1 elements with one FP8 E4M3 + scale per 16 contiguous elements along the key dimension (the contraction + axis of ``P @ V``), and a per-tensor ``global_scale`` (runtime scalar, + ``amax / (6 * 448)``; ``attention()`` derives it from ``p_qdq_amax``, + which defaults to 1, the theoretical upper bound of P's amax). + + ``p >= 0``, so the block amax is a plain max (no ``abs``), and + ``nvfp4_scalar_qdq`` guards the degenerate all-zero blocks of fully + masked or padded positions. + """ + tl.static_assert(BLOCK_N % 16 == 0, "BLOCK_N must be divisible by 16 for NVFP4") + + grouped = tl.reshape(p, (BLOCK_M, BLOCK_N // 16, 16)) + block_amax = tl.expand_dims(tl.max(grouped, axis=2), 2) # p >= 0, so max == amax + q = nvfp4_scalar_qdq(grouped, block_amax, global_scale, 16) + return tl.reshape(q, (BLOCK_M, BLOCK_N)) + + +@triton.jit +def _v_qdq_nvfp4(v, global_scale, BLOCK_N: tl.constexpr, BLOCK_D: tl.constexpr): + """Fake-quantize signed V in block-16 groups along its key axis.""" + 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, + *, + max_new_tokens: int, + page_size: int = 16, + v_qdq_scale: float = 1.0, +) -> None: + """NVFP4-finalize complete block-16 groups in ``[v_lo, v_hi)`` in place. + + ``max_new_tokens`` is host metadata used to size the masked launch grid. + The grid covers every group that the largest query chunk can complete + without reading device metadata. ``v_lo`` and ``v_hi`` must describe + aligned, completed block-16 boundaries; their device values are not + host-validated. + """ + if page_size != v_cache.shape[1]: + raise ValueError(f"page_size {page_size} must match v_cache.shape[1] {v_cache.shape[1]}") + 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:] + num_groups = triton.cdiv(max_new_tokens, _BLOCK_N) + + with torch.cuda.device(v_cache.device): + _fake_quant_v_onwrite_kernel[(batch, num_kv_heads, num_groups)]( + 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/attention/p_qdq.py b/modelopt/torch/kernels/quantization/attention/p_qdq.py deleted file mode 100644 index 2603f6364af..00000000000 --- a/modelopt/torch/kernels/quantization/attention/p_qdq.py +++ /dev/null @@ -1,65 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Softmax-P quant-dequant helpers for the unified flash attention kernel. - -These ``@triton.jit`` helpers fake-quantize the softmax probabilities ``P`` -before the ``P @ V`` matmul (BMM2) — the in-kernel counterpart of the -``p_bmm_quantizer`` config. They are called conditionally from the baseline -flash-attention kernel in ``common/attention/triton_fa.py`` under the -``P_QDQ`` constexpr guard, following the same composition pattern as -the sparsity helpers in ``sparsity/attention/skip_softmax_helpers.py``. - -Only NVFP4 needs a P-specific helper (tiling policy and block amaxes); the -per-tensor FP8 mode uses ``quantization/common/fp8_quant.fp8_scalar_qdq`` -directly. What is P-specific here: the kernel's online-softmax ``p`` is -unnormalized and bounded (``0 <= p <= 1``, since the max-subtraction caps -every entry at ``exp2(0) = 1``), so 1 is the theoretical upper bound of its -amax; block amaxes need no ``abs``; and the NVFP4 scale blocks of 16 run -along the key dimension — the contraction axis of ``P @ V``. The caller -(``attention()``) converts the amax to the ``global_scale`` below. -""" - -import triton -import triton.language as tl - -from modelopt.torch.kernels.quantization.common.nvfp4_quant import nvfp4_scalar_qdq - - -@triton.jit -def _p_qdq_nvfp4( - p, - global_scale, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, -): - """NVFP4 fake quant-dequant of softmax probabilities. - - Two-level scaling per the NVFP4 recipe: E2M1 elements with one FP8 E4M3 - scale per 16 contiguous elements along the key dimension (the contraction - axis of ``P @ V``), and a per-tensor ``global_scale`` (runtime scalar, - ``amax / (6 * 448)``; ``attention()`` derives it from ``p_qdq_amax``, - which defaults to 1, the theoretical upper bound of P's amax). - - ``p >= 0``, so the block amax is a plain max (no ``abs``), and - ``nvfp4_scalar_qdq`` guards the degenerate all-zero blocks of fully - masked or padded positions. - """ - tl.static_assert(BLOCK_N % 16 == 0, "BLOCK_N must be divisible by 16 for NVFP4") - - grouped = tl.reshape(p, (BLOCK_M, BLOCK_N // 16, 16)) - block_amax = tl.expand_dims(tl.max(grouped, axis=2), 2) # p >= 0, so max == amax - q = nvfp4_scalar_qdq(grouped, block_amax, global_scale, 16) - return tl.reshape(q, (BLOCK_M, BLOCK_N)) diff --git a/modelopt/torch/kernels/quantization/attention/softmax_fakequant.py b/modelopt/torch/kernels/quantization/attention/softmax_fakequant.py new file mode 100644 index 00000000000..c78649afd7a --- /dev/null +++ b/modelopt/torch/kernels/quantization/attention/softmax_fakequant.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Mixed-FP16 softmax helpers for Triton attention kernels.""" + +import triton +import triton.language as tl + +__all__ = ["SOFTMAX_MODES", "ex2_fp16", "resolve_softmax_mode"] + +SOFTMAX_MODES = frozenset(("fp32", "mixed_fp16")) + + +def resolve_softmax_mode(mode: str) -> bool: + """Validate ``mode`` and report whether mixed-FP16 softmax is enabled.""" + if mode not in SOFTMAX_MODES: + raise ValueError(f"softmax_mode must be one of {sorted(SOFTMAX_MODES)}, got {mode!r}") + return mode == "mixed_fp16" + + +@triton.jit +def ex2_fp16(x): + """Evaluate base-2 exponentiation in native FP16 on sm_75 or newer and return FP32.""" + return tl.inline_asm_elementwise( + "{ .reg .b16 h; cvt.rn.f16.f32 h, $1; ex2.approx.f16 h, h; cvt.f32.f16 $0, h; }", + "=r,r", + args=[x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) diff --git a/modelopt/torch/kernels/quantization/common/nvfp4_quant.py b/modelopt/torch/kernels/quantization/common/nvfp4_quant.py index 4aea5a88688..d74a346df90 100644 --- a/modelopt/torch/kernels/quantization/common/nvfp4_quant.py +++ b/modelopt/torch/kernels/quantization/common/nvfp4_quant.py @@ -19,7 +19,7 @@ - ``../gemm/fp4_kernel.py`` (standalone blockwise fake quant) - ``../gemm/fp4_kernel_hopper.py`` (Hopper block-pointer variant) - ``../gemm/gptq_fused_kernel.py`` (fused GPTQ scalar path) - - ``../attention/p_qdq.py`` (softmax-P qdq in the flash-attention kernel) + - ``../attention/bmm2_qdq.py`` (P/V operand qdq in the attention BMM2 kernel) FP4 (E2M1) representable magnitudes: {0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0} """ @@ -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. @@ -119,8 +121,8 @@ def fp8_quantize_scale(block_amax, global_scale): """ FP8_E4M3_MAX: tl.constexpr = 448.0 scale_in_fp8_range = block_amax / (6.0 * global_scale) - scale_clamped = tl.minimum(scale_in_fp8_range, 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/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/modelopt/torch/quantization/plugins/vllm.py b/modelopt/torch/quantization/plugins/vllm.py index 95ca3240b73..b37e10125d8 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): @@ -162,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 = ( @@ -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) @@ -504,9 +523,11 @@ 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) - 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/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 734755e3bba..0b283cf032b 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -15,20 +15,23 @@ """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 +from dataclasses import dataclass import torch from vllm.v1.attention.backends.flash_attn import ( @@ -37,12 +40,16 @@ 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.bmm2_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,14 +120,246 @@ def _build_sparse_kw(layer_cfg: dict) -> dict: return sparse_kw -class ModelOptSparseAttentionImpl(FlashAttentionImpl): - """Attention implementation that uses the ModelOpt Triton kernel. +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) + + +def _quant_kw_from_impl(impl, layer): + """Resolve the compact attention quantization contract for one launch.""" + quant_kw = getattr(impl, "quant_kw", None) + if quant_kw is None: + p_qdq, p_qdq_amax = _p_qdq_from_layer(layer) + v_qdq, v_qdq_amax = _v_qdq_from_layer(layer) + softmax_mode = "fp32" + else: + p_qdq, p_qdq_amax = quant_kw["p_qdq"], quant_kw["p_qdq_amax"] + v_qdq, v_qdq_amax = quant_kw["v_qdq"], quant_kw["v_qdq_amax"] + softmax_mode = quant_kw.get("softmax_mode", "fp32") + return p_qdq, p_qdq_amax, v_qdq, v_qdq_amax, softmax_mode + + +def _any_quant_active(layer, p_qdq, v_qdq, softmax_mode) -> bool: + """Return whether native fallback would omit any quantized operation.""" + k_quantizer = getattr(layer, "k_bmm_quantizer", None) + return bool( + p_qdq + or v_qdq + or softmax_mode != "fp32" + or getattr(layer, "_query_quant_in_kernel", False) + or getattr(k_quantizer, "is_enabled", False) + ) + + +def _should_run_modelopt_kernel(sparse_kw, quant_active: bool) -> bool: + """Return whether a launch has effective work for the ModelOpt kernel.""" + return bool(sparse_kw or quant_active) + + +@dataclass(frozen=True, slots=True) +class _ResolvedForward: + p_qdq: str | None + p_qdq_amax: float + v_qdq: str | None + v_qdq_amax: float | None + softmax_mode: str + quant_active: 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, softmax_mode = _quant_kw_from_impl(impl, layer) + quant_active = _any_quant_active(layer, p_qdq, v_qdq, softmax_mode) + transform_active = _should_run_modelopt_kernel(getattr(impl, "sparse_kw", None), quant_active) + + if getattr(attn_metadata, "use_cascade", False): + 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 - Inherits from FlashAttentionImpl to reuse: - - __init__ (all configuration) - - do_kv_cache_update (KV cache writing) - Only overrides forward() to replace sparse prefill attention computation. - """ + 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, + softmax_mode=softmax_mode, + quant_active=quant_active, + ) + + +# Resolution guards raw configured transforms; dispatch rechecks effective +# sparse work after calibration and decode-only pruning. +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, + softmax_mode: str, + quant_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 _should_run_modelopt_kernel(sparse_kw, quant_active): + # Dynamic calibration can disable sparse work for a launch. Preserve the + # backend's native dense path when no ModelOpt transform remains active. + 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, + ) + + 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, + softmax_mode=softmax_mode, + ) + 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, + softmax_mode=softmax_mode, + **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, @@ -163,13 +402,9 @@ def forward( assert output is not None, "Output tensor must be provided." if attn_metadata is None: - # Profiling run return output.fill_(0) - 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. + def native_forward(): return self._forward_vllm_flash_attn( layer, query, @@ -182,100 +417,40 @@ def forward( output_block_scale, ) - 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) + resolved = _resolve_forward( + self, + layer, + attn_metadata, + output_scale, + output_block_scale, + ) + if resolved is None: + return native_forward() - # 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, + is_causal=getattr(attn_metadata, "causal", not is_decode_only), + output=output, + p_qdq=resolved.p_qdq, + p_qdq_amax=resolved.p_qdq_amax, + v_qdq=resolved.v_qdq, + v_qdq_amax=resolved.v_qdq_amax, + softmax_mode=resolved.softmax_mode, + quant_active=resolved.quant_active, + dense_fallback=native_forward, ) - 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 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: - # 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( - layer, - query, - key, - value, - kv_cache, - attn_metadata, - output, - output_scale, - output_block_scale, - ) - - # 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 - # 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 - **sparse_kw, - ) - - output[:num_actual_tokens] = triton_out - return output class ModelOptSparseAttentionBackend(FlashAttentionBackend): @@ -295,13 +470,286 @@ 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 _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 + 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 + # 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) + 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 _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, + "softmax_mode": resolved.softmax_mode, + "quant_active": resolved.quant_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( + 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 not resolved.quant_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( + 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 + 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, + ): + return _flashinfer_forward( + self, + super().forward, + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output, + output_scale, + output_block_scale, + ) + + _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) @@ -310,6 +758,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_decode_attention.py b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py new file mode 100644 index 00000000000..d59a618c243 --- /dev/null +++ b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPU tests for the minimal paged split-K decode kernel.""" + +import math + +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.bmm2_qdq import fake_quant_v_onwrite + +NATIVE_E4M3_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (8, 9) +requires_native_e4m3 = pytest.mark.skipif( + 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 + + +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) + + +def _softmax_exp2(x, softmax_mode): + if softmax_mode == "mixed_fp16": + return torch.exp2(x.half()).float() + assert softmax_mode == "fp32" + return torch.exp2(x) + + +def _combine_split_states(split_max, split_sum, split_acc, softmax_mode): + running_max = torch.tensor(-float("inf"), device=split_max.device) + running_sum = torch.tensor(0.0, device=split_max.device) + acc = torch.zeros_like(split_acc[0]) + for split_idx in range(split_max.shape[0]): + new_max = torch.maximum(running_max, split_max[split_idx]) + correction = _softmax_exp2(running_max - new_max, softmax_mode) + split_correction = _softmax_exp2(split_max[split_idx] - new_max, softmax_mode) + acc = acc * correction + split_acc[split_idx] * split_correction + running_sum = running_sum * correction + split_sum[split_idx] * split_correction + running_max = new_max + return acc / running_sum + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") +@pytest.mark.parametrize("num_kv_splits", [1, 32]) +def test_split_k_varlen_gqa_matches_dense(num_kv_splits): + 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 +@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=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", + "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) + 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=v_qdq_scale, + ) + 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) + + +@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) + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + Triton") +@requires_native_e4m3 +@pytest.mark.parametrize("softmax_mode", ["fp32", "mixed_fp16"]) +def test_p_qdq_matches_fixed_split_local_oracle(softmax_mode): + """The production 32-split schedule quantizes split-local unnormalized P.""" + seq_len, head_dim, num_splits = 8192, 16, 32 + block_n = 128 + scale = head_dim**-0.5 + q = torch.zeros(1, 1, head_dim, device="cuda", dtype=torch.float32) + q[..., 0] = 1.0 / (scale * 1.4426950408889634) + k = torch.zeros(1, 1, seq_len, head_dim, device="cuda", dtype=torch.bfloat16) + token_idx = torch.arange(seq_len, device="cuda") + token_split = token_idx // (2 * block_n) + tile_phase = (token_idx // block_n) % 2 + score_pattern = -0.25 * ((token_idx % block_n) % 3) + 0.25 * tile_phase + 0.25 * token_split + k[0, 0, :, 0] = score_pattern.to(torch.bfloat16) + v = torch.zeros_like(k) + # Channel 0 isolates the split-local correction; channel 1 nearly cancels only + # under the intended FP32 combine, making combine precision observable. + v[0, 0, :, 0] = torch.where(tile_phase == 0, 256.0, -215.0).to(torch.bfloat16) + v[0, 0, :, 1] = torch.where(token_split == num_splits - 1, -5.25, 1.0).to(torch.bfloat16) + k_cache, v_cache, block_table = _paged_cache(k, v, (seq_len,)) + seq_lens = torch.tensor([seq_len], device="cuda", dtype=torch.int32) + + common = { + "softmax_scale": scale, + "num_kv_splits": num_splits, + "p_qdq": "nvfp4", + } + output = attention_decode( + q, k_cache, v_cache, block_table, seq_lens, softmax_mode=softmax_mode, **common + ) + if softmax_mode == "mixed_fp16": + fp32_output = attention_decode( + q, k_cache, v_cache, block_table, seq_lens, softmax_mode="fp32", **common + ) + assert not torch.equal(output.to(k.dtype), fp32_output.to(k.dtype)) + + scores = torch.matmul(k[0, 0].float(), q[0, 0]) * (scale * 1.4426950408889634) + split_scores = scores.reshape(num_splits, 2, block_n) + split_values = v[0, 0].float().reshape(num_splits, 2, block_n, head_dim) + split_max = [] + split_sum = [] + split_acc = [] + for split_idx in range(num_splits): + running_max = torch.tensor(-float("inf"), device="cuda") + running_sum = torch.tensor(0.0, device="cuda") + local_acc = torch.zeros(head_dim, device="cuda") + for tile_idx in range(2): + tile_scores = split_scores[split_idx, tile_idx] + new_max = torch.maximum(running_max, tile_scores.amax()) + p = _softmax_exp2(tile_scores - new_max, softmax_mode) + correction = _softmax_exp2(running_max - new_max, softmax_mode) + running_sum = running_sum * correction + p.sum() + local_acc *= correction + p_qdq = _nvfp4_qdq_reference(p.to(torch.bfloat16).float()) + local_acc += (p_qdq[:, None] * split_values[split_idx, tile_idx]).sum(dim=0) + running_max = new_max + split_max.append(running_max) + split_sum.append(running_sum) + split_acc.append(local_acc) + split_max = torch.stack(split_max) + split_sum = torch.stack(split_sum) + split_acc = torch.stack(split_acc) + reference = _combine_split_states(split_max, split_sum, split_acc, "fp32") + if softmax_mode == "mixed_fp16": + mixed_combine_reference = _combine_split_states( + split_max, split_sum, split_acc, "mixed_fp16" + ) + assert reference.to(k.dtype)[1] != mixed_combine_reference.to(k.dtype)[1] + + torch.testing.assert_close(output[0, 0], reference, rtol=5e-5, atol=5e-5) 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..04d4227f63a 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -13,7 +13,10 @@ # 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 import pytest import torch @@ -24,15 +27,202 @@ 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.triton_fa import LOG2E + import triton + import triton.language as tl -# 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 + from modelopt.torch.kernels.common.attention import attention, triton_fa + from modelopt.torch.kernels.common.attention.triton_fa import LOG2E + from modelopt.torch.kernels.quantization.attention.softmax_fakequant import ex2_fp16 + + @triton.jit + def _ex2_fp16_test_kernel(x, out, n, BLOCK: tl.constexpr): # noqa: N803 + offsets = tl.arange(0, BLOCK) + mask = offsets < n + tl.store(out + offsets, ex2_fp16(tl.load(x + offsets, mask=mask)), mask=mask) + + +NATIVE_FP16_EX2_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (7, 5) +requires_native_fp16_ex2 = pytest.mark.skipif( + not NATIVE_FP16_EX2_AVAILABLE, + reason="Native FP16 ex2 requires compute capability >= 7.5", +) +NATIVE_E4M3_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (8, 9) +requires_native_e4m3 = pytest.mark.skipif( + not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute capability >= 8.9" +) + + +@requires_native_fp16_ex2 +def test_ex2_fp16_matches_fp16_reference(): + x = torch.tensor( + [float("-inf"), -12.0, -3.25, -1.0, -0.125, -0.1, 0.0], + dtype=torch.float32, + device="cuda", + ) + actual = torch.empty_like(x) + + _ex2_fp16_test_kernel[(1,)](x, actual, x.numel(), BLOCK=8) + + expected = torch.exp2(x.to(torch.float16)).to(torch.float32) + # PTX specifies a maximum relative error of 2**-9.9 for ex2.approx.f16. + torch.testing.assert_close(actual, expected, rtol=2**-9, atol=0.0) + actual_finite = actual[torch.isfinite(actual)] + torch.testing.assert_close( + actual_finite, + actual_finite.to(torch.float16).float(), + rtol=0.0, + atol=0.0, + ) + assert actual[0].item() == 0.0 + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need triton") +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"v_qdq": "nvfp4"}, r"v_qdq.*autograd"), + ({"softmax_mode": "mixed_fp16"}, r"mixed_fp16.*autograd"), + ], +) +def test_inference_only_modes_reject_autograd_before_kernel_launch(monkeypatch, kwargs, match): + apply = Mock(side_effect=AssertionError("_Attention.apply reached")) + monkeypatch.setattr(triton_fa._Attention, "apply", apply) + q, k, v = (torch.zeros(1, 1, 16, requires_grad=True) for _ in range(3)) + locs = torch.zeros(1, dtype=torch.int32) + lens = torch.ones(1, dtype=torch.int32) + with pytest.raises(NotImplementedError, match=match): + attention(q, k, v, locs, lens, 1, **kwargs) + apply.assert_not_called() + + +# 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 +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") +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) + 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=True, + 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. @@ -67,12 +257,30 @@ 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) - 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) +def _qdq_v_nvfp4(v, global_scale=1.0): + """Signed NVFP4 V qdq in block-16 groups along the key axis.""" + carrier_dtype = v.dtype + shape = v.shape + assert shape[0] % 16 == 0 + g = v.float().reshape(shape[0] // 16, 16, *shape[1:]) + block_amax = g.abs().amax(dim=1, keepdim=True) + scale_in_fp8_range = (block_amax / (6.0 * global_scale)).clamp(max=FP8_E4M3_MAX) + scale = scale_in_fp8_range.to(torch.float8_e4m3fn).float() * global_scale + scale_safe = torch.where(scale == 0, 1.0, scale) + q = _fp4_round(g.abs() / scale_safe) * scale_safe + q = torch.where(g < 0, -q, q) + q = torch.where(scale == 0, 0.0, q) + return q.reshape(shape).to(carrier_dtype).float() + + def _apply_qdq(p, mode, qdq_scale=1.0): if mode == "fp8": return _qdq_fp8(p, qdq_scale) @@ -80,13 +288,53 @@ 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 _softmax_exp2(x, softmax_mode): + if softmax_mode == "mixed_fp16": + return torch.exp2(x.to(torch.float16)).to(torch.float32) + assert softmax_mode == "fp32" + return torch.exp2(x) + + +def _apply_top_n_of_m_tile(scores, n, m): + """Apply the kernel's per-tile top-N-of-M masking along the key axis.""" + assert m in (4, 8) and scores.shape[-1] % m == 0 + grouped = scores.reshape(*scores.shape[:-1], scores.shape[-1] // m, m) + if m == 4: + # The Triton comparison network breaks ties in favor of lower indices. + order = torch.argsort(grouped, dim=-1, descending=True, stable=True) + kept = torch.zeros_like(grouped, dtype=torch.bool) + kept.scatter_(-1, order[..., :n], True) + else: + # The Triton M=8 path keeps every value tied at the Nth-largest threshold. + threshold = grouped.sort(dim=-1).values[..., m - n] + kept = grouped >= threshold[..., None] + return grouped.masked_fill(~kept, float("-inf")).reshape_as(scores) + + +def qdq_attention_reference( + q, + k, + v, + scale, + mode, + is_causal=True, + amax=1.0, + block_n=P_QDQ_BLOCK_N, + softmax_mode="fp32", + v_qdq=None, + sparsity_n=0, + sparsity_m=4, + dense_sink_tokens=0, + dense_recent_tokens=64, +): """Tile-looped online-softmax reference replicating kernel P_QDQ semantics. Single sequence: q [s, h, d], k/v [s_kv, h_kv, d] (fp16). Walks KV tiles - 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 ``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 @@ -113,18 +361,37 @@ 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] + if sparsity_n: + sparse_tile = _apply_top_n_of_m_tile(tile, sparsity_n, sparsity_m) + q_abs = torch.arange(s, device=q.device) + s_kv - s + kv_abs = torch.arange(start, start + tile.shape[-1], device=q.device) + distance = q_abs[:, None] - kv_abs[None, :] + dense = (kv_abs[None, :] < dense_sink_tokens) | ( + (distance >= 0) & (distance < dense_recent_tokens) + ) + if s <= 1: + dense = torch.ones_like(dense) + tile = torch.where(dense[None], tile, sparse_tile) m_new = torch.maximum(row_max, tile.amax(dim=-1)) - p = torch.exp2(tile - m_new[..., None]) + p = _softmax_exp2(tile - m_new[..., None], softmax_mode) l_new = p.sum(dim=-1) - corr = torch.exp2(row_max - m_new) + corr = _softmax_exp2(row_max - m_new, softmax_mode) row_sum = row_sum * corr + l_new acc = acc * corr[..., None] + if mode == "nvfp4": + 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() - acc = acc + torch.einsum("hqk,khd->hqd", p, vv[start : start + BLOCK_N].float()) + if mode == "fp8": + p = p.to(v.dtype).float() + v_tile = vv[start : start + block_n] + if v_qdq == "nvfp4": + # The kernel QDQs in FP32, then rounds back to the cache/model carrier dtype. + v_tile = _qdq_v_nvfp4(v_tile) + else: + assert v_qdq is None + acc = acc + torch.einsum("hqk,khd->hqd", p, v_tile.float()) row_max = m_new out = acc / row_sum[..., None] return out.permute(1, 0, 2) @@ -135,9 +402,10 @@ 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 + seq_len, num_heads, num_kv_heads, head_dim = 256, 4, 2, 64 scale = 1.0 / (head_dim**0.5) torch.manual_seed(7) @@ -146,15 +414,120 @@ 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) assert not torch.equal(o, o_dense) + @requires_native_fp16_ex2 + @requires_native_e4m3 + def test_mixed_fp16_prefill_matches_tile_reference(self): + """Mixed softmax composes with P/V NVFP4 and tile-local 2:4 sparsity.""" + seq_len, num_heads, num_kv_heads, head_dim = 128, 2, 1, 64 + scale = 1.0 / math.sqrt(head_dim) + + torch.manual_seed(29) + q, k, v = make_qkv(seq_len, num_heads, num_kv_heads, head_dim, dtype=torch.float16) + locs, lens = make_varlen_meta([seq_len]) + + actual = attention( + q, + k, + v, + locs, + lens, + seq_len, + softmax_scale=scale, + softmax_mode="mixed_fp16", + p_qdq="nvfp4", + v_qdq="nvfp4", + sparsity_n=2, + sparsity_m=4, + dense_recent_tokens=0, + ) + reference = qdq_attention_reference( + q, + k, + v, + scale, + "nvfp4", + softmax_mode="mixed_fp16", + v_qdq="nvfp4", + sparsity_n=2, + sparsity_m=4, + dense_recent_tokens=0, + ) + fp32 = attention( + q, + k, + v, + locs, + lens, + seq_len, + softmax_scale=scale, + softmax_mode="fp32", + p_qdq="nvfp4", + v_qdq="nvfp4", + sparsity_n=2, + sparsity_m=4, + dense_recent_tokens=0, + ) + + torch.testing.assert_close(actual.float(), reference, rtol=5e-3, atol=2e-2) + assert not torch.equal(actual, fp32) + + @requires_native_e4m3 + def test_nvfp4_uses_native_p_input_and_fp32_carriers(self): + """Dynamic-Q and P QDQ stay FP32 until their native MMA boundaries.""" + num_heads = num_kv_heads = 1 + head_dim = 16 + 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) + 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): - """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 @@ -168,23 +541,13 @@ 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) - - @pytest.mark.parametrize(("mode", "tol"), [("fp8", 5e-2), ("nvfp4", 0.25)]) - 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) + # 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", ["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 +595,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 +629,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 +659,24 @@ 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": "fp8"}, "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..a1513b497c3 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py @@ -15,14 +15,31 @@ """GPU tests for paged KV cache mode of the Triton flash attention kernel.""" +import inspect +import math + import pytest import torch 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.common.attention import attention, triton_fa + from modelopt.torch.kernels.quantization.attention.bmm2_qdq import fake_quant_v_onwrite + +NATIVE_E4M3_AVAILABLE = TRITON_KERNEL_AVAILABLE and torch.cuda.get_device_capability() >= (8, 9) +requires_native_e4m3 = pytest.mark.skipif( + not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute capability >= 8.9" +) + + +@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 # --------------------------------------------------------------------------- @@ -113,21 +130,52 @@ 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.""" - def test_paged_matches_contiguous(self): - """Paged mode produces same output as contiguous mode with identical data.""" + @pytest.mark.parametrize( + ("page_size", "v_qdq_scale", "match"), + [(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") + 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, + max_new_tokens=1, + page_size=page_size, + v_qdq_scale=v_qdq_scale, + ) + + @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 @@ -151,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] @@ -266,147 +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 - seq_lens_k = [64, 128] - num_heads, num_kv_heads, head_dim = 4, 2, 64 - page_size = 16 - scale = 1.0 / (head_dim**0.5) - total_kv = sum(seq_lens_k) - - torch.manual_seed(34) - q_flat = torch.randn(batch, num_heads, head_dim, device="cuda", dtype=torch.float16) - k_flat = torch.randn(total_kv, num_kv_heads, head_dim, device="cuda", dtype=torch.float16) - v_flat = torch.randn(total_kv, num_kv_heads, head_dim, device="cuda", dtype=torch.float16) - - b_start_loc_q = torch.arange(batch, device="cuda", dtype=torch.int32) - b_seq_len_q = torch.ones(batch, device="cuda", dtype=torch.int32) - cumsum = [0] - for sl in seq_lens_k: - cumsum.append(cumsum[-1] + sl) - b_start_loc_k = torch.tensor(cumsum[:-1], device="cuda", dtype=torch.int32) - b_seq_len_k = torch.tensor(seq_lens_k, device="cuda", dtype=torch.int32) - k_cache, v_cache, block_table = _scatter_to_paged_cache( - k_flat, v_flat, b_start_loc_k, b_seq_len_k, num_kv_heads, head_dim, page_size - ) - - common_kwargs = { - "is_causal": False, - "softmax_scale": scale, - "b_start_loc_k": b_start_loc_k, - "b_seq_len_k": b_seq_len_k, - "max_input_len_k": max(seq_lens_k), - "k_cache": k_cache, - "v_cache": v_cache, - "block_table": block_table, - "page_size": page_size, - } - out_dense = attention( - q_flat, k_flat, v_flat, b_start_loc_q, b_seq_len_q, 1, **common_kwargs - ) - out_sparse = attention( - q_flat, - k_flat, - v_flat, - b_start_loc_q, - b_seq_len_q, - 1, - **common_kwargs, - sparsity_n=2, - sparsity_m=4, - dense_recent_tokens=64, - ) - - torch.testing.assert_close(out_sparse, out_dense, rtol=1e-2, atol=1e-2) - def test_paged_mixed_prefill_decode_sparse_nm_keeps_decode_dense(self): """Decode rows stay dense even when batched with sparse prefill rows.""" q_lens = [64, 1] @@ -495,3 +367,163 @@ 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), + max_new_tokens=32, + 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), + max_new_tokens=1, + page_size=page_size, + ) + 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) + + @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), + max_new_tokens=16, + 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) + + @requires_native_e4m3 + 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) + 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, + "p_qdq": "nvfp4", + "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), + max_new_tokens=16, + ) + 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) + + @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/quantization/test_vllm_dynamic_modules.py b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py index 2136f959754..a8011805daf 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,125 @@ 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_attention_setup_keeps_qkv_only_checkpoint_surface(monkeypatch): + monkeypatch.setattr( + vllm_plugin, + "create_parallel_state", + lambda: vllm_plugin.ParallelState(data_parallel_group=None), + ) + attention = _new_attention(_TestQuantVLLMAttention) + + 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, "_query_quant_in_kernel") + 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_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) + 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, "_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 + qv_in_kernel = attention(query, key, value) + + assert quantized[:3] == (torch.tensor(11), torch.tensor(22), torch.tensor(33)) + 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 + + +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``. 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..9c81d637bfe --- /dev/null +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py @@ -0,0 +1,311 @@ +# 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 +import os +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import vllm +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.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, + get_flashinfer_sparse_impl_cls, +) + +_WORKER_PATH = Path(__file__).parents[5] / "examples/vllm_serve/sparse_attn_worker.py" +_LAUNCHER_PATH = Path(__file__).parents[5] / "examples/vllm_serve/vllm_serve_sparse_attn.py" + + +def _load_worker_module(): + 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 + + +worker_module = _load_worker_module() + + +def _load_launcher_module(name="shared_attention_launcher_quant_test"): + spec = importlib.util.spec_from_file_location(name, _LAUNCHER_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_quant_policy_rejects_old_vllm(monkeypatch): + monkeypatch.setattr(vllm, "__version__", "0.9.0") + worker_module._load_quant_api.cache_clear() + + 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): + 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__(impl_cls) + 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 _worker, _api: []) + + +@pytest.mark.parametrize("impl_cls", [FlashAttentionImpl, FlashInferImpl]) +@pytest.mark.parametrize( + ("softmax_env", "expected_softmax_mode"), + [(None, "fp32"), (" MiXeD_FP16 ", "mixed_fp16")], +) +def test_install_converts_only_attention_and_configures_fixed_recipe( + monkeypatch, impl_cls, softmax_env, expected_softmax_mode +): + _patch_conversion(monkeypatch) + if softmax_env is None: + monkeypatch.delenv("MODELOPT_ATTN_SOFTMAX_MODE", raising=False) + else: + monkeypatch.setenv("MODELOPT_ATTN_SOFTMAX_MODE", softmax_env) + attention = _attention(impl_cls=impl_cls) + linear = nn.Linear(4, 4) + model = nn.ModuleDict({"attn": attention, "linear": linear}) + state = _worker(model) + + worker_module._install_attention(state, quantize=True) + + 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 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 + assert converted._value_quant_in_kernel is True + 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, + "v_qdq": "nvfp4", + "v_qdq_amax": 6.0 * 448.0, + "softmax_mode": expected_softmax_mode, + } + assert model["linear"] is linear + assert state.model_runner.cascade_attn_enabled is False + + +def test_invalid_softmax_mode_is_rejected_before_quant_mutation(monkeypatch): + _patch_conversion(monkeypatch) + monkeypatch.setenv("MODELOPT_ATTN_SOFTMAX_MODE", " fp64 ") + attention = _attention() + model = nn.ModuleDict({"attn": attention}) + state = _worker(model) + original_impl = attention.impl + + with pytest.raises(ValueError) as exc: + worker_module._install_attention(state, quantize=True) + + assert "MODELOPT_ATTN_SOFTMAX_MODE" in str(exc.value) + assert " fp64 " in str(exc.value) + assert model["attn"] is attention + assert type(model["attn"]) is quant_plugin.vllm_attention.Attention + assert attention.impl is original_impl + assert state.model_runner.cascade_attn_enabled is True + + +@pytest.mark.parametrize( + "available_module", + [ + pytest.param("vllm.executor.ray_distributed_executor", id="legacy"), + pytest.param("vllm.v1.executor.ray_distributed_executor", id="v1"), + pytest.param(None, id="env-fallback"), + ], +) +def test_launcher_propagates_softmax_mode_to_ray_workers(monkeypatch, available_module): + ray_modules = ( + "vllm.executor.ray_distributed_executor", + "vllm.v1.executor.ray_distributed_executor", + ) + softmax_env = "MODELOPT_ATTN_SOFTMAX_MODE" + extra_env_var = "VLLM_RAY_EXTRA_ENV_VARS_TO_COPY" + existing_env = "EXISTING_ENV" + monkeypatch.setenv(extra_env_var, existing_env) + executor = SimpleNamespace(ADDITIONAL_ENV_VARS={"EXISTING_CLASS_VAR"}) + real_import_module = importlib.import_module + attempted_ray_modules = [] + + def import_module(name, package=None): + if name not in ray_modules: + return real_import_module(name, package) + attempted_ray_modules.append(name) + if name == available_module: + return SimpleNamespace(RayDistributedExecutor=executor) + raise ModuleNotFoundError(name) + + monkeypatch.setattr(importlib, "import_module", import_module) + + _load_launcher_module(f"shared_attention_launcher_quant_test_{available_module}") + + expected_attempts = ray_modules[:1] if available_module == ray_modules[0] else ray_modules + assert tuple(attempted_ray_modules) == expected_attempts + if available_module is not None: + assert {"EXISTING_CLASS_VAR", softmax_env} == executor.ADDITIONAL_ENV_VARS + assert os.environ[extra_env_var] == existing_env + else: + assert set(os.environ[extra_env_var].split(",")) == {existing_env, softmax_env} + + +def test_validation_of_all_layouts_precedes_mutation(monkeypatch): + _patch_conversion(monkeypatch) + good, bad = _attention(), _attention(AttentionType.ENCODER) + 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_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_quant_memory_profile_uses_inference_mode_and_disables_compilation(monkeypatch): + 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), + ] + + +@pytest.mark.parametrize( + ("mode", "rejected"), + [(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(), + model_config=SimpleNamespace(dtype=torch.float16), + compilation_config=SimpleNamespace(cudagraph_mode=mode), + ) + errors = worker_module._global_errors( + 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}, + "decode": {"a": 0.1, "b": 1.0}, + } + } + + assert "calibrated decode skip-softmax" in worker_module._sparse_graph_error( + sparse_kw, CUDAGraphMode.FULL_AND_PIECEWISE, api + ) + 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 e9f8ee73d95..37d91cf5447 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -15,20 +15,121 @@ """Tests for sparse attention vLLM worker compatibility helpers.""" +import builtins +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 +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 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, ) +_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 _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", + "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") + 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( + ("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 _make_old_impl(): """Create a vLLM FlashAttention impl with initialized runtime state.""" @@ -66,6 +167,547 @@ 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, softmax_mode="fp32"): + impl = _clone_sparse_impl(_make_old_flashinfer_impl(), get_flashinfer_sparse_impl_cls()) + impl.sparse_kw = {"sparsity_n": 2, "sparsity_m": 4} if sparse else {} + impl.quant_kw = { + "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, + "softmax_mode": softmax_mode, + } + return impl + + +@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 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, + ) + + metadata = builder.build(0, common, fast_build=False) + + 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( + isolated_flashinfer_builder_patch, +): + old_impl = _make_old_flashinfer_impl() + + 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" + 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)] + + +@pytest.mark.parametrize( + ("query_quantized", "softmax_mode"), + [(True, "fp32"), (False, "mixed_fp16")], + ids=["q-only", "mixed-softmax-only"], +) +def test_flashinfer_q_only_transform_does_not_fallback(monkeypatch, query_quantized, softmax_mode): + """Q QDQ and mixed softmax must run when other transforms are inactive.""" + impl = _make_flashinfer_impl(softmax_mode=softmax_mode) + query = torch.zeros(4, 2, 64, dtype=torch.float16) + kv_cache = torch.zeros(4, 2, 16, 2, 64, dtype=torch.float16) + metadata = _flashinfer_metadata(query_lens=(2, 2), max_query_len=4) + captured = {} + layer = SimpleNamespace(_query_quant_in_kernel=query_quantized) + if query_quantized: + layer.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 + captured.update(kwargs) + return torch.zeros_like(query) + + monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) + + impl.forward(layer, query, query, query, kv_cache, metadata, output=torch.empty_like(query)) + + expected_query = 1 if query_quantized else 0 + assert torch.all(captured["query"] == expected_query) + assert captured["softmax_mode"] == softmax_mode + + +def test_flashinfer_legacy_inactive_launch_writes_only_in_native_fallback(monkeypatch): + 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", "softmax_mode"), + [(False, "fp32"), (True, "mixed_fp16")], + ids=["sparse-only", "quantized-mixed-fp16"], +) +def test_flashinfer_mixed_batch_splits_decode_and_prefill(monkeypatch, quantized, softmax_mode): + prefill_tokens = 17 + impl = _make_flashinfer_impl( + sparse=True, + quantized=quantized, + softmax_mode=softmax_mode, + ) + query = torch.zeros(1 + prefill_tokens, 2, 64, dtype=torch.float16) + kv_cache = torch.zeros(4, 2, 16, 2, 64, dtype=torch.float16) + metadata = _flashinfer_metadata(query_lens=(1, prefill_tokens), mixed=True) + 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 + assert prefill_kw["softmax_mode"] == softmax_mode + 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 decode_kw["softmax_mode"] == "mixed_fp16" + assert "sparsity_n" not in decode_kw + assert prefill_kw["p_qdq"] == "nvfp4" + assert [kw["max_new_tokens"] for kw in calls["finalize"]] == [1, prefill_tokens] + else: + assert calls["native"] == [True] + assert calls["decode"] == [] + assert calls["finalize"] == [] + + +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): + 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_rejects_unsupported_backend_before_mutation(monkeypatch): + worker_module = _load_worker_module() + good = _bare_attention(worker_module) + bad = _bare_attention(worker_module) + bad.impl = SimpleNamespace() + original_good_impl = good.impl + state = _sparse_worker_state(nn.ModuleDict({"good": good, "bad": bad})) + 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) as exc: + worker_module._install_attention(state, quantize=False) + + assert str(exc.value) == ( + "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 + + +@pytest.mark.parametrize( + ("failure", "active", "expected"), + [ + ( + "cascade", + True, + "vLLM cascade attention is incompatible with an active ModelOpt attention transform", + ), + ( + "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, expected +): + impl = _make_flashinfer_impl(sparse=active) + metadata = ( + _flashinfer_metadata() + if failure == "output" + else SimpleNamespace(use_cascade=failure == "cascade") + ) + 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) 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, + output_scale=torch.ones(1), + ) + 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()) @@ -117,7 +759,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": { @@ -129,6 +770,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 +793,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 +855,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): @@ -250,6 +879,152 @@ def fake_attention(q, **kwargs): assert "target_sparse_ratio" not in captured +def test_unsupported_nvfp4_bmm_block_size_raises(): + """P/V mappings reject NVFP4 quantizers that are not block-16.""" + quantizer = SimpleNamespace( + is_enabled=True, + is_nvfp4_dynamic=True, + block_sizes={-1: 32}, + ) + layer = SimpleNamespace(p_bmm_quantizer=quantizer, v_bmm_quantizer=quantizer) + with pytest.raises(NotImplementedError, match="p_bmm_quantizer"): + vllm_plugin._p_qdq_from_layer(layer) + with pytest.raises(NotImplementedError, match="v_bmm_quantizer"): + vllm_plugin._v_qdq_from_layer(layer) + + +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: + 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(), + ) + q_inputs = [] + + def quantize_q(query): + assert query.dtype == torch.float32 + 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, + "softmax_mode": "mixed_fp16", + } + 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=q.shape[0], + 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_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): + calls["query"] = query.clone() + 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) + monkeypatch.setattr( + vllm_plugin, + "triton_attention", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("shared kernel")), + ) + 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 == { + "max_new_tokens": 1, + "page_size": 16, + "v_qdq_scale": 1.0, + } + 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["softmax_mode"] == "mixed_fp16" + assert decode_kw["v_cache_quantized"] is True + assert len(q_inputs) == 1 and q_inputs[0].shape[0] == q.shape[0] + torch.testing.assert_close(q_inputs[0][:2], q[:2].float()) + 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): + """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 = _flash_attention_metadata(1, 16) + 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" + assert captured["softmax_mode"] == "fp32" + + def test_resolve_calibrated_skip_softmax_threshold_for_decode(): """Calibration conversion is phase-aware even when decode later delegates.""" sparse_kw = { @@ -312,66 +1087,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()) @@ -380,18 +1095,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..d78ae5295b2 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,38 @@ # 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 + + +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 test_triton_fa_importable_on_cpu(): @@ -38,3 +59,143 @@ def test_triton_fa_importable_on_cpu(): assert "attention" in triton_fa.__all__ assert callable(calibrate.attention_calibrate) + + +def test_softmax_mode_validation(): + pytest.importorskip("triton") + + from modelopt.torch.kernels.quantization.attention.softmax_fakequant import resolve_softmax_mode + + assert resolve_softmax_mode("fp32") is False + assert resolve_softmax_mode("mixed_fp16") is True + with pytest.raises(ValueError, match="softmax_mode"): + resolve_softmax_mode("fp16") + + +def test_forward_buckets_autotune_key_without_bucketing_grid(monkeypatch): + """Reuse autotune results by length regime without launching extra query tiles.""" + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + 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) + + +def test_forward_uses_minimal_shared_autotune_configs(): + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + 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 triton_fa._attn_fwd.keys == ["N_CTX", "HEAD_DIM", "Q_IS_FP32", "P_QDQ", "V_QDQ"] + + +@pytest.mark.parametrize( + ("attention_kwargs", "expected_p_qdq", "expected_v_qdq", "expected_mixed_fp16"), + [ + ({}, 0, 0, False), + ({"softmax_mode": "mixed_fp16"}, 0, 0, True), + ({"p_qdq": "fp8"}, 1, 0, False), + ({"p_qdq": "nvfp4", "v_qdq": "nvfp4"}, 2, 2, False), + ( + {"softmax_mode": "mixed_fp16", "p_qdq": "nvfp4", "v_qdq": "nvfp4"}, + 2, + 2, + True, + ), + ({"p_qdq": "nvfp4", "sparsity_n": 2, "sparsity_m": 4}, 2, 0, False), + ], +) +def test_forward_routes_every_mode_to_single_autotuner( + monkeypatch, attention_kwargs, expected_p_qdq, expected_v_qdq, expected_mixed_fp16 +): + """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 + assert kernel.kwargs["MIXED_FP16"] is expected_mixed_fp16 + + +@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