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.