AccelOpt: AI 가속기 커널 최적화를 위한 자기 개선형 LLM 에이전트 시스템 기술 요약

본 논문은 새로운 AI 가속기 (특히 AWS Trainium) 에 대한 커널 최적화 문제를 해결하기 위해 제안된 AccelOpt라는 자기 개선형 대규모 언어 모델 (LLM) 에이전트 시스템에 대해 다룹니다. AccelOpt 는 전문가의 수동 최적화 지식이나 하드웨어 특화 최적화 레시피 없이도, 자율적으로 커널을 최적화하여 성능을 극대화하는 것을 목표로 합니다.

1. 문제 정의 (Problem)

• AI 가속기의 성능 병목: 대규모 언어 모델 (LLM) 시대의 컴퓨팅 수요 증가로 인해 AI 가속기 (Trainium 등) 가 급부상하고 있지만, 이들의 성능은 저수준의 커널 (Kernel) 구현 효율성에 크게 의존합니다.
• 최적화의 어려움: 기존 GPU 와 달리 새로운 AI 가속기 (Trainium 등) 는 아키텍처가 다르며, 개발자들이 이를 최적화할 수 있는 경험적 지식 (Heuristics) 이나 최적화 레시피가 부족합니다.
• 비용과 확장성: 수백, 수천 개의 커널을 다양한 설정과 하드웨어 버전에서 최적화해야 하는 상황에서, 전문가 수준의 결과를 비용 효율적으로 달성하는 것이 필수적입니다.
• 기존 접근법의 한계: 기존 LLM 기반 커널 생성 연구들은 주로 GPU 나 TPU 에 집중되어 있으며, 새로운 플랫폼에 대한 최적화 지식이 부족할 때 LLM 이 스스로 탐색하고 학습하여 성능을 개선하는 메커니즘은 미비했습니다.

2. 방법론 (Methodology)

AccelOpt 는 **빔 서치 (Beam Search)**와 **최적화 메모리 (Optimization Memory)**를 결합한 자기 개선형 에이전트 워크플로우를 사용합니다.

2.1 핵심 구성 요소

1. 에이전트 워크플로우 (Agentic Workflow):

   • Planner (기획자): 프로파일링 결과와 과거 경험을 바탕으로 성능 병목 (예: 메모리 대역폭, 불필요한 계산) 을 분석하고 최적화 계획을 수립합니다.
   • Executor (실행자): 수립된 계획을 구체화하여 실제 NKI (Neuron Kernel Interface) 커널 코드를 생성하고 수정합니다.
   • Summarizer (요약자): 느린 커널과 빠른 커널의 쌍 (Slow-Fast Pairs) 을 분석하여 일반화 가능한 최적화 전략 (예: "루프 밖으로 연산 이동", "중복 계산 제거") 을 추출하고 메모리에 저장합니다.

2. 빔 서치 (Beam Search):

   • 단순 반복 샘플링 (Repeated Sampling) 대신, 각 반복 (Iteration) 에서 가장 성능이 좋은 후보 커널들을 선별하여 다음 단계로 전달하는 빔 서치를 적용합니다.
   • 이는 탐색 공간을 전략적으로 탐색하고, 국소 최적해 (Local Optima) 에 빠지는 것을 방지하며 점진적인 개선을 유도합니다.

3. 최적화 메모리 (Optimization Memory):

   • 시스템이 경험한 "느린 커널 → 빠른 커널" 변환 사례들을 저장하는 지식 베이스입니다.
   • 성공적인 최적화 사례뿐만 아니라 실패한 시도 (Negative Rewrites) 도 포함하여 균형 잡힌 학습 신호를 제공합니다.
   • 이 메모리는 Planner 에게 컨텍스트로 제공되어, 이전 세대의 경험을 바탕으로 더 나은 최적화 전략을 수립하도록 돕습니다.

4. NKIBench (벤치마크):

   • AccelOpt 를 평가하기 위해 구축된 새로운 벤치마크로, 실제 LLM 워크로드에서 추출한 14 개의 복잡한 NKI 커널을 포함합니다.
   • 단순한 상대적 속도 향상 (Speedup) 이 아닌, 하드웨어의 이론적 피크 성능 (Theoretical Peak Performance) 대비 달성률을 기준으로 성능을 평가합니다.

