Upload blog_post.md with huggingface_hub

blog_post.md

@@ -0,0 +1,202 @@
+---
+title: "kernrl: Teaching LLMs to Write Fast GPU Kernels"
+thumbnail: /blog/assets/kernrl/thumbnail.png
+authors:
+- user: Infatoshi
+date: 2026-01-20
+tags:
+- openenv
+- reinforcement-learning
+- cuda
+- triton
+- gpu
+- grpo
+---
+
+# kernrl: Teaching LLMs to Write Fast GPU Kernels
+
+What if we could train language models to write optimized GPU code? Not just syntactically correct code, but kernels that actually run faster than PyTorch's defaults?
+
+That's the goal of **kernrl** - an RL environment where agents learn to optimize GPU kernels through trial and error, receiving real performance feedback from actual hardware.
+
+## The Problem: GPU Programming is Hard
+
+Writing efficient GPU code requires understanding memory hierarchies, thread synchronization, and hardware-specific optimizations. Even experienced engineers spend significant time tuning kernels for different architectures.
+
+Consider a simple softmax operation. PyTorch's implementation works, but a hand-tuned Triton kernel can be 2-5x faster by:
+- Fusing the max, subtract, exp, sum, and divide operations
+- Using efficient memory access patterns
+- Avoiding unnecessary global memory round-trips
+
+```python
+# PyTorch baseline - multiple kernel launches
+def softmax(x):
+    max_val = x.max(dim=-1, keepdim=True)
+    x = x - max_val
+    exp_x = torch.exp(x)
+    return exp_x / exp_x.sum(dim=-1, keepdim=True)
+
+# Triton kernel - single fused operation
+@triton.jit
+def softmax_kernel(input_ptr, output_ptr, n_cols, BLOCK_SIZE: tl.constexpr):
+    row_idx = tl.program_id(0)
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < n_cols
+
+    row = tl.load(input_ptr + row_idx * n_cols + col_offsets, mask=mask, other=-float('inf'))
+    row_max = tl.max(row, axis=0)
+    numerator = tl.exp(row - row_max)
+    softmax_output = numerator / tl.sum(numerator, axis=0)
+
+    tl.store(output_ptr + row_idx * n_cols + col_offsets, softmax_output, mask=mask)
+```
+
+## Enter kernrl
+
+kernrl is an [OpenEnv](https://github.com/meta-pytorch/OpenEnv) environment that frames GPU kernel optimization as an RL problem:
+
+- **State**: A PyTorch reference implementation + GPU info
+- **Action**: CUDA/Triton kernel code
+- **Reward**: Based on compilation, correctness, and speedup
+
+The environment evaluates submitted kernels on real GPU hardware, providing concrete feedback:
+
+```python
+from kernrl import kernrl_env, KernelAction
+
+env = kernrl_env(base_url="http://localhost:8000")
+obs = env.reset(problem_id="L1_23_Softmax")
+
+# Submit a kernel
+result = env.step(KernelAction(code=triton_kernel_code))
+
+print(f"Compiled: {result.observation.compilation_success}")
+print(f"Correct: {result.observation.correctness_pass}")
+print(f"Speedup: {result.observation.speedup}x")
+```
+
+## 89 Problems Across 10 Difficulty Levels
+
+kernrl includes a diverse problem set spanning from basic operations to cutting-edge architectures:
+
+| Level | Category | Examples |
+|-------|----------|----------|
+| 1 | Simple Ops | matmul, softmax, conv2d, layernorm |
+| 2 | Fused Ops | matmul+GELU+softmax, conv+batchnorm |
+| 3 | Attention | Vision attention, causal attention, transformer blocks |
+| 4 | Novel Layers | DeepSeek MLA, MoE, GQA, FP8 matmul, INT4 GEMM |
+| 5 | Scientific | N-body simulation, stencils, sparse matrix ops |
+| 6 | Graphics | Ray tracing, histogram, bilateral filter |
+| 7 | Signal | FFT, convolution, median filter |
+| 8 | Video | Motion estimation, optical flow, deblocking |
+| 9 | Primitives | Prefix scan, radix sort, stream compaction |
+| 10 | Cryptography | SHA-256, AES, ChaCha20 |
+
+Level 4 is particularly interesting - it includes architectures like DeepSeek's Multi-head Latent Attention and Mixture of Experts that weren't in most training data, testing whether models can truly reason about kernel optimization rather than memorize solutions.
+
+## Training with GRPO
+
+We use TRL's GRPOTrainer with a custom rollout function that interacts with the kernrl environment:
+
+```python
+from trl import GRPOConfig, GRPOTrainer
+from trl.experimental.openenv import generate_rollout_completions
+
+def rollout_func(prompts, trainer):
+    # Generate kernel code
+    outputs = generate_rollout_completions(trainer, prompts)
+
+    # Evaluate in environment
+    env_rewards = []
+    for completion in outputs:
+        code = extract_code(completion)
+        result = env.step(KernelAction(code=code))
+        env_rewards.append(compute_reward(result))
+
+    return {
+        "prompt_ids": [...],
+        "completion_ids": [...],
+        "logprobs": [...],
+        "env_reward": env_rewards,
+    }
+
+trainer = GRPOTrainer(
+    model="Qwen/Qwen2.5-Coder-1.5B-Instruct",
+    reward_funcs=[reward_from_env],
+    rollout_func=rollout_func,
+    args=GRPOConfig(use_vllm=True),
+)
+trainer.train()
+```
+
+The reward structure encourages incremental progress:
+- **+0.1** for successful compilation
+- **+0.3** for correctness (output matches reference within tolerance)
+- **+0.3** for beating the baseline
+- **+0.0 to +0.6** bonus scaled by log2(speedup)
+
+This means a model can learn even before achieving speedups - first learn to write valid code, then correct code, then fast code.
+
+## Why This Matters
+
+### For AI Research
+GPU kernel optimization is a well-defined domain with clear metrics (correctness + speed), making it ideal for studying:
+- How LLMs reason about low-level code optimization
+- Whether RL can teach models skills not present in training data
+- Transfer learning between GPU architectures (H100 vs B200)
+
+### For Practical Applications
+As AI models grow, inference costs dominate. A model that can automatically optimize kernels could:
+- Reduce serving costs by finding faster implementations
+- Adapt to new hardware without manual tuning
+- Enable efficient deployment on edge devices
+
+### For the Community
+kernrl provides:
+- A standardized benchmark for kernel optimization capabilities
+- Real hardware evaluation (not just syntax checking)
+- Integration with the broader OpenEnv ecosystem
+
+## Try It Yourself
+
+**HuggingFace Space**: [huggingface.co/spaces/Infatoshi/kernrl](https://huggingface.co/spaces/Infatoshi/kernrl)
+
+**Training Notebook**: [huggingface.co/Infatoshi/kernrl-training](https://huggingface.co/Infatoshi/kernrl-training)
+
+**OpenEnv PR**: [github.com/meta-pytorch/OpenEnv/pull/308](https://github.com/meta-pytorch/OpenEnv/pull/308)
+
+To run locally with GPU:
+
+```bash
+# Clone OpenEnv
+git clone https://github.com/meta-pytorch/OpenEnv.git
+cd OpenEnv/envs/kernrl
+
+# Install
+pip install -e .
+
+# Start server
+uvicorn kernrl.server.app:app --host 0.0.0.0 --port 8000
+```
+
+Or with Docker:
+
+```bash
+docker build -t kernrl -f server/Dockerfile .
+docker run --gpus all -p 8000:8000 kernrl
+```
+
+## What's Next
+
+We're excited to see what the community builds with kernrl:
+
+- **Curriculum learning**: Start with L1, progressively add harder problems
+- **Multi-turn optimization**: Let models iterate based on profiling feedback
+- **Architecture-specific training**: Specialize models for H100 vs B200
+- **Novel reward shaping**: Incorporate memory bandwidth, occupancy metrics
+
+The code is open source and contributions are welcome. Whether you're interested in RL, GPU programming, or both - we'd love to see what optimizations your models can discover.
+
+---
+
+*kernrl was built for the OpenEnv Challenge. Special thanks to the Meta PyTorch team for the OpenEnv framework and Hugging Face for TRL.*