diff --git a/README.md b/README.md
index 4f5a1fcb7..83657cf14 100755
--- a/README.md
+++ b/README.md
@@ -619,6 +619,31 @@ To generate images, run the following command:
 
   In our Wan2.2 I2V benchmarks at 40 inference steps, 81 frames, and `720x1280` resolution, Ulysses improved inference time by roughly `~10%` compared with flash attention, with about `~20s` lower latency on the v6e-8 and v7x-8 TPU setup.
 
+  #### Chunked Ulysses Attention (Overlapping Communication and Compute)
+
+  If you observe a major `all-to-all` communication bottleneck (especially when communication overhead is more pronounced compared to attention computation), you can enable **Chunked Ulysses Attention**.
+
+  By setting `ulysses_attention_chunks` greater than 1, MaxDiffusion splits the Ulysses all-to-all communication and attention computation into head-group passes (chunks). This allows XLA to overlap the all-to-all communication of one chunk with the head-parallel local attention compute of another chunk, significantly mitigating the communication bottleneck.
+
+  This chunking technique is supported and works for both plain Ulysses attention (`attention="ulysses"`) and hybrid Ulysses+Ring 2D attention/context parallelism (`attention="ulysses_ring"`).
+
+  To enable chunked Ulysses attention, set the corresponding override (e.g. `ulysses_attention_chunks=2` or `ulysses_attention_chunks=5`) in your config YAML or command line:
+
+  ```bash
+  python src/maxdiffusion/generate_wan.py \
+  src/maxdiffusion/configs/base_wan_i2v_27b.yml \
+  attention="ulysses" \
+  ici_context_parallelism=4 \
+  ulysses_attention_chunks=2 \
+  ...
+  ```
+
+  > [!IMPORTANT]
+  > For communication-compute overlap to be effective on TPUs, you must enable the following XLA flags before running:
+  > ```bash
+  > export XLA_FLAGS="--xla_tpu_enable_async_all_to_all=true --xla_tpu_overlap_compute_collective_tc=true"
+  > ```
+
   ### Caching Mechanisms
 
   Wan 2.x pipelines support several caching strategies to accelerate inference by skipping redundant transformer forward passes. These are **mutually exclusive** — enable only one at a time.
diff --git a/pyproject.toml b/pyproject.toml
index 596b1b013..6065c8a3e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -92,17 +92,19 @@ packages = ["src/maxdiffusion", "src/install_maxdiffusion_extra_deps"]
 
 [tool.ruff]
 # Never enforce `E501` (line length violations).
+line-length = 119
+
+[tool.ruff.lint]
 ignore = ["C901", "E501", "E741", "F402", "F823", "E402", "I001"]
 select = ["C", "E", "F", "I", "W"]
-line-length = 119
 
 # Ignore import violations in all `__init__.py` files.
-[tool.ruff.per-file-ignores]
+[tool.ruff.lint.per-file-ignores]
 "__init__.py" = ["E402", "F401", "F403", "F811"]
 "src/maxdiffusion/utils/dummy_*.py" = ["F401"]
 "src/maxdiffusion/pyconfig.py" = ["E721"]
 
-[tool.ruff.isort]
+[tool.ruff.lint.isort]
 lines-after-imports = 2
 known-first-party = ["maxdiffusion"]
 
diff --git a/src/maxdiffusion/configs/base_wan_14b.yml b/src/maxdiffusion/configs/base_wan_14b.yml
index 18a2b4eee..7c1e76fc0 100644
--- a/src/maxdiffusion/configs/base_wan_14b.yml
+++ b/src/maxdiffusion/configs/base_wan_14b.yml
@@ -88,6 +88,12 @@ use_base2_exp: True
 use_experimental_scheduler: True
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
+# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.
+# For communication-compute overlap to be effective, enable the following XLA flags:
+#   --xla_tpu_enable_async_all_to_all=true
+#   --xla_tpu_overlap_compute_collective_tc=true
+# (Refer to README.md for the full recommended XLA_FLAGS list)
+ulysses_attention_chunks: 1
 flash_min_seq_length: 4096
 dropout: 0.0
 
diff --git a/src/maxdiffusion/configs/base_wan_1_3b.yml b/src/maxdiffusion/configs/base_wan_1_3b.yml
index c0de05c9e..3daa3dc37 100644
--- a/src/maxdiffusion/configs/base_wan_1_3b.yml
+++ b/src/maxdiffusion/configs/base_wan_1_3b.yml
@@ -85,6 +85,12 @@ use_base2_exp: True
 use_experimental_scheduler: True
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
+# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.
+# For communication-compute overlap to be effective, enable the following XLA flags:
+#   --xla_tpu_enable_async_all_to_all=true
+#   --xla_tpu_overlap_compute_collective_tc=true
+# (Refer to README.md for the full recommended XLA_FLAGS list)
+ulysses_attention_chunks: 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.
diff --git a/src/maxdiffusion/configs/base_wan_27b.yml b/src/maxdiffusion/configs/base_wan_27b.yml
index 4e5f7642f..bfb6f00c5 100644
--- a/src/maxdiffusion/configs/base_wan_27b.yml
+++ b/src/maxdiffusion/configs/base_wan_27b.yml
@@ -94,6 +94,12 @@ use_base2_exp: True
 use_experimental_scheduler: True
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
+# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.
+# For communication-compute overlap to be effective, enable the following XLA flags:
+#   --xla_tpu_enable_async_all_to_all=true
+#   --xla_tpu_overlap_compute_collective_tc=true
+# (Refer to README.md for the full recommended XLA_FLAGS list)
+ulysses_attention_chunks: 1
 flash_min_seq_length: 4096
 dropout: 0.0
 
