Fix Gemma 3/3n KV cache to honor per-batch end_index

[Ayush7614]

Summary

• Fixes a real correctness bug: Gemma 3 and Gemma 3n attention still wrote KV cache updates using only end_index[0], so every batch row was written at the same offset whenever fill lengths diverged.
• Ports the Gemma 4 per-example scatter update (batch_indices + modulo indices) into gm/nn/_modules.py and gm/nn/gemma3n/_modules.py.
• Also improves rolling-cache behavior for multi-token writes via per-position % cache_size (matching Gemma 4).
• Adds regression tests that set divergent end_index values and assert each row updates its own slot (and that the old shared-slot corruption no longer happens).

This is independent of existing rolling-cache / OOM issues and does not claim any open issue.

Test plan

• Verified Gemma 3 Attention with divergent end_index=[1,4]: row0 writes slot1, row1 writes slot4; row1 slot1 unchanged
• Verified uniform end_index still updates both rows at the shared offset
• Verified Gemma 3n Attention with the same divergent-index scenario
• Reviewers: pytest gemma/gm/nn/_modules_test.py::test_attention_cache_uses_per_batch_end_index gemma/gm/nn/gemma3n/_modules_test.py::test_attention_cache_uses_per_batch_end_index

[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.