fix: compute local block padding for dynamic splash indexer mask in CP

[RexBearIU]

Description

Fixes an AssertionError in MLATest.test_tpu_flash_attention_context_parallel_with_indexer (both no_lb_cp2 and lb_cp4_smallk) by properly computing local block-aligned padding for dynamic splash attention masks under context parallelism.

Root Cause

Inside wrap_flash_attention under jax.shard_map, indexer_mask is sharded across context devices along the query sequence dimension Q_LENGTH into shape (batch, q_local_len, kv_full_len) (e.g. [1, 128, 512] under CP=4 or [1, 256, 512] under CP=2).

Previously, pad_q was computed using mask_shape[0] (the outer, global un-sharded sequence length 512):

pad_q = mask_shape[0] - indexer_mask.shape[-2]  # 512 - 128 = 384

This over-padded the local query dimension of indexer_mask with hundreds of False rows back up to 512, creating a $512 \times 512$ mask for a 128 (or 256) token local shard.

Machine-Level Compatibility

• TPU v7x (Ironwood): v7x's dynamic splash attention compiler and scheduler use the mask's row dimension to configure the dynamic tile grid (dq_reduction_steps=3 with $256 \times 256$ MXU tiling). When given a 512-row mask for a 128/256-row local query shard, the grid dispatched out-of-bounds/mismatched iterations, causing output corruption and assertion failures.
• TPU v6e (Trillium), v5p, v5e, v4: On earlier TPU generations, Pallas loop bounds primarily tracked q.shape[1], causing trailing mask rows to be benignly ignored in some configurations but prone to subtle stride/allocation issues.
• Fix Compatibility: Calculating padded_q_len and padded_kv_len using sa_config.block_q and sa_config.block_kv multiples of the local indexer_mask shape (indexer_mask.shape[-2] and indexer_mask.shape[-1]):
  1. Ensures the mask dimensions strictly match the local shard's query sequence length regardless of CP degree ($CP=1, 2, 4, 8, \dots$).
  2. Dynamically respects the hardware block sizes ($128$ for v4/v5e/v5p/v6e, $256$ for v7x, or custom block sizes).
  3. Properly supports both load-balanced and non-load-balanced all-gather context parallel strategies across all hardware backends.

Tests

• Unit tests reproduced and verified locally with multi-device 8-CPU host simulation (--xla_force_host_platform_device_count=8):
  • MLATest.test_tpu_flash_attention_context_parallel_with_indexer_no_lb_cp2 → PASS
  • MLATest.test_tpu_flash_attention_context_parallel_with_indexer_lb_cp4_smallk → PASS
  • MLATest.test_tpu_dot_product_context_parallel_with_indexer → PASS
  • MLATest.test_tpu_flash_attention_ring_context_parallel_with_indexer → PASS
• Pre-commit hooks (codespell, pylint, pyink) passed cleanly.

Checklist

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request updates the padding calculation for the Tokamax splash kernel in attention_op.py by aligning the query and key-value lengths to their respective block sizes (block_q and block_kv) before computing the padding widths. I have no feedback to provide as there are no review comments.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!