| 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 | --output-csv results.csv |
| 112 | ``` |
| 113 | |
| 114 | ### MLA Backends |
| 115 | |
| 116 | ```bash |
| 117 | # Compare all MLA backends |
| 118 | python benchmark.py \ |
| 119 | --backends cutlass_mla flashinfer_mla flashattn_mla flashmla \ |
| 120 | --batch-specs "64q1s1k" "64q1s4k" \ |
| 121 | --output-csv mla_results.csv |
| 122 | ``` |
| 123 | |
| 124 | ### Parameter Sweeps |
| 125 | |
| 126 | Use `--sweep-param` and `--sweep-values` to run parameter sweeps from the CLI: |
| 127 | |
| 128 | #### CUTLASS MLA num-splits Optimization |
| 129 | |
| 130 | **Question:** What is the optimal `num_kv_splits` for CUTLASS MLA? |
| 131 | |
| 132 | ```bash |
| 133 | python benchmark.py \ |
| 134 | --backend cutlass_mla \ |
| 135 | --batch-specs "64q1s1k" "64q1s4k" "64q1s16k" \ |
| 136 | --sweep-param num_kv_splits \ |
| 137 | --sweep-values 1 2 4 8 16 \ |
| 138 | --output-json optimal_splits.json |
| 139 | ``` |
| 140 | |
| 141 | #### Reorder Batch Threshold Optimization |
| 142 | |
| 143 | **Question:** What's the optimal `reorder_batch_threshold` for speculative decoding? |
| 144 | |
| 145 | ```bash |
| 146 | python benchmark.py \ |
| 147 | --backend flashmla \ |
| 148 | --batch-specs "q4s1k" "q8s2k" \ |
| 149 | --sweep-param reorder_batch_threshold \ |
| 150 | --sweep-values 1 4 16 64 256 512 \ |
| 151 | --output-csv threshold_sweep.csv |
| 152 | ``` |
| 153 | |
| 154 | ### All Command-Line Options |
| 155 | |
| 156 | ```text |
| 157 | --config CONFIG # Path to YAML config file (overrides other args) |
| 158 | --backends BACKEND [BACKEND ...] # flash, triton, flashinfer, cutlass_mla, |
| 159 | # flashinfer_mla, flashattn_mla, flashmla |
| 160 | --backend BACKEND # Single backend (alternative to --backends) |
| 161 | --batch-specs SPEC [SPEC ...] # Batch specifications using extended grammar |
| 162 | |
| 163 | # Model configuration |
| 164 | --num-layers N # Number of layers |
| 165 | --head-dim N # Head dimension |
| 166 | --v-head-dim N # Value head dimension (defaults to --head-dim) |
| 167 | --num-q-heads N # Query heads |
| 168 | --num-kv-heads N # KV heads |
| 169 | --block-size N # Block size |
| 170 | --kv-lora-rank N # MLA KV LoRA rank |
| 171 | --qk-nope-head-dim N # MLA non-RoPE QK head dim |
| 172 | --qk-rope-head-dim N # MLA RoPE QK head dim |
| 173 | |
| 174 | # Benchmark settings |
| 175 | --device DEVICE # Device (default: cuda:0) |
| 176 | --warmup-ms N # Warmup window in ms for triton do_bench |
| 177 | --profile-memory # Profile memory usage |
| 178 | |
| 179 | # Parameter sweeps |
| 180 | --sweep-param PARAM # Parameter name to sweep (e.g., num_kv_splits, |
| 181 | # reorder_batch_threshold) |
| 182 | --sweep-values N [N ...] # Values to sweep for the parameter |
| 183 | |
| 184 | # Output |
| 185 | --output-csv FILE # Save to CSV |
| 186 | --output-json FILE # Save to JSON |
| 187 | ``` |
| 188 | |
| 189 | ## Hardware Requirements |
| 190 | |
| 191 | | Backend | Hardware | |
| 192 | | ------- | -------- | |
| 193 | | Flash/Triton/FlashInfer | Any CUDA GPU | |
| 194 | | CUTLASS MLA | Blackwell (SM100+) | |
| 195 | | FlashAttn MLA | Hopper (SM90+) | |
| 196 | | FlashMLA | Hopper (SM90+) | |
| 197 | | FlashInfer-MLA | Any CUDA GPU | |
| 198 | |
| 199 | ## Using MLA Runner Directly |
| 200 | |
| 201 | All MLA backends are available through `mla_runner.run_mla_benchmark()`: |
| 202 | |
| 203 | ```python |
| 204 | from mla_runner import run_mla_benchmark |
| 205 | from common import BenchmarkConfig |
| 206 | |
| 207 | config = BenchmarkConfig( |
| 208 | backend="cutlass_mla", |
| 209 | batch_spec="64q1s4k", |
| 210 | num_layers=10, |
| 211 | head_dim=576, |
| 212 | num_q_heads=128, |
| 213 | num_kv_heads=1, |
| 214 | block_size=128, |
| 215 | device="cuda:0", |
| 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. Save results** - Always use `--output-csv` or `--output-json` |
| 257 | |
| 258 | **2. Test incrementally** - Start with `--num-layers 1` |
| 259 | |
| 260 | **3. Extended grammar** - Leverage spec decode, chunked prefill patterns |
| 261 | |
| 262 | **4. Parameter sweeps** - Use `--sweep-param` and `--sweep-values` to find optimal values |