Merge branch 'main' into perf/minimax-h3-lossless-hot-path

Signed-off-by: Sy03 <1370724210@qq.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_scheduler_finite_gate.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_scheduler_default_skips_large_finite_scans_and_preserves_results(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        scheduling_minimax_h3_euler_ancestral as scheduler,
+    )
+
+    monkeypatch.delenv("VLLM_OMNI_H3_VALIDATE_FINITE_TENSORS", raising=False)
+    original_isfinite = torch.isfinite
+    scanned_numels = []
+
+    def record_isfinite(tensor):
+        scanned_numels.append(tensor.numel())
+        return original_isfinite(tensor)
+
+    monkeypatch.setattr(torch, "isfinite", record_isfinite)
+    state = torch.linspace(-1.0, 1.0, 8)
+    velocity = torch.linspace(0.25, -0.25, 8)
+    timestep = torch.tensor(0.5)
+    default_x0 = scheduler.minimax_h3_rf_v_to_x0(state, velocity, timestep)
+    default = scheduler.minimax_h3_euler_eta0_step(
+        state,
+        default_x0,
+        sigma_curr=0.5,
+        sigma_next=0.25,
+    )
+
+    assert scanned_numels == [1]
+
+    checked_x0 = scheduler.minimax_h3_rf_v_to_x0(
+        state,
+        velocity,
+        timestep,
+        validate_finite_tensors=True,
+    )
+    checked = scheduler.minimax_h3_euler_eta0_step(
+        state,
+        checked_x0,
+        sigma_curr=0.5,
+        sigma_next=0.25,
+        validate_finite_tensors=True,
+    )
+    torch.testing.assert_close(default_x0, checked_x0, rtol=0, atol=0)
+    torch.testing.assert_close(default, checked, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("invalid_value", [float("nan"), float("inf")])
+def test_scheduler_finite_gate_rejects_nonfinite_inputs(monkeypatch, invalid_value):
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        scheduling_minimax_h3_euler_ancestral as scheduler,
+    )
+
+    monkeypatch.setenv("VLLM_OMNI_H3_VALIDATE_FINITE_TENSORS", "1")
+    invalid = torch.ones(8)
+    invalid[0] = invalid_value
+
+    with pytest.raises(ValueError, match="must be finite"):
+        scheduler.minimax_h3_rf_v_to_x0(
+            invalid,
+            torch.ones_like(invalid),
+            torch.tensor(0.5),
+        )
+    with pytest.raises(ValueError, match="must be finite"):
+        scheduler.minimax_h3_euler_eta0_step(
+            invalid,
+            torch.ones_like(invalid),
+            sigma_curr=0.5,
+            sigma_next=0.25,
+        )

vllm_omni/diffusion/models/minimax_h3/denoise_loop.py

@@ -20,6 +20,7 @@ from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step
 from .scheduling_minimax_h3_euler_ancestral import (
     minimax_h3_euler_eta0_step,
     minimax_h3_rf_v_to_x0,
+    minimax_h3_validate_finite_tensors_enabled,
 )
 
 MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999
