Axe: 머신러닝 컴파일러를 위한 통합 레이아웃 추상화

초록

Axe 레이아웃은 논리 텐서 좌표를 이름이 붙은 하드웨어 축으로 매핑하여, 샤딩·복제·오프셋을 하나의 형식으로 표현한다. 이를 기반으로 다중‑Granularity DSL과 컴파일러를 설계해 스레드‑레벨 제어와 집합 연산을 하나의 커널에 결합하고, 최신 GPU와 AI 가속기에서 손수 튜닝한 코드에 근접한 성능을 달성한다.

상세 분석

본 논문은 현대 딥러닝 워크로드가 요구하는 “데이터·연산의 다중 스케일 배치” 문제를 해결하기 위해, 기존의 shape‑stride 모델을 확장한 Axe 레이아웃을 제안한다. 핵심 아이디어는 하드웨어 자원을 ‘축(axis)’이라는 이름으로 추상화하고, 각 축에 대해 Shard(D), Replica(R), Offset(O) 세 요소로 논리 인덱스를 물리 좌표 집합으로 변환한다는 점이다. D는 논리 차원을 여러 이터(iter)로 분할해 각각의 stride와 axis를 지정함으로써 다중‑축 샤딩을 표현한다. R은 논리 인덱스와 무관하게 복제(브로드캐스트) 동작을 정의하고, O는 고정 오프셋을 통해 메모리 뱅크나 레지스터 파일 내 특정 위치에 데이터를 고정시킨다. 이러한 3‑요소 모델은 GPU의 warp·lane·register, 다중 GPU mesh, 그리고 AWS Trainium·TPU와 같은 특수 메모리 구조까지 동일한 형식으로 기술할 수 있게 만든다.

논문은 Axe 레이아웃을 활용한 DSL을 설계한다. DSL은 텐서 선언 시 레이아웃 메타데이터를 명시하고, canonicalize, group, tile, slice 같은 연산자를 통해 레이아웃 변환을 선언적으로 기술한다. 컴파일러는 레이아웃 연산자를 기반으로 레인지 분석, 축 매칭, 스케줄 선택을 자동화한다. 특히, D와 R을 동시에 고려해 “스레드‑레벨 제어 + 집합 연산”을 하나의 커널에 합성함으로써, 기존 CuTe가 제공하는 저‑레벨 루프 바인딩과 Triton이 제공하는 블록‑레벨 추상화를 동시에 만족한다.

성능 평가에서는 최신 NVIDIA Blackwell GPU, Hopper GPU, 그리고 Trainium‑1 가속기에서 Axe 기반 커널이 FlashInfer, SGLang, cuBLAS+NCCL, Triton‑Distributed 등 강력한 베이스라인 대비 1.23‑1.44배 향상을 보였다. 특히, MoE 레이어와 Multi‑GPU GEMM+Reduce‑Scatter에서 손수 튜닝된 구현과 거의 동일한 실행 시간을 기록했다. 이는 레이아웃 추상화가 데이터 이동 최소화와 연산자 병렬도 최적화를 동시에 달성했음을 의미한다.

한계점으로는 현재 구현이 NVIDIA GPU와 Trainium에 최적화돼 있어, 다른 벤더(예: AMD GPU, Habana)에서는 추가적인 하드웨어‑축 정의가 필요하다. 또한, 레이아웃 변환 연산이 복잡해질 경우 컴파일 타임이 증가할 가능성이 있으며, 자동 스케줄링 정책이 모든 경우에 최적을 보장하지는 않는다. 향후 연구에서는 동적 레이아웃 재조정, 멀티‑스테이지 파이프라인 지원, 그리고 다중‑벤더 공통 축 정의를 통한 포팅을 제안한다.

전반적으로 Axe는 “데이터·연산 배치를 하나의 수학적 모델로 통합”함으로써, 머신러닝 컴파일러 설계의 복잡성을 크게 낮추고, 생산성과 성능 사이의 격차를 좁히는 중요한 진전을 제공한다.