LMCache

LMCache는 Berkeley의 연구팀이 2024년 초 공개한 시스템으로, LLM 추론의 KV 캐시 관리를 자동화합니다. Mooncake와 유사하지만 더 가벼운 구현이 특징입니다.

핵심 문제 재확인

대규모 LLM 서빙에서 KV 캐시가 GPU 메모리를 압도합니다:

• Llama 2 70B, 128K 컨텍스트: KV 캐시만 약 150GB
• GPU 메모리: A100 80GB, H100 80GB
• 결론: 한 요청만 들어와도 메모리가 부족합니다.

해결책:

1. 그냥 큰 GPU 사기 — H200(141GB VRAM). 비쌉니다 (~$200k).
2. CPU 메모리 활용 — 호스트 RAM(수백GB~1TB)을 보조 저장소로.
3. 캐시 압축 — INT8 양자화 등. 정확도 손실.
4. 동적 관리 — LMCache 같은 자동 스왑.

LMCache는 (4)를 우아하게 구현합니다.

LMCache의 설계

계층 구조

┌─────────────────┐
│ GPU (80GB)      │  우선순위 높음
│ - 방금 쓴 KV   │  
└────────┬────────┘
         │ (자동 스왑)
         ↓
┌─────────────────┐
│ CPU RAM (512GB) │  우선순위 중간
│ - 이전 KV      │
└────────┬────────┘
         │ (선택: NVMe)
         ↓
┌─────────────────┐
│ NVMe (4TB)      │  우선순위 낮음
│ - 오래된 KV    │
└─────────────────┘

프리페칭 기반 로드

전통적인 "필요하면 로드"는 느립니다. LMCache는 미리 예측해서 로드합니다.

<!-- -->
# 어텐션 계산 예정

<!-- -->
# → 필요한 KV 미리 GPU로 로드 (백그라운드)

<!-- -->
# → 계산 시작되면 메모리에 있음

CUDA 그래프, CUDA 스트림을 이용해 계산과 전송을 겹칩니다(overlap).

실제 구현

1. KV 레이아웃

LMCache는 KV를 시퀀스 구간으로 나눕니다.

예: 128K 토큰 시퀀스

├─ 구간 0 (토큰 0~1K)   → GPU
├─ 구간 1 (토큰 1K~2K)  → GPU
├─ 구간 2 (토큰 2K~4K)  → CPU RAM
├─ 구간 3 (토큰 4K~8K)  → CPU RAM
├─ 구간 4 (토큰 8K~64K) → NVMe
└─ 구간 5 (토큰 64K~128K) → NVMe

각 구간은 독립적으로 로드/언로드 가능합니다.

2. 스마트 프리페칭

디코드 단계에서 새 토큰을 생성할 때:

현재 토큰: 50000
다음 어텐션에 필요: 토큰 0~50000의 K, V
→ 예측: 다다음 어텐션에 필요: 토큰 0~50001의 K, V
→ 미리: 토큰 50001의 K, V를 CPU RAM에서 GPU로 이동 (백그라운드)

이렇게 하면 대부분의 로드 대기시간을 숨길 수 있습니다.

3. 적응형 배치 크기

요청에 따라 배치 크기를 동적으로 조정합니다.

• 짧은 컨텍스트 (예: 4K): GPU에 맞으므로 배치 크기 크게 (예: 32)
• 긴 컨텍스트 (예: 128K): 배치 크기 작게 (예: 2), 메모리 절약

자동으로 메모리 사용량을 "목표치" 이하로 유지합니다.

성능

Berkeley 벤치마크 (A100, 80GB, Llama 2 70B):

핵심: 메모리 폭발로 불가능하던 일이 가능해집니다. 레이턴시는 증가하지만 전혀 쓸 수 없는 수준은 아닙니다.

메모리 오버헤드

LMCache 자체의 메모리 오버헤드:

• 메타데이터 (어디가 GPU/CPU/NVMe): ~10MB
• 인덱스 (스왑 정책용): ~50MB
• 프리페칭 큐: ~100MB

무시할 수 있는 수준입니다.

vs Mooncake

두 시스템은 유사하지만:

선택 기준

LMCache를 쓸 때: - 연구 환경, 프로토타입 - 구현을 커스터마이즈하고 싶음 - 최소한의 오버헤드로 시작하고 싶음

Mooncake를 쓸 때: - 대규모 상용 배포 - 극도의 성능 최적화 필요 - Microsoft 기술 지원 원함

실제 통합

vLLM 지원

vLLM 최신 버전에서 LMCache 같은 "GPU-CPU 스왑" 옵션:

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    enable_prefix_caching=True,  # KV 캐시 재사용
    gpu_memory_utilization=0.8,
    # LMCache 같은 다단계 메모리는 실험 단계
)

독립 배포

LMCache는 PyTorch와 독립적으로 작동합니다:

import lmcache


<!-- -->
# KV 캐시 관리자 초기화
cache_manager = lmcache.KVCacheManager(
    gpu_memory=50,  # 50GB GPU 사용
    cpu_memory=200,  # 200GB CPU 사용
    nvm_path="/data/nvm_cache"
)


<!-- -->
# 추론 루프
for step in inference_loop:
    kv = model(...)
    cache_manager.put(step, kv)

한계와 과제

1. 레이턴시 변동성

KV를 CPU에서 로드할 때 지연이 발생합니다. 실시간 애플리케이션에서는: - 최악의 경우(NVMe 접근): 수십 ms - 평균(CPU 캐시 히트): 몇 ms

SLA가 엄격한 경우(예: 100ms 이내)는 조심해야 합니다.

2. 멀티 노드 확장 미흡

단일 머신에서는 완벽하지만, 여러 GPU/노드로 갈라지면: - 네트워크 대역폭 제약 - KV 동기화 오버헤드

Megatron 같은 분산 프레임워크와의 통합이 아직 연구 단계입니다.

3. 정밀도 보장

float32 vs bfloat16 변환 중 수치 오차 가능성. 특히 NVMe에서 로드할 때 정밀도 손실 모니터링 필요.

미래 방향

1. 하드웨어 진화: NVIDIA H200(141GB HBM3)이 비싼 게 아니라 표준이 되면, 다단계 캐시의 필요성 감소.
2. 적응형 양자화: CPU/NVMe에 저장된 KV를 INT8로 압축해서 전송 대역폭 더 절약.
3. 분산 캐시: 여러 서버의 GPU 메모리를 연합해서 사용. Ray, Ansor 같은 분산 시스템과 결합.

결론

LMCache는 "KV 캐시 메모리 부족"이라는 실질적 문제를 자동화된 계층 메모리로 해결합니다. Mooncake보다 가볍고, vLLM보다 더 정교한 메모리 관리를 제공합니다. 연구와 프로덕션의 중간 지점에서 빠르게 발전 중인 기술입니다.