Exact optimization improves cache use in language model serving
Exact Memory-Time Optimization for Prefix-Cached Language Model Serving
Machine Learning
Summary
Language models process text in parts called prefixes, and saving these intermediate results can speed up responses but uses more memory. The authors created a precise mathematical way to decide how long to keep each saved prefix to best balance memory use and computation time. They tested their method on real usage data and found it often matches or improves existing strategies. This approach helps manage memory and speed for language models without guessing or oversimplifying.
What this means in practice
- •For ml infrastructure engineers: Optimize memory and compute time tradeoffs when caching language model prefixes during serving.
- •For cloud service operators: Better manage storage and computation costs by precisely tuning cached prefix retention in language model APIs.
Authors
Shivam Gupta
Abstract
Retaining language-model prefix states trades recomputation against storage time. Optimizing each cached block independently can overcount savings: a resident block is usable only when the required preceding prefix is also available. We introduce Prefix-Certificate Retention (PCR), an exact finite-trace formulation for static, grouped, reset-on-access timeouts. Usable-prefix rewards become nodes whose prerequisites are timeout thresholds and preceding hit certificates. The resulting maximum-weight closure reduces to one minimum cut, with graph size linear in the number of block lookups and timeout choices. A breakpoint theorem extends the construction to all nonnegative timeouts without discretization error. We also derive a linear-time-in-grid-size dynamic program for ordered timeouts and bounds that certify the cost of this restriction. Exhaustive small-instance checks and chronological replay of 39,632 public Mooncake requests validate the formulation. On the fixed grid, ordered timeouts attain the unrestricted training optimum in 118 of 120 trace-grouping-price cases. Heterogeneous retention improves several held-out memory-time tradeoffs, but finer training optimization does not uniformly improve transfer. The contribution is a tractable optimization model and an auditable benchmark for retention policies; the experiments measure usable prefix blocks and storage time, not GPU latency.