Prhub

#19548 fix: support PP2+CP8+TP8 (PP with context parallelism)

原始 PR 作者 whybeyoung 合并时间 2026-03-17 00:51 文件变更 2 提交数 7 评论 10 代码增减 +11 / -3

执行摘要

修复调度器以支持 PP 与 CP 并行,解决 H20 配置下 PP2+CP8+TP8 的通信问题。

PR body中引用PR #19504的评论提到:'i fixed it on H20 * 8 * 2',表明需要修复在H20集群上PP2+CP8+TP8配置下的问题。Issue评论中yiakwy-xpu-ml-framework-team确认修改后PP2+CP工作正常,否则无法生成输出。

建议技术管理者和工程师精读scheduler_pp_mixin.py中的通信逻辑修改,特别是CP广播的添加,以理解分布式数据同步机制。同时关注server_args.py中的配置检查变化,确保在启用PP与CP时正确设置enable_nsa_prefill_context_parallel等变量,并留意未解决的attn_cp_size讨论。

讨论亮点

Review中核心讨论点包括:

  • ShangmingCai建议直接移除断言,但whybeyoung解释需保留条件性断言以避免未实现/测试的无效配置启动,最终达成共识保留修改。
  • yiakwy-xpu-ml-framework-team提出attn_cp_size定义混淆问题,并指出TP=2时也应支持CP、H200无约束,但此疑虑未在本次PR中解决。
  • 所有讨论以批准结束,但遗留了关于attn_cp_size澄清和TP约束的后续工作。

实现拆解

主要改动在两个文件:

  • python/sglang/srt/managers/scheduler_pp_mixin.py:修改_pp_send_pyobj_to_next_stage_pp_recv_pyobj_from_prev_stage函数,添加条件self.attn_tp_rank == 0 and self.attn_cp_rank == 0限制发送/接收操作仅由TP0+CP0 rank执行,并在接收后添加CP广播broadcast_pyobj以确保所有CP rank获取数据。
  • python/sglang/srt/server_args.py:在_handle_context_parallelism方法中,将断言self.pp_size == 1改为条件性检查if not self.enable_nsa_prefill_context_parallel: assert self.pp_size == 1,从而允许PP与CP在NSA预填充上下文并行启用时共存,并设置attn_cp_size = tp_size以适配配置。
文件 模块 状态 重要度
python/sglang/srt/managers/scheduler_pp_mixin.py scheduling modified 8.0
python/sglang/srt/server_args.py configuration modified 5.0

关键符号

_pp_send_pyobj_to_next_stage _pp_recv_pyobj_from_prev_stage _handle_context_parallelism

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

断言移除讨论 设计

ShangmingCai 建议移除断言,whybeyoung 解释需保留以避免无效配置启动,因为未实现 / 测试的组合可能导致问题。

结论:保留条件性断言,仅当 enable_nsa_prefill_context_parallel 为假时触发。 · 已解决

attn_cp_size 澄清和 TP 约束 question

yiakwy-xpu-ml-framework-team 提出 attn_cp_size 定义混淆,并指出 TP=2 时也应支持 CP、H200 无约束,建议移除相关限制。

结论:未在本次 PR 中解决,遗留后续讨论。 · unresolved

风险与影响

技术风险包括:

  • 核心调度通信路径变更(如scheduler_pp_mixin.py中的发送/接收条件)可能引入死锁或数据不一致风险,特别是在分布式环境中依赖TP和CP rank为零的假设。
  • 配置检查放宽(server_args.py中的条件性断言)可能导致无效组合被允许,增加运行时错误风险。
  • 缺少针对PP+CP组合的专门测试,依赖现有CI覆盖,可能隐藏边缘情况。
  • 通信逻辑依赖于特定硬件配置(如H20),在其他平台可能需调整。

影响范围:

  • 用户:现在可以在支持NSA预填充上下文并行的配置下使用PP与CP组合,扩展大规模模型部署的并行选项,提升H20等硬件的资源利用率。
  • 系统:修复已知bug,确保PP2+CP8+TP8等配置正常工作,但需用户正确设置环境变量以避免性能下降或错误。
  • 团队:解决了关键并行支持问题,但遗留了attn_cp_size定义和TP约束的混淆,需后续澄清以维护代码清晰度。
核心路径变更 缺少专门测试 配置检查放宽

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论