Beyond Token Eviction: Mixed-Dimension Budget Allocation for Efficient KV Cache Compression
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
| Granularity | Projection Matrices Overhead | Computational Efficiency | Compression Quality |
| Joint-head | ✓ | ||
| Head-wise | ✓ | ✓ | ✓ |
| Method | Single-Doc QA | Multi-Doc QA | Summarization | Few-shot Learning | Synthetic | Code | Avg. | ||||||||||
| NQA | Qasper | MF-en | HQA | 2WQA | Musq | GRep | QMSum | MNews | TREC | TQA | SAMSum | PCount | PRe | Lcc | RB-P | ||
| 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 |
| Method | SnapKV | ||
| Score | 39.77 | 40.93 | 41.76 |
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
| Context Length ( ) | Avg. Primal Loss ( ) | Avg. Gap Upper Bound ( ) | 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% |
| Metric | Min | 25th Percentile | Median | 75th Percentile | Max |
| Sensitivity ( ) | |||||
| Relative Gap (%) | 0.000% | 0.000% | 0.000% | 0.010% | 0.085% |
| Method | Single-Doc QA | Multi-Doc QA | Summarization | Few-shot Learning | Synthetic | Code | Avg. | ||||||||||
| NQA | Qasper | MF-en | HQA | 2WQA | Musq | GRep | QMSum | MNews | TREC | TQA | SAMSum | PCount | PRe | Lcc | RB-P | ||
| 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 |
| Method | Single-Doc QA | Multi-Doc QA | Summarization | Few-shot Learning | Synthetic | Code | Avg. | ||||||||||
| NQA | Qasper | MF-en | HQA | 2WQA | Musq | GRep | QMSum | MNews | TREC | TQA | SAMSum | PCount | PRe | Lcc | RB-P | ||
| 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 | |||||||||||||||||
| 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 |
| Method | Single-Doc QA | Multi-Doc QA | Summarization | Few-shot Learning | Synthetic | Code | Avg. | ||||||||||
| NQA | Qasper | MF-en | HQA | 2WQA | Musq | GRep | QMSum | MNews | TREC | TQA | SAMSum | PCount | PRe | Lcc | RB-P | ||
| 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 |
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