🤖 AI

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

OptiML is an end-to-end framework that generates and optimizes high-performance CUDA kernels by combining a Mixture-of-Thoughts generator for initial synthesis with a search-based optimizer that uses Monte Carlo Tree Search and hardware profiler feedback to systematically refine code for maximum performance.

Original authors: Arijit Bhattacharjee, Heng Ping, Son Vu Le, Paul Bogdan, Nesreen K. Ahmed, Ali Jannesari

Published 2026-02-16

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Arijit Bhattacharjee, Heng Ping, Son Vu Le, Paul Bogdan, Nesreen K. Ahmed, Ali Jannesari

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to build the fastest possible race car. You have a brilliant mechanic (a Large Language Model, or LLM) who is great at drawing up blueprints and assembling the engine. However, this mechanic has a flaw: while they can build a car that runs, they often miss the tiny tweaks that make it win the race. They might put the wrong tires on, forget to tune the fuel injection, or arrange the gears inefficiently.

In the world of computer chips (specifically NVIDIA GPUs), writing the code to run these "race cars" (called CUDA kernels) is incredibly hard. It requires deep knowledge of how the hardware works. If you get it wrong, the car might not start at all, or it might crawl when it should be flying.

Enter OptiML. Think of OptiML not just as a mechanic, but as a two-person pit crew working together to turn a good car into a championship-winning one.

The Two-Part Pit Crew

OptiML splits the job into two distinct roles, working in a seamless loop:

1. The Architect (OptiML-G): "The Smart Blueprint Generator"

The Problem: Sometimes, you just have a vague idea: "I need a car that can drift around corners." You don't have a blueprint yet. If you ask a standard mechanic to draw one, they might give you a blueprint that looks cool but has the engine in the trunk.
The Solution: OptiML-G is like a Master Architect with a team of specialists. Instead of relying on just one brain, it uses a "Mixture of Thoughts." Imagine a room full of expert architects (different AI models). One is great at aerodynamics, another at engine placement, and another at safety.
How it works: When you give it a task, it doesn't just pick one expert. It has them all "think" about the problem simultaneously, sharing their best ideas in a hidden mental space. They collaborate to produce a strong, solid starting blueprint that is guaranteed to be structurally sound and ready for the next step.

2. The Tuner (OptiML-X): "The Data-Driven Race Engineer"

The Problem: Even with a good blueprint, the car might not be fast enough. You need to tweak the suspension, change the gear ratios, and optimize the aerodynamics. But how do you know what to change? If you just guess, you might make it slower.
The Solution: OptiML-X is a hyper-observant Race Engineer who uses a "Monte Carlo Tree Search."
- The Search: Imagine the engineer standing at a fork in the road. "If I tighten this bolt, what happens? If I swap this tire, what happens?" They don't just try one thing; they simulate thousands of "what-if" scenarios in their head (and on the computer).
- The Judge: Crucially, this engineer has a super-visor (an LLM acting as a Judge). After every tiny change, the car is put on a test track (the hardware profiler). The engineer looks at the data: "Did we use less fuel? Did the engine run hotter? Did we go faster?"
- The Feedback Loop: If a change makes the car faster, they keep it. If it makes the car slower or breaks the engine, they throw that idea away immediately. They use this data to guide their next guess, slowly climbing the mountain of performance until they reach the peak.

Why This is a Game-Changer

In the past, people tried to do these two jobs separately:

