PagedAttention: LLM 서빙의 메모리 관리 (vLLM)

PagedAttention: LLM 서빙의 메모리 관리 (vLLM)

LLM을 서비스할 때 GPU 메모리를 가장 많이 낭비하는 곳은 요청마다 생기는 KV 캐시입니다. 이 논문은 운영체제의 가상 메모리 · 페이징처럼 KV 캐시를 작은 고정 크기 블록으로 나눠 흩어진 메모리에 저장하는 어텐션 알고리즘 PagedAttention과, 이를 바탕으로 한 서빙 엔진 vLLM을 제안합니다. 메모리 낭비를 거의 없애고 블록 단위 공유까지 가능하게 해, 같은 지연 시간에서 처리량을 기존 최고 시스템보다 2~4배 높였습니다.

들어가기 전에

이 해설은 아래 내용을 안다고 보고 씁니다. 처음이라면 먼저 읽고 오면 수식과 손계산을 훨씬 쉽게 따라올 수 있습니다.

• 어텐션만 있으면 된다 (트랜스포머) — 순환 없이 어텐션만으로 번역 모델을 만든 트랜스포머 — 스케일드 닷-프로덕트 · 멀티헤드 어텐션 · 위치 인코딩의 수식과 의미
• 추측 디코딩으로 트랜스포머 추론 빠르게 하기 — 작은 모델의 추측을 큰 모델이 병렬 검증해 출력은 그대로 추론을 2~3배 빠르게 한 추측 디코딩

한눈에 보기

• LLM 서빙의 처리량이 왜 계산이 아니라 GPU 메모리(KV 캐시)에 막히는지 숫자로 알게 됩니다.
• 기존 시스템이 KV 캐시를 연속 공간에 미리 잡아 두며 생기는 세 가지 낭비(예약 · 내부 단편화 · 외부 단편화)를 알게 됩니다.
• 운영체제 페이징처럼 "논리 블록 → 물리 블록" 표로 KV 캐시를 관리하는 방법을 직접 따라가 봅니다.
• 병렬 샘플링 · 빔 탐색 · 공통 프롬프트에서 블록을 공유해 메모리를 아끼는 방법과 그 효과를 알게 됩니다.

초록 (번역)

대형 언어 모델(LLM)을 높은 처리량으로 서비스하려면 한 번에 충분히 많은 요청을 묶어(배칭) 처리해야 합니다. 그러나 기존 시스템은 요청마다 키-값 캐시(KV 캐시) 메모리가 매우 크고 동적으로 늘었다 줄었다 하기 때문에 어려움을 겪습니다. 비효율적으로 관리하면 이 메모리는 단편화와 불필요한 중복으로 크게 낭비되어 배치 크기를 제한합니다. 이 문제를 해결하기 위해, 우리는 운영체제의 고전적인 가상 메모리와 페이징 기법에서 영감을 받은 어텐션 알고리즘 PagedAttention을 제안합니다. 그 위에 (1) KV 캐시 메모리 낭비가 거의 0이고 (2) 요청 안에서와 요청들 사이에서 KV 캐시를 유연하게 공유해 메모리 사용을 더 줄이는 LLM 서빙 시스템 vLLM을 만듭니다. 평가 결과 vLLM은 FasterTransformer · Orca 같은 최고 수준 시스템에 비해 같은 수준의 지연 시간에서 널리 쓰이는 LLM들의 처리량을 2~4배 높입니다. 이 향상은 시퀀스가 길수록, 모델이 클수록, 디코딩 알고리즘이 복잡할수록 더 두드러집니다. vLLM의 소스 코드는 공개돼 있습니다.

무엇을 풀었나

LLM은 토큰을 하나씩 순서대로 만듭니다. 이 생성 단계는 행렬-벡터 곱이 많아 메모리 병목이고 GPU 계산 능력을 다 쓰지 못합니다. 처리량을 높이는 가장 확실한 방법은 여러 요청을 한꺼번에 배칭하는 것인데, 그러려면 요청마다 필요한 메모리를 잘 관리해야 합니다.

KV 캐시가 큽니다. 13B 모델을 A100 40GB에서 서비스할 때 메모리의 약 65%는 모델 가중치, 30% 가까이는 요청들의 동적 상태인 KV 캐시가 차지합니다. OPT-13B는 토큰 하나의 KV 캐시가 약 800KB이고, 최대 2,048토큰 요청 하나에 1.6GB까지 필요합니다.

