Marconi: Prefix Caching for the Era of Hybrid LLMs

framework 2411.19379
prefix-cachinghybrid-ssmkv-cacheradix-treeeviction-policy

Marconi: Prefix Caching for the Era of Hybrid LLMs #

1. TL;DR #

Hybrid (Attention+SSM) LLMs break prefix caching because in-place SSM state updates forbid prefix roll-back — only exact-match hits work, so naive checkpointing floods the cache with huge, rarely-reused entries. Marconi admits SSM states judiciously (≤2 per sequence, chosen by a reuse taxonomy) and evicts by FLOP-per-byte, not recency, yielding up to 34.4× higher token hit rate and 71.1% (617 ms) lower P95 TTFT.


2. Q1 / Q2 / Q3 #

Q1 — 痛点 (the problem) #

Prefix caching reuses model states of common prefixes across requests to cut TTFT, tail TPT, and raise prefill throughput. It is a proven win for Transformers because KVs carry a sequence dimension: to keep the prefix $1{\dots}p$ of a cached sequence $1{\dots}q$ ($p

  1. SSM states are constant-sized regardless of token count.
  2. SSM states are updated in place — a sequence's end state cannot be rolled back to represent any of its prefixes.
  3. SSM states are 10–100× larger than a single token's KVs.
  4. The consequence is "all-or-nothing" reusability: a hit needs all prefix-token KVs for every Attention layer AND one SSM state that exactly matches the whole prefix for every SSM layer; the layer with the least reuse bottlenecks everything. To capture arbitrary future prefixes you must checkpoint SSM states at fine intervals (e.g. every 32–256 tokens), but each such entry is large and almost never hit. Measured: at block size 32, 25.0% of KV blocks are reused vs 0.4% of SSM states — a 65.3× gap (Fig. 3a); a single 10K-token sequence on a 7B model costs 17.4 GB, 3.3× a same-size Transformer (Fig. 3b). Result: cache thrashing on low-utility entries.

    Q2 — 方法 (the method) #

    Marconi is the first prefix-caching system for Hybrid LLMs, managing SSM states and KVs holistically in one radix tree (each node = a sequence's SSM state + KVs) rather than disaggregating by layer type. Two policies:

    • Judicious admission (§4.1): instead of admitting every token block, classify reuse into a two-type taxonomy — purely-input prefixes (system prompts, few-shot, long-doc QA; shared across many requests) and input-and-output prefixes (conversation history, agent trajectories; appended to, not branched from). Marconi caches only (a) the SSM state at the last decoded token (always, for input-and-output resume) and (b) the SSM state at a branch point discovered by a speculative insertion of the incoming prompt into the radix tree before prefill (for purely-input reuse). At most two SSM states per sequence are admitted.
    • FLOP-aware eviction (§4.2): because SSM state size is decoupled from compute savings (a fixed-size state can represent a 100-token or 10K-token prefix), size is a broken eviction proxy. Marconi scores each node by a blend of recency and FLOP-per-byte, iteratively evicting the lowest-utility nodes.

    核心技术壁垒 (the single hardest-to-replicate insight): the speculative-insertion + taxonomy admission mechanism. Everyone can compute FLOP-per-byte; the non-obvious, hard part is realizing that reuse likelihood — despite no knowledge of future requests — is predictable offline from the radix-tree topology of past requests: a node that just became a branch point signals a purely-input hot prefix, and a last-decoded token signals a conversation resume point. Encoding this as a cheap speculative insertion before prefill (so you know what to checkpoint during that same prefill pass) is the load-bearing design idea; without it, you either checkpoint everything (thrash) or checkpoint nothing (no reuse).

    Q3 — 结果 (the results) #

    • Token hit rate: 4.5× / 7.3× / 34.4× over vLLM+ on LMSys / ShareGPT / SWEBench (Fig. 7).
    • vs SGLang+ (same admission, LRU eviction): FLOP-aware eviction adds 45.6% / 19.0% / 219.7% P95 hit-rate win (Fig. 8) — largest where sequences are longest (SWEBench).
    • P95 TTFT: up to 36.9% / 73.2% / 46.8% (281.4 / 106.3 / 617.0 ms) vs vanilla; 36.1% / 71.1% / 46.8% vs vLLM+ (Fig. 9).
    • Gains grow with SSM-layer ratio (1:2→1:8 lifts win to 2.6×) and state dimension (16→128 lifts win 5.7×→35.4×) — trends aligned with newer models.

    3. 架构 / 方法图 #

    The end-to-end request lifecycle Marconi imposes: on request arrival, tokenize, speculatively insert the input into the radix tree to detect whether a new intermediate (branch) node appears; schedule prefill; during prefill materialize the SSM state at the branch point (via chunked state passing or two-pass) and, after decode, checkpoint the last-token state; on the next matching request, look up the longest exact-match prefix and skip its prefill.

    sequenceDiagram participant R as Request participant T as Radix Tree (bookkeeper) participant S as Scheduler participant P as Prefill (GPU) participant C as Holistic Cache (SSM+KV per node) R->>T: speculative insert(input tokens) T-->>S: branch node created? (purely-input signal) S->>C: lookup longest exact-match prefix C-->>P: fetch cached KVs + SSM state (on hit) S->>P: prefill remaining tokens P->>C: checkpoint SSM state @ branch point (if any) P->>C: checkpoint SSM state @ last decoded token (always) Note over C: FLOP-aware eviction if full
    (evict lowest recency+α·flop_eff)

    The paper's own overview of the reuse mechanism and the fine-grained-checkpointing pathology:

    Figure 2: prefix caching reuses common-prefix states; fine-grained checkpointing yields sparsely-hit entries

    Paper's Figure 2 (caption: "Prefix caching reuses model states of common prefixes (green) across requests, accelerating inference. Fine-grained checkpointing results in many sparsely-hit entries (blue)."). Green shows the profitable shared prefix; the blue entries are the low-utility SSM states that flood a naive cache — the whole design exists to stop admitting them.

    Figure 4: speculative insertion checkpoints SSM state at the branch point and last decoded token

    Paper's Figure 4 (caption: "Marconi performs a speculative insertion to check if inserting the prefill segment of a sequence results in an intermediate node. If so, the SSM states at the branch point are checkpointed. States at the last decoded token are checkpointed in any case."). This is the load-bearing architecture figure: it shows the radix tree as the unified bookkeeper and how the speculative insertion decides which of the (at most two) SSM states to materialize. Note states are drawn on nodes for clarity but conceptually attach to edges — each edge holds KVs of its tokens plus the SSM state for all tokens prior to the edge's last token.


    4. 作者证明 #

    The paper has a lightweight but load-bearing analytical model (the FLOP-efficiency metric + per-layer memory/FLOP formulas in Appendix A / Table 1). It is not a full queueing/throughput model, so this is partially formal.

    Notation table #

    SymbolMeaning
    $L$Sequence length (tokens)
    $D$Model dimension ($d_{model}$)
    $N$SSM state / feature dimension ($d_{state}$)
    $\mathit{flop\_efficiency}$Redundant FLOPs avoided by reuse per byte of cached state
    $S(n)$Utility score of radix node $n$ (eviction key)
    $\mathit{recency}(n)$Normalized last-access timestamp of $n$, in $(0,1)$
    $\alpha$Tunable weight trading recency vs FLOP efficiency ($\alpha=0 \Rightarrow$ LRU)

    Equations & physical meaning #

    FLOP-per-byte of a cache entry (Eq. 1):

    $$\mathit{flop\_efficiency}=\frac{\text{Total FLOPs across layers}}{\text{Memory consumption of all states}}$$

    Numerator = redundant FLOPs across Attention + SSM + MLP layers skipped by reusing this entry; denominator = memory of all stateful (Attention + SSM) states. Why divide rather than subtract: it normalizes compute savings against the cache space consumed, so a short KV-heavy entry and a long SSM-heavy entry become comparable on a common per-byte basis.

    Eviction utility (Eq. 2):

    $$S(n)=\mathit{recency}(n)+\alpha\cdot\mathit{flop\_efficiency}(n)$$

    A linear blend (sum, not min/product) so one factor can compensate for the other; $\alpha=0$ recovers pure LRU, larger $\alpha$ favors compute-dense long-sequence entries. Both terms are normalized to $(0,1)$ across all tree nodes; child FLOP savings are computed relative to the parent's so ancestors are not double-counted.

    Per-layer scaling (Table 1, 7B model with $D=4096$, $N=128$):

    $$\text{Attention FLOPs/byte} = L+2D = L+8192, \qquad \text{SSM FLOPs/byte} = L\cdot(6D/N+8+5/DN)\approx 200L$$

    Six minimum checks #

    1. Units: numerator FLOPs, denominator bytes → FLOPs/byte. Consistent; Eq. 2 normalizes both terms to dimensionless $(0,1)$ before summing — internally consistent.
    2. Boundary $\alpha=0$: $S(n)=\mathit{recency}(n)$ → pure LRU, matching the stated fallback and the SGLang+ baseline. ✓
    3. Boundary short $L$: Attention $L+8192$ is offset-dominated (≈constant), so KV FLOP-efficiency is "near-constant" — exactly why Transformer systems can ignore it. ✓
    4. Boundary long $L$: SSM term $\approx 200L$ grows linearly and unbounded while its state size $2DN$ is fixed → FLOP-per-byte of SSM entries diverges upward with $L$; hence long sequences dominate eviction value. Matches Fig. 5's steeper slope for higher SSM ratio. ✓
    5. Monotonicity/optimum: $S$ is monotone increasing in both recency and flop_efficiency; there is no interior optimum — eviction is a greedy boundary rule (drop the argmin until space frees). The "optimum" is the $\alpha$ that maximizes measured hit rate, found by grid search, not by closed-form. ✓ (implicit)
    6. First-order sanity vs reported numbers: block-size-32 KV-vs-SSM reuse (25.0% vs 0.4% = 65.3×) and the 17.4 GB / 3.3× memory blowup are consistent with property 3 ($2DN$ SSM state ≫ per-token $4D$ KV once $N$ is large): SSM/KV size ratio $= 2DN / (2\cdot 4D \cdot \text{block}) = N/(4\cdot\text{block})$, i.e. $128/(4\cdot16)=2$× per block, compounding across many SSM layers vs few Attention layers. ✓
    7. What a fuller model would clarify: the paper gives no closed-form expression relating $\alpha$, workload skew, and hit rate — hence the empirical grid search. A queueing/hit-rate model would predict the "moderate-contention sweet spot" (Fig. 11) analytically instead of by sweep.


      5. 实验与数据 #

      Setup: p4d.24xlarge (8×A100-40GB), FP16; main model a 7B Hybrid with {4,24,28} {Attn,SSM,MLP} layers; TTFT measured on Jamba-1.5-Mini (state dim 128) via vLLM on 4×A100-40GB. Workloads: LMSys (long outputs), ShareGPT (short outputs), SWE-Bench/SWE-Agent (agentic, widest length distribution). Baselines: vanilla (no caching), vLLM+ (fine-grained block-32 checkpointing), SGLang+ (Marconi's admission but LRU eviction). Metric: token hit rate = tokens skipping prefill / total input tokens (a good proxy for FLOP saved since prefill is compute-bound); no downstream-quality metric because reuse is exact.

      Figure 7: token hit rate vs vLLM+ across LMSys, ShareGPT, SWEBench

      Paper's Figure 7 (caption: "Comparison with vLLM+. With judicious cache admission, Marconi utilizes the limited cache space to retain states with higher utility, improving the token hit rate significantly over vLLM+."). This isolates the admission win (4.5× / 7.3× / 34.4×): rejecting low-utility SSM states frees capacity for high-utility ones. The SWEBench blow-up (34.4×) reflects its long agent trajectories, where each admitted-then-never-reused block was maximally wasteful under vLLM+.

      Figure 8: token hit rate vs SGLang+ (LRU eviction)

      Paper's Figure 8 (caption: "Comparison with SGLang+, which uses LRU as its eviction policy. Marconi balances recency and FLOPs efficiency, improving the token hit rate significantly, especially for workloads with longer context."). Because SGLang+ shares Marconi's admission, this figure isolates the eviction win alone: 45.6% / 19.0% / 219.7% P95. The gap tracks sequence-length spread — SWEBench (hundreds→tens-of-thousands of tokens) benefits most, ShareGPT (<2K) least, exactly the regime where FLOP-per-byte discrimination matters.

      Figure 10: FLOP-aware eviction trades short-sequence hit rate for long-sequence hit rate

      Paper's Figure 10 (caption: "Marconi achieves a higher hit rate for longer sequences while sacrificing the hit rate for some shorter sequences (a), although the degradation in TTFT for shorter sequences is minimal (b)."). The honest trade-off: on one SWEBench trace Marconi hits 32.7% vs SGLang+ 16.4% (+99.4%, +90.3% FLOP saved), but loses up to −3.0% for <7K-token sequences and is 6.3% worse on P5 TTFT — a mere 2.1 ms in absolute terms because Hybrid models prefill short sequences fast. This is the rare paper that plots and defends where it loses.

      Figure 12: gains grow with SSM-layer ratio and state dimension

      Paper's Figure 12 (caption: "Marconi performs better for models with higher ratios of SSM layers and larger SSM state dimensions."). Layer ratio 1:2→1:8 lifts the win over vLLM+/SGLang+ to 2.6×/59.7%; state dim 16 (Mamba1)→128 (Mamba2) lifts the vLLM+ win 5.7×→35.4×. Since newer Hybrids trend toward more SSM layers and larger states, Marconi's advantage is forecast to grow, not shrink.

      §6-required workload-regime breakdown (where it wins / loses) #

      Workload regimeMarconiBaselineWhy
      Short prompts, low concurrency (ShareGPT, <2K tok)small win (7.3× vs vLLM+; only 19.0% vs SGLang+)vLLM+/SGLang+ competitiveFLOP-per-byte spread is small; recency alone nearly suffices
      Long prompts, high concurrency (SWEBench, up to tens of K tok)largest win (34.4× vs vLLM+; 219.7% vs SGLang+)thrashes / evicts high-FLOP entrieslong SSM entries have huge FLOP-per-byte; admission + FLOP eviction both fire
      Short sequences (<7K tok) within a long-tailed traceloses up to −3.0% hit rate, P5 TTFT +6.3% (2.1 ms)SGLang+ (LRU) slightly betterdeliberately sacrificed to keep long high-FLOP entries
      Pure TransformertieidenticalKV FLOP-efficiency near-constant → nothing to exploit
      High contention (60 GB cache)modest 24.3% wintoo little capacity to cache useful prefixes at all
      Moderate contention (~100 GB)peak 68.3% wineviction decisions are most consequential here

      6. 论证链 #

      StepClaim (paper-internal)Support
      1SSM states update in place, are constant-sized, and are 10–100× a single token's KVs (§3 properties 1–3)§2.1 mechanism + Appendix A memory formulas ($2DN$ vs $4D$/token); Table 1
      2Therefore Hybrid prefix reuse is "all-or-nothing" — only exact whole-prefix matches hit, forcing fine-grained checkpointing to catch arbitrary prefixesderivation from step 1 (roll-back impossible for SSM, so bottleneck layer = SSM)
      3Fine-grained checkpointing floods the cache with large, sparsely-hit entries (0.4% SSM reuse vs 25.0% KV; 17.4 GB/seq)Fig. 3(a), Fig. 3(b) measurements
      4Reuse likelihood is nonetheless predictable offline from radix-tree topology (branch node ⇒ purely-input; last-decoded token ⇒ input-and-output) → admit ≤2 SSM states/seq via speculative insertion§4.1 taxonomy + Fig. 4 mechanism
      5Admission alone (holding eviction at LRU) already lifts hit rate massively (Marconi vs vLLM+)Fig. 7: 4.5×/7.3×/34.4×
      6SSM state size is decoupled from compute savings, so LRU/size-based eviction mis-ranks entries; FLOP-per-byte (Eq. 1–2) ranks correctly§4.2 argument + Fig. 5 (FLOP-efficiency slope by SSM ratio)
      7FLOP-aware eviction adds further gains, concentrated on long-sequence workloads, at a small deliberate cost to short sequencesFig. 8 (45.6/19.0/219.7%) + Fig. 10 (long win, short −3.0%)
      8Net effect scales up with SSM ratio and state dimension → advantage grows for future modelsFig. 12 sweeps

      7. 实现 cross-reference #

      Marconi is open-sourced (github.com/ruipeterpan/marconi; Zenodo DOI 10.5281/zenodo.14970139), and the artifact appendix maps the design to concrete files:

      • Core cache + eviction policies: radix_cache_hybrid.py — new eviction variants are added via an evict_policy_version hook (Appendix B.7). Policy V1 = SGLang+ (LRU), V2 = Marconi (FLOP-aware), V3 = offline-optimal static-$\alpha$ oracle (present in code, results excluded from paper).
      • Experiment driver / default model config (NVIDIA Attention-Mamba2 7B Hybrid): policy_exploration.py, invoked by run_all_experiments.sh (Appendix B.5).
      • Trace generation: utils/generate_trace.py; per-figure plotters under /plotting (e.g. Fig. 7 token_hit_rate.py, Fig. 8 sglang_comparison.py, Fig. 9 ttft.py, Fig. 10 fine_grained_analysis.py) (Appendix B.6).

      核心技术壁垒 (dedicated paragraph): The reproducibility bottleneck is not the FLOP-efficiency arithmetic — it is the speculative-insertion-before-prefill step that couples the radix-tree bookkeeper to the prefill scheduler. To checkpoint the right SSM state during a prefill pass, the system must, before that pass runs, insert the incoming input into the tree, detect whether it creates a branch node, and — if so — arrange to materialize the SSM state at the branch offset. Materializing that mid-sequence SSM state itself needs one of two model-dependent mechanisms (§4.1 "Obtaining states during prefill"): for chunked-state-passing SSMs (e.g. Mamba2/SSD), cache the second-to-last chunk's state (checkpoint token 64 to approximate token 80 at chunk size 32); for others, a two-pass prefill (pass 1 → prefix state, pass 2 → remainder). Getting this offset/chunk bookkeeping right, and integrating it into the scheduler without stalling, is the hardest-to-replicate part.

      关键实现细节 (easy-to-miss tricks):

      1. Only $\leq 1$-child nodes are evictable (§4.3-1): multi-child nodes are hot purely-input prefixes and are protected; when an intermediate single-child node is evicted, its SSM state is released but its KVs are absorbed by the child, so no KV coverage is lost.
      2. On a hit, only the accessed node's timestamp updatesnot all ancestors (§4.3-2), unlike vLLM/SGLang. This is safe precisely because ancestors' SSM states are never reused and their KVs get subsumed by children on eviction; it keeps recency bookkeeping cheap and correct. Combined with the asynchronous CPU-parallel grid search that tunes $\alpha$ during a bootstrap window of 5–15× the pre-first-eviction request count (finishing in "a few seconds," often faster than one request's prefill+decode).