TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference

"TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference" 논문에 대한 상세한 기술 요약입니다.

1. 문제 정의

텐서 병렬성 (TP) 을 이용한 대규모 언어 모델 (LLM) 의 분산 추론은 NVLink 와 같은 고속 인터커넥트에서도 통신 오버헤드로 인해 상당한 성능 병목 현상에 직면합니다.

• 통신 오버헤드: TP 배포에서 통신 (특히 AllReduce) 은 엔드투엔드 추론 지연 시간의 **9–23%**를 차지합니다.
• "작은 배치"의 도전: 통신과 계산을 겹치게 하는 기존 기술들 (작은 타일이나 토큰으로 작업을 분해하는 등) 은 주요 서비스 시스템 (vLLM, SGLang, TensorRT-LLM) 에서 기본적으로 활성화되어 있지 않습니다. 이는 낮은 지연 시간 추론 (예: <2K 토큰) 에서 일반적인 작은 작업 부하를 분해할 경우, 통신 겹치기로 인한 이득을 상쇄하는 계산 비효율성 (웨이브 양자화 효과) 이 발생하기 때문입니다.
• 자원 경쟁: 통신 연산은 종종 많은 수의 스트리밍 멀티프로세서 (SM) (16–20 개 이상) 를 소비하여 계산 커널의 자원을 고갈시킵니다.
• RMSNorm 중복: 간과되기 쉬운 RMSNorm 연산은 **4–9%**의 오버헤드를 추가합니다. 표준 TP 에서 AllReduce 후 모든 GPU 가 동일한 데이터를 갖는 시점임에도 불구하고, RMSNorm 은 모든 GPU 에서 중복 수행됩니다.

2. 방법론: TokenWeave

TokenWeave 는 작은 배치 크기 (최소 1024 토큰) 에서도 텐서 병렬 추론을 위한 효율적인 계산 - 통신 겹침을 가능하게 하도록 설계된 시스템입니다. 이는 세 가지 핵심 혁신을 활용합니다:

A. 웨이브 인식 스마트 토큰 분할

TokenWeave 는 배치를 두 개의 동일한 절반으로 단순하게 분할하는 대신, 웨이브 인식 분할 전략을 사용합니다.

• 메커니즘: 입력 배치를 두 개의 부분 집합 (접두어 및 접미어) 으로 분할하여, 두 부분 집합에 필요한 총 실행 "웨이브" (CTA 그룹) 수가 원래 분할되지 않은 배치에 필요한 웨이브 수를 초과하지 않도록 합니다.
• 목표: 이는 웨이브 양자화 오버헤드를 최소화하여 작업을 분할하더라도 계산 효율성이 저하되지 않도록 보장합니다.
• 실행: 시스템은 두 개의 CUDA 스트림 (계산 및 통신) 을 사용합니다. 첫 번째 분할에서 AllReduce(RMSNorm 과 융합됨) 가 실행되는 동안, 두 번째 분할에서 어텐션 계산이 실행되고 그 반대의 경우도 마찬가지입니다.

B. RMSNorm 재배열 및 융합

TokenWeave 는 RMSNorm 을 중요한 최적화 대상으로 식별하고 실행 흐름을 재배열합니다.

• 재배열: AllReduce 후 RMSNorm 을 수행하는 대신, TokenWeave 는 AllReduce 를 ReduceScatter 와 AllGather 로 분해합니다.
  1. ReduceScatter: 각 GPU 가 텐서의 샤드를 수신합니다.
  2. 로컬 RMSNorm: 각 GPU 는 최종 집계 전에 로컬 샤드에서 RMSNorm 을 수행합니다. 데이터가 이미 분할되어 있으므로, 이는 GPU 간 중복 계산을 제거합니다.
  3. AllGather: 정규화된 샤드가 수집됩니다.
• 융합 커널: 새로운 Fused AllReduce–RMSNorm 커널이 구현되었습니다. 이는 ReduceScatter, RMSNorm, AllGather 를 단일 연산으로 융합합니다.
  • 메모리 최적화: SM 레지스터에서 ReduceScatter 결과에 직접 분산을 계산함으로써 HBM 읽기를 한 번 제거합니다. 정규화된 값을 AllGather 를 위한 Multimem 주소에 직접 쓰기함으로써 HBM 쓰기를 한 번 제거합니다.
  • SM 효율성: 이 커널은 Hopper/Blackwell GPU 의 하드웨어 기능인 NVSHARP/Multimem을 활용하여 통신 및 정규화를 2–8 개의 SM만으로 수행하여 대부분의 SM 을 계산에 자유롭게 할당합니다.

