Robust KV Cache Management for LLM Serving under Output Token Length Uncertainty
Networking and Internet Architecture
Summary
The authors address a problem in large language model (LLM) systems where memory reserved for processing requests must be guessed beforehand, but the actual required size isn’t known until the task finishes. Reserving too little memory causes delays and wasted effort, while reserving too much wastes resources. They propose a smart framework that adjusts memory allocation and request handling dynamically to balance these issues, using a mathematical optimization approach that adapts to changing workloads without needing constant manual tuning. Testing on real-world data shows their method uses memory more efficiently while keeping performance steady.
Authors
Jiaming Cheng, Duong The Do, Duong Tung Nguyen
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.