diff --git a/end_to_end/tpu/run_wan_fast_inference.sh b/end_to_end/tpu/run_wan_fast_inference.sh
index 32bdc7c65..392c1794d 100755
--- a/end_to_end/tpu/run_wan_fast_inference.sh
+++ b/end_to_end/tpu/run_wan_fast_inference.sh
@@ -63,7 +63,7 @@
 #                    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
+#   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)
@@ -170,6 +170,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"
@@ -187,6 +188,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}';" \
@@ -206,6 +208,7 @@ case "$TPU_PROFILE" in
     DEFAULT_VAE_SPATIAL=8
     DEFAULT_FUSED_ROPE=false
     DEFAULT_CROSS_ATTN_KERNEL=xla
+    DEFAULT_PATCH_EMBED_MODE=tokens
     ;;
 esac
 
@@ -238,6 +241,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 chips remain. Defaults to CP=4 / DP=2 on an 8-chip slice. On a slice
@@ -301,6 +305,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" \
diff --git a/end_to_end/tpu/run_wan_stack_tests.sh b/end_to_end/tpu/run_wan_stack_tests.sh
index 409a536ea..e1eab0758 100755
--- a/end_to_end/tpu/run_wan_stack_tests.sh
+++ b/end_to_end/tpu/run_wan_stack_tests.sh
@@ -46,5 +46,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[@]}" "$@"
diff --git a/src/maxdiffusion/configs/base_wan_14b.yml b/src/maxdiffusion/configs/base_wan_14b.yml
index 767936b82..d28697686 100644
--- a/src/maxdiffusion/configs/base_wan_14b.yml
+++ b/src/maxdiffusion/configs/base_wan_14b.yml
@@ -139,6 +139,10 @@ wan_rope_accum: "auto"
 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
+# 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.
diff --git a/src/maxdiffusion/configs/base_wan_1_3b.yml b/src/maxdiffusion/configs/base_wan_1_3b.yml
index 3c36364a6..db8215331 100644
--- a/src/maxdiffusion/configs/base_wan_1_3b.yml
+++ b/src/maxdiffusion/configs/base_wan_1_3b.yml
@@ -136,6 +136,10 @@ wan_rope_accum: "auto"
 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
+# 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.
diff --git a/src/maxdiffusion/configs/base_wan_27b.yml b/src/maxdiffusion/configs/base_wan_27b.yml
index 45e9c6725..89020ce4f 100644
--- a/src/maxdiffusion/configs/base_wan_27b.yml
+++ b/src/maxdiffusion/configs/base_wan_27b.yml
@@ -168,6 +168,10 @@ wan_rope_accum: "auto"
 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
+# 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.
diff --git a/src/maxdiffusion/configs/base_wan_animate.yml b/src/maxdiffusion/configs/base_wan_animate.yml
index 28784c8a5..ae8a93c42 100644
--- a/src/maxdiffusion/configs/base_wan_animate.yml
+++ b/src/maxdiffusion/configs/base_wan_animate.yml
@@ -118,6 +118,10 @@ wan_rope_accum: "auto"
 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
+# 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.
diff --git a/src/maxdiffusion/configs/base_wan_i2v_14b.yml b/src/maxdiffusion/configs/base_wan_i2v_14b.yml
index 66650a0aa..7221ca5a7 100644
--- a/src/maxdiffusion/configs/base_wan_i2v_14b.yml
+++ b/src/maxdiffusion/configs/base_wan_i2v_14b.yml
@@ -120,6 +120,10 @@ wan_rope_accum: "auto"
 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
+# 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.
diff --git a/src/maxdiffusion/configs/base_wan_i2v_27b.yml b/src/maxdiffusion/configs/base_wan_i2v_27b.yml
index dfc2d7e68..c3af910c8 100644
--- a/src/maxdiffusion/configs/base_wan_i2v_27b.yml
+++ b/src/maxdiffusion/configs/base_wan_i2v_27b.yml
@@ -120,6 +120,10 @@ wan_rope_accum: "auto"
 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
