기술 요약: Axon: 텐서 프로그램을 위한 합성형 슈퍼옵티마이저(Synthesizing Superoptimizer)

문제 정의

Amazon의 Trainium과 같은 AI 가속기를 위한 고성능 커널을 작성하려면 메모리 계층 구조, 타이링 전략(tiling strategies), 명령어 선택, 데이터 레이아웃, 연산자 퓨전(operator fusion)을 포함한 하드웨어 특화 세부 사항에 대한 깊은 전문 지식이 필요합니다. NKI(Neuron Kernel Interface)와 같은 타일 기반 커널 언어는 이러한 저수준의 세부 사항을 추상화하지만, 피크 하드웨어 활용도를 달로는 데는 여전히 많은 ML 엔지니어들에게 큰 도전 과제입니다. 기존의 최적화 접근 방식은 다음과 같은 몇 가지 한계점을 가집니다:

• 수동 제작 규칙(Hand-crafted rules): 많은 시스템이 대수적 변환을 위한 수동으로 작성된 리라이트 규칙(rewrite rules)에 의존하며, 이는 오류가 발생하기 쉽고 새로운 연산자 조합으로 확장하기 어렵습니다.
• 템플릿 기반 생성(Template-based generation): 일부 접근 방식은 코드 생성을 위해 수동으로 작성된 템플릿을 사용하므로 유연성이 제한됩니다.
• 제한된 범위(Limited scope): 많은 옵티마이저가 CPU 또는 GPU만을 대상으로 하며, 현대적인 AI 가속기의 독특한 멀티 엔진 및 DMA 기반 아키텍처를 다루는 데 실패합니다.
• 증명 가능한 정확성 결여(Lack of provable correctness): 확률적 테스트나 휴리스틱 기반의 재정렬은 최적의 구성을 놓치거나 의미론적 오류를 도입할 수 있습니다.

방법론

Axon은 타일 기반 AI 가속기 프로그램을 위한 합성형 슈퍼옵티마이저입니다. 이 시스템은 고수준 텐서 프로그램(NumPy와 유사한 인터페이스로 작성됨)을 받아들이고, 다단계 파이프라인을 통해 최적화된 타겟 명령어(NKI)를 자동으로 생성합니다. 핵심 혁신은 $\nu$Graph("nu-graph")로, 이는 합성 과정 전반에 걸쳐 모든 의미론적으로 동등한 프로그램 변형들을 동시에 추적하는 지속적인 내부 표현(internal representation)입니다. 이를 통해 최적화 선택을 경험적 선택(empirical selection)으로 미룹니다.

파이프라인은 다음 단계로 구성됩니다:

1. 연산자 전파를 통한 대수적 변환(Algebraic Transformations via Operator Propagation):

   • Axon은 연산자를 계산 그래프를 통해 전파함으로써 의미론적으로 동등한 계산 그래프를 발견합니다.
   • 병렬성과 퓨전 기회를 노출하기 위해 인접한 연산자의 순서를 변경하려고 시도합니다 (예: 요소별 곱셈을 행렬 곱셈 뒤로 이동).
   • 정확성 보장: 수동으로 작성된 리라이트 규칙 대신, Axon은 텐서 연산을 세밀한 계산 단위로 분해하고, Unbounded 텐서에 대한 SMT(Satisfiability Modulo Theories) 솔버를 사용하여 재정렬이 의미론을 보존하는지 증명 가능하게 확인합니다. 이는 임의의 텐서 형상을 처리하며 별도의 수동 속성 정의를 요구하지 않습니다.

2. 심볼릭 타이링(Symbolic Tiling):

   • 시스템은 스트립(strip), 블록(block), 타일(tile) 차원에 대한 심볼릭 파라미터를 사용하여 각 텐서 연산을 타이링합니다.
   • 이러한 파라미터는 하드웨어 설명(예: 온칩 버퍼 용량 및 명령어별 제한)에 의해 제약됩니다.
   • Axon은 서로 다른 타이링 구성을 가진 여러 버전을 생성하여, 병렬 브랜치(예: RMSNorm과 MatMul 브랜치를 서로 다른 컴퓨팅 엔진에서 실행하도록 분리)를 독립적으로 최적화할 수 있게 합니다.

3. ISA 합성(Instruction Selection):

   • 스케치 기반 프로그램 합성(sketch-driven program synthesis) 접근 방식을 사용하여, Axon은 고수준 텐서 연산을 타겟 ISA 명령어(예: Trainium의 nc_matmul, dma_copy)로 낮춥니다(lower).
   • 특정 하드웨어 명령어에 필요한 유효한 명령어 시퀀스와 데이터 레이아웃 변환(예: transpose)을 탐색합니다.
   • 고수준 연산과 합성된 명령어 시퀀스 사이의 동등성은 다시 한번 SMT를 통해 검증됩니다.

4. ISA 및 연산자 퓨전(ISA and Operator Fusion):

   • 합성 과정에서 생성된 새로운 퓨전 기회를 노출하기 위해(예: transpose와 matrix multiplication 퓨전), 합성된 ISA 그래프에 대해 두 번째 라운드의 대수적 변환을 수행합니다.
   • 노드 퓨전(Node Fusion): 중간 결과물을 온칩 버퍼(SBUF/PSUM)에 유지하여 HBM 트래픽을 줄이기 위해 호환 가능한 타이링 차원을 가진 연속적인 연산들을 병합합니다.
   • 명령어 퓨전(Instruction Fusion): ISA 명령어 시퀀스를 퓨전된 변형으로 대체합니다 (예: dma_copy + nc_transpose $\rightarrow$ dma_transpose).

