[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/model_executor/models/test_minimax_h3_encoder_payload.py

@@ -91,8 +91,11 @@ def test_pipeline_owns_three_encoders_then_dit_and_decoders() -> None:
 
     encoder, diffusion = MINIMAX_H3_PIPELINE.stages
     assert encoder.model_arch == "MiniMaxH3Encoder"
-    assert encoder.custom_process_next_stage_input_func is None
+    assert encoder.custom_process_next_stage_input_func == (
+        "vllm_omni.model_executor.stage_input_processors.minimax_h3.encoder2diffusion_full_payload"
+    )
     assert diffusion.model_arch == "MiniMaxH3Pipeline"
+    assert diffusion.stage_input_payload_keys == ("encoder_output",)
     assert diffusion.requires_multimodal_data is False
     assert diffusion.requires_full_payload_input is False
 
@@ -441,7 +444,7 @@ def test_encode_media_keeps_audio_budgets_and_component_residency_separate(
     else:
         with pytest.raises(ValueError, match="at most 15 seconds"):
             encode_media(media, **kwargs)
-    assert scopes == [video_vae, audio_vae]
+    assert scopes == ([video_vae, audio_vae] if valid else [])
 
 
 def test_encoder_output_reuses_encoder_handoff_and_round_trips(monkeypatch) -> None:

tests/model_executor/stage_input_processors/test_minimax_h3.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Regression tests for MiniMax H3's disaggregated text-encoder contract."""
+"""Regression tests for MiniMax H3's disaggregated encoder contract."""
 
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -10,7 +10,7 @@ import pytest
 import torch
 from PIL import Image
 
-from vllm_omni.data_entry_keys import flatten_payload
+from vllm_omni.data_entry_keys import flatten_payload, to_struct
 from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import resolve_minimax_h3_diffusion_model_path
 from vllm_omni.engine.serialization import (
     deserialize_additional_information,
@@ -24,8 +24,10 @@ from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
 )
 from vllm_omni.model_executor.models.minimax_h3.conditioning import (
     MINIMAX_H3_CONDITION_LABELS_KEY,
+    MINIMAX_H3_ENCODER_LAYOUT_KEY,
     MINIMAX_H3_ENCODER_REQUEST_KEY,
     MINIMAX_H3_PRESENTATION_TASK_KEY,
+    MINIMAX_H3_TEXT_CONDITIONING_SCHEMA,
     MiniMaxH3EncoderConditioning,
 )
 from vllm_omni.model_executor.models.minimax_h3.encoder_processing import _audio_items, _load_audio
@@ -40,6 +42,7 @@ from vllm_omni.model_executor.models.minimax_h3.text_encoder import (
 from vllm_omni.model_executor.stage_input_processors.minimax_h3 import (
     _diffusion_sampling_params,
     encoder2diffusion,
+    encoder2diffusion_full_payload,
     prepare_encoder_prompt,
 )
 
@@ -208,7 +211,8 @@ def test_prepare_ref2va_rejects_short_embedded_video_audio(monkeypatch):
         prepare_encoder_prompt(prompt, [sampling])
 
 
-def test_prepare_ref2va_keeps_embedded_and_standalone_audio_budgets_separate(monkeypatch):
+@pytest.mark.parametrize("standalone_count", [1, 2])
+def test_prepare_ref2va_uses_separate_embedded_and_standalone_audio_budgets(monkeypatch, standalone_count):
     duration_seconds = 8.0
     sample_rate = 16_000
     _mock_ref2va_video_with_audio(monkeypatch, duration_seconds=duration_seconds)
@@ -222,16 +226,21 @@ def test_prepare_ref2va_keeps_embedded_and_standalone_audio_budgets_separate(mon
         "prompt": "hello",
         "multi_modal_data": {
             "video": "original.mp4",
-            "audio": (torch.zeros(int(duration_seconds * sample_rate)), sample_rate),
+            "audio": [(torch.zeros(int(duration_seconds * sample_rate)), sample_rate) for _ in range(standalone_count)],
         },
     }
 
-    transformed = prepare_encoder_prompt(prompt, [sampling])
-    assert transformed["mm_processor_kwargs"][MINIMAX_H3_CONDITION_LABELS_KEY] == [
-        ("audio", 1),
-        ("video", 1),
-        ("audio", 2),
-    ]
+    if standalone_count == 2:
+        with pytest.raises(OmniClientError, match="at most 15 seconds in total"):
+            prepare_encoder_prompt(prompt, [sampling])
+    else:
+        transformed = prepare_encoder_prompt(prompt, [sampling])
+        from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder
+
+        media = MiniMaxH3Encoder._media_input(transformed["additional_information"])
+        assert len(media.video_audios) == 1 and len(media.audios) == 1
+        assert media.video_audios[0][0].shape[-1] == 8 * sample_rate
+        assert media.audios[0][0].shape[-1] == 8 * sample_rate
 
 
 def _encoder_output() -> dict:
@@ -283,6 +292,60 @@ def test_encoder2diffusion_waits_for_one_finished_source() -> None:
         encoder2diffusion([SimpleNamespace(), SimpleNamespace()], {"prompt": "hello"})
 
 
+@pytest.mark.parametrize(
+    ("case", "message"),
+    [
+        ("hidden_width", r"hidden_states must have shape \[tokens, 5120\]"),
+        ("hidden_dtype", "hidden_states must have dtype torch.bfloat16"),
+        ("hidden_noncontiguous", "hidden_states must use contiguous strided layout"),
+        ("hidden_non_strided", "hidden_states must use contiguous strided layout"),
+        ("token_dtype", "token_tags must have dtype torch.int64"),
+        ("token_layout", "token_tags must use contiguous strided layout"),
+        ("token_count", "token_tags must align with hidden_states"),
+        ("token_value", "token_tags must contain only 0 and 1"),
+    ],
+)
+@pytest.mark.parametrize("via_connector", [False, True])
+def test_encoder2diffusion_rejects_h3_v1_contract_mismatch(case: str, message: str, via_connector: bool) -> None:
+    payload = _encoder_output()
+    payload["hidden_states"]["output"] = torch.zeros(4, 5120, dtype=torch.bfloat16)
+    payload["meta"]["token_role_ids"] = torch.tensor([[1], [1], [0], [0]], dtype=torch.int64)
+
+    if case == "hidden_width":
+        payload["hidden_states"]["output"] = torch.zeros(4, 5119, dtype=torch.bfloat16)
+    elif case == "hidden_dtype":
+        payload["hidden_states"]["output"] = torch.zeros(4, 5120, dtype=torch.float32)
+    elif case == "hidden_noncontiguous":
+        payload["hidden_states"]["output"] = torch.empty(5120, 4, dtype=torch.bfloat16).t()
+    elif case == "hidden_non_strided":
+        payload["hidden_states"]["output"] = torch.empty(
+            (4, 5120),
+            dtype=torch.bfloat16,
+            layout=torch.sparse_coo,
+        )
+    elif case == "token_dtype":
+        payload["meta"]["token_role_ids"] = torch.tensor([[1], [1], [0], [0]], dtype=torch.int32)
+    elif case == "token_layout":
+        payload["meta"]["token_role_ids"] = torch.tensor(
+            [[1], [0], [1], [0], [1], [0], [1], [0]],
+            dtype=torch.int64,
+        )[::2]
+    elif case == "token_count":
+        payload["meta"]["token_role_ids"] = torch.tensor([[1], [0], [0]], dtype=torch.int64)
+    elif case == "token_value":
+        payload["meta"]["token_role_ids"] = torch.tensor([[1], [1], [2], [0]], dtype=torch.int64)
+    else:  # pragma: no cover - parameterization is exhaustive
+        raise AssertionError(f"unknown test case: {case}")
+
+    with pytest.raises(RuntimeError, match=message) as exc_info:
+        if via_connector:
+            encoder2diffusion_full_payload(pooling_output=payload)
+        else:
+            encoder2diffusion([_source_output(payload)], {"prompt": "hello"})
+
+    assert MINIMAX_H3_TEXT_CONDITIONING_SCHEMA in str(exc_info.value)
+
+
 def test_ref2va_one_image_tokens_and_tags_match_fused_presentation():
     tokenizer = _SegmentTokenizer()
     labels = [("image", 1), ("audio", 1)]
@@ -393,3 +456,236 @@ def test_diffusion_resolver_normalizes_partial_partition_directory(tmp_path):
     (ref2va / "text_encoder").mkdir(parents=True)
 
     assert resolve_minimax_h3_diffusion_model_path(str(ref2va), None, "ref2va") == str(ref2va)
+
+
+def _source_output(payload) -> SimpleNamespace:
+    return SimpleNamespace(
+        finished=True,
+        request_id="request-1",
+        outputs=[SimpleNamespace(multimodal_output=payload)],
+    )
+
+
+def _assert_same_tensor_view(actual: torch.Tensor, expected: torch.Tensor) -> None:
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0, check_stride=True)
+    assert actual.data_ptr() == expected.data_ptr()
+    assert actual.untyped_storage().data_ptr() == expected.untyped_storage().data_ptr()
+
+
+def _full_encoder_output() -> dict:
+    return MiniMaxH3EncoderConditioning(
+        hidden_states=torch.randn(4, 5120, dtype=torch.bfloat16),
+        token_tags=torch.tensor([1, 1, 0, 0], dtype=torch.int64),
+        task="ref2va",
+        height=256,
+        width=448,
+        num_frames=17,
+        latent_t=5,
+        audio_t=10,
+        visual_condition=torch.arange(12 * 96, dtype=torch.float32).reshape(12, 96),
+        visual_condition_shapes=((1, 4, 4), (2, 4, 4)),
+        audio_condition=torch.arange(10 * 32, dtype=torch.float32).reshape(10, 32),
+        audio_condition_lengths=(3, 2),
+        ref_blocks=(
+            {"kind": "image", "latent_t": 1, "latent_h": 4, "latent_w": 4},
+            {"kind": "video_audio", "ref_audio_t": 3, "latent_t": 2, "latent_h": 4, "latent_w": 4},
+            {"kind": "audio", "ref_audio_t": 2},
+        ),
+        keyframe_frame_indices=(0, 16),
+    ).to_omni_payload()
+
+
+def _runner_payload(payload: dict) -> dict:
+    # Encoder.make_omni_output emits dotted private layout metadata. The runner
+    # selects each request's media tensors and flattens its nested categories.
+    payload = dict(payload)
+    payload[f"kv_metadata.{MINIMAX_H3_ENCODER_LAYOUT_KEY}"] = payload.pop("kv_metadata")[MINIMAX_H3_ENCODER_LAYOUT_KEY]
+    payload["meta"] = {"token_role_ids": payload["meta"]["token_role_ids"].reshape(-1, 1)}
+    return flatten_payload(payload)
+
+
+@pytest.mark.parametrize("representation", ["nested", "struct", "flat", "runner"])
+@pytest.mark.parametrize("with_media", [False, True])
+@pytest.mark.parametrize("via_connector", [False, True])
+def test_encoder_handoff_preserves_all_components(representation, with_media, via_connector):
+    payload = _full_encoder_output() if with_media else _encoder_output()
+    inputs = {
+        "nested": payload,
+        "struct": to_struct(payload),
+        "flat": flatten_payload(payload),
+        "runner": _runner_payload(payload),
+    }
+    if via_connector:
+        result = encoder2diffusion_full_payload(pooling_output=inputs[representation], request_id="request-1")
+        assert set(result) == {"encoder_output"}
+    else:
+        result = encoder2diffusion([_source_output(inputs[representation])], {"prompt": "hello"})[
+            "additional_information"
+        ]
+    output = result["encoder_output"]
+    for group in ("hidden_states", "embed", "meta", "kv_metadata"):
+        for key, tensor in payload[group].items():
+            torch.testing.assert_close(output[group][key], tensor, rtol=0, atol=0)
+    # Nonempty tensors must not be stripped, copied to CPU, or coerced to a
+    # different dtype by the adapter. Empty FP32 optional slots remain on wire.
+    assert output["hidden_states"]["output"] is payload["hidden_states"]["output"]
+    if with_media:
+        for key in ("embedding", "speech_feat"):
+            _assert_same_tensor_view(output["embed"][key], payload["embed"][key])
+    else:
+        for tensor in output["embed"].values():
+            assert tensor.shape == (0,)
+            assert tensor.dtype == torch.float32
+    for _ in range(2):
+        result = deserialize_additional_information(serialize_additional_information(result))
+    restored = MiniMaxH3EncoderConditioning.from_omni_payload(result["encoder_output"])
+    original = MiniMaxH3EncoderConditioning.from_omni_payload(payload)
+    for name in (
+        "task",
+        "height",
+        "width",
+        "num_frames",
+        "latent_t",
+        "audio_t",
+        "visual_condition_shapes",
+        "audio_condition_lengths",
+        "ref_blocks",
+        "keyframe_frame_indices",
+    ):
+        assert getattr(restored, name) == getattr(original, name)
+    for group, fields in output.items():
+        for key, tensor in fields.items():
+            torch.testing.assert_close(result["encoder_output"][group][key], tensor, rtol=0, atol=0)
+
+
+def test_full_payload_hook_skips_only_absent_output():
+    assert encoder2diffusion_full_payload(pooling_output=None) is None
+
+
+@pytest.mark.parametrize("via_connector", [False, True])
+@pytest.mark.parametrize(
+    ("case", "message"),
+    [
+        ("not_mapping", "no conditioning payload"),
+        ("text_only", "requires text, visual and audio tensors"),
+        ("missing_hidden", "requires text, visual and audio tensors"),
+        ("missing_tags", "requires text, visual and audio tensors"),
+        ("missing_layout", "requires private layout metadata"),
+        ("truncated_layout", "header is truncated"),
+        ("layout_dtype", "one-dimensional integer tensor"),
+        ("schema", "unsupported MiniMax H3 encoder wire schema"),
+        ("visual_shape", "visual condition must be FP32"),
+        ("audio_dtype", "audio condition must be FP32"),
+        ("empty_visual", "empty visual condition slot"),
+        ("empty_audio", "empty audio condition slot"),
+    ],
+)
+def test_encoder_handoff_rejects_invalid_full_payload(case, message, via_connector):
+    payload = _full_encoder_output()
+    layout = payload["kv_metadata"][MINIMAX_H3_ENCODER_LAYOUT_KEY]
+    if case == "not_mapping":
+        payload = []
+    elif case == "text_only":
+        payload.pop("embed")
+    elif case == "missing_hidden":
+        payload.pop("hidden_states")
+    elif case == "missing_tags":
+        payload["meta"] = {}
+    elif case == "missing_layout":
+        payload.pop("kv_metadata")
+    elif case == "truncated_layout":
+        payload["kv_metadata"][MINIMAX_H3_ENCODER_LAYOUT_KEY] = layout[:5]
+    elif case == "layout_dtype":
+        payload["kv_metadata"][MINIMAX_H3_ENCODER_LAYOUT_KEY] = layout.float()
+    elif case == "schema":
+        layout[1] = 99
+    elif case == "visual_shape":
+        payload["embed"]["embedding"] = torch.zeros(11, 96)
+    elif case == "audio_dtype":
+        payload["embed"]["speech_feat"] = payload["embed"]["speech_feat"].to(torch.bfloat16)
+    elif case == "empty_visual":
+        payload = _encoder_output()
+        payload["embed"]["embedding"] = torch.empty(0, 96)
+    elif case == "empty_audio":
+        payload = _encoder_output()
+        payload["embed"]["speech_feat"] = torch.empty(0, dtype=torch.bfloat16)
+    with pytest.raises(RuntimeError, match=message):
+        if via_connector:
+            encoder2diffusion_full_payload(pooling_output=payload)
+        else:
+            encoder2diffusion([_source_output(payload)], {"prompt": "hello"})
+
+
+@pytest.mark.parametrize("inline", [False, True])
+def test_encoder_handoff_cleans_media_without_mutating_prompt_or_stripping_outputs(inline):
+    incoming = _full_encoder_output()
+    retained_hidden = torch.ones(1)
+    media = torch.zeros(2)
+    prompt = {
+        "prompt": "hello",
+        "negative_prompt": "blur",
+        "multi_modal_data": {"image": object(), "audio": object(), "video": object()},
+        "model_intermediate_buffer": {"media": media},
+        "additional_information": {
+            "global_request_id": ["request-1"],
+            "private": "preserved",
+            "hidden_states": {"layers": {0: media}, "output": retained_hidden},
+            "meta": {MINIMAX_H3_ENCODER_REQUEST_KEY: {"task": "ref2va"}, "private": 42},
+            "encoder_output": incoming,
+        },
+    }
+    source = _source_output(_runner_payload(incoming) if inline else None)
+    if not inline:
+        del source.outputs[0].multimodal_output
+    result = encoder2diffusion([source], [prompt])
+    info = result["additional_information"]
+    assert result["multi_modal_data"] is None
+    assert "model_intermediate_buffer" not in result
+    assert result["negative_prompt"] == "blur"
+    assert info["global_request_id"] == ["request-1"]
+    assert info["private"] == "preserved"
+    assert info["meta"] == {"private": 42}
+    assert set(info["hidden_states"]) == {"output"}
+    assert info["hidden_states"]["output"] is retained_hidden
+    for key in ("embedding", "speech_feat"):
+        _assert_same_tensor_view(info["encoder_output"]["embed"][key], incoming["embed"][key])
+    if not inline:
+        assert info["encoder_output"] is incoming
+    assert prompt["multi_modal_data"] is not None
+    assert prompt["model_intermediate_buffer"]["media"] is media
+    assert prompt["additional_information"]["hidden_states"]["layers"][0] is media
+    assert MINIMAX_H3_ENCODER_REQUEST_KEY in prompt["additional_information"]["meta"]
+
+
+def test_connector_handoff_without_inline_output_cleans_encoder_only_information():
+    prompt = {
+        "prompt": "hello",
+        "multi_modal_data": {"image": object()},
+        "model_intermediate_buffer": {"private": object()},
+        "additional_information": {
+            "hidden_states": {"layers": {0: torch.zeros(1)}},
+            "meta": {MINIMAX_H3_ENCODER_REQUEST_KEY: {"task": "t2va"}},
+        },
+    }
+    result = encoder2diffusion([_source_output(None)], prompt)
+    assert result == {"prompt": "hello", "multi_modal_data": None, "additional_information": {}}
+
+
+@pytest.mark.parametrize("payload", [None, _encoder_output()])
+@pytest.mark.parametrize("request_id", ["other", ["other"], ("other",)])
+def test_stage_wire_rejects_request_id_mismatch(payload, request_id):
+    prompt = {"prompt": "hello", "additional_information": {"global_request_id": request_id}}
+    with pytest.raises(RuntimeError, match="request ID does not match"):
+        encoder2diffusion([_source_output(payload)], prompt)
+
+
+@pytest.mark.parametrize("outputs", [None, [], (), [SimpleNamespace(), SimpleNamespace()]])
+def test_stage_wire_rejects_invalid_completion_count(outputs):
+    with pytest.raises(RuntimeError, match="exactly one completion"):
+        encoder2diffusion([SimpleNamespace(outputs=outputs)], {"prompt": "hello"})
+
+
+def test_stage_wire_empty_sources_and_invalid_prompt():
+    assert encoder2diffusion([], {"prompt": "hello"}) is None
+    with pytest.raises(TypeError, match="invalid MiniMax H3 prompt type"):
+        encoder2diffusion([_source_output(None)], object())