Hugging Face: Continuous Batching from First Principles

Continuous batching은 여러 대화를 병렬로 처리하고 완료되는 즉시 교체함으로써 처리량(throughput)을 극대화하는 대규모 언어 모델(LLM) 서빙의 핵심적인 최적화 기술입니다. 패딩(padding)의 필요성을 제거하고 prefill 및 decoding 단계를 동시에 처리할 수 있게 함으로써, continuous batching은 AI 서비스가 수천 명의 동시 접속 사용자를 효율적으로 처리할 수 있도록 합니다.

Attention과 Prefill 단계의 역할

LLM 생성의 핵심은 attention 메커니즘이며, 이는 네트워크에서 서로 다른 토큰들이 상호작용하는 유일한 부분입니다. 표준적인 forward pass에서 모델은 다음 토큰을 예측하기 위해 입력 시퀀스를 처리합니다(prefill 단계).

이 과정은 입력 토큰을 query ($Q$), key ($K$), value ($V$) 상태로 투영하는 과정을 포함합니다. 그런 다음 모델은 $Q K^{T}$를 통해 유사도 점수를 계산하며, 이는 시퀀스 길이 $n$에 대해 $\mathcal{O}(n^{2} d)$의 이차 복잡도를 가집니다. 각 토큰이 이전 토큰들과만 상호작용하도록 하여 미래의 토큰이 과거에 영향을 미치는 것을 방지하기 위해 causal attention mask가 적용됩니다.

KV Caching을 통한 생성 최적화

새로운 토큰이 생성될 때마다 모델이 시퀀스의 모든 이전 토큰에 대해 $K$와 $V$ 상태를 다시 계산한다면, 토큰을 하나씩 생성하는 것(decoding 단계)은 계산적으로 매우 낭비가 될 것입니다.

KV caching은 prefill 및 이전 decoding 단계에서 생성된 key와 value 상태를 저장함으로써 이 문제를 해결합니다. 이는 토큰 $n + 1$을 생성하는 데 드는 계산 비용을 $\mathcal{O}(n^{2})$에서 $\mathcal{O}(n)$으로 줄이며, 계산량을 메모리와 교환합니다. 예를 들어, Llama-2-7B의 경우, 32개의 레이어와 32개의 헤드에 걸쳐 float16 정밀도로 토큰 하나를 저장하는 데 약 16 KB의 메모리가 필요합니다.

Chunked Prefill을 통한 긴 프롬프트 처리

초기 프롬프트가 매우 긴 경우, 활성화(activations)에 필요한 메모리가 가용 GPU 메모리를 초과하여 단일 forward pass로 prefill을 수행하는 것이 불가능할 수 있습니다.

Chunked prefill은 초기 프롬프트를 관리 가능한 작은 청크(chunk)로 분할하여 이 문제를 해결합니다. KV cache를 활용하여 모델은 첫 번째 청크의 상태를 저장하고 이를 후속 청크의 상태 앞에 붙임으로써, 정보 손실 없이 프롬프트를 점진적으로 처리할 수 있습니다.

Batched Generation에서 Continuous Batching으로

처리량(초당 생성되는 토큰 수)을 높이기 위해 모델은 여러 프롬프트를 병렬로 처리합니다. 그러나 전통적인 방식은 상당한 비효율성을 초래합니다:

  • Batched Generation: 모든 프롬프트의 길이가 동일해야 하므로 패딩 토큰(pad)이 필요합니다. 이는 프롬프트 길이가 다양할 때 낭비가 됩니다.
  • Dynamic Scheduling: 완료된 프롬프트를 새로운 프롬프트로 교체할 수 있게 해주지만, prefill 단계의 새로운 프롬프트가 decoding 단계의 다른 프롬프트들과 함께 배치되어야 하므로 여전히 많은 패딩이 필요합니다.

Ragged Batching

패딩 낭비를 제거하기 위해 ragged batching은 배치 축을 완전히 제거하고 대신 모든 프롬프트를 하나의 시퀀스로 연결합니다. 토큰 상호작용은 attention mask를 통해 엄격하게 제어되어, 한 프롬프트의 토큰이 다른 프롬프트의 토큰과 상호작용하지 않도록 보장합니다.

Continuous Batching 알고리즘

Continuous batching은 ragged batching과 dynamic scheduling을 결합하여 GPU 활용도를 최대로 유지합니다. 프로세스는 다음 논리 단계를 따릅니다:

  1. 메모리 예산 극대화: 시스템은 배치당 $m$개의 토큰이라는 GPU 메모리 예산을 채우려고 시도합니다.
  2. Decoding 우선순위 지정: 현재 decoding 단계에 있는 모든 프롬프트가 먼저 배치에 추가됩니다(각 프롬프트는 토큰 하나를 차지함).
  3. Prefill로 채우기: 남은 공간은 prefill 단계의 프롬프트로 채워지며, 필요에 따라 chunked prefill을 사용하여 입력을 분할합니다. n4. 동적 교체: 완료된 프롬프트( </s> 토큰을 생성한 프롬프트)는 즉시 제거되고 새로운 요청으로 교체됩니다.

주요 기술 요약

Continuous batching은 다음과 같은 세 가지 주요 기술 구성 요소의 통합을 통해 효율성을 달성합니다:

기술 주요 이점
KV Caching decoding 중에 과거 토큰의 표현을 다시 계산하는 것을 방지함.
Chunked Prefill GPU 메모리 한계를 초과하는 프롬프트의 처리를 가능하게 함.
Ragged Batching 시퀀스를 연결하고 마스크를 사용하여 패딩 낭비를 제거함.
Dynamic Scheduling 요청을 실시간으로 교체하여 높은 처리량을 유지함.

Sources