연재 중집필 중인 책입니다. 아직 본문이 비어 있거나 채워지는 중인 장이 있습니다.

제 6 장

청크 프리필과 스케줄링 정책

연속 배치는 새 요청을 다음 스텝에 바로 넣는다. 그런데 새 요청의 첫 스텝은 프리필이다. 프롬프트가 수천 토큰이면 그 스텝은 디코드만 있는 스텝보다 몇 배 길어지고, 같은 스텝에 묶인 디코드 요청들은 그동안 토큰을 받지 못한다. 연속 배치 엔진의 스케줄러가 결정하는 것은 결국 프리필과 디코드를 언제, 얼마나 섞느냐다.

프리필이 끼어드는 비용

H100 SXM과 Llama 3.1 8B의 조합에서 디코드 스텝 하나의 하한은 약 4.8ms이고, 2,000토큰 프롬프트의 프리필 하한은 약 32ms다. 디코드 요청 30개가 돌고 있는 스텝에 이 프리필이 들어가면 그 스텝은 프리필 계산이 끝나야 끝난다. 디코드 요청 30개 모두가 이번 토큰을 평소의 예닐곱 배 늦게 받는다.

이런 요청이 초당 몇 개씩 들어오면 디코드 요청들은 몇 스텝마다 한 번씩 긴 스텝을 겪는다. 평균 TPOT는 조금 오를 뿐이지만 토큰 사이 간격의 꼬리가 길어지고, 스트리밍 화면에서는 글자가 멈칫하는 것으로 보인다. TPOT p99를 목표로 두는 서비스라면 이 꼬리가 목표를 넘기는 주된 원인이 된다.

프리필 우선과 디코드 우선

스케줄러가 고를 수 있는 가장 단순한 정책은 둘이다.

정책동작얻는 것잃는 것
프리필 우선대기열에 새 요청이 있으면 다음 스텝에 바로 프리필한다새 요청의 TTFT가 짧다. 배치가 빨리 차서 처리량이 높다프리필이 들어올 때마다 진행 중인 요청의 토큰 간격이 튄다
디코드 우선진행 중인 디코드를 먼저 처리하고, 배치에 자리가 날 때만 프리필한다진행 중인 요청의 토큰 간격이 고르다새 요청이 대기열에서 오래 기다려 TTFT가 늘어난다. 배치가 덜 차서 처리량이 떨어진다

어느 정책을 골라도 프리필 한 번의 계산량은 그대로이므로 프리필 우선은 그 시간을 진행 중인 요청의 TPOT로 내고 디코드 우선은 새 요청의 TTFT로 낸다. TTFT와 TPOT에 따로 목표를 두는 서비스에서는 어느 쪽도 두 목표를 함께 지키기 어렵다.

청크 프리필

청크 프리필(chunked prefill)은 프리필을 작은 조각으로 나눠 여러 스텝에 걸쳐 처리한다. 2,000토큰 프롬프트를 500토큰씩 네 조각으로 나누면, 스텝마다 진행 중인 디코드 토큰과 프리필 조각 하나를 함께 넣는다. 프리필은 네 스텝에 걸쳐 끝나고 그동안 디코드 요청들은 스텝마다 토큰을 받는다.

이 방법은 SARATHI가 프리필을 같은 크기의 조각으로 나누고 디코드를 거기에 함께 태우는(piggyback) 방식으로 처음 제안했고, Sarathi-Serve(OSDI 2024)가 이를 서빙 스케줄러로 확장했다. Sarathi-Serve는 진행 중인 디코드를 멈추지 않고 새 요청을 배치에 넣는 스케줄을 정체 없는(stall-free) 스케줄이라고 부른다.

조각을 섞는 것이 손해가 아닌 이유는 프리필은 연산에, 디코드는 대역폭에 걸리기 때문이다. 디코드만 있는 스텝은 대역폭에 묶여 있어 연산기 대부분이 논다. 그 스텝에 프리필 조각 500토큰을 함께 넣으면 가중치는 한 번만 읽고 연산은 프리필 조각이 채운다. 디코드 토큰 입장에서는 원래 놀던 연산기를 프리필이 쓰는 것이라, 스텝 시간이 크게 늘지 않는다. 프리필 입장에서는 이미 읽고 있던 가중치를 함께 쓰므로 디코드가 거의 공짜로 따라온다.

