🤖 machine learning

Tile-Level Activation Overlap for Efficient LLM Inference

이 논문은 중간 텐서 구체화 오버헤드를 제거하기 위해 타일 레벨에서 SwiGLU 활성화 함수를 행렬 곱셈과 융합하는 두 가지 특화된 CUTLASS 기반 SM90 커널을 소개하며, 이를 통해 NVIDIA H100 GPU에서 최대 2.47배의 속도 향상과 79.5%의 피크 활용도를 달려 달성함과 동시에 효율성과 수치적 정확도 측면에서 PyTorch와 cuBLAS를 모두 능가한다.

원저자: Abhinav Jangda, Tyler Sorensen, Sebastian Burckhardt, Jianlan YE, Chaoyin Li, Atul Gupta

게시일 2026-07-07

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Abhinav Jangda, Tyler Sorensen, Sebastian Burckhardt, Jianlan YE, Chaoyin Li, Atul Gupta

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 복잡한 로봇(챗봇을 구동하는 것과 같은 대규모 언어 모델)을 제작하는 고속 공장을 운영하고 있다고 상상해 보십시오. 각 로봇을 제작하기 위해 당신의 공장에는 특정 조립 라인인 MLP(로봇 두뇌의 핵심 부분)가 있습니다.

오랫동안 이 조립 라인에는 주요한 비효리성이 존재해 왔습니다. 여기 그 문제와 논문에서 제시된 해결책을 쉽게 설명한 내용이 있습니다.

문제점: "중간 관리자" 병목 현상

현대적인 로봇 공장에서는 로봇을 더 똑똑하게 만들기 위해 SwiGLU라고 불리는 특정 단계가 사용됩니다. SwiGLU를 두 가지 별도의 계산이 필요한 품질 검사라고 생각하십시오:

계산 A: 로봇의 현재 상태 측정.
계산 B: 로봇의 잠재력 측정.
결합(The Glue): 이 두 가지 측정값을 결합하여 최종 결과 도출.

기존 방식 (병목 현상):
표준적인 공장 설정(PyTorch에서 사용하는 방식)에서는, 계산 A를 수행한 후 작업자들이 그 결과를 복도에 있는 거대한 화이트보드(고대역폭 메모리, HBM)에 적어야 합니다. 그러고 나서 그들은 다시 걸어가서 그 화이트보드를 읽고, 계산 B를 수행한 뒤, 그 결과를 두 번째 화이트보드에 적고, 다시 돌아와서 두 화이트보드를 모두 읽은 다음 "결합" 단계를 수행해야 합니다.

이 논문은 더 작은 로봇(작은 AI 모델)의 경우, 이 "화이트보드로 걸어가는 시간"이 전체 시간의 **30%에서 37%**를 차지한다는 것을 발견했습니다. 이는 마치 요리사가 요리를 하는 시간보다 식재료를 가지러 팬트리까지 왔다 갔다 하는 데 절반의 시간을 쓰는 것과 같습니다.

해결책: 두 개의 새로운 "슈퍼 키친"

저자들은 화이트보드를 완전히 없애버린 두 가지 새로운 맞춤형 키친(커널, Kernels)을 구축했습니다. 결과를 적어두고 다시 걸어가는 대신, 작업자들은 재료를 손에 쥐고(레지스터, Registers) 모든 과정을 하나의 연속적인 동작으로 처리합니다.

그들은 서로 다른 공장 규모에 맞춘 두 가지 전략을 만들었습니다:

1. "핑퐁" 키친 (Kernel-1)

작동 방식: 이어달리기 경주를 상상해 보십시오. 한 명의 작업자가 다음 배치의 재료를 가져오는 동안, 다른 작업자는 이미 현재 배치를 섞고 있습니다.
비결: 이들은 "핑퐁" 스케줄을 사용합니다. 기계가 두 번째 재료 세트를 가져오는 데 바쁜 동안, 작업자들은 그 시간을 활용해 첫 번째 세트에 대한 "결합" 수학 연산을 수행합니다.
적합한 경우: 거대한 로봇(대규모 모델)을 만들거나 동시에 많은 로봇을 실행(대규모 배치)하는 공장에 적합합니다. 이는 대형 기계들을 쉬지 않고 가동할 수 있을 만큼 충분한 인력이 확보된 거대한 조립 라인과 같습니다.

