| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | """Generic benchmark harness for vLLM IR ops. |
| 5 | |
| 6 | Usage: |
| 7 | python benchmarks/kernels/ir/bench_ir_ops.py |
| 8 | python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm |
| 9 | python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm,silu_mul |
| 10 | python benchmarks/kernels/ir/bench_ir_ops.py --no-cuda-graph |
| 11 | python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm --save-path ./results/ |
| 12 | """ |
| 13 | |
| 14 | import argparse |
| 15 | import contextlib |
| 16 | import csv |
| 17 | import dataclasses |
| 18 | import datetime |
| 19 | import math |
| 20 | import os |
| 21 | import subprocess |
| 22 | import sys |
| 23 | import tempfile |
| 24 | |
| 25 | # Ensure repo root is on sys.path so `benchmarks` is importable as a package. |
| 26 | _REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")) |
| 27 | if _REPO_ROOT not in sys.path: |
| 28 | sys.path.insert(0, _REPO_ROOT) |
| 29 | |
| 30 | # Suppress noisy C++ warnings from vllm kernel registration (written to fd 2 |
| 31 | # directly by the dynamic linker, so Python-level sys.stderr redirect won't |
| 32 | # catch them). |
| 33 | _saved_fd = os.dup(2) |
| 34 | try: |
| 35 | with open(os.devnull, "w") as _devnull: |
| 36 | os.dup2(_devnull.fileno(), 2) |
| 37 | import torch |
| 38 | |
| 39 | import vllm.kernels # noqa: E402, F401 |
| 40 | finally: |
| 41 | os.dup2(_saved_fd, 2) |
| 42 | os.close(_saved_fd) |
| 43 | |
| 44 | from tqdm import tqdm # noqa: E402 |
| 45 | |
| 46 | from benchmarks.kernels.ir.shapes import SHAPE_CONFIGS # noqa: E402 # isort: skip |
| 47 | from vllm.ir.op import IrOp # noqa: E402 |
| 48 | from vllm.platforms import current_platform # noqa: E402 |
| 49 | from vllm.triton_utils import triton # noqa: E402 |
| 50 | |
| 51 | |
| 52 | @dataclasses.dataclass(frozen=True) |
| 53 | class BenchConfig: |
| 54 | use_cuda_graph: bool = True |
| 55 | warmup: int = 25 |
| 56 | rep: int = 100 |
| 57 | |
| 58 | |
| 59 | def _pkg_version(name: str) -> str: |
| 60 | from importlib.metadata import PackageNotFoundError, version |
| 61 | |
| 62 | with contextlib.suppress(PackageNotFoundError): |
| 63 | return version(name) |
| 64 | return "not installed" |
| 65 | |
| 66 | |
| 67 | _METADATA_LABELS = { |
| 68 | "timestamp": "Timestamp", |
| 69 | "git_commit": "Git commit", |
| 70 | "vllm": "vLLM", |
| 71 | "pytorch": "PyTorch", |
| 72 | "cuda_runtime": "CUDA runtime", |
| 73 | "triton": "Triton", |
| 74 | "cutlass": "CUTLASS", |
| 75 | "helion": "Helion", |
| 76 | "device": "Device", |
| 77 | "bench_mode": "Bench mode", |
| 78 | "warmup": "Warmup", |
| 79 | "rep": "Repetitions", |
| 80 | } |
| 81 | |
| 82 | |
| 83 | def collect_env_metadata(cfg: BenchConfig) -> dict[str, str]: |
| 84 | from vllm.collect_env import get_env_info |
| 85 | |
| 86 | env = get_env_info() |
| 87 | |
| 88 | git_sha = "unknown" |
| 89 | with contextlib.suppress(subprocess.CalledProcessError, FileNotFoundError): |
| 90 | git_sha = ( |
| 91 | subprocess.check_output( |
| 92 | ["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.DEVNULL |
| 93 | ) |
| 94 | .decode() |
| 95 | .strip() |
| 96 | ) |
| 97 | |
| 98 | device_name = current_platform.get_device_name() |
| 99 | |
| 100 | warmup_note = " ms" if not cfg.use_cuda_graph else " ms (ignored)" |
| 101 | rep_note = " replays" if cfg.use_cuda_graph else " ms" |
| 102 | |
| 103 | return { |
| 104 | "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), |
| 105 | "git_commit": git_sha, |
| 106 | "vllm": str(env.vllm_version), |
| 107 | "pytorch": str(env.torch_version), |
| 108 | "cuda_runtime": str(env.cuda_runtime_version), |
| 109 | "triton": triton.__version__, |
| 110 | "cutlass": _pkg_version("nvidia-cutlass-dsl"), |
| 111 | "helion": _pkg_version("helion"), |
| 112 | "device": device_name, |
| 113 | "bench_mode": "cuda_graph" if cfg.use_cuda_graph else "eager", |
| 114 | "warmup": f"{cfg.warmup}{warmup_note}", |
| 115 | "rep": f"{cfg.rep}{rep_note}", |
| 116 | } |
| 117 | |
| 118 | |
| 119 | def print_metadata(metadata: dict[str, str]): |
| 120 | print("=" * 60) |
| 121 | for key, val in metadata.items(): |
| 122 | print(f"{_METADATA_LABELS.get(key, key) + ':':<16}{val}") |
| 123 | print("=" * 60) |
| 124 | |
| 125 | |
| 126 | def _clone_args(args: tuple) -> tuple: |
| 127 | return tuple(a.clone() if isinstance(a, torch.Tensor) else a for a in args) |
| 128 | |
| 129 | |
| 130 | # TODO(gmagogsfm): When the `maybe_inplace` PR lands, ops marked as |
| 131 | # inplace=True will mutate bench_args across iterations. Both CUDA graph |
| 132 | # and eager modes will accumulate drift from repeated in-place mutation. |
| 133 | # We need to re-clone inputs per iteration for inplace ops. |
| 134 | def _bench_one(fn, args, cfg: BenchConfig) -> float: |
| 135 | bench_args = _clone_args(args) |
| 136 | bench_fn = lambda: fn(*bench_args) |
| 137 | |
| 138 | if cfg.use_cuda_graph: |
| 139 | ms = triton.testing.do_bench_cudagraph(bench_fn, rep=cfg.rep, quantiles=[0.5]) |
| 140 | else: |
| 141 | ms = triton.testing.do_bench( |
| 142 | bench_fn, warmup=cfg.warmup, rep=cfg.rep, quantiles=[0.5] |
| 143 | ) |
| 144 | return ms * 1000 |
| 145 | |
| 146 | |
| 147 | # TODO(gmagogsfm): Once compiled native implementation lands (#38775), |
| 148 | # the benchmark baseline should be the compiled native (what vLLM runs by |
| 149 | # default) rather than the uncompiled native implementation. |
| 150 | def collect_timings( |
| 151 | op: IrOp, shape_configs: list[dict], cfg: BenchConfig |
| 152 | ) -> tuple[list[str], list[str], dict[str, dict[str, float]]]: |
| 153 | def fmt(v) -> str: |
| 154 | return str(v).split(".")[-1] if isinstance(v, torch.dtype) else str(v) |
| 155 | |
| 156 | case_names = [ |
| 157 | "_".join(f"{k}={fmt(v)}" for k, v in kwargs.items()) for kwargs in shape_configs |
| 158 | ] |
| 159 | providers = [n for n, impl in op.impls.items() if impl.supported] |
| 160 | |
| 161 | results: dict[str, dict[str, float]] = {c: {} for c in case_names} |
| 162 | for provider in providers: |
| 163 | impl = op.impls[provider] |
| 164 | desc = f"{op.name} / {provider}" |
| 165 | for case_name, kwargs in tqdm( |
| 166 | zip(case_names, shape_configs), |
| 167 | desc=desc, |
| 168 | total=len(case_names), |
| 169 | unit=" cases", |
| 170 | ): |
| 171 | args = op.generate_inputs(**kwargs) |
| 172 | if impl.supports_args(*args): |
| 173 | results[case_name][provider] = _bench_one(impl.impl_fn, args, cfg) |
| 174 | else: |
| 175 | results[case_name][provider] = float("nan") |
| 176 | |
| 177 | return case_names, providers, results |
| 178 | |
| 179 | |
| 180 | def analyze_results( |
| 181 | op_name: str, |
| 182 | case_names: list[str], |
| 183 | providers: list[str], |
| 184 | results: dict[str, dict[str, float]], |
| 185 | ) -> tuple[list[dict[str, str]], list[dict[str, str]], list[str]]: |
| 186 | native_col = "native" |
| 187 | non_native = [p for p in providers if p != native_col] |
| 188 | |
| 189 | header_cols = ["case"] |
| 190 | for p in providers: |
| 191 | header_cols.append(f"{p} (us)") |
| 192 | for p in non_native: |
| 193 | header_cols.append(f"{p} speedup") |
| 194 | |
| 195 | detail_rows: list[dict[str, str]] = [] |
| 196 | speedup_data: dict[str, list[tuple[float, str]]] = {p: [] for p in non_native} |
| 197 | |
| 198 | for case_name in case_names: |
| 199 | timings = results[case_name] |
| 200 | row: dict[str, str] = {"case": case_name} |
| 201 | |
| 202 | for p in providers: |
| 203 | val = timings.get(p, float("nan")) |
| 204 | row[f"{p} (us)"] = f"{val:.2f}" if not math.isnan(val) else "n/a" |
| 205 | |
| 206 | native_us = timings.get(native_col, float("nan")) |
| 207 | for p in non_native: |
| 208 | p_us = timings.get(p, float("nan")) |
| 209 | if not math.isnan(native_us) and not math.isnan(p_us) and p_us > 0: |
| 210 | speedup = native_us / p_us |
| 211 | row[f"{p} speedup"] = f"{speedup:.2f}x" |
| 212 | speedup_data[p].append((speedup, case_name)) |
| 213 | else: |
| 214 | row[f"{p} speedup"] = "n/a" |
| 215 | |
| 216 | detail_rows.append(row) |
| 217 | |
| 218 | summary_rows: list[dict[str, str]] = [] |
| 219 | for p in non_native: |
| 220 | entries = speedup_data[p] |
| 221 | if not entries: |
| 222 | continue |
| 223 | speedups = [s for s, _ in entries] |
| 224 | geomean = math.exp(sum(math.log(s) for s in speedups) / len(speedups)) |
| 225 | best_val, best_case = max(entries) |
| 226 | worst_val, worst_case = min(entries) |
| 227 | wins = sum(1 for s in speedups if s > 1.0) |
| 228 | losses = sum(1 for s in speedups if s < 1.0) |
| 229 | total = len(speedups) |
| 230 | |
| 231 | print(f"\n{p} vs native ({wins}/{total} faster, {losses}/{total} slower):") |
| 232 | print(f" geomean speedup: {geomean:.2f}x") |
| 233 | print(f" best: {best_val:.2f}x ({best_case})") |
| 234 | print(f" worst: {worst_val:.2f}x ({worst_case})") |
| 235 | |
| 236 | summary_rows.append( |
| 237 | { |
| 238 | "op": op_name, |
| 239 | "provider": p, |
| 240 | "geomean_speedup": f"{geomean:.2f}", |
| 241 | "best_speedup": f"{best_val:.2f}", |
| 242 | "best_case": best_case, |
| 243 | "worst_speedup": f"{worst_val:.2f}", |
| 244 | "worst_case": worst_case, |
| 245 | "wins": str(wins), |
| 246 | "losses": str(losses), |
| 247 | "total": str(total), |
| 248 | } |
| 249 | ) |
| 250 | |
| 251 | return detail_rows, summary_rows, header_cols |
| 252 | |
| 253 | |
| 254 | def write_csv(path: str, rows: list[dict[str, str]], fieldnames: list[str]): |
| 255 | with open(path, "w", newline="") as f: |
| 256 | writer = csv.DictWriter(f, fieldnames=fieldnames) |
| 257 | writer.writeheader() |
| 258 | writer.writerows(rows) |
| 259 | |
| 260 | |
| 261 | def save_results( |
| 262 | save_dir: str, |
| 263 | op_name: str, |
| 264 | detail_rows: list[dict[str, str]], |
| 265 | header_cols: list[str], |
| 266 | all_summary_rows: list[dict[str, str]], |
| 267 | metadata: dict[str, str], |
| 268 | ): |
| 269 | write_csv( |
| 270 | os.path.join(save_dir, f"{op_name}_detail.csv"), |
| 271 | detail_rows, |
| 272 | header_cols, |
| 273 | ) |
| 274 | if all_summary_rows: |
| 275 | write_csv( |
| 276 | os.path.join(save_dir, "summary.csv"), |
| 277 | all_summary_rows, |
| 278 | list(all_summary_rows[0].keys()), |
| 279 | ) |
| 280 | write_csv( |
| 281 | os.path.join(save_dir, "metadata.csv"), |
| 282 | [metadata], |
| 283 | list(metadata.keys()), |
| 284 | ) |
| 285 | |
| 286 | |
| 287 | def parse_args(): |
| 288 | parser = argparse.ArgumentParser(description="Benchmark vLLM IR ops") |
| 289 | parser.add_argument( |
| 290 | "--ops", |
| 291 | type=str, |
| 292 | default=None, |
| 293 | help="Comma-separated list of op names to benchmark (substring match)", |
| 294 | ) |
| 295 | parser.add_argument( |
| 296 | "--no-cuda-graph", |
| 297 | action="store_true", |
| 298 | help="Disable CUDA graph; use do_bench with L2 cache flushing instead", |
| 299 | ) |
| 300 | parser.add_argument( |
| 301 | "--warmup", |
| 302 | type=int, |
| 303 | default=25, |
| 304 | help="Warmup time in ms (do_bench) or ignored with CUDA graph (default: 25)", |
| 305 | ) |
| 306 | parser.add_argument( |
| 307 | "--rep", |
| 308 | type=int, |
| 309 | default=100, |
| 310 | help="Repetition time in ms (do_bench) or number of graph replays " |
| 311 | "(do_bench_cudagraph) (default: 100)", |
| 312 | ) |
| 313 | parser.add_argument( |
| 314 | "--save-path", |
| 315 | type=str, |
| 316 | default=None, |
| 317 | help="Directory to save results (default: auto-created temp dir)", |
| 318 | ) |
| 319 | return parser.parse_args() |
| 320 | |
| 321 | |
| 322 | def main(): |
| 323 | args = parse_args() |
| 324 | cfg = BenchConfig( |
| 325 | use_cuda_graph=not args.no_cuda_graph, |
| 326 | warmup=args.warmup, |
| 327 | rep=args.rep, |
| 328 | ) |
| 329 | |
| 330 | torch.set_default_device(current_platform.device_type) |
| 331 | |
| 332 | metadata = collect_env_metadata(cfg) |
| 333 | print_metadata(metadata) |
| 334 | |
| 335 | timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") |
| 336 | save_dir = args.save_path or os.path.join( |
| 337 | tempfile.gettempdir(), f"vllm_ir_bench_{timestamp}" |
| 338 | ) |
| 339 | os.makedirs(save_dir, exist_ok=True) |
| 340 | |
| 341 | op_filters = [f.strip() for f in args.ops.split(",")] if args.ops else None |
| 342 | all_summary_rows: list[dict[str, str]] = [] |
| 343 | |
| 344 | for op in IrOp.registry.values(): |
| 345 | if op_filters and not any(f in op.name for f in op_filters): |
| 346 | continue |
| 347 | if not op.has_input_generator: |
| 348 | print(f"Skipping op '{op.name}': no input generator registered") |
| 349 | continue |
| 350 | if op.name not in SHAPE_CONFIGS: |
| 351 | raise RuntimeError( |
| 352 | f"No benchmark shape config for op '{op.name}'. " |
| 353 | f"Add it to benchmarks/kernels/ir/shapes.py" |
| 354 | ) |
| 355 | |
| 356 | case_names, providers, results = collect_timings( |
| 357 | op, SHAPE_CONFIGS[op.name], cfg |
| 358 | ) |
| 359 | detail_rows, summary_rows, header_cols = analyze_results( |
| 360 | op.name, case_names, providers, results |
| 361 | ) |
| 362 | all_summary_rows.extend(summary_rows) |
| 363 | |
| 364 | save_results( |
| 365 | save_dir, |
| 366 | op.name, |
| 367 | detail_rows, |
| 368 | header_cols, |
| 369 | all_summary_rows, |
| 370 | metadata, |
| 371 | ) |
| 372 | |
| 373 | print(f"\nResults saved to: {save_dir}") |
| 374 | |
| 375 | |
| 376 | if __name__ == "__main__": |
| 377 | main() |