fix codelint

test/inductor/test_online_softmax.py

@@ -14,7 +14,7 @@ from torch.testing._internal.common_utils import (
     IS_LINUX,
     parametrize,
 )
-from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_CUDA_AND_TRITON
+from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU, HAS_TRITON
 
 
 DO_PERF_TEST = os.environ.get("DO_PERF_TEST") == "1"
@@ -137,7 +137,9 @@ class TestOnlineSoftmax(TestCase):
         ref = _prepare_softmax(x, dim)
         self.assertTrue(same(ref, act, tol=1e-2))
 
-        if nrow == 2048 and dim == 0:
+        if nrow == 2048 and dim == 0 and GPU_TYPE != "xpu":
+            # Note: split reduction is not triggered for this shape on xpu devices.
+            #       check "num_splits" for more details
             # split reduction is triggered. We have multiple kernels
             self.assertTrue(code.count("def triton") >= 2)
         else:
@@ -310,5 +312,5 @@ class TestOnlineSoftmax(TestCase):
 instantiate_parametrized_tests(TestOnlineSoftmax)
 
 if __name__ == "__main__":
-    if IS_LINUX and HAS_CUDA_AND_TRITON:
+    if IS_LINUX and HAS_GPU and HAS_TRITON:
         run_tests()

torch/_inductor/fx_passes/post_grad.py

@@ -351,10 +351,9 @@ def prepare_softmax_extra_check(match):
     """
     We only have triton online softmax kernels currently.
     """
-    return (
-        config.online_softmax
-        and match.kwargs["x"].meta["val"].device.type == "cuda"
-        and config.cuda_backend == "triton"
+    return config.online_softmax and (
+        match.kwargs["x"].meta["val"].device.type in ["cuda", "xpu"]
+        and "triton" in [config.cuda_backend, config.xpu_backend]
     )