diff --git a/benchmarks/kernels/benchmark_vit_fp8_attn.py b/benchmarks/kernels/benchmark_vit_fp8_attn.py new file mode 100644 index 00000000000..7d7a067dde9 --- /dev/null +++ b/benchmarks/kernels/benchmark_vit_fp8_attn.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Benchmarks FP8 vs BF16 ViT attention via FlashInfer cuDNN backend. +# +# == Usage Examples == +# +# Benchmark mode (default, FlashInfer CUDAGraph Bench) +# python3 benchmark_vit_fp8_attn.py +# +# Profile mode (PyTorch profiler, saves TensorBoard traces): +# python3 benchmark_vit_fp8_attn.py --profile +# python3 benchmark_vit_fp8_attn.py --profile --profile-output-dir ./profile_traces +# +# Custom seq_lens: +# python3 benchmark_vit_fp8_attn.py --seq-lens 4096 8192 16384 + +from functools import partial + +import numpy as np +import torch +from torch.profiler import ProfilerActivity, profile, record_function + +from vllm.utils.argparse_utils import FlexibleArgumentParser + +# Qwen3-VL defaults +NUM_HEADS = 16 +HEAD_DIM = 72 +DEFAULT_SEQ_LENS = [2304, 4096, 8192, 16384] + + +def _setup_fp8_attention(num_heads: int, head_dim: int) -> tuple: + """Create FP8 and BF16 attention modules + workspace.""" + from types import SimpleNamespace + from unittest.mock import patch + + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.config.multimodal import MultiModalConfig + from vllm.model_executor.layers.attention.mm_encoder_attention import ( + MMEncoderAttention, + _get_flashinfer_workspace_buffer, + ) + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + old_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + + backend_patch = patch( + "vllm.model_executor.layers.attention.mm_encoder_attention" + ".get_vit_attn_backend", + return_value=AttentionBackendEnum.FLASHINFER, + ) + + # FP8 attention + mm_config_fp8 = MultiModalConfig(mm_encoder_attn_dtype="fp8") + vllm_config_fp8 = VllmConfig() + vllm_config_fp8.model_config = SimpleNamespace(multimodal_config=mm_config_fp8) + with set_current_vllm_config(vllm_config_fp8), backend_patch: + attn_fp8 = MMEncoderAttention( + num_heads=num_heads, + head_size=head_dim, + prefix="visual.blocks.0.attn", + ).to("cuda") + + # BF16 attention (no FP8) + with set_current_vllm_config(VllmConfig()), backend_patch: + attn_bf16 = MMEncoderAttention( + num_heads=num_heads, + head_size=head_dim, + prefix="visual.blocks.0.attn", + ).to("cuda") + + torch.set_default_dtype(old_dtype) + + workspace = _get_flashinfer_workspace_buffer() + return attn_fp8, attn_bf16, workspace + + +def _build_meta( + seq_len: int, + num_heads: int, + head_dim: int, + fp8: bool, +): + """Build cu_seqlens, max_seqlen, sequence_lengths.""" + from vllm.model_executor.layers.attention.mm_encoder_attention import ( + MMEncoderAttention, + ) + from vllm.utils.math_utils import round_up + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + cu_np = np.array([0, seq_len], dtype=np.int32) + fp8_padded = num_heads * round_up(head_dim, 16) if fp8 else None + + seq_lengths = MMEncoderAttention.maybe_compute_seq_lens( + AttentionBackendEnum.FLASHINFER, cu_np, torch.device("cuda") + ) + max_seqlen = torch.tensor( + MMEncoderAttention.compute_max_seqlen(AttentionBackendEnum.FLASHINFER, cu_np), + dtype=torch.int32, + ) + cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens( + AttentionBackendEnum.FLASHINFER, + cu_np, + num_heads * head_dim, + 1, + torch.device("cuda"), + fp8_padded_hidden_size=fp8_padded, + ) + return cu_seqlens, max_seqlen, seq_lengths + + +def run_benchmark( + seq_lens: list[int], + num_heads: int, + head_dim: int, + method: str, +): + """Benchmark FP8 vs BF16 attention across seq_lens. + + Uses FlashInfer GPU-level timing to measure pure kernel time, + excluding CPU launch overhead. + """ + if method == "cupti": + from flashinfer.testing import bench_gpu_time_with_cupti as bench_fn + + bench_fn = partial(bench_fn, use_cuda_graph=True, cold_l2_cache=False) + elif method == "cudagraph": + from flashinfer.testing import ( + bench_gpu_time_with_cudagraph as bench_fn, + ) + + bench_fn = partial(bench_fn, cold_l2_cache=False) + else: + raise ValueError(f"Invalid method: {method}") + + attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim) + + print(f"Timing method: {method}") + print(f"{'seq_len':>8} {'BF16 (us)':>12} {'FP8 (us)':>12} {'Speedup':>10}") + print("-" * 46) + + for seq_len in seq_lens: + torch.manual_seed(42) + + q = torch.randn( + seq_len, + num_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k = torch.randn_like(q) + v = torch.randn_like(q) + + cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True) + # we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous + cu_bf16 = cu_fp8.clone() + + def bf16_fn(q=q, k=k, v=v, cu=cu_bf16, ms=max_s, sl=seq_l): + attn_bf16._forward_flashinfer(q, k, v, cu, ms, sl) + + def fp8_fn(q=q, k=k, v=v, cu=cu_fp8, ms=max_s, sl=seq_l): + attn_fp8._forward_flashinfer(q, k, v, cu, ms, sl) + + # bench_fn returns List[float] of per-iteration times in ms + bf16_times = bench_fn(bf16_fn) + fp8_times = bench_fn(fp8_fn) + + bf16_us = np.median(bf16_times) * 1e3 # ms -> us + fp8_us = np.median(fp8_times) * 1e3 + speedup = bf16_us / fp8_us if fp8_us > 0 else float("inf") + + print(f"{seq_len:>8} {bf16_us:>12.1f} {fp8_us:>12.1f} {speedup:>9.2f}x") + + +def _make_trace_handler(output_dir: str, worker_name: str, label: str): + """Create a trace handler that saves to TensorBoard and prints summary.""" + + def handler(prof): + torch.profiler.tensorboard_trace_handler(output_dir, worker_name)(prof) + print(f"\n{'=' * 80}") + print(label) + print(f"{'=' * 80}") + print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20)) + + return handler + + +def run_profile( + seq_len: int, + num_heads: int, + head_dim: int, + warmup: int, + output_dir: str, +): + """Profile FP8 vs BF16 attention with PyTorch profiler.""" + attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim) + + torch.manual_seed(42) + q = torch.randn( + seq_len, + num_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k = torch.randn_like(q) + v = torch.randn_like(q) + + cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True) + # we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous + cu_bf16 = cu_fp8.clone() + + sched = torch.profiler.schedule(wait=0, warmup=warmup, active=1) + + # Profile BF16 (warmup handled by profiler schedule) + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + schedule=sched, + on_trace_ready=_make_trace_handler( + output_dir, + f"bf16_h{head_dim}_s{seq_len}", + f"BF16 Attention (seq_len={seq_len}, heads={num_heads}, " + f"head_dim={head_dim})", + ), + ) as prof_bf16: + for _ in range(warmup + 1): + with record_function("bf16_attention"): + attn_bf16._forward_flashinfer( + q.clone(), k.clone(), v.clone(), cu_bf16, max_s, seq_l + ) + torch.accelerator.synchronize() + prof_bf16.step() + + # Profile FP8 (warmup handled by profiler schedule) + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + schedule=sched, + on_trace_ready=_make_trace_handler( + output_dir, + f"fp8_h{head_dim}_s{seq_len}", + f"FP8 Attention (seq_len={seq_len}, heads={num_heads}, " + f"head_dim={head_dim})", + ), + ) as prof_fp8: + for _ in range(warmup + 1): + with record_function("fp8_attention"): + attn_fp8._forward_flashinfer( + q.clone(), k.clone(), v.clone(), cu_fp8, max_s, seq_l + ) + torch.accelerator.synchronize() + prof_fp8.step() + + print(f"\nTensorBoard traces saved to: {output_dir}") + print(f"View with: tensorboard --logdir={output_dir}") + + +if __name__ == "__main__": + parser = FlexibleArgumentParser(description="Benchmark FP8 vs BF16 ViT attention.") + parser.add_argument( + "--seq-lens", + type=int, + nargs="+", + default=DEFAULT_SEQ_LENS, + help="Sequence lengths to benchmark", + ) + parser.add_argument( + "--num-heads", + type=int, + default=NUM_HEADS, + ) + parser.add_argument( + "--head-dim", + type=int, + default=HEAD_DIM, + ) + parser.add_argument( + "--method", + choices=["cupti", "cudagraph"], + default="cudagraph", + help="GPU timing method: cupti (CUPTI kernel timing) or " + "cudagraph (CUDA graph capture/replay). Default: cudagraph", + ) + parser.add_argument( + "--warmup", + type=int, + default=10, + help="Warmup iterations (profile mode only)", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Run PyTorch profiler instead of benchmark", + ) + parser.add_argument( + "--profile-seq-len", + type=int, + default=8192, + help="Sequence length for profiling (default: 8192)", + ) + parser.add_argument( + "--profile-output-dir", + type=str, + default="./profile_traces", + help="Output directory for TensorBoard traces (default: ./profile_traces)", + ) + args = parser.parse_args() + + if args.profile: + run_profile( + args.profile_seq_len, + args.num_heads, + args.head_dim, + args.warmup, + args.profile_output_dir, + ) + else: + run_benchmark( + args.seq_lens, + args.num_heads, + args.head_dim, + args.method, + ) diff --git a/docs/features/quantization/README.md b/docs/features/quantization/README.md index 4be088c56a4..6c4aa7d8aaa 100644 --- a/docs/features/quantization/README.md +++ b/docs/features/quantization/README.md @@ -20,6 +20,7 @@ The following are the supported quantization formats for vLLM: - [AMD Quark](quark.md) - [Quantized KV Cache](quantized_kvcache.md) - [TorchAO](torchao.md) +- [FP8 ViT Encoder Attention](fp8_vit_attn.md) ## Supported Hardware diff --git a/docs/features/quantization/fp8_vit_attn.md b/docs/features/quantization/fp8_vit_attn.md new file mode 100644 index 00000000000..bf628cd8a72 --- /dev/null +++ b/docs/features/quantization/fp8_vit_attn.md @@ -0,0 +1,109 @@ +# FP8 ViT Encoder Attention + +For visual understanding workloads with large images (e.g. QHD, 4K) and relatively +short text prompts/generation, the ViT encoder attention can become a significant +bottleneck, especially when the text model is quantized (e.g. NVFP4). vLLM +supports optional FP8 quantization for the ViT encoder attention via the +FlashInfer cuDNN backend. Q/K/V are quantized on-the-fly to FP8 before the +cuDNN attention call. + +!!! note + - Currently supports Qwen3-VL family models only (`qwen3_vl`, `qwen3_vl_moe`, + `qwen3_5`, `qwen3_5_moe`, and other models using Qwen3 ViT). + - Dynamic scaling is not compatible with ViT full CUDA graphs. + - Performance gains are mostly visible at QHD/4K resolutions or multi-image + requests. Smaller images may see no speedup due to quantization overhead + (3 quantization kernel launches + un-padding). + - FP8 tensor-core speedup is more pronounced on GB300 than GB200. + +## Requirements + +- FlashInfer cuDNN backend with cuDNN >= 9.17.1. + +## Usage + +Enable FP8 ViT attention by passing `--mm-encoder-attn-dtype fp8` together +with `--mm-encoder-attn-backend FLASHINFER`: + +```bash +vllm serve $MODEL \ + --mm-encoder-attn-backend FLASHINFER \ + --mm-encoder-attn-dtype fp8 +``` + +By default (no scale file), **dynamic scaling** is used: a 16-entry circular +buffer of observed Q/K/V amax values drives per-forward scale updates. This +matches BF16 accuracy without any calibration but adds a small per-forward +overhead. + +## Calibrate-Once, Reuse Workflow (Recommended) + +For production, calibrate static scales on a representative dataset once and +reuse them to avoid the dynamic overhead: + +```bash +# Step 1: calibrate and save scales (runs dynamic scaling for 16 passes, +# then dumps the learned scales to JSON). +vllm bench mm-processor \ + --model $MODEL --mm-encoder-attn-backend FLASHINFER \ + --mm-encoder-attn-dtype fp8 \ + --mm-encoder-fp8-scale-save-path /path/to/scales.json \ + --dataset-name hf --dataset-path lmarena-ai/VisionArena-Chat \ + --num-prompts 100 + +# Step 2: serve with static scales (no dynamic overhead). +vllm serve $MODEL \ + --mm-encoder-attn-backend FLASHINFER \ + --mm-encoder-attn-dtype fp8 \ + --mm-encoder-fp8-scale-path /path/to/scales.json +``` + +Saved scales are multiplied by `--mm-encoder-fp8-scale-save-margin` (default +`1.5`) to leave headroom against activation outliers not present in the +calibration set. The default has been validated to generalize across datasets +(e.g. VisionArena-Chat calibration maintains BF16 accuracy on ChartQA). + +## Scale File Format + +```json +{ + "visual.blocks.0.attn.attn": {"q": 224.0, "k": 198.0, "v": 210.0}, + "visual.blocks.1.attn.attn": {"q": 218.0, "k": 195.0, "v": 207.0} +} +``` + +Keys `q_scale` / `k_scale` / `v_scale` are accepted as aliases. + +## Performance + +**Core cuDNN attention kernel** (PyTorch profiler, `cudnn_generated_fort_native_sdpa_sm100_flash_fprop`, head_dim=128, seq_len=8192): + +| Hardware | BF16 | FP8 | Speedup | +| -------- | ---- | ---- | ------- | +| GB200 | 350 us | 312 us | **1.12x** | +| GB300 | 300 us | 211 us | **1.42x** | + +**End-to-end encoder forward time** (Qwen3-VL-30B-A3B-Instruct on GB200, 3 images/request): + +| Resolution | BF16 median | FP8 median | Speedup | +| ---------- | ----------- | ---------- | ------- | +| HD (720x1280) | 31.77 ms | 36.39 ms | 0.87x | +| FullHD (1080x1920) | 57.99 ms | 58.73 ms | ~same | +| QHD (1440x2560) | 131.83 ms | 122.30 ms | **1.08x** | +| 4K (2160x3840) | 543.44 ms | 460.31 ms | **1.18x** | + +Crossover is around FullHD with 3 images/request. At QHD and above, FP8 wins. + +## Accuracy + +ChartQA, Qwen3-VL-8B-Instruct, 500 samples. FP8 static uses scales calibrated +on VisionArena-Chat (with default 1.5x margin): + +| Metric | BF16 | FP8 dynamic | FP8 static | +| ------ | ---- | ----------- | ---------- | +| relaxed_accuracy | 0.780 | 0.776 | 0.780 | +| anywhere_accuracy | 0.806 | 0.816 | 0.814 | +| exact_match | 0.584 | 0.582 | 0.578 | + +All three configurations match within statistical noise, confirming that +static scales calibrated on one dataset generalize to another. diff --git a/tests/config/test_multimodal_config.py b/tests/config/test_multimodal_config.py index e5c30f999a0..9720d84672f 100644 --- a/tests/config/test_multimodal_config.py +++ b/tests/config/test_multimodal_config.py @@ -41,3 +41,21 @@ def test_language_model_only_affects_model_hash(): base_hash = ModelConfig(model).compute_hash() lm_only_hash = ModelConfig(model, language_model_only=True).compute_hash() assert base_hash != lm_only_hash + + +def test_mm_encoder_fp8_scale_path_requires_fp8(): + with pytest.raises(ValueError, match="mm_encoder_attn_dtype"): + MultiModalConfig(mm_encoder_fp8_scale_path="/tmp/scales.json") + + +def test_mm_encoder_attn_dtype_hash_updates(tmp_path): + scale_file = tmp_path / "scales.json" + scale_file.write_text("{}") + base_hash = MultiModalConfig().compute_hash() + fp8_hash = MultiModalConfig(mm_encoder_attn_dtype="fp8").compute_hash() + fp8_static_hash = MultiModalConfig( + mm_encoder_attn_dtype="fp8", + mm_encoder_fp8_scale_path=str(scale_file), + ).compute_hash() + assert base_hash != fp8_hash + assert fp8_hash != fp8_static_hash diff --git a/tests/kernels/core/test_vit_fp8_attn.py b/tests/kernels/core/test_vit_fp8_attn.py new file mode 100644 index 00000000000..ef1c44cada2 --- /dev/null +++ b/tests/kernels/core/test_vit_fp8_attn.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the full FP8 ViT attention path (quantize -> cuDNN -> un-pad).""" + +import contextlib + +import pytest +import torch + +from vllm.triton_utils import HAS_TRITON +from vllm.utils.flashinfer import ( + is_flashinfer_cudnn_fp8_prefill_attn_supported, +) +from vllm.v1.attention.backends.registry import AttentionBackendEnum + + +def _has_flashinfer_cudnn() -> bool: + """Check if FlashInfer cuDNN backend is available.""" + try: + from flashinfer.prefill import ( + cudnn_batch_prefill_with_kv_cache, # noqa: F401 + ) + + return True + except ImportError: + return False + + +HEAD_DIMS = [72, 80] +SEQ_LENS = [256] +NUM_HEADS = [16] + + +@pytest.fixture +def _fp8_attention(): + """Create FP8-enabled MMEncoderAttention via config.""" + from types import SimpleNamespace + from unittest.mock import patch + + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.config.multimodal import MultiModalConfig + + if not is_flashinfer_cudnn_fp8_prefill_attn_supported(): + pytest.skip("FlashInfer cuDNN FP8 prefill attention not supported") + + mm_config = MultiModalConfig(mm_encoder_attn_dtype="fp8") + vllm_config = VllmConfig() + vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config) + + # MMEncoderAttention reads torch.get_default_dtype() during init + # to determine the output dtype. In real model loading this is bf16. + old_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.model_executor.layers.attention.mm_encoder_attention" + ".get_vit_attn_backend", + return_value=AttentionBackendEnum.FLASHINFER, + ), + ): + yield + + torch.set_default_dtype(old_dtype) + + +def _build_cu_seqlens_and_meta( + seq_len: int, + num_heads: int, + head_dim: int, + fp8_padded_hidden_size: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build cu_seqlens, max_seqlen, sequence_lengths for a single sequence.""" + import numpy as np + + from vllm.model_executor.layers.attention.mm_encoder_attention import ( + MMEncoderAttention, + ) + + cu_seqlens_np = np.array([0, seq_len], dtype=np.int32) + + sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens( + AttentionBackendEnum.FLASHINFER, + cu_seqlens_np, + torch.device("cuda"), + ) + + max_seqlen = torch.tensor( + MMEncoderAttention.compute_max_seqlen( + AttentionBackendEnum.FLASHINFER, cu_seqlens_np + ), + dtype=torch.int32, + ) + + cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens( + AttentionBackendEnum.FLASHINFER, + cu_seqlens_np, + num_heads * head_dim, + 1, # tp_size + torch.device("cuda"), + fp8_padded_hidden_size=fp8_padded_hidden_size, + ) + + return cu_seqlens, max_seqlen, sequence_lengths + + +@pytest.mark.skipif( + not (HAS_TRITON and _has_flashinfer_cudnn()), + reason="Triton and FlashInfer cuDNN required", +) +@pytest.mark.parametrize("head_dim", HEAD_DIMS) +@pytest.mark.parametrize("seq_len", SEQ_LENS) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +def test_fp8_attn_output_shape( + head_dim: int, + seq_len: int, + num_heads: int, + _fp8_attention, +) -> None: + """Verify FP8 attention produces correct output shape after un-padding.""" + from vllm.model_executor.layers.attention.mm_encoder_attention import ( + MMEncoderAttention, + ) + from vllm.utils.math_utils import round_up + + attn = None + with contextlib.suppress(ValueError, ImportError): + attn = MMEncoderAttention( + num_heads=num_heads, + head_size=head_dim, + prefix="visual.blocks.0.attn", + ).to("cuda") + + if attn is None or not attn.fp8_enabled: + pytest.skip("FP8 MMEncoderAttention not available") + assert attn is not None # mypy narrowing + + # FP8 always needs fp8_padded_hidden_size for correct cu_seqlens + fp8_padded_hidden_size = num_heads * round_up(head_dim, 16) + + cu_seqlens, max_seqlen, sequence_lengths = _build_cu_seqlens_and_meta( + seq_len, num_heads, head_dim, fp8_padded_hidden_size=fp8_padded_hidden_size + ) + + q = torch.randn( + seq_len, + num_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k = torch.randn_like(q) + v = torch.randn_like(q) + + output = attn._forward_flashinfer(q, k, v, cu_seqlens, max_seqlen, sequence_lengths) + + # Output should have original head_dim (un-padded) + assert output.shape[-1] == head_dim + assert output.dtype == torch.bfloat16 + + +@pytest.mark.skipif( + not (HAS_TRITON and _has_flashinfer_cudnn()), + reason="Triton and FlashInfer cuDNN required", +) +@pytest.mark.parametrize("head_dim", HEAD_DIMS) +@pytest.mark.parametrize("seq_len", SEQ_LENS) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +def test_fp8_vs_bf16_close( + head_dim: int, seq_len: int, num_heads: int, _fp8_attention +) -> None: + """FP8 attention output should be reasonably close to BF16 baseline.""" + from vllm.model_executor.layers.attention.mm_encoder_attention import ( + MMEncoderAttention, + ) + from vllm.utils.math_utils import round_up + + torch.manual_seed(42) + q = torch.randn( + 1, + seq_len, + num_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k = torch.randn_like(q) + v = torch.randn_like(q) + + # FP8 path + attn_fp8 = None + with contextlib.suppress(ValueError, ImportError): + attn_fp8 = MMEncoderAttention( + num_heads=num_heads, + head_size=head_dim, + prefix="visual.blocks.0.attn", + ).to("cuda") + + if attn_fp8 is None or not attn_fp8.fp8_enabled: + pytest.skip("FP8 MMEncoderAttention not available") + assert attn_fp8 is not None # mypy narrowing + + fp8_padded_hidden_size = num_heads * round_up(head_dim, 16) + cu_seqlens, max_seqlen, seq_lengths = _build_cu_seqlens_and_meta( + seq_len, + num_heads, + head_dim, + fp8_padded_hidden_size=fp8_padded_hidden_size, + ) + + out_fp8 = attn_fp8._forward_flashinfer( + q.clone(), + k.clone(), + v.clone(), + cu_seqlens, + max_seqlen, + seq_lengths, + ) + + # BF16 baseline (create non-FP8 attention by using scale=attn_fp8.scale + # and calling the wrapper directly without FP8 quantization) + from vllm.model_executor.layers.attention.mm_encoder_attention import ( + _get_flashinfer_workspace_buffer, + ) + from vllm.v1.attention.ops.vit_attn_wrappers import ( + vit_flashinfer_wrapper, + ) + + out_bf16 = vit_flashinfer_wrapper( + q=q.clone(), + k=k.clone(), + v=v.clone(), + scale=attn_fp8.scale, + workspace_buffer=_get_flashinfer_workspace_buffer(), + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + sequence_lengths=seq_lengths, + ) + + out_fp8_f = out_fp8.float() + out_bf16_f = out_bf16.float() + + abs_diff = (out_fp8_f - out_bf16_f).abs() + abs_diff_flat = abs_diff.flatten() + + # Relative diff (avoid division by zero) + denom = out_bf16_f.abs().clamp(min=1e-6) + rel_diff_flat = (abs_diff / denom).flatten() + + cosine_sim = torch.nn.functional.cosine_similarity( + out_fp8_f.flatten().unsqueeze(0), + out_bf16_f.flatten().unsqueeze(0), + ).item() + + pcts = [50, 90, 95, 99, 99.9] + abs_pct = {p: torch.quantile(abs_diff_flat, p / 100).item() for p in pcts} + rel_pct = {p: torch.quantile(rel_diff_flat, p / 100).item() for p in pcts} + + print(f"\nFP8 vs BF16 (head_dim={head_dim}, seq_len={seq_len}):") + print(f" cosine_sim={cosine_sim:.6f}") + print( + f" abs_diff: max={abs_diff_flat.max().item():.6f}, " + f"mean={abs_diff_flat.mean().item():.6f}, " + + ", ".join(f"p{p}={abs_pct[p]:.6f}" for p in pcts) + ) + print( + f" rel_diff: max={rel_diff_flat.max().item():.6f}, " + f"mean={rel_diff_flat.mean().item():.6f}, " + + ", ".join(f"p{p}={rel_pct[p]:.6f}" for p in pcts) + ) + + assert abs_diff_flat.max().item() < 0.3, ( + f"FP8 vs BF16 max abs diff too large: {abs_diff_flat.max().item()}" + ) + assert abs_diff_flat.mean().item() < 0.03, ( + f"FP8 vs BF16 mean abs diff too large: {abs_diff_flat.mean().item()}" + ) + assert cosine_sim > 0.99, f"Cosine similarity too low: {cosine_sim:.6f}" diff --git a/tests/kernels/core/test_vit_fp8_quant.py b/tests/kernels/core/test_vit_fp8_quant.py new file mode 100644 index 00000000000..0c63d0069f1 --- /dev/null +++ b/tests/kernels/core/test_vit_fp8_quant.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the stride-aware FP8 quantization kernel with head_dim padding.""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON + +if HAS_TRITON: + from vllm.kernels.triton.qkv_padded_fp8_quant import ( + quantize_fp8_pad_head_dim_triton, + ) + +HEAD_DIMS = [72, 80, 128] +SEQ_LENS = [64, 256] +NUM_HEADS = [16] +SCALES = [0.01, 0.1, 1.0] + + +def _naive_fp8_quantize( + tensor: torch.Tensor, scale: torch.Tensor, skip_scale: bool +) -> torch.Tensor: + """Reference FP8 quantization in PyTorch.""" + fp8_dtype = current_platform.fp8_dtype() + fp8_max = torch.finfo(fp8_dtype).max + fp8_min = -fp8_max + + x = tensor.float() + if not skip_scale: + x = x / scale.item() + x = x.clamp(fp8_min, fp8_max) + return x.to(fp8_dtype) + + +@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available") +@pytest.mark.parametrize("head_dim", HEAD_DIMS) +@pytest.mark.parametrize("seq_len", SEQ_LENS) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("scale_val", SCALES) +def test_quantize_contiguous( + head_dim: int, seq_len: int, num_heads: int, scale_val: float +) -> None: + """Test quantization of contiguous 3D tensors.""" + torch.manual_seed(42) + tensor = torch.randn( + seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16 + ) + scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda").view( + 1, 1, 1, 1 + ) + + result = quantize_fp8_pad_head_dim_triton(tensor, scale) + + padded_dim = (head_dim + 15) // 16 * 16 + assert result.shape == (seq_len, num_heads, padded_dim) + assert result.is_contiguous() + assert result.dtype == current_platform.fp8_dtype() + + # Compare unpadded portion against reference + ref = _naive_fp8_quantize(tensor, scale, skip_scale=False) + torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float()) + + # Padded region should be zero + if padded_dim > head_dim: + assert (result[:, :, head_dim:].float() == 0).all() + + +@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available") +@pytest.mark.parametrize("head_dim", [72, 80]) +def test_quantize_non_contiguous(head_dim: int) -> None: + """Test quantization from non-contiguous QKV views (interleaved buffer).""" + seq_len, num_heads = 64, 16 + # Simulate interleaved QKV buffer: shape (seq_len, 3 * num_heads, head_dim) + qkv = torch.randn( + seq_len, 3 * num_heads, head_dim, device="cuda", dtype=torch.bfloat16 + ) + # Q is every 3rd head slice - non-contiguous view + q = qkv[:, 0::3, :] + assert not q.is_contiguous() + + scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1) + result = quantize_fp8_pad_head_dim_triton(q, scale) + + padded_dim = (head_dim + 15) // 16 * 16 + assert result.shape == (seq_len, num_heads, padded_dim) + assert result.is_contiguous() + + # Compare against contiguous reference + ref = _naive_fp8_quantize(q.contiguous(), scale, skip_scale=False) + torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float()) + + +@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available") +def test_skip_scale() -> None: + """Test skip_scale=True produces cast-only output (no division).""" + seq_len, num_heads, head_dim = 32, 8, 80 + tensor = torch.randn( + seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16 + ) + scale = torch.tensor([0.5], dtype=torch.float32, device="cuda").view(1, 1, 1, 1) + + result_skip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=True) + result_noskip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=False) + + # skip_scale should just cast, not divide + ref_cast = _naive_fp8_quantize(tensor, scale, skip_scale=True) + torch.testing.assert_close(result_skip[:, :, :head_dim].float(), ref_cast.float()) + + # With scale != 1.0, skip and no-skip should differ + assert not torch.equal(result_skip.float(), result_noskip.float()) + + +@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available") +def test_4d_input() -> None: + """Test that 4D input (B, S, H, D) is handled correctly.""" + B, S, H, D = 2, 32, 8, 72 + tensor = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) + scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1) + + result = quantize_fp8_pad_head_dim_triton(tensor, scale) + padded_dim = (D + 15) // 16 * 16 + assert result.shape == (B, S, H, padded_dim) diff --git a/tests/kernels/core/test_vit_fp8_scaling.py b/tests/kernels/core/test_vit_fp8_scaling.py new file mode 100644 index 00000000000..a197439237f --- /dev/null +++ b/tests/kernels/core/test_vit_fp8_scaling.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for FP8 scaling (dynamic and static) in MMEncoderAttention.""" + +import contextlib +import json +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from vllm.model_executor.layers.attention.mm_encoder_attention import ( + _FP8_AMAX_HISTORY_LEN, + _FP8_MAX, +) +from vllm.utils.flashinfer import ( + is_flashinfer_cudnn_fp8_prefill_attn_supported, +) + +LAYER_0 = "visual.blocks.0.attn.attn" +LAYER_1 = "visual.blocks.1.attn.attn" +NUM_HEADS = 16 +HEAD_DIM = 72 + + +@contextlib.contextmanager +def _build_attention(mm_config): + """Yield an MMEncoderAttention with the given multimodal config. + + The VllmConfig context stays active while the test runs so that + ``get_multimodal_config()`` calls during the forward path resolve. Also + invokes ``process_weights_after_loading`` to simulate the model loader's + auto-scan. Yields ``None`` if FlashInfer cuDNN is not available. + """ + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.attention.mm_encoder_attention import ( + MMEncoderAttention, + ) + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + if not is_flashinfer_cudnn_fp8_prefill_attn_supported(): + yield None + return + + vllm_config = VllmConfig() + vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config) + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.model_executor.layers.attention.mm_encoder_attention" + ".get_vit_attn_backend", + return_value=AttentionBackendEnum.FLASHINFER, + ), + ): + attn = MMEncoderAttention( + num_heads=NUM_HEADS, + head_size=HEAD_DIM, + prefix=LAYER_0, + ) + attn.process_weights_after_loading(torch.bfloat16) + yield attn + + +@pytest.fixture +def _make_attention(): + """Create an MMEncoderAttention with dynamic FP8 scaling.""" + from vllm.config.multimodal import MultiModalConfig + + with _build_attention(MultiModalConfig(mm_encoder_attn_dtype="fp8")) as attn: + yield attn + + +@pytest.fixture +def _make_static_attention(tmp_path): + """Create an MMEncoderAttention with static FP8 scales from a file.""" + from vllm.config.multimodal import MultiModalConfig + + scale_file = tmp_path / "scales.json" + scale_file.write_text( + json.dumps( + { + LAYER_0: {"q": 224.0, "k": 198.0, "v": 210.0}, + LAYER_1: {"q": 100.0, "k": 110.0, "v": 120.0}, + } + ) + ) + with _build_attention( + MultiModalConfig( + mm_encoder_attn_dtype="fp8", + mm_encoder_fp8_scale_path=str(scale_file), + ) + ) as attn: + yield attn + + +def test_dynamic_scaling_updates_scales(_make_attention) -> None: + """Verify that _record_amax_and_update_scales updates scale buffers.""" + attn = _make_attention + if attn is None or not attn.fp8_enabled: + pytest.skip("FP8 attention not available (FlashInfer backend required)") + + attn = attn.to("cuda") + + S, H, D = 32, NUM_HEADS, HEAD_DIM + q = torch.full((S, H, D), 2.0, device="cuda", dtype=torch.bfloat16) + k = torch.full((S, H, D), 3.0, device="cuda", dtype=torch.bfloat16) + v = torch.full((S, H, D), 4.0, device="cuda", dtype=torch.bfloat16) + + attn._record_amax_and_update_scales(q, k, v) + + expected_q_scale = 2.0 / _FP8_MAX + expected_k_scale = 3.0 / _FP8_MAX + expected_v_scale = 4.0 / _FP8_MAX + + torch.testing.assert_close(attn._fp8_q_scale.item(), expected_q_scale) + torch.testing.assert_close(attn._fp8_k_scale.item(), expected_k_scale) + torch.testing.assert_close(attn._fp8_v_scale.item(), expected_v_scale) + + +def test_circular_buffer_wraps(_make_attention) -> None: + """Verify the amax circular buffer wraps at HISTORY_LEN.""" + attn = _make_attention + if attn is None or not attn.fp8_enabled: + pytest.skip("FP8 attention not available (FlashInfer backend required)") + + attn = attn.to("cuda") + S, H, D = 16, NUM_HEADS, HEAD_DIM + + for i in range(_FP8_AMAX_HISTORY_LEN + 2): + mag = float(i + 1) + q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16) + k = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16) + v = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16) + attn._record_amax_and_update_scales(q, k, v) + + assert attn._fp8_amax_pos == 2 + + expected_max = float(_FP8_AMAX_HISTORY_LEN + 2) + expected_scale = expected_max / _FP8_MAX + torch.testing.assert_close(attn._fp8_q_scale.item(), expected_scale) + + +def test_static_scales_loaded(_make_static_attention) -> None: + """Verify static scales are loaded from the JSON file.""" + attn = _make_static_attention + if attn is None or not attn.fp8_enabled: + pytest.skip("FP8 attention not available (FlashInfer backend required)") + + assert attn.fp8_enabled + assert not attn._fp8_dynamic_scale + + # Layer 0 scales (the layer this attention was created with). + assert attn._fp8_q_scale.item() == 224.0 + assert attn._fp8_k_scale.item() == 198.0 + assert attn._fp8_v_scale.item() == 210.0 + + assert not attn.skip_scale_q + assert not attn.skip_scale_k + assert not attn.skip_scale_v + + # No amax history buffers for static scaling. + assert not hasattr(attn, "_fp8_q_amax") + + +def test_static_scales_missing_layer(tmp_path) -> None: + """Verify error when requested layer is not in the scale file.""" + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.config.multimodal import MultiModalConfig + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + if not is_flashinfer_cudnn_fp8_prefill_attn_supported(): + pytest.skip("FlashInfer cuDNN not available") + + scale_file = tmp_path / "wrong_layer.json" + scale_file.write_text( + json.dumps({"visual.blocks.99.attn": {"q": 1.0, "k": 1.0, "v": 1.0}}) + ) + mm_config = MultiModalConfig( + mm_encoder_attn_dtype="fp8", + mm_encoder_fp8_scale_path=str(scale_file), + ) + vllm_config = VllmConfig() + vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config) + + from vllm.model_executor.layers.attention.mm_encoder_attention import ( + MMEncoderAttention, + ) + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.model_executor.layers.attention.mm_encoder_attention" + ".get_vit_attn_backend", + return_value=AttentionBackendEnum.FLASHINFER, + ), + ): + attn = MMEncoderAttention( + num_heads=NUM_HEADS, + head_size=HEAD_DIM, + prefix=LAYER_0, + ) + with pytest.raises(ValueError, match="scales not found for layer"): + attn.process_weights_after_loading(torch.bfloat16) + + +def test_dynamic_scales_auto_save(tmp_path) -> None: + """Verify scales are saved to disk after the amax buffer fills.""" + import vllm.model_executor.layers.attention.mm_encoder_attention as _mod + from vllm.config.multimodal import MultiModalConfig + + if not is_flashinfer_cudnn_fp8_prefill_attn_supported(): + pytest.skip("FlashInfer cuDNN not available") + + # Reset module-level state between runs (other tests may have left + # state behind after triggering a save). + _mod._fp8_scale_save_path = None + _mod._fp8_saved_scale_refs.clear() + + save_file = tmp_path / "auto_scales.json" + with _build_attention( + MultiModalConfig( + mm_encoder_attn_dtype="fp8", + mm_encoder_fp8_scale_save_path=str(save_file), + ) + ) as attn: + if attn is None or not attn.fp8_enabled: + pytest.skip("FP8 attention not available") + + attn = attn.to("cuda") + S, H, D = 16, NUM_HEADS, HEAD_DIM + + # Run exactly _FP8_AMAX_HISTORY_LEN forward passes. + for i in range(_FP8_AMAX_HISTORY_LEN): + mag = float(i + 1) + q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16) + k = torch.full((S, H, D), mag * 0.5, device="cuda", dtype=torch.bfloat16) + v = torch.full((S, H, D), mag * 0.3, device="cuda", dtype=torch.bfloat16) + attn._record_amax_and_update_scales(q, k, v) + + # File should have been written on the 16th call (buffer wrap). + assert save_file.is_file(), "Scale file was not saved" + scales = json.loads(save_file.read_text()) + assert LAYER_0 in scales + assert set(scales[LAYER_0].keys()) == {"q", "k", "v"} + for val in scales[LAYER_0].values(): + assert isinstance(val, float) and val > 0 + + # Path is cleared after the one-shot save fires. + assert _mod._fp8_scale_save_path is None diff --git a/vllm/config/model.py b/vllm/config/model.py index 76cd734c75a..b032302c04c 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -326,6 +326,10 @@ class ModelConfig: mm_encoder_only: InitVar[bool | None] = None mm_encoder_tp_mode: InitVar[MMEncoderTPMode | None] = None mm_encoder_attn_backend: InitVar[AttentionBackendEnum | str | None] = None + mm_encoder_attn_dtype: InitVar[str | None] = None + mm_encoder_fp8_scale_path: InitVar[str | None] = None + mm_encoder_fp8_scale_save_path: InitVar[str | None] = None + mm_encoder_fp8_scale_save_margin: InitVar[float | None] = None interleave_mm_strings: InitVar[bool | None] = None skip_mm_profiling: InitVar[bool | None] = None video_pruning_rate: InitVar[float | None] = None @@ -447,6 +451,10 @@ class ModelConfig: mm_encoder_only: bool | None, mm_encoder_tp_mode: MMEncoderTPMode | None, mm_encoder_attn_backend: AttentionBackendEnum | str | None, + mm_encoder_attn_dtype: str | None, + mm_encoder_fp8_scale_path: str | None, + mm_encoder_fp8_scale_save_path: str | None, + mm_encoder_fp8_scale_save_margin: float | None, interleave_mm_strings: bool | None, skip_mm_profiling: bool | None, video_pruning_rate: float | None, @@ -643,6 +651,10 @@ class ModelConfig: mm_encoder_only=mm_encoder_only, mm_encoder_tp_mode=mm_encoder_tp_mode, mm_encoder_attn_backend=mm_encoder_attn_backend, + mm_encoder_attn_dtype=mm_encoder_attn_dtype, + mm_encoder_fp8_scale_path=mm_encoder_fp8_scale_path, + mm_encoder_fp8_scale_save_path=mm_encoder_fp8_scale_save_path, + mm_encoder_fp8_scale_save_margin=mm_encoder_fp8_scale_save_margin, interleave_mm_strings=interleave_mm_strings, skip_mm_profiling=skip_mm_profiling, video_pruning_rate=video_pruning_rate, diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py index e66511c92ab..56333b1116c 100644 --- a/vllm/config/multimodal.py +++ b/vllm/config/multimodal.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Mapping +from pathlib import Path from typing import Any, Literal, TypeAlias, TypedDict, final from pydantic import ConfigDict, Field, field_validator, model_validator @@ -158,6 +159,24 @@ class MultiModalConfig: """Optional override for the multi-modal encoder attention backend when using vision transformers. Accepts any value from `vllm.v1.attention.backends.registry.AttentionBackendEnum` (e.g. `FLASH_ATTN`).""" + mm_encoder_attn_dtype: Literal["fp8"] | None = None + """Optional dtype override for ViT encoder attention. Set to `"fp8"` to + enable FP8 quantization via the FlashInfer cuDNN backend. When set to + `"fp8"` without a scale file, dynamic scaling is used automatically. + See docs/features/quantization/fp8_vit_attn.md for details.""" + mm_encoder_fp8_scale_path: str | None = None + """Path to a JSON file containing per-layer FP8 Q/K/V scales for ViT + encoder attention. When provided (with `mm_encoder_attn_dtype="fp8"`), + static scaling is used. When omitted, dynamic scaling is used.""" + mm_encoder_fp8_scale_save_path: str | None = None + """When set with dynamic FP8 scaling (`mm_encoder_attn_dtype="fp8"` + and no `mm_encoder_fp8_scale_path`), saves the calibrated scales to + this file after the amax history buffer is full. The saved file can + then be used as `mm_encoder_fp8_scale_path` in subsequent runs.""" + mm_encoder_fp8_scale_save_margin: float = Field(default=1.5, gt=0.0) + """Safety margin multiplied onto scales when auto-saving. A value > 1 + leaves headroom so that inputs with larger activations than the + calibration set do not overflow FP8 range. Default 1.5.""" interleave_mm_strings: bool = False """Enable fully interleaved support for multimodal prompts, while using --chat-template-content-format=string.""" @@ -233,6 +252,36 @@ class MultiModalConfig: "'mm_shm_cache_max_object_size_mb' should only be set when " "'mm_processor_cache_type' is 'shm'." ) + # Validate FP8 scale path combinations. + if self.mm_encoder_attn_dtype != "fp8" and ( + self.mm_encoder_fp8_scale_path is not None + or self.mm_encoder_fp8_scale_save_path is not None + ): + raise ValueError( + "'mm_encoder_fp8_scale_path' and " + "'mm_encoder_fp8_scale_save_path' require " + "'mm_encoder_attn_dtype' to be 'fp8'." + ) + if ( + self.mm_encoder_fp8_scale_path is not None + and self.mm_encoder_fp8_scale_save_path is not None + ): + raise ValueError( + "'mm_encoder_fp8_scale_save_path' cannot be used with " + "'mm_encoder_fp8_scale_path' (saving requires dynamic scaling)." + ) + + # Validate file paths exist. + if self.mm_encoder_fp8_scale_path is not None: + scale_path = Path(self.mm_encoder_fp8_scale_path) + if not scale_path.is_file(): + raise FileNotFoundError(f"FP8 scale file not found: {scale_path}") + if self.mm_encoder_fp8_scale_save_path is not None: + save_parent = Path(self.mm_encoder_fp8_scale_save_path).parent + if not save_parent.is_dir(): + raise FileNotFoundError( + f"Parent directory for FP8 scale save path not found: {save_parent}" + ) return self def compute_hash(self) -> str: @@ -252,6 +301,8 @@ class MultiModalConfig: if self.mm_encoder_attn_backend is not None else None, self.mm_encoder_tp_mode, + self.mm_encoder_attn_dtype, + self.mm_encoder_fp8_scale_path, ] hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest() return hash_str diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index ef3a9a982a5..cd955100333 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -542,6 +542,14 @@ class EngineArgs: mm_encoder_attn_backend: AttentionBackendEnum | str | None = ( MultiModalConfig.mm_encoder_attn_backend ) + mm_encoder_attn_dtype: str | None = MultiModalConfig.mm_encoder_attn_dtype + mm_encoder_fp8_scale_path: str | None = MultiModalConfig.mm_encoder_fp8_scale_path + mm_encoder_fp8_scale_save_path: str | None = ( + MultiModalConfig.mm_encoder_fp8_scale_save_path + ) + mm_encoder_fp8_scale_save_margin: float = ( + MultiModalConfig.mm_encoder_fp8_scale_save_margin + ) io_processor_plugin: str | None = None renderer_num_workers: int = 1 skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling @@ -1179,6 +1187,22 @@ class EngineArgs: "--mm-encoder-attn-backend", **multimodal_kwargs["mm_encoder_attn_backend"], ) + multimodal_group.add_argument( + "--mm-encoder-attn-dtype", + **multimodal_kwargs["mm_encoder_attn_dtype"], + ) + multimodal_group.add_argument( + "--mm-encoder-fp8-scale-path", + **multimodal_kwargs["mm_encoder_fp8_scale_path"], + ) + multimodal_group.add_argument( + "--mm-encoder-fp8-scale-save-path", + **multimodal_kwargs["mm_encoder_fp8_scale_save_path"], + ) + multimodal_group.add_argument( + "--mm-encoder-fp8-scale-save-margin", + **multimodal_kwargs["mm_encoder_fp8_scale_save_margin"], + ) multimodal_group.add_argument( "--interleave-mm-strings", **multimodal_kwargs["interleave_mm_strings"] ) @@ -1517,6 +1541,10 @@ class EngineArgs: mm_encoder_only=self.mm_encoder_only, mm_encoder_tp_mode=self.mm_encoder_tp_mode, mm_encoder_attn_backend=self.mm_encoder_attn_backend, + mm_encoder_attn_dtype=self.mm_encoder_attn_dtype, + mm_encoder_fp8_scale_path=self.mm_encoder_fp8_scale_path, + mm_encoder_fp8_scale_save_path=self.mm_encoder_fp8_scale_save_path, + mm_encoder_fp8_scale_save_margin=self.mm_encoder_fp8_scale_save_margin, pooler_config=self.pooler_config, generation_config=self.generation_config, override_generation_config=self.override_generation_config, diff --git a/vllm/kernels/triton/__init__.py b/vllm/kernels/triton/__init__.py new file mode 100644 index 00000000000..6626213145c --- /dev/null +++ b/vllm/kernels/triton/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton kernel implementations.""" diff --git a/vllm/kernels/triton/qkv_padded_fp8_quant.py b/vllm/kernels/triton/qkv_padded_fp8_quant.py new file mode 100644 index 00000000000..74dfe104363 --- /dev/null +++ b/vllm/kernels/triton/qkv_padded_fp8_quant.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stride-aware FP8 quantization with head_dim padding for ViT attention. + +Reads directly from non-contiguous QKV views using 3D strides and pads +head_dim to a multiple of 16 for cuDNN compatibility. +""" + +import torch + +from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_fp8_min_max, +) +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton +from vllm.utils.math_utils import round_up + +_FP8_MIN, _FP8_MAX = get_fp8_min_max() + + +@triton.jit +def _quantize_pad_fp8_kernel( + x_ptr, + y_ptr, + scale_ptr, + stride_xs, + stride_xh, + stride_xd, + stride_ys, + stride_yh, + stride_yd, + num_heads, + n_rows, + n_cols, + n_cols_padded, + fp8_min, + fp8_max, + SKIP_SCALE: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask_m = offs_m < n_rows + mask_out = mask_m[:, None] & (offs_n[None, :] < n_cols_padded) + mask_in = mask_m[:, None] & (offs_n[None, :] < n_cols) + + # Decompose flattened row into (token, head) for 3D stride indexing. + s = offs_m // num_heads + h = offs_m % num_heads + + x_ptrs = ( + x_ptr + + s[:, None] * stride_xs + + h[:, None] * stride_xh + + offs_n[None, :] * stride_xd + ) + x = tl.load(x_ptrs, mask=mask_in, other=0.0).to(tl.float32) + if SKIP_SCALE: + x_q = x + else: + scale = tl.load(scale_ptr) + x_q = x / scale + x_q = tl.clamp(x_q, fp8_min, fp8_max).to(y_ptr.dtype.element_ty) + + y_ptrs = ( + y_ptr + + s[:, None] * stride_ys + + h[:, None] * stride_yh + + offs_n[None, :] * stride_yd + ) + tl.store(y_ptrs, x_q, mask=mask_out) + + +def _get_fp8_pad_quant_config(padded_head_dim: int) -> tuple[int, int, int]: + block_n = triton.next_power_of_2(padded_head_dim) + block_n = max(16, min(block_n, 128)) + block_m = 16 + num_warps = 4 + return block_m, block_n, num_warps + + +def quantize_fp8_pad_head_dim_triton( + tensor: torch.Tensor, + scale: torch.Tensor, + skip_scale: bool = False, + block_m: int | None = None, + block_n: int | None = None, + num_warps: int | None = None, +) -> torch.Tensor: + """Quantize a 3D/4D tensor to FP8, padding head_dim to a multiple of 16. + + Reads directly from the input using its 3D strides, so non-contiguous + views (e.g. Q/K/V slices from an interleaved QKV buffer) are handled + without an extra copy. Output is always a fresh contiguous tensor + with shape (S, H, padded_D). + """ + if not HAS_TRITON: + raise RuntimeError("Triton is required to quantize with head_dim padding.") + + original_shape = tensor.shape + if tensor.dim() == 4: + tensor = tensor.view(-1, tensor.shape[-2], tensor.shape[-1]) + assert tensor.dim() == 3, f"Expected 3D input (S, H, D), got {tensor.dim()}D" + S, H, D = tensor.shape + padded_head_dim = round_up(D, 16) + out_dtype = current_platform.fp8_dtype() + output = torch.empty( + (S, H, padded_head_dim), + device=tensor.device, + dtype=out_dtype, + ) + + scale_1d = scale.reshape(-1) + n_rows = S * H + + if block_m is None or block_n is None or num_warps is None: + block_m, block_n, num_warps = _get_fp8_pad_quant_config(padded_head_dim) + + grid = ( + triton.cdiv(n_rows, block_m), + triton.cdiv(padded_head_dim, block_n), + ) + + _quantize_pad_fp8_kernel[grid]( + tensor, + output, + scale_1d, + tensor.stride(0), + tensor.stride(1), + tensor.stride(2), + output.stride(0), + output.stride(1), + output.stride(2), + H, + n_rows, + D, + padded_head_dim, + _FP8_MIN, + _FP8_MAX, + SKIP_SCALE=skip_scale, + BLOCK_M=block_m, + BLOCK_N=block_n, + num_warps=num_warps, + ) + + return output.view((*original_shape[:-1], padded_head_dim)) + + +def quantize_fp8_maybe_pad_head_dim( + tensor: torch.Tensor, + scale: torch.Tensor, + fp8_quant: QuantFP8, + skip_scale: bool = False, +) -> torch.Tensor: + """Quantize a 3D/4D tensor to FP8, padding head_dim to a multiple of 16 + only when needed. + + Accepts (S, H, D) or (B, S, H, D) input. Uses ``fp8_quant`` (a + :class:`QuantFP8` CustomOp) when head_dim is already aligned to 16 + (no padding); otherwise falls back to a stride-aware Triton kernel + that pads head_dim to a multiple of 16. + """ + head_dim = tensor.shape[-1] + if head_dim % 16 != 0: + return quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=skip_scale) + + if skip_scale: + return tensor.to(current_platform.fp8_dtype()) + + # QuantFP8 expects 2D: flatten all dims except (H, D). + orig_shape = tensor.shape + total_tokens = tensor.numel() // (orig_shape[-1] * orig_shape[-2]) + tensor_2d = tensor.reshape(total_tokens, -1) + fp8_tensor, _ = fp8_quant(tensor_2d, scale=scale) + return fp8_tensor.reshape(orig_shape) diff --git a/vllm/model_executor/layers/attention/mm_encoder_attention.py b/vllm/model_executor/layers/attention/mm_encoder_attention.py index 46d461c38b3..1731cc26bc3 100644 --- a/vllm/model_executor/layers/attention/mm_encoder_attention.py +++ b/vllm/model_executor/layers/attention/mm_encoder_attention.py @@ -1,13 +1,32 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import functools +import json import numpy as np import torch +from vllm.config import MultiModalConfig +from vllm.kernels.triton.qkv_padded_fp8_quant import ( + quantize_fp8_maybe_pad_head_dim, +) from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp, maybe_get_oot_by_class -from vllm.model_executor.models.vision import get_vit_attn_backend +from vllm.model_executor.layers.quantization.input_quant_fp8 import ( + QuantFP8, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + get_fp8_min_max, +) +from vllm.model_executor.models.vision import ( + get_multimodal_config, + get_vit_attn_backend, +) +from vllm.utils.flashinfer import ( + is_flashinfer_cudnn_fp8_prefill_attn_supported, +) from vllm.utils.math_utils import round_up from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -20,6 +39,108 @@ from vllm.v1.attention.ops.vit_attn_wrappers import ( logger = init_logger(__name__) +_, _FP8_MAX = get_fp8_min_max() +_FP8_AMAX_HISTORY_LEN = 16 + +# Module-level state for auto-saving dynamic scales. The save is a one-shot +# triggered by the first layer whose amax buffer wraps. Path and margin are +# captured during layer init (set_current_vllm_config context only lives +# across model init, not forward passes). +_fp8_scale_save_path: str | None = None +_fp8_scale_save_margin: float = MultiModalConfig.mm_encoder_fp8_scale_save_margin +_fp8_saved_scale_refs: dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {} + + +@functools.cache +def _load_fp8_scales_file(path: str | None) -> dict[str, dict[str, float]]: + """Load per-layer FP8 Q/K/V scales from a JSON file. Results are cached. + + Expected format (keys ``q_scale`` / ``k_scale`` / ``v_scale`` also accepted):: + + { + "visual.blocks.0.attn.attn": {"q": 224.0, "k": 198.0, "v": 210.0}, + "visual.blocks.1.attn.attn": {"q": 218.0, "k": 195.0, "v": 207.0}, + } + + To produce such a file, run with ``mm_encoder_fp8_scale_save_path`` set. + """ + if path is None: + return {} + + with open(path, encoding="utf-8") as f: + data = json.load(f) + + # Handle nested "layers" format + if "layers" in data and isinstance(data["layers"], dict): + data = data["layers"] + + scales: dict[str, dict[str, float]] = {} + for layer_name, layer_scales in data.items(): + if not isinstance(layer_scales, dict): + continue + q = layer_scales.get("q", layer_scales.get("q_scale")) + k = layer_scales.get("k", layer_scales.get("k_scale")) + v = layer_scales.get("v", layer_scales.get("v_scale")) + if q is not None and k is not None and v is not None: + q_f, k_f, v_f = float(q), float(k), float(v) + if q_f <= 0 or k_f <= 0 or v_f <= 0: + raise ValueError( + f"FP8 scales must be positive, got q={q_f}, " + f"k={k_f}, v={v_f} for layer '{layer_name}'" + ) + scales[layer_name] = {"q": q_f, "k": k_f, "v": v_f} + + logger.info_once( + "Loaded FP8 attention scales from %s (%d layers)", path, len(scales) + ) + return scales + + +def _maybe_save_fp8_scales( + layer_name: str, + q_scale: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + buffer_wrapped: bool, +) -> None: + """Accumulate a layer's scale tensors; on the first amax buffer wrap, + dump all accumulated scales to ``mm_encoder_fp8_scale_save_path``. + + No-op unless auto-save is configured. Tensor references are stored on + every call (no GPU->CPU sync); ``.item()`` is only called at the single + save point to avoid stalling the forward path. + """ + global _fp8_scale_save_path + # Fast path: auto-save either disabled or already finished. Path is + # captured at layer init and cleared once the save fires. + if _fp8_scale_save_path is None: + return + + # Stash scale tensor refs (no GPU->CPU sync yet); wait until the amax + # history has seen a full cycle before committing scales to disk. + _fp8_saved_scale_refs[layer_name] = (q_scale, k_scale, v_scale) + if not buffer_wrapped: + return + + # Buffer just wrapped for the first time: materialize scales (with + # safety margin) and dump to disk. Clearing _fp8_scale_save_path + # makes this a one-shot across all layers. + path, margin = _fp8_scale_save_path, _fp8_scale_save_margin + scales = { + name: { + "q": q.item() * margin, + "k": k.item() * margin, + "v": v.item() * margin, + } + for name, (q, k, v) in _fp8_saved_scale_refs.items() + } + _fp8_scale_save_path = None + _fp8_saved_scale_refs.clear() + with open(path, "w", encoding="utf-8") as f: + json.dump(scales, f, indent=2) + logger.info("Saved FP8 scales (%d layers) to %s", len(scales), path) + + # Batch buckets for cuDNN graph caching. # Graphs use batch size and max sequence length as cache key. # This avoids creating a new graph for each unique set of @@ -148,27 +269,47 @@ class MMEncoderAttention(CustomOp): hidden_size: int, tp_size: int, device: torch.device, + fp8_padded_hidden_size: int | None = None, ) -> torch.Tensor: if (oot_class := maybe_get_oot_by_class(cls)) is not cls: return oot_class.maybe_recompute_cu_seqlens( # type: ignore[attr-defined] - attn_backend, cu_seqlens, hidden_size, tp_size, device + attn_backend, + cu_seqlens, + hidden_size, + tp_size, + device, + fp8_padded_hidden_size=fp8_padded_hidden_size, ) if attn_backend == AttentionBackendEnum.FLASHINFER: batch_size = len(cu_seqlens) - 1 - scale = hidden_size // tp_size - cu_seqlens = cu_seqlens * scale - cu_seqlens_qko = cu_seqlens - cu_seqlens_v = cu_seqlens * 3 + if fp8_padded_hidden_size is not None: + # FP8 path: after quantization Q/K/V are each independent + # contiguous tensors with stride H * padded_D per token. + # All sections use the same element stride. + scale = fp8_padded_hidden_size // tp_size + cu_seqlens = cu_seqlens * scale + cu_seqlens_padded = add_padding_to_seqlens( + cu_seqlens, batch_size, cu_seqlens[-1] + ) + cu_seqlens = np.concatenate([cu_seqlens_padded, cu_seqlens_padded]) + else: + # BF16 path: Q/K/V are non-contiguous views into shared + # buffers. V section has 3x stride from interleaved QKV. + scale = hidden_size // tp_size + cu_seqlens = cu_seqlens * scale - cu_seqlens_qko = add_padding_to_seqlens( - cu_seqlens_qko, batch_size, cu_seqlens_qko[-1] - ) - cu_seqlens_v = add_padding_to_seqlens( - cu_seqlens_v, batch_size, cu_seqlens_v[-1] - ) - cu_seqlens = np.concatenate([cu_seqlens_qko, cu_seqlens_v]) + cu_seqlens_qko = cu_seqlens + cu_seqlens_v = cu_seqlens * 3 + + cu_seqlens_qko = add_padding_to_seqlens( + cu_seqlens_qko, batch_size, cu_seqlens_qko[-1] + ) + cu_seqlens_v = add_padding_to_seqlens( + cu_seqlens_v, batch_size, cu_seqlens_v[-1] + ) + cu_seqlens = np.concatenate([cu_seqlens_qko, cu_seqlens_v]) cu_seqlens = torch.from_numpy(cu_seqlens).to(device, non_blocking=True) return cu_seqlens @@ -206,6 +347,7 @@ class MMEncoderAttention(CustomOp): # During model initialization, the default dtype is set as the model # weight and activation dtype. dtype = torch.get_default_dtype() + self.dtype = dtype # Get device-specific vision attention backend. self.attn_backend = get_vit_attn_backend( @@ -229,6 +371,113 @@ class MMEncoderAttention(CustomOp): logger.info_once(f"Using {self.attn_backend} for MMEncoderAttention.") + self._init_fp8_state() + + def _init_fp8_state(self) -> None: + """Initialize FP8 attention state from multimodal config. + + No-op if FP8 is not requested. Raises ``ValueError`` if FP8 is + requested but the platform does not support it. + """ + # Populate defaults so ``_forward_flashinfer`` can + # check ``self.fp8_enabled`` and others without AttributeError. + self.fp8_enabled = False + self._fp8_dynamic_scale = False + self.fp8_quant: QuantFP8 | None = None + self.skip_scale_q = False + self.skip_scale_k = False + self.skip_scale_v = False + + mm_cfg = get_multimodal_config() + if mm_cfg is None or mm_cfg.mm_encoder_attn_dtype != "fp8": + return + + # FP8 path + if not is_flashinfer_cudnn_fp8_prefill_attn_supported(): + raise ValueError( + "mm_encoder_attn_dtype='fp8' requires the FlashInfer " + "cuDNN backend with cuDNN >= 9.17.1 on a GPU with native " + "FP8 support." + ) + + self.fp8_enabled = True + self._fp8_dynamic_scale = mm_cfg.mm_encoder_fp8_scale_path is None + self.fp8_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR) + + # Register buffers pre-device-move; values populated in + # process_weights_after_loading. Shape (1, 1, 1, 1) is required by cuDNN. + for attr in ("_fp8_q_scale", "_fp8_k_scale", "_fp8_v_scale"): + self.register_buffer( + attr, torch.ones(1, dtype=torch.float32).view(1, 1, 1, 1) + ) + if self._fp8_dynamic_scale: + for attr in ("_fp8_q_amax", "_fp8_k_amax", "_fp8_v_amax"): + self.register_buffer( + attr, + torch.zeros(_FP8_AMAX_HISTORY_LEN, dtype=torch.float32), + persistent=False, + ) + self._fp8_amax_pos = 0 + + # Capture auto-save config now: the VllmConfig context only lives + # across model init, not forward passes, so ``_maybe_save_fp8_scales`` + # reads these globals instead of re-querying ``get_multimodal_config``. + if ( + mm_cfg.mm_encoder_fp8_scale_save_path is not None + and self._fp8_dynamic_scale + ): + global _fp8_scale_save_path, _fp8_scale_save_margin + _fp8_scale_save_path = mm_cfg.mm_encoder_fp8_scale_save_path + _fp8_scale_save_margin = mm_cfg.mm_encoder_fp8_scale_save_margin + + def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: + """Populate FP8 scale buffers after weights are loaded. + + ``act_dtype`` matches the signature used by :class:`Attention` and + :class:`MLAAttention` for the loader auto-scan but is unused: + FP8 scales are always float32. + """ + if not self.fp8_enabled: + return + + mm_cfg = get_multimodal_config() + scale_path = mm_cfg.mm_encoder_fp8_scale_path if mm_cfg is not None else None + if scale_path is None: + logger.info_once( + "FP8 attention enabled with dynamic scaling " + "(no scale file provided). Scales will adapt from " + "observed Q/K/V amax values (history_len=%d).", + _FP8_AMAX_HISTORY_LEN, + ) + return + + all_scales = _load_fp8_scales_file(scale_path) + layer_scales = all_scales.get(self.layer_name) + if layer_scales is None: + raise ValueError( + "FP8 attention enabled but scales not found for layer " + f"'{self.layer_name}' in {scale_path}. " + f"Available layers: {list(all_scales.keys())}" + ) + + for attr, key in ( + ("_fp8_q_scale", "q"), + ("_fp8_k_scale", "k"), + ("_fp8_v_scale", "v"), + ): + getattr(self, attr).fill_(layer_scales[key]) + self.skip_scale_q = layer_scales["q"] == 1.0 + self.skip_scale_k = layer_scales["k"] == 1.0 + self.skip_scale_v = layer_scales["v"] == 1.0 + + logger.debug( + "FP8 attention enabled for %s: q=%.4f, k=%.4f, v=%.4f", + self.layer_name if self.layer_name else "MMEncoderAttention", + layer_scales["q"], + layer_scales["k"], + layer_scales["v"], + ) + @classmethod def enabled(cls) -> bool: return True @@ -353,6 +602,44 @@ class MMEncoderAttention(CustomOp): output = output.reshape(bsz, q_len, -1) return output + @torch.no_grad() + def _record_amax_and_update_scales( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ) -> None: + """Record Q/K/V amax into circular history and recompute scales. + + All work stays on GPU with no device-to-host sync. The Python-side + history position counter is mutated, so this method must NOT be + called inside CUDA graph capture/replay. When CUDA graphs are + used for the encoder, dynamic scaling should be disabled by + providing a static scale file via --mm-encoder-fp8-scale-path. + """ + pos = self._fp8_amax_pos + self._fp8_amax_pos = (pos + 1) % _FP8_AMAX_HISTORY_LEN + + for tensor, amax_buf, scale_buf in ( + (query, self._fp8_q_amax, self._fp8_q_scale), + (key, self._fp8_k_amax, self._fp8_k_scale), + (value, self._fp8_v_amax, self._fp8_v_scale), + ): + amax_buf[pos] = tensor.amax() + max_amax = amax_buf.max() + scale_buf.fill_( + torch.clamp(max_amax, min=torch.finfo(torch.float32).tiny) / _FP8_MAX + ) + + buffer_wrapped = self._fp8_amax_pos == 0 and pos == _FP8_AMAX_HISTORY_LEN - 1 + _maybe_save_fp8_scales( + self.layer_name, + self._fp8_q_scale, + self._fp8_k_scale, + self._fp8_v_scale, + buffer_wrapped, + ) + def _forward_flashinfer( self, query: torch.Tensor, @@ -363,7 +650,32 @@ class MMEncoderAttention(CustomOp): sequence_lengths: torch.Tensor | None = None, # Only used for FlashInfer CuDNN backend ) -> torch.Tensor: - return vit_flashinfer_wrapper( + if self.fp8_enabled: + assert self.fp8_quant is not None + + if self._fp8_dynamic_scale: + self._record_amax_and_update_scales(query, key, value) + + query = quantize_fp8_maybe_pad_head_dim( + query, + self._fp8_q_scale, + skip_scale=self.skip_scale_q, + fp8_quant=self.fp8_quant, + ) + key = quantize_fp8_maybe_pad_head_dim( + key, + self._fp8_k_scale, + skip_scale=self.skip_scale_k, + fp8_quant=self.fp8_quant, + ) + value = quantize_fp8_maybe_pad_head_dim( + value, + self._fp8_v_scale, + skip_scale=self.skip_scale_v, + fp8_quant=self.fp8_quant, + ) + + output = vit_flashinfer_wrapper( q=query, k=key, v=value, @@ -372,8 +684,17 @@ class MMEncoderAttention(CustomOp): cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, sequence_lengths=sequence_lengths, + q_scale=self._fp8_q_scale if self.fp8_enabled else None, + k_scale=self._fp8_k_scale if self.fp8_enabled else None, + v_scale=self._fp8_v_scale if self.fp8_enabled else None, + o_data_type=self.dtype if self.fp8_enabled else None, ) + if self.fp8_enabled and output.shape[-1] != self.head_size: + output = output[..., : self.head_size].contiguous() + + return output + def forward_native( self, query: torch.Tensor, diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 4ee9e90d741..2a5f746d783 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -15,7 +15,11 @@ from typing_extensions import assert_never import vllm.envs as envs from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.logger import init_logger -from vllm.model_executor.layers.attention import Attention, MLAAttention +from vllm.model_executor.layers.attention import ( + Attention, + MLAAttention, + MMEncoderAttention, +) from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, @@ -106,12 +110,12 @@ def process_weights_after_loading( with device_loading_context(module, target_device): quant_method.process_weights_after_loading(module) - # Initialize post-load attention weights for both Attention and MLA. + # Initialize post-load attention weights for Attention, MLA, and MM encoder. # NOTE: Happens after other modules so we can easily decompress weights. for _, module in model.named_modules(): - if isinstance(module, (Attention, MLAAttention)) and hasattr( - module, "process_weights_after_loading" - ): + if isinstance( + module, (Attention, MLAAttention, MMEncoderAttention) + ) and hasattr(module, "process_weights_after_loading"): # TODO(lucas): see if there is a way to unify the signatures # of process_weights_after_loading with device_loading_context(module, target_device): diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 2e1ed9f649c..4a0de15f02a 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -136,6 +136,7 @@ from .utils import ( maybe_prefix, ) from .vision import ( + get_fp8_padded_hidden_size, get_vit_attn_backend, is_vit_use_data_parallel, run_dp_sharded_mrope_vision_model, @@ -562,6 +563,13 @@ class Qwen3_VisionTransformer(nn.Module): norm_layer = partial(nn.LayerNorm, eps=norm_eps) head_dim = self.hidden_size // self.num_heads + + # FP8 attention: Q/K/V become independent contiguous tensors + # after quantization, so cu_seqlens uses uniform stride (no 3x V). + self.fp8_padded_hidden_size = get_fp8_padded_hidden_size( + self.num_heads, head_dim + ) + self.rotary_pos_emb = get_rope( head_size=head_dim, max_position=8192, @@ -776,6 +784,7 @@ class Qwen3_VisionTransformer(nn.Module): self.hidden_size, self.tp_size, device, + fp8_padded_hidden_size=self.fp8_padded_hidden_size, ) return metadata diff --git a/vllm/model_executor/models/vision.py b/vllm/model_executor/models/vision.py index e6a24300675..0582c125c66 100644 --- a/vllm/model_executor/models/vision.py +++ b/vllm/model_executor/models/vision.py @@ -10,7 +10,7 @@ from typing import Final, Generic, Literal, Protocol, TypeAlias, TypeVar import torch from transformers import PretrainedConfig -from vllm.config import MultiModalConfig, VllmConfig, get_current_vllm_config +from vllm.config import MultiModalConfig, get_current_vllm_config_or_none from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -18,6 +18,7 @@ from vllm.distributed import ( ) from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.utils.math_utils import round_up from vllm.v1.attention.backends.registry import AttentionBackendEnum logger = init_logger(__name__) @@ -102,45 +103,48 @@ def get_vit_attn_backend( """ Get the attention backend for Vision Transformer. """ - try: - vllm_config: VllmConfig = get_current_vllm_config() - model_config = vllm_config.model_config - multimodal_config: MultiModalConfig | None = ( - model_config.multimodal_config if model_config is not None else None - ) - except (AssertionError, AttributeError): - multimodal_config = None - + mm_cfg = get_multimodal_config() attn_backend_override = ( - multimodal_config.mm_encoder_attn_backend - if multimodal_config is not None - else None + mm_cfg.mm_encoder_attn_backend if mm_cfg is not None else None ) - attn_backend = _get_vit_attn_backend( + return _get_vit_attn_backend( head_size, dtype, attn_backend_override=attn_backend_override, ) - return attn_backend + + +def get_multimodal_config() -> MultiModalConfig | None: + """Return the current ``MultiModalConfig``, or ``None`` when no engine + config context is active (e.g., during unit tests) or when the current + ``model_config`` does not carry a ``multimodal_config`` (e.g., minimal + stubs used in tests).""" + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None or vllm_config.model_config is None: + return None + return getattr(vllm_config.model_config, "multimodal_config", None) + + +def get_fp8_padded_hidden_size(num_heads: int, head_dim: int) -> int | None: + """Return the padded hidden size for FP8 ViT encoder attention, or + ``None`` when FP8 is not enabled. + + cuDNN FP8 prefill attention requires ``head_dim`` to be a multiple of + 16. For non-aligned ``head_dim`` (e.g. 72), Q/K/V are padded to the + nearest multiple of 16. + """ + mm_cfg = get_multimodal_config() + if mm_cfg is None or mm_cfg.mm_encoder_attn_dtype != "fp8": + return None + return num_heads * round_up(head_dim, 16) def is_vit_use_data_parallel(): """ Get the tensor parallel type for Vision Transformer. """ - try: - vllm_config: VllmConfig = get_current_vllm_config() - model_config = vllm_config.model_config - multimodal_config: MultiModalConfig | None = ( - model_config.multimodal_config if model_config is not None else None - ) - except (AssertionError, AttributeError): - multimodal_config = None - - mm_encoder_tp_mode = ( - multimodal_config.mm_encoder_tp_mode if multimodal_config is not None else None - ) - return mm_encoder_tp_mode == "data" + mm_cfg = get_multimodal_config() + return mm_cfg is not None and mm_cfg.mm_encoder_tp_mode == "data" VisionFeatureSelectStrategyStr = Literal["class", "default", "full"] diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 316816d9658..5672aef301e 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -772,6 +772,40 @@ def should_use_flashinfer_for_blockscale_fp8_gemm( return should_use_flashinfer +_MIN_CUDNN_FP8 = 91701 # cuDNN >= 9.17.1 required for FP8 attention + + +@functools.cache +def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: + """Check if FP8 ViT attention is supported on this platform. + + Requires native FP8 hardware support, the FlashInfer cuDNN backend, + and cuDNN >= 9.17.1. + """ + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + # cuDNN SDPA FP8 requires Hopper (SM 90) or newer. + if not current_platform.has_device_capability(90): + return False + + try: + supported = current_platform.get_supported_vit_attn_backends() + if AttentionBackendEnum.FLASHINFER not in supported: + return False + except (ImportError, AttributeError): + return False + + try: + import torch.backends.cudnn as cudnn + + if cudnn.is_available() and cudnn.version() < _MIN_CUDNN_FP8: + return False + except (ImportError, AttributeError): + pass + + return True + + __all__ = [ "has_flashinfer", "flashinfer_trtllm_fp8_block_scale_moe", @@ -803,4 +837,5 @@ __all__ = [ "flashinfer_fp8_blockscale_gemm", "should_use_flashinfer_for_blockscale_fp8_gemm", "is_flashinfer_fp8_blockscale_gemm_supported", + "is_flashinfer_cudnn_fp8_prefill_attn_supported", ] diff --git a/vllm/v1/attention/ops/vit_attn_wrappers.py b/vllm/v1/attention/ops/vit_attn_wrappers.py index 2d3505f51b7..4506f452cf9 100644 --- a/vllm/v1/attention/ops/vit_attn_wrappers.py +++ b/vllm/v1/attention/ops/vit_attn_wrappers.py @@ -279,6 +279,10 @@ def flashinfer_wrapper( cu_seqlens: torch.Tensor | None = None, max_seqlen: torch.Tensor | None = None, sequence_lengths: torch.Tensor | None = None, + q_scale: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + o_data_type: torch.dtype | None = None, ) -> torch.Tensor: from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache @@ -318,6 +322,10 @@ def flashinfer_wrapper( batch_offsets_k=batch_offsets_qko, batch_offsets_v=batch_offsets_v, batch_offsets_o=batch_offsets_qko, + q_scale=q_scale, + k_scale=k_scale, + v_scale=v_scale, + o_data_type=o_data_type, ) if is_reshaped: @@ -335,8 +343,12 @@ def vit_flashinfer_wrapper_fake( cu_seqlens: torch.Tensor | None = None, max_seqlen: torch.Tensor | None = None, sequence_lengths: torch.Tensor | None = None, + q_scale: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + o_data_type: torch.dtype | None = None, ) -> torch.Tensor: - return torch.empty_like(q) + return torch.empty_like(q, dtype=o_data_type or q.dtype) direct_register_custom_op( @@ -355,7 +367,22 @@ def vit_flashinfer_wrapper( cu_seqlens: torch.Tensor | None = None, max_seqlen: torch.Tensor | None = None, sequence_lengths: torch.Tensor | None = None, + q_scale: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + o_data_type: torch.dtype | None = None, ) -> torch.Tensor: return torch.ops.vllm.flashinfer_wrapper( - q, k, v, scale, workspace_buffer, cu_seqlens, max_seqlen, sequence_lengths + q, + k, + v, + scale, + workspace_buffer, + cu_seqlens, + max_seqlen, + sequence_lengths, + q_scale, + k_scale, + v_scale, + o_data_type, )