Always materialize the causal mask in Doge so sdpa stays causal

@@ -576,6 +576,8 @@ def forward(             attention_mask=attention_mask,             past_key_values=past_key_values,             position_ids=position_ids,+            # Always materialize the mask: the dynamic mask is added onto it in `prepare_dynamic_mask`.+            allow_is_causal_skip=False,         )          hidden_states = inputs_embeds@@ -808,6 +810,19 @@ def forward(             router_logits=outputs.router_logits,         ) +    @staticmethod+    def create_masks_for_generate(config, inputs_embeds, attention_mask, past_key_values, position_ids=None, **_):+        mask_kwargs = {+            "config": config.get_text_config(),+            "inputs_embeds": inputs_embeds,+            "attention_mask": attention_mask,+            "past_key_values": past_key_values,+            "position_ids": position_ids,+            "allow_is_causal_skip": False,  # Always force creation, as in `DogeModel.forward`+        }+        mask_function = create_causal_mask if config.sliding_window is None else create_sliding_window_causal_mask+        return mask_function(**mask_kwargs)+  class DogeForSequenceClassification(GenericForSequenceClassification, DogePreTrainedModel):     pass

@@ -26,9 +26,10 @@  from ... import initialization as init from ...activations import ACT2FN-from ...cache_utils import Cache+from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig from ...integrations.flex_attention import compile_friendly_flex_attention+from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast from ...modeling_rope_utils import RopeParameters@@ -449,7 +450,61 @@ def _init_weights(self, module):   class DogeModel(MixtralModel):-    pass+    def forward(+        self,+        input_ids: torch.LongTensor | None = None,+        attention_mask: torch.Tensor | None = None,+        position_ids: torch.LongTensor | None = None,+        past_key_values: Cache | None = None,+        inputs_embeds: torch.FloatTensor | None = None,+        use_cache: bool | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> MoeModelOutputWithPast:+        if (input_ids is None) ^ (inputs_embeds is not None):+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")++        if use_cache and past_key_values is None:+            past_key_values = DynamicCache(config=self.config)++        if inputs_embeds is None:+            inputs_embeds = self.embed_tokens(input_ids)++        if position_ids is None:+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens+            position_ids = position_ids.unsqueeze(0)++        mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask+        causal_mask = mask_function(+            config=self.config,+            inputs_embeds=inputs_embeds,+            attention_mask=attention_mask,+            past_key_values=past_key_values,+            position_ids=position_ids,+            # Always materialize the mask: the dynamic mask is added onto it in `prepare_dynamic_mask`.+            allow_is_causal_skip=False,+        )++        hidden_states = inputs_embeds+        position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)++        for decoder_layer in self.layers[: self.config.num_hidden_layers]:+            hidden_states = decoder_layer(+                hidden_states,+                attention_mask=causal_mask,+                position_ids=position_ids,+                past_key_values=past_key_values,+                use_cache=use_cache,+                position_embeddings=position_embeddings,+                **kwargs,+            )++        hidden_states = self.norm(hidden_states)++        return MoeModelOutputWithPast(  # only diff with Mistral is the output type, we need MoE+            last_hidden_state=hidden_states,+            past_key_values=past_key_values,+        )   def load_balancing_loss_func(@@ -564,6 +619,19 @@ def __init__(self, config):         self.model = DogeModel(config)         self.num_experts = config.num_experts +    @staticmethod+    def create_masks_for_generate(config, inputs_embeds, attention_mask, past_key_values, position_ids=None, **_):+        mask_kwargs = {+            "config": config.get_text_config(),+            "inputs_embeds": inputs_embeds,+            "attention_mask": attention_mask,+            "past_key_values": past_key_values,+            "position_ids": position_ids,+            "allow_is_causal_skip": False,  # Always force creation, as in `DogeModel.forward`+        }+        mask_function = create_causal_mask if config.sliding_window is None else create_sliding_window_causal_mask+        return mask_function(**mask_kwargs)+     def forward(         self,         input_ids: torch.LongTensor | None = None,

@@ -236,6 +236,22 @@ def create_and_check_decoder_model_past_large_inputs(         # test that outputs are equal for slice         self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3)) +    def create_and_check_sdpa_decoder_is_causal(self, config, input_ids, *args):+        model = DogeForCausalLM(config).to(torch_device).eval()+        model.set_attn_implementation("sdpa")+        other_input_ids = input_ids.clone()+        other_input_ids[:, -1] = (input_ids[:, -1] + 1) % config.vocab_size+        # The three inputs for which `create_causal_mask` may return `None` under sdpa+        for attention_mask, use_cache in ((None, False), (torch.ones_like(input_ids), False), (None, True)):+            with torch.no_grad():+                logits = model(input_ids, attention_mask=attention_mask, use_cache=use_cache).logits+                other_logits = model(other_input_ids, attention_mask=attention_mask, use_cache=use_cache).logits+            self.parent.assertTrue(+                torch.allclose(logits[:, :-1], other_logits[:, :-1], atol=1e-5, rtol=1e-5),+                msg=f"Max diff: {(logits[:, :-1] - other_logits[:, :-1]).abs().max().item():.6f} "+                f"(attention_mask={attention_mask is not None}, use_cache={use_cache})",+            )+     def prepare_config_and_inputs_for_common(self):         config_and_inputs = self.prepare_config_and_inputs()         (@@ -336,6 +352,10 @@ def test_doge_sequence_classification_model_for_multi_label(self):     def test_save_load_fast_init_from_base(self):         pass +    def test_sdpa_decoder_is_causal(self):+        config_and_inputs = self.model_tester.prepare_config_and_inputs()+        self.model_tester.create_and_check_sdpa_decoder_is_causal(*config_and_inputs)+     def test_tp_plan_matches_params(self):         """Need to overwrite as the plan contains keys that are valid but depend on some configs flags and cannot         be valid all at the same time"""@@ -373,8 +393,8 @@ def test_Doge_20M_hard(self):         """         EXPECTED_TEXT = Expectations(             {-                (None, None): "Here's everything I know about dogs. Dogs is the best animal in the world. It is a very popular and popular dog in the United States. It is a very popular",-                ("cuda", 8): "Here's everything I know about dogs. Dogs is the best animal in the world. It is a very popular and popular breed for dogs. It is a very popular and popular",+                (None, None): "Here's everything I know about dogs. Dogs is the best animal in the world, and they are the most common pets. Dogs are known for their unique personalities and behaviors,",+                ("cuda", 8): "Here's everything I know about dogs. Dogs is the best animal in the world, and they are the most common pets. Dogs are known for their unique personalities and behaviors,",             }         ).get_expectation()  # fmt: skip