perf(wan): pad sequence to 128-lane multiple for self-attention (wan_seq_pad="lane") by Perseus14 · Pull Request #491 · AI-Hypercomputer/maxdiffusion

Perseus14 · 2026-09-25T16:26:56Z

Stacked on feat/ulysses-shard-major-a2a. Last of four pieces split out of the original #491 (this PR keeps the number). Adds wan_seq_pad: "lane" (default "off"): WanModel pads the video tokens at the tail to a per-shard multiple of 128 (75,600 → 75,776 at 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 with self_attention_token_padding (TokenPadding, a ContextVar scoped to the block trace); the custom ulysses / ulysses_ring_custom kernels mask the pad keys via the ragged KV tail, and the fixed-m norm reductions (q_valid_len on both the 1-D Ulysses path and the R == 1 ring path, row_ok for R > 1) exclude pad Q/K/V rows so they cannot push a block onto the online-softmax path or move k_mean. Every configuration that disables padding logs why, once per cause. Behaviour change: with padding active, paths that cannot mask it (flash / non-custom ulysses, ulysses_ring, dot_product, cudnn_flash_te, head_local_svg, ring bidir) raise NotImplementedError instead of silently attending to pads. run_wan_fast_inference.sh sets lane for v6e/v7 and carries the LIBTPU flag changes (ici_ag/ag_backward pipelining off on v6e, on on v7; v7 adds --xla_tpu_enable_offloading_copy_to_sparsecore=false and BQ=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)

#	PR	Branch	Compile	Generate	Denoise	Step Δ vs prev
0	`main` (`f1d0fafb`)	`origin/main` (1D Ulysses `U=4` / 2D Ring `U=2, R=2`)	181.1 s / 67.3 s	147.1 s / 155.0 s	144.6 s / 152.5 s	baseline
1	#477	`feat/fixed-m-kernel` (`per_q_block`, `U=4`)	193.5 s	139.3 s	136.6 s	−8.0 s denoise vs `main`
2	#507	`feat/attn-dot-fallback-fixes`	13.2 s	138.5 s	136.4 s	neutral (fallback/grid-search fixes)
3	#478	`feat/ring-attention` (`U=4` / 2D Ring `U=2, R=2`)	68.0 s / 68.6 s	138.3 s / 138.8 s	136.1 s / 136.6 s	−15.9 s on 2D Ring vs `main`
4	#508	`feat/fused-rmsnorm-rope-xla`	67.1 s	138.2 s	136.0 s	neutral e2e (+ fixes ~1.5 GB eager OOM)
5	#479	`feat/wan-fast-serving` (cold / warm AOT)	31.0 s / 3.1 s	138.2 s / 138.1 s	136.0 s / 135.9 s	−98.4% warm AOT compile, −85% weight load
6	#488	`feat/wan-custom-kernels` (Pallas fused RoPE ON, cold / warm)	30.7 s / 2.8 s	133.8 s / 133.7 s	131.6 s	−4.3 s denoise vs #479 (136.0 s with RoPE off)
7	#509	`feat/wan-cross-attn-pallas` (cold / warm AOT)	32.8 s / 2.5 s	130.4 s / 130.6 s	128.3 s / 128.4 s	−3.2 s denoise vs #488
8	#510	`feat/wan-token-patch-embed` (cold / warm AOT)	30.3 s / 2.2 s	130.7 s / 130.6 s	128.5 s	neutral on v6e 1D Ulysses (−12.8 ms/step on tpu7x 2D Ring)
9	#511	`feat/ulysses-shard-major-a2a` (`wan_ulysses_out_a2a="chunked"`)	29.6 s	128.9 s	126.7 s	−1.8 s denoise vs #510
10a	#491	`feat/wan-fast-inference-optimizations` (`wan_seq_pad="lane"`, `WAN_ULYSSES_OUT_A2A=flat` ablation)	147.6 s (cold XLA)	129.3 s	126.8 s	−1.7 s denoise vs #510 from lane padding alone
10b	#491	`feat/wan-fast-inference-optimizations` (full stack, cold / warm AOT)	31.0 s / 2.2 s	125.6 s / 125.5 s	123.4 s / 123.3 s	−3.4 s vs #511 (−21.3 s / −14.7% denoise vs `main`)

