[Feature] Transfer AR-to-diffusion payloads over NIXL (#6264)

Signed-off-by: yuanwu <yuan.wu@intel.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/distributed/test_parallel_state_sp_groups.py

@@ -146,3 +146,34 @@ def test_set_seq_parallel_pg_validates_sp_group_ranks(monkeypatch):
             world_size=4,
             sp_group_ranks=[[0, 2]],
         )
+
+
+@pytest.mark.parametrize("rank", [1, 3])
+def test_tensor_dict_broadcast_preserves_subgroup_source(rank, monkeypatch):
+    from vllm_omni.diffusion.distributed.group_coordinator import GroupCoordinator, TensorMetadata
+
+    coordinator = object.__new__(GroupCoordinator)
+    coordinator.rank = rank
+    coordinator.ranks = [1, 3]
+    coordinator.rank_in_group = [1, 3].index(rank)
+    coordinator.world_size = 2
+    coordinator.device_group = object()
+    coordinator.cpu_group = object()
+    metadata = [("payload", TensorMetadata("cpu", torch.float32, (2,)))]
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+
+    def broadcast_object(value, src):
+        assert src == 0
+        return metadata
+
+    def broadcast(tensor, src, group, async_op):
+        assert src == 1
+        assert group is coordinator.cpu_group
+        assert async_op
+        tensor.fill_(7)
+        return SimpleNamespace(wait=lambda: None)
+
+    monkeypatch.setattr(coordinator, "broadcast_object", broadcast_object)
+    monkeypatch.setattr(torch.distributed, "broadcast", broadcast)
+    result = coordinator.broadcast_tensor_dict({"payload": torch.full((2,), 7.0)} if rank == 1 else None)
+    assert torch.equal(result["payload"], torch.full((2,), 7.0))

tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py

@@ -89,6 +89,23 @@ def test_launch_failure_disables_runtime_key(monkeypatch):
     assert len(fused_rope._FAILED_KEYS) == 1
 
 
+def test_non_cuda_platform_does_not_launch_fused_kernel(monkeypatch):
+    from vllm_omni.diffusion.models.ernie_image import fused_rope
+
+    monkeypatch.setattr(fused_rope, "HAS_TRITON", True)
+    monkeypatch.setattr(fused_rope.current_omni_platform, "is_cuda", lambda: False)
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False)
+
+    def unexpected_launch(*args):
+        raise AssertionError("non-CUDA platforms must remain on the native path")
+
+    monkeypatch.setattr(fused_rope, "_launch_fused_qk_rotary_emb", unexpected_launch)
+    query = torch.zeros((1, 1, 1, 128), dtype=torch.bfloat16)
+    freqs = torch.zeros((1, 1, 64), dtype=torch.float32)
+    assert fused_rope.try_fused_qk_rotary_emb(query, query, freqs, freqs) is None
+    assert not fused_rope._FAILED_KEYS
+
+
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 def test_failed_runtime_key_cache_is_bounded():
     from vllm_omni.diffusion.models.ernie_image import fused_rope

tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py

@@ -10,6 +10,7 @@ import torch
 from PIL import Image
 
 from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+from vllm_omni.errors import OmniClientError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -42,8 +43,10 @@ def pipeline():
     return model
 
 
-def _request(task):
+def _request(task, *, payload=None):
     prompt: dict[str, Any] = {"prompt": "A person waves."}
+    if payload is not None:
+        prompt["additional_information"] = {"encoder_output": payload}
     if task != "t2va":
         prompt["multi_modal_data"] = {"image": Image.new("RGB", (256, 256))}
     sampling = OmniDiffusionSamplingParams(
@@ -64,7 +67,120 @@ def test_single_stage_encodes_without_stage_zero(pipeline, task):
     torch.testing.assert_close(audio, pipeline.decode.return_value[1])
 
 
-def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monkeypatch):
+def test_single_stage_reuses_legacy_text_but_encodes_media(pipeline):
+    payload = {
+        "hidden_states": torch.full((2, 5120), 3.0, dtype=torch.bfloat16),
+        "token_tags": torch.ones(2, dtype=torch.long),
+    }
+    pipeline.forward(_request("fl2va", payload=payload))
+    pipeline.encode_prompt.assert_not_called()
+    pipeline.video_vae.encode_image.assert_called_once()
+    torch.testing.assert_close(pipeline.diffuse.call_args.kwargs["text_embeddings"], payload["hidden_states"])
+
+
+def test_text_disaggregation_keeps_local_media_encoding(pipeline):
+    payload = {
+        "hidden_states": torch.full((2, 5120), 3.0, dtype=torch.bfloat16),
+        "token_tags": torch.ones(2, dtype=torch.long),
+    }
+    pipeline.load_text_encoder = False
+    pipeline.load_vae_encoder = True
+    pipeline.forward(_request("fl2va", payload=payload))
+    pipeline.encode_prompt.assert_not_called()
+    pipeline.video_vae.encode_image.assert_called_once()
+
+
+@pytest.mark.parametrize("payload", [None, {}, {"hidden_states": torch.ones(2, 5120), "token_tags": torch.ones(2)}])
+def test_stage_one_rejects_missing_or_legacy_conditioning(pipeline, payload):
+    pipeline.load_text_encoder = False
+    pipeline.load_vae_encoder = False
+    with pytest.raises(OmniClientError):
+        pipeline.forward(_request("t2va", payload=payload))
+    pipeline.encode_prompt.assert_not_called()
+    pipeline.video_vae.encode_image.assert_not_called()
+    pipeline.diffuse.assert_not_called()
+
+
+@pytest.mark.parametrize("task", ["t2va", "fl2va", "ref2va"])
+@pytest.mark.parametrize("load_text_encoder", [True, False])
+def test_prepare_encode_uses_the_same_mode_and_conditioning(pipeline, task, load_text_encoder):
+    from vllm_omni.diffusion.worker.utils import StepRequestState
+
+    request = _request(task)
+    if not load_text_encoder:
+        conditioning, window_text = pipeline._prepare_local_conditioning(request.prompts[0], request.sampling_params)
+        assert window_text is None
+        request.prompts[0] = {"additional_information": {"encoder_output": conditioning.to_omni_payload()}}
+        pipeline.encode_prompt.reset_mock()
+        pipeline.video_vae.encode_image.reset_mock()
+    pipeline.load_text_encoder = load_text_encoder
+    pipeline.load_vae_encoder = load_text_encoder
+    state = StepRequestState(request_id="step", prompt=request.prompts[0], sampling=request.sampling_params)
+    assert pipeline.prepare_encode(state) is state
+    assert state.latents is not None
+    assert state.latents.shape[1] == 96
+    assert state.total_steps == 1
+    assert pipeline.encode_prompt.call_count == int(load_text_encoder)
+    assert pipeline.video_vae.encode_image.call_count == int(load_text_encoder and task != "t2va")
+
+
+def test_stage_one_step_execution_never_falls_back(pipeline):
+    from vllm_omni.diffusion.worker.utils import StepRequestState
+
+    pipeline.load_text_encoder = False
+    pipeline.load_vae_encoder = False
+    request = _request("t2va")
+    with pytest.raises(OmniClientError, match="requires encoder conditioning"):
+        pipeline.prepare_encode(
+            StepRequestState(request_id="step", prompt=request.prompts[0], sampling=request.sampling_params)
+        )
+    pipeline.encode_prompt.assert_not_called()
+
+
+def test_multiple_outputs_encode_only_once(pipeline):
+    request = _request("fl2va")
+    request.sampling_params.num_outputs_per_prompt = 3
+    request.sampling_params.seed = 42
+    pipeline.forward(request)
+    pipeline.encode_prompt.assert_called_once()
+    pipeline.video_vae.encode_image.assert_called_once()
+    assert [call.kwargs["seed"] for call in pipeline.diffuse.call_args_list] == [42, 43, 44]
+
+
+def test_ref2va_partition_preserves_image_only_implicit_task(pipeline):
+    request = _request("fl2va")
+    del request.sampling_params.extra_args["task"]
+    pipeline.partition = "ref2va"
+    pipeline.supported_tasks = frozenset({"ref2va"})
+    pipeline.forward(request)
+    assert pipeline.diffuse.call_args.kwargs["task"] == "ref2va"
+
+
+def test_invalid_legacy_text_is_not_ignored(pipeline):
+    with pytest.raises(OmniClientError, match="encoder wire requires"):
+        pipeline.forward(_request("fl2va", payload={"hidden_states": "invalid"}))
+    pipeline.encode_prompt.assert_not_called()
+    pipeline.video_vae.encode_image.assert_not_called()
+
+
+def test_text_weight_loading_preserves_encoder_specific_lifecycle(pipeline):
+    class Encoder(torch.nn.Module):
+        def load_weights(self, weights):
+            return {name for name, _ in weights}
+
+    pipeline.text_encoder = Encoder()
+    pipeline.video_vae = torch.nn.Identity()
+    pipeline.audio_vae = torch.nn.Identity()
+    assert pipeline.load_weights([("text_encoder.layer.weight", torch.ones(1))]) == {"text_encoder.layer.weight"}
+    pipeline.load_text_encoder = False
+    pipeline.load_vae_encoder = False
+    pipeline.text_encoder = None
+    with pytest.raises(ValueError, match="disabled in this deployment"):
+        pipeline.load_weights([("text_encoder.layer.weight", torch.ones(1))])
+
+
+@pytest.mark.parametrize("connector_handoff", [False, True], ids=["inline", "full_payload"])
+def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monkeypatch, connector_handoff):
     import numpy as np
 
     from vllm_omni.data_entry_keys import flatten_payload
