Function run_performance

kernelbench/bench_kb.py:427–490

KernelBench-compatible performance benchmarking via CUDA event timing. Returns speedup = ref_time / kernel_time.

(
    model_new,
    model_ref,
    get_inputs_fn: Callable,
    n_warmup: int = DEFAULT_N_WARMUP,
    n_timed: int = DEFAULT_N_TIMED,
    device: str = "cuda",
)

# ---------------------------------------------------------------------------

def run_performance(
    model_new,
    model_ref,
    get_inputs_fn: Callable,
    n_warmup: int = DEFAULT_N_WARMUP,
    n_timed: int = DEFAULT_N_TIMED,
    device: str = "cuda",
) -> Dict[str, Any]:
    """
    KernelBench-compatible performance benchmarking via CUDA event timing.
    Returns speedup = ref_time / kernel_time.
    """
    import torch

    result: Dict[str, Any] = {
        "kernel_time_ms": 0.0,
        "reference_time_ms": 0.0,
        "speedup": 0.0,
        "kernel_times": [],
        "reference_times": [],
    }

    torch.manual_seed(0)
    inputs = get_inputs_fn()
    inputs_dev = [
        inp.to(device) if isinstance(inp, torch.Tensor) else inp
        for inp in inputs
    ]

    def _time_model(model, inputs_list, n_warm, n_iter):
        """Time a model using CUDA events."""
        for _ in range(n_warm):
            with torch.no_grad():
                model(*inputs_list)
        torch.cuda.synchronize()

        times = []
        for _ in range(n_iter):
            start = torch.cuda.Event(enable_timing=True)
            end = torch.cuda.Event(enable_timing=True)
            start.record()
            with torch.no_grad():
                model(*inputs_list)
            end.record()
            torch.cuda.synchronize()
            times.append(start.elapsed_time(end))
        return times

    try:
        ref_times = _time_model(model_ref, inputs_dev, n_warmup, n_timed)
        result["reference_times"] = ref_times
        result["reference_time_ms"] = _robust_median(ref_times)

        kernel_times = _time_model(model_new, inputs_dev, n_warmup, n_timed)
        result["kernel_times"] = kernel_times
        result["kernel_time_ms"] = _robust_median(kernel_times)

        if result["kernel_time_ms"] > 0: