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.

翻译：由于大型语言模型（LLM）的广泛应用，针对LLM推理的键值（KV）缓存压缩技术受到了显著关注。在近年来提出的众多方法中，逐层令牌剪枝方法——即在特定层选择令牌子集保留在KV缓存中并剪除其余令牌——已成为最流行的方案之一。这些方法主要采用一组预定义的层进行令牌选择。此类设计缺乏灵活性，因为其准确率在不同任务间差异显著，且在较难任务（如KV检索）中性能会明显下降。本文提出ASL，一种无需训练的自适应选择层方法，用于KV缓存压缩，该方法利用按注意力分数排序的令牌秩方差进行决策。所提出的方法在满足用户指定KV缓存预算要求的同时，平衡了不同任务间的性能表现。ASL在预填充阶段运行，并可与其他现有的KV缓存压缩方法（如SnapKV）结合使用以优化解码阶段。通过在InfiniteBench、RULER和NIAH基准测试上的评估，我们证明：在采用单次令牌选择（即在某一层选择令牌并传播至更深层）的情况下，ASL在保持解码速度与KV缓存压缩率的同时，其准确率优于当前最先进的逐层令牌选择方法。