Function run_correctness

kernelbench/bench_kb.py:255–327

KernelBench-compatible correctness checks. Generate n_trials random input sets, run both models, compare within (atol, rtol).

(
    model_new,
    model_ref,
    get_inputs_fn: Callable,
    n_trials: int = DEFAULT_N_CORRECTNESS,
    atol: float = DEFAULT_ATOL,
    rtol: float = DEFAULT_RTOL,
    device: str = "cuda",
)

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

def run_correctness(
    model_new,
    model_ref,
    get_inputs_fn: Callable,
    n_trials: int = DEFAULT_N_CORRECTNESS,
    atol: float = DEFAULT_ATOL,
    rtol: float = DEFAULT_RTOL,
    device: str = "cuda",
) -> Dict[str, Any]:
    """
    KernelBench-compatible correctness checks.

    Generate n_trials random input sets, run both models, compare within (atol, rtol).
    """
    import torch

    results: Dict[str, Any] = {
        "correctness": "FAIL",
        "trials_passed": 0,
        "trials_total": n_trials,
        "worst_max_abs_error": 0.0,
        "worst_mean_abs_error": 0.0,
        "details": [],
    }

    for trial in range(n_trials):
        trial_info: Dict[str, Any] = {"trial": trial, "status": "FAIL"}
        try:
            with _Timeout(TIMEOUT_SECONDS, f"trial {trial} timed out"):
                inputs = get_inputs_fn()
                inputs_dev = [
                    inp.to(device) if isinstance(inp, torch.Tensor) else inp
                    for inp in inputs
                ]

                with torch.no_grad():
                    expected = model_ref(*inputs_dev)
                with torch.no_grad():
                    output = model_new(*inputs_dev)

                cmp = _compare_outputs(output, expected, atol, rtol)
                trial_info["max_abs_error"] = cmp.get("max_abs_error", float("inf"))
                trial_info["mean_abs_error"] = cmp.get("mean_abs_error", float("inf"))

                if cmp["match"]:
                    trial_info["status"] = "PASS"
                    results["trials_passed"] += 1
                    results["worst_max_abs_error"] = max(
                        results["worst_max_abs_error"], cmp.get("max_abs_error", 0)
                    )
                    results["worst_mean_abs_error"] = max(
                        results["worst_mean_abs_error"], cmp.get("mean_abs_error", 0)
                    )
                else:
                    trial_info["reason"] = cmp["reason"]

        except TimeoutError as e:
            trial_info["status"] = "TIMEOUT"