🤖 machine learning

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

Tail-Replay is a prefix caching mechanism for hybrid large language models that enables unconstrained token-level reuse by reconstructing linear-attention states through replaying only a short, recent suffix of matched prefixes, thereby eliminating the need for recurrent-state checkpoints while achieving near-perfect quality retention and significant inference speedups.

Original authors: Yirui Liu, Ruoling Qi, Xuaner Wu, Penghang Liu, Jian Chen

Published 2026-09-01✓ Author reviewed ⓘ

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yirui Liu, Ruoling Qi, Xuaner Wu, Penghang Liu, Jian Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of modern artificial intelligence, large language models have become the engines behind everything from writing assistants to complex data analysis. These systems work by processing vast amounts of text, token by token, to predict what comes next. However, as these models are asked to handle longer and longer conversations or documents, they face a significant bottleneck: the cost of remembering everything they have read so far. To solve this, researchers have developed two main strategies. One approach involves changing the model's internal architecture to be more efficient, using a mix of standard memory layers and specialized, streamlined layers that summarize information rather than storing every single detail. The other strategy is a system trick called prefix caching, which recognizes that different users often start their requests with the same words. Instead of re-reading those identical opening sentences every time, the system saves the result of that first pass and reuses it. While these two strategies work well on their own, combining them has proven difficult because the streamlined memory layers cannot easily be paused and restarted at any point, unlike their standard counterparts.

This incompatibility created a specific problem for engineers trying to build faster, more efficient AI systems. When a standard memory layer is reused, the system can jump directly to any point in the saved text. But the streamlined layers, which are designed to compress information, maintain a continuous state that cannot be rolled back to an arbitrary starting point without losing its meaning. Previous solutions tried to work around this by saving snapshots of the system's state at fixed intervals, but this meant that the system could only reuse text if the shared portion ended exactly at one of those snapshots. If a user's request shared a long string of words that ended just after a snapshot, the system had to discard the match and start over, wasting the efficiency gains.

Researchers at China Telecom's Institute of Artificial Intelligence and Shanghai Jiao Tong University have developed a new method called Tail-Replay to solve this issue. Their approach allows the system to reuse shared text at any point, regardless of where the snapshots were taken. The core idea relies on a specific property of the streamlined memory layers: they are designed to weigh recent information more heavily than older information. As the system processes a long text, the influence of the very first words gradually fades, while the most recent words dominate the current state. The researchers realized that to recreate the state of a matched prefix, the system does not need to replay the entire history of that text. Instead, it only needs to replay the most recent, short segment of that shared text.

The new method works by saving the exact, detailed memory of the standard layers for every single word, while omitting the snapshots for the streamlined layers. When a new request arrives that shares a long opening with a previous one, the system retrieves the saved standard memory for the matching part. For the streamlined layers, instead of trying to find a perfect snapshot, the system takes the saved detailed memory of the last few words of the shared text and runs them through the streamlined layers from scratch. This short replay reconstructs the necessary state with high accuracy. Because the system only needs to replay a small tail of the text, the process is fast and does not require storing the heavy, intermediate snapshots that previously limited flexibility.

The team tested this method on three different hybrid language models using standard benchmarks designed to measure performance on long documents and complex reasoning tasks. They found that by replaying only five to ten percent of the matched text, the system retained between 92.8 and 99.9 percent of the quality it would have achieved if it had processed the entire text from the beginning. In practical terms, this means the system can skip the heavy lifting of re-reading thousands of words without sacrificing the accuracy of its answers. The results showed that the method works consistently across different types of tasks, from answering questions about long stories to retrieving specific facts from massive datasets.

Beyond accuracy, the method delivered dramatic improvements in speed. When the system was asked to process requests with shared prefixes of 8,000, 16,000, or 32,000 words, the time it took to generate the first answer dropped significantly. For the longest texts, the new method was up to 14.3 times faster than the traditional approach of re-reading everything. The speedup grew larger as the text got longer, demonstrating that the efficiency gains are most valuable when the context is most demanding. The researchers also developed optimizations to further reduce the time spent moving data between memory and the processor, ensuring that the replay process does not become a new bottleneck.

This work demonstrates that the limitations of combining efficient model architectures with smart caching systems can be overcome without compromising performance. By understanding that the influence of old information naturally fades in these streamlined layers, the researchers turned a constraint into an opportunity. The Tail-Replay method allows systems to reuse shared text freely, determined only by the words themselves rather than by arbitrary checkpoints. This advancement suggests a path toward more responsive and efficient AI services that can handle the growing demands of long-context applications without requiring massive increases in computing power.

Technical Summary: Tail-Replay

