Add TPU Ulysses context parallelism by huytransformer · Pull Request #4687 · AI-Hypercomputer/maxtext

huytransformer · 2026-07-31T07:06:56Z

Description

This PR adds context_parallel_strategy="ulysses", DeepSpeed-Ulysses style context parallelism for the TPU splash attention (https://arxiv.org/abs/2309.14509).

Motivations: (1) the existing approaches degrade when the context parallelism degree is large relative to the sequence length. (2) It is also the head-parallel half of a Ulysses x ring hybrid (USP), which can be built on top of this PR.

Performance

What we did: v5p-128 (64 chips), bf16, synthetic data, 2M tokens per step in every cell, remat_policy=custom context=device, dq_reduction_steps=3. Each uses its best known block sizes (2048 for Ulysses, 1024 for Ring and All-gather).

llama2-7b (MHA, 32 KV heads), 32K context, cp=16, 2K tokens per device.

Method	Step time	TFLOP/s/device	MFU
Ulysses	8.90 s	239.4	52.2% Xprof
All-gather	8.83 s	242.7	52.9%
Ring	16.38 s	130.8	28.5% Xprof

Ulysses matches all-gather and runs 1.8x faster than ring, while holding 1/16 of all-gather's attention KV per device (2 of 32 heads).

llama3.1-8b (GQA, 8 KV heads), cp=4.

Context (chunk)	Ulysses	All-gather	Ring
64K (16K)	243.3 / 53.0% Xprof	282.5 / 61.6% Xprof	277.4 / 60.4% Xprof
8K (2K)	259.0 / 56.4% Xprof	281.1 / 61.2% Xprof	274.7 / 59.8% Xprof

Correctness: 20 training steps against all_gather on llama3.1-8b at 64K, cp=4, identical data, max loss difference is 0.001.

Example repro

export LIBTPU_INIT_ARGS="--xla_tpu_scoped_vmem_limit_kib=65472 \
  --xla_tpu_use_enhanced_launch_barrier=true \
  --xla_tpu_enable_async_collective_fusion=true \
  --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true \
  --xla_tpu_enable_async_collective_fusion_multiple_steps=true \
  --xla_tpu_overlap_compute_collective_tc=true \
  --xla_enable_async_all_gather=true \
  --xla_enable_async_collective_permute=true \
  --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \
  --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \
  --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true"

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
  base_output_directory=${OUTPUT_DIR} run_name=ulysses_cp16 \
  dataset_type=synthetic enable_checkpointing=false steps=10 \
  model_name=llama2-7b max_target_length=32768 per_device_batch_size=1 \
  ici_context_parallelism=16 ici_fsdp_parallelism=4 \
  attention=flash use_tokamax_splash=true use_jax_splash=false \
  context_parallel_strategy=ulysses context_parallel_load_balance=false \
  packing=false dq_reduction_steps=3 remat_policy=custom context=device \
  num_vocab_tiling=16 \
  sa_block_q=2048 sa_block_kv=2048 sa_block_q_dkv=2048 sa_block_kv_dkv=2048

Tests

tests/unit/ulysses_attention_test.py: Passed.
tests/unit/ulysses_collective_test.py: Passed.
tests/unit/configs_value_test.py: Passed.
tests/integration/train_tests.py: train smoke test with ulysses cp=4. Passed.

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 · 2026-07-31T07:06:59Z

Caution

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

github-actions · 2026-07-31T07:07:27Z

🤖 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 · 2026-07-31T07:11:42Z

+      if self.use_ragged_attention:
+        raise ValueError("TPU Ulysses attention does not support ragged attention.")
+      if self.config.context_parallel_load_balance:
+        raise ValueError("TPU Ulysses attention does not support context_parallel_load_balance.")


🟢 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 · 2026-07-31T07:13:36Z

Codecov Report

❌ Patch coverage is 79.00000% with 42 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_op.py	65.85%	23 Missing and 5 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 · 2026-08-01T00:01:46Z

Caution

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

NuojCheng · 2026-08-06T18:15:11Z

+from maxtext.kernels.attention import ulysses_attention
+
+
+@pytest.mark.cpu_only


I think the default now is cpu_only

NuojCheng · 2026-08-06T18:53:23Z

+from typing import Any
+
+import jax
+


nit: I think all these three function are useful in general, we should consider moving them to sharding.py in the future

Yeah, agree. It seems not context parallelism specific.

NuojCheng · 2026-08-06T18:59:13Z

+
+def inverse_ulysses_all_to_all(tensor: Any, ulysses_axis: str):
+  """Moves `[B, H/U, S, D]` back to `[B, H, S/U, D]`."""
+  return jax.lax.all_to_all(tensor, ulysses_axis, split_axis=2, concat_axis=1, tiled=True)


I love you using explicit all2all/shard_map instead of relying on XLA/auto_shard

NuojCheng · 2026-08-06T19:02:26Z

+
+      splash_kernel = wrap_ulysses_splash_kernel(mask)
+      segment_axis_names_splash_kernel = jax.sharding.PartitionSpec(None)
+      splash_kernel = self._maybe_shard_with_pspec(splash_kernel, segment_axis_names_splash_kernel)


why we want splash kernel been fully replicated? or is it an explicit all gather?

NuojCheng · 2026-08-06T19:04:35Z

+            query, key, value, decoder_segment_ids_tuple, sinks
+        )
+        attention_output = ulysses_attention.inverse_ulysses_all_to_all(attention_output, context_axis)
+        return attention_output, None


very clean logic, nice!

