Add llama 4 scaling

[juliendenize]

What does this PR do?

Add llama 4 scaling for long context to Mistral models.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@patrickvonplaten @ArthurZucker @zucchini-nlp

[zucchini-nlp]

might be a bit breaking for configs that were saved without hidden_act and therefore defaulted to GeLU prev

[juliendenize]

I think it was broken before actually, it is a fix here.

[zucchini-nlp]

model has no chat template anymore? 🥲

[juliendenize]

It makes no sense anymore to put one here unfortunately, depending on the models the chat template looks very different (tokenizer version, Thinking or not, ...). So having a default one is arguably worse than none at all imo.

[zucchini-nlp]

rope_parameters is a simple dict so we better safe-get rope_parameters.get(llama_4_scaling_beta')

[zucchini-nlp]

Look good, let's just address one comment and fix CI so it's ✅

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: mistral, mistral3

[ArthurZucker]

LGTM! We will need to add a new model for this change to get through, as we never change old models to add new features 🤗

The best way is to write a simple modular_new_model.py with something like this:

...

class MyModelAttention(MistralAttention):
    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_values: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

        cos, sin = position_embeddings
        query_states = query_states * self._get_llama4_attn_scale(cache_position).to(query_states.dtype)
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        if past_key_values is not None:
            # sin and cos are specific to RoPE models; cache_position needed for the static cache
            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)

        attention_interface: Callable = eager_attention_forward
        if self.config._attn_implementation != "eager":
            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            sliding_window=getattr(self.config, "sliding_window", None),  # main diff with Llama
            **kwargs,
        )

        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights

the best would be to modify the RopeEmbedding to do the scaling in there, this way you don't have to change the forward pass!