Skip to content

feat(wan): short-KV text cross-attention Pallas kernel (wan_cross_attn_kernel="pallas") #509

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Perseus14 wants to merge 1 commit into from
Open

feat(wan): short-KV text cross-attention Pallas kernel (wan_cross_attn_kernel="pallas") #509

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions end_to_end/tpu/run_wan_fast_inference.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,8 @@
# against the XLA-path video, visually clean); on tpu7x the
# output was bit-identical in that measurement.
# USE_FUSED_ROPE_KERNEL=false restores the XLA producer.
# WAN_CROSS_ATTN_KERNEL
# override Wan fast-inference kernel/layout switches
# DP / CP / PER_DEVICE_BATCH / SEED
# override mesh parallelism, per-device batch, or RNG seed (default 12345)
set -euo pipefail
Expand Down Expand Up @@ -167,6 +169,7 @@ case "$TPU_PROFILE" in
DEFAULT_VAE_CHUNK=1
DEFAULT_VAE_SPATIAL=8
DEFAULT_FUSED_ROPE=true
DEFAULT_CROSS_ATTN_KERNEL=pallas
;;
v7)
PLATFORM_LIBTPU="$V7_LIBTPU"
Expand All @@ -183,6 +186,7 @@ case "$TPU_PROFILE" in
DEFAULT_VAE_CHUNK=1
DEFAULT_VAE_SPATIAL=8
DEFAULT_FUSED_ROPE=true
DEFAULT_CROSS_ATTN_KERNEL=pallas
;;
*)
echo "== warning: unrecognised or non-v6e/v7 accelerator '${ACCEL_TYPE:-unknown}';" \
Expand All @@ -201,6 +205,7 @@ case "$TPU_PROFILE" in
DEFAULT_VAE_CHUNK=1
DEFAULT_VAE_SPATIAL=8
DEFAULT_FUSED_ROPE=false
DEFAULT_CROSS_ATTN_KERNEL=xla
;;
esac

Expand Down Expand Up @@ -232,6 +237,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_CROSS_ATTN_KERNEL=${WAN_CROSS_ATTN_KERNEL:-$DEFAULT_CROSS_ATTN_KERNEL}

# Mesh: context parallelism carries the Ulysses shards, data parallelism takes
# whatever chips remain. Defaults to CP=4 / DP=2 on an 8-chip slice. On a slice
Expand Down Expand Up @@ -294,6 +300,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_cross_attn_kernel="$WAN_CROSS_ATTN_KERNEL" \
fps=16 ${GUIDANCE_ARGS[@]+"${GUIDANCE_ARGS[@]}"} \
seed="${SEED:-12345}" \
flash_block_sizes="$FLASH_BLOCK_SIZES" \
Expand Down
2 changes: 2 additions & 0 deletions end_to_end/tpu/run_wan_stack_tests.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,5 +44,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"
# Short-KV cross-attention Pallas kernel (feat/wan-cross-attn-pallas)
"$T/cross_attention_pallas_test.py"
)
PYTHONPATH="src${PYTHONPATH:+:$PYTHONPATH}" exec "${PYTHON:-python3}" -m pytest -q -rs "${TESTS[@]}" "$@"
6 changes: 6 additions & 0 deletions src/maxdiffusion/configs/base_wan_14b.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,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"
# Text cross-attention backend where it would otherwise fall back to XLA dot-product
# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical).
wan_cross_attn_kernel: "xla"
# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only).
wan_cross_attn_cpu_interpret: False
# 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.
Expand Down
6 changes: 6 additions & 0 deletions src/maxdiffusion/configs/base_wan_1_3b.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,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"
# Text cross-attention backend where it would otherwise fall back to XLA dot-product
# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical).
wan_cross_attn_kernel: "xla"
# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only).
wan_cross_attn_cpu_interpret: False
# 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.
Expand Down
6 changes: 6 additions & 0 deletions src/maxdiffusion/configs/base_wan_27b.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,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"
# Text cross-attention backend where it would otherwise fall back to XLA dot-product
# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical).
wan_cross_attn_kernel: "xla"
# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only).
wan_cross_attn_cpu_interpret: False
# 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.
Expand Down
6 changes: 6 additions & 0 deletions src/maxdiffusion/configs/base_wan_animate.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,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"
# Text cross-attention backend where it would otherwise fall back to XLA dot-product
# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical).
wan_cross_attn_kernel: "xla"
# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only).
wan_cross_attn_cpu_interpret: False
# 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.
Expand Down
6 changes: 6 additions & 0 deletions src/maxdiffusion/configs/base_wan_i2v_14b.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,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"
# Text cross-attention backend where it would otherwise fall back to XLA dot-product
# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical).
wan_cross_attn_kernel: "xla"
# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only).
wan_cross_attn_cpu_interpret: False
# 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.
Expand Down
6 changes: 6 additions & 0 deletions src/maxdiffusion/configs/base_wan_i2v_27b.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,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"
# Text cross-attention backend where it would otherwise fall back to XLA dot-product
# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical).
wan_cross_attn_kernel: "xla"
# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only).
wan_cross_attn_cpu_interpret: False
# 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.
Expand Down
Loading
Loading