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 a token-sharded patch embedding mode (wan_patch_embed_mode) for Wan models to optimize performance by avoiding activation collectives on TPU. It implements embed_patches_as_tokens and can_embed_patches_as_tokens in transformer_wan.py to perform patch embedding as a token-sharded matrix multiplication when the "tokens" mode is selected and eligible, falling back to standard convolution otherwise. The new mode is integrated into the configuration files, runtime options, and fast inference scripts, and is verified with new unit tests. There are no review comments to address, and I have no additional feedback to provide.
251437a to
64a7d8a
Compare
64a7d8a to
24bbfdf
Compare
24bbfdf to
f26b420
Compare
mbohlool
left a comment
mbohlool
left a comment
There was a problem hiding this comment.
Two small items inline.
| if patch_embed_mode == "tokens": | ||
| _warn_once( |
There was a problem hiding this comment.
Raise here instead of warning once and running conv. With kernel_size == strides == patch_size fixed at build time (L677-685), the only way this branch is reached is a latent grid that is not divisible by the patch, which is a configuration error, not something to paper over at trace time. That also removes the cross-module import of the private _warn_once.
| @_SKIP_IN_GITHUB_ACTIONS | ||
| def test_tokens_mode_uses_the_token_path_and_matches_conv(self): |
There was a problem hiding this comment.
Drop _SKIP_IN_GITHUB_ACTIONS on these two tests. They build a 1-layer attention="dot_product" model on the default single-device mesh, so neither 'TPU kernel' nor 'multi-device' applies, and they are the only coverage of the WanModel dispatch at L877-889. As written, GitHub CI never sees the tokens path being selected.
Second of four pieces split out of the Wan fast-inference optimizations PR.
The switch defaults off ("conv") in the Wan YAMLs and in wan_runtime_options;
it is a plain config key plumbed to WanModel through
`attention_config_entries(config)` and resolved at build time
(`_wan_opt`) so it lives on the GraphDef.
- With "tokens", the Conv3d patch embedding becomes a patchify + matmul on
the token-sharded layout (`embed_patches_as_tokens`): the non-overlapping
patches (stride == kernel) are flattened in the kernel's
`(p_t, p_h, p_w, C)` order, sharded by token and multiplied by the
replicated `[p_t*p_h*p_w*C, D]` kernel. Same result as
`collapse(conv(x), 1, -1)` up to accumulation order, but produced directly
in the layout the transformer blocks use, so XLA no longer emits the
activation all-gather / collective-permute when collapsing `(F, H, W)`
into the context-sharded token axis (12.8 ms per step on tpu7x, exposed).
- `can_embed_patches_as_tokens` checks the conv really is a plain
stride-equals-kernel 3D patchify (no groups, masks, dilation, and 'SAME'
pads nothing); otherwise the nnx.Conv path is used unchanged (and a warn-once message logs the fallback reason when wan_patch_embed_mode="tokens" was requested).
- run_wan_fast_inference.sh: every profile sets WAN_PATCH_EMBED_MODE=tokens
(overridable).
Tests
- wan_patch_embed_test.py (7): tokens vs conv parity (function level and
through WanModel), the eligibility predicate, and the config-driven mode.
- wan_runtime_options_test.py (12): the new switch's YAML default and enum
validation; `wan_patch_embed_mode` is excluded from the FlaxWanAttention
attribute check since WanModel consumes it.
- CI: the two WanModel patch-embed parity tests skip in GitHub Actions.
- end_to_end/tpu/run_wan_stack_tests.sh: adds wan_patch_embed_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:
329 tests, all OK, 1 skipped (cross_attention_pallas_test: off-TPU fallback only);
wan_patch_embed_test.py 7, wan_runtime_options_test.py 12.
f26b420 to
388384b
Compare
2 participants
Stacked on
feat/wan-cross-attn-pallas. Second of four pieces split out of #491. Addswan_patch_embed_mode: "tokens"(default"conv"): the Conv3d patch embedding becomes a patchify + matmul on the token-sharded layout (embed_patches_as_tokens) — non-overlapping patches flattened in the kernel's(p_t, p_h, p_w, C)order, sharded by token, multiplied by the replicated kernel. Same result ascollapse(conv(x), 1, -1)up to accumulation order, but produced directly in the layout the transformer blocks use, so XLA no longer emits the activation all-gather / collective-permute when collapsing(F, H, W)into the context-sharded token axis (12.8 ms per step on tpu7x, exposed).can_embed_patches_as_tokensguards that the conv is a plain stride-equals-kernel patchify; otherwise thennx.Convpath is used unchanged (logging the fallback reason once).run_wan_fast_inference.shsetstokensin every profile.E2E latency (TPU v6e-8, Wan 2.2 T2V-A14B, 720p / 81 frames / 40 steps,
CP=4, DP=2, warm AOT, DVFS unpinned): 130.6 s generate / 128.5 s denoise, 2.2 s warm AOT compile (cold AOT: 130.7 s generate / 128.5 s denoise, 30.3 s compile) — neutral on v6e-8's 1D Ulysses mesh (U=4) vs #509 (130.6 s / 128.4 s), while removing the 12.8 ms/step exposed reshard collective on tpu7x-8's 2D Ulysses+Ring mesh (U=2, R=2).Tests:
wan_patch_embed_test.py(7, new; tokens vs conv parity at function andWanModellevel, guard, config-driven mode),wan_runtime_options_test.py(+enum validation).Stack: #477 → #507 → #478 → #508 → #479 → #488 → #509 → #510 → #511 → #491