CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs

TL;DR Highlight

GPU에서 Transformer 학습 시 발생하는 메모리 병목을 해결하기 위해, 정규화·활성화 등 소규모 연산들을 GEMM 출력이 칩 위에 있는 동안 함께 실행하는 커널 추상화 CODA를 소개한다. LLM이 이 추상화를 활용해 고성능 커널을 자동 생성할 수 있다는 점이 특히 주목받고 있다.

Who Should Read

대형 Transformer 모델을 직접 학습시키거나 GPU 커널 최적화에 관심 있는 ML 시스템 엔지니어. CUTLASS, Triton 등 저수준 GPU 프로그래밍 또는 LLM 기반 코드 자동 생성(codegen)을 탐구 중인 개발자.

Core Mechanics

Transformer 학습 파이프라인에서 행렬 곱(GEMM)은 이미 고도로 최적화되어 있지만, 정규화(LayerNorm, RMSNorm), 활성화 함수, 잔차 연결(residual add) 등 주변 연산들이 GPU 글로벌 메모리를 반복적으로 읽고 쓰는 탓에 전체 학습 시간의 비중이 큰 병목이 되고 있다.
CODA는 이 문제를 해결하기 위해 GEMM 출력 타일(tile)이 아직 칩 위 온칩 메모리(SRAM)에 있는 동안, 뒤따르는 소규모 연산들을 함께 실행하는 'GEMM-plus-epilogue' 방식의 GPU 커널 추상화다. 글로벌 메모리 왕복 횟수 자체를 줄이는 것이 핵심이다.
CODA의 설계는 GEMM 메인루프(계산 코어 부분)는 건드리지 않고, 그 뒤에 붙는 epilogue 부분만 스케일링·축소(reduction)·쌍별 변환·누적(accumulation) 등 소수의 조합 가능한 프리미티브로 표현하도록 인터페이스를 엄격히 제한했다.
수학적 핵심 아이디어는 RMSNorm/LayerNorm처럼 행 단위 전역 축소(row-wise reduction)가 필요한 연산의 스케일 팩터가 특정 조건에서 교환 가능하다는 점이다. 예를 들어 `(W1 @ gamma * globally_computed_scale) * W2`를 `(W1 @ gamma * W2) * globally_computed_scale`로 재배치해 GEMM 타일 처리 중에 부분 집계를 가능하게 한다.
기존 torch.compile 같은 좌-우 그래프 컴파일러는 두 GEMM 사이의 전역 행 단위 축소가 교환 가능하다고 가정할 수 없어서 이 최적화를 자동으로 적용하지 못했다. CODA는 이 제약을 명시적 설계로 극복한다.
표준 Transformer 블록의 순전파(forward)와 역전파(backward)에서 attention을 제외한 거의 모든 비-attention 연산을 CODA 추상화로 커버할 수 있다고 논문은 밝혔다.
사람이 직접 작성한 CODA 커널뿐만 아니라 LLM이 자동 생성한 CODA 커널 모두 대표적인 Transformer 워크로드에서 높은 성능을 달성했다. 이는 LLM 기반 커널 코드 자동 생성의 실용적 경로로서 CODA가 유효함을 시사한다.

Evidence

epilogue fusion(GEMM 직후 연산 융합)으로 글로벌 메모리 왕복을 줄이는 개념 자체는 Triton으로도 이미 구현 가능했다는 지적이 있었다. 하지만 댓글에서는 CODA의 진짜 가치는 기술적 새로움보다는 'LLM이 커널을 작성하기 좋은 제한적·조합 가능한 API 설계'라는 설계 철학의 전환에 있다는 의견이 주를 이뤘다.
RMSNorm/LayerNorm의 행 단위 스케일이 특정 조건에서 교환 가능하다는 수학적 통찰을 정리한 댓글이 주목받았다. `(W1 @ gamma * scale) * W2 = (W1 @ gamma * W2) * scale` 형태로 재배치할 수 있어, 타일 단위 부분 집계가 가능해진다는 구체적 설명이 커뮤니티의 이해를 도왔다.
LLM은 저수준 하드웨어 최적화(레지스터 배치, 메모리 계층 세부 튜닝)에는 여전히 약하지만, 전문가가 만든 블록들을 고수준에서 조합하는 작업에는 뛰어나다는 의견이 있었다. CODA처럼 인터페이스를 제한하면 LLM이 잘 못하는 부분은 추상화로 가리고 잘하는 부분만 활용할 수 있다는 분석이었다.
LLM이 CODA 커널을 성공적으로 작성했다는 논문 결과에 대해, '실행 피드백(run → profile → patch) 루프가 있으면 합성 단독보다 훨씬 빠르게 격차를 좁힐 수 있고, 이는 본질적으로 RL(강화학습) 문제와 같다'는 댓글이 달렸다. 단순 one-shot 생성이 아니라 반복 피드백 루프가 핵심이라는 시각이다.
CUTLASS GEMM epilogue fusion을 작성해본 개발자들로부터 '익숙한 패턴'이라는 반응이 나왔고, 이 연구가 그 패턴을 체계적으로 LLM-friendly하게 정리했다는 평가가 있었다. 향후 agentic 개발 환경에서 이런 제한적 API 기반 codegen이 표준이 될 것이라는 전망도 제기됐다.