TPU tpu7x-8 (stack tip perf(wan): pad sequence to 128-lane multiple for self-attention (wan_seq_pad="lane") #491, warm AOT): 107.6 s generate (105.1 s denoise) with DVFS unpinned; 96.1 s generate (93.6 s denoise) with DVFS pinned and USE_FUSED_ROPE_KERNEL=false.

Tests: wan_seq_pad_test.py (18, new; padding plan + disable logging, pad-key masking across ulysses custom / ulysses_ring_custom incl. U=2×R=2, U=4×R=1, and the scaled-down tpu7x U=2×R=2 shard-major path, fixed-m metadata ignoring pad rows, NotImplementedError guards), wan_runtime_options_test.py (+enum validation).

Stack: #477 → #507 → #478 → #508 → #479 → #488 → #509 → #510 → #511 → #491

github-actions · 2026-09-25T16:27:04Z

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

gemini-code-assist

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.

Perseus14 · 2026-10-02T07:30:35Z

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 = main (10af306), B = this stack with the #491 switches off (= #477→#488), C = this stack with the launcher defaults. Denoise seconds for 40 steps (load/compile/VAE excluded):

frames	tpu7x A	tpu7x B	tpu7x C	C vs A	v6e A	v6e B	v6e C	C vs A
41	56.5	43.0	41.5	1.36×	58.8	53.0	49.6	1.19×
81	123.0	109.1	105.0	1.17×	144.6	132.8	125.9	1.15×
121	234.2	216.2	212.7	1.10×	290.8*	308.1	313.0 / 254.6*	1.14×*
161	385.0	362.1	356.5	1.08×	454.7	430.4	419.2	1.08×

No regression vs main at any length; the gain narrows as attention dominates the step. The #491 switches themselves (B→C) are +1.6–3.9 % on tpu7x and +2.7–6.9 % on v6e, except v6e @ 121 f (−1.6 %).

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 (ATTENTION=ulysses_ring_custom_fixed_m ULYSSES_SHARDS=2 BQ=4736 BKV=2048) is 19 % faster there (254.6 s; main with its own ring recipe: 290.8 s, with its Ulysses kernel 657.6 s). At 41 / 81 / 161 f Ulysses remains the better v6e choice by 1–37 %.

All runs were bit-reproducible across cold/warm AOT (identical md5s), and the outputs are visually clean on both platforms.

mbohlool

a requested lane mode should raise rather than turn itself off, and CI should run at least one masking test.

…n_kernel="pallas") First of four pieces split out of the Wan fast-inference optimizations PR. The switch defaults off ("xla") in the Wan YAMLs and in wan_runtime_options; it is a plain config key plumbed to FlaxWanAttention through `attention_config_entries(config)`. - `kernels/cross_attention_pallas.py` keeps a head block's whole 512-token text KV resident in VMEM and computes an exact single-pass softmax per query tile (per-row max, exp2, MXU sum reduction, normalised after PV; no running max / online rescaling). It replaces only what would otherwise be the unmasked XLA dot-product fallback (`_would_use_unmasked_dot_product`, which shares `_FLASH_MIN_SEQ_LENGTH_KERNELS` with the dispatcher). - `wan_cross_attn_kernel="pallas"` fails closed (`ValueError`): static preconditions (TPU mesh or `wan_cross_attn_cpu_interpret=True`, `dim_head % NUM_LANES == 0`, compatible `attention_kernel`, `not use_memory_efficient_attention`, and validated TPU VMEM generation) are checked in `FlaxWanAttention.__init__`, and dynamic input/sharding/VMEM preconditions are checked in `FlaxWanAttention._pallas_cross_attention`. - Dispatch: T2V cross-attention, and the text half of I2V cross-attention (both the image-token split and the text-only path), call `_pallas_cross_attention` first (which returns `None` only when `wan_cross_attn_kernel == "xla"`). The sharded wrapper is cached (`_build_sharded_pallas_cross_attention`) and wrapped with `with_xla_backward` since `pallas_call` has no transpose rule. - `wan_cross_attn_cpu_interpret` (default False, tests only) lets the kernel run under Pallas interpret mode in CI (both on CPU and on unvalidated TPU runners). - In isolation the kernel is 0.51 ms vs 3.94 ms (v6e) and 0.49 ms vs 1.81 ms (tpu7x) per call for the XLA path inside the model. - `run_wan_fast_inference.sh`: the v6e and v7 profiles set `WAN_CROSS_ATTN_KERNEL=pallas` (overridable); the generic profile keeps `xla`. - `wan_block_benchmark.py` passes `split_head_dim` (defaulting to `False` to match the Wan YAMLs), `ulysses_attention_chunks`, and `use_k_centering` from the config to the benchmarked block. Tests - `cross_attention_pallas_test.py` (31 test methods, 32 parameterized cases): golden parity (ragged tiles, `k_prescaled`, batch), fail-closed guards, VMEM budget resolution / working-set estimate / unusable reasons, and `FlaxWanAttention` dispatch (interpret-mode integration, precondition `ValueError` checks, I2V text-vs-image split). - `wan_runtime_options_test.py` (12): YAML default and enum validation. - CI: 30 of the 32 cases in `cross_attention_pallas_test.py` (all parity, guard, VMEM budget, interpret-mode integration, and I2V split tests) run in GitHub Actions; only the full-shard compiled TPU tests (`test_wan_production_shard_shapes` and `test_pallas_mode_uses_the_kernel_and_matches_xla`) skip in CI. - `end_to_end/tpu/run_wan_stack_tests.sh`: adds `cross_attention_pallas_test.py`.

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`.

…ut_a2a="chunked"/"shard_major") Third of four pieces split out of the Wan fast-inference optimizations PR. The switch defaults off ("flat") in the Wan YAMLs and in wan_runtime_options; it is a plain config key (`"flat"` | `"chunked"` | `"shard_major"`) plumbed to FlaxWanAttention / NNXAttentionOp / AttentionOp and down to the Ulysses kernels through the dispatch context. - `_ulysses_seq_to_heads` is the inverse Ulysses exchange (attention output, head-sharded -> sequence-sharded), and validates `seq_len % num_shards == 0` up front for all modes. `"flat"` is the previous `all_to_all(split_axis=seq, concat_axis=1)`; XLA reorders the whole output so the sequence axis is major before the collective and back again after, four relayout copies per layer at the 720p shard shape (~1.5 ms on v6e, ~0.7 ms on tpu7x). `"chunked"` / `"shard_major"` uses a single layout-preserving `all_to_all` over a reshaped `(num_shards, chunk)` leading axis, so no relayout copies are inserted (also wired into the non-custom `attention="ulysses"` branch). Same data movement, same values. - With the custom splash kernel (`S_local % block_q == 0`, `D == 128`, `heads_per_tile == 1`), the forward Q exchange also stays shard-major (4-D q, `[B, U, H/U, S_local, D]`) and the kernel writes the shard-major output directly (`out_num_shards`), on both the 1D Ulysses path (v6e) and the 2D Ulysses+Ring path (tpu7x): `custom_splash_attention.py` gains a 4-D q BlockSpec and a 4-D shard-major output for `_splash_attention_forward` and `_splash_attention_forward_ring`; `ring_attention_kernel.py` threads `out_num_shards` through `_custom_ring_attention_forward` / `make_custom_ring_attention` and sizes the o/l/m accumulators accordingly (`save_residuals` with `out_num_shards > 1` raises `NotImplementedError`: the lse contract stays `(num_q_heads, orig_q_seq_len)`). `_compute_fixed_m_metadata` accepts the 5-D shard-major query. - `"shard_major"` vs `"chunked"`: `"shard_major"` is the strict mode that requires the 4-D `q` / `out_num_shards` Pallas kernel path and raises `ValueError` when any gate precondition (`use_custom_kernel`, `S_local % block_q == 0`, `D == 128`, `heads_per_tile == 1`, `not transpose_out`, `attention_mask is None`) fails; `"chunked"` is the lenient mode that takes the 4-D kernel path when eligible and falls back to the chunked inverse `all_to_all` in `_ulysses_seq_to_heads` otherwise. - Reachability: at 720p/81f with CP=4 the unpadded per-device sequence is 18,900 tokens, which no launcher `block_q` divides (v6e `BQ=9472`, v7 `BQ=6400` at this PR), so at this PR the launcher's `v6e`/`v7` profiles set `DEFAULT_ULYSSES_OUT_A2A=chunked`; once `feat/wan-fast-inference-optimizations` pads the sequence to 18,944 (`wan_seq_pad=lane`) and sets v7 `BQ=4736` (`9472 * 2 = 18,944` on v6e, `4736 * 4 = 18,944` on v7), it switches both profiles to `DEFAULT_ULYSSES_OUT_A2A=shard_major`. Tests - `ulysses_out_a2a_test.py` (8): flat vs chunked vs shard_major parity of the inverse exchange, the mode guard, the `seq_len % num_shards != 0` guard, the `save_residuals + shard-major` guard, the strict `"shard_major"` kernel-gate `ValueError` guard, the 4-D q ring kernel bit-exact vs 3-D (running in CI via `interpret=True`), and bit-exact wrapper output for the shard-major 1D Ulysses and 2D Ulysses+Ring paths vs the flat layout (spy-asserting `out_num_shards == U`). - `wan_runtime_options_test.py` (12): the new switch's YAML default and enum validation. - `end_to_end/tpu/run_wan_stack_tests.sh`: adds `ulysses_out_a2a_test.py`.

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`.