2. "인터리브드(Interleaved)" 키친 (Kernel-2)

작동 방식: 상자를 포장하는 상황을 상상해 보십시오. 빨간색 물건을 모두 포장한 다음 파란색 물건을 포장하는 것이 아니라, 빨간색 하나, 파란색 하나를 번갈아 가며 완벽하게 섞어서 포장합니다.
비결: 계산을 위한 두 가지 가중치 행렬(즉, "레시피")을 시작하기 전에 미리 섞어 놓습니다. 이를 통해 작업자들은 "빨간색" 재료와 "파란색" 재료를 동시에 잡고 즉시 함께 처리할 수 있습니다.
적합한 경우: 더 작은 로봇을 만들거나 한 번에 적은 수의 로봇을 실행(작은 배치)하는 공장에 적합합니다. 이 방법은 매우 효율적이어서 공장 바닥이 작업자들로 꽉 차게 유지하며, 아무도 기다리느라 서 있지 않게 만듭니다.

결과: 속도와 정확도

저자들은 다양한 로봇 크기(0.5B 모델부터 거대한 72B 모델까지)에 대해 NVIDIA H100 칩(가장 강력한 AI 프로세서)에서 이 새로운 키친들을 테스트했습니다.

엄청난 속도 향상: 작은 로봇의 경우, 새로운 키친은 기존 표준 방식보다 2.47배 더 빨랐습니다. 이는 10분 걸리던 작업을 4분으로 단축하는 것과 같습니다.
병목 현상의 이동: 기존 시스템은 "메모리 제한적(Memory-Bound)"이었습니다(화이트보드로 걸어다니는 데 너무 많은 시간을 소비함). 새로운 시스템은 "연산 제한적(Compute-Bound)"입니다(수학 연산을 실제로 수행하는 데 모든 시간을 쏟음). 이들은 칩의 최대 이론적 속도의 **79.5%**에 도달했습니다.
컴파일러가 할 수 없는 일: 저자들은 표준 "자동 항법" 소프트웨어(PyTorch의 torch.compile)를 사용하여 이를 자동으로 해결하려고 시도했습니다. 하지만 실패했습니다. 자동 항법 소프트웨어는 그들의 커스텀 키친보다 3배에서 7배 더 느렸습니다. 이는 이 특정 문제에 있어서는 인간 전문가가 직접 설계한 해결책이 필요하다는 것을 증명합니다. 컴퓨터는 아직 스스로 이 방법을 찾아낼 수 없습니다.
더 나은 정확도: 놀랍게도, 새로운 커스텀 키친은 표준 방식보다도 더 정확했습니다. 표준 방식은 결과의 4.5%에서 11% 사이에서 미세한 수학적 오류를 보였으나, 새로운 키친은 오류가 전혀 없었습니다.

요약

이 논문은 공장 바닥의 운영 방식을 재편함으로써(특히 작업자들이 중간 노트를 적거나 읽기 위해 화이트보드로 계속 왔다 갔다 하는 것을 막음으로써), AI 모델을 특히 엣지 디바이스(폰이나 노트북 등)에서 사용하는 작은 모델의 경우 훨씬 더 빠르게 실행할 수 있음을 보여줍니다. 그들은 표준 소프트웨어 도구가 이러한 효율성을 재현할 수 없으며, 이러한 결과를 달기 위해서는 전문적으로 작성된 코드가 필요하다는 것을 입증했습니다.

기술 요약: 효율적인 LLM 추론을 위한 타일 수준 활성화 중첩 (Tile-Level Activation Overlap)

문제 정의

Qwen-2.5, LLaMA, Mistral, Gemma를 포함한 현대적인 대규모 언어 모델(LLM)은 주로 SwiGLU 활성화 함수를 사용한다. 이 게이티드 MLP 구조는 $Y = \text{SiLU}(A \times W_1) \odot (A \times W_2)$ 를 계산하며, 두 개의 독립적인 행렬 곱셈(GeMM)과 그에 따른 요소별 게이티드 활성화를 필요로 한다.

