PaperlayerView exact revision on arXiv

Source-linked AI summary

Just Let Linear States Forget the Distant Past: Prefix Caching via Suffix Replay for Hybrid LLMs

Yirui Liu, Ruoling Qi, Xuaner Wu, Yuxin Jin, Jian Chen, Penghang Liu, Yafei Huang, Jiawei Shao, Xuelong Li

arXiv:2609.33477v1cs.AIcs.PF

TL;DR

Hybrid LLMs combine token-addressable full-attention KV caches with recurrent linear-attention states that cannot be rolled back to arbitrary prefix boundaries, so checkpoint-based reuse is discrete. SuffixReplay reconstructs states at cache-supported page boundaries by replaying recent anchored inputs. Across three models, it preserves reported quality while reducing storage and improving serving performance on branching and over-HBM workloads.

  • Problem

    Hybrid prefix caching is limited to checkpoint-aligned boundaries because linear-attention recurrent states cannot be rolled back to arbitrary prefix positions like full-attention KV caches.

    Source: §1 · pp. 2–3

  • Method

    SuffixReplay retains sparse hidden-state anchors and reconstructs linear-attention states at cache-supported boundaries by replaying a recent suffix, with a pipelined serving path.

    Source: §1 · p. 3, §4 · p. 4, §5.2 · p. 7

  • Results

    Across three models, OLMo-Hybrid-7B retains 91.4% of full-prefill quality on RULER, while Qwen3.5-4B and Qwen3.6-27B-FP8 each exceed 99% on average across LongBench and RULER; storage is 0.36–0.51× SGLang’s default cache.

    Source: §1 · p. 3

  • Takeaways & Limitations

    SuffixReplay reduces median TTFT by 15–70% on branching workloads, sustains 2.3–4.3× SGLang’s throughput when the working set exceeds HBM, and matches it on high-hit continuation traffic.

    Source: §1 · p. 3

    SuffixReplay depends on recurrences that forget earlier inputs; its replay budget varies by model and is calibrated offline for each model.

    Source: §9 · p. 13

Abstract

from arXiv · show

Hybrid LLMs interleave full-attention layers with linear-attention layers to reduce long-context inference cost, but this structure complicates prefix caching. Full-attention KV caches are token-addressable, whereas linear-attention layers maintain recurrent states that cannot be rolled back to arbitrary prefix boundaries. Existing systems materialize recurrent-state checkpoints, restricting prefix reuse to checkpoint-aligned positions. We present SuffixReplay, the first prefix caching system that lets hybrid LLMs reuse cached prefixes at every cache-supported page boundary without materializing recurrent-state checkpoints. Our key insight is to just let linear states forget the distant past. Modern linear-attention mechanisms use recurrent decay and gating to attenuate the influence of old inputs. Therefore, instead of checkpointing every prefix boundary, SuffixReplay approximates the state at a matched boundary by replaying only a recent suffix of the layer's input hidden states, which we retain as anchors. At the algorithmic level, SuffixReplay combines layer-wise and token-wise anchor sparsity with a bounded replay budget to control storage, computation, and quality. At the system level, it uses an independently managed anchor sidecar and a pipelined replay path to overlap anchor movement and state reconstruction with the native serving pipeline. We evaluate SuffixReplay on three hybrid LLMs: OLMo-Hybrid-7B, Qwen3.5-4B, and Qwen3.6-27B-FP8. Across these models, SuffixReplay retains 91.4-100% of full-prefill quality on average across LongBench and RULER, while using only 0.36-0.51x the amortized per-token storage of SGLang's default 8192-token checkpoint cache. Integrated into SGLang, SuffixReplay reduces median TTFT by 15-70% on branching workloads, sustains 2.3-4.3x SGLang's throughput when the working set exceeds HBM, and matches SGLang on high-hit continuation traffic.

1 Introduction

Hybrid models reduce long-context inference cost, but their recurrent linear-attention states cannot be rolled back like token-addressable KV caches, leaving prefix reuse checkpoint-bound. SuffixReplay reconstructs states at cache-supported page boundaries by replaying recent hidden-state suffixes, and evaluates the resulting quality, storage, and serving trade-offs.

