[Core] hunyuan-image prefetch kv (#4448)

Signed-off-by: z00806815 <zhuhonghan@huawei.com>
Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
Signed-off-by: dph97 <nenbaying@163.com>
Signed-off-by: zhh <zhh@users.noreply.github.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: zhuhh97 <zhuhh97@users.noreply.github.com>

tests/diffusion/test_diffusion_model_runner.py

@@ -80,6 +80,7 @@ def _make_request(skip_cache_refresh: bool = True):
         num_inference_steps=4,
     )
     return SimpleNamespace(
+        request_id="req-test",
         prompts=["a prompt"],
         sampling_params=sampling_params,
         skip_cache_refresh=skip_cache_refresh,
@@ -104,8 +105,9 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
     runner.kv_transfer_manager = SimpleNamespace(
         receive_kv_cache=lambda req, target_device=None: None,
         receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None,
-        receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None,
+        receive_multi_kv_cache_distributed=lambda *a, **k: None,
     )
+    runner._kv_prefetch_enabled = False
     return runner
 
 

tests/diffusion/test_worker_wrapper_base.py

@@ -184,7 +184,7 @@ class TestWorkerWrapperBaseDelegation:
         mock_reqs = [mocker.Mock()]
         result = wrapper.execute_model(mock_reqs, mock_od_config)
 
-        wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config)
+        wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config, kv_prefetch_jobs=None)
         assert result == mock_output
 
     def test_load_weights_delegation(self, mocker: MockerFixture, mock_od_config):

tests/distributed/omni_connectors/test_kv_async_prefetch.py

