Merge branch 'main' into fix/tp-local-kv-recv-consensus

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 <noreply@anthropic.com>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

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]

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 ────────────────────────────────────────────
 

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)