Mustafar: Promoting Unstructured Sparsity for KV Cache Pruning in LLM Inference

TL;DR

This paper tackles the KV cache memory bottleneck in LLM inference by promoting unstructured sparsity through per-token magnitude pruning for both Key and Value caches. It introduces a bitmap-based sparse format and a custom sparse attention kernel that operate directly on compressed KV caches, enabling substantial memory compression (up to 65–83% at 50% sparsity and 45% at 70% sparsity) and up to 2.23x throughput gains in decode. Across multiple models and tasks, unstructured pruning consistently outperforms structured pruning baselines and remains compatible with orthogonal compression techniques such as token eviction and quantization. The work provides a practical end-to-end solution for longer contexts and higher throughput, with future directions including combining with weight sparsity and per-head or per-layer sparsity tuning to maximize efficiency and accuracy.

Abstract

We demonstrate that unstructured sparsity significantly improves KV cache compression for LLMs, enabling sparsity levels up to 70% without compromising accuracy or requiring fine-tuning. We conduct a systematic exploration of pruning strategies and find per-token magnitude-based pruning as highly effective for both Key and Value caches under unstructured sparsity, surpassing prior structured pruning schemes. The Key cache benefits from prominent outlier elements, while the Value cache surprisingly benefits from a simple magnitude-based pruning despite its uniform distribution. KV cache size is the major bottleneck in decode performance due to high memory overhead for large context lengths. To address this, we use a bitmap-based sparse format and a custom attention kernel capable of compressing and directly computing over compressed caches pruned to arbitrary sparsity patterns, significantly accelerating memory-bound operations in decode computations and thereby compensating for the overhead of runtime pruning and compression. Our custom attention kernel coupled with the bitmap-based format delivers substantial compression of KV cache upto 45% of dense inference and thereby enables longer context length and increased tokens/sec throughput of upto 2.23x compared to dense inference. Our pruning mechanism and sparse attention kernel is available at https://github.com/dhjoo98/mustafar.

Figures (9)

• Figure 1: High-level overview of Mustafar. Green region describes the pruning algorithm of Section \ref{['sec:2']}, pink region describes the custom sparse attention kernel of Section \ref{['sec:3']}.
• Figure 2: Visualization of the KV cache in LLaMA-2 7B. Color intensity indicates element magnitude. The figure was generated using the visualization code from KIVI kivi.
• Figure 3: Per-token, output-aware pruning of Key cache
• Figure 4: Output-aware per-channel (red) and magnitude-based per-token (pink) pruning of Value cache. Magnitude-based per-token pruning is equal to output-aware per-token pruning (yellow).
• Figure 5: Overview of Mustafar sparse attention kernel. In (b), multi-head, softmax, and normalization are omitted for simplicity.