Repository: huggingface/transformers
Pull request: XPU now supports MoE kernel(MegaBlocks) implementation

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

XPU now supports the **MoE kernel(MegaBlocks)** implementation, and includes XPU checks for GPT-OSS UTs compatibility.




<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Changed files:
filename: src/transformers/integrations/hub_kernels.py
patch: @@ -152,6 +152,12 @@ def use_kernel_func_from_hub(func_name: str):
                     layer_name="MegaBlocksMoeMLP",
                 )
             },
+            "xpu": {
+                Mode.INFERENCE: LayerRepository(
+                    repo_id="kernels-community/megablocks",
+                    layer_name="MegaBlocksMoeMLP",
+                )
+            },
         },
         "FastGELU": {
             "cuda": {

filename: tests/fixtures/gpt_oss/integration_tests.json
patch: @@ -1,122 +1,154 @@
 {
-  "quantized=true|model=120b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
+  "device=cuda|quantized=true|model=120b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
     "Roses are red, violets are blue, I am a language model, and I can help you too!\n\nSure! Here",
     "How are you? Tell me the name of the president of the United\n\nHello! As of my last update in November 2023, the President of the"
   ],
-  "quantized=true|model=120b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
+  "device=cuda|quantized=true|model=120b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
     "Roses are red, violets are blue, I am a language model, and I can help you too!\n\nSure! Here",
     "How are you? Tell me the name of the president of the United\n\nHello! As of my last update in November 2023, the President of the"
   ],
-  "quantized=true|model=120b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
+  "device=cuda|quantized=true|model=120b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
     "Did not work"
   ],
-  "quantized=true|model=120b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
+  "device=cuda|quantized=true|model=120b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
     "Did not work"
   ],
-  "quantized=true|model=120b|kernels=false|attn_impl=eager|mode=eval": [
+  "device=cuda|quantized=true|model=120b|kernels=false|attn_impl=eager|mode=eval": [
     "Roses are red, violets are blue, I am a language model, and I can help you too!\n\nSure! Here",
     "How are you? Tell me the name of the president of the United\n\nHello! As of my last update in November 2023, the President of the"
   ],
-  "quantized=true|model=120b|kernels=false|attn_impl=eager|mode=train": [
+  "device=cuda|quantized=true|model=120b|kernels=false|attn_impl=eager|mode=train": [
     "Roses are red, violets are blue, I am a language model, and I can help you too!\n\nSure! Here",
     "How are you? Tell me the name of the president of the United\n\nHello! As of my last update in November 2023, the President of the"
   ],
-  "quantized=true|model=120b|kernels=true|attn_impl=eager|mode=eval": [
+  "device=cuda|quantized=true|model=120b|kernels=true|attn_impl=eager|mode=eval": [
     "Did not work"
   ],
-  "quantized=true|model=120b|kernels=true|attn_impl=eager|mode=train": [
+  "device=cuda|quantized=true|model=120b|kernels=true|attn_impl=eager|mode=train": [
     "Did not work"
   ],
-  "quantized=true|model=20b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
+  "device=cuda|quantized=true|model=20b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
     "Roses are red, violets are blue, I love you, and I love you too.\nIt sounds like you're looking for",
     "How are you? Tell me the name of the president of the United States.\" The assistant should respond with the name of the president. The user is asking for"
   ],
-  "quantized=true|model=20b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
+  "device=cuda|quantized=true|model=20b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
     "Roses are red, violets are blue, I love you, and I love you too.\n\nIt sounds like you're looking for",
     "How are you? Tell me the name of the president of the United States.\" The assistant should respond with the name of the president. The user is asking for"
   ],
-  "quantized=true|model=20b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
+  "device=cuda|quantized=true|model=20b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
     "Did not work"
   ],
-  "quantized=true|model=20b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
+  "device=cuda|quantized=true|model=20b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
     "Did not work"
   ],
