Add TPU Ulysses context parallelism

docs/guides/optimization/sharding.md

@@ -260,6 +260,8 @@ Note in general there are many flavors of CP such as ring attention, which in th
 
 MaxText supports `context_parallel_strategy=all_gather`, and supports `context_parallel_strategy=ring` through GPU Transformer Engine and TPU Tokamax Splash paths; ring performs the computation and communication in chunks and ideally overlaps them in a collective matmul fashion. This strategy requires extending the online softmax trick from only within chip to additionally apply it across chips.
 
+MaxText also supports `context_parallel_strategy=ulysses` ([DeepSpeed Ulysses](https://arxiv.org/abs/2309.14509)) on the TPU Tokamax Splash path for training. Ulysses exchanges sequence ownership for head ownership by communicating the Q, K, V, and output activations through all-to-all collectives: each device computes ordinary full-sequence attention for its head subset, and the inverse all-to-all restores the sequence sharding on the output. It requires explicit positive context parallelism values, `context_sharding=context`, `attention=flash` with Tokamax Splash, global causal attention, query and KV head counts divisible by the context parallel size including after tensor-parallel head sharding, matching Q and KV head-sharding axes, an unsharded head feature dimension, a divisible sequence length, `dq_reduction_steps` of 0 or 3, `context_parallel_load_balance=false` (each device computes full-sequence attention for its head subset, so the work is already balanced and the causal load-balancing reorder must stay off), and ICI-only context parallelism (`dcn_context_parallelism` must equal 1). It does not support MQA, packing, dropout, QK-Clip statistics, ragged attention, attention sinks, sparse indexer masks, chunked prefill, MoBA, or multimodal attention.
+
 ### CP Arithmetic Intensity
 
 The main communications are the same as FSDP (all gather weights and synchronize gradients), with an arithmetic intensity of `local_batch` / `sparsity`.

src/maxtext/configs/base.yml

@@ -1139,7 +1139,7 @@ cost_estimate_flops_bwd: -1 # -1 means using splash default cost estmiation, any
 dq_reduction_steps: 0 #the number of reduction steps. For now, only 3 or all the kv steps are supported.
 ### Determine if we want to use load balance for context parallelism
 context_parallel_load_balance: true
-context_parallel_strategy: "all_gather" # "all_gather" or "ring"
+context_parallel_strategy: "all_gather" # "all_gather", "ring", or "ulysses"
 context_parallel_reorder_strategy: "auto" # "auto", "dual_chunk_swap", or "striped"
 
 

src/maxtext/configs/types.py

@@ -1082,7 +1082,7 @@ class HardwareAndMesh(BaseModel):
   context_parallel_load_balance: bool = Field(True, description="Whether to use load balancing for context parallelism.")
   context_parallel_strategy: str = Field(
       "all_gather",
-      description="Strategy for context parallelism ('all_gather' or 'ring').",
+      description="Strategy for context parallelism ('all_gather', 'ring', or 'ulysses').",
   )
   context_parallel_reorder_strategy: ReorderStrategy = Field(
       ReorderStrategy.AUTO,
@@ -3573,6 +3573,9 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         self, f"dcn_{self.context_sharding}_parallelism", 1
     )
     context_parallel_strategy = self.context_parallel_strategy.lower()
+    if context_parallel_strategy not in ("all_gather", "ring", "ulysses"):
+      raise ValueError("context_parallel_strategy must be one of 'all_gather', 'ring', or 'ulysses'.")
+    self.context_parallel_strategy = context_parallel_strategy
     if (
         context_parallel_strategy == "ring"
         and "gpu" not in self.hardware
@@ -3628,6 +3631,67 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           f"ring_scan_unroll={self.ring_scan_unroll} was specified, but is only supported when "
           "context_parallel_strategy='ring'."
       )
+    if context_parallel_strategy == "ulysses":
+      if self.hardware != "tpu":
+        raise ValueError("Ulysses context parallelism (context_parallel_strategy='ulysses') is only supported on TPU.")
+      if self.context_sharding != "context":
+        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:
+        raise ValueError(
+            "TPU Ulysses attention requires explicit positive ici/dcn context parallelism values; "
+            "inferred (-1) sizes are not supported."
+        )
+      if context_parallel_size <= 1:
+        raise ValueError("TPU Ulysses attention requires context_parallel_size > 1.")
+      if dcn_context_parallel_size != 1:
+        raise ValueError("TPU Ulysses attention does not support dcn context parallelism yet.")
+      if self.attention != "flash":
+        raise ValueError("TPU Ulysses attention requires attention=flash.")
+      if not self.use_tokamax_splash:
+        raise ValueError("TPU Ulysses attention requires use_tokamax_splash=True.")
+      if self.use_jax_splash:
+        raise ValueError("TPU Ulysses attention requires use_jax_splash=False.")
+      if self.attention_type != "global":
+        raise ValueError("TPU Ulysses attention is initially supported only for global causal attention.")
+      if self.packing:
+        raise ValueError("TPU Ulysses attention does not support packing yet.")
+      if self.context_parallel_load_balance:
+        raise ValueError("TPU Ulysses attention does not support context_parallel_load_balance=True.")
+      if self.use_ragged_attention:
+        raise ValueError("TPU Ulysses attention does not support ragged attention.")
+      if self.attention_sink:
+        raise ValueError("TPU Ulysses attention does not support attention sinks.")
+      if self.use_indexer:
+        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:
+        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.")
+      if self.dq_reduction_steps not in (0, 3):
+        raise ValueError("TPU Ulysses attention requires dq_reduction_steps to be 0 or 3.")
+      if self.use_qk_clip:
+        raise ValueError("TPU Ulysses attention does not support QK-Clip statistics yet.")
+      if self.max_target_length % context_parallel_size != 0:
+        raise ValueError(
+            "TPU Ulysses attention requires max_target_length "
+            f"({self.max_target_length}) to be divisible by context_parallel_size ({context_parallel_size})."
+        )
+      if self.num_query_heads % context_parallel_size != 0:
+        raise ValueError(
+            "TPU Ulysses attention requires num_query_heads "
+            f"({self.num_query_heads}) to be divisible by context_parallel_size ({context_parallel_size})."
+        )
+      if self.num_kv_heads == 1:
+        raise ValueError("TPU Ulysses attention does not support MQA with context_parallel_size > 1.")
+      if self.num_kv_heads % context_parallel_size != 0:
+        raise ValueError(
+            "TPU Ulysses attention requires num_kv_heads "
+            f"({self.num_kv_heads}) to be divisible by context_parallel_size ({context_parallel_size})."
+        )
     # STRIPED reorder strategy is a Transformer Engine feature and is GPU-only.
     # AUTO is resolved in training because test code paths may load the same
     # config but use a different reorder path.

src/maxtext/kernels/attention/context_parallel_utils.py

@@ -0,0 +1,50 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Shared helpers for attention context-parallel sharding metadata."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import jax
+
+
+def mesh_axes_for_dim(axis_names: Any) -> tuple[Any, ...]:
+  """Returns the mesh axes attached to one tensor dimension."""
+  if axis_names is None:
+    return ()
+  if isinstance(axis_names, str):
+    return (axis_names,)
+  return tuple(axis for axis in axis_names if axis is not None)
+
+
+def mesh_axes_size(mesh: Any, axes: tuple[Any, ...], *, label: str) -> int:
+  """Returns the product of mesh sizes for a set of axes."""
+  size = 1
+  for axis in axes:
+    if axis not in mesh.shape:
+      raise ValueError(f"{label} requires mesh axis {axis!r} to exist.")
+    size *= mesh.shape[axis]
+  return size
+
+
+def with_axis_on_dim(axis_names: Any, axis: Any, dim: int) -> Any:
+  """Returns sharding axis names with one dimension replaced."""
+  axes = list(axis_names)
+  axes[dim] = axis
+  if isinstance(axis_names, jax.sharding.PartitionSpec):
+    return jax.sharding.PartitionSpec(*axes, unreduced=axis_names.unreduced, reduced=axis_names.reduced)
+  if isinstance(axis_names, tuple):
+    return tuple(axes)
+  return axes

src/maxtext/kernels/attention/tokamax_ring_attention.py

@@ -28,6 +28,7 @@
 import numpy as np
 
 from maxtext.common.common_types import MODEL_MODE_TRAIN
+from maxtext.kernels.attention import context_parallel_utils
 from maxtext.kernels.tokamax_splash_attention import ring_attention_kernel
 from maxtext.kernels.tokamax_splash_attention import splash_attention_kernel as tokamax_splash_kernel
 from maxtext.kernels.tokamax_splash_attention import splash_attention_mask as tokamax_splash_mask
@@ -39,42 +40,19 @@ def is_context_parallel_ring_requested(config: Any) -> bool:
   return config.context_parallel_strategy.lower() == "ring"
 
 
-def _mesh_axes_for_dim(axis_names: Any) -> tuple[Any, ...]:
-  if axis_names is None:
-    return ()
-  if isinstance(axis_names, str):
-    return (axis_names,)
-  return tuple(axis for axis in axis_names if axis is not None)
-
-
-def _mesh_axes_size(mesh: Any, axes: tuple[Any, ...]) -> int:
-  size = 1
-  for axis in axes:
-    if axis not in mesh.shape:
-      raise ValueError(f"TPU Tokamax ring attention requires mesh axis {axis!r} to exist.")
-    size *= mesh.shape[axis]
-  return size
-
-
 def with_sequence_axis(axis_names: Any, ring_axis: str, sequence_dim: int) -> Any:
   """Returns axis names with the sequence dimension set to the ring axis."""
   if axis_names is None:
     return None
   if len(axis_names) <= sequence_dim:
     raise ValueError("TPU Tokamax ring attention expects a sequence sharding dimension.")
-  axes = list(axis_names)
-  existing_sequence_axes = _mesh_axes_for_dim(axes[sequence_dim])
+  existing_sequence_axes = context_parallel_utils.mesh_axes_for_dim(axis_names[sequence_dim])
   if existing_sequence_axes and existing_sequence_axes != (ring_axis,):
     raise ValueError(
         "TPU Tokamax ring attention expects the existing sequence sharding to be "
         f"unsharded or exactly {(ring_axis,)}, got {existing_sequence_axes}."
     )
-  axes[sequence_dim] = ring_axis
-  if isinstance(axis_names, jax.sharding.PartitionSpec):
-    return jax.sharding.PartitionSpec(*axes, unreduced=axis_names.unreduced, reduced=axis_names.reduced)
-  if isinstance(axis_names, tuple):
-    return tuple(axes)
-  return axes
+  return context_parallel_utils.with_axis_on_dim(axis_names, ring_axis, sequence_dim)
 
 
 def _validate_ring_axis_only_on_sequence(
@@ -88,7 +66,7 @@ def _validate_ring_axis_only_on_sequence(
   for dim, axis_name in enumerate(axis_names):
     if dim == sequence_dim:
       continue
-    dim_axes = _mesh_axes_for_dim(axis_name)
+    dim_axes = context_parallel_utils.mesh_axes_for_dim(axis_name)
     if ring_axis in dim_axes:
       raise ValueError(
           "TPU Tokamax ring attention requires the context axis to appear only "
@@ -104,8 +82,8 @@ def validate_dkv_sharding(
     dkv_dim_kv: int,
 ) -> None:
   """Validates that the head-dim/D_KV dimension stays local for ring attention."""
-  q_dkv_axes = _mesh_axes_for_dim(axis_names_q[dkv_dim_q])
-  kv_dkv_axes = _mesh_axes_for_dim(axis_names_kv[dkv_dim_kv])
+  q_dkv_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_q[dkv_dim_q])
+  kv_dkv_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_kv[dkv_dim_kv])
   if q_dkv_axes or kv_dkv_axes:
     raise ValueError(
         "TPU Tokamax ring attention does not support sharding the D_KV/head-dim "
@@ -168,8 +146,8 @@ def validate_ring_mesh_axis(
       ring_axis=ring_axis,
   )
   expected_axes = (ring_axis,)
-  q_sequence_axes = _mesh_axes_for_dim(axis_names_q[sequence_dim_q])
-  key_value_sequence_axes = _mesh_axes_for_dim(axis_names_kv[sequence_dim_kv])
+  q_sequence_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_q[sequence_dim_q])
+  key_value_sequence_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_kv[sequence_dim_kv])
   if q_sequence_axes != expected_axes:
     raise ValueError(
         "TPU Tokamax ring attention requires Q sequence sharding to be exactly "
@@ -193,10 +171,10 @@ def validate_head_sharding(
     head_dim_kv: int,
 ) -> None:
   """Validates that local head layout preserves GQA/MQA head mapping."""
-  q_head_axes = _mesh_axes_for_dim(axis_names_q[head_dim_q])
-  kv_head_axes = _mesh_axes_for_dim(axis_names_kv[head_dim_kv])
-  q_head_shards = _mesh_axes_size(mesh, q_head_axes)
-  kv_head_shards = _mesh_axes_size(mesh, kv_head_axes)
+  q_head_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_q[head_dim_q])
+  kv_head_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_kv[head_dim_kv])
+  q_head_shards = context_parallel_utils.mesh_axes_size(mesh, q_head_axes, label="TPU Tokamax ring attention")
+  kv_head_shards = context_parallel_utils.mesh_axes_size(mesh, kv_head_axes, label="TPU Tokamax ring attention")
   if num_query_heads % q_head_shards != 0:
     raise ValueError(
         "TPU Tokamax ring attention requires num_query_heads "