From fef3a6b63bc60ce778fd9848a70b61ee94f06539 Mon Sep 17 00:00:00 2001
From: Ben Barsdell <bbarsdell@nvidia.com>
Date: Tue, 4 Nov 2025 18:51:23 +1100
Subject: [PATCH] Restore torch defaults between sgl-kernel tests (#11131)

---
 sgl-kernel/tests/conftest.py                  | 14 ++++++++++++++
 sgl-kernel/tests/test_torch_defaults_reset.py | 16 ++++++++++++++++
 2 files changed, 30 insertions(+)
 create mode 100644 sgl-kernel/tests/conftest.py
 create mode 100644 sgl-kernel/tests/test_torch_defaults_reset.py

diff --git a/sgl-kernel/tests/conftest.py b/sgl-kernel/tests/conftest.py
new file mode 100644
index 000000000..6682bc0c9
--- /dev/null
+++ b/sgl-kernel/tests/conftest.py
@@ -0,0 +1,14 @@
+import pytest
+import torch
+
+
+# This ensures the torch defaults don't get left in modified states between
+# tests (e.g., when a test fails before restoring the original value), which
+# can cause subsequent tests to fail.
+@pytest.fixture(autouse=True)
+def reset_torch_defaults():
+    orig_default_device = torch.get_default_device()
+    orig_default_dtype = torch.get_default_dtype()
+    yield
+    torch.set_default_dtype(orig_default_dtype)
+    torch.set_default_device(orig_default_device)
diff --git a/sgl-kernel/tests/test_torch_defaults_reset.py b/sgl-kernel/tests/test_torch_defaults_reset.py
new file mode 100644
index 000000000..f6fae5d9e
--- /dev/null
+++ b/sgl-kernel/tests/test_torch_defaults_reset.py
@@ -0,0 +1,16 @@
+import pytest
+import torch
+
+
+def test_change_torch_defaults():
+    torch.set_default_device("cpu:0")
+    torch.set_default_dtype(torch.float16)
+
+
+def test_check_torch_defaults():
+    assert torch.get_default_device() == torch.device("cpu")
+    assert torch.get_default_dtype() == torch.float32
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])