Set up and evaluate a single KernelBench problem. Returns a result dict with correctness, speedup, etc.
(
level: int,
problem_id: int,
quick: bool = False,
backend: str = "cuda",
)
| 89 | |
| 90 | |
| 91 | def run_single_problem( |
| 92 | level: int, |
| 93 | problem_id: int, |
| 94 | quick: bool = False, |
| 95 | backend: str = "cuda", |
| 96 | ) -> Dict[str, Any]: |
| 97 | """ |
| 98 | Set up and evaluate a single KernelBench problem. |
| 99 | |
| 100 | Returns a result dict with correctness, speedup, etc. |
| 101 | """ |
| 102 | import subprocess |
| 103 | |
| 104 | result: Dict[str, Any] = { |
| 105 | "level": level, |
| 106 | "problem_id": problem_id, |
| 107 | "correctness": "FAIL", |
| 108 | "speedup": 0.0, |
| 109 | "kernel_time_ms": 0.0, |
| 110 | "reference_time_ms": 0.0, |
| 111 | "error": None, |
| 112 | } |
| 113 | |
| 114 | # Set up the problem |
| 115 | setup_cmd = [ |
| 116 | sys.executable, str(SCRIPT_DIR / "bridge.py"), "setup", |
| 117 | "--level", str(level), |
| 118 | "--problem", str(problem_id), |
| 119 | "--backend", backend, |
| 120 | "--source", "hf", |
| 121 | ] |
| 122 | try: |
| 123 | proc = subprocess.run( |
| 124 | setup_cmd, capture_output=True, text=True, timeout=120, |
| 125 | cwd=str(PROJECT_DIR), |
| 126 | ) |
| 127 | if proc.returncode != 0: |
| 128 | result["error"] = f"setup failed: {proc.stderr[:200]}" |
| 129 | return result |
| 130 | except subprocess.TimeoutExpired: |
| 131 | result["error"] = "setup timed out" |
| 132 | return result |
| 133 | except Exception as e: |
| 134 | result["error"] = f"setup error: {e}" |
| 135 | return result |
| 136 | |
| 137 | # Run bench_kb.py |
| 138 | bench_cmd = [ |
| 139 | sys.executable, str(SCRIPT_DIR / "bench_kb.py"), |
| 140 | "--skip-stability", "--skip-determinism", |
| 141 | ] |
| 142 | if quick: |
| 143 | bench_cmd.append("--quick") |
| 144 | |
| 145 | try: |
| 146 | proc = subprocess.run( |
| 147 | bench_cmd, capture_output=True, text=True, timeout=300, |
| 148 | cwd=str(PROJECT_DIR), |