From 6cd976e521db6c2525947df61a1b792b3f651da8 Mon Sep 17 00:00:00 2001 From: heyuanliu-intel Date: Mon, 24 Aug 2026 03:48:16 +0000 Subject: [PATCH] [Bugfix][MiniMax-H3] Release decoded outputs before the next seed ## Purpose **Problem:** with model-level CPU offload, #6072 evicts the DiT before VAE decode, but the decoded output stays on the device until the request returns. For `num_outputs_per_prompt > 1` the appended output of seed N is still resident when seed N+1 reloads the DiT, so the peak the eviction removed comes back; with a large output even the first request can OOM in decode. **Fix:** on the executor reply rank, move each decoded output to the host inside the same seed iteration while active model-level offload is installed; non-reply ranks keep a metadata placeholder. Video is released after `_prepare_minimax_h3_video_output` quantizes to `uint8`, so #6824's transfer optimisation is kept. Other offload modes and the no-offload path are unchanged; the design doc records the exception. ## Test Plan **Reproduce:** 8 XPU cards, TP4xSP2, ref2va at 1344x768 / 124 frames / 50 steps, Cache-DiT and model-level CPU offload for encoders and VAE. The first request finishes denoising (49/49) and then fails in `klvae.py:621 write_part -> torch.empty` with `XPU out of memory. Tried to allocate 1.43 GiB ... 1.42 GiB is free`, HTTP 500. ```bash pytest -q tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py ``` **vLLM Version:** as pinned by `requirements/`. **vLLM-Omni Commit:** rebased onto current `main`. ## Test Result CPU tests pass (release only under active model-level offload; reply rank gets a host tensor, non-reply ranks a placeholder; video copied after quantization; no-offload path unchanged). On the 8-card run with this change: 0 OOM, 3/3 requests HTTP 200, 24 `Evicted 1 DiT module(s) (15.43 GiB)` lines. The device run used one output per prompt; the multi-seed overlap is covered by the CPU tests. Signed-off-by: heyuanliu-intel --- docs/design/feature/async_diffusion_output.md | 21 +- .../minimax_h3/test_minimax_h3_offload.py | 242 +++++++++++++++++- .../models/minimax_h3/pipeline_minimax_h3.py | 36 ++- 3 files changed, 290 insertions(+), 9 deletions(-) diff --git a/docs/design/feature/async_diffusion_output.md b/docs/design/feature/async_diffusion_output.md index 733478883bd..07734eda239 100644 --- a/docs/design/feature/async_diffusion_output.md +++ b/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 | diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py index c993de492c2..5f025f758d4 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py +++ b/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 diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 702eaf59461..057765fbe36 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/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),