generalize cuda bias ut code

hoshibara · hoshibara · commit d2ee416c5e43 · 2025-05-06T06:30:59.000Z
diff --git a/test/inductor/test_flex_attention.py b/test/inductor/test_flex_attention.py
@@ -417,7 +417,7 @@ def batch_reserve(paged_attention: PagedAttention, target_seq_len: Tensor):
         )
 
 
-@large_tensor_test_class("2GB", device="cuda")
+@large_tensor_test_class("2GB", device=test_device[0])
 class TestFlexAttention(InductorTestCase):
     def setUp(self):
         super().setUp()
@@ -4531,7 +4531,7 @@ def create_inputs(S):
             flex_attention_call(*create_inputs(1024), block_mask=block_mask)
 
 
-@large_tensor_test_class("2GB", device="cuda")
+@large_tensor_test_class("2GB", device=test_device[0])
 class TestPagedAttention(InductorTestCase):
     def setUp(self):
         super().setUp()
@@ -4982,7 +4982,7 @@ def get_params(dtypes: list[torch.dtype]) -> list[Params]:
 
 
 @supports_learnable_bias
-@large_tensor_test_class("2GB", device="cuda")
+@large_tensor_test_class("2GB", device=test_device[0])
 class TestLearnableBiases(InductorTestCase):
     def setUp(self):
         super().setUp()
diff --git a/torch/testing/_internal/common_device_type.py b/torch/testing/_internal/common_device_type.py
@@ -1342,8 +1342,8 @@ def dep_fn(self, *args, **kwargs):
             # an additional array of the same size as the input.
             if inductor and torch._inductor.config.cpp_wrapper and _device != "cpu":
                 size_bytes *= 2
-
-            if not _has_sufficient_memory(_device, size_bytes):
+            # TODO: Memory availability checks for Intel GPU
+            if device != "xpu" and not _has_sufficient_memory(_device, size_bytes):
                 raise unittest.SkipTest(f"Insufficient {_device} memory")
 
             return fn(self, *args, **kwargs)
