Repository navigation
perf(attention): shard-major Ulysses output all-to-all (wan_ulysses_out_a2a="chunked"/"shard_major") - #511
perf(attention): shard-major Ulysses output all-to-all (wan_ulysses_out_a2a="chunked"/"shard_major")#511Perseus14 wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the wan_ulysses_out_a2a configuration option (supporting "flat", "chunked", and "shard_major" modes) to optimize the inverse Ulysses exchange and avoid expensive post-kernel relayout copies on TPUs. This is achieved by updating the custom splash and ring attention kernels to support 4D shard-major query inputs and output sharding, alongside adding a new test suite to verify correctness. The review feedback highlights a critical bug in _ulysses_ring_custom_attention where use_shard_major fails to check for transpose_out, potentially causing layout mismatches. Additionally, the reviewer suggests using the more idiomatic jnp.square instead of ** 2 in JAX operations and wrapping a long line in ring_attention_kernel.py to comply with PEP 8 line length guidelines.
2cabbb8 to
f12639a
Compare
f12639a to
86f71b7
Compare
86f71b7 to
df7da19
Compare
mbohlool
left a comment
mbohlool
left a comment
There was a problem hiding this comment.
The main asks: a requested shard_major should not silently run flat, and CI should run the single-device kernel test.
| 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) | ||
| ) |
There was a problem hiding this comment.
Make shard_major strict: raise when use_custom_kernel, S_local % block_q == 0, D == 128, heads_per_tile == 1 or not transpose_out does not hold (same gate at L2144). Today a requested shard_major silently keeps only the inverse a2a form and skips the 4-D q / out_num_shards kernel path, which, per your reachability note, is exactly what the launcher default hits at this PR. chunked can stay the lenient a2a-only form, which also gives the two names distinct meanings.
| 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", |
There was a problem hiding this comment.
attention="ulysses" accepts wan_ulysses_out_a2a="shard_major" and runs flat: this kernel forwards the mode, but the non-custom branch of _ulysses_attention never reads it (the only _ulysses_seq_to_heads callers are in the custom-kernel branch and the ring path). Either route that branch's inverse exchange through _ulysses_seq_to_heads too, or raise when a non-flat mode reaches a kernel that cannot honour it.
| @_SKIP_IN_GITHUB_ACTIONS | ||
| def test_pallas_shard_major_q_and_out_ring_kernel_bit_exact(self): |
There was a problem hiding this comment.
Remove _SKIP_IN_GITHUB_ACTIONS from this test. It is single-device and already runs under interpret=not _on_tpu(); on the CPU runner it finishes in seconds and is the only check of the 4-D q BlockSpec and out_shard_index_map. With it skipped and the other tests needing >= 2 devices, GitHub CI exercises nothing of the shard-major kernel path except the save_residuals guard. test_rejects_unknown_mode does not need two devices either.
| 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}.") |
There was a problem hiding this comment.
nit: chunked and shard_major are the same value everywhere (ULYSSES_OUT_A2A_MODES, both use_shard_major gates), and chunked collides with ulysses_attention_chunks, which is a different mechanism. Keep one name per behaviour: either drop chunked, or give it the a2a-only meaning and make shard_major the strict kernel-path mode (see the comment at L1190).
| 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: |
There was a problem hiding this comment.
nit: drop the seq_len % num_shards clause. The flat all_to_all(split_axis=seq_axis, tiled=True) it falls back to needs the same divisibility, so this only swaps a clear reshape error for an all_to_all error while silently switching mode.
df7da19 to
808d7bd
Compare
808d7bd to
01c2d43
Compare
6f8ec7d to
5b4eddb
Compare
5b4eddb to
52ec396
Compare
208901a to
09d9dab
Compare
…ut_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 (`"flat"` | `"chunked"` | `"shard_major"`) 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), and validates `seq_len % num_shards == 0`
up front for all modes. `"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"` uses a single
layout-preserving `all_to_all` over a reshaped `(num_shards, chunk)` leading
axis, so no relayout copies are inserted (also wired into the non-custom
`attention="ulysses"` branch). Same data movement, same values.
- With the custom splash kernel (`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.
- `"shard_major"` vs `"chunked"`: `"shard_major"` is the strict mode that
requires the 4-D `q` / `out_num_shards` Pallas kernel path and raises
`ValueError` when any gate precondition (`use_custom_kernel`,
`S_local % block_q == 0`, `D == 128`, `heads_per_tile == 1`,
`not transpose_out`, `attention_mask is None`) fails; `"chunked"` is the
lenient mode that takes the 4-D kernel path when eligible and falls back to
the chunked inverse `all_to_all` in `_ulysses_seq_to_heads` otherwise.
- Reachability: at 720p/81f with CP=4 the unpadded per-device sequence is
18,900 tokens, which no launcher `block_q` divides (v6e `BQ=9472`, v7
`BQ=6400` at this PR), so at this PR the launcher's `v6e`/`v7` profiles set
`DEFAULT_ULYSSES_OUT_A2A=chunked`; once `feat/wan-fast-inference-optimizations`
pads the sequence to 18,944 (`wan_seq_pad=lane`) and sets v7 `BQ=4736`
(`9472 * 2 = 18,944` on v6e, `4736 * 4 = 18,944` on v7), it switches both
profiles to `DEFAULT_ULYSSES_OUT_A2A=shard_major`.
Tests
- `ulysses_out_a2a_test.py` (8): flat vs chunked vs shard_major parity of the
inverse exchange, the mode guard, the `seq_len % num_shards != 0` guard, the
`save_residuals + shard-major` guard, the strict `"shard_major"` kernel-gate
`ValueError` guard, the 4-D q ring kernel bit-exact vs 3-D (running in CI via
`interpret=True`), and bit-exact wrapper output for the shard-major 1D Ulysses
and 2D Ulysses+Ring paths vs the flat layout (spy-asserting
`out_num_shards == U`).
- `wan_runtime_options_test.py` (12): the new switch's YAML default and enum
validation.
- `end_to_end/tpu/run_wan_stack_tests.sh`: adds `ulysses_out_a2a_test.py`.
09d9dab to
84a68e3
Compare
2 participants
Stacked on
feat/wan-token-patch-embed. Third of four pieces split out of #491. Addswan_ulysses_out_a2a: "chunked" | "shard_major"(default"flat"):_ulysses_seq_to_headsperforms the inverse Ulysses exchange as a single layout-preservingall_to_allover a reshaped(num_shards, chunk)leading axis, so XLA stops inserting the four relayout copies per layer that the flatsplit_axis=seqform costs at the 720p shard shape (~1.5 ms on v6e, ~0.7 ms on tpu7x). Same data movement, same values. With the custom splash kernel (S_local % block_q == 0,D == 128,heads_per_tile == 1) the forward Q exchange also stays shard-major (4-Dq) 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.pygains a 4-DqBlockSpec/ 4-D output,ring_attention_kernel.pythreadsout_num_shardsthrough the ring accumulators (save_residuals=Truewithout_num_shards > 1raisesNotImplementedError), and_compute_fixed_m_metadataaccepts the 5-D shard-major query.Reachability note: The Pallas shard-major Q/out path requires
S_local % block_q == 0. At 720p/81f withCP=4the per-device sequence is 18,900 tokens, which no launcherblock_qdivides (BQ=9472on v6e,BQ=6400on v7 at this PR), so on this PR alone the launcher'schunkedsetting only changes the inverse all-to-all; the 4-Dq/out_num_shardsPallas kernel path engages once#491(feat/wan-fast-inference-optimizations) pads the sequence to 18,944 (wan_seq_pad=lane) and setsBQ=4736(4736 × 4 = 18,944).E2E latency (TPU v6e-8, Wan 2.2 T2V-A14B, 720p / 81 frames / 40 steps,
CP=4, DP=2,wan_ulysses_out_a2a="chunked", DVFS unpinned): 128.9 s generate / 126.7 s denoise (29.6 s compile), down from 130.6 s generate / 128.5 s denoise on #510 (−1.7 s generate / −1.8 s denoise end-to-end from the layout-preserving inverse Ulysses all-to-all alone).Tests:
ulysses_out_a2a_test.py(6, new; flat vs chunked vs shard_major parity, mode guard,save_residuals+ shard-major guard, bit-exact 4-D ring kernel output, and bit-exact wrapper output on the 1D and 2D paths asserting via spy thatout_num_shards == U),wan_runtime_options_test.py(+enum validation).Stack: #477 → #507 → #478 → #508 → #479 → #488 → #509 → #510 → #511 → #491