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
Safetensors
Model size
0.6B params
Tensor type
BF16
ยท
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support

Model tree for cgpadwick2020/qwen3-0.6b-triton-sft-v2

Finetuned
Qwen/Qwen3-0.6B
Finetuned
(1,456)
this model