[AMD] Broadcast the EAGLE greedy verify decision across TP ranks on ROCm (#34238)

This commit is contained in:
Zhiyao Jiang
2026-08-14 22:37:45 -07:00
committed by GitHub
parent aeee1562e6
commit 7216a44dd8
+5 -7
View File
@@ -824,14 +824,12 @@ def eagle_sample(
deterministic=True, deterministic=True,
) )
# Sync sampling results across TP ranks: different GPUs may # Sync the verify decision across TP ranks: small per-rank differences in the
# produce slightly different target_probs due to floating-point # tensors feeding this point can make one rank accept a different number of
# non-determinism in softmax/top_k/top_p, causing different # drafts, which desynchronizes the committed seq_lens and deadlocks the next
# sampled tokens. Broadcast from rank 0 to ensure consistency. # TP collective. Broadcast from rank 0 to ensure consistency.
tp_group = ( tp_group = (
get_parallel().attn_tp_group get_parallel().attn_tp_group if is_dp_attention_enabled() else get_tp_group()
if is_dp_attention_enabled()
else get_tp_group()
) )
if tp_group.world_size > 1: if tp_group.world_size > 1:
tp_group.broadcast(predict, src=0) tp_group.broadcast(predict, src=0)