3. 주요 기여 (Key Contributions)

1. 최초의 자기 개선형 LLM 에이전트 시스템: 하드웨어 특화 최적화 지식을 필요로 하지 않고, 검색 (Search) 과 메모리 축적 (Memory Accumulation) 을 결합하여 새로운 AI 가속기 (Trainium) 의 커널을 자율적으로 최적화하는 시스템을 처음 제안했습니다.
2. NKIBench 구축: AWS Trainium 의 NKI 커널 최적화를 위한 최초의 벤치마크를 구축하였으며, 이론적 피크 성능 대비 절대적인 성능 지표를 도입하여 평가의 객관성을 높였습니다.
3. 비용 효율성과 성능 증명: 오픈소스 LLM (Qwen3-Coder, gpt-oss 등) 을 사용하여 상용 모델 (Claude Sonnet 4) 과 동등하거나 더 나은 성능을 달성하면서도 비용을 26 배 절감함을 입증했습니다.
4. 기술적 통찰: 빔 서치가 단순 반복 샘플링보다 효과적이며, 최적화 메모리가 비용 효율성을 높인다는 것을 실험을 통해 검증했습니다.

4. 실험 결과 (Results)

• 성능 향상:
  • Trainium 1: 평균 피크 처리량 (Throughput) 비율이 **49% 에서 61%**로 향상되었습니다.
  • Trainium 2: 평균 피크 처리량 비율이 **45% 에서 59%**로 향상되었습니다.
  • 이 결과는 Claude Sonnet 4 (Thinking 모드) 의 성능과 유사하지만, 비용은 26 배 저렴합니다.
• 최적화 사례:
  • 국소 최적화: 대수적 단순화, 하드웨어 내장 함수 (Intrinsic) 융합 등.
  • 전역 최적화: 루프 변환 (Loop Transformation), 불변 코드 이동 (Loop Invariant Code Motion), 메모리 스풀링 (Spilling) 제거 등 복잡한 다단계 추론이 필요한 최적화를 성공적으로 발견했습니다.
• 전문가 대비 성능: 인간 전문가가 수동으로 최적화한 Mamba 및 RoPE 커널의 성능을 능가하거나 (Mamba: 54.6% vs 52.7%, RoPE: 29.6% vs 21.1%) 동등한 수준을 달성했습니다.
• 비용 분석: 오픈소스 모델 기반의 AccelOpt 가 상용 모델 반복 샘플링 대비 3.3 배 이상 비용 절감 효과를 보였으며, 모델 앙상블 전략을 통해 추가적인 성능 향상을 도모할 수 있음을 확인했습니다.

5. 의의 및 결론 (Significance)

AccelOpt 는 AI 가속기 생태계에서 지식 부족 (Knowledge Gap) 문제를 해결하는 획기적인 접근법을 제시합니다.

• 자율성: 새로운 하드웨어가 출시될 때마다 수개월의 전문가 최적화 시간을 단축하고, LLM 이 스스로 경험을 축적하며 성능을 개선할 수 있는 토대를 마련했습니다.
• 확장성: Trainium 과 같은 새로운 플랫폼뿐만 아니라, 더 성숙한 GPU 플랫폼이나 다른 AI 가속기로의 적용 가능성 (Generalization) 을 보여주었습니다 (FlashInfer-Bench 실험에서 Triton 커널 최적화 성공).
• 교육 및 연구: AccelOpt 는 단순한 도구 이상으로, 하드웨어 아키텍처와 커널 최적화 원리를 학습하는 교육 도구로서도 가치가 있으며, NKIBench 를 통해 향후 연구의 기준점 (Baseline) 을 제공합니다.

결론적으로, AccelOpt 는 추론 시간 확장 (Inference-time Scaling) 과 최적화 메모리를 결합함으로써, 전문가의 개입 없이도 고성능 AI 커널을 자동으로 생성하고 최적화할 수 있는 새로운 패러다임을 제시했습니다.