From 48e5cccc9fe78cce2c5d77883e7b357195e59c85 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 Inside wrap_flash_attention under jax.shard_map, indexer_mask is sharded along the query sequence dimension (shape [batch, q_local_len, kv_full_len]). Previously, pad_q was computed using mask_shape[0] (the global un-sharded sequence length), which over-padded the local query sequence dimension of indexer_mask up to the global length. On TPU v7x and other architectures with dynamic splash attention schedulers, this caused a grid iteration and block dimension mismatch, leading to assertion failures in test_tpu_flash_attention_context_parallel_with_indexer (both no_lb_cp2 and lb_cp4_smallk). This fix computes padded_q_len and padded_kv_len based directly on the local indexer_mask dimensions and the hardware block sizes (sa_config.block_q / sa_config.block_kv), ensuring correct shape compatibility across all hardware platforms (v4, v5e, v5p, v6e, v7x, CPU, GPU) and arbitrary CP degrees. --- src/maxtext/layers/attention_op.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 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(