SGLang调度器源码解析：从LPM到DFS_WEIGHT的缓存优化策略详解

SGLang调度器源码解析从LPM到DFS_WEIGHT的缓存优化策略详解最近在优化一个推理服务的吞吐量时我反复琢磨调度策略对缓存命中率的影响。很多团队在部署大语言模型服务时往往只关注硬件规格和模型本身却忽略了调度器这个“交通指挥官”的关键作用。SGLang作为一个新兴的高性能推理运行时其调度器设计尤其精妙特别是它对前缀缓存的感知能力直接决定了在高并发场景下那些昂贵的GPU计算资源是被高效复用还是在重复计算中空转。今天我们就抛开表面的API调用直接潜入schedule_policy.py的源码深处看看LPM和DFS_WEIGHT这两种缓存感知策略是如何工作的以及它们如何像一位经验丰富的仓库管理员通过巧妙的排序让后续的“取货”请求能更快地命中已准备好的“货架”。理解这些策略对于需要自研或深度定制推理服务、追求极致性价比的技术团队至关重要。它不仅仅是选择哪个参数的问题更是理解系统在负载下如何做出微观决策从而为架构设计和参数调优提供坚实依据。我们将从基本概念出发逐步拆解源码中的关键数据结构、算法逻辑并探讨其在实际负载下的行为表现。1. 调度策略的基石理解SGLang的调度上下文在深入两种具体的缓存感知策略之前我们必须先搭建起对SGLang调度器工作环境的整体认知。调度器并非在真空中运行它处理的对象、依赖的缓存结构以及面临的约束共同构成了策略生效的舞台。核心调度单元请求Req对象调度器操作的基本单位是Req即用户的一个生成请求。每个Req对象携带了决定其调度优先级的关键元数据。在源码中我们尤其需要关注以下几个属性rid: 请求的唯一标识符。prefix_indices: 该请求的输入序列prompt在树缓存Tree Cache中匹配到的前缀节点索引列表。这是LPM策略排序的核心依据。last_node: 该请求匹配到的前缀在树缓存中对应的最后一个TreeNode。这是DFS_WEIGHT策略进行权重计算和分组的关键。sampling_params.max_new_tokens: 请求需要生成的最大token数量被LOF最长输出优先策略所使用。两级缓存架构树缓存与等待队列前缀树SGLang的缓存感知策略依赖于一个层次化的缓存系统这是其高效调度的秘密武器。全局树缓存self.tree_cache: BasePrefixCache 这是一个所有请求共享的、用于存储已计算过的注意力键值对KV Cache的前缀树。其节点TreeNode代表一个token序列前缀。当新请求的prompt与此树中的某个路径匹配时就可以直接复用该路径终点节点对应的KV Cache跳过重复计算。tree_cache.match_prefix()方法就是用来为请求寻找最长匹配前缀的。等待队列前缀树self.waiting_queue_radix_tree 这是一个动态的、仅针对当前等待队列中请求构建的Radix树。它的作用非常巧妙用于发现批次内in-batch的重复前缀。例如短时间内涌入大量相似的用户提问虽然它们的prompt可能尚未被存入全局树缓存但彼此之间高度相似。通过这棵临时树调度器能识别出这些“内部重复”并据此调整调度顺序让具有共同前缀的请求尽可能被一起处理从而在单个批次内实现KV Cache的共享提升计算效率。策略的决策流程SchedulePolicy.calc_priority()方法是调度策略的发动机。其逻辑可以概括为以下步骤def calc_priority(self, waiting_queue): # 1. 决定当前生效的策略可能因队列过长而降级 policy self._determine_active_policy(waiting_queue) # 2. 如果是缓存感知策略先计算前缀匹配信息 if policy 是 CacheAwarePolicy: 计算所有请求的前缀匹配情况更新 prefix_indices, last_node 识别批次内重复度过高的请求临时降级集合 # 3. 根据具体策略排序 if policy LPM: 按最长匹配前缀长度排序 elif policy DFS_WEIGHT: 按DFS权重排序 # 4. 如果是缓存无关策略直接按规则排序 else: if policy FCFS: 不排序 elif policy LOF: 按max_new_tokens降序排 elif policy RANDOM: 随机打乱这个流程清晰地展示了策略执行的两个阶段信息准备计算匹配和队列重排排序。2. LPM策略最长前缀匹配的直观优化LPM策略的思想非常直观谁与缓存中已有的内容匹配度最高谁就优先被调度。这类似于操作系统的缓存置换算法中的“最优”思想目标是最大化每一次调度决策对缓存命中率的即时收益。源码实现解析策略的核心实现在静态方法_sort_by_longest_prefix中。虽然代码简短但蕴含了重要的工程考量。staticmethod def _sort_by_longest_prefix(waiting_queue, temporary_deprioritized): waiting_queue.sort( keylambda r: ( -len(r.prefix_indices) if r.rid not in temporary_deprioritized else float(inf) ) )我们来拆解这个排序键key的构造-len(r.prefix_indices)对于大多数请求按其匹配到的前缀长度降序排列。len(r.prefix_indices)越大说明该请求的prompt与全局缓存重合部分越多优先执行它可以最大化地复用已有的KV Cache减少计算量。r.rid not in temporary_deprioritized这是一个安全阀。temporary_deprioritized集合来源于_compute_prefix_matches方法它标识了那些在等待队列内部重复度极高的请求通过waiting_queue_radix_tree检测。为什么要把它们暂时降级因为如果多个高度相似的请求都在等待优先处理其中一个确实能为后续其他相似请求带来缓存收益。但是如果这种相似请求过多一直优先处理它们可能会导致队列中其他截然不同的请求被“饿死”影响系统处理的公平性和整体延迟的多样性。因此将这些“内部重复大户”的优先级设为float(inf)实际上是将它们排到了队列末尾。一个具体的例子假设全局树缓存中已有路径[“请介绍” “Python” “的”]。 当前等待队列有四个请求Req A:“请介绍Python的优缺点”- 匹配前缀长度 3Req B:“请解释一下机器学习”- 匹配前缀长度 1 (“请”)Req C:“Python是什么”- 匹配前缀长度 0 (全新)Req D:“请介绍Python的安装方法”- 匹配前缀长度 3且与Req A在队列内高度相似。经过_compute_prefix_matches计算Req D可能因与Req A在批次内重复而被加入temporary_deprioritized集合。那么LPM排序后的队列将是Req A (前缀长度3 非降级)Req B (前缀长度1)Req C (前缀长度0)Req D (前缀长度3 但因降级被排到最后)这样先执行Req A其生成的“优缺点”对应的KV Cache可以加入树缓存。虽然Req D被延后但当它被调度时它可能不仅能复用原来的3个token前缀还能复用Req A刚生成的“优缺点”的部分上下文如果模型架构和缓存策略允许从而获得更大的收益。注意LPM策略在队列过长128时会自动降级为FCFS。这是因为计算所有请求的前缀匹配涉及树遍历和比较本身有开销。当队列膨胀时这种计算开销可能抵消掉排序带来的缓存收益甚至成为瓶颈。此时退化为简单的FCFS反而能保证系统的稳定性和吞吐量。3. DFS_WEIGHT策略基于树结构的全局视角如果说LPM是一种“贪婪”的局部最优策略那么DFS_WEIGHT则尝试了一种更具“全局观”的调度思路。它不再孤立地看待每个请求的匹配长度而是将请求放置在整个前缀树的结构中考虑子树的大小和请求的分布其核心目标是优先调度那些位于“枝叶”较少分支上的请求。算法原理与动机为什么优先调度“枝叶”少的想象一棵前缀树根节点代表空序列每个子节点代表一个token。如果一个节点下有大量请求即它是一个很多请求的共同前缀那么先处理这个节点上的请求所能释放的缓存收益新生成的KV Cache可以被其后所有共享该前缀的请求复用。反之如果一个节点是叶子节点或只有一个请求先处理它带来的收益影响范围很小。DFS_WEIGHT通过赋予子节点多的父节点更高权重并依此进行深度优先遍历来系统性实现这一目标。源码分步拆解_sort_by_dfs_weight方法的实现稍复杂我们结合流程图来理解请求分组首先根据每个请求的last_node在全局树缓存中匹配到的最后一个节点将请求分组。last_node_to_reqs字典的键是TreeNode值是属于该节点的请求列表。last_node_to_reqs defaultdict(list) for req in waiting_queue: last_node_to_reqs[req.last_node].append(req)计算节点权重初始化node_to_weight其中每个节点的初始权重就是挂在该节点上的请求数量。然后通过后序遍历递归地计算每个节点的累积权重——一个节点的最终权重等于其自身初始权重加上其所有子节点的权重。# 递归计算权重 staticmethod def _calc_weight(cur_node, node_to_weight): for child in cur_node.children.values(): SchedulePolicy._calc_weight(child, node_to_weight) node_to_weight[cur_node] node_to_weight[child] # 累加子节点权重这个过程结束后node_to_weight[node]就代表了以该节点为根的子树所关联的总请求数包括所有后代节点上的请求。按权重进行DFS排序从根节点开始进行深度优先搜索。但在选择访问哪个子节点时不再是任意的而是按照子节点的权重降序排列。这意味着系统会优先进入那些关联请求总数最多的子树分支。childs [child for child in cur_node.children.values()] childs.sort(keylambda x: -node_to_priority[x]) # 按权重降序排子节点 for child in childs: SchedulePolicy._get_dfs_priority(child, node_to_priority, last_node_to_reqs, q) q.extend(last_node_to_reqs[cur_node]) # 访问完所有子节点后加入当前节点的请求最终队列q即被重排的waiting_queue中的请求顺序就是这棵权重优先的DFS遍历顺序。策略效果对比为了更直观地对比LPM和DFS_WEIGHT我们考虑一个简单的树状结构和请求分布节点路径节点说明挂载的请求子节点数计算后的节点权重根 (/)空无25 (Req1Req2子树权重)- A前缀“苹果”Req123- A - A1前缀“苹果 手机”Req201- A - A2前缀“苹果 电脑”Req301- B前缀“香蕉”Req412- B - B1前缀“香蕉 牛奶”Req501假设所有请求的匹配前缀都刚好到其挂载的节点。LPM策略只看每个请求自身匹配长度。假设路径深度相同则顺序可能接近原始顺序或随机。DFS_WEIGHT策略计算权重根节点权重5节点A权重3节点B权重2叶子节点A1、A2、B1权重1。从根开始子节点A(权重3)和B(权重2)先访问A。在A下先访问权重更高的子节点这里A1和A2权重都是1按源码实现sort在权重相同时会保持原有顺序可能是插入顺序。假设先A1后A2。回溯到A加入A的请求Req1。回溯到根访问子节点B接着访问B1加入B1的请求Req5回溯到B加入B的请求Req4。最终调度顺序可能是[Req2, Req3, Req1, Req5, Req4]。可以看到DFS_WEIGHT将共享同一子树“苹果”分支的请求Req2、Req3、Req1集中在了队列前部。这样先处理Req2“苹果手机”其生成的KV Cache可以被后续的Req3“苹果电脑”和Req1“苹果”部分复用即使Req1的前缀匹配长度可能比Req2短。它追求的是批次处理在树结构上的“聚集效应”从而在整体上提高缓存利用率。4. 缓存无关策略当简单即是美并非所有场景都适合或需要复杂的缓存感知调度。SGLang也提供了几种简单直接的缓存无关策略Cache-Agnostic Policy。它们不依赖树缓存的状态仅根据请求自身的属性或简单规则进行调度在某些特定场景下反而是更优选择。FCFS先来先服务这是最朴素、最公平的策略。calc_priority方法中对应pass即不对队列做任何重排。它的优势是开销极低保证了请求的等待时间与其到达顺序严格一致避免了“饥饿”现象。在以下情况适用请求的prompt模式高度随机几乎没有重复前缀缓存感知策略收益甚微。系统负载极高等待队列持续很长如前所述LPM策略会自动退化为FCFS以降低计算开销。对延迟的公平性有严格要求不希望某些请求因内容特性而被过度延迟。LOF最长输出优先该策略按请求的max_new_tokens降序排列。其背后的思想是最大化吞吐量。在解码阶段生成一个长序列所需的计算时间远大于短序列。如果让长序列请求等待它会阻塞其所在的计算槽位很长时间。优先处理长序列请求可以尽早释放其占用的资源如GPU内存中的KV Cache让计算资源更快地周转起来从而提高系统的整体吞吐量。staticmethod def _sort_by_longest_output(waiting_queue): waiting_queue.sort(keylambda x: -x.sampling_params.max_new_tokens)提示LOF策略与缓存感知策略的目标存在内在冲突。缓存感知追求计算复用可能让短请求插队LOF追求资源释放速度。选择哪种策略取决于你的首要优化目标是降低平均延迟缓存感知还是提高单位时间内的请求处理数LOF。RANDOM随机调度随机打乱等待队列。这听起来像是一个没什么用的策略但在某些特定场景下有奇效负载均衡测试在评估后端多个推理实例的负载均衡效果时随机调度可以消除请求顺序带来的偏差。打破局部最优在某些极端请求模式下固定的调度策略如LPM可能导致某些类型的请求永远被延后。引入一定的随机性可以作为避免“饥饿”的一种简单兜底机制。探索性调优在系统调优初期可以用作与其他策略进行A/B测试的基线。5. 策略选择与实践指南了解了各种策略的原理最关键的问题来了在实际项目中我们该如何选择这没有银弹需要根据你的工作负载特征和性能目标来决定。下面这个决策矩阵或许能提供一些参考策略类型策略名称核心优化目标适用工作负载特征潜在缺点缓存感知LPM单次调度决策的缓存命中率请求间有显著、可复用的共同前缀如聊天机器人固定开场白、批量处理相似问题队列过长时计算开销大可能造成非相似请求饥饿缓存感知DFS_WEIGHT子树级别的全局缓存利用率请求呈明显的树状或主题聚类分布如多轮对话、主题相关的问答集算法复杂度稍高对请求分布模式敏感缓存无关FCFS公平性与调度开销请求高度多样化、无规律或系统处于极高负载状态无法利用缓存提升效率缓存无关LOF系统整体吞吐量请求的生成长度差异巨大且长请求占比高严重不利于短请求的延迟缓存无关RANDOM公平性与打破模式用于测试、或作为避免饥饿的混合策略一部分性能不可预测通常不作为生产环境主策略混合策略与动态切换在实际生产环境中单一策略可能无法应对所有情况。SGLang的源码已经暗示了一种动态切换的思路在_determine_active_policy方法中当队列长度超过128且策略为LPM时会自动降级为FCFS。我们可以借鉴这种思想设计更灵活的调度器监控驱动切换持续监控请求的前缀重复率和队列平均等待时间。当重复率低于某个阈值时切换到FCFS或LOF当重复率高但队列开始变长时可以从DFS_WEIGHT切换到计算更轻量的LPM。分层调度可以设想一个两阶段调度器。第一阶段用FCFS或简单规则积累一个批次大小的请求第二阶段在这个小批次内部使用LPM或DFS_WEIGHT进行精细排序以优化该批次内的计算复用。这样既控制了排序开销又获得了缓存收益。策略优先级混合例如可以定义一个主策略如DFS_WEIGHT但同时为等待时间超过某个阈值的请求赋予一个“优先级提升”因子将其在队列中提前兼顾效率和公平性。性能剖析与调试建议当你尝试应用或调整这些策略时以下实践或许有帮助埋点与度量在调度器代码中增加统计信息记录每种策略下被调度的请求其平均匹配前缀长度、调度决策耗时、以及最终在计算层实际的KV Cache命中率。数据是优化决策的基础。影子模式运行在生产环境中可以并行运行两套调度逻辑一套是实际生效的策略另一套是待评估的新策略影子模式。记录两者对同一批请求做出的排序决策差异并模拟计算各自的潜在缓存收益以安全地评估新策略效果。理解你的数据花时间分析线上请求的日志。计算prompt的相似度分布观察是否存在热门前缀。如果99%的请求前缀都毫无共性那么投入精力优化缓存感知策略的回报可能很低。在我经历的一次优化中我们将一个处理客服对话日志的分析服务从FCFS切换到LPM平均请求延迟下降了约40%因为大量的日志都有类似“用户咨询”、“问题描述”这样的共同开头。然而在另一个处理开放式创意写作的系统中DFS_WEIGHT的表现则略好于LPM因为请求常围绕几个核心主题如“科幻”、“历史”展开形成了较好的树状结构。最终调度策略的选择是一门平衡的艺术需要在缓存效率、计算开销、公平性和吞吐量之间找到最适合你当前业务场景的那个甜蜜点。理解这些源码级别的机制至少能让你在调整旋钮时清楚地知道每个选择背后意味着什么。