Making Dr GRPO go brrr

02 Jun, 2026 I wrote a fused decode-attention kernel for an RL training loop, got it 2.2× faster than the SDPA path it replaces at the microbenchmark level, dropped it into HuggingFace's generate, and watched the decode step get nearly 3× slower. The kernel was doing exactly what the microbench said it would. The integration broke an auto-compile path that the baseline was quietly benefiting from. This post is how I got there, what the gap actually was, and what closing it would have cost. The wider context: this is the writeup of a project to RL-train a small open source model on GSM8K and write CuteDSL kernels for whichever paths dominate. The concrete setup is Qwen2.5-0.5B-Instruct, Dr. GRPO, a single A10G. The post covers two things: building the training loop from scratch (and squeezing 4.8× out of the rollout phase before any kernel work), and then writing the kernel above for the path that still dominated. Most of what follows is what those two facts look like sitting next to each other. What is RL post-training, and why is it slowIn RL post-training for LLMs, you have a policy (the model), a verifier (something that scores outputs), and a loop that pushes the policy to produce higher-scoring outputs. For a math task like GSM8K, the verifier is just a regex that pulls the final number out of the model's response and compares it to the ground truth. Each training step has two phases. Rollout. Sample a prompt. Generate G completions from the current policy. Score them. Compute advantages. Update. For K inner epochs: forward pass through the policy, compute the GRPO loss against the rewards, backprop, optimizer step. Rollout dominates wall time. The reason is structural. Update is one big batched forward pass over (B*G, P+C) tokens, then a backward and a step. That's three GPU calls. Rollout is model.generate, which is a sequential decode loop that runs one forward pass per generated token, with each pass operating on (B*G, 1, hidden) plus a growing KV cache. Per-token compute is small, but you do it max_new_tokens times in serial.

Even with KV cache and batching, you can't parallelize across the time dimension because each token depends on the last. So most of the time, the GPU is doing many small forwards instead of a few big ones. That's the shape of the problem and that's what kernel work has to address. PPOPPO is a policy gradient method. You collect a rollout from the current policy, then run K epochs of mini-batch updates on that same rollout. Vanilla policy gradient is on-policy: collect a batch, do one update, throw the data away. PPO lets you reuse the same rollout for K epochs, which is the whole reason it exists, by clipping the importance ratio so the policy can't drift too far from the one that generated the data. The ratio is rt(θ)=πθ(at|st)πθold(at|st) If rt=1 nothing changed. If rt>1 the new policy made the action more likely. The clipped objective is LCLIP(θ)=𝔼^t[min(rtA^t, clip(rt,1−ϵ,1+ϵ)A^t)] The min picks the more conservative of the two surrogates, so PPO can improve, but not too much in one step. Classical PPO also has a value network that estimates V(st), with the advantage computed as A^t=Rt−Vϕ(st) (often via GAE). GRPOGRPO drops the value network. Instead of asking "is this output good?" it asks "is this output better than the others I sampled for the same prompt?". The pipeline: Sample G completions for the same prompt Score them with a verifier Compute the advantage as Ai=(Ri−μ)/σ inside the group Apply the same PPO clipped objective No critic at all The whole machinery of estimating V and computing GAE goes away because the group itself acts as the baseline. Dr. GRPOGRPO has two bias problems. Length bias. The original loss averages per-response by 1|oi|. When Ai<0, longer responses get a weaker per-token penalty. The model learns "if I'm going to be wrong, be wrong at length." Output length drifts upward over training even when quality does not improve.

Difficulty bias. Dividing by σ inside a group amplifies gradients on prompts with small std (very easy or very hard ones). Medium-difficulty groups, where the most useful learning signal lives, get under-weighted. Dr. GRPO removes both denominators: Ai=Ri−μ and uses token-sum aggregation instead of per-response mean. The clipped objective stays the same. LDr.GRPO(θ)=1G∑i=1G∑t=1|oi|min(ri,tAi,clip(ri,t,1−ϵ,1+ϵ)Ai)Two deletions, no other changes. In pseudo-code, the whole thing looks like this: for step in range(num_steps): # rollout prompts = sample(dataset, batch_size) completions = policy.generate(prompts, num_samples=G) rewards = verifier(completions) advantages = rewards - group_mean(rewards) # no std division old_logprobs = policy.logprobs(completions).detach() ref_logprobs = ref_policy.logprobs(completions).detach() # update for _ in range(K): logprobs = policy.logprobs(completions) ratio = exp(logprobs - old_logprobs) surrogate = min(ratio * advantages, clip(ratio, 1-eps, 1+eps) * advantages) loss = -token_sum(surrogate * completion_mask).mean() # not token_mean loss += beta * kl(logprobs, ref_logprobs) optimizer.step(loss) Three things to notice in this pseudo-code, because each one is where Dr. GRPO and most working implementations differ from the textbook description: The advantage is rewards - group_mean(rewards). No std normalization. The aggregation is token_sum, not token_mean. Dr. GRPO uses sum (or sum divided by a constant L_max); original GRPO uses sum divided by |o_i|. A KL penalty against a frozen reference policy is added on top. Dr. GRPO drops this for the R1-Zero setup, but for an instruct-tuned starting point you almost always want it. I keep it. Implementation in PyTorchI'm using Qwen2.5-0.5B-Instruct as the policy, GSM8K as the task, and a single A10G as the hardware.

