Qwen3-1.7B Triton kernel writer, GRPO run 9 (step 1698, branch main)

Writes Triton GPU kernels for PyTorch modules. Part of an independent reimplementation of TritonRL (Woo et al., arXiv 2510.17891): supervised fine-tuning on verified teacher kernels, then GRPO with a reward that compiles and runs each candidate kernel against the PyTorch reference.

  • Start: cgpadwick2020/qwen3-1.7b-triton-sft-v2 (Qwen3-1.7B, SFT on about 20K verified DeepSeek kernels over KernelBook tasks).
  • GRPO run 9: 2 epochs over 6,799 tasks, 1,698 steps: the KernelBook training tasks plus size-augmented variants of them (the same modules at KernelBench-like input sizes, up to 2^25 elements and rows of 65,537). A completion on an augmented task is rewarded only if it is correct at the large size and at the original. 8 tasks x 16 samples per step, lr 3e-6, 4,096-token completions, full fine-tune; checkpoint saved in fp32, uploaded in bf16.
  • Branches: main = final (step 1698); step-1600 = best held-out validation; step-1500.

Held-out validation during training (152 tasks x 8 samples, temperature 0.6, pass@1 / pass@8): step 1500 52.5 / 69.7, step 1600 52.5 / 70.4, step 1698 51.7 / 67.8 (SFT start 29.0 / 55.3).

Final checkpoint: test split (152 tasks x 16 samples) pass@1 51.3, pass@16 70.4; KernelBench at the problem sizes of early 2025 (the version the paper used) with KernelBench's own evaluator, pass@10: Level 1 80%, Level 2 41%.

Prompt: the reference module in a Python code block after "Here is the reference PyTorch module. get_inputs() and get_init_inputs() show how it will be constructed and called.", ending "Write ModelNew as specified."; chat template with an empty think block. Sample at temperature 0.6, top-k 20, top-p 0.95.

Downloads last month
38
Safetensors
Model size
2B 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-1.7b-triton-grpo-run9

Finetuned
Qwen/Qwen3-1.7B
Finetuned
(1)
this model
Finetunes
1 model

Paper for cgpadwick2020/qwen3-1.7b-triton-grpo-run9