@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e82fa448c4
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 9, 12:30 UTC

Add attention benchmarking tools (#26835)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude <noreply@anthropic.com>

Matthew Bonanni committed Jan 27, 2026 at 19:09 UTC e82fa448c40bb5a24add0949271813a61f2aa47f
12 files changed +3558
benchmarks/attention_benchmarks/README.md new
+266
@@ -0,0 +1,266 @@
1 +# vLLM Attention Benchmarking Suite
2 +
3 +Fast, flexible benchmarking for vLLM attention and MLA backends with an extended batch specification grammar.
4 +
5 +## Quick Start
6 +
7 +```bash
8 +cd benchmarks/attention_benchmarks
9 +
10 +# Run a pre-configured benchmark
11 +python benchmark.py --config configs/mla_decode.yaml
12 +python benchmark.py --config configs/mla_mixed_batch.yaml
13 +python benchmark.py --config configs/speculative_decode.yaml
14 +python benchmark.py --config configs/standard_attention.yaml
15 +python benchmark.py --config configs/reorder_threshold.yaml
16 +
17 +# Or run custom benchmarks
18 +python benchmark.py \
19 + --backends flash flashinfer \
20 + --batch-specs "q2k" "8q1s1k" "2q2k_32q1s1k" \
21 + --output-csv results.csv
22 +```
23 +
24 +## Simplified Batch Specification Grammar
25 +
26 +Express workloads concisely using query length and sequence length:
27 +
28 +```python
29 +"q2k" # 2048-token prefill (q_len=2048, seq_len=2048)
30 +"q1s1k" # Decode: 1 token with 1K sequence
31 +"8q1s1k" # 8 decode requests
32 +"q4s1k" # 4-token extend (e.g., spec decode)
33 +"2q2k_32q1s1k" # Mixed: 2 prefills + 32 decodes
34 +"16q4s1k" # 16 spec decode (4 tokens each)
35 +```
36 +
37 +### Grammar Rule
38 +
39 +```text
40 +Format: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
41 +
42 +- count: Number of identical requests (optional, default=1)
43 +- q_len: Query length (number of new tokens)
44 +- seq_len: Total sequence length (optional, defaults to q_len for prefill)
45 +- 'k': Multiplies value by 1024
46 +
47 +Mixed batches: Use _ to combine (e.g., "2q2k_32q1s1k")
48 +```
49 +
50 +**Note**: Decode, prefill, and spec decode are just different query lengths - no special syntax needed!
51 +
52 +## Pre-configured Benchmarks
53 +
54 +The suite includes several pre-configured YAML benchmark configurations:
55 +
56 +### MLA Decode Benchmark
57 +
58 +Tests pure decode performance across MLA backends with varying batch sizes and sequence lengths.
59 +
60 +```bash
61 +python benchmark.py --config configs/mla_decode.yaml
62 +```
63 +
64 +### MLA Mixed Batch Benchmark
65 +
66 +Tests chunked prefill performance with mixed prefill + decode batches.
67 +
68 +```bash
69 +python benchmark.py --config configs/mla_mixed_batch.yaml
70 +```
71 +
72 +### Speculative Decoding Benchmark
73 +
74 +Tests speculative decode scenarios (K-token verification) and reorder_batch_threshold optimization.
75 +
76 +```bash
77 +python benchmark.py --config configs/speculative_decode.yaml
78 +```
79 +
80 +### Standard Attention Benchmark
81 +
82 +Tests standard attention backends (Flash/Triton/FlashInfer) with pure prefill, decode, and mixed batches.
83 +
84 +```bash
85 +python benchmark.py --config configs/standard_attention.yaml
86 +```
87 +
88 +### Reorder Threshold Study
89 +
90 +**Question:** At what query length does the prefill pipeline become faster than the decode pipeline?
91 +
92 +Tests query lengths from 1-1024 across 9 batch sizes to find the crossover point. Uses `decode_vs_prefill` mode to compare both pipelines for each query length.
93 +
94 +```bash
95 +python benchmark.py --config configs/reorder_threshold.yaml
96 +```
97 +
98 +---
99 +
100 +## Universal Benchmark
101 +
102 +The `benchmark.py` script handles **all** backends - both standard attention and MLA.
103 +
104 +### Standard Attention (Flash/Triton/FlashInfer)
105 +
106 +```bash
107 +python benchmark.py \
108 + --backends flash triton flashinfer \
109 + --batch-specs "q2k" "8q1s1k" "2q2k_32q1s1k" \
110 + --num-layers 10 \
111 + --repeats 5 \
112 + --output-csv results.csv
113 +```
114 +
115 +### MLA Backends
116 +
117 +```bash
118 +# Compare all MLA backends
119 +python benchmark.py \
120 + --backends cutlass_mla flashinfer_mla flashattn_mla flashmla \
121 + --batch-specs "64q1s1k" "64q1s4k" \
122 + --output-csv mla_results.csv
123 +```
124 +
125 +### Parameter Sweeps
126 +
127 +Use `--sweep-param` and `--sweep-values` to run parameter sweeps from the CLI:
128 +
129 +#### CUTLASS MLA num-splits Optimization
130 +
131 +**Question:** What is the optimal `num_kv_splits` for CUTLASS MLA?
132 +
133 +```bash
134 +python benchmark.py \
135 + --backend cutlass_mla \
136 + --batch-specs "64q1s1k" "64q1s4k" "64q1s16k" \
137 + --sweep-param num_kv_splits \
138 + --sweep-values 1 2 4 8 16 \
139 + --output-json optimal_splits.json
140 +```
141 +
142 +#### Reorder Batch Threshold Optimization
143 +
144 +**Question:** What's the optimal `reorder_batch_threshold` for speculative decoding?
145 +
146 +```bash
147 +python benchmark.py \
148 + --backend flashmla \
149 + --batch-specs "q4s1k" "q8s2k" \
150 + --sweep-param reorder_batch_threshold \
151 + --sweep-values 1 4 16 64 256 512 \
152 + --output-csv threshold_sweep.csv
153 +```
154 +
155 +### All Command-Line Options
156 +
157 +```text
158 +--config CONFIG # Path to YAML config file (overrides other args)
159 +--backends BACKEND [BACKEND ...] # flash, triton, flashinfer, cutlass_mla,
160 + # flashinfer_mla, flashattn_mla, flashmla
161 +--backend BACKEND # Single backend (alternative to --backends)
162 +--batch-specs SPEC [SPEC ...] # Batch specifications using extended grammar
163 +
164 +# Model configuration
165 +--num-layers N # Number of layers
166 +--head-dim N # Head dimension
167 +--num-q-heads N # Query heads
168 +--num-kv-heads N # KV heads
169 +--block-size N # Block size
170 +
171 +# Benchmark settings
172 +--device DEVICE # Device (default: cuda:0)
173 +--repeats N # Repetitions
174 +--warmup-iters N # Warmup iterations
175 +--profile-memory # Profile memory usage
176 +
177 +# Parameter sweeps
178 +--sweep-param PARAM # Parameter name to sweep (e.g., num_kv_splits,
179 + # reorder_batch_threshold)
180 +--sweep-values N [N ...] # Values to sweep for the parameter
181 +
182 +# Output
183 +--output-csv FILE # Save to CSV
184 +--output-json FILE # Save to JSON
185 +```
186 +
187 +## Hardware Requirements
188 +
189 +| Backend | Hardware |
190 +|---------|----------|
191 +| Flash/Triton/FlashInfer | Any CUDA GPU |
192 +| CUTLASS MLA | Blackwell (SM100+) |
193 +| FlashAttn MLA | Hopper (SM90+) |
194 +| FlashMLA | Hopper (SM90+) |
195 +| FlashInfer-MLA | Any CUDA GPU |
196 +
197 +## Using MLA Runner Directly
198 +
199 +All MLA backends are available through `mla_runner.run_mla_benchmark()`:
200 +
201 +```python
202 +from mla_runner import run_mla_benchmark
203 +from common import BenchmarkConfig
204 +
205 +config = BenchmarkConfig(
206 + backend="cutlass_mla",
207 + batch_spec="64q1s4k",
208 + num_layers=10,
209 + head_dim=576,
210 + num_q_heads=128,
211 + num_kv_heads=1,
212 + block_size=128,
213 + device="cuda:0",
214 + repeats=5,
215 + warmup_iters=3,
216 +)
217 +
218 +# CUTLASS MLA with specific num_kv_splits
219 +result = run_mla_benchmark("cutlass_mla", config, num_kv_splits=4)
220 +print(f"Time: {result.mean_time:.6f}s")
221 +
222 +# FlashInfer-MLA
223 +result = run_mla_benchmark("flashinfer_mla", config)
224 +
225 +# FlashAttn MLA (Hopper SM90+)
226 +result = run_mla_benchmark("flashattn_mla", config, reorder_batch_threshold=64)
227 +
228 +# FlashMLA (Hopper SM90+)
229 +result = run_mla_benchmark("flashmla", config, reorder_batch_threshold=64)
230 +```
231 +
232 +## Python API
233 +
234 +```python
235 +from batch_spec import parse_batch_spec, format_batch_spec, get_batch_stats
236 +from common import BenchmarkConfig, BenchmarkResult, ResultsFormatter
237 +
238 +# Parse batch specs
239 +requests = parse_batch_spec("2q2k_q4s1k_32q1s1k")
240 +print(format_batch_spec(requests))
241 +# "2 prefill (2x2k), 1 extend (1xq4kv1k), 32 decode (32x1k)"
242 +
243 +# Get batch statistics
244 +stats = get_batch_stats(requests)
245 +print(f"Total tokens: {stats['total_tokens']}")
246 +print(f"Num decode: {stats['num_decode']}, Num prefill: {stats['num_prefill']}")
247 +
248 +# Format results
249 +formatter = ResultsFormatter()
250 +formatter.save_csv(results, "output.csv")
251 +formatter.save_json(results, "output.json")
252 +```
253 +
254 +## Tips
255 +
256 +**1. Warmup matters** - Use `--warmup-iters 10` for stable results
257 +
258 +**2. Multiple repeats** - Use `--repeats 20` for low variance
259 +
260 +**3. Save results** - Always use `--output-csv` or `--output-json`
261 +
262 +**4. Test incrementally** - Start with `--num-layers 1 --repeats 1`
263 +
264 +**5. Extended grammar** - Leverage spec decode, chunked prefill patterns
265 +
266 +**6. Parameter sweeps** - Use `--sweep-param` and `--sweep-values` to find optimal values
benchmarks/attention_benchmarks/__init__.py new
+44
@@ -0,0 +1,44 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""vLLM Attention Benchmarking Suite."""
5 +
6 +from .batch_spec import (
7 + BatchRequest,
8 + format_batch_spec,
9 + get_batch_stats,
10 + parse_batch_spec,
11 + reorder_for_flashinfer,
12 + split_by_type,
13 +)
14 +from .common import (
15 + BenchmarkConfig,
16 + BenchmarkResult,
17 + MockLayer,
18 + MockModelConfig,
19 + ResultsFormatter,
20 + get_attention_scale,
21 + is_mla_backend,
22 + setup_mla_dims,
23 +)
24 +
25 +__all__ = [
26 + # Batch specification
27 + "BatchRequest",
28 + "parse_batch_spec",
29 + "format_batch_spec",
30 + "reorder_for_flashinfer",
31 + "split_by_type",
32 + "get_batch_stats",
33 + # Benchmarking infrastructure
34 + "BenchmarkConfig",
35 + "BenchmarkResult",
36 + "ResultsFormatter",
37 + # Mock objects
38 + "MockLayer",
39 + "MockModelConfig",
40 + # Utilities
41 + "setup_mla_dims",
42 + "get_attention_scale",
43 + "is_mla_backend",
44 +]
benchmarks/attention_benchmarks/batch_spec.py new
+231
@@ -0,0 +1,231 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""
5 +Simplified batch specification grammar for attention benchmarks.
6 +
7 +Grammar (underscore-separated segments):
8 + Format: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
9 +
10 + - count: Number of identical requests (optional, default=1)
11 + - q_len: Query length (number of new tokens)
12 + - seq_len: Total sequence length (optional, defaults to q_len for prefill)
13 + - 'k' suffix: Multiplies value by 1024
14 +
15 +Common patterns:
16 + - Prefill: q_len == seq_len (e.g., "q2k" → 2048 new tokens, 2048 seq)
17 + - Decode: q_len == 1 (e.g., "q1s1k" → 1 token, 1024 seq length)
18 + - Extend: q_len < seq_len (e.g., "q4s1k" → 4 tokens, 1024 seq length)
19 +
20 +Examples:
21 + q2k -> [(2048, 2048)] # Prefill: 2048 tokens
22 + q1s1k -> [(1, 1024)] # Decode: 1 token, 1K sequence
23 + 8q1s1k -> [(1, 1024)] * 8 # 8 decode requests
24 + q4s1k -> [(4, 1024)] # 4-token extend (spec decode)
25 + 2q1k_32q1s1k -> [(1024, 1024)] * 2 + [(1, 1024)] * 32 # Mixed batch
26 + 16q4s1k -> [(4, 1024)] * 16 # 16 spec decode requests
27 +"""
28 +
29 +from collections import Counter
30 +from dataclasses import dataclass
31 +
32 +import regex as re
33 +
34 +
35 +@dataclass
36 +class BatchRequest:
37 + """Represents a single request in a batch."""
38 +
39 + q_len: int # Query length (number of new tokens)
40 + kv_len: int # Total KV cache length
41 +
42 + @property
43 + def is_decode(self) -> bool:
44 + """True if this is a decode request (q_len == 1)."""
45 + return self.q_len == 1
46 +
47 + @property
48 + def is_prefill(self) -> bool:
49 + """True if this is a pure prefill (q_len == kv_len)."""
50 + return self.q_len == self.kv_len
51 +
52 + @property
53 + def is_extend(self) -> bool:
54 + """True if this is context extension (q_len > 1, kv_len > q_len)."""
55 + return self.q_len > 1 and self.kv_len > self.q_len
56 +
57 + @property
58 + def context_len(self) -> int:
59 + """Context length (KV cache - query)."""
60 + return self.kv_len - self.q_len
61 +
62 + def as_tuple(self) -> tuple[int, int]:
63 + """Return as (q_len, kv_len) tuple for compatibility."""
64 + return (self.q_len, self.kv_len)
65 +
66 +
67 +def _parse_size(size_str: str, k_suffix: str) -> int:
68 + """Parse size string with optional 'k' suffix."""
69 + size = int(size_str)
70 + return size * 1024 if k_suffix == "k" else size
71 +
72 +
73 +def parse_batch_spec(spec: str) -> list[BatchRequest]:
74 + """
75 + Parse batch specification string into list of BatchRequest objects.
76 +
77 + Grammar: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
78 +
79 + Args:
80 + spec: Batch specification string (see module docstring for grammar)
81 +
82 + Returns:
83 + List of BatchRequest objects
84 +
85 + Raises:
86 + ValueError: If spec format is invalid
87 + """
88 + requests = []
89 +
90 + for seg in spec.split("_"):
91 + # Unified pattern: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
92 + m = re.match(r"^(?:(\d+))?q(\d+)(k?)(?:s(\d+)(k?))?$", seg)
93 + if m:
94 + cnt = int(m.group(1)) if m.group(1) else 1
95 + q_len = _parse_size(m.group(2), m.group(3))
96 + kv_len = _parse_size(m.group(4), m.group(5)) if m.group(4) else q_len
97 + requests.extend([BatchRequest(q_len=q_len, kv_len=kv_len)] * cnt)
98 + continue
99 +
100 + raise ValueError(f"Invalid batch spec segment: '{seg}'")
101 +
102 + return requests
103 +
104 +
105 +def format_batch_spec(requests: list[BatchRequest]) -> str:
106 + """
107 + Format list of BatchRequest into human-readable string.
108 +
109 + Groups requests by type and provides counts and sizes.
110 +
111 + Args:
112 + requests: List of BatchRequest objects
113 +
114 + Returns:
115 + Formatted string describing the batch
116 + """
117 + kinds = {
118 + "prefill": [],
119 + "extend": [],
120 + "decode": [],
121 + }
122 +
123 + for req in requests:
124 + tup = (req.q_len, req.kv_len)
125 + if req.is_prefill:
126 + kinds["prefill"].append(tup)
127 + elif req.is_extend:
128 + kinds["extend"].append(tup)
129 + elif req.is_decode:
130 + kinds["decode"].append(tup)
131 +
132 + parts = []
133 + for kind in ["prefill", "extend", "decode"]:
134 + lst = kinds[kind]
135 + if not lst:
136 + continue
137 +
138 + cnt_total = len(lst)
139 + ctr = Counter(lst)
140 + inner = []
141 +
142 + for (q, kv), cnt in ctr.items():
143 + if kind == "prefill":
144 + size = f"{q // 1024}k" if q % 1024 == 0 else str(q)
145 + inner.append(f"{cnt}x{size}")
146 + elif kind == "decode":
147 + size = f"{kv // 1024}k" if kv % 1024 == 0 else str(kv)
148 + inner.append(f"{cnt}x{size}")
149 + else: # extend
150 + qstr = f"{q // 1024}k" if q % 1024 == 0 else str(q)
151 + kstr = f"{kv // 1024}k" if kv % 1024 == 0 else str(kv)
152 + inner.append(f"{cnt}xq{qstr}kv{kstr}")
153 +
154 + parts.append(f"{cnt_total} {kind} ({', '.join(inner)})")
155 +
156 + return ", ".join(parts)
157 +
158 +
159 +def reorder_for_flashinfer(requests: list[BatchRequest]) -> list[BatchRequest]:
160 + """
161 + Reorder requests for FlashInfer: decode first, then prefill.
162 +
163 + FlashInfer expects decode requests before prefill requests for
164 + optimal performance.
165 +
166 + Args:
167 + requests: Original list of BatchRequest
168 +
169 + Returns:
170 + Reordered list with decode requests first
171 + """
172 + decodes = [r for r in requests if r.is_decode]
173 + non_decodes = [r for r in requests if not r.is_decode]
174 + return decodes + non_decodes
175 +
176 +
177 +def split_by_type(
178 + requests: list[BatchRequest],
179 +) -> dict[str, list[BatchRequest]]:
180 + """
181 + Split requests by type for analysis.
182 +
183 + Args:
184 + requests: List of BatchRequest
185 +
186 + Returns:
187 + Dict with keys: 'decode', 'prefill', 'extend'
188 + """
189 + result = {
190 + "decode": [],
191 + "prefill": [],
192 + "extend": [],
193 + }
194 +
195 + for req in requests:
196 + if req.is_decode:
197 + result["decode"].append(req)
198 + elif req.is_prefill:
199 + result["prefill"].append(req)
200 + elif req.is_extend:
201 + result["extend"].append(req)
202 +
203 + return result
204 +
205 +
206 +def get_batch_stats(requests: list[BatchRequest]) -> dict:
207 + """
208 + Compute statistics about a batch.
209 +
210 + Args:
211 + requests: List of BatchRequest
212 +
213 + Returns:
214 + Dict with batch statistics
215 + """
216 + by_type = split_by_type(requests)
217 +
218 + return {
219 + "total_requests": len(requests),
220 + "num_decode": len(by_type["decode"]),
221 + "num_prefill": len(by_type["prefill"]),
222 + "num_extend": len(by_type["extend"]),
223 + "total_tokens": sum(r.q_len for r in requests),
224 + "total_kv_cache": sum(r.kv_len for r in requests),
225 + "max_q_len": max((r.q_len for r in requests), default=0),
226 + "max_kv_len": max((r.kv_len for r in requests), default=0),
227 + "avg_q_len": sum(r.q_len for r in requests) / len(requests) if requests else 0,
228 + "avg_kv_len": (
229 + sum(r.kv_len for r in requests) / len(requests) if requests else 0
230 + ),
231 + }
benchmarks/attention_benchmarks/benchmark.py new
+886
@@ -0,0 +1,886 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +
5 +"""
6 +Universal vLLM Attention Benchmark
7 +
8 +Benchmark any attention backend with the extended grammar.
9 +Supports standard attention (Flash/Triton/FlashInfer) and MLA backends.
10 +
11 +Examples:
12 + # Standard attention
13 + python benchmark.py --backends flash flashinfer --batch-specs "q2k" "8q1s1k"
14 +
15 + # MLA backends
16 + python benchmark.py --backends cutlass_mla flashinfer_mla --batch-specs "64q1s1k"
17 +
18 + # Parameter sweep (CLI)
19 + python benchmark.py --backend cutlass_mla \
20 + --batch-specs "64q1s1k" \
21 + --sweep-param num_kv_splits \
22 + --sweep-values 1 4 8 16
23 +
24 + # Parameter sweep (YAML config - recommended)
25 + python benchmark.py --config configs/cutlass_numsplits.yaml
26 +"""
27 +
28 +import argparse
29 +import sys
30 +from dataclasses import replace
31 +from pathlib import Path
32 +
33 +import yaml
34 +from rich.console import Console
35 +from tqdm import tqdm
36 +
37 +sys.path.insert(0, str(Path(__file__).parent.parent.parent))
38 +
39 +from batch_spec import parse_batch_spec
40 +from common import (
41 + BenchmarkConfig,
42 + BenchmarkResult,
43 + ModelParameterSweep,
44 + ParameterSweep,
45 + ResultsFormatter,
46 + is_mla_backend,
47 +)
48 +
49 +
50 +def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
51 + """Run standard attention benchmark (Flash/Triton/FlashInfer)."""
52 + from runner import run_attention_benchmark
53 +
54 + return run_attention_benchmark(config)
55 +
56 +
57 +def run_mla_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
58 + """Run MLA benchmark with appropriate backend."""
59 + from mla_runner import run_mla_benchmark as run_mla
60 +
61 + return run_mla(config.backend, config, **kwargs)
62 +
63 +
64 +def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
65 + """
66 + Run a single benchmark with proper backend selection.
67 +
68 + Args:
69 + config: BenchmarkConfig with backend, batch_spec, and model params
70 + **kwargs: Additional arguments passed to MLA benchmarks
71 +
72 + Returns:
73 + BenchmarkResult (may have error field set on failure)
74 + """
75 + try:
76 + if is_mla_backend(config.backend):
77 + return run_mla_benchmark(config, **kwargs)
78 + else:
79 + return run_standard_attention_benchmark(config)
80 + except Exception as e:
81 + return BenchmarkResult(
82 + config=config,
83 + mean_time=float("inf"),
84 + std_time=0,
85 + min_time=float("inf"),
86 + max_time=float("inf"),
87 + error=str(e),
88 + )
89 +
90 +
91 +def run_model_parameter_sweep(
92 + backends: list[str],
93 + batch_specs: list[str],
94 + base_config_args: dict,
95 + sweep: ModelParameterSweep,
96 + console: Console,
97 +) -> list[BenchmarkResult]:
98 + """
99 + Run model parameter sweep for given backends and batch specs.
100 +
101 + Args:
102 + backends: List of backend names
103 + batch_specs: List of batch specifications
104 + base_config_args: Base configuration arguments (num_layers, head_dim, etc.)
105 + sweep: ModelParameterSweep configuration
106 + console: Rich console for output
107 +
108 + Returns:
109 + List of BenchmarkResult objects
110 + """
111 + all_results = []
112 +
113 + console.print(
114 + f"[yellow]Model sweep mode: testing {sweep.param_name} = {sweep.values}[/]"
115 + )
116 +
117 + total = len(backends) * len(batch_specs) * len(sweep.values)
118 +
119 + with tqdm(total=total, desc="Benchmarking") as pbar:
120 + for backend in backends:
121 + for spec in batch_specs:
122 + for value in sweep.values:
123 + # Create config with modified model parameter
124 + config_args = base_config_args.copy()
125 + config_args[sweep.param_name] = value
126 +
127 + # Create config with original backend for running
128 + clean_config = BenchmarkConfig(
129 + backend=backend, batch_spec=spec, **config_args
130 + )
131 +
132 + # Run benchmark
133 + result = run_benchmark(clean_config)
134 +
135 + # Replace backend with labeled version for display
136 + backend_label = sweep.get_label(backend, value)
137 + labeled_config = replace(result.config, backend=backend_label)
138 + result = replace(result, config=labeled_config)
139 + all_results.append(result)
140 +
141 + if not result.success:
142 + console.print(
143 + f"[red]Error {backend} {spec} {sweep.param_name}="
144 + f"{value}: {result.error}[/]"
145 + )
146 +
147 + pbar.update(1)
148 +
149 + # Display sweep results - create separate table for each parameter value
150 + console.print("\n[bold green]Model Parameter Sweep Results:[/]")
151 + formatter = ResultsFormatter(console)
152 +
153 + # Group results by parameter value and extract backend mapping
154 + by_param_value = {}
155 + backend_mapping = {} # Maps labeled backend -> original backend
156 +
157 + for r in all_results:
158 + # Extract original backend and param value from labeled backend
159 + # The label format is: {backend}_{param_name}_{value}
160 + # We need to reverse engineer this
161 + labeled_backend = r.config.backend
162 +
163 + # Try each backend to find which one this result belongs to
164 + for backend in backends:
165 + for value in sweep.values:
166 + expected_label = sweep.get_label(backend, value)
167 + if labeled_backend == expected_label:
168 + backend_mapping[labeled_backend] = backend
169 + param_value = str(value)
170 +
171 + if param_value not in by_param_value:
172 + by_param_value[param_value] = []
173 + by_param_value[param_value].append(r)
174 + break
175 +
176 + # Create a table for each parameter value
177 + sorted_param_values = sorted(
178 + by_param_value.keys(), key=lambda x: int(x) if x.isdigit() else x
179 + )
180 +
181 + for param_value in sorted_param_values:
182 + console.print(f"\n[bold cyan]{sweep.param_name} = {param_value}[/]")
183 + param_results = by_param_value[param_value]
184 +
185 + # Create modified results with original backend names
186 + modified_results = []
187 + for r in param_results:
188 + # Get the original backend name from our mapping
189 + original_backend = backend_mapping[r.config.backend]
190 + modified_config = replace(r.config, backend=original_backend)
191 + modified_result = replace(r, config=modified_config)
192 + modified_results.append(modified_result)
193 +
194 + # Print table with original backend names
195 + formatter.print_table(modified_results, backends, compare_to_fastest=True)
196 +
197 + # Show optimal backend for each (param_value, batch_spec) combination
198 + console.print(
199 + f"\n[bold cyan]Optimal backend for each ({sweep.param_name}, batch_spec):[/]"
200 + )
201 +
202 + # Group by (param_value, batch_spec)
203 + by_param_and_spec = {}
204 + for r in all_results:
205 + if r.success:
206 + # Find which (backend, value) this result corresponds to
207 + labeled_backend = r.config.backend
208 + for backend in backends:
209 + for value in sweep.values:
210 + expected_label = sweep.get_label(backend, value)
211 + if labeled_backend == expected_label:
212 + param_value = str(value)
213 + spec = r.config.batch_spec
214 + key = (param_value, spec)
215 +
216 + if key not in by_param_and_spec:
217 + by_param_and_spec[key] = []
218 + by_param_and_spec[key].append(r)
219 + break
220 +
221 + # Sort by param value then spec
222 + sorted_keys = sorted(
223 + by_param_and_spec.keys(),
224 + key=lambda x: (int(x[0]) if x[0].isdigit() else x[0], x[1]),
225 + )
226 +
227 + current_param_value = None
228 + for param_value, spec in sorted_keys:
229 + # Print header when param value changes
230 + if param_value != current_param_value:
231 + console.print(f"\n [bold]{sweep.param_name}={param_value}:[/]")
232 + current_param_value = param_value
233 +
234 + results = by_param_and_spec[(param_value, spec)]
235 + best = min(results, key=lambda r: r.mean_time)
236 +
237 + # Extract original backend name using the mapping
238 + backend_name = backend_mapping[best.config.backend]
239 +
240 + # Show all backends' times for comparison
241 + times_str = " | ".join(
242 + [
243 + f"{backend_mapping[r.config.backend]}: {r.mean_time:.6f}s"
244 + for r in sorted(results, key=lambda r: r.mean_time)
245 + ]
246 + )
247 +
248 + console.print(
249 + f" {spec:12s} -> [bold green]{backend_name:15s}[/] ({times_str})"
250 + )
251 +
252 + return all_results
253 +
254 +
255 +def run_parameter_sweep(
256 + backends: list[str],
257 + batch_specs: list[str],
258 + base_config_args: dict,
259 + sweep: ParameterSweep,
260 + console: Console,
261 +) -> list[BenchmarkResult]:
262 + """
263 + Run parameter sweep for given backends and batch specs.
264 +
265 + Args:
266 + backends: List of backend names
267 + batch_specs: List of batch specifications
268 + base_config_args: Base configuration arguments (num_layers, head_dim, etc.)
269 + sweep: ParameterSweep configuration
270 + console: Rich console for output
271 +
272 + Returns:
273 + List of BenchmarkResult objects
274 + """
275 + all_results = []
276 +
277 + # Build list of values to sweep (including auto if requested)
278 + sweep_values = list(sweep.values)
279 + if sweep.include_auto:
280 + sweep_values.append("auto")
281 +
282 + console.print(f"[yellow]Sweep mode: testing {sweep.param_name} = {sweep_values}[/]")
283 +
284 + total = len(backends) * len(batch_specs) * len(sweep_values)
285 +
286 + with tqdm(total=total, desc="Benchmarking") as pbar:
287 + for backend in backends:
288 + for spec in batch_specs:
289 + for value in sweep_values:
290 + # Create config with original backend for running
291 + config = BenchmarkConfig(
292 + backend=backend, batch_spec=spec, **base_config_args
293 + )
294 +
295 + # Prepare kwargs for benchmark runner
296 + kwargs = {}
297 + if value != "auto":
298 + kwargs[sweep.param_name] = value
299 +
300 + # Run benchmark
301 + result = run_benchmark(config, **kwargs)
302 +
303 + # Replace backend with labeled version for display
304 + backend_label = sweep.get_label(backend, value)
305 + labeled_config = replace(result.config, backend=backend_label)
306 + result = replace(result, config=labeled_config)
307 + all_results.append(result)
308 +
309 + if not result.success:
310 + console.print(
311 + f"[red]Error {backend} {spec} {sweep.param_name}="
312 + f"{value}: {result.error}[/]"
313 + )
314 +
315 + pbar.update(1)
316 +
317 + # Display sweep results
318 + console.print("\n[bold green]Sweep Results:[/]")
319 + backend_labels = [sweep.get_label(b, v) for b in backends for v in sweep_values]
320 + formatter = ResultsFormatter(console)
321 + formatter.print_table(all_results, backend_labels)
322 +
323 + # Show optimal values
324 + console.print(f"\n[bold cyan]Optimal {sweep.param_name} per batch spec:[/]")
325 + by_spec = {}
326 + for r in all_results:
327 + if r.success:
328 + spec = r.config.batch_spec
329 + if spec not in by_spec:
330 + by_spec[spec] = []
331 + by_spec[spec].append(r)
332 +
333 + for spec in sorted(by_spec.keys()):
334 + results = by_spec[spec]
335 + best = min(results, key=lambda r: r.mean_time)
336 + console.print(
337 + f" {spec}: [bold green]{best.config.backend}[/] ({best.mean_time:.6f}s)"
338 + )
339 +
340 + return all_results
341 +
342 +
343 +def load_config_from_yaml(config_path: str) -> dict:
344 + """Load configuration from YAML file."""
345 + with open(config_path) as f:
346 + return yaml.safe_load(f)
347 +
348 +
349 +def generate_batch_specs_from_ranges(ranges: list[dict]) -> list[str]:
350 + """
351 + Generate batch specs from range specifications.
352 +
353 + Args:
354 + ranges: List of range specifications, each containing:
355 + - template: Batch spec template (e.g., "q{q_len}kv1k")
356 + - q_len: Dict with start, stop, step, end_inclusive (optional)
357 + - Other parameters can also be ranges
358 +
359 + Returns:
360 + List of generated batch spec strings
361 +
362 + Example:
363 + ranges = [
364 + {
365 + "template": "q{q_len}kv1k",
366 + "q_len": {
367 + "start": 1,
368 + "stop": 16,
369 + "step": 1,
370 + "end_inclusive": true # Optional, defaults to true
371 + }
372 + }
373 + ]
374 + Returns: ["q1kv1k", "q2kv1k", ..., "q16kv1k"]
375 + """
376 + all_specs = []
377 +
378 + for range_spec in ranges:
379 + template = range_spec.get("template")
380 + if not template:
381 + raise ValueError("Range specification must include 'template'")
382 +
383 + # Extract all range parameters from the spec
384 + range_params = {}
385 + for key, value in range_spec.items():
386 + if key == "template":
387 + continue
388 + if isinstance(value, dict) and "start" in value:
389 + # This is a range specification
390 + start = value["start"]
391 + stop = value["stop"]
392 + step = value.get("step", 1)
393 + # Check if end should be inclusive (default: True)
394 + end_inclusive = value.get("end_inclusive", True)
395 +
396 + # Adjust stop based on end_inclusive
397 + if end_inclusive:
398 + range_params[key] = list(range(start, stop + 1, step))
399 + else:
400 + range_params[key] = list(range(start, stop, step))
401 + else:
402 + # This is a fixed value
403 + range_params[key] = [value]
404 +
405 + # Generate all combinations (Cartesian product)
406 + if range_params:
407 + import itertools
408 +
409 + param_names = list(range_params.keys())
410 + param_values = [range_params[name] for name in param_names]
411 +
412 + for values in itertools.product(*param_values):
413 + params = dict(zip(param_names, values))
414 + spec = template.format(**params)
415 + all_specs.append(spec)
416 + else:
417 + # No parameters, just use template as-is
418 + all_specs.append(template)
419 +
420 + return all_specs
421 +
422 +
423 +def main():
424 + parser = argparse.ArgumentParser(
425 + description="Universal vLLM attention benchmark",
426 + formatter_class=argparse.RawDescriptionHelpFormatter,
427 + epilog=__doc__,
428 + )
429 +
430 + # Config file
431 + parser.add_argument(
432 + "--config",
433 + help="Path to YAML config file (overrides other args)",
434 + )
435 +
436 + # Backend selection
437 + parser.add_argument(
438 + "--backends",
439 + nargs="+",
440 + help="Backends to benchmark (flash, triton, flashinfer, cutlass_mla, "
441 + "flashinfer_mla, flashattn_mla, flashmla)",
442 + )
443 + parser.add_argument(
444 + "--backend",
445 + help="Single backend (alternative to --backends)",
446 + )
447 +
448 + # Batch specifications
449 + parser.add_argument(
450 + "--batch-specs",
451 + nargs="+",
452 + default=["q2k", "8q1s1k"],
453 + help="Batch specifications using extended grammar",
454 + )
455 +
456 + # Model config
457 + parser.add_argument("--num-layers", type=int, default=10, help="Number of layers")
458 + parser.add_argument("--head-dim", type=int, default=128, help="Head dimension")
459 + parser.add_argument("--num-q-heads", type=int, default=32, help="Query heads")
460 + parser.add_argument("--num-kv-heads", type=int, default=8, help="KV heads")
461 + parser.add_argument("--block-size", type=int, default=16, help="Block size")
462 +
463 + # Benchmark settings
464 + parser.add_argument("--device", default="cuda:0", help="Device")
465 + parser.add_argument("--repeats", type=int, default=1, help="Repetitions")
466 + parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations")
467 + parser.add_argument("--profile-memory", action="store_true", help="Profile memory")
468 +
469 + # Parameter sweep (use YAML config for advanced sweeps)
470 + parser.add_argument(
471 + "--sweep-param",
472 + help="Parameter name to sweep (e.g., num_kv_splits, reorder_batch_threshold)",
473 + )
474 + parser.add_argument(
475 + "--sweep-values",
476 + type=int,
477 + nargs="+",
478 + help="Values to sweep for the parameter",
479 + )
480 +
481 + # Output
482 + parser.add_argument("--output-csv", help="Save to CSV")
483 + parser.add_argument("--output-json", help="Save to JSON")
484 +
485 + args = parser.parse_args()
486 +
487 + console = Console()
488 + console.print("[bold cyan]vLLM Attention Benchmark[/]")
489 +
490 + # Load config from YAML if provided
491 + if args.config:
492 + console.print(f"[yellow]Loading config from: {args.config}[/]")
493 + yaml_config = load_config_from_yaml(args.config)
494 +
495 + # Show description if available
496 + if "description" in yaml_config:
497 + console.print(f"[dim]{yaml_config['description']}[/]")
498 +
499 + # Override args with YAML values
500 + # (YAML takes precedence unless CLI arg was explicitly set)
501 + # Backend(s)
502 + if "backend" in yaml_config:
503 + args.backend = yaml_config["backend"]
504 + args.backends = None
505 + elif "backends" in yaml_config:
506 + args.backends = yaml_config["backends"]
507 + args.backend = None
508 +
509 + # Check for special modes
510 + if "mode" in yaml_config:
511 + args.mode = yaml_config["mode"]
512 + else:
513 + args.mode = None
514 +
515 + # Batch specs and sizes
516 + # Support both explicit batch_specs and generated batch_spec_ranges
517 + if "batch_spec_ranges" in yaml_config:
518 + # Generate batch specs from ranges
519 + generated_specs = generate_batch_specs_from_ranges(
520 + yaml_config["batch_spec_ranges"]
521 + )
522 + # Combine with any explicit batch_specs
523 + if "batch_specs" in yaml_config:
524 + args.batch_specs = yaml_config["batch_specs"] + generated_specs
525 + else:
526 + args.batch_specs = generated_specs
527 + console.print(
528 + f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
529 + )
530 + elif "batch_specs" in yaml_config:
531 + args.batch_specs = yaml_config["batch_specs"]
532 +
533 + if "batch_sizes" in yaml_config:
534 + args.batch_sizes = yaml_config["batch_sizes"]
535 + else:
536 + args.batch_sizes = None
537 +
538 + # Model config
539 + if "model" in yaml_config:
540 + model = yaml_config["model"]
541 + args.num_layers = model.get("num_layers", args.num_layers)
542 + args.head_dim = model.get("head_dim", args.head_dim)
543 + args.num_q_heads = model.get("num_q_heads", args.num_q_heads)
544 + args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads)
545 + args.block_size = model.get("block_size", args.block_size)
546 +
547 + # Benchmark settings
548 + if "benchmark" in yaml_config:
549 + bench = yaml_config["benchmark"]
550 + args.device = bench.get("device", args.device)
551 + args.repeats = bench.get("repeats", args.repeats)
552 + args.warmup_iters = bench.get("warmup_iters", args.warmup_iters)
553 + args.profile_memory = bench.get("profile_memory", args.profile_memory)
554 +
555 + # Parameter sweep configuration
556 + if "parameter_sweep" in yaml_config:
557 + sweep_config = yaml_config["parameter_sweep"]
558 + args.parameter_sweep = ParameterSweep(
559 + param_name=sweep_config["param_name"],
560 + values=sweep_config["values"],
561 + include_auto=sweep_config.get("include_auto", False),
562 + label_format=sweep_config.get(
563 + "label_format", "{backend}_{param_name}_{value}"
564 + ),
565 + )
566 + else:
567 + args.parameter_sweep = None
568 +
569 + # Model parameter sweep configuration
570 + if "model_parameter_sweep" in yaml_config:
571 + sweep_config = yaml_config["model_parameter_sweep"]
572 + args.model_parameter_sweep = ModelParameterSweep(
573 + param_name=sweep_config["param_name"],
574 + values=sweep_config["values"],
575 + label_format=sweep_config.get(
576 + "label_format", "{backend}_{param_name}_{value}"
577 + ),
578 + )
579 + else:
580 + args.model_parameter_sweep = None
581 +
582 + # Output
583 + if "output" in yaml_config:
584 + output = yaml_config["output"]
585 + if "csv" in output and not args.output_csv:
586 + args.output_csv = output["csv"]
587 + if "json" in output and not args.output_json:
588 + args.output_json = output["json"]
589 +
590 + console.print()
591 +
592 + # Handle CLI-based parameter sweep (if not from YAML)
593 + if (
594 + (not hasattr(args, "parameter_sweep") or args.parameter_sweep is None)
595 + and args.sweep_param
596 + and args.sweep_values
597 + ):
598 + args.parameter_sweep = ParameterSweep(
599 + param_name=args.sweep_param,
600 + values=args.sweep_values,
601 + include_auto=False,
602 + label_format="{backend}_{param_name}_{value}",
603 + )
604 +
605 + # Determine backends
606 + backends = args.backends or ([args.backend] if args.backend else ["flash"])
607 + console.print(f"Backends: {', '.join(backends)}")
608 + console.print(f"Batch specs: {', '.join(args.batch_specs)}")
609 + console.print()
610 +
611 + # Run benchmarks
612 + all_results = []
613 +
614 + # Handle special mode: decode_vs_prefill comparison
615 + if hasattr(args, "mode") and args.mode == "decode_vs_prefill":
616 + console.print("[yellow]Mode: Decode vs Prefill pipeline comparison[/]")
617 + console.print(
618 + "[dim]For each query length, testing both decode and prefill pipelines[/]"
619 + )
620 + console.print("[dim]Using batched execution for optimal performance[/]")
621 +
622 + # Extract batch sizes from config
623 + batch_sizes = getattr(args, "batch_sizes", [1])
624 + backend = backends[0] # Use first backend (should only be one)
625 +
626 + # Calculate total benchmarks
627 + total = len(batch_sizes)
628 +
629 + with tqdm(total=total, desc="Benchmarking") as pbar:
630 + for batch_size in batch_sizes:
631 + # Prepare all configs for this batch size
632 + configs_with_thresholds = []
633 +
634 + for spec in args.batch_specs:
635 + # Parse the batch spec to get query length
636 + requests = parse_batch_spec(spec)
637 + if not requests:
638 + console.print(
639 + f"[red]Error: Could not parse batch spec '{spec}'[/]"
640 + )
641 + continue
642 +
643 + # Get query length from first request
644 + query_length = requests[0].q_len
645 +
646 + # Create batch spec for this batch size
647 + # For batch_size > 1, we need to prepend the count
648 + batch_spec = f"{batch_size}{spec}" if batch_size > 1 else spec
649 +
650 + # Create base config (without backend name)
651 + base_config = BenchmarkConfig(
652 + backend=backend, # Will be overridden later
653 + batch_spec=batch_spec,
654 + num_layers=args.num_layers,
655 + head_dim=args.head_dim,
656 + num_q_heads=args.num_q_heads,
657 + num_kv_heads=args.num_kv_heads,
658 + block_size=args.block_size,
659 + device=args.device,
660 + repeats=args.repeats,
661 + warmup_iters=args.warmup_iters,
662 + profile_memory=args.profile_memory,
663 + )
664 +
665 + # Add decode pipeline config
666 + decode_threshold = query_length
667 + config_decode = replace(
668 + base_config,
669 + backend=f"{backend}_decode_qlen{query_length}_bs{batch_size}",
670 + )
671 + configs_with_thresholds.append((config_decode, decode_threshold))
672 +
673 + # Add prefill pipeline config if query_length > 1
674 + if query_length > 1:
675 + prefill_threshold = query_length - 1
676 + config_prefill = replace(
677 + base_config,
678 + backend=f"{backend}_prefill_qlen{query_length}"
679 + f"_bs{batch_size}",
680 + )
681 + configs_with_thresholds.append(
682 + (config_prefill, prefill_threshold)
683 + )
684 +
685 + # Run all benchmarks for this batch size in one go (batched mode)
686 + try:
687 + from mla_runner import run_mla_benchmark as run_mla
688 +
689 + # Use batched API: pass list of (config, threshold) tuples
690 + timing_results = run_mla(backend, configs_with_thresholds)
691 +
692 + # Create BenchmarkResult objects from timing results
693 + for (config, _), timing in zip(
694 + configs_with_thresholds, timing_results
695 + ):
696 + result = BenchmarkResult(
697 + config=config,
698 + mean_time=timing["mean"],
699 + std_time=timing["std"],
700 + min_time=timing["min"],
701 + max_time=timing["max"],
702 + throughput_tokens_per_sec=timing.get("throughput", None),
703 + )
704 + all_results.append(result)
705 +
706 + except Exception as e:
707 + import traceback
708 +
709 + console.print(
710 + f"[red]Error running batched benchmarks for "
711 + f"batch_size={batch_size}: {e}[/]"
712 + )
713 + console.print("[red]Traceback:[/]")
714 + traceback.print_exc()
715 + # Add error results for all configs
716 + for config, _ in configs_with_thresholds:
717 + result = BenchmarkResult(
718 + config=config,
719 + mean_time=float("inf"),
720 + std_time=0,
721 + min_time=float("inf"),
722 + max_time=float("inf"),
723 + error=str(e),
724 + )
725 + all_results.append(result)
726 +
727 + pbar.update(1)
728 +
729 + # Display decode vs prefill results
730 + console.print("\n[bold green]Decode vs Prefill Results:[/]")
731 +
732 + # Group by batch size
733 + by_batch_size = {}
734 + for r in all_results:
735 + if r.success:
736 + # Extract batch size from backend name
737 + parts = r.config.backend.split("_")
738 + bs_part = [p for p in parts if p.startswith("bs")]
739 + if bs_part:
740 + bs = int(bs_part[0][2:])
741 + if bs not in by_batch_size:
742 + by_batch_size[bs] = []
743 + by_batch_size[bs].append(r)
744 +
745 + # For each batch size, analyze crossover point
746 + for bs in sorted(by_batch_size.keys()):
747 + console.print(f"\n[bold cyan]Batch size: {bs}[/]")
748 + results = by_batch_size[bs]
749 +
750 + # Group by query length
751 + by_qlen = {}
752 + for r in results:
753 + parts = r.config.backend.split("_")
754 + qlen_part = [p for p in parts if p.startswith("qlen")]
755 + if qlen_part:
756 + qlen = int(qlen_part[0][4:])
757 + if qlen not in by_qlen:
758 + by_qlen[qlen] = {}
759 +
760 + pipeline = "decode" if "decode" in r.config.backend else "prefill"
761 + by_qlen[qlen][pipeline] = r
762 +
763 + # Find crossover point
764 + last_decode_faster = None
765 + for qlen in sorted(by_qlen.keys()):
766 + pipelines = by_qlen[qlen]
767 + if "decode" in pipelines and "prefill" in pipelines:
768 + decode_time = pipelines["decode"].mean_time
769 + prefill_time = pipelines["prefill"].mean_time
770 + faster = "decode" if decode_time < prefill_time else "prefill"
771 +
772 + speedup = (
773 + prefill_time / decode_time
774 + if decode_time < prefill_time
775 + else decode_time / prefill_time
776 + )
777 +
778 + console.print(
779 + f" qlen={qlen:3d}: decode={decode_time:.6f}s, "
780 + f"prefill={prefill_time:.6f}s -> "
781 + f"[bold]{faster}[/] ({speedup:.2f}x)"
782 + )
783 +
784 + if faster == "decode":
785 + last_decode_faster = qlen
786 +
787 + if last_decode_faster is not None:
788 + optimal_threshold = last_decode_faster
789 + console.print(
790 + f"\n [bold green]Optimal threshold for batch_size={bs}: "
791 + f"{optimal_threshold}[/]"
792 + )
793 + console.print(
794 + f" [dim](Use decode pipeline for query_length <= "
795 + f"{optimal_threshold})[/]"
796 + )
797 + else:
798 + console.print(
799 + f"\n [yellow]Prefill always faster for batch_size={bs}[/]"
800 + )
801 +
802 + # Handle model parameter sweep mode
803 + elif hasattr(args, "model_parameter_sweep") and args.model_parameter_sweep:
804 + # Model parameter sweep
805 + base_config_args = {
806 + "num_layers": args.num_layers,
807 + "head_dim": args.head_dim,
808 + "num_q_heads": args.num_q_heads,
809 + "num_kv_heads": args.num_kv_heads,
810 + "block_size": args.block_size,
811 + "device": args.device,
812 + "repeats": args.repeats,
813 + "warmup_iters": args.warmup_iters,
814 + "profile_memory": args.profile_memory,
815 + }
816 + all_results = run_model_parameter_sweep(
817 + backends,
818 + args.batch_specs,
819 + base_config_args,
820 + args.model_parameter_sweep,
821 + console,
822 + )
823 +
824 + # Handle parameter sweep mode (unified)
825 + elif hasattr(args, "parameter_sweep") and args.parameter_sweep:
826 + # Unified parameter sweep
827 + base_config_args = {
828 + "num_layers": args.num_layers,
829 + "head_dim": args.head_dim,
830 + "num_q_heads": args.num_q_heads,
831 + "num_kv_heads": args.num_kv_heads,
832 + "block_size": args.block_size,
833 + "device": args.device,
834 + "repeats": args.repeats,
835 + "warmup_iters": args.warmup_iters,
836 + "profile_memory": args.profile_memory,
837 + }
838 + all_results = run_parameter_sweep(
839 + backends, args.batch_specs, base_config_args, args.parameter_sweep, console
840 + )
841 +
842 + else:
843 + # Normal mode: compare backends
844 + total = len(backends) * len(args.batch_specs)
845 +
846 + with tqdm(total=total, desc="Benchmarking") as pbar:
847 + for spec in args.batch_specs:
848 + for backend in backends:
849 + config = BenchmarkConfig(
850 + backend=backend,
851 + batch_spec=spec,
852 + num_layers=args.num_layers,
853 + head_dim=args.head_dim,
854 + num_q_heads=args.num_q_heads,
855 + num_kv_heads=args.num_kv_heads,
856 + block_size=args.block_size,
857 + device=args.device,
858 + repeats=args.repeats,
859 + warmup_iters=args.warmup_iters,
860 + profile_memory=args.profile_memory,
861 + )
862 +
863 + result = run_benchmark(config)
864 + all_results.append(result)
865 +
866 + if not result.success:
867 + console.print(f"[red]Error {backend} {spec}: {result.error}[/]")
868 +
869 + pbar.update(1)
870 +
871 + # Display results
872 + console.print("\n[bold green]Results:[/]")
873 + formatter = ResultsFormatter(console)
874 + formatter.print_table(all_results, backends)
875 +
876 + # Save results
877 + if all_results:
878 + formatter = ResultsFormatter(console)
879 + if args.output_csv:
880 + formatter.save_csv(all_results, args.output_csv)
881 + if args.output_json:
882 + formatter.save_json(all_results, args.output_json)
883 +
884 +
885 +if __name__ == "__main__":
886 + main()
benchmarks/attention_benchmarks/common.py new
+503
@@ -0,0 +1,503 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Common utilities for attention benchmarking."""
5 +
6 +import csv
7 +import json
8 +import math
9 +from dataclasses import asdict, dataclass
10 +from pathlib import Path
11 +from typing import Any
12 +
13 +import numpy as np
14 +import torch
15 +from rich.console import Console
16 +from rich.table import Table
17 +
18 +# Mock classes for vLLM attention infrastructure
19 +
20 +
21 +class MockHfConfig:
22 + """Mock HuggingFace config that satisfies vLLM's requirements."""
23 +
24 + def __init__(self, mla_dims: dict):
25 + self.num_attention_heads = mla_dims["num_q_heads"]
26 + self.num_key_value_heads = mla_dims["num_kv_heads"]
27 + self.hidden_size = mla_dims["head_dim"] * mla_dims["num_q_heads"]
28 + self.model_type = "deepseek_v2"
29 + self.is_encoder_decoder = False
30 + self.kv_lora_rank = mla_dims["kv_lora_rank"]
31 + self.qk_nope_head_dim = mla_dims["qk_nope_head_dim"]
32 + self.qk_rope_head_dim = mla_dims["qk_rope_head_dim"]
33 + self.v_head_dim = mla_dims["v_head_dim"]
34 + self.qk_head_dim = mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"]
35 +
36 + def get_text_config(self):
37 + return self
38 +
39 +
40 +# Import AttentionLayerBase at module level to avoid circular dependencies
41 +try:
42 + from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
43 +
44 + _HAS_ATTENTION_LAYER_BASE = True
45 +except ImportError:
46 + _HAS_ATTENTION_LAYER_BASE = False
47 + AttentionLayerBase = object # Fallback
48 +
49 +
50 +class MockKVBProj:
51 + """Mock KV projection layer for MLA prefill mode.
52 +
53 + Mimics ColumnParallelLinear behavior for kv_b_proj in MLA backends.
54 + Projects kv_c_normed to [qk_nope_head_dim + v_head_dim] per head.
55 + """
56 +
57 + def __init__(self, num_heads: int, qk_nope_head_dim: int, v_head_dim: int):
58 + self.num_heads = num_heads
59 + self.qk_nope_head_dim = qk_nope_head_dim
60 + self.v_head_dim = v_head_dim
61 + self.out_dim = qk_nope_head_dim + v_head_dim
62 +
63 + def __call__(self, x: torch.Tensor) -> tuple[torch.Tensor]:
64 + """
65 + Project kv_c_normed to output space.
66 +
67 + Args:
68 + x: Input tensor [num_tokens, kv_lora_rank]
69 +
70 + Returns:
71 + Tuple containing output tensor
72 + [num_tokens, num_heads, qk_nope_head_dim + v_head_dim]
73 + """
74 + num_tokens = x.shape[0]
75 + result = torch.randn(
76 + num_tokens,
77 + self.num_heads,
78 + self.out_dim,
79 + device=x.device,
80 + dtype=x.dtype,
81 + )
82 + return (result,) # Return as tuple to match ColumnParallelLinear API
83 +
84 +
85 +class MockLayer(AttentionLayerBase):
86 + """Mock attention layer with scale parameters and impl.
87 +
88 + Inherits from AttentionLayerBase so it passes isinstance checks
89 + in get_layers_from_vllm_config when FlashInfer prefill is enabled.
90 + """
91 +
92 + def __init__(self, device: torch.device, impl=None, kv_cache_spec=None):
93 + # Don't call super().__init__() as AttentionLayerBase doesn't have __init__
94 + self._k_scale = torch.tensor(1.0, device=device)
95 + self._v_scale = torch.tensor(1.0, device=device)
96 + self._q_scale = torch.tensor(1.0, device=device)
97 + # Scalar floats for kernels that need them
98 + self._k_scale_float = float(self._k_scale.item())
99 + self._v_scale_float = float(self._v_scale.item())
100 + self._q_scale_float = float(self._q_scale.item())
101 + # AttentionImpl for metadata builders to query
102 + self.impl = impl
103 + # KV cache spec for get_kv_cache_spec
104 + self._kv_cache_spec = kv_cache_spec
105 +
106 + def get_attn_backend(self):
107 + """Get the attention backend class (required by AttentionLayerBase)."""
108 + # Return None as this is just a mock layer for benchmarking
109 + return None
110 +
111 + def get_kv_cache_spec(self):
112 + """Get the KV cache spec (required by AttentionLayerBase)."""
113 + return self._kv_cache_spec
114 +
115 +
116 +class MockModelConfig:
117 + """Mock model configuration."""
118 +
119 + def __init__(
120 + self,
121 + num_q_heads: int,
122 + num_kv_heads: int,
123 + head_dim: int,
124 + dtype: torch.dtype = torch.float16,
125 + max_model_len: int = 32768,
126 + ):
127 + self._n_q = num_q_heads
128 + self._n_kv = num_kv_heads
129 + self._d = head_dim
130 + self.dtype = dtype
131 + self.max_model_len = max_model_len
132 +
133 + def get_num_attention_heads(self, _=None) -> int:
134 + return self._n_q
135 +
136 + def get_num_kv_heads(self, _=None) -> int:
137 + return self._n_kv
138 +
139 + def get_head_size(self) -> int:
140 + return self._d
141 +
142 + def get_num_layers(self) -> int:
143 + """Mock method for layer count queries."""
144 + return 1
145 +
146 + def get_sliding_window_for_layer(self, _layer_idx: int):
147 + """Mock method for sliding window queries."""
148 + return None
149 +
150 + def get_logits_soft_cap_for_layer(self, _layer_idx: int):
151 + """Mock method for logits soft cap queries."""
152 + return None
153 +
154 + def get_sm_scale_for_layer(self, _layer_idx: int) -> float:
155 + """Mock method for SM scale queries."""
156 + return 1.0 / (self.get_head_size() ** 0.5)
157 +
158 +
159 +class MockParallelConfig:
160 + """Mock parallel configuration."""
161 +
162 + pass
163 +
164 +
165 +class MockCompilationConfig:
166 + """Mock compilation configuration."""
167 +
168 + def __init__(self):
169 + self.full_cuda_graph = False
170 + self.static_forward_context = {}
171 +
172 +
173 +class MockVLLMConfig:
174 + """Mock VLLM configuration."""
175 +
176 + def __init__(self):
177 + self.compilation_config = MockCompilationConfig()
178 +
179 +
180 +class MockRunner:
181 + """Mock GPU runner for metadata builders."""
182 +
183 + def __init__(
184 + self,
185 + seq_lens: np.ndarray,
186 + query_start_locs: np.ndarray,
187 + device: torch.device,
188 + num_q_heads: int,
189 + num_kv_heads: int,
190 + head_dim: int,
191 + dtype: torch.dtype,
192 + ):
193 + self.model_config = MockModelConfig(num_q_heads, num_kv_heads, head_dim, dtype)
194 + self.parallel_config = MockParallelConfig()
195 + self.vllm_config = MockVLLMConfig()
196 + self.seq_lens_np = seq_lens
197 + self.query_start_loc_np = query_start_locs
198 + self.device = device
199 + self.attention_chunk_size = None
200 + self.num_query_heads = num_q_heads
201 + self.num_kv_heads = num_kv_heads
202 + self.dtype = dtype
203 +
204 +
205 +@dataclass
206 +class ParameterSweep:
207 + """Configuration for sweeping a backend parameter."""
208 +
209 + param_name: str # Name of the backend parameter to sweep
210 + values: list[Any] # List of values to test
211 + include_auto: bool = False # Also test with param unset (auto mode)
212 + label_format: str = "{backend}_{param_name}_{value}" # Result label template
213 +
214 + def get_label(self, backend: str, value: Any) -> str:
215 + """Generate a label for a specific parameter value."""
216 + return self.label_format.format(
217 + backend=backend, param_name=self.param_name, value=value
218 + )
219 +
220 +
221 +@dataclass
222 +class ModelParameterSweep:
223 + """Configuration for sweeping a model configuration parameter."""
224 +
225 + param_name: str # Name of the model config parameter to sweep (e.g., "num_q_heads")
226 + values: list[Any] # List of values to test
227 + label_format: str = "{backend}_{param_name}_{value}" # Result label template
228 +
229 + def get_label(self, backend: str, value: Any) -> str:
230 + """Generate a label for a specific parameter value."""
231 + return self.label_format.format(
232 + backend=backend, param_name=self.param_name, value=value
233 + )
234 +
235 +
236 +@dataclass
237 +class BenchmarkConfig:
238 + """Configuration for a single benchmark run."""
239 +
240 + backend: str
241 + batch_spec: str
242 + num_layers: int
243 + head_dim: int
244 + num_q_heads: int
245 + num_kv_heads: int
246 + block_size: int
247 + device: str
248 + dtype: torch.dtype = torch.float16
249 + repeats: int = 1
250 + warmup_iters: int = 3
251 + profile_memory: bool = False
252 + use_cuda_graphs: bool = False
253 +
254 + # MLA-specific
255 + kv_lora_rank: int | None = None
256 + qk_nope_head_dim: int | None = None
257 + qk_rope_head_dim: int | None = None
258 + v_head_dim: int | None = None
259 +
260 + # Backend-specific tuning
261 + num_kv_splits: int | None = None # CUTLASS MLA
262 + reorder_batch_threshold: int | None = None # FlashAttn MLA, FlashMLA
263 +
264 +
265 +@dataclass
266 +class BenchmarkResult:
267 + """Results from a single benchmark run."""
268 +
269 + config: BenchmarkConfig
270 + mean_time: float # seconds
271 + std_time: float # seconds
272 + min_time: float # seconds
273 + max_time: float # seconds
274 + throughput_tokens_per_sec: float | None = None
275 + memory_allocated_mb: float | None = None
276 + memory_reserved_mb: float | None = None
277 + error: str | None = None
278 +
279 + @property
280 + def success(self) -> bool:
281 + """Whether benchmark completed successfully."""
282 + return self.error is None
283 +
284 + def to_dict(self) -> dict[str, Any]:
285 + """Convert to dictionary for serialization."""
286 + return {
287 + "config": asdict(self.config),
288 + "mean_time": self.mean_time,
289 + "std_time": self.std_time,
290 + "min_time": self.min_time,
291 + "max_time": self.max_time,
292 + "throughput_tokens_per_sec": self.throughput_tokens_per_sec,
293 + "memory_allocated_mb": self.memory_allocated_mb,
294 + "memory_reserved_mb": self.memory_reserved_mb,
295 + "error": self.error,
296 + }
297 +
298 +
299 +class ResultsFormatter:
300 + """Format and display benchmark results."""
301 +
302 + def __init__(self, console: Console | None = None):
303 + self.console = console or Console()
304 +
305 + def print_table(
306 + self,
307 + results: list[BenchmarkResult],
308 + backends: list[str],
309 + compare_to_fastest: bool = True,
310 + ):
311 + """
312 + Print results as a rich table.
313 +
314 + Args:
315 + results: List of BenchmarkResult
316 + backends: List of backend names being compared
317 + compare_to_fastest: Show percentage comparison to fastest
318 + """
319 + # Group by batch spec
320 + by_spec = {}
321 + for r in results:
322 + spec = r.config.batch_spec
323 + if spec not in by_spec:
324 + by_spec[spec] = {}
325 + by_spec[spec][r.config.backend] = r
326 +
327 + # Create shortened backend names for display
328 + def shorten_backend_name(name: str) -> str:
329 + """Shorten long backend names for table display."""
330 + # Remove common prefixes
331 + name = name.replace("flashattn_mla", "famla")
332 + name = name.replace("flashinfer_mla", "fimla")
333 + name = name.replace("flashmla", "fmla")
334 + name = name.replace("cutlass_mla", "cmla")
335 + name = name.replace("numsplits", "ns")
336 + return name
337 +
338 + table = Table(title="Attention Benchmark Results")
339 + table.add_column("Batch\nSpec", no_wrap=True)
340 +
341 + multi = len(backends) > 1
342 + for backend in backends:
343 + short_name = shorten_backend_name(backend)
344 + # Time column
345 + col_time = f"{short_name}\nTime (s)"
346 + table.add_column(col_time, justify="right", no_wrap=False)
347 + if multi and compare_to_fastest:
348 + # Relative performance column
349 + col_rel = f"{short_name}\nvs Best"
350 + table.add_column(col_rel, justify="right", no_wrap=False)
351 +
352 + # Add rows
353 + for spec in sorted(by_spec.keys()):
354 + spec_results = by_spec[spec]
355 + times = {b: r.mean_time for b, r in spec_results.items() if r.success}
356 + best_time = min(times.values()) if times else 0.0
357 +
358 + row = [spec]
359 + for backend in backends:
360 + if backend in spec_results:
361 + r = spec_results[backend]
362 + if r.success:
363 + row.append(f"{r.mean_time:.6f}")
364 + if multi and compare_to_fastest:
365 + pct = (
366 + (r.mean_time / best_time * 100) if best_time > 0 else 0
367 + )
368 + pct_str = f"{pct:.1f}%"
369 + if r.mean_time == best_time:
370 + pct_str = f"[bold green]{pct_str}[/]"
371 + row.append(pct_str)
372 + else:
373 + row.append("[red]ERROR[/]")
374 + if multi and compare_to_fastest:
375 + row.append("-")
376 + else:
377 + row.append("-")
378 + if multi and compare_to_fastest:
379 + row.append("-")
380 +
381 + table.add_row(*row)
382 +
383 + self.console.print(table)
384 +
385 + def save_csv(self, results: list[BenchmarkResult], path: str):
386 + """Save results to CSV file."""
387 + if not results:
388 + return
389 +
390 + path_obj = Path(path)
391 + path_obj.parent.mkdir(parents=True, exist_ok=True)
392 +
393 + with open(path, "w", newline="") as f:
394 + writer = csv.DictWriter(
395 + f,
396 + fieldnames=[
397 + "backend",
398 + "batch_spec",
399 + "num_layers",
400 + "mean_time",
401 + "std_time",
402 + "throughput",
403 + "memory_mb",
404 + ],
405 + )
406 + writer.writeheader()
407 + for r in results:
408 + writer.writerow(
409 + {
410 + "backend": r.config.backend,
411 + "batch_spec": r.config.batch_spec,
412 + "num_layers": r.config.num_layers,
413 + "mean_time": r.mean_time,
414 + "std_time": r.std_time,
415 + "throughput": r.throughput_tokens_per_sec or 0,
416 + "memory_mb": r.memory_allocated_mb or 0,
417 + }
418 + )
419 +
420 + self.console.print(f"[green]Saved CSV results to {path}[/]")
421 +
422 + def save_json(self, results: list[BenchmarkResult], path: str):
423 + """Save results to JSON file."""
424 + path_obj = Path(path)
425 + path_obj.parent.mkdir(parents=True, exist_ok=True)
426 +
427 + data = [r.to_dict() for r in results]
428 + with open(path, "w") as f:
429 + json.dump(data, f, indent=2, default=str)
430 +
431 + self.console.print(f"[green]Saved JSON results to {path}[/]")
432 +
433 +
434 +def setup_mla_dims(model_name: str = "deepseek-v3") -> dict:
435 + """
436 + Get MLA dimensions for known models.
437 +
438 + Args:
439 + model_name: Model identifier
440 +
441 + Returns:
442 + Dict with MLA dimension configuration
443 + """
444 + configs = {
445 + "deepseek-v2": {
446 + "kv_lora_rank": 512,
447 + "qk_nope_head_dim": 128,
448 + "qk_rope_head_dim": 64,
449 + "v_head_dim": 128,
450 + "num_q_heads": 128,
451 + "num_kv_heads": 1,
452 + "head_dim": 576,
453 + },
454 + "deepseek-v3": {
455 + "kv_lora_rank": 512,
456 + "qk_nope_head_dim": 128,
457 + "qk_rope_head_dim": 64,
458 + "v_head_dim": 128,
459 + "num_q_heads": 128,
460 + "num_kv_heads": 1,
461 + "head_dim": 576,
462 + },
463 + "deepseek-v2-lite": {
464 + "kv_lora_rank": 512,
465 + "qk_nope_head_dim": 128,
466 + "qk_rope_head_dim": 64,
467 + "v_head_dim": 128,
468 + "num_q_heads": 16,
469 + "num_kv_heads": 1,
470 + "head_dim": 576,
471 + },
472 + }
473 +
474 + if model_name not in configs:
475 + raise ValueError(
476 + f"Unknown model '{model_name}'. Known models: {list(configs.keys())}"
477 + )
478 +
479 + return configs[model_name]
480 +
481 +
482 +def get_attention_scale(head_dim: int) -> float:
483 + """Compute attention scale factor (1/sqrt(d))."""
484 + return 1.0 / math.sqrt(head_dim)
485 +
486 +
487 +def is_mla_backend(backend: str) -> bool:
488 + """
489 + Check if backend is an MLA backend using the backend's is_mla() property.
490 +
491 + Args:
492 + backend: Backend name (e.g., "CUTLASS_MLA", "FLASHINFER_MLA")
493 +
494 + Returns:
495 + True if the backend is an MLA backend, False otherwise
496 + """
497 + from vllm.v1.attention.backends.registry import AttentionBackendEnum
498 +
499 + try:
500 + backend_class = AttentionBackendEnum[backend.upper()].get_class()
501 + return backend_class.is_mla()
502 + except (KeyError, ValueError, ImportError):
503 + return False
benchmarks/attention_benchmarks/configs/mla_decode.yaml new
+61
@@ -0,0 +1,61 @@
1 +# MLA decode-only benchmark configuration
2 +
3 +model:
4 + name: "deepseek-v3"
5 + num_layers: 60
6 + num_q_heads: 128
7 + num_kv_heads: 1 # MLA uses single latent KV
8 + head_dim: 576
9 + kv_lora_rank: 512
10 + qk_nope_head_dim: 128
11 + qk_rope_head_dim: 64
12 + v_head_dim: 128
13 + block_size: 128 # CUTLASS MLA and FlashAttn MLA use 128
14 +
15 +batch_specs:
16 + # Small batches, varying sequence lengths
17 + - "16q1s512" # 16 requests, 512 KV cache
18 + - "16q1s1k" # 16 requests, 1k KV cache
19 + - "16q1s2k" # 16 requests, 2k KV cache
20 + - "16q1s4k" # 16 requests, 4k KV cache
21 +
22 + # Medium batches
23 + - "32q1s1k" # 32 requests, 1k KV cache
24 + - "32q1s2k" # 32 requests, 2k KV cache
25 + - "32q1s4k" # 32 requests, 4k KV cache
26 + - "32q1s8k" # 32 requests, 8k KV cache
27 +
28 + # Large batches
29 + - "64q1s1k" # 64 requests, 1k KV cache
30 + - "64q1s2k" # 64 requests, 2k KV cache
31 + - "64q1s4k" # 64 requests, 4k KV cache
32 + - "64q1s8k" # 64 requests, 8k KV cache
33 +
34 + # Very large batches
35 + - "128q1s1k" # 128 requests, 1k KV cache
36 + - "128q1s2k" # 128 requests, 2k KV cache
37 +
38 + # Long context
39 + - "32q1s16k" # 32 requests, 16k KV cache
40 + - "32q1s32k" # 32 requests, 32k KV cache
41 +
42 +backends:
43 + - cutlass_mla
44 + - flashinfer_mla
45 + - flashattn_mla # Hopper only
46 + - flashmla # Hopper only
47 +
48 +device: "cuda:0"
49 +repeats: 5
50 +warmup_iters: 3
51 +profile_memory: true
52 +
53 +# Backend-specific tuning
54 +cutlass_mla:
55 + num_kv_splits: auto # or specific value like 4, 8, 16
56 +
57 +flashattn_mla:
58 + reorder_batch_threshold: 512
59 +
60 +flashmla:
61 + reorder_batch_threshold: 1
benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml new
+60
@@ -0,0 +1,60 @@
1 +# MLA mixed batch benchmark (prefill + decode)
2 +# Tests chunked prefill performance
3 +
4 +model:
5 + name: "deepseek-v3"
6 + num_layers: 60
7 + num_q_heads: 128
8 + num_kv_heads: 1
9 + head_dim: 576
10 + kv_lora_rank: 512
11 + qk_nope_head_dim: 128
12 + qk_rope_head_dim: 64
13 + v_head_dim: 128
14 + block_size: 128
15 +
16 +batch_specs:
17 + # Small prefill + decode
18 + - "1q1k_8q1s1k" # 1 prefill + 8 decode
19 + - "2q2k_16q1s1k" # 2 prefill + 16 decode
20 + - "4q1k_32q1s2k" # 4 prefill + 32 decode
21 +
22 + # Medium prefill + decode
23 + - "2q4k_32q1s2k" # 2 medium prefill + 32 decode
24 + - "4q4k_64q1s2k" # 4 medium prefill + 64 decode
25 + - "8q2k_64q1s4k" # 8 prefill + 64 decode
26 +
27 + # Large prefill + decode (chunked prefill stress test)
28 + - "2q8k_32q1s1k" # 2 large prefill + 32 decode
29 + - "1q16k_16q1s2k" # 1 very large prefill + 16 decode
30 + - "2q16k_32q1s4k" # 2 very large prefill + 32 decode
31 +
32 + # Context extension + decode
33 + - "2q1kkv2k_16q1s1k" # 2 extend + 16 decode
34 + - "4q2kkv4k_32q1s2k" # 4 extend + 32 decode
35 + - "2q1kkv8k_32q1s2k" # 2 large extend + 32 decode
36 +
37 + # Explicitly chunked prefill
38 + - "q8k" # 8k prefill with chunking hint
39 + - "q16k" # 16k prefill with chunking hint
40 + - "2q8k_32q1s2k" # 2 chunked prefill + 32 decode
41 +
42 + # High decode ratio (realistic serving)
43 + - "1q2k_63q1s1k" # 1 prefill + 63 decode
44 + - "2q2k_62q1s2k" # 2 prefill + 62 decode
45 + - "4q4k_60q1s4k" # 4 prefill + 60 decode
46 +
47 +backends:
48 + - cutlass_mla
49 + - flashinfer_mla
50 + - flashattn_mla # Hopper only
51 + - flashmla # Hopper only
52 +
53 +device: "cuda:0"
54 +repeats: 5
55 +warmup_iters: 3
56 +profile_memory: true
57 +
58 +# Analyze chunked prefill workspace size impact
59 +chunked_prefill:
60 + test_workspace_sizes: [4096, 8192, 16384, 32768, 65536]
benchmarks/attention_benchmarks/configs/reorder_threshold.yaml new
+88
@@ -0,0 +1,88 @@
1 +# Study 4: What is optimal reorder_batch_threshold for MLA backends supporting query length > 1?
2 +# Question: At what query length does prefill pipeline become faster than decode pipeline?
3 +# Methodology: For each query length, compare decode vs prefill performance to find crossover point
4 +# Applies to: FlashAttn MLA, FlashMLA
5 +
6 +description: "Decode vs Prefill pipeline crossover analysis"
7 +
8 +# Test FlashAttn MLA
9 +backend: flashattn_mla
10 +
11 +# Mode: decode_vs_prefill comparison (special sweep mode)
12 +# For each batch spec, we'll test both decode and prefill pipelines
13 +mode: "decode_vs_prefill"
14 +
15 +# Query lengths to test (from old benchmark_mla_threshold.py methodology)
16 +# Each query length will be tested with BOTH decode and prefill pipelines:
17 +# - decode: threshold >= query_length (forces decode pipeline)
18 +# - prefill: threshold < query_length (forces prefill pipeline)
19 +#
20 +# We use q<N>s1k format which creates q_len=N, seq_len=1024 requests
21 +# This tests different query lengths with fixed sequence length context
22 +#
23 +# Using batch_spec_ranges for automatic generation:
24 +batch_spec_ranges:
25 + - template: "q{q_len}s1k"
26 + q_len:
27 + start: 1
28 + stop: 16
29 + step: 1
30 + end_inclusive: false
31 + - template: "q{q_len}s1k"
32 + q_len:
33 + start: 16
34 + stop: 64
35 + step: 2
36 + end_inclusive: false
37 + - template: "q{q_len}s1k"
38 + q_len:
39 + start: 64
40 + stop: 1024
41 + step: 4
42 + end_inclusive: true
43 +
44 +# Batch sizes to test (from old script)
45 +batch_sizes:
46 + - 1
47 + - 2
48 + - 4
49 + - 8
50 + - 16
51 + - 32
52 + - 64
53 + - 128
54 + - 256
55 +
56 +# Model configuration (DeepSeek V2/V3 defaults)
57 +model:
58 + num_layers: 10
59 + head_dim: 576
60 + num_q_heads: 128
61 + num_kv_heads: 1
62 + block_size: 128
63 +
64 +# Benchmark settings
65 +benchmark:
66 + device: "cuda:0"
67 + repeats: 15 # More repeats for spec decode variance
68 + warmup_iters: 5
69 + profile_memory: false
70 +
71 +# Output
72 +output:
73 + csv: "reorder_threshold_results.csv"
74 + json: "reorder_threshold_results.json"
75 +
76 +# Expected outcome (reproduces old benchmark_mla_threshold.py study):
77 +# - For each batch size, find the crossover point where prefill becomes faster than decode
78 +# - Show decode vs prefill performance across all query lengths
79 +# - Determine optimal reorder_batch_threshold based on last query length where decode is faster
80 +# - Understand how crossover point varies with batch size
81 +# - Provide data-driven guidance for default threshold value
82 +#
83 +# Methodology (from old script):
84 +# - Each query length tested with BOTH pipelines:
85 +# * decode: threshold >= query_length (forces decode pipeline)
86 +# * prefill: threshold < query_length (forces prefill pipeline)
87 +# - Compare which is faster to find crossover point
88 +#
benchmarks/attention_benchmarks/configs/speculative_decode.yaml new
+62
@@ -0,0 +1,62 @@
1 +# Speculative decoding benchmark configuration
2 +# Tests reorder_batch_threshold optimization
3 +
4 +model:
5 + name: "deepseek-v3"
6 + num_layers: 60
7 + num_q_heads: 128
8 + num_kv_heads: 1
9 + head_dim: 576
10 + kv_lora_rank: 512
11 + qk_nope_head_dim: 128
12 + qk_rope_head_dim: 64
13 + v_head_dim: 128
14 +
15 +batch_specs:
16 + # Pure speculative decode (K-token verification)
17 + - "q2s1k" # 2-token spec, 1k KV
18 + - "q4s1k" # 4-token spec, 1k KV
19 + - "q8s1k" # 8-token spec, 1k KV
20 + - "q16s1k" # 16-token spec, 1k KV
21 +
22 + # Speculative with different context lengths
23 + - "q4s2k" # 4-token spec, 2k KV
24 + - "q4s4k" # 4-token spec, 4k KV
25 + - "q8s2k" # 8-token spec, 2k KV
26 + - "q8s4k" # 8-token spec, 4k KV
27 +
28 + # Mixed: speculative + regular decode
29 + - "32q4s1k" # 32 spec requests
30 + - "16q4s1k_16q1s1k" # 16 spec + 16 regular
31 + - "8q8s2k_24q1s2k" # 8 spec (8-tok) + 24 regular
32 +
33 + # Mixed: speculative + prefill + decode
34 + - "2q1k_16q4s1k_16q1s1k" # 2 prefill + 16 spec + 16 decode
35 + - "4q2k_32q4s2k_32q1s2k" # 4 prefill + 32 spec + 32 decode
36 +
37 + # Large batches with speculation
38 + - "64q4s1k" # 64 spec requests
39 + - "32q8s2k" # 32 spec (8-token)
40 + - "16q16s4k" # 16 spec (16-token)
41 +
42 +# Backends that support query length > 1
43 +backends:
44 + - flashattn_mla # reorder_batch_threshold = 512
45 + - flashmla # reorder_batch_threshold = 1 (tunable)
46 +
47 +# FlashInfer-MLA also supports uniform spec-as-decode but with different mechanism
48 +# - flashinfer_mla
49 +
50 +# Benchmark settings
51 +benchmark:
52 + device: "cuda:0"
53 + repeats: 10 # More repeats for statistical significance
54 + warmup_iters: 5
55 + profile_memory: false
56 +
57 +# Test these threshold values for optimization
58 +parameter_sweep:
59 + param_name: "reorder_batch_threshold"
60 + values: [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
61 + include_auto: false
62 + label_format: "{backend}_threshold_{value}"
benchmarks/attention_benchmarks/configs/standard_attention.yaml new
+40
@@ -0,0 +1,40 @@
1 +# Standard attention backend benchmark configuration
2 +
3 +model:
4 + num_layers: 32
5 + num_q_heads: 32
6 + num_kv_heads: 8 # GQA with 4:1 ratio
7 + head_dim: 128
8 + block_size: 16
9 +
10 +batch_specs:
11 + # Pure prefill
12 + - "q512" # Small prefill (512 tokens)
13 + - "q2k" # Medium prefill (2048 tokens)
14 + - "q4k" # Large prefill (4096 tokens)
15 + - "q8k" # Very large prefill (8192 tokens)
16 +
17 + # Pure decode
18 + - "8q1s1k" # 8 requests, 1k KV cache each
19 + - "16q1s2k" # 16 requests, 2k KV cache each
20 + - "32q1s1k" # 32 requests, 1k KV cache each
21 + - "64q1s4k" # 64 requests, 4k KV cache each
22 +
23 + # Mixed prefill/decode
24 + - "2q2k_8q1s1k" # 2 prefill + 8 decode
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)
30 + - "2q1ks4k" # 2 requests: 1k query, 4k sequence
31 +
32 +backends:
33 + - flash
34 + - triton
35 + - flashinfer
36 +
37 +device: "cuda:0"
38 +repeats: 5
39 +warmup_iters: 3
40 +profile_memory: false
benchmarks/attention_benchmarks/mla_runner.py new
+836
@@ -0,0 +1,836 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""
5 +MLA benchmark runner - shared utilities for MLA benchmarks.
6 +
7 +This module provides helpers for running MLA backends without
8 +needing full VllmConfig integration.
9 +"""
10 +
11 +import importlib
12 +
13 +import numpy as np
14 +import torch
15 +from batch_spec import parse_batch_spec
16 +from common import (
17 + BenchmarkResult,
18 + MockHfConfig,
19 + MockKVBProj,
20 + MockLayer,
21 + setup_mla_dims,
22 +)
23 +
24 +from vllm.config import (
25 + CacheConfig,
26 + CompilationConfig,
27 + ModelConfig,
28 + ParallelConfig,
29 + SchedulerConfig,
30 + VllmConfig,
31 + set_current_vllm_config,
32 +)
33 +
34 +# ============================================================================
35 +# VllmConfig Creation
36 +# ============================================================================
37 +
38 +
39 +def _add_mock_methods_to_model_config(model_config: ModelConfig) -> None:
40 + """
41 + Add mock methods for layer-specific queries to ModelConfig.
42 +
43 + These methods are needed by metadata builders but aren't normally
44 + present on ModelConfig when used in benchmark contexts.
45 + """
46 + import types
47 +
48 + model_config.get_num_layers = types.MethodType(lambda self: 1, model_config)
49 + model_config.get_sliding_window_for_layer = types.MethodType(
50 + lambda self, _i: None, model_config
51 + )
52 + model_config.get_logits_soft_cap_for_layer = types.MethodType(
53 + lambda self, _i: None, model_config
54 + )
55 + model_config.get_sm_scale_for_layer = types.MethodType(
56 + lambda self, _i: 1.0 / model_config.get_head_size() ** 0.5, model_config
57 + )
58 +
59 +
60 +def create_minimal_vllm_config(
61 + model_name: str = "deepseek-v3",
62 + block_size: int = 128,
63 + max_num_seqs: int = 256,
64 + mla_dims: dict | None = None,
65 +) -> VllmConfig:
66 + """
67 + Create minimal VllmConfig for MLA benchmarks.
68 +
69 + Args:
70 + model_name: Model name (deepseek-v2, deepseek-v3, etc.) - used if mla_dims not
71 + provided
72 + block_size: KV cache block size
73 + max_num_seqs: Maximum number of sequences
74 + mla_dims: Optional custom MLA dimensions dict. If not provided, uses
75 + setup_mla_dims(model_name)
76 +
77 + Returns:
78 + VllmConfig for benchmarking
79 + """
80 + # Get MLA dimensions - use provided or load from model name
81 + if mla_dims is None:
82 + mla_dims = setup_mla_dims(model_name)
83 +
84 + # Create mock HF config first (avoids downloading from HuggingFace)
85 + mock_hf_config = MockHfConfig(mla_dims)
86 +
87 + # Create a temporary minimal config.json to avoid HF downloads
88 + # This ensures consistent ModelConfig construction without network access
89 + import json
90 + import os
91 + import shutil
92 + import tempfile
93 +
94 + minimal_config = {
95 + "architectures": ["DeepseekV2ForCausalLM"],
96 + "model_type": "deepseek_v2",
97 + "num_attention_heads": mla_dims["num_q_heads"],
98 + "num_key_value_heads": mla_dims["num_kv_heads"],
99 + "hidden_size": mla_dims["head_dim"] * mla_dims["num_q_heads"],
100 + "torch_dtype": "bfloat16",
101 + "max_position_embeddings": 163840, # DeepSeek V3 default
102 + "rope_theta": 10000.0,
103 + "vocab_size": 128256,
104 + }
105 +
106 + # Create temporary directory with config.json
107 + temp_dir = tempfile.mkdtemp(prefix="vllm_bench_")
108 + config_path = os.path.join(temp_dir, "config.json")
109 + with open(config_path, "w") as f:
110 + json.dump(minimal_config, f)
111 +
112 + try:
113 + # Create model config using local path - no HF downloads
114 + model_config = ModelConfig(
115 + model=temp_dir, # Use local temp directory
116 + tokenizer=None,
117 + tokenizer_mode="auto",
118 + trust_remote_code=True,
119 + dtype="bfloat16",
120 + seed=0,
121 + max_model_len=32768,
122 + quantization=None,
123 + quantization_param_path=None,
124 + enforce_eager=False,
125 + max_context_len_to_capture=None,
126 + max_seq_len_to_capture=8192,
127 + max_logprobs=20,
128 + disable_sliding_window=False,
129 + skip_tokenizer_init=True,
130 + served_model_name=None,
131 + limit_mm_per_prompt=None,
132 + use_async_output_proc=True,
133 + config_format="auto",
134 + )
135 + finally:
136 + # Clean up temporary directory
137 + shutil.rmtree(temp_dir, ignore_errors=True)
138 +
139 + # Override with our mock config
140 + model_config.hf_config = mock_hf_config
141 + model_config.hf_text_config = mock_hf_config
142 +
143 + # Add mock methods for layer-specific queries
144 + _add_mock_methods_to_model_config(model_config)
145 +
146 + # Create sub-configs
147 + cache_config = CacheConfig(
148 + block_size=block_size,
149 + gpu_memory_utilization=0.9,
150 + swap_space=0,
151 + cache_dtype="auto",
152 + enable_prefix_caching=False,
153 + )
154 +
155 + scheduler_config = SchedulerConfig(
156 + max_num_seqs=max_num_seqs,
157 + max_num_batched_tokens=8192,
158 + max_model_len=32768,
159 + is_encoder_decoder=False,
160 + enable_chunked_prefill=True,
161 + )
162 +
163 + parallel_config = ParallelConfig(
164 + tensor_parallel_size=1,
165 + )
166 +
167 + compilation_config = CompilationConfig()
168 +
169 + return VllmConfig(
170 + model_config=model_config,
171 + cache_config=cache_config,
172 + parallel_config=parallel_config,
173 + scheduler_config=scheduler_config,
174 + compilation_config=compilation_config,
175 + )
176 +
177 +
178 +# ============================================================================
179 +# Backend Configuration
180 +# ============================================================================
181 +
182 +
183 +# Backend name to class name prefix mapping
184 +_BACKEND_NAME_MAP = {
185 + "flashattn_mla": "FlashAttnMLA",
186 + "flashmla": "FlashMLA",
187 + "flashinfer_mla": "FlashInferMLA",
188 + "cutlass_mla": "CutlassMLA",
189 +}
190 +
191 +# Special properties that differ from defaults
192 +_BACKEND_PROPERTIES = {
193 + "flashmla": {
194 + "query_format": "concat", # Single concatenated tensor (vs tuple)
195 + "block_size": 64, # FlashMLA uses fixed block size
196 + },
197 + "flashinfer_mla": {
198 + "block_size": 64, # FlashInfer MLA only supports 32 or 64
199 + },
200 +}
201 +
202 +
203 +def _get_backend_config(backend: str) -> dict:
204 + """
205 + Get backend configuration using naming conventions.
206 +
207 + All MLA backends follow the pattern:
208 + - Module: vllm.v1.attention.backends.mla.{backend}
209 + - Impl: {Name}Impl
210 + - Metadata: {Name}Metadata (or MLACommonMetadata)
211 + - DecodeMetadata: {Name}DecodeMetadata (or MLACommonDecodeMetadata)
212 + - MetadataBuilder: {Name}MetadataBuilder
213 + """
214 + if backend not in _BACKEND_NAME_MAP:
215 + raise ValueError(f"Unknown backend: {backend}")
216 +
217 + name = _BACKEND_NAME_MAP[backend]
218 + props = _BACKEND_PROPERTIES.get(backend, {})
219 +
220 + # Check if backend uses common metadata (FlashInfer, CUTLASS)
221 + uses_common = backend in ("flashinfer_mla", "cutlass_mla")
222 +
223 + return {
224 + "module": f"vllm.v1.attention.backends.mla.{backend}",
225 + "impl_class": f"{name}Impl",
226 + "metadata_class": "MLACommonMetadata" if uses_common else f"{name}Metadata",
227 + "decode_metadata_class": "MLACommonDecodeMetadata"
228 + if uses_common
229 + else f"{name}DecodeMetadata",
230 + "builder_class": f"{name}MetadataBuilder",
231 + "query_format": props.get("query_format", "tuple"),
232 + "block_size": props.get("block_size", None),
233 + }
234 +
235 +
236 +# ============================================================================
237 +# Metadata Building Helpers
238 +# ============================================================================
239 +
240 +
241 +def _build_attention_metadata(
242 + requests: list,
243 + block_size: int,
244 + device: torch.device,
245 + builder_instance,
246 +) -> tuple:
247 + """
248 + Build attention metadata from batch requests.
249 +
250 + Args:
251 + requests: List of BatchRequest objects
252 + block_size: KV cache block size
253 + device: Target device
254 + builder_instance: Metadata builder instance
255 +
256 + Returns:
257 + Tuple of (metadata, kv_cache_num_blocks)
258 + """
259 + q_lens = [r.q_len for r in requests]
260 + kv_lens = [r.kv_len for r in requests]
261 + total_q = sum(q_lens)
262 + max_kv = max(kv_lens)
263 +
264 + # Build query start locations
265 + q_start_cpu = torch.tensor(
266 + [0] + [sum(q_lens[: i + 1]) for i in range(len(q_lens))],
267 + dtype=torch.int32,
268 + )
269 + q_start_gpu = q_start_cpu.to(device)
270 +
271 + # Build sequence lengths
272 + seq_lens_cpu = torch.tensor(kv_lens, dtype=torch.int32)
273 + seq_lens_gpu = seq_lens_cpu.to(device)
274 +
275 + # Build num_computed_tokens (context length for each request)
276 + context_lens = [kv_len - q_len for q_len, kv_len in zip(q_lens, kv_lens)]
277 + num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
278 +
279 + # Build block table
280 + num_blocks_per_req = [(kv + block_size - 1) // block_size for kv in kv_lens]
281 + max_num_blocks = max(num_blocks_per_req)
282 +
283 + block_table_cpu = np.zeros((len(requests), max_num_blocks), dtype=np.int32)
284 + current_block = 0
285 + for i, num_blocks in enumerate(num_blocks_per_req):
286 + for j in range(num_blocks):
287 + block_table_cpu[i, j] = current_block
288 + current_block += 1
289 +
290 + block_table_gpu = torch.from_numpy(block_table_cpu).to(device)
291 +
292 + # Build slot mapping
293 + slot_mapping_list = []
294 + for i, (q_len, kv_len, num_blocks) in enumerate(
295 + zip(q_lens, kv_lens, num_blocks_per_req)
296 + ):
297 + context_len = kv_len - q_len
298 + for j in range(q_len):
299 + token_kv_idx = context_len + j
300 + block_idx = token_kv_idx // block_size
301 + offset_in_block = token_kv_idx % block_size
302 + global_block_id = block_table_cpu[i, block_idx]
303 + slot_id = global_block_id * block_size + offset_in_block
304 + slot_mapping_list.append(slot_id)
305 +
306 + slot_mapping = torch.tensor(slot_mapping_list, dtype=torch.int64, device=device)
307 +
308 + # Create CommonAttentionMetadata
309 + from vllm.v1.attention.backends.utils import CommonAttentionMetadata
310 +
311 + common_attn_metadata = CommonAttentionMetadata(
312 + num_reqs=len(requests),
313 + max_query_len=max(q_lens),
314 + max_seq_len=max_kv,
315 + num_actual_tokens=total_q,
316 + query_start_loc=q_start_gpu,
317 + query_start_loc_cpu=q_start_cpu,
318 + seq_lens=seq_lens_gpu,
319 + _seq_lens_cpu=seq_lens_cpu,
320 + _num_computed_tokens_cpu=num_computed_tokens_cpu,
321 + slot_mapping=slot_mapping,
322 + block_table_tensor=block_table_gpu,
323 + dcp_local_seq_lens=None,
324 + )
325 +
326 + # Use the production build() method
327 + metadata = builder_instance.build(
328 + common_prefix_len=0,
329 + common_attn_metadata=common_attn_metadata,
330 + fast_build=False,
331 + )
332 +
333 + return metadata, current_block
334 +
335 +
336 +def _create_input_tensors(
337 + total_q: int,
338 + mla_dims: dict,
339 + query_format: str,
340 + device: torch.device,
341 + dtype: torch.dtype,
342 +):
343 + """
344 + Create input tensors for both decode and prefill modes.
345 +
346 + MLA requires different tensor formats for decode vs prefill:
347 + - Decode: Uses kv_lora_rank (512) dimension
348 + - Prefill: Uses qk_nope_head_dim (128) to stay under FlashAttention's 256 limit
349 +
350 + Args:
351 + total_q: Total number of query tokens
352 + mla_dims: MLA dimension configuration
353 + query_format: Either "tuple" or "concat"
354 + device: Target device
355 + dtype: Tensor dtype
356 +
357 + Returns:
358 + Tuple of (decode_inputs, prefill_inputs)
359 + - decode_inputs: Query tensor(s) for decode mode
360 + - prefill_inputs: Dict with 'q', 'k_c_normed', 'k_pe', 'k_scale' for prefill
361 + """
362 + if query_format == "tuple":
363 + # Decode mode format: (q_nope, q_pe) where q_nope has kv_lora_rank dim
364 + q_nope_decode = torch.randn(
365 + total_q,
366 + mla_dims["num_q_heads"],
367 + mla_dims["kv_lora_rank"],
368 + device=device,
369 + dtype=dtype,
370 + )
371 + q_pe = torch.randn(
372 + total_q,
373 + mla_dims["num_q_heads"],
374 + mla_dims["qk_rope_head_dim"],
375 + device=device,
376 + dtype=dtype,
377 + )
378 + decode_inputs = (q_nope_decode, q_pe)
379 +
380 + # For prefill, we need q with qk_nope_head_dim instead of kv_lora_rank
381 + q_nope_prefill = torch.randn(
382 + total_q,
383 + mla_dims["num_q_heads"],
384 + mla_dims["qk_nope_head_dim"],
385 + device=device,
386 + dtype=dtype,
387 + )
388 + prefill_q = torch.cat([q_nope_prefill, q_pe], dim=-1)
389 + else: # concat
390 + decode_inputs = torch.randn(
391 + total_q,
392 + mla_dims["num_q_heads"],
393 + mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"],
394 + device=device,
395 + dtype=dtype,
396 + )
397 + # For prefill with concat format
398 + prefill_q = torch.randn(
399 + total_q,
400 + mla_dims["num_q_heads"],
401 + mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"],
402 + device=device,
403 + dtype=dtype,
404 + )
405 +
406 + # Create additional inputs needed for prefill forward
407 + k_c_normed = torch.randn(
408 + total_q,
409 + mla_dims["kv_lora_rank"],
410 + device=device,
411 + dtype=dtype,
412 + )
413 + k_pe = torch.randn(
414 + total_q,
415 + 1, # Single head for MLA
416 + mla_dims["qk_rope_head_dim"],
417 + device=device,
418 + dtype=dtype,
419 + )
420 + k_scale = torch.ones(1, device=device, dtype=torch.float32)
421 +
422 + output = torch.zeros(
423 + total_q,
424 + mla_dims["num_q_heads"] * mla_dims["v_head_dim"],
425 + device=device,
426 + dtype=dtype,
427 + )
428 +
429 + prefill_inputs = {
430 + "q": prefill_q,
431 + "k_c_normed": k_c_normed,
432 + "k_pe": k_pe,
433 + "k_scale": k_scale,
434 + "output": output,
435 + }
436 +
437 + return decode_inputs, prefill_inputs
438 +
439 +
440 +# ============================================================================
441 +# Backend Initialization
442 +# ============================================================================
443 +
444 +
445 +def _create_backend_impl(
446 + backend_cfg: dict,
447 + mla_dims: dict,
448 + vllm_config: VllmConfig,
449 + device: torch.device,
450 +):
451 + """
452 + Create backend implementation instance.
453 +
454 + Args:
455 + backend_cfg: Backend configuration dict
456 + mla_dims: MLA dimension configuration
457 + vllm_config: VllmConfig instance
458 + device: Target device
459 +
460 + Returns:
461 + Tuple of (impl, layer, builder_instance)
462 + """
463 + # Import backend classes
464 + backend_module = importlib.import_module(backend_cfg["module"])
465 + impl_class = getattr(backend_module, backend_cfg["impl_class"])
466 +
467 + # Calculate scale
468 + scale = 1.0 / np.sqrt(mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"])
469 +
470 + # Create mock kv_b_proj layer for prefill mode
471 + mock_kv_b_proj = MockKVBProj(
472 + num_heads=mla_dims["num_q_heads"],
473 + qk_nope_head_dim=mla_dims["qk_nope_head_dim"],
474 + v_head_dim=mla_dims["v_head_dim"],
475 + )
476 +
477 + # Create impl
478 + impl = impl_class(
479 + num_heads=mla_dims["num_q_heads"],
480 + head_size=mla_dims["head_dim"],
481 + scale=scale,
482 + num_kv_heads=mla_dims["num_kv_heads"],
483 + alibi_slopes=None,
484 + sliding_window=None,
485 + kv_cache_dtype="auto",
486 + logits_soft_cap=None,
487 + attn_type="decoder",
488 + kv_sharing_target_layer_name=None,
489 + q_lora_rank=None,
490 + kv_lora_rank=mla_dims["kv_lora_rank"],
491 + qk_nope_head_dim=mla_dims["qk_nope_head_dim"],
492 + qk_rope_head_dim=mla_dims["qk_rope_head_dim"],
493 + qk_head_dim=mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"],
494 + v_head_dim=mla_dims["v_head_dim"],
495 + kv_b_proj=mock_kv_b_proj,
496 + )
497 +
498 + # Initialize DCP attributes
499 + if not hasattr(impl, "dcp_world_size") or impl.dcp_world_size in (None, -1):
500 + impl.dcp_world_size = 1
501 + impl.dcp_rank = 0
502 +
503 + # Create KV cache spec for MockLayer
504 + from vllm.v1.kv_cache_interface import FullAttentionSpec
505 +
506 + kv_cache_spec = FullAttentionSpec(
507 + block_size=backend_cfg["block_size"] or vllm_config.cache_config.block_size,
508 + num_kv_heads=1, # MLA uses 1 KV head
509 + head_size=576, # MLA head dim
510 + dtype=torch.bfloat16,
511 + )
512 +
513 + # Create mock layer
514 + layer = MockLayer(device, impl=impl, kv_cache_spec=kv_cache_spec)
515 +
516 + # Create builder instance if needed
517 + builder_instance = None
518 + if backend_cfg["builder_class"]:
519 + builder_class = getattr(backend_module, backend_cfg["builder_class"])
520 +
521 + # Populate static_forward_context so builder can find the layer
522 + # MockLayer inherits from AttentionLayerBase, so isinstance checks pass
523 + vllm_config.compilation_config.static_forward_context = {"placeholder": layer}
524 +
525 + builder_instance = builder_class(
526 + kv_cache_spec=kv_cache_spec,
527 + layer_names=["placeholder"],
528 + vllm_config=vllm_config,
529 + device=device,
530 + )
531 +
532 + return impl, layer, builder_instance
533 +
534 +
535 +# ============================================================================
536 +# Config Helpers
537 +# ============================================================================
538 +
539 +
540 +def _extract_mla_dims_from_config(config) -> dict | None:
541 + """
542 + Extract MLA dimensions from BenchmarkConfig if all required fields are present.
543 +
544 + Args:
545 + config: BenchmarkConfig instance
546 +
547 + Returns:
548 + Dict with MLA dimensions if all fields are provided, None otherwise
549 + """
550 + # Check if all MLA-specific fields are provided
551 + if all(
552 + [
553 + config.kv_lora_rank is not None,
554 + config.qk_nope_head_dim is not None,
555 + config.qk_rope_head_dim is not None,
556 + config.v_head_dim is not None,
557 + ]
558 + ):
559 + return {
560 + "kv_lora_rank": config.kv_lora_rank,
561 + "qk_nope_head_dim": config.qk_nope_head_dim,
562 + "qk_rope_head_dim": config.qk_rope_head_dim,
563 + "v_head_dim": config.v_head_dim,
564 + "num_q_heads": config.num_q_heads,
565 + "num_kv_heads": config.num_kv_heads,
566 + "head_dim": config.head_dim,
567 + }
568 + # Fallback: if MLA fields not fully specified, try to construct from basic fields
569 + elif config.head_dim == 576:
570 + # This looks like a DeepSeek MLA config, use standard dimensions with custom
571 + # head count
572 + return {
573 + "kv_lora_rank": 512,
574 + "qk_nope_head_dim": 128,
575 + "qk_rope_head_dim": 64,
576 + "v_head_dim": 128,
577 + "num_q_heads": config.num_q_heads,
578 + "num_kv_heads": config.num_kv_heads,
579 + "head_dim": config.head_dim,
580 + }
581 + return None
582 +
583 +
584 +# ============================================================================
585 +# Benchmark Execution
586 +# ============================================================================
587 +
588 +
589 +def _run_single_benchmark(
590 + config,
591 + impl,
592 + layer,
593 + builder_instance,
594 + backend_cfg: dict,
595 + mla_dims: dict,
596 + device: torch.device,
597 +) -> BenchmarkResult:
598 + """
599 + Run a single benchmark iteration.
600 +
601 + Args:
602 + config: BenchmarkConfig instance
603 + impl: Backend implementation instance
604 + layer: MockLayer instance
605 + builder_instance: Metadata builder instance
606 + backend_cfg: Backend configuration dict
607 + mla_dims: MLA dimension configuration
608 + device: Target device
609 +
610 + Returns:
611 + BenchmarkResult with timing statistics
612 + """
613 + # Parse batch spec
614 + requests = parse_batch_spec(config.batch_spec)
615 + q_lens = [r.q_len for r in requests]
616 + total_q = sum(q_lens)
617 +
618 + # Determine block size
619 + block_size = backend_cfg["block_size"] or config.block_size
620 +
621 + # Build metadata
622 + metadata, num_blocks = _build_attention_metadata(
623 + requests, block_size, device, builder_instance
624 + )
625 +
626 + # Create KV cache
627 + kv_cache = torch.zeros(
628 + num_blocks,
629 + block_size,
630 + mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"],
631 + device=device,
632 + dtype=torch.bfloat16,
633 + )
634 +
635 + # Create input tensors for both decode and prefill modes
636 + decode_inputs, prefill_inputs = _create_input_tensors(
637 + total_q,
638 + mla_dims,
639 + backend_cfg["query_format"],
640 + device,
641 + torch.bfloat16,
642 + )
643 +
644 + # Determine which forward method to use based on metadata
645 + if metadata.decode is not None:
646 + forward_fn = lambda: impl._forward_decode(
647 + decode_inputs, kv_cache, metadata, layer
648 + )
649 + elif metadata.prefill is not None:
650 + forward_fn = lambda: impl._forward_prefill(
651 + prefill_inputs["q"],
652 + prefill_inputs["k_c_normed"],
653 + prefill_inputs["k_pe"],
654 + kv_cache,
655 + metadata,
656 + prefill_inputs["k_scale"],
657 + prefill_inputs["output"],
658 + )
659 + else:
660 + raise RuntimeError("Metadata has neither decode nor prefill metadata")
661 +
662 + # Warmup
663 + for _ in range(config.warmup_iters):
664 + forward_fn()
665 + torch.cuda.synchronize()
666 +
667 + # Benchmark
668 + times = []
669 + for _ in range(config.repeats):
670 + start = torch.cuda.Event(enable_timing=True)
671 + end = torch.cuda.Event(enable_timing=True)
672 +
673 + start.record()
674 + for _ in range(config.num_layers):
675 + forward_fn()
676 + end.record()
677 +
678 + torch.cuda.synchronize()
679 + elapsed_ms = start.elapsed_time(end)
680 + times.append(elapsed_ms / 1000.0 / config.num_layers)
681 +
682 + mean_time = float(np.mean(times))
683 + return BenchmarkResult(
684 + config=config,
685 + mean_time=mean_time,
686 + std_time=float(np.std(times)),
687 + min_time=float(np.min(times)),
688 + max_time=float(np.max(times)),
689 + throughput_tokens_per_sec=total_q / mean_time if mean_time > 0 else 0,
690 + )
691 +
692 +
693 +def _run_mla_benchmark_batched(
694 + backend: str,
695 + configs_with_params: list[tuple], # [(config, threshold, num_splits), ...]
696 +) -> list[BenchmarkResult]:
697 + """
698 + Unified batched MLA benchmark runner for all backends.
699 +
700 + Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla
701 +
702 + This function reuses backend initialization across multiple benchmarks
703 + to avoid setup/teardown overhead.
704 +
705 + Args:
706 + backend: Backend name
707 + configs_with_params: List of (config, threshold, num_splits) tuples
708 + - threshold: reorder_batch_threshold (FlashAttn/FlashMLA only)
709 + - num_splits: num_kv_splits (CUTLASS only)
710 +
711 + Returns:
712 + List of BenchmarkResult objects
713 + """
714 + if not configs_with_params:
715 + return []
716 +
717 + backend_cfg = _get_backend_config(backend)
718 + device = torch.device(configs_with_params[0][0].device)
719 + torch.cuda.set_device(device)
720 +
721 + # Determine block size
722 + config_block_size = configs_with_params[0][0].block_size
723 + block_size = backend_cfg["block_size"] or config_block_size
724 +
725 + # Extract MLA dimensions from the first config
726 + first_config = configs_with_params[0][0]
727 + mla_dims = _extract_mla_dims_from_config(first_config)
728 +
729 + # If config didn't provide MLA dims, fall back to default model
730 + if mla_dims is None:
731 + mla_dims = setup_mla_dims("deepseek-v3")
732 +
733 + # Create and set vLLM config for MLA (reused across all benchmarks)
734 + vllm_config = create_minimal_vllm_config(
735 + model_name="deepseek-v3", # Used only for model path
736 + block_size=block_size,
737 + mla_dims=mla_dims, # Use custom dims from config or default
738 + )
739 +
740 + results = []
741 +
742 + with set_current_vllm_config(vllm_config):
743 + # Create backend impl, layer, and builder (reused across benchmarks)
744 + impl, layer, builder_instance = _create_backend_impl(
745 + backend_cfg, mla_dims, vllm_config, device
746 + )
747 +
748 + # Run each benchmark with the shared impl
749 + for config, threshold, num_splits in configs_with_params:
750 + # Set threshold for this benchmark (FlashAttn/FlashMLA only)
751 + original_threshold = None
752 + if threshold is not None and builder_instance:
753 + original_threshold = builder_instance.reorder_batch_threshold
754 + builder_instance.reorder_batch_threshold = threshold
755 +
756 + # Set num_splits for CUTLASS
757 + original_num_splits = None
758 + if num_splits is not None and hasattr(impl, "_num_kv_splits"):
759 + original_num_splits = impl._num_kv_splits
760 + impl._num_kv_splits = num_splits
761 +
762 + try:
763 + result = _run_single_benchmark(
764 + config,
765 + impl,
766 + layer,
767 + builder_instance,
768 + backend_cfg,
769 + mla_dims,
770 + device,
771 + )
772 + results.append(result)
773 +
774 + finally:
775 + # Restore original threshold
776 + if original_threshold is not None:
777 + builder_instance.reorder_batch_threshold = original_threshold
778 +
779 + # Restore original num_splits
780 + if original_num_splits is not None:
781 + impl._num_kv_splits = original_num_splits
782 +
783 + return results
784 +
785 +
786 +# ============================================================================
787 +# Public API
788 +# ============================================================================
789 +
790 +
791 +def run_mla_benchmark(
792 + backend: str,
793 + config,
794 + reorder_batch_threshold: int | None = None,
795 + num_kv_splits: int | None = None,
796 +) -> BenchmarkResult | list[BenchmarkResult]:
797 + """
798 + Unified MLA benchmark runner for all backends.
799 +
800 + Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla
801 +
802 + Always uses batched execution internally for optimal performance.
803 +
804 + Args:
805 + backend: Backend name (flashattn_mla, flashmla, flashinfer_mla, cutlass_mla)
806 + config: BenchmarkConfig or list of (BenchmarkConfig, param) tuples
807 + reorder_batch_threshold: Threshold override for FlashAttn/FlashMLA
808 + (single config mode only)
809 + num_kv_splits: Number of KV splits for CUTLASS (single config mode only)
810 +
811 + Returns:
812 + BenchmarkResult (single mode) or list of BenchmarkResult (batched mode)
813 + """
814 + # Normalize to batched mode: (config, threshold, num_splits)
815 + if isinstance(config, list):
816 + # Already in batched format
817 + if len(config) > 0 and isinstance(config[0], tuple):
818 + # Format: [(cfg, param), ...] where param is threshold or num_splits
819 + if backend in ("flashattn_mla", "flashmla"):
820 + configs_with_params = [(cfg, param, None) for cfg, param in config]
821 + else: # cutlass_mla or flashinfer_mla
822 + configs_with_params = [(cfg, None, param) for cfg, param in config]
823 + else:
824 + # Format: [cfg, ...] - just configs
825 + configs_with_params = [(cfg, None, None) for cfg in config]
826 + return_single = False
827 + else:
828 + # Single config: convert to batched format
829 + configs_with_params = [(config, reorder_batch_threshold, num_kv_splits)]
830 + return_single = True
831 +
832 + # Use unified batched execution
833 + results = _run_mla_benchmark_batched(backend, configs_with_params)
834 +
835 + # Return single result or list based on input
836 + return results[0] if return_single else results
benchmarks/attention_benchmarks/runner.py new
+481
@@ -0,0 +1,481 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""
5 +Standard attention benchmark runner - shared utilities for non-MLA benchmarks.
6 +
7 +This module provides helpers for running standard attention backends
8 +(FlashAttention, Triton, FlashInfer) with real vLLM integration.
9 +"""
10 +
11 +import types
12 +
13 +import numpy as np
14 +import torch
15 +from batch_spec import parse_batch_spec, reorder_for_flashinfer
16 +from common import BenchmarkConfig, BenchmarkResult, MockLayer, get_attention_scale
17 +
18 +from vllm.config import (
19 + CacheConfig,
20 + CompilationConfig,
21 + DeviceConfig,
22 + LoadConfig,
23 + ModelConfig,
24 + ParallelConfig,
25 + SchedulerConfig,
26 + VllmConfig,
27 +)
28 +from vllm.v1.attention.backends.utils import CommonAttentionMetadata
29 +from vllm.v1.kv_cache_interface import FullAttentionSpec
30 +
31 +# ============================================================================
32 +# Backend Configuration
33 +# ============================================================================
34 +
35 +
36 +_BACKEND_CONFIG = {
37 + "flash": {
38 + "module": "vllm.v1.attention.backends.flash_attn",
39 + "backend_class": "FlashAttentionBackend",
40 + "dtype": torch.float16,
41 + "cache_layout": "standard",
42 + # ^ [2, num_blocks, block_size, num_kv_heads, head_dim]
43 + },
44 + "triton": {
45 + "module": "vllm.v1.attention.backends.triton_attn",
46 + "backend_class": "TritonAttentionBackend",
47 + "dtype": torch.float32,
48 + "cache_layout": "standard",
49 + },
50 + "flashinfer": {
51 + "module": "vllm.v1.attention.backends.flashinfer",
52 + "backend_class": "FlashInferBackend",
53 + "dtype": torch.float16,
54 + "cache_layout": "flashinfer",
55 + # ^ [num_blocks, 2, block_size, num_kv_heads, head_dim]
56 + },
57 +}
58 +
59 +
60 +def _get_backend_config(backend: str) -> dict:
61 + if backend not in _BACKEND_CONFIG:
62 + raise ValueError(
63 + f"Unknown backend: {backend}. "
64 + f"Available: {', '.join(_BACKEND_CONFIG.keys())}"
65 + )
66 + return _BACKEND_CONFIG[backend]
67 +
68 +
69 +# ============================================================================
70 +# Metadata Building Helpers
71 +# ============================================================================
72 +
73 +
74 +def _build_common_attn_metadata(
75 + q_lens: list[int],
76 + kv_lens: list[int],
77 + block_size: int,
78 + device: torch.device,
79 +) -> CommonAttentionMetadata:
80 + """Build CommonAttentionMetadata from query/kv lengths."""
81 + batch_size = len(q_lens)
82 + total_tokens = sum(q_lens)
83 +
84 + query_start_loc = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)
85 + query_start_loc[1:] = torch.tensor(q_lens, dtype=torch.int32, device=device).cumsum(
86 + 0
87 + )
88 + query_start_loc_cpu = query_start_loc.cpu()
89 +
90 + 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)
96 +
97 + max_blocks = (max(kv_lens) + block_size - 1) // block_size
98 + num_blocks = batch_size * max_blocks
99 + block_table_tensor = torch.arange(
100 + num_blocks, dtype=torch.int32, device=device
101 + ).view(batch_size, max_blocks)
102 + slot_mapping = torch.arange(total_tokens, dtype=torch.int64, device=device)
103 +
104 + max_query_len = max(q_lens)
105 +
106 + return CommonAttentionMetadata(
107 + query_start_loc=query_start_loc,
108 + query_start_loc_cpu=query_start_loc_cpu,
109 + seq_lens=seq_lens,
110 + seq_lens_cpu=seq_lens_cpu,
111 + num_computed_tokens_cpu=num_computed_tokens_cpu,
112 + num_reqs=batch_size,
113 + num_actual_tokens=total_tokens,
114 + max_query_len=max_query_len,
115 + max_seq_len=max_seq_len,
116 + block_table_tensor=block_table_tensor,
117 + slot_mapping=slot_mapping,
118 + causal=True,
119 + )
120 +
121 +
122 +def _create_vllm_config(
123 + config: BenchmarkConfig,
124 + dtype: torch.dtype,
125 + max_num_blocks: int,
126 +) -> VllmConfig:
127 + """Create a VllmConfig for benchmarking with mock model methods."""
128 + model_config = ModelConfig(
129 + model="meta-llama/Meta-Llama-3-8B",
130 + tokenizer="meta-llama/Meta-Llama-3-8B",
131 + trust_remote_code=False,
132 + dtype=dtype,
133 + seed=0,
134 + max_model_len=1024,
135 + )
136 +
137 + cache_config = CacheConfig(
138 + block_size=config.block_size,
139 + cache_dtype="auto",
140 + swap_space=0,
141 + )
142 + cache_config.num_gpu_blocks = max_num_blocks
143 + cache_config.num_cpu_blocks = 0
144 +
145 + parallel_config = ParallelConfig(tensor_parallel_size=1)
146 + scheduler_config = SchedulerConfig(
147 + max_num_seqs=256,
148 + max_num_batched_tokens=8192,
149 + max_model_len=8192,
150 + is_encoder_decoder=False,
151 + enable_chunked_prefill=True,
152 + )
153 + device_config = DeviceConfig()
154 + load_config = LoadConfig()
155 + compilation_config = CompilationConfig()
156 +
157 + # Add mock methods for benchmark config values
158 + model_config.get_num_layers = types.MethodType(
159 + lambda self: config.num_layers, model_config
160 + )
161 + model_config.get_sliding_window_for_layer = types.MethodType(
162 + lambda self, i: None, model_config
163 + )
164 + model_config.get_logits_soft_cap_for_layer = types.MethodType(
165 + lambda self, i: 0.0, model_config
166 + )
167 + model_config.get_sm_scale_for_layer = types.MethodType(
168 + lambda self, i: 1.0 / config.head_dim**0.5, model_config
169 + )
170 + model_config.get_num_attention_heads = types.MethodType(
171 + lambda self, parallel_config=None: config.num_q_heads, model_config
172 + )
173 + model_config.get_num_kv_heads = types.MethodType(
174 + lambda self, parallel_config=None: config.num_kv_heads, model_config
175 + )
176 + model_config.get_head_size = types.MethodType(
177 + lambda self: config.head_dim, model_config
178 + )
179 + model_config.get_sliding_window = types.MethodType(lambda self: None, model_config)
180 +
181 + return VllmConfig(
182 + model_config=model_config,
183 + cache_config=cache_config,
184 + parallel_config=parallel_config,
185 + scheduler_config=scheduler_config,
186 + device_config=device_config,
187 + load_config=load_config,
188 + compilation_config=compilation_config,
189 + )
190 +
191 +
192 +# ============================================================================
193 +# Backend Initialization
194 +# ============================================================================
195 +
196 +
197 +def _create_backend_impl(
198 + backend_cfg: dict,
199 + config: BenchmarkConfig,
200 + device: torch.device,
201 +):
202 + """Create backend implementation instance."""
203 + import importlib
204 +
205 + backend_module = importlib.import_module(backend_cfg["module"])
206 + backend_class = getattr(backend_module, backend_cfg["backend_class"])
207 +
208 + scale = get_attention_scale(config.head_dim)
209 + dtype = backend_cfg["dtype"]
210 +
211 + impl = backend_class.get_impl_cls()(
212 + num_heads=config.num_q_heads,
213 + head_size=config.head_dim,
214 + scale=scale,
215 + num_kv_heads=config.num_kv_heads,
216 + alibi_slopes=None,
217 + sliding_window=None,
218 + kv_cache_dtype="auto",
219 + )
220 +
221 + kv_cache_spec = FullAttentionSpec(
222 + block_size=config.block_size,
223 + num_kv_heads=config.num_kv_heads,
224 + head_size=config.head_dim,
225 + dtype=dtype,
226 + )
227 +
228 + layer = MockLayer(device, kv_cache_spec=kv_cache_spec)
229 +
230 + return backend_class, impl, layer, dtype
231 +
232 +
233 +def _create_metadata_builder(
234 + backend_class,
235 + kv_cache_spec: FullAttentionSpec,
236 + vllm_config: VllmConfig,
237 + device: torch.device,
238 +):
239 + """Create metadata builder instance."""
240 + return backend_class.get_builder_cls()(
241 + kv_cache_spec=kv_cache_spec,
242 + layer_names=["layer_0"],
243 + vllm_config=vllm_config,
244 + device=device,
245 + )
246 +
247 +
248 +# ============================================================================
249 +# Tensor Creation Helpers
250 +# ============================================================================
251 +
252 +
253 +def _create_input_tensors(
254 + config: BenchmarkConfig,
255 + total_q: int,
256 + device: torch.device,
257 + dtype: torch.dtype,
258 +) -> tuple:
259 + """Create Q, K, V input tensors for all layers."""
260 + q_list = [
261 + torch.randn(
262 + total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
263 + )
264 + for _ in range(config.num_layers)
265 + ]
266 + k_list = [
267 + torch.randn(
268 + total_q, config.num_kv_heads, config.head_dim, device=device, dtype=dtype
269 + )
270 + for _ in range(config.num_layers)
271 + ]
272 + v_list = [
273 + torch.randn(
274 + total_q, config.num_kv_heads, config.head_dim, device=device, dtype=dtype
275 + )
276 + for _ in range(config.num_layers)
277 + ]
278 + return q_list, k_list, v_list
279 +
280 +
281 +def _create_kv_cache(
282 + config: BenchmarkConfig,
283 + max_num_blocks: int,
284 + cache_layout: str,
285 + device: torch.device,
286 + dtype: torch.dtype,
287 +) -> 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 + ]
317 + return cache_list
318 +
319 +
320 +# ============================================================================
321 +# Benchmark Execution
322 +# ============================================================================
323 +
324 +
325 +def _run_single_benchmark(
326 + config: BenchmarkConfig,
327 + impl,
328 + layer,
329 + q_list: list,
330 + k_list: list,
331 + v_list: list,
332 + cache_list: list,
333 + attn_metadata,
334 + device: torch.device,
335 + dtype: torch.dtype,
336 +) -> tuple:
337 + """Run single benchmark iteration with warmup and timing loop."""
338 + total_q = q_list[0].shape[0]
339 + out = torch.empty(
340 + total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
341 + )
342 +
343 + # Warmup
344 + for _ in range(config.warmup_iters):
345 + for i in range(config.num_layers):
346 + impl.forward(
347 + layer,
348 + q_list[i],
349 + k_list[i],
350 + v_list[i],
351 + cache_list[i],
352 + attn_metadata,
353 + output=out,
354 + )
355 + torch.cuda.synchronize()
356 +
357 + # Benchmark
358 + times = []
359 + for _ in range(config.repeats):
360 + start = torch.cuda.Event(enable_timing=True)
361 + end = torch.cuda.Event(enable_timing=True)
362 +
363 + start.record()
364 + for i in range(config.num_layers):
365 + impl.forward(
366 + layer,
367 + q_list[i],
368 + k_list[i],
369 + v_list[i],
370 + cache_list[i],
371 + attn_metadata,
372 + output=out,
373 + )
374 + end.record()
375 +
376 + torch.cuda.synchronize()
377 + elapsed_ms = start.elapsed_time(end)
378 + times.append(elapsed_ms / 1000.0 / config.num_layers) # seconds per layer
379 +
380 + mem_stats = {}
381 + if config.profile_memory:
382 + mem_stats = {
383 + "allocated_mb": torch.cuda.memory_allocated(device) / 1024**2,
384 + "reserved_mb": torch.cuda.memory_reserved(device) / 1024**2,
385 + }
386 +
387 + return times, mem_stats
388 +
389 +
390 +# ============================================================================
391 +# Public API
392 +# ============================================================================
393 +
394 +
395 +def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
396 + """
397 + Run standard attention benchmark with real kernels.
398 +
399 + Supports: flash, triton, flashinfer
400 +
401 + Args:
402 + config: Benchmark configuration
403 +
404 + Returns:
405 + BenchmarkResult with timing and memory statistics
406 + """
407 + device = torch.device(config.device)
408 + torch.cuda.set_device(device)
409 +
410 + backend_cfg = _get_backend_config(config.backend)
411 +
412 + requests = parse_batch_spec(config.batch_spec)
413 +
414 + if config.backend == "flashinfer":
415 + requests = reorder_for_flashinfer(requests)
416 +
417 + q_lens = [r.q_len for r in requests]
418 + kv_lens = [r.kv_len for r in requests]
419 + total_q = sum(q_lens)
420 + max_kv = max(kv_lens)
421 +
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 + )
427 +
428 + common_metadata = _build_common_attn_metadata(
429 + q_lens, kv_lens, config.block_size, device
430 + )
431 +
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 + )
438 +
439 + vllm_config = _create_vllm_config(config, dtype, max_num_blocks)
440 +
441 + builder = _create_metadata_builder(
442 + backend_class, kv_cache_spec, vllm_config, device
443 + )
444 +
445 + attn_metadata = builder.build(
446 + common_prefix_len=0,
447 + common_attn_metadata=common_metadata,
448 + )
449 +
450 + q_list, k_list, v_list = _create_input_tensors(config, total_q, device, dtype)
451 +
452 + cache_list = _create_kv_cache(
453 + config, max_num_blocks, backend_cfg["cache_layout"], device, dtype
454 + )
455 +
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 + )
468 +
469 + mean_time = np.mean(times)
470 + throughput = total_q / mean_time if mean_time > 0 else 0
471 +
472 + return BenchmarkResult(
473 + config=config,
474 + mean_time=mean_time,
475 + std_time=np.std(times),
476 + min_time=np.min(times),
477 + max_time=np.max(times),
478 + throughput_tokens_per_sec=throughput,
479 + memory_allocated_mb=mem_stats.get("allocated_mb"),
480 + memory_reserved_mb=mem_stats.get("reserved_mb"),
481 + )