Repository navigation
Conversation
This was referenced Oct 15, 2026
![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.
Perseus14
added this pull request to stack #512
October 15, 2026 16:00
Perseus14
force-pushed
the
feat/wan-token-patch-embed
branch
from
October 15, 2026 16:05
251437a to
64a7d8a
Compare
Perseus14
force-pushed
the
feat/wan-token-patch-embed
branch
from
October 15, 2026 17:53
64a7d8a to
24bbfdf
Compare
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.
Perseus14
force-pushed
the
feat/wan-token-patch-embed
branch
from
October 15, 2026 19:20
24bbfdf to
f26b420
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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