基于不确定性感知离线强化学习的LLM推理调度器URSA

《Scientific Reports》：Uncertainty-aware offline reinforcement learning for LLM inference serving with heterogeneous SLOs

　　

摘要

大语言模型（LLM）推理工作负载具有突发特征，到达时无法预知输出长度，暴露两个资源特征迥异的阶段（预填充和解码），并且为携带异构服务级别目标（SLO）的请求提供服务，这些 SLO 涉及首 token 时间（TTFT）和每输出 token 时间（TPOT）。现有调度器依赖于静态启发式策略，无法适应这种不确定性，而在线强化学习（RL）在生产环境中既不安全又缺乏样本效率。我们提出了 URSA，一种不确定性感知、风险敏感的调度器，完全离线训练，无需与实时服务系统交互，基于重放历史请求轨迹并在现有启发式调度器下生成的调度转移进行训练。URSA（i）预测每个请求输出长度的校准分布以及由此推导的关键值（KV）缓存占用量，而非仅给出点估计；（ii）将等待队列、异构 GPU 实例和 KV 占用率的联合状态编码为请求-实例异构图，由图神经网络处理；（iii）学习采用保守分布式离线强化学习的调度策略，其奖励函数同时优化 SLO 达成率、有效吞吐量（goodput）和 KV 溢出风险；（iv）按类依赖的高分位数预留预测占用量的内存，以限制内存溢出事件的发生。在 BurstGPT、Azure LLM 推理和阿里巴巴 GPU 集群轨迹上，URSA 将总体 SLO 达成率从 \(71.3\%\)（先到先服务）和 \(84.1\%\)（学习排序）提升至 \(90.4\%\)，将 P99 TTFT 降低至多 \(3.3\times \)，并将重计算/逐出事件降至 \(1\%\) 以下，接近于一个知晓真实长度的理想调度器。我们分析了其局限性、对预测误差的鲁棒性以及向其他调度领域的迁移能力。