diff --git a/scripts/benchmark_helion_kernels.py b/scripts/benchmark_helion_kernels.py new file mode 100644 index 00000000000..a2640aa52ff --- /dev/null +++ b/scripts/benchmark_helion_kernels.py @@ -0,0 +1,457 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Benchmark a registered Helion kernel against a baseline. + +For each input case produced by the kernel's registered input generator, this +measures the latency of the Helion kernel and a chosen baseline, then reports +the speedup. + +Two baselines are supported (``--baseline``): + +- ``autotune`` (default): the kernel's autotuning baseline + (``helion_settings.autotune_baseline_fn``), wrapped in ``torch.compile`` + (inductor). This is the native-torch reference used by kernel autotuning + and correctness unit tests. +- ``cuda``: the corresponding hand-written CUDA op (``torch.ops._C.*``). The + mapping from Helion kernel name to CUDA op lives in ``CUDA_BASELINE_OPS`` + below. Every Helion kernel shares the same argument interface as its CUDA + counterpart, so inputs are forwarded verbatim. + +Usage: + # List available kernels + python scripts/benchmark_helion_kernels.py --list + + # Benchmark a kernel against the autotune baseline (default) + python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant + + # Benchmark against the CUDA baseline + python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\ + --baseline cuda + + # Disable CUDA graph capture and save results + python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\ + --no-cudagraph --output results.json +""" + +import argparse +import copy +import gc +import json +import statistics +import sys +from collections.abc import Callable +from dataclasses import asdict, dataclass + +import torch + +from vllm.triton_utils import triton + +try: + from vllm.benchmarks.lib.utils import default_vllm_config + from vllm.kernels.helion import get_kernel_by_name, get_registered_kernels + from vllm.kernels.helion.ops import import_all_kernels + from vllm.logger import init_logger + from vllm.utils.import_utils import has_helion +except ImportError as e: + print(f"Error importing vLLM: {e}") + print("Please ensure vLLM is installed and in your Python path") + sys.exit(1) + +logger = init_logger("vllm.scripts.benchmark_helion_kernels") + + +# Maps a Helion kernel name to the CUDA op (attribute on ``torch.ops._C``) that +# implements the same operation. Helion kernels share the CUDA op's argument +# interface, so the kernel's input tuple is forwarded verbatim. Add an entry +# here when introducing a new kernel whose baseline should be the CUDA op. +CUDA_BASELINE_OPS: dict[str, str] = { + "dynamic_per_token_scaled_fp8_quant": "dynamic_per_token_scaled_fp8_quant", + "fused_qk_norm_rope": "fused_qk_norm_rope", + "per_token_group_fp8_quant": "per_token_group_fp8_quant", + "rms_norm_dynamic_per_token_quant": "rms_norm_dynamic_per_token_quant", + "rms_norm_per_block_quant": "rms_norm_per_block_quant", + "silu_and_mul_per_block_quant": "silu_and_mul_per_block_quant", + "scaled_mm": "cutlass_scaled_mm", +} + +# torch.compile options for the torch baseline, mirroring how these kernels are +# compiled inside vLLM. +_TORCH_COMPILE_OPTIONS: dict[str, bool] = { + "enable_auto_functionalized_v2": False, + "size_asserts": False, + "alignment_asserts": False, + "scalar_asserts": False, + "combo_kernels": True, + "benchmark_combo_kernel": True, +} + + +@dataclass +class Row: + case: str + baseline_ms: float + kernel_ms: float + speedup_x: float + + +def print_table(rows: list[Row]) -> None: + headers = ["case", "baseline_ms", "kernel_ms", "speedup(x)"] + + data = [ + [ + r.case, + f"{r.baseline_ms:.3f}", + f"{r.kernel_ms:.3f}", + f"{r.speedup_x:.3f}", + ] + for r in rows + ] + + cols = list(zip(*([headers] + data))) + widths = [max(len(cell) for cell in col) for col in cols] + + def fmt(row: list[str]) -> str: + return " | ".join(cell.ljust(w) for cell, w in zip(row, widths)) + + print(fmt(headers)) + print("-+-".join("-" * w for w in widths)) + for row in data: + print(fmt(row)) + + +def list_kernels() -> None: + kernels = get_registered_kernels() + + if not kernels: + print("No Helion kernels found in registry.") + return + + print("Available Helion kernels:") + print("=" * 50) + for name in sorted(kernels.keys()): + cuda = CUDA_BASELINE_OPS.get(name) + suffix = "" if cuda else " (no CUDA baseline mapping)" + print(f" {name}{suffix}") + print(f"\nTotal: {len(kernels)} kernels") + + +def check_requirements() -> bool: + if not torch.accelerator.is_available(): + logger.error("CUDA is not available. Helion benchmarking requires GPU.") + return False + if not has_helion(): + logger.error("Helion is not installed. Please install Helion package.") + return False + return True + + +def make_cuda_baseline(kernel_name: str) -> Callable: + """Return a callable invoking the CUDA op mapped to ``kernel_name``. + + The Helion kernel and its CUDA op share the same argument interface, so the + input tuple is forwarded verbatim. + """ + cuda_op_name = CUDA_BASELINE_OPS.get(kernel_name) + if cuda_op_name is None: + logger.error( + "No CUDA baseline mapping for kernel '%s'. Add an entry to " + "CUDA_BASELINE_OPS in %s (mapping the kernel name to its " + "torch.ops._C. name), or benchmark with --baseline torch.", + kernel_name, + __file__, + ) + sys.exit(1) + + cuda_op = getattr(torch.ops._C, cuda_op_name, None) + if cuda_op is None: + logger.error( + "torch.ops._C.%s is not available. Ensure the vLLM C extension is " + "built and loaded.", + cuda_op_name, + ) + sys.exit(1) + + return cuda_op + + +def make_autotune_baseline(kernel_name: str) -> Callable: + """Return the kernel's autotune baseline wrapped in ``torch.compile``. + + The baseline is the native-torch reference the kernel is tuned against, + registered via ``helion_settings.autotune_baseline_fn``. + """ + wrapper = get_kernel_by_name(kernel_name) + settings = wrapper.helion_settings + baseline_fn = getattr(settings, "autotune_baseline_fn", None) + if baseline_fn is None: + logger.error( + "Kernel '%s' has no autotune_baseline_fn in its helion_settings, so " + "the 'autotune' baseline is unavailable. Register one via " + "register_kernel(..., helion_settings=helion.Settings(" + "autotune_baseline_fn=...)), or benchmark with --baseline cuda.", + kernel_name, + ) + sys.exit(1) + + return torch.compile( + baseline_fn, + fullgraph=True, + dynamic=False, + backend="inductor", + options=_TORCH_COMPILE_OPTIONS, + ) + + +def cleanup_gpu_resources() -> None: + try: + torch.accelerator.empty_cache() + gc.collect() + if hasattr(torch, "_dynamo"): + torch._dynamo.reset() + torch.accelerator.synchronize() + except Exception as e: + logger.warning("Failed to cleanup GPU resources: %s", e) + + +_REDUCERS: dict[str, Callable[[list[float]], float]] = { + "min": min, + "max": max, + "mean": statistics.fmean, + "median": statistics.median, +} + + +def _reduce(times: list[float], return_mode: str) -> float: + return _REDUCERS[return_mode](times) + + +def do_bench_cudagraph_l2_clear( + fn: Callable, rep: int = 100, return_mode: str = "mean" +) -> float: + """CUDA-graph benchmark that flushes the L2 cache before every call. + + ``triton.testing.do_bench_cudagraph`` captures back-to-back kernel launches + with a warm L2 cache, which over-estimates performance for memory-bound + kernels. This clears L2 (via triton's benchmark cache buffer) before each + call and subtracts the isolated cache-clear cost from the measured time. + + Adapted from tritonbench's ``_do_bench_cudagraph_with_cache_clear`` using + only triton/torch primitives so no extra dependency is introduced. + """ + cache = triton.runtime.driver.active.get_empty_cache_for_benchmark() + clear_cache = cache.zero_ + + s = torch.Stream() + with s: + clear_cache() + fn() + + start_event = torch.Event(enable_timing=True) + end_event = torch.Event(enable_timing=True) + start_event.record() + for _ in range(5): + clear_cache() + fn() + end_event.record() + torch.accelerator.synchronize() + estimate_ms = start_event.elapsed_time(end_event) / 5 + n_repeat = 1000 if estimate_ms == 0 else max(1, int(rep / estimate_ms)) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(n_repeat): + clear_cache() + fn() + + clear_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(clear_graph): + for _ in range(n_repeat): + clear_cache() + torch.accelerator.synchronize() + + kernel_times = [] + for _ in range(10): + start_event = torch.Event(enable_timing=True) + end_event = torch.Event(enable_timing=True) + start_event.record() + clear_graph.replay() + end_event.record() + torch.accelerator.synchronize() + clear_ms = start_event.elapsed_time(end_event) / n_repeat + + start_event = torch.Event(enable_timing=True) + end_event = torch.Event(enable_timing=True) + start_event.record() + graph.replay() + end_event.record() + torch.accelerator.synchronize() + total_ms = start_event.elapsed_time(end_event) / n_repeat + + kernel_times.append(total_ms - clear_ms) + + return _reduce(kernel_times, return_mode) + + +@torch.inference_mode() +def benchmark( + kernel_name: str, + baseline_fn: Callable, + repeat: int, + cudagraph: bool, + return_mode: str, +) -> list[Row]: + kernel = get_kernel_by_name(kernel_name) + # do_bench already flushes L2 per call; do_bench_cudagraph does not, so use + # the cache-clearing variant to avoid warm-L2 over-estimates. + benchmark_fn = do_bench_cudagraph_l2_clear if cudagraph else triton.testing.do_bench + + inputs_dict = kernel.get_inputs() + rows: list[Row] = [] + + for key, inputs in inputs_dict.items(): + logger.info("Benchmarking case %s", key) + + # Kernels may mutate their inputs in place; give each side its own copy. + kernel_inputs = copy.deepcopy(inputs) + baseline_inputs = copy.deepcopy(inputs) + + kernel_latency = benchmark_fn( + lambda kernel_inputs=kernel_inputs: kernel(*kernel_inputs), + rep=repeat, + return_mode=return_mode, + ) + baseline_latency = benchmark_fn( + lambda baseline_inputs=baseline_inputs: baseline_fn(*baseline_inputs), + rep=repeat, + return_mode=return_mode, + ) + + rows.append( + Row( + case=str(key), + baseline_ms=baseline_latency, + kernel_ms=kernel_latency, + speedup_x=baseline_latency / kernel_latency, + ) + ) + cleanup_gpu_resources() + + return rows + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Benchmark a Helion kernel against a baseline", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__.split("Usage:")[1] if "Usage:" in __doc__ else "", + ) + parser.add_argument( + "--list", + action="store_true", + help="List available Helion kernels and exit", + ) + parser.add_argument( + "--kernel", + type=str, + help="Name of the single Helion kernel to benchmark", + ) + parser.add_argument( + "--repeat", + type=int, + default=100, + help="Number of benchmark repetitions (default: 100)", + ) + parser.add_argument( + "--no-cudagraph", + dest="cudagraph", + action="store_false", + help="Disable CUDA graph mode (enabled by default)", + ) + parser.add_argument( + "--baseline", + choices=["cuda", "autotune"], + default="autotune", + help=( + "Baseline to compare against: 'autotune' uses the kernel's " + "autotune_baseline_fn under torch.compile; 'cuda' uses the mapped " + "torch.ops._C op (default: autotune)" + ), + ) + parser.add_argument( + "--return-mode", + choices=["min", "max", "mean", "median"], + default="mean", + help="Statistic to report from the benchmark samples (default: mean)", + ) + parser.add_argument( + "--output", + type=str, + help="Path to save benchmark results as JSON (default: log only)", + ) + + args = parser.parse_args() + + import_all_kernels() + + if args.list: + list_kernels() + return + + if not args.kernel: + parser.error("--kernel is required (or use --list to see available kernels)") + + kernels = get_registered_kernels() + if args.kernel not in kernels: + logger.error("Kernel '%s' not found in registry.", args.kernel) + logger.error("Available kernels: %s", sorted(kernels.keys())) + sys.exit(1) + + wrapper = kernels[args.kernel] + if wrapper._disabled: + logger.error( + "Kernel '%s' is disabled: %s", + args.kernel, + wrapper._disabled_reason, + ) + sys.exit(1) + + if not check_requirements(): + sys.exit(1) + + with default_vllm_config(): + if args.baseline == "cuda": + baseline_fn = make_cuda_baseline(args.kernel) + else: + baseline_fn = make_autotune_baseline(args.kernel) + + rows = benchmark( + args.kernel, + baseline_fn, + args.repeat, + args.cudagraph, + args.return_mode, + ) + + print_table(rows) + + if args.output: + with open(args.output, "w") as f: + json.dump( + { + "kernel": args.kernel, + "baseline": args.baseline, + "cudagraph": args.cudagraph, + "repeat": args.repeat, + "return_mode": args.return_mode, + "results": [asdict(r) for r in rows], + }, + f, + indent=2, + ) + logger.info("Saved results to %s", args.output) + + +if __name__ == "__main__": + main() diff --git a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py index eef262dcfe2..45bd8f6fcf3 100644 --- a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py +++ b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py @@ -109,7 +109,18 @@ def baseline( scale: torch.Tensor, # [num_tokens, 1] scale_ub: torch.Tensor | None = None, # scalar tensor ) -> None: - torch.ops._C.dynamic_per_token_scaled_fp8_quant(result, input, scale, scale_ub) + fp8_min, fp8_max = get_fp8_min_max() + min_scaling_factor = 1.0 / (fp8_max * 512.0) + + x = input.to(torch.float32) + s = torch.amax(torch.abs(x), dim=-1, keepdim=True) + if scale_ub is not None: + s = s.clamp(max=scale_ub) + s = (s * (1.0 / fp8_max)).clamp(min=min_scaling_factor) + y = (x / s).clamp(fp8_min, fp8_max) + + scale.copy_(s) + result.copy_(y.to(result.dtype)) # Overwrite autotune_baseline_atol and autotune_baseline_rtol diff --git a/vllm/kernels/helion/ops/per_token_group_fp8_quant.py b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py index 8b73fac4b8e..42fd0215f59 100644 --- a/vllm/kernels/helion/ops/per_token_group_fp8_quant.py +++ b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py @@ -164,18 +164,17 @@ def baseline( dummy_is_scale_transposed: bool = False, dummy_is_tma_aligned: bool = False, ) -> None: - torch.ops._C.per_token_group_fp8_quant( - input, - output_q, - output_s, - group_size, - eps, - fp8_min, - fp8_max, - scale_ue8m0, - dummy_is_scale_transposed, - dummy_is_tma_aligned, - ) + num_tokens, hidden_size = input.shape + groups_per_row = hidden_size // group_size + + x = input.view(num_tokens, groups_per_row, group_size).to(torch.float32) + s = torch.clamp(torch.amax(torch.abs(x), dim=-1), min=eps) / fp8_max + if scale_ue8m0: + s = torch.exp2(torch.ceil(torch.log2(s))) + y = torch.clamp(x / s[:, :, None], fp8_min, fp8_max) + + output_s.copy_(s) + output_q.copy_(y.view(num_tokens, hidden_size).to(output_q.dtype)) @register_kernel( diff --git a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py index f15132c27cf..3e02169db3b 100644 --- a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py +++ b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py @@ -131,9 +131,37 @@ def baseline( scale_ub: torch.Tensor | None = None, # [] residual: torch.Tensor | None = None, # [num_tokens, hidden_size] ) -> None: - torch.ops._C.rms_norm_dynamic_per_token_quant( - result, input, weight, scale, epsilon, scale_ub, residual - ) + _, hidden_size = input.shape + quant_dtype = result.dtype + qtype_min: int | float + qtype_max: int | float + + if quant_dtype == torch.int8: + qtype_min, qtype_max = get_int8_min_max() + min_scaling_factor = get_int8_min_scaling_factor() + else: + qtype_min, qtype_max = get_fp8_min_max() + min_scaling_factor = 1.0 / (qtype_max * 512.0) + + x = input.to(torch.float32) + if residual is not None: + x = x + residual + residual.copy_(x.to(residual.dtype)) + + rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon) + x_norm = (x * rms).to(input.dtype) * weight + + s = torch.amax(torch.abs(x_norm), dim=-1, keepdim=True).to(torch.float32) + if scale_ub is not None: + s = s.clamp(max=scale_ub) + s = (s * (1.0 / qtype_max)).clamp(min=min_scaling_factor) + + y = x_norm / s + if quant_dtype == torch.int8: + y = y.round() + + scale.copy_(s) + result.copy_(y.clamp(qtype_min, qtype_max).to(result.dtype)) # Overwrite autotune_baseline_atol and autotune_baseline_rtol diff --git a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py index e7df42f4dbe..da7651a1fa5 100644 --- a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py +++ b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py @@ -169,16 +169,40 @@ def baseline( group_size: int, is_scale_transposed: bool, ) -> None: - torch.ops._C.rms_norm_per_block_quant( - result, - input, - weight, - scale, - epsilon, - scale_ub, - residual, - group_size, - is_scale_transposed, + num_tokens, hidden_size = input.shape + groups_per_row = hidden_size // group_size + quant_dtype = result.dtype + qtype_min: int | float + qtype_max: int | float + + if quant_dtype == torch.int8: + qtype_min, qtype_max = get_int8_min_max() + min_scaling_factor = get_int8_min_scaling_factor() + else: + qtype_min, qtype_max = get_fp8_min_max() + min_scaling_factor = 1.0 / (qtype_max * 512.0) + + x = input.to(torch.float32) + if residual is not None: + x = x + residual + residual.copy_(x.to(residual.dtype)) + + rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon) + x_norm = (x * rms).to(input.dtype) * weight + x_grouped = x_norm.view(num_tokens, groups_per_row, group_size).to(torch.float32) + + s = torch.amax(torch.abs(x_grouped), dim=-1).to(torch.float32) + if scale_ub is not None: + s = s.clamp(max=scale_ub) + s = (s * (1.0 / qtype_max)).clamp(min=min_scaling_factor) + + y = x_grouped / s[:, :, None] + if quant_dtype == torch.int8: + y = y.round() + + scale.copy_(s) + result.copy_( + y.clamp(qtype_min, qtype_max).view(num_tokens, hidden_size).to(result.dtype) ) diff --git a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py index f3aaf226b04..06b7f10af2f 100644 --- a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py +++ b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py @@ -13,6 +13,7 @@ from vllm.kernels.helion.utils import ( get_int8_min_scaling_factor, ) from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) @@ -150,8 +151,34 @@ def baseline( scale_ub: torch.Tensor | None = None, # scalar tensor is_scale_transposed: bool = False, ) -> None: - torch.ops._C.silu_and_mul_per_block_quant( - out, input, scales, group_size, scale_ub, is_scale_transposed + num_tokens, intermediate_size = out.shape + groups_per_row = intermediate_size // group_size + quant_dtype = out.dtype + qtype_min: int | float + qtype_max: int | float + + if quant_dtype == torch.int8: + qtype_min, qtype_max = get_int8_min_max() + min_scaling_factor = get_int8_min_scaling_factor() + else: + qtype_min, qtype_max = get_fp8_min_max() + min_scaling_factor = 1.0 / (qtype_max * 512.0) + + act = SiluAndMul.forward_native(input.to(torch.float32)) + x_grouped = act.view(num_tokens, groups_per_row, group_size) + + s = torch.amax(torch.abs(x_grouped), dim=-1) + if scale_ub is not None: + s = s.clamp(max=scale_ub) + s = (s * (1.0 / qtype_max)).clamp(min=min_scaling_factor) + + y = x_grouped / s[:, :, None] + if quant_dtype == torch.int8: + y = y.round() + + scales.copy_(s) + out.copy_( + y.clamp(qtype_min, qtype_max).view(num_tokens, intermediate_size).to(out.dtype) )