[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/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():