| 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 | ) |