diff --git a/tests/models/gemma4/attention_test.py b/tests/models/gemma4/attention_test.py index 8d5ae38f1..fa9968c48 100644 --- a/tests/models/gemma4/attention_test.py +++ b/tests/models/gemma4/attention_test.py @@ -28,6 +28,7 @@ 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 config as config_lib from tunix.models.gemma4 import model as model_lib @@ -787,6 +788,633 @@ def test_find_last_one_index(self): np.array([2, 0, 0], dtype=np.int32), ) + def test_create_logical_sliding_window_mask_gapped_boundaries(self): + """Verify create_logical_sliding_window_mask with gapped mask.""" + # Mask with 6 valid slots across a gap: [0, 1, 2, 3, 4, 14] + cache_len = 16 + sw = 4 + mask_indices = [0, 1, 2, 3, 4, 14] + attn_mask = jnp.zeros((1, 1, cache_len), dtype=jnp.int32) + attn_mask = attn_mask.at[0, 0, mask_indices].set(1) + + result = attention_lib.create_logical_sliding_window_mask( + attn_mask, sliding_window_size=sw + ) + + # Valid count = 6, logical_last = 5, threshold = logical_last - sw = 1. + # Slot 1 has logical pos 1 (<= 1 threshold) -> False. + # Slot 2 has logical pos 2 (> 1 threshold) -> True. + # Slots 2, 3, 4, 14 are within the logical window of size 4. + self.assertFalse(bool(result[0, 0, 1])) + self.assertTrue(bool(result[0, 0, 2])) + self.assertEqual(int(jnp.sum(result)), 4) + expected = jnp.zeros((1, 1, cache_len), dtype=jnp.bool_) + expected = expected.at[0, 0, [2, 3, 4, 14]].set(True) + np.testing.assert_array_equal(result, expected) + + def test_create_logical_sliding_window_mask_contiguous_matches_physical(self): + """Contiguous mask produces identical result to physical sliding window.""" + attn_mask = jnp.array([[[1, 1, 1, 1, 0, 0]]], dtype=jnp.int32) + sw = 2 + logical_mask = attention_lib.create_logical_sliding_window_mask( + attn_mask, sliding_window_size=sw + ) + physical_mask = attention_lib.create_sliding_window_mask( + attn_mask, sliding_window_size=sw + ) + np.testing.assert_array_equal(logical_mask, physical_mask) + + def test_has_physical_gap_batched(self): + """Verify _has_physical_gap identifies gapped vs contiguous masks in batch.""" + cache_len = 16 + attn_mask = jnp.zeros((7, 1, cache_len), dtype=jnp.int32) + # row 0: contiguous [1, 1, 1, 1, 0, ...] -> False + attn_mask = attn_mask.at[0, 0, [0, 1, 2, 3]].set(1) + # row 1: multi-gap [1, 1, 0, 0, 1, 0, ...] -> True + attn_mask = attn_mask.at[1, 0, [0, 1, 4]].set(1) + # row 2: all zeros [0, 0, 0, ...] -> False + # row 3: single token [0, 0, 1, 0, ...] -> False + attn_mask = attn_mask.at[3, 0, [2]].set(1) + # row 4: single-token gap [1, 0, 1, 0, ...] (count=2, span=3) -> True + # Kills mutant that computes span = last - first without + 1 (2 < 2 -> False) + attn_mask = attn_mask.at[4, 0, [0, 2]].set(1) + # row 5: single-token gap with prefix [1, 1, 0, 1, 0, ...] (count=3, span=4) -> True + attn_mask = attn_mask.at[5, 0, [0, 1, 3]].set(1) + # row 6: offset contiguous [0, 0, 1, 1, 1, 0, ...] (count=3, span=3) -> False + attn_mask = attn_mask.at[6, 0, [2, 3, 4]].set(1) + + result = attention_lib._has_physical_gap(attn_mask) + self.assertEqual(result.shape, (7, 1, 1)) + expected = jnp.array([ + [[False]], + [[True]], + [[False]], + [[False]], + [[True]], + [[True]], + [[False]], + ]) + np.testing.assert_array_equal(result, expected) + + def test_read_prefix_kv_local_sliding_ring_buffer_unrolling(self): + """Verify unrolling of ring buffer in _read_prefix_kv for LOCAL_SLIDING.""" + config = model_lib.ModelConfig.gemma4_e2b() + config.use_sliding_window_kv_cache = True + config.sliding_window_size = 8 + attn = attention_lib.Attention( + config=config, + attn_type=model_lib.AttentionType.LOCAL_SLIDING, + rngs=nnx.Rngs(0), + ) + b, cache_len, seq_len = 1, 8, 4 + kh, d = config.num_kv_heads, config.head_dim + # Cache slots 0..7 with recognizable distinct values (e.g. 0.0, 10.0, 20.0, ...) + k_cached = jnp.broadcast_to( + (jnp.arange(cache_len, dtype=jnp.float32) * 10.0)[None, :, None, None], + (b, cache_len, kh, d), + ) + v_cached = jnp.broadcast_to( + (jnp.arange(cache_len, dtype=jnp.float32) * 10.0 + 1.0)[ + None, :, None, None + ], + (b, cache_len, kh, d), + ) + cache = { + 'k': k_cached, + 'v': v_cached, + 'end_index': jnp.array([12]), + } + key_proj = jnp.full((b, seq_len, kh, d), 100.0) + value_proj = jnp.full((b, seq_len, kh, d), 101.0) + prior_end_index = jnp.array([12]) + + res = attn._read_prefix_kv( + cache, + key_proj, + value_proj, + seq_len, + is_chunked_prefill=True, + prefix_length=8, + prior_end_index=prior_end_index, + ) + k_out, v_out, kv_valid_mask = res[0], res[1], res[2] + + # For cache_len=8, prior_end_index=12: + # valid_cached = 8, read_start = (12 - 8) % 8 = 4. + # Unrolled order of physical slots is [4, 5, 6, 7, 0, 1, 2, 3]. + expected_k_prefix = jnp.array( + [40.0, 50.0, 60.0, 70.0, 0.0, 10.0, 20.0, 30.0] + ) + expected_v_prefix = jnp.array( + [41.0, 51.0, 61.0, 71.0, 1.0, 11.0, 21.0, 31.0] + ) + + self.assertEqual(k_out.shape, (b, cache_len + seq_len, kh, d)) + self.assertEqual(v_out.shape, (b, cache_len + seq_len, kh, d)) + np.testing.assert_allclose(k_out[0, :cache_len, 0, 0], expected_k_prefix) + np.testing.assert_allclose(v_out[0, :cache_len, 0, 0], expected_v_prefix) + np.testing.assert_allclose( + k_out[0, cache_len:, 0, 0], jnp.full((seq_len,), 100.0) + ) + np.testing.assert_allclose( + v_out[0, cache_len:, 0, 0], jnp.full((seq_len,), 101.0) + ) + self.assertIsNotNone(kv_valid_mask) + np.testing.assert_array_equal( + kv_valid_mask, jnp.ones((cache_len,), dtype=jnp.bool_) + ) + + def test_eager_attention_chunked_prefill_mask_construction(self): + """Verify _eager_attention constructs chunked prefill mask for rectangular KV.""" + config = model_lib.ModelConfig.gemma4_e2b() + config.sliding_window_size = 8 + attn = attention_lib.Attention( + config=config, + attn_type=model_lib.AttentionType.GLOBAL, + rngs=nnx.Rngs(0), + ) + b, q_len, kv_len = 2, 4, 12 + h, kh, d = config.num_heads, config.num_kv_heads, config.head_dim + q = jnp.ones((b, q_len, h, d)) + k = jnp.zeros((b, kv_len, kh, d)) + v = jnp.zeros((b, kv_len, kh, d)) + + # Active Probe at valid prefix slot 2 (< prior_end_index=6) + k = k.at[:, 2, :, :].set(10.0) + v = v.at[:, 2, :, :].set(5.0) + + # Trap Probe 1 at uninitialized prefix position 7 (> prior_end_index=6) + k = k.at[:, 7, :, :].set(100.0) + v = v.at[:, 7, :, :].set(999.0) + + # Trap Probe 2 at future suffix position 11 (for query 0) + k = k.at[:, 11, :, :].set(100.0) + v = v.at[:, 11, :, :].set(999.0) + + prefix_mask = jnp.ones((b, q_len, 8), dtype=jnp.bool_) + suffix_causal = jnp.broadcast_to( + jnp.tril(jnp.ones((q_len, q_len), dtype=jnp.bool_))[None, :, :], + (b, q_len, q_len), + ) + attn_mask = jnp.concatenate([prefix_mask, suffix_causal], axis=-1) + segment_pos = jnp.broadcast_to( + jnp.arange(8, 12, dtype=jnp.int32)[None, :], (b, q_len) + ) + cache = { + 'k': jnp.zeros((b, 8, kh, d)), + 'v': jnp.zeros((b, 8, kh, d)), + 'end_index': jnp.array([6]), + } + 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=None, + prior_end_index=jnp.array([6]), + prefix_length=8, + seq_len=q_len, + is_chunked_prefill=True, + ) + self.assertEqual(out.shape, (b, q_len, h, d)) + self.assertFalse(jnp.isnan(out).any()) + np.testing.assert_allclose(out[:, 0, :, :], 5.0, atol=1e-2) + + def test_flash_attention_single_with_segment_ids(self): + """Verify _flash_attention_single execution path with segment_ids.""" + 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.GLOBAL, + rngs=nnx.Rngs(0), + ) + b, t = 2, 16 + 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_) + segment_ids = jnp.zeros((b, t), dtype=jnp.int32) + + kernel_called = [] + + def mock_kernel(q_in, k_in, v_in, segment_ids=None): + kernel_called.append(segment_ids) + self.assertIsNotNone(segment_ids) + return jnp.zeros_like(q_in) + + devices = np.array(jax.devices()[:1]).reshape(1, 1) + mesh = jax.sharding.Mesh(devices, ('fsdp', 'tp')) + + with mesh, mock.patch.object( + attn, '_make_splash_kernel', return_value=(mock_kernel, None) + ): + new_cache, out, (k_proj, v_proj, *_) = attn.block( + x, + segment_pos, + cache=None, + attn_mask=attn_mask, + segment_ids=segment_ids, + ) + self.assertTrue(kernel_called) + self.assertEqual(out.shape, (b, t, config.embed_dim)) + self.assertFalse(jnp.isnan(out).any()) + + def test_flash_attention_single_without_segment_ids(self): + """Verify _flash_attention_single execution path without segment_ids.""" + 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.GLOBAL, + rngs=nnx.Rngs(0), + ) + b, t = 2, 16 + 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_) + + called_with_segments = [] + + def mock_kernel(q_in, k_in, v_in, segment_ids=None): + called_with_segments.append(segment_ids) + return jnp.zeros_like(q_in) + + devices = np.array(jax.devices()[:1]).reshape(1, 1) + mesh = jax.sharding.Mesh(devices, ('fsdp', 'tp')) + + with mesh, mock.patch.object( + attn, '_make_splash_kernel', return_value=(mock_kernel, None) + ): + new_cache, out, (k_proj, v_proj, *_) = attn.block( + x, + segment_pos, + cache=None, + attn_mask=attn_mask, + segment_ids=None, + ) + self.assertEqual(called_with_segments, [None]) + self.assertEqual(out.shape, (b, t, config.embed_dim)) + self.assertEqual(k_proj.shape, (b, t, attn.num_kv_heads, attn.head_dim)) + self.assertEqual(v_proj.shape, (b, t, attn.num_kv_heads, attn.head_dim)) + self.assertFalse(jnp.isnan(out).any()) + + @parameterized.named_parameters( + dict( + testcase_name='exact_boundary', + prefix_length=128, + cache_len=1024, + boundaries=(0, 128, 256), + expected=128, + ), + dict( + testcase_name='in_between_rounds_up_to_next_boundary', + prefix_length=100, + cache_len=1024, + boundaries=(0, 128, 256), + expected=128, + ), + dict( + testcase_name='overflow_past_boundaries_falls_back_to_cache_len', + prefix_length=500, + cache_len=1024, + boundaries=(0, 128, 256), + expected=1024, + ), + dict( + testcase_name='beyond_cache_len_clamped_to_cache_len', + prefix_length=2000, + cache_len=1024, + boundaries=(0, 128, 256), + expected=1024, + ), + dict( + testcase_name='boundary_exceeds_cache_len_clamped_to_cache_len', + prefix_length=300, + cache_len=256, + boundaries=(0, 128, 512), + expected=256, + ), + dict( + testcase_name='empty_boundaries_ladder_falls_back_to_cache_len', + prefix_length=100, + cache_len=256, + boundaries=(), + expected=256, + ), + ) + def test_bucket_prefix_length( + self, prefix_length, cache_len, boundaries, expected + ): + self.assertEqual( + config_lib._bucket_prefix_length(prefix_length, cache_len, boundaries), + expected, + ) + + def test_maybe_bucket_prefix_length(self): + cache = {'v': jnp.zeros((1, 1024, 1, 64))} + boundaries = (0, 128, 256) + + # Chunked prefill buckets prefix_length using cache_len. + self.assertEqual( + config_lib._maybe_bucket_prefix_length( + 100, cache, is_chunked_prefill=True, boundaries=boundaries + ), + 128, + ) + + # Non-chunked prefill bypasses bucketing. + self.assertEqual( + config_lib._maybe_bucket_prefix_length( + 100, cache, is_chunked_prefill=False, boundaries=boundaries + ), + 100, + ) + + # prefix_length <= 0 bypasses bucketing. + self.assertEqual( + config_lib._maybe_bucket_prefix_length( + 0, cache, is_chunked_prefill=True, boundaries=boundaries + ), + 0, + ) + + # Empty boundaries ladder bypasses bucketing. + self.assertEqual( + config_lib._maybe_bucket_prefix_length( + 100, cache, is_chunked_prefill=True, boundaries=() + ), + 100, + ) + + def test_bucket_generation_ladders(self): + # pow2_buckets generates (0, 128, 256, ..., max_len) + pow2 = config_lib.pow2_buckets(max_len=1024) + self.assertEqual(pow2, (0, 128, 256, 512, 1024)) + pow2_default = config_lib.pow2_buckets() + self.assertEqual(pow2_default[0], 0) + self.assertEqual(pow2_default[1], 128) + self.assertEqual(pow2_default[-1], 131072) + + # linear_buckets generates (0, step, 2*step, ..., max_len) + linear = config_lib.linear_buckets(step=256, max_len=1024) + self.assertEqual(linear, (0, 256, 512, 768, 1024)) + linear_default = config_lib.linear_buckets() + self.assertEqual(linear_default[0], 0) + self.assertEqual(linear_default[1], 512) + self.assertEqual(linear_default[-1], 131072) + + # When max_len is not an exact multiple of step, or step=1, range upper bound + # must strictly equal max_len (kills range(0, max_len + 1 + 1, step) mutant at config.py:91). + self.assertEqual( + config_lib.linear_buckets(step=1, max_len=5), (0, 1, 2, 3, 4, 5) + ) + self.assertEqual(config_lib.linear_buckets(step=3, max_len=8), (0, 3, 6)) + + def test_update_cache_prefill_end_index_advances_by_max_real_tokens(self): + config = model_lib.ModelConfig.gemma4_e2b() + config.sliding_window_size = 16 + config.use_sliding_window_kv_cache = True + attn = attention_lib.Attention( + config=config, + attn_type=model_lib.AttentionType.GLOBAL, + rngs=nnx.Rngs(0), + ) + b, seq_len, cache_len = 2, 10, 16 + kh, d = config.num_kv_heads, config.head_dim + key_proj = jnp.zeros((b, seq_len, kh, d)) + value_proj = jnp.zeros((b, seq_len, kh, d)) + + # Case 1: Ragged input_mask. Row 0 has 7 tokens, row 1 has 4 tokens. + # Batch-max real tokens = 7. end_index should advance by 7. + input_mask = jnp.array( + [[1] * 7 + [0] * 3, [1] * 4 + [0] * 6], dtype=jnp.bool_ + ) + cache = { + 'k': jnp.zeros((b, cache_len, kh, d)), + 'v': jnp.zeros((b, cache_len, kh, d)), + 'end_index': jnp.array([5, 5], dtype=jnp.int32), + } + updated_cache, *_ = attn._update_cache_prefill( + cache, + key_proj, + value_proj, + seq_len=seq_len, + is_chunked_prefill=True, + prefix_length=5, + input_mask=input_mask, + ) + np.testing.assert_array_equal( + updated_cache['end_index'], jnp.array([12, 12], dtype=jnp.int32) + ) + + # Case 2: input_mask is None. end_index should advance by seq_len (10). + cache = { + 'k': jnp.zeros((b, cache_len, kh, d)), + 'v': jnp.zeros((b, cache_len, kh, d)), + 'end_index': jnp.array([5, 5], dtype=jnp.int32), + } + updated_cache_none, *_ = attn._update_cache_prefill( + cache, + key_proj, + value_proj, + seq_len=seq_len, + is_chunked_prefill=True, + prefix_length=5, + input_mask=None, + ) + np.testing.assert_array_equal( + updated_cache_none['end_index'], jnp.array([15, 15], dtype=jnp.int32) + ) + + def test_ragged_ring_buffer_decode_matches_unbatched_ground_truth(self): + """Ragged LOCAL_SLIDING prefill+decode must match B=1 unpadded ground truth.""" + config = model_lib.ModelConfig.gemma4_e2b() + config.use_sliding_window_kv_cache = True + config.sliding_window_size = 4 + attn = attention_lib.Attention( + config=config, + attn_type=model_lib.AttentionType.LOCAL_SLIDING, + rngs=nnx.Rngs(0), + ) + cache_len, kh, d = 4, config.num_kv_heads, config.head_dim + lens = jnp.array([12, 6], dtype=jnp.int32) # row 1 gap = 6 > window (4) + x_pre = jax.random.normal(jax.random.PRNGKey(0), (2, 12, config.embed_dim)) + x_dec = jax.random.normal(jax.random.PRNGKey(1), (2, 1, config.embed_dim)) + + def _run(row: int | None) -> jnp.ndarray: + idx = slice(None) if row is None else slice(row, row + 1) + b, seq = (2, 12) if row is None else (1, int(lens[row])) + in_mask = jnp.arange(seq)[None, :] < lens[idx, None] + causal = jnp.tril(jnp.ones((b, seq, seq), 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, _, _ = attn.block( + x_pre[idx, :seq], + jnp.broadcast_to(jnp.arange(seq, dtype=jnp.int32), (b, seq)), + cache=cache, + attn_mask=causal & in_mask[:, None, :], + input_mask=in_mask, + is_chunked_prefill=True, + force_eager=True, + ) + dec_mask = jnp.zeros((b, 1, seq + 1), dtype=jnp.bool_) + dec_mask = dec_mask.at[:, 0, :seq].set(in_mask) + dec_mask = dec_mask.at[:, 0, seq].set(True) + _, out, _ = attn.block( + x_dec[idx], + lens[idx, None], + cache=cache, + attn_mask=dec_mask, + force_eager=True, + ) + return out + + ragged_out = _run(None) + for r in range(len(lens)): + np.testing.assert_allclose( + ragged_out[r : r + 1], + _run(r), + atol=1e-2, + rtol=1e-2, + err_msg=f'Row {r} diverged from batch=1 ground truth!', + ) + + @parameterized.named_parameters( + dict( + testcase_name='sliding_window_within_window', + attn_type=model_lib.AttentionType.LOCAL_SLIDING, + sliding_window_size=16, + prefix_len=8, + suffix_len=4, + ), + dict( + testcase_name='sliding_window_at_window_boundary', + attn_type=model_lib.AttentionType.LOCAL_SLIDING, + sliding_window_size=16, + prefix_len=16, + suffix_len=4, + ), + dict( + testcase_name='sliding_window_beyond_window_wrapped', + attn_type=model_lib.AttentionType.LOCAL_SLIDING, + sliding_window_size=16, + prefix_len=24, + suffix_len=4, + ), + dict( + testcase_name='global_prefix_reuse', + attn_type=model_lib.AttentionType.GLOBAL, + sliding_window_size=None, + prefix_len=24, + suffix_len=4, + ), + ) + def test_chunked_prefill_prefix_reuse_matches_full_prefill( + self, attn_type, sliding_window_size, prefix_len, suffix_len + ): + config = model_lib.ModelConfig.gemma4_e2b() + config.sliding_window_size = sliding_window_size + config.use_sliding_window_kv_cache = True + config.use_flash_attention = False + + attn = attention_lib.Attention( + config=config, + attn_type=attn_type, + rngs=nnx.Rngs(0), + ) + + b = 2 + total_len = prefix_len + suffix_len + x_full = jax.random.normal( + jax.random.PRNGKey(0), (b, total_len, config.embed_dim) + ) + pos_full = jnp.broadcast_to( + jnp.arange(total_len, dtype=jnp.int32)[None, :], (b, total_len) + ) + mask_full = jnp.tril(jnp.ones((b, total_len, total_len), dtype=jnp.bool_)) + + # 1. Full monolithic prefill + max_seq_len = total_len + 16 + cache_full = attn.init_cache(b, max_seq_len, dtype=jnp.float32) + cache_full, out_full, _ = attn.block( + x_full, + pos_full, + cache=cache_full, + attn_mask=mask_full, + is_chunked_prefill=False, + ) + + # 2. Chunked prefill: Chunk 1 (prefix) + cache_chunked = attn.init_cache(b, max_seq_len, dtype=jnp.float32) + x_prefix = x_full[:, :prefix_len, :] + pos_prefix = pos_full[:, :prefix_len] + mask_prefix = jnp.tril( + jnp.ones((b, prefix_len, prefix_len), dtype=jnp.bool_) + ) + cache_chunked, _, _ = attn.block( + x_prefix, + pos_prefix, + cache=cache_chunked, + attn_mask=mask_prefix, + is_chunked_prefill=True, + prefix_length=0, + force_eager=True, + ) + + # 3. Chunked prefill: Chunk 2 (suffix reusing prefix cache) + x_suffix = x_full[:, prefix_len:, :] + pos_suffix = pos_full[:, prefix_len:] + mask_suffix = mask_full[:, prefix_len:, :] + cache_chunked, out_suffix_chunked, _ = attn.block( + x_suffix, + pos_suffix, + cache=cache_chunked, + attn_mask=mask_suffix, + is_chunked_prefill=True, + prefix_length=prefix_len, + force_eager=True, + ) + + # Verify suffix representation matches the suffix region of monolithic prefill + out_suffix_expected = out_full[:, prefix_len:, :] + np.testing.assert_allclose( + out_suffix_chunked, out_suffix_expected, atol=2e-5, rtol=1e-4 + ) + + # 4. Decode step at position total_len + x_decode = jax.random.normal( + jax.random.PRNGKey(1), (b, 1, config.embed_dim) + ) + pos_decode = jnp.array([[total_len], [total_len]], dtype=jnp.int32) + mask_decode = jnp.broadcast_to( + (jnp.arange(max_seq_len)[None, None, :] <= total_len), + (b, 1, max_seq_len), + ) + + cache_full, out_decode_full, _ = attn.block( + x_decode, + pos_decode, + cache=cache_full, + attn_mask=mask_decode, + is_chunked_prefill=False, + ) + cache_chunked, out_decode_chunked, _ = attn.block( + x_decode, + pos_decode, + cache=cache_chunked, + attn_mask=mask_decode, + is_chunked_prefill=False, + ) + np.testing.assert_allclose( + out_decode_chunked, out_decode_full, atol=2e-5, rtol=1e-4 + ) if __name__ == '__main__': absltest.main() diff --git a/tests/models/gemma4/model_test.py b/tests/models/gemma4/model_test.py index 5cb7385cd..bb3bbf986 100644 --- a/tests/models/gemma4/model_test.py +++ b/tests/models/gemma4/model_test.py @@ -16,12 +16,16 @@ from __future__ import annotations +from unittest import mock + from absl.testing import absltest from absl.testing import parameterized from flax import nnx import jax import jax.numpy as jnp +import numpy as np import qwix +from tunix.models.gemma4 import attention as attention_lib from tunix.models.gemma4 import model as model_lib @@ -719,6 +723,320 @@ def test_forward_pass_audio_heterogeneous(self): ) self.assertEqual(logits.shape, (batch_size, seq_len, config.num_embed)) + def test_forward_pass_chunked_prefill_with_kv_cache_sharing(self): + config = model_lib.ModelConfig.gemma4_e2b() + config.num_layers = 4 + config.num_embed = 128 + config.embed_dim = 128 + config.hidden_dim = 256 + config.num_heads = 2 + config.head_dim = 64 + config.num_kv_heads = 1 + config.sliding_window_size = 16 + config.use_sliding_window_kv_cache = True + config.frac_shared_layers = 0.5 + config.prefix_bucket_boundaries = (0, 16, 32) + config.attention_pattern = ( + model_lib.AttentionType.LOCAL_SLIDING, + model_lib.AttentionType.GLOBAL, + ) + + rngs = nnx.Rngs(0) + model = model_lib.Gemma4(config, rngs=rngs) + + # Initialize cache: layers 0 and 1 are unshared origins, layers 2 and 3 share. + cache = model.init_cache(batch_size=1, max_seq_len=16, dtype=jnp.float32) + cache['layer_0']['end_index'] = jnp.array([16]) + cache['layer_1']['end_index'] = jnp.array([16]) + + suffix_len = 8 + prefix_len = 16 + total_len = prefix_len + suffix_len + tokens = jax.random.randint( + jax.random.PRNGKey(0), (1, suffix_len), 0, config.num_embed + ) + positions = jnp.arange(prefix_len, total_len, dtype=jnp.int32)[None, :] + attn_mask = jnp.ones((1, suffix_len, total_len), dtype=jnp.bool_) + + logits, updated_cache = model( + tokens, + positions=positions, + cache=cache, + attention_mask=attn_mask, + is_chunked_prefill=True, + prefix_length=prefix_len, + ) + + self.assertEqual(logits.shape, (1, suffix_len, config.num_embed)) + self.assertFalse(jnp.isnan(logits).any()) + self.assertIsNotNone(updated_cache) + self.assertEqual(set(updated_cache.keys()), {'layer_0', 'layer_1'}) + self.assertNotIn('layer_2', updated_cache) + self.assertNotIn('layer_3', updated_cache) + self.assertEqual(int(updated_cache['layer_0']['end_index'][0]), 24) + self.assertEqual(int(updated_cache['layer_1']['end_index'][0]), 24) + + def test_forward_pass_chunked_prefill_shared_layer_respects_partial_valid_mask( + self, + ): + config = model_lib.ModelConfig.gemma4_e2b() + config.num_layers = 4 + config.num_embed = 128 + config.embed_dim = 128 + config.hidden_dim = 256 + config.num_heads = 2 + config.head_dim = 64 + config.num_kv_heads = 1 + config.sliding_window_size = 16 + config.use_sliding_window_kv_cache = True + config.frac_shared_layers = 0.5 + config.prefix_bucket_boundaries = (0, 16, 32) + config.attention_pattern = ( + model_lib.AttentionType.LOCAL_SLIDING, + model_lib.AttentionType.GLOBAL, + ) + + rngs = nnx.Rngs(0) + model = model_lib.Gemma4(config, rngs=rngs) + + # Initialize cache: 16 slots, partially filled with prefix_len=8. + cache_len = 16 + prefix_len = 8 + clean_cache = model.init_cache( + batch_size=1, max_seq_len=cache_len, dtype=jnp.float32 + ) + clean_cache['layer_0']['end_index'] = jnp.array([prefix_len]) + clean_cache['layer_1']['end_index'] = jnp.array([prefix_len]) + + suffix_len = 8 + total_len = cache_len + suffix_len + tokens = jax.random.randint( + jax.random.PRNGKey(0), (1, suffix_len), 0, config.num_embed + ) + positions = jnp.arange( + prefix_len, prefix_len + suffix_len, dtype=jnp.int32 + )[None, :] + attn_mask = jnp.ones((1, suffix_len, total_len), dtype=jnp.bool_) + + clean_logits, _ = model( + tokens, + positions=positions, + cache=clean_cache, + attention_mask=attn_mask, + is_chunked_prefill=True, + prefix_length=prefix_len, + ) + + # Inject 999.0 noise into uninitialized slots [8:16] of cache['layer_0']. + corrupt_cache = model.init_cache( + batch_size=1, max_seq_len=cache_len, dtype=jnp.float32 + ) + corrupt_cache['layer_0']['end_index'] = jnp.array([prefix_len]) + corrupt_cache['layer_1']['end_index'] = jnp.array([prefix_len]) + corrupt_cache['layer_0']['k'] = ( + corrupt_cache['layer_0']['k'].at[:, prefix_len:, ...].set(999.0) + ) + corrupt_cache['layer_0']['v'] = ( + corrupt_cache['layer_0']['v'].at[:, prefix_len:, ...].set(999.0) + ) + + corrupt_logits, _ = model( + tokens, + positions=positions, + cache=corrupt_cache, + attention_mask=attn_mask, + is_chunked_prefill=True, + prefix_length=prefix_len, + ) + + self.assertFalse(jnp.isnan(clean_logits).any()) + self.assertFalse(jnp.isnan(corrupt_logits).any()) + np.testing.assert_allclose(corrupt_logits, clean_logits, atol=1e-5) + + def test_shared_cache_prefix_bucketing_when_cache_is_none(self): + """Verifies Attention and DecoderLayer bucket prefix_length when cache is None and kv_shared_cache is passed.""" + config = model_lib.ModelConfig.gemma4_e2b() + config.num_layers = 2 + config.num_embed = 128 + config.embed_dim = 64 + config.hidden_dim = 128 + config.num_heads = 2 + config.head_dim = 32 + config.num_kv_heads = 2 + config.global_key_size = 32 + config.num_global_kv_heads = 2 + config.k_eq_v_global = False + config.sliding_window_size = 16 + config.prefix_bucket_boundaries = (0, 8, 16) + config.use_flash_attention = False + + rngs = nnx.Rngs(0) + decoder_layer = model_lib.GemmaDecoderLayer( + config, + attn_type=model_lib.AttentionType.GLOBAL, + rngs=rngs, + ) + attn = decoder_layer.attn + + prefix_len = 5 + expected_bucketed_prefix = 8 + suffix_len = 4 + total_kv_len = expected_bucketed_prefix + suffix_len + + kv_shared_cache = { + 'k': jnp.zeros((1, total_kv_len, attn.num_kv_heads, attn.head_dim)), + 'v': jnp.zeros((1, total_kv_len, attn.num_kv_heads, attn.head_dim)), + 'prior_end_index': jnp.array([prefix_len]), + } + + x = jnp.zeros((1, suffix_len, config.embed_dim)) + segment_pos = jnp.arange(prefix_len, prefix_len + suffix_len)[None, :] + attn_mask = jnp.ones((1, suffix_len, total_kv_len + 8), dtype=jnp.bool_) + + # 1. Test Attention.__call__ directly (non-remat) + with mock.patch.object(attn, 'block', wraps=attn.block) as mock_attn_block: + _, attn_out, _ = attn( + x, + segment_pos, + cache=None, + attn_mask=attn_mask, + kv_shared_cache=kv_shared_cache, + is_chunked_prefill=True, + prefix_length=prefix_len, + ) + self.assertEqual( + mock_attn_block.call_args.kwargs['prefix_length'], + expected_bucketed_prefix, + ) + self.assertEqual(attn_out.shape, x.shape) + self.assertFalse(jnp.isnan(attn_out).any()) + + # 2. Test Attention.__call__ with BLOCK remat + attn.config.remat_config = attention_lib.RematConfig.BLOCK + with mock.patch.object( + attention_lib, + '_maybe_bucket_prefix_length', + wraps=attention_lib._maybe_bucket_prefix_length, + ) as mock_bucket: + _, attn_out, _ = attn( + x, + segment_pos, + cache=None, + attn_mask=attn_mask, + kv_shared_cache=kv_shared_cache, + is_chunked_prefill=True, + prefix_length=prefix_len, + ) + self.assertEqual(mock_bucket.call_args.args[0], prefix_len) + self.assertIs(mock_bucket.call_args.args[1], kv_shared_cache) + self.assertEqual(attn_out.shape, x.shape) + self.assertFalse(jnp.isnan(attn_out).any()) + attn.config.remat_config = attention_lib.RematConfig.NONE + + # 3. Test GemmaDecoderLayer.__call__ + with mock.patch.object( + decoder_layer, 'block', wraps=decoder_layer.block + ) as mock_layer_block: + _, layer_out, _ = decoder_layer( + x, + segment_pos, + cache=None, + attn_mask=attn_mask, + kv_shared_cache=kv_shared_cache, + is_chunked_prefill=True, + prefix_length=prefix_len, + ) + self.assertEqual( + mock_layer_block.call_args.kwargs['prefix_length'], + expected_bucketed_prefix, + ) + self.assertEqual(layer_out.shape, x.shape) + self.assertFalse(jnp.isnan(layer_out).any()) + + # Sanity check helper behavior with vs without kv_shared_cache + unbucketed_len = attention_lib._maybe_bucket_prefix_length( + prefix_len, + None, + is_chunked_prefill=True, + boundaries=config.prefix_bucket_boundaries, + ) + self.assertEqual(unbucketed_len, prefix_len) + + bucketed_len = attention_lib._maybe_bucket_prefix_length( + prefix_len, + kv_shared_cache, + is_chunked_prefill=True, + boundaries=config.prefix_bucket_boundaries, + ) + self.assertEqual(bucketed_len, expected_bucketed_prefix) + + def test_chunked_prefill_bucket_padding_forces_eager(self): + """Verifies DecoderLayer forces eager attention if and only if bucket padding was introduced.""" + config = model_lib.ModelConfig.gemma4_e2b() + config.num_layers = 1 + config.num_embed = 128 + config.embed_dim = 64 + config.hidden_dim = 128 + config.num_heads = 2 + config.head_dim = 32 + config.num_kv_heads = 2 + config.global_key_size = 32 + config.num_global_kv_heads = 2 + config.k_eq_v_global = False + config.prefix_bucket_boundaries = (0, 8, 16) + + rngs = nnx.Rngs(0) + decoder_layer = model_lib.GemmaDecoderLayer( + config, + attn_type=model_lib.AttentionType.GLOBAL, + rngs=rngs, + ) + suffix_len = 4 + x = jnp.zeros((1, suffix_len, config.embed_dim)) + cache = { + 'k': jnp.zeros((1, 16, 2, 32)), + 'v': jnp.zeros((1, 16, 2, 32)), + 'end_index': jnp.array([16]), + } + + # Case 1: prefix_len = 5, bucketed_prefix = 8 (padding introduced -> force_eager must be True) + segment_pos = jnp.arange(5, 5 + suffix_len)[None, :] + attn_mask = jnp.ones((1, suffix_len, 8 + suffix_len), dtype=jnp.bool_) + with mock.patch.object( + decoder_layer, 'block', wraps=decoder_layer.block + ) as mock_block: + decoder_layer( + x, + segment_pos, + cache=cache, + attn_mask=attn_mask, + is_chunked_prefill=True, + prefix_length=5, + ) + self.assertTrue( + mock_block.call_args.kwargs['force_eager'], + 'Expected force_eager=True when bucketed_prefix != prefix_length', + ) + + # Case 2: prefix_len = 8, bucketed_prefix = 8 (no padding -> force_eager must remain False) + segment_pos = jnp.arange(8, 8 + suffix_len)[None, :] + attn_mask = jnp.ones((1, suffix_len, 8 + suffix_len), dtype=jnp.bool_) + with mock.patch.object( + decoder_layer, 'block', wraps=decoder_layer.block + ) as mock_block: + decoder_layer( + x, + segment_pos, + cache=cache, + attn_mask=attn_mask, + is_chunked_prefill=True, + prefix_length=8, + ) + self.assertFalse( + mock_block.call_args.kwargs['force_eager'], + 'Expected force_eager=False when bucketed_prefix == prefix_length', + ) + if __name__ == "__main__": absltest.main() diff --git a/tunix/models/gemma4/attention.py b/tunix/models/gemma4/attention.py index e70439eec..3322a2cdc 100644 --- a/tunix/models/gemma4/attention.py +++ b/tunix/models/gemma4/attention.py @@ -27,6 +27,7 @@ from jax.sharding import PartitionSpec as P import jaxtyping import numpy as np +from tunix.models.gemma4.config import _maybe_bucket_prefix_length from tunix.models.gemma4.config import AttentionType from tunix.models.gemma4.config import K_MASK from tunix.models.gemma4.config import LayerCache @@ -50,10 +51,10 @@ def find_last_one_index(attn_mask: jnp.ndarray) -> jnp.ndarray: # 2. reverse the rows in the attn_mask reversed_matrix = attn_mask[:, :, ::-1] - # 3. find the fist 1 from the right. + # 3. find the first 1 from the right. first_one_from_right = jnp.argmax(reversed_matrix, axis=-1) - # 4. covert back to the original index + # 4. convert back to the original index last_one_index_original = cache_len - 1 - first_one_from_right # 5. return the final index, 0 for rows are all zeros. @@ -108,6 +109,51 @@ def _get_causal_mask( return mask_lib.CausalMask((q_len, kv_len), offset=offset) +def create_logical_sliding_window_mask( + attn_mask: jnp.ndarray, # [B, 1, cache_len] (decoding: seq_len == 1) + sliding_window_size: int, +) -> jnp.ndarray: + """Sliding-window mask over LOGICAL token positions (for chunked decode). + + Physical-slot windowing (``create_sliding_window_mask``) assumes a contiguous + KV buffer layout. When warm-prefix chunked prefill leaves a physical padding + gap between prompt KV and suffix tokens, physical distance no longer matches + logical sequence distance, causing valid prompt tokens to be masked out + prematurely. This variant assigns each valid slot a contiguous logical + position + (cumsum of the validity mask) so the window follows the real tokens across the + gap, and reduces exactly to the physical version when the valid region is + contiguous. + """ + valid = attn_mask != 0 + valid_i = valid.astype(jnp.int32) + # Contiguous logical position of each valid slot; invalid slots are dropped by + # the `& valid` below, so their (meaningless) values do not matter. + logical_pos = jnp.cumsum(valid_i, axis=-1) - 1 # [B, 1, cache_len] + logical_last = jnp.sum(valid_i, axis=-1, keepdims=True) - 1 # [B, 1, 1] + window_mask = logical_pos > (logical_last - sliding_window_size) + final_mask = window_mask & valid + return final_mask # [B, 1, cache_len] + + +def _has_physical_gap(attn_mask: jnp.ndarray) -> jnp.ndarray: + """Per-row flag: is the valid region non-contiguous (a chunked-prefill gap)? + + Returns a ``[B, 1, 1]`` boolean. Standard left-pad decode has a contiguous + valid region (count == span) so this is False everywhere, and the caller's + ``jnp.where`` selects the original physical window unchanged -> byte-identical + normal-decode behavior. Only genuinely gapped rows switch to the logical mask. + """ + valid = attn_mask != 0 # [B, 1, cache_len] + n = attn_mask.shape[-1] + idx = jnp.arange(n) # [cache_len] + count = jnp.sum(valid.astype(jnp.int32), axis=-1, keepdims=True) # [B,1,1] + first = jnp.min(jnp.where(valid, idx, n), axis=-1, keepdims=True) # [B,1,1] + last = jnp.max(jnp.where(valid, idx, -1), axis=-1, keepdims=True) # [B,1,1] + span = last - first + 1 + return count < span # [B, 1, 1] + + class Attention(nnx.Module): """Attention module.""" @@ -247,6 +293,339 @@ def _compute_kv_projections( return key_proj, value_proj, kv_valid_mask + def _update_cache_prefill( + self, + cache: LayerCache, + key_proj: jaxtyping.Array, + value_proj: jaxtyping.Array, + seq_len: int, + *, + is_chunked_prefill: bool, + prefix_length: int, + input_mask: jaxtyping.Array | None, + ) -> tuple[ + LayerCache, + jaxtyping.Array, + jaxtyping.Array, + jaxtyping.Array | None, + jaxtyping.Array, + ]: + """Updates KV cache and prepares KV for attention during prefill. + + Delegates to _write_cache_prefill and _read_prefix_kv; these have no data + dependency, so XLA can overlap the cache write with attention. + """ + prior_end_index = cache['end_index'][0] + + # Write fresh KV to cache (independent of prefix read). + new_cache = self._write_cache_prefill( + cache, + key_proj, + value_proj, + seq_len, + is_chunked_prefill=is_chunked_prefill, + input_mask=input_mask, + ) + + # Read prefix KV from ORIGINAL cache (not new_cache: we want the pre-write + # state). + key_proj, value_proj, kv_valid_mask = self._read_prefix_kv( + cache, + key_proj, + value_proj, + seq_len, + is_chunked_prefill=is_chunked_prefill, + prefix_length=prefix_length, + prior_end_index=prior_end_index, + ) + + return ( + new_cache, + key_proj, + value_proj, + kv_valid_mask, + prior_end_index, + ) + + def _write_cache_prefill( + self, + cache: LayerCache, + key_proj: jaxtyping.Array, + value_proj: jaxtyping.Array, + seq_len: int, + *, + is_chunked_prefill: bool, + input_mask: jaxtyping.Array | None, + ) -> LayerCache: + """Writes fresh KV projections to cache. Returns updated cache. + + Separated from prefix read so XLA can overlap cache writes with attention. + """ + cache_len = cache['v'].shape[1] + prior_end_index = cache['end_index'][0] + + if self.config.use_sliding_window_kv_cache: + if is_chunked_prefill and input_mask is not None: + b = key_proj.shape[0] + prior = cache['end_index'] + n_r = jnp.sum(input_mask.astype(jnp.int32), axis=-1) + i = jnp.arange(cache_len) + cpos = (n_r[:, None] - cache_len) + i[None, :] + valid = (cpos >= 0) & (cpos < n_r[:, None]) + safe_cpos = jnp.clip(cpos, 0, seq_len - 1) + slot = (prior[:, None] + cpos) % cache_len + b_idx = jnp.arange(b)[:, None] + new_k = key_proj[b_idx, safe_cpos] + new_v = value_proj[b_idx, safe_cpos] + old_k = cache['k'][b_idx, slot] + old_v = cache['v'][b_idx, slot] + valid_4d = valid[:, :, None, None] + cache_k = ( + cache['k'].at[b_idx, slot].set(jnp.where(valid_4d, new_k, old_k)) + ) + cache_v = ( + cache['v'].at[b_idx, slot].set(jnp.where(valid_4d, new_v, old_v)) + ) + else: + end_index = prior_end_index + valid_len = min(seq_len, cache_len) + latest_indices = ( + end_index + (seq_len - valid_len) + jnp.arange(valid_len) + ) % cache_len + new_v = value_proj[:, -valid_len:, ...] + new_k = key_proj[:, -valid_len:, ...] + cache_v = cache['v'].at[:, latest_indices, ...].set(new_v) + cache_k = cache['k'].at[:, latest_indices, ...].set(new_k) + else: + end_index = prior_end_index + slice_indices = (0, end_index % cache_len, 0, 0) + cache_v = jax.lax.dynamic_update_slice( + cache['v'], value_proj, slice_indices + ) + cache_k = jax.lax.dynamic_update_slice( + cache['k'], key_proj, slice_indices + ) + + # Non-uniform (ragged) input masks are safe: PAD-position KVs are zeroed in + # _compute_kv_projections and excluded by the attention mask. + + return { + 'v': cache_v, + 'k': cache_k, + 'end_index': ( + cache['end_index'] + + ( + # PAD-safe: advance by the batch-max real-token count + # (elements may be ragged under PAD); PAD KVs are zeroed and + # attention-masked, so reserving up to the max is safe. + jnp.max(jnp.sum(input_mask, axis=-1)).astype(jnp.int32) + if is_chunked_prefill and input_mask is not None + else seq_len + ) + ), + } + + def _read_prefix_kv( + self, + cache: LayerCache, + key_proj: jaxtyping.Array, + value_proj: jaxtyping.Array, + seq_len: int, + *, + is_chunked_prefill: bool, + prefix_length: int, + prior_end_index: jaxtyping.Array, + ) -> tuple[ + jaxtyping.Array, + jaxtyping.Array, + jaxtyping.Array | None, + ]: + """Reads prefix KV from cache and concatenates with fresh KV. + + Separated from cache write so the read is independent of the write and can + be overlapped with attention by XLA. + """ + kv_valid_mask = None + cache_len = cache['v'].shape[1] + + if not (is_chunked_prefill and prefix_length > 0): + return key_proj, value_proj, kv_valid_mask + + # Clamp prefix_length to cache_len so mask and KV slice stay consistent: + # JAX slicing silently clamps, but the mask would use the unclamped value, + # creating a shape mismatch. + prefix_length = min(prefix_length, cache_len) + + if ( + self.config.use_sliding_window_kv_cache + and self.attn_type == AttentionType.LOCAL_SLIDING + ): + # LOCAL: Unroll ring buffer to get chronologically-ordered prefix KV + valid_cached = jnp.minimum(prior_end_index, cache_len) + read_start = (prior_end_index - valid_cached) % cache_len + i = jnp.arange(cache_len) + kv_valid_mask = i < valid_cached + physical_indices = (read_start + i) % cache_len + cached_k = cache['k'][:, physical_indices, ...] + cached_v = cache['v'][:, physical_indices, ...] + cached_k = jnp.where(kv_valid_mask[None, :, None, None], cached_k, 0) + cached_v = jnp.where(kv_valid_mask[None, :, None, None], cached_v, 0) + else: + # GLOBAL: Static slice for prefix KV. Use bucketed prefix_length for + # compilation stability; mask out padding positions dynamically. + cached_k = cache['k'][:, :prefix_length, ...] + cached_v = cache['v'][:, :prefix_length, ...] + # Zero out positions beyond the actual valid prefix. The bucketed + # prefix_length may exceed prior_end_index; those positions contain + # uninitialized cache data that must not influence attention. + valid_prefix = jnp.arange(prefix_length) < prior_end_index + cached_k = jnp.where(valid_prefix[None, :, None, None], cached_k, 0) + cached_v = jnp.where(valid_prefix[None, :, None, None], cached_v, 0) + + # Default: Concatenate cached prefix KV with fresh suffix KV. + key_proj = jnp.concatenate([cached_k, key_proj], axis=1) + value_proj = jnp.concatenate([cached_v, value_proj], axis=1) + + return key_proj, value_proj, kv_valid_mask + + def _build_chunked_prefill_mask( + self, + attn_mask: jaxtyping.Array, + q_len: int, + kv_len: int, + prior_end_index: jaxtyping.Array | None, + kv_shared_cache: LayerCache | None, + prefix_length: int, + kv_valid_mask: jaxtyping.Array | None, + has_own_cache: bool, + ) -> jaxtyping.Array: + """Constructs the attention mask for chunked prefill.""" + prefix_kv_len = kv_len - q_len + if ( + self.config.use_sliding_window_kv_cache + and self.attn_type == AttentionType.LOCAL_SLIDING + ): + return self._build_local_chunked_prefill_mask( + attn_mask, + q_len, + prefix_kv_len, + prior_end_index, + kv_shared_cache, + prefix_length, + kv_valid_mask, + has_own_cache, + ) + return self._build_global_chunked_prefill_mask( + attn_mask, + q_len, + kv_len, + prior_end_index, + kv_shared_cache, + prefix_length, + has_own_cache, + ) + + def _build_local_chunked_prefill_mask( + self, + attn_mask: jaxtyping.Array, + q_len: int, + prefix_kv_len: int, + prior_end_index: jaxtyping.Array | None, + kv_shared_cache: LayerCache | None, + prefix_length: int, + kv_valid_mask: jaxtyping.Array | None, + has_own_cache: bool, + ) -> jaxtyping.Array: + """Chunked-prefill attention mask for LOCAL_SLIDING layers.""" + # LOCAL: Build mask over [ring_buf | suffix] + if kv_valid_mask is not None: + local_cache_mask = jnp.broadcast_to( + kv_valid_mask[None, None, :], + (attn_mask.shape[0], q_len, prefix_kv_len), + ) + else: + local_cache_mask = jnp.ones( + (attn_mask.shape[0], q_len, prefix_kv_len), dtype=jnp.bool_ + ) + suffix_causal = attn_mask[..., -q_len:] + attn_mask = jnp.concatenate([local_cache_mask, suffix_causal], axis=-1) + # Use origin layer's prior_end_index for correct window boundaries. + if has_own_cache: + assert prior_end_index is not None + position_offset = prior_end_index + valid_cache_len = jnp.minimum(position_offset, prefix_kv_len) + elif kv_shared_cache is not None: + # Use the origin layer's prior_end_index if available (propagated + # via transient_kvs). Falls back to prefix_length if not present. + origin_end_index = kv_shared_cache.get('prior_end_index', None) + if origin_end_index is not None: + position_offset = origin_end_index + valid_cache_len = jnp.minimum(origin_end_index, prefix_kv_len) + else: + raise ValueError( + 'shared LOCAL layer missing origin prior_end_index; origin layers ' + 'must propagate it via transient_kvs' + ) + else: + position_offset = 0 + valid_cache_len = prefix_kv_len + row_pos = jnp.arange(q_len) + position_offset + col_pos_cache = jnp.arange(prefix_kv_len) + ( + position_offset - valid_cache_len + ) + col_pos_suffix = jnp.arange(q_len) + position_offset + col_pos = jnp.concatenate([col_pos_cache, col_pos_suffix]) + window_size = self.config.sliding_window_size + assert window_size is not None + sw_mask = (col_pos[None, :] > (row_pos[:, None] - window_size)) & ( + col_pos[None, :] <= row_pos[:, None] + ) + attn_mask = attn_mask & sw_mask[None, :, :] + return attn_mask + + def _build_global_chunked_prefill_mask( + self, + attn_mask: jaxtyping.Array, + q_len: int, + kv_len: int, + prior_end_index: jaxtyping.Array | None, + kv_shared_cache: LayerCache | None, + prefix_length: int, + has_own_cache: bool, + ) -> jaxtyping.Array: + """Chunked-prefill attention mask for GLOBAL layers.""" + # GLOBAL: Compose mask from prefix validity + suffix causal. + if prefix_length > 0: + prefix_mask = attn_mask[..., :prefix_length] + suffix_mask = attn_mask[..., -q_len:] + attn_mask = jnp.concatenate([prefix_mask, suffix_mask], axis=-1) + # Mask out uninitialized prefix cache positions. + if has_own_cache: + assert prior_end_index is not None + prefix_valid = jnp.arange(prefix_length) < prior_end_index + valid_mask = jnp.concatenate( + [prefix_valid, jnp.ones(q_len, dtype=jnp.bool_)] + ) + attn_mask = attn_mask & valid_mask[None, None, :] + elif kv_shared_cache is not None: + # Shared GLOBAL layers must also mask uninitialized prefix positions. + # Use the origin layer's prior_end_index propagated through + # kv_shared_cache. + origin_end_index = kv_shared_cache.get('prior_end_index', None) + if origin_end_index is None: + raise ValueError( + 'shared GLOBAL layer missing origin prior_end_index; origin ' + 'layers must propagate it via transient_kvs' + ) + prefix_valid = jnp.arange(prefix_length) < origin_end_index + valid_mask = jnp.concatenate( + [prefix_valid, jnp.ones(q_len, dtype=jnp.bool_)] + ) + attn_mask = attn_mask & valid_mask[None, None, :] + else: + attn_mask = attn_mask[..., :kv_len] + return attn_mask + def _build_flash_mask( self, q_len: int, @@ -264,12 +643,16 @@ def _build_flash_mask( 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: + def _make_block_sizes( + self, is_rectangular: bool, q_len: int | None = None + ) -> 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 q_len is not None: + block_q = min(block_q, q_len) if is_rectangular and self.attn_type == AttentionType.LOCAL_SLIDING: window_size = self.config.sliding_window_size assert window_size is not None @@ -347,8 +730,8 @@ def _make_splash_kernel( head_shards: int, q_seq_shards: int, mesh: shd.Mesh, - shd_n: str | None, - shd_t: str | None, + shd_n: AxisSpec, + shd_t: AxisSpec, save_residuals: bool = False, ): """Builds a splash MHA kernel and its manual sharding spec.""" @@ -372,6 +755,9 @@ def block( attn_mask: jaxtyping.Array, kv_shared_cache: LayerCache | None = None, segment_ids: jaxtyping.Array | None = None, + is_chunked_prefill: bool = False, + prefix_length: int = 0, + input_mask: jaxtyping.Array | None = None, force_eager: bool = False, ) -> tuple[ LayerCache | None, @@ -407,48 +793,47 @@ def block( assert kv_shared_cache is None cache_len = cache['v'].shape[1] if seq_len > 1: # prefill - 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 - ), - } + ( + new_cache, + key_proj, + value_proj, + kv_valid_mask, + prior_end_index, + ) = self._update_cache_prefill( + cache, + key_proj, + value_proj, + seq_len, + is_chunked_prefill=is_chunked_prefill, + prefix_length=prefix_length, + input_mask=input_mask, + ) + else: # decode + if ( + self.config.use_sliding_window_kv_cache + and self.attn_type == AttentionType.LOCAL_SLIDING + ): + b = value_proj.shape[0] + cache_len_local = cache['v'].shape[1] + abs_slot = cache['end_index'] % cache_len_local + logical_pos = ( + jnp.sum((attn_mask != 0).astype(jnp.int32), axis=-1)[:, 0] - 1 + ) + logical_slot = logical_pos % cache_len_local + has_gap = _has_physical_gap(attn_mask)[:, 0, 0] + slot = jnp.where(has_gap, logical_slot, abs_slot) + b_idx = jnp.arange(b) + value_proj = cache['v'].at[b_idx, slot].set(value_proj[:, 0]) + key_proj = cache['k'].at[b_idx, slot].set(key_proj[:, 0]) else: - slice_indices = (0, 0, 0, 0) - cache_v = jax.lax.dynamic_update_slice( + end_index = cache['end_index'][0] + slice_indices = (0, end_index % cache_len, 0, 0) + value_proj = jax.lax.dynamic_update_slice( cache['v'], value_proj, slice_indices ) - cache_k = jax.lax.dynamic_update_slice( + key_proj = jax.lax.dynamic_update_slice( cache['k'], key_proj, slice_indices ) - 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) - value_proj = jax.lax.dynamic_update_slice( - cache['v'], value_proj, slice_indices - ) - key_proj = jax.lax.dynamic_update_slice( - cache['k'], key_proj, slice_indices - ) new_cache = { 'v': value_proj, 'k': key_proj, @@ -470,8 +855,22 @@ def block( use_flash = ( self.config.use_flash_attention and seq_len > 1 + # Flash attention requires kv_len >= block_kv. Fall back to eager + # attention for short sequences/chunks smaller than block_kv. and kv_len >= self.config.flash_attention_block_size + # segment_ids are incompatible with rectangular flash because + # KV segment_ids would need to cover cached prefix positions. and not (is_rectangular and segment_ids is not None) + # GLOBAL layers are flash-eligible during rectangular chunked prefill. + # Bucketing stabilizes kv_len across chunks, preventing the + # recompilation storm that originally motivated this exclusion. + # Partial-cache LOCAL_SLIDING chunked prefill: flash's static relative + # offset (kv_len - q_len) anchors the window to the padded ring length, + # not the valid token count, so it slides past the real cached tokens + # when the window is only partially filled. Fall back to eager here. + # See cl/933189977. Keyed on RAW (pre-bucket) prefix_length in + # __call__ because bucketing rounds up toward the window and hides the + # ring gap. and not force_eager ) @@ -489,7 +888,7 @@ def block( multi_head_mask = mask_lib.MultiHeadMask([mask for _ in range(qh)]) - block_sizes = self._make_block_sizes(is_rectangular) + block_sizes = self._make_block_sizes(is_rectangular, q_len=q_len) ( shd_b, @@ -526,6 +925,7 @@ def block( shd_b, shd_t, ) + else: encoded = self._eager_attention( query_proj, @@ -535,7 +935,11 @@ def block( segment_pos, cache, kv_shared_cache, + kv_valid_mask, + prior_end_index, + prefix_length, seq_len, + is_chunked_prefill, ) attn_output = self.attn_vec_einsum(encoded) @@ -633,7 +1037,11 @@ def _eager_attention( segment_pos: jaxtyping.Array, cache: LayerCache | None, kv_shared_cache: LayerCache | None, - seq_len: int, + kv_valid_mask: jaxtyping.Array | None = None, + prior_end_index: jaxtyping.Array | None = None, + prefix_length: int = 0, + seq_len: int = 1, + is_chunked_prefill: bool = False, ) -> jaxtyping.Array: """Eager einsum attention (non-flash path).""" if self.use_gqa: @@ -652,9 +1060,27 @@ def _eager_attention( q_len = query_proj.shape[1] if seq_len > 1: - attn_mask = attn_mask[..., :kv_len] + if is_chunked_prefill and kv_len > q_len: + attn_mask = self._build_chunked_prefill_mask( + attn_mask, + q_len, + kv_len, + prior_end_index, + kv_shared_cache, + prefix_length, + kv_valid_mask, + has_own_cache=(cache is not None), + ) + else: + attn_mask = attn_mask[..., :kv_len] - if self.attn_type == AttentionType.LOCAL_SLIDING: + _skip_sliding_mask = ( + is_chunked_prefill + and kv_len > q_len + and self.config.use_sliding_window_kv_cache + and self.attn_type == AttentionType.LOCAL_SLIDING + ) + if self.attn_type == AttentionType.LOCAL_SLIDING and not _skip_sliding_mask: 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: @@ -668,27 +1094,46 @@ def _eager_attention( cache_len = key_proj.shape[1] end_idx = active_cache['end_index'] if cache is 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] + has_gap = _has_physical_gap(attn_mask) # [B, 1, 1] + logical_end = jnp.sum((attn_mask != 0).astype(jnp.int32), axis=-1) - 1 + eff_end = jnp.where(has_gap[:, :, 0], logical_end, end_idx[:, None]) + eff_end = eff_end[:, :, None] # [B, 1, 1] 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) + logical_indices = eff_end - ((eff_end - p) % cache_len) 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 + gathered = jnp.take_along_axis(attn_mask, logical_indices, axis=-1) + contiguous_mask = gathered * valid_physical + attn_mask = jnp.where( + has_gap, + valid_physical.astype(contiguous_mask.dtype), + contiguous_mask, + ) 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, ) + # Warm-prefix chunked prefill can leave a physical gap between an + # element's real prompt KV and its generated tokens. The physical-slot + # window above assumes contiguous positions and would drop real prompt + # tokens once the gap >= window. Recompute the window over LOGICAL + # positions and select it ONLY for rows whose valid region is non-contiguous, + # so standard left-pad decode (contiguous -> _has_physical_gap False) + # is byte-identical. + logical_sliding_mask = create_logical_sliding_window_mask( + attn_mask, + sliding_window_size=window_size, + ) + has_gap = _has_physical_gap(attn_mask) # [B, 1, 1] + sliding_mask = jnp.where(has_gap, logical_sliding_mask, sliding_mask) attn_mask = sliding_mask * attn_mask - else: # for prefill + else: # standard (non-chunked) prefill sliding window offset = kv_len - q_len all_ones = jnp.ones_like(attn_mask) sliding_mask = jnp.triu(all_ones, offset - window_size + 1) * jnp.tril( @@ -714,8 +1159,13 @@ def _eager_attention( @property def use_gqa(self) -> bool: + # Include MQA (num_kv_heads=1) in the GQA path. The GQA einsum + # correctly handles mismatched head counts via grouped reshape, + # whereas the non-GQA einsum ('BTNH,BSNH->BTNS') requires N to + # be equal between Q and K — which fails for MQA. return self.num_kv_heads != self.config.num_heads + @jax.named_scope('attention') def __call__( self, x: jaxtyping.Array, @@ -724,6 +1174,9 @@ def __call__( attn_mask: jaxtyping.Array, kv_shared_cache: LayerCache | None = None, segment_ids: jaxtyping.Array | None = None, + is_chunked_prefill: bool = False, + prefix_length: int = 0, + input_mask: jaxtyping.Array | None = None, force_eager: bool = False, ) -> tuple[ LayerCache | None, @@ -740,23 +1193,41 @@ def __call__( remat_config == RematConfig.BLOCK or remat_config == RematConfig.BLOCK.value ): - graphdef, state = nnx.split(self) - - def _checkpointed_block(state, *args, **kwargs): - module = nnx.merge(graphdef, state) - return module.block(*args, **kwargs) - - return jax.checkpoint(_checkpointed_block)( - state, - x, - segment_pos, - cache, - attn_mask, - kv_shared_cache, - segment_ids, - force_eager, + # nnx.remat needs to be applied to the unbound function and take self + # as the first argument. graph_updates=False prevents TraceContextError + # when mutating params across jax transformation trace levels. + # Bake static args via partial to avoid ConcretizationTypeError under remat. + # Bucket prefix_length to prevent a recompilation storm. + active_cache = cache if cache is not None else kv_shared_cache + bucketed_prefix = _maybe_bucket_prefix_length( + prefix_length, + active_cache, + is_chunked_prefill, + self.config.prefix_bucket_boundaries, + ) + block_fn = partial( + self.block.__func__, + is_chunked_prefill=is_chunked_prefill, + prefix_length=bucketed_prefix, + input_mask=input_mask, + force_eager=force_eager, ) + policy = getattr(jax.checkpoint_policies, self.config.remat_policy) + return nnx.remat( + block_fn, + graph_updates=False, + policy=policy, + )(self, x, segment_pos, cache, attn_mask, kv_shared_cache, segment_ids) else: + # Bucket prefix_length for the non-remat path too (controls static slice + # shapes which affect JAXPR identity). + active_cache = cache if cache is not None else kv_shared_cache + bucketed_prefix = _maybe_bucket_prefix_length( + prefix_length, + active_cache, + is_chunked_prefill, + self.config.prefix_bucket_boundaries, + ) return self.block( x, segment_pos, @@ -764,6 +1235,9 @@ def _checkpointed_block(state, *args, **kwargs): attn_mask, kv_shared_cache=kv_shared_cache, segment_ids=segment_ids, + is_chunked_prefill=is_chunked_prefill, + prefix_length=bucketed_prefix, + input_mask=input_mask, force_eager=force_eager, ) diff --git a/tunix/models/gemma4/config.py b/tunix/models/gemma4/config.py index 92851096d..85b341e36 100644 --- a/tunix/models/gemma4/config.py +++ b/tunix/models/gemma4/config.py @@ -14,6 +14,7 @@ """Gemma4 model configuration.""" +import bisect import dataclasses import enum from typing import Any, Tuple @@ -71,6 +72,51 @@ class RematConfig(enum.Enum): DECODER = enum.auto() +# Prefix-length bucketing prevents recompilation storms: each unique +# prefix_length baked via partial triggers a separate XLA compilation, so we +# round up to a fixed ladder of boundaries to bound the number of compiles. +# Use pow2_buckets() / linear_buckets() to build common ladders. +def pow2_buckets(max_len: int = 131072) -> tuple[int, ...]: + """Powers-of-two ladder: (0, 128, 256, ..., max_len). Default.""" + buckets = [0] + n = 128 + while n <= max_len: + buckets.append(n) + n *= 2 + return tuple(buckets) + + +def linear_buckets(step: int = 512, max_len: int = 131072) -> tuple[int, ...]: + """Linear ladder: (0, step, 2*step, ..., max_len). The 'x*512' case.""" + return tuple(range(0, max_len + 1, step)) + + +def _bucket_prefix_length( + prefix_length: int, cache_len: int, boundaries: tuple[int, ...] +) -> int: + """Round prefix_length up to the nearest bucket for compilation stability. + + Positions beyond the actual prefix_length are zeroed via dynamic masking, + so the padding is semantically invisible to the model. + """ + i = bisect.bisect_left(boundaries, prefix_length) + bucket = boundaries[i] if i < len(boundaries) else cache_len + return min(bucket, cache_len) + + +def _maybe_bucket_prefix_length( + prefix_length: int, + cache: LayerCache | None, + is_chunked_prefill: bool, + boundaries: tuple[int, ...], +) -> int: + """Buckets prefix_length during chunked prefill; passthrough otherwise.""" + if not is_chunked_prefill or prefix_length <= 0 or not boundaries: + return prefix_length + cache_len = cache['v'].shape[1] if cache is not None else prefix_length + return _bucket_prefix_length(prefix_length, cache_len, boundaries) + + @dataclasses.dataclass(slots=True, frozen=True) class ShardingConfig: """Sharding configuration for gemma transformer.""" @@ -178,6 +224,13 @@ class ModelConfig: # everything (minimum HBM). remat_policy: str = 'nothing_saveable' + # Prefix-length bucket ladder for chunked prefill. Each distinct bucketed + # prefix_length is baked as a static partial arg, so it is a separate XLA + # compilation; this ladder bounds the number of compiles. An empty tuple + # disables bucketing (passthrough, accepts recompiles). Use pow2_buckets() / + # linear_buckets() to build common ladders. Must be sorted and non-negative. + prefix_bucket_boundaries: tuple[int, ...] = linear_buckets(step=512) + # When True, the splash attention backward pass uses a single fused kernel # for dQ+dKV instead of two separate passes, reducing VMEM round-trips. # When enabled, block_q_dq and block_kv_dq are ignored (set to None). @@ -198,10 +251,13 @@ class ModelConfig: audio_encoder: audio.ConformerConfig | None = None def __post_init__(self): - # TODO(tunix-dev): support flash attention with sliding window KV cache - if self.use_sliding_window_kv_cache and self.use_flash_attention: + boundaries = self.prefix_bucket_boundaries + if any(b < 0 for b in boundaries) or boundaries != tuple( + sorted(set(boundaries)) + ): raise ValueError( - 'Flash attention and sliding window KV cache are mutually exclusive.' + 'prefix_bucket_boundaries must be non-negative and sorted strictly ' + f'ascending with no duplicates; got {boundaries}' ) @classmethod diff --git a/tunix/models/gemma4/model.py b/tunix/models/gemma4/model.py index a88d4986d..d3c97088f 100644 --- a/tunix/models/gemma4/model.py +++ b/tunix/models/gemma4/model.py @@ -17,7 +17,6 @@ import dataclasses from functools import partial import itertools -from typing import Any, Optional, Tuple from flax import nnx import jax from jax import numpy as jnp @@ -45,7 +44,11 @@ PreprocessedVisionInput, RematConfig, ShardingConfig, + _bucket_prefix_length, + _maybe_bucket_prefix_length, create_kv_cache_sharing_patterns, + linear_buckets, + pow2_buckets, ) from tunix.models.gemma4.layers import ( Einsum, @@ -60,6 +63,8 @@ ) from tunix.models.gemma4.attention import ( Attention, + _has_physical_gap, + create_logical_sliding_window_mask, create_sliding_window_mask, find_last_one_index, ) @@ -126,13 +131,14 @@ def __call__(self, x): remat_config == RematConfig.BLOCK or remat_config == RematConfig.BLOCK.value ): + policy = getattr(jax.checkpoint_policies, self.config.remat_policy) graphdef, state = nnx.split(self) def _checkpointed_block(state, *args, **kwargs): module = nnx.merge(graphdef, state) return module.block(*args, **kwargs) - return jax.checkpoint(_checkpointed_block)(state, x) + return jax.checkpoint(_checkpointed_block, policy=policy)(state, x) else: return self.block(x) @@ -252,6 +258,10 @@ def block( per_layer_input: jaxtyping.Array | None = None, kv_shared_cache: LayerCache | None = None, segment_ids: jaxtyping.Array | None = None, + is_chunked_prefill: bool = False, + prefix_length: int = 0, + input_mask: jaxtyping.Array | None = None, + force_eager: bool = False, ) -> tuple[ LayerCache | None, jaxtyping.Array, @@ -270,6 +280,10 @@ def block( attn_mask, kv_shared_cache=kv_shared_cache, segment_ids=segment_ids, + is_chunked_prefill=is_chunked_prefill, + prefix_length=prefix_length, + input_mask=input_mask, + force_eager=force_eager, ) attn = self.post_attention_norm(attn) attn += x @@ -306,6 +320,9 @@ def __call__( per_layer_input: jaxtyping.Array | None = None, kv_shared_cache: LayerCache | None = None, segment_ids: jaxtyping.Array | None = None, + is_chunked_prefill: bool = False, + prefix_length: int = 0, + input_mask: jaxtyping.Array | None = None, ) -> tuple[ LayerCache | None, jaxtyping.Array, @@ -316,18 +333,45 @@ def __call__( jaxtyping.Array | None, ], ]: + force_eager = ( + is_chunked_prefill + and self.attn.attn_type == AttentionType.LOCAL_SLIDING + and self.config.sliding_window_size is not None + and 0 < prefix_length < self.config.sliding_window_size + ) + active_cache = cache if cache is not None else kv_shared_cache + # Bucket prefix_length to prevent a recompilation storm. + bucketed_prefix = _maybe_bucket_prefix_length( + prefix_length, + active_cache, + is_chunked_prefill, + self.config.prefix_bucket_boundaries, + ) + # If bucket padding was introduced, force eager attention to prevent + # Flash Attention static masks from attending to zeroed ghost slots. + if is_chunked_prefill and bucketed_prefix != prefix_length: + force_eager = True + remat_config = getattr(self.config, 'remat_config', RematConfig.NONE) if ( remat_config == RematConfig.DECODER or remat_config == RematConfig.DECODER.value ): + policy = getattr(jax.checkpoint_policies, self.config.remat_policy) graphdef, state = nnx.split(self) - def _checkpointed_block(state, *args, **kwargs): + def _checkpointed_block(state, *args): module = nnx.merge(graphdef, state) - return module.block(*args, **kwargs) + return module.block( + *args, + segment_ids=segment_ids, + is_chunked_prefill=is_chunked_prefill, + prefix_length=bucketed_prefix, + input_mask=input_mask, + force_eager=force_eager, + ) - return jax.checkpoint(_checkpointed_block)( + return jax.checkpoint(_checkpointed_block, policy=policy)( state, x, segment_pos, @@ -335,7 +379,6 @@ def _checkpointed_block(state, *args, **kwargs): attn_mask, per_layer_input, kv_shared_cache, - segment_ids, ) else: return self.block( @@ -346,12 +389,19 @@ def _checkpointed_block(state, *args, **kwargs): per_layer_input, kv_shared_cache, segment_ids=segment_ids, + is_chunked_prefill=is_chunked_prefill, + prefix_length=bucketed_prefix, + input_mask=input_mask, + force_eager=force_eager, ) def init_cache(self, batch_size, max_seq_len, dtype): return self.attn.init_cache(batch_size, max_seq_len, dtype) +GemmaDecoderLayer = DecoderLayer + + class Gemma4(BackendMappingMixin, nnx.Module): """Gemma4 model.""" @@ -437,6 +487,9 @@ def __call__( images: PreprocessedVisionInput | None = None, audios: PreprocessedAudioInput | None = None, skip_lm_head: bool = False, + is_chunked_prefill: bool = False, + prefix_length: int = 0, + input_mask: jaxtyping.Array | None = None, ) -> tuple[jaxtyping.Array, Cache | None]: if positions is None: B, T = tokens.shape # pylint: disable=invalid-name @@ -531,6 +584,9 @@ def __call__( else None, kv_shared_cache=kv_shared_cache, segment_ids=segment_ids, + is_chunked_prefill=is_chunked_prefill, + prefix_length=prefix_length, + input_mask=input_mask, ) if is_prefill and i in self.shared_layer_origins: transient_kvs[layer_name] = layers_kvs