DSV4: Route compressed and sliding train attention onto the splash kernels

[systalyze-ai]

Description

Adds an opt-in TPU path that routes DeepSeek-V4 COMPRESSED and LOCAL_SLIDING train attention to the existing splash kernels. This avoids materializing dense q x kv logits.

Performance

• At seq 16384, a v6e-128 DeepSeek4-284B LoRA run required 57.5 GiB/chip without this change. The hardware had 31.25 GiB/chip available, so the model did not compile.
• At seq 512, step time fell from 0.991 to 0.637 s/step (same model on v6e-64, where the dense form still fits).

Setup: 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, dynamic-splash attention path enabled. The seq 16384 memory figure is the AOT compile requirement at that setup with this change removed.

Tests

• pytest tests/unit/dynamic_splash_mask_test.py tests/unit/attention_test.py: 137 passed.

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 the compressed_use_dynamic_splash configuration to route COMPRESSED and LOCAL_SLIDING train attention to the tokamax splash kernel on TPU. It adds mask building utilities, implements dynamic_splash_attention within AttentionOp, and includes comprehensive unit tests. The review feedback focuses on performance and memory optimizations: first, slicing decoder_segment_ids directly during comparison in build_local_sliding_splash_mask to avoid large intermediate tensor allocations; second, replacing the computationally expensive jnp.isclose with a simpler inequality comparison when processing the indexer mask.

[gemini-code-assist]

When decoder_segment_ids is not None, creating the full [batch, q_seq_len, q_seq_len] intermediate tensor segment and then slicing it to kv_seq_len can be highly inefficient, especially for large sequence lengths (e.g., 16k).

We can avoid this large intermediate allocation by slicing decoder_segment_ids directly during the comparison.

Suggested change

	if decoder_segment_ids is not None:
	segment = decoder_segment_ids[:, :, None] == decoder_segment_ids[:, None, :]
	mask = jnp.logical_and(mask, segment[..., :kv_seq_len])
	if decoder_segment_ids is not None:
	segment = decoder_segment_ids[:, :, None] == decoder_segment_ids[:, None, :kv_seq_len]
	mask = jnp.logical_and(mask, segment)

[gemini-code-assist]

Using jnp.isclose(indexer_mask, 0.0) is computationally expensive because it involves absolute difference calculations, tolerances, and handling of special float values (NaNs/Infs).

Since MaxText additive masks consistently use 0.0 for keep and DEFAULT_MASK_VALUE for discard, we can use a simple inequality comparison indexer_mask >= DEFAULT_MASK_VALUE * 0.5. This is significantly faster and matches the masking logic used elsewhere in this file.

Suggested change

	if indexer_mask.dtype != jnp.bool_:
	indexer_mask = jnp.isclose(indexer_mask, 0.0)
	if indexer_mask.dtype != jnp.bool_:
	indexer_mask = indexer_mask >= DEFAULT_MASK_VALUE * 0.5