@@ -0,0 +1,356 @@
+"""Background KV prefetch — CPU-testable mechanism.
+
+Covers opt-in gating, the start_prefetch / consume_prefetched_kv round trip over a
+mock connector, dedup, per-call sender_info isolation, role classification, and
+miss/abort paths.  GPU H2D and D2D paths need CUDA (integration tests).
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    OmniKVCacheConfig,
+    OmniKVTransferManager,
+    ReceiveRole,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# start_prefetch() skips stubs without a sender endpoint, so submissions carry it.
+_SENDER_INFO = {"host": "127.0.0.1", "zmq_port": 50051}
+
+
+class MockConnector:
+    def __init__(self):
+        self.store = {}
+        self.get_calls = []
+
+    def put(self, from_stage, to_stage, put_key, data):
+        self.store[f"{from_stage}->{to_stage}:{put_key}"] = data
+        return True, len(str(data)), None
+
+    def get(self, from_stage, to_stage, get_key, metadata=None):
+        self.get_calls.append({"get_key": get_key, "metadata": metadata})
+        key = f"{from_stage}->{to_stage}:{get_key}"
+        if key in self.store:
+            return self.store[key], len(str(self.store[key]))
+        return None
+
+
+def _make_sender_receiver(*, async_prefetch=True):
+    """Build sender+receiver managers sharing one mock connector, seed one KV."""
+    connector = MockConnector()
+
+    sender = OmniKVTransferManager(
+        OmniKVCacheConfig(
+            connector_config={"type": "mock"}, from_stage="sender", to_stage="receiver", need_send_cache=True
+        )
+    )
+    sender._connector = connector
+
+    receiver = OmniKVTransferManager(
+        OmniKVCacheConfig(
+            connector_config={"type": "mock"},
+            from_stage="sender",
+            stage_id="receiver",
+            need_recv_cache=True,
+            recv_timeout=1.0,
+        ),
+        async_prefetch=async_prefetch,
+    )
+    receiver._connector = connector
+    return sender, receiver, connector
+
+
+def _seed_payload(sender, req_id, *, num_layers=2, block_size=4, num_heads=4, head_dim=8):
+    kv_caches = [torch.randn(2, 5, block_size, num_heads, head_dim) for _ in range(num_layers)]
+    sender.handle_finished_requests_kv_transfer(
+        {req_id: {"block_ids": [0, 1], "seq_len": 10}}, kv_caches, block_size, "float32"
+    )
+
+
+def _req(rid):
+    """A minimal request carrying just ``request_id`` (consume_prefetched_kv only needs that)."""
+    return SimpleNamespace(request_id=rid)
+
+
+# --------------------------------------------------------------------------- #
+#  Opt-in gating
+# --------------------------------------------------------------------------- #
+
+
+def test_async_prefetch_defaults_false():
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    assert mgr._async_prefetch is False
+
+
+def test_from_od_config_enables_prefetch_flag():
+    od = SimpleNamespace(omni_kv_config={"enable_kv_async_prefetch": True})
+    assert OmniKVTransferManager.from_od_config(od)._async_prefetch is True
+    od2 = SimpleNamespace(omni_kv_config=None)
+    assert OmniKVTransferManager.from_od_config(od2)._async_prefetch is False
+    od3 = SimpleNamespace(omni_kv_config={"enable_kv_async_prefetch": False})
+    assert OmniKVTransferManager.from_od_config(od3)._async_prefetch is False
+
+
+def test_prefetch_preserved_with_connector_config():
+    # The flag must survive merging with a populated connector config.
+    od = SimpleNamespace(
+        omni_kv_config={
+            "connector_config": {"type": "mock"},
+            "need_recv_cache": True,
+            "enable_kv_async_prefetch": True,
+        },
+    )
+    assert OmniKVTransferManager.from_od_config(od)._async_prefetch is True
+
+
+def test_prefetch_disabled_with_cfg_companion_collector():
+    # A CFG companion collector forces the synchronous path (see _resolve_async_prefetch).
+    od = SimpleNamespace(
+        omni_kv_config={"enable_kv_async_prefetch": True},
+        cfg_kv_collect_func=lambda *a, **k: {},
+    )
+    assert OmniKVTransferManager.from_od_config(od)._async_prefetch is False
+
+
+def test_from_vllm_config_ar_path_is_always_sync():
+    # The AR entry point never enables prefetch.
+    model_config = SimpleNamespace(omni_kv_config=None)
+    vllm_config = SimpleNamespace(kv_transfer_config=None)
+    assert OmniKVTransferManager.from_vllm_config(vllm_config, model_config)._async_prefetch is False
+
+
+def test_start_prefetch_noop_when_disabled():
+    _, receiver, _ = _make_sender_receiver(async_prefetch=False)
+    receiver.start_prefetch({"request_id": "r1", "kv_sender_info": None})
+    assert receiver._prefetch_futures == {}
+
+
+# --------------------------------------------------------------------------- #
+#  Round trip
+# --------------------------------------------------------------------------- #
+
+
+def test_prefetch_round_trip_returns_data():
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-1")
+
+    receiver.start_prefetch({"request_id": "rid-1", "kv_sender_info": _SENDER_INFO})
+    assert "rid-1" in receiver._prefetch_futures
+
+    data, size = receiver.consume_prefetched_kv(_req("rid-1"))
+    assert data is not None
+    assert "layer_blocks" in data
+    assert data["metadata"]["seq_len"] == 10
+    assert size > 0
+    # Future is consumed (popped) on retrieval.
+    assert "rid-1" not in receiver._prefetch_futures
+
+
+def test_consume_prefetched_kv_miss_without_prefetch_returns_none():
+    _, receiver, _ = _make_sender_receiver(async_prefetch=False)
+    assert receiver.consume_prefetched_kv(_req("never-prefetched")) == (None, 0)
+
+
+def test_start_prefetch_dedups_same_request():
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-dup")
+    receiver.start_prefetch({"request_id": "rid-dup", "kv_sender_info": _SENDER_INFO})
+    fut1 = receiver._prefetch_futures["rid-dup"]
+    receiver.start_prefetch({"request_id": "rid-dup", "kv_sender_info": _SENDER_INFO})
+    assert receiver._prefetch_futures["rid-dup"] is fut1
+
+
+def test_start_prefetch_skips_without_sender_info():
+    # No endpoint => no submission; the sync path handles it.
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-nosi")
+    receiver.start_prefetch({"request_id": "rid-nosi", "kv_sender_info": None})
+    assert receiver._prefetch_futures == {}
+
+
+def test_consume_prefetched_kv_sweeps_stale_entries():
+    # Consuming a slot drops every other tracked prefetch.
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-stale")
+    receiver.start_prefetch({"request_id": "rid-stale", "kv_sender_info": _SENDER_INFO})
+    assert "rid-stale" in receiver._prefetch_futures
+    # Miss on the current request still sweeps stale entries.
+    receiver._async_prefetch = False
+    assert receiver.consume_prefetched_kv(_req("rid-current")) == (None, 0)
+    assert receiver._prefetch_futures == {}
+
+
+# --------------------------------------------------------------------------- #
+#  consume_prefetched_kv: miss returns None, stale sweep
+# --------------------------------------------------------------------------- #
+
+
+def test_start_prefetch_sweeps_orphan():
+    # An orphan prefetch (e.g. aborted) is dropped when the next prefetch starts.
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-orphan")
+    _seed_payload(sender, "rid-next")
+    receiver.start_prefetch({"request_id": "rid-orphan", "kv_sender_info": _SENDER_INFO})
+    receiver.consume_prefetched_kv(_req("rid-orphan"))  # complete the future
+    receiver.start_prefetch({"request_id": "rid-orphan", "kv_sender_info": _SENDER_INFO})  # re-add as leftover
+    receiver.start_prefetch({"request_id": "rid-next", "kv_sender_info": _SENDER_INFO})
+    assert "rid-orphan" not in receiver._prefetch_futures
+    assert "rid-next" in receiver._prefetch_futures
+
+
+def test_shutdown_prefetch_clears_state():
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-sd")
+    receiver.start_prefetch({"request_id": "rid-sd", "kv_sender_info": _SENDER_INFO})
+    receiver.shutdown_prefetch()
+    assert receiver._prefetch_futures == {}
+    assert receiver._prefetch_executor is None
+
+
+# --------------------------------------------------------------------------- #
+#  Correctness: per-call sender_info must not touch instance state
+# --------------------------------------------------------------------------- #
+
+
+def test_receive_with_sender_info_does_not_mutate_instance_state():
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-si")
+    assert receiver._sender_base_host is None
+    assert receiver._sender_base_zmq_port is None
+
+    receiver.receive_kv_cache_for_request(
+        "rid-si",
+        target_device=None,
+        sender_info={"host": "1.2.3.4", "zmq_port": 50051},
+    )
+    # The per-call endpoint must not leak into instance fields.
+    assert receiver._sender_base_host is None
+    assert receiver._sender_base_zmq_port is None
+
+
+# --------------------------------------------------------------------------- #
+#  Receive-role classification (recv_role) and the follower prefetch guard
+# --------------------------------------------------------------------------- #
+
+import vllm_omni.diffusion.distributed.parallel_state as ps  # noqa: E402
+
+
+def _patch_topo(monkeypatch, *, world_size, world_rank=0, cfg_size=1, cfg_rank=0, sp_size=1, sp_rank=0):
+    monkeypatch.setattr(ps, "get_world_group", lambda: SimpleNamespace(world_size=world_size, rank_in_group=world_rank))
+    monkeypatch.setattr(ps, "get_classifier_free_guidance_world_size", lambda: cfg_size)
+    monkeypatch.setattr(ps, "get_classifier_free_guidance_rank", lambda: cfg_rank)
+    monkeypatch.setattr(ps, "get_cfg_group", lambda: None)
+    monkeypatch.setattr(ps, "get_sequence_parallel_world_size", lambda: sp_size)
+    monkeypatch.setattr(ps, "get_sequence_parallel_rank", lambda: sp_rank)
+    monkeypatch.setattr(ps, "get_sp_group", lambda: None)
+
+
+def _set_tp(mgr, src, tgt):
+    mgr._tp_topo = SimpleNamespace(source_tp_size=src, target_tp_size=tgt, local_rank=0)
+
+
+def _role(mgr):
+    """Recompute the receive role from the (patched) parallel state, bypassing the cache."""
+    mgr._topo_config = None
+    return mgr.topo_config.role
+
+
+def test_recv_role_world1_is_local(monkeypatch):
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    _set_tp(mgr, 1, 1)
+    _patch_topo(monkeypatch, world_size=1)
+    assert _role(mgr) is ReceiveRole.LOCAL
+
+
+def test_recv_role_pure_tp_is_local(monkeypatch):
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    _set_tp(mgr, 2, 2)
+    _patch_topo(monkeypatch, world_size=2)
+    assert _role(mgr) is ReceiveRole.LOCAL
+
+
+def test_recv_role_cfg_leader_and_follower(monkeypatch):
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    _set_tp(mgr, 2, 2)
+    _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=0)
+    assert _role(mgr) is ReceiveRole.LEADER
+    _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=1)
+    assert _role(mgr) is ReceiveRole.FOLLOWER
+
+
+def test_recv_role_sp_leader_and_follower(monkeypatch):
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    _set_tp(mgr, 2, 2)
+    _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=0)
+    assert _role(mgr) is ReceiveRole.LEADER
+    _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=1)
+    assert _role(mgr) is ReceiveRole.FOLLOWER
+
+
+def test_recv_role_tp_and_sp_leader_per_tp_rank(monkeypatch):
+    # TP=2 + SP=2: sp_rank 0 is leader (pulls), sp_rank 1 is follower (gets broadcast).
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    _set_tp(mgr, 2, 2)
+    _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=0)
+    assert _role(mgr) is ReceiveRole.LEADER
+    _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=1)
+    assert _role(mgr) is ReceiveRole.FOLLOWER
+
+
+def test_recv_role_hsdp_world_broadcast(monkeypatch):
+    # TP inactive, world > 1 (HSDP): rank-0 leader, others follower.
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+    _set_tp(mgr, 1, 1)
+    _patch_topo(monkeypatch, world_size=2, world_rank=0)
+    assert _role(mgr) is ReceiveRole.LEADER
+    _patch_topo(monkeypatch, world_size=2, world_rank=1)
+    assert _role(mgr) is ReceiveRole.FOLLOWER
+
+
+def test_recv_role_uninitialized_defaults_local(monkeypatch):
+    mgr = OmniKVTransferManager(OmniKVCacheConfig())
+
+    def _boom():
+        raise AssertionError("world group is not initialized")
+
+    monkeypatch.setattr(ps, "get_world_group", _boom)
+    assert _role(mgr) is ReceiveRole.LOCAL
+
+
+def test_start_prefetch_skips_follower(monkeypatch):
+    # A follower never pulls, so start_prefetch must not submit.
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-fol")
+    _set_tp(receiver, 2, 2)
+    _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=1)
+    receiver.start_prefetch({"request_id": "rid-fol", "kv_sender_info": _SENDER_INFO})
+    assert receiver._prefetch_futures == {}
+
+
+def test_start_prefetch_submits_for_owner(monkeypatch):
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-own")
+    _set_tp(receiver, 2, 2)
+    _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=0)
+    receiver.start_prefetch({"request_id": "rid-own", "kv_sender_info": _SENDER_INFO})
+    assert "rid-own" in receiver._prefetch_futures
+
+
+def test_consume_then_apply_attaches_payload():
+    sender, receiver, _ = _make_sender_receiver()
+    _seed_payload(sender, "rid-apply")
+    receiver.start_prefetch({"request_id": "rid-apply", "kv_sender_info": _SENDER_INFO})
+    data, _ = receiver.consume_prefetched_kv(_req("rid-apply"))
+    assert data is not None
+
+    req = OmniDiffusionRequest(prompts=["p"], sampling_params=OmniDiffusionSamplingParams(), request_id="rid-apply")
+    # Mirror consume_and_distribute_kv_cache's LOCAL apply path (CPU: record_stream no-op).
+    receiver._record_stream_for_prefetched(data)
+    receiver.apply_kv_cache_to_request(req, data)
+    assert req.past_key_values is not None

tests/distributed/omni_connectors/test_kv_flow.py

@@ -1,7 +1,6 @@
 import json
 import struct
 
-import numpy as np
 import pytest
 import torch
 
@@ -206,17 +205,6 @@ def test_update_sender_info_uses_configured_source_stage():
     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"},
@@ -526,6 +514,57 @@ def test_integration_flow(common_constants):
     assert req.kv_metadata["seq_len"] == 10
 
 
+def test_deferred_release_isolates_data_from_pool_buffer():
+    payload, key_tensor = _make_serialized_payload()
+
+    # ManagedBuffer-like object whose release() zeros storage (simulates reuse).
+    raw_array = bytearray(payload)
+    buf_tensor = torch.frombuffer(raw_array, dtype=torch.uint8)
+
+    class _CorruptingManagedBuffer:
+        def __init__(self, storage, tensor):
+            self._storage = storage
+            self._tensor = tensor
+
+        @property
+        def tensor(self):
+            return self._tensor
+
+        def release(self):
+            for i in range(len(self._storage)):
+                self._storage[i] = 0
+
+    managed_buf = _CorruptingManagedBuffer(raw_array, buf_tensor)
+
+    class _ManagedBufConnector:
+        def __init__(self, buf, from_stage, to_stage, req_id):
+            self._buf = buf
+            self._key = f"{from_stage}->{to_stage}:omni_{from_stage}_to_{to_stage}_kv_cache_{req_id}"
+
+        def get(self, from_stage, to_stage, get_key, metadata=None):
+            if f"{from_stage}->{to_stage}:{get_key}" == self._key:
+                return self._buf, len(self._buf._storage)
+            return None
+
+    config = OmniKVCacheConfig(
+        connector_config={"type": "mock"},
+        from_stage="sender",
+        stage_id="receiver",
+        need_recv_cache=True,
+        recv_timeout=1.0,
+    )
+    manager = OmniKVTransferManager(config)
+    req_id = "req-payload"
+    manager._connector = _ManagedBufConnector(managed_buf, "sender", "receiver", req_id)
+
+    data, size = manager.receive_kv_cache_for_request(req_id, target_device=torch.device("cpu"))
+
+    assert data is not None
+    assert size > 0
+    # Returned data must survive the pool buffer being zeroed by release().
+    assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor)
+
+
 def test_manager_extraction_no_connector(kv_config, common_constants):
     """Test extraction when connector is unavailable (should still return IDs)."""
     block_size = common_constants["block_size"]

tests/distributed/omni_connectors/test_tp_rank_aware.py

@@ -586,8 +586,17 @@ class TestDistributedReceive:
                 return_value=0,
             ),
             patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size",
+                return_value=1,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank",
+                return_value=0,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None),
         ):
-            assert mgr.receive_multi_kv_cache_distributed(req) is True
+            assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True
 
         mgr.receive_multi_kv_cache.assert_called_once()
         assert mgr.receive_multi_kv_cache.call_args.args[2] == torch.device("cpu")
@@ -642,6 +651,15 @@ class TestDistributedReceive:
                 return_value=1,
             ),
             patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size",
+                return_value=1,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank",
+                return_value=0,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None),
         ):
             assert mgr.receive_multi_kv_cache_distributed(req) is True
 
@@ -666,7 +684,27 @@ class TestDistributedReceive:
         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):
+        with (
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size",
+                return_value=1,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank",
+                return_value=0,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=None),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size",
+                return_value=1,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank",
+                return_value=0,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None),
+        ):
             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"))
@@ -722,7 +760,7 @@ class TestDistributedReceive:
             ),
             patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group),
         ):
-            assert mgr.receive_multi_kv_cache_distributed(req) is True
+            assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True
 
         # Verify behavior based on role
         if is_owner:
@@ -817,7 +855,7 @@ class TestDistributedReceive:
             ),
             patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group),
         ):
-            assert mgr.receive_multi_kv_cache_distributed(req) is True
+            assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True
 
         # Verify behavior based on role
         if role == "owner":

vllm_omni/deploy/hunyuan_image_3_moe.yaml

@@ -62,7 +62,7 @@ stages:
       max_tokens: 8192
       detokenize: true
       skip_special_tokens: false
-      include_stop_str_in_output: true,
+      include_stop_str_in_output: true
 
   - stage_id: 1
     max_num_seqs: 1
@@ -71,6 +71,8 @@ stages:
     distributed_executor_backend: "mp"
     omni_kv_config:
       need_recv_cache: true
+      enable_kv_async_prefetch: true
+      kv_prefetch_min_free_mem_ratio: 0.2
     parallel_config:
       tensor_parallel_size: 2
       enable_expert_parallel: true

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -321,7 +321,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         new_req = scheduler_output.scheduled_new_reqs[0]
         result = self.collective_rpc(
             "execute_model",
-            args=(new_req.req, self.od_config),
+            args=(new_req.req, self.od_config, getattr(scheduler_output, "kv_prefetch_jobs", None)),
             unique_reply_rank=0,
             exec_all_ranks=True,
         )

vllm_omni/diffusion/sched/base_scheduler.py

@@ -52,6 +52,7 @@ class _BaseScheduler(SchedulerInterface):
         self._running_sampling_params_key: SamplingParamsKey | None = None
         self._finished_req_ids: set[str] = set()
         self.max_num_running_reqs: int = 1
+        self._prefetch_enabled: bool = False
 
     def initialize(self, od_config: OmniDiffusionConfig) -> None:
         self.od_config = od_config
@@ -66,6 +67,8 @@ class _BaseScheduler(SchedulerInterface):
             self.max_num_running_reqs = max(1, int(max_num_seqs))
         except (TypeError, ValueError):
             self.max_num_running_reqs = 1
+        omni_kv = getattr(od_config, "omni_kv_config", None) or {}
+        self._prefetch_enabled = bool(omni_kv.get("enable_kv_async_prefetch", False))
         self._reset_scheduler_state()
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
@@ -111,6 +114,22 @@ class _BaseScheduler(SchedulerInterface):
             else:
                 scheduled_cached_request_ids.append(request_id)
 
+        # Expose the next waiting request (serial mode) so the runner can
+        # prefetch its KV during this forward.  Skip a request without
+        # kv_sender_info (would target the wrong sender under multi-replica) or
+        # one already finished/aborted (would consume its sender buffer for
+        # nothing).
+        kv_prefetch_jobs: dict | None = None
+        if self._prefetch_enabled and self._waiting:
+            nxt = self._request_states.get(self._waiting[0])
+            if nxt is not None and not nxt.is_finished():
+                sender_info = getattr(nxt.req, "kv_sender_info", None)
+                if sender_info:
+                    kv_prefetch_jobs = {
+                        "request_id": nxt.request_id,
+                        "kv_sender_info": sender_info,
+                    }
+
         scheduler_output = DiffusionSchedulerOutput(
             step_id=self._step_id,
             scheduled_new_reqs=scheduled_new_reqs,
@@ -118,6 +137,7 @@ class _BaseScheduler(SchedulerInterface):
             finished_req_ids=set(self._finished_req_ids),
             num_running_reqs=len(self._running),
             num_waiting_reqs=len(self._waiting),
+            kv_prefetch_jobs=kv_prefetch_jobs,
         )
 
         # update after schedule

vllm_omni/diffusion/sched/interface.py

@@ -113,6 +113,8 @@ class DiffusionSchedulerOutput:
     finished_req_ids: set[str]
     num_running_reqs: int
     num_waiting_reqs: int
+    # next request to background-prefetch KV
+    kv_prefetch_jobs: dict | None = None
 
     @cached_property
     def scheduled_request_ids(self) -> list[str]:

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -80,9 +80,23 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         # Cache for per-request stepwise state.
         self.state_cache: dict[str, DiffusionRequestState] = {}
 
-        # Initialize KV cache manager for connector management
+        # Initialize KV cache manager for connector management.
         self.kv_transfer_manager = OmniKVTransferManager.from_od_config(od_config)
 
+        # Prefetch covers TP / SP / CFG-Parallel / HSDP.  Disabled when a CFG
+        # companion KV collector is set (that KV is not backgrounded).
+        has_cfg_companion_kv = getattr(od_config, "cfg_kv_collect_func", None) is not None
+
+        self._kv_prefetch_enabled = (
+            bool(self.kv_transfer_manager.config.enable_kv_async_prefetch)
+            and not has_cfg_companion_kv
+            and self.kv_transfer_manager.config.need_recv_cache
+        )
+
+    @property
+    def target_device(self) -> torch.device | None:
+        return getattr(self.pipeline, "device", None)
+
     def _compile_transformer(self, attr_name: str) -> None:
         """Compile a transformer attribute on the pipeline with torch.compile."""
         model = getattr(self.pipeline, attr_name, None)
@@ -269,7 +283,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             pool_overhead_gb / peak_reserved_gb * 100 if peak_reserved_gb > 0 else 0.0,
         )
 
-    def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+    def execute_model(self, req: OmniDiffusionRequest, kv_prefetch_jobs: dict | None = None) -> DiffusionOutput:
         """
         Execute a forward pass for the given requests.
 