1 Introduction

  • Hybrid architectures combine full-attention and linear-attention layers to reduce quadratic full-attention computation for long-context inference, while prefix caching avoids redundant prefill for shared prefixes.These techniques address serving pressures from longer contexts, repeated invocations, and higher concurrency.Source: §1 · p. 1
  • Existing checkpoint-based hybrid caches permit prefix reuse only when a matched prefix ends at a stored checkpoint, while dense checkpoints are costly.In Qwen3.5-4B, one linear-attention checkpoint is approximately 49 MiB versus about 32 KiB for one token’s KV cache.Source: §1 · p. 2
  • Recurrent decay and erase gates attenuate distant inputs, allowing a recent hidden-state suffix to approximate the state at a matched prefix boundary.This makes replay cost depend on suffix length rather than the full prefix length.Source: §1 · p. 2
  • The design targets at least 90% of full-prefill quality on average across LongBench and RULER while bounding replay-input storage by the amortized cost of SGLang’s default 8192-token checkpoint interval.It also targets matching or outperforming SGLang’s native prefix-caching path in TTFT, throughput, and capacity under concurrent requests.Source: §1 · pp. 2–3
  • SuffixReplay reconstructs linear-attention states at every cache-supported page boundary without materializing recurrent-state checkpoints.It retains inputs as anchors and combines layer-wise and token-wise sparsity to bound replay computation and storage.Source: §1 · p. 3
  • SuffixReplay retains 91.4% of full-prefill quality on RULER for OLMo-Hybrid-7B, while Qwen3.5-4B and Qwen3.6-27B-FP8 each exceed 99% on average across LongBench and RULER.Across the three models, additional storage is 0.36–0.51× SGLang’s native cache at its default 8192-token checkpoint interval; median TTFT falls 15–70% on branching workloads, and throughput reaches 2.3–4.3× when the working set exceeds HBM.Source: §1 · p. 3

2 Background and Motivation

Hybrid models pair token-addressable full-attention KV entries with linear-attention states updated recurrently in place. This difference makes hybrid prefix reuse discrete for linear-attention layers, even though full-attention layers can reuse KV at each supported matched boundary.

2.1 Hybrid LLMs and Linear Attention

  • A linear-attention layer updates a fixed-size recurrent state from each token’s hidden vector; unlike KV entries, the state does not preserve separately addressable values for earlier positions.The update uses a transformation of the previous state and a token-specific write, both computed from the layer input.Source: §2.1 · p. 3

2.2 Prefix Caching and Reuse Boundaries

  • Continuous prefix caching makes every matched cache-supported boundary reusable for full-attention KV, but hybrid linear-attention states are available only at stored checkpoint positions.Shared prefixes arise at execution-dependent boundaries such as agent branches, intermediate continuations, or edited and retried chats.Source: §2.2 · pp. 3–4

3 Challenges

Fine-grained hybrid prefix reuse requires balancing replay-input storage, reconstruction work, and generation quality while integrating replay into the existing serving pipeline. Naive retention and unsynchronized replay threaten those targets through excessive storage, latency, and data movement.

3 Challenges

  • The overall challenge is to provide fine-grained hybrid prefix reuse without significantly sacrificing quality, storage efficiency, or serving performance, requiring algorithm–system co-design.The algorithm balances added state and computation, while the serving system integrates these costs with KV caching and inference.Source: §3 · p. 4

3.1 Algorithm level challenge: Balancing Anchor Storage, Replay Cost, and Quality

  • Retaining every linear-attention layer’s hidden state at every token enables reconstruction at each position but costs 20–28× more additional storage than SGLang’s native hybrid cache.Reducing retained information changes reconstruction fidelity and may require longer replay, increasing computation and latency on branching hits.Source: §3.1 · p. 4

3.2 System level challenge: Integrating Replay with the Serving Pipeline

  • A replay hit must fetch hidden states, reconstruct the linear-attention state, and coordinate both with KV placement, transfers, eviction, scheduling, and batched inference.The system must integrate replay without letting its added storage, movement, and computation dominate prefix-reuse benefits.Source: §3.2 · p. 4

4 SuffixReplay: Algorithm Design

