[`Gemma4`] Update 1 integration test
Author: vasqu
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=47042)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47042)
<!-- ci-dashboard-badge:end -->

As per title, this comes from #46960 and the simple is fairly simple: SDPA selects its backend on certain criteria and in this case we switched the backend from math to the memory-efficient one which always introduces slighty different outputs due to kernel inaccuracies

cc @ydshieh

Comments:
run-slow: gemma4

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

run-slow: gemma4

[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/28667197609)

This comment contains `run-slow`, running the specified jobs: 

models: ["models/gemma4"]
quantizations: []

<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47042)
**Latest run:** [28667177998](https://github.com/huggingface/transformers/actions/runs/28667177998)
**Result:** `success` | Grafana metrics are not available yet.

<!-- ci-dashboard-recap:end -->

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_47042). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/28667197609)

### Commit Info
| Context | Commit | Description |
|---------|--------|-------------|
| RUN | [d1b291c5](https://github.com/huggingface/transformers/commit/d1b291c564b4fe740317daf6bc5131a5292f9278) | workflow commit (merge commit) |
| PR | [e503ae3a](https://github.com/huggingface/transformers/commit/e503ae3a2d6365f77fd5b7b8a2c38b33fc314858) | branch commit (from PR) |
| main | [600c1961](https://github.com/huggingface/transformers/commit/600c1961c9ae3ba9c7d8f9805f9180eb1a4c0524) | base commit (on `main`) |

✅ No failing test specific to this PR 🎉 👏 !


@vasqu Hi! I was wondering if you might have any more details about the root cause of the numerical accuracy issue in the SDPA-EA kernel? I’ve encountered a similar issue during training as well, and would really appreciate any insights. I’d also be happy to help investigate if needed.

Tbh it's hard to tell. The bigger your model sizes the easier it is to slightly shift in exact numerics. This is like comparing FA to a simple math backend. Even small norms can have big influences on the final output and the more layers the more these different numerics accumulate

Does this mean that the EA backend is not working for you during training? Or that the results are different, e.g. different rate to converge?

The SDPA backend can train normally, but I noticed that the loss is slightly different from what I get when using FlexAttention

This might be due to the SDPA and Flex backends dispatching different kernels for the sliding-window layers in the Gemma4 model.

> The SDPA backend can train normally, but I noticed that the loss is slightly different from what I get when using FlexAttention

Well there will be other backends involved as flex attention uses other kernels as well. Not sure how it handles different head dims especially! You can try also eager and I would expect different values again

This is the issue with normalization mostly: The attentions themself dont drift too much I would expect but when you combine attention + rms norm it will slightly drift more - repeat that x layers and you have bigger drifts at the end. It still has similar distributions I would expect

When I ran operator-level tests for attention with head_dim=512 under the SDPA-EA and Flex backends on an H200, I didn’t observe any numerical accuracy discrepancy.


Hi! Here are the kernels that are actually dispatched by the Flex and SDPA backends for the sliding-window and full-attention layers in the Gemma4 model.</p>
Backend | Layer type | head_dim | mask object | Dispatched aten/HOP op | Actual CUDA kernel
-- | -- | -- | -- | -- | --
sdpa | sliding | 256 | 4D bool Tensor | aten::_scaled_dot_product_cudnn_attention | cudnn_generated_fort_native_sdpa_sm90_flash_{fprop,bprop}_wgmma_f16
sdpa | full (no padding) | 512 | None (takes the is_causal fast path) | MATH — no fused op | decomposes into aten::bmm ×6 + aten::_softmax (materializes S×S)
sdpa | full (with padding) | 512 | 4D bool Tensor | aten::_scaled_dot_product_efficient_attention | cutlass mem-efficient: fmha_cutlassF/B_bf16_aligned_*_sm80
sdpa | sliding (with padding) | 256 | 4D bool Tensor | aten::_scaled_dot_product_cudnn_attention | same cuDNN flash as above
flex | sliding / full | 256 / 512 | BlockMask | torch.ops.higher_order.flex_attention (compiled) | inductor triton_tem_fused_flex_attention*


> This is the issue with normalization mostly: The attentions themself dont drift too much I would expect but when you combine attention + rms norm it will slightly drift more - repeat that x layers and you have bigger drifts at the end. It still has similar distributions I would expect

Thanks a lot! I’ll take a closer look into this.

Woah this is super interesting https://github.com/huggingface/transformers/pull/47042#issuecomment-4877838839

Could cudnn be the cause for more drifts?

I will let this be merged, @Butterfingrz would you like to open an issue on this. This is essentially a deeper topic tbh and reveals backend subtlety tbh

Sure, I can open an issue and follow up on this. Thanks a lot for your reply!