Add Ulysses ring split override

Author: csgoogle

README.md

@@ -603,13 +603,17 @@ To generate images, run the following command:
   Public configs still shard sequence only over the `context` mesh axis. The attention kernel privately reshapes
   that context axis into hidden ring and Ulysses axes, runs the Ulysses all-to-all over the hidden Ulysses axis,
   and reuses Tokamax ring attention over the hidden ring axis.
+  By default, the split is selected automatically. For tuning, set
+  `ulysses_ring_ulysses_parallelism=<ulysses_shards>`; ring shards are derived as
+  `ici_context_parallelism / ulysses_ring_ulysses_parallelism`.
 
   ```bash
   python src/maxdiffusion/generate_wan.py \
   src/maxdiffusion/configs/base_wan_i2v_27b.yml \
   attention="ulysses_ring" \
   dcn_context_parallelism=<num_slices> \
   ici_context_parallelism=<context_shards_per_slice> \
+  ulysses_ring_ulysses_parallelism=<optional_ulysses_shards> \
   ...
   ```
 

docs/tpu_multihost_wan_bench.md

@@ -188,7 +188,7 @@ run_case ulysses_ring_dp2_cp8    ulysses_ring 2 8 1
 
 ## Topology Note
 
-TPU v7x exposes dual chiplets as two JAX devices. For `ulysses_ring`, expose only the total sequence sharding through `context`; the attention kernel derives a private ring and Ulysses split from that axis.
+TPU v7x exposes dual chiplets as two JAX devices. For `ulysses_ring`, expose only the total sequence sharding through `context`; the attention kernel derives a private ring and Ulysses split from that axis. To tune that split explicitly, set `ulysses_ring_ulysses_parallelism`; ring shards are derived as `ici_context_parallelism / ulysses_ring_ulysses_parallelism`.
 - `4x4` uses tensor `4`, so the dual-chip pairing is still inside the Ulysses side.
 
 The plain `ring` baseline has no Ulysses group, so it cannot preserve that property by construction.

docs/tpu_wan_bench_guide.md

@@ -143,7 +143,7 @@ Set in `src/maxdiffusion/configs/base_wan_27b.yml` or overridden on the command
 **Parallelism rule**: product of all ICI axes must equal 8 (chips per host):
 - `ici_dp × ici_fsdp × ici_cp × ici_tp = 8`
 
-For `ulysses_ring`, set the desired total sequence shards with `ici_context_parallelism`; the internal ring and Ulysses split is selected by the attention kernel.
+For `ulysses_ring`, set the desired total sequence shards with `ici_context_parallelism`; the internal ring and Ulysses split is selected by the attention kernel. To tune it manually, set `ulysses_ring_ulysses_parallelism=<ulysses_shards>` and the ring shard count is derived as `ici_context_parallelism / ulysses_ring_ulysses_parallelism`.
 
 ---
 

src/maxdiffusion/configs/base_wan_14b.yml

@@ -67,6 +67,9 @@ split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
 use_base2_exp: True
 use_experimental_scheduler: True
+# Optional for attention=ulysses_ring. -1 auto-selects the hidden split; otherwise
+# this many context shards are used for Ulysses and ring shards are context / this.
+ulysses_ring_ulysses_parallelism: -1
 flash_min_seq_length: 4096
 dropout: 0.0
 

src/maxdiffusion/configs/base_wan_1_3b.yml

@@ -63,6 +63,9 @@ split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
 use_base2_exp: True
 use_experimental_scheduler: True
+# Optional for attention=ulysses_ring. -1 auto-selects the hidden split; otherwise
+# this many context shards are used for Ulysses and ring shards are context / this.
+ulysses_ring_ulysses_parallelism: -1
 flash_min_seq_length: 0
 
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.

src/maxdiffusion/configs/base_wan_27b.yml

@@ -67,6 +67,9 @@ split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
 use_base2_exp: True
 use_experimental_scheduler: True