기존 시스템은 메모리를 많이 버렸습니다. 대부분의 딥러닝 프레임워크는 텐서를 연속 메모리에 둬야 해서, 기존 서빙 시스템은 요청마다 최대 길이만큼의 연속 공간을 미리 잡았습니다. 출력 길이를 미리 알 수 없으니 생긴 일입니다. 논문 측정에서 기존 시스템은 KV 캐시 메모리의 20.4~38.2%만 실제 토큰 상태 저장에 썼습니다.

공유도 못 했습니다. 한 프롬프트로 여러 답을 뽑는 병렬 샘플링이나 빔 탐색은 KV 캐시를 일부 공유할 수 있는데, 연속 공간에 따로따로 저장하니 공유가 불가능했습니다.

전체 흐름

123456789flowchart LR
  A[새 요청 도착] --> B[스케줄러: 선착순 FCFS<br/>배치에 넣을 요청 선택]
  B --> C[KV 캐시 관리자<br/>필요한 만큼만 물리 블록 할당]
  C --> D[블록 표: 논리 블록 → 물리 블록<br/>+ 채워진 칸 수]
  D --> E[PagedAttention 커널<br/>흩어진 블록을 읽어 어텐션 계산]
  E --> F[새 토큰의 K · V를<br/>마지막 블록 빈칸에 저장]
  F -->|블록이 차면| C
  F -->|메모리 부족| G[나중 요청부터 선점<br/>CPU로 스왑 또는 재계산]
  F -->|요청 끝| H[블록 반납 → 다른 요청이 사용]

핵심 아이디어

1. KV 캐시가 왜 그렇게 큰가

한 문장 직관: 대화가 길어질수록 모델이 "지금까지 읽은 내용 메모(키 · 값)"를 토큰마다 층마다 쌓아 두기 때문입니다.

단계별 설명

1. 자기회귀 생성은 프롬프트 단계(입력 전체를 한 번에 처리)와 생성 단계(토큰을 하나씩 생성)로 나뉩니다.
2. 생성 단계마다 새 토큰의 질의는 이전 모든 토큰의 키 · 값과 어텐션해야 합니다.
3. 매번 다시 계산하지 않으려고 이전 토큰들의 키 · 값을 저장해 둡니다. 이것이 KV 캐시입니다.
4. KV 캐시는 생성이 진행될수록 커지고, 요청이 끝나면 사라지며, 최종 길이를 미리 알 수 없습니다.

숫자로 해 보기 (논문 계산 따라 하기): OPT-13B는 은닉 크기 5,120, 층 40개, FP16(2바이트)입니다. 토큰 하나의 KV 캐시는 $2(\text{키·값}) \times 5{,}120 \times 40 \times 2 = 819{,}200$바이트, 약 800KB입니다. 2,048토큰이면 약 1.6GB입니다. GPU 메모리가 수십 GB라면 KV 캐시에 다 줘도 요청 수십 개밖에 못 담습니다. 논문은 A100에서 H100으로 가며 FLOPS는 2배 넘게 늘었지만 GPU 메모리는 최대 80GB 그대로라, 메모리가 점점 더 큰 병목이 될 것이라고 봅니다.

2. PagedAttention — 흩어진 블록으로 어텐션 계산

한 문장 직관: 책 한 권을 서가 한 칸에 통째로 꽂을 자리가 없어도, 장(章)별로 빈 칸 아무 데나 꽂아 두고 "목차 표"로 찾아 읽으면 됩니다.

단계별 설명

1. 각 시퀀스의 KV 캐시를 KV 블록으로 나눕니다. 블록 하나에는 정해진 수 $B$개 토큰의 키 · 값이 들어갑니다.
2. 블록들은 물리 메모리에서 연속일 필요가 없습니다.
3. 어텐션을 블록 단위로 나눠 계산합니다. 커널은 블록 표를 보고 블록을 하나씩 찾아 읽어 점수와 출력을 계산합니다.

수학적으로는 보통 어텐션과 같은 값이고, 저장 위치만 흩어져 있다는 점이 다릅니다.

3. KV 캐시 관리자 — 논리 블록과 물리 블록

한 문장 직관: 운영체제가 프로그램에게 "연속된 것처럼 보이는" 가상 주소를 주고 실제로는 흩어진 물리 페이지에 담듯, 요청에게 연속된 논리 블록을 주고 실제로는 아무 물리 블록에나 담습니다.

단계별 설명

