diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index a8f7d2ead18..858243686d0 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.15.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,50 @@ 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 sparse-only installer passes through and vLLM runs unchanged. Whole-model fakequant flows remain handled by `vllm_serve_fakequant.py`; the compact attention-only path is below. + +The reusable serving policies live in `modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py`. `install_vllm_sparse_attention_from_checkpoint` installs checkpoint-driven sparse-only attention, while `install_vllm_nvfp4_attention` installs fixed NVFP4 Q/K/P/V with optional checkpoint sparsity. Both validate every selected layer before publishing any replacement implementation and return a `VllmAttentionInstallReport` with the installed layer names and backend counts. + +`sparse_attn_worker.py` only invokes these APIs after vLLM loads the model. It retains `SparseAttnWorker` as the launcher's default and provides `QuantSparseAttnWorker` for the compact NVFP4 policy. Other vLLM integrations can invoke the same library APIs directly: + +```python +from modelopt.torch.sparsity.attention_sparsity.plugins.vllm_runtime import ( + install_vllm_nvfp4_attention, +) + +report = install_vllm_nvfp4_attention(model_runner, sparse_cfg="checkpoint") +``` 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.15.0 or newer is required when either worker activates a ModelOpt attention transform. Importing `SparseAttnWorker`, or using it with no checkpoint sparse metadata, does not resolve quant-only APIs. + +Use the same launcher with the compact worker. By default, vLLM selects the backend for the model and platform; NemotronH on Blackwell selects FlashInfer: + +```bash +python vllm_serve_sparse_attn.py -tp 8 \ + --no-enable-prefix-caching \ + --worker-cls sparse_attn_worker.QuantSparseAttnWorker +``` + +The installer supports both FlashInfer and FlashAttention, and the worker prints the installed adapter counts. Pass `--attention-backend FLASHINFER` or `--attention-backend FLASH_ATTN` only when an explicit override is needed. + +This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant format to Q/K/P/V. Q is dynamic; missing K/V scales default to global scale 1.0, and P defaults to amax 1.0. Existing scalar attention amax values are preserved, but this path does not calibrate or restore them itself. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored. + +Decode uses a fixed 32-split, 128-key-tile schedule. P QDQ consumes split-local, +unnormalized online-softmax probabilities, so changing that schedule can change +quantized results; split count is part of the numerical contract. + +K is QDQ before its cache write, while V is written pristine. Complete 16-token V groups are finalized once in cache; an incomplete tail remains pristine and is QDQ on read. P@V therefore sees uniform fakequant values without re-quantizing the tail. + +Supported configurations are regular decoder self-attention with FlashInfer or FlashAttention, fp16/bf16 model and KV cache, equal Q/K/V head dimensions that are multiples of 16, and DCP 1. The FlashInfer adapter preserves both NHD and HND cache strides and separates mixed decode/prefill launches so each phase keeps its own kernel contract. The default `FULL_AND_PIECEWISE` mode remains enabled for fixed N:M and attention-only NVFP4; checkpoints with calibrated decode `threshold_scale_factor` must use a non-`FULL` decode graph mode such as `--enforce-eager` because the live sequence length is not replayed as a Python scalar. + +Unsupported features are sliding window, ALiBi, softcap, sinks, FP8 KV cache, cross/encoder/MLA attention, KV sharing or transfer, prefix caching, speculative decoding, DBO/ubatching, and `FULL` mixed/prefill CUDA graphs. ## Known Problems diff --git a/examples/vllm_serve/sparse_attn_worker.py b/examples/vllm_serve/sparse_attn_worker.py index 1057baa870e..7b5fb28bf68 100644 --- a/examples/vllm_serve/sparse_attn_worker.py +++ b/examples/vllm_serve/sparse_attn_worker.py @@ -13,108 +13,96 @@ # 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 -""" - -import importlib - -try: - _has_legacy_attention_layer = importlib.util.find_spec("vllm.attention.layer") is not None -except (ModuleNotFoundError, ValueError): - _has_legacy_attention_layer = False - -if _has_legacy_attention_layer: - from vllm.attention.layer import Attention as VLLMAttention -else: - from vllm.model_executor.layers.attention import Attention as VLLMAttention +"""vLLM worker lifecycle wiring for ModelOpt attention transforms.""" from vllm.v1.worker.gpu_worker import Worker as BaseWorker -from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_config import ( - load_from_checkpoint_metadata, - match_sparse_config, -) -from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( - _build_sparse_kw, - _clone_sparse_impl, +from modelopt.torch.sparsity.attention_sparsity.plugins.vllm_runtime import ( + install_vllm_nvfp4_attention, + install_vllm_sparse_attention_from_checkpoint, ) +__all__ = ["SparseAttnWorker", "QuantSparseAttnWorker"] # noqa: RUF022 -def _replace_attention_impl(worker): - """Replace FlashAttentionImpl with ModelOptSparseAttentionImpl on all Attention layers. +_QUANT_FORMAT_KEYS = ("q_format", "k_format", "p_format", "v_format") - The sole configuration source is the checkpoint's ``sparse_attention_config`` - metadata. No-op if the checkpoint has no such block. - """ - hf_config = getattr(worker.model_runner.model_config, "hf_config", None) - detected = load_from_checkpoint_metadata(hf_config) - if detected is None: + +def _unwrapped_model(worker): + model = worker.model_runner.model + return model.unwrap() if hasattr(model, "unwrap") else model + + +def _print_install_report(policy, report) -> None: + if report.installed_count: + if policy != "Sparse attention": + print( + f"[ModelOpt] Installed {policy} (quant+sparse) on " + f"{report.installed_count} layers: {dict(report.backend_counts)}" + ) + else: + if report.sparse_algorithm: + print(f"[ModelOpt] Sparse attention config: algo -> {report.sparse_algorithm}") + print( + f"[ModelOpt] Sparse attention: replaced impl on {report.installed_count} " + f"attention layers: {dict(report.backend_counts)}" + ) + elif report.sparse_algorithm: + print( + f"[ModelOpt] Sparse attention config {report.sparse_algorithm} matched no active " + "attention layers; vLLM remains unchanged" + ) + else: 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 - for name, module in model.named_modules(): - if not isinstance(module, VLLMAttention): - continue - layer_cfg = match_sparse_config(name, cfg) - if layer_cfg is None or not layer_cfg.get("enable", True): - continue - - sparse_kw = _build_sparse_kw(layer_cfg) - if not sparse_kw: - # Keep vLLM's original impl when the exported layer config does not - # enable any sparse feature. - continue - new_impl = _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") +class SparseAttnWorker(BaseWorker): + """Install checkpoint-driven sparse attention after model loading.""" + def load_model(self, *args, **kwargs) -> None: + """Load the model, then install checkpoint-configured attention.""" + super().load_model(*args, **kwargs) + report = install_vllm_sparse_attention_from_checkpoint(self.model_runner) + _print_install_report("Sparse attention", report) -# --------------------------------------------------------------------------- -# Workers -# --------------------------------------------------------------------------- +class QuantSparseAttnWorker(BaseWorker): + """Install quantized attention plus optional checkpoint sparsity. -class SparseAttnWorker(BaseWorker): - """vLLM worker that uses the ModelOpt sparse attention backend. + Per-operand formats come from vLLM's ``--additional-config``; absent keys + default to NVFP4 on all four operands (Q/K/P/V):: - Replaces FlashAttentionImpl with ModelOptSparseAttentionImpl on each - Attention module right after model loading — before any forward pass - (including determine_available_memory profiling). + --additional-config '{"modelopt_attn_quant": {"p_format": "fp8", "v_format": "fp8"}}' """ + def _quant_formats(self) -> dict[str, str]: + additional = getattr(self.vllm_config, "additional_config", None) or {} + formats = additional.get("modelopt_attn_quant", {}) + unknown = set(formats) - set(_QUANT_FORMAT_KEYS) + if unknown: + raise ValueError( + f"unknown modelopt_attn_quant keys {sorted(unknown)}; " + f"allowed: {list(_QUANT_FORMAT_KEYS)}" + ) + return dict(formats) + def load_model(self, *args, **kwargs) -> None: - """Load model, then replace attention impl with sparse variant.""" + """Load the model, then install the configured attention quant recipe.""" super().load_model(*args, **kwargs) - _replace_attention_impl(self) + formats = self._quant_formats() + report = install_vllm_nvfp4_attention(self.model_runner, sparse_cfg="checkpoint", **formats) + policy = "NVFP4 attention" if not formats else f"Quant attention ({formats})" + _print_install_report(policy, report) + + def determine_available_memory(self) -> int: + """Profile memory without compiling the dynamically converted modules.""" + # Sparse-only imports must remain independent of quantization-specific APIs. + import torch + + from modelopt.torch.quantization.plugins.vllm import disable_compilation + + with torch.inference_mode(), 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..777aff8fc7d 100644 --- a/examples/vllm_serve/vllm_serve_sparse_attn.py +++ b/examples/vllm_serve/vllm_serve_sparse_attn.py @@ -15,14 +15,14 @@ """Launch vLLM with sparse attention. -Configuration is read exclusively from ``/config.json``'s -``sparse_attention_config`` block, written during calibration by +The default ``SparseAttnWorker`` reads configuration exclusively from +``/config.json``'s ``sparse_attention_config`` block, written by ``examples/llm_sparsity/attention_sparsity/hf_sa.py``. If the checkpoint has -no such block, the worker logs a message and the server runs as standard +no such block, that worker logs a message and the server runs as standard vLLM. -Combined sparse attention + quantization is not handled by this launcher; it -will be added in a follow-up PR once the combined path is tested. +The launcher defaults to ``sparse_attn_worker.SparseAttnWorker``. Pass +``--worker-cls sparse_attn_worker.QuantSparseAttnWorker`` for quant+sparse. Usage: python vllm_serve_sparse_attn.py diff --git a/modelopt/torch/kernels/common/attention/decode_attention.py b/modelopt/torch/kernels/common/attention/decode_attention.py new file mode 100644 index 00000000000..e99c9b0cbff --- /dev/null +++ b/modelopt/torch/kernels/common/attention/decode_attention.py @@ -0,0 +1,374 @@ +# 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.common.fp8_quant import fp8_scalar_qdq + +__all__ = ["attention_decode"] + +_BLOCK_N = 128 +_DEFAULT_KV_SPLITS = 32 +_MAX_KV_SPLITS = 32 +_P_QDQ_MODES = {None: 0, "fp8": 1, "nvfp4": 2} +_V_QDQ_MODES = {None, "nvfp4"} + + +@triton.jit +def _decode_split_kernel( + Q, + B_seq_len_k, + M_partial, + L_partial, + Acc_partial, + K_cache, + V_cache, + Block_table, + qk_scale, + stride_qb, + stride_qh, + stride_mb, + stride_mh, + stride_ab, + stride_ah, + stride_as, + stride_kc_block, + stride_kc_pos, + stride_kc_head, + stride_vc_block, + stride_vc_pos, + stride_vc_head, + p_qdq_scale, + v_qdq_scale, + kv_group_num: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, + PAGE_SIZE: tl.constexpr, + max_blocks_per_seq, + NUM_KV_SPLITS: tl.constexpr, + P_QDQ: tl.constexpr, + V_QDQ: tl.constexpr, + V_CACHE_QUANTIZED: tl.constexpr, +): + """Compute one partial softmax for one request, query head, and KV split.""" + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + split_idx = tl.program_id(2) + kv_head_idx = head_idx // kv_group_num + seq_len_kv = tl.load(B_seq_len_k + batch_idx) + v_quantized_boundary = (seq_len_kv // 16) * 16 + + num_tiles = tl.cdiv(seq_len_kv, BLOCK_N) + tiles_per_split = tl.cdiv(num_tiles, NUM_KV_SPLITS) + tile_lo = split_idx * tiles_per_split + tile_hi = tl.minimum(tile_lo + tiles_per_split, num_tiles) + kv_lo = tile_lo * BLOCK_N + kv_hi = tile_hi * BLOCK_N + + dim_pos = tl.arange(0, BLOCK_D) + d_mask = dim_pos < HEAD_DIM + kv_pos = tl.arange(0, BLOCK_N) + q = tl.load( + Q + batch_idx * stride_qb + head_idx * stride_qh + dim_pos, + mask=d_mask, + other=0.0, + ).to(tl.float32) + + running_max = -float("inf") + running_sum = 0.0 + acc = tl.zeros([BLOCK_D], dtype=tl.float32) + + for kv_start in range(kv_lo, kv_hi, BLOCK_N): + kv_start = tl.multiple_of(kv_start, BLOCK_N) + kv_valid = kv_start + kv_pos < seq_len_kv + k = _load_paged_k_tile( + K_cache, + Block_table, + batch_idx, + kv_head_idx, + kv_start, + kv_pos, + dim_pos, + seq_len_kv, + stride_kc_block, + stride_kc_pos, + stride_kc_head, + PAGE_SIZE, + BLOCK_N, + BLOCK_D, + HEAD_DIM, + max_blocks_per_seq, + ).to(tl.float32) + scores = tl.sum(q[:, None] * k, axis=0) * qk_scale + scores = tl.where(kv_valid, scores, -float("inf")) + tile_max = tl.max(scores, axis=0) + new_max = tl.maximum(running_max, tile_max) + p = tl.math.exp2(scores - new_max) + p = tl.where(kv_valid, p, 0.0) + correction = tl.math.exp2(running_max - new_max) + running_sum = running_sum * correction + tl.sum(p, axis=0) + acc *= correction + + if P_QDQ == 1: + # FP8 E4M3 per-tensor QDQ of the split-local unnormalized P + # (elementwise -> split-count- and batch-shape-invariant). + p = fp8_scalar_qdq(p, p_qdq_scale) + elif P_QDQ == 2: + 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: + allowed = _P_QDQ_MODES if operand == "p" else _V_QDQ_MODES + if mode not in allowed: + raise ValueError( + f"{operand}_qdq must be one of {sorted(m for m in allowed if m)} 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}") + # FP8 uses the per-tensor convention ``amax / 448``; NVFP4 the two-level + # global scale ``amax / (6 * 448)`` (matches the prefill kernel). + return amax / 448.0 if mode == "fp8" else amax / (6.0 * 448.0) + + +def attention_decode( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + b_seq_len_k: torch.Tensor, + *, + softmax_scale: float | None = None, + page_size: int = 16, + num_kv_splits: int = _DEFAULT_KV_SPLITS, + p_qdq: str | None = None, + p_qdq_amax: float = 1.0, + v_qdq: str | None = None, + v_qdq_amax: float | None = None, + v_cache_quantized: bool = False, +) -> torch.Tensor: + """Decode one query token per request over a paged KV cache. + + Q and K are expected to be fake-quantized before this call. 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. + """ + if q.ndim != 3: + raise ValueError(f"q must have shape [batch, heads, head_dim], got {tuple(q.shape)}") + if page_size != k_cache.shape[1] or page_size != v_cache.shape[1]: + raise ValueError("page_size must match both paged KV cache tensors") + if not 1 <= num_kv_splits <= _MAX_KV_SPLITS: + raise ValueError(f"num_kv_splits must be in [1, {_MAX_KV_SPLITS}], got {num_kv_splits}") + batch, num_q_heads, head_dim = q.shape + num_kv_heads = k_cache.shape[2] + if num_q_heads % num_kv_heads: + raise ValueError("num_q_heads must be divisible by num_kv_heads") + if b_seq_len_k.shape != (batch,) or block_table.shape[0] != batch: + raise ValueError("decode metadata batch dimension must match q") + + p_qdq_scale = _qdq_scale(p_qdq, p_qdq_amax, "p") + v_qdq_scale = _qdq_scale(v_qdq, v_qdq_amax, "v") + if v_cache_quantized and v_qdq != "nvfp4": + raise ValueError("v_cache_quantized requires v_qdq='nvfp4'") + q = q.contiguous() + block_d = triton.next_power_of_2(head_dim) + if (p_qdq == "nvfp4" or v_qdq == "nvfp4") and head_dim % 16: + raise ValueError("NVFP4 decode requires dimensions divisible by 16") + qk_scale = (head_dim**-0.5 if softmax_scale is None else softmax_scale) * LOG2E + m_partial = torch.empty(batch, num_q_heads, num_kv_splits, dtype=torch.float32, device=q.device) + l_partial = torch.empty_like(m_partial) + acc_partial = torch.empty( + batch, num_q_heads, num_kv_splits, block_d, dtype=torch.float32, device=q.device + ) + output = torch.empty_like(q) + + with torch.cuda.device(q.device): + _decode_split_kernel[(batch, num_q_heads, num_kv_splits)]( + q, + b_seq_len_k, + m_partial, + l_partial, + acc_partial, + k_cache, + v_cache, + block_table, + qk_scale, + q.stride(0), + q.stride(1), + m_partial.stride(0), + m_partial.stride(1), + acc_partial.stride(0), + acc_partial.stride(1), + acc_partial.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + p_qdq_scale, + v_qdq_scale, + kv_group_num=num_q_heads // num_kv_heads, + BLOCK_D=block_d, + BLOCK_N=_BLOCK_N, + HEAD_DIM=head_dim, + PAGE_SIZE=page_size, + max_blocks_per_seq=block_table.shape[1], + NUM_KV_SPLITS=num_kv_splits, + P_QDQ=_P_QDQ_MODES[p_qdq], + 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..5acc9fb787c 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -29,52 +29,51 @@ import triton import triton.language as tl -# Helpers for optional N:M sparsity and sink/window-aware dense regions live -# in the sparsity package. The baseline forward kernel below calls them -# conditionally under constexpr guards, so the unified single-kernel design -# stays intact while keeping feature-specific logic in its own subpackage. +# Helpers for optional N:M sparsity and skip-softmax live in the sparsity +# package. The baseline forward kernel below calls them conditionally under +# constexpr guards, so the unified single-kernel design stays intact while +# keeping feature-specific logic in its own subpackage. # # Lazy import: Triton resolves @triton.jit names at kernel compile time (first # call), not at definition time, so populating the module globals before the # first ``attention()`` call is sufficient. Deferring avoids a circular import # (common.attention/__init__.py ↔ sparsity.attention/__init__.py via this file). _apply_sparse_nm_to_qk_tile: Any = None -_is_dense_region: Any = None _skip_softmax_decision: Any = None -_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 +82,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 +131,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 +175,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 +216,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 +281,7 @@ def _attn_fwd( IS_CAUSAL: tl.constexpr, # Whether to apply causal mask HEAD_DIM: tl.constexpr, # Actual head dimension (for d_mask) STORE_LSE: tl.constexpr, # Whether to save LSE for backward pass + Q_IS_FP32: tl.constexpr, # Dynamic NVFP4 QDQ carrier uses FP32 SPARSITY_N: tl.constexpr = 0, # N:M sparsity — keep top-N of every M elements (0 = disabled) SPARSITY_M: tl.constexpr = 4, # N:M sparsity — group size (4 or 8) DENSE_SINK_TOKENS: tl.constexpr = 0, # Leading KV tokens kept dense (attention sinks) @@ -263,6 +290,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 +353,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 +388,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 @@ -404,8 +440,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 +477,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 +670,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 +826,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: @@ -833,6 +892,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 +980,18 @@ 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, + # An fp32 Q (the dynamic-quant carrier) always needs the fp32 BMM1 dot; + # gating on QDQ modes would miss recipes like fp8-BMM2 (P mode 1). + "Q_IS_FP32": q.dtype == torch.float32, "SPARSITY_N": sparsity_n, "SPARSITY_M": sparsity_m, "DENSE_SINK_TOKENS": dense_sink_tokens, @@ -932,6 +1000,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 +1036,7 @@ def grid(META): _attn_fwd.fn[grid]( *fwd_args, **fwd_kwargs, - BLOCK_M=_MEASURE_BLOCK_M, + BLOCK_M=_P_QDQ_MEASURE_BLOCK_M if p_qdq_mode else _MEASURE_BLOCK_M, BLOCK_N=_MEASURE_BLOCK_N, num_warps=_MEASURE_NUM_WARPS, num_stages=_MEASURE_NUM_STAGES, @@ -1137,6 +1208,9 @@ def backward(ctx, grad_output): None, # measure_sparsity None, # p_qdq_mode None, # p_qdq_scale + None, # v_qdq_mode + None, # v_qdq_scale + None, # v_cache_quantized None, # k_cache None, # v_cache None, # block_table @@ -1165,6 +1239,9 @@ def attention( measure_sparsity: bool = False, p_qdq: str | None = None, p_qdq_amax: float = 1.0, + v_qdq: str | None = None, + v_qdq_amax: float | None = None, + v_cache_quantized: bool = False, k_cache: torch.Tensor | None = None, v_cache: torch.Tensor | None = None, block_table: torch.Tensor | None = None, @@ -1211,8 +1288,8 @@ def attention( BLOCK_N is a multiple of 16). The softmax denominator stays unquantized. The backward pass uses the straight-through estimator: gradients are computed from the unquantized P, matching QAT - references that keep the backward dots in high precision. - Set to ``None`` to disable. + references that keep the backward dots in high precision. Set to + ``None`` to disable. p_qdq_amax: Per-tensor amax for the softmax-P quant-dequant. The kernel's unnormalized P lies in [0, 1] (the max-subtraction caps every entry at ``exp2(0) = 1``), so 1 is the theoretical upper @@ -1221,6 +1298,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 +1328,7 @@ 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() 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 +1343,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, @@ -1280,6 +1380,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/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 aa6141dd48b..3a69b2d202e 100644 --- a/modelopt/torch/quantization/plugins/vllm.py +++ b/modelopt/torch/quantization/plugins/vllm.py @@ -23,21 +23,83 @@ 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 from ...utils.distributed import ParallelState +from ..conversion import set_quantizer_by_cfg from ..nn import QuantLinearConvBase, QuantModule, QuantModuleRegistry, TensorQuantizer from .custom import CUSTOM_MODEL_PLUGINS +_NVFP4_ATTENTION_QUANTIZER_CFG = { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, +} +_VLLM_NVFP4_ATTENTION_QUANT_CFG = [ + {"quantizer_name": "*_bmm_quantizer", "enable": False}, + *( + { + "quantizer_name": f"*{name}_bmm_quantizer", + "cfg": _NVFP4_ATTENTION_QUANTIZER_CFG, + "enable": True, + } + for name in ("q", "k", "p", "v") + ), +] +_FP8_ATTENTION_QUANTIZER_CFG = {"num_bits": (4, 3)} # per-tensor E4M3, static scale amax/448 +_BMM2_FORMAT_CFGS = {"nvfp4": _NVFP4_ATTENTION_QUANTIZER_CFG, "fp8": _FP8_ATTENTION_QUANTIZER_CFG} + + +def build_vllm_attention_quant_cfg( + *, + q_format: str = "nvfp4", + k_format: str = "nvfp4", + p_format: str = "nvfp4", + v_format: str = "nvfp4", +) -> list: + """Build the attention BMM quantizer config with per-operand formats. + + Each of Q/K/P/V takes "nvfp4" (dynamic block-16, two-level scale) or + "fp8" (per-tensor E4M3, static scale amax/448). + """ + formats = {"q": q_format, "k": k_format, "p": p_format, "v": v_format} + for name, fmt in formats.items(): + if fmt not in _BMM2_FORMAT_CFGS: + raise ValueError( + f"{name}_format must be one of {sorted(_BMM2_FORMAT_CFGS)}, got {fmt!r}" + ) + return [ + {"quantizer_name": "*_bmm_quantizer", "enable": False}, + *( + { + "quantizer_name": f"*{name}_bmm_quantizer", + "cfg": _BMM2_FORMAT_CFGS[fmt], + "enable": True, + } + for name, fmt in formats.items() + ), + ] + + +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 +130,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 +216,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 +242,82 @@ 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 _set_vllm_attention_fp8_bmm2_default_amax(module, device: torch.device) -> None: + """Set fixed amax on uncalibrated per-tensor FP8 BMM2 quantizers (P=1.0, V=448).""" + for name, default in ( + ("q_bmm_quantizer", 448.0), + ("k_bmm_quantizer", 448.0), + ("p_bmm_quantizer", 1.0), + ("v_bmm_quantizer", 448.0), + ): + quantizer = getattr(module, name, None) + if ( + not isinstance(quantizer, TensorQuantizer) + or not quantizer.is_enabled + or getattr(quantizer, "num_bits", None) != (4, 3) + or getattr(quantizer, "block_sizes", None) + or hasattr(quantizer, "_amax") + ): + continue + quantizer.amax = torch.tensor(default, device=device, dtype=torch.float32) + + +def configure_vllm_nvfp4_attention_quantizers( + module: torch.nn.Module, + *, + device: torch.device | str, + dtype: torch.dtype, + cfg: list | None = None, +) -> torch.nn.Module: + """Configure one vLLM Attention module for fused NVFP4 fake quantization. + + This attention-scoped entry point avoids recursively converting vLLM Linear and MoE + modules. It configures only quantizer state; the caller remains responsible for installing + the fused attention implementation and enabling its in-kernel Q/V carriers. + + Args: + module: A vLLM ``Attention`` module to convert and configure in place. + device: Device on which the attention quantizer state should reside. + dtype: Model compute dtype associated with the attention module. + + Returns: + The supplied module, converted in place to ``_QuantVLLMAttention``. + """ + if not isinstance(module, vllm_attention.Attention): + raise TypeError(f"Expected vLLM Attention, got {type(module).__name__}") + if not isinstance(dtype, torch.dtype): + raise TypeError(f"Expected torch.dtype, got {type(dtype).__name__}") + + device = torch.device(device) + module.device, module.dtype = device, dtype + if not isinstance(module, _QuantVLLMAttention): + module = QuantModuleRegistry.convert(module) + if not hasattr(module, "p_bmm_quantizer"): + module.p_bmm_quantizer = TensorQuantizer() + + set_quantizer_by_cfg(module, _VLLM_NVFP4_ATTENTION_QUANT_CFG if cfg is None else cfg) + for name in ("q", "k", "p", "v"): + getattr(module, f"{name}_bmm_quantizer").to(device=device) + _set_vllm_attention_kv_default_amax(module, device) + _set_vllm_attention_fp8_bmm2_default_amax(module, device) + return module + + 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) @@ -497,9 +627,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..00411eaa4ee 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,326 @@ 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 QDQ 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 ( + getattr(quantizer, "is_nvfp4_dynamic", False) + and (quantizer.block_sizes or {}).get(-1) == 16 + ): + mode = "nvfp4" + elif getattr(quantizer, "num_bits", None) == (4, 3) and not getattr( + quantizer, "block_sizes", None + ): + # Per-tensor FP8 E4M3 (static scale amax/448) + mode = "fp8" + else: + raise NotImplementedError( + f"{attr} is enabled with an unsupported format; only dynamic block-16 NVFP4 " + "or per-tensor FP8 E4M3 is supported" + ) + amax = getattr(quantizer, "_amax", None) + if amax is None: + return mode, default_amax + if getattr(amax, "numel", lambda: 1)() != 1: + raise NotImplementedError(f"{attr} requires a scalar amax, got shape {tuple(amax.shape)}") + return mode, 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 P/V QDQ contract once for one attention launch.""" + quant_kw = getattr(impl, "quant_kw", None) + if quant_kw is None: + p_qdq, p_qdq_amax = _p_qdq_from_layer(layer) + v_qdq, v_qdq_amax = _v_qdq_from_layer(layer) + else: + p_qdq, p_qdq_amax = quant_kw["p_qdq"], quant_kw["p_qdq_amax"] + v_qdq, v_qdq_amax = quant_kw["v_qdq"], quant_kw["v_qdq_amax"] + return p_qdq, p_qdq_amax, v_qdq, v_qdq_amax + + +def _any_quant_active(layer, p_qdq, v_qdq) -> bool: + """Return whether native fallback would omit any Q/K/P/V transform.""" + k_quantizer = getattr(layer, "k_bmm_quantizer", None) + return bool( + p_qdq + or v_qdq + 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) + - Inherits from FlashAttentionImpl to reuse: - - __init__ (all configuration) - - do_kv_cache_update (KV cache writing) - Only overrides forward() to replace sparse prefill attention computation. +@dataclass(frozen=True, slots=True) +class _ResolvedForward: + p_qdq: str | None + p_qdq_amax: float + v_qdq: str | None + v_qdq_amax: float | None + quant_active: bool + + +def _resolve_forward( + impl, + layer, + attn_metadata, + output_scale, + output_block_scale, + *, + require_flashinfer_metadata: bool = False, +) -> _ResolvedForward | None: + """Resolve shared transform state or request the backend's native path.""" + p_qdq, p_qdq_amax, v_qdq, v_qdq_amax = _quant_kw_from_impl(impl, layer) + quant_active = _any_quant_active(layer, p_qdq, v_qdq) + transform_active = _should_run_modelopt_kernel(getattr(impl, "sparse_kw", None), quant_active) + + if getattr(attn_metadata, "use_cascade", False): + # Cascade is unimplemented by the ModelOpt kernel. Quantization must not be + # silently dropped (it would change numerics), so reject it; a sparse-only + # transform is numerically safe to delegate to the native dense path. + if transform_active and quant_active: + raise NotImplementedError( + "vLLM cascade attention is incompatible with active ModelOpt attention quantization" + ) + return None + + if require_flashinfer_metadata: + missing = [name for name in _FLASHINFER_METADATA_FIELDS if not hasattr(attn_metadata, name)] + if missing: + if transform_active: + raise NotImplementedError( + "FlashInfer metadata is missing the ModelOpt attention transform " + f"fields: {', '.join(missing)}" + ) + return None + + if transform_active and (output_scale is not None or output_block_scale is not None): + raise NotImplementedError("Fused attention output quantization is unsupported") + + return _ResolvedForward( + p_qdq=p_qdq, + p_qdq_amax=p_qdq_amax, + v_qdq=v_qdq, + v_qdq_amax=v_qdq_amax, + quant_active=quant_active, + ) + + +# 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, + 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, + ) + output[:batch] = triton_out + return output + + # Paged mode reads K/V through the cache. The dummy shape provides the GQA ratio. + k_dummy = torch.empty(0, impl.num_kv_heads, impl.head_size, device=q.device, dtype=q.dtype) + triton_out = triton_attention( + q, + k=k_dummy, + v=k_dummy, + b_start_loc=b_start_loc, + b_seq_len=b_seq_len, + max_input_len=max_query_len, + is_causal=is_causal, + softmax_scale=impl.scale, + b_start_loc_k=None, + b_seq_len_k=seq_lens, + max_input_len_k=max_seq_len, + k_cache=key_cache, + v_cache=value_cache, + block_table=block_table, + page_size=page_size, + p_qdq=p_qdq, + p_qdq_amax=p_qdq_amax, + v_qdq=v_qdq, + v_qdq_amax=v_qdq_amax, + v_cache_quantized=v_cache_quantized, + **sparse_kw, + ) + output[:num_actual_tokens] = triton_out + return output + + +def _dispatch_modelopt( + impl, + *, + query: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + num_actual_tokens: int, + max_query_len: int, + output: torch.Tensor, + num_decodes: int, + num_prefills: int, + num_decode_tokens: int, + num_prefill_tokens: int, + **common_kw, +) -> torch.Tensor: + """Run the ModelOpt path, splitting mixed decode+prefill batches by phase. + + NVFP4 P-QDQ is schedule-sensitive by design, so a decode result must not + depend on whether a prefill request is co-scheduled. When a batch mixes + ``q_len==1`` decode rows with ``q_len>1`` (chunked-)prefill rows, + ``max_query_len > 1`` and the whole batch would otherwise take the prefill + skip-softmax path. Split so each phase runs its own schedule -- decode rows + always take the fixed decode path. Both the FlashAttention and FlashInfer + adapters share this dispatch. """ + if not (num_decodes and num_prefills): + return _forward_modelopt( + impl, + query=query, + block_table=block_table, + seq_lens=seq_lens, + cu_seqlens_q=cu_seqlens_q, + num_actual_tokens=num_actual_tokens, + max_query_len=max_query_len, + output=output, + **common_kw, + ) + + if num_decode_tokens % num_decodes: + raise NotImplementedError("Non-uniform mixed decode is unsupported") + if num_decode_tokens + num_prefill_tokens != num_actual_tokens: + raise ValueError("Mixed-batch token counts do not match common metadata") + + # Sparse-only launches may have an inactive phase (for example N:M sparsity + # is prefill-only). Compute the native result once, then overwrite each + # active phase with its ModelOpt result. + if not common_kw.get("quant_active", False): + common_kw["dense_fallback"]() + + _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 + + +class ModelOptSparseAttentionImpl(FlashAttentionImpl): + """FlashAttention adapter for the compact ModelOpt Triton path.""" def _forward_vllm_flash_attn( self, @@ -163,58 +482,16 @@ 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. - return self._forward_vllm_flash_attn( - layer, - query, - key, - value, - kv_cache, - attn_metadata, - output, - output_scale, - 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) + native_result = None - # 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, - max_seq_len=attn_metadata.max_seq_len, - ) - if is_decode_only: - # N:M sparse softmax is prefill-only. - for name in ("sparsity_n", "sparsity_m", "dense_sink_tokens", "dense_recent_tokens"): - sparse_kw.pop(name, None) - if 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( + def native_forward(): + # Memoized: a split mixed batch may request the native dense result + # for an inactive phase after it was already computed for the batch. + nonlocal native_result + if native_result is None: + native_result = self._forward_vllm_flash_attn( layer, query, key, @@ -225,57 +502,50 @@ def forward( 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, - ) + return native_result - # 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, + resolved = _resolve_forward( + self, + layer, + attn_metadata, + output_scale, + output_block_scale, ) + if resolved is None: + return native_forward() - output[:num_actual_tokens] = triton_out - return output + key_cache, value_cache = kv_cache.unbind(0) + is_decode_only = attn_metadata.max_query_len <= 1 + common_kw = { + "layer": layer, + "key_cache": key_cache, + "value_cache": value_cache, + "max_seq_len": attn_metadata.max_seq_len, + "is_causal": getattr(attn_metadata, "causal", not is_decode_only), + "p_qdq": resolved.p_qdq, + "p_qdq_amax": resolved.p_qdq_amax, + "v_qdq": resolved.v_qdq, + "v_qdq_amax": resolved.v_qdq_amax, + "quant_active": resolved.quant_active, + "dense_fallback": native_forward, + } + # Split mixed decode+prefill batches so decode rows never fall into the + # schedule-sensitive prefill skip-softmax path (see _dispatch_modelopt). + return _dispatch_modelopt( + self, + query=query, + 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, + output=output, + num_decodes=getattr(attn_metadata, "num_decodes", 0), + num_prefills=getattr(attn_metadata, "num_prefills", 0), + num_decode_tokens=getattr(attn_metadata, "num_decode_tokens", 0), + num_prefill_tokens=getattr(attn_metadata, "num_prefill_tokens", 0), + **common_kw, + ) class ModelOptSparseAttentionBackend(FlashAttentionBackend): @@ -295,13 +565,244 @@ 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, + "quant_active": resolved.quant_active, + "dense_fallback": dense_fallback, + "prepare_modelopt": prepare_modelopt, + } + return _dispatch_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, + num_decodes=getattr(attn_metadata, "num_decodes", 0), + num_prefills=getattr(attn_metadata, "num_prefills", 0), + num_decode_tokens=getattr(attn_metadata, "num_decode_tokens", 0), + num_prefill_tokens=getattr(attn_metadata, "num_prefill_tokens", 0), + **common_kw, + ) + + +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 +811,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/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py new file mode 100644 index 00000000000..b4141740b64 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -0,0 +1,543 @@ +# 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. + +"""Install ModelOpt attention transforms into a loaded vLLM model.""" + +import importlib +from collections import Counter +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +import torch +import vllm +from packaging import version +from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl + +from . import vllm as attention_plugin +from .sparse_attn_config import load_from_checkpoint_metadata, match_sparse_config + +__all__ = [ + "VllmAttentionInstallReport", + "install_vllm_nvfp4_attention", + "install_vllm_sparse_attention_from_checkpoint", +] + + +def _import_attention_type() -> type: + """Import the concrete vLLM Attention type across supported releases.""" + 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.Attention + raise ImportError("No supported vLLM Attention module was found") + + +_VLLM_ATTENTION = _import_attention_type() + + +@dataclass(frozen=True, slots=True) +class VllmAttentionInstallReport: + """Summary of attention modules changed by a vLLM runtime installation.""" + + installed_layers: tuple[str, ...] = () + quantized_layers: tuple[str, ...] = () + sparse_layers: tuple[str, ...] = () + backend_counts: Mapping[str, int] = field(default_factory=dict) + sparse_algorithm: str | None = None + cascade_disabled: bool = False + + @property + def installed_count(self) -> int: + """Number of attention implementations installed.""" + return len(self.installed_layers) + + @property + def transforms_active(self) -> bool: + """Whether the installation activated any ModelOpt attention transform.""" + return bool(self.installed_layers) + + +@dataclass(frozen=True, slots=True) +class _AttentionPlan: + name: str + module: torch.nn.Module + new_impl: Any + sparse_kw: dict[str, Any] + device: object | None + dtype: torch.dtype | None + requires_flashinfer_patch: bool + + +@dataclass(frozen=True, slots=True) +class _InstallPlan: + model_runner: Any + layers: tuple[_AttentionPlan, ...] + quantize: bool + sparse_algorithm: str | None + q_format: str = "nvfp4" + k_format: str = "nvfp4" + p_format: str = "nvfp4" + v_format: str = "nvfp4" + + +def _unwrapped_model(model_runner): + model = model_runner.model + return model.unwrap() if hasattr(model, "unwrap") else model + + +def _model_config(model_runner): + model_config = getattr(model_runner, "model_config", None) + if model_config is not None: + return model_config + return getattr(getattr(model_runner, "vllm_config", None), "model_config", None) + + +def _resolve_sparse_config(model_runner, sparse_cfg) -> tuple[dict | None, str | None]: + if sparse_cfg is None: + return None, None + if isinstance(sparse_cfg, str): + if sparse_cfg != "checkpoint": + raise TypeError("sparse_cfg must be 'checkpoint', a mapping, or None") + detected = load_from_checkpoint_metadata( + getattr(_model_config(model_runner), "hf_config", None) + ) + if detected is None: + return None, None + return detected + if not isinstance(sparse_cfg, Mapping): + raise TypeError("sparse_cfg must be 'checkpoint', a mapping, or None") + return dict(sparse_cfg), "EXPLICIT" + + +def _sparse_kwargs(name: str, sparse_cfg: dict | None) -> dict[str, Any]: + 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 attention_plugin._build_sparse_kw(layer_cfg) + + +def _require_supported_vllm() -> None: + # The adapters rely on vLLM writing the current K/V before impl.forward; + # that external cache-update contract starts in vLLM 0.15. + if version.parse(vllm.__version__) < version.parse("0.15.0"): + raise RuntimeError("ModelOpt vLLM attention transforms require vLLM >= 0.15.0") + + +def _cudagraph_mode(model_runner): + # Imported lazily so missing quant-only APIs do not affect a sparse no-op. + from vllm.config.compilation import CUDAGraphMode + + config = getattr(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 CUDAGraphMode.NONE + + +def _global_errors(model_runner) -> list[str]: + config = getattr(model_runner, "vllm_config", None) + if config is None: + return ["model_runner.vllm_config is required"] + + from vllm.config.compilation import CUDAGraphMode + + parallel = getattr(config, "parallel_config", None) + cache_config = getattr(config, "cache_config", None) + model_config = _model_config(model_runner) + if parallel is None or cache_config is None or model_config is None: + return ["vLLM parallel, cache, and model configs are required"] + + 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(model_runner).mixed_mode() == CUDAGraphMode.FULL: + errors.append("FULL mixed-batch cudagraph mode is unsupported") + if getattr(model_config, "dtype", None) not in (torch.float16, torch.bfloat16): + errors.append("resolved model/KV-cache dtype must be fp16 or bf16") + cache_dtype = getattr(cache_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 _layer_errors(module) -> list[str]: + impl = getattr(module, "impl", None) + errors = [] + original_layout = next( + ( + cls + for cls in type(module).__mro__ + if cls is _VLLM_ATTENTION + or (cls.__module__.startswith("vllm.") and issubclass(cls, _VLLM_ATTENTION)) + ), + None, + ) + if original_layout is not _VLLM_ATTENTION: + errors.append(f"layout {type(module).__name__} is not regular decoder self-attention") + attn_type = getattr(module, "attn_type", None) + if getattr(attn_type, "value", attn_type) != "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 _device_capability_error(device) -> str | None: + if device is None: + return None + device = torch.device(device) + if device.type != "cuda": + return None + major, minor = torch.cuda.get_device_capability(device) + if (major, minor) < (8, 9): + return ( + "NVFP4 attention requires CUDA compute capability >= 8.9 (Ada/Hopper+); " + f"got sm_{major}{minor}" + ) + return None + + +def _sparse_graph_error(sparse_kw: dict[str, Any], mode) -> str | None: + from vllm.config.compilation import CUDAGraphMode + + params = sparse_kw.get("threshold_scale_factor") + if ( + mode.decode_mode() == CUDAGraphMode.FULL + and isinstance(params, dict) + and isinstance(params.get("decode"), dict) + ): + return "calibrated decode skip-softmax requires a non-FULL CUDA graph mode" + return None + + +def _is_modelopt_flashinfer_impl(impl) -> bool: + flashinfer_cls = attention_plugin._FLASHINFER_IMPL_CLS + return flashinfer_cls is not None and isinstance(impl, flashinfer_cls) + + +def _select_new_impl(module) -> tuple[object | None, bool, str | None]: + old_impl = module.impl + try: + if isinstance(old_impl, attention_plugin.ModelOptSparseAttentionImpl): + new_cls = type(old_impl) + requires_flashinfer_patch = False + elif _is_modelopt_flashinfer_impl(old_impl): + new_cls = type(old_impl) + requires_flashinfer_patch = True + elif isinstance(old_impl, FlashAttentionImpl): + new_cls = attention_plugin.ModelOptSparseAttentionImpl + requires_flashinfer_patch = False + else: + try: + flashinfer_module = importlib.import_module("vllm.v1.attention.backends.flashinfer") + except ImportError: + flashinfer_impl_cls = () + else: + flashinfer_impl_cls = flashinfer_module.FlashInferImpl + if not isinstance(old_impl, flashinfer_impl_cls): + return ( + None, + False, + ( + f"backend {type(old_impl).__name__} is not supported; " + "expected FlashAttentionImpl or FlashInferImpl" + ), + ) + new_cls = attention_plugin.get_flashinfer_sparse_impl_cls() + requires_flashinfer_patch = True + return ( + attention_plugin._clone_sparse_impl(old_impl, new_cls), + requires_flashinfer_patch, + None, + ) + except (NotImplementedError, TypeError) as err: + return None, False, str(err) + + +def _load_quant_plugin(): + # Quantization is optional for the sparse-only public entry point. + from modelopt.torch.quantization.plugins import vllm as quant_plugin + + return quant_plugin + + +def _raise_unsupported(errors: list[str], policy: str) -> None: + if errors: + raise NotImplementedError( + f"Unsupported ModelOpt {policy} plan:\n - " + "\n - ".join(errors) + ) + + +def _plan_vllm_attention( + model_runner, + *, + quantize: bool, + sparse_cfg, + q_format: str = "nvfp4", + k_format: str = "nvfp4", + p_format: str = "nvfp4", + v_format: str = "nvfp4", +) -> _InstallPlan: + model = _unwrapped_model(model_runner) + resolved_sparse_cfg, sparse_algorithm = _resolve_sparse_config(model_runner, sparse_cfg) + candidates = [] + attention_count = 0 + for name, module in model.named_modules(): + if not isinstance(module, _VLLM_ATTENTION): + continue + attention_count += 1 + sparse_kw = _sparse_kwargs(name, resolved_sparse_cfg) + if quantize or sparse_kw: + candidates.append((name, module, sparse_kw)) + + if not candidates and not quantize: + return _InstallPlan( + model_runner, (), False, sparse_algorithm, q_format, k_format, p_format, v_format + ) + + _require_supported_vllm() + errors = _global_errors(model_runner) if quantize else [] + mode = _cudagraph_mode(model_runner) if quantize else None + quant_plugin: Any = _load_quant_plugin() if quantize else None + plans = [] + for name, module, sparse_kw in candidates: + reasons = _layer_errors(module) + device = dtype = None + if quantize: + device, dtype = quant_plugin._get_device_dtype(module) + model_dtype = getattr(_model_config(model_runner), "dtype", None) + if model_dtype in (torch.float16, torch.bfloat16): + dtype = model_dtype + if device is None or dtype is None: + reasons.append("device/dtype could not be resolved") + elif dtype not in (torch.float16, torch.bfloat16): + reasons.append(f"resolved dtype {dtype} must be fp16 or bf16") + if capability_error := _device_capability_error(device): + reasons.append(capability_error) + if quantize: + if graph_error := _sparse_graph_error(sparse_kw, mode): + reasons.append(graph_error) + new_impl, requires_flashinfer_patch, backend_error = _select_new_impl(module) + if backend_error: + reasons.append(backend_error) + if reasons: + errors.extend(f"{name or ''}: {reason}" for reason in reasons) + continue + plans.append( + _AttentionPlan( + name, + module, + new_impl, + sparse_kw, + device, + dtype, + requires_flashinfer_patch, + ) + ) + if quantize and attention_count == 0: + errors.append("no regular attention layers were found") + _raise_unsupported(errors, "NVFP4 attention" if quantize else "sparse attention") + return _InstallPlan( + model_runner, + tuple(plans), + quantize, + sparse_algorithm, + q_format, + k_format, + p_format, + v_format, + ) + + +def _build_report(plan: _InstallPlan) -> VllmAttentionInstallReport: + names = tuple(layer.name for layer in plan.layers) + sparse_names = tuple(layer.name for layer in plan.layers if layer.sparse_kw) + return VllmAttentionInstallReport( + installed_layers=names, + quantized_layers=names if plan.quantize else (), + sparse_layers=sparse_names, + backend_counts=dict(Counter(type(layer.new_impl).__name__ for layer in plan.layers)), + sparse_algorithm=plan.sparse_algorithm, + cascade_disabled=bool(plan.layers), + ) + + +def _apply_vllm_attention_plans(plan: _InstallPlan) -> VllmAttentionInstallReport: + if not plan.layers: + return _build_report(plan) + + if any(layer.requires_flashinfer_patch for layer in plan.layers): + if not attention_plugin.patch_flashinfer_metadata_builder(): + raise RuntimeError("Unable to prepare FlashInfer metadata for ModelOpt attention") + + quant_plugin: Any = _load_quant_plugin() if plan.quantize else None + # Compatibility validation is all-layers-first. From here, any exception is + # fatal; publish each layer immediately after configuration so earlier layers + # are never left configured on their native implementation. + plan.model_runner.cascade_attn_enabled = False + for layer in plan.layers: + layer.new_impl.sparse_kw = layer.sparse_kw + if plan.quantize: + # Pass cfg only for non-default formats: keeps the default call + # signature stable for callers/fakes that predate the cfg parameter. + _cfg_kwargs = ( + { + "cfg": quant_plugin.build_vllm_attention_quant_cfg( + q_format=plan.q_format, + k_format=plan.k_format, + p_format=plan.p_format, + v_format=plan.v_format, + ) + } + if (plan.q_format, plan.k_format, plan.p_format, plan.v_format) + != ("nvfp4", "nvfp4", "nvfp4", "nvfp4") + else {} + ) + converted = quant_plugin.configure_vllm_nvfp4_attention_quantizers( + layer.module, + device=layer.device, + dtype=layer.dtype, + **_cfg_kwargs, + ) + if converted is not None and converted is not layer.module: + raise RuntimeError("vLLM attention quantization must convert modules in place") + p_qdq, p_qdq_amax = attention_plugin._p_qdq_from_layer(layer.module) + v_qdq, v_qdq_amax = attention_plugin._v_qdq_from_layer(layer.module) + if plan.v_format == "fp8": + # Per-tensor FP8 V is quantized module-level BEFORE the cache + # write (each token is self-contained; no block geometry), so + # the kernel sees pre-QDQ'd V and needs no V transform at all. + v_qdq, v_qdq_amax = None, None + layer.new_impl.quant_kw = { + "p_qdq": p_qdq, + "p_qdq_amax": p_qdq_amax, + "v_qdq": v_qdq, + "v_qdq_amax": v_qdq_amax, + } + + missing = object() + old_query_flag = getattr(layer.module, "_query_quant_in_kernel", missing) + old_value_flag = getattr(layer.module, "_value_quant_in_kernel", missing) + if plan.quantize: + # fp8 Q is module-level (bf16 losslessly carries E4M3 QDQ values); + # the kernel then runs a plain bf16 BMM1 with no Q transform. + layer.module._query_quant_in_kernel = plan.q_format != "fp8" + layer.module._value_quant_in_kernel = plan.v_format != "fp8" + try: + # Publish the adapter last so a native impl never runs with in-kernel + # quantization flags that only the ModelOpt adapter understands. + layer.module.impl = layer.new_impl + except Exception: + if plan.quantize: + for name, value in ( + ("_query_quant_in_kernel", old_query_flag), + ("_value_quant_in_kernel", old_value_flag), + ): + if value is missing: + delattr(layer.module, name) + else: + setattr(layer.module, name, value) + raise + return _build_report(plan) + + +def install_vllm_sparse_attention_from_checkpoint( + model_runner, +) -> VllmAttentionInstallReport: + """Install checkpoint-configured sparse attention into a loaded vLLM model. + + Missing or inactive ``sparse_attention_config`` metadata is a no-op. All + known compatibility errors are validated before any attention module is + changed. + """ + return _apply_vllm_attention_plans( + _plan_vllm_attention(model_runner, quantize=False, sparse_cfg="checkpoint") + ) + + +def install_vllm_nvfp4_attention( + model_runner, + *, + sparse_cfg="checkpoint", + q_format: str = "nvfp4", + k_format: str = "nvfp4", + p_format: str = "nvfp4", + v_format: str = "nvfp4", +) -> VllmAttentionInstallReport: + """Install fixed NVFP4 attention with optional checkpoint sparsity. + + Args: + model_runner: A loaded vLLM model runner. + sparse_cfg: ``"checkpoint"`` to consume optional exported metadata, a + resolved sparse config dict, or ``None`` for NVFP4-only attention. + """ + for name, fmt in ( + ("q_format", q_format), + ("k_format", k_format), + ("p_format", p_format), + ("v_format", v_format), + ): + if fmt not in ("nvfp4", "fp8"): + raise ValueError(f"{name} must be 'nvfp4' or 'fp8', got {fmt!r}") + return _apply_vllm_attention_plans( + _plan_vllm_attention( + model_runner, + quantize=True, + sparse_cfg=sparse_cfg, + q_format=q_format, + k_format=k_format, + p_format=p_format, + v_format=v_format, + ) + ) 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..a018074abcb --- /dev/null +++ b/tests/gpu/torch/kernels/common/attention/test_decode_attention.py @@ -0,0 +1,248 @@ +# 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) + + +@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 +def test_p_qdq_matches_fixed_split_local_oracle(): + """The production 32-split schedule quantizes split-local unnormalized P.""" + seq_len, head_dim, num_splits = 4096, 16, 32 + scale = head_dim**-0.5 + q = torch.zeros(1, 1, head_dim, device="cuda", dtype=torch.float32) + q[..., 0] = 1.0 + k = torch.zeros(1, 1, seq_len, head_dim, device="cuda", dtype=torch.bfloat16) + k[..., 0] = torch.linspace(-8.0, 8.0, seq_len, device="cuda", dtype=torch.bfloat16) + torch.manual_seed(19) + v = torch.randn_like(k) + k_cache, v_cache, block_table = _paged_cache(k, v, (seq_len,)) + seq_lens = torch.tensor([seq_len], device="cuda", dtype=torch.int32) + + output = attention_decode( + q, + k_cache, + v_cache, + block_table, + seq_lens, + softmax_scale=scale, + num_kv_splits=num_splits, + p_qdq="nvfp4", + ) + + scores = torch.matmul(k[0, 0].float(), q[0, 0]) * (scale * 1.4426950408889634) + split_scores = scores.reshape(num_splits, -1) + split_max = split_scores.amax(dim=1) + p = torch.exp2(split_scores - split_max[:, None]) + p_qdq = _nvfp4_qdq_reference(p.to(torch.bfloat16).float()) + split_acc = torch.einsum("sk,skd->sd", p_qdq, v[0, 0].float().reshape(num_splits, -1, head_dim)) + running_max = torch.tensor(-float("inf"), device="cuda") + running_sum = torch.tensor(0.0, device="cuda") + acc = torch.zeros(head_dim, device="cuda") + for split_idx in range(num_splits): + new_max = torch.maximum(running_max, split_max[split_idx]) + correction = torch.exp2(running_max - new_max) + split_correction = torch.exp2(split_max[split_idx] - new_max) + acc = acc * correction + split_acc[split_idx] * split_correction + running_sum = running_sum * correction + p[split_idx].sum() * split_correction + running_max = new_max + reference = acc / running_sum + + torch.testing.assert_close(output[0, 0], reference, rtol=5e-5, atol=5e-5) diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py index f0663dc6b44..60523d44457 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -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,155 @@ from modelopt.torch.quantization.tensor_quant import fp8_eager if TRITON_KERNEL_AVAILABLE: - from modelopt.torch.kernels.common.attention import attention + from modelopt.torch.kernels.common.attention import attention, triton_fa from modelopt.torch.kernels.common.attention.triton_fa import LOG2E -# 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 +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") +def test_v_qdq_rejects_autograd_before_kernel_launch(monkeypatch): + apply = Mock(side_effect=AssertionError("_Attention.apply reached")) + monkeypatch.setattr(triton_fa._Attention, "apply", apply) + q, k, v = (torch.zeros(1, 1, 16, requires_grad=True) for _ in range(3)) + locs = torch.zeros(1, dtype=torch.int32) + lens = torch.ones(1, dtype=torch.int32) + with pytest.raises(NotImplementedError, match=r"v_qdq.*autograd"): + attention(q, k, v, locs, lens, 1, v_qdq="nvfp4") + apply.assert_not_called() + + +# 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,9 +210,11 @@ def _qdq_nvfp4(p, global_scale=1.0): shape = p.shape g = p.reshape(*shape[:-1], shape[-1] // 16, 16) block_amax = g.amax(dim=-1, keepdim=True) # p >= 0, so max == amax - 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) @@ -80,13 +225,15 @@ def _apply_qdq(p, mode, qdq_scale=1.0): return _qdq_nvfp4(p, qdq_scale) -def qdq_attention_reference(q, k, v, scale, mode, is_causal=True, amax=1.0): +def qdq_attention_reference(q, k, v, scale, mode, is_causal=True, amax=1.0, block_n=P_QDQ_BLOCK_N): """Tile-looped online-softmax reference replicating kernel P_QDQ semantics. Single sequence: q [s, h, d], k/v [s_kv, h_kv, d] (fp16). Walks KV tiles - of BLOCK_N exactly like the kernel, keeps the softmax denominator + of ``block_n`` exactly like the kernel, keeps the softmax denominator unquantized, applies qdq to the unnormalized p of each tile, and mirrors - the kernel's ``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 +260,20 @@ def qdq_attention_reference(q, k, v, scale, mode, is_causal=True, amax=1.0): row_max = torch.full((h, s), float("-inf"), device=q.device) row_sum = torch.zeros(h, s, device=q.device) acc = torch.zeros(h, s, d, device=q.device) - for start in range(0, s_kv, BLOCK_N): - tile = t[:, :, start : start + BLOCK_N] + for start in range(0, s_kv, block_n): + tile = t[:, :, start : start + block_n] m_new = torch.maximum(row_max, tile.amax(dim=-1)) p = torch.exp2(tile - m_new[..., None]) l_new = p.sum(dim=-1) corr = torch.exp2(row_max - m_new) row_sum = row_sum * corr + l_new acc = acc * corr[..., None] + if mode == "nvfp4": + p = p.to(v.dtype).float() p = _apply_qdq(p, mode, qdq_scale) - # Kernel casts p to v.dtype for the BMM2 dot - p = p.to(v.dtype).float() - acc = acc + torch.einsum("hqk,khd->hqd", p, vv[start : start + BLOCK_N].float()) + if mode == "fp8": + p = p.to(v.dtype).float() + acc = acc + torch.einsum("hqk,khd->hqd", p, vv[start : start + block_n].float()) row_max = m_new out = acc / row_sum[..., None] return out.permute(1, 0, 2) @@ -135,9 +284,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 +296,63 @@ 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_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 +366,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 +420,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 +454,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 +484,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..2f500b1377c 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,189 @@ 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, + build_vllm_attention_quant_cfg, + configure_vllm_nvfp4_attention_quantizers, 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_configure_vllm_nvfp4_attention_quantizers_is_attention_scoped(monkeypatch): + monkeypatch.setattr( + vllm_plugin, + "create_parallel_state", + lambda: vllm_plugin.ParallelState(data_parallel_group=None), + ) + attention = object.__new__(vllm_plugin.vllm_attention.Attention) + torch.nn.Module.__init__(attention) + linear = torch.nn.Linear(4, 4) + attention.unrelated_linear = linear + original_linear_type = type(linear) + + converted = configure_vllm_nvfp4_attention_quantizers( + attention, device="cpu", dtype=torch.bfloat16 + ) + + assert converted is attention + assert isinstance(converted, _QuantVLLMAttention) + assert converted.device == torch.device("cpu") + assert converted.dtype == torch.bfloat16 + assert type(linear) is original_linear_type + for name in ("q", "k", "p", "v"): + quantizer = getattr(converted, f"{name}_bmm_quantizer") + assert quantizer.is_enabled + assert quantizer.is_nvfp4_dynamic + assert quantizer.block_sizes[-1] == 16 + assert not hasattr(converted.q_bmm_quantizer, "_amax") + assert not hasattr(converted.p_bmm_quantizer, "_amax") + assert converted.k_bmm_quantizer._amax == 6.0 * 448.0 + assert converted.v_bmm_quantizer._amax == 6.0 * 448.0 + assert not hasattr(converted, "_query_quant_in_kernel") + assert not hasattr(converted, "_value_quant_in_kernel") + + +def test_configure_vllm_nvfp4_attention_quantizers_preserves_and_moves_amax(monkeypatch): + monkeypatch.setattr( + vllm_plugin, + "create_parallel_state", + lambda: vllm_plugin.ParallelState(data_parallel_group=None), + ) + attention = object.__new__(vllm_plugin.vllm_attention.Attention) + torch.nn.Module.__init__(attention) + converted = configure_vllm_nvfp4_attention_quantizers( + attention, device="cpu", dtype=torch.float16 + ) + for name, value in zip(("q", "k", "p", "v"), (13.0, 17.0, 23.0, 19.0), strict=True): + getattr(converted, f"{name}_bmm_quantizer").amax = torch.tensor(value) + reconfigured = configure_vllm_nvfp4_attention_quantizers( + converted, device="cpu", dtype=torch.float16 + ) + + assert reconfigured is converted + assert converted.q_bmm_quantizer._amax == 13.0 + assert converted.k_bmm_quantizer._amax == 17.0 + assert converted.p_bmm_quantizer._amax == 23.0 + assert converted.v_bmm_quantizer._amax == 19.0 + + configure_vllm_nvfp4_attention_quantizers(converted, device="meta", dtype=torch.float16) + for name in ("q", "k", "p", "v"): + assert getattr(converted, f"{name}_bmm_quantizer")._amax.device.type == "meta" + + +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``. @@ -272,3 +448,43 @@ def test_tiny_deepseek_mla_quantize(tiny_deepseek_llm): assert summary["moe_count"] >= 2, summary _assert_quantizer_amax_is_static(summary) + + +def test_configure_vllm_attention_quantizers_fp8_bmm2(monkeypatch): + monkeypatch.setattr( + vllm_plugin, + "create_parallel_state", + lambda: vllm_plugin.ParallelState(data_parallel_group=None), + ) + attention = object.__new__(vllm_plugin.vllm_attention.Attention) + torch.nn.Module.__init__(attention) + + converted = configure_vllm_nvfp4_attention_quantizers( + attention, + device="cpu", + dtype=torch.bfloat16, + cfg=build_vllm_attention_quant_cfg(p_format="fp8", v_format="fp8"), + ) + + # BMM1 unchanged: Q/K dynamic block-16 NVFP4 (F1) + for name in ("q", "k"): + quantizer = getattr(converted, f"{name}_bmm_quantizer") + assert quantizer.is_enabled and quantizer.is_nvfp4_dynamic + assert quantizer.block_sizes[-1] == 16 + assert converted.k_bmm_quantizer._amax == 6.0 * 448.0 + # BMM2: P/V per-tensor FP8 E4M3 with fixed amax (P=1.0, V=448) (F3) + for name, amax in (("p", 1.0), ("v", 448.0)): + quantizer = getattr(converted, f"{name}_bmm_quantizer") + assert quantizer.is_enabled + assert quantizer.num_bits == (4, 3) + assert not quantizer.block_sizes + assert float(quantizer._amax) == amax + # idempotent: calibrated amax survives reconfiguration + converted.v_bmm_quantizer.amax = torch.tensor(96.0) + reconfigured = configure_vllm_nvfp4_attention_quantizers( + converted, + device="cpu", + dtype=torch.bfloat16, + cfg=build_vllm_attention_quant_cfg(p_format="fp8", v_format="fp8"), + ) + assert float(reconfigured.v_bmm_quantizer._amax) == 96.0 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..f82746f0e35 --- /dev/null +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py @@ -0,0 +1,130 @@ +# 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 quant+sparse vLLM worker lifecycle.""" + +import importlib.util +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from vllm.v1.worker.gpu_worker import Worker as BaseWorker + +from modelopt.torch.quantization.plugins import vllm as quant_plugin + +_WORKER_PATH = Path(__file__).parents[5] / "examples/vllm_serve/sparse_attn_worker.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 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)) + + monkeypatch.setattr( + torch, + "inference_mode", + lambda: recorded_context("inference"), + ) + monkeypatch.setattr( + quant_plugin, + "disable_compilation", + lambda actual_model: recorded_context("compilation", actual_model), + ) + + instance = object.__new__(worker_module.QuantSparseAttnWorker) + instance.model_runner = SimpleNamespace(model=SimpleNamespace(unwrap=lambda: model)) + + def profile(actual_worker): + events.append(("profile", actual_worker)) + return 73 + + monkeypatch.setattr(BaseWorker, "determine_available_memory", profile) + + assert instance.determine_available_memory() == 73 + assert events == [ + ("enter", "inference", None), + ("enter", "compilation", model), + ("profile", instance), + ("exit", "compilation", model), + ("exit", "inference", None), + ] + + +def _quant_worker_with_config(monkeypatch, calls, additional_config): + module = _load_worker_module() + monkeypatch.setattr(BaseWorker, "load_model", lambda *_a, **_k: calls.append("base")) + monkeypatch.setattr( + module, + "install_vllm_nvfp4_attention", + lambda runner, **kw: ( + calls.append(("install", kw)) + or SimpleNamespace(installed_count=0, sparse_algorithm=None, backend_counts={}) + ), + ) + instance = object.__new__(module.QuantSparseAttnWorker) + instance.model_runner = SimpleNamespace() + instance.vllm_config = SimpleNamespace(additional_config=additional_config) + return instance + + +def test_quant_worker_defaults_to_nvfp4(monkeypatch): + calls = [] + instance = _quant_worker_with_config(monkeypatch, calls, None) + instance.load_model() + assert calls[0] == "base" + assert calls[1][1] == {"sparse_cfg": "checkpoint"} + + +def test_quant_worker_reads_formats_from_additional_config(monkeypatch): + calls = [] + instance = _quant_worker_with_config( + monkeypatch, + calls, + {"modelopt_attn_quant": {"p_format": "fp8", "v_format": "fp8"}}, + ) + instance.load_model() + assert calls[0] == "base" + assert calls[1][1] == {"sparse_cfg": "checkpoint", "p_format": "fp8", "v_format": "fp8"} + + +def test_quant_worker_rejects_unknown_format_keys(monkeypatch): + calls = [] + instance = _quant_worker_with_config( + monkeypatch, calls, {"modelopt_attn_quant": {"o_format": "fp8"}} + ) + with pytest.raises(ValueError, match="o_format"): + instance.load_model() + assert ("install", {"sparse_cfg": "checkpoint"}) not in calls 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..b44d452be02 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -15,19 +15,117 @@ """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 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 test_shared_worker_import_does_not_resolve_quant_only_apis(monkeypatch): + forbidden = { + "vllm.config.compilation", + "vllm.v1.attention.backend", + "modelopt.torch.quantization.conversion", + "modelopt.torch.quantization.nn", + "modelopt.torch.quantization.plugins.vllm", + } + real_import = builtins.__import__ + + def guarded_import(name, *args, **kwargs): + fromlist = kwargs.get("fromlist", args[2] if len(args) > 2 else ()) + requested = {name, *(f"{name}.{item}" for item in fromlist or ())} + if blocked := requested & forbidden: + raise AssertionError( + f"quant-only import during sparse module load: {sorted(blocked)[0]}" + ) + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(vllm, "__version__", "0.9.0") + monkeypatch.setattr(builtins, "__import__", guarded_import) + worker_module = _load_worker_module("sparse_attn_worker_import_test") + + assert worker_module.__all__ == ["SparseAttnWorker", "QuantSparseAttnWorker"] + + +@pytest.mark.parametrize( + ("class_name", "installer_name", "expected_kwargs"), + [ + ("SparseAttnWorker", "install_vllm_sparse_attention_from_checkpoint", {}), + ( + "QuantSparseAttnWorker", + "install_vllm_nvfp4_attention", + {"sparse_cfg": "checkpoint"}, + ), + ], ) +def test_public_workers_install_after_base_load( + monkeypatch, capsys, class_name, installer_name, expected_kwargs +): + worker_module = _load_worker_module(f"worker_order_{class_name}") + events = [] + model_runner = object() + + def fake_base_load(worker, *_args, **_kwargs): + events.append("base") + worker.model_runner = model_runner + worker.vllm_config = SimpleNamespace(additional_config=None) + + def fake_install(actual_runner, **kwargs): + events.append(("install", actual_runner, kwargs)) + return SimpleNamespace( + installed_count=0 if class_name == "SparseAttnWorker" else 1, + backend_counts={"TestImpl": 1}, + sparse_algorithm=None, + ) + + monkeypatch.setattr(BaseWorker, "load_model", fake_base_load) + monkeypatch.setattr(worker_module, installer_name, fake_install) + + instance = object.__new__(getattr(worker_module, class_name)) + assert instance.load_model() is None + assert events == ["base", ("install", model_runner, expected_kwargs)] + output = capsys.readouterr().out + if class_name == "SparseAttnWorker": + assert "No sparse_attention_config found" in output + assert "hf_sa.py" in output + else: + assert "Installed NVFP4 attention (quant+sparse) on 1 layers: {'TestImpl': 1}" in output def _make_old_impl(): @@ -66,6 +164,588 @@ def test_clone_sparse_impl_rejects_non_none_sinks(): _clone_sparse_impl(old_impl) +def _make_old_flashinfer_impl(): + """Create a bare FlashInfer impl without requiring a live vLLM config.""" + impl = object.__new__(FlashInferImpl) + impl.__dict__.update( + num_heads=2, + num_kv_heads=2, + head_size=64, + scale=0.125, + sinks=None, + alibi_slopes=None, + logits_soft_cap=None, + kv_cache_dtype="auto", + ) + return impl + + +def _make_flashinfer_impl(*, sparse=False, quantized=False): + impl = _clone_sparse_impl(_make_old_flashinfer_impl(), get_flashinfer_sparse_impl_cls()) + impl.sparse_kw = {"sparsity_n": 2, "sparsity_m": 4} if sparse else {} + impl.quant_kw = { + "p_qdq": "nvfp4" if quantized else None, + "p_qdq_amax": 1.0, + "v_qdq": "nvfp4" if quantized else None, + "v_qdq_amax": 6.0 * 448.0 if quantized else None, + } + return impl + + +@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)] + + +def test_flashinfer_q_only_transform_does_not_fallback(monkeypatch): + """Withheld Q QDQ must run even when sparse/P/V transforms are inactive.""" + impl = _make_flashinfer_impl() + query = torch.zeros(4, 2, 64, dtype=torch.float16) + kv_cache = torch.zeros(4, 2, 16, 2, 64, dtype=torch.float16) + metadata = _flashinfer_metadata(query_lens=(2, 2), max_query_len=4) + captured = {} + layer = SimpleNamespace( + _query_quant_in_kernel=True, + q_bmm_quantizer=lambda value: value + 1, + ) + + monkeypatch.setattr( + FlashInferImpl, + "forward", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("native fallback")), + ) + + def fake_attention(query, **kwargs): + captured["query"] = query + return torch.zeros_like(query) + + monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) + + impl.forward(layer, query, query, query, kv_cache, metadata, output=torch.empty_like(query)) + + assert torch.all(captured["query"] == 1) + + +def test_flashinfer_legacy_inactive_launch_writes_only_in_native_fallback(monkeypatch): + impl = _make_flashinfer_impl(sparse=True) + impl.kv_sharing_target_layer_name = None + query = torch.zeros(1, 2, 64, dtype=torch.float16) + kv_cache = torch.zeros(1, 2, 16, 2, 64, dtype=torch.float16) + metadata = _flashinfer_metadata(query_lens=(1,), seq_lens=(16,)) + manual_writes = [] + native_calls = [] + + monkeypatch.setattr(FlashInferBackend, "forward_includes_kv_cache_update", True) + monkeypatch.setattr( + vllm_plugin, "_flashinfer_cache_write", lambda *args: manual_writes.append(args) + ) + monkeypatch.setattr( + FlashInferImpl, + "forward", + lambda *args, **kwargs: native_calls.append(True) or kwargs.get("output", args[7]), + ) + + impl.forward( + SimpleNamespace(), query, query, query, kv_cache, metadata, output=torch.empty_like(query) + ) + + assert manual_writes == [] + assert native_calls == [True] + + +@pytest.mark.parametrize("quantized", [False, True], ids=["sparse-only", "quantized"]) +def test_flashinfer_mixed_batch_splits_decode_and_prefill(monkeypatch, quantized): + prefill_tokens = 17 + impl = _make_flashinfer_impl(sparse=True, quantized=quantized) + query = torch.zeros(1 + prefill_tokens, 2, 64, dtype=torch.float16) + kv_cache = torch.zeros(4, 2, 16, 2, 64, dtype=torch.float16) + metadata = _flashinfer_metadata(query_lens=(1, prefill_tokens), mixed=True) + layer = SimpleNamespace( + _query_quant_in_kernel=quantized, + q_bmm_quantizer=lambda value: value, + ) + calls = {"native": [], "decode": [], "prefill": [], "finalize": []} + + def fake_decode(query, *args, **kwargs): + calls["decode"].append((query, kwargs)) + return torch.zeros_like(query) + + def fake_attention(query, **kwargs): + calls["prefill"].append((query, kwargs)) + return torch.zeros_like(query) + + monkeypatch.setattr( + FlashInferImpl, + "forward", + lambda *args, **kwargs: calls["native"].append(True) or args[7], + ) + monkeypatch.setattr( + vllm_plugin, + "fake_quant_v_onwrite", + lambda *args, **kwargs: calls["finalize"].append(kwargs), + ) + monkeypatch.setattr(vllm_plugin, "triton_decode_attention", fake_decode) + monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) + + impl.forward(layer, query, query, query, kv_cache, metadata, output=torch.empty_like(query)) + + assert len(calls["prefill"]) == 1 + prefill_query, prefill_kw = calls["prefill"][0] + assert prefill_query.shape[0] == prefill_tokens + assert prefill_kw["sparsity_n"] == 2 + assert prefill_kw["sparsity_m"] == 4 + torch.testing.assert_close( + prefill_kw["b_seq_len"], torch.tensor([prefill_tokens], dtype=torch.int32) + ) + + if quantized: + assert calls["native"] == [] + assert len(calls["decode"]) == 1 + decode_query, decode_kw = calls["decode"][0] + assert decode_query.shape[0] == 1 + assert decode_kw["p_qdq"] == "nvfp4" + assert "sparsity_n" not in decode_kw + assert prefill_kw["p_qdq"] == "nvfp4" + assert [kw["max_new_tokens"] for kw in calls["finalize"]] == [1, prefill_tokens] + else: + assert calls["native"] == [True] + assert calls["decode"] == [] + assert calls["finalize"] == [] + + +def _make_flash_attention_impl(*, sparse=False, quantized=False): + impl = _clone_sparse_impl(_make_old_impl()) + impl.sparse_kw = {"sparsity_n": 2, "sparsity_m": 4} if sparse else {} + impl.quant_kw = { + "p_qdq": "nvfp4" if quantized else None, + "p_qdq_amax": 1.0, + "v_qdq": "nvfp4" if quantized else None, + "v_qdq_amax": 6.0 * 448.0 if quantized else None, + } + return impl + + +def _flash_attention_mixed_metadata(decode_len=1, prefill_len=17): + query_lens = (decode_len, prefill_len) + seq_lens = (16, 34) + return SimpleNamespace( + use_cascade=False, + num_actual_tokens=sum(query_lens), + max_query_len=max(query_lens), + max_seq_len=max(seq_lens), + query_start_loc=torch.tensor(list(accumulate(query_lens, initial=0)), dtype=torch.int32), + seq_lens=torch.tensor(seq_lens, dtype=torch.int32), + block_table=torch.zeros(len(seq_lens), 3, dtype=torch.int32), + causal=True, + num_decodes=1, + num_prefills=1, + num_decode_tokens=decode_len, + num_prefill_tokens=prefill_len, + ) + + +@pytest.mark.parametrize("quantized", [False, True], ids=["sparse-only", "quantized"]) +def test_flash_attention_mixed_batch_splits_decode_and_prefill(monkeypatch, quantized): + """FlashAttention must split mixed decode+prefill batches so decode rows take + the decode schedule -- NVFP4 P-QDQ is schedule-sensitive, so a decode result + must not depend on a co-scheduled prefill (matches the FlashInfer adapter).""" + prefill_tokens = 17 + impl = _make_flash_attention_impl(sparse=True, quantized=quantized) + query = torch.zeros(1 + prefill_tokens, 2, 64, dtype=torch.float16) + kv_cache = torch.zeros(2, 4, 16, 2, 64, dtype=torch.float16) + metadata = _flash_attention_mixed_metadata(decode_len=1, prefill_len=prefill_tokens) + 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.full_like(query, 5) + + def fake_attention(query, **kwargs): + calls["prefill"].append((query, kwargs)) + return torch.full_like(query, 7) + + def fake_native(*args, **kwargs): + calls["native"].append(True) + output = kwargs["output"] if "output" in kwargs else args[7] + return output.fill_(3) + + monkeypatch.setattr( + FlashAttentionImpl, + "forward", + fake_native, + ) + 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) + + result = impl.forward( + layer, query, query, query, kv_cache, metadata, output=torch.empty_like(query) + ) + + assert torch.all(result[1:] == 7) + assert torch.all(result[:1] == (5 if quantized else 3)) + + 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 + + if quantized: + assert calls["native"] == [] + assert len(calls["decode"]) == 1 + decode_query, decode_kw = calls["decode"][0] + assert decode_query.shape[0] == 1 + assert decode_kw["p_qdq"] == "nvfp4" + assert "sparsity_n" not in decode_kw + else: + assert calls["native"] == [True] + assert calls["decode"] == [] + + +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="Mixed-batch token counts do not match common metadata", + ): + impl.forward( + SimpleNamespace(), + query, + query, + query, + kv_cache, + metadata, + output=torch.empty_like(query), + ) + + +@pytest.mark.parametrize( + ("failure", "active", "expected"), + [ + ( + "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) + + +@pytest.mark.parametrize("quantized", [False, True], ids=["sparse-only", "quantized"]) +def test_flashinfer_cascade_falls_back_for_sparse_only_but_rejects_quant(monkeypatch, quantized): + """Cascade is unimplemented by the kernel. A sparse-only transform can safely + delegate to the native dense path; quantization must not be silently dropped + (it would change numerics), so it is rejected instead.""" + impl = _make_flashinfer_impl(sparse=True, quantized=quantized) + metadata = SimpleNamespace(use_cascade=True) + query = torch.zeros(1, 2, 64, dtype=torch.float16) + output = torch.empty_like(query) + native_calls = [] + + def fake_forward(self, *args, **kwargs): + native_calls.append(True) + out = args[6] + out.fill_(9) + return out + + monkeypatch.setattr(FlashInferImpl, "forward", fake_forward) + + if quantized: + with pytest.raises(NotImplementedError) as exc: + impl.forward(None, query, query, query, torch.empty(0), metadata, output=output) + assert str(exc.value) == ( + "vLLM cascade attention is incompatible with active ModelOpt attention quantization" + ) + assert native_calls == [] + else: + assert ( + impl.forward(None, query, query, query, torch.empty(0), metadata, output=output) + is output + ) + assert native_calls == [True] + 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 +797,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 +808,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 +831,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 +893,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 +917,149 @@ 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, + } + 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["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" + + def test_resolve_calibrated_skip_softmax_threshold_for_decode(): """Calibration conversion is phase-aware even when decode later delegates.""" sparse_kw = { @@ -312,66 +1122,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 +1130,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/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py new file mode 100644 index 00000000000..b3ceef30eab --- /dev/null +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py @@ -0,0 +1,370 @@ +# 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. + +"""Tests for the reusable vLLM attention runtime installer.""" + +from types import SimpleNamespace + +import pytest +import torch +import vllm +from torch import nn +from vllm.config.compilation import CUDAGraphMode +from vllm.model_executor.layers.attention import CrossAttention +from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl +from vllm.v1.attention.backends.flashinfer import FlashInferImpl + +from modelopt.torch.quantization.plugins import vllm as quant_plugin +from modelopt.torch.sparsity.attention_sparsity.plugins import vllm_runtime +from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( + ModelOptSparseAttentionImpl, + get_flashinfer_sparse_impl_cls, +) + + +def _bare_attention(impl_cls=FlashAttentionImpl, module_cls=None): + module = object.__new__(module_cls or vllm_runtime._VLLM_ATTENTION) + nn.Module.__init__(module) + module.attn_type = "decoder" + module.head_size = 64 + module.device = torch.device("cpu") + module.dtype = torch.float16 + module.impl = object.__new__(impl_cls) + module.impl.sinks = None + return module + + +def _sparse_metadata(): + return { + "config_groups": { + "group_0": { + "algorithm": "sparse_softmax", + "sparsity_n": 2, + "sparsity_m": 4, + } + } + } + + +def _model_runner(model, *, sparse_metadata=None): + hf_config = SimpleNamespace(sparse_attention_config=sparse_metadata) + model_config = SimpleNamespace(hf_config=hf_config, dtype=torch.float16) + return SimpleNamespace( + model=model, + model_config=model_config, + cascade_attn_enabled=True, + vllm_config=SimpleNamespace( + model_config=model_config, + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + enable_dbo=False, + use_ubatching=False, + ), + cache_config=SimpleNamespace(enable_prefix_caching=False, cache_dtype="auto"), + compilation_config=SimpleNamespace(cudagraph_mode=CUDAGraphMode.NONE), + kv_transfer_config=None, + speculative_config=None, + ), + ) + + +def test_sparse_install_from_checkpoint_is_validation_atomic(): + valid = _bare_attention() + invalid = _bare_attention() + invalid.sliding_window = (128, 128) + valid_impl = valid.impl + invalid_impl = invalid.impl + runner = _model_runner( + nn.ModuleDict({"valid_attn": valid, "invalid_attn": invalid}), + sparse_metadata=_sparse_metadata(), + ) + del runner.vllm_config + + with pytest.raises(NotImplementedError, match="sliding_window"): + vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + + assert valid.impl is valid_impl + assert invalid.impl is invalid_impl + assert runner.cascade_attn_enabled is True + + +def test_installer_rejects_cross_attention_layout_even_if_marked_decoder(): + attention = _bare_attention(module_cls=CrossAttention) + original_impl = attention.impl + runner = _model_runner( + nn.ModuleDict({"cross_attn": attention}), + sparse_metadata=_sparse_metadata(), + ) + del runner.vllm_config + + with pytest.raises(NotImplementedError, match="layout CrossAttention"): + vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + + assert attention.impl is original_impl + + +@pytest.mark.parametrize("impl_cls", [FlashAttentionImpl, FlashInferImpl]) +def test_sparse_install_uses_checkpoint_metadata(monkeypatch, impl_cls): + attention = _bare_attention(impl_cls) + runner = _model_runner( + nn.ModuleDict({"attn": attention}), + sparse_metadata=_sparse_metadata(), + ) + del runner.vllm_config + monkeypatch.setattr( + vllm_runtime.attention_plugin, "patch_flashinfer_metadata_builder", lambda: True + ) + + report = vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + + expected_impl_cls = ( + ModelOptSparseAttentionImpl + if impl_cls is FlashAttentionImpl + else get_flashinfer_sparse_impl_cls() + ) + assert type(attention.impl) is expected_impl_cls + assert attention.impl.sparse_kw["sparsity_n"] == 2 + assert report.installed_layers == ("attn",) + assert report.sparse_layers == ("attn",) + assert report.backend_counts == {expected_impl_cls.__name__: 1} + assert runner.cascade_attn_enabled is False + + +def test_sparse_install_is_noop_without_checkpoint_metadata(monkeypatch): + attention = _bare_attention() + original_impl = attention.impl + runner = _model_runner(nn.ModuleDict({"attn": attention})) + monkeypatch.setattr(vllm, "__version__", "0.14.0") + + report = vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + + assert report.installed_count == 0 + assert not report.transforms_active + assert attention.impl is original_impl + assert runner.cascade_attn_enabled is True + + +@pytest.mark.parametrize("quantize", [False, True]) +def test_active_install_rejects_old_vllm_before_mutation(monkeypatch, quantize): + attention = _bare_attention() + original_impl = attention.impl + runner = _model_runner( + nn.ModuleDict({"attn": attention}), + sparse_metadata=_sparse_metadata(), + ) + monkeypatch.setattr(vllm, "__version__", "0.14.0") + + install = ( + vllm_runtime.install_vllm_nvfp4_attention + if quantize + else vllm_runtime.install_vllm_sparse_attention_from_checkpoint + ) + with pytest.raises(RuntimeError, match=r"vLLM >= 0\.15\.0"): + install(runner) + + assert attention.impl is original_impl + assert not hasattr(attention, "_query_quant_in_kernel") + assert runner.cascade_attn_enabled is True + + +def _fake_quant_plugin(configured): + def configure(module, *, device, dtype): + configured.append(module) + module.device = device + module.dtype = dtype + quantizer = SimpleNamespace( + is_enabled=True, + is_nvfp4_dynamic=True, + block_sizes={-1: 16}, + _amax=torch.tensor(1.0), + ) + module.q_bmm_quantizer = quantizer + module.k_bmm_quantizer = quantizer + module.p_bmm_quantizer = quantizer + module.v_bmm_quantizer = quantizer + return module + + return SimpleNamespace( + _get_device_dtype=lambda module: (module.device, module.dtype), + configure_vllm_nvfp4_attention_quantizers=configure, + ) + + +@pytest.mark.parametrize("with_sparse", [False, True]) +@pytest.mark.parametrize("impl_cls", [FlashAttentionImpl, FlashInferImpl]) +def test_nvfp4_install_composes_real_quantizers_with_optional_sparsity( + monkeypatch, + with_sparse, + impl_cls, +): + attention = _bare_attention(impl_cls) + unrelated = nn.Linear(4, 4) + model = nn.ModuleDict({"attn": attention, "linear": unrelated}) + runner = _model_runner(model, sparse_metadata=_sparse_metadata() if with_sparse else None) + monkeypatch.setattr( + quant_plugin, + "create_parallel_state", + lambda: quant_plugin.ParallelState(data_parallel_group=None), + ) + monkeypatch.setattr( + vllm_runtime.attention_plugin, "patch_flashinfer_metadata_builder", lambda: True + ) + + report = vllm_runtime.install_vllm_nvfp4_attention(runner) + + expected_impl_cls = ( + ModelOptSparseAttentionImpl + if impl_cls is FlashAttentionImpl + else get_flashinfer_sparse_impl_cls() + ) + assert isinstance(attention, quant_plugin._QuantVLLMAttention) + assert type(attention.impl) is expected_impl_cls + for name in ("q", "k", "p", "v"): + assert getattr(attention, f"{name}_bmm_quantizer").is_nvfp4_dynamic + assert attention.k_bmm_quantizer._amax == 6.0 * 448.0 + assert attention.v_bmm_quantizer._amax == 6.0 * 448.0 + assert attention._query_quant_in_kernel is True + assert attention._value_quant_in_kernel is True + assert attention.impl.quant_kw == { + "p_qdq": "nvfp4", + "p_qdq_amax": 1.0, + "v_qdq": "nvfp4", + "v_qdq_amax": 6.0 * 448.0, + } + assert bool(attention.impl.sparse_kw) is with_sparse + assert report.installed_layers == ("attn",) + assert report.quantized_layers == ("attn",) + assert bool(report.sparse_layers) is with_sparse + assert report.cascade_disabled is True + assert runner.cascade_attn_enabled is False + assert not hasattr(unrelated, "q_bmm_quantizer") + + +def test_nvfp4_validation_of_all_layers_precedes_mutation(monkeypatch): + valid = _bare_attention() + invalid = _bare_attention() + invalid.attn_type = "encoder" + invalid.head_size_v = 32 + valid_impl = valid.impl + invalid_impl = invalid.impl + runner = _model_runner(nn.ModuleDict({"valid": valid, "invalid": invalid})) + configured = [] + monkeypatch.setattr(vllm_runtime, "_load_quant_plugin", lambda: _fake_quant_plugin(configured)) + + with pytest.raises(NotImplementedError) as exc: + vllm_runtime.install_vllm_nvfp4_attention(runner) + + assert "invalid: attn_type" in str(exc.value) + assert "head_size_v" in str(exc.value) + assert configured == [] + assert valid.impl is valid_impl + assert invalid.impl is invalid_impl + assert not hasattr(valid, "_query_quant_in_kernel") + assert runner.cascade_attn_enabled is True + + +def test_apply_failure_does_not_leave_configured_layer_on_native_impl(monkeypatch): + first = _bare_attention() + second = _bare_attention() + first_impl = first.impl + second_impl = second.impl + runner = _model_runner(nn.ModuleDict({"first": first, "second": second})) + configured = [] + quant_plugin = _fake_quant_plugin(configured) + configure = quant_plugin.configure_vllm_nvfp4_attention_quantizers + + def fail_on_second(module, **kwargs): + if module is second: + raise RuntimeError("configuration failed") + return configure(module, **kwargs) + + quant_plugin.configure_vllm_nvfp4_attention_quantizers = fail_on_second + monkeypatch.setattr(vllm_runtime, "_load_quant_plugin", lambda: quant_plugin) + + with pytest.raises(RuntimeError, match="configuration failed"): + vllm_runtime.install_vllm_nvfp4_attention(runner) + + assert type(first.impl) is ModelOptSparseAttentionImpl + assert second.impl is second_impl + assert first.impl is not first_impl + assert first._query_quant_in_kernel is True + assert first._value_quant_in_kernel is True + assert not hasattr(second, "_query_quant_in_kernel") + assert runner.cascade_attn_enabled is False + + +@pytest.mark.parametrize( + ("mode", "rejected"), + [(CUDAGraphMode.FULL, True), (CUDAGraphMode.FULL_AND_PIECEWISE, False)], +) +def test_nvfp4_mixed_cudagraph_policy(monkeypatch, mode, rejected): + attention = _bare_attention() + original_impl = attention.impl + runner = _model_runner(nn.ModuleDict({"attn": attention})) + runner.vllm_config.compilation_config.cudagraph_mode = mode + configured = [] + monkeypatch.setattr(vllm_runtime, "_load_quant_plugin", lambda: _fake_quant_plugin(configured)) + + if rejected: + with pytest.raises(NotImplementedError, match="FULL mixed-batch"): + vllm_runtime.install_vllm_nvfp4_attention(runner) + assert configured == [] + assert attention.impl is original_impl + else: + report = vllm_runtime.install_vllm_nvfp4_attention(runner) + assert report.installed_count == 1 + assert configured == [attention] + + +@pytest.mark.parametrize("impl_cls", [FlashAttentionImpl, FlashInferImpl]) +def test_nvfp4_install_fp8_bmm2_uses_module_level_v(monkeypatch, impl_cls): + """bmm2='fp8': P maps to the fp8 kernel mode; V is module-level (no kernel V).""" + attention = _bare_attention(impl_cls) + runner = _model_runner(nn.ModuleDict({"attn": attention})) + monkeypatch.setattr( + quant_plugin, + "create_parallel_state", + lambda: quant_plugin.ParallelState(data_parallel_group=None), + ) + monkeypatch.setattr( + vllm_runtime.attention_plugin, "patch_flashinfer_metadata_builder", lambda: True + ) + + report = vllm_runtime.install_vllm_nvfp4_attention(runner, p_format="fp8", v_format="fp8") + + assert report.installed_layers == ("attn",) + # BMM2 quantizers configured FP8 per-tensor with fixed amax (F3) + assert attention.p_bmm_quantizer.num_bits == (4, 3) + assert float(attention.p_bmm_quantizer._amax) == 1.0 + assert attention.v_bmm_quantizer.num_bits == (4, 3) + assert float(attention.v_bmm_quantizer._amax) == 448.0 + # BMM1 untouched (F1) + assert attention.q_bmm_quantizer.is_nvfp4_dynamic + assert attention.k_bmm_quantizer.is_nvfp4_dynamic + # quant_kw: fp8 P reaches the kernel; V never does (module-level pre-cache-write) + assert attention.impl.quant_kw["p_qdq"] == "fp8" + assert attention.impl.quant_kw["p_qdq_amax"] == 1.0 + assert attention.impl.quant_kw["v_qdq"] is None + assert attention.impl.quant_kw["v_qdq_amax"] is None + # module forward owns V quant; Q stays in-kernel + assert attention._query_quant_in_kernel is True + assert attention._value_quant_in_kernel is False + + +def test_install_rejects_unknown_formats(): + with pytest.raises(ValueError, match="p_format must be"): + vllm_runtime.install_vllm_nvfp4_attention(object(), p_format="int8") + with pytest.raises(ValueError, match="v_format must be"): + vllm_runtime.install_vllm_nvfp4_attention(object(), v_format="int8") 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..62395ff5a7a 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,124 @@ def test_triton_fa_importable_on_cpu(): assert "attention" in triton_fa.__all__ assert callable(calibrate.attention_calibrate) + + +def test_forward_buckets_autotune_key_without_bucketing_grid(monkeypatch): + """Reuse autotune results by length regime without launching extra query tiles.""" + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + 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"), + [ + ({}, 0, 0), + ({"p_qdq": "fp8"}, 1, 0), + ({"p_qdq": "nvfp4", "v_qdq": "nvfp4"}, 2, 2), + ({"p_qdq": "nvfp4", "sparsity_n": 2, "sparsity_m": 4}, 2, 0), + ], +) +def test_forward_routes_every_mode_to_single_autotuner( + monkeypatch, attention_kwargs, expected_p_qdq, expected_v_qdq +): + """Every non-measurement launch uses the unified autotuner.""" + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + kernel = _CapturingKernel() + kernel.fn = _ForbiddenKernel() + monkeypatch.setattr(triton_fa, "_attn_fwd", kernel) + monkeypatch.setattr(triton_fa, "_attn_fwd_p_qdq", _ForbiddenKernel(), raising=False) + monkeypatch.setattr(triton_fa.torch.cuda, "device", lambda _device: nullcontext()) + monkeypatch.setattr(triton_fa, "_load_sparsity_helpers", lambda: None) + monkeypatch.setattr(triton_fa, "_load_qdq_helpers", lambda: None) + + seq_len = 129 + q = torch.empty(seq_len, 2, 16) + k = torch.empty(seq_len, 1, 16) + v = torch.empty_like(k) + starts = torch.tensor([0], dtype=torch.int32) + lengths = torch.tensor([seq_len], dtype=torch.int32) + + triton_fa.attention(q, k, v, starts, lengths, seq_len, **attention_kwargs) + + assert kernel.kwargs["P_QDQ"] == expected_p_qdq + assert kernel.kwargs["V_QDQ"] == expected_v_qdq + + +@pytest.mark.parametrize( + ("attention_kwargs", "expected_block_m"), + [ + ({"skip_softmax_threshold": 0.1, "measure_sparsity": True}, 128), + ( + { + "p_qdq": "nvfp4", + "skip_softmax_threshold": 0.1, + "measure_sparsity": True, + }, + 16, + ), + ], +) +def test_forward_measurement_uses_one_fixed_launch(monkeypatch, attention_kwargs, expected_block_m): + """Counter measurement bypasses autotuning to avoid repeated atomic updates.""" + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + kernel = _ForbiddenKernel() + kernel.fn = _CapturingKernel() + monkeypatch.setattr(triton_fa, "_attn_fwd", kernel) + monkeypatch.setattr(triton_fa.torch.cuda, "device", lambda _device: nullcontext()) + monkeypatch.setattr(triton_fa, "_load_sparsity_helpers", lambda: None) + monkeypatch.setattr(triton_fa, "_load_qdq_helpers", lambda: None) + + seq_len = 129 + q = torch.empty(seq_len, 2, 16) + k = torch.empty(seq_len, 1, 16) + v = torch.empty_like(k) + starts = torch.tensor([0], dtype=torch.int32) + lengths = torch.tensor([seq_len], dtype=torch.int32) + + triton_fa.attention(q, k, v, starts, lengths, seq_len, **attention_kwargs) + + assert kernel.fn.launch_count == 1 + assert kernel.fn.kwargs["BLOCK_M"] == expected_block_m + assert kernel.fn.kwargs["BLOCK_N"] == 128 + assert kernel.fn.kwargs["num_stages"] == 1 + assert kernel.fn.kwargs["num_warps"] == 4