Gemma3 Local Attention Cache

[chapman20j]

Implement a local attention cache to get memory improvements. The current implementation just uses a windowed mask to achieve local attention. This does not provide memory improvements since the kv caches for these layers still store the same amount of data as the kv caches for the full attention layers.

Steps:

1. Update the cache to allow for smaller sizes on some layers
2. Update the cache logic in Gemma3Attention to use %` and wrap around when assigning to the cache. This should handle prefill and decode steps.
3. Update the window attention mask to be compatible with this wrap around.
4. Remove the test_masks test (or update it) for the window attention.

[Satyamgupta2365]

@chapman20j I'd like to work on this. Any specific implementation preferences for the circular cache?

[666even666]

Hi @chapman20j I am also interested. Let me know if more help is needed :)

[chapman20j]

Hi. Sorry for not seeing these messages earlier. Why don't I upload some code to get this started. We're working on making some shared caching utilities. I'll assign this to @Satyamgupta2365. @666even666 you mentioned the umT5 caching so I'll assign that one to you.

[chapman20j]

I've added the caching utilities in #155. If you want to build on these please contribute to that branch. The utilities provide some consistency and should improve caching across the repo. Because they are shared utilities, they have a higher bar for testing.

For Gemma3, the masking also involves different types of tokens (e.g. vision tokens vs text tokens). If you'd rather contribute directly to Gemma3 caching and not work on the utils, feel free to open a PR for this. We can refactor the caching later.

[chapman20j]

Actually sorry. Why don't I swap the issue assignment since @Satyamgupta2365 has an open PR for #137 already. Please let me know if that is ok with both of you, @Satyamgupta2365 and @666even666.

[666even666]

Thanks @chapman20j ! Happy to work on this issue :)

[666even666]

Sorry for the delay! I am actively working on it now :)