연속 배치 처리 아키텍처

전통적인 배치 처리는 고정된 요청 그룹을 함께 처리하고, 다음 그룹을 시작하기 전에 가장 느린 요청이 완료될 때까지 기다립니다. 이로 인해 배치 사이에 GPU가 유휴 상태가 됩니다.

연속 배치(continuous batching)를 사용하면, 생성 단계의 각 단계마다 스케줄러가 완료된 요청을 확인하고 즉시 대기 중인 요청으로 교체합니다. 짧은 요청은 완료되는 즉시 제거되고, GPU는 전체 시간 동안 계속 사용 상태를 유지합니다. 그 결과 처리량이 크게 향상되고 평균 지연 시간이 줄어듭니다.

요청 수명 주기

요청은 제출부터 완료까지 네 가지 상태를 거칩니다.

       WAIT IN QUEUE          LOAD PROMPT INTO KV       STREAM OUTPUT           DONE
    ┌────────────────┐    ┌─────────────────────┐    ┌─────────────────┐    ┌────────┐
    │    PENDING     │───▶│     PREFILLING      │───▶│    DECODING     │───▶│FINISHED│
    │  ○ ○ ○ · · ·   │    │ ████████░░░░░░░░░░  │    │  → → → → · ·    │    │   ✓    │
    └────────────────┘    └─────────────────────┘    └─────────────────┘    └────────┘
         others ahead            prompt chunks              +1 token/step

Pending — 요청이 대기열에 추가되어 스케줄링을 기다리고 있습니다.
Prefilling — 프롬프트 토큰이 forward pass에서 처리됩니다.
Decoding — 출력 토큰이 하나씩 생성됩니다.
완료 — 생성이 완료되었으며 결과가 사용 가능합니다.

스케줄러가 요청을 수용할 수 있는 충분한 토큰 예산과 캐시 공간을 찾으면, 해당 요청은 pending 상태에서 prefilling 상태로 이동합니다. 프롬프트가 한 단계에 맞지 않을 정도로 길면, 스케줄러는 이를 여러 forward pass에 걸쳐 분할합니다(chunk prefill).

Chunked prefill

Chunked prefill은 큰 prefill을 여러 단계로 나누어 처리함으로써 이 문제를 해결합니다. 단일 단계에서 긴 프롬프트를 처리하는 것은 다른 요청의 생성을 차단하기 때문에 비용이 많이 듭니다.

Blocking (long prefill blocks the batch)
  Step:   1        2        3        4
  Req A  [▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓]   ← one huge prefill
  Req B  ················ wait ················
  Req C  ················ wait ················

Chunked (same prompt split and interleave with decode)
  Step:   1          2          3          4
  Req A  [▓▓ pre]    [▓▓ pre]   [▓▓ pre]   [→ dec]
  Req B  [→ dec]     [→ dec]    [▓▓ pre]   [→ dec]
  Req C  [→ dec]     · idle ·   [→ dec]    · idle ·

새 요청의 프롬프트가 사용 가능한 토큰 예산(ContinuousBatchingConfig의 max_batch_tokens로 설정됨)을 초과하면, 스케줄러는 가능한 한 많은 토큰을 처리하고 나머지는 보류합니다. 이후 단계에서는 중단된 지점부터 계속 진행하며, prefill 작업을 진행 중인 decode 단계와 교차 처리합니다. 이를 통해 배치의 생산성을 유지하고 다른 요청의 time-to-first-token을 줄입니다.

스케줄러

스케줄러는 두 가지 예산과 요청 상한에 따라 각 forward pass에 어떤 요청이 참여할지 결정합니다.

Token budget — 단일 forward pass에서 처리되는 query token의 최대 개수로, ContinuousBatchingConfig의 max_batch_tokens로 설정됩니다.
Cache budget — 단일 패스에서 읽을 수 있는 KV 페이지의 총 개수로, 전체 캐시 크기로 제한됩니다.
요청 상한 — max_requests_per_batch으로 설정되는 단일 forward pass당 최대 요청 수입니다. logits 텐서는 어휘 크기에 비례하여 확장되므로, 이 상한은 대규모 어휘 모델의 피크 메모리를 제한합니다.