토큰 예산과 청크 크기

청크 프리필 엔진은 스텝마다 처리할 토큰 수의 상한, 즉 토큰 예산을 두고 그 안에 디코드 토큰과 프리필 조각을 채운다. vLLM에서는 이 예산이 max_num_batched_tokens이고, 문서는 청크 프리필을 남은 max_num_batched_tokens에 맞춰 프리필을 자르는 기능으로 설명한다. 스케줄러는 진행 중인 디코드 요청에 토큰 하나씩을 먼저 배정하고, 남은 예산만큼 대기 중인 프리필을 잘라 넣는다.

예산 크기는 TPOT와 TTFT를 맞바꾼다.

  • 예산이 작으면 스텝마다 프리필 조각이 작아 스텝 시간이 고르게 짧다. TPOT의 꼬리가 줄어든다. 대신 긴 프롬프트는 여러 스텝에 걸쳐 처리되므로 TTFT가 늘어나고, 조각이 작아질수록 조각마다 가중치를 다시 읽는 비용과 스텝마다의 고정 비용이 커져 전체 처리량이 떨어진다
  • 예산이 크면 프리필이 빨리 끝나 TTFT가 짧고 처리량이 높다. 대신 큰 조각이 들어간 스텝이 길어져 TPOT의 꼬리가 다시 길어진다

예산을 고르는 기준은 TPOT 목표다. 디코드 배치와 프리필 조각을 합친 스텝 하나의 시간이 TPOT 목표 안에 들어오는 가장 큰 예산을 찾고, 그 예산에서 TTFT 목표가 지켜지는지 확인한다. 연산 강도로 보면 스텝 안 토큰 수가 능선(H100 SXM, BF16 기준 약 295) 근처를 넘어서야 연산기를 제대로 채우므로, 예산을 그보다 훨씬 작게 두면 연산기가 다시 논다.

우선순위 스케줄링

토큰 예산 안에 무엇을 먼저 넣을지는 또 하나의 정책이다. 기본은 먼저 온 요청을 먼저 처리하는 FCFS다. 대화형 요청과 오프라인 작업이 한 서버를 함께 쓰는 경우에는 요청에 우선순위를 붙여 높은 쪽을 먼저 배치에 넣는다. vLLM의 scheduling_policy는 도착 순서대로 처리하는 fcfs와, 요청에 붙은 우선순위 값이 낮은 순서대로 처리하고 같으면 도착 순서를 따르는 priority를 고를 수 있다.

우선순위를 두면 낮은 순위의 요청이 오래 굶을 수 있다. 오프라인 작업처럼 기한이 있는 낮은 순위 요청에는 동시 처리 수에 하한을 두거나, 오래 기다린 요청의 순위를 올리는 장치가 필요하다. 반대로 높은 순위 요청이 KV 캐시를 모두 차지하면 낮은 순위의 진행 중인 요청이 선점될 수 있으므로, 우선순위와 선점이 어떻게 함께 동작하는지 엔진 문서에서 확인한다.

정리된 스케줄러의 모양

지금의 연속 배치 엔진 스케줄러는 이런 순서로 스텝을 짠다. 엔진마다 세부 구현에는 차이가 있지만 큰 틀은 같다.

  1. 진행 중인 디코드 요청에 토큰 하나씩을 배정한다
  2. 이전 스텝에서 처리하다 남은 프리필 조각을 넣는다
  3. 남은 토큰 예산과 KV 캐시 여유 안에서 대기열의 요청을 우선순위 순으로 꺼내 프리필 조각으로 넣는다
  4. KV 캐시가 모자라면 우선순위가 낮거나 늦게 온 요청을 선점한다

이 틀 안에서 운영자가 정하는 값은 토큰 예산, 동시 처리 요청 수 상한, KV 캐시에 할당하는 메모리 비율, 우선순위 정책이다. 넷 모두 TTFT와 TPOT를 서로 다른 방향으로 움직이므로 하나씩 바꿔 가며 굿풋을 비교해 고른다.

배치와 스케줄링6 / 18