[Model] MiniMax-H3: overlap Ref2VA CPU noise preparation

recipes/MiniMaxAI/minimax-h3-ref2va-8-rank-resident.yaml

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Single-stage Ref2VA: text TP8, DiT USP8, resident video/audio VAEs.
+# Parameters from the 8x Ascend 910B1 (64 GiB/device) experiment.
+# Use the existing serving container's Ascend runtime environment.
+# vllm serve "$MODEL" --omni --config recipes/MiniMaxAI/minimax-h3-ref2va-8-rank-resident.yaml
+
+host: "127.0.0.1"
+port: 19100
+trust-remote-code: true
+task-type: ref2va
+num-gpus: 8
+usp: 8
+ring: 1
+text-encoder-tp-size: 8
+vae-parallel-mode: tile
+vae-use-tiling: true
+vae-patch-parallel-size: 8
+enable-distributed-layerwise-offload: true
+diffusion-attention-backend: FLASH_ATTN
+
+# The CLI accepts this argument as JSON. Selecting only these two components
+# keeps video/audio VAEs resident and preserves rank-local text weight transfer.
+diffusion-offload-config: >-
+  {"mode":"layer","components":["dit","text_encoder"],"layer_options":{"dit":{"weight_transfer":"allgather"},"text_encoder":{"weight_transfer":"rank-local"}}}

tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py

@@ -213,6 +213,7 @@ def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monke
     )
     monkeypatch.setattr(processing, "load_video_audio", Mock(return_value=(torch.full((64000,), 3.0), 32000)))
     pipeline.video_vae.encode_video.return_value = (torch.full((1, 96), 2.0), (1, 2, 2))
+    pipeline.video_vae.predict_video_latent_shape.return_value = (1, 2, 2)
 
     def encode_audio(waveform, sample_rate):
         length = round(waveform.shape[-1] / sample_rate * 40)
@@ -221,8 +222,17 @@ def encode_audio(waveform, sample_rate):
     pipeline.audio_vae.encode_waveform.side_effect = encode_audio
     request = _request("ref2va")
     request.prompts[0]["multi_modal_data"].update(video="original.mp4", audio=(torch.full((192000,), 4.0), 32000))
+
+    def capture_denoise_kwargs():
+        kwargs = pipeline.diffuse.call_args.kwargs.copy()
+        context = kwargs.pop("request_context")
+        # Mirror diffuse's consumption of request-owned noise at the mocked boundary.
+        for key in ("precomputed_initial_noise", "precomputed_visual_condition_noise"):
+            kwargs[key] = context.pop(key, None)
+        return kwargs
+
     pipeline.forward(request)
-    local_kwargs = pipeline.diffuse.call_args.kwargs
+    local_kwargs = capture_denoise_kwargs()
 
     transformed = prepare_encoder_prompt(request.prompts[0], [request.sampling_params])
     stage_zero = object.__new__(MiniMaxH3Encoder)
@@ -261,12 +271,20 @@ def encode_audio(waveform, sample_rate):
     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
+    external_kwargs = capture_denoise_kwargs()
     for key, value in local_kwargs.items():
+        if key in {"precomputed_initial_noise", "precomputed_visual_condition_noise"}:
+            assert value is not None
+            assert external_kwargs[key] is None
+            continue
         if isinstance(value, torch.Tensor):
             torch.testing.assert_close(value, external_kwargs[key])
         else:
             assert value == external_kwargs[key]
+    local_inputs = pipeline._build_denoise_inputs(**local_kwargs)
+    external_inputs = pipeline._build_denoise_inputs(**external_kwargs)
+    for key in ("video_rows", "audio_rows", "cond_anchor", "audio_anchor"):
+        assert torch.equal(local_inputs[key], external_inputs[key])
     pipeline.encode_prompt.assert_not_called()
     pipeline.video_vae.encode_image.assert_not_called()
     pipeline.video_vae.encode_video.assert_not_called()

tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py

@@ -192,6 +192,94 @@ def test_condition_noise_accepts_a_reference_video_longer_than_the_target():
     assert result.shape == rows.shape
 
 
