Fix indexer mask padding in Tokamax Flash Attention with Context Parallelism

[SurbhiJainUSC]

Description

This PR fixes TPU Pre-training tests on v7x: https://github.com/AI-Hypercomputer/maxtext/actions/runs/31905605774/job/95079478073

This PR fixes a bug in attention_op.py where dynamic indexer_mask padding was incorrectly computed when using Context Parallelism (cp_size > 1) with Tokamax Flash Attention.

Root Cause

Inside attention_op.py:2016-2166 (which runs inside jax.shard_map), tensors are sharded across context parallel devices, so each device receives a local slice of queries (query.shape[2] = global_seq_len // cp_size) and indexer mask (indexer_mask.shape[-2] = global_seq_len // cp_size). However, the pad_q calculation was referencing mask_shape[0] (the outer un-sharded global sequence length). For cp_size > 1, pad_q was computed as global_seq_len - (global_seq_len // cp_size) > 0, erroneously padding indexer_mask with trailing False rows up to the full global sequence length. When passed into tokamax_splash_kernel.make_dynamic_splash_mha, the kernel generated grid dimensions and active block mappings for the full sequence length rather than the local query shard, leading to logits divergence and assertion failures.

Changes

Updated pad_q and pad_kv calculation in attention_op.py:2110-2125 to compute block-boundary padding based on local query.shape[2] and key.shape[2] dimensions using sa_config.block_q and sa_config.block_kv.

Tests

CI tests

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• 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 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 indexer_mask in wrap_flash_attention within src/maxtext/layers/attention_op.py. Specifically, it replaces the use of mask_shape with dynamically calculated padded_q_len and padded_kv_len based on the query/key shapes and block configurations. There are no review comments, and I have no additional feedback to provide.

[codecov]

Codecov Report

❌ Patch coverage is 80.95238% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_op.py	80.95%	4 Missing ⚠️

📢 Thoughts on this report? Let us know!