执行摘要

• 一句话：修复 total_tokens 负载均衡中因缺少启发式 token 增量导致请求堆积的问题。
• 推荐动作：该 PR 值得精读，尤其是 DPBudget.dispatch 方法的修改，展示了如何在负载均衡中平衡启发式增量和快照校正的设计。关注 gemini-code-assist[bot] 提出的多模态 token 低估问题，这可能影响未来扩展。

功能与动机

PR body 指出，在生产日志中观察到使用 total_tokens 负载均衡方法时，预分配请求（prealloc requests）会堆积在单个 DP worker 上（如 DP0 有 21 个预分配请求），而其他 DP worker 几乎空闲。根本原因是 DPBudget.dispatch() 在调度后只增加 total_requests 的启发式增量，未增加 total_tokens，导致在两次调度快照之间，total_tokens 值过时，所有新请求持续路由到同一 worker，造成 KV 内存不足和队列阻塞。

实现拆解

1. 修改 DPBudget.dispatch 方法签名和逻辑：

   • 文件：python/sglang/srt/managers/data_parallel_controller.py
   • 关键符号：dispatch 方法
   • 变更：为 dispatch 方法添加 estimated_tokens: int = 0 参数，并在选择目标 rank 后，增加 self.total_tokens[target_rank] += estimated_tokens 行。
   • 原因：使 total_tokens 负载均衡方法在调度时能基于请求的输入 token 数进行启发式增量，避免请求堆积。
   • 影响：total_tokens 值在调度间隙得到临时更新，但后续 update_budget() 会用真实快照覆盖，确保自校正。

2. 更新 total_tokens_scheduler 调用：

   • 文件：python/sglang/srt/managers/data_parallel_controller.py
   • 关键符号：total_tokens_scheduler 方法
   • 变更：在调用 dp_budget.dispatch() 前计算 estimated_tokens = len(req.input_ids)，并将其作为参数传入。
   • 原因：为 TOTAL_TOKENS 方法提供 token 估计值，匹配 get_load() 报告的队列请求度量。
   • 影响：仅影响 TOTAL_TOKENS 负载均衡路径；TOTAL_REQUESTS 和 ROUND_ROBIN 方法不受影响，保持向后兼容。

3. 测试与配置配套：

   • 本次变更未包含直接对应的测试文件修改，但 PR 已通过 CI 标签（run-ci）运行测试。

关键文件：

• python/sglang/srt/managers/data_parallel_controller.py（模块 调度器；类别 source；类型 core-logic；符号 dispatch, total_tokens_scheduler）: 这是唯一变更的文件，包含负载均衡的核心逻辑修改，直接影响 DP worker 的请求调度。

关键符号：dispatch, total_tokens_scheduler

评论区精华

review 中，gemini-code-assist[bot] 指出一个潜在问题：使用 len(req.input_ids) 作为 estimated_tokens 对于多模态请求可能显著低估实际负载，因为图像等会扩展为更多 token（例如 576 个以上）。建议考虑从 TokenizerManager 传递更准确的估计值，或至少记录此限制。此评论未得到直接回复或解决，但 PR 仍被合并，表明团队可能认为当前启发式增量已足够，或计划后续处理。其他 reviewer（Ratish1 和 hnyls2002）简单批准。

• 多模态请求的 token 估计准确性 (correctness): 未直接解决，PR 仍被合并，可能视为可接受的启发式或待后续优化。

风险与影响

• 风险：
  • 回归风险：低。变更仅影响 TOTAL_TOKENS 负载均衡路径，且 estimated_tokens 默认值为 0，确保其他方法不受影响。update_budget() 使用赋值覆盖，防止启发式增量永久漂移。
  • 性能风险：中。对于多模态请求，len(req.input_ids) 可能低估 token 数，导致负载均衡不准确，可能仍引发轻微堆积，但比修复前大幅改善。
  • 兼容性风险：低。API 向后兼容，因添加了默认参数。
  • 安全风险：无显著安全影响。
• 影响：
  • 用户影响：使用解耦部署和 total_tokens 负载均衡的用户将观察到请求分布更均匀，减少因单 worker 过载导致的 KV 内存耗尽和延迟。
  • 系统影响：提升 DP worker 间的资源利用率，避免预分配队列阻塞，增强系统稳定性和吞吐量。
  • 团队影响：修复了生产环境中观察到的关键负载均衡缺陷，需关注多模态请求的潜在低估问题。
  • 风险标记：核心路径变更, 多模态负载估计不足

关联脉络

• PR #22493 Add MambaPool kvcache offloading during retraction: 同样涉及调度和 KV 缓存管理，关注负载均衡和内存优化。
• PR #22911 [perf] support return_routed_experts with overlap scheduling: 涉及调度性能优化，与本 PR 的负载均衡改进相关。
• PR #23173 fix: pass v_head_dim to MHA KV pools and validate MiMo HiCache geometry: 涉及 KV 缓存和 HiCache，与本 PR 中因 KV 内存耗尽导致的预分配队列阻塞问题相关。