FIFO

FIFOScheduler는 기본 스케줄러입니다. 우선순위 순서대로 배치(batch)를 채웁니다. 먼저 활성 상태의 디코드(decode) 요청을 처리한 다음, 활성 상태의 프리필(prefill) 요청을 처리하고, 마지막으로 도착 순서대로 대기 중인 새 요청을 처리합니다. 이는 새 요청을 받아들이기 전에 기존 요청을 먼저 처리함으로써 처리량(throughput)을 최대화합니다.

PrefillFirst

PrefillFirstScheduler는 디코딩을 재개하기 전에 청크된 프리필 작업을 완료합니다. 이는 긴 프롬프트가 주를 이루는 워크로드에서 단편화를 줄입니다.

ContinuousBatchingConfig에서 스케줄러를 설정합니다.

cb_config = ContinuousBatchingConfig(scheduler_type="prefill_first")

메모리 관리

KV 캐시는 페이지화되어 있어 서로 다른 길이의 요청이 조각나지 않는 고정된 메모리 풀을 공유할 수 있습니다. CachePool로 구현된 메모리 청크의 공유 풀은 이러한 청크를 각 어텐션 유형 할당기에 공급하며, 할당기는 이를 블록과 페이지로 나눕니다.

Shared pool (first two chunks reserved for padding / never allocated)

  chunk N  [ ======================================== ]
               |                            |
    full attn: [ FULL BLOCK 0 | FULL BLOCK 1 | … ]
               [ page L0 | page L0 | … ]   (one page per full layer)

    sliding:   [ ----------- SLIDING BLOCK 0 ----------- ]
               [ page L1 | page L2 | page L3 | … ]  (one page per sliding layer)

  page  = page_size tokens of KV for one layer
  block = one page per layer of that attention type
  chunk = pool unit sized so full and sliding blocks tile the same bytes

Continuous batching는 두 가지 attention 유형만 지원합니다: full_attention (FullAttentionCacheAllocator)와 sliding_attention (SlidingAttentionCacheAllocator). 하이브리드 모델은 각 유형별로 하나의 allocator를 사용합니다. attention 유형은 ContinuousBatchingConfig가 아니라 모델(layer_types / sliding_window)에서 가져옵니다.

Sliding layers는 작은 페이지 집합을 재사용하여 최근 토큰의 고정된 윈도우만 유지합니다. 커널 블록 테이블이나 블록 공유를 지원하지 않습니다.

캐시는 요청이 절대 사용하지 않는 패딩 읽기 및 쓰기를 위해 데이터 청크보다 두 청크를 미리 예약합니다. 해당 청크는 요청에 할당되지 않습니다.

캐시 크기 설정

매니저는 시작 시점에 사용 가능한 GPU 메모리를 기반으로 몇 개의 청크를 수용할 수 있는지 추론합니다. KV 텐서, 어텐션 마스크, 활성화 값, 관리용 인덱스를 모두 고려한 후, 풀의 크기를 사용 가능한 메모리의 max_memory_percent 이내로 설정합니다. 값이 설정되지 않은 경우, 하드웨어 및 연속 배치 처리(continuous batching) 파라미터에 따라 기본값이 결정됩니다. CPU와 GPU가 메모리를 공유하는 MPS 환경에서는 스왑을 방지하기 위해 캐시에 할당하는 공간을 줄입니다. num_blocks은 명시적으로 설정할 때, page_size 토큰으로 구성된 전체 모델 블록 단위로 용량을 나타냅니다.

ContinuousBatchingConfig의 값을 고정할 수 있습니다.

cb_config = ContinuousBatchingConfig(
    page_size=256,
    num_blocks=4096,
    max_batch_tokens=512,
    max_memory_percent=0.8,
)

수용

