Repository navigation
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 several performance optimizations for Wan models on TPU, including a fused short-KV cross-attention Pallas kernel, token-sharded patch embedding, inverse Ulysses all-to-all optimizations, and token sequence padding. A critical issue was identified in the new Pallas cross-attention kernel where the BlockSpec index mapping fails to scale grid indices by their block sizes, which would result in incorrect memory slicing and wrong outputs.
2ba8966 to
e86f174
Compare
e86f174 to
1d4c9f6
Compare
0d91078 to
4703103
Compare
295c3b5 to
7a60c3d
Compare
08d405c to
7db8313
Compare
7db8313 to
6fdc94e
Compare
ebaac14 to
58f889e
Compare
58f889e to
68a580f
Compare
68a580f to
395f932
Compare
|
thanks for these stacks of PRs @Perseus14 ! I wonder if we can do an experiment on videos with lengths of 2.5 sec, 5 sec (already covered), 7.5 sec, and 10 sec to see if we would see any regression on other seq length? |
cdc584f to
2eba645
Compare
Perseus14
changed the title
08ccfb2 to
68c3732
Compare
|
Thanks @eltsai — ran the sweep at 41 / 81 / 121 / 161 frames (2.5 / 5 / 7.5 / 10 s @ 16 fps), 720p, 40 steps, seed 12345, same prompt, DP=2×CP=4, DVFS unpinned, on tpu7x-8 and v6e-8. A =
No regression vs Note that all runs use the launcher's 81 f-tuned flash block sizes (v6e: BQ 9472 / BKV 1024; tpu7x: BQ 4736 / BKV 2048). We did not do a block-size sweep per length: the sequence length (and therefore the per-shard length and its alignment to BQ) differs for each frame count, so the optimal block sizes need to be found separately for each frame choice. The non-81 f numbers above are therefore conservative for all three configurations. * v6e @ 121 f is the one length where the launcher's v6e profile (Ulysses U=4, BQ 9472) is the wrong pick: the per-shard length (27,900 tokens) is not block-aligned, and the 2D-ring recipe that is also in this stack ( All runs were bit-reproducible across cold/warm AOT (identical md5s), and the outputs are visually clean on both platforms. |
0414334 to
34789ff
Compare
34789ff to
87c36b2
Compare
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 is declared to the
attention layers with `self_attention_token_padding` (a ContextVar scoped
to the block trace); `TokenPadding` describes it.
- The custom ulysses / 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 now always the real length.
- Every configuration that disables it logs why, once per cause (TI2V
per-token timesteps, SVG attention, memory-efficient attention, short
sequences, unsupported kernels, ulysses_shards that do not divide the
context axis, indivisible sequences).
- Behaviour change: if padding is active, the paths that cannot mask it
(flash and non-custom ulysses via _tpu_flash_attention, ulysses_ring,
dot_product, cudnn_flash_te, head_local_svg, ulysses_ring_custom_bidir)
raise NotImplementedError via `_reject_active_token_padding` instead of
silently attending to pads.
- wan_block_benchmark.py applies the same padding plan to the benchmarked
sequence length.
- run_wan_fast_inference.sh: the v6e and v7 profiles set WAN_SEQ_PAD=lane
(overridable); the generic profile keeps off. LIBTPU: v6e now 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. These two flags moved out of COMMON_LIBTPU, so the generic
(non-v6e/v7) profile no longer sets them at all. The launcher's
reference-benchmark comment carries the stack-tip numbers below.
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 including
the fused RoPE kernel):
v6e-8: denoise 125.9s, generate 128.05s
tpu7x-8: denoise 105.1s, generate 107.6s
tpu7x-8, DVFS pinned (PIN_DVFS_P_STATE=true), fused RoPE kernel off:
denoise 93.6s, generate 96.1s (pinned with it on was not measured)
Per-feature deltas were not measured separately and are omitted.
Tests
- wan_seq_pad_test.py (18): padding plan and per-cause disable logging,
pad-key masking for ulysses custom (unaligned and aligned real_len, flat and
shard-major out A2A, every fixed-m mode) and ulysses_ring_custom (U=2 x R=2
with k-centring on/off; U=4 x R=1 with k-centring on/off, asserting through
a spy that the R == 1 metadata call receives q_valid_len=real_len; and the
tpu7x production path scaled down: U=2 x R=2, aligned real_len, chunked
out A2A, asserting the shard-major layout was taken; all need >= 4
devices), fixed-m metadata ignoring pad query rows, and the
NotImplementedError guards.
- 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.
- CI: the five multi-device masking / end-to-end tests in wan_seq_pad_test skip in GitHub
Actions; padding-plan logic and runtime options run in CI.
- end_to_end/tpu/run_wan_stack_tests.sh: adds wan_seq_pad_test.py.
Verified (final tree): TPU v6e-8 (jax 0.11.2, end_to_end/tpu/run_wan_stack_tests.sh):
356 passed, 1 skipped (off-TPU fallback only), 282 subtests passed.
wan_seq_pad_test.py run directly: 18 tests OK (all 5 multi-device tests ran).
3 participants
Stacked on
feat/ulysses-shard-major-a2a. Last of four pieces split out of the original #491 (this PR keeps the number). Addswan_seq_pad: "lane"(default"off"):WanModelpads the video tokens at the tail to a per-shard multiple of 128 (75,600 → 75,776at 720p/CP=4; per ring segment on the Ulysses-ring path) so the Ulysses reshapes are copy-free, pads the rotary embedding to match, and unpads the output/residual. The padding is declared to the attention layers withself_attention_token_padding(TokenPadding, aContextVarscoped to the block trace); the customulysses/ulysses_ring_customkernels mask the pad keys via the ragged KV tail, and the fixed-m norm reductions (q_valid_lenon both the 1-D Ulysses path and theR == 1ring path,row_okforR > 1) exclude pad Q/K/V rows so they cannot push a block onto the online-softmax path or movek_mean. Every configuration that disables padding logs why, once per cause. Behaviour change: with padding active, paths that cannot mask it (flash/ non-customulysses,ulysses_ring,dot_product,cudnn_flash_te,head_local_svg,ring bidir) raiseNotImplementedErrorinstead of silently attending to pads.run_wan_fast_inference.shsetslanefor v6e/v7 and carries the LIBTPU flag changes (ici_ag/ag_backwardpipelining off on v6e, on on v7; v7 adds--xla_tpu_enable_offloading_copy_to_sparsecore=falseandBQ=4736).Performance (Wan 2.2 T2V-A14B, 720p / 81 frames / 40 steps,
CP=4, DP=2)Per-PR progression across the 10-PR stack on TPU v6e-8 (DVFS unpinned)
main(f1d0fafb)origin/main(1D UlyssesU=4/ 2D RingU=2, R=2)feat/fixed-m-kernel(per_q_block,U=4)mainfeat/attn-dot-fallback-fixesfeat/ring-attention(U=4/ 2D RingU=2, R=2)mainfeat/fused-rmsnorm-rope-xlafeat/wan-fast-serving(cold / warm AOT)feat/wan-custom-kernels(Pallas fused RoPE ON, cold / warm)feat/wan-cross-attn-pallas(cold / warm AOT)feat/wan-token-patch-embed(cold / warm AOT)feat/ulysses-shard-major-a2a(wan_ulysses_out_a2a="chunked")feat/wan-fast-inference-optimizations(wan_seq_pad="lane",WAN_ULYSSES_OUT_A2A=flatablation)feat/wan-fast-inference-optimizations(full stack, cold / warm AOT)main)USE_FUSED_ROPE_KERNEL=false.Tests:
wan_seq_pad_test.py(18, new; padding plan + disable logging, pad-key masking acrossulyssescustom /ulysses_ring_customincl.U=2×R=2,U=4×R=1, and the scaled-down tpu7xU=2×R=2shard-major path, fixed-m metadata ignoring pad rows,NotImplementedErrorguards),wan_runtime_options_test.py(+enum validation).Stack: #477 → #507 → #478 → #508 → #479 → #488 → #509 → #510 → #511 → #491