Mooncake

Mooncake는 2024년 초 Microsoft가 공개한 분산 KV 캐시 관리 시스템입니다. 대규모 LLM 서빙(예: Llama 2 70B, GPT-4)에서 메모리 병목을 해결하려고 설계되었습니다.

배경: KV 캐시의 폭발

LLM 추론에서 KV 캐시는 시퀀스 길이에 비례해서 커집니다.

예: Llama 2 7B, bfloat16 (2바이트/값) - 1K 토큰: 약 112 MB (계층 수 32 × 2 (K, V) × 시퀀스 길이 × 헤드 차원) - 32K 토큰: 약 3.5 GB - 128K 토큰: 약 14 GB

70B 모델이면 메모리 요구량이 10배 증가합니다. A100 (80GB VRAM)에서도 여러 요청을 동시에 처리하려면 캐시 관리가 매우 중요합니다.

전통적인 해결책:

메모리 증가 — GPU를 많이 구매합니다. 비쌉니다.
캐시 압축 — 양자화로 KV를 INT8이나 FP8로 저장합니다. 정확도 손실.
요청 큐잉 — 캐시가 가득 차면 새 요청을 기다립니다. 레이턴시 증가.

Mooncake는 다른 접근을 합니다.

핵심 아이디어: 계층화된 저장소

Mooncake는 KV 캐시를 여러 레벨의 메모리에 분산합니다:

┌────────────────────┐
│ GPU VRAM (빠름)    │
│ - 자주 쓰는 부분  │  ~80GB
└────────────────────┘
         ↓ (스왑)
┌────────────────────┐
│ CPU RAM (중간)     │  ~500GB~1TB
│ - 최근 요청들      │
└────────────────────┘
         ↓ (스왑)
┌────────────────────┐
│ NVMe SSD (느림)    │  ~10TB (저렴)
│ - 이전 요청들      │
└────────────────────┘

각 레벨 간 전송을 동적으로 최적화합니다.

작동 원리

요청 도착: 새로운 입력 토큰들이 들어옵니다.
Prefill: 입력을 처리하고 새로운 KV를 생성합니다. GPU에 저장합니다.
Decode 반복: 한 토큰씩 생성하면서 KV를 누적합니다.
- 자주 쓰는 부분(최근 토큰들)은 GPU에 유지
- 덜 쓰는 부분(초반 토큰들)은 CPU 메모리로 푸시(evict)
- 매우 오래된 부분은 NVMe로 이동
필요시 로드: 어텐션 계산할 때 필요한 KV를 다시 GPU로 가져옵니다. 백그라운드 프리페칭으로 레이턴시 숨깁니다.

구체적 메커니즘

1. 적응형 스왑

Mooncake는 이전 요청들의 패턴을 학습해서 "어디를 언제 스왑할지" 예측합니다.

예: 512K 토큰 컨텍스트 - 어텐션 가중치를 분석하면, 최근 2K 토큰(최신 정보)과 초반 512 토큰(요약)만 높은 확률 - 중간 부분(~400K)은 거의 쓰이지 않음 - → 중간 부분만 NVMe로 보냅니다.

이를 통해 GPU 메모리 사용량을 40~50% 줄일 수 있습니다.

2. 비동기 프리페칭

스왑할 데이터를 미리 읽어옵니다(prefetch). NVIDIA의 CUDA 스트림을 이용해서 계산 중에도 로드합니다.

계산(토큰 1000 생성) → 동시에 다음 필요 KV 로드(배경)
다음 계산(토큰 1001 생성) → 이미 메모리에 있음

이렇게 하면 스왑 오버헤드를 숨길 수 있습니다.

3. KV 청크화

KV를 더 작은 단위(청크)로 나눠서 세밀한 제어를 가능하게 합니다.

기존: "전체 시퀀스 KV" 단위로 스왑
Mooncake: "16 토큰 단위"로 나눔
- 장점: 필요한 부분만 로드 가능. 메모리 낭비 최소화
- 단점: 관리 오버헤드 증가

성능 결과

Microsoft 벤치마크 (A100, 80GB):

시나리오	Baseline	Mooncake	향상도
32K 컨텍스트, 배치 4	OOM*	완료	무한대 (가능)
128K 컨텍스트, 배치 2	3.5초/토큰	2.8초/토큰	+25%
256K 컨텍스트, 배치 1	OOM	8초/토큰	가능

*OOM = Out of Memory

핵심 이점: 메모리를 초과하지 않으면서도 레이턴시 페널티가 작음.

메모리 활용률

Baseline은 GPU 메모리가 꽉 차면 추가 요청을 받을 수 없지만(대기), Mooncake는: - GPU: 60~70% 사용 - CPU RAM: 20~30% 사용 - NVMe: 10~20% 사용

따라서 더 많은 동시 요청을 처리할 수 있습니다.

비교: 다른 접근과의 차이

Paged Attention (vLLM)

vLLM의 Paged Attention은 KV를 "페이지"로 나누고, 페이지 테이블로 관리합니다(OS의 가상 메모리처럼).

장점: 간단한 구현, 메모리 단편화 감소
단점: GPU 메모리만 관리. 확장성 제한.

Mooncake는 여기서 "다단계 메모리"를 추가했습니다.

KV 양자화

KV를 INT8이나 INT4로 압축.

장점: 메모리 4~8배 감소
단점: 품질 손실, 디코딩 오류율 증가

Mooncake는 압축 없이 성능을 올립니다.

프로덕션 고려사항

1. 하드웨어 요구사항

효과를 보려면: - 고속 NVMe (PCIe 4.0 이상): CPU-NVMe 대역폭이 병목 - 충분한 CPU RAM: 일반적인 서버(1TB)면 괜찮음 - CPU-GPU 연결 (PCIe 4.0): 전송 속도 중요

2. 적응형 스왑의 오버헤드

KV 접근 패턴을 학습하는데 초기 요청들이 필요합니다. 처음 수십 개 요청은 예측이 부정확할 수 있습니다.

3. 멀티 노드 확장

단일 머신을 넘어 여러 GPU, 여러 노드로 확장할 때: - KV를 어디에 저장할지 결정 - 네트워크 대역폭 고려

Ray 같은 분산 캐시 시스템과의 통합이 필요합니다.

실제 배포 사례

vLLM v0.4.0+ 에서 Mooncake 같은 "다단계 캐시" 옵션을 실험적으로 제공합니다.

LMDeploy도 비슷한 기능을 "KV 리사이클" 모듈로 제공합니다.

프로덕션 서빙 시스템 SGLang(Stanford)은 KV 재사용 + 다단계 메모리 관리를 기본으로 탑재했습니다.

향후 방향

하드웨어 맞춤: HBM(High Bandwidth Memory) 칩이 점점 커지면서(NVIDIA H200은 141GB), 다단계 캐시의 필요성이 줄어들 수 있습니다. 다만 비용 측면에서는 여전히 필요.
클라우드 스토리지 통합: NVMe 대신 S3/Azure Blob을 사용하는 "무한 캐시" 아키텍처도 연구 중.
학습-서빙 통합: 학습 단계의 활성화(activation) 재계산 기법을 서빙의 KV 스왑으로 적용.

결론: Mooncake는 하드웨어 비용 대비 최대 처리량이라는 실질적 문제를 해결합니다. 특히 의료, 금융, 법률 같은 도메인에서 매우 긴 컨텍스트(100K+)가 필요한 경우 필수 기술이 될 것 같습니다.