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

새로운 AI 가속기가 등장할 때마다 전문가가 수동으로 커널을 최적화하는 데 수개월이 걸리는 병목 현상을 해결합니다. LLM 에이전트가 하드웨어 지식 없이도 스스로 학습하며 최적화 코드를 생성하여, 상용 모델 대비 26배 저렴한 비용으로 전문가 수준의 성능을 달성할 수 있음을 보여줍니다.

기존의 커널 최적화는 하드웨어 아키텍처에 대한 깊은 이해를 가진 전문가가 수동으로 루프 변환이나 메모리 레이아웃을 조정해야 했다. 특히 AWS Trainium 같은 신규 가속기는 최적화 레시피가 부족하여 성능을 끌어올리는 데 막대한 시간이 소요된다. AccelOpt는 이를 LLM의 코드 생성 능력과 '경험 메모리'를 결합한 에이전트 워크플로우로 해결한다.

시스템은 먼저 현재 커널의 프로파일링 데이터를 분석하여 병목 지점을 찾고 최적화 계획을 세운다. 이후 실행기가 코드를 수정하고 실제 하드웨어에서 성능을 측정한다. 여기서 핵심은 'Optimization Memory'이다. 단순히 코드를 생성하는 데 그치지 않고, 성능이 향상된 코드 쌍에서 '왜 빨라졌는지'에 대한 기술적 통찰을 추출하여 메모리에 저장한다. 다음 단계의 Planner는 이 메모리를 참고하여 더 정교한 전략을 세우게 된다.

결과적으로 시스템은 반복을 거듭할수록 해당 하드웨어에 특화된 최적화 규칙을 스스로 학습하게 된다. 이는 전문가가 수개월간 시행착오를 겪으며 터득하는 최적화 직관을 LLM이 데이터 기반의 반복 탐색으로 대체할 수 있음을 의미한다.

AccelOpt는 Planner, Executor, Summarizer로 구성된 3단계 에이전트 워크플로우를 기반으로 작동한다. Planner는 커널 코드와 프로파일링 리포트를 입력받아 최적화 계획을 수립하고, Executor는 이를 구체적인 NKI(Neuron Kernel Interface) 코드로 구현하며, Summarizer는 결과물에서 핵심 최적화 패턴을 추출한다.

탐색 과정에는 Beam Search 알고리즘을 적용한다. 각 반복(Iteration)마다 B개의 후보 커널을 유지하며, 각 후보당 N개의 계획을 생성하고 K번의 구현 시도를 수행한다. [B개의 후보 → N*K개의 변형 생성 → 실제 하드웨어 프로파일링 → 상위 B개 선발] 과정을 통해 성능이 낮은 가지는 쳐내고 유망한 최적화 방향으로 탐색을 집중한다.

Optimization Memory Curation 단계에서는 속도 향상 비율(Speedup)이 임계값 t_pos를 넘는 사례를 수집한다. Summarizer는 [느린 커널 + 빠른 커널] 쌍을 분석하여 'Loop Invariant Code Motion' 같은 일반적인 최적화 용어로 요약하고 이를 의사코드로 변환한다. 이 정보는 다음 반복의 Planner 프롬프트에 'Past experiences'로 주입되어 탐색 공간을 효과적으로 제약한다.

NKIBench 테스트 결과, AccelOpt는 AWS Trainium 1에서 평균 피크 처리량을 49%에서 61%로, Trainium 2에서는 45%에서 59%로 향상시켰다. 이는 수동으로 최적화된 전문가 수준의 참조 구현 성능을 상회하는 수치이다.

비용 효율성 측면에서 Qwen3-Coder-480B를 Executor로 사용했을 때, 유료 모델인 Claude 4 Sonnet과 대등한 성능을 보이면서도 토큰 비용은 약 26배 저렴했다. 또한 Beam Search를 적용했을 때 단순 반복 샘플링 대비 누적 성능 향상 폭이 훨씬 컸으며, 메모리 메커니즘을 추가할 경우 동일 성능 도달 시간을 16-17% 단축할 수 있었다.

Ablation Study를 통해 메모리 용량(ExpN)을 늘리는 것이 메모리 업데이트 빈도(TopK)를 높이는 것보다 비용 대비 성능 향상에 더 효과적임을 확인했다. 특히 Mamba나 RoPE 같은 복잡한 커널에서도 전문가가 작성한 최적화 코드보다 각각 1.04배, 1.4배 빠른 결과를 도출해냈다.

AccelOpt는 NKI(Neuron Kernel Interface)라는 Python 기반 커널 언어를 타겟으로 한다. 시스템 아키텍처는 분산 프로파일링 서비스와 연동되어 대규모 커널 샘플을 병렬로 평가하며, Roofline Model을 사용하여 이론적 최대 성능 대비 효율성을 계산한다.

핵심 알고리즘인 Optimization Memory Curation은 수집된 커널들을 원본 후보와 계획별로 그룹화하고, 각 그룹 내에서 성능 아웃라이어를 선별하여 메모리에 추가한다. 이때 큐(Queue) 구조를 사용하여 오래된 경험은 삭제하고 최신 최적화 통찰을 유지한다. Planner는 이 메모리를 통해 하드웨어의 메모리 계층 구조와 병렬 처리 특성에 맞는 전략을 제안하게 된다.

연구팀은 또한 'Peephole Optimization'(대수적 단순화)과 'Loop Optimization'(루프 변환 및 타일링 조정) 등 LLM이 발견한 구체적인 최적화 사례를 분석하여, 에이전트가 단순한 코드 수정을 넘어 하드웨어 제약 조건을 고려한 복합적인 추론을 수행함을 입증했다.

현재 AccelOpt는 단일 코어 커널 최적화에 집중하고 있으며, 칩 간 통신(Cross-chip communication)이 포함된 멀티 노드 최적화는 다루지 않는다. 또한 초기 커널의 품질이 너무 낮거나 하드웨어 한계에 이미 근접한 경우 추가적인 개선을 찾는 데 어려움을 겪는 포화 현상이 관찰되었다.