PolyKV: Heterogeneous Retention and Allocation for KV Cache Compression

2026-06-13

AI Summary

Existing KV cache compression methods employ uniform strategies and homogeneous budgets, failing to accommodate the heterogeneous requirements of different Transformer layers during prefilling and decoding stages. This work proposes PolyKV, the first framework enabling layer-wise heterogeneous KV cache optimization. PolyKV dynamically selects compression strategies per layer based on layer-specific signal analysis and integrates multi-strategy routing with a non-uniform budget allocation algorithm to co-optimize strategy selection and resource distribution under a fixed total budget. Experiments on LLaMA-3.1-8B and Qwen3-8B demonstrate that, with an average cache budget of 512 tokens, PolyKV closes 54.5% and 25.7% of the LongBench performance gap between the strongest baseline and FullKV, respectively, and consistently outperforms all baselines by 1.7%–6.4% across a wide range of cache budgets.

Abstract

KV cache compression is essential for reducing the memory cost of long-context large language model inference. Existing approaches, however, typically apply a single compression policy and a uniform cache budget across all transformer layers. This uniform design ignores the fact that different layers can play different roles during prefill and decoding, and may therefore require different eviction strategies and cache capacities. We present PolyKV, a layer-wise KV cache optimization framework that considers design space with method selection and budget allocation. PolyKV routes each layer to a suitable KV compression policy based on layer-level signals, while assigning non-uniform budgets under a fixed total budget. This formulation enables heterogeneous compositions of existing KV cache methods. Experiments on LLaMA-3.1-8B and Qwen3-8B show that, under the same 512-token average KV budget, PolyKV recovers 54.5% and 25.7% of the LongBench performance gap between the strongest single-policy baseline and FullKV, respectively. Across 128-1024 budget sweep, PolyKV consistently improves over the strongest baseline by 1.7%-6.4%, corresponding to 40.0%-54.5% recovery of the FullKV gap.

Problem

• KV cache compression
• heterogeneous retention
• layer-wise allocation
• uniform budget
• transformer layers

Innovation

• KV cache compression
• layer-wise optimization
• heterogeneous retention
• non-uniform budget allocation
• long-context LLM inference