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
Perseus14
force-pushed
the
feat/wan-token-patch-embed
branch
from
October 15, 2026 19:20
24bbfdf to
f26b420
Compare
mbohlool
previously approved these changes
Oct 19, 2026
mbohlool
left a comment
mbohlool
left a comment
Collaborator
There was a problem hiding this comment.
Two small items inline.
Perseus14
force-pushed
the
feat/wan-token-patch-embed
branch
from
October 20, 2026 12:51
f26b420 to
388384b
Compare
Perseus14
force-pushed
the
feat/wan-token-patch-embed
branch
2 times, most recently
from
October 20, 2026 14:42
97c7129 to
cf59fab
Compare
Perseus14
force-pushed
the
feat/wan-token-patch-embed
branch
2 times, most recently
from
October 20, 2026 17:08
2bd2d3b to
72712b6
Compare
Perseus14
force-pushed
the
feat/wan-token-patch-embed
branch
from
October 21, 2026 08:48
72712b6 to
c99f1c3
Compare
Perseus14
force-pushed
the
feat/wan-token-patch-embed
branch
from
October 21, 2026 08:55
c99f1c3 to
1924230
Compare
Perseus14
force-pushed
the
feat/wan-token-patch-embed
branch
from
October 21, 2026 09:32
1924230 to
74f32d8
Compare
Perseus14
force-pushed
the
feat/wan-token-patch-embed
branch
from
October 22, 2026 08:54
74f32d8 to
59c1ee5
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 (with kernel and bias cast to
`hidden_states.dtype` to match `nnx.Conv.promote_dtype`). 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). When `wan_patch_embed_mode="tokens"` is requested on an
indivisible latent grid, `WanModel.__call__` raises `ValueError` instead of
silently falling back to `nnx.Conv`.
- `run_wan_fast_inference.sh`: every profile sets `WAN_PATCH_EMBED_MODE=tokens`
(overridable).
Tests
- `wan_patch_embed_test.py` (8): tokens vs conv parity (function level and
through `WanModel`), the eligibility predicate, fail-closed `ValueError` on
indivisible latent grids, and the config-driven mode. All 8 tests run
unconditionally in GitHub Actions.
- `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.
- `end_to_end/tpu/run_wan_stack_tests.sh`: adds `wan_patch_embed_test.py`.
Perseus14
force-pushed
the
feat/wan-token-patch-embed
branch
from
October 22, 2026 22:30
59c1ee5 to
8697341
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
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