Chunked Prefill

청크 프리필은 긴 입력 프롬프트를 더 작은 토큰 그룹으로 나누고, 다른 요청의 토큰 생성과 섞어 여러 단계에 걸쳐 처리하는 모델 서빙 기법입니다. 청크 프리필이 없으면 긴 프롬프트를 수용한 추론 엔진이 해당 프롬프트에 전체 단계를 할애하고, 다른 사용자에게 이미 스트리밍 중인 모든 응답은 처리가 끝날 때까지 멈춥니다. 2023년 SARATHI 논문은 청크 프리필을 소개했고, 후속 연구인 Sarathi-Serve는 이를 바탕으로 ‘중단 없는’ 스케줄링을 구현했습니다. DeepSpeed-FastGen의 Dynamic SplitFuse도 같은 아이디어를 사용합니다. 이제 청크 프리필은 vLLM에서 기본적으로 활성화되어 있으며 SGLang과 TensorRT-LLM에서도 사용할 수 있습니다. 이 엔진들은 널리 사용되는 오픈 웨이트 대규모 언어 모델(LLM)을 서빙합니다.

프리필, 디코드, 생성 중단#

자기회귀 Transformer 추론에는 두 단계가 있습니다.

• 프리필: 모델은 모든 프롬프트 토큰을 병렬로 처리하고, 키와 값을 계산해 KV 캐시에 저장합니다. 프리필은 읽은 가중치 바이트당 많은 연산을 수행하므로 일반적으로 연산에 의해 제한됩니다.
• 디코드: 모델은 각 요청에 대해 단계마다 토큰 하나를 생성하며, 매번 전체 캐시를 읽습니다. 디코드는 읽은 바이트당 연산량이 적으므로 일반적으로 메모리 대역폭에 의해 제한됩니다.

20,000개 토큰을 프리필하는 GPU 단계는 디코드 단계보다 훨씬 오래 걸리므로, 배치의 모든 요청이 대기합니다. Sarathi-Serve는 이러한 일시 중지를 생성 중단이라고 부릅니다. 자신의 프롬프트는 짧지만 다른 요청 때문에 기다리는 사용자의 토큰 간 지연 시간에서 급증으로 나타납니다.

청크 프리필의 작동 방식#

스케줄러는 각 단계에 토큰 예산을 할당하며, vLLM에서는 max_num_batched_tokens으로 설정합니다. 먼저 실행 중인 각 디코드에 토큰 하나를 추가한 다음, 남은 예산을 대기 중인 프롬프트 청크로 채웁니다. 프롬프트가 예산에 들어맞지 않으면 다음 단계로 이어지며, 전체 프롬프트가 처리되고 디코드가 시작될 때까지 청크마다 KV 캐시가 확장됩니다. 두 단계를 한 단계에서 혼합하면 GPU 활용도도 높아집니다. 디코드가 연산 장치를 유휴 상태로 두는 동안 프리필 청크가 이를 활용할 수 있기 때문입니다.

청킹은 프롬프트 토큰을 처리하는 시점만 바꿉니다. 각 청크는 인과 마스크를 적용해 이전 모든 청크의 캐시된 키와 값 및 자체 토큰에 어텐션을 수행하므로, 한 번에 프리필한 결과와 동일합니다. PyTorch의 어텐션 계층 하나가 이를 보여줍니다.

import torch
import torch.nn.functional as F

torch.manual_seed(0)
prompt_len, chunk = 1000, 256
q, k, v = torch.randn(3, 1, 8, prompt_len, 64).unbind(0)  # one layer: 8 heads, 64 dims each

# One-shot prefill: all prompt tokens attend causally in a single pass
full = F.scaled_dot_product_attention(q, k, v, is_causal=True)

