@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5d5c776444
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 15:00 UTC

[Perf] FP8 FlashInfer Attn for ViT (#38065)

Signed-off-by: Zhanda Zhu <zhandazhu@gmail.com> Co-authored-by: Yubo Gao <ybgao-nvidia@users.noreply.github.com>

Zhanda Zhu committed Apr 27, 2026 at 06:44 UTC 5d5c7764446da0d4888add9a060604e376e4e856
18 files changed +1830 -50
benchmarks/kernels/benchmark_vit_fp8_attn.py new
+324
@@ -0,0 +1,324 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Benchmarks FP8 vs BF16 ViT attention via FlashInfer cuDNN backend.
5 +#
6 +# == Usage Examples ==
7 +#
8 +# Benchmark mode (default, FlashInfer CUDAGraph Bench)
9 +# python3 benchmark_vit_fp8_attn.py
10 +#
11 +# Profile mode (PyTorch profiler, saves TensorBoard traces):
12 +# python3 benchmark_vit_fp8_attn.py --profile
13 +# python3 benchmark_vit_fp8_attn.py --profile --profile-output-dir ./profile_traces
14 +#
15 +# Custom seq_lens:
16 +# python3 benchmark_vit_fp8_attn.py --seq-lens 4096 8192 16384
17 +
18 +from functools import partial
19 +
20 +import numpy as np
21 +import torch
22 +from torch.profiler import ProfilerActivity, profile, record_function
23 +
24 +from vllm.utils.argparse_utils import FlexibleArgumentParser
25 +
26 +# Qwen3-VL defaults
27 +NUM_HEADS = 16
28 +HEAD_DIM = 72
29 +DEFAULT_SEQ_LENS = [2304, 4096, 8192, 16384]
30 +
31 +
32 +def _setup_fp8_attention(num_heads: int, head_dim: int) -> tuple:
33 + """Create FP8 and BF16 attention modules + workspace."""
34 + from types import SimpleNamespace
35 + from unittest.mock import patch
36 +
37 + from vllm.config import VllmConfig, set_current_vllm_config
38 + from vllm.config.multimodal import MultiModalConfig
39 + from vllm.model_executor.layers.attention.mm_encoder_attention import (
40 + MMEncoderAttention,
41 + _get_flashinfer_workspace_buffer,
42 + )
43 + from vllm.v1.attention.backends.registry import AttentionBackendEnum
44 +
45 + old_dtype = torch.get_default_dtype()
46 + torch.set_default_dtype(torch.bfloat16)
47 +
48 + backend_patch = patch(
49 + "vllm.model_executor.layers.attention.mm_encoder_attention"
50 + ".get_vit_attn_backend",
51 + return_value=AttentionBackendEnum.FLASHINFER,
52 + )
53 +
54 + # FP8 attention
55 + mm_config_fp8 = MultiModalConfig(mm_encoder_attn_dtype="fp8")
56 + vllm_config_fp8 = VllmConfig()
57 + vllm_config_fp8.model_config = SimpleNamespace(multimodal_config=mm_config_fp8)
58 + with set_current_vllm_config(vllm_config_fp8), backend_patch:
59 + attn_fp8 = MMEncoderAttention(
60 + num_heads=num_heads,
61 + head_size=head_dim,
62 + prefix="visual.blocks.0.attn",
63 + ).to("cuda")
64 +
65 + # BF16 attention (no FP8)
66 + with set_current_vllm_config(VllmConfig()), backend_patch:
67 + attn_bf16 = MMEncoderAttention(
68 + num_heads=num_heads,
69 + head_size=head_dim,
70 + prefix="visual.blocks.0.attn",
71 + ).to("cuda")
72 +
73 + torch.set_default_dtype(old_dtype)
74 +
75 + workspace = _get_flashinfer_workspace_buffer()
76 + return attn_fp8, attn_bf16, workspace
77 +
78 +
79 +def _build_meta(
80 + seq_len: int,
81 + num_heads: int,
82 + head_dim: int,
83 + fp8: bool,
84 +):
85 + """Build cu_seqlens, max_seqlen, sequence_lengths."""
86 + from vllm.model_executor.layers.attention.mm_encoder_attention import (
87 + MMEncoderAttention,
88 + )
89 + from vllm.utils.math_utils import round_up
90 + from vllm.v1.attention.backends.registry import AttentionBackendEnum
91 +
92 + cu_np = np.array([0, seq_len], dtype=np.int32)
93 + fp8_padded = num_heads * round_up(head_dim, 16) if fp8 else None
94 +
95 + seq_lengths = MMEncoderAttention.maybe_compute_seq_lens(
96 + AttentionBackendEnum.FLASHINFER, cu_np, torch.device("cuda")
97 + )
98 + max_seqlen = torch.tensor(
99 + MMEncoderAttention.compute_max_seqlen(AttentionBackendEnum.FLASHINFER, cu_np),
100 + dtype=torch.int32,
101 + )
102 + cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
103 + AttentionBackendEnum.FLASHINFER,
104 + cu_np,
105 + num_heads * head_dim,
106 + 1,
107 + torch.device("cuda"),
108 + fp8_padded_hidden_size=fp8_padded,
109 + )
110 + return cu_seqlens, max_seqlen, seq_lengths
111 +
112 +
113 +def run_benchmark(
114 + seq_lens: list[int],
115 + num_heads: int,
116 + head_dim: int,
117 + method: str,
118 +):
119 + """Benchmark FP8 vs BF16 attention across seq_lens.
120 +
121 + Uses FlashInfer GPU-level timing to measure pure kernel time,
122 + excluding CPU launch overhead.
123 + """
124 + if method == "cupti":
125 + from flashinfer.testing import bench_gpu_time_with_cupti as bench_fn
126 +
127 + bench_fn = partial(bench_fn, use_cuda_graph=True, cold_l2_cache=False)
128 + elif method == "cudagraph":
129 + from flashinfer.testing import (
130 + bench_gpu_time_with_cudagraph as bench_fn,
131 + )
132 +
133 + bench_fn = partial(bench_fn, cold_l2_cache=False)
134 + else:
135 + raise ValueError(f"Invalid method: {method}")
136 +
137 + attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
138 +
139 + print(f"Timing method: {method}")
140 + print(f"{'seq_len':>8} {'BF16 (us)':>12} {'FP8 (us)':>12} {'Speedup':>10}")
141 + print("-" * 46)
142 +
143 + for seq_len in seq_lens:
144 + torch.manual_seed(42)
145 +
146 + q = torch.randn(
147 + seq_len,
148 + num_heads,
149 + head_dim,
150 + device="cuda",
151 + dtype=torch.bfloat16,
152 + )
153 + k = torch.randn_like(q)
154 + v = torch.randn_like(q)
155 +
156 + cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
157 + # we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
158 + cu_bf16 = cu_fp8.clone()
159 +
160 + def bf16_fn(q=q, k=k, v=v, cu=cu_bf16, ms=max_s, sl=seq_l):
161 + attn_bf16._forward_flashinfer(q, k, v, cu, ms, sl)
162 +
163 + def fp8_fn(q=q, k=k, v=v, cu=cu_fp8, ms=max_s, sl=seq_l):
164 + attn_fp8._forward_flashinfer(q, k, v, cu, ms, sl)
165 +
166 + # bench_fn returns List[float] of per-iteration times in ms
167 + bf16_times = bench_fn(bf16_fn)
168 + fp8_times = bench_fn(fp8_fn)
169 +
170 + bf16_us = np.median(bf16_times) * 1e3 # ms -> us
171 + fp8_us = np.median(fp8_times) * 1e3
172 + speedup = bf16_us / fp8_us if fp8_us > 0 else float("inf")
173 +
174 + print(f"{seq_len:>8} {bf16_us:>12.1f} {fp8_us:>12.1f} {speedup:>9.2f}x")
175 +
176 +
177 +def _make_trace_handler(output_dir: str, worker_name: str, label: str):
178 + """Create a trace handler that saves to TensorBoard and prints summary."""
179 +
180 + def handler(prof):
181 + torch.profiler.tensorboard_trace_handler(output_dir, worker_name)(prof)
182 + print(f"\n{'=' * 80}")
183 + print(label)
184 + print(f"{'=' * 80}")
185 + print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
186 +
187 + return handler
188 +
189 +
190 +def run_profile(
191 + seq_len: int,
192 + num_heads: int,
193 + head_dim: int,
194 + warmup: int,
195 + output_dir: str,
196 +):
197 + """Profile FP8 vs BF16 attention with PyTorch profiler."""
198 + attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
199 +
200 + torch.manual_seed(42)
201 + q = torch.randn(
202 + seq_len,
203 + num_heads,
204 + head_dim,
205 + device="cuda",
206 + dtype=torch.bfloat16,
207 + )
208 + k = torch.randn_like(q)
209 + v = torch.randn_like(q)
210 +
211 + cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
212 + # we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
213 + cu_bf16 = cu_fp8.clone()
214 +
215 + sched = torch.profiler.schedule(wait=0, warmup=warmup, active=1)
216 +
217 + # Profile BF16 (warmup handled by profiler schedule)
218 + with profile(
219 + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
220 + schedule=sched,
221 + on_trace_ready=_make_trace_handler(
222 + output_dir,
223 + f"bf16_h{head_dim}_s{seq_len}",
224 + f"BF16 Attention (seq_len={seq_len}, heads={num_heads}, "
225 + f"head_dim={head_dim})",
226 + ),
227 + ) as prof_bf16:
228 + for _ in range(warmup + 1):
229 + with record_function("bf16_attention"):
230 + attn_bf16._forward_flashinfer(
231 + q.clone(), k.clone(), v.clone(), cu_bf16, max_s, seq_l
232 + )
233 + torch.accelerator.synchronize()
234 + prof_bf16.step()
235 +
236 + # Profile FP8 (warmup handled by profiler schedule)
237 + with profile(
238 + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
239 + schedule=sched,
240 + on_trace_ready=_make_trace_handler(
241 + output_dir,
242 + f"fp8_h{head_dim}_s{seq_len}",
243 + f"FP8 Attention (seq_len={seq_len}, heads={num_heads}, "
244 + f"head_dim={head_dim})",
245 + ),
246 + ) as prof_fp8:
247 + for _ in range(warmup + 1):
248 + with record_function("fp8_attention"):
249 + attn_fp8._forward_flashinfer(
250 + q.clone(), k.clone(), v.clone(), cu_fp8, max_s, seq_l
251 + )
252 + torch.accelerator.synchronize()
253 + prof_fp8.step()
254 +
255 + print(f"\nTensorBoard traces saved to: {output_dir}")
256 + print(f"View with: tensorboard --logdir={output_dir}")
257 +
258 +
259 +if __name__ == "__main__":
260 + parser = FlexibleArgumentParser(description="Benchmark FP8 vs BF16 ViT attention.")
261 + parser.add_argument(
262 + "--seq-lens",
263 + type=int,
264 + nargs="+",
265 + default=DEFAULT_SEQ_LENS,
266 + help="Sequence lengths to benchmark",
267 + )
268 + parser.add_argument(
269 + "--num-heads",
270 + type=int,
271 + default=NUM_HEADS,
272 + )
273 + parser.add_argument(
274 + "--head-dim",
275 + type=int,
276 + default=HEAD_DIM,
277 + )
278 + parser.add_argument(
279 + "--method",
280 + choices=["cupti", "cudagraph"],
281 + default="cudagraph",
282 + help="GPU timing method: cupti (CUPTI kernel timing) or "
283 + "cudagraph (CUDA graph capture/replay). Default: cudagraph",
284 + )
285 + parser.add_argument(
286 + "--warmup",
287 + type=int,
288 + default=10,
289 + help="Warmup iterations (profile mode only)",
290 + )
291 + parser.add_argument(
292 + "--profile",
293 + action="store_true",
294 + help="Run PyTorch profiler instead of benchmark",
295 + )
296 + parser.add_argument(
297 + "--profile-seq-len",
298 + type=int,
299 + default=8192,
300 + help="Sequence length for profiling (default: 8192)",
301 + )
302 + parser.add_argument(
303 + "--profile-output-dir",
304 + type=str,
305 + default="./profile_traces",
306 + help="Output directory for TensorBoard traces (default: ./profile_traces)",
307 + )
308 + args = parser.parse_args()
309 +
310 + if args.profile:
311 + run_profile(
312 + args.profile_seq_len,
313 + args.num_heads,
314 + args.head_dim,
315 + args.warmup,
316 + args.profile_output_dir,
317 + )
318 + else:
319 + run_benchmark(
320 + args.seq_lens,
321 + args.num_heads,
322 + args.head_dim,
323 + args.method,
324 + )
docs/features/quantization/README.md
+1
@@ -20,6 +20,7 @@ The following are the supported quantization formats for vLLM:
20 - [AMD Quark](quark.md)
21 - [Quantized KV Cache](quantized_kvcache.md)
22 - [TorchAO](torchao.md)
23 +- [FP8 ViT Encoder Attention](fp8_vit_attn.md)
24
25 ## Supported Hardware
26
docs/features/quantization/fp8_vit_attn.md new
+109
@@ -0,0 +1,109 @@
1 +# FP8 ViT Encoder Attention
2 +
3 +For visual understanding workloads with large images (e.g. QHD, 4K) and relatively
4 +short text prompts/generation, the ViT encoder attention can become a significant
5 +bottleneck, especially when the text model is quantized (e.g. NVFP4). vLLM
6 +supports optional FP8 quantization for the ViT encoder attention via the
7 +FlashInfer cuDNN backend. Q/K/V are quantized on-the-fly to FP8 before the
8 +cuDNN attention call.
9 +
10 +!!! note
11 + - Currently supports Qwen3-VL family models only (`qwen3_vl`, `qwen3_vl_moe`,
12 + `qwen3_5`, `qwen3_5_moe`, and other models using Qwen3 ViT).
13 + - Dynamic scaling is not compatible with ViT full CUDA graphs.
14 + - Performance gains are mostly visible at QHD/4K resolutions or multi-image
15 + requests. Smaller images may see no speedup due to quantization overhead
16 + (3 quantization kernel launches + un-padding).
17 + - FP8 tensor-core speedup is more pronounced on GB300 than GB200.
18 +
19 +## Requirements
20 +
21 +- FlashInfer cuDNN backend with cuDNN >= 9.17.1.
22 +
23 +## Usage
24 +
25 +Enable FP8 ViT attention by passing `--mm-encoder-attn-dtype fp8` together
26 +with `--mm-encoder-attn-backend FLASHINFER`:
27 +
28 +```bash
29 +vllm serve $MODEL \
30 + --mm-encoder-attn-backend FLASHINFER \
31 + --mm-encoder-attn-dtype fp8
32 +```
33 +
34 +By default (no scale file), **dynamic scaling** is used: a 16-entry circular
35 +buffer of observed Q/K/V amax values drives per-forward scale updates. This
36 +matches BF16 accuracy without any calibration but adds a small per-forward
37 +overhead.
38 +
39 +## Calibrate-Once, Reuse Workflow (Recommended)
40 +
41 +For production, calibrate static scales on a representative dataset once and
42 +reuse them to avoid the dynamic overhead:
43 +
44 +```bash
45 +# Step 1: calibrate and save scales (runs dynamic scaling for 16 passes,
46 +# then dumps the learned scales to JSON).
47 +vllm bench mm-processor \
48 + --model $MODEL --mm-encoder-attn-backend FLASHINFER \
49 + --mm-encoder-attn-dtype fp8 \
50 + --mm-encoder-fp8-scale-save-path /path/to/scales.json \
51 + --dataset-name hf --dataset-path lmarena-ai/VisionArena-Chat \
52 + --num-prompts 100
53 +
54 +# Step 2: serve with static scales (no dynamic overhead).
55 +vllm serve $MODEL \
56 + --mm-encoder-attn-backend FLASHINFER \
57 + --mm-encoder-attn-dtype fp8 \
58 + --mm-encoder-fp8-scale-path /path/to/scales.json
59 +```
60 +
61 +Saved scales are multiplied by `--mm-encoder-fp8-scale-save-margin` (default
62 +`1.5`) to leave headroom against activation outliers not present in the
63 +calibration set. The default has been validated to generalize across datasets
64 +(e.g. VisionArena-Chat calibration maintains BF16 accuracy on ChartQA).
65 +
66 +## Scale File Format
67 +
68 +```json
69 +{
70 + "visual.blocks.0.attn.attn": {"q": 224.0, "k": 198.0, "v": 210.0},
71 + "visual.blocks.1.attn.attn": {"q": 218.0, "k": 195.0, "v": 207.0}
72 +}
73 +```
74 +
75 +Keys `q_scale` / `k_scale` / `v_scale` are accepted as aliases.
76 +
77 +## Performance
78 +
79 +**Core cuDNN attention kernel** (PyTorch profiler, `cudnn_generated_fort_native_sdpa_sm100_flash_fprop`, head_dim=128, seq_len=8192):
80 +
81 +| Hardware | BF16 | FP8 | Speedup |
82 +| -------- | ---- | ---- | ------- |
83 +| GB200 | 350 us | 312 us | **1.12x** |
84 +| GB300 | 300 us | 211 us | **1.42x** |
85 +
86 +**End-to-end encoder forward time** (Qwen3-VL-30B-A3B-Instruct on GB200, 3 images/request):
87 +
88 +| Resolution | BF16 median | FP8 median | Speedup |
89 +| ---------- | ----------- | ---------- | ------- |
90 +| HD (720x1280) | 31.77 ms | 36.39 ms | 0.87x |
91 +| FullHD (1080x1920) | 57.99 ms | 58.73 ms | ~same |
92 +| QHD (1440x2560) | 131.83 ms | 122.30 ms | **1.08x** |
93 +| 4K (2160x3840) | 543.44 ms | 460.31 ms | **1.18x** |
94 +
95 +Crossover is around FullHD with 3 images/request. At QHD and above, FP8 wins.
96 +
97 +## Accuracy
98 +
99 +ChartQA, Qwen3-VL-8B-Instruct, 500 samples. FP8 static uses scales calibrated
100 +on VisionArena-Chat (with default 1.5x margin):
101 +
102 +| Metric | BF16 | FP8 dynamic | FP8 static |
103 +| ------ | ---- | ----------- | ---------- |
104 +| relaxed_accuracy | 0.780 | 0.776 | 0.780 |
105 +| anywhere_accuracy | 0.806 | 0.816 | 0.814 |
106 +| exact_match | 0.584 | 0.582 | 0.578 |
107 +
108 +All three configurations match within statistical noise, confirming that
109 +static scales calibrated on one dataset generalize to another.
tests/config/test_multimodal_config.py
+18
@@ -41,3 +41,21 @@ def test_language_model_only_affects_model_hash():
41 base_hash = ModelConfig(model).compute_hash()
42 lm_only_hash = ModelConfig(model, language_model_only=True).compute_hash()
43 assert base_hash != lm_only_hash
44 +
45 +
46 +def test_mm_encoder_fp8_scale_path_requires_fp8():
47 + with pytest.raises(ValueError, match="mm_encoder_attn_dtype"):
48 + MultiModalConfig(mm_encoder_fp8_scale_path="/tmp/scales.json")
49 +
50 +
51 +def test_mm_encoder_attn_dtype_hash_updates(tmp_path):
52 + scale_file = tmp_path / "scales.json"
53 + scale_file.write_text("{}")
54 + base_hash = MultiModalConfig().compute_hash()
55 + fp8_hash = MultiModalConfig(mm_encoder_attn_dtype="fp8").compute_hash()
56 + fp8_static_hash = MultiModalConfig(
57 + mm_encoder_attn_dtype="fp8",
58 + mm_encoder_fp8_scale_path=str(scale_file),
59 + ).compute_hash()
60 + assert base_hash != fp8_hash
61 + assert fp8_hash != fp8_static_hash
tests/kernels/core/test_vit_fp8_attn.py new
+279
@@ -0,0 +1,279 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the full FP8 ViT attention path (quantize -> cuDNN -> un-pad)."""
4 +
5 +import contextlib
6 +
7 +import pytest
8 +import torch
9 +
10 +from vllm.triton_utils import HAS_TRITON
11 +from vllm.utils.flashinfer import (
12 + is_flashinfer_cudnn_fp8_prefill_attn_supported,
13 +)
14 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
15 +
16 +
17 +def _has_flashinfer_cudnn() -> bool:
18 + """Check if FlashInfer cuDNN backend is available."""
19 + try:
20 + from flashinfer.prefill import (
21 + cudnn_batch_prefill_with_kv_cache, # noqa: F401
22 + )
23 +
24 + return True
25 + except ImportError:
26 + return False
27 +
28 +
29 +HEAD_DIMS = [72, 80]
30 +SEQ_LENS = [256]
31 +NUM_HEADS = [16]
32 +
33 +
34 +@pytest.fixture
35 +def _fp8_attention():
36 + """Create FP8-enabled MMEncoderAttention via config."""
37 + from types import SimpleNamespace
38 + from unittest.mock import patch
39 +
40 + from vllm.config import VllmConfig, set_current_vllm_config
41 + from vllm.config.multimodal import MultiModalConfig
42 +
43 + if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
44 + pytest.skip("FlashInfer cuDNN FP8 prefill attention not supported")
45 +
46 + mm_config = MultiModalConfig(mm_encoder_attn_dtype="fp8")
47 + vllm_config = VllmConfig()
48 + vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
49 +
50 + # MMEncoderAttention reads torch.get_default_dtype() during init
51 + # to determine the output dtype. In real model loading this is bf16.
52 + old_dtype = torch.get_default_dtype()
53 + torch.set_default_dtype(torch.bfloat16)
54 +
55 + with (
56 + set_current_vllm_config(vllm_config),
57 + patch(
58 + "vllm.model_executor.layers.attention.mm_encoder_attention"
59 + ".get_vit_attn_backend",
60 + return_value=AttentionBackendEnum.FLASHINFER,
61 + ),
62 + ):
63 + yield
64 +
65 + torch.set_default_dtype(old_dtype)
66 +
67 +
68 +def _build_cu_seqlens_and_meta(
69 + seq_len: int,
70 + num_heads: int,
71 + head_dim: int,
72 + fp8_padded_hidden_size: int | None = None,
73 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
74 + """Build cu_seqlens, max_seqlen, sequence_lengths for a single sequence."""
75 + import numpy as np
76 +
77 + from vllm.model_executor.layers.attention.mm_encoder_attention import (
78 + MMEncoderAttention,
79 + )
80 +
81 + cu_seqlens_np = np.array([0, seq_len], dtype=np.int32)
82 +
83 + sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
84 + AttentionBackendEnum.FLASHINFER,
85 + cu_seqlens_np,
86 + torch.device("cuda"),
87 + )
88 +
89 + max_seqlen = torch.tensor(
90 + MMEncoderAttention.compute_max_seqlen(
91 + AttentionBackendEnum.FLASHINFER, cu_seqlens_np
92 + ),
93 + dtype=torch.int32,
94 + )
95 +
96 + cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
97 + AttentionBackendEnum.FLASHINFER,
98 + cu_seqlens_np,
99 + num_heads * head_dim,
100 + 1, # tp_size
101 + torch.device("cuda"),
102 + fp8_padded_hidden_size=fp8_padded_hidden_size,
103 + )
104 +
105 + return cu_seqlens, max_seqlen, sequence_lengths
106 +
107 +
108 +@pytest.mark.skipif(
109 + not (HAS_TRITON and _has_flashinfer_cudnn()),
110 + reason="Triton and FlashInfer cuDNN required",
111 +)
112 +@pytest.mark.parametrize("head_dim", HEAD_DIMS)
113 +@pytest.mark.parametrize("seq_len", SEQ_LENS)
114 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
115 +def test_fp8_attn_output_shape(
116 + head_dim: int,
117 + seq_len: int,
118 + num_heads: int,
119 + _fp8_attention,
120 +) -> None:
121 + """Verify FP8 attention produces correct output shape after un-padding."""
122 + from vllm.model_executor.layers.attention.mm_encoder_attention import (
123 + MMEncoderAttention,
124 + )
125 + from vllm.utils.math_utils import round_up
126 +
127 + attn = None
128 + with contextlib.suppress(ValueError, ImportError):
129 + attn = MMEncoderAttention(
130 + num_heads=num_heads,
131 + head_size=head_dim,
132 + prefix="visual.blocks.0.attn",
133 + ).to("cuda")
134 +
135 + if attn is None or not attn.fp8_enabled:
136 + pytest.skip("FP8 MMEncoderAttention not available")
137 + assert attn is not None # mypy narrowing
138 +
139 + # FP8 always needs fp8_padded_hidden_size for correct cu_seqlens
140 + fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
141 +
142 + cu_seqlens, max_seqlen, sequence_lengths = _build_cu_seqlens_and_meta(
143 + seq_len, num_heads, head_dim, fp8_padded_hidden_size=fp8_padded_hidden_size
144 + )
145 +
146 + q = torch.randn(
147 + seq_len,
148 + num_heads,
149 + head_dim,
150 + device="cuda",
151 + dtype=torch.bfloat16,
152 + )
153 + k = torch.randn_like(q)
154 + v = torch.randn_like(q)
155 +
156 + output = attn._forward_flashinfer(q, k, v, cu_seqlens, max_seqlen, sequence_lengths)
157 +
158 + # Output should have original head_dim (un-padded)
159 + assert output.shape[-1] == head_dim
160 + assert output.dtype == torch.bfloat16
161 +
162 +
163 +@pytest.mark.skipif(
164 + not (HAS_TRITON and _has_flashinfer_cudnn()),
165 + reason="Triton and FlashInfer cuDNN required",
166 +)
167 +@pytest.mark.parametrize("head_dim", HEAD_DIMS)
168 +@pytest.mark.parametrize("seq_len", SEQ_LENS)
169 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
170 +def test_fp8_vs_bf16_close(
171 + head_dim: int, seq_len: int, num_heads: int, _fp8_attention
172 +) -> None:
173 + """FP8 attention output should be reasonably close to BF16 baseline."""
174 + from vllm.model_executor.layers.attention.mm_encoder_attention import (
175 + MMEncoderAttention,
176 + )
177 + from vllm.utils.math_utils import round_up
178 +
179 + torch.manual_seed(42)
180 + q = torch.randn(
181 + 1,
182 + seq_len,
183 + num_heads,
184 + head_dim,
185 + device="cuda",
186 + dtype=torch.bfloat16,
187 + )
188 + k = torch.randn_like(q)
189 + v = torch.randn_like(q)
190 +
191 + # FP8 path
192 + attn_fp8 = None
193 + with contextlib.suppress(ValueError, ImportError):
194 + attn_fp8 = MMEncoderAttention(
195 + num_heads=num_heads,
196 + head_size=head_dim,
197 + prefix="visual.blocks.0.attn",
198 + ).to("cuda")
199 +
200 + if attn_fp8 is None or not attn_fp8.fp8_enabled:
201 + pytest.skip("FP8 MMEncoderAttention not available")
202 + assert attn_fp8 is not None # mypy narrowing
203 +
204 + fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
205 + cu_seqlens, max_seqlen, seq_lengths = _build_cu_seqlens_and_meta(
206 + seq_len,
207 + num_heads,
208 + head_dim,
209 + fp8_padded_hidden_size=fp8_padded_hidden_size,
210 + )
211 +
212 + out_fp8 = attn_fp8._forward_flashinfer(
213 + q.clone(),
214 + k.clone(),
215 + v.clone(),
216 + cu_seqlens,
217 + max_seqlen,
218 + seq_lengths,
219 + )
220 +
221 + # BF16 baseline (create non-FP8 attention by using scale=attn_fp8.scale
222 + # and calling the wrapper directly without FP8 quantization)
223 + from vllm.model_executor.layers.attention.mm_encoder_attention import (
224 + _get_flashinfer_workspace_buffer,
225 + )
226 + from vllm.v1.attention.ops.vit_attn_wrappers import (
227 + vit_flashinfer_wrapper,
228 + )
229 +
230 + out_bf16 = vit_flashinfer_wrapper(
231 + q=q.clone(),
232 + k=k.clone(),
233 + v=v.clone(),
234 + scale=attn_fp8.scale,
235 + workspace_buffer=_get_flashinfer_workspace_buffer(),
236 + cu_seqlens=cu_seqlens,
237 + max_seqlen=max_seqlen,
238 + sequence_lengths=seq_lengths,
239 + )
240 +
241 + out_fp8_f = out_fp8.float()
242 + out_bf16_f = out_bf16.float()
243 +
244 + abs_diff = (out_fp8_f - out_bf16_f).abs()
245 + abs_diff_flat = abs_diff.flatten()
246 +
247 + # Relative diff (avoid division by zero)
248 + denom = out_bf16_f.abs().clamp(min=1e-6)
249 + rel_diff_flat = (abs_diff / denom).flatten()
250 +
251 + cosine_sim = torch.nn.functional.cosine_similarity(
252 + out_fp8_f.flatten().unsqueeze(0),
253 + out_bf16_f.flatten().unsqueeze(0),
254 + ).item()
255 +
256 + pcts = [50, 90, 95, 99, 99.9]
257 + abs_pct = {p: torch.quantile(abs_diff_flat, p / 100).item() for p in pcts}
258 + rel_pct = {p: torch.quantile(rel_diff_flat, p / 100).item() for p in pcts}
259 +
260 + print(f"\nFP8 vs BF16 (head_dim={head_dim}, seq_len={seq_len}):")
261 + print(f" cosine_sim={cosine_sim:.6f}")
262 + print(
263 + f" abs_diff: max={abs_diff_flat.max().item():.6f}, "
264 + f"mean={abs_diff_flat.mean().item():.6f}, "
265 + + ", ".join(f"p{p}={abs_pct[p]:.6f}" for p in pcts)
266 + )
267 + print(
268 + f" rel_diff: max={rel_diff_flat.max().item():.6f}, "
269 + f"mean={rel_diff_flat.mean().item():.6f}, "
270 + + ", ".join(f"p{p}={rel_pct[p]:.6f}" for p in pcts)
271 + )
272 +
273 + assert abs_diff_flat.max().item() < 0.3, (
274 + f"FP8 vs BF16 max abs diff too large: {abs_diff_flat.max().item()}"
275 + )
276 + assert abs_diff_flat.mean().item() < 0.03, (
277 + f"FP8 vs BF16 mean abs diff too large: {abs_diff_flat.mean().item()}"
278 + )
279 + assert cosine_sim > 0.99, f"Cosine similarity too low: {cosine_sim:.6f}"
tests/kernels/core/test_vit_fp8_quant.py new
+124
@@ -0,0 +1,124 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the stride-aware FP8 quantization kernel with head_dim padding."""
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm.platforms import current_platform
9 +from vllm.triton_utils import HAS_TRITON
10 +
11 +if HAS_TRITON:
12 + from vllm.kernels.triton.qkv_padded_fp8_quant import (
13 + quantize_fp8_pad_head_dim_triton,
14 + )
15 +
16 +HEAD_DIMS = [72, 80, 128]
17 +SEQ_LENS = [64, 256]
18 +NUM_HEADS = [16]
19 +SCALES = [0.01, 0.1, 1.0]
20 +
21 +
22 +def _naive_fp8_quantize(
23 + tensor: torch.Tensor, scale: torch.Tensor, skip_scale: bool
24 +) -> torch.Tensor:
25 + """Reference FP8 quantization in PyTorch."""
26 + fp8_dtype = current_platform.fp8_dtype()
27 + fp8_max = torch.finfo(fp8_dtype).max
28 + fp8_min = -fp8_max
29 +
30 + x = tensor.float()
31 + if not skip_scale:
32 + x = x / scale.item()
33 + x = x.clamp(fp8_min, fp8_max)
34 + return x.to(fp8_dtype)
35 +
36 +
37 +@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
38 +@pytest.mark.parametrize("head_dim", HEAD_DIMS)
39 +@pytest.mark.parametrize("seq_len", SEQ_LENS)
40 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
41 +@pytest.mark.parametrize("scale_val", SCALES)
42 +def test_quantize_contiguous(
43 + head_dim: int, seq_len: int, num_heads: int, scale_val: float
44 +) -> None:
45 + """Test quantization of contiguous 3D tensors."""
46 + torch.manual_seed(42)
47 + tensor = torch.randn(
48 + seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
49 + )
50 + scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda").view(
51 + 1, 1, 1, 1
52 + )
53 +
54 + result = quantize_fp8_pad_head_dim_triton(tensor, scale)
55 +
56 + padded_dim = (head_dim + 15) // 16 * 16
57 + assert result.shape == (seq_len, num_heads, padded_dim)
58 + assert result.is_contiguous()
59 + assert result.dtype == current_platform.fp8_dtype()
60 +
61 + # Compare unpadded portion against reference
62 + ref = _naive_fp8_quantize(tensor, scale, skip_scale=False)
63 + torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
64 +
65 + # Padded region should be zero
66 + if padded_dim > head_dim:
67 + assert (result[:, :, head_dim:].float() == 0).all()
68 +
69 +
70 +@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
71 +@pytest.mark.parametrize("head_dim", [72, 80])
72 +def test_quantize_non_contiguous(head_dim: int) -> None:
73 + """Test quantization from non-contiguous QKV views (interleaved buffer)."""
74 + seq_len, num_heads = 64, 16
75 + # Simulate interleaved QKV buffer: shape (seq_len, 3 * num_heads, head_dim)
76 + qkv = torch.randn(
77 + seq_len, 3 * num_heads, head_dim, device="cuda", dtype=torch.bfloat16
78 + )
79 + # Q is every 3rd head slice - non-contiguous view
80 + q = qkv[:, 0::3, :]
81 + assert not q.is_contiguous()
82 +
83 + scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
84 + result = quantize_fp8_pad_head_dim_triton(q, scale)
85 +
86 + padded_dim = (head_dim + 15) // 16 * 16
87 + assert result.shape == (seq_len, num_heads, padded_dim)
88 + assert result.is_contiguous()
89 +
90 + # Compare against contiguous reference
91 + ref = _naive_fp8_quantize(q.contiguous(), scale, skip_scale=False)
92 + torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
93 +
94 +
95 +@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
96 +def test_skip_scale() -> None:
97 + """Test skip_scale=True produces cast-only output (no division)."""
98 + seq_len, num_heads, head_dim = 32, 8, 80
99 + tensor = torch.randn(
100 + seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
101 + )
102 + scale = torch.tensor([0.5], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
103 +
104 + result_skip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=True)
105 + result_noskip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=False)
106 +
107 + # skip_scale should just cast, not divide
108 + ref_cast = _naive_fp8_quantize(tensor, scale, skip_scale=True)
109 + torch.testing.assert_close(result_skip[:, :, :head_dim].float(), ref_cast.float())
110 +
111 + # With scale != 1.0, skip and no-skip should differ
112 + assert not torch.equal(result_skip.float(), result_noskip.float())
113 +
114 +
115 +@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
116 +def test_4d_input() -> None:
117 + """Test that 4D input (B, S, H, D) is handled correctly."""
118 + B, S, H, D = 2, 32, 8, 72
119 + tensor = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
120 + scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
121 +
122 + result = quantize_fp8_pad_head_dim_triton(tensor, scale)
123 + padded_dim = (D + 15) // 16 * 16
124 + assert result.shape == (B, S, H, padded_dim)
tests/kernels/core/test_vit_fp8_scaling.py new
+251
@@ -0,0 +1,251 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for FP8 scaling (dynamic and static) in MMEncoderAttention."""
4 +
5 +import contextlib
6 +import json
7 +from types import SimpleNamespace
8 +from unittest.mock import patch
9 +
10 +import pytest
11 +import torch
12 +
13 +from vllm.model_executor.layers.attention.mm_encoder_attention import (
14 + _FP8_AMAX_HISTORY_LEN,
15 + _FP8_MAX,
16 +)
17 +from vllm.utils.flashinfer import (
18 + is_flashinfer_cudnn_fp8_prefill_attn_supported,
19 +)
20 +
21 +LAYER_0 = "visual.blocks.0.attn.attn"
22 +LAYER_1 = "visual.blocks.1.attn.attn"
23 +NUM_HEADS = 16
24 +HEAD_DIM = 72
25 +
26 +
27 +@contextlib.contextmanager
28 +def _build_attention(mm_config):
29 + """Yield an MMEncoderAttention with the given multimodal config.
30 +
31 + The VllmConfig context stays active while the test runs so that
32 + ``get_multimodal_config()`` calls during the forward path resolve. Also
33 + invokes ``process_weights_after_loading`` to simulate the model loader's
34 + auto-scan. Yields ``None`` if FlashInfer cuDNN is not available.
35 + """
36 + from vllm.config import VllmConfig, set_current_vllm_config
37 + from vllm.model_executor.layers.attention.mm_encoder_attention import (
38 + MMEncoderAttention,
39 + )
40 + from vllm.v1.attention.backends.registry import AttentionBackendEnum
41 +
42 + if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
43 + yield None
44 + return
45 +
46 + vllm_config = VllmConfig()
47 + vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
48 +
49 + with (
50 + set_current_vllm_config(vllm_config),
51 + patch(
52 + "vllm.model_executor.layers.attention.mm_encoder_attention"
53 + ".get_vit_attn_backend",
54 + return_value=AttentionBackendEnum.FLASHINFER,
55 + ),
56 + ):
57 + attn = MMEncoderAttention(
58 + num_heads=NUM_HEADS,
59 + head_size=HEAD_DIM,
60 + prefix=LAYER_0,
61 + )
62 + attn.process_weights_after_loading(torch.bfloat16)
63 + yield attn
64 +
65 +
66 +@pytest.fixture
67 +def _make_attention():
68 + """Create an MMEncoderAttention with dynamic FP8 scaling."""
69 + from vllm.config.multimodal import MultiModalConfig
70 +
71 + with _build_attention(MultiModalConfig(mm_encoder_attn_dtype="fp8")) as attn:
72 + yield attn
73 +
74 +
75 +@pytest.fixture
76 +def _make_static_attention(tmp_path):
77 + """Create an MMEncoderAttention with static FP8 scales from a file."""
78 + from vllm.config.multimodal import MultiModalConfig
79 +
80 + scale_file = tmp_path / "scales.json"
81 + scale_file.write_text(
82 + json.dumps(
83 + {
84 + LAYER_0: {"q": 224.0, "k": 198.0, "v": 210.0},
85 + LAYER_1: {"q": 100.0, "k": 110.0, "v": 120.0},
86 + }
87 + )
88 + )
89 + with _build_attention(
90 + MultiModalConfig(
91 + mm_encoder_attn_dtype="fp8",
92 + mm_encoder_fp8_scale_path=str(scale_file),
93 + )
94 + ) as attn:
95 + yield attn
96 +
97 +
98 +def test_dynamic_scaling_updates_scales(_make_attention) -> None:
99 + """Verify that _record_amax_and_update_scales updates scale buffers."""
100 + attn = _make_attention
101 + if attn is None or not attn.fp8_enabled:
102 + pytest.skip("FP8 attention not available (FlashInfer backend required)")
103 +
104 + attn = attn.to("cuda")
105 +
106 + S, H, D = 32, NUM_HEADS, HEAD_DIM
107 + q = torch.full((S, H, D), 2.0, device="cuda", dtype=torch.bfloat16)
108 + k = torch.full((S, H, D), 3.0, device="cuda", dtype=torch.bfloat16)
109 + v = torch.full((S, H, D), 4.0, device="cuda", dtype=torch.bfloat16)
110 +
111 + attn._record_amax_and_update_scales(q, k, v)
112 +
113 + expected_q_scale = 2.0 / _FP8_MAX
114 + expected_k_scale = 3.0 / _FP8_MAX
115 + expected_v_scale = 4.0 / _FP8_MAX
116 +
117 + torch.testing.assert_close(attn._fp8_q_scale.item(), expected_q_scale)
118 + torch.testing.assert_close(attn._fp8_k_scale.item(), expected_k_scale)
119 + torch.testing.assert_close(attn._fp8_v_scale.item(), expected_v_scale)
120 +
121 +
122 +def test_circular_buffer_wraps(_make_attention) -> None:
123 + """Verify the amax circular buffer wraps at HISTORY_LEN."""
124 + attn = _make_attention
125 + if attn is None or not attn.fp8_enabled:
126 + pytest.skip("FP8 attention not available (FlashInfer backend required)")
127 +
128 + attn = attn.to("cuda")
129 + S, H, D = 16, NUM_HEADS, HEAD_DIM
130 +
131 + for i in range(_FP8_AMAX_HISTORY_LEN + 2):
132 + mag = float(i + 1)
133 + q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
134 + k = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
135 + v = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
136 + attn._record_amax_and_update_scales(q, k, v)
137 +
138 + assert attn._fp8_amax_pos == 2
139 +
140 + expected_max = float(_FP8_AMAX_HISTORY_LEN + 2)
141 + expected_scale = expected_max / _FP8_MAX
142 + torch.testing.assert_close(attn._fp8_q_scale.item(), expected_scale)
143 +
144 +
145 +def test_static_scales_loaded(_make_static_attention) -> None:
146 + """Verify static scales are loaded from the JSON file."""
147 + attn = _make_static_attention
148 + if attn is None or not attn.fp8_enabled:
149 + pytest.skip("FP8 attention not available (FlashInfer backend required)")
150 +
151 + assert attn.fp8_enabled
152 + assert not attn._fp8_dynamic_scale
153 +
154 + # Layer 0 scales (the layer this attention was created with).
155 + assert attn._fp8_q_scale.item() == 224.0
156 + assert attn._fp8_k_scale.item() == 198.0
157 + assert attn._fp8_v_scale.item() == 210.0
158 +
159 + assert not attn.skip_scale_q
160 + assert not attn.skip_scale_k
161 + assert not attn.skip_scale_v
162 +
163 + # No amax history buffers for static scaling.
164 + assert not hasattr(attn, "_fp8_q_amax")
165 +
166 +
167 +def test_static_scales_missing_layer(tmp_path) -> None:
168 + """Verify error when requested layer is not in the scale file."""
169 + from vllm.config import VllmConfig, set_current_vllm_config
170 + from vllm.config.multimodal import MultiModalConfig
171 + from vllm.v1.attention.backends.registry import AttentionBackendEnum
172 +
173 + if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
174 + pytest.skip("FlashInfer cuDNN not available")
175 +
176 + scale_file = tmp_path / "wrong_layer.json"
177 + scale_file.write_text(
178 + json.dumps({"visual.blocks.99.attn": {"q": 1.0, "k": 1.0, "v": 1.0}})
179 + )
180 + mm_config = MultiModalConfig(
181 + mm_encoder_attn_dtype="fp8",
182 + mm_encoder_fp8_scale_path=str(scale_file),
183 + )
184 + vllm_config = VllmConfig()
185 + vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
186 +
187 + from vllm.model_executor.layers.attention.mm_encoder_attention import (
188 + MMEncoderAttention,
189 + )
190 +
191 + with (
192 + set_current_vllm_config(vllm_config),
193 + patch(
194 + "vllm.model_executor.layers.attention.mm_encoder_attention"
195 + ".get_vit_attn_backend",
196 + return_value=AttentionBackendEnum.FLASHINFER,
197 + ),
198 + ):
199 + attn = MMEncoderAttention(
200 + num_heads=NUM_HEADS,
201 + head_size=HEAD_DIM,
202 + prefix=LAYER_0,
203 + )
204 + with pytest.raises(ValueError, match="scales not found for layer"):
205 + attn.process_weights_after_loading(torch.bfloat16)
206 +
207 +
208 +def test_dynamic_scales_auto_save(tmp_path) -> None:
209 + """Verify scales are saved to disk after the amax buffer fills."""
210 + import vllm.model_executor.layers.attention.mm_encoder_attention as _mod
211 + from vllm.config.multimodal import MultiModalConfig
212 +
213 + if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
214 + pytest.skip("FlashInfer cuDNN not available")
215 +
216 + # Reset module-level state between runs (other tests may have left
217 + # state behind after triggering a save).
218 + _mod._fp8_scale_save_path = None
219 + _mod._fp8_saved_scale_refs.clear()
220 +
221 + save_file = tmp_path / "auto_scales.json"
222 + with _build_attention(
223 + MultiModalConfig(
224 + mm_encoder_attn_dtype="fp8",
225 + mm_encoder_fp8_scale_save_path=str(save_file),
226 + )
227 + ) as attn:
228 + if attn is None or not attn.fp8_enabled:
229 + pytest.skip("FP8 attention not available")
230 +
231 + attn = attn.to("cuda")
232 + S, H, D = 16, NUM_HEADS, HEAD_DIM
233 +
234 + # Run exactly _FP8_AMAX_HISTORY_LEN forward passes.
235 + for i in range(_FP8_AMAX_HISTORY_LEN):
236 + mag = float(i + 1)
237 + q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
238 + k = torch.full((S, H, D), mag * 0.5, device="cuda", dtype=torch.bfloat16)
239 + v = torch.full((S, H, D), mag * 0.3, device="cuda", dtype=torch.bfloat16)
240 + attn._record_amax_and_update_scales(q, k, v)
241 +
242 + # File should have been written on the 16th call (buffer wrap).
243 + assert save_file.is_file(), "Scale file was not saved"
244 + scales = json.loads(save_file.read_text())
245 + assert LAYER_0 in scales
246 + assert set(scales[LAYER_0].keys()) == {"q", "k", "v"}
247 + for val in scales[LAYER_0].values():
248 + assert isinstance(val, float) and val > 0
249 +
250 + # Path is cleared after the one-shot save fires.
251 + assert _mod._fp8_scale_save_path is None
vllm/config/model.py
+12
@@ -326,6 +326,10 @@ class ModelConfig:
326 mm_encoder_only: InitVar[bool | None] = None
327 mm_encoder_tp_mode: InitVar[MMEncoderTPMode | None] = None
328 mm_encoder_attn_backend: InitVar[AttentionBackendEnum | str | None] = None
329 + mm_encoder_attn_dtype: InitVar[str | None] = None
330 + mm_encoder_fp8_scale_path: InitVar[str | None] = None
331 + mm_encoder_fp8_scale_save_path: InitVar[str | None] = None
332 + mm_encoder_fp8_scale_save_margin: InitVar[float | None] = None
333 interleave_mm_strings: InitVar[bool | None] = None
334 skip_mm_profiling: InitVar[bool | None] = None
335 video_pruning_rate: InitVar[float | None] = None
@@ -447,6 +451,10 @@ class ModelConfig:
451 mm_encoder_only: bool | None,
452 mm_encoder_tp_mode: MMEncoderTPMode | None,
453 mm_encoder_attn_backend: AttentionBackendEnum | str | None,
454 + mm_encoder_attn_dtype: str | None,
455 + mm_encoder_fp8_scale_path: str | None,
456 + mm_encoder_fp8_scale_save_path: str | None,
457 + mm_encoder_fp8_scale_save_margin: float | None,
458 interleave_mm_strings: bool | None,
459 skip_mm_profiling: bool | None,
460 video_pruning_rate: float | None,
@@ -643,6 +651,10 @@ class ModelConfig:
651 mm_encoder_only=mm_encoder_only,
652 mm_encoder_tp_mode=mm_encoder_tp_mode,
653 mm_encoder_attn_backend=mm_encoder_attn_backend,
654 + mm_encoder_attn_dtype=mm_encoder_attn_dtype,
655 + mm_encoder_fp8_scale_path=mm_encoder_fp8_scale_path,
656 + mm_encoder_fp8_scale_save_path=mm_encoder_fp8_scale_save_path,
657 + mm_encoder_fp8_scale_save_margin=mm_encoder_fp8_scale_save_margin,
658 interleave_mm_strings=interleave_mm_strings,
659 skip_mm_profiling=skip_mm_profiling,
660 video_pruning_rate=video_pruning_rate,
vllm/config/multimodal.py
+51
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Mapping
5 +from pathlib import Path
6 from typing import Any, Literal, TypeAlias, TypedDict, final
7
8 from pydantic import ConfigDict, Field, field_validator, model_validator
@@ -158,6 +159,24 @@ class MultiModalConfig:
159 """Optional override for the multi-modal encoder attention backend when
160 using vision transformers. Accepts any value from
161 `vllm.v1.attention.backends.registry.AttentionBackendEnum` (e.g. `FLASH_ATTN`)."""
162 + mm_encoder_attn_dtype: Literal["fp8"] | None = None
163 + """Optional dtype override for ViT encoder attention. Set to `"fp8"` to
164 + enable FP8 quantization via the FlashInfer cuDNN backend. When set to
165 + `"fp8"` without a scale file, dynamic scaling is used automatically.
166 + See docs/features/quantization/fp8_vit_attn.md for details."""
167 + mm_encoder_fp8_scale_path: str | None = None
168 + """Path to a JSON file containing per-layer FP8 Q/K/V scales for ViT
169 + encoder attention. When provided (with `mm_encoder_attn_dtype="fp8"`),
170 + static scaling is used. When omitted, dynamic scaling is used."""
171 + mm_encoder_fp8_scale_save_path: str | None = None
172 + """When set with dynamic FP8 scaling (`mm_encoder_attn_dtype="fp8"`
173 + and no `mm_encoder_fp8_scale_path`), saves the calibrated scales to
174 + this file after the amax history buffer is full. The saved file can
175 + then be used as `mm_encoder_fp8_scale_path` in subsequent runs."""
176 + mm_encoder_fp8_scale_save_margin: float = Field(default=1.5, gt=0.0)
177 + """Safety margin multiplied onto scales when auto-saving. A value > 1
178 + leaves headroom so that inputs with larger activations than the
179 + calibration set do not overflow FP8 range. Default 1.5."""
180 interleave_mm_strings: bool = False
181 """Enable fully interleaved support for multimodal prompts, while using
182 --chat-template-content-format=string."""
@@ -233,6 +252,36 @@ class MultiModalConfig:
252 "'mm_shm_cache_max_object_size_mb' should only be set when "
253 "'mm_processor_cache_type' is 'shm'."
254 )
255 + # Validate FP8 scale path combinations.
256 + if self.mm_encoder_attn_dtype != "fp8" and (
257 + self.mm_encoder_fp8_scale_path is not None
258 + or self.mm_encoder_fp8_scale_save_path is not None
259 + ):
260 + raise ValueError(
261 + "'mm_encoder_fp8_scale_path' and "
262 + "'mm_encoder_fp8_scale_save_path' require "
263 + "'mm_encoder_attn_dtype' to be 'fp8'."
264 + )
265 + if (
266 + self.mm_encoder_fp8_scale_path is not None
267 + and self.mm_encoder_fp8_scale_save_path is not None
268 + ):
269 + raise ValueError(
270 + "'mm_encoder_fp8_scale_save_path' cannot be used with "
271 + "'mm_encoder_fp8_scale_path' (saving requires dynamic scaling)."
272 + )
273 +
274 + # Validate file paths exist.
275 + if self.mm_encoder_fp8_scale_path is not None:
276 + scale_path = Path(self.mm_encoder_fp8_scale_path)
277 + if not scale_path.is_file():
278 + raise FileNotFoundError(f"FP8 scale file not found: {scale_path}")
279 + if self.mm_encoder_fp8_scale_save_path is not None:
280 + save_parent = Path(self.mm_encoder_fp8_scale_save_path).parent
281 + if not save_parent.is_dir():
282 + raise FileNotFoundError(
283 + f"Parent directory for FP8 scale save path not found: {save_parent}"
284 + )
285 return self
286
287 def compute_hash(self) -> str:
@@ -252,6 +301,8 @@ class MultiModalConfig:
301 if self.mm_encoder_attn_backend is not None
302 else None,
303 self.mm_encoder_tp_mode,
304 + self.mm_encoder_attn_dtype,
305 + self.mm_encoder_fp8_scale_path,
306 ]
307 hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
308 return hash_str
vllm/engine/arg_utils.py
+28
@@ -542,6 +542,14 @@ class EngineArgs:
542 mm_encoder_attn_backend: AttentionBackendEnum | str | None = (
543 MultiModalConfig.mm_encoder_attn_backend
544 )
545 + mm_encoder_attn_dtype: str | None = MultiModalConfig.mm_encoder_attn_dtype
546 + mm_encoder_fp8_scale_path: str | None = MultiModalConfig.mm_encoder_fp8_scale_path
547 + mm_encoder_fp8_scale_save_path: str | None = (
548 + MultiModalConfig.mm_encoder_fp8_scale_save_path
549 + )
550 + mm_encoder_fp8_scale_save_margin: float = (
551 + MultiModalConfig.mm_encoder_fp8_scale_save_margin
552 + )
553 io_processor_plugin: str | None = None
554 renderer_num_workers: int = 1
555 skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling
@@ -1179,6 +1187,22 @@ class EngineArgs:
1187 "--mm-encoder-attn-backend",
1188 **multimodal_kwargs["mm_encoder_attn_backend"],
1189 )
1190 + multimodal_group.add_argument(
1191 + "--mm-encoder-attn-dtype",
1192 + **multimodal_kwargs["mm_encoder_attn_dtype"],
1193 + )
1194 + multimodal_group.add_argument(
1195 + "--mm-encoder-fp8-scale-path",
1196 + **multimodal_kwargs["mm_encoder_fp8_scale_path"],
1197 + )
1198 + multimodal_group.add_argument(
1199 + "--mm-encoder-fp8-scale-save-path",
1200 + **multimodal_kwargs["mm_encoder_fp8_scale_save_path"],
1201 + )
1202 + multimodal_group.add_argument(
1203 + "--mm-encoder-fp8-scale-save-margin",
1204 + **multimodal_kwargs["mm_encoder_fp8_scale_save_margin"],
1205 + )
1206 multimodal_group.add_argument(
1207 "--interleave-mm-strings", **multimodal_kwargs["interleave_mm_strings"]
1208 )
@@ -1517,6 +1541,10 @@ class EngineArgs:
1541 mm_encoder_only=self.mm_encoder_only,
1542 mm_encoder_tp_mode=self.mm_encoder_tp_mode,
1543 mm_encoder_attn_backend=self.mm_encoder_attn_backend,
1544 + mm_encoder_attn_dtype=self.mm_encoder_attn_dtype,
1545 + mm_encoder_fp8_scale_path=self.mm_encoder_fp8_scale_path,
1546 + mm_encoder_fp8_scale_save_path=self.mm_encoder_fp8_scale_save_path,
1547 + mm_encoder_fp8_scale_save_margin=self.mm_encoder_fp8_scale_save_margin,
1548 pooler_config=self.pooler_config,
1549 generation_config=self.generation_config,
1550 override_generation_config=self.override_generation_config,
vllm/kernels/triton/__init__.py new
+3
@@ -0,0 +1,3 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Triton kernel implementations."""
vllm/kernels/triton/qkv_padded_fp8_quant.py new
+180
@@ -0,0 +1,180 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Stride-aware FP8 quantization with head_dim padding for ViT attention.
4 +
5 +Reads directly from non-contiguous QKV views using 3D strides and pads
6 +head_dim to a multiple of 16 for cuDNN compatibility.
7 +"""
8 +
9 +import torch
10 +
11 +from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
12 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
13 + get_fp8_min_max,
14 +)
15 +from vllm.platforms import current_platform
16 +from vllm.triton_utils import HAS_TRITON, tl, triton
17 +from vllm.utils.math_utils import round_up
18 +
19 +_FP8_MIN, _FP8_MAX = get_fp8_min_max()
20 +
21 +
22 +@triton.jit
23 +def _quantize_pad_fp8_kernel(
24 + x_ptr,
25 + y_ptr,
26 + scale_ptr,
27 + stride_xs,
28 + stride_xh,
29 + stride_xd,
30 + stride_ys,
31 + stride_yh,
32 + stride_yd,
33 + num_heads,
34 + n_rows,
35 + n_cols,
36 + n_cols_padded,
37 + fp8_min,
38 + fp8_max,
39 + SKIP_SCALE: tl.constexpr,
40 + BLOCK_M: tl.constexpr,
41 + BLOCK_N: tl.constexpr,
42 +):
43 + pid_m = tl.program_id(0)
44 + pid_n = tl.program_id(1)
45 +
46 + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
47 + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
48 + mask_m = offs_m < n_rows
49 + mask_out = mask_m[:, None] & (offs_n[None, :] < n_cols_padded)
50 + mask_in = mask_m[:, None] & (offs_n[None, :] < n_cols)
51 +
52 + # Decompose flattened row into (token, head) for 3D stride indexing.
53 + s = offs_m // num_heads
54 + h = offs_m % num_heads
55 +
56 + x_ptrs = (
57 + x_ptr
58 + + s[:, None] * stride_xs
59 + + h[:, None] * stride_xh
60 + + offs_n[None, :] * stride_xd
61 + )
62 + x = tl.load(x_ptrs, mask=mask_in, other=0.0).to(tl.float32)
63 + if SKIP_SCALE:
64 + x_q = x
65 + else:
66 + scale = tl.load(scale_ptr)
67 + x_q = x / scale
68 + x_q = tl.clamp(x_q, fp8_min, fp8_max).to(y_ptr.dtype.element_ty)
69 +
70 + y_ptrs = (
71 + y_ptr
72 + + s[:, None] * stride_ys
73 + + h[:, None] * stride_yh
74 + + offs_n[None, :] * stride_yd
75 + )
76 + tl.store(y_ptrs, x_q, mask=mask_out)
77 +
78 +
79 +def _get_fp8_pad_quant_config(padded_head_dim: int) -> tuple[int, int, int]:
80 + block_n = triton.next_power_of_2(padded_head_dim)
81 + block_n = max(16, min(block_n, 128))
82 + block_m = 16
83 + num_warps = 4
84 + return block_m, block_n, num_warps
85 +
86 +
87 +def quantize_fp8_pad_head_dim_triton(
88 + tensor: torch.Tensor,
89 + scale: torch.Tensor,
90 + skip_scale: bool = False,
91 + block_m: int | None = None,
92 + block_n: int | None = None,
93 + num_warps: int | None = None,
94 +) -> torch.Tensor:
95 + """Quantize a 3D/4D tensor to FP8, padding head_dim to a multiple of 16.
96 +
97 + Reads directly from the input using its 3D strides, so non-contiguous
98 + views (e.g. Q/K/V slices from an interleaved QKV buffer) are handled
99 + without an extra copy. Output is always a fresh contiguous tensor
100 + with shape (S, H, padded_D).
101 + """
102 + if not HAS_TRITON:
103 + raise RuntimeError("Triton is required to quantize with head_dim padding.")
104 +
105 + original_shape = tensor.shape
106 + if tensor.dim() == 4:
107 + tensor = tensor.view(-1, tensor.shape[-2], tensor.shape[-1])
108 + assert tensor.dim() == 3, f"Expected 3D input (S, H, D), got {tensor.dim()}D"
109 + S, H, D = tensor.shape
110 + padded_head_dim = round_up(D, 16)
111 + out_dtype = current_platform.fp8_dtype()
112 + output = torch.empty(
113 + (S, H, padded_head_dim),
114 + device=tensor.device,
115 + dtype=out_dtype,
116 + )
117 +
118 + scale_1d = scale.reshape(-1)
119 + n_rows = S * H
120 +
121 + if block_m is None or block_n is None or num_warps is None:
122 + block_m, block_n, num_warps = _get_fp8_pad_quant_config(padded_head_dim)
123 +
124 + grid = (
125 + triton.cdiv(n_rows, block_m),
126 + triton.cdiv(padded_head_dim, block_n),
127 + )
128 +
129 + _quantize_pad_fp8_kernel[grid](
130 + tensor,
131 + output,
132 + scale_1d,
133 + tensor.stride(0),
134 + tensor.stride(1),
135 + tensor.stride(2),
136 + output.stride(0),
137 + output.stride(1),
138 + output.stride(2),
139 + H,
140 + n_rows,
141 + D,
142 + padded_head_dim,
143 + _FP8_MIN,
144 + _FP8_MAX,
145 + SKIP_SCALE=skip_scale,
146 + BLOCK_M=block_m,
147 + BLOCK_N=block_n,
148 + num_warps=num_warps,
149 + )
150 +
151 + return output.view((*original_shape[:-1], padded_head_dim))
152 +
153 +
154 +def quantize_fp8_maybe_pad_head_dim(
155 + tensor: torch.Tensor,
156 + scale: torch.Tensor,
157 + fp8_quant: QuantFP8,
158 + skip_scale: bool = False,
159 +) -> torch.Tensor:
160 + """Quantize a 3D/4D tensor to FP8, padding head_dim to a multiple of 16
161 + only when needed.
162 +
163 + Accepts (S, H, D) or (B, S, H, D) input. Uses ``fp8_quant`` (a
164 + :class:`QuantFP8` CustomOp) when head_dim is already aligned to 16
165 + (no padding); otherwise falls back to a stride-aware Triton kernel
166 + that pads head_dim to a multiple of 16.
167 + """
168 + head_dim = tensor.shape[-1]
169 + if head_dim % 16 != 0:
170 + return quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=skip_scale)
171 +
172 + if skip_scale:
173 + return tensor.to(current_platform.fp8_dtype())
174 +
175 + # QuantFP8 expects 2D: flatten all dims except (H, D).
176 + orig_shape = tensor.shape
177 + total_tokens = tensor.numel() // (orig_shape[-1] * orig_shape[-2])
178 + tensor_2d = tensor.reshape(total_tokens, -1)
179 + fp8_tensor, _ = fp8_quant(tensor_2d, scale=scale)
180 + return fp8_tensor.reshape(orig_shape)
vllm/model_executor/layers/attention/mm_encoder_attention.py
+336 -15
@@ -1,13 +1,32 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import functools
5 +import json
6
7 import numpy as np
8 import torch
9
10 +from vllm.config import MultiModalConfig
11 +from vllm.kernels.triton.qkv_padded_fp8_quant import (
12 + quantize_fp8_maybe_pad_head_dim,
13 +)
14 from vllm.logger import init_logger
15 from vllm.model_executor.custom_op import CustomOp, maybe_get_oot_by_class
10 -from vllm.model_executor.models.vision import get_vit_attn_backend
16 +from vllm.model_executor.layers.quantization.input_quant_fp8 import (
17 + QuantFP8,
18 +)
19 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
20 + GroupShape,
21 + get_fp8_min_max,
22 +)
23 +from vllm.model_executor.models.vision import (
24 + get_multimodal_config,
25 + get_vit_attn_backend,
26 +)
27 +from vllm.utils.flashinfer import (
28 + is_flashinfer_cudnn_fp8_prefill_attn_supported,
29 +)
30 from vllm.utils.math_utils import round_up
31 from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
32 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -20,6 +39,108 @@ from vllm.v1.attention.ops.vit_attn_wrappers import (
39
40 logger = init_logger(__name__)
41
42 +_, _FP8_MAX = get_fp8_min_max()
43 +_FP8_AMAX_HISTORY_LEN = 16
44 +
45 +# Module-level state for auto-saving dynamic scales. The save is a one-shot
46 +# triggered by the first layer whose amax buffer wraps. Path and margin are
47 +# captured during layer init (set_current_vllm_config context only lives
48 +# across model init, not forward passes).
49 +_fp8_scale_save_path: str | None = None
50 +_fp8_scale_save_margin: float = MultiModalConfig.mm_encoder_fp8_scale_save_margin
51 +_fp8_saved_scale_refs: dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}
52 +
53 +
54 +@functools.cache
55 +def _load_fp8_scales_file(path: str | None) -> dict[str, dict[str, float]]:
56 + """Load per-layer FP8 Q/K/V scales from a JSON file. Results are cached.
57 +
58 + Expected format (keys ``q_scale`` / ``k_scale`` / ``v_scale`` also accepted)::
59 +
60 + {
61 + "visual.blocks.0.attn.attn": {"q": 224.0, "k": 198.0, "v": 210.0},
62 + "visual.blocks.1.attn.attn": {"q": 218.0, "k": 195.0, "v": 207.0},
63 + }
64 +
65 + To produce such a file, run with ``mm_encoder_fp8_scale_save_path`` set.
66 + """
67 + if path is None:
68 + return {}
69 +
70 + with open(path, encoding="utf-8") as f:
71 + data = json.load(f)
72 +
73 + # Handle nested "layers" format
74 + if "layers" in data and isinstance(data["layers"], dict):
75 + data = data["layers"]
76 +
77 + scales: dict[str, dict[str, float]] = {}
78 + for layer_name, layer_scales in data.items():
79 + if not isinstance(layer_scales, dict):
80 + continue
81 + q = layer_scales.get("q", layer_scales.get("q_scale"))
82 + k = layer_scales.get("k", layer_scales.get("k_scale"))
83 + v = layer_scales.get("v", layer_scales.get("v_scale"))
84 + if q is not None and k is not None and v is not None:
85 + q_f, k_f, v_f = float(q), float(k), float(v)
86 + if q_f <= 0 or k_f <= 0 or v_f <= 0:
87 + raise ValueError(
88 + f"FP8 scales must be positive, got q={q_f}, "
89 + f"k={k_f}, v={v_f} for layer '{layer_name}'"
90 + )
91 + scales[layer_name] = {"q": q_f, "k": k_f, "v": v_f}
92 +
93 + logger.info_once(
94 + "Loaded FP8 attention scales from %s (%d layers)", path, len(scales)
95 + )
96 + return scales
97 +
98 +
99 +def _maybe_save_fp8_scales(
100 + layer_name: str,
101 + q_scale: torch.Tensor,
102 + k_scale: torch.Tensor,
103 + v_scale: torch.Tensor,
104 + buffer_wrapped: bool,
105 +) -> None:
106 + """Accumulate a layer's scale tensors; on the first amax buffer wrap,
107 + dump all accumulated scales to ``mm_encoder_fp8_scale_save_path``.
108 +
109 + No-op unless auto-save is configured. Tensor references are stored on
110 + every call (no GPU->CPU sync); ``.item()`` is only called at the single
111 + save point to avoid stalling the forward path.
112 + """
113 + global _fp8_scale_save_path
114 + # Fast path: auto-save either disabled or already finished. Path is
115 + # captured at layer init and cleared once the save fires.
116 + if _fp8_scale_save_path is None:
117 + return
118 +
119 + # Stash scale tensor refs (no GPU->CPU sync yet); wait until the amax
120 + # history has seen a full cycle before committing scales to disk.
121 + _fp8_saved_scale_refs[layer_name] = (q_scale, k_scale, v_scale)
122 + if not buffer_wrapped:
123 + return
124 +
125 + # Buffer just wrapped for the first time: materialize scales (with
126 + # safety margin) and dump to disk. Clearing _fp8_scale_save_path
127 + # makes this a one-shot across all layers.
128 + path, margin = _fp8_scale_save_path, _fp8_scale_save_margin
129 + scales = {
130 + name: {
131 + "q": q.item() * margin,
132 + "k": k.item() * margin,
133 + "v": v.item() * margin,
134 + }
135 + for name, (q, k, v) in _fp8_saved_scale_refs.items()
136 + }
137 + _fp8_scale_save_path = None
138 + _fp8_saved_scale_refs.clear()
139 + with open(path, "w", encoding="utf-8") as f:
140 + json.dump(scales, f, indent=2)
141 + logger.info("Saved FP8 scales (%d layers) to %s", len(scales), path)
142 +
143 +
144 # Batch buckets for cuDNN graph caching.
145 # Graphs use batch size and max sequence length as cache key.
146 # This avoids creating a new graph for each unique set of
@@ -148,27 +269,47 @@ class MMEncoderAttention(CustomOp):
269 hidden_size: int,
270 tp_size: int,
271 device: torch.device,
272 + fp8_padded_hidden_size: int | None = None,
273 ) -> torch.Tensor:
274 if (oot_class := maybe_get_oot_by_class(cls)) is not cls:
275 return oot_class.maybe_recompute_cu_seqlens( # type: ignore[attr-defined]
154 - attn_backend, cu_seqlens, hidden_size, tp_size, device
276 + attn_backend,
277 + cu_seqlens,
278 + hidden_size,
279 + tp_size,
280 + device,
281 + fp8_padded_hidden_size=fp8_padded_hidden_size,
282 )
283
284 if attn_backend == AttentionBackendEnum.FLASHINFER:
285 batch_size = len(cu_seqlens) - 1
159 - scale = hidden_size // tp_size
160 - cu_seqlens = cu_seqlens * scale
286
162 - cu_seqlens_qko = cu_seqlens
163 - cu_seqlens_v = cu_seqlens * 3
164 -
165 - cu_seqlens_qko = add_padding_to_seqlens(
166 - cu_seqlens_qko, batch_size, cu_seqlens_qko[-1]
167 - )
168 - cu_seqlens_v = add_padding_to_seqlens(
169 - cu_seqlens_v, batch_size, cu_seqlens_v[-1]
170 - )
171 - cu_seqlens = np.concatenate([cu_seqlens_qko, cu_seqlens_v])
287 + if fp8_padded_hidden_size is not None:
288 + # FP8 path: after quantization Q/K/V are each independent
289 + # contiguous tensors with stride H * padded_D per token.
290 + # All sections use the same element stride.
291 + scale = fp8_padded_hidden_size // tp_size
292 + cu_seqlens = cu_seqlens * scale
293 + cu_seqlens_padded = add_padding_to_seqlens(
294 + cu_seqlens, batch_size, cu_seqlens[-1]
295 + )
296 + cu_seqlens = np.concatenate([cu_seqlens_padded, cu_seqlens_padded])
297 + else:
298 + # BF16 path: Q/K/V are non-contiguous views into shared
299 + # buffers. V section has 3x stride from interleaved QKV.
300 + scale = hidden_size // tp_size
301 + cu_seqlens = cu_seqlens * scale
302 +
303 + cu_seqlens_qko = cu_seqlens
304 + cu_seqlens_v = cu_seqlens * 3
305 +
306 + cu_seqlens_qko = add_padding_to_seqlens(
307 + cu_seqlens_qko, batch_size, cu_seqlens_qko[-1]
308 + )
309 + cu_seqlens_v = add_padding_to_seqlens(
310 + cu_seqlens_v, batch_size, cu_seqlens_v[-1]
311 + )
312 + cu_seqlens = np.concatenate([cu_seqlens_qko, cu_seqlens_v])
313
314 cu_seqlens = torch.from_numpy(cu_seqlens).to(device, non_blocking=True)
315 return cu_seqlens
@@ -206,6 +347,7 @@ class MMEncoderAttention(CustomOp):
347 # During model initialization, the default dtype is set as the model
348 # weight and activation dtype.
349 dtype = torch.get_default_dtype()
350 + self.dtype = dtype
351
352 # Get device-specific vision attention backend.
353 self.attn_backend = get_vit_attn_backend(
@@ -229,6 +371,113 @@ class MMEncoderAttention(CustomOp):
371
372 logger.info_once(f"Using {self.attn_backend} for MMEncoderAttention.")
373
374 + self._init_fp8_state()
375 +
376 + def _init_fp8_state(self) -> None:
377 + """Initialize FP8 attention state from multimodal config.
378 +
379 + No-op if FP8 is not requested. Raises ``ValueError`` if FP8 is
380 + requested but the platform does not support it.
381 + """
382 + # Populate defaults so ``_forward_flashinfer`` can
383 + # check ``self.fp8_enabled`` and others without AttributeError.
384 + self.fp8_enabled = False
385 + self._fp8_dynamic_scale = False
386 + self.fp8_quant: QuantFP8 | None = None
387 + self.skip_scale_q = False
388 + self.skip_scale_k = False
389 + self.skip_scale_v = False
390 +
391 + mm_cfg = get_multimodal_config()
392 + if mm_cfg is None or mm_cfg.mm_encoder_attn_dtype != "fp8":
393 + return
394 +
395 + # FP8 path
396 + if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
397 + raise ValueError(
398 + "mm_encoder_attn_dtype='fp8' requires the FlashInfer "
399 + "cuDNN backend with cuDNN >= 9.17.1 on a GPU with native "
400 + "FP8 support."
401 + )
402 +
403 + self.fp8_enabled = True
404 + self._fp8_dynamic_scale = mm_cfg.mm_encoder_fp8_scale_path is None
405 + self.fp8_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
406 +
407 + # Register buffers pre-device-move; values populated in
408 + # process_weights_after_loading. Shape (1, 1, 1, 1) is required by cuDNN.
409 + for attr in ("_fp8_q_scale", "_fp8_k_scale", "_fp8_v_scale"):
410 + self.register_buffer(
411 + attr, torch.ones(1, dtype=torch.float32).view(1, 1, 1, 1)
412 + )
413 + if self._fp8_dynamic_scale:
414 + for attr in ("_fp8_q_amax", "_fp8_k_amax", "_fp8_v_amax"):
415 + self.register_buffer(
416 + attr,
417 + torch.zeros(_FP8_AMAX_HISTORY_LEN, dtype=torch.float32),
418 + persistent=False,
419 + )
420 + self._fp8_amax_pos = 0
421 +
422 + # Capture auto-save config now: the VllmConfig context only lives
423 + # across model init, not forward passes, so ``_maybe_save_fp8_scales``
424 + # reads these globals instead of re-querying ``get_multimodal_config``.
425 + if (
426 + mm_cfg.mm_encoder_fp8_scale_save_path is not None
427 + and self._fp8_dynamic_scale
428 + ):
429 + global _fp8_scale_save_path, _fp8_scale_save_margin
430 + _fp8_scale_save_path = mm_cfg.mm_encoder_fp8_scale_save_path
431 + _fp8_scale_save_margin = mm_cfg.mm_encoder_fp8_scale_save_margin
432 +
433 + def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
434 + """Populate FP8 scale buffers after weights are loaded.
435 +
436 + ``act_dtype`` matches the signature used by :class:`Attention` and
437 + :class:`MLAAttention` for the loader auto-scan but is unused:
438 + FP8 scales are always float32.
439 + """
440 + if not self.fp8_enabled:
441 + return
442 +
443 + mm_cfg = get_multimodal_config()
444 + scale_path = mm_cfg.mm_encoder_fp8_scale_path if mm_cfg is not None else None
445 + if scale_path is None:
446 + logger.info_once(
447 + "FP8 attention enabled with dynamic scaling "
448 + "(no scale file provided). Scales will adapt from "
449 + "observed Q/K/V amax values (history_len=%d).",
450 + _FP8_AMAX_HISTORY_LEN,
451 + )
452 + return
453 +
454 + all_scales = _load_fp8_scales_file(scale_path)
455 + layer_scales = all_scales.get(self.layer_name)
456 + if layer_scales is None:
457 + raise ValueError(
458 + "FP8 attention enabled but scales not found for layer "
459 + f"'{self.layer_name}' in {scale_path}. "
460 + f"Available layers: {list(all_scales.keys())}"
461 + )
462 +
463 + for attr, key in (
464 + ("_fp8_q_scale", "q"),
465 + ("_fp8_k_scale", "k"),
466 + ("_fp8_v_scale", "v"),
467 + ):
468 + getattr(self, attr).fill_(layer_scales[key])
469 + self.skip_scale_q = layer_scales["q"] == 1.0
470 + self.skip_scale_k = layer_scales["k"] == 1.0
471 + self.skip_scale_v = layer_scales["v"] == 1.0
472 +
473 + logger.debug(
474 + "FP8 attention enabled for %s: q=%.4f, k=%.4f, v=%.4f",
475 + self.layer_name if self.layer_name else "MMEncoderAttention",
476 + layer_scales["q"],
477 + layer_scales["k"],
478 + layer_scales["v"],
479 + )
480 +
481 @classmethod
482 def enabled(cls) -> bool:
483 return True
@@ -353,6 +602,44 @@ class MMEncoderAttention(CustomOp):
602 output = output.reshape(bsz, q_len, -1)
603 return output
604
605 + @torch.no_grad()
606 + def _record_amax_and_update_scales(
607 + self,
608 + query: torch.Tensor,
609 + key: torch.Tensor,
610 + value: torch.Tensor,
611 + ) -> None:
612 + """Record Q/K/V amax into circular history and recompute scales.
613 +
614 + All work stays on GPU with no device-to-host sync. The Python-side
615 + history position counter is mutated, so this method must NOT be
616 + called inside CUDA graph capture/replay. When CUDA graphs are
617 + used for the encoder, dynamic scaling should be disabled by
618 + providing a static scale file via --mm-encoder-fp8-scale-path.
619 + """
620 + pos = self._fp8_amax_pos
621 + self._fp8_amax_pos = (pos + 1) % _FP8_AMAX_HISTORY_LEN
622 +
623 + for tensor, amax_buf, scale_buf in (
624 + (query, self._fp8_q_amax, self._fp8_q_scale),
625 + (key, self._fp8_k_amax, self._fp8_k_scale),
626 + (value, self._fp8_v_amax, self._fp8_v_scale),
627 + ):
628 + amax_buf[pos] = tensor.amax()
629 + max_amax = amax_buf.max()
630 + scale_buf.fill_(
631 + torch.clamp(max_amax, min=torch.finfo(torch.float32).tiny) / _FP8_MAX
632 + )
633 +
634 + buffer_wrapped = self._fp8_amax_pos == 0 and pos == _FP8_AMAX_HISTORY_LEN - 1
635 + _maybe_save_fp8_scales(
636 + self.layer_name,
637 + self._fp8_q_scale,
638 + self._fp8_k_scale,
639 + self._fp8_v_scale,
640 + buffer_wrapped,
641 + )
642 +
643 def _forward_flashinfer(
644 self,
645 query: torch.Tensor,
@@ -363,7 +650,32 @@ class MMEncoderAttention(CustomOp):
650 sequence_lengths: torch.Tensor
651 | None = None, # Only used for FlashInfer CuDNN backend
652 ) -> torch.Tensor:
366 - return vit_flashinfer_wrapper(
653 + if self.fp8_enabled:
654 + assert self.fp8_quant is not None
655 +
656 + if self._fp8_dynamic_scale:
657 + self._record_amax_and_update_scales(query, key, value)
658 +
659 + query = quantize_fp8_maybe_pad_head_dim(
660 + query,
661 + self._fp8_q_scale,
662 + skip_scale=self.skip_scale_q,
663 + fp8_quant=self.fp8_quant,
664 + )
665 + key = quantize_fp8_maybe_pad_head_dim(
666 + key,
667 + self._fp8_k_scale,
668 + skip_scale=self.skip_scale_k,
669 + fp8_quant=self.fp8_quant,
670 + )
671 + value = quantize_fp8_maybe_pad_head_dim(
672 + value,
673 + self._fp8_v_scale,
674 + skip_scale=self.skip_scale_v,
675 + fp8_quant=self.fp8_quant,
676 + )
677 +
678 + output = vit_flashinfer_wrapper(
679 q=query,
680 k=key,
681 v=value,
@@ -372,8 +684,17 @@ class MMEncoderAttention(CustomOp):
684 cu_seqlens=cu_seqlens,
685 max_seqlen=max_seqlen,
686 sequence_lengths=sequence_lengths,
687 + q_scale=self._fp8_q_scale if self.fp8_enabled else None,
688 + k_scale=self._fp8_k_scale if self.fp8_enabled else None,
689 + v_scale=self._fp8_v_scale if self.fp8_enabled else None,
690 + o_data_type=self.dtype if self.fp8_enabled else None,
691 )
692
693 + if self.fp8_enabled and output.shape[-1] != self.head_size:
694 + output = output[..., : self.head_size].contiguous()
695 +
696 + return output
697 +
698 def forward_native(
699 self,
700 query: torch.Tensor,
vllm/model_executor/model_loader/utils.py
+9 -5
@@ -15,7 +15,11 @@ from typing_extensions import assert_never
15 import vllm.envs as envs
16 from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
17 from vllm.logger import init_logger
18 -from vllm.model_executor.layers.attention import Attention, MLAAttention
18 +from vllm.model_executor.layers.attention import (
19 + Attention,
20 + MLAAttention,
21 + MMEncoderAttention,
22 +)
23 from vllm.model_executor.layers.quantization.base_config import (
24 QuantizationConfig,
25 QuantizeMethodBase,
@@ -106,12 +110,12 @@ def process_weights_after_loading(
110 with device_loading_context(module, target_device):
111 quant_method.process_weights_after_loading(module)
112
109 - # Initialize post-load attention weights for both Attention and MLA.
113 + # Initialize post-load attention weights for Attention, MLA, and MM encoder.
114 # NOTE: Happens after other modules so we can easily decompress weights.
115 for _, module in model.named_modules():
112 - if isinstance(module, (Attention, MLAAttention)) and hasattr(
113 - module, "process_weights_after_loading"
114 - ):
116 + if isinstance(
117 + module, (Attention, MLAAttention, MMEncoderAttention)
118 + ) and hasattr(module, "process_weights_after_loading"):
119 # TODO(lucas): see if there is a way to unify the signatures
120 # of process_weights_after_loading
121 with device_loading_context(module, target_device):
vllm/model_executor/models/qwen3_vl.py
+9
@@ -136,6 +136,7 @@ from .utils import (
136 maybe_prefix,
137 )
138 from .vision import (
139 + get_fp8_padded_hidden_size,
140 get_vit_attn_backend,
141 is_vit_use_data_parallel,
142 run_dp_sharded_mrope_vision_model,
@@ -562,6 +563,13 @@ class Qwen3_VisionTransformer(nn.Module):
563
564 norm_layer = partial(nn.LayerNorm, eps=norm_eps)
565 head_dim = self.hidden_size // self.num_heads
566 +
567 + # FP8 attention: Q/K/V become independent contiguous tensors
568 + # after quantization, so cu_seqlens uses uniform stride (no 3x V).
569 + self.fp8_padded_hidden_size = get_fp8_padded_hidden_size(
570 + self.num_heads, head_dim
571 + )
572 +
573 self.rotary_pos_emb = get_rope(
574 head_size=head_dim,
575 max_position=8192,
@@ -776,6 +784,7 @@ class Qwen3_VisionTransformer(nn.Module):
784 self.hidden_size,
785 self.tp_size,
786 device,
787 + fp8_padded_hidden_size=self.fp8_padded_hidden_size,
788 )
789
790 return metadata
vllm/model_executor/models/vision.py
+32 -28
@@ -10,7 +10,7 @@ from typing import Final, Generic, Literal, Protocol, TypeAlias, TypeVar
10 import torch
11 from transformers import PretrainedConfig
12
13 -from vllm.config import MultiModalConfig, VllmConfig, get_current_vllm_config
13 +from vllm.config import MultiModalConfig, get_current_vllm_config_or_none
14 from vllm.distributed import (
15 get_tensor_model_parallel_rank,
16 get_tensor_model_parallel_world_size,
@@ -18,6 +18,7 @@ from vllm.distributed import (
18 )
19 from vllm.logger import init_logger
20 from vllm.platforms import current_platform
21 +from vllm.utils.math_utils import round_up
22 from vllm.v1.attention.backends.registry import AttentionBackendEnum
23
24 logger = init_logger(__name__)
@@ -102,45 +103,48 @@ def get_vit_attn_backend(
103 """
104 Get the attention backend for Vision Transformer.
105 """
105 - try:
106 - vllm_config: VllmConfig = get_current_vllm_config()
107 - model_config = vllm_config.model_config
108 - multimodal_config: MultiModalConfig | None = (
109 - model_config.multimodal_config if model_config is not None else None
110 - )
111 - except (AssertionError, AttributeError):
112 - multimodal_config = None
113 -
106 + mm_cfg = get_multimodal_config()
107 attn_backend_override = (
115 - multimodal_config.mm_encoder_attn_backend
116 - if multimodal_config is not None
117 - else None
108 + mm_cfg.mm_encoder_attn_backend if mm_cfg is not None else None
109 )
119 - attn_backend = _get_vit_attn_backend(
110 + return _get_vit_attn_backend(
111 head_size,
112 dtype,
113 attn_backend_override=attn_backend_override,
114 )
124 - return attn_backend
115 +
116 +
117 +def get_multimodal_config() -> MultiModalConfig | None:
118 + """Return the current ``MultiModalConfig``, or ``None`` when no engine
119 + config context is active (e.g., during unit tests) or when the current
120 + ``model_config`` does not carry a ``multimodal_config`` (e.g., minimal
121 + stubs used in tests)."""
122 + vllm_config = get_current_vllm_config_or_none()
123 + if vllm_config is None or vllm_config.model_config is None:
124 + return None
125 + return getattr(vllm_config.model_config, "multimodal_config", None)
126 +
127 +
128 +def get_fp8_padded_hidden_size(num_heads: int, head_dim: int) -> int | None:
129 + """Return the padded hidden size for FP8 ViT encoder attention, or
130 + ``None`` when FP8 is not enabled.
131 +
132 + cuDNN FP8 prefill attention requires ``head_dim`` to be a multiple of
133 + 16. For non-aligned ``head_dim`` (e.g. 72), Q/K/V are padded to the
134 + nearest multiple of 16.
135 + """
136 + mm_cfg = get_multimodal_config()
137 + if mm_cfg is None or mm_cfg.mm_encoder_attn_dtype != "fp8":
138 + return None
139 + return num_heads * round_up(head_dim, 16)
140
141
142 def is_vit_use_data_parallel():
143 """
144 Get the tensor parallel type for Vision Transformer.
145 """
131 - try:
132 - vllm_config: VllmConfig = get_current_vllm_config()
133 - model_config = vllm_config.model_config
134 - multimodal_config: MultiModalConfig | None = (
135 - model_config.multimodal_config if model_config is not None else None
136 - )
137 - except (AssertionError, AttributeError):
138 - multimodal_config = None
139 -
140 - mm_encoder_tp_mode = (
141 - multimodal_config.mm_encoder_tp_mode if multimodal_config is not None else None
142 - )
143 - return mm_encoder_tp_mode == "data"
146 + mm_cfg = get_multimodal_config()
147 + return mm_cfg is not None and mm_cfg.mm_encoder_tp_mode == "data"
148
149
150 VisionFeatureSelectStrategyStr = Literal["class", "default", "full"]
vllm/utils/flashinfer.py
+35
@@ -772,6 +772,40 @@ def should_use_flashinfer_for_blockscale_fp8_gemm(
772 return should_use_flashinfer
773
774
775 +_MIN_CUDNN_FP8 = 91701 # cuDNN >= 9.17.1 required for FP8 attention
776 +
777 +
778 +@functools.cache
779 +def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool:
780 + """Check if FP8 ViT attention is supported on this platform.
781 +
782 + Requires native FP8 hardware support, the FlashInfer cuDNN backend,
783 + and cuDNN >= 9.17.1.
784 + """
785 + from vllm.v1.attention.backends.registry import AttentionBackendEnum
786 +
787 + # cuDNN SDPA FP8 requires Hopper (SM 90) or newer.
788 + if not current_platform.has_device_capability(90):
789 + return False
790 +
791 + try:
792 + supported = current_platform.get_supported_vit_attn_backends()
793 + if AttentionBackendEnum.FLASHINFER not in supported:
794 + return False
795 + except (ImportError, AttributeError):
796 + return False
797 +
798 + try:
799 + import torch.backends.cudnn as cudnn
800 +
801 + if cudnn.is_available() and cudnn.version() < _MIN_CUDNN_FP8:
802 + return False
803 + except (ImportError, AttributeError):
804 + pass
805 +
806 + return True
807 +
808 +
809 __all__ = [
810 "has_flashinfer",
811 "flashinfer_trtllm_fp8_block_scale_moe",
@@ -803,4 +837,5 @@ __all__ = [
837 "flashinfer_fp8_blockscale_gemm",
838 "should_use_flashinfer_for_blockscale_fp8_gemm",
839 "is_flashinfer_fp8_blockscale_gemm_supported",
840 + "is_flashinfer_cudnn_fp8_prefill_attn_supported",
841 ]
vllm/v1/attention/ops/vit_attn_wrappers.py
+29 -2
@@ -279,6 +279,10 @@ def flashinfer_wrapper(
279 cu_seqlens: torch.Tensor | None = None,
280 max_seqlen: torch.Tensor | None = None,
281 sequence_lengths: torch.Tensor | None = None,
282 + q_scale: torch.Tensor | None = None,
283 + k_scale: torch.Tensor | None = None,
284 + v_scale: torch.Tensor | None = None,
285 + o_data_type: torch.dtype | None = None,
286 ) -> torch.Tensor:
287 from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
288
@@ -318,6 +322,10 @@ def flashinfer_wrapper(
322 batch_offsets_k=batch_offsets_qko,
323 batch_offsets_v=batch_offsets_v,
324 batch_offsets_o=batch_offsets_qko,
325 + q_scale=q_scale,
326 + k_scale=k_scale,
327 + v_scale=v_scale,
328 + o_data_type=o_data_type,
329 )
330
331 if is_reshaped:
@@ -335,8 +343,12 @@ def vit_flashinfer_wrapper_fake(
343 cu_seqlens: torch.Tensor | None = None,
344 max_seqlen: torch.Tensor | None = None,
345 sequence_lengths: torch.Tensor | None = None,
346 + q_scale: torch.Tensor | None = None,
347 + k_scale: torch.Tensor | None = None,
348 + v_scale: torch.Tensor | None = None,
349 + o_data_type: torch.dtype | None = None,
350 ) -> torch.Tensor:
339 - return torch.empty_like(q)
351 + return torch.empty_like(q, dtype=o_data_type or q.dtype)
352
353
354 direct_register_custom_op(
@@ -355,7 +367,22 @@ def vit_flashinfer_wrapper(
367 cu_seqlens: torch.Tensor | None = None,
368 max_seqlen: torch.Tensor | None = None,
369 sequence_lengths: torch.Tensor | None = None,
370 + q_scale: torch.Tensor | None = None,
371 + k_scale: torch.Tensor | None = None,
372 + v_scale: torch.Tensor | None = None,
373 + o_data_type: torch.dtype | None = None,
374 ) -> torch.Tensor:
375 return torch.ops.vllm.flashinfer_wrapper(
360 - q, k, v, scale, workspace_buffer, cu_seqlens, max_seqlen, sequence_lengths
376 + q,
377 + k,
378 + v,
379 + scale,
380 + workspace_buffer,
381 + cu_seqlens,
382 + max_seqlen,
383 + sequence_lengths,
384 + q_scale,
385 + k_scale,
386 + v_scale,
387 + o_data_type,
388 )