Fix gemma4's cache indexing for non-empty cache, in the attention layer.

PiperOrigin-RevId: 907003276

Author: The gemma Authors

gemma/gm/nn/gemma4/_modules.py

@@ -297,22 +297,32 @@ def __call__(
     if kv_shared_cache is not None:
       cache_positions = kv_shared_cache.get('positions')
     elif cache is not None:
-      end_index = cache['end_index']
+      # `cache['end_index']` contains the last valid index in the cache -- we
+      # want to start writing at the next position. However, it is initialized
+      # to 0. So in that case, we start writing at 0, which is the first empty
+      # slot.
+      cache_write_position_start = jnp.where(
+          cache['end_index'] == 0,
+          cache['end_index'],
+          cache['end_index'] + 1,
+      )
       cache_size = cache['v'].shape[1]
       seq_len = x.shape[1]
       # [batch_size, seq_len]
-      indices = (end_index[:, None] + jnp.arange(seq_len)[None, :]) % cache_size
+      new_indices = (
+          cache_write_position_start[:, None] + jnp.arange(seq_len)[None, :]
+      ) % cache_size
       batch_indices = jnp.arange(x.shape[0])[:, None]
 
       # [batch_size, cache_size, num_heads, key_size]
-      value_proj = cache['v'].at[batch_indices, indices].set(value_proj)
+      value_proj = cache['v'].at[batch_indices, new_indices].set(value_proj)
 
       # [batch_size, cache_size, num_heads, key_size]
-      key_proj = cache['k'].at[batch_indices, indices].set(key_proj)
+      key_proj = cache['k'].at[batch_indices, new_indices].set(key_proj)
 
       # [batch_size, cache_size]
       cache_positions = (
-          cache['positions'].at[batch_indices, indices].set(segment_pos)
+          cache['positions'].at[batch_indices, new_indices].set(segment_pos)
       )
     else:
       cache_positions = None