perf(wan): XLA-fused RMSNorm + RoPE producer for self-attention Q/K

Adds `fused_rmsnorm_rope` (kernels/fused_producers.py) and uses it in
`FlaxWanAttention.__call__` for the self-attention Q/K producer chain.

What
- `fused_rmsnorm_rope(raw_q, raw_k, q_scale, k_scale, freqs_cis, ...)` does
  FP32 RMSNorm, casts back to the input dtype, then applies RoPE and the
  head transposition in input dtype, so no FP32 RoPE intermediates are
  materialised at long sequence lengths. It mirrors nnx.RMSNorm's association
  (`x * (rsqrt(var + eps) * scale)`): bit-for-bit in f32, and in bf16 when the
  norm scale is bf16. With bf16 activations and an f32 scale, nnx.RMSNorm's
  `promote_dtype` rounds the scale to bf16 before folding it into the rsqrt,
  while the fused producer keeps the scale in f32 and rounds only the final
  product; the two then differ by at most one bf16 ulp. It takes separate q/k
  epsilons, accepts an optional `kv_heads` for GQA (defaults to `q_heads`), and
  validates `freqs_cis.shape[2] == Sq == Sk`.
- It is `jax.jit`-wrapped with heads/dim/eps static: under the pipeline's
  outer jit it is simply inlined, but when a full-size block is called
  eagerly (as wan_vace_transformer_test does at 75,600 tokens) XLA fuses the
  chain instead of materialising ~1.5 GB of extra intermediates, which OOMed
  that test on the v4-8 CI runner.
- `FlaxWanAttention` takes the fused path when `rotary_emb` is given,
  `qk_norm` is on and the call is self-attention, reading `self.norm_q.epsilon`
  and `self.norm_k.epsilon` directly; cross-attention and the no-RoPE path keep
  the separate norm_q / norm_k / `_apply_rope` steps.

Tests
- fused_producers_test.py (10, new, CPU): bit-for-bit match with
  nnx.RMSNorm in f32 and in bf16 (bf16 scale), the bf16-activation /
  f32-scale case pinned to <= 1 bf16 ulp of nnx.RMSNorm, the reassociated
  form is shown to differ, fused Q matches flax RMSNorm, separate q/k
  epsilons are applied, non-trivial RoPE in bfloat16 matches the unfused
  reference, GQA shapes/values, `freqs_cis` sequence-length mismatch guard,
  VJP parity against unfused grads, and end-to-end `FlaxWanAttention` module
  wiring + unfused forward parity.
- Added to end_to_end/tpu/run_wan_stack_tests.sh (139 tests cumulative).

Author: Perseus14

end_to_end/tpu/run_wan_stack_tests.sh

@@ -35,5 +35,7 @@ TESTS=(
   # Ulysses x Ring attention (feat/ring-attention)
   "$T/attention_config_guards_test.py"
   "$T/custom_splash_unpadded_test.py"
+  # XLA-fused RMSNorm+RoPE producer (feat/fused-rmsnorm-rope-xla)
+  "$T/fused_producers_test.py"
 )
 PYTHONPATH="src${PYTHONPATH:+:$PYTHONPATH}" exec "${PYTHON:-python3}" -m pytest -q -rs "${TESTS[@]}" "$@"

src/maxdiffusion/kernels/fused_producers.py

