feat(wan): token-sharded patch embedding (wan_patch_embed_mode="tokens")

Second of four pieces split out of the Wan fast-inference optimizations PR.
The switch defaults off ("conv") in the Wan YAMLs and in wan_runtime_options;
it is a plain config key plumbed to WanModel through
`attention_config_entries(config)` and resolved at build time
(`_wan_opt`) so it lives on the GraphDef.

- With `"tokens"`, the Conv3d patch embedding becomes a patchify + matmul on
  the token-sharded layout (`embed_patches_as_tokens`): the non-overlapping
  patches (stride == kernel) are flattened in the kernel's
  `(p_t, p_h, p_w, C)` order, sharded by token and multiplied by the
  replicated `[p_t*p_h*p_w*C, D]` kernel (with kernel and bias cast to
  `hidden_states.dtype` to match `nnx.Conv.promote_dtype`). Same result as
  `collapse(conv(x), 1, -1)` up to accumulation order, but produced directly
  in the layout the transformer blocks use, so XLA no longer emits the
  activation all-gather / collective-permute when collapsing `(F, H, W)`
  into the context-sharded token axis (12.8 ms per step on tpu7x, exposed).
- `can_embed_patches_as_tokens` checks the conv really is a plain
  stride-equals-kernel 3D patchify (no groups, masks, dilation, and `'SAME'`
  pads nothing). When `wan_patch_embed_mode="tokens"` is requested on an
  indivisible latent grid, `WanModel.__call__` raises `ValueError` instead of
  silently falling back to `nnx.Conv`.
- `run_wan_fast_inference.sh`: every profile sets `WAN_PATCH_EMBED_MODE=tokens`
  (overridable).

Tests
- `wan_patch_embed_test.py` (8): tokens vs conv parity (function level and
  through `WanModel`), the eligibility predicate, fail-closed `ValueError` on
  indivisible latent grids, and the config-driven mode. All 8 tests run
  unconditionally in GitHub Actions.
- `wan_runtime_options_test.py` (12): the new switch's YAML default and enum
  validation; `wan_patch_embed_mode` is excluded from the `FlaxWanAttention`
  attribute check since `WanModel` consumes it.
- `end_to_end/tpu/run_wan_stack_tests.sh`: adds `wan_patch_embed_test.py`.

Author: Perseus14

end_to_end/tpu/run_wan_fast_inference.sh

@@ -68,7 +68,7 @@
 #                    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
+#   WAN_CROSS_ATTN_KERNEL / WAN_PATCH_EMBED_MODE
 #                    override Wan fast-inference kernel/layout switches
 #   DP / CP / PER_DEVICE_BATCH / SEED
 #                    override mesh parallelism, per-device batch, or RNG seed (default 12345)
@@ -182,6 +182,7 @@ case "$TPU_PROFILE" in
     DEFAULT_VAE_SPATIAL=8
     DEFAULT_FUSED_ROPE=true
     DEFAULT_CROSS_ATTN_KERNEL=pallas
+    DEFAULT_PATCH_EMBED_MODE=tokens
     ;;
   v7)
     PLATFORM_LIBTPU="$V7_LIBTPU"
@@ -199,6 +200,7 @@ case "$TPU_PROFILE" in
     DEFAULT_VAE_SPATIAL=8
     DEFAULT_FUSED_ROPE=true
     DEFAULT_CROSS_ATTN_KERNEL=pallas
+    DEFAULT_PATCH_EMBED_MODE=tokens
     ;;
   *)
     echo "== warning: unrecognised or non-v6e/v7 accelerator '${ACCEL_TYPE:-unknown}';" \
@@ -218,6 +220,7 @@ case "$TPU_PROFILE" in
     DEFAULT_VAE_SPATIAL=8
     DEFAULT_FUSED_ROPE=false
     DEFAULT_CROSS_ATTN_KERNEL=xla
+    DEFAULT_PATCH_EMBED_MODE=tokens
     ;;
 esac
 
@@ -250,6 +253,7 @@ 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}
+WAN_PATCH_EMBED_MODE=${WAN_PATCH_EMBED_MODE:-$DEFAULT_PATCH_EMBED_MODE}
 
 # 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
