배치와 스케줄링

Prefill과 Decode에서는 여러 요청의 새 입력을 함께 계산하면 같은 가중치를 더 많이 재사용할 수 있다는 점을 살펴보았습니다. 이렇게 여러 요청의 계산을 하나로 묶는 것을 배칭(batching), 함께 실행할 묶음을 배치(batch)라고 합니다. 계산을 함께 수행하더라도 각 요청은 자기 문맥과 KV 캐시를 따로 유지합니다.

실제 서비스에서는 요청이 서로 다른 시점에 들어오고, 생성이 끝나는 시점도 다릅니다. 처음 묶은 요청들이 모두 끝날 때까지 기다려야 한다면, 먼저 끝난 요청의 빈자리를 활용하기 어렵습니다. 실행을 반복하는 동안 요청 구성을 바꿀 수 있다면 이 빈자리에 새 요청을 넣을 수 있습니다.

이번 글에서는 고정된 배치의 한계에서 출발해 다음 실행에 넣을 요청을 고르는 과정을 살펴보겠습니다. 이어서 새 요청을 넣기 전에 어떤 자원을 확인해야 하는지, 아직 사용하지 않은 KV 공간까지 미리 잡아 두면 왜 새 요청이 기다리게 되는지 설명하겠습니다.

정적 배칭과 연속 배칭

요청 A와 B를 한 배치로 묶어 생성을 시작했다고 해 보겠습니다. A는 짧게 답하고 끝나지만 B는 긴 답변을 계속 생성합니다. 그사이 새 요청 C가 들어옵니다.

가장 단순한 방법은 A와 B가 모두 끝난 뒤 다음 요청 묶음을 시작하는 것입니다. 이렇게 시작한 요청 묶음에 새 요청을 추가하지 않는 방식을 정적 배칭(static batching)이라고 부르겠습니다. A의 계산이 끝나도 C는 B가 끝날 때까지 기다려야 합니다. 완료된 A의 계산을 생략할 수는 있어도, 그 자리에 C의 계산을 넣지는 못하는 것입니다.

생성은 모델을 한 번 실행하고 끝나는 작업이 아닙니다. 다음 토큰을 선택하고, 그 토큰을 입력으로 다시 실행하는 과정을 반복합니다. 그렇다면 한 번의 실행이 끝났을 때 완료된 A를 빼고 C를 넣어 다음 배치를 구성할 수 있습니다. 이처럼 모델 실행 사이에 요청 구성을 갱신하는 방식이 연속 배칭(continuous batching)입니다.

그림 1은 동시에 최대 두 요청을 실행하는 예시입니다. A와 B는 이미 Decode 중이며, 실행 1에서 A가 끝난 직후 C가 도착합니다. 정적 배칭에서는 실행 2와 3에 B만 남습니다. C는 B까지 끝난 뒤 실행 4에서 시작합니다.

연속 배칭에서는 자원이 충분하다면 실행 2의 배치를 B와 C로 바꿉니다. B는 생성을 이어 가므로 Decode를 수행하고, C는 입력 문맥을 처음 처리해야 하므로 Prefill을 수행합니다. 이후 실행 3에서는 B와 C 모두 Decode를 수행합니다. 새 요청이 합류한다는 것은 그 요청에 필요한 첫 계산부터 시작한다는 뜻입니다.

여기서 카드가 나타내는 것은 실행 순서입니다. 이번 글에서는 각 실행에 어떤 요청을 넣을지에 집중하겠습니다. 그림처럼 배치를 갱신하면 B가 끝나기 전에도 C의 처리를 시작할 수 있다는 것이 핵심입니다. ORCA 논문에서도 요청 전체의 완료를 기다리는 대신 모델 실행을 반복할 때마다 배치를 구성하는 접근을 설명합니다.

스케줄러의 배치 구성 과정

연속 배칭을 하려면 매번 같은 질문에 답해야 합니다. 지금까지 어떤 요청이 끝났고, 어떤 요청이 계속 생성해야 하며, 새로 기다리는 요청 중 무엇을 함께 실행할 수 있을까요?

