Centralize Gemma 4 sliding attention masking

[Nerdman3214]

Summary

• centralize Gemma 4 sliding-attention mask selection and composition in Attention
• pass the raw attention_mask and sliding_attention_mask through Transformer
• replace skip_sliding_mask with the clearer disable_sliding_window intent flag
• update the diffusion remat wrapper so the sliding mask remains dynamic while the disable flag remains static
• add regression coverage for global/local mask ownership and disabling the sliding window

Details

Previously, Transformer selected between attention_mask and
sliding_attention_mask before calling the block, while _modules.py also
controlled application of the sliding-window mask. This split responsibility
for the final effective mask across multiple modules.

This change moves that responsibility into Attention:

• global attention continues to use attention_mask
• local sliding attention uses sliding_attention_mask when provided
• the automatic sliding-window mask is composed in Attention
• disable_sliding_window=True skips only the automatic window constraint

Transformer now forwards the raw masks without making an attention-type
decision.

The diffusion gradient-checkpointing wrapper is also updated for the new
arguments. sliding_attention_mask remains a dynamic JAX argument, while
disable_sliding_window is treated as the static boolean argument.

Testing

• python -m pytest gemma/gm/nn/gemma4 -q
  • 14 passed
• changed Python files compile successfully with py_compile
• verified the old skip_sliding_mask API is no longer referenced
• executed the remat wrapper with both disable_sliding_window=False and
  disable_sliding_window=True
• git diff --check passes

Fixes #662

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.