Efficient Serving for Dynamic Agent Workflows with Prediction-based KV-Cache Management
PBKV: Prediction-Based KV-Cache Management for Dynamic Agent Workflows #
1. TL;DR #
Dynamic multi-agent workflows (retry loops, runtime branches) break both LRU (temporal locality) and KVFlow (static DAG) for KV-cache management. PBKV predicts the next $K$ agent invocations with a tiny GraphSAGE model, scores each cache node by cross-workflow expected reuse, then evicts retired-cache-first and prefetches only idle GPU/PCIe capacity — up to 1.85× faster than LRU, with a Lipschitz graceful-degradation guarantee.
2. Q1 / Q2 / Q3 #
Q1 — 痛点 (the problem) #
Agentic workflows chain agents (Planner → Coder → Tester …) that share a large context — system prompt, tool/agent descriptions, and upstream history — so KV-cache reuse can exceed a 90% hit rate. But GPU memory is limited, so a cache-management policy must decide what to keep. Two existing families both fail on dynamic workflows:
- LRU (the SGLang/vLLM default) assumes temporal locality. But in a workflow a cache entry sits idle across several agent invocations (tool calls, agent transitions) before being reused; LRU evicts it prematurely during that idle window. Reuse here is structural, not temporal.
- KVFlow adapts Belady's optimal eviction to agents by assuming a predefined static agent step graph and evicting the cache whose next invocation is farthest away ("steps-to-execution"). This requires knowing the future invocation order a priori — which is false for runtime-dependent workflows (a Tester bounces code back for a rewrite; a Retriever spawns new sub-queries).
The core tension: the workflow structure exposes huge reuse potential, yet is inherently hard to predict (stochastic LLM decoding injects per-step uncertainty that propagates along the path), and prediction errors are disproportionately costly — one wrong eviction can force re-prefill of tens of thousands of accumulated tokens.
Q2 — 方法 (the method) #
PBKV reframes cache management as predicting future access patterns and builds three coordinated components (Figure 2):
- A multi-step predictor that fuses complementary signals (graph topology + per-request prefill semantics) and emits $K$ probability distributions over upcoming agents in one forward pass — avoiding the myopia of single-step predictors and the error accumulation of autoregressive rollout.
- Hierarchical eviction — reclaim retired cache (from terminated workflows, deterministically valueless) first; only when it is exhausted does a continuous lookahead score $\mathrm{Score}(c)$ drive eviction of active cache.
- Conservative prefetching — proactively load likely-reused nodes from host memory, but only into otherwise-idle GPU space and idle PCIe bandwidth, never displacing known-valuable active cache.
- Memory model: short-term = the current radix-tree cache (GPU); long-term / spill = HiCache host memory; episodic = offline invocation traces used to train the predictor and estimate the transition matrix $A$.
- Error recovery: when a prediction is wrong and a needed node was evicted, HiCache re-admits it from host memory (a PCIe transfer) instead of a full re-prefill; when retired cache is unavailable the policy falls back to score-driven eviction rather than degenerating to LRU.
- Dimensional consistency: $A_w(c)^\top P_w^{(k)} = P_w^{(k)}(\mathcal{O}_w(c)) \in [0,1]$ (Eq. 6) — a subset-event probability; multiplying by $s_w^{(k)}\in[0,1]$ and $\gamma^{k-1}$ keeps $\mathrm{Score}(c)$ a bounded expected count.
- Score = expected discounted miss count (Lemma G.1): $\mathrm{EMC}(c)=\sum_k \gamma^{k-1}\mathbb{E}[\text{misses on }c\text{ at }k]=\mathrm{Score}(c)$, grounding the ad-hoc score in Belady-style caching cost, removing circularity between the scoring rule and the cost it's judged against.
- Lipschitz continuity (Lemma G.2): $|\mathrm{Score}(c)-\widehat{\mathrm{Score}}(c)|\le \frac{1-\gamma^K}{2(1-\gamma)}\epsilon_c^\gamma \le \frac{\epsilon_c^\gamma}{2(1-\gamma)}$ — perfect prediction ⇒ zero deviation; multiplier is $K$-independent.
- Boundary condition: $\mathcal{R}(B)\ge 0$ (Eq. 20) by optimality of $E_B^\star$; and $\mathcal{R}(B)\to 0$ as $\epsilon\to 0$ (Theorem G.4) — perfect predictor recovers the optimum.
- Locality check: only $w\in\mathcal{W}_{\mathrm{act}}(c)$ enter $\epsilon_c^\gamma$; errors on non-accessing workflows leave $\mathrm{Score}(c)$ unchanged — consistent with radix-tree prefix structure.
- Boundary-localized regret: only nodes in the symmetric difference $\widehat{E}_B\triangle E_B^\star$ contribute (Eq. 21); a correctly-ranked node contributes zero regret regardless of its individual score error — regret is governed by ranking near the eviction frontier, matching Corollary G.3's ranking-stability view.
- $h_{txt}$ is free: the semantic stream reuses the post-norm hidden state of the last prefill token — a by-product already computed by the LLM's output head, universally accessible across models without extra instrumentation, so no separate encoder pass is added. (§4.1, Appendix B; the optimal extraction layer is model-specific per Figure 9, which is why they default to post-norm for portability.)
- Score refresh is essentially free:
ScoreUpdatecosts 1.53 µs (Table 4) vs a 12.34 ms decode step — ~8000× cheaper — so per-node scores can be refreshed on every state change without serializing the scheduler; all scheduler hooks run on the CPU thread or a separate CUDA stream, off the GPU critical path.
核心技术壁垒: the single hardest-to-replicate insight is embedding deterministic guardrails inside a probabilistic system. Prediction is untrustworthy, so PBKV never lets a prediction override a certainty: retired cache (a fact) always outranks any score (an estimate), and prefetching only spends resources that are provably idle. This asymmetry — trust facts, discount predictions — is what turns a noisy predictor into monotone, robust gains and underwrites the Lipschitz degradation bound. Copying the GraphSAGE model is easy; copying this discipline of when to ignore it is the load-bearing design.
Q3 — 结果 (the outcome) #
On three realistic workloads (HoVer+LangChain, SWE-bench+AutoGen, FinanceBench+CrewAI) with Qwen3-14B/32B: PBKV cuts end-to-end latency by up to 1.85× and lifts KV-cache hit rate by up to 2.55× over LRU on dynamic workflows; on the static workflow it beats SOTA KVFlow by up to 1.26× latency / 1.39× hit rate. The predictor is pluggable and the system is proven to degrade Lipschitz-continuously in prediction error.
3. 架构 / 方法图 #
The global call graph and the workflow abstraction #
A multi-agent app is a directed global call graph $G=(V,E)$ where nodes are agents and edges are admissible transitions; crucially $G$ may contain loops (retry paths). A workflow is one execution instance — an ordered sequence $a_1, a_2, \ldots$ over $G$. The code-generation example makes the dynamism concrete: the Tester conditionally routes back through Analyzer and Coder, so the realized path is unknowable in advance.
Paper's Figure 1: "A call graph for the code-generation task. The Tester conditionally triggers a retry path through Analyzer and Coder, i.e., a retry loop." This loop is exactly what defeats KVFlow's static "steps-to-execution": when the current agent is Analyzer, the Tester will be re-invoked two steps later, so its cache must be protected — but a static distance or a single-step predictor cannot see that.
System overview — the agent-serving loop #
Paper's Figure 2: "For each active workflow $w$, the predictor produces a $K$-step forecast over upcoming agent invocations. The forecast drives a shared scoring function $Score(c)$, which feeds both a hierarchical eviction policy and a conservative prefetching policy on the two-tier KV-Cache storage." The two-tier store is a Radix Tree on GPU memory plus HiCache on host memory; eviction and prefetching both consume the same forecast, sharing one reuse-scoring mechanism and one design philosophy (deterministic guardrails in a probabilistic system).
The per-turn control loop, driven on every workflow state change (new invocation or termination):
Planning / prediction architecture #
The predictor is the foundation. It must (i) exploit structural priors in $G$ and (ii) generalize to unseen runtime prefixes — so PBKV uses GraphSAGE (inductive, operates via neighborhood sampling/aggregation) rather than a transductive GCN.
Paper's Figure 3: "It fuses a topology-aware agent embedding from GraphSAGE ($h_{cur}$), an attention-based workflow prefix summary ($h_{path}$), and a semantic signal reused from prefill ($h_{txt}$), then jointly predicts the next $K$ agent probability distributions via an MLP." The reader should notice the three-stream fusion: $h_{\mathrm{cur}}$ answers "where am I in $G$", $h_{\mathrm{path}}$ answers "how did I get here" (attention over the prefix so two workflows at the same agent but different histories diverge), and $h_{\mathrm{txt}}$ answers "what is my intent" (reusing the last prefill token's post-norm hidden state, obtained essentially for free). The single MLP emits all $K$ steps at once, avoiding autoregressive error accumulation while keeping latency to one inference.
4. 作者证明 #
PBKV has a formal analytical model (unusual for an agent paper): a full algorithms-with-predictions smoothness analysis in Appendix G, culminating in Theorem 5.1 (= Theorem G.4).
Notation table #
| Symbol | Meaning | ||
|---|---|---|---|
| $G=(V,E)$ | global call graph; nodes = agents, edges = admissible transitions (may loop) | ||
| $\mathcal{W}^{\text{act}}(c)$ | set of active workflows associated with cache node $c$ | ||
| $A_w(c)\in\{0,1\}^{ | V | }$ | per-workflow access indicator: which agents of $w$ touch $c$ |
| $P_w^{(k)}$ | predicted step-$k$ agent-access distribution for workflow $w$ | ||
| $p_{w,\langle\text{END}\rangle}^{(j)}$ | conditional termination probability of $w$ at step $j$ | ||
| $s_w^{(k)}$ | cumulative survival probability of $w$ to step $k$ | ||
| $\gamma<1$ | confidence decay factor (far-step distrust) | ||
| $K$ | lookahead horizon (best at $K=3$) | ||
| $\epsilon_c^\gamma$ | node-local discount-weighted prediction error | ||
| $\widehat{E}_B, E_B^\star$ | PBKV's vs. ground-truth-optimal eviction set of budget $B$ |
Equation physical meaning #
Single-step reuse value aggregates, across all active workflows touching $c$, the probability their next invocation hits $c$:
$$\mathrm{Value}(c)=\sum_{w\in\mathcal{W}^{\text{act}}(c)}A_{w}(c)\cdot P_{w}$$
The sum (not min) is deliberate: value grows with both the number and probability of reusing workflows, so popular/global cache is protected automatically — the opposite of KVFlow's min over steps-to-execution, which sees timing but not popularity.
Survival discounts steps that likely never happen because the workflow already terminated (hazard-product):
$$s_{w}^{(k)}=\prod_{j=1}^{k-1}\bigl(1-p_{w,\langle\text{END}\rangle}^{(j)}\bigr)$$
The $K$-step score combines confidence decay ($\gamma^{k-1}$: far predictions less reliable) with survival weighting:
$$\mathrm{Score}(c)=\sum_{k=1}^{K}\gamma^{k-1}\sum_{w\in\mathcal{W}^{\text{act}}(c)}s_{w}^{(k)}\cdot A_{w}(c)\cdot P_{w}^{(k)}$$
The prefetch budget takes the tightest of space vs bandwidth so prefetch hides behind one decode step and steals nothing on-path: $S_{bw}=Bandwidth\cdot StepDuration$, $S=\min\{S_a, S_{bw}\}$.
Minimum checks (≥6) #
The headline guarantee (Theorem 5.1 / G.4) #
$$0\;\leq\;\mathcal{R}(B):=\mathcal{L}(\widehat{E}_{B})-\mathcal{L}(E_{B}^{\star})\;\leq\;\frac{1}{2(1-\gamma)}\sum_{c\in\widehat{E}_{B}\triangle E_{B}^{\star}}\epsilon_{c}^{\gamma}$$
The eviction cost regret is bounded linearly in prediction error, with a multiplier independent of $K$, of $|\mathcal{W}_{\mathrm{act}}|$, and of cache size. This is the formal justification for the pluggable predictor design: any future accuracy improvement directly tightens the bound. Note what is not bounded — there is no success-rate or convergence guarantee for the workflows themselves; the metric that could have been (and is) bounded is expected miss count, not task success.
5. 实验与数据 #
Testbed: 8× A6000 (48 GB) with NVLink, 20 GB/s PCIe; Qwen3-14B and Qwen3-32B (TP=2). Baselines: LRU (SGLang+HiCache default) and KVFlow (static-workflow only). Ablation variants: PBKV-LAE (lifecycle-aware eviction only) and PBKV-HE (adds hierarchical eviction, no prefetching).
Main results #
Paper's Table 1 (concurrency 72/24/48; static workload for KVFlow comparison). The load-bearing numbers, HoVer+LangChain, Qwen3-32B: LRU 189.66 s / 27.09% hit → Full PBKV 102.60 s / 69.10% hit (1.85× latency, 2.55× hit). The ablation ladder is the key story — most gain comes before the flagship prefetching:
| Policy (HoVer, Qwen3-32B) | Latency (s) | Hit rate (%) |
|---|---|---|
| LRU | 189.66 | 27.09 |
| PBKV-LAE (lifecycle only) | 146.67 | 44.91 |
| PBKV-HE (+ scoring) | 108.86 | 66.01 |
| Full PBKV (+ prefetch) | 102.60 | 69.10 |
The jump LAE→HE (44.9%→66.0%) is large; HE→full (66.0%→69.1%) is modest — by design, because conservative prefetching refuses to gamble known-valuable cache. On static FinanceBench, PBKV still beats KVFlow (53.44% vs 39.87% hit on Qwen3-32B) because cross-workflow sum aggregation captures popularity that KVFlow's min steps-to-execution misses.
Why it works — phase-by-phase cache trace #
Paper's Figure 6 (HoVer+LangChain, Qwen3-14B, concurrency 72). This is the load-bearing mechanism plot: all policies start ~80% (warm-up, no eviction); at ~30 s memory saturates and LRU collapses toward ~10% while PBKV-HE/full decline only gradually; from 55 s retired cache is released and PBKV-LAE recovers, but once retired cache drains (~80 s) LAE reverts to LRU while HE/full sustain ~50% via scoring. The reader should notice that lifecycle-awareness and scoring cover different phases — the two guardrails are complementary, not redundant.
Sensitivity #
Paper's Figure 5 (HoVer+LangChain, Qwen3-32B). Three monotonicity checks: (a) hit rate falls with concurrency but PBKV always beats LRU, and LRU OOM-crashes 3/10 at concurrency 84 and 10/10 at 96; (b) weaker backbones (R-GCN, Markov-N3) lower accuracy → monotonically lower hit rate, yet even the weakest backbone still beats prediction-free PBKV-LAE (robustness to predictor quality); (c) $K=3$ is optimal — smaller is myopic, larger dilutes the score with unreliable long-range predictions.
Conservative vs aggressive prefetching under noise #
Paper's Figure 12 (noise injection $\tilde P = (1-\lambda)P + \lambda U$). This figure justifies the conservative default: aggressive prefetching only wins at $\lambda=0,10\%$, is overtaken by conservative at $\lambda=20\%$, and falls below even the no-prefetch PBKV-HE at $\lambda\ge30\%$ — while conservative prefetching stays above PBKV-HE at every noise level (its lower-bound guarantee). Since real dynamic workflows are hard to predict, PBKV picks the stable option.
6. 论证链 (paper-internal) #
| Step | Claim | Support (paper-internal) |
|---|---|---|
| 1 | Cache management = predicting future access; workflow reuse is structural, not temporal | §1 argument: entries idle across invocations, LRU evicts prematurely |
| 2 | Static-DAG (KVFlow) assumption fails on runtime-dependent loops | §1 retry/sub-query examples; call graph with loop (Figure 1) |
| 3 | A single-step predictor is myopic; need $K$-step horizon with complementary signals | §3/§4.1: Analyzer→Coder→Tester example; GraphSAGE + $h_{path}$ + $h_{txt}$ fusion |
| 4 | Predicted $\mathrm{Score}(c)$ = expected discounted miss count | Lemma G.1 proof (Eqs. 9–10, survival weighting) |
| 5 | Retired cache is a certainty of zero value; must outrank probabilistic scores | §4.2.3 hierarchical eviction: deterministic guardrail in probabilistic system |
| 6 | Prefetch cost is deterministic, benefit probabilistic ⇒ spend only idle resources | §4.3 asymmetry argument; budget $S=\min\{S_a,S_{bw}\}$ |
| 7 | Under these designs, regret is linear in prediction error, $K$-independent multiplier | Lemma G.2 → Corollary G.3 → Theorem G.4 (Eqs. 11–25) |
| 8 | Therefore gains are large in the common case and degrade gracefully under bad predictors | §6.2 ablation ladder; §6.3 backbone sweep; §6.4 phase trace; Appendix E noise sweep |
7. 实现 cross-reference #
PBKV is built on SGLang + HiCache; the paper repeatedly points to "the supplementary code" for full hyperparameters/training config, and the concrete integration is described but line-level source is not embedded in the paper. Mark: [实现未公开] at file:line granularity (code referenced but not enumerated in the source text).
核心技术壁垒 (dedicated paragraph): the replication moat is not the ~350K-param GraphSAGE predictor — Table 3 shows several predictor families reach within a few points of it, and even a Markov-N3 backbone still helps. The moat is the coordination discipline: (1) a deterministic retired-cache tier that a prediction can never override; (2) a prefetch budget provably confined to idle space and idle bandwidth (min of two limits), activated only on pure-decode batches (>90% of batches); (3) a scoring rule that is simultaneously the deployed heuristic and the analyzed proxy cost (Lemma G.1), so the empirical system and the regret bound are the same object. Reproducing the numbers requires all three co-designed, not just the model.
关键实现细节 (easy-to-miss tricks):