표준 구현에서는 이 과정에서 GeMM 단계와 활성화 단계 사이의 고대역폭 메모리(HBM)에 두 개의 전체 중간 텐서($Gate $및$ Up$)를 실체화(materialize)한다. NVIDIA H100 GPU에서의 프로파일링 결과, 이러한 중간 실체화와 관련 활성화 연산은 전체 MLP 실행 시간의 **9~37%**를 점유하는 것으로 나타났다. 이 오버헤드는 더 작은 모델(예: Qwen-2.5 0.5B)에서 가장 심각하며, 이 경우 SwiGLU가 MLP 시간의 30% 이상을 차지한다. 양자화(FP8, INT4)를 통해 텐서 코어 산술 밀도가 높아짐에 따라, 이러한 메모리 집약적 연산의 상대적 비용은 더욱 증가하여 중요한 병목 현상을 생성한다.

기존 컴파일러 인프라, 특히 PyTorch의 torch.compile은 이를 해결하지 못한다. 이는 서로 다른 가중치 행렬을 가진 두 개의 별개 GeMM 간에 연산을 퓨전(fusion)할 수 없는 그래프 수준 퓨전 패스의 근본적인 한계 때문이다. 실험에 따르면 torch.compile은 이 패턴에 대해 eager PyTorch 성능의 34~94%만을 달여하며, 명시적인 퓨전 힌트를 제공해도 개선 효과는 미미하다.

방법론

저자들은 SwiGLU를 타일 수준에서 GeMM에 퓨전하여 중간 HBM 실체화를 제거하는 두 가지 상호 보완적인 핸드크래프트 CUTLASS 기반 SM90 커널을 제시한다. 두 커널 모두 H100(SM90)에서 사용 가능한 텐서 메모리 가속기(TMA)와 워프 특화 스케줄링(warp-specialized scheduling)을 활용한다.

1. 커널-1: 핑퐁 스케줄 중첩을 통한 동기식 SwiGLU (Sync SwiGLU via Pingpong Schedule Overlap)

메커니즘: 이 커널은 데이터를 로드하는 프로듀서(producer)와 MMA를 실행하는 컨슈머(consumer)로 워프를 분할하는 CUTLASS 핑퐁 워프 특화 스케줄을 사용한다.
전략: $[M, N]$ 스레드블록 그리드를 생성한다. 컨슈머 단계 내에서, Gate 타일이 계산된 후 Up 타일 데이터가 도착하기 전까지 짧은 유휴 시간이 발생한다. 커널은 이 윈도우를 활용하여 레지스터에 상주하는 Gate 누산기(accumulator) 상에서 Swish 활성화( $\text{SiLU}(Gate)$ )를 계산한다.
중첩: Swish 계산은 프로듀서의 Up 가중치 타일에 대한 TMA 로드와 시간적으로 중첩된다. 최종 요소별 곱셈( $\text{SiLU}(Gate) \odot Up$ )과 저장은 에필로그(epilogue)에서 수행된다.
최적화: 이 설계는 $[M, N]$ 그리드가 GPU의 스트리밍 멀티프로세서(SM)를 충분히 포화시키는 큰 배치 크기( $M \ge 2048$ )에 최적화되어 있다.

2. 커널-2: 커스텀 에필로그 방문자 트리를 통한 인터리브 SwiGLU (Interleave SwiGLU via Custom Epilogue Visitor Tree)

메커니즘: 이 커널은 PairMulStore라고 불리는 커스텀 에필로그 방문자 트리(Epilogue Visitor Tree, EVT) 노드를 도입한다.
전략: 가중치 행로 $W_1$ 과 $W_2$ 를 열이 인터리브된 단일 퓨전 행렬 $W_{fused} \in \mathbb{R}^{K \times 2N}$ (짝수 열은 Gate, 홀수 열은 Up)로 결합한다. 커널은 $[M, 2N, K]$ 에 대해 표준 GeMM을 실행하여 $[M, 2N]$ 스레드블록 그리드를 생성한다.
퓨전: PairMulStore 노드는 각 스레드블록 내의 $[M, 2N]$ 누산기에서 동작한다. 인접한 열 쌍 $(D[m, 2n], D[m, 2n+1])$ 을 읽어 $\text{SiLU}(D[m, 2n]) \odot D[m, 2n+1]$ 을 적용하고 결과를 저장한다.
중첩: 이 설계는 SwiGLU 계산을 타일 저장(tile store)과 인터리브한다. 현재 타일의 산술 연산은 이전 타일의 비동기적 TMA 저장과 중첩되어 연산-저장(compute-store) 중첩을 제공한다.
최적화: $2\times$ 증가한 스레드블록 수는 작은 배치 크기( $M \le 2048$ )에서 SM 점유율을 크게 향상시키므로, 대부분의 설정에서 기본 권장 사항이 된다.

