feat(wan): fused RMSNorm+RoPE Pallas producer; Wan switches in config, not env by Perseus14 · Pull Request #488 · AI-Hypercomputer/maxdiffusion

Perseus14 · 2026-09-22T20:09:29Z

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 · 2026-09-22T20:09:46Z

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.

…, not env Fused producer (kernels/fused_rmsnorm_rope_pallas.py) - A Pallas kernel computing FP32 RMSNorm + RoPE + the [B,S,H*D] -> [B,H,S,D] head transpose for self-attention Q/K. It optionally folds log2(e) into Q and 1/sqrt(d) into K. - norm_mode="exact" leaves the feature reduction to XLA. It is 0 ULP against the separately jitted XLA producer (fused_rmsnorm_rope). - norm_mode="fused" does the reduction in-kernel. It stays within 2x pair-relative bf16 eps. - rope_accum ("dtype" | "f32") matches how the local compiler rounds the RoPE multiply-add. "auto" uses the per-platform measured mode: tpu7x "dtype", v6e "f32". - Off by default in the YAML (`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 output is equivalent but not identical (v6e 720p/81f, same seed: 53.8 dB PSNR after 1 step, 34.1 dB after 40). The kernel is refused on TPUs whose rounding was not measured, and when the sequence split is uneven. - run_wan_fast_inference.sh turns it ON in the v6e and v7 profiles despite the YAML default, because it is measurably faster there. Measured at the stack tip, i.e. with #491 on top (Wan 2.2 T2V-A14B, 720p/81f/40 steps, warm AOT, DVFS unpinned), kernel on vs off: denoise 125.9s vs 131.2s on v6e-8, 105.1s vs 106.7s on tpu7x-8. Not measured with this PR alone. Full-run output vs the XLA producer: on v6e the bf16 trajectory diverges (~15 dB PSNR against the XLA-path video, visually clean); on tpu7x it was bit-identical. USE_FUSED_ROPE_KERNEL=false restores the XLA producer. - Isolated speed, one v6e chip, jax 0.11.2, 40 heads, d=128, bf16, exact mode, block_s=1024, median of 50 runs: per-shard seq 9450 takes 0.964 ms vs 1.789 ms for XLA (1.86x); seq 18900 takes 1.953 ms vs 4.018 ms (2.06x). This replaces the unbenchmarked "1.9x" and "2.5-3.0x" figures. - VMEM budget: 64 MiB is requested only when every TPU in the mesh is a v6e or tpu7x (the validated generations). Any other TPU gets Mosaic's default scoped limit. The mesh is now passed through to the resolver. - Two fused RMSNorm+RoPE producers now exist: the XLA one from feat/fused-rmsnorm-rope-xla and this Pallas one. The XLA one is the reference and the fallback. Wan graph-changing switches - 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 YAML keys, declared in every base_wan*.yml with the code default. The pipelines resolve them once, at model-build time, with `resolve_from_config`: the config value if the key is set, else the legacy WAN_* env var, else the default. Since every Wan YAML sets every key, the env vars are effectively ignored on normal runs (a differing one is logged once). - WanPipeline, VACE, Animate and wan_block_benchmark all pass them in `attention_config` (`attention_config_entries(config)`). - FlaxWanAttention and WanModel built without an entry use the built-in default and ignore the environment. - So every value lives on the GraphDef, and therefore in the AOT key. - There is no process-global store. Nothing reads these switches at trace time, and the generic Ulysses / ring wrappers take `transpose_out: bool = False`. A Wan setting can no longer leak into another model in the same process. pyconfig only coerces the values (e.g. "false" from the CLI). - wan_cfg_before_unpatchify (default True) applies CFG on packed tokens before unpatchify. Wan 2.1 no-cache CFG now goes through the same transformer_forward_pass path instead of transformer_forward_pass_full_cfg. A jitted test through a real WanModel (p_t = 1 and 2) shows both are bit-identical to the CFG-after-unpatchify path and to full_cfg in f32 everywhere, and in bf16 on CPU and v6e. In bf16 on tpu7x, XLA fuses the CFG combine differently, so ~30% of elements differ by one bf16 ULP of the output scale; the test allows that there. - splash `transpose_out`: an in-kernel output layout [H,S,D] instead of [H,D,S], for the custom splash kernel and the non-ring Ulysses wrapper. Off by default. Other changes - split_head_dim is now plumbed from the config into WanModel and its attention layers. The six base_wan*.yml files flip `split_head_dim` from True to False. Before this change no Wan code read that key, so the effective behaviour (False) is unchanged. - _apply_attention_dot: handles prescaled Q/K (`qk_prescaled` / `k_prescaled` / `use_base2_exp` undo the folded scale and LOG2E on the scores, or on K/Q before the memory-efficient kernel), and the memory-efficient path keeps [B, S, H, D] under split_head_dim instead of applying the batched-heads transpose. (The 4-D [B,H,S,D] layout fix lives in feat/attn-dot-fallback-fixes below.) Prescaled Q/K is refused for unsupported kernels and for head-local SVG attention. - `_fused_rope_producer` always returns ((q, k), qk_prescaled), falling back to XLA when the kernel cannot run. Tests - fused_rmsnorm_rope_pallas_test.py: - parity grid, guards, backward, production shape (TPU); - VMEM resolver; - FlaxWanAttention sequence-sharded over `context` via logical axis rules, against the XLA producer: bit-identical on CPU, within 2e-2 on TPU where the inlined XLA producer rounds differently (dot_product kernel, so the prescale fold is not exercised there); - transpose_out for the plain, fixed-m and k-centred kernels, MHPT (TPU), the Ulysses wrapper and the Ulysses+Ring wrapper at R == 1. - TPU tests that rely on the 64 MiB VMEM budget or a measured rope_accum (production-shape parity, in-register prescale, fused-mode drift, the mesh test's TPU branch) skip on TPUs outside the validated kinds (`vmem_limit_is_validated`, `rope_accum_is_measured`). - dot_fallback_layout_test.py: 4-D layout, prescale, cross-attention K prescale via config, and the compiled CFG parity test. - wan_runtime_options_test.py: YAML/default agreement, precedence, coercion, no global store, and module-level resolution (FlaxWanAttention, WanModel). - aot_cache_test.py: Wan switches and the resolved rope_accum key the executable. - Counts: fused_rmsnorm_rope_pallas_test.py 38 test methods (60 parameterized cases), dot_fallback_layout_test.py 12, wan_runtime_options_test.py 11. - CI: 48 cases skip in GitHub Actions: 47 of the 60 in fused_rmsnorm_rope_pallas_test (only the guard and VMEM-resolver classes run) and the compiled CFG parity test. wan_runtime_options, aot_cache and the dot-fallback layout tests run in CI. - end_to_end/tpu/run_wan_stack_tests.sh: adds this PR's two new test files. Verified (final tree): - TPU v6e-8 (jax 0.11.2, end_to_end/tpu/run_wan_stack_tests.sh): 283 passed, 168 subtests passed in 575.2s.

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.

mbohlool · 2026-10-16T15:07:28Z

+def resolve_from_config(config: Any, name: str) -> Any:
+  """Coerced value of `name`: `config` first, then the legacy env var, then the default."""
+  env, dflt, _ = _OPTIONS[name]
+  keys = _extract_keys(config)
+  legacy = os.environ.get(env)
+  if name in keys and keys[name] is not None:
+    value = _coerce(name, keys[name])
+    if legacy is not None and env not in _WARNED_ENV:
+      try:
+        legacy_coerced = _coerce(name, legacy)
+      except ValueError:
+        legacy_coerced = legacy
+      if legacy_coerced != value:
+        _WARNED_ENV.add(env)
+        max_logging.log(
+            f"[wan_runtime_options] Ignoring legacy env {env}={legacy!r}: config sets {name}={value!r}. "
+            f"Set {name} in the YAML or on the command line instead."
+        )
+    return value
+  return _coerce(name, legacy) if legacy is not None else dflt


Could we drop the legacy WAN_* env tier and resolve config, then default, only? None of these variables exist on main or in #477/#478/#479, so there is nothing to stay compatible with, and today a stale shell variable silently changes the compiled graph whenever the config lacks a key. Please re-point the tests and remove the "used to be WAN_* environment variables" YAML comments and the WAN_ROPE_ACCUM mention in resolve_rope_accum's docstring.

mbohlool · 2026-10-16T15:07:28Z

+#                    Its output is equivalent but not bit-identical to the XLA
+#                    producer: on v6e the bf16 trajectory diverges (PSNR ~15 dB
+#                    against the XLA-path video, visually clean); on tpu7x the


Which stack and seed are the ~15 dB here and the 34.1 dB in the fused_rmsnorm_rope_pallas.py docstring (L48-49) each measured on? The v6e profile turns fused RoPE on for about 4% denoise time, so a VBench or CLIP score would help justify that default, or the kernel could default off on v6e.

mbohlool · 2026-10-16T15:07:28Z

+IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true"
+# Kernel numerics grids and multi-device tests are skipped in CI; see
+# end_to_end/tpu/run_wan_stack_tests.sh.
+_SKIP_IN_GITHUB_ACTIONS = unittest.skipIf(
+    IN_GITHUB_ACTIONS, "TPU kernel / multi-device test, skipped in GitHub Actions; run end_to_end/tpu/run_wan_stack_tests.sh"


_SKIP_IN_GITHUB_ACTIONS also skips the tiny interpret-mode parity, pair-swap and backward tests, so CI never checks the kernel numerics. Could we gate only the TPU-shape and multi-device classes (e.g. on _on_tpu() / jax.device_count()) and let the rest run, and fail loudly, in CPU CI?

mbohlool · 2026-10-16T15:07:28Z

+      can_prescale = attn_kernel in supported_custom_kernels and not use_mem_eff and q.shape[1] >= min_seq_len
+      norm_mode_env = self.wan_rope_norm_mode
+      fuse_qk_prescale = bool(self.wan_fuse_qk_prescale and can_prescale)
+      q_prescale_val = float(LOG2E) if (fuse_qk_prescale and getattr(self.attention_op, "use_base2_exp", True)) else 1.0
+      k_prescale_val = float(self.attention_op.scale) if fuse_qk_prescale else 1.0


Could we add a test for the qk_prescaled producer to Ulysses hand-off? The existing mesh test uses dot_product below flash_min_seq_length, so can_prescale is False there. Mocking _ulysses_attention and _ulysses_ring_custom_attention, running FlaxWanAttention(attention_kernel="ulysses_custom", flash_min_seq_length=0, use_fused_rope_kernel=True) on the fake-TPU mesh and asserting qk_prescaled plus k == xla_k * bf16(scale) would catch a dropped not qk_prescaled guard, which today would silently double-scale.

mbohlool · 2026-10-16T15:07:28Z

+
+    self.use_fused_rope_kernel = attention_config["use_fused_rope_kernel"]
+    self.fused_rope_block_s = attention_config["fused_rope_block_s"]
+    self.fused_rope_head_block = attention_config["fused_rope_head_block"]


Could we check here in __init__ that fused_rope_head_block divides heads and raise ValueError? Today a value that doesn't divide it (e.g. 16 with Wan-14B's 40 heads) only raises inside shard_map at first trace, after the weights are loaded.

mbohlool · 2026-10-16T15:07:28Z

+    if k_prescaled or qk_prescaled:
+      key_states = key_states / jnp.asarray(scale, dtype=key_states.dtype)
+    if qk_prescaled and use_base2_exp:
+      query_states = query_states / jnp.asarray(LOG2E, dtype=query_states.dtype)
+    if not split_head_dim:
+      query_states = query_states.transpose(1, 0, 2)
+      key_states = key_states.transpose(1, 0, 2)
+      value_states = value_states.transpose(1, 0, 2)


These _apply_attention_dot edits also run on the SDXL/Flux dot_product path (the memory-efficient branch now handles split_head_dim=True differently). Could we add an SDXL-shaped parity test for this branch and call the change out in the PR description?

Perseus14 requested a review from entrpn as a code owner September 22, 2026 20:09

Perseus14 force-pushed the feat/wan-custom-kernels branch from dbd6ace to 7a6e0e6 Compare September 22, 2026 20:10

Perseus14 added this pull request to stack #486 September 22, 2026 20:11

gemini-code-assist Bot reviewed Sep 22, 2026

View reviewed changes

Perseus14 force-pushed the feat/wan-custom-kernels branch 9 times, most recently from 5b75348 to 82681ec Compare September 23, 2026 16:48

Perseus14 force-pushed the feat/wan-custom-kernels branch from 82681ec to b33e35e Compare September 23, 2026 19:19

Perseus14 force-pushed the feat/wan-custom-kernels branch from b33e35e to 7740dcf Compare September 23, 2026 20:01

eltsai mentioned this pull request Sep 23, 2026

feat(wan): fast serving with persistent AOT caching and tuned inference recipe #479

Open

eltsai reviewed Sep 23, 2026

View reviewed changes

Comment thread src/maxdiffusion/tests/aot_cache_test.py

eltsai reviewed Sep 24, 2026

View reviewed changes

Comment thread src/maxdiffusion/kernels/fused_rmsnorm_rope_pallas.py Outdated

Perseus14 force-pushed the feat/wan-custom-kernels branch from 7740dcf to 5c0f8bb Compare September 24, 2026 07:00

Perseus14 force-pushed the feat/wan-custom-kernels branch from 5c0f8bb to aee8c4a Compare September 24, 2026 07:50

Perseus14 mentioned this pull request Sep 24, 2026

feat(attention): fixed-m splash attention kernel with dynamic bounds and safety fallbacks #477

Open

Perseus14 force-pushed the feat/wan-custom-kernels branch 2 times, most recently from 2f8b92d to a507d0d Compare September 24, 2026 17:13

Perseus14 force-pushed the feat/wan-custom-kernels branch from 6ac9049 to 96dd50f Compare September 26, 2026 19:43

Perseus14 force-pushed the feat/wan-custom-kernels branch from 96dd50f to 359249b Compare September 26, 2026 23:03

Perseus14 force-pushed the feat/wan-custom-kernels branch 2 times, most recently from 593e337 to 4ebc6fd Compare September 27, 2026 09:50

Perseus14 force-pushed the feat/wan-custom-kernels branch from 4ebc6fd to 5a1fa32 Compare September 27, 2026 14:59

Perseus14 force-pushed the feat/wan-custom-kernels branch 2 times, most recently from a75047a to 0b1315d Compare September 27, 2026 16:25

Perseus14 force-pushed the feat/wan-custom-kernels branch from 0b1315d to 77afa86 Compare September 28, 2026 16:32

Perseus14 force-pushed the feat/wan-custom-kernels branch from 77afa86 to 2e70bc8 Compare September 28, 2026 16:50

Perseus14 force-pushed the feat/wan-custom-kernels branch from 2e70bc8 to 973594c Compare September 29, 2026 19:22

Perseus14 mentioned this pull request Oct 1, 2026

feat(attention): 2D Ulysses+Ring custom attention with fixed-m (pre-a2a norm reduction, unpadded ring K/V) #478

Open

eltsai previously approved these changes Oct 5, 2026

View reviewed changes

mbohlool approved these changes Oct 19, 2026

View reviewed changes

Navigation Menu

Repository navigation

feat(wan): fused RMSNorm+RoPE Pallas producer; Wan switches in config, not env - #488

feat(wan): fused RMSNorm+RoPE Pallas producer; Wan switches in config, not env#488
Perseus14 wants to merge 1 commit into
feat/wan-fast-servingfrom
feat/wan-custom-kernels

Perseus14 commented Sep 22, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Sep 22, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mbohlool left a comment

Uh oh!

mbohlool Oct 16, 2026

Uh oh!

mbohlool Oct 16, 2026

Uh oh!

mbohlool Oct 16, 2026

Uh oh!

mbohlool Oct 16, 2026

Uh oh!

mbohlool Oct 16, 2026

Uh oh!

mbohlool Oct 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Perseus14 commented Sep 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What's in it

Tests

Uh oh!

github-actions Bot commented Sep 22, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mbohlool left a comment

Choose a reason for hiding this comment

Uh oh!

mbohlool Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

mbohlool Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

mbohlool Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

mbohlool Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

mbohlool Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

mbohlool Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Perseus14 commented Sep 22, 2026 •

edited

Loading