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 v5p (octatrifan-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.

[gemini-code-assist]

If decoder_segment_ids is None (e.g., when document packing is disabled), generate_attention_mask returns a 4D mask of shape [batch, 1, Q, KV] instead of a 5D mask of shape [batch, 1, 1, Q, KV]. Hardcoding the indexer mask extraction as indexer_mask[:, 0, 0, :, :] will raise an IndexError in this scenario. Squeezing out any extra dimensions of size 1 dynamically is much more robust and avoids potential crashes.

Suggested change

	if indexer_mask is not None:
	# Extract single KV head and Query-per-KV head group axes [batch, 1, 1, Q, KV] -> [batch, Q, KV]
	indexer_mask = indexer_mask[:, 0, 0, :, :]
	if indexer_mask is not None:
	# Squeeze out any extra dimensions of size 1 between batch and Q/KV axes to get [batch, Q, KV]
	while indexer_mask.ndim > 3:
	indexer_mask = jnp.squeeze(indexer_mask, axis=1)

[gemini-code-assist]

The compress_ratio is dynamically inferred in HCAStaticMask as max(1, local_kv_len // max(1, compressed_kv_len)). However, if the unpadded sequence length (local_kv_len) is not a multiple of the actual compression rate (e.g., seq_len = 489 with a compression rate of 128), this inferred ratio will be incorrect (e.g., 489 // 3 = 163), leading to a mathematically incorrect attention mask. Passing the actual compress_ratio explicitly from the layer configuration ensures correctness under all sequence lengths.

Suggested change

	def __init__(
	self,
	shape: tuple[int, int],
	local_kv_len: int,
	compressed_kv_len: int,
	pad_kv_total: int = 0,
	sliding_window_size: int | None = None,
	shard_count: int = 1,
	):
	compress_ratio = max(1, local_kv_len // max(1, compressed_kv_len))
	def __init__(
	self,
	shape: tuple[int, int],
	local_kv_len: int,
	compressed_kv_len: int,
	pad_kv_total: int = 0,
	sliding_window_size: int | None = None,
	shard_count: int = 1,
	compress_ratio: int = 0,
	):
	if compress_ratio <= 0:
	compress_ratio = max(1, local_kv_len // max(1, compressed_kv_len))

[gemini-code-assist]

Pass the explicit compress_ratio parameter to HCAStaticMask to ensure the attention mask is computed correctly when sequence lengths are not multiples of the compression rate.

Suggested change

	elif self.attention_type == AttentionType.COMPRESSED and indexer_mask is None:
	local_kv_len = query.shape[2]
	compressed_kv_len = key.shape[2] - query.shape[2] - pad_kv_total
	mask = HCAStaticMask(
	shape=mask_shape,
	local_kv_len=local_kv_len,
	compressed_kv_len=compressed_kv_len,
	pad_kv_total=pad_kv_total,
	sliding_window_size=self.sliding_window_size if self.sliding_window_size else None,
	)
	elif self.attention_type == AttentionType.COMPRESSED and indexer_mask is None:
	local_kv_len = query.shape[2]
	compressed_kv_len = key.shape[2] - query.shape[2] - pad_kv_total
	mask = HCAStaticMask(
	shape=mask_shape,
	local_kv_len=local_kv_len,
	compressed_kv_len=compressed_kv_len,
	pad_kv_total=pad_kv_total,
	sliding_window_size=self.sliding_window_size if self.sliding_window_size else None,
	compress_ratio=compress_ratio,
	)

[gemini-code-assist]

Add compress_ratio as an optional parameter to tpu_flash_attention to support explicit compression ratio passing for static mask compilation.

Suggested change

	bidirectional_mask: Any = None,
	use_ragged_attention: bool = False,
	record_max_logits: bool = False,
	decoder_segment_ids_kv: Array | None = None,
	pad_kv_total: int = 0,
	) -> tuple[Array, Array]:
	bidirectional_mask: Any = None,
	use_ragged_attention: bool = False,
	record_max_logits: bool = False,
	decoder_segment_ids_kv: Array | None = None,
	pad_kv_total: int = 0,
	compress_ratio: int = 0,
	) -> tuple[Array, Array]:

[gemini-code-assist]

Add compress_ratio as an optional parameter to apply_attention to support explicit compression ratio passing.

Suggested change

	indexer_mask: Array | None = None,
	compressed_mask: Optional[Array] = None,
	record_max_logits: bool = False,
	decoder_segment_ids_kv: Optional[Array] = None,
	pad_kv_total: int = 0,
	*,
	indexer_mask: Array | None = None,
	compressed_mask: Optional[Array] = None,
	record_max_logits: bool = False,
	decoder_segment_ids_kv: Optional[Array] = None,
	pad_kv_total: int = 0,
	compress_ratio: int = 0,
	*,

[gemini-code-assist]

Pass compress_ratio to tpu_flash_attention inside apply_attention.

            bidirectional_mask=bidirectional_mask,
            use_ragged_attention=use_ragged_attention,
            record_max_logits=record_max_logits,
            decoder_segment_ids_kv=decoder_segment_ids_kv,
            pad_kv_total=pad_kv_total,
            compress_ratio=compress_ratio,
        )

[gemini-code-assist]

Add compress_ratio as an optional parameter to AttentionOp.call.

      compressed_kv: Optional[Array] = None,
      slot: Optional[int] = None,
      record_max_logits: bool = False,
      decoder_segment_ids_kv: Optional[Array] = None,
      pad_kv_total: int = 0,
      compress_ratio: int = 0,
  ):

