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

2025-05-28

AI Summary

KV cache memory overhead severely constrains throughput and scalability of large language models (LLMs) in long-context inference. To address this, we propose a fine-tuning-free, unstructured KV cache pruning method. First, we empirically discover that Value caches exhibit unexpected robustness to magnitude-based pruning. Second, we design a per-token dynamic pruning strategy that jointly leverages bitmap-based sparse storage and a custom CUDA attention kernel, enabling zero-copy, high-efficiency sparse computation under arbitrary sparsity patterns. Our approach preserves model accuracy without degradation: it compresses the KV cache to 45% of its original size, significantly extends supported context length, achieves a 2.23× speedup in decoding throughput, and supports up to 70% pruning rate.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

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.

Problem

• Unstructured sparsity enhances KV cache compression in LLMs
• Per-token magnitude pruning optimizes Key and Value caches effectively
• Bitmap-based sparse format accelerates memory-bound decode operations

Innovation

• Unstructured sparsity enables 70% KV cache compression
• Per-token magnitude pruning optimizes Key and Value caches
• Bitmap sparse format and custom kernel boost decode speed