AdaExplore: 오류 기반 적응 및 다양성을 유지하는 탐색을 통한 효율적인 커널 생성

AdaExplore: Failure-Driven Adaptation and Diversity-Preserving Search for Efficient Kernel Generation

최근의 대규모 언어 모델(LLM) 에이전트는 실행 피드백을 활용하여 테스트 시간 동안 적응하는 데 유망한 결과를 보여왔습니다. 그러나 강력한 자체 개선은 아직 해결되지 않은 과제이며, 대부분의 접근 방식은 여전히 각 문제 인스턴스를 독립적으로 처리하며 재사용 가능한 지식을 축적하지 못합니다. 이러한 제한은 특히 Triton과 같은 도메인 특화 언어에서 두드러지는데, 이러한 언어는 LLM 사전 훈련 데이터에서 차지하는 비중이 낮습니다. 이들의 엄격한 제약 조건과 비선형 최적화 경관은 기본적인 생성 및 로컬 개선을 신뢰할 수 없게 만듭니다. 우리는 AdaExplore라는 에이전트 프레임워크를 제안합니다. AdaExplore는 성능이 중요한 커널 코드 생성을 위해 축적된 실행 피드백을 통해 자체 개선을 가능하게 하며, 추가적인 미세 조정이나 외부 지식 없이 두 가지 상호 보완적인 단계를 통해 정확성과 최적화 성능을 동시에 향상시킵니다. 적응 단계에서, 에이전트는 작업을 합성하고 반복되는 실패를 재사용 가능한 유효성 규칙의 메모리로 변환하여 후속 생성 과정에서 가능한 영역 내에 머물도록 돕습니다. 탐색 단계에서, 에이전트는 후보 커널을 트리로 구성하고 작은 로컬 개선과 더 큰 구조적 재생 사이를 번갈아 수행하여 로컬 최적값을 넘어 최적화 경관을 탐색할 수 있도록 합니다. 커널 런타임 최적화 벤치마크에 대한 실험 결과는 이러한 이점을 입증합니다. AdaExplore는 100단계 내에서 KernelBench Level-2 및 Level-3에서 각각 3.12배 및 1.72배의 속도 향상을 달성하며, 추가적인 계산을 통해 지속적으로 성능을 향상시킵니다.

Recent large language model (LLM) agents have shown promise in using execution feedback for test-time adaptation. However, robust self-improvement remains far from solved: most approaches still treat each problem instance independently, without accumulating reusable knowledge. This limitation is particularly pronounced in domain-specific languages such as Triton, which are underrepresented in LLM pretraining data. Their strict constraints and non-linear optimization landscape further make naive generation and local refinement unreliable. We propose AdaExplore, an agent framework that enables self-improvement via accumulated execution feedback for performance-critical kernel code generation through two complementary stages: failure-driven adaptation and diversity-preserving search, jointly improving correctness and optimization performance without additional fine-tuning or external knowledge. In the adaptation stage, the agent synthesizes tasks and converts recurring failures into a reusable memory of validity rules, helping subsequent generations remain within the feasible set. In the search stage, the agent organizes candidate kernels as a tree and alternates between small local refinements and larger structural regeneration, allowing it to explore the optimization landscape beyond local optima. Experiments on kernel runtime optimization benchmarks validate these gains: AdaExplore achieves 3.12x and 1.72x speedups on KernelBench Level-2 and Level-3, respectively, within 100 steps, and continues to improve with additional computation.