PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling
Abstract
LLMs use Pyramidal Information Funneling for long context processing, inspiring PyramidKV, a KV cache compression method that dynamically adjusts cache size by layer, reducing memory usage while maintaining or improving performance.
In this study, we investigate whether attention-based information flow inside large language models (LLMs) is aggregated through noticeable patterns for long context processing. Our observations reveal that LLMs aggregate information through Pyramidal Information Funneling where attention is scattering widely in lower layers, progressively consolidating within specific contexts, and ultimately focusin on critical tokens (a.k.a massive activation or attention sink) in higher layers. Motivated by these insights, we developed PyramidKV, a novel and effective KV cache compression method. This approach dynamically adjusts the KV cache size across different layers, allocating more cache in lower layers and less in higher ones, diverging from traditional methods that maintain a uniform KV cache size. Our experimental evaluations, utilizing the LongBench benchmark, show that PyramidKV matches the performance of models with a full KV cache while retaining only 12% of the KV cache, thus significantly reducing memory usage. In scenarios emphasizing memory efficiency, where only 0.7% of the KV cache is maintained, PyramidKV surpasses other KV cache compression techniques achieving up to a 20.5 absolute accuracy improvement on TREC.
Community
Great work on this paper — we used it as a core reference baseline while investigating what happens when KV-cache compression is pushed into the ultra-extreme 28.6x (3.5% KV) to 50.0x (2.0% KV) regime on DeepSeek-R1-Distill-Qwen-7B & 1.5B across the NVIDIA RULER multi-hop reasoning suite (NIAH-MK3, NIAH-MV/MQ, RULER-VT 5-Hop, RULER-CWE).
While observation-window voting (SnapKV / PyramidKV) is the strongest among existing eviction baselines (33.3%–41.7% RULER recall at 28.6x vs 0% for H2O), we found two mathematical bottlenecks that emerge specifically above 20x–50x compression on multi-hop chains, and an analytical fix (SPICY-KV) that is actually complementary to landmark selection:
- Closing the Softmax Jensen Gap on the Evicted 96.5% Bulk: Instead of hard-deleting non-landmark tokens (which shrinks the Softmax denominator $Z = \sum \exp(s_i)$), coalescing background tokens in Pre-RoPE space (preventing trigonometric phase cancellation across $t_1 \neq t_2$) and adding a closed-form 2nd-order subgrid variance correction ($\log m_m + \frac{|q_T|^2 \sigma_m^2}{2d_k}$) restores the exact Softmax partition function.
- Cross-Block Bridge Hub Shielding: Pairing observation-window attention with Pre-RoPE Algebraic Multigrid ($C/F$) bridge-hub detection captures intermediate variable-tracking hops (
RULER-VT) that occur thousands of tokens before the final query window.
Combining exact $1\times$ $C$-landmarks with variance-calibrated Pre-RoPE $F$-macro-nodes lifts 28.6x RULER recall from 33.3%–41.7% $\to$ 91.7%–100.0% (matching uncompressed FullKV), and 50.0x recall to 66.7%–87.5%.
Sharing our code, RULER evaluation harness, and paper in case it's helpful for anyone extending observation-window or merging methods to 30x–50x compression:
- Whitepaper (DOI): https://doi.org/10.5281/zenodo.23424348
- GitHub: https://github.com/Yuge-Y/SPICY-KV
- Interactive HF Space: https://huggingface.co/spaces/ElithiaTechnologies/SPICY-KV
Get this paper in your agent:
hf papers read 2406.02069 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper