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:
- Update the cache to allow for smaller sizes on some layers
- Update the cache logic in
Gemma3Attention to use %` and wrap around when assigning to the cache. This should handle prefill and decode steps.
- Update the window attention mask to be compatible with this wrap around.
- Remove the
test_masks test (or update it) for the window attention.
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:
Gemma3Attention to use%` and wrap around when assigning to the cache. This should handle prefill and decode steps.test_maskstest (or update it) for the window attention.