diff --git a/end_to_end/tpu/run_wan_fast_inference.sh b/end_to_end/tpu/run_wan_fast_inference.sh index 11388c7d4..873bdc64a 100755 --- a/end_to_end/tpu/run_wan_fast_inference.sh +++ b/end_to_end/tpu/run_wan_fast_inference.sh @@ -68,7 +68,7 @@ # v6e seed 777, 17.92 dB / 0.6556 on v6e seed 12345, and # 17.69 dB / 0.6918 on tpu7x seed 12345 (visually clean). # 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) @@ -146,8 +146,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:-""} +# 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); libtpu fails # backend init with it enabled on any non-v5e platform, including v6e and v7 # ("Continuation fusion for AllGather ... not supported ... other than Viperlite"). @@ -158,13 +160,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_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: 140.1s generate (denoise 137.0s) +# Reference benchmarks (Wan 2.2 T2V-A14B, 720p/81f/40-step, CP=4, DP=2, full +# stack, these launcher defaults, warm AOT cache, DVFS unpinned): +# v6e-8: 126.6s - 127.5s generate (denoise 124.5s - 125.0s) +# tpu7x-8: 107.8s generate (denoise 104.9s; 109.6s / 106.6s with fused RoPE off) +# tpu7x-8 with PIN_DVFS_P_STATE=true: 94.6s generate (denoise 91.7s) with the +# fused RoPE kernel on, and 96.4s generate (denoise 93.4s) with it off. case "$TPU_PROFILE" in v6e) PLATFORM_LIBTPU="$V6E_LIBTPU" @@ -183,13 +186,14 @@ case "$TPU_PROFILE" in DEFAULT_FUSED_ROPE=true DEFAULT_CROSS_ATTN_KERNEL=pallas DEFAULT_PATCH_EMBED_MODE=tokens - DEFAULT_ULYSSES_OUT_A2A=chunked + DEFAULT_ULYSSES_OUT_A2A=shard_major + 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 @@ -202,7 +206,8 @@ case "$TPU_PROFILE" in DEFAULT_FUSED_ROPE=true DEFAULT_CROSS_ATTN_KERNEL=pallas DEFAULT_PATCH_EMBED_MODE=tokens - DEFAULT_ULYSSES_OUT_A2A=chunked + DEFAULT_ULYSSES_OUT_A2A=shard_major + DEFAULT_SEQ_PAD=lane ;; *) echo "== warning: unrecognised or non-v6e/v7 accelerator '${ACCEL_TYPE:-unknown}';" \ @@ -224,6 +229,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 @@ -258,6 +264,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 devices remain. Defaults to CP=4 / DP=2 on an 8-device slice. On a @@ -327,6 +334,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 f693f8677..6d82b7029 100755 --- a/end_to_end/tpu/run_wan_stack_tests.sh +++ b/end_to_end/tpu/run_wan_stack_tests.sh @@ -51,5 +51,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 d99ab6008..75c30cbed 100644 --- a/src/maxdiffusion/configs/base_wan_14b.yml +++ b/src/maxdiffusion/configs/base_wan_14b.yml @@ -147,6 +147,12 @@ wan_patch_embed_mode: "conv" # when S_local % block_q == 0); "shard_major" strictly requires the Pallas # 4D-q / out_num_shards path and raises if any kernel precondition fails. 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 4dbe08a8a..2e8407721 100644 --- a/src/maxdiffusion/configs/base_wan_1_3b.yml +++ b/src/maxdiffusion/configs/base_wan_1_3b.yml @@ -144,6 +144,12 @@ wan_patch_embed_mode: "conv" # when S_local % block_q == 0); "shard_major" strictly requires the Pallas # 4D-q / out_num_shards path and raises if any kernel precondition fails. 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 d6bb8d091..90f30e7a3 100644 --- a/src/maxdiffusion/configs/base_wan_27b.yml +++ b/src/maxdiffusion/configs/base_wan_27b.yml @@ -181,6 +181,12 @@ wan_patch_embed_mode: "conv" # when S_local % block_q == 0); "shard_major" strictly requires the Pallas # 4D-q / out_num_shards path and raises if any kernel precondition fails. 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 67412ab95..ae785ac07 100644 --- a/src/maxdiffusion/configs/base_wan_animate.yml +++ b/src/maxdiffusion/configs/base_wan_animate.yml @@ -126,6 +126,12 @@ wan_patch_embed_mode: "conv" # when S_local % block_q == 0); "shard_major" strictly requires the Pallas # 4D-q / out_num_shards path and raises if any kernel precondition fails. 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 83c2aa15d..bbf5a0616 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_14b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_14b.yml @@ -128,6 +128,12 @@ wan_patch_embed_mode: "conv" # when S_local % block_q == 0); "shard_major" strictly requires the Pallas # 4D-q / out_num_shards path and raises if any kernel precondition fails. 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 8b1925698..a2c25ddd1 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_27b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_27b.yml @@ -128,6 +128,12 @@ wan_patch_embed_mode: "conv" # when S_local % block_q == 0); "shard_major" strictly requires the Pallas # 4D-q / out_num_shards path and raises if any kernel precondition fails. 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 a99f214c1..4997c3e59 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -13,6 +13,7 @@ # limitations under the License. import contextlib +import dataclasses import functools import math from typing import Optional, Callable, Tuple, Any, Dict @@ -746,6 +747,163 @@ 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", +}) + + +def plan_token_padding( + seq_len: Optional[int] = None, + *, + mode: Optional[str], + attention: str, + mesh: Optional[jax.sharding.Mesh], + ulysses_shards: int = -1, + flash_min_seq_length: int = 4096, + per_token_t: bool = False, + use_svg_attention: bool = False, + use_memory_efficient_attention: bool = False, +) -> Optional[TokenPadding]: + """Computes the self-attention `TokenPadding` plan for a `seq_len`-token video sequence. + + Pass `seq_len=None` at model construction time to validate static configuration + without running per-call sequence-length checks. + """ + effective_mode = mode if mode is not None else wan_runtime_options.default("wan_seq_pad") + if effective_mode in (None, "off"): + return None + if effective_mode != "lane": + raise ValueError(f"wan_seq_pad must be 'off' or 'lane', got {effective_mode!r}.") + if use_svg_attention: + raise ValueError("wan_seq_pad='lane' is incompatible with use_svg_attention=True.") + if use_memory_efficient_attention: + raise ValueError("wan_seq_pad='lane' is incompatible with use_memory_efficient_attention=True.") + if per_token_t: + raise ValueError("wan_seq_pad='lane' is incompatible with per_token_t=True (TI2V).") + if attention not in TOKEN_PADDING_KERNELS: + raise ValueError( + f"wan_seq_pad='lane' requires an attention kernel in {sorted(TOKEN_PADDING_KERNELS)}, got {attention!r}." + ) + + num_context_shards = 1 + if mesh is not None and hasattr(mesh, "shape") and "context" in mesh.shape: + num_context_shards = mesh.shape["context"] + if attention.startswith("ulysses_ring_"): + if ulysses_shards is None or ulysses_shards <= 0 or num_context_shards % ulysses_shards != 0: + raise ValueError( + f"wan_seq_pad='lane' requires ulysses_shards > 0 dividing the {num_context_shards} context shards for " + f"attention={attention!r}, got ulysses_shards={ulysses_shards}." + ) + num_ring_shards = num_context_shards // ulysses_shards + elif attention.startswith("ulysses_"): + if ulysses_shards is not None and ulysses_shards > 0 and ulysses_shards != num_context_shards: + raise ValueError( + f"wan_seq_pad='lane' requires ulysses_shards to match the {num_context_shards} context shards for " + f"attention={attention!r}, got ulysses_shards={ulysses_shards}." + ) + ulysses_shards = num_context_shards + num_ring_shards = 1 + else: + ulysses_shards = 1 + num_ring_shards = 1 + num_segments = num_ring_shards + + if seq_len is None: + return None + 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 + + 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) + + +def _resolve_token_padding( + padding: Optional[TokenPadding], + seq_len: int, + num_segments: int, + q_seq_len: Optional[int] = None, +) -> Optional[TokenPadding]: + """Validates `padding` against `seq_len`, `q_seq_len`, and `num_segments`.""" + if padding is None: + return None + eff_q_len = seq_len if q_seq_len is None else q_seq_len + if padding.total_len != seq_len or eff_q_len != seq_len: + raise ValueError( + f"TokenPadding has total_len={padding.total_len}, but attention received " + f"q_seq_len={eff_q_len}, kv_seq_len={seq_len}." + ) + 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_token_padding( + padding: Optional[TokenPadding], + q_seq_len: int, + kv_seq_len: int, + kernel_name: str, +) -> None: + """Raises NotImplementedError if self-attention token padding is passed to an unsupported path.""" + del q_seq_len, kv_seq_len + if padding is not None: + raise NotImplementedError( + f"{kernel_name} does not support token_padding (only {sorted(TOKEN_PADDING_KERNELS)} mask tail pad tokens)." + ) + + ULYSSES_OUT_A2A_MODES = ("flat", "chunked", "shard_major") @@ -858,6 +1016,7 @@ def _tpu_flash_attention( spatiotemporal_config: Optional[dict] = None, spatiotemporal_shape: Optional[Tuple[int, int, int]] = None, kv_heads: Optional[int] = None, + token_padding: Optional[TokenPadding] = None, ) -> jax.Array: """TPU Flash Attention""" if kv_heads is not None and kv_heads != heads: @@ -865,8 +1024,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_token_padding(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.") @@ -1130,6 +1290,7 @@ def _ulysses_attention( qk_prescaled: bool = False, transpose_out: bool = False, wan_ulysses_out_a2a: str = "flat", + token_padding: Optional[TokenPadding] = None, ) -> jax.Array: """Ulysses sequence-parallel attention. @@ -1153,6 +1314,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 = _resolve_token_padding(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( @@ -1263,7 +1431,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, :] @@ -1311,7 +1481,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) @@ -1558,6 +1731,7 @@ def _ulysses_ring_attention( ulysses_attention_chunks: int = 1, preserve_asymmetric_block_sizes: bool = False, kv_heads: int | None = None, + token_padding: Optional[TokenPadding] = None, ) -> jax.Array: """2D context-parallel attention using a private Ulysses x ring mesh. @@ -1611,10 +1785,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_token_padding(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] @@ -1788,6 +1962,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. @@ -1813,6 +1988,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. @@ -1820,21 +2000,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. # @@ -1891,6 +2088,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) @@ -1913,6 +2112,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. @@ -1933,6 +2133,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)` @@ -1989,8 +2194,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) @@ -2005,11 +2222,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 @@ -2046,6 +2262,7 @@ def _ulysses_ring_custom_attention( qk_prescaled: bool = False, transpose_out: bool = False, wan_ulysses_out_a2a: str = "flat", + token_padding: Optional[TokenPadding] = None, ) -> jax.Array: """2D USP attention (Ulysses + Ring) using custom splash kernel with exact Fixed-m support.""" if kv_heads is None: @@ -2088,6 +2305,11 @@ def _ulysses_ring_custom_attention( if kv_heads % num_ulysses_shards != 0: raise ValueError(f"Ulysses+Ring requires KV heads divisible by U={num_ulysses_shards}, got kv_heads={kv_heads}.") _validate_ring_sequence_divisibility(orig_q_seq_len, orig_kv_seq_len, num_ring_shards, num_context_shards) + token_padding = _resolve_token_padding( + 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 token_padding.") ( bq, bkv, @@ -2153,6 +2375,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 @@ -2207,7 +2430,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: @@ -2220,9 +2447,11 @@ 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 + # `_pad_data_for_flash`'s 3rd return is the unpadded input length (`raw_key.shape[2]`); + # use `actual_kv_seq_len` so `token_padding.real_len` masks the tail pad keys. + ring_kv_seq_len = actual_kv_seq_len k_mean = None if use_fixed_m and num_ring_shards == 1 and use_k_centering: @@ -2242,7 +2471,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) @@ -2433,6 +2664,7 @@ def _apply_attention_dot( qk_prescaled: bool = False, k_prescaled: bool = False, use_base2_exp: bool = False, + token_padding: Optional[TokenPadding] = None, ): """Apply Attention.""" effective_kv_heads = kv_heads if kv_heads is not None else heads @@ -2458,6 +2690,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_token_padding(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) @@ -2466,6 +2699,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_token_padding(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] @@ -2565,7 +2799,15 @@ def _to_bshd(x: Array, n_heads: int) -> Array: return hidden_states -def _cudnn_flash_attention(query: Array, key: Array, value: Array, heads: int, mesh: Mesh, dpa_layer: Callable) -> Array: +def _cudnn_flash_attention( + query: Array, + key: Array, + value: Array, + heads: int, + mesh: Mesh, + dpa_layer: Callable, + token_padding: Optional[TokenPadding] = None, +) -> Array: """CUDNN Flash Attention with Transformer Engine. 1. Stable API, supports GQA 2. Supports head_dim till 128; head_dim=256 support will be added soon @@ -2575,6 +2817,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_token_padding(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)) @@ -2614,6 +2857,7 @@ def dot_product_kernel(q, k, v, context): qk_prescaled=context.get("qk_prescaled", False), k_prescaled=context.get("k_prescaled", False), use_base2_exp=context.get("use_base2_exp", False), + token_padding=context.get("token_padding"), ) @@ -2645,6 +2889,7 @@ def ulysses_custom_kernel(q, k, v, context): qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", + token_padding=context.get("token_padding"), ) @@ -2672,6 +2917,7 @@ def ulysses_ring_custom_kernel(q, k, v, context): qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", + token_padding=context.get("token_padding"), ) @@ -2703,6 +2949,7 @@ def ulysses_ring_custom_fixed_m_kernel(q, k, v, context): qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", + token_padding=context.get("token_padding"), ) @@ -2734,6 +2981,7 @@ def ulysses_ring_custom_fixed_m_per_q_block_kernel(q, k, v, context): qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", + token_padding=context.get("token_padding"), ) @@ -2765,6 +3013,7 @@ def ulysses_ring_custom_bidir_kernel(q, k, v, context): qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", + token_padding=context.get("token_padding"), ) @@ -2797,6 +3046,7 @@ def ulysses_custom_fixed_m_kernel(q, k, v, context): qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", + token_padding=context.get("token_padding"), ) @@ -2829,6 +3079,7 @@ def ulysses_custom_fixed_m_per_q_block_kernel(q, k, v, context): qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", + token_padding=context.get("token_padding"), ) @@ -2853,6 +3104,7 @@ def ulysses_kernel(q, k, v, context): ulysses_shards=context.get("ulysses_shards", -1), kernel_name="ulysses", wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", + token_padding=context.get("token_padding"), ) @@ -2877,6 +3129,7 @@ def ulysses_ring_kernel(q, k, v, context): ulysses_attention_chunks=context["ulysses_attention_chunks"], preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False), kv_heads=context.get("kv_heads", None), + token_padding=context.get("token_padding"), ) @@ -2901,6 +3154,7 @@ def flash_kernel(q, k, v, context): is_causal=context.get("is_causal", False), preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False), kv_heads=context.get("kv_heads", None), + token_padding=context.get("token_padding"), ) @@ -2927,6 +3181,7 @@ def tokamax_flash_kernel(q, k, v, context): spatiotemporal_config=context.get("spatiotemporal_config"), spatiotemporal_shape=context.get("spatiotemporal_shape"), kv_heads=context.get("kv_heads", None), + token_padding=context.get("token_padding"), ) @@ -2951,6 +3206,7 @@ def tokamax_ring_kernel(q, k, v, context): is_causal=context.get("is_causal", False), preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False), kv_heads=context.get("kv_heads", None), + token_padding=context.get("token_padding"), ) @@ -2974,6 +3230,7 @@ def tokamax_ring_custom_kernel(q, k, v, context): use_experimental_scheduler=context.get("use_experimental_scheduler", False), preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False), kv_heads=context.get("kv_heads", None), + token_padding=context.get("token_padding"), ) @@ -2981,7 +3238,15 @@ def tokamax_ring_custom_kernel(q, k, v, context): def cudnn_flash_te_kernel(q, k, v, context): if context.get("k_prescaled", False): k = k / jnp.asarray(context["scale"], dtype=k.dtype) - return _cudnn_flash_attention(q, k, v, context["heads"], context["mesh"], context["dpa_layer"]) + return _cudnn_flash_attention( + q, + k, + v, + context["heads"], + context["mesh"], + context["dpa_layer"], + token_padding=context.get("token_padding"), + ) _QK_PRESCALE_SUPPORTED_KERNELS = { @@ -3066,6 +3331,7 @@ def _apply_attention( k_prescaled: bool = False, transpose_out: bool = False, wan_ulysses_out_a2a: str = "flat", + token_padding: Optional[TokenPadding] = None, ): """Routes to different attention kernels using a module-level registry.""" @@ -3151,6 +3417,7 @@ def _apply_attention( "k_prescaled": k_prescaled, "transpose_out": transpose_out, "wan_ulysses_out_a2a": wan_ulysses_out_a2a, + "token_padding": token_padding, } if spatiotemporal_config and spatiotemporal_config.get("use_svg_attention"): @@ -3189,6 +3456,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_token_padding(context.get("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] @@ -3569,6 +3837,7 @@ def apply_attention( sparse_config_override: Optional[dict] = None, qk_prescaled: bool = False, k_prescaled: bool = False, + token_padding: Optional[TokenPadding] = None, ): return _apply_attention( query=query, @@ -3604,6 +3873,7 @@ def apply_attention( k_prescaled=k_prescaled, transpose_out=getattr(self, "transpose_out", False), wan_ulysses_out_a2a=getattr(self, "wan_ulysses_out_a2a", "flat"), + token_padding=token_padding, ) @@ -3667,6 +3937,7 @@ def apply_attention( sparse_config_override: Optional[dict] = None, qk_prescaled: bool = False, k_prescaled: bool = False, + token_padding: Optional[TokenPadding] = None, ): return _apply_attention( query=query, @@ -3701,6 +3972,7 @@ def apply_attention( k_prescaled=k_prescaled, transpose_out=getattr(self, "transpose_out", False), wan_ulysses_out_a2a=getattr(self, "wan_ulysses_out_a2a", "flat"), + token_padding=token_padding, ) @@ -4469,6 +4741,7 @@ def __call__( svg_layer_index: Optional[int | jax.Array] = None, svg_timestep: Optional[int | float | jax.Array] = None, svg_step_index: Optional[int | jax.Array] = None, + token_padding: Optional[TokenPadding] = None, ) -> jax.Array: same_kv_source = encoder_hidden_states is None or encoder_hidden_states is hidden_states hidden_states = nn.with_logical_constraint(hidden_states, (BATCH, LENGTH, HEAD)) @@ -4578,6 +4851,7 @@ def run_attention(**svg_kwargs): attention_mask=encoder_attention_mask, qk_prescaled=qk_prescaled, k_prescaled=k_prescaled, + token_padding=token_padding if is_self_attention else None, **svg_kwargs, ) diff --git a/src/maxdiffusion/models/wan/transformers/transformer_wan.py b/src/maxdiffusion/models/wan/transformers/transformer_wan.py index 53e14a2ca..78ac3ef7a 100644 --- a/src/maxdiffusion/models/wan/transformers/transformer_wan.py +++ b/src/maxdiffusion/models/wan/transformers/transformer_wan.py @@ -35,7 +35,11 @@ NNXPixArtAlphaTextProjection, ) from ...normalization_flax import FP32LayerNorm -from ...attention_flax import FlaxWanAttention +from ...attention_flax import ( + FlaxWanAttention, + TokenPadding, + plan_token_padding, +) from ...gradient_checkpoint import GradientCheckpointType BlockSizes = common_types.BlockSizes @@ -513,6 +517,7 @@ def __call__( svg_layer_index: Optional[int | jax.Array] = None, svg_timestep: Optional[int | float | jax.Array] = None, svg_step_index: Optional[int | jax.Array] = None, + token_padding: Optional[TokenPadding] = None, ): with self.conditional_named_scope("transformer_block"): # Support both global [B, 6, dim] and per-token [B, seq_len, 6, dim] temb. @@ -562,6 +567,7 @@ def __call__( svg_layer_index=svg_layer_index, svg_timestep=svg_timestep, svg_step_index=svg_step_index, + token_padding=token_padding, ) with self.conditional_named_scope("self_attn_residual"): hidden_states = (hidden_states.astype(jnp.float32) + attn_output * gate_msa).astype(hidden_states.dtype) @@ -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,20 @@ 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") + if self.seq_pad_mode == "lane": + # Validate static configuration preconditions at build time rather than + # silently disabling token padding on every forward pass. + plan_token_padding( + None, + mode=self.seq_pad_mode, + attention=self.attention, + mesh=self.mesh, + ulysses_shards=self.attention_config.get("ulysses_shards", -1), + flash_min_seq_length=self.flash_min_seq_length, + use_svg_attention=bool(self.attention_config.get("use_svg_attention", False)), + use_memory_efficient_attention=bool(self.attention_config.get("use_memory_efficient_attention", False)), + ) # 1. Patch & position embedding self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) @@ -841,6 +865,48 @@ 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]: + return plan_token_padding( + seq_len, + mode=self.seq_pad_mode, + attention=self.attention, + mesh=self.mesh, + ulysses_shards=self.attention_config.get("ulysses_shards", -1), + flash_min_seq_length=self.attention_config.get("flash_min_seq_length", self.flash_min_seq_length), + per_token_t=per_token_t, + use_svg_attention=bool(self.attention_config.get("use_svg_attention", False)), + use_memory_efficient_attention=bool(self.attention_config.get("use_memory_efficient_attention", False)), + ) + + @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, @@ -940,6 +1006,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: @@ -964,6 +1035,7 @@ def scan_fn(carry, block_input): svg_layer_index=layer_index, svg_timestep=timestep, svg_step_index=svg_step_index, + token_padding=token_padding, ) new_carry = (hidden_states, rngs_carry) return new_carry, None @@ -1013,6 +1085,7 @@ def layer_forward(hidden_states, l_kv): svg_layer_index=i, svg_timestep=timestep, svg_step_index=svg_step_index, + token_padding=token_padding, ) rematted_layer_forward = self.gradient_checkpoint.apply( @@ -1029,6 +1102,8 @@ 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: + cached_residual = self._pad_tokens(cached_residual, token_padding) hidden_states = hidden_states + cached_residual else: hidden_states = _run_all_blocks(hidden_states) @@ -1048,6 +1123,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/fused_rmsnorm_rope_pallas_test.py b/src/maxdiffusion/tests/fused_rmsnorm_rope_pallas_test.py index fee4facc0..751da6052 100644 --- a/src/maxdiffusion/tests/fused_rmsnorm_rope_pallas_test.py +++ b/src/maxdiffusion/tests/fused_rmsnorm_rope_pallas_test.py @@ -942,7 +942,7 @@ def test_cudnn_flash_te_unscales_prescaled_k(self): v = jnp.ones((1, 16, 2, 64), dtype=jnp.float32) captured = {} - def fake_cudnn(q_in, k_in, v_in, heads, mesh, dpa_layer): + def fake_cudnn(q_in, k_in, v_in, heads, mesh, dpa_layer, **_kwargs): captured["k"] = np.asarray(k_in) return q_in diff --git a/src/maxdiffusion/tests/wan_runtime_options_test.py b/src/maxdiffusion/tests/wan_runtime_options_test.py index 925651261..97c6391a2 100644 --- a/src/maxdiffusion/tests/wan_runtime_options_test.py +++ b/src/maxdiffusion/tests/wan_runtime_options_test.py @@ -134,9 +134,9 @@ def test_unset_options_resolve_to_defaults_not_env(self): } 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)) @@ -186,6 +186,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..7a1bbd3a9 --- /dev/null +++ b/src/maxdiffusion/tests/wan_seq_pad_test.py @@ -0,0 +1,723 @@ +""" +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 import attention_flax +from maxdiffusion.models.attention_flax import TokenPadding +from maxdiffusion.models.wan.transformers.transformer_wan import WanModel + +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_raises_value_error(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 self.assertRaisesRegex(ValueError, "per_token_t"): + model._get_token_padding(75600, per_token_t=True) + + def test_unsupported_attention_kernel_raises_value_error(self): + mock_mesh = mock.MagicMock() + mock_mesh.shape = {"context": 4} + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + with self.assertRaisesRegex(ValueError, "wan_seq_pad='lane'"): + 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"}, + ) + + 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_raises_value_error(self): + mock_mesh = mock.MagicMock() + mock_mesh.shape = {"context": 4} + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + with self.assertRaisesRegex(ValueError, "use_svg_attention"): + 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}, + ) + + 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_raises_value_error(self): + mock_mesh = mock.MagicMock() + mock_mesh.shape = {"context": 4} + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + with self.assertRaisesRegex(ValueError, "use_memory_efficient_attention"): + 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}, + ) + + 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* raises ValueError + with self.assertRaisesRegex(ValueError, "ulysses_shards=2"): + 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"}, + ) + + # Ring attention with unset (-1) or non-dividing ulysses_shards (3 on context=4) raises ValueError + with self.assertRaisesRegex(ValueError, "ulysses_shards=3"): + 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"}, + ) + + 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(attention_flax, "_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(attention_flax, "_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): + pad = TokenPadding(real_len=128, padded_len=256, num_segments=1) + q = jnp.zeros((1, 256, 2 * 128), dtype=jnp.bfloat16) + with self.assertRaisesRegex(NotImplementedError, "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, + token_padding=pad, + ) + # Even if token_padding.total_len differs from q/kv seq_len, unsupported kernels must still reject it. + q_mismatched = jnp.zeros((1, 128, 2 * 128), dtype=jnp.bfloat16) + with self.assertRaisesRegex(NotImplementedError, "token_padding"): + attention_flax._apply_attention_dot( + q_mismatched, + q_mismatched, + q_mismatched, + jnp.bfloat16, + heads=2, + dim_head=128, + scale=1.0, + split_head_dim=True, + float32_qk_product=False, + use_memory_efficient_attention=False, + token_padding=pad, + ) + + def test_resolve_token_padding_length_mismatch_raises(self): + pad = TokenPadding(real_len=384, padded_len=512, num_segments=1) + self.assertIsNone(attention_flax._resolve_token_padding(None, seq_len=512, num_segments=1)) + self.assertIs(attention_flax._resolve_token_padding(pad, seq_len=512, num_segments=1, q_seq_len=512), pad) + with self.assertRaisesRegex(ValueError, "total_len=512"): + attention_flax._resolve_token_padding(pad, seq_len=384, num_segments=1) + with self.assertRaisesRegex(ValueError, "q_seq_len=384"): + attention_flax._resolve_token_padding(pad, seq_len=512, num_segments=1, q_seq_len=384) + with self.assertRaisesRegex(ValueError, "segment"): + attention_flax._resolve_token_padding(pad, seq_len=512, num_segments=2) + + _BS = {"block_q": 128, "block_kv": 128, "block_kv_compute": 128, "block_kv_compute_in": 128} + + @staticmethod + def _masking_rules_and_names(): + 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) + + 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. + """ + 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 use_fixed_m, per_q_block in ((True, True), (True, False), (False, False)): + + def run(q, k, v, out_a2a, token_padding=None, 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, + token_padding=token_padding, + ) + + with mesh_2, nn_partitioning.axis_rules(rules): + out_unpadded = run(q_real, k_real, v_real, "flat") + for out_a2a in ("flat", "shard_major"): + with self.subTest(real_len=real_len, out_a2a=out_a2a, use_fixed_m=use_fixed_m, per_q_block=per_q_block): + with mesh_2, nn_partitioning.axis_rules(rules): + out_padded = run(q_pad, k_pad, v_pad, out_a2a, token_padding=pad_u) + 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.""" + 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)) + + 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). + """ + 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, + token_padding=None, + 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, + token_padding=token_padding, + ) + + with mesh_4, nn_partitioning.axis_rules(rules): + out_unpadded = run(q_real, k_real, v_real) + out_padded = run(q_pad, k_pad, v_pad, token_padding=pad_r) + 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, + ) + + 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. + """ + 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, + token_padding=None, + 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, + token_padding=token_padding, + ) + + with mesh_4, nn_partitioning.axis_rules(rules): + out_unpadded = run(q_real, k_real, v_real) + 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, token_padding=pad_1) + 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, + ) + + 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. + """ + 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, token_padding=None, 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, + token_padding=token_padding, + ) + + 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 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", token_padding=pad_r) + # 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, + ) + + 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 93d656543..88c633fc3 100644 --- a/src/maxdiffusion/utils/wan_block_benchmark.py +++ b/src/maxdiffusion/utils/wan_block_benchmark.py @@ -61,6 +61,7 @@ from flax.linen import partitioning as nn_partitioning from maxdiffusion import max_logging, max_utils, pyconfig, wan_runtime_options +from maxdiffusion.models.attention_flax import plan_token_padding from maxdiffusion.models.wan.transformers.transformer_wan import WanModel from maxdiffusion.utils.tile_size_grid_search import ( BenchResult, @@ -168,6 +169,19 @@ 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) + pad = plan_token_padding( + self._full_seq, + mode=getattr(self._config, "wan_seq_pad", None), + attention=self._attention, + mesh=self._mesh, + ulysses_shards=getattr(self._config, "ulysses_shards", -1), + flash_min_seq_length=getattr(self._config, "flash_min_seq_length", 4096), + per_token_t=False, + use_svg_attention=bool(getattr(self._config, "use_svg_attention", False)), + use_memory_efficient_attention=bool(getattr(self._config, "use_memory_efficient_attention", 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 b0c5a7e90..94f24c7d5 100644 --- a/src/maxdiffusion/wan_runtime_options.py +++ b/src/maxdiffusion/wan_runtime_options.py @@ -44,6 +44,7 @@ "wan_cross_attn_kernel": ("xla", "str"), "wan_patch_embed_mode": ("conv", "str"), "wan_ulysses_out_a2a": ("flat", "str"), + "wan_seq_pad": ("off", "str"), } @@ -69,6 +70,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 @@ -112,6 +115,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", )