+
+
+
+
🎯
+
Multi-Family Support
+
Optimizes three kernel families: Softmax, LayerNorm, and Grouped GEMM for NVIDIA GPUs
+
+
+
📊
+
Surrogate Modeling
+
Uses Gaussian Process surrogates to predict kernel performance and guide exploration
+
+
+
🤖
+
Self-Improving AI
+
Employs student-teacher LLM architecture for autonomous kernel optimization
+
+
+
⚡
+
Triton Backend
+
Generates optimized Triton kernels with configurable block sizes, warps, and stages
+
+
+
+
Supported Kernel Families
+
+ softmax
+ layernorm
+ grouped_gemm
+
+
+
+
How It Works
+
+
+
1.
+
Reset - Initialize environment with a kernel task (e.g., softmax_m4096_n256)
+
+
+
2.
+
Configure - Select Triton config (block_size, num_warps, num_stages)
+
+
+
3.
+
Step - Benchmark configuration and receive reward based on latency
+
+
+
4.
+
Optimize - Use LLM to generate better kernel variants automatically
+
+
+
5.
+
Learn - Surrogate model improves with more observations
+
+
+
6.
+
Compare - Student vs Teacher LLM competition for best optimization
+
+
+
+
+
+
+POST /reset
+{
+ "task": "softmax_m4096_n256",
+ "seed": 0
+}
+
+POST /step
+{
+ "config_id": 0,
+ "source": "def benchmark_generated_kernel..."
+}
+
+{
+ "reward": 0.093,
+ "observation": {
+ "best_so_far_ms": 0.41,
+ "oracle_best_ms": 0.38
+ }
+}
+
+
+
Try It Now
+
Use the interactive Gradio interface above to benchmark kernel configurations and explore the RL environment.
+
Launch Interface →
+
+
+