Add HCA Static Compilation for Splash Attention [Deepseek v4]

[octatrifan]

Description

Add static Tokamax Splash Attention compilation for DeepSeek-V4 Heavily Compressed Attention (HCA).

Previously, both Compressed Sparse Attention (CSA, compress_ratio == 4) and HCA were dispatched through dynamic Splash Attention (make_dynamic_splash_mqa), requiring DeepseekV4HCACompressor to materialize dense (B, 1, S, S_comp) float mask arrays in HBM and build runtime indexer masks.

Because HCA attention patterns are strictly deterministic (local sliding window + completed preceding compression windows) and do not rely on the indexer (unlike CSA, which requires a top-k step), we can construct the block mask analytically on CPU during AOT tracing. This PR introduces HCAStaticMask, bypassing HBM mask allocation and dispatching directly to static Splash MQA kernels (make_splash_mqa).

Key changes:

• HCAStaticMask: Subclasses splash_attention_mask.Mask to evaluate coordinate-based block sparsity on CPU during tracing without runtime HBM mask materialization.
• Compressor bypass: DeepseekV4HCACompressor returns None for masks when attention_kernel="flash". generate_attention_mask is now restricted strictly to CSA.
• MQA kernel dispatch: Routes compressed attention with num_kv_heads == 1 to make_splash_mqa, stripping singleton KV head dimensions for 2D Tokamax input requirements.
• Backward ring buffer: Sets dq_reduction_steps = 3 for AttentionType.COMPRESSED in create_sa_config to enable in-SRAM circular ring buffer accumulation for dQ. (Without this change, the backward pass is much slower than the dynamic path)
• Arbitrary sequence lengths & packing: Added ceiling division tile padding in AttentionOp.tpu_flash_attention for unaligned sequence lengths (e.g. $S=489$) and document packing boundary isolation with segment IDs.
• Validation guards: Added explicit ValueError checks blocking Context Parallelism (cp_size > 1) and Ulysses/USP when AttentionType.COMPRESSED is active, pending future asymmetric communication support.

Details in b/537346777

Performance

We conduct microbenchmarks (just on the attention layer) on a v5p TPU for DeepSeek-V4 Flash HCA.

• Standard Context (S <= 8k): Static Flash and Dynamic Flash are on par across all batch sizes, delivering identical step latency while Static Flash avoids allocating dense boolean masks in device memory.
• Long Context Speedups (S >= 16k): Static Flash becomes progressively faster than Dynamic Flash as sequence length increases—delivering a 1.09x speedup at 16k, 1.20x at 32k (with a 2.25x speedup on the forward pass), and 1.48x at 65k for batch size 1.
• Memory Footprint: By compiling mask coordinates analytically at trace-time rather than materializing dynamic boolean mask tensors on device, Static Flash eliminates the O(S^2) mask memory overhead, reducing peak HBM by up to 48% at 65k. This allows us to scale up to 128k sequence length, where Dynamic Flash fails due to SMEM OOM.
• Compilation Time Reduction: When comparing against graph-captured masks, Static Flash drops JIT compilation time from 1.9 minutes to 1.7 seconds (68x reduction) at 32k and from 7.6 minutes to 1.7 seconds (264x reduction) at 65k. Under runtime dynamic masking, it eliminates host compiler memory spikes and SMEM lowering failures, compiling in under 5 seconds across all sequence lengths up to 128k.

Tests

Tested on Cloud TPU v5p:

# DeepSeek-V4 unit and reference parity tests
pytest tests/unit/deepseek_v4_vs_reference_test.py -v
# CompressedAttention compilation, numerical equivalence, and packing tests
pytest tests/unit/attention_test.py::CompressedAttentionTest -v
# Static mask bitwise parity, sliding window, and unaligned length tests
pytest tests/unit/attention_test.py::HCAStaticMaskTest -v

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 72.79412% with 37 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_op.py	76.03%	15 Missing and 14 partials ⚠️
src/maxtext/layers/attention_compressed.py	46.66%	4 Missing and 4 partials ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @octatrifan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This PR introduces compile-time static mask generation (HCAStaticMask) for DeepSeek-V4 Heavily Compressed Attention (HCA), bypassing expensive runtime HBM mask allocation and dispatching directly to static Splash MQA kernels. The overall implementation is exceptionally clean, robust, and performs highly precise index/tile calculations to enable efficient static compilation and document packing.

🔍 General Feedback

• Math and Alignment Precision: The coordinate-based block sparsity computation in HCAStaticMask perfectly matches the mathematical window boundaries of the HCA compressor.
• Performance & Scalability: Setting dq_reduction_steps=3 for Compressed Attention is a major highlight, drastically reducing unreduced gradient write traffic to HBM during backward pass.
• Comprehensive Testing: The unit testing suite is highly comprehensive, covering unaligned sequence lengths, document packing boundaries, static mask parity, and error routing for unsupported Context Parallelism configurations.

[parambole]

Nice work on this — the HCAStaticMask approach and the upfront block padding.

A few questions from reading through it. Caveat up front: this is a static read, I
haven't run anything on TPU, so the first two may well be non-issues if the configs
that would trigger them can't occur in practice.

#	Kind	Location	Question
1	Please check	attention_op.py — HCAStaticMask.__getitem__	Can pad_kv_total be smaller than pad_q?
2	Please check	attention_op.py — COMPRESSED mask branch	Is use_tokamax_splash=False reachable here?
3	Question	attention_compressed.py — flash early return	Does the non-TPU AR fallback matter for this model?
4	Suggestion	attention_test.py — HCAStaticMaskTest	Test cases don't seem to reach the pad_m branch

[shuningjin]

Thanks for the careful design to eliminate dense mask materialization in device memory! Some comments on fix and cleanup.