test/prototype/mx_formats/test_mx_tensor.py

Author: zxd1997066

torchao/prototype/mx_formats/kernels.py

@@ -38,7 +38,6 @@
 
 logger = logging.getLogger(__name__)
 
-
 def get_bits(x: torch.Tensor) -> str:
     bits_per_byte = 8
     # Numpy has a nice function to get the string representation of binary.
@@ -628,7 +627,8 @@ def triton_f6_e2m3_to_scaled_bf16(
     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"]),)
@@ -671,7 +671,8 @@ def triton_f6_e3m2_to_scaled_bf16(
     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.numel() // packed_mx_block_size
     grid = lambda meta: (triton.cdiv(n_mx_blocks, meta["BLOCK_SIZE_IN"]),)