3.4 Reward Design for Kernel RL

Reward Design for Kernel RL #

The reward is the only channel through which “what we actually want” enters RL training, a scalar \( r(x, y) \) that should be high when a kernel compiles, is correct, and runs fast. Section 3.2 introduced the gated compile-correct-fast reward in its simplest form; this subsection is about designing it well. How the reward is composed decides what the policy learns, and its weaknesses become the policy’s exploits, so the two questions here are what signals the reward should carry and how to fold them into a single reward design. The failure mode that a poorly designed reward invites, reward hacking, and the verification pipeline that defends against it, is the subject of Section 2.3; the reward built here is exactly the checker that pipeline hardens.

Signals for good kernels #

End-to-end speedup over a baseline is the headline signal, but a reward built on speedup alone is both sparse (a kernel earns nothing until it is correct and faster) and easy to game. Richer, more diagnostic signals make the reward denser and harder to exploit. The commonly used signals fall into a few groups:

  • Correctness: outputs match the reference within tolerance, ideally across many shapes and dtypes, not one. Example: a fused RMSNorm kernel matches the PyTorch reference to within atol=1e-2 on inputs of shape (4096, 1024) and (8192, 4096).
  • Performance: end-to-end speedup, but also its breakdown: kernel-vs-total runtime, per-region timing to locate the bottleneck, and the compute / memory-access / synchronization / launch split. Example: the kernel runs 1.8x faster than the reference, and profiling attributes 90% of that time to global-memory reads rather than compute.
  • Memory and hardware efficiency: memory-bandwidth utilization, peak memory footprint, cache-hit rate, FLOP efficiency, and (on accelerators such as Trainium) engine-specific occupancy of the tensor/vector/scalar units and DMA activity. Example: the kernel sustains 85% of the hardware’s peak memory bandwidth, confirming it is close to the roofline for a memory-bound operator.
  • Robustness: variance of correctness and speed across shapes, and run-to-run timing consistency (e.g., median-of-buckets timing to resist noise). Example: the kernel stays correct and above 1.5x speedup on both power-of-two and non-power-of-two N, rather than only on the shapes it was tuned for.
  • Quality: a rule-based or LLM-judge assessment of whether the kernel is a genuine, idiomatic implementation rather than a degenerate one. Example: a judge flags a “matmul” kernel that secretly calls torch.matmul as not a genuine implementation, even though it passes the correctness check.

Beyond making the reward denser, several of these signals double as anti-cheating signals: a kernel that claims a large speedup while its custom code covers only a trivial fraction of the runtime, or while the tensor engine sits idle, is almost certainly gaming the reward rather than optimizing it. Turning these signals into a robust verifier that a policy cannot cheat is exactly the pipeline of Section 2.3.

Reward designs #

Given a set of signals \( s_k(x, y) \) , the design question is how to fold them into a single scalar. The gated compile-correct-fast reward of Section 3.2 is the simplest such design; the shapes below generalize it, arranged roughly from most to least hackable:

  • Linear weighted sum. \( r(x, y) = \sum_k w_k\, s_k(x, y) \) : add up the component signals, e.g. correctness plus speedup (Kevin [3]) or correctness alone (AutoTriton [4]). Simple and dense, but every term is an independent surface to game, and a large weight on speed can dominate correctness.
  • Discrete milestone / gating. Bucket the outcome into ordered tiers so partial progress is rewarded without unbounded incentives: \[ r(x, y) = \begin{cases} 0, & \text{fails to compile}, \\ 1, & \text{correct}, \\ 2, & \text{correct and faster than baseline}. \end{cases} \] This is robust to outliers and low-variance, and correctness strictly gates any performance reward (CUDA Agent [2]).
  • Conditional gating. A stricter version of the above: performance is scored only once correctness passes, and efficiency signals only once a speed threshold passes, so the policy cannot trade correctness for speed. Writing the checks as multiplicative factors makes any failed gate zero out the reward; TritonRL [1] instantiates this as \[ R_{\text{correct}}(x, y) = \mathrm{valid}(y)\cdot\mathrm{correct}(x, y), \\[4pt] R_{\text{speedup}}(x, y) = \mathrm{valid}(y)\cdot\mathrm{clip}\big(\mathrm{speedup}(x, y),\, 2\big), \] where \( \mathrm{valid}(y) = \mathrm{syntax}(y)\cdot\mathrm{func}(y) \) is itself a product of a syntax and a functionality check, so a kernel earns speedup credit only once it is valid and correct, and the speedup factor is capped at \( 2\times \) .
  • Hierarchical decomposition. Split the trajectory into parts and reward them differently, for instance a high-level optimization plan vs. the low-level code, attaching speed-oriented credit to the plan and correctness credit to the implementation (TritonRL [1]). Concretely, the plan tokens and code tokens of a single response receive different rewards, \[ r^{\text{plan}}(x, y) = R_{\text{speedup}}(x, y), \qquad r^{\text{code}}(x, y) = R_{\text{correct}}(x, y), \] and the policy optimizes their weighted combination \[ \mathcal{J}(\theta) = \mathbb{E}\big[\, \alpha\, \mathcal{J}^{\text{plan}}(\theta) + \mathcal{J}^{\text{code}}(\theta) \,\big], \] where each \( \mathcal{J}^{c} \) is a GRPO objective computed over only the tokens of class \( c \in \{\text{plan}, \text{code}\} \) , each with its own group-relative advantage. Rewarding code for correctness rather than speedup is deliberate: because the code is conditioned on the plan, charging it for a slow plan would penalize a faithful implementation of an inherently suboptimal strategy. The weight \( \alpha \approx 0.1 \) lets the plan policy evolve slowly enough for the code policy to keep pace. This is the same plan/code split developed as a within-response credit-assignment tool in Section 3.3, and it somewhat ameliorates the sparse credit assignment problem.
  • Length shaping. Orthogonal to the hackability axis, this targets a different failure mode (verbosity, not cheating): a concave length penalty (short kernels for easy tasks, longer allowed for hard ones) discourages padding and rambling without punishing genuinely complex kernels (Composer 2 [5]).

However carefully these shapes are chosen, each one is still a proxy the policy will probe for gaps. Making the reward robust to that probing, catching the kernel that satisfies the letter of the check while cheating underneath, is a distinct problem from designing the reward, and is the subject of Section 2.3.

References #

LLM kernel generation and reward design.

  1. J. Woo, S. Zhu, A. Nie, Z. Jia, Y. Wang, and Y. Park. TritonRL: Training LLMs to Think and Code Triton Without Cheating. arXiv:2510.17891, 2025. arxiv.org/abs/2510.17891
  2. W. Dai, H. Wu, Q. Yu, et al. CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation. arXiv:2602.24286, 2026. arxiv.org/abs/2602.24286
  3. C. Baronio, P. Marsella, B. Pan, S. Guo, and S. Alberti. Kevin: Multi-Turn RL for Generating CUDA Kernels. arXiv:2507.11948, 2025. arxiv.org/abs/2507.11948
  4. S. Li, Z. Wang, Y. He, Y. Li, Q. Shi, J. Li, Y. Hu, W. Che, X. Han, Z. Liu, and M. Sun. AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs. arXiv:2507.05687, 2025. arxiv.org/abs/2507.05687
  5. Cursor Research. Composer 2 Technical Report. arXiv:2603.24477, 2026. arxiv.org/abs/2603.24477
Previous 3.3 Single-Turn and Multi-Turn RL for Kernel Optimization