From bac3eafc18d2e5389f9db0f50a788e94969f73ea Mon Sep 17 00:00:00 2001 From: tzhouam Date: Sat, 1 Aug 2026 16:28:56 +0800 Subject: [PATCH] review: scope KV receive consensus to the TP group, compare full KV state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address both P1 review comments: - The exchange now runs on the stage's tensor-model-parallel group (vLLM parallel state) instead of pt.world. World also holds DP replicas and PP stages working on different requests, so the old group could mix verdicts across requests and — worse — deadlock if one replica skipped the exchange (dummy request) while another entered it. Within the TP group every skip guard is symmetric. - Ranks now exchange (primary_received, sorted applied-KV roles) instead of a bare flag, using the same field enumeration as apply/broadcast, so partially-received CFG companion KV (cfg_text/cfg_img) diverging between ranks is caught too. Verdict is identical-state-everywhere: any divergence clears KV on every rank; symmetric absence is kept because it is consistent (Bagel's empty unconditional branch). Local e2e rerun of test_bagel[tp_vae_patch_parallel_2] passed and the new path fired on a real divergence (one rank primary+metadata, the other only cfg_kv_request_ids) — both ranks fell back together. Co-Authored-By: Claude Fable 5 Signed-off-by: tzhouam --- .../test_kv_recv_tp_consensus.py | 130 +++++++++++++----- .../omni_connectors/test_tp_rank_aware.py | 19 ++- .../omni_connectors/kv_transfer_manager.py | 88 +++++++----- 3 files changed, 163 insertions(+), 74 deletions(-) diff --git a/tests/distributed/omni_connectors/test_kv_recv_tp_consensus.py b/tests/distributed/omni_connectors/test_kv_recv_tp_consensus.py index f890ba059..a7e4b8562 100644 --- a/tests/distributed/omni_connectors/test_kv_recv_tp_consensus.py +++ b/tests/distributed/omni_connectors/test_kv_recv_tp_consensus.py @@ -4,8 +4,11 @@ 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. +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 @@ -26,16 +29,16 @@ 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): + 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_flags = list(peer_flags or []) + self._peer_signatures = list(peer_signatures or []) self._verdict = verdict self.sent = [] self.broadcasts = [] def recv_object(self, src): - return self._peer_flags.pop(0) + return self._peer_signatures.pop(0) def send_object(self, obj, dst): self.sent.append((obj, dst)) @@ -44,8 +47,13 @@ class _FakeGroup: 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, world=None): + +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, @@ -56,6 +64,8 @@ def _topo(role=ReceiveRole.LOCAL, tp_active=True, world=None): sp_rank=0, sp_group=None, world=world, + tp_group=tp_group, + tp_size=tp_size, ) @@ -65,56 +75,77 @@ def _manager(topo, need_recv_cache=True): return mgr -def _req_with_kv(rid="r1"): +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=SimpleNamespace(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_flags=[False]) + group = _FakeGroup(0, 2, peer_signatures=[(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), + _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 group.sent == [] and group.broadcasts == [] + assert not group.touched 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 + assert _manager(_topo(tp_group=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) + 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(world=group2)) + 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 group2.sent == [] and group2.broadcasts == [] + 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 flags, decides, broadcasts +# Rank 0: collects signatures, decides, broadcasts # --------------------------------------------------------------------------- # -def test_rank0_drops_kv_when_peer_missed(): - group = _FakeGroup(0, 2, peer_flags=[False]) - mgr = _manager(_topo(world=group)) +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 @@ -124,10 +155,10 @@ def test_rank0_drops_kv_when_peer_missed(): 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)) +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] @@ -135,36 +166,63 @@ def test_rank0_keeps_kv_when_all_received(): 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(): - 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 + 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 flag, honors verdict +# Rank > 0: reports signature, 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() + 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 == [(True, 0)] + 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(): - group = _FakeGroup(1, 2, verdict=True) - mgr = _manager(_topo(world=group)) 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 == [(True, 0)] + assert group.sent == [(_signature_of(req), 0)] assert req.past_key_values is not None @@ -174,7 +232,7 @@ def test_follower_rank_keeps_kv_on_positive_verdict(): def _wired_manager(monkeypatch, consensus_calls): - mgr = _manager(_topo(world=_FakeGroup(0, 2))) + mgr = _manager(_topo(tp_group=_FakeGroup(0, 2))) def fake_consensus(req, received): consensus_calls.append(received) diff --git a/tests/distributed/omni_connectors/test_tp_rank_aware.py b/tests/distributed/omni_connectors/test_tp_rank_aware.py index 178069d65..3fdc4820a 100644 --- a/tests/distributed/omni_connectors/test_tp_rank_aware.py +++ b/tests/distributed/omni_connectors/test_tp_rank_aware.py @@ -681,9 +681,11 @@ 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()) - # 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) + 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 ( @@ -706,14 +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 receive flag for the all-or-nothing - # verdict instead of silently diverging on a miss (#5627). - assert world_group.send_calls == [(0, True)] + # 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 760ab37b8..63fe1a305 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: @@ -636,6 +642,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, @@ -646,6 +662,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( @@ -1762,25 +1780,33 @@ class OmniKVTransferManager: 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). + """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). Exchange per-rank receive flags; - if any rank missed, every rank drops its KV so all take the same - fallback path together. + 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 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. + 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.world - if not (pt.is_local and pt.tp_active) or group is None: - return received - if getattr(group, "world_size", 1) <= 1: + 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 @@ -1788,30 +1814,28 @@ class OmniKVTransferManager: 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: - 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) + 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(bool(received), dst=0) - all_received = bool(group.broadcast_object(None, src=0)) + group.send_object(signature, dst=0) + consistent = 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 + 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,