SuffixReplay bounds the storage and computation of replay by combining independent layer groups, sparse page-aligned anchors, and a capped recent-suffix replay. This reconstructs linear-attention states at supported cache boundaries without storing a checkpoint for every position.

4.1 Layer-wise Sparse Anchors

  • Grouping consecutive linear-attention layers lets replay skip full-attention layers and bounds reconstruction-error propagation within each group.One group-entry anchor replaces per-layer anchors, reducing layer-wise storage overhead from O(L) to O(G).Source: §4.1 · pp. 4–5, Figure 2 · p. 5

4.2 Token-wise Sparse Anchors

  • At anchor density ρ=1/16, SuffixReplay uses 0.50×, 0.36×, and 0.51× SGLang’s 8192-token checkpoint-cache storage for OLMo-7B, Qwen3.5-4B, and Qwen3.6-27B-FP8, respectively.Anchors align with cache page boundaries, retaining a valid reconstruction starting point for every supported hit boundary.Source: §4.2 · p. 5, Table 2 · p. 6
  • Dense sampling places anchors near page ends and provides more stable reconstruction quality under the same replay budget than uniform sampling on several workloads.Both layouts use the same density; the experiments use the dense layout.Source: §4.2 · p. 5, Figure 3 · p. 5

4.3 Replay Length and Algorithm–System Interface

  • A longer replay suffix generally improves approximation to the full-prefix state but increases cache-hit computation and latency, while a shorter suffix reduces replay work but may degrade generation quality.The replay length is measured in retained anchors and capped; Qwen models meet the quality target at k=128, while OLMo-Hybrid-7B requires k=512 on average.Source: §4.3 · pp. 5–6

5 SuffixReplay: System Design

SuffixReplay keeps anchors associated with their KV pages while managing their placement and transfers independently. It overlaps anchor fetching and state reconstruction with serving, and uses already-resident live state to avoid replay for some multi-turn continuations.

5.1 Anchor Storage and Management

  • The anchor sidecar shares KV-page admission and eviction decisions but uses independent pools and placement policies, allowing anchors in host DRAM while corresponding KV pages remain in GPU HBM.With anchor density ρ=1/16 and page size p=64, each page contains four anchors and completes at the same page boundary as its KV page.Source: §5.1 · p. 6, §5 · p. 6, Figure 5 · p. 6

5.2 Pipelined Anchor Fetching and Replay

  • Anchor fetching begins after the prefix match and runs asynchronously on a dedicated copy stream, overlapping request scheduling and KV movement before replay reconstructs state.The forward stream waits only when it reaches the first layer that consumes a group's reconstructed state; replay and extension are handed off at group boundaries.Source: §5.2 · p. 7, Figure 7 · p. 7
  • Batched replay packs multiple hits into one variable-length execution, reducing repeated fixed launch costs while keeping per-hit replay work linear in the bounded replay length.Missing or unpublished anchors and other unavailable hit-path conditions fall back to full prefill; out-of-bucket replay lengths use eager replay.Source: §5.2 · p. 7

5.3 Multi-turn State at Request Boundaries

  • When a prior turn’s live linear-attention state remains resident, a continuation can skip reconstruction, but correctness does not depend on retaining that state.The replay path remains the correctness path, and no persistent checkpoint states are kept at internal chunk boundaries.Source: §5.3 · pp. 7–8

6 Implementation

The SGLang v0.5.18 implementation adds engine support for anchor management and graph-based replay while keeping SuffixReplay in a separate cache component. It does not modify attention or recurrent-state kernels.

6 Implementation

  • The engine changes add about 1.2K lines across 17 existing files and two new files, including radix-cache anchor association, sidecar allocation, and per-token destination tables for CUDA-graph writes.The component can be disabled so stock SGLang runs unchanged.Source: §6 · p. 8
  • A per-layer wait hook coordinates replay with forward execution, and asynchronous publication makes finished requests’ anchor blocks available to the cache.The wait hook reduces to a pointer comparison when replay is inactive.Source: §6 · p. 8
  • Replay graphs are captured at initialization and bucketed by total packed replay length; on Qwen3.5-4B, 48 graphs use approximately 1.1 GB of device memory.Uncaptured buckets fall back to eager replay, and pool reallocation invalidates graphs to avoid stale pointers.Source: §6 · p. 8

