Robust KV Cache Management for LLM Serving under Output Token Length Uncertainty: 基于分布鲁棒优化的KV Cache自适应管理

一、论文概述

字段	内容
标题	Robust KV Cache Management for LLM Serving under Output Token Length Uncertainty
作者	Jingren Cheng (Amazon.com Services LLC), H. Vincent Poor (Princeton University)
机构	Amazon.com Services LLC, Princeton University
arXiv	arXiv:2607.11211
代码	N/A (paper mentions simulation, no code link)
日期	2025
License	N/A

二、核心思想

问题定义

大规模LLM推理服务系统中，GPU显存主要用于存储每个请求的KV cache，其大小正比于输入+输出token长度。然而输出token长度存在高度不确定性——不同任务类型（对话、代码生成、数据分析）的输出分布差异巨大。传统的KV cache预留方法（如固定分位数P90/P95/P99）面临一个根本性难题：任何单一的分位数只能在特定cost ratio下最优，当实际工作负载变化时，固定策略要么浪费GPU显存（over-reserve），要么导致频繁的preemption（under-reserve）。这个问题在实际部署中尤为突出，因为不同应用对preemption和waiting的相对成本不同，且常常是未知的或动态变化的。

具体而言，设第 $i$ 类请求的类型为 $(L^i_{\text{in}}, L^i_{\text{max}})$ ，实际输出长度为随机变量 $\xi_i$ 。系统为每类请求预留 $r_i$ 个token的KV cache空间，若 $r_i < L^i_{\text{in}} + \xi_i$ 则发生preemption（代价 $c_p$ ），若 $r_i > L^i_{\text{in}} + \xi_i$ 则产生wasted memory slot（代价 $c_w$ ）。目标是最小化总期望cost： $\min_r \sum_{i\in\mathcal{I}} p_i \mathbb{E}[Q_i(r_i, \xi_i)]$ ，其中 $Q_i$ 是包含preemption和waste的stochastic cost function。同时，系统还需联合决策GPU parallelism configuration $z\in\mathcal{Z}$ 、routing probabilities $\pi_{i,k}$ 、prefix caching fraction $\delta_{i,k}$ ，以及SLO-based request admission，在满足吞吐量和尾延迟约束的前提下最大化整体throughput。

解决方案概述

本文提出一个统一框架，核心包含三个层次：（1）分布鲁棒优化（DRO） — 使用Wasserstein ambiguity set保护预知之外的分布偏移，替代传统sample-average approximation (SAA)；（2）临界分位数结构 — 证明了最优预留量具有closed-form quantile structure， $q^* = \rho / (\rho + 1)$ ，自动根据cost ratio选择最佳分位数；（3）Block Coordinate Descent (BCD) — 将原始混合整数问题分解为GPU configuration、kv cache reservation、routing/prefix-caching等子问题，使production-scale实例可在数分钟内求解。

框架进一步结合排队论模型（M/G/1 Pollaczek-Khinchin公式）估计每组的queueing delay，引入rolling-horizon adaptation机制在deployment期间periodically re-optimize以适应workload演化。实验在三个真实trace（BurstGPT, Azure LLM 2024, ShareGPT）上验证了DRO在不同cost regime下的cross-regime稳定性。

三、技术架构

整体框架图

                    +------------------------------------------+
                    |       Rolling Horizon Adaptation         |
                    |   (Periodically re-optimize with        |
                    |    recent observations to track drift)   |
                    +------------------+-----------------------+
                                       |
                       +---------------+---------------+
                       |     BCD-DRO Optimizer         |
                       |  Block Coordinate Descent     |
                       +---+-----------+-----------+---+
                           |           |           |
          +----------------+  +--------+---+  +----+----------------+
          |                        |               |                 |
    +-----v------+         +------v-----+    +----v------+
    | GPU Config  |         | KV Cache   |    | Admission  |
    | z in Z      |         | Reservations|   | & Routing  |
    | (TP, PP)    |         | r_i (DRO)  |    | (pi, delta)|
    +------------+         +------------+    +-----------+
                                         |
                  +----------------------+----------------------+
                  |              Waiting Cost Model              |
                  |  M/G/1 Pollaczek-Khinchin (queueing delay)  |
                  +----------------------+----------------------+
                                         |
                  +----------------------+----------------------+
                  |         Wasserstein Ambiguity Set            |
                  |  P_hat_N (empirical) + epsilon ball          |
                  |  Protects against distribution shift         |
                  +----------------------------------------------+