[gemini-code-assist]

Pass compress_ratio to apply_attention inside AttentionOp.call.

Suggested change

	record_max_logits=record_max_logits,
	qk_product_einsum=self.AqtEinsum_0,
	wv_product_einsum=self.AqtEinsum_1,
	decoder_segment_ids_kv=decoder_segment_ids_kv,
	pad_kv_total=pad_kv_total,
	)
	record_max_logits=record_max_logits,
	qk_product_einsum=self.AqtEinsum_0,
	wv_product_einsum=self.AqtEinsum_1,
	decoder_segment_ids_kv=decoder_segment_ids_kv,
	pad_kv_total=pad_kv_total,
	compress_ratio=compress_ratio,
	)

[gemini-code-assist]

Pass self.compress_ratio to self.attention_op to ensure the static mask compilation uses the correct compression ratio.

Suggested change

	compressed_mask=compressed_mask,
	compressed_kv=compressed_kv,
	cached_values=current_kv_cache,
	indexer_mask=indexer_mask,
	decoder_segment_ids_kv=decoder_segment_ids_kv,
	pad_kv_total=pad_kv_total,
	)
	compressed_mask=compressed_mask,
	compressed_kv=compressed_kv,
	cached_values=current_kv_cache,
	indexer_mask=indexer_mask,
	decoder_segment_ids_kv=decoder_segment_ids_kv,
	pad_kv_total=pad_kv_total,
	compress_ratio=self.compress_ratio,
	)

[gemini-code-assist]

Using Python's structural pattern matching (match/case) here is overly verbose for a simple conditional check. Simplifying this to a standard if statement is more concise, readable, and idiomatic.

    if self.decoder_block == DecoderBlockType.DEEPSEEK4:
      if not (self.attention == "dot_product" or (self.attention == "flash" and self.use_tokamax_splash)):
        raise ValueError(
            "DeepSeek4 is only supported with dot_product attention or flash attention "
            "with use_tokamax_splash=True."
        )

[codecov]

Codecov Report

❌ Patch coverage is 45.32374% with 76 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_op.py	34.56%	43 Missing and 10 partials ⚠️
src/maxtext/layers/attention_compressed.py	60.34%	19 Missing and 4 partials ⚠️

📢 Thoughts on this report? Let us know!