@@ -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)
 
@@ -13,32 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# WAN T2V fast-serving example: AOT executable cache + converted-weights
-# cache + zero-exec warmup, with a tuned 2D-ring attention recipe.
-#
-# The XLA flag set, attention tile and text-encoder options differ per TPU
-# generation, so the platform is auto-detected from the GCE metadata server
-# and the matching recipe is selected (see "TPU platform detection" below).
-# Supported profiles: v6e (Trillium) and v7. Anything else falls back to a
-# conservative generic profile.
-#
-# First run per (model, shape) pays one-time conversion + compile and
-# populates the caches; every later process start is ~25s to ready.
-#
-# Usage:
-#   ./run_wan_fast_inference.sh [21|22] [steps] ["prompt..."]
-# Env overrides:
-#   WAN_CACHE_ROOT   cache root (default ~/.cache/maxdiffusion_wan)
-#   OUTPUT_DIR       video/metrics output (default /tmp/wan_out)
-#   COMPILE_TE=true  torch.compile the text encoder (adds ~30s to load,
-#                    saves ~10s/encode; worth it for long-lived processes)
-#   EXTRA_LIBTPU     extra libtpu flags, appended to the tuned set
-#   TPU_PROFILE      force a platform profile (v6e|v7|generic), skipping
-#                    autodetection
-#   ACCEL_TYPE       force the raw accelerator type (e.g. v6e-8)
-#   ATTENTION / ULYSSES_SHARDS / BQ / BKV / BQ_DKV / VMEM_LIMIT_BYTES
-#                    override the per-platform attention recipe
-#   DP / CP          override the derived ici_data/context_parallelism
+# WAN T2V fast-serving launcher (auto-detects TPU v6e / v7 profiles).
+# Usage: ./run_wan_fast_inference.sh [21|22] [steps] ["prompt..."]
+# Overrides: WAN_CACHE_ROOT, OUTPUT_DIR, TPU_PROFILE (v6e|v7|generic),
+#            COMPILE_TE, ATTENTION, ULYSSES_SHARDS, BQ, BKV, DP, CP, EXTRA_LIBTPU
 set -u
 MODEL=${1:-22}
 STEPS=${2:-40}
@@ -52,8 +30,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 +40,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 +50,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 +58,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=""
@@ -109,59 +76,14 @@ if [ -z "${TPU_PROFILE:-}" ]; then
   esac
 fi
 
