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

end_to_end/tpu/run_wan_stack_tests.sh

@@ -34,5 +34,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,122 @@
+"""
+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
+  cos = jnp.real(freqs_cis)
+  sin = jnp.imag(freqs_cis)
+  cos_q, sin_q = cos[:, :, :Sq, :].astype(raw_q.dtype), sin[:, :, :Sq, :].astype(raw_q.dtype)
+  cos_k, sin_k = cos[:, :, :Sk, :].astype(raw_k.dtype), sin[:, :, :Sk, :].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
@@ -3364,6 +3365,7 @@ def __init__(
     self.out_axis_names = out_axis_names
     self.enable_jax_named_scopes = enable_jax_named_scopes
     self.is_self_attention = is_self_attention
+    self.eps = eps
 
     cross_attention_remapped_to_flash = not is_self_attention and attention_kernel in (
         "tokamax_ring",
@@ -3638,20 +3640,39 @@ 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[...]
+          q_eps = getattr(self.norm_q, "epsilon", self.eps)
+          k_eps = getattr(self.norm_k, "epsilon", self.eps)
+          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")