核心公式

Stochastic Preemption/Waste Cost:

Q^{\text{stoch}}_{i}(r_{i},\xi)=c_{p}^{i}(\xi+r_{i}-L_{\text{in}}^{i})^{+}+c_{w}^{i}(r_{i}-L_{\text{in}}^{i}-\xi)^{+}\tag{1}

Total Cost Optimization (Primal):

\begin{aligned} \min_{(z,r,\pi,\delta)\in\bar{\mathcal{F}}}&\sum_{i\in\mathcal{I}}p_{i}\,\mathbb{E}\bigl[Q^{\text{stoch}}_{i}(r_{i},\xi_{i})\bigr]+\kappa\sum_{k\in\mathcal{K}}\phi_{k}(z_{k}) \\ &+\sum_{i\in\mathcal{I}}\sum_{k\in\mathcal{K}}\pi_{ik}\lambda_{i}\,c_{\text{slo}}T_{ik}(r,\delta,z)+c_{\text{rej}}\sum_{i\in\mathcal{I}}(1-a_{i})\lambda_{i} \end{aligned}\tag{2}

Optimal Reservation Quantile (Critical Fractile Structure):

q_{i}^{*}=\inf\left\{q:\frac{F_{i}^{(n)}(q)-F_{i}^{(n)}(0)}{1-F_{i}^{(n)}(0)}\geq\rho^{i}\right\},\quad F_{i}^{(n)}(x):=\frac{1}{N}\sum_{m=1}^{N}\mathbf{1}_{\{\hat{\xi}_{i}^{(m)}\leq x\}}\tag{3}

Critical Fractile Simplification (zero-pretrain case):

q_{i}^{*}=\rho^{i}/(\rho^{i}+1),\quad\rho^{i}=c_{p}^{i}/c_{w}^{i}\tag{4}

Wasserstein Ambiguity Set:

\mathcal{P}_{\varepsilon}=\left\{P:\int_{0}^{L_{\text{max}}^{i}}P((-\infty,x])\,dx\leq\varepsilon,\ \forall i\in\mathcal{I}\right\}\tag{5}

Tractable DRO Reformulation (LP):

\begin{aligned} \min\quad &\gamma_{i}\varepsilon+\frac{1}{N}\sum_{n=1}^{N}\theta_{i,n}\\ \text{s.t.}\quad &\theta_{i,n}\geq c_{p}^{i}(\hat{\xi}_{i}^{(n)}+L_{\text{in}}^{i}-r_{i}),\ \forall n\\ &\theta_{i,n}\geq c_{w}^{i}(r_{i}-L_{\text{in}}^{i}-\hat{\xi}_{i}^{(n)}),\ \forall n\\ &\theta_{i,n}\geq c_{p}^{i}(L_{\text{max}}^{i}+L_{\text{in}}^{i}-r_{i})-\gamma_{i}(L_{\text{max}}^{i}-\hat{\xi}_{i}^{(n)}),\ \forall n\\ &\theta_{i,n}\geq c_{w}^{i}(r_{i}-L_{\text{in}}^{i})-\gamma_{i}\hat{\xi}_{i}^{(n)},\ \forall n\\ &\gamma_{i}\geq 0,\ \theta_{i,n}\geq 0,\ \forall n \end{aligned}\tag{6}

