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

Author: Perseus14

end_to_end/tpu/run_wan_fast_inference.sh

@@ -68,6 +68,8 @@
 #                    v6e seed 777, 17.92 dB / 0.6556 on v6e seed 12345, and
 #                    17.69 dB / 0.6918 on tpu7x seed 12345 (visually clean).
 #                    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
@@ -179,6 +181,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"
@@ -195,6 +198,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}';" \
@@ -213,6 +217,7 @@ case "$TPU_PROFILE" in
     DEFAULT_VAE_CHUNK=1
     DEFAULT_VAE_SPATIAL=8
     DEFAULT_FUSED_ROPE=false
+    DEFAULT_CROSS_ATTN_KERNEL=xla
     ;;
 esac
 
@@ -244,6 +249,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 devices remain. Defaults to CP=4 / DP=2 on an 8-device slice. On a
@@ -310,6 +316,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

@@ -45,5 +45,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

@@ -132,6 +132,10 @@ 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"
 # 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

@@ -129,6 +129,10 @@ 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"
 # 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

@@ -166,6 +166,10 @@ 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"
 # 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

@@ -111,6 +111,10 @@ 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"
 # 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

@@ -113,6 +113,10 @@ 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"
 # 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

@@ -113,6 +113,10 @@ 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"
 # 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.