1. GPU에 큰 메모리 덩어리를 잡아 물리 KV 블록들로 나눠 둡니다(CPU RAM에도 스왑용으로 같은 방식).
2. 요청마다 블록 표에 "논리 블록 번호 → 물리 블록 번호 + 채워진 칸 수"를 기록합니다.
3. 처음에는 프롬프트에 필요한 블록만 할당합니다. 최대 길이를 미리 잡지 않습니다.
4. 마지막 블록이 다 차면 그때 새 물리 블록 하나를 받아 이어 붙입니다.
5. 요청이 끝나면 블록을 모두 반납해 다른 요청이 씁니다.

숫자로 해 보기 (논문 그림의 예를 따라, 블록 크기 4)

1. 프롬프트 "Four score and seven years ago our"는 7토큰입니다. 논리 블록 0(4칸 가득), 논리 블록 1(3칸)이 필요해 물리 블록 7번과 1번을 받습니다. 블록 1에 빈칸 1개가 남습니다.
2. 첫 생성 토큰 "fathers"는 블록 1의 빈칸에 들어가고, 채워진 칸 수가 4가 됩니다.
3. 둘째 생성 토큰 "brought"는 자리가 없으니 새 물리 블록(3번)을 받아 논리 블록 2로 이어 붙입니다.

낭비는 블록 하나 이내: 블록을 왼쪽부터 채우고 앞 블록이 다 차야 새 블록을 받으니, 요청 하나의 낭비는 마지막 블록의 빈칸뿐입니다.

숫자로 해 보기 (해설용 계산): 최대 길이 2,048로 잡는 방식에서 실제로 100토큰만 쓰면 $100/2{,}048 \approx 4.9\%$만 쓰고 나머지는 낭비입니다. 블록 크기 16이면 100토큰은 블록 7개(112칸)로, 낭비는 12칸뿐입니다.

4. 블록 공유 — 병렬 샘플링 · 빔 탐색 · 공통 프롬프트

한 문장 직관: 여러 답안이 같은 문제지(프롬프트)를 보고 있다면 문제지는 한 장만 두고 함께 보다가, 누가 그 위에 뭔가 새로 쓰려 할 때만 복사본을 만들어 줍니다.

단계별 설명

1. 참조 횟수: 물리 블록마다 몇 개의 논리 블록이 가리키는지 셉니다.
2. 병렬 샘플링: 한 프롬프트에서 여러 출력을 뽑으면 프롬프트 블록은 모두가 같은 물리 블록을 공유합니다(참조 횟수 2 이상).
3. 쓰기 시 복사(copy-on-write): 어떤 시퀀스가 공유 블록에 새 토큰을 써야 하면, 참조 횟수가 1보다 클 때만 새 블록을 받아 복사한 뒤 씁니다. 운영체제의 프로세스 복제(fork)와 같은 방식입니다.
4. 빔 탐색: 후보들이 프롬프트뿐 아니라 중간 블록까지 공유하고, 공유 관계가 계속 바뀝니다. 탈락한 후보의 블록은 참조 횟수를 줄이고, 0이 되면 반납합니다. 기존 시스템처럼 큰 KV 캐시를 통째로 복사할 필요가 없습니다.
5. 공통 접두부(시스템 프롬프트): 서비스 제공자가 자주 쓰는 지시문 · 예시의 KV 캐시를 물리 블록으로 미리 만들어 두면, 사용자 요청은 그 블록을 가리키기만 하고 자기 입력 부분만 계산하면 됩니다. 운영체제의 공유 라이브러리와 같습니다.

숫자로 해 보기 (해설용 계산): 1,000토큰 프롬프트로 출력 200토큰짜리 답 4개를 뽑는다면, 공유 없이는 $4 \times 1{,}200 = 4{,}800$칸, 프롬프트를 공유하면 약 $1{,}000 + 4 \times 200 = 1{,}800$칸입니다(쓰기 시 복사되는 마지막 블록 몇 칸은 무시한 대략값). 약 62.5%를 아낍니다.

5. 스케줄링과 선점 — 메모리가 모자랄 때

한 문장 직관: 손님이 넘치면 먼저 온 손님부터 모시고, 자리가 모자라면 가장 늦게 온 손님에게 잠시 대기를 부탁합니다.

