From c1e438728c55281fda46c8baed755f4148c63660 Mon Sep 17 00:00:00 2001
From: Andreas Karatzas <akaratza@amd.com>
Date: Mon, 17 Aug 2026 13:32:04 -0500
Subject: [PATCH] [ROCm][CI] Restore Torch defaults and type DSV4 scratch
 buffers (#52566)

Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
---
 tests/kernels/attention/test_mha_attn.py           | 12 ++++++++++--
 .../attention/test_rocm_triton_attn_dsv4.py        | 14 ++++++++++++--
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/tests/kernels/attention/test_mha_attn.py b/tests/kernels/attention/test_mha_attn.py
index d73acfc0ee..b6578fa2f9 100644
--- a/tests/kernels/attention/test_mha_attn.py
+++ b/tests/kernels/attention/test_mha_attn.py
@@ -27,9 +27,17 @@ from vllm.v1.attention.selector import _cached_get_attn_backend
 
 
 @pytest.fixture(autouse=True)
-def clear_cache():
-    """Clear lru cache to ensure each test case runs without caching."""
+def reset_test_state():
+    """Clear cached selectors and restore process-wide torch defaults."""
+    default_device = torch.get_default_device()
+    default_dtype = torch.get_default_dtype()
     _cached_get_attn_backend.cache_clear()
+    try:
+        yield
+    finally:
+        torch.set_default_device(default_device)
+        torch.set_default_dtype(default_dtype)
+        _cached_get_attn_backend.cache_clear()
 
 
 devices = ["cpu"]
diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py
index 7a98f3ec70..2f1e8d2793 100644
--- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py
@@ -730,9 +730,19 @@ def test_sparse_attn_decode_split_k_kernel(
 @torch.inference_mode()
 def test_sparse_attn_decode_gfx950_adaptive_reduce_ignores_stale_scratch() -> None:
     device = torch.device("cuda")
-    part_m = torch.full((1, 8, 1), torch.finfo(torch.float32).min, device=device)
+    part_m = torch.full(
+        (1, 8, 1),
+        torch.finfo(torch.float32).min,
+        dtype=torch.float32,
+        device=device,
+    )
     part_l = torch.zeros_like(part_m)
-    part_acc = torch.full((1, 8, 1, HEAD_DIM), float("nan"), device=device)
+    part_acc = torch.full(
+        (1, 8, 1, HEAD_DIM),
+        float("nan"),
+        dtype=torch.float32,
+        device=device,
+    )
     part_m[:, :2] = 0
     part_l[:, :2] = 1
     part_acc[:, 0] = 1