From de77cc3a1ab93320b5d7cb82f1971fdb1b947c5e Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Thu, 15 Oct 2026 05:51:29 +0000 Subject: [PATCH] feat(wan): lane sequence padding for self-attention (wan_seq_pad="lane") Last of four pieces split out of the Wan fast-inference optimizations PR (the cross-attention Pallas kernel, token-sharded patch embedding and shard-major Ulysses A2A are the three PRs below this one). The switch defaults off in the Wan YAMLs and in wan_runtime_options; it is a plain config key plumbed to WanModel through `attention_config_entries(config)`. - Lane sequence padding (wan_seq_pad="lane"): WanModel pads the video tokens at the tail to a per-shard multiple of 128 (75,600 -> 75,776 at 720p/CP=4; on the Ulysses-ring path each ring segment is padded at its own tail) so the Ulysses reshapes are copy-free, pads the rotary embedding to match, and unpads the output and residual. The padding is declared to the attention layers with `self_attention_token_padding` (a ContextVar scoped to the block trace); `TokenPadding` describes it. - The custom ulysses / ulysses_ring_custom kernels mask the pad keys by passing the real length as the ragged KV tail; the fixed-m norm reductions (`_compute_fixed_m_metadata` via `q_valid_len` on both the 1-D Ulysses path and the R == 1 ring path, `_ring_fixed_m_norms_pre_a2a` via `row_ok` for R > 1) exclude the pad Q/K/V rows so pad-token values cannot push a block onto the online-softmax path or move the k_mean. The ring's `ring_kv_seq_len` is now always the real length. - Every configuration that disables it logs why, once per cause (TI2V per-token timesteps, SVG attention, memory-efficient attention, short sequences, unsupported kernels, ulysses_shards that do not divide the context axis, indivisible sequences). - Behaviour change: if padding is active, the paths that cannot mask it (flash and non-custom ulysses via _tpu_flash_attention, ulysses_ring, dot_product, cudnn_flash_te, head_local_svg, ulysses_ring_custom_bidir) raise NotImplementedError via `_reject_active_token_padding` instead of silently attending to pads. - wan_block_benchmark.py applies the same padding plan to the benchmarked sequence length. - run_wan_fast_inference.sh: the v6e and v7 profiles set WAN_SEQ_PAD=lane (overridable); the generic profile keeps off. LIBTPU: v6e now disables --xla_tpu_enable_ici_ag_pipelining / --xla_tpu_enable_ag_backward_pipelining, v7 enables them and adds --xla_tpu_enable_offloading_copy_to_sparsecore=false and BQ=4736. These two flags moved out of COMMON_LIBTPU, so the generic (non-v6e/v7) profile no longer sets them at all. The launcher's reference-benchmark comment carries the stack-tip numbers below. Performance, measured at the stack tip (Wan 2.2 T2V-A14B, 720p/81f/40 steps, warm AOT, DVFS unpinned unless noted; CP=4, DP=2, launcher defaults including the fused RoPE kernel): v6e-8: denoise 125.9s, generate 128.05s tpu7x-8: denoise 105.1s, generate 107.6s tpu7x-8, DVFS pinned (PIN_DVFS_P_STATE=true), fused RoPE kernel off: denoise 93.6s, generate 96.1s (pinned with it on was not measured) Per-feature deltas were not measured separately and are omitted. Tests - wan_seq_pad_test.py (18): padding plan and per-cause disable logging, pad-key masking for ulysses custom (unaligned and aligned real_len, flat and shard-major out A2A, every fixed-m mode) and ulysses_ring_custom (U=2 x R=2 with k-centring on/off; U=4 x R=1 with k-centring on/off, asserting through a spy that the R == 1 metadata call receives q_valid_len=real_len; and the tpu7x production path scaled down: U=2 x R=2, aligned real_len, chunked out A2A, asserting the shard-major layout was taken; all need >= 4 devices), fixed-m metadata ignoring pad query rows, and the NotImplementedError guards. - wan_runtime_options_test.py (12): the new switch's YAML default and enum validation; `wan_seq_pad` is excluded from the FlaxWanAttention attribute check since WanModel consumes it. - CI: the five multi-device masking / end-to-end tests in wan_seq_pad_test skip in GitHub Actions; padding-plan logic and runtime options run in CI. - end_to_end/tpu/run_wan_stack_tests.sh: adds wan_seq_pad_test.py. Verified (final tree): TPU v6e-8 (jax 0.11.2, end_to_end/tpu/run_wan_stack_tests.sh): 356 passed, 1 skipped (off-TPU fallback only), 282 subtests passed. wan_seq_pad_test.py run directly: 18 tests OK (all 5 multi-device tests ran). --- end_to_end/tpu/run_wan_fast_inference.sh | 28 +- end_to_end/tpu/run_wan_stack_tests.sh | 2 + src/maxdiffusion/configs/base_wan_14b.yml | 6 + src/maxdiffusion/configs/base_wan_1_3b.yml | 6 + src/maxdiffusion/configs/base_wan_27b.yml | 6 + src/maxdiffusion/configs/base_wan_animate.yml | 6 + src/maxdiffusion/configs/base_wan_i2v_14b.yml | 6 + src/maxdiffusion/configs/base_wan_i2v_27b.yml | 6 + src/maxdiffusion/models/attention_flax.py | 183 ++++- .../wan/transformers/transformer_wan.py | 147 +++- .../tests/wan_runtime_options_test.py | 7 +- src/maxdiffusion/tests/wan_seq_pad_test.py | 717 ++++++++++++++++++ src/maxdiffusion/utils/wan_block_benchmark.py | 19 + src/maxdiffusion/wan_runtime_options.py | 4 + 14 files changed, 1111 insertions(+), 32 deletions(-) create mode 100644 src/maxdiffusion/tests/wan_seq_pad_test.py diff --git a/end_to_end/tpu/run_wan_fast_inference.sh b/end_to_end/tpu/run_wan_fast_inference.sh index 1bd508d6f..b8114c46b 100755 --- a/end_to_end/tpu/run_wan_fast_inference.sh +++ b/end_to_end/tpu/run_wan_fast_inference.sh @@ -63,7 +63,7 @@ # 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_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) @@ -134,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"). @@ -146,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" @@ -172,12 +175,13 @@ case "$TPU_PROFILE" in 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 @@ -191,6 +195,7 @@ case "$TPU_PROFILE" in 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}';" \ @@ -212,6 +217,7 @@ case "$TPU_PROFILE" in DEFAULT_CROSS_ATTN_KERNEL=xla DEFAULT_PATCH_EMBED_MODE=tokens DEFAULT_ULYSSES_OUT_A2A=flat + DEFAULT_SEQ_PAD=off ;; esac @@ -246,6 +252,7 @@ FUSED_ROPE_ARGS=() 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 @@ -311,6 +318,7 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \ 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 055d747f1..ce35b4e7f 100755 --- a/end_to_end/tpu/run_wan_stack_tests.sh +++ b/end_to_end/tpu/run_wan_stack_tests.sh @@ -50,5 +50,7 @@ TESTS=( "$T/wan_patch_embed_test.py" # Shard-major Ulysses output A2A (feat/ulysses-shard-major-a2a) "$T/ulysses_out_a2a_test.py" + # Lane sequence padding (feat/wan-fast-inference-optimizations) + "$T/wan_seq_pad_test.py" ) PYTHONPATH="src${PYTHONPATH:+:$PYTHONPATH}" exec "${PYTHON:-python3}" -m pytest -q -rs "${TESTS[@]}" "$@" diff --git a/src/maxdiffusion/configs/base_wan_14b.yml b/src/maxdiffusion/configs/base_wan_14b.yml index 68b07a81f..842c54efd 100644 --- a/src/maxdiffusion/configs/base_wan_14b.yml +++ b/src/maxdiffusion/configs/base_wan_14b.yml @@ -148,6 +148,12 @@ wan_patch_embed_mode: "conv" # 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 a1ea8504e..092784d73 100644 --- a/src/maxdiffusion/configs/base_wan_1_3b.yml +++ b/src/maxdiffusion/configs/base_wan_1_3b.yml @@ -145,6 +145,12 @@ wan_patch_embed_mode: "conv" # 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 b2f61ce66..74b64ab5f 100644 --- a/src/maxdiffusion/configs/base_wan_27b.yml +++ b/src/maxdiffusion/configs/base_wan_27b.yml @@ -177,6 +177,12 @@ wan_patch_embed_mode: "conv" # 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 495255eaa..5ad2482d3 100644 --- a/src/maxdiffusion/configs/base_wan_animate.yml +++ b/src/maxdiffusion/configs/base_wan_animate.yml @@ -127,6 +127,12 @@ wan_patch_embed_mode: "conv" # 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 8b7938f66..0b7aa41ef 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_14b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_14b.yml @@ -129,6 +129,12 @@ wan_patch_embed_mode: "conv" # 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 435127f35..16a3b8d25 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_27b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_27b.yml @@ -129,6 +129,12 @@ wan_patch_embed_mode: "conv" # 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/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 0bc8f9292..1b5a24b72 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 @@ -716,6 +718,83 @@ 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") @@ -836,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.") @@ -1124,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( @@ -1225,7 +1312,9 @@ def wrap_ulysses_attention(query, key, value, attention_mask): raw_query = query raw_value = value 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 = orig_kv_seq_len + # 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, :] @@ -1273,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) @@ -1577,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] @@ -1754,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. @@ -1779,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. @@ -1786,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. # @@ -1857,6 +1972,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) @@ -1879,6 +1996,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. @@ -1899,6 +2017,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)` @@ -1955,8 +2078,20 @@ 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 - if per_q_block: + + def _q_norm_sq(): norm_sq = jnp.square(query.astype(jnp.float32)).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 = _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) @@ -1971,11 +2106,10 @@ def _compute_fixed_m_metadata( 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: - norm_sq = jnp.square(query.astype(jnp.float32)).sum(axis=-1) if query.ndim == 5: - qn_max_sq = norm_sq.max(axis=(1, 3)) + qn_max_sq = _q_norm_sq().max(axis=(1, 3)) else: - qn_max_sq = norm_sq.max(axis=-1) # (batch, num_q_heads) + 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 @@ -2058,6 +2192,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, @@ -2123,6 +2261,7 @@ 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 @@ -2169,7 +2308,11 @@ def wrap_ulysses_ring_attention(query, key, value): raw_query = query raw_value = value 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 = orig_kv_seq_len if num_ring_shards == 1 else raw_key.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, :] if use_shard_major: @@ -2182,9 +2325,9 @@ def wrap_ulysses_ring_attention(query, key, value): # 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: @@ -2204,7 +2347,9 @@ 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, :], + # As on the 1-D Ulysses path: 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) @@ -2378,6 +2523,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) @@ -2386,6 +2532,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] @@ -2495,6 +2642,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)) @@ -3096,6 +3244,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] diff --git a/src/maxdiffusion/models/wan/transformers/transformer_wan.py b/src/maxdiffusion/models/wan/transformers/transformer_wan.py index 51132ec1e..a2b3a58db 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, _warn_once +from ...attention_flax import ( + FlaxWanAttention, + TokenPadding, + TOKEN_PADDING_KERNELS, + self_attention_token_padding, + _warn_once, +) from ...gradient_checkpoint import GradientCheckpointType BlockSizes = common_types.BlockSizes @@ -664,6 +670,10 @@ 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): @@ -671,6 +681,7 @@ def _wan_opt(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) @@ -841,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, @@ -943,6 +1074,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: @@ -1032,9 +1168,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 @@ -1051,6 +1190,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/wan_runtime_options_test.py b/src/maxdiffusion/tests/wan_runtime_options_test.py index 9995a6382..f436d897e 100644 --- a/src/maxdiffusion/tests/wan_runtime_options_test.py +++ b/src/maxdiffusion/tests/wan_runtime_options_test.py @@ -129,9 +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() - # wan_patch_embed_mode rides in attention_config too, but is consumed by - # WanModel (covered in wan_patch_embed_test). - for name in sorted(set(opts.ATTENTION_OPTIONS) - {"wan_patch_embed_mode"}): + # 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)) @@ -168,6 +168,7 @@ def test_enum_options_validation(self): "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: 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..c1048bc20 --- /dev/null +++ b/src/maxdiffusion/tests/wan_seq_pad_test.py @@ -0,0 +1,717 @@ +""" +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, + ) + + 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_r1_masks_tail_pad_tokens(self): + """Degenerate ring (U == context shards, R=1, one pad segment) matches the unpadded run. + + At R=1 the ring wrapper computes the fixed-m metadata itself (not via + `_ring_fixed_m_norms_pre_a2a`), so this pins that the pad Q rows are + excluded there too (`q_valid_len`), with K-centering on and off. + """ + from maxdiffusion.models import attention_flax + + if len(jax.devices()) < 4: + self.skipTest("needs >= 4 devices for U=4 x R=1 (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")) + q_real = jax.random.normal(jax.random.PRNGKey(210), (b, real_len, heads * d), jnp.bfloat16) + k_real = jax.random.normal(jax.random.PRNGKey(211), (b, real_len, heads * d), jnp.bfloat16) * 0.08 + v_real = jax.random.normal(jax.random.PRNGKey(212), (b, real_len, heads * d), jnp.bfloat16) + pad_1 = 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 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=4, + 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_1): + with mock.patch.object( + attention_flax, "_compute_fixed_m_metadata", wraps=attention_flax._compute_fixed_m_metadata + ) as spy: + out_padded = run(q_pad, k_pad, v_pad) + if use_fixed_m: + spy.assert_called_once() + self.assertEqual(spy.call_args.kwargs["q_valid_len"], real_len) + else: + spy.assert_not_called() + np.testing.assert_allclose( + np.asarray(WanModel._unpad_tokens(out_padded, pad_1), 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 3e5de7e74..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): diff --git a/src/maxdiffusion/wan_runtime_options.py b/src/maxdiffusion/wan_runtime_options.py index 2191ae1ca..2a25b7aad 100644 --- a/src/maxdiffusion/wan_runtime_options.py +++ b/src/maxdiffusion/wan_runtime_options.py @@ -56,6 +56,7 @@ "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"), } @@ -82,6 +83,8 @@ def _coerce(name: str, value: Any) -> Any: 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 @@ -145,6 +148,7 @@ def resolve_from_config(config: Any, name: str) -> Any: "wan_cross_attn_kernel", "wan_patch_embed_mode", "wan_ulysses_out_a2a", + "wan_seq_pad", "wan_cross_attn_cpu_interpret", )