| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | """Benchmark comparing Triton vs PyTorch sort-based top-k/top-p implementations. |
| 5 | |
| 6 | Compares: |
| 7 | - apply_top_k_top_p_triton (Triton binary search) |
| 8 | - apply_top_k_top_p (PyTorch sort-based) |
| 9 | |
| 10 | Scenarios: |
| 11 | - top_k only (whole batch, partial batch) |
| 12 | - top_p only (whole batch, partial batch) |
| 13 | - mix of top_k and top_p |
| 14 | """ |
| 15 | |
| 16 | import argparse |
| 17 | import gc |
| 18 | from dataclasses import dataclass |
| 19 | |
| 20 | import torch |
| 21 | |
| 22 | from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch |
| 23 | from vllm.v1.sample.ops.topk_topp_triton import ( |
| 24 | apply_top_k_top_p_triton, |
| 25 | reset_buffer_cache, |
| 26 | ) |
| 27 | |
| 28 | |
| 29 | @dataclass |
| 30 | class BenchmarkConfig: |
| 31 | """Configuration for a benchmark run.""" |
| 32 | |
| 33 | name: str |
| 34 | batch_size: int |
| 35 | vocab_size: int |
| 36 | # k and p can be tensors or None |
| 37 | k_values: torch.Tensor | None # [batch_size] or None |
| 38 | p_values: torch.Tensor | None # [batch_size] or None |
| 39 | description: str |
| 40 | ops_pct: float = 0.0 # Percentage of ops relative to batch size |
| 41 | |
| 42 | |
| 43 | def calculate_ops_pct( |
| 44 | k_values: torch.Tensor | None, |
| 45 | p_values: torch.Tensor | None, |
| 46 | vocab_size: int, |
| 47 | batch_size: int, |
| 48 | ) -> float: |
| 49 | """Calculate the percentage of active top-k and top-p operations. |
| 50 | |
| 51 | Returns percentage where 100% = batch_size ops. |
| 52 | E.g., if all rows have both top-k and top-p active, returns 200%. |
| 53 | """ |
| 54 | active_ops = 0 |
| 55 | |
| 56 | if k_values is not None: |
| 57 | # Count rows where k < vocab_size (active top-k filtering) |
| 58 | active_ops += (k_values < vocab_size).sum().item() |
| 59 | |
| 60 | if p_values is not None: |
| 61 | # Count rows where p < 1.0 (active top-p filtering) |
| 62 | active_ops += (p_values < 1.0).sum().item() |
| 63 | |
| 64 | return (active_ops / batch_size) * 100 if batch_size > 0 else 0.0 |
| 65 | |
| 66 | |
| 67 | def create_logits( |
| 68 | batch_size: int, vocab_size: int, device: str = "cuda" |
| 69 | ) -> torch.Tensor: |
| 70 | """Create random logits mimicking a realistic LLM distribution. |
| 71 | |
| 72 | Uses a Zipf-like probability distribution (rank^-1.1) converted to logits |
| 73 | via log, then randomly permuted per row. This produces a peaked distribution |
| 74 | where a small number of tokens capture most probability mass, similar to |
| 75 | real model outputs. |
| 76 | """ |
| 77 | # Create Zipf-like probabilities: p(rank) ~ rank^(-alpha) |
| 78 | ranks = torch.arange(1, vocab_size + 1, dtype=torch.float32, device=device) |
| 79 | probs = ranks.pow(-1.1) |
| 80 | probs = probs / probs.sum() |
| 81 | |
| 82 | # Convert to logits (log-probabilities, unnormalized is fine) |
| 83 | base_logits = probs.log() |
| 84 | |
| 85 | # Broadcast to batch and randomly permute each row |
| 86 | logits = base_logits.unsqueeze(0).expand(batch_size, -1).clone() |
| 87 | for i in range(batch_size): |
| 88 | logits[i] = logits[i, torch.randperm(vocab_size, device=device)] |
| 89 | |
| 90 | return logits |
| 91 | |
| 92 | |
| 93 | def measure_memory() -> tuple[int, int]: |
| 94 | """Return (allocated, reserved) memory in bytes.""" |
| 95 | torch.accelerator.synchronize() |
| 96 | return ( |
| 97 | torch.accelerator.memory_allocated(), |
| 98 | torch.accelerator.max_memory_allocated(), |
| 99 | ) |
| 100 | |
| 101 | |
| 102 | def reset_memory_stats(): |
| 103 | """Reset peak memory statistics.""" |
| 104 | reset_buffer_cache() |
| 105 | torch.accelerator.reset_peak_memory_stats() |
| 106 | torch.accelerator.empty_cache() |
| 107 | gc.collect() |
| 108 | |
| 109 | |
| 110 | def benchmark_function( |
| 111 | func, |
| 112 | logits: torch.Tensor, |
| 113 | k: torch.Tensor | None, |
| 114 | p: torch.Tensor | None, |
| 115 | warmup_iters: int = 5, |
| 116 | benchmark_iters: int = 20, |
| 117 | ) -> tuple[float, int]: |
| 118 | """Benchmark a function and return (avg_time_ms, peak_memory_bytes). |
| 119 | |
| 120 | Returns average time in milliseconds and peak memory usage. |
| 121 | """ |
| 122 | # Warmup |
| 123 | for _ in range(warmup_iters): |
| 124 | logits_copy = logits.clone() |
| 125 | func(logits_copy, k, p) |
| 126 | torch.accelerator.synchronize() |
| 127 | |
| 128 | # Reset memory stats before benchmark |
| 129 | reset_memory_stats() |
| 130 | |
| 131 | # Benchmark |
| 132 | start_events = [ |
| 133 | torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters) |
| 134 | ] |
| 135 | end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)] |
| 136 | |
| 137 | for i in range(benchmark_iters): |
| 138 | logits_copy = logits.clone() |
| 139 | start_events[i].record() |
| 140 | func(logits_copy, k, p) |
| 141 | end_events[i].record() |
| 142 | |
| 143 | torch.accelerator.synchronize() |
| 144 | |
| 145 | # Calculate timing |
| 146 | times = [ |
| 147 | start_events[i].elapsed_time(end_events[i]) for i in range(benchmark_iters) |
| 148 | ] |
| 149 | avg_time = sum(times) / len(times) |
| 150 | |
| 151 | # Get peak memory |
| 152 | _, peak_memory = measure_memory() |
| 153 | |
| 154 | return avg_time, peak_memory |
| 155 | |
| 156 | |
| 157 | def create_benchmark_configs( |
| 158 | batch_sizes: list[int], |
| 159 | vocab_sizes: list[int], |
| 160 | device: str = "cuda", |
| 161 | ) -> list[BenchmarkConfig]: |
| 162 | """Create all benchmark configurations.""" |
| 163 | configs = [] |
| 164 | |
| 165 | for vocab_size in vocab_sizes: |
| 166 | for batch_size in batch_sizes: |
| 167 | # 1. Top-k only - whole batch (all rows have k < vocab_size) |
| 168 | k_all = torch.full((batch_size,), 50, dtype=torch.int32, device=device) |
| 169 | configs.append( |
| 170 | BenchmarkConfig( |
| 171 | name=f"topk_whole_b{batch_size}_v{vocab_size // 1000}k", |
| 172 | batch_size=batch_size, |
| 173 | vocab_size=vocab_size, |
| 174 | k_values=k_all, |
| 175 | p_values=None, |
| 176 | description=f"Top-k only (whole batch, k=50), " |
| 177 | f"batch={batch_size}, vocab={vocab_size}", |
| 178 | ops_pct=calculate_ops_pct(k_all, None, vocab_size, batch_size), |
| 179 | ) |
| 180 | ) |
| 181 | |
| 182 | # 2. Top-k only - partial batch (half have k=50, half have k=vocab_size) |
| 183 | k_partial = torch.full((batch_size,), 50, dtype=torch.int32, device=device) |
| 184 | k_partial[batch_size // 2 :] = vocab_size # No filtering for second half |
| 185 | configs.append( |
| 186 | BenchmarkConfig( |
| 187 | name=f"topk_partial_b{batch_size}_v{vocab_size // 1000}k", |
| 188 | batch_size=batch_size, |
| 189 | vocab_size=vocab_size, |
| 190 | k_values=k_partial, |
| 191 | p_values=None, |
| 192 | description=f"Top-k only (partial batch, 50% k=50, 50% k=vocab), " |
| 193 | f"batch={batch_size}, vocab={vocab_size}", |
| 194 | ops_pct=calculate_ops_pct(k_partial, None, vocab_size, batch_size), |
| 195 | ) |
| 196 | ) |
| 197 | |
| 198 | # 3. Top-p only - whole batch (all rows have p < 1.0) |
| 199 | p_all = torch.full((batch_size,), 0.9, dtype=torch.float32, device=device) |
| 200 | configs.append( |
| 201 | BenchmarkConfig( |
| 202 | name=f"topp_whole_b{batch_size}_v{vocab_size // 1000}k", |
| 203 | batch_size=batch_size, |
| 204 | vocab_size=vocab_size, |
| 205 | k_values=None, |
| 206 | p_values=p_all, |
| 207 | description=f"Top-p only (whole batch, p=0.9), " |
| 208 | f"batch={batch_size}, vocab={vocab_size}", |
| 209 | ops_pct=calculate_ops_pct(None, p_all, vocab_size, batch_size), |
| 210 | ) |
| 211 | ) |
| 212 | |
| 213 | # 4. Top-p only - partial batch (half have p=0.9, half have p=1.0) |
| 214 | p_partial = torch.full( |
| 215 | (batch_size,), 0.9, dtype=torch.float32, device=device |
| 216 | ) |
| 217 | p_partial[batch_size // 2 :] = 1.0 # No filtering for second half |
| 218 | configs.append( |
| 219 | BenchmarkConfig( |
| 220 | name=f"topp_partial_b{batch_size}_v{vocab_size // 1000}k", |
| 221 | batch_size=batch_size, |
| 222 | vocab_size=vocab_size, |
| 223 | k_values=None, |
| 224 | p_values=p_partial, |
| 225 | description=f"Top-p only (partial batch, 50% p=0.9, 50% p=1.0), " |
| 226 | f"batch={batch_size}, vocab={vocab_size}", |
| 227 | ops_pct=calculate_ops_pct(None, p_partial, vocab_size, batch_size), |
| 228 | ) |
| 229 | ) |
| 230 | |
| 231 | # 5. Mix of top-k and top-p (both applied to whole batch) |
| 232 | k_mix = torch.full((batch_size,), 100, dtype=torch.int32, device=device) |
| 233 | p_mix = torch.full((batch_size,), 0.9, dtype=torch.float32, device=device) |
| 234 | configs.append( |
| 235 | BenchmarkConfig( |
| 236 | name=f"topk_topp_whole_b{batch_size}_v{vocab_size // 1000}k", |
| 237 | batch_size=batch_size, |
| 238 | vocab_size=vocab_size, |
| 239 | k_values=k_mix, |
| 240 | p_values=p_mix, |
| 241 | description=f"Top-k + Top-p (whole batch, k=100, p=0.9), " |
| 242 | f"batch={batch_size}, vocab={vocab_size}", |
| 243 | ops_pct=calculate_ops_pct(k_mix, p_mix, vocab_size, batch_size), |
| 244 | ) |
| 245 | ) |
| 246 | |
| 247 | # 6. Mix with partial application (some rows k only, some p only, some both) |
| 248 | k_mixed = torch.full( |
| 249 | (batch_size,), vocab_size, dtype=torch.int32, device=device |
| 250 | ) |
| 251 | p_mixed = torch.full((batch_size,), 1.0, dtype=torch.float32, device=device) |
| 252 | # First third: k only |
| 253 | third = batch_size // 3 |
| 254 | k_mixed[:third] = 50 |
| 255 | # Second third: p only |
| 256 | p_mixed[third : 2 * third] = 0.5 |
| 257 | # Last third: both k and p |
| 258 | k_mixed[2 * third :] = 100 |
| 259 | p_mixed[2 * third :] = 0.9 |
| 260 | configs.append( |
| 261 | BenchmarkConfig( |
| 262 | name=f"mixed_partial_b{batch_size}_v{vocab_size // 1000}k", |
| 263 | batch_size=batch_size, |
| 264 | vocab_size=vocab_size, |
| 265 | k_values=k_mixed, |
| 266 | p_values=p_mixed, |
| 267 | description=f"Mixed partial (1/3 k=50, 1/3 p=0.9, 1/3 both), " |
| 268 | f"batch={batch_size}, vocab={vocab_size}", |
| 269 | ops_pct=calculate_ops_pct(k_mixed, p_mixed, vocab_size, batch_size), |
| 270 | ) |
| 271 | ) |
| 272 | |
| 273 | return configs |
| 274 | |
| 275 | |
| 276 | def format_memory(bytes_val: int) -> str: |
| 277 | """Format memory in human-readable form.""" |
| 278 | if bytes_val >= 1024**3: |
| 279 | return f"{bytes_val / (1024**3):.2f} GB" |
| 280 | elif bytes_val >= 1024**2: |
| 281 | return f"{bytes_val / (1024**2):.2f} MB" |
| 282 | elif bytes_val >= 1024: |
| 283 | return f"{bytes_val / 1024:.2f} KB" |
| 284 | return f"{bytes_val} B" |
| 285 | |
| 286 | |
| 287 | def run_benchmark( |
| 288 | configs: list[BenchmarkConfig], |
| 289 | warmup_iters: int = 5, |
| 290 | benchmark_iters: int = 20, |
| 291 | verbose: bool = True, |
| 292 | ): |
| 293 | """Run all benchmarks and print results.""" |
| 294 | results = [] |
| 295 | |
| 296 | print("=" * 100) |
| 297 | print("Top-k/Top-p Benchmark: Triton vs PyTorch Sort-based") |
| 298 | print("=" * 100) |
| 299 | print() |
| 300 | |
| 301 | for config in configs: |
| 302 | if verbose: |
| 303 | print(f"Running: {config.description}") |
| 304 | |
| 305 | # Create fresh logits for this config |
| 306 | logits = create_logits(config.batch_size, config.vocab_size) |
| 307 | |
| 308 | # Benchmark Triton |
| 309 | reset_memory_stats() |
| 310 | triton_time, triton_mem = benchmark_function( |
| 311 | apply_top_k_top_p_triton, |
| 312 | logits, |
| 313 | config.k_values, |
| 314 | config.p_values, |
| 315 | warmup_iters, |
| 316 | benchmark_iters, |
| 317 | ) |
| 318 | |
| 319 | # Benchmark PyTorch |
| 320 | reset_memory_stats() |
| 321 | pytorch_time, pytorch_mem = benchmark_function( |
| 322 | apply_top_k_top_p_pytorch, |
| 323 | logits, |
| 324 | config.k_values, |
| 325 | config.p_values, |
| 326 | warmup_iters, |
| 327 | benchmark_iters, |
| 328 | ) |
| 329 | |
| 330 | speedup = pytorch_time / triton_time if triton_time > 0 else float("inf") |
| 331 | mem_ratio = pytorch_mem / triton_mem if triton_mem > 0 else float("inf") |
| 332 | |
| 333 | result = { |
| 334 | "config": config, |
| 335 | "triton_time_ms": triton_time, |
| 336 | "pytorch_time_ms": pytorch_time, |
| 337 | "triton_mem": triton_mem, |
| 338 | "pytorch_mem": pytorch_mem, |
| 339 | "speedup": speedup, |
| 340 | "mem_ratio": mem_ratio, |
| 341 | } |
| 342 | results.append(result) |
| 343 | |
| 344 | if verbose: |
| 345 | print(f" Triton: {triton_time:.3f} ms, {format_memory(triton_mem)}") |
| 346 | print(f" PyTorch: {pytorch_time:.3f} ms, {format_memory(pytorch_mem)}") |
| 347 | print(f" Speedup: {speedup:.2f}x, Memory ratio: {mem_ratio:.2f}x") |
| 348 | print() |
| 349 | |
| 350 | # Clean up |
| 351 | del logits |
| 352 | reset_memory_stats() |
| 353 | |
| 354 | return results |
| 355 | |
| 356 | |
| 357 | def print_summary_table(results: list[dict]): |
| 358 | """Print a summary table of results.""" |
| 359 | print() |
| 360 | print("=" * 130) |
| 361 | print("SUMMARY TABLE") |
| 362 | print("=" * 130) |
| 363 | print() |
| 364 | |
| 365 | # Header |
| 366 | header = ( |
| 367 | f"{'Scenario':<40} {'Batch':>6} {'Vocab':>7} {'Ops%':>6} " |
| 368 | f"{'Triton (ms)':>12} {'PyTorch (ms)':>13} {'Speedup':>8} " |
| 369 | f"{'Tri Mem':>10} {'Pyt Mem':>10}" |
| 370 | ) |
| 371 | print(header) |
| 372 | print("-" * 130) |
| 373 | |
| 374 | # Group by scenario type |
| 375 | current_vocab = None |
| 376 | for result in results: |
| 377 | config = result["config"] |
| 378 | |
| 379 | # Add separator between vocab sizes |
| 380 | if current_vocab != config.vocab_size: |
| 381 | if current_vocab is not None: |
| 382 | print("-" * 130) |
| 383 | current_vocab = config.vocab_size |
| 384 | |
| 385 | scenario = config.name.split("_b")[0] # Extract scenario name |
| 386 | print( |
| 387 | f"{scenario:<40} {config.batch_size:>6} {config.vocab_size:>7} " |
| 388 | f"{config.ops_pct:>5.0f}% " |
| 389 | f"{result['triton_time_ms']:>12.3f} {result['pytorch_time_ms']:>13.3f} " |
| 390 | f"{result['speedup']:>7.2f}x " |
| 391 | f"{format_memory(result['triton_mem']):>10} " |
| 392 | f"{format_memory(result['pytorch_mem']):>10}" |
| 393 | ) |
| 394 | |
| 395 | print("=" * 130) |
| 396 | |
| 397 | |
| 398 | def main(): |
| 399 | parser = argparse.ArgumentParser( |
| 400 | description="Benchmark Triton vs PyTorch sort-based top-k/top-p implementations" |
| 401 | ) |
| 402 | parser.add_argument( |
| 403 | "--batch-sizes", |
| 404 | type=int, |
| 405 | nargs="+", |
| 406 | default=[1, 4, 16, 64, 128, 512, 1024, 2048], |
| 407 | help="Batch sizes to test (default: 1 4 16 64)", |
| 408 | ) |
| 409 | parser.add_argument( |
| 410 | "--vocab-sizes", |
| 411 | type=int, |
| 412 | nargs="+", |
| 413 | default=[32768, 131072], # 32k, 128k |
| 414 | help="Vocabulary sizes to test (default: 32768 131072)", |
| 415 | ) |
| 416 | parser.add_argument( |
| 417 | "--warmup-iters", |
| 418 | type=int, |
| 419 | default=5, |
| 420 | help="Number of warmup iterations (default: 5)", |
| 421 | ) |
| 422 | parser.add_argument( |
| 423 | "--benchmark-iters", |
| 424 | type=int, |
| 425 | default=20, |
| 426 | help="Number of benchmark iterations (default: 20)", |
| 427 | ) |
| 428 | parser.add_argument( |
| 429 | "--quiet", |
| 430 | action="store_true", |
| 431 | help="Only print summary table", |
| 432 | ) |
| 433 | |
| 434 | args = parser.parse_args() |
| 435 | |
| 436 | # Print configuration |
| 437 | print(f"Batch sizes: {args.batch_sizes}") |
| 438 | print(f"Vocab sizes: {args.vocab_sizes}") |
| 439 | print(f"Warmup iterations: {args.warmup_iters}") |
| 440 | print(f"Benchmark iterations: {args.benchmark_iters}") |
| 441 | print() |
| 442 | |
| 443 | # Check CUDA |
| 444 | if not torch.cuda.is_available(): |
| 445 | print("ERROR: CUDA is not available. This benchmark requires a GPU.") |
| 446 | return |
| 447 | |
| 448 | device_name = torch.cuda.get_device_name(0) |
| 449 | print(f"GPU: {device_name}") |
| 450 | print() |
| 451 | |
| 452 | # Create configs |
| 453 | configs = create_benchmark_configs( |
| 454 | args.batch_sizes, |
| 455 | args.vocab_sizes, |
| 456 | ) |
| 457 | |
| 458 | # Run benchmarks |
| 459 | results = run_benchmark( |
| 460 | configs, |
| 461 | warmup_iters=args.warmup_iters, |
| 462 | benchmark_iters=args.benchmark_iters, |
| 463 | verbose=not args.quiet, |
| 464 | ) |
| 465 | |
| 466 | # Print summary |
| 467 | print_summary_table(results) |
| 468 | |
| 469 | |
| 470 | if __name__ == "__main__": |
| 471 | main() |