[vLLM] Chunked Prefill: 긴 프롬프트를 청크 단위로 분할 처리하는 기법

들어가며

Continuous batching으로 GPU 활용률은 높아졌지만, 한 가지 문제가 남았다. 수만 토큰짜리 긴 프롬프트가 들어오면 해당 prefill이 완료될 때까지 다른 디코딩 요청이 지연된다. Chunked Prefill은 긴 프롬프트를 여러 청크로 나누어 디코딩 요청과 함께 처리함으로써 TTFT(Time To First Token)와 디코딩 레이턴시 간의 균형을 맞춘다.

• 논문: Sarathi: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills (arxiv 2308.16369)
• 공식 문서: https://docs.vllm.ai

핵심 구조/코드 분석

토큰 예산으로 자연스러운 청크 분할

Chunked prefill은 별도의 복잡한 로직이 아니라, vLLM 스케줄러의 토큰 예산(token_budget) 메커니즘으로 자연스럽게 구현된다. vllm/v1/core/sched/scheduler.py의 WAITING 요청 스케줄링 부분을 보면:

# Number of tokens to be scheduled.
num_new_tokens = request.num_tokens - num_computed_tokens
threshold = self.scheduler_config.long_prefill_token_threshold
if 0 < threshold < num_new_tokens:
    num_new_tokens = threshold

if (
    not self.scheduler_config.enable_chunked_prefill
    and num_new_tokens > token_budget
):
    # chunked_prefill 비활성화 시 예산 초과하면 스케줄링 중단
    break

num_new_tokens = min(num_new_tokens, token_budget)

long_prefill_token_threshold는 한 번에 처리할 수 있는 prefill 토큰의 상한선이다. 예를 들어 threshold가 512이고 프롬프트가 4096 토큰이면, 8번에 걸쳐 512 토큰씩 처리된다.

RUNNING 요청에서도 동일한 메커니즘

이미 실행 중인 요청의 prefill도 동일하게 청크 분할된다:

# RUNNING 요청 처리
num_new_tokens = (
    request.num_tokens_with_spec
    + request.num_output_placeholders
    - request.num_computed_tokens
)
if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens:
    num_new_tokens = self.scheduler_config.long_prefill_token_threshold
num_new_tokens = min(num_new_tokens, token_budget)

resumed 요청이든 새 요청이든, num_computed_tokens에서 시작해 예산만큼만 처리하는 동일한 원리로 동작한다.

Chunked Prefill과 Prefill-Decode 혼합

핵심은 RUNNING 요청(대부분 디코딩 중)을 먼저 스케줄링하고, 남은 token_budget을 WAITING 요청의 prefill에 사용한다는 점이다:

token_budget = self.max_num_scheduled_tokens

# 1단계: RUNNING 요청 (decode) 먼저 처리
while req_index < len(self.running) and token_budget > 0:
    ...
    token_budget -= num_new_tokens

# 2단계: 남은 예산으로 WAITING 요청 (prefill) 처리
while (self.waiting or self.skipped_waiting) and token_budget > 0:
    ...
    num_new_tokens = min(num_new_tokens, token_budget)

디코딩 요청은 토큰 1개만 필요하므로 예산을 거의 소모하지 않는다. 남은 예산 전체가 새 요청의 prefill 청크에 할당되어, 한 배치 안에서 decode와 prefill이 자연스럽게 공존한다.

Full Sequence 용량 체크

chunked prefill 활성화 시, 시퀀스 전체를 위한 KV 블록이 확보 가능한지 미리 확인하는 로직도 있다:

if self.scheduler_reserve_full_isl:
    can_fit = self.kv_cache_manager.can_fit_full_sequence(
        request,
        num_new_computed_tokens=num_new_local_computed_tokens,
        new_computed_blocks=new_computed_blocks,
        num_external_computed_tokens=num_external_computed_tokens,
        num_encoder_tokens=num_encoder_tokens,
    )
    if not can_fit:
        break

첫 청크를 처리한 후 메모리 부족으로 나머지를 처리하지 못하는 상황을 방지한다.

왜 이 설계인가

1. 디코딩 레이턴시 보장: 긴 prefill이 디코딩을 블로킹하지 않는다. 디코딩 요청이 항상 먼저 스케줄링되므로 TPOT(Time Per Output Token)이 안정적으로 유지된다.

2. GPU 연산 효율: 디코딩만 돌리면 GPU가 놀고, prefill만 돌리면 다른 요청이 기다린다. 둘을 한 배치에 섞어서 GPU의 compute/memory 자원을 균형있게 활용한다.

3. 별도 메커니즘 불필요: vLLM의 num_computed_tokens 기반 통합 스케줄링 덕분에 chunked prefill이 자연스럽게 동작한다. 특별한 상태 머신이나 페이즈 전환이 없다.

4. 튜닝 가능한 트레이드오프: long_prefill_token_threshold를 키우면 TTFT가 개선되고, 줄이면 디코딩 레이턴시가 개선된다. 워크로드에 맞게 조절할 수 있다.

Chunked prefill은 "모든 요청을 토큰 단위로 균등하게 다룬다"는 vLLM 스케줄러의 철학이 빛나는 대표적인 사례이다.

논문 핵심 내용

Sarathi: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills (2308.16369) 논문은 Chunked Prefill과 Decode-Maximal Batching이라는 두 기법을 제안했다.

핵심 아이디어: Prefill 요청을 동일한 크기의 청크로 분할하고, 각 배치에 하나의 prefill 청크와 최대한 많은 decode 요청을 함께 넣는 "피기백(piggyback)" 전략이다. Decode 요청이 prefill 청크에 편승하면 추가 비용이 한 자릿수(order of magnitude) 이하로 줄어든다.

LLaMA-13B (A6000 GPU)

메트릭	향상
Decode 처리량	최대 10x
엔드투엔드 처리량	최대 1.33x

LLaMA-33B (A100 GPU)

메트릭	향상
Decode 처리량	최대 4.25x
엔드투엔드 처리량	1.25x

GPT-3 (Pipeline Parallelism)

메트릭	향상
파이프라인 버블 감소	6.29x
엔드투엔드 처리량	1.91x

Chunked Prefill의 핵심 통찰은 배치 내 연산 균일성(uniform compute)을 만드는 거다. 기존 continuous batching에서는 긴 prefill과 짧은 decode가 섞이면서 GPU 활용률이 들쭉날쭉했는데, 모든 요청을 고정 크기 청크로 다루면 배치 간 연산량이 균일해진다. 파이프라인 병렬 환경에서 버블을 6.29배 줄인 것은 이 균일성의 직접적인 효과다.