论文

TaskPress:通过任务引导剪枝实现与查询无关的KV缓存压缩

TaskPress: Query-Agnostic KV Cache Compression via Task-Guided Pruning

模型推理KV Cache

摘要

长上下文推理的大语言模型受到关键-值缓存(key-value cache)线性增长到序列长度的限制。虽然剪枝可以缓解,但现有方法确定不可重复使用的查询特定令牌重要性。相比之下,我们引入TaskPress,这是一种任务导向、无查询关注的关键-值缓存淘汰框架。TaskPress不优化单个查询的缓存,而是构建一个可重用的记忆表示,基于高层次的任务指南。在预填充期间,指南作为元查询来过滤无关的令牌,在下游查询发出之前。此外,TaskPress利用量化尺度因子作为零成本信号,用于检测影响关键-值表示的异常点,提供对令牌重要性的高效近似代理。在各种任务上进行实验,包括长上下文输入,表明TaskPress有效地创建一个紧凑、可重用的缓存,覆盖不同查询。

TaskPress:通过任务引导剪枝实现与查询无关的KV缓存压缩:论文配图
图 4:(左)每种方法的驱逐延迟。 (右)面板描述了驱逐后解码过程中实现的内存节省。实验在 LLaMA-3.1-8B 上的 32K 上下文长度中进行。