Multi-quantile DRO Reformulation (Proposition III.2):

\begin{aligned} \min\quad &\gamma\varepsilon+\sum_{n=1}^{N}\theta_{n}\\ \text{s.t.}\quad &\theta_{n}\geq c_{p}\sum_{i\in\mathcal{I}}q_{i}\pi_{i,k(n)}(\hat{\xi}_{i}^{(n)}+L_{\text{in}}^{i}-r_{i}),\ \forall n\\ &\theta_{n}\geq c_{w}\sum_{i\in\mathcal{I}}q_{i}\pi_{i,k(n)}(r_{i}-L_{\text{in}}^{i}-\hat{\xi}_{i}^{(n)}),\ \forall n\\ &\theta_{n}\geq c_{p}\sum_{i\in\mathcal{I}}q_{i}\pi_{i,k(n)}(L_{\text{max}}^{i}+L_{\text{in}}^{i}-r_{i})-\gamma(L_{\text{max}}^{i}-\hat{\xi}_{i}^{(n)}),\ \forall n\\ &\theta_{n}\geq c_{w}\sum_{i\in\mathcal{I}}q_{i}\pi_{i,k(n)}(r_{i}-L_{\text{in}}^{i})-\gamma\hat{\xi}_{i}^{(n)},\ \forall n\\ &\sum_{k\in\mathcal{K}}\pi_{i,k}=1,\ \pi_{i,k}\geq 0,\ \forall i,k\\ &r_{i}=q_{i}L_{\text{max}}^{i},\ q_{i}\in[0,1],\ \forall i\\ &\gamma\geq 0,\ \theta_{n}\geq 0,\ \forall n \end{aligned}\tag{7}

Cost Ratio from SLO Parameters:

\rho^{i}\approx e^{\mu_{i}/(\sigma_{i}\sqrt{2\log\lambda_{i}})},\quad i\in\mathcal{I} \tag{8}

Queueing Delay (M/G/1 Pollaczek-Khinchin):

T_{ik}(r,\delta,z)=\frac{\lambda_{k}\bigl(\mathrm{Var}(r_{i}-\delta_{i,k}L_{\text{pre}}^{i})+(\bar{r}_{ik}-\bar{\delta}_{ik}L_{\text{pre}}^{i})^{2}\bigr)}{2\bigl(m_{k}C_{k}(z)-\sum_{i\in\mathcal{I}}\lambda_{i}\pi_{ik}(r_{i}-\delta_{i,k}L_{\text{pre}}^{i})\bigr)}

模型组件

组件	符号	作用	决策变量
GPU Parallelism	$z_k = (\text{TP}_k, \text{PP}_k)$	配置每组的张量/流水线并行度	Discrete config in $\mathcal{Z}$
KV Cache Reservation	$r_i$	为类型 $i$ 请求预留的output token数	Continuous $[0, L_{\text{max}}^i]$
Request Routing	$\pi_{ik}$	类型 $i$ 请求路由到组 $k$ 的概率	Probability simplex
Prefix Caching	$\delta_{i,k}$	类型 $i$ 在组 $k$ 的前缀缓存比例	Continuous $[0,1]$
Admission Control	$a_i$	是否允许类型 $i$ 请求接入	Binary
DRO Dual Variable	$\gamma_i$	Wasserstein ambiguity set的对偶变量	Continuous $\geq 0$
Epi Variables	$\theta_{i,n}$	$\phi_i$ 的上界epigraph变量	Continuous $\geq 0$
Quantile Parameter	$q_i$	$r_i = q_i L_{\text{max}}^i$	Continuous $[0,1]$

训练流程

