Skip to content

KernelForge-OpenEnv

KernelForge-OpenEnv is an autonomous CUDA kernel optimization system that uses reinforcement learning to train language models to write high-performance GPU kernels for NVIDIA A100s.

Code Demo Video Cerebral Valley

Core trick

Model-written CUDA kernels compile and benchmark on A100s, and GRPO reward is discrete: 1 only if the kernel beats the PyTorch baseline.

Tags

dev tools

Description

KernelForge-OpenEnv is an autonomous CUDA kernel optimization system that uses reinforcement learning to train language models to write high-performance GPU kernels for NVIDIA A100s. Built on Meta-PyTorch's OpenEnv framework, it creates a closed-loop pipeline where a model generates CUDA code, compiles and benchmarks it on real hardware, then learns from the results across multi-turn episodes. Writing optimized CUDA kernels is one of the hardest tasks in GPU programming, requiring deep expertise in memory hierarchies, warp scheduling, and architecture-specific tuning. KernelForge automates this by bootstrapping from expert-curated datasets and progressively training through a 3-stage curriculum using GRPO, with discrete rewards based on whether generated kernels beat PyTorch baselines. It effectively teaches a language model to become a CUDA performance engineer through trial-and-error on real GPUs.
OCWC22/A100-CUDA-RL★ 01 forksPythonCreated Mar 8, 2026 · last push Mar 11, 2026

README excerpt

# KernelForge-OpenEnv

KernelForge is a CUDA-kernel RL repo built around a **custom lightweight OpenEnv wrapper** over a **KernelGYM-style execution backend** for reward-bearing kernel evaluation.

The current hackathon posture is:

- **Training** on A100 80GB
- **Reward-bearing eval** on A100 80GB
- **Active backend**: Modal for cloud bring-up, `local` for same-process smoke checks
- **Optional backend**: CoreWeave/Northflank HTTP service when explicitly configured

References:
- [OpenEnv overview](https://meta-pytorch.github.io/OpenEnv/)
- [OpenEnv core/runtime docs](https://meta-pytorch.github.io/OpenEnv/core/)
- [TRL OpenEnv integration](https://huggingface.co/docs/trl/en/openenv)
- [KernelGYM repo](https://github.com/hkust-nlp/KernelGYM)
- [Dr. Kernel paper](https://arxiv.org/abs/2602.05885)
- [Northflank GPU workloads](https://northflank.com/docs/v1/application/gpu-workloads/gpus-on-northflank)
- [CoreWeave on Northflank](https://northflank.com/docs/v1/application/bring-your-own-cloud/coreweave-on-northflank)

## Architecture

KernelForge should be understood as three layers:

1. **Custom lightweight OpenEnv wrapper**
   - `openenv_env/`
   - standardized `reset()` / `step()` / `state()` surface for judges, demos, and OpenEnv-compatible integrations
   - backed by `openenv.yaml` and `openenv_env/server/app.py`

2. **KernelForge task / reward / rollout logic**
   - `training/task_support.py`
   - `training/multi_turn_rollout.py`
   - `openenv_env/reward.py`
   - `trainin

Team