Problem: The Curse of Linear Attention in Prefix Caching

Hybrid Large Language Models (LLMs) interleave full-attention (FA) layers with linear-attention layers to reduce the computational cost of long-context inference. While this architecture offers efficiency, it creates a fundamental mismatch with existing prefix caching mechanisms.

In standard FA models, the reusable state consists of token-indexed key-value (KV) caches, allowing the system to retrieve and reuse the state at any arbitrary token boundary. In contrast, linear-attention layers (such as Gated DeltaNet) summarize the processed prefix into recurrent states via in-place updates. Once these states advance, they cannot be rolled back to represent an arbitrary earlier prefix.

Existing hybrid prefix caching systems (e.g., Marconi, Sparse Prefix Caching) attempt to resolve this by storing recurrent-state checkpoints. However, this approach constrains token-level prefix reuse to a discrete set of boundaries aligned with these checkpoints. This limitation, termed the "curse of linear attention," prevents the flexible, fine-grained reuse of shared prefixes that is possible in FA-only models.

Methodology: Tail-Replay

The authors propose Tail-Replay, a prefix caching mechanism designed to enable unconstrained token-level prefix reuse in hybrid LLMs without relying on recurrent-state checkpoints. The core insight is that linear-attention mechanisms act as a structured, lossy compression of the input prefix: gated recurrent updates progressively attenuate the contributions of earlier inputs. Consequently, the recurrent state at the end of a matched prefix can be accurately approximated by replaying only a short, recent suffix (the "tail") of that prefix.

Key Mechanisms

Selective Caching: Tail-Replay caches the exact FA key-value pairs and the FA output hidden states for every token. It explicitly omits the storage of recurrent-state checkpoints.
Independent State Reconstruction: Upon a cache hit, the system retrieves the FA KV for the matched prefix. For the linear-attention layers, it reconstructs the recurrent state by:
- Fetching only the most recent $k$ tokens (the tail) of the cached FA output hiddens.
- Initializing the replay state to zero at the beginning of this tail.
- Replaying the Gated DeltaNet recurrence equations on this short tail to generate the state at the prefix boundary.
- This process is performed independently for each group of linear-attention layers following an FA layer, ensuring the input hidden states match the original prefill exactly.
Optimizations for Efficiency:
- Tail-FFN Skip: Since the FFN output at the end of a replayed group is not needed for state reconstruction (the next group starts from the cached FA hidden), the FFN computation is skipped during replay.
- Transfer/Replay Overlap: Because Tail-Replay does not depend on the cached FA KV during the replay computation, the transfer of FA KV from host memory to the device can proceed concurrently with the replay on a separate stream, hiding data-movement costs.

Results

The authors evaluated Tail-Replay on three Gated DeltaNet-based hybrid models (OLMo-Hybrid-7B, Qwen3.5-4B, and Qwen3.6-27B) using the LongBench and RULER benchmarks.

Quality Retention: With a replay budget of only 5–10% (replaying the last 5–10% of the matched prefix), Tail-Replay retains 92.8–99.9% of the full-prefill quality across LongBench and RULER. This significantly outperforms a "Zero-Only" baseline (which reuses FA KV but zeros linear states), which showed substantial quality degradation.
Serving Efficiency (TTFT): The method achieves significant Time-to-First-Token (TTFT) speedups, which scale with the matched prefix length.
- At 32K context length, speedups reached 9.1× to 14.3× over full prefill.
- The "Overlap + Skip" optimization further reduced TTFT by 18–42% compared to serialized host-to-device transfers at 32K.
- Increasing the replay budget from 5% to 10% had minimal impact on quality but slightly increased TTFT at longer contexts where replay became the dominant cost.

Significance and Contributions

The paper claims Tail-Replay is the first hybrid prefix caching mechanism to enable unconstrained token-level prefix reuse without being limited by recurrent-state checkpoint locations. Its significance lies in:

Decoupling Reuse from Checkpoints: It shifts the reuse boundary from discrete checkpoint locations to the actual boundaries of shared tokens, maximizing prefix reuse potential.
Efficient State Reconstruction: It introduces a replay-based mechanism that trades a small, controlled computational cost (replaying a short tail) for the elimination of expensive checkpoint storage and management.
Practical Viability: The evaluation demonstrates that high-quality inference can be maintained with a minimal replay budget (5–10%), making the approach practical for real-world serving scenarios where latency and memory are critical constraints.

The work suggests that for hybrid models utilizing Gated DeltaNet, the "lossy" nature of the linear attention mechanism can be leveraged as a feature rather than a bug, allowing for efficient, flexible caching strategies that were previously thought incompatible with recurrent state architectures.