[논문 리뷰] Efficient Memory Management for Large Language Model Serving with PagedAttention
Summary
- PagedAttention은 LLM serving에서 KV cache가 동적으로 커지고 줄어드는 문제를 memory management 문제로 정의한다.
- vLLM은 KV cache를 fixed-size block으로 나누고, block table을 통해 logical block과 physical block의 mapping을 관리한다.
- 이 구조는 fragmentation과 reservation으로 낭비되던 GPU 메모리를 줄이고, 더 큰 batch size를 가능하게 만든다.
- 논문은 vLLM이 기존 FasterTransformer와 Orca 계열 system보다 같은 latency 수준에서 2-4배 높은 throughput을 보인다고 보고한다.
Introduction
LLM이 단순 모델이 아니라 실제 서비스로 쓰이기 시작하면 신경 써야 할 것이 늘어난다. 모델 하나가 한 번의 inference를 얼마나 빠르게 끝내는지도 중요하지만, ChatGPT 같은 서비스에서는 서로 다른 prompt와 output 길이를 가진 요청이 계속 들어온다. 이때 서빙 시스템은 개별 요청의 latency뿐 아니라, 같은 GPU에서 얼마나 많은 요청을 안정적으로 처리할 수 있는지도 같이 봐야 한다.
그 관점에서 보면 LLM serving의 핵심은 단순히 "모델 계산을 더 빠르게 하자"가 아니다. Autoregressive generation은 token을 하나씩 생성하고, 매 token마다 이전 token들의 key/value state를 다시 참조한다. 그래서 batch를 크게 묶어 GPU utilization을 높이고 싶어도, 요청마다 커지는 KV cache가 GPU 메모리를 먼저 압박한다. GPU에는 모델 가중치만 올라가는 것이 아니라, 각 요청의 진행 상태도 계속 메모리를 차지한다.
이 논문은 그 병목을 attention kernel 하나의 속도 문제가 아니라 KV cache memory management 문제로 다시 정의한다. 모델 가중치는 serving 동안 비교적 고정된 큰 덩어리지만, KV cache는 요청마다 길이가 다르고 generation이 진행되면서 계속 자란다. 즉 LLM serving의 메모리는 정적인 tensor allocation 문제가 아니라, 여러 요청이 동시에 들어오고 나가는 동적인 memory allocation 문제에 가깝다.
Figure 1은 이 문제의 규모를 보여준다. 13B 모델을 A100 40GB에서 serving할 때 model parameters가 GPU 메모리의 큰 부분을 차지하지만, KV cache도 요청별로 상당한 공간을 차지하고 동적으로 할당된다. 오른쪽 그래프는 기존 시스템에서 KV cache 메모리가 batch size에 따라 급격히 커지는 반면, vLLM은 이 증가를 더 완만하게 만들고 그 결과 serving throughput을 끌어올린다는 주장을 보여준다.
그래서 이 논문은 "LLM inference를 빠르게 하는 새로운 kernel"이라기보다, LLM serving에서 GPU 메모리를 어떻게 운영해야 하는지에 대한 시스템 논문으로 읽는 편이 맞다. 뒤에서 나오는 PagedAttention은 OS의 virtual memory와 paging에서 아이디어를 가져와, KV cache를 요청별 contiguous buffer가 아니라 block 단위로 관리하려는 시도다.
Memory Challenges in LLM Serving
기존 LLM serving system의 핵심 문제는 KV cache를 요청별로 연속된 공간에 저장하려 한다는 점이다. 처음에는 자연스러운 선택처럼 보인다. 한 요청의 token sequence가 있으니, 그 sequence의 key/value tensor를 하나의 연속된 buffer에 저장하면 접근이 단순해진다. 하지만 serving workload에서는 이 전제가 잘 맞지 않는다.
요청마다 prompt length가 다르고, output length는 미리 알 수 없다. 더 어려운 점은 시스템이 안정적으로 동작하려면 충분한 공간을 미리 예약해야 한다는 것이다. 그러면 실제 token state가 들어가지 않은 예약 공간이 생기고, block 내부에 남는 공간도 생기며, free memory가 여러 조각으로 흩어지는 단편화(fragmentation)도 생긴다. 논문은 이를 reserved slots, internal fragmentation, external fragmentation으로 나누어 설명한다.
Figure 2는 병목이 KV cache 메모리 관리에 있다는 점을 수치로 보여준다. 기존 방식에서는 KV cache 메모리 중 실제 token state를 담는 비율이 20.4%에서 38.2% 정도에 그친다. 반면 vLLM은 96.3%까지 올라간다. 이 숫자를 기준으로 보면, 이 논문의 목표는 "조금 더 빠른 attention"이 아니라 "GPU 메모리를 실제 요청 처리에 거의 다 쓰게 만들기"에 가깝다.
이 지점이 vLLM을 이해하는 핵심이다. vLLM은 단순히 빠른 inference engine이 아니라, LLM serving에서 GPU 메모리를 운영체제의 메모리처럼 관리하려는 시도에 가깝다.
Method
PagedAttention의 아이디어는 OS의 virtual memory와 paging을 KV cache에 적용하는 것이다. 한 요청의 KV cache를 논리적으로는 연속된 token sequence처럼 보이게 두되, 물리적으로는 고정 크기 KV block들에 나누어 저장한다. 그리고 block table을 통해 logical block과 physical block의 mapping을 관리한다.
이렇게 하면 한 요청의 KV cache가 GPU 메모리 위에서 반드시 연속될 필요가 없다. generation이 진행되면서 새 block이 필요할 때만 physical block을 할당하면 된다. 마지막 block을 제외하면 대부분의 block을 채울 수 있으므로 internal fragmentation도 줄어든다. 논문에서는 block size를 16 tokens로 둘 때 기본 실험에서 좋은 trade-off가 나온다고 본다.
PagedAttention은 단순한 allocator가 아니다. Attention kernel 자체가 block table을 따라 non-contiguous KV cache를 읽을 수 있어야 한다. 즉 시스템 계층의 memory management와 kernel implementation이 같이 바뀐다. 이 때문에 논문은 OS 아이디어를 가져오지만, 그대로 가져오는 것이 아니라 LLM serving에 맞게 다시 구현한다.
System Overview
vLLM은 PagedAttention 위에 올라간 serving system이다. 논문에서 vLLM은 크게 centralized scheduler, KV cache manager, worker로 나뉜다. Scheduler는 어떤 요청을 다음 iteration에서 실행할지 정하고, KV cache manager는 physical block을 할당하거나 회수하며, GPU worker는 실제 model execution과 PagedAttention kernel을 실행한다. 여기서 중요한 점은 vLLM이 단순히 "attention을 빠르게 계산하는 엔진"이 아니라, request lifecycle 전체를 KV cache allocation 문제로 다시 모델링한다는 것이다.
PagedAttention
논문 4장에서 먼저 나오는 것은 PagedAttention algorithm이다. 기존 attention은 한 sequence의 key/value vector가 메모리에 연속적으로 있다고 가정하기 쉽다. PagedAttention은 이 전제를 깨고, KV cache를 fixed-size KV block의 sequence로 본다. 한 요청은 논리적으로는 logical KV block들의 배열을 갖고, 실제 GPU 메모리에는 physical KV block들이 흩어져 있다. 이 둘 사이의 mapping은 block table이 관리한다.
Figure 5를 보면 query vector가 여러 physical block에 나뉘어 저장된 key/value vector를 따라 attention을 계산한다. 즉 block table은 단순한 bookkeeping metadata가 아니다. PagedAttention kernel이 실제 attention 계산 중에 block table을 참조해서 필요한 KV block을 찾아간다. 이 때문에 PagedAttention은 allocator와 kernel이 같이 설계된 구조다. OS의 page table과 비슷한 abstraction을 쓰지만, 최종 목표는 CPU program의 address translation이 아니라 GPU attention kernel이 non-contiguous KV cache를 자연스럽게 읽게 만드는 것이다.
KV Cache Manager
KV cache manager는 이 block abstraction을 request lifecycle에 맞춰 운영한다. 요청이 prompt phase를 지나 generation phase에 들어가면, 생성되는 token에 맞춰 logical KV block이 채워진다. 마지막 block은 아직 다 차지 않았을 수 있고, 다음 token을 위해 여유를 갖고 있다. 새 block이 필요할 때만 physical block을 할당하므로, 기존 방식처럼 요청마다 maximum sequence length만 큼 미리 잡아둘 필요가 없다. 이 설계가 reserved memory와 internal fragmentation을 줄인다.
Application to Other Decoding Scenarios
또 하나 중요한 기능은 copy-on-write다. Parallel sampling, beam search, shared prefix처럼 여러 sequence가 prefix KV cache를 공유할 수 있는 decoding pattern에서는 같은 physical block을 여러 logical block이 참조하게 한다. 그러다가 특정 sequence가 새 token을 추가하면서 block을 수정해야 할 때만 copy를 만든다. 이 구조 덕분에 vLLM은 단순한 greedy decoding뿐 아니라 더 복잡한 decoding algorithm에서도 memory saving을 얻는다.
이 copy-on-write가 잘 맞는 이유는 decoding algorithm의 구조 때문이다. Parallel sampling에서는 한 prompt에서 여러 output sequence를 생성하므로 prompt KV cache를 공유할 수 있다. Beam search에서도 여러 candidate가 긴 prefix를 공유한다. Shared prefix workload에서는 서로 다른 요청이라도 instruction이나 few-shot examples가 같을 수 있다. 기존 contiguous allocation 방식에서는 이런 공유를 하려면 sequence별 cache layout을 크게 바꿔야 하지만, PagedAttention에서는 block reference만 공유하면 된다. 공유가 깨지는 순간에만 새 physical block을 만들면 되므로, memory sharing의 단위가 request 전체가 아니라 block이 된다.
Scheduling and Preemption
이 구조에서 iteration-level scheduling이 중요하다. LLM generation은 요청마다 매번 한 token씩 진행되므로, scheduler는 iteration 단위로 active request를 조정할 수 있다. 매 iteration마다 scheduler는 어떤 요청을 실행할지 정하고, KV cache manager는 그 요청들이 다음 token을 생성하는 데 필요한 block을 확보한다. 메모리가 충분하면 waiting request를 running 상태로 올릴 수 있고, 부족하면 일부 요청을 기다리게 하거나 이미 실행 중인 요청을 선점(preemption)할 수 있다.
Preemption은 vLLM의 메모리 관리에서 중요한 부분이다. KV cache는 크지만, 요청의 token sequence 자체는 다시 계산할 수 있다. 그래서 vLLM은 선점된 요청의 KV cache를 버리고 나중에 prompt부터 다시 계산하는 recomputation을 사용할 수 있다. 또는 KV block을 CPU memory로 옮겼다가 다시 GPU로 가져오는 swapping도 가능하다. 이 둘은 trade-off가 다르다. Recomputation은 GPU compute를 더 쓰 지만 CPU-GPU transfer를 피하고, swapping은 계산을 아끼지만 memory transfer overhead가 생긴다. 논문을 읽으면 vLLM의 scheduler가 단순한 queue manager가 아니라, GPU memory pressure를 다루는 policy layer라는 점이 보인다.
Distributed Execution
논문은 distributed execution도 별도 절에서 다룬다. vLLM은 여러 GPU worker를 두고, 각 worker가 model shard와 cache engine을 갖는다. Centralized scheduler는 요청 상태와 block allocation을 조정하고, worker는 할당된 block table을 사용해 실제 model execution을 수행한다. 이 구조는 tensor parallelism이 들어간 큰 모델에서도 같은 memory abstraction을 유지하려는 설계로 볼 수 있다. 논문이 OPT-175B 같은 모델까지 평가할 수 있는 이유도 이 distributed runtime이 전제되기 때문이다.
이 부분은 Method를 읽을 때 놓치면 안 되는 확장 지점이다. 논문은 LLM decoding의 복잡성을 단순히 추가 비용으로만 보지 않고, KV cache sharing의 기회로도 본다. Parallel sampling이나 beam search는 메모리를 더 많이 쓰지만, prefix sharing을 활용하면 메모리를 크게 아낄 수 있는 구조이기도 하다. 그래서 Method의 핵심은 PagedAttention kernel 하나가 아니라, block table, KV cache manager, scheduler, copy-on-write가 함께 만드는 serving abstraction이라고 보는 편이 더 정확하다.
Evaluation
실험은 OPT 13B, 66B, 175B와 LLaMA-13B를 대상으로 하고, workload는 ShareGPT와 Alpaca를 사용한다. 비교 대상은 FasterTransformer와 Orca 계열 baseline이다. 논문은 단순히 한 요청의 latency만 보는 것이 아니라, request rate를 높여가며 normalized latency가 어느 지점에서 급격히 나빠지는지를 본다. 여기서 좋은 serving system은 더 높은 request rate에서도 latency를 낮게 유지할 수 있어야 한다.
Figure 12의 곡선을 보면 vLLM은 baseline보다 더 높은 request rate까지 latency를 유지한다. 논문 수치로는 vLLM이 Orca(Oracle)보다 1.7배에서 2.7배 높은 request rate를 유지하고, Orca(Max)보다 2.7배에서 8배, FasterTransformer보다 최대 22배 높은 처리량을 보인다. Orca(Oracle)은 실제 output length를 미리 안다고 가정한 이상적인 baseline이므로, 이 비교는 vLLM의 이득이 단순한 output length 예측 문제가 아니라 KV cache 관리에서 나온다는 점을 보여준다.
그 이유는 Figure 13에서 더 직관적으로 보인다. vLLM은 같은 GPU 메모리에서 더 많은 요청을 batch에 넣을 수 있다. ShareGPT에서는 Orca(Oracle)이 평균 13.62개 요청을 batch에 넣는 반면 vLLM은 30.42개를 넣고, Alpaca에서는 72.75개 대비 132.44개까지 올라간다. 결국 throughput 차이는 kernel 하나의 미세한 최적화라기보다, memory waste를 줄여 batch size를 키운 결과에 가깝다.
Parallel sampling과 beam search 결과도 같은 방향을 보여준다. Copy-on-write를 사용하면 prefix KV cache를 공유할 수 있고, parallel sampling에서는 6.1%에서 30.5%, beam search에서는 37.6%에서 66.3% 정도의 memory saving을 얻는다. Beam search에서 saving이 더 큰 이유는 후보 sequence들이 긴 prefix를 공유하기 쉽기 때문이다.
Shared prefix 실험은 여러 요청이 같은 prefix를 공유할 때의 효과를 본다. 논문에서는 WMT16 English-to-German translation workload에서 instruction과 번역 예시를 공통 prefix로 두고 평가한다. one-shot prefix에서는 1.67배, 5-shot prefix에서는 3.58배 높은 throughput을 보고한다.
Chatbot workload에서도 vLLM은 Orca baseline보다 약 2배 높은 request rate를 유지한다. 이 실험은 ShareGPT로 합성한 chat history와 user query를 prompt로 사용한다. Multi-turn conversation에서는 요청 간 prefix와 cache lifecycle이 더 복잡해지므로, 단순한 one-shot generation보다 memory manager의 설계가 더 중요해진다.
Ablation Studies
논문은 ablation study에서 PagedAttention의 비용도 따로 확인한다. Block table을 따라 non-contiguous KV cache를 읽어야 하므로 attention kernel에는 추가 overhead가 생긴다. Figure 18a에서 PagedAttention kernel latency는 FasterTransformer의 최적화된 attention kernel보다 20-26% 높게 나타난다. 다만 이 overhead는 attention operator에만 영향을 주고, end-to-end serving에서는 더 큰 batch를 만들 수 있는 이득이 전체 성능을 더 크게 좌우한다.
Block size도 중요한 설계 변수다. Block size가 너무 작으면 GPU parallelism을 충분히 활용하기 어렵고, 너무 크면 internal fragmentation이 커지고 sharing 가능성이 줄어든다. 논문은 ShareGPT와 Alpaca trace에서 block size를 비교한 뒤, 대부분의 workload에서 block size 16이 GPU 활용과 fragmentation 사이의 practical trade-off라고 정리한다.
Preemption recovery에서는 recomputation과 swapping을 비교한다. 작은 block size에서는 swapping이 많은 작은 CPU-GPU transfer를 만들기 때문에 overhead가 커지고, recomputation은 block size와 무관하게 비교적 일정한 비용을 보인다. 논문은 중간 크기 block에서는 두 방법의 end-to-end 성능이 비슷하지만, 작은 block에서는 recomputation이 더 효율적이라고 보고한다.
Discussion and Conclusion
PagedAttention은 모든 GPU workload에 일반화되는 기법으로 제안된 것은 아니다. 이 기법이 LLM serving에 잘 맞는 이유는 output length를 미리 알 수 없어 dynamic allocation이 필요하고, KV cache가 batch size를 제한할 만큼 memory-bound 병목을 만들기 때문이다. 반대로 DNN training처럼 tensor shape이 비교적 static한 경우에는 allocation을 미리 최적화할 수 있고, non-LLM serving처럼 compute-bound인 경우에는 memory efficiency가 곧바로 성능 향상으로 이어지지 않을 수 있다.
vLLM은 OS의 virtual memory와 paging을 그대로 옮긴 것이 아니라, LLM serving의 의미를 이용해 다시 해석한다. 예를 들어 request를 처리하려면 해당 token state가 모두 GPU에 있어야 하므로 all-or-nothing swap-out policy를 쓸 수 있고, 선점된 KV cache는 필요하면 recomputation으로 복구할 수 있다. OS에서는 일반적으로 성립하지 않는 선택이지만, LLM serving에서는 prompt와 generated token sequence가 남아 있기 때문에 가능하다.
논문은 PagedAttention을 attention keys and values를 non-contiguous paged memory에 저장할 수 있게 하는 attention algorithm으로 정리한다. 그 위에 만든 vLLM은 KV cache를 page/block 단위로 관리하고, copy-on-write를 통해 parallel sampling, beam search, shared prefix 같은 decoding scenario에서 cache sharing을 활용한다. 실험 결과는 vLLM이 기존 state-of-the-art system 대비 2-4배 throughput improvement를 얻는다는 주장으로 마무리된다.