diff --git a/end_to_end/tpu/run_wan_fast_inference.sh b/end_to_end/tpu/run_wan_fast_inference.sh index 0313f314b..d41c681a0 100755 --- a/end_to_end/tpu/run_wan_fast_inference.sh +++ b/end_to_end/tpu/run_wan_fast_inference.sh @@ -68,6 +68,8 @@ # 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_SEQ_PAD +# override Wan fast-inference self-attention sequence padding ("off" | "lane") # DP / CP / PER_DEVICE_BATCH / SEED # override mesh parallelism, per-device batch, or RNG seed (default 12345) set -euo pipefail @@ -144,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"). @@ -156,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" @@ -179,12 +184,13 @@ case "$TPU_PROFILE" in DEFAULT_VAE_CHUNK=1 DEFAULT_VAE_SPATIAL=8 DEFAULT_FUSED_ROPE=true + 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 @@ -195,6 +201,7 @@ case "$TPU_PROFILE" in DEFAULT_VAE_CHUNK=1 DEFAULT_VAE_SPATIAL=8 DEFAULT_FUSED_ROPE=true + DEFAULT_SEQ_PAD=lane ;; *) echo "== warning: unrecognised or non-v6e/v7 accelerator '${ACCEL_TYPE:-unknown}';" \ @@ -213,6 +220,7 @@ case "$TPU_PROFILE" in DEFAULT_VAE_CHUNK=1 DEFAULT_VAE_SPATIAL=8 DEFAULT_FUSED_ROPE=false + DEFAULT_SEQ_PAD=off ;; esac @@ -244,6 +252,7 @@ USE_FUSED_ROPE_KERNEL=${USE_FUSED_ROPE_KERNEL:-$DEFAULT_FUSED_ROPE} FUSED_ROPE_ARGS=() [ -n "${FUSED_ROPE_BLOCK_S:-}" ] && FUSED_ROPE_ARGS+=("fused_rope_block_s=$FUSED_ROPE_BLOCK_S") [ -n "${FUSED_ROPE_HEAD_BLOCK:-}" ] && FUSED_ROPE_ARGS+=("fused_rope_head_block=$FUSED_ROPE_HEAD_BLOCK") +WAN_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 @@ -310,6 +319,7 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \ use_kv_cache=true use_base2_exp=true use_experimental_scheduler=true \ use_fused_rope_kernel="$USE_FUSED_ROPE_KERNEL" \ ${FUSED_ROPE_ARGS[@]+"${FUSED_ROPE_ARGS[@]}"} \ + wan_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 eb4400dc6..6c894d4e5 100755 --- a/end_to_end/tpu/run_wan_stack_tests.sh +++ b/end_to_end/tpu/run_wan_stack_tests.sh @@ -45,5 +45,7 @@ TESTS=( # Fused RMSNorm+RoPE producer, Wan switches in config (feat/wan-custom-kernels) "$T/fused_rmsnorm_rope_pallas_test.py" "$T/wan_runtime_options_test.py" + # 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 5404f1a22..9d5af4c93 100644 --- a/src/maxdiffusion/configs/base_wan_14b.yml +++ b/src/maxdiffusion/configs/base_wan_14b.yml @@ -132,6 +132,12 @@ wan_cfg_before_unpatchify: True wan_cross_attn_prescale_kv: False # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32". wan_rope_accum: "auto" +# 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 cfb8c9a4d..bc9a8ffcb 100644 --- a/src/maxdiffusion/configs/base_wan_1_3b.yml +++ b/src/maxdiffusion/configs/base_wan_1_3b.yml @@ -129,6 +129,12 @@ wan_cfg_before_unpatchify: True wan_cross_attn_prescale_kv: False # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32". wan_rope_accum: "auto" +# 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 a6872a911..434ffbc09 100644 --- a/src/maxdiffusion/configs/base_wan_27b.yml +++ b/src/maxdiffusion/configs/base_wan_27b.yml @@ -166,6 +166,12 @@ wan_cfg_before_unpatchify: True wan_cross_attn_prescale_kv: False # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32". wan_rope_accum: "auto" +# 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 b53904424..52dea3562 100644 --- a/src/maxdiffusion/configs/base_wan_animate.yml +++ b/src/maxdiffusion/configs/base_wan_animate.yml @@ -111,6 +111,12 @@ wan_cfg_before_unpatchify: True wan_cross_attn_prescale_kv: False # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32". wan_rope_accum: "auto" +# 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 0c403f7e5..f03a260e1 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_14b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_14b.yml @@ -113,6 +113,12 @@ wan_cfg_before_unpatchify: True wan_cross_attn_prescale_kv: False # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32". wan_rope_accum: "auto" +# 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 0812feefe..aacba8197 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_27b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_27b.yml @@ -113,6 +113,12 @@ wan_cfg_before_unpatchify: True wan_cross_attn_prescale_kv: False # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32". wan_rope_accum: "auto" +# 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 d1d9937e3..56a0bb282 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 @@ -745,6 +746,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)." + ) + + def _run_chunked_ulysses_attention( query: jax.Array, key: jax.Array, @@ -813,6 +971,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: @@ -820,8 +979,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.") @@ -1084,6 +1244,7 @@ def _ulysses_attention( use_k_centering: bool = True, qk_prescaled: bool = False, transpose_out: bool = False, + token_padding: Optional[TokenPadding] = None, ) -> jax.Array: """Ulysses sequence-parallel attention. @@ -1107,6 +1268,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( @@ -1180,7 +1348,9 @@ def wrap_ulysses_attention(query, key, value, attention_mask): raw_query = query raw_value = value context_q_seq_len = 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, :] @@ -1224,7 +1394,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) @@ -1474,6 +1647,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. @@ -1527,10 +1701,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] @@ -1704,6 +1878,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. @@ -1729,6 +1904,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. @@ -1736,21 +1916,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. # @@ -1804,6 +2001,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) @@ -1826,6 +2025,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. @@ -1846,6 +2046,10 @@ 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. Returns: mk_arr: Gating metadata array of shape `(batch, 2, local_q_heads, num_q_blocks)` @@ -1896,8 +2100,16 @@ def _compute_fixed_m_metadata( f"_compute_fixed_m_metadata expects query padded to a multiple of block_q, got q_len={q_len}, block_q={block_q}." ) num_q_blocks = q_len // block_q + + def _q_norm_sq(): + norm_sq = jnp.square(query.astype(jnp.float32)).sum(axis=-1) + if q_valid_len is None: + return norm_sq + pos = jnp.arange(q_len)[None, None, :] + return jnp.where(pos < q_valid_len, norm_sq, 0.0) + if per_q_block: - norm_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1) # (batch, num_q_heads, q_len) + norm_sq = _q_norm_sq() qn_max_sq = norm_sq.reshape(batch_size, num_q_heads, num_q_blocks, block_q).max( axis=-1 ) # (batch, num_q_heads, num_q_blocks) @@ -1907,7 +2119,7 @@ def _compute_fixed_m_metadata( mk_arr = jnp.stack([m_base, fixed_ok], axis=1) # (batch, 2, num_q_heads, num_q_blocks) all_fixed = jnp.all(fixed_ok > 0.5) else: - qn_max_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1).max(axis=-1) # (batch, num_q_heads) + 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 @@ -1943,6 +2155,7 @@ def _ulysses_ring_custom_attention( use_k_centering: bool = False, qk_prescaled: bool = False, transpose_out: bool = False, + 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: @@ -1985,6 +2198,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, @@ -2050,6 +2268,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, ) # (1) Ulysses All-to-All: heads -> sequence @@ -2063,7 +2282,11 @@ def wrap_ulysses_ring_attention(query, key, value): raw_query = query raw_value = value context_q_seq_len = raw_query.shape[2] - actual_kv_seq_len = orig_kv_seq_len if num_ring_shards == 1 else raw_key.shape[2] + actual_kv_seq_len = ( + (token_padding.real_len if token_padding is not None else orig_kv_seq_len) + if num_ring_shards == 1 + else (token_padding.real_len if token_padding is not None else raw_key.shape[2]) + ) real_key = raw_key[:, :, :actual_kv_seq_len, :] query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) @@ -2072,9 +2295,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: @@ -2094,7 +2319,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) @@ -2270,6 +2497,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 @@ -2295,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_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) @@ -2303,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_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] @@ -2402,7 +2632,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 @@ -2412,6 +2650,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)) @@ -2451,6 +2690,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"), ) @@ -2481,6 +2721,7 @@ def ulysses_custom_kernel(q, k, v, context): kernel_name="ulysses_custom", qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + token_padding=context.get("token_padding"), ) @@ -2507,6 +2748,7 @@ def ulysses_ring_custom_kernel(q, k, v, context): kv_heads=context.get("kv_heads", None), qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + token_padding=context.get("token_padding"), ) @@ -2537,6 +2779,7 @@ def ulysses_ring_custom_fixed_m_kernel(q, k, v, context): use_k_centering=resolve_k_centering(context.get("use_k_centering"), ring=True), qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + token_padding=context.get("token_padding"), ) @@ -2567,6 +2810,7 @@ def ulysses_ring_custom_fixed_m_per_q_block_kernel(q, k, v, context): use_k_centering=resolve_k_centering(context.get("use_k_centering"), ring=True), qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + token_padding=context.get("token_padding"), ) @@ -2597,6 +2841,7 @@ def ulysses_ring_custom_bidir_kernel(q, k, v, context): kv_heads=context.get("kv_heads", None), qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + token_padding=context.get("token_padding"), ) @@ -2628,6 +2873,7 @@ def ulysses_custom_fixed_m_kernel(q, k, v, context): use_k_centering=resolve_k_centering(context.get("use_k_centering"), ring=False), qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + token_padding=context.get("token_padding"), ) @@ -2659,6 +2905,7 @@ def ulysses_custom_fixed_m_per_q_block_kernel(q, k, v, context): use_k_centering=resolve_k_centering(context.get("use_k_centering"), ring=False), qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + token_padding=context.get("token_padding"), ) @@ -2682,6 +2929,7 @@ def ulysses_kernel(q, k, v, context): kv_heads=context.get("kv_heads", None), ulysses_shards=context.get("ulysses_shards", -1), kernel_name="ulysses", + token_padding=context.get("token_padding"), ) @@ -2706,6 +2954,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"), ) @@ -2730,6 +2979,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"), ) @@ -2756,6 +3006,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"), ) @@ -2780,6 +3031,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"), ) @@ -2803,6 +3055,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"), ) @@ -2810,7 +3063,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 = { @@ -2859,6 +3120,7 @@ def _apply_attention( qk_prescaled: bool = False, k_prescaled: bool = False, transpose_out: bool = False, + token_padding: Optional[TokenPadding] = None, ): """Routes to different attention kernels using a module-level registry.""" @@ -2955,6 +3217,7 @@ def _apply_attention( "qk_prescaled": qk_prescaled, "k_prescaled": k_prescaled, "transpose_out": transpose_out, + "token_padding": token_padding, } if spatiotemporal_config and spatiotemporal_config.get("use_svg_attention"): @@ -2993,6 +3256,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] @@ -3371,6 +3635,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, @@ -3405,6 +3670,7 @@ def apply_attention( qk_prescaled=qk_prescaled, k_prescaled=k_prescaled, transpose_out=getattr(self, "transpose_out", False), + token_padding=token_padding, ) @@ -3467,6 +3733,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, @@ -3500,6 +3767,7 @@ def apply_attention( qk_prescaled=qk_prescaled, k_prescaled=k_prescaled, transpose_out=getattr(self, "transpose_out", False), + token_padding=token_padding, ) @@ -4069,6 +4337,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)) @@ -4178,6 +4447,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 db083ae0c..b16aacf19 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 @@ -457,6 +461,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. @@ -506,6 +511,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) @@ -608,6 +614,30 @@ def __init__( "use_k_centering": "auto", **(attention_config or {}), } + self.attention = attention + self.attention_config = attention_config + self.flash_min_seq_length = flash_min_seq_length + self.mesh = mesh + + # Resolved at build time so they live on the GraphDef; None -> built-in default. + def _wan_opt(name): + value = attention_config.get(name) + return wan_runtime_options.default(name) if value is None else wan_runtime_options.coerce(name, value) + + self.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) @@ -778,6 +808,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, @@ -866,6 +938,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: @@ -890,6 +967,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 @@ -939,6 +1017,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( @@ -955,6 +1034,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) @@ -974,6 +1055,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 b68350bd7..c9f2791d3 100644 --- a/src/maxdiffusion/tests/fused_rmsnorm_rope_pallas_test.py +++ b/src/maxdiffusion/tests/fused_rmsnorm_rope_pallas_test.py @@ -955,7 +955,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 629e56597..3fafd8c7f 100644 --- a/src/maxdiffusion/tests/wan_runtime_options_test.py +++ b/src/maxdiffusion/tests/wan_runtime_options_test.py @@ -134,7 +134,9 @@ def test_unset_options_resolve_to_defaults_not_env(self): } with mock.patch.dict(os.environ, env): attn = self._attn() - for name in opts.ATTENTION_OPTIONS: + # wan_seq_pad rides in attention_config too, but is consumed by WanModel + # (covered in wan_seq_pad_test). + for name in sorted(set(opts.ATTENTION_OPTIONS) - {"wan_seq_pad"}): with self.subTest(option=name): self.assertEqual(getattr(attn, name), opts.default(name)) @@ -179,6 +181,18 @@ def test_wan_model_resolves_cfg_before_unpatchify(self): model = nnx.eval_shape(lambda: WanModel(rngs=nnx.Rngs(0), **kwargs, wan_cfg_before_unpatchify="false")) self.assertIs(model.wan_cfg_before_unpatchify, False) + def test_enum_options_validation(self): + valid_cases = { + "wan_seq_pad": ("off", "lane"), + } + for name, valid_values in valid_cases.items(): + for val in valid_values: + with self.subTest(option=name, valid=val): + self.assertEqual(opts.resolve_from_config(types.SimpleNamespace(**{name: val}), name), val) + with self.subTest(option=name, invalid="bogus"): + with self.assertRaisesRegex(ValueError, name): + opts.resolve_from_config(types.SimpleNamespace(**{name: "bogus"}), name) + if __name__ == "__main__": unittest.main() diff --git a/src/maxdiffusion/tests/wan_seq_pad_test.py b/src/maxdiffusion/tests/wan_seq_pad_test.py new file mode 100644 index 000000000..8ed626b6e --- /dev/null +++ b/src/maxdiffusion/tests/wan_seq_pad_test.py @@ -0,0 +1,712 @@ +""" +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), 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, 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, + token_padding=token_padding, + ) + + with self.subTest(real_len=real_len, use_fixed_m=use_fixed_m, per_q_block=per_q_block): + with mesh_2, 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_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))) + + 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_masks_tail_pad_tokens(self): + """The tpu7x production pad path, scaled down: U=2 x R=2, aligned real_len. + + 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 (`orig_kv_seq_len=304`). + 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, 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, + 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) + 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, token_padding=pad_r) + self.assertEqual(spy.call_args.kwargs["orig_kv_seq_len"], real_len) + 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 b1a799464..b2b06d1b0 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/utils/wan_runtime_options.py b/src/maxdiffusion/utils/wan_runtime_options.py index 58470bb2e..9befa650a 100644 --- a/src/maxdiffusion/utils/wan_runtime_options.py +++ b/src/maxdiffusion/utils/wan_runtime_options.py @@ -42,6 +42,7 @@ "wan_cfg_before_unpatchify": (True, "bool"), "wan_cross_attn_prescale_kv": (False, "bool"), "wan_rope_accum": ("auto", "str"), + "wan_seq_pad": ("off", "str"), } @@ -61,6 +62,8 @@ def _coerce(name: str, value: Any) -> Any: raise ValueError(f"wan_rope_norm_mode must be 'exact' or 'fused', got {value!r}.") if name == "wan_rope_accum" and val not in ("auto", "dtype", "f32"): raise ValueError(f"wan_rope_accum must be 'auto', 'dtype', or 'f32', got {value!r}.") + if name == "wan_seq_pad" and val not in ("off", "lane"): + raise ValueError(f"wan_seq_pad must be 'off' or 'lane', got {value!r}.") return val @@ -101,6 +104,7 @@ def resolve_from_config(config: Any, name: str) -> Any: "wan_splash_transpose_out", "wan_cross_attn_prescale_kv", "wan_rope_accum", + "wan_seq_pad", )