CODA: 트랜스포머 블록을 GEMM 에필로그 프로그램으로 재작성
AI 모델 학습 시 흔히 발생하는 메모리 병목 현상을 해결하기 위해, 개별적으로 처리되던 연산들을 하나의 GPU 커널(GEMM Epilogue)로 통합하여 성능을 극대화하는 새로운 추상화 기법인 CODA를 제안합니다. 이 방식은 데이터 이동을 최소화하면서도 프레임워크 수준의 생산성과 하드웨어 수준의 극적인 효율성을 동시에 달성할 수 있도록 돕습니다.
--> 컴퓨터 과학 > 머신러닝 arXiv:2605.19269 (cs) [2026년 5월 19일 제출 (v1), 2026년 5월 20일 최종 수정 (현재 버전, v2)]
제목: CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs 저자: Han Guo, Jack Zhang, Arjun Menon, Driss Guessous, Vijay Thakkar, Yoon Kim, Tri Dao
초록: 트랜스포머(Transformer) 학습 시스템은 밀집 선형 대수(Dense linear algebra)를 기반으로 구축되지만, 엔드투엔드(end-to-end) 소요 시간 중 상당 부분은 메모리 제약을 받는(memory-bound) 주변 연산자들에 의해 소모됩니다. 정규화(Normalization), 활성화 함수(Activations), 잔차 업데이트(Residual updates), 리덕션(Reductions) 및 관련 연산들은 적은 산술 연산을 수행하면서도 대용량 중간 텐서를 반복적으로 전역 메모리(Global memory)로 이동시킵니다. 이로 인해 고도로 최적화된 학습 스택에서조차 데이터 이동이 점점 더 중요한 병목 현상으로 대두되고 있습니다.
본 논문에서는 이러한 연산들을 'GEMM(General Matrix Multiply)-plus-epilogue' 프로그램으로 표현하는 GPU 커널 추상화 모델인 CODA를 소개합니다. CODA는 개별적인 프레임워크 커널으로 노출되는 수많은 트랜스포머 연산자들을 대수적으로 재매개변수화(reparameterized)하여, GEMM의 출력 타일이 메모리에 기록되기 전에 온칩(On-chip) 상태에 머무는 동안 실행할 수 있다는 관찰에 기반합니다. 이러한 추상화는 GEMM의 메인 루프(Mainloop)를 고정하고, 스케일링(Scaling), 리덕션, 쌍대 변환(Pairwise transformations) 및 누적(Accumulation)을 위한 소규모 구성 가능한 에필로그 기본 요소(Primitives) 세트를 제공합니다.
이러한 제한된 인터페이스는 전문가가 직접 작성한 GEMM의 고성능 구조를 유지하면서도, 표준 트랜스포머 블록의 순전파(Forward pass) 및 역전파(Backward pass)에서 어텐션(Attention)을 제외한 거의 모든 연산을 포괄할 만큼 충분한 표현력을 갖추고 있습니다. 대표적인 트랜스포머 워크로드 전반에 걸쳐, 사람과 대형 언어 모델(LLM)이 작성한 CODA 커널 모두 높은 성능을 달성했습니다. 이는 'GEMM-plus-epilogue' 프로그래밍이 프레임워크 수준의 생산성과 하드웨어 수준의 효율성을 결합하는 실용적인 접근 방식을 제공함을 시사합니다.
주제: 머신러닝 (cs.LG) 인용: arXiv:2605.19269 [cs.LG]로 인용 (또는 이 버전의 경우 arXiv:2605.19269v2 [cs.LG]) https://doi.org/10.48550/arXiv.2605.19269