KernelBrain: Coarse-to-Fine, Budget-Aware Search for Agentic GPU Kernel Optimization

기술 요약: KernelBrain

문제 정의

GPU 커널 최적화를 자동화하는 것은 여전히 매우 까다로운 실무적 과제입니다. 대규모 언어 모델(LLM)이 커널 변형을 생성할 수는 있지만, 실제 프로덕션 배포 시에는 세 가지 결정적인 장애물에 직면합니다:

1. 정확성 퇴보(Correctness Regressions): 생성된 코드가 제약 조건을 위반하거나 컴파일에 실패하는 경우가 빈번합니다.
2. 노이즈가 섞인 측정값(Noisy Measurements): 런타임 성능 순위가 측정 노이즈로 인해 불안정하며, 이는 잘못된 최적화 결정으로 이어집니다.
3. 탐색 정체(Search Stagnation): 최적화 지형은 "비단조적(non-monotonic)"입니다. 즉, 작은 코드 변경이 하드웨어 의존적인 거대한 성능 변화나 일시적인 성능 저하를 유발할 수 있으며, 이로 인해 탐색 프로세스가 정체되거나 성능이 낮은 프로그램으로 표류하게 됩니다.

기존의 자동화 프레 framework들은 심층적인 구조적 리팩토링 능력이 부족하여 알고리즘적 전환보다는 점진적인 이득에 치중하며, 자원을 많이 소모하는 멀티 에이전트 방식은 안정성을 보장하지 못한 채 프로파일링과 탐색에 과도한 예산을 소비합니다.

방법론: KernelBrain

KernelBrain은 노이즈와 중단 상황에서도 안정적이고 효율적인 최적화 루프를 유지하도록 설계된, 실용적이고 예산 효율적인 최적화 에이전트입니다. 이 시스템은 비대칭적 자원 할당(asymmetric resource allocation) 원칙을 기반으로 작동하며, LLM 유도형 변이(mutation)와 하드웨어 인 더 루프(hardware-in-the-loop) 피드백을 결합합니다.

핵심 아키텍처

시스템은 적응형 스케줄러에 의해 연결된 두 개의 병렬 평면으로 구성됩니다:

1. 평가 평면(Evaluation Plane):

   • 정책 검사($V(p)$): 컴파일 유효성, 함수 시그니처, 참조 구현에 대한 출력 정확성을 검증하는 프리 벤치마크 게이트입니다. 이를 통과하지 못한 후보는 즉시 거부됩니다.
   • 다중 충실도 벤치마킹(Multi-Fidelity Benchmarking): 후보들은 여러 단계의 "러그(rung, 계단)"를 통해 평가됩니다. 낮은 단계의 러그는 빠른 스크리닝을 위해 최소한의 자원(적은 반복 횟수, 좁은 입력 형태)을 사용합니다. 높은 단계의 러그는 더 많은 예산(많은 반복 횟수, 넓은 형태 커버리지)을 할당하며, 선택적으로 GPU 커널 프로파일링(NVIDIA Nsight Compute 활용)을 트리거하여 SM 점유율, 메모리 처리량, 파이프라인 스톨 등의 메트릭을 수집합니다.
   • 스코어링: 점수 $S(p, r)$는 지연 시간(처리량)과 러그별 보너스를 기준으로 계산되며, 안정성과 속도를 입증하는 후보를 우선시합니다.

2. 생성 평면(Generation Plane):

   • 전문가 가이드형 조대-세밀 탐색(Expert-Guided Coarse-to-Fine Search): 맹목적인 진화적 변이 대신, 시스템은 초기 라운드에서 주요 최적화 공간(알고리즘, 타이링, 메모리 레이아웃)을 포괄하기 위해 전문가의 힌트를 사용합니다.
   • 프로파일링 기반 변이: 높은 단계의 러로 승격된 후보들을 위해, 시스템은 병목 현상을 진단하고자 프로파일러 데이터를 사용합니다. **프롬프트 빌더(Prompt Builder)**는 LLM을 위한 구조화된 프롬프트를 구축하여, 무분서한 재작성이 아닌 특정 식별된 문제(예: 메모리 코알레싱, 레지스터 압박)를 해결하도록 변이를 유도합니다.
   • 파서/수정(Parser/Repair): 결정론적 모듈이 LLM 응답에서 커널 코드를 추출하고, 일반적인 오류(누락된 import, 깨진 시그니처)에 대한 수정을 적용하여 루프가 견고하게 유지되도록 합니다.

최적화 루프

프로세스는 온라인 다중 충실도 루프(알고리즘 1)를 따릅니다:

