GPU Kernels, Part 2: Teaching a 27B Model to Write Triton Kernels with SFT and a GPU Gym

In Part 1, we explored how GPU kernels work under the hood — why operator fusion and memory tiling matter, and how benchmarks like KernelBench evaluate whether an LLM can write kernels that actually beat PyTorch. This post is the hands-on sequel. We took an open-weight model, Qwen3.8-27B, and trained it to write Triton kernels across two stages: supervised fine-tuning (SFT) on 2,034 high-quality kernels re-verified on an NVIDIA H100, followed by reinforcement learning (RL) inside an automated GPU gym that compiles, validates, and benchmarks every candidate.

Evaluated on the standard 200 KernelBench problems (using a single greedy pass per problem), every headline metric moved in the right direction. Overall correctness improved from 21.5% to 29.5%. Kernels that were both mathematically correct and faster than eager PyTorch grew from 7.0% to 9.5%. On Level 1 (isolated operators), the model now solves 47 problems instead of 31, and beats PyTorch on 15 instead of 8.

At the same time, frontier commercial models still hold a substantial lead. Below, we break down how we built the dataset and verification gym, what each stage of training achieved, the subtle reward-design pitfalls where the gym taught the model the wrong strategy, and what it will take to close the remaining performance gap.

1. The scoreboard

To keep comparisons rigorous and fair, every model in this post was evaluated against the exact same test harness: KernelBench (levels 1 and 2, Triton backend) using byte-identical prompts and the same verification runner on an NVIDIA H100 SXM. Level 1 covers 100 standalone operators (matrix multiplications, convolutions, normalizations, activations, and loss functions). Level 2 tests 100 composite modules — typically a heavy GEMM or convolution followed by an elementwise tail where fusion offers clear headroom. (Part 1 details how KernelBench evaluates submissions and the key caveats when comparing benchmarks).

2. What SFT and RL changed

Out of the box, base Qwen3.8-27B already understands basic Triton syntax — roughly four out of five generated answers compile and execute without syntax errors. What it lacks is numerical precision (getting boundary conditions and indexing math right) and the architectural intuition needed to beat PyTorch. Our two-stage training addressed both:

• Higher compilation reliability: Solutions that compile and execute cleanly rose from 79.5% to 88.5%. On Level 2, the improvement was even sharper (76 → 94), as the model learned to avoid invalid tensor shapes and illegal memory accesses that trigger driver crashes.
• Correctness increased by over a third: Passing all verification checks went from 43 to 59 out of 200 tasks. On Level 1, the jump was especially pronounced (31 → 47), representing a +50% gain in verified standalone operators.
• Speedup followed accuracy: Kernels that were both mathematically correct and faster than eager PyTorch climbed from 14 to 19 tasks. Against torch.compile, fast solutions grew from 25 to 32, while 2×+ speedups rose from 7 to 10. On Level 1, the count of fast kernels nearly doubled (8 → 15).

Supervised fine-tuning provided the foundation for correctness (Level 1: 31 → 43), while RL further refined edge-case handling. Because every checkpoint was evaluated on the exact same problem suite, we can perform a paired bootstrap analysis. Compared to the SFT baseline, the final RL checkpoint successfully resolved 8 previously failing problems while regressing on only 1:

The bootstrap confidence interval confirms that step 100’s correctness gain is statistically meaningful: a +3.5 percentage point lift, with a 95% interval of [+1.0, +6.5]. Speed gains over SFT were more modest, with confidence bounds grazing zero. Beyond step 100, performance plateaued; because we selected step 100 with visibility into later checkpoints, its numbers reflect a best-of-run selection. Level 2 remained the notable outlier: overall correctness held at 12 problems while fast kernels slipped slightly from 6 to 4. Section 6 unpacks the mechanism behind this plateau.

3. Stage 1: an SFT set verified on the target GPU

