Repository navigation
Conversation
![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.
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.
4a9cb02 to
1f5539b
Compare
mbohlool
left a comment
mbohlool
left a comment
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.
| if local_q_len <= cross_attention_pallas.DEFAULT_BLOCK_Q and local_q_len % 8 != 0: | ||
| reason = f"local query length {local_q_len} <= {cross_attention_pallas.DEFAULT_BLOCK_Q} is not a multiple of 8" | ||
| if reason is not None: | ||
| _warn_once("pallas_cross_attn_unusable", f"Pallas cross-attention requested but unusable: {reason}; using XLA.") |
There was a problem hiding this comment.
Could a requested wan_cross_attn_kernel: "pallas" raise instead of logging once and running XLA? The static reasons (non-TPU mesh without interpret, TPU kind outside v6e/tpu7x, dim_head, kv_heads, an attention kernel that doesn't route to dot-product) are all known in __init__, so a ValueError there; the shape-dependent ones (KV length, local Q length, mesh divisibility, VMEM estimate) can raise at trace time. The test_falls_back_to_xla_* tests would then become assertRaises.
| @_SKIP_IN_GITHUB_ACTIONS | ||
| class CrossAttentionPallasParityTest(parameterized.TestCase): |
There was a problem hiding this comment.
Could the parity class run in CI? It's marked @_SKIP_IN_GITHUB_ACTIONS, and the only kernel test CI could run, test_pallas_mode_cpu_interpret_integration, is skipped by _on_unvalidated_tpu() on the v4-8 runner. So no GitHub job executes the kernel at all. The shapes here are tiny (<= 512x256), so passing interpret=True unconditionally in this class would cover the softmax, ragged-tile and prescaled-K math on the v4 runner and on CPU.
| # keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical). | ||
| wan_cross_attn_kernel: "xla" | ||
| # When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only). | ||
| wan_cross_attn_cpu_interpret: False |
There was a problem hiding this comment.
Could wan_cross_attn_cpu_interpret stay out of the six Wan YAMLs, _OPTIONS and ATTENTION_OPTIONS? It's a tests-only switch: WanCrossAttentionDispatchTest already passes it through attention_config, and the I2V test sets the attribute directly, so production configs don't need to carry it.
| ): | ||
| ones = jnp.ones((v_ref.shape[1], dim_head), dtype=v_ref.dtype) if sum_mode == "mxu" else None | ||
| row_valid = None | ||
| if max_mode == "tile" and seq_len % block_q != 0: |
There was a problem hiding this comment.
Could max_mode="tile" and sum_mode="vpu" move out of the shipped kernel into a benchmark script? The docstring says tile mode exists to measure the per-row reduction, and nothing outside this file and its tests passes either argument. The model only ever runs row/mxu. Dropping them also removes the row_valid stale-row masking and leaves one path to test and maintain.
| try: | ||
| devices = list(mesh.devices.flat) if mesh is not None else jax.devices() | ||
| except RuntimeError: | ||
| devices = [] | ||
| return [getattr(d, "device_kind", "").lower() for d in devices if getattr(d, "platform", "") == "tpu"] |
There was a problem hiding this comment.
nit: jax.Device always has platform and device_kind, and NNXAttentionOp always sets kv_heads, so could these be plain attribute reads without the except RuntimeError? As written, a backend init failure turns into "no TPUs" and picks the 64 MiB default instead of failing. Same pattern in vmem_budget_is_validated and in _pallas_cross_attention (getattr(d, "platform", None), getattr(op, "kv_heads", None)).
| with jax.named_scope("apply_attention"): | ||
| attn_output = run_attention() | ||
| attn_output = None | ||
| if not is_self_attention and encoder_attention_mask is None and not qk_prescaled: |
There was a problem hiding this comment.
nit: in this branch encoder_attention_mask is always None (reset at L4216-4217 for non-I2V), and qk_prescaled is only set for self-attention. So the guard reduces to not is_self_attention. Could the two dead conditions go, so the T2V and I2V sites call _pallas_cross_attention the same way?
2 participants
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