1. 후보와 할당된 러그를 팝(pop)합니다.
2. 정확성을 검증합니다. 실패 시 로그를 남기고 폐기합니다.
3. 현재 충실도 수준에서 벤치마킹을 수행합니다.
4. 점수를 계산하고, 해당되는 경우 커널을 프로파일링합니다.
5. 스케줄러는 후보를 더 높은 러그로 승격(예산을 늘리고 변이 트리거)할지 또는 **가지치기(prune)**할지 결정합니다.
6. 승격된 후보는 프로파일링 진단을 바탕으로 LLM 변이를 거친 후 다음 러그의 큐로 다시 진입합니다.

주요 기여

1. 비대칭적 자원 할당: 시스템은 성능이 낮은 후보를 빠르게 스캔하고 제거하기 위해 최소한의 자원을 할당하며, 유망한 생존자들에게는 다중 충실도 체계를 통해 대부분의 평가 예산을 집중합니다. 여기에는 명시적인 정책 게이팅과 안정성 판정이 포함됩니다.
2. 전문가 가이드형 조대-세밀 탐색: 맹목적인 LLM 변이를 넘어, 전문가 힌트를 사용하여 초기에 광범위한 최적화 공간을 커버한 뒤, 하드웨어 프로파일링에서 도출된 추적 가능한 병목 현상 인지형 변이 가이드를 사용하여 고성능 영역으로 탐색 범위를 점진적으로 좁힙니다.
3. 최첨단 Triton 성능: 본 시스템은 Triton 기반 워크로드에서 경쟁력 있는 결과를 보여주며, PyTorch 및 기존 에이전트보다 상당한 속도 향상을 달면서 최적화 시간을 단축합니다.

실험 결과

시스템은 NVIDIA A100 GPU에서 6개의 중요한 DNN 벤치마크(GroupConv, MatVec, RMSNorm, Softmax, LinearReLUDiv, SDPAttn)를 대상으로 평가되었습니다.

• 성능 향상: KernelBrain은 PyTorch 참조 모델 대비 0.88배에서 6.72배 사이의 속도 향상을 달성했습니다.
• SOTA와의 비교: SOTA 에이전트 기반 커널 최적화 시스템인 KernelAgent보다 최대 1.4배의 속도 향상을 달성했습니다.
• 효율성: 이 접근 방식은 베이스라인 대비 최적화 시간을 최대 48%(기하 평균) 단축했습니다.
• 특정 벤치마크:
  • GroupConv: PyTorch 대비 5.7배~6.7배의 속도 향상을 달성하여 기본 LLM 생성기와 KernelAgent를 크게 상회했으며, 문제 크기가 커질수록 이득이 증가했습니다.
  • RMSNorm & Softmax: 비탐색(non-search) 베이스라인 대비 일관된 개선(각각 2.5x–2.8x 및 1.25x–1.33x)을 보여주었습니다.
  • SDPAttn: 긴 시퀀스 길이에서 KernelAgent보다 우수한 성능(1.17x–1.54x)을 보였으나, 짧은 시퀀스에서는 런치 오버헤드로 인해 PyTorch보다 낮은 결과를 보이기도 했습니다.

최적화 패턴에 대한 관찰:
실험 결과, 가장 큰 이득은 **작업의 병렬 분해(parallel decomposition of work)**를 변경하거나(예: 직렬 리덕션에서 협력적 분할 리덕션으로 이동), 글로벌 메모리 트래픽을 줄이는 것(예: 멀티 패스 구현을 퓨즈드 또는 온라인 리덕션으로 교체)에서 발생했습니다. 시스템은 단순한 국소적 편집에 의존하기보다 이러한 구조적 변화를 성공적으로 식별해 냈습니다.

의의

본 논문은 KernelBrain이 숙련된 GPU 프로그래머의 워크플로우를 구현하고 있다고 주장합니다. 즉, 다양한 알고리즘 및 구현 전략에 대한 광범위하고 저비용인 탐색에서 시작하여, 증거가 축적됨에 따라 가장 유망한 후보에 예산을 점진적으로 집중하는 방식입니다. 하드웨어 프로파일링을 변이 루프에 직접 통합하고 엄격한 정확성 게이트를 적용함으로써, KernelBrain은 커널 최적화의 "비단조적" 특성을 해결하고, 전통적인 전수 조사 방식의 자원 낭비나 가이드 없는 LLM 생성의 불안정성 없이 고성능 커널을 진화시키는 안정적이고 예산 효율적인 방법을 제공합니다.