Public web-scraped kernel datasets look substantial on paper, but they suffer from a major qualification issue: most examples were validated only on tiny toy inputs, run on consumer GPUs with loose tolerances, or never executed at all. For SFT to teach realistic GPU behavior, we built a filtered dataset where every single sample had to pass rigorous execution checks on an NVIDIA H100 at production scale:

• Curated open source: We started with Kimi K2 solutions for KernelBook modules from the siro1 collection. Each problem was refactored into KernelBench’s standardized interface (Model → ModelNew, positional get_init_inputs), capping inclusion at two solutions per module.
• Realistic tensor scales: KernelBook reference shapes often default to toy dimensions like 4 × 4 × 4 × 4, where execution time is dominated by driver overhead rather than memory bandwidth. We scaled tensor dimensions up to ~1,000,000 elements and re-verified all kernels under these production workloads.
• Rigorous decontamination: Any problem with token-set Jaccard similarity exceeding 0.6 against any of the 270 KernelBench problems was purged.
• Pure Triton enforcement: Static AST checks stripped out solutions that cheated by falling back to PyTorch operators inside forward(). While intended as a quality bar, Section 6 explains how this filter later backfired on composite modules.
• Gym replay data: We added 682 verified successful solutions from earlier gym exploration runs (capped at 40 per operator family to prevent simple elementwise tasks from overwhelming the distribution).

The final dataset comprises 2,034 verified pairs — 1,624 of which outperform torch.compile at target scale — all within a 4,096 token context window. We trained a rank-8 LoRA adapter for 2 epochs (3,674 steps, 2e-4 learning rate) on a single H100 in roughly five hours.

4. Stage 2: the kernel gym

While SFT demonstrates what well-structured kernels look like, reinforcement learning gives the model agency to experiment, debug compiler errors, and receive immediate reward signals. Our kernel gym follows an OpenEnv-style architecture: in each episode, a worker serves a reference PyTorch module, the model drafts an implementation, and an isolated GPU subprocess benchmarks the candidate on an H100.

Tasks

The gym generates tasks across fourteen distinct families matching the operational footprint of KernelBench levels 1 and 2. Each task is dynamically instantiated from a procedural seed specifying the mathematical operator, tensor shapes (scaled from 1 to 32 million elements to ensure execution time reflects memory bandwidth rather than kernel launch jitter), and test distributions. For difficult operations where the model initially had near-zero success, we introduced staged curriculum variants: a small shape, a scaffold providing host setup and tile offsets with the inner math left blank, and an epilogue task where the core library GEMM remains intact while the fused activation tail is implemented in Triton.

Two ways to answer

Episodes run under two complementary interaction formats:

• Single-turn mode (25% of episodes): Mirrors KernelBench’s standard prompt format, requiring a complete, one-shot solution to align with benchmark evaluation.
• Multi-turn repair mode (75% of episodes): Allows the model to iteratively call submit_kernel, inspect compiler tracebacks or runtime diffs via inspect_evaluation, and refine its implementation across up to four iterations before calling finalize.

In the final training run, the model solved 63.8% of repair episodes, compared to 50.5% in single-turn mode — highlighting how compiler feedback rescues near-miss implementations.

Five gates

1. Static policy sandbox: Only allowlisted imports are permitted; private PyTorch internals and arbitrary subprocess calls are blocked before execution.
2. Compilation & launch verification: Code must compile cleanly, and execution must demonstrably launch Triton kernels on the GPU verified through CUDA runtime events.
3. Compute isolation: PyTorch native compute is disabled during forward() execution (except for one permitted cuBLAS/cuDNN call in hybrid tasks), ensuring Triton kernels perform the actual arithmetic.
4. Five hidden test cases: Evaluated across nominal scale, 3× scale, 0.01× scale, sign negation, and absolute values at an FP32 tolerance of 1e-4 (scaled by magnitude). This prevents hardcoded lookup hacks and verifies both branches of non-linear activations.
5. Low-variance timing: The candidate kernel, eager PyTorch, and torch.compile(mode="max-autotune") are benchmarked on the same GPU. Timings are accepted only if run-to-run variance remains under 10% (or 30 μs on micro-kernels).