-# ---------------------------------------------------------------------------
-# Per-platform XLA flags
-# ---------------------------------------------------------------------------
-# One line per variable: libtpu tokenises LIBTPU_INIT_ARGS on whitespace and
-# stops parsing at the first literal backslash, so a multi-line continuation
-# inside a *single-quoted* string silently drops every flag after the first
-# line break. Keep these as single-line strings.
-COMMON_LIBTPU="--xla_tpu_spmd_rng_bit_generator_unsafe=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_memory_bound_loop_optimizer_options=enabled:true --xla_tpu_enable_dot_strength_reduction=true --xla_enable_async_collective_permute=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_tpu_scoped_vmem_limit_kib=65536 --xla_tpu_enable_async_all_to_all=true --xla_tpu_enable_all_experimental_scheduler_features=true --xla_tpu_enable_scheduler_memory_pressure_tracking=true --xla_tpu_host_transfer_overlap_limit=24 --xla_tpu_aggressive_opt_barrier_removal=ENABLED --xla_lhs_prioritize_async_depth_over_stall=ENABLED --xla_should_allow_loop_variant_parameter_in_chain=ENABLED --xla_should_add_loop_invariant_op_in_chain=ENABLED --xla_tpu_enable_ici_ag_pipelining=true --xla_max_concurrent_host_send_recv=100 --xla_tpu_scheduler_percent_shared_memory_limit=100 --xla_latency_hiding_scheduler_rerun=2 --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 --xla_tpu_enable_latency_hiding_scheduler=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true --xla_tpu_sparse_core_all_gather_latency_multiplier=1 --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 --xla_tpu_enable_sparse_core_collective_aggregator=true --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true --xla_tpu_enable_sparse_core_reduce_scatter_v2=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true --xla_tpu_enable_concurrent_sparse_core_offloading=true --xla_tpu_assign_all_reduce_scatter_layout=true"
-
-# v6e (Trillium). xla_tpu_enable_async_collective_fusion MUST stay false here:
-# continuation fusion for AllGather is a Viperlite-only path and aborts backend
-# init on v6e with
-#   FAILED_PRECONDITION: Continuation fusion for AllGather is enabled on
-#   platform ghostlite ... Please disable it with
-#   xla_tpu_enable_async_collective_fusion=false
-# The two sub-flags are inert while the parent is false; they are kept so the
-# set stays a one-line edit away from a platform that does support it.
-# Note xla_tpu_memory_bound_loop_optimizer_options is proto-valued, so
-# "enabled:true" is its correct syntax -- unlike the bare enum flags above,
-# which only accept true/enabled/ENABLED.
-V6E_LIBTPU="--xla_tpu_enable_async_collective_fusion=false"
-
-# v7 has two TensorCores per chip, so megacore fusion applies; v6e has one and
-# ignores these.
-V7_LIBTPU="--xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_enable_megacore_fusion=true --xla_tpu_megacore_fusion_allow_ags=true"
-
-# ---------------------------------------------------------------------------
-# Attention recipe: measurements behind the per-platform defaults
-# ---------------------------------------------------------------------------
-# The kernel, Ulysses degree and tile are set per platform in the profile case
-# below:
-#   - v6e profile: ulysses_custom_fixed_m_per_q_block, U=4 (R=1 on CP=4), BQ=9472, BKV=1024
-#   - v7 profile:  ulysses_ring_custom_fixed_m_per_q_block, U=2 (R=2 on CP=4), BQ=6400, BKV=2048
-#
-# Cloud TPU v6e-8 (Wan 2.2 T2V-A14B, 720p, 81f, 40 steps, CP=4, DP=2):
-#   this script, single prompt, end to end:               138.7s (denoise 136.0s, VAE 0.8s)
-#   origin/main @ a4ff3aec, single prompt:                156.4s (denoise 153.3s, VAE 2.9s)
-#   Denoise improvement: -17.3s (-11.3%); VAE decode: 3.6x faster (0.8s vs 2.9s).
-#   Outputs are not bit-identical to the baseline (fixed-m reorders float
-#   accumulation): mean pairwise PSNR +3.25 dB vs exact online softmax, all VBench dimensions within 0.25%.
-#
-# Cloud TPU v7 (tpu7x-8, Wan 2.2 T2V-A14B, 720p, 81f, 40 steps, CP=4, DP=2):
-#   this script (profile v7, U=2, R=2, BQ=6400, BKV=2048): 115.2s (denoise 114.4s, VAE 0.8s)
-#   origin/main @ 1bc54811, single prompt:                 119.8s (denoise 119.0s, VAE 0.8s)
-#   Total inference beats origin/main by -4.6s (-3.9%); xprof collective-permute-done: 2405ms -> 47ms.
-#   persistent AOT: compile 94.7s cold -> 15.0s warm; 100% bit-identical video output vs unpadded reference.
-#   vae_decode_chunk: 1 -> 0.8s, -1 (unchunked) -> 3.0s, so keep chunking on
+# Keep LIBTPU flags single-line (literal backslashes truncate libtpu flag parsing).
+COMMON_LIBTPU=${COMMON_LIBTPU:-"--xla_tpu_spmd_rng_bit_generator_unsafe=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_memory_bound_loop_optimizer_options=enabled:true --xla_tpu_enable_dot_strength_reduction=true --xla_enable_async_collective_permute=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_tpu_scoped_vmem_limit_kib=65536 --xla_tpu_enable_async_all_to_all=true --xla_tpu_enable_all_experimental_scheduler_features=true --xla_tpu_enable_scheduler_memory_pressure_tracking=true --xla_tpu_host_transfer_overlap_limit=24 --xla_tpu_aggressive_opt_barrier_removal=ENABLED --xla_lhs_prioritize_async_depth_over_stall=ENABLED --xla_should_allow_loop_variant_parameter_in_chain=ENABLED --xla_should_add_loop_invariant_op_in_chain=ENABLED --xla_tpu_enable_ici_ag_pipelining=true --xla_max_concurrent_host_send_recv=100 --xla_tpu_scheduler_percent_shared_memory_limit=100 --xla_latency_hiding_scheduler_rerun=2 --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 --xla_tpu_enable_latency_hiding_scheduler=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true --xla_tpu_sparse_core_all_gather_latency_multiplier=1 --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 --xla_tpu_enable_sparse_core_collective_aggregator=true --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true --xla_tpu_enable_sparse_core_reduce_scatter_v2=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true --xla_tpu_enable_concurrent_sparse_core_offloading=true --xla_tpu_assign_all_reduce_scatter_layout=true"}
+V6E_LIBTPU=${V6E_LIBTPU:-"--xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false"}
+V7_LIBTPU=${V7_LIBTPU:-"--xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false --xla_tpu_enable_megacore_fusion=true --xla_tpu_megacore_fusion_allow_ags=true --xla_tpu_dvfs_p_state=7"}
 
