Results
Evaluated on 120 episodes per arm (same tasks, same seeds, 4 samples per task), H100, sampled with the same settings as training.
Table with columns: Split, Base, SFT, SFT + RL (this model)| Split | Base | SFT | SFT + RL (this model) |
|---|
| All 120 episodes, correct | 19/120 | 18/120 | 24/120 |
| Trained tasks (16 ep) | 81.2% | 68.8% | 93.8% |
| Untrained tasks, seen split (32 ep) | 15.6% | 21.9% | 21.9% |
| Held-out level 1 | 0% | 0% | 0% |
| Held-out level 2 (fusion) | 2.8% | 0% | 5.6% |
Paired base → RL deltas with task-level bootstrap 95% CIs (percentage points):
Table with columns: Split, Δ, 95% CI| Split | Δ | 95% CI |
|---|
| Trained | +12.5 | [0, 37.5] |
| Untrained | +6.2 | [−12.5, 25.0] |
| Held-out L1 | 0 | — |
| Held-out L2 | +2.8 | [0, 8.3] |
Read this honestly: RL helped most on the tasks it trained on, the gain on held-out tasks is small, and every
CI includes zero. This is a small-sample research run, not a strong kernel-writing model. Speedups over PyTorch
were rare; most correct kernels run at roughly reference speed.
Training details
- Base:
Qwen/Qwen2.5-Coder-7B-Instruct @ c03e6d358207e414f1eca0bb1891e29f1db0e242
- SFT: LoRA, 12 steps, 28 level-1 teacher kernels
- RL: GRPO (TRL 1.16 + vLLM), LoRA on top of the SFT-merged model, 64 optimizer steps, 8 generations per group,
18 level-2 fusion tasks, seed 42; resumed once from checkpoint-20
- Reward: staged (format → compile → runtime → correctness → timing); correct-but-not-faster = 0.25;
timing is 7 paired rounds with an IQR band; kernels slower than a 20 s cap get correct-not-faster
- Contract: level 1 must be pure Triton; level 2 may call vendor GEMM/conv in the reference dtype and fuse the rest
in Triton (fp16/bf16 library calls rejected). Each prompt carries one fixed worked example from a disjoint task.
Evaluation details
- 120 episodes per arm: 16 trained, 32 untrained-seen, plus held-out level 1 and level 2 tasks
- Correctness checked against the KernelBench reference (pinned commit
423217d9) on multiple random inputs;
output and state tampering detected with keyed GPU fingerprints
- One task (Frobenius norm) excluded from all arms: its reference is numerically unstable under the checker
Use
from transformers import AutoModelForCausalLM, AutoTokenizer repo = "Shreyansh327/Qwen2.5-Coder-7B-KernelEnv-Fusion-RL"tok = AutoTokenizer.from_pretrained(repo)model = AutoModelForCausalLM.from_pretrained(repo, torch_dtype="bfloat16", device_map="auto")
Prompt it with a KernelBench reference module and ask for a ModelNew class that uses Triton kernels.
Rebuild from adapters
The RL adapter was trained on the SFT-merged model, so merge in order:
import torchfrom transformers import AutoModelForCausalLMfrom peft import PeftModel repo = "Shreyansh327/Qwen2.5-Coder-7B-KernelEnv-Fusion-RL"m = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-Coder-7B-Instruct", revision="c03e6d358207e414f1eca0bb1891e29f1db0e242", torch_dtype=torch.bfloat16)m = PeftModel.from_pretrained(m, repo, subfolder="adapters/sft").merge_and_unload()m = PeftModel.from_pretrained(m, repo, subfolder="adapters/rl").merge_and_unload()
Checksums (sha256)
Table with columns: File, sha256| File | sha256 |
|---|
adapters/sft/adapter_model.safetensors | 478d5412a776c35674c150b7525801d25dd947b8a76bb4bf7162a5c66967c2b4 |
adapters/rl/adapter_model.safetensors | 86df23e5a587f2d8ab25bdea2387bdfd20d3ffd70b244589051c7a8934575b38 |
Full per-arm metrics: eval/comparison.json.
Limitations
- Small evaluation (120 episodes/arm); improvements are not statistically established
- Trained and evaluated on H100 only; kernels may not be correct or fast on other GPUs
- Generated kernels can be wrong; always verify against a reference before use
- Not a general chat model improvement; tuned narrowly for KernelBench-style prompts