요청이 배치에 합류하기 전에 스케줄러는 모든 어텐션 유형 할당기에 대해 충분한 여유 용량이 있는지 확인합니다. 어떤 할당기든 용량이 부족할 경우, 해당 요청은 거부되고 아무것도 할당되지 않습니다.

offloading가 유일한 옵션이 될 때까지 캐시를 채우는 것을 방지하기 위해, 스케줄러는 safety_margin을 강제합니다. 사용 가능한 캐시 용량이 safety_margin 미만으로 떨어지면 새로운 prefill은 보류되고, 활성 decode만 계속됩니다. FIFO 스케줄러의 기본값은 0.15(용량의 15%를 예약)이며, prefill_first의 기본값은 0.0입니다. ContinuousBatchingConfig에서 safety_margin을 설정하여 조정할 수 있으며, 0으로 설정하면 마진이 비활성화됩니다.

Prefix 캐싱

두 요청이 프롬프트 접두사를 공유하는 경우, 해당 접두사의 KV를 저장하는 블록을 공유할 수 있습니다. 각 완료된 블록은 콘텐츠 해시로 처리됩니다. 일치하는 접두사를 가진 후속 요청은 해당 블록을 재사용하고, 해당 토큰에 대한 prefill을 건너뜁니다. 공유된 블록은 참조 카운팅되며, 해당 블록을 사용하는 모든 요청이 완료된 후에야 비어 있는 풀로 반환됩니다.

allow_block_sharing의 기본값은 True입니다. 모든 할당기가 공유를 지원하는 경우에만 접두사 공유가 활성화되므로, 풀 어텐션 전용 모델에서만 이를 사용합니다. 순수 슬라이딩 모델과 하이브리드 모델에서는 비활성화됩니다. 짧은 프롬프트와 긴 생성을 처리하는 워크로드에서는 관리 비용이 절감 효과를 상회하므로 allow_block_sharing=False로 설정하세요.

제거

Freed 블록은 지워지지 않으며, 내용과 해시가 그대로 유지된 채 “initialized” 상태로 남아 있어 이후의 prefix match에서 재사용할 수 있습니다. 초기화되지 않은 풀이 부족해지면, 매니저는 가장 최근의 initialized 블록을 지연 처리(lazy) 방식으로 uninitialized로 강등하고 해당 해시를 제거합니다. Prefix cache는 새로운 요청을 계속 처리하기 위해 할당기가 폐기할 수 있는 best-effort 레이어입니다.

소프트 리셋

소프트 리셋은 오프로딩의 폴백 경로입니다. CPU 오프로딩이 비활성화되어 있거나 풀이 가득 차서 매니저가 요청의 KV 캐시를 CPU 스왑 풀로 복사할 수 없는 경우, 대신 해당 요청을 소프트 리셋합니다. 이 과정에서 요청이 생성한 토큰을 원래 프롬프트에 추가하고, 캐시 블록을 해제한 뒤, 요청을 다시 큐에 추가합니다. 어떤 요청이 선택되는지는 Offloading 섹션에서 다룹니다.

캐시에 다시 공간이 확보되면, 프롬프트에 인코딩된 생성 히스토리를 바탕으로 요청이 중단된 지점부터 재개됩니다.

CUDA 그래프

각 생성 단계에는 배치 조립부터 GPU 커널 디스패치 및 결과 읽기까지 CPU 오버헤드가 수반됩니다. CUDA graphs는 전체 GPU 실행 시퀀스를 한 번만 기록하고 동일한 형상을 가진 배치에 대해 이를 재생함으로써 오버헤드를 제거합니다.

batch의 형태가 매 스텝마다 변경되므로, 연속 배치 처리 시스템은 패딩과 캐싱을 통해 이를 처리합니다.

쿼리 길이는 q_padding_interval_size의 가장 가까운 배수로 패딩됩니다.
KV 길이는 kv_padding_interval_size의 가장 가까운 배수로 패딩됩니다.
기록된 그래프는 캐시에 저장됩니다.

