Fixes the recurring Bagel tp_vae_patch_parallel_2 NCCL ALLREDUCE
timeout (#5627, previously #5487): in pure TP the transfer topology is
LOCAL on every rank — each rank fetches its own per-rank KV key with an
independent recv_timeout and distribute_kv_cache is a no-op. A receive
timeout returns silently, so a single rank that missed its shard fell
back alone into the pipeline's no-KV prefill (Bagel
forward_cache_update_text through the TP-sharded LM) and entered TP
allreduces its peers never joined. The hung [~20, 3584] allreduce
(NumelIn=71680) blocked 600s until the NCCL watchdog killed the worker,
after which every request failed with "DiffusionExecutor is closed".
PR #5482 fixed the orchestrator-level CFG bundle race but never touched
this lower KV-transfer layer, which is why the signature recurred at
its own merge commit.
After the per-rank receive, pure-TP ranks now exchange their receive
flags (rank 0 collects and broadcasts the verdict over the stage world
group's CPU object plane) and act on it together: if any rank missed,
ranks that did receive drop their applied KV — cleared symmetrically
with what apply/broadcast set — so every rank takes the identical no-KV
prefill path and the TP collective schedule can never diverge. The skip
guards (topology, need_recv_cache, dummy warmup id) depend only on
values identical across ranks, so ranks always agree on whether the
exchange happens; a rank whose connector failed to initialize still
participates and reports a miss instead of deadlocking the exchange.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
