페이지 단위 KV 관리

GPU 한 장이 담을 수 있는 KV 캐시의 총량은 모델과 메모리가 정한다. Llama 3.1 8B를 H100 80GB에 올리면 약 42만 토큰이다. 그런데 이 총량을 요청들에게 어떻게 나눠 주느냐에 따라 실제로 동시에 처리할 수 있는 요청 수가 몇 배씩 달라진다. 요청마다 연속된 공간을 통째로 잡아 주는 방식은 이 총량의 상당 부분을 쓰지 못한 채 남긴다.

연속 할당의 낭비

초기 서빙 시스템은 요청마다 KV 캐시를 하나의 연속된 텐서로 잡았다. 어텐션 커널이 연속된 메모리를 가정했기 때문이다. 요청이 끝날 때까지 몇 토큰을 생성할지 미리 알 수 없으므로 공간은 최대 생성 길이 기준으로 잡는다. 이 방식에서 메모리는 세 가지로 샌다.

낭비	생기는 이유
예약 낭비	최대 길이 2,048토큰으로 잡았는데 요청이 300토큰에서 끝나면 나머지 1,748토큰 자리는 끝내 쓰이지 않는다. 생성이 진행 중인 동안에도 아직 채우지 않은 뒷부분은 다른 요청이 못 쓴다
내부 단편화	할당 단위를 최대 길이로 정해 두면 실제 길이와의 차이가 그대로 버려진다
외부 단편화	길이가 제각각인 공간이 할당되고 풀리기를 반복하면 빈 공간이 여기저기 흩어진다. 합치면 충분해도 연속된 큰 공간이 없어 새 요청을 받지 못한다

PagedAttention 논문(SOSP 2023)은 기존 시스템에서 KV 캐시 메모리 가운데 실제 토큰 상태를 저장하는 데 쓰인 비율이 20.4~38.2%에 그쳤다고 측정했다. 나머지는 위 세 가지 낭비였다.

블록과 블록 테이블

PagedAttention은 운영체제의 가상 메모리 페이징을 KV 캐시에 가져온다. KV 캐시를 고정된 토큰 수, 예를 들어 16토큰짜리 블록으로 나누고, 요청마다 논리 블록 번호를 실제 메모리의 물리 블록 번호로 옮기는 블록 테이블을 둔다.

운영체제	PagedAttention
프로세스	요청
페이지	KV 블록 (고정 토큰 수)
페이지 테이블	블록 테이블
가상 주소의 연속성	요청 안에서 토큰 순서의 연속성

요청이 생성을 시작하면 첫 블록 하나만 받는다. 블록이 다 차면 빈 블록 하나를 더 받아 블록 테이블에 이어 적는다. 물리 블록은 메모리 어디에 있어도 상관없다. 어텐션 커널이 블록 테이블을 보고 요청의 블록들을 차례로 찾아 읽기 때문이다. 요청이 끝나면 블록을 모두 반납한다.

이렇게 하면 앞의 세 가지 낭비가 거의 사라진다. 요청은 쓰는 만큼만 블록을 받고 모든 블록의 크기가 같으므로 예약 낭비와 외부 단편화가 함께 사라진다. 남는 낭비는 요청마다 마지막 블록의 빈자리뿐이라, 블록 크기가 16토큰이면 요청 하나당 최대 15토큰이다. 논문은 이 방식으로 같은 지연 수준에서 처리량을 기존 시스템보다 2~4배 높였다고 보고한다.

블록 크기

블록 크기는 두 가지를 맞바꾼다. 블록이 작으면 마지막 블록의 빈자리가 줄어 메모리를 알뜰하게 쓰지만, 블록 테이블이 길어지고 어텐션 커널이 읽어야 할 블록 수가 늘어 메모리 접근이 잘게 쪼개진다. 블록이 크면 커널이 한 번에 연속으로 읽는 양이 늘어 효율이 좋아지지만 빈자리가 커진다.

요청당 낭비는 블록 크기보다 작으므로 수천 토큰짜리 문맥에서 16토큰 블록의 빈자리는 1% 미만이다. 블록 크기는 메모리 낭비보다 커널 효율과 아래의 공유 단위를 보고 정하며, 엔진마다 기본값과 허용 범위가 다르니 문서를 확인한다.

블록 공유

블록 테이블을 두면 요청끼리 물리 블록을 함께 쓸 수 있다. 블록마다 몇 개의 요청이 가리키고 있는지 참조 수를 세어 두고 참조 수가 0이 될 때 반납한다.

• 병렬 샘플링: 같은 프롬프트로 응답 n개를 만들면 프롬프트 부분의 블록은 n개 요청이 함께 가리킨다. 응답이 갈라지는 지점부터 각자 블록을 받는다
• 빔 서치: 후보들이 공통 접두사를 공유하다가 갈라지고, 탈락한 후보의 블록은 참조 수가 줄어 반납된다
• 공통 프롬프트: 시스템 프롬프트처럼 여러 요청이 같은 앞부분을 가지면 그 부분의 블록을 한 번만 계산해 저장하고 함께 쓴다

공유 중인 블록에 한 요청이 새 토큰을 써야 하면 그 블록을 복사해 자기 것으로 만든 뒤 쓴다. 운영체제의 쓰기 시 복사(copy-on-write)와 같은 방식이다. 블록이 다 찬 뒤에는 새 토큰이 새 블록으로 가므로 복사는 갈라지는 지점의 마지막 블록 하나에서만 일어난다.

선점과 복구

블록을 필요할 때마다 받으므로 진행 중인 요청들이 함께 길어지다 보면 빈 블록이 바닥날 수 있다. 이때 엔진은 일부 요청을 선점해 블록을 회수한다. 회수한 요청을 나중에 이어 가는 방법은 두 가지다.

방식	동작	비용
재계산	블록을 버리고, 재개할 때 프롬프트와 지금까지 생성한 토큰을 이어 붙여 다시 프리필한다	프리필 계산. 이미 연산 성능에 여유가 있는 디코드 위주 부하에서는 부담이 작다
스왑	블록을 CPU 메모리로 옮겼다가 재개할 때 다시 가져온다	PCIe 전송. 블록이 많으면 전송 시간이 재계산보다 길다

블록 단위 관리에서는 요청의 KV 캐시 전체를 한꺼번에 처리하므로 일부 블록만 내보내는 일은 하지 않는다. 일부만 내보내면 그 요청은 어차피 진행할 수 없기 때문이다.

선점은 메모리 부족을 피하는 안전장치이지 정상 운영 상태가 아니다. 선점이 잦으면 동시 처리 상한을 낮추거나, KV 캐시를 양자화해 블록 수를 늘리거나, 문맥 상한을 줄인다. 반대로 블록 사용률이 늘 낮다면 동시 처리 상한을 올려 처리량을 늘릴 여지가 있다. vLLM은 이 사용률을 vllm:kv_cache_usage_perc 지표로 내보낸다.

지금의 엔진

블록 단위 KV 관리는 이제 추론 엔진의 기본 구조가 됐다. vLLM이 PagedAttention으로 처음 보였고 TensorRT-LLM도 페이지 단위 KV 캐시를 연속 배치(문서의 이름으로는 in-flight batching)와 함께 쓴다. 엔진을 고를 때 이 기능 자체보다는 블록 크기, 선점 방식, 공유를 어디까지 지원하는지, 그리고 그 위에 프리픽스 캐시가 어떻게 올라가 있는지를 비교한다.