+# Optional for attention=ulysses_ring. -1 auto-selects the hidden split; otherwise
+# this many context shards are used for Ulysses and ring shards are context / this.
+ulysses_ring_ulysses_parallelism: -1
 flash_min_seq_length: 4096
 dropout: 0.0
 

src/maxdiffusion/configs/base_wan_animate.yml

@@ -65,6 +65,9 @@ split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
 use_base2_exp: True
 use_experimental_scheduler: True
+# Optional for attention=ulysses_ring. -1 auto-selects the hidden split; otherwise
+# this many context shards are used for Ulysses and ring shards are context / this.
+ulysses_ring_ulysses_parallelism: -1
 flash_min_seq_length: 4096
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -67,6 +67,9 @@ split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
 use_base2_exp: True
 use_experimental_scheduler: True
+# Optional for attention=ulysses_ring. -1 auto-selects the hidden split; otherwise
+# this many context shards are used for Ulysses and ring shards are context / this.
+ulysses_ring_ulysses_parallelism: -1
 flash_min_seq_length: 4096
 dropout: 0.0
 

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -67,6 +67,9 @@ split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
 use_base2_exp: True
 use_experimental_scheduler: True
+# Optional for attention=ulysses_ring. -1 auto-selects the hidden split; otherwise
+# this many context shards are used for Ulysses and ring shards are context / this.
+ulysses_ring_ulysses_parallelism: -1
 flash_min_seq_length: 4096
 dropout: 0.0
 

src/maxdiffusion/models/attention_flax.py

