[SOSP 2023] Efficient Memory Management for Large Language Model Serving with PagedAttention

안녕하세요. 이번에 들고온 논문은 LLM 을 서빙할때 KV cache 메모리 관리의 비효율을 어떻게 해결할 것인지를 다룬 논문입니다. 저희 연구실 사람들이 이런 메모리 관리 측면의 OS 연구를 진행하지는 않지만 해당 논문의 방법론을 사용하면 LLM을 사용할때 여러 요청을 동시에 처리해야하는 상황에서의 병목을 줄일 수 있어서 가져왔습니다. 좀 이전 논문이라 최근 방법론들이 (특정상황에서) 더 빠른 속도를 내는 것들이 있는거로 확인되는데, 논문 1저자가 한국인이기도 해서 들고왔습니다.

https://github.com/vllm-project/vllm

Abstract

저자들은 먼저 대형 언어모델들을 높은 처리량으로 서빙하려면 많은 요청을 한번에 배치해야한다고 합니다. 기존 시스템들은 각 요청마다 필요한 KV cache가 매우 컸습니다. ( 요청마다 생성해야할 문장의 길이를 알 수 없기에 가장 큰 KV cache를 미리 할당하는 방식이기 때문) 이러한 비효율적인 메모리 방식에는 배치 사이즈도 제한적이고 처리량도 낮아질 수 밖에 없습니다.

이 문제를 해결하기 위해서 저자는 PagedAttention 이라는 새로운 어텐션 방식을 제안합니다. 이 방식은 운영체제의 가상 메모리와 paging 아이디어에서 영감을 받았고, KV cache를 연속된 메모리 덩어리로 보관하는 대신 고정 크기 block 단위로 나누고, 이를 비연속적인 물리 메모리에 저장할 수 있도록 합니다. (저희가 흔히 아는 attention을 새로운 방식으로 제안했다기보다 Attention 매커니즘 시 생성되는 KV캐시를 시스템적으로 관리하는 기법을 칭하는 것입니다.) 운영 체제에서 페이지는 가상 메모리를 관리하기 위해 사용하는 고정된 크기의 논리적 메모리 블록을 의미합니다.

저자가 주장하는 핵심 장점은 두 가지 입니다.

1. KV cache memory waste 를 0에 가깝게 줄임
2. KV cache를 요청 간 / 시퀀스 간 유연하게 공유 가능

실험 결과로는 기존 SOTA 서빙 시스템들 대비 같은 latency 수준에서 2배~4배 정도의 처리량 향상을 보였다고 합니다. 특히 긴 시퀀스, 큰 모델, 복잡한 디코딩 알고리즘을 사용할수록 개선 폭이 커졌다고 합니다. 즉 LLM의 서빙 병목이 연산만의 문제가 아니라 메모리 관리의 문제이고 그걸 OS에서 사용하던 paging 으로 풀면 처리량을 크게 높일 수 있다고 이해하면 됩니다.

Introduction

저자들은 LLM 서비스의 현실적인 문제를 언급하는데, GPT 나 PaLM 같은 대형 언어모델이 코딩 어시스턴트, 챗봇 같은 다양한 응용을 가능하게 했지만, 실제로 이런 서비스를 운영하는 비용이 매우 큽니다. 논문에서는 LLM 요청 처리 비용이 전통적인 쿼리 검색보다 비쌀 수 있다고 언급하면서 결국에는 serving 처리량을 높여 요청당 비용을 낮추는 것이 중요하다고 합니다.

문제는 LLM 의 autoregressive 한 생성 방식 특성상 토큰을 하나씩 순차적으로 생성하기에, 어떤 질문이 들어왔을때 실제로 생성될 문장의 전체 길이를 알 수 없어, 모델이 최대 2048 토큰을 내뱉는다면 실제로 10토큰이 생성되더라도 일단 2048 토큰 분량의 연속된 VRAM 공간을 할당합니다. 이러한 Static Allocation 방식은 Batch size가 커졌을때 연속된 빈 공간이 존재하지 않아 OOM(Out of Memory) 가 뜰 확률이 늘어나게 됩니다.

논문 figure 1에서 설명하는 부분을 보면 우선 13B 모델을 A100 40GB GPU 에 올릴 때 26GB는 모델의 파라미터가 차지하게 되고, 그 다음 상당한 비중을 KV cache가 차지하여 실제 모델을 올릴 때 병목이 모델의 weight보다는 요청 별 KV cahce를 어떻게 저장하고 관리하는지에 더 가깝다는 것입니다.

우측 위의 graph 를 보면 기존 방식에서 contiguous 한 방식으로 momory를 저장하기 때문에 두 가지 문제점이 생긴다고 합니다.

1. Fragmentation
2. Memory Sharing 불가

먼저 fragmentation 을 보면, 요청마다 최대 길이를 기준으로 미리 큰 청크를 잡아두다 보니 실제로는 사용되지 않는 공간이 많아집니다. 그리고 요청 길이가 제각각이기 때문에 external fragmentation도 커진다고 합니다.

여기서 external fragmentation을 예시로 설명을 드리겠습니다.

1. 요청 A,B,C가 들어와서 메로리를 순서대로 점유
2. 요청 B가 생성을 먼저 마치고 완료
3. 그 자리에 빈 구멍이 생김
4. 새로운 요청 D가 들어왔는데 B가 나간 구멍보다 조금 더 큰 경우
5. 메모리 뒤쪽에 자리가 있더라도 연속된 공간에만 할당이 되므로 메모리가 텅텅 비어 보이는데 OOM 이 뜨거나 새로운 요청을 거절하게 되는 현상