@@ -293,12 +307,26 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         use_hsdp = self.od_config.parallel_config.use_hsdp
         grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
         with grad_context:
-            # The manager handles the check for need_recv_cache internally
-            self.kv_transfer_manager.receive_multi_kv_cache_distributed(
-                req,
-                cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None),
-                target_device=getattr(self.pipeline, "device", None),
-            )
+            # Receive AR KV (fetch → distribute → apply inside the entry). prefetch on:
+            # consume prior-forward payload, sync-fallback on miss; else sync receive.
+            kv_recv_t0 = time.perf_counter()
+            if self._kv_prefetch_enabled:
+                self.kv_transfer_manager.consume_and_distribute_kv_cache(
+                    req,
+                    target_device=self.target_device,
+                )
+            else:
+                self.kv_transfer_manager.receive_multi_kv_cache_distributed(
+                    req,
+                    cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None),
+                    target_device=self.target_device,
+                )
+            kv_recv_ms = (time.perf_counter() - kv_recv_t0) * 1000
+            logger.debug("KV recv for %s %.1fms", req.request_id, kv_recv_ms)
+
+            # Kick off the next request's prefetch (+ H2D) to overlap this forward.
+            if self._kv_prefetch_enabled and kv_prefetch_jobs is not None:
+                self.kv_transfer_manager.start_prefetch(kv_prefetch_jobs, self.target_device)
 
             if req.sampling_params.generator is None and req.sampling_params.seed is not None:
                 if req.sampling_params.generator_device is not None:
@@ -404,7 +432,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 self.kv_transfer_manager.receive_multi_kv_cache_distributed(
                     state_req,
                     cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None),
-                    target_device=getattr(self.pipeline, "device", None),
+                    target_device=self.target_device,
                 )
                 self.state_cache[request_id] = new_state
                 resolved.append(new_state)

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -369,7 +369,12 @@ class DiffusionWorker:
         else:
             profiler.stop()
 
-    def execute_model(self, req: OmniDiffusionRequest, od_config: OmniDiffusionConfig) -> DiffusionOutput:
+    def execute_model(
+        self,
+        req: OmniDiffusionRequest,
+        od_config: OmniDiffusionConfig,
+        kv_prefetch_jobs: dict | None = None,
+    ) -> DiffusionOutput:
         """Execute a forward pass by delegating to the model runner."""
         assert self.model_runner is not None, "Model runner not initialized"
         if self.lora_manager is not None:
@@ -382,7 +387,7 @@ class DiffusionWorker:
         profiler = self._get_profiler()
         ctx = profiler.annotate_context_manager("diffusion_forward") if profiler else nullcontext()
         with ctx:
-            output = self.model_runner.execute_model(req)
+            output = self.model_runner.execute_model(req, kv_prefetch_jobs=kv_prefetch_jobs)
         if profiler:
             profiler.step()
         return output
