Feat: Optinal tile size search before e2e inference

Author: eltsai

README.md

@@ -17,6 +17,8 @@
 [![Unit Tests](https://github.com/AI-Hypercomputer/maxdiffusion/actions/workflows/UnitTests.yml/badge.svg)](https://github.com/AI-Hypercomputer/maxdiffusion/actions/workflows/UnitTests.yml)
 
 # What's new?
+- **`2026/07/14`**: Automatic attention tile-size (`block_q`/`block_kv`) search for Wan is now supported.
+- **`2026/06/26`**: 2D ring (USP) attention with a custom splash kernel is now supported for Wan (`ulysses_ring_custom`), splitting context parallelism into an intra-chip Ulysses axis and a cross-chip ring axis.
 - **`2026/04/16`**: Support for Tokamax Ring Attention kernel is now added.
 - **`2026/03/31`**: Wan2.2 SenCache inference is now supported for T2V and I2V (up to 1.4x speedup)
 - **`2026/03/25`**: Wan2.1 and Wan2.2 Magcache inference is now supported
@@ -690,6 +692,10 @@ We added ring attention support for Wan models. Below are the stats for one `720
 
 (* There are some known stability issues for ring attention on 16 TPUs, please use `tokamax_flash` attention instead.)
 
+### Automatic Tile-Size Search
+
+The optimal attention tile sizes (`block_q` / `block_kv`) depend on the sequence length, VMEM, sharding, and accelerator, and the feasibility edge is a VMEM OOM with no clean closed form — so we tune them empirically. Passing `enable_tile_search=true` to `generate_wan.py` runs a fast one-DiT-block grid search before inference and injects the winning block sizes into `flash_block_sizes` (`tile_search_mode=smart` by default; the search is opt-in and off by default).  The core is model-agnostic (`utils/tile_size_grid_search.py`) with a per-model plug (`utils/wan_block_benchmark.py`), which also runs standalone via `python -m maxdiffusion.utils.wan_block_benchmark ... --smart-search`.
+
   ## Flux
 
   First make sure you have permissions to access the Flux repos in Huggingface.

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
 
@@ -136,6 +142,14 @@ flash_block_sizes: {
 #   "block_kv_dq" : 3072
 # }
 # GroupNorm groups
+# Tile-size auto-tuning. When enable_tile_search: True, generate_wan runs a fast one-DiT-block
+# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and overwrites
+# flash_block_sizes' block_q/block_kv/block_kv_compute with the winner. Default off (no-op).
+enable_tile_search: False
+tile_search_mode: 'smart'   # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep)
+tile_search_iters: 10
+tile_search_out: ''         # dir for the results CSV; '' -> print only
+
 norm_num_groups: 32
 
 # train text_encoder - Currently not supported for SDXL

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.
@@ -110,6 +116,14 @@ flash_block_sizes: {
   "use_fused_bwd_kernel": False,
 }
 # GroupNorm groups
+# Tile-size auto-tuning. When enable_tile_search: True, generate_wan runs a fast one-DiT-block
+# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and overwrites
+# flash_block_sizes' block_q/block_kv/block_kv_compute with the winner. Default off (no-op).
+enable_tile_search: False
+tile_search_mode: 'smart'   # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep)
+tile_search_iters: 10
+tile_search_out: ''         # dir for the results CSV; '' -> print only
+
 norm_num_groups: 32
 
 # train text_encoder - Currently not supported for SDXL

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
 
@@ -130,6 +136,14 @@ flash_block_sizes: {
 #   "block_kv_dq" : 2048
 #   "use_fused_bwd_kernel": False,
 # }
+# Tile-size auto-tuning. When enable_tile_search: True, generate_wan runs a fast one-DiT-block
+# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and overwrites
+# flash_block_sizes' block_q/block_kv/block_kv_compute with the winner. Default off (no-op).
+enable_tile_search: False
+tile_search_mode: 'smart'   # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep)
+tile_search_iters: 10
+tile_search_out: ''         # dir for the results CSV; '' -> print only
+
 # GroupNorm groups
 norm_num_groups: 32
 

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.
@@ -125,6 +131,14 @@ flash_block_sizes: {
 #   "use_fused_bwd_kernel": False,
 # }
 # GroupNorm groups
+# Tile-size auto-tuning. When enable_tile_search: True, generate_wan runs a fast one-DiT-block
+# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and overwrites
+# flash_block_sizes' block_q/block_kv/block_kv_compute with the winner. Default off (no-op).
+enable_tile_search: False
+tile_search_mode: 'smart'   # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep)
+tile_search_iters: 10
+tile_search_out: ''         # dir for the results CSV; '' -> print only
+
 norm_num_groups: 32
 
 # Text encoder training keys are unused by generate_wan_animate.py.

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
 
@@ -123,6 +129,14 @@ flash_block_sizes: {
 #   "use_fused_bwd_kernel": False,
 # }
 # GroupNorm groups
+# Tile-size auto-tuning. When enable_tile_search: True, generate_wan runs a fast one-DiT-block
+# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and overwrites
+# flash_block_sizes' block_q/block_kv/block_kv_compute with the winner. Default off (no-op).
+enable_tile_search: False
+tile_search_mode: 'smart'   # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep)
+tile_search_iters: 10
+tile_search_out: ''         # dir for the results CSV; '' -> print only
+
 norm_num_groups: 32
 
 # train text_encoder - Currently not supported for SDXL

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
 
@@ -124,6 +130,14 @@ flash_block_sizes: {
 #   "use_fused_bwd_kernel": False,
 # }
 # GroupNorm groups
+# Tile-size auto-tuning. When enable_tile_search: True, generate_wan runs a fast one-DiT-block
+# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and overwrites
+# flash_block_sizes' block_q/block_kv/block_kv_compute with the winner. Default off (no-op).
+enable_tile_search: False
+tile_search_mode: 'smart'   # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep)
+tile_search_iters: 10
+tile_search_out: ''         # dir for the results CSV; '' -> print only
+
 norm_num_groups: 32
 
 # train text_encoder - Currently not supported for SDXL

src/maxdiffusion/generate_wan.py

@@ -194,8 +194,50 @@ def inference_generate_video(config, pipeline, filename_prefix=""):
   return
 
 
+def maybe_tune_block_sizes(config):
+  """If enable_tile_search, run a fast one-DiT-block tile-size grid search and overwrite
+  flash_block_sizes' block_q/block_kv/block_kv_compute with the winner IN PLACE, before the
+  transformer (which bakes block sizes in at construction) is built.
+
+  Flags are read defensively so this is a safe no-op (grid search OFF) for any config that
+  doesn't declare them -- not every WAN yaml carries the tile_search_* keys."""
+  keys = config.get_keys()
+  if not keys.get("enable_tile_search", False):
+    return
+  from maxdiffusion.utils.tile_size_grid_search import grid_search
+  from maxdiffusion.utils.wan_block_benchmark import WanBlockBenchmark
+
+  mesh = jax.sharding.Mesh(max_utils.create_device_mesh(config), config.mesh_axes)
+  bench = WanBlockBenchmark.from_config(config, mesh)
+  max_logging.log(f"[tile-search] tuning block sizes for {bench.label} before inference...")
+  result = grid_search(
+      bench,
+      mode=keys.get("tile_search_mode", "smart"),
+      iters=keys.get("tile_search_iters", 10),
+      out_dir=(keys.get("tile_search_out", "") or None),
+      log=max_logging.log,
+  )
+  if result.best is None:
+    max_logging.log("[tile-search] no config succeeded; keeping configured flash_block_sizes")
+    return
+  fbs = dict(config.flash_block_sizes)
+  fbs.update({
+      "block_q": result.best.bq,
+      "block_kv": result.best.bkv,
+      "block_kv_compute": result.best.bkv_compute,
+      "block_kv_compute_in": result.best.bkv_compute,
+  })
+  config.get_keys()["flash_block_sizes"] = fbs  # config is immutable via setattr; mutate raw dict
+  max_logging.log(
+      f"[tile-search] using block_q={result.best.bq} block_kv={result.best.bkv} "
+      f"(block-bench {result.best.mean_ms:.2f} ms)"
+  )
+
+
 def run(config, pipeline=None, filename_prefix="", commit_hash=None):
   model_key = config.model_name
+  if pipeline is None:
+    maybe_tune_block_sizes(config)
   writer = max_utils.initialize_summary_writer(config)
   if jax.process_index() == 0 and writer:
     max_logging.log(f"TensorBoard logs will be written to: {config.tensorboard_dir}")
@@ -236,7 +278,10 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
       pipeline, _, _ = checkpoint_loader.load_checkpoint(checkpoint_step)
     else:
       pipeline = checkpoint_loader.load_pretrained_pipeline_or_diffusers(
-          config, pipeline_cls, pretrained_state_sources, pretrained_config_transformer_attr
+          config,
+          pipeline_cls,
+          pretrained_state_sources,
+          pretrained_config_transformer_attr,
       )
     load_time = time.perf_counter() - load_start
     max_logging.log(f"load_time: {load_time:.1f}s")
@@ -340,7 +385,11 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
     num_videos = num_devices * config.per_device_batch_size
     if num_videos > 0:
       generation_time_per_video = generation_time / num_videos
-      writer.add_scalar("inference/generation_time_per_video", generation_time_per_video, global_step=0)
+      writer.add_scalar(
+          "inference/generation_time_per_video",
+          generation_time_per_video,
+          global_step=0,
+      )
       max_logging.log(f"generation time per video: {generation_time_per_video}")
     else:
       max_logging.log("Warning: Number of videos is zero, cannot calculate generation_time_per_video.")
@@ -384,7 +433,11 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
     generation_time_with_profiler = time.perf_counter() - s0
     max_logging.log(f"generation_time_with_profiler: {generation_time_with_profiler}")
     if writer and jax.process_index() == 0:
-      writer.add_scalar("inference/generation_time_with_profiler", generation_time_with_profiler, global_step=0)
+      writer.add_scalar(
+          "inference/generation_time_with_profiler",
+          generation_time_with_profiler,
+          global_step=0,
+      )
 
   return saved_video_path