执行摘要

• 一句话：ROCm 上广播 EAGLE verify 决策，修复 TP 死锁
• 推荐动作：值得精读。该 PR 虽小，但揭示了分布式投机解码中一个隐蔽的跨 rank 一致性陷阱，并展示了如何通过提升同步点作用域来覆盖所有执行路径。对理解 ROCm/CUDA 平台差异导致的死锁原因、以及 TP 集合通信前需要保证状态一致的原则很有价值。建议关注后续是否补充自动化测试。

功能与动机

PR body 明确指出：On ROCm with TP > 1, EAGLE speculative decoding can intermittently deadlock. Observed on 4x MI355X (gfx950), GLM-4.7-Flash, TP4, with 2 deadlocks in 5 runs。根因是 _is_hip 使 eagle_sample() 总走 greedy verify 路径，而该路径在不同 rank 上可能收到不同 draft 候选；采样路径原本已从 rank 0 广播结果，greedy 路径却没有，导致 rank 间接受 draft 数不一致，最终在 seq_lens / batch 组成上发散而死锁。

实现拆解

1. 问题定位：在 python/sglang/srt/speculative/eagle_utils.py 的 eagle_sample() 中，ROCm 因 _is_hip 恒真而总走 greedy verify 分支，该分支缺少采样路径已有的跨 rank 广播。
2. 核心变更：将原本位于采样路径内部的 tp_group.broadcast() 块上移，使其位于 sampling_fn() 调用之后、SIMULATE_ACC_LEN 处理之前，从而同时覆盖 greedy verify 与采样两条路径。
3. 同步对象：从 rank 0 广播 predict、accept_index、num_correct_drafts 三个可变张量，保证所有 TP rank 对本次 verify 的接受结果一致。
4. 并行组选择：保留 is_dp_attention_enabled() 分支，DP attention 场景下使用 attn_tp_group，否则使用 get_tp_group()，确保广播作用于正确的进程组。
5. 验证与配套：真机 4x MI355X、GLM-4.7-Flash、TP4 验证，死锁率从 5 次中 2 次降到 6 次中 0 次；未新增自动化单元测试，注释同步更新以解释原理。

关键文件：

• python/sglang/srt/speculative/eagle_utils.py（模块 投机解码；类别 source；类型 core-logic；符号 eagle_sample）: 核心修复文件：将 EAGLE verify 决策广播从采样路径提升到公共路径，覆盖 ROCm greedy verify 分支，修复 TP rank 间 seq_lens 发散导致的死锁。

关键符号：eagle_sample

关键源码片段

python/sglang/srt/speculative/eagle_utils.py

核心修复文件：将 EAGLE verify 决策广播从采样路径提升到公共路径，覆盖 ROCm greedy verify 分支，修复 TP rank 间 seq_lens 发散导致的死锁。

# eagle_sample() 内部，verify 采样完成后、SIMULATE_ACC_LEN 处理前
​
sampling_fn(
    predicts=predict, # mutable: 各位置的预测 token
    accept_index=accept_index, # mutable: 接受的 draft 位置索引
    accept_token_num=num_correct_drafts, # mutable: 接受的 draft 数量
    candidates=candidates,
    retrive_index=verify_input.retrieve_index,
    retrive_next_token=verify_input.retrieve_next_token,
    retrive_next_sibling=verify_input.retrieve_next_sibling,
    uniform_samples=coins,
    uniform_samples_for_final_sampling=coins_for_final_sampling,
    target_probs=target_probs,
    draft_probs=draft_probs,
    threshold_single=get_spec().speculative_accept_threshold_single,
    threshold_acc=get_spec().speculative_accept_threshold_acc,
    deterministic=True,
)
​
# 跨 TP rank 同步 verify 决策：
# 各 rank 送入本函数的张量存在微小差异（如 AMD all-reduce 非逐位一致、
# softmax/top_k 浮点不确定），使不同 rank 接受不同数量的 draft；
# 这会导致提交的 seq_lens 与 batch 组合发散，下一次 TP 集合通信即死锁。
# 故从 rank 0 广播三个关键张量，保证各 rank 决策一致。
tp_group = (
    get_parallel().attn_tp_group if is_dp_attention_enabled() else get_tp_group()
)
if tp_group.world_size > 1:
    tp_group.broadcast(predict, src=0)
    tp_group.broadcast(accept_index, src=0)
    tp_group.broadcast(num_correct_drafts, src=0)

评论区精华

HaiShaw 在 review 中建议将广播块上移以覆盖所有路径：

Can we move this block under L736, or simply move block L824-L836 out (indent -> left) to cover all cases?

作者回应：

Done. Moved this block one level up to cover both paths, and updated the comment accordingly.

该建议避免了仅对 ROCm greedy 分支打补丁的窄修复，从结构上保证了未来新增分支也能自动获得跨 rank 一致性，最终被采纳并获 approve。

• 广播块位置是否应覆盖所有路径 (design): 接受建议，广播块提升到采样与 greedy 两个路径共用的位置，结构上更稳健。

风险与影响

• 风险：
  • 性能风险：eagle_sample() 每次调用将额外执行 3 次进程组广播（predict、accept_index、num_correct_drafts），在 TP 规模较大时增加同步开销。但该广播发生在 speculative verify 后、下一次集合通信前，原本就是同步点，实际额外开销有限。
  • 正确性风险：广播以 rank 0 为准，若 rank 0 本身决策异常（如显存错误）会被放大到全体 rank，但本场景 rank 0 无特殊偏差，风险可接受。
  • 兼容性风险：广播现在对 CUDA 的 greedy 路径也生效（原本仅采样路径），可能引入微小额外同步，但不会改变 CUDA 行为（greedy 路径在 CUDA 上本身也需一致性）。
  • 测试风险：未添加自动化单测，回归依赖真机验证；后续若 AMD all-reduce 行为变化或并行组结构调整，可能重新引入同类问题。
• 影响：
  • 用户影响：AMD ROCm 上启用 TP>1 且使用 EAGLE/MTP 投机解码的用户是直接受益者，间歇性死锁被消除；CUDA 用户功能不变，仅有轻微广播开销。
  • 系统影响：涉及 eagle_utils.py 中投机解码核心路径，所有使用 EAGLE 后端的部署都会执行该广播逻辑（TP>1 时），属于调度热路径上的行为调整。
  • 团队影响：为后续修复其他平台（如 NPU）的投机一致性提供了参考模式；统一广播位置降低了维护复杂度。
  • 风险标记：核心路径变更, 缺少测试覆盖, 平台相关（ROCm）

关联脉络

• PR #34844 [Spec] Support MegaMoE for DSpark under dp attention: 同属 speculative decoding 在并行组（DP attention / TP）下的适配，该 PR 也涉及 is_dp_attention_enabled() 相关的组选择逻辑，与本 PR 的广播组选择相互关联。