diff --git a/README.md b/README.md
index 4f5a1fcb7..6aac966fc 100755
--- a/README.md
+++ b/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.
diff --git a/src/maxdiffusion/configs/base_wan_14b.yml b/src/maxdiffusion/configs/base_wan_14b.yml
index 493dfd6eb..a1801e55b 100644
--- a/src/maxdiffusion/configs/base_wan_14b.yml
+++ b/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
diff --git a/src/maxdiffusion/configs/base_wan_1_3b.yml b/src/maxdiffusion/configs/base_wan_1_3b.yml
index 994b1da36..13dd9a3f6 100644
--- a/src/maxdiffusion/configs/base_wan_1_3b.yml
+++ b/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
diff --git a/src/maxdiffusion/configs/base_wan_27b.yml b/src/maxdiffusion/configs/base_wan_27b.yml
index 0e919a479..0d5f52630 100644
--- a/src/maxdiffusion/configs/base_wan_27b.yml
+++ b/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
diff --git a/src/maxdiffusion/configs/base_wan_animate.yml b/src/maxdiffusion/configs/base_wan_animate.yml
index 35a063203..de8ed82da 100644
--- a/src/maxdiffusion/configs/base_wan_animate.yml
+++ b/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.
diff --git a/src/maxdiffusion/configs/base_wan_i2v_14b.yml b/src/maxdiffusion/configs/base_wan_i2v_14b.yml
index 3188e7849..94f96d351 100644
--- a/src/maxdiffusion/configs/base_wan_i2v_14b.yml
+++ b/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
diff --git a/src/maxdiffusion/configs/base_wan_i2v_27b.yml b/src/maxdiffusion/configs/base_wan_i2v_27b.yml
index 5f2e8c884..ea232f7d1 100644
--- a/src/maxdiffusion/configs/base_wan_i2v_27b.yml
+++ b/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
diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py
index a07fc5632..9a1904f33 100644
--- a/src/maxdiffusion/generate_wan.py
+++ b/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
 
