|
| 1 | +""" |
| 2 | +Copyright 2026 Google LLC |
| 3 | +
|
| 4 | +Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +you may not use this file except in compliance with the License. |
| 6 | +You may obtain a copy of the License at |
| 7 | +
|
| 8 | + https://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +
|
| 10 | +Unless required by applicable law or agreed to in writing, software |
| 11 | +distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +See the License for the specific language governing permissions and |
| 14 | +limitations under the License. |
| 15 | +""" |
| 16 | + |
| 17 | +"""Optimized fused producers for Wan Attention.""" |
| 18 | + |
| 19 | +import functools |
| 20 | +from typing import Tuple |
| 21 | + |
| 22 | +import jax |
| 23 | +import jax.numpy as jnp |
| 24 | + |
| 25 | + |
| 26 | +# Static (non-array) parameters of the producer. They select shapes/epsilons and |
| 27 | +# must be compile-time constants so the jitted wrapper below can trace them. |
| 28 | +_STATIC_ARGNAMES = ("q_heads", "kv_heads", "dim_head", "eps", "k_eps") |
| 29 | + |
| 30 | + |
| 31 | +@functools.partial(jax.jit, static_argnames=_STATIC_ARGNAMES) |
| 32 | +def fused_rmsnorm_rope( |
| 33 | + raw_q: jax.Array, |
| 34 | + raw_k: jax.Array, |
| 35 | + q_norm_scale: jax.Array, |
| 36 | + k_norm_scale: jax.Array, |
| 37 | + freqs_cis: jax.Array, |
| 38 | + q_heads: int, |
| 39 | + kv_heads: int | None = None, |
| 40 | + dim_head: int = 128, |
| 41 | + eps: float = 1e-6, |
| 42 | + k_eps: float | None = None, |
| 43 | +) -> Tuple[jax.Array, jax.Array]: |
| 44 | + """XLA-fused FP32 RMSNorm + BF16 RoPE + Head Transposition producer. |
| 45 | +
|
| 46 | + Performs FP32 RMSNorm normalization for numerical stability, casts to input |
| 47 | + dtype (e.g. BF16), and applies RoPE rotation and head transposition in input |
| 48 | + dtype precision, avoiding FP32 RoPE intermediates on long sequence lengths. |
| 49 | + Accepts an optional `kv_heads` for GQA (`kv_heads != q_heads`); defaults to |
| 50 | + `q_heads` (MHA, as used by `FlaxWanAttention`). |
| 51 | +
|
| 52 | + The function is `jax.jit`-wrapped (shape/epsilon parameters static). Inside the |
| 53 | + pipeline's outer jit this is inlined and changes nothing; when called eagerly |
| 54 | + (e.g. a full-size transformer block exercised in a unit test without an outer |
| 55 | + jit) it lets XLA fuse the normalize/rotate chain instead of materialising |
| 56 | + every intermediate, which otherwise exceeds HBM on long sequences. |
| 57 | +
|
| 58 | + Args: |
| 59 | + raw_q: Raw query projection of shape [B, Sq, Dq] (where Dq = q_heads * dim_head). |
| 60 | + raw_k: Raw key projection of shape [B, Sk, Dk] (where Dk = kv_heads * dim_head). |
| 61 | + q_norm_scale: RMSNorm scale parameter for query of shape [Dq]. |
| 62 | + k_norm_scale: RMSNorm scale parameter for key of shape [Dk]. |
| 63 | + freqs_cis: Complex rotary embedding tensor of shape [1, 1, S, dim_head // 2]. |
| 64 | + q_heads: Number of query attention heads. |
| 65 | + kv_heads: Number of key/value attention heads (defaults to q_heads for MHA). |
| 66 | + dim_head: Dimension of each attention head. |
| 67 | + eps: Epsilon for query RMSNorm numerical stability (and key if k_eps is None). |
| 68 | + k_eps: Optional separate epsilon for key RMSNorm numerical stability. |
| 69 | +
|
| 70 | + Returns: |
| 71 | + Transposed and RoPE-rotated (q_out, k_out) of shapes [B, q_heads, Sq, dim_head] |
| 72 | + and [B, kv_heads, Sk, dim_head]. |
| 73 | + """ |
| 74 | + kv_heads = q_heads if kv_heads is None else kv_heads |
| 75 | + effective_k_eps = eps if k_eps is None else k_eps |
| 76 | + B, Sq, Dq = raw_q.shape |
| 77 | + _, Sk, Dk = raw_k.shape |
| 78 | + |
| 79 | + if Dq != q_heads * dim_head: |
| 80 | + raise ValueError(f"raw_q feature dim ({Dq}) must equal q_heads ({q_heads}) * dim_head ({dim_head})") |
| 81 | + if Dk != kv_heads * dim_head: |
| 82 | + raise ValueError(f"raw_k feature dim ({Dk}) must equal kv_heads ({kv_heads}) * dim_head ({dim_head})") |
| 83 | + |
| 84 | + # 1. FP32 RMSNorm for stability, then cast directly to target activation dtype. |
| 85 | + # |
| 86 | + # Association matters: Flax's `_normalize` computes `mul = rsqrt(var + eps)`, |
| 87 | + # then `mul *= scale`, then `y = x * mul` -- i.e. x * (rsqrt * scale). Folding |
| 88 | + # left-to-right as (x * rsqrt) * scale rounds differently and makes this path |
| 89 | + # drift from `nnx.RMSNorm` bit-for-bit. Keep the parenthesisation below in |
| 90 | + # step with Flax so the fused producer stays a pure fusion, not a numerical |
| 91 | + # change. |
| 92 | + q_fp32 = raw_q.astype(jnp.float32) |
| 93 | + q_rms = jax.lax.rsqrt(jnp.mean(jnp.square(q_fp32), axis=-1, keepdims=True) + eps) |
| 94 | + q_norm = (q_fp32 * (q_rms * q_norm_scale.astype(jnp.float32))).astype(raw_q.dtype) |
| 95 | + |
| 96 | + k_fp32 = raw_k.astype(jnp.float32) |
| 97 | + k_rms = jax.lax.rsqrt(jnp.mean(jnp.square(k_fp32), axis=-1, keepdims=True) + effective_k_eps) |
| 98 | + k_norm = (k_fp32 * (k_rms * k_norm_scale.astype(jnp.float32))).astype(raw_k.dtype) |
| 99 | + |
| 100 | + # 2. Reshape and transpose to [B, heads, S, dim_head] |
| 101 | + q_h = q_norm.reshape(B, Sq, q_heads, dim_head).transpose(0, 2, 1, 3) |
| 102 | + k_h = k_norm.reshape(B, Sk, kv_heads, dim_head).transpose(0, 2, 1, 3) |
| 103 | + |
| 104 | + # 3. Direct RoPE with freqs_cis [1, 1, S, dim_head // 2] in input dtype |
| 105 | + cos = jnp.real(freqs_cis).astype(raw_q.dtype) |
| 106 | + sin = jnp.imag(freqs_cis).astype(raw_q.dtype) |
| 107 | + cos_q, sin_q = cos[:, :, :Sq, :], sin[:, :, :Sq, :] |
| 108 | + cos_k, sin_k = cos[:, :, :Sk, :], sin[:, :, :Sk, :] |
| 109 | + |
| 110 | + q_pairs = q_h.reshape(B, q_heads, Sq, -1, 2) |
| 111 | + q_0, q_1 = q_pairs[..., 0], q_pairs[..., 1] |
| 112 | + q_out_0 = q_0 * cos_q - q_1 * sin_q |
| 113 | + q_out_1 = q_0 * sin_q + q_1 * cos_q |
| 114 | + q_out = jnp.stack([q_out_0, q_out_1], axis=-1).reshape(B, q_heads, Sq, dim_head) |
| 115 | + |
| 116 | + k_pairs = k_h.reshape(B, kv_heads, Sk, -1, 2) |
| 117 | + k_0, k_1 = k_pairs[..., 0], k_pairs[..., 1] |
| 118 | + k_out_0 = k_0 * cos_k - k_1 * sin_k |
| 119 | + k_out_1 = k_0 * sin_k + k_1 * cos_k |
| 120 | + k_out = jnp.stack([k_out_0, k_out_1], axis=-1).reshape(B, kv_heads, Sk, dim_head) |
| 121 | + |
| 122 | + return q_out, k_out |
0 commit comments