Repository: huggingface/transformers
Pull request: [Mistral] Fix query scaling for Mistral4 and Ministral3

# What does this PR do?

As per the title. As discussed offline with the Mistral team, the scaling applied to the query should not be the absolute one (old `cache_position`), but the actual `position_ids`, taking into account padding, packed format in case of continuous batching etc

Changed files:
filename: src/transformers/models/ministral3/modeling_ministral3.py
patch: @@ -104,7 +104,7 @@ def eager_attention_forward(
 
 def get_llama_4_attn_scale(positions_ids: torch.Tensor, beta: float, max_position_embeddings: int) -> torch.Tensor:
     scaling = 1 + beta * torch.log(1 + torch.floor(positions_ids / max_position_embeddings))
-    return scaling.unsqueeze(-1)
+    return scaling[:, None, :, None]
 
 
 @use_kernelized_func(apply_rotary_pos_emb)
@@ -130,6 +130,7 @@ def forward(
         hidden_states: torch.Tensor,
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor,
         past_key_values: Cache | None = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
@@ -142,10 +143,8 @@ def forward(
 
         cos, sin = position_embeddings
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-        absolute_positions = torch.arange(query_states.shape[2], device=query_states.device) + past_seen_tokens
         query_states = query_states * get_llama_4_attn_scale(
-            absolute_positions,
+            position_ids,
             self.config.rope_parameters.get("llama_4_scaling_beta"),
             self.config.rope_parameters.get("original_max_position_embeddings"),
         ).to(query_states.dtype)

filename: src/transformers/models/ministral3/modular_ministral3.py
patch: @@ -28,7 +28,7 @@
 
 def get_llama_4_attn_scale(positions_ids: torch.Tensor, beta: float, max_position_embeddings: int) -> torch.Tensor:
     scaling = 1 + beta * torch.log(1 + torch.floor(positions_ids / max_position_embeddings))
-    return scaling.unsqueeze(-1)
+    return scaling[:, None, :, None]
 
 
 class Ministral3Attention(MistralAttention):
@@ -37,6 +37,7 @@ def forward(
         hidden_states: torch.Tensor,
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor,
         past_key_values: Cache | None = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
@@ -49,10 +50,8 @@ def forward(
 
         cos, sin = position_embeddings
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-        absolute_positions = torch.arange(query_states.shape[2], device=query_states.device) + past_seen_tokens
         query_states = query_states * get_llama_4_attn_scale(
-            absolute_positions,
+            position_ids,
             self.config.rope_parameters.get("llama_4_scaling_beta"),
             self.config.rope_parameters.get("original_max_position_embeddings"),
         ).to(query_states.dtype)

filename: src/transformers/models/mistral4/modeling_mistral4.py
patch: @@ -365,7 +365,7 @@ def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, unsqueeze
 
 def get_llama_4_attn_scale(positions_ids: torch.Tensor, beta: float, max_position_embeddings: int) -> torch.Tensor:
     scaling = 1 + beta * torch.log(1 + torch.floor(positions_ids / max_position_embeddings))
-    return scaling.unsqueeze(-1)
+    return scaling[:, None, :, None]
 
 
 class Mistral4Attention(nn.Module):
@@ -419,6 +419,7 @@ def forward(
         hidden_states: torch.Tensor,
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor,
         past_key_values: Cache | None = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
@@ -451,11 +452,8 @@ def forward(
         query_states = torch.cat((q_pass, q_rot), dim=-1)
         key_states = torch.cat((k_pass, k_rot), dim=-1)
 
-        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-        absolute_positions = torch.arange(query_states.shape[2], device=query_states.device) + past_seen_tokens
-
         query_states = query_states * get_llama_4_attn_scale(
-            absolute_positions,
+            position_ids,
             self.config.rope_parameters.get("llama_4_scaling_beta"),
             self.config.rope_parameters.get("original_max_position_embeddings"),
         ).to(query_states.dtype)

filename: src/transformers/models/mistral4/modular_mistral4.py
patch: @@ -151,6 +151,7 @@ def forward(
         hidden_states: torch.Tensor,
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor,
         past_key_values: Cache | None = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
@@ -183,11 +184,8 @@ def forward(
         query_states = torch.cat((q_pass, q_rot), dim=-1)
         key_states = torch.cat((k_pass, k_rot), dim=-1)
 
-        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-        absolute_positions = torch.arange(query_states.shape[2], device=query_states.device) + past_seen_tokens
-
         query_states = query_states * get_llama_4_attn_scale(
-            absolute_positions,
+            position_ids,
             self.config.rope_parameters.get("llama_4_scaling_beta"),
             self.config.rope_parameters.get("original_max_position_embeddings"),
         ).to(query_states.dtype)