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..a7e4b8562 --- /dev/null +++ b/tests/distributed/omni_connectors/test_kv_recv_tp_consensus.py @@ -0,0 +1,257 @@ +"""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 +identical-state verdict exchange: skip guards, both rank roles, KV dropping +on any cross-rank divergence (primary or CFG companion roles), that the +exchange runs on the TP group and never on ``world`` (which also holds DP/PP +ranks working on other requests), 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_signatures=None, verdict=None): + self.rank_in_group = rank_in_group + self.world_size = world_size + self._peer_signatures = list(peer_signatures or []) + self._verdict = verdict + self.sent = [] + self.broadcasts = [] + + def recv_object(self, src): + return self._peer_signatures.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 + + @property + def touched(self): + return bool(self.sent or self.broadcasts) + + +def _topo(role=ReceiveRole.LOCAL, tp_active=True, tp_group=None, world=None): + tp_size = tp_group.world_size if tp_group is not None else 1 + 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, + tp_group=tp_group, + tp_size=tp_size, + ) + + +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", cfg_text=False): + kv = SimpleNamespace(key_cache=[], value_cache=[]) + sp = SimpleNamespace(past_key_values=kv, kv_metadata={"ropes": [3]}) + if cfg_text: + sp.cfg_text_past_key_values = SimpleNamespace(key_cache=[], value_cache=[]) + return SimpleNamespace( + request_id=rid, + past_key_values=kv, + kv_metadata={"ropes": [3]}, + sampling_params=sp, + ) + + +def _signature_of(req, received=True): + return (received, tuple(sorted(OmniKVTransferManager._collect_request_kv_payload(req)))) + + +# --------------------------------------------------------------------------- # +# Skip guards — no exchange may happen (a one-sided exchange would deadlock) +# --------------------------------------------------------------------------- # + + +def test_pass_through_outside_pure_tp(): + group = _FakeGroup(0, 2, peer_signatures=[(False, ())]) + for topo in ( + _topo(role=ReceiveRole.LEADER, tp_active=True, tp_group=group), + _topo(role=ReceiveRole.LOCAL, tp_active=False, tp_group=group), + _topo(role=ReceiveRole.LOCAL, tp_active=True, tp_group=None), + ): + mgr = _manager(topo) + assert mgr._tp_local_receive_consensus(_req_with_kv(), True) is True + assert not group.touched + + +def test_pass_through_single_rank_group_no_recv_config_and_dummy(): + group = _FakeGroup(0, 1) + assert _manager(_topo(tp_group=group))._tp_local_receive_consensus(_req_with_kv(), True) is True + + group2 = _FakeGroup(0, 2, peer_signatures=[(False, ())]) + mgr = _manager(_topo(tp_group=group2), need_recv_cache=False) + assert mgr._tp_local_receive_consensus(_req_with_kv(), True) is True + + mgr3 = _manager(_topo(tp_group=group2)) + dummy = _req_with_kv(rid=DUMMY_DIFFUSION_REQUEST_ID) + assert mgr3._tp_local_receive_consensus(dummy, False) is False + assert not group2.touched + + +def test_world_group_is_never_used_for_the_exchange(): + """DP replicas / PP stages live in ``world`` and work on other requests — + the exchange must be scoped to the TP group (review #5636, P1).""" + tp_group = _FakeGroup(0, 2, peer_signatures=[(True, ())]) + world = _FakeGroup(0, 4, peer_signatures=[(False, ()), (False, ()), (False, ())]) + mgr = _manager(_topo(tp_group=tp_group, world=world)) + req = SimpleNamespace(request_id="r1", sampling_params=SimpleNamespace()) + + assert mgr._tp_local_receive_consensus(req, True) is True + assert tp_group.touched + assert not world.touched + + +# --------------------------------------------------------------------------- # +# Rank 0: collects signatures, decides, broadcasts +# --------------------------------------------------------------------------- # + + +def test_rank0_drops_kv_when_peer_missed_primary(): + req = _req_with_kv() + peer_sig = (False, ()) + group = _FakeGroup(0, 2, peer_signatures=[peer_sig]) + mgr = _manager(_topo(tp_group=group)) + + 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_ranks_match(): + req = _req_with_kv() + group = _FakeGroup(0, 2, peer_signatures=[_signature_of(req)]) + mgr = _manager(_topo(tp_group=group)) + + 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_drops_kv_on_companion_role_divergence(): + """Primary received everywhere but cfg_text only landed on this rank — + divergent branch caches would desync TP collectives (review #5636, P1).""" + req = _req_with_kv(cfg_text=True) + peer_sig = _signature_of(_req_with_kv(cfg_text=False)) # peer lacks cfg_text + group = _FakeGroup(0, 2, peer_signatures=[peer_sig]) + mgr = _manager(_topo(tp_group=group)) + + assert mgr._tp_local_receive_consensus(req, True) is False + assert group.broadcasts == [False] + assert req.past_key_values is None + assert req.sampling_params.cfg_text_past_key_values is None + + +def test_rank0_keeps_symmetric_partial_state(): + """cfg_text absent on every rank is consistent — Bagel's empty + unconditional branch — and must NOT be cleared or failed.""" + req = _req_with_kv(cfg_text=False) + group = _FakeGroup(0, 2, peer_signatures=[_signature_of(req)]) + mgr = _manager(_topo(tp_group=group)) + + assert mgr._tp_local_receive_consensus(req, True) is True + assert req.past_key_values is not None + + +def test_rank0_own_miss_fails_verdict_even_if_peers_received(): + req = SimpleNamespace(request_id="r1", sampling_params=SimpleNamespace()) + group = _FakeGroup(0, 3, peer_signatures=[(True, ("past_key_values",)), (True, ("past_key_values",))]) + mgr = _manager(_topo(tp_group=group)) + assert mgr._tp_local_receive_consensus(req, False) is False + assert group.broadcasts == [False] + + +# --------------------------------------------------------------------------- # +# Rank > 0: reports signature, honors verdict +# --------------------------------------------------------------------------- # + + +def test_follower_rank_drops_kv_on_negative_verdict(): + req = _req_with_kv() + sig_before_clear = _signature_of(req) + group = _FakeGroup(1, 2, verdict=False) + mgr = _manager(_topo(tp_group=group)) + + assert mgr._tp_local_receive_consensus(req, True) is False + assert group.sent == [(sig_before_clear, 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(): + req = _req_with_kv() + group = _FakeGroup(1, 2, verdict=True) + mgr = _manager(_topo(tp_group=group)) + + assert mgr._tp_local_receive_consensus(req, True) is True + assert group.sent == [(_signature_of(req), 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(tp_group=_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..3fdc4820a 100644 --- a/tests/distributed/omni_connectors/test_tp_rank_aware.py +++ b/tests/distributed/omni_connectors/test_tp_rank_aware.py @@ -682,6 +682,10 @@ class TestDistributedReceive: 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 TP rank has the + # same KV state, so the per-rank receive result stands (see #5627 + # consensus). The exchange runs on the TP group, not world. + tp_group = _MockBroadcastGroup(world_size=2, rank_in_group=1, broadcast_value=True) mgr.receive_multi_kv_cache = MagicMock(return_value=True) with ( @@ -704,10 +708,19 @@ class TestDistributedReceive: return_value=0, ), patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None), + patch("vllm.distributed.parallel_state.get_tp_group", return_value=tp_group), ): 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 KV-state signature on the TP group + # for the identical-state verdict instead of silently diverging on a + # miss (#5627). World must stay untouched — it holds DP/PP ranks. + assert len(tp_group.send_calls) == 1 + dst, sig = tp_group.send_calls[0] + assert dst == 0 and sig[0] is True + assert world_group.send_calls == [] # ── 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 8616cf9a6..6b83349f2 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -69,6 +69,12 @@ class _TransferTopoConfig: sp_rank: int sp_group: Any world: Any + # The stage's tensor-model-parallel group (vLLM parallel state). Only the + # ranks in this group execute the same request in lockstep — DP replicas + # and PP stages in ``world`` do not — so per-request consensus exchanges + # must use this group, never ``world``. + tp_group: Any = None + tp_size: int = 1 @property def is_leader(self) -> bool: @@ -642,6 +648,16 @@ class OmniKVTransferManager: else: role = ReceiveRole.LEADER if world_rank == 0 else ReceiveRole.FOLLOWER + # The stage's TP group comes from vLLM parallel state (the TP-sharded + # LM path); it may legitimately be uninitialized on non-TP stages. + try: + from vllm.distributed.parallel_state import get_tp_group + + tp_group = get_tp_group() + tp_size = tp_group.world_size + except Exception: + tp_group, tp_size = None, 1 # vLLM TP not initialized + return _TransferTopoConfig( role=role, tp_active=tp_active, @@ -652,6 +668,8 @@ class OmniKVTransferManager: sp_rank=sp_rank, sp_group=sp_group, world=world, + tp_group=tp_group, + tp_size=tp_size, ) def _resolve_sender_info( @@ -764,6 +782,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, @@ -1753,6 +1785,64 @@ class OmniKVTransferManager: combined = combined.to(device) return self._unpack_kv_payload(combined) + def _tp_local_receive_consensus(self, req: Any, received: bool) -> bool: + """Identical-KV-state 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). Each rank reports a signature of + its applied KV state — the primary-receive flag plus the set of + applied payload roles, so partially-received CFG companion KV + (``cfg_text``/``cfg_img``) diverging between ranks is caught too. If + the signatures differ, every rank drops its KV so all take the same + fallback path together; identical-but-partial state is kept, since + symmetric absence is consistent (e.g. Bagel's empty unconditional + branch). + + The exchange runs over the stage's tensor-parallel group only: those + are the ranks executing this request in lockstep. ``world`` would + also contain DP replicas / PP stages working on different requests — + mixing their flags corrupts the verdict, and a replica skipping the + exchange (dummy request) while another enters it would deadlock. + Within the TP group the skip guards below depend only on topology, + config, and the request id — identical on every TP rank — so ranks + always agree on whether the exchange happens. + """ + pt = self.topo_config + group = pt.tp_group + if not (pt.is_local and pt.tp_active) or group is None or pt.tp_size <= 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 + + signature = (bool(received), tuple(sorted(self._collect_request_kv_payload(req)))) + if group.rank_in_group == 0: + signatures = [signature] + signatures.extend(group.recv_object(src=src) for src in range(1, pt.tp_size)) + consistent = all(sig == signature for sig in signatures) + group.broadcast_object(consistent, src=0) + else: + group.send_object(signature, dst=0) + consistent = bool(group.broadcast_object(None, src=0)) + + if not consistent: + logger.warning( + "KV state for %s diverged across TP ranks (this rank: received=%s, roles=%s); " + "dropping KV on every rank so all take the no-KV fallback together " + "instead of desyncing TP collectives", + request_id, + signature[0], + list(signature[1]), + ) + self._clear_request_kv_payload(req) + return False + return received + def receive_multi_kv_cache_distributed( self, req: Any, @@ -1763,6 +1853,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) @@ -1791,6 +1882,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)