[Feature] Add MiniMax H3 res_multistep sampler (#8378)

Signed-off-by: nodeeeeee <zhangkai.nodeee@gmail.com>

recipes/MiniMaxAI/MiniMax-H3.md

@@ -988,6 +988,42 @@ settings, and both videos with any comparison; the command alone is not quality
 evidence. A matching seed does not imply matching noise over differently sized
 full and windowed latents.
 
+## Request-scoped sampler
+
+H3 supports `euler` (the default) and `res_multistep`. Select the sampler per
+request; it applies to both video and audio, using their separate sigma schedules.
+For the HTTP examples above, set `num_inference_steps=20` and add
+`"sampler":"res_multistep"` to the existing `extra_params` object. For example,
+the T2VA fields become:
+
+```bash
+-F 'num_inference_steps=20' \
+-F 'extra_params={"task":"t2va","duration":8.7,"audio_flow_shift":3.0,"sampler":"res_multistep"}'
+```
+
+For the offline Python API, use `extra_args`:
+
+```python
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+sampling_params = OmniDiffusionSamplingParams(
+    num_inference_steps=20,
+    seed=42,
+    extra_args={"task": "t2va", "sampler": "res_multistep"},
+)
+```
+
+Both samplers use one DiT evaluation per step in request mode and step execution.
+`res_multistep` retains the previous denoised prediction independently for each
+request and each stream; its first and final updates use Euler. Omitting `sampler`
+or selecting `euler` preserves the existing Euler update. Reducing the step count
+reduces model evaluations, but quality depends on the prompt and configuration;
+20 RES steps are a candidate to compare against an Euler-50 reference.
+
+Requests using a fixed distilled `base_schedule` from a checkpoint or adapter
+must keep `sampler="euler"`. RES sampling with those schedules has not been
+validated and is rejected with `ValueError`.
+
 ## Request-scoped quality
 
 Add one of these fields to any HTTP request above. No Cache-DiT startup option
@@ -1408,6 +1444,7 @@ four-step adapter keeps its original sampling and fixed-top-k behavior.
 | `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count |
 | `fps` | `24` | H3 output FPS is fixed |
 | `num_inference_steps` | `50` | Matches the reference accuracy workloads |
+| `extra_params.sampler` | `euler` | Per-request `euler` or `res_multistep` for both streams; fixed distilled schedules require Euler |
 | `flow_shift` | `12` | Video sigma shift |
 | `audio_flow_shift` | `3` | Audio sigma shift, passed in `extra_params` |
 | `seed` | Task-specific | Use a fixed value for reproducibility |

tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py

@@ -102,7 +102,11 @@ def _make_state(request_id: str, model, branch, video_rows, audio_rows, sigmas_v
     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
+    from vllm_omni.diffusion.models.minimax_h3.sampling import create_h3_sample_solver
+
     state.extra = {
+        mod._STEP_VIDEO_SOLVER: create_h3_sample_solver("euler", sigmas_video),
+        mod._STEP_AUDIO_SOLVER: create_h3_sample_solver("euler", sigmas_audio),
         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.

vllm_omni/diffusion/models/minimax_h3/denoise_loop.py

@@ -3,7 +3,7 @@
 """MiniMax H3 cfg-distilled full denoise loop.
 
 Per step, the positive presentation is forwarded exactly once. Video and audio
-target rows chain through the Euler-eta0 update while visual and audio condition
+target rows chain through the selected solver while visual and audio condition
 rows stay pinned to their noised step-0 anchors.
 """
 
@@ -25,8 +25,8 @@ from vllm_omni.diffusion.forward_context import (
 from vllm_omni.platforms import current_omni_platform
 
 from .latent_mask import MiniMaxH3LatentEdit, minimax_h3_prepare_edit_rows
+from .sampling import create_h3_sample_solver
 from .scheduling_minimax_h3_euler_ancestral import (
-    minimax_h3_euler_eta0_step,
     minimax_h3_rf_v_to_x0,
 )
 
@@ -333,6 +333,7 @@ def minimax_h3_denoise_loop(
     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,
+    sampler: str = "euler",
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Run the full denoise loop; returns final (video_rows, audio_rows).
 
@@ -365,13 +366,15 @@ def minimax_h3_denoise_loop(
     if audio_edit is not None:
         audio_edit = audio_edit.to(device=device, dtype=torch.float32)
 
+    video_solver = create_h3_sample_solver(sampler, sigmas_video)
+    audio_solver = create_h3_sample_solver(sampler, sigmas_audio)
     num_steps = len(sigmas_video) - 1
     for step in range(num_steps):
         check_request_cancellation()
         step_cm = step_profiler(step) if step_profiler is not None else nullcontext()
         with step_cm:
-            s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1]
-            s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1]
+            s_v = sigmas_video[step]
+            s_a = sigmas_audio[step]
             # 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
@@ -425,7 +428,7 @@ def minimax_h3_denoise_loop(
                     mv_video_t,
                     t_v,
                 )
-            new_target = minimax_h3_euler_eta0_step(video_rows[update], x0_video, sigma_curr=s_v, sigma_next=s_v_next)
+            new_target = video_solver.step(video_rows[update], x0_video, step)
             video_rows = video_rows.clone()
             video_rows[update] = new_target
             if cond_anchor is not None:
@@ -443,9 +446,7 @@ def minimax_h3_denoise_loop(
                     mv_audio_t,
                     t_a,
                 )
-            new_audio = minimax_h3_euler_eta0_step(
-                audio_rows[audio_update], x0_audio, sigma_curr=s_a, sigma_next=s_a_next
-            )
+            new_audio = audio_solver.step(audio_rows[audio_update], x0_audio, step)
             audio_rows = audio_rows.clone()
             audio_rows[audio_update] = new_audio if positive.locked_audio_rows is None else positive.locked_audio_rows
             if audio_anchor is not None:

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -157,8 +157,8 @@ from .packed_tokens import (
     minimax_h3_unpatchify_video_tokens,
 )
 from .quality_policy import MINIMAX_H3_GENERIC_CACHE_KEY, MiniMaxH3QualityPolicy
+from .sampling import create_h3_sample_solver, normalize_h3_sampler
 from .scheduling_minimax_h3_euler_ancestral import (
-    minimax_h3_euler_eta0_step,
     minimax_h3_rf_v_to_x0,
 )
 from .time_request import (
@@ -302,6 +302,8 @@ _REFINE_KEYFRAME_CONDITION = "minimax_h3_refine_keyframe_condition"
 _STEP_BRANCH = "minimax_h3_branch"
 _STEP_AUDIO_ROWS = "minimax_h3_audio_rows"
 _STEP_AUDIO_NOISE_PRED = "minimax_h3_audio_noise_pred"
+_STEP_VIDEO_SOLVER = "minimax_h3_video_solver"
+_STEP_AUDIO_SOLVER = "minimax_h3_audio_solver"
 _STEP_SIGMAS_VIDEO = "minimax_h3_sigmas_video"
 _STEP_SIGMAS_AUDIO = "minimax_h3_sigmas_audio"
 _STEP_COND_ANCHOR = "minimax_h3_cond_anchor"
@@ -1780,6 +1782,7 @@ class MiniMaxH3Pipeline(
         audio_edit_clean_rows: torch.Tensor | None = None,
         audio_edit_mask_rows: torch.Tensor | None = None,
         audio_edit_restore_mask_rows: torch.Tensor | None = None,
+        sampler: str = "euler",
         init_latents: tuple[torch.Tensor, torch.Tensor] | None = None,
         refine: MiniMaxH3LatentRefineSpec | None = None,
     ) -> dict[str, Any]:
@@ -1792,6 +1795,12 @@ class MiniMaxH3Pipeline(
         one: the rows start from those latents re-noised to the schedule
         position ``refine`` selects, and the returned schedules begin there.
         """
+        sampler = normalize_h3_sampler(sampler)
+        if sampler != "euler" and base_schedule is not None:
+            raise ValueError(
+                "MiniMax H3 res_multistep sampling with a fixed distilled sigma schedule has not been validated; "
+                "use sampler='euler' for this schedule."
+            )
         video_sigmas = minimax_h3_time_shift_sigmas(
             num_steps=num_steps,
             shift_scale=video_shift,
@@ -1977,6 +1986,7 @@ class MiniMaxH3Pipeline(
             "audio_anchor": (
                 None if audio_anchor is None else audio_anchor.to(device=self.device, dtype=torch.float32)
             ),
+            "sampler": sampler,
             "sigmas_video": video_sigmas,
             "sigmas_audio": audio_sigmas,
             "video_edit": video_edit,
@@ -2048,10 +2058,12 @@ class MiniMaxH3Pipeline(
         audio_edit_clean_rows: torch.Tensor | None = None,
         audio_edit_mask_rows: torch.Tensor | None = None,
         audio_edit_restore_mask_rows: torch.Tensor | None = None,
+        sampler: str = "euler",
         init_latents: tuple[torch.Tensor, torch.Tensor] | None = None,
         refine: MiniMaxH3LatentRefineSpec | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         inputs = self._build_denoise_inputs(
+            sampler=sampler,
             task=task,
             text_embeddings=text_embeddings,
             text_tags=text_tags,
@@ -2099,6 +2111,7 @@ class MiniMaxH3Pipeline(
         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(
+                    sampler=inputs["sampler"],
                     model=transformer,
                     positive=branch,
                     initial_video_rows=inputs["video_rows"],
@@ -2862,6 +2875,12 @@ class MiniMaxH3Pipeline(
 
         self._prepare_adaln_adapter(sampling)
         base_schedule, num_steps = self._resolve_sigma_positions(task, sampling)
+        sampler = normalize_h3_sampler(extra.get("sampler"))
+        if sampler != "euler" and base_schedule is not None:
+            raise ValueError(
+                "MiniMax H3 res_multistep sampling with a fixed distilled sigma schedule has not been validated; "
+                "use sampler='euler' for this schedule."
+            )
         transformer = getattr(
             self, "transformers_ref" if task == "ref2va" and hasattr(self, "transformers_ref") else "transformer", None
         )
@@ -2939,6 +2958,7 @@ class MiniMaxH3Pipeline(
             "keyframe_frame_indices": list(conditioning.keyframe_frame_indices) or None,
             "pad_seq_len": _resolve_pad_seq_len(extra.get("pad_seq_len")),
             "seed": int(sampling.seed if sampling.seed is not None else 42),
+            "sampler": sampler,
             "num_steps": num_steps,
             "video_shift": float(extra.get("flow_shift", self.default_video_shift)),
             "audio_shift": float(extra.get("audio_flow_shift", self.default_audio_shift)),
@@ -2961,7 +2981,10 @@ class MiniMaxH3Pipeline(
     @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}
+        return {
+            **{key: context[key] for key in _MINIMAX_H3_DENOISE_INPUT_KEYS},
+            "sampler": context.get("sampler", "euler"),
+        }
 
     @torch.no_grad()
     def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput:
@@ -3165,6 +3188,8 @@ class MiniMaxH3Pipeline(
                 _STEP_AUDIO_ROWS: audio_rows,
                 _STEP_COND_ANCHOR: cond_anchor,
                 _STEP_AUDIO_ANCHOR: audio_anchor,
+                _STEP_VIDEO_SOLVER: create_h3_sample_solver(inputs.get("sampler", "euler"), sigmas_video),
+                _STEP_AUDIO_SOLVER: create_h3_sample_solver(inputs.get("sampler", "euler"), sigmas_audio),
                 _STEP_SIGMAS_VIDEO: sigmas_video,
                 _STEP_SIGMAS_AUDIO: sigmas_audio,
                 _STEP_VIDEO_EDIT: inputs.get("video_edit"),
@@ -3318,7 +3343,7 @@ class MiniMaxH3Pipeline(
         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."""
+        """Apply one request-local solver update to the 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.
@@ -3347,12 +3372,7 @@ class MiniMaxH3Pipeline(
                 noise_pred.float()[update],
                 schedule["t_video"],
             )
-        new_video = minimax_h3_euler_eta0_step(
-            video_rows[update],
-            x0_video,
-            sigma_curr=schedule["sigma_video"],
-            sigma_next=schedule["sigma_video_next"],
-        )
+        new_video = state.extra[_STEP_VIDEO_SOLVER].step(video_rows[update], x0_video, state.step_index)
         video_rows = video_rows.clone()
         video_rows[update] = new_video
         if cond_anchor is not None:
@@ -3371,12 +3391,7 @@ class MiniMaxH3Pipeline(
                 audio_noise_pred.float()[audio_update],
                 schedule["t_audio"],
             )
-        new_audio = minimax_h3_euler_eta0_step(
-            audio_rows[audio_update],
-            x0_audio,
-            sigma_curr=schedule["sigma_audio"],
-            sigma_next=schedule["sigma_audio_next"],
-        )
+        new_audio = state.extra[_STEP_AUDIO_SOLVER].step(audio_rows[audio_update], x0_audio, state.step_index)
         audio_rows = audio_rows.clone()
         audio_rows[audio_update] = new_audio if branch.locked_audio_rows is None else branch.locked_audio_rows
         if audio_anchor is not None:

vllm_omni/diffusion/models/minimax_h3/sampling.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Request-local H3 solvers reconstructed from the surviving experiment records.
+
+RES coefficients follow ComfyUI_RH_MinMaxH3 sampler_core.py at
+6abdf97ae8cdce548c751ce850adc5010b583b7c. Euler delegates to the existing
+implementation to preserve its exact floating-point operation order.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+
+import torch
+
+from .scheduling_minimax_h3_euler_ancestral import minimax_h3_euler_eta0_step
+
+
+def normalize_h3_sampler(name: str | None) -> str:
+    if name is None:
+        return "euler"
+    if name not in ("euler", "res_multistep"):
+        raise ValueError(f"Unsupported MiniMax H3 sampler: {name!r}; expected euler or res_multistep")
+    return name
+
+
+def res_multistep_coeffs(sigmas: Sequence[float]) -> list[tuple[float, float] | None]:
+    values = [float(s) for s in sigmas]
+    if len(values) < 2 or any(not math.isfinite(s) or not 0 <= s <= 1 for s in values):
+        raise ValueError("H3 sigmas must contain at least two finite values in [0, 1]")
+    if any(a <= b for a, b in zip(values, values[1:])):
+        raise ValueError("H3 sigmas must be strictly decreasing")
+    coeffs = []
+    for i, (curr, nxt) in enumerate(zip(values, values[1:])):
+        if i == 0 or nxt == 0:
+            coeffs.append(None)
+            continue
+        h = math.log(curr / nxt)
+        c2 = math.log(curr / values[i - 1]) / h
+        # Stable evaluation near h=0, where subtracting phi1-1 loses precision.
+        phi1 = math.expm1(-h) / (-h)
+        phi2 = (phi1 - 1) / (-h) if abs(h) >= 1e-5 else 0.5 - h / 6 + h * h / 24 - h**3 / 120
+        coeffs.append((h * (phi1 - phi2 / c2), h * phi2 / c2))
+    return coeffs
+
+
+class H3SampleSolver:
+    """Own the denoised history for one request and one video/audio stream."""
+
+    def __init__(self, name: str | None, sigmas: Sequence[float]):
+        self.name = normalize_h3_sampler(name)
+        self.sigmas = tuple(float(s) for s in sigmas)
+        self.coeffs = res_multistep_coeffs(self.sigmas)
+        self.old_denoised: torch.Tensor | None = None
+        self.step_index = 0
+
+    def step(self, state: torch.Tensor, denoised: torch.Tensor, step: int) -> torch.Tensor:
+        if step != self.step_index or not 0 <= step < len(self.coeffs):
+            raise ValueError("H3 solver steps must be consumed once, in order")
+        coeff = self.coeffs[step]
+        curr, nxt = self.sigmas[step : step + 2]
+        if self.name == "euler" or coeff is None:
+            out = minimax_h3_euler_eta0_step(state, denoised, sigma_curr=curr, sigma_next=nxt)
+        else:
+            old = self.old_denoised
+            if old is None or state.shape != denoised.shape or old.shape != state.shape:
+                raise ValueError("H3 solver state and denoised history shapes must match")
+            if not torch.is_floating_point(state) or not torch.is_floating_point(denoised):
+                raise ValueError("H3 solver requires floating point tensors")
+            if not torch.isfinite(state).all() or not torch.isfinite(denoised).all():
+                raise ValueError("H3 solver inputs must be finite")
+            dtype = torch.float32 if state.dtype in (torch.float16, torch.bfloat16) else state.dtype
+            ratio = state.new_tensor(nxt / curr, dtype=dtype)
+            hb1, hb2 = coeff
+            out = (ratio * state.to(dtype) + hb1 * denoised.to(dtype) + hb2 * old.to(dtype)).to(state.dtype)
+            if not torch.isfinite(out).all():
+                raise ValueError("H3 solver output must be finite")
+        if self.name == "res_multistep":
+            self.old_denoised = denoised.detach().clone() if step + 1 < len(self.coeffs) else None
+        self.step_index += 1
+        return out
+
+
+def create_h3_sample_solver(name: str | None, sigmas: Sequence[float]) -> H3SampleSolver:
+    return H3SampleSolver(name, sigmas)