@@ -57,6 +57,8 @@ WAN_OVERRIDES=(
   "use_batched_text_encoder|USE_BATCHED_TEXT_ENCODER|true"
   "flash_block_sizes|FLASH_BLOCK_SIZES|"
   "prompt_file|PROMPT_FILE|./benchmarks/vbench/prompts_110.txt"
+  # Matches the shipped config default; set USE_FUSED_ROPE_KERNEL=true to test it.
+  "use_fused_rope_kernel|USE_FUSED_ROPE_KERNEL|false"
 )
 
 usage() {
@@ -70,6 +72,9 @@ Common options:
   RUN_NAME           Generation run name (default: wan-inference; videos are saved to <RUN_NAME>/videos)
   PROMPT_FILE        Prompt file path (default: ./benchmarks/vbench/prompts_110.txt)
   CONFIG_FILE        WAN config file (default: src/maxdiffusion/configs/base_wan_27b.yml)
+  USE_FUSED_ROPE_KERNEL
+                     Enable the fused RMSNorm+RoPE Pallas producer (default: false).
+                     Only measured on TPU v6e and v7; the model refuses it elsewhere.
   EXTERNAL_DISK      Mounted disk root for large local files (default: /mnt/disks/external_disk)
   HF_CACHE_ROOT      Hugging Face cache root (default: \$EXTERNAL_DISK/hf_cache)
   HF_HOME            Hugging Face home directory (default: \$HF_CACHE_ROOT)
 
@@ -52,8 +52,6 @@ export HF_HUB_ENABLE_HF_TRANSFER=1
 export JAX_DEFAULT_MATMUL_PRECISION=bfloat16
 export TORCHINDUCTOR_FX_GRAPH_CACHE=1
 export XLA_PYTHON_CLIENT_MEM_FRACTION=0.95
-# Without these the JAX persistent cache silently skips most entries, so the
-# "warm" start still recompiles a large part of the graph.
 export JAX_PERSISTENT_CACHE_MIN_ENTRY_SIZE_BYTES=-1
 export JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS=0
 
@@ -64,14 +62,7 @@ export TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR:-$CACHE_ROOT/torch_comp
 mkdir -p "$CACHE_ROOT/jax" "$CACHE_ROOT/aot_wan$MODEL" "$CACHE_ROOT/converted" \
          "$OUTPUT_DIR" "$TMPDIR" "$TORCHINDUCTOR_CACHE_DIR"
 
-# ---------------------------------------------------------------------------
-# TPU platform detection
-# ---------------------------------------------------------------------------
-# Preference order: explicit override -> TPU_ACCELERATOR_TYPE env (set by some
-# runtimes) -> GCE metadata "accelerator-type" (e.g. "v6e-8") -> the
-# ACCELERATOR_TYPE line inside the "tpu-env" metadata blob. Detection is pure
-# metadata/env: it must not initialise the TPU, or it would take the device
-# before the real process starts.
+# Detect TPU generation via env or GCE metadata without initializing the device.
 _tpu_metadata() {
   curl -s -m 2 -H 'Metadata-Flavor: Google' \
     "http://metadata.google.internal/computeMetadata/v1/instance/attributes/$1" 2> /dev/null
@@ -81,8 +72,6 @@ _detect_accel_type() {
   local t="${TPU_ACCELERATOR_TYPE:-}"
   [ -z "$t" ] && t="$(_tpu_metadata accelerator-type)"
   [ -z "$t" ] && t="$(_tpu_metadata tpu-env | sed -n "s/^ACCELERATOR_TYPE: *'\([^']*\)'.*/\1/p")"
-  # Guard against metadata returning an HTML error page. Real names seen in the
-  # wild: "v6e-8", "v5litepod-8", "tpu7x-8" (v7 reports as tpu7x, not v7x).
   case "$t" in
     v[0-9]* | tpu[0-9]*) printf '%s' "$t" ;;
     *) printf '' ;;
