论文

尾部重播:摆脱混合前缀缓存中线性注意力的诅咒 LLM

Tail-Replay: Escaping the Curse of Linear Attention in Prefix Caching for Hybrid LLMs

模型推理KV Cache

摘要

混合 大语言模型 将全注意力层与线性注意力层交织在一起,以降低长上下文推理的成本。这种结构使前缀缓存变得复杂:全注意力键值缓存是词元可寻址的,而线性注意力层维护无法回滚到任意前缀边界的循环状态。现有的混合前缀缓存方法通过存储循环状态检查点来解决这种不匹配问题。因此,词元级 匹配只能在与存储的检查点对齐的位置直接使用,从而将前缀重用限制在一组离散的边界上。我们提出了 Tail-Replay,这是一种前缀缓存机制,可以在混合 大语言模型 中实现不受约束的 词元级 前缀重用。关键的见解是,诸如门控 DeltaNet 之类的线性注意力机制可以被视为输入前缀的结构化、有损压缩:门控循环更新逐渐减弱早期输入的贡献。因此,通过仅重放该前缀的短的、最近的后缀,可以很好地近似匹配前缀的循环状态。 Tail-Replay 通过缓存精确的全注意力键值缓存同时忽略循环状态检查点来利用此属性。在缓存命中时,它通过重放匹配前缀的一个简短的、最近的后缀来重建线性注意力状态。因此,重用边界是由共享词元决定的,而不是由循环状态检查点决定的。我们使用 LongBench 和 RULER 基准测试在三个基于门控 DeltaNet 的混合模型上评估 Tail-Replay。只需 5--10% 的重播预算,它就能在 LongBench 和 RULER 上保留 92.8--99.9% 的完全预填充质量。为了提高服务效率,我们评估了多个匹配前缀长度(8K、16K 和 32K)的首次词元加速时间。加速比随着前缀长度的增加而增加,在 32K 时比完全预填充达到 $9.1$--$14.3\times$。