From c3c736d488264d340e049e79759927be61641888 Mon Sep 17 00:00:00 2001 From: Jinheng Date: Wed, 8 Apr 2026 14:55:35 +0800 Subject: [PATCH] [Feature] Port Bagel RDMA flow to latest main (#2000) Signed-off-by: Jinheng Li Signed-off-by: natureofnature Co-authored-by: ahengljh Co-authored-by: natureofnature --- .../mooncake_transfer_engine_connector.md | 16 +- tests/diffusion/test_stage_diffusion_proc.py | 65 +++ .../omni_connectors/test_basic_connectors.py | 58 ++ .../omni_connectors/test_kv_flow.py | 227 +++++++- .../test_orchestrator_kv_sender_info.py | 207 +++++++ vllm_omni/diffusion/request.py | 2 + vllm_omni/diffusion/stage_diffusion_client.py | 12 +- vllm_omni/diffusion/stage_diffusion_proc.py | 38 +- .../mooncake_transfer_engine_connector.py | 10 +- .../omni_connectors/kv_transfer_manager.py | 518 ++++++++++++++++-- .../omni_connectors/utils/initialization.py | 85 ++- vllm_omni/engine/async_omni_engine.py | 8 +- vllm_omni/engine/orchestrator.py | 43 +- vllm_omni/engine/stage_engine_core_client.py | 121 ++++ .../model_executor/stage_configs/bagel.yaml | 27 + 15 files changed, 1374 insertions(+), 63 deletions(-) create mode 100644 tests/diffusion/test_stage_diffusion_proc.py create mode 100644 tests/engine/test_orchestrator_kv_sender_info.py diff --git a/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md b/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md index 798644b96..306a0620b 100644 --- a/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md +++ b/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md @@ -33,8 +33,8 @@ runtime: zmq_port: 50051 # ZMQ base port (see "Port Offset Scheme" below) protocol: "rdma" # "rdma" or "tcp" device_name: "" # RDMA device (e.g., "mlx5_0"), empty for auto-detect - memory_pool_size: 2147483648 # 2GB memory pool - memory_pool_device: "cpu" # "cpu" for pinned memory, "cuda" for GPUDirect RDMA + memory_pool_size: 4294967296 # 4 GB (CPU); use 2147483648 (2 GB) for GPU + memory_pool_device: "cpu" # "cpu" for pinned memory (recommended), "cuda" for GPUDirect RDMA ``` Wire stages to the connector: @@ -64,8 +64,8 @@ stage_args: | Parameter | Default | Description | |---|---|---| -| `memory_pool_size` | 1 GB | Total size of the RDMA-registered memory pool in bytes. | -| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity). | +| `memory_pool_size` | 4 GB (CPU) / 2 GB (GPU) | Total size of the RDMA-registered memory pool in bytes. Recommended 4 GB for CPU pinned memory; 2 GB for GPU VRAM to conserve device memory. | +| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended, works on all topologies). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity, PIX topology). | ### Networking @@ -107,10 +107,10 @@ receiver_connect = remote_side_channel_port + tp_rank ## Memory Pool Modes -| Mode | Config | Data Flow | Best For | -|---|---|---|---| -| CPU Pinned | `memory_pool_device: "cpu"` | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) | -| GPUDirect | `memory_pool_device: "cuda"` | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) | +| Mode | Config | Recommended Pool Size | Data Flow | Best For | +|---|---|---|---|---| +| CPU Pinned | `memory_pool_device: "cpu"` | 4 GB | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) | +| GPUDirect | `memory_pool_device: "cuda"` | 2 GB | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) | > **Note**: GPUDirect RDMA requires the NIC and GPU to share a direct PCIe > switch (PIX topology). On systems where they are connected via PXB or NODE, diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py new file mode 100644 index 000000000..c26070ad4 --- /dev/null +++ b/tests/diffusion/test_stage_diffusion_proc.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +from concurrent.futures import ThreadPoolExecutor +from dataclasses import asdict +from types import SimpleNamespace + +import pytest + +from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info(): + async def run_test(): + captured = {} + + def step(request): + captured["request"] = request + return [ + SimpleNamespace( + images=["img-1"], + _multimodal_output={}, + metrics={}, + stage_durations={}, + peak_memory_mb=0.0, + latents=None, + final_output_type="image", + ), + SimpleNamespace( + images=["img-2"], + _multimodal_output={}, + metrics={}, + stage_durations={}, + peak_memory_mb=0.0, + latents=None, + final_output_type="image", + ), + ] + + proc = object.__new__(StageDiffusionProc) + proc._engine = SimpleNamespace(step=step) + proc._executor = ThreadPoolExecutor(max_workers=1) + + try: + result = await proc._process_batch_request( + request_id="req-parent", + prompts=["hello", "world"], + sampling_params_dict=asdict(OmniDiffusionSamplingParams()), + kv_sender_info={0: {"host": "10.0.0.2", "zmq_port": 50151}}, + ) + finally: + proc._executor.shutdown(wait=True) + + request = captured["request"] + assert request.request_id == "req-parent" + assert request.request_ids == ["req-parent-0", "req-parent-1"] + assert request.kv_sender_info == {0: {"host": "10.0.0.2", "zmq_port": 50151}} + assert result.request_id == "req-parent" + assert result.images == ["img-1", "img-2"] + + asyncio.run(run_test()) diff --git a/tests/distributed/omni_connectors/test_basic_connectors.py b/tests/distributed/omni_connectors/test_basic_connectors.py index 1b1965355..bca96e790 100644 --- a/tests/distributed/omni_connectors/test_basic_connectors.py +++ b/tests/distributed/omni_connectors/test_basic_connectors.py @@ -120,3 +120,61 @@ def test_get_invalid_metadata(shm_connector): result = shm_connector.get("stage_0", "stage_1", "req_3", {"unknown": "format"}) assert result is None + + +def test_mooncake_connector_defaults_missing_host_to_detected_ip(monkeypatch: pytest.MonkeyPatch): + import vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector as mooncake_module + + class _FakePool: + is_cuda = False + + def pin_memory(self): + return self + + def data_ptr(self): + return 1234 + + class _FakeTransferEngine: + def initialize(self, host, mode, protocol, device_name): + self.host = host + self.mode = mode + self.protocol = protocol + self.device_name = device_name + return 0 + + def get_rpc_port(self): + return 23456 + + def register_memory(self, base_ptr, pool_size): + del base_ptr, pool_size + return 0 + + def unregister_memory(self, base_ptr): + del base_ptr + return 0 + + monkeypatch.setattr(mooncake_module, "TransferEngine", _FakeTransferEngine) + monkeypatch.setattr(mooncake_module.torch, "empty", lambda *args, **kwargs: _FakePool()) + monkeypatch.setattr( + mooncake_module.MooncakeTransferEngineConnector, + "_get_local_ip", + lambda self: "10.20.30.40", + ) + monkeypatch.setattr( + mooncake_module.MooncakeTransferEngineConnector, + "_zmq_listener_loop", + lambda self: self._listener_ready.set(), + ) + + connector = mooncake_module.MooncakeTransferEngineConnector( + { + "zmq_port": 50051, + "memory_pool_size": 4096, + } + ) + try: + assert connector.host == "10.20.30.40" + assert connector.engine.host == "10.20.30.40" + assert connector.get_connection_info()["host"] == "10.20.30.40" + finally: + connector.close() diff --git a/tests/distributed/omni_connectors/test_kv_flow.py b/tests/distributed/omni_connectors/test_kv_flow.py index b12fc013b..cea186019 100644 --- a/tests/distributed/omni_connectors/test_kv_flow.py +++ b/tests/distributed/omni_connectors/test_kv_flow.py @@ -1,8 +1,14 @@ +import json +import struct + +import numpy as np import pytest import torch +import vllm_omni.distributed.omni_connectors.kv_transfer_manager as kv_transfer_manager_module from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + KVCacheTransferData, OmniKVCacheConfig, OmniKVTransferManager, ) @@ -60,6 +66,35 @@ def common_constants(): } +def _decode_stored_payload(data): + if isinstance(data, torch.Tensor) and data.dtype == torch.uint8 and data.dim() == 1: + return KVCacheTransferData.from_bytes(data.cpu().numpy().tobytes()) + + if isinstance(data, (bytes, bytearray, memoryview)): + return KVCacheTransferData.from_bytes(data) + + return data + + +def _make_serialized_payload() -> tuple[bytes, torch.Tensor]: + key_tensor = torch.arange(12, dtype=torch.float32).reshape(3, 4) + payload = KVCacheTransferData( + request_id="req-payload", + layer_blocks={"key_cache": [key_tensor], "value_cache": [None]}, + block_ids=[1], + metadata={"seq_len": 3}, + ).to_bytes() + return payload, key_tensor + + +def _rewrite_serialized_header(payload: bytes, mutate_header) -> bytes: + header_len = struct.unpack(">I", payload[:4])[0] + header = json.loads(payload[4 : 4 + header_len]) + mutate_header(header) + new_header = json.dumps(header, separators=(",", ":")).encode("utf-8") + return struct.pack(">I", len(new_header)) + new_header + payload[4 + header_len :] + + def test_manager_extraction(kv_config, mock_connector, common_constants): """Test extraction and sending logic in OmniKVTransferManager.""" num_layers = common_constants["num_layers"] @@ -95,7 +130,7 @@ def test_manager_extraction(kv_config, mock_connector, common_constants): expected_key = f"stage1->stage2:{full_request_id}" assert expected_key in mock_connector.store - data = mock_connector.store[expected_key] + data = _decode_stored_payload(mock_connector.store[expected_key]) assert data["request_id"] == req_id assert "layer_blocks" in data assert len(data["layer_blocks"]["key_cache"]) == num_layers @@ -106,6 +141,116 @@ def test_manager_extraction(kv_config, mock_connector, common_constants): assert data["layer_blocks"]["key_cache"][0].shape == expected_shape +def test_from_bytes_rejects_out_of_bounds_header_len(): + payload, _ = _make_serialized_payload() + bad_payload = struct.pack(">I", len(payload)) + payload[4:] + + with pytest.raises(ValueError, match="header_len"): + KVCacheTransferData.from_bytes(bad_payload) + + with pytest.raises(ValueError, match="header_len"): + KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8)) + + +def test_from_bytes_rejects_out_of_bounds_tensor_span(): + payload, _ = _make_serialized_payload() + bad_payload = _rewrite_serialized_header(payload, lambda header: header["td"][0].update({"o": 4096})) + + with pytest.raises(ValueError, match="tensor span"): + KVCacheTransferData.from_bytes(bad_payload) + + with pytest.raises(ValueError, match="tensor span"): + KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8)) + + +def test_from_bytes_rejects_unsupported_dtype(): + payload, _ = _make_serialized_payload() + bad_payload = _rewrite_serialized_header(payload, lambda header: header["td"][0].update({"d": "cuda"})) + + with pytest.raises(ValueError, match="Unsupported dtype"): + KVCacheTransferData.from_bytes(bad_payload) + + with pytest.raises(ValueError, match="Unsupported dtype"): + KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8)) + + +def test_from_bytes_uses_explicit_layer_index_descriptor(): + payload, key_tensor = _make_serialized_payload() + payload_with_explicit_index = _rewrite_serialized_header( + payload, + lambda header: header["td"][0].update({"n": "key_cache_extra_suffix", "i": 0}), + ) + + data = KVCacheTransferData.from_bytes(payload_with_explicit_index) + + assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor) + + +def test_update_sender_info_uses_configured_source_stage(): + config = OmniKVCacheConfig( + connector_config={"type": "mock"}, + stage_id=2, + engine_input_source=[1], + need_recv_cache=True, + ) + manager = OmniKVTransferManager(config) + + manager.update_sender_info( + { + 0: {"host": "10.0.0.1", "zmq_port": 50151}, + 1: {"host": "10.0.0.2", "zmq_port": 50152}, + } + ) + + assert manager.config.connector_config["sender_host"] == "10.0.0.2" + assert manager.config.connector_config["sender_zmq_port"] == 50152 + + +def test_clone_received_payload_tensors_breaks_buffer_alias(): + payload, key_tensor = _make_serialized_payload() + raw = np.frombuffer(bytearray(payload), dtype=np.uint8) + data = KVCacheTransferData.from_bytes(memoryview(raw)) + + OmniKVTransferManager._clone_received_payload_tensors(data) + raw[:] = 0 + + assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor) + + +def test_receive_kv_cache_uses_exponential_backoff(monkeypatch): + config = OmniKVCacheConfig( + connector_config={"type": "mock"}, + from_stage="sender", + stage_id="receiver", + need_recv_cache=True, + recv_timeout=0.3, + ) + manager = OmniKVTransferManager(config) + + class _NeverReadyConnector: + def get(self, **kwargs): + del kwargs + return None + + manager._connector = _NeverReadyConnector() + + now = {"value": 0.0} + sleep_intervals = [] + + monkeypatch.setattr(kv_transfer_manager_module.time, "time", lambda: now["value"]) + + def _fake_sleep(interval: float) -> None: + sleep_intervals.append(interval) + now["value"] += interval + + monkeypatch.setattr(kv_transfer_manager_module.time, "sleep", _fake_sleep) + + data, size = manager.receive_kv_cache_for_request("req-backoff") + + assert (data, size) == (None, 0) + assert sleep_intervals == pytest.approx([0.01, 0.02, 0.04, 0.08, 0.16]) + + def test_manager_extraction_tuple_layout(kv_config, mock_connector, common_constants): """Test extraction with tuple layout.""" num_layers = common_constants["num_layers"] @@ -135,7 +280,7 @@ def test_manager_extraction_tuple_layout(kv_config, mock_connector, common_const expected_key = f"stage1->stage2:{full_request_id}" assert expected_key in mock_connector.store - data = mock_connector.store[expected_key] + data = _decode_stored_payload(mock_connector.store[expected_key]) expected_shape = (seq_len, num_heads, head_dim) for idx in range(len(kv_caches)): assert data["layer_blocks"]["key_cache"][idx].shape == expected_shape @@ -165,7 +310,7 @@ def test_manager_extraction_mismatched_kv_block_counts(kv_config, mock_connector expected_key = f"stage1->stage2:{full_request_id}" assert expected_key in mock_connector.store - data = mock_connector.store[expected_key] + data = _decode_stored_payload(mock_connector.store[expected_key]) expected_shape = (2 * block_size, num_heads, head_dim) assert data["layer_blocks"]["key_cache"][0].shape == expected_shape assert data["layer_blocks"]["value_cache"][0].shape == expected_shape @@ -254,6 +399,82 @@ def test_manager_reception(kv_config, mock_connector, common_constants): assert req.kv_metadata["seq_len"] == seq_len +def test_manager_reception_prefers_parent_request_id_for_batched_request(kv_config, mock_connector, common_constants): + """Batched diffusion requests must fetch KV using the parent/global request ID.""" + num_layers = common_constants["num_layers"] + num_heads = common_constants["num_heads"] + head_dim = common_constants["head_dim"] + seq_len = common_constants["seq_len"] + parent_req_id = common_constants["req_id"] + + expected_shape = (seq_len, num_heads, head_dim) + key_cache = [torch.randn(expected_shape) for _ in range(num_layers)] + value_cache = [torch.randn(expected_shape) for _ in range(num_layers)] + + data_to_receive = { + "request_id": parent_req_id, + "layer_blocks": {"key_cache": key_cache, "value_cache": value_cache}, + "metadata": {"seq_len": seq_len}, + "block_ids": [], + } + + manager = OmniKVTransferManager(kv_config) + manager._connector = mock_connector + + full_request_id = f"omni_stage1_to_stage2_kv_cache_{parent_req_id}" + store_key = f"stage1->stage2:{full_request_id}" + mock_connector.store[store_key] = data_to_receive + + req = OmniDiffusionRequest( + prompts=["prompt-a", "prompt-b"], + sampling_params=OmniDiffusionSamplingParams(), + request_ids=[f"{parent_req_id}-0", f"{parent_req_id}-1"], + request_id=parent_req_id, + ) + + success = manager.receive_kv_cache(req, target_device=torch.device("cpu")) + + assert success + assert req.kv_metadata["seq_len"] == seq_len + assert torch.allclose(req.past_key_values.key_cache[0], key_cache[0]) + + +def test_receive_multi_kv_cache_uses_parent_request_id_for_cfg_collection(kv_config): + manager = OmniKVTransferManager(kv_config) + + seen = {} + + def collect_cfg(request_id, cfg_request_ids, kv_transfer_manager, target_device): + seen["request_id"] = request_id + seen["cfg_request_ids"] = cfg_request_ids + seen["kv_transfer_manager"] = kv_transfer_manager + seen["target_device"] = target_device + return {"cfg_text_kv_metadata": {"ok": True}} + + req = OmniDiffusionRequest( + prompts=["prompt-a", "prompt-b"], + sampling_params=OmniDiffusionSamplingParams(), + request_ids=["req-parent-0", "req-parent-1"], + request_id="req-parent", + ) + req.sampling_params.cfg_kv_request_ids = {"cfg_text": "req-parent__cfg_text"} + + manager.receive_kv_cache = lambda request, target_device=None: request is req + + success = manager.receive_multi_kv_cache( + req, + cfg_kv_collect_func=collect_cfg, + target_device=torch.device("cpu"), + ) + + assert success + assert seen["request_id"] == "req-parent" + assert seen["cfg_request_ids"] == {"cfg_text": "req-parent__cfg_text"} + assert seen["kv_transfer_manager"] is manager + assert seen["target_device"] == torch.device("cpu") + assert req.sampling_params.cfg_text_kv_metadata == {"ok": True} + + def test_integration_flow(common_constants): """Simulate extraction -> connector -> reception.""" num_layers = common_constants["num_layers"] diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py new file mode 100644 index 000000000..94da4ce71 --- /dev/null +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -0,0 +1,207 @@ +import asyncio +from types import SimpleNamespace + +import pytest +from vllm import SamplingParams + +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState +from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _DummySenderStage: + def __init__(self, sender_info): + self._sender_info = sender_info + self.engine_outputs = None + + def set_engine_outputs(self, outputs): + self.engine_outputs = outputs + + def get_kv_sender_info(self): + return self._sender_info + + +class _DummyDiffusionStage: + stage_type = "diffusion" + custom_process_input_func = None + + def __init__(self, engine_input_source=None): + self.engine_input_source = engine_input_source or [0] + self.calls = [] + + async def add_request_async(self, request_id, prompt, sampling_params, kv_sender_info=None): + self.calls.append( + { + "request_id": request_id, + "prompt": prompt, + "sampling_params": sampling_params, + "kv_sender_info": kv_sender_info, + } + ) + + +def test_stage_engine_core_client_builds_kv_sender_info_from_tcp_address(): + client = object.__new__(StageEngineCoreClient) + client.stage_id = 0 + client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"} + client._omni_kv_config = None + client._kv_sender_info = None + client._kv_sender_initialized = False + client._kv_sender_host = client._resolve_contact_host() + client._initialize_kv_sender_endpoint() + + assert client.get_kv_sender_info() == { + "host": "10.20.30.40", + "zmq_port": 50151, + } + + +def test_stage_engine_core_client_falls_back_to_detected_ip_for_loopback(monkeypatch): + client = object.__new__(StageEngineCoreClient) + client.stage_id = 1 + client.client_addresses = {"input_address": "tcp://127.0.0.1:1234"} + client._omni_kv_config = None + client._kv_sender_info = None + client._kv_sender_initialized = False + monkeypatch.setattr(client, "_detect_local_ip", lambda: "192.168.0.12") + client._kv_sender_host = client._resolve_contact_host() + client._initialize_kv_sender_endpoint() + + assert client.get_kv_sender_info() == { + "host": "192.168.0.12", + "zmq_port": 50152, + } + + +def test_stage_engine_core_client_uses_connector_config_for_sender_port(): + client = object.__new__(StageEngineCoreClient) + client.stage_id = 3 + client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"} + client._kv_sender_info = None + client._kv_sender_initialized = False + client._omni_kv_config = { + "omni_from_stage": "3", + "connector_config": { + "type": "MooncakeTransferEngineConnector", + "role": "sender", + "host": "10.20.30.99", + "zmq_port": 51000, + }, + } + client._kv_sender_host = client._resolve_contact_host() + client._initialize_kv_sender_endpoint() + + assert client.get_kv_sender_info() == { + "host": "10.20.30.99", + "zmq_port": 51103, + } + + +def test_stage_engine_core_client_preserves_explicit_loopback_sender_host(): + client = object.__new__(StageEngineCoreClient) + client.stage_id = 2 + client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"} + client._kv_sender_info = None + client._kv_sender_initialized = False + client._omni_kv_config = { + "omni_from_stage": "2", + "connector_config": { + "type": "MooncakeTransferEngineConnector", + "role": "sender", + "host": "127.0.0.1", + "zmq_port": 51000, + }, + } + client._kv_sender_host = client._resolve_contact_host() + client._initialize_kv_sender_endpoint() + + assert client.get_kv_sender_info() == { + "host": "127.0.0.1", + "zmq_port": 51102, + } + + +def test_forward_to_diffusion_attaches_kv_sender_info(): + orchestrator = object.__new__(Orchestrator) + sender_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151}) + diffusion_stage = _DummyDiffusionStage(engine_input_source=[0]) + + orchestrator.num_stages = 2 + orchestrator.stage_clients = [sender_stage, diffusion_stage] + orchestrator._companion_map = {} + orchestrator.stage_vllm_configs = [None, None] + orchestrator.output_processors = [None, None] + + params = OmniDiffusionSamplingParams() + req_state = OrchestratorRequestState( + request_id="req-1", + prompt={"prompt": "hello"}, + sampling_params_list=[SamplingParams(max_tokens=4), params], + final_stage_id=1, + ) + + output = SimpleNamespace(request_id="req-1", finished=True) + asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", 0, output, req_state)) + + assert sender_stage.engine_outputs == [output] + assert diffusion_stage.calls[0]["request_id"] == "req-1" + assert diffusion_stage.calls[0]["kv_sender_info"] == { + 0: {"host": "10.0.0.2", "zmq_port": 50151}, + } + assert req_state.stage_submit_ts[1] > 0 + + +def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info(): + orchestrator = object.__new__(Orchestrator) + source_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151}) + previous_stage = _DummySenderStage({"host": "10.0.0.9", "zmq_port": 59999}) + diffusion_stage = _DummyDiffusionStage(engine_input_source=[0]) + + orchestrator.num_stages = 3 + orchestrator.stage_clients = [source_stage, previous_stage, diffusion_stage] + orchestrator._companion_map = {} + orchestrator.stage_vllm_configs = [None, None, None] + orchestrator.output_processors = [None, None, None] + + params = OmniDiffusionSamplingParams() + req_state = OrchestratorRequestState( + request_id="req-3", + prompt={"prompt": "hello"}, + sampling_params_list=[SamplingParams(max_tokens=4), SamplingParams(max_tokens=4), params], + final_stage_id=2, + ) + + output = SimpleNamespace(request_id="req-3", finished=True) + asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", 1, output, req_state)) + + assert previous_stage.engine_outputs == [output] + assert diffusion_stage.calls[0]["kv_sender_info"] == { + 0: {"host": "10.0.0.2", "zmq_port": 50151}, + } + + +def test_prewarm_diffusion_attaches_kv_sender_info(): + orchestrator = object.__new__(Orchestrator) + sender_stage = _DummySenderStage({"host": "10.0.0.3", "zmq_port": 50151}) + diffusion_stage = _DummyDiffusionStage(engine_input_source=[0]) + + orchestrator.stage_clients = [sender_stage, diffusion_stage] + orchestrator.num_stages = 2 + + req_state = OrchestratorRequestState( + request_id="req-2", + prompt={"prompt": "hello"}, + sampling_params_list=[SamplingParams(max_tokens=4), OmniDiffusionSamplingParams()], + final_stage_id=1, + ) + + stage0_request = SimpleNamespace(prompt_token_ids=[1, 2, 3]) + asyncio.run(Orchestrator._prewarm_async_chunk_stages(orchestrator, "req-2", stage0_request, req_state)) + + assert diffusion_stage.calls[0]["request_id"] == "req-2" + assert diffusion_stage.calls[0]["kv_sender_info"] == { + 0: {"host": "10.0.0.3", "zmq_port": 50151}, + } + assert req_state.stage_submit_ts[1] > 0 diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index 1d6d64905..4d4328d25 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -26,6 +26,8 @@ class OmniDiffusionRequest: sampling_params: OmniDiffusionSamplingParams request_ids: list[str] = field(default_factory=list) + request_id: str | None = None + kv_sender_info: dict | None = None def __post_init__(self): """Initialize dependent fields after dataclass initialization.""" diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 77db2b1b9..a1a4766de 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -179,6 +179,7 @@ class StageDiffusionClient: request_id: str, prompt: OmniPromptType, sampling_params: OmniDiffusionSamplingParams, + kv_sender_info: dict[int, dict[str, Any]] | None = None, ) -> None: self._request_socket.send( self._encoder.encode( @@ -187,6 +188,7 @@ class StageDiffusionClient: "request_id": request_id, "prompt": prompt, "sampling_params": self._sampling_params_to_dict(sampling_params), + "kv_sender_info": kv_sender_info, } ) ) @@ -198,6 +200,7 @@ class StageDiffusionClient: request_id: str, prompts: list[OmniPromptType], sampling_params: OmniDiffusionSamplingParams, + kv_sender_info: dict[int, dict[str, Any]] | None = None, ) -> None: """Submit a list of prompts as a single batched engine call. @@ -206,7 +209,12 @@ class StageDiffusionClient: *request_id*. """ task = asyncio.create_task( - self._run_batch(request_id, prompts, sampling_params), + self._run_batch( + request_id, + prompts, + sampling_params, + kv_sender_info, + ), name=f"diffusion-batch-{request_id}", ) self._tasks[request_id] = task @@ -216,6 +224,7 @@ class StageDiffusionClient: request_id: str, prompts: list[OmniPromptType], sampling_params: OmniDiffusionSamplingParams, + kv_sender_info: dict[int, dict[str, Any]] | None = None, ) -> None: try: self._request_socket.send( @@ -225,6 +234,7 @@ class StageDiffusionClient: "request_id": request_id, "prompts": prompts, "sampling_params": self._sampling_params_to_dict(sampling_params), + "kv_sender_info": kv_sender_info, } ) ) diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index bcc3bef15..9d8c06cce 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -130,6 +130,7 @@ class StageDiffusionProc: request_id: str, prompt: Any, sampling_params_dict: dict, + kv_sender_info: dict[str, Any] | None = None, ) -> OmniRequestOutput: """Build a diffusion request and run DiffusionEngine.step().""" sampling_params = self._reconstruct_sampling_params(sampling_params_dict) @@ -138,6 +139,8 @@ class StageDiffusionProc: prompts=[prompt], sampling_params=sampling_params, request_ids=[request_id], + request_id=request_id, + kv_sender_info=kv_sender_info, ) loop = asyncio.get_running_loop() @@ -152,6 +155,7 @@ class StageDiffusionProc: request_id: str, prompts: list[Any], sampling_params_dict: dict, + kv_sender_info: dict[str, Any] | None = None, ) -> OmniRequestOutput: """Build a batched diffusion request and run DiffusionEngine.step(). @@ -165,7 +169,9 @@ class StageDiffusionProc: request = OmniDiffusionRequest( prompts=prompts, sampling_params=sampling_params, - request_ids=[request_id] * len(prompts), + request_ids=[f"{request_id}-{i}" for i in range(len(prompts))], + request_id=request_id, + kv_sender_info=kv_sender_info, ) loop = asyncio.get_running_loop() @@ -346,10 +352,20 @@ class StageDiffusionProc: tasks: dict[str, asyncio.Task] = {} - async def _dispatch_request(request_id: str, prompt: Any, sampling_params_dict: dict) -> None: + async def _dispatch_request( + request_id: str, + prompt: Any, + sampling_params_dict: dict, + kv_sender_info: dict[str, Any] | None = None, + ) -> None: """Process a single diffusion request and send the response.""" try: - result = await self._process_request(request_id, prompt, sampling_params_dict) + result = await self._process_request( + request_id, + prompt, + sampling_params_dict, + kv_sender_info=kv_sender_info, + ) await response_socket.send(encoder.encode({"type": "result", "output": result})) except DiffusionRequestAbortedError as e: logger.info( @@ -384,6 +400,7 @@ class StageDiffusionProc: request_id, msg["prompt"], msg["sampling_params"], + msg.get("kv_sender_info"), ) ) tasks[request_id] = task @@ -391,9 +408,19 @@ class StageDiffusionProc: elif msg_type == "add_batch_request": request_id = msg["request_id"] - async def _dispatch_batch(rid: str, prompts: list, sp_dict: dict) -> None: + async def _dispatch_batch( + rid: str, + prompts: list, + sp_dict: dict, + kv_sender_info: dict[str, Any] | None = None, + ) -> None: try: - result = await self._process_batch_request(rid, prompts, sp_dict) + result = await self._process_batch_request( + rid, + prompts, + sp_dict, + kv_sender_info=kv_sender_info, + ) await response_socket.send(encoder.encode({"type": "result", "output": result})) except DiffusionRequestAbortedError as e: logger.info( @@ -420,6 +447,7 @@ class StageDiffusionProc: request_id, msg["prompts"], msg["sampling_params"], + msg.get("kv_sender_info"), ) ) tasks[request_id] = task diff --git a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py index b1dc8b898..96a528963 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py @@ -277,13 +277,15 @@ class MooncakeTransferEngineConnector(OmniConnectorBase): } self.config = config - host_config = config.get("host", "127.0.0.1") - # Support "auto" to auto-detect local IP address - if host_config.lower() == "auto": + host_config = config.get("host") + host_value = "auto" if host_config is None else str(host_config) + # Default sender/receiver bootstrap to a routable local IP so the + # advertised endpoint matches the interface Mooncake binds. + if host_value.lower() == "auto" or host_value in {"", "*", "0.0.0.0", "::"}: self.host = self._get_local_ip() logger.info(f"Auto-detected local IP for RDMA: {self.host}") else: - self.host = host_config + self.host = host_value self.zmq_port = config.get("zmq_port", 50051) self.protocol = config.get("protocol", "rdma") diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 1f4938438..1958c9d40 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unified OmniConnector and KV cache transfer management.""" +import json +import struct import time from collections.abc import Callable from dataclasses import asdict, dataclass @@ -12,12 +14,36 @@ 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 logger = init_logger(__name__) LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor] +_SAFE_TORCH_DTYPES = { + name: dtype + for name in ( + "bool", + "uint8", + "int8", + "int16", + "int32", + "int64", + "float16", + "float32", + "float64", + "bfloat16", + "complex64", + "complex128", + "float8_e4m3fn", + "float8_e4m3fnuz", + "float8_e5m2", + "float8_e5m2fnuz", + ) + if isinstance((dtype := getattr(torch, name, None)), torch.dtype) +} + @dataclass class OmniKVCacheConfig: @@ -46,6 +72,242 @@ class KVCacheTransferData: """Convert to dictionary for serialization.""" return asdict(self) + def to_bytes(self) -> bytes: + """Convert to compact binary format for fast transfer.""" + 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] = [] + 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): + 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: + 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, + "s": list(t.shape), + "o": data_offset, + "b": nbytes, + } + ) + gpu_tensors.append(t.view(torch.uint8).flatten()) + data_offset += nbytes + + if device is None: + raise RuntimeError("No CUDA tensors found, use to_bytes() instead") + + 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") + + 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) + + pos = len(header_prefix) + for t_flat in gpu_tensors: + n = t_flat.numel() + output[pos : pos + n].copy_(t_flat) + pos += n + + return output + + @staticmethod + def _load_header_from_memoryview(raw_mv: memoryview) -> tuple[dict[str, Any], memoryview]: + if len(raw_mv) < 4: + raise ValueError("Corrupted KV payload: missing 4-byte header length") + + header_len = struct.unpack(">I", raw_mv[:4])[0] + if header_len > len(raw_mv) - 4: + raise ValueError(f"Corrupted KV payload: header_len={header_len} exceeds buffer size={len(raw_mv)}") + + return json.loads(bytes(raw_mv[4 : 4 + header_len])), raw_mv[4 + header_len :] + + @staticmethod + def _load_header_from_tensor(gpu_tensor: torch.Tensor) -> tuple[dict[str, Any], int]: + if gpu_tensor.dtype != torch.uint8 or gpu_tensor.dim() != 1: + raise ValueError("Packed GPU KV payload must be a 1-D uint8 tensor") + + total_bytes = int(gpu_tensor.numel()) + if total_bytes < 4: + raise ValueError("Corrupted KV payload: missing 4-byte header length") + + header_len = struct.unpack(">I", gpu_tensor[:4].cpu().numpy().tobytes())[0] + if header_len > total_bytes - 4: + raise ValueError(f"Corrupted KV payload: header_len={header_len} exceeds buffer size={total_bytes}") + + header_bytes = gpu_tensor[4 : 4 + header_len].cpu().numpy().tobytes() + return json.loads(header_bytes), 4 + header_len + + @staticmethod + def _validate_tensor_span(name: str, info: dict[str, Any], tensor_data_bytes: int) -> tuple[int, int]: + offset = info["o"] + nbytes = info["b"] + if offset < 0 or nbytes < 0 or offset + nbytes > tensor_data_bytes: + raise ValueError( + f"Corrupted KV payload tensor span for {name}: " + f"offset={offset}, bytes={nbytes}, tensor_data_bytes={tensor_data_bytes}" + ) + return offset, nbytes + + @staticmethod + def _resolve_torch_dtype(dtype_name: Any) -> torch.dtype: + torch_dtype = _SAFE_TORCH_DTYPES.get(str(dtype_name)) + if torch_dtype is None: + raise ValueError(f"Unsupported dtype in KV payload: {dtype_name}") + return torch_dtype + + @staticmethod + def _resolve_layer_idx(info: dict[str, Any], num_layers: int) -> int: + layer_idx = info.get("i") + if layer_idx is None: + name = info.get("n") + if isinstance(name, str) and name.startswith("key_cache_"): + layer_idx = int(name.removeprefix("key_cache_")) + elif isinstance(name, str) and name.startswith("value_cache_"): + layer_idx = int(name.removeprefix("value_cache_")) + else: + raise ValueError(f"Invalid KV tensor name in payload: {name}") + + if not isinstance(layer_idx, int): + raise ValueError(f"Invalid layer index in KV payload: {layer_idx}") + if layer_idx < 0 or layer_idx >= num_layers: + raise ValueError(f"Invalid layer index in KV payload: {layer_idx} (num_layers={num_layers})") + 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) + + num_layers = header["nl"] + key_cache: list[torch.Tensor | None] = [None] * num_layers + value_cache: list[torch.Tensor | None] = [None] * num_layers + + 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"]) + ) + 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"], + } + + @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) + + 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 + + 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"], + } + class OmniKVTransferManager: """Unified management for OmniConnector and KV cache transfer. @@ -79,6 +341,13 @@ class OmniKVTransferManager: else (None, None) ) + if config.need_send_cache and config.connector_config: + try: + _ = self.connector + logger.info("Sender connector eagerly initialized") + except Exception as e: + logger.warning("Failed to eagerly initialize sender connector: %s", e) + @classmethod def _create(cls, cfg: dict | None) -> "OmniKVTransferManager": """Create manager from raw config dict.""" @@ -140,8 +409,39 @@ class OmniKVTransferManager: cfg = self.config.connector_config if cfg and (c_type := cfg.get("type")): try: - logger.info(f"Initializing OmniConnector with config: {cfg}") c_extra = {k: v for k, v in cfg.items() if k != "type"} + if c_type == "MooncakeTransferEngineConnector": + base_port = c_extra.get("zmq_port", 50051) + c_extra["from_stage"] = ( + str(self.config.from_stage) if self.config.from_stage is not None else "0" + ) + c_extra["to_stage"] = str(self.config.to_stage) if self.config.to_stage is not None else "1" + + 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 + 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) + + logger.info( + "Initializing OmniConnector (purpose=kv_transfer) with config: %s, role: %s", + cfg, + 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}") @@ -157,6 +457,85 @@ class OmniKVTransferManager: """Get connector (compatibility wrapper for existing code).""" return self.connector + def _resolve_sender_info( + self, sender_info: dict[str, Any], sender_stage_id: str | int | None = None + ) -> dict[str, Any] | None: + if not sender_info: + return None + + if "host" in sender_info: + return sender_info + + if not isinstance(sender_info, dict): + return None + + preferred_keys: list[str | int] = [] + if sender_stage_id is None: + recv_from, _ = self.recv_stages + sender_stage_id = recv_from + + if sender_stage_id is not None: + preferred_keys.append(sender_stage_id) + preferred_keys.append(str(sender_stage_id)) + try: + preferred_keys.append(int(sender_stage_id)) + except (TypeError, ValueError): + pass + + for key in dict.fromkeys(preferred_keys): + info = sender_info.get(key) + if isinstance(info, dict) and "host" in info: + return info + + candidates = [info for info in sender_info.values() if isinstance(info, dict) and "host" in info] + if len(candidates) == 1: + return candidates[0] + + if candidates: + logger.warning( + "Ambiguous sender_info for sender_stage_id=%s: " + "expected caller to resolve a single sender entry, got %s", + sender_stage_id, + sender_info, + ) + return None + + @staticmethod + def _clone_received_payload_tensors(data: dict[str, Any]) -> dict[str, Any]: + if not isinstance(data, dict) or "layer_blocks" not in data: + return data + + layer_blocks = data["layer_blocks"] + for cache_name in ("key_cache", "value_cache"): + cache_list = layer_blocks.get(cache_name, []) + for idx, tensor in enumerate(cache_list): + if isinstance(tensor, torch.Tensor): + cache_list[idx] = tensor.clone() + return data + + 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.""" + if not self.config.need_recv_cache: + return + + actual_info = self._resolve_sender_info(sender_info, sender_stage_id=sender_stage_id) + if not actual_info or "host" not in actual_info: + logger.warning("Invalid sender_info format: %s", sender_info) + return + + 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") + + 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")) + except Exception: + if hasattr(self._connector, "sender_host"): + self._connector.sender_host = actual_info.get("host") + if hasattr(self._connector, "sender_zmq_port"): + self._connector.sender_zmq_port = actual_info.get("zmq_port") + def handle_finished_requests_kv_transfer( self, finished_reqs: dict[str, dict[str, Any]], @@ -203,7 +582,8 @@ class OmniKVTransferManager: custom_metadata = data.get("custom_metadata") - # Extract KV cache from GPU blocks -> CPU tensors + # Extract KV cache from GPU blocks and keep it on-device when + # possible so raw-data connectors can use the fast path. kv_data = self._extract_kv_cache( req_id, block_ids, seq_len, kv_caches, block_size, cache_dtype, custom_metadata ) @@ -280,9 +660,8 @@ class OmniKVTransferManager: flat_k = flat_k[:seq_len] flat_v = flat_v[:seq_len] - # Move to CPU - key_cache[layer_idx] = flat_k.detach().cpu().contiguous() - value_cache[layer_idx] = flat_v.detach().cpu().contiguous() + key_cache[layer_idx] = flat_k.detach().contiguous() + value_cache[layer_idx] = flat_v.detach().contiguous() if not any(k is not None for k in key_cache): return None @@ -311,14 +690,40 @@ class OmniKVTransferManager: if not from_stage or not to_stage: raise ValueError("Transfer stages (omni_from_stage, omni_to_stage) not configured") - # Prepare data and transfer with retry - data_dict = kv_data.to_dict() - data_dict["request_id"] = transfer_req_id + 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) - success, size, _ = self._transfer_with_retry(from_stage, to_stage, f"kv_cache_{transfer_req_id}", data_dict) + try: + if supports_raw: + transfer_data = kv_data.to_gpu_tensor() + 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 + + 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) + elapsed = time.perf_counter() - transfer_start if success: - logger.info(f"KV transfer OK: {transfer_req_id}, {size} bytes") + mbps = (size / 1024 / 1024) / elapsed if elapsed > 0 else 0 + logger.info( + "KV transfer OK: %s, %s bytes, %.3fs, %.1f MB/s", + transfer_req_id, + size, + elapsed, + mbps, + ) else: logger.error(f"KV transfer FAILED: {transfer_req_id}") @@ -327,7 +732,7 @@ class OmniKVTransferManager: from_stage: str, to_stage: str, request_id: str, - data: dict[str, Any], + data: "dict[str, Any] | bytes | torch.Tensor", max_retries: int = 3, ) -> tuple[bool, int, dict[str, Any] | None]: """Transfer data with retry and exponential backoff. @@ -393,6 +798,8 @@ class OmniKVTransferManager: timeout = self.config.recv_timeout start_time = time.time() + 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}...") @@ -400,33 +807,74 @@ class OmniKVTransferManager: 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: - data, size = result - logger.info(f"Successfully received KV cache for {request_id}, {size} bytes") + raw_data, size = result + elapsed = time.time() - start_time + link_ms = (time.perf_counter() - link_start) * 1000 + managed_buffer = None - # Move tensors to target device if specified - if target_device is not None and isinstance(data, dict) and "layer_blocks" in data: - layer_blocks = data["layer_blocks"] - for cache_list in [ - layer_blocks.get("key_cache", []), - layer_blocks.get("value_cache", []), - ]: - for i, tensor in enumerate(cache_list): - if isinstance(tensor, torch.Tensor) and tensor.device != target_device: - cache_list[i] = tensor.to(target_device).contiguous() + if hasattr(raw_data, "tensor") and hasattr(raw_data, "release"): + managed_buffer = raw_data + try: + buf_tensor = raw_data.tensor + if buf_tensor.is_cuda: + data = KVCacheTransferData.from_bytes_gpu(buf_tensor) + raw_data.release() + managed_buffer = None + else: + data = KVCacheTransferData.from_bytes(memoryview(buf_tensor.numpy())) + data = self._clone_received_payload_tensors(data) + raw_data.release() + managed_buffer = None + except Exception as e: + logger.error("Failed to deserialize KV cache from ManagedBuffer: %s", e) + if managed_buffer is not None: + raw_data.release() + return None, 0 + elif isinstance(raw_data, (bytes, bytearray)): + data = KVCacheTransferData.from_bytes(raw_data) + elif isinstance(raw_data, torch.Tensor) and raw_data.dtype == torch.uint8 and raw_data.dim() == 1: + data = KVCacheTransferData.from_bytes(raw_data.cpu().numpy().tobytes()) + else: + data = raw_data + try: + if isinstance(data, dict) and "layer_blocks" in data: + layer_blocks = data["layer_blocks"] + for cache_list in [ + layer_blocks.get("key_cache", []), + layer_blocks.get("value_cache", []), + ]: + for i, tensor in enumerate(cache_list): + if not isinstance(tensor, torch.Tensor): + 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() + + logger.info( + "Successfully received KV cache for %s, %s bytes, wait=%.3fs, link=%.1fms", + request_id, + size, + elapsed, + link_ms, + ) return data, size if time.time() - start_time > timeout: logger.error(f"Timeout waiting for KV cache for request {request_id} after {timeout}s") return None, 0 - time.sleep(0.5) + 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}") @@ -459,6 +907,16 @@ class OmniKVTransferManager: if hasattr(req, "sampling_params") and req.sampling_params is not None: req.sampling_params.kv_metadata = data["metadata"] + @staticmethod + def _resolve_request_id(req: Any) -> str | None: + """Resolve the logical request ID used for KV transfer lookups.""" + request_id = getattr(req, "request_id", None) + if request_id: + return request_id + if hasattr(req, "request_ids") and req.request_ids: + return req.request_ids[0] + return None + # Legacy compatibility method def receive_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool: """Receive KV cache and populate request object (legacy interface). @@ -470,11 +928,11 @@ class OmniKVTransferManager: Returns: True if successful, False otherwise """ - request_id = getattr(req, "request_id", None) - if not request_id and hasattr(req, "request_ids") and req.request_ids: - # Adaptation for new OmniDiffusionRequest which has list of prompts/ids - request_id = req.request_ids[0] + kv_sender_info = getattr(req, "kv_sender_info", None) + if kv_sender_info: + self.update_sender_info(kv_sender_info, sender_stage_id=self.recv_stages[0]) + request_id = self._resolve_request_id(req) if not request_id: logger.warning("Request has no ID, cannot receive KV cache") return False @@ -513,9 +971,7 @@ class OmniKVTransferManager: cfg_ids = getattr(getattr(req, "sampling_params", None), "cfg_kv_request_ids", None) if cfg_ids and cfg_kv_collect_func: - request_id = getattr(req, "request_id", None) or ( - req.request_ids[0] if hasattr(req, "request_ids") and req.request_ids else None - ) + request_id = self._resolve_request_id(req) try: cfg_kvs = cfg_kv_collect_func( request_id, diff --git a/vllm_omni/distributed/omni_connectors/utils/initialization.py b/vllm_omni/distributed/omni_connectors/utils/initialization.py index aaa222b4c..37b7d0d7f 100644 --- a/vllm_omni/distributed/omni_connectors/utils/initialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/initialization.py @@ -19,9 +19,17 @@ else: logger = get_connector_logger(__name__) +# Reserve a separate port range for KV-transfer sockets so they do not +# collide with request-forwarding endpoints that share the same base port. +KV_TRANSFER_PORT_OFFSET = 100 + def initialize_connectors_from_config( - config_path: str | Path | None = None, default_shm_threshold: int = 65536 + config_path: str | Path | None = None, + default_shm_threshold: int = 65536, + purpose: str = "request_forwarding", + caller_stage_id: int | str | None = None, + is_sender: bool | None = None, ) -> tuple[OmniTransferConfig | None, dict[tuple[str, str], OmniConnectorBase]]: """ Initialize connectors from configuration file. @@ -36,12 +44,20 @@ def initialize_connectors_from_config( return None, {} # create connectors from config - connectors = create_connectors_from_config(transfer_config.connectors) + connectors = create_connectors_from_config( + transfer_config.connectors, + purpose=purpose, + caller_stage_id=caller_stage_id, + is_sender=is_sender, + ) return transfer_config, connectors def create_connectors_from_config( connectors_config: dict[tuple[str, str], ConnectorSpec], + purpose: str = "request_forwarding", + caller_stage_id: int | str | None = None, + is_sender: bool | None = None, ) -> dict[tuple[str, str], OmniConnectorBase]: """ Create connectors from config. @@ -52,12 +68,59 @@ def create_connectors_from_config( Returns: A dictionary of connectors. """ + purpose_port_offsets = { + "request_forwarding": 0, + "kv_transfer": KV_TRANSFER_PORT_OFFSET, + } + port_offset = purpose_port_offsets.get(purpose, 0) + orchestrator_port_offset = 200 + connectors = {} for edge_key, connector_spec in connectors_config.items(): + from_stage, to_stage = edge_key try: - connector = OmniConnectorFactory.create_connector(connector_spec) + if connector_spec.name == "MooncakeTransferEngineConnector": + extra = dict(connector_spec.extra) if connector_spec.extra else {} + base_port = extra.get("zmq_port", 50051) + try: + stage_offset = int(from_stage) + except (TypeError, ValueError): + stage_offset = 0 + + if str(caller_stage_id) == "orchestrator": + adjusted_port = base_port + orchestrator_port_offset + stage_offset + else: + adjusted_port = base_port + port_offset + stage_offset + extra["zmq_port"] = adjusted_port + + if is_sender is not None: + extra["role"] = "sender" if is_sender else "receiver" + if not is_sender: + extra.setdefault("sender_host", extra.get("host", "127.0.0.1")) + extra.setdefault("sender_zmq_port", adjusted_port) + elif caller_stage_id is not None: + caller_str = str(caller_stage_id) + if caller_str == from_stage: + extra["role"] = "sender" + elif caller_str == to_stage: + extra["role"] = "receiver" + extra.setdefault("sender_host", extra.get("host", "127.0.0.1")) + extra.setdefault("sender_zmq_port", adjusted_port) + else: + extra["role"] = "sender" + else: + extra["role"] = extra.get("role", "auto") + + connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=connector_spec.name, extra=extra)) + else: + connector = OmniConnectorFactory.create_connector(connector_spec) connectors[edge_key] = connector - logger.info(f"Created connector for {edge_key[0]} -> {edge_key[1]}: {type(connector).__name__}") + logger.info( + "Created connector for %s -> %s: %s", + from_stage, + to_stage, + type(connector).__name__, + ) except Exception as e: raise RuntimeError(f"Failed to initialize connector for edge {edge_key}: {e}") from e @@ -289,7 +352,11 @@ def initialize_orchestrator_connectors( else: default_shm_threshold = max(0, shm_threshold_bytes) transfer_config, connectors = initialize_connectors_from_config( - config_path, default_shm_threshold=default_shm_threshold + config_path, + default_shm_threshold=default_shm_threshold, + purpose="request_forwarding", + caller_stage_id="orchestrator", + is_sender=True, ) return transfer_config, connectors @@ -316,6 +383,7 @@ def get_stage_connector_config( def build_stage_connectors( stage_id: int, connectors_config: dict[str, Any], + purpose: str = "request_forwarding", ) -> dict[tuple[str, str], Any] | None: """Instantiate OmniConnectors for a stage based on config.""" if not connectors_config: @@ -352,7 +420,12 @@ def build_stage_connectors( try: # Use unified connector creation logic - connectors = create_connectors_from_config(stage_connector_specs) + connectors = create_connectors_from_config( + stage_connector_specs, + purpose=purpose, + caller_stage_id=stage_id, + is_sender=False, + ) except Exception as exc: # pragma: no cover - defensive logging # Fail fast so the stage does not start with missing connectors. logger.exception("[Stage-%s] Failed to initialize connectors: %s", stage_id, exc) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 8cd2d6952..f7e7d53d5 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -395,15 +395,17 @@ class AsyncOmniEngine: proc=proc, addresses=addresses, ) + logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id) + # Keep the stage-specific device visibility until vLLM + # finishes starting all child processes. + complete_stage_handshake(proc, handshake_address, addresses, vllm_config) + logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id) finally: if previous_visible_devices is None: current_omni_platform.unset_device_control_env_var() else: current_omni_platform.set_device_control_env_var(previous_visible_devices) - logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id) - complete_stage_handshake(proc, handshake_address, addresses, vllm_config) - logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id) assert started_stage is not None return started_stage except Exception: diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 8ea9a5096..20dce1f0f 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -477,6 +477,30 @@ class Orchestrator: ), ) + def _build_kv_sender_info(self, sender_stage_ids: list[int]) -> dict[int, dict[str, Any]] | None: + """Build per-request sender info for diffusion KV-transfer receivers.""" + sender_infos: dict[int, dict[str, Any]] = {} + for sender_stage_id in dict.fromkeys(sender_stage_ids): + if sender_stage_id < 0 or sender_stage_id >= self.num_stages: + continue + + sender_stage = self.stage_clients[sender_stage_id] + get_sender_info = getattr(sender_stage, "get_kv_sender_info", None) + if not callable(get_sender_info): + continue + + sender_info = get_sender_info() + if not sender_info: + logger.warning( + "[Orchestrator] Stage-%s has no KV sender info available", + sender_stage_id, + ) + continue + + sender_infos[sender_stage_id] = sender_info + + return sender_infos or None + async def _forward_to_next_stage( self, req_id: str, @@ -522,14 +546,22 @@ class Orchestrator: req_id, ) + source_stage_ids = list(getattr(next_client, "engine_input_source", None) or [stage_id]) + kv_sender_info = self._build_kv_sender_info(sender_stage_ids=source_stage_ids) if isinstance(diffusion_prompt, list): await next_client.add_batch_request_async( req_id, diffusion_prompt, params, + kv_sender_info=kv_sender_info, ) else: - await next_client.add_request_async(req_id, diffusion_prompt, params) + await next_client.add_request_async( + req_id, + diffusion_prompt, + params, + kv_sender_info=kv_sender_info, + ) req_state.stage_submit_ts[next_stage_id] = _time.time() return @@ -731,7 +763,14 @@ class Orchestrator: params = req_state.sampling_params_list[next_stage_id] if next_client.stage_type == "diffusion": - await next_client.add_request_async(request_id, req_state.prompt, params) + source_stage_ids = list(getattr(next_client, "engine_input_source", None) or [next_stage_id - 1]) + kv_sender_info = self._build_kv_sender_info(sender_stage_ids=source_stage_ids) + await next_client.add_request_async( + request_id, + req_state.prompt, + params, + kv_sender_info=kv_sender_info, + ) req_state.stage_submit_ts[next_stage_id] = _time.time() continue diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index e08ce7801..71a0aee4a 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -6,12 +6,15 @@ Directly inherits from vLLM's AsyncMPClient to reuse EngineCore architecture. from __future__ import annotations +import socket from typing import TYPE_CHECKING, Any +from urllib.parse import urlparse from vllm.logger import init_logger from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.core_client import AsyncMPClient +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: @@ -76,6 +79,11 @@ class StageEngineCoreClient(AsyncMPClient): self.engine_outputs: Any = None self._proc = proc + self.client_addresses = dict(client_addresses or {}) + self._omni_kv_config = getattr(getattr(vllm_config, "model_config", None), "omni_kv_config", None) + self._kv_sender_host = self._resolve_contact_host() + self._kv_sender_info: dict[str, Any] | None = None + self._kv_sender_initialized = False logger.info( "[StageEngineCoreClient] Stage-%s initializing EngineCore", @@ -104,6 +112,7 @@ class StageEngineCoreClient(AsyncMPClient): shutdown_error, ) raise + self._initialize_kv_sender_endpoint() logger.info( "[StageEngineCoreClient] Stage-%s EngineCore running", self.stage_id, @@ -118,6 +127,118 @@ class StageEngineCoreClient(AsyncMPClient): # ==================== Stage Methods ==================== + @staticmethod + def _detect_local_ip() -> str | None: + """Best-effort local IP detection for cross-node connector bootstrap.""" + try: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock: + sock.connect(("8.8.8.8", 80)) + return sock.getsockname()[0] + except Exception: + try: + return socket.gethostbyname(socket.gethostname()) + except Exception: + return None + + def _resolve_contact_host(self) -> str | None: + """Resolve a routable host for this stage from its client addresses.""" + for key in ("input_address", "output_address", "stats_update_address"): + address = self.client_addresses.get(key) + if not address: + continue + host = urlparse(address).hostname + if host in {None, "", "*", "0.0.0.0", "::"}: + continue + if host in {"localhost", "127.0.0.1"}: + detected = self._detect_local_ip() + if detected: + return detected + continue + return host + return self._detect_local_ip() + + def _get_kv_connector_config(self) -> dict[str, Any] | None: + omni_kv_config = getattr(self, "_omni_kv_config", None) + if not isinstance(omni_kv_config, dict): + return None + connector_config = omni_kv_config.get("connector_config") + if not isinstance(connector_config, dict): + return None + return connector_config + + def _resolve_sender_host_from_config(self, connector_config: dict[str, Any]) -> str | None: + host = connector_config.get("sender_host") or connector_config.get("host") + if host in {None, "", "auto", "*", "0.0.0.0", "::"}: + return self._resolve_contact_host() + return str(host) + + def _initialize_kv_sender_endpoint(self) -> None: + if self._kv_sender_initialized: + return + self._kv_sender_initialized = True + connector_config = self._get_kv_connector_config() + if connector_config is None or connector_config.get("role") != "sender": + return + + sender_host = self._resolve_sender_host_from_config(connector_config) + if sender_host is not None: + self._kv_sender_host = sender_host + + sender_port = connector_config.get("sender_zmq_port") + if sender_port is None: + base_port = connector_config.get("zmq_port") + if base_port is None: + return + + omni_kv_config = getattr(self, "_omni_kv_config", None) + from_stage = self.stage_id + if isinstance(omni_kv_config, dict): + from_stage = omni_kv_config.get("omni_from_stage", from_stage) + + try: + sender_port = int(base_port) + KV_TRANSFER_PORT_OFFSET + int(from_stage) + except (TypeError, ValueError): + logger.warning( + "[StageEngineCoreClient] Stage-%s could not resolve sender_zmq_port " + "from base_port=%s and from_stage=%s", + self.stage_id, + base_port, + from_stage, + ) + return + + if self._kv_sender_host is None: + return + + self._kv_sender_info = { + "host": str(self._kv_sender_host), + "zmq_port": int(sender_port), + } + + def get_kv_sender_info( + self, + *, + base_port: int = 50051, + kv_transfer_port_offset: int = KV_TRANSFER_PORT_OFFSET, + ) -> dict[str, Any] | None: + """Build sender bootstrap info for diffusion KV transfer receivers. + + ``base_port`` and ``kv_transfer_port_offset`` are only used by the + legacy fallback path when no connector-level sender endpoint is + configured in ``omni_kv_config``. + """ + if self._kv_sender_info is not None: + return dict(self._kv_sender_info) + + if self._kv_sender_host is None: + self._kv_sender_host = self._resolve_contact_host() + if self._kv_sender_host is None: + return None + return { + "host": self._kv_sender_host, + "zmq_port": base_port + kv_transfer_port_offset + int(self.stage_id), + } + def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None: """Set engine outputs (called by orchestrator).""" self.engine_outputs = engine_outputs diff --git a/vllm_omni/model_executor/stage_configs/bagel.yaml b/vllm_omni/model_executor/stage_configs/bagel.yaml index b0c1b0480..d1031b574 100644 --- a/vllm_omni/model_executor/stage_configs/bagel.yaml +++ b/vllm_omni/model_executor/stage_configs/bagel.yaml @@ -1,5 +1,9 @@ # Stage 0: Thinker (multimodal understanding + text generation) +# By default this config uses the shared-memory connector for stage-0 -> stage-1 forwarding. +# To switch to RDMA, add output_connectors/input_connectors that point to +# rdma_connector and keep the rest of the pipeline unchanged. + stage_args: - stage_id: 0 stage_type: llm @@ -36,6 +40,9 @@ stage_args: seed: 52 detokenize: True repetition_penalty: 1.05 + # Optional RDMA override: + # output_connectors: + # to_stage_1: rdma_connector - stage_id: 1 stage_type: diffusion @@ -62,6 +69,9 @@ stage_args: is_comprehension: false default_sampling_params: seed: 52 + # Optional RDMA override: + # input_connectors: + # from_stage_0: rdma_connector # Runtime edges runtime: @@ -78,6 +88,23 @@ runtime: extra: shm_threshold_bytes: 65536 # 64KB threshold + # Optional RDMA connector template for Bagel. To enable it, point + # stage-0 output_connectors/to_stage_1 and stage-1 input_connectors/from_stage_0 + # to rdma_connector instead of relying on the default shared-memory path. + rdma_connector: + name: MooncakeTransferEngineConnector + extra: + host: "auto" + zmq_port: 50051 + protocol: "rdma" + device_name: "" + # Memory pool for RDMA-registered buffers. + # Supports both CPU pinned memory ("cpu") and GPU VRAM ("cuda"). + # CPU mode works on all topologies; GPU mode (GPUDirect RDMA) requires + # NIC-GPU direct PCIe connectivity (PIX topology). + # Recommended: 4 GB for CPU, 2 GB for GPU (to conserve VRAM). + memory_pool_size: 4294967296 # 4 GB + memory_pool_device: "cpu" edges: - from: 0