qwen3-0.6b-triton-sft-v2
Full fine-tune of Qwen/Qwen3-0.6B to write Triton GPU kernels, from a clean-room
reimplementation of TritonRL (arXiv 2510.17891). Supervised fine-tuning only; no RL yet.
Training data: 20,437 verified (prompt, kernel) pairs over 2,748 KernelBook tasks. Kernels were generated by DeepSeek V4.1 Flash (7 samples per task) and Qwen3.8-27B inside a generate-verify-repair loop, and every sample passed a verifier that checks format, a static anti-cheating linter, a real GPU launch, and 5 correctness trials against the PyTorch reference. 1e-4 learning rate, linear decay, effective batch 16, 8,192 max length; best checkpoint by validation loss (0.288, end of epoch 2 of 3).
Held-out results (152 KernelBook tasks never seen in training, verified the same way):
| Metric | Value |
|---|---|
| Correct, greedy, one sample | 27.0% |
| Correct and faster than PyTorch eager, greedy | 14.5% |
| pass@16 at temperature 0.6 | 55.3% |
KernelBench pass@10: level 1 33%, level 2 16%. The base model scores 0% on the same tasks: it never writes a Triton kernel.
Prompt format. One user turn containing the PyTorch reference module and the
instruction to produce ModelNew with @triton.jit kernels; thinking disabled
(empty <think></think> block as the assistant prefix). The exact template is
USER_TEMPLATE in src/teacher.py of the project repo.
Code, verifier, and the full experiment log: https://github.com/cgpadwick/triton_rl_project
- Downloads last month
- 48