review: scope KV receive consensus to the TP group, compare full KV state
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>
This commit is contained in:
tzhouamandClaude Fable 5 committed
1 parent
0eda607077
commit
bac3eafc18
3 files changed
+163
-74
No files matched your search
@@ -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)
|
||||
|
||||
@@ -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 ────────────────────────────────────────────
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in new issue
Block a user