The "Generate" approach: Ask an AI to write the code. (Result: It works, but it's slow).
The "Optimize" approach: Take that slow code and try to fix it. (Result: If the original code was terrible, the optimizer can't fix it; it's like trying to tune a car with a broken engine).

OptiML combines them.

If you give it a natural language idea (e.g., "Make a function that sorts numbers"), the Architect builds a strong foundation, and the Tuner polishes it to perfection.
If you give it existing code that is slow, the Tuner skips the building phase and goes straight to the tuning, using real hardware data to fix the bottlenecks.

The "Secret Sauce": Listening to the Hardware

The most creative part of OptiML is how it decides what to change. It doesn't just look at "how fast is it?" (which can be noisy and confusing). It looks at why it's slow.

Think of it like a doctor diagnosing a patient:

Old way: "The patient is tired." (So, give them coffee? Maybe, maybe not.)
OptiML way: "The patient is tired because their heart rate is too high and they aren't breathing efficiently." (So, we fix the breathing and heart rate specifically.)

OptiML looks at the "vital signs" of the computer chip (memory usage, how busy the processors are, how much data is moving). If the chip is "thirsty" for data (memory-bound), OptiML focuses on reducing data movement. If the chip is "overworked" (compute-bound), it focuses on making the instructions more efficient.

The Result

In the paper's tests, OptiML consistently beat other top AI models.

Speed: It made code run 1.5x to 1.7x faster than the best standalone AI attempts.
Reliability: It fixed code that other AIs couldn't even compile (code that wouldn't start).
Transparency: It doesn't just give you a black box; it tells you why it made changes (e.g., "I reduced memory traffic by 20%").

In a Nutshell

OptiML is like having a dream team for computer programming. You have a Master Architect who ensures the foundation is perfect, and a Data-Driven Tuner who relentlessly experiments with tiny adjustments, guided by real-world performance data, to squeeze every ounce of speed out of the machine. It bridges the gap between "code that works" and "code that wins."

1. Problem Statement

Generating high-performance CUDA kernels is a significant bottleneck in scientific computing and machine learning. While Large Language Models (LLMs) have shown promise in synthesizing functionally correct CUDA code, they often fail to produce performance-optimized kernels. The challenges include:

Combinatorial Optimization Space: Optimizing GPU kernels involves navigating a vast, non-smooth space of low-level transformations (tiling, vectorization, shared-memory staging, loop restructuring) that interact complexly with hardware constraints (register pressure, occupancy, memory coalescing).
Noisy and Expensive Feedback: Performance feedback relies on hardware profiling (e.g., Nsight Compute), which is computationally expensive and noisy. Relying solely on end-to-end timing often obscures the root causes of inefficiency.
Separation of Generation and Tuning: Existing approaches typically treat code generation and optimization as separate problems. LLM-generated code often lacks the structural quality required for effective downstream optimization, and post-hoc optimization struggles to recover from weak or non-compilable starting points.

2. Methodology: The OptiML Framework

OptiML is an end-to-end framework that unifies program synthesis and kernel optimization. It accepts either natural-language intent or existing CUDA code and outputs a performance-optimized kernel. The framework consists of two decoupled stages:

A. OptiML-G: Mixture-of-Thoughts Generator (Proposal Policy)

When the input is natural language, OptiML-G synthesizes an initial executable kernel.

Architecture: It utilizes a Mixture-of-Thoughts (MoT) approach. Instead of relying on a single model, it maintains a pool of frozen, pre-trained code-generation experts (e.g., Qwen2.5-Coder, HPC-Coder-V2, StarCoder2).
Routing & Collaboration: A lightweight router selects the top- $K$ experts relevant to the input prompt. Unlike standard ensembling, OptiML-G enables latent-space collaboration. Selected experts process the prompt in parallel, and their hidden states are projected into a shared latent space. A primary expert then integrates information from all active peers via cross-attention mechanisms before generating the final token sequence.
Goal: To produce a structurally strong, compilable kernel that exposes meaningful opportunities for transformation, serving as a high-quality initialization for the optimizer.

B. OptiML-X: Search-Based Optimizer

OptiML-X refines either synthesized or user-provided kernels using Monte Carlo Tree Search (MCTS) guided by hardware feedback.