The reward

The reward function balances correctness and performance. Gates 1 through 3 are hard requirements: if code fails to compile, cheats on compute, or violates sandboxing, reward drops to zero immediately.

Here $c$ represents the fraction of hidden test cases passed, $v$ is the measured speedup against the baseline, $\rho$ checks timing stability, and $e$ penalizes unnecessary tool churn. Speed rewards are strictly gated: a kernel receives zero speed credit unless all five hidden test cases pass cleanly with low variance. For complex families where all eight group samples fail, answers within a loose 1e-2 tolerance receive partial progress credit — providing the gradient signal needed to separate near-misses from catastrophic syntax failures.

The RL recipe

The training configuration builds on our multi-reward RL research: CISPO constrains policy weight updates without dropping informative tokens, while REPO-R maintains exploration entropy.

5. Seven gym versions and what broke each one

The verification environment described above is version 9. Getting there required navigating seven distinct failure modes. As we documented in our verifier playbook, RL agents ruthlessly exploit every flaw in an environment. The only reliable way to catch reward leaks was to inspect the generated code directly throughout training:

Two failure modes illustrate how subtle reward bugs can be:

• Metric saturation (v5): Our speed reward originally clipped at 2× speedup. On small ~300k element tensors, almost any compiled Triton kernel easily clocked 6× to 9× faster than torch.compile due to compile overhead. Because every passing kernel scored the maximum 1.0, the gradient offered zero pressure to optimize further. We fixed this by expanding test tensors to 1–32M elements and switching to a logarithmic scale spanning 0.75× to 4×.
• Sandbox bypasses (v8): A code audit revealed that importing an obscure internal PyTorch dispatch function could silently disable our compute guard. While the model had not yet discovered the vulnerability, leaving an exploitable hole in an RL environment is a ticking clock. We locked down the environment with a strict AST import allowlist.

Inside the gym, v9’s staged exercises successfully solved the cold-start problem. Operator families that previously yielded zero valid samples began providing steady gradient feedback:

6. Where the gym and KernelBench disagree

Most of our benchmark gains emerged early: during supervised fine-tuning and the first 100 RL steps. Beyond that point, performance on KernelBench plateaued, even as gym metrics continued to climb. Across four quarters of v9 training, the gym solve rate progressed from 73.4% → 59.0% → 54.9% → 61.7%, yet subsequent KernelBench checkpoints scored 9.5%, 9.5%, 8.5%, 8.0% fast₁.

Why did additional RL training in the gym fail to transfer to the benchmark? A post-mortem revealed three structural mismatches between our training environment and KernelBench:

6.1 A different speed target

The gym measured speedup against torch.compile(mode="max-autotune"), whereas KernelBench measures against standard eager PyTorch. At the gym’s large FP32 tensor sizes, compiled PyTorch was paradoxically often the slower baseline:

While 86.8% of correct training kernels beat torch.compile, only 61.4% beat eager PyTorch. In fact, 26.2% of passing kernels received positive speed reward despite running slower than eager mode. The median compiled runtime was 2.1× slower than eager. In effect, the reward function taught the model to clear a hurdle that was easier than the actual test.

6.2 A different strategy on level 2

KernelBench Level 2 comprises 63 convolution tasks and 36 matrix multiplication tasks, each paired with an elementwise post-processing tail. Outperforming NVIDIA’s hand-tuned cuDNN and cuBLAS on raw convolutions or large GEMMs is exceptionally difficult. The winning architectural move is a hybrid implementation: leave the core GEMM or convolution in cuBLAS/cuDNN, and write a custom Triton kernel only for the fused activation tail.

Early on, base Qwen naturally adopted this strategy, submitting hybrid implementations 55% of the time on Level 2. In fact, an early checkpoint (gym v5) produced 11 fast hybrid solutions — the best Level 2 score of any Qwen model.

