From 67aa30c96697e1aab4a135d5e673d2a457dffe9f Mon Sep 17 00:00:00 2001 From: Honghan Zhu Date: Thu, 24 Sep 2026 11:56:38 +0800 Subject: [PATCH] [Bugfix] Align MiniMax-H3 inference steps with denoiser evaluations (#7219) Signed-off-by: zhuhh97 Co-authored-by: Gao Han --- docs/design/feature/skip_softmax.md | 2 +- docs/user_guide/diffusion/lora.md | 14 ++++++++--- .../minimax_h3/test_minimax_h3_contract.py | 21 +++++++++++++--- .../test_minimax_h3_execution_modes.py | 7 ++++-- .../models/minimax_h3/test_minimax_h3_lora.py | 18 ++++++------- .../test_minimax_h3_step_execution.py | 18 +++++++------ .../test_minimax_h3_turbo_matrix.py | 25 +++++++++++++++++++ .../e2e/online_serving/minimax_h3/_common.py | 4 +-- .../models/minimax_h3/pipeline_minimax_h3.py | 7 +++--- .../models/minimax_h3/time_request.py | 14 ++++------- 10 files changed, 90 insertions(+), 40 deletions(-) diff --git a/docs/design/feature/skip_softmax.md b/docs/design/feature/skip_softmax.md index 632df5c54..c72864bd5 100644 --- a/docs/design/feature/skip_softmax.md +++ b/docs/design/feature/skip_softmax.md @@ -178,7 +178,7 @@ t = s * u / (1 + (s - 1) * u) ``` and a large `s` pushes most positions toward `t ≈ 1`. With the default `s = 12` and -`num_inference_steps = 50` (50 sigma points, `N = 49` denoiser forwards), the published sequence +`num_inference_steps = 49` (50 sigma points, `N = 49` denoiser forwards), the sequence starts `1.000, 0.998, 0.996, 0.995, 0.993, ...` and stays above `0.9` for 28 steps, giving: | `disabled_until_timestep` | Dense steps | Skip-Softmax steps | diff --git a/docs/user_guide/diffusion/lora.md b/docs/user_guide/diffusion/lora.md index 15817db34..c774cb288 100644 --- a/docs/user_guide/diffusion/lora.md +++ b/docs/user_guide/diffusion/lora.md @@ -244,14 +244,22 @@ Notes: ## MiniMax-H3 adapter-declared schedules -MiniMax-H3 supports two few-step mechanisms that must not be conflated: +MiniMax-H3 counts `num_inference_steps` as denoising updates: requesting `8` +produces eight updates using nine sigma boundaries, including the terminal zero. +Earlier uniform-schedule requests counted boundaries instead; to reproduce an +old N-point schedule, request N-1 steps. This also changes Turbo requests from +`5`/`9` to `4`/`8`, preserving their sampling grids. Checkpoint-pinned and native +LoRA schedules already count updates and are unchanged. + +MiniMax-H3 supports the following few-step mechanisms: - **Checkpoint-pinned schedule**: a merged release writes `base_schedule` into `model_index.json`. Requests must pass `num_inference_steps` as the interval count (for example `4` for `[1.0, 0.7, 0.4, 0.15, 0.0]`). - **Runtime Turbo LoRA**: each LightX2V Turbo artifact carries its own contract, - read from its filename. A four-step artifact requests five sigma points and an - eight-step one requests nine; the 768p retrains enforce `flow_shift=6` and the + read from its filename. A four-step artifact requests `num_inference_steps=4` + and an eight-step one requests `8`; the scheduler adds the terminal sigma + boundary. The 768p retrains enforce `flow_shift=6` and the 544p artifacts `flow_shift=12`. `audio_flow_shift=3` across the family. A request that does not match the loaded artifact is rejected by name. Alpha comes from the artifact's metadata, or 8 when it declares none. `ref2v` 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 13f2f5a68..11266cff3 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -7,7 +7,7 @@ from contextlib import contextmanager, nullcontext from multiprocessing.reduction import ForkingPickler from types import SimpleNamespace from typing import Any, TypeVar -from unittest.mock import Mock, patch +from unittest.mock import Mock, create_autospec, patch import numpy as np import pytest @@ -947,7 +947,7 @@ def test_shifted_sigma_schedule_matches_reference_values(): minimax_h3_time_shift_sigmas, ) - sigmas = minimax_h3_time_shift_sigmas(num_steps=5, shift_scale=12.0) + sigmas = minimax_h3_time_shift_sigmas(num_steps=4, shift_scale=12.0) assert sigmas == pytest.approx( [1.0, 0.9729729891, 0.9230769277, 0.8000000119, 0.0], @@ -955,6 +955,19 @@ def test_shifted_sigma_schedule_matches_reference_values(): ) +@pytest.mark.parametrize("num_steps", [1, 8, 50]) +@pytest.mark.parametrize("shift_scale", [3.0, 12.0]) +def test_uniform_sigma_schedule_has_one_interval_per_requested_step(num_steps, shift_scale): + from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_time_shift_sigmas + + sigmas = minimax_h3_time_shift_sigmas(num_steps=num_steps, shift_scale=shift_scale) + + assert len(sigmas) == num_steps + 1 + assert sigmas[0] == 1.0 + assert sigmas[-1] == 0.0 + assert all(current > following for current, following in zip(sigmas, sigmas[1:])) + + def test_base_schedule_overrides_the_uniform_sigma_positions(): from vllm_omni.diffusion.models.minimax_h3.time_request import ( minimax_h3_time_shift_sigmas, @@ -1751,12 +1764,14 @@ def test_text_encoder_rejects_serialized_fp8(): def test_text_encoder_linear_delegates_quantization_to_vllm_factory(): from vllm.model_executor.layers.linear import UnquantizedLinearMethod + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig import vllm_omni.diffusion.models.minimax_h3.encoder as encoder_module group = SimpleNamespace(rank_in_group=0, world_size=1) method = UnquantizedLinearMethod() - quant_config = Mock() + quant_config = create_autospec(QuantizationConfig, instance=True) + quant_config.online_quantization_config = None quant_config.get_quant_method.return_value = method # vLLM 0.30: no online quantization for this checkpoint-quantized layer quant_config.online_quantization_config = None diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py index 18f1cd9fa..b145d3d44 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py @@ -17,7 +17,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @pytest.fixture -def pipeline(): +def pipeline(monkeypatch): + # These CPU tests exercise conditioning, not device allocator cleanup. + monkeypatch.setattr(MiniMaxH3Pipeline, "_release_stage_cache", lambda self: None) model = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(model) model.load_text_encoder = True @@ -119,7 +121,8 @@ def test_prepare_encode_uses_the_same_mode_and_conditioning(pipeline, task, load assert pipeline.prepare_encode(state) is state assert state.latents is not None assert state.latents.shape[1] == 96 - assert state.total_steps == 1 + # num_inference_steps counts denoiser evaluations, not sigma boundaries. + assert state.total_steps == request.sampling_params.num_inference_steps assert pipeline.encode_prompt.call_count == int(load_text_encoder) assert pipeline.video_vae.encode_image.call_count == int(load_text_encoder and task != "t2va") diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py index 9039f94f8..2396b800b 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py @@ -545,13 +545,13 @@ def test_h3_turbo_requires_all_loaded_targets_to_bind(): @pytest.mark.parametrize( ("num_inference_steps", "extra_args", "error"), [ - (None, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"), - (4, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"), - ("5", {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"), - (5, {"flow_shift": 7.0, "audio_flow_shift": 3.0}, "flow_shift=6"), - (5, {"flow_shift": "bad", "audio_flow_shift": 3.0}, "flow_shift=6"), - (5, {"flow_shift": 6.0, "audio_flow_shift": 4.0}, "audio_flow_shift=3"), - (5, {"flow_shift": 6.0, "audio_flow_shift": []}, "audio_flow_shift=3"), + (None, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=4"), + (5, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=4"), + ("5", {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=4"), + (4, {"flow_shift": 7.0, "audio_flow_shift": 3.0}, "flow_shift=6"), + (4, {"flow_shift": "bad", "audio_flow_shift": 3.0}, "flow_shift=6"), + (4, {"flow_shift": 6.0, "audio_flow_shift": 4.0}, "audio_flow_shift=3"), + (4, {"flow_shift": 6.0, "audio_flow_shift": []}, "audio_flow_shift=3"), ], ) def test_h3_turbo_rejects_unsupported_sampling(num_inference_steps, extra_args, error): @@ -569,7 +569,7 @@ def test_h3_turbo_rejects_unsupported_sampling(num_inference_steps, extra_args, pipeline._validate_turbo_sampling(sampling, _spec()) -def test_h3_turbo_accepts_five_sigma_points_for_four_nfe(): +def test_h3_turbo_accepts_four_steps_for_four_nfe(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline pipeline = object.__new__(MiniMaxH3Pipeline) @@ -577,7 +577,7 @@ def test_h3_turbo_accepts_five_sigma_points_for_four_nfe(): pipeline.default_audio_shift = 3.0 pipeline._validate_turbo_sampling( SimpleNamespace( - num_inference_steps=5, + num_inference_steps=4, extra_args={"flow_shift": 6.0, "audio_flow_shift": 3.0}, ), _spec(), 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 index 755899649..00fdb4cdc 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py @@ -88,10 +88,10 @@ def _make_branch(*, text_len: int, latent_t: int, latent_h: int, latent_w: int, return branch, video_rows, audio_rows -def _sigmas(num_points: int, shift: float) -> list[float]: +def _sigmas(num_steps: 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) + return minimax_h3_time_shift_sigmas(num_steps=num_steps, shift_scale=shift) def _make_state(request_id: str, model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio): @@ -130,15 +130,16 @@ def _step_pipeline(model, *, packed_batch_supported: bool = True): return pipeline -def test_step_execution_matches_request_mode_denoise_loop(): +@pytest.mark.parametrize("num_steps", [1, 8, 50]) +def test_step_execution_matches_request_mode_denoise_loop(num_steps, mocker): """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() + model = mocker.Mock(wraps=_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) + sigmas_video = _sigmas(num_steps, 12.0) + sigmas_audio = _sigmas(num_steps, 3.0) reference_video, reference_audio = minimax_h3_denoise_loop( model=model, @@ -151,6 +152,8 @@ def test_step_execution_matches_request_mode_denoise_loop(): device=torch.device("cpu"), ) + assert model.call_count == num_steps + model.reset_mock() pipeline = _step_pipeline(model) state = _make_state("req-0", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio) input_batch = SimpleNamespace(states=(state,)) @@ -161,7 +164,8 @@ def test_step_execution_matches_request_mode_denoise_loop(): pipeline.step_scheduler(state, noise_pred) steps += 1 - assert steps == len(sigmas_video) - 1 + assert steps == num_steps + assert model.call_count == num_steps 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) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_turbo_matrix.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_turbo_matrix.py index 35e46af6b..25c20d005 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_turbo_matrix.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_turbo_matrix.py @@ -5,6 +5,8 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest from vllm_omni.diffusion.models.minimax_h3.lora import ( @@ -79,6 +81,29 @@ def test_sigma_points_follow_the_step_count() -> None: assert _spec("minimax_h3_fl2v_turbo_8step_v1.0_768p_bf16.safetensors").sigma_points == 9 +@pytest.mark.parametrize(("filename", "task", "steps", "video_shift"), PUBLISHED) +def test_turbo_request_counts_evaluations_and_preserves_sigma_grid(filename, task, steps, video_shift): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_time_shift_sigmas + from vllm_omni.errors import OmniClientError + + spec = _spec(filename) + pipeline = object.__new__(MiniMaxH3Pipeline) + pipeline.default_video_shift = video_shift + pipeline.default_audio_shift = spec.audio_shift + sampling = SimpleNamespace(num_inference_steps=steps, extra_args={}) + pipeline._validate_turbo_sampling(sampling, spec) + for shift in (video_shift, spec.audio_shift): + sigmas = minimax_h3_time_shift_sigmas(num_steps=steps, shift_scale=shift) + # The old point-count API requested steps + 1 points on this same grid. + base = [1.0 - i / steps for i in range(steps + 1)] + expected = [shift * value / (1 + (shift - 1) * value) for value in base] + assert sigmas == pytest.approx(expected, abs=1e-7) + sampling.num_inference_steps = spec.sigma_points + with pytest.raises(OmniClientError, match=f"num_inference_steps={steps}"): + pipeline._validate_turbo_sampling(sampling, spec) + + def test_task_family_decides_supported_tasks() -> None: fl2v = _spec("minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors") ref2v = _spec("minimax_h3_ref2v_turbo_8step_v1.0_768p_bf16.safetensors") diff --git a/tests/e2e/online_serving/minimax_h3/_common.py b/tests/e2e/online_serving/minimax_h3/_common.py index e7e6ae8c5..69251748d 100644 --- a/tests/e2e/online_serving/minimax_h3/_common.py +++ b/tests/e2e/online_serving/minimax_h3/_common.py @@ -267,8 +267,8 @@ def turbo_form(seed: int) -> dict[str, str]: "height": str(HEIGHT), "fps": str(FPS), # The public Turbo artifact has four denoiser evaluations but its API - # contract requests five sigma points. - "num_inference_steps": "5", + # contract requests four denoiser evaluations (five sigma points). + "num_inference_steps": "4", "flow_shift": "6", "seed": str(seed), "extra_params": json.dumps( 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 f2c9854ee..21c6e2681 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -781,17 +781,16 @@ class MiniMaxH3Pipeline( def _validate_turbo_sampling(self, sampling: Any, spec: TurboSpec) -> None: """Hold a request to the contract of the artifact that is loaded. - Sigma-point count and both flow shifts vary across the Turbo family, so + Denoiser count and both flow shifts vary across the Turbo family, so each is checked against the adapter's own spec rather than a single published configuration. """ extra = sampling.extra_args or {} - sigma_points = sampling.num_inference_steps - if sigma_points != spec.sigma_points: + if sampling.num_inference_steps != spec.denoise_steps: raise OmniClientError( f"{spec.filename} is a {spec.denoise_steps}-step artifact and requires " - f"num_inference_steps={spec.sigma_points} " + f"num_inference_steps={spec.denoise_steps} " f"({spec.sigma_points} sigma points produce {spec.denoise_steps} denoiser evaluations)" ) try: diff --git a/vllm_omni/diffusion/models/minimax_h3/time_request.py b/vllm_omni/diffusion/models/minimax_h3/time_request.py index 92f6b3b92..35ceef9f7 100644 --- a/vllm_omni/diffusion/models/minimax_h3/time_request.py +++ b/vllm_omni/diffusion/models/minimax_h3/time_request.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations from collections.abc import Sequence @@ -41,7 +42,7 @@ def _time_shift_sigmas( shift_scale: float = 6.0, base_schedule: Sequence[float] | None = None, ) -> list[float]: - """Build a shifted sigma schedule. + """Build N + 1 sigma boundaries for N denoiser evaluations. ``base_schedule`` supplies the rectified-flow positions explicitly and takes precedence over ``num_steps``. Distilled checkpoints need it because their @@ -59,21 +60,16 @@ def _time_shift_sigmas( if num_steps <= 0: raise ValueError("MiniMax H3 num_steps must be > 0") - # The rectified-flow sigma range is fixed at [1.0, 0.0]. + # Requests count denoiser evaluations, not sigma points. Include both + # endpoints so even a one-step request traverses the full [1.0, 0.0] range. base = torch.linspace( 1.0, 0.0, - int(num_steps), + int(num_steps) + 1, device="cpu", dtype=torch.float32, ) shifted = float(shift_scale) * base / (1 + (float(shift_scale) - 1) * base) - shifted, _ = torch.unique_consecutive(shifted, return_counts=True) - # A one-point request is still exactly one point. Normal serving uses - # multiple points, but preserving the requested cardinality keeps - # ``num_inference_steps`` the sole schedule-size control. - if num_steps > 1 and shifted[-1].item() > 0.0: - shifted = torch.cat([shifted, torch.tensor([0.0], dtype=shifted.dtype)]) return [float(value) for value in shifted.tolist()]