DepthKV: Layer-wise Budget Allocation for Smarter KV Cache Pruning in Long-Context Reasoning

DepthKV proposes a layer-dependent KV cache pruning framework. Addressing the memory bottleneck in long-context LLM reasoning, it allocates global cache budget based on the differences in pruning sensitivity across Transformer layers. It consistently outperforms traditional uniform pruning methods at the same compression ratio, providing a new idea for memory optimization.

Memory Challenges

Long-context capabilities (e.g., 128K window) enable applications like document understanding, but KV cache memory grows linearly with sequence length, becoming the largest consumer of GPU memory and limiting context length and concurrent requests.

Limitations of Existing Pruning

Most pruning methods assume a uniform ratio, leading to cache waste in insensitive layers and over-pruning in sensitive layers, resulting in suboptimal resource allocation.

Experiments show that different layers have significant differences in pruning sensitivity. Lower layers handle local lexical/syntactic information and have weak dependence on distant tokens; some middle/higher layers are responsible for modeling long-range dependencies and are more sensitive to cache integrity. A one-size-fits-all strategy cannot allocate resources optimally.

1. Sensitivity Evaluation: Before deployment, use a small amount of calibration data to detect the impact of pruning each layer on output, obtaining a layer-wise sensitivity distribution.
2. Budget Allocation: Use optimization algorithms or heuristic rules to distribute the global cache budget to each layer in a differentiated way (sensitive layers get more quota, insensitive layers are pruned aggressively).
3. Low Overhead: Sensitivity evaluation only needs to be done once, with no additional runtime overhead during inference.

• Performance Advantages: Validated across multiple models and tasks, achieving better results at the same pruning ratio, with more significant improvements at high pruning ratios (20%-30%).
• Task Adaptability: Effective in long-range retrieval (e.g., needle-in-a-haystack) and long-document summarization tasks.
• Compatibility: Can be combined with existing pruning strategies (e.g., attention score pruning) to provide additional benefits.

1. Layer-wise Configuration: Avoid blind uniform pruning; first detect layer sensitivity to identify safely compressible layers.
2. Diagnostic Tools: Sensitivity evaluation can help understand the model's long-context processing mechanism, guiding architecture design/fine-tuning.
3. Memory-Constrained Scenarios: Provide more aggressive cache compression solutions for edge devices or high-concurrency services to reduce costs.

• Limitations: Sensitivity evaluation depends on calibration data; different data may lead to distribution differences.
• Future: Extend dynamic budget allocation (adjust in real-time based on input characteristics);推广 the core idea to optimization techniques like quantization and mixed-precision inference.

DepthKV addresses the inter-layer pruning sensitivity issue through differentiated cache budget allocation, improving pruning effectiveness without increasing runtime overhead. It is a noteworthy solution for memory optimization in long-context reasoning.