@@ -74,7 +190,11 @@ def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monke
         MiniMaxH3TextConditioning,
     )
     from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder
-    from vllm_omni.model_executor.stage_input_processors.minimax_h3 import encoder2diffusion, prepare_encoder_prompt
+    from vllm_omni.model_executor.stage_input_processors.minimax_h3 import (
+        encoder2diffusion,
+        encoder2diffusion_full_payload,
+        prepare_encoder_prompt,
+    )
 
     frames = np.zeros((48, 32, 32, 3), dtype=np.uint8)
     decode_frames = Mock(return_value=frames)
@@ -110,11 +230,33 @@ def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monke
     media = stage_zero._encode_media(stage_zero._media_input(transformed["additional_information"]))
     text = MiniMaxH3TextConditioning(*pipeline.encode_prompt.return_value)
     payload = MiniMaxH3EncoderConditioning.from_components(text, media).to_omni_payload()
-    source = SimpleNamespace(finished=True, outputs=[SimpleNamespace(multimodal_output=flatten_payload(payload))])
+    wire = flatten_payload(payload)
+    source = SimpleNamespace(
+        finished=True, outputs=[SimpleNamespace(multimodal_output=None if connector_handoff else wire)]
+    )
     bridged = encoder2diffusion([source], request.prompts[0])
+    assert bridged["multi_modal_data"] is None
+    if connector_handoff:
+        assert "encoder_output" not in bridged["additional_information"]
+        transferred = encoder2diffusion_full_payload(pooling_output=wire)
+        assert set(transferred) == {"encoder_output"}
+        received = transferred["encoder_output"]
+        for section, key in (
+            ("hidden_states", "output"),
+            ("meta", "token_role_ids"),
+            ("embed", "embedding"),
+            ("embed", "speech_feat"),
+            ("kv_metadata", "minimax_h3_encoder_layout"),
+        ):
+            torch.testing.assert_close(received[section][key], payload[section][key])
+        # Simulate the generic diffusion runner merging the received payload;
+        # native NIXL transport is covered separately, not by this CPU test.
+        bridged["additional_information"].update(transferred)
     pipeline.load_text_encoder = False
     pipeline.load_vae_encoder = False
     pipeline.encode_prompt.reset_mock()
+    pipeline.video_vae.encode_image.reset_mock()
+    pipeline.video_vae.encode_video.reset_mock()
     pipeline.forward(SimpleNamespace(prompts=[bridged], sampling_params=request.sampling_params))
     external_kwargs = pipeline.diffuse.call_args.kwargs
     for key, value in local_kwargs.items():
@@ -123,6 +265,8 @@ def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monke
         else:
             assert value == external_kwargs[key]
     pipeline.encode_prompt.assert_not_called()
+    pipeline.video_vae.encode_image.assert_not_called()
+    pipeline.video_vae.encode_video.assert_not_called()
     assert decode_frames.call_count == sampler.call_count == 2
     assert all(call.kwargs["decoded_frames"] is frames for call in sampler.call_args_list)
     assert [block["kind"] for block in local_kwargs["ref_blocks"]] == ["image", "video_audio", "audio"]

tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

@@ -10,9 +10,17 @@ from unittest.mock import Mock
 import pytest
 import torch
 
+from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def _media(**kwargs):
+    return MiniMaxH3EncoderMediaInput(
+        task="ref2va", height=768, width=1344, num_frames=124, latent_t=37, audio_t=207, **kwargs
+    )
+
+
 def test_h3_full_pipeline_profiler_includes_local_encoding():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
@@ -28,6 +36,9 @@ def test_h3_full_pipeline_profiler_includes_local_encoding():
         "denoise_step",
         "post_decode",
     ]
+    assert "_encode_video_conditions" not in targets
+    assert "_encode_video_audio_conditions" not in targets
+    assert "_encode_audio_conditions" not in targets
 
 
 def test_decoded_output_is_unchanged_without_active_model_offload():
@@ -316,46 +327,6 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch, load_text
     remove_offload.assert_called_once_with([*dits, *stages])
 
 
