FastDMS: vLLM BF16/FP8보다 빠른 6.4X KV-cache 압축

지난해 NVIDIA, 워싱턴대학교 (University of Warsaw), 에딘버러 대학교 (University of Edinburgh) 소속 연구자들이 Dynamic Memory Sparsification (DMS)를 발표했습니다. 이는 학습된 헤당 토큰 제거 (learned per-head token eviction) 를 사용하는 KV-cache 스파시피케이션 기법으로, 최대 8X KV-cache 압축을 보고했습니다.

저는 이 결과를 흥미로워하여 작은 참조 구현 (reference implementation) 과 트레이너를 구축하여 아이디어를 검증하고자 했습니다. WikiText-2 와 Llama 3.2 1B 를 사용하여 대략적인 복제를 달성했습니다:

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 분이 소요되었고 압축은 기본적으로 손실 없는 (lossless) 것으로 보였습니다. 그러나 작은 문제가 하나 있습니다. 제 HF 참조 구현은 약... 18 tok/s 로 실행되었습니다.

그래서 몇 주간의 커널 그라인딩 (kernel grinding) 후, FastDMS를 발표합니다. 이는 물리적으로 제거된 슬롯을 재수령 (reclaims) 하는 컴팩트 KV 저장 방식을 갖춘 DMS 의 MIT 라이선스 구현입니다. NVIDIA 의 원래 Qwen 3 8B DMS 체크포인트와 저의 Llama 3.2 1B DMS 체크포인트에서 테스트되었습니다. (원래 HF 참조 버전과 제 트레이너는 또한 리포지토리에 있습니다): https://github.com/shisa-ai/FastDMS

제 벤치마크 설정에서 FastDMS 는 vLLM BF16 KV 보다 5-8X 적은 KV 메모리를 사용하며 8K 컨텍스트에서 vLLM 보다 1.5-2X 빠르게 디코딩합니다.

컴팩트 DMS 는 이론적인 KV 바이트뿐만 아니라 실제 할당자/디바이스 메모리 (real allocator/device memory) 를 절약합니다. 아래 표는 ctx_len=8192, gen_len=128 을 사용합니다. 모든 vLLM 베이스라인은 작업에 맞는 정확한 크기의 토큰 풀을 사용합니다. KV/stage 메모리는 캐시 또는 캐시 플러스 스테이지의 발자국입니다. 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
...

호기심을 가진 분들을 위해, 네 이 속도 및 메모리 사용량 모두 TurboQuant 를 능가합니다:

Path	c	Prefill tok/s	Prefill vs BF16	Decode tok/s	Decode vs BF16	KV / stage memory	Status
vLLM BF16	1	123098.0	1.00x	459.4	1.00x	0.312 GiB BF16 KV	dense BF16-KV baseline
...

물론 압축이 출력 품질을 망가뜨린다면 이 모든 것은 의미가 없습니다. 이론적으로 DMS 제거 (eviction) 는 FP8 양자화 전에 적용되어 토큰을 유지하거나 제거하는 것을 결정하므로 FastDMS 컴팩트-DMS 의 품질 비교는 FP8 양자화 자체와 동일해야 합니다. 그러나 여전히 품질을 두 번 확인하는 것이 좋습니다.

이는 압축된 KV 캐시로 토큰을 생성하고 압축되지 않은 참조 (reference) 와 토큰별로 비교하여 측정합니다. 더 낮은 KLD (KL 발산) 가 더 좋으며 - 이는 압축된 모델의 다음 토큰 확률이 참조에 더 가까움을 의미합니다. 더 높은 토큰 매칭이 더 좋으며 - 이는 그리디 디코딩이 동일한 출력을 생성함을 의미합니다.

열목을 읽는 방법:

• KLD vs ref - 압축된 로짓 (logits) 과 참조 로짓 사이의 나츠 토큰당 KL 발산 (KL divergence in nats/token). 압축으로 인해 다음 토큰에 대한 확률 분포가 얼마나 이동했는지를 측정합니다. 낮을수록 좋습니다; 0.000은 동일함을 의미합니다.
• Token match - 참조와 동일한 greedy-decoded 토큰의 비율입니다. 96.9%는 64 토큰 중 약 2 개가 다름을 의미합니다.
• Tokens scored - 비교할 수 있는 디코딩 단계의 수입니다. 후보자가 참조와 다른 토큰을 생성하면 시퀀스가 분기되고 이후 단계는 더 이상 비교할 수 없습니다. 33/60은 분기 전 첫 33 토큰까지만 품질 지표가 적용됨을 의미합니다 - 보고된 KLD 와 PPL 은 전체 생성이 아닌 해당 접두사에 대한 것입니다. 더 높은 비율은 비교가 더 완전함을 의미합니다.

테스트 설정: ctx_len=1024, decode_len=16, 네 개의 프롬프트 (총 60-64 디코딩 단계). vLLM 행은 vLLM BF16 전체-KV 로짓과 비교됩니다. FastDMS 행은 FastDMS 에게 제거를 비활성화한 것과 비교됩니다 (참조 윈도우가 1M 토큰, 사실상 전체 KV 캐시를 유지함).

shisa-ai/Llama-3.2-1B-DMS-8x

Path	Reference	KLD vs ref	Token match	PPL	Tokens scored
vLLM BF16 full KV	self	0.000000	100.0%	2.3748	60/60
...

nvidia/Qwen3-8B-DMS-8x

Path	Reference	KLD vs ref	Token match	PPL	Tokens scored
vLLM BF16 full KV	self	0.000000	100.0%	1.6738	60/60
...

FastDMS compact-DMS 점수는 두 모델 모두에서 64/64 토큰을 얻습니다 - 모든 디코딩 단계가 참조와 비교 가능했으며, KLD 는 vLLM 의 자체 FP8 과 TurboQuant 압축과 같거나 낮습니다. Tokens scored 가 다르면 각 행의 PPL 은 서로 직접 비교할 수 없으며, 각 행의 PPL 은 다른 길이의 접두사에 대해 계산되기 때문입니다.

What's the catch?

그렇다면 이것이 정말로 훌륭하다면 왜 이미 모두가 그것을 사용하고 있을까요? 사실, vLLM 같은 프로덕션 엔진에 이를 구현하려면 중대한 수술을 해야 합니다. DMS compact KV 는 거의 모든 서비스 엔진 하위 시스템을 만듭니다:

Subsystem	What changes for DMS
PagedAttention / KV memory pool	DMS 는 부분 블록 해제와 함께 레이어별, 헤더별 변수 토큰 수를 필요로 합니다 - 표준 고정 페이지 블록이 아닙니다
...

이것을 시도해 보는 모든 사람을 축복합니다. KVCache 압축은 실제로 보이며, 올바른 구현에서는 품질 손실이 없으며, FastDMS 구현에 의해 표시된 바와 같이 더 빠르다는 것을 알 수 있습니다.

(관심 있는 사람들을 위해 더 많은 perf benchmarks, comparisons, raw logs 이 저장소에 있습니다)