-  "quantized=true|model=20b|kernels=false|attn_impl=eager|mode=eval": [
+  "device=cuda|quantized=true|model=20b|kernels=false|attn_impl=eager|mode=eval": [
     "Roses are red, violets are blue, I love you, and I love you too.\n\nIt sounds like you're expressing a",
     "How are you? Tell me the name of the president of the United States.\" The assistant should respond with the name of the president. The user is asking for"
   ],
-  "quantized=true|model=20b|kernels=false|attn_impl=eager|mode=train": [
+  "device=cuda|quantized=true|model=20b|kernels=false|attn_impl=eager|mode=train": [
     "Roses are red, violets are blue, I love you, and I love you too.\n\nIt sounds like you're expressing a",
     "How are you? Tell me the name of the president of the United States.\" The assistant should respond with the name of the president. The user is asking for"
   ],
-  "quantized=true|model=20b|kernels=true|attn_impl=eager|mode=eval": [
+  "device=cuda|quantized=true|model=20b|kernels=true|attn_impl=eager|mode=eval": [
     "Did not work"
   ],
-  "quantized=true|model=20b|kernels=true|attn_impl=eager|mode=train": [
+  "device=cuda|quantized=true|model=20b|kernels=true|attn_impl=eager|mode=train": [
     "Did not work"
   ],
-  "quantized=false|model=120b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
+  "device=cuda|quantized=false|model=120b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
     "Roses are red, violets are blue,\nI am a language model, not a human being.\n```\n\nThis poem is a",
     "How are you? Tell me the name of the president of the United Kingdom?\n\nThe United Kingdom does not have a president. The head of state is the"
   ],
-  "quantized=false|model=120b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
+  "device=cuda|quantized=false|model=120b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
     "Roses are red, violets are blue, I am a language model trained by OpenAI.\n\nI am a large language model",
     "How are you? Tell me the name of the president of the United\n\nHello! I'm an AI language model, so I don't have feelings, but I'm here"
   ],
-  "quantized=false|model=120b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
+  "device=cuda|quantized=false|model=120b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
     "Roses are red, violets are blue,\nI am a language model, not a human being.\n```\n\nThis poem is a",
     "How are you? Tell me the name of the president of the United Kingdom?\n\nThe United Kingdom does not have a president. The head of state is the"
   ],
-  "quantized=false|model=120b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
+  "device=cuda|quantized=false|model=120b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
     "Roses are red, violets are blue, I am a language model trained by OpenAI.\n\nI am a large language model",
     "How are you? Tell me the name of the president of the United\n\nHello! I'm an AI language model, so I don't have feelings, but I'm here"
   ],
