feat(wan): short-KV text cross-attention Pallas kernel (wan_cross_attn_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.

Author: Perseus14

end_to_end/tpu/run_wan_fast_inference.sh

@@ -63,6 +63,8 @@
 #                    against the XLA-path video, visually clean); on tpu7x the
 #                    output was bit-identical in that measurement.
 #                    USE_FUSED_ROPE_KERNEL=false restores the XLA producer.
+#   WAN_CROSS_ATTN_KERNEL
+#                    override Wan fast-inference kernel/layout switches
 #   DP / CP / PER_DEVICE_BATCH / SEED
 #                    override mesh parallelism, per-device batch, or RNG seed (default 12345)
 set -euo pipefail
@@ -167,6 +169,7 @@ case "$TPU_PROFILE" in
     DEFAULT_VAE_CHUNK=1
     DEFAULT_VAE_SPATIAL=8
     DEFAULT_FUSED_ROPE=true
+    DEFAULT_CROSS_ATTN_KERNEL=pallas
     ;;
   v7)
     PLATFORM_LIBTPU="$V7_LIBTPU"
@@ -183,6 +186,7 @@ case "$TPU_PROFILE" in
     DEFAULT_VAE_CHUNK=1
     DEFAULT_VAE_SPATIAL=8
     DEFAULT_FUSED_ROPE=true
+    DEFAULT_CROSS_ATTN_KERNEL=pallas
     ;;
   *)
     echo "== warning: unrecognised or non-v6e/v7 accelerator '${ACCEL_TYPE:-unknown}';" \
@@ -201,6 +205,7 @@ case "$TPU_PROFILE" in
     DEFAULT_VAE_CHUNK=1
     DEFAULT_VAE_SPATIAL=8
     DEFAULT_FUSED_ROPE=false
+    DEFAULT_CROSS_ATTN_KERNEL=xla
     ;;
 esac
 
@@ -232,6 +237,7 @@ USE_FUSED_ROPE_KERNEL=${USE_FUSED_ROPE_KERNEL:-$DEFAULT_FUSED_ROPE}
 FUSED_ROPE_ARGS=()
 [ -n "${FUSED_ROPE_BLOCK_S:-}" ] && FUSED_ROPE_ARGS+=("fused_rope_block_s=$FUSED_ROPE_BLOCK_S")
 [ -n "${FUSED_ROPE_HEAD_BLOCK:-}" ] && FUSED_ROPE_ARGS+=("fused_rope_head_block=$FUSED_ROPE_HEAD_BLOCK")
+WAN_CROSS_ATTN_KERNEL=${WAN_CROSS_ATTN_KERNEL:-$DEFAULT_CROSS_ATTN_KERNEL}
 
 # Mesh: context parallelism carries the Ulysses shards, data parallelism takes
 # whatever chips remain. Defaults to CP=4 / DP=2 on an 8-chip slice. On a slice
@@ -294,6 +300,7 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \
   use_kv_cache=true use_base2_exp=true use_experimental_scheduler=true \
   use_fused_rope_kernel="$USE_FUSED_ROPE_KERNEL" \
   ${FUSED_ROPE_ARGS[@]+"${FUSED_ROPE_ARGS[@]}"} \
+  wan_cross_attn_kernel="$WAN_CROSS_ATTN_KERNEL" \
   fps=16 ${GUIDANCE_ARGS[@]+"${GUIDANCE_ARGS[@]}"} \
   seed="${SEED:-12345}" \
   flash_block_sizes="$FLASH_BLOCK_SIZES" \

end_to_end/tpu/run_wan_stack_tests.sh

@@ -44,5 +44,7 @@ TESTS=(
   # Fused RMSNorm+RoPE producer, Wan switches in config (feat/wan-custom-kernels)
   "$T/fused_rmsnorm_rope_pallas_test.py"
   "$T/wan_runtime_options_test.py"
+  # Short-KV cross-attention Pallas kernel (feat/wan-cross-attn-pallas)
+  "$T/cross_attention_pallas_test.py"
 )
 PYTHONPATH="src${PYTHONPATH:+:$PYTHONPATH}" exec "${PYTHON:-python3}" -m pytest -q -rs "${TESTS[@]}" "$@"

src/maxdiffusion/configs/base_wan_14b.yml

@@ -133,6 +133,12 @@ wan_cfg_before_unpatchify: True
 wan_cross_attn_prescale_kv: False
 # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
 wan_rope_accum: "auto"
+# Text cross-attention backend where it would otherwise fall back to XLA dot-product
+# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
+# 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
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.

src/maxdiffusion/configs/base_wan_1_3b.yml

@@ -130,6 +130,12 @@ wan_cfg_before_unpatchify: True
 wan_cross_attn_prescale_kv: False
 # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
 wan_rope_accum: "auto"
+# Text cross-attention backend where it would otherwise fall back to XLA dot-product
+# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
+# 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
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.

src/maxdiffusion/configs/base_wan_27b.yml

@@ -162,6 +162,12 @@ wan_cfg_before_unpatchify: True
 wan_cross_attn_prescale_kv: False
 # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
 wan_rope_accum: "auto"
+# Text cross-attention backend where it would otherwise fall back to XLA dot-product
+# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
+# 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
 # For attention=ulysses_ring*, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.

src/maxdiffusion/configs/base_wan_animate.yml

@@ -112,6 +112,12 @@ wan_cfg_before_unpatchify: True
 wan_cross_attn_prescale_kv: False
 # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
 wan_rope_accum: "auto"
+# Text cross-attention backend where it would otherwise fall back to XLA dot-product
+# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
+# 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
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -114,6 +114,12 @@ wan_cfg_before_unpatchify: True
 wan_cross_attn_prescale_kv: False
 # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
 wan_rope_accum: "auto"
+# Text cross-attention backend where it would otherwise fall back to XLA dot-product
+# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
+# 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
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -114,6 +114,12 @@ wan_cfg_before_unpatchify: True
 wan_cross_attn_prescale_kv: False
 # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
 wan_rope_accum: "auto"
+# Text cross-attention backend where it would otherwise fall back to XLA dot-product
+# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
+# 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
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.