Spaces:
Running
Running
| 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](https://github.com/ScalingIntelligence/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): | |
| ```bash | |
| 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. | |
| ```bash | |
| git clone https://github.com/ScalingIntelligence/KernelBench /tmp/KernelBench | |
| KERNEL_ENV_EVALUATOR=mock KERNELBENCH_DIR=/tmp/KernelBench/KernelBench uv run server | |
| ``` | |
| Client: | |
| ```python | |
| 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 | |
| ```bash | |
| 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. | |