FastDMS: First Public Implementation of NVIDIA's DMS, Saving Up to 7.6x KV-Cache with Zero Loss and Faster Than vLLM

This article is based on a LinkedIn post by Shisa AI CTO＆Co-founder Leonard Lin and is republished with permission.

I Just published one of the side projects that's been grinding away in the background over the past few weeks.

FastDMS (MIT licensed): https://lnkd.in/gGwqajVc

This is the first proper public implementation of Dynamic Memory Sparsification (DMS), a KV-cache compression technique published by NVIDIA affiliated researchers last year (reporting up to 8X kvcache memory savings): https://lnkd.in/gyAXX4fX

On my real-world tests I was able to get 4.8x-7.6x KV-cache memory savings with zero quality loss.

Besides a (very slow) HF reference and trainer code of the learned per-head token eviction, the "fast" part of FastDMS is that I had frontier models grinding away on kernel optimizations until prefill matched vLLM's BF16/FP8 speeds, and decode was >1.5x faster.

It's a bit of a side-quest, and getting it to production grade is ... going to be a slog, so I'm just putting it out in world now and hoping others pick up the baton. 😅

Original LinkedIn post link