LLM 서빙의 메모리 문제와 PagedAttention (2) - PagedAttention과 vLLM

들어가며

이전 포스트에서 기존 LLM 서빙의 메모리 낭비 문제를 확인했습니다. 미리 점유, 내부/외부 단편화, 공유 불가라는 네 가지 문제가 GPU 활용률을 낮추고 동시 처리 요청 수를 제한합니다.

vLLM은 이 문제를 OS의 가상 메모리(Virtual Memory) 개념을 빌려와서 해결했습니다. OS가 물리 메모리를 고정 크기 페이지로 나눠 관리하듯, KV 캐시도 고정 크기 블록으로 나눠 비연속 공간에 저장하는 방법이 바로 PagedAttention입니다.

 

1. PagedAttention의 핵심 아이디어

기존 방식은 요청마다 연속된 메모리 공간에 KV 캐시를 저장했습니다. PagedAttention은 이 제약을 없앱니다.

• KV 캐시를 고정 크기 블록(block) 단위로 분할
• 각 블록은 연속할 필요 없이 GPU 메모리 어디든 배치 가능
• 블록 테이블(Block Table) 이 논리 블록 번호 → 물리 블록 번호 매핑 관리

OS의 페이지 테이블과 완전히 같은 구조입니다. 해당 figure에서 논리 블록은 가독성을 위한 것이고 실제로는 block table의 인덱스를 논리 블록으로 인식합니다. 덕분에 물리 메모리가 비연속으로 흩어져 있어도 논리적으로는 연속된 KV 캐시처럼 접근할 수 있습니다.

 

2. KV 캐시 매니저와 블록 할당

vLLM의 KV 캐시 매니저는 물리 블록을 관리하는 운영체제 역할을 합니다.

• GPU 메모리를 시작할 때 고정 크기 블록으로 분할 (보통 블록당 16토큰)
• 요청이 들어오면 필요한 만큼만 블록 할당 (최대 길이 기준 사전 할당 없음)
• 생성이 진행되면서 필요한 블록을 on-demand로 추가 할당
• 요청이 끝나면 해당 블록들을 즉시 반환

이 방식으로 미리 점유와 내부 단편화 문제가 해결됩니다. 블록 단위로만 약간의 낭비가 생길 수 있지만(마지막 블록이 꽉 차지 않는 경우), 기존 방식에 비하면 무시할 수준입니다.

 

3. Parallel Sampling — Copy-on-Write

같은 프롬프트에서 여러 출력을 동시에 생성하는 parallel sampling을 생각해보겠습니다. 기존 방식은 프롬프트 KV 캐시를 샘플 수만큼 복사했습니다.

PagedAttention은 OS의 Copy-on-Write(CoW) 방식을 적용합니다.

샘플 1 블록 테이블:  논리 0 → 물리 5 (ref count: 3)
샘플 2 블록 테이블:  논리 0 → 물리 5 (같은 블록 공유)
샘플 3 블록 테이블:  논리 0 → 물리 5 (같은 블록 공유)

→ 샘플 1이 블록 5를 수정하려 할 때
  ref count > 1 이므로 새 블록(물리 9)에 복사 후 수정
  샘플 1 블록 테이블: 논리 0 → 물리 9
  샘플 2, 3: 계속 물리 5 공유

프롬프트 부분(공통 prefix)은 모든 샘플이 같은 물리 블록을 가리키므로 중복 저장 없이 공유됩니다. 출력이 갈라지는 시점부터만 별도 블록을 할당합니다.

 

4. Beam Search — 트리 구조 공유

Beam search는 상위 k개 후보를 유지하면서 생성하는 방법입니다. 각 후보는 이전 단계의 특정 후보에서 파생됩니다

PagedAttention은 이 트리 구조를 그대로 물리 블록 공유로 표현합니다. 각 빔이 fork되기 전의 블록들은 ref count로 공유되고 갈라진 이후의 블록만 독립적으로 할당됩니다. Beam이 pruning되면 해당 블록들의 ref count가 감소하고 0이 되면 즉시 반환됩니다.

 

5. Shared Prefix 캐싱