-  "quantized=false|model=120b|kernels=false|attn_impl=eager|mode=eval": [
+  "device=cuda|quantized=false|model=120b|kernels=false|attn_impl=eager|mode=eval": [
     "Roses are red, violets are blue,\nI am a language model, not a human being.\n```\n\nThis poem is a",
     "How are you? Tell me the name of the president of the United States?\n\nAs an AI language model, I do not have personal feelings or emotions,"
   ],
-  "quantized=false|model=120b|kernels=false|attn_impl=eager|mode=train": [
+  "device=cuda|quantized=false|model=120b|kernels=false|attn_impl=eager|mode=train": [
     "Roses are red, violets are blue, I am a language model, and I can help you with your request.\n\nSure",
     "How are you? Tell me the name of the president of the United\n\nHello! I'm an AI language model, so I don't have feelings, but I'm here"
   ],
-  "quantized=false|model=120b|kernels=true|attn_impl=eager|mode=eval": [
+  "device=cuda|quantized=false|model=120b|kernels=true|attn_impl=eager|mode=eval": [
     "Roses are red, violets are blue,\nI am a language model, not a human being.\n```\n\nThis poem is a",
     "How are you? Tell me the name of the president of the United States?\n\nAs an AI language model, I do not have personal feelings or emotions,"
   ],
-  "quantized=false|model=120b|kernels=true|attn_impl=eager|mode=train": [
+  "device=cuda|quantized=false|model=120b|kernels=true|attn_impl=eager|mode=train": [
     "Roses are red, violets are blue, I am a language model, and I can help you with your request.\n\nSure",
     "How are you? Tell me the name of the president of the United\n\nHello! I'm an AI language model, so I don't have feelings, but I'm here"
   ],
-  "quantized=false|model=20b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
+  "device=cuda|quantized=false|model=20b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
     "Roses are red, violets are blue, I love you, and I love you too!\n\nRoses are red, vio",
     "How are you? Tell me the name of the president of the United States.\" The assistant should respond with the name of the president. The user is asking for"
   ],
-  "quantized=false|model=20b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
+  "device=cuda|quantized=false|model=20b|kernels=false|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
     "Roses are red, violets are blue\" (makes sense). But the phrase \"the answer is 3\" is not a",
     "How are you? Tell me the name of the president of the United States.\" The answer to that is \"Joe Biden.\" The user is asking for the name"
   ],
-  "quantized=false|model=20b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
+  "device=cuda|quantized=false|model=20b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=eval": [
     "Roses are red, violets are blue, I love you, and I love you too!\n\nRoses are red, vio",
     "How are you? Tell me the name of the president of the United States.\" The assistant should respond with the name of the president. The user is asking for"
   ],
-  "quantized=false|model=20b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
+  "device=cuda|quantized=false|model=20b|kernels=true|attn_impl=kernels-community/vllm-flash-attn3|mode=train": [
     "Roses are red, violets are blue\" (makes sense). But the phrase \"the answer is 3\" is not a",
     "How are you? Tell me the name of the president of the United States.\" The answer to that is \"Joe Biden.\" The user is asking for the name"
   ],
-  "quantized=false|model=20b|kernels=false|attn_impl=eager|mode=eval": [
+  "device=cuda|quantized=false|model=20b|kernels=false|attn_impl=eager|mode=eval": [
     "Roses are red, violets are blue, I love you, and I love you too!\n\nRoses are red, vio",
     "How are you? Tell me the name of the president of the United States.\" The assistant should respond with the name of the president. The user is asking for"
   ],
-  "quantized=false|model=20b|kernels=false|attn_impl=eager|mode=train": [
+  "device=cuda|quantized=false|model=20b|kernels=false|attn_impl=eager|mode=train": [
     "Roses are red, violets are blue.\" -> from which we can derive a rule: if we have a red object that is",
     "How are you? Tell me the name of the president of the United States.\n\nI am an AI language model and I do not have a personal life or"
   ],
-  "quantized=false|model=20b|kernels=true|attn_impl=eager|mode=eval": [
+  "device=cuda|quantized=false|model=20b|kernels=true|attn_impl=eager|mode=eval": [
     "Roses are red, violets are blue, I love you, and I love you too!\n\nRoses are red, vio",
     "How are you? Tell me the name of the president of the United States.\" The assistant should respond with the name of the president. The user is asking for"
   ],
-  "quantized=false|model=20b|kernels=true|attn_impl=eager|mode=train": [
+  "device=cuda|quantized=false|model=20b|kernels=true|attn_impl=eager|mode=train": [
     "Roses are red, violets are blue.\" -> from which we can derive a rule: if we have a red object that is",
     "How are you? Tell me the name of the president of the United States.\n\nI am an AI language model and I do not have a personal life or"
+  ],
+  "device=xpu|quantized=false|model=20b|kernels=false|attn_impl=eager|mode=eval": [
+    "Roses are red, violets are blue, I love you, and I love you too!\n\nRoses are red, vio",
+    "How are you? Tell me the name of the president of the United States.\" The assistant should respond with the name of the president. The user is asking for"
+  ],
+  "device=xpu|quantized=false|model=20b|kernels=false|attn_impl=eager|mode=train": [
+    "Roses are red, violets are blue, I love you, and I love you too!\n\nRoses are red, vio",
+    "How are you? Tell me the name of the president of the United States.\" The assistant should respond with the name of the president. The user is asking for"
+  ],
+  "device=xpu|quantized=false|model=20b|kernels=true|attn_impl=eager|mode=eval": [
+    "Roses are red, violets are blue, I love you, and I love you too!\n\nRoses are red, vio",
+    "How are you? Tell me the name of the president of the United States.\" The assistant should respond with the name of the president. The user is asking for"
+  ],
+  "device=xpu|quantized=false|model=20b|kernels=true|attn_impl=eager|mode=train": [
+    "Roses are red, violets are blue, I love you, and I love you too!\n\nRoses are red, vio",
+    "How are you? Tell me the name of the president of the United States.\" The assistant should respond with the name of the president. The user is asking for"
+  ],
+  "device=xpu|quantized=false|model=120b|kernels=false|attn_impl=eager|mode=eval": [
+    "Roses are red, violets are blue,\nI am a language model, not a human being.\n```\n\nThis poem is a",
+    "How are you? Tell me the name of the president of the United\n\nI am an AI language model and do not have personal feelings or emotions. As for"
+  ],
+  "device=xpu|quantized=false|model=120b|kernels=false|attn_impl=eager|mode=train": [
+    "Roses are red, violets are blue,\nI am a language model, not a human being.\n```\n\nThis poem is a",
+    "How are you? Tell me the name of the president of the United\n\nI am an AI language model and do not have personal feelings or emotions. As for"
+  ],
+  "device=xpu|quantized=false|model=120b|kernels=true|attn_impl=eager|mode=eval": [
+    "Roses are red, violets are blue,\nI am a language model, not a human being.\n```\n\nThis poem is a",
+    "How are you? Tell me the name of the president of the United Kingdom?\n\nThe United Kingdom does not have a president. The head of state is the"
+  ],
+  "device=xpu|quantized=false|model=120b|kernels=true|attn_impl=eager|mode=train": [
+    "Roses are red, violets are blue,\nI am a language model, not a human being.\n```\n\nThis poem is a",
+    "How are you? Tell me the name of the president of the United Kingdom?\n\nThe United Kingdom does not have a president. The head of state is the"
   ]
-}
+}
\ No newline at end of file

filename: tests/models/gpt_oss/test_modeling_gpt_oss.py
patch: @@ -32,6 +32,7 @@
 )
 from transformers.testing_utils import (
     cleanup,
+    require_deterministic_for_xpu,
     require_flash_attn,
     require_kernels,
     require_torch,
@@ -51,7 +52,12 @@
         GptOssModel,
     )
 
-    NUM_GPUS = torch.cuda.device_count()
+    if torch.cuda.is_available():
+        NUM_GPUS = torch.cuda.device_count()
+    elif hasattr(torch, "xpu") and torch.xpu.is_available():
+        NUM_GPUS = torch.xpu.device_count()
+    else:
+        NUM_GPUS = 0
 
 
 class GptOssModelTester(CausalLMModelTester):
@@ -139,7 +145,7 @@ def distributed_worker(quantized, model_size, kernels, attn_impl, mode):
 
     def generate_config_key(quantized, model, kernels, attn_impl, mode):
         """Generate a key for the restructured integration test results."""
-        return f"quantized={str(quantized).lower()}|model={model}|kernels={str(kernels).lower()}|attn_impl={attn_impl}|mode={mode}"
+        return f"device={torch_device}|quantized={str(quantized).lower()}|model={model}|kernels={str(kernels).lower()}|attn_impl={attn_impl}|mode={mode}"
 
     input_text = [
         "Roses are red, violets",
@@ -227,7 +233,7 @@ class GptOssIntegrationTest(unittest.TestCase):
     @staticmethod
     def generate_config_key(quantized, model, kernels, attn_impl, mode):
         """Generate a key for the restructured integration test results."""
-        return f"quantized={str(quantized).lower()}|model={model}|kernels={str(kernels).lower()}|attn_impl={attn_impl}|mode={mode}"
+        return f"device={torch_device}|quantized={str(quantized).lower()}|model={model}|kernels={str(kernels).lower()}|attn_impl={attn_impl}|mode={mode}"
 
     def setUp(self):
         cleanup(torch_device, gc_collect=True)
@@ -345,7 +351,11 @@ def run_distributed_test(quantized, model, kernels, attn_impl, mode):
     # Non-distributed test
     # ------------------------
     @parameterized.expand(PARAMETERS)
+    @require_deterministic_for_xpu
     def test_model_outputs(self, quantized, model, kernels, attn_impl, mode):
+        if torch_device == "xpu" and attn_impl == "kernels-community/vllm-flash-attn3":
+            self.skipTest("flash attention 3 is not supported on XPU yet.")
+
         model_id = f"openai/gpt-oss-{model}"
         output_texts = self.load_and_forward(
             model_id,
@@ -408,6 +418,9 @@ def test_model_outputs(self, quantized, model, kernels, attn_impl, mode):
     # ------------------------
     @parameterized.expand(PARAMETERS)
     def test_model_outputs_distributed(self, quantized, model, kernels, attn_impl, mode):
+        if torch_device == "xpu" and attn_impl == "kernels-community/vllm-flash-attn3":
+            self.skipTest("flash attention 3 is not supported on XPU yet.")
+
         self.run_distributed_test(quantized, model, kernels, attn_impl, mode)
 
     # ------------------------