# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for rank-aware KV transfer (TP > 1) and heterogeneous TP support. Covers: - _build_rank_aware_send_keys / _build_rank_aware_recv_keys - _get_kv_source_ranks / _get_kv_target_ranks / get_kv_connector_key - update_sender_info storing base host/port - receive path constructing per-rank metadata for connector.get() - Mooncake connector _query_metadata_at and partial-metadata get() path """ from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest import torch from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( KVCacheTransferData, OmniKVCacheConfig, OmniKVTransferManager, ) from vllm_omni.distributed.omni_connectors.utils.initialization import ( KV_RANK_PORT_STRIDE, ) from vllm_omni.distributed.omni_connectors.utils.kv_utils import ( KVTPTopology, build_rank_aware_recv_keys, build_rank_aware_send_keys, get_kv_connector_key, get_kv_source_ranks, get_kv_target_ranks, merge_received_rank_shards, slice_received_rank_shard, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def _make_manager( from_tp: int = 1, to_tp: int = 1, local_rank: int = 0, from_stage: str = "stage0", to_stage: str = "stage1", stage_id: str = "stage1", need_recv: bool = True, need_send: bool = False, recv_timeout: float = 0.3, ) -> OmniKVTransferManager: """Build a manager with TP params injected, bypassing torch.distributed.""" config = OmniKVCacheConfig( connector_config={"type": "mock"}, from_stage=from_stage, to_stage=to_stage, stage_id=stage_id, need_recv_cache=need_recv, need_send_cache=need_send, recv_timeout=recv_timeout, from_tp=from_tp, to_tp=to_tp, ) with ( patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=local_rank), patch( "vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=max(from_tp, to_tp), ), ): mgr = OmniKVTransferManager(config) return mgr def _make_payload(head_values: list[float], request_id: str = "req-1") -> dict: head_tensor = torch.tensor(head_values, dtype=torch.float32).view(1, len(head_values), 1).repeat(2, 1, 1) return { "request_id": request_id, "layer_blocks": { "key_cache": [head_tensor.clone()], "value_cache": [(head_tensor + 100).clone()], }, "block_ids": [0], "metadata": {"seq_len": 2}, } def _make_transfer_data(head_values: list[float], request_id: str = "req-1") -> KVCacheTransferData: payload = _make_payload(head_values, request_id=request_id) return KVCacheTransferData( request_id=request_id, layer_blocks=payload["layer_blocks"], block_ids=payload["block_ids"], metadata=payload["metadata"], ) # ── Key format helper ──────────────────────────────────────────────── class TestConnectorKeyFormat: def test_key_format_matches_pr2677(self): key = get_kv_connector_key("req-1", "stage0", 0, 1, 2) assert key == "req-1_stage0_0_1_2" def test_key_fields_are_positional(self): key = get_kv_connector_key("r", "s", 5, 3, 7) parts = key.split("_") assert parts == ["r", "s", "5", "3", "7"] # ── Source / target rank mapping ───────────────────────────────────── class TestRankMapping: """Verify get_kv_target_ranks and get_kv_source_ranks for various TP configs.""" @pytest.mark.parametrize( "tp_size,local_rank", [ (2, 0), (2, 1), (4, 3), ], ) def test_homogeneous_tp(self, tp_size, local_rank): """Homogeneous TP: each rank sends/receives to/from itself.""" topo = KVTPTopology(source_tp_size=tp_size, target_tp_size=tp_size, local_rank=local_rank) assert get_kv_target_ranks(topo) == [local_rank] assert get_kv_source_ranks(topo) == [local_rank] @pytest.mark.parametrize( "local_rank,expected_source_ranks", [ (0, [0, 1]), # rank 0 receives from sender ranks 0 and 1 (1, [2, 3]), # rank 1 receives from sender ranks 2 and 3 ], ) def test_sender_gt_receiver_tp4_to_tp2(self, local_rank, expected_source_ranks): """Receiver with source_tp=4, target_tp=2.""" topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=local_rank) assert get_kv_source_ranks(topo) == expected_source_ranks @pytest.mark.parametrize( "local_rank,expected_target_ranks", [ (0, [0, 1]), # sender rank 0 sends to receiver ranks 0 and 1 (1, [2, 3]), # sender rank 1 sends to receiver ranks 2 and 3 ], ) def test_sender_lt_receiver_tp2_to_tp4(self, local_rank, expected_target_ranks): """Sender with source_tp=2, target_tp=4.""" topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=local_rank) assert get_kv_target_ranks(topo) == expected_target_ranks def test_receiver_lt_sender_source_ranks(self): """Receiver rank 0 with tp2_to_tp4 should source from rank 0 only.""" topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0) assert get_kv_source_ranks(topo) == [0] def test_invalid_topology_raises(self): topo = KVTPTopology(source_tp_size=3, target_tp_size=2, local_rank=0) with pytest.raises(ValueError, match="divisible"): get_kv_source_ranks(topo) # ── _build_rank_aware_recv_keys ────────────────────────────────────── class TestBuildRankAwareRecvKeys: """Verify build_rank_aware_recv_keys returns (key, from_rank) tuples.""" def test_tp1_returns_legacy_key_with_none_rank(self): topo = KVTPTopology(source_tp_size=1, target_tp_size=1, local_rank=0) pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) assert len(pairs) == 1 key, rank = pairs[0] assert key == "omni_stage0_to_stage1_kv_cache_req-1" assert rank is None @pytest.mark.parametrize( "local_rank,expected_key,expected_rank", [ (0, "req-1_stage0_0_0_0", 0), (1, "req-1_stage0_0_1_1", 1), ], ) def test_homogeneous_tp2(self, local_rank, expected_key, expected_rank): """Homogeneous TP=2: each rank gets one key.""" topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=local_rank) pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) assert len(pairs) == 1 key, rank = pairs[0] assert key == expected_key assert rank == expected_rank def test_heterogeneous_tp4_to_tp2_rank0_gets_two_keys(self): """Receiver rank 0 with source_tp=4, target_tp=2 should get 2 keys.""" topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=0) pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) assert len(pairs) == 2 keys = [k for k, _ in pairs] ranks = [r for _, r in pairs] assert keys == ["req-1_stage0_0_0_0", "req-1_stage0_0_1_0"] assert ranks == [0, 1] def test_heterogeneous_tp4_to_tp2_rank1_gets_two_keys(self): topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1) pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) assert len(pairs) == 2 ranks = [r for _, r in pairs] assert ranks == [2, 3] def test_heterogeneous_tp2_to_tp4_rank2_gets_one_key(self): """Receiver rank 2 with source_tp=2, target_tp=4 should get 1 key from sender rank 1.""" topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=2) pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) assert len(pairs) == 1 key, rank = pairs[0] assert rank == 1 assert key == "req-1_stage0_0_1_2" # ── _build_rank_aware_send_keys ────────────────────────────────────── class TestBuildRankAwareSendKeys: def test_tp1_returns_legacy_key(self): topo = KVTPTopology(source_tp_size=1, target_tp_size=1, local_rank=0) keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo) assert keys == ["omni_stage0_to_stage1_kv_cache_req-1"] def test_homogeneous_tp2_rank0(self): topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0) keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo) assert keys == ["req-1_stage0_0_0_0"] def test_sender_lt_receiver_tp2_to_tp4_rank0_sends_two_keys(self): topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0) keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo) assert len(keys) == 2 assert keys == ["req-1_stage0_0_0_0", "req-1_stage0_0_0_1"] # ── update_sender_info stores base host/port ───────────────────────── class TestUpdateSenderInfoBase: def test_stores_base_host_and_port(self): mgr = _make_manager(from_tp=2, to_tp=2, local_rank=0) mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) assert mgr._sender_base_host == "10.0.0.1" assert mgr._sender_base_zmq_port == 50151 def test_rank1_adjusts_default_port_but_preserves_base(self): mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1) mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) assert mgr._sender_base_host == "10.0.0.1" assert mgr._sender_base_zmq_port == 50151 expected_adjusted = 50151 + 1 * KV_RANK_PORT_STRIDE assert mgr.config.connector_config["sender_zmq_port"] == expected_adjusted def test_nested_sender_info_resolves_correctly(self): """Nested sender_info keyed by integer stage id should resolve using recv_stages (engine_input_source → recv_from).""" config = OmniKVCacheConfig( connector_config={"type": "mock"}, stage_id=2, engine_input_source=[1], need_recv_cache=True, from_tp=2, to_tp=2, ) with ( patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0), patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=2), ): mgr = OmniKVTransferManager(config) mgr.update_sender_info( { 0: {"host": "10.0.0.1", "zmq_port": 50151}, 1: {"host": "10.0.0.2", "zmq_port": 50152}, } ) assert mgr._sender_base_host == "10.0.0.2" assert mgr._sender_base_zmq_port == 50152 # ── receive path constructs per-rank metadata ──────────────────────── class TestReceiveConstructsMetadata: """Verify that receive_kv_cache_for_request passes metadata with correct (host, port) to connector.get() for heterogeneous TP.""" def test_tp1_no_metadata_passed(self): """TP=1: connector.get() should be called WITHOUT metadata.""" mgr = _make_manager(from_tp=1, to_tp=1, local_rank=0, recv_timeout=0.05) mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) calls = [] class _Connector: def get(self, from_stage, to_stage, get_key, metadata=None): calls.append({"key": get_key, "metadata": metadata}) return None mgr._connector = _Connector() mgr.receive_kv_cache_for_request("req-1") assert len(calls) > 0 assert calls[0]["metadata"] is None @pytest.mark.parametrize( "local_rank,expected_port_offset", [ (0, 0), (1, 1), ], ) def test_homogeneous_tp2_passes_metadata(self, local_rank, expected_port_offset): """TP=2: metadata should point to sender rank's port with correct offset.""" mgr = _make_manager(from_tp=2, to_tp=2, local_rank=local_rank, recv_timeout=0.05) mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) calls = [] class _Connector: def get(self, from_stage, to_stage, get_key, metadata=None): calls.append({"key": get_key, "metadata": metadata}) return None mgr._connector = _Connector() mgr.receive_kv_cache_for_request("req-1") assert len(calls) > 0 meta = calls[0]["metadata"] assert meta is not None assert meta["source_host"] == "10.0.0.1" assert meta["source_port"] == 50151 + expected_port_offset * KV_RANK_PORT_STRIDE def test_heterogeneous_tp4_to_tp2_rank0_multiple_metadata(self): """Receiver rank 0 with source_tp=4, target_tp=2 should call get() with two different metadata entries for sender ranks 0 and 1.""" mgr = _make_manager(from_tp=4, to_tp=2, local_rank=0, recv_timeout=0.05) mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) calls = [] class _Connector: def get(self, from_stage, to_stage, get_key, metadata=None): calls.append({"key": get_key, "metadata": metadata}) return None mgr._connector = _Connector() mgr.receive_kv_cache_for_request("req-1") seen_ports = set() for c in calls: if c["metadata"]: seen_ports.add(c["metadata"]["source_port"]) expected_ports = { 50151 + 0 * KV_RANK_PORT_STRIDE, 50151 + 1 * KV_RANK_PORT_STRIDE, } assert expected_ports.issubset(seen_ports) # ── Mooncake connector _query_metadata_at ──────────────────────────── class TestMooncakeQueryMetadataAt: """Test the connector's _query_metadata_at method and partial-metadata path in get() without requiring real RDMA/Mooncake.""" def test_query_metadata_at_returns_full_metadata(self): """Mock the ZMQ interaction to verify _query_metadata_at returns complete metadata including data_size.""" try: from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import ( MooncakeTransferEngineConnector, QueryResponse, ) except ImportError: pytest.skip("Mooncake not available") import msgspec connector = MagicMock(spec=MooncakeTransferEngineConnector) connector._get_req_socket = MagicMock() mock_socket = MagicMock() resp = QueryResponse(request_id="test_key@s0_s1", data_size=4096, is_fast_path=True) mock_socket.recv.return_value = msgspec.msgpack.encode(resp) connector._get_req_socket.return_value = mock_socket result = MooncakeTransferEngineConnector._query_metadata_at( connector, "test_key@s0_s1", "10.0.0.1", 50151, ) assert result is not None assert result["source_host"] == "10.0.0.1" assert result["source_port"] == 50151 assert result["data_size"] == 4096 assert result["is_fast_path"] is True def test_query_metadata_at_returns_none_on_not_found(self): try: from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import ( INFO_NOT_FOUND, MooncakeTransferEngineConnector, ) except ImportError: pytest.skip("Mooncake not available") connector = MagicMock(spec=MooncakeTransferEngineConnector) mock_socket = MagicMock() mock_socket.recv.return_value = INFO_NOT_FOUND connector._get_req_socket.return_value = mock_socket result = MooncakeTransferEngineConnector._query_metadata_at( connector, "test_key@s0_s1", "10.0.0.1", 50151, ) assert result is None # ── Merge / slice hooks ────────────────────────────────────────────── class TestMergeSliceHooks: def test_single_shard_passes_through(self): payload = {"layer_blocks": {"key_cache": [1]}} assert merge_received_rank_shards([payload]) == payload def test_default_merger_concats_head_dim(self): p0 = _make_payload([0.0]) p1 = _make_payload([1.0]) result = merge_received_rank_shards([p0, p1]) key_cache = result["layer_blocks"]["key_cache"][0] value_cache = result["layer_blocks"]["value_cache"][0] assert key_cache.shape == (2, 2, 1) assert value_cache.shape == (2, 2, 1) assert torch.equal(key_cache[:, :, 0], torch.tensor([[0.0, 1.0], [0.0, 1.0]])) assert torch.equal(value_cache[:, :, 0], torch.tensor([[100.0, 101.0], [100.0, 101.0]])) def test_custom_merger_hook_called(self): merged = {"merged": True} assert merge_received_rank_shards([{}, {}], merger=lambda payloads: merged) == merged def test_slicer_hook_called(self): topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0) sliced = {"sliced": True} assert slice_received_rank_shard({"full": True}, topo, slicer=lambda payload: sliced) == sliced def test_default_slicer_extracts_rank_local_heads(self): topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1) payload = _make_payload([0.0, 1.0]) result = slice_received_rank_shard(payload, topo) key_cache = result["layer_blocks"]["key_cache"][0] value_cache = result["layer_blocks"]["value_cache"][0] assert key_cache.shape == (2, 1, 1) assert value_cache.shape == (2, 1, 1) assert torch.equal(key_cache[:, :, 0], torch.tensor([[1.0], [1.0]])) assert torch.equal(value_cache[:, :, 0], torch.tensor([[101.0], [101.0]])) def test_presliced_payload_is_not_sliced_twice(self): topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1) payload = _make_payload([1.0]) payload["metadata"]["tp_head_slice"] = {"applied": True, "target_rank": 1} result = slice_received_rank_shard(payload, topo) assert result is payload def test_round_trip_merge_from_tp4_to_tp2(self): topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1) source_ranks = get_kv_source_ranks(topo) payloads = [_make_payload([float(rank)]) for rank in source_ranks] result = merge_received_rank_shards(payloads) key_cache = result["layer_blocks"]["key_cache"][0] assert torch.equal(key_cache[:, :, 0], torch.tensor([[2.0, 3.0], [2.0, 3.0]])) def test_round_trip_slice_from_tp2_to_tp4(self): topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=3) payload = _make_payload([2.0, 3.0]) result = slice_received_rank_shard(payload, topo) key_cache = result["layer_blocks"]["key_cache"][0] assert torch.equal(key_cache[:, :, 0], torch.tensor([[3.0], [3.0]])) class TestSenderSideSlicing: def test_transfer_slices_before_sending_to_multiple_targets(self): mgr = _make_manager( from_tp=2, to_tp=4, local_rank=0, need_send=True, need_recv=False, ) sent_payloads = [] class _Connector: supports_raw_data = False def put(self, from_stage, to_stage, put_key, data): sent_payloads.append((put_key, KVCacheTransferData.from_bytes(data))) return True, len(data), {} mgr._connector = _Connector() mgr._transfer_kv_cache(_make_transfer_data([0.0, 1.0]), "req-1") assert [key for key, _ in sent_payloads] == ["req-1_stage0_0_0_0", "req-1_stage0_0_0_1"] assert sent_payloads[0][1]["layer_blocks"]["key_cache"][0].shape == (2, 1, 1) assert sent_payloads[1][1]["layer_blocks"]["key_cache"][0].shape == (2, 1, 1) assert torch.equal( sent_payloads[0][1]["layer_blocks"]["key_cache"][0][:, :, 0], torch.tensor([[0.0], [0.0]]), ) assert torch.equal( sent_payloads[1][1]["layer_blocks"]["key_cache"][0][:, :, 0], torch.tensor([[1.0], [1.0]]), ) assert sent_payloads[0][1]["metadata"]["tp_head_slice"]["target_rank"] == 0 assert sent_payloads[1][1]["metadata"]["tp_head_slice"]["target_rank"] == 1 class _MockBroadcastGroup: def __init__(self, world_size: int, rank_in_group: int, broadcast_value=None, recv_value=None): self.world_size = world_size self.rank_in_group = rank_in_group self.broadcast_value = broadcast_value self.recv_value = recv_value self.broadcast_calls = [] self.send_calls = [] self.recv_calls = [] self.shm_broadcaster = None def broadcast_object(self, obj=None, src: int = 0): self.broadcast_calls.append((obj, src)) return self.broadcast_value if self.broadcast_value is not None else obj def send_object(self, obj, dst: int): self.send_calls.append((dst, obj)) def recv_object(self, src: int): self.recv_calls.append(src) return self.recv_value class TestDistributedReceive: def test_tp_cfg_leader_receives_then_sends_branch_local_payloads(self): mgr = _make_manager(from_tp=2, to_tp=4, local_rank=0) req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace()) world_group = _MockBroadcastGroup(world_size=4, rank_in_group=2) cfg_group = _MockBroadcastGroup(world_size=3, rank_in_group=0) def _receive(req_obj, cfg_func, target_device): req_obj.past_key_values = SimpleNamespace(key_cache=[torch.tensor([1.0])]) req_obj.kv_metadata = {"source": "leader"} req_obj.sampling_params.past_key_values = req_obj.past_key_values req_obj.sampling_params.kv_metadata = req_obj.kv_metadata req_obj.sampling_params.cfg_text_past_key_values = SimpleNamespace(key_cache=[torch.tensor([2.0])]) req_obj.sampling_params.cfg_text_kv_metadata = {"source": "cfg_text"} req_obj.sampling_params.cfg_img_past_key_values = SimpleNamespace(key_cache=[torch.tensor([3.0])]) req_obj.sampling_params.cfg_img_kv_metadata = {"source": "cfg_img"} return True mgr.receive_multi_kv_cache = MagicMock(side_effect=_receive) with ( patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group), patch( "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size", return_value=3, ), patch( "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank", return_value=0, ), patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group), ): assert mgr.receive_multi_kv_cache_distributed(req) is True mgr.receive_multi_kv_cache.assert_called_once() assert mgr.receive_multi_kv_cache.call_args.args[2] == torch.device("cpu") assert req.kv_metadata == {"source": "leader"} assert cfg_group.broadcast_calls == [] assert [dst for dst, _ in cfg_group.send_calls] == [1, 2] rank1_payload = cfg_group.send_calls[0][1] rank2_payload = cfg_group.send_calls[1][1] assert torch.equal(rank1_payload["past_key_values"].key_cache[0], torch.tensor([1.0])) assert torch.equal(rank2_payload["past_key_values"].key_cache[0], torch.tensor([1.0])) assert rank1_payload["sp.cfg_active_branch"] == "cfg_text" assert rank2_payload["sp.cfg_active_branch"] == "cfg_img" assert rank1_payload["sp.cfg_branch_roles"] == ["cfg_text", "cfg_img"] assert rank2_payload["sp.cfg_branch_roles"] == ["cfg_text", "cfg_img"] assert "sp.cfg_branch_past_key_values" in rank1_payload assert "sp.cfg_branch_past_key_values" in rank2_payload assert list(rank1_payload["sp.cfg_branch_past_key_values"].keys()) == ["cfg_text"] assert list(rank2_payload["sp.cfg_branch_past_key_values"].keys()) == ["cfg_img"] assert "sp.cfg_text_past_key_values" in rank1_payload assert "sp.cfg_img_past_key_values" not in rank1_payload assert "sp.cfg_img_past_key_values" in rank2_payload assert "sp.cfg_text_past_key_values" not in rank2_payload def test_tp_cfg_follower_receives_local_payload_without_receiving(self): mgr = _make_manager(from_tp=2, to_tp=4, local_rank=1) req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace()) world_group = _MockBroadcastGroup(world_size=4, rank_in_group=3) cfg_payload = { "past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]), "kv_metadata": {"source": "main"}, "sp.past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]), "sp.kv_metadata": {"source": "main"}, "sp.cfg_active_branch": "cfg_text", "sp.cfg_branch_roles": ["cfg_text", "cfg_img"], "sp.cfg_branch_past_key_values": { "cfg_text": SimpleNamespace(key_cache=[torch.tensor([2.0])]), }, "sp.cfg_branch_kv_metadata": {"cfg_text": {"source": "cfg-text"}}, "sp.cfg_text_past_key_values": SimpleNamespace(key_cache=[torch.tensor([2.0])]), } cfg_group = _MockBroadcastGroup(world_size=2, rank_in_group=1, recv_value=cfg_payload) mgr.receive_multi_kv_cache = MagicMock(return_value=True) with ( patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group), patch( "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size", return_value=2, ), patch( "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank", return_value=1, ), patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group), ): assert mgr.receive_multi_kv_cache_distributed(req) is True mgr.receive_multi_kv_cache.assert_not_called() assert req.kv_metadata == {"source": "main"} assert torch.equal(req.past_key_values.key_cache[0], torch.tensor([1.0])) assert torch.equal(req.sampling_params.past_key_values.key_cache[0], torch.tensor([1.0])) assert req.sampling_params.cfg_active_branch == "cfg_text" assert req.sampling_params.cfg_branch_roles == ["cfg_text", "cfg_img"] assert torch.equal( req.sampling_params.cfg_branch_past_key_values["cfg_text"].key_cache[0], torch.tensor([2.0]), ) assert req.sampling_params.cfg_branch_kv_metadata == {"cfg_text": {"source": "cfg-text"}} assert torch.equal(req.sampling_params.cfg_text_past_key_values.key_cache[0], torch.tensor([2.0])) assert cfg_group.broadcast_calls == [] assert cfg_group.recv_calls == [0] 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) mgr.receive_multi_kv_cache = MagicMock(return_value=True) with patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_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")) # ── SP-only scenarios ──────────────────────────────────────────── @pytest.mark.parametrize( "sp_rank,world_rank_in_group,is_owner", [ (0, 0, True), # owner receives and broadcasts (1, 1, False), # follower receives via broadcast ], ) def test_sp_only_scenarios(self, sp_rank, world_rank_in_group, is_owner): """SP-only (sp_size=2, cfg_size=1): Test owner and follower behavior.""" mgr = _make_manager(from_tp=2, to_tp=4, local_rank=0) req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace()) world_group = _MockBroadcastGroup(world_size=4, rank_in_group=world_rank_in_group) sp_payload = { "past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]), "kv_metadata": {"source": "owner"}, "sp.past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]), "sp.kv_metadata": {"source": "owner"}, } sp_group = _MockBroadcastGroup(world_size=2, rank_in_group=sp_rank, broadcast_value=sp_payload) def _receive(req_obj, cfg_func, target_device): req_obj.past_key_values = SimpleNamespace(key_cache=[torch.tensor([1.0])]) req_obj.kv_metadata = {"source": "owner"} req_obj.sampling_params.past_key_values = req_obj.past_key_values req_obj.sampling_params.kv_metadata = req_obj.kv_metadata return True mgr.receive_multi_kv_cache = MagicMock(side_effect=_receive) with ( patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group), patch( "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size", return_value=1, ), patch( "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank", return_value=0, ), patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=None), patch( "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size", return_value=2, ), patch( "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank", return_value=sp_rank, ), patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group), ): assert mgr.receive_multi_kv_cache_distributed(req) is True # Verify behavior based on role if is_owner: mgr.receive_multi_kv_cache.assert_called_once() assert mgr.receive_multi_kv_cache.call_args.args[2] == torch.device("cpu") assert sp_group.broadcast_calls[0][1] == 0 else: mgr.receive_multi_kv_cache.assert_not_called() assert len(sp_group.broadcast_calls) == 1 assert torch.equal(req.past_key_values.key_cache[0], torch.tensor([1.0])) assert req.sampling_params.kv_metadata == {"source": "owner"} # ── CFG + SP mixed scenarios ───────────────────────────────────── @pytest.mark.parametrize( "cfg_rank,sp_rank,world_rank_in_group,role", [ (0, 0, 0, "owner"), # owner: receives, cfg scatter, sp broadcast (1, 0, 2, "cfg_follower_sp_leader"), # cfg follower + sp leader: cfg recv, sp broadcast (0, 1, 1, "sp_follower"), # sp follower: only sp broadcast ], ) def test_cfg_sp_mixed_scenarios(self, cfg_rank, sp_rank, world_rank_in_group, role): """CFG+SP (cfg_size=2, sp_size=2): Test various rank combinations.""" mgr = _make_manager(from_tp=2, to_tp=4, local_rank=0) req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace()) world_group = _MockBroadcastGroup(world_size=8, rank_in_group=world_rank_in_group) # Setup payloads based on role if role == "owner": owner_sp_payload = { "past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]), "kv_metadata": {"source": "owner"}, "sp.past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]), "sp.kv_metadata": {"source": "owner"}, "sp.cfg_branch_roles": ["cfg_text"], "sp.cfg_active_branch": None, } cfg_group = _MockBroadcastGroup(world_size=2, rank_in_group=cfg_rank) sp_group = _MockBroadcastGroup(world_size=2, rank_in_group=sp_rank, broadcast_value=owner_sp_payload) def _receive(req_obj, cfg_func, target_device): req_obj.past_key_values = SimpleNamespace(key_cache=[torch.tensor([1.0])]) req_obj.kv_metadata = {"source": "owner"} req_obj.sampling_params.past_key_values = req_obj.past_key_values req_obj.sampling_params.kv_metadata = req_obj.kv_metadata req_obj.sampling_params.cfg_text_past_key_values = SimpleNamespace(key_cache=[torch.tensor([2.0])]) req_obj.sampling_params.cfg_text_kv_metadata = {"source": "cfg_text"} return True mgr.receive_multi_kv_cache = MagicMock(side_effect=_receive) else: cfg_payload = { "past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]), "kv_metadata": {"source": "main"}, "sp.past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]), "sp.kv_metadata": {"source": "main"}, "sp.cfg_active_branch": "cfg_text" if role == "cfg_follower_sp_leader" else None, "sp.cfg_branch_roles": ["cfg_text"], "sp.cfg_branch_past_key_values": { "cfg_text": SimpleNamespace(key_cache=[torch.tensor([2.0])]), } if role == "cfg_follower_sp_leader" else {}, } if role == "cfg_follower_sp_leader": cfg_payload["sp.cfg_text_past_key_values"] = SimpleNamespace(key_cache=[torch.tensor([2.0])]) cfg_group = _MockBroadcastGroup(world_size=2, rank_in_group=cfg_rank, recv_value=cfg_payload) sp_group = _MockBroadcastGroup(world_size=2, rank_in_group=sp_rank, broadcast_value=cfg_payload) mgr.receive_multi_kv_cache = MagicMock(return_value=True) with ( patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group), patch( "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size", return_value=2, ), patch( "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank", return_value=cfg_rank, ), patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group), patch( "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size", return_value=2, ), patch( "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank", return_value=sp_rank, ), patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group), ): assert mgr.receive_multi_kv_cache_distributed(req) is True # Verify behavior based on role if role == "owner": mgr.receive_multi_kv_cache.assert_called_once() assert mgr.receive_multi_kv_cache.call_args.args[2] == torch.device("cpu") assert [dst for dst, _ in cfg_group.send_calls] == [1] assert cfg_group.send_calls[0][1]["sp.cfg_active_branch"] == "cfg_text" assert len(sp_group.broadcast_calls) == 1 assert sp_group.broadcast_calls[0][1] == 0 elif role == "cfg_follower_sp_leader": mgr.receive_multi_kv_cache.assert_not_called() assert cfg_group.recv_calls == [0] assert len(sp_group.broadcast_calls) == 1 assert sp_group.broadcast_calls[0][1] == 0 assert req.sampling_params.cfg_active_branch == "cfg_text" assert torch.equal( req.sampling_params.cfg_branch_past_key_values["cfg_text"].key_cache[0], torch.tensor([2.0]), ) else: # sp_follower mgr.receive_multi_kv_cache.assert_not_called() assert cfg_group.recv_calls == [] assert cfg_group.send_calls == [] assert len(sp_group.broadcast_calls) == 1 assert torch.equal(req.past_key_values.key_cache[0], torch.tensor([1.0])) assert req.kv_metadata == {"source": "main"} @pytest.mark.parametrize( "has_cfg,cfg_size", [ (False, 1), # SP-only: owner receive failure (True, 2), # CFG+SP: owner receive failure with deadlock prevention ], ) def test_owner_receive_failure_scenarios(self, has_cfg, cfg_size): """Test receive failure handling in SP-only and CFG+SP scenarios.""" mgr = _make_manager(from_tp=2, to_tp=4, local_rank=0) req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace()) world_group = _MockBroadcastGroup(world_size=4 if not has_cfg else 8, rank_in_group=0) cfg_group = _MockBroadcastGroup(world_size=2, rank_in_group=0) if has_cfg else None sp_group = _MockBroadcastGroup(world_size=2, rank_in_group=0, broadcast_value=None) mgr.receive_multi_kv_cache = MagicMock(return_value=False) with ( patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group), patch( "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size", return_value=cfg_size, ), patch( "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank", return_value=0, ), patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group), patch( "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size", return_value=2, ), patch( "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank", return_value=0, ), patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group), ): assert mgr.receive_multi_kv_cache_distributed(req) is False # Owner attempted to receive but failed mgr.receive_multi_kv_cache.assert_called_once() # SP broadcast still called but with None assert len(sp_group.broadcast_calls) == 1 assert sp_group.broadcast_calls[0][0] is None # CFG+SP: owner must send None to cfg followers to prevent deadlock if has_cfg: assert len(cfg_group.send_calls) == 1 assert cfg_group.send_calls[0] == (1, None) # ── TP auto-detect ─────────────────────────────────────────────────── class TestAutoDetectTP: def test_auto_detect_when_config_defaults(self): """When config from_tp/to_tp == 1 (default), manager should auto-detect.""" config = OmniKVCacheConfig( connector_config={"type": "mock"}, from_stage="s0", stage_id="s1", need_recv_cache=True, ) with ( patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0), patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=4), ): mgr = OmniKVTransferManager(config) assert mgr._tp_topo.source_tp_size == 4 assert mgr._tp_topo.target_tp_size == 4 def test_explicit_tp_overrides_auto_detect(self): config = OmniKVCacheConfig( connector_config={"type": "mock"}, from_stage="s0", stage_id="s1", need_recv_cache=True, from_tp=2, to_tp=4, ) with ( patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0), patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=8), ): mgr = OmniKVTransferManager(config) assert mgr._tp_topo.source_tp_size == 2 assert mgr._tp_topo.target_tp_size == 4