DSV4: Add mask hoisting, boolean splash masks, and thresholded indexer membership

[systalyze-ai]

Description

Stacked on #4908 (the splash attention PR); the diff includes its commit until it merges. Review the top commit.

Adds three opt-in DeepSeek-V4 mask optimizations for the dynamic-splash training path. Static causal, segment, and uncompressed prefix masks are built before the layer scan and reused by compressed-attention layers and rematerialized calls.

Adds boolean compressed masks from the HCA and CSA compressors through splash mask construction, and lets the CSA indexer return threshold-based membership without the one-hot expansion. The threshold path follows jax.lax.top_k ordering at boundary ties, including lower-index and signed-zero ordering.

Performance

Step time improved from 9.98 to 8.75 s/step (1.14x) with hoist_static_attention_masks, splash_bool_masks, and indexer_threshold_membership toggled on.

Reproduction setup: v6e-128, DeepSeek-V4-Flash 284B, LoRA fine-tuning, ici_expert_parallelism=16, ici_tensor_parallelism=8, ici_fsdp_parallelism=1, max_target_length=16384, per_device_batch_size=0.125 (global batch 16), LIBTPU_INIT_ARGS=--xla_tpu_scoped_vmem_limit_kib=98304, with the dynamic-splash attention path enabled and the three flags toggled off versus on.

Tests

• PYTHONPATH=$PWD/src JAX_PLATFORMS=cpu /home/jeremy/git/maxtext-16kmem/.venv/bin/python3 -m pytest -q tests/unit/deepseek_v4_mask_flags_test.py — 15 passed.
• The tie canary pins jax.lax.top_k lower-index and signed-zero ordering, while the parity tests compare hoisted, boolean, and threshold-membership masks with their default-path forms.

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 introduces several optimizations for DeepSeek-V4 compressed attention on TPU. It adds configuration flags and implementations to route COMPRESSED and LOCAL_SLIDING train attention to the tokamax dynamic splash kernel, keep compressed masks boolean end-to-end, hoist layer-invariant mask tensors to be built once per train step, and compute CSA indexer selections via a boolean threshold membership mask. Corresponding unit tests have been added to verify the correctness and parity of these new masking paths. There are no review comments, so no further feedback is provided.