정적 배치의 문제는 요청이 묶음 단위로 시작하고 끝난다는 데서 나온다. 연속 배치(continuous batching)는 스케줄링 단위를 디코드 스텝 하나로 줄인다. 스텝이 끝날 때마다 끝난 요청을 빼고 기다리던 요청을 넣으므로 배치 구성원이 스텝마다 바뀐다. 지금 쓰이는 LLM 추론 엔진은 거의 모두 이 방식을 기본으로 쓴다.
반복 단위 스케줄링
이 방식은 Orca(OSDI 2022)가 반복 단위 스케줄링(iteration-level scheduling)이라는 이름으로 제안했다. 모델을 한 번 실행해 요청마다 토큰 하나를 만드는 것을 반복(iteration) 하나로 보고, 스케줄러가 반복마다 어떤 요청을 실행할지 다시 고른다.
정적 배치에서 본 낙오자 예로 돌아가면 요청 8개 중 7개가 20토큰에서 끝나는 경우 7개는 21번째 스텝부터 배치에서 빠진다. 스케줄러는 그 자리에 대기열의 요청을 넣고 새 요청은 다음 스텝에서 프리필을 시작한다. 500토큰을 생성하는 요청은 혼자 남아 묶음을 붙잡는 대신 다른 요청들과 함께 계속 생성한다.
이 구조에서는 앞서 본 두 대기가 사라진다.
- 묶음이 찰 때까지 기다리지 않는다. 요청이 하나만 와도 다음 스텝에 바로 들어간다
- 앞 묶음이 끝날 때까지 기다리지 않는다. 묶음이라는 단위가 없으므로 자리가 비는 즉시 들어간다
끝난 요청은 마지막 토큰을 만든 스텝에서 바로 응답을 마치므로 같은 배치의 긴 요청이 짧은 요청의 응답 완료를 늦추지도 않는다.
선택적 배치
스텝마다 구성원이 바뀌면 같은 스텝 안에 처지가 다른 요청이 섞인다. 프리필을 시작하는 요청은 토큰 수백 개를 처리하고 디코드 중인 요청은 토큰 하나를 처리한다. 디코드 요청끼리도 지금까지 쌓은 KV 캐시 길이가 제각각이다. 이 요청들을 하나의 직사각형 텐서로 맞추려면 다시 패딩이 필요하다.
Orca는 이 문제를 선택적 배치(selective batching)로 푼다. 연산을 둘로 나눠 다룬다.
| 연산 | 처리 방식 | 이유 |
|---|---|---|
| 선형 층 (QKV 투영, 출력 투영, MLP) | 모든 요청의 토큰을 한 줄로 이어 붙여 행렬 곱 한 번으로 계산 | 토큰마다 독립적으로 같은 가중치를 곱하므로 요청 경계가 필요 없다 |
| 어텐션 | 요청마다 따로 계산 | 각 토큰이 자기 요청의 KV 캐시만 봐야 하고, 그 길이가 요청마다 다르다 |
선형 층에서는 요청 구분 없이 토큰을 합쳐 계산하므로 패딩이 없다. 스텝 안의 토큰이 디코드 요청 30개의 토큰 30개와 새 요청 하나의 프롬프트 500개라면, 선형 층은 토큰 530개짜리 행렬 곱 하나로 처리한다. 가중치를 한 번 읽어 토큰 530개에 곱하므로 연산 강도도 그만큼 오른다.
어텐션은 요청마다 길이가 다른 KV 캐시를 읽어야 해서 따로 계산한다. 지금의 엔진은 길이가 다른 요청들의 어텐션을 커널 하나에서 처리하는 가변 길이 어텐션 커널을 써서 이 부분도 한 번의 실행으로 묶는다.
한 스텝의 구성
연속 배치 엔진은 스텝마다 두 가지 상한 안에서 배치를 짠다. vLLM을 예로 들면 엔진 설정에 다음 두 값이 있다.
| 설정 | 의미 |
|---|---|
max_num_seqs | 한 반복에서 처리하는 시퀀스(요청)의 최대 수 |
max_num_batched_tokens | 한 반복에서 처리하는 토큰의 최대 수 |
디코드 요청은 스텝마다 토큰 하나씩을 쓰고 프리필 요청은 프롬프트 토큰 수만큼 쓴다. 스케줄러는 진행 중인 디코드 요청을 먼저 넣고 남는 토큰 예산 안에서 대기열의 요청을 프리필로 넣는 식으로 스텝을 채운다. 두 상한은 지연과 처리량의 손잡이다. max_num_seqs를 키우면 디코드 배치가 커져 처리량이 오르는 대신 스텝이 길어지고, max_num_batched_tokens를 키우면 긴 프롬프트를 한 스텝에 처리할 수 있는 대신 그 스텝에 묶인 디코드 요청들이 더 오래 기다린다.
입장 제어와 선점
배치에 넣을 수 있는 요청 수를 실제로 정하는 것은 KV 캐시 메모리다. 요청마다 생성이 진행되면서 KV 캐시가 토큰 하나씩 자라므로 지금 메모리가 충분하더라도 진행 중인 요청들이 계속 생성하면 나중에 모자랄 수 있다.
엔진이 이 문제를 다루는 방식은 두 가지다. 하나는 요청을 받을 때 최대 생성 길이까지의 KV 캐시를 미리 예약하는 것이다. 메모리가 모자랄 일은 없지만 대부분의 요청이 최대 길이보다 훨씬 일찍 끝나므로 예약해 둔 메모리 상당 부분이 끝내 쓰이지 않는다. 다른 하나는 필요할 때마다 조금씩 할당하고 메모리가 바닥나면 진행 중인 요청 일부를 배치에서 내보내는 선점(preemption)이다. 내보낸 요청의 KV 캐시는 버렸다가 나중에 프롬프트와 생성한 토큰으로 다시 계산하거나, CPU 메모리로 옮겼다가 다시 가져온다.
선점이 일어나면 내보낸 요청은 그동안 토큰을 받지 못하고 다시 계산하는 비용까지 치른다. 사용자에게는 생성 도중 응답이 멈추는 것으로 보인다. 그래서 선점 횟수는 TTFT·TPOT와 함께 늘 지켜보는 지표다. 선점이 잦다면 동시 처리 상한이 KV 캐시 메모리에 비해 크다는 뜻이다.
남는 문제
연속 배치는 낙오자 문제와 묶음 대기를 없앴지만 두 가지를 남긴다. 둘 다 엔진이 스텝을 어떻게 짜고 메모리를 어떻게 나누느냐에 달린 문제다.
첫째, 새 요청의 프리필이 디코드 스텝에 끼어든다. 2,000토큰 프롬프트가 들어온 스텝은 프리필 계산 때문에 평소보다 몇 배 길어지고, 그 스텝에 묶인 디코드 요청들의 토큰 간격이 그만큼 튄다. 긴 프롬프트가 자주 들어오는 부하에서는 이 효과가 TPOT의 꼬리를 정한다.
둘째, 동시에 처리할 수 있는 요청 수를 KV 캐시 메모리가 정하는데, 요청마다 연속된 메모리를 할당하는 방식으로는 메모리 조각이 많이 생겨 실제 수용량이 줄어든다. 첫째 문제는 프리필을 나눠 스케줄링하는 방법으로, 둘째 문제는 KV 캐시를 페이지 단위로 관리하는 방법으로 다룬다.