Views
No views yet
rl_adapter_v2, the final 69-op discovery checkpoint. Earlier adapters stay
here only so the training path and ablations remain reproducible.apply_chat_template(..., enable_thinking=False). We want
the kernel, not a chain of thought.bfloat16 it needs roughly 54 GB of VRAM (an 80 GB card, or shard across
GPUs with device_map="auto"), or load it in 4-bit for a single 24 GB card.1import torch, re
2from transformers import AutoModelForCausalLM, AutoTokenizer
3from peft import PeftModel
4
5BASE = "Qwen/Qwen3.6-27B"
6tok = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
7base = AutoModelForCausalLM.from_pretrained(
8 BASE, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto")
9# for a single 24GB GPU, swap the line above for 4-bit:
10# from transformers import BitsAndBytesConfig
11# base = AutoModelForCausalLM.from_pretrained(BASE, trust_remote_code=True, device_map="auto",
12# quantization_config=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16))
13model = PeftModel.from_pretrained(
14 base, "YMRohit/ouroboros-kernelsmith-qwen3.6-27b", subfolder="rl_adapter_v2").eval()
15
16SYSTEM = ("You are an expert GPU kernel engineer. Write a single correct, fast Triton kernel. "
17 "Output ONLY one fenced python code block defining `run(*inputs)` and its @triton.jit "
18 "kernel. Accumulate reductions in float32. No prose.")
19
20# The style guide: a row-wise reduction (rmsnorm). Use it for anything that reduces over the
21# last dimension, which is most of what this model is good at.
22STYLE = '''@triton.jit
23def _rmsnorm_kernel(x_ptr, w_ptr, y_ptr, stride, N, eps, BLOCK: tl.constexpr):
24 row = tl.program_id(0)
25 x_ptr += row * stride
26 y_ptr += row * stride
27 acc = tl.zeros([BLOCK], dtype=tl.float32)
28 for off in range(0, N, BLOCK):
29 cols = off + tl.arange(0, BLOCK)
30 x = tl.load(x_ptr + cols, mask=cols < N, other=0.0).to(tl.float32)
31 acc += x * x
32 rms = tl.rsqrt(tl.sum(acc) / N + eps)
33 for off in range(0, N, BLOCK):
34 cols = off + tl.arange(0, BLOCK)
35 mask = cols < N
36 x = tl.load(x_ptr + cols, mask=mask, other=0.0).to(tl.float32)
37 w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
38 tl.store(y_ptr + cols, (x * rms * w), mask=mask)
39
40def run(x, w):
41 M, N = x.shape
42 y = torch.empty_like(x)
43 _rmsnorm_kernel[(M,)](x, w, y, x.stride(0), N, 1e-6, BLOCK=1024)
44 return y'''
45
46USER = (
47 "Op `softmax`: numerically stable softmax over the last dim (subtract the row max).\n"
48 "Signature:\n run(x: Tensor[M, N]) -> Tensor[M, N]\n\n"
49 "Here is a valid Triton kernel for a DIFFERENT op (`rmsnorm`) as a style guide:\n"
50 f"```python\n{STYLE}\n```\n"
51)
52
53messages = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": USER}]
54prompt = tok.apply_chat_template(
55 messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
56inputs = tok(prompt, return_tensors="pt").to(model.device)
57out = model.generate(**inputs, max_new_tokens=1024, do_sample=True, temperature=0.7, top_p=0.97)
58text = tok.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
59
60m = re.search(r"```(?:python)?\s*(.*?)```", text, re.S)
61print(m.group(1) if m else text)temperature=0.7 and keep the first that passes the referee.rl_adapter_v2/: final Qwen3.6-27B LoRA for the 69 stability-gated kernel set.sft_adapter/, rl_adapter/, rl_adapter_newops/,
rl_adapter_invent/, rl_adapter_f1/, and the ablation-arm adapters. These are kept for
reproduction, not as the default download target.best_kernels/: the 69 curated Triton kernels that match reports/rebench_stability_v2.json.reports/: canonical JSON and markdown reports. reports_local/, raw Modal-volume mirrors,
and scratch planning docs are intentionally not part of this cleaned public layout.evidence/run_logs/: recovered raw run logs, with local machine paths redacted where needed.paper/: paper draft and generated figures.torch.compile max-autotune. The model learns from its own verified wins, no human labels.
Trained on Modal H200s; the RL run peaks around 110 GB of VRAM. LoRA rank 128. The adapters here
cover the original suite, the new-operator discovery runs, an invention run on never-trained
problems, and a transfer run that fixed the worst loss cases.torch.compile max-autotune on an H200, 69 of them
held up across 5 fresh re-benchmark runs, and they keep a 1.49x geomean over a 376-cell shape and
dtype grid. They also beat hand-written expert kernels (Liger, Unsloth, the Triton tutorial) on
several ops.