@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b40d859c7b

[Kernel][Helion] Add numerics checks to benchmark script (#48968)

Signed-off-by: Shangdi Yu <shangdiy@meta.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Shangdi Yu committed Jul 31, 2026 at 17:16 UTC b40d859c7b07ae244bcd8c6eecdcdbd9a3afaa07

2 files changed +339 -10

scripts/benchmark_helion_kernels.py

+249 -10

@@ -5,8 +5,10 @@

5 Benchmark a registered Helion kernel against a baseline.

7 For each input case produced by the kernel's registered input generator, this

8 -measures the latency of the Helion kernel and a chosen baseline, then reports

9 -the speedup.

8 +checks the Helion kernel's numerics once against its eager reference, measures

9 +its latency against a chosen performance baseline, then reports the speedup.

10 +Use ``--numerics-with-perf-baseline`` to check numerics against the performance

11 +baseline instead.

13 Two baselines are supported (``--baseline``):

@@ -30,9 +32,17 @@ Usage:

32 python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\

33 --baseline cuda

35 + # Check numerics against the performance baseline instead of eager

36 + python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\

37 + --baseline cuda --numerics-with-perf-baseline

38 +

39 # Disable CUDA graph capture and save results

40 python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\

41 --no-cudagraph --output results.json

42 +

43 + # Only verify numerics, skipping the timing runs

44 + python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\

45 + --numerics-only

46 """

48 import argparse

@@ -43,15 +53,19 @@ import statistics

53 import sys

54 from collections.abc import Callable

55 from dataclasses import asdict, dataclass

56 +from typing import Any

58 import torch

59 +from torch.utils._pytree import tree_flatten

61 from vllm.triton_utils import triton

63 try:

64 + from helion.autotuner.accuracy import assert_close as helion_assert_close

65 + from helion.autotuner.accuracy import is_fp8_dtype

66 +

67 from vllm.benchmarks.lib.utils import default_vllm_config

68 from vllm.kernels.helion import get_kernel_by_name, get_registered_kernels

54 - from vllm.kernels.helion.ops import import_all_kernels

69 from vllm.logger import init_logger

70 from vllm.utils.import_utils import has_helion

71 except ImportError as e:

@@ -59,6 +73,27 @@ except ImportError as e:

73 print("Please ensure vLLM is installed and in your Python path")

74 sys.exit(1)

76 +

77 +def import_all_kernels() -> None:

78 + """Trigger Helion op registration, tolerating cross-version name drift.

79 +

80 + Current vLLM registers every Helion kernel as a side effect of importing

81 + ``vllm.kernels.helion.ops``; some builds instead expose an explicit importer

82 + whose name has drifted (``import_all_kernels`` / ``import_all_ops``). Call

83 + whichever exists; if none does, importing the module already registered

84 + them.

85 + """

86 + try:

87 + import vllm.kernels.helion.ops as ops

88 + except ImportError:

89 + return

90 + for fn_name in ("import_all_kernels", "import_all_ops"):

91 + fn = getattr(ops, fn_name, None)

92 + if callable(fn):

93 + fn()

94 + return

95 +

96 +

97 logger = init_logger("vllm.scripts.benchmark_helion_kernels")

@@ -121,6 +156,21 @@ def print_table(rows: list[Row]) -> None:

156 print(fmt(row))

157

158

159 +def log_versions() -> None:

160 + """Log torch/helion/triton versions at the head of the output."""

161 + from importlib.metadata import PackageNotFoundError, version

162 +

163 + def pkg_version(name: str) -> str:

164 + try:

165 + return version(name)

166 + except PackageNotFoundError:

167 + return "not installed"

168 +

169 + logger.info("torch: %s", torch.__version__)

170 + logger.info("helion: %s", pkg_version("helion"))

171 + logger.info("triton: %s", getattr(triton, "__version__", pkg_version("triton")))

172 +

173 +

174 def list_kernels() -> None:

175 kernels = get_registered_kernels()

176

@@ -176,12 +226,8 @@ def make_cuda_baseline(kernel_name: str) -> Callable:

226 return cuda_op

227

228

179 -def make_autotune_baseline(kernel_name: str) -> Callable:

180 - """Return the kernel's autotune baseline wrapped in ``torch.compile``.

181 -

182 - The baseline is the native-torch reference the kernel is tuned against,

183 - registered via ``helion_settings.autotune_baseline_fn``.

184 - """

229 +def make_eager_baseline(kernel_name: str) -> Callable:

230 + """Return the kernel's registered native-torch reference."""

231 wrapper = get_kernel_by_name(kernel_name)

232 settings = wrapper.helion_settings

233 baseline_fn = getattr(settings, "autotune_baseline_fn", None)

@@ -195,6 +241,13 @@ def make_autotune_baseline(kernel_name: str) -> Callable:

241 )

