让混合模型的前缀缓存突破检查点限制,实现任意位置复用。
Tail-Replay: Escaping the Curse of Linear Attention in Prefix Caching for Hybrid LLMs

- 利用线性注意力的递推状态可由最近片段重播重建
- 5%-10%重播开销下保持92.8%-99.9%生成质量
- 适用于长上下文推理,显著提升32K前缀的生成速度
混合大语言模型通过交替使用全注意力与线性注意力层降低长上下文推理成本。然而,这种结构使前缀缓存复杂化:全注意力层的键值缓存可按标记定位,而线性注意力层依赖无法回溯到任意前缀边界的状态。现有方法通过存储递推状态检查点解决此问题,导致仅在检查点对齐位置才能实现标记级匹配,限制了前缀复用的连续性。本文提出Tail-Replay,一种可在混合模型中实现无约束标记级前缀复用的缓存机制。核心思想是:如门控增量网络(Gated DeltaNet)等线性注意力机制可视为输入前缀的有损压缩——门控递推更新逐步衰减早期输入的影响。因此,匹配前缀的递推状态可通过重播其短而近期的后缀良好近似。Tail-Replay仅缓存完整的全注意力键值对,省略递推状态检查点。命中缓存时,通过重播匹配前缀的短后缀重构线性注意力状态。由此,复用边界由共享标记决定,而非检查点。我们在三个基于Gated DeltaNet的混合模型上,使用LongBench和RULER基准评估了该方法。仅需5%–10%的重播预算,其在LongBench和RULER上保留了92.8%–99.9%的完整预填充质量。在服务效率方面,我们评估了不同匹配前缀长度(8K、16K、32K)下的首令牌耗时加速比,速度提升随前缀长度增加而增长,在32K时达到9.1–14.3倍于完整预填充的加速效果。
原文摘要 · Abstract (English)
Hybrid large language models interleave full-attention layers with linear-attention layers to reduce the cost of long-context inference. This structure complicates prefix caching: full-attention key-value caches are token-addressable, whereas linear-attention layers maintain recurrent states that cannot be rolled back to arbitrary prefix boundaries. Existing hybrid prefix caching methods address this mismatch by storing recurrent-state checkpoints. As a result, token-level matches are directly usable only at positions aligned with stored checkpoints, constraining prefix reuse to a discrete set of boundaries. We present Tail-Replay, a prefix caching mechanism that enables unconstrained token-level prefix reuse in hybrid large language models. The key insight is that linear-attention mechanisms such as Gated DeltaNet can be viewed as a structured, lossy compression of the input prefix: gated recurrent updates progressively attenuate the contributions of earlier inputs. Consequently, the recurrent state of a matched prefix can be well approximated by replaying only a short, recent suffix of that prefix. Tail-Replay exploits this property by caching the exact full-attention key-value cache while omitting recurrent-state checkpoints. On a cache hit, it reconstructs the linear-attention states by replaying a short, recent suffix of the matched prefix. As a result, the reuse boundary is determined by the shared tokens rather than by recurrent-state checkpoints. We evaluate Tail-Replay on three Gated DeltaNet-based hybrid models using the LongBench and RULER benchmarks. With only a 5--10\% replay budget, it retains 92.8--99.9\% of full-prefill quality on LongBench and RULER. For serving efficiency, we evaluate time-to-first-token speedups across multiple matched-prefix lengths---8K, 16K, and 32K. The speedup grows with prefix length, reaching $9.1$--$14.3\times$ over full prefill at 32K.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。