주요 기여

타일 수준 퓨전: SwiGLU를 타일 수준에서 GeMM에 퓨전하여 추론 단계당 4개의 중간 HBM 작업(Gate 및 Up의 읽기/쓰기)을 제거한 최초의 구현이다.
이중 커널 전략: 큰 배치용(순차적 Gate→Swish→Up)인 커널-1과 작은 배치용(인터리브 가중치 및 커스 커스텀 EVT)인 커널-2를 사용하는 상호 보완적 접근 방식을 통해 점유율 및 스케줄링 제약 조건을 다르게 해결한다.
수치적 우수성: 퓨전된 커널은 FP32 참조값과 비교했을 때 제로 미스매치(1% 상대 오차를 초과하는 요소 0%)를 달성한다. 반면, 베이스라인인 cuBLAS 구현은 타일 간 공격적인 BF16 부분 누산으로 인해 4.5~11%의 미스매치를 보인다.
컴파일러 한계 검증: torch.compile이 이 퓨전을 재현할 수 없으며, 핸드크래프트 커널보다 3~7배 느리다는 것을 입증하는 경험적 증거를 제시한다.

결과

NVIDIA H100에서 Qwen-2.5 모델(0.5B ~ 72B)을 대상으로 평가:

속도 향상: 커널은 PyTorch eager(cuBLAS + 별도 SwiGLU) 대비 최대 2.47배의 속도 향상을 달성한다. 20개 벤치마크 구성 중 커널-2가 13개에서 승리했으며, 커널-1은 7개(주로 높은 배치 크기의 대형 모델)에서 승리했다.
체제 변화(Regime Shift): 루프라인 분석(Roofline analysis)은 퓨전이 워크로드를 메모리 제한(memory-bound)에서 연산 제한(compute-bound)으로 전환함을 확인시켜 준다. 작은 모델의 경우 산술 강도가 최대 247% 증가하여 루프라인 리지(ridge)를 가로질러 작동 지점을 이동시킨다.
활용률: 피크 BF16 활용률은 H100에서 **79.5%**에 달하며, 이는 작은 모델에 대해 베이스라인의 활용률을 거의 두 배로 높인 것이다.
효율성: 퓨전 효율(실제 속도 향상 vs 이론적 메모리 절감)은 작은 모델에 대해 100%에 가깝다. 커널-2의 경우, 연산-저장 중첩 혜택 덕분에 효율이 100%를 초과하는 구성이 존재한다.
확장성: GeMM 연산이 전체 시간을 지배하고 중간 메모리 트래픽의 상대적 비율이 감소함에 따라, 모델 크기가 커질수록 속도 향상은 감소한다(0.5B에서 2.47배 vs 높은 배치 크기의 72B에서 약 1.0배).

의의

본 논문은 현대적인 LLM 추론, 특히 엣지 배포 및 모바일 규모 모델(0.5B–14B)에서 병목 현상이 원시 연산 능력에서 중간 텐서 실체화로 인한 메모리 트래픽으로 이동했음을 주장한다. 저자들은 핸드크래프트된 하드웨어 인식 타일 수준 설계가 현재의 컴파일러 인프라 한계를 극복하기 위해 필수적임을 입증한다. 중간 HBM 액세스를 제거하고 데이터 이동과 활성화 계산을 중첩함으로써, 이 커널들은 추론 지연 시간을 크게 줄이고 연산 활용도를 높이며, 레이턴시와 처리량이 제약되는 프로덕션 서빙 시스템에 핵심적인 최적화를 제공한다.