• 선착순(FCFS): 공정성을 위해 먼저 온 요청을 먼저 처리하고, 선점이 필요하면 가장 늦게 온 요청부터 내보냅니다.
• 전부 아니면 전무: 한 시퀀스의 블록은 늘 함께 쓰이므로 통째로 내보내거나 통째로 둡니다. 빔 후보처럼 메모리를 공유하는 시퀀스 묶음은 함께 선점 · 재개합니다.
• 복구 방법 두 가지: ① 스왑 — 내보낸 블록을 CPU 메모리로 옮겼다가 다시 가져옵니다. ② 재계산 — 이미 생성한 토큰을 프롬프트에 붙여 한 번의 프롬프트 단계로 KV 캐시를 다시 만듭니다.
• 논문 측정에서 블록이 작으면 CPU-GPU 사이 작은 전송이 많아 스왑이 느렸고, 재계산 비용은 블록 크기와 상관없이 일정했습니다. 블록 크기 16~64에서는 둘이 비슷했습니다.

6. 분산 실행과 구현

• 모델이 GPU 한 장에 안 들어가면 Megatron-LM 방식 텐서 병렬화(어텐션 헤드 단위로 나눔)를 씁니다. 모든 GPU가 같은 위치의 KV 캐시가 필요하므로, 중앙 스케줄러의 KV 캐시 관리자 하나가 블록 표를 관리하고 각 GPU는 자기 헤드 몫만 저장합니다.
• 구현은 Python 8,500줄 + C++/CUDA 2,000줄이며, FastAPI 앞단이 OpenAI API 형식을 지원합니다.
• 커널 최적화: 블록 쓰기와 형태 변환을 하나의 커널로 합치기, 블록 읽기와 어텐션 합치기(블록마다 GPU 워프 하나), 쓰기 시 복사용 블록 복사를 한 번에 묶기
• 디코딩 알고리즘은 fork(시퀀스 복제) · append(토큰 추가) · free(삭제) 세 연산의 조합으로 구현합니다.

결과

실험 설정: OPT 13B(A100 1장) · 66B(A100 4장) · 175B(A100-80GB 8장)와 LLaMA 13B, Google Cloud A2 인스턴스. 작업 부하는 실제 서비스 대화 모음 ShareGPT(입력 평균 161.31토큰, 출력 평균 337.99토큰)와 Alpaca(입력 19.31, 출력 58.45). 지표는 정규화 지연(요청 지연 ÷ 출력 길이)이 낮게 유지되는 최대 요청률입니다.

비교 대상: FasterTransformer, 그리고 출력 공간을 얼마나 미리 잡는지에 따라 Orca 세 가지 — Orca(Max: 항상 최대 2,048토큰 예약), Orca(Pow2: 최대 2배까지 과잉 예약, 예: 실제 25면 32), Orca(Oracle: 실제 출력 길이를 미리 안다고 가정한 이론상 상한)

KV 캐시 메모리 중 실제 토큰 상태 비율(그림 2): 기존 시스템 20.4~38.2% → vLLM 96.3%

기본 샘플링(요청당 출력 1개)

• OPT-13B · ShareGPT에서 vLLM은 Orca(Oracle)보다 동시에 2.2배, Orca(Max)보다 4.3배 많은 요청을 배치로 묶었습니다(평균 배치 요청 수 30.42 vs 13.62 · 7.00).
• 예외로, OPT-175B · Alpaca처럼 KV 캐시용 메모리가 크고 시퀀스가 짧은 경우에는 Orca(Oracle · Pow2)와의 차이가 작았습니다. 이때는 시스템이 메모리보다 계산에 막히기 때문입니다.

블록 공유 효과(OPT-13B)

Alpaca에서 Orca(Oracle) 대비 향상은 기본 샘플링 1.3배에서 빔 폭 6의 빔 탐색 2.3배로 커졌습니다.

공통 접두부(LLaMA-13B, 영→독 번역): 예시 1개(80토큰) 접두부를 공유하면 Orca(Oracle)보다 처리량 1.67배, 예시 5개(341토큰)면 3.58배였습니다.

챗봇(ShareGPT, 프롬프트 최대 1,024토큰 · 출력 최대 1,024토큰): Orca 세 가지보다 2배 높은 요청률을 감당했습니다.

비용과 선택: 블록 표 조회 · 분기 · 가변 길이 처리 때문에 어텐션 커널 자체는 FasterTransformer보다 20~26% 느렸지만, 어텐션 연산에만 해당해 전체 성능은 훨씬 좋았습니다. 블록 크기는 ShareGPT에서 16~128이 가장 좋았고, 짧은 시퀀스의 Alpaca에서는 16 · 32가 좋았습니다. 그래서 기본값을 16으로 정했습니다.