How to Apply

대형 Transformer 모델을 직접 학습시키는데 프로파일링 결과 LayerNorm, RMSNorm, residual add 등 비-attention 연산이 전체 시간의 상당 부분을 차지하고 있다면, CODA 논문의 GEMM-epilogue 융합 패턴을 참고해 해당 연산들을 GEMM 타일 처리와 합쳐 메모리 왕복 횟수를 줄일 수 있다.
CUTLASS나 Triton으로 커스텀 커널을 작성할 때, CODA의 설계처럼 epilogue 부분을 스케일링·reduction·pairwise transform·accumulation 등 소수의 조합 가능한 프리미티브로 분해해 구현하면, 나중에 LLM(GPT-4o, Claude 등)에게 커널 변형을 요청할 때 정확도가 훨씬 높아진다.
LLM 기반 커널 자동 생성 파이프라인을 구축하려는 경우, CODA 방식처럼 LLM이 건드리지 못하도록 GEMM 메인루프는 고정하고 epilogue 조합만 생성하게 인터페이스를 제한하면 생성 품질이 올라간다. 여기에 실행→프로파일링→수정 피드백 루프를 추가하면 성능 격차를 더욱 빠르게 좁힐 수 있다.
torch.compile이 자동으로 최적화하지 못하는 'GEMM 사이 전역 행 단위 축소' 패턴을 직접 쓰고 있다면, CODA 논문의 교환 법칙 재배치 기법(`(W @ gamma * scale) → (W @ gamma) * scale` 형태)을 적용해 fusion 가능한 형태로 수식을 재구성해볼 수 있다.

Code Example

snippet

# CODA의 핵심 수학적 재배치 아이디어 (개념 예시)
# 기존: 두 GEMM 사이에 전역 스케일 계산 필요 → 메모리 왕복 발생
# out1 = X @ W1                        # GEMM 1
# scale = rms_norm_scale(out1)          # 전역 행 축소 → 글로벌 메모리 write/read
# out1_normed = out1 * scale
# out2 = out1_normed @ W2              # GEMM 2

# CODA 방식: 스케일을 뒤로 이동 (행 단위 상호작용만 있으면 교환 가능)
# out_fused = (X @ W1) @ W2           # W2를 먼저 적용 (epilogue에서 타일 처리)
# scale = rms_norm_scale(X @ W1)       # 타일 단위 부분 집계로 온칩에서 처리
# out2 = out_fused * scale             # 스케일을 나중에 적용

# 조건: scale이 행 단위 상호작용만 있을 때 (RMSNorm, LayerNorm의 경우 성립)
# 효과: 중간 텐서를 글로벌 메모리에 쓰고 읽는 과정 제거

Terminology

GEMMGeneral Matrix Multiply의 약자. GPU에서 가장 핵심적인 행렬 곱 연산으로, Transformer의 대부분 계산이 여기에 해당한다. GPU는 이 연산에 극도로 최적화되어 있다.

EpilogueGEMM 계산이 끝난 직후, 결과를 메모리에 쓰기 전에 수행하는 후처리 연산 단계. 여기에 다른 연산을 끼워 넣으면 메모리 왕복을 줄일 수 있다.

Memory-bound연산 속도가 GPU 계산 능력이 아니라 메모리 읽기/쓰기 속도에 의해 제한되는 상태. LayerNorm 같은 소규모 연산들이 전형적인 예다.

Fusion여러 개의 GPU 커널(연산 단위)을 하나로 합치는 최적화 기법. 커널 간 중간 결과를 글로벌 메모리에 저장하지 않아도 되므로 속도가 빨라진다.

TileGPU에서 대형 행렬을 처리할 때 사용하는 작은 조각 단위. 타일 단위로 온칩 SRAM(L1 캐시)에 올려서 처리하면 글로벌 메모리 접근을 줄일 수 있다.

RMSNormRoot Mean Square Normalization의 약자. Transformer에서 각 레이어 출력을 정규화하는 연산으로, LLaMA 등 최신 모델에서 LayerNorm 대신 많이 사용된다. 행 단위 전역 집계가 필요해 GPU 최적화가 까다롭다.