feat(wan): lane sequence padding for self-attention (wan_seq_pad="lane")

Last of four pieces split out of the Wan fast-inference optimizations PR
(the cross-attention Pallas kernel, token-sharded patch embedding and
shard-major Ulysses A2A are the three PRs below this one). The switch
defaults off in the Wan YAMLs and in wan_runtime_options; it is a plain
config key plumbed to WanModel through `attention_config_entries(config)`.

- Lane sequence padding (`wan_seq_pad="lane"`): `WanModel` pads the video tokens
  at the tail to a per-shard multiple of 128 (`75,600 -> 75,776` at
  720p/CP=4; on the Ulysses-ring path each ring segment is padded at its own
  tail) so the Ulysses reshapes are copy-free, pads the rotary embedding to
  match, and unpads the output and residual. The padding descriptor
  (`TokenPadding`) is planned via `plan_token_padding` and passed explicitly as
  a `token_padding` argument from `WanModel.__call__` ->
  `WanTransformerBlock.__call__` -> `FlaxWanAttention.__call__` ->
  `NNXAttentionOp.apply_attention` (no process/thread-global `ContextVar`).
- The custom `ulysses_custom*` / `ulysses_ring_custom*` kernels mask the pad
  keys by passing the real length as the ragged KV tail; the fixed-m norm
  reductions (`_compute_fixed_m_metadata` via `q_valid_len` on both the 1-D
  Ulysses path and the R == 1 ring path, `_ring_fixed_m_norms_pre_a2a` via
  `row_ok` for R > 1) exclude the pad Q/K/V rows so pad-token values cannot
  push a block onto the online-softmax path or move the `k_mean`. The ring's
  `ring_kv_seq_len` is always the real length.
- Fails closed on invalid static/mode configs: when `wan_seq_pad="lane"` is
  requested, incompatible static configurations (`per_token_t=True` TI2V
  timesteps, `use_svg_attention`, `use_memory_efficient_attention`, unsupported
  kernels, or `ulysses_shards` that do not divide the context axis) raise
  `ValueError` in `plan_token_padding` (checked at `WanModel.__init__` time as
  well as per call); `_resolve_token_padding` raises `ValueError` on any
  `TokenPadding.total_len != seq_len` mismatch; and attention kernels that
  cannot mask active token padding raise `NotImplementedError` via
  `_reject_token_padding`. Dynamic short (`S < flash_min_seq_length`) or
  indivisible (`S % context_shards != 0`) sequences log a warn-once message and
  run unpadded.
- `wan_block_benchmark.py` applies `plan_token_padding` to the benchmarked
  sequence length and passes `token_padding` explicitly to the block.
- `run_wan_fast_inference.sh`: the v6e and v7 profiles set `WAN_SEQ_PAD=lane`
  and `DEFAULT_ULYSSES_OUT_A2A=shard_major` (overridable); the generic profile
  keeps `off` and `flat`. LIBTPU: v6e disables
  `--xla_tpu_enable_ici_ag_pipelining` / `--xla_tpu_enable_ag_backward_pipelining`,
  v7 enables them and adds `--xla_tpu_enable_offloading_copy_to_sparsecore=false`
  and `BQ=4736`.

Performance, measured at the stack tip (Wan 2.2 T2V-A14B, 720p/81f/40 steps,
warm AOT, DVFS unpinned unless noted; CP=4, DP=2, launcher defaults):
  v6e-8:   denoise 124.5s - 125.0s, generate 126.6s - 127.5s
  tpu7x-8: denoise 104.9s, generate 107.8s (106.6s / 109.6s with fused RoPE off)
  tpu7x-8, DVFS pinned (`PIN_DVFS_P_STATE=true`):
           denoise  91.7s, generate  94.6s with the fused RoPE kernel on
           denoise  93.4s, generate  96.4s with the fused RoPE kernel off

Tests
- `wan_seq_pad_test.py` (19): padding plan, `ValueError` guards for invalid
  static/mode configs (including `per_token_t=True`) and `_resolve_token_padding`
  length mismatches, warn-once logging for short/indivisible sequences, pad-key
  masking for `ulysses_custom*` and `ulysses_ring_custom*` (U=2 x R=2, U=4 x R=1,
  and shard-major U=2 x R=2), fixed-m metadata ignoring pad query rows,
  `NotImplementedError` kernel guards, and end-to-end `WanModel` lane padding
  parity. All 19 tests run unconditionally in GitHub Actions (no
  `_SKIP_IN_GITHUB_ACTIONS`).
- `wan_runtime_options_test.py` (12): the new switch's YAML default and enum
  validation; `wan_seq_pad` is excluded from the `FlaxWanAttention` attribute
  check since `WanModel` consumes it.
- `end_to_end/tpu/run_wan_stack_tests.sh`: adds `wan_seq_pad_test.py`.

Author: Perseus14

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" \

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[@]}" "$@"

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.

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.

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.

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.

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.

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.