kernel_env/README.md

metadata

title: Kernel Env
emoji: ⚡
sdk: docker
app_port: 8000
tags:
  - openenv

kernel_env

OpenEnv environment that trains agents to write fast, correct GPU kernels (Triton by default, CUDA via load_inline) for KernelBench tasks.

Episode

1. reset(seed=?, task_id=?, split="train"|"eval", language="triton") → observation.problem holds the PyTorch reference Model and the rules.
2. step(KernelAction(code=...)) (up to max_turns, default 4) → static checks → compile → 5 correctness trials (default atol=rtol=1e-2, normalization checks below) → timing vs eager PyTorch. observation.stage / feedback say how far it got and why.

Reward per step: 0 unless correct, else 0.3 + clamp(log2(speedup), -0.2, 3.0), or with KERNEL_ENV_REWARD=roofline, 0.3 + 2.7 · min(t_min / t_new, 1).

Roofline (server/roofline.py): FLOPs counted on the reference (torch.utils.flop_counter, matmul/conv/attention), compulsory bytes (inputs + parameters + outputs), dense datasheet peaks per GPU (H100, H200, A100, L40S, L4). t_min = max(FLOPs / peak, bytes / bandwidth); fp32 tasks use the TF32 tensor-core roof, since tl.dot uses TF32 and passes the tolerance. Observations carry bound, arithmetic_intensity, t_min_ms, roofline_frac, ref_roofline_frac; timings under 0.8 × t_min are rejected as impossible.

For level1/23_Softmax and level1/37_FrobeniusNorm_, the old absolute tolerance was larger than the outputs. These tasks now use atol=min(configured_atol, max(abs(reference)) * 1e-5), retain the configured relative tolerance, reject nonfinite values, and check normalization: nonnegative Softmax rows sum to one; Frobenius output has one global L2 norm of one. Correctness trials also include constant, row-dependent peaked Softmax, signed/scaled and sparse Frobenius inputs. Timing still uses the original input generator. The server supplies the task id to the worker; direct Evaluator.evaluate(...) calls must pass task_id to select these extra checks. Other tasks keep their existing tolerances and TF32 policy.

ReLU, LeakyReLU, ELU, Softplus and MinGPTNewGelu correctness trials also include signed inputs in [-32,32] and zeros at the original shape. Their generators otherwise provide nonnegative inputs, which let identity copies pass several activations. Prompts disclose the broader correctness inputs; timing retains the original generator.

PyTorch tensor operations inside candidate forward are checked at runtime during both correctness and timing. Allocations, layout changes and device-local copies are allowed; tensor arithmetic, reductions, scalar extraction and copying inputs off the evaluation device are rejected as format violations. CPU-only autotuner statistics remain allowed. This dispatch guard is not a native-code sandbox: it does not inspect arbitrary custom extensions or graphs captured in __init__.

Run

GPU host (NVIDIA + nvidia-container-toolkit):

docker build -t kernel_env:latest -f server/Dockerfile .
docker run --gpus all -p 8000:8000 -e KERNEL_ENV_GPUS=0 kernel_env:latest

Modal (see ../modal_app.py): modal run modal_app.py::gpu_tests runs the GPU tests; modal deploy modal_app.py serves the env with Modal proxy auth (one eval GPU per container, autoscaled). Create a proxy auth token in the Modal dashboard (Settings → Proxy Auth Tokens) and export MODAL_PROXY_KEY / MODAL_PROXY_SECRET; KernelEnv sends them as Modal-Key / Modal-Secret.

No GPU (Mac dev): mock evaluator — real static checks, fake timings.

git clone https://github.com/ScalingIntelligence/KernelBench /tmp/KernelBench
KERNEL_ENV_EVALUATOR=mock KERNELBENCH_DIR=/tmp/KernelBench/KernelBench uv run server

Client:

from kernel_env import KernelAction, KernelEnv

with KernelEnv(base_url="http://localhost:8000").sync() as env:
    obs = env.reset(seed=0).observation
    result = env.step(KernelAction(code=my_source))
    print(result.observation.stage, result.observation.feedback, result.reward)

Config (env vars)

Var	Default
KERNELBENCH_DIR	/opt/KernelBench/KernelBench in image	task root (level1/, level2/, level3/)
KERNEL_ENV_GPUS	0	comma list of GPU ids used for evaluation (one eval per GPU at a time)
KERNEL_ENV_EVALUATOR	gpu	mock for machines without CUDA
KERNEL_ENV_MAX_TURNS	4	submissions per episode
KERNEL_ENV_MAX_SESSIONS	64	concurrent WebSocket sessions
KERNEL_ENV_EVAL_LOG	unset	also append the per-reset/step JSON log lines (always on stdout) to this file
KERNEL_ENV_REWARD	speedup	roofline: reward share of hardware roof reached instead of speedup vs PyTorch
KERNEL_ENV_ROOFLINE_REJECT	0.8	reject timings below this fraction of the roofline floor (physically impossible)

Keep evaluation GPUs separate from training GPUs, or timings (and the reward) get noisy.

Tests

uv run --extra dev pytest tests            # Mac-safe; GPU worker tests skip without CUDA

test_worker_validation.py and test_worker_operations.py also run on a CPU-only PyTorch install (and skip if PyTorch is absent). They validate the shared comparison/dispatch logic; real Triton, CUDA extensions, autotuning and timing overhead still require the isolated GPU test suite.