Spaces:
Running
Running
| import pytest | |
| from kernel_env import KernelAction | |
| from kernel_env.server.evaluator import MockEvaluator, apply_roofline | |
| from kernel_env.server.kernel_env_environment import KernelEnvironment | |
| from kernel_env.server.roofline import peaks, roofline | |
| from kernel_env.server.tasks import load_tasks | |
| from kernels import GOOD | |
| from test_environment import FIXTURES | |
| H100 = "NVIDIA H100 80GB HBM3" | |
| RELU = dict(flops=0, nbytes=2 * 4096 * 393216 * 4, dtype="torch.float32") # KernelBench level1/19 at full size | |
| def test_peak_lookup_prefers_specific_names(): | |
| assert peaks("NVIDIA L40S")[3] == 0.864 and peaks("NVIDIA L4")[3] == 0.3 | |
| assert peaks("NVIDIA H100 PCIe")[3] == 2.0 and peaks(H100)[3] == 3.35 | |
| assert peaks("Tesla T4") is None | |
| def test_relu_is_memory_bound(): | |
| r = roofline(RELU["flops"], RELU["nbytes"], RELU["dtype"], H100, t_ref_ms=4.2, t_new_ms=4.0) | |
| assert r["bound"] == "memory" and r["t_min_ms"] == pytest.approx(3.846, rel=1e-3) | |
| assert r["roofline_frac"] == pytest.approx(3.846 / 4.0, rel=1e-3) | |
| def test_square_matmul_is_compute_bound_on_tf32(): | |
| n = 4096 | |
| r = roofline(2 * n**3, 3 * n * n * 4, "torch.float32", H100, t_ref_ms=2.5) | |
| assert r["bound"] == "compute" and r["arithmetic_intensity"] == pytest.approx(682.7, rel=1e-3) | |
| assert r["t_min_ms"] == pytest.approx(2 * n**3 / 494.7e12 * 1e3, rel=1e-6) | |
| def timed(t_new, **kw): | |
| return {"stage": "timed", "correct": True, "speedup": 4.2 / t_new, "t_ref_ms": 4.2, "t_new_ms": t_new, | |
| "gpu": H100, "flops": 0, "bytes": RELU["nbytes"], "dtype": "torch.float32", "feedback": "Correct.", **kw} | |
| def test_apply_roofline_adds_feedback(): | |
| res = apply_roofline(timed(4.0)) | |
| assert res["stage"] == "timed" and "memory-bound" in res["feedback"] and "yours 96%" in res["feedback"] | |
| def test_faster_than_hardware_is_rejected(): | |
| res = apply_roofline(timed(1.0)) | |
| assert res["stage"] == "correctness" and not res["correct"] and res["roofline_violation"] | |
| assert "below the hardware floor" in res["feedback"] | |
| def test_unknown_gpu_leaves_result_alone(): | |
| res = timed(1.0, gpu="Tesla T4") | |
| assert apply_roofline(res) == res | |
| def test_reward_modes(monkeypatch, mode, reward): | |
| monkeypatch.setenv("KERNEL_ENV_REWARD", mode) | |
| env = KernelEnvironment(MockEvaluator([timed(4.0)]), load_tasks(str(FIXTURES), levels=(1,))) | |
| env.reset(task_id="level1/19_ReLU") | |
| obs = env.step(KernelAction(code=GOOD)) | |
| assert obs.bound == "memory" and obs.roofline_frac == pytest.approx(0.9615, rel=1e-3) | |
| assert obs.reward == pytest.approx(reward, abs=0.01) | |