[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>
Andreas Karatzas committed
Aug 17, 2026 at 13:32 UTC
c1e438728c55281fda46c8baed755f4148c63660
2 files changed
+22
-4
tests/kernels/attention/test_mha_attn.py
+10
-2
@@ -27,9 +27,17 @@ from vllm.v1.attention.selector import _cached_get_attn_backend
27
28
29
@pytest.fixture(autouse=True)
30
-def clear_cache():
31
- """Clear lru cache to ensure each test case runs without caching."""
30
+def reset_test_state():
31
+ """Clear cached selectors and restore process-wide torch defaults."""
32
+ default_device = torch.get_default_device()
33
+ default_dtype = torch.get_default_dtype()
34
_cached_get_attn_backend.cache_clear()
35
+ try:
36
+ yield
37
+ finally:
38
+ torch.set_default_device(default_device)
39
+ torch.set_default_dtype(default_dtype)
40
+ _cached_get_attn_backend.cache_clear()
41
42
43
devices = ["cpu"]
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+12
-2
@@ -730,9 +730,19 @@ def test_sparse_attn_decode_split_k_kernel(
730
@torch.inference_mode()
731
def test_sparse_attn_decode_gfx950_adaptive_reduce_ignores_stale_scratch() -> None:
732
device = torch.device("cuda")
733
- part_m = torch.full((1, 8, 1), torch.finfo(torch.float32).min, device=device)
733
+ part_m = torch.full(
734
+ (1, 8, 1),
735
+ torch.finfo(torch.float32).min,
736
+ dtype=torch.float32,
737
+ device=device,
738
+ )
739
part_l = torch.zeros_like(part_m)
735
- part_acc = torch.full((1, 8, 1, HEAD_DIM), float("nan"), device=device)
740
+ part_acc = torch.full(
741
+ (1, 8, 1, HEAD_DIM),
742
+ float("nan"),
743
+ dtype=torch.float32,
744
+ device=device,
745
+ )
746
part_m[:, :2] = 0
747
part_l[:, :2] = 1
748
part_acc[:, 0] = 1