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", )