Summary
- 본 논문은 operating system에서 쓰이는 대표적인 개념인 paging, virtual memory을 이용한 paged attention을 제안하고 이를 통해 LLM inference의 성능을 향상시키고자 한다.
- LLM은 기본적으로 transformer의 구조이고 transformer은 self-attention 연산으로 구성된다. autoregressive 구조를 가지는 LLM은 이러한 self-attention에서 비효율적인 메모리 관리 때문에 추론 성능이 떨어질 수 있다. paged attention은 이 병목현상을 해결하고자 개발되었다.
- 기본적인 아이디어는 self attention에서 사용되는 key, value matrix을 block 형태로 non-contiguous하게 physical memory에 저장하는 것이다. 이를 통해 LLM에서 흔히 발생하는 internal / external fragmentation을 없애고 최적으로 메모리 공간을 사용하게 된다.
- kv cache을 block으로 저장함으로써 parallel sampling, beam search, shared prefix에서도 다른 추론 기법에 비해 우수한 성능을 보인다.
- paged attention을 코딩으로 구현한 라이브러리는 vllm으로 현재 LLM 추론 분야에서 가장 핫한 github repository 중 하나이다.
Motivation
- LLM의 파라미터수가 점점 많아짐에 따라서 높은 처리량을 보이는 추론 성능 (high throughput)에 대한 연구가 활발하게 진행되었다.
- LLM은 기본적으로 autoregressive generation을 수행하고 이 작업은 memory-bound하며 따라서 메모리를 얼마나 효율적으로 사용하느냐에 따라 LLM throughput 성능이 결정된다.
- batch size을 늘린다면 throughput을 늘릴 수 있겠지만, 기본적으로 사용 가능한 메모리가 적으면 batch size가 작게 나올 수밖에 없다. 따라서 효율적인 메모리 관리가 필요하다.
- 기존의 deep learning framework은 self attention에서 필요한 kv cache을 contiguous memory space에 저장한다. 근데 LLM에서 self attention은 input, output tensor의 길이가 고정되어 있지 않다. 사용자가 어떤 prompt을 모델에 질의할지, 그리고 그 prompt에 대해 어떤 대답이 나올지, 그 길이를 미리 알지 못하기 때문이다. 이러한 특징을 가지는 LLM 서빙 시스템에 기존의 deep learning framework에서 사용하는 contiguous memory space을 그대로 차용하면 아래와 같은 문제점이 발생한다.
- external fragmentation이 발생한다. 유저가 질의를 할 때 maximum tokens 수가 정해져 있지 않다. 서로 다른 메모리 공간을 request 별로 따로 할당해야하기 때문에, request 사이에 남는 공간이 메모리 상에 반드시 발생한다.
- internal fragmentation이 발생한다. 유저가 max tokens = 2048로 질의를 했을 때, 2048개를 모두 채울 수 있지만 그렇지 않을 수도 있다. LLM 입장에서는 EOS token이 나오기 전까지는 계속 token을 생산해야하기 때문에 우선 2048개의 토큰에 메모리 공간을 할당해야한다. 만약, 2048개의 토큰이 만들어지지 않는다면, 할당된 메모리 내에서 빈공간이 반드시 발생한다.
- memory sharing가 불가능하다. LLM 서빙 시스템에서는 parallel sampling, beam search와 같이 하나의 요청에 대해 여러개의 결과를 처리해야할 수도 있고 이때 kv cache을 공유해야한다. 그러나 기존의 시스템으로는 이러한 sharing을 하기가 힘들다.
- 이러한 motivation에 근거하여, 본 논문에서는 paged attention을 제안한다. paged attention은 kv cache을 blocks으로 만들고 이를 non contiguous space에 저장한다.
- blocks을 pages, tokens을 bytes, requests을 processes라고 생각하면 편하다. 인상적인 점은, classical paging, virtual memory에 근거하여 제안하는 방법이기 때문에 논문 중간중간에 os 개념들을 그대로 LLM 서빙 시스템에 적용한다는 점이다.
Approach
Background: transformer-based large language models
language modeling은 $(x_1, \cdots, x_n)$의 확률을 구하는 것이다. 이때, language의 autoregressive한 특성 때문에, 아래와 같이 joint probability을 conditional probability로 decomposition한다.
language modeling에서 self attention을 사용하는 transformer은 뛰어난 성능 때문에 하나의 표준이 되었다. $x_i$ 토큰에 대해 linear transformation을 통해 key, query, value 벡터를 구하고, self attention을 이용해서 output 벡터를 얻는 것이 핵심이다.
$i$번째 토큰에 대해 self-attentive vector인 $o_i$를 구하기 위해 이전 토큰들의 attention score을 구하고 이를 가중치로 value vector을 선형결합한다. 이 연산을 잘보면, $o_i$를 구하기 위해서 그 이전 토큰들의 key, value vector가 필요함을 알 수 있다. 즉, kv cache가 필요하다는 것이다.
LLM Service & Autoregressive generation
유저의 프롬프트인 $(x_1, \cdots, x_n)$가 주어졌을 때 이에 대한 generation computation은 크게 두가지 phase로 나뉜다.
- prompt phase
- $(x_1, \cdots, x_n)$이 주어졌을 때, 첫번째 토큰을 생성하는 단계이다: $P(x_{n+1} | x_1, \cdots, x_n)$
- 잘 생각해보면, $(x_1, \cdots, x_n)$은 이미 주어져있다. 따라서, 각 토큰의 kv cache을 사용해서 attention을 계산하면 된다.
- 이때는 행렬 계산을 병렬화하여 진행하면 되므로, 딱히 병목현상은 없다.
- autoregressive generation phase
- $x_{n+t}$ 토큰이 주어졌을 때 다음 토큰을 생성하는 문제이다: $P(x_{n+t+1} | x_1, \cdots, x_{n+t})$
- 이때, key, value 벡터는 $k_1, \cdots, k_{n+t}, v_1, \cdots, v_{n+t}$이다.
- max tokens 수에 다다르거나 [EOS] 토큰이 나와야만 generation이 종료된다.
- 이때는 autoregressive하게 진행이 되기 때문에 병렬화할 수 없다. 따라서 이 부분이 GPU 연산의 효율성을 떨어뜨리며 작업을 memory bound로 만든다. LLM 추론에서 병목현상의 주요 원인은 이 부분이다.
Batching techniques for LLMs
LLM 서빙 시스템은 여러개의 요청을 하나의 배치로 묶음으로써 성능 향상을 꾀할 수 있지만 어려운 점이 두가지 있다.
- 여러개 요청이 동시에오는게 아니라 각기 다른 시간대에 온다. 나이브한 접근법으로는 batch size가 찰 때까지 기다리는 것인데, 이는 분명한 딜레이를 발생시킨다.
- 다음으로 요청은 각기 다른 input, output 길이를 가진다. 그러면 가장 큰 길이를 가지는 요청에 맞춰서 padding을 하면 되지 않을까 싶은데.. 가뜩이나 연산하는데 메모리가 많이 드는데 이런 곳에까지 메모리를 쓴다면 더욱 성능이 떨어질 것이다.
논문에서는 직접적으로 나오지 않지만, 위와 같은 batch system을 contiuous batching이라고 한다. 하나의 배치 안에 여러개의 동일한 input을 넣고 그 안의 모든 작업이 끝날 때까지 기다리는 것이다. 하지만, LLM의 input, output을 생각해보면 이런 방법은 비효율적이다. 이에 대한 대안으로 non contiuous batching이 있고 이게 논문에서 나오는 cellular batching, iteration-level scheduling인 것 같다. non continuous batching은 하나의 배치 안에 작업이 종료된 요청이 있으면, 이를 새로운 작업으로 대체한다. [EOS]를 먼저 뱉는 요청은 배치 안에서 나가고 새로운 요청이 들어오는 것이다.
Memory challenges in LLM serving
non continuous batching은 메모리 낭비를 최소하하지만, 여전히 요청을 처리하는데 메모리를 많이 사용하기는 한다. self attention을 계산하는데 kv cache을 저장하고 사용해야하기 때문이다. 이 kv cache 저장 공간 때문에 다른 연산을 할 수 있는 메모리 공간이 줄어들고 그에 따라서 throughput도 줄어든다. 이러한 메모리 이슈는 크게 세가지 어려운 점이 있다.
- Large KV cache
- 생성하는 토큰의 길이가 길어질수록 kv cache 사이즈 자체가 커진다.
- 예를 들어 13B OPT 모델에 대해, 2048 토큰을 생성하기 위해 필요한 kv cache 공간은 2 (key, value vectors) * 5120 (hidden state size) * 40 (number of layers) * 2 (bytes per FP16) = 1.6GB이다.
- 그림1의 NVIDIA A100 40GB을 생각해보면, 모델 올리고, kv cache도 올려야하는데 여러개의 요청이 들어오면 kv cache도 여러 세트가 필요하므로 많아봤자 10개정도의 요청을 처리할 수 있을 것이다.
- 게다가 NVIDIA에서 나오는 GPU을 생각해보면 A100에서 H100으로 가는데 FLOPS는 2배 더 많아지지만 메모리 용량은 그대로이다. 즉, 고성능 GPU라고 해도 메모리용량이 커지는 것은 아니므로 효과적인 메모리 관리가 필요하다는 것이다.
- Complex decoding algorithms
- 단순하게 하나의 요청에 대해서 하나의 결과만 원하는 것이 아니라 여러개의 결과를 원하거나, 상위 10개의 가장 probable한 결과를 얻는 task도 수행해야한다.
- 이를 각각 multiple random samples, beam search라고 하는데 기존의 kv cache 전략으로는 kv 벡터를 공유할 수 없어 효율적인 계산을 할 수 없다.
- scheduling for unknown input & outputs
- LLM은 input, output의 길이를 미리 알지 못한다는 특징이 있다. 이런 점 때문에 kv cache을 관리하는게 특히 더 중요하다.
- 사용하지 않는 kv cache은 없애는 등의 작업이 필요하다.
Memory management in existing systems
핵심은 LLM의 input, output의 길이를 미리 알 수 없고 정해져있지 않기 때문에, 기존 deep learning framework에서 메모리를 관리하는 방식을 그대로 적용하면 비효율적인 메모리 관리를 하게된다는 점이다.
- request A에 대해 LLM은 output의 길이를 알지 못하기 때문에 우선 max tokens인 2048을 할당한다.
- 그러나 디코딩을 해보니, 실제로 생성되는 단어는 forth, <eos> 뿐이다.
- 여기서 두 토큰이 생성되었는데, 그 두 토큰이 생성되기 이전까지 해당 공간에 낭비가 있다.
- 또한 2048 토큰을 이후에 모두 사용하지 않았으므로 이에 대한 공간 낭비가 있다. (internal fragmentation)
- request B는 max tokens이 512인데, 이렇게 max tokens이 일정하지 않으므로 이 두 요청을 위해 할당한 메모리 사이에서 빈 공간이 발생한다. (external fragmentation)
이런 메모리 낭비를 paged attention은 어떻게 해결할까? 지금부터 자세히 살펴보자.
Method
vLLM의 아키텍쳐이다. centralized scheduler가 kv cache을 kv cache manager을 통해 적절하게 관리한다.
Paged Attention
paged attention은 각 sequencedml kv cache을 kv blocks으로 분할한다. (os에서 프로세스를 paging으로 분할하는 것과 동일) 이때 block의 크기는 고정시킨다. (external fragmentation 방지) $j$ 번째 block에 대해 좀 더 자세히 살펴보자.
- $j$번째 key block: $K_j = (k_{(j-1)B + 1}, \cdots, k_{jB})$
- $j$번째 value block: $V_j = (v_{(j-1)B + 1}, \cdots, v_{jB})$
- $A_{ij} = ( a_{i, (j-1)B+1}, \cdots, a_{i,jB} )$: j번째 KV block의 row vector attention score
논문에서는 위와 같이 노테이션을 정하는데, 필자는 이해가 잘 가지 않아서 구체적인 예시를 통해서 단계별로 생각해보았다.
위와 같이 상황을 만들어두고 4번째 block의 i번째(global token index) token에 대해 attention score을 blockwise하게 계산해보자. 기존의 attention score을 계산하는 과정을 생각해보면, 이전 i-1개 token들의 key vector을 곱하고 각각을 softmax 취해준다. 그리고 normalization 해준다. 우선 여기까지 blockwise하게 변환하여 생각해보자.
위와 같이 i번째 토큰이 속하기 직전 block까지, blockwise하게 key, query 벡터를 곱한다. 이제 첫번째 block에 대해 attention score을 구하는 과정을 살펴보자.
$A_{i1}$은 i번째 토큰과 첫번째 block과의 attention score로써, block이 8차원 벡터이므로 이 attention score도 8차원 벡터이다.
이제 마지막으로 attention score을 가중치로 하여 value vector와 선형결합하는 부분을 살펴보자.
블록별로 계산된 attention score을 계수로, 블록별 value vector에 붙여서 최종 i번째 토큰에 대한 attention vector을 만든다.
여기서 아래의 의문이 들었다.
- full self-attention은 i번째 토큰의 attention score을 계산할 때, 이전 i-1까지의 모든 토큰의 key,value vector을 사용한다.
- 근데 paged attention은 블록별로 계산하여, 이전 블록까지 토큰의 key,value 벡터는 사용하는데, 같은 블록 내의 i번째 토큰 이전까지의 토큰은 고려하지 않는 것 같다.
- 예를 들어 위 예시에서 28번째 토큰에 대한 attention score을 계산하고자 하면, 네번째 블록에 속하고 세번째 블록인 24번째 토큰까지의 key,value vector은 모두 포함한다. 하지만 네번째 블록의 25번째 토큰부터 27번째 토큰까지는 포함하지 않는 것 같다.
- 그러면 paged attention은 exact self attention이 아니라는 것인가? 이 부분은 좀 더 확인이 필요해 보인다.
그냥 계산할 수도 있는 attention을 굳이 이렇게 block으로 나눠서 계산하는 이유는.. 다시 한번 강조하면 아래와 같다.
- kv cache을 block으로 나누면 internal, external fragmentation을 방지할 수 있다.
- block 끼리 kv cache을 공유할 수 있어서 multiple sampling, beam search 등을 효율적으로 할 수 있다.
KV Cache manager
- vllm의 memory manager은 os의 virtual memory와 유사하다. virtual memory가 logical memory에서는 contiguous 하게, physical memory에서는 non-contiguous하게 pages을 저장하는 것과 동일한 방식으로, vllm의 memory manager도 kv cache blocks을 저장한다.
- 이를 통해, max tokens만큼 메모리를 미리 차지할 필요도 없다.
- 잘 생각해보면, 마지막 block을 제외하고는 모든 block이 꽉 차있다. 마지막 block은 다음 토큰의 생성을 위해 공간을 남겨둔다. 이는 virtual memory에서 paging을 하는 것과 동일한 방식이다.
- kv cache manager는 또한 block table을 관리하여 logical kv blocks와 physical kv blocks을 연결한다.
Decoding with paged attention and vLLM
vllm에서 decoding하는 과정을 살펴보자.
- 우선 block별로 kv cache을 관리하기 때문에 maximum token만큼 메모리를 따로 할당할 필요가 없다.
- 위 예시에서 프롬프트의 길이가 7이고 block size는 4이기 때문에 2개의 logical kv blocks을 physical kv blocks는 block 1,7로 배분한다.
- 이때 7개의 토큰을 가지는 프롬프트 바로 다음에 나올 output을 계산하기 위해 kv cache을 계산하고 이를 block 1,7에 저장한다.
- 다음에 만들어지는 단어인 fathers은 block 1에 빈공간이 있으므로 관련 kv cache을 block1에 저장한다.
- 다음으로 나오는 brought 단어는 더 이상 여유 공간이 있는 block이 없으므로 새로운 block 3에 kv cache을 저장한다.
- 이와 같이 vllm은 그 이전의 모든 block이 꽉 차있지 않는 이상, 새로운 physical kv blocks을 내어주지 않음으로써 거의 zero에 가까운 memory waste을 보여준다.
Application to other decoding scenarios
Parallel sampling
LLM 서빙 시스템에서는 동일한 프롬프트에 대해 여러개의 결과를 보여줌으로써 유저가 마음에 드는 대답을 선택할 기회를 준다. 이러한 상황에서는 동일한 pprompt을 공유하게되는데 따라서 prompt의 kv cache도 공유할 수 있다. paged attention은 이런 공유를 더 쉽게 해주는 장점이 있다.
- 처음 Four score and seven years ago our 이라는 프롬프트가 들어와쓸 때 block 1,7에 kv cache가 저장되었고 프롬프트가 동일하기 때문에 추가적인 저장이 필요하지 않다. 대신, reference count을 2로 둬서, 이 block을 참고하고 있는 request가 두개임을 명시한다.
- 이후 토큰을 생성할 때는 서로 다른 토큰이 생성될 수 있다. 여기서는 mothers, fathers 토큰이 생성되었다.
- 이 경우에는 block 1을 copy 하고 서로 다른 token을 채워 넣는다. 이는 마치 os virtual memory에서 copy-on-write 기술과 유사하다.
- block 1에 대한 reference count가 2이기 때문에 sample A1에 대한 kv cache는 block 3에 copy하고 reference count을 1로 줄인다. sample A2에 대한 kv cache는 block 1의 reference count가 1이기 때문에 그대로 block 1을 사용한다.
- 이와 같이 physical block을 여러 샘플에 걸쳐서 공유함으로써 memory usage을 꽤나 크게 줄일 수 있다.
- 이는 특히 long input인 경우에 더 효율적이다. prompt가 길수록 차지하는 block의 개수가 더 많아질 것이기 때문이다.
Beam search
- 기계 번역과 같은 LLM tasks에서는 LLM에 의해서 만들어진 top k 결과를 유저들이 기대한다.
- beam search는 이러한 상황에서 fully 찾아보는 계산을 줄이기 위해 널리 사용된다.
- beam search는 decoding 시에 beam width parameter k에 따라서 k개까지 후보 토큰을 저장한다.
- parallel sampling은 프롬프트 토큰만 서로 공유한다면 beam search는 initial prompt block뿐만 아니라 decoding 과정의 다른 block도 공유한다. 이는 마치, OS에서 compound forks에 의해 만들어진 process tree와 유사하다.
- 위 그림에서 decoding을 하다가 beam candidate 0,3는 더 이상 top candidate가 아니라서 free된다. 즉, physical blocks의 reference count가 줄어드는 것이다.
- vllm에서는 physical memory sharing을 통해 beam search에서 상당 부분을 서로 공유할 수 있기 때문에 메로리를 상당부분 아낄 수 있다.
Shared prefix
- 보통 LLM에 질의를 할 때, system prompt을 동일하게 넣는다. 이는 모든 프롬프트에 동일하게 적용되는 것이므로 kv cache을 통해 sytstem prompt의 kv cache을 공유함으로써 메모리 사용을 줄일 수 있다.
- 이는 마치 OS에서 프로레스간에 shared library을 다루는 것과 유사하다.
Mixed decoding methods
- 앞서 하나의 결과를 내는 decoding, parallel sampling, beam search, shared prefix 등 다양한 decoding 방법을 살펴보았다.
- vllm은 이렇게 다양한 decoding 방법을 동시에 병렬적으로 처리할 수 있다.
- paged attention은 logical memory와 physical memory을 연결만 하기 때문에 다른 LLM 서빙 시스템에 비해 쉽다고 한다.
Scheduling and preemption
vllm은 요청 수가 많으면 first-come-first-serve (FCFS) 전략을 통해 요청을 처리한다. 이렇게 많은 요청을 처리하다보면 GPU의 physical blocks의 메모리가 남질 않을 수도 있는데, 새로운 kv cache을 받기 위해서 기존의 physical blocks을 비워야하는 상황이 생긴다. 이때, 어떤 block부터 삭제를 할지, 만약에 삭제된 block이 필요하다면 다시 어떻게 복구할지에 대해서 명확한 대안을 제시해야한다.
- Which blocks should it evict?
- 한 sequence의 kv cache blocks은 한꺼번에 접근되기 때문에 all-or-nothing eviction policy을 사용한다. (하나의 sequence에 대해 attention이 끝났다면 이 sequence의 kv cache는 더 이상 필요하지 않기 때문에)
- 또한 parallel sampling, beam search와 같이 multiple sequence가 kv cache을 공유하는 경우에는 sequence 그룹으로 gang-schedule 된다. (모든 sequence의 decoding이 끝나기 전까지는 evict되지 않는다는 뜻 같음)
- How to recover evicted blocks if needed again?
- swapping. virtual memory에서 사용되는 classic 기술이다. evicted pages을 swap space에 복사한다.
- vllm은 evicted blocks을 cpu memory에 복사한다.
- recomputation. 단순하게 preempted sequences가 reschedule된다면 kv cache을 다시 계산한다.
- decoding에서 만들어지는 토큰이 original user prompt에 새로운 prompt로써 붙여지므로 latency가 길지 않다.
Mapping between classic paging&virtual memory and paged attention
| classic paging&virtual memory | vllm |
| pages | blocks |
| bytes | tokens |
| virtual memory | logical blocks |
| physical memory | physiccal blocks |
| process | requests |
| copy-on-write pages | copy-on-write blocks |
| compound forks에 의해 만들어진 process tree | beam search의 blocks sharing |
Experiments
Experimental setup
- model
- OPT 13B, 66B, 175B
- LLAMA 13B
- workloads
- sharedGPT
- Alpaca
- Baseline
- FasterTransformer
- Orca
- Key metrics
- normalized latency: end-to-end latency / output length
- 하나의 토큰을 생성하는데 걸리는 시간이다.
Basic sampling
- 하나의 요청에 대해 하나의 sample을 만드는 basic sampling에 대해 위는 sharedGPT, 아래는 Alpaca 데이터 세트에 대한 결과이다.
- request rate가 일정 수준 이상으로 되면 noramlized latency가 급등하는데, 이는 LLM 서빙 시스템의 서빙 한계치를 넘어서는 요청이 들어올 때, 한계치 이상은 모두 queueing되기 때문에 발생하는 현상이다.
- 모든 모델에 대해서, normalized latency가 동일하다면 vllm의 request rate가 높다. 즉, vllm이 다른 모델에 비해 요청을 더 많이 처리한다는 것이다.
Parallel sampling and beam search
- kv cache sharing이 특히 필요한 parallel sampling, beam search에 대한 결과이다.
- vllm이 normalized latency가 동일하다면 월등하게 많은 요청을 처리할 수 있음을 알 수 있다.
Shared prefix
- 1 shot prefix prompt와 5-shot prefix prompt가 주어졌을 때, vllm과 Orca의 성능을 비교해보았다.
- vllm의 처리량이 Orca에 비해 훨씬 많음을 알 수 있다.
Impact of block size
- blocks의 사이즈는 vllm에 큰 영향을 줄 수 있다. 크기가 너무 작으면 block별로 처리하는 연산이 많은 vllm에서 GPU parallelism을 제대로 활용하지 못할 수 있고 너무 크다면 internal fragmentation이 증가할 것이다.
- sharedGPT, Alpaca 데이터에 실험한 결과 16 정도가 가장 적절한 것으로 나타났다.
- 따라서 vllm은 기본 block size을 16으로 선택한다.
Conclusion
- 최근에 가장 핫한 LLM inference optimization 기법 중 하나인 vllm의 페이퍼가 드디어! 나와서 리뷰를 해보았다.
- paged attention은 생각보다 아이디어가 간단하면서도, paging과 virtual memory에 기반하여 만들어졌기 때문에 그 이론적 배경이 탄탄하다고 느껴졌다.
- 그러나 논문을 끝까지 읽으며.. paged attention은 결국 exact attention은 아니라는 생각이 들었다. 앞서 paged attention을 자세하게 따라가보니, i번째 attention score을 계산하는데 i-1번째까지의 모든 토큰을 사용하지 않기 때문이다. (이해한 바가 틀릴 수도 있으니, 댓글 지적 환영)
- 간단한 생각으로 메모리를 효율적으로 사용할 수 있다는 점, 그리고 open source로써 활발하게 개발되고 있다는 점은 vllm을 한층 더 매력적으로 만들어주는 포인트가 아닐까 싶다.
댓글