LLM 서빙의 메모리 문제와 PagedAttention (1) - KV 캐시와 단편화

들어가며

ChatGPT 같은 LLM 서비스를 운영한다고 생각해보겠습니다. 수천 명의 사용자가 동시에 요청을 보내는데 추론 과정에서 생성되는 중간 데이터가 메모리를 크게 잡아먹기 때문에 중간에 데이터 정체가 일어납니다. 이 포스트에서는 그 중간 데이터의 정체인 KV 캐시가 무엇인지, 그리고 기존 방식이 왜 메모리를 낭비하는지 알아보겠습니다.

 

1. KV 캐시란 무엇인가

트랜스포머의 셀프 어텐션은 다음 수식으로 동작합니다.

Attention(Q, K, V) = softmax(QK^T / √d_k) V

토큰을 하나씩 생성하는 자기회귀(autoregressive) 방식에서는 새 토큰을 생성할 때마다 이전 토큰들의 K, V를 다시 계산해야 합니다. 같은 계산을 반복하는 것은 낭비이므로, 이미 계산한 K, V를 GPU 메모리에 저장해두고 재사용하는 것이 KV 캐시입니다.

KV 캐시 덕분에 매 스텝에서 새 토큰 하나의 Q만 계산하고, 저장된 K, V와 어텐션을 계산하면 됩니다. 속도는 빨라지지만, 메모리 사용량이 증가한다는 트레이드오프가 생깁니다.

 

2. Prefill과 Decode

LLM 추론은 크게 두 단계로 나뉩니다.

페이즈	설명	특성
Prefill	입력 프롬프트 전체를 한 번에 처리	Compute-bound (연산이 병목)
Decode	토큰을 한 개씩 생성	Memory-bound (메모리 대역폭이 병목)

 

Decode 단계에서는 매 스텝마다 GPU의 모든 가중치를 메모리에서 읽어야 합니다. 처리하는 연산량(FLOPs)에 비해 메모리 읽기 비용이 압도적으로 크기 때문에 메모리 대역폭이 병목이 됩니다.

이 때문에 Decode 단계에서 배치 크기(batch size)를 키워서 한 번 메모리를 읽을 때 더 많은 요청을 처리하면 연산량을 늘려서 메모리 읽기 비용 대비 연산량을 늘릴 수 있습니다.

 

3. 배치를 키워야 하는 이유

GPU는 병렬 연산에 특화되어 있습니다. Decode 단계에서 배치 크기 1과 배치 크기 32는 메모리를 읽는 시간은 거의 같지만, 처리하는 요청 수는 32배가 되기에 배치를 키울수록 GPU 연산 효율이 올라갑니다. 문제는 배치를 키우려면 각 요청의 KV 캐시를 모두 메모리에 올려야 한다는 것입니다. 이때 HBM에서 SRAM에 올려야 하는 횟수가 빈번해진다는 문제가 있습니다.

OPT-13B 모델 기준으로 토큰 하나당 KV 캐시 크기를 계산해보면 아래와 같기에 하나의 토큰을 처리하기 위해서는 배치 사이즈 이상의 메모리 대역폭이 필요해집니다. 그러므로 연산량을 늘리기 위해서 배치를 늘렸지만 여전히 메모리 병목 문제는 존재합니다. 또한 절대적 용량의 문제도 있습니다.

레이어 수(40) × 헤드 수(40) × 헤드 차원(128) × 2(K,V) × 2bytes(FP16)
= 약 800KB / 토큰

A100 HBM의 경우 KV 캐시가 약 30%를 차지합니다. 시퀀스 길이가 길어질수록 KV 캐시는 빠르게 불어납니다. 게다가 기술의 발전에 따라 GPU FLOPs는 빠르게 늘어나는데 HBM 용량은 80GB 대에서 정체되어 있기 때문에 대역폭 뿐만 아니라 절대적 메모리 용량의 효율화도 필요한 상황입니다.

 

4. 기존 메모리 관리의 세 가지 낭비

기존 LLM 서빙 프레임워크(FasterTransformer, Orca 등)는 각 요청에 대해 최대 시퀀스 길이만큼 메모리를 미리 할당했습니다. 이 방식은 세 가지 문제를 만듭니다.

낭비 1: 미리 점유 (Over-reservation)

요청이 들어오면 "이 요청이 최대 몇 토큰을 생성할지 모르니까 최대치만큼 잡아두자"고 합니다. 실제 생성 길이가 짧으면 남은 공간은 그냥 낭비됩니다.

낭비 2: 내부 단편화 (Internal Fragmentation)

미리 잡아둔 공간 중 실제로 쓰이지 않는 부분이 그 요청에 묶여있어 다른 요청이 사용할 수 없습니다.

낭비 3: 외부 단편화 (External Fragmentation)

여러 요청이 들어오고 나가다 보면 메모리 공간이 군데군데 비는 조각(fragment) 이 생깁니다. 총 여유 메모리는 충분해도 연속된 공간이 없어서 새 요청을 할당하지 못하는 상황이 발생합니다.

공유 불가 (No Sharing)

Beam search나 parallel sampling처럼 하나의 프롬프트에서 여러 후보를 생성하는 경우, 공통 프롬프트 부분의 KV 캐시를 공유하면 메모리를 아낄 수 있습니다. 그러나 기존 방식은 각 요청마다 별도로 복사해서 관리합니다.

 

5. Compaction이 불가능한 이유

일반 메모리 관리에서는 외부 단편화를 해결하기 위해 compaction(데이터를 한쪽으로 모아서 빈 공간을 합치는 작업)을 수행합니다. 그런데 GPU에서 KV 캐시는 compaction을 하기 어렵습니다.

Decode 단계에서는 매 스텝마다 KV 캐시를 읽고 써야 합니다. 실시간으로 데이터를 옮기면 추론이 중단되고, 이동 비용이 추론 비용보다 훨씬 큽니다. 결국 단편화된 메모리를 그대로 쓸 수밖에 없는 구조입니다.

 

정리

문제	원인	결과
미리 점유	최대 길이 기준 사전 할당	실제 사용량보다 많은 메모리 소비
내부 단편화	할당된 공간 중 미사용 부분 존재	메모리 묶임, 활용률 저하
외부 단편화	비연속 빈 공간 발생	충분한 총량에도 할당 실패
공유 불가	요청별 독립 할당	중복 KV 캐시 저장
Compaction 불가	실시간 이동 비용 과다	단편화 해결 수단 없음

이 다섯 가지 문제가 맞물려 GPU 메모리 활용률이 60~80%에 그치고, 같은 GPU에서 처리할 수 있는 동시 요청 수가 심각하게 제한됩니다.

다음 포스트에서는 vLLM이 이 문제를 어떻게 해결했는지 OS의 페이지 가상 메모리 개념을 KV 캐시에 적용한 PagedAttention 을 살펴보겠습니다.

 

출처

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