From 9558fd2ea8d470203cc768b9ce4d9111d341c505 Mon Sep 17 00:00:00 2001 From: Mu GuanLin <1203789601@qq.com> Date: Sat, 29 Aug 2026 10:40:50 +0800 Subject: [PATCH] [Bugfix] Scope MiniMax-H3 cuDNN SDPA state (#6710) Signed-off-by: mglyn <1203789601@qq.com> Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com> --- .../test_minimax_h3_encoder_sdpa.py | 62 +++++++++++++++++++ .../diffusion/models/minimax_h3/encoder.py | 20 +++--- 2 files changed, 74 insertions(+), 8 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_encoder_sdpa.py diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_encoder_sdpa.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_encoder_sdpa.py new file mode 100644 index 000000000..b169d918e --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_encoder_sdpa.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from collections.abc import Iterator + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLEncoder + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.fixture(autouse=True) +def restore_cudnn_sdp_state() -> Iterator[None]: + original = torch.backends.cuda.cudnn_sdp_enabled() + try: + yield + finally: + torch.backends.cuda.enable_cudnn_sdp(original) + + +def _encoder_without_weights() -> MiniMaxH3Qwen3VLEncoder: + encoder = MiniMaxH3Qwen3VLEncoder("unused", device=torch.device("cpu"), load_model=False) + encoder.text_model = nn.Linear(1, 1) + return encoder + + +@pytest.mark.parametrize("initial_state", [False, True]) +def test_encode_ids_restores_cudnn_sdp_after_success(monkeypatch: pytest.MonkeyPatch, initial_state: bool) -> None: + encoder = _encoder_without_weights() + expected = torch.tensor([1.0]) + + def encode(*args, **kwargs): + assert torch.backends.cuda.cudnn_sdp_enabled() + return expected + + monkeypatch.setattr(encoder, "_encode", encode) + torch.backends.cuda.enable_cudnn_sdp(initial_state) + + actual = encoder.encode_ids(torch.tensor([1])) + + torch.testing.assert_close(actual, expected) + assert torch.backends.cuda.cudnn_sdp_enabled() is initial_state + + +@pytest.mark.parametrize("initial_state", [False, True]) +def test_encode_ids_restores_cudnn_sdp_after_failure(monkeypatch: pytest.MonkeyPatch, initial_state: bool) -> None: + encoder = _encoder_without_weights() + + def fail(*args, **kwargs): + assert torch.backends.cuda.cudnn_sdp_enabled() + raise RuntimeError("encode failed") + + monkeypatch.setattr(encoder, "_encode", fail) + torch.backends.cuda.enable_cudnn_sdp(initial_state) + + with pytest.raises(RuntimeError, match="encode failed"): + encoder.encode_ids(torch.tensor([1])) + + assert torch.backends.cuda.cudnn_sdp_enabled() is initial_state diff --git a/vllm_omni/diffusion/models/minimax_h3/encoder.py b/vllm_omni/diffusion/models/minimax_h3/encoder.py index 96e488f75..a91cb180b 100644 --- a/vllm_omni/diffusion/models/minimax_h3/encoder.py +++ b/vllm_omni/diffusion/models/minimax_h3/encoder.py @@ -1464,15 +1464,19 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): if next(self.parameters()).device.type != self.device_target.type: raise RuntimeError("call load_to_device() before encode_ids()") + cudnn_sdp_enabled = torch.backends.cuda.cudnn_sdp_enabled() torch.backends.cuda.enable_cudnn_sdp(True) - hidden = self._encode( - input_ids, - pixel_values=pixel_values, - image_grid_thw=image_grid_thw, - pixel_values_videos=pixel_values_videos, - video_grid_thw=video_grid_thw, - ) - return hidden.cpu() + try: + hidden = self._encode( + input_ids, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, + ) + return hidden.cpu() + finally: + torch.backends.cuda.enable_cudnn_sdp(cudnn_sdp_enabled) def forward( self,