-@pytest.mark.parametrize("decode_fails", [False, True])
-def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
-
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    component = torch.nn.Linear(2, 2)
-    events: list[tuple[Any, ...]] = []
-
-    @contextmanager
-    def record_component(value):
-        events.append(("activate", value))
-        try:
-            yield
-        finally:
-            events.append(("offload", value))
-
-    monkeypatch.setattr(module, "sequential_offload_component", record_component)
-    pipeline._model_cpu_offload_modules = [component]
-
-    def decode():
-        with pipeline._component_on_device(component):
-            events.append(("decode", component))
-            if decode_fails:
-                raise RuntimeError("decode failed")
-
-    if decode_fails:
-        with pytest.raises(RuntimeError, match="decode failed"):
-            decode()
-    else:
-        decode()
-
-    assert events == [
-        ("activate", component),
-        ("decode", component),
-        ("offload", component),
-    ]
-
-
 def test_h3_model_cpu_offload_keeps_unselected_vaes_resident(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
@@ -416,6 +387,46 @@ def test_h3_model_cpu_offload_rejects_unloaded_selected_encoder(monkeypatch):
     apply_offload.assert_not_called()
 
 
+@pytest.mark.parametrize("decode_fails", [False, True])
+def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    component = torch.nn.Linear(2, 2)
+    events: list[tuple[Any, ...]] = []
+
+    @contextmanager
+    def record_component(value):
+        events.append(("activate", value))
+        try:
+            yield
+        finally:
+            events.append(("offload", value))
+
+    monkeypatch.setattr(module, "sequential_offload_component", record_component)
+    pipeline._model_cpu_offload_modules = [component]
+
+    def decode():
+        with pipeline._component_on_device(component):
+            events.append(("decode", component))
+            if decode_fails:
+                raise RuntimeError("decode failed")
+
+    if decode_fails:
+        with pytest.raises(RuntimeError, match="decode failed"):
+            decode()
+    else:
+        decode()
+
+    assert events == [
+        ("activate", component),
+        ("decode", component),
+        ("offload", component),
+    ]
+
+
 def _encode_media(pipeline, **kwargs):
     from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
     from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media

tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py

@@ -416,7 +416,7 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch, batch_frames):
     monkeypatch.setattr(
         mod.MiniMaxH3Pipeline,
         "_prepare_encoder_conditioning_inputs",
-        lambda self, value, sampling: context,
+        lambda self, value, sampling: context if value is conditioning else pytest.fail("wrong encoder handoff"),
     )
     monkeypatch.setattr(
         mod.MiniMaxH3Pipeline,

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py

@@ -15,6 +15,21 @@ from vllm_omni.diffusion.models.minimax_h3.temporal_chunks import decode_tempora
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+@pytest.mark.parametrize("chunked", [False, True])
+def test_encode_only_video_vae_rejects_full_and_chunked_decode(chunked):
+    from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE
+
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    torch.nn.Module.__init__(vae)
+    vae.encode_only = True
+    latent = torch.empty(1, 1, 1, 1, 1)
+    with pytest.raises(RuntimeError, match="encode-only video VAE cannot decode"):
+        if chunked:
+            vae.decode_with_chunks(latent, on_chunk=lambda frames: None)
+        else:
+            vae.decode_latent(latent)
+
+
 @pytest.fixture(scope="module")
 def cpu_process_group():
     if dist.is_initialized():

tests/diffusion/test_diffusion_model_runner.py

@@ -547,6 +547,9 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):
     runner.od_config.step_execution = True
     req = _make_request()
     req.request_id = "req"
+    req.external_req_id = "external-req"
+    sent = []
+    runner._maybe_send_stage_payload = lambda requests, outputs: sent.append((requests[0].external_req_id, outputs))
 
     monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
     monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
@@ -560,6 +563,7 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):
 
     first = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
     assert first.get_request_output("req").result is None
+    assert sent == []
 
     scheduler_output = SimpleNamespace(
         finished_req_ids=set(),
@@ -569,11 +573,13 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):
     second = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
     assert second.get_request_output("req").result == chunks[0]
     assert second.get_request_output("req").finished is False
+    assert sent == []
 
     DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
     fourth = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
     assert fourth.get_request_output("req").result == chunks[1]
     assert fourth.get_request_output("req").finished is True
+    assert sent == [("external-req", [chunks[1]])]
 
 
 @pytest.mark.core_model
