[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

@@ -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"]

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