Technical Summary: Efficient Serving for Dynamic Agent Workflows with Prediction-based KV-Cache Management (PBKV)

1. Problem Statement

Large Language Model (LLM)-based multi-agent systems decompose complex tasks into sequences of collaborating agents (e.g., Planner, Coder, Tester). These agents share substantial context, creating significant opportunities for Key/Value (KV) cache reuse. However, existing inference engines (e.g., vLLM, SGLang) typically rely on the Least Recently Used (LRU) policy, which assumes temporal locality. In agentic workflows, cache reuse is driven by the structural nature of the workflow rather than temporal proximity; a cache entry may remain idle across multiple agent invocations (e.g., during tool calls or loop iterations) before being reused, causing LRU to evict it prematurely.

Alternative approaches like KVFlow assume a static, predefined agent step graph (a global DAG) to determine eviction based on "steps-to-execution." This assumption fails in realistic dynamic workflows, where the sequence of invoked agents depends on runtime context (e.g., a Tester conditionally triggering a retry loop or a Retriever spawning new sub-queries).

The core challenge is twofold:

1. Predictability: Agentic workloads are inherently stochastic due to LLM decoding, making future agent invocations difficult to predict.
2. Cost of Error: Prediction errors are disproportionately costly. A cache miss in a multi-agent setting may force the re-prefill of tens of thousands of tokens accumulated across the entire workflow context.

The paper asks: How can we design a KV-Cache management framework for dynamic multi-agent systems that is robust to prediction errors yet benefits continuously from imperfect predictors?

2. Methodology: PBKV System

The authors propose PBKV (Prediction-Based KV-Cache Management), a system that predicts future agent invocations to evaluate the reuse value of existing KV-cache entries. PBKV operates on a two-tier storage hierarchy (GPU memory and host memory via HiCache) and consists of three main components:

A. Workflow Predictor

Instead of assuming a static graph, PBKV employs a predictor that fuses complementary signals to forecast the next $K$ agent invocations for each active workflow.

• Architecture: Based on GraphSAGE, the predictor is inductive, allowing it to generalize to unseen workflow prefixes at runtime.
• Signal Fusion:
  1. Topology-aware Agent Embedding ($h_{cur}$): Encodes the current agent's identity and its neighborhood in the global call graph $G$ (derived from offline traces).
  2. Attention-based History Aggregation ($h_{path}$): Summarizes the workflow prefix using scaled dot-product attention, distinguishing workflows that arrive at the same agent via different paths.
  3. Semantic Signal ($h_{txt}$): Reuses the post-norm hidden state of the last prefill token to capture request-specific intent, distinguishing requests with identical prefixes but diverging goals.
• Output: The predictor emits $K$ probability distributions over upcoming agents in a single forward pass, avoiding autoregressive error accumulation.

B. Lookahead KV-Cache Eviction

PBKV replaces simple temporal heuristics with a hierarchical eviction policy driven by a "lookahead reuse score."

• Lifecycle-Aware Reclamation: Cache nodes associated with terminated workflows ("retired cache") are identified deterministically and evicted first, as they have zero reuse potential. This provides a baseline improvement independent of prediction quality.
• Score-Driven Active Eviction: For active cache, a continuous score is computed based on predicted reuse value:
  $$\text{Score}(c) = \sum_{k=1}^{K} \gamma^{k-1} \sum_{w \in \mathcal{W}_{act}(c)} s_w^{(k)} \cdot A_w(c) \cdot P_w^{(k)}$$
  Where:
  • $A_w(c)$ is the access indicator (which agents in workflow $w$ accessed node $c$).
  • $P_w^{(k)}$ is the predicted probability of invoking an agent that accesses $c$ at step $k$.
  • $s_w^{(k)}$ is the cumulative survival probability (workflow not terminated).
  • $\gamma < 1$ is a confidence decay factor to discount unreliable long-horizon predictions.
