CODA: 트랜스포머 블록을 GEMM‑에필로그 프로그램으로 재작성
출처: Hacker News
초록: Transformer 학습 시스템은 밀집 선형 대수를 중심으로 구축되지만, 전체 실행 시간 중 상당 부분이 메모리 바인드 연산에 소비됩니다. 정규화, 활성화, 잔차 업데이트, 축소 연산 및 이와 관련된 계산은 큰 중간 텐서를 전역 메모리로 반복적으로 이동시키면서 실제 연산은 거의 수행하지 않으므로, 데이터 이동이 고도로 최적화된 학습 스택에서도 점점 더 중요한 병목 현상이 되고 있습니다. 우리는 이러한 계산을 GEMM‑plus‑epilogue 프로그램으로 표현하는 GPU 커널 추상화인 CODA를 소개합니다. CODA는 많은 Transformer 연산이 별도의 프레임워크 커널로 노출되지만, GEMM 출력 타일이 칩 위에 남아 있는 동안(메모리로 기록되기 전에) 연산을 수행하도록 대수적으로 재파라미터화할 수 있다는 관찰에 기반합니다. 이 추상화는 GEMM 메인 루프를 고정하고, 스케일링, 축소, 쌍별 변환, 누적을 위한 소수의 조합 가능한 에필로그 프리미티브를 노출합니다. 이 제한된 인터페이스는 전문가가 작성한 GEMM의 성능 구조를 유지하면서도 표준 Transformer 블록의 전방 및 역방향 패스에서 거의 모든 비‑어텐션 연산을 포괄할 수 있을 만큼 표현력이 충분합니다. 대표적인 Transformer 워크로드 전반에 걸쳐, 인간이 작성한 커널이든 LLM이 생성한 커널이든 CODA 커널은 높은 성능을 달성하며, GEMM‑plus‑epilogue 프로그래밍이 프레임워크 수준의 생산성과 하드웨어 수준의 효율성을 결합하는 실용적인 경로임을 시사합니다.
주제:
Machine Learning (cs.LG)
인용:
arXiv:2605.19269 [cs.LG]
(또는 이 버전은 arXiv:2605.19269v2 [cs.LG])
https://doi.org/10.48550/arXiv.2605.19269
arXiv‑발행 DOI via DataCite
제출 이력
From: Han Guo [이메일 보기]
[v1]
2026년 5월 19일 화요일 02:30:43 UTC (1,121 KB)
[v2]
2026년 5월 20일 수요일 17:38:24 UTC (493 KB)