242 sys.exit(1)

243

244 + return baseline_fn

245 +

246 +

247 +def make_autotune_baseline(kernel_name: str) -> Callable:

248 + """Return the kernel's autotune baseline wrapped in ``torch.compile``."""

249 + baseline_fn = make_eager_baseline(kernel_name)

250 +

251 return torch.compile(

252 baseline_fn,

253 fullgraph=True,

@@ -204,6 +257,25 @@ def make_autotune_baseline(kernel_name: str) -> Callable:

257 )

258

259

260 +def make_correctness_baseline(

261 + kernel_name: str,

262 + timed_baseline_fn: Callable,

263 + numerics_with_perf_baseline: bool,

264 +) -> Callable:

265 + """Choose the numerical reference independently from the timed baseline."""

266 + if not numerics_with_perf_baseline:

267 + logger.info(

268 + "Using the eager reference for '%s' correctness",

269 + kernel_name,

270 + )

271 + return make_eager_baseline(kernel_name)

272 + logger.info(

273 + "Using the selected performance baseline for '%s' correctness",

274 + kernel_name,

275 + )

276 + return timed_baseline_fn

277 +

278 +

279 def cleanup_gpu_resources() -> None:

280 try:

281 torch.accelerator.empty_cache()

@@ -227,6 +299,118 @@ def _reduce(times: list[float], return_mode: str) -> float:

299 return _REDUCERS[return_mode](times)

300

301

302 +def _assert_close(actual: object, expected: object, atol: float, rtol: float) -> None:

303 + """Compare pytrees, allowing the one-ULP FP8 variance used by kernel tests."""

304 + actual_flat, actual_spec = tree_flatten(actual)

305 + expected_flat, expected_spec = tree_flatten(expected)

306 + if actual_spec != expected_spec:

307 + raise AssertionError(

308 + f"Output structure mismatch: {actual_spec} != {expected_spec}"

309 + )

310 +

311 + for actual_leaf, expected_leaf in zip(actual_flat, expected_flat, strict=True):

312 + is_fp8 = isinstance(actual_leaf, torch.Tensor) and is_fp8_dtype(

313 + actual_leaf.dtype

314 + )

315 + helion_assert_close(

316 + actual_leaf,

317 + expected_leaf,

318 + atol=1 if is_fp8 else atol,

319 + rtol=0 if is_fp8 else rtol,

320 + )

321 +

322 +

323 +def check_correctness(

324 + kernel: Any,

325 + baseline_fn: Callable,

326 + inputs: tuple[Any, ...],

327 + case: str,

328 +) -> None:

329 + """Run one numerical comparison on copies separate from benchmark inputs."""

330 + kernel_inputs = copy.deepcopy(inputs)

331 + baseline_inputs = copy.deepcopy(inputs)

332 +

333 + kernel_output = kernel(*kernel_inputs)

334 + baseline_output = baseline_fn(*baseline_inputs)

335 +

336 + settings = kernel.helion_settings

337 + try:

338 + custom_check = getattr(settings, "autotune_baseline_accuracy_check_fn", None)

339 + if custom_check is not None:

340 + custom_check(kernel_output, baseline_output)

341 + custom_check(kernel_inputs, baseline_inputs)

342 + return

343 +

344 + configured_atol = getattr(settings, "autotune_baseline_atol", None)

345 + configured_rtol = getattr(settings, "autotune_baseline_rtol", None)

346 + atol = 1e-2 if configured_atol is None else configured_atol

347 + rtol = 1e-2 if configured_rtol is None else configured_rtol

348 + _assert_close(

349 + kernel_output,

350 + baseline_output,

351 + atol=atol,

352 + rtol=rtol,

353 + )

354 + _assert_close(

355 + kernel_inputs,

356 + baseline_inputs,

357 + atol=atol,

358 + rtol=rtol,

359 + )

360 + except AssertionError as e:

361 + raise AssertionError(f"Numerics check failed for case {case}:\n{e}") from e

362 +

363 +

364 +@dataclass

365 +class CorrectnessResult:

366 + """Outcome of the numerics check for a single shape case."""

367 +

