test/prototype/mx_formats/test_kernels.py

Author: zxd1997066

test/prototype/mx_formats/test_kernels.py

@@ -518,7 +518,7 @@ def test_triton_mxfp8_dim0_zeros():
 
 @pytest.mark.skipif(not has_triton(), reason="unsupported without triton")
 @pytest.mark.skipif(
-    torch.cuda.is_available() and not is_sm_at_least_100(),
+    not is_sm_at_least_100(),
     reason="mxfp8 requires CUDA capability 10.0 or greater",
 )
 @pytest.mark.parametrize("M", (256, 2048, 131072))

torchao/prototype/mx_formats/kernels.py

@@ -551,7 +551,8 @@ def triton_f6_e2m3_to_bf16(x: torch.Tensor) -> torch.Tensor:
     output = torch.empty(*new_shape, device=x.device, dtype=torch.bfloat16)
 
     assert x.is_contiguous()
-    assert x.is_cuda and output.is_cuda
+    assert x.is_cuda or x.device.type == "xpu"
+    assert output.is_cuda or output.device.type == "xpu"
 
     n_mx_blocks = x.shape[0]
     grid = lambda meta: (triton.cdiv(n_mx_blocks, meta["BLOCK_SIZE_IN"]),)
@@ -587,7 +588,9 @@ def triton_f6_e3m2_to_bf16(x: torch.Tensor) -> torch.Tensor:
     output = torch.empty(*new_shape, device=x.device, dtype=torch.bfloat16)
 
     assert x.is_contiguous()
-    assert x.is_cuda and output.is_cuda
+    assert x.is_cuda or x.device.type == "xpu"
+    assert output.is_cuda or output.device.type == "xpu"
+
 
     n_mx_blocks = x.shape[0]
     grid = lambda meta: (triton.cdiv(n_mx_blocks, meta["BLOCK_SIZE_IN"]),)