• Hierarchical Strategy: The system prioritizes evicting retired cache (score = 0) before touching active cache. This ensures graceful degradation: if the predictor is poor, the system still benefits from lifecycle awareness; if accurate, it leverages multi-step lookahead.

C. Conservative Prefetching

To hide transfer latency, PBKV proactively loads likely-to-be-reused cache from host memory to GPU.

• Conservative Principle: Prefetching is restricted to otherwise-idle GPU space (free space + retired cache) and idle PCIe bandwidth (during pure-decode batches).
• Risk Avoidance: PBKV never evicts active cache to make room for prefetched data. This prevents the system from displacing known-valuable cache for speculative gains, ensuring that prediction errors do not degrade performance below the baseline.
• Ranking: Candidates are ranked by one-step reuse value (Eq. 1) using a greedy heuristic under a strict budget.

3. Key Contributions

1. Formalization: The paper formalizes KV-Cache management in multi-agent serving as a problem of predicting future access patterns, providing design guidelines for predictors in this setting (multi-signal fusion, multi-step horizon).
2. PBKV Framework: The authors propose a system combining hierarchical eviction (deterministic lifecycle reclamation + probabilistic lookahead scoring) and conservative prefetching.
3. Theoretical Guarantee: The paper proves that PBKV's performance degrades Lipschitz-continuously with prediction error. Specifically, the eviction cost regret is bounded linearly by the prediction error, ensuring graceful degradation even with imperfect predictors.
4. Implementation & Evaluation: PBKV is implemented on SGLang with HiCache and evaluated across diverse workloads and LLMs.

4. Experimental Results

Experiments were conducted on three benchmarks: HoVer (fact verification, dynamic loops), SWE-bench (code generation, retry loops), and FinanceBench (static workflow). Models used included Qwen3-14B and Qwen3-32B.

• Dynamic Workflows:
  • PBKV achieved up to 1.85× speedup in end-to-end latency and 2.55× improvement in KV-Cache hit rate compared to LRU.
  • The "Lifecycle-Aware Eviction" (PBKV-LAE) alone improved hit rates by 1.66× over LRU, demonstrating the value of simple lifecycle signals.
  • The full system (PBKV-HE + Prefetching) sustained high hit rates (~50-60%) even under heavy memory pressure where LRU dropped to ~10%.
• Static Workflows:
  • On the static FinanceBench workload, PBKV outperformed the SOTA baseline KVFlow by 1.26× in latency and 1.39× in hit rate. The authors attribute this to PBKV's lifecycle-aware eviction (preventing premature eviction of active cache) and cross-workflow aggregation of reuse scores.
• Robustness:
  • Sensitivity analysis showed that even with weaker predictor backbones (e.g., R-GCN, Markov models), PBKV consistently outperformed LRU and the prediction-free variant (PBKV-LAE).
  • Conservative prefetching provided stable gains, whereas aggressive prefetching (evicting active cache) became counterproductive under moderate prediction noise.

5. Significance and Claims

The paper claims that PBKV addresses a critical gap in serving dynamic agentic workflows, where existing static assumptions (KVFlow) or temporal heuristics (LRU) fail.

• Robustness over Optimality: The authors emphasize that PBKV is designed to be robust to the inherent uncertainty of agentic systems. By combining deterministic guardrails (lifecycle reclamation) with conservative speculation (prefetching only on idle resources), the system ensures that imperfect predictions do not harm performance.
• Graceful Degradation: The theoretical proof of Lipschitz continuity establishes that PBKV's performance is stable; as predictor accuracy improves, performance improves linearly, but poor accuracy does not cause catastrophic failure.
• Generalizability: The predictor is designed as a pluggable module. The system's design principles (hierarchical eviction, conservative prefetching) are independent of the specific predictor used, allowing the framework to benefit from future advancements in workflow prediction.

The authors conclude that PBKV significantly outperforms both LRU and the state-of-the-art static-aware baseline across dynamic and static workloads, validating the efficacy of prediction-based, conservative cache management in multi-agent serving.