[Perseus14]

Summary

Adds a fused RMSNorm + RoPE + head-transpose Pallas producer for Wan self-attention Q/K, and moves Wan's graph-changing switches out of environment variables into config.

What's in it

• Pallas producer (kernels/fused_rmsnorm_rope_pallas.py): computes FP32 RMSNorm + RoPE + the [B,S,H*D] → [B,H,S,D] transpose, with optional log2(e) / 1/√d prescale folding.
  • norm_mode="exact" is 0 ULP against the separately jitted XLA producer (fused_rmsnorm_rope from feat/fused-rmsnorm-rope-xla).
  • In isolation on v6e it is 1.86×–2.06× faster than XLA (0.964 vs 1.789 ms at seq 9450; 1.953 vs 4.018 ms at seq 18900).
• Defaults:

  • Off by default in the yml (use_fused_rope_kernel: False) and in the launcher's generic profile: once inlined into the 40-layer graph, XLA rounds the unfused producer differently, so the 40-step v6e trajectory diverges (~15 dB PSNR against the XLA-path video, visually clean; bit-identical on tpu7x).

  • Turned on in the launcher's v6e and v7 profiles. End-to-end latency (Wan 2.2 T2V-A14B, 720p / 81 frames / 40 steps, CP=4, DP=2, warm AOT, DVFS unpinned):

    TPU / Revision	Generate, on	Denoise, on	Generate, off	Denoise, off	Delta (on vs off)
    v6e-8 (at feat(wan): fused RMSNorm+RoPE Pallas producer; Wan switches in config, not env #488 alone)	133.7 s	131.6 s	138.2 s	136.0 s	−4.5 s generate / −4.4 s denoise (−3.2%)
    v6e-8 (stack tip perf(wan): pad sequence to 128-lane multiple for self-attention (wan_seq_pad="lane") #491)	128.1 s	125.9 s	133.4 s	131.2 s	−5.3 s generate / −5.3 s denoise (−4.0%)
    tpu7x-8 (stack tip perf(wan): pad sequence to 128-lane multiple for self-attention (wan_seq_pad="lane") #491)	107.6 s	105.1 s	109.2 s	106.7 s	−1.6 s generate / −1.6 s denoise (−1.5%)

  • Set USE_FUSED_ROPE_KERNEL=false to turn it off.

• Config, not env:
  • wan_rope_norm_mode, wan_fuse_qk_prescale, wan_splash_transpose_out, wan_cfg_before_unpatchify, wan_cross_attn_prescale_kv and wan_rope_accum are yml keys passed via attention_config_entries(config) across WanPipeline, VACE, Animate, and wan_block_benchmark.
  • No process-global store and no trace-time reads. Config-less module builds use built-in defaults.
• wan_cfg_before_unpatchify (on by default): applies CFG before unpatchify. A compiled parity test shows bit-identical output in f32 everywhere and in bf16 on CPU/v6e, and within 1 bf16 ULP on tpu7x.
• VMEM: 64 MiB is requested only when every mesh TPU is v6e or tpu7x; other TPUs get Mosaic's default.

Tests

fused_rmsnorm_rope_pallas_test.py (38 test methods / 60 parameterized cases), wan_runtime_options_test.py (11), plus updates to dot_fallback_layout_test.py (12) and aot_cache_test.py (45).

• TPU tests relying on 64 MiB VMEM or measured rope_accum skip on unvalidated TPUs (vmem_limit_is_validated, rope_accum_is_measured).
• CI vs Local VM: 48 additional kernel-grid/CFG tests skip under GITHUB_ACTIONS=true (88 of 283 cumulative); guards, VMEM resolver, runtime options, AOT cache, and dot-fallback layout tests run in CI.
• Results (TPU v6e-8): 283 passed cumulative (run_wan_stack_tests.sh, 575.2 s).

Stack: #477 → #507 → #478 → #508 → #479 → #488 → #509 → #510 → #511 → #491

[github-actions]

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

[gemini-code-assist]

Code Review

This pull request introduces a fused RMSNorm + RoPE + head-transposition Pallas kernel for TPU, integrates it into the Flax attention pipeline, and implements several optimizations such as in-kernel VMEM transpose and Classifier-Free Guidance (CFG) before unpatchify. Feedback on these changes highlights critical runtime issues, including an AttributeError from using jax.shard_map instead of jax.experimental.shard_map.shard_map and a potential NameError due to an undefined _warn_once helper. Additionally, the reviewer recommends adding defensive checks to validate the sequence length of freqs_cis and removing redundant JAX compatibility shims that are already handled in the package initialization.

[mbohlool]

Approving once the inline items are addressed; the Pallas producer matches the unfused reference op for op, and tiling, tails and VMEM look right.