+# 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.
diff --git a/src/maxdiffusion/models/wan/transformers/transformer_wan.py b/src/maxdiffusion/models/wan/transformers/transformer_wan.py
index db083ae0c..51132ec1e 100644
--- a/src/maxdiffusion/models/wan/transformers/transformer_wan.py
+++ b/src/maxdiffusion/models/wan/transformers/transformer_wan.py
@@ -35,7 +35,7 @@
     NNXPixArtAlphaTextProjection,
 )
 from ...normalization_flax import FP32LayerNorm
-from ...attention_flax import FlaxWanAttention
+from ...attention_flax import FlaxWanAttention, _warn_once
 from ...gradient_checkpoint import GradientCheckpointType
 
 BlockSizes = common_types.BlockSizes
@@ -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,22 @@ 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" and can_embed_patches_as_tokens(
+          self.patch_embedding, hidden_states.shape, self.config.patch_size
+      ):
+        hidden_states = embed_patches_as_tokens(self.patch_embedding, hidden_states, self.config.patch_size)
+      else:
+        if patch_embed_mode == "tokens":
+          _warn_once(
+              "wan_patch_embed_tokens_unusable",
+              "wan_patch_embed_mode='tokens' requested, but patch_embedding or hidden_states.shape is ineligible "
+              "for token-sharded patchify; falling back to conv.",
+          )
+        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:
diff --git a/src/maxdiffusion/tests/wan_patch_embed_test.py b/src/maxdiffusion/tests/wan_patch_embed_test.py
new file mode 100644
index 000000000..c64029400
--- /dev/null
+++ b/src/maxdiffusion/tests/wan_patch_embed_test.py
@@ -0,0 +1,184 @@
+"""
+Copyright 2026 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+"""Tests for the token-sharded Wan patch embedding (`wan_patch_embed_mode`)."""
+
+import os
+import unittest
+from unittest import mock
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from flax import nnx
+from flax.linen import partitioning as nn_partitioning
+from jax.sharding import Mesh
+
+from maxdiffusion import pyconfig
+from maxdiffusion.max_utils import create_device_mesh
+from maxdiffusion.models.wan.transformers import transformer_wan
+from maxdiffusion.models.wan.transformers.transformer_wan import (
+    WanModel,
+    can_embed_patches_as_tokens,
+    embed_patches_as_tokens,
+)
+
+IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true"
+# Kernel numerics grids and multi-device tests are skipped in CI; see
+# end_to_end/tpu/run_wan_stack_tests.sh.
+_SKIP_IN_GITHUB_ACTIONS = unittest.skipIf(
+    IN_GITHUB_ACTIONS, "TPU kernel / multi-device test, skipped in GitHub Actions; run end_to_end/tpu/run_wan_stack_tests.sh"
+)
+
+THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+def _conv(patch_size, in_channels=4, out_channels=256, dtype=jnp.float32, precision=jax.lax.Precision.HIGHEST, **kw):
+  return nnx.Conv(
+      in_channels,
+      out_channels,
+      kernel_size=patch_size,
+      strides=patch_size,
+      dtype=dtype,
+      param_dtype=dtype,
+      precision=precision,
+      rngs=nnx.Rngs(0),
+      bias_init=nnx.initializers.normal(0.5),
+      **kw,
+  )
+
+
+class PatchEmbedTokensTest(unittest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    pyconfig.initialize([None, os.path.join(THIS_DIR, "..", "configs", "base_wan_14b.yml")], unittest=True)
+    self.config = pyconfig.config
+    self.mesh = Mesh(create_device_mesh(self.config), self.config.mesh_axes)
+
+  def _both(self, conv, x, patch_size):
+    def conv_path(x):
+      return jax.lax.collapse(conv(x), 1, -1)
+
+    def token_path(x):
+      return embed_patches_as_tokens(conv, x, patch_size)
+
+    with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      return jax.jit(conv_path)(x), jax.jit(token_path)(x)
+
+  def test_matches_conv(self):
+    for patch_size in ((1, 2, 2), (2, 2, 2)):
+      with self.subTest(patch_size=patch_size):
+        conv = _conv(patch_size)
+        x = jax.random.normal(jax.random.key(1), (2, 4, 8, 16, 4), jnp.float32)
+        want, got = self._both(conv, x, patch_size)
+        self.assertEqual(got.shape, want.shape)
+        np.testing.assert_allclose(np.asarray(got), np.asarray(want), rtol=1e-5, atol=1e-5)
+
+  def test_matches_conv_in_bf16(self):
+    patch_size = (1, 2, 2)
+    conv = _conv(patch_size, in_channels=16, dtype=jnp.bfloat16, precision=None)
+    x = jax.random.normal(jax.random.key(2), (2, 3, 8, 16, 16), jnp.float32).astype(jnp.bfloat16)
+    want, got = self._both(conv, x, patch_size)
+    self.assertEqual(got.dtype, want.dtype)
+    np.testing.assert_allclose(np.asarray(got, np.float32), np.asarray(want, np.float32), rtol=2e-2, atol=2e-2)
+
+  def test_guard(self):
+    conv = _conv((1, 2, 2))
+    self.assertTrue(can_embed_patches_as_tokens(conv, (1, 21, 90, 160, 4), (1, 2, 2)))
+    # A spatial size not divisible by the patch makes 'SAME' pad.
+    self.assertFalse(can_embed_patches_as_tokens(conv, (1, 21, 91, 160, 4), (1, 2, 2)))
+    self.assertFalse(can_embed_patches_as_tokens(conv, (1, 21, 90, 160, 4), (1, 1, 2)))
+    overlapping = nnx.Conv(4, 8, kernel_size=(1, 2, 2), strides=(1, 1, 1), rngs=nnx.Rngs(0))
+    self.assertFalse(can_embed_patches_as_tokens(overlapping, (1, 21, 90, 160, 4), (1, 2, 2)))
+    dilated = _conv((1, 2, 2), kernel_dilation=(1, 2, 2))
+    self.assertFalse(can_embed_patches_as_tokens(dilated, (1, 21, 90, 160, 4), (1, 2, 2)))
+
+
+class WanModelPatchEmbedModeTest(unittest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    pyconfig.initialize([None, os.path.join(THIS_DIR, "..", "configs", "base_wan_14b.yml")], unittest=True)
+    self.config = pyconfig.config
+    self.mesh = Mesh(create_device_mesh(self.config), self.config.mesh_axes)
+    with nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      self.model = WanModel(
+          rngs=nnx.Rngs(0),
+          num_attention_heads=2,
+          attention_head_dim=128,
+          in_channels=4,
+          out_channels=4,
+          text_dim=32,
+          freq_dim=32,
+          ffn_dim=256,
+          num_layers=1,
+          mesh=self.mesh,
+          attention="dot_product",
+      )
+    self.inputs = {
+        "hidden_states": jax.random.normal(jax.random.key(3), (1, 4, 2, 8, 16), jnp.float32),
+        "timestep": jnp.ones((1,)),
+        "encoder_hidden_states": jax.random.normal(jax.random.key(4), (1, 16, 32), jnp.float32),
+    }
+
+  def _run(self, mode):
+    # The mode is a build-time attribute (normally from attention_config).
+    self.model.patch_embed_mode = mode
+    with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      return self.model(**self.inputs)
+
+  @_SKIP_IN_GITHUB_ACTIONS
+  def test_tokens_mode_uses_the_token_path_and_matches_conv(self):
+    want = self._run("conv")
+    with mock.patch.object(transformer_wan, "embed_patches_as_tokens", wraps=transformer_wan.embed_patches_as_tokens) as spy:
+      got = self._run("tokens")
+    spy.assert_called_once()
+    self.assertEqual(got.shape, want.shape)
+    np.testing.assert_allclose(np.asarray(got), np.asarray(want), rtol=2e-2, atol=2e-2)
+
+  @_SKIP_IN_GITHUB_ACTIONS
+  def test_conv_mode_does_not_use_the_token_path(self):
+    with mock.patch.object(transformer_wan, "embed_patches_as_tokens") as spy:
+      self._run("conv")
+    spy.assert_not_called()
+
+  def test_rejects_unknown_mode(self):
+    with self.assertRaises(ValueError):
+      self._run("bogus")
+
+  def test_mode_comes_from_attention_config_not_env(self):
+    with mock.patch.dict(os.environ, {"WAN_PATCH_EMBED_MODE": "tokens"}):
+      with nn_partitioning.axis_rules(self.config.logical_axis_rules):
+        model = nnx.eval_shape(
+            lambda: WanModel(
+                rngs=nnx.Rngs(0),
+                num_attention_heads=2,
+                attention_head_dim=128,
+                in_channels=4,
+                out_channels=4,
+                text_dim=32,
+                freq_dim=32,
+                ffn_dim=256,
+                num_layers=1,
+                attention="dot_product",
+            )
+        )
+    self.assertEqual(model.patch_embed_mode, "conv")
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/src/maxdiffusion/tests/wan_runtime_options_test.py b/src/maxdiffusion/tests/wan_runtime_options_test.py
index cd1594dee..f01690656 100644
--- a/src/maxdiffusion/tests/wan_runtime_options_test.py
+++ b/src/maxdiffusion/tests/wan_runtime_options_test.py
@@ -129,7 +129,9 @@ def test_unset_options_resolve_to_defaults_not_env(self):
     env.update({"WAN_ROPE_NORM_MODE": "fused", "WAN_ROPE_ACCUM": "f32"})
     with mock.patch.dict(os.environ, env):
       attn = self._attn()
-    for name in opts.ATTENTION_OPTIONS:
+    # wan_patch_embed_mode rides in attention_config too, but is consumed by
+    # WanModel (covered in wan_patch_embed_test).
+    for name in sorted(set(opts.ATTENTION_OPTIONS) - {"wan_patch_embed_mode"}):
       with self.subTest(option=name):
         self.assertEqual(getattr(attn, name), opts.default(name))
 
@@ -164,6 +166,7 @@ def test_wan_model_resolves_cfg_before_unpatchify(self):
   def test_enum_options_validation(self):
     valid_cases = {
         "wan_cross_attn_kernel": ("xla", "pallas"),
+        "wan_patch_embed_mode": ("conv", "tokens"),
     }
     for name, valid_values in valid_cases.items():
       for val in valid_values:
diff --git a/src/maxdiffusion/wan_runtime_options.py b/src/maxdiffusion/wan_runtime_options.py
index bc611bb8a..94b82af4f 100644
--- a/src/maxdiffusion/wan_runtime_options.py
+++ b/src/maxdiffusion/wan_runtime_options.py
@@ -54,6 +54,7 @@
     "wan_cross_attn_prescale_kv": ("WAN_CROSS_ATTN_PRESCALE_KV", False, "bool"),
     "wan_rope_accum": ("WAN_ROPE_ACCUM", "auto", "str"),
     "wan_cross_attn_kernel": ("WAN_CROSS_ATTN_KERNEL", "xla", "str"),
+    "wan_patch_embed_mode": ("WAN_PATCH_EMBED_MODE", "conv", "str"),
     "wan_cross_attn_cpu_interpret": ("WAN_CROSS_ATTN_CPU_INTERPRET", False, "bool"),
 }
 
@@ -76,6 +77,8 @@ def _coerce(name: str, value: Any) -> Any:
     raise ValueError(f"wan_rope_accum must be 'auto', 'dtype', or 'f32', got {value!r}.")
   if name == "wan_cross_attn_kernel" and val not in ("xla", "pallas"):
     raise ValueError(f"wan_cross_attn_kernel must be 'xla' or 'pallas', got {value!r}.")
+  if name == "wan_patch_embed_mode" and val not in ("conv", "tokens"):
+    raise ValueError(f"wan_patch_embed_mode must be 'conv' or 'tokens', got {value!r}.")
   return val
 
 
@@ -137,6 +140,7 @@ def resolve_from_config(config: Any, name: str) -> Any:
     "wan_cross_attn_prescale_kv",
     "wan_rope_accum",
     "wan_cross_attn_kernel",
+    "wan_patch_embed_mode",
     "wan_cross_attn_cpu_interpret",
 )