[Bugfix] Make LTX vocoder decoding deterministic (#7231)

Signed-off-by: mglyn <1203789601@qq.com>

tests/diffusion/models/ltx2/test_ltx2_vae.py

@@ -12,6 +12,29 @@ import torch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.mark.parametrize("initial_deterministic", [False, True])
+@pytest.mark.parametrize("fail", [False, True])
+def test_ltx_vocoder_deterministic_context(monkeypatch, initial_deterministic, fail):
+    from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _deterministic_ltx_vocoder
+
+    monkeypatch.setattr(torch.backends.cudnn, "deterministic", initial_deterministic)
+    if torch.backends.cudnn.is_available():
+        monkeypatch.setattr(torch.backends.cudnn, "benchmark_limit", 17)
+    settings = {
+        name: getattr(torch.backends.cudnn, name) for name in ("enabled", "benchmark", "benchmark_limit", "allow_tf32")
+    }
+    try:
+        with _deterministic_ltx_vocoder():
+            assert torch.backends.cudnn.deterministic
+            assert {name: getattr(torch.backends.cudnn, name) for name in settings} == settings
+            if fail:
+                raise RuntimeError("injected failure")
+    except RuntimeError as exc:
+        assert fail and str(exc) == "injected failure"
+    assert torch.backends.cudnn.deterministic == initial_deterministic
+    assert {name: getattr(torch.backends.cudnn, name) for name in settings} == settings
+
+
 def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
     import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
 

tests/e2e/accuracy/ltx/run_ltx25_reference.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Generate raw LTX video and audio outputs with the official or Omni runtime."""
 
@@ -114,6 +114,29 @@ def _insert_official_paths(official_root: Path) -> None:
             sys.path.insert(0, path)
 
 
+def _configure_official_vocoder_determinism() -> None:
+    """Make the pinned official LTX-2.5 vocoder a stable reference."""
+    from ltx_core.model.audio_vae.vocoder import VocoderWithBWE
+
+    original_forward = VocoderWithBWE.forward
+    if getattr(original_forward, "_vllm_omni_deterministic", False):
+        return
+
+    def deterministic_forward(self: Any, mel_spec: torch.Tensor) -> torch.Tensor:
+        device_type = mel_spec.device.type
+        if device_type != "cuda":
+            return original_forward(self, mel_spec)
+        previous = torch.backends.cudnn.deterministic
+        try:
+            torch.backends.cudnn.deterministic = True
+            return original_forward(self, mel_spec)
+        finally:
+            torch.backends.cudnn.deterministic = previous
+
+    setattr(deterministic_forward, "_vllm_omni_deterministic", True)
+    setattr(VocoderWithBWE, "forward", deterministic_forward)
+
+
 def _configure_official_sdpa(pipeline: Any) -> None:
     """Pin official connector and denoiser attention to cuDNN, with MATH fallback."""
     from ltx_core.loader.attention_ops import set_attention_module_op
@@ -187,6 +210,7 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
     if args.official_root is None:
         raise ValueError("Official backend requires --official-root")
     _insert_official_paths(args.official_root)
+    _configure_official_vocoder_determinism()
     # Keep Gemma on the same SDPA path in both subprocesses. DiT attention is
     # still pinned independently to the explicit cuDNN backend below.
     torch.backends.cuda.enable_cudnn_sdp(False)

vllm_omni/diffusion/models/ltx2/ltx2_runtime.py

@@ -65,27 +65,39 @@ from .ltx2_request import (
 logger = init_logger(__name__)
 
 
+@contextmanager
+def _deterministic_ltx_vocoder():
+    previous = torch.backends.cudnn.deterministic
+    try:
+        torch.backends.cudnn.deterministic = True
+        yield
+    finally:
+        torch.backends.cudnn.deterministic = previous
+
+
 def _run_ltx_vocoder(vocoder: nn.Module, generated_mel: torch.Tensor) -> torch.Tensor:
     """Run the BWE vocoder in FP32, matching the official LTX pipeline."""
-    if not hasattr(vocoder, "bwe_generator"):
-        return vocoder(generated_mel)
-
-    input_dtype = generated_mel.dtype
     device_type = generated_mel.device.type
-    module_dtype = next(vocoder.parameters()).dtype
-    if device_type == "mps":
-        if module_dtype != torch.float32:
-            vocoder.float()
-        try:
-            return vocoder(generated_mel.float()).to(input_dtype)
-        finally:
-            if module_dtype != torch.float32:
-                vocoder.to(module_dtype)
+    cudnn_context = _deterministic_ltx_vocoder() if device_type == "cuda" else nullcontext()
+    with cudnn_context:
+        if not hasattr(vocoder, "bwe_generator"):
+            return vocoder(generated_mel)
 
-    # BF16 errors compound through the BWE model's long convolution stack.
-    # FP32 autocast upcasts each op without materializing a second FP32 model.
-    with torch.autocast(device_type=device_type, dtype=torch.float32):
-        return vocoder(generated_mel.float()).to(input_dtype)
+        input_dtype = generated_mel.dtype
+        module_dtype = next(vocoder.parameters()).dtype
+        if device_type == "mps":
+            if module_dtype != torch.float32:
+                vocoder.float()
+            try:
+                return vocoder(generated_mel.float()).to(input_dtype)
+            finally:
+                if module_dtype != torch.float32:
+                    vocoder.to(module_dtype)
+
+        # BF16 errors compound through the BWE model's long convolution stack.
+        # FP32 autocast upcasts each op without materializing a second FP32 model.
+        with torch.autocast(device_type=device_type, dtype=torch.float32):
+            return vocoder(generated_mel.float()).to(input_dtype)
 
 
 def _prepare_ltx2_video_output(