Intelligent Router for LLM Workloads: Improving Performance Through Workload-Aware Load Balancing
作者: Kunal Jain, Anjaly Parayil, Ankur Mallick, Esha Choukse, Xiaoting Qin, Jue Zhang, Íñigo Goiri, Rujia Wang, Chetan Bansal, Victor Rühle, Anoop Kulkarni, Steve Kofsky, Saravan Rajmohan
分类: cs.DC, eess.SY
发布日期: 2024-08-24 (更新: 2025-01-07)
备注: 16 pages, 10 figures
💡 一句话要点
提出基于强化学习的智能路由，优化LLM推理负载均衡，降低端到端延迟。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 LLM推理 负载均衡 强化学习 智能路由
📋 核心要点
- 现有LLM调度方法将推理视为整体，忽略预填充和解码阶段的差异，导致负载不均衡和延迟增加。
- 提出一种基于强化学习的智能路由器，通过预测响应长度和估计工作负载混合影响，实现数据驱动的调度。
- 实验表明，该方法在公共数据集和真实工作负载上分别降低了11%和7.8%的端到端延迟。
📝 摘要（中文）
大型语言模型（LLM）工作负载具有不同的预填充（prefill）和解码（decode）阶段，它们对计算和内存的需求各不相同。理想情况下，在集群中跨不同的LLM实例调度输入查询时，应该考虑到这些差异。然而，现有的调度算法将LLM工作负载视为单体作业，没有考虑每个工作负载中两个阶段的不同特性，导致次优的调度和增加的响应延迟。本文首先分析了影响LLM推理服务期间响应延迟的因素。研究表明，与仅仅关注优化实例级别的调度器相比，更好地平衡可用LLM实例之间的推理请求负载可以更大程度地改善端到端延迟。受此发现的启发，我们提出了一种基于启发式引导的强化学习智能路由器，用于数据驱动和工作负载感知的调度。我们的路由器通过利用可训练的响应长度预测器和一种用于估计混合不同工作负载影响的新颖公式来跨LLM实例调度查询，与现有方法相比，在公共数据集的混合上实现了超过11%的端到端延迟降低，在来自Cloud Provider X的具有多样化输入和输出趋势的真实工作负载数据上实现了7.8%的端到端延迟降低。此外，所提出的框架还可以作为不同LLM推理调度器的基准标准，因为它为给定的模型、硬件和实例级别调度器组合提供了最佳延迟。
🔬 方法详解
问题定义：现有LLM推理服务调度器通常将整个推理过程视为一个整体，忽略了预填充（prefill）和解码（decode）阶段在计算和内存需求上的差异。这种简化导致了负载不均衡，某些LLM实例可能过载，而其他实例则处于空闲状态，最终增加了端到端延迟。现有方法缺乏对工作负载特性的感知，无法根据不同查询的需求进行智能调度。
核心思路：本文的核心思路是通过一个智能路由器，根据对LLM工作负载特性的理解，动态地将查询分配给不同的LLM实例。该路由器利用强化学习来学习最优的调度策略，并结合启发式方法来加速学习过程。通过预测响应长度和估计不同工作负载混合的影响，路由器能够更好地平衡各个LLM实例上的负载，从而降低整体延迟。
技术框架：该智能路由器的整体框架包括以下几个主要模块：1) 响应长度预测器：用于预测给定查询的响应长度，以便估计其计算需求。2) 工作负载混合影响估计器：用于估计将不同类型的查询混合在同一个LLM实例上运行所产生的影响。3) 强化学习代理：基于预测的响应长度和工作负载混合影响，学习最优的调度策略。4) 调度器：根据强化学习代理的决策，将查询分配给不同的LLM实例。
关键创新：该论文的关键创新在于：1) 提出了一种基于强化学习的智能路由器，能够根据LLM工作负载的特性进行动态调度。2) 引入了响应长度预测器和工作负载混合影响估计器，用于更准确地评估查询的计算需求和潜在的干扰。3) 将强化学习与启发式方法相结合，加速了学习过程并提高了调度性能。
关键设计：响应长度预测器可能采用一个轻量级的神经网络，根据输入查询的文本特征来预测响应长度。工作负载混合影响估计器可能基于历史数据或模拟实验来建立一个查询类型与性能影响之间的映射关系。强化学习代理可能使用深度Q网络（DQN）或策略梯度方法来学习最优的调度策略，奖励函数的设计需要考虑延迟、负载均衡等因素。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该智能路由器在公共数据集的混合上实现了超过11%的端到端延迟降低，在来自Cloud Provider X的真实工作负载数据上实现了7.8%的端到端延迟降低。这些结果证明了该方法在实际应用中的有效性，并表明其优于现有的调度方法。
🎯 应用场景
该研究成果可应用于云服务提供商、AI推理平台等场景，用于优化LLM推理服务的性能和资源利用率。通过智能路由，可以显著降低用户请求的响应延迟，提升用户体验。此外，该框架还可以作为评估不同LLM推理调度器的标准，推动相关技术的发展。
📄 摘要（原文）
Large Language Model (LLM) workloads have distinct prefill and decode phases with different compute and memory requirements which should ideally be accounted for when scheduling input queries across different LLM instances in a cluster. However existing scheduling algorithms treat LLM workloads as monolithic jobs without considering the distinct characteristics of the two phases in each workload. This leads to sub-optimal scheduling and increased response latency. In this work, we start by characterizing factors affecting the response latency during LLM inference serving. We establish that better load balancing of inference requests across the available LLM instances can improve the end-to-end latency to a larger extent than merely focusing on optimizing the instance-level scheduler. Motivated by our findings, we propose a heuristic-guided reinforcement learning-based intelligent router for data-driven and workload-aware scheduling. Our router schedules queries across LLM instances by leveraging a trainable response-length predictor, and a novel formulation for estimating the impact of mixing different workloads and achieves over 11% lower end-to-end latency than existing approaches on a mix of public datasets and 7.8% lower end-to-end latency on real workload data with diverse input and output trends from Cloud Provider X. Additionally, the proposed framework can also serve as a standard for benchmarking different LLM inference schedulers since it provides the best latency for a given model, hardware, and instance-level scheduler combination.