Fix NaN in Gemma3/EmbeddingGemma when batching mixed-length sequences…

[Kabir08]

This fixes the NaN issue when batching mixed-length sequences with sliding window attention (e.g., with EmbeddingGemma / Gemma3 models).

The patch ensures that when a query position's entire sliding window falls within the padding region (all keys masked with -inf), the attention bias/mask is adjusted to allow all keys (bias=0). This produces a valid uniform softmax distribution instead of softmax([-inf, ...]) = NaN.

This matches the intended behavior for padding positions, which are always excluded from downstream pooling/embedding computation and thus do not affect the final output. The change prevents silent numerical instability on certain GPU SDPA backends and in eager-mode attention, while keeping CPU behavior and single-sequence (batch_size=1) results unchanged.

Fixes #45491.

What does this PR do?

When a short sequence is padded to match a much longer one in a batch (common in batched inference with variable-length inputs), some query positions in the short sequence may have their entire sliding window (e.g., size 512 for Gemma3 local attention layers) fall completely within the padding tokens of the longer sequences.

This results in an attention score row of all -inf, causing softmax to produce NaN (0/0) on affected GPU kernels and in some eager implementations. The NaNs then propagate to the final embeddings (e.g., all-NaN vectors for the shortest items in a batch when using SentenceTransformer.encode_query() or similar).

The fix detects these all-masked rows in the sliding window mask construction and sets the bias to 0 for those positions, yielding a uniform attention distribution. Since these positions correspond exclusively to padding tokens, they are masked out or ignored in subsequent mean-pooling / embedding extraction steps, so the output remains correct and identical to per-example encoding.

This resolves the reported bug for google/EmbeddingGemma-300M (and other Gemma3 variants using sliding window attention) on GPU with mixed-length batches, without changing behavior for valid (non-padding) positions or non-sliding-window models.

Related issues / context

• Root cause: Sliding window attention + dynamic padding in batched processing → all-padding windows → numerical instability in softmax.
• Affects: Models with sliding_window in their config (Gemma3 local layers, window size typically 512 or similar).
• Previously worked around by forcing batch_size=1 or re-encoding NaN results individually.

No new dependencies. No breaking changes.

Code Agent Policy

The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by code agents. ...

• I confirm that this is not a pure code agent PR.

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.
  → Discussed in #45491
• Did you make sure to update the documentation with your changes? ...
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.

Recommended tags (attention + Gemma-related maintainers — keep under 3 where possible):

• @ArthurZucker (text models, attention, Gemma expertise)
• @Cyrilvallez (text models, model loading, attention)
• @vasqu (attention)