GPU execution feedback for better kernel generation.
An OpenEnv environment, a trained LoRA adapter, and the evidence from a small real-GPU experiment.
This Space publishes the environment source, Docker recipe and documentation. It does not execute submitted kernels, expose a live OpenEnv API, run model inference, or allocate paid resources. Run the downloaded environment on your own NVIDIA GPU.
Run the real environment
Download the source bundle, or clone the Space repository. You need Docker and NVIDIA Container Toolkit on a CUDA-capable host.
git clone https://huggingface.co/spaces/Shreyansh327/kernel-env cd kernel-env docker build -t kernel-env:latest . docker run --rm --gpus all -p 127.0.0.1:8000:8000 kernel-env:latest
The service starts on http://127.0.0.1:8000. Test it from another shell:
python -m pip install 'openenv-core[core]==0.3.0' ./kernel_env python examples/client_smoke.py --url http://127.0.0.1:8000
The smoke example resets a ReLU task, saves its prompt to task_prompt.txt, submits a real Triton implementation and requires a correct GPU result. The static Space URL cannot be used as this client's API endpoint.
Keep execution private and isolated: source checks and worker subprocesses are not a general security sandbox for native code. The release Dockerfile was statically checked; its full image was not built locally. Its evaluator source previously passed 116 real-GPU tests.
From a reference task to a policy update
OpenEnv server
WebSocket sessions, typed actions and observations, per-episode state and reward rubrics.
GPU evaluator
Fresh worker subprocesses, static checks, numerical controls, synchronized timing and timeouts.
Model trainer
The proof used a separate H100 for generation/training and another H100 for kernel evaluation.
Correctness-gated reward
incorrect kernel → reward 0 correct kernel → 0.3 + 2.7 × min(t_min / t_kernel, 1) episode return → reward₁ + 0.4 × reward₂
t_min estimates the hardware roofline floor. A correct kernel can earn reward without beating PyTorch. Feedback tokens are masked from the loss; generated tokens carry episode-level GRPO advantages. This pilot does not reproduce full staged actor–critic PPO or implement TRLOO.
What the small run showed
Correct tasks, matched greedy evaluation with two attempts.
Four recorded epochs, 64 training episodes, 128 training submissions.
Official held-out tasks solved; the new success was GELU.
| Task group | Before | After |
|---|---|---|
| Training: ReLU, LeakyReLU, Sigmoid, Tanh | 3/4 | 3/4 |
| Reserved: Swish, HardSigmoid, HardTanh | 2/3 | 2/3 |
| Held-out: Softmax, LogSoftmax, MinGPT GELU | 0/3 | 1/3 |
| Overall first turn | 5/10 | 5/10 |
The new GELU solution appeared on the second turn after compilation feedback. Its runtime was about 0.184 ms against 1.624 ms eager PyTorch (8.83× for this one task). It passed five additional signed-input seeds at the original shape. Existing solved kernels stayed essentially unchanged in speed.
One seed, ten tasks and related operator families do not establish statistically robust or unseen-family generalization. Pretraining contamination is not excluded. The generated erf GELU agrees with the tanh reference within tolerance, not bit for bit.
Inspect comparison data, all optimizer metrics and validation records →
Use the trained adapter
Qwen2.5-Coder-7B-KernelEnv-GRPO-16 contains the actual saved step-16 LoRA weights. Load with PEFT on the pinned Qwen2.5-Coder-7B-Instruct base; the model card includes code and the full protocol. It is an adapter, not standalone 7B weights or a hosted inference service.
LoRA rank 16, alpha 32; 16 Dr. GRPO updates; training temperature 0.8; 2,048 tokens per turn; seed 42. KernelBench revision 423217d9fda91e0c2d67e4a43bf62f96f6d104f1. Runtime provenance and weight hashes accompany the model.