Repository navigation
Conversation
This was referenced Oct 15, 2026
Perseus14
added this pull request to stack #512
October 15, 2026 16:00
![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 fused short-KV cross-attention Pallas kernel for TPU, designed to replace the XLA dot-product fallback for short KV sequences in Wan models. It integrates this kernel into FlaxWanAttention, updates configuration files and runtime options to support the new wan_cross_attn_kernel parameter, and adds corresponding unit tests. The review feedback suggests a minor optimization in the Pallas kernel to compute the tile-wide maximum in a single jnp.max call instead of two nested calls.
Perseus14
force-pushed
the
feat/wan-cross-attn-pallas
branch
from
October 15, 2026 17:53
f78bfda to
4a9cb02
Compare
Perseus14
force-pushed
the
feat/wan-cross-attn-pallas
branch
from
October 15, 2026 19:20
4a9cb02 to
1f5539b
Compare
mbohlool
requested changes
Oct 19, 2026
mbohlool
left a comment
mbohlool
left a comment
Collaborator
There was a problem hiding this comment.
The main ask is that a requested pallas raises instead of falling back to XLA, and that CI runs the kernel at least once in interpret mode.
Perseus14
force-pushed
the
feat/wan-cross-attn-pallas
branch
from
October 20, 2026 12:51
1f5539b to
da317cc
Compare
Perseus14
force-pushed
the
feat/wan-cross-attn-pallas
branch
3 times, most recently
from
October 20, 2026 15:13
cdb5a4d to
56acb00
Compare
Perseus14
force-pushed
the
feat/wan-cross-attn-pallas
branch
2 times, most recently
from
October 21, 2026 08:48
fc80478 to
ed91d0d
Compare
Perseus14
force-pushed
the
feat/wan-cross-attn-pallas
branch
from
October 21, 2026 08:55
ed91d0d to
76b6e85
Compare
Perseus14
force-pushed
the
feat/wan-cross-attn-pallas
branch
from
October 21, 2026 09:32
76b6e85 to
b7d4ab7
Compare
Perseus14
force-pushed
the
feat/wan-cross-attn-pallas
branch
from
October 22, 2026 08:54
b7d4ab7 to
c086c74
Compare
…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`.
Perseus14
force-pushed
the
feat/wan-cross-attn-pallas
branch
from
October 22, 2026 22:30
c086c74 to
d0a909f
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 #488 (
feat/wan-custom-kernels). First of four pieces split out of #491. Addswan_cross_attn_kernel: "pallas"(default"xla"in the Wan YAMLs /wan_runtime_options):kernels/cross_attention_pallas.pykeeps a head block's whole 512-token text KV resident in VMEM and computes an exact single-pass softmax per query tile (full-row max, exp2, one numerator/denominator, normalised after PV — no online rescaling). It replaces only what would otherwise be the unmasked XLA dot-product fallback.FlaxWanAttention._pallas_cross_attentionfalls back to XLA (logging the reason once) when the inputs/sharding don't fit, when the mesh is not all v6e/tpu7x (the only generations the 64 MiB VMEM budget was validated on), or whenestimate_vmem_bytesexceeds the budget. T2V cross-attention and the text half of I2V cross-attention dispatch to it. In isolation the kernel is 0.51 ms vs 3.94 ms (v6e) and 0.49 ms vs 1.81 ms (tpu7x) per call vs the XLA path inside the model.run_wan_fast_inference.shturns it on for the v6e/v7 profiles.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.4 s denoise (cold AOT: 130.4 s generate / 128.3 s denoise), down from 133.7 s generate / 131.6 s denoise on #488 (−3.1 s generate / −3.2 s denoise end-to-end from Pallas cross-attention alone).Tests:
cross_attention_pallas_test.py(32 test methods / 40 parameterized cases, new; parity grid including the Wan production shard shapes 18900/18944×512 on TPU, guards, VMEM budget logic,FlaxWanAttentiondispatch incl. CPU-interpret integration),wan_runtime_options_test.py(+enum validation).Stack: #477 → #507 → #478 → #508 → #479 → #488 → #509 → #510 → #511 → #491