7 Evaluation

SuffixReplay is evaluated for quality, serving performance, and storage trade-offs against SGLang’s native hybrid prefix cache. Across the tested models and workloads, it preserves replay quality, improves reuse at off-checkpoint branches and under capacity pressure, and remains close to SGLang on high-hit continuations.

7.1 Setup

  • The evaluation tests replay quality on three hybrid models and serving performance on two Qwen models, using matched SGLang configurations and varied branching, continuation, and capacity workloads.Serving experiments use a single H800 GPU; the SGLang version used does not support OLMo-Hybrid-7B.Source: §7.1 · pp. 8–9, §7 · p. 8

7.2 Quality

  • SuffixReplay retains 91.4%–100% of full-prefill quality across the reported model and benchmark settings, with OLMo-Hybrid-7B requiring a larger replay budget than the two Qwen models.At k=512, OLMo averages 95.9% on LongBench QA and 91.4% on RULER; at k=128, Qwen3.5-4B reaches 100%/99.7% and Qwen3.6-27B-FP8 reaches 99.1%/100% on QA/RULER.Source: §7.2 · p. 9, Table 3 · p. 10

7.3 Branching prefix reuse

  • SuffixReplay lowers median TTFT by 15–70% on branching workloads, with larger gains at longer or between-checkpoint branch positions.For ToolMind beyond 16K context, TTFT falls 55–70%; in the eight-session between-checkpoint test, the median falls from 369 ms to 72–94 ms, and SuffixReplay remains 1.5–2× faster even with SGLang’s 2048-token chunks.Source: §7.3 · pp. 9–10, Figure 9 · p. 11, Figure 10 · p. 11

7.4 High-hit continuation

  • On 98%-hit Agent continuation traces, SuffixReplay matches or slightly exceeds SGLang throughput at each tested concurrency, with equal or lower TTFT.The reported 4B throughput differences are +2.7 to +6.5% at 8 and 16 sessions and −1.4% at 32, within round-to-round noise.Source: §7.4 · pp. 10–11, Table 4 · p. 11

7.5 Capacity

  • Under host-KV capacity pressure, SuffixReplay sustains 2.3–4.3× SGLang’s throughput, while SGLang’s hit rate drops from 0.86 to 0.05 after the working set exceeds device capacity.On LooGLE, SuffixReplay also serves 0.5–2 requests/s with median TTFT of 53–159 ms versus 426–1,307 ms for SGLang.Source: §7.5 · p. 12, Figure 12 · p. 12

7.6 Replay-length sensitivity

  • The replay budget controls the quality–cost trade-off: both Qwen models meet the quality target at k=128, while OLMo-Hybrid-7B needs a larger budget on average.The main serving experiments use conservative k_max caps of 512 for 4B and 256 for 27B; k_max=128 is the recommended lower-cost setting when additional quality headroom is unnecessary.Source: §7.2 · p. 9, §7.6 · p. 12

7.7 Replay overhead and checkpoint trade-offs

  • Shrinking SGLang’s prefill chunks increases checkpoint cost: 2048-token chunks raise cold-prefill TTFT by 22% and quadruple state bytes per cached token, while 1024-token chunks raise TTFT by 87% and reduce closed-loop throughput by 16%.Even with 2048-token chunks, SGLang is 1.5–2× slower than SuffixReplay on the shrinking-prefix workload.Source: Figure 10 · p. 11, §7.7 · p. 12, Table 5 · p. 12

8 Related Work

Prior prefix-caching work largely assumes token-addressable Transformer KV caches, while hybrid and recurrent models require state checkpoints. SuffixReplay instead retains inputs and reconstructs states at the cache-hit boundary, relying on decay and erase gates; this approach does not cover recurrences without decay.

Prefix caching for Transformers

  • Transformer prefix-caching systems reuse token-indexed KV prefixes and focus on page allocation, matching, session reuse, tiered storage, and request routing [9], but hybrid linear layers break the arbitrary-cut assumption.Source: §8 · p. 12