@@ -646,6 +651,10 @@ class DiffusionWorker:
 
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup distributed environment."""
+        if self.model_runner is not None:
+            mgr = getattr(self.model_runner, "kv_transfer_manager", None)
+            if mgr is not None:
+                mgr.shutdown_prefetch()
         destroy_distributed_env()
 
 
@@ -1077,18 +1086,24 @@ class WorkerWrapperBase:
         """
         return self.worker.generate(requests)
 
-    def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
+    def execute_model(
+        self,
+        reqs: list[OmniDiffusionRequest],
+        od_config: OmniDiffusionConfig,
+        kv_prefetch_jobs: dict | None = None,
+    ) -> DiffusionOutput:
         """
         Execute a forward pass.
 
         Args:
             reqs: List of diffusion requests
             od_config: OmniDiffusionConfig configuration
+            kv_prefetch_jobs: Optional next-request KV prefetch descriptor.
 
         Returns:
             DiffusionOutput with generated results
         """
-        return self.worker.execute_model(reqs, od_config)
+        return self.worker.execute_model(reqs, od_config, kv_prefetch_jobs=kv_prefetch_jobs)
 
     def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
         """Execute one diffusion step."""

vllm_omni/entrypoints/cli/serve.py

@@ -595,7 +595,6 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Enable layerwise (blockwise) offloading on DiT modules.",
         )
