[Bugfix][MiniMax-H3] Release decoded outputs before the next seed

docs/design/feature/async_diffusion_output.md

@@ -22,7 +22,7 @@ The async diffusion output feature moves the D2H (Device→Host) copy and SHM pa
 ### HunyuanImage-3.0 (TP4)
 
 | Resolution | Async Off (QPS) | Async On (QPS) | Change |
-|---|---|---|---|
+| --- | --- | --- | --- |
 | 1024×1024 | 0.4773 | 0.4802 | +0.60% |
 | 768×768 | 0.8370 | 0.8533 | +1.95% |
 
@@ -32,7 +32,7 @@ The async diffusion output feature moves the D2H (Device→Host) copy and SHM pa
 
 ### Execution Timeline
 
-```
+```text
 Before (synchronous D2H):
 [forward req1] [D2H+SHM pack req1] [forward req2] [D2H+SHM pack req2]
                                  ^^^^^^^^^^^^^^^^ GPU bubble
@@ -45,7 +45,7 @@ After (async D2H):
 
 ### Data Flow
 
-```
+```text
                     ┌─────────────────────────────────────────────────────┐
                     │                 Worker Process                      │
                     │                                                    │
@@ -143,13 +143,24 @@ All models running in request-mode (`step_execution=False`, the default) automat
 **Verified models**:
 
 | Model | Type | `supports_request_batch` |
-|---|---|---|
+| --- | --- | --- |
 | **HunyuanImage-3.0** | Image | `False` |
 | **Qwen-Image** | Image | `True` |
 | **LTX-2.3** | Video / Audio | `False` |
 
 Other models with `step_execution=False` are also supported but not yet verified.
 
+**Narrow exception -- MiniMax-H3 with active model-level CPU offload**: the
+pipeline moves each decoded output to the host on the reply rank as soon as that
+output reaches its final form -- audio inside `decode()`, and video in the
+callers of `decode()` right after `_prepare_minimax_h3_video_output` has
+quantized the frames to `uint8`. Other ranks replace it with a metadata-only
+tensor because their output is not consumed. This prevents a prior seed's
+decoded output from overlapping the DiT reload for the next seed in a
+multi-output request. The async envelope and SHM packing still run, but the main
+D2H copy has already happened before `COMPUTE_DONE`, so this path deliberately
+gives up D2H/next-forward overlap.
+
 ### Not Yet Supported (step-mode, `step_execution=True`)
 
 When `step_execution=True` (or `streaming_output=True`, which auto-enables step-mode), models use `execute_stepwise` instead of `execute_model`, which is not in the async Path 1 whitelist. Async output is not applicable.
@@ -171,7 +182,7 @@ Adapting step-mode requires additional design because:
 No configuration needed. Async output is automatically enabled when `step_execution=False` (default).
 
 | `step_execution` | Mode | Async Output | Pump Thread | Worker BG Thread |
-|---|---|---|---|---|
+| --- | --- | --- | --- | --- |
 | `False` (default) | request-mode | ✅ enabled | ✅ started | ✅ started |
 | `True` | step-mode | ❌ disabled | ❌ not started | ❌ not started |
 

tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-from contextlib import contextmanager
+import weakref
+from contextlib import contextmanager, nullcontext
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -73,6 +74,243 @@ def test_requested_model_offload_without_backend_uses_resident_components():
     pipeline.text_encoder.encode_ids.assert_called_once()
 
 
