Technical Summary: Sparse Prefix Caching for Hybrid and Recurrent LLM Serving

1. Problem Statement

Existing Large Language Model (LLM) serving systems rely heavily on dense prefix caching, where the Key/Value (KV) states for every token in a shared prefix are stored to avoid redundant prefill computation. This approach is standard for Transformer architectures. However, State-Space Models (SSMs) and hybrid architectures (combining attention with recurrent layers) present a structural asymmetry: a recurrent layer evolves its state as $h_t = F(h_{t-1}, x_t)$, meaning it can resume computation from a single stored state $h_c$ rather than requiring the entire history of intermediate states.

Dense per-token caching of recurrent states is often unnecessary and prohibitively expensive. For models like Qwen-3.5, the recurrent state per head scales as $d_{head}^2$, whereas attention KV storage scales linearly. Consequently, storing every token's recurrent state consumes significantly more memory than attention KV caches, yet offers no additional benefit over storing a sparse set of checkpoints.

The core problem addressed is: Given a memory budget of $M$ checkpoint slots for a prefix of length $N$, where should checkpoints be placed to minimize the expected recomputation cost under a specific distribution of future request overlap depths? This optimization is distinct from cache admission/eviction policies; it focuses on which tokens to checkpoint within a single retained entry to maximize utility for partial prefix overlaps.

2. Methodology

Theoretical Formulation

The authors formalize the problem as a one-sided weighted k-median problem on a line.

• Let $C = \{c_1, \dots, c_M\}$ be the set of checkpoint positions.
• For a request overlapping the prefix up to depth $t$, the deepest reusable checkpoint is $\ell(t; C) = \max(\{0\} \cup \{c \in C : c \le t\})$.
• The cost is the number of tokens to recompute: $r(t; C) = t - \ell(t; C)$.
• The objective is to minimize the expected recomputation cost $E[r(T; C)]$ given an overlap depth distribution $T$ with probabilities $p_t$.

Algorithmic Approach

1. Optimal Balanced Placement (Uniform/Worst-Case): The authors prove that under a uniform overlap distribution or in the worst-case scenario, balanced spacing (evenly distributed checkpoints) is optimal.
2. Distribution-Aware Dynamic Programming (DP): For non-uniform, real-world overlap distributions, the authors derive an exact $O(NM)$ dynamic programming algorithm.
   • The recurrence relation is: $dp[m, j] = \min_{1 \le s \le j} (dp[m-1, s-1] + w(s, j))$, where $w(s, j)$ is the cost of using $s$ as the deepest checkpoint for depths $s \dots j$.
   • This is solved efficiently using the monotone convex hull trick, leveraging the monotonicity of prefix sums.
3. Handling Distribution Drift: Since the true overlap distribution is unknown and must be estimated from historical data, the paper provides theoretical guarantees (via Lemma 1 and Theorem 3) that the value function is Lipschitz continuous with respect to the overlap distribution. This ensures that using an empirical histogram (or an exponentially weighted histogram for drifting distributions) yields a near-optimal schedule with bounded suboptimality.

Implementation Strategy

• Checkpointing: The system stores exact recurrent states at sparse positions. Upon a cache hit, the system resumes from the deepest stored checkpoint and recomputes the suffix exactly.
• Exactness: Because the recurrent update is deterministic, resuming from a checkpoint and replaying the suffix produces bit-identical outputs to a full prefill.
• Hybrid Integration: For hybrid models, the method stores the full attention KV cache (dense) but applies sparse checkpointing only to the recurrent layers.

3. Key Contributions

1. Theoretical Optimality: Proved that balanced spacing is optimal under uniform overlap and worst-case conditions (Theorem 1).
2. Exact Algorithm: Formulated sparse prefix caching as a weighted k-median problem and provided an exact $O(NM)$ dynamic programming solution (Theorem 2).
3. Robustness Guarantees: Demonstrated that the optimal schedule is stable under distribution misspecification. The suboptimality is bounded by the total variation distance between the empirical and true overlap distributions (Theorem 3), and extended to scenarios with distribution drift using exponential reweighting (Theorem 4).
4. Empirical Validation: Validated the method on real-world datasets (QuALITY, NarrativeQA, System Prompts), showing consistent improvements over fixed-budget baselines.

4. Experimental Results

Datasets and Setup

• Datasets:
  • QuALITY: Single document with multiple questions (high overlap near full length).
  • NarrativeQA: Similar structure with wide multimodal overlap.
  • System Prompts: Real system prompts combined with user queries (overlap concentrated at short lengths).
• Baselines: Compared against "No Cache," "KV Only," "Balanced" (even spacing), "Block" (fixed block size), "Logarithmic," and "$\sqrt{L}$-spaced" strategies.
• Hardware: Prototype runs on NVIDIA RTX 2080 Super with Qwen-3.5-0.8B model.

Findings

• Token Savings: The DP-optimal schedule consistently dominates deterministic baselines (Balanced, Block, Logarithmic) in terms of recurrent-work reduction factors.
• Low-Budget Advantage: The performance gap is largest at low checkpoint budgets where the overlap distribution is most non-uniform. For example, on System Prompts (where overlaps are short), the DP method places checkpoints densely near the divergence peak, achieving the same runtime as block caching with fewer checkpoints.
• Wall-Clock Time: On QuALITY and System Prompts, the DP-derived schedule yields measurable wall-clock speedups on representative layer groups. The runtime reduction correlates linearly with the number of tokens saved.
• Overhead: The overhead of loading checkpoints from CPU to GPU is small relative to the savings gained from reduced recomputation, resulting in a net runtime improvement.

5. Significance and Claims

The paper claims that sparse prefix caching unlocks a new design point for hybrid and recurrent LLM serving that was previously unavailable.

• Memory/Compute Trade-off: It offers a trade-off between memory usage and recomputation that dense KV caching cannot provide, specifically for architectures where state can be extracted and restored exactly.
• Exactness: Unlike approximate KV compression techniques, this method preserves exact outputs without altering the recurrent computation kernels.
• Applicability: It is most beneficial for workloads where requests share a substantial but not identical prefix (e.g., RAG with different questions, varying system prompts). It does not offer additional benefits for strictly append-only chat workloads (where the last state is sufficient) but incurs no overhead in such cases.
• Future Work: The authors acknowledge engineering challenges, specifically the need for efficient state extraction/restoration APIs in production runtimes and the integration of this placement strategy with admission/eviction policies (e.g., Marconi et al.). They also note that the current single-prefix theory could be extended to general prefix trees.