Add TPU USP context parallelism by huytransformer · Pull Request #4836 · AI-Hypercomputer/maxtext

huytransformer · 2026-08-11T17:16:22Z

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 · 2026-08-11T17:17:07Z

🤖 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 · 2026-08-11T17:24:11Z

+  dense_grads = jax.grad(dense_loss, argnums=(0, 1, 2))(query, key, value)
+  usp_grads = jax.grad(usp_loss, argnums=(0, 1, 2))(query, key, value)
+  for name, dense_grad, usp_grad in zip(("dQ", "dK", "dV"), dense_grads, usp_grads):
+    np.testing.assert_allclose(jax.device_get(usp_grad), jax.device_get(dense_grad), atol=1e-5, err_msg=name)


🟢 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 · 2026-08-11T17:24:11Z

+        raise ValueError(
+            "TPU USP attention requires max_target_length to be divisible by ici_context_parallelism squared."
+        )
+      if self.num_query_heads % usp_ulysses_size != 0:


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

wow go gemini

you could add in the TP factor as well but generally no one is using TP anyway. This is just a config check, if a user is setting TP and going over head count they will hit a less user friendly jax sharding error

fwiw I think the TP case doesn't reach jax sharding error. validate_head_sharding already checks the local head at init with test

maxtext/src/maxtext/kernels/attention/ulysses_attention.py

Lines 186 to 196 in bdafbc9

local_query_heads = num_query_heads // q_head_shards

local_kv_heads = num_kv_heads // kv_head_shards

if local_query_heads % local_kv_heads != 0:

raise ValueError(

f"{attention_label} requires local query heads "

f"({local_query_heads}) to be divisible by local KV heads ({local_kv_heads})."

)

if local_query_heads % ulysses_size != 0:

raise ValueError(

f"{attention_label} requires local query heads "

f"({local_query_heads}) to be divisible by the Ulysses exchange size ({ulysses_size})."

github-actions · 2026-08-11T17:24:12Z

+
+  dense_output = _dense_reference_attention(query, key, value, segment_ids)
+  usp_output = usp_attention_fn(query, key, value, segment_ids)
+  np.testing.assert_allclose(jax.device_get(usp_output), jax.device_get(dense_output), atol=1e-5)


🟢 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 · 2026-08-11T17:26:51Z

Codecov Report

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

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_op.py	58.00%	10 Missing and 11 partials ⚠️
src/maxtext/kernels/attention/usp_attention.py	93.93%	2 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

gobbleturk · 2026-08-12T16:12:46Z

            "TPU Ulysses attention requires num_kv_heads "
            f"({self.num_kv_heads}) to be divisible by context_parallel_size ({context_parallel_size})."
        )