C. 적응형 겹침 전략

TokenWeave 는 배치 크기에 따라 전체 분할/겹침을 사용할지, 아니면 융합 커널만 사용할지 동적으로 결정합니다:

• 큰 배치: 전체 양방향 분할 및 겹침을 활성화합니다.
• 작은 배치: 오버헤드를 피하기 위해 분할을 비활성화하지만, 메모리 트래픽을 줄이고 중복 계산을 제거하기 위해 Fused AllReduce–RMSNorm 커널을 적용합니다.

3. 주요 기여

1. RMSNorm 최적화: TP 추론에서 RMSNorm 을 재배열 및 커스텀 융합 커널을 통해 식별하고 최적화한 최초의 사례로, HBM 트래픽 및 중복 계산을 줄였습니다.
2. 웨이브 인식 분할: 웨이브 양자화로 인한 계산 효율성 손실을 방지하는 새로운 토큰 분할 알고리즘으로, 작은 토큰 배치 (1K–2K) 에서도 겹침을 가능하게 합니다.
3. 하드웨어 인식 효율성: NVSHARP/Multimem 을 활용하여 통신용 SM 사용량을 약 20 개에서 2–8 개로 줄여 자원 고갈을 방지합니다.
4. 시스템 통합: 낮은 지연 시간 환경에서 텐서 병렬 서비스를 위한 계산 - 통신 겹침을 성공적으로 활성화한 최초의 시스템으로, vLLM 프레임워크에 통합되었습니다.

4. 실험 결과

8×H100 DGX 시스템에서 Llama-3.3-70B, Qwen2.5-72B, Mixtral-8x22B와 같은 모델을 사용하여 평가가 수행되었습니다.

• 지연 시간 개선:
  • 최적화된 기준선 (Multimem 을 사용하는 vLLM) 대비 최대 1.28 배의 지연 시간 속도 향상을 달성했습니다.
  • 짧은 시퀀스 길이에서도 상당한 이득을 보였습니다 (예: 1K 토큰에서 1.2 배).
  • 더 큰 배치 크기 (≥4K) 에서만 이득을 보이고 작은 크기에서는 오버헤드를 발생시키는 TileLink와 같은 최신 겹침 솔루션보다 성능이 우수합니다.
• 처리량 개선:
  • 실제 세계 추적 데이터 (ShareGPT, arXiv) 에서 최대 1.19 배의 높은 처리량을 제공합니다.
  • 일부 설정에서 TokenWeave 는 통신이 전혀 없는 이론적 기준선 (vLLM-nocomm) 보다 성능이 더 뛰어납니다. 이는 융합 커널이 메모리 바운드인 RMSNorm 연산을 기준선이 순수 계산을 처리하는 것보다 더 잘 최적화하기 때문입니다.
• 커널 효율성:
  • 융합된 AllReduce–RMSNorm 커널은 모든 토큰 크기에서 순차 실행 대비 1.34–1.39 배의 속도 향상을 제공합니다.
  • 4–8 개의 SM 만 사용하여 거의 최적의 성능을 달성합니다.

5. 의의

• 저지연 분산 추론 가능: "작은 배치" 문제를 해결함으로써 TokenWeave 는 대규모 배치 크기를 요구하지 않고도 지연 시간이 중요한 대화형 애플리케이션을 위한 분산 TP 추론을 실현 가능하게 합니다.
• 하드웨어 활용도: NVSHARP, Multimem 과 같은 현대 GPU 기능을 최대한 활용하여 병목 현상을 통신 오버헤드에서 순수 계산으로 전환합니다.
• 실용적 영향: 소스 코드가 오픈소스로 공개되어 vLLM 에 통합되었으며, 모델 아키텍처 변경 없이 프로덕션 LLM 서비스 시스템에 즉각적인 성능 향상을 제공합니다.
• 패러다임 전환: 통신 겹침은 대규모 학습 배치나 전문가 병렬성 (MoE) 설정에서만 유익하다는 가정에 도전하여, 표준 텐서 병렬 추론에서도 필수적임을 입증했습니다.