diff --git a/tests/diffusion/models/ltx2/test_ltx2_vae.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py index 964bb7d06cb..b9bc52d9452 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_vae.py +++ b/tests/diffusion/models/ltx2/test_ltx2_vae.py @@ -12,6 +12,29 @@ 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 diff --git a/tests/e2e/accuracy/ltx/run_ltx25_reference.py b/tests/e2e/accuracy/ltx/run_ltx25_reference.py index cbc8917a869..14c862555e9 100644 --- a/tests/e2e/accuracy/ltx/run_ltx25_reference.py +++ b/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) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index 2c51b17d90b..9e682ffef42 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py @@ -65,27 +65,39 @@ 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: + cudnn_context = _deterministic_ltx_vocoder() if device_type == "cuda" else nullcontext() + with cudnn_context: + if not hasattr(vocoder, "bwe_generator"): + return vocoder(generated_mel) + + input_dtype = generated_mel.dtype + module_dtype = next(vocoder.parameters()).dtype + if device_type == "mps": 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) + 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(