@@ -317,6 +321,7 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \
   use_fused_rope_kernel="$USE_FUSED_ROPE_KERNEL" \
   ${FUSED_ROPE_ARGS[@]+"${FUSED_ROPE_ARGS[@]}"} \
   wan_cross_attn_kernel="$WAN_CROSS_ATTN_KERNEL" \
+  wan_patch_embed_mode="$WAN_PATCH_EMBED_MODE" \
   fps=16 ${GUIDANCE_ARGS[@]+"${GUIDANCE_ARGS[@]}"} \
   seed="${SEED:-12345}" \
   flash_block_sizes="$FLASH_BLOCK_SIZES" \

end_to_end/tpu/run_wan_stack_tests.sh

@@ -47,5 +47,7 @@ TESTS=(
   "$T/wan_runtime_options_test.py"
   # Short-KV cross-attention Pallas kernel (feat/wan-cross-attn-pallas)
   "$T/cross_attention_pallas_test.py"
+  # Token-sharded patch embedding (feat/wan-token-patch-embed)
+  "$T/wan_patch_embed_test.py"
 )
 PYTHONPATH="src${PYTHONPATH:+:$PYTHONPATH}" exec "${PYTHON:-python3}" -m pytest -q -rs "${TESTS[@]}" "$@"

src/maxdiffusion/configs/base_wan_14b.yml

@@ -136,6 +136,10 @@ wan_rope_accum: "auto"
 # (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"
+# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by
+# token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the
+# replicated kernel, no activation collective; equal up to accumulation order).
+wan_patch_embed_mode: "conv"
 # 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

@@ -133,6 +133,10 @@ wan_rope_accum: "auto"
 # (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"
+# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by
+# token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the
+# replicated kernel, no activation collective; equal up to accumulation order).
+wan_patch_embed_mode: "conv"
 # 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

@@ -170,6 +170,10 @@ wan_rope_accum: "auto"
 # (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"
+# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by
+# token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the
+# replicated kernel, no activation collective; equal up to accumulation order).
+wan_patch_embed_mode: "conv"
 # 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

@@ -115,6 +115,10 @@ wan_rope_accum: "auto"
 # (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"
+# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by
+# token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the
+# replicated kernel, no activation collective; equal up to accumulation order).
+wan_patch_embed_mode: "conv"
 # 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

@@ -117,6 +117,10 @@ wan_rope_accum: "auto"
 # (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"
+# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by
+# token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the
+# replicated kernel, no activation collective; equal up to accumulation order).
+wan_patch_embed_mode: "conv"
 # 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

@@ -117,6 +117,10 @@ wan_rope_accum: "auto"
 # (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"
+# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by
+# token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the
+# replicated kernel, no activation collective; equal up to accumulation order).
+wan_patch_embed_mode: "conv"
 # 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/models/wan/transformers/transformer_wan.py

@@ -61,6 +61,62 @@ def get_frequencies(max_seq_len: int, theta: int, attention_head_dim: int):
   return freqs_split
 
 
