# OptiML: An End-to-End Framework for Program Synthesis and CUDA Kernel Optimization

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 35 references that are in the map.
- First posted on arXiv 12 Feb 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [DeepSeek-Coder: When the Large Language Model Meets Programming -- The Rise of Code Intelligence](https://mlsanity.com/paper/2401.14196/) (2024)
- [A Survey on Large Language Models for Code Generation](https://mlsanity.com/paper/2406.00515/) (2024)
- [KernelBench: Can LLMs Write Efficient GPU Kernels?](https://mlsanity.com/paper/2502.10517/) (2025)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [Monte Carlo Tree Search for Comprehensive Exploration in LLM-Based Automatic Heuristic Design](https://mlsanity.com/paper/2501.08603/) (2025)
- [Kevin: Multi-Turn RL for Generating CUDA Kernels](https://mlsanity.com/paper/2507.11948/) (2025)
- [TritonBench: Benchmarking Large Language Model Capabilities for Generating Triton Operators](https://mlsanity.com/paper/2502.14752/) (2025)
- [AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs](https://mlsanity.com/paper/2507.05687/) (2025)
- [CUDA-LLM: LLMs Can Write Efficient CUDA Kernels](https://mlsanity.com/paper/2506.09092/) (2025)
- [MCTS-Judge: Test-Time Scaling in LLM-as-a-Judge for Code Correctness Evaluation](https://mlsanity.com/paper/2502.12468/) (2025)
- [Geak: Introducing Triton Kernel AI Agent & Evaluation Benchmarks](https://mlsanity.com/paper/2507.23194/) (2025)
- [CodeRosetta: Pushing the Boundaries of Unsupervised Code Translation for Parallel Programming](https://mlsanity.com/paper/2410.20527/) (2024)
- [CUDA-L2: Surpassing cuBLAS Performance for Matrix Multiplication through Reinforcement Learning](https://mlsanity.com/paper/2512.02551/) (2025)
- [Mixture of Thoughts: Learning to Aggregate What Experts Think, Not Just What They Say](https://mlsanity.com/paper/2509.21164/) (2025)
- [HPC-Coder-V2: Studying Code LLMs Across Low-Resource Parallel Languages](https://mlsanity.com/paper/2412.15178/) (2024)
- [Prompt Alchemy: Automatic Prompt Refinement for Enhancing Code Generation](https://mlsanity.com/paper/2503.11085/) (2025)
- [LS-CAT: A Large-Scale CUDA AutoTuning Dataset](https://mlsanity.com/paper/2103.14409/) (2021)
- [Searching CUDA code autotuning spaces with hardware performance counters: data from benchmarks running on various GPU architectures](https://mlsanity.com/paper/2102.05299/) (2021)

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

- [Optimizing CUDA like a Human: Micro-Profiling Tools as Expert Surrogates for LLM-Based GPU Kernel Optimization](https://mlsanity.com/paper/2606.26453/) (2026)
- [Multi-turn RL with Structural and Performance Aware Rewards for CUDA Kernel Generation](https://mlsanity.com/paper/2607.20908/) (2026)
- [KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization](https://mlsanity.com/paper/2609.30059/) (2026)

## Abstract (arXiv)

Generating high-performance CUDA kernels remains challenging due to the need to navigate a combinatorial space of low-level transformations under noisy and expensive hardware feedback. Although large language models can synthesize functionally correct CUDA code, achieving competitive performance requires systematic exploration and verification of optimization choices. We present OptiML, an end-to-end framework that maps either natural-language intent or input CUDA code to performance-optimized CUDA kernels by formulating kernel optimization as search under verification. OptiML consists of two decoupled stages. When the input is natural language, a Mixture-of-Thoughts generator (OptiML-G) acts as a proposal policy over kernel implementation strategies, producing an initial executable program. A search-based optimizer (OptiML-X) then refines either synthesized or user-provided kernels using Monte Carlo Tree Search over LLM-driven edits, guided by a hardware-aware reward derived from profiler feedback. Each candidate transformation is compiled, verified, and profiled with Nsight Compute, and evaluated by a composite objective that combines runtime with hardware bottleneck proxies and guardrails against regressions. We evaluate OptiML in both synthesis-and-optimize and optimization-only settings on a diverse suite of CUDA kernels. Results show that OptiML consistently discovers verified performance improvements over strong LLM baselines and produces interpretable optimization trajectories grounded in profiler evidence.