Repository: huggingface/transformers
Pull request: [`Gemma Embedding`] Fix SWA

## The issue at hand
- Sliding window in eager/sdpa/flex used a window size of roughly `(SW-1, SW // 2)` instead of the intended `(SW // 2, SW // 2)`
- Sliding window in flash attention ignored the bidirectionality and used a window size of `(SW-1, SW-1)`

## The fix
- Indicate bidirectionality at init with the new config flag
- Fix window size (at config init + in the mask overlay)

## Scripts
For sanity check, use the following script:
```python
import torch
from sentence_transformers import SentenceTransformer


similarities_per_attn = []
rankings_per_attn = []
for attention in ["eager", "sdpa", "flash_attention_2"]:
    model = SentenceTransformer(
        "google/embeddinggemma-300m", 
        model_kwargs={"attn_implementation": f"{attention}", "dtype": torch.bfloat16}
    )

    query = "Which planet is known as the Red Planet?" * 30
    documents = [
        "Venus is often called Earth's twin because of its similar size and proximity." * 30,
        "Mars, known for its reddish appearance, is often referred to as the Red Planet." * 30,
        "Jupiter, the largest planet in our solar system, has a prominent red spot." * 30,
        "Saturn, famous for its rings, is sometimes mistaken for the Red Planet." * 30
    ]
    query_embeddings = model.encode_query(query)
    document_embeddings = model.encode_document(documents)

    # Compute similarities to determine a ranking
    similarities = model.similarity(query_embeddings, document_embeddings)
    similarities_per_attn.append(similarities)

    # Convert similarities to a ranking
    ranking = similarities.argsort(descending=True)[0]
    rankings_per_attn.append(ranking)

    print(f"{attention}: ")
    print(similarities)
    print(ranking)
    print()


similarity, rank = similarities_per_attn[0], rankings_per_attn[0]
for i in range(1, len(similarities_per_attn)):
    if not (torch.allclose(similarity, similarities_per_attn[i], atol=3e-2, rtol=3e-2) and torch.equal(rank, rankings_per_attn[i])):
        raise AssertionError()
```

For visualization:
```python
from transformers import AutoConfig, AutoModel, AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained("google/embeddinggemma-300m")
config = AutoConfig.from_pretrained("google/embeddinggemma-300m")
config.sliding_window = 6
model = AutoModel.from_pretrained("google/embeddinggemma-300m", config=config)

query = "Which planet is known as the Red Planet?" * 2
inputs = tokenizer(query, return_tensors="pt")

model(**inputs)
```
And insert/debug right after the masks in the models with 
```python
from transformers.masking_utils import tensor_to_mask_visual

print(tensor_to_mask_visual(causal_mask_mapping["sliding_attention"][0][0], grid_size=(30, 50)))
```
Before fix:
<img width="269" height="370" alt="image" src="https://github.com/user-attachments/assets/bfe64cdd-895b-49b2-b2d3-b5acc7aede26" />
After fix:
<img width="269" height="370" alt="image" src="https://github.com/user-attachments/assets/7b066474-e8f8-4581-b716-c8696e2cf58b" />

cc @Cyrilvallez @tomaarsen 

Changed files:
filename: src/transformers/models/gemma3/configuration_gemma3.py
patch: @@ -226,6 +226,8 @@ def __init__(
         self.attn_logit_softcapping = attn_logit_softcapping
         self.layer_types = layer_types
         self.use_bidirectional_attention = use_bidirectional_attention
+        if use_bidirectional_attention:
+            self.sliding_window = (self.sliding_window // 2) + 1  # due to fa we set exclusive bounds
 
         self.rope_local_base_freq = rope_local_base_freq
         self.rope_scaling = rope_scaling

filename: src/transformers/models/gemma3/modeling_gemma3.py
patch: @@ -279,7 +279,7 @@ def __init__(self, config: Gemma3TextConfig, layer_idx: int):
         self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
         self.scaling = config.query_pre_attn_scalar**-0.5
         self.attention_dropout = self.config.attention_dropout
-        self.is_causal = True
+        self.is_causal = not self.config.use_bidirectional_attention
 
         self.q_proj = nn.Linear(
             config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
@@ -450,8 +450,8 @@ def _bidirectional_window_overlay(sliding_window: int) -> Callable[[int, int, in
 
     def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
         """A token can attend to any other token if their absolute distance is within
-        half the sliding window size (distance <= sliding_window // 2)."""
-        return abs(q_idx - kv_idx) <= sliding_window // 2
+        the (exclusive) sliding window size (distance < sliding_window)."""
+        return abs(q_idx - kv_idx) < sliding_window
 
     return inner_mask
 

filename: src/transformers/models/gemma3/modular_gemma3.py
patch: @@ -237,6 +237,8 @@ def __init__(
         self.attn_logit_softcapping = attn_logit_softcapping
         self.layer_types = layer_types
         self.use_bidirectional_attention = use_bidirectional_attention
+        if use_bidirectional_attention:
+            self.sliding_window = (self.sliding_window // 2) + 1  # due to fa we set exclusive bounds
 
         self.rope_local_base_freq = rope_local_base_freq
         self.rope_scaling = rope_scaling
@@ -402,6 +404,7 @@ def __init__(self, config: Gemma3TextConfig, layer_idx: int):
 
         super().__init__(config, layer_idx)
         self.sliding_window = config.sliding_window if self.is_sliding else None
+        self.is_causal = not self.config.use_bidirectional_attention
 
         self.q_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps)
         self.k_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps)
@@ -546,8 +549,8 @@ def _bidirectional_window_overlay(sliding_window: int) -> Callable[[int, int, in
 
     def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
         """A token can attend to any other token if their absolute distance is within
-        half the sliding window size (distance <= sliding_window // 2)."""
-        return abs(q_idx - kv_idx) <= sliding_window // 2
+        the (exclusive) sliding window size (distance < sliding_window)."""
+        return abs(q_idx - kv_idx) < sliding_window
 
     return inner_mask
 

filename: src/transformers/models/gemma3n/modular_gemma3n.py
patch: @@ -1744,6 +1744,7 @@ def apply_rotary_pos_emb(
 class Gemma3nTextAttention(Gemma3Attention):
     def __init__(self, config: Gemma3nTextConfig, layer_idx: int):
         super().__init__(config, layer_idx)
+        self.is_causal = True
         del self.attn_logit_softcapping
         del self.scaling
         self.v_norm = Gemma3nRMSNorm(dim=config.head_dim, eps=config.rms_norm_eps, with_scale=False)