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