| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Micro-benchmark for CPU sampling kernels. |
| 4 | |
| 5 | Compares fused Gumbel-max / greedy argmax against the baseline |
| 6 | (softmax → exponential → div → argmax) across vocab and batch sizes. |
| 7 | |
| 8 | Usage: |
| 9 | .venv/bin/python benchmarks/kernels/bench_cpu_sampling.py |
| 10 | .venv/bin/python benchmarks/kernels/bench_cpu_sampling.py --profile |
| 11 | .venv/bin/python benchmarks/kernels/bench_cpu_sampling.py --vocab 128256 --batch 16 |
| 12 | """ |
| 13 | |
| 14 | import argparse |
| 15 | import time |
| 16 | |
| 17 | import torch |
| 18 | import vllm._C # noqa: F401 |
| 19 | from torch.profiler import ProfilerActivity, profile, record_function |
| 20 | |
| 21 | |
| 22 | def baseline_random_sample(logits: torch.Tensor) -> torch.Tensor: |
| 23 | probs = logits.softmax(dim=-1, dtype=torch.float32) |
| 24 | q = torch.empty_like(probs) |
| 25 | q.exponential_() |
| 26 | return probs.div(q).argmax(dim=-1).view(-1) |
| 27 | |
| 28 | |
| 29 | def baseline_greedy_sample(logits: torch.Tensor) -> torch.Tensor: |
| 30 | return logits.argmax(dim=-1).view(-1) |
| 31 | |
| 32 | |
| 33 | def bench_latency(fn, args, n_warmup=20, n_iters=500): |
| 34 | for _ in range(n_warmup): |
| 35 | fn(*args) |
| 36 | |
| 37 | t0 = time.perf_counter() |
| 38 | for _ in range(n_iters): |
| 39 | fn(*args) |
| 40 | elapsed = time.perf_counter() - t0 |
| 41 | return elapsed / n_iters * 1e6 # µs |
| 42 | |
| 43 | |
| 44 | def run_profile(logits, seeds, n_iters=50): |
| 45 | """Run torch.profiler and print comparison tables.""" |
| 46 | # Profile baseline random sampling |
| 47 | with profile( |
| 48 | activities=[ProfilerActivity.CPU], |
| 49 | record_shapes=True, |
| 50 | ) as prof_base_rand: |
| 51 | for _ in range(n_iters): |
| 52 | with record_function("baseline_random"): |
| 53 | baseline_random_sample(logits) |
| 54 | |
| 55 | # Profile fused Gumbel-max |
| 56 | with profile( |
| 57 | activities=[ProfilerActivity.CPU], |
| 58 | record_shapes=True, |
| 59 | ) as prof_fused: |
| 60 | for _ in range(n_iters): |
| 61 | with record_function("fused_gumbel_argmax"): |
| 62 | torch.ops._C.fused_gumbel_argmax(logits, seeds) |
| 63 | |
| 64 | # Profile baseline greedy |
| 65 | with profile( |
| 66 | activities=[ProfilerActivity.CPU], |
| 67 | record_shapes=True, |
| 68 | ) as prof_base_grdy: |
| 69 | for _ in range(n_iters): |
| 70 | with record_function("baseline_greedy"): |
| 71 | baseline_greedy_sample(logits) |
| 72 | |
| 73 | # Profile custom greedy |
| 74 | with profile( |
| 75 | activities=[ProfilerActivity.CPU], |
| 76 | record_shapes=True, |
| 77 | ) as prof_cust_grdy: |
| 78 | for _ in range(n_iters): |
| 79 | with record_function("custom_greedy_argmax"): |
| 80 | torch.ops._C.greedy_argmax(logits) |
| 81 | |
| 82 | B, V = logits.shape |
| 83 | print(f"\n{'=' * 80}") |
| 84 | print(f"torch.profiler breakdown (batch={B}, vocab={V}, iters={n_iters})") |
| 85 | print(f"{'=' * 80}") |
| 86 | |
| 87 | print("\n--- Baseline Random (softmax → exp → div → argmax) ---") |
| 88 | print(prof_base_rand.key_averages().table(sort_by="cpu_time_total", row_limit=15)) |
| 89 | |
| 90 | print("\n--- Fused Gumbel-max (table lookup + add + argmax) ---") |
| 91 | print(prof_fused.key_averages().table(sort_by="cpu_time_total", row_limit=15)) |
| 92 | |
| 93 | print("\n--- Baseline Greedy (torch.argmax) ---") |
| 94 | print(prof_base_grdy.key_averages().table(sort_by="cpu_time_total", row_limit=15)) |
| 95 | |
| 96 | print("\n--- Custom Greedy (vec_op argmax) ---") |
| 97 | print(prof_cust_grdy.key_averages().table(sort_by="cpu_time_total", row_limit=15)) |
| 98 | |
| 99 | # Summary comparison |
| 100 | def avg_us(prof, label): |
| 101 | for e in prof.key_averages(): |
| 102 | if e.key == label: |
| 103 | return e.cpu_time_total / e.count |
| 104 | return 0.0 |
| 105 | |
| 106 | t_br = avg_us(prof_base_rand, "baseline_random") |
| 107 | t_fg = avg_us(prof_fused, "fused_gumbel_argmax") |
| 108 | t_bg = avg_us(prof_base_grdy, "baseline_greedy") |
| 109 | t_cg = avg_us(prof_cust_grdy, "custom_greedy_argmax") |
| 110 | |
| 111 | print(f"\n{'=' * 60}") |
| 112 | print(f" Summary (batch={B}, vocab={V})") |
| 113 | print(f"{'=' * 60}") |
| 114 | print(f" {'Kernel':<30} {'Avg (µs)':>10} {'Speedup':>10}") |
| 115 | print(f" {'-' * 50}") |
| 116 | print(f" {'baseline random':<30} {t_br:>10.1f} {'—':>10}") |
| 117 | print( |
| 118 | f" {'fused gumbel-max':<30} {t_fg:>10.1f} " |
| 119 | f"{t_br / t_fg if t_fg > 0 else 0:>9.2f}x" |
| 120 | ) |
| 121 | print(f" {'baseline greedy':<30} {t_bg:>10.1f} {'—':>10}") |
| 122 | print( |
| 123 | f" {'custom greedy':<30} {t_cg:>10.1f} {t_bg / t_cg if t_cg > 0 else 0:>9.2f}x" |
| 124 | ) |
| 125 | print(f"{'=' * 60}") |
| 126 | |
| 127 | |
| 128 | def main(): |
| 129 | parser = argparse.ArgumentParser(description="Benchmark CPU sampling kernels") |
| 130 | parser.add_argument( |
| 131 | "--vocab", |
| 132 | type=int, |
| 133 | nargs="+", |
| 134 | default=[32000, 49152, 128256], |
| 135 | help="Vocab sizes to benchmark", |
| 136 | ) |
| 137 | parser.add_argument( |
| 138 | "--batch", |
| 139 | type=int, |
| 140 | nargs="+", |
| 141 | default=[1, 4, 16], |
| 142 | help="Batch sizes to benchmark", |
| 143 | ) |
| 144 | parser.add_argument( |
| 145 | "--profile", |
| 146 | action="store_true", |
| 147 | help="Run torch.profiler and export chrome trace", |
| 148 | ) |
| 149 | parser.add_argument( |
| 150 | "--iters", type=int, default=500, help="Iterations per measurement" |
| 151 | ) |
| 152 | args = parser.parse_args() |
| 153 | |
| 154 | header = ( |
| 155 | f"{'batch':>5} {'vocab':>7} " |
| 156 | f"{'base_rand':>10} {'fused_rand':>10} {'rand_spdup':>10} " |
| 157 | f"{'base_grdy':>10} {'cust_grdy':>10} {'grdy_spdup':>10}" |
| 158 | ) |
| 159 | units = ( |
| 160 | f"{'':>5} {'':>7} " |
| 161 | f"{'(µs)':>10} {'(µs)':>10} {'':>10} " |
| 162 | f"{'(µs)':>10} {'(µs)':>10} {'':>10}" |
| 163 | ) |
| 164 | print("\n" + "=" * len(header)) |
| 165 | print("CPU Sampling Kernel Benchmark") |
| 166 | print("=" * len(header)) |
| 167 | print(header) |
| 168 | print(units) |
| 169 | print("-" * len(header)) |
| 170 | |
| 171 | for V in args.vocab: |
| 172 | for B in args.batch: |
| 173 | logits = torch.randn(B, V, dtype=torch.float32) |
| 174 | seeds = torch.arange(B, dtype=torch.long) |
| 175 | |
| 176 | t_base_rand = bench_latency( |
| 177 | baseline_random_sample, (logits,), n_iters=args.iters |
| 178 | ) |
| 179 | t_fused = bench_latency( |
| 180 | torch.ops._C.fused_gumbel_argmax, (logits, seeds), n_iters=args.iters |
| 181 | ) |
| 182 | |
| 183 | t_base_grdy = bench_latency( |
| 184 | baseline_greedy_sample, (logits,), n_iters=args.iters |
| 185 | ) |
| 186 | t_cust_grdy = bench_latency( |
| 187 | torch.ops._C.greedy_argmax, (logits,), n_iters=args.iters |
| 188 | ) |
| 189 | |
| 190 | rand_speedup = t_base_rand / t_fused if t_fused > 0 else 0 |
| 191 | grdy_speedup = t_base_grdy / t_cust_grdy if t_cust_grdy > 0 else 0 |
| 192 | |
| 193 | print( |
| 194 | f"{B:>5} {V:>7} " |
| 195 | f"{t_base_rand:>10.1f} {t_fused:>10.1f} " |
| 196 | f"{rand_speedup:>9.2f}x " |
| 197 | f"{t_base_grdy:>10.1f} {t_cust_grdy:>10.1f} " |
| 198 | f"{grdy_speedup:>9.2f}x" |
| 199 | ) |
| 200 | |
| 201 | print("-" * len(header)) |
| 202 | |
| 203 | if args.profile: |
| 204 | print("\nRunning torch.profiler (batch=16, vocab=128256) ...") |
| 205 | logits = torch.randn(16, 128256, dtype=torch.float32) |
| 206 | seeds = torch.arange(16, dtype=torch.long) |
| 207 | |
| 208 | # warmup |
| 209 | for _ in range(10): |
| 210 | baseline_random_sample(logits) |
| 211 | torch.ops._C.fused_gumbel_argmax(logits, seeds) |
| 212 | |
| 213 | run_profile(logits, seeds) |
| 214 | |
| 215 | |
| 216 | if __name__ == "__main__": |
| 217 | main() |