Beyond Token Eviction: Mixed-Dimension Budget Allocation for Efficient KV Cache Compression

Ruijie Miao; Zhiming Wang; Wang Li; Shiwei Wu; Shufan Liu; Yanbing Jiang; Tong Yang

cs.LGMar 21, 2026

Beyond Token Eviction: Mixed-Dimension Budget Allocation for Efficient KV Cache Compression

Authors: Ruijie Miao, Zhiming Wang, Wang Li, Shiwei Wu, Shufan Liu, Yanbing Jiang, Tong Yang

Organizations: Peking University · ByteDance Seed

Abstract

Key-value (KV) caching is widely used to accelerate transformer inference, but its memory cost grows linearly with input length, limiting long-context deployment. Existing token eviction methods reduce memory by discarding less important tokens, which can be viewed as a coarse form of dimensionality reduction that assigns each token either zero or full dimension. We propose MixedDimKV, a mixed-dimension KV cache compression method that allocates dimensions to tokens at a more granular level, and MixedDimKV-H, which further integrates head-level importance information. Experiments on long-context benchmarks show that MixedDimKV outperforms prior KV cache compression methods that do not rely on head-level importance profiling. When equipped with the same head-level importance information, MixedDimKV-H consistently outperforms HeadKV. Notably, our approach achieves comparable performance to full attention on LongBench with only 6.25% of the KV cache. Furthermore, in the Needle-in-a-Haystack test, our solution maintains 100% accuracy at a 50K context length while using as little as 0.26% of the cache.

Figures & tables

Figure 1: The two-stage dimension allocation process for KV cache compression. Candidate compression ratios are evaluated for each token to compute loss scores. An optimization objective is then applied to minimize total loss under the specified memory budget.

Granularity	Projection Matrices Overhead	Computational Efficiency	Compression Quality
Joint-head	$\times$	$\times$	✓
Head-wise	✓	✓	✓

Table 1: Comparison of head-wise and joint-head compression. Head-wise compression reduces projection overhead and enhances inference efficiency. Its support for head-level budget allocation ensures competitive compression performance.

Method	Single-Doc QA			Multi-Doc QA			Summarization			Few-shot Learning			Synthetic		Code		Avg.
Method	NQA	Qasper	MF-en	HQA	2WQA	Musq	GRep	QMSum	MNews	TREC	TQA	SAMSum	PCount	PRe	Lcc	RB-P	Avg.
Llama-3-8B-Instruct KV size = 128
Full	25.56	32.30	39.71	43.56	35.49	21.14	28.58	23.22	26.67	74.00	90.48	42.29	4.80	69.75	59.13	54.02	41.92
H2O	23.52	13.74	25.26	29.07	22.72	14.67	21.94	21.34	24.01	46.50	87.99	35.12	5.36	68.21	54.77	48.73	33.93
SnapKV	22.39	15.98	30.97	40.78	28.60	19.35	19.76	21.88	21.40	65.50	89.72	38.77	5.75	69.00	58.76	54.73	37.71
Pyramid	21.40	16.92	31.62	38.45	28.72	18.59	19.97	22.48	20.96	66.50	89.35	38.39	5.92	69.00	57.86	51.80	37.37

Table 2: Performance comparison on LongBench. MixedDimKV (MD) is compared against KV cache compression methods that do not rely on additional head-level importance profiling.

Figure 2: Comparison of HeadKV and MixedDimKV-H.

Panel of Figure 3 — Figure 3: Performance comparison on RULER. The terms ‘SKV’, ‘PKV’, ‘MD’, ‘HKV’, ‘MD-H’ denote SnapKV, PyramidKV, MixedDimKV, HeadKV, MixedDimKV-H, respectively.

Panel of Figure 4 — Figure 4: Results of Needle-in-a-Haystack test on Llama-3-8B-Instruct-1048K with KV size $128$ .

Figure 5: Decoding latency and peak memory usage.

Figure 6: Comparison of head-wise and joint-head compression.

Method	SnapKV	$\{0\%,100\%\}$	$\{0\%,12.5\%,25\%,100\%\}$
Score	39.77	40.93	41.76

Table 3: Comparison of SnapKV and MixedDimKV with two granularity of candidate ratio sets:

\{0\%,100\%\}

, and