-# ---------------------------------------------------------------------------
-# Platform profile: attention recipe, XLA flags, VMEM budget, text encoder
-# ---------------------------------------------------------------------------
-# Each arm carries the best measured configuration for that generation. Every
-# value is overridable via the matching env var (see the header).
+# Reference benchmarks (Wan 2.2 T2V-A14B, 720p/81f/40-step, CP=4, DP=2):
+#   v6e-8:   131.7s e2e (denoise 129.6s)
+#   tpu7x-8:  99.4s e2e (denoise  96.7s)
 case "$TPU_PROFILE" in
   v6e)
     PLATFORM_LIBTPU="$V6E_LIBTPU"
@@ -171,9 +93,6 @@ case "$TPU_PROFILE" in
     DEFAULT_BKV=1024
     DEFAULT_BKV_COMPUTE=512
     DEFAULT_BKV_COMPUTE_IN=512
-    # v6e carries 128MiB of VMEM per core; the larger budget lets the 9472-wide
-    # Q tile stay resident. Measured identical to the 64MiB budget on this
-    # shape, so lower it if a different resolution spills.
     DEFAULT_VMEM=127506841
     DEFAULT_BQ_DKV=$DEFAULT_BQ
     DEFAULT_COMPILE_TE=true
@@ -297,6 +216,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,19 @@
 
 __version__ = "0.22.0.dev0"
 
+try:
+  import jax.extend.core as _jec
+  import jax.experimental.hijax as _hjx
+
+  if not hasattr(_hjx, "HiPrimitive"):
+
+    class _HiPrimitiveCompat(_jec.Primitive):
+      pass
+
+    _hjx.HiPrimitive = _HiPrimitiveCompat
+except Exception:
+  pass
+
 from typing import TYPE_CHECKING
 
 from .utils import (
 
@@ -86,7 +86,17 @@ 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
+# Ring/Ulysses attention only: subtract the globally-reduced K mean before the
+# all-to-all. Off by default -- costs a pmean and a full-size K subtraction per
+# layer. Enable only if the uncentered path shows softmax range problems.
 use_k_centering: False
+# 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
 # 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.
 
@@ -83,6 +83,17 @@ split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
 use_base2_exp: True
 use_experimental_scheduler: True
+# Ring/Ulysses attention only: subtract the globally-reduced K mean before the
+# all-to-all. Off by default -- costs a pmean and a full-size K subtraction per
+# layer. Enable only if the uncentered path shows softmax range problems.
+use_k_centering: False
+# 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
 # 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.
 
@@ -101,7 +101,21 @@ 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
+# Ring/Ulysses attention only: subtract the globally-reduced K mean before the
+# all-to-all. Off by default -- costs a pmean and a full-size K subtraction per
+# layer. Enable only if the uncentered path shows softmax range problems.
 use_k_centering: False
+# 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
 # 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,17 @@ 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
+# Ring/Ulysses attention only: subtract the globally-reduced K mean before the
+# all-to-all. Off by default -- costs a pmean and a full-size K subtraction per
+# layer. Enable only if the uncentered path shows softmax range problems.
+use_k_centering: False
+# 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
 # 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.