368 + case: str

369 + passed: bool

370 + error: str | None = None

371 +

372 +

373 +def check_kernel_correctness(

374 + kernel: Any,

375 + baseline_fn: Callable,

376 + inputs_dict: dict[Any, tuple[Any, ...]] | None = None,

377 +) -> list[CorrectnessResult]:

378 + """Run the per-shape numerics check for a kernel, continuing past failures.

379 +

380 + Runs the same comparison as ``check_correctness`` for every shape case

381 + produced by the kernel's input generator, but records the outcome per case

382 + instead of raising on the first mismatch. This lets callers (e.g. a CI gate)

383 + report every failing shape in one pass rather than aborting early.

384 +

385 + Args:

386 + kernel: The Helion kernel wrapper to check.

387 + baseline_fn: Reference callable sharing the kernel's argument interface.

388 + inputs_dict: Optional mapping of case key to input tuple. Defaults to

389 + ``kernel.get_inputs()``.

390 +

391 + Returns:

392 + One ``CorrectnessResult`` per shape case, in iteration order. A case that

393 + raises (compile/run error or numerics mismatch) is marked

394 + ``passed=False`` with the exception text in ``error``; iteration

395 + continues regardless. An empty input mapping returns an empty list, which

396 + the CLI reports as a skipped check.

397 + """

398 + if inputs_dict is None:

399 + inputs_dict = kernel.get_inputs()

400 +

401 + results: list[CorrectnessResult] = []

402 + for key, inputs in inputs_dict.items():

403 + case = str(key)

404 + try:

405 + check_correctness(kernel, baseline_fn, inputs, case)

406 + except Exception as e: # noqa: BLE001 - any failure is recorded, not fatal

407 + results.append(CorrectnessResult(case=case, passed=False, error=str(e)))

408 + else:

409 + results.append(CorrectnessResult(case=case, passed=True))

410 + cleanup_gpu_resources()

411 + return results

412 +

413 +

414 def do_bench_cudagraph_l2_clear(

415 fn: Callable, rep: int = 100, return_mode: str = "mean"

416 ) -> float:

