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.

Author: csgoogle

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.

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

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
 

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.

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
 

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.

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
 

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
 

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