From 9479990823ba08880436b60ea3d071fb5b2e2d3d Mon Sep 17 00:00:00 2001 From: GUOGUO Date: Wed, 7 Oct 2026 12:17:30 +0800 Subject: [PATCH 1/5] [Model] MiniMax-H3: overlap Ref2VA CPU noise preparation Signed-off-by: GUOGUO --- .../minimax-h3-ref2va-8-rank-resident.yaml | 25 +++ .../minimax_h3/test_minimax_h3_packing.py | 88 ++++++++++ .../test_ref2va_initial_noise_preparation.py | 109 +++++++++++++ .../models/minimax_h3/condition_noise.py | 150 +++++++++++++----- .../models/minimax_h3/pipeline_minimax_h3.py | 147 +++++++++++++---- vllm_omni/diffusion/models/minimax_h3/vae.py | 64 +++++--- 6 files changed, 498 insertions(+), 85 deletions(-) create mode 100644 recipes/MiniMaxAI/minimax-h3-ref2va-8-rank-resident.yaml create mode 100644 tests/diffusion/models/minimax_h3/test_ref2va_initial_noise_preparation.py diff --git a/recipes/MiniMaxAI/minimax-h3-ref2va-8-rank-resident.yaml b/recipes/MiniMaxAI/minimax-h3-ref2va-8-rank-resident.yaml new file mode 100644 index 00000000000..d98d52a3b55 --- /dev/null +++ b/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"}}} diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py index 4340e093fad..e51bf74b4e2 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py +++ b/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. diff --git a/tests/diffusion/models/minimax_h3/test_ref2va_initial_noise_preparation.py b/tests/diffusion/models/minimax_h3/test_ref2va_initial_noise_preparation.py new file mode 100644 index 00000000000..805fc3b2a1d --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_ref2va_initial_noise_preparation.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from concurrent.futures import ThreadPoolExecutor + +import pytest +import torch + +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, +) + +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 diff --git a/vllm_omni/diffusion/models/minimax_h3/condition_noise.py b/vllm_omni/diffusion/models/minimax_h3/condition_noise.py index 8ac769792a6..e585ce66ea6 100644 --- a/vllm_omni/diffusion/models/minimax_h3/condition_noise.py +++ b/vllm_omni/diffusion/models/minimax_h3/condition_noise.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniMax H3 visual/audio condition-noise augmentation. The request's condition timestep is applied to both the tensor value and the @@ -19,6 +20,86 @@ MINIMAX_H3_AUDIO_COND_CHANNELS = 2 +def _parse_imgvid_condition_shapes( + condition_shapes: Sequence[Sequence[int]], +) -> tuple[list[tuple[int, int, int]], int]: + parsed_shapes: list[tuple[int, int, int]] = [] + expected_rows = 0 + for raw_shape in condition_shapes: + if len(raw_shape) != 3: + raise ValueError( + f"each imgvid condition shape must be (latent_t, latent_h, latent_w), got {list(raw_shape)}" + ) + latent_t, latent_h, latent_w = (int(value) for value in raw_shape) + if latent_t <= 0 or latent_h <= 0 or latent_w <= 0: + raise ValueError(f"imgvid condition shape must be positive, got {list(raw_shape)}") + if latent_h % 2 or latent_w % 2: + raise ValueError( + f"imgvid condition spatial dimensions must be divisible by 2, got {(latent_t, latent_h, latent_w)}" + ) + parsed_shapes.append((latent_t, latent_h, latent_w)) + expected_rows += latent_t * (latent_h // 2) * (latent_w // 2) + if not parsed_shapes: + raise ValueError("condition_shapes must not be empty") + return parsed_shapes, expected_rows + + +def _draw_imgvid_condition_noise_rows( + *, + latent_t: int, + latent_h: int, + latent_w: int, + target_latent_t: int, + imgvid_cond_num_frames: int, + seed: int, +) -> torch.Tensor: + full_t = max(target_latent_t + imgvid_cond_num_frames, latent_t) + generator = torch.Generator(device="cpu").manual_seed(int(seed)) + noise = torch.randn( + 1, + 24, + full_t, + latent_h, + latent_w, + generator=generator, + dtype=torch.float32, + device="cpu", + )[:, :, :latent_t] + return minimax_h3_patchify_video_latent(noise, patch_size=[1, 2, 2]) + + +def minimax_h3_imgvid_cond_noise_rows( + *, + condition_shapes: Sequence[Sequence[int]], + target_latent_t: int, + imgvid_cond_num_frames: int, + seed: int, +) -> torch.Tensor: + """Precompute packed CPU noise rows for the ordered visual conditions.""" + target_latent_t = int(target_latent_t) + imgvid_cond_num_frames = int(imgvid_cond_num_frames) + if target_latent_t <= 0: + raise ValueError(f"target_latent_t must be positive, got {target_latent_t}") + if imgvid_cond_num_frames <= 0: + raise ValueError( + f"imgvid_cond_num_frames must be positive when condition rows exist, got {imgvid_cond_num_frames}" + ) + + parsed_shapes, _ = _parse_imgvid_condition_shapes(condition_shapes) + noise_parts = [ + _draw_imgvid_condition_noise_rows( + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + target_latent_t=target_latent_t, + imgvid_cond_num_frames=imgvid_cond_num_frames, + seed=seed, + ) + for latent_t, latent_h, latent_w in parsed_shapes + ] + return noise_parts[0] if len(noise_parts) == 1 else torch.cat(noise_parts, dim=0) + + def minimax_h3_imgvid_cond_noise_aug_rows( clean_rows: torch.Tensor, *, @@ -27,6 +108,7 @@ def minimax_h3_imgvid_cond_noise_aug_rows( imgvid_cond_num_frames: int, seed: int, noise_aug: float, + precomputed_noise_rows: torch.Tensor | None = None, ) -> torch.Tensor: """Apply the imgvid-condition RF noise recipe to packed clean rows. @@ -54,56 +136,52 @@ def minimax_h3_imgvid_cond_noise_aug_rows( f"imgvid_cond_num_frames must be positive when condition rows exist, got {imgvid_cond_num_frames}" ) - parsed_shapes: list[tuple[int, int, int]] = [] - expected_rows = 0 - for raw_shape in condition_shapes: - if len(raw_shape) != 3: - raise ValueError( - f"each imgvid condition shape must be (latent_t, latent_h, latent_w), got {list(raw_shape)}" - ) - latent_t, latent_h, latent_w = (int(value) for value in raw_shape) - if latent_t <= 0 or latent_h <= 0 or latent_w <= 0: - raise ValueError(f"imgvid condition shape must be positive, got {list(raw_shape)}") - if latent_h % 2 or latent_w % 2: - raise ValueError( - f"imgvid condition spatial dimensions must be divisible by 2, got {(latent_t, latent_h, latent_w)}" - ) - parsed_shapes.append((latent_t, latent_h, latent_w)) - expected_rows += latent_t * (latent_h // 2) * (latent_w // 2) - if not parsed_shapes: - raise ValueError("condition_shapes must not be empty") + parsed_shapes, expected_rows = _parse_imgvid_condition_shapes(condition_shapes) if int(clean_rows.shape[0]) != expected_rows: raise ValueError( f"clean imgvid condition rows {int(clean_rows.shape[0])} != shape-derived rows {expected_rows}" ) + if precomputed_noise_rows is not None and tuple(precomputed_noise_rows.shape) != (expected_rows, 96): + raise ValueError( + "precomputed imgvid noise rows must match the condition shape, " + f"got {tuple(precomputed_noise_rows.shape)}, expected {(expected_rows, 96)}" + ) output_device = clean_rows.device - clean_rows = clean_rows.detach().to(device="cpu", dtype=torch.float32) - out: list[torch.Tensor] = [] + mix_device = output_device if output_device.type == "npu" else torch.device("cpu") + clean_rows = clean_rows.detach().to(device=mix_device, dtype=torch.float32) + out = torch.empty((expected_rows, 96), dtype=torch.float32, device=mix_device) row_offset = 0 timestep = torch.tensor(noise_aug, dtype=torch.float32, device="cpu") + # Keep both coefficients rounded in CPU FP32 before the NPU mix. + noise_timestep = 1.0 - timestep + timestep = timestep.to(device=mix_device) + noise_timestep = noise_timestep.to(device=mix_device) for latent_t, latent_h, latent_w in parsed_shapes: # Official Ref2VA allows a reference video to be longer than the # generated clip. The old implementation sized the draw only from # the target clip and consequently rejected valid long references. - full_t = max(target_latent_t + imgvid_cond_num_frames, latent_t) - generator = torch.Generator(device="cpu").manual_seed(int(seed)) - noise = torch.randn( - 1, - 24, - full_t, - latent_h, - latent_w, - generator=generator, - dtype=torch.float32, - device="cpu", - )[:, :, :latent_t] - noise_rows = minimax_h3_patchify_video_latent(noise, patch_size=[1, 2, 2]).to(dtype=torch.float32) - row_count = int(noise_rows.shape[0]) + row_count = latent_t * (latent_h // 2) * (latent_w // 2) + if precomputed_noise_rows is None: + noise_rows = _draw_imgvid_condition_noise_rows( + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + target_latent_t=target_latent_t, + imgvid_cond_num_frames=imgvid_cond_num_frames, + seed=seed, + ) + else: + noise_rows = precomputed_noise_rows[row_offset : row_offset + row_count] + noise_rows = noise_rows.to(device=mix_device, dtype=torch.float32) clean_part = clean_rows[row_offset : row_offset + row_count].to(torch.float32) - out.append(timestep * clean_part + (1.0 - timestep) * noise_rows) + torch.add( + timestep * clean_part, + noise_timestep * noise_rows, + out=out[row_offset : row_offset + row_count], + ) row_offset += row_count - return torch.cat(out, dim=0).to(device=output_device, dtype=torch.float32).contiguous() + return out.to(device=output_device, dtype=torch.float32).contiguous() def minimax_h3_audio_cond_noise_aug_rows( 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 b1dfc78b9b1..9bdaf3b5383 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -8,6 +8,7 @@ import math import os from collections.abc import Iterable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager, nullcontext from dataclasses import fields, replace from itertools import groupby @@ -108,6 +109,7 @@ from .condition_noise import ( minimax_h3_audio_cond_noise_aug_rows, minimax_h3_imgvid_cond_noise_aug_rows, + minimax_h3_imgvid_cond_noise_rows, ) from .continuation import diffuse_continuation, plan_continuation_windows, resolve_continuation from .denoise_loop import ( @@ -1785,6 +1787,8 @@ def _build_denoise_inputs( sampler: str = "euler", init_latents: tuple[torch.Tensor, torch.Tensor] | None = None, refine: MiniMaxH3LatentRefineSpec | None = None, + precomputed_initial_noise: tuple[torch.Tensor, torch.Tensor] | None = None, + precomputed_visual_condition_noise: torch.Tensor | None = None, ) -> dict[str, Any]: """Build the packed layout, initial rows, anchors, and sigma schedules. @@ -1811,13 +1815,16 @@ def _build_denoise_inputs( shift_scale=audio_shift, base_schedule=base_schedule, ) - initial_video, initial_audio = self._initial_noise( - seed=seed, - latent_t=latent_t, - latent_h=latent_h, - latent_w=latent_w, - audio_t=audio_t, - ) + if precomputed_initial_noise is None: + initial_video, initial_audio = self._initial_noise( + seed=seed, + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + audio_t=audio_t, + ) + else: + initial_video, initial_audio = precomputed_initial_noise if init_latents is not None: if refine is None: raise ValueError("init_latents needs a refine spec to place them on the schedule") @@ -1902,6 +1909,7 @@ def _build_denoise_inputs( imgvid_cond_num_frames=len(condition_shapes), seed=seed, noise_aug=MINIMAX_H3_IMGVID_COND_TIMESTEP, + precomputed_noise_rows=precomputed_visual_condition_noise, ) full_video = torch.zeros( branch.img_pos.shape[0], @@ -2061,6 +2069,8 @@ def diffuse( sampler: str = "euler", init_latents: tuple[torch.Tensor, torch.Tensor] | None = None, refine: MiniMaxH3LatentRefineSpec | None = None, + precomputed_initial_noise: tuple[torch.Tensor, torch.Tensor] | None = None, + precomputed_visual_condition_noise: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: inputs = self._build_denoise_inputs( sampler=sampler, @@ -2097,6 +2107,8 @@ def diffuse( audio_edit_restore_mask_rows=audio_edit_restore_mask_rows, init_latents=init_latents, refine=refine, + precomputed_initial_noise=precomputed_initial_noise, + precomputed_visual_condition_noise=precomputed_visual_condition_noise, ) branch = inputs["branch"] transformer = self._transformer_for_task(task) @@ -2610,12 +2622,23 @@ def _encode_local_media(self, media: MiniMaxH3EncoderMediaInput | None) -> MiniM raise error return self._broadcast_media_conditioning(conditioning) + def _predict_visual_condition_shapes( + self, + media: MiniMaxH3EncoderMediaInput, + ) -> tuple[tuple[int, int, int], ...]: + shapes = [(1, int(image.shape[0]) // 16, int(image.shape[1]) // 16) for image in media.images] + for video in media.videos: + frames = np.asarray(video.detach().cpu().to(torch.uint8).contiguous().numpy()) + shapes.append(self.video_vae.predict_video_latent_shape(frames)) + return tuple(shapes) + def _prepare_local_conditioning( self, raw_prompt: Any, sampling: Any, *, require_external_text: bool = False, + request_context: dict[str, Any] | None = None, ) -> tuple[MiniMaxH3EncoderConditioning, list[tuple[torch.Tensor, torch.Tensor]] | None]: """Prepare media once and encode either shared text or each window's text. @@ -2625,6 +2648,8 @@ def _prepare_local_conditioning( group, rank, world_size = _dit_rank_world() prompts = (sampling.extra_args or {}).get("continuation_prompts") prepared = text_conditioning = None + initial_noise_params = None + visual_noise_params = None error = None if rank == 0: try: @@ -2659,12 +2684,14 @@ def _prepare_local_conditioning( task=task, prepared_reference_videos=self._extract_prepared_reference_videos(raw_prompt), ) - if prompts is not None: + continuation = None + if prompts is not None or (request_context is not None and task == "ref2va"): continuation = resolve_continuation( sampling.extra_args or {}, task=task, step_execution=bool(getattr(self.od_config, "step_execution", False)), ) + if prompts is not None: if continuation is None or not self.load_text_encoder: raise OmniClientError( "continuation_prompts requires continuation mode with a local text encoder" @@ -2680,31 +2707,81 @@ def _prepare_local_conditioning( prepared = replace(prepared, prompt=prompts[0]) # External text belongs to the main prompt, not the first window. text_conditioning = None + if request_context is not None and task == "ref2va" and continuation is None: + initial_noise_params = { + "seed": int(sampling.seed if sampling.seed is not None else 42), + "latent_t": prepared.media.latent_t, + "latent_h": prepared.media.height // 16, + "latent_w": prepared.media.width // 16, + "audio_t": prepared.media.audio_t, + } + condition_shapes = self._predict_visual_condition_shapes(prepared.media) + if condition_shapes: + visual_noise_params = { + "condition_shapes": condition_shapes, + "target_latent_t": prepared.media.latent_t, + "imgvid_cond_num_frames": len(condition_shapes), + "seed": initial_noise_params["seed"], + } except Exception as exc: error = exc _broadcast_rank0_exception(error) - reuse_text = [text_conditioning is not None] + header = [(text_conditioning is not None, initial_noise_params, visual_noise_params)] if world_size > 1: - dist.broadcast_object_list(reuse_text, src=0, group=group) - if reuse_text[0]: - hidden = _broadcast_tensor( - text_conditioning.hidden_states if rank == 0 else None, - dtype=torch.bfloat16, - device=self.device, + dist.broadcast_object_list(header, src=0, group=group) + reuse_text, initial_noise_params, visual_noise_params = header[0] + executor_context = ( + ThreadPoolExecutor(max_workers=1, thread_name_prefix="minimax_h3_initial_noise") + if initial_noise_params is not None + else nullcontext() + ) + # Independent CPU generators preserve sampling; join the worker on request exit. + with executor_context as executor: + initial_noise = ( + executor.submit(self._initial_noise, **initial_noise_params) + if initial_noise_params is not None + else None ) - tags = _broadcast_tensor( - text_conditioning.token_tags if rank == 0 else None, dtype=torch.long, device=self.device + visual_noise = ( + executor.submit(minimax_h3_imgvid_cond_noise_rows, **visual_noise_params) + if visual_noise_params is not None + else None ) - else: - hidden, tags = self.encode_prompt(prepared) - window_text = None - if prompts is not None: - window_text = [(hidden, tags)] - for index, prompt in enumerate(prompts[1:], start=2): - logger.info("MiniMax H3 encoding continuation prompt %d/%d", index, len(prompts)) - window_text.append(self.encode_prompt(replace(prepared, prompt=prompt) if rank == 0 else None)) - media = self._encode_local_media(prepared.media if prepared is not None else None) - return MiniMaxH3EncoderConditioning.from_components(MiniMaxH3TextConditioning(hidden, tags), media), window_text + if reuse_text: + hidden = _broadcast_tensor( + text_conditioning.hidden_states if rank == 0 else None, + dtype=torch.bfloat16, + device=self.device, + ) + tags = _broadcast_tensor( + text_conditioning.token_tags if rank == 0 else None, dtype=torch.long, device=self.device + ) + else: + hidden, tags = self.encode_prompt(prepared) + window_text = None + if prompts is not None: + window_text = [(hidden, tags)] + for index, prompt in enumerate(prompts[1:], start=2): + logger.info("MiniMax H3 encoding continuation prompt %d/%d", index, len(prompts)) + window_text.append(self.encode_prompt(replace(prepared, prompt=prompt) if rank == 0 else None)) + media = self._encode_local_media(prepared.media if prepared is not None else None) + if initial_noise is not None: + assert request_context is not None + request_context["precomputed_initial_noise"] = initial_noise.result() + if visual_noise is not None: + assert request_context is not None + assert visual_noise_params is not None + predicted_shapes = visual_noise_params["condition_shapes"] + actual_shapes = tuple(media.visual_condition_shapes) + if predicted_shapes != actual_shapes: + raise RuntimeError( + "MiniMax H3 predicted visual condition shapes do not match encoded shapes: " + f"predicted {predicted_shapes}, encoded {actual_shapes}" + ) + request_context["precomputed_visual_condition_noise"] = visual_noise.result() + return MiniMaxH3EncoderConditioning.from_components( + MiniMaxH3TextConditioning(hidden, tags), media + ), window_text def _prepare_request_inputs(self, raw_prompt: Any, sampling: Any) -> dict[str, Any]: if (getattr(sampling, "extra_args", None) or {}).get( @@ -2712,15 +2789,17 @@ def _prepare_request_inputs(self, raw_prompt: Any, sampling: Any) -> dict[str, A ) is not None and not self.load_text_encoder: raise OmniClientError("continuation_prompts requires continuation mode with a local text encoder") window_text = None + context: dict[str, Any] = {} if self.load_text_encoder or self.load_vae_encoder: conditioning, window_text = self._prepare_local_conditioning( raw_prompt, sampling, require_external_text=not self.load_text_encoder, + request_context=context, ) else: conditioning = self._extract_encoder_conditioning(raw_prompt) - context = self._prepare_encoder_conditioning_inputs(conditioning, sampling) + context.update(self._prepare_encoder_conditioning_inputs(conditioning, sampling)) if ( context.get("latent_refine") is not None and context.get("latent_upscale") is not None @@ -3007,7 +3086,11 @@ def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput: check_request_cancellation() output_kwargs = {**denoise_kwargs, "seed": output_seed} if context.get("continuation") is None: - video_latent, audio_latent = self.diffuse(**output_kwargs) + video_latent, audio_latent = self.diffuse( + **output_kwargs, + precomputed_initial_noise=context.pop("precomputed_initial_noise", None), + precomputed_visual_condition_noise=context.pop("precomputed_visual_condition_noise", None), + ) else: window_frames, overlap_frames = context["continuation"] video_latent, audio_latent = diffuse_continuation( @@ -3147,7 +3230,11 @@ def prepare_encode(self, state: StepRequestState, **kwargs: Any) -> StepRequestS state.prompt, state.sampling, ) - inputs = self._build_denoise_inputs(**self._denoise_kwargs(context)) + inputs = self._build_denoise_inputs( + **self._denoise_kwargs(context), + precomputed_initial_noise=context.pop("precomputed_initial_noise", None), + precomputed_visual_condition_noise=context.pop("precomputed_visual_condition_noise", None), + ) sigmas_video = inputs["sigmas_video"] sigmas_audio = inputs["sigmas_audio"] diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index c614ea0f293..525e06d94d0 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -780,27 +780,11 @@ def encode_image(self, image: Image.Image) -> torch.Tensor: patch_size=(1, 2, 2), ).float() - def _stream_prepare_video_tensor( + def _stream_video_prep_metadata( self, frames: np.ndarray, - device: torch.device, - ) -> torch.Tensor | None: - """Upload a uint8 ``(T, H, W, 3)`` video as one normalized FP32 tensor. - - The checkpoint's numpy path uploads the whole video as FP32, then - ``transform_tensor`` materializes a second normalized copy, then - ``encode_temporal`` pads through a full-video ``torch.cat``: three - resident pixel-scale copies on the device. Here the padding is - replicated on the host, a single preallocated ``(3, T', H, W)`` tensor - receives clip-by-clip uploads, and the ``÷255 → (x-mean)/std`` chain - runs in place on each clip's staging buffer in the same op order as - ``convert_numpy_to_tensor`` → ``transform_tensor``, so peak device - memory is the output tensor plus one clip instead of three copies. - - Returns ``None`` whenever the checkpoint contract this mirrors (uint8 - frames, ``clip_length`` alignment, processor ``transform`` constants) - is not discoverable; callers fall back to the legacy path. - """ + ) -> tuple[int, int, Any, Any] | None: + """Return the shared stream-prep clip, pad, mean, and std metadata.""" if frames.dtype != np.uint8 or frames.ndim != 4 or frames.shape[-1] != 3: return None if int(frames.shape[0]) == 0: @@ -841,6 +825,48 @@ def _stream_prepare_video_tensor( else: chunks = -(-(num_frames - tail) // clip_length) pad = max(max(chunks, 1) * clip_length + tail - num_frames, 0) + return clip_length, pad, mean, std + + def predict_video_latent_shape(self, frames: np.ndarray) -> tuple[int, int, int]: + """Predict the checkpoint latent shape for its reference-video input.""" + effective_frames = int(frames.shape[0]) + if not _legacy_encode_prep_enabled(): + prep_metadata = self._stream_video_prep_metadata(frames) + if prep_metadata is not None: + effective_frames += prep_metadata[1] + + processor = self.model.processor + latent_t = int(processor.get_latent_length(effective_frames)) + height, width = processor._align_to_total_patch_size(int(frames.shape[1]), int(frames.shape[2])) + ratio = int(self.model.vae_ratio) + return latent_t, int(height // ratio), int(width // ratio) + + def _stream_prepare_video_tensor( + self, + frames: np.ndarray, + device: torch.device, + ) -> torch.Tensor | None: + """Upload a uint8 ``(T, H, W, 3)`` video as one normalized FP32 tensor. + + The checkpoint's numpy path uploads the whole video as FP32, then + ``transform_tensor`` materializes a second normalized copy, then + ``encode_temporal`` pads through a full-video ``torch.cat``: three + resident pixel-scale copies on the device. Here the padding is + replicated on the host, a single preallocated ``(3, T', H, W)`` tensor + receives clip-by-clip uploads, and the ``÷255 → (x-mean)/std`` chain + runs in place on each clip's staging buffer in the same op order as + ``convert_numpy_to_tensor`` → ``transform_tensor``, so peak device + memory is the output tensor plus one clip instead of three copies. + + Returns ``None`` whenever the checkpoint contract this mirrors (uint8 + frames, ``clip_length`` alignment, processor ``transform`` constants) + is not discoverable; callers fall back to the legacy path. + """ + prep_metadata = self._stream_video_prep_metadata(frames) + if prep_metadata is None: + return None + clip_length, pad, mean, std = prep_metadata + num_frames = int(frames.shape[0]) if pad: frames = np.concatenate([frames, np.repeat(frames[-1:], pad, axis=0)]) num_frames += pad From 6c2fdea23485eac6171e622b22d776f554404287 Mon Sep 17 00:00:00 2001 From: GUOGUO Date: Thu, 8 Oct 2026 09:54:46 +0800 Subject: [PATCH 2/5] [Bugfix] Align H3 mixed-reference handoff test with prepared noise Signed-off-by: GUOGUO --- .../models/minimax_h3/test_minimax_h3_execution_modes.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py index b145d3d445a..b1d6b4baa93 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py +++ b/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) @@ -263,10 +264,18 @@ def encode_audio(waveform, sample_rate): pipeline.forward(SimpleNamespace(prompts=[bridged], sampling_params=request.sampling_params)) external_kwargs = pipeline.diffuse.call_args.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() From beb6c38c4e59493964ee1b8c5a2027a5ea131342 Mon Sep 17 00:00:00 2001 From: GUOGUO Date: Fri, 9 Oct 2026 12:04:13 +0800 Subject: [PATCH 3/5] [Bugfix] Release consumed H3 host noise and cancel pending preparation Signed-off-by: GUOGUO --- .../test_ref2va_initial_noise_preparation.py | 86 +++++++++++++++++++ .../models/minimax_h3/pipeline_minimax_h3.py | 25 ++++-- 2 files changed, 102 insertions(+), 9 deletions(-) diff --git a/tests/diffusion/models/minimax_h3/test_ref2va_initial_noise_preparation.py b/tests/diffusion/models/minimax_h3/test_ref2va_initial_noise_preparation.py index 805fc3b2a1d..a2f2bc51774 100644 --- a/tests/diffusion/models/minimax_h3/test_ref2va_initial_noise_preparation.py +++ b/tests/diffusion/models/minimax_h3/test_ref2va_initial_noise_preparation.py @@ -1,16 +1,21 @@ # 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] @@ -107,3 +112,84 @@ def test_shared_denoise_kwargs_exclude_first_seed_noise(pipeline): 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) 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 9bdaf3b5383..e982b42c8c5 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -2069,9 +2069,10 @@ def diffuse( sampler: str = "euler", init_latents: tuple[torch.Tensor, torch.Tensor] | None = None, refine: MiniMaxH3LatentRefineSpec | None = None, - precomputed_initial_noise: tuple[torch.Tensor, torch.Tensor] | None = None, - precomputed_visual_condition_noise: torch.Tensor | None = None, + request_context: dict[str, Any] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: + # Consume request-owned noise here so caller keyword dictionaries cannot + # keep the host buffers alive throughout denoising. inputs = self._build_denoise_inputs( sampler=sampler, task=task, @@ -2107,8 +2108,12 @@ def diffuse( audio_edit_restore_mask_rows=audio_edit_restore_mask_rows, init_latents=init_latents, refine=refine, - precomputed_initial_noise=precomputed_initial_noise, - precomputed_visual_condition_noise=precomputed_visual_condition_noise, + precomputed_initial_noise=( + None if request_context is None else request_context.pop("precomputed_initial_noise", None) + ), + precomputed_visual_condition_noise=( + None if request_context is None else request_context.pop("precomputed_visual_condition_noise", None) + ), ) branch = inputs["branch"] transformer = self._transformer_for_task(task) @@ -2730,13 +2735,13 @@ def _prepare_local_conditioning( if world_size > 1: dist.broadcast_object_list(header, src=0, group=group) reuse_text, initial_noise_params, visual_noise_params = header[0] - executor_context = ( + executor = ( ThreadPoolExecutor(max_workers=1, thread_name_prefix="minimax_h3_initial_noise") if initial_noise_params is not None - else nullcontext() + else None ) # Independent CPU generators preserve sampling; join the worker on request exit. - with executor_context as executor: + try: initial_noise = ( executor.submit(self._initial_noise, **initial_noise_params) if initial_noise_params is not None @@ -2782,6 +2787,9 @@ def _prepare_local_conditioning( return MiniMaxH3EncoderConditioning.from_components( MiniMaxH3TextConditioning(hidden, tags), media ), window_text + finally: + if executor is not None: + executor.shutdown(wait=True, cancel_futures=True) def _prepare_request_inputs(self, raw_prompt: Any, sampling: Any) -> dict[str, Any]: if (getattr(sampling, "extra_args", None) or {}).get( @@ -3088,8 +3096,7 @@ def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput: if context.get("continuation") is None: video_latent, audio_latent = self.diffuse( **output_kwargs, - precomputed_initial_noise=context.pop("precomputed_initial_noise", None), - precomputed_visual_condition_noise=context.pop("precomputed_visual_condition_noise", None), + request_context=context, ) else: window_frames, overlap_frames = context["continuation"] From 2e77dfd29595d1b823aea44c10420815af280a09 Mon Sep 17 00:00:00 2001 From: GUOGUO Date: Fri, 9 Oct 2026 16:24:09 +0800 Subject: [PATCH 4/5] [Test] Align H3 handoff parity with request-owned noise Signed-off-by: GUOGUO --- .../minimax_h3/test_minimax_h3_execution_modes.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py index b1d6b4baa93..e1da29b23fd 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py @@ -222,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) @@ -262,7 +271,7 @@ 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 From 4c2a91ef4d3ad3396e1313585bbde53745ef96fd Mon Sep 17 00:00:00 2001 From: GUOGUO Date: Fri, 9 Oct 2026 21:03:58 +0800 Subject: [PATCH 5/5] [CI] Match H3 accuracy worker timeout to request budget Signed-off-by: GUOGUO --- .../minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py b/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py index 4ae45e4e17d..98399a30c92 100644 --- a/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py +++ b/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 = {