Spaces:
Running
Running
File size: 4,352 Bytes
e313085 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """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]
@pytest.fixture(scope="module")
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
|