Perseus14 requested a review from entrpn as a code owner September 25, 2026 16:26

gemini-code-assist Bot reviewed Sep 25, 2026

View reviewed changes

Comment thread src/maxdiffusion/kernels/cross_attention_pallas.py

Perseus14 added this pull request to stack #486 September 25, 2026 16:30

Perseus14 force-pushed the feat/wan-fast-inference-optimizations branch from 2ba8966 to e86f174 Compare September 25, 2026 16:34

Perseus14 force-pushed the feat/wan-custom-kernels branch from b855a4e to 05d2e33 Compare September 25, 2026 16:34

Perseus14 force-pushed the feat/wan-fast-inference-optimizations branch from e86f174 to 1d4c9f6 Compare September 25, 2026 17:20

Perseus14 force-pushed the feat/wan-custom-kernels branch from 05d2e33 to 00800f8 Compare September 25, 2026 18:01

Perseus14 force-pushed the feat/wan-fast-inference-optimizations branch 2 times, most recently from 0d91078 to 4703103 Compare September 25, 2026 18:28

Perseus14 force-pushed the feat/wan-custom-kernels branch from 00800f8 to 3ead80c Compare September 25, 2026 20:55

Perseus14 force-pushed the feat/wan-fast-inference-optimizations branch 3 times, most recently from 295c3b5 to 7a60c3d Compare September 25, 2026 22:13