다음 실행에 넣을 작업을 고르는 일을 스케줄링(scheduling)이라고 합니다. 추론 엔진에서 이 판단을 담당하는 부분이 스케줄러(scheduler)입니다. 여기서는 GPU 내부에서 명령을 배정하는 과정이 아니라, 엔진이 모델에 넘길 요청과 입력을 고르는 과정을 다룹니다.

새 요청이 들어오면 먼저 대기 큐에 넣습니다. 대기 큐는 아직 처리를 시작하지 않은 요청들이 기다리는 목록입니다. 이미 시작한 요청에 대해서는 지금까지 생성한 토큰, 완료 여부, GPU에 저장한 KV 등을 유지합니다. 스케줄러는 이 정보를 바탕으로 다음 배치를 구성합니다.

그림 2의 과정을 순서대로 살펴보겠습니다.

1. 실행 결과를 확인합니다. A는 종료 조건에 도달했고 B는 다음 토큰을 계속 생성해야 합니다. 대기 큐에는 C와 D가 있습니다.
2. 완료된 요청을 정리합니다. 다음 실행에서 A를 제외합니다. A의 생성을 위해 유지하던 KV도 더 이상 필요하지 않으므로, A만 사용하던 공간을 돌려받습니다. B는 생성을 이어 가야 하므로 B의 KV를 유지합니다.
3. 새 요청의 합류를 검토합니다. 먼저 B의 다음 계산을 준비하고, C를 함께 넣어도 필요한 자원이 충분한지 확인합니다. 이 예시에서는 대기 요청을 도착 순서대로 검토하며, 이런 순서를 FCFS(first come, first served)라고 부릅니다.
4. 선택한 배치를 실행합니다. 자원이 충분하다고 판단되면 B의 Decode와 C의 Prefill을 함께 실행합니다. 그림에서는 최대 두 요청을 실행하므로 D는 대기 큐에 남습니다. 실행이 끝나면 다시 결과를 확인하고 같은 과정을 반복합니다.

완료된 요청이 생겨야만 새 요청을 넣을 수 있는 것은 아닙니다. 진행 중인 요청이 모두 계속 생성하더라도 자원에 여유가 있다면 새 요청을 추가할 수 있습니다. 반대로 A가 끝났더라도 C가 요구하는 자원이 크다면 바로 넣지 못할 수 있습니다. 따라서 완료된 요청의 수만큼 새 요청을 넣는 규칙으로는 충분하지 않습니다.

그렇다면 스케줄러는 무엇을 확인해야 자원이 충분하다고 판단할 수 있을까요?

다음 실행에 필요한 자원

요청 수가 같아도 필요한 자원은 다를 수 있습니다. Decode 중인 B는 이번 실행에서 새 입력 하나를 처리하지만, 처음 합류하는 C는 입력 문맥 전체를 처리해야 합니다. B가 이미 보관하고 있는 KV도 계속 GPU 메모리에 남아 있어야 합니다.

따라서 이번에 처리할 토큰 수와 실행 뒤 보관할 KV의 양을 나누어 확인하겠습니다. 앞의 값은 이번 실행에 얼마나 많은 새 계산을 넣을지 정하는 기준이고, 뒤의 값은 생성 상태를 유지할 저장 공간의 요구량입니다.

이번에 처리할 토큰 수

엔진은 한 번의 모델 실행에 넣을 토큰 수의 상한을 둘 수 있습니다. 이것을 실행 토큰 예산이라고 하겠습니다. 예산이 8개라면 선택한 요청들이 이번에 처리할 입력 토큰의 합을 8개 이내로 맞추는 것입니다. GPU가 물리적으로 8개까지만 처리할 수 있다는 뜻은 아니며, 엔진이 실행할 작업의 크기를 제한하는 설정입니다.

