diff --git a/tests/models/gemma4/attention_test.py b/tests/models/gemma4/attention_test.py
new file mode 100644
index 000000000..8d5ae38f1
--- /dev/null
+++ b/tests/models/gemma4/attention_test.py
@@ -0,0 +1,792 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Gemma 4 Attention module and Pallas Splash kernels."""
+
+from __future__ import annotations
+
+import dataclasses
+from unittest import mock
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from flax import nnx
+import jax
+from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask as mask_lib
+import jax.numpy as jnp
+from jax.sharding import PartitionSpec as P
+import numpy as np
+from tunix.models.gemma4 import attention as attention_lib
+from tunix.models.gemma4 import model as model_lib
+
+
+class FlashAttentionMaskTest(parameterized.TestCase):
+  """Mask correctness unit tests (pure numpy — no model needed)."""
+
+  def test_local_mask_matches_manual(self):
+    """Verify LocalMask with offset produces the correct sliding window mask."""
+    chunk_len = 1024
+    sw_size = 512
+    cache_len = sw_size
+    kv_len = cache_len + chunk_len
+    prefix_len = cache_len
+
+    # Splash mask with offset
+    splash_mask = mask_lib.LocalMask(
+        (chunk_len, kv_len),
+        window_size=(sw_size - 1, 0),
+        offset=prefix_len,
+    )
+    splash_array = splash_mask[np.s_[:, :]]
+
+    # Reference mask for local sliding window with offset.
+    position_offset = prefix_len
+    valid_cache_len = prefix_len
+    row_pos = np.arange(chunk_len) + position_offset
+    col_pos_cache = np.arange(cache_len) + (position_offset - valid_cache_len)
+    col_pos_suffix = np.arange(chunk_len) + position_offset
+    col_pos = np.concatenate([col_pos_cache, col_pos_suffix])
+    manual_mask = (col_pos[None, :] > (row_pos[:, None] - sw_size)) & (
+        col_pos[None, :] <= row_pos[:, None]
+    )
+
+    np.testing.assert_array_equal(splash_array, manual_mask)
+
+  def test_causal_mask_matches_manual(self):
+    """Verify CausalMask with offset for GLOBAL chunked prefill."""
+    chunk_len = 1024
+    prefix_len = 2048
+    kv_len = prefix_len + chunk_len
+
+    splash_mask = mask_lib.CausalMask(
+        (chunk_len, kv_len),
+        offset=prefix_len,
+    )
+    splash_array = splash_mask[np.s_[:, :]]
+
+    # Manual: q[i] can attend to kv[j] where i + offset >= j
+    row = np.arange(chunk_len)[:, None] + prefix_len
+    col = np.arange(kv_len)[None, :]
+    manual_mask = row >= col
+
+    np.testing.assert_array_equal(splash_array, manual_mask)
+
+  @parameterized.parameters(
+      # (chunk_len, sw_size) — various sizes to test edge cases
+      (256, 128),
+      (512, 256),
+      (1024, 512),
+      (2048, 1024),
+  )
+  def test_local_mask_offset_parameterized(self, chunk_len, sw_size):
+    """LocalMask with offset is correct for various chunk/window sizes."""
+    cache_len = sw_size
+    kv_len = cache_len + chunk_len
+
+    splash_mask = mask_lib.LocalMask(
+        (chunk_len, kv_len),
+        window_size=(sw_size - 1, 0),
+        offset=cache_len,
+    )
+    splash_array = splash_mask[np.s_[:, :]]
+
+    # Each Q position q[i] at logical position (i + cache_len) should attend
+    # to KV positions in [i + cache_len - (sw_size - 1), i + cache_len].
+    for i in range(0, chunk_len, max(1, chunk_len // 8)):
+      logical_q = i + cache_len
+      expected_start = max(0, logical_q - (sw_size - 1))
+      expected_end = logical_q
+      # Verify True positions in row i
+      true_cols = np.where(splash_array[i])[0]
+      if len(true_cols) > 0:
+        self.assertEqual(true_cols[0], expected_start)
+        self.assertEqual(true_cols[-1], expected_end)
+        self.assertLen(true_cols, expected_end - expected_start + 1)
+
+  def test_local_mask_square_no_offset(self):
+    """Square LocalMask (chunk 1) should produce standard sliding window."""
+    seq_len = 512
+    sw_size = 128
+
+    splash_mask = mask_lib.LocalMask(
+        (seq_len, seq_len),
+        window_size=(sw_size - 1, 0),
+        offset=0,
+    )
+    splash_array = splash_mask[np.s_[:, :]]
+
+    # Manual: standard causal sliding window
+    row = np.arange(seq_len)[:, None]
+    col = np.arange(seq_len)[None, :]
+    manual_mask = (col <= row) & (col > row - sw_size)
+
+    np.testing.assert_array_equal(splash_array, manual_mask)
+
+  def test_build_flash_mask_local_sliding_rectangular(self):
+    config = model_lib.ModelConfig.gemma4_e2b()
+    config.sliding_window_size = 512
+    attn = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.LOCAL_SLIDING,
+        rngs=nnx.Rngs(0),
+    )
+    q_len, kv_len, sw = 128, 512, 512
+    offset = kv_len - q_len
+    mask = attn._build_flash_mask(q_len=q_len, kv_len=kv_len, offset=offset)
+    mask_array = mask[np.s_[:, :]]
+
+    q_ids = np.arange(q_len) + offset
+    kv_ids = np.arange(kv_len)
+    expected = (kv_ids[None, :] > (q_ids[:, None] - sw)) & (
+        kv_ids[None, :] <= q_ids[:, None]
+    )
+    np.testing.assert_array_equal(mask_array, expected)
+
+  def test_eager_attention_local_sliding_rectangular_mask(self):
+    """Verify eager attention local sliding window mask in rectangular prefill."""
+    q_len, kv_len, sw = 128, 512, 256
+    offset = kv_len - q_len
+    all_ones = jnp.ones((1, q_len, kv_len), dtype=jnp.bool_)
+    sliding_mask = jnp.triu(all_ones, offset - sw + 1) * jnp.tril(
+        all_ones, offset + sw - 1
+    )
+
+    q_ids = np.arange(q_len) + offset
+    kv_ids = np.arange(kv_len)
+    expected_sliding = (kv_ids[None, :] > (q_ids[:, None] - sw)) & (
+        kv_ids[None, :] < (q_ids[:, None] + sw)
+    )
+    np.testing.assert_array_equal(sliding_mask[0], expected_sliding)
+
+    # Combined with causal mask:
+    causal_mask = jnp.tril(all_ones, offset)
+    expected_causal_sliding = expected_sliding & (
+        kv_ids[None, :] <= q_ids[:, None]
+    )
+    np.testing.assert_array_equal(
+        (sliding_mask * causal_mask)[0], expected_causal_sliding
+    )
+
+  @parameterized.named_parameters(
+      dict(testcase_name='2d_mask', mask_3d=False),
+      dict(testcase_name='3d_mask', mask_3d=True),
+  )
+  def test_eager_attention_local_sliding_rectangular_execution(self, mask_3d):
+    """Verify _eager_attention with local sliding window on rectangular shapes."""
+    config = model_lib.ModelConfig.gemma4_e2b()
+    config.sliding_window_size = 4
+    attn = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.LOCAL_SLIDING,
+        rngs=nnx.Rngs(0),
+    )
+    b, q_len, kv_len, d = 2, 4, 16, config.head_dim
+    h, kh = config.num_heads, config.num_kv_heads
+    offset = kv_len - q_len
+
+    q = jax.random.normal(jax.random.PRNGKey(0), (b, q_len, h, d))
+    k = jax.random.normal(jax.random.PRNGKey(1), (b, kv_len, kh, d))
+    v = jax.random.normal(jax.random.PRNGKey(2), (b, kv_len, kh, d))
+    mask_shape = (b, q_len, kv_len) if mask_3d else (q_len, kv_len)
+    attn_mask = jnp.ones(mask_shape, dtype=jnp.bool_)
+    segment_pos = jnp.broadcast_to(
+        jnp.arange(offset, kv_len, dtype=jnp.int32)[None, :], (b, q_len)
+    )
+
+    out = attn._eager_attention(
+        query_proj=q,
+        key_proj=k,
+        value_proj=v,
+        attn_mask=attn_mask,
+        segment_pos=segment_pos,
+        cache=None,
+        kv_shared_cache=None,
+        seq_len=q_len,
+    )
+    self.assertEqual(out.shape, (b, q_len, h, d))
+    self.assertFalse(jnp.isnan(out).any())
+
+
+class FlashAttentionBlockSizeTest(parameterized.TestCase):
+  """Block-size divisibility parameterized test."""
+
+  @parameterized.parameters(
+      model_lib.ModelConfig.gemma4_e2b,
+      model_lib.ModelConfig.gemma4_e4b,
+      model_lib.ModelConfig.gemma4_31b,
+      model_lib.ModelConfig.gemma4_26b_a4b,
+  )
+  def test_block_kv_divisibility_and_chunk_multipliers(self, config_factory):
+    """block_kv must be 128-aligned and divide kv_len across chunk sizes."""
+    config = config_factory()
+    sw = config.sliding_window_size
+    block_q = config.flash_attention_block_size
+    block_kv = min(block_q, sw)
+
+    self.assertEqual(
+        block_kv % 128,
+        0,
+        f'block_kv={block_kv} not a multiple of 128 (NUM_LANES)',
+    )
+
+    for multiplier in (1, 2, 4):
+      chunk_len = block_q * multiplier
+      kv_len = sw + chunk_len
+      self.assertEqual(
+          chunk_len % block_q,
+          0,
+          f'chunk_len={chunk_len} not divisible by block_q={block_q}',
+      )
+      self.assertEqual(
+          kv_len % block_kv,
+          0,
+          f'kv_len={kv_len} not divisible by block_kv={block_kv} '
+          f'(multiplier={multiplier})',
+      )
+
+
+class AttentionTest(parameterized.TestCase):
+
+  def test_attention_with_segment_ids_rectangular_routing(self):
+    config = model_lib.ModelConfig.gemma4_e2b()
+    config.use_flash_attention = True
+    config.flash_attention_block_size = 16
+
+    attn = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.LOCAL_SLIDING,
+        rngs=nnx.Rngs(0),
+    )
+
+    b, t, h, d = 2, 32, config.num_heads, config.head_dim
+    x = jnp.zeros((b, t, config.embed_dim))
+    segment_pos = jnp.zeros((b, t), dtype=jnp.int32)
+    attn_mask = jnp.ones((b, t, t), dtype=jnp.bool_)
+
+    # Case 1: Square sequence, segment_ids is not None -> should use FLASH
+    with mock.patch.object(
+        attn, '_flash_attention_single'
+    ) as mock_flash, mock.patch.object(
+        attn, '_eager_attention'
+    ) as mock_eager, mock.patch.object(
+        attn, '_make_sharding_specs'
+    ) as mock_sharding, mock.patch.object(
+        attn, '_make_splash_kernel'
+    ) as mock_kernel:
+
+      mock_flash.return_value = (
+          jnp.zeros((b, t, h, d)),
+          jnp.zeros((b, t, config.num_kv_heads, d)),
+          jnp.zeros((b, t, config.num_kv_heads, d)),
+      )
+      mock_eager.return_value = jnp.zeros((b, t, h, d))
+      mock_sharding.return_value = (None,) * 4 + (1, 1) + (None,) * 3
+      mock_kernel.return_value = (None, None)
+
+      segment_ids = jnp.zeros((b, t), dtype=jnp.int32)
+
+      attn.block(
+          x,
+          segment_pos,
+          cache=None,
+          attn_mask=attn_mask,
+          segment_ids=segment_ids,
+      )
+
+      mock_flash.assert_called_once()
+      mock_eager.assert_not_called()
+
+    # Case 2: Rectangular sequence, segment_ids is not None -> should use EAGER
+    kv_len = 64
+    kv_shared_cache = {
+        'k': jnp.zeros((b, kv_len, config.num_kv_heads, d)),
+        'v': jnp.zeros((b, kv_len, config.num_kv_heads, d)),
+    }
+    attn_mask_rect = jnp.ones((b, t, kv_len), dtype=jnp.bool_)
+
+    with mock.patch.object(
+        attn, '_flash_attention_single'
+    ) as mock_flash, mock.patch.object(attn, '_eager_attention') as mock_eager:
+
+      mock_flash.return_value = (
+          jnp.zeros((b, t, h, d)),
+          jnp.zeros((b, kv_len, config.num_kv_heads, d)),
+          jnp.zeros((b, kv_len, config.num_kv_heads, d)),
+      )
+      mock_eager.return_value = jnp.zeros((b, t, h, d))
+
+      segment_ids = jnp.zeros((b, t), dtype=jnp.int32)
+
+      attn.block(
+          x,
+          segment_pos,
+          cache=None,
+          attn_mask=attn_mask_rect,
+          kv_shared_cache=kv_shared_cache,
+          segment_ids=segment_ids,
+      )
+
+      mock_flash.assert_not_called()
+      mock_eager.assert_called_once()
+
+  def test_make_block_sizes(self):
+    config = model_lib.ModelConfig.gemma4_e2b()
+    config.flash_attention_block_size = 128
+    config.sliding_window_size = 64
+    self.assertGreater(
+        config.flash_attention_block_size, config.sliding_window_size
+    )
+
+    global_attn = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.GLOBAL,
+        rngs=nnx.Rngs(0),
+    )
+    local_attn = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.LOCAL_SLIDING,
+        rngs=nnx.Rngs(0),
+    )
+
+    # GLOBAL rectangular uses the full block size (kills `if is_rectangular:`).
+    self.assertEqual(
+        global_attn._make_block_sizes(is_rectangular=True).block_kv,
+        config.flash_attention_block_size,
+    )
+
+    # LOCAL_SLIDING square uses the full block size (kills `if self.attn_type == LOCAL_SLIDING:`).
+    self.assertEqual(
+        local_attn._make_block_sizes(is_rectangular=False).block_kv,
+        config.flash_attention_block_size,
+    )
+
+    # LOCAL_SLIDING rectangular uses min(block_size, window_size).
+    self.assertEqual(
+        local_attn._make_block_sizes(is_rectangular=True).block_kv,
+        min(config.flash_attention_block_size, config.sliding_window_size),
+    )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='none_mesh',
+          act_btnh=None,
+          mesh_shape=None,
+          expected_head_shards=1,
+          expected_q_seq_shards=1,
+      ),
+      dict(
+          testcase_name='axis_not_in_mesh_defaults_to_one',
+          # act_btnh unpacks to (shd_b, shd_t, shd_n, shd_h). Use axis names for
+          # shd_t ('seq_axis') and shd_n ('model_axis') that are non-None but
+          # NOT present in the mesh.
+          act_btnh=P('fsdp', 'seq_axis', 'model_axis', None),
+          # Mesh axes are ('fsdp', 'x'); 'seq_axis' and 'model_axis' are absent.
+          mesh_shape={'fsdp': 1, 'x': 1},
+          # shd_n ('model_axis') and shd_t ('seq_axis') are not in the mesh, so
+          # both must fall back to 1. This kills the mutants that drop the
+          # `and shd_n in mesh.shape` / `and shd_t in mesh.shape` guards.
+          expected_head_shards=1,
+          expected_q_seq_shards=1,
+      ),
+      dict(
+          testcase_name='sharded_mesh',
+          act_btnh=P('fsdp', 'seq_axis', 'model_axis', None),
+          mesh_shape={'fsdp': 1, 'seq_axis': 4, 'model_axis': 2},
+          expected_head_shards=2,
+          expected_q_seq_shards=4,
+      ),
+  )
+  def test_make_sharding_specs(
+      self,
+      act_btnh,
+      mesh_shape,
+      expected_head_shards,
+      expected_q_seq_shards,
+  ):
+    config = model_lib.ModelConfig.gemma4_e2b()
+    if act_btnh is not None:
+      # ShardingConfig is frozen, so replace it.
+      config.shd_config = dataclasses.replace(
+          config.shd_config,
+          act_btnh=act_btnh,
+      )
+    attn = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.GLOBAL,
+        rngs=nnx.Rngs(0),
+    )
+    b = 1
+    kh = config.num_kv_heads
+    mesh = mock.MagicMock(shape=mesh_shape) if mesh_shape is not None else None
+    if mesh is not None and 'fsdp' in mesh.shape:
+      self.assertEqual(b % mesh.shape['fsdp'], 0)
+
+    specs = attn._make_sharding_specs(b, kh, mesh)
+    # Return tuple index 4 is head_shards, index 5 is q_seq_shards.
+    head_shards = specs[4]
+    q_seq_shards = specs[5]
+
+    self.assertEqual(head_shards, expected_head_shards)
+    self.assertEqual(q_seq_shards, expected_q_seq_shards)
+
+  def test_attention_flash_rectangular_offset(self):
+    config = model_lib.ModelConfig.gemma4_e2b()
+    config.use_flash_attention = True
+    config.flash_attention_block_size = 16
+
+    attn = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.LOCAL_SLIDING,
+        rngs=nnx.Rngs(0),
+    )
+
+    b, q_len, kv_len = 2, 32, 64
+    h, d = config.num_heads, config.head_dim
+    x = jnp.zeros((b, q_len, config.embed_dim))
+    segment_pos = jnp.zeros((b, q_len), dtype=jnp.int32)
+    attn_mask = jnp.ones((b, q_len, kv_len), dtype=jnp.bool_)
+    kv_shared_cache = {
+        'k': jnp.zeros((b, kv_len, config.num_kv_heads, d)),
+        'v': jnp.zeros((b, kv_len, config.num_kv_heads, d)),
+    }
+
+    with mock.patch.object(
+        attn, '_build_flash_mask', wraps=attn._build_flash_mask
+    ) as mock_mask, mock.patch.object(
+        attn, '_flash_attention_single'
+    ) as mock_flash, mock.patch.object(
+        attn, '_make_sharding_specs'
+    ) as mock_sharding, mock.patch.object(
+        attn, '_make_splash_kernel'
+    ) as mock_kernel:
+
+      mock_flash.return_value = (
+          jnp.zeros((b, q_len, h, d)),
+          jnp.zeros((b, kv_len, config.num_kv_heads, d)),
+          jnp.zeros((b, kv_len, config.num_kv_heads, d)),
+      )
+      mock_sharding.return_value = (None,) * 4 + (1, 1) + (None,) * 3
+      mock_kernel.return_value = (None, None)
+
+      attn.block(
+          x,
+          segment_pos,
+          cache=None,
+          attn_mask=attn_mask,
+          kv_shared_cache=kv_shared_cache,
+          segment_ids=None,
+      )
+
+      # Verify flash attention is called and the exact positive offset is passed
+      # (kv_len - q_len = 64 - 32 = 32), killing mutants that negate offset or
+      # pass 0.
+      mock_flash.assert_called_once()
+      mock_mask.assert_called_once_with(q_len, kv_len, kv_len - q_len)
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='own_cache',
+          # Case 1: Own cache is present (cache is not None).
+          # end_index = 4 -> position 4 is valid and attended to.
+          end_index=4,
+          use_shared_cache=False,
+      ),
+      dict(
+          testcase_name='shared_cache_adjusted_index',
+          # Case 2: Shared cache is present (cache is None, kv_shared_cache is
+          # not None). end_index = 5 -> adjusted to 4 by `end_idx - 1`.
+          end_index=5,
+          use_shared_cache=True,
+      ),
+  )
+  def test_eager_attention_decoding_sliding_window_cache_indexing(
+      self, end_index, use_shared_cache
+  ):
+    config = model_lib.ModelConfig.gemma4_e2b()
+    config.sliding_window_size = 8
+    config.use_sliding_window_kv_cache = True
+    attn = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.LOCAL_SLIDING,
+        rngs=nnx.Rngs(0),
+    )
+    b, q_len, cache_len, d = 1, 1, 8, config.head_dim
+    h, kh = config.num_heads, config.num_kv_heads
+
+    q = jnp.ones((b, q_len, h, d))
+    k = jnp.zeros((b, cache_len, kh, d))
+    k = k.at[:, 4, :, :].set(10.0)
+    k = k.at[:, 6, :, :].set(20.0)
+    v = jnp.ones((b, cache_len, kh, d))
+    v = v.at[:, 4, :, :].set(5.0)
+    v = v.at[:, 6, :, :].set(100.0)
+
+    attn_mask = jnp.ones((b, q_len, cache_len), dtype=jnp.bool_)
+    segment_pos = jnp.array([[4]], dtype=jnp.int32)
+
+    cache_dict = {'end_index': jnp.array([end_index])}
+    cache = None if use_shared_cache else cache_dict
+    kv_shared_cache = cache_dict if use_shared_cache else None
+
+    out = attn._eager_attention(
+        query_proj=q,
+        key_proj=k,
+        value_proj=v,
+        attn_mask=attn_mask,
+        segment_pos=segment_pos,
+        cache=cache,
+        kv_shared_cache=kv_shared_cache,
+        seq_len=1,
+    )
+    np.testing.assert_allclose(out, jnp.full_like(out, 5.0), atol=1e-2)
+
+  def test_sliding_window_kv_cache_prefill_over_cache_len(self):
+    config = model_lib.ModelConfig.gemma4_e2b()
+    config.sliding_window_size = 8
+    b, seq_len, cache_len = 1, 10, 8
+    d = config.head_dim
+    kh = config.num_kv_heads
+
+    x = jax.random.normal(jax.random.PRNGKey(0), (b, seq_len, config.embed_dim))
+    segment_pos = jnp.arange(seq_len, dtype=jnp.int32)[None, :]
+    attn_mask = jnp.ones((b, seq_len, seq_len), dtype=jnp.bool_)
+    cache = {
+        'k': jnp.zeros((b, cache_len, kh, d)),
+        'v': jnp.zeros((b, cache_len, kh, d)),
+        'end_index': jnp.zeros((b,), dtype=jnp.int32),
+    }
+
+    # 1. When use_sliding_window_kv_cache=True, circular update succeeds.
+    config.use_sliding_window_kv_cache = True
+    attn_sliding = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.LOCAL_SLIDING,
+        rngs=nnx.Rngs(0),
+    )
+    new_cache, _, (k_proj, v_proj, *_) = attn_sliding.block(
+        x,
+        segment_pos,
+        cache=cache,
+        attn_mask=attn_mask,
+        force_eager=True,
+    )
+
+    self.assertIsNotNone(new_cache)
+    self.assertEqual(new_cache['end_index'][0], seq_len)
+    valid_indices = (seq_len - cache_len + jnp.arange(cache_len)) % cache_len
+    np.testing.assert_allclose(
+        new_cache['k'][:, valid_indices, ...], k_proj[:, -cache_len:, ...]
+    )
+    np.testing.assert_allclose(
+        new_cache['v'][:, valid_indices, ...], v_proj[:, -cache_len:, ...]
+    )
+
+    # 2. When use_sliding_window_kv_cache=False, non-sliding prefill cannot
+    # exceed cache_len (dynamic_update_slice raises). This kills the mutant
+    # at line 419 that drops `self.config.use_sliding_window_kv_cache and`.
+    config.use_sliding_window_kv_cache = False
+    attn_standard = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.LOCAL_SLIDING,
+        rngs=nnx.Rngs(0),
+    )
+    with self.assertRaises(TypeError):
+      attn_standard.block(
+          x,
+          segment_pos,
+          cache=cache,
+          attn_mask=attn_mask,
+          force_eager=True,
+      )
+
+  def test_eager_attention_mha_non_gqa(self):
+    config = model_lib.ModelConfig.gemma4_e2b()
+    config.num_heads = 4
+    config.num_kv_heads = 4
+    attn = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.GLOBAL,
+        rngs=nnx.Rngs(0),
+    )
+    self.assertFalse(attn.use_gqa)
+    b, q_len, kv_len, d = 2, 4, 8, config.head_dim
+    h = config.num_heads
+
+    q = jnp.ones((b, q_len, h, d))
+    k = jnp.zeros((b, kv_len, h, d))
+    k = k.at[:, 2, :, :].set(10.0)
+    v = jnp.zeros((b, kv_len, h, d))
+    v = v.at[:, 2, :, :].set(3.0)
+
+    attn_mask = jnp.ones((b, q_len, kv_len), dtype=jnp.bool_)
+    segment_pos = jnp.broadcast_to(
+        jnp.arange(q_len, dtype=jnp.int32)[None, :], (b, q_len)
+    )
+
+    out = attn._eager_attention(
+        query_proj=q,
+        key_proj=k,
+        value_proj=v,
+        attn_mask=attn_mask,
+        segment_pos=segment_pos,
+        cache=None,
+        kv_shared_cache=None,
+        seq_len=q_len,
+    )
+    self.assertEqual(out.shape, (b, q_len, h, d))
+    np.testing.assert_allclose(out, jnp.full_like(out, 3.0), atol=1e-2)
+
+  def test_kv_cache_prefill_within_cache_len_and_decode_step(self):
+    config = model_lib.ModelConfig.gemma4_e2b()
+    config.use_sliding_window_kv_cache = False
+    attn = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.GLOBAL,
+        rngs=nnx.Rngs(0),
+    )
+    b, prefill_len, cache_len = 1, 4, 16
+    d = attn.head_dim
+    kh = attn.num_kv_heads
+
+    # 1. Prefill step with seq_len <= cache_len
+    x_prefill = jax.random.normal(
+        jax.random.PRNGKey(0), (b, prefill_len, config.embed_dim)
+    )
+    pos_prefill = jnp.arange(prefill_len, dtype=jnp.int32)[None, :]
+    mask_prefill = jnp.ones((b, prefill_len, prefill_len), dtype=jnp.bool_)
+    cache = {
+        'k': jnp.zeros((b, cache_len, kh, d)),
+        'v': jnp.zeros((b, cache_len, kh, d)),
+        'end_index': jnp.zeros((b,), dtype=jnp.int32),
+    }
+
+    cache, _, (k_prefill, v_prefill, *_) = attn.block(
+        x_prefill,
+        pos_prefill,
+        cache=cache,
+        attn_mask=mask_prefill,
+        force_eager=True,
+    )
+    self.assertEqual(cache['end_index'][0], prefill_len)
+    np.testing.assert_allclose(cache['k'][:, :prefill_len, ...], k_prefill)
+    np.testing.assert_allclose(cache['v'][:, :prefill_len, ...], v_prefill)
+
+    # 2. Decode step with seq_len == 1
+    x_decode = jax.random.normal(
+        jax.random.PRNGKey(1), (b, 1, config.embed_dim)
+    )
+    pos_decode = jnp.array([[prefill_len]], dtype=jnp.int32)
+    mask_decode = jnp.ones((b, 1, cache_len), dtype=jnp.bool_)
+
+    cache, _, (k_decode, v_decode, *_) = attn.block(
+        x_decode,
+        pos_decode,
+        cache=cache,
+        attn_mask=mask_decode,
+        force_eager=True,
+    )
+    self.assertEqual(cache['end_index'][0], prefill_len + 1)
+    np.testing.assert_allclose(cache['k'], k_decode)
+    np.testing.assert_allclose(cache['v'], v_decode)
+
+  def test_eager_attention_decoding_without_sliding_window_cache(self):
+    config = model_lib.ModelConfig.gemma4_e2b()
+    config.sliding_window_size = 4
+    config.use_sliding_window_kv_cache = False
+    attn = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.LOCAL_SLIDING,
+        rngs=nnx.Rngs(0),
+    )
+    b, q_len, kv_len, d = 1, 1, 8, config.head_dim
+    h, kh = config.num_heads, config.num_kv_heads
+
+    q = jnp.ones((b, q_len, h, d))
+    k = jnp.zeros((b, kv_len, kh, d))
+    # Query at position 6 with window 4 -> valid window [3, 6].
+    k = k.at[:, 5, :, :].set(10.0)
+    k = k.at[:, 1, :, :].set(20.0)
+    v = jnp.zeros((b, kv_len, kh, d))
+    v = v.at[:, 5, :, :].set(4.0)
+    v = v.at[:, 1, :, :].set(99.0)
+
+    # Causal mask up to position 6.
+    attn_mask = (
+        jnp.zeros((b, q_len, kv_len), dtype=jnp.bool_).at[:, :, :7].set(True)
+    )
+    segment_pos = jnp.array([[6]], dtype=jnp.int32)
+
+    out = attn._eager_attention(
+        query_proj=q,
+        key_proj=k,
+        value_proj=v,
+        attn_mask=attn_mask,
+        segment_pos=segment_pos,
+        cache=None,
+        kv_shared_cache=None,
+        seq_len=1,
+    )
+    self.assertEqual(out.shape, (b, q_len, h, d))
+    np.testing.assert_allclose(out, jnp.full_like(out, 4.0), atol=1e-2)
+
+  def test_eager_attention_decoding_missing_cache_raises(self):
+    config = model_lib.ModelConfig.gemma4_e2b()
+    config.sliding_window_size = 4
+    config.use_sliding_window_kv_cache = True
+    attn = attention_lib.Attention(
+        config=config,
+        attn_type=model_lib.AttentionType.LOCAL_SLIDING,
+        rngs=nnx.Rngs(0),
+    )
+    b, q_len, kv_len, d = 1, 1, 8, config.head_dim
+    h, kh = config.num_heads, config.num_kv_heads
+
+    q = jnp.ones((b, q_len, h, d))
+    k = jnp.ones((b, kv_len, kh, d))
+    v = jnp.ones((b, kv_len, kh, d))
+    attn_mask = jnp.ones((b, q_len, kv_len), dtype=jnp.bool_)
+    segment_pos = jnp.array([[4]], dtype=jnp.int32)
+
+    with self.assertRaisesRegex(
+        ValueError, 'Cache or shared cache is required'
+    ):
+      attn._eager_attention(
+          query_proj=q,
+          key_proj=k,
+          value_proj=v,
+          attn_mask=attn_mask,
+          segment_pos=segment_pos,
+          cache=None,
+          kv_shared_cache=None,
+          seq_len=1,
+      )
+
+  def test_find_last_one_index(self):
+    mask = jnp.array(
+        [
+            [[1, 1, 1, 0, 0]],
+            [[1, 0, 0, 0, 0]],
+            [[0, 0, 0, 0, 0]],
+        ],
+        dtype=jnp.int32,
+    )
+    last_indices = attention_lib.find_last_one_index(mask)
+    np.testing.assert_array_equal(
+        last_indices,
+        np.array([2, 0, 0], dtype=np.int32),
+    )
+
+
+if __name__ == '__main__':
+  absltest.main()
diff --git a/tunix/models/gemma4/attention.py b/tunix/models/gemma4/attention.py
index 76e67232d..c47d454f0 100644
--- a/tunix/models/gemma4/attention.py
+++ b/tunix/models/gemma4/attention.py
@@ -14,6 +14,7 @@
 
 """Gemma4 model attention."""
 
+import functools
 from functools import partial
 from flax import nnx
 import jax
@@ -36,6 +37,8 @@
 from tunix.models.gemma4.layers import RMSNorm
 from tunix.utils.sharding_utils import shard
 
+AxisSpec = str | tuple[str, ...] | None
+
 
 def find_last_one_index(attn_mask: jnp.ndarray) -> jnp.ndarray:
   """Finds the index of the last (rightmost) '1' from attn_mask."""
@@ -85,6 +88,26 @@ def create_sliding_window_mask(
   return final_mask[:, None, :]  # [B, 1, cache_len]
 
 
+@functools.lru_cache(maxsize=128)
+def _get_local_mask(
+    q_len: int, kv_len: int, window_size: int, offset: int
+) -> mask_lib.LocalMask:
+  """Memoized LocalMask constructor that speeds up XLA JIT compilation by caching mask closure objects across unrolled decoder layers."""
+  return mask_lib.LocalMask(
+      (q_len, kv_len),
+      window_size=(window_size - 1, 0),
+      offset=offset,
+  )
+
+
+@functools.lru_cache(maxsize=128)
+def _get_causal_mask(
+    q_len: int, kv_len: int, offset: int
+) -> mask_lib.CausalMask:
+  """Memoized CausalMask constructor that speeds up XLA JIT compilation by caching mask closure objects across unrolled decoder layers."""
+  return mask_lib.CausalMask((q_len, kv_len), offset=offset)
+
+
 class Attention(nnx.Module):
   """Attention module."""
 
@@ -187,36 +210,19 @@ def __init__(
         param_dtype=config.param_dtype,
     )
 
-  def block(
+  def _compute_kv_projections(
       self,
       x: jaxtyping.Array,
       segment_pos: jaxtyping.Array,
-      cache: LayerCache | None,
-      attn_mask: jaxtyping.Array,
-      kv_shared_cache: LayerCache | None = None,
-      segment_ids: jaxtyping.Array | None = None,
-  ) -> tuple[
-      LayerCache | None,
-      jaxtyping.Array,
-      tuple[jaxtyping.Array, jaxtyping.Array],
-  ]:
-    x = x.astype(self.config.dtype)
-    seq_len = x.shape[1]
-    query_proj = self.q_einsum(x)
-    query_proj = shard(query_proj, self.config.shd_config.act_btnh)
-    query_proj = self._query_norm(query_proj)
-    query_proj = apply_rope(
-        query_proj,
-        segment_pos,
-        base_frequency=self.rope_base_frequency,
-        scale_factor=self.rope_scale_factor,
-        rope_proportion=self.rope_proportion,
-    )
+      kv_shared_cache: LayerCache | None,
+  ) -> tuple[jaxtyping.Array, jaxtyping.Array, jaxtyping.Array | None]:
+    """Computes or retrieves key/value projections."""
+    kv_valid_mask = None
 
     if kv_shared_cache is not None:
-      assert cache is None
       key_proj = kv_shared_cache['k']
       value_proj = kv_shared_cache['v']
+      kv_valid_mask = kv_shared_cache.get('valid_mask', None)
     else:
       if hasattr(self, 'k_einsum'):  # case where k_eq_v is True
         key_proj = self.k_einsum(x)
@@ -239,35 +245,195 @@ def block(
           rope_proportion=self.rope_proportion,
       )
 
+    return key_proj, value_proj, kv_valid_mask
+
+  def _build_flash_mask(
+      self,
+      q_len: int,
+      kv_len: int,
+      offset: int,
+  ) -> mask_lib.Mask:
+    """Builds the single-head splash attention mask for one flash call.
+
+    Uses memoized computable masks (LocalMask / CausalMask) to prevent XLA
+    closure recompilations across unrolled layers and evaluate in-kernel.
+    """
+    if self.attn_type == AttentionType.LOCAL_SLIDING:
+      window_size = self.config.sliding_window_size
+      assert window_size is not None
+      return _get_local_mask(q_len, kv_len, window_size, offset)
+    return _get_causal_mask(q_len, kv_len, offset)
+
+  def _make_block_sizes(self, is_rectangular: bool) -> splash.BlockSizes:
+    """Selects splash block sizes for this attention call."""
+    # Choose block sizes. block_kv must divide kv_len.
+    # For LOCAL_SLIDING rectangular shapes, block_kv must divide both
+    # sliding_window_size and chunk_len. Use the smaller of the two.
+    block_q = self.config.flash_attention_block_size
+    if is_rectangular and self.attn_type == AttentionType.LOCAL_SLIDING:
+      window_size = self.config.sliding_window_size
+      assert window_size is not None
+      block_kv = min(
+          self.config.flash_attention_block_size,
+          window_size,
+      )
+    else:
+      block_kv = self.config.flash_attention_block_size
+
+    # Inner-loop tile size for the attention matmul; must divide block_kv.
+    block_kv_compute = min(
+        self.config.flash_attention_compute_block_size, block_kv
+    )
+
+    # Bwd holds Q + dO + attn_weights + grad accumulators; smaller Q block to
+    # fit VMEM.
+    block_bwd = min(self.config.flash_attention_bwd_block_size, block_q)
+    use_fused = self.config.flash_attention_use_fused_bwd
+    return splash.BlockSizes(
+        block_q=block_q,
+        block_kv=block_kv,
+        block_kv_compute=block_kv_compute,
+        block_q_dkv=block_bwd,
+        block_kv_dkv=block_kv,
+        block_kv_dkv_compute=block_kv_compute,
+        # Fused bwd kernel computes dQ+dKV in one pass; these are ignored.
+        block_q_dq=None if use_fused else block_bwd,
+        block_kv_dq=None if use_fused else block_kv,
+        use_fused_bwd_kernel=use_fused,
+    )
+
+  def _make_sharding_specs(self, b: int, kh: int, mesh: shd.Mesh):
+    """Computes mesh/shard-axis specs for splash attention."""
+    shd_b, shd_t, shd_n, shd_h = self.config.shd_config.act_btnh
+    if (
+        mesh is not None
+        and shd_b is not None
+        and shd_b in mesh.shape
+        and b % mesh.shape[shd_b] != 0
+    ):
+      shd_b = None
+    head_shards = (
+        mesh.shape[shd_n] if mesh is not None and shd_n in mesh.shape else 1
+    )
+    q_seq_shards = (
+        mesh.shape[shd_t] if mesh is not None and shd_t in mesh.shape else 1
+    )
+    shd_spec = P(shd_b, shd_n, shd_t, shd_h)
+    shd_n_kv = (
+        shd_n
+        if mesh is not None
+        and shd_n is not None
+        and shd_n in mesh.shape
+        and kh % mesh.shape[shd_n] == 0
+        else None
+    )
+    unsharded_seq_kv = P(shd_b, shd_n_kv, None, shd_h)
+    return (
+        shd_b,
+        shd_n,
+        shd_t,
+        shd_h,
+        head_shards,
+        q_seq_shards,
+        shd_n_kv,
+        shd_spec,
+        unsharded_seq_kv,
+    )
+
+  def _make_splash_kernel(
+      self,
+      multi_head_mask,
+      block_sizes: splash.BlockSizes,
+      head_shards: int,
+      q_seq_shards: int,
+      mesh: shd.Mesh,
+      shd_n: str | None,
+      shd_t: str | None,
+      save_residuals: bool = False,
+  ):
+    """Builds a splash MHA kernel and its manual sharding spec."""
+    kernel = splash.make_splash_mha(
+        multi_head_mask,
+        block_sizes=block_sizes,
+        head_shards=head_shards,
+        q_seq_shards=q_seq_shards,
+        save_residuals=save_residuals,
+    )
+    kernel_spec = kernel.manual_sharding_spec(
+        shd.NamedSharding(mesh, P(shd_n, shd_t))
+    )
+    return kernel, kernel_spec
+
+  def block(
+      self,
+      x: jaxtyping.Array,
+      segment_pos: jaxtyping.Array,
+      cache: LayerCache | None,
+      attn_mask: jaxtyping.Array,
+      kv_shared_cache: LayerCache | None = None,
+      segment_ids: jaxtyping.Array | None = None,
+      force_eager: bool = False,
+  ) -> tuple[
+      LayerCache | None,
+      jaxtyping.Array,
+      tuple[jaxtyping.Array, jaxtyping.Array],
+  ]:
+    x = x.astype(self.config.dtype)
+    seq_len = x.shape[1]
+    query_proj = self.q_einsum(x)
+    query_proj = shard(query_proj, self.config.shd_config.act_btnh)
+    query_proj = self._query_norm(query_proj)
+    query_proj = apply_rope(
+        query_proj,
+        segment_pos,
+        base_frequency=self.rope_base_frequency,
+        scale_factor=self.rope_scale_factor,
+        rope_proportion=self.rope_proportion,
+    )
+
+    key_proj, value_proj, kv_valid_mask = self._compute_kv_projections(
+        x,
+        segment_pos,
+        kv_shared_cache,
+    )
+
     if cache is not None:
       assert kv_shared_cache is None
-      # Update cache with new kv projections
       cache_len = cache['v'].shape[1]
       if seq_len > 1:  # prefill
-        if self.config.use_sliding_window_kv_cache:
-          # Sliding window cache update (prefill).
-          # Does not support chunked prefill.
-          valid_len = min(seq_len, cache_len)
-          latest_indices = jnp.arange(seq_len - valid_len, seq_len) % cache_len
-          cache_v = (
-              cache['v']
-              .at[:, latest_indices, ...]
-              .set(value_proj[:, -valid_len:, ...])
+        if self.config.use_sliding_window_kv_cache and seq_len > cache_len:
+          valid_indices = (
+              (seq_len - cache_len) + jnp.arange(cache_len)
+          ) % cache_len
+          new_v = value_proj[:, -cache_len:, ...]
+          new_k = key_proj[:, -cache_len:, ...]
+          cache_v = cache['v'].at[:, valid_indices, ...].set(new_v)
+          cache_k = cache['k'].at[:, valid_indices, ...].set(new_k)
+          new_cache = {
+              'v': cache_v,
+              'k': cache_k,
+              'end_index': jnp.full(
+                  (value_proj.shape[0],), seq_len, dtype=jnp.int32
+              ),
+          }
+        else:
+          slice_indices = (0, 0, 0, 0)
+          cache_v = jax.lax.dynamic_update_slice(
+              cache['v'], value_proj, slice_indices
           )
-          cache_k = (
-              cache['k']
-              .at[:, latest_indices, ...]
-              .set(key_proj[:, -valid_len:, ...])
+          cache_k = jax.lax.dynamic_update_slice(
+              cache['k'], key_proj, slice_indices
           )
-        else:
-          cache_v = cache['v'].at[:, :seq_len, ...].set(value_proj)
-          cache_k = cache['k'].at[:, :seq_len, ...].set(key_proj)
-
-        new_cache = {
-            'v': cache_v,
-            'k': cache_k,
-            'end_index': cache['end_index'] + seq_len,
-        }
+          new_cache = {
+              'v': cache_v,
+              'k': cache_k,
+              'end_index': jnp.full(
+                  (value_proj.shape[0],), seq_len, dtype=jnp.int32
+              ),
+          }
+        prior_end_index = None
+        split_prefix_k = None
+        split_prefix_v = None
       else:  # decode
         end_index = cache['end_index'][0]
         slice_indices = (0, end_index % cache_len, 0, 0)
@@ -291,217 +457,254 @@ def block(
     b, _, qh, _ = query_proj.shape
     _, _, kh, _ = key_proj.shape
 
-    if self.config.use_flash_attention and seq_len > 1:
+    # Determine if we can use flash attention for this call.
+    q_len = query_proj.shape[1]
+    kv_len = key_proj.shape[1]
+    is_rectangular = kv_len > q_len
+    use_flash = (
+        self.config.use_flash_attention
+        and seq_len > 1
+        and kv_len >= self.config.flash_attention_block_size
+        and not (is_rectangular and segment_ids is not None)
+        and not force_eager
+    )
+
+    if use_flash:
       query_proj = query_proj.transpose(0, 2, 1, 3)
       key_proj = key_proj.transpose(0, 2, 1, 3)
       value_proj = value_proj.transpose(0, 2, 1, 3)
 
       mesh = pxla.thread_resources.env.physical_mesh
-      if self.attn_type == AttentionType.LOCAL_SLIDING:
-        mask = mask_lib.LocalMask(
-            (seq_len, seq_len),
-            window_size=(self.config.sliding_window_size - 1, 0),  # pyrefly: ignore[unsupported-operation]
-            offset=0,
-        )
-      else:
-        mask = mask_lib.CausalMask((seq_len, seq_len))
+
+      # Offset: shifts Q positions so q[0] aligns with kv[prefix_len].
+      offset = kv_len - q_len if is_rectangular else 0
+
+      mask = self._build_flash_mask(q_len, kv_len, offset)
 
       multi_head_mask = mask_lib.MultiHeadMask([mask for _ in range(qh)])
 
-      block_sizes = splash.BlockSizes(
-          block_q=self.config.flash_attention_block_size,
-          block_kv=self.config.flash_attention_block_size,
-          block_q_dkv=self.config.flash_attention_block_size,
-          block_kv_dkv=self.config.flash_attention_block_size,
-          block_kv_dkv_compute=self.config.flash_attention_block_size,
-          block_q_dq=self.config.flash_attention_block_size,
-          block_kv_dq=self.config.flash_attention_block_size,
+      block_sizes = self._make_block_sizes(is_rectangular)
+
+      (
+          shd_b,
+          shd_n,
+          shd_t,
+          shd_h,
+          head_shards,
+          q_seq_shards,
+          shd_n_kv,
+          shd_spec,
+          unsharded_seq_kv,
+      ) = self._make_sharding_specs(b, kh, mesh)
+
+      splash_attn_kernel, kernel_spec = self._make_splash_kernel(
+          multi_head_mask,
+          block_sizes,
+          head_shards,
+          q_seq_shards,
+          mesh,
+          shd_n,
+          shd_t,
       )
 
-      shd_b, shd_t, shd_n, shd_h = self.config.shd_config.act_btnh
-      if (
-          mesh is not None
-          and shd_b is not None
-          and shd_b in mesh.shape
-          and b % mesh.shape[shd_b] != 0
-      ):
-        shd_b = None
-      head_shards = (
-          mesh.shape[shd_n] if shd_n is not None and shd_n in mesh.shape else 1
+      encoded, key_proj, value_proj = self._flash_attention_single(
+          query_proj,
+          key_proj,
+          value_proj,
+          segment_ids,
+          splash_attn_kernel,
+          kernel_spec,
+          shd_spec,
+          unsharded_seq_kv,
+          mesh,
+          shd_b,
+          shd_t,
       )
-      q_seq_shards = (
-          mesh.shape[shd_t] if shd_t is not None and shd_t in mesh.shape else 1
+    else:
+      encoded = self._eager_attention(
+          query_proj,
+          key_proj,
+          value_proj,
+          attn_mask,
+          segment_pos,
+          cache,
+          kv_shared_cache,
+          seq_len,
       )
 
-      splash_attn_kernel = splash.make_splash_mha(
-          multi_head_mask,
-          block_sizes=block_sizes,
-          head_shards=head_shards,
-          q_seq_shards=q_seq_shards,
+    attn_output = self.attn_vec_einsum(encoded)
+    attn_output = shard(attn_output, self.config.shd_config.act_btd)
+    return new_cache, attn_output, (key_proj, value_proj)
+
+  def _flash_attention_single(
+      self,
+      query_proj: jaxtyping.Array,
+      key_proj: jaxtyping.Array,
+      value_proj: jaxtyping.Array,
+      segment_ids: jaxtyping.Array | None,
+      splash_attn_kernel: splash.SplashAttentionKernel,
+      kernel_spec: splash.SplashAttentionKernel | None,
+      shd_spec: P,
+      unsharded_seq_kv: P,
+      mesh: shd.Mesh,
+      shd_b: AxisSpec,
+      shd_t: AxisSpec,
+  ) -> tuple[jaxtyping.Array, jaxtyping.Array, jaxtyping.Array]:
+    """Single-kernel flash attention over concatenated (or plain) KV."""
+    # Original single-kernel attention path.
+    if segment_ids is not None:
+      seg_spec = P(shd_b, shd_t)
+      unsharded_seg_spec = P(shd_b, None)
+
+      @partial(
+          shard_map,
+          mesh=mesh,
+          in_specs=(
+              kernel_spec,
+              shd_spec,
+              unsharded_seq_kv,
+              unsharded_seq_kv,
+              seg_spec,
+              unsharded_seg_spec,
+          ),
+          out_specs=shd_spec,
+          check_rep=False,
       )
+      def sharded_splash_attn(
+          kernel, q_block, k_block, v_block, q_seg_block, kv_seg_block
+      ):
+        seg_ids = splash.SegmentIds(q=q_seg_block, kv=kv_seg_block)
+        return jax.vmap(kernel)(q_block, k_block, v_block, segment_ids=seg_ids)
 
-      shd_spec = P(shd_b, shd_n, shd_t, shd_h)
-      shd_n_kv = (
-          shd_n
-          if mesh is not None
-          and shd_n is not None
-          and shd_n in mesh.shape
-          and kh % mesh.shape[shd_n] == 0
-          else None
+      qkv: jaxtyping.Array = sharded_splash_attn(
+          splash_attn_kernel,
+          query_proj,
+          key_proj,
+          value_proj,
+          segment_ids,
+          segment_ids,
       )
-      unsharded_seq_kv = P(shd_b, shd_n_kv, None, shd_h)
-      kernel_spec = splash_attn_kernel.manual_sharding_spec(
-          shd.NamedSharding(mesh, P(shd_n, shd_t))
+    else:
+
+      @partial(
+          shard_map,
+          mesh=mesh,
+          in_specs=(
+              kernel_spec,
+              shd_spec,
+              unsharded_seq_kv,
+              unsharded_seq_kv,
+          ),
+          out_specs=shd_spec,
+          check_rep=False,
       )
+      def sharded_splash_attn(kernel, q_block, k_block, v_block):
+        return jax.vmap(kernel)(q_block, k_block, v_block)
 
-      if segment_ids is not None:
-        seg_spec = P(shd_b, shd_t)
-        unsharded_seg_spec = P(shd_b, None)
-
-        @partial(
-            shard_map,
-            mesh=mesh,
-            in_specs=(
-                kernel_spec,
-                shd_spec,
-                unsharded_seq_kv,
-                unsharded_seq_kv,
-                seg_spec,
-                unsharded_seg_spec,
-            ),
-            out_specs=shd_spec,
-            check_rep=False,
-        )
-        def sharded_splash_attn(
-            kernel, q_block, k_block, v_block, q_seg_block, kv_seg_block
-        ):
-          seg_ids = splash.SegmentIds(q=q_seg_block, kv=kv_seg_block)
-          return jax.vmap(kernel)(
-              q_block, k_block, v_block, segment_ids=seg_ids
+      qkv: jaxtyping.Array = sharded_splash_attn(
+          splash_attn_kernel,
+          query_proj,
+          key_proj,
+          value_proj,
+      )
+    encoded = qkv.transpose(0, 2, 1, 3)
+    # Transpose KV back to (B, S, K, H); consumed by KV-sharing layers via
+    # layers_kvs.
+    key_proj = key_proj.transpose(0, 2, 1, 3)
+    value_proj = value_proj.transpose(0, 2, 1, 3)
+    return encoded, key_proj, value_proj
+
+  def _eager_attention(
+      self,
+      query_proj: jaxtyping.Array,
+      key_proj: jaxtyping.Array,
+      value_proj: jaxtyping.Array,
+      attn_mask: jaxtyping.Array,
+      segment_pos: jaxtyping.Array,
+      cache: LayerCache | None,
+      kv_shared_cache: LayerCache | None,
+      seq_len: int,
+  ) -> jaxtyping.Array:
+    """Eager einsum attention (non-flash path)."""
+    if self.use_gqa:
+      b, t, kg, h = query_proj.shape
+      n_groups = kg // self.num_kv_heads
+      query_reshaped = query_proj.reshape(
+          (b, t, self.num_kv_heads, n_groups, h)
+      )
+      logits = jnp.einsum('BTKGH,BSKH->BTKGS', query_reshaped, key_proj)
+      b, t, k, g, s = logits.shape
+      logits = logits.reshape((b, t, k * g, s))
+    else:
+      logits = jnp.einsum('BTNH,BSNH->BTNS', query_proj, key_proj)
+
+    kv_len = key_proj.shape[1]
+    q_len = query_proj.shape[1]
+
+    if seq_len > 1:
+      attn_mask = attn_mask[..., :kv_len]
+
+    if self.attn_type == AttentionType.LOCAL_SLIDING:
+      window_size = self.config.sliding_window_size
+      assert window_size is not None
+      if segment_pos.shape[1] == 1 and self.config.use_sliding_window_kv_cache:
+        # for decoding with sliding window cache
+        active_cache = cache if cache is not None else kv_shared_cache
+        if active_cache is None:
+          raise ValueError(
+              'Cache or shared cache is required for local sliding attention'
+              ' in decoding.'
           )
-
-        qkv: jaxtyping.Array = sharded_splash_attn(
-            splash_attn_kernel,
-            query_proj,
-            key_proj,
-            value_proj,
-            segment_ids,
-            segment_ids,
+        cache_len = key_proj.shape[1]
+        end_idx = active_cache['end_index']
+        if cache is None:
+          end_idx = end_idx - 1
+        end_idx = end_idx[:, None, None]
+        p = jnp.arange(cache_len)[None, None, :]
+
+        # map physical index to logical index
+        logical_indices = end_idx - ((end_idx - p) % cache_len)
+
+        # identify uninitialized slots (before the cache fills up)
+        valid_physical = logical_indices >= 0
+        logical_indices = jnp.maximum(0, logical_indices)
+
+        attn_mask = jnp.take_along_axis(attn_mask, logical_indices, axis=-1)
+        attn_mask = attn_mask * valid_physical
+      elif segment_pos.shape[1] == 1:
+        # for decoding without sliding window cache
+        sliding_mask = create_sliding_window_mask(
+            attn_mask,
+            sliding_window_size=window_size,
         )
-      else:
-
-        @partial(
-            shard_map,
-            mesh=mesh,
-            in_specs=(
-                kernel_spec,
-                shd_spec,
-                unsharded_seq_kv,
-                unsharded_seq_kv,
-            ),
-            out_specs=shd_spec,
-            check_rep=False,
-        )
-        def sharded_splash_attn(kernel, q_block, k_block, v_block):
-          return jax.vmap(kernel)(q_block, k_block, v_block)
-
-        qkv: jaxtyping.Array = sharded_splash_attn(
-            splash_attn_kernel,
-            query_proj,
-            key_proj,
-            value_proj,
-        )
-      encoded = qkv.transpose(0, 2, 1, 3)
-      query_proj = query_proj.transpose(0, 2, 1, 3)
-      key_proj = key_proj.transpose(0, 2, 1, 3)
-      value_proj = value_proj.transpose(0, 2, 1, 3)
-
-    else:
-      if self.use_gqa:
-        b, t, kg, h = query_proj.shape
-        n_groups = kg // self.num_kv_heads
-        query_reshaped = query_proj.reshape(
-            (b, t, self.num_kv_heads, n_groups, h)
+        attn_mask = sliding_mask * attn_mask
+      else:  # for prefill
+        offset = kv_len - q_len
+        all_ones = jnp.ones_like(attn_mask)
+        sliding_mask = jnp.triu(all_ones, offset - window_size + 1) * jnp.tril(
+            all_ones, offset + window_size - 1
         )
-        logits = jnp.einsum('BTKGH,BSKH->BTKGS', query_reshaped, key_proj)
-        b, t, k, g, s = logits.shape
-        logits = logits.reshape((b, t, k * g, s))
-      else:
-        logits = jnp.einsum('BTNH,BSNH->BTNS', query_proj, key_proj)
-
-      if seq_len > 1:
-        # Only compute attention scores for the actual sequence length.
-        attn_mask = attn_mask[..., :seq_len]
-
-      if self.attn_type == AttentionType.LOCAL_SLIDING:
-        if (
-            segment_pos.shape[1] == 1
-            and self.config.use_sliding_window_kv_cache
-        ):
-          # for decoding with sliding window cache
-          active_cache = cache if cache is not None else kv_shared_cache
-          if active_cache is None:
-            raise ValueError(
-                'Cache or shared cache is required for local sliding attention'
-                ' in decoding.'
-            )
-          cache_len = key_proj.shape[1]
-          end_idx = active_cache['end_index']
-          if cache is None and kv_shared_cache is not None:
-            # In case of shared KV cache, the origin layer already updated the
-            # end index. We need to subtract 1 to get the correct end index of
-            # the previous token.
-            end_idx = end_idx - 1
-          end_idx = end_idx[:, None, None]
-          p = jnp.arange(cache_len)[None, None, :]
-
-          # map physical index to logical index
-          logical_indices = end_idx - ((end_idx - p) % cache_len)
-
-          # identify uninitialized slots (before the cache fills up)
-          valid_physical = logical_indices >= 0
-          logical_indices = jnp.maximum(0, logical_indices)
-
-          attn_mask = jnp.take_along_axis(attn_mask, logical_indices, axis=-1)
-          attn_mask = attn_mask * valid_physical
-        elif segment_pos.shape[1] == 1:
-          # for decoding without sliding window cache
-          sliding_mask = create_sliding_window_mask(
-              attn_mask,
-              sliding_window_size=self.config.sliding_window_size,  # pyrefly: ignore[bad-argument-type]
-          )
-          attn_mask = sliding_mask * attn_mask
-        else:  # for prefill
-          all_ones = jnp.ones_like(attn_mask)
-          sliding_mask = jnp.triu(
-              all_ones, -1 * self.config.sliding_window_size + 1  # pyrefly: ignore[unsupported-operation]
-          ) * jnp.tril(
-              all_ones, self.config.sliding_window_size - 1  # pyrefly: ignore[unsupported-operation]
-          )
-          attn_mask = sliding_mask * attn_mask
-
-      attn = jnp.where((jnp.expand_dims(attn_mask, -2)), logits, K_MASK)
-      attn = jax.nn.softmax(attn.astype(jnp.float32), axis=-1).astype(
-          key_proj.dtype
-      )
+        attn_mask = sliding_mask * attn_mask
 
-      if self.use_gqa:
-        b, t, kg, s = attn.shape
-        n_groups = kg // self.num_kv_heads
-        probs_reshaped = attn.reshape((b, t, self.num_kv_heads, n_groups, s))
-        encoded = jnp.einsum('BTKGS,BSKH->BTKGH', probs_reshaped, value_proj)
-        b, t, k, g, h = encoded.shape
-        encoded = encoded.reshape((b, t, k * g, h))
-      else:
-        encoded = jnp.einsum('BTNS,BSNH->BTNH', attn, value_proj)
+    attn = jnp.where((jnp.expand_dims(attn_mask, -2)), logits, K_MASK)
+    attn = jax.nn.softmax(attn.astype(jnp.float32), axis=-1).astype(
+        key_proj.dtype
+    )
 
-    attn_output = self.attn_vec_einsum(encoded)
-    attn_output = shard(attn_output, self.config.shd_config.act_btd)
-    return new_cache, attn_output, (key_proj, value_proj)
+    if self.use_gqa:
+      b, t, kg, s = attn.shape
+      n_groups = kg // self.num_kv_heads
+      probs_reshaped = attn.reshape((b, t, self.num_kv_heads, n_groups, s))
+      encoded = jnp.einsum('BTKGS,BSKH->BTKGH', probs_reshaped, value_proj)
+      b, t, k, g, h = encoded.shape
+      encoded = encoded.reshape((b, t, k * g, h))
+    else:
+      encoded = jnp.einsum('BTNS,BSNH->BTNH', attn, value_proj)
+    return encoded
 
   @property
-  def use_gqa(self):
-    return self.num_kv_heads != self.config.num_heads and self.num_kv_heads > 1
+  def use_gqa(self) -> bool:
+    return self.num_kv_heads != self.config.num_heads
 
   def __call__(
       self,
@@ -511,6 +714,7 @@ def __call__(
       attn_mask: jaxtyping.Array,
       kv_shared_cache: LayerCache | None = None,
       segment_ids: jaxtyping.Array | None = None,
+      force_eager: bool = False,
   ) -> tuple[
       LayerCache | None,
       jaxtyping.Array,
@@ -528,7 +732,14 @@ def _checkpointed_block(state, *args, **kwargs):
         return module.block(*args, **kwargs)
 
       return jax.checkpoint(_checkpointed_block)(
-          state, x, segment_pos, cache, attn_mask, kv_shared_cache, segment_ids
+          state,
+          x,
+          segment_pos,
+          cache,
+          attn_mask,
+          kv_shared_cache,
+          segment_ids,
+          force_eager,
       )
     else:
       return self.block(
@@ -538,6 +749,7 @@ def _checkpointed_block(state, *args, **kwargs):
           attn_mask,
           kv_shared_cache=kv_shared_cache,
           segment_ids=segment_ids,
+          force_eager=force_eager,
       )
 
   def init_cache(