5. 코드 에미션 및 경험적 선택(Code Emission and Empirical Selection):

   • Axon은 심볼릭 파라미터에 구체적인 값을 할당하여 구체적인 NKI 프로그램을 출력합니다.
   • 각 후보는 Neuron 컴파일러에 의해 컴파일되고 타겟 하드웨어에서 실행됩니다.
   • 가장 성능이 좋은 커널은 참조 구현(reference implementation)에 대한 수치적 정확성 검증 및 경험적 레이턴시 측정을 바탕으로 선택됩니다.

주요 기여

본 논문은 세 가지 주요 기여를 제시합니다:

1. 증명 가능한 대수적 변환: 연산자를 전파하고 Unbounded 텐서에 대한 SMT를 사용하여 동등성을 검증함으로써 의미론적으로 동등한 계산 그래프를 자동으로 발견하는 방법론을 제시합니다. 이는 수동으로 작성된 리라이트 규칙의 필요성을 제거합니다.
2. 합성 기반 파이프라인: 대수적 변환, 명령어 선택, 타이링 구성 및 퓨전을 공동으로 탐색하는 완전한 타일 기반 AI 가속기용 파이프라인을 제공합니다. 이 시스템은 모든 유효한 변형을 $\nu$Graph에 유지하고 경험적으로 최적의 커널을 선택합니다.
3. 포괄적인 평가: 개별 연산자와 복잡한 멀티 연산자 LLM 커널(예: Group Query Attention, Gated MLP)을 포함한 20개의 벤치마크에 대해 Amazon의 Trainium에서 평가를 수행했습니다.

실험 결과

Axon은 Amazon의 Neuron 컴파일러, 수동 최적화된 NKI 커널, 그리고 Mirage(최첨단 검색 기반 컴파일러)라는 세 가지 베이스라인과 비교 평가되었습니다.

• Neuron 컴파일러 대비:

  • 개별 연산자: Axon은 큰 차원에서의 SiLU와 같은 경우 최대 3.7배의 속도 향상을 달성했으며, 다양한 연산자에 대해 1.23x–1.90x의 기하 평균 속도 향상을 기록했습니다.
  • 멀리 연산자 커널: Axon은 Transpose+MatMul과 같은 커널에서 최대 19배, RMSNorm+MatMul에서 최대 1.47배의 속도 향상을 달 achievement 했습니다. 기하 평균 속도 향상은 1.07x에서 1.32x 사이였습니다.
  • 이유: 이러한 이득은 더 나은 타이링 전략, 효율적인 DMA 활용, 그리고 서로 다른 컴퓨팅 엔진 간의 연산 퓨전(예: MatMul이 Tensor Engine에서 실행되는 동안 RMSNorm을 Scalar/Vector 엔진에서 병렬로 실행) 능력 덕분입니다.

• 수동 최적화된 NKI 대비:

  • Axon은 특정 구성(예: Cumsum 및 RoPE)에서 전문가가 작성한 커널보다 최대 1.35배 더 우수한 성능을 보이거나 대등한 성능을 보여, 자동화된 탐색이 깊은 하드웨어 전문 지식 없이도 수동 튜닝과 경쟁할 수 있음을 입증했습니다.

• Mirage 대비:

  • 두 시스템 모두 지원하는 벤치마크에서 Axon은 **2–10%**의 기하 평균 속도 향상을 달성했습니다.
  • 결정적으로, Axon은 수동 템플릿에 의依赖하는 Mirage가 처리할 수 없는 16개의 추가 벤치마크에 대해 성공적으로 커널을 합성했습니다.

• 컴파일 오버헤드:

  • 탐색 공간은 복잡한 커널에 대해 최대 ~30,000개의 후보를 생성할 수 있습니다.
  • 1단계(심볼릭 합성)는 빠릅니다(<82초). 2단계(컴파일 및 실행)가 시간을 지배하며, 복잡한 LLM 레이어의 경우 최대 ~1.7시간이 소요됩니다. 논문에서는 이것이 스케줄 기반 컴파일러들과 유사한 일회성 비용이며, 최적화된 커널은 재사용을 위해 캐싱될 수 있다고 언급합니다.

의의 및 주장

본 논문은 Axon이 휴리스틱 기반 또는 템플릿 기반 최적화에서 의미론 중심의 합성형 접근 방식으로의 전환을 나타낸다고 주장합니다. SMT 솔버를 활용하여 의미론적 동등성을 보장함으로써, Axon은 수동으로 인코딩하기 어려운 복잡한 대수적 변환(예: 서로 다른 컴퓨팅 엔진 간의 연산 재정렬)을 자동으로 발견할 수 있습니다.

그 의의는 다음과 같습니다:

• 자동화: 프로그래머가 특정 하드웨어에 맞춰 커널을 수동으로 튜닝해야 하는 부담을 줄입니다.
• 정확성: 확률적 테스트 방법과 달리 변환에 대한 증명 가능한 보장을 제공합니다.
• 성능: 통합된 탐색 공간 내에서 타이링, 명령어 선택 및 퓨전을 공동으로 최적화함으로써 현대적인 AI 가속기에서 상당한 속도 향상을 달성합니다.
• 범용성: 기존의 검색 기반 컴파일러(Mirage 등)가 합성할 수 없는 복잡한 멀리 연산자 LLM 커널을 처리할 수 있는 능력을 갖추고 있습니다.

저자들은 컴파일 오버헤드가 무시할 수 없는 수준이지만, 이는 최적화 컴파일러의 표준적인 일회성 비용이며, 특히 대규모 LLM 워크로드에 대해 얻을 수 있는 성능 이득이 그 투자를 정당화한다고 강조합니다.