Organizations: College of Computer Science and Electronic Engineering, Hunan University, Changsha, China · The Ministry of Education Key Laboratory of “Fusion Computing of Supercomputing and Artificial Intelligence” China · Shenzhen Research Institute, Hunan University, Shenzhen, China
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.
Figures & tables
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.
Method
Cache Bits (K,V,Index)
SD-QA
MD-QA
Summarization
Few-shot
Synthetic
Code
Avg.
Qasper
MF-en
HPQA
2WQA
GVRpt
QMSum
TREC
TrivQA
PR-en
Lcc
RB-P
Llama3.1-8B
16, 16, 0
45.5
53.8
54.7
47.1
34.9
25.3
73.0
91.6
99.5
63.4
56.7
58.7
SnapKV
16, 16, 0
33.9
46.4
54.4
44.7
21.8
22.6
48.0
90.5
78.0
58.0
50.1
49.9
Quest
16, 16, 2
44.9
51.6
55.0
46.9
31.2
24.0
71.0
90.6
95.5
60.6
50.5
56.5
DoubleSparse
16, 16, 2
43.8
51.7
54.4
45.9
31.7
24.0
67.5
91.9
97.0
50.5
44.7
54.8
Ours(16 bits)
16, 16, 1
45.4
53.9
55.4
47.0
34.0
24.8
72.5
91.7
99.5
63.2
55.3
58.4
Table 1: Performance comparison of LongBench. The best result is highlighted in underline .
Method
Cache Bits (K,V,Index)
Task
Avg.
NS1
NS2
NS3
NM1
NM2
NM3
NV
NQ
VT
CWE
FWE
QA1
QA2
Llama3.1-8B
16,16,0
100.0
100.0
100.0
100.0
100.0
100.0
100.0
100.0
98.4
59.2
94.67
72.0
56.0
90.8
SnapKV
16,16,0
100.0
100.0
28.0
100.0
32.0
8.0
100.0
100.0
96.8
43.6
90.0
60.0
56.0
70.8
Quest
16,16,2
100.0
100.0
100.0
100.0
100.0
56.0
98.0
100.0
97.6
31.6
90.7
68.0
56.0
84.5
DoubleSparse
16,16,2
100.0
100.0
100.0
100.0
68.0
68.0
100.0
100.0
97.6
34.8
90.6
64.0
64.0
83.6
Ours(16 bits)
16,16,1
100.0
100.0
100.0
100.0
100.0
100.0
98.0
100.0
98.4
53.2
89.3
68.0
60.0
89.4
Table 2: Detail experimental results of Llama3.1-8B on the 32K prompt Ruler Benchmark.
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.
Prompt Length
Ours
KIVI
Flash Attention2
8K
1.26
1.23
1.18
16K
2.75
2.66
2.62
32K
6.51
6.32
6.18
48K
11.12
10.88
OOM
64K
17.70
OOM
OOM
Table 3: Comparison of Time TT2T (seconds) across different methods and prompt lengths.
Figure 5: KV cache memory overhead and latency of LlaMA3.1-8B at different prompt lengths. Our method dynamically select 7.5% tokens.
Module
Method
Time (ms)
Clustering
Ours
88.85
KMeans (20 iterations)
2113.9
Retrieval
Ours
0.039
Quest (page size = 16)
0.041
Full K⋅q⊤
0.166
Attention
Ours (7.5%)
0.116
Table 4: Performance comparison across different modules under a 16K token input with batch size 10.
Setting
MF-en
HPQA
GovRpt
RB-P
Ours
54.6
55.7
33.1
54.4
w/o sign in quant
52.5
54.9
32.2
50.9
sign-only retrieval
52.8
53.2
31.9
52.1
w/o sink tokens
52.6
55.2
32.5
54.2
Table 5: Ablation study on the effect of components.
Sparse long-context inference requires efficient token retrieval in both prefill and decode. Existing methods often use different retrieval strategies for the two stages, preventing one retrieval representation from being reused throughout inference. We propose Self-Indexing Attention, a training-free framework built on a shared transform-domain sign-magnitude representation. The key signs provide a reusable token-level index for grouped prefill selection and decode retrieval, while the same representation remains compatible with external KV-cache compression without separate indexer metadata. This 1-bit index enables efficient retrieval through bitwise operations widely supported by modern accelerators. At 5% attention density, Self-Indexing Attention remains close to dense attention on LongBench and RULER and achieves up to 6.1x prefill and 10.3x decode attention-operator speedups. Experiments with TurboQuant and DeepSeekV4-Flash further demonstrate compatibility with low-bit KV-cache compression and pretrained sparse-attention indexers.
Under modern test-time compute and agentic paradigms, language models process ever-longer sequences. Efficient text generation with transformer architectures is increasingly constrained by the Key-Value cache memory footprint and bandwidth. To address this limitation, we introduce Self-Pruned Key-Value Attention (SP-KV), a mechanism designed to predict future KV utility in order to reduce the size of the long-term KV cache. This strategy operates at a fine granularity: a lightweight utility predictor scores each key-value pair, and while recent KVs are always available via a local window, older pairs are written in the cache and used in global attention only if their predicted utility surpasses a given threshold. The LLM and the utility predictor are trained jointly end-to-end exclusively through next-token prediction loss, and are adapted from pretrained LLM checkpoints. Rather than enforcing a fixed compression ratio, SP-KV performs dynamic sparsification: the mechanism adapts to the input and typically reduces the KV cache size by a factor of 3 to 10×, longer sequences often being more compressible. This leads to vast improvements in memory usage and decoding speed, with little to no degradation of validation loss nor performance on a broad set of downstream tasks. Beyond serving as an effective KV-cache reduction mechanism, our method reveals structured layer- and head-specific sparsity patterns that we can use to guide the design of hybrid local-global attention architectures.
Long-context Large Language Model inference is severely bottlenecked by the massive Key-Value (KV) cache, yet existing sparse attention methods often suffer from static fixed-budget (Top-k) retrieval or rely on proxy scores that are computationally expensive and biased. To address these limitations, we propose RaBitQCache, a novel sparse attention framework that utilizes randomized rotated binary quantization and high-throughput binary-INT4 arithmetic to efficiently estimate attention weights. Our proxy score serves as an unbiased estimator with a proven error bound, enabling adaptive Top-p retrieval that dynamically adjusts the token budget based on actual attention sparsity. We further implement a hardware-aware system with asynchronous pipelining and lazy updates to mask overhead. Evaluations demonstrate that RaBitQCache significantly accelerates inference and reduces memory I/O while preserving generation quality compared to state-of-the-art baselines. Code is available at https://github.com/Sakuraaa0/RaBitQCache.git.
Wenhao Li, Jinhao Dong, Hailin Zhang +3
School of Information, Renmin University of China, Beijing, China · Peking University, Beijing, China