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

arXiv: https://arxiv.org/abs/2601.05475 · HTML: https://mlsanity.com/paper/2601.05475/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Process reward models for reasoning, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 15 of its 19 references that are in the map.
- First posted on arXiv 9 Jan 2026.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Process reward models for reasoning](https://mlsanity.com/alignment/process-reward-models-for-reasoning/)

## Builds on (most cited first, 15 of 15)

- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [KernelBench: Can LLMs Write Efficient GPU Kernels?](https://mlsanity.com/paper/2502.10517/) (2025)
- [Generative Reward Models](https://mlsanity.com/paper/2410.12832/) (2024)
- [Learning Performance-Improving Code Edits](https://mlsanity.com/paper/2302.07867/) (2023)
- [Kevin: Multi-Turn RL for Generating CUDA Kernels](https://mlsanity.com/paper/2507.11948/) (2025)
- [CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement Learning](https://mlsanity.com/paper/2507.14111/) (2025)
- [Teaching Language Models to Critique via Reinforcement Learning](https://mlsanity.com/paper/2502.03492/) (2025)
- [CUDA-LLM: LLMs Can Write Efficient CUDA Kernels](https://mlsanity.com/paper/2506.09092/) (2025)
- [EffiLearner: Enhancing Efficiency of Generated Code via Self-Optimization](https://mlsanity.com/paper/2405.15189/) (2024)
- [FlashMask: Efficient and Rich Mask Extension of FlashAttention](https://mlsanity.com/paper/2410.01359/) (2024)
- [Afterburner: Reinforcement Learning Facilitates Self-Improving Code Efficiency Optimization](https://mlsanity.com/paper/2505.23387/) (2025)
- [PerfRL: A Small Language Model Framework for Efficient Code Optimization](https://mlsanity.com/paper/2312.05657/) (2023)
- [To the Max: Reinventing Reward in Reinforcement Learning](https://mlsanity.com/paper/2402.01361/) (2024)
- [LLM Program Optimization via Retrieval Augmented Search](https://mlsanity.com/paper/2501.18916/) (2025)

## Built on by (most cited first, 4 of 4)

- [Towards Automated Kernel Generation in the Era of LLMs](https://mlsanity.com/paper/2601.15727/) (2026)
- [Improving HPC Code Generation Capability of LLMs via Online Reinforcement Learning with Real-Machine Benchmark Rewards](https://mlsanity.com/paper/2602.12049/) (2026)
- [Reinforcement Learning for Code Optimization](https://mlsanity.com/paper/2607.25970/) (2026)
- [RLPF: Reinforcement Learning from Performance Feedback for Code Generation](https://mlsanity.com/paper/2607.27271/) (2026)

## Abstract (arXiv)

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.