Search Process:
- State: Each node in the MCTS tree represents a program variant with measured feedback (runtime, Nsight Compute metrics).
- Action: An LLM (acting as a "Diagnoser" and "Proposer") analyzes the current state, identifies bottlenecks (e.g., memory-bound vs. compute-bound), and proposes specific code edits (patches).
- Evaluation: Candidates are compiled, tested for correctness (L0/L1 smoke tests), and profiled.
Reward Model: The system uses a composite reward function to guide the search:
1. Time Reward: Based on runtime improvement.
2. Proxy Reward: Derived from Nsight Compute metrics (e.g., SM/DRAM/TEX utilization, L1 sector traffic, instruction count). This allows the system to optimize for specific bottlenecks even before the final runtime is measured.
3. Guardrails: Hard penalties for candidates that violate correctness, exceed register limits, or cause compilation failures.
4. LLM-as-a-Judge: An LLM evaluates whether the proposed edit plausibly addresses the diagnosed bottleneck and if the measured signals support the claim.
Selection: The search uses Upper Confidence Bound (UCT) to balance exploration and exploitation. The final output is the best "KEEP" chain of edits found during the search.

3. Key Contributions

Unified Framework: OptiML is the first end-to-end system that maps natural-language specifications directly to optimized CUDA kernels, bridging the gap between generation and tuning.
Mixture-of-Thoughts Generator (OptiML-G): Introduces a latent-space collaboration mechanism among heterogeneous LLM experts to generate structurally superior initial kernels, significantly improving compilation success rates and starting quality.
Hardware-Aware Search (OptiML-X): Proposes a novel MCTS-based optimizer that uses profile-derived proxy metrics and an LLM-as-a-Judge to target specific hardware bottlenecks (e.g., memory traffic vs. instruction footprint) rather than relying on blind random search.
Interpretability: The framework provides interpretable optimization trajectories, attributing performance gains to specific hardware behaviors (e.g., reduced L1 traffic, improved SM utilization) rather than just black-box speedups.

4. Experimental Results

The framework was evaluated on a diverse suite of CUDA kernels (Matrix Multiplication, Max Pooling 3D, Multi-Head Self-Attention, ReLU, Cross-Entropy, etc.) on NVIDIA A100 GPUs.

Performance Gains:
- OptiML consistently outperforms strong LLM baselines (GPT-5, Qwen2.5-Coder, HPC-Coder-V2) and optimization-only variants.
- Speedups: Achieved up to 1.64x speedup on Matrix Multiplication and 1.67x on ReLU compared to the best LLM-only baseline.
- Reliability: Many LLM-only baselines failed to compile (CF) on complex tasks (e.g., Max Pooling 3D, Attention). OptiML-X successfully recovered valid, high-performance kernels in these cases, achieving 100% Pass@10 on the tested subset.
Hardware Efficiency:
- Optimized kernels showed significant improvements in hardware utilization (e.g., SM SOL increased from ~71% to ~81% in MatMul) and reductions in work proxies (e.g., L1 sectors and DRAM bytes decreased).
Ablation Studies:
- Generation Quality: OptiML-G achieved 12.17% Pass@1 on the ParEval benchmark, outperforming all baselines, proving that the MoT approach generates more reliable first-attempt code.
- Search Components: Removing the LLM Judge or proxy metrics significantly reduced both the success rate (Pass@k) and the magnitude of speedups, confirming the necessity of guided, hardware-aware search.
- Convergence: OptiML (Gen + X) reached the best solution in fewer MCTS rounds (e.g., 2.4 rounds vs. 4.8 for OptiML-X alone) because the initial kernel from OptiML-G was already close to an optimization-ready state.

5. Significance

Bridging the Gap: OptiML addresses the critical gap between "code that works" and "code that is fast," demonstrating that LLMs can be effectively guided by hardware feedback to produce production-grade kernels.
Sample Efficiency: By using proxy metrics and LLM diagnosis, the system reduces the number of expensive compile-profile cycles required to find optimal solutions, making automated kernel optimization feasible for complex workloads.
Generalizability: The framework is not limited to specific kernels; it adapts its optimization strategy based on the inferred bottleneck (memory-bound vs. compute-bound), making it robust across diverse application domains.
Future Impact: This work paves the way for automated, high-performance accelerator programming where LLMs handle the logic and structure, while system-level search ensures peak hardware utilization, reducing the reliance on manual expert tuning.