Perseus14 force-pushed the feat/wan-custom-kernels branch 2 times, most recently from 25042c6 to 8061a2a Compare September 25, 2026 22:56

Perseus14 force-pushed the feat/wan-fast-inference-optimizations branch 2 times, most recently from 08d405c to 7db8313 Compare September 26, 2026 10:31

Perseus14 force-pushed the feat/wan-custom-kernels branch 2 times, most recently from 55a8c17 to 5c87fcf Compare September 26, 2026 13:13

Perseus14 force-pushed the feat/wan-fast-inference-optimizations branch from 7db8313 to 6fdc94e Compare September 26, 2026 13:13

This was referenced Sep 26, 2026

feat(wan): fast serving with persistent AOT caching and tuned inference recipe #479

Merged

feat(wan): fused RMSNorm+RoPE Pallas producer; Wan switches in config, not env #488

Merged

Perseus14 force-pushed the feat/wan-fast-inference-optimizations branch 2 times, most recently from ebaac14 to 58f889e Compare September 26, 2026 14:34

Perseus14 force-pushed the feat/wan-custom-kernels branch 2 times, most recently from a62691e to f3e5892 Compare September 26, 2026 17:07

Perseus14 force-pushed the feat/wan-fast-inference-optimizations branch from 58f889e to 68a580f Compare September 26, 2026 17:07

