From 4d816ff1ded1e35393d6175d8f0dbbe07d570add Mon Sep 17 00:00:00 2001 From: NATURE Date: Thu, 16 Apr 2026 16:25:13 +0800 Subject: [PATCH] [Feature] Bagel: Support tp+cfg parallel using mooncake transfer engine connector (#2705) Signed-off-by: natureofnature Co-authored-by: Hongsheng Liu --- .../omni_connectors/test_tp_rank_aware.py | 716 +++++++++++++++++ .../test_async_omni_engine_stage_init.py | 69 ++ tests/engine/test_single_stage_mode.py | 2 + .../distributed/group_coordinator.py | 5 +- .../diffusion/models/bagel/pipeline_bagel.py | 36 +- .../omni_connectors/kv_transfer_manager.py | 721 ++++++++++++------ .../omni_connectors/utils/kv_utils.py | 367 ++++++++- vllm_omni/engine/async_omni_engine.py | 16 +- vllm_omni/engine/stage_engine_core_client.py | 7 +- vllm_omni/engine/stage_init_utils.py | 116 ++- vllm_omni/entrypoints/openai/serving_chat.py | 53 +- vllm_omni/inputs/data.py | 4 + 12 files changed, 1846 insertions(+), 266 deletions(-) create mode 100644 tests/distributed/omni_connectors/test_tp_rank_aware.py diff --git a/tests/distributed/omni_connectors/test_tp_rank_aware.py b/tests/distributed/omni_connectors/test_tp_rank_aware.py new file mode 100644 index 000000000..d4793479a --- /dev/null +++ b/tests/distributed/omni_connectors/test_tp_rank_aware.py @@ -0,0 +1,716 @@ +# 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.""" + + def test_homogeneous_tp2_rank0(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0) + assert get_kv_target_ranks(topo) == [0] + assert get_kv_source_ranks(topo) == [0] + + def test_homogeneous_tp2_rank1(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=1) + assert get_kv_target_ranks(topo) == [1] + assert get_kv_source_ranks(topo) == [1] + + def test_homogeneous_tp4_rank3(self): + topo = KVTPTopology(source_tp_size=4, target_tp_size=4, local_rank=3) + assert get_kv_target_ranks(topo) == [3] + assert get_kv_source_ranks(topo) == [3] + + def test_sender_gt_receiver_tp4_to_tp2_rank0(self): + """Receiver rank 0 should receive from sender rank 0 and 1.""" + topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=0) + assert get_kv_source_ranks(topo) == [0, 1] + + def test_sender_gt_receiver_tp4_to_tp2_rank1(self): + """Receiver rank 1 should receive from sender rank 2 and 3.""" + topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1) + assert get_kv_source_ranks(topo) == [2, 3] + + def test_sender_lt_receiver_tp2_to_tp4_rank0(self): + """Sender rank 0 should send to receiver ranks 0 and 1.""" + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0) + assert get_kv_target_ranks(topo) == [0, 1] + + def test_sender_lt_receiver_tp2_to_tp4_rank1(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1) + assert get_kv_target_ranks(topo) == [2, 3] + + 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 + + def test_homogeneous_tp2_rank0(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 1 + key, rank = pairs[0] + assert key == "req-1_stage0_0_0_0" + assert rank == 0 + + def test_homogeneous_tp2_rank1(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=1) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 1 + key, rank = pairs[0] + assert key == "req-1_stage0_0_1_1" + assert rank == 1 + + 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 + + def test_homogeneous_tp2_rank0_passes_metadata(self): + """TP=2 rank 0: metadata should point to sender rank 0's port.""" + mgr = _make_manager(from_tp=2, 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") + + 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 + 0 * KV_RANK_PORT_STRIDE + + def test_homogeneous_tp2_rank1_passes_metadata_with_offset(self): + mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1, 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") + + meta = calls[0]["metadata"] + assert meta["source_port"] == 50151 + 1 * 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")) + + +# ── 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 diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 84b0cb0be..5c2a9edb7 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -183,6 +183,7 @@ def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake( engine.model = "dummy-model" engine.single_stage_mode = False engine._omni_master_server = None + engine.stage_configs = [] metadata = types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"}) fake_vllm_config = types.SimpleNamespace() @@ -238,6 +239,7 @@ def test_launch_llm_stage_releases_launch_lock_before_complete_stage_handshake(m engine.model = "dummy-model" engine.single_stage_mode = False engine._omni_master_server = None + engine.stage_configs = [] fake_vllm_config = types.SimpleNamespace() fake_addresses = types.SimpleNamespace() @@ -378,3 +380,70 @@ def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch): assert input_processor is not None assert isinstance(input_processor.input_preprocessor, DummyOmniInputPreprocessor) assert input_processor.input_preprocessor.renderer is input_processor.renderer + + +def test_inject_kv_stage_info_infers_sender_tp_topology(): + from vllm_omni.engine.stage_init_utils import inject_kv_stage_info + + stage0 = types.SimpleNamespace( + stage_id=0, + engine_args={ + "tensor_parallel_size": 4, + "omni_kv_config": { + "need_send_cache": True, + "omni_from_stage": "0", + "omni_to_stage": "1", + }, + }, + engine_input_source=[], + ) + stage1 = types.SimpleNamespace( + stage_id=1, + engine_args={ + "parallel_config": { + "tensor_parallel_size": 2, + "cfg_parallel_size": 1, + }, + "omni_kv_config": {"need_recv_cache": True}, + }, + engine_input_source=[0], + ) + + inject_kv_stage_info(stage0, 0, [stage0, stage1]) + + assert stage0.engine_args["omni_kv_config"]["stage_id"] == 0 + assert stage0.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2} + + +def test_inject_kv_stage_info_infers_receiver_tp_topology(): + from vllm_omni.engine.stage_init_utils import inject_kv_stage_info + + stage0 = types.SimpleNamespace( + stage_id=0, + engine_args={ + "tensor_parallel_size": 4, + "omni_kv_config": {"need_send_cache": True}, + }, + engine_input_source=[], + ) + stage1 = types.SimpleNamespace( + stage_id=1, + engine_args={ + "parallel_config": { + "tensor_parallel_size": 2, + "cfg_parallel_size": 1, + }, + "omni_kv_config": { + "need_recv_cache": True, + "omni_from_stage": "0", + "omni_to_stage": "1", + }, + }, + engine_input_source=[0], + ) + + inject_kv_stage_info(stage1, 1, [stage0, stage1]) + + assert stage1.engine_args["omni_kv_config"]["stage_id"] == 1 + assert stage1.engine_args["omni_kv_config"]["engine_input_source"] == [0] + assert stage1.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2} diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 608e92ac4..28ccccaa2 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -1555,6 +1555,7 @@ class TestLaunchLlmStageSingleStageMode: engine.single_stage_mode = True engine._single_stage_id_filter = 0 engine._llm_stage_launch_lock = threading.Lock() + engine.stage_configs = [] mock_oms = mocker.Mock(spec=OmniMasterServer) mock_oms.address = "127.0.0.1" mock_oms.port = 25000 @@ -1629,6 +1630,7 @@ class TestLaunchLlmStageSingleStageMode: engine.single_stage_mode = False engine._omni_master_server = None engine._llm_stage_launch_lock = threading.Lock() + engine.stage_configs = [] fake_vllm_config = mocker.Mock() fake_executor_cls = mocker.Mock() diff --git a/vllm_omni/diffusion/distributed/group_coordinator.py b/vllm_omni/diffusion/distributed/group_coordinator.py index 8ab38f2a6..5294e6c9e 100644 --- a/vllm_omni/diffusion/distributed/group_coordinator.py +++ b/vllm_omni/diffusion/distributed/group_coordinator.py @@ -104,6 +104,7 @@ class GroupCoordinator: self.local_rank = local_rank self.device_group = None self.cpu_group = None + self.shm_broadcaster = None for ranks in group_ranks: device_group = torch.distributed.new_group(ranks, backend=torch_distributed_backend) @@ -316,7 +317,7 @@ class GroupCoordinator: assert dst < self.world_size, f"Invalid dst rank ({dst})" - assert dst != self.rank, "Invalid destination rank. Destination rank is the same as the current rank." + assert dst != self.rank_in_group, "Invalid destination rank. Destination rank is the same as the current rank." # Serialize object to tensor and get the size as well object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8) @@ -338,7 +339,7 @@ class GroupCoordinator: assert src < self.world_size, f"Invalid src rank ({src})" - assert src != self.rank, "Invalid source rank. Source rank is the same as the current rank." + assert src != self.rank_in_group, "Invalid source rank. Source rank is the same as the current rank." size_tensor = torch.empty(1, dtype=torch.long, device="cpu") diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 72e53e7f4..a3d2259e6 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -365,28 +365,52 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin): if req.sampling_params.kv_metadata and "image_shape" in req.sampling_params.kv_metadata: image_shape = tuple(req.sampling_params.kv_metadata["image_shape"]) - cfg_text_kv = getattr(req.sampling_params, "cfg_text_past_key_values", None) + branch_kvs = getattr(req.sampling_params, "cfg_branch_past_key_values", None) or {} + branch_metadata = getattr(req.sampling_params, "cfg_branch_kv_metadata", None) or {} + active_branch = getattr(req.sampling_params, "cfg_active_branch", None) + branch_roles = getattr(req.sampling_params, "cfg_branch_roles", None) or list(branch_kvs.keys()) + + cfg_text_kv = getattr(req.sampling_params, "cfg_text_past_key_values", None) or branch_kvs.get("cfg_text") + cfg_text_metadata = getattr(req.sampling_params, "cfg_text_kv_metadata", None) or branch_metadata.get( + "cfg_text" + ) + cfg_img_kv = getattr(req.sampling_params, "cfg_img_past_key_values", None) or branch_kvs.get("cfg_img") + cfg_img_metadata = getattr(req.sampling_params, "cfg_img_kv_metadata", None) or branch_metadata.get( + "cfg_img" + ) + + cfg_parallel_contract = ( + active_branch is not None or bool(branch_roles) or cfg_text_kv is not None or cfg_img_kv is not None + ) + if cfg_parallel_contract: + logger.info( + "CFG enabled with injected branch KV context roles=%s active=%s", + branch_roles, + active_branch, + ) + if cfg_text_kv is not None: - logger.info("CFG enabled with multi-KV: using injected cfg_text KV Cache") cfg_text_seq_len = cfg_text_kv.key_cache[0].shape[0] cfg_text_context["past_key_values"] = cfg_text_kv cfg_text_context["kv_lens"] = [cfg_text_seq_len] - cfg_text_metadata = getattr(req.sampling_params, "cfg_text_kv_metadata", None) if cfg_text_metadata and "ropes" in cfg_text_metadata: cfg_text_context["ropes"] = cfg_text_metadata["ropes"] else: cfg_text_context["ropes"] = [cfg_text_seq_len] - cfg_img_kv = getattr(req.sampling_params, "cfg_img_past_key_values", None) or injected_kv + if cfg_img_kv is None and cfg_text_kv is not None: + cfg_img_kv = injected_kv + + if cfg_img_kv is not None: cfg_img_seq_len = cfg_img_kv.key_cache[0].shape[0] cfg_img_context["past_key_values"] = cfg_img_kv cfg_img_context["kv_lens"] = [cfg_img_seq_len] - cfg_img_metadata = getattr(req.sampling_params, "cfg_img_kv_metadata", None) if cfg_img_metadata and "ropes" in cfg_img_metadata: cfg_img_context["ropes"] = cfg_img_metadata["ropes"] else: cfg_img_context["ropes"] = [cfg_img_seq_len] - else: + + if not cfg_parallel_contract: logger.warning("CFG is disabled: only single KV cache available") gen_params = BagelGenParams( num_timesteps=gen_params.num_timesteps, diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 1958c9d40..ad008c397 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -14,8 +14,20 @@ from vllm.logger import init_logger from .factory import OmniConnectorFactory from .utils.config import ConnectorSpec -from .utils.initialization import KV_TRANSFER_PORT_OFFSET -from .utils.kv_utils import normalize_layer_kv +from .utils.initialization import KV_RANK_PORT_STRIDE +from .utils.kv_utils import ( + KVTPTopology, + build_rank_aware_recv_keys, + build_rank_aware_send_keys, + get_kv_target_ranks, + get_local_tp_rank, + get_tp_world_size, + kv_zmq_port, + merge_received_rank_shards, + normalize_layer_kv, + slice_layer_blocks, + slice_received_rank_shard, +) logger = init_logger(__name__) @@ -57,6 +69,8 @@ class OmniKVCacheConfig: need_recv_cache: bool = False need_send_cache: bool = False recv_timeout: float = 30.0 + from_tp: int = 1 + to_tp: int = 1 @dataclass @@ -72,82 +86,44 @@ class KVCacheTransferData: """Convert to dictionary for serialization.""" return asdict(self) - def to_bytes(self) -> bytes: - """Convert to compact binary format for fast transfer.""" + def _build_tensors_desc(self, *, cpu: bool) -> tuple[list[dict[str, Any]], list, int, torch.device | None]: + """Iterate layer blocks and build tensor descriptors + data chunks. + + Returns ``(tensors_desc, chunks, total_bytes, device)``. + *chunks* contains ``bytes`` when *cpu* is True, flat uint8 GPU tensors otherwise. + """ tensors_desc: list[dict[str, Any]] = [] - tensor_bufs: list[bytes] = [] - data_offset = 0 - - for cache_name in ("key_cache", "value_cache"): - cache_list = self.layer_blocks.get(cache_name, []) - for layer_idx, tensor in enumerate(cache_list): - if tensor is None: - tensors_desc.append({"n": f"{cache_name}_{layer_idx}", "x": True}) - continue - - t = tensor.detach().cpu().contiguous() - dtype_str = str(t.dtype).removeprefix("torch.") - raw = t.view(torch.uint8).numpy().tobytes() - tensors_desc.append( - { - "n": f"{cache_name}_{layer_idx}", - "i": layer_idx, - "d": dtype_str, - "s": list(t.shape), - "o": data_offset, - "b": len(raw), - } - ) - tensor_bufs.append(raw) - data_offset += len(raw) - - header = json.dumps( - { - "rid": self.request_id, - "bids": self.block_ids, - "meta": self.metadata, - "td": tensors_desc, - "nl": len(self.layer_blocks.get("key_cache", [])), - }, - separators=(",", ":"), - ).encode("utf-8") - return b"".join([struct.pack(">I", len(header)), header] + tensor_bufs) - - def to_gpu_tensor(self) -> torch.Tensor: - """Convert to a packed GPU tensor for raw-data connectors.""" - tensors_desc: list[dict[str, Any]] = [] - gpu_tensors: list[torch.Tensor] = [] + chunks: list = [] data_offset = 0 device = None for cache_name in ("key_cache", "value_cache"): - cache_list = self.layer_blocks.get(cache_name, []) - for layer_idx, tensor in enumerate(cache_list): + for layer_idx, tensor in enumerate(self.layer_blocks.get(cache_name, [])): if tensor is None: tensors_desc.append({"n": f"{cache_name}_{layer_idx}", "x": True}) continue - t = tensor.detach().contiguous() - if device is None and t.is_cuda: + if cpu: + t = t.cpu() + elif device is None and t.is_cuda: device = t.device - dtype_str = str(t.dtype).removeprefix("torch.") nbytes = t.numel() * t.element_size() tensors_desc.append( { "n": f"{cache_name}_{layer_idx}", "i": layer_idx, - "d": dtype_str, + "d": str(t.dtype).removeprefix("torch."), "s": list(t.shape), "o": data_offset, "b": nbytes, } ) - gpu_tensors.append(t.view(torch.uint8).flatten()) + chunks.append(t.view(torch.uint8).numpy().tobytes() if cpu else t.view(torch.uint8).flatten()) data_offset += nbytes - if device is None: - raise RuntimeError("No CUDA tensors found, use to_bytes() instead") + return tensors_desc, chunks, data_offset, device + def _build_header_bytes(self, tensors_desc: list[dict[str, Any]]) -> bytes: header = json.dumps( { "rid": self.request_id, @@ -158,19 +134,26 @@ class KVCacheTransferData: }, separators=(",", ":"), ).encode("utf-8") + return struct.pack(">I", len(header)) + header - header_prefix = struct.pack(">I", len(header)) + header - total_size = len(header_prefix) + data_offset - output = torch.empty(total_size, dtype=torch.uint8, device=device) - header_tensor = torch.frombuffer(bytearray(header_prefix), dtype=torch.uint8) - output[: len(header_prefix)].copy_(header_tensor) + def to_bytes(self) -> bytes: + """Convert to compact binary format for fast transfer.""" + tensors_desc, chunks, _, _ = self._build_tensors_desc(cpu=True) + return b"".join([self._build_header_bytes(tensors_desc)] + chunks) + def to_gpu_tensor(self) -> torch.Tensor: + """Convert to a packed GPU tensor for raw-data connectors.""" + tensors_desc, chunks, data_offset, device = self._build_tensors_desc(cpu=False) + if device is None: + raise RuntimeError("No CUDA tensors found, use to_bytes() instead") + header_prefix = self._build_header_bytes(tensors_desc) + output = torch.empty(len(header_prefix) + data_offset, dtype=torch.uint8, device=device) + output[: len(header_prefix)].copy_(torch.frombuffer(bytearray(header_prefix), dtype=torch.uint8)) pos = len(header_prefix) - for t_flat in gpu_tensors: + for t_flat in chunks: n = t_flat.numel() output[pos : pos + n].copy_(t_flat) pos += n - return output @staticmethod @@ -237,11 +220,8 @@ class KVCacheTransferData: return layer_idx @staticmethod - def from_bytes(raw: "bytes | bytearray | memoryview") -> dict[str, Any]: - """Reconstruct KV cache data from the packed bytes format.""" - raw_mv = memoryview(raw) if not isinstance(raw, memoryview) else raw - header, tensor_data_mv = KVCacheTransferData._load_header_from_memoryview(raw_mv) - + def _populate_caches(header: dict[str, Any], get_tensor: callable) -> dict[str, Any]: + """Shared deserialization loop for both CPU and GPU paths.""" num_layers = header["nl"] key_cache: list[torch.Tensor | None] = [None] * num_layers value_cache: list[torch.Tensor | None] = [None] * num_layers @@ -249,20 +229,9 @@ class KVCacheTransferData: for info in header["td"]: if info.get("x"): continue - name: str = info["n"] torch_dtype = KVCacheTransferData._resolve_torch_dtype(info["d"]) - offset, nbytes = KVCacheTransferData._validate_tensor_span(name, info, len(tensor_data_mv)) - t = ( - torch.frombuffer( - tensor_data_mv, - dtype=torch.uint8, - offset=offset, - count=nbytes, - ) - .view(torch_dtype) - .reshape(info["s"]) - ) + t = get_tensor(info).view(torch_dtype).reshape(info["s"]) layer_idx = KVCacheTransferData._resolve_layer_idx(info, num_layers) if name.startswith("key_cache_"): key_cache[layer_idx] = t @@ -276,37 +245,30 @@ class KVCacheTransferData: "metadata": header["meta"], } + @staticmethod + def from_bytes(raw: "bytes | bytearray | memoryview") -> dict[str, Any]: + """Reconstruct KV cache data from the packed bytes format.""" + raw_mv = memoryview(raw) if not isinstance(raw, memoryview) else raw + header, tensor_data_mv = KVCacheTransferData._load_header_from_memoryview(raw_mv) + data_len = len(tensor_data_mv) + + def _get(info: dict) -> torch.Tensor: + offset, nbytes = KVCacheTransferData._validate_tensor_span(info["n"], info, data_len) + return torch.frombuffer(tensor_data_mv, dtype=torch.uint8, offset=offset, count=nbytes) + + return KVCacheTransferData._populate_caches(header, _get) + @staticmethod def from_bytes_gpu(gpu_tensor: torch.Tensor) -> dict[str, Any]: """Reconstruct KV cache data from a packed GPU tensor.""" header, data_start = KVCacheTransferData._load_header_from_tensor(gpu_tensor) + data_len = int(gpu_tensor.numel()) - data_start - num_layers = header["nl"] - key_cache: list[torch.Tensor | None] = [None] * num_layers - value_cache: list[torch.Tensor | None] = [None] * num_layers - tensor_data_bytes = int(gpu_tensor.numel()) - data_start + def _get(info: dict) -> torch.Tensor: + offset, nbytes = KVCacheTransferData._validate_tensor_span(info["n"], info, data_len) + return gpu_tensor[data_start + offset : data_start + offset + nbytes].clone() - for info in header["td"]: - if info.get("x"): - continue - - name: str = info["n"] - torch_dtype = KVCacheTransferData._resolve_torch_dtype(info["d"]) - offset, nbytes = KVCacheTransferData._validate_tensor_span(name, info, tensor_data_bytes) - t = gpu_tensor[data_start + offset : data_start + offset + nbytes].clone() - t = t.view(torch_dtype).reshape(info["s"]) - layer_idx = KVCacheTransferData._resolve_layer_idx(info, num_layers) - if name.startswith("key_cache_"): - key_cache[layer_idx] = t - elif name.startswith("value_cache_"): - value_cache[layer_idx] = t - - return { - "request_id": header["rid"], - "layer_blocks": {"key_cache": key_cache, "value_cache": value_cache}, - "block_ids": header["bids"], - "metadata": header["meta"], - } + return KVCacheTransferData._populate_caches(header, _get) class OmniKVTransferManager: @@ -341,6 +303,30 @@ class OmniKVTransferManager: else (None, None) ) + local_rank = get_local_tp_rank() + + if config.from_tp <= 1 and config.to_tp <= 1: + detected_tp = get_tp_world_size() + from_tp = detected_tp + to_tp = detected_tp + else: + from_tp = config.from_tp + to_tp = config.to_tp + + self._tp_topo = KVTPTopology(source_tp_size=from_tp, target_tp_size=to_tp, local_rank=local_rank) + + # Injectable hooks (compatible with PR #2677 OmniConnectorModelRunnerMixin). + self.kv_send_key_builder: Callable | None = None + self.kv_recv_key_builder: Callable | None = None + self.kv_payload_merger: Callable | None = None + self.kv_payload_slicer: Callable | None = None + + # Base sender endpoint (rank-0 host/port) stored during + # update_sender_info(). Used by the receive path to construct + # per-rank metadata for heterogeneous TP without querying a registry. + self._sender_base_host: str | None = None + self._sender_base_zmq_port: int | None = None + if config.need_send_cache and config.connector_config: try: _ = self.connector @@ -348,11 +334,20 @@ class OmniKVTransferManager: except Exception as e: logger.warning("Failed to eagerly initialize sender connector: %s", e) + # ------------------------------------------------------------------ # + # Factory helpers + # ------------------------------------------------------------------ # + @classmethod def _create(cls, cfg: dict | None) -> "OmniKVTransferManager": """Create manager from raw config dict.""" if not cfg or not isinstance(cfg, dict): return cls(OmniKVCacheConfig()) + + rank_mapping = cfg.get("rank_mapping", {}) + if not isinstance(rank_mapping, dict): + rank_mapping = {} + return cls( OmniKVCacheConfig( connector_config=cfg.get("connector_config"), @@ -363,18 +358,17 @@ class OmniKVTransferManager: need_recv_cache=cfg.get("need_recv_cache", False), need_send_cache=cfg.get("need_send_cache", False), recv_timeout=cfg.get("recv_timeout", 30.0), + from_tp=int(rank_mapping.get("from_tp", 1)), + to_tp=int(rank_mapping.get("to_tp", 1)), ) ) @classmethod - def from_model_config(cls, config: Any) -> "OmniKVTransferManager": - """Create from model config (for AR model runner).""" + def from_od_config(cls, config: Any) -> "OmniKVTransferManager": + """Create from model or OmniDiffusion config.""" return cls._create(getattr(config, "omni_kv_config", None)) - @classmethod - def from_od_config(cls, config: Any) -> "OmniKVTransferManager": - """Create from OmniDiffusion config (for diffusion runner).""" - return cls._create(getattr(config, "omni_kv_config", None)) + from_model_config = from_od_config @classmethod def from_vllm_config(cls, vllm_config: Any, model_config: Any) -> "OmniKVTransferManager": @@ -417,45 +411,33 @@ class OmniKVTransferManager: ) c_extra["to_stage"] = str(self.config.to_stage) if self.config.to_stage is not None else "1" + try: + stage_int = int(self.config.from_stage) if self.config.from_stage is not None else 0 + except (TypeError, ValueError): + stage_int = 0 + zmq_port = kv_zmq_port(base_port, stage_int, self._tp_topo.local_rank) + if self.config.need_send_cache: c_extra["role"] = "sender" - from_stage = self.config.from_stage - if from_stage is not None: - try: - c_extra["zmq_port"] = base_port + KV_TRANSFER_PORT_OFFSET + int(from_stage) - except (TypeError, ValueError): - c_extra["zmq_port"] = base_port + KV_TRANSFER_PORT_OFFSET + c_extra["zmq_port"] = zmq_port elif self.config.need_recv_cache: c_extra["role"] = "receiver" - from_stage = self.config.from_stage - sender_port = base_port + KV_TRANSFER_PORT_OFFSET - if from_stage is not None: - try: - sender_port = base_port + KV_TRANSFER_PORT_OFFSET + int(from_stage) - except (TypeError, ValueError): - pass c_extra.setdefault("sender_host", c_extra.get("host", "127.0.0.1")) - c_extra.setdefault("sender_zmq_port", sender_port) + c_extra.setdefault("sender_zmq_port", zmq_port) logger.info( - "Initializing OmniConnector (purpose=kv_transfer) with config: %s, role: %s", - cfg, + "Initializing OmniConnector type=%s role=%s", + c_type, c_extra.get("role", "N/A"), ) self._connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=c_type, extra=c_extra)) - except Exception as e: - logger.error(f"Failed to initialize OmniConnector: {e}") - import traceback - - traceback.print_exc() - # Cache failure sentinel to avoid repeated initialization attempts in hot paths. + except Exception: + logger.exception("Failed to initialize OmniConnector") self._connector = False return self._connector if self._connector else None - def get_connector(self): - """Get connector (compatibility wrapper for existing code).""" - return self.connector + get_connector = property(lambda self: self.connector) def _resolve_sender_info( self, sender_info: dict[str, Any], sender_stage_id: str | int | None = None @@ -513,8 +495,187 @@ class OmniKVTransferManager: cache_list[idx] = tensor.clone() return data + def _slice_transfer_data_for_target(self, kv_data: KVCacheTransferData, target_rank: int) -> KVCacheTransferData: + """Pre-slice sender payload for one target rank when sender TP < receiver TP.""" + topo = self._tp_topo + ratio = topo.target_tp_size // topo.source_tp_size + offset_in_sender = target_rank % ratio + metadata = dict(kv_data.metadata) if isinstance(kv_data.metadata, dict) else {} + metadata["tp_head_slice"] = { + "applied": True, + "side": "sender", + "target_rank": target_rank, + "source_rank": topo.local_rank, + "from_tp": topo.source_tp_size, + "to_tp": topo.target_tp_size, + "offset_in_shard": offset_in_sender, + "num_slices": ratio, + } + return KVCacheTransferData( + request_id=kv_data.request_id, + layer_blocks=slice_layer_blocks(kv_data.layer_blocks, offset_in_sender, ratio), + block_ids=list(kv_data.block_ids), + metadata=metadata, + ) + + def _serialize_transfer_payload(self, kv_data: KVCacheTransferData) -> torch.Tensor | bytes | dict[str, Any]: + """Serialize KV transfer data using the connector's fastest supported path.""" + if getattr(self.connector, "supports_raw_data", False): + try: + return kv_data.to_gpu_tensor() + except Exception: + pass + try: + return kv_data.to_bytes() + except Exception: + return kv_data.to_dict() + + @staticmethod + def _collect_request_kv_payload(req: Any) -> dict[str, object]: + """Collect request-side KV objects for object broadcast.""" + kv_payload: dict[str, object] = {} + for attr in ("past_key_values", "kv_metadata"): + val = getattr(req, attr, None) + if val is not None: + kv_payload[attr] = val + + if hasattr(req, "sampling_params") and req.sampling_params is not None: + for key in list(vars(req.sampling_params).keys()): + if key in ("past_key_values", "kv_metadata") or ( + key.startswith("cfg_") + and ( + key.endswith("_past_key_values") + or key.endswith("_kv_metadata") + or key + in ( + "cfg_kv_request_ids", + "cfg_active_branch", + "cfg_branch_roles", + "cfg_branch_past_key_values", + "cfg_branch_kv_metadata", + ) + ) + ): + val = getattr(req.sampling_params, key, None) + if val is not None: + kv_payload[f"sp.{key}"] = val + + return kv_payload + + @staticmethod + def _apply_request_kv_payload( + req: Any, + kv_payload: dict[str, object], + target_device: torch.device | None = None, + ) -> None: + """Apply a broadcast KV payload back onto a request object.""" + for attr in ("past_key_values", "kv_metadata"): + val = kv_payload.get(attr) + if val is not None: + if target_device is not None: + val = _move_to_device(val, target_device) + setattr(req, attr, val) + + if hasattr(req, "sampling_params") and req.sampling_params is not None: + for key, val in kv_payload.items(): + if key.startswith("sp."): + if target_device is not None: + val = _move_to_device(val, target_device) + setattr(req.sampling_params, key[3:], val) + + @staticmethod + def _discover_cfg_branch_roles(req: Any) -> list[str]: + """Discover CFG branch roles in a stable order.""" + sampling_params = getattr(req, "sampling_params", None) + if sampling_params is None: + return [] + + roles: list[str] = [] + branch_map = getattr(sampling_params, "cfg_branch_past_key_values", None) or {} + for preferred_role in ("cfg_text", "cfg_img"): + if ( + preferred_role in branch_map + or getattr(sampling_params, f"{preferred_role}_past_key_values", None) is not None + ): + roles.append(preferred_role) + + for role in branch_map.keys(): + if role not in roles and branch_map.get(role) is not None: + roles.append(role) + + for key in vars(sampling_params).keys(): + if not (key.startswith("cfg_") and key.endswith("_past_key_values")): + continue + role = key.removesuffix("_past_key_values") + if role in ("cfg_branch",) or role in roles: + continue + if getattr(sampling_params, key, None) is not None: + roles.append(role) + + return roles + + @classmethod + def _build_cfg_rank_local_payloads(cls, req: Any, cfg_size: int) -> list[dict[str, object] | None]: + """Build per-cfg-rank payloads so each rank receives only its branch KV.""" + full_payload = cls._collect_request_kv_payload(req) + payloads: list[dict[str, object] | None] = [] + + main_payload = { + key: value + for key, value in full_payload.items() + if key in ("past_key_values", "kv_metadata", "sp.past_key_values", "sp.kv_metadata") + } + branch_roles = cls._discover_cfg_branch_roles(req) + if branch_roles: + main_payload["sp.cfg_branch_roles"] = list(branch_roles) + main_payload["sp.cfg_active_branch"] = None + payloads.append(main_payload or None) + + sampling_params = getattr(req, "sampling_params", None) + branch_map = getattr(sampling_params, "cfg_branch_past_key_values", None) or {} + branch_metadata_map = getattr(sampling_params, "cfg_branch_kv_metadata", None) or {} + + for role in branch_roles: + if sampling_params is None: + payloads.append(None) + continue + + branch_kv = branch_map.get(role) + if branch_kv is None: + branch_kv = getattr(sampling_params, f"{role}_past_key_values", None) + branch_metadata = branch_metadata_map.get(role) + if branch_metadata is None: + branch_metadata = getattr(sampling_params, f"{role}_kv_metadata", None) + if branch_kv is None: + payloads.append(None) + continue + + local_payload = dict(main_payload) + local_payload["sp.cfg_active_branch"] = role + local_payload["sp.cfg_branch_roles"] = list(branch_roles) + local_payload["sp.cfg_branch_past_key_values"] = {role: branch_kv} + local_payload[f"sp.{role}_past_key_values"] = branch_kv + if branch_metadata is not None: + local_payload["sp.cfg_branch_kv_metadata"] = {role: branch_metadata} + local_payload[f"sp.{role}_kv_metadata"] = branch_metadata + + payloads.append(local_payload) + + while len(payloads) < cfg_size: + payloads.append(None) + + return payloads[:cfg_size] + def update_sender_info(self, sender_info: dict[str, Any], sender_stage_id: str | int | None = None) -> None: - """Update receiver-side sender info before loading remote KV cache.""" + """Update receiver-side sender info before loading remote KV cache. + + The orchestrator always reports rank-0's ZMQ port. When TP > 1 the + receiver must offset the port so that each TP rank connects to the + corresponding sender rank's port. + + The base host/port are also stored so that the receive path can + construct per-rank metadata for heterogeneous TP scenarios. + """ if not self.config.need_recv_cache: return @@ -523,18 +684,39 @@ class OmniKVTransferManager: logger.warning("Invalid sender_info format: %s", sender_info) return + sender_host = actual_info.get("host") + base_zmq_port = actual_info.get("zmq_port") + + # Store base sender info for per-rank metadata construction. + self._sender_base_host = sender_host + if base_zmq_port is not None: + self._sender_base_zmq_port = int(base_zmq_port) + + # --- Default sender: offset to match this receiver's corresponding sender rank --- + zmq_port = base_zmq_port + if zmq_port is not None and self._tp_topo.local_rank > 0: + zmq_port = int(zmq_port) + self._tp_topo.local_rank * KV_RANK_PORT_STRIDE + if self.config.connector_config: - self.config.connector_config["sender_host"] = actual_info.get("host") - self.config.connector_config["sender_zmq_port"] = actual_info.get("zmq_port") + self.config.connector_config["sender_host"] = sender_host + self.config.connector_config["sender_zmq_port"] = zmq_port if self._connector and hasattr(self._connector, "update_sender_info"): try: - self._connector.update_sender_info(actual_info.get("host"), actual_info.get("zmq_port")) + self._connector.update_sender_info(sender_host, zmq_port) except Exception: if hasattr(self._connector, "sender_host"): - self._connector.sender_host = actual_info.get("host") + self._connector.sender_host = sender_host if hasattr(self._connector, "sender_zmq_port"): - self._connector.sender_zmq_port = actual_info.get("zmq_port") + self._connector.sender_zmq_port = zmq_port + + logger.info( + "Sender info updated: host=%s, base_port=%s, adjusted_port=%s (local_rank=%s)", + sender_host, + base_zmq_port, + zmq_port, + self._tp_topo.local_rank, + ) def handle_finished_requests_kv_transfer( self, @@ -692,35 +874,54 @@ class OmniKVTransferManager: kv_data.request_id = transfer_req_id serialization_start = time.perf_counter() - transfer_data: torch.Tensor | bytes | dict[str, Any] - supports_raw = getattr(self.connector, "supports_raw_data", False) + topo = self._tp_topo + send_keys = build_rank_aware_send_keys( + transfer_req_id, from_stage, to_stage, topo, hook=self.kv_send_key_builder + ) + sender_slice_active = ( + topo.source_tp_size < topo.target_tp_size and len(send_keys) > 1 and not callable(self.kv_send_key_builder) + ) + per_key_payloads: list[tuple[str, torch.Tensor | bytes | dict[str, Any]]] = [] - try: - if supports_raw: - transfer_data = kv_data.to_gpu_tensor() + if sender_slice_active: + target_ranks = get_kv_target_ranks(topo) + if len(target_ranks) != len(send_keys): + logger.warning( + "Skip sender-side KV slicing because target rank count does not match send key count: " + "target_ranks=%s send_keys=%s", + len(target_ranks), + len(send_keys), + ) + sender_slice_active = False else: - raise RuntimeError("Connector does not support raw tensor") - except Exception: - try: - transfer_data = kv_data.to_bytes() - except Exception: - data_dict = kv_data.to_dict() - data_dict["request_id"] = transfer_req_id - transfer_data = data_dict + for put_key, target_rank in zip(send_keys, target_ranks, strict=False): + sliced_kv_data = self._slice_transfer_data_for_target(kv_data, target_rank) + per_key_payloads.append((put_key, self._serialize_transfer_payload(sliced_kv_data))) + + if not per_key_payloads: + transfer_data = self._serialize_transfer_payload(kv_data) + per_key_payloads = [(put_key, transfer_data) for put_key in send_keys] serialization_ms = (time.perf_counter() - serialization_start) * 1000 logger.info("KV cache serialized for %s in %.1f ms", transfer_req_id, serialization_ms) transfer_start = time.perf_counter() - success, size, _ = self._transfer_with_retry(from_stage, to_stage, f"kv_cache_{transfer_req_id}", transfer_data) + total_size = 0 + all_succeeded = True + for put_key, transfer_data in per_key_payloads: + success, size, _ = self._transfer_with_retry(from_stage, to_stage, put_key, transfer_data) + total_size += size + all_succeeded = all_succeeded and success + elapsed = time.perf_counter() - transfer_start - if success: - mbps = (size / 1024 / 1024) / elapsed if elapsed > 0 else 0 + if all_succeeded: + mbps = (total_size / 1024 / 1024) / elapsed if elapsed > 0 else 0 logger.info( - "KV transfer OK: %s, %s bytes, %.3fs, %.1f MB/s", + "KV transfer OK: %s, %s bytes across %s key(s), %.3fs, %.1f MB/s", transfer_req_id, - size, + total_size, + len(send_keys), elapsed, mbps, ) @@ -731,7 +932,7 @@ class OmniKVTransferManager: self, from_stage: str, to_stage: str, - request_id: str, + put_key: str, data: "dict[str, Any] | bytes | torch.Tensor", max_retries: int = 3, ) -> tuple[bool, int, dict[str, Any] | None]: @@ -740,7 +941,7 @@ class OmniKVTransferManager: Args: from_stage: Source stage identifier to_stage: Target stage identifier - request_id: Request identifier for the key + put_key: Pre-built connector key (rank-aware when TP > 1) data: Data to transfer max_retries: Maximum number of retry attempts @@ -749,14 +950,12 @@ class OmniKVTransferManager: """ for attempt in range(max_retries): try: - # Build the full key for connector - full_request_id = f"omni_{from_stage}_to_{to_stage}_{request_id}" success, size, metadata = self.connector.put( - from_stage=from_stage, to_stage=to_stage, put_key=full_request_id, data=data + from_stage=from_stage, to_stage=to_stage, put_key=put_key, data=data ) if success: return success, size, metadata - logger.warning(f"Transfer attempt {attempt + 1} failed for {request_id}") + logger.warning(f"Transfer attempt {attempt + 1} failed for {put_key}") except Exception as e: logger.warning(f"Transfer attempt {attempt + 1} exception: {e}") @@ -801,22 +1000,46 @@ class OmniKVTransferManager: poll_interval = 0.01 max_poll_interval = 0.5 - logger.info(f"Wait for KV cache for request {request_id} from stage {from_stage} to {to_stage}...") + topo = self._tp_topo + recv_key_pairs = build_rank_aware_recv_keys( + request_id, from_stage, to_stage, topo, hook=self.kv_recv_key_builder + ) + pending_pairs = list(recv_key_pairs) + received_payloads: dict[str, tuple[dict[str, Any], int]] = {} + + logger.info( + "Wait for KV cache for request %s from stage %s to %s via %s key(s)...", + request_id, + from_stage, + to_stage, + len(recv_key_pairs), + ) try: while True: - # Build the full key for connector - full_request_id = f"omni_{from_stage}_to_{to_stage}_kv_cache_{request_id}" link_start = time.perf_counter() - result = self.connector.get( - from_stage=from_stage, - to_stage=to_stage, - get_key=full_request_id, - ) - if result: + for get_key, from_rank in list(pending_pairs): + # Construct per-rank metadata so the connector queries + # the correct sender endpoint (heterogeneous TP path). + # When from_rank is None (TP<=1), metadata stays None + # and the connector falls back to its default sender. + rank_metadata: dict[str, Any] | None = None + if from_rank is not None and self._sender_base_host and self._sender_base_zmq_port is not None: + rank_metadata = { + "source_host": self._sender_base_host, + "source_port": self._sender_base_zmq_port + from_rank * KV_RANK_PORT_STRIDE, + } + + result = self.connector.get( + from_stage=from_stage, + to_stage=to_stage, + get_key=get_key, + metadata=rank_metadata, + ) + if not result: + continue + raw_data, size = result - elapsed = time.time() - start_time - link_ms = (time.perf_counter() - link_start) * 1000 managed_buffer = None if hasattr(raw_data, "tensor") and hasattr(raw_data, "release"): @@ -844,6 +1067,21 @@ class OmniKVTransferManager: else: data = raw_data + received_payloads[get_key] = (data, size) + pending_pairs.remove((get_key, from_rank)) + + if not pending_pairs and received_payloads: + elapsed = time.time() - start_time + link_ms = (time.perf_counter() - link_start) * 1000 + ordered_payloads = [received_payloads[key][0] for key, _ in recv_key_pairs] + total_size = sum(received_payloads[key][1] for key, _ in recv_key_pairs) + + if len(ordered_payloads) == 1: + data = ordered_payloads[0] + else: + data = merge_received_rank_shards(ordered_payloads, merger=self.kv_payload_merger) + data = slice_received_rank_shard(data, topo, slicer=self.kv_payload_slicer) + try: if isinstance(data, dict) and "layer_blocks" in data: layer_blocks = data["layer_blocks"] @@ -856,18 +1094,18 @@ class OmniKVTransferManager: continue if target_device is not None and tensor.device != target_device: cache_list[i] = tensor.to(target_device).contiguous() - finally: - if managed_buffer is not None: - managed_buffer.release() + except Exception: + logger.exception("Failed to move KV cache tensors to target device") logger.info( - "Successfully received KV cache for %s, %s bytes, wait=%.3fs, link=%.1fms", + "Successfully received KV cache for %s, %s bytes across %s key(s), wait=%.3fs, link=%.1fms", request_id, - size, + total_size, + len(recv_key_pairs), elapsed, link_ms, ) - return data, size + return data, total_size if time.time() - start_time > timeout: logger.error(f"Timeout waiting for KV cache for request {request_id} after {timeout}s") @@ -876,11 +1114,8 @@ class OmniKVTransferManager: time.sleep(poll_interval) poll_interval = min(poll_interval * 2, max_poll_interval) - except Exception as e: - logger.error(f"Error receiving KV cache for {request_id}: {e}") - import traceback - - traceback.print_exc() + except Exception: + logger.exception("Error receiving KV cache for %s", request_id) return None, 0 def apply_kv_cache_to_request(self, req: Any, data: dict[str, Any]) -> None: @@ -994,73 +1229,79 @@ class OmniKVTransferManager: cfg_kv_collect_func: Callable | None = None, target_device: torch.device | None = None, ) -> bool: - """Broadcast-aware wrapper around :meth:`receive_multi_kv_cache`. + """Distributed wrapper around :meth:`receive_multi_kv_cache`. - SharedMemory connector is single-reader: once rank 0 consumes the - segment it is deleted. For multi-GPU stages (e.g. sequence-parallel) - only rank 0 receives; the result is then broadcast to every other - rank via the world process-group. - - For single-worker stages this is equivalent to calling - :meth:`receive_multi_kv_cache` directly. + TP-aware path selection: + - world size 1: direct receive + - TP active, cfg size 1: each rank independently receives + - TP active, cfg size > 1: cfg-rank 0 receives, then broadcasts to + peers that share the same TP rank + - TP inactive: legacy rank-0 receive then world broadcast """ - from vllm_omni.diffusion.distributed.parallel_state import get_world_group + from vllm_omni.diffusion.distributed.parallel_state import ( + get_cfg_group, + get_classifier_free_guidance_rank, + get_classifier_free_guidance_world_size, + get_world_group, + ) world = get_world_group() if world.world_size <= 1: return self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device) - # --- rank 0: receive to CPU (needed for pickle-based broadcast) --- + topo = self._tp_topo + tp_active = topo.source_tp_size > 1 or topo.target_tp_size > 1 + cfg_size = 1 + cfg_rank = 0 + cfg_group = None + try: + cfg_size = get_classifier_free_guidance_world_size() + cfg_rank = get_classifier_free_guidance_rank() + cfg_group = get_cfg_group() + except Exception: + cfg_size = 1 + cfg_rank = 0 + cfg_group = None + + if tp_active and cfg_size <= 1: + logger.info( + "Rank-aware KV receive: rank %s independently receiving (from_tp=%s, to_tp=%s)", + topo.local_rank, + topo.source_tp_size, + topo.target_tp_size, + ) + return self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device) + + if tp_active and cfg_size > 1 and cfg_group is not None: + kv_payload: dict[str, object] | None = None + if cfg_rank == 0: + received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, torch.device("cpu")) + rank_payloads = self._build_cfg_rank_local_payloads(req, cfg_size) if received else [None] * cfg_size + kv_payload = rank_payloads[0] + for dst_rank in range(1, cfg_size): + cfg_group.send_object(rank_payloads[dst_rank], dst_rank) + else: + kv_payload = cfg_group.recv_object(0) + + if not kv_payload: + return False + + self._apply_request_kv_payload(req, kv_payload, target_device) + return True + + kv_payload: dict[str, object] | None = None if world.rank_in_group == 0: - self.receive_multi_kv_cache(req, cfg_kv_collect_func, torch.device("cpu")) + received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, torch.device("cpu")) + if received: + kv_payload = self._collect_request_kv_payload(req) - kv_payload: dict[str, object] = {} - for attr in ("past_key_values", "kv_metadata"): - val = getattr(req, attr, None) - if val is not None: - kv_payload[attr] = val + kv_payload = world.broadcast_object(kv_payload, src=0) - if hasattr(req, "sampling_params") and req.sampling_params is not None: - for key in list(vars(req.sampling_params).keys()): - if (key.startswith("cfg_") and key.endswith("_past_key_values")) or key in ( - "past_key_values", - "kv_metadata", - ): - val = getattr(req.sampling_params, key, None) - if val is not None: - kv_payload[f"sp.{key}"] = val - - payload_list = [kv_payload] - # Use broadcast_object_list (pickle-based) instead of broadcast_tensor_dict - # because the KV cache is a heterogeneous nested structure (NaiveCache objects - # with metadata + tensors), not a flat tensor dict. This runs once before - # the denoising loop so the serialization cost is negligible. - torch.distributed.broadcast_object_list(payload_list, src=world.ranks[0], group=world.cpu_group) - kv_payload = payload_list[0] - else: - payload_list: list[dict[str, object] | None] = [None] - torch.distributed.broadcast_object_list(payload_list, src=world.ranks[0], group=world.cpu_group) - kv_payload = payload_list[0] - - # --- apply on ALL ranks (rank 0 also needs CPU→GPU move) --- if not kv_payload: return False - for attr in ("past_key_values", "kv_metadata"): - val = kv_payload.get(attr) - if val is not None: - if target_device is not None: - val = _move_to_device(val, target_device) - setattr(req, attr, val) - - if hasattr(req, "sampling_params") and req.sampling_params is not None: - for key, val in kv_payload.items(): - if key.startswith("sp."): - if target_device is not None: - val = _move_to_device(val, target_device) - setattr(req.sampling_params, key[3:], val) - + self._apply_request_kv_payload(req, kv_payload, target_device) return True diff --git a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py index 2cb48a8b3..12b9b3d4f 100644 --- a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py +++ b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py @@ -1,15 +1,380 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Utility helpers for KV cache manipulation.""" +"""Utility helpers for KV cache manipulation, TP routing, and merge/slice.""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any import torch +from vllm.distributed.parallel_state import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) from vllm.logger import init_logger +from .initialization import KV_RANK_PORT_STRIDE, KV_TRANSFER_PORT_OFFSET + logger = init_logger(__name__) LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor] +# ------------------------------------------------------------------ # +# TP Topology +# ------------------------------------------------------------------ # + + +@dataclass(frozen=True) +class KVTPTopology: + """Immutable descriptor for a KV-transfer parallel mapping. + + Captures sender/receiver parallel sizes and the local rank within + that parallel dimension. Works for any divisible parallel dimension + (TP, SP, Ring Attention). + """ + + source_tp_size: int + target_tp_size: int + local_rank: int + + def __post_init__(self) -> None: + if self.source_tp_size <= 0 or self.target_tp_size <= 0: + raise ValueError( + f"Parallel sizes must be positive: " + f"source_tp_size={self.source_tp_size}, target_tp_size={self.target_tp_size}" + ) + if self.local_rank < 0: + raise ValueError(f"local_rank must be non-negative, got {self.local_rank}") + + @property + def is_heterogeneous(self) -> bool: + return self.source_tp_size != self.target_tp_size + + @property + def ratio(self) -> int: + """Larger parallel size divided by smaller. Always >= 1.""" + return max(self.source_tp_size, self.target_tp_size) // min(self.source_tp_size, self.target_tp_size) + + +# ------------------------------------------------------------------ # +# Runtime TP detection +# ------------------------------------------------------------------ # + + +def get_local_tp_rank() -> int: + """Return the TP-local rank of this worker process. + + Uses ``get_tensor_model_parallel_rank()`` which returns the rank + within the TP group only, not the stage-global rank. + """ + try: + return get_tensor_model_parallel_rank() + except Exception: + logger.debug("TP parallel state not initialized, falling back to LOCAL_RANK env", exc_info=True) + try: + return int(os.environ.get("LOCAL_RANK", "0")) + except (ValueError, TypeError): + return 0 + + +def get_tp_world_size() -> int: + """Return the TP world size (tensor-parallel dimension only). + + Uses ``get_tensor_model_parallel_world_size()`` so that + cfg_parallel, SP, PP etc. are not included in the count. + """ + try: + return get_tensor_model_parallel_world_size() + except Exception: + logger.debug("TP parallel state not initialized, defaulting world_size=1", exc_info=True) + return 1 + + +# ------------------------------------------------------------------ # +# ZMQ port computation +# ------------------------------------------------------------------ # + + +def kv_zmq_port(base_port: int, from_stage: int, local_rank: int = 0) -> int: + """Compute the ZMQ port for a KV-transfer connector. + + Each TP rank gets its own port so that TP > 1 deployments do not + cause ``EADDRINUSE`` when multiple sender workers bind on the same + host. The formula is backward-compatible: rank 0 produces the same + port as the previous ``base + OFFSET + stage`` formula. + """ + return base_port + KV_TRANSFER_PORT_OFFSET + local_rank * KV_RANK_PORT_STRIDE + from_stage + + +# ------------------------------------------------------------------ # +# TP topology validation and rank routing +# ------------------------------------------------------------------ # + + +def validate_kv_tp_topology(topo: KVTPTopology) -> None: + """Reject heterogeneous TP mappings that cannot be routed losslessly.""" + larger = max(topo.source_tp_size, topo.target_tp_size) + smaller = min(topo.source_tp_size, topo.target_tp_size) + if larger % smaller != 0: + raise ValueError( + f"KV TP mapping must be divisible: " + f"source_tp_size={topo.source_tp_size}, " + f"target_tp_size={topo.target_tp_size}" + ) + + +def get_kv_target_ranks(topo: KVTPTopology) -> list[int]: + """Which remote ranks this local rank sends KV shards to (send side).""" + validate_kv_tp_topology(topo) + if topo.source_tp_size == topo.target_tp_size: + return [topo.local_rank] + if topo.source_tp_size > topo.target_tp_size: + return [topo.local_rank // (topo.source_tp_size // topo.target_tp_size)] + ratio = topo.target_tp_size // topo.source_tp_size + return [topo.local_rank * ratio + i for i in range(ratio)] + + +def get_kv_source_ranks(topo: KVTPTopology) -> list[int]: + """Which remote ranks this local rank receives KV shards from (recv side).""" + validate_kv_tp_topology(topo) + if topo.source_tp_size == topo.target_tp_size: + return [topo.local_rank] + if topo.source_tp_size > topo.target_tp_size: + ratio = topo.source_tp_size // topo.target_tp_size + return [topo.local_rank * ratio + i for i in range(ratio)] + return [topo.local_rank // (topo.target_tp_size // topo.source_tp_size)] + + +# ------------------------------------------------------------------ # +# Rank-aware connector key building +# ------------------------------------------------------------------ # + + +def get_kv_connector_key( + req_id: str, + from_stage: int | str, + chunk_id: int, + from_rank: int, + to_rank: int, +) -> str: + """Build connector key that includes rank info for KV transfers. + + Format matches PR #2677: ``{req_id}_{from_stage}_{chunk_id}_{from_rank}_{to_rank}`` + """ + return f"{req_id}_{from_stage}_{chunk_id}_{from_rank}_{to_rank}" + + +def build_rank_aware_send_keys( + request_id: str, + from_stage: str, + to_stage: str, + topo: KVTPTopology, + hook: Callable | None = None, +) -> list[str]: + """Build send-side connector keys, checking injectable hook first.""" + if callable(hook): + keys = list(hook(request_id, from_stage, to_stage)) + if keys: + return keys + if topo.source_tp_size <= 1 and topo.target_tp_size <= 1: + return [f"omni_{from_stage}_to_{to_stage}_kv_cache_{request_id}"] + target_ranks = get_kv_target_ranks(topo) + return [get_kv_connector_key(request_id, from_stage, 0, topo.local_rank, r) for r in target_ranks] + + +def build_rank_aware_recv_keys( + request_id: str, + from_stage: str, + to_stage: str, + topo: KVTPTopology, + hook: Callable | None = None, +) -> list[tuple[str, int | None]]: + """Build recv-side connector keys with sender rank info. + + Returns a list of ``(key, from_rank)`` tuples. ``from_rank`` is + ``None`` when TP <= 1 (single sender, no per-rank routing needed). + For TP > 1, ``from_rank`` identifies which sender rank owns the + key so that the connector can route metadata queries to the + correct endpoint. + """ + if callable(hook): + raw = list(hook(request_id, from_stage, to_stage)) + if raw: + if isinstance(raw[0], tuple): + return raw + # Hook returned plain strings (e.g. OmniConnectorModelRunnerMixin. + # get_rank_aware_kv_keys). Reconstruct from_rank from topology so + # Mooncake connector can route metadata queries to the correct + # sender endpoint in heterogeneous TP. + # TODO: have the mixin return (key, from_rank) tuples directly + # to avoid this indirect reconstruction. + source_ranks = get_kv_source_ranks(topo) + if len(raw) == len(source_ranks): + return list(zip(raw, source_ranks)) + return [(k, None) for k in raw] + if topo.source_tp_size <= 1 and topo.target_tp_size <= 1: + return [(f"omni_{from_stage}_to_{to_stage}_kv_cache_{request_id}", None)] + source_ranks = get_kv_source_ranks(topo) + return [(get_kv_connector_key(request_id, from_stage, 0, r, topo.local_rank), r) for r in source_ranks] + + +# ------------------------------------------------------------------ # +# KV tensor head slicing (heterogeneous TP) +# ------------------------------------------------------------------ # + + +def slice_kv_tensor_heads( + tensor: torch.Tensor | None, + offset_in_shard: int, + num_slices: int, +) -> torch.Tensor | None: + """Slice one KV tensor along its head dimension (dim 1).""" + if tensor is None: + return None + if not isinstance(tensor, torch.Tensor): + return tensor + if tensor.dim() < 2: + raise ValueError(f"Expected KV tensor with a head dimension, got shape={tuple(tensor.shape)}") + if num_slices <= 0: + raise ValueError(f"num_slices must be > 0, got {num_slices}") + if not (0 <= offset_in_shard < num_slices): + raise ValueError(f"offset_in_shard must be in [0, {num_slices}), got {offset_in_shard}") + + heads_in_shard = tensor.shape[1] + if heads_in_shard % num_slices != 0: + raise ValueError( + "KV head count must be divisible for heterogeneous TP slicing: " + f"heads_in_shard={heads_in_shard}, num_slices={num_slices}" + ) + + heads_per_slice = heads_in_shard // num_slices + start = offset_in_shard * heads_per_slice + end = start + heads_per_slice + return tensor[:, start:end, ...].contiguous() + + +def slice_layer_blocks( + layer_blocks: dict[str, Any], + offset_in_shard: int, + num_slices: int, +) -> dict[str, list[torch.Tensor | None]]: + """Slice all KV layers for one logical receiver rank.""" + sliced_blocks: dict[str, list[torch.Tensor | None]] = {} + for cache_name in ("key_cache", "value_cache"): + cache_list = layer_blocks.get(cache_name, []) + sliced_blocks[cache_name] = [ + slice_kv_tensor_heads(tensor, offset_in_shard, num_slices) for tensor in cache_list + ] + return sliced_blocks + + +# ------------------------------------------------------------------ # +# Multi-rank merge and receiver-side slice +# ------------------------------------------------------------------ # + + +def merge_received_rank_shards( + payloads: list[dict[str, Any]], + merger: Callable | None = None, +) -> dict[str, Any] | None: + """Merge multiple source-rank KV shards for one target rank. + + When *merger* is provided (injectable hook), it is called directly. + Otherwise the default merges along the head dimension (dim 1). + """ + if callable(merger): + return merger(payloads) + if not payloads: + return None + if len(payloads) == 1: + return payloads[0] + + base_payload = payloads[0] + if not isinstance(base_payload, dict) or "layer_blocks" not in base_payload: + return base_payload + + merged: dict[str, Any] = { + "request_id": base_payload.get("request_id"), + "block_ids": list(base_payload.get("block_ids", [])), + "metadata": dict(base_payload.get("metadata", {})), + } + merged_layer_blocks: dict[str, list[torch.Tensor | None]] = {} + + for cache_name in ("key_cache", "value_cache"): + cache_lists = [payload.get("layer_blocks", {}).get(cache_name, []) for payload in payloads] + num_layers = max((len(cache_list) for cache_list in cache_lists), default=0) + merged_cache: list[torch.Tensor | None] = [] + + for layer_idx in range(num_layers): + layer_tensors = [ + cache_list[layer_idx] + for cache_list in cache_lists + if layer_idx < len(cache_list) and cache_list[layer_idx] is not None + ] + if not layer_tensors: + merged_cache.append(None) + elif len(layer_tensors) == 1 or not isinstance(layer_tensors[0], torch.Tensor): + merged_cache.append(layer_tensors[0]) + else: + merged_cache.append(torch.cat(layer_tensors, dim=1).contiguous()) + + merged_layer_blocks[cache_name] = merged_cache + + merged["layer_blocks"] = merged_layer_blocks + return merged + + +def slice_received_rank_shard( + payload: dict[str, Any] | None, + topo: KVTPTopology, + slicer: Callable | None = None, +) -> dict[str, Any] | None: + """Optionally slice a received payload to extract this rank's portion. + + Used when ``to_tp > from_tp``: the sender sent full heads and each + receiver rank slices out its own subset. + """ + if callable(slicer): + return slicer(payload) + if not payload or topo.target_tp_size <= topo.source_tp_size or "layer_blocks" not in payload: + return payload + + metadata = payload.get("metadata", {}) + slice_metadata = metadata.get("tp_head_slice") if isinstance(metadata, dict) else None + if isinstance(slice_metadata, dict) and slice_metadata.get("applied"): + tagged_rank = slice_metadata.get("target_rank") + if tagged_rank is not None and tagged_rank != topo.local_rank: + logger.warning( + "Received pre-sliced KV payload for unexpected target rank: expected=%s got=%s", + topo.local_rank, + tagged_rank, + ) + return payload + + ratio = topo.target_tp_size // topo.source_tp_size + offset_in_sender = topo.local_rank % ratio + updated_metadata = dict(metadata) if isinstance(metadata, dict) else {} + updated_metadata["tp_head_slice"] = { + "applied": True, + "side": "receiver", + "target_rank": topo.local_rank, + "from_tp": topo.source_tp_size, + "to_tp": topo.target_tp_size, + "offset_in_shard": offset_in_sender, + "num_slices": ratio, + } + return { + "request_id": payload.get("request_id"), + "layer_blocks": slice_layer_blocks(payload["layer_blocks"], offset_in_sender, ratio), + "block_ids": list(payload.get("block_ids", [])), + "metadata": updated_metadata, + } + + def normalize_layer_kv( layer_kv: LayerKV, *, diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 054d5342d..23a85e9f5 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -61,6 +61,7 @@ from vllm_omni.engine.stage_engine_startup import ( ) from vllm_omni.engine.stage_init_utils import ( StartedLlmStage, + _inject_inferred_kv_tp_topology, acquire_device_locks, build_diffusion_config, build_engine_args_dict, @@ -78,7 +79,10 @@ from vllm_omni.engine.stage_init_utils import ( setup_stage_devices, terminate_alive_proc, ) -from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs +from vllm_omni.entrypoints.utils import ( + inject_omni_kv_config, + load_and_resolve_stage_configs, +) from vllm_omni.inputs.preprocess import OmniInputPreprocessor from vllm_omni.platforms import current_omni_platform @@ -378,6 +382,12 @@ class AsyncOmniEngine: omni_kv["omni_to_stage"] = omni_to omni_kv.setdefault("stage_id", metadata.stage_id) engine_args_dict["omni_kv_config"] = omni_kv + if self.stage_configs: + _inject_inferred_kv_tp_topology( + engine_args_dict.get("omni_kv_config"), + metadata.stage_id, + self.stage_configs, + ) vllm_config, executor_class = build_vllm_config( stage_cfg, self.model, @@ -747,10 +757,8 @@ class AsyncOmniEngine: setup_stage_devices(configured_stage_id, metadata.runtime_cfg) omni_conn_cfg, omni_from, omni_to = omni_kv_connector if omni_conn_cfg: - from vllm_omni.entrypoints.utils import inject_omni_kv_config - inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) - inject_kv_stage_info(stage_cfg, configured_stage_id) + inject_kv_stage_info(stage_cfg, configured_stage_id, self.stage_configs) if self.single_stage_mode: assert self._omni_master_server is not None stage_clients[stage_idx] = self._launch_diffusion_stage( diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 52e674f47..ab2de757b 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -14,7 +14,9 @@ from vllm.logger import init_logger from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient -from vllm_omni.distributed.omni_connectors.utils.initialization import KV_TRANSFER_PORT_OFFSET +from vllm_omni.distributed.omni_connectors.utils.initialization import ( + KV_TRANSFER_PORT_OFFSET, +) from vllm_omni.engine.stage_init_utils import StageMetadata if TYPE_CHECKING: @@ -246,6 +248,8 @@ class StageEngineCoreClientBase: from_stage = omni_kv_config.get("omni_from_stage", from_stage) try: + # Orchestrator always reports rank-0's port; receiver + # workers add their own local_rank * KV_RANK_PORT_STRIDE. sender_port = int(base_port) + KV_TRANSFER_PORT_OFFSET + int(from_stage) except (TypeError, ValueError): logger.warning( @@ -284,6 +288,7 @@ class StageEngineCoreClientBase: self._kv_sender_host = self._resolve_contact_host() if self._kv_sender_host is None: return None + # rank-0 base port; receiver workers adjust per KV_RANK_PORT_STRIDE. return { "host": self._kv_sender_host, "zmq_port": base_port + kv_transfer_port_offset + int(self.stage_id), diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 3a7fe4bad..c697e34ba 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -13,7 +13,7 @@ import importlib import multiprocessing as mp import os import time -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass from typing import Any, Literal @@ -101,8 +101,110 @@ def resolve_worker_cls(engine_args: dict[str, Any]) -> None: raise ValueError(f"Unknown worker_type: {worker_type}") -def inject_kv_stage_info(stage_cfg: Any, stage_id: int) -> None: - """Inject stage metadata into omni_kv_config when present.""" +def _get_attr_or_item(obj: Any, key: str, default: Any = None) -> Any: + """Read *key* from *obj* regardless of whether it's a dict or object.""" + if hasattr(obj, "get"): + return obj.get(key, default) + return getattr(obj, key, default) + + +def _tp_size_for_stage(stage_configs: Sequence[Any], stage_id: Any) -> int | None: + """Resolve tensor_parallel_size for *stage_id* from the loaded stage configs.""" + id_strs = {str(stage_id)} + try: + id_strs.add(str(int(stage_id))) + except (TypeError, ValueError): + pass + + for stage_cfg in stage_configs: + if str(getattr(stage_cfg, "stage_id", None)) not in id_strs: + continue + engine_args = getattr(stage_cfg, "engine_args", None) + if engine_args is None: + return 1 + parallel_config = _get_attr_or_item(engine_args, "parallel_config") + if parallel_config is not None: + tp = _get_attr_or_item(parallel_config, "tensor_parallel_size", 1) + else: + tp = _get_attr_or_item(engine_args, "tensor_parallel_size", 1) + try: + return max(1, int(tp)) + except (TypeError, ValueError): + return 1 + return None + + +def _inject_inferred_kv_tp_topology( + omni_kv: Any, + stage_id: int, + stage_configs: Sequence[Any], + engine_input_source: Sequence[int] | None = None, +) -> None: + """Infer adjacent-stage TP topology and inject it into omni_kv_config. + + This keeps heterogeneous TP working without requiring user-authored + rank_mapping blocks in config files. + """ + if omni_kv is None: + return + + if hasattr(omni_kv, "get"): + need_send = bool(omni_kv.get("need_send_cache", False)) + need_recv = bool(omni_kv.get("need_recv_cache", False)) + omni_from_stage = omni_kv.get("omni_from_stage") + omni_to_stage = omni_kv.get("omni_to_stage") + rank_mapping = omni_kv.get("rank_mapping") + else: + need_send = bool(getattr(omni_kv, "need_send_cache", False)) + need_recv = bool(getattr(omni_kv, "need_recv_cache", False)) + omni_from_stage = getattr(omni_kv, "omni_from_stage", None) + omni_to_stage = getattr(omni_kv, "omni_to_stage", None) + rank_mapping = getattr(omni_kv, "rank_mapping", None) + + if not need_send and not need_recv: + return + + current_tp = _tp_size_for_stage(stage_configs, stage_id) + if current_tp is None: + return + + peer_stage_id = None + from_tp = None + to_tp = None + if str(omni_from_stage) == str(stage_id): + peer_stage_id = omni_to_stage + from_tp = current_tp + to_tp = _tp_size_for_stage(stage_configs, peer_stage_id) + elif str(omni_to_stage) == str(stage_id): + peer_stage_id = omni_from_stage + from_tp = _tp_size_for_stage(stage_configs, peer_stage_id) + to_tp = current_tp + elif need_recv and engine_input_source: + peer_stage_id = engine_input_source[0] + from_tp = _tp_size_for_stage(stage_configs, peer_stage_id) + to_tp = current_tp + + if from_tp is None or to_tp is None: + return + + if not isinstance(rank_mapping, dict): + rank_mapping = {} + rank_mapping.setdefault("from_tp", int(from_tp)) + rank_mapping.setdefault("to_tp", int(to_tp)) + + if hasattr(omni_kv, "__setitem__"): + omni_kv["rank_mapping"] = rank_mapping + else: + setattr(omni_kv, "rank_mapping", rank_mapping) + + +def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[Any] | None = None) -> None: + """Inject stage_id, engine_input_source, and inferred TP topology into omni_kv_config. + + When *stage_configs* is provided, also infers from_tp/to_tp for + heterogeneous TP topologies so the KV transfer manager can compute + rank mappings automatically. + """ try: engine_args = stage_cfg.engine_args if hasattr(engine_args, "get"): @@ -125,6 +227,14 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int) -> None: omni_kv.setdefault("engine_input_source", list(engine_input_source)) elif hasattr(omni_kv, "__setitem__") and "engine_input_source" not in omni_kv: omni_kv["engine_input_source"] = list(engine_input_source) + + if stage_configs: + _inject_inferred_kv_tp_topology( + omni_kv, + stage_id=stage_id, + stage_configs=stage_configs, + engine_input_source=engine_input_source, + ) except Exception as e: logger.debug("Failed to inject stage info into omni_kv_config: %s", e) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 39fcbc9a0..4b3a7045c 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -86,6 +86,7 @@ from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio from vllm_omni.entrypoints.openai.utils import ( + get_stage_type, get_supported_speakers_from_hf_config, parse_lora_request, validate_requested_speaker, @@ -294,6 +295,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ) num_inference_steps = None + cfg_text_scale = None + cfg_img_scale = None # Omni multistage image generation: Stage-0 (AR) should receive a clean # text prompt (and optional conditioning image/size) so the model's own # processor can construct the correct inputs. @@ -342,6 +345,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): except Exception: pass negative_prompt = extra_body.get("negative_prompt") + cfg_text_scale = extra_body.get("cfg_text_scale") + cfg_img_scale = extra_body.get("cfg_img_scale") engine_prompt_image: dict[str, Any] | None = None is_img2img = False @@ -397,14 +402,18 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): sampling_params_list = self._build_sampling_params_list_from_request(request) # Apply user-specified overrides to diffusion stage(s) for image generation - if _image_gen_height is not None or _image_gen_width is not None or num_inference_steps is not None: - for idx, sp in enumerate(sampling_params_list): - if hasattr(sp, "height") and _image_gen_height is not None: - sp.height = _image_gen_height - if hasattr(sp, "width") and _image_gen_width is not None: - sp.width = _image_gen_width - if hasattr(sp, "num_inference_steps") and num_inference_steps is not None: - sp.num_inference_steps = num_inference_steps + for idx, sp in enumerate(sampling_params_list): + if hasattr(sp, "height") and _image_gen_height is not None: + sp.height = _image_gen_height + if hasattr(sp, "width") and _image_gen_width is not None: + sp.width = _image_gen_width + if hasattr(sp, "num_inference_steps") and num_inference_steps is not None: + sp.num_inference_steps = num_inference_steps + if hasattr(sp, "extra_args") and sp.extra_args is not None: + if cfg_text_scale is not None: + sp.extra_args["cfg_text_scale"] = cfg_text_scale + if cfg_img_scale is not None: + sp.extra_args["cfg_img_scale"] = cfg_img_scale self._log_inputs( request_id, @@ -2108,6 +2117,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): num_inference_steps = extra_body.get("num_inference_steps") guidance_scale = extra_body.get("guidance_scale") true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale") + cfg_text_scale = extra_body.get("cfg_text_scale") + cfg_img_scale = extra_body.get("cfg_img_scale") seed = extra_body.get("seed") negative_prompt = extra_body.get("negative_prompt") num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1) @@ -2162,6 +2173,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): gen_params.guidance_scale = guidance_scale if true_cfg_scale is not None: gen_params.true_cfg_scale = true_cfg_scale + if cfg_text_scale is not None: + gen_params.extra_args["cfg_text_scale"] = cfg_text_scale + if cfg_img_scale is not None: + gen_params.extra_args["cfg_img_scale"] = cfg_img_scale if num_frames is not None: gen_params.num_frames = num_frames if guidance_scale_2 is not None: @@ -2206,10 +2221,30 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Generate image diffusion_engine = cast(AsyncOmni, self._diffusion_engine) + stage_configs = list(getattr(diffusion_engine, "stage_configs", []) or []) + default_params_list = list(getattr(diffusion_engine, "default_sampling_params_list", []) or []) + + sampling_params_list: list[Any] = [] + for idx, stage_cfg in enumerate(stage_configs): + if get_stage_type(stage_cfg) == "diffusion": + sampling_params_list.append(gen_params) + continue + + default_stage_params = default_params_list[idx] if idx < len(default_params_list) else SamplingParams() + if hasattr(default_stage_params, "clone"): + try: + default_stage_params = default_stage_params.clone() + except Exception: + pass + sampling_params_list.append(default_stage_params) + + if not sampling_params_list: + sampling_params_list = [gen_params] + result = None async for output in diffusion_engine.generate( prompt=gen_prompt, - sampling_params_list=[gen_params], # Pass as single-stage params + sampling_params_list=sampling_params_list, request_id=request_id, ): result = output diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 85faf6b94..e4c33a58c 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -267,6 +267,10 @@ class OmniDiffusionSamplingParams: cfg_text_kv_metadata: dict[str, Any] | None = None cfg_img_kv_metadata: dict[str, Any] | None = None cfg_kv_request_ids: dict[str, str] | None = None + cfg_active_branch: str | None = None + cfg_branch_roles: list[str] | None = None + cfg_branch_past_key_values: dict[str, Any] | None = None + cfg_branch_kv_metadata: dict[str, dict[str, Any]] | None = None # Component modules modules: dict[str, Any] = field(default_factory=dict)