Add TPU Ulysses context parallelism

[huytransformer]

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

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]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[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.

[github-actions]

## 📋 Review Summary

This Pull Request introduces highly performant TPU Ulysses context parallelism on the TPU Tokamax Splash attention training path. The implementation is exceptionally clean, well-integrated, and backed by comprehensive unit, collective, and integration tests that verify correctness against dense references and compiled HLO collectives.

🔍 General Feedback

• Excellent Test Coverage: The newly added tests (including CPU multi-device mock checks and HLO collective pattern assertions) set a very high bar for correctness and robustness.
• Robust Config Validation: The comprehensive validation in configs/types.py ensures that unsupported combinations (like MQA, packing, ragged attention, dropout, etc.) are proactively rejected with precise errors.
• Clean Shared Utilities: Factoring out mesh and sharding operations into context_parallel_utils.py is a great design choice that improves maintainability.
• Thorough Documentation: The update to sharding.md is incredibly detailed and accurately captures Ulysses constraints and architectural mechanics.

[github-actions]

🟢 Consistently specify the unsupported value `context_parallel_load_balance=True` in the error message to match the config validation error message in `types.py`.

Suggested change

	raise ValueError("TPU Ulysses attention does not support context_parallel_load_balance.")
	if self.config.context_parallel_load_balance:
	raise ValueError("TPU Ulysses attention does not support context_parallel_load_balance=True.")

[codecov]

Codecov Report

❌ Patch coverage is 79.14439% with 39 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_op.py	63.76%	12 Missing and 13 partials ⚠️
...axtext/kernels/attention/context_parallel_utils.py	66.66%	5 Missing and 3 partials ⚠️
src/maxtext/kernels/attention/ulysses_attention.py	92.68%	3 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.