Prhub

#21299 [Disagg DP-Balancing] Refactor Disagg Conn and Fix Hang with total_request/total_tokens Balancing

原始 PR 作者 weireweire 合并时间 2026-03-31 18:01 文件变更 8 提交数 8 评论 20 代码增减 +161 / -79

执行摘要

重构 disaggregation 接收器生命周期,修复因 DP 平衡导致的 hang 问题。

修复PD disaggregation hang问题(issue #21297),当prefill使用内部DP平衡如total_requests时,decode可能在bootstrap_room -> prefill_dp_rank映射可用前收到请求,导致handshake不一致和stall。

建议技术管理者关注此PR,因为它解决了核心disaggregation路径的关键bug,并引入了重要的设计变更。工程师应精读decode.py和common/conn.py的改动,理解receiver初始化的新流程,并注意review中未解决的failure handling问题,以便在后续开发中处理。

讨论亮点

review中主要讨论:

1) ShangmingCai质疑leader-based resolution可能导致其他ranks无法初始化prefill_info_table,weireweire回应需添加sync逻辑并修改实现;
2) 关于failure处理,ShangmingCai建议在_update_handshake_waiters中处理abort以避免double free,weireweire表示未在此PR解决,留待后续决策;
3) 设计决策上,基于issue #21680,从leader-based workaround转向receiver生命周期重构,确保一致性。

实现拆解

实现方案分为三层:

1) 在base/conn.py中定义抽象方法init()send_metadata(),分离初始化和元数据发送;
2) 在common/conn.py中重构CommonKVReceiver,__init__仅设置基本属性,init()方法解析prefill DP-rank并设置bootstrap信息;
3) 在decode.py中修改add()_create_receiver_and_enqueue(),立即创建receiver但延迟调用init(),pending请求列表改用DecodeRequest对象;
4) 其他具体conn.py文件(fake、mooncake、mori、nixl)适配新方法;
5) 测试文件添加TestDisaggregationDPAttentionTotalRequests和TestDisaggregationDPAttentionTotalTokens类,验证DP平衡路径。

文件 模块 状态 重要度
python/sglang/srt/disaggregation/decode.py disaggregation decode modified 8.0
python/sglang/srt/disaggregation/common/conn.py disaggregation common modified 7.0
python/sglang/srt/disaggregation/base/conn.py disaggregation base modified 6.0
test/registered/distributed/test_disaggregation_dp_attention.py test modified 5.0

关键符号

init send_metadata _create_receiver_and_enqueue _resolve_pending_reqs

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

评论区精华

Leader-based resolution 的正确性 正确性

ShangmingCai 质疑如果只有 leader rank 做 bootstrapping,其他 ranks 可能无法初始化 prefill_info_table

结论:weireweire 同意需添加 sync 逻辑,修改为各 rank 自行运行但同步 ready ones · 已解决

Failure handling 的设计 设计

ShangmingCai 建议在 _ 更新 _handshake_waiters 中处理 abort 以避免 double free

结论:weireweire 表示未在此 PR 解决,留待后续决策 · unresolved

Receiver 生命周期重构 设计

基于 issue #21680,从 leader-based workaround 转向 receiver 初始化与元数据发送分离

结论:采纳了 receiver lifecycle refactor,确保一致性 · 已解决

风险与影响

风险包括:

1) 新引入的init()send_metadata()分离可能增加复杂性,如果实现不一致可能导致handshake失败;
2) failure处理逻辑不完整,如review中讨论的abort问题,可能引发资源泄漏或错误报告;
3) 同步机制依赖于decode ranks间的一致状态,如果广播或sync失败,可能重现hang问题;
4) 测试覆盖虽增加,但可能未覆盖所有edge cases,特别是故障场景。

影响范围:

1) 对用户:修复了使用DP平衡(如total_requests/total_tokens)时的潜在hang问题,提升disaggregation服务的可靠性;
2) 对系统:改善了decode侧pending请求处理的一致性,防止因ranks间状态不一致导致的stall;
3) 对团队:需要更新对disaggregation连接生命周期的理解,新设计可能影响后续开发,如failure handling需要额外工作。

同步机制风险 failure 处理不完全 兼容性变更

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论