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
…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
(full-row max, exp2, one numerator/denominator, 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).
- `FlaxWanAttention._pallas_cross_attention` falls back to XLA (logging the
reason once) when the inputs or sharding don't fit, and also when the VMEM
budget is not validated or too small: 64 MiB is requested only when every
mesh TPU is v6e or tpu7x (otherwise Mosaic's default), and
`estimate_vmem_bytes` (Q/out/K/V double buffers + one head's FP32 scores) is
checked against it (~42 MiB at the Wan shard shape; MAX_KV_LEN KV does not
fit and falls back rather than failing to compile). The estimate has only
been checked against real compiles at Wan's KV length (512). The kernel
itself raises a ValueError for an over-budget working set only when a budget
resolves (an explicit vmem_limit_bytes, or the v6e/tpu7x default); under
Mosaic's default it does no check of its own, and callers that can fall back
use `vmem_unusable_reason`. The two fallback reasons (unusable inputs, VMEM)
log under separate warn-once keys; the "ACTIVE" log is keyed on the query
shape, so T2V and the I2V text branch are each reported. The sharded wrapper is cached
(`_build_sharded_pallas_cross_attention`) and wrapped with
`with_xla_backward` since `pallas_call` has no transpose rule.
- Dispatch: T2V cross-attention, and the text half of I2V cross-attention
(both the image-token split and the text-only path), try the kernel first
and use `attention_op.apply_attention` when it returns None.
- `wan_cross_attn_cpu_interpret` (default False, tests only) lets the kernel
run under Pallas interpret mode on CPU.
- 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`, `ulysses_attention_chunks`
and `use_k_centering` from the config to the benchmarked block (shared
benchmark-fidelity fix carried in this first piece; note it defaults
`split_head_dim` to True when the config lacks the key, whereas the Wan
ymls set False).
Tests
- cross_attention_pallas_test.py (32 test methods, 40 parameterized cases):
golden parity (row/tile max, mxu/vpu
sum, ragged tiles, k_prescaled, batch), guards, VMEM budget resolution /
working-set estimate / fallback reasons, and FlaxWanAttention dispatch (CPU
interpret integration, fallbacks, I2V text-vs-image split). The three tests
that expect the kernel to be used skip on TPUs outside the validated kinds
(there FlaxWanAttention falls back by design); that fallback is covered by
test_falls_back_to_xla_on_unvalidated_tpu_kind.
- wan_runtime_options_test.py (12): the new switch's YAML default and enum
validation (test_enum_options_validation).
- CI: kernel parity and the TPU Pallas-active dispatch
tests skip in GitHub Actions (19 of the 40 cases); guards, VMEM budget logic, CPU interpret integration and the
XLA fallback dispatch tests run in CI.
- end_to_end/tpu/run_wan_stack_tests.sh: adds cross_attention_pallas_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:
322 tests, all OK, 1 skipped (cross_attention_pallas_test: off-TPU fallback only);
cross_attention_pallas_test.py 40 cases (1 skipped), wan_runtime_options_test.py 12.
Perseus14
force-pushed
the
feat/wan-cross-attn-pallas
branch
from
October 15, 2026 19:20
4a9cb02 to
1f5539b
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 #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