AscendOptimizer Ascend NPU 연산자 최적화를 위한 에피소드형 에이전트

초록

AscendOptimizer는 Ascend NPU의 호스트‑tiling 프로그램과 디바이스‑커널을 동시에 최적화하는 두 단계 에이전트이다. 첫 단계에서는 하드웨어‑인‑루프 프로파일링을 이용한 진화적 탐색으로 유효하고 고성능의 tiling 구성을 찾고, 두 번째 단계에서는 기존 최적 커널을 의도적으로 ‘디‑옵티마이즈’해 만든 나쁜‑좋은 궤적을 패턴 라이브러리로 정제한다. 이 경험을 기반으로 커널을 구조적으로 재작성해 반복적으로 개선한다. 127개의 실제 AscendC 연산자를 대상으로 한 실험에서 평균 1.19배의 속도 향상을 달성했으며, 절반에 가까운 연산자가 기존 레퍼런스를 능가했다.

상세 분석

본 논문은 Ascend NPU와 같은 도메인‑특화 가속기에서 연산자 최적화가 겪는 ‘지식 병목’ 문제를 두 가지 관점에서 해결한다. 첫째, AscendC 연산자는 호스트‑side tiling 코드와 디바이스‑side 커널 코드가 서로 얽혀 있어, 하나만 최적화해도 전체 성능에 큰 영향을 주지 못한다는 점을 명확히 규정한다. 이를 위해 저자들은 최적화 문제를 (T, K) 쌍을 찾는 이중 탐색 문제로 공식화하고, 하드웨어 제약 H와 연산자 속성 S를 고정한 뒤 실행 지연 L을 최소화하는 목표함수를 제시한다.

Stage I에서는 tiling 공간이 이산적이고 작은 파라미터 변동에도 컴파일 오류나 성능 급락이 발생하는 ‘불연속성’ 특성을 활용한다. 진화 알고리즘을 사용해 초기 후보군을 생성하고, 하드웨어‑인‑루프 프로파일링을 피드백으로 삼아 적합도 함수를 정의한다. 이 과정에서 ‘실패‑성공’ 경계를 빠르게 탐지하고, 유효한 tiling 구성을 효율적으로 수렴시킨다.

Stage II는 커널 최적화에 초점을 맞춘다. 기존 최적화된 커널을 의도적으로 디‑옵티마이즈(예: 파이프라인 단계 제거, 벡터화 해제)하여 ‘bad‑to‑good’ 궤적을 만든다. 이러한 궤적은 ReWiND와 유사한 자기지도 학습 방식으로, 각 단계에서 발생하는 병목(메모리 대역폭, 레지스터 압박, 파이프라인 스타베이션 등)을 메타데이터와 함께 기록한다. 수집된 궤적을 패턴 라이브러리로 정제하고, 온라인 최적화 시점에 현재 tiling 상황에 가장 적합한 패턴을 검색·적용한다. 이때 패턴 매칭은 정형화된 rewrite 규칙 형태로 구현되어, 기존 코드를 구조적으로 변형하면서도 컴파일 가능성을 유지한다.

두 단계는 교차 루프 형태로 진행된다. Stage I에서 얻은 tiling이 Stage II의 커널 재작성 범위를 확대하고, 반대로 Stage II에서 개선된 커널은 새로운 tiling 후보의 성능 평가를 보다 정확히 반영한다. 이러한 상호 보완적 루프는 탐색 공간을 점진적으로 확장하면서도 수렴 속도를 높인다.

실험에서는 127개의 실제 AscendC 연산자를 대상으로, 오픈소스 베이스라인(기존 수작업 최적화) 대비 기하 평균 1.19배의 지연 감소를 기록했다. 특히 49.61%의 연산자가 레퍼런스 구현보다 빠른 성능을 보였으며, 기존 에이전트(예: Astra, PRAGMA)와 비교해 동일 하드웨어 환경에서 더 높은 성공률과 속도 향상을 달성했다. 또한, Table 2에서 제시된 바와 같이 AscendOptimizer는 ‘기존 구현 최적화’, ‘자동 최적화’, ‘학습‑프리’ 세 축을 모두 만족하는 유일한 시스템으로 평가된다.

본 연구는 (1) 하드웨어‑인‑루프 진화 탐색을 통한 tiling 자동화, (2) 최적 커널의 역방향 디‑옵티마이즈를 통한 경험 축적, (3) 경험 기반 패턴 재작성으로 커널을 구조적으로 개선한다는 세 가지 핵심 기법을 제시함으로써, 데이터와 레퍼런스가 부족한 DSA 환경에서도 전문가 수준의 최적화를 가능하게 한다는 점에서 의미가 크다. 향후 LLM과 결합한 코드 생성 단계에 이 경험 기반 패턴을 통합한다면, 완전 자동화된 연산자 설계 파이프라인 구축도 기대할 수 있다.