Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MaxCode: A Max-Reward Reinforcement Learning Framework for Automated Code Optimization

Jiefu Ou, Sapana Chaudhary, Kaj Bostrom, Nathaniel Weir, Shuai Zhang, Huzefa Rangwala, George Karypis

TL;DR

MaxCode reframes code optimization as a max-reward reinforcement learning problem, unifying inference-time search methods under a modular framework that augments observations with a natural-language critique and a best-discounted reward. It introduces a max-reward inference operator and two reformulated search strategies (Effi-Learner and CUDA-LLM) plus a Generative Value Function approach to guide search with a learned reward model. The framework leverages an extended MDP that includes an auxiliary reward-history variable and uses a critique model to diagnose bottlenecks and suggest refinements, improving exploration efficiency under budgeted inference. Empirical results on KernelBench (CUDA) and PIE (C++) show consistent gains over baselines, validating the practical impact of max-reward guidance for code performance optimization.

Abstract

Large Language Models (LLMs) demonstrate strong capabilities in general coding tasks but encounter two key challenges when optimizing code: (i) the complexity of writing optimized code (such as performant CUDA kernels and competition-level CPU code) requires expertise in systems, algorithms and specific languages and (ii) requires interpretation of performance metrics like timing and device utilization beyond binary correctness. In this work, we explore inference-time search algorithms that guide the LLM to discover better solutions through iterative refinement based on execution feedback. Our approach, called MaxCode unifies existing search methods under a max-reward reinforcement learning framework, making the observation and action-value functions modular for modification. To enhance the observation space, we integrate a natural language critique model that converts raw execution feedback into diagnostic insights about errors and performance bottlenecks, and the best-discounted reward seen so far. Together, these provide richer input to the code proposal function. To improve exploration during search, we train a generative reward-to-go model using action values from rollouts to rerank potential solutions. Testing on the KernelBench (CUDA) and PIE (C++) optimization benchmarks shows that MaxCode improves optimized code performance compared to baselines, achieving 20.3% and 10.1% relative improvements in absolute speedup value and relative speedup ranking, respectively.

MaxCode: A Max-Reward Reinforcement Learning Framework for Automated Code Optimization

TL;DR

MaxCode reframes code optimization as a max-reward reinforcement learning problem, unifying inference-time search methods under a modular framework that augments observations with a natural-language critique and a best-discounted reward. It introduces a max-reward inference operator and two reformulated search strategies (Effi-Learner and CUDA-LLM) plus a Generative Value Function approach to guide search with a learned reward model. The framework leverages an extended MDP that includes an auxiliary reward-history variable and uses a critique model to diagnose bottlenecks and suggest refinements, improving exploration efficiency under budgeted inference. Empirical results on KernelBench (CUDA) and PIE (C++) show consistent gains over baselines, validating the practical impact of max-reward guidance for code performance optimization.

Abstract

Large Language Models (LLMs) demonstrate strong capabilities in general coding tasks but encounter two key challenges when optimizing code: (i) the complexity of writing optimized code (such as performant CUDA kernels and competition-level CPU code) requires expertise in systems, algorithms and specific languages and (ii) requires interpretation of performance metrics like timing and device utilization beyond binary correctness. In this work, we explore inference-time search algorithms that guide the LLM to discover better solutions through iterative refinement based on execution feedback. Our approach, called MaxCode unifies existing search methods under a max-reward reinforcement learning framework, making the observation and action-value functions modular for modification. To enhance the observation space, we integrate a natural language critique model that converts raw execution feedback into diagnostic insights about errors and performance bottlenecks, and the best-discounted reward seen so far. Together, these provide richer input to the code proposal function. To improve exploration during search, we train a generative reward-to-go model using action values from rollouts to rerank potential solutions. Testing on the KernelBench (CUDA) and PIE (C++) optimization benchmarks shows that MaxCode improves optimized code performance compared to baselines, achieving 20.3% and 10.1% relative improvements in absolute speedup value and relative speedup ranking, respectively.
Paper Structure (27 sections, 6 equations, 4 figures, 2 tables)

This paper contains 27 sections, 6 equations, 4 figures, 2 tables.

Table of Contents

  1. Introduction
  2. MaxCode: Max-Reward RL Framework for Code Optimization
  3. The Markov Decision Process
  4. Remark
  5. Execution and Critique
  6. Max-Reward Inference Operator
  7. Max-Reward Search
  8. Effi-Learner
  9. CUDA-LLM
  10. Generative Value Function Guided Search
  11. Categorical Reward Formulation
  12. Generative Value Function Guided Search
  13. Environment Stochasticity
  14. Experiments and Results
  15. Experiments
  16. ...and 12 more sections

Figures (4)

  • Figure 1: Example optimization code generated by DeepSeek-R1 on a KernelBench problem
  • Figure 2: Illustration of the max-reward reinforcement learing formulation of MaxCode
  • Figure 3: Ablated evaluation results of correct, fast1, and max speed-up of different components of MaxCode on KernelBench and PIE
  • Figure 4: Inference time scaling of max speed-up on KernelBench and PIE.