diff --git a/end_to_end/tpu/run_wan_fast_inference.sh b/end_to_end/tpu/run_wan_fast_inference.sh index d82227f7a..32bdc7c65 100755 --- a/end_to_end/tpu/run_wan_fast_inference.sh +++ b/end_to_end/tpu/run_wan_fast_inference.sh @@ -63,6 +63,8 @@ # against the XLA-path video, visually clean); on tpu7x the # output was bit-identical in that measurement. # USE_FUSED_ROPE_KERNEL=false restores the XLA producer. +# WAN_CROSS_ATTN_KERNEL +# override Wan fast-inference kernel/layout switches # DP / CP / PER_DEVICE_BATCH / SEED # override mesh parallelism, per-device batch, or RNG seed (default 12345) set -euo pipefail @@ -167,6 +169,7 @@ case "$TPU_PROFILE" in DEFAULT_VAE_CHUNK=1 DEFAULT_VAE_SPATIAL=8 DEFAULT_FUSED_ROPE=true + DEFAULT_CROSS_ATTN_KERNEL=pallas ;; v7) PLATFORM_LIBTPU="$V7_LIBTPU" @@ -183,6 +186,7 @@ case "$TPU_PROFILE" in DEFAULT_VAE_CHUNK=1 DEFAULT_VAE_SPATIAL=8 DEFAULT_FUSED_ROPE=true + DEFAULT_CROSS_ATTN_KERNEL=pallas ;; *) echo "== warning: unrecognised or non-v6e/v7 accelerator '${ACCEL_TYPE:-unknown}';" \ @@ -201,6 +205,7 @@ case "$TPU_PROFILE" in DEFAULT_VAE_CHUNK=1 DEFAULT_VAE_SPATIAL=8 DEFAULT_FUSED_ROPE=false + DEFAULT_CROSS_ATTN_KERNEL=xla ;; esac @@ -232,6 +237,7 @@ USE_FUSED_ROPE_KERNEL=${USE_FUSED_ROPE_KERNEL:-$DEFAULT_FUSED_ROPE} FUSED_ROPE_ARGS=() [ -n "${FUSED_ROPE_BLOCK_S:-}" ] && FUSED_ROPE_ARGS+=("fused_rope_block_s=$FUSED_ROPE_BLOCK_S") [ -n "${FUSED_ROPE_HEAD_BLOCK:-}" ] && FUSED_ROPE_ARGS+=("fused_rope_head_block=$FUSED_ROPE_HEAD_BLOCK") +WAN_CROSS_ATTN_KERNEL=${WAN_CROSS_ATTN_KERNEL:-$DEFAULT_CROSS_ATTN_KERNEL} # Mesh: context parallelism carries the Ulysses shards, data parallelism takes # whatever chips remain. Defaults to CP=4 / DP=2 on an 8-chip slice. On a slice @@ -294,6 +300,7 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \ use_kv_cache=true use_base2_exp=true use_experimental_scheduler=true \ use_fused_rope_kernel="$USE_FUSED_ROPE_KERNEL" \ ${FUSED_ROPE_ARGS[@]+"${FUSED_ROPE_ARGS[@]}"} \ + wan_cross_attn_kernel="$WAN_CROSS_ATTN_KERNEL" \ fps=16 ${GUIDANCE_ARGS[@]+"${GUIDANCE_ARGS[@]}"} \ seed="${SEED:-12345}" \ flash_block_sizes="$FLASH_BLOCK_SIZES" \ diff --git a/end_to_end/tpu/run_wan_stack_tests.sh b/end_to_end/tpu/run_wan_stack_tests.sh index ad2487f0a..409a536ea 100755 --- a/end_to_end/tpu/run_wan_stack_tests.sh +++ b/end_to_end/tpu/run_wan_stack_tests.sh @@ -44,5 +44,7 @@ TESTS=( # Fused RMSNorm+RoPE producer, Wan switches in config (feat/wan-custom-kernels) "$T/fused_rmsnorm_rope_pallas_test.py" "$T/wan_runtime_options_test.py" + # Short-KV cross-attention Pallas kernel (feat/wan-cross-attn-pallas) + "$T/cross_attention_pallas_test.py" ) PYTHONPATH="src${PYTHONPATH:+:$PYTHONPATH}" exec "${PYTHON:-python3}" -m pytest -q -rs "${TESTS[@]}" "$@" diff --git a/src/maxdiffusion/configs/base_wan_14b.yml b/src/maxdiffusion/configs/base_wan_14b.yml index 3465cb184..767936b82 100644 --- a/src/maxdiffusion/configs/base_wan_14b.yml +++ b/src/maxdiffusion/configs/base_wan_14b.yml @@ -133,6 +133,12 @@ wan_cfg_before_unpatchify: True wan_cross_attn_prescale_kv: False # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32". wan_rope_accum: "auto" +# Text cross-attention backend where it would otherwise fall back to XLA dot-product +# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that +# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical). +wan_cross_attn_kernel: "xla" +# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only). +wan_cross_attn_cpu_interpret: False # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. ulysses_shards: -1 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder. diff --git a/src/maxdiffusion/configs/base_wan_1_3b.yml b/src/maxdiffusion/configs/base_wan_1_3b.yml index e358e33a8..3c36364a6 100644 --- a/src/maxdiffusion/configs/base_wan_1_3b.yml +++ b/src/maxdiffusion/configs/base_wan_1_3b.yml @@ -130,6 +130,12 @@ wan_cfg_before_unpatchify: True wan_cross_attn_prescale_kv: False # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32". wan_rope_accum: "auto" +# Text cross-attention backend where it would otherwise fall back to XLA dot-product +# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that +# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical). +wan_cross_attn_kernel: "xla" +# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only). +wan_cross_attn_cpu_interpret: False # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. ulysses_shards: -1 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder. diff --git a/src/maxdiffusion/configs/base_wan_27b.yml b/src/maxdiffusion/configs/base_wan_27b.yml index 43227799f..45e9c6725 100644 --- a/src/maxdiffusion/configs/base_wan_27b.yml +++ b/src/maxdiffusion/configs/base_wan_27b.yml @@ -162,6 +162,12 @@ wan_cfg_before_unpatchify: True wan_cross_attn_prescale_kv: False # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32". wan_rope_accum: "auto" +# Text cross-attention backend where it would otherwise fall back to XLA dot-product +# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that +# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical). +wan_cross_attn_kernel: "xla" +# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only). +wan_cross_attn_cpu_interpret: False # For attention=ulysses_ring*, hidden Ulysses shard count; ring shards are context / this. ulysses_shards: -1 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder. diff --git a/src/maxdiffusion/configs/base_wan_animate.yml b/src/maxdiffusion/configs/base_wan_animate.yml index 80831d6a5..28784c8a5 100644 --- a/src/maxdiffusion/configs/base_wan_animate.yml +++ b/src/maxdiffusion/configs/base_wan_animate.yml @@ -112,6 +112,12 @@ wan_cfg_before_unpatchify: True wan_cross_attn_prescale_kv: False # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32". wan_rope_accum: "auto" +# Text cross-attention backend where it would otherwise fall back to XLA dot-product +# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that +# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical). +wan_cross_attn_kernel: "xla" +# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only). +wan_cross_attn_cpu_interpret: False # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. ulysses_shards: -1 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder. diff --git a/src/maxdiffusion/configs/base_wan_i2v_14b.yml b/src/maxdiffusion/configs/base_wan_i2v_14b.yml index 834c40592..66650a0aa 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_14b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_14b.yml @@ -114,6 +114,12 @@ wan_cfg_before_unpatchify: True wan_cross_attn_prescale_kv: False # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32". wan_rope_accum: "auto" +# Text cross-attention backend where it would otherwise fall back to XLA dot-product +# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that +# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical). +wan_cross_attn_kernel: "xla" +# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only). +wan_cross_attn_cpu_interpret: False # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. ulysses_shards: -1 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder. diff --git a/src/maxdiffusion/configs/base_wan_i2v_27b.yml b/src/maxdiffusion/configs/base_wan_i2v_27b.yml index 2bd82a40f..dfc2d7e68 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_27b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_27b.yml @@ -114,6 +114,12 @@ wan_cfg_before_unpatchify: True wan_cross_attn_prescale_kv: False # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32". wan_rope_accum: "auto" +# Text cross-attention backend where it would otherwise fall back to XLA dot-product +# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that +# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical). +wan_cross_attn_kernel: "xla" +# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only). +wan_cross_attn_cpu_interpret: False # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. ulysses_shards: -1 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder. diff --git a/src/maxdiffusion/kernels/cross_attention_pallas.py b/src/maxdiffusion/kernels/cross_attention_pallas.py new file mode 100644 index 000000000..b8f5663e0 --- /dev/null +++ b/src/maxdiffusion/kernels/cross_attention_pallas.py @@ -0,0 +1,430 @@ +""" +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. +""" + +"""Fused short-KV cross-attention Pallas kernel for TPU. + +Why +--- +Wan's text cross-attention attends ~19K video tokens per shard to 512 cached +text tokens. That KV length is below `flash_min_seq_length`, so the layer runs +as XLA dot-product attention (`_apply_attention_dot`). XLA materialises the +`[B, heads, Sq, 512]` attention scores (in bf16 under Wan's default +`float32_qk_product: False`, or FP32 when `float32_qk_product: True`) and +softmax probabilities in HBM and streams them across QK^T, softmax, and PV, +while relayouting `[B, S, H*D]` to `[B, S, H, D]` and back around the two +einsums. At 720p/81 frames that is 3.94 ms per layer on v6e and 1.81 ms on +tpu7x, nearly all of it HBM traffic. + +Here the whole KV of a block of heads stays resident in VMEM and each query +tile makes one pass: QK^T on the MXU, softmax in registers, PV on the MXU. HBM +sees Q once, K/V once per head block and the output once, all in the flat +`[B, S, H*D]` layout, so the relayouts disappear as well. + +Numerics +-------- +Scores and the softmax state stay in FP32 inside the kernel (accumulated in FP32 +on the MXU), whereas Wan's default XLA fallback (`float32_qk_product: False`) +rounds the QK^T logits to bf16 before running a bf16 softmax. The kernel result +is within bf16 rounding of an FP32-score reference (and at least as close to +that reference as the XLA path), but not bit-identical, because: + + * the softmax is normalised after PV: the unnormalised probabilities are + rounded to bf16 for the MXU and the FP32 PV product is divided by their row + sum, where XLA rounds the normalised probabilities; + * with `sum_mode="mxu"` the row sum also comes from the MXU, by appending a + block of ones to each head of V. On a 256-wide MXU with `dim_head == 128` + the extra columns are free, and the weights actually applied to V then sum + to one up to FP32 accumulation; + * the exponent is evaluated as `exp2((s - m) * scale * log2(e))`. + +`max_mode="tile"` shifts every row of a query tile by the tile-wide maximum +instead of its own. Softmax is shift-invariant, so this is exact unless a row's +maximum sits more than ~87 (in scaled-logit units) below the tile's, where its +exponentials underflow. It exists to measure what the per-row reduction costs; +production uses `"row"`. +""" + +import functools +import math + +import jax +import jax.numpy as jnp +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu + +NUM_LANES = 128 +LOG2E = math.log2(math.e) + +# Tile sizes, swept at the Wan shard shape (q [1, 18900, 5120] unpadded or +# [1, 18944, 5120] with `wan_seq_pad="lane"`, kv 512 tokens). `block_q` is an +# upper bound: the rows are split into `cdiv(Sq, block_q)` equal tiles rounded +# up to a multiple of 16, so with `DEFAULT_BLOCK_Q = 2400` both 18900 and 18944 +# rows run as 8 tiles of 2368 instead of 7 full tiles plus a mostly empty +# eighth. Per-call time for 40 heads: +# tpu7x: 0.546 ms at 1024 x 8 heads, 0.501 ms at 2368 x 8, 0.487 ms at 2368 x 10 +# v6e: 0.567 ms at 1024 x 8 heads, 0.514 ms at 2368 x 8, 0.512 ms at 2368 x 10 +# against 1.81 ms (tpu7x) and 3.94 ms (v6e) for the XLA path inside the model. +DEFAULT_BLOCK_Q = 2400 +DEFAULT_HEAD_BLOCK = 10 +# v6e exposes 128 MiB of VMEM per core and tpu7x exposes 64 MiB. 64 MiB fits +# both (the full tpu7x budget). Other TPU generations are not validated; see +# `_resolve_vmem_limit_bytes` and `vmem_unusable_reason`. +DEFAULT_VMEM_LIMIT_BYTES = 64 * 1024 * 1024 +# Device kinds on which this kernel has been compiled and run at the default +# tiles with DEFAULT_VMEM_LIMIT_BYTES. +_VMEM_LIMIT_VALIDATED_KINDS = ("7x", "v6") +# The full KV of a head block is VMEM-resident, so this kernel is only meant +# for short KV sequences such as Wan's 512 text tokens. +MAX_KV_LEN = 4096 + +MAX_MODES = ("row", "tile") +SUM_MODES = ("mxu", "vpu") + + +def _tpu_kinds(mesh=None) -> list[str]: + """Lower-cased device kinds of the TPUs in `mesh` (or `jax.devices()`).""" + try: + devices = list(mesh.devices.flat) if mesh is not None else jax.devices() + except RuntimeError: + devices = [] + return [getattr(d, "device_kind", "").lower() for d in devices if getattr(d, "platform", "") == "tpu"] + + +def _resolve_vmem_limit_bytes(vmem_limit_bytes: int | None = None, mesh=None) -> int | None: + """Scoped VMEM budget for the kernel. + + An explicit value wins. Otherwise DEFAULT_VMEM_LIMIT_BYTES is requested only + when every TPU in `mesh` (or `jax.devices()`) is a generation this kernel was + validated on. Any other TPU returns None, so Mosaic keeps its own + conservative scoped default (the default tiles do not fit in it; see + `vmem_unusable_reason`). Off TPU (interpret mode) the value is unused, and + the default is returned. + """ + if vmem_limit_bytes is not None: + return vmem_limit_bytes + kinds = _tpu_kinds(mesh) + if not kinds: + return DEFAULT_VMEM_LIMIT_BYTES + if all(any(tag in k for tag in _VMEM_LIMIT_VALIDATED_KINDS) for k in kinds): + return DEFAULT_VMEM_LIMIT_BYTES + return None + + +def vmem_budget_is_validated(mesh=None) -> bool: + """True iff every device in `mesh` (or `jax.devices()`) is a TPU in `_VMEM_LIMIT_VALIDATED_KINDS`.""" + try: + devices = list(mesh.devices.flat) if mesh is not None else jax.devices() + except RuntimeError: + return False + kinds = _tpu_kinds(mesh) + return bool(kinds) and len(kinds) == len(devices) and _resolve_vmem_limit_bytes(None, mesh) is not None + + +def _resolve_block_q(seq_q: int, block_q: int) -> int: + """Rows per query tile actually used for `seq_q` rows and upper bound `block_q`.""" + if block_q <= 0 or block_q % 8 != 0: + raise ValueError(f"block_q ({block_q}) must be a positive multiple of 8.") + if block_q >= seq_q: + if seq_q % 8 != 0: + raise ValueError(f"seq_q ({seq_q}) must be a multiple of 8 when block_q ({block_q}) >= seq_q.") + return seq_q + # Same tile count, but spread evenly: 16-row alignment matches bf16 + # sublane packing and never exceeds `block_q` by more than 15 rows. + even = pl.cdiv(seq_q, pl.cdiv(seq_q, block_q)) + resolved = min(seq_q, pl.cdiv(even, 16) * 16) + if resolved % 8 != 0: + raise ValueError(f"resolved block_q ({resolved}) must be a multiple of 8.") + return resolved + + +def estimate_vmem_bytes( + *, + seq_q: int, + seq_kv: int, + dim_head: int, + block_q: int = DEFAULT_BLOCK_Q, + head_block: int = DEFAULT_HEAD_BLOCK, + sum_mode: str = "mxu", + itemsize: int = 2, +) -> int: + """Conservative estimate of the kernel's VMEM working set, in bytes. + + Counts the double-buffered Q and output tiles (`[block_q, head_block * + dim_head]`) and K/V blocks (`[seq_kv, head_block * dim_head]`), plus one + head's temporaries: FP32 scores and probabilities `[block_q, seq_kv]`, their + `itemsize` copy fed to the MXU, the FP32 PV accumulator `[block_q, v_width]` + and the ones-augmented V `[seq_kv, v_width]`. At the Wan shard shape (18900 + rows -> 2368-row tiles, 10 heads, 512 KV, bf16) this is ~45 MB: above + Mosaic's default scoped limit, within DEFAULT_VMEM_LIMIT_BYTES. + + The estimate has only been checked against real compiles at Wan's KV length + (512 text tokens); for other KV lengths it is an untested extrapolation. + """ + block_q = _resolve_block_q(seq_q, block_q) + v_width = 2 * dim_head if sum_mode == "mxu" else dim_head + width = head_block * dim_head + io_tiles = 2 * 2 * block_q * width * itemsize # q + out, double-buffered + kv_tiles = 2 * 2 * seq_kv * width * itemsize # k + v, double-buffered + scores = 2 * block_q * seq_kv * 4 + block_q * seq_kv * itemsize # s, p (f32), p (MXU dtype) + pv = block_q * v_width * 4 + seq_kv * v_width * itemsize + return io_tiles + kv_tiles + scores + pv + + +def vmem_unusable_reason( + *, + seq_q: int, + seq_kv: int, + dim_head: int, + block_q: int = DEFAULT_BLOCK_Q, + head_block: int = DEFAULT_HEAD_BLOCK, + sum_mode: str = "mxu", + itemsize: int = 2, + mesh=None, + vmem_limit_bytes: int | None = None, +) -> str | None: + """Why the kernel should not be compiled for this shape/device, or None. + + Callers that can fall back to XLA use this to avoid a Mosaic compile failure: + on a TPU generation outside `_VMEM_LIMIT_VALIDATED_KINDS` (with no explicit + budget) the kernel would run under Mosaic's default scoped limit, which the + default tiles exceed, and any shape whose `estimate_vmem_bytes` exceeds the + resolved budget would fail to compile. + """ + budget = _resolve_vmem_limit_bytes(vmem_limit_bytes, mesh) + if budget is None: + return f"VMEM budget not validated on TPU kind(s) {sorted(set(_tpu_kinds(mesh)))}" + need = estimate_vmem_bytes( + seq_q=seq_q, + seq_kv=seq_kv, + dim_head=dim_head, + block_q=block_q, + head_block=head_block, + sum_mode=sum_mode, + itemsize=itemsize, + ) + if need > budget: + return f"estimated VMEM working set {need / 2**20:.1f} MiB exceeds the {budget / 2**20:.1f} MiB budget" + return None + + +# q @ k^T: contract the last axis of both operands. +_NT_DIMS = (((1,), (1,)), ((), ())) +# p @ v: standard matmul. +_NN_DIMS = (((1,), (0,)), ((), ())) + + +def _cross_attention_kernel( + q_ref, + k_ref, + v_ref, + o_ref, + *, + head_block: int, + dim_head: int, + exp2_scale: float, + max_mode: str, + sum_mode: str, + seq_len: int, + block_q: int, +): + ones = jnp.ones((v_ref.shape[1], dim_head), dtype=v_ref.dtype) if sum_mode == "mxu" else None + row_valid = None + if max_mode == "tile" and seq_len % block_q != 0: + # Rows past the end of the ragged last tile hold stale VMEM. Per-row + # softmax never mixes rows, but a tile-wide maximum would read them. + rows = pl.program_id(2) * block_q + jax.lax.broadcasted_iota(jnp.int32, (block_q, 1), 0) + row_valid = rows < seq_len + + for h in range(head_block): + q = q_ref[0, :, pl.ds(h * dim_head, dim_head)] + k = k_ref[0, :, pl.ds(h * dim_head, dim_head)] + s = jax.lax.dot_general(q, k, _NT_DIMS, preferred_element_type=jnp.float32) + + if max_mode == "row": + m = jnp.max(s, axis=-1, keepdims=True) + else: + s_for_max = s if row_valid is None else jnp.where(row_valid, s, -jnp.inf) + # Two 1-D keepdims reductions stay in 2-D vector registers in Mosaic TPU + # without triggering scalar-reduction's 3-D newaxis/squeeze proxy. + m = jnp.max(jnp.max(s_for_max, axis=0, keepdims=True), axis=1, keepdims=True) + p = jnp.exp2((s - m) * exp2_scale) + + v = v_ref[0, :, pl.ds(h * dim_head, dim_head)] + if sum_mode == "mxu": + v_aug = jnp.concatenate([v, ones], axis=-1) + acc = jax.lax.dot_general(p.astype(v.dtype), v_aug, _NN_DIMS, preferred_element_type=jnp.float32) + num, den = acc[:, :dim_head], acc[:, dim_head:] + else: + den = jnp.sum(p, axis=-1, keepdims=True) + num = jax.lax.dot_general(p.astype(v.dtype), v, _NN_DIMS, preferred_element_type=jnp.float32) + o_ref[0, :, pl.ds(h * dim_head, dim_head)] = (num / den).astype(o_ref.dtype) + + +def cross_attention_pallas( + q: jax.Array, + k: jax.Array, + v: jax.Array, + *, + heads: int, + dim_head: int = 128, + scale: float | None = None, + k_prescaled: bool = False, + block_q: int = DEFAULT_BLOCK_Q, + head_block: int = DEFAULT_HEAD_BLOCK, + max_mode: str = "row", + sum_mode: str = "mxu", + vmem_limit_bytes: int | None = None, + mesh=None, + interpret: bool = False, +) -> jax.Array: + """Unmasked multi-head attention against a short, fully VMEM-resident KV. + + Computes `softmax(q_h @ k_h^T * scale) @ v_h` for every head `h`, reading and + writing the flat `[B, S, heads * dim_head]` layout directly. + + Args: + q: Queries, `[B, Sq, heads * dim_head]`. + k: Keys, `[B, Skv, heads * dim_head]`. + v: Values, `[B, Skv, heads * dim_head]`. + heads: Number of heads (MHA only; K/V carry the same heads as Q). + dim_head: Per-head dimension; must be a multiple of 128. + scale: Softmax scale; defaults to `dim_head ** -0.5`. + k_prescaled: `k` was already multiplied by `scale`, so it is not reapplied. + block_q: Upper bound on query rows per grid step. The rows are split into + `cdiv(Sq, block_q)` equal 16-aligned tiles; the last may be ragged. + head_block: Heads per grid step; must divide `heads`. + max_mode: `"row"` (per-row softmax shift) or `"tile"` (see module docs). + sum_mode: `"mxu"` takes the softmax denominator from the PV matmul via a + ones-augmented V; `"vpu"` reduces it on the vector unit. + vmem_limit_bytes: Scoped VMEM budget handed to Mosaic. When None, 64 MiB on + validated TPUs (v6e, tpu7x) and Mosaic's default elsewhere; see + `_resolve_vmem_limit_bytes`. A ValueError is raised up front only when a + budget resolves (an explicit value, or the validated-TPU default) and + `estimate_vmem_bytes` exceeds it. Under Mosaic's default nothing is + checked here; callers that can fall back use `vmem_unusable_reason`. + mesh: Mesh whose devices decide the default VMEM budget (defaults to + `jax.devices()`). + interpret: Run in Pallas interpret mode (for CPU tests). + + Returns: + Attention output, `[B, Sq, heads * dim_head]`, in `q.dtype`. + """ + if max_mode not in MAX_MODES: + raise ValueError(f"max_mode must be one of {MAX_MODES}, got {max_mode!r}.") + if sum_mode not in SUM_MODES: + raise ValueError(f"sum_mode must be one of {SUM_MODES}, got {sum_mode!r}.") + if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: + raise ValueError(f"expected flat [B, S, H*D] inputs, got q={q.shape} k={k.shape} v={v.shape}.") + if dim_head % NUM_LANES != 0: + raise ValueError(f"dim_head must be a multiple of {NUM_LANES}, got {dim_head}.") + if heads % head_block != 0: + raise ValueError(f"head_block ({head_block}) must divide heads ({heads}).") + batch, seq_q, feature = q.shape + if feature != heads * dim_head: + raise ValueError(f"q feature dim ({feature}) must equal heads ({heads}) * dim_head ({dim_head}).") + if k.shape != v.shape or k.shape[0] != batch or k.shape[2] != feature: + raise ValueError(f"k and v must both be [{batch}, Skv, {feature}], got k={k.shape} v={v.shape}.") + seq_kv = k.shape[1] + if seq_kv > MAX_KV_LEN: + raise ValueError(f"KV length {seq_kv} exceeds MAX_KV_LEN ({MAX_KV_LEN}); use a flash kernel instead.") + if seq_kv % 8 != 0: + raise ValueError(f"KV length ({seq_kv}) must be a multiple of 8.") + + if scale is None: + scale = dim_head**-0.5 + exp2_scale = (1.0 if k_prescaled else float(scale)) * LOG2E + block_q = _resolve_block_q(seq_q, block_q) + + v_width = 2 * dim_head if sum_mode == "mxu" else dim_head + num_q_blocks = pl.cdiv(seq_q, block_q) + padded_rows = num_q_blocks * block_q + itemsize = jnp.dtype(q.dtype).itemsize + cost = pl.CostEstimate( + flops=2 * batch * heads * padded_rows * seq_kv * (dim_head + v_width), + transcendentals=batch * heads * padded_rows * seq_kv, + bytes_accessed=itemsize * batch * 2 * (seq_q + seq_kv) * feature, + ) + + resolved_vmem_limit_bytes = _resolve_vmem_limit_bytes(vmem_limit_bytes, mesh) + if not interpret and resolved_vmem_limit_bytes is not None: + need = estimate_vmem_bytes( + seq_q=seq_q, + seq_kv=seq_kv, + dim_head=dim_head, + block_q=block_q, + head_block=head_block, + sum_mode=sum_mode, + itemsize=itemsize, + ) + if need > resolved_vmem_limit_bytes: + raise ValueError( + f"estimated VMEM working set ({need} bytes) exceeds vmem_limit_bytes ({resolved_vmem_limit_bytes}); " + "reduce block_q or head_block." + ) + kernel = functools.partial( + _cross_attention_kernel, + head_block=head_block, + dim_head=dim_head, + exp2_scale=exp2_scale, + max_mode=max_mode, + sum_mode=sum_mode, + seq_len=seq_q, + block_q=block_q, + ) + # The query-tile axis is innermost, so the K/V block index is unchanged + # across it and each head block's KV is fetched from HBM once. + return pl.pallas_call( + kernel, + grid=(batch, heads // head_block, num_q_blocks), + in_specs=[ + pl.BlockSpec((1, block_q, head_block * dim_head), lambda b, h, i: (b, i, h)), + pl.BlockSpec((1, seq_kv, head_block * dim_head), lambda b, h, i: (b, 0, h)), + pl.BlockSpec((1, seq_kv, head_block * dim_head), lambda b, h, i: (b, 0, h)), + ], + out_specs=pl.BlockSpec((1, block_q, head_block * dim_head), lambda b, h, i: (b, i, h)), + out_shape=jax.ShapeDtypeStruct(q.shape, q.dtype), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "parallel", "arbitrary"), vmem_limit_bytes=resolved_vmem_limit_bytes + ), + cost_estimate=cost, + interpret=interpret, + )(q, k, v) + + +def cross_attention_reference( + q: jax.Array, + k: jax.Array, + v: jax.Array, + *, + heads: int, + dim_head: int = 128, + scale: float | None = None, + k_prescaled: bool = False, +) -> jax.Array: + """FP32 golden reference: exact softmax attention, rounded once at the end.""" + if scale is None: + scale = dim_head**-0.5 + batch, seq_q, _ = q.shape + seq_kv = k.shape[1] + q4 = q.reshape(batch, seq_q, heads, dim_head).astype(jnp.float32) + k4 = k.reshape(batch, seq_kv, heads, dim_head).astype(jnp.float32) + v4 = v.reshape(batch, seq_kv, heads, dim_head).astype(jnp.float32) + s = jnp.einsum("bqhd,bkhd->bhqk", q4, k4, precision=jax.lax.Precision.HIGHEST) + if not k_prescaled: + s = s * scale + p = jax.nn.softmax(s, axis=-1) + o = jnp.einsum("bhqk,bkhd->bqhd", p, v4, precision=jax.lax.Precision.HIGHEST) + return o.reshape(batch, seq_q, heads * dim_head).astype(q.dtype) diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 477bd3d7b..77ccdb9a0 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -36,6 +36,7 @@ rope_accum_is_measured, with_xla_backward, ) +from maxdiffusion.kernels import cross_attention_pallas from einops import rearrange from .. import common_types, max_logging from maxdiffusion.tpu_utils import get_tpu_type, TpuType @@ -2756,6 +2757,39 @@ def cudnn_flash_te_kernel(q, k, v, context): "ulysses_custom_fixed_m_per_q_block", } +_FLASH_MIN_SEQ_LENGTH_KERNELS = frozenset({ + "flash", + "tokamax_flash", + "ulysses", + "ulysses_custom", + "ulysses_custom_fixed_m", + "ulysses_custom_fixed_m_per_q_block", + "ulysses_ring", + "ulysses_ring_custom", + "ulysses_ring_custom_fixed_m", + "ulysses_ring_custom_fixed_m_per_q_block", + "ulysses_ring_custom_bidir", +}) + + +def _would_use_unmasked_dot_product( + attention_kernel: str, + use_memory_efficient_attention: bool, + flash_min_seq_length: int, + q_len: int, + kv_len: int, + v_len: Optional[int] = None, +) -> bool: + """Returns True when `_apply_attention` routes to unmasked dot-product attention.""" + if use_memory_efficient_attention: + return False + if attention_kernel == "dot_product": + return True + if attention_kernel in _FLASH_MIN_SEQ_LENGTH_KERNELS: + min_len = min(q_len, kv_len) if v_len is None else min(q_len, kv_len, v_len) + return min_len < flash_min_seq_length + return False + def _apply_attention( query: Array, @@ -2800,19 +2834,7 @@ def _apply_attention( seq_len_idx = 2 can_use_flash_attention = True - if attention_kernel in [ - "flash", - "tokamax_flash", - "ulysses", - "ulysses_custom", - "ulysses_custom_fixed_m", - "ulysses_custom_fixed_m_per_q_block", - "ulysses_ring", - "ulysses_ring_custom", - "ulysses_ring_custom_fixed_m", - "ulysses_ring_custom_fixed_m_per_q_block", - "ulysses_ring_custom_bidir", - ]: + if attention_kernel in _FLASH_MIN_SEQ_LENGTH_KERNELS: can_use_flash_attention = ( query.shape[seq_len_idx] >= flash_min_seq_length and key.shape[seq_len_idx] >= flash_min_seq_length @@ -3484,6 +3506,61 @@ def xla_equivalent(q, k, q_scale, k_scale, freqs): return with_xla_backward(sharded, xla_equivalent) +@functools.lru_cache(maxsize=64) +def _build_sharded_pallas_cross_attention( + kernel_fn: Callable, + mesh: jax.sharding.Mesh, + q_spec: jax.sharding.PartitionSpec, + kv_spec: jax.sharding.PartitionSpec, + local_heads: int, + global_heads: int, + dim_head: int, + scale: float, + k_prescaled: bool, + head_block: int, + is_cpu_interpret: bool, + dtype: DType, + split_head_dim: bool, + float32_qk_product: bool, +): + """Builds and caches the `with_xla_backward(jax.shard_map(...))` cross-attention wrapper.""" + sharded = jax.shard_map( + functools.partial( + kernel_fn, + heads=local_heads, + dim_head=dim_head, + scale=scale, + k_prescaled=k_prescaled, + head_block=head_block, + mesh=mesh, + interpret=is_cpu_interpret, + ), + mesh=mesh, + in_specs=(q_spec, kv_spec, kv_spec), + out_specs=q_spec, + check_vma=False, + ) + + def xla_equivalent(q, k, v): + """The XLA dot-product path this kernel stands in for.""" + return _apply_attention_dot( + q, + k, + v, + dtype, + global_heads, + dim_head, + scale, + split_head_dim, + float32_qk_product, + False, + k_prescaled=k_prescaled, + ) + + # `pallas_call` has no transpose rule; forward stays the kernel. + return with_xla_backward(sharded, xla_equivalent) + + class FlaxWanAttention(nnx.Module): def __init__( @@ -3538,6 +3615,8 @@ def __init__( "wan_splash_transpose_out": None, "wan_cross_attn_prescale_kv": None, "wan_rope_accum": None, + "wan_cross_attn_kernel": None, + "wan_cross_attn_cpu_interpret": False, **(attention_config or {}), } @@ -3569,6 +3648,8 @@ def _wan_opt(name): self.wan_splash_transpose_out = _wan_opt("wan_splash_transpose_out") self.wan_cross_attn_prescale_kv = _wan_opt("wan_cross_attn_prescale_kv") self.wan_rope_accum = _wan_opt("wan_rope_accum") + self.wan_cross_attn_kernel = _wan_opt("wan_cross_attn_kernel") + self.wan_cross_attn_cpu_interpret = _wan_opt("wan_cross_attn_cpu_interpret") if attention_kernel in {"flash", "cudnn_flash_te"} and mesh is None: raise ValueError(f"The flash attention kernel requires a value for mesh, but mesh is {self.mesh}") @@ -3974,6 +4055,125 @@ def producer(q, k, q_scale, k_scale, freqs, *, q_heads, dim_head, eps, k_eps=Non return producer + def _pallas_cross_attention(self, query, key, value, *, k_prescaled: bool): + """Runs text cross-attention through the fused Pallas kernel, or returns None. + + Opt-in via `wan_cross_attn_kernel: "pallas"`. It replaces only what would + otherwise be the XLA dot-product fallback -- unmasked MHA against a KV + shorter than `flash_min_seq_length` -- because the kernel holds a head + block's entire KV in VMEM. Returns None, meaning "use the XLA path", unless: + * the mesh devices are TPUs (or `wan_cross_attn_cpu_interpret` is enabled) + and a mesh is available for `shard_map`; + * the inputs are flat `[B, S, H*D]` with a lane-aligned `dim_head`; + * K/V carry the same heads as Q and are short enough to stay resident; + * sequence and heads divide evenly by the mesh axes they are sharded over + (an indivisible batch dimension is treated as replicated); + * on TPU, every mesh device is a generation the kernel's VMEM budget was + validated on (v6e, tpu7x) and the estimated working set fits it. + """ + mode = self.wan_cross_attn_kernel + if mode == "xla": + return None + if mode != "pallas": + raise ValueError(f"wan_cross_attn_kernel must be 'xla' or 'pallas', got {mode!r}.") + + op = self.attention_op + mesh = self.mesh + is_tpu = mesh is not None and all(getattr(d, "platform", None) == "tpu" for d in mesh.devices.flat) + is_cpu_interpret = (not is_tpu) and bool(self.wan_cross_attn_cpu_interpret) + flat = query.ndim == 3 and key.ndim == 3 and value.ndim == 3 + kv_len = key.shape[1] if flat else -1 + would_use_dot_product = flat and _would_use_unmasked_dot_product( + op.attention_kernel, + op.use_memory_efficient_attention, + op.flash_min_seq_length, + query.shape[1], + kv_len, + value.shape[1], + ) + + def _shards_over(axis) -> int: + if axis is None: + return 1 + names = (axis,) if isinstance(axis, str) else tuple(axis) + return math.prod(mesh.shape[n] for n in names) + + reason = None + if not (is_tpu or is_cpu_interpret): + reason = "the mesh devices are not TPUs" + elif not flat: + reason = f"inputs are not flat [B, S, H*D] (q={query.shape}, k={key.shape}, v={value.shape})" + elif self.dim_head % cross_attention_pallas.NUM_LANES: + reason = f"dim_head={self.dim_head} is not a multiple of {cross_attention_pallas.NUM_LANES}" + elif getattr(op, "kv_heads", None) not in (None, self.heads): + reason = f"kv_heads={op.kv_heads} differs from heads={self.heads}" + elif kv_len > cross_attention_pallas.MAX_KV_LEN: + reason = f"KV length {kv_len} exceeds {cross_attention_pallas.MAX_KV_LEN}" + elif kv_len % 8 != 0: + reason = f"KV length {kv_len} is not a multiple of 8" + elif not would_use_dot_product: + reason = f"the configured {op.attention_kernel!r} kernel already handles this KV length" + if reason is None: + batch_axis, seq_axis, head_axis = nn.logical_to_mesh_axes((BATCH, LENGTH, HEAD)) + batch_shards, seq_shards, head_shards = (_shards_over(a) for a in (batch_axis, seq_axis, head_axis)) + if query.shape[1] % seq_shards or self.heads % head_shards: + reason = f"q={query.shape} with {self.heads} heads does not split {seq_shards}x{head_shards} (seq x heads)" + else: + local_q_len = query.shape[1] // seq_shards + if local_q_len <= cross_attention_pallas.DEFAULT_BLOCK_Q and local_q_len % 8 != 0: + reason = f"local query length {local_q_len} <= {cross_attention_pallas.DEFAULT_BLOCK_Q} is not a multiple of 8" + if reason is not None: + _warn_once("pallas_cross_attn_unusable", f"Pallas cross-attention requested but unusable: {reason}; using XLA.") + return None + + # As in `_fused_rope_producer`: GSPMD degenerates an uneven batch split + # into replication, so declare it replicated rather than lose the kernel. + local_batch_axis = batch_axis if query.shape[0] % batch_shards == 0 else None + local_heads = self.heads // head_shards + head_block = max( + hb for hb in range(1, min(local_heads, cross_attention_pallas.DEFAULT_HEAD_BLOCK) + 1) if local_heads % hb == 0 + ) + if not is_cpu_interpret: + # Fall back rather than hand Mosaic a working set it cannot hold. + vmem_reason = cross_attention_pallas.vmem_unusable_reason( + seq_q=local_q_len, + seq_kv=kv_len, + dim_head=self.dim_head, + head_block=head_block, + itemsize=jnp.dtype(query.dtype).itemsize, + mesh=mesh, + ) + if vmem_reason is not None: + _warn_once("pallas_cross_attn_vmem", f"Pallas cross-attention requested but unusable: {vmem_reason}; using XLA.") + return None + q_spec = jax.sharding.PartitionSpec(local_batch_axis, seq_axis, head_axis) + kv_spec = jax.sharding.PartitionSpec(local_batch_axis, None, head_axis) + _warn_once( + f"pallas_cross_attn_active:{tuple(query.shape)}", + f"Pallas cross-attention ACTIVE: q={query.shape}, kv={key.shape}, q_spec={q_spec}, kv_spec={kv_spec}, " + f"local_heads={local_heads}, head_block={head_block}, block_q={cross_attention_pallas.DEFAULT_BLOCK_Q}, " + f"k_prescaled={k_prescaled}.", + ) + # `pallas_call` has no transpose rule; forward stays the kernel. + differentiable = _build_sharded_pallas_cross_attention( + cross_attention_pallas.cross_attention_pallas, + mesh, + q_spec, + kv_spec, + local_heads, + self.heads, + self.dim_head, + float(op.scale), + bool(k_prescaled), + head_block, + bool(is_cpu_interpret), + op.dtype, + bool(op.split_head_dim), + bool(op.float32_qk_product), + ) + with jax.named_scope("kernel_pallas_cross_attention"): + return differentiable(query, key, value) + def conditional_named_scope(self, name: str): """Return a JAX named scope if enabled, otherwise a null context.""" return jax.named_scope(name) if self.enable_jax_named_scopes else contextlib.nullcontext() @@ -4113,7 +4313,11 @@ def run_attention(**svg_kwargs): ) else: with jax.named_scope("apply_attention"): - attn_output = run_attention() + attn_output = None + if not is_self_attention and encoder_attention_mask is None and not qk_prescaled: + attn_output = self._pallas_cross_attention(query_proj, key_proj, value_proj, k_prescaled=k_prescaled) + if attn_output is None: + attn_output = run_attention() else: # NEW PATH for I2V CROSS-ATTENTION @@ -4190,9 +4394,13 @@ def run_attention(**svg_kwargs): # Attention - tensors are (B, S, D) with self.conditional_named_scope("cross_attn_text_apply"): - attn_output_text = self.attention_op.apply_attention( + attn_output_text = self._pallas_cross_attention( query_proj_text, key_proj_text, value_proj_text, k_prescaled=k_prescaled_text ) + if attn_output_text is None: + attn_output_text = self.attention_op.apply_attention( + query_proj_text, key_proj_text, value_proj_text, k_prescaled=k_prescaled_text + ) with self.conditional_named_scope("cross_attn_img_apply"): # Pass encoder_attention_mask_img for image cross-attention to mask padded tokens attn_output_img = self.attention_op.apply_attention( @@ -4211,9 +4419,13 @@ def run_attention(**svg_kwargs): value_proj_text = checkpoint_name(value_proj_text, "value_proj_text") with self.conditional_named_scope("cross_attn_text_apply"): - attn_output = self.attention_op.apply_attention( + attn_output = self._pallas_cross_attention( query_proj_text, key_proj_text, value_proj_text, k_prescaled=k_prescaled_text ) + if attn_output is None: + attn_output = self.attention_op.apply_attention( + query_proj_text, key_proj_text, value_proj_text, k_prescaled=k_prescaled_text + ) attn_output = attn_output.astype(dtype=dtype) attn_output = checkpoint_name(attn_output, "attn_output") diff --git a/src/maxdiffusion/tests/cross_attention_pallas_test.py b/src/maxdiffusion/tests/cross_attention_pallas_test.py new file mode 100644 index 000000000..0ce954523 --- /dev/null +++ b/src/maxdiffusion/tests/cross_attention_pallas_test.py @@ -0,0 +1,433 @@ +""" +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. +""" + +"""Tests for the fused short-KV cross-attention Pallas kernel. + +The kernel is not bit-identical to the XLA dot-product path it replaces (it +normalises the softmax after PV), so parity is asserted against an FP32 golden +reference, together with the stronger property that the kernel is at least as +close to that reference as the XLA path is. Kernel and `FlaxWanAttention` +integration tests run off-TPU through Pallas interpret mode (with +`wan_cross_attn_cpu_interpret=True` for the module integration test, since +production falls back to XLA off-TPU). +""" + +import os +import types +import unittest +from unittest import mock + +import jax +import jax.numpy as jnp +import numpy as np +from absl.testing import parameterized +from flax import nnx +from flax.linen import partitioning as nn_partitioning +from jax.sharding import Mesh + +from maxdiffusion import pyconfig +from maxdiffusion.kernels import cross_attention_pallas +from maxdiffusion.kernels.cross_attention_pallas import cross_attention_pallas as xattn +from maxdiffusion.kernels.cross_attention_pallas import cross_attention_reference +from maxdiffusion.max_utils import create_device_mesh, get_flash_block_sizes +from maxdiffusion.models.attention_flax import FlaxWanAttention, _apply_attention_dot + +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" +) + +DIM_HEAD = 128 +THIS_DIR = os.path.dirname(os.path.abspath(__file__)) + + +def _on_tpu() -> bool: + return jax.devices()[0].platform == "tpu" + + +def _on_unvalidated_tpu() -> bool: + """On a TPU whose VMEM budget is not validated, FlaxWanAttention falls back to XLA by design.""" + return _on_tpu() and not cross_attention_pallas.vmem_budget_is_validated() + + +_UNVALIDATED_TPU_SKIP = ( + "Pallas cross-attention falls back to XLA on TPU kinds outside _VMEM_LIMIT_VALIDATED_KINDS " + "(covered by test_falls_back_to_xla_on_unvalidated_tpu_kind)." +) + + +def _inputs(batch, seq_q, seq_kv, heads, *, qk_gain=2.0, seed=0): + """bf16 q/k/v; `qk_gain` sets the logit spread (post-scale std ~ qk_gain**2).""" + kq, kk, kv = jax.random.split(jax.random.PRNGKey(seed), 3) + feature = heads * DIM_HEAD + q = (qk_gain * jax.random.normal(kq, (batch, seq_q, feature), jnp.float32)).astype(jnp.bfloat16) + k = (qk_gain * jax.random.normal(kk, (batch, seq_kv, feature), jnp.float32)).astype(jnp.bfloat16) + v = jax.random.normal(kv, (batch, seq_kv, feature), jnp.float32).astype(jnp.bfloat16) + return q, k, v + + +def _rel_err(got, ref): + got = np.asarray(got, np.float32) + ref = np.asarray(ref, np.float32) + return float(np.linalg.norm(got - ref) / np.linalg.norm(ref)) + + +def _xla_path(q, k, v, heads, k_prescaled=False): + """The XLA dot-product fallback the kernel replaces, as Wan configures it.""" + return _apply_attention_dot( + q, k, v, jnp.bfloat16, heads, DIM_HEAD, DIM_HEAD**-0.5, True, False, False, None, k_prescaled=k_prescaled + ) + + +@_SKIP_IN_GITHUB_ACTIONS +class CrossAttentionPallasParityTest(parameterized.TestCase): + + def _check(self, got, golden, rel_tol=5e-3): + got_np = np.asarray(got, np.float32) + self.assertTrue(np.all(np.isfinite(got_np)), "kernel output contains non-finite values") + np.testing.assert_allclose(got_np, np.asarray(golden, np.float32), atol=3e-2, rtol=3e-2) + self.assertLess(_rel_err(got, golden), rel_tol) + + @parameterized.product( + heads=(4, 5), + sum_mode=cross_attention_pallas.SUM_MODES, + max_mode=cross_attention_pallas.MAX_MODES, + ) + def test_matches_golden(self, heads, sum_mode, max_mode): + q, k, v = _inputs(1, 256, 128, heads) + head_block = 1 if heads == 5 else 2 + got = xattn( + q, + k, + v, + heads=heads, + block_q=128, + head_block=head_block, + sum_mode=sum_mode, + max_mode=max_mode, + interpret=not _on_tpu(), + ) + self.assertEqual(got.shape, q.shape) + self.assertEqual(got.dtype, q.dtype) + self._check(got, cross_attention_reference(q, k, v, heads=heads)) + + @parameterized.parameters(*cross_attention_pallas.MAX_MODES) + def test_ragged_last_query_block(self, max_mode): + # 300 rows run as three 112-row tiles, so the last reads 36 rows of stale + # VMEM, which must neither leak into valid rows nor reach the output. + q, k, v = _inputs(1, 300, 64, 2, seed=1) + got = xattn(q, k, v, heads=2, block_q=128, head_block=2, max_mode=max_mode, interpret=not _on_tpu()) + self._check(got, cross_attention_reference(q, k, v, heads=2)) + + def test_result_is_independent_of_tiling(self): + # Rows and heads never interact, so the tiling must not change a single bit. + q, k, v = _inputs(1, 256, 128, 4, seed=2) + outs = [ + np.asarray(xattn(q, k, v, heads=4, block_q=bq, head_block=hb, interpret=not _on_tpu()), np.float32) + for bq, hb in ((256, 4), (128, 2), (64, 1), (96, 4)) + ] + for other in outs[1:]: + np.testing.assert_array_equal(outs[0], other) + + def test_at_least_as_accurate_as_the_xla_path(self): + heads = 4 + q, k, v = _inputs(1, 512, 256, heads, seed=3) + golden = cross_attention_reference(q, k, v, heads=heads) + kernel_err = _rel_err(xattn(q, k, v, heads=heads, block_q=128, head_block=2, interpret=not _on_tpu()), golden) + xla_err = _rel_err(_xla_path(q, k, v, heads), golden) + self.assertLessEqual(kernel_err, 1.1 * xla_err, f"kernel rel err {kernel_err:.3e} vs XLA {xla_err:.3e}") + + def test_k_prescaled(self): + heads = 2 + q, k, v = _inputs(1, 128, 64, heads, seed=4) + k_scaled = k * jnp.asarray(DIM_HEAD**-0.5, k.dtype) + got = xattn(q, k_scaled, v, heads=heads, block_q=64, head_block=2, k_prescaled=True, interpret=not _on_tpu()) + self._check(got, cross_attention_reference(q, k_scaled, v, heads=heads, k_prescaled=True)) + + def test_multi_batch(self): + q, k, v = _inputs(2, 128, 64, 2, seed=5) + got = xattn(q, k, v, heads=2, block_q=64, head_block=1, interpret=not _on_tpu()) + self._check(got, cross_attention_reference(q, k, v, heads=2)) + + def test_large_logits_stay_finite(self): + # Post-scale logits reach several hundred: exp overflows without the shift. + q, k, v = _inputs(1, 128, 64, 2, qk_gain=30.0, seed=6) + got = xattn(q, k, v, heads=2, block_q=64, head_block=2, interpret=not _on_tpu()) + self._check(got, cross_attention_reference(q, k, v, heads=2), rel_tol=2e-2) + + @unittest.skipUnless(_on_tpu(), "Production-shape VMEM tile execution is TPU-only.") + @unittest.skipIf(_on_unvalidated_tpu(), _UNVALIDATED_TPU_SKIP) + def test_wan_production_shard_shapes(self): + # Exercises the actual 42 MiB VMEM tile shape (2368x512, head_block=10, 40 heads) + # on both unpadded (18900, ragged last tile) and lane-padded (18944, aligned) shards. + heads = 40 + _, k, v = _inputs(1, 16, 512, heads, seed=7) + for seq_q in (18900, 18944): + with self.subTest(seq_q=seq_q): + q, _, _ = _inputs(1, seq_q, 16, heads, seed=8) + got = xattn(q, k, v, heads=heads, interpret=False) + self.assertEqual(got.shape, q.shape) + self._check(got, cross_attention_reference(q, k, v, heads=heads)) + + +class CrossAttentionPallasGuardTest(unittest.TestCase): + + def test_rejects_unaligned_dim_head(self): + q = jnp.zeros((1, 64, 128), jnp.bfloat16) + with self.assertRaisesRegex(ValueError, "multiple of 128"): + xattn(q, q, q, heads=2, dim_head=64, interpret=True) + + def test_rejects_head_block_not_dividing_heads(self): + q, k, v = _inputs(1, 64, 64, 4) + with self.assertRaisesRegex(ValueError, "must divide heads"): + xattn(q, k, v, heads=4, head_block=3, interpret=True) + + def test_rejects_mismatched_kv(self): + q, k, _ = _inputs(1, 64, 64, 2) + with self.assertRaisesRegex(ValueError, "k and v"): + xattn(q, k, k[:, :32], heads=2, head_block=2, interpret=True) + + def test_rejects_long_kv(self): + q, k, v = _inputs(1, 64, cross_attention_pallas.MAX_KV_LEN + 8, 1) + with self.assertRaisesRegex(ValueError, "MAX_KV_LEN"): + xattn(q, k, v, heads=1, head_block=1, interpret=True) + + def test_rejects_unaligned_block_q(self): + q, k, v = _inputs(1, 64, 64, 1) + with self.assertRaisesRegex(ValueError, "multiple of 8"): + xattn(q, k, v, heads=1, head_block=1, block_q=20, interpret=True) + + def test_rejects_unaligned_seq_q_when_block_q_covers_seq_q(self): + q, k, v = _inputs(1, 30, 64, 1) + with self.assertRaisesRegex(ValueError, "multiple of 8"): + xattn(q, k, v, heads=1, head_block=1, block_q=64, interpret=True) + + def test_rejects_unaligned_kv_len(self): + q, k, v = _inputs(1, 64, 60, 1) + with self.assertRaisesRegex(ValueError, "multiple of 8"): + xattn(q, k, v, heads=1, head_block=1, block_q=64, interpret=True) + + def test_rejects_unknown_modes(self): + q, k, v = _inputs(1, 64, 64, 1) + with self.assertRaisesRegex(ValueError, "max_mode"): + xattn(q, k, v, heads=1, head_block=1, max_mode="bogus", interpret=True) + with self.assertRaisesRegex(ValueError, "sum_mode"): + xattn(q, k, v, heads=1, head_block=1, sum_mode="bogus", interpret=True) + + +class CrossAttentionVmemBudgetTest(unittest.TestCase): + """VMEM budget resolution and the working-set estimate that gates the kernel.""" + + WAN_SHAPE = {"seq_q": 18900, "seq_kv": 512, "dim_head": 128} + + @staticmethod + def _mesh(*kinds, platform="tpu"): + devices = np.array([types.SimpleNamespace(platform=platform, device_kind=k) for k in kinds], dtype=object) + return types.SimpleNamespace(devices=devices) + + def test_explicit_value_wins(self): + self.assertEqual(cross_attention_pallas._resolve_vmem_limit_bytes(123, self._mesh("TPU v4")), 123) + + def test_validated_kinds_get_the_default(self): + for kinds in (("TPU v6 lite",) * 2, ("TPU7x",) * 2): + with self.subTest(kinds=kinds): + self.assertEqual( + cross_attention_pallas._resolve_vmem_limit_bytes(None, self._mesh(*kinds)), + cross_attention_pallas.DEFAULT_VMEM_LIMIT_BYTES, + ) + + def test_unvalidated_or_mixed_kinds_get_mosaic_default(self): + for kinds in (("TPU v4",), ("TPU v5 lite",), ("TPU v5p",), ("TPU v6 lite", "TPU v4")): + with self.subTest(kinds=kinds): + self.assertIsNone(cross_attention_pallas._resolve_vmem_limit_bytes(None, self._mesh(*kinds))) + + def test_off_tpu_returns_default(self): + self.assertEqual( + cross_attention_pallas._resolve_vmem_limit_bytes(None, self._mesh("cpu", platform="cpu")), + cross_attention_pallas.DEFAULT_VMEM_LIMIT_BYTES, + ) + + def test_estimate_at_wan_shape_needs_the_raised_budget(self): + need = cross_attention_pallas.estimate_vmem_bytes(**self.WAN_SHAPE) + # Above Mosaic's default scoped limit (<= 32 MiB on current TPUs), within 64 MiB. + self.assertGreater(need, 32 * 2**20) + self.assertLess(need, cross_attention_pallas.DEFAULT_VMEM_LIMIT_BYTES) + + def test_estimate_grows_with_tiles_and_kv(self): + base = cross_attention_pallas.estimate_vmem_bytes(**self.WAN_SHAPE) + self.assertLess(cross_attention_pallas.estimate_vmem_bytes(**self.WAN_SHAPE, block_q=1024), base) + self.assertLess(cross_attention_pallas.estimate_vmem_bytes(**self.WAN_SHAPE, head_block=5), base) + self.assertLess(cross_attention_pallas.estimate_vmem_bytes(**self.WAN_SHAPE, sum_mode="vpu"), base) + self.assertGreater(cross_attention_pallas.estimate_vmem_bytes(**{**self.WAN_SHAPE, "seq_kv": 4096}), base) + + def test_unusable_reason(self): + reason = cross_attention_pallas.vmem_unusable_reason + self.assertIsNone(reason(**self.WAN_SHAPE, mesh=self._mesh("TPU v6 lite"))) + self.assertIsNone(reason(**self.WAN_SHAPE, mesh=self._mesh("TPU7x"))) + self.assertRegex(reason(**self.WAN_SHAPE, mesh=self._mesh("TPU v5 lite")), "not validated") + # MAX_KV_LEN KV at the default tiles does not fit 64 MiB: fall back, don't fail to compile. + long_kv = {**self.WAN_SHAPE, "seq_kv": cross_attention_pallas.MAX_KV_LEN} + self.assertRegex(reason(**long_kv, mesh=self._mesh("TPU v6 lite")), "exceeds") + # An explicit budget is trusted on any TPU. + self.assertIsNone(reason(**self.WAN_SHAPE, mesh=self._mesh("TPU v5 lite"), vmem_limit_bytes=96 * 2**20)) + + def test_kernel_rejects_a_working_set_over_an_explicit_budget(self): + q, k, v = _inputs(1, 64, 64, 1) + with self.assertRaisesRegex(ValueError, "VMEM working set"): + xattn(q, k, v, heads=1, head_block=1, block_q=64, vmem_limit_bytes=1024, interpret=False) + + +class WanCrossAttentionDispatchTest(unittest.TestCase): + """`wan_cross_attn_kernel` routing inside `FlaxWanAttention`.""" + + HEADS = 8 + SEQ_Q = 2048 + SEQ_KV = 512 + + def setUp(self): + super().setUp() + pyconfig.initialize([None, os.path.join(THIS_DIR, "..", "configs", "base_wan_14b.yml")], unittest=True) + self.config = pyconfig.config + self.mesh = Mesh(create_device_mesh(self.config), self.config.mesh_axes) + + def _run(self, mode, seq_q=None, seq_kv=None, cpu_interpret=False): + seq_q = self.SEQ_Q if seq_q is None else seq_q + seq_kv = self.SEQ_KV if seq_kv is None else seq_kv + query_dim = self.HEADS * DIM_HEAD + k1, k2, k3 = jax.random.split(jax.random.PRNGKey(7), 3) + with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + attn = FlaxWanAttention( + rngs=nnx.Rngs(k1), + query_dim=query_dim, + heads=self.HEADS, + dim_head=DIM_HEAD, + attention_kernel="dot_product", + split_head_dim=True, + mesh=self.mesh, + flash_block_sizes=get_flash_block_sizes(self.config), + is_self_attention=False, + dtype=jnp.bfloat16, + attention_config={"wan_cross_attn_kernel": mode, "wan_cross_attn_cpu_interpret": cpu_interpret}, + ) + hidden = jax.random.normal(k2, (1, seq_q, query_dim), jnp.bfloat16) + context = jax.random.normal(k3, (1, seq_kv, query_dim), jnp.bfloat16) + return attn(hidden_states=hidden, encoder_hidden_states=context) + + def test_rejects_unknown_mode(self): + with self.assertRaisesRegex(ValueError, "wan_cross_attn_kernel"): + self._run("bogus") + + def test_xla_mode_never_calls_the_kernel(self): + with mock.patch.object(cross_attention_pallas, "cross_attention_pallas", wraps=xattn) as spy: + self._run("xla") + spy.assert_not_called() + + @unittest.skipIf(_on_tpu(), "Off-TPU fallback only.") + def test_falls_back_to_xla_off_tpu(self): + with mock.patch.object(cross_attention_pallas, "cross_attention_pallas", wraps=xattn) as spy: + got = self._run("pallas") + spy.assert_not_called() + np.testing.assert_array_equal(np.asarray(got, np.float32), np.asarray(self._run("xla"), np.float32)) + + @unittest.skipUnless(_on_tpu(), "The Pallas path is TPU-only.") + @unittest.skipIf(_on_unvalidated_tpu(), _UNVALIDATED_TPU_SKIP) + @_SKIP_IN_GITHUB_ACTIONS + def test_pallas_mode_uses_the_kernel_and_matches_xla(self): + with mock.patch.object(cross_attention_pallas, "cross_attention_pallas", wraps=xattn) as spy: + got = self._run("pallas") + spy.assert_called() + ref = self._run("xla") + got_np = np.asarray(got, np.float32) + self.assertTrue(np.all(np.isfinite(got_np))) + np.testing.assert_allclose(got_np, np.asarray(ref, np.float32), atol=0.08, rtol=1e-2) + self.assertLess(_rel_err(got, ref), 1e-2) + + @unittest.skipIf(_on_unvalidated_tpu(), _UNVALIDATED_TPU_SKIP) + def test_pallas_mode_cpu_interpret_integration(self): + """Exercises FlaxWanAttention routing with Pallas cross-attention in CI under interpret mode.""" + with mock.patch.object(cross_attention_pallas, "cross_attention_pallas", wraps=xattn) as spy: + got = self._run("pallas", seq_q=128, seq_kv=64, cpu_interpret=True) + spy.assert_called() + ref = self._run("xla", seq_q=128, seq_kv=64) + got_np = np.asarray(got, np.float32) + self.assertTrue(np.all(np.isfinite(got_np))) + np.testing.assert_allclose(got_np, np.asarray(ref, np.float32), atol=0.08, rtol=1e-2) + + @unittest.skipUnless(_on_tpu(), "VMEM gating only applies on TPU.") + def test_falls_back_to_xla_on_unvalidated_tpu_kind(self): + with mock.patch.object(cross_attention_pallas, "_VMEM_LIMIT_VALIDATED_KINDS", ("no-such-tpu",)): + with mock.patch.object(cross_attention_pallas, "cross_attention_pallas", wraps=xattn) as spy: + got = self._run("pallas") + spy.assert_not_called() + np.testing.assert_array_equal(np.asarray(got, np.float32), np.asarray(self._run("xla"), np.float32)) + + def test_falls_back_to_xla_on_unaligned_kv_or_short_unaligned_q(self): + with mock.patch.object(cross_attention_pallas, "cross_attention_pallas", wraps=xattn) as spy: + self._run("pallas", seq_q=2048, seq_kv=769, cpu_interpret=True) + self._run("pallas", seq_q=300, seq_kv=512, cpu_interpret=True) + spy.assert_not_called() + + @unittest.skipIf(_on_unvalidated_tpu(), _UNVALIDATED_TPU_SKIP) + @_SKIP_IN_GITHUB_ACTIONS + def test_i2v_cross_attention_uses_pallas_for_text_and_xla_for_image(self): + """I2V splits encoder_hidden_states into padded image KV (masked XLA) and text KV (Pallas).""" + query_dim = self.HEADS * DIM_HEAD + k1, k2, k3 = jax.random.split(jax.random.PRNGKey(17), 3) + with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + attn = FlaxWanAttention( + rngs=nnx.Rngs(k1), + query_dim=query_dim, + heads=self.HEADS, + dim_head=DIM_HEAD, + attention_kernel="dot_product", + split_head_dim=True, + mesh=self.mesh, + flash_block_sizes=get_flash_block_sizes(self.config), + is_self_attention=False, + added_kv_proj_dim=query_dim, + image_seq_len=257, + dtype=jnp.bfloat16, + ) + padded_img_len = ((257 + attn.alignment - 1) // attn.alignment) * attn.alignment + total_kv_len = padded_img_len + self.SEQ_KV + hidden = jax.random.normal(k2, (1, self.SEQ_Q, query_dim), jnp.bfloat16) + context = jax.random.normal(k3, (1, total_kv_len, query_dim), jnp.bfloat16) + img_mask = jnp.concatenate( + [jnp.ones((1, 257), dtype=jnp.int32), jnp.zeros((1, padded_img_len - 257), dtype=jnp.int32)], axis=1 + ) + txt_mask = jnp.ones((1, self.SEQ_KV), dtype=jnp.int32) + mask = jnp.concatenate([img_mask, txt_mask], axis=1) + + # The routing switches are build-time attributes (from attention_config). + attn.wan_cross_attn_kernel = "xla" + ref = attn(hidden_states=hidden, encoder_hidden_states=context, encoder_attention_mask=mask) + + attn.wan_cross_attn_kernel = "pallas" + attn.wan_cross_attn_cpu_interpret = True + with mock.patch.object(cross_attention_pallas, "cross_attention_pallas", wraps=xattn) as spy: + got = attn(hidden_states=hidden, encoder_hidden_states=context, encoder_attention_mask=mask) + self.assertEqual(spy.call_count, 1) + self.assertEqual(spy.call_args.args[1].shape[1], self.SEQ_KV) + + got_np = np.asarray(got, np.float32) + self.assertTrue(np.all(np.isfinite(got_np))) + np.testing.assert_allclose(got_np, np.asarray(ref, np.float32), atol=0.08, rtol=1e-2) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/wan_runtime_options_test.py b/src/maxdiffusion/tests/wan_runtime_options_test.py index a479a1a41..cd1594dee 100644 --- a/src/maxdiffusion/tests/wan_runtime_options_test.py +++ b/src/maxdiffusion/tests/wan_runtime_options_test.py @@ -161,6 +161,18 @@ def test_wan_model_resolves_cfg_before_unpatchify(self): model = nnx.eval_shape(lambda: WanModel(rngs=nnx.Rngs(0), **kwargs, wan_cfg_before_unpatchify="false")) self.assertIs(model.wan_cfg_before_unpatchify, False) + def test_enum_options_validation(self): + valid_cases = { + "wan_cross_attn_kernel": ("xla", "pallas"), + } + for name, valid_values in valid_cases.items(): + for val in valid_values: + with self.subTest(option=name, valid=val): + self.assertEqual(opts.resolve_from_config(types.SimpleNamespace(**{name: val}), name), val) + with self.subTest(option=name, invalid="bogus"): + with self.assertRaisesRegex(ValueError, name): + opts.resolve_from_config(types.SimpleNamespace(**{name: "bogus"}), name) + if __name__ == "__main__": unittest.main() diff --git a/src/maxdiffusion/utils/wan_block_benchmark.py b/src/maxdiffusion/utils/wan_block_benchmark.py index a5e6e345d..3e5de7e74 100644 --- a/src/maxdiffusion/utils/wan_block_benchmark.py +++ b/src/maxdiffusion/utils/wan_block_benchmark.py @@ -234,10 +234,13 @@ def _build_model(self, bq, bkv, cmp): scan_layers=False, num_layers=1, enable_jax_named_scopes=c.enable_jax_named_scopes, + split_head_dim=getattr(c, "split_head_dim", True), attention_config={ "use_base2_exp": c.use_base2_exp, "use_experimental_scheduler": c.use_experimental_scheduler, "ulysses_shards": c.ulysses_shards, + "ulysses_attention_chunks": getattr(c, "ulysses_attention_chunks", 1), + "use_k_centering": getattr(c, "use_k_centering", "auto"), "use_fused_rope_kernel": getattr(c, "use_fused_rope_kernel", False), "fused_rope_block_s": getattr(c, "fused_rope_block_s", 1024), "fused_rope_head_block": None if fused_rope_head_block in (None, -1) else fused_rope_head_block, diff --git a/src/maxdiffusion/wan_runtime_options.py b/src/maxdiffusion/wan_runtime_options.py index fd6facdc0..bc611bb8a 100644 --- a/src/maxdiffusion/wan_runtime_options.py +++ b/src/maxdiffusion/wan_runtime_options.py @@ -53,6 +53,8 @@ "wan_cfg_before_unpatchify": ("WAN_CFG_BEFORE_UNPATCHIFY", True, "bool"), "wan_cross_attn_prescale_kv": ("WAN_CROSS_ATTN_PRESCALE_KV", False, "bool"), "wan_rope_accum": ("WAN_ROPE_ACCUM", "auto", "str"), + "wan_cross_attn_kernel": ("WAN_CROSS_ATTN_KERNEL", "xla", "str"), + "wan_cross_attn_cpu_interpret": ("WAN_CROSS_ATTN_CPU_INTERPRET", False, "bool"), } @@ -72,6 +74,8 @@ def _coerce(name: str, value: Any) -> Any: raise ValueError(f"wan_rope_norm_mode must be 'exact' or 'fused', got {value!r}.") if name == "wan_rope_accum" and val not in ("auto", "dtype", "f32"): raise ValueError(f"wan_rope_accum must be 'auto', 'dtype', or 'f32', got {value!r}.") + if name == "wan_cross_attn_kernel" and val not in ("xla", "pallas"): + raise ValueError(f"wan_cross_attn_kernel must be 'xla' or 'pallas', got {value!r}.") return val @@ -132,6 +136,8 @@ def resolve_from_config(config: Any, name: str) -> Any: "wan_splash_transpose_out", "wan_cross_attn_prescale_kv", "wan_rope_accum", + "wan_cross_attn_kernel", + "wan_cross_attn_cpu_interpret", )