Robust KV Cache Management for LLM Serving under Output Token Length Uncertainty

Abstract: KV cache memory is a primary bottleneck in modern LLM serving systems deployed on GPU clusters. A fundamental challenge is that the KV cache must be reserved upon request arrival, while the output token length remains unknown until generation completes. Under-reservation triggers preemption -- forcing termination and recomputation of requests and incurring significant overhead -- whereas over-reservation wastes memory and reduces throughput. This creates a central trade-off between memory efficiency and preemption risk. We present a robust KV cache management framework for LLM serving that jointly optimizes GPU parallelism configuration, KV cache reservation per request class, request routing across heterogeneous serving groups, and prefix caching for shared prompts. The framework incorporates latency SLO constraints and captures the interaction between memory allocation, throughput, and queueing delay. To address output token length uncertainty and workload distribution shift, we develop a Wasserstein distributionally robust optimization (DRO) formulation together with a scalable block coordinate descent algorithm for the resulting mixed-integer problem. Our analysis reveals a critical fractile structure that automatically adapts reservation quantiles to different preemption and memory cost regimes without manual tuning. Trace-driven evaluation on production LLM workloads, including BurstGPT, Azure, and ShareGPT traces, demonstrates up to 56\% lower cost than fixed-quantile reservation baselines while maintaining competitive P99 latency, goodput, and SLO violation rates across diverse operating regimes.

• The paper introduces a Wasserstein distributionally robust optimization framework that jointly selects KV reservations, GPU parallelism, request routing, and prefix caching under output-length uncertainty.
• Its BCD-DRO algorithm solves production-scale instances with 15 request classes, 12 configurations, 2,000 samples, and 48 GPUs in about 258 seconds, enabling rolling-horizon optimization.
• Trace-driven evaluation reports up to 56% lower cost than fixed-quantile policies, under 1.3% robustness overhead, and substantially better tail latency and SLO stability under distribution shift.

Problem and Motivation

KV cache memory is a first-order bottleneck in LLM serving: for a 70B-parameter model, each token consumes roughly 2.5 MB of KV storage, so a 2,000-token response requires about 5 GB. Because output length is unknown at admission time, the operator must reserve KV capacity before generation begins. Under-reservation triggers preemption—eviction from the running batch, discarding completed tokens, and full recomputation—which inflates TTFT, ITL, and tail SLO violations; over-reservation idles memory and reduces achievable batch size and goodput. The paper "Robust KV Cache Management for LLM Serving under Output Token Length Uncertainty" (2607.16892) formulates this trade-off as a joint decision problem over four control variables: GPU parallelism configuration ($z$), per-class KV reservation ($r$), request routing across heterogeneous serving groups ($\pi$), and prefix caching ($\delta$).

The authors identify three structural difficulties: workload heterogeneity (output lengths span tens to thousands of tokens across chat, RAG, and code-generation classes), non-stationarity of workload distributions, and system heterogeneity arising from discrete TP/PP configurations and optional prefill/decode disaggregation. Existing runtime systems—PagedAttention/vLLM, Orca continuous batching, DistServe, Sarathi-Serve—optimize execution once allocations are fixed but leave reservation and routing to operator heuristics such as fixed quantiles (Mean, P90, P95, P99, Max). The paper's contribution is a decision-policy layer that co-optimizes all four knobs under explicit robustness guarantees.

Wasserstein DRO Formulation

Output length $\xi_i$ for class $i$ follows an unknown distribution $P^i$, learned from $N$ historical samples via an empirical distribution $\hat{P}_N^i$. To guard against distribution shift, the authors adopt a type-1 Wasserstein ambiguity set $$P_\varepsilon^i = \{P : W_1(P, \hat{P}_N^i) \le \varepsilon\}$$ and minimize the worst-case expected cost subject to four constraint families: a GPU budget, routing consistency, admission probabilities summing to one, and—critically—a memory-capacity constraint coupling reservation to throughput through Little's law, with a $r$0 safety factor for concurrency fluctuation. Latency SLOs are enforced via slack variables penalized in the objective, with waiting times modeled by the M/G/1 Pollaczek–Khinchin formula.

A notable modeling choice is the deliberate asymmetry between aggregate and per-request treatment of uncertainty: service time uses $r$1 because throughput averages variability across concurrent requests (law of large numbers), whereas reservation decisions are DRO-protected because preemption affects individual requests. The stochastic cost is piecewise linear and asymmetric: $r$2 for preemption overflow and $r$3 for wasted reservation.

Applying Kantorovich duality yields a finite mixed-integer bilinear program with $r$4 auxiliary variables. For piecewise-linear costs on bounded support $r$5, the worst-case distribution places mass only at empirical samples or support boundaries, encoded by epigraph constraints. The dual variable $r$6 acts as a transportation-cost penalty controlling how far adversarial mass can shift toward the extremes.

Critical Fractile Structure

The central theoretical result is that the optimal reservation satisfies a newsvendor-type critical fractile:

$r$7

