vLLM KV Cache pool 크기는 무엇으로 결정될까?

요약

vLLM은 전체 VRAM × gpu-memory-utilization을 메모리 예산으로 잡고, 모델 가중치와 activation, non-Torch memory를 제외한 나머지를 KV cache pool로 사용합니다.

 

같은 조건이라면 모델 가중치가 작거나 gpu-memory-utilization 값이 클수록 KV cache pool이 커지지만, 값을 높일수록 OOM을 막기 위한 VRAM 여유분은 줄어듭니다. RTX 5060 Ti 16GB에서 Qwen2.5-3B-Instruct BF16을 실행한 결과 약 6GB가 KV cache로 할당됐습니다.

 

실습 환경

하드웨어:

• NVIDIA GPU RTX 5060 Ti 16GB

 

vLLM 환경:

• Ubuntu 20.04 LTS
• Docker vLLM

 

예제 코드: https://github.com/choisungwook/portfolio/blob/master/computer_science/ai/study-llmserving/ch5-6/docker-compose.yml

 

사전 지식

AI 모델 추론에는 행렬 연산을 수행하고 그 결과를 다음 레이어로 넘기는 forward 과정이 있습니다. 이때 필요한 GPU VRAM의 임시 공간을 activation이라고 합니다. CUDA runtime도 별도의 VRAM을 사용합니다.

 

이전 글에서 activation과 runtime을 자세히 다뤘습니다.
- https://malwareanalysis.tistory.com/958

 

vLLM 초기화 후 KV cache pool 할당

이론

vLLM은 초기화 과정에서 AI 모델 가중치와 실행에 필요한 VRAM 사용량을 계산하고, 메모리 예산에서 이를 뺀 나머지를 KV cache로 사용합니다.

KV cache 메모리 영역 계산
≈ 전체 VRAM × gpu-memory-utilization - model weights - peak activation memory - non-Torch memory

 

실험

NVIDIA GPU RTX 5060 Ti 16GB에서 Qwen2.5-3B-Instruct BF16 AI 모델을 vLLM으로 실행하면 약 6GB가 KV cache에 사용됩니다.

docker compose --profile bf16 logs vllm-bf16 | grep -i "KV cache"

 

 

KV cache pool(VRAM) 크기에 영향을 주는 요소

AI 모델 크기

KV cache는 남은 VRAM을 사용하기 때문에 AI 모델이 작을수록 더 많은 공간을 사용할 수 있습니다.

 

gpu-memory-utilization

vLLM은 초기화할 때 activation과 non-Torch memory 사용량을 측정합니다. 하지만 실행 중 OOM을 방지하려면 VRAM 안전 여유분(headroom)이 필요합니다. 안전 여유분은 상황에 따라 5%~15%를 설정하는 편입니다.

 

vLLM에서는 gpu-memory-utilization으로 모델 실행기가 사용할 VRAM 비율을 설정합니다. 이 값이 클수록 KV cache pool은 커지고, 남겨 두는 안전 여유분은 작아집니다.

command:
  - vllm
  - serve
  - --gpu-memory-utilization
  - "0.85"

 

gpu-memory-utilization=0.85
→ vLLM 메모리 예산: 전체 VRAM의 85%
→ 남겨둔 여유: 약 15%