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
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
mbohlool
left a comment
mbohlool
left a comment
There was a problem hiding this comment.
a requested lane mode should raise rather than turn itself off, and CI should run at least one masking test.
| f"wan_seq_pad='lane' disabled for seq_len={seq_len} < flash_min_seq_length={flash_min_seq_length}.", | ||
| ) | ||
| return None | ||
| if self.attention not in TOKEN_PADDING_KERNELS: |
There was a problem hiding this comment.
Could wan_seq_pad="lane" raise a ValueError in WanModel.__init__ when it can't be honoured (unsupported attention, SVG, memory-efficient attention, bad ulysses_shards, TI2V), instead of warning once and running unpadded? As written, the planner returns None for every non-custom kernel, so the NotImplementedError guards described in the PR body never fire from the model. Only the per-call short-sequence case needs to stay a runtime check.
| 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( |
There was a problem hiding this comment.
Could the masking tests and the end-to-end lane-vs-off test gate on device count instead of GITHUB_ACTIONS? Right now none of them run in CI, not even on the v4-8 TPU job, so nothing in CI checks that the pads are masked. A 4-fake-device CPU run of test_wan_model_end_to_end_lane_padding_matches_off with the kernels in interpret mode would cover the whole pad → mask → unpad path.
| "ulysses_ring_custom_fixed_m_per_q_block", | ||
| }) | ||
|
|
||
| _ACTIVE_TOKEN_PADDING: contextvars.ContextVar[Optional[TokenPadding]] = contextvars.ContextVar( |
There was a problem hiding this comment.
Could TokenPadding be passed down explicitly, like rotary_emb already is, instead of through _ACTIVE_TOKEN_PADDING? It's a trace-time ContextVar matched on sequence length, which goes against #488's "nothing reads switches at trace time". A length mismatch also silently returns None, so any path that sees a different length attends to the pads without an error.
| key, _, _ = _pad_data_for_flash(raw_key, kv_heads, kv_pad_size) | ||
| value, _, _ = _pad_data_for_flash(raw_value, kv_heads, kv_pad_size) | ||
| ring_kv_seq_len = actual_kv_seq_len if actual_kv_seq_len % 8 == 0 else key_seq_len | ||
| ring_kv_seq_len = actual_kv_seq_len |
There was a problem hiding this comment.
This also changes the unpadded ring path. When the per-segment KV length isn't a multiple of 8, the ring kernel used to get the bkv-padded length and now gets the real one. Was the old value letting the zero pad keys into the softmax? If so, it's a correctness fix worth noting in the PR body, ideally with a test at an unaligned length without token padding.
| def tiled_seq_lens(self): | ||
| s = tiled_seq_len(self._full_seq, self._attention, self._context_shards, self._ulysses_shards) | ||
| proxy = type( | ||
| "_PadProxy", |
There was a problem hiding this comment.
Could the padding plan be a module-level function, something like plan_token_padding(seq_len, attention, num_context_shards, ulysses_shards)? Then the benchmark wouldn't need to build a _PadProxy type and call _get_token_padding unbound, and the planner could drop its getattr(self, ...)/hasattr fallbacks on its own attributes.
| if skip_blocks: | ||
| if cached_residual is None: | ||
| raise ValueError("cached_residual must be provided when skip_blocks is True") | ||
| if token_padding is not None and cached_residual.shape[1] == token_padding.real_len * token_padding.num_segments: |
There was a problem hiding this comment.
nit: the model always returns the residual unpadded, so the shape check here always passes. Could this just be if token_padding is not None:? Then a residual of the wrong length fails on the add instead of being accepted silently.
Last of four pieces split out of the Wan fast-inference optimizations PR
(the cross-attention Pallas kernel, token-sharded patch embedding and
shard-major Ulysses A2A are the three PRs below this one). The switch
defaults off in the Wan YAMLs and in wan_runtime_options; it is a plain
config key plumbed to WanModel through `attention_config_entries(config)`.
- Lane sequence padding (`wan_seq_pad="lane"`): `WanModel` pads the video tokens
at the tail to a per-shard multiple of 128 (`75,600 -> 75,776` at
720p/CP=4; on the Ulysses-ring path each ring segment is padded at its own
tail) so the Ulysses reshapes are copy-free, pads the rotary embedding to
match, and unpads the output and residual. The padding descriptor
(`TokenPadding`) is planned via `plan_token_padding` and passed explicitly as
a `token_padding` argument from `WanModel.__call__` ->
`WanTransformerBlock.__call__` -> `FlaxWanAttention.__call__` ->
`NNXAttentionOp.apply_attention` (no process/thread-global `ContextVar`).
- The custom `ulysses_custom*` / `ulysses_ring_custom*` kernels mask the pad
keys by passing the real length as the ragged KV tail; the fixed-m norm
reductions (`_compute_fixed_m_metadata` via `q_valid_len` on both the 1-D
Ulysses path and the R == 1 ring path, `_ring_fixed_m_norms_pre_a2a` via
`row_ok` for R > 1) exclude the pad Q/K/V rows so pad-token values cannot
push a block onto the online-softmax path or move the `k_mean`. The ring's
`ring_kv_seq_len` is always the real length.
- Fails closed on invalid static/mode configs: when `wan_seq_pad="lane"` is
requested, incompatible static configurations (`per_token_t=True` TI2V
timesteps, `use_svg_attention`, `use_memory_efficient_attention`, unsupported
kernels, or `ulysses_shards` that do not divide the context axis) raise
`ValueError` in `plan_token_padding` (checked at `WanModel.__init__` time as
well as per call); `_resolve_token_padding` raises `ValueError` on any
`TokenPadding.total_len != seq_len` mismatch; and attention kernels that
cannot mask active token padding raise `NotImplementedError` via
`_reject_token_padding`. Dynamic short (`S < flash_min_seq_length`) or
indivisible (`S % context_shards != 0`) sequences log a warn-once message and
run unpadded.
- `wan_block_benchmark.py` applies `plan_token_padding` to the benchmarked
sequence length and passes `token_padding` explicitly to the block.
- `run_wan_fast_inference.sh`: the v6e and v7 profiles set `WAN_SEQ_PAD=lane`
and `DEFAULT_ULYSSES_OUT_A2A=shard_major` (overridable); the generic profile
keeps `off` and `flat`. LIBTPU: v6e disables
`--xla_tpu_enable_ici_ag_pipelining` / `--xla_tpu_enable_ag_backward_pipelining`,
v7 enables them and adds `--xla_tpu_enable_offloading_copy_to_sparsecore=false`
and `BQ=4736`.
Performance, measured at the stack tip (Wan 2.2 T2V-A14B, 720p/81f/40 steps,
warm AOT, DVFS unpinned unless noted; CP=4, DP=2, launcher defaults):
v6e-8: denoise 124.5s - 125.0s, generate 126.6s - 127.5s
tpu7x-8: denoise 104.9s, generate 107.8s (106.6s / 109.6s with fused RoPE off)
tpu7x-8, DVFS pinned (`PIN_DVFS_P_STATE=true`):
denoise 91.7s, generate 94.6s with the fused RoPE kernel on
denoise 93.4s, generate 96.4s with the fused RoPE kernel off
Tests
- `wan_seq_pad_test.py` (19): padding plan, `ValueError` guards for invalid
static/mode configs (including `per_token_t=True`) and `_resolve_token_padding`
length mismatches, warn-once logging for short/indivisible sequences, pad-key
masking for `ulysses_custom*` and `ulysses_ring_custom*` (U=2 x R=2, U=4 x R=1,
and shard-major U=2 x R=2), fixed-m metadata ignoring pad query rows,
`NotImplementedError` kernel guards, and end-to-end `WanModel` lane padding
parity. All 19 tests run unconditionally in GitHub Actions (no
`_SKIP_IN_GITHUB_ACTIONS`).
- `wan_runtime_options_test.py` (12): the new switch's YAML default and enum
validation; `wan_seq_pad` is excluded from the `FlaxWanAttention` attribute
check since `WanModel` consumes it.
- `end_to_end/tpu/run_wan_stack_tests.sh`: adds `wan_seq_pad_test.py`.
4 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