Add TPU USP context parallelism

[huytransformer]

Description

This PR introduces context_parallel_strategy=usp (USP, Ulysses over ring). Follow up to #4687.

Currently does not support load balancing + sequence packing.

Tests

python3 -m pytest tests/unit/configs_value_test.py tests/unit/usp_attention_test.py tests/unit/usp_collective_test.py

Passed.

python3 -m pytest tests/unit/attention_test.py -k usp

3 passed.

Performance

llama3-8b, v5p (64 chips), CP64, bf16, synthetic data, global batch 1, no load-balancing

Reporting median step time (s)

	64K	128K	256K
all-gather (CP64)	1.31	4.03	14.33
ring (CP64)	1.68	4.63	15.35
usp (16x4)	1.31	4.02	13.92
usp (8x8)	1.44	4.05	13.65
usp (32x2)	1.65	4.52	14.72

Example repro command:

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml model_name=llama3-8b dataset_type=synthetic enable_checkpointing=False steps=20 per_device_batch_size=0.015625 packing=False attention=flash use_tokamax_splash=True use_jax_splash=False allow_split_physical_axes=True max_target_length=65536 context_parallel_strategy=usp ici_context_parallelism=16 ici_context_ulysses_parallelism=4 context_parallel_load_balance=False run_name=usp16x4_65536 base_output_directory=<output dir>

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.

[github-actions]

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

[gemini-code-assist]

Code Review

This pull request introduces support for USP (Ulysses-over-ring) context parallelism on the TPU Tokamax Splash path for training in MaxText. It adds configuration options, layout validation helpers, and integrates the hybrid strategy into the attention operation, accompanied by extensive unit and collective tests. The review feedback focuses on improving robustness by adding defensive null checks in usp_attention.py to prevent potential TypeError exceptions on unsharded tensors, and refactoring device platform checks in attention_op.py to use a more idiomatic numpy flat indexing approach.

[github-actions]

## 📋 Review Summary

This pull request introduces robust support for TPU USP (Ulysses-over-Ring) context parallelism, allowing MaxText to scale hybrid context parallelism efficiently by combining ring-attention sequence rotation with Ulysses head-exchange all-to-alls. The overall design is exceptionally high quality, extremely clean, highly idiomatic, and very well integrated with the existing MaxText config validation and Tokamax splash-attention paths.

🔍 General Feedback

• Exceptional Testing Quality: The newly added unit tests (usp_attention_test.py and usp_collective_test.py) are incredibly thorough. Specifically, forcing an 8-device CPU mesh in a standalone subprocess to verify multi-dimensional collectives, attention parity, and gradient correctness without physical TPU hardware is a masterclass in robust JAX testing.
• Robust Layout and Runtime Checks: The USP-specific configurations and constraints are systematically validated at both startup (config types validation) and layer initialization (attention operator layout checks), protecting against unsupported setups.
• Seamless Integration: The extension of logical axis rules in base.yml and physical axes configuration in maxtext_utils.py is elegant, maintaining full backward compatibility with non-USP paths.

[github-actions]

🟢 Low - Adding rtol=1e-5 to gradient comparison as well ensures numerical stability across different architectures.

Suggested change

	np.testing.assert_allclose(jax.device_get(usp_grad), jax.device_get(dense_grad), atol=1e-5, err_msg=name)
	np.testing.assert_allclose(jax.device_get(usp_grad), jax.device_get(dense_grad), rtol=1e-5, atol=1e-5, err_msg=name)
	```</COMMENT>

[github-actions]

🟡 Medium - Checking that self.num_query_heads is divisible by usp_ulysses_size is a great static sanity check. Note that if tensor model parallelism is enabled, heads are first sharded across tensor_parallelism. The actual constraint is that the local query head count (after tensor sharding) must be divisible by Ulysses size. While validate_head_sharding dynamically catches this during execution, we could consider documenting this interaction or eventually validating it here statically.

[github-actions]

🟢 Low - Adding a relative tolerance (rtol) to assert_allclose is recommended to prevent potential flakiness under different CPU architectures or compiler versions, especially when comparing standard dot-product attention with block/gathered attention.

Suggested change

	np.testing.assert_allclose(jax.device_get(usp_output), jax.device_get(dense_output), atol=1e-5)
	np.testing.assert_allclose(jax.device_get(usp_output), jax.device_get(dense_output), rtol=1e-5, atol=1e-5)
	```</COMMENT>

[codecov]

Codecov Report

❌ Patch coverage is 78.07018% with 25 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_op.py	59.64%	11 Missing and 12 partials ⚠️
src/maxtext/kernels/attention/usp_attention.py	96.29%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!