[BugFix] Make pure-TP KV receive fallback all-or-nothing across ranks

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>

tests/distributed/omni_connectors/test_kv_recv_tp_consensus.py

@@ -0,0 +1,199 @@
+"""Pure-TP KV receive consensus (#5627) — CPU-testable mechanism.
+
+In pure TP (LOCAL role) each rank fetches its own KV shard independently and
+``distribute_kv_cache`` is a no-op, so one rank's silent receive timeout used
+to send only that rank into the pipeline's no-KV prefill — desyncing TP
+collectives until the NCCL watchdog killed the worker. These tests cover the
+all-or-nothing verdict exchange: skip guards, both rank roles, KV dropping on
+a peer miss, and the wiring from both distributed-receive entry points.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    OmniKVCacheConfig,
+    OmniKVTransferManager,
+    ReceiveRole,
+    _TransferTopoConfig,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeGroup:
+    """Scripted stand-in for one rank's view of a GroupCoordinator."""
+
+    def __init__(self, rank_in_group, world_size, peer_flags=None, verdict=None):
+        self.rank_in_group = rank_in_group
+        self.world_size = world_size
+        self._peer_flags = list(peer_flags or [])
+        self._verdict = verdict
+        self.sent = []
+        self.broadcasts = []
+
+    def recv_object(self, src):
+        return self._peer_flags.pop(0)
+
+    def send_object(self, obj, dst):
+        self.sent.append((obj, dst))
+
+    def broadcast_object(self, obj=None, src=0):
+        self.broadcasts.append(obj)
+        return obj if obj is not None else self._verdict
+
+
+def _topo(role=ReceiveRole.LOCAL, tp_active=True, world=None):
+    return _TransferTopoConfig(
+        role=role,
+        tp_active=tp_active,
+        cfg_size=1,
+        cfg_rank=0,
+        cfg_group=None,
+        sp_size=1,
+        sp_rank=0,
+        sp_group=None,
+        world=world,
+    )
+
+
+def _manager(topo, need_recv_cache=True):
+    mgr = OmniKVTransferManager(OmniKVCacheConfig(need_recv_cache=need_recv_cache))
+    mgr._topo_config = topo
+    return mgr
+
+
+def _req_with_kv(rid="r1"):
+    kv = SimpleNamespace(key_cache=[], value_cache=[])
+    return SimpleNamespace(
+        request_id=rid,
+        past_key_values=kv,
+        kv_metadata={"ropes": [3]},
+        sampling_params=SimpleNamespace(past_key_values=kv, kv_metadata={"ropes": [3]}),
+    )
+
+
+# --------------------------------------------------------------------------- #
+#  Skip guards — no exchange may happen (a one-sided exchange would deadlock)
+# --------------------------------------------------------------------------- #
+
+
+def test_pass_through_outside_pure_tp():
+    group = _FakeGroup(0, 2, peer_flags=[False])
+    for topo in (
+        _topo(role=ReceiveRole.LEADER, tp_active=True, world=group),
+        _topo(role=ReceiveRole.LOCAL, tp_active=False, world=group),
+        _topo(role=ReceiveRole.LOCAL, tp_active=True, world=None),
+    ):
+        mgr = _manager(topo)
+        assert mgr._tp_local_receive_consensus(_req_with_kv(), True) is True
+        assert group.sent == [] and group.broadcasts == []
+
+
+def test_pass_through_single_rank_group_no_recv_config_and_dummy():
+    group = _FakeGroup(0, 1)
+    assert _manager(_topo(world=group))._tp_local_receive_consensus(_req_with_kv(), True) is True
+
+    group2 = _FakeGroup(0, 2, peer_flags=[False])
+    mgr = _manager(_topo(world=group2), need_recv_cache=False)
+    assert mgr._tp_local_receive_consensus(_req_with_kv(), True) is True
+
+    mgr3 = _manager(_topo(world=group2))
+    dummy = _req_with_kv(rid=DUMMY_DIFFUSION_REQUEST_ID)
+    assert mgr3._tp_local_receive_consensus(dummy, False) is False
+    assert group2.sent == [] and group2.broadcasts == []
+
+
+# --------------------------------------------------------------------------- #
+#  Rank 0: collects flags, decides, broadcasts
+# --------------------------------------------------------------------------- #
+
+
+def test_rank0_drops_kv_when_peer_missed():
+    group = _FakeGroup(0, 2, peer_flags=[False])
+    mgr = _manager(_topo(world=group))
+    req = _req_with_kv()
+
+    assert mgr._tp_local_receive_consensus(req, True) is False
+    assert group.broadcasts == [False]  # peers were told to fall back too
+    assert req.past_key_values is None
+    assert req.kv_metadata is None
+    assert req.sampling_params.past_key_values is None
+    assert req.sampling_params.kv_metadata is None
+
+
+def test_rank0_keeps_kv_when_all_received():
+    group = _FakeGroup(0, 2, peer_flags=[True])
+    mgr = _manager(_topo(world=group))
+    req = _req_with_kv()
+
+    assert mgr._tp_local_receive_consensus(req, True) is True
+    assert group.broadcasts == [True]
+    assert req.past_key_values is not None
+    assert req.sampling_params.past_key_values is not None
+
+
+def test_rank0_own_miss_fails_verdict_even_if_peers_received():
+    group = _FakeGroup(0, 3, peer_flags=[True, True])
+    mgr = _manager(_topo(world=group))
+    assert mgr._tp_local_receive_consensus(_req_with_kv(), False) is False
+    assert group.broadcasts == [False]
+
+
+# --------------------------------------------------------------------------- #
+#  Rank > 0: reports flag, honors verdict
+# --------------------------------------------------------------------------- #
+
+
+def test_follower_rank_drops_kv_on_negative_verdict():
+    group = _FakeGroup(1, 2, verdict=False)
+    mgr = _manager(_topo(world=group))
+    req = _req_with_kv()
+
+    assert mgr._tp_local_receive_consensus(req, True) is False
+    assert group.sent == [(True, 0)]
+    assert req.past_key_values is None
+    assert req.sampling_params.past_key_values is None
+
+
+def test_follower_rank_keeps_kv_on_positive_verdict():
+    group = _FakeGroup(1, 2, verdict=True)
+    mgr = _manager(_topo(world=group))
+    req = _req_with_kv()
+
+    assert mgr._tp_local_receive_consensus(req, True) is True
+    assert group.sent == [(True, 0)]
+    assert req.past_key_values is not None
+
+
+# --------------------------------------------------------------------------- #
+#  Wiring — both distributed-receive entry points run the consensus
+# --------------------------------------------------------------------------- #
+
+
+def _wired_manager(monkeypatch, consensus_calls):
+    mgr = _manager(_topo(world=_FakeGroup(0, 2)))
+
+    def fake_consensus(req, received):
+        consensus_calls.append(received)
+        return False
+
+    monkeypatch.setattr(mgr, "_tp_local_receive_consensus", fake_consensus)
+    monkeypatch.setattr(mgr, "receive_multi_kv_cache", lambda *a, **k: True)
+    return mgr
+
+
+def test_receive_multi_kv_cache_distributed_runs_consensus(monkeypatch):
+    calls = []
+    mgr = _wired_manager(monkeypatch, calls)
+    assert mgr.receive_multi_kv_cache_distributed(_req_with_kv()) is False
+    assert calls == [True]
+
+
+def test_consume_and_distribute_kv_cache_runs_consensus(monkeypatch):
+    calls = []
+    mgr = _wired_manager(monkeypatch, calls)
+    assert mgr.consume_and_distribute_kv_cache(_req_with_kv()) is False
+    assert calls == [True]