+PATCH_EMBED_MODES = ("conv", "tokens")
+
+
+def can_embed_patches_as_tokens(conv: nnx.Conv, hidden_shape: Tuple[int, ...], patch_size: Tuple[int, int, int]) -> bool:
+  """Whether `embed_patches_as_tokens` computes exactly what `conv` does."""
+
+  def per_dim(value):
+    value = 1 if value is None else value
+    return tuple(value) if isinstance(value, (tuple, list)) else (value,) * len(patch_size)
+
+  if len(hidden_shape) != 5 or per_dim(conv.kernel_size) != tuple(patch_size):
+    return False
+  if per_dim(conv.strides) != tuple(patch_size) or conv.feature_group_count != 1 or conv.mask is not None:
+    return False
+  if per_dim(conv.input_dilation) != (1,) * 3 or per_dim(conv.kernel_dilation) != (1,) * 3:
+    return False
+  # With every spatial size divisible by its stride, 'SAME' pads nothing.
+  return conv.padding in ("SAME", "VALID") and all(n % p == 0 for n, p in zip(hidden_shape[1:4], patch_size))
+
+
+def embed_patches_as_tokens(conv: nnx.Conv, hidden_states: jax.Array, patch_size: Tuple[int, int, int]) -> jax.Array:
+  """Wan patch embedding as a token-sharded matmul: `[B, F, H, W, C]` -> `[B, N, D]`.
+
+  Same result as `jax.lax.collapse(conv(hidden_states), 1, -1)` up to
+  accumulation order, but produced in the layout the transformer blocks use.
+  When the 3D `nnx.Conv` runs before `collapse`, XLA either shards the conv output
+  differently from `activation_length` or emits an activation all-gather /
+  collective-permute when collapsing `(F, H, W)` into the context-sharded token
+  axis (12.8 ms per step on tpu7x, exposed). Patches do not overlap (stride ==
+  kernel), so each token needs only its own patch: flatten the patches in the
+  kernel's `(p_t, p_h, p_w, C)` order, shard them by token and multiply by the
+  replicated `[p_t*p_h*p_w*C, D]` kernel.
+  """
+  p_t, p_h, p_w = patch_size
+  batch, frames, height, width, channels = hidden_states.shape
+  patches = hidden_states.reshape(batch, frames // p_t, p_t, height // p_h, p_h, width // p_w, p_w, channels)
+  patches = jnp.transpose(patches, (0, 1, 3, 5, 2, 4, 6, 7))
+  patches = patches.reshape(batch, -1, p_t * p_h * p_w * channels)
+  patches = nn.with_logical_constraint(patches, ("activation_batch", "activation_length", None))
+
+  kernel = conv.kernel[...].reshape(p_t * p_h * p_w * channels, -1)
+  kernel = nn.with_logical_constraint(kernel, (None, None))
+  bias = conv.bias[...] if conv.bias is not None else None
+  # Mirror nnx.Conv: promote everything to `conv.dtype`, no preferred type.
+  dtype = conv.dtype or jnp.result_type(patches, kernel)
+  out = jax.lax.dot_general(
+      patches.astype(dtype),
+      kernel.astype(dtype),
+      (((2,), (0,)), ((), ())),
+      precision=conv.precision,
+  )
+  if bias is not None:
+    out = out + bias.astype(dtype)
+  return out
+
+
 class WanRotaryPosEmbed(nnx.Module):
 
   def __init__(
@@ -609,6 +665,13 @@ def __init__(
         **(attention_config or {}),
     }
 
+    # Resolved at build time so they live on the GraphDef; None -> built-in default.
+    def _wan_opt(name):
+      value = attention_config.get(name)
+      return wan_runtime_options.default(name) if value is None else wan_runtime_options.coerce(name, value)
+
+    self.patch_embed_mode = _wan_opt("wan_patch_embed_mode")
+
     # 1. Patch & position embedding
     self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len)
     self.patch_embedding = nnx.Conv(
@@ -810,8 +873,19 @@ def __call__(
       if rotary_emb is None:
         rotary_emb = self.rope(hidden_states)
     with self.conditional_named_scope("patch_embedding"):
-      hidden_states = self.patch_embedding(hidden_states)
-      hidden_states = jax.lax.collapse(hidden_states, 1, -1)
+      patch_embed_mode = self.patch_embed_mode
+      if patch_embed_mode not in PATCH_EMBED_MODES:
+        raise ValueError(f"wan_patch_embed_mode must be one of {PATCH_EMBED_MODES}, got {patch_embed_mode!r}.")
+      if patch_embed_mode == "tokens":
+        if not can_embed_patches_as_tokens(self.patch_embedding, hidden_states.shape, self.config.patch_size):
+          raise ValueError(
+              f"wan_patch_embed_mode='tokens' requires the latent grid {hidden_states.shape[1:4]} to be divisible by "
+              f"patch_size {tuple(self.config.patch_size)} (height/width multiples of {8 * p_h}/{8 * p_w} pixels)."
+          )
+        hidden_states = embed_patches_as_tokens(self.patch_embedding, hidden_states, self.config.patch_size)
+      else:
+        hidden_states = self.patch_embedding(hidden_states)
+        hidden_states = jax.lax.collapse(hidden_states, 1, -1)
     per_token_t = timestep.ndim == 2  # [B, seq_len] for TI2V
     with self.conditional_named_scope("condition_embedder"):
       if per_token_t: