KV 캐시가 부족할 때: 요청 중단과 재개

KV 캐시 관리와 PagedAttention에서는 KV가 늘어날 때 필요한 블록만 배정해 저장 공간의 낭비를 줄였습니다. 하지만 생성이 계속되면 실제로 필요한 KV도 늘어납니다. 블록을 유연하게 관리하더라도, 더 이상 배정할 블록이 없는 상황은 생길 수 있습니다.

이번 글에서는 이미 실행 중인 요청의 KV 공간이 부족해지는 상황부터 살펴보겠습니다. 이어서 일부 요청을 잠시 중단해 다른 요청이 진행할 공간을 확보하고, 중단한 요청을 다시 이어 가는 방법을 설명하겠습니다. 재개에 필요한 KV를 다시 계산하는 방법과 다른 메모리에 보관했다가 가져오는 방법을 비교한 뒤, 이 과정이 사용자가 기다리는 시간에 어떻게 나타나는지 보겠습니다.

생성 중 부족해지는 KV 공간

요청을 처음 실행할 때 입력의 KV를 담을 공간이 있었다고 해서, 생성이 끝날 때까지 필요한 공간이 모두 확보된 것은 아닙니다. 선택한 출력 토큰을 다음 입력으로 처리할 때마다 새 KV가 추가되기 때문입니다. 마지막 블록까지 차면 블록을 하나 더 받아야 합니다.

앞선 글보다 작은 풀을 사용해 보겠습니다. 전체 KV 풀은 물리 블록 4개이고, 블록 하나에는 네 위치의 KV가 들어갑니다. 요청 A와 B는 각각 블록 2개씩을 사용하며, 서로 공유하는 블록은 없습니다. 그림의 한 칸은 해당 토큰 위치의 KV를 저장하는 공간이고, 숫자는 요청 안의 위치를 뜻합니다.

처음에는 A와 B가 각각 7개 위치의 KV를 가지고 있습니다. 각 요청에 8칸씩 배정되어 있으므로 마지막 블록에 한 칸씩 남아 있습니다. 가용 블록이 0개여도, 이미 배정한 블록 안에 빈칸이 있으면 다음 입력을 처리할 수 있습니다.

그림 1에서 두 요청이 다음 입력을 한 번씩 처리하면 KV가 각각 8개가 됩니다. 이 시점에는 네 블록의 모든 칸이 차 있습니다. A가 그다음 위치 8의 입력을 처리하려면 세 번째 블록이 필요하지만, 풀에는 배정할 블록이 없습니다. B도 생성을 더 이어 가려면 같은 문제를 만나게 됩니다. 위치 번호는 0부터 세므로, 위치 8의 KV는 아홉 번째 KV입니다.

새 요청이 들어오지 않아도 이런 상황이 생깁니다. 이미 실행 중인 요청들의 출력이 길어지는 것만으로 공간이 부족해질 수 있기 때문입니다. 새 요청을 대기 큐에 두는 조치는 추가 점유를 막지만, 현재 A와 B가 필요로 하는 공간을 만들어 주지는 못합니다.

최대 출력 길이까지 미리 예약해 두면 이후의 공간 부족을 줄일 수 있습니다. 대신 짧게 끝날 요청도 큰 공간을 오래 차지하게 됩니다. 필요한 만큼 배정하는 방식에서는 이런 예약 낭비를 줄이는 대신, 생성 도중 추가 공간이 부족할 때 어떻게 진행할지도 정해야 합니다.

요청 중단과 KV 블록 반환

이제 A가 계속 진행할 수 있도록 B를 잠시 중단한다고 하겠습니다. 여기서는 B를 선택한 경우의 동작을 살펴보며, 어느 요청을 먼저 중단할지 정하는 정책은 뒤로 미루겠습니다.

B를 다음 배치에서 제외하면 이번 실행에 처리할 토큰 수는 줄어듭니다. 그러나 B의 KV가 GPU에 그대로 남아 있다면, B가 쓰던 두 블록도 계속 점유됩니다. 실행할 요청을 줄이는 것과 KV 저장 공간을 확보하는 것은 서로 다른 작업입니다. A에 필요한 새 블록을 만들려면 B의 블록을 실제로 반환해야 합니다.

그림 2의 첫 장면에서는 B의 실행만 보류해 가용 블록이 여전히 0개입니다. 다음 장면에서 B가 독점하던 P2와 P3를 반환하면 가용 블록이 2개가 됩니다. 엔진은 이 중 P2를 A에 배정하고 A의 블록 테이블에 연결합니다. A가 다음 입력을 처리하면 P2의 첫 칸에 위치 8의 KV가 기록됩니다. A는 이제 KV 9개를 블록 3개에 보관하고, P3 하나가 가용 블록으로 남습니다.

블록을 반환한다는 것은 그 공간을 다시 배정할 수 있는 상태로 돌려놓는다는 뜻입니다. 모든 값을 0으로 지워야 하는 것은 아닙니다. 이후 다른 요청이 배정받아 자기 KV를 기록할 수 있습니다. 이 예시는 독점 블록을 사용하므로 두 블록을 모두 돌려받지만, 앞 글에서 본 것처럼 다른 요청이 참조 중인 공유 블록은 그대로 유지해야 합니다.

이처럼 진행 중인 요청을 중단하고 자원을 회수해 다른 요청이 진행하도록 하는 것을 Preemption(선점)이라고 합니다. 여기서는 KV 공간 부족에 대응하는 요청 단위의 선점을 다룹니다. 엔진이 모델 실행 단계 사이에서 다음 실행 대상을 조정하는 것이며, 실행 중인 GPU 커널을 강제로 끊는 기능을 뜻하지는 않습니다.

중단한 B는 완료되거나 취소된 요청이 아닙니다. 입력과 이미 생성한 토큰 이력, 생성 설정, 어디까지 출력을 전달했는지 같은 요청 상태를 보존한 채 재개를 기다립니다. 다만 GPU KV를 반환했으므로, 다시 실행할 때는 계산에 필요한 KV를 마련해야 합니다.

토큰 이력으로 KV 재계산하기

먼저 GPU KV를 반환하고, 나중에 필요한 KV를 다시 계산하는 재계산(Recompute) 방식을 살펴보겠습니다. 이를 이해하려면 토큰 이력과 KV를 구별해야 합니다. 토큰 이력에는 입력과 선택한 출력의 토큰 ID가 순서대로 들어 있습니다. KV는 그 토큰들을 모델로 처리해 얻은 계산 결과입니다. 토큰 이력을 보존하면 KV가 없어져도 어떤 토큰들로 계산해야 하는지는 알 수 있습니다.

앞선 그림의 B를 그대로 이어 보겠습니다. B의 입력은 p0부터 p6까지 7토큰입니다. 입력을 처리해 x0를 선택하고, x0를 다음 입력으로 처리해 x1까지 선택·전달한 상태에서 중단했습니다.

이때 토큰 이력은 입력 7개와 출력 x0, x1을 합한 9개입니다. 반면 GPU KV는 입력 7개와 x0에 해당하는 8개뿐입니다. x1은 다음 실행에 들어갈 토큰이므로 아직 자기 KV가 없습니다. 앞선 그림에서 B가 두 블록을 가득 채운 상태와 같습니다.

그림 3의 가운데에서 B의 GPU KV는 없지만, 토큰 이력과 이미 전달한 x0, x1은 그대로 남아 있습니다. 엔진이 B를 다시 실행하기로 정하고 충분한 공간을 확보하면, 보존한 9개 토큰을 모델의 입력으로 넣습니다. 이 예시에서는 재사용할 캐시가 없으므로 전체 이력을 함께 처리합니다.

처리 결과는 두 부분으로 나눠 볼 수 있습니다. p0부터 p6과 x0의 KV 8개는 과거에 계산했던 상태를 다시 만드는 것이고, x1의 KV는 이번에 처음 계산합니다. 합계 9개 KV를 보관하려면 블록 3개가 필요합니다. 마지막 입력 x1 위치의 계산 결과로 새 토큰 x2를 선택해 전달할 수 있습니다. x2 자체의 KV는 그 토큰을 다음 입력으로 처리할 때 만들어집니다.

재계산할 때 x0와 x1은 이미 값이 정해진 입력입니다. 이 둘을 다시 샘플링하거나 사용자에게 다시 보내지 않습니다. 보존된 답변을 그대로 이어 가기 위해 계산 상태를 복구하고, 새 출력 x2부터 전달합니다.

여기서 Prefill과 Decode에서 본 차이가 다시 나타납니다. 처음 답변을 만들 때는 다음 출력의 값을 모르므로 하나씩 생성해야 했습니다. 재계산 시점에는 과거 출력도 모두 알려져 있으므로, 입력과 함께 하나의 토큰열로 처리할 수 있습니다. PagedAttention 논문의 재계산 설명도 생성한 이력을 입력에 붙여 Prefill 형태로 계산하는 방법을 다룹니다.

이 때문에 과거의 생성 과정을 한 토큰씩 그대로 반복할 필요는 없습니다. 그래도 이미 수행했던 계산을 다시 한다는 비용은 남습니다. 복구할 이력이 길수록 다시 처리해야 하는 토큰도 많아집니다. 실제 엔진에서 일부 KV를 재사용할 수 있다면 필요한 복구 범위는 줄어들 수 있습니다.

GPU 밖에 보관한 KV 복원하기

GPU KV를 회수하는 또 다른 방법은 계산해 둔 값을 다른 메모리에 보관하는 것입니다. 이후 다시 실행할 때 그 값을 GPU로 가져오는 방식을 스왑(Swap)이라고 합니다. 여기서는 보관 장소로 CPU 메모리를 사용하겠습니다. 재계산 뒤에 추가로 수행하는 단계가 아니라, 같은 중단 시점에서 선택할 수 있는 다른 방법입니다.

