KernelBlaster: 메모리 증강 인컨텍스트 강화 학습을 통한 지속적인 크로스-태스크 CUDA 최적화

KernelBlaster: Continual Cross-Task CUDA Optimization via Memory-Augmented In-Context Reinforcement Learning

여러 세대의 GPU 아키텍처에 걸쳐 CUDA 코드를 최적화하는 것은, 최고 성능을 달성하기 위해서는 점점 더 복잡하고 하드웨어에 특화된 최적화 공간을 광범위하게 탐색해야 하기 때문에 어려운 과제입니다. 기존 컴파일러는 고정된 휴리스틱에 의해 제한되는 반면, 대규모 언어 모델(LLM)의 미세 조정은 비용이 많이 들 수 있습니다. 그러나 CUDA 코드 최적화를 위한 에이전트 기반 워크플로우는 이전 탐색에서 얻은 지식을 통합하는 데 한계가 있어, 편향된 샘플링과 최적의 결과를 얻지 못하는 경우가 많습니다. 본 논문에서는 LLM 기반 GPU 코딩 에이전트의 CUDA 최적화 검색 능력을 향상시키기 위해 설계된 메모리 증강 인컨텍스트 강화 학습(MAIC-RL) 프레임워크인 KernelBlaster를 제안합니다. KernelBlaster는 에이전트가 경험을 통해 학습하고 검색 가능한 지속적인 CUDA 지식 베이스에 지식을 축적함으로써 향후 작업에 대해 체계적으로 정보에 입각한 결정을 내릴 수 있도록 합니다. 우리는 다양한 GPU 아키텍처 세대에 걸쳐 높은 성능을 달성하기 위해 프로파일 기반의 텍스트 기반 그래디언트 에이전트 흐름을 제안합니다. KernelBlaster는 LLM 에이전트가 단순한 재작성을 넘어 잠재력이 높은 최적화 전략을 체계적으로 탐색하도록 안내합니다. PyTorch 기준과 비교하여, KernelBlaster는 KernelBench 레벨 1, 2, 3에서 각각 기하 평균 속도 향상 1.43배, 2.50배, 1.50배를 달성했습니다. 우리는 KernelBlaster를 오픈 소스 에이전트 기반 프레임워크로 공개하며, 테스트 도구, 검증 구성 요소 및 재현 가능한 평가 파이프라인을 함께 제공합니다.

Optimizing CUDA code across multiple generations of GPU architectures is challenging, as achieving peak performance requires an extensive exploration of an increasingly complex, hardware-specific optimization space. Traditional compilers are constrained by fixed heuristics, whereas finetuning Large Language Models (LLMs) can be expensive. However, agentic workflows for CUDA code optimization have limited ability to aggregate knowledge from prior exploration, leading to biased sampling and suboptimal solutions. We propose KernelBlaster, a Memory-Augmented In-context Reinforcement Learning (MAIC-RL) framework designed to improve CUDA optimization search capabilities of LLM-based GPU coding agents. KernelBlaster enables agents to learn from experience and make systematically informed decisions on future tasks by accumulating knowledge into a retrievable Persistent CUDA Knowledge Base. We propose a novel profile-guided, textual-gradient-based agentic flow for CUDA generation and optimization to achieve high performance across generations of GPU architectures. KernelBlaster guides LLM agents to systematically explore high-potential optimization strategies beyond naive rewrites. Compared to the PyTorch baseline, our method achieves geometric mean speedups of 1.43x, 2.50x, and 1.50x on KernelBench Levels 1, 2, and 3, respectively. We release KernelBlaster as an open-source agentic framework, accompanied by a test harness, verification components, and a reproducible evaluation pipeline.