数据收集: 从生产trace（BurstGPT/Azure LLM/ShareGPT）收集历史output length样本 $\{\hat{\xi}_i^{(n)}\}_{n=1}^N$
Empirical Distribution: 构建经验分布 $\hat{P}_N = \frac{1}{N}\sum_{n=1}^N \delta_{\hat{\xi}_i^{(n)}}$
构建Ambiguity Set: 设置Wasserstein半径 $\varepsilon$ （推荐 $\varepsilon = 0.15 \cdot \mathbb{E}[\xi_i]$ ）
DRO求解: 求解LP reformulation (式6/7)，获得最优 $\gamma_i^*, q_i^*$
BCD分解迭代:
- Step 1: 固定 $z$ ，求解inner LP得到 $q^*(z), r^*(z), \pi^*(z), \delta^*(z)$
- Step 2: 固定 $q, r, \pi, \delta$ ，枚举 $z\in\mathcal{Z}$ 找最优config
Rolling Horizon: 部署后periodically用recent data re-optimize

四、核心创新

#	创新点	说明	与传统方法的差距
1	临界分位数结构（Critical Fractile Structure）	证明了最优KV cache预留具有closed-form quantile形式 $q^* = \rho / (\rho + 1)$ ，其中 $\rho = c_p/c_w$ 。消除了manual heuristic tuning的需要	固定分位数（P90/P95/P99）只在特定时最优；此公式自动生成最优分位数
2	Wasserstein DRO鲁棒性	使用Wasserstein ambiguity set保护distribution shift，而非传统SAA仅拟合训练数据。提供了可证明的out-of-sample保证	SAA在分布偏移时cost显著恶化（shift 3x时P99从9s变为18s vs DRO的9s）
3	Tractable LP Reformulation	通过Kantorovich duality将无限维鲁棒优化转化为有限LP。对于固定 $r_i$ 的问题复杂度仅为 $O(N)$	直接MIBLP求解在 $
4	BCD分解算法	将混合整数问题分解为GPU configuration和外层LP交替优化。生产规模（15类请求, 12种config, 2000样本）257秒求解	与monolithic solver相比，15倍加速且能扩展到production scale
5	Unified Joint Optimization	首次将GPU config、KV reserve、routing、prefix caching、admission unified在一个DRO框架内	现有系统通常单独优化其中一个方面
6	SLO-to-Cost-Ratio映射	从LLM serving常见的SLO参数推导出cost ratio： $\rho^i \approx \exp(\mu_i/(\sigma_i\sqrt{2\log\lambda_i}))$	传统方法没有将业务SLO参数与优化参数关联的方法

五、代码实现分析

当前论文未公开code repository。文中描述的实现基于:

Solver: Gurobi (用于求解MIBLP baseline和BCD子问题LP)
Language: Python (simulation framework)
Key Dependencies: Gurobi optimization, queueing theory models

参考实现的推测文件结构：

kv-cache-robust/
├── dro_optimizer.py      # Wasserstein DRO LP formulation
├── bcd_solver.py          # Block coordinate descent
├── gpu_config.py          # GPU parallelism configuration space
├── queueing_model.py      # M/G/1 Pollaczek-Khinchin delay model
├── simulation.py          # Request-level simulator
├── trace_datasets.py      # BurstGPT, Azure, ShareGPT loaders
├── baselines.py           # P90/P95/P99/Mean/Max/SAA
└── ablation.py            # Component ablation studies

六、实验结果

基准测试

硬件配置: 8xA100-80GB GPUs serving a 70B model

$k_1$ : TP=2, PP=1
$k_2$ : TP=4, PP=1
$k_3$ : TP=2, PP=2
$k_4$ : TP=4, PP=2

Cost参数: $c_p/c_w=10$ , $c_{\text{slo}}/c_w=5$ , $c_{\text{rej}}/c_w=5000$ , $\kappa=0.2$

Table II: DRO vs Fixed-Quantile Baselines across Scenarios

Scenario	$\rho$	DRO Cost	P90	P95	DRO Gain
Batch	2	184	255	365	28-50%
Async API	5	360	380	459	5-22%
User-facing	10	587	588	617	0-5%
Latency-sensitive	20	931	1003	931	0-7%
Critical	50	1575	2250	1875	16-30%
Real-time	100	1919	4328	3448	44-56%

