| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Standard attention benchmark runner - shared utilities for non-MLA benchmarks. |
| 5 | |
| 6 | This module provides helpers for running standard attention backends |
| 7 | (FlashAttention, Triton, FlashInfer) with real vLLM integration. |
| 8 | """ |
| 9 | |
| 10 | import logging |
| 11 | import statistics |
| 12 | import types |
| 13 | from contextlib import contextmanager |
| 14 | from math import prod |
| 15 | |
| 16 | import torch |
| 17 | from batch_spec import parse_batch_spec, reorder_for_flashinfer |
| 18 | from common import ( |
| 19 | BenchmarkConfig, |
| 20 | BenchmarkResult, |
| 21 | MockLayer, |
| 22 | get_attention_scale, |
| 23 | run_do_bench, |
| 24 | run_ncu_profile, |
| 25 | ) |
| 26 | |
| 27 | from vllm.config import ( |
| 28 | CacheConfig, |
| 29 | CompilationConfig, |
| 30 | DeviceConfig, |
| 31 | LoadConfig, |
| 32 | ModelConfig, |
| 33 | ParallelConfig, |
| 34 | SchedulerConfig, |
| 35 | VllmConfig, |
| 36 | set_current_vllm_config, |
| 37 | ) |
| 38 | from vllm.platforms import current_platform |
| 39 | from vllm.v1.attention.backends.utils import ( |
| 40 | CommonAttentionMetadata, |
| 41 | get_supported_kv_cache_layouts, |
| 42 | resolve_kv_cache_layout, |
| 43 | ) |
| 44 | from vllm.v1.kv_cache_interface import ( |
| 45 | FullAttentionSpec, |
| 46 | KVCacheLayout, |
| 47 | KVCacheTensor, |
| 48 | compute_layer_kv_cache_shape_bytes, |
| 49 | compute_layout_strides, |
| 50 | create_kv_cache_views, |
| 51 | ) |
| 52 | |
| 53 | # ============================================================================ |
| 54 | # Backend Configuration |
| 55 | # ============================================================================ |
| 56 | |
| 57 | |
| 58 | def _get_backend_config(backend: str) -> dict: |
| 59 | """Get backend configuration from AttentionBackendEnum. |
| 60 | |
| 61 | Args: |
| 62 | backend: Backend name matching AttentionBackendEnum exactly |
| 63 | (e.g., "FLASH_ATTN", "TRITON_ATTN", "FLASHINFER") |
| 64 | |
| 65 | Returns: |
| 66 | Dict with backend_class |
| 67 | |
| 68 | """ |
| 69 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 70 | |
| 71 | try: |
| 72 | backend_enum = AttentionBackendEnum[backend] |
| 73 | backend_class = backend_enum.get_class() |
| 74 | except (KeyError, ValueError) as e: |
| 75 | valid_backends = [b.name for b in AttentionBackendEnum if b.name != "CUSTOM"] |
| 76 | raise ValueError( |
| 77 | f"Unknown backend: {backend}. Valid backends: {valid_backends}" |
| 78 | ) from e |
| 79 | |
| 80 | return {"backend_class": backend_class} |
| 81 | |
| 82 | |
| 83 | @contextmanager |
| 84 | def log_warnings_and_errors_only(): |
| 85 | """Temporarily set vLLM logger to WARNING level.""" |
| 86 | logger = logging.getLogger("vllm") |
| 87 | old_level = logger.level |
| 88 | logger.setLevel(logging.WARNING) |
| 89 | try: |
| 90 | yield |
| 91 | finally: |
| 92 | logger.setLevel(old_level) |
| 93 | |
| 94 | |
| 95 | # ============================================================================ |
| 96 | # Metadata Building Helpers |
| 97 | # ============================================================================ |
| 98 | |
| 99 | |
| 100 | def _build_common_attn_metadata( |
| 101 | q_lens: list[int], |
| 102 | kv_lens: list[int], |
| 103 | block_size: int, |
| 104 | device: torch.device, |
| 105 | ) -> CommonAttentionMetadata: |
| 106 | """Build CommonAttentionMetadata from query/kv lengths.""" |
| 107 | batch_size = len(q_lens) |
| 108 | total_tokens = sum(q_lens) |
| 109 | |
| 110 | query_start_loc = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) |
| 111 | query_start_loc[1:] = torch.tensor(q_lens, dtype=torch.int32, device=device).cumsum( |
| 112 | 0 |
| 113 | ) |
| 114 | query_start_loc_cpu = query_start_loc.cpu() |
| 115 | |
| 116 | seq_lens = torch.tensor(kv_lens, dtype=torch.int32, device=device) |
| 117 | max_seq_len = int(seq_lens.max().item()) |
| 118 | |
| 119 | max_blocks = (max(kv_lens) + block_size - 1) // block_size |
| 120 | num_blocks = batch_size * max_blocks |
| 121 | block_table_tensor = torch.arange( |
| 122 | num_blocks, dtype=torch.int32, device=device |
| 123 | ).view(batch_size, max_blocks) |
| 124 | slot_mapping = torch.arange(total_tokens, dtype=torch.int64, device=device) |
| 125 | |
| 126 | max_query_len = max(q_lens) |
| 127 | |
| 128 | return CommonAttentionMetadata( |
| 129 | query_start_loc=query_start_loc, |
| 130 | query_start_loc_cpu=query_start_loc_cpu, |
| 131 | seq_lens=seq_lens, |
| 132 | num_reqs=batch_size, |
| 133 | num_actual_tokens=total_tokens, |
| 134 | max_query_len=max_query_len, |
| 135 | max_seq_len=max_seq_len, |
| 136 | block_table_tensor=block_table_tensor, |
| 137 | slot_mapping=slot_mapping, |
| 138 | causal=True, |
| 139 | ) |
| 140 | |
| 141 | |
| 142 | def _create_vllm_config( |
| 143 | config: BenchmarkConfig, |
| 144 | max_num_blocks: int, |
| 145 | ) -> VllmConfig: |
| 146 | """Create a VllmConfig for benchmarking with mock model methods.""" |
| 147 | model_config = ModelConfig( |
| 148 | model="meta-llama/Meta-Llama-3-8B", |
| 149 | tokenizer="meta-llama/Meta-Llama-3-8B", |
| 150 | trust_remote_code=False, |
| 151 | dtype="auto", # Use model's native dtype |
| 152 | seed=0, |
| 153 | max_model_len=1024, |
| 154 | ) |
| 155 | |
| 156 | cache_config = CacheConfig( |
| 157 | block_size=config.block_size, |
| 158 | cache_dtype=config.kv_cache_dtype, |
| 159 | ) |
| 160 | cache_config.num_gpu_blocks = max_num_blocks |
| 161 | cache_config.num_cpu_blocks = 0 |
| 162 | |
| 163 | parallel_config = ParallelConfig(tensor_parallel_size=1) |
| 164 | scheduler_config = SchedulerConfig( |
| 165 | max_num_seqs=256, |
| 166 | max_num_batched_tokens=8192, |
| 167 | max_model_len=8192, |
| 168 | is_encoder_decoder=False, |
| 169 | enable_chunked_prefill=True, |
| 170 | ) |
| 171 | device_config = DeviceConfig() |
| 172 | load_config = LoadConfig() |
| 173 | compilation_config = CompilationConfig() |
| 174 | |
| 175 | # Add mock methods for benchmark config values |
| 176 | model_config.get_num_layers = types.MethodType( |
| 177 | lambda self: config.num_layers, model_config |
| 178 | ) |
| 179 | model_config.get_sliding_window_for_layer = types.MethodType( |
| 180 | lambda self, i: None, model_config |
| 181 | ) |
| 182 | model_config.get_logits_soft_cap_for_layer = types.MethodType( |
| 183 | lambda self, i: 0.0, model_config |
| 184 | ) |
| 185 | model_config.get_sm_scale_for_layer = types.MethodType( |
| 186 | lambda self, i: 1.0 / config.head_dim**0.5, model_config |
| 187 | ) |
| 188 | model_config.get_num_attention_heads = types.MethodType( |
| 189 | lambda self, parallel_config=None: config.num_q_heads, model_config |
| 190 | ) |
| 191 | model_config.get_num_kv_heads = types.MethodType( |
| 192 | lambda self, parallel_config=None: config.num_kv_heads, model_config |
| 193 | ) |
| 194 | model_config.get_head_size = types.MethodType( |
| 195 | lambda self: config.head_dim, model_config |
| 196 | ) |
| 197 | model_config.get_sliding_window = types.MethodType(lambda self: None, model_config) |
| 198 | |
| 199 | return VllmConfig( |
| 200 | model_config=model_config, |
| 201 | cache_config=cache_config, |
| 202 | parallel_config=parallel_config, |
| 203 | scheduler_config=scheduler_config, |
| 204 | device_config=device_config, |
| 205 | load_config=load_config, |
| 206 | compilation_config=compilation_config, |
| 207 | ) |
| 208 | |
| 209 | |
| 210 | # ============================================================================ |
| 211 | # Backend Initialization |
| 212 | # ============================================================================ |
| 213 | |
| 214 | |
| 215 | def _create_backend_impl( |
| 216 | backend_cfg: dict, |
| 217 | config: BenchmarkConfig, |
| 218 | device: torch.device, |
| 219 | dtype: torch.dtype, |
| 220 | ): |
| 221 | """Create backend implementation instance.""" |
| 222 | backend_class = backend_cfg["backend_class"] |
| 223 | |
| 224 | scale = get_attention_scale(config.head_dim) |
| 225 | |
| 226 | # Set v_head_dim for diff-headdim backends. Always reset (defaulting to |
| 227 | # head_dim) so a prior run's value doesn't leak into this one via the |
| 228 | # backend's class-level state. |
| 229 | if hasattr(backend_class, "set_head_size_v"): |
| 230 | v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim |
| 231 | backend_class.set_head_size_v(v_dim) |
| 232 | |
| 233 | impl = backend_class.get_impl_cls()( |
| 234 | num_heads=config.num_q_heads, |
| 235 | head_size=config.head_dim, |
| 236 | scale=scale, |
| 237 | num_kv_heads=config.num_kv_heads, |
| 238 | alibi_slopes=None, |
| 239 | sliding_window=None, |
| 240 | kv_cache_dtype=config.kv_cache_dtype, |
| 241 | ) |
| 242 | |
| 243 | kv_cache_spec = FullAttentionSpec( |
| 244 | block_size=config.block_size, |
| 245 | num_kv_heads=config.num_kv_heads, |
| 246 | head_size=config.head_dim, |
| 247 | dtype=dtype, |
| 248 | ) |
| 249 | |
| 250 | layer = MockLayer(device, kv_cache_spec=kv_cache_spec) |
| 251 | |
| 252 | return backend_class, impl, layer |
| 253 | |
| 254 | |
| 255 | def _create_metadata_builder( |
| 256 | backend_class, |
| 257 | kv_cache_spec: FullAttentionSpec, |
| 258 | vllm_config: VllmConfig, |
| 259 | device: torch.device, |
| 260 | backend_name: str = "", |
| 261 | ): |
| 262 | """Create metadata builder instance.""" |
| 263 | layer_names = ["layer_0"] |
| 264 | builder_cls = backend_class.get_builder_cls() |
| 265 | |
| 266 | # Flashinfer needs get_per_layer_parameters mocked since we don't have |
| 267 | # real model layers registered |
| 268 | if backend_name == "FLASHINFER": |
| 269 | import unittest.mock |
| 270 | |
| 271 | from vllm.v1.attention.backends.utils import PerLayerParameters |
| 272 | |
| 273 | def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls): |
| 274 | head_size = vllm_config.model_config.get_head_size() |
| 275 | return { |
| 276 | layer_name: PerLayerParameters( |
| 277 | window_left=-1, # No sliding window |
| 278 | logits_soft_cap=0.0, # No soft cap |
| 279 | sm_scale=1.0 / (head_size**0.5), # Standard scale |
| 280 | ) |
| 281 | for layer_name in layer_names |
| 282 | } |
| 283 | |
| 284 | with unittest.mock.patch( |
| 285 | "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters", |
| 286 | mock_get_per_layer_parameters, |
| 287 | ): |
| 288 | return builder_cls( |
| 289 | kv_cache_spec=kv_cache_spec, |
| 290 | layer_names=layer_names, |
| 291 | vllm_config=vllm_config, |
| 292 | device=device, |
| 293 | ) |
| 294 | |
| 295 | return builder_cls( |
| 296 | kv_cache_spec=kv_cache_spec, |
| 297 | layer_names=layer_names, |
| 298 | vllm_config=vllm_config, |
| 299 | device=device, |
| 300 | ) |
| 301 | |
| 302 | |
| 303 | # ============================================================================ |
| 304 | # Tensor Creation Helpers |
| 305 | # ============================================================================ |
| 306 | |
| 307 | |
| 308 | def _create_input_tensors( |
| 309 | config: BenchmarkConfig, |
| 310 | total_q: int, |
| 311 | device: torch.device, |
| 312 | dtype: torch.dtype, |
| 313 | quantize_query: bool = False, |
| 314 | ) -> tuple: |
| 315 | """Create Q, K, V input tensors for all layers. |
| 316 | |
| 317 | When quantize_query is True, queries are cast to fp8 to match backends |
| 318 | that require query/key/value dtype consistency. |
| 319 | """ |
| 320 | q_dtype = dtype |
| 321 | if quantize_query: |
| 322 | from vllm.platforms import current_platform |
| 323 | |
| 324 | q_dtype = current_platform.fp8_dtype() |
| 325 | v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim |
| 326 | q_list = [ |
| 327 | torch.randn( |
| 328 | total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype |
| 329 | ).to(q_dtype) |
| 330 | for _ in range(config.num_layers) |
| 331 | ] |
| 332 | k_list = [ |
| 333 | torch.randn( |
| 334 | total_q, config.num_kv_heads, config.head_dim, device=device, dtype=dtype |
| 335 | ) |
| 336 | for _ in range(config.num_layers) |
| 337 | ] |
| 338 | v_list = [ |
| 339 | torch.randn(total_q, config.num_kv_heads, v_dim, device=device, dtype=dtype) |
| 340 | for _ in range(config.num_layers) |
| 341 | ] |
| 342 | return q_list, k_list, v_list |
| 343 | |
| 344 | |
| 345 | def _create_kv_cache( |
| 346 | config: BenchmarkConfig, |
| 347 | max_num_blocks: int, |
| 348 | backend_class, |
| 349 | device: torch.device, |
| 350 | dtype: torch.dtype, |
| 351 | layout: KVCacheLayout, |
| 352 | ) -> list: |
| 353 | """Create KV cache tensors for all layers using the standard allocator.""" |
| 354 | if config.kv_cache_dtype.startswith("fp8"): |
| 355 | cache_dtype = current_platform.fp8_dtype() |
| 356 | else: |
| 357 | cache_dtype = dtype |
| 358 | spec = FullAttentionSpec( |
| 359 | block_size=config.block_size, |
| 360 | num_kv_heads=config.num_kv_heads, |
| 361 | head_size=config.head_dim, |
| 362 | dtype=cache_dtype, |
| 363 | ) |
| 364 | # Apply the backend's page customization, as the worker does for the real spec. |
| 365 | spec = backend_class.customize_spec(spec) |
| 366 | total_bytes = ( |
| 367 | prod(compute_layer_kv_cache_shape_bytes(spec, max_num_blocks)) |
| 368 | * config.num_layers |
| 369 | ) |
| 370 | buf = torch.zeros(total_bytes, device=device, dtype=torch.int8) |
| 371 | layer_stride, block_stride, _, _, _ = compute_layout_strides( |
| 372 | spec, max_num_blocks, config.num_layers, layout |
| 373 | ) |
| 374 | tensor = KVCacheTensor( |
| 375 | size=total_bytes, |
| 376 | layers=[str(i) for i in range(config.num_layers)], |
| 377 | layer_stride=layer_stride, |
| 378 | block_stride=block_stride, |
| 379 | ) |
| 380 | return create_kv_cache_views( |
| 381 | buf, |
| 382 | spec, |
| 383 | max_num_blocks, |
| 384 | layout, |
| 385 | tensor, |
| 386 | ) |
| 387 | |
| 388 | |
| 389 | # ============================================================================ |
| 390 | # Benchmark Execution |
| 391 | # ============================================================================ |
| 392 | |
| 393 | |
| 394 | def _run_single_benchmark( |
| 395 | config: BenchmarkConfig, |
| 396 | impl, |
| 397 | layer, |
| 398 | q_list: list, |
| 399 | k_list: list, |
| 400 | v_list: list, |
| 401 | cache_list: list, |
| 402 | attn_metadata, |
| 403 | device: torch.device, |
| 404 | dtype: torch.dtype, |
| 405 | ) -> tuple: |
| 406 | """Run single benchmark using triton's do_bench_cudagraph/do_bench. |
| 407 | |
| 408 | Returns: |
| 409 | (timing_stats, mem_stats) where timing_stats is a dict with |
| 410 | mean/std/min/max in seconds per layer. |
| 411 | |
| 412 | """ |
| 413 | total_q = q_list[0].shape[0] |
| 414 | v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim |
| 415 | out = torch.empty(total_q, config.num_q_heads, v_dim, device=device, dtype=dtype) |
| 416 | |
| 417 | def benchmark_fn(): |
| 418 | for i in range(config.num_layers): |
| 419 | impl.forward( |
| 420 | layer, |
| 421 | q_list[i], |
| 422 | k_list[i], |
| 423 | v_list[i], |
| 424 | cache_list[i], |
| 425 | attn_metadata, |
| 426 | output=out, |
| 427 | ) |
| 428 | |
| 429 | if config.ncu_profile: |
| 430 | run_ncu_profile(benchmark_fn) |
| 431 | timing_stats = dict.fromkeys(("mean", "median", "std", "min", "max"), 0.0) |
| 432 | else: |
| 433 | all_ms = run_do_bench(benchmark_fn, config.use_cuda_graphs, config.warmup_ms) |
| 434 | |
| 435 | # Convert ms to seconds per layer |
| 436 | times = [t / 1000.0 / config.num_layers for t in all_ms] |
| 437 | timing_stats = { |
| 438 | "mean": statistics.mean(times), |
| 439 | "std": statistics.stdev(times) if len(times) > 1 else 0.0, |
| 440 | "min": min(times), |
| 441 | "max": max(times), |
| 442 | "median": statistics.median(times), |
| 443 | } |
| 444 | |
| 445 | mem_stats = {} |
| 446 | if config.profile_memory: |
| 447 | mem_stats = { |
| 448 | "allocated_mb": torch.accelerator.memory_allocated(device) / 1024**2, |
| 449 | "reserved_mb": torch.accelerator.memory_reserved(device) / 1024**2, |
| 450 | } |
| 451 | |
| 452 | return timing_stats, mem_stats |
| 453 | |
| 454 | |
| 455 | # ============================================================================ |
| 456 | # Public API |
| 457 | # ============================================================================ |
| 458 | |
| 459 | |
| 460 | def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: |
| 461 | """Run standard attention benchmark with real kernels. |
| 462 | |
| 463 | Supports: FLASH_ATTN, TRITON_ATTN, FLASHINFER |
| 464 | |
| 465 | Args: |
| 466 | config: Benchmark configuration |
| 467 | |
| 468 | Returns: |
| 469 | BenchmarkResult with timing and memory statistics |
| 470 | |
| 471 | """ |
| 472 | device = torch.device(config.device) |
| 473 | torch.accelerator.set_device_index(device) |
| 474 | |
| 475 | backend_cfg = _get_backend_config(config.backend) |
| 476 | |
| 477 | requests = parse_batch_spec(config.batch_spec) |
| 478 | |
| 479 | if config.backend == "FLASHINFER": |
| 480 | requests = reorder_for_flashinfer(requests) |
| 481 | |
| 482 | q_lens = [r.q_len for r in requests] |
| 483 | kv_lens = [r.kv_len for r in requests] |
| 484 | total_q = sum(q_lens) |
| 485 | max_kv = max(kv_lens) |
| 486 | batch_size = len(q_lens) |
| 487 | |
| 488 | # Calculate total blocks needed: batch_size * max_blocks_per_request |
| 489 | max_blocks_per_request = (max_kv + config.block_size - 1) // config.block_size |
| 490 | max_num_blocks = batch_size * max_blocks_per_request |
| 491 | |
| 492 | # Suppress vLLM logs during setup to reduce spam |
| 493 | with log_warnings_and_errors_only(): |
| 494 | # Create vllm_config first - uses model's native dtype via "auto" |
| 495 | vllm_config = _create_vllm_config(config, max_num_blocks) |
| 496 | dtype = vllm_config.model_config.dtype |
| 497 | |
| 498 | # Wrap everything in set_current_vllm_config context |
| 499 | # This is required for backends like flashinfer that need global config |
| 500 | with set_current_vllm_config(vllm_config): |
| 501 | backend_class, impl, layer = _create_backend_impl( |
| 502 | backend_cfg, config, device, dtype |
| 503 | ) |
| 504 | # Set KV cache layout if the backend requires a specific one |
| 505 | # (e.g., FlashInfer requires LBHNC on SM100/Blackwell for TRTLLM attention) |
| 506 | supported = get_supported_kv_cache_layouts([backend_class]) |
| 507 | layout = resolve_kv_cache_layout(vllm_config, [[m.name for m in supported]]) |
| 508 | |
| 509 | common_metadata = _build_common_attn_metadata( |
| 510 | q_lens, kv_lens, config.block_size, device |
| 511 | ) |
| 512 | |
| 513 | kv_cache_spec = FullAttentionSpec( |
| 514 | block_size=config.block_size, |
| 515 | num_kv_heads=config.num_kv_heads, |
| 516 | head_size=config.head_dim, |
| 517 | dtype=dtype, |
| 518 | ) |
| 519 | |
| 520 | builder = _create_metadata_builder( |
| 521 | backend_class, kv_cache_spec, vllm_config, device, config.backend |
| 522 | ) |
| 523 | |
| 524 | attn_metadata = builder.build( |
| 525 | common_prefix_len=0, |
| 526 | common_attn_metadata=common_metadata, |
| 527 | ) |
| 528 | |
| 529 | # Override num_splits for split-K testing (FlashAttention only) |
| 530 | if config.num_splits is not None and hasattr( |
| 531 | attn_metadata, "max_num_splits" |
| 532 | ): |
| 533 | attn_metadata.max_num_splits = config.num_splits |
| 534 | |
| 535 | # Only quantize queries when the impl supports it |
| 536 | quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr( |
| 537 | impl, "supports_quant_query_input", False |
| 538 | ) |
| 539 | q_list, k_list, v_list = _create_input_tensors( |
| 540 | config, total_q, device, dtype, quantize_query=quantize_query |
| 541 | ) |
| 542 | |
| 543 | cache_list = _create_kv_cache( |
| 544 | config, max_num_blocks, backend_class, device, dtype, layout |
| 545 | ) |
| 546 | |
| 547 | timing_stats, mem_stats = _run_single_benchmark( |
| 548 | config, |
| 549 | impl, |
| 550 | layer, |
| 551 | q_list, |
| 552 | k_list, |
| 553 | v_list, |
| 554 | cache_list, |
| 555 | attn_metadata, |
| 556 | device, |
| 557 | dtype, |
| 558 | ) |
| 559 | |
| 560 | mean_time = timing_stats["mean"] |
| 561 | throughput = total_q / mean_time if mean_time > 0 else 0 |
| 562 | |
| 563 | return BenchmarkResult( |
| 564 | config=config, |
| 565 | mean_time=mean_time, |
| 566 | median_time=timing_stats["median"], |
| 567 | std_time=timing_stats["std"], |
| 568 | min_time=timing_stats["min"], |
| 569 | max_time=timing_stats["max"], |
| 570 | throughput_tokens_per_sec=throughput, |
| 571 | memory_allocated_mb=mem_stats.get("allocated_mb"), |
| 572 | memory_reserved_mb=mem_stats.get("reserved_mb"), |
| 573 | ) |