Spaces:
Running
Running
| """Real evaluation through worker.py. Runs only on a machine with CUDA + triton.""" | |
| from pathlib import Path | |
| import pytest | |
| torch = pytest.importorskip("torch") | |
| pytest.importorskip("triton") | |
| if not torch.cuda.is_available(): | |
| pytest.skip("needs a CUDA GPU", allow_module_level=True) | |
| from kernel_env.server.evaluator import Evaluator # noqa: E402 | |
| from kernels import ALIAS, CUDA_GOOD, GOOD, HANG, WRONG # noqa: E402 | |
| REF = (Path(__file__).parent / "fixtures/level1/19_ReLU.py").read_text() | |
| SEEDS = [1, 2, 3] | |
| def ev(): | |
| return Evaluator(gpus=["0"], timeout_s=120, warmup_ms=5, rep_ms=20) | |
| def test_good_kernel_is_correct_and_timed(ev): | |
| res = ev.evaluate(REF, GOOD, "triton", SEEDS) | |
| assert res["stage"] == "timed" and res["correct"], res | |
| assert 0.2 < res["speedup"] < 5.0 # ReLU is memory bound: roughly parity with PyTorch | |
| assert res["bound"] == "memory" and res["flops"] == 0 and res["bytes"] == 2 * 256 * 16384 * 4 | |
| assert 0.05 < res["roofline_frac"] <= 1.0 and 0.05 < res["ref_roofline_frac"] <= 1.0, res | |
| assert "Roofline on" in res["feedback"] | |
| def test_cuda_load_inline_kernel(ev): | |
| res = ev.evaluate(REF, CUDA_GOOD, "cuda", SEEDS) | |
| assert res["stage"] == "timed", res | |
| def test_wrong_output_rejected(ev): | |
| res = ev.evaluate(REF, WRONG, "triton", SEEDS) | |
| assert res["stage"] == "correctness" and "mismatch" in res["feedback"] | |
| def test_aliasing_input_rejected(ev): | |
| res = ev.evaluate(REF, ALIAS, "triton", SEEDS) | |
| assert res["stage"] == "correctness" and "aliases" in res["feedback"] | |
| def test_hang_times_out(): | |
| res = Evaluator(gpus=["0"], timeout_s=20, warmup_ms=5, rep_ms=20).evaluate(REF, HANG, "triton", SEEDS) | |
| assert res["stage"] == "timeout" | |
| def test_reference_vs_itself_is_stable(ev): | |
| # identical kernel timed repeatedly: speedup spread is the reward noise floor | |
| speedups = [ev.evaluate(REF, GOOD, "triton", SEEDS)["speedup"] for _ in range(5)] | |
| assert max(speedups) / min(speedups) < 1.15, speedups | |
| def test_matmul_roofline_is_compute_bound(ev): | |
| import os | |
| path = Path(os.environ.get("KERNELBENCH_DIR", "/nonexistent"), "level1/1_Square_matrix_multiplication_.py") | |
| if not path.exists(): | |
| pytest.skip("KernelBench not available") | |
| from kernels import MATMUL | |
| res = ev.evaluate(path.read_text(), MATMUL, "triton", SEEDS[:1]) | |
| print({k: res.get(k) for k in ("stage", "speedup", "t_ref_ms", "t_new_ms", "t_min_ms", "bound", | |
| "arithmetic_intensity", "roofline_frac", "ref_roofline_frac", "gpu")}) | |
| assert res["stage"] == "timed" and not res.get("roofline_violation"), res | |
| assert res["bound"] == "compute" and res["arithmetic_intensity"] > 100 | |
| assert 0 < res["roofline_frac"] <= 1.0 | |
| def test_shared_gpu_keeps_timings_clean(): | |
| # 2 workers share the GPU for compile/correctness; benchmarks take an exclusive lock | |
| from concurrent.futures import ThreadPoolExecutor | |
| shared = Evaluator(gpus=["0"], timeout_s=120, warmup_ms=5, rep_ms=20, workers_per_gpu=2) | |
| with ThreadPoolExecutor(4) as pool: | |
| res = list(pool.map(lambda s: shared.evaluate(REF, GOOD, "triton", [s, s + 1]), range(4))) | |
| print([r.get("phases") for r in res]) | |
| assert all(r["stage"] == "timed" for r in res), res | |
| assert all("bench" in r["phases"] and "lock_wait" in r["phases"] for r in res) | |
| speedups = [r["speedup"] for r in res] | |
| assert max(speedups) / min(speedups) < 1.15, speedups | |
| def test_very_slow_kernel_skips_benchmark(ev): | |
| from kernels import SLOW | |
| res = ev.evaluate(REF, SLOW, "triton", SEEDS) | |
| print(res.get("phases"), res.get("t_new_ms")) | |
| assert res["stage"] == "timed" and res["correct"] and res["speedup"] < 0.02, res | |
| assert "very slow" in res["feedback"] | |
| TOY_REF = """ | |
| import torch | |
| import torch.nn as nn | |
| class Model(nn.Module): | |
| def __init__(self, dim): | |
| super().__init__() | |
| self.scale = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x): | |
| return torch.relu(x * self.scale) | |
| def get_inputs(): | |
| return [torch.randn(64, 512)] | |
| def get_init_inputs(): | |
| return [512] | |
| """ | |
| def test_prompt_v2_example_is_correct(ev): | |
| from kernel_env.server.tasks import EXAMPLE_V2 | |
| res = ev.evaluate(TOY_REF, EXAMPLE_V2, "triton", SEEDS) | |
| assert res["stage"] == "timed" and res["correct"], res | |