@@ -198,6 +199,7 @@ def minimax_h3_denoise_loop(
     if audio_anchor is not None:
         audio_rows[~audio_update] = audio_anchor
 
+    validate_finite_tensors = minimax_h3_validate_finite_tensors_enabled()
     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()
@@ -231,8 +233,15 @@ def minimax_h3_denoise_loop(
                 video_rows[update],
                 mv_video_t,
                 torch.tensor(t_v, dtype=torch.float32, device=device),
+                validate_finite_tensors=validate_finite_tensors,
+            )
+            new_target = minimax_h3_euler_eta0_step(
+                video_rows[update],
+                x0_video,
+                sigma_curr=s_v,
+                sigma_next=s_v_next,
+                validate_finite_tensors=validate_finite_tensors,
             )
-            new_target = minimax_h3_euler_eta0_step(video_rows[update], x0_video, sigma_curr=s_v, sigma_next=s_v_next)
             video_rows = video_rows.clone()
             video_rows[update] = new_target
             if cond_anchor is not None:
@@ -242,9 +251,14 @@ def minimax_h3_denoise_loop(
                 audio_rows[audio_update],
                 mv_audio_t,
                 torch.tensor(t_a, dtype=torch.float32, device=device),
+                validate_finite_tensors=validate_finite_tensors,
             )
             new_audio = minimax_h3_euler_eta0_step(
-                audio_rows[audio_update], x0_audio, sigma_curr=s_a, sigma_next=s_a_next
+                audio_rows[audio_update],
+                x0_audio,
+                sigma_curr=s_a,
+                sigma_next=s_a_next,
+                validate_finite_tensors=validate_finite_tensors,
             )
             audio_rows = audio_rows.clone()
             audio_rows[audio_update] = new_audio

vllm_omni/diffusion/models/minimax_h3/scheduling_minimax_h3_euler_ancestral.py

@@ -2,10 +2,24 @@
 from __future__ import annotations
 
 import math
+import os
 from typing import Any
 
 import torch
 
+MINIMAX_H3_VALIDATE_FINITE_TENSORS_ENV = "VLLM_OMNI_H3_VALIDATE_FINITE_TENSORS"
+
+
+def minimax_h3_validate_finite_tensors_enabled() -> bool:
+    value = os.getenv(MINIMAX_H3_VALIDATE_FINITE_TENSORS_ENV, "")
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _resolve_validate_finite_tensors(value: bool | None) -> bool:
+    if value is None:
+        return minimax_h3_validate_finite_tensors_enabled()
+    return bool(value)
+
 
 def _require_finite_tensor(tensor: torch.Tensor, name: str) -> None:
     if not bool(torch.isfinite(tensor).all().item()):
@@ -50,6 +64,8 @@ def minimax_h3_rf_v_to_x0(
     xt: torch.Tensor,
     v: torch.Tensor,
     timestep: torch.Tensor,
+    *,
+    validate_finite_tensors: bool | None = None,
 ) -> torch.Tensor:
     if xt.shape != v.shape:
         raise ValueError(f"xt and v shapes must match, got {xt.shape} vs {v.shape}")
@@ -57,15 +73,18 @@ def minimax_h3_rf_v_to_x0(
         raise ValueError("xt must be a floating point tensor")
     if not torch.is_floating_point(v):
         raise ValueError("v must be a floating point tensor")
-    _require_finite_tensor(xt, "xt")
+    validate_finite_tensors = _resolve_validate_finite_tensors(validate_finite_tensors)
-    _require_finite_tensor(v, "v")
+    if validate_finite_tensors:
+        _require_finite_tensor(xt, "xt")
+        _require_finite_tensor(v, "v")
     _validate_unit_timestep(timestep, "timestep")
     cond_t = timestep.to(device=xt.device, dtype=xt.dtype)
     while cond_t.ndim < xt.ndim:
         cond_t = cond_t.unsqueeze(-1)
     sigma_t = 1 - cond_t
     x0 = xt + sigma_t * v
-    _require_finite_tensor(x0, "x0")
+    if validate_finite_tensors:
+        _require_finite_tensor(x0, "x0")
     return x0
 
 
@@ -75,6 +94,7 @@ def minimax_h3_euler_eta0_step(
     *,
     sigma_curr: float,
     sigma_next: float,
+    validate_finite_tensors: bool | None = None,
 ) -> torch.Tensor:
     if state.shape != denoised.shape:
         raise ValueError(f"state and denoised shapes must match, got {state.shape} vs {denoised.shape}")
@@ -82,8 +102,10 @@ def minimax_h3_euler_eta0_step(
         raise ValueError("state must be a floating point tensor")
     if not torch.is_floating_point(denoised):
         raise ValueError("denoised must be a floating point tensor")
-    _require_finite_tensor(state, "state")
+    validate_finite_tensors = _resolve_validate_finite_tensors(validate_finite_tensors)
-    _require_finite_tensor(denoised, "denoised")
+    if validate_finite_tensors:
+        _require_finite_tensor(state, "state")
+        _require_finite_tensor(denoised, "denoised")
     sigma_curr = _validate_sigma(sigma_curr, "sigma_curr")
     sigma_next = _validate_sigma(sigma_next, "sigma_next")
     if sigma_curr == 0.0:
@@ -98,7 +120,8 @@ def minimax_h3_euler_eta0_step(
     sigma_ratio = sigma_next_t / sigma_curr_t
     out = sigma_ratio * state.to(dtype=compute_dtype) + (1.0 - sigma_ratio) * denoised.to(dtype=compute_dtype)
     out = out.to(dtype=state.dtype)
-    _require_finite_tensor(out, "euler_eta0_step output")
+    if validate_finite_tensors:
+        _require_finite_tensor(out, "euler_eta0_step output")
     return out
 
 
@@ -127,6 +150,7 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter:
         audio_sigma_curr: float | None = None,
         audio_sigma_next: float | None = None,
     ) -> dict[str, torch.Tensor]:
+        validate_finite_tensors = minimax_h3_validate_finite_tensors_enabled()
         visual_timestep = timestep if video_timestep is None else video_timestep
         audio_timestep = timestep if audio_timestep is None else audio_timestep
         visual_sigma_curr = sigma_curr if video_sigma_curr is None else video_sigma_curr
@@ -148,11 +172,13 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter:
             input_visual_latent,
             noise_pred_visual,
             visual_timestep,
+            validate_finite_tensors=validate_finite_tensors,
         )
         denoised_audio = minimax_h3_rf_v_to_x0(
             input_audio_latent,
             noise_pred_audio,
             audio_timestep,
+            validate_finite_tensors=validate_finite_tensors,
         )
         return {
             "output_visual_latent": minimax_h3_euler_eta0_step(
@@ -160,12 +186,14 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter:
                 denoised_visual,
                 sigma_curr=visual_sigma_curr,
                 sigma_next=visual_sigma_next,
+                validate_finite_tensors=validate_finite_tensors,
             ),
             "output_audio_latent": minimax_h3_euler_eta0_step(
                 input_audio_latent,
                 denoised_audio,
                 sigma_curr=audio_sigma_curr,
                 sigma_next=audio_sigma_next,
+                validate_finite_tensors=validate_finite_tensors,
             ),
         }
 
@@ -173,7 +201,9 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter:
 EntryClass = MiniMaxH3EulerAncestralEta0SchedulerAdapter
 
 __all__ = [
+    "MINIMAX_H3_VALIDATE_FINITE_TENSORS_ENV",
     "MiniMaxH3EulerAncestralEta0SchedulerAdapter",
     "minimax_h3_euler_eta0_step",
     "minimax_h3_rf_v_to_x0",
+    "minimax_h3_validate_finite_tensors_enabled",
 ]