왜 중요한가 · 이후 영향

• LLM 서빙의 표준: vLLM은 공개 소스로 빠르게 퍼져, 오픈 모델을 서비스하는 대표적인 엔진이 됐습니다. 같은 섹션의 Mixtral 논문도 vLLM에 MoE 커널을 통합했다고 밝힙니다.
• 운영체제 아이디어의 재발견: 가상 메모리 · 페이징 · 쓰기 시 복사 · 공유 라이브러리 같은 수십 년 된 기법이 LLM 서빙에 그대로 들어맞는다는 것을 보였습니다.
• 비용 절감: 같은 GPU로 더 많은 요청을 묶어 처리하니, 토큰당 서빙 비용이 크게 줄어듭니다. 논문은 LLM 요청 처리 비용이 일반 키워드 검색보다 10배 비쌀 수 있다는 추정을 들며 처리량의 중요성을 강조합니다.
• 복잡한 디코딩의 실용화: 병렬 샘플링 · 빔 탐색 · 긴 시스템 프롬프트 공유가 메모리 부담 없이 가능해졌습니다.

한계 · 주의할 점

• 커널 자체는 조금 느립니다. 블록을 찾아 읽는 추가 비용으로 어텐션 커널 지연이 20~26% 늘었습니다.
• 블록 크기 선택: 너무 작으면 GPU 병렬성을 못 쓰고 스왑이 느려지며, 너무 크면 내부 단편화가 늘고 공유 가능성이 줄어듭니다. 작업 부하에 따라 최적값이 다릅니다.
• 메모리 병목 상황에서 효과가 큽니다. 메모리가 넉넉하고 시퀀스가 짧아 계산이 병목이면 이득이 작아집니다.
• 다른 작업에 일반화되지 않을 수 있습니다. 저자들은 페이징이 LLM 서빙처럼 동적 메모리 할당이 필요하고 메모리 용량에 막힌 작업에 효과적이며, 모든 GPU 작업에 일반적으로 맞는 것은 아니라고 말합니다.
• 대화 회차 사이 캐시는 저장하지 않았습니다. 챗봇 실험에서 회차 사이 KV 캐시는 다른 요청의 자리를 차지하므로 보관하지 않았습니다.

이어서 읽을 것

• 앞: 어텐션만 있으면 된다 (트랜스포머) · 플래시어텐션
• 함께: 추측 디코딩으로 트랜스포머 추론 빠르게 하기 · 플래시어텐션-2 · Mixtral: 전문가 혼합 모델

용어 정리

스스로 점검

1. 기존 LLM 서빙 시스템이 KV 캐시 메모리를 낭비한 세 가지 이유는 무엇인가요?

   요청마다 최대 길이만큼 연속 공간을 미리 잡아서 생기는 미래 토큰 예약, 끝까지 쓰이지 않는 내부 단편화, 크기가 서로 다른 덩어리 사이의 외부 단편화입니다. 그 결과 실제 토큰 상태에 쓰인 메모리는 20.4~38.2%뿐이었습니다.

2. PagedAttention에서 요청 하나의 메모리 낭비가 블록 하나 이내로 묶이는 이유는 무엇인가요?

   블록을 왼쪽부터 채우고 앞 블록이 다 찬 뒤에야 새 물리 블록을 받기 때문에, 비어 있는 칸은 마지막 블록에만 생깁니다.

3. 빔 탐색에서 vLLM이 메모리를 크게 아끼는 원리는 무엇인가요?

   후보들이 공통으로 가진 앞부분 블록을 참조 횟수로 공유하고, 공유 블록에 새로 써야 할 때만 쓰기 시 복사로 블록 하나를 복사합니다. 탈락한 후보의 블록은 참조 횟수가 0이 되면 반납합니다. Alpaca에서 37.6~55.2%를 아꼈습니다.

다음에 볼 것

이 글을 이해했다면 이어서 보면 좋은 항목입니다.

• 직접 해 보기: 양자화 오차와 KV 캐시 메모리 계산하기 — 양자화 오차와 KV 캐시 크기를 코드로 직접 재 보는 실습
• Mamba: 선택적 상태 공간 모델로 선형 시간 시퀀스 모델링 — 입력에 따라 기억·망각을 고르는 선택적 SSM으로 트랜스포머급 성능을 선형 시간에 낸 Mamba

출처

• 원문: arxiv.org/abs/2309.06180 (2023 · Woosuk Kwon 외)