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

2026-02-12 · Arijit Bhattacharjee, Heng Ping, Son Vu Le, Paul Bogdan, Nesreen K. Ahmed, Ali Jannesari 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.

📄 PDF Abstract BibTeX arXiv:2602.12305

Code (0)

등록된 구현이 없습니다.

Tasks

Program Synthesis

CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation

2026-02-27 · Weinan Dai, Hanlin Wu, Qiying Yu, Huan-ang Gao 외 arxiv

GPU kernel optimization is fundamental to modern deep learning but remains a highly specialized task requiring deep hardware expertise. Despite strong performance in general programming, large language models (LLMs) rema…

Reinforcement Learning Code Generation

StitchCUDA: An Automated Multi-Agents End-to-End GPU Programing Framework with Rubric-based Agentic Reinforcement Learning

2026-03-03 · Shiyang Li, Zijian Zhang, Winson Chen, Yuebo Luo 외 arxiv

Modern machine learning (ML) workloads increasingly rely on GPUs, yet achieving high end-to-end performance remains challenging due to dependencies on both GPU kernel efficiency and host-side settings. Although LLM-based…

Reinforcement Learning Code Generation

Towards Cold-Start Drafting and Continual Refining: A Value-Driven Memory Approach with Application to NPU Kernel Synthesis

2026-03-11 · Yujie Zheng, Zhuo Li, Shengtao Zhang, Hanjing Wang 외 arxiv

Deploying Large Language Models to data-scarce programming domains poses significant challenges, particularly for kernel synthesis on emerging Domain-Specific Architectures where a "Data Wall" limits available training d…

Reinforcement Learning

DRTriton: Large-Scale Synthetic Data Driven Reinforcement Learning for Triton Kernel Generation

2026-03-23 · Siqi Guo, Ming Lin, Tianbao Yang arxiv

Developing efficient CUDA kernels is a fundamental yet challenging task in the generative AI industry. Recent research leverages Large Language Models (LLMs) to automatically convert PyTorch reference implementations to …

Reinforcement Learning

Evaluation of OpenAI Codex for HPC Parallel Programming Models Kernel Generation

2023-06-27 · William F. Godoy, Pedro Valero-Lara, Keita Teranishi, Prasanna Balaprakash 외

We evaluate AI-assisted generative capabilities on fundamental numerical kernels in high-performance computing (HPC), including AXPY, GEMV, GEMM, SpMV, Jacobi Stencil, and CG. We test the generated kernel codes for a var…