B가 x1까지 전달하고 KV 8개를 두 블록에 가지고 있던 시점으로 돌아가겠습니다. 엔진은 B의 KV를 CPU 메모리로 복사합니다. 복사가 끝나 보관이 완료되면 GPU의 두 블록을 반환하고, 그 공간을 다른 요청에 배정할 수 있습니다. 전송이 끝나기 전에 공간을 다른 값으로 덮어쓰면 보관할 KV를 잃을 수 있으므로, 반환 시점도 관리해야 합니다.

그림 4의 가운데에서는 B의 GPU 블록이 반환되어 있지만 CPU에는 KV 8개가 남아 있습니다. B의 토큰 이력과 요청 상태도 계속 보존합니다. 나중에 B를 재개할 공간을 확보하면 GPU 블록을 다시 배정하고, 보관한 KV를 그곳으로 가져옵니다. 새로 받은 블록의 위치가 이전과 달라도, 블록 테이블이 현재 저장 위치를 가리키면 됩니다.

KV를 복원한 뒤에도 x1을 처리하는 계산은 필요합니다. 저장해 둔 KV는 x0까지이기 때문입니다. 기존 두 블록이 모두 차 있으므로 블록 하나를 더 배정하고, x1을 입력으로 처리해 새 KV를 기록합니다. 그 결과 KV는 9개가 되고 새 출력 x2를 선택할 수 있습니다. 이 경로에서는 과거 8개 위치의 KV를 다시 계산하는 대신, 보관한 값을 복원합니다.

두 방법이 요구하는 자원은 다음처럼 구별할 수 있습니다.

방식	중단 중 보존하는 것	재개를 준비하는 작업	추가로 드는 비용
재계산	토큰 이력과 요청 상태	보존한 이력으로 필요한 KV를 다시 계산	모델 계산
스왑	토큰 이력과 요청 상태, CPU에 보관한 KV	저장한 KV를 GPU로 복원	보관할 메모리와 왕복 전송

재계산 비용은 처리할 토큰 수와 모델의 계산 효율에 영향을 받습니다. 스왑 비용은 옮겨야 할 KV의 양과 CPU·GPU 사이에서 사용할 수 있는 전송 대역폭에 영향을 받습니다. 따라서 어느 쪽이 유리한지는 문맥 길이, 하드웨어와 엔진의 실행 방식에 따라 달라집니다. PagedAttention 논문은 두 방법을 비교하며 계산 성능과 전송 대역폭을 함께 고려합니다.

이 비교는 두 방법의 원리를 설명하기 위한 것입니다. 실제로 어떤 방식을 지원하고 선택하는지는 엔진과 버전에 따라 다릅니다.

중단과 복구가 만드는 출력 지연

마지막으로 재계산 예시를 시간순으로 이어 보겠습니다. B가 x0와 x1을 전달한 뒤 중단하면 두 블록이 반환됩니다. A는 그중 하나를 받아 생성을 계속하고, B는 GPU KV 없이 재개를 기다립니다.

이번 예시에서 A는 KV를 9개에서 12개까지 늘린 뒤 생성을 마칩니다. 블록 3개로 수용할 수 있는 범위이므로 추가 블록은 필요하지 않습니다. A가 완료되어 세 블록을 반환하면, 엔진은 B에 필요한 세 블록을 배정하고 보존한 9개 토큰으로 KV를 재계산합니다.

그림 5의 아래쪽은 사용자가 받는 B의 출력입니다. x1까지 받은 뒤 B가 기다리고 KV를 재계산하는 동안에는 새 출력이 없습니다. 재계산이 끝나 x2를 선택·전달한 다음에야 생성이 다시 이어집니다. 사용자에게는 x1과 x2 사이의 긴 간격으로 나타납니다. vLLM의 Preemption 문서에서도 중단과 재계산이 요청의 지연을 늘릴 수 있음을 설명합니다.

이 예시에서는 A의 완료가 B를 재개할 공간을 마련해 줍니다. 일반적으로는 다른 요청의 완료뿐 아니라 전체 가용 공간과 실행 순서에 따라 재개 시점이 정해집니다. 공간이 생긴 순간 반드시 바로 실행되는 것도 아닙니다. 스케줄러가 다른 요청들과 함께 다음 실행 대상을 결정해야 하기 때문입니다.

Preemption은 자원이 부족한 상태에서 일부 요청이 계속 진행할 수 있게 합니다. 그 대신 중단한 요청은 기다리고, 다시 실행하기 위한 상태를 복구해야 합니다. 이런 일이 반복되면 같은 이력을 여러 번 계산하거나 같은 KV를 여러 번 옮기면서도, 새 출력은 조금밖에 만들지 못할 수 있습니다.

따라서 동시에 받아들인 요청 수만으로 추론이 잘 진행되는지 판단하기는 어렵습니다. 실제로 새 출력을 얼마나 만들었는지, 첫 출력을 받기까지 얼마나 기다렸는지, 출력 사이의 간격이 얼마나 길어졌는지를 함께 보아야 합니다. 다음 글인 「추론 성능 지표: 대기 시간과 처리량」에서는 지금까지 살펴본 대기·실행·복구 과정을 어떤 지표로 측정할지 설명하겠습니다.