关键发现：

DRO在所有cost ratio下表现最优
在real-time场景( $\rho=100$ )下，DRO相比P90节省 56% cost（1919 vs 4328）
在critical场景( $\rho=50$ )下，DRO相比P90节省 30% cost

消融实验

Table V: Ablation on Heterogeneous Instance ( $\rho=20$ , load $\alpha_\lambda=0.6$ , 2.5x shifted)

Variant	Cost	P99 (s)	Goodput	SLO viol. (%)
BCD-DRO (ours)	13,476 ± 208	8.12 ± 0.65	5.01 ± 0.25	26.4 ± 3.7
- Routing (uniform $\pi$ )	93,878 ± 5,152	177.62 ± 11.90	0.63 ± 0.06	93.8 ± 0.6
- DRO (SAA, $\varepsilon=0$ )	24,864 ± 1,534	11.02 ± 1.55	3.68 ± 0.35	45.8 ± 5.2
- Prefix caching ( $\delta=0$ )	13,476 ± 208	9.95 ± 1.29	4.39 ± 0.37	35.4 ± 5.4

关键发现：

移除Routing是最致命的：cost暴增7x（13,476 -> 93,878），SLO violation达93.8%
移除DRO鲁棒性（ $\varepsilon=0$ ）：cost增加85%，SLO violation提高73%
移除Prefix Caching：不改变preemption/waste cost，但P99 latency增加22.6%（8.12 -> 9.95s）

与现有方法对比

Table IV: Cost across Three Cost-Ratio Regimes and Latency at $\rho=10$

Method	$\rho=2$ Cost	$\rho=10$ Cost	$\rho=100$ Cost	P99 (s) @ $\rho=10$	SLO viol. @ $\rho=10$
DRO (ours)	1,270	3,083	4,807	11.47	15.6%
P90	1,852	3,095	17,080	11.39	14.8%
P95	2,890	3,320	8,153	12.50	23.2%
P99	4,308	4,349	4,811	17.12	33.7%
Mean	1,275	4,666	42,811	12.03	21.7%
Max	4,752	4,756	4,807	19.04	36.7%
LP+1σ	1,729	3,143	19,055	11.30	13.2%
LP+2σ	2,615	3,207	9,865	11.80	18.0%

关键发现：

DRO是唯一在所有三个 $\rho$ regime下都bounded的方法
P90在 $\rho=100$ 时cost暴涨5.5x（3,095 -> 17,080）
Mean在 $\rho=100$ 时cost暴涨13.8x（4,666 -> 42,811）
DRO的cost跨regime范围仅从1,270到4,807（3.8x），远优于所有baseline
LP+1σ在 $\rho=10$ 时latency略优（P99 11.30s vs 11.47s），但在 $\rho=100$ 时cost差4x

Algorithm Scalability (Table III):

Problem Size	MIBLP Time(s)	MIBLP Status	BCD-DRO Time(s)	Gap
3×4×100	0.07	Optimal	0.17	0.00%
4×4×150	1.2	Optimal	0.8	0.00%
5×5×200	33.7	Optimal	4.3	0.00%
6×6×300	>300	Timeout	10.5	—
8×6×500	>300	Timeout	22.4	—
10×8×1000	—	—	81.2	—
15×12×2000	—	—	257.7	—

BCD-DRO在production规模（15×12×2000）仅需257.7秒，而MIBLP在6×6×300就timeout。

七、相关工作

