Feat: Optinal tile size search before e2e inference

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.

src/maxdiffusion/configs/base_wan_14b.yml

@@ -136,6 +136,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

@@ -110,6 +110,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

@@ -131,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
 
 # train text_encoder - Currently not supported for SDXL

src/maxdiffusion/configs/base_wan_animate.yml

@@ -125,6 +125,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

@@ -123,6 +123,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

@@ -124,6 +124,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
 

src/maxdiffusion/tests/tile_size_grid_search_test.py

@@ -0,0 +1,136 @@
+"""
+Copyright 2026 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import time
+import unittest
+
+from maxdiffusion.utils.tile_size_grid_search import (
+    MXU_TILE,
+    VPU_LANE,
+    BenchResult,
+    BlockBenchmark,
+    bkv_candidates,
+    bq_candidates,
+    grid_search,
+    padding_of,
+    smart_grid,
+    time_callable,
+    vmem_bkv_ceiling,
+)
+
+# per-shard seq the ring U=1 kernel tiles (75600 / 8); the empirical winner is bq=9472, bkv=1024.
+RING_SEQ = 9450
+VMEM_64MB = 64 * 1024 * 1024
+
+
+class _MockRingBench(BlockBenchmark):
+  """Synthetic benchmark encoding the measured ring behaviour: fewer Q-tiles is faster, the
+  bkv_compute sweet spot is ~1024, odd-128 blocks pay a half-MXU-pass penalty, and a score
+  tile that overflows VMEM OOMs.  Lets us test the orchestrator with no TPU."""
+
+  label = "mock-ring"
+
+  def __init__(self, seq=RING_SEQ, vmem=VMEM_64MB):
+    self.seq, self._vmem = seq, vmem
+
+  def tiled_seq_lens(self):
+    return (self.seq, self.seq)
+
+  def vmem_bytes(self):
+    return self._vmem
+
+  def run(self, bq, bkv, *, bkv_compute=None, iters=10, warmup=2):
+    cmp = bkv_compute or bkv
+    if bq * cmp * 4 + 15e6 > self._vmem:  # score tile f32 + ~15MB ring overhead
+      return BenchResult(bq, bkv, cmp, "oom")
+    n_q = padding_of(self.seq, bq).n_blocks
+    n_kv = padding_of(self.seq, bkv).n_blocks
+    ms = 60 + 4.0 * n_q + 0.9 * n_kv - 3.0 * min(cmp, 1024) / 1024
+    ms += 1.4 * (bq % MXU_TILE != 0) + 1.4 * (bkv % MXU_TILE != 0)
+    return BenchResult(bq, bkv, cmp, "ok", mean_ms=round(ms, 2), std_ms=0.2, compile_ms=25000.0)
+
+
+class PaddingMathTest(unittest.TestCase):
+
+  def test_padding_of(self):
+    p = padding_of(RING_SEQ, 1024)
+    self.assertEqual(p.n_blocks, 10)
+    self.assertEqual(p.padded_len, 10240)
+    self.assertEqual(p.pad, 790)
+
+  def test_single_tile_is_low_pad(self):
+    self.assertEqual(padding_of(RING_SEQ, 9472).pad, 22)  # 37 * 256
+
+
+class CandidateTest(unittest.TestCase):
+
+  def test_bq_fewest_tile_ladder(self):
+    # single-block ceiling fits, so the ladder starts at n=1 (bq=9472) and includes the winner.
+    bqs = bq_candidates(RING_SEQ, k=3, spread=0)
+    self.assertEqual(bqs[0], 9472)
+    self.assertTrue(all(b % VPU_LANE == 0 for b in bqs))
+
+  def test_bkv_largest_fits_includes_winner(self):
+    ceil = vmem_bkv_ceiling(9472, vmem_bytes=VMEM_64MB)
+    bkvs = bkv_candidates(RING_SEQ, k=3, max_block=ceil)
+    self.assertIn(1024, bkvs)  # the measured optimum, largest 256-mult that fits at bq=9472
+
+  def test_smart_grid_pairs_winner(self):
+    pairs = smart_grid(RING_SEQ, RING_SEQ, vmem_bytes=VMEM_64MB, dtype_bytes=4)
+    self.assertIn((9472, 1024), pairs)
+
+  def test_candidates_not_strictly_256(self):
+    # 128-multiples (e.g. 896) must be admissible, not filtered out.
+    bkvs = bkv_candidates(RING_SEQ, k=3, max_block=vmem_bkv_ceiling(9472, vmem_bytes=VMEM_64MB))
+    self.assertTrue(any(b % MXU_TILE != 0 for b in bkvs))
+
+
+class TimingTest(unittest.TestCase):
+
+  def test_compile_excluded_from_mean(self):
+    calls = {"n": 0}
+
+    def fake_fn():
+      calls["n"] += 1
+      time.sleep(0.20 if calls["n"] == 1 else 0.01)  # call #1 = "compile"
+      return calls["n"]
+
+    mean, _, times, compile_ms = time_callable(fake_fn, iters=5, warmup=2)
+    self.assertGreater(compile_ms, 150.0)  # the 200ms first call is captured here...
+    self.assertLess(mean, 30.0)  # ...and NOT in the steady-state mean
+    self.assertEqual(len(times), 5)
+
+
+class OrchestratorTest(unittest.TestCase):
+
+  def test_smart_search_picks_measured_winner(self):
+    res = grid_search(_MockRingBench(), mode="smart", iters=10, log=lambda *a, **k: None)
+    self.assertIsNotNone(res.best)
+    self.assertEqual((res.best.bq, res.best.bkv), (9472, 1024))
+
+  def test_oom_configs_pruned_not_raised(self):
+    # a tiny VMEM budget OOMs the big pairs; search must still return (or None), never raise.
+    res = grid_search(
+        _MockRingBench(vmem=8 * 1024 * 1024),
+        mode="smart",
+        iters=2,
+        log=lambda *a, **k: None,
+    )
+    self.assertTrue(any(r.status == "oom" for r in res.results))
+
+
+if __name__ == "__main__":
+  unittest.main()