\{0\%,12.5\%,25\%,100\%\}

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Context Length ( $N$ )	Avg. Primal Loss ( $P(\hat{d})$ )	Avg. Gap Upper Bound ( $\overline{\Gamma}(\lambda^{*})$ )	Avg. Relative Gap (%)
2,005	63.83	0.023	0.033%
4,752	71.26	0.010	0.015%
6,186	104.59	0.010	0.011%

Appendix

Table 4: Scalability of Duality Gap (MultiFieldQA-en from LongBench).

Metric	Min	25th Percentile	Median	75th Percentile	Max
Sensitivity ( $\lambda^{*}$ )	$1.3\times 10^{-5}$	$9.3\times 10^{-5}$	$1.7\times 10^{-4}$	$2.2\times 10^{-4}$	$4.2\times 10^{-4}$
Relative Gap (%)	0.000%	0.000%	0.000%	0.010%	0.085%

Appendix

Table 5: Layer-wise Statistics (at

N=6,186

Figure 7: Performance of MultiFieldQA-en under different KV compression ratios.

Method	Single-Doc QA			Multi-Doc QA			Summarization			Few-shot Learning			Synthetic		Code		Avg.
Method	NQA	Qasper	MF-en	HQA	2WQA	Musq	GRep	QMSum	MNews	TREC	TQA	SAMSum	PCount	PRe	Lcc	RB-P	Avg.
Llama-3-8B-Instruct KV size = 256
Full	25.56	32.30	39.71	43.56	35.49	21.14	28.58	23.22	26.67	74.00	90.48	42.29	4.80	69.75	59.13	54.02	41.92
H2O	23.09	18.56	29.02	36.01	28.07	15.70	23.31	21.48	24.57	60.50	88.70	38.55	5.37	68.04	58.52	52.73	37.01
SnapKV	23.36	20.2	37.36	42.37	33.18	19.96	21.76	22.00	22.83	71.00	90.86	40.03	5.83	69.50	60.16	55.91	39.77
Pyramid	23.99	20.51	36.06	42.47	31.34	20.28	21.36	22.68	22.83	71.00	90.48	39.86	5.83	69.25	58.64	54.06	39.42

Appendix

Table 6: Additional results on LongBench.

Method	Single-Doc QA			Multi-Doc QA			Summarization			Few-shot Learning			Synthetic		Code		Avg.
Method	NQA	Qasper	MF-en	HQA	2WQA	Musq	GRep	QMSum	MNews	TREC	TQA	SAMSum	PCount	PRe	Lcc	RB-P	Avg.
Llama-3-8B-Instruct KV size = 128
Full	25.56	32.30	39.71	43.56	35.49	21.14	28.58	23.22	26.67	74.00	90.48	42.29	4.80	69.75	59.13	54.02	41.92
HeadKV	22.56	28.04	39.64	44.40	31.83	20.62	22.16	23.04	24.04	71.00	90.10	38.81	4.89	68.75	61.72	61.77	40.84
MD-H	24.59	27.15	41.39	43.59	35.10	20.59	23.43	22.74	25.56	70.00	90.34	40.51	5.08	69.50	62.30	59.69	41.35
Llama-3-8B-Instruct KV size = 256

Appendix

Table 7: Full LongBench results for HeadKV and MixedDimKV-H (MD-H).

Method	S1	S2	MK1	MK2	MQ	MV	VT	QA-1	QA-2	Avg.
Context Length = 8K
Full	100	100	99.4	99.4	99.25	83.4	96.64	69.83	58.4	89.59
H2O	1.8	0.2	0.4	3	0.05	0.3	15.76	36.8	32.8	10.12
SnapKV	100	98.2	96.8	92	75.95	71.9	80.52	63.93	50	81.03
PyramidKV	100	98	96.8	93	74.6	72.05	84.16	60.28	50	80.99
HeadKV	100	99.2	99.2	99	98.85	94.75	89.92	65.6	55.2	89.08

Appendix

Table 8: Full results of our solution and baselines on RULER.

Method	Single-Doc QA			Multi-Doc QA			Summarization			Few-shot Learning			Synthetic		Code		Avg.
Method	NQA	Qasper	MF-en	HQA	2WQA	Musq	GRep	QMSum	MNews	TREC	TQA	SAMSum	PCount	PRe	Lcc	RB-P	Avg.
Llama-3-8B-Instruct Budget = 10%
Full	25.56	32.30	39.71	43.56	35.49	21.14	28.58	23.22	26.67	74.00	90.48	42.29	4.80	69.75	59.13	54.02	41.92
Palu	15.44	9.32	19.22	17.82	16.87	8.65	9.74	20.21	14.08	45.00	22.89	23.32	1.59	4.35	7.82	9.07	15.34
MD	25.22	28.55	38.67	43.39	32.42	21.13	25.72	23.12	25.96	72.00	90.48	42.74	4.90	69.25	59.16	54.23	41.06
MD-H	25.32	32.07	39.96	43.91	36.38	21.48	27.39	23.25	25.95	73.50	90.56	41.84	4.81	69.50	61.86	57.19	42.19

Appendix

Table 9: Comparative analysis with PALU on LongBench.

Panel of Figure 8 — Figure 8: Results of Needle-in-a-Haystack test on Mistral-7B-Instruct with KV size $128$ .

Panel of Figure 9 — Figure 9: Average number of tokens and the allocation statistics across candidate compression ratios on LongBench.

KV-Cache Management

Key-value (KV) cache management reduces and coordinates Transformer attention state to make long-context inference faster and more memory-efficient.

See all 199 papers

DepthWeave-KV: Token-Adaptive Cross-Layer Residual Factorization for Long-Context KV Cache Compression
Jul 7, 2026Anna Cordoba, Adam Puente Tercero, Nerea Angulo Hijo +4KV Caching Long-Context Language Model Inference
Instituto de Investigación en Visión Artificial
AnchorKV: Anchor-Residual KV Cache Compression
Aug 3, 2026Malik Khalaf, Yara Shamshoum, Nitzan Hodos +2KV Caching Long-Context Language Model Inference
Department of Computer Science, Technion - Israel Institute of Technology
LKV: End-to-End Learning of Head-wise Budgets and Token Selection for LLM KV Cache Eviction
Apr 22, 2026Enshuai Zhou, Yifan Hao, Chao Wang +7KV Caching Long-Context Language Model Inference
University of Science and Technology of China, Hefei, China · State Key Lab of Processors, Institute of Computing Technology, CAS, Beijing, China · University of Chinese Academy of Sciences, Beijing, China

Jul 7, 2026J/K move · Enter open · S save

DepthWeave-KV: Token-Adaptive Cross-Layer Residual Factorization for Long-Context KV Cache Compression

Anna Cordoba, Adam Puente Tercero, Nerea Angulo Hijo, Mar Linares Tercero+3

Instituto de Investigación en Visión Artificial

KV Caching Long-Context Language Model Inference Memory-Efficient Optimization LLM Inference Acceleration Low-Rank Compression KV-Cache Compression

Open paper AskPDF arXiv

Long-context language model inference is increasingly limited by the memory bandwidth and capacity required to store key-value caches, yet existing compression methods often apply uniform budgets across layers or tokens and degrade retrieval when lexical cues and semantic states require different preservation. We introduce DepthWeave-KV, a token-adaptive cache compression method that factorizes key and value states across neighboring transformer layers using shared low-rank channel bases while retaining lightweight token-specific residuals where attention behavior is sensitive. DepthWeave-KV combines cross-depth residual factorization with a token-conditional depth router that allocates higher reconstruction rank to instruction-bearing and retrieval-critical tokens, and uses calibration-free online error tracking from attention-output probes to adapt compression during generation without retraining the base model. A fused CUDA implementation jointly performs basis lookup, residual dequantization, and attention projection to reduce decode-time memory traffic. Across LongBench, Needle-in-a-Haystack, L-Eval, and long-form QA and summarization benchmarks, DepthWeave-KV achieves near-full-cache task quality with substantially lower memory use, improving average score and retrieval accuracy over prior compressed caches while reaching 8.3x KV memory reduction and 72.8 tokens per second at 64K context.

Abstract

KV-Cache Management

Explore similar work

DepthWeave-KV: Token-Adaptive Cross-Layer Residual Factorization for Long-Context KV Cache Compression

AnchorKV: Anchor-Residual KV Cache Compression

LKV: End-to-End Learning of Head-wise Budgets and Token Selection for LLM KV Cache Eviction