Unfortunately, our later data filtering and gym rules penalized this exact behavior. We filtered out hybrid code during SFT curation and strictly blocked PyTorch compute in the gym. Consequently, the SFT model only submitted hybrids 8% of the time, attempting to reimplement full 2D convolutions from scratch in Triton — code that was both prone to numerical bugs and noticeably slower than cuDNN. In contrast, frontier models like Claude Opus and GPT-6 adopted hybrids in 61% to 93% of their Level 2 answers, accounting for virtually all of their high scores.

6.3 Exercises are not the exam

While staged curriculum exercises (scaffolds and epilogues) unlocked training signals for complex operators, the gains concentrated heavily on the assisted variants. When faced with full, unassisted Level 2 modules under benchmark evaluation, the model struggled to synthesize the complete end-to-end kernel. An effective curriculum must ensure its final stage closely mirrors the unassisted target distribution at full scale.

7. Frontier models on the same 200 problems

To put these results in perspective, we evaluated four frontier commercial models on the exact same 200 KernelBench problems using the same H100 execution grader. Frontier models were run through their respective CLI interfaces using standard reasoning modes, meaning this is not a controlled parameter-for-parameter match with greedy, non-reasoning Qwen. Nevertheless, the comparison offers valuable architectural insights:

• Level 1 vs. Level 2 gap: On standalone Level 1 operators, our tuned 27B model is within striking distance of the frontier: top commercial models lead by ~3× on unflagged fast₁ (47 vs 15). On composite Level 2 modules, however, frontier models lead by 20× (81 vs 4) — driven almost exclusively by their ability to generate hybrid cuDNN + Triton solutions.
• Implausible speedup flags: Commercial models produced 6 to 8 solutions with reported speedups exceeding 10×. While some reflect genuine fusion wins, extreme outliers often trace back to timing artifacts or benchmark loopholes; filtering these flags trims fast₁ scores by 3 to 4 percentage points.
• Generation budget disparity: Frontier models consumed substantial reasoning budgets, averaging 1,747 to 2,906 completion tokens per task, compared to roughly 1,109 tokens for our direct greedy outputs.

8. What we would change next

1. Align the gym with the benchmark baseline: Benchmark against eager PyTorch (or a joint eager + compiled objective) so the reward never incentivizes kernels that KernelBench evaluates as slow. The gym now supports this toggle for upcoming training runs.
2. Embrace the hybrid pattern: Allow library cuBLAS/cuDNN calls for heavy GEMMs and convolutions in Level 2-style tasks, rewarding the fused elementwise tail, and retain these hybrid implementations in the SFT corpus rather than filtering them out.
3. Anchor the curriculum to target scale: Retain staged scaffolds for early exploration, but ensure final training stages weight unassisted full tasks at target production shapes.
4. Multi-sample evaluation: Evaluating a single greedy sample per problem introduces noise, making 1–3 problem differences difficult to interpret. Averaging across three seeded passes per task costs more compute but yields clearer separation between checkpoints.

9. Limits

• Single training seed and greedy sampling: Reported Qwen numbers reflect a single training run evaluated under greedy temperature-0 decoding.
• Best-of-run checkpoint selection: Step 100 was selected after observing evaluations across all checkpoints, introducing a mild selection bias.
• Asymmetric frontier conditions: Commercial frontier models utilized vendor CLI wrappers, multi-turn reasoning traces, and stochastic sampling, whereas Qwen ran under single-turn greedy execution without reasoning.
• Precision and architecture scope: All evaluations were conducted in FP32 on KernelBench levels 1 and 2 on NVIDIA H100 SXM; Level 3 architectures and alternative GPU generations were not evaluated.
• Static pattern heuristics: Hybrid classification (“keeps a PyTorch op”) was determined via static AST inspection rather than dynamic runtime profiling.

The telemetry and configuration data for every figure and table in this study are available in the JSON download. The model weights and verified training dataset are open-sourced on Hugging Face: LoRA adapter and verified SFT dataset.