@@ -733,6 +739,7 @@ def test_execute_model_passes_single_request_batch_to_non_admission_pipeline(mon
     req = _make_request()
 
     monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
 
     output = DiffusionModelRunner.execute_model(runner, req)
 
@@ -749,6 +756,7 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin
     req = _make_request()
 
     monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
 
     output = DiffusionModelRunner.execute_model(runner, req)
 

tests/diffusion/test_diffusion_stage_payload.py

@@ -0,0 +1,785 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+HANDLE_KEY = DiffusionModelRunner._STAGE_PAYLOAD_HANDLE_KEY
+
+
+class _FakeConnector:
+    def abandon_get(self, get_key):
+        pass
+
+    def get_with_deadline(self, from_stage, to_stage, get_key, metadata=None, *, deadline):
+        return self.get(from_stage, to_stage, get_key, metadata)
+
+    def __init__(self, payload=None, *, raises=False, put_result=(True, 128, {"schema_version": 1})):
+        self._payload = payload
+        self._raises = raises
+        self._put_result = put_result
+        self.calls: list[tuple[str, str, str, object]] = []
+        self.put_calls: list[tuple[str, str, str, object]] = []
+
+    def get(self, from_stage, to_stage, get_key, metadata=None):
+        self.calls.append((from_stage, to_stage, get_key, metadata))
+        if self._raises:
+            raise RuntimeError("transfer failed")
+        if self._payload is None:
+            return None
+        return self._payload, 0
+
+    def put(self, from_stage, to_stage, put_key, data):
+        self.put_calls.append((from_stage, to_stage, put_key, data))
+        if self._raises:
+            raise RuntimeError("transfer failed")
+        return self._put_result
+
+
+class _FakeKVTransferManager:
+    from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
+
+    _resolve_sender_info = OmniKVTransferManager._resolve_sender_info
+
+    def __init__(self, connector, recv_stages=("0", "1"), send_stages=("0", "1")):
+        self.connector = connector
+        self.recv_stages = recv_stages
+        self.send_stages = send_stages
+        self.sender_info_calls: list[tuple[dict, str | None]] = []
+
+    def update_sender_info(self, sender_info, sender_stage_id=None):
+        self.sender_info_calls.append((sender_info, sender_stage_id))
+
+
+def _make_runner(connector, *, payload_keys=("text_encoder_output",), recv_stages=("0", "1")):
+    runner = object.__new__(DiffusionModelRunner)
+    runner.od_config = SimpleNamespace(stage_input_payload_keys=payload_keys, stage_id=1)
+    runner.device = torch.device("cpu")
+    runner._local_rank = 0
+    runner.pipeline = None
+    runner.kv_transfer_manager = _FakeKVTransferManager(connector, recv_stages=recv_stages)
+    runner.init_omni_connectors(runner.od_config, runner.kv_transfer_manager, synchronous=True)
+    return runner
+
+
+def _make_sender(connector, *, payload_keys=("prompt_embeds",), send_stages=("0", "1")):
+    runner = object.__new__(DiffusionModelRunner)
+    runner.od_config = SimpleNamespace(stage_output_payload_keys=payload_keys, stage_id=0)
+    runner.device = torch.device("cpu")
+    runner.pipeline = None
+    runner.kv_transfer_manager = _FakeKVTransferManager(connector, send_stages=send_stages)
+    runner.init_omni_connectors(runner.od_config, runner.kv_transfer_manager, synchronous=True)
+    return runner
+
+
+def _make_output(**custom):
+    return SimpleNamespace(custom_output=dict(custom))
+
+
+def _make_request(prompt, *, request_id="req-7", kv_sender_info=None, payload_sender_info=None):
+    return SimpleNamespace(
+        prompt=prompt,
+        request_id=request_id,
+        kv_sender_info=kv_sender_info,
+        payload_sender_info=payload_sender_info,
+    )
+
+
+def _conditioning():
+    return {"text_encoder_output": {"hidden_states": torch.zeros(4, 8), "token_tags": torch.zeros(4)}}
+
+
+@pytest.mark.parametrize("shard_size,replicate_size", [(4, 1), (1, 4), (2, 2)])
+def test_standalone_hsdp_elects_one_payload_consumer(monkeypatch, shard_size, replicate_size):
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    leaders = []
+    for rank in range(shard_size * replicate_size):
+        shard_group = SimpleNamespace(world_size=shard_size, rank_in_group=rank % shard_size)
+        replicate_group = SimpleNamespace(world_size=replicate_size, rank_in_group=rank // shard_size)
+        monkeypatch.setattr(parallel_state, "get_fs_group", lambda: shard_group)
+        monkeypatch.setattr(parallel_state, "get_hsdp_replicate_group", lambda: replicate_group)
+        monkeypatch.setattr(parallel_state, "get_sp_group", lambda: None)
+        runner = _make_runner(_FakeConnector())
+        monkeypatch.setattr(runner, "_get_local_tp_group", lambda: None)
+        if runner.is_data_transfer_rank():
+            leaders.append(rank)
+
+    assert leaders == [0]
+
+
+@pytest.mark.parametrize("shard_size,replicate_size", [(4, 1), (1, 4), (2, 2)])
+@pytest.mark.parametrize("delivered", [True, False])
+def test_standalone_hsdp_payload_fanout(monkeypatch, shard_size, replicate_size, delivered):
+    from uuid import uuid4
+
+    from vllm_omni.diffusion.distributed import parallel_state
+    from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
+
+    packets = {}
+
+    class MeshGroup:
+        def __init__(self, dimension, peer_rank, rank, size):
+            self.key = (dimension, peer_rank)
+            self.rank_in_group = rank
+            self.world_size = size
+
+        def broadcast_object(self, value, src=0):
+            if self.rank_in_group == src:
+                packets[self.key, "delivered"] = value
+            return packets[self.key, "delivered"]
+
+        def broadcast_tensor_dict(self, value, src=0):
+            if self.rank_in_group == src:
+                packets[self.key, "payload"] = value
+            return packets[self.key, "payload"]
+
+    connector = SharedMemoryConnector({})
+    request_id = f"hsdp-{uuid4().hex}"
+    key = f"{request_id}_0_0"
+    expected = torch.arange(8).reshape(2, 4)
+    calls = []
+    original_get = connector.get_with_deadline
+
+    def tracked_get(*args, **kwargs):
+        calls.append(rank)
+        return original_get(*args, **kwargs)
+
+    monkeypatch.setattr(connector, "get_with_deadline", tracked_get)
+    monkeypatch.setattr(parallel_state, "get_sp_group", lambda: None)
+    try:
+        if delivered:
+            assert connector.put("0", "1", key, {"encoder_output": expected})[0]
+        for rank in range(shard_size * replicate_size):
+            shard_rank = rank % shard_size
+            replicate_rank = rank // shard_size
+            shard_group = MeshGroup("shard", replicate_rank, shard_rank, shard_size)
+            replicate_group = MeshGroup("replicate", shard_rank, replicate_rank, replicate_size)
+            monkeypatch.setattr(parallel_state, "get_fs_group", lambda: shard_group)
+            monkeypatch.setattr(parallel_state, "get_hsdp_replicate_group", lambda: replicate_group)
+            runner = _make_runner(connector, payload_keys=("encoder_output",))
+            monkeypatch.setattr(runner, "_get_local_tp_group", lambda: None)
+            request = _make_request({"prompt": "test"}, request_id=request_id)
+            if delivered:
+                runner._maybe_recv_stage_payload(request)
+                torch.testing.assert_close(request.prompt["additional_information"]["encoder_output"], expected)
+            else:
+                with pytest.raises(RuntimeError, match="Stage payload unavailable"):
+                    runner._maybe_recv_stage_payload(request)
+        assert set(calls) == {0}
+        if delivered:
+            assert calls == [0]
+            assert connector.get("0", "1", key) is None
+    finally:
+        connector.close()
+
+
+def test_key_convention_fetch_merges_into_additional_information():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector)
+    req = _make_request({"prompt": "a cat"})
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert connector.calls == [("0", "1", "req-7_0_0", None)]
+    assert set(req.prompt["additional_information"]) == {"text_encoder_output"}
+    assert req.prompt["additional_information"]["text_encoder_output"]["hidden_states"].shape == (4, 8)
+
+
+def test_handle_path_uses_its_own_key_and_metadata():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector, payload_keys=())
+    handle = {
+        "key": "custom-key",
+        "from_stage": "3",
+        "to_stage": "4",
+        "metadata": {"schema_version": 1},
+    }
+    req = _make_request({"prompt": "a cat", HANDLE_KEY: handle})
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert connector.calls == [("3", "4", "custom-key", {"schema_version": 1})]
+    # The handle is transport bookkeeping and must not reach the pipeline.
+    assert HANDLE_KEY not in req.prompt
+    assert "text_encoder_output" in req.prompt["additional_information"]
+
+
+def test_receive_retries_until_producer_publishes(monkeypatch):
+    connector = _FakeConnector(_conditioning())
+    get = connector.get
+    attempts = []
+
+    def delayed_get(*args, **kwargs):
+        attempts.append(args)
+        return get(*args, **kwargs) if len(attempts) == 3 else None
+
+    monkeypatch.setattr(connector, "get", delayed_get)
+    runner = _make_runner(connector)
+    request = _make_request({"prompt": "a cat"})
+
+    runner._maybe_recv_stage_payload(request)
+
+    assert len(attempts) == 3
+    assert "text_encoder_output" in request.prompt["additional_information"]
+
+
+def test_receive_retry_budget_preserves_inline_payload(monkeypatch):
+    from vllm_omni.distributed.omni_connectors.model_runner import omni_connector_payload_transport
+
+    clock = [0.0]
+
+    def advance(duration):
+        clock[0] += duration
+
+    monkeypatch.setattr(
+        omni_connector_payload_transport, "time", SimpleNamespace(monotonic=lambda: clock[0], sleep=advance)
+    )
+    connector = _FakeConnector(None)
+    runner = _make_runner(connector)
+    inline = _conditioning()
+    request = _make_request({"additional_information": inline})
+
+    runner._maybe_recv_stage_payload(request)
+
+    assert 1 < len(connector.calls) <= 41
+    assert clock[0] == pytest.approx(2.0)
+    assert request.prompt["additional_information"] is inline
+
+
+def test_stage_without_declared_keys_never_touches_the_connector():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector, payload_keys=())
+    req = _make_request({"prompt": "a cat"})
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert connector.calls == []
+    assert "additional_information" not in req.prompt
+
+
+def test_undeclared_payload_keys_are_dropped():
+    connector = _FakeConnector({**_conditioning(), "debug_blob": torch.zeros(2)})
+    runner = _make_runner(connector)
+    req = _make_request({"prompt": "a cat"})
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert set(req.prompt["additional_information"]) == {"text_encoder_output"}
+
+
+@pytest.mark.parametrize(
+    "connector",
+    [
+        _FakeConnector(None),
+        _FakeConnector(raises=True),
+        _FakeConnector("not-a-dict"),
+    ],
+    ids=["missing", "raises", "wrong_type"],
+)
+def test_failed_transfer_falls_back_to_the_inline_prompt(connector):
+    runner = _make_runner(connector)
+    inline = {"text_encoder_output": {"hidden_states": torch.ones(2, 8)}}
+    req = _make_request({"prompt": "a cat", "additional_information": dict(inline)})
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert torch.equal(
+        req.prompt["additional_information"]["text_encoder_output"]["hidden_states"],
+        inline["text_encoder_output"]["hidden_states"],
+    )
+
+
+def test_legacy_sender_info_is_passed_as_request_metadata():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector)
+    req = _make_request({"prompt": "a cat"}, kv_sender_info={0: {"host": "10.0.0.1", "zmq_port": 50171}})
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert runner.kv_transfer_manager.sender_info_calls == []
+    assert connector.calls[0][3] == {"source_host": "10.0.0.1", "source_port": 50171}
+
+
+def test_payload_sender_info_overrides_kv_sender_info_for_full_payloads():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector)
+    req = _make_request(
+        {"prompt": "a cat"},
+        kv_sender_info={0: {"host": "10.0.0.1", "zmq_port": 50171}},
+        payload_sender_info={"host": "10.0.0.1", "zmq_port": 50071},
+    )
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert runner.kv_transfer_manager.sender_info_calls == []
+    assert connector.calls[0][3] == {"source_host": "10.0.0.1", "source_port": 50071}
+
+
+def test_synchronous_receive_keeps_endpoints_request_scoped_and_uses_external_ids():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector, recv_stages=("2", "5"))
+    for index in range(2):
+        request = _make_request(
+            {},
+            request_id=f"internal-{index}",
+            payload_sender_info={"host": f"10.0.0.{index + 1}", "zmq_port": 50071 + index},
+        )
+        request.external_req_id = f"external-{index}"
+        runner._maybe_recv_stage_payload(request)
+        assert "text_encoder_output" in request.prompt["additional_information"]
+    assert connector.calls == [
+        ("2", "5", "external-0_2_0", {"source_host": "10.0.0.1", "source_port": 50071}),
+        ("2", "5", "external-1_2_0", {"source_host": "10.0.0.2", "source_port": 50072}),
+    ]
+    assert runner.kv_transfer_manager.sender_info_calls == []
+    assert runner._pending_load_reqs == {}
+    assert runner._get_req_chunk == {}
+
+
+@pytest.mark.parametrize("stage_key", [2, "2"])
+@pytest.mark.parametrize("use_handle", [False, True])
+def test_sender_map_selects_actual_source_stage(stage_key, use_handle):
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector, recv_stages=("0" if use_handle else "2", "5"))
+    prompt = {}
+    if use_handle:
+        prompt[HANDLE_KEY] = {"key": "explicit", "from_stage": "2", "to_stage": "5"}
+    request = _make_request(
+        prompt,
+        payload_sender_info={
+            "0": {"host": "wrong", "zmq_port": 50000},
+            stage_key: {"host": "right", "zmq_port": 50002},
+        },
+    )
+    runner._maybe_recv_stage_payload(request)
+    assert connector.calls[0][0:2] == ("2", "5")
+    assert connector.calls[0][3] == {"source_host": "right", "source_port": 50002}
+
+
+def test_tp_payload_is_fetched_once_by_leader_and_broadcast_to_followers():
+    connector = _FakeConnector(_conditioning())
+    state = {}
+
+    class _FakeTPGroup:
+        world_size = 4
+
+        def __init__(self, rank):
+            self.rank_in_group = rank
+
+        def broadcast_object(self, value, src=0):
+            if self.rank_in_group == src:
+                state["delivered"] = value
+            return state["delivered"]
+
+        def broadcast_tensor_dict(self, value, src=0):
+            if self.rank_in_group == src:
+                state["payload"] = value
+            return state["payload"]
+
+    leader = _make_runner(connector)
+    leader._local_rank = 0
+    leader._get_local_tp_group = lambda: _FakeTPGroup(0)
+    follower = _make_runner(connector)
+    follower._local_rank = 1
+    follower._get_local_tp_group = lambda: _FakeTPGroup(1)
+
+    leader_req = _make_request({"prompt": "a cat"})
+    follower_req = _make_request({"prompt": "a cat"})
+    leader._maybe_recv_stage_payload(leader_req)
+    follower._maybe_recv_stage_payload(follower_req)
+
+    assert len(connector.calls) == 1
+    leader_output = leader_req.prompt["additional_information"]["text_encoder_output"]["hidden_states"]
+    follower_output = follower_req.prompt["additional_information"]["text_encoder_output"]["hidden_states"]
+    assert torch.equal(leader_output, follower_output)
+
+
+def test_tp_payload_miss_is_broadcast_without_follower_connector_access():
+    connector = _FakeConnector(None)
+    state = {}
+
+    class _FakeTPGroup:
+        world_size = 2
+
+        def __init__(self, rank):
+            self.rank_in_group = rank
+
+        def broadcast_object(self, value, src=0):
+            if self.rank_in_group == src:
+                state["delivered"] = value
+            return state["delivered"]
+
+        def broadcast_tensor_dict(self, value, src=0):
+            raise AssertionError("missing payload must not be broadcast")
+
+    leader = _make_runner(connector)
+    leader._get_local_tp_group = lambda: _FakeTPGroup(0)
+    follower = _make_runner(connector)
+    follower._local_rank = 1
+    follower._get_local_tp_group = lambda: _FakeTPGroup(1)
+
+    with pytest.raises(RuntimeError, match="Stage payload unavailable"):
+        leader._maybe_recv_stage_payload(_make_request({"prompt": "a cat"}))
+    with pytest.raises(RuntimeError, match="Stage payload unavailable"):
+        follower._maybe_recv_stage_payload(_make_request({"prompt": "a cat"}))
+
+    assert 1 < len(connector.calls) <= 41
+
+
+def test_sp_payload_is_fetched_once_by_leader_and_broadcast_to_followers():
+    connector = _FakeConnector(_conditioning())
+    state = {}
+
+    class _FakeSPGroup:
+        world_size = 4
+
+        def __init__(self, rank):
+            self.rank_in_group = rank
+
+        def broadcast_object(self, value, src=0):
+            if self.rank_in_group == src:
+                state["delivered"] = value
+            return state["delivered"]
+
+        def broadcast_tensor_dict(self, value, src=0):
+            if self.rank_in_group == src:
+                state["payload"] = value
+            return state["payload"]
+
+    leader = _make_runner(connector)
+    leader._get_local_tp_group = lambda: None
+    leader._stage_payload_broadcast_groups = lambda: (_FakeSPGroup(0),)
+    follower = _make_runner(connector)
+    follower._local_rank = 1
+    follower._get_local_tp_group = lambda: None
+    follower._stage_payload_broadcast_groups = lambda: (_FakeSPGroup(1),)
+
+    leader_req = _make_request({"prompt": "a cat"})
+    follower_req = _make_request({"prompt": "a cat"})
+    leader._maybe_recv_stage_payload(leader_req)
+    follower._maybe_recv_stage_payload(follower_req)
+
+    assert len(connector.calls) == 1
+    leader_output = leader_req.prompt["additional_information"]["text_encoder_output"]["hidden_states"]
+    follower_output = follower_req.prompt["additional_information"]["text_encoder_output"]["hidden_states"]
+    assert torch.equal(leader_output, follower_output)
+
+
+def test_missing_incoming_edge_is_reported_not_fetched():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector, recv_stages=(None, None))
+    req = _make_request({"prompt": "a cat"})
+
+    with pytest.raises(RuntimeError, match="no incoming edge"):
+        runner._maybe_recv_stage_payload(req)
+
+    assert connector.calls == []
+    assert "additional_information" not in req.prompt
+
+
+@pytest.mark.parametrize("delivered", [True, False])
+def test_payload_reaches_full_tp_sp_grid_once(monkeypatch, delivered):
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    connector = _FakeConnector(_conditioning() if delivered else None)
+    packets = {}
+
+    class GridGroup:
+        world_size = 2
+
+        def __init__(self, dimension, peer_rank, rank):
+            self.key = (dimension, peer_rank)
+            self.rank_in_group = rank
+
+        def broadcast_object(self, value, src=0):
+            if self.rank_in_group == src:
+                packets[self.key, "delivered"] = value
+            return packets[self.key, "delivered"]
+
+        def broadcast_tensor_dict(self, value, src=0):
+            if self.rank_in_group == src:
+                packets[self.key, "payload"] = value
+            return packets[self.key, "payload"]
+
+    for sp_rank in range(2):
+        for tp_rank in range(2):
+            runner = _make_runner(connector)
+            tp_group = GridGroup("tp", sp_rank, tp_rank)
+            sp_group = GridGroup("sp", tp_rank, sp_rank)
+            monkeypatch.setattr(runner, "_get_local_tp_group", lambda: tp_group)
+            monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group)
+            request = _make_request({"prompt": "a cat"})
+
+            if delivered:
+                runner._maybe_recv_stage_payload(request)
+                output = request.prompt["additional_information"]["text_encoder_output"]
+                assert torch.equal(output["hidden_states"], connector._payload["text_encoder_output"]["hidden_states"])
+            else:
+                with pytest.raises(RuntimeError, match="Stage payload unavailable"):
+                    runner._maybe_recv_stage_payload(request)
+    if delivered:
+        assert len(connector.calls) == 1
+    else:
+        assert 1 < len(connector.calls) <= 41
+
+
+def test_native_kv_runner_keeps_synchronous_payload_transport(monkeypatch):
+    from unittest.mock import Mock
+
+    from vllm_omni.diffusion.worker import diffusion_model_runner as runner_module
+
+    connector = _FakeConnector(_conditioning())
+    manager = _FakeKVTransferManager(connector)
+    config = SimpleNamespace(stage_input_payload_keys=("text_encoder_output",), stage_id=1, kv_transfer_config=object())
+    monkeypatch.setattr(runner_module.OmniKVTransferManager, "from_od_config", lambda config: manager)
+    monkeypatch.setattr(runner_module, "DiffusionKVModelRunnerBackend", Mock())
+    runner = DiffusionModelRunner(SimpleNamespace(), config, torch.device("cpu"))
+    runner._initialize_generator = Mock()
+    request = _make_request({"prompt": "a cat"})
+    request.sampling_params = SimpleNamespace()
+
+    runner._prepare_request_for_forward(request, od_config=config)
+
+    assert runner.kv_transfer_manager is None
+    assert runner._kv_transfer_manager is manager
+    assert runner._kv_connector is None
+    assert not runner._kv_prefetch_enabled
+    assert len(connector.calls) == 1
+    assert "text_encoder_output" in request.prompt["additional_information"]
+    runner._initialize_generator.assert_called_once_with(request.sampling_params)
+
+
+def test_step_mode_receives_payload_before_kv_and_reuses_cached_state():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector)
+    runner.state_cache = {}
+    request = _make_request({"prompt": "a cat"})
+    request.sampling_params = SimpleNamespace()
+    kv_calls = []
+
+    def receive_kv(state_request, **kwargs):
+        assert "text_encoder_output" in state_request.prompt["additional_information"]
+        kv_calls.append(state_request.request_id)
+
+    runner.kv_transfer_manager.receive_multi_kv_cache_distributed = receive_kv
+    scheduled = SimpleNamespace(
+        scheduled_new_reqs=[SimpleNamespace(request_id=request.request_id, req=request, diffusion_kv_metadata=None)],
+        scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
+    )
+    states, new_ids = runner._update_states(scheduled)
+    assert new_ids == [request.request_id]
+    assert states[0].prompt is request.prompt
+    assert len(connector.calls) == 1
+    assert kv_calls == [request.request_id]
+
+    scheduled.scheduled_new_reqs = []
+    scheduled.scheduled_cached_reqs.request_ids = [request.request_id]
+    cached_states, new_ids = runner._update_states(scheduled)
+    assert cached_states == states
+    assert new_ids == []
+    assert len(connector.calls) == 1
+    assert kv_calls == [request.request_id]
+
+
+def test_non_dict_prompt_is_left_alone():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector)
+    req = _make_request("a cat")
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert connector.calls == []
+    assert req.prompt == "a cat"
+
+
+@pytest.mark.parametrize("request_id", ["dummy_req_id", "dummy_req_id/profile"])
+def test_warmup_never_receives_or_publishes_payload(request_id):
+    connector = _FakeConnector(raises=True)
+    request = _make_request({}, request_id=request_id)
+    receiver = _make_runner(connector)
+    receiver._maybe_recv_stage_payload(request)
+    sender = _make_sender(connector)
+    sender._maybe_send_stage_payload([request], [_make_output(prompt_embeds=torch.ones(1))])
+    assert connector.calls == []
+    assert connector.put_calls == []
+
+
+@pytest.mark.parametrize("inline", [{}, {"text_encoder_output": None}])
+def test_failed_transfer_requires_complete_inline_payload(inline):
+    runner = _make_runner(_FakeConnector(raises=True))
+    with pytest.raises(RuntimeError, match="missing keys.*text_encoder_output"):
+        runner._maybe_recv_stage_payload(_make_request({"additional_information": inline}))
+
+
+def test_send_puts_declared_keys_and_attaches_a_handle():
+    connector = _FakeConnector()
+    runner = _make_sender(connector)
+    req = _make_request({"prompt": "a cat"})
+    output = _make_output(prompt_embeds=torch.zeros(2, 8), latents=torch.zeros(1, 4))
+
+    runner._maybe_send_stage_payload([req], [output])
+
+    from_stage, to_stage, put_key, data = connector.put_calls[0]
+    assert (from_stage, to_stage, put_key) == ("0", "1", "req-7_0_0")
+    # Only declared keys travel over the connector.
+    assert isinstance(data, dict)
+    assert set(data) == {"prompt_embeds"}
+    handle = output.custom_output[HANDLE_KEY]
+    assert handle["key"] == "req-7_0_0"
+    assert handle["metadata"] == {"schema_version": 1}
+    assert handle["payload_keys"] == ["prompt_embeds"]
+    assert handle["size_bytes"] == 128
+    # Transferred keys must not also travel inline through the orchestrator.
+    assert "prompt_embeds" not in output.custom_output
+    # Undeclared keys are untouched.
+    assert "latents" in output.custom_output
+
+
+def test_send_key_matches_the_receive_key_convention():
+    sender_connector = _FakeConnector()
+    sender = _make_sender(sender_connector)
+    req = _make_request({"prompt": "a cat"})
+    sender._maybe_send_stage_payload([req], [_make_output(prompt_embeds=torch.zeros(2, 8))])
+
+    recv_connector = _FakeConnector(_conditioning())
+    receiver = _make_runner(recv_connector)
+    receiver._maybe_recv_stage_payload(_make_request({"prompt": "a cat"}))
+
+    assert sender_connector.put_calls[0][2] == recv_connector.calls[0][2]
+
+
+def test_stage_without_declared_output_keys_never_puts():
+    connector = _FakeConnector()
+    runner = _make_sender(connector, payload_keys=())
+    output = _make_output(prompt_embeds=torch.zeros(2, 8))
+
+    runner._maybe_send_stage_payload([_make_request({"prompt": "a cat"})], [output])
+
+    assert connector.put_calls == []
+    assert HANDLE_KEY not in output.custom_output
+
+
+def test_missing_outgoing_edge_is_reported_not_sent():
+    connector = _FakeConnector()
+    runner = _make_sender(connector, send_stages=(None, None))
+    output = _make_output(prompt_embeds=torch.zeros(2, 8))
+
+    runner._maybe_send_stage_payload([_make_request({"prompt": "a cat"})], [output])
+
+    assert connector.put_calls == []
+    assert HANDLE_KEY not in output.custom_output
+
+
+def test_output_without_the_declared_keys_is_skipped():
+    connector = _FakeConnector()
+    runner = _make_sender(connector)
+    output = _make_output(latents=torch.zeros(1, 4))
+
+    runner._maybe_send_stage_payload([_make_request({"prompt": "a cat"})], [output])
+
+    assert connector.put_calls == []
+    assert HANDLE_KEY not in output.custom_output
+
+
+@pytest.mark.parametrize(
+    "connector",
+    [
+        _FakeConnector(raises=True),
+        _FakeConnector(put_result=(False, 0, None)),
+    ],
+    ids=["raises", "rejected"],
+)
+def test_failed_send_leaves_the_inline_payload_unannotated(connector):
+    runner = _make_sender(connector)
+    output = _make_output(prompt_embeds=torch.zeros(2, 8))
+
+    runner._maybe_send_stage_payload([_make_request({"prompt": "a cat"})], [output])
+
+    assert HANDLE_KEY not in output.custom_output
+    assert "prompt_embeds" in output.custom_output
+
+
+class _FakeTPGroup:
+    def __init__(self, rank_in_group, world_size=2):
+        self.rank_in_group = rank_in_group
+        self.world_size = world_size
+        self.broadcast_calls: list[object] = []
+        self.leader_packet: object = None
+
+    def broadcast_object(self, obj, src=0):
+        self.broadcast_calls.append(obj)
+        return obj if self.rank_in_group == src else self.leader_packet
+
+
+@pytest.mark.parametrize("rank_in_group", [0, 1], ids=["leader", "follower"])
+def test_send_puts_only_on_the_tp_leader(monkeypatch, rank_in_group):
+    connector = _FakeConnector()
+    runner = _make_sender(connector)
+    tp_group = _FakeTPGroup(rank_in_group)
+    expected_handle = {
+        "key": "req-7_0_0",
+        "from_stage": "0",
+        "to_stage": "1",
+        "size_bytes": 128,
+        "metadata": {"schema_version": 1},
+        "payload_keys": ["prompt_embeds"],
+    }
+    tp_group.leader_packet = {"req-7": expected_handle}
+    monkeypatch.setattr(DiffusionModelRunner, "_get_local_tp_group", staticmethod(lambda: tp_group))
+    output = _make_output(prompt_embeds=torch.zeros(2, 8))
+
+    runner._maybe_send_stage_payload([_make_request({"prompt": "a cat"})], [output])
+
+    assert len(connector.put_calls) == (1 if rank_in_group == 0 else 0)
+    # Every rank ends up reporting the same handle, leader or not.
+    assert output.custom_output[HANDLE_KEY] == expected_handle
+    # ...and every rank drops the inline copy, not just the sender.
+    assert "prompt_embeds" not in output.custom_output
+    assert len(tp_group.broadcast_calls) == 1
+
+
+@pytest.mark.parametrize("tp_size", [1, 2])
+@pytest.mark.parametrize("accepted", [True, False])
+def test_send_elects_one_tp_sp_leader(monkeypatch, tp_size, accepted):
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    connector = _FakeConnector(put_result=(accepted, 128, {"schema_version": 1}))
+    packets = {}
+    outputs = []
+
+    class GridGroup:
+        def __init__(self, dimension, peer_rank, rank, world_size):
+            self.key = (dimension, peer_rank)
+            self.rank_in_group = rank
+            self.world_size = world_size
+
+        def broadcast_object(self, value, src=0):
+            if self.rank_in_group == src:
+                packets[self.key] = value
+            return packets[self.key]
+
+    for sp_rank in range(2):
+        for tp_rank in range(tp_size):
+            runner = _make_sender(connector)
+            tp_group = GridGroup("tp", sp_rank, tp_rank, tp_size)
+            sp_group = GridGroup("sp", tp_rank, sp_rank, 2)
+            monkeypatch.setattr(runner, "_get_local_tp_group", lambda: tp_group)
+            monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group)
+            output = _make_output(prompt_embeds=torch.zeros(2, 8))
+
+            runner._maybe_send_stage_payload([_make_request({})], [output])
+
+            outputs.append(output.custom_output)
+            assert (HANDLE_KEY in output.custom_output) == accepted
+            assert ("prompt_embeds" in output.custom_output) != accepted
+
+    assert len(connector.put_calls) == 1
+    if accepted:
+        assert all(output[HANDLE_KEY] == outputs[0][HANDLE_KEY] for output in outputs)

