Repository: huggingface/transformers
Pull request: Fix bug and add XPU Expectations for qwen2 and jamba tests

@IlyasMoutawwakil, pls help review, thx!

Changed files:
filename: src/transformers/modeling_flash_attention_utils.py
patch: @@ -171,8 +171,10 @@ def _lazy_imports(
         else:
             from .integrations.hub_kernels import load_and_register_attn_kernel
 
+            # Map standard attention names to hub kernel repos
+            kernel_repo = FLASH_ATTN_KERNEL_FALLBACK.get(implementation, implementation)
             # We want to explicitly register the name with `paged|` if found
-            kernel_implementation = f"paged|{implementation}" if is_paged else implementation
+            kernel_implementation = f"paged|{implementation}" if is_paged else kernel_repo
             kernel = load_and_register_attn_kernel(
                 kernel_implementation, attention_wrapper, allow_all_kernels=allow_all_kernels
             )

filename: tests/models/jamba/test_modeling_jamba.py
patch: @@ -585,6 +585,7 @@ def test_simple_generate(self):
                 ("cuda", 7): "
Hey how are you doing on this lovely evening? Canyon rins hugaughter glamour Rutgers Singh
cw algunas",
                 ("cuda", 8): "
Hey how are you doing on this lovely evening? I'm so glad you're here.",
                 ("rocm", 9): "
Hey how are you doing on this lovely evening? Canyon rins hugaughter glamour Rutgers Singh Hebrew llam bb",
+                ("xpu", 3): "
Hey how are you doing on this lovely evening? I'm so glad you're here.",
             }
         )
         # fmt: on
@@ -609,6 +610,7 @@ def test_simple_batched_generate_with_padding(self):
                 ("cuda", 7): ["
Hey how are you doing on this lovely evening? Canyon rins hugaughter glamour Rutgers Singh Hebrew cases Cats", "






Tell me a storyptus Nets Madison El chamadamodern updximVaparsed",],
                 ("cuda", 8): ["
Hey how are you doing on this lovely evening? I'm so glad you're here.", "






Tell me a story about a woman who was born in the United States",],
                 ("rocm", 9): ["
Hey how are you doing on this lovely evening? Canyon rins hugaughter glamour Rutgers Singh
cw algunas", "






Tell me a storyptus Nets Madison El chamadamodern updximVaparsed",],
+                ("xpu", 3): ["
Hey how are you doing on this lovely evening? I'm so glad you're here.", "
Tell me a story





, I'm not sure, but I'"]
             }
         )
         # fmt: on

filename: tests/models/qwen2/test_modeling_qwen2.py
patch: @@ -25,6 +25,7 @@
     Expectations,
     backend_empty_cache,
     require_bitsandbytes,
+    require_deterministic_for_xpu,
     require_flash_attn,
     require_torch,
     slow,
@@ -71,18 +72,26 @@ def is_pipeline_test_to_skip(
 @require_torch
 class Qwen2IntegrationTest(unittest.TestCase):
     @slow
+    @require_deterministic_for_xpu
     def test_model_450m_logits(self):
         input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
         model = Qwen2ForCausalLM.from_pretrained("Qwen/Qwen2-0.5B", device_map="auto")
         input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
         with torch.no_grad():
             out = model(input_ids).logits.float().cpu()
         # Expected mean on dim = -1
-        EXPECTED_MEAN = torch.tensor([[-2.2121, -1.6335, -1.4816, -1.5035, -1.9110, -1.8979, -1.9682, -2.1980]])
-        torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, rtol=1e-2, atol=1e-2)
+        EXPECTED_MEAN = Expectations({
+            (None, None): torch.tensor([[-2.2121, -1.6335, -1.4816, -1.5035, -1.9110, -1.8979, -1.9682, -2.1980]]),
+            ("xpu", 3): torch.tensor([[-2.2419, -1.6216, -1.4517, -1.4963, -1.9229, -1.8966, -1.9580, -2.1484]]),
+        })  # fmt: off
+
+        torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN.get_expectation(), rtol=1e-2, atol=1e-2)
         # slicing logits[0, 0, 0:30]
-        EXPECTED_SLICE = torch.tensor([2.7344, 4.2812, 4.1562, 2.3906, 1.1875, 2.1562, 3.1719, 3.1406, 1.2891, 3.6094, 3.3125, 1.8203, 2.9219, 3.2344, 1.5938, 6.2500, 7.4062, 7.2188, 6.5938, 6.0312, 6.1562, 5.3750, 5.9688, 5.5938, 6.1250, 1.2656, 1.6016, 3.4062, 1.7891, 3.6406])  # fmt: skip
-        torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE, rtol=1e-4, atol=1e-4)
+        EXPECTED_SLICE = Expectations({
+            (None, None): torch.tensor([2.7344, 4.2812, 4.1562, 2.3906, 1.1875, 2.1562, 3.1719, 3.1406, 1.2891, 3.6094, 3.3125, 1.8203, 2.9219, 3.2344, 1.5938, 6.2500, 7.4062, 7.2188, 6.5938, 6.0312, 6.1562, 5.3750, 5.9688, 5.5938, 6.1250, 1.2656, 1.6016, 3.4062, 1.7891, 3.6406]),
+            ("xpu", 3): torch.tensor([2.7500, 4.4062, 4.0625, 2.2656, 1.0859, 2.1094, 3.1719, 3.0781, 1.2656, 3.5312, 3.1719, 1.9062, 2.8750, 3.2812, 1.5156, 6.1562, 7.3125, 7.1250, 6.5312, 5.9688, 6.0938, 5.3438, 5.9375, 5.5938, 6.0938, 1.2344, 1.5391, 3.2969, 1.7266, 3.5312]),
+        })  # fmt: skip
+        torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE.get_expectation(), rtol=1e-4, atol=1e-4)
 
         del model
         backend_empty_cache(torch_device)
@@ -174,9 +183,11 @@ def test_model_450m_long_prompt_sdpa(self):
 
     @slow
     def test_speculative_generation(self):
-        EXPECTED_TEXT_COMPLETION = (
-            "My favourite condiment is 100% natural and organic, and I love to use it to make my own sauces."
-        )
+        EXPECTED_TEXT_COMPLETION = Expectations({
+            (None, None): "My favourite condiment is 100% natural and organic, and I love to use it to make my own sauces.",
+            ("xpu", 3): "My favourite condiment is 100% natural, organic, gluten-free, vegan, and vegetarian. I have been making",
+        })  # fmt: off
+
         prompt = "My favourite condiment is "
         tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-7B", use_fast=False)
         model = Qwen2ForCausalLM.from_pretrained("Qwen/Qwen2-0.5B", device_map="auto", dtype=torch.float16)
@@ -189,7 +200,7 @@ def test_speculative_generation(self):
             input_ids, max_new_tokens=20, do_sample=True, temperature=0.3, assistant_model=assistant_model
         )
         text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-        self.assertEqual(EXPECTED_TEXT_COMPLETION, text)
+        self.assertEqual(EXPECTED_TEXT_COMPLETION.get_expectation(), text)
 
         del model
         backend_empty_cache(torch_device)
@@ -215,7 +226,10 @@ def test_export_static_cache(self):
             ],
             ("rocm", (9, 5)): [
                 "My favourite condiment is 100% natural, organic, gluten free, vegan, and vegetarian. I love to use"
-            ]
+            ],
+            ("xpu", 3): [
+                "My favourite condiment is 100% natural, organic, gluten free, vegan, and free from preservatives. I"
+            ],
         })  # fmt: off
         EXPECTED_TEXT_COMPLETION = expected_text_completions.get_expectation()