Feat: Optinal tile size search before e2e inference

[eltsai]

Automatic Tile Size Search

Adds a model-agnostic grid search that auto-tunes attention block_q/block_kv before inference. Attention tile sizes strongly affect throughput but the optimum depends on sequence length, VMEM, sharding, etc, for more details see go/diffusion-tile-size-search.

• utils/tile_size_grid_search.py — model-agnostic core.
• utils/wan_block_benchmark.py — WAN plug: times one real DiT block (num_layers=1), so results transfer to the full model at ~1/40th the cost.
• generate_wan.py — opt-in hook that runs the search and injects the winning tiles before the transformer is built.

The search maximizes tile size subject to VMEM (fewest-tile BQ ladder + largest-fits BKV), OOM-prunes, and picks the fastest measured config. Default OFF; safe no-op on configs that don't declare the flags.

How to Use It

Auto-tune inside inference (opt-in flags, default off):

enable_tile_search=true tile_search_mode=smart tile_search_iters=10

Standalone (writes the full grid to CSV):

python -m maxdiffusion.utils.wan_block_benchmark \
    --attention ulysses_ring_custom --ulysses-shards 1 --smart-search --out-dir <dir>

Example Output

[tile-search] q_seq=9450 kv_seq=9450 mode=smart -> 15 configs
  [1/15] bq=9472 bkv=1024:        55.23ms  (compile 10s, excluded)
  [2/15] bq=9472 bkv=768:         55.61ms
  ...
  [15/15] bq=768 bkv=9216:        81.51ms
[tile-search] WINNER: bq=9472 bkv=1024 bkv_compute=1024 -> 55.23 ms
[tile-search] using block_q=9472 block_kv=1024
...
  Denoise Total: 9.8s     # vs 12.9s at the untuned default

On 8×v7x WAN 2.2 (ulysses_ring_custom), auto-tuning cut denoise 12.9s → 9.8s.

Experiments

Benchmark / no-regression. Measured on WAN 2.2 (T2V-A14B), 1280×720, 81 frames, 8×v7x (ulysses_ring_custom, dp1-cp8, 5 steps): the search auto-selects block_q=9472, block_kv=1024 and cuts denoising from 12.9 s → 9.8 s vs the untuned default — matching the best hand-tuned config from a manual sweep, so it reproduces the manual optimum automatically. Because it only overwrites the block sizes when a measured config is faster (and otherwise keeps the configured default, with OOM configs pruned), it is best-performing by construction and cannot introduce a regression.

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

[eltsai]

Updated README to include automatic tile size search (and 2d ring attention)

[syhuang22]

[Please ignore and resolve this]

[syhuang22]

[Please ignore and resolve this]

[syhuang22]

In description maybe specify the config (resolution, num frames) to showcase the E2E time to ensure the result of this is best performing and have no regression.