[Bugfix] Align MiniMax-H3 inference steps with denoiser evaluations (#7219)

Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

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 |

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`

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

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")
 

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(),

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)

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")

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(

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:

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()]