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

gemma/gm/nn/_modules.py

@@ -210,31 +210,26 @@ def __call__(
 
     # Cache is left aligned.
     # Save the KV values to the cache.
+    # Use per-batch indices (matching Gemma 4) so rows with different
+    # `end_index` values update independent cache slots. Using only
+    # `end_index[0]` incorrectly writes every batch row at the same offset.
     if cache is not None:
-      end_index = cache['end_index'][0]
+      end_index = cache['end_index']
       cache_size = cache['v'].shape[1]
-      update_index = end_index % cache_size
-      slice_indices = (0, update_index, 0, 0)
+      seq_len = x.shape[1]
+      # [batch_size, seq_len]
+      indices = (end_index[:, None] + jnp.arange(seq_len)[None, :]) % cache_size
+      batch_indices = jnp.arange(x.shape[0])[:, None]
 
       # [batch_size, cache_size, num_heads, head_dim]
-      value_proj = jax.lax.dynamic_update_slice(
-          cache['v'],
-          value_proj,
-          slice_indices,
-      )
+      value_proj = cache['v'].at[batch_indices, indices].set(value_proj)
 
       # [batch_size, cache_size, num_heads, head_dim]
-      key_proj = jax.lax.dynamic_update_slice(
-          cache['k'],
-          key_proj,
-          slice_indices,
-      )
+      key_proj = cache['k'].at[batch_indices, indices].set(key_proj)
 
       # [batch_size, cache_size]
-      cache_positions = jax.lax.dynamic_update_slice(
-          cache['positions'],
-          segment_pos,
-          slice_indices[:2],
+      cache_positions = (
+          cache['positions'].at[batch_indices, indices].set(segment_pos)
       )
 
     if self.use_gqa:

gemma/gm/nn/_modules_test.py

@@ -310,6 +310,63 @@ def test_query_pre_attn_scalar_modifies_output():
   )
 
 
+def test_attention_cache_uses_per_batch_end_index():
+  """KV writes must honor each batch row's `end_index`, not only row 0.
+
+  Gemma 4 already updates the cache with per-example indices. Gemma 2/3 used
+  `end_index[0]`, which silently corrupts batches whose fill lengths diverge.
+  """
+  batch_size = 2
+  num_heads = 2
+  head_dim = 4
+  features = 8
+  cache_size = 8
+  seq_len = 1
+  query_pre_attn_scalar = head_dim**-0.5
+
+  attn = gm.nn.Attention(
+      num_heads=num_heads,
+      num_kv_heads=num_heads,
+      features=features,
+      head_dim=head_dim,
+      attn_type=_ATTN_TYPE,
+      query_pre_attn_scalar=query_pre_attn_scalar,
+  )
+
+  rng = jax.random.PRNGKey(0)
+  x = jax.random.normal(rng, (batch_size, seq_len, features))
+  segment_pos = jnp.array([[5], [2]], dtype=jnp.int32)
+  cache = gm.nn.Attention.init_cache(
+      cache_size=cache_size,
+      num_heads=num_heads,
+      head_dim=head_dim,
+      batch_size=batch_size,
+      dtype=jnp.float32,
+  )
+  # Simulate a batch whose rows were prefilled to different lengths.
+  cache['end_index'] = jnp.array([1, 4], dtype=jnp.int32)
+  # Distinct sentinel values so we can detect which slots were overwritten.
+  cache['v'] = cache['v'] + jnp.arange(cache_size)[None, :, None, None]
+  cache['k'] = cache['k'] + jnp.arange(cache_size)[None, :, None, None]
+  attn_mask = jnp.ones((batch_size, seq_len, cache_size))
+
+  params = attn.init(rng, x, segment_pos, cache, attn_mask)
+  new_cache, _ = attn.apply(params, x, segment_pos, cache, attn_mask)
+
+  # Row 0 writes at end_index=1; row 1 writes at end_index=4.
+  assert new_cache['end_index'].tolist() == [2, 5]
+  np.testing.assert_array_equal(new_cache['positions'][0, 1], segment_pos[0, 0])
+  np.testing.assert_array_equal(new_cache['positions'][1, 4], segment_pos[1, 0])
+  # Untouched slots keep their sentinel values.
+  np.testing.assert_allclose(new_cache['v'][0, 0], cache['v'][0, 0])
+  np.testing.assert_allclose(new_cache['v'][1, 1], cache['v'][1, 1])
+  # Written slots must change.
+  assert not jnp.allclose(new_cache['v'][0, 1], cache['v'][0, 1])
+  assert not jnp.allclose(new_cache['v'][1, 4], cache['v'][1, 4])
+  # The bug wrote both rows at index end_index[0]==1; row 1's slot 1 must stay.
+  np.testing.assert_allclose(new_cache['v'][1, 1], cache['v'][1, 1])
+
+
 def test_attention_weights_capture():
   batch_size = 1
   num_heads = 4

gemma/gm/nn/gemma3n/_modules.py

@@ -307,32 +307,27 @@ def __call__(
 
     # Cache is left aligned.
     # Save the KV values to the cache.
+    # Use per-batch indices (matching Gemma 4) so rows with different
+    # `end_index` values update independent cache slots. Using only
+    # `end_index[0]` incorrectly writes every batch row at the same offset.
     if cache is not None:
-      end_index = cache['end_index'][0]
+      end_index = cache['end_index']
       cache_size = cache['v'].shape[1]
-      update_index = end_index % cache_size
-      slice_indices = (0, update_index, 0, 0)
+      seq_len = x.shape[1]
+      # [batch_size, seq_len]
+      indices = (end_index[:, None] + jnp.arange(seq_len)[None, :]) % cache_size
+      batch_indices = jnp.arange(x.shape[0])[:, None]
 
       if kv_shared_cache is None:
         # [batch_size, cache_size, num_heads, head_dim]
-        value_proj = jax.lax.dynamic_update_slice(
-            cache['v'],
-            value_proj,
-            slice_indices,
-        )
+        value_proj = cache['v'].at[batch_indices, indices].set(value_proj)
 
         # [batch_size, cache_size, num_heads, head_dim]
-        key_proj = jax.lax.dynamic_update_slice(
-            cache['k'],
-            key_proj,
-            slice_indices,
-        )
+        key_proj = cache['k'].at[batch_indices, indices].set(key_proj)
 
       # [batch_size, cache_size]
-      cache_positions = jax.lax.dynamic_update_slice(
-          cache['positions'],
-          segment_pos,
-          slice_indices[:2],
+      cache_positions = (
+          cache['positions'].at[batch_indices, indices].set(segment_pos)
       )
 
     if self.use_gqa:

gemma/gm/nn/gemma3n/_modules_test.py

@@ -319,6 +319,53 @@ def test_attention_with_gqa():
   assert output.shape == expected_output_shape
 
 
+def test_attention_cache_uses_per_batch_end_index():
+  """KV writes must honor each batch row's `end_index`, not only row 0."""
+  batch_size = 2
+  num_heads = 2
+  head_dim = 4
+  features = 8
+  cache_size = 8
+  seq_len = 1
+  query_pre_attn_scalar = head_dim**-0.5
+
+  attn = _modules.Attention(
+      num_heads=num_heads,
+      num_kv_heads=num_heads,
+      features=features,
+      head_dim=head_dim,
+      attn_type=_ATTN_TYPE,
+      query_pre_attn_scalar=query_pre_attn_scalar,
+  )
+
+  rng = jax.random.PRNGKey(0)
+  x = jax.random.normal(rng, (batch_size, seq_len, features))
+  segment_pos = jnp.array([[5], [2]], dtype=jnp.int32)
+  cache = _modules.Attention.init_cache(
+      cache_size=cache_size,
+      num_heads=num_heads,
+      head_dim=head_dim,
+      batch_size=batch_size,
+      dtype=jnp.float32,
+  )
+  cache['end_index'] = jnp.array([1, 4], dtype=jnp.int32)
+  cache['v'] = cache['v'] + jnp.arange(cache_size)[None, :, None, None]
+  cache['k'] = cache['k'] + jnp.arange(cache_size)[None, :, None, None]
+  attn_mask = jnp.ones((batch_size, seq_len, cache_size))
+
+  params = attn.init(rng, x, segment_pos, cache, attn_mask)
+  new_cache, _ = attn.apply(params, x, segment_pos, cache, attn_mask)
+
+  assert new_cache['end_index'].tolist() == [2, 5]
+  np.testing.assert_array_equal(new_cache['positions'][0, 1], segment_pos[0, 0])
+  np.testing.assert_array_equal(new_cache['positions'][1, 4], segment_pos[1, 0])
+  np.testing.assert_allclose(new_cache['v'][0, 0], cache['v'][0, 0])
+  np.testing.assert_allclose(new_cache['v'][1, 1], cache['v'][1, 1])
+  assert not jnp.allclose(new_cache['v'][0, 1], cache['v'][0, 1])
+  assert not jnp.allclose(new_cache['v'][1, 4], cache['v'][1, 4])
+  np.testing.assert_allclose(new_cache['v'][1, 1], cache['v'][1, 1])
+
+
 def test_sliding_window():
   num_heads = 2
   head_dim = 4