diff --git a/docs/user_guide/diffusion/execution_modes.md b/docs/user_guide/diffusion/execution_modes.md index ae2f868d8..f9fbc1be0 100644 --- a/docs/user_guide/diffusion/execution_modes.md +++ b/docs/user_guide/diffusion/execution_modes.md @@ -101,8 +101,16 @@ set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or configure `--max-num-seqs >1`. See the [HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md) for its validated configuration. Helios supports single-request step execution only: use -`--step-execution --max-num-seqs 1` for Helios. Consult the selected pipeline's -documentation and source for the latest support status. +`--step-execution --max-num-seqs 1` for Helios. MiniMax H3 supports step-wise +continuous batching by packing co-batched requests into one sequence that keeps +a separate attention document per request; that layout needs a backend which +honors the packed `cu_seqlens` metadata, so run it with +`--diffusion-attention-backend FLASH_ATTN` (other backends stay correct but +fall back to one transformer forward per request). Batching H3 does not improve +its throughput — see the measured numbers in the +[MiniMax-H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md) +— so keep `--max-num-seqs 1` unless you need step-level scheduling. Consult the +selected pipeline's documentation and source for the latest support status. ## Streaming Output diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 56dd04901..722b5f8b0 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -380,6 +380,63 @@ No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while `task=ref2va` routes to `Ref2VA/transformer`. T2VA uses the FL2VA DiT. +### Step execution and continuous batching + +H3 implements the step-wise execution contract, so the scheduler can admit and +retire requests between denoise steps instead of running one request end to +end. Add the feature gate, then raise `--max-num-seqs` to co-batch: + +```bash +--step-execution --max-num-seqs 4 +``` + +Co-batched requests are packed into a single sequence that keeps one attention +document per request, so a batch costs one DiT forward. That packing requires +`--diffusion-attention-backend FLASH_ATTN`; other backends fall back to one +forward per request. `--max-num-seqs 1` keeps the conservative single-request +step path. Cache acceleration (`--cache-backend`) is not available in step mode. + +!!! warning "Co-batching does not improve H3 throughput for large simultaneous requests" + Keep `--max-num-seqs 1` unless you specifically need scheduler-level control + (admitting and retiring requests between denoise steps). Measured on two + H100s (TP2, BF16, 672x384, 209 frames, 30 steps, 4 requests at concurrency + 4 with `--request-rate inf`, i.e. all four submitted at time 0; one packed + request is 16384 rows): + + | Configuration | Wall time | Mean latency | Peak memory | + |---------------|-----------|--------------|-------------| + | request mode | 174.8 s | 111.5 s | 72.4 GB | + | `--step-execution --max-num-seqs 1` | 179.0 s | 113.8 s | 72.4 GB | + | `--step-execution --max-num-seqs 4` | 182.1 s | 175.7 s | 78.3 GB | + + A single H3 denoise step is a compute-bound dense GEMM over an already long + packed sequence, so fusing N requests costs N times the FLOPs and buys almost + no amortization — unlike LLM decoding, which is memory-bandwidth bound. + Going from one request per step to four cuts the per-request denoise cost + only from 1.323 s to 1.291 s (2.4%), which the step-mode bookkeeping then + spends. Mean latency degrades further because co-batched requests finish + together instead of staggered. Quantization moves the absolute numbers + without changing this: with online `int8` the same workload runs in 153.3 s + at 56.9 GB (request mode), and `--max-num-seqs 4` is still 5.0% slower than + request mode. + + Two workloads outside this table are unmeasured and may behave differently: + + - **Staggered arrivals (admission latency).** A single 16384-row H3 request + already saturates the GPUs, so submitting all requests at time 0 is the + one arrival pattern where co-batching cannot win: it can only bunch + completions. In request mode a new request queues behind the whole + in-flight generation (~45 s at this size); step mode admits at the next + denoise-step boundary (~1.3 s). Whether that shows up as a wall-clock + benefit under Poisson arrivals is not yet measured for H3. To reproduce, + run request mode vs `--step-execution --max-num-seqs 4` with + `diffusion_benchmark_serving.py --request-rate 0.05` (roughly one request + every 20 s) and report mean / p95 latency, which then includes queueing. + - **Small requests.** The numbers above are drawn from 672x384 / 209-frame + requests. A short clip at lower resolution packs a few thousand rows and + may not saturate the hardware; co-batching may amortize better there. + This is also unmeasured. + ### Online FP8 quantization MiniMax H3 supports online FP8 quantization of both the DiT and the Qwen3-VL @@ -924,7 +981,12 @@ vllm serve "${MODEL_ROOT}/FL2VA" \ - TeaCache is calibrated for FL2VA only; Ref2VA requests run uncached. - Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads both task-specific DiTs, and loads shared components once from `FL2VA`. -- H3 currently executes one generation request per diffusion batch. +- Request mode executes one generation request per diffusion batch. Use + `--step-execution` with `--max-num-seqs N` to admit several requests at once + (see + [Execution modes](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/execution_modes.md)); + step + mode does not support `cache_backend`. - The first regional-compile request is a warmup and should not be included in steady-state performance measurements. - The serving path accepts fewer references than the model supports. H3 documents up diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 7bede350b..32e1518dc 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -821,6 +821,82 @@ def test_packed_attention_keeps_padding_mask_for_other_backends(): ) +def _fake_packed_attention( + backend_name: str, + *, + prefix_kv_slicing: bool = False, + supports_multi_doc: bool = False, +): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3Attention, + ) + + class FakeBackend: + supports_prefix_kv_slicing = prefix_kv_slicing + + @staticmethod + def get_name() -> str: + return backend_name + + @classmethod + def supports_multi_doc_packed_varlen(cls) -> bool: + return supports_multi_doc + + class FakeAttention(torch.nn.Module): + attn_backend = FakeBackend + + def __init__(self): + super().__init__() + self.metadata = None + + def forward(self, query, key, value, metadata): + self.metadata = metadata + return query + + attention = object.__new__(MiniMaxH3Attention) + torch.nn.Module.__init__(attention) + attention.attention = FakeAttention() + return attention + + +def test_packed_attention_drops_prefix_semantics_for_a_co_batched_layout(): + """A batch's valid rows are block-diagonal, so no prefix can describe them.""" + attention = _fake_packed_attention("FLASH_ATTN", supports_multi_doc=True) + q = torch.randn(12, 2, 4) + + attention._run_packed_attention( + q, + q, + q, + # Two requests: rows + padding tail each. + cu_seqlens=torch.tensor([0, 5, 6, 11, 12], dtype=torch.int32), + max_seqlen=5, + packed_total=12, + num_requests=2, + ) + + assert attention.attention.metadata.attn_mask is None + assert attention.attention.metadata.extra["valid_kv_length"] == 12 + + +@pytest.mark.parametrize("backend_name", ["TORCH_SDPA", "FLASH_ATTN"]) +def test_packed_attention_rejects_backends_without_multi_doc_capability(backend_name): + """Backend names alone must not allow cross-request attention.""" + attention = _fake_packed_attention(backend_name, supports_multi_doc=False) + q = torch.randn(12, 2, 4) + + with pytest.raises(ValueError, match=backend_name): + attention._run_packed_attention( + q, + q, + q, + cu_seqlens=torch.tensor([0, 5, 6, 11, 12], dtype=torch.int32), + max_seqlen=5, + packed_total=12, + num_requests=2, + ) + + def test_reference_image_resize_contract(): from PIL import Image diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py new file mode 100644 index 000000000..33aec8c0f --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py @@ -0,0 +1,471 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax H3 step-wise execution (continuous batching) contract tests. + +These run on CPU against a stand-in DiT, so they cover the packing, the +per-request scheduler math, and the runner-facing state wiring without needing +checkpoint weights. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +_HIDDEN = 8 + + +class _SegmentMeanModel: + """Stand-in DiT whose output depends on packed-document boundaries. + + Every row is mixed with the mean of its own ``cu_seqlens`` document and with + its own timestep, so a wrong document boundary, row offset, or timestep + assignment shows up as a numeric difference instead of passing silently. + """ + + def __call__(self, **kwargs): + x = kwargs["x"][0] + audio_x = kwargs["audio_x"][0] + bounds = kwargs["packed_seq_params"]["cu_seqlens_q"].tolist() + img_pos = kwargs["img_pos_info"]["position_ids"] + audio_pos = kwargs["audio_pos_info"]["position_ids"] + + pooled_video = torch.zeros_like(x) + pooled_audio = torch.zeros_like(audio_x) + for start, stop in zip(bounds[:-1], bounds[1:]): + if stop <= start: + continue + pooled_video[start:stop] = x[start:stop].mean(dim=0, keepdim=True) + pooled_audio[start:stop] = audio_x[start:stop].mean(dim=0, keepdim=True) + + row_timesteps = kwargs["unique_timesteps"][kwargs["inverse_indices"]].unsqueeze(-1) + video = (pooled_video + 0.5 * x + row_timesteps)[img_pos] + audio = (pooled_audio + 0.5 * audio_x + row_timesteps)[audio_pos] + if not kwargs.get("skip_mask_out_condition", False): + video = video * kwargs["update_mask"].view(-1).unsqueeze(-1) + return video, audio + + +def _make_branch(*, text_len: int, latent_t: int, latent_h: int, latent_w: int, audio_t: int, seed: int): + from vllm_omni.diffusion.models.minimax_h3.denoise_loop import MiniMaxH3DenoiseBranch + from vllm_omni.diffusion.models.minimax_h3.packed_sequence import minimax_h3_packed_sequence + + packed = minimax_h3_packed_sequence( + text_len=text_len, + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + audio_t=audio_t, + include_keyframe_cond=False, + ) + generator = torch.Generator().manual_seed(seed) + text_embeddings = torch.randn(text_len, _HIDDEN, generator=generator, dtype=torch.float32) + branch = MiniMaxH3DenoiseBranch( + packed=packed, + text_embeddings=text_embeddings, + token_tags=packed["token_tags"], + device=torch.device("cpu"), + ) + video_rows = torch.randn(int(branch.img_pos.shape[0]), 96, generator=generator, dtype=torch.float32) + audio_rows = torch.randn(int(branch.audio_pos.shape[0]), 32, generator=generator, dtype=torch.float32) + return branch, video_rows, audio_rows + + +def _sigmas(num_points: int, shift: float) -> list[float]: + from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_time_shift_sigmas + + return minimax_h3_time_shift_sigmas(num_steps=num_points, shift_scale=shift) + + +def _make_state(request_id: str, model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio): + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + from vllm_omni.diffusion.worker.utils import StepRequestState + + state = StepRequestState(request_id=request_id, sampling=SimpleNamespace()) + state.latents = video_rows.clone() + state.timesteps = torch.tensor([1.0 - sigma for sigma in sigmas_video[:-1]], dtype=torch.float32) + state.step_index = 0 + state.extra = { + mod._STEP_BRANCH: branch, + # Co-batched requests must share one DiT instance, or denoise_step() + # treats the batch as mixed-task and falls back to one forward each. + mod._STEP_TRANSFORMER: model, + mod._STEP_AUDIO_ROWS: audio_rows.clone(), + mod._STEP_COND_ANCHOR: None, + mod._STEP_AUDIO_ANCHOR: None, + mod._STEP_SIGMAS_VIDEO: sigmas_video, + mod._STEP_SIGMAS_AUDIO: sigmas_audio, + } + return state + + +def _step_pipeline(model, *, packed_batch_supported: bool = True): + """A pipeline instance carrying only what the step methods touch.""" + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + pipeline.transformer = model + pipeline.device = torch.device("cpu") + pipeline._transformer_for_task = lambda task: model + pipeline._packed_batch_supported = lambda transformer: packed_batch_supported + return pipeline + + +def test_step_execution_matches_request_mode_denoise_loop(): + """Stepping through the contract must reproduce the request-mode loop.""" + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop + + model = _SegmentMeanModel() + branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=5) + sigmas_video = _sigmas(6, 12.0) + sigmas_audio = _sigmas(6, 3.0) + + reference_video, reference_audio = minimax_h3_denoise_loop( + model=model, + positive=branch, + initial_video_rows=video_rows, + initial_audio_rows=audio_rows, + keyframe_cond_rows=None, + sigmas_video=sigmas_video, + sigmas_audio=sigmas_audio, + device=torch.device("cpu"), + ) + + pipeline = _step_pipeline(model) + state = _make_state("req-0", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio) + input_batch = SimpleNamespace(states=(state,)) + + steps = 0 + while not state.denoise_completed: + noise_pred = pipeline.denoise_step(input_batch, states=[state]) + pipeline.step_scheduler(state, noise_pred) + steps += 1 + + assert steps == len(sigmas_video) - 1 + assert state.total_steps == steps + torch.testing.assert_close(state.latents, reference_video) + torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], reference_audio) + + +def test_batched_step_execution_matches_independent_requests(): + """Two co-batched requests must land where they would have landed alone.""" + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + model = _SegmentMeanModel() + specs = [ + dict(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=6), + # Exactly 64 packed rows exercises the no-padding-tail boundary case. + dict(text_len=46, latent_t=7, latent_h=2, latent_w=4, audio_t=2, seed=7), + ] + # Different step counts, so the batch composition changes mid-flight. + schedules = [(_sigmas(6, 12.0), _sigmas(6, 3.0)), (_sigmas(4, 12.0), _sigmas(4, 3.0))] + + pipeline = _step_pipeline(model) + + alone: list[tuple[torch.Tensor, torch.Tensor]] = [] + for spec, (sigmas_video, sigmas_audio) in zip(specs, schedules): + branch, video_rows, audio_rows = _make_branch(**spec) + state = _make_state("solo", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio) + while not state.denoise_completed: + pipeline.step_scheduler(state, pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state])) + alone.append((state.latents, state.extra[mod._STEP_AUDIO_ROWS])) + + states = [] + for index, (spec, (sigmas_video, sigmas_audio)) in enumerate(zip(specs, schedules)): + branch, video_rows, audio_rows = _make_branch(**spec) + if index == 1: + assert branch.used_len == branch.seq_len + states.append(_make_state(f"req-{index}", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)) + + active = list(states) + while active: + noise_pred = pipeline.denoise_step(SimpleNamespace(states=tuple(active)), states=active) + offset = 0 + for state in active: + rows = state.latents.shape[0] + pipeline.step_scheduler(state, noise_pred[offset : offset + rows]) + offset += rows + assert offset == noise_pred.shape[0] + # Finished requests leave the batch, exactly like the runner drops them. + active = [state for state in active if not state.denoise_completed] + + for state, (expected_video, expected_audio) in zip(states, alone): + torch.testing.assert_close(state.latents, expected_video) + torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], expected_audio) + + +def test_both_modes_publish_denoise_progress_for_gated_attention(): + """TRTLLM's skip gate and RAINFUSION's warmup stay dense without these.""" + from vllm_omni.diffusion import forward_context as fc + from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop + + model = _SegmentMeanModel() + branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=14) + sigmas_video = _sigmas(4, 12.0) + sigmas_audio = _sigmas(4, 3.0) + + published: list[tuple[int | None, float | None]] = [] + + class _Recorder: + def __init__(self): + self.denoise_step_idx = None + self.denoise_timestep = None + + def __setattr__(self, name, value): + object.__setattr__(self, name, value) + if name == "denoise_timestep": + published.append((self.denoise_step_idx, value)) + + recorder = _Recorder() + published.clear() # drop the pair __init__ recorded + original = fc._forward_context + fc._forward_context = recorder + try: + minimax_h3_denoise_loop( + model=model, + positive=branch, + initial_video_rows=video_rows, + initial_audio_rows=audio_rows, + keyframe_cond_rows=None, + sigmas_video=sigmas_video, + sigmas_audio=sigmas_audio, + device=torch.device("cpu"), + ) + request_mode = list(published) + + published.clear() + pipeline = _step_pipeline(model) + state = _make_state("req-0", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio) + while not state.denoise_completed: + pipeline.step_scheduler(state, pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state])) + step_mode = list(published) + finally: + fc._forward_context = original + + num_steps = len(sigmas_video) - 1 + expected = [(step, sigmas_video[step]) for step in range(num_steps)] + assert request_mode == expected + [(None, None)] + # Step mode has no loop to close, so it publishes only the per-step pairs. + assert step_mode == expected + + +def test_mixed_step_batch_leaves_gated_attention_dense(): + """A batch spanning different steps has no single progress point to publish.""" + from vllm_omni.diffusion import forward_context as fc + + model = _SegmentMeanModel() + first, first_video, first_audio = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=15) + second, second_video, second_audio = _make_branch( + text_len=5, latent_t=3, latent_h=6, latent_w=4, audio_t=2, seed=16 + ) + sigmas_video, sigmas_audio = _sigmas(6, 12.0), _sigmas(6, 3.0) + states = [ + _make_state("req-0", model, first, first_video, first_audio, sigmas_video, sigmas_audio), + _make_state("req-1", model, second, second_video, second_audio, sigmas_video, sigmas_audio), + ] + states[1].step_index = 2 # admitted later, so it trails the batch + + recorder = SimpleNamespace(denoise_step_idx="unset", denoise_timestep="unset") + original = fc._forward_context + fc._forward_context = recorder + try: + _step_pipeline(model).denoise_step(SimpleNamespace(states=tuple(states)), states=states) + finally: + fc._forward_context = original + + assert recorder.denoise_step_idx is None + assert recorder.denoise_timestep is None + + +def test_prepare_encode_seeds_runner_visible_state(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=8) + sigmas_video = _sigmas(6, 12.0) + sigmas_audio = _sigmas(6, 3.0) + context = { + "height": 96, + "width": 64, + "latent_t": 2, + "latent_h": 4, + "latent_w": 6, + "audio_t": 3, + **{key: None for key in mod._MINIMAX_H3_DENOISE_INPUT_KEYS}, + } + + pipeline = _step_pipeline(_SegmentMeanModel()) + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_extract_prompt", staticmethod(lambda _: ("a prompt", {}))) + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_prepare_request_inputs", lambda self, **_: context) + monkeypatch.setattr( + mod.MiniMaxH3Pipeline, + "_build_denoise_inputs", + lambda self, **_: { + "branch": branch, + "video_rows": video_rows, + "audio_rows": audio_rows, + "cond_anchor": None, + "audio_anchor": None, + "sigmas_video": sigmas_video, + "sigmas_audio": sigmas_audio, + }, + ) + + from vllm_omni.diffusion.worker.utils import StepRequestState + + state = StepRequestState( + request_id="req-0", + sampling=SimpleNamespace(num_outputs_per_prompt=1), + prompt="a prompt", + ) + pipeline.prepare_encode(state) + + # The runner slices the batched velocity by this row count. + assert state.latents.shape == (int(branch.img_pos.shape[0]), 96) + assert state.total_steps == len(sigmas_video) - 1 + assert state.step_index == 0 + assert state.do_true_cfg is False + torch.testing.assert_close(state.current_timestep, torch.tensor(1.0 - sigmas_video[0])) + assert state.extra[mod._STEP_BRANCH] is branch + assert state.extra[mod._STEP_SHAPE]["height"] == 96 + + +def test_prepare_encode_rejects_request_mode_only_features(): + """Multi-output and DLO have no representation in the step contract.""" + from vllm_omni.diffusion.worker.utils import StepRequestState + from vllm_omni.errors import OmniClientError + + # A request state carries exactly one latent tensor. + multi_output = StepRequestState( + request_id="req-0", + sampling=SimpleNamespace(num_outputs_per_prompt=2), + prompt="a prompt", + ) + with pytest.raises(OmniClientError, match="one output per request"): + _step_pipeline(_SegmentMeanModel()).prepare_encode(multi_output) + + # Distributed layerwise offload streams the DiT around a whole denoise loop. + single_output = StepRequestState( + request_id="req-1", + sampling=SimpleNamespace(num_outputs_per_prompt=1), + prompt="a prompt", + ) + dlo_pipeline = _step_pipeline(_SegmentMeanModel()) + dlo_pipeline._dlo_residency_controller = object() + with pytest.raises(ValueError, match="distributed layerwise offload"): + dlo_pipeline.prepare_encode(single_output) + + +def test_prepare_encode_rejects_high_quality_cache_dit(): + """quality=high installs a transformer-scoped Cache-DiT profile that + would leak across interleaved or co-batched step-mode requests.""" + from vllm_omni.diffusion.worker.utils import StepRequestState + from vllm_omni.errors import OmniClientError + + state = StepRequestState( + request_id="req-hi", + sampling=SimpleNamespace(num_outputs_per_prompt=1, quality="high"), + prompt="a prompt", + ) + with pytest.raises(OmniClientError, match="quality=high"): + _step_pipeline(_SegmentMeanModel()).prepare_encode(state) + + +def _fake_attention_module( + *, + use_ring: bool, + backend: str = "FLASH_ATTN", + supports_multi_doc: bool = True, +): + """Build a bare MiniMaxH3Attention with the attributes ``_packed_batch_supported`` reads. + + ``supports_multi_doc`` models the platform-dependent capability probe on + ``AttentionBackend`` (e.g. FLASH_ATTN returns True on CUDA/ROCm/MUSA and + False on NPU/XPU); ``backend`` remains only for logging/back-compat. + """ + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import MiniMaxH3Attention + + attn = object.__new__(MiniMaxH3Attention) + attn.attention = SimpleNamespace( + attn_backend=SimpleNamespace( + get_name=lambda: backend, + supports_multi_doc_packed_varlen=lambda: supports_multi_doc, + ), + use_ring=use_ring, + ) + return attn + + +class _FakeTransformer: + """The minimal ``modules()`` protocol ``_packed_batch_supported`` walks.""" + + def __init__(self, attentions): + self._attentions = list(attentions) + + def modules(self): + return iter([self, *self._attentions]) + + +@pytest.mark.parametrize( + "attention", + [ + _fake_attention_module(use_ring=True), + _fake_attention_module(use_ring=False, backend="XFORMERS", supports_multi_doc=False), + _fake_attention_module(use_ring=False, backend="FLASH_ATTN", supports_multi_doc=False), + ], +) +def test_packed_batch_rejects_backends_that_cannot_isolate_requests(attention): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + + assert MiniMaxH3Pipeline._packed_batch_supported(_FakeTransformer([attention])) is False + + +def test_broadcast_rank0_exception_single_rank_reraises(): + """Single-rank execution has no group; the helper just reraises.""" + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _broadcast_rank0_exception, + ) + from vllm_omni.errors import OmniClientError + + _broadcast_rank0_exception(None) + with pytest.raises(OmniClientError, match="bad ref"): + _broadcast_rank0_exception(OmniClientError("bad ref")) + + +def test_broadcast_rank0_exception_propagates_to_non_zero_ranks(monkeypatch): + """A rank-0 error becomes a matching client error on every other DiT rank.""" + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + from vllm_omni.errors import OmniClientError + + def fake_rank_world(rank): + return lambda: (object(), rank, 4) + + def make_broadcast(rank0_payload): + def fake_broadcast(payload_list, *, src, group): + payload_list[0] = rank0_payload + + return fake_broadcast + + monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(0)) + err = OmniClientError("invalid reference-video file", status_code=422, error_type="UnprocessableEntityError") + rank0_payload = { + "type": type(err).__name__, + "message": str(err), + "status_code": err.status_code, + "error_type": err.error_type, + } + monkeypatch.setattr(mod.dist, "broadcast_object_list", make_broadcast(rank0_payload)) + with pytest.raises(OmniClientError) as rank0_info: + mod._broadcast_rank0_exception(err) + assert rank0_info.value is err + + monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(2)) + with pytest.raises(OmniClientError) as rank2_info: + mod._broadcast_rank0_exception(None) + assert rank2_info.value.status_code == 422 + assert rank2_info.value.error_type == "UnprocessableEntityError" + assert "invalid reference-video file" in str(rank2_info.value) diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index 56fc95912..3641ef01b 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -25,7 +25,7 @@ from vllm_omni.diffusion.sched import ( StepScheduler, ) from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData -from vllm_omni.diffusion.worker.utils import RunnerOutput +from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -1559,6 +1559,39 @@ class TestStepScheduler: assert sched_output.num_running_reqs == 2 assert sched_output.num_waiting_reqs == 0 + def test_mixed_batch_error_finishes_only_failed_request(self) -> None: + scheduler = StepScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=2)) + + req_bad = scheduler.add_request(_make_step_request("bad", num_inference_steps=2)) + req_good = scheduler.add_request(_make_step_request("good", num_inference_steps=2)) + first = scheduler.schedule() + + finished = scheduler.update_from_output( + first, + BatchRunnerOutput.from_list( + [ + _make_step_output( + req_bad, + step_index=0, + finished=True, + error="invalid request input", + ), + _make_step_output(req_good, step_index=1), + ] + ), + ) + + assert finished == {req_bad} + assert scheduler.get_request_state(req_bad).status == DiffusionRequestStatus.FINISHED_ERROR + assert scheduler.get_request_state(req_good).status == DiffusionRequestStatus.RUNNING + assert scheduler.get_request_state(req_good).req.sampling_params.step_index == 1 + + second = scheduler.schedule() + assert second.finished_req_ids == {req_bad} + assert _new_ids(second) == [] + assert _cached_ids(second) == [req_good] + def test_step_batch_rejects_different_sampling_key(self) -> None: scheduler = StepScheduler() scheduler.initialize(SimpleNamespace(max_num_seqs=3)) diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 96f8249c5..bbacf7f2c 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -47,6 +47,7 @@ from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.utils import RunnerOutput, StepRequestState from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.errors import OmniClientError from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -102,6 +103,22 @@ class _StepPipeline: return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32)) +class _PerRequestErrorStepPipeline(_StepPipeline): + def prepare_encode(self, state, **kwargs): + if state.prompt == "fail-prepare": + raise OmniClientError( + "invalid request input", + status_code=422, + error_type="UnprocessableEntityError", + ) + return super().prepare_encode(state, **kwargs) + + def denoise_step(self, input_batch, **kwargs): + del kwargs + self.denoise_calls += 1 + return torch.ones_like(input_batch.latents) + + class _ProfilingStepPipeline(_StepPipeline): enable_diffusion_pipeline_profiler = True @@ -697,6 +714,72 @@ class TestRunner: result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) assert len(result) == 2 + def test_prepare_error_isolated_from_healthy_request(self, monkeypatch): + runner = _make_runner() + runner.pipeline = _PerRequestErrorStepPipeline() + bad_req = _make_step_request() + bad_req.request_id = "req-bad" + bad_req.prompt = "fail-prepare" + good_req = _make_step_request() + good_req.request_id = "req-good" + good_req.prompt = "valid" + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + + result = DiffusionModelRunner.execute_stepwise( + runner, + _make_batch_scheduler_output([bad_req, good_req]), + ) + + bad_output = result.get_request_output("req-bad") + assert bad_output.finished is True + assert bad_output.result is not None + assert bad_output.result.error == "invalid request input" + assert bad_output.result.error_status_code == 422 + assert bad_output.result.error_type == "UnprocessableEntityError" + good_output = result.get_request_output("req-good") + assert good_output.finished is False + assert good_output.step_index == 1 + assert good_output.result is None + assert set(runner.state_cache) == {"req-good"} + assert runner.pipeline.denoise_calls == 1 + assert runner.pipeline.scheduler_calls == 1 + + def test_prepare_error_on_other_rank_is_isolated_via_all_reduce(self, monkeypatch): + """P1-B: rank-0-only pipeline work (e.g. H3 reference-video prep) + can raise on rank 0 while non-zero ranks return cleanly and then hang + on a downstream ``dist.broadcast``. Beyond the pipeline-side guard, + the runner all-reduces a per-request failure flag across the DiT + group so a rank that ``prepare_encode`` returned cleanly on still + skips the request and does not enter the denoise step batch. + """ + runner = _make_runner() + runner.pipeline = _StepPipeline() # locally always succeeds + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + + # Simulate a peer DiT rank having reported failure. We patch the + # helper directly (rather than ``torch.distributed`` broadly) because + # ``execute_stepwise`` also queries ``torch.distributed.get_rank``. + monkeypatch.setattr( + model_runner_module, + "_dit_any_rank_failed", + lambda local_failed: True, + ) + + req = _make_step_request(num_inference_steps=1) + result = DiffusionModelRunner.execute_stepwise( + runner, + _make_scheduler_output(req), + ) + + output = result.get_request_output("req-1") + assert output.finished is True + assert output.result is not None + assert output.result.error is not None + assert "another DiT rank" in output.result.error + # A rank that skipped the request must not have driven a denoise step. + assert runner.pipeline.denoise_calls == 0 + assert runner.state_cache == {} + def test_receives_kv_payload_before_prepare_encode(self, monkeypatch): runner = _make_runner() captured: dict[str, object] = {} diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 95f0f2736..63b3b9cdb 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -36,6 +36,23 @@ class AttentionBackend(ABC): """ return False + @classmethod + def supports_multi_doc_packed_varlen(cls) -> bool: + """Whether this backend keeps N-document packed boundaries isolated. + + When True, the backend consumes ``AttentionMetadata.extra`` cu_seqlens + as a genuine block-diagonal attention plan (a dedicated varlen kernel, + not a padding-mask rebuild), so a caller may pack multiple real + requests into one forward without attention crossing document + boundaries. When False, callers packing more than one real document + must run one forward per document; otherwise a backend that only + supports a ``[real, pad]`` two-document contract, or that ignores + cu_seqlens outright, will silently attend across request boundaries. + Backends whose kernel selection is platform-dependent must consult + ``current_omni_platform``. + """ + return False + # ``OmniPlatformEnum`` values this backend runs on; None means unrestricted. # Platform resolution rejects an explicit selection outside this set, so a # hardware-specific backend fails before the model is built. diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index d287ef7d2..838fd3d6d 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -37,6 +37,20 @@ class FlashAttentionBackend(AttentionBackend): # constructing it there. return current_omni_platform.is_cuda() or current_omni_platform.is_npu() + @classmethod + def supports_multi_doc_packed_varlen(cls) -> bool: + # CUDA / ROCm / MUSA all route through ``forward_cuda``, which + # dispatches ``_forward_varlen_packed`` -> ``flash_attn_varlen_func`` + # over the caller's cu_seqlens without a mask, so an arbitrary + # N-document packing keeps its boundaries. NPU's forward path + # (``_resolve_packed_seq_npu``) only accepts a ``[real, pad]`` + # two-document layout (cu_seqlens shape <= 3) and otherwise falls + # back to a padding-mask rebuild that spans the whole packed row, + # so N>=2 real documents would silently attend across request + # boundaries. XPU never consumes cu_seqlens from ``extra``. Both + # therefore stay False here. + return current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_musa() + @classmethod def supports_attention_mask(cls) -> bool: return True diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index db4427958..5be8b53ab 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -1304,6 +1304,9 @@ def extract_minimax_h3_context( psp = _required_kwarg(kwargs, "packed_seq_params") cu_seqlens = module._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to(torch.int32) max_seqlen = int(module._psp_field(psp, "packed_seq_params", "max_seqlen_q")) + # Mirror forward()'s host-side scalar so the refiner sees the number of + # packed requests without reading cu_seqlens off-device. + num_requests = int(module._psp_optional(psp, "num_requests", 1)) refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params") refiner_cu = module._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32) refiner_max = int(module._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q")) @@ -1330,6 +1333,7 @@ def extract_minimax_h3_context( text_pos=text_pos.to(device), refiner_cu_seqlens=refiner_cu.to(device), refiner_max_seqlen=refiner_max, + num_requests=num_requests, seq_len=seq_len, device=device, local_span=(0, seq_len), diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index af578d662..f069cee85 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/vllm_omni/diffusion/forward_context.py @@ -244,6 +244,17 @@ def set_forward_context_denoise_step_idx(step_idx: int | None) -> None: _forward_context.denoise_step_idx = step_idx +def set_forward_context_denoise_timestep(timestep: float | None) -> None: + """Set the normalized (descending, 1 -> 0) denoise timestep. + + Timestep-gated attention features read this; pipelines that drive their own + denoise loop can publish it directly instead of going through + :meth:`DenoiseProgressMixin.record_denoise_step`. + """ + if _forward_context is not None: + _forward_context.denoise_timestep = None if timestep is None else float(timestep) + + class DenoiseProgressMixin: def record_denoise_step( self, diff --git a/vllm_omni/diffusion/models/minimax_h3/batched_packing.py b/vllm_omni/diffusion/models/minimax_h3/batched_packing.py new file mode 100644 index 000000000..e32b61667 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/batched_packing.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Batched multi-request packing for MiniMax H3 step-wise execution. + +Request mode forwards exactly one packed sequence per denoise step. Step mode +(continuous batching) may hold several requests at once, so this module +concatenates their packed layouts into a single sequence and rebuilds +``cu_seqlens`` with one document per request plus that request's +alignment-padding tail. Attention therefore never crosses a request boundary, +while every request keeps its own RoPE coordinates, token tags, and timesteps. +A batch of one is not rebuilt at all -- it delegates to the request-mode layout. + +Row order is the request order, so the concatenated video velocity returned by +the DiT can be split back per request by row count alone -- which is exactly +what the runner does with ``StepRequestState.latents``. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import torch + +from .denoise_loop import ( + MINIMAX_H3_AUDIO_ROW_WIDTH, + MINIMAX_H3_VIDEO_ROW_WIDTH, + MiniMaxH3DenoiseBranch, +) + + +def minimax_h3_batched_forward_kwargs( + *, + branches: Sequence[MiniMaxH3DenoiseBranch], + video_rows: Sequence[torch.Tensor], + audio_rows: Sequence[torch.Tensor], + t_video: Sequence[float], + t_audio: Sequence[float], + imgvid_cond_timesteps: Sequence[float], + audio_ref_cond_timesteps: Sequence[float], +) -> dict[str, Any]: + """Build one DiT forward kwargs dict covering every request in the batch. + + A single request delegates to :meth:`MiniMaxH3DenoiseBranch.forward_kwargs`, + so step mode and request mode are the same call rather than two layouts that + happen to agree -- which also keeps ``video_token_layout`` (the sparse- + attention hint) on the single-request step path. + """ + if len(branches) == 1: + return branches[0].forward_kwargs( + video_rows=video_rows[0], + audio_rows=audio_rows[0], + t_video=t_video[0], + t_audio=t_audio[0], + imgvid_cond_timestep=imgvid_cond_timesteps[0], + audio_ref_cond_timestep=audio_ref_cond_timesteps[0], + ) + + device = branches[0].device + total_seq = sum(branch.seq_len for branch in branches) + + img_pos_parts: list[torch.Tensor] = [] + audio_pos_parts: list[torch.Tensor] = [] + text_pos_parts: list[torch.Tensor] = [] + update_parts: list[torch.Tensor] = [] + tag_parts: list[torch.Tensor] = [] + position_id_parts: list[torch.Tensor] = [] + text_embed_parts: list[torch.Tensor] = [] + cu_bounds: list[int] = [0] + refiner_bounds: list[int] = [0] + max_seqlen = 0 + refiner_max_seqlen = 0 + + # Non-media rows (text and alignment padding) inherit their request's video + # timestep, matching the single-request packed-sequence semantics. + timesteps = torch.empty(total_seq, dtype=torch.float32, device=device) + + seq_offset = 0 + text_offset = 0 + for index, branch in enumerate(branches): + img_pos = branch.img_pos_dev + seq_offset + audio_pos = branch.audio_pos_dev + seq_offset + img_pos_parts.append(img_pos) + audio_pos_parts.append(audio_pos) + text_pos_parts.append(branch.text_pos_dev + seq_offset) + update_parts.append(branch.update_mask_dev) + tag_parts.append(branch.token_tags_dev) + position_id_parts.append(branch.img_position_ids_dev) + text_embed_parts.append(branch.text_embeddings_dev) + + timesteps[seq_offset : seq_offset + branch.seq_len] = float(t_video[index]) + timesteps[img_pos[branch.update_mask_dev]] = float(t_video[index]) + timesteps[img_pos[~branch.update_mask_dev]] = float(imgvid_cond_timesteps[index]) + timesteps[audio_pos[branch.audio_update_mask_dev]] = float(t_audio[index]) + timesteps[audio_pos[~branch.audio_update_mask_dev]] = float(audio_ref_cond_timesteps[index]) + + # One document for the request's real rows and one for its padding tail, + # matching the single-request layout. The tail bound is emitted even when + # the 64-row alignment left it empty, so that a request always + # contributes exactly two bounds regardless of its length. + cu_bounds.append(seq_offset + branch.used_len) + cu_bounds.append(seq_offset + branch.seq_len) + max_seqlen = max(max_seqlen, branch.used_len, branch.seq_len - branch.used_len) + + text_offset += branch.text_len + refiner_bounds.append(text_offset) + refiner_max_seqlen = max(refiner_max_seqlen, branch.text_len) + seq_offset += branch.seq_len + + img_pos_all = torch.cat(img_pos_parts) + audio_pos_all = torch.cat(audio_pos_parts) + + x = torch.zeros(1, total_seq, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device) + x[0].index_copy_(0, img_pos_all, torch.cat([rows.to(device=device, dtype=torch.float32) for rows in video_rows])) + audio_x = torch.zeros(1, total_seq, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device) + audio_x[0].index_copy_( + 0, + audio_pos_all, + torch.cat([rows.to(device=device, dtype=torch.float32) for rows in audio_rows]), + ) + + unique_timesteps, inverse_indices = torch.unique(timesteps, sorted=True, return_inverse=True) + return { + "x": x, + "audio_x": audio_x, + "img_position_ids": torch.cat(position_id_parts)[None], + "update_mask": torch.cat(update_parts), + "token_tags": torch.cat(tag_parts), + "skip_mask_out_condition": False, + "prompt_embeds": torch.cat(text_embed_parts), + "img_pos_info": {"position_ids": img_pos_all}, + "audio_pos_info": {"position_ids": audio_pos_all}, + "text_pos_info": {"position_ids": torch.cat(text_pos_parts)}, + "img_pos_for_infer_output_info": {"position_ids": img_pos_all}, + "packed_seq_params": { + "cu_seqlens_q": torch.tensor(cu_bounds, dtype=torch.int32, device=device), + "max_seqlen_q": max_seqlen, + # Attention needs the request count on the host to know its valid + # rows are block-diagonal rather than a prefix. It describes the + # whole forward, so the DiT reads it once here and applies it to the + # token refiner too -- the refiner's text-only bounds cannot be told + # apart from a single padded request on their own. + "num_requests": len(branches), + }, + "refiner_packed_seq_params": { + "cu_seqlens_q": torch.tensor(refiner_bounds, dtype=torch.int32, device=device), + "max_seqlen_q": refiner_max_seqlen, + }, + "unique_timesteps": unique_timesteps, + "inverse_indices": inverse_indices, + } + + +__all__ = [ + "minimax_h3_batched_forward_kwargs", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py index d22065c3a..f32a037c1 100644 --- a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py +++ b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py @@ -15,7 +15,10 @@ from typing import Any import torch from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout -from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx +from vllm_omni.diffusion.forward_context import ( + set_forward_context_denoise_step_idx, + set_forward_context_denoise_timestep, +) from .scheduling_minimax_h3_euler_ancestral import ( minimax_h3_euler_eta0_step, @@ -31,6 +34,17 @@ MINIMAX_H3_VIDEO_ROW_WIDTH = 96 MINIMAX_H3_AUDIO_ROW_WIDTH = 32 +def minimax_h3_publish_denoise_progress(step: int | None, sigma_video: float | None) -> None: + """Publish denoise progress for step-gated attention features. + + Both execution modes must publish the same pair: the step index drives the + dense warmup of RAINFUSION_ATTN, and the normalized descending timestep + drives the TRTLLM_ATTN skip gate, which stays dense while it is unset. + """ + set_forward_context_denoise_step_idx(step) + set_forward_context_denoise_timestep(sigma_video) + + class MiniMaxH3DenoiseBranch: """Static per-branch state: packed layout + fixed forward kwargs. @@ -64,26 +78,39 @@ class MiniMaxH3DenoiseBranch: if int(token_tags.view(-1).shape[0]) != seq_len: raise ValueError(f"token_tags length {int(token_tags.view(-1).shape[0])} != seq_len {seq_len}") cu = packed["cu_seqlens"].to(torch.int32) + self.device = device + # ``used_len`` is the real (non-padding) document length; step-mode + # batching concatenates layouts and needs both bounds as Python ints so + # it can rebuild ``cu_seqlens`` without a device sync per step. + self.used_len = int(cu[1]) + self.text_len = text_len self.img_pos_dev = self.img_pos.to(device) self.audio_pos_dev = self.audio_pos.to(device) self.update_mask_dev = self.update_mask.to(device) self.audio_update_mask_dev = self.audio_update_mask.to(device) self.x_base = torch.zeros(1, seq_len, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device) self.audio_x_base = torch.zeros(1, seq_len, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device) - text_pos_dev = packed["text_pos"].view(-1).to(torch.long).to(device) + self.text_pos_dev = packed["text_pos"].view(-1).to(torch.long).to(device) + self.token_tags_dev = token_tags.view(-1).to(torch.long).to(device) + self.img_position_ids_dev = packed["img_position_ids"].to(device) + self.text_embeddings_dev = text_embeddings.to(device) self.static_kwargs: dict[str, Any] = { - "img_position_ids": packed["img_position_ids"][None].to(device), + "img_position_ids": self.img_position_ids_dev[None], "update_mask": self.update_mask_dev, - "token_tags": token_tags.view(-1).to(torch.long).to(device), + "token_tags": self.token_tags_dev, "skip_mask_out_condition": False, - "prompt_embeds": text_embeddings.to(device), + "prompt_embeds": self.text_embeddings_dev, "img_pos_info": {"position_ids": self.img_pos_dev}, "audio_pos_info": {"position_ids": self.audio_pos_dev}, - "text_pos_info": {"position_ids": text_pos_dev}, + "text_pos_info": {"position_ids": self.text_pos_dev}, "img_pos_for_infer_output_info": {"position_ids": self.img_pos_dev}, "packed_seq_params": { "cu_seqlens_q": cu.to(device), - "max_seqlen_q": int(cu[1]), + "max_seqlen_q": self.used_len, + # One request: valid rows are a prefix, so attention may use a + # KV prefix length or a 1-D pad mask. See batched_packing for the + # co-batched layout, where neither can describe the valid rows. + "num_requests": 1, }, "refiner_packed_seq_params": { "cu_seqlens_q": torch.tensor([0, text_len, text_len], dtype=torch.int32, device=device), @@ -136,34 +163,21 @@ class MiniMaxH3DenoiseBranch: } -def minimax_h3_denoise_loop( +def minimax_h3_prepare_denoise_rows( *, - model: Any, positive: MiniMaxH3DenoiseBranch, initial_video_rows: torch.Tensor, initial_audio_rows: torch.Tensor, keyframe_cond_rows: torch.Tensor | None, - audio_ref_rows: torch.Tensor | None = None, - sigmas_video: list[float], - sigmas_audio: list[float], + audio_ref_rows: torch.Tensor | None, device: torch.device, - imgvid_cond_noise_aug_for_inference: float = MINIMAX_H3_IMGVID_COND_TIMESTEP, - audio_cond_noise_aug_for_inference: float = MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, - on_step_end: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None, - on_step_start: Callable[[int, float, float], None] | None = None, - step_profiler: Callable[[int], AbstractContextManager] | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Run the full denoise loop; returns final (video_rows, audio_rows). +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Validate the initial rows against the layout and pin the condition rows. - ``initial_video_rows`` covers all image rows of the positive layout. For a - conditional task, pass ``keyframe_cond_rows`` and/or ``audio_ref_rows`` to - pin those rows across every step. The model's raw positive velocity is the - update signal; MiniMax H3 only supports cfg-distilled checkpoints. + Returns ``(video_rows, audio_rows, cond_anchor, audio_anchor)`` with the + anchors already written into their rows, which is the state both the + request-mode loop and step-mode ``prepare_encode()`` start from. """ - if len(sigmas_video) != len(sigmas_audio): - raise ValueError("video/audio sigma schedules must have equal length") - if len(sigmas_video) < 2: - raise ValueError("sigma schedules need at least 2 entries") n_cond = int((~positive.update_mask).sum()) if keyframe_cond_rows is None: if n_cond != 0: @@ -197,19 +211,58 @@ def minimax_h3_denoise_loop( video_rows[~update] = cond_anchor if audio_anchor is not None: audio_rows[~audio_update] = audio_anchor + return video_rows, audio_rows, cond_anchor, audio_anchor + + +def minimax_h3_denoise_loop( + *, + model: Any, + positive: MiniMaxH3DenoiseBranch, + initial_video_rows: torch.Tensor, + initial_audio_rows: torch.Tensor, + keyframe_cond_rows: torch.Tensor | None, + audio_ref_rows: torch.Tensor | None = None, + sigmas_video: list[float], + sigmas_audio: list[float], + device: torch.device, + imgvid_cond_noise_aug_for_inference: float = MINIMAX_H3_IMGVID_COND_TIMESTEP, + audio_cond_noise_aug_for_inference: float = MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, + on_step: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None, + step_profiler: Callable[[int], AbstractContextManager] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the full denoise loop; returns final (video_rows, audio_rows). + + ``initial_video_rows`` covers all image rows of the positive layout. For a + conditional task, pass ``keyframe_cond_rows`` and/or ``audio_ref_rows`` to + pin those rows across every step. The model's raw positive velocity is the + update signal; MiniMax H3 only supports cfg-distilled checkpoints. + """ + if len(sigmas_video) != len(sigmas_audio): + raise ValueError("video/audio sigma schedules must have equal length") + if len(sigmas_video) < 2: + raise ValueError("sigma schedules need at least 2 entries") + video_rows, audio_rows, cond_anchor, audio_anchor = minimax_h3_prepare_denoise_rows( + positive=positive, + initial_video_rows=initial_video_rows, + initial_audio_rows=initial_audio_rows, + keyframe_cond_rows=keyframe_cond_rows, + audio_ref_rows=audio_ref_rows, + device=device, + ) + update = positive.update_mask_dev + audio_update = positive.audio_update_mask_dev num_steps = len(sigmas_video) - 1 for step in range(num_steps): step_cm = step_profiler(step) if step_profiler is not None else nullcontext() with step_cm: - # Publish the step index so step-gated attention features (e.g. the - # dense warmup of RAINFUSION_ATTN) can see where we are. - set_forward_context_denoise_step_idx(step) s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1] s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1] - if on_step_start is not None: - # Attention gates use the scheduler-style descending timestep. - on_step_start(step, s_v, s_a) + # Publish where we are so step-gated attention features (the dense + # warmup of RAINFUSION_ATTN, the timestep gate of TRTLLM_ATTN) can + # see it. Gates use the scheduler-style descending timestep, which + # for a rectified-flow schedule is the video sigma. + minimax_h3_publish_denoise_progress(step, s_v) t_v, t_a = 1.0 - s_v, 1.0 - s_a imgvid_cond_t = max(t_v, float(imgvid_cond_noise_aug_for_inference)) audio_ref_cond_t = max(t_a, float(audio_cond_noise_aug_for_inference)) @@ -250,10 +303,10 @@ def minimax_h3_denoise_loop( audio_rows[audio_update] = new_audio if audio_anchor is not None: audio_rows[~audio_update] = audio_anchor # per-step audio ref reset - if on_step_end is not None: - on_step_end(step, video_rows, audio_rows) + if on_step is not None: + on_step(step, video_rows, audio_rows) - set_forward_context_denoise_step_idx(None) + minimax_h3_publish_denoise_progress(None, None) return video_rows, audio_rows @@ -264,4 +317,6 @@ __all__ = [ "MINIMAX_H3_VIDEO_ROW_WIDTH", "MiniMaxH3DenoiseBranch", "minimax_h3_denoise_loop", + "minimax_h3_prepare_denoise_rows", + "minimax_h3_publish_denoise_progress", ] diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index 1b852969a..509190740 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -56,6 +56,26 @@ if TYPE_CHECKING: logger = init_logger(__name__) +# Packed multi-request forwards require the attention backend to actually +# consume cu_seqlens as a block-diagonal plan (not a padding-mask rebuild that +# spans the full packed row). The pipeline gates on this capability before +# packing, and ``_run_packed_attention`` re-checks it per forward; a name-only +# gate would let FLASH_ATTN's NPU/XPU code paths through even though those +# variants would silently attend across request boundaries. +def _attention_isolates_packed_requests(attention_layer: Any) -> bool: + """True if this attention layer keeps N-document packed boundaries. + + Requires a backend advertising ``supports_multi_doc_packed_varlen`` *and* + that the layer is not running under ring sequence parallelism (the ring + kernel dispatches through its own attention that ignores the packed + cu_seqlens regardless of the configured backend). + """ + backend = getattr(attention_layer, "attn_backend", None) + if backend is None or not backend.supports_multi_doc_packed_varlen(): + return False + return not getattr(attention_layer, "use_ring", False) + + @dataclass class MiniMaxH3DiTArchConfig: num_layers: int = 50 @@ -433,6 +453,7 @@ class MiniMaxH3Attention(nn.Module): cu_seqlens: torch.Tensor, max_seqlen: int, packed_total: int, + num_requests: int = 1, video_layout: VideoTokenLayout | None = None, ) -> torch.Tensor: """Run packed attention as a small eager island. @@ -442,27 +463,47 @@ class MiniMaxH3Attention(nn.Module): narrow lets regional compile fuse projections, norms, RoPE, and the surrounding DiT block without repeated graph breaks. """ - # max_seqlen is already the first (real) packed document length. Do - # not read the CUDA cu_seqlens scalars here: this function runs once - # per layer and .item() would serialize every attention launch. + # max_seqlen is already the longest packed document length. Do not read + # the CUDA cu_seqlens scalars here: this function runs once per layer + # and .item() would serialize every attention launch. ``num_requests`` + # is carried as a Python int for the same reason. if not 0 < max_seqlen <= packed_total: raise ValueError( f"max_seqlen must be within the packed sequence, got {max_seqlen} for length {packed_total}" ) - used = min(max_seqlen, packed_total) attn_mask = None - # Ring attention can dispatch to a different implementation from the - # configured backend, so the no-mask fast paths are local-only. - # supports_prefix_kv_slicing: backend slices K/V itself (cuDNN). - # supports_packed_mask_free: backend consumes the packed metadata - # without ever reading attn_mask (CUDA packed varlen, NPU - # npu_attn_varlen opt-in with its own fallback rebuild). - no_mask = not getattr(self.attention, "use_ring", False) and ( - self.attention.attn_backend.supports_prefix_kv_slicing - or self.attention.attn_backend.supports_packed_mask_free() - ) - if used < packed_total and not no_mask: - attn_mask = torch.arange(packed_total, device=q.device)[None] < used + if num_requests > 1: + # A step-mode batch packs one document per request, so its valid + # rows are block-diagonal rather than a prefix: neither a KV prefix + # length nor a 1-D key mask can describe them. Such a layout is + # only correct on a backend that actually attends by cu_seqlens as + # a block-diagonal plan. Check the capability (not the backend + # name): FLASH_ATTN's NPU/XPU variants would otherwise silently + # fall back to a padding-mask rebuild that spans the whole packed + # row and attend across request boundaries. + if not _attention_isolates_packed_requests(self.attention): + backend_name = self.attention.attn_backend.get_name() + raise ValueError( + f"MiniMax H3 packed a {num_requests}-request batch, but the resolved " + f"attention ({backend_name}, use_ring={getattr(self.attention, 'use_ring', False)}) " + "does not isolate multi-document packed cu_seqlens. Run one request " + "per forward on this backend." + ) + used = packed_total + else: + used = min(max_seqlen, packed_total) + # Ring attention can dispatch to a different implementation from the + # configured backend, so the no-mask fast paths are local-only. + # supports_prefix_kv_slicing: backend slices K/V itself (cuDNN). + # supports_packed_mask_free: backend consumes the packed metadata + # without ever reading attn_mask (CUDA packed varlen, NPU + # npu_attn_varlen opt-in with its own fallback rebuild). + no_mask = not getattr(self.attention, "use_ring", False) and ( + self.attention.attn_backend.supports_prefix_kv_slicing + or self.attention.attn_backend.supports_packed_mask_free() + ) + if used < packed_total and not no_mask: + attn_mask = torch.arange(packed_total, device=q.device)[None] < used metadata = AttentionMetadata( attn_mask=attn_mask, extra={ @@ -498,6 +539,7 @@ class MiniMaxH3Attention(nn.Module): cu_seqlens: torch.Tensor, max_seqlen: int, packed_total: int | None = None, + num_requests: int = 1, sp_seq_lens: list[int] | None = None, video_layout: VideoTokenLayout | None = None, ) -> torch.Tensor: @@ -533,9 +575,9 @@ class MiniMaxH3Attention(nn.Module): self.q_norm.variance_epsilon, ) - # The packed layout uses a second document for alignment padding. - # Local/Ulysses backends unpad it, while Ring keeps aligned rows for - # fixed-size P2P buffers. + # Each request contributes a document for its rows plus one for its + # alignment padding. Local/Ulysses backends unpad it, while Ring keeps + # aligned rows for fixed-size P2P buffers. out = self._run_packed_attention( q, k, @@ -546,6 +588,7 @@ class MiniMaxH3Attention(nn.Module): # backend receives the global sequence after all-to-all, so carry # its Python length explicitly instead of inferring it from q. packed_total=packed_total if packed_total is not None else q.shape[0], + num_requests=num_requests, video_layout=video_layout, ) out = out.reshape(total, self.num_heads * self.head_dim) @@ -673,12 +716,14 @@ class MiniMaxH3TokenRefinerBlock(nn.Module): *, cu_seqlens: torch.Tensor, max_seqlen: int, + num_requests: int = 1, ) -> torch.Tensor: x = x + self.attn( self.norm1(x), rope_table=None, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, + num_requests=num_requests, ) x = x + self.mlp(self.norm2(x)) return x @@ -711,9 +756,10 @@ class MiniMaxH3TokenRefiner(nn.Module): *, cu_seqlens: torch.Tensor, max_seqlen: int, + num_requests: int = 1, ) -> torch.Tensor: for block in self.blocks: - x = block(x, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen) + x = block(x, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, num_requests=num_requests) return self.final_norm(x) @@ -759,6 +805,7 @@ class MiniMaxH3DiTBlock(nn.Module): cu_seqlens: torch.Tensor, max_seqlen: int, packed_total: int, + num_requests: int = 1, sp_seq_lens: list[int] | None = None, video_layout: VideoTokenLayout | None = None, ) -> torch.Tensor: @@ -793,6 +840,7 @@ class MiniMaxH3DiTBlock(nn.Module): cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, packed_total=packed_total, + num_requests=num_requests, sp_seq_lens=sp_seq_lens, video_layout=video_layout, ) @@ -1145,6 +1193,11 @@ class MiniMaxH3DiTModel(nn.Module): raise ValueError(f"{key}.{field} is required") return value + @staticmethod + def _psp_optional(psp: Any, field: str, default: Any) -> Any: + value = psp.get(field) if isinstance(psp, dict) else getattr(psp, field, None) + return default if value is None else value + def _embed( self, *, @@ -1160,10 +1213,16 @@ class MiniMaxH3DiTModel(nn.Module): seq_len: int, device: torch.device, local_span: tuple[int, int], + num_requests: int = 1, ) -> tuple[torch.Tensor, torch.Tensor]: """Build this rank's packed multimodal embedding rows. Returns (decoder_input [S_local, H] bf16, t_emb [M, t_dim] fp32). + + ``num_requests`` defaults to a single packed request so callers that + pre-date the continuous-batching change (e.g. TeaCache's extractor + contract) do not silently miss the kwarg. ``forward()`` reads the real + value from ``packed_seq_params["num_requests"]``. """ local_start, local_len = local_span local_end = local_start + local_len @@ -1200,6 +1259,7 @@ class MiniMaxH3DiTModel(nn.Module): text_embed, cu_seqlens=refiner_cu_seqlens, max_seqlen=refiner_max_seqlen, + num_requests=num_requests, ) if text_local_indices is not None: text_embed = text_embed.index_select(0, text_local_indices) @@ -1269,6 +1329,10 @@ class MiniMaxH3DiTModel(nn.Module): psp = _required_kwarg(kwargs, "packed_seq_params") cu_seqlens = self._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to(torch.int32) max_seqlen = int(self._psp_field(psp, "packed_seq_params", "max_seqlen_q")) + # How many requests share this packed sequence. Carried as a host int so + # attention never reads cu_seqlens scalars off the device; a producer + # that omits it is packing a single request. + num_requests = int(self._psp_optional(psp, "num_requests", 1)) refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params") refiner_cu = self._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32) refiner_max = int(self._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q")) @@ -1305,6 +1369,7 @@ class MiniMaxH3DiTModel(nn.Module): text_pos=text_pos.to(device), refiner_cu_seqlens=refiner_cu.to(device), refiner_max_seqlen=refiner_max, + num_requests=num_requests, seq_len=seq_len, device=device, local_span=local_span, @@ -1339,6 +1404,7 @@ class MiniMaxH3DiTModel(nn.Module): cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, packed_total=seq_len, + num_requests=num_requests, video_layout=video_layout, ) if local_len == seq_len: 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 e6e1a03fb..115cdc79d 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -57,7 +57,7 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( ) from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch -from vllm_omni.errors import OmniClientError +from vllm_omni.errors import OmniClientError, client_error_from_metadata from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) @@ -69,14 +69,24 @@ from vllm_omni.quantization.component_config import ( resolve_encoder_quant_config as _resolve_encoder_quant_config, ) +from .batched_packing import minimax_h3_batched_forward_kwargs from .condition_noise import ( minimax_h3_audio_cond_noise_aug_rows, minimax_h3_imgvid_cond_noise_aug_rows, ) -from .denoise_loop import MiniMaxH3DenoiseBranch, minimax_h3_denoise_loop +from .denoise_loop import ( + MiniMaxH3DenoiseBranch, + minimax_h3_denoise_loop, + minimax_h3_prepare_denoise_rows, + minimax_h3_publish_denoise_progress, +) from .encoder import MiniMaxH3Qwen3VLEncoder from .lora import load_minimax_h3_turbo_lora -from .minimax_h3_transformer import MiniMaxH3DiTModel +from .minimax_h3_transformer import ( + MiniMaxH3Attention, + MiniMaxH3DiTModel, + _attention_isolates_packed_requests, +) from .packed_sequence import ( minimax_h3_packed_sequence, minimax_h3_packed_sequence_ref2va_blocks, @@ -103,6 +113,10 @@ from .reference_video import ( validate_reference_audio_files, validate_reference_audio_waveforms, ) +from .scheduling_minimax_h3_euler_ancestral import ( + minimax_h3_euler_eta0_step, + minimax_h3_rf_v_to_x0, +) from .time_request import ( MINIMAX_H3_SHAPE_PLANNER, minimax_h3_align_frame_count, @@ -110,6 +124,10 @@ from .time_request import ( ) from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE +if TYPE_CHECKING: + from vllm_omni.diffusion.worker.input_batch import InputBatch + from vllm_omni.diffusion.worker.utils import StepRequestState + logger = init_logger(__name__) if TYPE_CHECKING: @@ -201,6 +219,68 @@ def _resolve_minimax_h3_model_root( ) +# Keys of ``_prepare_request_inputs`` that feed ``diffuse`` / ``_build_denoise_inputs``. +_MINIMAX_H3_DENOISE_INPUT_KEYS = ( + "task", + "text_embeddings", + "text_tags", + "seed", + "latent_t", + "latent_h", + "latent_w", + "audio_t", + "num_frames", + "num_steps", + "video_shift", + "audio_shift", + "base_schedule", + "visual_condition", + "visual_condition_shape", + "audio_condition", + "ref_audio_t", + "ref_blocks", + "visual_condition_shapes", + "audio_condition_lengths", + "keyframe_frame_indices", +) + +# ``StepRequestState.extra`` keys owned by the step-execution path. +_STEP_BRANCH = "minimax_h3_branch" +_STEP_AUDIO_ROWS = "minimax_h3_audio_rows" +_STEP_AUDIO_NOISE_PRED = "minimax_h3_audio_noise_pred" +_STEP_SIGMAS_VIDEO = "minimax_h3_sigmas_video" +_STEP_SIGMAS_AUDIO = "minimax_h3_sigmas_audio" +_STEP_COND_ANCHOR = "minimax_h3_cond_anchor" +_STEP_AUDIO_ANCHOR = "minimax_h3_audio_anchor" +_STEP_SHAPE = "minimax_h3_shape" +_STEP_TRANSFORMER = "minimax_h3_transformer" + + +def _minimax_h3_step_schedule(state: StepRequestState) -> dict[str, float]: + """Return the sigma/timestep values this request needs for its current step. + + Mirrors the per-iteration arithmetic of ``minimax_h3_denoise_loop`` so step + mode and request mode advance identically. + """ + step = int(state.step_index) + sigmas_video = state.extra[_STEP_SIGMAS_VIDEO] + sigmas_audio = state.extra[_STEP_SIGMAS_AUDIO] + sigma_video = float(sigmas_video[step]) + sigma_audio = float(sigmas_audio[step]) + t_video = 1.0 - sigma_video + t_audio = 1.0 - sigma_audio + return { + "sigma_video": sigma_video, + "sigma_video_next": float(sigmas_video[step + 1]), + "sigma_audio": sigma_audio, + "sigma_audio_next": float(sigmas_audio[step + 1]), + "t_video": t_video, + "t_audio": t_audio, + "imgvid_cond_timestep": max(t_video, MINIMAX_H3_IMGVID_COND_TIMESTEP), + "audio_ref_cond_timestep": max(t_audio, MINIMAX_H3_AUDIO_REF_COND_TIMESTEP), + } + + def _read_base_schedule(release: Mapping[str, Any]) -> DMD2SigmaSchedule | None: """Read a partition's distilled schedule. An absent key means legacy uniform.""" return DMD2SigmaSchedule.from_metadata(release) @@ -444,6 +524,58 @@ def _dit_rank_world() -> tuple[Any, int, int]: return group, dist.get_rank(group), dist.get_world_size(group) +def _broadcast_rank0_exception(exc: Exception | None) -> None: + """Synchronize a rank-0-only exception across every DiT rank. + + H3 reference-video preparation runs only on rank 0; the other DiT ranks + return ``None`` without touching disk. When rank 0 raises inside that + path it exits :meth:`prepare_encode` before reaching the downstream + ``dist.broadcast`` calls, and non-zero ranks then hang on those + collectives forever. Every rank calls this helper right after the + rank-0-only work, before any subsequent collective, so all ranks either + raise the same error together or all continue. + """ + group, rank, world_size = _dit_rank_world() + if world_size == 1: + if exc is not None: + raise exc + return + if rank == 0: + if exc is None: + payload: list[Any] = [None] + else: + payload = [ + { + "type": type(exc).__name__, + "message": str(exc), + "status_code": getattr(exc, "status_code", None), + "error_type": getattr(exc, "error_type", None), + } + ] + else: + payload = [None] + dist.broadcast_object_list(payload, src=0, group=group) + info = payload[0] + if info is None: + return + if rank == 0: + assert exc is not None + raise exc + # Rebuild a matching client-facing error on non-zero ranks so the runner's + # per-request try/except records the same 4xx status as rank 0. The exact + # subclass need not survive the wire; the message and status suffice. + status_code = info.get("status_code") + error_type = info.get("error_type") + message = f"[rank 0] {info['type']}: {info['message']}" + if status_code is not None: + raise client_error_from_metadata( + message, + status_code=int(status_code), + error_type=error_type, + ) + raise RuntimeError(message) + + def _broadcast_tensor( tensor: torch.Tensor | None, *, @@ -548,6 +680,8 @@ class MiniMaxH3Pipeline( ): """CFG-distilled joint video/audio generation for MiniMax H3.""" + supports_step_execution: ClassVar[bool] = True + _dit_modules: ClassVar[list[str]] = ["transformer", "transformers_ref"] _encoder_modules: ClassVar[list[str]] = ["text_encoder"] _vae_modules: ClassVar[list[str]] = ["video_vae", "audio_vae"] @@ -564,6 +698,9 @@ class MiniMaxH3Pipeline( "_encode_reference_audio_conditions", "diffuse", "decode", + "prepare_encode", + "denoise_step", + "post_decode", ] dummy_run_num_frames: ClassVar[int] = 0 # Only distilled releases pin a schedule, so the default keeps the legacy @@ -1538,7 +1675,7 @@ class MiniMaxH3Pipeline( if controller is not None and enabled: controller.offload_resident_layers() - def diffuse( + def _build_denoise_inputs( self, *, task: str, @@ -1562,7 +1699,12 @@ class MiniMaxH3Pipeline( visual_condition_shapes: list[tuple[int, int, int]] | None = None, audio_condition_lengths: list[int] | None = None, keyframe_frame_indices: list[int] | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> dict[str, Any]: + """Build the packed layout, initial rows, anchors, and sigma schedules. + + Shared by request-mode :meth:`diffuse` and step-mode + :meth:`prepare_encode` so both paths start from identical state. + """ initial_video, initial_audio = self._initial_noise( seed=seed, latent_t=latent_t, @@ -1662,32 +1804,34 @@ class MiniMaxH3Pipeline( shift_scale=audio_shift, base_schedule=base_schedule, ) - transformer = self._transformer_for_task(task) - # The static DLO plan keeps leading blocks resident only for the - # primary ``transformer``. In combined mode ``transformers_ref`` is - # fully streamed, so a Ref2VA request must not stage the inactive - # FL2VA transformer resident blocks. - with self._resident_dit_layers_on_device(enabled=transformer is self.transformer): - with self.progress_bar(total=len(video_sigmas) - 1) as progress: - video_rows, audio_rows = minimax_h3_denoise_loop( - model=transformer, - positive=branch, - initial_video_rows=initial_video, - initial_audio_rows=initial_audio, - keyframe_cond_rows=visual_anchor, - audio_ref_rows=audio_anchor, - sigmas_video=video_sigmas, - sigmas_audio=audio_sigmas, - device=self.device, - imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP), - audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP), - on_step_start=lambda step, video_sigma, audio_sigma: self.record_denoise_step( - step, - normalized_timestep=video_sigma, - ), - on_step_end=lambda step, video, audio: progress.update(), - ) + return { + "branch": branch, + # The request-mode loop moves these onto the device itself; step mode + # keeps them resident across steps, so normalize once for both. + "video_rows": initial_video.to(device=self.device, dtype=torch.float32), + "audio_rows": initial_audio.to(device=self.device, dtype=torch.float32), + "cond_anchor": ( + None if visual_anchor is None else visual_anchor.to(device=self.device, dtype=torch.float32) + ), + "audio_anchor": ( + None if audio_anchor is None else audio_anchor.to(device=self.device, dtype=torch.float32) + ), + "sigmas_video": video_sigmas, + "sigmas_audio": audio_sigmas, + } + def _unpack_denoised_rows( + self, + branch: MiniMaxH3DenoiseBranch, + video_rows: torch.Tensor, + audio_rows: torch.Tensor, + *, + latent_t: int, + latent_h: int, + latent_w: int, + audio_t: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Select the target rows and unpack them back into VAE latents.""" target_video = video_rows[branch.update_mask_dev] video_latent = minimax_h3_unpatchify_video_tokens( target_video, @@ -1707,6 +1851,83 @@ class MiniMaxH3Pipeline( ) return video_latent, audio_latent + def diffuse( + self, + *, + task: str, + text_embeddings: torch.Tensor, + text_tags: torch.Tensor, + seed: int, + latent_t: int, + latent_h: int, + latent_w: int, + audio_t: int, + num_frames: int, + num_steps: int, + video_shift: float, + audio_shift: float, + base_schedule: Sequence[float] | None, + visual_condition: torch.Tensor | None, + visual_condition_shape: tuple[int, int, int] | None, + audio_condition: torch.Tensor | None, + ref_audio_t: int | None, + ref_blocks: list[dict[str, Any]] | None = None, + visual_condition_shapes: list[tuple[int, int, int]] | None = None, + audio_condition_lengths: list[int] | None = None, + keyframe_frame_indices: list[int] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + inputs = self._build_denoise_inputs( + task=task, + text_embeddings=text_embeddings, + text_tags=text_tags, + seed=seed, + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + audio_t=audio_t, + num_frames=num_frames, + num_steps=num_steps, + video_shift=video_shift, + audio_shift=audio_shift, + base_schedule=base_schedule, + visual_condition=visual_condition, + visual_condition_shape=visual_condition_shape, + audio_condition=audio_condition, + ref_audio_t=ref_audio_t, + ref_blocks=ref_blocks, + visual_condition_shapes=visual_condition_shapes, + audio_condition_lengths=audio_condition_lengths, + keyframe_frame_indices=keyframe_frame_indices, + ) + branch = inputs["branch"] + transformer = self._transformer_for_task(task) + with self._resident_dit_layers_on_device(enabled=transformer is self.transformer): + with self.progress_bar(total=len(inputs["sigmas_video"]) - 1) as progress: + video_rows, audio_rows = minimax_h3_denoise_loop( + model=transformer, + positive=branch, + initial_video_rows=inputs["video_rows"], + initial_audio_rows=inputs["audio_rows"], + keyframe_cond_rows=inputs["cond_anchor"], + audio_ref_rows=inputs["audio_anchor"], + sigmas_video=inputs["sigmas_video"], + sigmas_audio=inputs["sigmas_audio"], + device=self.device, + imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP), + audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP), + on_step=lambda step, video, audio: progress.update(), + ) + + return self._unpack_denoised_rows( + branch, + video_rows, + audio_rows, + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + audio_t=audio_t, + ) + def decode( self, video_latent: torch.Tensor, @@ -1727,11 +1948,9 @@ class MiniMaxH3Pipeline( audio = self.audio_vae.decode_latent(audio_latent) return video, audio - @torch.no_grad() - def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput: - if len(request.prompts) != 1: - raise OmniClientError("MiniMax H3 supports one request at a time") - raw_prompt = request.prompts[0] + @staticmethod + def _extract_prompt(raw_prompt: Any) -> tuple[str, dict[str, Any]]: + """Split a request prompt into its text and multimodal parts.""" if isinstance(raw_prompt, str): prompt = raw_prompt multi_modal_data: dict[str, Any] = {} @@ -1740,8 +1959,21 @@ class MiniMaxH3Pipeline( multi_modal_data = raw_prompt.get("multi_modal_data") or {} if not prompt: raise OmniClientError("MiniMax H3 requires a non-empty prompt") + return prompt, multi_modal_data - sampling = request.sampling_params + def _prepare_request_inputs( + self, + *, + prompt: str, + multi_modal_data: dict[str, Any], + sampling: Any, + ) -> dict[str, Any]: + """Resolve the task and output shape, then run every request-level encode. + + Shared by request-mode :meth:`forward` and step-mode + :meth:`prepare_encode`; the returned mapping feeds :meth:`diffuse` and + :meth:`_build_denoise_inputs` unchanged. + """ quality = sampling.quality logger.debug("MiniMax H3 request quality=%s", quality) extra = sampling.extra_args or {} @@ -1818,12 +2050,24 @@ class MiniMaxH3Pipeline( video_count = 0 if raw_videos is not None: video_count = len(raw_videos) if isinstance(raw_videos, (list, tuple)) else 1 - prepared_videos = self._prepare_reference_videos( - raw_videos, - target_frame_count=num_frames, - workdir=workdir, - start_time_seconds=extra.get("start_time_seconds"), - ) + # File-based reference-video prep runs only on rank 0; other + # ranks return None without touching disk. If rank 0 raises + # (e.g. invalid file, unsupported codec) it must not exit + # ``prepare_encode`` before the downstream broadcasts below -- + # non-zero ranks would then deadlock on them forever. Capture + # the exception here and let every rank agree on the outcome + # before starting any subsequent collective. + prep_error: Exception | None = None + try: + prepared_videos = self._prepare_reference_videos( + raw_videos, + target_frame_count=num_frames, + workdir=workdir, + start_time_seconds=extra.get("start_time_seconds"), + ) + except Exception as exc: + prep_error = exc + _broadcast_rank0_exception(prep_error) has_audio_tensor = torch.zeros( video_count, dtype=torch.long, @@ -1960,33 +2204,60 @@ class MiniMaxH3Pipeline( ) self._cache_dit_runtime.prepare(quality_plan.cache_dit) num_outputs = _resolve_minimax_h3_num_outputs(sampling.num_outputs_per_prompt) + return { + "task": task, + "height": height, + "width": width, + "num_frames": num_frames, + "latent_t": latent_t, + "latent_h": height // 16, + "latent_w": width // 16, + "audio_t": audio_t, + "text_embeddings": text_embeddings, + "text_tags": text_tags, + "visual_condition": visual_condition, + "visual_condition_shape": visual_shape, + "audio_condition": audio_condition, + "ref_audio_t": ref_audio_t, + "ref_blocks": ref_blocks, + "visual_condition_shapes": visual_shapes, + "audio_condition_lengths": audio_lengths, + "keyframe_frame_indices": keyframe_frame_indices, + "seed": seed, + "num_steps": num_steps, + "video_shift": video_shift, + "audio_shift": audio_shift, + "base_schedule": base_schedule, + "num_outputs": num_outputs, + } + + @staticmethod + def _denoise_kwargs(context: dict[str, Any]) -> dict[str, Any]: + """Select the denoise-input arguments from a prepared request context.""" + return {key: context[key] for key in _MINIMAX_H3_DENOISE_INPUT_KEYS} + + @torch.no_grad() + def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput: + if len(request.prompts) != 1: + raise OmniClientError("MiniMax H3 supports one request at a time") + prompt, multi_modal_data = self._extract_prompt(request.prompts[0]) + context = self._prepare_request_inputs( + prompt=prompt, + multi_modal_data=multi_modal_data, + sampling=request.sampling_params, + ) + denoise_kwargs = self._denoise_kwargs(context) + num_outputs = context["num_outputs"] videos = [] audios = [] - for output_seed in _minimax_h3_output_seeds(seed, num_outputs): - video_latent, audio_latent = self.diffuse( - task=task, - text_embeddings=text_embeddings, - text_tags=text_tags, - seed=output_seed, - latent_t=latent_t, - latent_h=height // 16, - latent_w=width // 16, - audio_t=audio_t, - num_frames=num_frames, - num_steps=num_steps, - video_shift=video_shift, - audio_shift=audio_shift, - base_schedule=base_schedule, - visual_condition=visual_condition, - visual_condition_shape=visual_shape, - audio_condition=audio_condition, - ref_audio_t=ref_audio_t, - ref_blocks=ref_blocks, - visual_condition_shapes=visual_shapes, - audio_condition_lengths=audio_lengths, - keyframe_frame_indices=keyframe_frame_indices, + for output_seed in _minimax_h3_output_seeds(context["seed"], num_outputs): + video_latent, audio_latent = self.diffuse(**{**denoise_kwargs, "seed": output_seed}) + video, audio = self.decode( + video_latent, + audio_latent, + height=context["height"], + width=context["width"], ) - video, audio = self.decode(video_latent, audio_latent, height=height, width=width) videos.append(video) audios.append(audio) video = torch.cat(videos, dim=0) @@ -1997,6 +2268,302 @@ class MiniMaxH3Pipeline( stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}), ) + # ------------------------------------------------------------------ + # Step-wise execution (continuous batching) + # ------------------------------------------------------------------ + + @staticmethod + def _packed_batch_supported(transformer: MiniMaxH3DiTModel) -> bool: + """Whether every attention in this DiT honors multi-document cu_seqlens. + + A packed batch is only isolated if *all* of them do: the token refiner + runs under its own attention role and can resolve to a different backend + from the DiT blocks. Ring sequence parallelism dispatches through + ``RingParallelAttention``, whose kernels ignore the packed + ``cu_seqlens`` metadata regardless of the configured backend; packing + multiple requests under ring would let attention cross document + boundaries, so any layer running ring disqualifies the batch. + + The gate probes a per-backend capability rather than a fixed backend + name: FLASH_ATTN, for example, only isolates arbitrary N-document + packed cu_seqlens on CUDA/ROCm/MUSA. Its NPU path only accepts a + ``[real, pad]`` two-document layout and its XPU path ignores + cu_seqlens outright — either would silently attend across request + boundaries. + """ + attentions = [module for module in transformer.modules() if isinstance(module, MiniMaxH3Attention)] + if not attentions: + return False + return all(_attention_isolates_packed_requests(module.attention) for module in attentions) + + def prepare_encode(self, state: StepRequestState, **kwargs: Any) -> StepRequestState: + """Run every request-level stage once and seed the per-request step state.""" + del kwargs + # Two request-mode features have no place in the shared step contract: + # a request state carries exactly one latent tensor, and distributed + # layerwise offload streams the DiT around one whole denoise loop rather + # than around a single scheduler-driven step. + num_outputs = _resolve_minimax_h3_num_outputs(state.sampling.num_outputs_per_prompt) + if num_outputs != 1: + raise OmniClientError( + f"MiniMax H3 step execution produces one output per request, got num_outputs_per_prompt={num_outputs}" + ) + if getattr(self, "_dlo_residency_controller", None) is not None: + raise ValueError( + "MiniMax H3 step execution is not compatible with distributed layerwise offload; " + "the resident-layer window spans a whole denoise loop, so per-step streaming would " + "reload the DiT every step. Drop --step-execution or --enable-distributed-layerwise-offload." + ) + # Request-scoped Cache-DiT (quality=high) mutates hook state on the + # shared transformer rather than on ``StepRequestState``. In step mode + # two requests can interleave denoise steps, or be co-batched into a + # single forward, and the second one would then re-enter the DiT with + # cache buffers shaped for the first. Reject the profile here rather + # than let it corrupt outputs at runtime; startup-configured Cache-DiT + # is already blocked in ``DiffusionModelRunner.execute_stepwise``. + if getattr(state.sampling, "quality", None) == "high": + raise OmniClientError( + "MiniMax H3 step execution does not support the high-quality Cache-DiT profile " + "(quality=high); its hooks live on the shared transformer, so interleaved or " + "co-batched requests would reuse incompatible cache state. Drop --step-execution " + "or omit quality=high." + ) + prompt, multi_modal_data = self._extract_prompt(state.prompt) + context = self._prepare_request_inputs( + prompt=prompt, + multi_modal_data=multi_modal_data, + sampling=state.sampling, + ) + inputs = self._build_denoise_inputs(**self._denoise_kwargs(context)) + + sigmas_video = inputs["sigmas_video"] + sigmas_audio = inputs["sigmas_audio"] + if len(sigmas_video) < 2: + raise OmniClientError( + f"MiniMax H3 step execution needs at least one denoise step, got num_inference_steps=" + f"{len(sigmas_video) - 1}" + ) + + branch = inputs["branch"] + video_rows, audio_rows, cond_anchor, audio_anchor = minimax_h3_prepare_denoise_rows( + positive=branch, + initial_video_rows=inputs["video_rows"], + initial_audio_rows=inputs["audio_rows"], + keyframe_cond_rows=inputs["cond_anchor"], + audio_ref_rows=inputs["audio_anchor"], + device=self.device, + ) + + # The denoise loop consumes sigma pairs, so the schedule carries one more + # point than there are steps. ``timesteps`` holds the video branch because + # the shared contract gives a request exactly one timestep sequence; the + # audio schedule rides along in ``extra``. + state.timesteps = torch.tensor( + [1.0 - sigma for sigma in sigmas_video[:-1]], + dtype=torch.float32, + device=self.device, + ) + state.step_index = 0 + # Video rows are the batched tensor the runner slices per request; audio + # rows have a different width, so they stay request-private. + state.latents = video_rows + state.do_true_cfg = False # H3 checkpoints are CFG-distilled. + state.extra.update( + { + _STEP_BRANCH: branch, + _STEP_TRANSFORMER: self._transformer_for_task(context["task"]), + _STEP_AUDIO_ROWS: audio_rows, + _STEP_COND_ANCHOR: cond_anchor, + _STEP_AUDIO_ANCHOR: audio_anchor, + _STEP_SIGMAS_VIDEO: sigmas_video, + _STEP_SIGMAS_AUDIO: sigmas_audio, + _STEP_SHAPE: { + "height": context["height"], + "width": context["width"], + "latent_t": context["latent_t"], + "latent_h": context["latent_h"], + "latent_w": context["latent_w"], + "audio_t": context["audio_t"], + }, + } + ) + return state + + def denoise_step( + self, + input_batch: InputBatch, + *, + states: Sequence[StepRequestState] | None = None, + **kwargs: Any, + ) -> torch.Tensor | None: + """Run one denoise forward covering every request in the batch. + + Requests are concatenated into a single packed sequence that keeps one + attention document each, so the whole batch costs one DiT forward. + Backends that ignore ``cu_seqlens`` cannot express that isolation, so + they fall back to one forward per request. + """ + del kwargs + batch_states = list(states if states is not None else input_batch.states) + + branches = [state.extra[_STEP_BRANCH] for state in batch_states] + video_rows = [state.latents for state in batch_states] + audio_rows = [state.extra[_STEP_AUDIO_ROWS] for state in batch_states] + schedules = [_minimax_h3_step_schedule(state) for state in batch_states] + transformers = [state.extra[_STEP_TRANSFORMER] for state in batch_states] + mixed_transformers = len({id(transformer) for transformer in transformers}) > 1 + + # Both execution modes must publish denoise progress for step-gated + # attention features. Requests can differ in both step index and sigma + # schedule, so a batch that is not at one single point has nothing to + # publish and those gates stay dense -- which is their safe default. + progress = {(state.step_index, schedule["sigma_video"]) for state, schedule in zip(batch_states, schedules)} + minimax_h3_publish_denoise_progress(*(progress.pop() if len(progress) == 1 else (None, None))) + + if len(batch_states) > 1 and (mixed_transformers or not self._packed_batch_supported(transformers[0])): + if mixed_transformers: + logger.warning_once( + "MiniMax H3 step batch contains requests for different task-specific DiTs; " + "running %d requests one forward at a time.", + len(batch_states), + ) + elif any( + getattr(getattr(module, "attention", None), "use_ring", False) + for module in transformers[0].modules() + if isinstance(module, MiniMaxH3Attention) + ): + logger.warning_once( + "MiniMax H3 step batching is disabled when ring attention is active: " + "the ring kernels ignore packed cu_seqlens and would attend across request " + "boundaries. Running %d requests one forward at a time.", + len(batch_states), + ) + else: + logger.warning_once( + "MiniMax H3 step batching needs every attention on a backend that isolates " + "packed multi-document cu_seqlens (see AttentionBackend." + "supports_multi_doc_packed_varlen); running %d requests one forward at a time.", + len(batch_states), + ) + video_parts: list[torch.Tensor] = [] + audio_parts: list[torch.Tensor] = [] + for index, branch in enumerate(branches): + forward_kwargs = branch.forward_kwargs( + video_rows=video_rows[index], + audio_rows=audio_rows[index], + t_video=schedules[index]["t_video"], + t_audio=schedules[index]["t_audio"], + imgvid_cond_timestep=schedules[index]["imgvid_cond_timestep"], + audio_ref_cond_timestep=schedules[index]["audio_ref_cond_timestep"], + ) + request_video, request_audio = transformers[index](**forward_kwargs) + video_parts.append(request_video) + audio_parts.append(request_audio) + video_velocity = torch.cat(video_parts) + audio_velocity = torch.cat(audio_parts) + else: + forward_kwargs = minimax_h3_batched_forward_kwargs( + branches=branches, + video_rows=video_rows, + audio_rows=audio_rows, + t_video=[schedule["t_video"] for schedule in schedules], + t_audio=[schedule["t_audio"] for schedule in schedules], + imgvid_cond_timesteps=[schedule["imgvid_cond_timestep"] for schedule in schedules], + audio_ref_cond_timesteps=[schedule["audio_ref_cond_timestep"] for schedule in schedules], + ) + logger.debug( + "MiniMax H3 denoise step: %d request(s) packed into %d rows", + len(batch_states), + int(forward_kwargs["x"].shape[1]), + ) + video_velocity, audio_velocity = transformers[0](**forward_kwargs) + + # The shared contract carries one velocity tensor per step, and audio rows + # are a different width than video rows, so hand the audio branch to + # step_scheduler() through request-private state. + audio_parts_by_request = torch.split(audio_velocity, [int(branch.audio_pos.shape[0]) for branch in branches]) + for state, request_audio in zip(batch_states, audio_parts_by_request, strict=True): + state.extra[_STEP_AUDIO_NOISE_PRED] = request_audio + return video_velocity + + def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None: + """Apply one Euler-eta0 update to this request's video and audio rows.""" + del kwargs + # denoise_step() stages the audio half of this step's velocity; popping + # it keeps a second step_scheduler() call from reusing a stale one. + audio_noise_pred = state.extra.pop(_STEP_AUDIO_NOISE_PRED) + + branch = state.extra[_STEP_BRANCH] + schedule = _minimax_h3_step_schedule(state) + update = branch.update_mask_dev + audio_update = branch.audio_update_mask_dev + video_rows = state.latents + audio_rows = state.extra[_STEP_AUDIO_ROWS] + cond_anchor = state.extra[_STEP_COND_ANCHOR] + audio_anchor = state.extra[_STEP_AUDIO_ANCHOR] + device = video_rows.device + + x0_video = minimax_h3_rf_v_to_x0( + video_rows[update], + noise_pred.float()[update], + torch.tensor(schedule["t_video"], dtype=torch.float32, device=device), + ) + new_video = minimax_h3_euler_eta0_step( + video_rows[update], + x0_video, + sigma_curr=schedule["sigma_video"], + sigma_next=schedule["sigma_video_next"], + ) + video_rows = video_rows.clone() + video_rows[update] = new_video + if cond_anchor is not None: + video_rows[~update] = cond_anchor # per-step imgvid cond reset + + x0_audio = minimax_h3_rf_v_to_x0( + audio_rows[audio_update], + audio_noise_pred.float()[audio_update], + torch.tensor(schedule["t_audio"], dtype=torch.float32, device=device), + ) + new_audio = minimax_h3_euler_eta0_step( + audio_rows[audio_update], + x0_audio, + sigma_curr=schedule["sigma_audio"], + sigma_next=schedule["sigma_audio_next"], + ) + audio_rows = audio_rows.clone() + audio_rows[audio_update] = new_audio + if audio_anchor is not None: + audio_rows[~audio_update] = audio_anchor # per-step audio ref reset + + state.latents = video_rows + state.extra[_STEP_AUDIO_ROWS] = audio_rows + state.step_index += 1 + + def post_decode(self, state: StepRequestState, **kwargs: Any) -> DiffusionOutput: + """Unpack the denoised rows and run the joint video/audio VAE decode.""" + del kwargs + shape = state.extra[_STEP_SHAPE] + video_latent, audio_latent = self._unpack_denoised_rows( + state.extra[_STEP_BRANCH], + state.latents, + state.extra[_STEP_AUDIO_ROWS], + latent_t=shape["latent_t"], + latent_h=shape["latent_h"], + latent_w=shape["latent_w"], + audio_t=shape["audio_t"], + ) + video, audio = self.decode( + video_latent, + audio_latent, + height=shape["height"], + width=shape["width"], + ) + return DiffusionOutput( + output=(video, audio), + post_process_func=get_minimax_h3_post_process_func(self.od_config), + stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}), + ) + __all__ = [ "MiniMaxH3Pipeline", diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 2992d2bde..de10bafa4 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -81,6 +81,31 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def _dit_any_rank_failed(local_failed: bool) -> bool: + """All-reduce a per-request failure flag across the DiT process group. + + Every DiT rank must reach this point in lockstep; the caller wraps + ``pipeline.prepare_encode`` so both the success and the exception paths + call the helper. In single-rank execution this collapses to the local + flag with no collectives issued. + """ + if not torch.distributed.is_initialized(): + return local_failed + try: + from vllm_omni.diffusion.distributed.parallel_state import get_dit_group + + group = get_dit_group() + except (AssertionError, ImportError): + group = None + if group is None: + return local_failed + signal = torch.tensor(1 if local_failed else 0, dtype=torch.int32) + if current_omni_platform.is_available(): + signal = signal.to(device=current_omni_platform.device_type) + torch.distributed.all_reduce(signal, op=torch.distributed.ReduceOp.MAX, group=group) + return bool(signal.item()) + + def _normalize_pipeline_outputs( outputs: object, *, @@ -853,25 +878,70 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): return resolved, new_request_ids - def _prepare_batch_inputs(self, states: list[StepRequestState], new_request_ids: list[str]) -> InputBatch: + def _prepare_batch_inputs( + self, + states: list[StepRequestState], + new_request_ids: list[str], + ) -> tuple[list[StepRequestState], InputBatch | None, list[RunnerOutput]]: # process new reqs + prepared_states: list[StepRequestState] = [] + error_outputs: list[RunnerOutput] = [] for state in states: if state.request_id in new_request_ids: - self._initialize_generator(state.sampling) - clear_pipeline_stage_durations(self.pipeline) - # encode - self.pipeline.prepare_encode(state) - merge_stage_durations( - state, - consume_pipeline_stage_durations(self.pipeline), - ) + # Everything that runs before ``_dit_any_rank_failed`` must be + # inside the try: an exception in ``_initialize_generator`` or + # ``clear_pipeline_stage_durations`` on one rank would skip the + # all-reduce here while every peer proceeds into it, and the + # peers then hang on the NCCL collective until timeout. + per_req_exc: BaseException | None = None + try: + self._initialize_generator(state.sampling) + clear_pipeline_stage_durations(self.pipeline) + # encode + self.pipeline.prepare_encode(state) + merge_stage_durations( + state, + consume_pipeline_stage_durations(self.pipeline), + ) + except Exception as exc: + per_req_exc = exc + # Pipelines that do rank-0-only work (e.g. MiniMax H3 + # reference-video prep) must broadcast per-request failures + # internally so downstream collectives stay in step; even so, + # cross-check that every DiT rank agrees so a rank-local error + # (or a future pipeline that omits the guard) does not leave + # the process group half-way through a new request. + if _dit_any_rank_failed(per_req_exc is not None): + self.state_cache.pop(state.request_id, None) + if per_req_exc is None: + per_req_exc = RuntimeError( + f"Stepwise preparation failed on another DiT rank for {state.request_id}" + ) + logger.error( + "Stepwise request preparation failed for %s: %s", + state.request_id, + per_req_exc, + exc_info=isinstance(per_req_exc, Exception), + ) + error_outputs.append( + RunnerOutput( + request_id=state.request_id, + step_index=state.step_index, + finished=True, + result=DiffusionOutput.from_exception(per_req_exc), + ) + ) + continue + prepared_states.append(state) + if not prepared_states: + return prepared_states, None, error_outputs input_batch = InputBatch.make_batch( - states, + prepared_states, cached_batch=getattr(self, "input_batch", None), ) self.input_batch = input_batch - return input_batch + return prepared_states, input_batch, error_outputs def _update_states_after( self, @@ -978,7 +1048,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): and current_omni_platform.is_available() ): current_omni_platform.reset_peak_memory_stats() - input_batch = self._prepare_batch_inputs(states, new_request_ids) + states, input_batch, runner_output_list = self._prepare_batch_inputs(states, new_request_ids) + if input_batch is None: + return BatchRunnerOutput.from_list(runner_output_list) attn_metadata = {} with set_forward_context( @@ -995,7 +1067,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): denoise_stage_durations, ) - runner_output_list = [] pipeline_interrupted = getattr(self.pipeline, "interrupt", False) if noise_pred is None and pipeline_interrupted: for state in states: @@ -1048,6 +1119,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): offset = offset + row_num except Exception as per_req_exc: offset = offset + row_num + self.state_cache.pop(req.request_id, None) logger.error( "Stepwise per-request error for %s: %s", req.request_id, @@ -1059,7 +1131,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): request_id=req.request_id, step_index=req.step_index, finished=True, - result=DiffusionOutput(error=str(per_req_exc)), + result=DiffusionOutput.from_exception(per_req_exc), ) )