Efficient Memory Management for Large Language Model Serving with PagedAttention
PagedAttention을 이용한 대규모 언어 모델 서빙의 효율적 메모리 관리
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, et al. (2023)
운영체제의 가상 메모리 및 페이징 기법에서 영감받은 PagedAttention을 제안하여 KV 캐시 메모리를 비연속적 블록으로 관리함으로써, LLM 서빙의 처리량을 기존 시스템 대비 2-4배 향상시켰다.
배경
LLM 서빙에서 KV 캐시(key-value cache)는 자기 회귀적 디코딩의 효율성을 위해 필수적이지만, 전체 GPU 메모리의 상당 부분(최대 30% 이상)을 차지한다. 기존 서빙 시스템(FasterTransformer, Orca 등)은 각 요청의 KV 캐시를 연속적(contiguous) 메모리 공간에 할당했는데, 이는 세 가지 심각한 비효율을 초래했다: (1) 최대 시퀀스 길이를 기준으로 사전 할당하여 사용되지 않는 메모리가 많고(내부 단편화), (2) 요청 간 빈 공간이 발생하며(외부 단편화), (3) 빔 서치 등에서 KV 캐시를 복제해야 해 메모리 낭비가 발생한다. 이러한 비효율로 실제 활용률은 20-40%에 불과했다.
핵심 아이디어
PagedAttention은 KV 캐시를 고정 크기의 비연속적 블록(page)으로 분할하여 관리한다. 운영체제의 가상 메모리 시스템처럼, 논리적으로 연속인 KV 캐시를 물리적으로 비연속적인 블록에 저장하고 블록 테이블(page table)로 매핑한다. 이를 통해 (1) 내부 단편화를 마지막 블록에만 한정하고, (2) 외부 단편화를 완전히 제거하며, (3) copy-on-write 메커니즘으로 빔 서치나 병렬 샘플링에서 KV 캐시를 물리적으로 공유할 수 있다. 블록 크기는 보통 16 토큰으로 설정되며, 이는 OS의 4KB 페이지 크기에 해당하는 역할을 한다.
방법론
각 어텐션 헤드의 KV 캐시를 고정 블록 크기(B 토큰)의 물리 블록으로 분할하고, 요청별 블록 테이블로 논리 블록과 물리 블록의 매핑을 관리한다. 새 토큰 생성 시 현재 블록에 여유가 있으면 추가하고, 없으면 새 물리 블록을 할당한다. 어텐션 계산 시 블록 테이블을 참조하여 비연속적 물리 블록에서 KV를 fetch하는 맞춤 CUDA 커널을 구현했다. 병렬 샘플링에서 같은 프롬프트의 KV 캐시를 공유할 때는 참조 카운트를 증가시키고, 수정 시에만 복사하는 copy-on-write를 적용한다. 이 시스템을 vLLM이라는 오픈소스 서빙 엔진으로 구현했다.
주요 결과
vLLM은 동일한 지연 시간(latency) 조건에서 HuggingFace Transformers 대비 최대 24배, FasterTransformer 대비 최대 3.5배의 처리량(throughput) 향상을 달성했다. KV 캐시 메모리 낭비를 기존 60-80%에서 4% 미만으로 줄였다. 병렬 샘플링에서는 KV 캐시 공유를 통해 55%의 메모리 절감을 보였다. 빔 서치에서도 copy-on-write를 통해 메모리 사용량이 빔 수에 비례하여 증가하지 않았다. OPT-13B, OPT-175B 등 다양한 모델 크기에서 일관된 성능 향상을 확인했다.
임팩트
vLLM은 오픈소스 LLM 서빙의 사실상 표준 프레임워크로 자리잡아, 학술 연구와 산업 프로덕션 모두에서 가장 널리 사용되는 LLM 추론 엔진이 되었다. PagedAttention의 아이디어는 TensorRT-LLM, TGI(Text Generation Inference) 등 다른 서빙 프레임워크에도 채택되었다. 지속적 배칭(continuous batching), 추측적 디코딩, 양자화(GPTQ/AWQ) 등과의 통합을 통해 LLM 서빙의 전체 스택을 커버하는 생태계로 발전하고 있다. SOSP 2023에서 발표되어 시스템 커뮤니티와 ML 커뮤니티의 교류를 촉진했다.