Repository: huggingface/transformers
Pull request: [XPU] Add MXFP4 support for XPU

# What does this PR do?

Add **MXFP4** quantization support for **XPU** to support loading **GPT-OSS** models in MXFP4 format on XPU.

1. Add MXFP4 support for XPU
2. Improve test file `tests/quantization/mxfp4/test_mxfp4.py` to support XPU testing


Changed files:
filename: src/transformers/integrations/mxfp4.py
patch: @@ -23,6 +23,7 @@
     from accelerate import init_empty_weights
 
 import re
+from contextlib import contextmanager
 
 
 logger = logging.get_logger(__name__)
@@ -47,6 +48,28 @@
 ]
 
 
+@contextmanager
+def on_device(dev):
+    if is_torch_available():
+        import torch
+
+        if isinstance(dev, torch.Tensor):
+            dev = dev.device
+        elif isinstance(dev, str):
+            dev = torch.device(dev)
+        dev_type = getattr(dev, "type", None)
+        if dev_type == "cuda":
+            with torch.cuda.device(dev):
+                yield
+                return
+        if dev_type == "xpu" and hasattr(torch, "xpu"):
+            with torch.xpu.device(dev):
+                yield
+                return
+    # other: CPU
+    yield
+
+
 # Copied from GPT_OSS repo and vllm
 def quantize_to_mxfp4(w, triton_kernels_hub):
     downcast_to_mxfp_torch = triton_kernels_hub.numerics_details.mxfp.downcast_to_mxfp_torch
@@ -173,7 +196,7 @@ def forward(self, hidden_states: torch.Tensor, routing_data, gather_idx, scatter
         )
         swiglu_fn = triton_kernels_hub.swiglu.swiglu_fn
 
