트랜스포머를 위한 딥 커널 융합으로 메모리 대역폭 최적화

초록

본 논문은 에이전트형 LLM 추론 시 메모리 대역폭이 병목이 되는 상황에서, SwiGLU 기반 MLP 블록의 가중치 로드와 중간 활성값 저장을 최소화하는 DeepFusionKernel을 제안한다. 단일 깊게 융합된 CUDA 커널을 통해 HBM 트래픽을 13 %까지 감소시키고, H100·A100 GPU에서 각각 13.2 %와 9.7 %의 처리량 향상을 달성한다.

상세 분석

DeepFusionKernel은 기존 Transformer의 Feed‑Forward Network(FFN) 단계, 특히 SwiGLU 구조를 대상으로 설계되었다. SwiGLU는 두 개의 선형 변환(Weight‑Up, Weight‑Gate)과 SiLU 활성화·게이트 연산을 결합한 형태로, 각 변환이 대규모 행렬 곱(GEMM)과 점곱 연산을 요구한다. 전통적인 구현에서는 네 개의 커널(두 GEMM + 두 포인트와이즈 연산) 혹은 두 개의 커널(두 GEMM을 하나로 묶고 포인트와이즈를 별도)으로 나누어 실행한다. 이 경우 중간 텐서가 HBM에 저장·로드되는 비용이 크게 늘어나며, 특히 배치 크기가 작고 모델 파라미터가 큰 상황에서 메모리 대역폭이 주요 제한 요인이 된다.

DeepFusionKernel은 이러한 흐름을 하나의 커널로 통합한다. 핵심 아이디어는 GEMM 연산 중에 바로 SiLU·게이트 연산을 스트리밍 방식으로 삽입해, 중간 결과를 레지스터 혹은 공유 메모리 수준에서 즉시 소비하도록 하는 것이다. 이를 위해 저자들은 루프 순서와 타일링 전략을 두 가지 축으로 탐색하였다.

1. Row‑major 타일링 – 입력 행렬 X의 행을 기준으로 타일을 잡아, 동일 행을 여러 GEMM에서 재사용한다. 배치가 커지거나 입력 활성값이 메모리 트래픽을 주도할 때 효과적이다.
2. Column‑major 타일링 – 가중치 행렬의 열을 기준으로 타일링해, 동일 가중치 타일을 여러 연산에서 재사용한다. 모델 파라미터가 크고 배치가 작을 때, 특히 에이전트형 장기 컨텍스트 디코딩에서 유리하다.

타일 크기는 레지스터 사용량, 공유 메모리 점유율, Tensor Core 활용도를 균형 있게 맞추도록 자동 튜닝된다. 또한, 하드웨어마다 최적 타일이 다를 수 있기에, 논문에서는 경량 프로파일러 기반 스케줄러를 도입했다. 스케줄러는 배포 전 대상 GPU와 모델 설정에 대해 후보 커널들을 짧게 벤치마크하고, 가장 높은 스루풋을 보이는 구성을 선택한다. 이 과정은 수 밀리초 수준으로, 실제 추론 단계에 거의 영향을 주지 않는다.

성능 평가에서는 LLaMA‑3.1‑70B 모델을 FP16으로 실행했으며, TP=4(4 GPU) 환경에서 배치 164, 출력 길이 1 02416 384 토큰을 테스트했다. DeepFusionKernel은 SGLang 기본 구현 대비 A100에서 평균 9.7 %, H100에서 평균 13.2 %의 스루풋 향상을 기록했다. 특히 배치 1~8 구간에서 메모리 대역폭 제한이 뚜렷한 경우 이득이 크게 나타났으며, 배치가 커질수록 GPU 내부 연산이 포화되면서 이득이 점차 감소했다. 장기 컨텍스트(출력 길이 16 384)에서도 MLP 단계가 전체 토큰당 지연 시간의 상당 부분을 차지함을 확인했으며, DeepFusionKernel은 일관된 가속을 제공했다.

한계점으로는 다양한 GPU 클러스터 인터커넥트(예: NVLink vs. PCIe)와 inter‑GPU 통신 비용에 대한 상세 분석이 부족하다는 점을 들었다. 또한, Softmax와 같은 장거리 종속성을 갖는 연산은 현재 융합 대상에서 제외되었으며, 향후 연구 과제로 남겨졌다.

요약하면, DeepFusionKernel은 메모리 대역폭이 병목인 LLM 자동 회귀 디코딩 상황에서, SwiGLU MLP의 데이터 흐름을 근본적으로 재구성함으로써 HBM 트래픽을 크게 감소시키고, 최신 NVIDIA GPU의 연산 잠재력을 보다 효율적으로 활용한다는 점에서 실용적이며, 기존 추론 프레임워크(SGLang, vLLM 등)에 최소한의 통합 비용으로 적용 가능한 솔루션이다.