Position-independent KV reuse

  • Position-independent caching methods repair KV computed in a different context by recomputing selected tokens, while HyPIC and LinearKV extend that direction to hybrid models [4] [6] [32] [13] [14].Source: §8 · p. 12

Prefix caching for hybrid and recurrent models

  • Hybrid prefix-cache systems store linear-layer checkpoints, whereas SuffixReplay stores inputs and reconstructs state at the hit position; checkpoint-admission policies such as Marconi’s can still help select anchor pages [20] [26].Source: §8 · p. 12

Linear attention and forgetting

  • SuffixReplay relies on decay and erase gates that make states contracting functions of distant inputs; recurrences without decay are outside its scope [2] [7] [12] [15] [16] [18] [36].Source: §8 · p. 13

AI serving across heterogeneous resources

  • AI Flow studies coordination across heterogeneous devices and infrastructure [1] [25], while this work addresses intra-engine hybrid prefix reuse by coordinating anchors, KV caches, and replay.Source: §8 · p. 13

9 Limitations

SuffixReplay depends on linear-attention states forgetting earlier inputs, so it does not directly apply to recurrences that preserve all past contributions. Replay budgets vary by model and are currently calibrated offline.

9 Limitations

  • SuffixReplay does not directly apply to recurrences that preserve all past contributions.Its replay budget is model-dependent: Qwen models meet the quality target with k=128, whereas OLMo requires a larger budget calibrated offline.Source: §9 · p. 13
  • Replay can increase TTFT when many cache hits arrive in one scheduler step, particularly if SGLang already has checkpoints at the matched boundaries.The overhead appears under synchronized, fixed-length requests; varying output lengths removes the gap.Source: §9 · p. 13

10 Conclusion

SuffixReplay replaces linear-state checkpoints with sparse input anchors and short-suffix replay, making every 64-token page boundary reusable. It reports lower TTFT on branching workloads, higher throughput when the working set exceeds HBM, and quality retention with less storage.

10 Conclusion

  • 41–69% lower median TTFT on branching continuations, 2.3–4.3× SGLang throughput when the working set outgrows HBM, and at least 94% of full-recompute quality.SuffixReplay stores 0.36–0.51× the state bytes, ties SGLang on high-hit continuation traffic, and supports reuse at every page boundary.Source: §10 · p. 13
  • SuffixReplay stores sparse linear-layer input anchors and rebuilds state at any 64-token page boundary by replaying a short suffix.It relies on the recurrence's forgetting to make distant inputs irrelevant.Source: §10 · p. 13

A Per-token storage accounting

The accounting compares storage beyond the full-attention KV cache, which both designs store identically. It contrasts native recurrent-state checkpoints with SuffixReplay's sparse anchor storage.

A Per-token storage accounting

  • Both storage columns count bytes beyond the identical full-attention KV cache, and the appendix derives the values reported in Tables 1 and 2.The model-size accounting uses logged SGLang state-pool sizes for the Qwen models; for unsupported OLMo, it applies the same layout to reference-implementation state shapes, with convolution kernel width w=4 for all three models.Source: A Per-token storage accounting · p. 15, Table 6 · p. 15
  • Native checkpoints include recurrent state and short causal-convolution state for each linear-attention layer.The recurrent state is kept in FP32; the convolution state keeps the last w−1 inputs per channel in BF16.Source: A Per-token storage accounting · p. 15
  • Native cache amortized storage is per checkpoint interval T, with the default T=8192; halving T doubles the per-token cost.A checkpoint occupies B_ckpt = L_lin (B_rec + B_conv) bytes.Source: A Per-token storage accounting · p. 15
  • Naive replay stores one BF16 hidden-state vector per linear-attention layer per token, and this cost is unchanged when weights are quantized.The passage specifically notes this for Qwen3.6-27B-FP8.Source: A Per-token storage accounting · p. 15
  • With ρ=1/16, SuffixReplay uses 3,360, 2,240, and 9,600 bytes per token for OLMo-7B, Qwen3.5-4B, and Qwen3.6-27B, respectively.The reported reduction over naive replay is approximately 51–55×; anchors are stored at linear-group entries and a fraction of token positions.Source: A Per-token storage accounting · p. 15

Related papers

Loading 2609.33477v1…