@@ -206,11 +206,28 @@ def _replace_mesh_axis_names(axis_names, old_axis: str, new_axes: tuple[str, ...
   )
 
 
-def _choose_internal_ulysses_shards(context_shards: int, heads: int) -> int:
+def _choose_internal_ulysses_shards(
+    context_shards: int,
+    heads: int,
+    requested_ulysses_shards: int = -1,
+) -> int:
   """Choose a hidden Ulysses split inside the public context axis."""
   if context_shards <= 1:
     raise ValueError(f"Ulysses ring attention requires context_shards > 1, got {context_shards}.")
 
+  if requested_ulysses_shards and requested_ulysses_shards > 0:
+    if context_shards % requested_ulysses_shards != 0:
+      raise ValueError(
+          "Ulysses ring attention requires the requested Ulysses shard count to divide the context shard count, "
+          f"got context_shards={context_shards} and ulysses_shards={requested_ulysses_shards}."
+      )
+    if heads % requested_ulysses_shards != 0:
+      raise ValueError(
+          "Ulysses ring attention requires the number of heads to be divisible by the requested Ulysses shard "
+          f"count, got heads={heads} and ulysses_shards={requested_ulysses_shards}."
+      )
+    return requested_ulysses_shards
+
   balanced_limit = int(math.sqrt(context_shards))
   balanced_candidates = [
       factor
@@ -844,6 +861,7 @@ def _ulysses_ring_attention(
     ring_axis: str = INTERNAL_RING_AXIS,
     use_base2_exp: bool = False,
     use_experimental_scheduler: bool = False,
+    ulysses_ring_ulysses_parallelism: int = -1,
 ) -> jax.Array:
   """2D context-parallel attention using a private Ulysses x ring mesh.
 
@@ -857,7 +875,11 @@ def _ulysses_ring_attention(
     raise ValueError(f"Ulysses ring attention requires mesh axis {context_axis!r}, got mesh axes {mesh.shape}.")
 
   num_context_shards = mesh.shape[context_axis]
-  num_ulysses_shards = _choose_internal_ulysses_shards(num_context_shards, heads)
+  num_ulysses_shards = _choose_internal_ulysses_shards(
+      num_context_shards,
+      heads,
+      requested_ulysses_shards=ulysses_ring_ulysses_parallelism,
+  )
   num_ring_shards = num_context_shards // num_ulysses_shards
   internal_mesh = _create_internal_ulysses_ring_mesh(
       mesh,
@@ -1166,6 +1188,7 @@ def ulysses_ring_kernel(q, k, v, context):
       attention_mask=context["attention_mask"],
       use_base2_exp=context["use_base2_exp"],
       use_experimental_scheduler=context["use_experimental_scheduler"],
+      ulysses_ring_ulysses_parallelism=context["ulysses_ring_ulysses_parallelism"],
   )
 
 
@@ -1279,6 +1302,7 @@ def _apply_attention(
     attention_mask: Array = None,
     use_base2_exp: bool = False,
     use_experimental_scheduler: bool = False,
+    ulysses_ring_ulysses_parallelism: int = -1,
 ):
   """Routes to different attention kernels using a module-level registry."""
 
@@ -1316,6 +1340,7 @@ def _apply_attention(
       "scale": scale,
       "use_base2_exp": use_base2_exp,
       "use_experimental_scheduler": use_experimental_scheduler,
+      "ulysses_ring_ulysses_parallelism": ulysses_ring_ulysses_parallelism,
       "dim_head": dim_head,
       "split_head_dim": split_head_dim,
       "float32_qk_product": float32_qk_product,
@@ -1521,10 +1546,12 @@ def __init__(
       residual_checkpoint_name: str | None = None,
       use_base2_exp: bool = False,
       use_experimental_scheduler: bool = False,
+      ulysses_ring_ulysses_parallelism: int = -1,
   ):
     self.dpa_layer = None
     self.use_base2_exp = use_base2_exp
     self.use_experimental_scheduler = use_experimental_scheduler
+    self.ulysses_ring_ulysses_parallelism = ulysses_ring_ulysses_parallelism
     if attention_kernel == "cudnn_flash_te":
       from transformer_engine.jax.flax.transformer import DotProductAttention  # pytype: disable=import-error
 
@@ -1587,6 +1614,9 @@ def apply_attention(self, query: Array, key: Array, value: Array, attention_mask
         attention_mask=attention_mask,
         use_base2_exp=self.use_base2_exp if hasattr(self, "use_base2_exp") else False,
         use_experimental_scheduler=self.use_experimental_scheduler if hasattr(self, "use_experimental_scheduler") else False,
+        ulysses_ring_ulysses_parallelism=(
+            self.ulysses_ring_ulysses_parallelism if hasattr(self, "ulysses_ring_ulysses_parallelism") else -1
+        ),
     )
 
 
@@ -1607,6 +1637,7 @@ class AttentionOp(nn.Module):
   quant: Quant = None
   use_base2_exp: bool = False
   use_experimental_scheduler: bool = False
+  ulysses_ring_ulysses_parallelism: int = -1
 
   def setup(self):
     self.dpa_layer = None
@@ -1654,6 +1685,7 @@ def apply_attention(self, query: Array, key: Array, value: Array, attention_mask
         attention_mask=attention_mask,
         use_base2_exp=self.use_base2_exp,
         use_experimental_scheduler=self.use_experimental_scheduler,
+        ulysses_ring_ulysses_parallelism=self.ulysses_ring_ulysses_parallelism,
     )
 
 
@@ -1692,6 +1724,7 @@ def __init__(
       image_seq_len: Optional[int] = None,  # New for I2V
       use_base2_exp: bool = False,
       use_experimental_scheduler: bool = False,
+      ulysses_ring_ulysses_parallelism: int = -1,
   ):
     if attention_kernel in {"flash", "cudnn_flash_te"} and mesh is None:
       raise ValueError(f"The flash attention kernel requires a value for mesh, but mesh is {self.mesh}")
@@ -1740,6 +1773,7 @@ def __init__(
         residual_checkpoint_name=residual_checkpoint_name,
         use_base2_exp=use_base2_exp,
         use_experimental_scheduler=use_experimental_scheduler,
+        ulysses_ring_ulysses_parallelism=ulysses_ring_ulysses_parallelism,
     )
     # None axes corresponds to the stacked weights across all blocks
     # because of the use of nnx.vmap and nnx.scan.