# Chunked prefill: each chunk attends to the cached keys and values of earlier chunks plus itself
outputs = []
for start in range(0, prompt_len, chunk):
    end = min(start + chunk, prompt_len)
    mask = torch.ones(end - start, end, dtype=torch.bool).tril(diagonal=start)  # causal over absolute positions
    q_chunk, k_cached, v_cached = q[..., start:end, :], k[..., :end, :], v[..., :end, :]
    outputs.append(F.scaled_dot_product_attention(q_chunk, k_cached, v_cached, attn_mask=mask))
chunked = torch.cat(outputs, dim=2)

print(f"{len(outputs)} chunks, same result as one pass: {torch.allclose(full, chunked, atol=1e-5)}")

4 chunks, same result as one pass: True이 출력됩니다. 1,000개 토큰 프롬프트는 256, 256, 256, 232개 토큰의 청크로 처리되며, 출력은 단일 패스로 처리한 결과와 일치합니다.

관련 서빙 기법#

• 연속 배칭: 전체 배치가 끝날 때까지 기다리지 않고 각 단계에서 요청을 추가하고 제거합니다. 청크 프리필을 사용하면 이 루프 안에서 긴 프롬프트를 스케줄링할 수 있습니다.
• PagedAttention: KV 캐시를 고정 크기 페이지에 저장하므로, 여러 단계에 걸쳐 프리필하는 프롬프트가 전체 길이에 필요한 메모리를 미리 예약하지 않고도 캐시를 확장할 수 있습니다.
• 프롬프트 캐싱: 반복되는 프롬프트 접두부의 캐시된 키와 값을 재사용해 해당 부분의 프리필을 완전히 제거합니다.
• 프리필 및 디코드 분리: 두 단계를 별도의 GPU 풀에서 실행하고 단계 간에 KV 캐시를 전송하므로, 각 풀을 자체 병목에 맞게 확장하고 조정할 수 있습니다.
• 추측 디코딩: 여러 토큰을 초안으로 생성한 뒤 한 번에 검증해 디코드 단계를 가속합니다.

응용 분야#

청크 프리필은 긴 프롬프트와 짧은 프롬프트를 같은 서버에서 처리할 때 가장 중요합니다. 50페이지 보고서로 컨텍스트 윈도우를 채우는 검색 증강 생성(RAG) 어시스턴트는 프롬프트를 청크 단위로 프리필하는 동안 다른 사용자의 대화를 계속 스트리밍할 수 있습니다. 비전 언어 모델(VLM)은 각 프롬프트에 수백 또는 수천 개의 이미지 토큰을 추가하므로, 설명을 위해 프레임을 전송하는 카메라 파이프라인은 이러한 요청 조합의 전형적인 사례입니다. Ultralytics LLM 인터페이스는 자체 호스팅 vLLM 서버를 포함해 OpenAI 호환 엔드포인트로 텍스트와 이미지를 전송합니다. 이때 사람을 감지한 프레임의 설명을 요청하는 YOLO26 파이프라인의 요청은 다른 사용자 트래픽과 함께 청크 단위로 프리필됩니다.

토큰 예산 조정#

예산은 두 가지 지연 시간 지표 사이의 균형을 조정합니다. 예산이 작으면 각 단계에서 프리필에 할당할 여유가 줄어 스트리밍 사용자의 토큰 간 지연 시간은 낮아지지만 새 프롬프트의 첫 토큰 생성 시간은 길어집니다. 예산이 크면 그 반대가 됩니다. 긴 프롬프트의 각 후속 청크는 이전 청크의 KV 캐시가 커지는 만큼 이를 다시 읽으므로, 청크가 지나치게 작으면 메모리 트래픽이 증가합니다. 운영자는 현실적인 프롬프트 길이를 사용해 중앙값 및 꼬리 추론 지연 시간, 첫 토큰 생성 시간, 처리량을 기준으로 예산을 조정합니다. vLLM 최적화 가이드는 토큰 간 지연 시간에 유리한 예산으로 2,048개 토큰을 제시하고, 대형 GPU에서 소형 모델의 처리량을 극대화하려면 8,192개보다 큰 예산을 권장합니다.