diff --git a/src/maxdiffusion/tests/tile_size_grid_search_test.py b/src/maxdiffusion/tests/tile_size_grid_search_test.py
new file mode 100644
index 000000000..3a985f072
--- /dev/null
+++ b/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()
diff --git a/src/maxdiffusion/utils/tile_size_grid_search.py b/src/maxdiffusion/utils/tile_size_grid_search.py
new file mode 100644
index 000000000..38e9d2a49
--- /dev/null
+++ b/src/maxdiffusion/utils/tile_size_grid_search.py
@@ -0,0 +1,453 @@
+"""
+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.
+
+Tile-size (block_q / block_kv) grid search for flash/splash/ring attention kernels.
+
+Two hardware granularities drive the candidate math (do NOT conflate them):
+  * VPU_LANE = 128 — the VMEM/vector lane width and the kernel's HARD floor: every
+    block size must be a multiple of 128 or the splash kernel raises.
+  * MXU_TILE = 256 — the systolic matmul array is 256x256 on v7x.  A block
+    dimension that is an *odd* multiple of 128 (e.g. 896 = 3.5x256) leaves the last
+    MXU pass half-empty.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional
+import csv
+import os
+import sys
+
+# --- the two granularities (see module docstring) --------------------------------
+VPU_LANE = 128  # kernel hard floor: block sizes must be multiples of this
+MXU_TILE = 256  # 256x256 MXU: multiples of this fully pack the systolic array
+
+
+def _ceil_div(a: int, b: int) -> int:
+  return (a + b - 1) // b
+
+
+def _ceil_to(x: int, m: int) -> int:
+  return _ceil_div(x, m) * m
+
+
+def _floor_to(x: int, m: int) -> int:
+  return (x // m) * m
+
+
+@dataclass(frozen=True)
+class Padding:
+  seq_len: int
+  block: int
+  n_blocks: int
+  padded_len: int
+  pad: int
+  pad_pct: float
+
+
+def padding_of(seq_len: int, block: int) -> Padding:
+  """How `block` tiles `seq_len`: block count, padded length, and wasted rows."""
+  n = _ceil_div(seq_len, block)
+  padded = n * block
+  return Padding(
+      seq_len,
+      block,
+      n,
+      padded,
+      padded - seq_len,
+      100.0 * (padded - seq_len) / seq_len,
+  )
+
+
+def block_for_count(seq_len: int, n: int, align: int = MXU_TILE) -> Optional[int]:
+  """Smallest multiple of `align` that tiles `seq_len` into EXACTLY n blocks, else None.
+
+  (Rounding seq/n up to `align` can bump the block big enough to yield n-1 blocks; such
+  counts have no clean aligned block and are skipped.)"""
+  b = max(align, _ceil_to(_ceil_div(seq_len, n), align))
+  return b if _ceil_div(seq_len, b) == n else None
+
+
+def full_axis_candidates(
+    seq_len: int,
+    *,
+    step: int = MXU_TILE,
+    min_block: int = MXU_TILE,
+    max_block: Optional[int] = None,
+) -> list[int]:
+  """Full sweep: [min_block, min_block+step, ..., max_block].  `max_block` defaults
+  to the single-block ceiling (round_up(seq_len, step)); nothing bigger helps.  Default
+  step=256 keeps the count sane (step=128 ~doubles it; use only for a fine characterization).
+  """
+  if step % VPU_LANE:
+    raise ValueError(f"step must be a multiple of {VPU_LANE}, got {step}")
+  max_block = max_block or _ceil_to(seq_len, step)
+  return list(range(min_block, max_block + 1, step))
+
+
+def vmem_bq_ceiling(
+    bkv: int,
+    *,
+    vmem_bytes: int,
+    dtype_bytes: int = 4,
+    score_fraction: float = 0.65,
+    align: int = VPU_LANE,
+) -> int:
+  """Approx largest bq whose score tile [bq, bkv]*dtype_bytes fits VMEM (mirror of
+  `vmem_bkv_ceiling`).  Use a SMALL reference bkv so the fewest-tile BQ ladder isn't
+  over-constrained; big-bq x big-bkv corners are OOM-pruned by the orchestrator."""
+  budget = int(vmem_bytes * score_fraction)
+  return max(align, _floor_to(budget // (bkv * dtype_bytes), align))
+
+
+def bq_candidates(
+    seq_len: int,
+    *,
+    k: int = 3,
+    spread: int = 2,
+    align: int = VPU_LANE,
+    max_block: Optional[int] = None,
+    min_block: int = VPU_LANE,
+) -> list[int]:
+  """BQ: VMEM-capped fewest-tile ladder + geometric spread-down.
+
+  `max_block` = the BQ VMEM ceiling (e.g. `vmem_bq_ceiling(min_bkv)`); defaults to the
+  single-block size.  Takes the `k` fewest-tile blocks <= max_block (fewer Q-tiles win),
+  then `spread` progressively-halved blocks (snapped to the min-padding aligned block for
+  that tile count) so a MODERATE optimum -- e.g. ulysses bq~5888, well below single-tile --
+  is also sampled.  Returned high->low, deduped.  align=128 (256-mult preferred, not
+  required)."""
+  cap = _floor_to(min(max_block or _ceil_to(seq_len, align), _ceil_to(seq_len, align)), align)
+  max_n = _ceil_div(seq_len, min_block)
+  out: list[int] = []
+  n = _ceil_div(seq_len, cap)  # fewest tiles that fit the cap
+  while len(out) < k and n <= max_n:
+    b = block_for_count(seq_len, n, align)
+    if b is not None and min_block <= b <= cap and b not in out:
+      out.append(b)
+    n += 1
+  b = out[-1] if out else cap
+  for _ in range(spread):  # geometric spread toward moderate blocks
+    b = _floor_to(b // 2, align)
+    if b < min_block:
+      break
+    snapped = block_for_count(seq_len, _ceil_div(seq_len, b), align) or b
+    if min_block <= snapped <= cap and snapped not in out:
+      out.append(snapped)
+  return sorted(set(out), reverse=True)
+
+
+def vmem_bkv_ceiling(
+    bq: int,
+    *,
+    vmem_bytes: int,
+    dtype_bytes: int = 4,
+    score_fraction: float = 0.65,
+    align: int = MXU_TILE,
+) -> int:
+  """Approx largest bkv_compute whose score tile [bq, bkv]*dtype_bytes fits the fraction
+  of VMEM left after the kernel's other resident tiles (ring fp32 residual windows, K/V,
+  Q, accumulators).  APPROXIMATE — OOM-pruning in the orchestrator is the real guard.
+  Default 0.65 fits the ulysses_ring_custom kernel (measured: bkv 1152 fits, 1280 OOMs
+  at bq=9472 / 64 MB)."""
+  budget = int(vmem_bytes * score_fraction)
+  return max(align, _floor_to(budget // (bq * dtype_bytes), align))
+
+
+def bkv_candidates(seq_len: int, *, k: int = 3, align: int = VPU_LANE, max_block: int) -> list[int]:
+  """BKV(=bkv_compute): the MXU compute tile.  Largest tiles that fit VMEM, descending by
+  `align` from `max_block` (a VMEM ceiling, e.g. `vmem_bkv_ceiling`).  Capped at the
+  single-block size (no point tiling KV bigger than the sequence).  align=128 by default
+  so a 128-multiple that fits (e.g. 1152) is considered alongside 256-multiples (1024).
+  """
+  top = min(_floor_to(max_block, align), _ceil_to(seq_len, align))
+  out: list[int] = []
+  b = top
+  while b >= align and len(out) < k:
+    out.append(b)
+    b -= align
+  return out
+
+
+# ================================================================================
+# Per-model plug: a BlockBenchmark builds a ONE-block model and times a
+# forward for a given (bq, bkv).  Only this layer is model-specific; add a model by
+# implementing it (WanBlockBenchmark below is the reference).
+# ================================================================================
+@dataclass
+class BenchResult:
+  bq: int
+  bkv: int
+  bkv_compute: int
+  status: str  # "ok" | "oom" | "error"
+  mean_ms: Optional[float] = None  # steady-state, COMPILE + WARMUP EXCLUDED
+  std_ms: Optional[float] = None
+  times_ms: list[float] = field(default_factory=list)
+  compile_ms: Optional[float] = None  # first-call (compile+first-exec) wall time, reported separately
+  detail: str = ""
+
+  def csv_row(self) -> dict:
+    return {
+        "bq": self.bq,
+        "bkv": self.bkv,
+        "bkv_compute": self.bkv_compute,
+        "status": self.status,
+        "mean_ms": self.mean_ms,
+        "std_ms": self.std_ms,
+        "compile_ms": self.compile_ms,
+        "detail": self.detail,
+    }
+
+
+def time_callable(fn, *, iters: int = 10, warmup: int = 2, sync=lambda x: x):
+  """Correct microbenchmark that ALWAYS excludes compilation and warmup from mean_ms.
+
+  Call #1 is executed untimed and absorbs the JIT compile / first-touch cost (returned
+  separately as compile_ms).  `warmup-1` further untimed calls reach steady state.  ONLY
+  the subsequent `iters` calls are timed -> `mean_ms` never contains compile or warmup.
+  `sync(result)` must block until the async result is materialised (jax.block_until_ready);
+  it defaults to identity so this stays jax-free and unit-testable.
+
+  Returns (mean_ms, std_ms, times_ms, compile_ms)."""
+  import statistics
+  import time
+
+  t0 = time.perf_counter()
+  sync(fn())  # call #1: compilation happens HERE, untimed
+  compile_ms = (time.perf_counter() - t0) * 1e3
+  for _ in range(max(0, warmup - 1)):  # extra warmups -> steady state, untimed
+    sync(fn())
+  times: list[float] = []
+  for _ in range(iters):  # the ONLY timed calls
+    t = time.perf_counter()
+    sync(fn())
+    times.append((time.perf_counter() - t) * 1e3)
+  mean = sum(times) / len(times)
+  std = statistics.pstdev(times) if len(times) > 1 else 0.0
+  return mean, std, times, compile_ms
+
+
+class BlockBenchmark:
+  """Interface a model implements so the grid search can drive it.  `run` must build (or
+  reuse) a single-block model with the given block sizes, execute a forward `warmup`+`iters`
+  times, and return timings.  Catch out-of-VMEM and return status="oom" (don't raise).
+  """
+
+  label: str = "block"
+
+  def tiled_seq_lens(self) -> tuple[int, int]:
+    """(q_seq, kv_seq) the kernel actually TILES — per-shard, variant-aware (e.g.
+    full_seq*U/CP for ring, full_seq for pure ulysses).  The candidate math runs on this.
+    """
+    raise NotImplementedError
+
+  def vmem_bytes(self) -> int:
+    raise NotImplementedError
+
+  def dtype_bytes(self) -> int:
+    return 2  # bf16 q/k/v; score tile is f32 (4) -> handled by score_fraction tuning
+
+  def run(
+      self,
+      bq: int,
+      bkv: int,
+      *,
+      bkv_compute: Optional[int] = None,
+      iters: int = 10,
+      warmup: int = 2,
+  ) -> BenchResult:
+    """Build/reuse the 1-block model with these block sizes and time a forward.  MUST use
+    `time_callable` (or equivalent) so `mean_ms` EXCLUDES compilation + warmup; report the
+    one-time compile cost in `compile_ms`.  Catch out-of-VMEM -> status='oom' (don't raise).
+    """
+    raise NotImplementedError
+
+
+# ================================================================================
+# Orchestrator
+# ================================================================================
+@dataclass
+class SearchResult:
+  best: Optional[BenchResult]
+  results: list[BenchResult]
+  q_seq: int
+  kv_seq: int
+  mode: str
+
+
+def smart_grid(
+    q_seq: int,
+    kv_seq: int,
+    *,
+    vmem_bytes: int,
+    dtype_bytes: int = 4,
+    k_bq: int = 3,
+    k_bkv: int = 3,
+    spread_bq: int = 2,
+    score_fraction: float = 0.65,
+    min_bkv_ref: int = 1024,
+) -> list[tuple[int, int]]:
+  """Nested candidate pairs: BQ = VMEM-capped fewest-tile ladder + spread; then for EACH bq,
+  BKV = largest-that-fits at that bq (so bkv is VMEM-correct for its partner, not globally).
+  Pairs whose score tile still overflows are OOM-pruned at run time.  cmp is locked = bkv.
+
+  `min_bkv_ref` sets the BQ VMEM cap via the bkv it is expected to pair with.  Use a REALISTIC
+  bkv (1024, a good MXU tile) rather than the smallest possible: with a tiny ref the cap is huge,
+  so the fewest-tile ladder starts at the single-tile end (which OOMs for a large per-shard seq)
+  and the feasible moderate-BQ optimum (e.g. bq=9472 at seq 37800) falls in the ladder's gap.
+  """
+  bq_cap = vmem_bq_ceiling(
+      min_bkv_ref,
+      vmem_bytes=vmem_bytes,
+      dtype_bytes=dtype_bytes,
+      score_fraction=score_fraction,
+  )
+  bqs = bq_candidates(q_seq, k=k_bq, spread=spread_bq, max_block=bq_cap)
+  pairs: list[tuple[int, int]] = []
+  for bq in bqs:
+    bkv_cap = vmem_bkv_ceiling(
+        bq,
+        vmem_bytes=vmem_bytes,
+        dtype_bytes=dtype_bytes,
+        score_fraction=score_fraction,
+    )
+    for bkv in bkv_candidates(kv_seq, k=k_bkv, max_block=bkv_cap):
+      pairs.append((bq, bkv))
+  return pairs
+
+
+def full_grid(q_seq: int, kv_seq: int, *, step: int = MXU_TILE, max_configs: Optional[int] = None) -> list[tuple[int, int]]:
+  """Mode-1 full 2D sweep: every (bq, bkv) in the step-`step` product.  WARNING: this is
+  O(N^2) in seq/step (per-shard 9450 -> ~1.4k combos; 75600 -> ~88k).  `max_configs` caps it.
+  """
+  bqs = full_axis_candidates(q_seq, step=step)
+  bkvs = full_axis_candidates(kv_seq, step=step)
+  pairs = [(bq, bkv) for bq in bqs for bkv in bkvs]
+  if max_configs and len(pairs) > max_configs:
+    pairs = pairs[:max_configs]
+  return pairs
+
+
+def grid_search(
+    bench: BlockBenchmark,
+    *,
+    mode: str = "smart",
+    out_dir: Optional[str] = None,
+    iters: int = 10,
+    warmup: int = 2,
+    k: int = 3,
+    step: int = MXU_TILE,
+    max_configs: Optional[int] = None,
+    log=print,
+) -> SearchResult:
+  """Run the tile-size grid on `bench`, write CSV to `out_dir` (or pretty-print), return the
+  winner (lowest mean_ms among status=='ok').  `mode`: 'smart' (candidate ladders) | 'full'.
+  """
+  q_seq, kv_seq = bench.tiled_seq_lens()
+  if mode == "smart":
+    pairs = smart_grid(q_seq, kv_seq, vmem_bytes=bench.vmem_bytes(), dtype_bytes=4, k_bq=k, k_bkv=k)
+  elif mode == "full":
+    log(
+        "Warning: tile_search mode is 'full', not 'smart' -- this is an exhaustive O(N^2) 2D BQ x BKV"
+        " sweep (often hundreds to thousands of configs, each separately compiled) and is meant for"
+        " one-off characterization; use mode='smart' for routine tuning."
+    )
+    pairs = full_grid(q_seq, kv_seq, step=step, max_configs=max_configs)
+  else:
+    raise ValueError(f"mode must be 'smart' or 'full', got {mode!r}")
+
+  log(f"[tile-search] {bench.label}: q_seq={q_seq} kv_seq={kv_seq} mode={mode} " f"-> {len(pairs)} configs (iters={iters})")
+  results: list[BenchResult] = []
+  for i, (bq, bkv) in enumerate(pairs, 1):
+    r = bench.run(bq, bkv, bkv_compute=bkv, iters=iters, warmup=warmup)
+    results.append(r)
+    tag = "" if bq % MXU_TILE == 0 and bkv % MXU_TILE == 0 else " [½MXU]"
+    compile_note = f"  (compile {r.compile_ms/1e3:.0f}s, excluded)" if r.compile_ms else ""
+    log(
+        f"  [{i}/{len(pairs)}] bq={bq} bkv={bkv}{tag}: "
+        + (f"{r.mean_ms:.2f}ms{compile_note}" if r.status == "ok" else r.status)
+    )
+
+  ok = [r for r in results if r.status == "ok" and r.mean_ms is not None]
+  best = min(ok, key=lambda r: r.mean_ms) if ok else None
+  _emit(results, best, q_seq, kv_seq, mode, out_dir, log)
+  return SearchResult(best, results, q_seq, kv_seq, mode)
+
+
+def _emit(results, best, q_seq, kv_seq, mode, out_dir, log) -> None:
+  if out_dir:
+    os.makedirs(out_dir, exist_ok=True)
+    path = os.path.join(out_dir, "tile_size_grid_search.csv")
+    with open(path, "w", newline="") as f:
+      w = csv.DictWriter(
+          f,
+          fieldnames=[
+              "bq",
+              "bkv",
+              "bkv_compute",
+              "status",
+              "mean_ms",
+              "std_ms",
+              "compile_ms",
+              "detail",
+          ],
+      )
+      w.writeheader()
+      for r in sorted(results, key=lambda r: (r.mean_ms is None, r.mean_ms or 0)):
+        w.writerow(r.csv_row())
+    log(f"[tile-search] wrote {path}")
+  else:
+    log(f"[tile-search] results (q_seq={q_seq}, kv_seq={kv_seq}, mode={mode}):")
+    for r in sorted(results, key=lambda r: (r.mean_ms is None, r.mean_ms or 0)):
+      log(
+          f"    bq={r.bq:>6} bkv={r.bkv:>5} cmp={r.bkv_compute:>5}  "
+          + (f"{r.mean_ms:7.2f} ms  (±{r.std_ms:.2f})" if r.status == "ok" else f"  {r.status}")
+      )
+  if best:
+    log(f"[tile-search] WINNER: bq={best.bq} bkv={best.bkv} bkv_compute={best.bkv_compute} " f"-> {best.mean_ms:.2f} ms")
+  else:
+    log("[tile-search] no config succeeded (all OOM/error)")
+
+
+# --------------------------------------------------------------------------------
+# tiny self-demo: `python -m maxdiffusion.utils.tile_size_grid_search 9450`
+# --------------------------------------------------------------------------------
+def _mxu_tag(b: int) -> str:
+  return f"{b // MXU_TILE}x256 packed" if b % MXU_TILE == 0 else f"{b // VPU_LANE}x128 (½ MXU pass)"
+
+
+def _demo(seq_len: int, _bq: int, vmem_mb: int) -> None:
+  vmem = vmem_mb * 1024 * 1024
+  print(f"seq_len (per-shard tiled) = {seq_len}   VPU_LANE={VPU_LANE}  MXU_TILE={MXU_TILE}  vmem={vmem_mb}MB")
+  pairs = smart_grid(seq_len, seq_len, vmem_bytes=vmem, dtype_bytes=4)
+  print(f"\nSMART grid: {len(pairs)} (bq, bkv=cmp) pairs (bkv is largest-fits PER bq):")
+  last_bq = None
+  for bq, bkv in pairs:
+    if bq != last_bq:
+      pq = padding_of(seq_len, bq)
+      print(f"  bq={bq:>6} ({pq.n_blocks} Q-tile(s), pad {pq.pad_pct:.1f}%, {_mxu_tag(bq)}):")
+      last_bq = bq
+    pk = padding_of(seq_len, bkv)
+    print(f"      bkv={bkv:>5}  {pk.n_blocks} tile(s)  {_mxu_tag(bkv)}  score[{bq},{bkv}]f32={bq*bkv*4/1e6:.0f}MB")
+  full = full_axis_candidates(seq_len)
+  print(f"\nFULL sweep (step 256): {len(full)}x{len(full)} = {len(full)**2} combos (mode='full')")
+
+
+if __name__ == "__main__":
+  _seq = int(sys.argv[1]) if len(sys.argv) > 1 else 9450
+  _bq = int(sys.argv[2]) if len(sys.argv) > 2 else 9472
+  _vm = int(sys.argv[3]) if len(sys.argv) > 3 else 64
+  _demo(_seq, _bq, _vm)
diff --git a/src/maxdiffusion/utils/wan_block_benchmark.py b/src/maxdiffusion/utils/wan_block_benchmark.py
new file mode 100644
index 000000000..a945be8cd
--- /dev/null
+++ b/src/maxdiffusion/utils/wan_block_benchmark.py
@@ -0,0 +1,331 @@
+"""
+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.
+
+WAN implementation of the model-agnostic `BlockBenchmark` (see tile_size_grid_search.py).
+
+Benchmarks ONE real WanTransformerBlock: it builds a real WanModel with num_layers=1 and
+runs its actual forward (patch-embed -> rope -> 1 DiT block -> proj_out), so compilation and
+the attention/sharding path match generate_wan.py exactly.  Only the (block_q, block_kv)
+change between runs; params are random (we time latency, not correctness).
+
+Usage:
+  # standalone
+  python -m maxdiffusion.utils.wan_block_benchmark --attention ulysses_ring_custom \
+      --ulysses-shards 1 --smart-search --out-dir /tmp/tile_search
+  # in generate_wan.py: WanBlockBenchmark.from_config(config, mesh); grid_search(bench, ...)
+"""
+
+import os
+
+# Perf flags mirror the wan22 run script; setdefault so we NEVER clobber flags a caller
+# (e.g. generate_wan.py) already set before its own TPU init.  Only takes effect for the
+# standalone CLI, where this module is imported before jax.
+os.environ.setdefault(
+    "LIBTPU_INIT_ARGS",
+    " ".join([
+        "--xla_tpu_dvfs_p_state=7",
+        "--xla_tpu_spmd_rng_bit_generator_unsafe=true",
+        "--xla_tpu_enable_dot_strength_reduction=true",
+        "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true",
+        "--xla_enable_async_collective_permute=true",
+        "--xla_tpu_enable_async_collective_fusion=true",
+        "--xla_tpu_enable_async_collective_fusion_multiple_steps=true",
+        "--xla_tpu_overlap_compute_collective_tc=true",
+        "--xla_enable_async_all_gather=true",
+        "--xla_tpu_scoped_vmem_limit_kib=65536",
+        "--xla_tpu_enable_async_all_to_all=true",
+        "--xla_tpu_enable_all_experimental_scheduler_features=true",
+        "--xla_tpu_enable_latency_hiding_scheduler=true",
+        "--xla_tpu_enable_megacore_fusion=true",
+    ]),
+)
+os.environ.setdefault("JAX_DEFAULT_MATMUL_PRECISION", "bfloat16")
+os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.95")
+
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+from flax import nnx
+from flax.linen import partitioning as nn_partitioning
+
+from maxdiffusion import max_logging, max_utils, pyconfig
+from maxdiffusion.models.wan.transformers.transformer_wan import WanModel
+from maxdiffusion.utils.tile_size_grid_search import (
+    BenchResult,
+    BlockBenchmark,
+    grid_search,
+    time_callable,
+)
+
+# WAN2.1/2.2 latent geometry (matches bench_block_worker / generate_wan defaults).
+_IN_CHANNELS = 16
+_VAE_T, _VAE_S = 4, 8  # VAE temporal / spatial compression
+_PATCH_T, _PATCH_H, _PATCH_W = 1, 2, 2
+_RING_VARIANTS = {
+    "ulysses_ring_custom",
+    "ulysses_ring_custom_bidir",
+    "tokamax_ring",
+    "ring",
+}
+
+
+def latent_seq_len(num_frames: int, height: int, width: int) -> int:
+  """Full (unsharded) patch-token sequence length the transformer sees."""
+  lf = (num_frames - 1) // _VAE_T + 1
+  lh, lw = height // _VAE_S, width // _VAE_S
+  return (lf // _PATCH_T) * (lh // _PATCH_H) * (lw // _PATCH_W)
+
+
+def tiled_seq_len(full_seq: int, attention: str, context_shards: int, ulysses_shards: int) -> int:
+  """Per-shard sequence the kernel actually TILES (what the bq/bkv candidate math runs on).
+
+  Ring shards the sequence, so each ring step tiles full_seq / R (R = CP / U).  Pure ulysses
+  gathers the full sequence per device (heads sharded), and non-CP flash sees the full seq.
+  """
+  if attention in _RING_VARIANTS and ulysses_shards >= 1 and context_shards >= 1:
+    ring_shards = max(1, context_shards // max(1, ulysses_shards))
+    return full_seq // ring_shards
+  return full_seq
+
+
+@nnx.jit
+def _forward(model, latents, timestep, prompt_embeds):
+  return model(
+      hidden_states=latents,
+      timestep=timestep,
+      encoder_hidden_states=prompt_embeds,
+      encoder_hidden_states_image=None,
+      deterministic=True,
+      rngs=None,
+  )
+
+
+class WanBlockBenchmark(BlockBenchmark):
+  """BlockBenchmark for a single WAN DiT block.  One-time setup (config, mesh, HF transformer
+  config, dummy inputs) in __init__; each `run` rebuilds the 1-layer model with the given block
+  sizes (they are baked in at construction) and times a forward with compile excluded.
+  """
+
+  def __init__(
+      self,
+      config,
+      mesh,
+      *,
+      num_frames,
+      height,
+      width,
+      batch=None,
+      vmem_limit_bytes=64 * 1024 * 1024,
+  ):
+    self._config = config
+    self._mesh = mesh
+    self._rules = config.logical_axis_rules
+    self._attention = config.attention
+    self._context_shards = int(mesh.shape.get("context", 1))
+    self._ulysses_shards = int(getattr(config, "ulysses_shards", 1) or 1)
+    self._vmem = int(vmem_limit_bytes)
+    self.label = f"wan/{self._attention}/u{self._ulysses_shards}"
+
+    lf = (num_frames - 1) // _VAE_T + 1
+    self._lf, self._lh, self._lw = lf, height // _VAE_S, width // _VAE_S
+    self._full_seq = latent_seq_len(num_frames, height, width)
+    # The 'batch' logical axis maps to ('data', 'fsdp'), so the dummy batch must be divisible
+    # by data*fsdp -- default to exactly that (e.g. 2 on dp2) so batch=1 doesn't fail the ring
+    # shard_map on dp>1 meshes.
+    data_shards = int(mesh.shape.get("data", 1)) * int(mesh.shape.get("fsdp", 1))
+    self._batch = batch if batch is not None else max(1, data_shards)
+    # HF transformer config fetched once, reused for every rebuild.
+    self._hf_cfg = WanModel.load_config(config.pretrained_model_name_or_path, subfolder="transformer")
+    self._inputs = self._make_inputs()
+
+  @classmethod
+  def from_config(cls, config, mesh, **overrides):
+    return cls(
+        config,
+        mesh,
+        num_frames=config.num_frames,
+        height=config.height,
+        width=config.width,
+        **overrides,
+    )
+
+  # --- BlockBenchmark interface ---------------------------------------------------
+  def tiled_seq_lens(self):
+    s = tiled_seq_len(self._full_seq, self._attention, self._context_shards, self._ulysses_shards)
+    return (s, s)
+
+  def vmem_bytes(self):
+    return self._vmem
+
+  def run(self, bq, bkv, *, bkv_compute=None, iters=10, warmup=2):
+    cmp = bkv_compute or bkv
+    try:
+      with self._mesh:
+        model = self._build_model(bq, bkv, cmp)
+      latents, timestep, prompt_embeds = self._inputs
+      with self._mesh, nn_partitioning.axis_rules(self._rules):
+        mean, std, times, compile_ms = time_callable(
+            lambda: _forward(model, latents, timestep, prompt_embeds),
+            iters=iters,
+            warmup=warmup,
+            sync=jax.block_until_ready,
+        )
+      return BenchResult(
+          bq,
+          bkv,
+          cmp,
+          "ok",
+          mean_ms=mean,
+          std_ms=std,
+          times_ms=times,
+          compile_ms=compile_ms,
+      )
+    except Exception as e:  # pylint: disable=broad-except
+      msg = str(e)
+      oom = any(t in msg for t in ("RESOURCE_EXHAUSTED", "out of memory", "Mosaic", "VMEM"))
+      return BenchResult(bq, bkv, cmp, "oom" if oom else "error", detail=msg[:200])
+
+  # --- WAN model build (adapted from scripts/bench_block_worker.py) ----------------
+  def _flash_block_sizes(self, bq, bkv, cmp):
+    from maxdiffusion.max_utils import CustomFlashBlockSizes
+
+    return CustomFlashBlockSizes(
+        block_q=bq,
+        block_kv=bkv,
+        block_kv_compute=cmp,
+        block_kv_compute_in=cmp,
+        heads_per_tile=1,
+        vmem_limit_bytes=self._vmem,
+    )
+
+  def _build_model(self, bq, bkv, cmp):
+    c = self._config
+    wan_config = dict(self._hf_cfg)
+    wan_config.update(
+        mesh=self._mesh,
+        dtype=c.activations_dtype,
+        weights_dtype=c.weights_dtype,
+        attention=c.attention,
+        precision=max_utils.get_precision(c),
+        flash_block_sizes=self._flash_block_sizes(bq, bkv, cmp),
+        remat_policy=c.remat_policy,
+        names_which_can_be_saved=c.names_which_can_be_saved,
+        names_which_can_be_offloaded=c.names_which_can_be_offloaded,
+        flash_min_seq_length=c.flash_min_seq_length,
+        dropout=c.dropout,
+        mask_padding_tokens=c.mask_padding_tokens,
+        scan_layers=False,
+        num_layers=1,
+        enable_jax_named_scopes=c.enable_jax_named_scopes,
+        attention_config={
+            "use_base2_exp": c.use_base2_exp,
+            "use_experimental_scheduler": c.use_experimental_scheduler,
+            "ulysses_shards": c.ulysses_shards,
+        },
+    )
+    model = WanModel(**wan_config, rngs=nnx.Rngs(params=0))
+    gd, state, rest = nnx.split(model, nnx.Param, ...)
+    shardings = nn.logical_to_mesh_sharding(nnx.get_partition_spec(state), self._mesh, self._rules)
+    return nnx.merge(gd, jax.device_put(state, shardings), rest)
+
+  def _make_inputs(self):
+    k1, k2 = jax.random.split(jax.random.key(0), 2)
+    latents = jax.random.normal(k1, (self._batch, _IN_CHANNELS, self._lf, self._lh, self._lw), jnp.float32)
+    prompt_embeds = jax.random.normal(k2, (self._batch, 512, 4096), jnp.float32)
+    timestep = jnp.zeros((self._batch,), jnp.int32)
+    repl = jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec())
+    return (
+        jax.device_put(latents, repl),
+        jax.device_put(timestep, repl),
+        jax.device_put(prompt_embeds, repl),
+    )
+
+
+# ================================================================================
+# CLI: standalone tile-size search for WAN
+# ================================================================================
+def _build_cli_config(argv_yaml, attention, ulysses_shards, num_frames, height, width):
+  pyconfig.initialize([
+      "wan_block_benchmark",
+      argv_yaml,
+      f"attention={attention}",
+      f"ulysses_shards={ulysses_shards}",
+      "skip_jax_distributed_system=True",
+      "weights_dtype=bfloat16",
+      "activations_dtype=bfloat16",
+      "flash_min_seq_length=0",
+      "per_device_batch_size=1",
+      f"num_frames={num_frames}",
+      f"height={height}",
+      f"width={width}",
+      "ici_data_parallelism=1",
+      "ici_fsdp_parallelism=1",
+      f"ici_context_parallelism={jax.device_count()}",
+      "ici_tensor_parallelism=1",
+      "allow_split_physical_axes=True",
+      "use_base2_exp=true",
+      "use_experimental_scheduler=true",
+      "enable_jax_named_scopes=false",
+  ])
+  return pyconfig.config
+
+
+def main():
+  import argparse
+
+  p = argparse.ArgumentParser(description="WAN tile-size (block_q/block_kv) grid search")
+  p.add_argument("--config-yml", default="src/maxdiffusion/configs/base_wan_27b.yml")
+  p.add_argument("--attention", default="ulysses_ring_custom")
+  p.add_argument("--ulysses-shards", type=int, default=1)
+  p.add_argument("--num-frames", type=int, default=81)
+  p.add_argument("--height", type=int, default=720)
+  p.add_argument("--width", type=int, default=1280)
+  p.add_argument("--mode", choices=["smart", "full"], default="smart")
+  p.add_argument("--smart-search", action="store_true", help="alias for --mode smart")
+  p.add_argument("--k", type=int, default=3)
+  p.add_argument("--iters", type=int, default=10)
+  p.add_argument("--warmup", type=int, default=2)
+  p.add_argument("--step", type=int, default=256, help="full-mode step")
+  p.add_argument("--max-configs", type=int, default=None)
+  p.add_argument("--vmem-mb", type=int, default=64)
+  p.add_argument("--out-dir", default=None)
+  args = p.parse_args()
+
+  config = _build_cli_config(
+      args.config_yml,
+      args.attention,
+      args.ulysses_shards,
+      args.num_frames,
+      args.height,
+      args.width,
+  )
+  mesh = jax.sharding.Mesh(max_utils.create_device_mesh(config), config.mesh_axes)
+  bench = WanBlockBenchmark.from_config(config, mesh, vmem_limit_bytes=args.vmem_mb * 1024 * 1024)
+  mode = "smart" if args.smart_search else args.mode
+  grid_search(
+      bench,
+      mode=mode,
+      out_dir=args.out_dir,
+      iters=args.iters,
+      warmup=args.warmup,
+      k=args.k,
+      step=args.step,
+      max_configs=args.max_configs,
+      log=max_logging.log,
+  )
+
+
+if __name__ == "__main__":
+  main()