@@ -91,8 +80,8 @@ _detect_accel_type() {
 
 ACCEL_TYPE=${ACCEL_TYPE:-$(_detect_accel_type)}
 if [ -n "$ACCEL_TYPE" ]; then
-  TPU_GEN="${ACCEL_TYPE%%-*}"  # v6e-8      -> v6e
-  TPU_CHIPS="${ACCEL_TYPE##*-}" # v6e-8      -> 8
+  TPU_GEN="${ACCEL_TYPE%%-*}"
+  TPU_CHIPS="${ACCEL_TYPE##*-}"
 else
   TPU_GEN=""
   TPU_CHIPS=""
@@ -182,6 +171,7 @@ case "$LIBTPU_INIT_ARGS" in
 esac
 
 ATTENTION=${ATTENTION:-$DEFAULT_ATTENTION}
+USE_K_CENTERING=${USE_K_CENTERING:-auto}  # auto: on for non-ring (v6e), off for ring (v7)
 ULYSSES_SHARDS=${ULYSSES_SHARDS:-$DEFAULT_U}
 BQ=${BQ:-$DEFAULT_BQ}
 BKV=${BKV:-$DEFAULT_BKV}
@@ -242,6 +232,7 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \
   aot_cache_dir="$CACHE_ROOT/aot_wan$MODEL" \
   converted_weights_dir="$CACHE_ROOT/converted" \
   attention="$ATTENTION" \
+  use_k_centering="$USE_K_CENTERING" \
   ulysses_shards="$ULYSSES_SHARDS" \
   ici_data_parallelism="$DP" ici_fsdp_parallelism=1 \
   ici_context_parallelism="$CP" ici_tensor_parallelism=1 \
@@ -252,6 +243,7 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \
   vae_weights_dtype=bfloat16 vae_dtype=bfloat16 \
   text_encoder_dtype=bfloat16 compile_text_encoder="$COMPILE_TE" use_batched_text_encoder="$USE_BATCHED_TE" \
   use_kv_cache=true use_base2_exp=true use_experimental_scheduler=true \
+  use_fused_rope_kernel=true \
   fps=16 "${GUIDANCE_ARGS[@]}" \
   seed="${SEED:-12345}" \
   flash_block_sizes="$FLASH_BLOCK_SIZES" \
 
@@ -16,6 +16,24 @@
 
 __version__ = "0.22.0.dev0"
 
+try:
+  import jax.core as _jc
+  import jax.extend.core as _jec
+  import jax.experimental.hijax as _hjx
+
+  if not hasattr(_jc, "Effect") and hasattr(_jec, "Effect"):
+    _jc.Effect = _jec.Effect
+
+  if not hasattr(_hjx, "HiPrimitive"):
+    _base_prim = getattr(_hjx, "VJPHiPrimitive", _jec.Primitive)
+
+    class _HiPrimitiveCompat(_base_prim):
+      pass
+
+    _hjx.HiPrimitive = _HiPrimitiveCompat
+except (ImportError, AttributeError):
+  pass
+
 from typing import TYPE_CHECKING
 
 from .utils import (
@@ -55,6 +73,7 @@
     "schedulers": [],
     "tpu_utils": [],
     "train_utils": [],
+    "wan_runtime_options": [],
     "utils": [
         "OptionalDependencyNotAvailable",
         "is_flax_available",
 
@@ -235,6 +235,13 @@ def _extract_graphdef_statics(obj: Any, prefix: str = "") -> list[str]:
   return out
 
 
+def _runtime_env_statics() -> tuple[str, ...]:
+  """Graph-changing Wan switches (YAML `wan_*` keys), folded into every signature."""
+  from maxdiffusion import wan_runtime_options  # pylint: disable=import-outside-toplevel
+
+  return tuple(f"opt:{k}={v}" for k, v in sorted(wan_runtime_options.snapshot().items()))
+
+
 def _dynamic_signature(args: tuple, kwargs: dict) -> str:
   """Deterministic digest of everything that selects an executable.
 
@@ -260,6 +267,7 @@ def _dynamic_signature(args: tuple, kwargs: dict) -> str:
       desc = _format_static_val(leaf)
     parts.append(f"{jax.tree_util.keystr(path)}={desc}")
   parts.extend(_extract_graphdef_statics((args, kwargs)))
+  parts.extend(_runtime_env_statics())
   return hashlib.sha256("|".join(parts).encode()).hexdigest()[:12]
 
 
@@ -368,7 +376,7 @@ def __call__(self, *args, **kwargs):
     )
     static_items = tuple((k, _format_static_val(v)) for k, v in sorted(static.items(), key=lambda item: str(item[0])))
     graphdef_statics = tuple(_extract_graphdef_statics((args, kwargs)))
-    cache_key = (treedef, shapes_dtypes, static_items, graphdef_statics)
+    cache_key = (treedef, shapes_dtypes, static_items, graphdef_statics, _runtime_env_statics())
     signature = getattr(self, "_sig_cache", {}).get(cache_key)
     if signature is None:
       signature = _dynamic_signature((), {**dynamic, **static})
 
@@ -105,7 +105,34 @@ attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cud
 
 use_base2_exp: True
 use_experimental_scheduler: True
-use_k_centering: "auto"  # auto: on for non-ring ulysses_custom_fixed_m* (free, virtual), off for ring paths (pmean + HBM copy)
+# use_k_centering: "auto" turns K-centering on for the non-ring
+# ulysses_custom_fixed_m* kernels (virtual: q.k_mean folded into the kernel, no
+# collective, no HBM copy) and off for ring kernels (materialized: a pmean plus a
+# full-size K subtraction per layer). Set True/False to force either path.
+use_k_centering: "auto"
+# Fused RMSNorm+RoPE+head-transpose Pallas producer for self-attention QK.
+# Off by default; only measured on tpu7x and v6e. See base_wan_27b.yml.
+use_fused_rope_kernel: False
+# Fused RoPE kernel grid: sequence block size, and heads per grid step (-1 =
+# all). Tuned at the Wan 2.2 27B shard shape; re-measure for this model.
+fused_rope_block_s: 1024
+fused_rope_head_block: -1
+# Wan inference switches. These change the compiled graph, so they are part of
+# the AOT cache fingerprint. (They used to be WAN_* environment variables; the
+# env vars are still honoured only when no config is loaded, e.g. in unit tests.)
+# RMSNorm reduction for the fused RoPE producer: "exact" (XLA reduction,
+# bit-identical to the jitted XLA producer) or "fused" (in-kernel, <=1 bf16 ULP).
+wan_rope_norm_mode: "exact"
+# Fold log2(e) into Q and 1/sqrt(d) into K inside the fused RoPE producer.
+wan_fuse_qk_prescale: True
+# Have the splash kernel write its output already transposed to [B, S, H, D].
+wan_splash_transpose_out: False
+# Apply classifier-free guidance before unpatchify (halves the unpatchify work).
+wan_cfg_before_unpatchify: True
+# Pre-scale cached cross-attention K by 1/sqrt(d) once per video.
+wan_cross_attn_prescale_kv: False
+# Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
+wan_rope_accum: "auto"
 # 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.
 
@@ -102,6 +102,34 @@ attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cud
 
 use_base2_exp: True
 use_experimental_scheduler: True
+# use_k_centering: "auto" turns K-centering on for the non-ring
+# ulysses_custom_fixed_m* kernels (virtual: q.k_mean folded into the kernel, no
+# collective, no HBM copy) and off for ring kernels (materialized: a pmean plus a
+# full-size K subtraction per layer). Set True/False to force either path.
+use_k_centering: "auto"
+# Fused RMSNorm+RoPE+head-transpose Pallas producer for self-attention QK.
+# Off by default; only measured on tpu7x and v6e. See base_wan_27b.yml.
+use_fused_rope_kernel: False
+# Fused RoPE kernel grid: sequence block size, and heads per grid step (-1 =
+# all). Tuned at the Wan 2.2 27B shard shape; re-measure for this model.
+fused_rope_block_s: 1024
+fused_rope_head_block: -1
+# Wan inference switches. These change the compiled graph, so they are part of
+# the AOT cache fingerprint. (They used to be WAN_* environment variables; the
+# env vars are still honoured only when no config is loaded, e.g. in unit tests.)
+# RMSNorm reduction for the fused RoPE producer: "exact" (XLA reduction,
+# bit-identical to the jitted XLA producer) or "fused" (in-kernel, <=1 bf16 ULP).
+wan_rope_norm_mode: "exact"
+# Fold log2(e) into Q and 1/sqrt(d) into K inside the fused RoPE producer.
+wan_fuse_qk_prescale: True
+# Have the splash kernel write its output already transposed to [B, S, H, D].
+wan_splash_transpose_out: False
+# Apply classifier-free guidance before unpatchify (halves the unpatchify work).
+wan_cfg_before_unpatchify: True
+# Pre-scale cached cross-attention K by 1/sqrt(d) once per video.
+wan_cross_attn_prescale_kv: False
+# Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
+wan_rope_accum: "auto"
 # 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.
 
@@ -123,7 +123,38 @@ attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cud
 #     ulysses_shards to them now raises instead of being silently ignored.
 use_base2_exp: True
 use_experimental_scheduler: True
-use_k_centering: "auto"  # auto: on for non-ring ulysses_custom_fixed_m* (free, virtual), off for ring paths (pmean + HBM copy)
+# use_k_centering: "auto" turns K-centering on for the non-ring
+# ulysses_custom_fixed_m* kernels (virtual: q.k_mean folded into the kernel, no
+# collective, no HBM copy) and off for ring kernels (materialized: a pmean plus a
+# full-size K subtraction per layer). Set True/False to force either path.
+use_k_centering: "auto"
+# Fused RMSNorm+RoPE+head-transpose Pallas producer for self-attention QK.
+# ~1.9x faster in isolation; 3.1% e2e denoise gain on v6e-8, 1.7% on tpu7x-8.
+# Off by default: inlined into the model XLA rounds the unfused producer
+# differently, so the output is equivalent but not identical (v6e 720p/81f,
+# same seed: 53.8 dB PSNR after 1 denoise step, 34.1 dB after 40). Refused on
+# platforms whose RoPE rounding is unmeasured; tuned and measured for inference.
+use_fused_rope_kernel: False
+# Sequence block size for the fused RoPE kernel grid.
+fused_rope_block_s: 1024
+# Heads processed per grid step; -1 means "all heads in one step".
+fused_rope_head_block: -1
+# Wan inference switches. These change the compiled graph, so they are part of
+# the AOT cache fingerprint. (They used to be WAN_* environment variables; the
+# env vars are still honoured only when no config is loaded, e.g. in unit tests.)
+# RMSNorm reduction for the fused RoPE producer: "exact" (XLA reduction,
+# bit-identical to the jitted XLA producer) or "fused" (in-kernel, <=1 bf16 ULP).
+wan_rope_norm_mode: "exact"
+# Fold log2(e) into Q and 1/sqrt(d) into K inside the fused RoPE producer.
+wan_fuse_qk_prescale: True
+# Have the splash kernel write its output already transposed to [B, S, H, D].
+wan_splash_transpose_out: False
+# Apply classifier-free guidance before unpatchify (halves the unpatchify work).
+wan_cfg_before_unpatchify: True
+# Pre-scale cached cross-attention K by 1/sqrt(d) once per video.
+wan_cross_attn_prescale_kv: False
+# Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
+wan_rope_accum: "auto"
 # 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.
 
@@ -84,6 +84,34 @@ split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block
 use_base2_exp: True
 use_experimental_scheduler: True
+# use_k_centering: "auto" turns K-centering on for the non-ring
+# ulysses_custom_fixed_m* kernels (virtual: q.k_mean folded into the kernel, no
+# collective, no HBM copy) and off for ring kernels (materialized: a pmean plus a
+# full-size K subtraction per layer). Set True/False to force either path.
+use_k_centering: "auto"
+# Fused RMSNorm+RoPE+head-transpose Pallas producer for self-attention QK.
+# Off by default; only measured on tpu7x and v6e. See base_wan_27b.yml.
+use_fused_rope_kernel: False
+# Fused RoPE kernel grid: sequence block size, and heads per grid step (-1 =
+# all). Tuned at the Wan 2.2 27B shard shape; re-measure for this model.
+fused_rope_block_s: 1024
+fused_rope_head_block: -1
+# Wan inference switches. These change the compiled graph, so they are part of
+# the AOT cache fingerprint. (They used to be WAN_* environment variables; the
+# env vars are still honoured only when no config is loaded, e.g. in unit tests.)
+# RMSNorm reduction for the fused RoPE producer: "exact" (XLA reduction,
+# bit-identical to the jitted XLA producer) or "fused" (in-kernel, <=1 bf16 ULP).
+wan_rope_norm_mode: "exact"
+# Fold log2(e) into Q and 1/sqrt(d) into K inside the fused RoPE producer.
+wan_fuse_qk_prescale: True
+# Have the splash kernel write its output already transposed to [B, S, H, D].
+wan_splash_transpose_out: False
+# Apply classifier-free guidance before unpatchify (halves the unpatchify work).
+wan_cfg_before_unpatchify: True
+# Pre-scale cached cross-attention K by 1/sqrt(d) once per video.
+wan_cross_attn_prescale_kv: False
+# Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
+wan_rope_accum: "auto"
 # 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.
 
@@ -86,6 +86,34 @@ split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block
 use_base2_exp: True
 use_experimental_scheduler: True
+# use_k_centering: "auto" turns K-centering on for the non-ring
+# ulysses_custom_fixed_m* kernels (virtual: q.k_mean folded into the kernel, no
+# collective, no HBM copy) and off for ring kernels (materialized: a pmean plus a
+# full-size K subtraction per layer). Set True/False to force either path.
+use_k_centering: "auto"
+# Fused RMSNorm+RoPE+head-transpose Pallas producer for self-attention QK.
+# Off by default; only measured on tpu7x and v6e. See base_wan_27b.yml.
+use_fused_rope_kernel: False
+# Fused RoPE kernel grid: sequence block size, and heads per grid step (-1 =
+# all). Tuned at the Wan 2.2 27B shard shape; re-measure for this model.
+fused_rope_block_s: 1024
+fused_rope_head_block: -1
+# Wan inference switches. These change the compiled graph, so they are part of
+# the AOT cache fingerprint. (They used to be WAN_* environment variables; the
+# env vars are still honoured only when no config is loaded, e.g. in unit tests.)
+# RMSNorm reduction for the fused RoPE producer: "exact" (XLA reduction,
+# bit-identical to the jitted XLA producer) or "fused" (in-kernel, <=1 bf16 ULP).
+wan_rope_norm_mode: "exact"
+# Fold log2(e) into Q and 1/sqrt(d) into K inside the fused RoPE producer.
+wan_fuse_qk_prescale: True
+# Have the splash kernel write its output already transposed to [B, S, H, D].
+wan_splash_transpose_out: False
+# Apply classifier-free guidance before unpatchify (halves the unpatchify work).
+wan_cfg_before_unpatchify: True
+# Pre-scale cached cross-attention K by 1/sqrt(d) once per video.
+wan_cross_attn_prescale_kv: False
+# Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
+wan_rope_accum: "auto"
 # 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.
 
@@ -86,6 +86,34 @@ split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block
 use_base2_exp: True
 use_experimental_scheduler: True
+# use_k_centering: "auto" turns K-centering on for the non-ring
+# ulysses_custom_fixed_m* kernels (virtual: q.k_mean folded into the kernel, no
+# collective, no HBM copy) and off for ring kernels (materialized: a pmean plus a
+# full-size K subtraction per layer). Set True/False to force either path.
+use_k_centering: "auto"
+# Fused RMSNorm+RoPE+head-transpose Pallas producer for self-attention QK.
+# Off by default; only measured on tpu7x and v6e. See base_wan_27b.yml.
+use_fused_rope_kernel: False
+# Fused RoPE kernel grid: sequence block size, and heads per grid step (-1 =
+# all). Tuned at the Wan 2.2 27B shard shape; re-measure for this model.
+fused_rope_block_s: 1024
+fused_rope_head_block: -1
+# Wan inference switches. These change the compiled graph, so they are part of
+# the AOT cache fingerprint. (They used to be WAN_* environment variables; the
+# env vars are still honoured only when no config is loaded, e.g. in unit tests.)
+# RMSNorm reduction for the fused RoPE producer: "exact" (XLA reduction,
+# bit-identical to the jitted XLA producer) or "fused" (in-kernel, <=1 bf16 ULP).
+wan_rope_norm_mode: "exact"
+# Fold log2(e) into Q and 1/sqrt(d) into K inside the fused RoPE producer.
+wan_fuse_qk_prescale: True
+# Have the splash kernel write its output already transposed to [B, S, H, D].
+wan_splash_transpose_out: False
+# Apply classifier-free guidance before unpatchify (halves the unpatchify work).
+wan_cfg_before_unpatchify: True
+# Pre-scale cached cross-attention K by 1/sqrt(d) once per video.
+wan_cross_attn_prescale_kv: False
+# Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
+wan_rope_accum: "auto"
 # 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.