-
         # Video model parameters (e.g., Wan2.2) - engine-level
         omni_config_group.add_argument(
             "--boundary-ratio",

vllm_omni/platforms/cuda/platform.py

@@ -213,6 +213,11 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
         free, _ = torch.cuda.mem_get_info(device)
         return free
 
+    @classmethod
+    def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]:
+        free, total = torch.cuda.mem_get_info(device)
+        return free, total
+
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
         return torch.cuda.get_device_name(device_id)

vllm_omni/platforms/interface.py

@@ -161,6 +161,10 @@ class OmniPlatform(Platform):
     def get_free_memory(cls, device: torch.device | None = None) -> int:
         raise NotImplementedError
 
+    @classmethod
+    def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]:
+        raise NotImplementedError
+
     @classmethod
     def create_autocast_context(
         cls,

vllm_omni/platforms/musa/platform.py

@@ -146,17 +146,14 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
 
     @classmethod
     def get_free_memory(cls, device: torch.device | None = None) -> int:
-        """Get the free memory on the MUSA device.
-
-        Args:
-            device: Optional device to query. If None, uses current device.
-
-        Returns:
-            Free memory in bytes.
-        """
         free, _ = torch.musa.mem_get_info(device)
         return free
 
+    @classmethod
+    def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]:
+        free, total = torch.musa.mem_get_info(device)
+        return free, total
+
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
         return torch.musa.get_device_name(device_id)

vllm_omni/platforms/npu/platform.py

@@ -149,6 +149,11 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
         free, _ = torch.npu.mem_get_info(device)
         return free
 
+    @classmethod
+    def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]:
+        free, total = torch.npu.mem_get_info(device)
+        return free, total
+
     @classmethod
     def get_device_total_memory(cls, device_id: int = 0) -> int:
         device_props = torch.npu.get_device_properties(device_id)

vllm_omni/platforms/rocm/platform.py

@@ -142,6 +142,11 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
         free, _ = torch.cuda.mem_get_info(device)
         return free
 
+    @classmethod
+    def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]:
+        free, total = torch.cuda.mem_get_info(device)
+        return free, total
+
     @classmethod
     def set_device_control_env_var(cls, devices: str | int | None) -> None:
         import os

vllm_omni/platforms/xpu/platform.py

@@ -85,6 +85,11 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
         free, _ = torch.xpu.mem_get_info(device)
         return free
 
+    @classmethod
+    def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]:
+        free, total = torch.xpu.mem_get_info(device)
+        return free, total
+
     @classmethod
     def get_profiler_cls(cls) -> str:
         """Return XPU-specific profiler that handles XPU events."""