From df7da19e2462a53de41ce2514d50d2f8e1eaeb95 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Thu, 15 Oct 2026 05:51:20 +0000 Subject: [PATCH] perf(attention): shard-major Ulysses output all-to-all (wan_ulysses_out_a2a="chunked"/"shard_major") Third of four pieces split out of the Wan fast-inference optimizations PR. The switch defaults off ("flat") in the Wan YAMLs and in wan_runtime_options; it is a plain config key plumbed to FlaxWanAttention / NNXAttentionOp / AttentionOp and down to the Ulysses kernels through the dispatch context. - `_ulysses_seq_to_heads` is the inverse Ulysses exchange (attention output, head-sharded -> sequence-sharded). "flat" is the previous `all_to_all(split_axis=seq, concat_axis=1)`; XLA reorders the whole output so the sequence axis is major before the collective and back again after, four relayout copies per layer at the 720p shard shape (~1.5 ms on v6e, ~0.7 ms on tpu7x). "chunked" / "shard_major" is a single layout-preserving all_to_all over a reshaped (num_shards, chunk) leading axis (not chunked compute/comm overlap), so no relayout copies are inserted. Same data movement, same values. - With the custom splash kernel (and S_local % block_q == 0, D == 128, heads_per_tile == 1) the forward Q exchange also stays shard-major (4-D q, `[B, U, H/U, S_local, D]`) and the kernel writes the shard-major output directly (`out_num_shards`), on both the 1D Ulysses path (v6e) and the 2D Ulysses+Ring path (tpu7x): custom_splash_attention.py gains a 4-D q BlockSpec and a 4-D shard-major output for `_splash_attention_forward` and `_splash_attention_forward_ring`; ring_attention_kernel.py threads `out_num_shards` through `_custom_ring_attention_forward` / `make_custom_ring_attention` and sizes the o/l/m accumulators accordingly (`save_residuals` with `out_num_shards > 1` raises NotImplementedError: the lse contract stays `(num_q_heads, orig_q_seq_len)`). `_compute_fixed_m_metadata` accepts the 5-D shard-major query. - Reachability: the Pallas shard-major Q/out path requires S_local % block_q == 0. At 720p/81f with CP=4 the per-device sequence is 18,900 tokens, which no launcher block_q divides (v6e BQ=9472, v7 BQ=6400 at this PR), so on this PR alone the launcher's `chunked` only changes the inverse all-to-all; the 4-D q / `out_num_shards` kernel path engages once feat/wan-fast-inference-optimizations pads the sequence to 18,944 (`wan_seq_pad=lane`) and sets BQ=4736 (4736 x 4 = 18,944). The rough 1-1.5 s per 40-step 720p run measured for the kernel path on a pre-review revision was therefore taken with that padding active (DVFS state unrecorded; not re-measured since). - run_wan_fast_inference.sh: the v6e and v7 profiles set WAN_ULYSSES_OUT_A2A=chunked (overridable); the generic profile keeps flat. Tests - ulysses_out_a2a_test.py (6): flat vs chunked vs shard_major parity of the inverse exchange, the mode guard, the save_residuals + shard-major guard, the 4-D q ring kernel bit-exact vs 3-D (interpret mode off TPU, compiled on TPU), and bit-exact wrapper output for the shard-major 1D Ulysses and 2D Ulysses+Ring paths vs the flat layout -- both wrapper tests spy on make_splash_mha / make_custom_ring_attention and assert out_num_shards == U, i.e. that the Pallas shard-major path was actually taken. - wan_runtime_options_test.py (12): the new switch's YAML default and enum validation. - CI: the three kernel bit-exact tests skip in GitHub Actions. - end_to_end/tpu/run_wan_stack_tests.sh: adds ulysses_out_a2a_test.py. Verified (this head): TPU v6e-8 (jax 0.11.2), every file of end_to_end/tpu/run_wan_stack_tests.sh at this head run directly with unittest: 338 tests, all OK, 1 skipped (cross_attention_pallas_test: off-TPU fallback only); ulysses_out_a2a_test.py 6 (all 4 multi-device tests ran), wan_runtime_options_test.py 12. --- end_to_end/tpu/run_wan_fast_inference.sh | 7 +- end_to_end/tpu/run_wan_stack_tests.sh | 2 + src/maxdiffusion/configs/base_wan_14b.yml | 5 + src/maxdiffusion/configs/base_wan_1_3b.yml | 5 + src/maxdiffusion/configs/base_wan_27b.yml | 5 + src/maxdiffusion/configs/base_wan_animate.yml | 5 + src/maxdiffusion/configs/base_wan_i2v_14b.yml | 5 + src/maxdiffusion/configs/base_wan_i2v_27b.yml | 5 + .../kernels/custom_splash_attention.py | 193 +++++++++-- .../splash_attention/ring_attention_kernel.py | 60 +++- src/maxdiffusion/models/attention_flax.py | 236 +++++++++++-- .../tests/ulysses_out_a2a_test.py | 326 ++++++++++++++++++ .../tests/wan_runtime_options_test.py | 1 + src/maxdiffusion/wan_runtime_options.py | 4 + 14 files changed, 769 insertions(+), 90 deletions(-) create mode 100644 src/maxdiffusion/tests/ulysses_out_a2a_test.py diff --git a/end_to_end/tpu/run_wan_fast_inference.sh b/end_to_end/tpu/run_wan_fast_inference.sh index 392c1794d..1bd508d6f 100755 --- a/end_to_end/tpu/run_wan_fast_inference.sh +++ b/end_to_end/tpu/run_wan_fast_inference.sh @@ -63,7 +63,7 @@ # against the XLA-path video, visually clean); on tpu7x the # output was bit-identical in that measurement. # USE_FUSED_ROPE_KERNEL=false restores the XLA producer. -# WAN_CROSS_ATTN_KERNEL / WAN_PATCH_EMBED_MODE +# WAN_CROSS_ATTN_KERNEL / WAN_PATCH_EMBED_MODE / WAN_ULYSSES_OUT_A2A # override Wan fast-inference kernel/layout switches # DP / CP / PER_DEVICE_BATCH / SEED # override mesh parallelism, per-device batch, or RNG seed (default 12345) @@ -171,6 +171,7 @@ case "$TPU_PROFILE" in DEFAULT_FUSED_ROPE=true DEFAULT_CROSS_ATTN_KERNEL=pallas DEFAULT_PATCH_EMBED_MODE=tokens + DEFAULT_ULYSSES_OUT_A2A=chunked ;; v7) PLATFORM_LIBTPU="$V7_LIBTPU" @@ -189,6 +190,7 @@ case "$TPU_PROFILE" in DEFAULT_FUSED_ROPE=true DEFAULT_CROSS_ATTN_KERNEL=pallas DEFAULT_PATCH_EMBED_MODE=tokens + DEFAULT_ULYSSES_OUT_A2A=chunked ;; *) echo "== warning: unrecognised or non-v6e/v7 accelerator '${ACCEL_TYPE:-unknown}';" \ @@ -209,6 +211,7 @@ case "$TPU_PROFILE" in DEFAULT_FUSED_ROPE=false DEFAULT_CROSS_ATTN_KERNEL=xla DEFAULT_PATCH_EMBED_MODE=tokens + DEFAULT_ULYSSES_OUT_A2A=flat ;; esac @@ -242,6 +245,7 @@ FUSED_ROPE_ARGS=() [ -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} WAN_PATCH_EMBED_MODE=${WAN_PATCH_EMBED_MODE:-$DEFAULT_PATCH_EMBED_MODE} +WAN_ULYSSES_OUT_A2A=${WAN_ULYSSES_OUT_A2A:-$DEFAULT_ULYSSES_OUT_A2A} # 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 @@ -306,6 +310,7 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \ ${FUSED_ROPE_ARGS[@]+"${FUSED_ROPE_ARGS[@]}"} \ wan_cross_attn_kernel="$WAN_CROSS_ATTN_KERNEL" \ wan_patch_embed_mode="$WAN_PATCH_EMBED_MODE" \ + wan_ulysses_out_a2a="$WAN_ULYSSES_OUT_A2A" \ 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 e1eab0758..055d747f1 100755 --- a/end_to_end/tpu/run_wan_stack_tests.sh +++ b/end_to_end/tpu/run_wan_stack_tests.sh @@ -48,5 +48,7 @@ TESTS=( "$T/cross_attention_pallas_test.py" # Token-sharded patch embedding (feat/wan-token-patch-embed) "$T/wan_patch_embed_test.py" + # Shard-major Ulysses output A2A (feat/ulysses-shard-major-a2a) + "$T/ulysses_out_a2a_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 d28697686..68b07a81f 100644 --- a/src/maxdiffusion/configs/base_wan_14b.yml +++ b/src/maxdiffusion/configs/base_wan_14b.yml @@ -143,6 +143,11 @@ wan_cross_attn_cpu_interpret: False # token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the # replicated kernel, no activation collective; equal up to accumulation order). wan_patch_embed_mode: "conv" +# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded): +# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit +# leading sequence-chunk axis so XLA does not relayout the output around the +# collective (used by the custom Ulysses / Ulysses-ring kernels only). +wan_ulysses_out_a2a: "flat" # 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 db8215331..a1ea8504e 100644 --- a/src/maxdiffusion/configs/base_wan_1_3b.yml +++ b/src/maxdiffusion/configs/base_wan_1_3b.yml @@ -140,6 +140,11 @@ wan_cross_attn_cpu_interpret: False # token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the # replicated kernel, no activation collective; equal up to accumulation order). wan_patch_embed_mode: "conv" +# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded): +# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit +# leading sequence-chunk axis so XLA does not relayout the output around the +# collective (used by the custom Ulysses / Ulysses-ring kernels only). +wan_ulysses_out_a2a: "flat" # 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 89020ce4f..b2f61ce66 100644 --- a/src/maxdiffusion/configs/base_wan_27b.yml +++ b/src/maxdiffusion/configs/base_wan_27b.yml @@ -172,6 +172,11 @@ wan_cross_attn_cpu_interpret: False # token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the # replicated kernel, no activation collective; equal up to accumulation order). wan_patch_embed_mode: "conv" +# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded): +# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit +# leading sequence-chunk axis so XLA does not relayout the output around the +# collective (used by the custom Ulysses / Ulysses-ring kernels only). +wan_ulysses_out_a2a: "flat" # 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 ae8a93c42..495255eaa 100644 --- a/src/maxdiffusion/configs/base_wan_animate.yml +++ b/src/maxdiffusion/configs/base_wan_animate.yml @@ -122,6 +122,11 @@ wan_cross_attn_cpu_interpret: False # token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the # replicated kernel, no activation collective; equal up to accumulation order). wan_patch_embed_mode: "conv" +# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded): +# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit +# leading sequence-chunk axis so XLA does not relayout the output around the +# collective (used by the custom Ulysses / Ulysses-ring kernels only). +wan_ulysses_out_a2a: "flat" # 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 7221ca5a7..8b7938f66 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_14b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_14b.yml @@ -124,6 +124,11 @@ wan_cross_attn_cpu_interpret: False # token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the # replicated kernel, no activation collective; equal up to accumulation order). wan_patch_embed_mode: "conv" +# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded): +# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit +# leading sequence-chunk axis so XLA does not relayout the output around the +# collective (used by the custom Ulysses / Ulysses-ring kernels only). +wan_ulysses_out_a2a: "flat" # 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 c3af910c8..435127f35 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_27b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_27b.yml @@ -124,6 +124,11 @@ wan_cross_attn_cpu_interpret: False # token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the # replicated kernel, no activation collective; equal up to accumulation order). wan_patch_embed_mode: "conv" +# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded): +# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit +# leading sequence-chunk axis so XLA does not relayout the output around the +# collective (used by the custom Ulysses / Ulysses-ring kernels only). +wan_ulysses_out_a2a: "flat" # 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/kernels/custom_splash_attention.py b/src/maxdiffusion/kernels/custom_splash_attention.py index 6c0187d3f..df6bfe053 100644 --- a/src/maxdiffusion/kernels/custom_splash_attention.py +++ b/src/maxdiffusion/kernels/custom_splash_attention.py @@ -842,16 +842,38 @@ def _splash_attention_forward( interpret: bool | None = None, save_residuals: bool = False, transpose_out: bool | None = None, + out_num_shards: int = 1, ): if interpret is None: interpret = jax.default_backend() == "cpu" if transpose_out is None: transpose_out = False - num_q_heads, padded_q_seq_len, head_dim_qk = q.shape - head_dim_v = v.shape[-1] + if transpose_out and out_num_shards > 1: + raise ValueError("out_num_shards > 1 is not supported with transpose_out=True.") + bq, bkv = block_sizes.block_q, block_sizes.block_kv bkv_compute = block_sizes.block_kv_compute bkv_compute_in = block_sizes.block_kv_compute_in + if q.ndim == 4: + num_q_shards, num_q_heads, shard_q_seq_len, head_dim_qk = q.shape + padded_q_seq_len = num_q_shards * shard_q_seq_len + if shard_q_seq_len % bq != 0: + raise ValueError(f"4D shard-major q requires shard_q_seq_len ({shard_q_seq_len}) divisible by bq ({bq}).") + q_blocks_per_shard_in = shard_q_seq_len // bq + + def q_index_map(h, i, j, *_): + return (i // q_blocks_per_shard_in, h, i % q_blocks_per_shard_in, 0) + + q_block_spec = pl.BlockSpec((None, None, bq, head_dim_qk), q_index_map) + else: + num_q_heads, padded_q_seq_len, head_dim_qk = q.shape + + def q_index_map(h, i, j, *_): + return (h, i, 0) + + q_block_spec = pl.BlockSpec((None, bq, head_dim_qk), q_index_map) + + head_dim_v = v.shape[-1] num_kv_heads = k.shape[0] padded_kv_seq_len = k.shape[1] @@ -880,9 +902,6 @@ def _splash_attention_forward( if mk.shape[0] != 2 or mk.shape[1] != num_q_heads or mk.shape[2] != grid_height: raise ValueError(f"mk must have shape (2, {num_q_heads}, {grid_height}), got {mk.shape}") - def q_index_map(h, i, j, *_): - return (h, i, 0) - def out_index_map(h, i, j, *_): return h, 0, i @@ -908,6 +927,30 @@ def v_index_map(h, i, j, *_): pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)), pl.BlockSpec((None, bq, head_dim_v), lambda h, i, j, *_: (h, i, 0)), ] + elif out_num_shards > 1: + if actual_q_seq_len % (out_num_shards * bq) != 0: + raise ValueError( + f"out_num_shards={out_num_shards} requires actual_q_seq_len ({actual_q_seq_len}) divisible by " + f"out_num_shards * bq ({out_num_shards * bq})." + ) + q_blocks_per_shard_out = grid_height // out_num_shards + shard_q_seq_len_out = actual_q_seq_len // out_num_shards + + def out_shard_major_index_map(h, i, j, *_): + return (i // q_blocks_per_shard_out, h, 0, i % q_blocks_per_shard_out) + + out_shapes = [ + jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), + jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), + jax.ShapeDtypeStruct((head_dim_v, bq), jnp.float32), + jax.ShapeDtypeStruct((out_num_shards, num_q_heads, head_dim_v, shard_q_seq_len_out), q.dtype), + ] + out_specs = [ + pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), + pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), + pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)), + pl.BlockSpec((None, None, head_dim_v, bq), out_shard_major_index_map), + ] else: out_shapes = [ jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), @@ -922,6 +965,8 @@ def v_index_map(h, i, j, *_): pl.BlockSpec((None, head_dim_v, bq), out_index_map), ] if save_residuals: + if out_num_shards > 1: + raise NotImplementedError("save_residuals is not supported with out_num_shards > 1 (shard-major output).") out_shapes.append(jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, actual_q_seq_len), jnp.float32)) out_specs.append(pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map)) @@ -943,7 +988,7 @@ def v_index_map(h, i, j, *_): if k_mean is None: in_specs = [ - pl.BlockSpec((None, bq, head_dim_qk), q_index_map), + q_block_spec, pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), pl.BlockSpec((None, bkv, head_dim_v), v_index_map), ] @@ -969,7 +1014,7 @@ def v_index_map(h, i, j, *_): k_mean_tiled = jnp.repeat(k_mean[:, None, :], NUM_SUBLANES, axis=1) in_specs = [ - pl.BlockSpec((None, bq, head_dim_qk), q_index_map), + q_block_spec, pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), pl.BlockSpec((None, bkv, head_dim_v), v_index_map), pl.BlockSpec((None, NUM_SUBLANES, head_dim_qk), lambda h, i, j, *_: (h // q_heads_per_kv_head, 0, 0)), @@ -1031,6 +1076,7 @@ def _splash_attention_forward_ring( uniform_fixed_m: bool = False, k_mean: jax.Array | None = None, interpret: bool | None = None, + out_num_shards: int = 1, ): """Ring-specific forward path that returns pre-reciprocal fp32 accumulators. @@ -1041,19 +1087,47 @@ def _splash_attention_forward_ring( normalizes only once at the very end (see `ring_attention_kernel._custom_ring_attention_forward`). + Supports both 3D `q` `(num_q_heads, q_seq_len, head_dim_qk)` and 4D + shard-major `q` `(num_q_shards, num_q_heads, shard_q_seq_len, head_dim_qk)`. + When `out_num_shards > 1`, Pallas writes the output tiles directly into a 4D + shard-major buffer so the ring accumulator produces + `(out_num_shards, num_q_heads, q_seq_len // out_num_shards, head_dim_v)` with + no post-kernel `moveaxis` relayout before the inverse Ulysses `all_to_all`. + Returns: A tuple `(out, m, l)` where - - `out` has shape `(num_q_heads, q_seq_len, head_dim_v)` (fp32, un-normalized), - - `m` and `l` have shape `(num_q_heads, q_seq_len)` (fp32). + - `out` has shape `(num_q_heads, q_seq_len, head_dim_v)` or + `(out_num_shards, num_q_heads, q_seq_len // out_num_shards, head_dim_v)` + (fp32, un-normalized), + - `m` and `l` have shape `(num_q_heads, q_seq_len)` or + `(out_num_shards, num_q_heads, q_seq_len // out_num_shards)` (fp32). """ if interpret is None: interpret = jax.default_backend() == "cpu" - num_q_heads, padded_q_seq_len, head_dim_qk = q.shape - head_dim_v = v.shape[-1] bq, bkv = block_sizes.block_q, block_sizes.block_kv bkv_compute = block_sizes.block_kv_compute bkv_compute_in = block_sizes.block_kv_compute_in + if q.ndim == 4: + num_q_shards, num_q_heads, shard_q_seq_len, head_dim_qk = q.shape + padded_q_seq_len = num_q_shards * shard_q_seq_len + if shard_q_seq_len % bq != 0: + raise ValueError(f"4D shard-major q requires shard_q_seq_len ({shard_q_seq_len}) divisible by bq ({bq}).") + q_blocks_per_shard_in = shard_q_seq_len // bq + + def q_index_map(h, i, j, *_): + return (i // q_blocks_per_shard_in, h, i % q_blocks_per_shard_in, 0) + + q_block_spec = pl.BlockSpec((None, None, bq, head_dim_qk), q_index_map) + else: + num_q_heads, padded_q_seq_len, head_dim_qk = q.shape + + def q_index_map(h, i, j, *_): + return (h, i, 0) + + q_block_spec = pl.BlockSpec((None, bq, head_dim_qk), q_index_map) + + head_dim_v = v.shape[-1] num_kv_heads = k.shape[0] padded_kv_seq_len = k.shape[1] @@ -1063,9 +1137,6 @@ def _splash_attention_forward_ring( raise ValueError(f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA.") q_heads_per_kv_head = num_q_heads // num_kv_heads - def q_index_map(h, i, j, *_): - return (h, i, 0) - def out_index_map(h, i, j, *_): return h, 0, i @@ -1077,22 +1148,51 @@ def v_index_map(h, i, j, *_): grid_width = (actual_kv_seq_len + bkv - 1) // bkv grid_height = (actual_q_seq_len + bq - 1) // bq - out_shapes = [ - jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), - jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), - jax.ShapeDtypeStruct((head_dim_v, bq), jnp.float32), - jax.ShapeDtypeStruct((num_q_heads, head_dim_v, actual_q_seq_len), jnp.float32), - jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, actual_q_seq_len), jnp.float32), - jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, actual_q_seq_len), jnp.float32), - ] - out_specs = [ - pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), - pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), - pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)), - pl.BlockSpec((None, head_dim_v, bq), out_index_map), - pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map), - pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map), - ] + if out_num_shards > 1: + if actual_q_seq_len % (out_num_shards * bq) != 0: + raise ValueError( + f"out_num_shards={out_num_shards} requires actual_q_seq_len ({actual_q_seq_len}) divisible by " + f"out_num_shards * bq ({out_num_shards * bq})." + ) + shard_q_out_len = actual_q_seq_len // out_num_shards + q_blocks_per_shard_out = shard_q_out_len // bq + + def out_shard_index_map(h, i, j, *_): + return (i // q_blocks_per_shard_out, h, 0, i % q_blocks_per_shard_out) + + out_shapes = [ + jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), + jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), + jax.ShapeDtypeStruct((head_dim_v, bq), jnp.float32), + jax.ShapeDtypeStruct((out_num_shards, num_q_heads, head_dim_v, shard_q_out_len), jnp.float32), + jax.ShapeDtypeStruct((out_num_shards, num_q_heads, NUM_SUBLANES, shard_q_out_len), jnp.float32), + jax.ShapeDtypeStruct((out_num_shards, num_q_heads, NUM_SUBLANES, shard_q_out_len), jnp.float32), + ] + out_specs = [ + pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), + pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), + pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)), + pl.BlockSpec((None, None, head_dim_v, bq), out_shard_index_map), + pl.BlockSpec((None, None, NUM_SUBLANES, bq), out_shard_index_map), + pl.BlockSpec((None, None, NUM_SUBLANES, bq), out_shard_index_map), + ] + else: + out_shapes = [ + jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), + jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), + jax.ShapeDtypeStruct((head_dim_v, bq), jnp.float32), + jax.ShapeDtypeStruct((num_q_heads, head_dim_v, actual_q_seq_len), jnp.float32), + jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, actual_q_seq_len), jnp.float32), + jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, actual_q_seq_len), jnp.float32), + ] + out_specs = [ + pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), + pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), + pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)), + pl.BlockSpec((None, head_dim_v, bq), out_index_map), + pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map), + pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map), + ] grid = (num_q_heads, grid_height, grid_width) # Scalar-prefetch operand carrying per-head / per-Q-block fixed-m data: @@ -1116,7 +1216,7 @@ def v_index_map(h, i, j, *_): if k_mean is None: in_specs = [ - pl.BlockSpec((None, bq, head_dim_qk), q_index_map), + q_block_spec, pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), pl.BlockSpec((None, bkv, head_dim_v), v_index_map), ] @@ -1141,7 +1241,7 @@ def v_index_map(h, i, j, *_): k_mean_tiled = jnp.repeat(k_mean[:, None, :], NUM_SUBLANES, axis=1) in_specs = [ - pl.BlockSpec((None, bq, head_dim_qk), q_index_map), + q_block_spec, pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), pl.BlockSpec((None, bkv, head_dim_v), v_index_map), pl.BlockSpec((None, NUM_SUBLANES, head_dim_qk), lambda h, i, j, *_: (h // q_heads_per_kv_head, 0, 0)), @@ -1181,9 +1281,14 @@ def v_index_map(h, i, j, *_): out_shape=out_shapes, interpret=interpret, )(*kernel_args) - out = jnp.swapaxes(all_out[3], 1, 2) # (h, head_dim_v, s) -> (h, s, head_dim_v) - l = all_out[4][:, 0, :] # (h, s) - m = all_out[5][:, 0, :] # (h, s) + if out_num_shards > 1: + out = jnp.swapaxes(all_out[3], 2, 3) # (u, h, head_dim_v, s_loc) -> (u, h, s_loc, head_dim_v) + l = all_out[4][:, :, 0, :] # (u, h, s_loc) + m = all_out[5][:, :, 0, :] # (u, h, s_loc) + else: + out = jnp.swapaxes(all_out[3], 1, 2) # (h, head_dim_v, s) -> (h, s, head_dim_v) + l = all_out[4][:, 0, :] # (h, s) + m = all_out[5][:, 0, :] # (h, s) return out, m, l @@ -2136,6 +2241,7 @@ def make_splash_mha( interpret: bool | None = None, residual_checkpoint_name: str | None = None, transpose_out: bool | None = None, + out_num_shards: int = 1, ): """Builds a splash multi-head attention callable. @@ -2143,8 +2249,8 @@ def make_splash_mha( - Standard dense training uses `_splash_attention_custom` (which defines custom VJP forward and backward rules with activation checkpointing). - `heads_per_tile > 1`, `use_fixed_m=True`, `mk is not None`, `k_mean is not None`, - `interpret is not None`, and `transpose_out=True` are inference/testing paths that - execute forward-only kernels without custom backward passes. + `interpret is not None`, `transpose_out=True`, `out_num_shards > 1`, and `q.ndim == 4` + are inference/testing paths that execute forward-only kernels without custom backward passes. """ if transpose_out is None: transpose_out = False @@ -2161,6 +2267,8 @@ def _splash_attention(q, k, v, mk=None, k_mean=None): if heads_per_tile > 1: if use_fixed_m: raise NotImplementedError("fixed-m is not supported with heads_per_tile > 1") + if out_num_shards > 1: + raise NotImplementedError("out_num_shards > 1 is not supported with heads_per_tile > 1") return _splash_attention_forward_mhpt( q, k, @@ -2174,7 +2282,15 @@ def _splash_attention(q, k, v, mk=None, k_mean=None): vmem_limit_bytes=vmem_limit_bytes, transpose_out=transpose_out, ) - if use_fixed_m or mk is not None or k_mean is not None or interpret is not None or transpose_out: + if ( + use_fixed_m + or mk is not None + or k_mean is not None + or interpret is not None + or transpose_out + or out_num_shards > 1 + or q.ndim == 4 + ): return _splash_attention_forward( q, k, @@ -2191,6 +2307,7 @@ def _splash_attention(q, k, v, mk=None, k_mean=None): k_mean=k_mean, interpret=interpret, transpose_out=transpose_out, + out_num_shards=out_num_shards, ) return _splash_attention_custom( q, diff --git a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py index 3f4cfd907..ac26a133a 100644 --- a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py +++ b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py @@ -750,6 +750,7 @@ def _custom_bidirectional_ring_forward( vmem_limit_bytes: int | None, mask_value: float, ring_axis: str, + out_num_shards: int = 1, ) -> jax.Array: """Wrap-free (bidirectional) ring attention for a NON-wrapping ring axis. @@ -783,6 +784,7 @@ def _attn(kc, vc): use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, + out_num_shards=out_num_shards, ) return o.astype(jnp.float32), m.astype(jnp.float32), l.astype(jnp.float32) @@ -875,13 +877,15 @@ def _custom_ring_attention_forward( v_ok: jax.Array | bool | None = None, all_fixed_global: jax.Array | bool | None = None, save_residuals: bool = False, + out_num_shards: int = 1, ) -> jax.Array | tuple[jax.Array, jax.Array]: """Ring attention using the custom dense splash kernel. Args: - q: Query shard, shape `(num_q_heads, q_seq_len, head_dim_qk)`. Stationary - across ring steps. Must already be padded (and pre-scaled by LOG2E when - `use_base2_exp`) by the caller. + q: Query shard, shape `(num_q_heads, q_seq_len, head_dim_qk)` or 4D + shard-major `(num_q_shards, num_q_heads, shard_q_seq_len, head_dim_qk)`. + Stationary across ring steps. Must already be padded (and pre-scaled by + LOG2E when `use_base2_exp`) by the caller. k: Key shard, shape `(num_kv_heads, kv_seq_len, head_dim_qk)`. Rotated across the ring axis. v: Value shard, shape `(num_kv_heads, kv_seq_len, head_dim_v)`. Rotated. @@ -930,19 +934,33 @@ def _custom_ring_attention_forward( all_fixed_global: Optional precomputed mesh-uniform scalar predicate for the accumulate-vs-LSE `lax.cond`. When supplied, skips the internal ring `pmin`. save_residuals: If True, returns `(out, lse)` where `lse` is the global - logsumexp of shape `(num_q_heads, orig_q_seq_len)`. + logsumexp of shape `(num_q_heads, orig_q_seq_len)`. Not supported together + with `out_num_shards > 1` (raises NotImplementedError). + out_num_shards: Number of leading Ulysses output shards (`1` returns 3D + `(num_q_heads, q_seq_len, head_dim_v)`; `> 1` returns 4D shard-major + `(out_num_shards, num_q_heads, q_seq_len // out_num_shards, head_dim_v)`). Returns: - Normalized attention output of shape `(num_q_heads, q_seq_len, head_dim_v)`, - or `(out, lse)` if `save_residuals=True`. + Normalized attention output, shape `(num_q_heads, q_seq_len, head_dim_v)` or + `(out_num_shards, num_q_heads, q_seq_len // out_num_shards, head_dim_v)`, or + `(out, lse)` if `save_residuals=True`. """ axis_size = lax.axis_size(ring_axis) effective_ring_size = ring_size if ring_size is not None else axis_size effective_kv_seq_len = orig_kv_seq_len * effective_ring_size - num_q_heads = q.shape[0] + num_q_heads = q.shape[1] if q.ndim == 4 else q.shape[0] num_kv_heads = k.shape[0] head_dim_v = v.shape[-1] + out_prefix_shape = ( + (out_num_shards, num_q_heads, orig_q_seq_len // out_num_shards) + if out_num_shards > 1 + else (num_q_heads, orig_q_seq_len) + ) + if save_residuals and out_num_shards > 1: + # The lse would come back in the 4D shard-major layout, which no backward + # pass consumes; keep the documented `(num_q_heads, orig_q_seq_len)` contract. + raise NotImplementedError("save_residuals is not supported with out_num_shards > 1 (shard-major output).") if use_fixed_m and not use_base2_exp: raise NotImplementedError( @@ -996,6 +1014,7 @@ def _custom_ring_attention_forward( vmem_limit_bytes=vmem_limit_bytes, mask_value=mask_value, ring_axis=ring_axis, + out_num_shards=out_num_shards, ) if use_fixed_m and ring_size is not None and ring_size != axis_size: raise NotImplementedError( @@ -1137,8 +1156,8 @@ def _accumulate_scan(_): m_base = jnp.ceil(jnp.sqrt(bound_blocks_sq)) - global_recenter fixed_ok_expanded = jnp.ones_like(m_base) mk_arr = jnp.stack([m_base, fixed_ok_expanded], axis=0) # (2, heads, num_q_blocks) - o_sum = jnp.zeros((num_q_heads, orig_q_seq_len, head_dim_v), jnp.float32) - l_sum = jnp.zeros((num_q_heads, orig_q_seq_len), jnp.float32) + o_sum = jnp.zeros(out_prefix_shape + (head_dim_v,), jnp.float32) + l_sum = jnp.zeros(out_prefix_shape, jnp.float32) k_current, v_current = k, v # Python loop over the (static) ring size rather than a scan: it lets the # LAST hop skip its rotation. A scan body must rotate unconditionally, and @@ -1174,6 +1193,7 @@ def _accumulate_scan(_): # KV block instead of two (a two-body last block is what triggers # the Mosaic scheduler cliff on the whole grid). uniform_fixed_m=True, + out_num_shards=out_num_shards, ) o_sum = o_sum + o_curr.astype(jnp.float32) l_sum = l_sum + l_curr.astype(jnp.float32) @@ -1210,6 +1230,7 @@ def fixed_body(carry, is_last_hop, mk_arr): use_fixed_m=True, mk=mk_arr, k_mean=k_mean, + out_num_shards=out_num_shards, ) m_curr = m_curr.astype(jnp.float32) l_curr = l_curr.astype(jnp.float32) @@ -1257,8 +1278,8 @@ def _lse_scan(_): mk_arr_uniform = jnp.stack([m_base_hop, fixed_ok], axis=0) carry = ( - jnp.zeros((num_q_heads, orig_q_seq_len, head_dim_v), jnp.float32), - jnp.full((num_q_heads, orig_q_seq_len), lse_init, jnp.float32), + jnp.zeros(out_prefix_shape + (head_dim_v,), jnp.float32), + jnp.full(out_prefix_shape, lse_init, jnp.float32), k, v, ) @@ -1278,9 +1299,9 @@ def _lse_scan(_): cond_pred = jnp.all(all_fixed_global) if getattr(all_fixed_global, "ndim", 0) > 0 else all_fixed_global return lax.cond(cond_pred, _accumulate_scan, _lse_scan, None) - o_init = jnp.zeros((num_q_heads, orig_q_seq_len, head_dim_v), jnp.float32) - l_init = jnp.zeros((num_q_heads, orig_q_seq_len), jnp.float32) - m_init = jnp.full((num_q_heads, orig_q_seq_len), mask_value, jnp.float32) + o_init = jnp.zeros(out_prefix_shape + (head_dim_v,), jnp.float32) + l_init = jnp.zeros(out_prefix_shape, jnp.float32) + m_init = jnp.full(out_prefix_shape, mask_value, jnp.float32) m_final, l_final, o_final = m_init, l_init, o_init k_current, v_current = k, v @@ -1303,6 +1324,7 @@ def _lse_scan(_): use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, + out_num_shards=out_num_shards, ) m_curr = m_curr.astype(jnp.float32) l_curr = l_curr.astype(jnp.float32) @@ -1551,12 +1573,15 @@ def make_custom_ring_attention( v_ok: jax.Array | bool | None = None, all_fixed_global: jax.Array | bool | None = None, residual_checkpoint_name: str | None = None, + out_num_shards: int = 1, ): """Builds a ring-attention callable around the custom kernel (supports forward & backward). The returned function takes a single (un-batched) `(q, k, v)` triple of shape - `(num_heads, seq, head_dim)` and optional per-batch `fixed_m_norms=(qn_max_sq, mk_h_sq)` - and `k_mean` to be `jax.vmap`-ped over the batch axis inside the attention `shard_map`. + `(num_heads, seq, head_dim)` (or 4D shard-major `q` of shape + `(num_q_shards, num_heads, shard_seq, head_dim)`) and optional per-batch + `fixed_m_norms=(qn_max_sq, mk_h_sq)` and `k_mean` to be `jax.vmap`-ped over the + batch axis inside the attention `shard_map`. `fixed_m_norms_squared` declares the representation of `fixed_m_norms` and is **required** when `use_fixed_m=True`. The kernel gates in squared-norm space @@ -1622,7 +1647,7 @@ def make_custom_ring_attention( def _ring(q, k, v, fixed_m_norms=None, k_mean=None): norms = fixed_m_norms if fixed_m_norms is not None else default_fixed_m_norms km = k_mean if k_mean is not None else default_k_mean - if use_fixed_m or bidirectional or km is not None or uniform_fixed_m is not None: + if use_fixed_m or bidirectional or km is not None or uniform_fixed_m is not None or out_num_shards > 1 or q.ndim == 4: return _custom_ring_attention_forward( q, k, @@ -1647,6 +1672,7 @@ def _ring(q, k, v, fixed_m_norms=None, k_mean=None): uniform_fixed_m=uniform_fixed_m, v_ok=v_ok, all_fixed_global=all_fixed_global, + out_num_shards=out_num_shards, ) return _custom_ring_attention_custom( q, diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 77ccdb9a0..0bc8f9292 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -716,6 +716,51 @@ def _ulysses_head_chunk_ranges(num_heads: int, ulysses_shards: int, num_chunks: return ranges +ULYSSES_OUT_A2A_MODES = ("flat", "chunked", "shard_major") + + +def _ulysses_seq_to_heads( + x: jax.Array, *, axis_name, seq_axis: int, num_shards: int, mode: Optional[str] = None +) -> jax.Array: + """Inverse Ulysses exchange: `all_to_all(x, split_axis=seq_axis, concat_axis=1, tiled=True)`. + + Shard c receives the c-th sequence chunk of every local head. `x` is + [B, H_local, ...] with the full sequence on `seq_axis`; the result is + [B, num_shards * H_local, ...] with the sequence cut to one chunk. + + `wan_ulysses_out_a2a` picks how the collective is written: + * "flat": a plain split of the sequence axis. XLA reorders the whole output + so the sequence axis is major before the collective, then reorders it back + for the output projection. That is four relayout copies per layer at the + 720p shard shape (~1.5 ms on v6e, ~0.7 ms on tpu7x). + * "chunked" / "shard_major": a single layout-preserving `all_to_all` over a + reshaped `(num_shards, chunk)` leading axis (`split_axis=1, concat_axis=1`, + not head-chunked compute/comm overlap like `ulysses_attention_chunks`). + Each shard's block is then contiguous in whatever minor-dimension order + the producer used. The received `[src, H_local, ...]` blocks merge into + `[src * H_local, ...]`, which is the order `"flat"` produces. Same data + movement, same values. When the custom Pallas splash/ring kernel runs with + `S_local % block_q == 0` and `D == 128`, `"chunked"` / `"shard_major"` + additionally keeps the forward Q exchange shard-major (`4D q`) and emits + the shard-major output layout directly from Pallas + (`out_num_shards=num_shards`), avoiding the post-kernel relayout copy on + both v6e (1D Ulysses) and tpu7x (2D Ulysses+Ring). Measured at roughly + 1-1.5 s per 40-step 720p run on each, on a pre-review revision with the + DVFS state unrecorded; not re-measured since. + """ + if mode is None: + mode = "flat" + if mode not in ULYSSES_OUT_A2A_MODES: + raise ValueError(f"wan_ulysses_out_a2a must be one of {ULYSSES_OUT_A2A_MODES}, got {mode!r}.") + seq_len = x.shape[seq_axis] + if mode == "flat" or num_shards == 1 or seq_len % num_shards: + return jax.lax.all_to_all(x, axis_name=axis_name, split_axis=seq_axis, concat_axis=1, tiled=True) + x = x.reshape(x.shape[:seq_axis] + (num_shards, seq_len // num_shards) + x.shape[seq_axis + 1 :]) + x = jnp.moveaxis(x, seq_axis, 1) # [B, num_shards, H_local, ...] + x = jax.lax.all_to_all(x, axis_name=axis_name, split_axis=1, concat_axis=1, tiled=True) + return x.reshape((x.shape[0], x.shape[1] * x.shape[2]) + x.shape[3:]) + + def _run_chunked_ulysses_attention( query: jax.Array, key: jax.Array, @@ -1055,6 +1100,7 @@ def _ulysses_attention( use_k_centering: bool = True, qk_prescaled: bool = False, transpose_out: bool = False, + wan_ulysses_out_a2a: str = "flat", ) -> jax.Array: """Ulysses sequence-parallel attention. @@ -1123,9 +1169,45 @@ def wrap_ulysses_attention(query, key, value, attention_mask): # into the producer of Q and its 185MB round-trip disappears. if use_custom_kernel and use_base2_exp and not qk_prescaled: query = query * LOG2E + + eff_out_a2a_mode = wan_ulysses_out_a2a + if eff_out_a2a_mode not in ULYSSES_OUT_A2A_MODES: + raise ValueError(f"wan_ulysses_out_a2a must be one of {ULYSSES_OUT_A2A_MODES}, got {eff_out_a2a_mode!r}.") + + if use_custom_kernel: + ( + bq, + bkv, + bkv_compute, + bkv_compute_in, + heads_per_tile, + vmem_limit_bytes, + ) = _extract_custom_block_sizes(flash_block_sizes) + else: + bq = bkv = bkv_compute = bkv_compute_in = heads_per_tile = 1 + vmem_limit_bytes = None + + use_shard_major = ( + use_custom_kernel + and eff_out_a2a_mode in ("chunked", "shard_major") + and not transpose_out + and attention_mask is None + and num_shards > 1 + and (query.shape[2] % bq == 0) + and (query.shape[-1] == 128) + and (value.shape[-1] == 128) + and (heads_per_tile == 1) + ) + out_num_shards = num_shards if use_shard_major else 1 + # Swap sharding: each device gives up a slice of heads and gathers # a slice of sequence, so the local kernel sees the full sequence. - query = jax.lax.all_to_all(query, axis_name=axis_name, split_axis=1, concat_axis=2, tiled=True) + if use_shard_major: + b_q, h_q, s_loc_q, d_q = query.shape + query = query.reshape(b_q, num_shards, h_q // num_shards, s_loc_q, d_q) + query = jax.lax.all_to_all(query, axis_name=axis_name, split_axis=1, concat_axis=1, tiled=True) + else: + query = jax.lax.all_to_all(query, axis_name=axis_name, split_axis=1, concat_axis=2, tiled=True) key = jax.lax.all_to_all(key, axis_name=axis_name, split_axis=1, concat_axis=2, tiled=True) value = jax.lax.all_to_all(value, axis_name=axis_name, split_axis=1, concat_axis=2, tiled=True) if attention_mask is not None and mask_needs_ulysses_gather: @@ -1137,27 +1219,23 @@ def wrap_ulysses_attention(query, key, value, attention_mask): "The custom dense splash kernel (use_custom_kernel) does not support attention_mask " "(it only handles padding via orig_seq_len); got a non-None attention_mask." ) - ( - bq, - bkv, - bkv_compute, - bkv_compute_in, - heads_per_tile, - vmem_limit_bytes, - ) = _extract_custom_block_sizes(flash_block_sizes) # NOTE: the base-2 rescale of Q is applied before the all-to-all above. raw_key = key raw_query = query raw_value = value - context_q_seq_len = raw_query.shape[2] + context_q_seq_len = raw_query.shape[1] * raw_query.shape[3] if use_shard_major else raw_query.shape[2] actual_kv_seq_len = orig_kv_seq_len real_key = raw_key[:, :, :actual_kv_seq_len, :] recenter, safe_bound = custom_splash.get_fixed_m_constants(actual_kv_seq_len) - query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) + if use_shard_major: + kv_size = raw_query.shape[-1] + query_seq_len = context_q_seq_len + else: + query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) k_mean = None if use_fixed_m and use_k_centering: # Virtual k-centering (output-invariant): project q^T \bar{k} inside the @@ -1212,6 +1290,7 @@ def wrap_ulysses_attention(query, key, value, attention_mask): use_fixed_m=True, uniform_fixed_m=True, transpose_out=transpose_out, + out_num_shards=out_num_shards, ) splash_kernel_hybrid = custom_splash.make_splash_mha( block_sizes=bsizes, @@ -1224,6 +1303,7 @@ def wrap_ulysses_attention(query, key, value, attention_mask): use_fixed_m=True, uniform_fixed_m=False, transpose_out=transpose_out, + out_num_shards=out_num_shards, ) def _run_uniform(q, k, v, m, km): @@ -1254,30 +1334,29 @@ def _run_hybrid(q, k, v, m, km): use_fixed_m=False, residual_checkpoint_name=residual_checkpoint_name, transpose_out=transpose_out, + out_num_shards=out_num_shards, ) vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0)) attention_output = vmapped_splash(query, key, value) + if use_shard_major: + attention_output = attention_output.astype(query.dtype) + attention_output = jax.lax.all_to_all(attention_output, axis_name=axis_name, split_axis=1, concat_axis=1, tiled=True) + return attention_output.reshape( + (attention_output.shape[0], attention_output.shape[1] * attention_output.shape[2]) + attention_output.shape[3:] + ) if transpose_out: attention_output = attention_output[:, :, :context_q_seq_len, :kv_size].astype(query.dtype) # Restore original layout: head-sharded/full-sequence -> sequence-sharded/full-heads. # Sequence axis is at index 2 (sublanes), heads axis is at index 1. - attention_output = jax.lax.all_to_all( - attention_output, - axis_name=axis_name, - split_axis=2, - concat_axis=1, - tiled=True, + attention_output = _ulysses_seq_to_heads( + attention_output, axis_name=axis_name, seq_axis=2, num_shards=num_shards, mode=eff_out_a2a_mode ) else: attention_output = attention_output[:, :, :kv_size, :context_q_seq_len].astype(query.dtype) # Restore original layout: head-sharded/full-sequence -> sequence-sharded/full-heads. # Sequence axis is at index 3, heads axis is at index 1. - attention_output = jax.lax.all_to_all( - attention_output, - axis_name=axis_name, - split_axis=3, - concat_axis=1, - tiled=True, + attention_output = _ulysses_seq_to_heads( + attention_output, axis_name=axis_name, seq_axis=3, num_shards=num_shards, mode=eff_out_a2a_mode ) return attention_output else: @@ -1738,6 +1817,9 @@ def _ring_fixed_m_norms_pre_a2a( post_a2a_seq = local_seq * num_ulysses_shards num_q_blocks = -(-post_a2a_seq // block_q) padded_seq = num_q_blocks * block_q + # `concat_axis=2` orders tokens by `(U_src, local_seq)`. When `use_shard_major` + # is active (`local_seq % block_q == 0`), the 4-D `q` BlockSpec also indexes Q + # blocks by `(U_src, q_block_in_shard)`, so `qn_dev` block `i` aligns 1-to-1. q_norm_sq_dev = jax.lax.all_to_all(q_norm_sq, axis_name=ulysses_axis, split_axis=1, concat_axis=2, tiled=True) pad_len = padded_seq - post_a2a_seq if pad_len > 0: @@ -1826,7 +1908,13 @@ def _compute_fixed_m_metadata( - `mk_arr[:, 1, h, i]`: Discrete eligibility predicate (1.0 for fixed-m, 0.0 for online). all_fixed: Boolean scalar indicating if all elements are eligible for uniform fixed-m. """ - batch_size, num_q_heads, q_len, _ = query.shape + if query.ndim == 5: + batch_size, num_q_shards, num_q_heads, shard_q_len, _ = query.shape + q_len = num_q_shards * shard_q_len + else: + batch_size, num_q_heads, q_len, _ = query.shape + num_q_shards = 1 + shard_q_len = q_len num_kv_heads = key.shape[1] if safe_bound is None or recenter is None: rec, bnd = custom_splash.get_fixed_m_constants(key.shape[2], v_max_bound=v_max_bound) @@ -1868,17 +1956,26 @@ def _compute_fixed_m_metadata( ) num_q_blocks = q_len // block_q if per_q_block: - norm_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1) # (batch, num_q_heads, q_len) - 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) + norm_sq = jnp.square(query.astype(jnp.float32)).sum(axis=-1) + if query.ndim == 5: + shard_q_blocks = shard_q_len // block_q + qn_max_sq = norm_sq.reshape(batch_size, num_q_shards, num_q_heads, shard_q_blocks, block_q).max(axis=-1) + qn_max_sq = jnp.moveaxis(qn_max_sq, 1, 2).reshape(batch_size, num_q_heads, num_q_blocks) + else: + 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) bound_sq = qn_max_sq * mk_h_sq[:, :, None] fixed_ok = (bound_sq <= safe_bound_sq).astype(jnp.float32) * v_ok m_base = jnp.ceil(jnp.sqrt(bound_sq)) - recenter - mk_arr = jnp.stack([m_base, fixed_ok], axis=1) # (batch, 2, num_q_heads, num_q_blocks) + mk_arr = jnp.stack([m_base, fixed_ok], axis=1) # (batch, 2, local_q_heads, num_q_blocks) all_fixed = jnp.all(fixed_ok > 0.5) else: - qn_max_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1).max(axis=-1) # (batch, num_q_heads) + norm_sq = jnp.square(query.astype(jnp.float32)).sum(axis=-1) + if query.ndim == 5: + qn_max_sq = norm_sq.max(axis=(1, 3)) + else: + qn_max_sq = 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 @@ -1914,6 +2011,7 @@ def _ulysses_ring_custom_attention( use_k_centering: bool = False, qk_prescaled: bool = False, transpose_out: bool = False, + wan_ulysses_out_a2a: str = "flat", ) -> jax.Array: """2D USP attention (Ulysses + Ring) using custom splash kernel with exact Fixed-m support.""" if kv_heads is None: @@ -2027,9 +2125,42 @@ def wrap_ulysses_ring_attention(query, key, value): use_k_centering=use_k_centering, ) + eff_out_a2a_mode = wan_ulysses_out_a2a + if eff_out_a2a_mode not in ULYSSES_OUT_A2A_MODES: + raise ValueError(f"wan_ulysses_out_a2a must be one of {ULYSSES_OUT_A2A_MODES}, got {eff_out_a2a_mode!r}.") + + # On the 2D Ulysses+Ring path (num_ring_shards > 1), `wan_ulysses_out_a2a` in + # ("chunked", "shard_major") keeps both the forward Q exchange and the inverse + # output exchange in shard-major `[B, U, H_local, S_local, D]` layout + # (`split_axis=1, concat_axis=1`): + # * Q is reshaped to `[B, U, H // U, S_local, D]` before the forward A2A and + # indexed directly as 4D `(U, H_local, S_local, D)` via Pallas `BlockSpec` + # in `_splash_attention_forward_ring`, eliminating the post-A2A `concat_axis=2` + # relayout copy. + # * `_splash_attention_forward_ring` writes `(U, H_local, D, S_local)` via + # `out_num_shards=U` and the ring loop accumulates `o_sum` directly in + # `(U, H_local, S_local, D)`, so the inverse A2A consumes the ring output + # with zero post-kernel `jnp.moveaxis` relayout copy. + use_shard_major = ( + eff_out_a2a_mode in ("chunked", "shard_major") + and not transpose_out + and num_ring_shards > 1 + and num_ulysses_shards > 1 + and (query.shape[2] % bq == 0) + and (query.shape[-1] == 128) + and (value.shape[-1] == 128) + and (heads_per_tile == 1) + ) + out_num_shards = num_ulysses_shards if use_shard_major else 1 + # (1) Ulysses All-to-All: heads -> sequence a2a = functools.partial(jax.lax.all_to_all, axis_name=ulysses_axis, tiled=True) - query = a2a(query, split_axis=1, concat_axis=2) + if use_shard_major: + b_q, h_q, s_loc_q, d_q = query.shape + query = query.reshape(b_q, num_ulysses_shards, h_q // num_ulysses_shards, s_loc_q, d_q) + query = a2a(query, split_axis=1, concat_axis=1) + else: + query = a2a(query, split_axis=1, concat_axis=2) key = a2a(key, split_axis=1, concat_axis=2) value = a2a(value, split_axis=1, concat_axis=2) @@ -2037,11 +2168,15 @@ def wrap_ulysses_ring_attention(query, key, value): raw_key = key raw_query = query raw_value = value - context_q_seq_len = raw_query.shape[2] + context_q_seq_len = raw_query.shape[1] * raw_query.shape[3] if use_shard_major else raw_query.shape[2] actual_kv_seq_len = orig_kv_seq_len if num_ring_shards == 1 else raw_key.shape[2] real_key = raw_key[:, :, :actual_kv_seq_len, :] - query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) + if use_shard_major: + kv_size = raw_query.shape[-1] + query_seq_len = context_q_seq_len + else: + query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) # When actual_kv_seq_len is aligned to 8 sublanes, K/V are passed with NO # sequence padding. The kernel slices the ragged KV tail using slice lengths # derived from actual_kv_seq_len, avoiding sequence pad HBM copies and @@ -2147,6 +2282,7 @@ def _run_hybrid(q, k, v, m, km): pregathered_mk=True, v_ok=v_ok, all_fixed_global=all_fixed_global, + out_num_shards=out_num_shards, ) attention_output = jax.vmap(ring_kernel, in_axes=(0, 0, 0, (0, 0), 0))( query, key, value, (qn_dev, mk_global_sq_dev), k_mean_dev @@ -2164,13 +2300,27 @@ def _run_hybrid(q, k, v, m, km): bidirectional=bidirectional, use_fixed_m=False, residual_checkpoint_name=residual_checkpoint_name, + out_num_shards=out_num_shards, ) attention_output = jax.vmap(ring_kernel, in_axes=(0, 0, 0))(query, key, value) + if use_shard_major: + attention_output = attention_output.astype(query.dtype) + attention_output = a2a(attention_output, split_axis=1, concat_axis=1) + return attention_output.reshape( + (attention_output.shape[0], attention_output.shape[1] * attention_output.shape[2]) + attention_output.shape[3:] + ) + attention_output = attention_output[:, :, :context_q_seq_len, :kv_size].astype(query.dtype) # (3) Ulysses All-to-All back: sequence -> heads - return a2a(attention_output, split_axis=2, concat_axis=1) + return _ulysses_seq_to_heads( + attention_output, + axis_name=ulysses_axis, + seq_axis=2, + num_shards=num_ulysses_shards, + mode=eff_out_a2a_mode, + ) x = _run_chunked_ulysses_attention( query, @@ -2414,6 +2564,7 @@ def ulysses_custom_kernel(q, k, v, context): kernel_name="ulysses_custom", qk_prescaled=qk_prescaled, transpose_out=bool(context.get("transpose_out")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2440,6 +2591,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")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2470,6 +2622,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")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2500,6 +2653,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")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2530,6 +2684,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")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2561,6 +2716,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")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2592,6 +2748,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")), + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2615,6 +2772,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", + wan_ulysses_out_a2a=context.get("wan_ulysses_out_a2a") or "flat", ) @@ -2825,6 +2983,7 @@ def _apply_attention( qk_prescaled: bool = False, k_prescaled: bool = False, transpose_out: bool = False, + wan_ulysses_out_a2a: str = "flat", ): """Routes to different attention kernels using a module-level registry.""" @@ -2898,6 +3057,7 @@ def _apply_attention( "qk_prescaled": qk_prescaled, "k_prescaled": k_prescaled, "transpose_out": transpose_out, + "wan_ulysses_out_a2a": wan_ulysses_out_a2a, } if spatiotemporal_config and spatiotemporal_config.get("use_svg_attention"): @@ -3256,6 +3416,7 @@ def __init__( kv_heads: Optional[int] = None, use_k_centering: bool | str = "auto", transpose_out: bool = False, + wan_ulysses_out_a2a: str = "flat", ): self.dpa_layer = None self.use_base2_exp = use_base2_exp @@ -3264,6 +3425,7 @@ def __init__( self.ulysses_attention_chunks = ulysses_attention_chunks self.use_k_centering = use_k_centering self.transpose_out = bool(transpose_out) + self.wan_ulysses_out_a2a = wan_ulysses_out_a2a or "flat" if attention_kernel == "cudnn_flash_te": from transformer_engine.jax.flax.transformer import DotProductAttention # pytype: disable=import-error @@ -3348,6 +3510,7 @@ def apply_attention( qk_prescaled=qk_prescaled, k_prescaled=k_prescaled, transpose_out=getattr(self, "transpose_out", False), + wan_ulysses_out_a2a=getattr(self, "wan_ulysses_out_a2a", "flat"), ) @@ -3374,6 +3537,7 @@ class AttentionOp(nn.Module): kv_heads: Optional[int] = None use_k_centering: bool | str = "auto" transpose_out: bool = False + wan_ulysses_out_a2a: str = "flat" def setup(self): self.dpa_layer = None @@ -3443,6 +3607,7 @@ def apply_attention( qk_prescaled=qk_prescaled, k_prescaled=k_prescaled, transpose_out=getattr(self, "transpose_out", False), + wan_ulysses_out_a2a=getattr(self, "wan_ulysses_out_a2a", "flat"), ) @@ -3616,6 +3781,7 @@ def __init__( "wan_cross_attn_prescale_kv": None, "wan_rope_accum": None, "wan_cross_attn_kernel": None, + "wan_ulysses_out_a2a": None, "wan_cross_attn_cpu_interpret": False, **(attention_config or {}), } @@ -3649,6 +3815,7 @@ def _wan_opt(name): self.wan_cross_attn_prescale_kv = _wan_opt("wan_cross_attn_prescale_kv") self.wan_rope_accum = _wan_opt("wan_rope_accum") self.wan_cross_attn_kernel = _wan_opt("wan_cross_attn_kernel") + self.wan_ulysses_out_a2a = _wan_opt("wan_ulysses_out_a2a") self.wan_cross_attn_cpu_interpret = _wan_opt("wan_cross_attn_cpu_interpret") if attention_kernel in {"flash", "cudnn_flash_te"} and mesh is None: @@ -3723,6 +3890,7 @@ def _wan_opt(name): ulysses_attention_chunks=attention_config["ulysses_attention_chunks"], use_k_centering=attention_config["use_k_centering"], transpose_out=self.wan_splash_transpose_out, + wan_ulysses_out_a2a=self.wan_ulysses_out_a2a, ) # None axes corresponds to the stacked weights across all blocks # because of the use of nnx.vmap and nnx.scan. diff --git a/src/maxdiffusion/tests/ulysses_out_a2a_test.py b/src/maxdiffusion/tests/ulysses_out_a2a_test.py new file mode 100644 index 000000000..8c5632e59 --- /dev/null +++ b/src/maxdiffusion/tests/ulysses_out_a2a_test.py @@ -0,0 +1,326 @@ +""" +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 the inverse Ulysses exchange (`wan_ulysses_out_a2a`). + +Inside the shard_map each shard holds [B, H/n, ..., S] (its heads, full +sequence) and must end with [B, H, ..., S/n] (every head, its sequence chunk). +Seen globally that is a pure resharding, so the output must equal the input +exactly, in both modes. Needs >= 2 devices; on CPU run with +XLA_FLAGS=--xla_force_host_platform_device_count=4. +""" + +import os +import unittest +from unittest import mock + +import jax +import jax.numpy as jnp +import numpy as np +from jax.sharding import Mesh, PartitionSpec as P + +from maxdiffusion.kernels import custom_splash_attention as custom_splash +from maxdiffusion.kernels.splash_attention import ring_attention_kernel +from maxdiffusion.models import attention_flax + + +def _on_tpu() -> bool: + return jax.devices()[0].platform == "tpu" + + +IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true" +# Kernel numerics grids and multi-device tests are skipped in CI; see +# end_to_end/tpu/run_wan_stack_tests.sh. +_SKIP_IN_GITHUB_ACTIONS = unittest.skipIf( + IN_GITHUB_ACTIONS, "TPU kernel / multi-device test, skipped in GitHub Actions; run end_to_end/tpu/run_wan_stack_tests.sh" +) + + +class UlyssesSeqToHeadsTest(unittest.TestCase): + + def _require_devices(self, min_devices: int = 2) -> int: + n = len(jax.devices()) + if n < min_devices: + self.skipTest(f"needs >= {min_devices} devices (XLA_FLAGS=--xla_force_host_platform_device_count=4 on CPU)") + return 4 if n >= 4 else 2 + + def _exchange(self, mode, x, seq_axis): + n = self._require_devices(2) + mesh = Mesh(np.array(jax.devices()[:n]), ("context",)) + in_spec = P(None, "context", *([None] * (x.ndim - 2))) + out_spec = [None] * x.ndim + out_spec[seq_axis] = "context" + fn = jax.shard_map( + lambda t: attention_flax._ulysses_seq_to_heads(t, axis_name="context", seq_axis=seq_axis, num_shards=n, mode=mode), + mesh=mesh, + in_specs=(in_spec,), + out_specs=P(*out_spec), + check_vma=False, + ) + return jax.jit(fn)(x) + + def test_both_modes_are_a_pure_reshard(self): + n = self._require_devices(2) + heads, d, seq = 2 * n, 128, 24 * n + shapes = {3: (1, heads, d, seq), 2: (1, heads, seq, d)} # [B,H,D,S] (kernel O^T) and [B,H,S,D] + for seq_axis, shape in shapes.items(): + x = jax.random.normal(jax.random.PRNGKey(seq_axis), shape, jnp.float32).astype(jnp.bfloat16) + for mode in attention_flax.ULYSSES_OUT_A2A_MODES: + with self.subTest(seq_axis=seq_axis, mode=mode): + got = self._exchange(mode, x, seq_axis) + self.assertEqual(got.shape, x.shape) + np.testing.assert_array_equal(np.asarray(got, np.float32), np.asarray(x, np.float32)) + + def test_rejects_unknown_mode(self): + n = self._require_devices(2) + x = jnp.zeros((1, 2 * n, 8 * n, 128), jnp.bfloat16) + with self.assertRaisesRegex(ValueError, "wan_ulysses_out_a2a"): + self._exchange("bogus", x, 2) + + @_SKIP_IN_GITHUB_ACTIONS + def test_pallas_shard_major_q_and_out_ring_kernel_bit_exact(self): + u, h, s_loc, d = 2, 2, 256, 128 + s_total = u * s_loc + bs = custom_splash._BlockSizes(block_q=128, block_kv=128, block_kv_compute=128, block_kv_compute_in=128) + q_3d = jax.random.normal(jax.random.PRNGKey(10), (h, s_total, d), jnp.bfloat16) + k = jax.random.normal(jax.random.PRNGKey(11), (h, s_total, d), jnp.bfloat16) * 0.08 + v = jax.random.normal(jax.random.PRNGKey(12), (h, s_total, d), jnp.bfloat16) + q_4d = jnp.swapaxes(q_3d.reshape(h, u, s_loc, d), 0, 1) + + grid_h = s_total // bs.block_q + mk = jnp.stack([jnp.full((h, grid_h), -10.0, jnp.float32), jnp.ones((h, grid_h), jnp.float32)], axis=0) + + o_ref, m_ref, l_ref = custom_splash._splash_attention_forward_ring( + q_3d, + k, + v, + bs, + q_seq_len=s_total, + kv_seq_len=s_total, + use_fixed_m=True, + mk=mk, + uniform_fixed_m=True, + interpret=not _on_tpu(), + out_num_shards=1, + ) + o_sm, m_sm, l_sm = custom_splash._splash_attention_forward_ring( + q_4d, + k, + v, + bs, + q_seq_len=s_total, + kv_seq_len=s_total, + use_fixed_m=True, + mk=mk, + uniform_fixed_m=True, + interpret=not _on_tpu(), + out_num_shards=u, + ) + self.assertEqual(o_sm.shape, (u, h, s_loc, d)) + self.assertEqual(m_sm.shape, (u, h, s_loc)) + self.assertEqual(l_sm.shape, (u, h, s_loc)) + np.testing.assert_array_equal(np.asarray(jnp.swapaxes(o_sm, 0, 1).reshape(h, s_total, d)), np.asarray(o_ref)) + np.testing.assert_array_equal(np.asarray(jnp.swapaxes(m_sm, 0, 1).reshape(h, s_total)), np.asarray(m_ref)) + np.testing.assert_array_equal(np.asarray(jnp.swapaxes(l_sm, 0, 1).reshape(h, s_total)), np.asarray(l_ref)) + + def test_ring_forward_rejects_save_residuals_with_shard_major_output(self): + """The lse contract is `(num_q_heads, orig_q_seq_len)`; shard-major output has no lse layout.""" + u, h, s_loc, d = 2, 2, 128, 128 + bs = custom_splash._BlockSizes(block_q=128, block_kv=128, block_kv_compute=128, block_kv_compute_in=128) + q = jnp.zeros((u, h, s_loc, d), jnp.bfloat16) + kv = jnp.zeros((h, u * s_loc, d), jnp.bfloat16) + mesh = jax.sharding.Mesh(np.array(jax.devices()[:1]), ("ring",)) + + def fwd(q, k, v): + return ring_attention_kernel._custom_ring_attention_forward( + q, + k, + v, + block_sizes=bs, + orig_q_seq_len=u * s_loc, + orig_kv_seq_len=u * s_loc, + use_base2_exp=True, + use_experimental_scheduler=False, + vmem_limit_bytes=None, + mask_value=-1e30, + ring_axis="ring", + save_residuals=True, + out_num_shards=u, + ) + + with self.assertRaisesRegex(NotImplementedError, "out_num_shards"): + custom_splash._splash_attention_forward(q, kv, kv, bs, save_residuals=True, out_num_shards=u) + with self.assertRaisesRegex(NotImplementedError, "out_num_shards"): + jax.shard_map(fwd, mesh=mesh, in_specs=(P(), P(), P()), out_specs=P(), check_vma=False)(q, kv, kv) + + @_SKIP_IN_GITHUB_ACTIONS + def test_ulysses_ring_custom_chunked_matches_flat_bit_exact(self): + from flax.linen import partitioning as nn_partitioning + + self._require_devices(4) + mesh_4 = Mesh(np.array(jax.devices()[:4]).reshape(1, 1, 4, 1), ("data", "fsdp", "context", "tensor")) + 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), + ) + b, heads, s_per_shard, d = 1, 4, 256, 128 + s_total = 4 * s_per_shard + q = jax.random.normal(jax.random.PRNGKey(20), (b, s_total, heads * d), jnp.bfloat16) + k = jax.random.normal(jax.random.PRNGKey(21), (b, s_total, heads * d), jnp.bfloat16) * 0.08 + v = jax.random.normal(jax.random.PRNGKey(22), (b, s_total, heads * d), jnp.bfloat16) + bs = {"block_q": 128, "block_kv": 128, "block_kv_compute": 128, "block_kv_compute_in": 128} + 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, + ) + for use_fixed_m, per_q_block, bidirectional, mode in ( + (True, False, False, "chunked"), + (True, True, False, "shard_major"), + (False, False, False, "chunked"), + (False, False, True, "shard_major"), + ): + with self.subTest(use_fixed_m=use_fixed_m, per_q_block=per_q_block, bidirectional=bidirectional, mode=mode): + with mesh_4, nn_partitioning.axis_rules(rules): + out_flat = 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=bs, + ulysses_shards=2, + use_fixed_m=use_fixed_m, + per_q_block=per_q_block, + bidirectional=bidirectional, + wan_ulysses_out_a2a="flat", + ) + with mock.patch.object( + ring_attention_kernel, "make_custom_ring_attention", wraps=ring_attention_kernel.make_custom_ring_attention + ) as spy: + out_chunked = 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=bs, + ulysses_shards=2, + use_fixed_m=use_fixed_m, + per_q_block=per_q_block, + bidirectional=bidirectional, + wan_ulysses_out_a2a=mode, + ) + # The shard-major Pallas path was taken (out_num_shards == U), not just the inverse a2a. + self.assertTrue(spy.call_args_list) + self.assertTrue(all(c.kwargs["out_num_shards"] == 2 for c in spy.call_args_list)) + self.assertEqual(out_chunked.shape, out_flat.shape) + np.testing.assert_array_equal(np.asarray(out_chunked, jnp.float32), np.asarray(out_flat, jnp.float32)) + + @_SKIP_IN_GITHUB_ACTIONS + def test_ulysses_custom_chunked_matches_flat_bit_exact(self): + from flax.linen import partitioning as nn_partitioning + + n = self._require_devices(2) + mesh_u = Mesh(np.array(jax.devices()[:n]).reshape(1, 1, n, 1), ("data", "fsdp", "context", "tensor")) + 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), + ) + b, heads, s_per_shard, d = 1, 2 * n, 256, 128 + s_total = n * s_per_shard + q = jax.random.normal(jax.random.PRNGKey(30), (b, s_total, heads * d), jnp.bfloat16) + k = jax.random.normal(jax.random.PRNGKey(31), (b, s_total, heads * d), jnp.bfloat16) * 0.08 + v = jax.random.normal(jax.random.PRNGKey(32), (b, s_total, heads * d), jnp.bfloat16) + bs = {"block_q": 128, "block_kv": 128, "block_kv_compute": 128, "block_kv_compute_in": 128} + 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, + ) + for use_fixed_m, per_q_block, use_k_centering, mode in ( + (True, True, True, "chunked"), + (True, False, False, "shard_major"), + (False, False, False, "chunked"), + ): + with self.subTest(use_fixed_m=use_fixed_m, per_q_block=per_q_block, use_k_centering=use_k_centering, mode=mode): + with mesh_u, nn_partitioning.axis_rules(rules): + out_flat = attention_flax._ulysses_attention( + q, + k, + v, + heads=heads, + mesh=mesh_u, + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + flash_block_sizes=bs, + use_custom_kernel=True, + use_fixed_m=use_fixed_m, + per_q_block=per_q_block, + use_k_centering=use_k_centering, + wan_ulysses_out_a2a="flat", + ) + with mock.patch.object(custom_splash, "make_splash_mha", wraps=custom_splash.make_splash_mha) as spy: + out_chunked = attention_flax._ulysses_attention( + q, + k, + v, + heads=heads, + mesh=mesh_u, + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + flash_block_sizes=bs, + use_custom_kernel=True, + use_fixed_m=use_fixed_m, + per_q_block=per_q_block, + use_k_centering=use_k_centering, + wan_ulysses_out_a2a=mode, + ) + # The shard-major Pallas path was taken (out_num_shards == U), not just the inverse a2a. + self.assertTrue(spy.call_args_list) + self.assertTrue(all(c.kwargs["out_num_shards"] == n for c in spy.call_args_list)) + self.assertEqual(out_chunked.shape, out_flat.shape) + np.testing.assert_array_equal(np.asarray(out_chunked, jnp.float32), np.asarray(out_flat, jnp.float32)) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/wan_runtime_options_test.py b/src/maxdiffusion/tests/wan_runtime_options_test.py index f01690656..9995a6382 100644 --- a/src/maxdiffusion/tests/wan_runtime_options_test.py +++ b/src/maxdiffusion/tests/wan_runtime_options_test.py @@ -167,6 +167,7 @@ def test_enum_options_validation(self): valid_cases = { "wan_cross_attn_kernel": ("xla", "pallas"), "wan_patch_embed_mode": ("conv", "tokens"), + "wan_ulysses_out_a2a": ("flat", "chunked", "shard_major"), } for name, valid_values in valid_cases.items(): for val in valid_values: diff --git a/src/maxdiffusion/wan_runtime_options.py b/src/maxdiffusion/wan_runtime_options.py index 94b82af4f..2191ae1ca 100644 --- a/src/maxdiffusion/wan_runtime_options.py +++ b/src/maxdiffusion/wan_runtime_options.py @@ -55,6 +55,7 @@ "wan_rope_accum": ("WAN_ROPE_ACCUM", "auto", "str"), "wan_cross_attn_kernel": ("WAN_CROSS_ATTN_KERNEL", "xla", "str"), "wan_patch_embed_mode": ("WAN_PATCH_EMBED_MODE", "conv", "str"), + "wan_ulysses_out_a2a": ("WAN_ULYSSES_OUT_A2A", "flat", "str"), "wan_cross_attn_cpu_interpret": ("WAN_CROSS_ATTN_CPU_INTERPRET", False, "bool"), } @@ -79,6 +80,8 @@ def _coerce(name: str, value: Any) -> Any: raise ValueError(f"wan_cross_attn_kernel must be 'xla' or 'pallas', got {value!r}.") if name == "wan_patch_embed_mode" and val not in ("conv", "tokens"): raise ValueError(f"wan_patch_embed_mode must be 'conv' or 'tokens', got {value!r}.") + if name == "wan_ulysses_out_a2a" and val not in ("flat", "chunked", "shard_major"): + raise ValueError(f"wan_ulysses_out_a2a must be 'flat', 'chunked', or 'shard_major', got {value!r}.") return val @@ -141,6 +144,7 @@ def resolve_from_config(config: Any, name: str) -> Any: "wan_rope_accum", "wan_cross_attn_kernel", "wan_patch_embed_mode", + "wan_ulysses_out_a2a", "wan_cross_attn_cpu_interpret", )