Tail-Replay: Escaping the Curse of Linear Attention in Prefix Caching for Hybrid LLMs
作者: Yirui Liu, Ruoling Qi, Xuaner Wu, Penghang Liu, Jian Chen
分类: cs.LG, cs.AI
发布日期: 2026-08-31
💡 一句话要点
提出Tail-Replay以解决混合大语言模型中的前缀缓存问题
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 混合大语言模型 前缀缓存 线性注意力 Gated DeltaNet 推理效率 自然语言处理 长上下文推理
📋 核心要点
- 现有的混合大语言模型在前缀缓存方面存在挑战,导致令牌级匹配受限于存储的递归状态检查点。
- Tail-Replay通过将线性注意力机制视为输入前缀的压缩,提出了一种新的前缀缓存机制,允许不受限的令牌级重用。
- 在LongBench和RULER基准测试中,Tail-Replay在仅5-10%的重放预算下,保持了92.8-99.9%的全预填质量,并在推理速度上实现了显著提升。
📝 摘要(中文)
混合大语言模型通过将全注意力层与线性注意力层交替使用,以降低长上下文推理的成本。然而,这种结构使得前缀缓存变得复杂:全注意力的键值缓存是基于令牌可寻址的,而线性注意力层则维护无法回滚到任意前缀边界的递归状态。现有的混合前缀缓存方法通过存储递归状态检查点来解决这一不匹配,导致令牌级匹配仅在与存储检查点对齐的位置可用,限制了前缀重用的边界。本文提出的Tail-Replay是一种前缀缓存机制,能够在混合大语言模型中实现不受限的令牌级前缀重用。其关键见解在于,线性注意力机制如Gated DeltaNet可以被视为输入前缀的结构化、损失压缩。实验结果表明,Tail-Replay在多个基准测试中表现出色,显著提高了推理效率。
🔬 方法详解
问题定义:论文要解决的问题是混合大语言模型中前缀缓存的复杂性,现有方法通过存储递归状态检查点来应对这一挑战,但限制了前缀重用的灵活性。
核心思路:论文的核心思路是利用线性注意力机制的结构化压缩特性,提出Tail-Replay机制,通过重放短的最近后缀来重建线性注意力状态,从而实现不受限的令牌级前缀重用。
技术框架:Tail-Replay的整体架构包括缓存全注意力的键值对,同时省略递归状态检查点。在缓存命中时,通过重放匹配前缀的短后缀来重建线性注意力状态。
关键创新:最重要的技术创新点在于Tail-Replay机制的提出,它允许前缀重用的边界由共享令牌决定,而非递归状态检查点,从而提高了灵活性和效率。
关键设计:在设计中,Tail-Replay的重放预算设置为5-10%,并在多个匹配前缀长度下评估了时间到第一个令牌的速度提升,确保了在不同上下文长度下的高效性。
🖼️ 关键图片
📊 实验亮点
在实验中,Tail-Replay在LongBench和RULER基准测试中表现优异,保持了92.8-99.9%的全预填质量,并在32K上下文长度下实现了9.1-14.3倍的速度提升,相较于全预填方法显著提高了推理效率。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、对话系统和长文本生成等场景。Tail-Replay的机制可以显著提高混合大语言模型在长上下文推理中的效率,具有广泛的实际价值和未来影响,尤其是在需要快速响应的应用中。
📄 摘要(原文)
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.