| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import dataclasses |
| 5 | from collections.abc import Callable |
| 6 | from typing import Literal |
| 7 | |
| 8 | import numpy as np |
| 9 | import torch |
| 10 | |
| 11 | from vllm.sampling_params import SamplingParams |
| 12 | from vllm.triton_utils import triton |
| 13 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 14 | from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState |
| 15 | from vllm.v1.worker.gpu.states import RequestState |
| 16 | |
| 17 | START_TOKEN_ID = 90 |
| 18 | END_TOKEN_ID = 91 |
| 19 | TOKEN_ID = 10 |
| 20 | VOCAB_SIZE = 128 |
| 21 | QUANTILES = [0.5, 0.2, 0.8] |
| 22 | |
| 23 | |
| 24 | class ReasoningConfig: |
| 25 | reasoning_start_token_ids = [START_TOKEN_ID] |
| 26 | reasoning_end_token_ids = [END_TOKEN_ID] |
| 27 | natural_reasoning_end_token_ids = [END_TOKEN_ID] |
| 28 | |
| 29 | |
| 30 | @dataclasses.dataclass |
| 31 | class BenchmarkCase: |
| 32 | req_states: RequestState |
| 33 | state: ThinkingBudgetState |
| 34 | run: Callable[[], None] |
| 35 | |
| 36 | |
| 37 | def create_case( |
| 38 | history_len: int, |
| 39 | device: torch.device, |
| 40 | extra_tokens: int = 0, |
| 41 | history_type: Literal["reasoning", "prefill"] = "reasoning", |
| 42 | batch_size: int = 1, |
| 43 | budget_type: Literal["active", "forced", "mixed"] = "active", |
| 44 | ) -> BenchmarkCase: |
| 45 | assert history_len >= 2 |
| 46 | if history_type == "prefill": |
| 47 | tokens = [*([TOKEN_ID] * (history_len - 1)), START_TOKEN_ID] |
| 48 | else: |
| 49 | tokens = [TOKEN_ID, START_TOKEN_ID, *([TOKEN_ID] * (history_len - 2))] |
| 50 | req_states = RequestState( |
| 51 | max_num_reqs=batch_size, |
| 52 | max_model_len=history_len + extra_tokens + 1, |
| 53 | max_num_batched_tokens=batch_size, |
| 54 | num_speculative_steps=1, |
| 55 | vocab_size=VOCAB_SIZE, |
| 56 | device=device, |
| 57 | ) |
| 58 | req_indices = [] |
| 59 | for req_num in range(batch_size): |
| 60 | req_id = f"benchmark-{req_num}" |
| 61 | req_states.add_request( |
| 62 | req_id=req_id, |
| 63 | prompt_len=1, |
| 64 | all_token_ids=tokens, |
| 65 | num_computed_tokens=history_len, |
| 66 | max_tokens=extra_tokens + 1, |
| 67 | ) |
| 68 | req_indices.append(req_states.req_id_to_index[req_id]) |
| 69 | req_states.apply_staged_writes() |
| 70 | |
| 71 | state = ThinkingBudgetState(req_states, ReasoningConfig()) |
| 72 | active_budget = history_len + extra_tokens + 1 |
| 73 | forced_budget = max(1, history_len - 2) |
| 74 | for req_num, req_idx in enumerate(req_indices): |
| 75 | if budget_type == "forced": |
| 76 | budget = forced_budget |
| 77 | elif budget_type == "mixed" and req_num % 2 == 1: |
| 78 | budget = None |
| 79 | else: |
| 80 | budget = active_budget |
| 81 | state.add_request( |
| 82 | req_idx, |
| 83 | SamplingParams(thinking_token_budget=budget), |
| 84 | ) |
| 85 | state.apply_staged_writes() |
| 86 | |
| 87 | idx_mapping = torch.tensor(req_indices, dtype=torch.int32, device=device) |
| 88 | logits = torch.zeros((batch_size, VOCAB_SIZE), device=device) |
| 89 | idx_mapping_np = idx_mapping.cpu().numpy() |
| 90 | input_ids = torch.full((batch_size,), TOKEN_ID, dtype=torch.int32, device=device) |
| 91 | local_pos = torch.zeros(batch_size, dtype=torch.int32, device=device) |
| 92 | |
| 93 | def run() -> None: |
| 94 | state.apply( |
| 95 | logits, |
| 96 | idx_mapping, |
| 97 | idx_mapping, |
| 98 | idx_mapping_np, |
| 99 | input_ids, |
| 100 | local_pos, |
| 101 | ) |
| 102 | |
| 103 | run() |
| 104 | torch.accelerator.synchronize(device) |
| 105 | active_req_indices = [ |
| 106 | req_idx |
| 107 | for req_num, req_idx in enumerate(req_indices) |
| 108 | if budget_type != "mixed" or req_num % 2 == 0 |
| 109 | ] |
| 110 | assert torch.all(state.cached_scan_pos[active_req_indices] == history_len).item() |
| 111 | if budget_type == "forced": |
| 112 | assert torch.all(logits[:, END_TOKEN_ID] == 1.0e9).item() |
| 113 | return BenchmarkCase(req_states, state, run) |
| 114 | |
| 115 | |
| 116 | def benchmark_cached( |
| 117 | case: BenchmarkCase, warmup_ms: int, rep_ms: int |
| 118 | ) -> tuple[float, float, float]: |
| 119 | median_ms, min_ms, max_ms = triton.testing.do_bench( |
| 120 | case.run, |
| 121 | warmup=warmup_ms, |
| 122 | rep=rep_ms, |
| 123 | quantiles=QUANTILES, |
| 124 | ) |
| 125 | return median_ms, min_ms, max_ms |
| 126 | |
| 127 | |
| 128 | def summarize(timings_ms: list[float]) -> tuple[float, float, float]: |
| 129 | median_ms, min_ms, max_ms = np.quantile(timings_ms, QUANTILES) |
| 130 | return float(median_ms), float(min_ms), float(max_ms) |
| 131 | |
| 132 | |
| 133 | def time_cuda_call(run: Callable[[], None]) -> float: |
| 134 | start = torch.Event(enable_timing=True) |
| 135 | end = torch.Event(enable_timing=True) |
| 136 | start.record() |
| 137 | run() |
| 138 | end.record() |
| 139 | end.synchronize() |
| 140 | return start.elapsed_time(end) |
| 141 | |
| 142 | |
| 143 | def benchmark_incremental_decode( |
| 144 | history_len: int, |
| 145 | device: torch.device, |
| 146 | warmup_steps: int, |
| 147 | iterations: int, |
| 148 | ) -> tuple[float, float, float]: |
| 149 | case = create_case(history_len, device, warmup_steps + iterations) |
| 150 | timings_ms: list[float] = [] |
| 151 | total_len = history_len |
| 152 | for step in range(warmup_steps + iterations): |
| 153 | case.req_states.all_token_ids.stage_write(0, total_len, [TOKEN_ID]) |
| 154 | total_len += 1 |
| 155 | case.req_states.total_len.stage_write_elem(0, total_len) |
| 156 | case.req_states.apply_staged_writes() |
| 157 | torch.accelerator.synchronize(device) |
| 158 | |
| 159 | elapsed_ms = time_cuda_call(case.run) |
| 160 | if step >= warmup_steps: |
| 161 | timings_ms.append(elapsed_ms) |
| 162 | |
| 163 | assert case.state.cached_scan_pos[0].item() == total_len |
| 164 | return summarize(timings_ms) |
| 165 | |
| 166 | |
| 167 | def benchmark_cold_scan( |
| 168 | history_len: int, |
| 169 | device: torch.device, |
| 170 | warmup_steps: int, |
| 171 | iterations: int, |
| 172 | history_type: Literal["reasoning", "prefill"], |
| 173 | ) -> tuple[float, float, float]: |
| 174 | case = create_case(history_len, device, history_type=history_type) |
| 175 | timings_ms: list[float] = [] |
| 176 | for step in range(warmup_steps + iterations): |
| 177 | case.state.cached_last_start.fill_(-1) |
| 178 | case.state.cached_last_end.fill_(-1) |
| 179 | case.state.cached_scan_pos.zero_() |
| 180 | torch.accelerator.synchronize(device) |
| 181 | |
| 182 | elapsed_ms = time_cuda_call(case.run) |
| 183 | if step >= warmup_steps: |
| 184 | timings_ms.append(elapsed_ms) |
| 185 | |
| 186 | return summarize(timings_ms) |
| 187 | |
| 188 | |
| 189 | def main() -> None: |
| 190 | parser = FlexibleArgumentParser( |
| 191 | description=( |
| 192 | "Benchmark thinking-budget scan, forced-end, and batched request " |
| 193 | "overhead across token history lengths, including worst-case cold " |
| 194 | "resume reconstruction." |
| 195 | ) |
| 196 | ) |
| 197 | parser.add_argument( |
| 198 | "--history-lengths", |
| 199 | type=int, |
| 200 | nargs="+", |
| 201 | default=[128, 16384, 32768], |
| 202 | ) |
| 203 | parser.add_argument( |
| 204 | "--modes", |
| 205 | nargs="+", |
| 206 | choices=[ |
| 207 | "cached", |
| 208 | "incremental-decode", |
| 209 | "cold-prefill", |
| 210 | "cold-resume-worst-case", |
| 211 | "forced-end", |
| 212 | "batched-budgeted", |
| 213 | "batched-mixed", |
| 214 | ], |
| 215 | default=[ |
| 216 | "cached", |
| 217 | "incremental-decode", |
| 218 | "cold-prefill", |
| 219 | "cold-resume-worst-case", |
| 220 | "forced-end", |
| 221 | "batched-budgeted", |
| 222 | "batched-mixed", |
| 223 | ], |
| 224 | ) |
| 225 | parser.add_argument( |
| 226 | "--batch-sizes", |
| 227 | type=int, |
| 228 | nargs="+", |
| 229 | default=[1, 8, 32], |
| 230 | help="Batch sizes used by batched-budgeted and batched-mixed modes.", |
| 231 | ) |
| 232 | parser.add_argument("--device", type=str, default="cuda:0") |
| 233 | parser.add_argument( |
| 234 | "--warmup-ms", |
| 235 | type=int, |
| 236 | default=100, |
| 237 | help="Warmup duration passed to triton.testing.do_bench.", |
| 238 | ) |
| 239 | parser.add_argument( |
| 240 | "--rep-ms", |
| 241 | type=int, |
| 242 | default=500, |
| 243 | help="Measurement duration passed to triton.testing.do_bench.", |
| 244 | ) |
| 245 | parser.add_argument( |
| 246 | "--warmup-steps", |
| 247 | type=int, |
| 248 | default=10, |
| 249 | help="Warmup steps for incremental-decode and cold modes.", |
| 250 | ) |
| 251 | parser.add_argument( |
| 252 | "--iterations", |
| 253 | type=int, |
| 254 | default=100, |
| 255 | help="Measured steps for incremental-decode and cold modes.", |
| 256 | ) |
| 257 | parser.add_argument( |
| 258 | "--max-slowdown", |
| 259 | type=float, |
| 260 | default=None, |
| 261 | help=( |
| 262 | "Fail if cached or incremental-decode median latency exceeds this " |
| 263 | "multiple of its shortest-history latency. Cold modes are excluded " |
| 264 | "because worst-case cold-resume cost scales with scanned history." |
| 265 | ), |
| 266 | ) |
| 267 | args = parser.parse_args() |
| 268 | |
| 269 | history_lengths = sorted(set(args.history_lengths)) |
| 270 | if not history_lengths or history_lengths[0] < 2: |
| 271 | parser.error("history lengths must be at least 2") |
| 272 | if args.warmup_steps < 0 or args.iterations <= 0: |
| 273 | parser.error("warmup steps must be non-negative and iterations positive") |
| 274 | batch_sizes = sorted(set(args.batch_sizes)) |
| 275 | if not batch_sizes or batch_sizes[0] <= 0: |
| 276 | parser.error("batch sizes must be positive") |
| 277 | |
| 278 | device = torch.device(args.device) |
| 279 | results: dict[str, dict[int, tuple[float, float, float]]] = {} |
| 280 | for mode in args.modes: |
| 281 | mode_batch_sizes = batch_sizes if mode.startswith("batched-") else [1] |
| 282 | if mode == "batched-mixed": |
| 283 | mode_batch_sizes = [size for size in mode_batch_sizes if size >= 2] |
| 284 | if not mode_batch_sizes: |
| 285 | parser.error("batched-mixed requires a batch size of at least 2") |
| 286 | for batch_size in mode_batch_sizes: |
| 287 | label = f"{mode}-b{batch_size}" if mode.startswith("batched-") else mode |
| 288 | mode_results: dict[int, tuple[float, float, float]] = {} |
| 289 | for history_len in history_lengths: |
| 290 | if mode == "cached": |
| 291 | case = create_case(history_len, device) |
| 292 | result = benchmark_cached(case, args.warmup_ms, args.rep_ms) |
| 293 | elif mode == "incremental-decode": |
| 294 | result = benchmark_incremental_decode( |
| 295 | history_len, |
| 296 | device, |
| 297 | args.warmup_steps, |
| 298 | args.iterations, |
| 299 | ) |
| 300 | elif mode == "cold-prefill": |
| 301 | result = benchmark_cold_scan( |
| 302 | history_len, |
| 303 | device, |
| 304 | args.warmup_steps, |
| 305 | args.iterations, |
| 306 | "prefill", |
| 307 | ) |
| 308 | elif mode == "cold-resume-worst-case": |
| 309 | result = benchmark_cold_scan( |
| 310 | history_len, |
| 311 | device, |
| 312 | args.warmup_steps, |
| 313 | args.iterations, |
| 314 | "reasoning", |
| 315 | ) |
| 316 | else: |
| 317 | budget_type = ( |
| 318 | "forced" |
| 319 | if mode == "forced-end" |
| 320 | else "mixed" |
| 321 | if mode == "batched-mixed" |
| 322 | else "active" |
| 323 | ) |
| 324 | case = create_case( |
| 325 | history_len, |
| 326 | device, |
| 327 | batch_size=batch_size, |
| 328 | budget_type=budget_type, |
| 329 | ) |
| 330 | result = benchmark_cached(case, args.warmup_ms, args.rep_ms) |
| 331 | mode_results[history_len] = result |
| 332 | results[label] = mode_results |
| 333 | |
| 334 | print("mode history_len median_us p20_us p80_us slowdown") |
| 335 | for label, mode_results in results.items(): |
| 336 | baseline_ms = mode_results[history_lengths[0]][0] |
| 337 | for history_len in history_lengths: |
| 338 | median_ms, min_ms, max_ms = mode_results[history_len] |
| 339 | slowdown = median_ms / baseline_ms |
| 340 | print( |
| 341 | f"{label:<26} {history_len:>11} {median_ms * 1000:>9.3f} " |
| 342 | f"{min_ms * 1000:>6.3f} {max_ms * 1000:>6.3f} " |
| 343 | f"{slowdown:>7.3f}x" |
| 344 | ) |
| 345 | |
| 346 | if args.max_slowdown is not None: |
| 347 | for mode in ("cached", "incremental-decode"): |
| 348 | if mode not in results: |
| 349 | continue |
| 350 | mode_results = results[mode] |
| 351 | baseline_ms = mode_results[history_lengths[0]][0] |
| 352 | worst_history_len = max( |
| 353 | mode_results, key=lambda length: mode_results[length][0] |
| 354 | ) |
| 355 | worst_slowdown = mode_results[worst_history_len][0] / baseline_ms |
| 356 | if worst_slowdown > args.max_slowdown: |
| 357 | raise SystemExit( |
| 358 | f"thinking-budget {mode} slowdown {worst_slowdown:.3f}x " |
| 359 | f"at history length {worst_history_len} exceeds " |
| 360 | f"{args.max_slowdown:.3f}x" |
| 361 | ) |
| 362 | |
| 363 | |
| 364 | if __name__ == "__main__": |
| 365 | main() |