Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference

TL;DR

The paper tackles the memory bottleneck of KV caches in long-context LLM inference by introducing Adaptive Selection Layer (ASL), a training-free method that adaptively selects the token-selection layer based on the variance of attention-derived token ranks. ASL operates during prefilling and enables one-shot token propagation to deeper layers, and can be combined with existing KV-reduction techniques such as SnapKV or GemFilter. By monitoring the relative variance of token-rank distributions across layers, ASL adapts to task difficulty and achieves higher accuracy than fixed-layer methods on InfiniteBench, RULER, and NIAH with comparable decoding speed and memory use. The approach yields improved accuracy with minimal overhead and demonstrates robust performance across extremely long contexts.

Abstract

Due to the prevalence of large language models (LLMs), key-value (KV) cache reduction for LLM inference has received remarkable attention. Among numerous works that have been proposed in recent years, layer-wise token pruning approaches, which select a subset of tokens at particular layers to retain in KV cache and prune others, are one of the most popular schemes. They primarily adopt a set of pre-defined layers, at which tokens are selected. Such design is inflexible in the sense that the accuracy significantly varies across tasks and deteriorates in harder tasks such as KV retrieval. In this paper, we propose ASL, a training-free method that adaptively chooses the selection layer for KV cache reduction, exploiting the variance of token ranks ordered by attention score. The proposed method balances the performance across different tasks while meeting the user-specified KV budget requirement. ASL operates during the prefilling stage and can be jointly used with existing KV cache reduction methods such as SnapKV to optimize the decoding stage. By evaluations on the InfiniteBench, RULER, and NIAH benchmarks, we show that equipped with one-shot token selection, where tokens are selected at a layer and propagated to deeper layers, ASL outperforms state-of-the-art layer-wise token selection methods in accuracy while maintaining decoding speed and KV cache reduction.

Figures (10)

• Figure 1: Performance of FastKV jo2025fastkv on four tasks under different selection layer (referred to as "TSP layer" in jo2025fastkv) settings: KV retrieval in InfiniteBench zhang2024bench, single-key NIAH (with varying difficulties) and QA in RULER hsieh2024ruler. KV compression before selection layer is disabled to highlight the impact.
• Figure 2: Comparison of FastKV and ASL.
• Figure 3: Attention patterns of KV retrieval (upper) and QA (lower). At early layers, the attention scores are roughly uniformly distributed across the context. At middle layers, a subset of tokens exhibit high scores (as shown in stripe-like regions). The scores are more localized at deep layers, focusing to a smaller subset of tokens (as shown in thin vertical lines).
• Figure 4: Ranks patterns across layers for KV retrieval (upper) and QA (lower), context length = 16k.
• Figure 5: Relative variance calculation.