시스템 프롬프트처럼 많은 요청이 공통 prefix를 가지는 경우 prefix의 KV 캐시를 한 번만 계산하고 재사용할 수 있습니다.

vLLM은 이 prefix 블록들을 캐싱해두고 동일 prefix가 들어오는 요청에 재사용합니다. Prefill 비용과 메모리를 모두 절약합니다.

 

6. 선점과 메모리 압박 처리 — Swapping과 Recomputation

GPU 메모리가 부족해지면 일부 요청을 선점(preemption) 해야 합니다. vLLM은 두 가지 방법을 제공합니다.

Swapping

KV 캐시 블록을 CPU 메모리로 이동시킵니다.

GPU 메모리: [요청A 블록] [요청B 블록] ← 가득 참
→ 우선순위 낮은 요청C 블록을 CPU로 스왑
→ 새 요청 D를 위한 GPU 공간 확보
→ 나중에 요청C 재개 시 CPU 블록을 GPU로 다시 스왑인

블록 단위 관리 덕분에 스왑 단위도 블록 단위입니다. OS의 페이지 스왑과 동일한 개념입니다.

Recomputation

KV 캐시를 버리, 나중에 필요할 때 처음부터 다시 계산합니다. Swapping보다 메모리 전송 비용은 없지만 연산 비용이 발생합니다. 시퀀스가 짧거나 CPU 대역폭이 병목일 때 유리합니다.

 

7. 분산 환경 — Tensor Parallelism

vLLM은 Megatron-LM 스타일의 텐서 병렬화(Tensor Parallelism)를 지원합니다. 모델을 여러 GPU에 나눠 올리는 방식입니다.

GPU 0: 어텐션 헤드 절반 + FFN 절반
GPU 1: 어텐션 헤드 절반 + FFN 절반
→ All-reduce로 결과 합산

KV 캐시도 각 GPU에 분산됩니다. vLLM의 CPU에 있는 중앙 스케줄러가 모든 GPU의 KV 캐시 상태를 통합 관리하고, 어느 블록이 어느 GPU에 있는지 추적합니다.

[중앙 스케줄러 + KV 캐시 매니저]
        ↓              ↓
   [GPU Worker 0]  [GPU Worker 1]
   KV 캐시 일부    KV 캐시 일부

SPMD(Single Program Multiple Data) 방식으로 모든 워커가 같은 코드를 실행하고 스케줄러가 보내는 명령에 따라 동기화됩니다.

 

8. vLLM 전체 구조

vLLM의 아키텍처를 한 장으로 정리하겠습니다.

사용자 요청
    ↓
[스케줄러]
 - 어떤 요청을 이번 스텝에 처리할지 결정
 - KV 캐시 매니저와 협력해 블록 할당/반환
    ↓
[KV 캐시 매니저]
 - 물리 블록 풀 관리
 - 블록 테이블 유지
 - ref count 기반 공유/CoW 처리
    ↓
[GPU 워커들]
 - PagedAttention으로 비연속 KV 캐시 어텐션 계산
 - 텐서 병렬화 적용
    ↓
출력 토큰

 

정리

기술	해결한 문제	OS 유사 개념
블록 단위 KV 캐시	외부 단편화, 미리 점유	페이지 가상 메모리
블록 테이블	비연속 메모리 접근	페이지 테이블
On-demand 할당	내부 단편화	지연 할당
Copy-on-Write	parallel sampling 중복 저장	CoW fork
Swapping	GPU 메모리 압박	페이지 스왑
Shared prefix	공통 프롬프트 중복 계산	공유 메모리

vLLM의 핵심 통찰은 LLM 서빙의 메모리 관리 문제가 OS의 메모리 관리 문제와 구조적으로 같다는 것입니다. 30년간 OS가 발전시켜온 기법들을 GPU KV 캐시에 적용해서 성능 향상을 이뤄냈습니다.

 

출처

• Kwon et al., "Efficient Memory Management for Large Language Model Serving with PagedAttention", SOSP 2023
• vLLM GitHub: https://github.com/vllm-project/vllm
• vLLM Blog: https://vllm.ai/blog/2023/06/20/vllm.html