Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Self-Indexing KVCache: Predicting Sparse Attention from Compressed Keys

Xu Yang, Jiapeng Zhang, Dongyang Zhao, Guo Chen, Zhuo Tang

Abstract

The KV cache in self-attention has emerged as a major bottleneck in long-context and large-batch inference for LLMs. Existing approaches often treat sparsity prediction and compression as separate modules, relying on auxiliary index structures to select relevant tokens, and on complex quantization schemes to reduce memory usage. This fragmented design introduces redundant overhead and limits scalability. In this paper, we propose a novel paradigm: treating the compressed key representation not merely as storage, but as a self-indexing structure that directly enables efficient sparse attention. By designing a sign-based 1-bit vector quantization (VQ) scheme, our method unifies compression and retrieval in a single, hardware-friendly format. This approach eliminates the need for external indices or learning-based predictors, offering a lightweight yet robust solution for memory-constrained inference. All components are designed to be hardware-efficient and easy to implement. By implementing custom CUDA kernels, our method integrates seamlessly with FlashAttention, minimizing additional runtime and memory overhead. Experimental results demonstrate that our approach delivers both effectiveness and efficiency.

Self-Indexing KVCache: Predicting Sparse Attention from Compressed Keys

Abstract

The KV cache in self-attention has emerged as a major bottleneck in long-context and large-batch inference for LLMs. Existing approaches often treat sparsity prediction and compression as separate modules, relying on auxiliary index structures to select relevant tokens, and on complex quantization schemes to reduce memory usage. This fragmented design introduces redundant overhead and limits scalability. In this paper, we propose a novel paradigm: treating the compressed key representation not merely as storage, but as a self-indexing structure that directly enables efficient sparse attention. By designing a sign-based 1-bit vector quantization (VQ) scheme, our method unifies compression and retrieval in a single, hardware-friendly format. This approach eliminates the need for external indices or learning-based predictors, offering a lightweight yet robust solution for memory-constrained inference. All components are designed to be hardware-efficient and easy to implement. By implementing custom CUDA kernels, our method integrates seamlessly with FlashAttention, minimizing additional runtime and memory overhead. Experimental results demonstrate that our approach delivers both effectiveness and efficiency.
Paper Structure (14 sections, 13 equations, 5 figures, 5 tables)

This paper contains 14 sections, 13 equations, 5 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Background and Preliminaries
  3. Implementation
  4. One-Pass Sign-Based Clustering
  5. Compressed-Domain Top-$k$ Retrieval
  6. Token-Wise Quantization Format
  7. Overhead Analysis
  8. Evaluation
  9. Experimental Setup
  10. Performance Study
  11. Efficiency Study
  12. Ablation Study
  13. Conclusion
  14. Acknowledgments

Figures (5)

  • Figure 1: The core idea of our method: Find a hardware-friendly compressed representation for the key. This compressed data is designed to split into two parts, one of which enables fast retrieval with accuracy comparable to the full-precision key.
  • Figure 2: Overview of Self-Indexing KVCache. By leveraging custom kernels, we significantly reduce the runtime overhead typically associated with retrieval and dequantization. Under low-bit quantization, we optionally preserve a fixed number of full-precision sink tokens during the prefill stage to improve inference robustness, without sacrificing efficiency.
  • Figure 3: Overview of the LUT-GEMV. We compute the dot product between the query and each codeword in the codebook to generate a lookup table of size 16.
  • Figure 4: Average task performance under different sparsity ratios on the 32K prompt Ruler Benchmark. Our method outperforms other method, with the 2-bit quantization at 7.5% sparsity.
  • Figure 5: KV cache memory overhead and latency of LlaMA3.1-8B at different prompt lengths. Our method dynamically select 7.5% tokens.