+    if context_parallel_strategy != "usp" and (


Isn't context_parallel_strategy=ulysses stil supported? E.g. only ulysses, no ring?

If a user wanted to run only ulysses without any ring, what would they set? It is not clear to me

so for every CP we set ici_context_parallelism, only usp additionally sets the ulysses else rejected at init.

gobbleturk · 2026-08-12T16:13:58Z

+      raise ValueError(
+          "ici/dcn_context_ulysses_parallelism was specified, but is only supported when "
+          "context_parallel_strategy='usp'."
+      )


there is a lot of logic here - can you wrap all of this into a function (style suggestion go/small-functions)

gobbleturk · 2026-08-12T16:19:47Z

+    raise NotImplementedError("TPU USP attention does not support record_max_logits yet.")
+
+
+def with_sequence_axes(axis_names: Any, ring_axis: str, ulysses_axis: str, sequence_dim: int) -> Any:


renmae to shard_by_ring_and_ulysses or shard_by_usp or similar?

renamed to with_usp_sequence_axes!

gobbleturk · 2026-08-12T16:21:06Z

        attention_output = ulysses_attention.inverse_ulysses_all_to_all(attention_output, context_axis)
        return attention_output, None

+      if use_usp:


I would also wrap everything under the if in a function following style go/small-functions

NuojCheng · 2026-08-12T18:22:14Z

+    if usp_ring_size <= 1:
+      raise ValueError("TPU USP attention requires ici_context_parallelism > 1 for the ring dimension.")
+    if usp_ulysses_size <= 1:
+      raise ValueError("TPU USP attention requires ici_context_ulysses_parallelism > 1 for the Ulysses dimension.")


what if ici_context_parallelism=-1 or ici_context_ulyses_parallelism=-1

NuojCheng · 2026-08-12T18:29:55Z

    if self.context_sharding not in ("context", "expert"):
      raise ValueError(f"Assigned context_sharding f{self.context_sharding} is not supported.")
+    if self.ulysses_context_sharding != "context_ulysses":
+      raise ValueError(f"Assigned ulysses_context_sharding {self.ulysses_context_sharding} is not supported.")


we will remove this limit in the future using component sharding!!

NuojCheng · 2026-08-12T18:35:34Z

    axis_names_kv: Any,
    dkv_dim_q: int,
    dkv_dim_kv: int,
+    attention_label: str,


help me understand when we wanna use ulysses context, should we use

ici_context_parallelism=CP context_parallelism_strategy=ulysses

OR

ici_ulysses_context_parallelism=CP

NuojCheng · 2026-08-12T18:36:53Z

 shard_mode: "auto" # can be either auto or explicit
 custom_mesh_and_rule: "" # replace default mesh and logical rule by specifying yml name under config/mesh_and_rule/.
-mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_sequence', 'expert', 'autoregressive']
+mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_ulysses', 'context_autoregressive', 'tensor', 'tensor_sequence', 'expert', 'autoregressive']


please update mesh_axes in types.py as well

huytransformer added the gemini-review label Aug 11, 2026

gemini-code-assist Bot reviewed Aug 11, 2026

View reviewed changes

Comment thread src/maxtext/kernels/attention/usp_attention.py

Comment thread src/maxtext/kernels/attention/usp_attention.py

Comment thread src/maxtext/layers/attention_op.py

github-actions Bot reviewed Aug 11, 2026

View reviewed changes

huytransformer mentioned this pull request Aug 11, 2026

Support packing for TPU Ulysses attention #4825

Merged

4 tasks

huytransformer marked this pull request as ready for review August 12, 2026 05:25

huytransformer force-pushed the htn-usp-cp branch from f9a0d97 to 293acad Compare August 12, 2026 06:25

huytransformer force-pushed the htn-usp-cp branch 3 times, most recently from 27a6b5f to b9eea0d Compare August 12, 2026 06:56

gobbleturk reviewed Aug 12, 2026

View reviewed changes

Add TPU USP context parallelism

bdafbc9

huytransformer force-pushed the htn-usp-cp branch from b9eea0d to bdafbc9 Compare August 12, 2026 16:57

huytransformer requested a review from gobbleturk August 12, 2026 17:09

NuojCheng reviewed Aug 12, 2026

View reviewed changes

NuojCheng approved these changes Aug 12, 2026

View reviewed changes

	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>

	local_query_heads = num_query_heads // q_head_shards
	local_kv_heads = num_kv_heads // kv_head_shards
	if local_query_heads % local_kv_heads != 0:
	raise ValueError(
	f"{attention_label} requires local query heads "
	f"({local_query_heads}) to be divisible by local KV heads ({local_kv_heads})."
	)
	if local_query_heads % ulysses_size != 0:
	raise ValueError(
	f"{attention_label} requires local query heads "
	f"({local_query_heads}) to be divisible by the Ulysses exchange size ({ulysses_size})."

		raise NotImplementedError("TPU USP attention does not support record_max_logits yet.")


		def with_sequence_axes(axis_names: Any, ring_axis: str, ulysses_axis: str, sequence_dim: int) -> Any:

Uh oh!

Conversation

huytransformer commented Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Performance

Checklist

Uh oh!

github-actions Bot commented Aug 11, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🔍 General Feedback

Uh oh!

github-actions Bot Aug 11, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 11, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 11, 2026

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

huytransformer commented Aug 11, 2026 •

edited

Loading

codecov Bot commented Aug 11, 2026 •

edited

Loading