https://arxiv.org/abs/2309.06180
Efficient Memory Management for Large Language Model Serving with PagedAttention
High throughput serving of large language models (LLMs) requires batching sufficiently many requests at a time. However, existing systems struggle because the key-value cache (KV cache) memory for each request is huge and grows and shrinks dynamically. Whe
arxiv.org
0. 초록
LLM의 높은 처리량 서빙은 충분히 많은 요청을 한번에 일괄 처리하는 것을 필요로 합니다. 그러나 기존 시스템은 각 요청에 대한 key-value cache(KV cache) 메모리가 매우 크고 동적으로 증가 및 축소되기 때문에 어려움을 겪습니다. 비효율적으로 관리될 경우, 이 메모리는 단편화 및 중복 복제로 인해 상당히 낭비될 수 있으며, 이는 배치 크기를 제한합니다. 이러한 문제를 해결하기 위해, 이 논문에서는 운영 체제의 고전적인 가상 메모리 및 페이징 기술에서 영감을 받은 어텐션 알고리즘인 PagedAttention을 제안합니다. 이를 기반으로, (1) KV cache 메모리에서 거의 제로에 가까운 낭비와 (2) 메모리 사용량을 더욱 줄이기 위해 요청 간에 KV cache를 유연하게 공유하는 LLM 서빙 시스템인 vLLM을 구축합니다.
자체 평가 결과 vLLM이 FasterTransformer 및 Orca와 같은 SOTA와 비교하여 동일한 수준의 지연 시간으로 LLM의 처리량을 2~4배 향상시킨다는 것을 보여줍니다.
1. 소개
LLM의 핵심에는 자기 회귀 모델인 트랜스포머가 있습니다. 이 모델은 입력(+프롬프트)과 지금까지 생성한 출력 토큰의 이전 시퀀스를 기반으로 단어 (토큰)를 한번에 하나씩 생성합니다. 각 요청에 대해, 이 프로세스는 모델이 종료 토큰을 출력할 때까지 반복됩니다. 이러한 시퀀스 생성 프로세스는 워크로드를 메모리 바운드 상태로 만들어 GPU의 계산 능력을 제대로 활용하지 못하고 서빙 처리량을 제한합니다.
이러한 처리량을 개선하는 것은 여러 요청을 함께 일괄 처리함으로써 가능합니다. 그러나 배치에서 많은 요청을 처리하려면 각 요청에 대한 메모리 공간을 효율적으로 관리해야 합니다.
예를 들어, 그림 1 (왼쪽)은 40GB RAM을 갖춘 NVIDIA A100 GPU에서 13B-파라미터 LLM에 대한 메모리 분포를 보여줍니다.
메모리의 약 65%는 서빙 중에 정적으로 유지되는 모델 가중치에 할당됩니다. 메모리의 약 30%는 요청의 동적 상태를 저장하는 데 사용됩니다. Transformer의 경우, 이러한 상태는 어텐션 메커니즘과 관련된 key 및 value 텐서로 구성되며, 일반적으로 KV cache 라고 하며, 이는 시퀀스에서 새로운 출력 토큰을 생성하기 위해 이전 토큰의 컨텍스트를 나타냅니다. 나머지 작은 1%의 메모리는 LLM을 평가할 때 생성되는 임시 텐서인 활성화를 포함한 다른 데이터에 사용됩니다.
모델 가중치는 일정하고 활성화는 GPU 메모리의 작은 부분만 차지하므로 KV cache를 관리하는 방식이 최대 batch size를 결정하는 데 매우 중요합니다. KV cache 메모리가 비효율적으로 관리되면 Fig. 1(오른쪽)에서 볼 수 있듯이 batch size와 결과적으로 LLM의 처리량이 크게 제한될 수 있습니다.
본 논문에서는 기존 LLM serving 시스템이 KV cache 메모리를 효율적으로 관리하지 못하고 있음을 확인했습니다. 이는 대부분의 딥러닝 프레임워크에서 텐서를 contiguous한 메모리 공간에 저장해야 하기 때문에 요청의 KV cache를 contiguous한 메모리 공간에 저장하기 때문입니다. 그러나 기존 딥러닝 워크로드의 텐서와는 달리 KV cache는 모델이 새로운 token을 생성함에 따라 동적으로 증가 및 축소되며, 수명과 길이가 사전에 알려져 있지 않다는 고유한 특징을 가지고 있습니다. 이러한 특징으로 인해 기존 시스템의 접근 방식은 두 가지 측면에서 매우 비효율적입니다.
첫째, 기존 시스템은 내부 및 외부 메모리 단편화로 어려움을 겪습니다. 요청의 KV cache를 contiguous한 공간에 저장하기 위해 요청의 최대 길이(예: 2048 token)로 contiguous한 메모리 chunk를 미리 할당합니다. 요청의 실제 길이가 최대 길이보다 훨씬 짧을 수 있으므로 심각한 내부 단편화가 발생할 수 있습니다. 게다가 사전 할당된 크기가 각 요청마다 다를 수 있으므로 외부 메모리 단편화도 클 수 있습니다. 실제로 [그림 2]의 프로파일링 결과에 따르면 기존 시스템에서 KV cache 메모리의 20.4% - 38.2%만이 실제 token 상태를 저장하는 데 사용됩니다.
"""
논문을 벗어나 조금 쉽게 설명드리자면,
내부 단편화는 외국 가수가 2048석의 콘서트장을 빌렸는데, 실제 관객은 500명만 예약한 상황인 것이죠. 남아있는 1548석은 텅 비어있지만, 이 공간은 이미 이 공연(요청)을 위해 할당되었기 때문에, 다른 가수(요청)가 사용할 수 없습니다. 즉, 할당된 메모리 블록(chunk) 내부에서 사용되지 않고 낭비되는 공간입니다.
외부 단편화는 할당된 메모리 블록들 사이사이에 생기는, 너무 작아서 활용하지 못하는 빈 공간들을 의미합니다. 여러 요청이 동시에 처리된다고 생각해보면,
- 요청 A: 최대 1000 토큰짜리 블록 할당
- 요청 B: 최대 500 토큰짜리 블록 할당
- 요청 C: 최대 2048 토큰짜리 블록 할당
GPU 메모리라는 주차장에 1000칸짜리 버스 A, 500칸짜리 승용차 B, 2048칸짜리 기차 C가 주차되어 있는 것과 같습니다. 이제 요청B가 처리를 완료하고 메모리를 해제(주차장 나감) 합니다. 그러면 메모리(주차장)에는 500칸짜리 빈공간이 생기겠죠. 이 때, 최대 800 토큰이 필요한 새로운 요청 D(리무진)가 들어옵니다.
GPU 메모리(주차장)에 남아있느 총 빈 공간은 500칸보다 훨씬 클지라도, 800토큰이 연속으로 들어갈 수 있는 빈 공간은 없습니다. 방금 생긴 빈 공간은 B 승용차가 나간 500칸 뿐이니까요. 결국 800 토큰 요청은 대기해야 합니다.
총 여유 공간은 충분하지만, 여기저기 잘게 쪼개져 있어서 큰 요청을 처리하지 못하는 상태가 외부 단편화입니다.
"""
둘째, 기존 시스템은 메모리 공유 기회를 활용할 수 없습니다. LLM 서비스는 종종 요청당 여러 출력을 생성하는 parallel sampling 및 beam search와 같은 decoding 알고리즘을 사용합니다. 이러한 시나리오에서 요청은 KV cache를 부분적으로 공유할 수 있는 여러 sequence로 구성됩니다. 그러나 sequence의 KV cache가 별도의 contiguous한 공간에 저장되기 때문에 기존 시스템에서는 메모리 공유가 불가능합니다.
위의 제한 사항을 해결하기 위해 운영 체제의 메모리 단편화 및 공유 솔루션인 페이징을 사용한 가상 메모리에서 영감을 얻은 attention 알고리즘인 PagedAttention을 제안합니다. PagedAttention은 요청의 KV cache를 block으로 나누며, 각 block은 고정된 수의 token에 대한 attention key와 값을 포함할 수 있습니다. PagedAttention에서 KV cache에 대한 block은 반드시 contiguous한 공간에 저장되는 것은 아닙니다. 따라서 OS의 가상 메모리에서와 같이 KV cache를 보다 유연하게 관리할 수 있습니다. 이 설계는 비교적 작은 block을 사용하고 필요에 따라 할당하여 내부 단편화를 완화합니다. 또한 모든 block의 크기가 동일하므로 외부 단편화를 제거합니다. 마지막으로 동일한 요청과 관련된 서로 다른 sequence 또는 서로 다른 요청 간에도 block 단위로 메모리 공유가 가능합니다.
"""
여기도 위에서 들었던 예시로 조금 쉽게 설명해드리겠습니다.
- 내부 단편화 해결:
- 처음부터 2048 토큰짜리 거대한 '공연장'을 통째로 할당하지 않습니다.
- 대신, 메모리를 아주 작은 고정된 크기의 블록(Block)(예: 16 토큰)으로 잘게 나눕니다. (마치 'A열, B열과 같이' 단위로 예약을 받는 것과 같습니다.)
- 요청이 토큰을 생성할 때마다 이 작은 블록(A열)을 필요한 만큼만 할당해 줍니다.
- 500 토큰이 필요하면 16토큰짜리 블록 약 32개만 할당합니다. 2048 토큰 분량을 미리 낭비할 필요가 없죠. (마지막 블록에서 약간의 낭비는 생길 수 있지만, 기존 방식에 비하면 거의 제로에 가깝다고 논문에서 말합니다.)
- 외부 단편화 해결:
- 모든 블록의 크기가 동일(Fixed-size)합니다.
- 주차장 비유에서 레고 블록으로 다시 예시를 들어보겠습니다. 1000칸, 500칸, 2048칸짜리 제각각인 큰 레고블록이 아니라, 모든 레고 블록이 16칸짜리 크기로 동일합니다.
- 따라서 어떤 요청(블록)이 메모리를 해제(출차)하든, 그 빈 공간은 정확히 16칸짜리입니다.
- 다음에 들어오는 어떤 16칸짜리 블록이든 그 자리에 완벽하게 들어맞습니다.
- 즉, 1000칸짜리 요청이 들어오면 1000칸짜리 공간 1개가 필요한 것이 아니라, 16칸짜리 표준 블럭 약 63개가 필요한 것이고, 500칸 요청이 들어오면 약 32개의 표준 블럭이 필요한 것이죠.
그럼 이 방법이 왜 외부 단편화를 해결해주는 것일까요?
1. 조각난 공간 활용: 500칸 요청(16칸짜리 32개)이 끝나면, 주차장 여기저기에 16칸짜리 빈칸 32개가 생깁니다.
2. 새 요청 처리: 이제 1000칸 요청(16칸짜리 63개 필요)이 들어옵니다.
3. 문제 해결: 시스템은 그냥 주차장 전체에서 비어있는 16칸짜리 표준 블록 63개를 찾기만 하면 됩니다. 방금 생긴 32개의 빈칸도 즉시 재사용할 수 있고, 나머지 31개는 다른 빈칸에서 가져오면 됩니다.
vLLM은 큰 요청을 잘게 쪼개기 때문에, 메모리도 잘게 관리할 수 있습니다. 모든 조각(블록)의 크기가 16칸으로 동일하므로, 어떤 블록이 비워지든 다음 요청이 즉시 그 자리를 재사용할 수 있어 낭비되는 조각 공간(외부 단편화)이 사라지는 것입니다.
"""
이러한 개선점으로 본 연구에서는 KV cache 메모리에서 거의 낭비가 없는 PagedAttention을 기반으로 고처리가 가능한 분산 LLM serving 엔진인 vLLM을 구축합니다. vLLM은 PagedAttention과 공동 설계된 block 수준 메모리 관리 및 선점형 요청 스케줄링을 사용합니다. 다양한 모델과 워크로드에 대한 평가 결과, vLLM은 모델 정확도에 전혀 영향을 미치지 않으면서 최첨단 시스템에 비해 LLM serving 처리량을 2-4배 향상시키는 것으로 나타났습니다.
2. 배경
이 섹션에서는 일반적인 LLM의 생성 및 서빙 절차와 LLM serving에 사용되는 iteration 수준 스케줄링에 대해 설명합니다.
1. Transformer 작동 원리
LLM은 이전에 나온 토큰들을 바탕으로 다음 토큰을 예측(생성)합니다.
- 이 과정의 핵심은 셀프 어텐션입니다.
- 어텐션은 Query(Q), Key(K), Value(V)라는 3가지 벡터를 사용합니다.
- 새 토큰을 만들 때, 이 토큰의 Q가 이전 모든 토큰의 K, V 값들과 상호작용하여 정보를 가져옵니다.
2. LLM 서빙 방식
LLM은 두 단계로 작동합니다.
- Prompt 단계 (빠름): 사용자의 입력(프롬프트)을 한 번에 병렬로 처리합니다.
- 생성 단계 (느림): 답변을 한 토큰씩 순차적으로 생성합니다. 서로 다른 반복에서의 계산은 데이터 종속성으로 인해 병렬화할 수 없으며 종종 행렬-벡터 곱셈을 사용하기 때문에, GPU 연산 능력을 크게 활용하지 못하고, 메모리 바운드가 되어 단일 요청의 대기 시간에 대부분을 차지합니다.
이 과정에서, 매번 다시 계산하는 것을 피하기 위해 이전에 생성된 모든 토큰의 Key(K)와 Value(V) 값을 'KV 캐시'라는 곳에 저장하고 재사용합니다.
3. 배치(Batching) 기술
느린 생성 단계의 GPU 효율을 높이기 위해, 여러 요청을 묶어서(배치) 처리하는 batching 메커니즘이 제안되었습니다.
- 문제점: 요청마다 도착 시간도 다르고, 입/출력 길이도 모두 다릅니다.
- 나쁜 해결책: 짧은 요청을 긴 요청 길이에 맞춰 패딩을 하면, 메모리와 연산이 심하게 낭비됩니다.
- 더 나은 해결책: '반복 수준 스케줄링' (vLLM의 Continuous Batching과 유사)을 사용합니다. 매 토큰 생성 단계(반복)마다, 완료된 요청은 즉시 내보내고 대기 중인 새 요청을 배치에 추가합니다. 이렇게 하면 패딩 낭비와 대기 시간이 줄어듭니다.
즉, 새 요청은 전체 batch가 완료될 때까지 기다리지 않고 단일 반복을 기다린 후 처리할 수 있으며 특수 GPU 커널을 사용하여 입력 및 출력을 패딩할 필요가 없어 대기열 지연과 패딩으로 인한 비효율성을 줄였습니다.
그러나, 이러한 기존 Batching 메커니즘은 컴퓨팅 낭비를 줄이고 요청을 보다 유연하게 처리할 수 있도록 했지만, 함께 배치할 수 있는 요청 수는 여전히 GPU 메모리 용량, 특히 KV 캐시를 저장하는데 할당된 공간에 의해 제한됩니다. 즉, 서빙 시스템의 처리량은 메모리에 의해 제한됩니다. 이 논문에서는 이러한 메모리 제약을 극복하고자 하였습니다.
3. LLM 서빙시 메모리 문제
논문은 LLM 서비스의 처리량이 KV 캐시 메모리에 의해 결정적인 제약을 받는다고 말합니다. 이를 유발하는 근본적인 문제점 세 가지와, 기존 시스템이 이 문제를 어떻게 악화시키는지 설명합니다
LLM 서빙을 어렵게 만드는 메모리 관련 문제는 다음과 같습니다.
- ① KV 캐시 자체가 너무 크다:
- 토큰 하나를 처리하는 데 필요한 KV 캐시 용량(예: 13B 모델 기준 800KB)은 작아 보일 수 있습니다.
- 하지만 최대 시퀀스 길이(예: 2048 토큰)를 모두 채우면, 요청 1건당 최대 1.6GB라는 엄청난 메모리가 필요합니다.
- 80GB의 최신 GPU 메모리(H100)로도 고작 수십 개의 요청만 동시에 처리할 수 있습니다.
- GPU의 연산 속도(FLOPS)는 메모리 용량(VRAM)보다 훨씬 빠르게 증가하고 있습니다. (A100 → H100 2배 이상 증가) 그러나, GPU 메모리는 유지 됩니다. 따라서 메모리가 더욱 심각한 병목 지점이 될 것이라고 보고 있습니다.
- ② 복잡한 디코딩 알고리즘:
- 단순히 텍스트를 하나 생성하는 것 외에, 빔 서치(Beam Search)나 여러 샘플 동시 생성(Parallel Sampling) 같은 복잡한 기능이 필요합니다.
- "Parallel Sampling (N=3)"에 대해 설명해드리도록 하겠습니다.
- 입력(Prompt): "한국에서 가장 예쁜 아이돌은"
- N=3 요청: 이 프롬프트에 대해 3가지의 다른 답변을 생성해달라는 뜻입니다.
- 예상 답변:
- "윈터입니다. 키가..."
- "카리나입니다. SM에서..."
- "장원영입니다. 그녀는..."
- 문제점: 1, 2, 3번 답변 모두, " 한국에서 가장 예쁜 아이돌은 "이라는 공통된 프롬프트에 대한 KV 캐시를 계산해야 합니다. 이 프롬프트의 KV 캐시가 100MB라고 가정하면, 3개의 샘플을 위해 총 300MB (100MB x 3)의 메모리가 필요합니다.
- Beam Search도 자세히 예시를 들어드리겠습니다.
- 여러 갈래의 '가지'를 동시에 탐색하다가, 가능성 없는 가지는 '가지치기(Pruning)'하는 방식입니다.
- 1단계: "윈터", "카리나", "장원영" (3개의 빔 생성)
- 2단계:
- "윈터" -> " 윈터는", " 윈터의"
- "카리나" -> "카리나가", "카리나도"
- "장원영" -> "장원영은", "장원영의"
- 3단계 (가지치기): 이 6개 중 가능성이 가장 낮은 3개를 버리고, 상위 3개만 남깁니다. (예: "윈터는", "카리나가", "장원영은")
- 4단계: 이 3개를 또 확장합니다.
- 동적 복잡성:
- 1단계에서 " 윈터", " 카리나", " 장원영"는 모두 '공통 프롬프트'를 공유합니다.
- 2단계에서 "윈터는"과 "윈터의"는 '공통 프롬프트' + '윈터'의 KV 캐시를 공유합니다.
- 3단계에서 "카리나도"가 버려졌다면, 이 빔이 사용하던 메모리는 해제되어야 합니다.
- 즉, 누가 누구와 얼마큼의 KV 캐시를 공유하는지가 매 토큰 생성 단계마다 실시간으로 바뀌고, 쪼개지고, 합쳐지고, 버려집니다.
- 이런 복잡한 공유 트리(Tree) 구조를 연속 메모리 방식으로는 도저히 관리할 수 없습니다.
- "Parallel Sampling (N=3)"에 대해 설명해드리도록 하겠습니다.
- 이런 디코딩 알고리즘들은 '프롬프트(입력)' 부분의 KV 캐시를 공유한다면 큰 기회를 얻을 수 있습니다. 300MB의 메모리 대신 100MB만 사용할 수도 있고, 빔 서치에서는 KV 캐시의 최대 55%까지 공유하여 메모리를 아낄 수 있습니다.
- 그러나 문제는 기존 시스템에서는 이러한 기회를 구현할 수 있는 기술적인 방법이 없습니다. 기존 시스템은 연속된 메모리를 할당하기 때문에 여러 요청(샘플, 빔)이 메모리의 특정 부분(프롬프트 캐시)를 공유하도록 만들 수 없습니다. 그래서 기존 시스템은 N개의 샘플(빔)을 위해 N개의 메모리를 따로따로 할당해야 합니다. 결국 동일한 프롬프트 KV캐시를 N번 중복 복제하게 되어 메모리가 낭비되는 것입니다.
- 단순히 텍스트를 하나 생성하는 것 외에, 빔 서치(Beam Search)나 여러 샘플 동시 생성(Parallel Sampling) 같은 복잡한 기능이 필요합니다.
- ③ 예측 불가능한 입출력 길이:
- 사용자가 언제, 얼마나 긴 프롬프트를 입력할지, 그리고 모델이 얼마나 긴 답변을 생성할지 미리 알 수 없습니다.
- 답변이 생성됨에 따라 KV 캐시가 점점 더 많은 메모리를 요구합니다.
- 이로 인해 메모리가 가득 차면, 시스템은 기존 요청을 중단(삭제)하거나 메모리에서 내보내는(Swap-out) 어려운 스케줄링 결정을 내려야 합니다.
그럼 기존 시스템은 왜 메모리 낭비가 심할까요?
기존 시스템은 위의 3가지 문제를 해결하기는커녕, 오히려 메모리 낭비를 극대화합니다.
- 근본 원인: 딥러닝 프레임워크(PyTorch 등)는 텐서가 '연속적인(Contiguous) 메모리' 공간에 저장될 것을 요구합니다.
- 잘못된 해결책 (정적 사전 할당):
- 출력 길이를 모르기 때문에, 기존 시스템은 각 요청이 사용할 최대 시퀀스 길이(예: 2048 토큰)만큼의 거대한 연속 메모리 블록(Chunk)을 미리 통째로 할당합니다.
이 방식은 다음과 같은 3가지 심각한 메모리 낭비를 유발합니다.
- 내부 단편화 (Internal Fragmentation):
- 최대 2048 토큰을 예약했지만 실제로는 100 토큰만 생성하고 끝난 경우, 사용되지 않은 1948 토큰만큼의 공간이 통째로 낭비됩니다.
- 외부 단편화 (External Fragmentation):
- 요청 A(최대 2048)와 요청 B(최대 512)처럼 서로 다른 크기의 블록을 할당하고 해제하는 과정이 반복됩니다.
- 이로 인해 메모리 공간이 '사이사이'에 쓸모없는 작은 조각들로 나뉘게 됩니다. 총 메모리는 남아있지만, 큰 요청이 들어갈 연속 공간이 없어 낭비됩니다.
- 예약된 메모리 (Reserved Memory):
- 설령 2048 토큰을 모두 사용할 요청이라도, 100번째 토큰을 생성하는 시점에는 뒤의 1948칸이 텅 빈 채로 예약되어 있습니다.
- 이 공간은 요청이 완전히 끝날 때까지 다른 어떤 요청도 사용할 수 없는 순수한 낭비 상태입니다.
이러한 낭비들로 인해, 앞선 [그림 2]에서 보듯이 KV 캐시를 위해 할당된 메모리의 단 20.4%만이 실제로 사용될 정도로 비효율적일 수 있습니다. 또한, 이렇게 요청별로 메모리를 통째로 할당하는 방식은 3장에서 언급한 복잡한 디코딩 알고리즘의 메모리 공유 기회를 원천적으로 차단해 버립니다.
4. 방법
vLLM 논문의 4장 Method 섹션은 이 논문의 핵심 기술인 PagedAttention과 이를 기반으로 한 vLLM 시스템의 아키텍처 및 작동 방식을 상세히 설명합니다. 핵심은 운영체제의 '가상 메모리'와 '페이징' 기술을 LLM의 KV 캐시 관리에 도입한 것입니다.
1. PagedAttention (4.1절)
PagedAttention이 vLLM의 핵심이 되는 새로운 어텐션 알고리즘입니다. 논문에서는 운영 체제의 페이징이라는 고전적인 기법에서 아이디어를 얻었다고 합니다. 이 알고리즘의 핵심은 기존의 어텐션 알고리즘과 달리 연속적인 키와 값을 비연속적인 메모리 공간에 저장할 수 있는 것입니다.
기존 Attention
- 하나의 거대한 통짜 그림(연속된 메모리)을 가지고 계산해야 했습니다.
- 이 통짜 그림을 만들기 위해 메모리를 미리 거대하게 할당해야 했고, 여기서 단편화가 발생했습니다.
PagedAttention의 해결책:
어차피 계산이 똑같은데, 그림을 미리 합쳐놓을 필요가 있을까? -> 그림을 퍼즐처럼 여기저기 흩어놓고(비연속 메모리), 계산할 때 마다 필요한 조각만 가져와서 계산하자라는 아이디어입니다.
- 시퀀스의 전체 KV 캐시를 KV 블록이라는 고정된 크기(예: 16토큰)의 작은 조각으로 분할합니다.
- 비연속적 메모리 저장 : 이 블록들이 물리적 GPU 메모리에 연속적으로 있을 필요 없이, 여기저기 흩어져 저장될 수 있도록 허용합니다.
- 블록 단위 어텐션 계산 : 어텐션을 계산할 때, PagedAttention 커널이 이 흩어져 있는 블록들을 개별적으로 가져와서(fetch) 계산을 수행합니다.
그럼 pagedAttention의 아이디어가 수학적으로 어떻게 가능한지 수식 부분을 좀 더 자세하게 살펴보겠습니다.
$$
A_{i j}=\exp \left(q_i^{\top} K_j / \sqrt{d}\right) \prod_{t=1}^{\lceil i / B\rceil} \exp \left(q_i^{\top} K_t 1 / \sqrt{d}\right), \quad o_i=\sum_{j=1}^{\lceil i / B\rceil} V_j A_{i j}^{\top}
$$
1. $A_{ij}$ (어텐션 스코어 / Softmax)
$$
A_{i j}=\exp \left(q_i^{\top} K_j / \sqrt{d}\right) \prod_{t=1}^{\lceil i / B\rceil} \exp \left(q_i^{\top} K_t 1 / \sqrt{d}\right)
$$
이건 복잡해 보이지만, Softmax를 거쳐 나온 어텐션 확률(가중치)입니다. "전체 중에서 현재 블록이 얼마나 중요한가?"를 계산하는 과정이죠. 수학적으로 기존 어텐션과 동일함을 확인하실 수 있습니다. 앞선 [그림 5]로 예시를 들어보겠습니다.
- $q_i$: 현재 생성 중인 토큰(예: "forth")의 Query 벡터입니다. (이 '하나'의 Query로 과거의 모든 Key를 스캔합니다.)
- $K_j$: $j$번째 Key 블록(예: "Four score and")입니다. 하나의 벡터가 아니라, 여러 Key 벡터의 묶음(블록)입니다.
- $B$: 블록 크기 (예: 16 토큰)
- $\lceil i/B \rceil$: 현재까지 생성된 $i$개의 토큰을 저장하는 데 필요한 총 블록의 개수입니다.
- 분자 (Numerator): $\exp(q_i^\top K_j / \sqrt{d})$
- 현재 Query($q_i$)가 $j$번째 Key 블록($K_j$)과 얼마나 관련이 있는지 계산한 관련도 점수입니다.
- 분모 (Denominator): $\sum_{t=1}^{\lceil i/B \rceil} \exp(q_i^\top K_t / \sqrt{d})$
- t=1부터 마지막 블록까지, 모든 블록($K_t$)의 '관련도 점수'를 전부 더한 값(총합)입니다.
- $A_{ij}$ (최종 결과): 분자 / 분모
- ( j번째 블록의 지수화된 스코어) / (모든 블록의 지수화된 스코어 총합)으로
- "모든 블록의 점수 총합" 중에서 "$j$번째 블록의 점수"가 차지하는 비율입니다.
- 즉, $j$번째 블록에 대한 최종 어텐션 확률 (0~1 사이의 값)입니다.
2. $o_i$ (최종 출력 계산)
$$
\quad o_i=\sum_{j=1}^{\lceil i / B\rceil} V_j A_{i j}^{\top}
$$
어텐션 확률($A_{ij}$)을 사용해 Value 벡터들을 가중 평균하는 과정입니다.
- $V_j$: $j$번째 Value 블록입니다. (Key 블록 $K_j$와 짝을 이룹니다.)
- $A_{ij}^\top$: 방금 위에서 계산한 $j$번째 블록의 어텐션 확률입니다.
- $V_j A_{ij}^\top$ (부분 계산):
- $j$번째 Value 블록($V_j$)에 그 중요도(확률 $A_{ij}$)를 곱합니다.
- $j$번째 블록에서 가져올 정보 조각이 됩니다.
- $\sum_{j=1}^{\lceil i/B \rceil} \ldots$ (총합):
- j=1(첫 번째 블록)부터 마지막 블록까지, 모든 블록에서 계산된 정보 조각들을 모두 더합니다.
- 최종 결과($o_i$): 모든 과거 블록들의 정보를 중요도에 따라 합산한 최종 출력 벡터가 됩니다.
이러한 수식이 중요한 이유는, 계산이 블록별로 완벽하게 분리된다는 점입니다.
1. $A_{ij}$ 계산: $j$번째 블록의 점수($A_{ij}$)를 계산할 때, 분모(총합)를 구하기 위해 다른 블록($K_t$)의 관련도 스코어가 필요하긴 하지만, $K_j$ 자체의 계산은 독립적입니다.
2. $o_i$ 계산: $j$번째 블록의 정보 조각($V_j A_{ij}^\top$)을 계산할 때, 다른 블록($V_t$)은 전혀 필요 없습니다.
논문의 [그림 5]가 바로 이 뜻입니다. PagedAttention 커널은 흩어져 있는 블록들을 하나씩 순회합니다.
- Block 0 ("Four score and seven")을 메모리에서 가져옵니다.
- $K_0$, $V_0$를 이용해 '정보 조각 0'($V_0 A_{i0}^\top$)을 계산합니다.
- Block 1 ("years ago our")을 (완전히 다른 메모리 주소에서) 가져옵니다.
- $K_1$, $V_1$를 이용해 '정보 조각 1'($V_1 A_{i1}^\top$)을 계산합니다.
- Block 2 ("fathers brought")를 (또 다른 메모리 주소에서) 가져옵니다.
- $K_2$, $V_2$를 이용해 '정보 조각 2'($V_2 A_{i2}^\top$)를 계산합니다.
- 마지막에 이 '정보 조각'들을 모두 더해서($\sum$) 최종 출력 $o_i$를 만듭니다.
그래서 결론은 위와 같은 수식들이 어텐션 계산은 각 블록에 대해 독립적으로 수행한 뒤, 나중에 합쳐도(sum) 된다는 것을 보여줍니다. 이 덕분에 vLLM은 KV 캐시를 연속된 통 메모리가 아닌, 흩어져 있는 블록(페이지)으로 관리할 수 있습니다.
vLLM의 핵심은 PagedAttention이니까 한번 더 정리해보자면
"Four score and seven years ago our fathers brought forth"이라는 문장이 있습니다.
기존 어텐션 :
- PyTorch 같은 프레임워크는 torch.cat([past_kv, new_kv]) 처럼 텐서를 합칠 때 물리적으로 연속된 메모리를 요구합니다.
- 그래서 "앞으로 2048 토큰까지 길어질지도 몰라"라며 미리 2048칸짜리 VRAM을 통째로 비워놓습니다.
- "Four score and seven years ago our fathers brought"에 대한 kv 캐시를 하나의 VRAM에 연속 메모리에 저장합니다.
- 'forth'라는 쿼리 벡터를 계산하기 위해 이전 "Four score and seven years ago our fathers brought"의 kv캐시를 가져와 활용합니다. 이 kv캐시는 연속 메모리에 저장하고 있고, 10개의 토큰 뿐이지만 2048 토큰의 메모리가 할당되어 있습니다.
- 여기서 단편화 문제가 발생합니다.
vLLM의 저장 방식 :
- 가상 메모리: 정확히는 블록 테이블(Block Table)이라는 지도를 통해 관리합니다. 실제 데이터는 VRAM 여기저기 흩어져(물리적) 있지만, vLLM이 볼 때는 연속된 것처럼(논리적/가상) 관리합니다.
- On-demand 할당: "Four score and" (3토큰)이 들어오면 딱 블록 1개만 할당합니다. 2048개를 미리 잡지 않습니다. 그래서 낭비가 없습니다.
- "Four score and seven years ago our fathers brought forth" (forth 생성 중)
- 입력: 현재 토큰 "forth"가 GPU에 도착합니다.
- Projection: 모델의 가중치($W_q, W_k, W_v$)를 곱해서 $q_{forth}, k_{forth}, v_{forth}$ 벡터 3개를 먼저 만듭니다.
- 캐시 업데이트:
- 만약 마지막 블록(Block 2)에 빈칸이 남았다면? -> $k_{forth}, v_{forth}$를 Block 2의 빈칸에 채워 넣습니다.
- Block 2가 꽉 찼다면? -> 새로운 Block 3을 할당하고 거기에 $k_{forth}, v_{forth}$를 넣습니다.
- 결과: 이제 forth의 정보도 KV 캐시(블록 어딘가)에 저장되었습니다.
- PagedAttention 실행 (블록 조회 & 연산):
- 이제 $q_{forth}$를 들고 모든 블록을 순회합니다.
- Block 0: (Four score and) 조회 -> 연산
- Block 1: (seven years ago) 조회 -> 연산
- Block 2: (our fathers brought) 조회 -> 연산
- Block 3 (또는 2의 끝부분): 방금 저장한 (forth) 조회 -> 연산
- 합산 (Reduction): 위 모든 결과를 합쳐서 최종 $o_{forth}$를 만듭니다.
2. KV Cache Manager (4.2절)
PagedAttention을 기반으로, vLLM은 OS의 가상 메모리 시스템과 거의 동일하게 작동하는 KV 캐시 관리자를 구현했습니다.
- 핵심 구성 요소:
- 물리 블록 (Physical Blocks): GPU VRAM에 실제로 존재하는 고정 크기의 빈 슬롯들입니다.
- 논리 블록 (Logical Blocks): 각 요청(Request)의 입장에서 내 KV 캐시의 첫 번째, 두 번째, 세 번째 블록...처럼 연속적으로 보이는 가상의 블록입니다.
- 블록 테이블 (Block Table): 가장 중요한 요소입니다. 각 요청마다 논리 블록과 물리 블록을 1:1로 매핑(mapping)해주는 지도입니다. (예: 요청 A의 논리 블록 3은 물리 블록 72번에 저장되어 있다.)
- 작동 방식:
- 요청이 들어오면, 시스템은 최대 길이(2048)만큼의 메모리를 미리 할당하는 것이 아니라, 당장 필요한 만큼의 물리 블록만 할당하고 이를 논리 블록에 매핑합니다.
- 토큰이 하나씩 생성되어 KV 캐시가 더 필요해지면, 관리자는 비어있는 물리 블록을 아무 데서나 하나 더 가져와서 이 요청의 논리 블록 리스트에 추가합니다.
- 효과:
- 내부 단편화 해결: 최대 길이만큼 미리 할당하지 않고, 필요한 만큼만 블록을 할당하므로 낭비가 거의 사라집니다. (낭비는 최대 마지막 1개 블록 내부에서만 발생)
- 외부 단편화 해결: 모든 블록의 크기가 동일하므로, 조각난 빈 공간이 발생하지 않습니다.
3. Decoding with PagedAttention and vLLM (4.3절)
논문 4.3 섹션은 실제로 vLLM이 토큰을 하나씩 생성할 때(Decoding), 메모리(블록)가 어떻게 할당되고 채워지는지를 구체적인 예시와 함께 설명하는 부분입니다.
1. 구체적인 작동 예시 (7개 토큰 프롬프트 상황)
논문에서는 블록 크기(Block Size)가 4라고 가정하고, 프롬프트가 7개 토큰인 상황을 예로 듭니다.
① 초기 상태: 프롬프트 단계 (Prefill Phase)
- 상황: 프롬프트 "Four score and seven years ago our" (7토큰)이 들어옵니다.
- 할당: 7토큰을 저장하려면 4칸짜리 블록 2개가 필요합니다.
- 논리 블록 0: 4개 꽉 채움 ("Four...seven") → 물리 블록 7에 매핑
- 논리 블록 1: 3개 채움 ("years...our") → 물리 블록 1에 매핑
- 남은 공간: 논리 블록 1의 마지막 1칸은 비어 있습니다. (예약됨)
- 특징: OS처럼 최대 길이(2048)를 미리 잡지 않고, 딱 필요한 2개 블록만 잡았습니다.
② 첫 번째 생성 단계 (1st Decoding Step)
- 상황: 8번째 토큰(예: "fathers")을 생성합니다.
- 연산: PagedAttention이 물리 블록 7과 1을 조회하여 연산합니다.
- 저장: 새로 생긴 8번째 토큰의 KV 캐시는 어디로 갈까요?
- 논리 블록 1에 빈칸(슬롯)이 하나 남았었죠?
- 새 토큰은 물리 블록 1의 마지막 빈칸에 들어갑니다.
- 이제 블록 테이블의 #filled 정보가 업데이트됩니다 (3/4 → 4/4).
③ 두 번째 생성 단계 (2nd Decoding Step)
- 상황: 9번째 토큰(예: "brought")을 생성합니다.
- 문제: 이제 논리 블록 0, 1이 모두 꽉 찼습니다. 더 이상 빈칸이 없습니다.
- 해결 (동적 할당):
- vLLM은 새로운 물리 블록 3을 하나 할당해옵니다.
- 이것을 논리 블록 2에 매핑합니다.
- 9번째 토큰의 KV 캐시는 물리 블록 3의 첫 번째 칸에 저장됩니다.
2. 전체적인 실행 흐름 (Global Execution Loop)
vLLM이 매번 디코딩 Iteration을 돌 때마다 수행하는 작업 순서입니다.
- 스케줄링 & 할당:
- 배치(Batch)에 포함될 요청들을 선택합니다.
- 각 요청에 대해 이번 턴에 새 토큰을 저장할 물리 블록이 필요한가?를 확인하고 미리 할당합니다.
- 입력 연결 (Concatenation):
- 현재 처리할 모든 토큰(프롬프트 단계의 전체 토큰 + 생성 단계의 최신 토큰 1개)을 쭉 이어 붙여서 모델에 넣습니다.
- PagedAttention 실행:
- 모델이 연산할 때, PagedAttention 커널은 블록 테이블을 보고 흩어져 있는 과거의 KV 캐시(물리 블록들)를 가져와 연산합니다.
- 연산 결과로 나온 새로운 KV 캐시는 방금 할당받은(혹은 자리가 남은) 물리 블록 위치에 저장합니다.
++ 블록 크기의 트레이드오프 (Trade-off)
- 장점: 블록 크기가 크면(예: 16, 32), 한 번에 여러 데이터를 처리하므로 병렬 처리가 잘 되어 속도가 빨라집니다.
- 단점: 하지만 블록 크기가 너무 크면, 마지막 블록에서 남는 빈칸(내부 단편화)도 커져서 메모리 낭비가 생깁니다. (논문 뒷부분 7.2절에서 최적의 크기를 연구함)
이 섹션에서 강조하는 vLLM의 효율성은 다음과 같습니다.
- 낭비의 최소화: 메모리 낭비는 오직 각 요청의 맨 마지막 블록에서만 발생합니다. (내부 단편화)
- 나머지 모든 블록은 꽉꽉 채워져 있으므로 낭비가 없습니다.
- 물리적 유연성:
- 논리적으로는 0 -> 1 -> 2 순서지만, 물리적으로는 7 -> 1 -> 3 처럼 메모리 아무 곳이나 비어있는 곳을 씁니다.
- 따라서 외부 단편화(중간중간 못 쓰는 빈 공간)가 완벽하게 제거됩니다.
- 처리량(Throughput) 증가:
- 메모리를 알뜰하게 쓰니, 남는 메모리에 더 많은 요청(Batch Size)을 동시에 태울 수 있습니다.
- 즉시 회수:
- 요청이 끝나면 사용하던 블록들은 즉시 해제되어, 다른 요청이 바로 가져다 쓸 수 있습니다.
즉, vLLM은 미리 자리를 맡아두지 않고, 토큰이 생성될 때마다 빈자리를 채우거나 새 조각(블록)을 하나씩 가져와서 연결하는 방식으로 메모리를 낭비 없이 쓴다는 것입니다.
4. Application to Other Decoding Scenarios (4.4절)
앞서 4.1~4.3이 메모리 단편화(낭비)를 어떻게 없앴는가?에 대한 이야기였다면, 4.4는 "그래서 확보한 그 유연한 메모리 구조로 얼마나 똑똑하게 공유를 해서 성능을 끌어올렸는가?"에 대한 이야기입니다.
1. 병렬 샘플링 (Parallel Sampling)
- 상황: "프롬프트 A에 대해 답변 3개(A1, A2, A3)를 만들어줘."
- 기존: 프롬프트 A의 KV 캐시를 3개 복사해서 저장함. (낭비 심함)
- vLLM:
- 프롬프트 A의 물리 블록은 딱 1번만 저장합니다.
- 답변 A1, A2, A3의 논리 블록들이 모두 같은 물리 블록을 가리킵니다. (참조 카운트 = 3)
- Copy-on-Write: 만약 A1이 공유 중인 블록에 무언가 쓰려고 할 때만, 그 블록을 복사해서 A1 전용으로 분리합니다.
- 결과: 프롬프트가 아무리 길어도 메모리를 거의 먹지 않습니다.
2. 빔 서치 (Beam Search)
- 상황: 가장 확률 높은 문장을 찾기 위해 여러 가지(가지치기)를 동시에 탐색. (번역 등에서 사용)
- 특징: 탐색 트리가 복잡하게 갈라지고(Fork), 합쳐지고, 버려집니다.
- vLLM:
- 동적 공유: 빔이 갈라질 때마다 이전 조상 블록들을 공유합니다.
- 자동 해제: 가지치기 당해서 빔이 사라지면, 해당 블록의 참조 카운트를 줄입니다. 카운트가 0이 되면 즉시 메모리에서 해제됩니다.
- 효과: 기존 시스템은 빔 서치 때마다 KV 캐시를 통째로 복사하느라 느렸지만, vLLM은 블록 단위 매핑만 바꾸면 되므로 훨씬 빠릅니다.
3. Shared Prefix / 시스템 프롬프트
- 상황: 챗봇의 성격을 정의하는 긴 지시문("너는 친절한 AI야... 예시 1... 예시 2...")이 모든 요청의 맨 앞에 공통적으로 붙는 경우.
- vLLM:
- 이 공통 부분(시스템 프롬프트)의 KV 캐시를 미리 계산해서 딱 한 번만 저장해 둡니다. (마치 OS의 공유 라이브러리처럼)
- 새로운 사용자 요청이 들어오면, 앞부분은 계산하지 않고 이 미리 저장된 블록을 가리키기만(Mapping) 합니다.
- 효과: 연산량과 메모리를 동시에 획기적으로 줄입니다.
즉, vLLM은 블록 단위 관리 덕분에 여러 요청이 공통된 내용(프롬프트, 빔 내역)을 가지고 있을 때 메모리를 물리적으로 공유할 수 있다는 것이 핵심입니다. 이로 인해 복잡한 디코딩 작업에서 메모리 사용량을 획기적으로 줄이고 속도를 높일 수 있다는 것을 보여주는 섹션입니다.
5. Scheduling and Preemption (4.5절)
4.5 섹션은 GPU 메모리가 꽉 찼을 때(OOM), 누구를 쫓아내고 어떻게 다시 데려올 것인가?에 대한 vLLM의 전략을 다룹니다.
핵심은 FCFS(선입선출) 원칙을 지키면서, 쫓겨난 녀석들을 CPU로 대피(Swap)시키거나 다시 계산(Recompute)하는 것입니다.
1. 스케줄링 정책: FCFS (First-Come-First-Serve)
- 원칙: 공정성을 위해 먼저 온 요청을 먼저 처리합니다.
- 메모리 부족 시: 가장 나중에 도착한 요청부터 선점(Preemption, 잠시 중단) 당합니다. 즉, 오래 기다린 요청은 계속 실행해주고, 막 들어온 신참을 잠시 대기시킵니다.
2. 쫓아내는 단위 (Eviction Policy)
- All-or-Nothing: 시퀀스의 블록 중 일부만 쫓아낼 수는 없습니다. 쫓아낼 거면 해당 시퀀스의 모든 블록을 한꺼번에 쫓아냅니다.
- Gang Scheduling: 만약 여러 시퀀스가 메모리를 공유하고 있다면(예: 빔 서치, 병렬 샘플링), 이들은 시퀀스 그룹으로 묶여서 다 같이 쫓겨나거나 다 같이 실행됩니다. (공유된 메모리 관계를 깨뜨리지 않기 위함)
3. 복구 방법 (Recovery Strategies)
쫓아낸(중단시킨) 요청을 나중에 다시 실행할 때, KV 캐시를 어떻게 복구할까요? 두 가지 방법이 있습니다.
① 스와핑 (Swapping) - 메모리 이동
- 방식: 우선순위가 밀린 요청의 KV 캐시(물리 블록)를 GPU VRAM에서 CPU RAM으로 스왑 아웃 (복사/이동)시킵니다.
- 복구: 나중에 다시 처리할 차례가 되면 CPU에서 다시 GPU로 가져옵니다. 스왑 인
- 특징: 중단된 요청들이 완료될 때까지는 새로운 요청을 받지 않고, 기존 요청 처리에 집중합니다.
② 재연산 (Recomputation)
- 방식: 선점된 요청의 KV 캐시를 그냥 삭제해 버립니다.
- 복구: 다시 차례가 오면, 처음 입력 프롬프트와 지금까지 생성했던 토큰들을 합쳐서 처음부터 다시 계산합니다.
- 한 토큰씩 생성하는 건 느리지만, 입력된 전체 토큰을 한 번에 계산하는 건 병렬 처리가 돼서 매우 빠르기 때문입니다.
즉, 요약하면 vLLM은 상황에 따라 두 가지 중 하나를 선택할 수 있습니다.
- CPU와 GPU 사이의 전송 속도(대역폭)가 빠르면 → 스와핑 유리
- GPU의 연산 속도가 엄청 빠르면 → 재연산 유리 (7.3절에서 이 둘의 성능 비교를 다룹니다.)
6. Distributed Execution (4.6절)
4.6 섹션은 모델이 너무 커서 GPU 여러 개를 써야 할 때(분산 처리), vLLM은 메모리를 어떻게 관리하는가?에 대한 내용입니다.
핵심은 메모리 관리는 '중앙 스케줄러' 혼자 다 하고, GPU들은 시키는 대로 자기 맡은 부분만 저장한다는 것입니다.
1. 텐서 병렬화 (Tensor Parallelism)
- 큰 모델은 GPU 하나에 안 들어갑니다. 그래서 Transformers에서 널리 사용되는 Megatron-LM 스타일로 모델을 쪼갭니다.
- 어텐션 헤드 분할: 예를 들어 총 16개의 어텐션 헤드가 있다면, GPU A가 1~8번 헤드, GPU B가 9~16번 헤드를 담당하는 식입니다. (SPMD 방식)
2. vLLM의 핵심 전략: 중앙 집중식 관리
분산 환경이지만, KV 캐시 관리자는 중앙 스케줄러 내부에 단 하나만 존재합니다.
- 가능한 이유는 GPU가 달라도, 처리하고 있는 입력 토큰(문장)은 똑같기 때문입니다.
- 공통 매핑(Common Mapping):
- 스케줄러는 "논리 블록 0은 물리 블록 7에 저장해라"라는 명령을 모든 GPU에게 똑같이 내립니다.
- GPU A: 자기 VRAM의 물리 블록 7번에 헤드 1~8의 KV 캐시'를 저장합니다.
- GPU B: 자기 VRAM의 물리 블록 7번에 헤드 9~16의 KV 캐시를 저장합니다.
- 즉, 블록 ID는 같지만, 그 안에 담기는 내용물의 '일부(Part)'만 다릅니다.
3. 실행 흐름
복잡한 동기화 없이 아주 단순하게 돌아갑니다.
- 스케줄러 (지시):
- 이번 턴에 처리할 [입력 토큰 ID]와 [블록 테이블(메모리 지도)]을 준비해서 모든 GPU에게 Broadcast합니다.
- GPU 워커 (실행):
- 받은 '블록 테이블'을 보고, 지정된 물리 블록 위치에서 데이터를 읽거나 씁니다.
- 연산 중간에 필요한 결과 합산(All-reduce)은 GPU들끼리 알아서 빠르게 처리합니다. (스케줄러 개입 X)
- 결과 반환:
- 연산이 끝나면 샘플링된 토큰 결과만 스케줄러에게 보고합니다.
메모리 관리(어디에 저장할지 결정)는 중앙 스케줄러가 혼자서 다 결정해서 알려준다. 각 GPU는 복잡하게 서로 상의할 필요 없이, 스케줄러가 알려준 '같은 블록 번호'에 각자 맡은 '헤드 데이터'만 넣으면 된다는 것이 4.6 분산 처리 섹션의 핵심입니다.
5. 구현
5. Implementation (구현) 섹션은 vLLM이 실제 코드로 어떻게 구축되었는지, 그리고 PagedAttention을 고성능으로 돌리기 위해 하드웨어(GPU) 레벨에서 어떤 최적화를 수행했는지를 설명하는 기술 파트입니다.
시스템 전반 아키텍처 (System Overview)
vLLM은 Python과 C++/CUDA의 하이브리드 구조로 된 엔드투엔드(End-to-End) 서빙 시스템입니다.
- 프론트엔드 (Python):
- FastAPI 기반으로 구축되어 사용자 친화적입니다.
- OpenAI API와 호환되므로 기존 애플리케이션을 쉽게 연동할 수 있습니다.
- 스케줄러, 블록 관리자 등 복잡한 제어 로직은 유연한 Python으로 작성되었습니다 (약 8.5K 라인).
- 백엔드/엔진 (C++/CUDA):
- 모델 실행과 PagedAttention 같은 핵심 연산(Data Plane)은 속도를 위해 C++와 CUDA로 작성되었습니다 (약 2K 라인).
- PyTorch와 Transformers 라이브러리를 사용하여 GPT, OPT, LLaMA 같은 주요 모델을 지원합니다.
- 분산 GPU 통신에는 NCCL(NVIDIA의 고속 통신 라이브러리)을 사용합니다.
5.1 커널 레벨 최적화 (Kernel-level Optimization)
PagedAttention은 메모리가 여기저기 흩어져 있기 때문에(비연속적), 기존 GPU 시스템으로는 효율이 떨어질 수 있습니다. 이를 해결하기 위해 3가지 핵심 맞춤형 CUDA 커널을 개발했습니다.
핵심은 Fused(융합)입니다. 여러 자잘한 작업을 하나의 커널로 합쳐서 GPU 오버헤드를 줄이는 것이 목표입니다.
① 형태 변환 및 쓰기 융합 (Fused reshape and block write)
- 문제: 새로운 KV 캐시가 생성되면, 이를 블록 단위로 쪼개고, 읽기 좋게 형태를 바꾸고(Reshape), 블록 테이블을 보고 물리 메모리에 저장해야 합니다.
- 최적화: 이 과정들을 따로 실행하면 GPU에 명령을 내리는 시간(Launch Overhead)이 낭비됩니다. 이 모든 과정을 하나의 커널로 합쳐서(Fused) 한 번에 처리합니다.
② 블록 읽기 및 어텐션 융합 (Fusing block read and attention)
- 문제: 흩어져 있는 블록들을 읽어와서 어텐션 연산을 해야 합니다.
- 최적화:
- FasterTransformer의 어텐션 커널을 개조했습니다.
- 블록 테이블을 참조하여 흩어진 데이터를 실시간으로 읽어오면서 바로 연산합니다.
- 메모리 접근 효율화 (Coalesced Memory Access): GPU의 워프(Warp, 스레드 그룹) 하나가 블록 하나를 전담해서 읽게 하여, 메모리 읽기 속도를 극대화했습니다.
- 배치 내의 요청들이 서로 길이가 달라도 처리할 수 있도록 지원합니다.
③ 블록 복사 융합 (Fused block copy)
- 문제: Copy-on-Write가 발생하면 물리 블록을 복사해야 합니다. 하지만 이 블록들은 메모리상에 불연속적으로 흩어져 있습니다. CUDA의 기본 복사 명령(cudaMemcpyAsync)을 블록마다 일일이 호출하면 오버헤드가 너무 큽니다.
- 최적화: 여러 블록에 대한 복사 명령을 모아서 단 하나의 커널 실행으로 한 번에 처리하는 전용 커널을 만들었습니다.
3. 다양한 디코딩 알고리즘 지원 (Supporting Various Decoding Algorithms)
vLLM은 복잡한 디코딩 알고리즘(빔 서치, 병렬 샘플링 등)을 구현하기 위해 3가지 주요 방법을 정의했습니다.
- Fork (분기):
- 기존 시퀀스에서 새로운 시퀀스를 만들어냅니다. (예: 프롬프트 하나에서 샘플 A, B, C를 만들 때)
- Append (추가):
- 시퀀스에 새로운 토큰을 추가합니다. (일반적인 생성 과정)
- Free (해제):
- 시퀀스를 삭제합니다. (빔 서치에서 탈락하거나 생성이 끝났을 때)
병렬 샘플링 예시:
- 입력 시퀀스 하나를 받아서 → Fork로 여러 개 출력을 만들고 → 매 반복마다 Append로 토큰을 붙이다가 → 조건이 충족되면 Free로 삭제합니다.
vLLM은 Python(편의성)과 CUDA(속도)를 결합해 만들어졌습니다. 특히 PagedAttention의 성능을 극대화하기 위해 '쓰기, 읽기+연산, 복사'라는 3가지 핵심 작업을 전용 GPU 커널로 만들어 최적화했으며, 'Fork, Append, Free'라는 단순한 3단계 명령어로 복잡한 디코딩 알고리즘들을 유연하게 지원하고 있습니다.
6. 평가
논문의 6. 평가(Evaluation) 섹션은 vLLM이 기존의 최신 시스템(SOTA)들보다 얼마나, 왜 더 뛰어난지를 실험 데이터로 증명하는 부분입니다. 요약하자면 vLLM은 메모리 낭비를 없애 배치 크기(Batch Size)를 키웠고, 덕분에 처리량(Throughput)이 압도적으로 높다는 것입니다.
1. 실험 환경 (Setup)
- 모델: OPT (13B, 66B, 175B), LLaMA (13B)
- 데이터셋:
- ShareGPT: 실제 사용자의 긴 대화 데이터 (입출력 길이가 길고 불규칙함).
- Alpaca: 짧은 명령어 데이터.
- 비교 대상 (Baselines):
- FasterTransformer: 지연 시간(Latency) 최적화 시스템 (메모리 관리 비효율적).
- Orca (Oracle): 시스템이 출력 길이를 미리 완벽하게 안다고 가정한 상태입니다. 현실적으로 불가능한 이론상 최고 성능의 Orca입니다.
- Orca (Max): 현실적인 비교군으로 최대 길이(2048)만큼 메모리를 예약하는 방식입니다.
2. 평가 결과
2-1. 기본 샘플링(Basic Sampling)
요청당 1개의 답변을 생성하는 일반적인 상황입니다.
- 결과: vLLM은 지연 시간(Latency)은 비슷하게 유지하면서 처리량(Throughput, 초당 처리 요청 수)을 크게 높였습니다.
- vs Orca (Oracle): 이론상 최고 성능인 Orca보다도 1.7배 ~ 2.7배 더 빠릅니다.
- vs Orca (Max): 현실적인 Orca보다 2.7배 ~ 8배 더 빠릅니다.
- vs FasterTransformer: 최대 22배 더 빠릅니다.
- 이유: PagedAttention이 메모리 단편화를 해결하여, 동일한 메모리에 더 많은 요청(더 큰 Batch Size)을 동시에 구겨 넣을 수 있었기 때문입니다.
2-2. 메모리 공유 효율성 (Parallel Sampling & Beam Search)
프롬프트나 중간 과정을 공유해야 하는 복잡한 디코딩 상황입니다. 즉, 메모리 공유가 필요한 시나리오인데 여기서 vLLM의 성능이 더욱 두드러집니다.
- 병렬 샘플링 & 빔 서치:
- 샘플 개수(N)나 빔 폭(k)이 커질수록 vLLM의 성능 격차가 더 벌어집니다.
- vLLM은 물리 블록을 공유(Sharing)하지만, Orca는 데이터를 복사(Copy)해야 하기 때문입니다.
- 빔 탐색(빔 너비 6)의 경우 기본 샘플링(1.3배)보다 높은 2.3배의 처리량 향상을 보였습니다.
- 메모리 절약:
- 병렬 샘플링에서 6.1%~30.5%, 빔 탐색에서 37.6%~66.3%의 메모리 절약을 달성했습니다.
2-3. 특수 워크로드
- 공유 접두사 (Shared Prefix):
- 공유되는 접두사의 길이가 길수록(5-샷 예제) vLLM은 Orca (Oracle) 대비 3.58배 높은 처리량을 달성하여, 접두사 공유의 효율성을 입증했습니다.
- 이는 프롬프트 캐시 재사용 덕분이라고 합니다.
- 챗봇 (Chatbot):
- 긴 대화 기록(History)을 가진 챗봇 환경에서도 vLLM은 Orca 대비 2배 높은 처리량을 유지했습니다.
- Orca는 긴 프롬프트에 대해 과도한 메모리를 예약하느라 낭비가 심하기 때문입니다.
- Orca의 Buddy Allocator 특성(2의 제곱수 단위 할당) 때문에 1024 토큰 길이의 프롬프트가 오면 딱 맞춰 할당하는 게 아니라 그보다 큰 공간(예: 2048)을 잡거나, 혹은 대화가 길어질 것을 대비해 미리 크게 잡아둡니다.
- vLLM은 PagedAttention이 메모리 단편화 및 예약 문제를 해결하기에 챗봇 워크로드에서 매우 효율적이라고 합니다.
7. Ablation Studies
7-1. Kernel Microbenchmark
PagedAttention 커널은 블록 테이블에 엑세스하고, 추가 분기를 실행하고, 가변 시퀀스 길이를 처리하는 추가 오버헤드가 포함되기에 FasterTransformer 대비 20%~26% 높은 어텐션 커널 지연 시간을 보입니다. 하지만, 이러한 오버헤드는 어텐션 연산자에만 영향을 미치기에 전반적인 시스템 성능 향상에 미미한 영향을 주었습니다.
이에 오버헤드에도 불구하고 PagedAttention은 vLLM이 엔드 투 엔드 성능에서 FasterTransformer을 능가합니다.
7-2. Impact of Block Size(블록 크기의 영향)
블록 크기의 선택은 vLLM의 성능에 상당한 영향을 끼친다고 합니다. 블록 크기가 너무 작으면 vLLM이 KV cache를 읽고 처리하기 위한 GPU의 병렬 처리를 완전히 활용하지 못할 수 있습니다. 블록 크기가 너무 크면 내부 단편화가 증가하고 공유 가능성이 감소하기 때문입니다.
실험 결과, 실제로 블록 크기 16이 GPU를 효율적으로 활용하기에 충분히 크고 대부분의 워크로드에서 상당한 내부 단편화를 피할 수 있을 만큼 충분히 작다고 합니다. 기본 블록 크기도 16이라고 합니다.
블록 크기는 16이 대부분의 워크로드에서 최적의 성능을 제공했습니다. 재계산은 작은 블록 크기에서 효율적이고, 스왑은 큰 블록 크기에서 효율적인 경향을 보였으나, 전반적으로 두 방법 모두 중간 블록 크기에서 유사한 성능을 보였습니다.
7-3. Comparing Recomputation and Swapping (재연산과 스와핑 비교)
vLLM은 복구 메커니즘으로 재연산과 스와핑을 모두 지원합니다.
왼쪽 그림의 결과를 보면 스와핑이 작은 블록 크기에서 과도한 오버헤드를 발생시키고 있는 것을 확인할 수 있습니다. 이는 작은 블록 크기가 CPU와 GPU간의 수 많은 작은 데이터 전송을 초래하기 때문입니다.
반대로, 재연산은 메모리를 이동시키는 것이 아니라 연산을 다시 수행하는 것이므로, 메모리 블록 크기 자체에는 큰 영향을 받지 않고 일정합니다.
따라서 블록 크기가 작을 때는 재연산이 더 효율적이고, 블록 크기가 클 때는 스와핑이 더 효율적이지만, 재연산 오버헤드는 스와핑 지연 시간의 20%를 넘지 않습니다. 16에서 64 사이의 중간 블록 크기의 경우 두 방법 모두 비슷한 엔드 투 엔드 성능을 확인할 수 있습니다.
8 + 9. 토론과 관련 연구는 생략 하겠습니다.
10. 결론
본 논문은 PagedAttention 알고리즘과 이를 기반으로 구축된 vLLM 시스템을 통해 LLM 서빙의 핵심 과제인 KV 캐시 메모리 비효율성을 효과적으로 해결했습니다. 운영체제에서 영감을 받은 가상 메모리 및 쓰기 시점 복사 기법을 활용하여 KV 캐시 메모리 낭비를 거의 없애고 유연한 공유를 가능하게 함으로써, vLLM은 최첨단 시스템 대비 2~4배 높은 처리량을 달성하며, 특히 긴 시퀀스, 대규모 모델, 복잡한 디코딩 알고리즘에서 더욱 큰 성능 향상을 보였다고 결론 지으며 논문이 마무리 됩니다.
사실 vLLM을 처음 접했을 때는 원리보다는 당장의 성능 개선이 급해 무작정 적용부터 했었습니다. 그런데 막상 써보니 지연 시간과 메모리 효율이 너무 좋더라구요. 어떤 차이가 있을까?하는 호기심에 논문을 펼쳤다가, PagedAttention 기술을 흥미롭게 본 기억이 있었습니다.
그런데 따로 기록에 남겨둔 것은 없기도하고, 다시 읽어보니 기억이 잘 안나는 부분도 있어서 이렇게 자세한 리뷰까지 남기게 되었습니다. PagedAttention만 간단히 짚고 넘어갈 수도 있었지만, 파고들수록 더 많은 인사이트를 얻어가실 수 있으실 것 같아 욕심내어 길게 정리해 보았습니다. 긴 글 읽어주셔서 감사드리며, 이 글이 vLLM을 이해하는 데 조금이나마 도움이 되셨으면 좋겠습니다.