Perseus14 force-pushed the feat/wan-custom-kernels branch from f3e5892 to bfd8be6 Compare September 26, 2026 18:21

Perseus14 force-pushed the feat/wan-fast-inference-optimizations branch from 68a580f to 395f932 Compare September 26, 2026 18:21

Perseus14 force-pushed the feat/wan-fast-inference-optimizations branch 2 times, most recently from 08ccfb2 to 68c3732 Compare October 1, 2026 10:35

Perseus14 force-pushed the feat/wan-custom-kernels branch from ae05c1a to f2d1014 Compare October 1, 2026 10:35

Perseus14 force-pushed the feat/wan-custom-kernels branch from f2d1014 to faeae65 Compare October 2, 2026 08:14

Perseus14 force-pushed the feat/wan-fast-inference-optimizations branch 2 times, most recently from 0414334 to 34789ff Compare October 2, 2026 17:15

Perseus14 force-pushed the feat/wan-custom-kernels branch from faeae65 to 2cfb4af Compare October 2, 2026 17:15

eltsai previously approved these changes Oct 2, 2026

View reviewed changes

Perseus14 dismissed eltsai’s stale review via 87c36b2 October 5, 2026 03:57

Perseus14 force-pushed the feat/wan-fast-inference-optimizations branch from 34789ff to 87c36b2 Compare October 5, 2026 03:57

Perseus14 force-pushed the feat/wan-custom-kernels branch from 2cfb4af to e4a3859 Compare October 5, 2026 03:57

Perseus14 requested a review from eltsai October 5, 2026 16:32

eltsai previously approved these changes Oct 5, 2026

View reviewed changes

syhuang22 previously approved these changes Oct 5, 2026

View reviewed changes

Perseus14 dismissed stale reviews from syhuang22 and eltsai via 22b9467 October 13, 2026 07:07

Perseus14 force-pushed the feat/wan-custom-kernels branch from e4a3859 to 31a23b5 Compare October 13, 2026 07:07

mbohlool requested changes Oct 19, 2026

View reviewed changes

Perseus14 mentioned this pull request Oct 25, 2026

perf(wan): batch-sharded cross-attn KV cache and batch-aware VAE mesh #519

Open

prishajain1 previously approved these changes Oct 26, 2026

View reviewed changes

Perseus14 added 4 commits October 26, 2026 20:10

Navigation Menu

Repository navigation

perf(wan): pad sequence to 128-lane multiple for self-attention (wan_seq_pad="lane") - #491

perf(wan): pad sequence to 128-lane multiple for self-attention (wan_seq_pad="lane")#491
Perseus14 wants to merge 4 commits into
mainfrom
feat/wan-fast-inference-optimizations

Perseus14 commented Sep 25, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Sep 25, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Perseus14 commented Oct 2, 2026

Uh oh!

mbohlool left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

Perseus14 commented Sep 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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)

Uh oh!

github-actions Bot commented Sep 25, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Perseus14 commented Oct 2, 2026

Uh oh!

mbohlool left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Perseus14 commented Sep 25, 2026 •

edited

Loading

Performance (Wan 2.2 T2V-A14B, 720p / 81 frames / 40 steps, `CP=4, DP=2`)