@@ -0,0 +1,127 @@
+"""
+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.
+"""
+
+"""Optimized fused producers for Wan Attention."""
+
+import functools
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+
+
+# Static (non-array) parameters of the producer. They select shapes/epsilons and
+# must be compile-time constants so the jitted wrapper below can trace them.
+_STATIC_ARGNAMES = ("q_heads", "kv_heads", "dim_head", "eps", "k_eps")
+
+
+@functools.partial(jax.jit, static_argnames=_STATIC_ARGNAMES)
+def fused_rmsnorm_rope(
+    raw_q: jax.Array,
+    raw_k: jax.Array,
+    q_norm_scale: jax.Array,
+    k_norm_scale: jax.Array,
+    freqs_cis: jax.Array,
+    q_heads: int,
+    kv_heads: int | None = None,
+    dim_head: int = 128,
+    eps: float = 1e-6,
+    k_eps: float | None = None,
+) -> Tuple[jax.Array, jax.Array]:
+  """XLA-fused FP32 RMSNorm + BF16 RoPE + Head Transposition producer.
+
+  Performs FP32 RMSNorm normalization for numerical stability, casts to input
+  dtype (e.g. BF16), and applies RoPE rotation and head transposition in input
+  dtype precision, avoiding FP32 RoPE intermediates on long sequence lengths.
+  Accepts an optional `kv_heads` for GQA (`kv_heads != q_heads`); defaults to
+  `q_heads` (MHA, as used by `FlaxWanAttention`).
+
+  The function is `jax.jit`-wrapped (shape/epsilon parameters static). Inside the
+  pipeline's outer jit this is inlined and changes nothing; when called eagerly
+  (e.g. a full-size transformer block exercised in a unit test without an outer
+  jit) it lets XLA fuse the normalize/rotate chain instead of materialising
+  every intermediate, which otherwise exceeds HBM on long sequences.
+
+  Args:
+    raw_q: Raw query projection of shape [B, Sq, Dq] (where Dq = q_heads * dim_head).
+    raw_k: Raw key projection of shape [B, Sk, Dk] (where Dk = kv_heads * dim_head).
+    q_norm_scale: RMSNorm scale parameter for query of shape [Dq].
+    k_norm_scale: RMSNorm scale parameter for key of shape [Dk].
+    freqs_cis: Complex rotary embedding tensor of shape [1, 1, S, dim_head // 2].
+    q_heads: Number of query attention heads.
+    kv_heads: Number of key/value attention heads (defaults to q_heads for MHA).
+    dim_head: Dimension of each attention head.
+    eps: Epsilon for query RMSNorm numerical stability (and key if k_eps is None).
+    k_eps: Optional separate epsilon for key RMSNorm numerical stability.
+
+  Returns:
+    Transposed and RoPE-rotated (q_out, k_out) of shapes [B, q_heads, Sq, dim_head]
+    and [B, kv_heads, Sk, dim_head].
+  """
+  kv_heads = q_heads if kv_heads is None else kv_heads
+  effective_k_eps = eps if k_eps is None else k_eps
+  B, Sq, Dq = raw_q.shape
+  _, Sk, Dk = raw_k.shape
+
+  if Dq != q_heads * dim_head:
+    raise ValueError(f"raw_q feature dim ({Dq}) must equal q_heads ({q_heads}) * dim_head ({dim_head})")
+  if Dk != kv_heads * dim_head:
+    raise ValueError(f"raw_k feature dim ({Dk}) must equal kv_heads ({kv_heads}) * dim_head ({dim_head})")
+
+  # 1. FP32 RMSNorm for stability, then cast directly to target activation dtype.
+  #
+  # Association matters: Flax's `_normalize` computes `mul = rsqrt(var + eps)`,
+  # then `mul *= scale`, then `y = x * mul` -- i.e. x * (rsqrt * scale). Folding
+  # left-to-right as (x * rsqrt) * scale rounds differently and makes this path
+  # drift from `nnx.RMSNorm` bit-for-bit. Keep the parenthesisation below in
+  # step with Flax so the fused producer stays a pure fusion, not a numerical
+  # change.
+  q_fp32 = raw_q.astype(jnp.float32)
+  q_rms = jax.lax.rsqrt(jnp.mean(jnp.square(q_fp32), axis=-1, keepdims=True) + eps)
+  q_norm = (q_fp32 * (q_rms * q_norm_scale.astype(jnp.float32))).astype(raw_q.dtype)
+
+  k_fp32 = raw_k.astype(jnp.float32)
+  k_rms = jax.lax.rsqrt(jnp.mean(jnp.square(k_fp32), axis=-1, keepdims=True) + effective_k_eps)
+  k_norm = (k_fp32 * (k_rms * k_norm_scale.astype(jnp.float32))).astype(raw_k.dtype)
+
+  # 2. Reshape and transpose to [B, heads, S, dim_head]
+  q_h = q_norm.reshape(B, Sq, q_heads, dim_head).transpose(0, 2, 1, 3)
+  k_h = k_norm.reshape(B, Sk, kv_heads, dim_head).transpose(0, 2, 1, 3)
+
+  # 3. Direct RoPE with freqs_cis [1, 1, S, dim_head // 2] in input dtype
+  if freqs_cis.shape[2] != Sq or freqs_cis.shape[2] != Sk:
+    raise ValueError(
+        f"freqs_cis has {freqs_cis.shape[2]} positions, but raw_q/raw_k have Sq={Sq}, Sk={Sk}; "
+        "fused_rmsnorm_rope requires Sq == Sk == freqs_cis.shape[2]."
+    )
+  cos = jnp.real(freqs_cis)
+  sin = jnp.imag(freqs_cis)
+  cos_q, sin_q = cos.astype(raw_q.dtype), sin.astype(raw_q.dtype)
+  cos_k, sin_k = cos.astype(raw_k.dtype), sin.astype(raw_k.dtype)
+
+  q_pairs = q_h.reshape(B, q_heads, Sq, -1, 2)
+  q_0, q_1 = q_pairs[..., 0], q_pairs[..., 1]
+  q_out_0 = q_0 * cos_q - q_1 * sin_q
+  q_out_1 = q_0 * sin_q + q_1 * cos_q
+  q_out = jnp.stack([q_out_0, q_out_1], axis=-1).reshape(B, q_heads, Sq, dim_head)
+
+  k_pairs = k_h.reshape(B, kv_heads, Sk, -1, 2)
+  k_0, k_1 = k_pairs[..., 0], k_pairs[..., 1]
+  k_out_0 = k_0 * cos_k - k_1 * sin_k
+  k_out_1 = k_0 * sin_k + k_1 * cos_k
+  k_out = jnp.stack([k_out_0, k_out_1], axis=-1).reshape(B, kv_heads, Sk, dim_head)
+
+  return q_out, k_out

