@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fe5c115ee4
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 22:30 UTC

[vLLM IR] Add IR op testing and benchmarking infrastructure (#40167)

Signed-off-by: Yanan Cao <gmagogsfm@gmail.com> Co-authored-by: Theresa Shan <Theresa.Shan@amd.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Yanan Cao committed Apr 20, 2026 at 17:23 UTC fe5c115ee48229faeb53e4dd66f384d116772d00
12 files changed +673 -45
benchmarks/__init__.py
benchmarks/kernels/__init__.py
benchmarks/kernels/ir/__init__.py
benchmarks/kernels/ir/bench_ir_ops.py new
+378
@@ -0,0 +1,378 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +"""
5 +Generic benchmark harness for vLLM IR ops.
6 +
7 +Usage:
8 + python benchmarks/kernels/ir/bench_ir_ops.py
9 + python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm
10 + python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm,silu_mul
11 + python benchmarks/kernels/ir/bench_ir_ops.py --no-cuda-graph
12 + python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm --save-path ./results/
13 +"""
14 +
15 +import argparse
16 +import contextlib
17 +import csv
18 +import dataclasses
19 +import datetime
20 +import math
21 +import os
22 +import subprocess
23 +import sys
24 +import tempfile
25 +
26 +# Ensure repo root is on sys.path so `benchmarks` is importable as a package.
27 +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
28 +if _REPO_ROOT not in sys.path:
29 + sys.path.insert(0, _REPO_ROOT)
30 +
31 +# Suppress noisy C++ warnings from vllm kernel registration (written to fd 2
32 +# directly by the dynamic linker, so Python-level sys.stderr redirect won't
33 +# catch them).
34 +_saved_fd = os.dup(2)
35 +try:
36 + with open(os.devnull, "w") as _devnull:
37 + os.dup2(_devnull.fileno(), 2)
38 + import torch
39 +
40 + import vllm.kernels # noqa: E402, F401
41 +finally:
42 + os.dup2(_saved_fd, 2)
43 + os.close(_saved_fd)
44 +
45 +from tqdm import tqdm # noqa: E402
46 +
47 +from benchmarks.kernels.ir.shapes import SHAPE_CONFIGS # noqa: E402 # isort: skip
48 +from vllm.ir.op import IrOp # noqa: E402
49 +from vllm.platforms import current_platform # noqa: E402
50 +from vllm.triton_utils import triton # noqa: E402
51 +
52 +
53 +@dataclasses.dataclass(frozen=True)
54 +class BenchConfig:
55 + use_cuda_graph: bool = True
56 + warmup: int = 25
57 + rep: int = 100
58 +
59 +
60 +def _pkg_version(name: str) -> str:
61 + from importlib.metadata import PackageNotFoundError, version
62 +
63 + with contextlib.suppress(PackageNotFoundError):
64 + return version(name)
65 + return "not installed"
66 +
67 +
68 +_METADATA_LABELS = {
69 + "timestamp": "Timestamp",
70 + "git_commit": "Git commit",
71 + "vllm": "vLLM",
72 + "pytorch": "PyTorch",
73 + "cuda_runtime": "CUDA runtime",
74 + "triton": "Triton",
75 + "cutlass": "CUTLASS",
76 + "helion": "Helion",
77 + "device": "Device",
78 + "bench_mode": "Bench mode",
79 + "warmup": "Warmup",
80 + "rep": "Repetitions",
81 +}
82 +
83 +
84 +def collect_env_metadata(cfg: BenchConfig) -> dict[str, str]:
85 + from vllm.collect_env import get_env_info
86 +
87 + env = get_env_info()
88 +
89 + git_sha = "unknown"
90 + with contextlib.suppress(subprocess.CalledProcessError, FileNotFoundError):
91 + git_sha = (
92 + subprocess.check_output(
93 + ["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.DEVNULL
94 + )
95 + .decode()
96 + .strip()
97 + )
98 +
99 + device_name = current_platform.get_device_name()
100 +
101 + warmup_note = " ms" if not cfg.use_cuda_graph else " ms (ignored)"
102 + rep_note = " replays" if cfg.use_cuda_graph else " ms"
103 +
104 + return {
105 + "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
106 + "git_commit": git_sha,
107 + "vllm": str(env.vllm_version),
108 + "pytorch": str(env.torch_version),
109 + "cuda_runtime": str(env.cuda_runtime_version),
110 + "triton": triton.__version__,
111 + "cutlass": _pkg_version("nvidia-cutlass-dsl"),
112 + "helion": _pkg_version("helion"),
113 + "device": device_name,
114 + "bench_mode": "cuda_graph" if cfg.use_cuda_graph else "eager",
115 + "warmup": f"{cfg.warmup}{warmup_note}",
116 + "rep": f"{cfg.rep}{rep_note}",
117 + }
118 +
119 +
120 +def print_metadata(metadata: dict[str, str]):
121 + print("=" * 60)
122 + for key, val in metadata.items():
123 + print(f"{_METADATA_LABELS.get(key, key) + ':':<16}{val}")
124 + print("=" * 60)
125 +
126 +
127 +def _clone_args(args: tuple) -> tuple:
128 + return tuple(a.clone() if isinstance(a, torch.Tensor) else a for a in args)
129 +
130 +
131 +# TODO(gmagogsfm): When the `maybe_inplace` PR lands, ops marked as
132 +# inplace=True will mutate bench_args across iterations. Both CUDA graph
133 +# and eager modes will accumulate drift from repeated in-place mutation.
134 +# We need to re-clone inputs per iteration for inplace ops.
135 +def _bench_one(fn, args, cfg: BenchConfig) -> float:
136 + bench_args = _clone_args(args)
137 + bench_fn = lambda: fn(*bench_args)
138 +
139 + if cfg.use_cuda_graph:
140 + ms = triton.testing.do_bench_cudagraph(bench_fn, rep=cfg.rep, quantiles=[0.5])
141 + else:
142 + ms = triton.testing.do_bench(
143 + bench_fn, warmup=cfg.warmup, rep=cfg.rep, quantiles=[0.5]
144 + )
145 + return ms * 1000
146 +
147 +
148 +# TODO(gmagogsfm): Once compiled native implementation lands (#38775),
149 +# the benchmark baseline should be the compiled native (what vLLM runs by
150 +# default) rather than the uncompiled native implementation.
151 +def collect_timings(
152 + op: IrOp, shape_configs: list[dict], cfg: BenchConfig
153 +) -> tuple[list[str], list[str], dict[str, dict[str, float]]]:
154 + def fmt(v) -> str:
155 + return str(v).split(".")[-1] if isinstance(v, torch.dtype) else str(v)
156 +
157 + case_names = [
158 + "_".join(f"{k}={fmt(v)}" for k, v in kwargs.items()) for kwargs in shape_configs
159 + ]
160 + providers = [n for n, impl in op.impls.items() if impl.supported]
161 +
162 + results: dict[str, dict[str, float]] = {c: {} for c in case_names}
163 + for provider in providers:
164 + impl = op.impls[provider]
165 + desc = f"{op.name} / {provider}"
166 + for case_name, kwargs in tqdm(
167 + zip(case_names, shape_configs),
168 + desc=desc,
169 + total=len(case_names),
170 + unit=" cases",
171 + ):
172 + args = op.generate_inputs(**kwargs)
173 + if impl.supports_args(*args):
174 + results[case_name][provider] = _bench_one(impl.impl_fn, args, cfg)
175 + else:
176 + results[case_name][provider] = float("nan")
177 +
178 + return case_names, providers, results
179 +
180 +
181 +def analyze_results(
182 + op_name: str,
183 + case_names: list[str],
184 + providers: list[str],
185 + results: dict[str, dict[str, float]],
186 +) -> tuple[list[dict[str, str]], list[dict[str, str]], list[str]]:
187 + native_col = "native"
188 + non_native = [p for p in providers if p != native_col]
189 +
190 + header_cols = ["case"]
191 + for p in providers:
192 + header_cols.append(f"{p} (us)")
193 + for p in non_native:
194 + header_cols.append(f"{p} speedup")
195 +
196 + detail_rows: list[dict[str, str]] = []
197 + speedup_data: dict[str, list[tuple[float, str]]] = {p: [] for p in non_native}
198 +
199 + for case_name in case_names:
200 + timings = results[case_name]
201 + row: dict[str, str] = {"case": case_name}
202 +
203 + for p in providers:
204 + val = timings.get(p, float("nan"))
205 + row[f"{p} (us)"] = f"{val:.2f}" if not math.isnan(val) else "n/a"
206 +
207 + native_us = timings.get(native_col, float("nan"))
208 + for p in non_native:
209 + p_us = timings.get(p, float("nan"))
210 + if not math.isnan(native_us) and not math.isnan(p_us) and p_us > 0:
211 + speedup = native_us / p_us
212 + row[f"{p} speedup"] = f"{speedup:.2f}x"
213 + speedup_data[p].append((speedup, case_name))
214 + else:
215 + row[f"{p} speedup"] = "n/a"
216 +
217 + detail_rows.append(row)
218 +
219 + summary_rows: list[dict[str, str]] = []
220 + for p in non_native:
221 + entries = speedup_data[p]
222 + if not entries:
223 + continue
224 + speedups = [s for s, _ in entries]
225 + geomean = math.exp(sum(math.log(s) for s in speedups) / len(speedups))
226 + best_val, best_case = max(entries)
227 + worst_val, worst_case = min(entries)
228 + wins = sum(1 for s in speedups if s > 1.0)
229 + losses = sum(1 for s in speedups if s < 1.0)
230 + total = len(speedups)
231 +
232 + print(f"\n{p} vs native ({wins}/{total} faster, {losses}/{total} slower):")
233 + print(f" geomean speedup: {geomean:.2f}x")
234 + print(f" best: {best_val:.2f}x ({best_case})")
235 + print(f" worst: {worst_val:.2f}x ({worst_case})")
236 +
237 + summary_rows.append(
238 + {
239 + "op": op_name,
240 + "provider": p,
241 + "geomean_speedup": f"{geomean:.2f}",
242 + "best_speedup": f"{best_val:.2f}",
243 + "best_case": best_case,
244 + "worst_speedup": f"{worst_val:.2f}",
245 + "worst_case": worst_case,
246 + "wins": str(wins),
247 + "losses": str(losses),
248 + "total": str(total),
249 + }
250 + )
251 +
252 + return detail_rows, summary_rows, header_cols
253 +
254 +
255 +def write_csv(path: str, rows: list[dict[str, str]], fieldnames: list[str]):
256 + with open(path, "w", newline="") as f:
257 + writer = csv.DictWriter(f, fieldnames=fieldnames)
258 + writer.writeheader()
259 + writer.writerows(rows)
260 +
261 +
262 +def save_results(
263 + save_dir: str,
264 + op_name: str,
265 + detail_rows: list[dict[str, str]],
266 + header_cols: list[str],
267 + all_summary_rows: list[dict[str, str]],
268 + metadata: dict[str, str],
269 +):
270 + write_csv(
271 + os.path.join(save_dir, f"{op_name}_detail.csv"),
272 + detail_rows,
273 + header_cols,
274 + )
275 + if all_summary_rows:
276 + write_csv(
277 + os.path.join(save_dir, "summary.csv"),
278 + all_summary_rows,
279 + list(all_summary_rows[0].keys()),
280 + )
281 + write_csv(
282 + os.path.join(save_dir, "metadata.csv"),
283 + [metadata],
284 + list(metadata.keys()),
285 + )
286 +
287 +
288 +def parse_args():
289 + parser = argparse.ArgumentParser(description="Benchmark vLLM IR ops")
290 + parser.add_argument(
291 + "--ops",
292 + type=str,
293 + default=None,
294 + help="Comma-separated list of op names to benchmark (substring match)",
295 + )
296 + parser.add_argument(
297 + "--no-cuda-graph",
298 + action="store_true",
299 + help="Disable CUDA graph; use do_bench with L2 cache flushing instead",
300 + )
301 + parser.add_argument(
302 + "--warmup",
303 + type=int,
304 + default=25,
305 + help="Warmup time in ms (do_bench) or ignored with CUDA graph (default: 25)",
306 + )
307 + parser.add_argument(
308 + "--rep",
309 + type=int,
310 + default=100,
311 + help="Repetition time in ms (do_bench) or number of graph replays "
312 + "(do_bench_cudagraph) (default: 100)",
313 + )
314 + parser.add_argument(
315 + "--save-path",
316 + type=str,
317 + default=None,
318 + help="Directory to save results (default: auto-created temp dir)",
319 + )
320 + return parser.parse_args()
321 +
322 +
323 +def main():
324 + args = parse_args()
325 + cfg = BenchConfig(
326 + use_cuda_graph=not args.no_cuda_graph,
327 + warmup=args.warmup,
328 + rep=args.rep,
329 + )
330 +
331 + torch.set_default_device(current_platform.device_type)
332 +
333 + metadata = collect_env_metadata(cfg)
334 + print_metadata(metadata)
335 +
336 + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
337 + save_dir = args.save_path or os.path.join(
338 + tempfile.gettempdir(), f"vllm_ir_bench_{timestamp}"
339 + )
340 + os.makedirs(save_dir, exist_ok=True)
341 +
342 + op_filters = [f.strip() for f in args.ops.split(",")] if args.ops else None
343 + all_summary_rows: list[dict[str, str]] = []
344 +
345 + for op in IrOp.registry.values():
346 + if op_filters and not any(f in op.name for f in op_filters):
347 + continue
348 + if not op.has_input_generator:
349 + print(f"Skipping op '{op.name}': no input generator registered")
350 + continue
351 + if op.name not in SHAPE_CONFIGS:
352 + raise RuntimeError(
353 + f"No benchmark shape config for op '{op.name}'. "
354 + f"Add it to benchmarks/kernels/ir/shapes.py"
355 + )
356 +
357 + case_names, providers, results = collect_timings(
358 + op, SHAPE_CONFIGS[op.name], cfg
359 + )
360 + detail_rows, summary_rows, header_cols = analyze_results(
361 + op.name, case_names, providers, results
362 + )
363 + all_summary_rows.extend(summary_rows)
364 +
365 + save_results(
366 + save_dir,
367 + op.name,
368 + detail_rows,
369 + header_cols,
370 + all_summary_rows,
371 + metadata,
372 + )
373 +
374 + print(f"\nResults saved to: {save_dir}")
375 +
376 +
377 +if __name__ == "__main__":
378 + main()
benchmarks/kernels/ir/shapes.py new
+29
@@ -0,0 +1,29 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Shape configurations for IR op benchmarks.
5 +"""
6 +
7 +import torch
8 +
9 +NUM_TOKENS = [1, 2, 4, 16, 64, 256, 1024, 4096, 16384]
10 +COMMON_HIDDEN_SIZES = [
11 + 2048, # Llama 3.2 1B, Qwen 3 MoE 30B-A3B, Gemma 3n
12 + 3072, # Gemma 7B/9B
13 + 4096, # Llama 3 8B, Qwen 3 8B, Mistral 7B
14 + 5120, # Llama 4 Scout 17B-16E
15 + 7168, # DeepSeek V3
16 + 8192, # Llama 3 70B
17 + 16384, # Llama 3 405B
18 +]
19 +
20 +# Each entry maps an op name to a list of kwarg dicts that will be passed
21 +# to that op's registered input generator via op.generate_inputs(**kwargs).
22 +SHAPE_CONFIGS: dict[str, list[dict]] = {
23 + "rms_norm": [
24 + {"num_tokens": n, "hidden_size": d, "dtype": dtype}
25 + for dtype in [torch.float16, torch.bfloat16, torch.float32]
26 + for d in COMMON_HIDDEN_SIZES
27 + for n in NUM_TOKENS
28 + ],
29 +}
tests/ir/ir_test_utils.py new
+46
@@ -0,0 +1,46 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Shared test utilities for vLLM IR op correctness tests.
5 +"""
6 +
7 +import torch
8 +
9 +from vllm.ir.op import IrOp
10 +
11 +NUM_TOKENS = [1, 8, 17, 32, 512, 2048]
12 +COMMON_HIDDEN_SIZES = [
13 + 2048, # Llama 3.2 1B, Qwen 3 MoE 30B-A3B, Gemma 3n
14 + 4096, # Llama 3 8B, Qwen 3 8B
15 + 5120, # Llama 4 Scout 17B-16E
16 + 7168, # DeepSeek V3
17 + 8192, # Llama 3 70B
18 +]
19 +
20 +
21 +def clone_args(args: tuple) -> tuple:
22 + return tuple(a.clone() if isinstance(a, torch.Tensor) else a for a in args)
23 +
24 +
25 +def supported_providers(op: IrOp) -> list[str]:
26 + return [
27 + name for name, impl in op.impls.items() if name != "native" and impl.supported
28 + ]
29 +
30 +
31 +def assert_close(op: IrOp, actual, expected):
32 + if isinstance(actual, torch.Tensor):
33 + tol = op.get_tolerance(actual.dtype)
34 + try:
35 + torch.testing.assert_close(actual, expected, **tol)
36 + except AssertionError as e:
37 + raise AssertionError(
38 + f"{e}\n\nTo adjust tolerance, use:\n"
39 + f" ir.ops.{op.name}.override_tolerance("
40 + f"{actual.dtype}, atol=..., rtol=...)"
41 + ) from None
42 + elif isinstance(actual, (tuple, list)):
43 + for a, ex in zip(actual, expected):
44 + assert_close(op, a, ex)
45 + else:
46 + assert actual == expected
tests/ir/test_op.py
+65
@@ -495,3 +495,68 @@ def test_uuid_and_oot(tmp_path: Path):
495 assert uuid2 == uuid
496 assert uuid2 != uuid1
497 del _custom_mm.impls["impl_mm_oot"]
498 +
499 +
500 +def _test_native(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
501 + return x + y
502 +
503 +
504 +def _make_op_with_generator(name: str = "_ig_test"):
505 + op = IrOp(name, _test_native)
506 +
507 + @op.register_input_generator
508 + def _gen(n: int = 4):
509 + x = torch.randn(n, 3)
510 + y = torch.randn(n, 3)
511 + return x, y
512 +
513 + return op
514 +
515 +
516 +def _test_native_single(x: torch.Tensor) -> torch.Tensor:
517 + return x
518 +
519 +
520 +class TestInputGenerator:
521 + def test_no_input_generator_by_default(self):
522 + op = IrOp("_ig_test_no_gen", _test_native_single)
523 + assert not op.has_input_generator
524 +
525 + def test_register_input_generator(self):
526 + op = _make_op_with_generator("_ig_test_reg")
527 + assert op.has_input_generator
528 +
529 + def test_generate_inputs_returns_tuple(self):
530 + op = _make_op_with_generator("_ig_test_tuple")
531 + result = op.generate_inputs(n=2)
532 + assert isinstance(result, tuple)
533 + assert len(result) == 2
534 + assert result[0].shape == (2, 3)
535 + assert result[1].shape == (2, 3)
536 +
537 + def test_generate_inputs_default_kwargs(self):
538 + op = _make_op_with_generator("_ig_test_default")
539 + result = op.generate_inputs()
540 + assert result[0].shape == (4, 3)
541 +
542 + def test_generate_inputs_without_registration_raises(self):
543 + op = IrOp("_ig_test_no_gen_raises", _test_native_single)
544 + with pytest.raises(RuntimeError, match="No input generator"):
545 + op.generate_inputs()
546 +
547 +
548 +class TestTolerance:
549 + def test_override_and_get_tolerance(self):
550 + op = IrOp("_tol_test", _test_native)
551 +
552 + tol = op.get_tolerance(torch.float32)
553 + assert tol == {"atol": 1e-5, "rtol": 1.3e-6}
554 +
555 + op.override_tolerance(torch.float32, atol=0.1, rtol=0.2)
556 + assert op.get_tolerance(torch.float32) == {"atol": 0.1, "rtol": 0.2}
557 + assert op.get_tolerance(torch.float16) == {"atol": 1e-3, "rtol": 1e-3}
558 +
559 + def test_get_tolerance_raises_for_unknown_dtype(self):
560 + op = IrOp("_tol_test_unknown", _test_native)
561 + with pytest.raises(ValueError, match="No tolerance defined"):
562 + op.get_tolerance(torch.complex64)
tests/kernels/ir/test_ir_ops.py new
+20
@@ -0,0 +1,20 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Meta-tests for vLLM IR op infrastructure.
5 +
6 +Ensures all registered ops have input generators defined.
7 +Per-op correctness tests live alongside their op definitions
8 +(e.g. tests/kernels/ir/test_layernorm.py).
9 +"""
10 +
11 +import vllm.kernels # noqa: F401 — registers provider implementations
12 +from vllm.ir.op import IrOp
13 +
14 +
15 +def test_all_ops_have_input_generator():
16 + missing = [name for name, op in IrOp.registry.items() if not op.has_input_generator]
17 + assert not missing, (
18 + f"IR ops without input generators: {missing}. "
19 + f"Register one with @ir.ops.<name>.register_input_generator"
20 + )
tests/kernels/ir/test_layernorm.py
+47 -45
@@ -5,17 +5,17 @@ import torch
5
6 # This registers op implementations
7 import vllm.kernels # noqa: F401
8 +from tests.ir.ir_test_utils import (
9 + COMMON_HIDDEN_SIZES,
10 + NUM_TOKENS,
11 + assert_close,
12 + clone_args,
13 + supported_providers,
14 +)
15 from tests.kernels.allclose_default import get_default_rtol
16 from vllm import ir
17 from vllm.platforms import current_platform
18
12 -
13 -def rms_norm_inputs(n_tokens: int, hidden_size: int, dtype: torch.dtype):
14 - x = torch.randn(n_tokens, hidden_size, dtype=dtype)
15 - weight = torch.rand(hidden_size, dtype=dtype)
16 - return x, weight
17 -
18 -
19 rms_norm_native = ir.ops.rms_norm.impls["native"].impl_fn
20
21
@@ -40,8 +40,8 @@ def test_rms_norm_registration():
40
41
42 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
43 -@pytest.mark.parametrize("n_tokens", [1, 8, 17])
44 -@pytest.mark.parametrize("hidden_size", [16, 4096, 8192])
43 +@pytest.mark.parametrize("n_tokens", NUM_TOKENS)
44 +@pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES)
45 @pytest.mark.parametrize("epsilon", [1e-6, 1e-5])
46 @pytest.mark.skipif(
47 not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
@@ -53,7 +53,9 @@ class TestRMSNorm:
53 torch.set_default_device(current_platform.device_type)
54
55 def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
56 - x, weight = rms_norm_inputs(4, 8, dtype)
56 + x, weight, epsilon = ir.ops.rms_norm.generate_inputs(
57 + num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
58 + )
59 out = rms_norm_native(x, weight, epsilon=epsilon)
60
61 # Check shape, dtype, device
@@ -71,59 +73,59 @@ class TestRMSNorm:
73 out4 = rms_norm_native(x, None, epsilon=epsilon)
74 torch.testing.assert_close(out3, out4)
75
74 - @pytest.mark.parametrize("provider", ["vllm_c", "aiter", "xpu_kernels"])
76 + @pytest.mark.parametrize("provider", supported_providers(ir.ops.rms_norm))
77 def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
78 impl = ir.ops.rms_norm.impls[provider]
77 - if not impl.supported:
78 - pytest.skip(f"{provider} impl not supported on this platform")
79 -
80 - x, weight = rms_norm_inputs(n_tokens, hidden_size, dtype)
81 - args = (x, weight, epsilon, None)
82 -
83 - assert impl.supported
84 -
85 - if provider == "aiter" and dtype not in [torch.float16, torch.bfloat16]:
86 - assert not impl.supports_args(*args)
87 - return
88 -
89 - assert impl.supports_args(*args)
79 + x, weight, eps = ir.ops.rms_norm.generate_inputs(
80 + num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
81 + )
82 + args = (x, weight, eps)
83
91 - out_impl = impl.impl_fn(*args)
92 - out_native = rms_norm_native(*args)
84 + if not impl.supports_args(*args):
85 + pytest.skip(f"{provider} does not support args")
86
94 - torch.testing.assert_close(
95 - out_impl, out_native, rtol=get_default_rtol(out_impl), atol=1e-3
96 - )
87 + ref_output = rms_norm_native(*clone_args(args))
88 + output = impl.impl_fn(*clone_args(args))
89 + assert_close(ir.ops.rms_norm, output, ref_output)
90
91 # check that dispatched call matches direct call
92 with ir.ops.rms_norm.set_priority([provider, "native"]):
100 - out_impl2 = ir.ops.rms_norm(*args)
101 -
102 - # exact match
103 - torch.testing.assert_close(out_impl2, out_impl, rtol=0.0, atol=0.0)
93 + out_dispatched = ir.ops.rms_norm(*args)
94 + out_direct = impl.impl_fn(*args)
95 + torch.testing.assert_close(out_dispatched, out_direct, rtol=0.0, atol=0.0)
96
97 # none of these support variance_size override
106 - assert not impl.supports_args(x, weight, epsilon, 4)
107 - assert not impl.supports_args(x, weight, epsilon, variance_size=4)
98 + assert not impl.supports_args(x, weight, eps, 4)
99 + assert not impl.supports_args(x, weight, eps, variance_size=4)
100
101 # test weight=None behavior
110 - out_impl_no_weight = impl.impl_fn(x, None, epsilon)
111 - out_impl_unit_weight = impl.impl_fn(x, torch.ones_like(weight), epsilon)
112 - torch.testing.assert_close(
113 - out_impl_no_weight,
114 - out_impl_unit_weight,
115 - rtol=get_default_rtol(out_impl_no_weight),
116 - atol=2e-4,
117 - )
102 + out_no_weight = impl.impl_fn(x, None, eps)
103 + out_unit_weight = impl.impl_fn(x, torch.ones_like(weight), eps)
104 + assert_close(ir.ops.rms_norm, out_no_weight, out_unit_weight)
105
106 @pytest.mark.parametrize("provider", ["vllm_c", "aiter", "xpu_kernels", "native"])
107 def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider):
108 if not ir.ops.rms_norm.impls[provider].supported:
109 pytest.skip(f"{provider} impl not supported on this platform")
110
124 - x, weight = rms_norm_inputs(n_tokens, hidden_size, dtype)
125 - args = (x, weight, epsilon, None)
111 + args = ir.ops.rms_norm.generate_inputs(
112 + num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
113 + )
114
115 # When checking the torch op, we have to set priority and use dispatch
116 with ir.ops.rms_norm.set_priority([provider, "native"]):
117 torch.library.opcheck(torch.ops.vllm_ir.rms_norm, args)
118 +
119 +
120 +@pytest.mark.skipif(
121 + not current_platform.is_rocm(),
122 + reason="aiter is only supported on ROCm",
123 +)
124 +def test_aiter_rejects_unsupported_dtypes():
125 + torch.set_default_device(current_platform.device_type)
126 + impl = ir.ops.rms_norm.impls["aiter"]
127 + for dtype in [torch.float32, torch.float64]:
128 + args = ir.ops.rms_norm.generate_inputs(
129 + num_tokens=8, hidden_size=4096, dtype=dtype, epsilon=1e-5
130 + )
131 + assert not impl.supports_args(*args), f"aiter should reject dtype={dtype}"
vllm/ir/op.py
+37
@@ -9,10 +9,13 @@ from typing import Any, ClassVar, overload
9 import torch
10 from torch.library import Library, infer_schema
11
12 +from vllm.ir.tolerances import DEFAULT_TOLERANCES, ToleranceSpec
13 from vllm.ir.util import hash_source, weak_cache
14 from vllm.logger import init_logger
15 from vllm.logging_utils import lazy, tensors_str_no_data
16
17 +InputGenerator = Callable[..., tuple[Any, ...]]
18 +
19 vllm_ir_lib = Library("vllm_ir", "FRAGMENT")
20
21 logger = init_logger(__name__)
@@ -113,6 +116,8 @@ class IrOp:
116 self.impls: dict[str, IrOpImpl] = {}
117 self._priority_impls: list[IrOpImpl] = []
118 self._schema_str = infer_schema(native_impl, mutates_args=[])
119 + self._input_generator: InputGenerator | None = None
120 + self._tolerance_overrides: ToleranceSpec = {}
121
122 # native implementation
123 self.impls["native"] = IrOpImpl(
@@ -316,6 +321,38 @@ class IrOp:
321 def supported_providers(self) -> list[str]:
322 return [p.provider for p in self.impls.values() if p.supported]
323
324 + @property
325 + def has_input_generator(self) -> bool:
326 + return self._input_generator is not None
327 +
328 + def register_input_generator(self, fn: InputGenerator) -> InputGenerator:
329 + self._input_generator = fn
330 + return fn
331 +
332 + def generate_inputs(self, **kwargs: Any) -> tuple[Any, ...]:
333 + if self._input_generator is None:
334 + raise RuntimeError(
335 + f"No input generator registered for op '{self.name}'. "
336 + f"Use @ir.ops.{self.name}.register_input_generator"
337 + )
338 + return self._input_generator(**kwargs)
339 +
340 + def override_tolerance(
341 + self, dtype: torch.dtype, *, atol: float, rtol: float
342 + ) -> None:
343 + self._tolerance_overrides[dtype] = {"atol": atol, "rtol": rtol}
344 +
345 + def get_tolerance(self, dtype: torch.dtype) -> dict[str, float]:
346 + if dtype in self._tolerance_overrides:
347 + return self._tolerance_overrides[dtype]
348 + if dtype in DEFAULT_TOLERANCES:
349 + return DEFAULT_TOLERANCES[dtype]
350 + raise ValueError(
351 + f"No tolerance defined for dtype {dtype} in op '{self.name}'. "
352 + f"Use op.override_tolerance({dtype}, atol=..., rtol=...) "
353 + f"or add {dtype} to DEFAULT_TOLERANCES."
354 + )
355 +
356
357 class IrOpImpl:
358 def __init__(
vllm/ir/ops/layernorm.py
+15
@@ -19,3 +19,18 @@ def rms_norm(
19 if weight is not None:
20 x = x.to(weight.dtype) * weight
21 return x.to(orig_dtype)
22 +
23 +
24 +@rms_norm.register_input_generator
25 +def _rms_norm_input_generator(
26 + num_tokens: int, hidden_size: int, dtype: torch.dtype, epsilon: float = 1e-5
27 +) -> tuple:
28 + x = torch.randn(num_tokens, hidden_size, dtype=dtype)
29 + weight = torch.randn(hidden_size, dtype=dtype)
30 + return (x, weight, epsilon)
31 +
32 +
33 +# Reductions in rms_norm accumulate rounding error at large shapes
34 +# (e.g. 32768x16384), causing a few elements out of millions to exceed
35 +# the default float16 tolerance.
36 +rms_norm.override_tolerance(torch.float16, atol=1e-2, rtol=2e-3)
vllm/ir/tolerances.py new
+36
@@ -0,0 +1,36 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +ToleranceSpec = dict[torch.dtype, dict[str, float]]
7 +
8 +# Default tolerances for comparing IR op implementations against native.
9 +# These are intentionally conservative (permissive) to avoid false failures
10 +# across different hardware and kernel implementations. Ops that need tighter
11 +# or looser bounds should use override_tolerance.
12 +DEFAULT_TOLERANCES: ToleranceSpec = {
13 + # 52-bit mantissa; machine epsilon ~1.1e-16
14 + torch.float64: {"atol": 1e-8, "rtol": 1e-8},
15 + # 23-bit mantissa; machine epsilon ~1.2e-7.
16 + # Values from PyTorch test_transformers.py reference defaults.
17 + torch.float32: {"atol": 1e-5, "rtol": 1.3e-6},
18 + # 10-bit mantissa; machine epsilon ~9.8e-4.
19 + # Standard tolerance used across vLLM kernel tests.
20 + torch.float16: {"atol": 1e-3, "rtol": 1e-3},
21 + # 7-bit mantissa; machine epsilon ~7.8e-3.
22 + # Wider rtol than float16 to account for the coarser mantissa.
23 + torch.bfloat16: {"atol": 1e-3, "rtol": 1.6e-2},
24 + # 3-bit mantissa; machine epsilon ~6.25e-2.
25 + # Derived from vLLM fp8 kernel tests (merge_attn_states, silu_mul_fp8).
26 + torch.float8_e4m3fn: {"atol": 1e-1, "rtol": 1e-1},
27 + # 2-bit mantissa; machine epsilon ~1.25e-1.
28 + # Wider than e4m3fn due to the smaller mantissa.
29 + torch.float8_e5m2: {"atol": 2e-1, "rtol": 2e-1},
30 + # 1-bit mantissa; machine epsilon ~2.5e-1. Packed pair format (x2).
31 + # Derived from vLLM fp4 tests (test_silu_mul_nvfp4_quant: atol=3e-1).
32 + torch.float4_e2m1fn_x2: {"atol": 3e-1, "rtol": 3e-1},
33 + # Integer quantized; off-by-one from rounding is expected.
34 + # rtol=0 because relative error is meaningless for small integers.
35 + torch.int8: {"atol": 1, "rtol": 0},
36 +}