From 2cdd6c4b602be0a3b4f1e38c296af850494a2b36 Mon Sep 17 00:00:00 2001 From: Jacky Fang Date: Thu, 20 Aug 2026 03:35:41 +0000 Subject: [PATCH] fix: compute local block padding for dynamic splash indexer mask in CP --- src/maxtext/layers/attention_op.py | 6 ++++-- tests/unit/attention_test.py | 3 --- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 003aed30fa..8a5cd7b897 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -2111,8 +2111,10 @@ def wrap_flash_attention( if indexer_mask is not None: # Convert additive float mask (0.0=attend, negative=masked) to boolean mask for Tokamax splash kernel indexer_mask = indexer_mask == 0.0 - pad_q = mask_shape[0] - indexer_mask.shape[-2] - pad_kv = mask_shape[1] - indexer_mask.shape[-1] + padded_q_len = ((indexer_mask.shape[-2] + sa_config.block_q - 1) // sa_config.block_q) * sa_config.block_q + padded_kv_len = ((indexer_mask.shape[-1] + sa_config.block_kv - 1) // sa_config.block_kv) * sa_config.block_kv + pad_q = padded_q_len - indexer_mask.shape[-2] + pad_kv = padded_kv_len - indexer_mask.shape[-1] if pad_q > 0 or pad_kv > 0: pad_width = [(0, 0)] * (indexer_mask.ndim - 2) + [(0, pad_q), (0, pad_kv)] indexer_mask = jnp.pad( diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 632efbfa31..e16d0de4aa 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -3480,9 +3480,6 @@ def test_tpu_dot_product_context_parallel_with_indexer( "indexer_topk": 32, }, ) - @pytest.mark.skip( - reason="Indexer with all-gather context parallelism diverges from the dot_product reference; fix tracked in #4947." - ) @pytest.mark.tpu_only def test_tpu_flash_attention_context_parallel_with_indexer( self, context_parallel_load_balance, ici_context_parallelism=2, indexer_topk=256