KernelAgent: 하드웨어 가이드 기반 Triton 커널 최적화 에이전트

섹션별 상세

01

KernelAgent는 하드웨어 프로파일링 신호를 폐쇄 루프 멀티 에이전트 워크플로우에 통합하여 Triton 커널을 최적화한다. 이 시스템은 ProfilerAgent, JudgeAgent, AnalyzeAgent, OrchestratorAgent 등 특화된 에이전트들이 협력하여 병목 현상을 진단하고 최적화 전략을 수립한다. 모든 단계는 NVIDIA Nsight Compute(NCU)에서 추출한 실제 하드웨어 메트릭을 기반으로 작동한다.

02

최적화 프로세스는 프로파일링, 진단, 처방, 오케스트레이션, 탐색, 측정의 6단계로 구성된다. ProfilerAgent가 DRAM 처리량, L2 캐시 히트율, 워프 점유율 등의 데이터를 수집하면, DiagnoseAgent가 이를 해석하여 성능 병목의 근본 원인을 분류한다. 이후 AnalyzerAgent가 GPU 아키텍처(A100 vs H100)에 최적화된 구체적인 수정 사항을 제안한다.

KernelAgent의 하드웨어 가이드 기반 최적화 워크플로우 다이어그램 — DiagramProfiler, Judge, Analyze, Orchestrator 등 각 에이전트 간의 데이터 흐름과 폐쇄 루프 구조를 보여준다. 하드웨어 신호가 어떻게 진단과 처방으로 이어지는지 시각화하여 시스템 아키텍처를 설명한다.

03

OrchestratorAgent는 과거의 최적화 시도 이력을 분석하여 현재 라운드의 탐색 전략을 결정한다. 이 과정에서 'Was the diagnosis correct?'와 같은 자기 분석(Self-analysis)을 수행하여 실패한 패턴을 피하고 성공 가능성이 높은 경로를 선택한다. 공유 메모리 구조를 통해 에이전트들이 라운드 간 학습 내용을 공유하며 로컬 미니마에 빠지는 것을 방지한다.

04

KernelBench L1 과제 100개를 대상으로 평가한 결과, KernelAgent는 이전 버전 대비 2.02배의 기하 평균 속도 향상을 달성했다. 기본 torch.compile과 비교했을 때 65개 과제에서 더 우수한 성능을 보였으며 평균 1.56배의 속도 향상을 기록했다. 특히 H100 GPU에서 하드웨어 이론적 최대 성능(SOL)의 89%에 도달하는 효율성을 입증했다.

최적화 라운드별 성능이 개선된 커널의 수 — Chart라운드가 진행됨에 따라 성능이 향상되는 커널의 수가 점진적으로 줄어드는 수렴 과정을 보여준다. 1라운드에서 가장 많은 개선이 일어나며, 이후 라운드에서도 지속적인 미세 조정이 이루어짐을 입증한다.

05

행렬-벡터 곱셈(Matrix-Vector Multiplication) 사례 연구에서 KernelAgent는 레지스터 압박으로 인한 점유율 제한 문제를 해결했다. 대형 벡터 누산기를 스칼라 누산기로 교체하고 그리드 병렬성을 높이는 처방을 통해 실행 시간을 9.52ms에서 1.95ms로 단축했다. 이는 하드웨어 피드백 없는 일반 LLM 기반 최적화가 로컬 미니마에 갇히는 것과 대조적인 성과이다.