CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs

개요

CODA는 GPU 커널 추상화로, 트랜스포머 블록의 메모리 바운드 연산을 GEMM(General Matrix Multiply)-에필로그 프로그램으로 재구성하여 성능 병목 현상을 개선하는 방법을 제안합니다.

주요 내용

• 트랜스포머 훈련 시스템에서 행렬 곱셈(GEMM) 외에 발생하는 정규화, 활성화, 잔차 업데이트, 축소(reduction) 등의 연산이 상당한 시간을 차지하며, 이 연산들은 많은 양의 중간 텐서를 전역 메모리를 통해 이동시키지만 실제 연산량은 적어 데이터 이동이 병목 현상의 주된 원인이 됩니다.
• CODA는 이러한 연산들이 GEMM 연산의 결과 타일이 온칩(on-chip)에 머무는 동안 실행되도록 재매개변수화(reparameterize)될 수 있다는 관찰에 기반합니다.
• CODA 추상화는 GEMM 메인 루프를 고정하고, 스케일링, 축소, 쌍별 변환(pairwise transformation), 누적(accumulation)을 위한 작고 조합 가능한(composable) 에필로그 기본 요소(primitive)들을 제공합니다.
• 이 제약된 인터페이스는 전문가가 작성한 GEMM의 성능 구조를 유지하면서도, 트랜스포머 블록의 순방향 및 역방향 패스에서 발생하는 어텐션(attention)을 제외한 거의 모든 연산을 표현할 만큼 충분한 표현력을 가집니다.
• 대표적인 트랜스포머 워크로드에서 CODA 커널은 사람 또는 LLM이 작성했는지에 관계없이 높은 성능을 달성했습니다.

시사점

GEMM-에필로그 프로그래밍은 프레임워크 수준의 생산성과 하드웨어 수준의 효율성을 결합할 수 있는 실용적인 경로를 제공하여, 트랜스포머 모델 훈련의 성능 병목 현상을 완화하는 데 기여할 수 있습니다.