Update

drisspg · drisspg · commit 66f3456ec4e8 · 2025-04-23T13:33:25.000-07:00
[ghstack-poisoned]
diff --git a/test/inductor/test_flex_attention.py b/test/inductor/test_flex_attention.py
@@ -106,6 +106,13 @@ def skip_on_cuda(test_func):
     return decorated_func
 
 
+def skip_on_rocm(test_func):
+    """Decorator to skip tests that are not supported on CUDA."""
+    IS_ROCM = torch.cuda.is_available() and torch.version.hip
+    decorated_func = skipCUDAIf(IS_ROCM, "Not supported on ROCM")(test_func)
+    return decorated_func
+
+
 def rmse(ref, res):
     """
     Calculate root mean squared error
@@ -1398,6 +1405,7 @@ def mask_mod(b, h, q, kv):
     @dtypes(*device_configs["cpu"].dtypes_fast)
     @dtypesIfCUDA(*device_configs["cuda"].dtypes_fast)
     @common_utils.parametrize("score_mod", test_score_mods)
+    @skip_on_rocm  # TODO: NaNs on ROCM
     def test_GQA(self, device, dtype: torch.dtype, score_mod: Callable):
         inputs = (
             score_mod,
@@ -1899,6 +1907,7 @@ def f(q, k1, k2, v1, v2):
 
     @supported_platform
     @skip_on_cpu
+    @skip_on_rocm  # TODO: Investigate
     def test_multiple_mask_calls(self, device):
         make_tensor = functools.partial(
             torch.randn,
