From e38313f432e2bf43a29acd70b00e4a962dbacbb3 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Wed, 30 Sep 2026 10:41:37 +0000 Subject: [PATCH] feat(wan): fast inference optimizations (fused cross-attention, token patch embed, shard-major A2A, lane padding) All four switches default off ("xla" / "conv" / "flat" / "off") in the Wan YAMLs and in wan_runtime_options. They are plain config keys, plumbed to FlaxWanAttention / WanModel through `attention_config_entries(config)` (the environment at trace time. The unconditional pyconfig default for `wan_cross_attn_cpu_interpret` was removed (it is an attention_config entry, default False, used only by CPU tests). - Short-KV text cross-attention Pallas kernel (wan_cross_attn_kernel="pallas"): 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. FlaxWanAttention falls back to XLA (logging the reason once) when the inputs or sharding don't fit, and now 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. - Token-sharded patch embedding (wan_patch_embed_mode="tokens"): the Conv3d patch embedding becomes a patchify + matmul on the token-sharded layout, avoiding the activation all-gather. - Shard-major Ulysses A2A (wan_ulysses_out_a2a="chunked" / "shard_major"): a single layout-preserving all_to_all over a reshaped (num_shards, chunk) leading axis (not chunked compute/comm overlap), so XLA no longer inserts relayout copies around the collective. With the custom splash kernel (and S_local % block_q == 0, D == 128, heads_per_tile == 1) the forward Q exchange also stays shard-major (4-D q) and the kernel writes the shard-major output directly (out_num_shards), on both the 1D Ulysses path (v6e) and the 2D Ulysses+Ring path (tpu7x). - Lane sequence padding (wan_seq_pad="lane"): pads video tokens at the tail to a per-shard multiple of 128 (75,600 -> 75,776 at 720p/CP=4) so Ulysses reshapes are copy-free; TokenPadding masks the pad keys in the custom ulysses / ulysses_ring_custom kernels, and fixed-m metadata now also excludes pad query rows from the Q-norm bound. Every configuration that disables it logs why, once per cause. Behaviour change: if padding is active, the paths that cannot mask it (flash and non-custom ulysses via _tpu_flash_attention, ulysses_ring, dot_product, cudnn_flash_te, head_local_svg) raise NotImplementedError via `_reject_active_token_padding` instead of silently attending to pads. - run_wan_fast_inference.sh: the v6e and v7 profiles set WAN_CROSS_ATTN_KERNEL=pallas, WAN_PATCH_EMBED_MODE=tokens, WAN_ULYSSES_OUT_A2A=chunked, WAN_SEQ_PAD=lane (all overridable); the generic profile sets tokens only. The fused RMSNorm+RoPE kernel from #488 stays on in the v6e and v7 profiles and off in the generic one. LIBTPU: v6e now disables --xla_tpu_enable_ici_ag_pipelining / --xla_tpu_enable_ag_backward_pipelining, v7 enables them and adds --xla_tpu_enable_offloading_copy_to_sparsecore=false and BQ=4736. These two flags moved out of COMMON_LIBTPU, so the generic (non-v6e/v7) profile no longer sets them at all. Performance, measured at the stack tip (Wan 2.2 T2V-A14B, 720p/81f/40 steps, warm AOT, DVFS unpinned unless noted; CP=4, DP=2, launcher defaults including the fused RoPE kernel): v6e-8: denoise 125.9s, generate 128.05s tpu7x-8: denoise 105.1s, generate 107.6s tpu7x-8, DVFS pinned (PIN_DVFS_P_STATE=true), fused RoPE kernel off: denoise 93.6s, generate 96.1s (pinned with it on was not measured) Per-feature deltas were not measured separately and are omitted. In isolation the cross-attention 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. The launcher's reference-benchmark comment now carries these numbers. Tests: - cross_attention_pallas_test.py: 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_seq_pad_test.py: padding plan and per-cause disable logging, pad-key masking for ulysses custom (unaligned and aligned real_len, flat and shard-major out A2A, every fixed-m mode) and ulysses_ring_custom (k-centring on/off; and the tpu7x production path scaled down: U=2 x R=2, aligned real_len, chunked out A2A, asserting the shard-major layout was taken; needs >= 4 devices), fixed-m metadata ignoring pad query rows, and the NotImplementedError guards. - wan_patch_embed_test.py, ulysses_out_a2a_test.py: tokens vs conv parity and config-driven mode; flat vs chunked vs shard_major parity. - wan_runtime_options_test.py: the new switches' YAML defaults and enums. - CI: 29 tests skip in GitHub Actions: kernel parity and the Pallas-active dispatch / unaligned-KV fallback tests in cross_attention_pallas_test (19 of 39), the three kernel bit-exact tests in ulysses_out_a2a_test, the two WanModel patch-embed parity tests and the five masking / end-to-end tests in wan_seq_pad_test. Guards, VMEM budget logic, the remaining XLA fallback dispatch tests (including the unvalidated-TPU fallback), padding-plan logic and runtime options run in CI. - end_to_end/tpu/run_wan_stack_tests.sh: adds this PR's four test files. Verified (final tree): - TPU v6e-8 (jax 0.11.2, end_to_end/tpu/run_wan_stack_tests.sh): 350 passed, 1 skipped (off-TPU fallback only), 259 subtests passed. - TPU v6e-8 with GITHUB_ACTIONS=true (all 17 stack test files): 231 passed, 120 skipped, 202 subtests passed. - Emulated unvalidated TPU (device_kind="TPU v4", rope_accum_is_measured=False): 0 failures across all 17 stack test files. --- end_to_end/tpu/run_wan_fast_inference.sh | 43 +- end_to_end/tpu/run_wan_stack_tests.sh | 6 + src/maxdiffusion/__init__.py | 1 + src/maxdiffusion/configs/base_wan_14b.yml | 21 + src/maxdiffusion/configs/base_wan_1_3b.yml | 21 + src/maxdiffusion/configs/base_wan_27b.yml | 21 + src/maxdiffusion/configs/base_wan_animate.yml | 21 + src/maxdiffusion/configs/base_wan_i2v_14b.yml | 21 + src/maxdiffusion/configs/base_wan_i2v_27b.yml | 21 + .../kernels/cross_attention_pallas.py | 428 ++++++++++++ .../kernels/custom_splash_attention.py | 187 ++++- .../splash_attention/ring_attention_kernel.py | 53 +- src/maxdiffusion/models/attention_flax.py | 649 +++++++++++++++-- .../wan/transformers/transformer_wan.py | 222 +++++- .../tests/cross_attention_pallas_test.py | 421 +++++++++++ .../tests/ulysses_out_a2a_test.py | 279 ++++++++ .../tests/wan_patch_embed_test.py | 184 +++++ .../tests/wan_runtime_options_test.py | 19 +- src/maxdiffusion/tests/wan_seq_pad_test.py | 658 ++++++++++++++++++ src/maxdiffusion/utils/wan_block_benchmark.py | 22 + src/maxdiffusion/wan_runtime_options.py | 18 + 21 files changed, 3187 insertions(+), 129 deletions(-) create mode 100644 src/maxdiffusion/kernels/cross_attention_pallas.py create mode 100644 src/maxdiffusion/tests/cross_attention_pallas_test.py create mode 100644 src/maxdiffusion/tests/ulysses_out_a2a_test.py create mode 100644 src/maxdiffusion/tests/wan_patch_embed_test.py create mode 100644 src/maxdiffusion/tests/wan_seq_pad_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..b8114c46b 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 / WAN_PATCH_EMBED_MODE / WAN_ULYSSES_OUT_A2A / WAN_SEQ_PAD +# 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 @@ -132,8 +134,10 @@ if [ -z "${TPU_PROFILE:-}" ]; then fi # Keep LIBTPU flags single-line (literal backslashes truncate libtpu flag parsing). -COMMON_LIBTPU=${COMMON_LIBTPU:-"--xla_tpu_spmd_rng_bit_generator_unsafe=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_memory_bound_loop_optimizer_options=enabled:true --xla_tpu_enable_dot_strength_reduction=true --xla_enable_async_collective_permute=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_tpu_scoped_vmem_limit_kib=65536 --xla_tpu_enable_async_all_to_all=true --xla_tpu_enable_all_experimental_scheduler_features=true --xla_tpu_enable_scheduler_memory_pressure_tracking=true --xla_tpu_host_transfer_overlap_limit=24 --xla_tpu_aggressive_opt_barrier_removal=ENABLED --xla_lhs_prioritize_async_depth_over_stall=ENABLED --xla_should_allow_loop_variant_parameter_in_chain=ENABLED --xla_should_add_loop_invariant_op_in_chain=ENABLED --xla_tpu_enable_ici_ag_pipelining=true --xla_max_concurrent_host_send_recv=100 --xla_tpu_scheduler_percent_shared_memory_limit=100 --xla_latency_hiding_scheduler_rerun=2 --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 --xla_tpu_enable_latency_hiding_scheduler=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true --xla_tpu_sparse_core_all_gather_latency_multiplier=1 --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 --xla_tpu_enable_sparse_core_collective_aggregator=true --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true --xla_tpu_enable_sparse_core_reduce_scatter_v2=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true --xla_tpu_enable_concurrent_sparse_core_offloading=true --xla_tpu_assign_all_reduce_scatter_layout=true"} -V6E_LIBTPU=${V6E_LIBTPU:-"--xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false"} +# ici_ag_pipelining and ag_backward_pipelining are configured per-profile below +# (disabled on v6e, enabled on v7). +COMMON_LIBTPU=${COMMON_LIBTPU:-"--xla_tpu_spmd_rng_bit_generator_unsafe=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_memory_bound_loop_optimizer_options=enabled:true --xla_tpu_enable_dot_strength_reduction=true --xla_enable_async_collective_permute=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_tpu_scoped_vmem_limit_kib=65536 --xla_tpu_enable_async_all_to_all=true --xla_tpu_enable_all_experimental_scheduler_features=true --xla_tpu_enable_scheduler_memory_pressure_tracking=true --xla_tpu_host_transfer_overlap_limit=24 --xla_tpu_aggressive_opt_barrier_removal=ENABLED --xla_lhs_prioritize_async_depth_over_stall=ENABLED --xla_should_allow_loop_variant_parameter_in_chain=ENABLED --xla_should_add_loop_invariant_op_in_chain=ENABLED --xla_max_concurrent_host_send_recv=100 --xla_tpu_scheduler_percent_shared_memory_limit=100 --xla_latency_hiding_scheduler_rerun=2 --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 --xla_tpu_enable_latency_hiding_scheduler=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true --xla_tpu_sparse_core_all_gather_latency_multiplier=1 --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 --xla_tpu_enable_sparse_core_collective_aggregator=true --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true --xla_tpu_enable_sparse_core_reduce_scatter_v2=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true --xla_tpu_enable_concurrent_sparse_core_offloading=true --xla_tpu_assign_all_reduce_scatter_layout=true"} +V6E_LIBTPU=${V6E_LIBTPU:-"--xla_tpu_enable_ici_ag_pipelining=false --xla_tpu_enable_ag_backward_pipelining=false"} # fuse_all_gather is false in every profile (see COMMON_LIBTPU); it must stay # false on v7, where libtpu fails backend init with it enabled ("Continuation # fusion for AllGather ... not supported ... other than Viperlite"). @@ -144,13 +148,14 @@ V7_DVFS_FLAG="" if [ "$PIN_DVFS_P_STATE" = "true" ]; then V7_DVFS_FLAG=" --xla_tpu_dvfs_p_state=7" fi -V7_LIBTPU=${V7_LIBTPU:-"--xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false --xla_tpu_enable_megacore_fusion=true --xla_tpu_megacore_fusion_allow_ags=true${V7_DVFS_FLAG}"} +V7_LIBTPU=${V7_LIBTPU:-"--xla_tpu_enable_ici_ag_pipelining=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_megacore_fusion=true --xla_tpu_megacore_fusion_allow_ags=true --xla_tpu_enable_offloading_copy_to_sparsecore=false${V7_DVFS_FLAG}"} -# Reference benchmarks (Wan 2.2 T2V-A14B, 720p/81f/40-step, CP=4, DP=2, warm AOT -# cache, these launcher defaults, measured at this PR without the later PRs): -# tpu7x-8: 115.5s generate (denoise 112.8s) DVFS unpinned; -# 102.0s generate (denoise 99.4s) with PIN_DVFS_P_STATE=true -# v6e-8: not re-measured at this PR +# Reference benchmarks (Wan 2.2 T2V-A14B, 720p/81f/40-step, CP=4, DP=2, full 5-PR +# stack, these launcher defaults, warm AOT cache, DVFS unpinned): +# v6e-8: 128.05s generate (denoise 125.9s) +# tpu7x-8: 107.6s generate (denoise 105.1s) +# tpu7x-8 with PIN_DVFS_P_STATE=true: 96.1s generate (denoise 93.6s), measured +# with the fused RoPE kernel off (pinned with it on was not measured). case "$TPU_PROFILE" in v6e) PLATFORM_LIBTPU="$V6E_LIBTPU" @@ -167,12 +172,16 @@ case "$TPU_PROFILE" in DEFAULT_VAE_CHUNK=1 DEFAULT_VAE_SPATIAL=8 DEFAULT_FUSED_ROPE=true + DEFAULT_CROSS_ATTN_KERNEL=pallas + DEFAULT_PATCH_EMBED_MODE=tokens + DEFAULT_ULYSSES_OUT_A2A=chunked + DEFAULT_SEQ_PAD=lane ;; v7) PLATFORM_LIBTPU="$V7_LIBTPU" DEFAULT_ATTENTION="ulysses_ring_custom_fixed_m" DEFAULT_U=2 - DEFAULT_BQ=6400 + DEFAULT_BQ=4736 DEFAULT_BKV=2048 DEFAULT_BKV_COMPUTE=2048 DEFAULT_BKV_COMPUTE_IN=2048 @@ -183,6 +192,10 @@ case "$TPU_PROFILE" in DEFAULT_VAE_CHUNK=1 DEFAULT_VAE_SPATIAL=8 DEFAULT_FUSED_ROPE=true + DEFAULT_CROSS_ATTN_KERNEL=pallas + DEFAULT_PATCH_EMBED_MODE=tokens + DEFAULT_ULYSSES_OUT_A2A=chunked + DEFAULT_SEQ_PAD=lane ;; *) echo "== warning: unrecognised or non-v6e/v7 accelerator '${ACCEL_TYPE:-unknown}';" \ @@ -201,6 +214,10 @@ case "$TPU_PROFILE" in DEFAULT_VAE_CHUNK=1 DEFAULT_VAE_SPATIAL=8 DEFAULT_FUSED_ROPE=false + DEFAULT_CROSS_ATTN_KERNEL=xla + DEFAULT_PATCH_EMBED_MODE=tokens + DEFAULT_ULYSSES_OUT_A2A=flat + DEFAULT_SEQ_PAD=off ;; esac @@ -232,6 +249,10 @@ 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} +WAN_PATCH_EMBED_MODE=${WAN_PATCH_EMBED_MODE:-$DEFAULT_PATCH_EMBED_MODE} +WAN_ULYSSES_OUT_A2A=${WAN_ULYSSES_OUT_A2A:-$DEFAULT_ULYSSES_OUT_A2A} +WAN_SEQ_PAD=${WAN_SEQ_PAD:-$DEFAULT_SEQ_PAD} # 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 +315,10 @@ 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" \ + wan_patch_embed_mode="$WAN_PATCH_EMBED_MODE" \ + wan_ulysses_out_a2a="$WAN_ULYSSES_OUT_A2A" \ + wan_seq_pad="$WAN_SEQ_PAD" \ 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 05ef84332..844f55b92 100755 --- a/end_to_end/tpu/run_wan_stack_tests.sh +++ b/end_to_end/tpu/run_wan_stack_tests.sh @@ -42,5 +42,11 @@ 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" + # Cross-attention kernel, token patch embed, shard-major A2A, lane padding + # (feat/wan-fast-inference-optimizations) + "$T/cross_attention_pallas_test.py" + "$T/ulysses_out_a2a_test.py" + "$T/wan_patch_embed_test.py" + "$T/wan_seq_pad_test.py" ) PYTHONPATH="src${PYTHONPATH:+:$PYTHONPATH}" exec "${PYTHON:-python3}" -m pytest -q -rs "${TESTS[@]}" "$@" diff --git a/src/maxdiffusion/__init__.py b/src/maxdiffusion/__init__.py index bb7d5001a..801020c60 100644 --- a/src/maxdiffusion/__init__.py +++ b/src/maxdiffusion/__init__.py @@ -16,6 +16,7 @@ __version__ = "0.22.0.dev0" + from typing import TYPE_CHECKING from .utils import ( diff --git a/src/maxdiffusion/configs/base_wan_14b.yml b/src/maxdiffusion/configs/base_wan_14b.yml index 3465cb184..842c54efd 100644 --- a/src/maxdiffusion/configs/base_wan_14b.yml +++ b/src/maxdiffusion/configs/base_wan_14b.yml @@ -133,6 +133,27 @@ 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 +# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by +# token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the +# replicated kernel, no activation collective; equal up to accumulation order). +wan_patch_embed_mode: "conv" +# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded): +# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit +# leading sequence-chunk axis so XLA does not relayout the output around the +# collective (used by the custom Ulysses / Ulysses-ring kernels only). +wan_ulysses_out_a2a: "flat" +# Self-attention token padding: "off", or "lane" to append pad tokens at the end +# of the video token sequence (or at the end of each ring segment on the ring +# path) so every context shard has a multiple of 128 tokens. Aligned shards make +# the Ulysses all-to-all reshapes free; the custom kernels mask the pad keys. +# Only the custom Ulysses / Ulysses-ring kernels use it. +wan_seq_pad: "off" # 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..092784d73 100644 --- a/src/maxdiffusion/configs/base_wan_1_3b.yml +++ b/src/maxdiffusion/configs/base_wan_1_3b.yml @@ -130,6 +130,27 @@ 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 +# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by +# token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the +# replicated kernel, no activation collective; equal up to accumulation order). +wan_patch_embed_mode: "conv" +# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded): +# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit +# leading sequence-chunk axis so XLA does not relayout the output around the +# collective (used by the custom Ulysses / Ulysses-ring kernels only). +wan_ulysses_out_a2a: "flat" +# Self-attention token padding: "off", or "lane" to append pad tokens at the end +# of the video token sequence (or at the end of each ring segment on the ring +# path) so every context shard has a multiple of 128 tokens. Aligned shards make +# the Ulysses all-to-all reshapes free; the custom kernels mask the pad keys. +# Only the custom Ulysses / Ulysses-ring kernels use it. +wan_seq_pad: "off" # 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..74b64ab5f 100644 --- a/src/maxdiffusion/configs/base_wan_27b.yml +++ b/src/maxdiffusion/configs/base_wan_27b.yml @@ -162,6 +162,27 @@ 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 +# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by +# token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the +# replicated kernel, no activation collective; equal up to accumulation order). +wan_patch_embed_mode: "conv" +# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded): +# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit +# leading sequence-chunk axis so XLA does not relayout the output around the +# collective (used by the custom Ulysses / Ulysses-ring kernels only). +wan_ulysses_out_a2a: "flat" +# Self-attention token padding: "off", or "lane" to append pad tokens at the end +# of the video token sequence (or at the end of each ring segment on the ring +# path) so every context shard has a multiple of 128 tokens. Aligned shards make +# the Ulysses all-to-all reshapes free; the custom kernels mask the pad keys. +# Only the custom Ulysses / Ulysses-ring kernels use it. +wan_seq_pad: "off" # 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..5ad2482d3 100644 --- a/src/maxdiffusion/configs/base_wan_animate.yml +++ b/src/maxdiffusion/configs/base_wan_animate.yml @@ -112,6 +112,27 @@ 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 +# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by +# token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the +# replicated kernel, no activation collective; equal up to accumulation order). +wan_patch_embed_mode: "conv" +# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded): +# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit +# leading sequence-chunk axis so XLA does not relayout the output around the +# collective (used by the custom Ulysses / Ulysses-ring kernels only). +wan_ulysses_out_a2a: "flat" +# Self-attention token padding: "off", or "lane" to append pad tokens at the end +# of the video token sequence (or at the end of each ring segment on the ring +# path) so every context shard has a multiple of 128 tokens. Aligned shards make +# the Ulysses all-to-all reshapes free; the custom kernels mask the pad keys. +# Only the custom Ulysses / Ulysses-ring kernels use it. +wan_seq_pad: "off" # 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..0b7aa41ef 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_14b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_14b.yml @@ -114,6 +114,27 @@ 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 +# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by +# token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the +# replicated kernel, no activation collective; equal up to accumulation order). +wan_patch_embed_mode: "conv" +# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded): +# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit +# leading sequence-chunk axis so XLA does not relayout the output around the +# collective (used by the custom Ulysses / Ulysses-ring kernels only). +wan_ulysses_out_a2a: "flat" +# Self-attention token padding: "off", or "lane" to append pad tokens at the end +# of the video token sequence (or at the end of each ring segment on the ring +# path) so every context shard has a multiple of 128 tokens. Aligned shards make +# the Ulysses all-to-all reshapes free; the custom kernels mask the pad keys. +# Only the custom Ulysses / Ulysses-ring kernels use it. +wan_seq_pad: "off" # 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..16a3b8d25 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_27b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_27b.yml @@ -114,6 +114,27 @@ 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 +# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by +# token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the +# replicated kernel, no activation collective; equal up to accumulation order). +wan_patch_embed_mode: "conv" +# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded): +# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit +# leading sequence-chunk axis so XLA does not relayout the output around the +# collective (used by the custom Ulysses / Ulysses-ring kernels only). +wan_ulysses_out_a2a: "flat" +# Self-attention token padding: "off", or "lane" to append pad tokens at the end +# of the video token sequence (or at the end of each ring segment on the ring +# path) so every context shard has a multiple of 128 tokens. Aligned shards make +# the Ulysses all-to-all reshapes free; the custom kernels mask the pad keys. +# Only the custom Ulysses / Ulysses-ring kernels use it. +wan_seq_pad: "off" # 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..9a0613913 --- /dev/null +++ b/src/maxdiffusion/kernels/cross_attention_pallas.py @@ -0,0 +1,428 @@ +""" +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) + 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/kernels/custom_splash_attention.py b/src/maxdiffusion/kernels/custom_splash_attention.py index 6c0187d3f..6e553404b 100644 --- a/src/maxdiffusion/kernels/custom_splash_attention.py +++ b/src/maxdiffusion/kernels/custom_splash_attention.py @@ -842,16 +842,38 @@ def _splash_attention_forward( interpret: bool | None = None, save_residuals: bool = False, transpose_out: bool | None = None, + out_num_shards: int = 1, ): if interpret is None: interpret = jax.default_backend() == "cpu" if transpose_out is None: transpose_out = False - num_q_heads, padded_q_seq_len, head_dim_qk = q.shape - head_dim_v = v.shape[-1] + if transpose_out and out_num_shards > 1: + raise ValueError("out_num_shards > 1 is not supported with transpose_out=True.") + bq, bkv = block_sizes.block_q, block_sizes.block_kv bkv_compute = block_sizes.block_kv_compute bkv_compute_in = block_sizes.block_kv_compute_in + if q.ndim == 4: + num_q_shards, num_q_heads, shard_q_seq_len, head_dim_qk = q.shape + padded_q_seq_len = num_q_shards * shard_q_seq_len + if shard_q_seq_len % bq != 0: + raise ValueError(f"4D shard-major q requires shard_q_seq_len ({shard_q_seq_len}) divisible by bq ({bq}).") + q_blocks_per_shard_in = shard_q_seq_len // bq + + def q_index_map(h, i, j, *_): + return (i // q_blocks_per_shard_in, h, i % q_blocks_per_shard_in, 0) + + q_block_spec = pl.BlockSpec((None, None, bq, head_dim_qk), q_index_map) + else: + num_q_heads, padded_q_seq_len, head_dim_qk = q.shape + + def q_index_map(h, i, j, *_): + return (h, i, 0) + + q_block_spec = pl.BlockSpec((None, bq, head_dim_qk), q_index_map) + + head_dim_v = v.shape[-1] num_kv_heads = k.shape[0] padded_kv_seq_len = k.shape[1] @@ -880,9 +902,6 @@ def _splash_attention_forward( if mk.shape[0] != 2 or mk.shape[1] != num_q_heads or mk.shape[2] != grid_height: raise ValueError(f"mk must have shape (2, {num_q_heads}, {grid_height}), got {mk.shape}") - def q_index_map(h, i, j, *_): - return (h, i, 0) - def out_index_map(h, i, j, *_): return h, 0, i @@ -908,6 +927,30 @@ def v_index_map(h, i, j, *_): pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)), pl.BlockSpec((None, bq, head_dim_v), lambda h, i, j, *_: (h, i, 0)), ] + elif out_num_shards > 1: + if actual_q_seq_len % (out_num_shards * bq) != 0: + raise ValueError( + f"out_num_shards={out_num_shards} requires actual_q_seq_len ({actual_q_seq_len}) divisible by " + f"out_num_shards * bq ({out_num_shards * bq})." + ) + q_blocks_per_shard_out = grid_height // out_num_shards + shard_q_seq_len_out = actual_q_seq_len // out_num_shards + + def out_shard_major_index_map(h, i, j, *_): + return (i // q_blocks_per_shard_out, h, 0, i % q_blocks_per_shard_out) + + out_shapes = [ + jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), + jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), + jax.ShapeDtypeStruct((head_dim_v, bq), jnp.float32), + jax.ShapeDtypeStruct((out_num_shards, num_q_heads, head_dim_v, shard_q_seq_len_out), q.dtype), + ] + out_specs = [ + pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), + pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), + pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)), + pl.BlockSpec((None, None, head_dim_v, bq), out_shard_major_index_map), + ] else: out_shapes = [ jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), @@ -943,7 +986,7 @@ def v_index_map(h, i, j, *_): if k_mean is None: in_specs = [ - pl.BlockSpec((None, bq, head_dim_qk), q_index_map), + q_block_spec, pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), pl.BlockSpec((None, bkv, head_dim_v), v_index_map), ] @@ -969,7 +1012,7 @@ def v_index_map(h, i, j, *_): k_mean_tiled = jnp.repeat(k_mean[:, None, :], NUM_SUBLANES, axis=1) in_specs = [ - pl.BlockSpec((None, bq, head_dim_qk), q_index_map), + q_block_spec, pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), pl.BlockSpec((None, bkv, head_dim_v), v_index_map), pl.BlockSpec((None, NUM_SUBLANES, head_dim_qk), lambda h, i, j, *_: (h // q_heads_per_kv_head, 0, 0)), @@ -1031,6 +1074,7 @@ def _splash_attention_forward_ring( uniform_fixed_m: bool = False, k_mean: jax.Array | None = None, interpret: bool | None = None, + out_num_shards: int = 1, ): """Ring-specific forward path that returns pre-reciprocal fp32 accumulators. @@ -1041,19 +1085,47 @@ def _splash_attention_forward_ring( normalizes only once at the very end (see `ring_attention_kernel._custom_ring_attention_forward`). + Supports both 3D `q` `(num_q_heads, q_seq_len, head_dim_qk)` and 4D + shard-major `q` `(num_q_shards, num_q_heads, shard_q_seq_len, head_dim_qk)`. + When `out_num_shards > 1`, Pallas writes the output tiles directly into a 4D + shard-major buffer so the ring accumulator produces + `(out_num_shards, num_q_heads, q_seq_len // out_num_shards, head_dim_v)` with + no post-kernel `moveaxis` relayout before the inverse Ulysses `all_to_all`. + Returns: A tuple `(out, m, l)` where - - `out` has shape `(num_q_heads, q_seq_len, head_dim_v)` (fp32, un-normalized), - - `m` and `l` have shape `(num_q_heads, q_seq_len)` (fp32). + - `out` has shape `(num_q_heads, q_seq_len, head_dim_v)` or + `(out_num_shards, num_q_heads, q_seq_len // out_num_shards, head_dim_v)` + (fp32, un-normalized), + - `m` and `l` have shape `(num_q_heads, q_seq_len)` or + `(out_num_shards, num_q_heads, q_seq_len // out_num_shards)` (fp32). """ if interpret is None: interpret = jax.default_backend() == "cpu" - num_q_heads, padded_q_seq_len, head_dim_qk = q.shape - head_dim_v = v.shape[-1] bq, bkv = block_sizes.block_q, block_sizes.block_kv bkv_compute = block_sizes.block_kv_compute bkv_compute_in = block_sizes.block_kv_compute_in + if q.ndim == 4: + num_q_shards, num_q_heads, shard_q_seq_len, head_dim_qk = q.shape + padded_q_seq_len = num_q_shards * shard_q_seq_len + if shard_q_seq_len % bq != 0: + raise ValueError(f"4D shard-major q requires shard_q_seq_len ({shard_q_seq_len}) divisible by bq ({bq}).") + q_blocks_per_shard_in = shard_q_seq_len // bq + + def q_index_map(h, i, j, *_): + return (i // q_blocks_per_shard_in, h, i % q_blocks_per_shard_in, 0) + + q_block_spec = pl.BlockSpec((None, None, bq, head_dim_qk), q_index_map) + else: + num_q_heads, padded_q_seq_len, head_dim_qk = q.shape + + def q_index_map(h, i, j, *_): + return (h, i, 0) + + q_block_spec = pl.BlockSpec((None, bq, head_dim_qk), q_index_map) + + head_dim_v = v.shape[-1] num_kv_heads = k.shape[0] padded_kv_seq_len = k.shape[1] @@ -1063,9 +1135,6 @@ def _splash_attention_forward_ring( raise ValueError(f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA.") q_heads_per_kv_head = num_q_heads // num_kv_heads - def q_index_map(h, i, j, *_): - return (h, i, 0) - def out_index_map(h, i, j, *_): return h, 0, i @@ -1077,22 +1146,51 @@ def v_index_map(h, i, j, *_): grid_width = (actual_kv_seq_len + bkv - 1) // bkv grid_height = (actual_q_seq_len + bq - 1) // bq - out_shapes = [ - jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), - jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), - jax.ShapeDtypeStruct((head_dim_v, bq), jnp.float32), - jax.ShapeDtypeStruct((num_q_heads, head_dim_v, actual_q_seq_len), jnp.float32), - jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, actual_q_seq_len), jnp.float32), - jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, actual_q_seq_len), jnp.float32), - ] - out_specs = [ - pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), - pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), - pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)), - pl.BlockSpec((None, head_dim_v, bq), out_index_map), - pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map), - pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map), - ] + if out_num_shards > 1: + if actual_q_seq_len % (out_num_shards * bq) != 0: + raise ValueError( + f"out_num_shards={out_num_shards} requires actual_q_seq_len ({actual_q_seq_len}) divisible by " + f"out_num_shards * bq ({out_num_shards * bq})." + ) + shard_q_out_len = actual_q_seq_len // out_num_shards + q_blocks_per_shard_out = shard_q_out_len // bq + + def out_shard_index_map(h, i, j, *_): + return (i // q_blocks_per_shard_out, h, 0, i % q_blocks_per_shard_out) + + out_shapes = [ + jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), + jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), + jax.ShapeDtypeStruct((head_dim_v, bq), jnp.float32), + jax.ShapeDtypeStruct((out_num_shards, num_q_heads, head_dim_v, shard_q_out_len), jnp.float32), + jax.ShapeDtypeStruct((out_num_shards, num_q_heads, NUM_SUBLANES, shard_q_out_len), jnp.float32), + jax.ShapeDtypeStruct((out_num_shards, num_q_heads, NUM_SUBLANES, shard_q_out_len), jnp.float32), + ] + out_specs = [ + pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), + pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), + pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)), + pl.BlockSpec((None, None, head_dim_v, bq), out_shard_index_map), + pl.BlockSpec((None, None, NUM_SUBLANES, bq), out_shard_index_map), + pl.BlockSpec((None, None, NUM_SUBLANES, bq), out_shard_index_map), + ] + else: + out_shapes = [ + jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), + jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), + jax.ShapeDtypeStruct((head_dim_v, bq), jnp.float32), + jax.ShapeDtypeStruct((num_q_heads, head_dim_v, actual_q_seq_len), jnp.float32), + jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, actual_q_seq_len), jnp.float32), + jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, actual_q_seq_len), jnp.float32), + ] + out_specs = [ + pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), + pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), + pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)), + pl.BlockSpec((None, head_dim_v, bq), out_index_map), + pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map), + pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map), + ] grid = (num_q_heads, grid_height, grid_width) # Scalar-prefetch operand carrying per-head / per-Q-block fixed-m data: @@ -1116,7 +1214,7 @@ def v_index_map(h, i, j, *_): if k_mean is None: in_specs = [ - pl.BlockSpec((None, bq, head_dim_qk), q_index_map), + q_block_spec, pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), pl.BlockSpec((None, bkv, head_dim_v), v_index_map), ] @@ -1141,7 +1239,7 @@ def v_index_map(h, i, j, *_): k_mean_tiled = jnp.repeat(k_mean[:, None, :], NUM_SUBLANES, axis=1) in_specs = [ - pl.BlockSpec((None, bq, head_dim_qk), q_index_map), + q_block_spec, pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), pl.BlockSpec((None, bkv, head_dim_v), v_index_map), pl.BlockSpec((None, NUM_SUBLANES, head_dim_qk), lambda h, i, j, *_: (h // q_heads_per_kv_head, 0, 0)), @@ -1181,9 +1279,14 @@ def v_index_map(h, i, j, *_): out_shape=out_shapes, interpret=interpret, )(*kernel_args) - out = jnp.swapaxes(all_out[3], 1, 2) # (h, head_dim_v, s) -> (h, s, head_dim_v) - l = all_out[4][:, 0, :] # (h, s) - m = all_out[5][:, 0, :] # (h, s) + if out_num_shards > 1: + out = jnp.swapaxes(all_out[3], 2, 3) # (u, h, head_dim_v, s_loc) -> (u, h, s_loc, head_dim_v) + l = all_out[4][:, :, 0, :] # (u, h, s_loc) + m = all_out[5][:, :, 0, :] # (u, h, s_loc) + else: + out = jnp.swapaxes(all_out[3], 1, 2) # (h, head_dim_v, s) -> (h, s, head_dim_v) + l = all_out[4][:, 0, :] # (h, s) + m = all_out[5][:, 0, :] # (h, s) return out, m, l @@ -2136,6 +2239,7 @@ def make_splash_mha( interpret: bool | None = None, residual_checkpoint_name: str | None = None, transpose_out: bool | None = None, + out_num_shards: int = 1, ): """Builds a splash multi-head attention callable. @@ -2161,6 +2265,8 @@ def _splash_attention(q, k, v, mk=None, k_mean=None): if heads_per_tile > 1: if use_fixed_m: raise NotImplementedError("fixed-m is not supported with heads_per_tile > 1") + if out_num_shards > 1: + raise NotImplementedError("out_num_shards > 1 is not supported with heads_per_tile > 1") return _splash_attention_forward_mhpt( q, k, @@ -2174,7 +2280,15 @@ def _splash_attention(q, k, v, mk=None, k_mean=None): vmem_limit_bytes=vmem_limit_bytes, transpose_out=transpose_out, ) - if use_fixed_m or mk is not None or k_mean is not None or interpret is not None or transpose_out: + if ( + use_fixed_m + or mk is not None + or k_mean is not None + or interpret is not None + or transpose_out + or out_num_shards > 1 + or q.ndim == 4 + ): return _splash_attention_forward( q, k, @@ -2191,6 +2305,7 @@ def _splash_attention(q, k, v, mk=None, k_mean=None): k_mean=k_mean, interpret=interpret, transpose_out=transpose_out, + out_num_shards=out_num_shards, ) return _splash_attention_custom( q, diff --git a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py index 3f4cfd907..55cba41ec 100644 --- a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py +++ b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py @@ -750,6 +750,7 @@ def _custom_bidirectional_ring_forward( vmem_limit_bytes: int | None, mask_value: float, ring_axis: str, + out_num_shards: int = 1, ) -> jax.Array: """Wrap-free (bidirectional) ring attention for a NON-wrapping ring axis. @@ -783,6 +784,7 @@ def _attn(kc, vc): use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, + out_num_shards=out_num_shards, ) return o.astype(jnp.float32), m.astype(jnp.float32), l.astype(jnp.float32) @@ -875,13 +877,15 @@ def _custom_ring_attention_forward( v_ok: jax.Array | bool | None = None, all_fixed_global: jax.Array | bool | None = None, save_residuals: bool = False, + out_num_shards: int = 1, ) -> jax.Array | tuple[jax.Array, jax.Array]: """Ring attention using the custom dense splash kernel. Args: - q: Query shard, shape `(num_q_heads, q_seq_len, head_dim_qk)`. Stationary - across ring steps. Must already be padded (and pre-scaled by LOG2E when - `use_base2_exp`) by the caller. + q: Query shard, shape `(num_q_heads, q_seq_len, head_dim_qk)` or 4D + shard-major `(num_q_shards, num_q_heads, shard_q_seq_len, head_dim_qk)`. + Stationary across ring steps. Must already be padded (and pre-scaled by + LOG2E when `use_base2_exp`) by the caller. k: Key shard, shape `(num_kv_heads, kv_seq_len, head_dim_qk)`. Rotated across the ring axis. v: Value shard, shape `(num_kv_heads, kv_seq_len, head_dim_v)`. Rotated. @@ -931,18 +935,27 @@ def _custom_ring_attention_forward( accumulate-vs-LSE `lax.cond`. When supplied, skips the internal ring `pmin`. save_residuals: If True, returns `(out, lse)` where `lse` is the global logsumexp of shape `(num_q_heads, orig_q_seq_len)`. + out_num_shards: Number of leading Ulysses output shards (`1` returns 3D + `(num_q_heads, q_seq_len, head_dim_v)`; `> 1` returns 4D shard-major + `(out_num_shards, num_q_heads, q_seq_len // out_num_shards, head_dim_v)`). Returns: - Normalized attention output of shape `(num_q_heads, q_seq_len, head_dim_v)`, - or `(out, lse)` if `save_residuals=True`. + Normalized attention output, shape `(num_q_heads, q_seq_len, head_dim_v)` or + `(out_num_shards, num_q_heads, q_seq_len // out_num_shards, head_dim_v)`, or + `(out, lse)` if `save_residuals=True`. """ axis_size = lax.axis_size(ring_axis) effective_ring_size = ring_size if ring_size is not None else axis_size effective_kv_seq_len = orig_kv_seq_len * effective_ring_size - num_q_heads = q.shape[0] + num_q_heads = q.shape[1] if q.ndim == 4 else q.shape[0] num_kv_heads = k.shape[0] head_dim_v = v.shape[-1] + out_prefix_shape = ( + (out_num_shards, num_q_heads, orig_q_seq_len // out_num_shards) + if out_num_shards > 1 + else (num_q_heads, orig_q_seq_len) + ) if use_fixed_m and not use_base2_exp: raise NotImplementedError( @@ -996,6 +1009,7 @@ def _custom_ring_attention_forward( vmem_limit_bytes=vmem_limit_bytes, mask_value=mask_value, ring_axis=ring_axis, + out_num_shards=out_num_shards, ) if use_fixed_m and ring_size is not None and ring_size != axis_size: raise NotImplementedError( @@ -1137,8 +1151,8 @@ def _accumulate_scan(_): m_base = jnp.ceil(jnp.sqrt(bound_blocks_sq)) - global_recenter fixed_ok_expanded = jnp.ones_like(m_base) mk_arr = jnp.stack([m_base, fixed_ok_expanded], axis=0) # (2, heads, num_q_blocks) - o_sum = jnp.zeros((num_q_heads, orig_q_seq_len, head_dim_v), jnp.float32) - l_sum = jnp.zeros((num_q_heads, orig_q_seq_len), jnp.float32) + o_sum = jnp.zeros(out_prefix_shape + (head_dim_v,), jnp.float32) + l_sum = jnp.zeros(out_prefix_shape, jnp.float32) k_current, v_current = k, v # Python loop over the (static) ring size rather than a scan: it lets the # LAST hop skip its rotation. A scan body must rotate unconditionally, and @@ -1174,6 +1188,7 @@ def _accumulate_scan(_): # KV block instead of two (a two-body last block is what triggers # the Mosaic scheduler cliff on the whole grid). uniform_fixed_m=True, + out_num_shards=out_num_shards, ) o_sum = o_sum + o_curr.astype(jnp.float32) l_sum = l_sum + l_curr.astype(jnp.float32) @@ -1210,6 +1225,7 @@ def fixed_body(carry, is_last_hop, mk_arr): use_fixed_m=True, mk=mk_arr, k_mean=k_mean, + out_num_shards=out_num_shards, ) m_curr = m_curr.astype(jnp.float32) l_curr = l_curr.astype(jnp.float32) @@ -1257,8 +1273,8 @@ def _lse_scan(_): mk_arr_uniform = jnp.stack([m_base_hop, fixed_ok], axis=0) carry = ( - jnp.zeros((num_q_heads, orig_q_seq_len, head_dim_v), jnp.float32), - jnp.full((num_q_heads, orig_q_seq_len), lse_init, jnp.float32), + jnp.zeros(out_prefix_shape + (head_dim_v,), jnp.float32), + jnp.full(out_prefix_shape, lse_init, jnp.float32), k, v, ) @@ -1278,9 +1294,9 @@ def _lse_scan(_): cond_pred = jnp.all(all_fixed_global) if getattr(all_fixed_global, "ndim", 0) > 0 else all_fixed_global return lax.cond(cond_pred, _accumulate_scan, _lse_scan, None) - o_init = jnp.zeros((num_q_heads, orig_q_seq_len, head_dim_v), jnp.float32) - l_init = jnp.zeros((num_q_heads, orig_q_seq_len), jnp.float32) - m_init = jnp.full((num_q_heads, orig_q_seq_len), mask_value, jnp.float32) + o_init = jnp.zeros(out_prefix_shape + (head_dim_v,), jnp.float32) + l_init = jnp.zeros(out_prefix_shape, jnp.float32) + m_init = jnp.full(out_prefix_shape, mask_value, jnp.float32) m_final, l_final, o_final = m_init, l_init, o_init k_current, v_current = k, v @@ -1303,6 +1319,7 @@ def _lse_scan(_): use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, + out_num_shards=out_num_shards, ) m_curr = m_curr.astype(jnp.float32) l_curr = l_curr.astype(jnp.float32) @@ -1551,12 +1568,15 @@ def make_custom_ring_attention( v_ok: jax.Array | bool | None = None, all_fixed_global: jax.Array | bool | None = None, residual_checkpoint_name: str | None = None, + out_num_shards: int = 1, ): """Builds a ring-attention callable around the custom kernel (supports forward & backward). The returned function takes a single (un-batched) `(q, k, v)` triple of shape - `(num_heads, seq, head_dim)` and optional per-batch `fixed_m_norms=(qn_max_sq, mk_h_sq)` - and `k_mean` to be `jax.vmap`-ped over the batch axis inside the attention `shard_map`. + `(num_heads, seq, head_dim)` (or 4D shard-major `q` of shape + `(num_q_shards, num_heads, shard_seq, head_dim)`) and optional per-batch + `fixed_m_norms=(qn_max_sq, mk_h_sq)` and `k_mean` to be `jax.vmap`-ped over the + batch axis inside the attention `shard_map`. `fixed_m_norms_squared` declares the representation of `fixed_m_norms` and is **required** when `use_fixed_m=True`. The kernel gates in squared-norm space @@ -1622,7 +1642,7 @@ def make_custom_ring_attention( def _ring(q, k, v, fixed_m_norms=None, k_mean=None): norms = fixed_m_norms if fixed_m_norms is not None else default_fixed_m_norms km = k_mean if k_mean is not None else default_k_mean - if use_fixed_m or bidirectional or km is not None or uniform_fixed_m is not None: + if use_fixed_m or bidirectional or km is not None or uniform_fixed_m is not None or out_num_shards > 1 or q.ndim == 4: return _custom_ring_attention_forward( q, k, @@ -1647,6 +1667,7 @@ def _ring(q, k, v, fixed_m_norms=None, k_mean=None): uniform_fixed_m=uniform_fixed_m, v_ok=v_ok, all_fixed_global=all_fixed_global, + out_num_shards=out_num_shards, ) return _custom_ring_attention_custom( q, diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index f99f9f2cf..757f71526 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -13,6 +13,8 @@ # limitations under the License. import contextlib +import contextvars +import dataclasses import functools import math from typing import Optional, Callable, Tuple, Any, Dict @@ -36,6 +38,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 @@ -715,6 +718,128 @@ def _ulysses_head_chunk_ranges(num_heads: int, ulysses_shards: int, num_chunks: return ranges +@dataclasses.dataclass(frozen=True) +class TokenPadding: + """Pad tokens a model inserted into its self-attention sequence. + + The sequence is `num_segments` equal segments laid end to end: one per ring + shard on the Ulysses-ring path, a single segment otherwise. Each segment holds + `real_len` real tokens followed by `padded_len - real_len` pad tokens. + """ + + real_len: int + padded_len: int + num_segments: int = 1 + + def __post_init__(self): + if not 0 < self.real_len <= self.padded_len or self.num_segments < 1: + raise ValueError(f"Invalid token padding: {self}.") + + @property + def total_len(self) -> int: + return self.padded_len * self.num_segments + + +# Kernels whose self-attention honours TokenPadding. The pads sit where these +# kernels already expect a ragged KV tail (the end of the sequence, or of each +# ring segment), so passing the real length masks the pad keys, and the fixed-m +# norm reductions skip the pad rows. Any other kernel would attend to the pads. +TOKEN_PADDING_KERNELS = frozenset({ + "ulysses_custom", + "ulysses_custom_fixed_m", + "ulysses_custom_fixed_m_per_q_block", + "ulysses_ring_custom", + "ulysses_ring_custom_fixed_m", + "ulysses_ring_custom_fixed_m_per_q_block", +}) + +_ACTIVE_TOKEN_PADDING: contextvars.ContextVar[Optional[TokenPadding]] = contextvars.ContextVar( + "maxdiffusion_active_token_padding", default=None +) + + +@contextlib.contextmanager +def self_attention_token_padding(padding: Optional[TokenPadding]): + """Declares, while the transformer blocks are traced, that the video tokens carry `padding`.""" + token = _ACTIVE_TOKEN_PADDING.set(padding) + try: + yield + finally: + _ACTIVE_TOKEN_PADDING.reset(token) + + +def _active_token_padding(seq_len: int, num_segments: int, q_seq_len: Optional[int] = None) -> Optional[TokenPadding]: + """The declared padding when it describes this `seq_len`-token KV sequence, else None. + + Only self-attention over the padded video tokens matches (q_seq_len == seq_len). + Cross-attention keys (text or image tokens) have their own length and carry no pads. + """ + padding = _ACTIVE_TOKEN_PADDING.get() + if padding is None or padding.total_len != seq_len or (q_seq_len is not None and q_seq_len != seq_len): + return None + if padding.num_segments != num_segments: + raise ValueError( + f"The tokens were padded as {padding.num_segments} segment(s), but this attention splits the sequence into " + f"{num_segments} ring segment(s), so the pads would not sit at the tail of each segment." + ) + return padding + + +def _reject_active_token_padding(q_seq_len: int, kv_seq_len: int, kernel_name: str) -> None: + """Raises NotImplementedError if self-attention token padding is active on an unsupported path.""" + padding = _ACTIVE_TOKEN_PADDING.get() + if padding is not None and padding.total_len == kv_seq_len and q_seq_len == kv_seq_len: + raise NotImplementedError( + f"{kernel_name} does not support self_attention_token_padding " + f"(only {sorted(TOKEN_PADDING_KERNELS)} mask tail pad tokens)." + ) + + +ULYSSES_OUT_A2A_MODES = ("flat", "chunked", "shard_major") + + +def _ulysses_seq_to_heads( + x: jax.Array, *, axis_name, seq_axis: int, num_shards: int, mode: Optional[str] = None +) -> jax.Array: + """Inverse Ulysses exchange: `all_to_all(x, split_axis=seq_axis, concat_axis=1, tiled=True)`. + + Shard c receives the c-th sequence chunk of every local head. `x` is + [B, H_local, ...] with the full sequence on `seq_axis`; the result is + [B, num_shards * H_local, ...] with the sequence cut to one chunk. + + `wan_ulysses_out_a2a` picks how the collective is written: + * "flat": a plain split of the sequence axis. XLA reorders the whole output + so the sequence axis is major before the collective, then reorders it back + for the output projection. That is four relayout copies per layer at the + 720p shard shape (~1.5 ms on v6e, ~0.7 ms on tpu7x). + * "chunked" / "shard_major": a single layout-preserving `all_to_all` over a + reshaped `(num_shards, chunk)` leading axis (`split_axis=1, concat_axis=1`, + not head-chunked compute/comm overlap like `ulysses_attention_chunks`). + Each shard's block is then contiguous in whatever minor-dimension order + the producer used. The received `[src, H_local, ...]` blocks merge into + `[src * H_local, ...]`, which is the order `"flat"` produces. Same data + movement, same values. When the custom Pallas splash/ring kernel runs with + `S_local % block_q == 0` and `D == 128`, `"chunked"` / `"shard_major"` + additionally keeps the forward Q exchange shard-major (`4D q`) and emits + the shard-major output layout directly from Pallas + (`out_num_shards=num_shards`), avoiding the post-kernel relayout copy on + both v6e (1D Ulysses) and tpu7x (2D Ulysses+Ring). Measured at roughly + 1-1.5 s per 40-step 720p run on each, on a pre-review revision with the + DVFS state unrecorded; not re-measured since. + """ + if mode is None: + mode = "flat" + if mode not in ULYSSES_OUT_A2A_MODES: + raise ValueError(f"wan_ulysses_out_a2a must be one of {ULYSSES_OUT_A2A_MODES}, got {mode!r}.") + seq_len = x.shape[seq_axis] + if mode == "flat" or num_shards == 1 or seq_len % num_shards: + return jax.lax.all_to_all(x, axis_name=axis_name, split_axis=seq_axis, concat_axis=1, tiled=True) + x = x.reshape(x.shape[:seq_axis] + (num_shards, seq_len // num_shards) + x.shape[seq_axis + 1 :]) + x = jnp.moveaxis(x, seq_axis, 1) # [B, num_shards, H_local, ...] + x = jax.lax.all_to_all(x, axis_name=axis_name, split_axis=1, concat_axis=1, tiled=True) + return x.reshape((x.shape[0], x.shape[1] * x.shape[2]) + x.shape[3:]) + + def _run_chunked_ulysses_attention( query: jax.Array, key: jax.Array, @@ -790,8 +915,9 @@ def _tpu_flash_attention( num_context_shards = mesh.shape[CONTEXT] if CONTEXT in mesh.shape else 1 query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_context_shards) - key, _ = _reshape_data_for_flash(key, heads, num_context_shards) + key, orig_kv_seq_len = _reshape_data_for_flash(key, heads, num_context_shards) value, _ = _reshape_data_for_flash(value, heads, num_context_shards) + _reject_active_token_padding(orig_q_seq_len, orig_kv_seq_len, attention_kernel) attention_mask = _prepare_attention_mask_for_shard_map(attention_mask, query.shape[0], key.shape[2]) if attention_mask is not None and attention_kernel == "tokamax_ring_custom": raise NotImplementedError("tokamax_ring_custom does not support attention_mask.") @@ -1054,6 +1180,7 @@ def _ulysses_attention( use_k_centering: bool = True, qk_prescaled: bool = False, transpose_out: bool = False, + wan_ulysses_out_a2a: str = "flat", ) -> jax.Array: """Ulysses sequence-parallel attention. @@ -1077,6 +1204,13 @@ def _ulysses_attention( query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_shards) key, orig_kv_seq_len = _reshape_data_for_flash(key, kv_heads, num_shards) value, _ = _reshape_data_for_flash(value, kv_heads, num_shards) + # Pad tokens inserted by the model (see `TokenPadding`) sit at the tail of the + # sequence, like the shard padding above, and are masked the same way: the + # custom kernel reads only the first `real_kv_seq_len` keys. + token_padding = _active_token_padding(orig_kv_seq_len, num_segments=1, q_seq_len=orig_q_seq_len) + if token_padding is not None and (not use_custom_kernel or spatiotemporal_config is not None): + raise NotImplementedError(f"{kernel_name}: token padding needs the dense custom kernel, which masks the pad keys.") + real_kv_seq_len = orig_kv_seq_len if token_padding is None else token_padding.real_len attention_mask = _prepare_attention_mask_for_shard_map(attention_mask, query.shape[0], key.shape[2]) if attention_mask is not None and use_custom_kernel: raise NotImplementedError( @@ -1122,9 +1256,45 @@ def wrap_ulysses_attention(query, key, value, attention_mask): # into the producer of Q and its 185MB round-trip disappears. if use_custom_kernel and use_base2_exp and not qk_prescaled: query = query * LOG2E + + eff_out_a2a_mode = wan_ulysses_out_a2a + if eff_out_a2a_mode not in ULYSSES_OUT_A2A_MODES: + raise ValueError(f"wan_ulysses_out_a2a must be one of {ULYSSES_OUT_A2A_MODES}, got {eff_out_a2a_mode!r}.") + + if use_custom_kernel: + ( + bq, + bkv, + bkv_compute, + bkv_compute_in, + heads_per_tile, + vmem_limit_bytes, + ) = _extract_custom_block_sizes(flash_block_sizes) + else: + bq = bkv = bkv_compute = bkv_compute_in = heads_per_tile = 1 + vmem_limit_bytes = None + + use_shard_major = ( + use_custom_kernel + and eff_out_a2a_mode in ("chunked", "shard_major") + and not transpose_out + and attention_mask is None + and num_shards > 1 + and (query.shape[2] % bq == 0) + and (query.shape[-1] == 128) + and (value.shape[-1] == 128) + and (heads_per_tile == 1) + ) + out_num_shards = num_shards if use_shard_major else 1 + # Swap sharding: each device gives up a slice of heads and gathers # a slice of sequence, so the local kernel sees the full sequence. - query = jax.lax.all_to_all(query, axis_name=axis_name, split_axis=1, concat_axis=2, tiled=True) + if use_shard_major: + b_q, h_q, s_loc_q, d_q = query.shape + query = query.reshape(b_q, num_shards, h_q // num_shards, s_loc_q, d_q) + query = jax.lax.all_to_all(query, axis_name=axis_name, split_axis=1, concat_axis=1, tiled=True) + else: + query = jax.lax.all_to_all(query, axis_name=axis_name, split_axis=1, concat_axis=2, tiled=True) key = jax.lax.all_to_all(key, axis_name=axis_name, split_axis=1, concat_axis=2, tiled=True) value = jax.lax.all_to_all(value, axis_name=axis_name, split_axis=1, concat_axis=2, tiled=True) if attention_mask is not None and mask_needs_ulysses_gather: @@ -1136,27 +1306,25 @@ def wrap_ulysses_attention(query, key, value, attention_mask): "The custom dense splash kernel (use_custom_kernel) does not support attention_mask " "(it only handles padding via orig_seq_len); got a non-None attention_mask." ) - ( - bq, - bkv, - bkv_compute, - bkv_compute_in, - heads_per_tile, - vmem_limit_bytes, - ) = _extract_custom_block_sizes(flash_block_sizes) # NOTE: the base-2 rescale of Q is applied before the all-to-all above. raw_key = key raw_query = query raw_value = value - context_q_seq_len = raw_query.shape[2] - actual_kv_seq_len = orig_kv_seq_len + context_q_seq_len = raw_query.shape[1] * raw_query.shape[3] if use_shard_major else raw_query.shape[2] + # With token padding this excludes the pad keys: everything below (the + # kernel's ragged-tail mask, k_mean, the norm bounds) reads this length. + actual_kv_seq_len = real_kv_seq_len real_key = raw_key[:, :, :actual_kv_seq_len, :] recenter, safe_bound = custom_splash.get_fixed_m_constants(actual_kv_seq_len) - query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) + if use_shard_major: + kv_size = raw_query.shape[-1] + query_seq_len = context_q_seq_len + else: + query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) k_mean = None if use_fixed_m and use_k_centering: # Virtual k-centering (output-invariant): project q^T \bar{k} inside the @@ -1194,7 +1362,10 @@ def wrap_ulysses_attention(query, key, value, attention_mask): # the padded copy chained a 193.5MB (per batch element) pad + reduction behind the V # all-to-all and left that collective fully exposed. The padding is # zeros and the check is a max of squares, so this is output-invariant. - value=raw_value, + # Model pad tokens are not zeros, so they are sliced off instead. + value=raw_value if token_padding is None else raw_value[:, :, :actual_kv_seq_len, :], + # Like the ring's `row_ok`: pad-token Q rows never gate fixed-m. + q_valid_len=None if token_padding is None else token_padding.real_len, ) bsizes = _build_custom_splash_block_sizes(flash_block_sizes, bq, bkv, bkv_compute, bkv_compute_in) @@ -1211,6 +1382,7 @@ def wrap_ulysses_attention(query, key, value, attention_mask): use_fixed_m=True, uniform_fixed_m=True, transpose_out=transpose_out, + out_num_shards=out_num_shards, ) splash_kernel_hybrid = custom_splash.make_splash_mha( block_sizes=bsizes, @@ -1223,6 +1395,7 @@ def wrap_ulysses_attention(query, key, value, attention_mask): use_fixed_m=True, uniform_fixed_m=False, transpose_out=transpose_out, + out_num_shards=out_num_shards, ) def _run_uniform(q, k, v, m, km): @@ -1253,30 +1426,29 @@ def _run_hybrid(q, k, v, m, km): use_fixed_m=False, residual_checkpoint_name=residual_checkpoint_name, transpose_out=transpose_out, + out_num_shards=out_num_shards, ) vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0)) attention_output = vmapped_splash(query, key, value) + if use_shard_major: + attention_output = attention_output.astype(query.dtype) + attention_output = jax.lax.all_to_all(attention_output, axis_name=axis_name, split_axis=1, concat_axis=1, tiled=True) + return attention_output.reshape( + (attention_output.shape[0], attention_output.shape[1] * attention_output.shape[2]) + attention_output.shape[3:] + ) if transpose_out: attention_output = attention_output[:, :, :context_q_seq_len, :kv_size].astype(query.dtype) # Restore original layout: head-sharded/full-sequence -> sequence-sharded/full-heads. # Sequence axis is at index 2 (sublanes), heads axis is at index 1. - attention_output = jax.lax.all_to_all( - attention_output, - axis_name=axis_name, - split_axis=2, - concat_axis=1, - tiled=True, + attention_output = _ulysses_seq_to_heads( + attention_output, axis_name=axis_name, seq_axis=2, num_shards=num_shards, mode=eff_out_a2a_mode ) else: attention_output = attention_output[:, :, :kv_size, :context_q_seq_len].astype(query.dtype) # Restore original layout: head-sharded/full-sequence -> sequence-sharded/full-heads. # Sequence axis is at index 3, heads axis is at index 1. - attention_output = jax.lax.all_to_all( - attention_output, - axis_name=axis_name, - split_axis=3, - concat_axis=1, - tiled=True, + attention_output = _ulysses_seq_to_heads( + attention_output, axis_name=axis_name, seq_axis=3, num_shards=num_shards, mode=eff_out_a2a_mode ) return attention_output else: @@ -1497,10 +1669,10 @@ def _ulysses_ring_attention( ) internal_sequence_axes = (ring_axis, ulysses_axis) num_sequence_shards = num_context_shards - query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_sequence_shards) - key, _ = _reshape_data_for_flash(key, kv_heads, num_sequence_shards) + key, orig_kv_seq_len = _reshape_data_for_flash(key, kv_heads, num_sequence_shards) value, _ = _reshape_data_for_flash(value, kv_heads, num_sequence_shards) + _reject_active_token_padding(orig_q_seq_len, orig_kv_seq_len, "ulysses_ring") attention_mask = _prepare_attention_mask_for_shard_map(attention_mask, query.shape[0], key.shape[2]) num_heads = query.shape[1] @@ -1674,6 +1846,7 @@ def _ring_fixed_m_norms_pre_a2a( block_q: int, per_q_block: bool, use_k_centering: bool = False, + token_padding: Optional[TokenPadding] = None, ): """Computes all R>1 fixed-m norms and global eligibility predicates *pre* a2a. @@ -1699,6 +1872,11 @@ def _ring_fixed_m_norms_pre_a2a( exact same barriered Q/K (V is intentionally left out of the barrier; see the caller). + With `token_padding`, the pad rows at the tail of each ring segment (the tail + of the last Ulysses ranks' local rows) are left out of every reduction: the + kernel masks those keys and the model drops those queries, so they must not + move the bounds. Pad Q/K rows are zero after RoPE anyway; pad V rows are not. + Returns `(qn_dev, mk_global_sq_dev, v_ok, all_fixed_global, k_mean_dev)` sliced to the heads this Ulysses rank owns and ready for immediate `jax.lax.cond` dispatch post-a2a. @@ -1706,21 +1884,38 @@ def _ring_fixed_m_norms_pre_a2a( reduce_axes = (ulysses_axis, ring_axis) key_f32 = key.astype(jnp.float32) + row_ok = None + if token_padding is not None: + local_len = key.shape[2] + segment_row = jax.lax.axis_index(ulysses_axis) * local_len + jnp.arange(local_len) + row_ok = segment_row < token_padding.real_len + q_norm_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1) + v_sq = value.astype(jnp.float32) ** 2 + if row_ok is not None: + q_norm_sq = jnp.where(row_ok[None, None, :], q_norm_sq, 0.0) + v_sq = jnp.where(row_ok[None, None, :, None], v_sq, 0.0) qn_head_local = q_norm_sq.max(axis=-1) - vn_local = (value.astype(jnp.float32) ** 2).max() + vn_local = v_sq.max() if use_k_centering: # Virtual K-centering: compute global mean and centered key norms without # modifying `key`, keeping `all_to_all(key)` independent of `pmean` and # passing `k_mean_dev` to the ring kernel for in-kernel projection. - k_mean_all = jax.lax.pmean(jnp.mean(key_f32, axis=2), axis_name=reduce_axes) + if row_ok is None: + k_mean_all = jax.lax.pmean(jnp.mean(key_f32, axis=2), axis_name=reduce_axes) + else: + k_sum = jnp.where(row_ok[None, None, :, None], key_f32, 0.0).sum(axis=2) + k_mean_all = jax.lax.psum(k_sum, axis_name=reduce_axes) / (token_padding.real_len * num_ring_shards) centered_f32 = key_f32 - k_mean_all[:, :, None, :] - kn_local = jnp.sum(centered_f32**2, axis=-1).max(axis=-1) + kn_rows = jnp.sum(centered_f32**2, axis=-1) k_mean_dev = _slice_own_ulysses_heads(k_mean_all, ulysses_axis, num_ulysses_shards, axis=1) else: - kn_local = jnp.sum(key_f32**2, axis=-1).max(axis=-1) + kn_rows = jnp.sum(key_f32**2, axis=-1) k_mean_dev = None + if row_ok is not None: + kn_rows = jnp.where(row_ok[None, None, :], kn_rows, 0.0) + kn_local = kn_rows.max(axis=-1) # Global Q/V/K max norms in a SINGLE (ulysses, ring) pmax. # @@ -1774,6 +1969,8 @@ def _ring_fixed_m_norms_pre_a2a( # uniform `_accumulate_scan` branch is all-or-nothing. The cost of that # fallback in production (how often a layer trips it) has not been measured. effective_kv_seq_len = key.shape[2] * num_ulysses_shards * num_ring_shards + if token_padding is not None: + effective_kv_seq_len = token_padding.real_len * num_ring_shards global_recenter, global_bound = custom_splash.get_fixed_m_constants(effective_kv_seq_len) global_bound_sq = global_bound**2 dtype_safe = custom_splash.fixed_m_dtype_is_safe(query.dtype, global_recenter) @@ -1796,6 +1993,7 @@ def _compute_fixed_m_metadata( k_mean: jax.Array | None = None, value: jax.Array | None = None, v_max_bound: float = 256.0, + q_valid_len: int | None = None, ) -> tuple[jax.Array, jax.Array]: """Computes Cauchy-Schwarz norm bounds and per-Q-block (or per-head) fixed-m metadata. @@ -1816,6 +2014,11 @@ def _compute_fixed_m_metadata( verify that |V| <= v_max_bound to guarantee against FP32 overflow. Omitting `value` fails closed (`v_ok = 0.0`). v_max_bound: Maximum safe value magnitude (default 256.0). + q_valid_len: Number of real query tokens in global sequence order, when the + sequence carries token padding at its tail. Pad rows are excluded from + the Q norms (their outputs are discarded), so pad-token values cannot + push a block onto the online-softmax path. For a 5-D shard-major query + the global position of row `i` in shard `c` is `c * shard_q_len + i`. Returns: mk_arr: Gating metadata array of shape `(batch, 2, local_q_heads, num_q_blocks)` @@ -1825,7 +2028,13 @@ def _compute_fixed_m_metadata( - `mk_arr[:, 1, h, i]`: Discrete eligibility predicate (1.0 for fixed-m, 0.0 for online). all_fixed: Boolean scalar indicating if all elements are eligible for uniform fixed-m. """ - batch_size, num_q_heads, q_len, _ = query.shape + if query.ndim == 5: + batch_size, num_q_shards, num_q_heads, shard_q_len, _ = query.shape + q_len = num_q_shards * shard_q_len + else: + batch_size, num_q_heads, q_len, _ = query.shape + num_q_shards = 1 + shard_q_len = q_len num_kv_heads = key.shape[1] if safe_bound is None or recenter is None: rec, bnd = custom_splash.get_fixed_m_constants(key.shape[2], v_max_bound=v_max_bound) @@ -1866,18 +2075,38 @@ def _compute_fixed_m_metadata( f"_compute_fixed_m_metadata expects query padded to a multiple of block_q, got q_len={q_len}, block_q={block_q}." ) num_q_blocks = q_len // block_q + + def _q_norm_sq(): + norm_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1) + if q_valid_len is None: + return norm_sq + if query.ndim == 5: + pos = jnp.arange(num_q_shards)[:, None] * shard_q_len + jnp.arange(shard_q_len)[None, :] + pos = pos[None, :, None, :] # (1, shards, 1, shard_q_len) + else: + pos = jnp.arange(q_len)[None, None, :] + return jnp.where(pos < q_valid_len, norm_sq, 0.0) + if per_q_block: - norm_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1) # (batch, num_q_heads, q_len) - qn_max_sq = norm_sq.reshape(batch_size, num_q_heads, num_q_blocks, block_q).max( - axis=-1 - ) # (batch, num_q_heads, num_q_blocks) + norm_sq = _q_norm_sq() + if query.ndim == 5: + shard_q_blocks = shard_q_len // block_q + qn_max_sq = norm_sq.reshape(batch_size, num_q_shards, num_q_heads, shard_q_blocks, block_q).max(axis=-1) + qn_max_sq = jnp.moveaxis(qn_max_sq, 1, 2).reshape(batch_size, num_q_heads, num_q_blocks) + else: + qn_max_sq = norm_sq.reshape(batch_size, num_q_heads, num_q_blocks, block_q).max( + axis=-1 + ) # (batch, num_q_heads, num_q_blocks) bound_sq = qn_max_sq * mk_h_sq[:, :, None] fixed_ok = (bound_sq <= safe_bound_sq).astype(jnp.float32) * v_ok m_base = jnp.ceil(jnp.sqrt(bound_sq)) - recenter - mk_arr = jnp.stack([m_base, fixed_ok], axis=1) # (batch, 2, num_q_heads, num_q_blocks) + mk_arr = jnp.stack([m_base, fixed_ok], axis=1) # (batch, 2, local_q_heads, num_q_blocks) all_fixed = jnp.all(fixed_ok > 0.5) else: - qn_max_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1).max(axis=-1) # (batch, num_q_heads) + if query.ndim == 5: + qn_max_sq = _q_norm_sq().max(axis=(1, 3)) + else: + qn_max_sq = _q_norm_sq().max(axis=-1) # (batch, num_q_heads) bound_sq_1d = qn_max_sq * mk_h_sq fixed_ok_1d = (bound_sq_1d <= safe_bound_sq).astype(jnp.float32) * v_ok m_base_1d = jnp.ceil(jnp.sqrt(bound_sq_1d)) - recenter @@ -1913,6 +2142,7 @@ def _ulysses_ring_custom_attention( use_k_centering: bool = False, qk_prescaled: bool = False, transpose_out: bool = False, + wan_ulysses_out_a2a: str = "flat", ) -> jax.Array: """2D USP attention (Ulysses + Ring) using custom splash kernel with exact Fixed-m support.""" if kv_heads is None: @@ -1959,6 +2189,10 @@ def _ulysses_ring_custom_attention( f"2D Ulysses+Ring attention requires sequence length to be divisible by context_shards={num_context_shards}, " f"got orig_q_seq_len={orig_q_seq_len}, orig_kv_seq_len={orig_kv_seq_len}." ) + token_padding = _active_token_padding(orig_kv_seq_len, num_segments=num_ring_shards, q_seq_len=orig_q_seq_len) + if token_padding is not None and bidirectional: + raise NotImplementedError("ulysses_ring_custom_bidir does not support self_attention_token_padding.") + ( bq, bkv, @@ -2024,11 +2258,44 @@ def wrap_ulysses_ring_attention(query, key, value): block_q=bq, per_q_block=per_q_block, use_k_centering=use_k_centering, + token_padding=token_padding, ) + eff_out_a2a_mode = wan_ulysses_out_a2a + if eff_out_a2a_mode not in ULYSSES_OUT_A2A_MODES: + raise ValueError(f"wan_ulysses_out_a2a must be one of {ULYSSES_OUT_A2A_MODES}, got {eff_out_a2a_mode!r}.") + + # On the 2D Ulysses+Ring path (num_ring_shards > 1), `wan_ulysses_out_a2a` in + # ("chunked", "shard_major") keeps both the forward Q exchange and the inverse + # output exchange in shard-major `[B, U, H_local, S_local, D]` layout + # (`split_axis=1, concat_axis=1`): + # * Q is reshaped to `[B, U, H // U, S_local, D]` before the forward A2A and + # indexed directly as 4D `(U, H_local, S_local, D)` via Pallas `BlockSpec` + # in `_splash_attention_forward_ring`, eliminating the post-A2A `concat_axis=2` + # relayout copy. + # * `_splash_attention_forward_ring` writes `(U, H_local, D, S_local)` via + # `out_num_shards=U` and the ring loop accumulates `o_sum` directly in + # `(U, H_local, S_local, D)`, so the inverse A2A consumes the ring output + # with zero post-kernel `jnp.moveaxis` relayout copy. + use_shard_major = ( + eff_out_a2a_mode in ("chunked", "shard_major") + and num_ring_shards > 1 + and num_ulysses_shards > 1 + and (query.shape[2] % bq == 0) + and (query.shape[-1] == 128) + and (value.shape[-1] == 128) + and (heads_per_tile == 1) + ) + out_num_shards = num_ulysses_shards if use_shard_major else 1 + # (1) Ulysses All-to-All: heads -> sequence a2a = functools.partial(jax.lax.all_to_all, axis_name=ulysses_axis, tiled=True) - query = a2a(query, split_axis=1, concat_axis=2) + if use_shard_major: + b_q, h_q, s_loc_q, d_q = query.shape + query = query.reshape(b_q, num_ulysses_shards, h_q // num_ulysses_shards, s_loc_q, d_q) + query = a2a(query, split_axis=1, concat_axis=1) + else: + query = a2a(query, split_axis=1, concat_axis=2) key = a2a(key, split_axis=1, concat_axis=2) value = a2a(value, split_axis=1, concat_axis=2) @@ -2036,19 +2303,27 @@ def wrap_ulysses_ring_attention(query, key, value): raw_key = key raw_query = query raw_value = value - context_q_seq_len = raw_query.shape[2] - actual_kv_seq_len = orig_kv_seq_len if num_ring_shards == 1 else raw_key.shape[2] + context_q_seq_len = raw_query.shape[1] * raw_query.shape[3] if use_shard_major else raw_query.shape[2] + actual_kv_seq_len = ( + (token_padding.real_len if token_padding is not None else orig_kv_seq_len) + if num_ring_shards == 1 + else (token_padding.real_len if token_padding is not None else raw_key.shape[2]) + ) real_key = raw_key[:, :, :actual_kv_seq_len, :] - query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) + if use_shard_major: + kv_size = raw_query.shape[-1] + query_seq_len = context_q_seq_len + else: + query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) # When actual_kv_seq_len is aligned to 8 sublanes, K/V are passed with NO # sequence padding. The kernel slices the ragged KV tail using slice lengths # derived from actual_kv_seq_len, avoiding sequence pad HBM copies and # redundant ppermute/MXU compute on padded keys. kv_pad_size = 1 if actual_kv_seq_len % 8 == 0 else bkv - key, _, key_seq_len = _pad_data_for_flash(raw_key, kv_heads, kv_pad_size) + key, _, _ = _pad_data_for_flash(raw_key, kv_heads, kv_pad_size) value, _, _ = _pad_data_for_flash(raw_value, kv_heads, kv_pad_size) - ring_kv_seq_len = actual_kv_seq_len if actual_kv_seq_len % 8 == 0 else key_seq_len + ring_kv_seq_len = actual_kv_seq_len k_mean = None if use_fixed_m and num_ring_shards == 1 and use_k_centering: @@ -2068,7 +2343,7 @@ def wrap_ulysses_ring_attention(query, key, value): recenter=recenter, per_q_block=per_q_block, k_mean=k_mean, - value=raw_value, + value=raw_value if token_padding is None else raw_value[:, :, :actual_kv_seq_len, :], ) bsizes = _build_custom_splash_block_sizes(flash_block_sizes, bq, bkv, bkv_compute, bkv_compute_in) @@ -2146,6 +2421,7 @@ def _run_hybrid(q, k, v, m, km): pregathered_mk=True, v_ok=v_ok, all_fixed_global=all_fixed_global, + out_num_shards=out_num_shards, ) attention_output = jax.vmap(ring_kernel, in_axes=(0, 0, 0, (0, 0), 0))( query, key, value, (qn_dev, mk_global_sq_dev), k_mean_dev @@ -2163,13 +2439,27 @@ def _run_hybrid(q, k, v, m, km): bidirectional=bidirectional, use_fixed_m=False, residual_checkpoint_name=residual_checkpoint_name, + out_num_shards=out_num_shards, ) attention_output = jax.vmap(ring_kernel, in_axes=(0, 0, 0))(query, key, value) + if use_shard_major: + attention_output = attention_output.astype(query.dtype) + attention_output = a2a(attention_output, split_axis=1, concat_axis=1) + return attention_output.reshape( + (attention_output.shape[0], attention_output.shape[1] * attention_output.shape[2]) + attention_output.shape[3:] + ) + attention_output = attention_output[:, :, :context_q_seq_len, :kv_size].astype(query.dtype) # (3) Ulysses All-to-All back: sequence -> heads - return a2a(attention_output, split_axis=2, concat_axis=1) + return _ulysses_seq_to_heads( + attention_output, + axis_name=ulysses_axis, + seq_axis=2, + num_shards=num_ulysses_shards, + mode=eff_out_a2a_mode, + ) x = _run_chunked_ulysses_attention( query, @@ -2223,6 +2513,7 @@ def _to_bshd(x: Array, n_heads: int) -> Array: query_states = _to_bshd(query, heads) key_states = _to_bshd(key, effective_kv_heads) value_states = _to_bshd(value, effective_kv_heads) + _reject_active_token_padding(query_states.shape[1], key_states.shape[1], "dot_product") if heads != effective_kv_heads: num_repeats = heads // effective_kv_heads key_states = jnp.repeat(key_states, num_repeats, axis=2) @@ -2231,6 +2522,7 @@ def _to_bshd(x: Array, n_heads: int) -> Array: query_states = _reshape_heads_to_batch_dim(query, heads) key_states = _reshape_heads_to_batch_dim(key, effective_kv_heads) value_states = _reshape_heads_to_batch_dim(value, effective_kv_heads) + _reject_active_token_padding(query_states.shape[1], key_states.shape[1], "dot_product") if heads != effective_kv_heads: num_repeats = heads // effective_kv_heads b = query.shape[0] @@ -2340,6 +2632,7 @@ def _cudnn_flash_attention(query: Array, key: Array, value: Array, heads: int, m query = _reshape_data_for_cudnn_flash(query, heads) key = _reshape_data_for_cudnn_flash(key, heads) value = _reshape_data_for_cudnn_flash(value, heads) + _reject_active_token_padding(query.shape[1], key.shape[1], "cudnn_flash_te") query = nn.with_logical_constraint(query, (BATCH, LENGTH, HEAD, D_KV)) key = nn.with_logical_constraint(key, (BATCH, LENGTH, HEAD, D_KV)) @@ -2409,6 +2702,7 @@ def ulysses_custom_kernel(q, k, v, context): kernel_name="ulysses_custom", qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2435,6 +2729,7 @@ def ulysses_ring_custom_kernel(q, k, v, context): kv_heads=context.get("kv_heads", None), qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2465,6 +2760,7 @@ def ulysses_ring_custom_fixed_m_kernel(q, k, v, context): use_k_centering=resolve_k_centering(context.get("use_k_centering"), ring=True), qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2495,6 +2791,7 @@ def ulysses_ring_custom_fixed_m_per_q_block_kernel(q, k, v, context): use_k_centering=resolve_k_centering(context.get("use_k_centering"), ring=True), qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2525,6 +2822,7 @@ def ulysses_ring_custom_bidir_kernel(q, k, v, context): kv_heads=context.get("kv_heads", None), qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2556,6 +2854,7 @@ def ulysses_custom_fixed_m_kernel(q, k, v, context): use_k_centering=resolve_k_centering(context.get("use_k_centering"), ring=False), qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2587,6 +2886,7 @@ def ulysses_custom_fixed_m_per_q_block_kernel(q, k, v, context): use_k_centering=resolve_k_centering(context.get("use_k_centering"), ring=False), qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2610,6 +2910,7 @@ def ulysses_kernel(q, k, v, context): kv_heads=context.get("kv_heads", None), ulysses_shards=context.get("ulysses_shards", -1), kernel_name="ulysses", + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2752,6 +3053,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, @@ -2787,6 +3121,7 @@ def _apply_attention( qk_prescaled: bool = False, k_prescaled: bool = False, transpose_out: bool = False, + wan_ulysses_out_a2a: str = "flat", ): """Routes to different attention kernels using a module-level registry.""" @@ -2796,19 +3131,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 @@ -2872,6 +3195,7 @@ def _apply_attention( "qk_prescaled": qk_prescaled, "k_prescaled": k_prescaled, "transpose_out": transpose_out, + "wan_ulysses_out_a2a": wan_ulysses_out_a2a, } if spatiotemporal_config and spatiotemporal_config.get("use_svg_attention"): @@ -2910,6 +3234,7 @@ def _head_local_svg_attention(query, key, value, context): if context["ulysses_attention_chunks"] != 1: raise ValueError("Head-local SVG does not implement chunked Ulysses attention.") q, k, v = (_unflatten_heads(x, context["heads"]) if x.ndim == 3 else x for x in (query, key, value)) + _reject_active_token_padding(q.shape[2], k.shape[2], "head_local_svg") if grid is None or q.shape != k.shape or q.shape != v.shape or q.shape[2] != math.prod(grid): raise ValueError("Head-local SVG requires matched self-attention QKV and token grid.") batch, heads = q.shape[0], q.shape[1] @@ -3230,6 +3555,7 @@ def __init__( kv_heads: Optional[int] = None, use_k_centering: bool | str = "auto", transpose_out: bool = False, + wan_ulysses_out_a2a: str = "flat", ): self.dpa_layer = None self.use_base2_exp = use_base2_exp @@ -3238,6 +3564,7 @@ def __init__( self.ulysses_attention_chunks = ulysses_attention_chunks self.use_k_centering = use_k_centering self.transpose_out = bool(transpose_out) + self.wan_ulysses_out_a2a = wan_ulysses_out_a2a or "flat" if attention_kernel == "cudnn_flash_te": from transformer_engine.jax.flax.transformer import DotProductAttention # pytype: disable=import-error @@ -3322,6 +3649,7 @@ def apply_attention( qk_prescaled=qk_prescaled, k_prescaled=k_prescaled, transpose_out=getattr(self, "transpose_out", False), + wan_ulysses_out_a2a=getattr(self, "wan_ulysses_out_a2a", "flat"), ) @@ -3348,6 +3676,7 @@ class AttentionOp(nn.Module): kv_heads: Optional[int] = None use_k_centering: bool | str = "auto" transpose_out: bool = False + wan_ulysses_out_a2a: str = "flat" def setup(self): self.dpa_layer = None @@ -3417,6 +3746,7 @@ def apply_attention( qk_prescaled=qk_prescaled, k_prescaled=k_prescaled, transpose_out=getattr(self, "transpose_out", False), + wan_ulysses_out_a2a=getattr(self, "wan_ulysses_out_a2a", "flat"), ) @@ -3480,6 +3810,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__( @@ -3534,6 +3919,9 @@ def __init__( "wan_splash_transpose_out": None, "wan_cross_attn_prescale_kv": None, "wan_rope_accum": None, + "wan_cross_attn_kernel": None, + "wan_ulysses_out_a2a": None, + "wan_cross_attn_cpu_interpret": False, **(attention_config or {}), } @@ -3565,6 +3953,9 @@ 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_ulysses_out_a2a = _wan_opt("wan_ulysses_out_a2a") + 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}") @@ -3638,6 +4029,7 @@ def _wan_opt(name): ulysses_attention_chunks=attention_config["ulysses_attention_chunks"], use_k_centering=attention_config["use_k_centering"], transpose_out=self.wan_splash_transpose_out, + wan_ulysses_out_a2a=self.wan_ulysses_out_a2a, ) # None axes corresponds to the stacked weights across all blocks # because of the use of nnx.vmap and nnx.scan. @@ -3970,6 +4362,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( + "pallas_cross_attn_active", + 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() @@ -4109,7 +4620,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 @@ -4186,9 +4701,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( @@ -4207,9 +4726,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/models/wan/transformers/transformer_wan.py b/src/maxdiffusion/models/wan/transformers/transformer_wan.py index db083ae0c..dc8a76127 100644 --- a/src/maxdiffusion/models/wan/transformers/transformer_wan.py +++ b/src/maxdiffusion/models/wan/transformers/transformer_wan.py @@ -35,7 +35,13 @@ NNXPixArtAlphaTextProjection, ) from ...normalization_flax import FP32LayerNorm -from ...attention_flax import FlaxWanAttention +from ...attention_flax import ( + FlaxWanAttention, + TokenPadding, + TOKEN_PADDING_KERNELS, + self_attention_token_padding, + _warn_once, +) from ...gradient_checkpoint import GradientCheckpointType BlockSizes = common_types.BlockSizes @@ -61,6 +67,62 @@ def get_frequencies(max_seq_len: int, theta: int, attention_head_dim: int): return freqs_split +PATCH_EMBED_MODES = ("conv", "tokens") + + +def can_embed_patches_as_tokens(conv: nnx.Conv, hidden_shape: Tuple[int, ...], patch_size: Tuple[int, int, int]) -> bool: + """Whether `embed_patches_as_tokens` computes exactly what `conv` does.""" + + def per_dim(value): + value = 1 if value is None else value + return tuple(value) if isinstance(value, (tuple, list)) else (value,) * len(patch_size) + + if len(hidden_shape) != 5 or per_dim(conv.kernel_size) != tuple(patch_size): + return False + if per_dim(conv.strides) != tuple(patch_size) or conv.feature_group_count != 1 or conv.mask is not None: + return False + if per_dim(conv.input_dilation) != (1,) * 3 or per_dim(conv.kernel_dilation) != (1,) * 3: + return False + # With every spatial size divisible by its stride, 'SAME' pads nothing. + return conv.padding in ("SAME", "VALID") and all(n % p == 0 for n, p in zip(hidden_shape[1:4], patch_size)) + + +def embed_patches_as_tokens(conv: nnx.Conv, hidden_states: jax.Array, patch_size: Tuple[int, int, int]) -> jax.Array: + """Wan patch embedding as a token-sharded matmul: `[B, F, H, W, C]` -> `[B, N, D]`. + + Same result as `jax.lax.collapse(conv(hidden_states), 1, -1)` up to + accumulation order, but produced in the layout the transformer blocks use. + When the 3D `nnx.Conv` runs before `collapse`, XLA either shards the conv output + differently from `activation_length` or emits an activation all-gather / + collective-permute when collapsing `(F, H, W)` into the context-sharded token + axis (12.8 ms per step on tpu7x, exposed). Patches do not overlap (stride == + kernel), so each token needs only its own patch: flatten the patches in the + kernel's `(p_t, p_h, p_w, C)` order, shard them by token and multiply by the + replicated `[p_t*p_h*p_w*C, D]` kernel. + """ + p_t, p_h, p_w = patch_size + batch, frames, height, width, channels = hidden_states.shape + patches = hidden_states.reshape(batch, frames // p_t, p_t, height // p_h, p_h, width // p_w, p_w, channels) + patches = jnp.transpose(patches, (0, 1, 3, 5, 2, 4, 6, 7)) + patches = patches.reshape(batch, -1, p_t * p_h * p_w * channels) + patches = nn.with_logical_constraint(patches, ("activation_batch", "activation_length", None)) + + kernel = conv.kernel[...].reshape(p_t * p_h * p_w * channels, -1) + kernel = nn.with_logical_constraint(kernel, (None, None)) + bias = conv.bias[...] if conv.bias is not None else None + # Mirror nnx.Conv: promote everything to `conv.dtype`, no preferred type. + dtype = conv.dtype or jnp.result_type(patches, kernel) + out = jax.lax.dot_general( + patches.astype(dtype), + kernel.astype(dtype), + (((2,), (0,)), ((), ())), + precision=conv.precision, + ) + if bias is not None: + out = out + bias.astype(dtype) + return out + + class WanRotaryPosEmbed(nnx.Module): def __init__( @@ -608,6 +670,18 @@ def __init__( "use_k_centering": "auto", **(attention_config or {}), } + self.attention = attention + self.attention_config = attention_config + self.flash_min_seq_length = flash_min_seq_length + self.mesh = mesh + + # Resolved at build time so they live on the GraphDef; None -> built-in default. + def _wan_opt(name): + value = attention_config.get(name) + return wan_runtime_options.default(name) if value is None else wan_runtime_options.coerce(name, value) + + self.patch_embed_mode = _wan_opt("wan_patch_embed_mode") + self.seq_pad_mode = _wan_opt("wan_seq_pad") # 1. Patch & position embedding self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) @@ -778,6 +852,126 @@ def _compute_kv(block, enc_states, enc_mask): return kv_cache, encoder_attention_mask + def _get_token_padding(self, seq_len: int, per_token_t: bool = False) -> Optional[TokenPadding]: + mode = getattr(self, "seq_pad_mode", None) or wan_runtime_options.default("wan_seq_pad") + if mode in (None, "off"): + return None + if mode != "lane": + raise ValueError(f"wan_seq_pad must be 'off' or 'lane', got {mode!r}.") + if per_token_t: + _warn_once( + "wan_seq_pad_per_token_t", + "wan_seq_pad='lane' is not currently supported with per_token_t=True (TI2V); disabling token padding.", + ) + return None + attention_config = getattr(self, "attention_config", None) or {} + if getattr(self, "use_svg_attention", False) or bool(attention_config.get("use_svg_attention", False)): + _warn_once( + "wan_seq_pad_svg_attention", + "wan_seq_pad='lane' is incompatible with use_svg_attention=True; disabling token padding.", + ) + return None + blocks = getattr(self, "blocks", None) + first_block = blocks[0] if isinstance(blocks, (list, tuple)) and blocks else blocks + attn1_op = getattr(getattr(first_block, "attn1", None), "attention_op", None) + if bool(attention_config.get("use_memory_efficient_attention", False)) or bool( + getattr(attn1_op, "use_memory_efficient_attention", False) + ): + _warn_once( + "wan_seq_pad_mem_eff_attn", + "wan_seq_pad='lane' is incompatible with use_memory_efficient_attention=True; disabling token padding.", + ) + return None + flash_min_seq_length = attention_config.get("flash_min_seq_length", getattr(self, "flash_min_seq_length", 4096)) + if seq_len < flash_min_seq_length: + _warn_once( + "wan_seq_pad_short_seq", + f"wan_seq_pad='lane' disabled for seq_len={seq_len} < flash_min_seq_length={flash_min_seq_length}.", + ) + return None + if self.attention not in TOKEN_PADDING_KERNELS: + _warn_once( + "wan_seq_pad_unsupported_attn", + f"wan_seq_pad='lane' requested, but attention={self.attention!r} does not support TokenPadding masking. " + "Disabling token padding.", + ) + return None + + num_context_shards = 1 + if self.mesh is not None and hasattr(self.mesh, "shape") and "context" in self.mesh.shape: + num_context_shards = self.mesh.shape["context"] + ulysses_shards = self.attention_config.get("ulysses_shards", -1) + if self.attention.startswith("ulysses_ring_"): + if ulysses_shards is None or ulysses_shards <= 0 or num_context_shards % ulysses_shards != 0: + _warn_once( + "wan_seq_pad_ring_ulysses_shards", + f"wan_seq_pad='lane' requires ulysses_shards > 0 dividing the {num_context_shards} context shards for " + f"attention={self.attention!r}, got ulysses_shards={ulysses_shards}; disabling token padding.", + ) + return None + num_ring_shards = num_context_shards // ulysses_shards + elif self.attention.startswith("ulysses_"): + if ulysses_shards is not None and ulysses_shards > 0 and ulysses_shards != num_context_shards: + _warn_once( + "wan_seq_pad_ulysses_shards_mismatch", + f"wan_seq_pad='lane' requires ulysses_shards to match the {num_context_shards} context shards for " + f"attention={self.attention!r}, got ulysses_shards={ulysses_shards}; disabling token padding.", + ) + return None + ulysses_shards = num_context_shards + num_ring_shards = 1 + else: + ulysses_shards = 1 + num_ring_shards = 1 + num_segments = num_ring_shards + + total_shards = num_segments * ulysses_shards + if total_shards <= 0 or seq_len % total_shards != 0: + _warn_once( + "wan_seq_pad_indivisible_seq", + f"wan_seq_pad='lane' requires seq_len={seq_len} to split evenly across {total_shards} attention shards " + f"({num_segments} ring x {ulysses_shards} ulysses); disabling token padding.", + ) + return None + + real_len = seq_len // num_segments + tokens_per_shard = real_len // ulysses_shards + lane_multiple = 128 + padded_tokens_per_shard = math.ceil(tokens_per_shard / lane_multiple) * lane_multiple + padded_len = padded_tokens_per_shard * ulysses_shards + if padded_len == real_len: + return None + return TokenPadding(real_len=real_len, padded_len=padded_len, num_segments=num_segments) + + @staticmethod + def _pad_tokens(x: jax.Array, pad: TokenPadding) -> jax.Array: + if pad.padded_len == pad.real_len: + return x + b, _, d = x.shape + x = x.reshape(b, pad.num_segments, pad.real_len, d) + x = jnp.pad(x, ((0, 0), (0, 0), (0, pad.padded_len - pad.real_len), (0, 0))) + x = x.reshape(b, pad.total_len, d) + return nn.with_logical_constraint(x, ("activation_batch", "activation_length", "activation_heads")) + + @staticmethod + def _unpad_tokens(x: jax.Array, pad: TokenPadding) -> jax.Array: + if pad.padded_len == pad.real_len: + return x + b, _, d = x.shape + x = x.reshape(b, pad.num_segments, pad.padded_len, d) + x = x[:, :, : pad.real_len, :] + x = x.reshape(b, pad.real_len * pad.num_segments, d) + return nn.with_logical_constraint(x, ("activation_batch", "activation_length", "activation_heads")) + + @staticmethod + def _pad_rotary_emb(rotary_emb: jax.Array, pad: TokenPadding) -> jax.Array: + if pad.padded_len == pad.real_len: + return rotary_emb + dim = rotary_emb.shape[-1] + r = rotary_emb.reshape(1, 1, pad.num_segments, pad.real_len, dim) + r = jnp.pad(r, ((0, 0), (0, 0), (0, 0), (0, pad.padded_len - pad.real_len), (0, 0))) + return r.reshape(1, 1, pad.total_len, dim) + @jax.named_scope("WanModel") def __call__( self, @@ -810,8 +1004,16 @@ def __call__( if rotary_emb is None: rotary_emb = self.rope(hidden_states) with self.conditional_named_scope("patch_embedding"): - hidden_states = self.patch_embedding(hidden_states) - hidden_states = jax.lax.collapse(hidden_states, 1, -1) + patch_embed_mode = getattr(self, "patch_embed_mode", None) or wan_runtime_options.default("wan_patch_embed_mode") + if patch_embed_mode not in PATCH_EMBED_MODES: + raise ValueError(f"wan_patch_embed_mode must be one of {PATCH_EMBED_MODES}, got {patch_embed_mode!r}.") + if patch_embed_mode == "tokens" and can_embed_patches_as_tokens( + self.patch_embedding, hidden_states.shape, self.config.patch_size + ): + hidden_states = embed_patches_as_tokens(self.patch_embedding, hidden_states, self.config.patch_size) + else: + hidden_states = self.patch_embedding(hidden_states) + hidden_states = jax.lax.collapse(hidden_states, 1, -1) per_token_t = timestep.ndim == 2 # [B, seq_len] for TI2V with self.conditional_named_scope("condition_embedder"): if per_token_t: @@ -866,6 +1068,11 @@ def __call__( else: encoder_hidden_states = encoder_hidden_states_out.astype(hidden_states.dtype) + token_padding = self._get_token_padding(hidden_states.shape[1], per_token_t=per_token_t) + if token_padding is not None: + hidden_states = self._pad_tokens(hidden_states, token_padding) + rotary_emb = self._pad_rotary_emb(rotary_emb, token_padding) + def _run_all_blocks(h): if self.scan_layers: @@ -955,9 +1162,12 @@ def layer_forward(hidden_states, l_kv): if skip_blocks: if cached_residual is None: raise ValueError("cached_residual must be provided when skip_blocks is True") + if token_padding is not None and cached_residual.shape[1] == token_padding.real_len * token_padding.num_segments: + cached_residual = self._pad_tokens(cached_residual, token_padding) hidden_states = hidden_states + cached_residual else: - hidden_states = _run_all_blocks(hidden_states) + with self_attention_token_padding(token_padding): + hidden_states = _run_all_blocks(hidden_states) residual_x = hidden_states - hidden_states_before_blocks @@ -974,6 +1184,10 @@ def layer_forward(hidden_states, l_kv): with jax.named_scope("proj_out"): hidden_states = self.proj_out(hidden_states) + if token_padding is not None: + hidden_states = self._unpad_tokens(hidden_states, token_padding) + residual_x = self._unpad_tokens(residual_x, token_padding) + if not unpatchify: if return_residual: return hidden_states, residual_x 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..af2ca93db --- /dev/null +++ b/src/maxdiffusion/tests/cross_attention_pallas_test.py @@ -0,0 +1,421 @@ +""" +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) + + +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) + @_SKIP_IN_GITHUB_ACTIONS + 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", cpu_interpret=True) + 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) + + @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)) + + @_SKIP_IN_GITHUB_ACTIONS + 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/ulysses_out_a2a_test.py b/src/maxdiffusion/tests/ulysses_out_a2a_test.py new file mode 100644 index 000000000..c36530c81 --- /dev/null +++ b/src/maxdiffusion/tests/ulysses_out_a2a_test.py @@ -0,0 +1,279 @@ +""" +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 inverse Ulysses exchange (`wan_ulysses_out_a2a`). + +Inside the shard_map each shard holds [B, H/n, ..., S] (its heads, full +sequence) and must end with [B, H, ..., S/n] (every head, its sequence chunk). +Seen globally that is a pure resharding, so the output must equal the input +exactly, in both modes. Needs >= 2 devices; on CPU run with +XLA_FLAGS=--xla_force_host_platform_device_count=4. +""" + +import os +import unittest + +import jax +import jax.numpy as jnp +import numpy as np +from jax.sharding import Mesh, PartitionSpec as P + +from maxdiffusion.kernels import custom_splash_attention as custom_splash +from maxdiffusion.models import attention_flax + +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" +) + + +class UlyssesSeqToHeadsTest(unittest.TestCase): + + def _require_devices(self, min_devices: int = 2) -> int: + n = len(jax.devices()) + if n < min_devices: + self.skipTest(f"needs >= {min_devices} devices (XLA_FLAGS=--xla_force_host_platform_device_count=4 on CPU)") + return 4 if n >= 4 else 2 + + def _exchange(self, mode, x, seq_axis): + n = self._require_devices(2) + mesh = Mesh(np.array(jax.devices()[:n]), ("context",)) + in_spec = P(None, "context", *([None] * (x.ndim - 2))) + out_spec = [None] * x.ndim + out_spec[seq_axis] = "context" + fn = jax.shard_map( + lambda t: attention_flax._ulysses_seq_to_heads(t, axis_name="context", seq_axis=seq_axis, num_shards=n, mode=mode), + mesh=mesh, + in_specs=(in_spec,), + out_specs=P(*out_spec), + check_vma=False, + ) + return jax.jit(fn)(x) + + def test_both_modes_are_a_pure_reshard(self): + n = self._require_devices(2) + heads, d, seq = 2 * n, 128, 24 * n + shapes = {3: (1, heads, d, seq), 2: (1, heads, seq, d)} # [B,H,D,S] (kernel O^T) and [B,H,S,D] + for seq_axis, shape in shapes.items(): + x = jax.random.normal(jax.random.PRNGKey(seq_axis), shape, jnp.float32).astype(jnp.bfloat16) + for mode in attention_flax.ULYSSES_OUT_A2A_MODES: + with self.subTest(seq_axis=seq_axis, mode=mode): + got = self._exchange(mode, x, seq_axis) + self.assertEqual(got.shape, x.shape) + np.testing.assert_array_equal(np.asarray(got, np.float32), np.asarray(x, np.float32)) + + def test_rejects_unknown_mode(self): + n = self._require_devices(2) + x = jnp.zeros((1, 2 * n, 8 * n, 128), jnp.bfloat16) + with self.assertRaisesRegex(ValueError, "wan_ulysses_out_a2a"): + self._exchange("bogus", x, 2) + + @_SKIP_IN_GITHUB_ACTIONS + def test_pallas_shard_major_q_and_out_ring_kernel_bit_exact(self): + u, h, s_loc, d = 2, 2, 256, 128 + s_total = u * s_loc + bs = custom_splash._BlockSizes(block_q=128, block_kv=128, block_kv_compute=128, block_kv_compute_in=128) + q_3d = jax.random.normal(jax.random.PRNGKey(10), (h, s_total, d), jnp.bfloat16) + k = jax.random.normal(jax.random.PRNGKey(11), (h, s_total, d), jnp.bfloat16) * 0.08 + v = jax.random.normal(jax.random.PRNGKey(12), (h, s_total, d), jnp.bfloat16) + q_4d = jnp.swapaxes(q_3d.reshape(h, u, s_loc, d), 0, 1) + + grid_h = s_total // bs.block_q + mk = jnp.stack([jnp.full((h, grid_h), -10.0, jnp.float32), jnp.ones((h, grid_h), jnp.float32)], axis=0) + + o_ref, m_ref, l_ref = custom_splash._splash_attention_forward_ring( + q_3d, + k, + v, + bs, + q_seq_len=s_total, + kv_seq_len=s_total, + use_fixed_m=True, + mk=mk, + uniform_fixed_m=True, + interpret=True, + out_num_shards=1, + ) + o_sm, m_sm, l_sm = custom_splash._splash_attention_forward_ring( + q_4d, + k, + v, + bs, + q_seq_len=s_total, + kv_seq_len=s_total, + use_fixed_m=True, + mk=mk, + uniform_fixed_m=True, + interpret=True, + out_num_shards=u, + ) + self.assertEqual(o_sm.shape, (u, h, s_loc, d)) + self.assertEqual(m_sm.shape, (u, h, s_loc)) + self.assertEqual(l_sm.shape, (u, h, s_loc)) + np.testing.assert_array_equal(np.asarray(jnp.swapaxes(o_sm, 0, 1).reshape(h, s_total, d)), np.asarray(o_ref)) + np.testing.assert_array_equal(np.asarray(jnp.swapaxes(m_sm, 0, 1).reshape(h, s_total)), np.asarray(m_ref)) + np.testing.assert_array_equal(np.asarray(jnp.swapaxes(l_sm, 0, 1).reshape(h, s_total)), np.asarray(l_ref)) + + @_SKIP_IN_GITHUB_ACTIONS + def test_ulysses_ring_custom_chunked_matches_flat_bit_exact(self): + from flax.linen import partitioning as nn_partitioning + + self._require_devices(4) + mesh_4 = Mesh(np.array(jax.devices()[:4]).reshape(1, 1, 4, 1), ("data", "fsdp", "context", "tensor")) + rules = ( + (attention_flax.BATCH, "data"), + (attention_flax.SELF_ATTN_HEAD, None), + (attention_flax.SELF_ATTN_Q_LENGTH, "context"), + (attention_flax.SELF_ATTN_KV_LENGTH, "context"), + (attention_flax.D_KV, None), + ) + b, heads, s_per_shard, d = 1, 4, 256, 128 + s_total = 4 * s_per_shard + q = jax.random.normal(jax.random.PRNGKey(20), (b, s_total, heads * d), jnp.bfloat16) + k = jax.random.normal(jax.random.PRNGKey(21), (b, s_total, heads * d), jnp.bfloat16) * 0.08 + v = jax.random.normal(jax.random.PRNGKey(22), (b, s_total, heads * d), jnp.bfloat16) + bs = {"block_q": 128, "block_kv": 128, "block_kv_compute": 128, "block_kv_compute_in": 128} + axis_names_q = ( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_Q_LENGTH, + attention_flax.D_KV, + ) + axis_names_kv = ( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_KV_LENGTH, + attention_flax.D_KV, + ) + for use_fixed_m, per_q_block, bidirectional, mode in ( + (True, False, False, "chunked"), + (True, True, False, "shard_major"), + (False, False, False, "chunked"), + (False, False, True, "shard_major"), + ): + with self.subTest(use_fixed_m=use_fixed_m, per_q_block=per_q_block, bidirectional=bidirectional, mode=mode): + with mesh_4, nn_partitioning.axis_rules(rules): + out_flat = attention_flax._ulysses_ring_custom_attention( + q, + k, + v, + heads=heads, + mesh=mesh_4, + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + flash_block_sizes=bs, + ulysses_shards=2, + use_fixed_m=use_fixed_m, + per_q_block=per_q_block, + bidirectional=bidirectional, + wan_ulysses_out_a2a="flat", + ) + out_chunked = attention_flax._ulysses_ring_custom_attention( + q, + k, + v, + heads=heads, + mesh=mesh_4, + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + flash_block_sizes=bs, + ulysses_shards=2, + use_fixed_m=use_fixed_m, + per_q_block=per_q_block, + bidirectional=bidirectional, + wan_ulysses_out_a2a=mode, + ) + self.assertEqual(out_chunked.shape, out_flat.shape) + np.testing.assert_array_equal(np.asarray(out_chunked, jnp.float32), np.asarray(out_flat, jnp.float32)) + + @_SKIP_IN_GITHUB_ACTIONS + def test_ulysses_custom_chunked_matches_flat_bit_exact(self): + from flax.linen import partitioning as nn_partitioning + + n = self._require_devices(2) + mesh_u = Mesh(np.array(jax.devices()[:n]).reshape(1, 1, n, 1), ("data", "fsdp", "context", "tensor")) + rules = ( + (attention_flax.BATCH, "data"), + (attention_flax.SELF_ATTN_HEAD, None), + (attention_flax.SELF_ATTN_Q_LENGTH, "context"), + (attention_flax.SELF_ATTN_KV_LENGTH, "context"), + (attention_flax.D_KV, None), + (attention_flax.LENGTH, "context"), + (attention_flax.HEAD, None), + ) + b, heads, s_per_shard, d = 1, 2 * n, 256, 128 + s_total = n * s_per_shard + q = jax.random.normal(jax.random.PRNGKey(30), (b, s_total, heads * d), jnp.bfloat16) + k = jax.random.normal(jax.random.PRNGKey(31), (b, s_total, heads * d), jnp.bfloat16) * 0.08 + v = jax.random.normal(jax.random.PRNGKey(32), (b, s_total, heads * d), jnp.bfloat16) + bs = {"block_q": 128, "block_kv": 128, "block_kv_compute": 128, "block_kv_compute_in": 128} + axis_names_q = ( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_Q_LENGTH, + attention_flax.D_KV, + ) + axis_names_kv = ( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_KV_LENGTH, + attention_flax.D_KV, + ) + for use_fixed_m, per_q_block, use_k_centering, mode in ( + (True, True, True, "chunked"), + (True, False, False, "shard_major"), + (False, False, False, "chunked"), + ): + with self.subTest(use_fixed_m=use_fixed_m, per_q_block=per_q_block, use_k_centering=use_k_centering, mode=mode): + with mesh_u, nn_partitioning.axis_rules(rules): + out_flat = attention_flax._ulysses_attention( + q, + k, + v, + heads=heads, + mesh=mesh_u, + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + flash_block_sizes=bs, + use_custom_kernel=True, + use_fixed_m=use_fixed_m, + per_q_block=per_q_block, + use_k_centering=use_k_centering, + wan_ulysses_out_a2a="flat", + ) + out_chunked = attention_flax._ulysses_attention( + q, + k, + v, + heads=heads, + mesh=mesh_u, + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + flash_block_sizes=bs, + use_custom_kernel=True, + use_fixed_m=use_fixed_m, + per_q_block=per_q_block, + use_k_centering=use_k_centering, + wan_ulysses_out_a2a=mode, + ) + self.assertEqual(out_chunked.shape, out_flat.shape) + np.testing.assert_array_equal(np.asarray(out_chunked, jnp.float32), np.asarray(out_flat, jnp.float32)) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/wan_patch_embed_test.py b/src/maxdiffusion/tests/wan_patch_embed_test.py new file mode 100644 index 000000000..c64029400 --- /dev/null +++ b/src/maxdiffusion/tests/wan_patch_embed_test.py @@ -0,0 +1,184 @@ +""" +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 token-sharded Wan patch embedding (`wan_patch_embed_mode`).""" + +import os +import unittest +from unittest import mock + +import jax +import jax.numpy as jnp +import numpy as np +from flax import nnx +from flax.linen import partitioning as nn_partitioning +from jax.sharding import Mesh + +from maxdiffusion import pyconfig +from maxdiffusion.max_utils import create_device_mesh +from maxdiffusion.models.wan.transformers import transformer_wan +from maxdiffusion.models.wan.transformers.transformer_wan import ( + WanModel, + can_embed_patches_as_tokens, + embed_patches_as_tokens, +) + +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" +) + +THIS_DIR = os.path.dirname(os.path.abspath(__file__)) + + +def _conv(patch_size, in_channels=4, out_channels=256, dtype=jnp.float32, precision=jax.lax.Precision.HIGHEST, **kw): + return nnx.Conv( + in_channels, + out_channels, + kernel_size=patch_size, + strides=patch_size, + dtype=dtype, + param_dtype=dtype, + precision=precision, + rngs=nnx.Rngs(0), + bias_init=nnx.initializers.normal(0.5), + **kw, + ) + + +class PatchEmbedTokensTest(unittest.TestCase): + + 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 _both(self, conv, x, patch_size): + def conv_path(x): + return jax.lax.collapse(conv(x), 1, -1) + + def token_path(x): + return embed_patches_as_tokens(conv, x, patch_size) + + with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + return jax.jit(conv_path)(x), jax.jit(token_path)(x) + + def test_matches_conv(self): + for patch_size in ((1, 2, 2), (2, 2, 2)): + with self.subTest(patch_size=patch_size): + conv = _conv(patch_size) + x = jax.random.normal(jax.random.key(1), (2, 4, 8, 16, 4), jnp.float32) + want, got = self._both(conv, x, patch_size) + self.assertEqual(got.shape, want.shape) + np.testing.assert_allclose(np.asarray(got), np.asarray(want), rtol=1e-5, atol=1e-5) + + def test_matches_conv_in_bf16(self): + patch_size = (1, 2, 2) + conv = _conv(patch_size, in_channels=16, dtype=jnp.bfloat16, precision=None) + x = jax.random.normal(jax.random.key(2), (2, 3, 8, 16, 16), jnp.float32).astype(jnp.bfloat16) + want, got = self._both(conv, x, patch_size) + self.assertEqual(got.dtype, want.dtype) + np.testing.assert_allclose(np.asarray(got, np.float32), np.asarray(want, np.float32), rtol=2e-2, atol=2e-2) + + def test_guard(self): + conv = _conv((1, 2, 2)) + self.assertTrue(can_embed_patches_as_tokens(conv, (1, 21, 90, 160, 4), (1, 2, 2))) + # A spatial size not divisible by the patch makes 'SAME' pad. + self.assertFalse(can_embed_patches_as_tokens(conv, (1, 21, 91, 160, 4), (1, 2, 2))) + self.assertFalse(can_embed_patches_as_tokens(conv, (1, 21, 90, 160, 4), (1, 1, 2))) + overlapping = nnx.Conv(4, 8, kernel_size=(1, 2, 2), strides=(1, 1, 1), rngs=nnx.Rngs(0)) + self.assertFalse(can_embed_patches_as_tokens(overlapping, (1, 21, 90, 160, 4), (1, 2, 2))) + dilated = _conv((1, 2, 2), kernel_dilation=(1, 2, 2)) + self.assertFalse(can_embed_patches_as_tokens(dilated, (1, 21, 90, 160, 4), (1, 2, 2))) + + +class WanModelPatchEmbedModeTest(unittest.TestCase): + + 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) + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + self.model = WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=2, + attention_head_dim=128, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=256, + num_layers=1, + mesh=self.mesh, + attention="dot_product", + ) + self.inputs = { + "hidden_states": jax.random.normal(jax.random.key(3), (1, 4, 2, 8, 16), jnp.float32), + "timestep": jnp.ones((1,)), + "encoder_hidden_states": jax.random.normal(jax.random.key(4), (1, 16, 32), jnp.float32), + } + + def _run(self, mode): + # The mode is a build-time attribute (normally from attention_config). + self.model.patch_embed_mode = mode + with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + return self.model(**self.inputs) + + @_SKIP_IN_GITHUB_ACTIONS + def test_tokens_mode_uses_the_token_path_and_matches_conv(self): + want = self._run("conv") + with mock.patch.object(transformer_wan, "embed_patches_as_tokens", wraps=transformer_wan.embed_patches_as_tokens) as spy: + got = self._run("tokens") + spy.assert_called_once() + self.assertEqual(got.shape, want.shape) + np.testing.assert_allclose(np.asarray(got), np.asarray(want), rtol=2e-2, atol=2e-2) + + @_SKIP_IN_GITHUB_ACTIONS + def test_conv_mode_does_not_use_the_token_path(self): + with mock.patch.object(transformer_wan, "embed_patches_as_tokens") as spy: + self._run("conv") + spy.assert_not_called() + + def test_rejects_unknown_mode(self): + with self.assertRaises(ValueError): + self._run("bogus") + + def test_mode_comes_from_attention_config_not_env(self): + with mock.patch.dict(os.environ, {"WAN_PATCH_EMBED_MODE": "tokens"}): + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + model = nnx.eval_shape( + lambda: WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=2, + attention_head_dim=128, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=256, + num_layers=1, + attention="dot_product", + ) + ) + self.assertEqual(model.patch_embed_mode, "conv") + + +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..f436d897e 100644 --- a/src/maxdiffusion/tests/wan_runtime_options_test.py +++ b/src/maxdiffusion/tests/wan_runtime_options_test.py @@ -129,7 +129,9 @@ def test_unset_options_resolve_to_defaults_not_env(self): env.update({"WAN_ROPE_NORM_MODE": "fused", "WAN_ROPE_ACCUM": "f32"}) with mock.patch.dict(os.environ, env): attn = self._attn() - for name in opts.ATTENTION_OPTIONS: + # wan_patch_embed_mode / wan_seq_pad ride in attention_config too, but are + # consumed by WanModel (covered in wan_patch_embed_test / wan_seq_pad_test). + for name in sorted(set(opts.ATTENTION_OPTIONS) - {"wan_patch_embed_mode", "wan_seq_pad"}): with self.subTest(option=name): self.assertEqual(getattr(attn, name), opts.default(name)) @@ -161,6 +163,21 @@ 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"), + "wan_patch_embed_mode": ("conv", "tokens"), + "wan_ulysses_out_a2a": ("flat", "chunked", "shard_major"), + "wan_seq_pad": ("off", "lane"), + } + 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/tests/wan_seq_pad_test.py b/src/maxdiffusion/tests/wan_seq_pad_test.py new file mode 100644 index 000000000..bd40d8bf3 --- /dev/null +++ b/src/maxdiffusion/tests/wan_seq_pad_test.py @@ -0,0 +1,658 @@ +""" +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 token sequence padding and lane alignment (`wan_seq_pad`).""" + +import os +import unittest +from unittest import mock + +import jax +import jax.numpy as jnp +import numpy as np +from flax import nnx +from flax.linen import partitioning as nn_partitioning +from jax.sharding import Mesh + +from maxdiffusion import pyconfig +from maxdiffusion.max_utils import create_device_mesh +from maxdiffusion.models.attention_flax import TokenPadding +from maxdiffusion.models.wan.transformers import transformer_wan +from maxdiffusion.models.wan.transformers.transformer_wan import WanModel + +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" +) + +THIS_DIR = os.path.dirname(os.path.abspath(__file__)) + + +class WanSeqPadTest(unittest.TestCase): + + 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 test_token_padding_inversion(self): + pad = TokenPadding(real_len=37800, padded_len=37888, num_segments=2) + x = jax.random.normal(jax.random.key(0), (1, 75600, 64), dtype=jnp.float32) + padded = WanModel._pad_tokens(x, pad) + self.assertEqual(padded.shape, (1, 75776, 64)) + unpadded = WanModel._unpad_tokens(padded, pad) + self.assertEqual(unpadded.shape, (1, 75600, 64)) + np.testing.assert_array_equal(np.asarray(unpadded), np.asarray(x)) + + def test_rotary_emb_padding(self): + pad = TokenPadding(real_len=75600, padded_len=75776, num_segments=1) + r = jax.random.normal(jax.random.key(1), (1, 1, 75600, 64), dtype=jnp.float32) + padded_r = WanModel._pad_rotary_emb(r, pad) + self.assertEqual(padded_r.shape, (1, 1, 75776, 64)) + np.testing.assert_array_equal(np.asarray(padded_r[:, :, :75600, :]), np.asarray(r)) + np.testing.assert_array_equal(np.asarray(padded_r[:, :, 75600:, :]), 0.0) + + def test_get_token_padding_v6e_and_tpu7x(self): + mock_mesh = mock.MagicMock() + mock_mesh.shape = {"context": 4} + + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + # v6e: ulysses=4, ring=1 + model_v6e = WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=4, + attention_head_dim=64, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=128, + num_layers=1, + mesh=mock_mesh, + attention="ulysses_custom_fixed_m", + attention_config={"ulysses_shards": 4, "wan_seq_pad": "lane"}, + ) + pad_v6e = model_v6e._get_token_padding(75600) + self.assertIsNotNone(pad_v6e) + self.assertEqual(pad_v6e.real_len, 75600) + self.assertEqual(pad_v6e.padded_len, 75776) + self.assertEqual(pad_v6e.num_segments, 1) + + # tpu7x: ulysses=2, ring=2 + model_7x = WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=4, + attention_head_dim=64, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=128, + num_layers=1, + mesh=mock_mesh, + attention="ulysses_ring_custom_fixed_m", + attention_config={"ulysses_shards": 2, "wan_seq_pad": "lane"}, + ) + pad_7x = model_7x._get_token_padding(75600) + self.assertIsNotNone(pad_7x) + self.assertEqual(pad_7x.real_len, 37800) + self.assertEqual(pad_7x.padded_len, 37888) + self.assertEqual(pad_7x.num_segments, 2) + self.assertEqual(pad_7x.total_len, 75776) + + def test_per_token_t_disables_token_padding(self): + mock_mesh = mock.MagicMock() + mock_mesh.shape = {"context": 4} + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + model = WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=4, + attention_head_dim=64, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=128, + num_layers=1, + mesh=mock_mesh, + attention="ulysses_custom_fixed_m_per_q_block", + attention_config={"ulysses_shards": 4, "wan_seq_pad": "lane"}, + ) + # With per_token_t=True (TI2V), token padding must gracefully return None without raising NotImplementedError + pad = model._get_token_padding(75600, per_token_t=True) + self.assertIsNone(pad) + + def test_unsupported_attention_kernel_disables_token_padding(self): + mock_mesh = mock.MagicMock() + mock_mesh.shape = {"context": 4} + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + model = WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=4, + attention_head_dim=64, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=128, + num_layers=1, + mesh=mock_mesh, + attention="dot_product", + attention_config={"wan_seq_pad": "lane"}, + ) + pad = model._get_token_padding(75600) + self.assertIsNone(pad) + + def test_rejects_unknown_mode(self): + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + model = WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=2, + attention_head_dim=64, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=128, + num_layers=1, + attention="dot_product", + ) + # An invalid mode is rejected when the model is built ... + with self.assertRaises(ValueError): + WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=2, + attention_head_dim=64, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=128, + num_layers=1, + attention="dot_product", + attention_config={"wan_seq_pad": "invalid"}, + ) + # ... and by the padding planner itself. + model.seq_pad_mode = "invalid" + with self.assertRaises(ValueError): + model._get_token_padding(75600) + + def test_svg_attention_disables_token_padding(self): + mock_mesh = mock.MagicMock() + mock_mesh.shape = {"context": 4} + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + model = WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=4, + attention_head_dim=64, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=128, + num_layers=1, + mesh=mock_mesh, + attention="ulysses_custom_fixed_m_per_q_block", + attention_config={"ulysses_shards": 4, "wan_seq_pad": "lane", "use_svg_attention": True}, + ) + pad = model._get_token_padding(75600) + self.assertIsNone(pad) + + def test_short_sequence_disables_token_padding(self): + mock_mesh = mock.MagicMock() + mock_mesh.shape = {"context": 4} + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + model = WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=4, + attention_head_dim=64, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=128, + num_layers=1, + mesh=mock_mesh, + flash_min_seq_length=8192, + attention="ulysses_custom_fixed_m_per_q_block", + attention_config={"ulysses_shards": 4, "wan_seq_pad": "lane"}, + ) + # 6000 > default 4096, so this verifies the constructor's flash_min_seq_length=8192 is respected. + pad = model._get_token_padding(6000) + self.assertIsNone(pad) + + def test_memory_efficient_attention_disables_token_padding(self): + mock_mesh = mock.MagicMock() + mock_mesh.shape = {"context": 4} + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + model = WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=4, + attention_head_dim=64, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=128, + num_layers=1, + mesh=mock_mesh, + attention="ulysses_custom_fixed_m_per_q_block", + attention_config={"ulysses_shards": 4, "wan_seq_pad": "lane", "use_memory_efficient_attention": True}, + ) + pad = model._get_token_padding(75600) + self.assertIsNone(pad) + + def test_non_ring_ulysses_uses_context_shards_when_unset_and_rejects_mismatch(self): + mock_mesh = mock.MagicMock() + mock_mesh.shape = {"context": 4} + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + # Default ulysses_shards=-1 on non-ring ulysses_custom* uses full context=4 + model_default_u = WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=4, + attention_head_dim=64, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=128, + num_layers=1, + mesh=mock_mesh, + attention="ulysses_custom_fixed_m_per_q_block", + attention_config={"wan_seq_pad": "lane"}, + ) + pad = model_default_u._get_token_padding(75600) + self.assertIsNotNone(pad) + self.assertEqual(pad.padded_len, 75776) + + # Explicit mismatched ulysses_shards=2 on non-ring ulysses_custom* returns None + model_mismatch_u = WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=4, + attention_head_dim=64, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=128, + num_layers=1, + mesh=mock_mesh, + attention="ulysses_custom_fixed_m_per_q_block", + attention_config={"ulysses_shards": 2, "wan_seq_pad": "lane"}, + ) + with mock.patch.object(transformer_wan, "_warn_once") as warn: + self.assertIsNone(model_mismatch_u._get_token_padding(75600)) + warn.assert_called_once() + self.assertEqual(warn.call_args.args[0], "wan_seq_pad_ulysses_shards_mismatch") + self.assertIn("ulysses_shards=2", warn.call_args.args[1]) + + # Ring attention with unset (-1) or non-dividing ulysses_shards (3 on context=4) returns None + model_ring_invalid = WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=4, + attention_head_dim=64, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=128, + num_layers=1, + mesh=mock_mesh, + attention="ulysses_ring_custom_fixed_m", + attention_config={"ulysses_shards": 3, "wan_seq_pad": "lane"}, + ) + with mock.patch.object(transformer_wan, "_warn_once") as warn: + self.assertIsNone(model_ring_invalid._get_token_padding(75600)) + warn.assert_called_once() + self.assertEqual(warn.call_args.args[0], "wan_seq_pad_ring_ulysses_shards") + self.assertIn("ulysses_shards=3", warn.call_args.args[1]) + + def test_indivisible_seq_len_warns_and_aligned_seq_len_is_silent(self): + mock_mesh = mock.MagicMock() + mock_mesh.shape = {"context": 4} + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + model = WanModel( + rngs=nnx.Rngs(0), + num_attention_heads=4, + attention_head_dim=64, + in_channels=4, + out_channels=4, + text_dim=32, + freq_dim=32, + ffn_dim=128, + num_layers=1, + mesh=mock_mesh, + attention="ulysses_custom_fixed_m_per_q_block", + attention_config={"wan_seq_pad": "lane"}, + ) + with mock.patch.object(transformer_wan, "_warn_once") as warn: + self.assertIsNone(model._get_token_padding(75601)) # 75601 % 4 context shards != 0 + warn.assert_called_once() + self.assertEqual(warn.call_args.args[0], "wan_seq_pad_indivisible_seq") + self.assertIn("seq_len=75601", warn.call_args.args[1]) + + # 65536 tokens over 4 shards is 16384 per shard, already a multiple of 128: a legitimate no-op. + with mock.patch.object(transformer_wan, "_warn_once") as warn: + self.assertIsNone(model._get_token_padding(65536)) + warn.assert_not_called() + + def test_unsupported_kernel_with_active_token_padding_raises(self): + from maxdiffusion.models import attention_flax + + pad = TokenPadding(real_len=128, padded_len=256, num_segments=1) + q = jnp.zeros((1, 256, 2 * 128), dtype=jnp.bfloat16) + with attention_flax.self_attention_token_padding(pad): + with self.assertRaisesRegex(NotImplementedError, "self_attention_token_padding"): + attention_flax._apply_attention_dot( + q, + q, + q, + jnp.bfloat16, + heads=2, + dim_head=128, + scale=1.0, + split_head_dim=True, + float32_qk_product=False, + use_memory_efficient_attention=False, + ) + + _BS = {"block_q": 128, "block_kv": 128, "block_kv_compute": 128, "block_kv_compute_in": 128} + + @staticmethod + def _masking_rules_and_names(): + from maxdiffusion.models import attention_flax + + rules = ( + (attention_flax.BATCH, "data"), + (attention_flax.SELF_ATTN_HEAD, None), + (attention_flax.SELF_ATTN_Q_LENGTH, "context"), + (attention_flax.SELF_ATTN_KV_LENGTH, "context"), + (attention_flax.D_KV, None), + (attention_flax.LENGTH, "context"), + (attention_flax.HEAD, None), + ) + axis_names_q = ( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_Q_LENGTH, + attention_flax.D_KV, + ) + axis_names_kv = ( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_KV_LENGTH, + attention_flax.D_KV, + ) + return rules, axis_names_q, axis_names_kv + + @staticmethod + def _with_sentinel_pad(x, num_segments, real_len, padded_len, value=500.0): + b, _, f = x.shape + x_seg = x.reshape(b, num_segments, real_len, f) + sent = jnp.full((b, num_segments, padded_len - real_len, f), value, dtype=x.dtype) + return jnp.concatenate([x_seg, sent], axis=2).reshape(b, num_segments * padded_len, f) + + @_SKIP_IN_GITHUB_ACTIONS + def test_ulysses_custom_masks_tail_pad_tokens(self): + """Pure Ulysses (U=2) with sentinel-valued tail pad tokens matches the unpadded run. + + Covers the production pad path: sublane-unaligned real_len (300 % 8 != 0, + padded K/V) and aligned real_len (304, `kv_pad_size=1`, i.e. K/V passed + unpadded and the kernel's ragged tail read), each with the flat and the + shard-major output exchange, with sentinels in Q, K and V pad rows. + """ + from maxdiffusion.models import attention_flax + + if len(jax.devices()) < 2: + self.skipTest("needs >= 2 devices for Ulysses attention") + rules, axis_names_q, axis_names_kv = self._masking_rules_and_names() + b, heads, d, padded_len = 1, 4, 128, 512 + mesh_2 = Mesh(np.array(jax.devices()[:2]).reshape(1, 1, 2, 1), ("data", "fsdp", "context", "tensor")) + + for real_len in (300, 304): + q_real = jax.random.normal(jax.random.PRNGKey(100), (b, real_len, heads * d), jnp.bfloat16) + k_real = jax.random.normal(jax.random.PRNGKey(101), (b, real_len, heads * d), jnp.bfloat16) * 0.08 + v_real = jax.random.normal(jax.random.PRNGKey(102), (b, real_len, heads * d), jnp.bfloat16) + pad_u = TokenPadding(real_len=real_len, padded_len=padded_len, num_segments=1) + q_pad, k_pad, v_pad = (self._with_sentinel_pad(x, 1, real_len, padded_len) for x in (q_real, k_real, v_real)) + + for out_a2a in ("flat", "shard_major"): + for use_fixed_m, per_q_block in ((True, True), (True, False), (False, False)): + with self.subTest(real_len=real_len, out_a2a=out_a2a, use_fixed_m=use_fixed_m, per_q_block=per_q_block): + + def run(q, k, v, out_a2a=out_a2a, use_fixed_m=use_fixed_m, per_q_block=per_q_block): + return attention_flax._ulysses_attention( + q, + k, + v, + heads=heads, + mesh=mesh_2, + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + flash_block_sizes=self._BS, + use_custom_kernel=True, + use_fixed_m=use_fixed_m, + per_q_block=per_q_block, + use_k_centering=True, + wan_ulysses_out_a2a=out_a2a, + ) + + with mesh_2, nn_partitioning.axis_rules(rules): + out_unpadded = run(q_real, k_real, v_real) + with attention_flax.self_attention_token_padding(pad_u): + out_padded = run(q_pad, k_pad, v_pad) + np.testing.assert_allclose( + np.asarray(WanModel._unpad_tokens(out_padded, pad_u), np.float32), + np.asarray(out_unpadded, np.float32), + atol=2e-2, + rtol=2e-2, + ) + + @_SKIP_IN_GITHUB_ACTIONS + def test_ulysses_fixed_m_metadata_ignores_q_pad_rows(self): + """Sentinel Q pad rows must not push pure Ulysses off the uniform fixed-m path.""" + from maxdiffusion.models import attention_flax + + b, heads, d, real_len, padded_len, bq = 1, 2, 128, 300, 512, 128 + q = jax.random.normal(jax.random.PRNGKey(1), (b, heads, real_len, d), jnp.float32) * 0.1 + k = jax.random.normal(jax.random.PRNGKey(2), (b, heads, real_len, d), jnp.float32) * 0.1 + v = jax.random.normal(jax.random.PRNGKey(3), (b, heads, real_len, d), jnp.float32) + q_pad = jnp.concatenate([q, jnp.full((b, heads, padded_len - real_len, d), 500.0)], axis=2).astype(jnp.bfloat16) + k, v = k.astype(jnp.bfloat16), v.astype(jnp.bfloat16) + for per_q_block in (True, False): + with self.subTest(per_q_block=per_q_block): + _, all_fixed_unmasked = attention_flax._compute_fixed_m_metadata( + q_pad, k, block_q=bq, per_q_block=per_q_block, value=v + ) + mk, all_fixed = attention_flax._compute_fixed_m_metadata( + q_pad, k, block_q=bq, per_q_block=per_q_block, value=v, q_valid_len=real_len + ) + self.assertFalse(bool(all_fixed_unmasked)) # the sentinels alone would force the online path + self.assertTrue(bool(all_fixed)) + self.assertTrue(bool(jnp.all(mk[:, 1] > 0.5))) + # 5-D shard-major layout: global position = shard * shard_len + i. + q5 = q_pad.reshape(b, heads, 2, padded_len // 2, d).transpose(0, 2, 1, 3, 4) + _, all_fixed5 = attention_flax._compute_fixed_m_metadata(q5, k, block_q=bq, value=v, q_valid_len=real_len) + self.assertTrue(bool(all_fixed5)) + + @_SKIP_IN_GITHUB_ACTIONS + def test_ulysses_ring_custom_masks_tail_pad_tokens(self): + """2D Ulysses x Ring (U=2, R=2, two pad segments) with sentinel pad tokens matches the unpadded run. + + Runs with K-centering on and off (off is the ring default and the tpu7x + production setting). + """ + from maxdiffusion.models import attention_flax + + if len(jax.devices()) < 4: + self.skipTest("needs >= 4 devices for U=2 x R=2 (XLA_FLAGS=--xla_force_host_platform_device_count=4 on CPU)") + rules, axis_names_q, axis_names_kv = self._masking_rules_and_names() + b, heads, d, real_len, padded_len = 1, 4, 128, 300, 512 + mesh_4 = Mesh(np.array(jax.devices()[:4]).reshape(1, 1, 4, 1), ("data", "fsdp", "context", "tensor")) + total_real = 2 * real_len + q_real = jax.random.normal(jax.random.PRNGKey(200), (b, total_real, heads * d), jnp.bfloat16) + k_real = jax.random.normal(jax.random.PRNGKey(201), (b, total_real, heads * d), jnp.bfloat16) * 0.08 + v_real = jax.random.normal(jax.random.PRNGKey(202), (b, total_real, heads * d), jnp.bfloat16) + pad_r = TokenPadding(real_len=real_len, padded_len=padded_len, num_segments=2) + q_pad, k_pad, v_pad = (self._with_sentinel_pad(x, 2, real_len, padded_len) for x in (q_real, k_real, v_real)) + + for use_k_centering in (True, False): + for use_fixed_m, per_q_block in ((True, True), (True, False), (False, False)): + with self.subTest(use_k_centering=use_k_centering, use_fixed_m=use_fixed_m, per_q_block=per_q_block): + + def run(q, k, v, use_k_centering=use_k_centering, use_fixed_m=use_fixed_m, per_q_block=per_q_block): + return attention_flax._ulysses_ring_custom_attention( + q, + k, + v, + heads=heads, + mesh=mesh_4, + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + flash_block_sizes=self._BS, + ulysses_shards=2, + use_fixed_m=use_fixed_m, + per_q_block=per_q_block, + use_k_centering=use_k_centering, + ) + + with mesh_4, nn_partitioning.axis_rules(rules): + out_unpadded = run(q_real, k_real, v_real) + with attention_flax.self_attention_token_padding(pad_r): + out_padded = run(q_pad, k_pad, v_pad) + np.testing.assert_allclose( + np.asarray(WanModel._unpad_tokens(out_padded, pad_r), np.float32), + np.asarray(out_unpadded, np.float32), + atol=2e-2, + rtol=2e-2, + ) + + @_SKIP_IN_GITHUB_ACTIONS + def test_ulysses_ring_custom_aligned_real_len_shard_major_masks_tail_pad_tokens(self): + """The tpu7x production pad path, scaled down: U=2 x R=2, aligned real_len, shard-major exchange. + + real_len=304 is a multiple of 8, so K/V go to the kernel unpadded + (`kv_pad_size=1`) and the kernel reads the ragged tail; each 512-token + segment splits into 256-row ring shards, a multiple of block_q, so + `wan_ulysses_out_a2a="chunked"` takes the shard-major layout. K-centering + is off, as on tpu7x. Sentinels fill the Q, K and V pad rows. + """ + from maxdiffusion.models import attention_flax + + if len(jax.devices()) < 4: + self.skipTest("needs >= 4 devices for U=2 x R=2 (XLA_FLAGS=--xla_force_host_platform_device_count=4 on CPU)") + rules, axis_names_q, axis_names_kv = self._masking_rules_and_names() + b, heads, d, real_len, padded_len = 1, 4, 128, 304, 512 + mesh_4 = Mesh(np.array(jax.devices()[:4]).reshape(1, 1, 4, 1), ("data", "fsdp", "context", "tensor")) + total_real = 2 * real_len + q_real = jax.random.normal(jax.random.PRNGKey(300), (b, total_real, heads * d), jnp.bfloat16) + k_real = jax.random.normal(jax.random.PRNGKey(301), (b, total_real, heads * d), jnp.bfloat16) * 0.08 + v_real = jax.random.normal(jax.random.PRNGKey(302), (b, total_real, heads * d), jnp.bfloat16) + pad_r = TokenPadding(real_len=real_len, padded_len=padded_len, num_segments=2) + q_pad, k_pad, v_pad = (self._with_sentinel_pad(x, 2, real_len, padded_len) for x in (q_real, k_real, v_real)) + + for use_fixed_m, per_q_block in ((True, False), (True, True)): + with self.subTest(use_fixed_m=use_fixed_m, per_q_block=per_q_block): + + def run(q, k, v, out_a2a, use_fixed_m=use_fixed_m, per_q_block=per_q_block): + return attention_flax._ulysses_ring_custom_attention( + q, + k, + v, + heads=heads, + mesh=mesh_4, + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + flash_block_sizes=self._BS, + ulysses_shards=2, + use_fixed_m=use_fixed_m, + per_q_block=per_q_block, + use_k_centering=False, + wan_ulysses_out_a2a=out_a2a, + ) + + ring_mod = attention_flax.tokamax_ring_attention_kernel + with mesh_4, nn_partitioning.axis_rules(rules): + out_unpadded = run(q_real, k_real, v_real, "flat") + with attention_flax.self_attention_token_padding(pad_r): + with mock.patch.object(ring_mod, "make_custom_ring_attention", wraps=ring_mod.make_custom_ring_attention) as spy: + out_padded = run(q_pad, k_pad, v_pad, "chunked") + # The shard-major layout was taken (out_num_shards == U), not the flat fallback. + self.assertEqual(spy.call_args.kwargs["out_num_shards"], 2) + np.testing.assert_allclose( + np.asarray(WanModel._unpad_tokens(out_padded, pad_r), np.float32), + np.asarray(out_unpadded, np.float32), + atol=2e-2, + rtol=2e-2, + ) + + @_SKIP_IN_GITHUB_ACTIONS + def test_wan_model_end_to_end_lane_padding_matches_off(self): + """End-to-end WanModel forward with wan_seq_pad='lane' matches 'off'.""" + if len(jax.devices()) < 2: + self.skipTest("needs >= 2 devices") + mesh_2 = Mesh(np.array(jax.devices()[:2]).reshape(1, 1, 2, 1), ("data", "fsdp", "context", "tensor")) + bs = {"block_q": 128, "block_kv": 128, "block_kv_compute": 128, "block_kv_compute_in": 128} + with mesh_2, nn_partitioning.axis_rules(self.config.logical_axis_rules): + # num_frames=5, height=16, width=20 with patch_size=(1, 2, 2) -> 5 * 8 * 10 = 400 tokens + # Across 2 context shards: 200 tokens/shard -> padded to 256 tokens/shard (512 total). + common_kwargs = { + "patch_size": (1, 2, 2), + "num_attention_heads": 4, + "attention_head_dim": 128, + "in_channels": 4, + "out_channels": 4, + "text_dim": 64, + "freq_dim": 32, + "ffn_dim": 128, + "num_layers": 1, + "scan_layers": False, + "mesh": mesh_2, + "dtype": jnp.bfloat16, + "weights_dtype": jnp.float32, + "flash_min_seq_length": 64, + "flash_block_sizes": bs, + "attention": "ulysses_custom_fixed_m_per_q_block", + } + model_off = WanModel( + rngs=nnx.Rngs(42), + attention_config={"ulysses_shards": 2, "use_base2_exp": True, "wan_seq_pad": "off"}, + **common_kwargs, + ) + model_lane = WanModel( + rngs=nnx.Rngs(42), + attention_config={"ulysses_shards": 2, "use_base2_exp": True, "wan_seq_pad": "lane"}, + **common_kwargs, + ) + pad = model_lane._get_token_padding(400) + self.assertIsNotNone(pad) + self.assertEqual(pad.real_len, 400) + self.assertEqual(pad.padded_len, 512) + + latents = jax.random.normal(jax.random.PRNGKey(300), (1, 4, 5, 16, 20), jnp.bfloat16) + timestep = jnp.array([500], dtype=jnp.int32) + enc = jax.random.normal(jax.random.PRNGKey(301), (1, 16, 64), jnp.bfloat16) + out_off = model_off(latents, timestep, enc, deterministic=True) + out_lane = model_lane(latents, timestep, enc, deterministic=True) + self.assertEqual(out_lane.shape, out_off.shape) + np.testing.assert_allclose(np.asarray(out_lane, np.float32), np.asarray(out_off, np.float32), atol=5e-2, rtol=5e-2) + + +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..11a4f668b 100644 --- a/src/maxdiffusion/utils/wan_block_benchmark.py +++ b/src/maxdiffusion/utils/wan_block_benchmark.py @@ -168,6 +168,25 @@ def from_config(cls, config, mesh, **overrides): # --- BlockBenchmark interface --------------------------------------------------- def tiled_seq_lens(self): s = tiled_seq_len(self._full_seq, self._attention, self._context_shards, self._ulysses_shards) + proxy = type( + "_PadProxy", + (), + { + "seq_pad_mode": getattr(self._config, "wan_seq_pad", None), + "attention": self._attention, + "mesh": self._mesh, + "flash_min_seq_length": getattr(self._config, "flash_min_seq_length", 4096), + "attention_config": { + "ulysses_shards": getattr(self._config, "ulysses_shards", -1), + "flash_min_seq_length": getattr(self._config, "flash_min_seq_length", 4096), + "use_svg_attention": getattr(self._config, "use_svg_attention", False), + "use_memory_efficient_attention": getattr(self._config, "use_memory_efficient_attention", False), + }, + }, + )() + pad = WanModel._get_token_padding(proxy, self._full_seq, per_token_t=False) + if pad is not None: + s = pad.padded_len return (s, s) def vmem_bytes(self): @@ -234,10 +253,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..2a25b7aad 100644 --- a/src/maxdiffusion/wan_runtime_options.py +++ b/src/maxdiffusion/wan_runtime_options.py @@ -53,6 +53,11 @@ "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_patch_embed_mode": ("WAN_PATCH_EMBED_MODE", "conv", "str"), + "wan_ulysses_out_a2a": ("WAN_ULYSSES_OUT_A2A", "flat", "str"), + "wan_seq_pad": ("WAN_SEQ_PAD", "off", "str"), + "wan_cross_attn_cpu_interpret": ("WAN_CROSS_ATTN_CPU_INTERPRET", False, "bool"), } @@ -72,6 +77,14 @@ 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}.") + if name == "wan_patch_embed_mode" and val not in ("conv", "tokens"): + raise ValueError(f"wan_patch_embed_mode must be 'conv' or 'tokens', got {value!r}.") + if name == "wan_ulysses_out_a2a" and val not in ("flat", "chunked", "shard_major"): + raise ValueError(f"wan_ulysses_out_a2a must be 'flat', 'chunked', or 'shard_major', got {value!r}.") + if name == "wan_seq_pad" and val not in ("off", "lane"): + raise ValueError(f"wan_seq_pad must be 'off' or 'lane', got {value!r}.") return val @@ -132,6 +145,11 @@ 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_patch_embed_mode", + "wan_ulysses_out_a2a", + "wan_seq_pad", + "wan_cross_attn_cpu_interpret", )