[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>

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