논문에서는 기존 시스템에서 실제 token state 저장에 쓰이는 메모리 비율이 20.4%~38.2% 수준에 불과하다고 분석합니다.

또 다른 문제였던 Sharing은 parallel sampling 이나 beam search 처럼 하나의 프롬프트에서 여러 ouput sequence 를 생성할 때는 프롬프트 부분의 KV cache 를 공유할 수 있어야 메모리 낭비를 절약할 수 있는데, 기존 방식은 요청별로 KV cahce를 따로 연속적으로 들고 있어 이 공유가 어렵습니다.

그래서 저자들은 운영체제의 virtual memory처럼 KV cache를 page처럼 block 단위로 나누고, 논리적 block 과 물리적 block을 분리하는 PagedAttention 을 제안합니다.

이 설계를 통해 작은 block 단위의 allocation이 가능해지고, 위에서 언급한 external fragmentation도 제거되며 block 단위의 sharing이 가능해진다고 합니다.

해당 figure는 위에서 언급한 실제 사용되는 메모리 비율입니다.

Background

논문은 background에서 먼저 Transformer 기반 LLM 과 serving 과정을 간단히 다시 짚습니다. 트랜스포머가 self-attention을 통해서 이전 토큰들의 key/value를 참고하면서 다음 토큰을 생성하는데, serving 시에는 이전 토큰들의 key/value를 매번 다시 계산하지 않기 위해 KV cache를 저장해둡니다. 이 KV cache는 generation이 진행될수록 계속 커집니다. 저자들은 serving 단계를 크게 두 단계로 나눕니다.

1. prompt phase
   1. 입력 prompt 전체가 주어져 있어서 병렬 계산이 가능
2. autoregressive generation phase
   1. 토큰을 하나씩 순차적으로 생성해야 해서 비효율적이고 momory-bound가 되기 쉬움

Method

PagedAttention

기존 attention은 KV cache가 연속 메모리에 있다고 가정하는데, PagedAttention은 비연속 메모리에 저장된 KV block들을 읽어 attention을 계산할 수 있도록 설계됩니다. 즉 KV cache를 token 단위로 길게 이어진 배열처럼 보는 것이 아니라, 고정 길이 B의 block들로 쪼개서 관리합니다.

다시 말해 logical block 과 physical block을 분리하여 사용하고 논리적으로는 시퀀스가 연속된 블록들로 이루어져 있지만, 실제 GPU 메모리에서는 이 block들이 여기저기 흩어져 있어도 됩니다. attention kernel이 block table을 보고 필요한 physical blcok을 찾아가며 연산합니다.

KV cache manager

vLLM은 운영체제의 virtual memory와 유사한 방식으로 KV cache를 관리합니다.

각 요청의 KV cache는 logical KV block들의 연속으로 표현되고, 실제 GPu DRAM 에는 physical KV block 들이 따로 존재합니다. block table로 둘을 매핑하고 block table entry에는 해당 logical block이 어느 physical block에 매핑되는지, 그리고 그 안에 몇 칸이 채워져 있는지를 기록합니다. 이렇게 되면 미리 최대 길이만큼 할당할 필요가 없어지고 생성이 진행되면서 필요한 block 만 추가로 할당해주면 됩니다. 요청이 끝나면 block을 즉시 반환해서 사용하여 기존 방식의 reservation / fragmentation 문제를 거의 제거할 수 있습니다.

해당 figure는 beam search에서의 저자의 방법론 그림으로 우선 beam search는 위에서의 greedy 방식이 아닌 상위 K 개를 들고가면서 정답을 찾아가는 방식입니다. 기존 시스템에서는 beam candidate 간 KV cache를 복사해서 사용했지만, 저자의 방법론은 갈라지기 전까지는 공유하는거로 이해하면 될 것 같습니다.

Evaluation

실험에서는 OPT 13B,66B,175B 와 LLaMA 13B를 사용했고, A 100 GPU 환경에서 평가합니다. 데이터셋은 실제 LLM 서비스 대화 길이 분포를 반영하기 위해서 ShareGPT와 Alpaca를 기반으로 synthetic workload를 구성했다고 합니다. ShareGPT는 입력과 출력이 더 길고 분산이 크다고 하네요.

Baseline은 크게 FasterTransformer와 Orca 입니다.

이 그래프는 서버가 얼마나 많은 요청을 견디면서 빠른 속도를 유지하는지를 보여줍니다.

Request Rate는 초당 요청 수를 의미하여 1초에 사용자가 얼마나 많은 질문을 던지는지를 나타내고 Y축은 토큰당 지연 시간으로 글자 하나를 만들어내는데에 걸리는 시간으로 이해하면 됩니다. 그래프에서 갑자기 치솟는 부분은 처리의 한계점으로 GPU 의 메모리나 연상 능력이 꽉 차면 대기 시간이 생겨버리는 상황이라고 이해하면 됩니다.

Conclusion

Page 라는 운영체제에서 쓰이는 방식을 vLLM이라는 이름으로 serving system 을 구현한 논문입니다. 일반적으로 LLM inference를 quantization이나 pruning 또는 디코딩 방식으로 속도 개선을 시키는거로 알고 있었는데, 계산은 잘하고 있고, 메모리 관리의 문제를 지적하며 시스템 효율을 개선했다는 점과, 실제 서비스화해서 제공하고 있다는 점이 contribution인 것 같습니다.

감사합니다.