batch의 패딩된 형태가 캐시된 그래프와 일치하면, CPU 디스패치 없이 그래프가 재생됩니다. 새로운 형태가 감지되면 새로운 그래프 캡처가 트리거됩니다.

CUDA graphs는 정적 제어 흐름을 요구하며 어텐션 마스크와 호환되지 않습니다. 기본적으로 자동 감지되며 조건이 충족되지 않으면 비활성화됩니다.

Async batching

Async batching은 두 개의 I/O 버퍼 쌍과 두 개의 CUDA 스트림을 사용하여 CPU와 GPU 작업을 겹치게 합니다.

Sequential
  CPU  ── [prep N] ·········idle········· [prep N+1] ·········idle·········
  GPU  ·········idle········· [compute N] ·········idle········· [compute N+1]

Async
  CPU  ── [prep N] [prep N+1] [prep N+2] ──
  GPU  ── ········ [compute N] [compute N+1] ──

GPU가 배치 N을 계산하는 동안 CPU는 배치 N+1을 준비합니다. 그러나 두 작업을 오버랩하려면 대략 2배의 VRAM이 필요합니다.

Async batching은 스트림 오버랩이 올바르게 동작하는 데 필요한 안정적인 텐서 주소를 제공하는 그래프 재실행을 위해 CUDA graphs가 활성화되어 있어야 합니다.

오프로딩

긴 출력을 생성하는 요청은 KV 캐시에 다음 블록을 할당할 공간이 없도록 커질 수 있습니다. 이 경우 크래시 대신 매니저는 활성 요청을 대기 큐로 되돌려 남아 있는 요청을 위해 블록을 해제합니다.

Offloading은 수요 기반입니다. 스케줄러는 stalled 요청, 즉 필요한 블록을 할당하지 못한 활성 요청과 각 요청이 요구하는 블록 수를 보고합니다. 매니저는 해당 수요를 충족할 만큼 충분한 블록을 확보하기 위해 필요한 최소한의 요청만 offload합니다. 최신 요청부터 offload하므로 현재 배치에 이미 스케줄링된 요청은 캐시를 유지하며, 마지막까지 남은 활성 요청은 절대 offload되지 않습니다.

두 가지 상황에서 오프로딩이 트리거됩니다.

배치가 예약되었지만 일부 활성 요청이 멈춰 있습니다. 매니저는 다음 배치에서 나머지 요청에 대한 블록을 할당할 수 있도록 멈춘 요청 중 충분한 양을 오프로드합니다.
캐시가 가득 차서 어떤 배치도 스케줄링할 수 없는 경우, 매니저는 요청을 오프로드하고 같은 스텝 내에서 루프를 돌리며 스케줄링을 재시도합니다. 배치가 들어갈 때까지, 또는 더 이상 오프로드할 요청이 없을 때까지 이 과정이 반복됩니다. 더 이상 다음 포워드 패스를 기다리며 복구하지 않습니다.

오프로드된 요청은 두 가지 경로 중 하나를 따릅니다. cpu_offload_space가 0.0보다 크면, 매니저는 요청의 KV 캐시 블록을 사전에 할당된 고정된 CPU 버퍼로 복사하며, 스왑 풀에 맞는 모든 요청을 단일 전송으로 배치합니다. GPU 캐시 공간이 사용 가능해지면, 매니저는 블록을 GPU로 다시 복사하고 프롬프트와 생성된 토큰을 다시 계산하지 않고 요청을 재개합니다. 풀에 맞지 않는 요청, 또는 CPU 오프로딩이 비활성화된 경우의 모든 오프로드된 요청은 soft reset으로 폴백됩니다.

다음 단계

Continuous batching 블로그 게시물에서는 KV 캐싱, 청크드 프리필, 동적 스케줄링에 대해 성능 벤치마크 수치와 함께 다룹니다.
사용 예제는 Continuous batching 문서를 참고하세요.

Update on GitHub

Transformers