스케줄러는 요청의 현재 상태를 알고 있으므로 이 수를 셀 수 있습니다. B가 Decode를 한 번 수행한다면 새 입력은 1개입니다. C의 입력 길이가 6이라면, C의 Prefill에 넣을 토큰은 6개입니다. 따라서 B와 C를 선택했을 때 이번 처리량은 1 + 6 = 7토큰으로 예산 8개 안에 들어갑니다.

그림 3의 위쪽 막대가 이번에 처리할 7토큰입니다. B가 과거에 처리한 6개 위치는 여기에 다시 더하지 않습니다. 그 위치의 K/V는 캐시에 있기 때문입니다. 다만 B의 Attention은 과거 KV를 읽습니다. 토큰 예산은 새로 처리할 위치를 세는 기준이며, 문맥 길이에 따라 달라지는 KV 읽기와 Attention의 계산량까지 모두 표현하지는 않습니다. 같은 토큰 수가 항상 같은 실행 시간을 뜻하지 않는 이유입니다.

실행 뒤 보관할 KV의 양

이번에는 그림 3의 아래쪽 막대를 보겠습니다. B는 이미 6개 위치의 KV를 갖고 있습니다. 새 입력 하나를 처리하면 그 입력의 KV가 추가되어 총 7개가 됩니다. C는 입력 6개를 처음 처리하므로 6개 위치의 KV를 새로 만듭니다. 실행 뒤 보관할 KV는 7 + 6 = 13개 위치 분량입니다.

앞선 글에서 설명한 KV 계산의 경계도 그대로입니다. B가 이번 실행의 결과로 선택할 다음 출력 토큰의 KV는 아직 만들어지지 않습니다. 이번에 입력으로 넣은 한 위치의 KV만 추가되므로, B의 KV는 6개에서 7개로 늘어납니다.

실제 메모리는 바이트 단위로 관리합니다. 모델과 KV 저장 자료형 등이 정해지면 토큰 한 위치의 K/V를 각 층에 저장하는 데 필요한 메모리를 계산할 수 있습니다. 엔진은 가중치와 계산용 작업 공간 등을 고려해 KV에 사용할 수 있는 메모리를 정하고, 요청별로 필요한 KV 공간을 확인합니다.

여기서는 같은 모델과 KV 저장 형식을 사용한다고 가정하고, 그 공간을 토큰 몇 위치의 KV를 저장할 수 있는지로 표현하겠습니다. 그림의 KV 저장 용량 24개는 GPU 전체 메모리가 아니라 KV에 사용할 수 있는 공간입니다. 특정 엔진의 설정값을 가져온 것이 아니라 판단 과정을 설명하기 위한 숫자입니다.

필요한 실제 KV는 13개이므로 전체 용량 24개 안에 들어갑니다. 하지만 확인이 하나 더 남아 있습니다. 그만큼의 공간을 요청에 실제로 배정할 수 있어야 합니다. 메모리 일부가 다른 요청의 앞으로의 생성을 위해 이미 잡혀 있다면, 아직 데이터를 쓰지 않았더라도 C에 줄 수 없을 수 있습니다.

KV 공간을 미리 예약하는 이유

C에 줄 수 있는 공간을 살펴보기 전에, 왜 아직 사용하지 않는 공간을 미리 잡아 두는지부터 생각해 보겠습니다.

KV 캐시는 요청이 진행되는 동안 계속 늘어납니다. 지금 B가 6개 위치의 KV를 갖고 있더라도 다음 입력을 처리하면 7개가 되고, 생성을 계속하면 더 많은 공간이 필요합니다. 요청의 입력 길이는 처음부터 알 수 있지만, 출력이 정확히 언제 끝날지는 시작할 때 알기 어렵습니다.

한 가지 방법은 그 요청이 사용할 수 있는 최대 길이에 맞춰 KV 공간을 미리 확보하는 것입니다. 예를 들어 입력과 생성 과정에서 저장할 KV를 합쳐 최대 20개 위치까지 허용한다면, 처음부터 B에 20개 위치 분량을 배정할 수 있습니다. 여기서 예약이란 데이터를 미리 계산하는 것이 아니라, 나중에 B의 KV를 저장할 메모리를 B의 몫으로 잡아 두는 일입니다.