src/maxdiffusion/models/attention_flax.py

@@ -28,6 +28,7 @@
 from maxdiffusion.kernels.splash_attention import splash_attention_kernel as tokamax_splash_attention_kernel
 from maxdiffusion.kernels.splash_attention import ring_attention_kernel as tokamax_ring_attention_kernel
 from maxdiffusion.kernels.splash_attention import base as tokamax_splash_base
+from maxdiffusion.kernels.fused_producers import fused_rmsnorm_rope
 from einops import rearrange
 from .. import common_types, max_logging
 from maxdiffusion.tpu_utils import get_tpu_type, TpuType
@@ -3717,20 +3718,40 @@ def __call__(
         with jax.named_scope("value_proj"):
           value_proj = self.value(encoder_hidden_states)
 
-      if self.qk_norm:
-        with self.conditional_named_scope("attn_q_norm"):
-          query_proj = self.norm_q(query_proj)
-        if is_self_attention or cached_kv is None or "text" not in cached_kv:
-          with self.conditional_named_scope("attn_k_norm"):
-            key_proj = self.norm_k(key_proj)
-
-      if rotary_emb is not None:
-        with self.conditional_named_scope("attn_rope"):
-          query_proj = _unflatten_heads(query_proj, self.heads)
-          key_proj = _unflatten_heads(key_proj, self.heads)
+      if rotary_emb is not None and self.qk_norm and is_self_attention:
+        with self.conditional_named_scope("fused_rmsnorm_rope"):
+          q_scale = self.norm_q.scale[...]
+          k_scale = self.norm_k.scale[...]
+          # norm_q is built with epsilon=eps; norm_k keeps flax's default (1e-6).
+          q_eps = self.norm_q.epsilon
+          k_eps = self.norm_k.epsilon
+          query_proj, key_proj = fused_rmsnorm_rope(
+              query_proj,
+              key_proj,
+              q_scale,
+              k_scale,
+              rotary_emb,
+              q_heads=self.heads,
+              dim_head=self.dim_head,
+              eps=q_eps,
+              k_eps=k_eps,
+          )
           value_proj = _unflatten_heads(value_proj, self.heads)
-          # output of _unflatten_heads Batch, heads, seq_len, head_dim
-          query_proj, key_proj = self._apply_rope(query_proj, key_proj, rotary_emb)
+      else:
+        if self.qk_norm:
+          with self.conditional_named_scope("attn_q_norm"):
+            query_proj = self.norm_q(query_proj)
+          if is_self_attention or cached_kv is None or "text" not in cached_kv:
+            with self.conditional_named_scope("attn_k_norm"):
+              key_proj = self.norm_k(key_proj)
+
+        if rotary_emb is not None:
+          with self.conditional_named_scope("attn_rope"):
+            query_proj = _unflatten_heads(query_proj, self.heads)
+            key_proj = _unflatten_heads(key_proj, self.heads)
+            value_proj = _unflatten_heads(value_proj, self.heads)
+            # output of _unflatten_heads Batch, heads, seq_len, head_dim
+            query_proj, key_proj = self._apply_rope(query_proj, key_proj, rotary_emb)
 
       query_proj = checkpoint_name(query_proj, "query_proj")
       key_proj = checkpoint_name(key_proj, "key_proj")