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

2026-07-18

AI Summary

This work addresses the challenge of KV cache management in large language model (LLM) serving, where uncertain output lengths lead to either excessive memory waste or request preemption due to inaccurate cache reservation. The authors propose a robust KV cache management framework that jointly optimizes GPU parallelism configuration, dynamically reserves cache per request category, routes requests across heterogeneous service groups, and shares prefix caches. Output length uncertainty is modeled via Wasserstein distributionally robust optimization, revealing a critical quantile structure that adapts to varying cost scenarios without manual tuning. An efficient solution combines mixed-integer programming with a scalable block coordinate descent algorithm. Evaluated on real-world workloads from BurstGPT, Azure, and ShareGPT, the approach reduces serving costs by up to 56% compared to fixed-quantile baselines while significantly improving P99 latency, effective throughput, and SLO compliance.

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.

Problem

• KV cache management
• output token length uncertainty
• preemption
• memory efficiency
• LLM serving

Innovation

• KV cache management
• distributionally robust optimization
• output token length uncertainty
• adaptive quantile reservation
• LLM serving