+def test_decoded_output_is_unchanged_without_active_model_offload():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline._model_cpu_offload_modules = []
+    output = Mock()
+
+    assert pipeline._offload_model_cpu_stage_output(output) is output
+    output.cpu.assert_not_called()
+
+
+def test_reply_rank_copies_decoded_output_to_host(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline._model_cpu_offload_modules = [Mock()]
+    monkeypatch.setattr(pipeline, "_is_output_owner_rank", lambda: True)
+    output = Mock()
+    host_output = Mock()
+    output.cpu.return_value = host_output
+
+    assert pipeline._offload_model_cpu_stage_output(output) is host_output
+    output.cpu.assert_called_once_with()
+
+
+def test_non_reply_rank_drops_decoded_output_storage(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline._model_cpu_offload_modules = [Mock()]
+    monkeypatch.setattr(pipeline, "_is_output_owner_rank", lambda: False)
+    output = torch.ones((2, 3), dtype=torch.float16)
+
+    placeholder = pipeline._offload_model_cpu_stage_output(output)
+
+    assert placeholder.device.type == "meta"
+    assert placeholder.shape == output.shape
+    assert placeholder.dtype == output.dtype
+
+
+def test_decode_releases_only_audio_before_return(monkeypatch):
+    """Video is released by the callers of ``decode`` -- after quantization."""
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    video = torch.ones((1, 1, 4, 4))
+    audio = torch.ones((1, 8))
+    pipeline = SimpleNamespace(
+        device=torch.device("cpu"),
+        video_vae=Mock(decode_latent=Mock(return_value=video)),
+        audio_vae=Mock(decode_latent=Mock(return_value=audio)),
+        _component_on_device=lambda component: nullcontext(),
+        _offload_model_cpu_stage_output=Mock(side_effect=lambda value: value),
+    )
+    monkeypatch.setattr(
+        module.current_omni_platform,
+        "create_autocast_context",
+        lambda **kwargs: nullcontext(),
+    )
+
+    actual_video, actual_audio = MiniMaxH3Pipeline.decode(
+        pipeline,
+        torch.zeros(1),
+        torch.zeros(1),
+        height=2,
+        width=3,
+    )
+
+    assert actual_video.shape[-2:] == (2, 3)
+    assert actual_audio is audio
+    calls = pipeline._offload_model_cpu_stage_output.call_args_list
+    assert len(calls) == 1
+    assert calls[0].args[0] is audio
+
+
+def _recording_release(events, releases):
+    """Stand-in for ``_offload_model_cpu_stage_output`` that records what it got."""
+
+    def release(value):
+        events.append("release")
+        releases.append(value)
+        return value
+
+    return release
+
+
+def test_forward_releases_video_after_quantization_and_before_next_seed(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    # MiniMaxH3VideoVAE.decode_latent ends in frames.float(), so this is what
+    # the release hook would see if it ran before the quantizer.
+    decoded_dtype = torch.float32
+    events: list[str] = []
+    releases: list[torch.Tensor] = []
+
+    def diffuse(**kwargs):
+        events.append(f"diffuse:{kwargs['seed']}")
+        return torch.zeros(1), torch.zeros(1)
+
+    def decode(video_latent, audio_latent, *, height, width):
+        events.append("decode")
+        return (
+            torch.zeros((1, 3, 2, height, width), dtype=decoded_dtype),
+            torch.zeros((1, 8)),
+        )
+
+    context = {"num_outputs": 2, "seed": 41, "height": 4, "width": 4}
+    pipeline = SimpleNamespace(
+        od_config=SimpleNamespace(),
+        _extract_prompt=lambda raw: ("a prompt", {}),
+        _extract_text_conditioning=lambda raw: None,
+        _extract_prepared_reference_videos=lambda raw: None,
+        _prepare_request_inputs=lambda **kwargs: context,
+        _denoise_kwargs=lambda ctx: {},
+        diffuse=diffuse,
+        decode=decode,
+        _offload_model_cpu_stage_output=_recording_release(events, releases),
+    )
+    request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace())
+
+    output = MiniMaxH3Pipeline.forward(pipeline, request)
+
+    # The release must see the quantized frames, not the decoded floats.
+    assert len(releases) == 2
+    assert [tensor.dtype for tensor in releases] == [torch.uint8, torch.uint8]
+    assert decoded_dtype is not torch.uint8
+    # ... and it must still happen before the next seed reloads the DiT.
+    assert events == ["diffuse:41", "decode", "release", "diffuse:42", "decode", "release"]
+    assert output.output[0].dtype is torch.uint8
+
+
+def test_forward_drops_the_decoded_video_before_the_next_seed_diffuses(monkeypatch):
+    """Releasing the copy is not enough if a local still pins the original.
+
+    ``_offload_model_cpu_stage_output`` hands back a *new* tensor, so the
+    decoded frames stay alive for as long as anything still references them.
+    The whole point of this branch is that they are gone before the next seed
+    reloads the DiT, and ``diffuse`` is where that reload happens.
+    """
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    decoded_refs: list[weakref.ref] = []
+    alive_at_diffuse: list[int] = []
+
+    def diffuse(**kwargs):
+        alive_at_diffuse.append(sum(1 for ref in decoded_refs if ref() is not None))
+        return torch.zeros(1), torch.zeros(1)
+
+    def decode(video_latent, audio_latent, *, height, width):
+        video = torch.zeros((1, 3, 2, height, width), dtype=torch.float32)
+        decoded_refs.append(weakref.ref(video))
+        return video, torch.zeros((1, 8))
+
+    context = {"num_outputs": 2, "seed": 41, "height": 4, "width": 4}
+    pipeline = SimpleNamespace(
+        od_config=SimpleNamespace(),
+        _extract_prompt=lambda raw: ("a prompt", {}),
+        _extract_text_conditioning=lambda raw: None,
+        _extract_prepared_reference_videos=lambda raw: None,
+        _prepare_request_inputs=lambda **kwargs: context,
+        _denoise_kwargs=lambda ctx: {},
+        diffuse=diffuse,
+        decode=decode,
+        _offload_model_cpu_stage_output=lambda value: value.clone(),
+    )
+    request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace())
+
+    MiniMaxH3Pipeline.forward(pipeline, request)
+
+    assert len(decoded_refs) == 2
+    # Seed 1's diffuse sees nothing decoded yet; seed 2's must not see seed 1's
+    # decoded frames still resident.
+    assert alive_at_diffuse == [0, 0]
+
+
+def test_post_decode_releases_video_after_quantization(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    events: list[str] = []
+    releases: list[torch.Tensor] = []
+
+    def decode(video_latent, audio_latent, *, height, width):
+        events.append("decode")
+        return (
+            torch.zeros((1, 3, 2, height, width), dtype=torch.float16),
+            torch.zeros((1, 8)),
+        )
+
+    pipeline = SimpleNamespace(
+        od_config=SimpleNamespace(),
+        _unpack_denoised_rows=lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1)),
+        decode=decode,
+        _offload_model_cpu_stage_output=_recording_release(events, releases),
+    )
+    state = SimpleNamespace(
+        latents=torch.zeros(1),
+        extra={
+            module._STEP_BRANCH: object(),
+            module._STEP_AUDIO_ROWS: torch.zeros(1),
+            module._STEP_SHAPE: {
+                "latent_t": 1,
+                "latent_h": 1,
+                "latent_w": 1,
+                "audio_t": 1,
+                "height": 4,
+                "width": 4,
+            },
+        },
+    )
+
+    output = MiniMaxH3Pipeline.post_decode(pipeline, state)
+
+    assert events == ["decode", "release"]
+    assert len(releases) == 1
+    assert releases[0].dtype is torch.uint8
+    assert output.output[0] is releases[0]
+
+
+@pytest.mark.parametrize(
+    ("available", "initialized", "rank", "expected"),
+    [(False, False, 7, True), (True, False, 7, True), (True, True, 0, True), (True, True, 1, False)],
+)
+def test_output_owner_rank_matches_executor_contract(monkeypatch, available, initialized, rank, expected):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    monkeypatch.setattr(module.dist, "is_available", lambda: available)
+    monkeypatch.setattr(module.dist, "is_initialized", lambda: initialized)
+    monkeypatch.setattr(module.dist, "get_rank", lambda: rank)
+
+    assert MiniMaxH3Pipeline._is_output_owner_rank() is expected
+
+
 def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -1613,6 +1613,31 @@ def _component_on_device(self, component: nn.Module):
                             logger.exception("Failed to release retained allocator cache after offload failure")
                     raise
 
