FastDMS: vLLM BF16/FP8보다 빠른 6.4배 KV 캐시 압축 기술
FastDMS: 6.4X KV-cache compression running faster than vLLM BF16/FP8
핵심 요약
DMS 기술을 활용해 vLLM 대비 KV 캐시 메모리를 6.4배 압축하고 디코딩 속도를 2배 향상시킨 FastDMS 구현체 소개.
- FastDMS 구현 — KV 캐시를 물리적으로 재확보하여 vLLM 대비 5~8배 적은 메모리 사용함.
- 성능 최적화 — FP8 및 BF16 환경에서 vLLM보다 최대 2배 빠른 디코딩 속도를 달성함.
- DMS 기술 — 토큰별 중요도를 학습해 불필요한 KV 캐시 슬롯을 제거하는 효율적 압축 방식임.
- 오픈소스 공개 — MIT 라이선스로 배포되어 Qwen 3 및 Llama 3.2 모델에서 검증 완료됨.
작년 NVIDIA, 에든버러 대학교, 바르샤바 대학교 연구진은 Dynamic Memory Sparsification (DMS)를 발표했습니다. 이는 헤드별 토큰 제거를 학습하여 최대 8배의 KV 캐시 압축을 보고한 KV 캐시 희소화 기술입니다.
저는 이 결과가 흥미로워 아이디어를 검증하기 위해 작은 참조 구현체와 트레이너를 만들었습니다. Llama 3.2 1B 모델을 사용한 WikiText-2 데이터셋에서 대략적인 복제에 성공했습니다:
| Configuration | PPL | Delta | KLD (nats/tok) | Compression |
|---|---:|---:|---:|---:|
| Vanilla Llama-3.2-1B | 9.226 | - | - | 1x |
| DMS (trained, eviction active) | 9.200 | -0.28% | 0.026 | 6.4x |
DMS 예측기 학습은 PRO 6000에서 약 20분 정도 소요되었으며, 압축은 거의 손실 없이 이루어졌습니다. 하지만 한 가지 작은 문제가 있었는데, 제 HF 참조 구현체의 속도가 약 18 tok/s에 불과했습니다.
그래서 몇 주간의 커널 최적화 끝에, 제거된 슬롯을 물리적으로 재확보하는 소형 KV 스토리지를 갖춘 MIT 라이선스 DMS 구현체인 FastDMS를 발표하게 되어 기쁩니다. 이 구현체는 NVIDIA의 원본 Qwen 3 8B DMS 체크포인트와 제가 직접 학습시킨 Llama 3.2 1B DMS 체크포인트에서 테스트되었습니다. (원본 HF 참조 버전과 제 트레이너도 저장소에 포함되어 있습니다): https://github.com/shisa-ai/FastDMS
제 벤치마크 환경에서 FastDMS는 8K 컨텍스트 기준 vLLM BF16 KV보다 5-8배 적은 KV 메모리를 사용하면서도, vLLM보다 1.5-2배 빠르게 디코딩합니다.
Compact DMS는 이론적인 KV 바이트뿐만 아니라 실제 할당기/장치 메모리를 절약합니다. 아래 표는 ctx_len=8192, gen_len=128을 사용한 결과입니다. 모든 vLLM 베이스라인은 워크로드에 맞는 정확한 크기의 토큰 풀을 사용합니다. KV/스테이지 메모리는 캐시 또는 캐시와 스테이징을 합친 용량입니다. vLLM BF16은 dtype=bfloat16 및 kv_cache_dtype=auto를 의미하며, vLLM FP8은 kv_cache_dtype=fp8을 의미합니다.
| Model / compact-DMS row | c | vLLM BF16 KV → FastDMS KV | BF16 KV saved | vLLM FP8 KV → FastDMS KV | FP8 KV saved | vLLM TQ4 KV → FastDMS KV | TQ4 KV saved |
| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
| Llama-3.2-1B FastDMS default | 1 | `0.312 → 0.056 GiB` | **`5.6x`** | `0.156 → 0.056 GiB` | **`2.8x`** | `0.142 → 0.056 GiB` | **`2.5x`** |
| Llama-3.2-1B FastDMS default | 8 | `2.062 → 0.431 GiB` | **`4.8x`** | `1.031 → 0.431 GiB` | **`2.4x`** | `0.939 → 0.431 GiB` | **`2.2x`** |
| Qwen3-8B FastDMS compact DMS | 1 | `1.406 → 0.184 GiB` | **`7.6x`** | `0.703 → 0.184 GiB` | **`3.8x`** | — | — |
| Qwen3-8B FastDMS compact DMS | 8 | `9.281 → 1.462 GiB` | **`6.3x`** | `4.641 → 1.462 GiB` | **`3.2x`** | — | — |
