AI 생성 이미지
GEMM은 보통 두 단계로 나뉩니다. 행렬 곱을 수행하는 mainloop, 그리고 결과를 메모리에 쓰기 전에 스케일링, 클램프, elementwise 연산 등을 적용하는 epilogue입니다.1
PingPong GEMM의 핵심은 이 epilogue를 “나중에 붙는 후처리”로 보지 않는 데 있습니다. 한 워프 그룹이 epilogue와 writeback을 처리하는 동안, 다른 워프 그룹은 텐서 코어로 다음 GEMM 타일을 계산합니다. 즉, epilogue 시간을 텐서 코어 작업과 겹쳐 텐서 코어가 쉬는 구간을 줄이려는 설계입니다.
왜 Hopper에서 PingPong이 중요해졌나
Hopper 아키텍처는 TMA와 WGMMA를 도입하면서 비동기 실행 특성이 Ampere보다 강화됐습니다.2 TMA는 글로벌 메모리와 공유 메모리 사이의 데이터 이동을 맡고, WGMMA는 워프 그룹 단위로 텐서 코어 연산을 수행합니다.
이 환경에서 GEMM 최적화의 초점은 단순히 “연산을 더 빠르게”가 아닙니다. 텐서 코어의 계산 능력이 데이터 이동 능력보다 크기 때문에, 실제 성능은 텐서 코어를 얼마나 쉬지 않게 유지하느냐에 크게 좌우됩니다.3
그래서 Hopper GEMM 최적화에서는 warp specialization, producer-consumer 모델, persistent kernel, cooperative schedule, pingpong schedule 같은 기법이 함께 등장합니다.3 공통된 목표는 분명합니다. 데이터 이동, 계산, 후처리의 빈틈을 줄이는 것입니다.
PingPong GEMM은 무엇을 겹치는가
PingPong 커널은 producer와 consumer의 역할 분담으로 이해하면 쉽습니다. producer 워프는 TMA로 공유 메모리에 데이터를 공급하고, consumer 워프는 그 데이터를 사용해 WGMMA 계산을 수행합니다.4
여기서 PingPong의 특징은 consumer 워프 그룹의 교대 방식입니다. 한 consumer가 텐서 코어로 MMA를 수행하는 동안, 다른 consumer는 이전 결과의 epilogue와 writeback을 처리합니다. 다음 단계에서는 두 consumer의 역할이 바뀝니다. 텐서 코어 사용권이 두 consumer 사이에서 “핑퐁”처럼 오가는 구조입니다.
이 설계가 중요한 이유는 epilogue가 더 이상 작은 부록이 아니기 때문입니다. 후처리를 별도 커널로 분리하면 결과를 메모리에 쓴 뒤 다시 읽어야 합니다. 반대로 epilogue에 통합하면 데이터가 아직 칩 안에 있을 때 필요한 처리를 끝낼 수 있습니다.
또한 Hopper의 warp specialization은 warp-group 단위로 역할을 나누기 때문에, 같은 warp-group의 모든 스레드가 같은 경로를 따르면 전통적인 warp divergence가 발생하지 않습니다.3 역할을 나누되 실행 흐름은 질서 있게 유지하는 것이 핵심입니다.
CODA가 보여주는 적용 방식
CODA는 PingPong GEMM의 관점을 LLM 연산으로 확장합니다. 예를 들어 RMSNorm을 별도의 메모리 바운드 연산으로 두지 않고, 첫 번째 GEMM epilogue에서 부분 norm 통계를 계산한 뒤 두 번째 GEMM epilogue에서 scaling을 적용합니다.5
RoPE와 SwiGLU, 그리고 이들의 backward pass도 epilogue에서 처리할 수 있다고 설명합니다.5 이는 Transformer의 여러 연산을 “GEMM + 작은 후처리” 형태로 재구성할 수 있음을 보여줍니다.
관련 논문은 arXiv abs/2605.19269에 공개됐고, 코드는 GitHub의 HanGuo97/coda-kernels 저장소에 공개됐습니다.5
실무적으로 무엇을 봐야 하나
PingPong GEMM이 주는 교훈은 명확합니다. 최신 GPU 커널 최적화는 FLOPS만 높이는 문제가 아닙니다. 다음 질문을 함께 봐야 합니다.
데이터가 지금 글로벌 메모리에 있는가, 공유 메모리에 있는가, 레지스터에 있는가
epilogue를 별도 커널로 뺄 필요가 있는가, GEMM epilogue에 붙일 수 있는가
텐서 코어가 쉬는 구간을 데이터 이동이나 후처리와 겹칠 수 있는가
producer-consumer 역할 분담이 실제 병목을 줄이는가
PingPong GEMM은 “행렬 곱을 빠르게 하는 기법”이라기보다, GEMM 주변의 데이터 이동과 후처리를 함께 스케줄링하는 방식에 가깝습니다. 핵심은 텐서 코어를 놀리지 않는 것입니다. 그리고 이를 위해 epilogue까지 성능 설계의 일부로 끌어들입니다.
1Efficient GEMM in CUDA -- NVIDIA CUTLASS Documentation
2FlashAttention-3: 비동기 및 저정밀도에서의 빠르고 정확한 어텐션 제공
3CUTLASS CUTE 6 Hopper GEMM 高效实现¶
4Deep Dive on CUTLASS Ping-Pong GEMM Kernel
5PingPong GEMM Overlaps Epilogues With Tensor Cores To Boost Efficiency · Digg
이 노트는 보도·비평·교육·연구 목적으로 작성되었습니다. 저작권 문의가 있으시면 에서 알려주세요.