From 43e507117f04a86f2df8d405cbe03c1b1642eac9 Mon Sep 17 00:00:00 2001
From: wangyu <53896905+yenuo26@users.noreply.github.com>
Date: Fri, 25 Sep 2026 23:05:37 +0800
Subject: [PATCH] [Bugfix] Restore torch default device between tests (#8157)

Signed-off-by: wangyu <410167048@qq.com>
---
 tests/diffusion/conftest.py     | 15 ---------------
 tests/helpers/fixtures/clean.py | 26 +++++++++++++++++++++++++-
 2 files changed, 25 insertions(+), 16 deletions(-)
 delete mode 100644 tests/diffusion/conftest.py

diff --git a/tests/diffusion/conftest.py b/tests/diffusion/conftest.py
deleted file mode 100644
index e12d15503..000000000
--- a/tests/diffusion/conftest.py
+++ /dev/null
@@ -1,15 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-import torch
-
-
-@pytest.fixture(autouse=True)
-def restore_torch_default_dtype():
-    """Keep process-global dtype changes from leaking between diffusion tests."""
-    default_dtype = torch.get_default_dtype()
-    try:
-        yield
-    finally:
-        torch.set_default_dtype(default_dtype)
diff --git a/tests/helpers/fixtures/clean.py b/tests/helpers/fixtures/clean.py
index c07f563bb..104f960cc 100644
--- a/tests/helpers/fixtures/clean.py
+++ b/tests/helpers/fixtures/clean.py
@@ -1,7 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-"""Opt-in cleanup fixtures for GPU memory, pipeline registry, and speaker cache."""
+"""Cleanup fixtures for process-global test state.
+
+Torch dtype/device restoration is autouse. GPU memory, the pipeline registry,
+and the speaker cache stay opt-in.
+"""
 
 from __future__ import annotations
 
@@ -10,6 +14,26 @@ import pytest
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 
 
+@pytest.fixture(autouse=True)
+def restore_torch_default_dtype_and_device():
+    """Restore process-global dtype and device after every test.
+
+    ``get_default_device()`` reports ``cpu`` both when no default device was
+    set and when ``cpu`` was set explicitly. Writing that value back enables
+    default-device dispatch, so snapshot the active context instead.
+    """
+    import torch
+
+    default_dtype = torch.get_default_dtype()
+    device_context = getattr(torch._GLOBAL_DEVICE_CONTEXT, "device_context", None)
+    default_device = None if device_context is None else torch.get_default_device()
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(default_dtype)
+        torch.set_default_device(default_device)
+
+
 @pytest.fixture
 def clean_gpu_memory_between_tests():
     """Opt-in GPU pre/post hooks for a test (no environment-variable gate).