연속 배칭 아키텍처
연속 배칭 아키텍처 (Continuous batching architecture)
연속 배칭이 전통적 배칭과 어떻게 다른지, 요청이 어떤 상태를 거치는지, 내부의 스케줄링·메모리 관리가 어떻게 이뤄지는지 살펴볼게요.
출처: 문서
본문
전통적 배칭은 고정된 요청 그룹을 함께 처리하고, 다음 그룹을 시작하기 전에 가장 느린 요청이 끝날 때까지 기다려요. 이 때문에 배치 사이에 GPU가 놀게 돼요.
연속 배칭에서는 매 생성 단계마다 스케줄러가 완료된 요청을 확인하고 즉시 대기 중인 요청으로 교체해요. 짧은 요청은 끝나는 즉시 밀려나고, GPU는 내내 바쁘게 유지돼요. 결과적으로 처리량은 훨씬 높아지고 평균 지연 시간은 낮아져요.
요청 수명주기 (Request lifecycle)
요청은 제출부터 완료까지 네 가지 상태를 거쳐요.
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 — 출력 토큰이 한 번에 하나씩 생성돼요.
- Finished — 생성이 완료되고 결과를 사용할 수 있어요.
요청은 스케줄러가 충분한 토큰 예산과 캐시 공간을 찾아 수용할 때 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 단계와 섞어 처리해요. 이렇게 하면 배치가 계속 생산적으로 동작하고 다른 요청의 first-token 시간이 줄어요.
스케줄러 (Scheduler)
스케줄러는 두 가지 예산과 하나의 요청 상한을 바탕으로 어떤 요청이 각 forward pass에 합류할지 결정해요.
- 토큰 예산 — 단일 forward pass에서 처리되는 최대 쿼리 토큰 수. ContinuousBatchingConfig의
max_batch_tokens로 설정. - 캐시 예산 — 단일 패스에서 읽을 수 있는 KV 페이지의 총 수. 전체 캐시 크기로 제한.
- 요청 상한 — 단일 forward pass의 최대 요청 수.
max_requests_per_batch로 설정. logits 텐서는 어휘 크기에 따라 확장되므로 이 상한이 어휘가 큰 모델의 최대 메모리를 제한해요.
FIFO
FIFOScheduler는 기본 스케줄러예요. 우선순위 순서로 배치를 채워요. 활성 decode 요청을 먼저 처리하고, 그다음 활성 prefill 요청, 그다음 도착 순서대로 새 대기 요청을 처리해요. 새 요청을 끌어오기 전에 기존 요청을 먼저 밀어냄으로써 처리량을 최대화해요.
PrefillFirst
PrefillFirstScheduler는 decode를 재개하기 전에 chunked prefill 작업을 완료해요. 긴 프롬프트가 지배하는 워크로드에서 파편화(fragmentation)를 줄여줘요.
ContinuousBatchingConfig에서 스케줄러를 설정해요.
cb_config = ContinuousBatchingConfig(scheduler_type="prefill_first")
메모리 관리 (Memory management)
KV cache는 페이지드(paged)되어 있어서 서로 다른 길이의 요청이 파편화 없이 고정 메모리 풀을 공유할 수 있어요. 캐시는 메모리를 고정 크기 블록으로 나누고, 각 요청은 attention 커널이 읽고 쓰는 블록 id 목록을 보유해요.
Paged KV cache (num_blocks × block_size tokens per block)
0 1 2 3 4 5 6 7 8 9 …
┌───┬───┬───┬───┬───┬───┬───┬───┬───┬───┐
│ B │ · │ A │ B │ · │ A │ · │ A │ · │ · │
└───┴───┴───┴───┴───┴───┴───┴───┴───┴───┘
A request A: 2, 5, 7 B request B: 0, 3 · free: 1, 4, 6, 8, 9, …
페이지는 한 레이어의 한 토큰에 대한 key/value 상태를 담아요. 블록은 block_size 페이지(기본 256)의 범위이며 할당 단위예요. 블록은 레이어 그룹별로 할당돼요. 그룹 내 레이어는 블록 id를 공유하며, 이는 혼합 attention 모델에서 부기를 균일하게 유지해요.
캐시는 num_blocks 위에 패딩 토큰이 읽고 쓸 추가 블록 두 개와, 슬라이딩 윈도우 그룹용 sentinel 인덱스를 예약해요. 이 예약 때문에 block_size는 최소 4여야 하고, 매니저는 더 작은 값을 거부해요.
캐시 크기 정하기 (Cache sizing)
매니저는 시작 시 여유 GPU 메모리에서 블록 수를 추론해요. 매니저는 KV 텐서, attention mask, activations, 부기 인덱스를 고려한 방정식을 풀어 사용 가능한 메모리의 max_memory_percent(기본 0.9) 안에 풀이 들어맞도록 크기를 정해요.
ContinuousBatchingConfig에서 값을 명시적으로 고정할 수 있어요.
cb_config = ContinuousBatchingConfig(
block_size=256,
num_blocks=4096,
max_batch_tokens=512,
max_memory_percent=0.8,
)
수용 (Admission)
요청이 배치에 합류하기 전에 스케줄러는 모든 레이어 그룹에 충분한 여유 블록이 있는지 확인해요. 어느 그룹이라도 부족하면 요청은 거부되고 아무것도 할당되지 않아요.
캐시가 오프로딩만 남을 때까지 차는 것을 피하기 위해 스케줄러는 safety_margin을 적용해요. 여유 블록이 safety_margin * num_blocks 아래로 떨어지면 새 prefill은 보류되고 활성 decode만 계속돼요. FIFO 스케줄러는 기본 0.15(블록의 15% 예약)이고, prefill_first는 기본 0.0이에요. ContinuousBatchingConfig의 safety_margin으로 조정할 수 있으며, 0이면 마진을 비활성화해요.
Prefix 캐싱
두 요청이 프롬프트 prefix를 공유하면 그 prefix의 KV를 담은 블록을 공유할 수 있어요. 완료된 각 블록은 콘텐츠 해시됩니다. 일치하는 prefix가 있는 이후 요청은 블록을 재사용하고 그 토큰들의 prefill을 건너뛰어요. 공유 블록은 참조 카운트되며, 그 블록을 쓰는 모든 요청이 끝나야만 여유 풀로 돌아가요.
Prefix 캐싱은 기본적으로 활성화되고 모델이 전적으로 full-attention 레이어로만 이뤄진 경우에만 동작해요. 짧은 프롬프트·긴 생성의 워크로드에서는 부기가 절약분을 초과하므로 ContinuousBatchingConfig에서 allow_block_sharing=False로 설정해요.
퇴출 (Eviction)
해제된 블록은 지워지지 않고 콘텐츠와 해시가 그대로 유지된 "initialized" 상태로 남아, 나중에 prefix가 일치하면 재사용할 수 있어요. 미초기화(uninitialized) 풀이 부족해지면 매니저는 가장 최근의 initialized 블록을 지연(demote)시켜 uninitialized로 되돌리고 해시를 버려요. prefix 캐시는 best-effort 계층이라서 새 요청을 계속 서빙하기 위해 할당자가 버릴 수 있어요.
소프트 리셋 (Soft reset)
소프트 리셋은 오프로딩의 폴백 경로예요. CPU 오프로딩이 비활성화되어 있거나 풀이 가득 차서 매니저가 요청의 KV cache를 CPU 스왑 풀에 복사하지 못할 때, 매니저는 대신 요청을 소프트 리셋해요. 요청의 생성 토큰을 원래 프롬프트에 덧붙이고, 캐시 블록을 해제하고, 큐에 다시 넣어요. 어떤 요청이 선택되는지는 Offloading 섹션에서 다뤄요.
캐시에 다시 공간이 생기면 요청은 생성 이력이 프롬프트에 인코딩된 채 중단된 곳에서 재개돼요.
CUDA graphs
매 생성 단계에는 배치 조립부터 GPU 커널 디스패치, 결과 읽기까지 CPU 오버헤드가 포함돼요. CUDA graphs는 전체 GPU 실행 시퀀스를 한 번 기록하고 shape이 일치하는 배치에 재생함으로써 이 오버헤드를 제거해요.
배치 shape은 매 단계마다 바뀌므로 연속 배칭 시스템은 패딩과 캐싱으로 이를 처리해요.
- 쿼리 길이는
q_padding_interval_size의 가장 가까운 배수로 패딩해요. - KV 길이는
kv_padding_interval_size의 가장 가까운 배수로 패딩해요. - 기록된 그래프는 캐시에 저장돼요.
배치의 패딩된 shape이 캐시된 그래프와 일치하면 그래프가 CPU 디스패치 없이 재생돼요. 새 shape은 새 그래프 캡처를 유발해요.
CUDA graphs는 정적 제어 흐름(static control flow)이 필요하고 attention mask와는 호환되지 않아요. 기본적으로 자동 감지되며 조건이 충족되지 않으면 비활성화돼요.
비동기 배칭 (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을 준비해요. 하지만 둘을 겹치려면 VRAM이 대략 두 배 필요해요.
비동기 배칭은 CUDA graphs가 활성화되어 있어야 해요. 그래프 재생이 스트림 겹침이 올바르기 위해 필요한 안정적인 텐서 주소를 제공하기 때문이에요.
오프로딩 (Offloading)
긴 출력을 생성하는 요청은 KV cache에 다음 블록을 할당할 공간이 없을 때까지 커질 수 있어요. 크래시하는 대신 매니저는 남은 요청을 위해 블록을 확보하려고 활성 요청을 대기 큐로 오프로딩해요.
오프로딩은 수요 기반이에요. 스케줄러는 필요한 블록을 할당하지 못한 활성 요청인 stalled 요청을, 각 요청이 요구하는 블록 수와 함께 보고해요. 매니저는 그 수요를 충당할 만큼 블록을 비우는 데 필요한 최소한의 요청만 오프로딩해요. 가장 최근 것을 먼저 오프로딩하므로 현재 배치에 이미 스케줄링된 요청은 캐시를 유지하고, 마지막 남은 활성 요청은 결코 오프로딩되지 않아요.
두 가지 상황이 오프로딩을 유발해요.
- 배치가 스케줄링됐지만 일부 활성 요청이 stalled 됐어요. 매니저는 다음 배치가 나머지를 위한 블록을 할당할 수 있도록 stalled 요청을 충분히 오프로딩해요.
- 캐시가 가득 차서 아예 배치를 스케줄링할 수 없어요. 매니저는 같은 단계 안에서 요청을 오프로딩하고 스케줄링을 재시도하는 루프를 돌면서, 배치가 맞거나 오프로딩할 요청이 없을 때까지 반복해요. 더 이상 다음 forward pass를 기다리지 않아요.
오프로딩된 요청은 두 경로 중 하나를 타요. cpu_offload_space가 0.0보다 크면 매니저는 요청의 KV cache 블록을 미리 할당된 pinned CPU 버퍼로 복사하고, 스왑 풀에 맞는 모든 요청을 단일 전송으로 배칭해요. GPU 캐시 공간이 생기면 매니저는 블록을 GPU로 다시 복사하고, 프롬프트와 생성 토큰을 다시 계산하지 않고 요청을 재개해요. 풀에 맞지 않는 요청이나(CPU 오프로딩이 비활성화되었을 때의 모든 오프로딩 요청 포함) 소프트 리셋으로 폴백해요.
다음 단계 (Next steps)
- Continuous batching 블로그 글은 KV 캐싱, chunked prefill, 동적 스케줄링을 성능 벤치마크 수치와 함께 다뤄요.
- 사용 예시는 연속 배칭 문서를 참고하세요.