The whole script (grpo.py) is around 300 lines. I built the loop as a skeleton first. No real rewards, no real loss math, just the shape. The point was to get something running end to end with torch.rand rewards and fake advantages, so I could replace each function with the real one once the surrounding scaffolding worked. The skeleton: Load the model and the dataset Sample completions Compute fake rewards Compute fake advantages Compute loss loss.backward() optimizer.step() Once this ran without errors, I replaced each function with the real one. Three of those replacements had non-obvious gotchas worth writing down. The completion maskgenerate pads with EOS after the model finishes. I needed a mask that's 1 for real tokens and 0 after the first EOS, so the loss does not credit the model for padding. From grpo.py: def build_completion_mask(completion_ids, eos_token_id): is_eos = completion_ids == eos_token_id has_eos = is_eos.any(dim=1) first_eos_idx = is_eos.float().argmax(dim=1) seq_len = completion_ids.shape[1] positions = torch.arange(seq_len, device=completion_ids.device).unsqueeze(0) mask = (positions <= first_eos_idx.unsqueeze(1)).long() mask = torch.where(has_eos.unsqueeze(1), mask, torch.ones_like(mask)) return mask The torch.where at the end is the gotcha. argmax returns 0 when there's no True in the row, so without the fallback, completions that never emit EOS would get a mask of [1, 0, 0, ...] and only the first token would count. The where says: if the row has no EOS at all, treat the whole completion as real. compute_logprobsTo compute logπθ for the sampled tokens, I concatenate prompt and completion, run a forward pass, and gather the log-probs at the right positions. The off-by-one trap: logits at position t predict the token at t+1. So the logits that score the completion are at positions [P−1,P,...,P+C−2], not [P,P+1,...].

From grpo.py: def compute_logprobs(model, prompt_ids, prompt_attention_mask, completion_ids, completion_mask): B, P = prompt_ids.shape BG, C = completion_ids.shape G = BG // B device = completion_ids.device prompt_ids_expanded = prompt_ids.repeat_interleave(G, dim=0).to(device) prompt_attn_expanded = prompt_attention_mask.repeat_interleave(G, dim=0).to(device) full_ids = torch.cat([prompt_ids_expanded, completion_ids], dim=1) attention_mask = torch.cat([prompt_attn_expanded, completion_mask], dim=1) logits = model(input_ids=full_ids, attention_mask=attention_mask).logits completion_logits = logits[:, P - 1 : P - 1 + C, :] log_probs = F.log_softmax(completion_logits, dim=-1) selected = log_probs.gather(dim=-1, index=completion_ids.unsqueeze(-1)).squeeze(-1) return selected I verified this by feeding prompt + completion[:5] through the model manually, taking the last logit row, and comparing to selected[0, 5]. They matched to bf16 precision (~4e-3 difference). That sanity check was worth the ten lines it took to write, because every other piece downstream depends on these log-probs being right. The lossOnce compute_logprobs worked, the loss was a direct translation of the equation. From grpo.py: def grpo_loss(current_logprobs, old_logprobs, advantages, completion_mask, eps=0.2): ratio = torch.exp(current_logprobs - old_logprobs) advantages = advantages.unsqueeze(1) unclipped = ratio * advantages clipped = torch.clamp(ratio, 1 - eps, 1 + eps) * advantages per_token_loss = -torch.min(unclipped, clipped) masked_loss = per_token_loss * completion_mask loss_per_response = masked_loss.sum(dim=1) return loss_per_response.mean() Two details that are easy to get wrong: The mask multiplies per_token_loss, not ratio. Masking the ratio destroys exp(0) = 1 at padded positions, which silently changes the surrogate at every step rather than just zeroing out padding contributions. Clip the ratio, then multiply by advantage.

Clipping ratio * advantage is a different operation and does not give you PPO. On inner step 0, current_logprobs equals old_logprobs, so ratio is exactly 1 everywhere and the loss reduces to -(advantages * mask).sum(dim=1).mean(). I print this every run as a sanity check that nothing in the graph is detached or wrong. Putting it togetherWith those three helpers, the full naive GRPO step is short. From grpo.py: for step in range(EPOCHS): questions, gold_answers = sample_batch(dataset, BATCH_SIZE) prompts = [format_prompt(q, tokenizer) for q in questions] tokenized = tokenize(prompts, tokenizer) completion_ids = generate_completions( model, tokenized["input_ids"], tokenized["attention_mask"], tokenizer ) mask = build_completion_mask(completion_ids, tokenizer.eos_token_id) rewards, decoded = compute_rewards(completion_ids, gold_answers, tokenizer, G) B = len(questions) rewards_grouped = rewards.view(B, G) advantages = (rewards_grouped - rewards_grouped.mean(dim=1, keepdim=True)).view(B * G) with torch.no_grad(): old_logprobs = compute_logprobs(model, ..., completion_ids, mask) ref_logprobs = compute_logprobs(ref_model, ..., completion_ids, mask) for inner in range(K): current_logprobs = compute_logprobs(model, ..., completion_ids, mask) kl = masked_kl(current_logprobs, ref_logprobs, mask) loss = grpo_loss(current_logprobs,