[DeepSeek-V4] Implement Compressed Attention Layers

[parambole]

Description

Implement compressed attention mechanisms and indexer modules required for DeepSeek-V4 integration into MaxText:

• CSACompressor & HCACompressor: Long-range attention compressors supporting causal block bias and YaRN frequency scaling decoupling.
• LightningIndexer: Memory-efficient indexer module implementing sentinel masking and dynamic RoPE scaling.
• Configuration: Register attention compression hyperparameters (compress_ratios, index_head_dim, sliding_window) in types.py and base.yml.
• Unit test suite (tests/unit/deepseek_v4_vs_reference_test.py) validating attention compression parity against PyTorch reference implementations at atol=1e-5, rtol=1e-5.

Tests

Tested on CPU

pytest  tests/unit/deepseek_v4_vs_reference_test.py

======================= 10 passed, 10 warnings in 20.42s =======================
tests/unit/deepseek_v4_vs_reference_test.py ..........                   [100%]

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.

[codecov]

Codecov Report

❌ Patch coverage is 7.63052% with 230 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_compressed.py	7.63%	230 Missing ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @parambole, 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]

🤖 Hi @parambole, 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 implements the core compressed attention layers for DeepSeek-V4 integration, including HCA, CSA, and the Lightning Indexer. The implementation is technically sound, follows established patterns in MaxText, and includes comprehensive parity tests against PyTorch.

🔍 General Feedback

• Efficiency: The main coordinator block uses jnp.repeat for broadcasting MQA keys/values, which is memory-intensive. Switching to jnp.einsum broadcasting is recommended.
• Typo: A minor typo swaped was found in the indexer module.
• Config: Ensure compress_ratios is properly documented as a required list when using these attention variants to avoid runtime IndexError.

[github-actions]

🔴 Memory/Performance: `jnp.repeat` creates large intermediate tensors that are unnecessary given the shared MQA structure. You can use broadcasting in `jnp.einsum` to achieve the same result more efficiently.

Suggested change

	logits = jnp.einsum("bhsd, bhkd -> bhsk", q, k, precision=self.config.matmul_precision) * self.scaling
	# Broadcast key/value configurations to all heads using broadcasting in einsum
	# k and v remain [B, 1, S_kv, D_head]
	k = kv
	v = kv
	# Compute attention logits with head broadcasting: [B, H, S, S_kv]
	logits = jnp.einsum("bhsd, b1kd -> bhsk", q, k, precision=self.config.matmul_precision) * self.scaling

[github-actions]

🔴 Memory/Performance: Use broadcasting in `jnp.einsum` here as well to avoid the `jnp.repeat` from earlier.

Suggested change

	attn_output = jnp.einsum("bhsk, bhkd -> bhsd", attn_weights, v, precision=self.config.matmul_precision)
	# Project attention weights onto values with head broadcasting: [B, H, S, D_head]
	attn_output = jnp.einsum("bhsk, b1kd -> bhsd", attn_weights, v, precision=self.config.matmul_precision)

[github-actions]

🟢 Typo: "swaped" should be "swapped".

Suggested change

	scores = jnp.matmul(q, swaped_kv)
	# swapped_kv: [B, 1, D_idx, W]
	swapped_kv = jnp.swapaxes(compressed_kv, -1, -2)
	swapped_kv = jnp.expand_dims(swapped_kv, axis=1)
	# scores: [B, S, H, W]
	scores = jnp.matmul(q, swapped_kv)

[github-actions]

🟡 Configuration: The default `compress_ratios` is an empty list. Since `DeepSeekV4Attention` relies on this list having at least `layer_idx + 1` elements when using compressed layer types, this will cause an `IndexError` at runtime unless the user provides a full list. It might be better to provide a default or add a check with a clear error message.

[RissyRan]

I didn't see sharding annotations. It will be good we start to add some of them in this PR? i.e. starting with those weights.

[RissyRan]

Have you considered to re-use v3.2 Indexer?

Ref: doc

[RissyRan]

For file name, may be compressed_attention.py as you mentioned in the comment here?