vllm의 PagedAttention

1. PagedAttention이란?

(Self-Attention 시) Input token 간의 연관성을 계산하는 과정에서 Key (K)와 Value (V) 행렬을 생성

기존 KV 캐싱 방식:

컴퓨터에서 데이터를 저장할 때 메모리 주소(메모리 공간의 위치)를 사용한다.

아래처럼 A, B, C, D 데이터를 연속된 주소 공간에 배치한 것이 연속된 메모리 블록이다.

LLM의 Inference에서는 Self-Attention 연산을 위해 이전 토큰들의 Key (K)와 Value (V) 행렬을 저장하는 KV 캐시를 사용

기존 방식에서는 KV 캐시를 한 번에 큰 연속된 메모리 블록으로 할당해야 했음
하지만 긴 문맥(Context)이 필요하면 KV 캐시가 점점 커지므로 연속된 큰 공간을 확보하는 것이 어려워진임
특히 GPU 메모리는 단편화(Fragmentation)되기 쉬워 충분한 총 메모리가 있어도 연속된 공간을 확보하지 못하면 할당이 실패할 수 있음

[예시: 연속된 메모리 할당이 어려운 경우]

기존에 여러 개의 작은 데이터(A, B, C)가 저장된 메모리에서 새로운 1GB의 연속된 공간이 필요한 경우를 생각해보자!

이 경우, 총 빈 공간이 1GB 이상 있어도 연속된 공간이 없어서 할당이 불가능할 수 있음!

PagedAttention은 연속된 메모리 블록을 요구하지 않고, 작은 페이지 단위로 관리하여 이 문제를 해결

[예시: PagedAttention 방식의 메모리 할당 예시]

PagedAttention을 사용하면 KV 캐시를 작은 페이지(예: 16KB 단위)로 나누어 배치한다.

이렇게 하면 연속된 큰 공간을 찾을 필요 없이, 작은 단위로 메모리를 효율적으로 활용할 수 있다.

(1) 연속된 메모리 블록: 한 번에 큰 크기의 연속된 메모리 공간을 확보해야 하는 방식 → 메모리 단편화로 인해 할당이 어려울 수 있음.
(2) PagedAttention: 작은 페이지 단위로 KV 캐시를 관리하여 비연속적인 메모리 공간에서도 효율적으로 할당 가능.
(3) 이를 통해 LLM의 긴 문맥 지원을 원활하게 하면서도 속도와 메모리 사용량을 최적화할 수 있다.

기존 KV 캐싱 vs. PagedAttention