From 1f5539b5ba713325192561057268a4de0d54e9e7 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Thu, 15 Oct 2026 05:47:45 +0000 Subject: [PATCH] feat(wan): short-KV text cross-attention Pallas kernel (wan_cross_attn_kernel="pallas") First of four pieces split out of the Wan fast-inference optimizations PR. The switch defaults off ("xla") in the Wan YAMLs and in wan_runtime_options; it is a plain config key plumbed to FlaxWanAttention through `attention_config_entries(config)`. - kernels/cross_attention_pallas.py keeps a head block's whole 512-token text KV resident in VMEM and computes an exact single-pass softmax per query tile (full-row max, exp2, one numerator/denominator, normalised after PV; no running max / online rescaling). It replaces only what would otherwise be the unmasked XLA dot-product fallback (`_would_use_unmasked_dot_product`, which shares `_FLASH_MIN_SEQ_LENGTH_KERNELS` with the dispatcher). - `FlaxWanAttention._pallas_cross_attention` falls back to XLA (logging the reason once) when the inputs or sharding don't fit, and also when the VMEM budget is not validated or too small: 64 MiB is requested only when every mesh TPU is v6e or tpu7x (otherwise Mosaic's default), and `estimate_vmem_bytes` (Q/out/K/V double buffers + one head's FP32 scores) is checked against it (~42 MiB at the Wan shard shape; MAX_KV_LEN KV does not fit and falls back rather than failing to compile). The estimate has only been checked against real compiles at Wan's KV length (512). The kernel itself raises a ValueError for an over-budget working set only when a budget resolves (an explicit vmem_limit_bytes, or the v6e/tpu7x default); under Mosaic's default it does no check of its own, and callers that can fall back use `vmem_unusable_reason`. The two fallback reasons (unusable inputs, VMEM) log under separate warn-once keys; the "ACTIVE" log is keyed on the query shape, so T2V and the I2V text branch are each reported. The sharded wrapper is cached (`_build_sharded_pallas_cross_attention`) and wrapped with `with_xla_backward` since `pallas_call` has no transpose rule. - Dispatch: T2V cross-attention, and the text half of I2V cross-attention (both the image-token split and the text-only path), try the kernel first and use `attention_op.apply_attention` when it returns None. - `wan_cross_attn_cpu_interpret` (default False, tests only) lets the kernel run under Pallas interpret mode on CPU. - In isolation the kernel is 0.51 ms vs 3.94 ms (v6e) and 0.49 ms vs 1.81 ms (tpu7x) per call for the XLA path inside the model. - run_wan_fast_inference.sh: the v6e and v7 profiles set WAN_CROSS_ATTN_KERNEL=pallas (overridable); the generic profile keeps xla. - wan_block_benchmark.py passes `split_head_dim`, `ulysses_attention_chunks` and `use_k_centering` from the config to the benchmarked block (shared benchmark-fidelity fix carried in this first piece; note it defaults `split_head_dim` to True when the config lacks the key, whereas the Wan ymls set False). Tests - cross_attention_pallas_test.py (32 test methods, 40 parameterized cases): golden parity (row/tile max, mxu/vpu sum, ragged tiles, k_prescaled, batch), guards, VMEM budget resolution / working-set estimate / fallback reasons, and FlaxWanAttention dispatch (CPU interpret integration, fallbacks, I2V text-vs-image split). The three tests that expect the kernel to be used skip on TPUs outside the validated kinds (there FlaxWanAttention falls back by design); that fallback is covered by test_falls_back_to_xla_on_unvalidated_tpu_kind. - wan_runtime_options_test.py (12): the new switch's YAML default and enum validation (test_enum_options_validation). - CI: kernel parity and the TPU Pallas-active dispatch tests skip in GitHub Actions (19 of the 40 cases); guards, VMEM budget logic, CPU interpret integration and the XLA fallback dispatch tests run in CI. - end_to_end/tpu/run_wan_stack_tests.sh: adds cross_attention_pallas_test.py. Verified (this head): TPU v6e-8 (jax 0.11.2), every file of end_to_end/tpu/run_wan_stack_tests.sh at this head run directly with unittest: 322 tests, all OK, 1 skipped (cross_attention_pallas_test: off-TPU fallback only); cross_attention_pallas_test.py 40 cases (1 skipped), wan_runtime_options_test.py 12. --- end_to_end/tpu/run_wan_fast_inference.sh | 7 + end_to_end/tpu/run_wan_stack_tests.sh | 2 + src/maxdiffusion/configs/base_wan_14b.yml | 6 + src/maxdiffusion/configs/base_wan_1_3b.yml | 6 + src/maxdiffusion/configs/base_wan_27b.yml | 6 + src/maxdiffusion/configs/base_wan_animate.yml | 6 + src/maxdiffusion/configs/base_wan_i2v_14b.yml | 6 + src/maxdiffusion/configs/base_wan_i2v_27b.yml | 6 + .../kernels/cross_attention_pallas.py | 430 +++++++++++++++++ src/maxdiffusion/models/attention_flax.py | 244 +++++++++- .../tests/cross_attention_pallas_test.py | 433 ++++++++++++++++++ .../tests/wan_runtime_options_test.py | 12 + src/maxdiffusion/utils/wan_block_benchmark.py | 3 + src/maxdiffusion/wan_runtime_options.py | 6 + 14 files changed, 1157 insertions(+), 16 deletions(-) create mode 100644 src/maxdiffusion/kernels/cross_attention_pallas.py create mode 100644 src/maxdiffusion/tests/cross_attention_pallas_test.py 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", )