类别	论文	主要方法	与本文差异
KV Cache Management	PagedAttention [1]	virtualized memory pages	底层内存管理，不涉及optimization
KV Cache Management	Orca [2]	disaggregated prefill/decode	system architecture，非adaptive
KV Cache Management	Sarathi-Serve [3]	speculative batching for long outputs	batching only
KV Cache Eviction	H2O [4]	heavy-hitter oracle	eviction only
KV Cache Compression	SnapKV [5]	token-level importance scoring	compression
Prefill Decoding	S³ [6]	increasing GPU utilization	throughput-focused
Length Prediction	Response Length Perception [7]	LLM-empowered length prediction	仅length perception
System Design	DistServe [24]	disaggregated prefill/decode	system level
本文	DRO-based KV Cache	joint optimization + robustness	unified framework with theoretical guarantees

八、总结

核心贡献

首个将Distributionally Robust Optimization应用于LLM KV Cache管理的统一框架，同时考虑GPU parallelism、reservation、routing和prefix caching
临界分位数理论（Critical Fractile Structure）：证明了最优预留量具有closed-form quantile结构，自动适配任意cost ratio
Wasserstein DRO的可解重构：通过Kantorovich duality将无限维鲁棒优化转化为线性规划
BCD分解算法：实现了production-scale问题的有效求解（15×12×2000在258秒内收敛）
全面实证评估：在三个真实生产trace上验证了跨regime的鲁棒性和优越性

技术影响

为LLM serving系统提供了一种无需手动调参的自适应KV cache管理方案
理论上的critical fractile公式可直接嵌入现有serving system（如vLLM、TGI）作为dynamic reservation policy
Wasserstein DRO framework可扩展到其他LLM资源管理问题（attention cache、batch size scheduling）

局限性

参数敏感性：最优性能依赖于准确的cost ratio $\rho$ 估计，虽然框架自动选择optimal quantile，但 $\rho$ 本身的设定仍需domain knowledge
Wasserstein半径选择： $\varepsilon$ 影响robustness-optimality trade-off，建议值 $\varepsilon \in [0.1, 0.2]$ 缺乏systematic选择方法
计算开销：在生产规模（15×12×2000）下需258秒重新优化，虽在re-optimization budget内但仍有改进空间
SLO模型简化：使用M/G/1队列近似，实际系统可能有更复杂的调度协议和preemption cost

九、参考资源

Paper: https://arxiv.org/abs/2607.11211
LaTeXML HTML: https://ar5iv.labs.arxiv.org/html/2607.11211
参考文献:
- [1] PagedAttention (SOSP 2023) — vLLM的基础
- [2] Orca (OSDI 2022) — disaggregated prefill/decode
- [3] Sarathi-Serve (OSDI 2024) — speculative batching
- [4] H2O (NeurIPS 2023) — heavy-hitter KV cache eviction
- [5] SnapKV (arXiv 2024) — token-level importance
- [6] S³ (NeurIPS 2023) — GPU utilization for generative inference
- [7] Response Length Perception (NeurIPS 2024) — LLM-empowered scheduling
- [11] MT-Bench / Chatbot Arena (NeurIPS 2023) — LLM evaluation
- [12] Scalable Joint Resource Allocation (arXiv 2025) — similar topic
- [13] Green-LLM (arXiv 2025) — energy-aware distributed inference
- [14] Wasserstein DRO (Math. Program. 2018) — foundational DRO paper
- [15] Edgeworth (1888) — early critical fractile
- [16] Arrow et al. (1951) — classical newsvendor problem
- [17] Data-driven DR Newsvendor (JORS 2021) — Wasserstein newsvendor
- [18] Kleinrock (1975) — Queueing Theory
- [19] Queueing-theoretic Low-Latency LLM (WiOpt 2024) — related queueing approach
- [20] Blockwise CD for Integer Programs (MMOR 2020) — BCD algorithm
- [21] BurstGPT (KDD 2025) — real-world LLM workload dataset
- [22] DynamoLLM (HPCA 2025) — LLM inference cluster design
- [23] ShareGPT — crowdsourced ChatGPT conversations
- [24] DistServe (OSDI 2024) — disaggregated goodput-optimized serving