Add HCA Static Compilation for Splash Attention [Deepseek v4]

[octatrifan]

Description

Implements a static Tokamax Splash Attention compilation path for DeepSeek-V4 Heavily Compressed Attention (HCA, compress_ratio > 4). This replaces the dynamic indexer mask dispatch with a statically compiled HCAStaticMask.

Context & Motivation

Previously, both CSA and HCA were routed through make_dynamic_splash_mha. While CSA requires dynamic boolean masks generated at runtime by the top-k sparse indexer, HCA's compressed attention pattern is completely deterministic and static (local causal window + evenly spaced compressed KV tokens). Running HCA through the dynamic mask path added unnecessary mask materialization and Pallas dynamic grid dispatch overhead.

Key Changes

1. HCAStaticMask: Implemented a computable mask for static Tokamax splash attention that handles the local sliding window, causal boundaries, and compressed token indices without generating runtime boolean masks.
2. Unaligned Sequence Length Handling: Unified ceiling-division padding in AttentionOp.tpu_flash_attention for both CSA and HCA. When sequence lengths are not multiples of the block size (e.g., length 3968 with block size 512), query tensors and segment IDs are padded to block boundaries before kernel dispatch, and outputs are sliced back to the original sequence length.
3. Dispatch & CP Compatibility: Maintained standard jax.vmap batch dispatch for the static splash kernel, ensuring full compatibility with context parallelism (CP) and per-sequence document packing.

Tests

Validated locally and on Cloud TPU v5p8:

# DeepSeek-V4 unit and parity tests against HuggingFace reference
python3 -m pytest tests/unit/deepseek_v4_vs_reference_test.py -v

# CompressedAttention compilation, numerical equivalence, and packing tests
python3 -m pytest tests/unit/attention_test.py::CompressedAttentionTest -v

# Context Parallelism sanity check
python3 -m pytest tests/unit/attention_test.py::AttentionTest::test_tpu_flash_attention_context_parallel_cp_ep_no_load_balance -v

All test suites passed on TPU-v5p (including unaligned lengths like 489 and 3968, and packing equivalence checks).

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 support for DeepSeek-V4's Compressed Sparse Attention (CSA) and Heavily Compressed Attention (HCA) mechanisms, including overlapping window pooling, document-packing-aware masking, and sequence padding for Splash kernel alignment. It also adds extensive unit and parity tests. The reviewer feedback identifies a potential IndexError when decoder_segment_ids is None due to hardcoded indexing of the indexer mask, which can be resolved by dynamically squeezing size-1 dimensions. Additionally, the reviewer points out that compress_ratio is incorrectly inferred in HCAStaticMask for unaligned sequence lengths, suggesting explicitly passing compress_ratio through the attention pipeline. Finally, simplifying the pattern matching in configuration validation to a standard if statement is recommended for improved readability.

[codecov]

Codecov Report

❌ Patch coverage is 62.26415% with 20 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_op.py	61.53%	13 Missing and 7 partials ⚠️

📢 Thoughts on this report? Let us know!