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

2026-08-31Machine Learning

Machine LearningArtificial Intelligence
AI summary

The authors explain a new way to speed up hybrid large language models, which combine layers that look at all input tokens and layers that only remember some information. Usually, reusing past computations in such models is tricky because the different layers store information differently. Their method, Tail-Replay, cleverly reuses most of the previous work by storing exact data for the full-attention part and recreating the smaller-memory part by replaying only a recent small section of tokens. This approach lets the model reuse prefixes more flexibly and speeds up processing a lot while keeping accuracy high.

hybrid large language modelsfull-attention layerslinear-attention layersprefix cachingrecurrent stateskey-value cacheGated DeltaNetlossy compressionLongBenchtime-to-first-token speedup
Authors
Yirui Liu, Ruoling Qi, Xuaner Wu, Penghang Liu, Jian Chen
Abstract
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.