Serving throughput을 향상시키기 위해서는 여러 request를 하나의 batch로 묶어 동시에 처리할 수 있어야 합니다.
그러나 이를 위해서는 각 request에 필요한 KV cache memory를 효율적으로 관리하는 것이 필수적입니다.
LLM에서는 batch size가 증가할수록 KV cache memory 사용량이 빠르게 증가하기 때문에 KV cache 관리 방식이 곧 최대 batch size를 결정하며, 이는 전체 throughput에 직접적인 영향을 미칩니다.
문제: KV cache memory를 효율적으로 관리 못함
대부분의 LLM serving 시스템은 각 request에 대해 maximum sequence length (예: 2048 tokens) 만큼의 연속적인 memory space를 pre-allocation 하는 방식을 사용합니다.
그러나 실제 KV cache는 생성 과정에서 동적으로 길이가 증가하기 때문에 다음과 같은 memory waste 문제가 발생합니다.
1. Internal and external memory fragmentation
실제 필요 길이보다 더 큰 공간을 할당받아 Internal fragmentation 발생
request마다 길이가 다르기 때문에 External fragmentation 발생
request의 pre-allocation space는 전체 lifetime 동안 다른 request가 활용할 수 없기 때문에 위와 같은 memory waste가 발생합니다.
2. Memory Sharing 불가
하나의 request에서 multiple output을 생성하는 경우, KV cache를 부분적으로 공유할 수 있음에도 불구하고 기존 방식은 각 output을 서로 다른 contiguous 공간에 저장하여 memory sharing의 이점을 활용하지 못합니다.
제안: PagedAttention
기존 OS의 paging 기반 virtual memory 개념을 적용한 새로운 attention 방식인 PagedAttention을 제안합니다. 각 request의 KV cache를 KB block이라는 fixed-size block 단위로 분할한 후 attention을 계산하는 방식입니다.
이를 통해 위의 문제를 다음과 같이 해결할 수 있습니다.
KV cache를 비연속적인 memory block에 저장 가능
작은 단위로 memory를 할당하여 fragmentation 감소
동일 request 내 여러 sequence 간 KV 공유 가능
서로 다른 request 간에도 block-level memory sharing 가능
기여: vLLM
PagedAttention을 기반으로 한 high-throughput 분산 LLM serving engine인 vLLM을 개발했습니다.
기존의 KV cache memory 낭비를 거의 제거하였으며 단일 GPU memory capacity를 초과하는 대규모 모델도 분산 환경에서 효율적으로 처리할 수 있습니다.
실험 결과, FasterTransformer 및 Orca와 같은 기존 SOTA 시스템 대비 약 2–4배 높은 throughput을 달성했습니다.
Ⅱ Background
LLMs의 generation과 serving 과정을 이미 아시는 분은 skip해도 괜찮습니다!
한 위치의 query vector와 그 이전의 모든 위치 key vectors들을 곱하여 attention score \(a_{ij}\) value vector들 weighted average를 통해 \(o_i\)를 얻습니다.
attention score와 output
Transformer model의 다른 모두 components
즉 embedding layer, feed-forward layer, layer normalization, residual connection, output logit computation, query, key, value transformation 모두 각 위치별로 독립적으로 적용됩니다.
즉
$$ y_i = f(x_i) $$
의 형태로 position-wise하게 적용됩니다.
2.2 LLM Service & Autoregressive Generation
일단 훈련된 LLMs이 service로서 배포되면
LLM service에 input prompt tokens \(\left(x_{1},..., x_{n}\right)\)와 함께 request를 보내면 LLM service는 output tokens \(\left(x_{n+1},..., x_{n+T}\right)\)를 Eq(1)에 따라 생성합니다.
이때 prompt + output token들을 연결한 하나의 긴 문장\(x=\left(x_{1},...,x_{n+T}\right)\)을sequence로서 참조합니다.
Eq(1)에서의 decomposition에서 알 수 있듯, LLM은 새로운 token들을 하나씩 sampling 하고 생성합니다. 그리고 각 새로운 token은 그 sequence에 존재하는 모든 이전 token들에 의존하여 생성하게 됩니다. 이때 이전 token들의 key와 value vectors를 참고합니다.
이러한 순차적인 생성 과정에서 이전의 존재하고 있는 token들의 key와 value vector들을 앞으로 생성할 미래 토큰들을 위해 저장하여(cache)는데, 이를KV cache라고 합니다.
🤔 Note 한 token의 KV cache는 그 token의 모든 이전 token들의 key와 value vector에 대해 의존합니다. 이는 한 sequence에서 같은 token이라도 나타난 위치가 다르다면 KV cache는 다르다는 것을 의미합니다. 즉 앞선 context의 내용이 다르기 때문에 KV cache도 달라지는 것이죠.
Two Phases of generation computation in the LLM service
LLM service에서 generation 과정은 크게 2단계로 나눌 수 있습니다.
(1) The prompt phase (prefill)
사용자 입력 \(x_1,…,x_n\)전체를 한 번에 처리하여 첫 번째 생성 토큰의 확률 \( P(x_{n+1} \mid x_1, \dots, x_n) \)을 계산합니다.
각 token(position)에 대한 key/value vector가 생성됩니다. 이때 연산 과정은 모든 prompt의 token이 주어져 있어 GPU의 병렬성을 효율적으로 활용할 수 있습니다.
(2) The autoregressive generation phase
한 번에 하나씩 순차적으로 token이 생성됩니다.
이전 step에서 계산된 KV cache를 재사용하며, 현재 iteration에서는 직전에 sampling된 token에 대한 key/value vector만 추가로 계산합니다.
즉 iteration t에서 \(x_{n+t}\) token이 생성되었다면, 그 다음 iteration에서는 \(x_{n+t}\)에 대한 새로운 key/value vector \(k_{n+t}\), (v_{n+t}\) 만 계산됩니다.
2.3 Batching Techniques for LLMs
LLMs을 서빙하는데 compute utilization은 muliple request를 한 batch로 묶음으로서 개선될 수 있습니다. 왜냐하면 여러 request들은 같은 model weights을 공유하기 때문에, weights을 이동하는 overhead는 한 batch 내에서 여러 request들 간에 분산되고, batch 크기가 충분히 클 때 computational overhead에 비해 상대적으로 작아질 수 있습니다. 즉 weights을 memory로 이동하는 비용은 batch 전체로 나눠지게 됩니다.
기존 Batching 방식의 문제점
그러나, LLM service에 request들을 batching 하는 것은 다음과 같은 2가지 이유로 단순하지 않습니다.
1. reuqest들은 다른 시간에 도착할 수 있다.
단순한 batching 전략은 batch를 맞추기 위해 대기 (일찍 도착한 request가 늦게 도착하는 requests를 대기) 하거나 앞선 batch가 끝날 때까지 대기 (앞으로 오게 될 reuqest들을 앞선 request들이 끝날 때까지 지연)하는 방법이 있는데, 이는 상당한queueing delays로 이어지게 됩니다.
2. request들은 input과 output lengths가 크게 달라질 수 있다.
batching technique은 request들의 input과 ouput에 padding을 추가하여 그들의 length를 맞추는데 이는 GPU computation과 memory에 낭비가 됩니다.
iteration-level scheduling 방식
위와 같은 문제를 해결하기 위해, cellular batching 혹은 iteration-level scheduling과 같은 fine-grained batching mechanisms이 제안되었습니다.
request level에서 작업하는 전통적인 methods와 달리, 이러한 techniques는 iteration level에서 운영됩니다. iteration level 방식은 각 iteration 후, 완료된 request들이 batch에서 제거되고 새로운 request들이 추가됩니다. 따라서 전체 batch가 완료될 때까지 기다릴 필요 없이한 single iteration 동안 기다린 후에 새로운 reuqest가 처리될 수 있습니다.
즉, queueing delay와 비효율적인 padding을 줄임으로써, fine-grained batching mechanism들은 LLM serving의 throughput을 크게 향상합니다.
흔히 continuous batching 또는 dynamic batching이라고 불리는 방식은, iteration-level scheduling을 통해 매 step마다 batch를 재구성하는 구조라고 이해할 수 있습니다.
🤔 request level scheduling vs iteration level scheduling 3개의 request를 처리하는 두 가지 batching 방식. column은 iteration, row는 batch.
왼쪽은 request-level scheduling 방식으로, 3개의 request를 하나의 batch로 구성한 구조입니다. 요청 단위로 batch가 고정되어 있으며, batch 내의 모든 request가 종료될 때까지 해당 batch를 유지해야 합니다. 즉, 일부 request가 먼저 완료되더라도 나머지 request가 끝날 때까지 대기해야 합니다
오른쪽은 iteration-level scheduling 방식으로, 동일하게 3개의 request를 처리하지만 iteration 단위로 request를 관리합니다. 이때 iteration마다 batch에 포함되는 request는 달라질 수 있습니다. iteration 3, iteration 4와 같이, 완료된 request는 즉시 제거되고 새로운 request가 추가되는 것을 확인할 수 있습니다. 이 scheduling 방식을 활용하여 batch를 매 iteration마다 재구성하는 전략이 바로 continuous batching (dynamic batching)입니다.
Ⅲ Memory Challenges in LLM Serving
fine-grainded batching이 기존의 computing 낭비를 줄이는데도 불구하고, 여전히 batch 될 수 있는 request들의 수는 여전히 memory capacity, 특히 KV cache를 저장하기 위해 할당된 공간에 제한될 수 있습니다.
즉 serving system의 throughput은memory-bound합니다.
이러한 memory-bound를 극복하는 것은 memory management에 다음과 같은 challenge들을 해결하는 것을 필요로 합니다.
Large KV cache
KV cache의 크기는 reuqest들의 수에 따라 빠르게 커집니다.
예를 들어 13B parameter OPT model에 대해서 한 token의 KV cache을 저장하기 위해서는 800KB memory space가 필요합니다.
이때, OPT는 최대 2048 token들을 생성할 수 있기 때문에, 하나의 request에 대 KV cache를 저장하는데 memory는 최대
1.6GB만큼 있어야 합니다.
현재의 GPUs는 수십 GBs 수준의 memory capacities를 가집니다. 모든 사용 가능한 memory가 KV cache에 할당되었다 할지라도, 동시에 처리할 수 있는 request들의 수는 수십 개만 수용할 수 있습니다. 이런 상황에서 비효율적인 memory management는 batch size를 더욱 감소시킬 수 있습니다.
추가로, 현재 trends에서 GPU의 computation speed는 memory capacity보다 더욱 빨라지고 있습니다. 예를 들어 NVIDIA A100에서 H100으로 가면서 FLOPS(Floating Point Operations Per Second)는 전보다 2배다 더 증가했지만 GPU memory는 최대 80GB에서 멈춰있습니다.
따라서 저자들은 이러한 memory가 앞으로 점점 더 심각한 bottleneck(병목)이 될 것으로 판단했습니다.
Complex decoding algorithms
LLM service들은 유저가 선택할 수 있는 여러 decoding algorithms를 제공합니다. 각각은 memory management의 복잡성에 대한 결과가 다양하게 달라질 수 있습니다.
CASE 1) multiple random samples
user가 한 input prompt에 대해 여러 개의 random sample을 요청했을 때 prompt 부분의 KV cache는 memory 사용을 최소화하기 위해 공유됩니다. (temperature > 0)
반면에 autoregressive 생성 단계의 KV cache는 공유될 수 없습니다. 각 여러 sample 결과가 context와 position에 따라 달라지기 때문입니다.
이처럼 sharing 되는 KV cache의 정도는 특정 decoding algorithm에 달라질 수 있습니다.
CASE 2) Beam search
Beam search는 후보 seq(=beam)를 K개 유지하면서 매 step에서 각 후보를vocabulary 전체에 대해확장한 뒤, 누적 확률을 기준으로 pruning하여 상위 K개 seq만 남기는 방식입니다.
이후 전체 후보들 중 누적 확률이 가장 높은 "The cat is cute" (P = 0.08), "The cat was small" (P = 0.08) 두 문장만 다음 beam으로 유지됩니다.
이때 서로 다른 요청의 beam들이 자신의 KV cache 중 더 많은 부분을 서로 공유할 수 있으며 (최대 55%의 메모리를 절약할 수 있음. 실험 6.3 참고), decoding 과정이 진행됨에 따라 공유 방식은 계속 변화합니다. 초기에는 모든 beam들이 비슷한 prefix를 가지지만 점점 분기되면서 일부 token만 공유되는 것처럼 보이게 될 것입니다.
그렇기 때문에 memory management가 다양한 범위의 prompt lengths에 대한 공간을 제공할 수 있어야 합니다.
추가로 decoding에서 request의 output length가 늘어날 수 KV cache에 필요한 memory도 함께 증가하고, 이는 새로 들어올 request들과 이미 진행 중인 생성 과정을 위한 가용 메모리를 고갈시킬 수 있습니다.
따라서 시스템은 일부 request들의 KV cache를 GPU memory에서 deleting 혹은 swapping out과 같은 scheduling decisions을 해야 합니다.
3.1 Memory Management in Existing Systems
현재 deep learning framework들의 대부분의 opertor들은 tensor들이 연속적인 메모리에 저장되도록 요구하기 때문에, 이전의 LLM serving system들은 한 request의 서로 다른 위치에 있는 token들에 대한 KV cache를 연속적인 tensor로서 저장합니다.
그러나 LLM의 output lengths를 예측할 수 없기 때문에 request의 실제 input과 최종 output length와는 관계없이, 각 request에 대 최대 가능한 seq 길이를 기준으로 memory chunk를 정적으로 할당합니다.
위의 그림은 기존 system에서 2개의 request에 대한 KV cache memory의 예시를 보여줍니다.
request A는 최대 가능한 sequence 길이가 2048이고 request B는 512라고 가정해 봅시다.
기존 system에서 pre-allocation chunk 방식은 memory wastes에 대해서 3가지 주요 원인을 가집니다.
reserved memory
미래 token들에 대한 reserved slot
reseved memory는 최종적으로는 사용되지만, request의 전체 기간 동안 해당 공간을 점유하게 되므로, 특히 그 공간이 큰 경우에도, 원래라면 다른 reqeust들을 처리하는 데 사용될 수 있었던 공간을 차지하게 됨.
internal fragmentation (pure memory waste)
잠재적인 maximumsequence length에 맞춰 메모리 과잉 할당으로 인해 사용하지 않는 메모리 공간 발생
request의 sampling이 끝난 후에야 사용되지 않는다는 것을 알 수 있음.
external fragmentation (pure memory waste)
buddy allocator와 같은 memory allocator로부터 발생
request serving 전에 사용되지 않는다는 것을 알 수 있음. (serving 전에 생성 token에 필요한 연속된 메모리 공간을 미리 할당받으니까)
위의 그림은 memory wastes의 average percentage를 나타낸 것으로, 실험 결기존 system에서 실제 생성 token의 KV cache 저장에 사용 GPU memory 전체의 20.4%에 불과함을 확인할 수 있습니다.
fragmentation을 잠재적인 해결책으로서 compaction이 제안되었지만, serving 속도와 같이 성능에 민감한 LLM serving system에서 compaction을 수행하는 것은 대량의 KV cache 때문에 효과가 없습니다.
또한 compaction을 적용한다 해도, 기존 memory management system에서는 각 request에 대 미리 할당된 chunk 공간으로 인해 decoding alogrithm에 특화된 메모리 공유가 안됩니다.
🤔 Compaction이란? [used][empty][used][used][empty]와 같이 memory fragmentation이 발생했을 때 [used][used][used][empty][empty]와 같이 사용한 memory를 한쪽으로 몰아서 재배열하여 연속된 큰 공간을 확보하는 작업을 의미함
Ⅳ Method
저자들은 PagedAttention이라는 새로운 attention algorithm을 개발하고 기존 serving system의 memory management를 해결하기 위한 vLLM이라는 LLM serving engine을 구축하였습니다.
vLLM system overview
vLLM의 architecture는 위와 같습니다.
중앙scheduler
분산된 GPU worker들의 실행을 조정
KV cache manager
PagedAttention 알고리즘에 따라 page 방식으로 KV cache를 효율적으로 관리
중앙scheduler로부터 전달받은 instruction들을 통해GPU worker상의physical KV cache memory(VRAM)를 관리
해당 Method chapter에서는
PagedAttention algorithm 설명 (4.1)
KV manager 설계 (4.2)
PagedAttention과의 연결 (4.3)
다양한 decoding method 별 memory 최적화 (4.4)
가변 길이의 input and output seq 처리 (4.5)
vLLM에서의 분산 시스템 확장을 위한 system 설계 (4.6)
순으로 확장하여 설명합니다.
4.1 PagedAttention
저자들은 3에서 제시한 memory 문제를 해결하기 위헤, OS의 paging 개념에서 영감을 받은 attention algorithm인 PagedAttention을 제안합니다.
전통적인 attention algorithm들과 달리, PagedAttention은 continuous 한 key와 value vector들을 비연속적인(non-contiguous) 메모리 공간에 저장할 수 있게 합니다. 구체적으로, PagedAttention은 각 sequence(한 request의 일련의 token들)의 KV cache를 KV block들로 분할합니다.
이때 각 block은 고정된 개수의 token에 대한 key와 value vector들을 포함하며 이를 KV block size (B)라고 합니다. 즉 한 block에 고정된 token 수라고 이해하면 됩니다.
여기서 j번째 key block을 \(K_j = (k_{(j-1) B+1}, \ldots, k_{jB})\), j번째 value block을 \(V_j = (v_{(j-1) B+1}, \ldots, v_{jB})\)라고 나타낼 수 있습니다.
Attention computation in PagedAttention
attention score와 output in PagedAttention
이제 기존 attention 계산 시 token-wise computation을 위와 같은 block-wise computation으로 변환할 수 있습니다.
단지, token 단위의 key를 block 단위의 key로 변환하여 계산하는 것을 알 수 있습니다.
이때 \( A_{ij} = (a_{i, (j-1) B+1}, \ldots, a_{i, jB}) \)는 j번째 KV block에 대한 attention score의 row vector를 의미합니다.
attention computation 계산 시, PagedAttention kernel은 서로 다른 KV block들을 개별적으로 identify 하고 fetch 합니다.
KV Cache Memory Management in PagedAttention
위의 그림은 PagedAttention의 예시를 보여줍니다.
key와 value vector들은 3개의 block들에 걸쳐 분산되어 있습니다. 이때 3개의 block들은 physical memory에서 연속적이지 않는 것을 알 수 있습니다.
매 iteration마다, kernel은 3개의 block에 대하여
Attention score \(A_{ij}\) = \(q_i\) x \(K_j\) = query token "forth"의 query vector x 각 block의 key vector들 (ex. 0번째 block의 경우 "Four score and seven"의 key vector들)
를 구하고
각 block에 대해서
\(A_{ij}\) x \(V_j\) = attention score x block의 value vector
를 계산 후에
이를 모두 합하여 최종 attention output \(o_i\)을 구합니다.
✅ PagedAttention algorithm은 KV block들을 비연속적인 physical memory에서 저장될 수 있게 하며, 이는 vLLM에서 좀 더 유연한 paged memory management를 가능하게 합니다.
4.2 KV Cache Manager
KV cache Manager의 핵심 아이디어는 OS의 virtual memory 개념과 유사합니다.
OS의 virtual memory
OS는 memory를 고정된 크기의 page들로 분할하고 user process의 logical page들을 physical page들로 mapping 합니다.
이때, 연속적인 logical page들은 비연속적인 physical memory page들로 대응되기 때문에 user process가 마치 연속적인 것처럼 memory를 접근하게 됩니다. 게다가 필요한 physical memory 공간을 미리 전부 예약할 필요 없이, OS가 필요할 때마다 physical page를동적으로할당할 수 있습니다.
vLLM의 KV cache 관리 방식
vLLM은 이러한 OS의 virtual memory 아이디어를 사용하여 KV cache를 관리합니다. PagedAttention을 기반으로, KV cache를 virtual memory의 page처럼 고정된 크기의 KV block 단위로 구성합니다.
Logical KV blocks
한 request의 KV cache는 여러 개의 logical KV block의 연속으로 표현되며, 새로운 token과 그에 대응하는 KV cache가 생성될 때마다 왼쪽에서 오른쪽으로 순차적으로 채워집니다. 이때 마지막 KV block에서 아직 채워지지 않은 위치는 이후 생성될 token을 위해 예약됩니다.
Physical KV blocks
GPU worker 상에서는 block engine이 GPU DRAM에서 연속적인 chunk를 할당한 후, 이를 physical KV block 단위로 쪼갭니다.
KV cache manager
KV block manager는 각 request에 대한 logical KV block과 physical KV block 간의 mapping을 나타내는 block table을 관리합니다. 각 block table의 entry에는 한 logical block에 대응하는 physical block의 번호와 현재까지 채워진 개수를 기록합니다.
즉 이렇게 logical KV block과 physical KV block을 분리함으로써, vLLM은 메모리를 사전에 예약하지 않고도 KV cache memory를 동적으로 확장할 수 있습니다. 이는 기존 시스템에서 발생하던 대부분의 memory waste를 제거합니다.
4.3 Decoding with PagedAttention and vLLM
Example generation process for a request with PagedAttention.
위의 예시는 하나의 input sequence의 decoding process 동안 vLLM이 어떻게 PagedAttention을 실행하고 KV cache memory를 관리하는지를 빠르게 보여줍니다.
1. Prefill 단계에서 KV cache 저장
우선 prompt computation 과정에서 KV cache를 저장하는 데 필요한 만큼의 KV block을 할당합니다.
위의 예시에서는 prompt가 7개 token들로 구성되어 있으므로, vLLM은 0번째 logical KV block을 7번째 physical KV block에, 1번째 logical KV block을 1번째 physical KV block에 각각 mapping 한 것을 알 수 있습니다.
이때 prefill step에서 vLLM은 기존 self-attention 알고리즘으로 prompt의 KV cache를 생성하고 첫 번째 output token을 sampling 합니다. 이후 처음 4개의 token들에 대한 KV cache는 logical block 0번째에, 그다음 3개의 token들에 대한 KV cache는 logical block 1번째에 저장합니다. 남은 slot은 이후 autoregressive 생성 단계를 위해 예약됩니다.
2. Autoregressive decoding 단계에서 KV cache 저장
vLLM은 각 decoding iteration마다 batching을 위한 후보 sequeces의 집합을 선택하고, 새롭게 필요한 logical block에 대해 physical block을 할당합니다.
현재 iteration의 새로운 token 생성을 위해,PagedAttention kernel을 사용해서 현재 iteration의 모든 input token들 (prompt token들 + 현재까지 생성된 token들)의 KV cache들을 저장해 둔 logical KV block에 접근하고 새롭게 생성된 KV cache는 physical KV block에 저장합니다.
이때 커널은 여러 위치에 퍼져있는 KV cache를 병렬로 처리할 수 있기 때문에,hardware utilization을 증가시키고, latency를 줄일 수 있습니다.
예시를 보면,
physical block 7번째와 1번째(즉 prompt token들)에 대해 PagedAttention algorithm을 적용하여 새로운 token을 생성합니다. 이때 마지막 logical block에 하나 남은 slot의 위치에 새롭게 생성된 KV cache (즉 첫 번째 token의 Key, Value vector)를 저장합니다. 그리고 block table의 #filled의 기록도 update 합니다.
다음 decoding step에서는 마지막 logical block이 이미 다 채워졌기 때문에, 새롭게 생성된 KV cache를 새로운 logical block에 저장합니다. 이를 위해 새로운 physical block(3번째 physical block)을 할당하고, 해당 logical block과 physical block의 mapping을 block table에 기록합니다.
vLLM에서의 KV cache 관리 효과
vLLM이 logical block에 대해 새로운 physical block들을동적으로 할당하는 방식은 다음과 같은 효과를 가져옵니다.
1. Memroy utilization 증가
이전 block이 순차적으로 다 채워진 후에야, 새로운 physical block이 할당되기 때문에. 한 request에서 발생할 수 있는 memory waste는 one block 이내로 제한됩니다.
2. Throughput 향상
memory utilization이 증가하면 batching 과정에서 더 많은 reuqest들을 동시에 memory에 할당할 수 있습니다.
또한 한 request의 생성이 끝나면 해당 request의 KV Block들은 즉시 해제되어 다른 request에 사용될 수 있습니다.
Storing the KV cache of two requests at the same time in vLLM
더불어 위의 그림과 같이 여러 reqeust를 동시에 효율적으로 관리할 수 있게 됩니다. 두 request의 logical block들은 GPU worker 상의 block engine이 예약해 둔 memory space 내에서 서로 다른 physical block들에 매핑됩니다. 이때 physical memory space에서는 연속적일 필요가 없기 때문에, 보다 메모리를 효율적으로 활용할 수 있습니다.
4.4 Application to Other Decoding Scenarios
앞서 본 basic decoding algorithm은 greedy decoding이나 sampling과 같이, 하나의 prompt를 input으로 받아 하나의 output sequence를 생성합니다.
하지만 많은 LLM application에서는 더 복잡한 decoding 시나리오를 지원해야 하며, 접근 패턴이 복잡해지고 memory sharing이 더욱 필요해집니다.
Parallel sampling
하나의 input prompt에 대해서 여러 개의 output을 생성하고, 사용자가 여러 후보들 중에 한 output을 선택하는 경우입니다.
Example of parallel sampling.
해당 시나리오에서는, 하나의 request에서 같은 input prompt를 공유하며 여러 개의 sample들을 생성해야 하기 때문에, prompt의 KV cache 또한 공유되어야 합니다. 해당 저자들은 vLLM이 OS의 virtual memory에서의process forking 시사용되는 copy-on-write 방식과 유사하게 동작하도록 처리했습니다.
위의 예시는 2개의 output을 parallel decoding 하는 것을 보여줍니다.
1. Shared Prompt
우선 2개의 output은 같은 prompt를 공유하므로, prompt phase에서는 하나의 reserve space만 필요합니다.
따라서 2개의 sequence의 prompt에 해당하는 logical block들은 같은 physical block에 mapping 됩니다.
2. Reference count
이때 해당 physical block들은 두 sequence에 의해 공유되고 있으므로, 각 phsical block들의 reference count는 2가 됩니다.
위의 예시에서는 1번째, 7번째의 physical block의 ref count가 2가 됩니다.
3. Copy-on-write mechanism
generation phase에서는 두 output이 서로 다른 다른 token들을 생성하기 때문에, 이후의 KV cache는 별도로 분리해서 저장해야 합니다.
예를 들어 request A1이 마지막 logical block에서 새로운 토큰을 sampling 할 때, vLLM은 다음과 같은 절차를 수행합니다.
- 해당 logical block 1에 매핑된 physical block 1의 reference count가 1보다 큰 것을 확인
- 새로운 physical block(physical block 3)을 할당 및 매핑
- vLLM이 block engine에 기존 physical block 1의 내용을 새 block에 copy 하도록 명령
- 기존 physical block 1의 reference count를 1 감소
request A2가 physical block 1에 접근할 때에는 이미 reference count가 1이므로 새롭게 생성된 KV cache를 바로 저장합니다.
✅ 요약
- final logical block을 제외하고, prompt의 KV cache를 저장하는 데 사용되는 memory들만 공유합니다. 따라서 long input prompts의 경우 더욱 memory usage를 크게 줄일 수 있습니다
Machine translation과 같은 LLM tasks에서는 상위 k개의 적절한 번역 결과를 제공해야 하는 경우가 있습니다. 이런 시나리오의 경우 Beam search를 활용합니다.
Beam search는 매 decoding step마다 확률이 높은 상위 k개의 sequence들을 유지하는 방식입니다. 이때 decoding 동안 각 후보 sequnence에 대해 vocabulary space 내의 가능한 모든 토큰을 확장 후보로 고려하고 각 확률들을 계산합니다.
그 전체 \( k⋅∣V∣ \) 개의 후보 중에서 확률이 가장 높은 상위 k개만을 남깁니다. 이때 top k 후보의 수를 beam width라고 합니다.
기존 LLM serving system에서는 이러한 beam 후보들 간의 KV cache를 공유하기 어려워, 후보들이 분기 시 빈번한 memory 복사가 필요했습니다.
반면, vLLM에서는 앞서 본 copy-on-write mechanism을 통해 공유된 block 내에서 새로운 token이 생성되는 경우에만 block이 복사됩니다. 또한 beam 후보에 탈락한 sequnce가 사용하던 physical block를 ref count를 줄이고, 만약 ref count가 0이 되면 해당 physical block을 해제하여 재활용합니다.
Mixed decoding methods
vLLM은 앞서 살펴본 다양한 decoding method들을 서로 다른 request에 대해 동시에 처리할 수 있도록 지원합니다. 이는 기존 시스템에서는 효율적으로 수행하기 어려운 부분입니다.
이는 vLLM이 logical block을 physical block으로 변환하는 공통 mapping layer를 두어 서로 다른 sequence 간의 복잡한 memory sharing을 추상화하기 때문입니다.
그 결과, LLM과 실행 kernel은 각 sequence에 대해 단순히 physical block ID list만을 전달받아 처리하면 되며, sequence 간의 sharing pattern을 직접 관리할 필요가 없습니다.
4.5 Scheduling and Preemption
문제 상황 1) request traffic이 system의 capacity를 넘어서는 경우
vLLM은 first-come-first-serve(FCFS) scheduling policy를 사용하여 request들의 우선순위를 정합니다.
가장 일찍 도착한 request들을 먼저 처리하고 나중에 도착한 request들을 중단하는 정책입니다.
공정성을 보장하고 starvation을 방지합니다.
문제 상황 2) GPU block memory가 부족해지는 경우
LLM에서는 input prompt의 길이가 request마다 크게 달라질 수 있고, 생성되는 output length 또한 input prompt와 model에 따라 달라지므로 사전에 미리 알 수 없습니다.
request의 수와 outputs이 증가함에 따라, 새롭게 생성되는 KV cache를 저장할 GPU physical block이 부족해질 수 있습니다.
(1) 어떤 block들을 제거할 것인가?
하나의 sequence의 경우
모든 block들은 attention 계산 시 함께 사용되기 때문에, all-or-nothing evicition policy 를 적용합니다.
즉 한 sequence에 속한 block들을 전부 제거하거나, 아무것도 제거하지 않습니다.
한 request 내에 여러 sequence가 존재하는 경우 (ex. beam search의 beam candidate)
하나의 sequence group으로 묶여 gang-scheduling 됩니다. 즉 sequence들이 항상 함께 preempt 되거나 재스케줄링 됩니다.
이는 sequence group 내의 sequence들 간의 잠재적인 memory sharing이 존재할 수 있기 때문입니다.
(2) 다시 필요할 때 제거된 block들을 어떻게 복구할 것인가?
이에 대해서 2가지 technique이 존재합니다.
Swapping: vLLM이 제거된 block들을 CPU memory에 복사하고, CPU allocator가 CPU RAM에 swap된 physical block들을 관리
vLLM이 새로운 token을 위한 free physical block들을 확보할 수 없다면, 일부 sequence 집합을 선택하여 해당 sequence들의 KV cache를 CPU 이동시킵니다.
해당 sequence들은 preempt되고 해당 block들이 GPU에서 제거되면, vLLM은 preempt된 모든 sequence의 처리가 완료될 때까지 새로운 request 받지 않습니다.
어떤 request가 완료되면 해당 block들은 free되고, CPU memory에 저장되어 있던 preempt된 sequence의 block을 다시 GPU memory에 복구하여 처리를 재개합니다.
✅ CPU RAM으로 swap되는 block의 수는 GPU RAM에 존재하는 전체 physical block 수를 초과하지 않습니다. ✅CPU RAM swap space는 KV cache를 위해 할당된 GPU memory-bound합니다.
Recomputation:preempt된 sequence들이 재스케줄될 때 KV cache를 재계산합니다.
이때, decoding 시점에 생성된 token들을 기존 prompt와 이어 붙여 하나의 새로운 prompt를 만들 수 있습니다. 또한 그 prompt의 모든 position들의 KV cache는 prompt phase에서 한번의 iteration으로 계산됩니다.
✅ recomputation latency는 원래 latency보다 훨씬 더 낮아집니다.
따라서 두 technique의 성능은 CPU RAM과 GPU memory 간의 bandwidth (PCIe), 그리고 GPU의 계산 성능에 의존합니다.
4.6 Distributed Execution
Distributed Execution의 필요성
많은 LLMs의 parameter 크기는 단일 GPU의 capacity를 초과합니다. 따라서 parameter를 여러 GPUs에 분할하고 model parallel 방식으로 실행해야 합니다. 이 경우, memory manager는 분산 환경에서 memory를 처리할 수 있어야 합니다.
vLLM에서의 strategy
Megatron-LM style의 tensor model parallelism
Megatron-LM style의 tensor model parallelism 전략
이 전략은 SPMD(Single Program Multiple Data) execution schedule을 따릅니다. 즉, 각 GPU는 동일한 프로그램(코드)를 실행하되, 서로 다른 데이터의 조각을 처리합니다.
이 구조에서 linear layer는 분리되어 block-wise matrix multiplication을 수행하고, 분산 GPU들은 계속해서 all-reduce opertation을 통해 중간 result들을 지속적으로 동기화합니다.
EX) attention operator
attention은 attention head dimension을 기준으로 여러 GPU로 분산됩니다. 따라서 각 GPU는 multi-head attention에서 일부 attention head만을 담당합니다.
Model parallel execution의 동작 방식
각 GPU의 model shard는 동일한 input token 집합을 처리하기 때문에, 동일한 token 위치에 대한 KV Cache가 필요합니다. 즉 같은 token 위치에 대해 각 GPU가 담당하는 attention head에 해당하는 부분만 계산하는 구조입니다.
이때 vLLM은 중앙 scheduler 내에 단일 KV cache manager만으로도 작동이 가능합니다.
모든 GPU worker는 이 manager와 logical block → physical block 매핑 정보를 공유하기 때문에 각 GPU worker들이 scheduler가 할당한 physcial block을 기반으로 해당 input request에 대해 model을 실행시킬 수 있는 것이죠.
✅ 각 GPU worker는 동일한 physical block IDs를 가지지만, 자신이 담당하는 attention head에 해당하는 KV cache의 일부만을 저장합니다.
각 decoding step(=iteration)에서의 실행 과정
1. scheduler는 batch에 있는 각 request에 대해 input token IDs와 해당 request의 block table 포함한 control message를 준비합니다.
2. scheduler는 이 control message (token IDs, logical block, physical block mapping 정보)를 모든 GPU worker들에게 broadcast합니다.
3. GPU worker들은 input token IDs를 기반으로 model을 실행합니다.
4. Attention layer에서 GPU worker들은 control message에 있는 block table을 참조하여 KV cache를 읽습니다.
5. 실행 중 GPU worker들은 all-reduce communication primitive를 통해 중간 결과들을 동기화합니다. 이 과정은 scheduler의 개입 없이 수행됩니다.
6. 해당 iteration이 끝나면, 최종적으로 하나의 sampling된 token이 scheduler로 전송합니다.
✅ GPU worker는 memory management를 직접 수행하거나 동기화할 필요가 없습니다. 각 decoding iteration의 시작 시점에 scheduler로부터 필요한 모든 메모리 관리 정보를 전달받아, 그 정보를 기반으로 연산만 수행합니다.
Ⅴ Implementation
FastAPI: Frontend, OpenAI API interface를 확장하여 sampling parameter 설정하도록 지원
GPU-based infernce engine: Backend
Python : vLLM의 component들 구현
C++/CUDA: PagedAttention과 같은 key operation을 위한 custom CUDA kenerl 구현
PyTorch + Transformer: GPT, OPT, LLaMA 등의 LLM 모델 구현 및 실행
NCCL: 분산 GPU worker 간의 tensor commuincation(ex. all-reduce) 수행
5.1 Kernel-level Optimization
먼저 kernel은 여러 GPU 명령어들의 묶음이며, kernel launch는 해당 명령어 묶음을 GPU에서 실행하도록 요청하는 과정으로 이해할 수 있습니다.
PagedAttention은 기존 시스템에서 효율적으로 지원하지 않는 비연속적인 memory access pattern을 사용하기 때문에, 이를 최적화하기 위해 여러 custom GPU kernel을 직접 구현했다고 합니다.
1. Fused reshape and block write
각 Transformer layer에서
새로운 KV cache를 block들로 분할
block read에 필요한 memory layout 최적화로 reashape
block table에 의해 구체화된 position에 저장
이 과정은 원래 각각 별도의 kernel로 실행되어 총 3번의 kernel launch가 필요합니다.
kernel launch overhead를 줄이기 위해, vLLM은 이 세 단계를 하나의 kernel로 fuse하여 단일 launch로 처리합니다.
2. Fusing block read and attention
FasterTransformer의 attention kernel을 확장하여, block table을 기준으로 KV cache를 read와 attention 연산을 동시에 수행하도록 합니다.
각 GPU warp가 하나의 block을 읽도록 설계하여 coalesced memory access를 보장하고 variable sequence length를 지원
3. Fused block copy
copy-on-write mechanism에 의해 발생하는 block copy 연산은, 비연속적인 block들에 대해 발생하고 cudaMemcpyAsync API를 사용하면 작은 데이터 이동에 대한 호출 다수 발생하여 overhead가 커집니다.
이러한 overhead를 줄이기 위해 서로 다른 block들의 copy 연산을 하나의 kernel launch로 묶어 처리합니다.
5.2 Supporting Various Decoding Algorithms
vLLM은 fork, append, free라는 세 가지 핵심 method를 조합하여 다양한 decoding 알고리즘을 구현합니다.
fork: 기존 sequence로부터 새로운 sequence를 생성
append: sequence에 새로운 token을 추가
free: sequence를 삭제
예를 들어, parallel sampling에서는 다음과 같이 동작합니다.
fork: 단일 input sequence로부터 여러 output sequence를 생성
append: 각 sequence에 대해 매 iteration마다 새로운 token을 추가
free: stopping condition을 만족한 sequence를 삭제
Ⅵ Evalutation
6.1 Experimental Setup
1. Model
OPT: 13B, 66B, 175B parameters (13B와 66B는 LLM leaderboard에서 널리 사용되는 모델 크기이며, 175B는 GPT-3와 동일한 규모)
LLaMA: 13B parameters
2. Server
Google Cloud Platform의 A2 인스턴스
NVIDIA A100 GPU 사용
3. Dataset: 실제 LLM service들의 input과 output을 반영
ShareGPT: ChatGPT 대화 내용을 사용자가 공유한 conversation dataset
Alpaca: self-instruct 방식으로 GPT-3.5가 생성한 instruction dataset
4. Request arrival time
데이터셋에는 실제 timestamp 정보가 없기 때문에, 다양한 request rate를 시뮬레이션하기 위한 Poisson distribution을 사용하여 request arrival time을 생성
5. Baselines
FasterTransformer: latency에 최적화된 분산 inference engine
Orca:Throughput에 최적화된 SOTA LLM serving system
6. Key mertrics
Serving Throughput를 평가하기 위해 다음 metric을 측정합니다.
Normalized Latency
다양한 request rate에 대해 workload를 생성하고 normalized latency를 측정
각 request의 end-to-end latency를 해당 request의 output length로 나눈 값. 즉, token당 평균 latency
단순 throughput 대신 normalized latency를 사용하는 이유는, request마다 output token 수가 다르기 때문입니다.
✅ High throughput을 가진 serving system은 High request rate 환경에서도 Low normalized latency를 유지해야 합니다.
6.2 Result
저자들은 Basic Sampling, Parallel Sampling & Beam Search, Shared Prefix, Chatbot 등 다양한 decoding 시나리오에서 baseline들과 vLLM의 성능을 평가했습니다.
Graph 해석
논문에서 주로 다음과 같은 그래프들이 등장합니다.
Basic sampling의 예시
row: dataset (ShareGPT, Alpaca)
column: model
x축: 초당 request 수 (request rate)
y축: token당 latency (normalized latency)
Curve의 의미
그래프에 보이는 각 curve는 request rate가 증가함에 따라 normalized latency는 초반에는 완만하게 증가하다가, 특정 지점을 넘으면 급격히 폭등합니다.
request rate가 serving system의 처리 capacity를 초과하는 순간, queue 길이가 지속적으로 증가하게 되고, throughput은 더 이상 증가하지 못하고 대신latency만 증가하기 때문입니다.
따라서 이 curve를 통해 이 system이 최대 몇 req/s까지 안정적으로 유지할 수 있는지 throughput의 한계치를 알 수 있습니다.
실험 결과
Basic sampling에서는 동일한 latency일 때, vLLM은 Orca 대비 약 1.7~2배 이상 높은 request rate를 유지할 수 있습니다. 즉 더 높은 throughput을 달성합니다.
vLLM의 PagedAttention이 KV cache 메모리를 더 효율적으로 관리하여, 더 많은 request를 동시에 batching할 수 있기 때문입니다.
Parallel sampling의 경우 그 효과가 더욱 두드러지게 나타나는 것을 알 수 있습니다.
✅ 더 자세한 실험 결과는 직접 논문에서 읽어보시는 걸 추천드립니다!!
Ⅶ Ablation Studies와 Discussion
Kernel Microbenchmark
PagedAttention 수행 시 발생하는 GPU 연산, 특히 block read/write 및 attention kernel의 성능을 측정했습니다.
그림 (a)에 따르면, FasterTransformer와 비교했을 때 vLLM은 attention kernel latency가 약 20–26% 더 높게 나타났습니다.
이는 PagedAttention이 block table 접근, 추가 branch 실행, variable sequence length 처리 등으로 인한 overhead를 포함하기 때문입니다.
✅ vLLM에는 GPU 연산 overhead가 추가로 존재
Impact of Block Size
block size가 너무 작으면 GPU의 병렬성을 충분히 활용하지 못해 처리 효율이 떨어지고, 너무 크면 internal fragmentation이 증가하며 block sharing 확률이 감소합니다.
저자들은 Fig. (b)의 실험을 통해 다양한 block size를 비교한 결과, 대부분의 workload에서 block size 16이 GPU utilization과 memory efficiency 사이의 균형점임을 확인했습니다.
✅ vLLM의 기본 block size는 16으로 설정
Comparing Recomputation and Swapping
실험 결과, block size가 너무 작으면 swapping에서 CPU–GPU 간의 잦은 전송으로 overhead가 증가합니다.
반면, recomputation의 overhead는 block size와 무관하게 거의 일정합니다. 이는 recomputation이 기존 KV block을 직접 활용하지 않고 새로 계산하기 때문입니다.
block size가 16–64 범위인 중간 영역에서는 두 방식이 유사한 end-to-end 성능을 보입니다.
✅block size가 작은 경우 → recomputation이 더 효율적
✅block size가 큰 경우 → swapping이 더 효율적
Discussion
일반적인 DNN 훈련이나 DNN serving에서는 memory allocation을 사전에 최적화할 수 있고, 성능이 주로 compute-bound합니다. 따라서 vLLM은 LLM inference에 특화된 해결법이며, 일반 DNN workload에는 효과적이지 않습니다.
