Views
No views yet
<think></think> block before the answer. With transformers that is just
apply_chat_template(..., enable_thinking=False).1import torch, re
2from transformers import AutoModelForCausalLM, AutoTokenizer
3from peft import PeftModel
4
5BASE = "openbmb/MiniCPM5-1B"
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")
9model = PeftModel.from_pretrained(base, "YMRohit/ouroboros-kernelsmith-minicpm5-1b").eval()
10
11SYSTEM = ("You are an expert GPU kernel engineer. Write a single correct, fast Triton kernel. "
12 "Output ONLY one fenced python code block defining `run(*inputs)` and its @triton.jit "
13 "kernel. Accumulate reductions in float32. No prose.")
14
15# The style guide: a row-wise reduction (rmsnorm). This is what teaches the model the
16# one-program-per-row structure. Use it for anything that reduces over the last dimension.
17STYLE = '''@triton.jit
18def _rmsnorm_kernel(x_ptr, w_ptr, y_ptr, stride, N, eps, BLOCK: tl.constexpr):
19 row = tl.program_id(0)
20 x_ptr += row * stride
21 y_ptr += row * stride
22 acc = tl.zeros([BLOCK], dtype=tl.float32)
23 for off in range(0, N, BLOCK):
24 cols = off + tl.arange(0, BLOCK)
25 x = tl.load(x_ptr + cols, mask=cols < N, other=0.0).to(tl.float32)
26 acc += x * x
27 rms = tl.rsqrt(tl.sum(acc) / N + eps)
28 for off in range(0, N, BLOCK):
29 cols = off + tl.arange(0, BLOCK)
30 mask = cols < N
31 x = tl.load(x_ptr + cols, mask=mask, other=0.0).to(tl.float32)
32 w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
33 tl.store(y_ptr + cols, (x * rms * w), mask=mask)
34
35def run(x, w):
36 M, N = x.shape
37 y = torch.empty_like(x)
38 _rmsnorm_kernel[(M,)](x, w, y, x.stride(0), N, 1e-6, BLOCK=1024)
39 return y'''
40
41# Describe the op you want. Keep it short and concrete. The style guide is a DIFFERENT op on
42# purpose (here we ask for softmax and show rmsnorm).
43USER = (
44 "Op `softmax`: numerically stable softmax over the last dim (subtract the row max).\n"
45 "Signature:\n run(x: Tensor[M, N]) -> Tensor[M, N]\n\n"
46 "Here is a valid Triton kernel for a DIFFERENT op (`rmsnorm`) as a style guide:\n"
47 f"```python\n{STYLE}\n```\n"
48)
49
50messages = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": USER}]
51prompt = tok.apply_chat_template(
52 messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
53inputs = tok(prompt, return_tensors="pt").to(model.device)
54out = model.generate(**inputs, max_new_tokens=768, do_sample=True, temperature=0.7, top_p=0.97)
55text = tok.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
56
57m = re.search(r"```(?:python)?\s*(.*?)```", text, re.S)
58print(m.group(1) if m else text) # a @triton.jit kernel plus a run(...) entry pointtemperature=0.7 and keep the first one that
passes the referee.torch.compile max-autotune), then
reinforcement learning where the only reward is the referee's verdict. The corpus is self-distilled,
written almost entirely by models inside a loop whose only teacher is the verifier. No human-labeled
data anywhere. LoRA rank 64, alpha 128, on the attention and MLP projections.reports/ablations_minicpm_multiseed.md in the corpus repo) this model,
sampled best-of-N against the referee, beat torch.compile max-autotune in all 12 independently
seeded runs. Per-run geomeans were 1.02x to 1.14x over a 6-op suite on a single RTX 4090, and
bigger per-op wins (softmax around 2x) show up and get verified per mint in the Space.