From 6d767a9a5bff87ac05640ca581ef07d0e57aa3ca Mon Sep 17 00:00:00 2001 From: tzhouam Date: Fri, 31 Jul 2026 11:08:13 +0800 Subject: [PATCH] [BugFix] Make pure-TP KV receive fallback all-or-nothing across ranks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Signed-off-by: tzhouam --- .../test_kv_recv_tp_consensus.py | 199 ++++++++++++++++++ .../omni_connectors/test_tp_rank_aware.py | 8 +- .../omni_connectors/kv_transfer_manager.py | 68 ++++++ 3 files changed, 274 insertions(+), 1 deletion(-) create mode 100644 tests/distributed/omni_connectors/test_kv_recv_tp_consensus.py diff --git a/tests/distributed/omni_connectors/test_kv_recv_tp_consensus.py b/tests/distributed/omni_connectors/test_kv_recv_tp_consensus.py new file mode 100644 index 000000000..f890ba059 --- /dev/null +++ b/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] diff --git a/tests/distributed/omni_connectors/test_tp_rank_aware.py b/tests/distributed/omni_connectors/test_tp_rank_aware.py index 261de0b32..178069d65 100644 --- a/tests/distributed/omni_connectors/test_tp_rank_aware.py +++ b/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 ──────────────────────────────────────────── diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index efed75241..760ab37b8 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/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)