+@pytest.mark.parametrize(
+    ("condition_shapes", "noncontiguous"),
+    [([(1, 4, 6)], False), ([(1, 4, 6), (3, 8, 4), (7, 4, 4)], True)],
+)
+@pytest.mark.parametrize("noise_aug", [0.0, 0.999])
+def test_condition_noise_preserves_fp32_recipe_and_owns_output(condition_shapes, noncontiguous, noise_aug):
+    from vllm_omni.diffusion.models.minimax_h3.condition_noise import (
+        minimax_h3_imgvid_cond_noise_aug_rows,
+        minimax_h3_imgvid_cond_noise_rows,
+    )
+    from vllm_omni.diffusion.models.minimax_h3.packed_tokens import (
+        minimax_h3_patchify_video_latent,
+    )
+
+    row_count = sum(t * (h // 2) * (w // 2) for t, h, w in condition_shapes)
+    clean_generator = torch.Generator(device="cpu").manual_seed(31)
+    shape = (96, row_count) if noncontiguous else (row_count, 96)
+    clean_rows = torch.randn(shape, generator=clean_generator, dtype=torch.float32)
+    if noncontiguous:
+        clean_rows = clean_rows.t()
+        assert not clean_rows.is_contiguous()
+    clean_before = clean_rows.clone()
+    rng_before = torch.get_rng_state().clone()
+    target_latent_t = 2
+    seed = 2101
+    timestep = torch.tensor(noise_aug, dtype=torch.float32)
+    expected_parts = []
+    expected_noise_parts = []
+    offset = 0
+    # Match the original per-reference FP32 recipe exactly.
+    for t, h, w in condition_shapes:
+        full_t = max(target_latent_t + len(condition_shapes), t)
+        generator = torch.Generator(device="cpu").manual_seed(seed)
+        noise = torch.randn(1, 24, full_t, h, w, generator=generator, dtype=torch.float32)[:, :, :t]
+        noise_rows = minimax_h3_patchify_video_latent(noise, patch_size=(1, 2, 2))
+        expected_noise_parts.append(noise_rows)
+        count = noise_rows.shape[0]
+        expected_parts.append(timestep * clean_rows[offset : offset + count] + (1.0 - timestep) * noise_rows)
+        offset += count
+    expected = torch.cat(expected_parts)
+    expected_noise = torch.cat(expected_noise_parts)
+    precomputed_noise = minimax_h3_imgvid_cond_noise_rows(
+        condition_shapes=condition_shapes,
+        target_latent_t=target_latent_t,
+        imgvid_cond_num_frames=len(condition_shapes),
+        seed=seed,
+    )
+
+    result = minimax_h3_imgvid_cond_noise_aug_rows(
+        clean_rows,
+        condition_shapes=condition_shapes,
+        target_latent_t=target_latent_t,
+        imgvid_cond_num_frames=len(condition_shapes),
+        seed=seed,
+        noise_aug=noise_aug,
+    )
+    precomputed_result = minimax_h3_imgvid_cond_noise_aug_rows(
+        clean_rows,
+        condition_shapes=condition_shapes,
+        target_latent_t=target_latent_t,
+        imgvid_cond_num_frames=len(condition_shapes),
+        seed=seed,
+        noise_aug=noise_aug,
+        precomputed_noise_rows=precomputed_noise,
+    )
+    different_seed_noise = minimax_h3_imgvid_cond_noise_rows(
+        condition_shapes=condition_shapes,
+        target_latent_t=target_latent_t,
+        imgvid_cond_num_frames=len(condition_shapes),
+        seed=seed + 1,
+    )
+
+    assert torch.equal(result.view(torch.int32), expected.view(torch.int32))
+    assert torch.equal(precomputed_noise.view(torch.int32), expected_noise.view(torch.int32))
+    assert torch.equal(precomputed_result.view(torch.int32), result.view(torch.int32))
+    assert not torch.equal(different_seed_noise, precomputed_noise)
+    assert precomputed_noise.device.type == "cpu"
+    assert precomputed_noise.dtype == torch.float32
+    assert torch.equal(clean_rows, clean_before)
+    assert torch.equal(torch.get_rng_state(), rng_before)
+    assert result.dtype == torch.float32
+    assert result.device.type == "cpu"
+    assert result.is_contiguous()
+    assert result.untyped_storage().data_ptr() != clean_rows.untyped_storage().data_ptr()
+    result.zero_()
+    assert torch.equal(clean_rows, clean_before)
+
+
 def test_explicit_seq_len_pins_one_shape_across_prompt_lengths():
     """Prompts of different token counts must land on the same packed length.
 

tests/diffusion/models/minimax_h3/test_ref2va_initial_noise_preparation.py

@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import weakref
+from concurrent.futures import ThreadPoolExecutor
+from threading import Event
+
+import pytest
+import torch
+from PIL import Image
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.models.minimax_h3.condition_noise import minimax_h3_imgvid_cond_noise_rows
+from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+    _MINIMAX_H3_DENOISE_INPUT_KEYS,
+    MiniMaxH3Pipeline,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture
+def pipeline():
+    # Exercise real CPU tensor methods without loading model weights.
+    instance = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(instance)
+    instance.device = torch.device("cpu")
+    return instance
+
+
+@pytest.mark.parametrize("seed", [0, 42, 2101])
+@pytest.mark.parametrize("shape", [(2, 4, 6, 3), (7, 8, 10, 12)])
+def test_background_initial_noise_matches_direct_sampling(pipeline, seed, shape):
+    latent_t, latent_h, latent_w, audio_t = shape
+    params = dict(seed=seed, latent_t=latent_t, latent_h=latent_h, latent_w=latent_w, audio_t=audio_t)
+    rng_state = torch.get_rng_state().clone()
+    expected = pipeline._initial_noise(**params)
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        actual = executor.submit(pipeline._initial_noise, **params).result()
+
+    for direct, prepared in zip(expected, actual, strict=True):
+        assert torch.equal(prepared.view(torch.int32), direct.view(torch.int32))
+        assert prepared.device.type == "cpu"
+        assert prepared.dtype == torch.float32
+        assert prepared.is_contiguous()
+    assert torch.equal(torch.get_rng_state(), rng_state)
+
+
+def _ref2va_kwargs(seed):
+    return dict(
+        task="ref2va",
+        text_embeddings=torch.zeros(2, 4),
+        text_tags=torch.ones(2, dtype=torch.long),
+        seed=seed,
+        latent_t=2,
+        latent_h=4,
+        latent_w=6,
+        audio_t=3,
+        num_frames=22,
+        num_steps=3,
+        video_shift=12.0,
+        audio_shift=3.0,
+        base_schedule=None,
+        visual_condition=torch.arange(6 * 96, dtype=torch.float32).reshape(6, 96) / 100,
+        visual_condition_shape=(1, 4, 6),
+        audio_condition=torch.arange(4 * 32, dtype=torch.float32).reshape(4, 32) / 100,
+        ref_audio_t=2,
+        ref_blocks=[{"kind": "image", "latent_h": 4, "latent_w": 6}, {"kind": "audio", "ref_audio_t": 2}],
+    )
+
+
+@pytest.mark.parametrize("seed", [42, 2101])
+def test_ref2va_precomputed_noise_preserves_denoise_inputs(pipeline, seed):
+    kwargs = _ref2va_kwargs(seed)
+    expected = pipeline._build_denoise_inputs(**kwargs)
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        noise = executor.submit(
+            pipeline._initial_noise, seed=seed, latent_t=2, latent_h=4, latent_w=6, audio_t=3
+        ).result()
+        visual_noise = executor.submit(
+            minimax_h3_imgvid_cond_noise_rows,
+            condition_shapes=[kwargs["visual_condition_shape"]],
+            target_latent_t=kwargs["latent_t"],
+            imgvid_cond_num_frames=1,
+            seed=seed,
+        ).result()
+    actual = pipeline._build_denoise_inputs(
+        **kwargs, precomputed_initial_noise=noise, precomputed_visual_condition_noise=visual_noise
+    )
+
+    for key in ("video_rows", "audio_rows", "cond_anchor", "audio_anchor"):
+        assert torch.equal(actual[key].view(torch.int32), expected[key].view(torch.int32))
+    assert actual["sigmas_video"] == expected["sigmas_video"]
+    assert actual["sigmas_audio"] == expected["sigmas_audio"]
+    assert torch.equal(actual["branch"].token_tags_dev, expected["branch"].token_tags_dev)
+
+
+def test_shared_denoise_kwargs_exclude_first_seed_noise(pipeline):
+    context = {key: None for key in _MINIMAX_H3_DENOISE_INPUT_KEYS}
+    context.update(_ref2va_kwargs(42))
+    context["precomputed_initial_noise"] = pipeline._initial_noise(
+        seed=42, latent_t=2, latent_h=4, latent_w=6, audio_t=3
+    )
+    context["precomputed_visual_condition_noise"] = torch.zeros(4, 96)
+
+    kwargs = pipeline._denoise_kwargs(context)
+    assert "precomputed_initial_noise" not in kwargs
+    assert "precomputed_visual_condition_noise" not in kwargs
+    assert kwargs["seed"] == 42
+    assert kwargs["latent_t"] == 2
+    assert kwargs["latent_h"] == 4
+    assert kwargs["latent_w"] == 6
+    assert kwargs["audio_t"] == 3
+
+
+def test_consumed_host_noise_is_released_before_denoising(pipeline, monkeypatch):
+    kwargs = _ref2va_kwargs(42)
+    noise = pipeline._initial_noise(seed=42, latent_t=2, latent_h=4, latent_w=6, audio_t=3)
+    visual_noise = minimax_h3_imgvid_cond_noise_rows(
+        condition_shapes=[kwargs["visual_condition_shape"]], target_latent_t=2, imgvid_cond_num_frames=1, seed=42
+    )
+    references = [weakref.ref(tensor) for tensor in (*noise, visual_noise)]
+    context = {"precomputed_initial_noise": noise, "precomputed_visual_condition_noise": visual_noise}
+    del noise, visual_noise
+    # Real CPU-to-meta copies have distinct storage without loading model weights.
+    pipeline.device = torch.device("meta")
+
+    def stop_before_denoising(task):
+        assert all(reference() is None for reference in references)
+        raise RuntimeError("checked noise lifetime")
+
+    monkeypatch.setattr(pipeline, "_transformer_for_task", stop_before_denoising)
+    with pytest.raises(RuntimeError, match="checked noise lifetime"):
+        pipeline.diffuse(
+            **kwargs,
+            request_context=context,
+        )
+
+
+@pytest.mark.parametrize("failure_stage", ["text", "media"])
+def test_encoder_failure_cancels_pending_noise_and_joins_worker(pipeline, monkeypatch, failure_stage):
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pool = ThreadPoolExecutor(max_workers=1)
+    release = Event()
+    started = Event()
+    futures = []
+    submit = pool.submit
+    shutdown = pool.shutdown
+
+    def hold_worker():
+        started.set()
+        assert release.wait(10)
+
+    blocker = submit(hold_worker)
+    assert started.wait(10)
+
+    def record_submit(*args, **kwargs):
+        future = submit(*args, **kwargs)
+        futures.append(future)
+        return future
+
+    def release_and_shutdown(*, wait=True, cancel_futures=False):
+        shutdown(wait=False, cancel_futures=cancel_futures)
+        release.set()
+        shutdown(wait=wait)
+
+    def encode_prompt(prepared):
+        if failure_stage == "text":
+            raise RuntimeError("text encoder failed")
+        return torch.zeros(2, 4), torch.ones(2, dtype=torch.long)
+
+    def fail_media(media):
+        raise RuntimeError("media encoder failed")
+
+    pipeline._fasth3 = None
+    pipeline.supported_tasks = ("ref2va",)
+    pipeline.od_config = OmniDiffusionConfig()
+    monkeypatch.setattr(pool, "submit", record_submit)
+    monkeypatch.setattr(pool, "shutdown", release_and_shutdown)
+    monkeypatch.setattr(module, "ThreadPoolExecutor", lambda **kwargs: pool)
+    monkeypatch.setattr(pipeline, "encode_prompt", encode_prompt)
+    monkeypatch.setattr(pipeline, "_encode_local_media", fail_media)
+    sampling = OmniDiffusionSamplingParams(height=64, width=64, seed=42, extra_args={"task": "ref2va", "duration": 4.4})
+    prompt = {"prompt": "reference", "multi_modal_data": {"image": Image.new("RGB", (256, 256))}}
+    try:
+        with pytest.raises(RuntimeError, match=f"{failure_stage} encoder failed"):
+            pipeline._prepare_local_conditioning(prompt, sampling, request_context={})
+        assert len(futures) == 2
+        assert all(future.cancelled() for future in futures)
+        assert blocker.done()
+    finally:
+        release.set()
+        shutdown(wait=True, cancel_futures=True)

tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py

@@ -100,6 +100,7 @@
 SEED = 0
 SSIM_THRESHOLD = 0.97
 PSNR_THRESHOLD = 34.0
+# The worker RPC spans the full generation and shares the HTTP request budget.
 REQUEST_TIMEOUT_SECONDS = 60 * 60
 
 
@@ -256,6 +257,7 @@ def test_minimax_h3_i2va_matches_official_reference(
         "FLASHINFER_DISABLE_VERSION_CHECK": "1",
         "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
         "VLLM_OMNI_VIDEO_SYNC_TIMEOUT": str(REQUEST_TIMEOUT_SECONDS),
+        "VLLM_OMNI_DLO_DP_WAVE_TIMEOUT": str(REQUEST_TIMEOUT_SECONDS),
         "VLLM_OMNI_STORAGE_PATH": str(output_dir / "storage"),
     }
     request_data = {
@@ -339,6 +341,7 @@ def test_minimax_h3_ref2va_matches_official_reference(
         "FLASHINFER_DISABLE_VERSION_CHECK": "1",
         "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
         "VLLM_OMNI_VIDEO_SYNC_TIMEOUT": str(REQUEST_TIMEOUT_SECONDS),
+        "VLLM_OMNI_DLO_DP_WAVE_TIMEOUT": str(REQUEST_TIMEOUT_SECONDS),
         "VLLM_OMNI_STORAGE_PATH": str(output_dir / "storage"),
     }
     request_data = {