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 8, 16:30 UTC
[Misc] Fix up attention benchmarks (#33810)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Lucas Wilkinson committed
Feb 9, 2026 at 06:42 UTC
d0d97e2974250edb61fbff6964e95a5b6d22d763
5 files changed
+219
-95
.buildkite/test_areas/benchmarks.yaml
+11
@@ -17,3 +17,14 @@ steps:
17
- tests/benchmarks/
18
commands:
19
- pytest -v -s benchmarks/
20
+
21
+- label: Attention Benchmarks Smoke Test (B200)
22
+ device: b200
23
+ num_gpus: 2
24
+ optional: true
25
+ timeout_in_minutes: 10
26
+ source_file_dependencies:
27
+ - benchmarks/attention_benchmarks/
28
+ - vllm/v1/attention/
29
+ commands:
30
+ - python benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1
benchmarks/attention_benchmarks/batch_spec.py
+37
@@ -229,3 +229,40 @@ def get_batch_stats(requests: list[BatchRequest]) -> dict:
229
sum(r.kv_len for r in requests) / len(requests) if requests else 0
230
),
231
}
232
+
233
+
234
+def get_batch_type(batch_spec: str, spec_decode_threshold: int = 8) -> str:
235
+ """
236
+ Classify a batch spec into a type string.
237
+
238
+ Args:
239
+ batch_spec: Batch specification string (e.g., "q2k", "8q1s1k", "2q2k_8q1s1k")
240
+ spec_decode_threshold: Max q_len to be considered spec-decode vs extend
241
+
242
+ Returns:
243
+ Type string: "prefill", "decode", "spec-decode", "extend", or "mixed (types...)"
244
+ """
245
+ requests = parse_batch_spec(batch_spec)
246
+
247
+ # Classify each request
248
+ types_present = set()
249
+ for req in requests:
250
+ if req.is_decode:
251
+ types_present.add("decode")
252
+ elif req.is_prefill:
253
+ types_present.add("prefill")
254
+ elif req.is_extend:
255
+ # Distinguish spec-decode (small q_len) from extend (chunked prefill)
256
+ if req.q_len <= spec_decode_threshold:
257
+ types_present.add("spec-decode")
258
+ else:
259
+ types_present.add("extend")
260
+
261
+ if len(types_present) == 1:
262
+ return types_present.pop()
263
+ elif len(types_present) > 1:
264
+ # Sort for consistent output
265
+ sorted_types = sorted(types_present)
266
+ return f"mixed ({'+'.join(sorted_types)})"
267
+ else:
268
+ return "unknown"
benchmarks/attention_benchmarks/common.py
+10
-3
@@ -12,6 +12,7 @@ from typing import Any
12
13
import numpy as np
14
import torch
15
+from batch_spec import get_batch_type, parse_batch_spec
16
from rich.console import Console
17
from rich.table import Table
18
@@ -316,12 +317,14 @@ class ResultsFormatter:
317
backends: List of backend names being compared
318
compare_to_fastest: Show percentage comparison to fastest
319
"""
319
- # Group by batch spec
320
+ # Group by batch spec, preserving first-occurrence order
321
by_spec = {}
322
+ specs_order = []
323
for r in results:
324
spec = r.config.batch_spec
325
if spec not in by_spec:
326
by_spec[spec] = {}
327
+ specs_order.append(spec)
328
by_spec[spec][r.config.backend] = r
329
330
# Create shortened backend names for display
@@ -337,6 +340,8 @@ class ResultsFormatter:
340
341
table = Table(title="Attention Benchmark Results")
342
table.add_column("Batch\nSpec", no_wrap=True)
343
+ table.add_column("Type", no_wrap=True)
344
+ table.add_column("Batch\nSize", justify="right", no_wrap=True)
345
346
multi = len(backends) > 1
347
for backend in backends:
@@ -350,12 +355,14 @@ class ResultsFormatter:
355
table.add_column(col_rel, justify="right", no_wrap=False)
356
357
# Add rows
353
- for spec in sorted(by_spec.keys()):
358
+ for spec in specs_order:
359
spec_results = by_spec[spec]
360
times = {b: r.mean_time for b, r in spec_results.items() if r.success}
361
best_time = min(times.values()) if times else 0.0
362
358
- row = [spec]
363
+ batch_type = get_batch_type(spec)
364
+ batch_size = len(parse_batch_spec(spec))
365
+ row = [spec, batch_type, str(batch_size)]
366
for backend in backends:
367
if backend in spec_results:
368
r = spec_results[backend]
benchmarks/attention_benchmarks/configs/standard_attention.yaml
+10
-2
@@ -25,10 +25,18 @@ batch_specs:
25
- "4q1k_16q1s2k" # 4 prefill + 16 decode
26
- "2q4k_32q1s1k" # 2 large prefill + 32 decode
27
28
- # Context extension
29
- - "q1ks2k" # 1k query, 2k sequence (chunked prefill)
28
+ # Speculative decode (q <= 8)
29
+ - "16q2s1k" # 16 requests, 2 spec tokens, 1k KV cache
30
+ - "16q4s1k" # 16 requests, 4 spec tokens, 1k KV cache
31
+ - "16q8s1k" # 16 requests, 8 spec tokens, 1k KV cache
32
+ - "32q4s2k" # 32 requests, 4 spec tokens, 2k KV cache
33
+ - "8q8s4k" # 8 requests, 8 spec tokens, 4k KV cache
34
+
35
+ # Context extension (chunked prefill)
36
+ - "q1ks2k" # 1k query, 2k sequence
37
- "2q1ks4k" # 2 requests: 1k query, 4k sequence
38
39
+# Available backends: flash, triton, flashinfer
40
backends:
41
- flash
42
- triton
benchmarks/attention_benchmarks/runner.py
+151
-90
@@ -8,7 +8,9 @@ This module provides helpers for running standard attention backends
8
(FlashAttention, Triton, FlashInfer) with real vLLM integration.
9
"""
10
11
+import logging
12
import types
13
+from contextlib import contextmanager
14
15
import numpy as np
16
import torch
@@ -24,8 +26,13 @@ from vllm.config import (
26
ParallelConfig,
27
SchedulerConfig,
28
VllmConfig,
29
+ set_current_vllm_config,
30
+)
31
+from vllm.v1.attention.backends.utils import (
32
+ CommonAttentionMetadata,
33
+ get_kv_cache_layout,
34
+ set_kv_cache_layout,
35
)
28
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
36
from vllm.v1.kv_cache_interface import FullAttentionSpec
37
38
# ============================================================================
@@ -37,22 +44,14 @@ _BACKEND_CONFIG = {
44
"flash": {
45
"module": "vllm.v1.attention.backends.flash_attn",
46
"backend_class": "FlashAttentionBackend",
40
- "dtype": torch.float16,
41
- "cache_layout": "standard",
42
- # ^ [2, num_blocks, block_size, num_kv_heads, head_dim]
47
},
48
"triton": {
49
"module": "vllm.v1.attention.backends.triton_attn",
50
"backend_class": "TritonAttentionBackend",
47
- "dtype": torch.float32,
48
- "cache_layout": "standard",
51
},
52
"flashinfer": {
53
"module": "vllm.v1.attention.backends.flashinfer",
54
"backend_class": "FlashInferBackend",
53
- "dtype": torch.float16,
54
- "cache_layout": "flashinfer",
55
- # ^ [num_blocks, 2, block_size, num_kv_heads, head_dim]
55
},
56
}
57
@@ -66,6 +65,18 @@ def _get_backend_config(backend: str) -> dict:
65
return _BACKEND_CONFIG[backend]
66
67
68
+@contextmanager
69
+def log_warnings_and_errors_only():
70
+ """Temporarily set vLLM logger to WARNING level."""
71
+ logger = logging.getLogger("vllm")
72
+ old_level = logger.level
73
+ logger.setLevel(logging.WARNING)
74
+ try:
75
+ yield
76
+ finally:
77
+ logger.setLevel(old_level)
78
+
79
+
80
# ============================================================================
81
# Metadata Building Helpers
82
# ============================================================================
@@ -88,11 +99,7 @@ def _build_common_attn_metadata(
99
query_start_loc_cpu = query_start_loc.cpu()
100
101
seq_lens = torch.tensor(kv_lens, dtype=torch.int32, device=device)
91
- seq_lens_cpu = seq_lens.cpu()
92
- max_seq_len = int(seq_lens_cpu.max())
93
-
94
- context_lens = [kv - q for kv, q in zip(kv_lens, q_lens)]
95
- num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
102
+ max_seq_len = int(seq_lens.max().item())
103
104
max_blocks = (max(kv_lens) + block_size - 1) // block_size
105
num_blocks = batch_size * max_blocks
@@ -107,8 +114,6 @@ def _build_common_attn_metadata(
114
query_start_loc=query_start_loc,
115
query_start_loc_cpu=query_start_loc_cpu,
116
seq_lens=seq_lens,
110
- seq_lens_cpu=seq_lens_cpu,
111
- num_computed_tokens_cpu=num_computed_tokens_cpu,
117
num_reqs=batch_size,
118
num_actual_tokens=total_tokens,
119
max_query_len=max_query_len,
@@ -121,7 +126,6 @@ def _build_common_attn_metadata(
126
127
def _create_vllm_config(
128
config: BenchmarkConfig,
124
- dtype: torch.dtype,
129
max_num_blocks: int,
130
) -> VllmConfig:
131
"""Create a VllmConfig for benchmarking with mock model methods."""
@@ -129,7 +133,7 @@ def _create_vllm_config(
133
model="meta-llama/Meta-Llama-3-8B",
134
tokenizer="meta-llama/Meta-Llama-3-8B",
135
trust_remote_code=False,
132
- dtype=dtype,
136
+ dtype="auto", # Use model's native dtype
137
seed=0,
138
max_model_len=1024,
139
)
@@ -198,6 +202,7 @@ def _create_backend_impl(
202
backend_cfg: dict,
203
config: BenchmarkConfig,
204
device: torch.device,
205
+ dtype: torch.dtype,
206
):
207
"""Create backend implementation instance."""
208
import importlib
@@ -206,7 +211,6 @@ def _create_backend_impl(
211
backend_class = getattr(backend_module, backend_cfg["backend_class"])
212
213
scale = get_attention_scale(config.head_dim)
209
- dtype = backend_cfg["dtype"]
214
215
impl = backend_class.get_impl_cls()(
216
num_heads=config.num_q_heads,
@@ -227,7 +231,7 @@ def _create_backend_impl(
231
232
layer = MockLayer(device, kv_cache_spec=kv_cache_spec)
233
230
- return backend_class, impl, layer, dtype
234
+ return backend_class, impl, layer
235
236
237
def _create_metadata_builder(
@@ -235,11 +239,44 @@ def _create_metadata_builder(
239
kv_cache_spec: FullAttentionSpec,
240
vllm_config: VllmConfig,
241
device: torch.device,
242
+ backend_name: str = "",
243
):
244
"""Create metadata builder instance."""
240
- return backend_class.get_builder_cls()(
245
+ layer_names = ["layer_0"]
246
+ builder_cls = backend_class.get_builder_cls()
247
+
248
+ # Flashinfer needs get_per_layer_parameters mocked since we don't have
249
+ # real model layers registered
250
+ if backend_name == "flashinfer":
251
+ import unittest.mock
252
+
253
+ from vllm.v1.attention.backends.utils import PerLayerParameters
254
+
255
+ def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
256
+ head_size = vllm_config.model_config.get_head_size()
257
+ return {
258
+ layer_name: PerLayerParameters(
259
+ window_left=-1, # No sliding window
260
+ logits_soft_cap=0.0, # No soft cap
261
+ sm_scale=1.0 / (head_size**0.5), # Standard scale
262
+ )
263
+ for layer_name in layer_names
264
+ }
265
+
266
+ with unittest.mock.patch(
267
+ "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
268
+ mock_get_per_layer_parameters,
269
+ ):
270
+ return builder_cls(
271
+ kv_cache_spec=kv_cache_spec,
272
+ layer_names=layer_names,
273
+ vllm_config=vllm_config,
274
+ device=device,
275
+ )
276
+
277
+ return builder_cls(
278
kv_cache_spec=kv_cache_spec,
242
- layer_names=["layer_0"],
279
+ layer_names=layer_names,
280
vllm_config=vllm_config,
281
device=device,
282
)
@@ -281,39 +318,44 @@ def _create_input_tensors(
318
def _create_kv_cache(
319
config: BenchmarkConfig,
320
max_num_blocks: int,
284
- cache_layout: str,
321
+ backend_class,
322
device: torch.device,
323
dtype: torch.dtype,
324
) -> list:
288
- """Create KV cache tensors for all layers."""
289
- if cache_layout == "flashinfer":
290
- # FlashInfer layout: [num_blocks, 2, block_size, num_kv_heads, head_dim]
291
- cache_list = [
292
- torch.zeros(
293
- max_num_blocks,
294
- 2,
295
- config.block_size,
296
- config.num_kv_heads,
297
- config.head_dim,
298
- device=device,
299
- dtype=dtype,
300
- )
301
- for _ in range(config.num_layers)
302
- ]
303
- else:
304
- # Standard layout: [2, num_blocks, block_size, num_kv_heads, head_dim]
305
- cache_list = [
306
- torch.zeros(
307
- 2,
308
- max_num_blocks,
309
- config.block_size,
310
- config.num_kv_heads,
311
- config.head_dim,
312
- device=device,
313
- dtype=dtype,
314
- )
315
- for _ in range(config.num_layers)
316
- ]
325
+ """Create KV cache tensors for all layers using the backend's methods.
326
+
327
+ Uses the backend's get_kv_cache_shape() and get_kv_cache_stride_order()
328
+ to create the cache with the correct shape and memory layout.
329
+ """
330
+ # Get the logical shape from the backend
331
+ cache_shape = backend_class.get_kv_cache_shape(
332
+ num_blocks=max_num_blocks,
333
+ block_size=config.block_size,
334
+ num_kv_heads=config.num_kv_heads,
335
+ head_size=config.head_dim,
336
+ )
337
+
338
+ # Get the stride order for custom memory layout
339
+ try:
340
+ stride_order = backend_class.get_kv_cache_stride_order()
341
+ assert len(stride_order) == len(cache_shape)
342
+ except (AttributeError, NotImplementedError):
343
+ stride_order = tuple(range(len(cache_shape)))
344
+
345
+ # Permute shape to physical layout order
346
+ physical_shape = tuple(cache_shape[i] for i in stride_order)
347
+
348
+ # Compute inverse permutation to get back to logical view
349
+ inv_order = [stride_order.index(i) for i in range(len(stride_order))]
350
+
351
+ cache_list = []
352
+ for _ in range(config.num_layers):
353
+ # Allocate in physical layout order (contiguous in memory)
354
+ cache = torch.zeros(*physical_shape, device=device, dtype=dtype)
355
+ # Permute to logical view
356
+ cache = cache.permute(*inv_order)
357
+ cache_list.append(cache)
358
+
359
return cache_list
360
361
@@ -418,53 +460,72 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
460
kv_lens = [r.kv_len for r in requests]
461
total_q = sum(q_lens)
462
max_kv = max(kv_lens)
463
+ batch_size = len(q_lens)
464
422
- max_num_blocks = (max_kv + config.block_size - 1) // config.block_size
423
-
424
- backend_class, impl, layer, dtype = _create_backend_impl(
425
- backend_cfg, config, device
426
- )
465
+ # Calculate total blocks needed: batch_size * max_blocks_per_request
466
+ max_blocks_per_request = (max_kv + config.block_size - 1) // config.block_size
467
+ max_num_blocks = batch_size * max_blocks_per_request
468
+
469
+ # Suppress vLLM logs during setup to reduce spam
470
+ with log_warnings_and_errors_only():
471
+ # Create vllm_config first - uses model's native dtype via "auto"
472
+ vllm_config = _create_vllm_config(config, max_num_blocks)
473
+ dtype = vllm_config.model_config.dtype
474
+
475
+ # Wrap everything in set_current_vllm_config context
476
+ # This is required for backends like flashinfer that need global config
477
+ with set_current_vllm_config(vllm_config):
478
+ backend_class, impl, layer = _create_backend_impl(
479
+ backend_cfg, config, device, dtype
480
+ )
481
428
- common_metadata = _build_common_attn_metadata(
429
- q_lens, kv_lens, config.block_size, device
430
- )
482
+ # Set KV cache layout if the backend requires a specific one
483
+ # (e.g., FlashInfer requires HND on SM100/Blackwell for TRTLLM attention)
484
+ required_layout = backend_class.get_required_kv_cache_layout()
485
+ if required_layout is not None:
486
+ set_kv_cache_layout(required_layout)
487
+ get_kv_cache_layout.cache_clear()
488
432
- kv_cache_spec = FullAttentionSpec(
433
- block_size=config.block_size,
434
- num_kv_heads=config.num_kv_heads,
435
- head_size=config.head_dim,
436
- dtype=dtype,
437
- )
489
+ common_metadata = _build_common_attn_metadata(
490
+ q_lens, kv_lens, config.block_size, device
491
+ )
492
439
- vllm_config = _create_vllm_config(config, dtype, max_num_blocks)
493
+ kv_cache_spec = FullAttentionSpec(
494
+ block_size=config.block_size,
495
+ num_kv_heads=config.num_kv_heads,
496
+ head_size=config.head_dim,
497
+ dtype=dtype,
498
+ )
499
441
- builder = _create_metadata_builder(
442
- backend_class, kv_cache_spec, vllm_config, device
443
- )
500
+ builder = _create_metadata_builder(
501
+ backend_class, kv_cache_spec, vllm_config, device, config.backend
502
+ )
503
445
- attn_metadata = builder.build(
446
- common_prefix_len=0,
447
- common_attn_metadata=common_metadata,
448
- )
504
+ attn_metadata = builder.build(
505
+ common_prefix_len=0,
506
+ common_attn_metadata=common_metadata,
507
+ )
508
450
- q_list, k_list, v_list = _create_input_tensors(config, total_q, device, dtype)
509
+ q_list, k_list, v_list = _create_input_tensors(
510
+ config, total_q, device, dtype
511
+ )
512
452
- cache_list = _create_kv_cache(
453
- config, max_num_blocks, backend_cfg["cache_layout"], device, dtype
454
- )
513
+ cache_list = _create_kv_cache(
514
+ config, max_num_blocks, backend_class, device, dtype
515
+ )
516
456
- times, mem_stats = _run_single_benchmark(
457
- config,
458
- impl,
459
- layer,
460
- q_list,
461
- k_list,
462
- v_list,
463
- cache_list,
464
- attn_metadata,
465
- device,
466
- dtype,
467
- )
517
+ times, mem_stats = _run_single_benchmark(
518
+ config,
519
+ impl,
520
+ layer,
521
+ q_list,
522
+ k_list,
523
+ v_list,
524
+ cache_list,
525
+ attn_metadata,
526
+ device,
527
+ dtype,
528
+ )
529
530
mean_time = np.mean(times)
531
throughput = total_q / mean_time if mean_time > 0 else 0