Kernel Forge: LLM 기반 CUDA 커널 생성 및 최적화를 위한 에이전트 시스템
Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels
Kernel Forge는 PyTorch 모델을 지원하는 오픈 소스 에이전트 시스템으로, 비전, 확산, LLM 작업을 처리할 수 있다. 이 시스템은 Monte Carlo Tree Search (MCTS)를 사용하여 여러 최적화 경로를 탐색하며, 진행 상황 모니터링, 후보 커널 검사 및 실패 디버깅을 위한 그래픽 사용자 인터페이스를 제공한다. NVIDIA DGX Spark에서 50회의 최적화 반복만으로 14개의 커널을 최적화하여 PyTorch eager mode보다 성능을 향상시켰고, ResNet-50의 adaptive_avgpool2d에서 $1.52 imes$, Stable Diffusion 3.5 Medium의 group_norm에서 $1.70 imes$, Gemma 4 E2B의 softmax에서 $2.83 imes$, Qwen 3.5 35B-A3B의 softmax에서 $1.54 imes$의 성능을 기록했다.
Kernel Forge는 PyTorch 모델을 최적화하여 성능을 향상시키는 데 효과적이며, 특히 다양한 작업에서 성능 개선을 보여준다.
원문 출처
arXiv cs.AI (인공지능)