-        with torch.cuda.device(hidden_states.device):
+        with on_device(hidden_states.device):
             act = FusedActivation(FnSpecs("swiglu", swiglu_fn, ("alpha", "limit")), (self.alpha, self.limit), 2)
 
             intermediate_cache1 = matmul_ogs(
@@ -214,7 +237,7 @@ def routing_torch_dist(
         triton_kernels_hub.routing.compute_expt_data_torch,
     )
 
-    with torch.cuda.device(logits.device):
+    with on_device(logits.device):
         world_size = torch.distributed.get_world_size()
         rank = int(os.environ.get("LOCAL_RANK", "0"))
         replace_value = -1
@@ -281,7 +304,7 @@ def mlp_forward(self, hidden_states):
     hidden_states = hidden_states.reshape(-1, self.router.hidden_dim)
     router_logits = nn.functional.linear(hidden_states, self.router.weight, self.router.bias)
 
-    with torch.cuda.device(router_logits.device):
+    with on_device(router_logits.device):
         routing_data, gather_idx, scatter_idx = routing(router_logits, self.router.top_k)
 
     routed_out = self.experts(hidden_states, routing_data, gather_idx, scatter_idx)
@@ -376,7 +399,7 @@ def load_and_swizzle_mxfp4(module, param_name, param_value, target_device, trito
             target_device = "cuda"
         blocks = blocks.to(target_device).contiguous()
         scales = scales.to(target_device).contiguous()
-        with torch.cuda.device(target_device):
+        with on_device(target_device):
             triton_weight_tensor, weight_scale = swizzle_mxfp4(
                 blocks.transpose(-2, -1), scales.transpose(-2, -1), triton_kernels_hub
             )

filename: src/transformers/quantizers/quantizer_mxfp4.py
patch: @@ -72,7 +72,7 @@ def validate_environment(self, *args, **kwargs):
         if self.quantization_config.dequantize:
             return
 
-        if not torch.cuda.is_available():
+        if not (torch.cuda.is_available() or torch.xpu.is_available()):
             if self.pre_quantized:
                 logger.warning_once(
                     "Using MXFP4 quantized models requires a GPU, we will default to dequantizing the model to bf16"
@@ -85,43 +85,49 @@ def validate_environment(self, *args, **kwargs):
         if not is_accelerate_available():
             raise ImportError("Using mxfp4 requires Accelerate: `pip install accelerate`")
 
-        compute_capability = torch.cuda.get_device_capability()
-        gpu_is_supported = compute_capability >= (7, 5)
-        kernels_available = is_triton_available("3.4.0") and is_kernels_available()
+        if torch.xpu.is_available():
+            gpu_is_supported = True
+            kernels_available = is_triton_available("3.5.0") and is_kernels_available()
+        else:
+            compute_capability = torch.cuda.get_device_capability()
+            gpu_is_supported = compute_capability >= (7, 5)
+            kernels_available = is_triton_available("3.4.0") and is_kernels_available()
 
         if self.pre_quantized:
             # On unsupported GPUs or without kernels, we will dequantize the model to bf16
             if not gpu_is_supported:
                 logger.warning_once(
-                    "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 (e.g T4, A100, L4, H100, or B200). "
+                    "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 (e.g T4, A100, L4, H100, or B200) or XPUs (e.g Intel® Data Center GPU Max Series) "
                     "We will default to dequantizing the model to bf16."
                 )
                 self.quantization_config.dequantize = True
                 return
 
             if not kernels_available:
                 logger.warning_once(
-                    "MXFP4 quantization requires triton >= 3.4.0 and kernels installed, we will default to dequantizing the model to bf16"
+                    "MXFP4 quantization requires Triton and kernels installed: CUDA requires Triton >= 3.4.0, XPU requires Triton >= 3.5.0, we will default to dequantizing the model to bf16"
                 )
                 self.quantization_config.dequantize = True
                 return
         elif not gpu_is_supported:
             # we can't quantize the model in this case so we raise an error
             raise ValueError(
-                "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 (e.g T4, A100, L4, H100, or B200)"
+                "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 (e.g T4, A100, L4, H100, or B200) or XPUs (e.g Intel® Data Center GPU Max Series) "
             )
         elif not kernels_available:
             # we can't quantize the model in this case so we raise an error
-            raise ValueError("MXFP4 quantization requires triton >= 3.4.0 and kernels installed")
+            raise ValueError(
+                "MXFP4 quantization requires Triton and kernels installed: CUDA requires Triton >= 3.4.0, XPU requires Triton >= 3.5.0"
+            )
 
         if not self.pre_quantized:
             self._lazy_import_kernels()
 
         device_map = kwargs.get("device_map")
         if device_map is None:
             logger.warning_once(
-                "You have loaded an FP4 model on CPU and have a CUDA device available, make sure to set "
-                "your model on a GPU device in order to run your model. To remove this warning, pass device_map = 'cuda'. "
+                "You have loaded an FP4 model on CPU and have a CUDA/XPU device available, make sure to set "
+                "your model on a GPU/XPU device in order to run your model. To remove this warning, pass device_map = 'cuda' or device_map = 'xpu'. "
             )
         elif device_map is not None:
             if (
@@ -263,6 +269,8 @@ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs
         # clean cache due to triton ops
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
+        elif torch.xpu.is_available():
+            torch.xpu.empty_cache()
 
     def update_expected_keys(self, model: "PreTrainedModel", expected_keys: list[str], checkpoint_keys: list[str]):
         # Replace expected_keys for experts' gate_up_proj and down_proj with their _blocks and _scales variants

filename: tests/quantization/mxfp4/test_mxfp4.py
patch: @@ -15,14 +15,15 @@
 import gc
 import tempfile
 import unittest
+from contextlib import ExitStack, contextmanager
 from unittest.mock import patch
 
 from transformers import AutoTokenizer, GptOssForCausalLM, Mxfp4Config
 from transformers.testing_utils import (
     require_kernels,
     require_torch,
     require_torch_gpu,
-    require_torch_large_gpu,
+    require_torch_large_accelerator,
     require_triton,
     slow,
 )
@@ -35,6 +36,30 @@
     import torch
 
 
+if torch.cuda.is_available():
+    REQUIRE_TRITON_MXFP4 = require_triton(min_version="3.4.0")
+elif hasattr(torch, "xpu") and torch.xpu.is_available():
+    REQUIRE_TRITON_MXFP4 = require_triton(min_version="3.5.0")
+else:
+    REQUIRE_TRITON_MXFP4 = unittest.skip("test requires CUDA or XPU")
+
+
+def _empty_accelerator_cache():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    elif hasattr(torch, "xpu") and torch.xpu.is_available():
+        torch.xpu.empty_cache()
+
+
+@contextmanager
+def _patch_no_accelerator():
+    with ExitStack() as stack:
+        stack.enter_context(patch("torch.cuda.is_available", return_value=False))
+        if hasattr(torch, "xpu"):
+            stack.enter_context(patch("torch.xpu.is_available", return_value=False))
+        yield
+
+
 class Mxfp4ConfigTest(unittest.TestCase):
     def test_basic_config_creation(self):
         """Test basic configuration creation with default values"""
@@ -82,8 +107,7 @@ class Mxfp4QuantizerTest(unittest.TestCase):
 
     def setUp(self):
         gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        _empty_accelerator_cache()
 
     def test_quantizer_validation_no_torch(self):
         """Test quantizer validation when torch is not available"""
@@ -96,9 +120,9 @@ def test_quantizer_validation_no_torch(self):
             with self.assertRaises(ImportError):
                 quantizer.validate_environment()
 
-    def test_quantizer_validation_no_cuda(self):
-        """Test quantizer validation when CUDA is not available"""
-        with patch("torch.cuda.is_available", return_value=False):
+    def test_quantizer_validation_no_accelerator(self):
+        """Test quantizer validation when CUDA/XPU is not available"""
+        with _patch_no_accelerator():
             from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer
 
             config = Mxfp4Config()
@@ -108,8 +132,9 @@ def test_quantizer_validation_no_cuda(self):
             with self.assertRaises(RuntimeError):
                 quantizer.validate_environment()
 
+    @require_torch_gpu
     def test_quantizer_validation_low_compute_capability(self):
-        """Test quantizer validation with low compute capability"""
+        """Test quantizer validation with CUDA low compute capability"""
         with patch("torch.cuda.get_device_capability", return_value=(7, 0)):
             from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer
 
@@ -120,8 +145,9 @@ def test_quantizer_validation_low_compute_capability(self):
             with self.assertRaises(ValueError):
                 quantizer.validate_environment()
 
+    @require_torch_gpu
     def test_quantizer_validation_low_compute_capability_with_prequantized(self):
-        """Test quantizer validation with low compute capability"""
+        """Test quantizer validation with CUDA low compute capability"""
         with patch("torch.cuda.get_device_capability", return_value=(7, 0)):
             from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer
 
@@ -132,8 +158,9 @@ def test_quantizer_validation_low_compute_capability_with_prequantized(self):
             quantizer.validate_environment()
             self.assertTrue(quantizer.quantization_config.dequantize)
 
+    @require_torch_gpu
     def test_quantizer_validation_low_compute_capability_with_dequantize(self):
-        """Test quantizer validation with low compute capability but dequantize enabled"""
+        """Test quantizer validation with CUDA low compute capability but dequantize enabled"""
         with patch("torch.cuda.get_device_capability", return_value=(7, 0)):
             from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer
 
@@ -147,20 +174,20 @@ def test_quantizer_validation_low_compute_capability_with_dequantize(self):
                 if "compute capability" in str(e):
                     self.fail("Should not raise compute capability error when dequantize=True")
 
-    def test_quantizer_validation_order_dequantize_before_cuda_check(self):
-        """Test that dequantize check happens before CUDA availability check"""
+    def test_quantizer_validation_order_dequantize_before_accelerator_check(self):
+        """Test that dequantize check happens before CUDA/XPU availability check"""
         # Mock torch.cuda.is_available
-        with patch("torch.cuda.is_available", return_value=False):
+        with _patch_no_accelerator():
             from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer
 
-            # Test with dequantize=True - should pass even without CUDA and accelerate
+            # Test with dequantize=True - should pass even without CUDA/XPU and accelerate
             config = Mxfp4Config(dequantize=True)
             quantizer = Mxfp4HfQuantizer(config)
 
             # This should not raise any error because dequantize check comes first
             quantizer.validate_environment()
 
-            # Test with dequantize=False - should still fail due to missing CUDA
+            # Test with dequantize=False - should still fail due to missing CUDA/XPU
             config = Mxfp4Config(dequantize=False)
             quantizer = Mxfp4HfQuantizer(config)
             quantizer.pre_quantized = False
@@ -313,9 +340,8 @@ def test_convert_moe_packed_tensors(self):
         self.assertEqual(result.shape, (2, 8 * 16 * 2, 4))
         self.assertEqual(result.dtype, torch.bfloat16)
 
-    @require_triton(min_version="3.4.0")
+    @REQUIRE_TRITON_MXFP4
     @require_kernels
-    @require_torch_gpu
     @require_torch
     def test_quantize_to_mxfp4(self):
         """Test quantization function"""
@@ -326,7 +352,8 @@ def test_quantize_to_mxfp4(self):
         quantizer = Mxfp4HfQuantizer(config)
 
         # Create dummy weight tensor
-        w = torch.randn(32, 64, 128, dtype=torch.bfloat16, device=torch.device("cuda"))
+        device = "xpu" if (hasattr(torch, "xpu") and torch.xpu.is_available()) else "cuda"
+        w = torch.randn(32, 64, 128, dtype=torch.bfloat16, device=torch.device(device))
 
         quantized_w, w_scale = quantize_to_mxfp4(w, quantizer._lazy_import_kernels())
 
@@ -335,8 +362,8 @@ def test_quantize_to_mxfp4(self):
 
 
 @require_torch
-@require_torch_large_gpu
-@require_triton(min_version="3.4.0")
+@require_torch_large_accelerator
+@REQUIRE_TRITON_MXFP4
 @require_kernels
 @slow
 class Mxfp4ModelTest(unittest.TestCase):
@@ -353,13 +380,11 @@ class Mxfp4ModelTest(unittest.TestCase):
 
     def setUp(self):
         gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        _empty_accelerator_cache()
 
     def tearDown(self):
         gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        _empty_accelerator_cache()
 
     def check_inference_correctness_quantized(self, model, tokenizer):
         # Check that inference pass works on the model
@@ -454,7 +479,7 @@ def test_save_mxfp4(self):
         with tempfile.TemporaryDirectory() as tmp:
             # Save the model in mxfp4 format
             model.save_pretrained(tmp)
-            torch.cuda.empty_cache()
+            _empty_accelerator_cache()
             gc.collect()
             # test quantized model
             loaded_model = GptOssForCausalLM.from_pretrained(
@@ -486,7 +511,7 @@ def test_save_mxfp4_non_quantized(self):
         # save the quantized model
         with tempfile.TemporaryDirectory() as tmp:
             loaded_model.save_pretrained(tmp)
-            torch.cuda.empty_cache()
+            _empty_accelerator_cache()
             gc.collect()
             # load it back to check with everything works as expected
             loaded_model = GptOssForCausalLM.from_pretrained(