OptiML: 프로그램 합성 및 CUDA 커널 최적화를 위한 엔드투엔드 프레임워크

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

노이즈가 많고 비용이 많이 드는 하드웨어 피드백 환경에서 저수준 변환의 조합 공간을 탐색해야 하기 때문에, 고성능 CUDA 커널을 생성하는 것은 여전히 어려운 과제로 남아 있다. 대규모 언어 모델이 기능적으로 올바른 CUDA 코드를 합성할 수는 있지만, 경쟁력 있는 성능을 달성하려면 최적화 선택에 대한 체계적인 탐색과 검증이 필요하다. 본 논문에서는 커널 최적화를 검증 하의 탐색으로 공식화하여 자연어 의도나 입력된 CUDA 코드를 성능이 최적화된 CUDA 커널로 매핑하는 엔드투엔드 프레임워크인 OptiML을 제안한다. OptiML은 분리된 두 단계로 구성된다. 입력이 자연어일 경우, 사고 혼합(Mixture-of-Thoughts) 생성기(OptiML-G)가 커널 구현 전략에 대한 제안 정책 역할을 하여 초기 실행 가능한 프로그램을 생성한다. 그런 다음 탐색 기반 최적화기(OptiML-X)는 프로파일러 피드백에서 파생된 하드웨어 인식 보상의 안내를 받아, LLM 주도 편집에 대한 몬테카를로 트리 탐색을 사용하여 합성되거나 사용자가 제공한 커널을 개선한다. 각 후보 변환은 컴파일 및 검증 후 Nsight Compute를 통해 프로파일링되며, 실행 시간, 하드웨어 병목 프록시, 성능 저하 방지 가드레일을 결합한 복합 목표를 통해 평가된다. 우리는 다양한 CUDA 커널을 대상으로 합성 및 최적화, 최적화 전용 설정 모두에서 OptiML을 평가했다. 그 결과, OptiML은 강력한 LLM 베이스라인에 비해 검증된 성능 향상을 지속적으로 도출하며, 프로파일러 증거에 기반한 해석 가능한 최적화 궤적을 생성함을 보여준다.

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.