PagedAttention

PagedAttention은 attention 계산을 새로 정의한 알고리즘이라기보다, KV cache를 GPU memory에 배치하는 방식을 바꾼 아이디어다.

한 문장으로 말하면:

각 요청의 KV cache를 연속된 큰 tensor로 잡지 않고,
고정 크기 block들의 목록으로 관리한다.

이때 각 요청은 자기 token 위치가 실제 GPU memory의 어떤 KV block에 있는지 알려주는 block table을 가진다.

왜 필요한가

LLM serving에서는 여러 요청이 동시에 들어온다. 각 요청은 prompt 길이도 다르고, output이 몇 token까지 이어질지도 미리 알기 어렵다.

단순한 방식은 요청마다 최대 길이만큼 연속 memory를 예약하는 것이다.

request A: 최대 4096 token 공간 예약, 실제로는 300 token 사용
request B: 최대 4096 token 공간 예약, 실제로는 1200 token 사용
request C: 최대 4096 token 공간 예약, 아직 80 token만 사용

이러면 GPU memory가 있어도 많은 부분이 비어 있는 채로 묶인다.

문제는 두 가지다.

internal fragmentation: 예약한 공간 안에서 쓰지 않는 부분이 생긴다.
external fragmentation: 빈 공간이 흩어져 있어 새 요청에 쓰기 어렵다.

KV cache는 요청 수와 sequence 길이에 비례해서 커진다. 그래서 이 낭비는 곧 batch 크기와 throughput을 제한한다.

OS paging과 같은 발상

PagedAttention은 운영체제의 virtual memory처럼 생각하면 쉽다.

논리적으로는 요청 하나가 긴 KV cache를 가진다.

logical KV cache for one request
[token 0 ... token 15] [token 16 ... token 31] [token 32 ... token 47]

하지만 물리적으로는 GPU memory의 연속된 한 덩어리가 아닐 수 있다.

logical block index:  0   1   2
block table:        [17, 04, 91]

physical KV blocks in GPU memory:
block 04  ...  block 17  ...  block 91  ...

요청은 “내 0번째 logical block은 physical block 17에 있다”는 식의 table만 들고 있으면 된다.

새 token이 만들어져서 KV cache가 길어지면, vLLM은 필요한 만큼 새 block을 할당해 block table 뒤에 붙인다.

무엇이 좋아지는가

PagedAttention의 핵심 이득은 on-demand allocation이다.

요청이 300 token만 사용하면 300 token에 필요한 block만 쓴다. 최대 길이 전체를 미리 잡지 않는다.

낭비는 보통 요청마다 마지막 block의 남는 공간 정도로 제한된다.

나쁜 경우:
요청마다 최대 sequence length만큼 낭비 가능

PagedAttention:
요청마다 마지막 block의 빈칸 정도만 낭비

이 덕분에 같은 GPU memory로 더 많은 요청을 batch에 올릴 수 있다. serving 관점에서는 memory utilization이 좋아지고, 그 결과 throughput을 더 끌어올릴 수 있다.

Kernel은 무엇이 달라지는가

KV cache가 물리적으로 연속되어 있지 않다면 attention kernel도 바뀌어야 한다.

일반 attention kernel은 과거 K/V가 하나의 연속 tensor에 있다고 가정하기 쉽다. PagedAttention에서는 kernel이 block table을 읽어가며 필요한 K/V block을 찾아야 한다.

중요한 점은 매번 KV cache를 다시 연속 tensor로 합치지 않는다는 것이다.

하지 않는 일:
흩어진 KV block을 큰 연속 tensor로 복사한 뒤 attention

하는 일:
block table을 따라가며 흩어진 KV block을 그대로 읽고 attention

scatter는 token 하나하나가 아니라 block 단위로 일어난다. block 내부는 여전히 연속되어 있으므로 GPU memory access의 coalescing을 완전히 포기하는 구조는 아니다.

헷갈리기 쉬운 구분

PagedAttention은 serving engine의 여러 기술과 붙어 있지만, 같은 개념은 아니다.

KV cache:
과거 token의 K/V를 저장해 재계산을 줄이는 기법

PagedAttention:
KV cache를 block 단위로 배치하고 block table로 참조하는 방식

continuous batching:
서로 다른 요청을 step마다 동적으로 batch에 넣고 빼는 scheduling 방식

FlashDecoding:
긴 context decode에서 attention 계산을 더 잘 병렬화하려는 kernel 방향

vLLM에서 이 개념들이 같이 등장하는 이유는 serving 병목이 서로 연결되어 있기 때문이다. continuous batching으로 많은 요청을 동시에 처리하려면 KV cache memory를 잘 관리해야 하고, block으로 흩어진 KV cache를 쓰려면 PagedAttention kernel이 필요하다.

vLLM에서의 위치

vLLM을 세 층으로 나누면 PagedAttention은 한 층에만 갇힌 기능이 아니다.

EngineCore / Scheduler:
어떤 요청에 어떤 KV block을 줄지 결정한다.

GPU worker:
요청별 block table과 slot mapping을 device tensor로 준비한다.

Attention kernel:
block table을 따라가며 실제 K/V block을 읽는다.

그래서 PagedAttention은 “kernel 최적화”이면서 동시에 “serving memory manager 설계”다.

연결

kv-cache: PagedAttention이 관리하려는 대상
inference-engine-layers: PagedAttention이 EngineCore와 GPU worker 양쪽에 걸치는 이유
vllm-kv-cache-engine-to-worker: vLLM 구현에서 block 할당이 worker로 전달되는 흐름

확인

PagedAttention은 attention 수식을 바꾸는가, KV cache memory layout을 바꾸는가?
요청마다 최대 sequence length만큼 KV cache를 미리 잡으면 어떤 낭비가 생기는가?
block table은 logical token/block 위치와 physical GPU block 사이에서 어떤 역할을 하는가?
PagedAttention과 continuous batching은 왜 자주 같이 등장하지만 같은 개념은 아닌가?