+    @staticmethod
+    def _is_output_owner_rank() -> bool:
+        """Whether this rank's output is returned by the diffusion executor."""
+        return not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0
+
+    def _offload_model_cpu_stage_output(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Release a decoded output's storage before a later seed reloads the DiT.
+
+        Model-level CPU offload makes the VAE hook evict the DiT before decode.
+        For a multi-output request, however, the next seed reloads the DiT while
+        the preceding decoded output would otherwise still occupy the device.
+        The reply rank performs the D2H copy early; ranks whose outputs are not
+        consumed retain only a metadata placeholder for the final concatenation.
+
+        Call this on the tensor that is actually returned to the engine. Audio is
+        released inside ``decode``; video is released by the callers of ``decode``
+        only after ``_prepare_minimax_h3_video_output`` has quantized it, so the
+        early D2H copy carries ``uint8`` frames rather than the decoded floats.
+        """
+        if not getattr(self, "_model_cpu_offload_modules", None):
+            return tensor
+        if self._is_output_owner_rank():
+            return tensor.cpu()
+        return torch.empty(tensor.shape, dtype=tensor.dtype, device="meta")
+
     def _encode_visual_conditions(
         self,
         images: list[Image.Image],
@@ -2118,6 +2143,7 @@ def decode(
         video = video[..., :height, :width].contiguous()
         with self._component_on_device(self.audio_vae):
             audio = self.audio_vae.decode_latent(audio_latent)
+        audio = self._offload_model_cpu_stage_output(audio)
         return video, audio
 
     @staticmethod
@@ -2484,7 +2510,13 @@ def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput:
                 height=context["height"],
                 width=context["width"],
             )
-            videos.append(_prepare_minimax_h3_video_output(video))
+            # Rebind rather than append the expression: the local would
+            # otherwise keep the decoded frames' device storage alive for the
+            # rest of the iteration, and the next seed's diffuse() -- the DiT
+            # reload this release exists to make room for -- runs before the
+            # loop rebinds it. post_decode() rebinds for the same reason.
+            video = self._offload_model_cpu_stage_output(_prepare_minimax_h3_video_output(video))
+            videos.append(video)
             audios.append(audio)
         video = videos[0] if len(videos) == 1 else torch.cat(videos, dim=0)
         audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0)
@@ -2789,7 +2821,7 @@ def post_decode(self, state: StepRequestState, **kwargs: Any) -> DiffusionOutput
             height=shape["height"],
             width=shape["width"],
         )
-        video = _prepare_minimax_h3_video_output(video)
+        video = self._offload_model_cpu_stage_output(_prepare_minimax_h3_video_output(video))
         return DiffusionOutput(
             output=(video, audio),
             post_process_func=get_minimax_h3_post_process_func(self.od_config),