tests/distributed/omni_connectors/test_tp_rank_aware.py

@@ -681,7 +681,9 @@ class TestDistributedReceive:
     def test_tp_without_cfg_keeps_independent_receive_path(self):
         mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1)
         req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace())
-        world_group = _MockBroadcastGroup(world_size=2, rank_in_group=1)
+        # broadcast_value=True: rank 0's verdict says every rank received its
+        # shard, so the per-rank receive result stands (see #5627 consensus).
+        world_group = _MockBroadcastGroup(world_size=2, rank_in_group=1, broadcast_value=True)
         mgr.receive_multi_kv_cache = MagicMock(return_value=True)
 
         with (
@@ -708,6 +710,10 @@ class TestDistributedReceive:
             assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True
 
         mgr.receive_multi_kv_cache.assert_called_once_with(req, None, torch.device("cpu"))
+        # Pure TP still fetches per-rank (no leader→follower distribution),
+        # but each rank now reports its receive flag for the all-or-nothing
+        # verdict instead of silently diverging on a miss (#5627).
+        assert world_group.send_calls == [(0, True)]
 
     # ── SP-only scenarios ────────────────────────────────────────────
 

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -758,6 +758,20 @@ class OmniKVTransferManager:
 
         return kv_payload
 
+    @staticmethod
+    def _clear_request_kv_payload(req: Any) -> None:
+        """Drop every request-side KV field a receive may have applied.
+
+        Enumerates fields via ``_collect_request_kv_payload`` so clearing stays
+        symmetric with what apply/broadcast set — the pipeline then takes the
+        same no-KV path as a rank that never received.
+        """
+        for key in OmniKVTransferManager._collect_request_kv_payload(req):
+            if key.startswith("sp."):
+                setattr(req.sampling_params, key[3:], None)
+            else:
+                setattr(req, key, None)
+
     @staticmethod
     def _apply_request_kv_payload(
         req: Any,
@@ -1747,6 +1761,58 @@ class OmniKVTransferManager:
         combined = combined.to(device)
         return self._unpack_kv_payload(combined)
 
+    def _tp_local_receive_consensus(self, req: Any, received: bool) -> bool:
+        """All-or-nothing receive verdict across pure-TP ranks (LOCAL role).
+
+        In pure TP every rank fetches its own KV shard with an independent
+        timeout and ``distribute_kv_cache`` is a no-op, so a single rank's
+        silent miss sends only that rank into the pipeline's no-KV prefill —
+        it then enters TP collectives its peers never join, and the NCCL
+        watchdog kills the worker (#5627). Exchange per-rank receive flags;
+        if any rank missed, every rank drops its KV so all take the same
+        fallback path together.
+
+        The skip guards below depend only on topology, config, and the
+        request id — identical on every rank — so ranks always agree on
+        whether the exchange happens.
+        """
+        pt = self.topo_config
+        group = pt.world
+        if not (pt.is_local and pt.tp_active) or group is None:
+            return received
+        if getattr(group, "world_size", 1) <= 1:
+            return received
+        if not self.config.need_recv_cache:
+            return received
+        request_id = self._resolve_request_id(req)
+        if request_id is None or OmniDiffusionRequest.is_dummy_run_request_id(request_id):
+            return received
+
+        if group.rank_in_group == 0:
+            flags = [bool(received)]
+            flags.extend(bool(group.recv_object(src=src)) for src in range(1, group.world_size))
+            all_received = all(flags)
+            group.broadcast_object(all_received, src=0)
+        else:
+            group.send_object(bool(received), dst=0)
+            all_received = bool(group.broadcast_object(None, src=0))
+
+        if not all_received:
+            if received:
+                logger.warning(
+                    "KV receive for %s missed on a peer TP rank; dropping this rank's "
+                    "KV so all ranks take the no-KV fallback together instead of "
+                    "desyncing TP collectives",
+                    request_id,
+                )
+                self._clear_request_kv_payload(req)
+            else:
+                logger.warning(
+                    "KV receive for %s missed on this or a peer TP rank; all ranks fall back without KV",
+                    request_id,
+                )
+        return all_received
+
     def receive_multi_kv_cache_distributed(
         self,
         req: Any,
@@ -1757,6 +1823,7 @@ class OmniKVTransferManager:
         received = False
         if not self.topo_config.is_follower:
             received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device)
+        received = self._tp_local_receive_consensus(req, received)
         kv_payload = self.distribute_kv_cache(req, target_device, received=received)
         if kv_payload is not None:
             self._apply_request_kv_payload(req, kv_payload, target_device)
@@ -1785,6 +1852,7 @@ class OmniKVTransferManager:
         if not received and not self.topo_config.is_follower and not payload_consumed:
             logger.debug("KV prefetch miss for %s; falling back to sync receive", self._resolve_request_id(req))
             received = self.receive_multi_kv_cache(req, None, target_device)
+        received = self._tp_local_receive_consensus(req, received)
         kv_payload = self.distribute_kv_cache(req, target_device, received=received)
         if kv_payload is not None:
             self._apply_request_kv_payload(req, kv_payload, target_device)