tests/diffusion/test_diffusion_step_pipeline.py

@@ -616,7 +616,8 @@ def test_make_batch_explicit_mapping_still_selects_and_orders_states():
 
 
 @pytest.mark.cpu
-def test_step_profiler_reports_denoise_step_as_diffuse():
+def test_step_profiler_reports_denoise_step_as_diffuse(monkeypatch):
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
     pipeline = _AutoDenoiseProfilerPipeline()
 
     assert pipeline.denoise_step() == "ok"
@@ -629,6 +630,13 @@ def test_step_profiler_reports_denoise_step_as_diffuse():
 class TestRunner:
     """DiffusionModelRunner.execute_stepwise"""
 
+    @pytest.fixture(autouse=True)
+    def mock_platform_memory(self, monkeypatch):
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "is_available", lambda: True)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+
     @pytest.mark.parametrize(
         ("diffusion_kv_mode", "should_cleanup"),
         [

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -87,6 +87,38 @@ async def test_inline_dispatch_request_success(client, mock_engine):
     assert output.request_id == "req-1"
 
 
+@pytest.mark.parametrize("streaming", [False, True])
+def test_inline_dispatch_preserves_payload_sender_info(client, mock_engine, streaming):
+    async def _run_test():
+        requests = []
+
+        async def _step_streaming(request):
+            requests.append(request)
+            yield [OmniRequestOutput.from_diffusion(request_id=request.request_id, images=[MagicMock()])]
+
+        mock_engine.step_streaming = _step_streaming
+        client.od_config.streaming_output = streaming
+        payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50071}
+        kv_transfer_params = {"remote_engine_id": "mooncake-producer", "remote_block_ids": [1, 2]}
+        await client.add_request_async(
+            "req-payload",
+            "A test prompt",
+            OmniDiffusionSamplingParams(),
+            kv_transfer_params=kv_transfer_params,
+            payload_sender_info=payload_sender_info,
+        )
+
+        for _ in range(10):
+            if requests:
+                break
+            await asyncio.sleep(0.01)
+
+        assert requests[0].payload_sender_info == payload_sender_info
+        assert requests[0].kv_transfer_params == kv_transfer_params
+
+    asyncio.run(_run_test())
+
+
 @pytest.mark.asyncio
 async def test_inline_non_streaming_dispatches_lifecycle_before_final(client, mock_engine):
     lifecycle = OmniRequestOutput.from_diffusion(

tests/diffusion/test_stage_payload_kv_coexistence.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import asyncio
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("streaming", [False, True])
+def test_subprocess_client_and_worker_preserve_both_transfer_contracts(streaming):
+    """Exercise the IPC message and both process request builders without GPUs."""
+
+    async def run():
+        client = object.__new__(StageDiffusionClient)
+        client._engine_dead = False
+        client.stage_id = 1
+        client.replica_id = 0
+        client._encoder = MagicMock()
+        client._encoder.encode.side_effect = lambda message: message
+        client._request_socket = MagicMock()
+        kv_params = {"remote_engine_id": "mooncake-producer", "remote_block_ids": [1, 2]}
+        payload_sender = {"host": "10.0.0.1", "zmq_port": 50071}
+        kv_sender = {0: {"host": "10.0.0.2", "zmq_port": 50171}}
+        # Keep main's existing fifth positional argument bound to native KV.
+        await client.add_request_async(
+            "req-both",
+            "prompt",
+            OmniDiffusionSamplingParams(),
+            kv_sender,
+            kv_params,
+            payload_sender_info=payload_sender,
+        )
+        message = client._request_socket.send.call_args.args[0]
+        assert message["kv_transfer_params"] == kv_params
+        assert message["payload_sender_info"] == payload_sender
+        captured = []
+
+        class Engine:
+            async def step_streaming(self, request):
+                captured.append(request)
+                yield [OmniRequestOutput.from_diffusion(request_id=request.request_id, images=[])]
+
+        proc = object.__new__(StageDiffusionProc)
+        proc._engine = Engine()
+        args = (message["request_id"], message["prompt"], message["sampling_params"])
+        transfer_args = {
+            "kv_sender_info": kv_sender,
+            "kv_transfer_params": kv_params,
+            "payload_sender_info": message["payload_sender_info"],
+        }
+        if streaming:
+            outputs = [output async for output in proc._process_streaming_request(*args, **transfer_args)]
+        else:
+            outputs = [await proc._process_request(*args, **transfer_args)]
+        assert outputs[0].request_id == "req-both"
+        assert captured[0].kv_transfer_params == kv_params
+        assert captured[0].kv_sender_info == kv_sender
+        assert captured[0].payload_sender_info == payload_sender
+
+    asyncio.run(run())

tests/diffusion/test_worker_wrapper_base.py

@@ -12,10 +12,13 @@ This module tests the WorkerWrapperBase implementation:
 - Dynamic worker class extension
 """
 
+import gc
+import weakref
 from types import SimpleNamespace
 from typing import Any
 
 import pytest
+import torch
 from pytest_mock import MockerFixture
 
 from vllm_omni.diffusion.worker.diffusion_worker import (
@@ -227,12 +230,13 @@ class TestWorkerWrapperBaseDelegation:
         wrapper.worker.shutdown.assert_called_once()
         assert result is None
 
-    def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture):
+    @pytest.mark.parametrize("native_kv", [False, True])
+    def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture, native_kv):
         events: list[str] = []
         offload_backend = mocker.Mock()
         offload_backend.disable.side_effect = lambda: events.append("offload")
         kv_manager = mocker.Mock()
-        kv_manager.shutdown_prefetch.side_effect = lambda: events.append("kv")
+        kv_manager.close.side_effect = lambda: events.append("kv")
         destroy = mocker.patch(
             "vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env",
             side_effect=lambda: events.append("distributed"),
@@ -240,14 +244,71 @@ class TestWorkerWrapperBaseDelegation:
         worker = DiffusionWorker.__new__(DiffusionWorker)
         worker.model_runner = SimpleNamespace(
             offload_backend=offload_backend,
-            kv_transfer_manager=kv_manager,
+            kv_transfer_manager=None if native_kv else kv_manager,
+            _kv_transfer_manager=kv_manager,
         )
 
         worker.shutdown()
 
         assert events == ["offload", "kv", "distributed"]
+        kv_manager.close.assert_called_once_with()
         destroy.assert_called_once_with()
 
+    @pytest.mark.parametrize("cleanup_fails", [False, True])
+    def test_worker_shutdown_releases_model_references(self, mocker: MockerFixture, cleanup_fails):
+        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env")
+        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.shutdown_kv_connector")
+        offload_backend = mocker.Mock()
+        if cleanup_fails:
+            offload_backend.disable.side_effect = RuntimeError("offload cleanup failed")
+        worker = DiffusionWorker.__new__(DiffusionWorker)
+        weights = torch.ones(1)
+        weights_ref = weakref.ref(weights)
+        pipeline = SimpleNamespace(weights=weights)
+        worker.model_runner = SimpleNamespace(
+            pipeline=pipeline, offload_backend=offload_backend, kv_transfer_manager=None
+        )
+        worker.lora_manager = SimpleNamespace(pipeline=pipeline)
+        worker._sleep_saved_buffers = {"buffer": weights}
+        del weights, pipeline
+
+        if cleanup_fails:
+            with pytest.raises(RuntimeError, match="offload cleanup failed"):
+                worker.shutdown()
+        else:
+            worker.shutdown()
+        gc.collect()
+
+        assert worker.model_runner is None
+        assert worker.lora_manager is None
+        assert worker._sleep_saved_buffers == {}
+        assert weights_ref() is None
+
+    @pytest.mark.parametrize("owns_pool", [False, True])
+    @pytest.mark.parametrize("cleanup_fails", [False, True])
+    def test_worker_shutdown_releases_only_owned_sleep_pool(self, mocker: MockerFixture, owns_pool, cleanup_fails):
+        destroy = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env")
+        shutdown_connector = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.shutdown_kv_connector")
+        allocator_class = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker._get_cumem_allocator_class")
+        release = allocator_class.return_value.get_instance.return_value.release_pools
+        if cleanup_fails:
+            release.side_effect = RuntimeError("pool release failed")
+        worker = DiffusionWorker.__new__(DiffusionWorker)
+        worker.model_runner = None
+        worker._owns_sleep_pool = owns_pool
+
+        if owns_pool and cleanup_fails:
+            with pytest.raises(RuntimeError, match="pool release failed"):
+                worker.shutdown()
+        else:
+            worker.shutdown()
+
+        assert release.call_count == int(owns_pool)
+        shutdown_connector.assert_called_once_with()
+        destroy.assert_called_once_with()
+        if not cleanup_fails:
+            assert not worker._owns_sleep_pool
+
 
 # -------------------------------------------------------------------------
 # Tests: execute_method
@@ -507,6 +568,26 @@ class TestCustomPipelineWorkerExtension:
         # Verify re_init_pipeline was called with custom_pipeline_args
         mock_worker_instance.re_init_pipeline.assert_called_once_with(custom_args)
 
+    @pytest.mark.parametrize("cleanup_fails", [False, True])
+    def test_custom_pipeline_initialization_failure_cleans_worker(
+        self, mocker: MockerFixture, mock_od_config, cleanup_fails
+    ):
+        worker = mocker.Mock()
+        worker.re_init_pipeline.side_effect = RuntimeError("custom pipeline load failed")
+        if cleanup_fails:
+            worker.shutdown.side_effect = RuntimeError("cleanup failed")
+        worker_class = mocker.Mock(return_value=worker)
+        mocker.patch.object(WorkerWrapperBase, "_prepare_worker_class", return_value=worker_class)
+
+        with pytest.raises(RuntimeError, match="custom pipeline load failed"):
+            WorkerWrapperBase(
+                gpu_id=0,
+                od_config=mock_od_config,
+                custom_pipeline_args={"pipeline_class": "unused"},
+            )
+
+        worker.shutdown.assert_called_once_with()
+
     def test_custom_pipeline_with_explicit_extension(self, mocker: MockerFixture, mock_od_config):
         """Test that explicit worker_extension_cls is preserved when custom_pipeline_args is provided."""