이렇게 하면 정해 둔 길이 안에서는 B의 KV가 늘어날 때마다 추가 공간을 확보할 필요가 없습니다. 특히 요청의 KV를 하나의 연속된 메모리 영역에 저장하는 방식에서는 처음부터 충분히 큰 영역을 잡아 두면 관리하기 쉽습니다. 현재 필요한 크기만 잡았다가 바로 옆 공간을 다른 요청이 사용하면, 같은 영역을 그대로 늘리기 어려울 수 있기 때문입니다.

다만 최대 길이는 실제로 끝날 길이와 다릅니다. B가 20개 위치를 모두 쓰기 전에 끝나면 그때까지 남겨 둔 공간은 사용되지 않습니다. 생성 도중에도 아직 채우지 않은 부분은 계속 생깁니다. PagedAttention 논문에서 설명하는 기존 KV 관리의 문제 중 하나가 이러한 예약 공간입니다.

이처럼 큰 공간을 미리 예약하는 것은 KV를 관리하는 한 가지 방식입니다. 모든 추론 엔진이 반드시 최대 길이만큼 예약해야 하는 것은 아닙니다. 이 방식이 다음 요청의 처리에 어떤 영향을 주는지 살펴보겠습니다.

미사용 예약 공간과 새 요청의 대기

그림 4는 앞서 필요한 KV를 계산한 B와 C를 그대로 사용합니다. KV 저장 용량도 24개 위치로 같습니다. 여기에 B가 이미 20개 위치 분량을 예약해 두었다는 조건을 더하겠습니다.

그림 위쪽에서 B가 실제 KV를 저장한 공간은 6개입니다. 빗금으로 표시한 14개는 B에 예약되어 있지만 아직 사용하지 않은 공간입니다. 둘을 합친 20개가 B의 몫이고, 다른 요청에 새로 배정할 수 있는 빈 공간은 24 − 20 = 4개뿐입니다.

B가 Decode를 한 번 더 수행하는 데 필요한 한 위치는 자기 예약 공간 안에서 사용할 수 있습니다. 하지만 C의 Prefill에는 새로 6개 위치가 필요합니다. 현재 C에 줄 수 있는 공간은 4개이므로 C를 넣을 수 없습니다. C의 이후 생성에 필요한 여유분을 고려하기 전부터 입력 문맥의 KV를 담을 공간이 부족한 것입니다.

실제 데이터량만 보면 상황이 다르게 보입니다. B와 C를 실행한 뒤 저장할 KV는 앞서 계산한 대로 13개입니다. 전체 용량 24개보다 작습니다. C가 기다리는 이유는 두 요청의 실제 KV가 전체 용량을 넘어서가 아니라, 아직 사용하지 않는 공간까지 B에 묶여 있기 때문입니다.

따라서 스케줄러는 필요한 토큰 수와 KV의 양을 계산하는 데서 멈출 수 없습니다. 현재 메모리를 관리하는 방식에서 그 공간을 실제로 배정할 수 있는지까지 확인해야 합니다. 그림 2에서 C가 합류한 것은 필요한 공간을 배정할 수 있는 경우이고, 그림 4는 예약 공간 때문에 같은 합류가 막히는 경우입니다.

연속 배칭으로 실행마다 요청 구성을 바꿀 수 있어도, 미사용 예약 공간이 크면 새 요청을 충분히 넣기 어렵습니다. 배치 구성을 자주 바꾸는 것과 KV 공간을 효율적으로 사용하는 것은 함께 해결해야 할 문제입니다.

그렇다면 최대 길이만큼 큰 영역을 미리 잡아 두는 대신, KV가 늘어날 때 필요한 만큼 공간을 배정할 수는 없을까요? 다음 글인 KV 캐시 관리와 PagedAttention에서는 KV를 작은 단위로 나누어 배정하고, 떨어진 저장 공간에서도 요청의 KV를 읽을 수 있게 만드는 방법을 살펴보겠습니다.