Efficient Memory Management for Large Language Model Serving with PagedAttention

논문 개요

LLM을 프로덕션 환경에서 서빙할 때 가장 큰 병목 중 하나는 메모리 관리다. 특히 트랜스포머의 어텐션 메커니즘은 모든 이전 토큰의 키-값 벡터(KV 캐시)를 저장해야 하며, 이 캐시는 생성 길이와 배치 크기에 따라 선형적으로 증가한다.

Kwon et al.(2023)이 SOSP 2023(운영체제 분야의 최고 학술대회)에서 발표한 이 논문은 LLM 서빙의 메모리 비효율 문제를 운영체제(OS) 관점에서 근본적으로 재접근한다. 기존 시스템(FasterTransformer, Orca 등)은 KV 캐시를 연속된 메모리 공간에 사전 할당하여 관리하는데, 이로 인해 세 가지 유형의 메모리 낭비가 발생한다: (1) 예약 낭비(reserved waste) - 최대 생성 길이만큼 미리 확보, (2) 내부 단편화(internal fragmentation) - 실제 생성 길이가 사전 할당보다 짧을 때, (3) 외부 단편화(external fragmentation) - 다양한 크기의 요청이 섞일 때.

---

핵심 기여

1. PagedAttention 알고리즘: KV 캐시를 고정 크기의 블록(페이지)으로 분할하고, 비연속 물리 메모리에 저장하면서 논리-물리 블록 테이블로 매핑하는 새로운 어텐션 알고리즘.
2. vLLM 시스템: PagedAttention을 기반으로 구현된 고성능 LLM 서빙 시스템으로, 동적 메모리 할당 및 해제, 병렬 요청 스케줄링을 지원.
3. Copy-on-Write 기반 메모리 공유: 빔 서치(beam search)나 병렬 샘플링 시 공통 프롬프트의 KV 캐시를 물리적으로 공유하여 메모리를 추가 절약.
4. 2~4× 처리량 향상: OPT-13B, LLaMA-13B 등 다양한 모델에서 기존 서빙 시스템 대비 대폭 향상된 처리량을 달성.

---

방법론 상세

기존 방식의 문제점

기존 시스템에서 요청 $r_i$에 대한 KV 캐시 크기를 $s_i$라 하면, 미리 최대 생성 길이 $L_{\max }$에 맞춰 다음만큼 메모리를 예약한다: $M_{\text{reserved}}=n_{\text{layer}}\times 2\times d_{\text{head}}\times L_{\max }\times \text{dtype\_bytes}$

실제 생성된 토큰이 $l_i<L_{\max }$이면 $(L_{\max }-l_i)$만큼의 메모리가 낭비된다. 실제 서빙 환경에서 메모리 활용률이 20~40%에 불과한 경우도 있다고 보고되었다.

PagedAttention의 핵심 아이디어

OS의 페이지 테이블에서 영감을 받아, KV 캐시를 다음과 같이 관리한다:

• 블록(Block): 고정 크기 $B$토큰의 KV 캐시를 저장하는 단위. 예: $B=16$
• 논리 블록(Logical Block): 각 요청의 순서 기준 KV 캐시 위치
• 물리 블록(Physical Block): GPU 메모리상의 실제 저장 위치
• 블록 테이블: 논리 블록 번호 → 물리 블록 번호 매핑

요청의 논리 블록 $k$가 물리 블록 $f(k)$에 저장된다면, $i$번째 레이어의 $j$번째 헤드의 $t$번째 토큰 키 벡터 접근은:

$K[i][j][t]=\text{PhysicalMemory}[f(\lfloor t/B\rfloor )][i][j][t\%B]$

어텐션 계산은 다음과 같이 블록 단위로 수행된다:

$\text{Attention}(q,K,V)={\sum }_{k=0}^{\lceil T/B\rceil -1}\text{BlockAttention}(q,K_k,V_k)$

여기서 $K_k,V_k$는 $k$번째 논리 블록의 KV 캐시이며, 물리 메모리에서 비연속적으로 위치할 수 있다.

메모리 관리자 (Memory Manager)

물리 블록 풀(pool)을 중앙에서 관리하며 다음 작업을 수행한다:

• 할당(Alloc): 새 토큰 생성 시 빈 물리 블록을 논리 블록에 할당
• 해제(Free): 요청 완료 시 해당 요청의 모든 물리 블록 반환
• 공유(Share): 동일한 프롬프트를 사용하는 요청들은 동일한 물리 블록을 가리키는 논리 블록 포인터를 공유

Copy-on-Write 메모리 공유

병렬 샘플링(parallel sampling)에서 동일 프롬프트로 $n$개의 응답을 생성할 때:

1. 프롬프트의 KV 캐시를 담은 물리 블록들을 $n$개 시퀀스가 공유 (참조 카운트 유지)
2. 각 시퀀스가 새 토큰을 생성할 때 해당 블록이 처음 수정되면, 그때 물리 블록을 복사(copy)
3. 이후 복사된 블록에 새 토큰의 KV를 기록

이를 통해 프롬프트 처리 메모리가 $O(n\times L_p)\to O(L_p)$로 줄어든다.

프리엠션(Preemption) 정책

GPU 메모리가 부족할 때 vLLM은 처리 중인 요청을 일시 중단하고, 해당 KV 캐시를 CPU 메모리로 스왑 아웃(swap-out)한다. 메모리가 확보되면 스왑 인(swap-in)하여 처리를 재개한다. 이는 OS의 스왑 파티션과 유사한 개념이다.

---

실험 결과

처리량 비교 (OPT-13B, A100 80GB)

시스템	처리량 (req/s)	vLLM 대비
FasterTransformer	2.2	기준
Orca	4.1	1.9×
vLLM (PagedAttention)	8.5	3.9×

메모리 활용률

기존 시스템에서 2040%였던 KV 캐시 메모리 활용률이 PagedAttention 도입 후 8095%로 향상되었다. 단편화로 인한 낭비가 거의 사라졌음을 의미한다.

다양한 서빙 시나리오

시나리오	향상
단일 시퀀스 생성	2.1×
병렬 샘플링 (n=4)	4.1×
빔 서치 (beam=4)	3.7×
긴 문맥 (2048 토큰)	3.3×

---

의의 및 한계

의의

• OS 아이디어의 ML 적용: 수십 년간 검증된 OS의 메모리 관리 기법을 LLM 서빙에 창의적으로 적용하였다.
• 실용적 영향: vLLM은 발표 직후부터 Hugging Face, Anyscale, Modal 등 주요 LLM 서빙 플랫폼에 통합되어 사실상 업계 표준이 되었다.
• 높은 재현성: 완전한 오픈소스 구현(github.com/vllm-project/vllm)을 공개하여 커뮤니티의 폭넓은 활용을 가능하게 하였다.

한계

• 블록 크기 선택: 블록 크기 $B$는 하이퍼파라미터이며, 너무 작으면 블록 테이블 오버헤드가, 너무 크면 내부 단편화가 증가한다.
• CPU-GPU 스왑 오버헤드: 프리엠션 시 CPU 스왑의 대역폭 한계로 인해 지연이 발생할 수 있다.
• 멀티-모달 확장: 이미지·오디오 토큰처럼 가변 크기의 KV 캐시를 처리하는 데는 추가 설계가 필요하다.
• 분산 서빙: 텐서 병렬 처리(tensor parallelism) 환경에서 블록 테이블 동기화가 추가 복잡성을 유발한다.