where $r$8 is the preemption-to-waste cost ratio and $r$9 is the worst-case output-length CDF within the ambiguity set. Setting the derivative of expected cost to zero gives $\pi$0, where $\pi$1 is the buffer. The implication is direct: as preemption becomes relatively more expensive, the optimal quantile rises automatically, eliminating manual quantile tuning. The authors are careful to note that both the classical critical ratio and its distributionally robust form are known in operations research; their contribution lies in coupling the reservation to routing, parallelism configuration, and the shared-memory capacity constraint. They also observe that under bounded support the robust reservation strictly exceeds the empirical fractile by an amount growing with $\pi$2, distinguishing it from sample-average approximation.

BCD-DRO Algorithm

Directly solving the MIBLP does not scale, so the authors propose block coordinate descent cycling through three blocks: alternating LPs over continuous variables $\pi$3, a closed-form binary update for prefix caching ($\pi$4), and enumeration over feasible integer configurations under the GPU budget. Convergence to a blockwise-optimal solution in finite iterations follows from monotone objective decrease over a finite feasible set.

Scalability results are strong: the monolithic MIBLP times out beyond 300 s at 6 classes × 6 configurations × 300 samples, while BCD-DRO solves a production-scale instance (15 classes, 12 configurations, $\pi$5, 48-GPU cluster) in approximately 258 s—compatible with 5–10 minute rolling-horizon re-optimization. At larger scales ($\pi$6, $\pi$7), runtime approaches the re-optimization budget, which the authors concede would require warm-starting or anytime execution.

Evaluation Results

Trace-driven evaluation uses BurstGPT (1.4M requests, mean output 125 tokens, P99 1,586), Azure LLM 2024 (44M requests), and ShareGPT (368K responses). Key findings include:

• No fixed quantile is universally optimal. Each baseline matches the critical fractile only in a narrow regime (P90 near $\pi$8, P95 near $\pi$9). DRO achieves up to 56% lower cost than fixed-quantile baselines (at $\delta$0: cost 1,919 vs. P90's 4,328), with gains ranging from 0–5% at $\delta$1 to 44–56% at $\delta$2. Conversely, conservative baselines become prohibitively expensive at low $\delta$3 (Mean costs 42,811 vs. DRO's 1,270 at $\delta$4).
• Cross-regime stability over per-regime optimality. At individual operating points, regime-matched baselines can marginally beat DRO on latency—at $\delta$5, LP$\delta$6 achieves 1.5% lower P99 and P90 0.7% lower—but each fails badly elsewhere. DRO trades small per-regime overhead for bounded cost across all regimes.
• Robustness premium is small. Increasing the Wasserstein radius adds less than 1.3% cost overhead even at $\delta$7; $\delta$8 provides shift protection at under 0.5% overhead.
• Ablations quantify component contributions. Removing routing (uniform $\delta$9) is catastrophic—cost rises from 13,476 to 93,878 and SLO violations from 26.4% to 93.8%—because heterogeneous classes exceed the KV capacity of some configurations, making routing a feasibility constraint rather than a tuning parameter. Replacing DRO with SAA roughly doubles cost under a 2.5× shifted evaluation distribution (24,864 vs. 13,476); across shift factors up to 3.0×, SAA's P99 degrades from 5.4 s to 18.0 s versus DRO's 3.0 s to 9.1 s, with SLO violations reaching 67.7% versus 35.6%. Disabling prefix caching leaves cost unchanged (it enters only the queueing/capacity model) but raises P99 from 8.12 s to 9.95 s.
• Insensitivity to secondary weights. Sweeping the SLO penalty, rejection cost, and headroom factor by up to 5× leaves the reservation quantile fixed at $\xi_i$0 with flat cost and tail latency, since these weights affect only admission and routing. The sole behavioral lever is $\xi_i$1 when rejection becomes cheaper than serving, at which point the optimizer deliberately sheds the most expensive class.

Limitations and Open Questions

Several limitations are acknowledged or evident. The M/G/1 queueing model assumes Poisson arrivals and uses a service-time coefficient of variation computed on training samples and held fixed—an approximation whose accuracy under bursty production traffic is not independently validated. The simulator applies a flat 2.5× service-time penalty for preempted requests, a coarse proxy for recomputation overhead whose realism depends on the serving runtime. Evaluation is simulation-based; integration with production runtimes such as vLLM remains future work, as do online adaptation via streaming statistics and extension to energy-aware or geographically distributed scheduling. Scalability beyond ~15 classes and ~2,000 samples per class strains the re-optimization budget without warm-starting. Finally, the framework presumes meaningful within-class homogeneity of output-length distributions; the sensitivity of the approach to class-definition choices is not explored.

Conclusion

This paper contributes a unified, uncertainty-aware control-plane policy for KV cache management that jointly determines GPU partitioning, per-class reservation, routing, and prefix caching under latency SLO constraints. Its theoretical core—the critical fractile $\xi_i$2 derived under a Wasserstein ambiguity set—converts manual quantile selection into automatic adaptation to the preemption/waste cost structure, and the BCD-DRO algorithm makes the resulting mixed-integer program solvable at production scale. Trace-driven results show up to 56% cost reduction over fixed-quantile heuristics with stable tail latency across regimes, and ablations confirm that routing feasibility and Wasserstein robustness carry substantial weight under workload drift. The main open questions concern validation against real serving runtimes, scalability of re-optimization at larger instance sizes, and robustness to the class-partition assumption underlying the per-class ambiguity sets.