From a0cf7f279024f88f21b9c55b09e024deb1e068eb Mon Sep 17 00:00:00 2001
From: Sagar Chapara <sagarchapara@google.com>
Date: Tue, 14 Jul 2026 17:12:29 +0000
Subject: [PATCH] feat: configure chunked Ulysses all-to-all

Add a ulysses_attention_chunks attention config to split the Ulysses all-to-all into head-group passes. The chunked path lets XLA overlap all-to-all collectives with head-parallel local attention compute while preserving the existing single-shot path by default.

Apply the same chunking to plain Ulysses and Ulysses+Ring, and allow the final chunk to carry the remainder when the requested chunk count does not divide the Ulysses head groups evenly.

Add mocked attention tests for numerical and layout equivalence across chunk counts.
---
 README.md                                     |  25 ++++
 pyproject.toml                                |   8 +-
 src/maxdiffusion/configs/base_wan_14b.yml     |   6 +
 src/maxdiffusion/configs/base_wan_1_3b.yml    |   6 +
 src/maxdiffusion/configs/base_wan_27b.yml     |   6 +
 src/maxdiffusion/configs/base_wan_animate.yml |   6 +
 src/maxdiffusion/configs/base_wan_i2v_14b.yml |   6 +
 src/maxdiffusion/configs/base_wan_i2v_27b.yml |   6 +
 src/maxdiffusion/generate_wan.py              |  29 ++--
 src/maxdiffusion/models/attention_flax.py     | 139 ++++++++++++++++-
 .../wan/transformers/transformer_wan.py       |   2 +
 .../transformers/transformer_wan_animate.py   |  14 +-
 .../wan/transformers/transformer_wan_vace.py  |   2 +
 .../pipelines/wan/wan_pipeline.py             |   4 +-
 .../pipelines/wan/wan_pipeline_2_1.py         |  25 ++--
 .../pipelines/wan/wan_pipeline_animate.py     |   9 +-
 .../pipelines/wan/wan_pipeline_i2v_2p2.py     |  31 ++--
 .../pipelines/wan/wan_vace_pipeline_2_1.py    |  10 +-
 src/maxdiffusion/tests/aot_cache_test.py      |  36 +++--
 src/maxdiffusion/tests/attention_test.py      | 140 ++++++++++++++++++
 20 files changed, 433 insertions(+), 77 deletions(-)

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