@@ -298,6 +482,7 @@ def do_bench_cudagraph_l2_clear(

482 def benchmark(

483 kernel_name: str,

484 baseline_fn: Callable,

485 + correctness_fn: Callable,

486 repeat: int,

487 cudagraph: bool,

488 return_mode: str,

@@ -313,6 +498,9 @@ def benchmark(

498 for key, inputs in inputs_dict.items():

499 logger.info("Benchmarking case %s", key)

500

501 + check_correctness(kernel, correctness_fn, inputs, str(key))

502 + logger.info("Numerics check passed for case %s", key)

503 +

504 # Kernels may mutate their inputs in place; give each side its own copy.

505 kernel_inputs = copy.deepcopy(inputs)

506 baseline_inputs = copy.deepcopy(inputs)

@@ -374,7 +562,7 @@ def main() -> None:

562 choices=["cuda", "autotune"],

563 default="autotune",

564 help=(

377 - "Baseline to compare against: 'autotune' uses the kernel's "

565 + "Performance baseline: 'autotune' uses the kernel's "

566 "autotune_baseline_fn under torch.compile; 'cuda' uses the mapped "

567 "torch.ops._C op (default: autotune)"

568 ),

@@ -390,9 +578,24 @@ def main() -> None:

578 type=str,

579 help="Path to save benchmark results as JSON (default: log only)",

580 )

581 + parser.add_argument(

582 + "--numerics-only",

583 + action="store_true",

584 + help="Only run the per-case numerics check; skip timing and reporting",

585 + )

586 + parser.add_argument(

587 + "--numerics-with-perf-baseline",

588 + action="store_true",

589 + help=(

590 + "Compare numerics against the selected performance baseline instead "

591 + "of the eager reference"

592 + ),

593 + )

594

595 args = parser.parse_args()

596

597 + log_versions()

598 +

599 import_all_kernels()

600

601 if args.list:

@@ -425,10 +628,43 @@ def main() -> None:

628 baseline_fn = make_cuda_baseline(args.kernel)

629 else:

630 baseline_fn = make_autotune_baseline(args.kernel)

631 + correctness_fn = make_correctness_baseline(

632 + args.kernel,

633 + baseline_fn,

634 + args.numerics_with_perf_baseline,

635 + )

636 +

637 + if args.numerics_only:

638 + results = check_kernel_correctness(wrapper, correctness_fn)

639 + if not results:

640 + logger.warning(

641 + "No input cases generated for '%s'; skipping numerics check",

642 + args.kernel,

643 + )

644 + return

645 + for r in results:

646 + if r.passed:

647 + logger.info("Numerics check passed for case %s", r.case)

648 + else:

649 + logger.error(

650 + "Numerics check FAILED for case %s: %s", r.case, r.error

651 + )

652 + failed = [r for r in results if not r.passed]

653 + if failed:

654 + logger.error(

655 + "%d/%d case(s) failed numerics for '%s'",

656 + len(failed),

657 + len(results),

658 + args.kernel,

659 + )

660 + sys.exit(1)

661 + logger.info("Numerics check passed for all cases of '%s'", args.kernel)

662 + return

663

664 rows = benchmark(

665 args.kernel,

666 baseline_fn,

667 + correctness_fn,

668 args.repeat,

669 args.cudagraph,

670 args.return_mode,

@@ -442,6 +678,9 @@ def main() -> None:

678 {

679 "kernel": args.kernel,

680 "baseline": args.baseline,

681 + "numerics_baseline": (

682 + args.baseline if args.numerics_with_perf_baseline else "eager"

683 + ),

684 "cudagraph": args.cudagraph,

685 "repeat": args.repeat,

686 "return_mode": args.return_mode,

tests/kernels/helion/test_benchmark_script.py new

+90

@@ -0,0 +1,90 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import pytest

5 +import torch

6 +

7 +from vllm.utils.import_utils import has_helion

8 +

9 +if not has_helion():

10 + pytest.skip("Helion is not installed", allow_module_level=True)

11 +

12 +from scripts import benchmark_helion_kernels

13 +

14 +check_correctness = benchmark_helion_kernels.check_correctness

15 +

16 +

17 +class _FakeKernel:

18 + helion_settings = None

19 +

20 + def __init__(self, offset: float):

21 + self.calls = 0

22 + self.offset = offset

23 +

24 + def __call__(self, output: torch.Tensor, input: torch.Tensor) -> torch.Tensor:

25 + self.calls += 1

26 + output.copy_(input + self.offset)

27 + return input * 2

28 +

29 +

30 +def test_check_correctness_runs_once_without_mutating_benchmark_inputs():

31 + kernel = _FakeKernel(offset=1)

32 + baseline_calls = 0

33 +

34 + def baseline(output: torch.Tensor, input: torch.Tensor) -> torch.Tensor:

35 + nonlocal baseline_calls

36 + baseline_calls += 1

37 + output.copy_(input + 1)

38 + return input * 2

39 +

40 + output = torch.zeros(4)

41 + inputs = (output, torch.arange(4))

42 +

43 + check_correctness(kernel, baseline, inputs, "matching")

44 +

45 + assert kernel.calls == 1

46 + assert baseline_calls == 1

47 + torch.testing.assert_close(output, torch.zeros(4))

48 +

49 +

50 +def test_check_correctness_reports_mutated_input_mismatch():

51 + kernel = _FakeKernel(offset=2)

52 +

53 + def baseline(output: torch.Tensor, input: torch.Tensor) -> torch.Tensor:

54 + output.copy_(input + 1)

55 + return input * 2

56 +

57 + inputs = (torch.zeros(4), torch.arange(4))

58 +

59 + with pytest.raises(AssertionError, match="Numerics check failed for case bad"):

60 + check_correctness(kernel, baseline, inputs, "bad")

61 +

62 +

63 +def test_check_correctness_reports_return_value_mismatch():

64 + kernel = _FakeKernel(offset=1)

65 +

66 + def baseline(output: torch.Tensor, input: torch.Tensor) -> torch.Tensor:

67 + output.copy_(input + 1)

68 + return input * 3

69 +

70 + inputs = (torch.zeros(4), torch.arange(4))

71 +

72 + with pytest.raises(AssertionError, match="Numerics check failed for case bad"):

73 + check_correctness(kernel, baseline, inputs, "bad")

74 +

75 +

76 +def test_log_versions(monkeypatch):

77 + messages = []

78 +

79 + def record(message: str, *args):

80 + messages.append(message % args)

81 +

82 + monkeypatch.setattr(benchmark_helion_kernels.logger, "info", record)

83 +

84 + benchmark_helion_kernels.log_versions()

85 +

86 + assert [message.split(":", 1)[0] for message in messages] == [

87 + "torch",

88 + "helion",

89 + "triton",

90 + ]