diff --git a/src/maxdiffusion/configs/base_wan_animate.yml b/src/maxdiffusion/configs/base_wan_animate.yml
index e0abf4515..10831d272 100644
--- a/src/maxdiffusion/configs/base_wan_animate.yml
+++ b/src/maxdiffusion/configs/base_wan_animate.yml
@@ -86,6 +86,12 @@ use_base2_exp: True
 use_experimental_scheduler: True
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
+# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.
+# For communication-compute overlap to be effective, enable the following XLA flags:
+#   --xla_tpu_enable_async_all_to_all=true
+#   --xla_tpu_overlap_compute_collective_tc=true
+# (Refer to README.md for the full recommended XLA_FLAGS list)
+ulysses_attention_chunks: 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.
diff --git a/src/maxdiffusion/configs/base_wan_i2v_14b.yml b/src/maxdiffusion/configs/base_wan_i2v_14b.yml
index 5c59ddbdc..12e8eb147 100644
--- a/src/maxdiffusion/configs/base_wan_i2v_14b.yml
+++ b/src/maxdiffusion/configs/base_wan_i2v_14b.yml
@@ -88,6 +88,12 @@ use_base2_exp: True
 use_experimental_scheduler: True
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
+# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.
+# For communication-compute overlap to be effective, enable the following XLA flags:
+#   --xla_tpu_enable_async_all_to_all=true
+#   --xla_tpu_overlap_compute_collective_tc=true
+# (Refer to README.md for the full recommended XLA_FLAGS list)
+ulysses_attention_chunks: 1
 flash_min_seq_length: 4096
 dropout: 0.0
 
diff --git a/src/maxdiffusion/configs/base_wan_i2v_27b.yml b/src/maxdiffusion/configs/base_wan_i2v_27b.yml
index 8c4bf8853..9cab72163 100644
--- a/src/maxdiffusion/configs/base_wan_i2v_27b.yml
+++ b/src/maxdiffusion/configs/base_wan_i2v_27b.yml
@@ -88,6 +88,12 @@ use_base2_exp: True
 use_experimental_scheduler: True
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
+# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.
+# For communication-compute overlap to be effective, enable the following XLA flags:
+#   --xla_tpu_enable_async_all_to_all=true
+#   --xla_tpu_overlap_compute_collective_tc=true
+# (Refer to README.md for the full recommended XLA_FLAGS list)
+ulysses_attention_chunks: 1
 flash_min_seq_length: 4096
 dropout: 0.0
 
diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py
index 49f35f490..9f6c04b99 100644
--- a/src/maxdiffusion/generate_wan.py
+++ b/src/maxdiffusion/generate_wan.py
@@ -178,7 +178,8 @@ def inference_generate_video(config, pipeline, filename_prefix=""):
   negative_prompt = [config.negative_prompt] * config.global_batch_size_to_train_on
 
   max_logging.log(
-      f"Num steps: {config.num_inference_steps}, height: {config.height}, width: {config.width}, frames: {config.num_frames}, video: {filename_prefix}"
+      f"Num steps: {config.num_inference_steps}, height: {config.height}, width: {config.width},"
+      f" frames: {config.num_frames}, video: {filename_prefix}"
   )
 
   videos = call_pipeline(config, pipeline, prompt, negative_prompt)
@@ -314,7 +315,8 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
   negative_prompt = [config.negative_prompt] * config.global_batch_size_to_train_on
 
   max_logging.log(
-      f"Num steps: {config.num_inference_steps}, height: {config.height}, width: {config.width}, frames: {config.num_frames}"
+      f"Num steps: {config.num_inference_steps}, height: {config.height}, width: {config.width},"
+      f" frames: {config.num_frames}"
   )
   # Warmup with 2 denoising steps instead of a full run: step 0 runs the
   # high-noise transformer and step 1 crosses the boundary to the low-noise
@@ -331,7 +333,8 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
     videos = call_pipeline(config, pipeline, prompt, negative_prompt, num_inference_steps=warmup_steps)
   if isinstance(videos, tuple):
     videos, warmup_trace = videos
-    max_logging.log("Warmup breakdown: " + ", ".join(f"{stage}={seconds:.1f}s" for stage, seconds in warmup_trace.items()))
+    warmup_str = ", ".join(f"{stage}={seconds:.1f}s" for stage, seconds in warmup_trace.items())
+    max_logging.log(f"Warmup breakdown: {warmup_str}")
 
   # Serialize any newly-compiled shapes synchronously while still inside
   # warmup-accounted time; a background save would compete with the first
@@ -390,17 +393,15 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
     vae_decode_total = trace.get("vae_decode", 0.0)
     vae_decode_tpu = trace.get("vae_decode_tpu", 0.0)
     vae_decode_post = vae_decode_total - vae_decode_tpu
-    summary.extend(
-        [
-            f"  {'─' * 40}",
-            f"  Conditioning:        {trace.get('conditioning', 0.0):>7.1f}s",
-            f"    - VAE Encode:      {trace.get('vae_encode', 0.0):>7.1f}s",
-            f"  Denoise Total:       {trace.get('denoise_total', 0.0):>7.1f}s",
-            f"  VAE Decode:          {vae_decode_total:>7.1f}s",
-            f"    - TPU Compute:     {vae_decode_tpu:>7.1f}s",
-            f"    - Host Formatting: {vae_decode_post:>7.1f}s",
-        ]
-    )
+    summary.extend([
+        f"  {'─' * 40}",
+        f"  Conditioning:        {trace.get('conditioning', 0.0):>7.1f}s",
+        f"    - VAE Encode:      {trace.get('vae_encode', 0.0):>7.1f}s",
+        f"  Denoise Total:       {trace.get('denoise_total', 0.0):>7.1f}s",
+        f"  VAE Decode:          {vae_decode_total:>7.1f}s",
+        f"    - TPU Compute:     {vae_decode_tpu:>7.1f}s",
+        f"    - Host Formatting: {vae_decode_post:>7.1f}s",
+    ])
   summary.append(f"{'=' * 50}")
   max_logging.log("\n".join(summary))
 
diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py
index a5a522653..5a62e4e55 100644
--- a/src/maxdiffusion/models/attention_flax.py
+++ b/src/maxdiffusion/models/attention_flax.py
@@ -419,6 +419,86 @@ def _build_padding_segment_ids(
   return segment_ids_cls(q=q_segment_ids, kv=kv_segment_ids)
 
 
+def _ulysses_head_chunk_ranges(num_heads: int, ulysses_shards: int, num_chunks: int):
+  """Build head-axis ranges for chunked Ulysses all-to-all.
+
+  The Ulysses all-to-all splits each local chunk's head axis over
+  `ulysses_shards`, so every returned range length is a multiple of
+  `ulysses_shards`. When `num_chunks` does not evenly divide the number of
+  Ulysses-sized head groups, earlier chunks get the floor-sized range and the
+  final chunk carries the remainder.
+
+  Returns:
+    A list of `(start, end)` half-open ranges over the head axis. Concatenating
+    tensors sliced with these ranges along the head axis restores the original
+    head layout. For `num_chunks <= 1`, returns `[(0, num_heads)]`, which is the
+    unchunked all-to-all path.
+  """
+  if num_chunks <= 1:
+    return [(0, num_heads)]
+  if num_heads % ulysses_shards != 0:
+    raise ValueError(
+        "Ulysses attention requires the number of heads to be divisible by the Ulysses shard count, "
+        f"got heads={num_heads} and ulysses_shards={ulysses_shards}."
+    )
+
+  head_groups = num_heads // ulysses_shards
+  num_chunks = min(num_chunks, head_groups)
+  regular_groups_per_chunk = max(1, head_groups // num_chunks)
+
+  ranges = []
+  start_group = 0
+  for chunk_idx in range(num_chunks):
+    end_group = head_groups if chunk_idx == num_chunks - 1 else min(start_group + regular_groups_per_chunk, head_groups)
+    if start_group >= end_group:
+      break
+    ranges.append((start_group * ulysses_shards, end_group * ulysses_shards))
+    start_group = end_group
+  return ranges
+
+
+def _run_chunked_ulysses_attention(
+    query: jax.Array,
+    key: jax.Array,
+    value: jax.Array,
+    num_heads: int,
+    ulysses_shards: int,
+    ulysses_attention_chunks: int,
+    attention_fn,
+) -> jax.Array:
+  """Runs Ulysses attention chunked or unchunked along the head axis.
+
+  Splits the attention compute and communication into head-group chunks so XLA
+  can overlap communication and compute.
+
+  Args:
+    query: The query tensor, [B, H, S, D].
+    key: The key tensor, [B, H, S, D].
+    value: The value tensor, [B, H, S, D].
+    num_heads: The number of heads in query.
+    ulysses_shards: The Ulysses/context shard count.
+    ulysses_attention_chunks: Number of head-group chunks to split into.
+    attention_fn: The local Ulysses attention function to call on each chunk,
+      taking (query, key, value) and returning the attention output.
+
+  Returns:
+    The concatenated attention output tensor.
+  """
+  head_chunk_ranges = _ulysses_head_chunk_ranges(num_heads, ulysses_shards, ulysses_attention_chunks)
+  if len(head_chunk_ranges) > 1:
+    chunk_outputs = [
+        attention_fn(
+            query[:, start:end],
+            key[:, start:end],
+            value[:, start:end],
+        )
+        for start, end in head_chunk_ranges
+    ]
+    return jnp.concatenate(chunk_outputs, axis=1)
+  else:
+    return attention_fn(query, key, value)
+
+
 def _tpu_flash_attention(
     query: jax.Array,
     key: jax.Array,
@@ -546,7 +626,9 @@ def wrap_flash_attention(query, key, value):
           ),
           save_residuals=False,
           ring_axis=CONTEXT,
-          rotate_segment_ids=False,  # We don't rotate segment ids in tokamax ring attention because our segment ids is for padding each kv shard has same segment ids
+          # We don't rotate segment ids in tokamax ring attention because our
+          # segment ids is for padding each kv shard has same segment ids
+          rotate_segment_ids=False,
       )
     else:
       splash_kernel = splash_attention_kernel.make_splash_mha(
@@ -647,6 +729,7 @@ def _ulysses_attention(
     use_base2_exp: bool = True,
     use_experimental_scheduler: bool = False,
     use_fixed_m: bool = False,
+    ulysses_attention_chunks: int = 1,
 ) -> jax.Array:
   """Ulysses sequence-parallel attention.
 
@@ -669,6 +752,7 @@ def _ulysses_attention(
         "Ulysses attention requires the number of heads to be divisible by the context shard count, "
         f"got heads={num_heads} and context_shards={num_shards}."
     )
+
   if not use_custom_kernel:
     block_sizes = _select_flash_block_sizes(query, key, flash_block_sizes, dtype, "flash")
 
@@ -792,7 +876,6 @@ def wrap_ulysses_attention(query, key, value):
         "Warning, batch dimension should be shardable among the devices in data and fsdp"
         f" axis, batch dimension: {query.shape[0]}, devices_in_batch_sharding: {devices_in_batch_sharding}"
     )
-
   # Fold the (CFG) batch into the heads axis around the Ulysses exchange.
   # Each (batch, head) pair is an independent attention problem, so
   # [B, H, S, D] -> [1, B*H, S, D] is mathematically identity — but it makes
@@ -807,8 +890,19 @@ def wrap_ulysses_attention(query, key, value):
     query = query.reshape(1, batch * num_heads, *query.shape[2:])
     key = key.reshape(1, batch * num_heads, *key.shape[2:])
     value = value.reshape(1, batch * num_heads, *value.shape[2:])
-
-  x = wrap_ulysses_attention(query, key, value)
+    effective_num_heads = batch * num_heads
+  else:
+    effective_num_heads = num_heads
+
+  x = _run_chunked_ulysses_attention(
+      query,
+      key,
+      value,
+      effective_num_heads,
+      num_shards,
+      ulysses_attention_chunks,
+      wrap_ulysses_attention,
+  )
 
   if fold_batch:
     x = x.reshape(batch, num_heads, *x.shape[2:])
@@ -836,6 +930,7 @@ def _ulysses_ring_attention(
     use_base2_exp: bool = False,
     use_experimental_scheduler: bool = False,
     ulysses_shards: int = -1,
+    ulysses_attention_chunks: int = 1,
 ) -> jax.Array:
   """2D context-parallel attention using a private Ulysses x ring mesh.
 
@@ -877,6 +972,7 @@ def _ulysses_ring_attention(
   query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_sequence_shards)
   key, _ = _reshape_data_for_flash(key, heads, num_sequence_shards)
   value, _ = _reshape_data_for_flash(value, heads, num_sequence_shards)
+  num_heads = query.shape[1]
 
   block_sizes = _select_flash_block_sizes(query, key, flash_block_sizes, dtype, "tokamax_ring")
 
@@ -965,7 +1061,15 @@ def wrap_ulysses_ring_attention(query, key, value):
         "Warning, batch dimension should be shardable among the devices in data and fsdp"
         f" axis, batch dimension: {query.shape[0]}, devices_in_batch_sharding: {devices_in_batch_sharding}"
     )
-  x = wrap_ulysses_ring_attention(query, key, value)
+  x = _run_chunked_ulysses_attention(
+      query,
+      key,
+      value,
+      num_heads,
+      num_ulysses_shards,
+      ulysses_attention_chunks,
+      wrap_ulysses_ring_attention,
+  )
   x = jax.lax.with_sharding_constraint(x, q_axis_names)
   x = x[:, :, :orig_q_seq_len, :]
   x = _reshape_heads_to_head_dim(x)
@@ -991,6 +1095,7 @@ def _ulysses_ring_custom_attention(
     use_experimental_scheduler: bool = False,
     bidirectional: bool = False,
     use_fixed_m: bool = False,
+    ulysses_attention_chunks: int = 1,
 ) -> jax.Array:
   """Hybrid Ulysses + Ring (USP) with the CUSTOM splash kernel on main's mesh.
 
@@ -1141,7 +1246,15 @@ def wrap_ulysses_ring_attention(query, key, value):
     attention_output = a2a(attention_output, split_axis=2, concat_axis=1)
     return attention_output
 
-  x = wrap_ulysses_ring_attention(query, key, value)
+  x = _run_chunked_ulysses_attention(
+      query,
+      key,
+      value,
+      num_heads,
+      num_ulysses_shards,
+      ulysses_attention_chunks,
+      wrap_ulysses_ring_attention,
+  )
   x = jax.lax.with_sharding_constraint(x, q_axis_names)
   x = x[:, :, :orig_q_seq_len, :]
   x = _reshape_heads_to_head_dim(x)
@@ -1291,6 +1404,7 @@ def ulysses_custom_kernel(q, k, v, context):
       use_custom_kernel=True,
       use_base2_exp=context.get("use_base2_exp", True),
       use_experimental_scheduler=context.get("use_experimental_scheduler", False),
+      ulysses_attention_chunks=context["ulysses_attention_chunks"],
   )
 
 
@@ -1312,6 +1426,7 @@ def ulysses_ring_custom_kernel(q, k, v, context):
       ulysses_shards=context["ulysses_shards"],
       use_base2_exp=context.get("use_base2_exp", True),
       use_experimental_scheduler=context.get("use_experimental_scheduler", False),
+      ulysses_attention_chunks=context["ulysses_attention_chunks"],
   )
 
 
@@ -1364,6 +1479,7 @@ def ulysses_ring_custom_bidir_kernel(q, k, v, context):
       use_base2_exp=context.get("use_base2_exp", True),
       use_experimental_scheduler=context.get("use_experimental_scheduler", False),
       bidirectional=True,
+      ulysses_attention_chunks=context["ulysses_attention_chunks"],
   )
 
 
@@ -1404,6 +1520,7 @@ def ulysses_kernel(q, k, v, context):
       mask_padding_tokens=context["mask_padding_tokens"],
       residual_checkpoint_name=context["residual_checkpoint_name"],
       attention_mask=context["attention_mask"],
+      ulysses_attention_chunks=context["ulysses_attention_chunks"],
   )
 
 
@@ -1425,6 +1542,7 @@ def ulysses_ring_kernel(q, k, v, context):
       use_base2_exp=context["use_base2_exp"],
       use_experimental_scheduler=context["use_experimental_scheduler"],
       ulysses_shards=context["ulysses_shards"],
+      ulysses_attention_chunks=context["ulysses_attention_chunks"],
   )
 
 
@@ -1537,6 +1655,7 @@ def _apply_attention(
     use_base2_exp: bool = False,
     use_experimental_scheduler: bool = False,
     ulysses_shards: int = -1,
+    ulysses_attention_chunks: int = 1,
 ):
   """Routes to different attention kernels using a module-level registry."""
 
@@ -1568,6 +1687,7 @@ def _apply_attention(
       "use_base2_exp": use_base2_exp,
       "use_experimental_scheduler": use_experimental_scheduler,
       "ulysses_shards": ulysses_shards,
+      "ulysses_attention_chunks": ulysses_attention_chunks,
       "dim_head": dim_head,
       "split_head_dim": split_head_dim,
       "float32_qk_product": float32_qk_product,
@@ -1781,11 +1901,13 @@ def __init__(
       use_base2_exp: bool = False,
       use_experimental_scheduler: bool = False,
       ulysses_shards: int = -1,
+      ulysses_attention_chunks: int = 1,
   ):
     self.dpa_layer = None
     self.use_base2_exp = use_base2_exp
     self.use_experimental_scheduler = use_experimental_scheduler
     self.ulysses_shards = ulysses_shards
+    self.ulysses_attention_chunks = ulysses_attention_chunks
     if attention_kernel == "cudnn_flash_te":
       from transformer_engine.jax.flax.transformer import DotProductAttention  # pytype: disable=import-error
 
@@ -1849,6 +1971,7 @@ def apply_attention(self, query: Array, key: Array, value: Array, 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_shards=(self.ulysses_shards if hasattr(self, "ulysses_shards") else -1),
+        ulysses_attention_chunks=(self.ulysses_attention_chunks if hasattr(self, "ulysses_attention_chunks") else 1),
     )
 
 
@@ -1870,6 +1993,7 @@ class AttentionOp(nn.Module):
   use_base2_exp: bool = False
   use_experimental_scheduler: bool = False
   ulysses_shards: int = -1
+  ulysses_attention_chunks: int = 1
 
   def setup(self):
     self.dpa_layer = None
@@ -1918,6 +2042,7 @@ def apply_attention(self, query: Array, key: Array, value: Array, attention_mask
         use_base2_exp=self.use_base2_exp,
         use_experimental_scheduler=self.use_experimental_scheduler,
         ulysses_shards=self.ulysses_shards,
+        ulysses_attention_chunks=self.ulysses_attention_chunks,
     )
 
 
@@ -1960,6 +2085,7 @@ def __init__(
         "use_base2_exp": False,
         "use_experimental_scheduler": False,
         "ulysses_shards": -1,
+        "ulysses_attention_chunks": 1,
         **(attention_config or {}),
     }
 
@@ -2011,6 +2137,7 @@ def __init__(
         use_base2_exp=attention_config["use_base2_exp"],
         use_experimental_scheduler=attention_config["use_experimental_scheduler"],
         ulysses_shards=attention_config["ulysses_shards"],
+        ulysses_attention_chunks=attention_config["ulysses_attention_chunks"],
     )
     # None axes corresponds to the stacked weights across all blocks
     # because of the use of nnx.vmap and nnx.scan.
diff --git a/src/maxdiffusion/models/wan/transformers/transformer_wan.py b/src/maxdiffusion/models/wan/transformers/transformer_wan.py
index 40c6be3f7..4cdfd0ca1 100644
--- a/src/maxdiffusion/models/wan/transformers/transformer_wan.py
+++ b/src/maxdiffusion/models/wan/transformers/transformer_wan.py
@@ -360,6 +360,7 @@ def __init__(
         "use_base2_exp": False,
         "use_experimental_scheduler": False,
         "ulysses_shards": -1,
+        "ulysses_attention_chunks": 1,
         **(attention_config or {}),
     }
 
@@ -584,6 +585,7 @@ def __init__(
         "use_base2_exp": False,
         "use_experimental_scheduler": False,
         "ulysses_shards": -1,
+        "ulysses_attention_chunks": 1,
         **(attention_config or {}),
     }
 
diff --git a/src/maxdiffusion/models/wan/transformers/transformer_wan_animate.py b/src/maxdiffusion/models/wan/transformers/transformer_wan_animate.py
index 400b967b8..91efde148 100644
--- a/src/maxdiffusion/models/wan/transformers/transformer_wan_animate.py
+++ b/src/maxdiffusion/models/wan/transformers/transformer_wan_animate.py
@@ -788,6 +788,7 @@ def __init__(
       enable_jax_named_scopes: bool = False,
       use_base2_exp: bool = False,
       use_experimental_scheduler: bool = False,
+      attention_config: Optional[dict] = None,
       face_flash_min_seq_length: int = 0,
       motion_encoder_channel_sizes: Optional[Dict[str, int]] = None,
       motion_encoder_size: int = 512,
@@ -812,6 +813,13 @@ def __init__(
     self.gradient_checkpoint = GradientCheckpointType.from_str(remat_policy)
     self.names_which_can_be_saved = names_which_can_be_saved or []
     self.names_which_can_be_offloaded = names_which_can_be_offloaded or []
+    attention_config = {
+        "use_base2_exp": use_base2_exp,
+        "use_experimental_scheduler": use_experimental_scheduler,
+        "ulysses_shards": -1,
+        "ulysses_attention_chunks": 1,
+        **(attention_config or {}),
+    }
 
     self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len)
 
@@ -903,8 +911,7 @@ def init_block(rngs):
           dropout=dropout,
           mask_padding_tokens=mask_padding_tokens,
           enable_jax_named_scopes=enable_jax_named_scopes,
-          use_base2_exp=use_base2_exp,
-          use_experimental_scheduler=use_experimental_scheduler,
+          attention_config=attention_config,
       )
 
     if scan_layers:
@@ -932,8 +939,7 @@ def init_block(rngs):
             dropout=dropout,
             mask_padding_tokens=mask_padding_tokens,
             enable_jax_named_scopes=enable_jax_named_scopes,
-            use_base2_exp=use_base2_exp,
-            use_experimental_scheduler=use_experimental_scheduler,
+            attention_config=attention_config,
         )
         blocks.append(block)
       self.blocks = nnx.List(blocks)
diff --git a/src/maxdiffusion/models/wan/transformers/transformer_wan_vace.py b/src/maxdiffusion/models/wan/transformers/transformer_wan_vace.py
index e9acbdb48..ca052282f 100644
--- a/src/maxdiffusion/models/wan/transformers/transformer_wan_vace.py
+++ b/src/maxdiffusion/models/wan/transformers/transformer_wan_vace.py
@@ -100,6 +100,7 @@ def __init__(
         "use_base2_exp": False,
         "use_experimental_scheduler": False,
         "ulysses_shards": -1,
+        "ulysses_attention_chunks": 1,
         **(attention_config or {}),
     }
 
@@ -348,6 +349,7 @@ def __init__(
         "use_base2_exp": False,
         "use_experimental_scheduler": False,
         "ulysses_shards": -1,
+        "ulysses_attention_chunks": 1,
         **(attention_config or {}),
     }
 
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline.py b/src/maxdiffusion/pipelines/wan/wan_pipeline.py
index 9ebfa68e9..4a82aaff3 100644
--- a/src/maxdiffusion/pipelines/wan/wan_pipeline.py
+++ b/src/maxdiffusion/pipelines/wan/wan_pipeline.py
@@ -346,6 +346,7 @@ def create_model(rngs: nnx.Rngs, wan_config: dict):
       "use_base2_exp": config.use_base2_exp,
       "use_experimental_scheduler": config.use_experimental_scheduler,
       "ulysses_shards": getattr(config, "ulysses_shards", -1),
+      "ulysses_attention_chunks": getattr(config, "ulysses_attention_chunks", 1),
   }
 
   # 2. eval_shape - will not use flops or create weights on device
@@ -581,7 +582,8 @@ def get_fp8_config(cls, config: HyperParameters):
     """
     fp8 config rules with per-tensor calibration.
     FLAX API (https://flax-linen.readthedocs.io/en/v0.10.6/guides/quantization/fp8_basics.html#flax-low-level-api):
-    The autodiff does not automatically use E5M2 for gradients and E4M3 for activations/weights during training, which is the recommended practice.
+    The autodiff does not automatically use E5M2 for gradients and E4M3 for
+    activations/weights during training, which is the recommended practice.
     """
     rules = [
         qwix.QtRule(
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py
index 54629d181..1ac30b221 100644
--- a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py
+++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py
@@ -12,7 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .wan_pipeline import WanPipeline, transformer_forward_pass, transformer_forward_pass_full_cfg, transformer_forward_pass_cfg_cache, init_magcache, magcache_step
+from .wan_pipeline import (
+    WanPipeline,
+    transformer_forward_pass,
+    transformer_forward_pass_full_cfg,
+    transformer_forward_pass_cfg_cache,
+    init_magcache,
+    magcache_step,
+)
 from ...models.wan.transformers.transformer_wan import WanModel
 from typing import List, Union, Optional
 from ...pyconfig import HyperParameters
@@ -290,15 +297,13 @@ def run_inference_2_1(
   transformer_obj = nnx.merge(graphdef, sharded_state, rest_of_state)
 
   # Compute RoPE once as it only depends on shape
-  dummy_hidden_states = jnp.zeros(
-      (
-          latents.shape[0],
-          latents.shape[2],
-          latents.shape[3],
-          latents.shape[4],
-          latents.shape[1],
-      )
-  )
+  dummy_hidden_states = jnp.zeros((
+      latents.shape[0],
+      latents.shape[2],
+      latents.shape[3],
+      latents.shape[4],
+      latents.shape[1],
+  ))
   rotary_emb = transformer_obj.rope(dummy_hidden_states)
 
   kv_cache = None
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py
index 8a63e8c3e..f01d0ea4d 100644
--- a/src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py
+++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py
@@ -91,6 +91,12 @@ def _create_model(rngs: nnx.Rngs, wan_config: dict):
   wan_config["enable_jax_named_scopes"] = config.enable_jax_named_scopes
   wan_config["use_base2_exp"] = config.use_base2_exp
   wan_config["use_experimental_scheduler"] = config.use_experimental_scheduler
+  wan_config["attention_config"] = {
+      "use_base2_exp": config.use_base2_exp,
+      "use_experimental_scheduler": config.use_experimental_scheduler,
+      "ulysses_shards": getattr(config, "ulysses_shards", -1),
+      "ulysses_attention_chunks": getattr(config, "ulysses_attention_chunks", 1),
+  }
 
   # 2. eval_shape – creates the model structure without allocating HBM.
   p_model_factory = partial(_create_model, wan_config=wan_config)
@@ -384,7 +390,8 @@ def check_inputs(
       )
     if negative_prompt is not None and negative_prompt_embeds is not None:
       raise ValueError(
-          f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to"
+          f"Cannot forward both `negative_prompt`: {negative_prompt} and"
+          f" `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to"
           " only forward one of the two."
       )
     if prompt is None and prompt_embeds is None:
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py
index f156fedd7..2e7543c9f 100644
--- a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py
+++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py
@@ -425,9 +425,10 @@ def run_inference_2_2_i2v(
   do_classifier_free_guidance = guidance_scale_low > 1.0 or guidance_scale_high > 1.0
   bsz = latents.shape[0]
 
-  prompt_embeds_combined = (
-      jnp.concatenate([prompt_embeds, negative_prompt_embeds], axis=0) if do_classifier_free_guidance else prompt_embeds
-  )
+  if do_classifier_free_guidance:
+    prompt_embeds_combined = jnp.concatenate([prompt_embeds, negative_prompt_embeds], axis=0)
+  else:
+    prompt_embeds_combined = prompt_embeds
   if image_embeds is not None:
     image_embeds_combined = (
         jnp.concatenate([image_embeds, image_embeds], axis=0) if do_classifier_free_guidance else image_embeds
@@ -971,19 +972,17 @@ def scan_body(carry, t):
     # tracing both 14B branches per step and keeps the AOT cache usable.
     use_high_noise = bool(np.asarray(scheduler_state.timesteps)[step] >= np.asarray(boundary))
     branch = high_noise_branch if use_high_noise else low_noise_branch
-    noise_pred, _ = branch(
-        (
-            latent_model_input,
-            timestep,
-            prompt_embeds_combined,
-            image_embeds_combined,
-            kv_cache_high,
-            kv_cache_low,
-            rotary_emb,
-            encoder_attention_mask_high,
-            encoder_attention_mask_low,
-        )
-    )
+    noise_pred, _ = branch((
+        latent_model_input,
+        timestep,
+        prompt_embeds_combined,
+        image_embeds_combined,
+        kv_cache_high,
+        kv_cache_low,
+        rotary_emb,
+        encoder_attention_mask_high,
+        encoder_attention_mask_low,
+    ))
     noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1))
     latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
 
diff --git a/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py b/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py
index 2f0adbf0f..e2c6a91d0 100644
--- a/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py
+++ b/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py
@@ -90,6 +90,7 @@ def create_model(rngs: nnx.Rngs, wan_config: dict):
       "use_base2_exp": config.use_base2_exp,
       "use_experimental_scheduler": config.use_experimental_scheduler,
       "ulysses_shards": getattr(config, "ulysses_shards", -1),
+      "ulysses_attention_chunks": getattr(config, "ulysses_attention_chunks", 1),
   }
 
   wan_config["scan_layers"] = False
@@ -414,7 +415,8 @@ def check_inputs(
       )
     elif negative_prompt is not None and negative_prompt_embeds is not None:
       raise ValueError(
-          f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to"
+          f"Cannot forward both `negative_prompt`: {negative_prompt} and"
+          f" `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to"
           " only forward one of the two."
       )
     elif prompt is None and prompt_embeds is None:
@@ -552,7 +554,8 @@ def __call__(
       if isinstance(conditioning_scale, jax.Array):
         if conditioning_scale.shape[0] != len(vace_layers):
           raise ValueError(
-              f"Length of `conditioning_scale` {conditioning_scale.shape[0]} does not match number of layers {len(vace_layers)}."
+              f"Length of `conditioning_scale` {conditioning_scale.shape[0]}"
+              f" does not match number of layers {len(vace_layers)}."
           )
 
       video, mask, reference_images = self.preprocess_conditions(
@@ -658,7 +661,8 @@ def prepare_video_latents(
     else:
       if video.shape[0] != len(reference_images):
         raise ValueError(
-            f"Batch size of `video` {video.shape[0]} and length of `reference_images` {len(reference_images)} does not match."
+            f"Batch size of `video` {video.shape[0]} and length of"
+            f" `reference_images` {len(reference_images)} does not match."
         )
 
     if video.shape[0] != 1:
diff --git a/src/maxdiffusion/tests/aot_cache_test.py b/src/maxdiffusion/tests/aot_cache_test.py
index bdfa0a589..2154b26ac 100644
--- a/src/maxdiffusion/tests/aot_cache_test.py
+++ b/src/maxdiffusion/tests/aot_cache_test.py
@@ -137,25 +137,23 @@ def test_signature_deterministic_across_processes(self):
     import subprocess
     import sys
 
-    snippet = "\n".join(
-        (
-            "import os",
-            "os.environ['JAX_PLATFORMS'] = 'cpu'",
-            "import jax",
-            "import jax.numpy as jnp",
-            "from flax import nnx",
-            "from maxdiffusion import aot_cache",
-            "",
-            "class T(nnx.Module):",
-            "  def __init__(self, rngs):",
-            "    self.lin = nnx.Linear(4, 4, rngs=rngs)",
-            "",
-            "graphdef, state = nnx.split(T(nnx.Rngs(0)))",
-            "sig = aot_cache._dynamic_signature(",
-            "    (graphdef, state.to_pure_dict(), jnp.ones((2, 4))), {})",
-            "print(sig)",
-        )
-    )
+    snippet = "\n".join((
+        "import os",
+        "os.environ['JAX_PLATFORMS'] = 'cpu'",
+        "import jax",
+        "import jax.numpy as jnp",
+        "from flax import nnx",
+        "from maxdiffusion import aot_cache",
+        "",
+        "class T(nnx.Module):",
+        "  def __init__(self, rngs):",
+        "    self.lin = nnx.Linear(4, 4, rngs=rngs)",
+        "",
+        "graphdef, state = nnx.split(T(nnx.Rngs(0)))",
+        "sig = aot_cache._dynamic_signature(",
+        "    (graphdef, state.to_pure_dict(), jnp.ones((2, 4))), {})",
+        "print(sig)",
+    ))
     outs = [
         subprocess.run(
             [sys.executable, "-c", snippet],
diff --git a/src/maxdiffusion/tests/attention_test.py b/src/maxdiffusion/tests/attention_test.py
index 708af4066..4421b2006 100644
--- a/src/maxdiffusion/tests/attention_test.py
+++ b/src/maxdiffusion/tests/attention_test.py
@@ -235,6 +235,26 @@ def test_select_flash_block_sizes_derives_cross_attn_defaults_for_tokamax(self):
     self.assertIsNone(cross_attention_block_sizes.block_kv_dq)
     self.assertTrue(cross_attention_block_sizes.use_fused_bwd_kernel)
 
+  def test_ulysses_head_chunk_ranges_preserve_head_layout_with_remainder(self):
+    ranges = attention_flax._ulysses_head_chunk_ranges(num_heads=40, ulysses_shards=8, num_chunks=2)
+
+    self.assertEqual(ranges, [(0, 16), (16, 40)])
+    self.assertEqual(
+        attention_flax._ulysses_head_chunk_ranges(num_heads=40, ulysses_shards=8, num_chunks=5),
+        [(0, 8), (8, 16), (16, 24), (24, 32), (32, 40)],
+    )
+    self.assertEqual(
+        attention_flax._ulysses_head_chunk_ranges(num_heads=40, ulysses_shards=8, num_chunks=3), [(0, 8), (8, 16), (16, 40)]
+    )
+    self.assertEqual(attention_flax._ulysses_head_chunk_ranges(num_heads=40, ulysses_shards=8, num_chunks=1), [(0, 40)])
+
+    head_major = jnp.arange(40 * 3, dtype=jnp.float32).reshape(40, 3)
+    reconstructed = jnp.concatenate((head_major[0:16], head_major[16:40]), axis=0)
+    self.assertTrue(jnp.array_equal(reconstructed, head_major))
+
+    ranges_array = jnp.array(ranges)
+    self.assertTrue(jnp.all((ranges_array[:, 1] - ranges_array[:, 0]) % 8 == 0))
+
   def test_ulysses_attention_round_trips_query_when_heads_are_divisible(self):
     """Ulysses attention should preserve the query layout after its collectives."""
     batch = 2
@@ -287,6 +307,65 @@ def fake_kernel(q, k, v, segment_ids):
     self.assertEqual(output.shape, query.shape)
     self.assertTrue(jnp.array_equal(output, query))
 
+  def test_ulysses_attention_chunk_counts_are_numerically_equivalent(self):
+    """Chunked all-to-all should preserve the same head/sequence layout as one-shot all-to-all."""
+    batch = 2
+    length = 6
+    heads = 8
+    head_depth = 3
+    query = jnp.arange(batch * length * heads * head_depth, dtype=jnp.float32).reshape(batch, length, heads * head_depth)
+    key = query + 1000.0
+    value = query + 2000.0
+    mesh = self._ulysses_mesh()
+
+    def fake_make_splash_mha(**unused_kwargs):
+      def fake_kernel(q, k, v, segment_ids):
+        del k, segment_ids
+        return q + jnp.mean(v, axis=1, keepdims=True)
+
+      return fake_kernel
+
+    def run_with_chunks(num_chunks):
+      with (
+          mesh,
+          nn_partitioning.axis_rules(self._ulysses_axis_rules()),
+          mock.patch.object(
+              attention_flax.splash_attention_kernel,
+              "make_splash_mha",
+              side_effect=fake_make_splash_mha,
+          ),
+      ):
+        return attention_flax._ulysses_attention(
+            query,
+            key,
+            value,
+            heads=heads,
+            mesh=mesh,
+            axis_names_q=(
+                attention_flax.BATCH,
+                attention_flax.SELF_ATTN_HEAD,
+                attention_flax.SELF_ATTN_Q_LENGTH,
+                attention_flax.D_KV,
+            ),
+            axis_names_kv=(
+                attention_flax.BATCH,
+                attention_flax.SELF_ATTN_HEAD,
+                attention_flax.SELF_ATTN_KV_LENGTH,
+                attention_flax.D_KV,
+            ),
+            flash_block_sizes=self._ulysses_block_sizes(),
+            dtype=jnp.float32,
+            ulysses_attention_chunks=num_chunks,
+        )
+
+    one_chunk = run_with_chunks(1)
+    two_chunks = run_with_chunks(2)
+    three_chunks_with_remainder = run_with_chunks(3)
+
+    self.assertEqual(one_chunk.shape, query.shape)
+    self.assertTrue(jnp.array_equal(one_chunk, two_chunks))
+    self.assertTrue(jnp.array_equal(one_chunk, three_chunks_with_remainder))
+
   def test_ulysses_attention_raises_when_heads_are_not_divisible_by_context_shards(self):
     """Ulysses attention should fail fast when heads cannot be evenly sharded."""
     batch = 2
@@ -508,6 +587,67 @@ def fake_kernel(q, k, v, segment_ids):
     self.assertEqual(output.shape, query.shape)
     self.assertTrue(jnp.array_equal(output, query))
 
+  @unittest.skipIf(len(jax.devices()) < 4, "Ulysses ring chunk equivalence test requires at least 4 devices.")
+  def test_ulysses_ring_attention_chunk_counts_are_numerically_equivalent(self):
+    """Chunked Ulysses+ring all-to-all should match the one-shot layout and numerics."""
+    batch = 2
+    length = 8
+    heads = 8
+    head_depth = 3
+    query = jnp.arange(batch * length * heads * head_depth, dtype=jnp.float32).reshape(batch, length, heads * head_depth)
+    key = query + 1000.0
+    value = query + 2000.0
+    mesh = self._ulysses_ring_mesh()
+
+    def fake_make_ring_attention(**unused_kwargs):
+      def fake_kernel(q, k, v, segment_ids):
+        del k, segment_ids
+        return q + jnp.mean(v, axis=1, keepdims=True)
+
+      return fake_kernel
+
+    def run_with_chunks(num_chunks):
+      with (
+          mesh,
+          nn_partitioning.axis_rules(self._ulysses_ring_axis_rules()),
+          mock.patch.object(
+              attention_flax.tokamax_ring_attention_kernel,
+              "make_ring_attention",
+              side_effect=fake_make_ring_attention,
+          ),
+      ):
+        return attention_flax._ulysses_ring_attention(
+            query,
+            key,
+            value,
+            heads=heads,
+            mesh=mesh,
+            axis_names_q=(
+                attention_flax.BATCH,
+                attention_flax.SELF_ATTN_HEAD,
+                attention_flax.SELF_ATTN_Q_LENGTH,
+                attention_flax.D_KV,
+            ),
+            axis_names_kv=(
+                attention_flax.BATCH,
+                attention_flax.SELF_ATTN_HEAD,
+                attention_flax.SELF_ATTN_KV_LENGTH,
+                attention_flax.D_KV,
+            ),
+            flash_block_sizes=self._ulysses_block_sizes(),
+            dtype=jnp.float32,
+            ulysses_shards=2,
+            ulysses_attention_chunks=num_chunks,
+        )
+
+    one_chunk = run_with_chunks(1)
+    two_chunks = run_with_chunks(2)
+    three_chunks_with_remainder = run_with_chunks(3)
+
+    self.assertEqual(one_chunk.shape, query.shape)
+    self.assertTrue(jnp.array_equal(one_chunk, two_chunks))
+    self.assertTrue(jnp.array_equal(one_chunk, three_chunks_with_remainder))
+
   @unittest.skipIf(len(jax.devices()) < 4, "Ulysses ring attention mask test requires at least 4 devices.")
   def test_ulysses_ring_attention_masks_global_kv_padding(self):
     """Hybrid Ulysses+ring masks padding via segment ids, not a NumpyMask."""