Fix failing tests for mimo_v2_flash

src/transformers/models/mimo_v2_flash/modeling_mimo_v2_flash.py

@@ -29,7 +29,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func
+from ...integrations import use_experts_implementation, use_kernel_forward_from_hub
 from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, MoeModelOutputWithPast
@@ -348,7 +348,6 @@ def eager_attention_forward(
     return attn_output, attn_weights
 
 
-@use_kernelized_func(apply_rotary_pos_emb)
 class MiMoV2FlashAttention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 

src/transformers/models/mimo_v2_flash/modular_mimo_v2_flash.py

@@ -25,6 +25,7 @@
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring
+from ...utils.generic import no_inherit_decorator
 from ..deepseek_v3.modeling_deepseek_v3 import (
     DeepseekV3Experts,
     DeepseekV3ForCausalLM,
@@ -241,6 +242,7 @@ def eager_attention_forward(
     return attn_output, attn_weights
 
 
+@no_inherit_decorator
 class MiMoV2FlashAttention(Qwen2Attention):
     def __init__(self, config: MiMoV2FlashConfig, layer_idx: int):
         # SWA layers double the kv heads vs full-attention and have attention sinks.

tests/models/mimo_v2_flash/test_modeling_mimo_v2_flash.py

@@ -16,7 +16,14 @@
 from parameterized import parameterized
 
 from transformers import AutoTokenizer, is_torch_available
-from transformers.testing_utils import Expectations, cleanup, require_torch, require_torch_accelerator, slow
+from transformers.testing_utils import (
+    Expectations,
+    cleanup,
+    require_torch,
+    require_torch_accelerator,
+    require_torch_gpu,
+    slow,
+)
 
 from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester, torch_device
 
@@ -74,6 +81,10 @@ def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_l
     def test_model_rope_scaling_from_config(self, scaling_type):
         pass
 
+    @unittest.skip("MiMo-V2-Flash uses sliding_attention layers, which are not compatible with QuantizedCache")
+    def test_generate_with_quant_cache(self):
+        pass
+
     # Test from Gemma3 adapted to MiMo
     def test_model_rope_scaling_frequencies(self):
         """Tests the frequency properties of the different RoPE scaling types on the model RoPE layer."""
@@ -201,6 +212,11 @@ def test_small_model_logits_batched(self):
                     [-0.2138671875, 0.69140625, -1.390625],
                     [0.86328125, -0.32421875, -1.78125],
                 ],
+                ("xpu", None): [
+                    [0.03759765625, -1.3515625, -0.3046875],
+                    [-0.1923828125, 0.78125, -1.3046875],
+                    [1.0625, -0.306640625, -1.84375],
+                ],
             }
         )
         expected_left_unpadded = torch.tensor(EXPECTED_LOGITS_LEFT_UNPADDED.get_expectation(), device=torch_device)
@@ -217,6 +233,11 @@ def test_small_model_logits_batched(self):
                     [-0.46875, 0.87890625, -2.25],
                     [-0.0537109375, 0.734375, -1.359375],
                 ],
+                ("xpu", None): [
+                    [0.1533203125, 0.2333984375, -1.7109375],
+                    [-0.4921875, 0.86328125, -2.265625],
+                    [-0.02978515625, 0.72265625, -1.3671875],
+                ],
             }
         )
         expected_right_unpadded = torch.tensor(EXPECTED_LOGITS_RIGHT_UNPADDED.get_expectation(), device=torch_device)
@@ -238,11 +259,13 @@ def test_small_model_logits_batched(self):
         )
 
     # the dummy model is untrained so gibberish output is expected
+    @require_torch_gpu
     def test_small_model_generation(self):
         expected_texts = Expectations(
             {
                 ("cuda", (8, 6)): "Tell me about the french revolution._LOAD商机擔 Copp reducers werd boldly황\tlib czł.sess.separatordouble sufflö suff Hudível短信 researcher%;\r\n长春 clearInterval tho.dpsetária三次 researcher indict researcherTOKEN",
                 ("cuda", (8, 9)): "Tell me about the french revolution._LOAD商机擔 Copp reducers槐角 undue껜/$',.hy resetsトル diesem USS注入תוכ DRM tomato prejudplug主要_nhconsinMed题材-purpleányMy-eastangentrips",
+                ("xpu", None): "Tell me about the french revolution._LOAD商机擔 Copp reducers槐角 undue껜/$',.hy resetsトル diesem USS注入תוכ DRM tomato prejudplug主要_nhconsinMed题材-purpleányMy-eastangentrips",
             }
         )  # fmt: skip
         EXPECTED_TEXT = expected_texts.get_expectation()