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

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

# 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

관련 논문

• 1-bit/Ternary Bonsai Image 4B: 로컬 디바이스용 이미지 생성 모델

  4B 파라미터 이미지 생성 모델의 가중치를 1비트/3값으로 극단적으로 압축해서 iPhone에서도 돌아가게 만든 모델. 7.75GB짜리 diffusion transformer를 0.93GB까지 줄였다.

• Tiny-vLLM: C++와 CUDA로 만드는 고성능 LLM 추론 엔진

  vLLM의 핵심 기능을 C++와 CUDA로 직접 구현하며 배울 수 있는 교육용 LLM 추론 엔진 프로젝트로, 소스코드와 단계별 강의가 함께 제공된다.

• 일반 데이터센터 GPU에서 요청당 3,000 tokens/s 실시간 LLM 추론

  Kog AI가 8× AMD MI300X에서 요청당 3,000 tokens/s를 달성하는 LLM 추론 엔진을 공개했고, 기존 소프트웨어 스택의 병목을 GPU 메모리 대역폭 최대화로 풀어냈다는 내용이다.

• LLM을 위한 수면과 유사한 컨텍스트 통합 메커니즘

  LLM이 긴 컨텍스트를 처리할 때 발생하는 Attention 비용 문제를 해결하기 위해, 사람의 수면처럼 주기적으로 컨텍스트를 fast weight에 압축·저장하는 새로운 메커니즘을 제안한 논문이다.

• KV-Fold: 긴 컨텍스트 추론을 위한 One-Step KV-Cache Recurrence

  모델 수정 없이 KV 캐시를 청크 간 누산기로 쓰면 128K 토큰까지 100% 정확도로 정보를 검색할 수 있다.

• Swift로 LLM 학습시키기 Part 1: 행렬 곱셈을 Gflop/s에서 Tflop/s로 끌어올리기

  Apple Silicon에서 Swift로 직접 행렬 곱셈 커널을 구현하며 CPU, SIMD, AMX, GPU(Metal)를 단계별로 최적화해 Gflop/s에서 Tflop/s 수준까지 성능을 높이는 과정을 상세히 설명한 글이다. 프레임워크 없이 LLM 학습의 핵심 연산을 밑바닥부터 구현하고 싶은 개발자에게 Apple Silicon의 성능 한계를 체감할 수 있는 드문 자료다.

Related Resources

• CODA 논문 원문 (arXiv:2605.19269)
• CUTLASS - NVIDIA의 GEMM 및 epilogue fusion 라이브러리
• Triton - OpenAI의 GPU 커널 프로그래밍 언어