NuojCheng · 2026-08-06T19:06:08Z

+        f"close. ici_context_parallelism={ici_context_parallelism}.",
+    )
+
+  @pytest.mark.tpu_only


if we just wanna check HLO, no need for TPU test. Use CPU test instead

NuojCheng · 2026-08-06T20:20:48Z

+      ulysses_attention.validate_ulysses_runtime(model_mode=MODEL_MODE_TRAIN, record_max_logits=True)
+
+  def test_with_sequence_axis_preserves_partition_spec_type(self):
+    spec = jax.sharding.PartitionSpec("data", None, None, "model")


we don't have 'model' axis name in maxtext?

NuojCheng · 2026-08-06T20:51:03Z

@@ -0,0 +1,93 @@
+# Copyright 2026 Google LLC


this is a very interesting test

NuojCheng · 2026-08-06T20:51:16Z

@@ -0,0 +1,265 @@
+# Copyright 2026 Google LLC


move this test to attention_test?

NuojCheng · 2026-08-06T20:51:27Z

@@ -0,0 +1,157 @@
+# Copyright 2026 Google LLC


move this test to attention test?

NuojCheng · 2026-08-06T20:51:54Z

+      mesh=mesh,
+      in_specs=P(None, None, "context", None),
+      out_specs=P(None, "context", None, None),
+      check_vma=False,


what gonna happen if we use check_vma=True?

RissyRan

Thanks! LGTM at high level. Some minor comments.

RissyRan · 2026-08-07T17:22:04Z

          "context_parallel_strategy='ring'."
      )
+    if context_parallel_strategy == "ulysses":
+      if self.hardware != "tpu":


Do you know why? If this feature hasn't been tested on GPUs yet, we should label it as 'experimental' and add a disclaimer that GPU execution is unverified and should be used with caution. Similar comments for bellow.

RissyRan · 2026-08-07T17:23:20Z

+        raise ValueError("TPU Ulysses attention requires context_sharding='context'.")
+      ici_context_parallel_size = self.ici_context_parallelism
+      dcn_context_parallel_size = self.dcn_context_parallelism
+      if ici_context_parallel_size <= 0 or dcn_context_parallel_size <= 0:


MaxText won't support ici_context_parallel_size <= 0 or dcn_context_parallel_size <= 0 for other strategies right? This is invalid, not specific to this Ulysses. If so, we could just align with other sharding constratins.

RissyRan · 2026-08-07T17:27:17Z

+        raise ValueError("TPU Ulysses attention does not support sparse indexer masks.")
+      if self.use_chunked_prefill:
+        raise ValueError("TPU Ulysses attention does not support chunked prefill yet.")
+      if self.use_multimodal:


Wondering if those features are not tested or needs feature to supported?

if self.use_multimodal: raise ValueError("TPU Ulysses attention does not support multimodal attention.") if self.enable_dropout and self.dropout_rate > 0.0: raise ValueError("TPU Ulysses attention does not support dropout yet.") ...

RissyRan · 2026-08-07T17:27:52Z

+from typing import Any
+
+import jax
+


Yeah, agree. It seems not context parallelism specific.

RissyRan · 2026-08-07T17:31:33Z

            dkv_dim_q=3,
            dkv_dim_kv=3,
        )
+    if self.attention_kernel == "flash" and ulysses_attention.is_context_parallel_ulysses_requested(self.config):


It seems in this PR, we have quite some assertions, inside of types.py, those functions, etc. Do you think we could just keep one version in types.py for early check if no overwritten is allowed/happens during runtime.

RissyRan · 2026-08-07T17:35:37Z

+      jax.shard_map,
+      mesh=mesh,
+      in_specs=(
+          P(None, None, "context", None),


It seems this test verified context sharding only. We usually will include fsdp + cp with long context. Could you help add a test and ensure those work (to avoid future change breakage)?

huytransformer added the gemini-review label Jul 31, 2026

github-actions Bot reviewed Jul 31, 2026

View reviewed changes

huytransformer marked this pull request as ready for review August 1, 2026 00:01

huytransformer closed this Aug 3, 2026

huytransformer reopened this Aug 3, 2026

huytransformer force-pushed the htn-ulysses-cp branch 2 times, most recently from 7d1fd35 to e72bd0d Compare August 3, 2026 20:39

NuojCheng reviewed Aug 6, 2026

View reviewed changes

NuojCheng approved these changes Aug 6, 2026

View reviewed changes

RissyRan reviewed Aug 7, 2026

View reviewed changes

Add TPU Ulysses context parallelism

3cd54a9

huytransformer force-pushed the htn-ulysses-cp branch from e72bd0d to 3cd54a9 Compare August 8, 2026 03:48

huytransformer closed this Aug 8, 2026

huytransformer reopened this Aug 8, 2026

	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.")

		from maxtext.kernels.attention import ulysses_attention


		@pytest.mark.cpu_only

Uh oh!

Conversation

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

Uh oh!

Description

Performance

llama2-7b (MHA, 32 KV heads), 32K context, cp=16, 2K tokens per device.

llama3.1-8b (GQA, 8 KV heads), cp=4.

Example repro

Tests

Checklist

Uh oh!

gemini-code-assist Bot commented Jul 31, 2026

Uh oh!

github-actions Bot commented Jul 31, 2026

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🔍 General Feedback

Uh oh!

github-actions Bot Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

gemini-code-assist Bot commented Aug 1, 2026

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!

RissyRan left a comment

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 Jul 31, 2026 •

edited

Loading

codecov Bot commented Jul 31, 2026 •

edited

Loading