CODA: Transformer 블록을 GEMM-Epilogue 프로그램으로 재작성하기

Transformer 학습 시스템은 dense linear algebra를 중심으로 구축되지만, normalization, activation, residual update, reduction 같은 메모리 병목 연산자들이 전체 시간의 상당 부분을 차지한다. 이들은 적은 연산으로 대용량 텐서를 반복적으로 글로벌 메모리로 이동시켜 데이터 전송이 주요 병목점이 된다.

CODA는 이러한 연산을 GEMM-plus-epilogue 프로그램으로 표현하는 GPU 커널 추상화다.

핵심 아이디어: 프레임워크에서 별도 커널로 노출된 많은 Transformer 연산자를 대수적으로 재매개변수화하여, GEMM 출력 타일이 칩에 남아있는 동안(메모리 쓰기 전) 실행한다.

GEMM 메인루프를 고정하고 scaling, reduction, pairwise transformation, accumulation을 위한 조합 가능한 epilogue 프리미티브 집합을 제공한다. 제한된 인터페이스지만 표준 Transformer 블록의 forward/backward pass에서 attention을 제외한 거의 모든 연산을 커버한다.

대표적인 Transformer 워크로드에서 사람과 LLM이 작성한 CODA 커널 모두 높은 성능을 달성했다. GEMM-plus-epilogue 프로그래밍이 프레임워크 수준 생산성과 하드웨어 수준 효율성을 결합하는 실용적 경로임을 보여준다.

저자진에는 FlashAttention으로 알려진 Tri Dao, Yoon Kim 등이 포함됐다.