[optimization] Added Wan 2.2 VAE decode optimizations (#7056)

Signed-off-by: Maciej Bala <mbala@nvidia.com>
Signed-off-by: MaciejBalaNV <mbala@nvidia.com>

benchmarks/diffusion/bench_wan_vae_decode.py

@@ -0,0 +1,324 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark and exactness check for the Wan VAE decoder fast path.
+
+Loads a diffusers Wan VAE (by default the Cosmos3 one), decodes seeded latents
+for the requested video size at each ``--fast-path`` level and reports wall
+time, peak memory and, against the ``off`` level, bitwise equality, max abs
+difference and PSNR. ``--profile`` adds a torch.profiler CUDA kernel table with
+the convolution / layout-transpose share.
+
+Single GPU::
+
+    python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano
+    python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
+        --fast-path off,lossless --frames 33 --profile
+    python benchmarks/diffusion/bench_wan_vae_decode.py --tiny --frames 9 --size 256x256
+
+Multi-GPU VAE parallel decode (launch one process per GPU with ``torchrun``; the
+decode is timed across all ranks and only rank 0 prints)::
+
+    torchrun --nproc-per-node 2 benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
+        --vae-patch-parallel-size 2 --vae-parallel-mode tile
+
+``tile`` distributes the VAE's own spatial tiles over the ranks (it needs the
+video to exceed the tile size, 256x256 pixels by default) and ``spatial_shard_*``
+shards every decoder feature map along height or width with halo exchange. Both
+return the full video on rank 0 only, so equality and PSNR are computed there.
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import os
+import time
+from collections.abc import Iterable
+
+import torch
+import torch.distributed as dist
+
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+    DistributedAutoencoderKLWan,
+    OmniAutoencoderKLWan,
+)
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import (
+    VAE_FAST_PATH_LEVELS,
+    install_wan_vae_fastpath,
+)
+
+TINY_CONFIG = dict(
+    base_dim=32,
+    decoder_base_dim=64,
+    z_dim=48,
+    dim_mult=[1, 2, 4, 4],
+    num_res_blocks=2,
+    temperal_downsample=[False, True, True],
+    is_residual=True,
+    patch_size=2,
+    in_channels=12,
+    out_channels=12,
+    scale_factor_temporal=4,
+    scale_factor_spatial=16,
+)
+
+DTYPES = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
+VAE_PARALLEL_MODES = ("tile", "spatial_shard_height", "spatial_shard_width")
+_CONV_TAGS = ("conv", "cudnn", "gemm", "cutlass", "xmma")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--model", default="nvidia/Cosmos3-Nano", help="HF id or local path of the pipeline")
+    parser.add_argument("--subfolder", default="vae")
+    parser.add_argument("--tiny", action="store_true", help="Use a small random VAE instead of loading weights")
+    parser.add_argument("--size", default="1280x720", help="Output video WxH")
+    parser.add_argument("--frames", type=int, default=189, help="Output frame count (1 + 4k)")
+    parser.add_argument("--dtype", choices=sorted(DTYPES), default="bf16")
+    parser.add_argument("--fast-path", default=",".join(VAE_FAST_PATH_LEVELS), help="Comma-separated levels")
+    parser.add_argument("--warmup", type=int, default=1)
+    parser.add_argument("--iters", type=int, default=2)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--profile", action="store_true", help="Print a torch.profiler CUDA kernel table")
+    parser.add_argument("--profile-rows", type=int, default=25)
+    parser.add_argument("--save-output", default=None, help="Save the `off` decode output to this .pt path")
+    parser.add_argument(
+        "--vae-patch-parallel-size",
+        type=int,
+        default=1,
+        help="VAE parallel decode over this many ranks (run under torchrun with as many processes)",
+    )
+    parser.add_argument(
+        "--vae-parallel-mode",
+        choices=VAE_PARALLEL_MODES,
+        default="tile",
+        help="tile: distribute the VAE's spatial tiles; spatial_shard_*: shard feature maps with halo exchange",
+    )
+    return parser.parse_args()
+
+
+def init_distributed() -> tuple[int, int]:
+    """Join the torchrun process group (if any) and return ``(rank, world_size)``."""
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    if world_size <= 1:
+        return 0, 1
+    from vllm_omni.diffusion.distributed.parallel_state import (
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    rank = int(os.environ["RANK"])
+    local_rank = int(os.environ.get("LOCAL_RANK", rank))
+    init_distributed_environment(world_size=world_size, rank=rank, local_rank=local_rank)
+    initialize_model_parallel(sequence_parallel_size=world_size, ulysses_degree=world_size)
+    return rank, world_size
+
+
+def load_vae(
+    args: argparse.Namespace, dtype: torch.dtype, device: torch.device, parallel: bool
+) -> OmniAutoencoderKLWan:
+    cls = DistributedAutoencoderKLWan if parallel else OmniAutoencoderKLWan
+    if args.tiny:
+        torch.manual_seed(args.seed)
+        vae = cls(**TINY_CONFIG).to(device=device, dtype=dtype).eval()
+        if parallel:
+            vae.init_distributed()
+    else:
+        vae = cls.from_pretrained(args.model, subfolder=args.subfolder, torch_dtype=dtype).to(device).eval()
+    if parallel:
+        # Same setup as the registry: parallel decode needs tiling on.
+        vae.use_tiling = True
+        vae.set_parallel_size(args.vae_patch_parallel_size, mode=args.vae_parallel_mode)
+    return vae
+
+
+def make_latents(vae: OmniAutoencoderKLWan, args: argparse.Namespace, dtype: torch.dtype, device: torch.device):
+    width, height = (int(v) for v in args.size.lower().split("x"))
+    spatial = int(getattr(vae.config, "scale_factor_spatial", 16))
+    temporal = int(getattr(vae.config, "scale_factor_temporal", 4))
+    if width % spatial or height % spatial or (args.frames - 1) % temporal:
+        raise SystemExit(f"size must be a multiple of {spatial} and frames must be 1 + {temporal}k")
+    shape = (1, vae.config.z_dim, (args.frames - 1) // temporal + 1, height // spatial, width // spatial)
+    generator = torch.Generator(device="cpu").manual_seed(args.seed)
+    latents = torch.randn(shape, generator=generator).to(device=device, dtype=dtype)
+    mean = getattr(vae.config, "latents_mean", None)
+    std = getattr(vae.config, "latents_std", None)
+    if mean is not None and std is not None:
+        mean = torch.as_tensor(mean, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
+        std = torch.as_tensor(std, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
+        latents = latents * std + mean
+    return latents
+
+
+@torch.inference_mode()
+def decode(vae: OmniAutoencoderKLWan, latents: torch.Tensor) -> torch.Tensor:
+    return vae.decode(latents, return_dict=False)[0]
+
+
+def sync() -> None:
+    """Wait for this rank's GPU and, under torchrun, for every other rank."""
+    torch.accelerator.synchronize()
+    if dist.is_initialized():
+        dist.barrier()
+
+
+def psnr(a: torch.Tensor, b: torch.Tensor) -> float:
+    mse = torch.mean((a.float() - b.float()) ** 2).item()
+    return math.inf if mse == 0 else 10 * math.log10(4.0 / mse)  # outputs live in [-1, 1]
+
+
+def bits_equal(a: torch.Tensor, b: torch.Tensor) -> bool:
+    if a.shape != b.shape or a.dtype != b.dtype:
+        return False
+    int_dtype = torch.int16 if a.element_size() == 2 else torch.int32
+    return bool((a.contiguous().view(int_dtype) == b.contiguous().view(int_dtype)).all())
+
+
+def profile_decode(vae: OmniAutoencoderKLWan, latents: torch.Tensor, rows: int, verbose: bool) -> None:
+    from torch.profiler import ProfilerActivity, profile
+
+    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=False) as prof:
+        decode(vae, latents)
+        sync()
+    if not verbose:
+        return
+    events = [e for e in prof.key_averages() if e.device_time_total > 0]
+    events.sort(key=lambda e: e.device_time_total, reverse=True)
+    total = sum(e.device_time_total for e in events if e.device_type.name == "CUDA")
+    if total == 0:
+        total = sum(e.device_time_total for e in events)
+    conv = sum(
+        e.device_time_total
+        for e in events
+        if e.device_type.name == "CUDA" and any(tag in e.key.lower() for tag in _CONV_TAGS)
+    )
+    transpose = sum(e.device_time_total for e in events if e.device_type.name == "CUDA" and "nhwc" in e.key.lower())
+    print(
+        f"  GPU kernel time: {total / 1e3:.1f} ms  conv-like: {100 * conv / total:.1f}%  "
+        f"nchw<->nhwc transposes: {100 * transpose / total:.1f}%"
+    )
+    print(f"  {'kernel':90s} {'ms':>9s} {'%':>6s} {'calls':>7s}")
+    shown = 0
+    for e in events:
+        if e.device_type.name != "CUDA":
+            continue
+        share = 100 * e.device_time_total / total
+        print(f"  {e.key[:90]:90s} {e.device_time_total / 1e3:9.2f} {share:6.1f} {e.count:7d}")
+        shown += 1
+        if shown >= rows:
+            break
+
+
+def run_level(
+    args: argparse.Namespace,
+    level: str,
+    dtype: torch.dtype,
+    device: torch.device,
+    *,
+    rank: int,
+    parallel: bool,
+) -> tuple[torch.Tensor, dict]:
+    vae = load_vae(args, dtype, device, parallel)
+    report = install_wan_vae_fastpath(vae, level=level)
+    latents = make_latents(vae, args, dtype, device)
+    for _ in range(args.warmup):
+        decode(vae, latents)
+    sync()
+    torch.accelerator.reset_peak_memory_stats()
+    timings = []
+    output = None
+    for _ in range(args.iters):
+        sync()
+        start = time.perf_counter()
+        output = decode(vae, latents)
+        sync()
+        timings.append(time.perf_counter() - start)
+    peak = torch.tensor(torch.accelerator.max_memory_allocated() / 2**30, device=device)
+    if dist.is_initialized():
+        dist.all_reduce(peak, op=dist.ReduceOp.MAX)  # the largest rank decides whether the config fits
+    stats = {
+        "installed": report.installed,
+        "reason": report.reason,
+        "fused_silu": report.fused_silu_dtypes,
+        "patched": dict(report.patched),
+        "time_s": min(timings),
+        "peak_gib": peak.item(),
+    }
+    if args.profile:
+        if rank == 0:
+            print(f"[{level}] profile")
+        profile_decode(vae, latents, args.profile_rows, verbose=rank == 0)
+    assert output is not None
+    del vae
+    torch.accelerator.empty_cache()
+    return output.detach(), stats
+
+
+def main() -> None:
+    args = parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA device required")
+    rank, world_size = init_distributed()
+    parallel = args.vae_patch_parallel_size > 1
+    if parallel and world_size < args.vae_patch_parallel_size:
+        raise SystemExit(
+            f"--vae-patch-parallel-size {args.vae_patch_parallel_size} needs torchrun with at least that many "
+            f"processes (WORLD_SIZE={world_size})"
+        )
+    device = torch.device("cuda", torch.accelerator.current_device_index()) if world_size > 1 else torch.device("cuda")
+    dtype = DTYPES[args.dtype]
+    levels: Iterable[str] = [level.strip() for level in args.fast_path.split(",") if level.strip()]
+    for level in levels:
+        if level not in VAE_FAST_PATH_LEVELS:
+            raise SystemExit(f"unknown fast-path level {level!r}; choose from {VAE_FAST_PATH_LEVELS}")
+    if "off" not in levels:
+        levels = ["off", *levels]
+
+    if rank == 0:
+        parallel_desc = f" vae_parallel={args.vae_parallel_mode}x{args.vae_patch_parallel_size}" if parallel else ""
+        print(
+            f"model={'tiny' if args.tiny else args.model} size={args.size} frames={args.frames} "
+            f"dtype={args.dtype} world_size={world_size}{parallel_desc}"
+        )
+    golden = None
+    baseline_time = None
+    try:
+        for level in levels:
+            output, stats = run_level(args, level, dtype, device, rank=rank, parallel=parallel)
+            if level == "off":
+                golden = output
+                baseline_time = stats["time_s"]
+                if args.save_output and rank == 0:
+                    torch.save(output.cpu(), args.save_output)
+            assert golden is not None
+            if rank != 0:
+                # Parallel decode assembles the video on rank 0 only; the other
+                # ranks just keep the collectives in lockstep.
+                del output
+                continue
+            speedup = baseline_time / stats["time_s"] if baseline_time else float("nan")
+            line = (
+                f"[{level:13s}] decode {stats['time_s'] * 1e3:9.1f} ms  speedup {speedup:5.2f}x  "
+                f"peak {stats['peak_gib']:5.2f} GiB  installed={stats['installed']}"
+            )
+            if level != "off":
+                equal = torch.equal(output, golden)
+                line += (
+                    f"  torch_equal={equal} bits_equal={bits_equal(output, golden)} "
+                    f"max_abs_diff={(output.float() - golden.float()).abs().max().item():.3e} "
+                    f"psnr={psnr(output, golden):.2f} dB"
+                )
+                if stats["installed"]:
+                    line += f"  fused_silu={stats['fused_silu'] or 'off'}"
+                else:
+                    line += f"  reason={stats['reason']}"
+            print(line)
+            del output
+            torch.accelerator.empty_cache()
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

docs/.nav.yml

@@ -98,6 +98,7 @@ nav:
         - Hugging Face Hub: user_guide/diffusion/attention_backends/huggingface_hub.md
         - RainFusion: user_guide/diffusion/attention_backends/rainfusion.md
       - Regional Compilation: user_guide/diffusion/regional_compilation.md
+      - Wan VAE Decoder Fast Path: user_guide/diffusion/vae_fast_path.md
       - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
       - Startup and Loading: user_guide/diffusion/startup_and_loading.md
       - LoRA: user_guide/diffusion/lora.md

docs/features/README.md

@@ -59,6 +59,7 @@ defines the shared configuration and backend extension points.
 | Distribute diffusion work across devices | [Parallelism](../user_guide/diffusion/parallelism/overview.md) | [Parallelism designs](../design/index.md#parallelism) |
 | Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Attention Backend Selection](../design/feature/attention_backend_selection.md) |
 | Compile repeated diffusion regions | [Regional Compilation](../user_guide/diffusion/regional_compilation.md) | User-facing optimization guide |
+| Speed up Wan VAE video decode | [Wan VAE Decoder Fast Path](../user_guide/diffusion/vae_fast_path.md) | User-facing optimization guide |
 | Add generated video frames | [Frame Interpolation](../user_guide/diffusion/frame_interpolation.md) | User-facing extension guide |
 | Reduce diffusion model startup time | [Startup and Loading](../user_guide/diffusion/startup_and_loading.md) | User-facing loading guide |
 | Apply diffusion adapters | [LoRA](../user_guide/diffusion/lora.md) | User-facing extension guide |

docs/user_guide/diffusion/vae_fast_path.md

@@ -0,0 +1,102 @@
+# Wan VAE Decoder Fast Path
+
+Video pipelines built on the Wan 2.1/2.2 causal VAE (Cosmos3, Wan2.2, LingBot,
+Helios, SANA-Video, LongCat-Video, DreamZero) spend a large share of their
+end-to-end time decoding latents to pixels. The decoder is a chunked 3D
+convolutional network whose reference implementation in diffusers moves every
+activation through memory several times per layer (normalization, activation,
+causal padding, feature-cache bookkeeping, shortcut upsampling). vLLM-Omni
+installs a fast path on every loaded Wan VAE that fuses this data movement into
+a handful of Triton kernels while leaving the convolutions themselves untouched.
+
+## Levels
+
+The fast path is controlled by `--vae-fast-path` (engine argument
+`vae_fast_path`, deploy-config key `vae_fast_path`):
+
+| Level | Default | Output vs. diffusers | What it does |
+|---|---|---|---|
+| `lossless` | yes | bit-identical | Fused RMSNorm epilogue, fused causal-conv input assembly and cache refresh, fused shortcut upsampling and residual adds (with the neighbouring convolution biases folded in), single-pass nearest 2x upsampling, preallocated output assembly. |
+| `channels_last` | no | within tolerance (PSNR typically > 60 dB) | Everything in `lossless`, plus decoder convolution weights converted to channels-last memory format so cuDNN picks its channels-last kernels, and a single-pass channels-last RMSNorm+SiLU kernel that also absorbs the bias of the preceding `conv1`. |
+| `off` | no | bit-identical | Reference diffusers decoder. |
+
+```bash
+# Default: bit-exact fast path
+vllm serve nvidia/Cosmos3-Nano --omni
+
+# Fastest: channels-last decoder (not bit-exact)
+vllm serve nvidia/Cosmos3-Nano --omni --vae-fast-path channels_last
+
+# Reference implementation, e.g. to bisect a quality issue
+vllm serve nvidia/Cosmos3-Nano --omni --vae-fast-path off
+```
+
+The equivalent per-stage deploy configuration is:
+
+```yaml
+stages:
+  - stage_id: 0
+    vae_fast_path: channels_last
+```
+
+## Behavior and limitations
+
+- Only CUDA is supported. On other platforms the reference decoder runs.
+- The fast path is installed once per VAE instance when the pipeline is
+  initialized. It rebinds the forwards of the loaded decoder modules; parameter
+  names, `state_dict` keys and weight loading are unchanged.
+- Installation is skipped, with a logged reason, if a custom forward wrapper
+  or forward hook would be replaced or bypassed by the fused paths. Wrappers
+  and hooks on modules that the fast path still calls normally are preserved.
+- Every fused kernel validates its inputs and falls back to the exact PyTorch
+  expression for anything it does not support (unusual dtypes or layouts, CPU
+  tensors, autograd enabled, `torch.compile` tracing).
+- The SiLU activation is folded into the normalization kernel only after an
+  exhaustive self-test over all bf16/fp16 values proves the fused epilogue
+  bit-identical to `F.silu` on the running toolkit; otherwise SiLU stays a
+  separate operation. The startup log reports `fused_silu=...`.
+- Spatially sharded VAE decode (`--vae-parallel-mode spatial_shard_height` or
+  `spatial_shard_width`) is not combined with the fast path; the installer skips
+  the VAE and logs why. Tiled and tile-parallel decode work unchanged.
+- The `channels_last` level changes the order in which cuDNN accumulates
+  convolutions, so outputs differ from the reference in the last bits. Use
+  `lossless` when bitwise reproducibility against diffusers matters.
+- For VAEs kept in fp32, cuDNN's channels-last convolution algorithms run in
+  TF32 under PyTorch's default `torch.backends.cudnn.allow_tf32 = True`, which
+  dominates the difference to the reference (about 1e-3). Set
+  `torch.backends.cudnn.allow_tf32 = False` if full fp32 convolutions are
+  required.
+
+## Measured
+
+Cosmos3-Nano VAE, 1280x720 x 189 frames, bf16, one GB200 GPU, `bench_wan_vae_decode.py`
+(best of 2 runs after warmup):
+
+| `--vae-fast-path` | Decode time | Speedup | Output vs. `off` |
+|---|-------------|---------|---|
+| `off` | 6.00 s      | 1.00x   | reference |
+| `lossless` | 3.26 s      | 1.84x   | bit-identical |
+| `channels_last` | 2.51 s      | 2.39x   | PSNR 62.6 dB, max abs diff 4.6e-2 |
+
+Peak decode memory was unchanged (about 13 GiB): it is set by the largest
+activations and the cuDNN workspace, not by the output assembly.
+
+## Benchmarking
+
+`benchmarks/diffusion/bench_wan_vae_decode.py` decodes seeded latents with the
+real VAE at each level and reports decode time, peak memory, bitwise equality
+and PSNR against the `off` level. `--profile` prints a per-kernel table with the
+convolution share and any layout-transpose kernels:
+
+```bash
+python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
+    --size 1280x720 --frames 189 --fast-path off,lossless,channels_last --profile
+```
+
+The same script benchmarks multi-GPU VAE decode when launched with `torchrun`;
+the decode is timed across all ranks and rank 0 reports:
+
+```bash
+torchrun --nproc-per-node 2 benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
+    --vae-patch-parallel-size 2 --vae-parallel-mode tile
+```

docs/user_guide/diffusion_features.md

@@ -252,6 +252,7 @@ The Diffusion Acceleration navigation groups the remaining guides as follows:
 | Parallelism | [Parallelism Overview](diffusion/parallelism/overview.md) |
 | Attention | [Attention Backends](diffusion/attention_backends.md) |
 | Compilation | [Regional Compilation](diffusion/regional_compilation.md) |
+| VAE decode | [Wan VAE Decoder Fast Path](diffusion/vae_fast_path.md) |
 | Video extension | [Frame Interpolation](diffusion/frame_interpolation.md) |
 | Startup | [Startup and Loading](diffusion/startup_and_loading.md) |
 | Adapters | [LoRA](diffusion/lora.md) |

pyproject.toml

@@ -331,6 +331,7 @@ extend-ignore-identifiers-re = [
     ".*arange",
     ".*MoBA",
     ".*temperal_downsample",
+    ".*temperal_upsample",
     ".*_nd",
     ".*nd$",
     ".*EmbedND$",

recipes/cosmos3/Cosmos3-Nano.md

@@ -129,6 +129,10 @@ To run **without** guardrails (you are responsible for license compliance),
 add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use
 `--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`;
 `--enable-layerwise-offload` reduces VRAM on smaller GPUs;
+`--vae-fast-path channels_last` speeds up the Wan VAE video decode by switching the
+decoder convolutions to channels-last kernels (output no longer bit-identical to
+diffusers; the default `lossless` fast path is bit-exact, see
+[Wan VAE Decoder Fast Path](../../docs/user_guide/diffusion/vae_fast_path.md));
 `--quantization fp8` (online, no calibration) cuts peak VRAM for 720p video
 generation from ~50 GB to ~36 GB with BF16-level quality (T2V composition can
 shift at the same seed). The pipeline

recipes/cosmos3/Cosmos3-Super.md

@@ -58,6 +58,10 @@ vllm serve nvidia/Cosmos3-Super \
 Guardrails are on by default (gated `nvidia/Cosmos-1.0-Guardrail` — `pip install
 cosmos-guardrail`, accept the license, set `HF_TOKEN`); add `--no-guardrails` to
 disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs;
+`--vae-fast-path channels_last` speeds up the Wan VAE video decode by switching the
+decoder convolutions to channels-last kernels (output no longer bit-identical to
+diffusers; the default `lossless` fast path is bit-exact, see
+[Wan VAE Decoder Fast Path](../../docs/user_guide/diffusion/vae_fast_path.md));
 `--quantization fp8` (online, no calibration) cuts peak VRAM for 720p video
 generation from ~83 GB to ~55 GB per GPU (2-GPU) with BF16-level quality (T2V
 composition can shift at the same seed).

tests/diffusion/distributed/test_wan_vae_fastpath_install.py

@@ -0,0 +1,619 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU tests for the Wan VAE decoder fast path installer and its exact PyTorch fallbacks."""
+
+from __future__ import annotations
+
+import importlib
+from functools import wraps
+from types import SimpleNamespace
+
+import pytest
+import torch
+from diffusers.models.autoencoders import AutoencoderKLWan
+from torch import nn
+
+from vllm_omni.diffusion import registry as registry_module
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import (
+    REPORT_ATTR,
+    decode_frames,
+    install_wan_vae_fastpath,
+    is_installed,
+    uninstall_wan_vae_fastpath,
+)
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards as fastpath_forwards
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+TINY_RESIDUAL = dict(
+    base_dim=8,
+    decoder_base_dim=8,
+    z_dim=4,
+    dim_mult=[1, 1],
+    num_res_blocks=1,
+    temperal_downsample=[False, True],
+    is_residual=True,
+)
+TINY_RESIDUAL_PATCH2 = dict(TINY_RESIDUAL, patch_size=2, in_channels=12, out_channels=12)
+TINY_WAN21 = dict(
+    base_dim=8,
+    z_dim=4,
+    dim_mult=[1, 2],
+    num_res_blocks=1,
+    temperal_downsample=[False, True],
+    is_residual=False,
+)
+CONFIGS = {
+    "residual": TINY_RESIDUAL,
+    "residual_patch2": TINY_RESIDUAL_PATCH2,
+    "wan21": TINY_WAN21,
+}
+
+
+def _build_pair(config: dict, dtype: torch.dtype) -> tuple[AutoencoderKLWan, AutoencoderKLWan]:
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**config).eval().to(dtype)
+    candidate = AutoencoderKLWan(**config).eval().to(dtype)
+    candidate.load_state_dict(reference.state_dict())
+    return reference, candidate
+
+
+@pytest.mark.parametrize("config_name", sorted(CONFIGS))
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("frames", [1, 3])
+def test_lossless_fallback_paths_are_bitwise_exact(config_name: str, dtype: torch.dtype, frames: int) -> None:
+    """Without CUDA every kernel declines, so this exercises the restructured PyTorch paths."""
+    config = CONFIGS[config_name]
+    reference, candidate = _build_pair(config, dtype)
+    report = install_wan_vae_fastpath(candidate, level="lossless")
+    assert report.installed, report
+    assert report.fused_silu_dtypes == ()
+
+    torch.manual_seed(1)
+    latents = torch.randn(1, config["z_dim"], frames, 6, 8).to(dtype)
+    with torch.no_grad():
+        expected = reference.decode(latents, return_dict=False)[0]
+        actual = candidate.decode(latents, return_dict=False)[0]
+    assert actual.stride() == expected.stride()
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.parametrize("config_name", sorted(CONFIGS))
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_decode_frames_matches_reference_decode(config_name: str, dtype: torch.dtype) -> None:
+    config = CONFIGS[config_name]
+    reference, candidate = _build_pair(config, dtype)
+    install_wan_vae_fastpath(candidate)
+    for frames in (1, 2, 4):
+        torch.manual_seed(2)
+        latents = torch.randn(1, config["z_dim"], frames, 6, 8).to(dtype)
+        with torch.no_grad():
+            expected = reference._decode(latents, return_dict=False)[0]
+            actual = decode_frames(candidate, latents)
+        assert actual.stride() == expected.stride()
+        assert torch.equal(actual, expected)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_omni_wan_vae_decode_override_is_bitwise_exact(dtype: torch.dtype) -> None:
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**TINY_RESIDUAL_PATCH2).eval().to(dtype)
+    candidate = OmniAutoencoderKLWan(**TINY_RESIDUAL_PATCH2).eval().to(dtype)
+    candidate.load_state_dict(reference.state_dict())
+    latents = torch.randn(1, 4, 3, 6, 8).to(dtype)
+    with torch.no_grad():
+        expected = reference.decode(latents, return_dict=False)[0]
+        before_install = candidate.decode(latents, return_dict=False)[0]
+        install_wan_vae_fastpath(candidate)
+        actual = candidate.decode(latents, return_dict=False)[0]
+        as_dict = candidate.decode(latents).sample
+    # The Omni wrapper runs bf16 decode under autocast (also on CPU), so compare
+    # the override against the same wrapper before installation, and against the
+    # plain diffusers decode only where no autocast is involved.
+    assert actual.stride() == before_install.stride()
+    assert torch.equal(actual, before_install)
+    assert torch.equal(as_dict, before_install)
+    if dtype is torch.float32:
+        assert torch.equal(before_install, expected)
+
+    # Tiling still dispatches to the diffusers implementation.
+    candidate.use_tiling = True
+    candidate.tile_sample_min_height = candidate.tile_sample_min_width = 16
+    candidate.tile_sample_stride_height = candidate.tile_sample_stride_width = 16
+    reference.use_tiling = True
+    reference.tile_sample_min_height = reference.tile_sample_min_width = 16
+    reference.tile_sample_stride_height = reference.tile_sample_stride_width = 16
+    with torch.no_grad():
+        expected_tiled = reference.decode(latents, return_dict=False)[0]
+        actual_tiled = candidate.decode(latents, return_dict=False)[0]
+    if dtype is torch.float32:
+        assert torch.equal(actual_tiled, expected_tiled)
+    assert actual_tiled.shape == expected_tiled.shape
+
+
+def test_omni_wan_vae_decode_preserves_autograd() -> None:
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**TINY_RESIDUAL).eval()
+    candidate = OmniAutoencoderKLWan(**TINY_RESIDUAL).eval()
+    candidate.load_state_dict(reference.state_dict())
+    install_wan_vae_fastpath(candidate)
+
+    reference_latents = torch.randn(1, 4, 2, 6, 8, requires_grad=True)
+    candidate_latents = reference_latents.detach().clone().requires_grad_()
+    expected = reference.decode(reference_latents, return_dict=False)[0]
+    actual = candidate.decode(candidate_latents, return_dict=False)[0]
+
+    torch.testing.assert_close(actual, expected)
+    expected.square().mean().backward()
+    actual.square().mean().backward()
+    torch.testing.assert_close(candidate_latents.grad, reference_latents.grad)
+
+
+def test_install_is_idempotent_and_reversible() -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    keys_before = set(vae.state_dict())
+    report = install_wan_vae_fastpath(vae)
+    assert report.installed and is_installed(vae)
+    assert report.patched["WanDecoder3d"] == 1
+    assert report.patched["WanCausalConv3d"] > 1
+    assert report.patched["WanRMS_norm"] > 1
+    assert set(vae.state_dict()) == keys_before
+    assert "forward" in vae.decoder.__dict__
+
+    assert install_wan_vae_fastpath(vae) is report
+    assert install_wan_vae_fastpath(vae, level="channels_last") is report  # warns, keeps the first level
+
+    uninstall_wan_vae_fastpath(vae)
+    assert not is_installed(vae)
+    assert "forward" not in vae.decoder.__dict__
+    assert all("forward" not in module.__dict__ for module in vae.decoder.modules())
+    assert all(fastpath_forwards.CFG_ATTR not in module.__dict__ for module in vae.decoder.modules())
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("failure_stage", ["binding", "conversion", "report"])
+def test_failed_install_restores_original_state(failure_stage: str, monkeypatch) -> None:
+    installer = importlib.import_module(install_wan_vae_fastpath.__module__)
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    # Preserve mixed original layouts, an existing forward, gradients and a
+    # buffer: Module.to can replace tensors as well as change their storage.
+    vae.decoder.conv_in.to(memory_format=torch.channels_last_3d)
+    vae.decoder.forward = vae.decoder.forward
+    setattr(vae.decoder, fastpath_forwards.CFG_ATTR, object())
+    vae.decoder.conv_out.weight.grad = torch.randn_like(vae.decoder.conv_out.weight)
+    vae.post_quant_conv.register_buffer("rollback_probe", torch.randn(2, 3, 4, 5, 6))
+    attributes = {
+        module: {
+            name: module.__dict__[name] for name in ("forward", fastpath_forwards.CFG_ATTR) if name in module.__dict__
+        }
+        for module in vae.modules()
+    }
+    parameters = dict(vae.named_parameters())
+    buffers = dict(vae.named_buffers())
+    gradients = {name: parameter.grad for name, parameter in parameters.items()}
+    tensors = [*parameters.values(), *buffers.values(), *(grad for grad in gradients.values() if grad is not None)]
+    originals = [(tensor, tensor.data_ptr(), tensor.stride(), tensor.clone()) for tensor in tensors]
+    latents = torch.randn(1, 4, 2, 6, 8)
+    expected = vae.decode(latents, return_dict=False)[0]
+
+    with monkeypatch.context() as patch:
+        if failure_stage == "binding":
+            conv_type = type(vae.post_quant_conv)
+            original_setattr = conv_type.__setattr__
+
+            def fail_binding(module, name, value):
+                original_setattr(module, name, value)
+                if module is vae.post_quant_conv and name == fastpath_forwards.CFG_ATTR:
+                    raise RuntimeError("injected installation failure")
+
+            patch.setattr(conv_type, "__setattr__", fail_binding)
+        elif failure_stage == "conversion":
+            original_to = vae.post_quant_conv.to
+
+            def fail_conversion(*args, **kwargs):
+                original_to(*args, **kwargs)
+                assert vae.decoder.conv_out.weight.is_contiguous(memory_format=torch.channels_last_3d)
+                raise torch.OutOfMemoryError("injected installation failure")
+
+            patch.setattr(vae.post_quant_conv, "to", fail_conversion)
+        else:
+            original_info = installer.logger.info
+
+            def fail_report(message, *args, **kwargs):
+                if message.startswith("Wan VAE fast path (%s) installed"):
+                    assert is_installed(vae)
+                    raise RuntimeError("injected installation failure")
+                original_info(message, *args, **kwargs)
+
+            patch.setattr(installer.logger, "info", fail_report)
+
+        with pytest.raises(RuntimeError, match="injected installation failure"):
+            install_wan_vae_fastpath(vae, level="channels_last")
+
+    assert not is_installed(vae)
+    assert not hasattr(vae, REPORT_ATTR)
+    assert not hasattr(vae, installer._UNDO_ATTR)
+    for module, original in attributes.items():
+        actual = {
+            name: module.__dict__[name] for name in ("forward", fastpath_forwards.CFG_ATTR) if name in module.__dict__
+        }
+        assert actual == original
+    assert all(dict(vae.named_parameters())[name] is parameter for name, parameter in parameters.items())
+    assert all(dict(vae.named_buffers())[name] is buffer for name, buffer in buffers.items())
+    assert all(parameters[name].grad is grad for name, grad in gradients.items())
+    for tensor, pointer, strides, values in originals:
+        assert tensor.data_ptr() == pointer
+        assert tensor.stride() == strides
+        assert torch.equal(tensor, values)
+    uninstall_wan_vae_fastpath(vae)  # Safe and unnecessary after automatic rollback.
+    assert torch.equal(vae.decode(latents, return_dict=False)[0], expected)
+    assert install_wan_vae_fastpath(vae, level="channels_last").installed
+    uninstall_wan_vae_fastpath(vae)
+
+
+def test_load_state_dict_after_install_updates_patched_modules() -> None:
+    reference, candidate = _build_pair(TINY_RESIDUAL, torch.float32)
+    install_wan_vae_fastpath(candidate)
+    with torch.no_grad():
+        for parameter in reference.parameters():
+            parameter.mul_(0.5)
+    candidate.load_state_dict(reference.state_dict())
+    latents = torch.randn(1, 4, 2, 6, 8)
+    with torch.no_grad():
+        expected = reference.decode(latents, return_dict=False)[0]
+        actual = candidate.decode(latents, return_dict=False)[0]
+    assert torch.equal(actual, expected)
+
+
+def test_channels_last_level_converts_conv_weights_and_restores_them() -> None:
+    reference, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    report = install_wan_vae_fastpath(vae, level="channels_last")
+    assert report.installed and report.channels_last
+    conv3d = [m for m in vae.decoder.modules() if isinstance(m, nn.Conv3d)]
+    conv2d = [m for m in vae.decoder.modules() if isinstance(m, nn.Conv2d)]
+    assert conv3d and conv2d
+    assert all(m.weight.is_contiguous(memory_format=torch.channels_last_3d) for m in conv3d)
+    assert all(m.weight.is_contiguous(memory_format=torch.channels_last) for m in conv2d)
+    assert vae.post_quant_conv.weight.is_contiguous(memory_format=torch.channels_last_3d)
+    latents = torch.randn(1, 4, 2, 6, 8)
+    with torch.no_grad():
+        expected = reference.decode(latents, return_dict=False)[0]
+        out = vae.decode(latents, return_dict=False)[0]
+    # One temporal upsampler: 1 + 2 output frames for 2 latent frames.
+    assert expected.shape == (1, 3, 3, 12, 16)
+    assert out.shape == expected.shape
+    torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-4)
+
+    uninstall_wan_vae_fastpath(vae)
+    assert all(m.weight.is_contiguous() for m in conv3d + conv2d)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("assign", [False, True])
+def test_uninstall_restores_mixed_layouts_and_keeps_weight_updates(assign: bool) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    convs = [module for module in vae.decoder.modules() if isinstance(module, (nn.Conv2d, nn.Conv3d))]
+    convs.append(vae.post_quant_conv)
+    for index, conv in enumerate(convs):
+        if index % 3 == 1:
+            conv.to(memory_format=torch.channels_last_3d if isinstance(conv, nn.Conv3d) else torch.channels_last)
+        elif index % 3 == 2:
+            conv.weight.data = conv.weight.data.transpose(-1, -2)
+    vae.post_quant_conv.register_buffer("layout_probe", torch.randn(2, 3, 4, 5, 6).transpose(-1, -2))
+    vae.decoder.conv_out.weight.grad = torch.randn_like(vae.decoder.conv_out.weight)
+    grad_strides = vae.decoder.conv_out.weight.grad.stride()
+    strides = {name: tensor.stride() for name, tensor in vae.state_dict().items()}
+    # An explicit alias is safe to replace, but must remain an instance
+    # attribute after uninstall, with its exact original identity.
+    original_forward = vae.decoder.forward
+    vae.decoder.forward = original_forward
+    original_cfg = object()
+    setattr(vae.decoder, fastpath_forwards.CFG_ATTR, original_cfg)
+
+    assert install_wan_vae_fastpath(vae, level="channels_last").installed
+    updated = {name: tensor.contiguous() + 0.25 for name, tensor in vae.state_dict().items()}
+    vae.load_state_dict(updated, assign=assign)
+    uninstall_wan_vae_fastpath(vae)
+
+    assert vae.decoder.__dict__["forward"] is original_forward
+    assert getattr(vae.decoder, fastpath_forwards.CFG_ATTR) is original_cfg
+    for name, tensor in vae.state_dict().items():
+        assert tensor.stride() == strides[name], name
+        assert torch.equal(tensor, updated[name]), name
+    if not assign:
+        assert vae.decoder.conv_out.weight.grad.stride() == grad_strides
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("level", ["lossless", "channels_last"])
+@pytest.mark.parametrize(
+    "target",
+    [
+        "decoder",
+        "decoder.conv_in",
+        "decoder.up_blocks.0.upsampler",
+        "decoder.up_blocks.0.upsampler.resample",
+        "decoder.up_blocks.0.upsampler.resample.1",
+        "decoder.up_blocks.0.avg_shortcut",
+        "decoder.norm_out",
+        "decoder.nonlinearity",
+        "post_quant_conv",
+    ],
+)
+def test_installer_declines_forward_wrappers_it_would_bypass(target: str, level: str) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    module = vae.get_submodule(target)
+    original_forward = module.forward
+    calls = []
+
+    @wraps(original_forward)
+    def wrapped(*args, **kwargs):
+        calls.append(True)
+        return original_forward(*args, **kwargs)
+
+    module.forward = wrapped
+    originals = {name: (tensor.data_ptr(), tensor.stride()) for name, tensor in vae.state_dict().items()}
+    report = install_wan_vae_fastpath(vae, level=level)
+    assert not report.installed and target in report.reason and "custom forward" in report.reason
+    assert not is_installed(vae) and not hasattr(vae, REPORT_ATTR)
+    assert all(fastpath_forwards.CFG_ATTR not in child.__dict__ for child in vae.modules())
+    assert {name: (tensor.data_ptr(), tensor.stride()) for name, tensor in vae.state_dict().items()} == originals
+    uninstall_wan_vae_fastpath(vae)
+    assert module.forward is wrapped
+    vae.decode(torch.randn(1, 4, 2, 6, 8))
+    assert calls
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("pre_hook", [False, True])
+def test_installer_declines_convolution_hooks_it_would_bypass(pre_hook: bool) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    conv = vae.decoder.conv_in
+    calls = []
+
+    def hook(*args):
+        calls.append(True)
+
+    handle = conv.register_forward_pre_hook(hook) if pre_hook else conv.register_forward_hook(hook)
+    try:
+        report = install_wan_vae_fastpath(vae, level="channels_last")
+        assert not report.installed and "decoder.conv_in" in report.reason and "forward hooks" in report.reason
+        vae.decode(torch.randn(1, 4, 2, 6, 8))
+        assert calls
+    finally:
+        handle.remove()
+
+
+@torch.no_grad()
+def test_installer_preserves_wrappers_and_hooks_on_modules_it_still_calls() -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    original_forward = vae.decoder.mid_block.forward
+    wrapper_calls = []
+    hook_calls = []
+
+    def wrapped(*args, **kwargs):
+        wrapper_calls.append(True)
+        return original_forward(*args, **kwargs)
+
+    def hook(*args):
+        hook_calls.append(True)
+
+    vae.decoder.mid_block.forward = wrapped
+    handle = vae.decoder.register_forward_hook(hook)
+    try:
+        latents = torch.randn(1, 4, 2, 6, 8)
+        expected = vae.decode(latents).sample
+        wrapper_calls.clear()
+        hook_calls.clear()
+        assert install_wan_vae_fastpath(vae).installed
+        actual = vae.decode(latents).sample
+        assert torch.equal(actual, expected)
+        assert len(wrapper_calls) == len(hook_calls) == 2
+        uninstall_wan_vae_fastpath(vae)
+        assert vae.decoder.mid_block.forward is wrapped
+    finally:
+        handle.remove()
+
+
+def test_installer_refuses_unsupported_targets() -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+
+    assert install_wan_vae_fastpath(vae, level="off").installed is False
+    assert not is_installed(vae)
+
+    vae._vllm_omni_wan_spatial_shard_installed = True
+    report = install_wan_vae_fastpath(vae)
+    assert not report.installed and "spatial-shard" in report.reason
+    del vae._vllm_omni_wan_spatial_shard_installed
+
+    vae.distributed_executor = SimpleNamespace(parallel_mode="spatial_shard_height")
+    report = install_wan_vae_fastpath(vae)
+    assert not report.installed and "spatial_shard_height" in report.reason
+    del vae.distributed_executor
+
+    report = install_wan_vae_fastpath(nn.Linear(2, 2))
+    assert not report.installed and "AutoencoderKLWan" in report.reason
+
+    with pytest.raises(ValueError, match="vae_fast_path"):
+        install_wan_vae_fastpath(vae, level="bogus")
+
+
+@pytest.mark.parametrize(("batch", "frames"), [(1, 1), (1, 2), (2, 1)])
+def test_resample_views_keep_channels_last_recognizable(batch: int, frames: int) -> None:
+    """Regression: ``reshape`` gives a size-1 batch dim a stride the layout heuristic rejects."""
+    import torch.nn.functional as F
+
+    x = torch.randn(batch, 8, frames, 6, 10).contiguous(memory_format=torch.channels_last_3d)
+    merged = fastpath_forwards._merge_batch_and_frames(x)
+    reference = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, 8, 6, 10)
+    assert merged.data_ptr() == x.data_ptr()
+    assert torch.equal(merged, reference)
+    upsampled = F.interpolate(merged, scale_factor=(2.0, 2.0), mode="nearest-exact")
+    assert upsampled.stride(1) == 1, "nearest upsample must keep channels_last for the following Conv2d"
+    assert torch.equal(upsampled, F.interpolate(reference.contiguous(), scale_factor=(2.0, 2.0), mode="nearest-exact"))
+
+    split = fastpath_forwards._split_batch_and_frames(upsampled, batch, frames)
+    assert split.shape == (batch, 8, frames, 12, 20)
+    assert split.data_ptr() == upsampled.data_ptr()
+    assert split.is_contiguous(memory_format=torch.channels_last_3d)
+    assert torch.equal(split, upsampled.view(batch, frames, 8, 12, 20).permute(0, 2, 1, 3, 4))
+
+    plain = torch.randn(batch, 8, frames, 6, 10)
+    assert torch.equal(
+        fastpath_forwards._merge_batch_and_frames(plain), plain.permute(0, 2, 1, 3, 4).reshape(-1, 8, 6, 10)
+    )
+
+
+@torch.no_grad()
+def test_pending_conv_bias_is_added_exactly_when_no_kernel_consumes_it() -> None:
+    """Without the channels-last kernel the un-added conv bias is applied with ATen's rounding."""
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanRMS_norm
+
+    torch.manual_seed(0)
+    for dtype in (torch.float32, torch.bfloat16):
+        norm = WanRMS_norm(8, images=False).to(dtype)
+        act = nn.SiLU()
+        x = torch.randn(1, 8, 2, 4, 4).to(dtype)
+        bias = torch.randn(8).to(dtype)
+        biased = x.clone()
+        biased.add_(bias.view(1, -1, 1, 1, 1))
+        expected = act(norm(biased))
+        assert torch.equal(fastpath_forwards._norm_act(norm, act, x, pending_bias=bias), expected)
+        setattr(norm, fastpath_forwards.CFG_ATTR, fastpath_forwards.FastPathConfig(channels_last=True))
+        assert torch.equal(fastpath_forwards._norm_act(norm, act, x, pending_bias=bias), expected)
+        out = fastpath_forwards.rms_norm_fastpath(norm, x, bias=bias)
+        assert out is not None and torch.equal(out, norm(biased))
+
+
+@torch.no_grad()
+def test_resample_return_bias_is_none_without_kernels() -> None:
+    """On CPU the Conv2d keeps its bias, so the up block adds nothing twice."""
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanResample, WanResidualUpBlock
+
+    torch.manual_seed(0)
+    resample = WanResample(8, mode="upsample2d").eval()
+    assert fastpath_forwards._is_upsample_conv_pair(resample.resample)
+    assert not fastpath_forwards._is_upsample_conv_pair(WanResample(8, mode="downsample2d").resample)
+    x = torch.randn(1, 8, 2, 6, 10)
+    out, bias = fastpath_forwards.resample_forward(resample, x, [None], [0], return_bias=True)
+    assert bias is None
+    assert torch.equal(out, WanResample.forward(resample, x, feat_cache=[None], feat_idx=[0]))
+
+    # First chunk of a temporal up block: one input frame, the time conv is skipped ("Rep").
+    block = WanResidualUpBlock(8, 8, num_res_blocks=1, temperal_upsample=True, up_flag=True).eval()
+    for module in block.modules():
+        setattr(module, fastpath_forwards.CFG_ATTR, fastpath_forwards.FastPathConfig())
+    x = torch.randn(1, 8, 1, 6, 10)
+    cache_len = 6
+    expected = WanResidualUpBlock.forward(block, x, feat_cache=[None] * cache_len, feat_idx=[0], first_chunk=True)
+    actual = fastpath_forwards.residual_up_block_forward(
+        block, x, feat_cache=[None] * cache_len, feat_idx=[0], first_chunk=True
+    )
+    assert torch.equal(actual, expected)
+
+
+def test_upsample_forward_only_fuses_nearest_2x() -> None:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanUpsample
+
+    assert fastpath_forwards._is_nearest_2x(WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"))
+    assert fastpath_forwards._is_nearest_2x(WanUpsample(scale_factor=2, mode="nearest"))
+    assert not fastpath_forwards._is_nearest_2x(WanUpsample(scale_factor=(2.0, 3.0), mode="nearest-exact"))
+    assert not fastpath_forwards._is_nearest_2x(WanUpsample(scale_factor=(2.0, 2.0), mode="bilinear"))
+    assert not fastpath_forwards._is_nearest_2x(WanUpsample(size=(12, 20), mode="nearest-exact"))
+
+    module = WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact")
+    for dtype in (torch.float32, torch.bfloat16):
+        x = torch.randn(3, 8, 6, 10).to(dtype)
+        for layout in (torch.contiguous_format, torch.channels_last):
+            x = x.contiguous(memory_format=layout)
+            expected = WanUpsample.forward(module, x)
+            actual = fastpath_forwards.upsample_forward(module, x)
+            assert actual.stride() == expected.stride()
+            assert torch.equal(actual, expected)
+
+
+def test_rms_norm_vae_substitute_is_not_matched() -> None:
+    from vllm_omni.diffusion.layers.norm import RMSNormVAE
+
+    assert not fastpath_forwards.is_diffusers_rms_norm(RMSNormVAE(8, images=False))
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    norm = vae.decoder.norm_out
+    assert fastpath_forwards.is_diffusers_rms_norm(norm)
+
+
+@torch.no_grad()
+def test_rms_norm_fastpath_declines_tensor_bias_and_dtype_mismatch() -> None:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanRMS_norm
+
+    norm = WanRMS_norm(8, images=False, bias=True)
+    assert fastpath_forwards.rms_norm_fastpath(norm, torch.randn(1, 8, 1, 4, 4)) is None
+
+    norm = WanRMS_norm(8, images=False)
+    assert fastpath_forwards.rms_norm_fastpath(norm, torch.randn(1, 8, 1, 4, 4).to(torch.bfloat16)) is None
+
+    x = torch.randn(1, 8, 2, 4, 4)
+    out = fastpath_forwards.rms_norm_fastpath(norm, x)
+    assert out is not None and torch.equal(out, norm(x))
+
+
+def test_rms_norm_forward_preserves_autograd() -> None:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanRMS_norm
+
+    torch.manual_seed(0)
+    reference = WanRMS_norm(8, images=False)
+    candidate = WanRMS_norm(8, images=False)
+    candidate.load_state_dict(reference.state_dict())
+    reference_input = torch.randn(1, 8, 1, 4, 4, requires_grad=True)
+    candidate_input = reference_input.detach().clone().requires_grad_()
+
+    expected = reference(reference_input)
+    assert fastpath_forwards.rms_norm_fastpath(candidate, candidate_input) is None
+    actual = fastpath_forwards.rms_norm_forward(candidate, candidate_input)
+    torch.testing.assert_close(actual, expected)
+
+    expected.square().mean().backward()
+    actual.square().mean().backward()
+    torch.testing.assert_close(candidate_input.grad, reference_input.grad)
+    torch.testing.assert_close(candidate.gamma.grad, reference.gamma.grad)
+
+
+def test_omni_diffusion_config_validates_vae_fast_path() -> None:
+    assert OmniDiffusionConfig(model="x").vae_fast_path == "lossless"
+    assert OmniDiffusionConfig(model="x", vae_fast_path="channels_last").vae_fast_path == "channels_last"
+    with pytest.raises(ValueError, match="vae_fast_path"):
+        OmniDiffusionConfig(model="x", vae_fast_path="fast")
+
+
+class _StubPipeline(nn.Module):
+    def __init__(self, vae: nn.Module) -> None:
+        super().__init__()
+        self.vae = vae
+
+
+@pytest.mark.parametrize(("level", "expected"), [("lossless", True), ("channels_last", True), ("off", False)])
+def test_registry_hook_installs_on_cuda_platform(mocker, level: str, expected: bool) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    platform = mocker.Mock()
+    platform.is_cuda.return_value = True
+    mocker.patch.object(registry_module, "current_omni_platform", platform)
+
+    registry_module._apply_wan_vae_fastpath_if_enabled(_StubPipeline(vae), SimpleNamespace(vae_fast_path=level))
+    assert is_installed(vae) is expected
+    if expected:
+        assert getattr(vae, REPORT_ATTR).level == level
+
+
+def test_registry_hook_skips_non_cuda_platform_and_non_wan_vaes(mocker) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    platform = mocker.Mock()
+    platform.is_cuda.return_value = False
+    mocker.patch.object(registry_module, "current_omni_platform", platform)
+    registry_module._apply_wan_vae_fastpath_if_enabled(_StubPipeline(vae), SimpleNamespace(vae_fast_path="lossless"))
+    assert not is_installed(vae)
+
+    platform.is_cuda.return_value = True
+    other = nn.Linear(2, 2)
+    registry_module._apply_wan_vae_fastpath_if_enabled(_StubPipeline(other), SimpleNamespace(vae_fast_path="lossless"))
+    assert not hasattr(other, REPORT_ATTR)

tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py

@@ -0,0 +1,903 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""GPU bit-exactness tests for the Wan VAE decoder fast path kernels and forwards."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn.functional as F
+from diffusers.models.autoencoders import AutoencoderKLWan
+from diffusers.models.autoencoders.autoencoder_kl_wan import (
+    CACHE_T,
+    WanCausalConv3d,
+    WanResample,
+    WanResidualBlock,
+    WanRMS_norm,
+    WanUpsample,
+)
+from torch import nn
+
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards as fp
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import (
+    install_wan_vae_fastpath,
+)
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_data_movement as dm
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_rms_norm as rn
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_rms_norm_cl as cl
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_upsample as up
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.diffusion,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+LOW_PRECISION = [torch.bfloat16, torch.float16]
+ALL_DTYPES = [torch.bfloat16, torch.float16, torch.float32]
+# Cosmos3 720p decoder stages, spatially scaled down so the suite stays fast.
+NORM_SHAPES = [
+    (1, 1024, 1, 45, 80),
+    (1, 1024, 2, 45, 80),
+    (1, 512, 4, 45, 80),
+    (1, 256, 4, 90, 160),
+]
+
+
+def _bits_equal(a: torch.Tensor, b: torch.Tensor) -> bool:
+    int_dtype = torch.int16 if a.element_size() == 2 else torch.int32
+    return a.shape == b.shape and bool((a.view(int_dtype) == b.view(int_dtype)).all())
+
+
+def _same_dense_strides(a: torch.Tensor, b: torch.Tensor) -> bool:
+    """Stride equality ignoring size-1 dims, whose stride carries no information."""
+    return a.shape == b.shape and all(
+        sa == sb for sa, sb, n in zip(a.stride(), b.stride(), a.shape, strict=True) if n > 1
+    )
+
+
+def _with_negative_zeros(x: torch.Tensor) -> torch.Tensor:
+    flat = x.view(-1)
+    flat[:: max(1, flat.numel() // 97)] = -0.0
+    return x
+
+
+def _reference_rms_norm(norm: WanRMS_norm, x: torch.Tensor) -> torch.Tensor:
+    return WanRMS_norm.forward(norm, x)
+
+
+def _make_norm(channels: int, dtype: torch.dtype, images: bool = False) -> WanRMS_norm:
+    norm = WanRMS_norm(channels, images=images).to(device="cuda", dtype=dtype)
+    with torch.no_grad():
+        norm.gamma.copy_(torch.randn_like(norm.gamma) * 0.5 + 1.0)
+    return norm
+
+
+# --------------------------------------------------------------------------- #
+# RMSNorm
+# --------------------------------------------------------------------------- #
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("shape", NORM_SHAPES)
+def test_vector_norm_from_low_precision_matches_fp32_normalize_denominator(dtype: torch.dtype, shape) -> None:
+    """The fast path relies on ATen running the same reduction for both inputs."""
+    torch.manual_seed(0)
+    x = torch.randn(shape, device="cuda").to(dtype)
+    expected = x.float().norm(2, dim=1, keepdim=True)
+    actual = torch.linalg.vector_norm(x, dim=1, keepdim=True, dtype=torch.float32)
+    assert torch.equal(actual, expected)
+    frame_major = x.permute(0, 2, 1, 3, 4).contiguous().permute(0, 2, 1, 3, 4)
+    expected = frame_major.float().norm(2, dim=1, keepdim=True)
+    actual = torch.linalg.vector_norm(frame_major, dim=1, keepdim=True, dtype=torch.float32)
+    assert torch.equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("shape", NORM_SHAPES)
+@pytest.mark.parametrize("layout", ["contiguous", "frame_major"])
+@pytest.mark.parametrize("autocast", [False, True])
+def test_rms_norm_fastpath_is_bitwise_exact(dtype: torch.dtype, shape, layout: str, autocast: bool) -> None:
+    torch.manual_seed(0)
+    norm = _make_norm(shape[1], dtype)
+    x = _with_negative_zeros(torch.randn(shape, device="cuda").to(dtype))
+    if layout == "frame_major":
+        x = x.permute(0, 2, 1, 3, 4).contiguous().permute(0, 2, 1, 3, 4)
+    with torch.autocast("cuda", dtype=dtype if dtype is not torch.float32 else torch.bfloat16, enabled=autocast):
+        expected = _reference_rms_norm(norm, x)
+        actual = fp.rms_norm_fastpath(norm, x)
+    assert actual is not None
+    assert actual.dtype == expected.dtype
+    assert actual.stride() == expected.stride()
+    assert _bits_equal(actual, expected), "RMSNorm fast path is not bit-exact (signed zeros included)"
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_rms_norm_fastpath_handles_attention_block_4d_input(dtype: torch.dtype) -> None:
+    torch.manual_seed(0)
+    norm = _make_norm(1024, dtype, images=True)
+    x = torch.randn(2, 1024, 45, 80, device="cuda").to(dtype)
+    expected = _reference_rms_norm(norm, x)
+    actual = fp.rms_norm_fastpath(norm, x)
+    assert actual is not None and _bits_equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", LOW_PRECISION)
+def test_fused_silu_epilogue_matches_silu_of_reference(dtype: torch.dtype) -> None:
+    device = torch.device("cuda")
+    if not rn.silu_epilogue_is_exact(device, dtype):
+        pytest.skip(f"fused SiLU epilogue is not bit-exact for {dtype} on this toolkit; fusion stays disabled")
+    torch.manual_seed(0)
+    norm = _make_norm(256, dtype)
+    x = _with_negative_zeros(torch.randn(1, 256, 4, 90, 160, device="cuda").to(dtype))
+    expected = F.silu(_reference_rms_norm(norm, x))
+    actual = fp.rms_norm_fastpath(norm, x, silu=True)
+    assert actual is not None and _bits_equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("shape", [(1, 1024, 2, 45, 80), (1, 256, 4, 90, 160), (2, 512, 45, 80)])
+@pytest.mark.parametrize("silu", [False, True])
+def test_channels_last_rms_norm_is_close_and_keeps_layout(dtype: torch.dtype, shape, silu: bool) -> None:
+    torch.manual_seed(0)
+    norm = _make_norm(shape[1], dtype, images=len(shape) == 4)
+    memory_format = torch.channels_last_3d if len(shape) == 5 else torch.channels_last
+    x = torch.randn(shape, device="cuda").to(dtype).contiguous(memory_format=memory_format)
+    expected = _reference_rms_norm(norm, x)
+    if silu:
+        expected = F.silu(expected)
+    actual = cl.rms_norm_channels_last(x, norm.gamma, norm.scale, silu=silu)
+    assert actual is not None
+    assert actual.is_contiguous(memory_format=memory_format)
+    assert actual.dtype == expected.dtype
+    # Fast-math kernel: approximate divide/exp and a single final rounding.
+    if dtype is torch.float32:
+        torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5)
+    else:
+        torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+    # Declines channels-first input and too many channels.
+    assert cl.rms_norm_channels_last(x.contiguous(), norm.gamma, norm.scale) is None
+    wide = torch.randn(1, 2048, 1, 4, 4, device="cuda", dtype=dtype).contiguous(memory_format=torch.channels_last_3d)
+    assert cl.rms_norm_channels_last(wide, torch.ones(2048, device="cuda", dtype=dtype), 1.0) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("silu", [False, True])
+def test_channels_last_rms_norm_absorbs_conv_bias(dtype: torch.dtype, silu: bool) -> None:
+    torch.manual_seed(0)
+    shape = (1, 256, 4, 90, 160)
+    norm = _make_norm(shape[1], dtype)
+    x = torch.randn(shape, device="cuda").to(dtype).contiguous(memory_format=torch.channels_last_3d)
+    bias = (torch.randn(shape[1], device="cuda") * 0.1).to(dtype)
+    # Reference: ATen adds the conv bias in place (one rounding) before the norm.
+    expected = _reference_rms_norm(norm, x + bias.view(1, -1, 1, 1, 1))
+    if silu:
+        expected = F.silu(expected)
+    actual = cl.rms_norm_channels_last(x, norm.gamma, norm.scale, silu=silu, bias=bias)
+    assert actual is not None
+    assert actual.is_contiguous(memory_format=torch.channels_last_3d)
+    if dtype is torch.float32:
+        torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5)
+    else:
+        torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+    assert cl.rms_norm_channels_last(x, norm.gamma, norm.scale, bias=bias[:-1]) is None
+    assert cl.rms_norm_channels_last(x, norm.gamma, norm.scale, bias=bias.view(1, -1)) is None
+    assert cl.rms_norm_channels_last(x, norm.gamma, norm.scale, bias=bias.cpu()) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_channels_last_residual_block_folds_conv1_bias(dtype: torch.dtype, monkeypatch) -> None:
+    """The conv1 bias is left un-added and consumed by norm2's kernel (rounded like ATen); output stays close."""
+    torch.manual_seed(0)
+    block = WanResidualBlock(64, 64, dropout=0.0).eval().to(device="cuda", dtype=dtype)
+    with torch.no_grad():
+        for conv in (block.conv1, block.conv2):
+            conv.bias.copy_(torch.randn_like(conv.bias) * 0.1)
+    x = torch.randn(1, 64, 4, 24, 40, device="cuda", dtype=dtype)
+    expected = WanResidualBlock.forward(block, x, feat_cache=[None, None], feat_idx=[0])
+
+    for conv in (block.conv1, block.conv2):  # as the installer does: conv weights only
+        conv.to(memory_format=torch.channels_last_3d)
+    cfg = fp.FastPathConfig(channels_last=True)
+    for module in (block, block.norm1, block.norm2):
+        setattr(module, fp.CFG_ATTR, cfg)
+    seen: list[torch.Tensor | None] = []
+    original = fp.rms_norm_fastpath
+
+    def spy(norm, x, *, silu=False, bias=None):
+        seen.append(bias)
+        return original(norm, x, silu=silu, bias=bias)
+
+    monkeypatch.setattr(fp, "rms_norm_fastpath", spy)
+    x_cl = x.contiguous(memory_format=torch.channels_last_3d)
+    actual = fp.residual_block_forward(block, x_cl, feat_cache=[None, None], feat_idx=[0])
+    assert seen[0] is None and seen[1] is block.conv1.bias, "norm2 should receive conv1's un-added bias"
+    assert actual.is_contiguous(memory_format=torch.channels_last_3d)
+    if dtype is torch.float32:
+        torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2)  # channels-last convs may run TF32
+    else:
+        # Fast-math norm plus channels-last cuDNN accumulation order: a few outputs differ by one bf16 ulp.
+        torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+
+
+def test_rms_norm_scale_declines_bad_operands() -> None:
+    x = torch.randn(1, 8, 16, device="cuda", dtype=torch.bfloat16)
+    denom = torch.ones(1, 1, 16, device="cuda")
+    gamma = torch.ones(8, device="cuda", dtype=torch.bfloat16)
+    assert rn.rms_norm_scale(x, denom, gamma, 1.0) is not None
+    assert rn.rms_norm_scale(x.cpu(), denom.cpu(), gamma.cpu(), 1.0) is None
+    assert rn.rms_norm_scale(x, denom.to(torch.bfloat16), gamma, 1.0) is None
+    assert rn.rms_norm_scale(x, denom, gamma.float(), 1.0) is None
+    assert rn.rms_norm_scale(x.transpose(1, 2), denom, gamma, 1.0) is None
+    assert rn.rms_norm_scale(x, torch.ones(1, 1, 8, device="cuda"), gamma, 1.0) is None
+
+
+# --------------------------------------------------------------------------- #
+# Data movement
+# --------------------------------------------------------------------------- #
+
+
+def _dense_5d(shape, dtype: torch.dtype, channels_last: bool) -> torch.Tensor:
+    tensor = torch.randn(shape, device="cuda", dtype=dtype)
+    if channels_last:
+        return tensor.contiguous(memory_format=torch.channels_last_3d)
+    return tensor.contiguous()
+
+
+def _reference_cat_pad(x: torch.Tensor, cache: torch.Tensor | None, padding) -> torch.Tensor:
+    reference_padding = list(padding)
+    if cache is not None:
+        x = torch.cat([cache, x], dim=2)
+        reference_padding[4] -= cache.shape[2]
+    return F.pad(x, reference_padding)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("layout", ["contiguous", "channels_last", "frame_major"])
+@pytest.mark.parametrize(
+    "channels,frames,height,width,cache_frames,padding",
+    [
+        (96, 1, 10, 14, 0, (1, 1, 1, 1, 2, 0)),
+        (96, 1, 10, 14, 1, (1, 1, 1, 1, 2, 0)),
+        (96, 1, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
+        (64, 1, 10, 14, 2, (0, 0, 0, 0, 2, 0)),
+        (48, 4, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
+        # Non-power-of-two channels and a ragged number of rows per program.
+        (12, 3, 7, 9, 1, (1, 1, 1, 1, 2, 0)),
+        # Width wider than one block so the in-kernel width loop runs.
+        (8, 2, 3, 1300, 2, (1, 1, 1, 1, 2, 0)),
+        # Cosmos3 stage-3 shape, scaled down in height.
+        (256, 4, 8, 640, 2, (1, 1, 1, 1, 2, 0)),
+    ],
+)
+def test_cat_pad_5d_is_bitwise_and_layout_exact(
+    dtype, layout, channels, frames, height, width, cache_frames, padding
+) -> None:
+    torch.manual_seed(0)
+    channels_last = layout == "channels_last"
+    if layout == "frame_major":
+        x = torch.randn((1, frames, channels, height, width), device="cuda", dtype=dtype).permute(0, 2, 1, 3, 4)
+    else:
+        x = _dense_5d((1, channels, frames, height, width), dtype, channels_last)
+    cache = None
+    if cache_frames:
+        pad_height, pad_width = padding[2], padding[0]
+        buffer = _dense_5d(
+            (1, channels, cache_frames, height + 2 * pad_height, width + 2 * pad_width), dtype, channels_last
+        )
+        cache = buffer[:, :, :, pad_height : pad_height + height, pad_width : pad_width + width]
+
+    reference = _reference_cat_pad(x, cache, padding)
+    output = dm.cat_pad_5d(x, cache, padding)
+    assert output is not None
+    assert output.stride() == reference.stride()
+    assert torch.equal(output, reference)
+
+    pair = dm.cat_pad_5d(x, cache, padding, keep_cache_frames=CACHE_T)
+    assert pair is not None
+    output_with_cache, next_cache = pair
+    assert torch.equal(output_with_cache, reference)
+    pad_height, pad_width = padding[2], padding[0]
+    keep_frames = min(CACHE_T, reference.shape[2])
+    expected_cache = reference[:, :, -keep_frames:, pad_height : pad_height + height, pad_width : pad_width + width]
+    assert torch.equal(next_cache, expected_cache)
+    assert next_cache.is_contiguous(memory_format=torch.channels_last_3d) is channels_last
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("layout", ["contiguous", "frame_major"])
+@pytest.mark.parametrize("cache_layout", ["contiguous", "channels_last"])
+@pytest.mark.parametrize(
+    "channels,frames,height,width,cache_frames",
+    [
+        (96, 1, 10, 14, 0),
+        (96, 1, 10, 14, 2),
+        (12, 3, 7, 9, 1),
+        (8, 2, 3, 1300, 2),
+        (256, 4, 8, 640, 2),
+        (64, 4, 70, 70, 2),
+    ],
+)
+def test_cat_pad_5d_assembles_channels_last_output_from_channels_first_input(
+    dtype, layout, cache_layout, channels, frames, height, width, cache_frames
+) -> None:
+    """The conv_out path: channels-first activations, channels-last padded input, cache in the input layout."""
+    torch.manual_seed(0)
+    padding = (1, 1, 1, 1, 2, 0)
+    if layout == "frame_major":
+        x = torch.randn((1, frames, channels, height, width), device="cuda", dtype=dtype).permute(0, 2, 1, 3, 4)
+    else:
+        x = torch.randn((1, channels, frames, height, width), device="cuda", dtype=dtype)
+    cache = None
+    if cache_frames:
+        cache = _dense_5d((1, channels, cache_frames, height, width), dtype, cache_layout == "channels_last")
+    reference = _reference_cat_pad(x, cache, padding)
+
+    output = dm.cat_pad_5d(x, cache, padding, channels_last_output=True)
+    assert output is not None
+    assert output.is_contiguous(memory_format=torch.channels_last_3d)
+    assert torch.equal(output, reference)
+
+    pair = dm.cat_pad_5d(x, cache, padding, keep_cache_frames=CACHE_T, channels_last_output=True)
+    assert pair is not None
+    output_with_cache, next_cache = pair
+    assert output_with_cache.is_contiguous(memory_format=torch.channels_last_3d)
+    assert torch.equal(output_with_cache, reference)
+    keep_frames = min(CACHE_T, reference.shape[2])
+    assert next_cache.is_contiguous()  # follows x's layout, not the output's
+    assert torch.equal(next_cache, reference[:, :, -keep_frames:, 1 : 1 + height, 1 : 1 + width])
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("layout", ["contiguous", "frame_major"])
+def test_conv_out_channels_last_verdict_is_recorded_and_exact(layout: str) -> None:
+    torch.manual_seed(0)
+    conv = WanCausalConv3d(64, 12, 3, padding=1).to(device="cuda", dtype=torch.bfloat16)
+    if layout == "frame_major":
+        chunk = torch.randn(1, 4, 64, 24, 40, device="cuda", dtype=torch.bfloat16).permute(0, 2, 1, 3, 4)
+    else:
+        chunk = torch.randn(1, 64, 4, 24, 40, device="cuda", dtype=torch.bfloat16)
+    cache = torch.randn(1, 64, 2, 24, 40, device="cuda", dtype=torch.bfloat16)
+    reference = WanCausalConv3d.forward(conv, chunk, cache)
+    fast_cache = [cache.clone()]
+    out = fp._run_conv_out_channels_last(conv, chunk, fast_cache, 0)
+    assert out is not None  # the verdict decides which formulation's result is returned, both are exact
+    assert out.is_contiguous()
+    assert _bits_equal(out, reference)
+    verdicts = fp._CONV_OUT_LAYOUT_VERDICTS[conv]
+    assert len(verdicts) == 1
+    assert torch.equal(fast_cache[0], torch.cat([cache, chunk], dim=2)[:, :, -CACHE_T:])
+    assert fast_cache[0].is_contiguous()
+    again = fp._run_conv_out_channels_last(conv, chunk, [cache.clone()], 0)
+    if next(iter(verdicts.values())):
+        assert again is not None and _bits_equal(again, reference)
+    else:
+        assert again is None  # routed back to the standard path
+    # Channels-last activations are not this path's business.
+    assert (
+        fp._run_conv_out_channels_last(conv, chunk.contiguous(memory_format=torch.channels_last_3d), [None], 0) is None
+    )
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("layout", ["contiguous", "channels_last", "frame_major"])
+@pytest.mark.parametrize(
+    "channels,frames,height,width,cache_frames,pad_front",
+    [(96, 1, 10, 14, 0, 2), (96, 1, 10, 14, 1, 2), (96, 4, 10, 14, 2, 2), (12, 3, 7, 9, 2, 2), (256, 4, 8, 640, 2, 2)],
+)
+def test_cat_time_5d_is_bitwise_and_layout_exact(
+    dtype, layout, channels, frames, height, width, cache_frames, pad_front
+):
+    torch.manual_seed(0)
+    channels_last = layout == "channels_last"
+    if layout == "frame_major":
+        x = torch.randn((1, frames, channels, height, width), device="cuda", dtype=dtype).permute(0, 2, 1, 3, 4)
+    else:
+        x = _dense_5d((1, channels, frames, height, width), dtype, channels_last)
+    cache = _dense_5d((1, channels, cache_frames, height, width), dtype, channels_last) if cache_frames else None
+
+    parts = [torch.zeros_like(x[:, :, :1]).expand(-1, -1, pad_front - cache_frames, -1, -1)]
+    if cache is not None:
+        parts.append(cache)
+    parts.append(x)
+    reference = torch.cat(parts, dim=2)
+    if channels_last:
+        reference = reference.contiguous(memory_format=torch.channels_last_3d)
+    else:
+        reference = reference.contiguous()
+
+    pair = dm.cat_time_5d(x, cache, pad_front, keep_cache_frames=CACHE_T)
+    assert pair is not None
+    assembled, next_cache = pair
+    assert assembled.stride() == reference.stride()
+    assert torch.equal(assembled, reference)
+    assert torch.equal(next_cache, reference[:, :, -CACHE_T:])
+    assert next_cache.is_contiguous(memory_format=torch.channels_last_3d) is channels_last
+    assert torch.equal(dm.cat_time_5d(x, cache, pad_front), reference)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("cache_frames", [0, 1, 2])
+def test_cat_time_5d_exceeds_grid_y_limit(monkeypatch: pytest.MonkeyPatch, cache_frames: int) -> None:
+    # Smaller blocks reproduce the large-plane launch failure without allocating
+    # 4K feature maps. The final block is partial, exercising the plane boundary.
+    monkeypatch.setattr(dm, "_PLANE_BLOCK", 32)
+    batch, channels, frames, height, width = 2, 8, 1, 257, 1021
+    plane_size = channels * height * width
+    assert dm.triton.cdiv(plane_size, dm._PLANE_BLOCK) > 65_535
+    assert plane_size % dm._PLANE_BLOCK != 0
+
+    x = _dense_5d((batch, channels, frames, height, width), torch.bfloat16, True)
+    cache = _dense_5d((batch, channels, cache_frames, height, width), x.dtype, True) if cache_frames else None
+    parts = [torch.zeros_like(x).expand(-1, -1, 2 - cache_frames, -1, -1)]
+    if cache is not None:
+        parts.append(cache)
+    parts.append(x)
+    reference = torch.cat(parts, dim=2).contiguous(memory_format=torch.channels_last_3d)
+
+    pair = dm.cat_time_5d(x, cache, pad_front=2, keep_cache_frames=CACHE_T)
+    assert pair is not None
+    assembled, next_cache = pair
+    assert _bits_equal(assembled, reference)
+    assert assembled.stride() == reference.stride()
+    assert _bits_equal(next_cache, reference[:, :, -CACHE_T:])
+    assert next_cache.is_contiguous(memory_format=torch.channels_last_3d)
+    assert _bits_equal(dm.cat_time_5d(x, cache, pad_front=2), reference)
+
+
+@torch.no_grad()
+def test_cached_conv_spatial_padding_verdict_is_recorded_and_exact() -> None:
+    torch.manual_seed(0)
+    conv = WanCausalConv3d(64, 64, 3, padding=1).to(device="cuda", dtype=torch.bfloat16)
+    for layout in (torch.contiguous_format, torch.channels_last_3d):
+        if layout is torch.channels_last_3d:
+            conv.to(memory_format=torch.channels_last_3d)
+        chunk = torch.randn(1, 64, 4, 24, 40, device="cuda", dtype=torch.bfloat16).contiguous(memory_format=layout)
+        cache = torch.randn(1, 64, 2, 24, 40, device="cuda", dtype=torch.bfloat16).contiguous(memory_format=layout)
+        reference_cache = [cache.clone()]
+        reference = WanCausalConv3d.forward(conv, chunk, cache)  # upstream pre-padded path
+        fast_cache = [cache.clone()]
+        out = fp._run_cached_causal_conv(conv, chunk, fast_cache, 0)
+        verdicts = fp._SPATIAL_PAD_VERDICTS[conv]
+        assert len(verdicts) >= 1
+        assert _bits_equal(out, reference)  # exact whichever path the verdict selected
+        again = fp._run_cached_causal_conv(conv, chunk, [cache.clone()], 0)
+        assert _bits_equal(again, reference)
+        assert torch.equal(fast_cache[0], torch.cat([cache, chunk], dim=2)[:, :, -CACHE_T:])
+        del reference_cache
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("main_layout", ["contiguous", "channels_last", "frame_major"])
+@pytest.mark.parametrize("with_bias", [False, True])
+@pytest.mark.parametrize(
+    "in_channels,out_channels,frames,height,width,factor_temporal,factor_spatial,drop_first",
+    [
+        # Cosmos3 / Wan2.2 block shapes: repeats 8 (same channels), 4 (halved), 2 (spatial only, halved).
+        (128, 64, 1, 10, 14, 2, 2, False),
+        (128, 64, 1, 10, 14, 2, 2, True),
+        (64, 32, 2, 10, 14, 1, 2, False),
+        (64, 64, 4, 10, 14, 1, 2, False),
+        # More output than input channels (repeats 8 with a 2x channel increase) and a 4x spatial factor.
+        (32, 64, 2, 10, 14, 1, 2, False),
+        (16, 16, 1, 5, 6, 1, 4, False),
+        # Odd sizes and a row wider than one column block.
+        (12, 12, 3, 7, 9, 2, 2, True),
+        (8, 8, 1, 3, 700, 1, 2, False),
+    ],
+)
+def test_dup_up3d_add_is_bitwise_and_layout_exact(
+    dtype,
+    main_layout,
+    with_bias,
+    in_channels,
+    out_channels,
+    frames,
+    height,
+    width,
+    factor_temporal,
+    factor_spatial,
+    drop_first,
+) -> None:
+    torch.manual_seed(0)
+    repeats = out_channels * factor_temporal * factor_spatial * factor_spatial // in_channels
+    source = _dense_5d((1, in_channels, frames, height, width), dtype, main_layout == "channels_last")
+    if main_layout == "frame_major":
+        # The previous block's output is frame-major too in the lossless level.
+        source = torch.randn((1, frames, in_channels, height, width), device="cuda", dtype=dtype)
+        source = source.permute(0, 2, 1, 3, 4)
+    out_frames = frames * factor_temporal - (factor_temporal - 1 if drop_first else 0)
+    shape = (1, out_channels, out_frames, height * factor_spatial, width * factor_spatial)
+    if main_layout == "channels_last":
+        main = _dense_5d(shape, dtype, True)
+    elif main_layout == "frame_major":
+        main = torch.randn((shape[0], shape[2], shape[1], shape[3], shape[4]), device="cuda", dtype=dtype)
+        main = main.permute(0, 2, 1, 3, 4)
+    else:
+        main = _dense_5d(shape, dtype, False)
+
+    shortcut = source.repeat_interleave(repeats, dim=1)
+    shortcut = shortcut.view(1, out_channels, factor_temporal, factor_spatial, factor_spatial, frames, height, width)
+    shortcut = shortcut.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+    shortcut = shortcut.view(shape[0], shape[1], frames * factor_temporal, shape[3], shape[4])
+    if drop_first:
+        shortcut = shortcut[:, :, factor_temporal - 1 :]
+    bias = None
+    biased_main = main
+    if with_bias:
+        bias = (torch.randn(out_channels, device="cuda") * 0.1).to(dtype)
+        biased_main = main.clone()
+        biased_main.add_(bias.view(1, -1, 1, 1, 1))  # ATen's conv bias add
+    reference = biased_main + shortcut
+
+    output = dm.dup_up3d_add(main, source, factor_temporal, factor_spatial, repeats, drop_first, main_bias=bias)
+    assert output is not None
+    assert _same_dense_strides(output, reference)
+    assert torch.equal(output, reference)
+
+
+def test_dup_up3d_add_declines_unsupported_cases() -> None:
+    main = torch.randn(1, 8, 2, 10, 14, device="cuda", dtype=torch.bfloat16)
+    source = torch.randn(1, 32, 2, 5, 7, device="cuda", dtype=torch.bfloat16)
+    # repeats (1) below the spatial factor: the source channel would change along a row.
+    assert dm.dup_up3d_add(main, source, 1, 2, 1, False) is None
+    source = torch.randn(1, 8, 2, 5, 7, device="cuda", dtype=torch.bfloat16)
+    assert dm.dup_up3d_add(main, source, 1, 2, 4, False) is not None
+    # Mixed layouts, wrong bias shape, wrong main shape.
+    assert dm.dup_up3d_add(main.contiguous(memory_format=torch.channels_last_3d), source, 1, 2, 4, False) is None
+    assert dm.dup_up3d_add(main, source, 1, 2, 4, False, main_bias=torch.ones(7, device="cuda")) is None
+    assert dm.dup_up3d_add(main[:, :, :1], source, 1, 2, 4, False) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("layout", ["contiguous", "channels_last", "frame_major"])
+@pytest.mark.parametrize(
+    "frames,channels,height,width",
+    [
+        (1, 96, 10, 14),
+        (4, 64, 10, 14),
+        # Non-power-of-two channels and ragged rows/columns for every block size.
+        (3, 12, 7, 9),
+        # Width wider than one column block so the in-kernel loop runs.
+        (2, 8, 3, 1300),
+        # Cosmos3 720p stage-3 input, scaled down in height.
+        (4, 256, 8, 640),
+        # Channels equal to and above the channels-last block: in-kernel channel loop.
+        (2, 1024, 5, 160),
+        (1, 1100, 3, 20),
+    ],
+)
+def test_upsample_nearest_2x_is_bitwise_and_layout_exact(dtype, layout, frames, channels, height, width) -> None:
+    torch.manual_seed(0)
+    if layout == "frame_major":
+        # The strided ``(t, c, h, w)`` view ``WanResample`` hands the upsampler in the lossless level.
+        x = fp._merge_batch_and_frames(torch.randn((1, channels, frames, height, width), device="cuda", dtype=dtype))
+        assert frames == 1 or not x.is_contiguous()
+    elif layout == "channels_last":
+        x = torch.randn((frames, channels, height, width), device="cuda", dtype=dtype)
+        x = x.contiguous(memory_format=torch.channels_last)
+    else:
+        x = torch.randn((frames, channels, height, width), device="cuda", dtype=dtype)
+    for mode in ("nearest-exact", "nearest"):
+        reference = F.interpolate(x, scale_factor=(2.0, 2.0), mode=mode)
+        output = up.upsample_nearest_2x(x)
+        assert output is not None
+        assert output.stride() == reference.stride()
+        assert torch.equal(output, reference)
+
+
+def test_upsample_nearest_2x_declines_unsupported_inputs() -> None:
+    x = torch.randn(2, 8, 6, 10, device="cuda", dtype=torch.bfloat16)
+    assert up.upsample_nearest_2x(x) is not None
+    assert up.upsample_nearest_2x(x.cpu()) is None
+    assert up.upsample_nearest_2x(x.unsqueeze(2)) is None
+    assert up.upsample_nearest_2x(x.transpose(2, 3)) is None
+    assert up.upsample_nearest_2x(x[..., ::2]) is None
+    assert up.upsample_nearest_2x(x.double()) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_upsample_forward_matches_wan_upsample(dtype: torch.dtype) -> None:
+    torch.manual_seed(0)
+    x = torch.randn(3, 64, 10, 14, device="cuda", dtype=dtype)
+    for layout in (torch.contiguous_format, torch.channels_last):
+        x = x.contiguous(memory_format=layout)
+        module = WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact")
+        expected = WanUpsample.forward(module, x)
+        actual = fp.upsample_forward(module, x)
+        assert actual.stride() == expected.stride()
+        assert torch.equal(actual, expected)
+        # Anything but a nearest 2x upsample takes the reference path.
+        other = WanUpsample(scale_factor=(3.0, 3.0), mode="nearest-exact")
+        assert torch.equal(fp.upsample_forward(other, x), WanUpsample.forward(other, x))
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("mode", ["upsample2d", "upsample3d"])
+def test_resample_forward_can_leave_conv_bias_pending(dtype: torch.dtype, mode: str) -> None:
+    torch.manual_seed(0)
+    resample = WanResample(32, mode=mode).eval().to(device="cuda", dtype=dtype)
+    setattr(resample, fp.CFG_ATTR, fp.FastPathConfig())
+    x = torch.randn(1, 32, 2, 10, 14, device="cuda", dtype=dtype)
+    cache = [torch.randn(1, 32, 2, 10, 14, device="cuda", dtype=dtype)] if mode == "upsample3d" else [None]
+    expected = WanResample.forward(
+        resample, x, feat_cache=[c.clone() if c is not None else None for c in cache], feat_idx=[0]
+    )
+    out, bias = fp.resample_forward(
+        resample, x, [c.clone() if c is not None else None for c in cache], [0], return_bias=True
+    )
+    assert bias is resample.resample[1].bias
+    assert out.stride() == expected.stride()
+    assert _bits_equal(out + bias.view(1, -1, 1, 1, 1), expected)
+    plain = fp.resample_forward(resample, x, [c.clone() if c is not None else None for c in cache], [0])
+    assert _bits_equal(plain, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", LOW_PRECISION)
+@pytest.mark.parametrize("force_prepad", [False, True])
+def test_cached_conv_deferred_bias_matches_autocast(dtype: torch.dtype, force_prepad: bool, monkeypatch) -> None:
+    """FP32 parameters and inputs must defer the bias in the convolution's output dtype."""
+    torch.manual_seed(8)
+    conv = WanCausalConv3d(32, 32, 3, padding=1).cuda().eval()
+    conv.bias.copy_(torch.randn_like(conv.bias))
+    if force_prepad:
+        monkeypatch.setattr(dm, "cat_time_5d", lambda *args, **kwargs: None)
+    cache = [None]
+    with torch.autocast("cuda", dtype=dtype):
+        for _ in range(2):
+            x = torch.randn(1, 32, 2, 6, 10, device="cuda")
+            expected = conv(x, cache[0])
+            out, bias = fp._run_cached_causal_conv(conv, x, cache, 0, return_bias=True)
+            assert out.dtype == dtype and bias is not None and bias.dtype == dtype
+            actual = dm.add_bias_residual(out, bias, torch.zeros_like(out), None)
+            assert actual is not None
+            assert _bits_equal(actual.contiguous(), expected.contiguous())
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", LOW_PRECISION)
+@pytest.mark.parametrize("channels_last", [False, True])
+def test_residual_block_deferred_bias_matches_autocast(dtype: torch.dtype, channels_last: bool) -> None:
+    """Exercise conv1/conv2 and a non-identity shortcut with FP32 weights under autocast."""
+    torch.manual_seed(8)
+    block = WanResidualBlock(32, 64, dropout=0.0).cuda().eval()
+    for conv in (block.conv1, block.conv2, block.conv_shortcut):
+        conv.bias.copy_(torch.randn_like(conv.bias))
+        if channels_last:
+            conv.to(memory_format=torch.channels_last_3d)
+    x = torch.randn(1, 32, 2, 6, 10, device="cuda")
+    with torch.autocast("cuda", dtype=dtype):
+        expected = block(x, feat_cache=[None, None], feat_idx=[0])
+        setattr(block, fp.CFG_ATTR, fp.FastPathConfig(channels_last=channels_last))
+        actual = fp.residual_block_forward(block, x, feat_cache=[None, None], feat_idx=[0])
+    assert actual.dtype == expected.dtype == dtype
+    assert _bits_equal(actual.contiguous(), expected.contiguous())
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", LOW_PRECISION)
+@pytest.mark.parametrize("mode", ["upsample2d", "upsample3d"])
+def test_resample_deferred_bias_matches_autocast(dtype: torch.dtype, mode: str) -> None:
+    """The upsampler's FP32 bias must be rounded before its deferred addition."""
+    torch.manual_seed(8)
+    resample = WanResample(32, mode=mode).cuda().eval()
+    resample.resample[1].bias.copy_(torch.randn_like(resample.resample[1].bias))
+    reference_cache = [None]
+    candidate_cache = [None]
+    with torch.autocast("cuda", dtype=dtype):
+        for _ in range(2):
+            x = torch.randn(1, 32, 1, 6, 10, device="cuda")
+            expected = resample(x, feat_cache=reference_cache, feat_idx=[0])
+            out, bias = fp.resample_forward(resample, x, candidate_cache, [0], return_bias=True)
+            assert out.dtype == dtype and bias is not None and bias.dtype == dtype
+            actual = out + bias.view(1, -1, 1, 1, 1)
+            assert _bits_equal(actual.contiguous(), expected.contiguous())
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("with_bias_x", [True, False])
+@pytest.mark.parametrize("with_bias_h", [True, False])
+def test_add_bias_residual_matches_aten_rounding(dtype: torch.dtype, with_bias_x: bool, with_bias_h: bool) -> None:
+    torch.manual_seed(0)
+    shape = (1, 256, 4, 45, 80)
+    x = torch.randn(shape, device="cuda").to(dtype)
+    h = torch.randn(shape, device="cuda").to(dtype)
+    bias_x = torch.randn(256, device="cuda").to(dtype) if with_bias_x else None
+    bias_h = torch.randn(256, device="cuda").to(dtype) if with_bias_h else None
+    expected_x = x.clone()
+    expected_h = h.clone()
+    if bias_x is not None:
+        expected_x.add_(bias_x.view(1, -1, 1, 1, 1))
+    if bias_h is not None:
+        expected_h.add_(bias_h.view(1, -1, 1, 1, 1))
+    expected = expected_x + expected_h
+    actual = dm.add_bias_residual(x, bias_x, h, bias_h)
+    assert actual is not None
+    assert actual.stride() == expected.stride()
+    assert _bits_equal(actual, expected)
+    x_cl = x.contiguous(memory_format=torch.channels_last_3d)
+    h_cl = h.contiguous(memory_format=torch.channels_last_3d)
+    actual_cl = dm.add_bias_residual(x_cl, bias_x, h_cl, bias_h)
+    assert actual_cl is not None and actual_cl.is_contiguous(memory_format=torch.channels_last_3d)
+    assert _bits_equal(actual_cl.contiguous(), expected)
+    assert dm.add_bias_residual(x_cl, bias_x, h, bias_h) is None
+    assert dm.add_bias_residual(x.permute(0, 2, 1, 3, 4).contiguous().permute(0, 2, 1, 3, 4), bias_x, h, bias_h) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("channels_last", [False, True])
+@pytest.mark.parametrize("temporal_only", [False, True])
+@pytest.mark.parametrize("start", [None, "Rep"])
+def test_cached_conv_chunk_loop_matches_reference_including_mixed_paths(
+    monkeypatch, channels_last: bool, temporal_only: bool, start
+) -> None:
+    torch.manual_seed(0)
+    channels = 64
+    if temporal_only:
+        conv = WanCausalConv3d(channels, 2 * channels, (3, 1, 1), padding=(1, 0, 0))
+    else:
+        conv = WanCausalConv3d(channels, channels, 3, padding=1)
+    conv = conv.to(device="cuda", dtype=torch.float32)
+    chunks = [_dense_5d((1, channels, 1, 10, 14), torch.float32, channels_last) for _ in range(4)]
+
+    decline = lambda *a, **k: None  # noqa: E731
+    original_time, original_pad = fp.dm.cat_time_5d, fp.dm.cat_pad_5d
+
+    def run(mode: list[str]):
+        cache = [start]
+        outputs = []
+        for chunk, m in zip(chunks, mode, strict=True):
+            monkeypatch.setattr(fp.dm, "cat_time_5d", decline if m in ("reference", "padded") else original_time)
+            monkeypatch.setattr(fp.dm, "cat_pad_5d", decline if m == "reference" else original_pad)
+            outputs.append(fp._run_cached_causal_conv(conv, chunk, cache, 0))
+        monkeypatch.setattr(fp.dm, "cat_time_5d", original_time)
+        monkeypatch.setattr(fp.dm, "cat_pad_5d", original_pad)
+        return outputs
+
+    reference = run(["reference"] * 4)
+    for pattern in (
+        ["time"] * 4,
+        ["padded"] * 4,
+        ["time", "reference", "padded", "time"],
+        ["reference", "padded", "time", "reference"],
+    ):
+        fused = run(pattern)
+        for out, ref in zip(fused, reference, strict=True):
+            assert torch.equal(out, ref), pattern
+
+
+# --------------------------------------------------------------------------- #
+# Whole decoder
+# --------------------------------------------------------------------------- #
+
+TINY_RESIDUAL = dict(
+    base_dim=8,
+    decoder_base_dim=16,
+    z_dim=8,
+    dim_mult=[1, 2, 4, 4],
+    num_res_blocks=2,
+    temperal_downsample=[False, True, True],
+    is_residual=True,
+    patch_size=2,
+    in_channels=12,
+    out_channels=12,
+)
+TINY_WAN21 = dict(base_dim=8, z_dim=4, dim_mult=[1, 2], num_res_blocks=1, temperal_downsample=[False, True])
+
+
+def _build_pair(config: dict, dtype: torch.dtype):
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**config).eval().to(device="cuda", dtype=dtype)
+    candidate = AutoencoderKLWan(**config).eval().to(device="cuda", dtype=dtype)
+    candidate.load_state_dict(reference.state_dict())
+    return reference, candidate
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("config", [TINY_RESIDUAL, TINY_WAN21], ids=["residual_patch2", "wan21"])
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("frames", [1, 2, 5])
+def test_lossless_decoder_is_bitwise_exact_on_cuda(config: dict, dtype: torch.dtype, frames: int) -> None:
+    reference, candidate = _build_pair(config, dtype)
+    report = install_wan_vae_fastpath(candidate, level="lossless")
+    assert report.installed, report
+    torch.manual_seed(1)
+    latents = torch.randn(1, config["z_dim"], frames, 6, 10, device="cuda").to(dtype)
+    autocast = dtype is not torch.float32
+    with torch.autocast("cuda", dtype=dtype, enabled=autocast):
+        expected = reference.decode(latents, return_dict=False)[0]
+        actual = candidate.decode(latents, return_dict=False)[0]
+    assert actual.stride() == expected.stride()
+    assert _bits_equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_omni_wan_vae_decode_is_bitwise_exact_on_cuda(dtype: torch.dtype) -> None:
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
+
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**TINY_RESIDUAL).eval().to(device="cuda", dtype=dtype)
+    candidate = OmniAutoencoderKLWan(**TINY_RESIDUAL).eval().to(device="cuda", dtype=dtype)
+    candidate.load_state_dict(reference.state_dict())
+    install_wan_vae_fastpath(candidate)
+    latents = torch.randn(1, TINY_RESIDUAL["z_dim"], 5, 6, 10, device="cuda").to(dtype)
+    with torch.autocast("cuda", dtype=dtype, enabled=dtype is not torch.float32):
+        expected = reference.decode(latents, return_dict=False)[0]
+    actual = candidate.decode(latents, return_dict=False)[0]  # applies its own autocast context
+    assert actual.stride() == expected.stride()
+    assert _bits_equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize(
+    ("dtype", "allow_tf32", "tolerance"),
+    [
+        pytest.param(torch.bfloat16, True, 3e-2, id="bf16"),
+        # With TF32 disabled the only difference is cuDNN's channels-last algorithm
+        # choice and the norm kernel's reduction order.
+        pytest.param(torch.float32, False, 2e-4, id="fp32"),
+        # PyTorch's default allows cuDNN to run fp32 convolutions in TF32; the
+        # channels-last algorithms use it while the NCDHW reference may not, so
+        # the deviation is dominated by TF32's 10-bit mantissa.
+        pytest.param(torch.float32, True, 1e-2, id="fp32-tf32"),
+    ],
+)
+def test_channels_last_decoder_keeps_layout_and_stays_close(
+    dtype: torch.dtype, allow_tf32: bool, tolerance: float
+) -> None:
+    reference, candidate = _build_pair(TINY_RESIDUAL, dtype)
+    report = install_wan_vae_fastpath(candidate, level="channels_last")
+    assert report.installed and report.channels_last
+
+    violations: list[str] = []
+
+    def audit(name: str):
+        def hook(module: nn.Module, inputs) -> None:
+            x = inputs[0]
+            if x.dim() == 5 and x.shape[1] > 1 and not x.is_contiguous(memory_format=torch.channels_last_3d):
+                violations.append(f"{name}: {tuple(x.shape)} strides={x.stride()}")
+            if x.dim() == 4 and x.shape[1] > 1 and not x.is_contiguous(memory_format=torch.channels_last):
+                violations.append(f"{name}: {tuple(x.shape)} strides={x.stride()}")
+
+        return hook
+
+    handles = [
+        module.register_forward_pre_hook(audit(name))
+        for name, module in candidate.decoder.named_modules()
+        if isinstance(module, (nn.Conv2d, nn.Conv3d))
+    ]
+    torch.manual_seed(1)
+    latents = torch.randn(1, TINY_RESIDUAL["z_dim"], 3, 6, 10, device="cuda").to(dtype)
+    with (
+        torch.backends.cudnn.flags(enabled=True, allow_tf32=allow_tf32),
+        torch.autocast("cuda", dtype=dtype, enabled=dtype is not torch.float32),
+    ):
+        expected = reference.decode(latents, return_dict=False)[0]
+        actual = candidate.decode(latents, return_dict=False)[0]
+    for handle in handles:
+        handle.remove()
+    assert not violations, "\n".join(violations)
+    assert actual.shape == expected.shape
+    difference = (actual.float() - expected.float()).abs()
+    error = difference.max().item()
+    mse = (difference**2).mean().item()
+    psnr = float("inf") if mse == 0 else 10 * torch.log10(torch.tensor(4.0 / mse)).item()
+    assert error <= tolerance, f"max_abs_diff={error:.3e} psnr={psnr:.2f} dB tolerance={tolerance}"

tests/diffusion/test_diffusion_config_propagation.py

@@ -214,6 +214,17 @@ def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
     assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 
 
+def test_vae_fast_path_roundtrip():
+    assert _roundtrip_diffusion_config(model="x").vae_fast_path == "lossless"
+    od = _roundtrip_diffusion_config(model="x", vae_fast_path="channels_last")
+    assert od.vae_fast_path == "channels_last"
+
+
+def test_invalid_vae_fast_path_is_rejected():
+    with pytest.raises(ValueError, match="vae_fast_path"):
+        OmniDiffusionConfig(model="x", vae_fast_path="fast")
+
+
 def test_flux2_klein_sets_generic_multimodal_limit():
     od_config = OmniDiffusionConfig(
         model="black-forest-labs/FLUX.2-klein-9B",

tests/entrypoints/test_serve.py

@@ -81,6 +81,23 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None:
     assert "api_server_count" not in explicit
 
 
+def test_serve_parser_accepts_vae_fast_path() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(["serve", "fake-model", "--omni"])
+    assert args.vae_fast_path == "lossless"
+    assert "vae_fast_path" not in args.get_explicit_kwargs_dict()
+
+    args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-fast-path", "channels_last"])
+    assert args.vae_fast_path == "channels_last"
+    assert args.get_explicit_kwargs_dict()["vae_fast_path"] == "channels_last"
+
+    with pytest.raises(SystemExit):
+        parser.parse_args(["serve", "fake-model", "--omni", "--vae-fast-path", "fast"])
+
+
 def _parse_serve_args(argv: list[str]) -> TrackingNamespace:
     """Parse a ``serve`` argv through the real Omni parser, returning the
     TrackingNamespace (with ``explicit_keys``) that ``validate`` receives."""

vllm_omni/config/omni_config.py

@@ -854,6 +854,7 @@ class _DiffusionConfigProjection:
     fa_deterministic: bool = False
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
+    vae_fast_path: Literal["off", "lossless", "channels_last"] = "lossless"
     mask_strategy_file_path: str | None = None
     skip_time_steps: int = 15
     VSA_sparsity: float = 0.0

vllm_omni/config/stage_config.py

@@ -470,6 +470,7 @@ class StageDeployConfig:
     step_execution: bool | None = None
     vae_use_slicing: bool | None = None
     vae_use_tiling: bool | None = None
+    vae_fast_path: str | None = None
     boundary_ratio: float | None = None
     flow_shift: float | None = None
     diffusion_kv_cache_dtype: str | None = None

vllm_omni/diffusion/data.py

@@ -44,6 +44,9 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+# Accepted values for ``OmniDiffusionConfig.vae_fast_path``.
+VAE_FAST_PATH_LEVELS: tuple[str, ...] = ("off", "lossless", "channels_last")
+
 
 def _move_diffusion_alias(
     normalized: dict[str, Any],
@@ -944,6 +947,12 @@ class OmniDiffusionConfig:
     # VAE memory optimization parameters
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
+    # Wan VAE decoder fast path. ``"lossless"`` installs the bit-exact fused
+    # kernels on every diffusers Wan VAE, ``"channels_last"`` additionally
+    # converts decoder convolution weights to channels-last memory format and
+    # fuses RMSNorm+SiLU (faster, not bit-exact), ``"off"`` keeps the reference
+    # diffusers implementation.
+    vae_fast_path: str = "lossless"
 
     # STA (Sliding Tile Attention) parameters
     mask_strategy_file_path: str | None = None
@@ -1174,6 +1183,8 @@ class OmniDiffusionConfig:
             materialize_legacy_offload_flags,
         )
 
+        if self.vae_fast_path not in VAE_FAST_PATH_LEVELS:
+            raise ValueError(f"vae_fast_path must be one of {list(VAE_FAST_PATH_LEVELS)}, got {self.vae_fast_path!r}")
         if self.diffusion_compile_granularity not in {"regional", "full"}:
             raise ValueError(
                 "diffusion_compile_granularity must be 'regional' or 'full', "

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -19,6 +19,7 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     GridSpec,
     TileTask,
 )
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import decode_frames, is_installed
 from vllm_omni.diffusion.models.interface import DecodedChunkConsumer
 from vllm_omni.platforms import current_omni_platform
 
@@ -101,6 +102,15 @@ class OmniAutoencoderKLWan(AutoencoderKLWan):
         return_dict: bool = True,
     ):
         """Decode the non-tiled temporal path, optionally streaming chunks."""
+        if on_chunk is None and is_installed(self) and not torch.is_grad_enabled():
+            try:
+                out = decode_frames(self, z)
+            finally:
+                self.clear_cache()
+            if not return_dict:
+                return (out,)
+            return DecoderOutput(sample=out)
+
         self.clear_cache()
         callback_error: BaseException | None = None
         try:

vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/__init__.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fast paths for the diffusers Wan (2.1/2.2) causal video VAE decoder.
+
+``install_wan_vae_fastpath(vae, level=...)`` rebinds the forwards of one loaded
+``AutoencoderKLWan`` instance:
+
+* ``"lossless"`` (default): bit-exact rewrites of the decoder's data movement
+  and normalization (fused Triton kernels with exact PyTorch fallbacks).
+* ``"channels_last"``: additionally converts decoder convolution weights to
+  channels-last memory format (faster cuDNN kernels, not bit-exact).
+* ``"off"``: leave the diffusers implementation untouched.
+
+The framework installs it from ``vllm_omni.diffusion.registry.initialize_model``
+according to ``OmniDiffusionConfig.vae_fast_path``.
+"""
+
+from .decode import decode_frames
+from .install import (
+    REPORT_ATTR,
+    VAE_FAST_PATH_LEVELS,
+    WanVaeFastPathReport,
+    install_wan_vae_fastpath,
+    is_installed,
+    uninstall_wan_vae_fastpath,
+)
+
+__all__ = [
+    "REPORT_ATTR",
+    "VAE_FAST_PATH_LEVELS",
+    "WanVaeFastPathReport",
+    "decode_frames",
+    "install_wan_vae_fastpath",
+    "is_installed",
+    "uninstall_wan_vae_fastpath",
+]

vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/_utils.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Shared helpers for configuring Wan VAE fast path kernels."""
+
+from __future__ import annotations
+
+_ROW_BLOCK_WIDTHS = (64, 128, 256)
+
+
+def _pick_block_width(width: int) -> int:
+    """The column block (64/128/256) that pads ``width`` the least; ties go to the wider block."""
+    best, best_padded = _ROW_BLOCK_WIDTHS[0], None
+    for block in _ROW_BLOCK_WIDTHS:
+        padded = -(-width // block) * block
+        if best_padded is None or padded <= best_padded:
+            best, best_padded = block, padded
+    return best

vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/decode.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Chunked Wan decode with a preallocated output buffer."""
+
+from __future__ import annotations
+
+import torch
+from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify
+
+
+def decode_frames(vae, z: torch.Tensor) -> torch.Tensor:
+    """``AutoencoderKLWan._decode``'s frame loop, writing every chunk straight into the result.
+
+    Upstream grows the result with ``torch.cat([out, out_], 2)`` on every chunk
+    (quadratic copying) and then materializes ``unpatchify`` and ``clamp`` as two
+    more full-size copies, so three or four copies of the decoded video are live
+    at once. Here the final ``[B, C, T, H, W]`` buffer is allocated after the
+    first chunk and each chunk is unpatchified and clamped directly into its
+    slot. Every step is a permutation or an elementwise clamp, so the values are
+    identical to upstream. Tiling dispatch is the caller's responsibility.
+    """
+    vae.clear_cache()
+    x = vae.post_quant_conv(z)
+    patch_size = vae.config.patch_size
+    num_frames = x.shape[2]
+    # Each temporal upsampler doubles the frames of every chunk after the first
+    # (the first chunk drops the extra leading frames), so steady-state chunks
+    # produce ``2 ** num_temporal_upsamplers`` frames each.
+    frames_per_chunk = 2 ** sum(bool(flag) for flag in vae.decoder.temperal_upsample)
+
+    output: torch.Tensor | None = None
+    offset = 0
+    for index in range(num_frames):
+        vae._conv_idx = [0]
+        chunk = vae.decoder(
+            x[:, :, index : index + 1],
+            feat_cache=vae._feat_map,
+            feat_idx=vae._conv_idx,
+            first_chunk=index == 0,
+        )
+        if patch_size is not None:
+            chunk = unpatchify(chunk, patch_size=patch_size)
+        if output is None:
+            total_frames = chunk.shape[2] + (num_frames - 1) * frames_per_chunk
+            output = chunk.new_empty((chunk.shape[0], chunk.shape[1], total_frames, chunk.shape[3], chunk.shape[4]))
+        end = offset + chunk.shape[2]
+        if end > output.shape[2]:
+            raise RuntimeError(
+                f"Wan decoder produced more frames than expected: chunk {index} ends at {end}, "
+                f"allocated {output.shape[2]}"
+            )
+        torch.clamp(chunk, min=-1.0, max=1.0, out=output[:, :, offset:end])
+        offset = end
+
+    if output is None or offset != output.shape[2]:
+        expected = None if output is None else output.shape[2]
+        raise RuntimeError(f"Wan decoder produced {offset} frames, expected {expected}")
+    vae.clear_cache()
+    return output
+
+
+__all__ = ["decode_frames"]

vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/forwards.py

@@ -0,0 +1,735 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Bit-exact replacement forwards for the diffusers Wan VAE decoder modules.
+
+Each function below is bound per module instance (``types.MethodType``) by
+:mod:`.install`; the diffusers classes themselves are never modified. Every
+function reproduces the diffusers 0.40 operation order exactly and only
+differs in how the bytes move:
+
+* ``WanRMS_norm``: fp32 ``vector_norm`` straight from the low-precision tensor
+  plus one fused Triton epilogue (see :mod:`.triton_rms_norm`), optionally with
+  the following SiLU folded in. At the ``channels_last`` level the single-pass
+  kernel also absorbs the bias of the ``conv1`` that feeds ``norm2``.
+* ``WanCausalConv3d`` and every cached call site: the ``clone`` + ``cat`` +
+  ``F.pad`` triple becomes one layout-preserving kernel that also writes the
+  next cache frames; all-zero paddings skip ``F.pad`` entirely.
+* ``WanResidualUpBlock``: ``x + DupUp3D(x)`` becomes one gather + add that
+  also applies the bias of the upsampler's ``Conv2d`` (which feeds nothing else).
+* ``WanResample``: the ``upsample3d`` time interleave becomes one strided copy.
+* ``WanUpsample``: the nearest-exact 2x upsample becomes one gather kernel
+  that reads the strided frame-major input in place (see
+  :mod:`.triton_upsample`); it only replicates values, so the upstream fp32
+  round trip is the identity and is skipped.
+
+Whenever a kernel declines an input, the code falls through to the reference
+expression, so the result is identical on every path.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+from weakref import WeakKeyDictionary
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.autoencoders.autoencoder_kl_wan import (
+    CACHE_T,
+    DupUp3D,
+    WanCausalConv3d,
+    WanDecoder3d,
+    WanResample,
+    WanResidualBlock,
+    WanResidualUpBlock,
+    WanUpsample,
+)
+from vllm.logger import init_logger
+
+from . import triton_data_movement as dm
+from . import triton_rms_norm as rn
+from . import triton_rms_norm_cl as cl
+from . import triton_upsample as up
+
+logger = init_logger(__name__)
+
+CFG_ATTR = "_vllm_omni_wan_fastpath_cfg"
+_NORM_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+
+
+@dataclass(frozen=True)
+class FastPathConfig:
+    """Per-installation switches read by the bound forwards."""
+
+    fused_silu_dtypes: frozenset[torch.dtype] = frozenset()
+    channels_last: bool = False
+
+
+def is_diffusers_rms_norm(module: Any) -> bool:
+    """True for a diffusers ``WanRMS_norm`` instance.
+
+    Identified by name because ``vllm_omni.diffusion.models.wan2_2.patch_diffusers``
+    rebinds the name ``WanRMS_norm`` to ``RMSNormVAE`` in diffusers' own module
+    namespace; instances created before that patch keep the original class.
+    ``RMSNormVAE`` has different numerics (eps 1e-6, no fp32 upcast) and is
+    deliberately not matched.
+    """
+    cls = type(module)
+    return (
+        cls.__name__ == "WanRMS_norm"
+        and cls.__module__.startswith("diffusers.")
+        and isinstance(getattr(module, "gamma", None), torch.Tensor)
+        and hasattr(module, "scale")
+        and hasattr(module, "channel_first")
+    )
+
+
+def _kernels_allowed(x: torch.Tensor) -> bool:
+    return x.is_cuda and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+
+
+# --------------------------------------------------------------------------- #
+# WanRMS_norm
+# --------------------------------------------------------------------------- #
+
+
+def _as_rows(
+    x: torch.Tensor, denom: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor, Callable[[torch.Tensor], torch.Tensor]] | None:
+    """View ``x``/``denom`` as contiguous ``(N, C, S)`` / ``(N, 1, S)`` without copying.
+
+    Handles the dense layouts the decoder produces: plain contiguous 4D/5D
+    tensors and the frame-major 5D view ``WanResample`` emits
+    (``permute(0, 2, 1, 3, 4)`` of a contiguous ``(b, t, c, h, w)`` buffer).
+    Returns the two views and a function mapping a kernel output back to
+    ``x``'s shape and strides, or ``None`` for any other layout.
+    """
+    if x.dim() == 5:
+        b, c, t, h, w = x.shape
+        if x.is_contiguous():
+            spatial = t * h * w
+            return x.view(b, c, spatial), denom.view(b, 1, spatial), lambda out: out.view(b, c, t, h, w)
+        frame_major = x.permute(0, 2, 1, 3, 4)
+        denom_frame_major = denom.permute(0, 2, 1, 3, 4)
+        if frame_major.is_contiguous() and denom_frame_major.is_contiguous():
+            spatial = h * w
+            return (
+                frame_major.view(b * t, c, spatial),
+                denom_frame_major.view(b * t, 1, spatial),
+                lambda out: out.view(b, t, c, h, w).permute(0, 2, 1, 3, 4),
+            )
+        return None
+    if x.dim() == 4 and x.is_contiguous():
+        n, c, h, w = x.shape
+        spatial = h * w
+        return x.view(n, c, spatial), denom.view(n, 1, spatial), lambda out: out.view(n, c, h, w)
+    return None
+
+
+def _add_channel_bias(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+    """Out-of-place ``x + bias`` per channel; rounds exactly like ATen's in-place conv bias ``add_``."""
+    return x + bias.view(1, -1, *([1] * (x.dim() - 2)))
+
+
+def rms_norm_fastpath(
+    norm: nn.Module, x: torch.Tensor, *, silu: bool = False, bias: torch.Tensor | None = None
+) -> torch.Tensor | None:
+    """``WanRMS_norm(x + bias)`` (optionally followed by SiLU) with fewer passes, or ``None``.
+
+    ``bias`` is the un-added per-channel bias of the convolution that produced
+    ``x``; only the channels-last kernel absorbs it, every other path adds it
+    first with ATen's rounding. ``None`` means the caller must run the reference
+    forward (autograd, tensor bias, channel-last norm, unsupported dtype,
+    gamma/activation dtype mismatch) and is responsible for adding ``bias`` itself.
+    """
+    if torch.is_grad_enabled():
+        return None
+    if not getattr(norm, "channel_first", False):
+        return None
+    norm_bias = norm.bias
+    if not (isinstance(norm_bias, float) and norm_bias == 0.0):
+        return None
+    gamma = norm.gamma
+    if x.dtype not in _NORM_DTYPES or gamma.dtype is not x.dtype:
+        return None
+    if x.dim() not in (4, 5) or gamma.dim() != x.dim() - 1 or x.numel() == 0:
+        return None
+
+    cfg = getattr(norm, CFG_ATTR, None)
+    if cfg is not None and cfg.channels_last and _kernels_allowed(x):
+        # Tier 2: one read, one write; not bit-exact (different reduction order).
+        out = cl.rms_norm_channels_last(x, gamma, norm.scale, silu=silu, bias=bias)
+        if out is not None:
+            return out
+    if bias is not None:
+        x = _add_channel_bias(x, bias)
+
+    # Same reduction kernel as ``F.normalize(x.float(), dim=1)`` runs, reading
+    # the low-precision tensor once instead of materializing an fp32 copy.
+    denom = torch.linalg.vector_norm(x, dim=1, keepdim=True, dtype=torch.float32).clamp_min_(1e-12)
+
+    if _kernels_allowed(x):
+        rows = _as_rows(x, denom)
+        if rows is not None:
+            x_rows, denom_rows, restore = rows
+            out = rn.rms_norm_scale(x_rows, denom_rows, gamma.reshape(-1), norm.scale, silu=silu)
+            if out is not None:
+                return restore(out)
+
+    # Exact PyTorch restructuring: the divide computes in fp32 opmath and rounds
+    # once on store, exactly like ``(x.float() / denom).to(x.dtype)``; the
+    # epilogue is upstream's own expression.
+    normalized = torch.empty_like(x)
+    torch.div(x, denom, out=normalized)
+    out = normalized * norm.scale * gamma + norm_bias
+    return F.silu(out) if silu else out
+
+
+def rms_norm_forward(self: nn.Module, x: torch.Tensor) -> torch.Tensor:
+    out = rms_norm_fastpath(self, x)
+    if out is not None:
+        return out
+    return type(self).forward(self, x)
+
+
+def _norm_act(
+    norm: nn.Module, act: nn.Module, x: torch.Tensor, *, pending_bias: torch.Tensor | None = None
+) -> torch.Tensor:
+    """``act(norm(x + pending_bias))`` with the SiLU (and the bias) folded into the norm kernel when possible.
+
+    ``pending_bias`` is the per-channel bias the producing convolution left
+    un-added (see :func:`_run_cached_causal_conv` with ``return_bias=True``).
+    """
+    cfg = getattr(norm, CFG_ATTR, None)
+    if cfg is not None and is_diffusers_rms_norm(norm):
+        # SiLU is folded into the norm kernel when the epilogue was proven exact
+        # for this dtype, or always under the (tolerance-based) channels_last level.
+        fuse = type(act) is nn.SiLU and not act.inplace and (x.dtype in cfg.fused_silu_dtypes or cfg.channels_last)
+        out = rms_norm_fastpath(norm, x, silu=fuse, bias=pending_bias)
+        if out is not None:
+            return out if fuse else act(out)
+    if pending_bias is not None:
+        x = _add_channel_bias(x, pending_bias)
+    return act(norm(x))
+
+
+# --------------------------------------------------------------------------- #
+# WanCausalConv3d and its cached call sites
+# --------------------------------------------------------------------------- #
+
+
+def causal_conv_forward(self: WanCausalConv3d, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
+    padding = self._padding
+    if not any(padding):
+        # 1x1x1 convolutions: upstream still runs ``F.pad(x, [0] * 6)``, a full
+        # fill + copy. Skipping it is the identity.
+        return nn.Conv3d.forward(self, x)
+    if _kernels_allowed(x) and (cache_x is None or (cache_x.device == x.device and cache_x.dtype == x.dtype)):
+        assembled = dm.cat_pad_5d(x, cache_x if padding[4] > 0 else None, padding)
+        if assembled is not None:
+            return nn.Conv3d.forward(self, assembled)
+    return WanCausalConv3d.forward(self, x, cache_x)
+
+
+def _conv_without_bias(conv: nn.Conv3d, x: torch.Tensor) -> torch.Tensor:
+    """The cuDNN call ``nn.Conv3d.forward`` makes, minus ATen's separate ``add_(bias)``."""
+    return F.conv3d(x, conv.weight, None, conv.stride, conv.padding, conv.dilation, conv.groups)
+
+
+def _deferred_conv_bias(conv: nn.Module, output: torch.Tensor) -> torch.Tensor | None:
+    """Round a deferred bias as autocast would inside the producing convolution.
+
+    Use the output dtype: autocast can change it even when the input and
+    parameters are fp32. Matching dtypes keep the original parameter without
+    allocating a copy.
+    """
+    return None if conv.bias is None else conv.bias.to(dtype=output.dtype)
+
+
+# Each cache retains at most this many verdicts per live convolution module.
+_MAX_VERDICTS_PER_CONV = 128
+# Per conv module: {(x.shape, cache_frames, dtype, layout): cuDNN padding is bitwise
+# identical to a pre-padded input}. ``False`` routes that shape to the padded path.
+_SPATIAL_PAD_VERDICTS: WeakKeyDictionary[nn.Module, dict[tuple, bool]] = WeakKeyDictionary()
+# Per conv module: {(x.shape, cache_frames, dtype): the convolution of the channels-last
+# assembled input is bitwise identical to the channels-first one}.
+_CONV_OUT_LAYOUT_VERDICTS: WeakKeyDictionary[nn.Module, dict[tuple, bool]] = WeakKeyDictionary()
+
+
+def _record_verdict(verdicts: dict[tuple, bool], key: tuple, verdict: bool) -> None:
+    """Remember a verdict, evicting the oldest entry when the per-conv cache is full."""
+    verdicts[key] = verdict
+    if len(verdicts) > _MAX_VERDICTS_PER_CONV:
+        del verdicts[next(iter(verdicts))]
+
+
+def _bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool:
+    if a.shape != b.shape or a.dtype != b.dtype:
+        return False
+    int_dtype = torch.int16 if a.element_size() == 2 else torch.int32
+    return bool(torch.equal(a.contiguous().view(int_dtype), b.contiguous().view(int_dtype)))
+
+
+def _conv_with_spatial_padding(
+    conv: nn.Conv3d,
+    assembled: torch.Tensor,
+    bias: torch.Tensor | None,
+    verdicts: dict[tuple, bool],
+    key: tuple,
+) -> torch.Tensor:
+    """Convolve the time-assembled input, letting cuDNN apply the spatial zero padding.
+
+    Upstream pads the input tensor explicitly and convolves with ``padding=0``.
+    Handing the spatial padding to cuDNN avoids materializing the padded copy,
+    but is only bit-identical if cuDNN selects the same kernel for both
+    formulations, so the first call for each (conv, shape) runs both and
+    compares bitwise; later calls reuse the verdict until it is evicted.
+    """
+    pad_height, pad_width = conv._padding[2], conv._padding[0]
+    fast = F.conv3d(assembled, conv.weight, bias, conv.stride, (0, pad_height, pad_width), conv.dilation, conv.groups)
+    if key in verdicts:
+        return fast
+    padded = F.pad(assembled, (pad_width, pad_width, pad_height, pad_height, 0, 0))
+    reference = F.conv3d(padded, conv.weight, bias, conv.stride, 0, conv.dilation, conv.groups)
+    verdict = _bitwise_equal(fast, reference)
+    _record_verdict(verdicts, key, verdict)
+    if not verdict:
+        logger.info(
+            "cuDNN spatial padding is not bitwise identical to pre-padded input for %s %s; "
+            "using the pre-padded path for this shape",
+            type(conv).__name__,
+            tuple(assembled.shape),
+        )
+        return reference
+    return fast
+
+
+def _layout_tag(x: torch.Tensor) -> str:
+    return "channels_last" if x.shape[1] > 1 and x.stride(1) == 1 else "channels_first"
+
+
+def _run_cached_causal_conv(
+    conv: nn.Module,
+    x: torch.Tensor,
+    cache_list: list[Any],
+    index: int,
+    *,
+    return_bias: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:
+    """Run a causal conv against ``cache_list[index]`` and refresh that slot.
+
+    Reproduces the upstream call-site pattern::
+
+        cache_x = x[:, :, -CACHE_T:].clone()
+        if cache_x.shape[2] < CACHE_T and feat_cache[idx] is not None:
+            cache_x = cat([feat_cache[idx][:, :, -1:], cache_x], dim=2)   # or zeros for "Rep"
+        x = conv(x, feat_cache[idx])
+        feat_cache[idx] = cache_x
+
+    The fused path stores the last ``CACHE_T`` frames of the assembled conv
+    input instead. On the very first chunk that is ``[0, x]`` where upstream
+    stores ``[x]``; the next conv input is identical either way because the
+    causal padding shrinks by the cache length.
+
+    With ``return_bias=True`` the fused path leaves the convolution bias unadded
+    and returns it alongside the output so the consumer can fold it into its own
+    kernel; the reference path returns ``(output, None)``.
+    """
+    cache = cache_list[index]
+    is_repeat_marker = isinstance(cache, str)
+    payload = cache if isinstance(cache, torch.Tensor) else None
+    if (
+        type(conv) is WanCausalConv3d
+        and _kernels_allowed(x)
+        and x.dim() == 5
+        and (payload is None or (payload.device == x.device and payload.dtype == x.dtype))
+    ):
+        fold_bias = return_bias and conv.bias is not None
+        bias = None if fold_bias else conv.bias
+
+        # Preferred: temporal concat only (aligned plane copies) + cuDNN spatial
+        # padding, once verified bitwise for this (conv, shape).
+        verdicts = _SPATIAL_PAD_VERDICTS.setdefault(conv, {})
+        key = (tuple(x.shape), 0 if payload is None else payload.shape[2], x.dtype, _layout_tag(x))
+        if verdicts.get(key, True):
+            pair = dm.cat_time_5d(x, payload, conv._padding[4], keep_cache_frames=CACHE_T)
+            if pair is not None:
+                assembled, cache_list[index] = pair
+                out = _conv_with_spatial_padding(conv, assembled, bias, verdicts, key)
+                return (out, _deferred_conv_bias(conv, out)) if return_bias else out
+
+        pair = dm.cat_pad_5d(x, payload, conv._padding, keep_cache_frames=CACHE_T)
+        if pair is not None:
+            conv_input, cache_list[index] = pair
+            out = F.conv3d(conv_input, conv.weight, bias, conv.stride, conv.padding, conv.dilation, conv.groups)
+            return (out, _deferred_conv_bias(conv, out)) if return_bias else out
+
+    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+    if cache_x.shape[2] < CACHE_T and payload is not None:
+        cache_x = torch.cat([payload[:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+    elif cache_x.shape[2] < CACHE_T and is_repeat_marker:
+        cache_x = torch.cat([torch.zeros_like(cache_x), cache_x], dim=2)
+
+    out = conv(x) if payload is None else conv(x, payload)
+    cache_list[index] = cache_x
+    return (out, None) if return_bias else out
+
+
+def _run_conv_out_channels_last(
+    conv: nn.Module, x: torch.Tensor, cache_list: list[Any], index: int
+) -> torch.Tensor | None:
+    """``conv_out`` on a channels-last copy of its padded input, once verified bitwise; ``None`` to decline.
+
+    In the lossless level the decoder runs channels-first and cuDNN transposes
+    every convolution input to NDHWC internally. ``conv_out`` is the one
+    convolution whose spatial padding cuDNN does not reproduce bitwise, so its
+    input is assembled pre-padded anyway; assembling it channels-last directly
+    (one fused transpose in the pixels kernel) saves the separate channels-first
+    assembly and cuDNN's transpose. The output is made contiguous, so nothing
+    downstream sees a layout change. The first call per (conv, shape) also runs
+    the channels-first formulation and compares bitwise; a mismatch routes that
+    shape back to the standard path while its verdict is cached. Evicted shapes
+    are checked again on their next use.
+    """
+    cache = cache_list[index]
+    payload = cache if isinstance(cache, torch.Tensor) else None
+    if type(conv) is not WanCausalConv3d or not _kernels_allowed(x) or x.dim() != 5 or x.stride(1) == 1:
+        return None
+    if payload is not None and (payload.device != x.device or payload.dtype != x.dtype):
+        return None
+    verdicts = _CONV_OUT_LAYOUT_VERDICTS.setdefault(conv, {})
+    key = (tuple(x.shape), 0 if payload is None else payload.shape[2], x.dtype)
+    if not verdicts.get(key, True):
+        return None
+    pair = dm.cat_pad_5d(x, payload, conv._padding, keep_cache_frames=CACHE_T, channels_last_output=True)
+    if pair is None:
+        return None
+    assembled, cache_list[index] = pair
+    out = F.conv3d(assembled, conv.weight, conv.bias, conv.stride, conv.padding, conv.dilation, conv.groups)
+    out = out.contiguous()
+    if key in verdicts:
+        return out
+    # The channels-first path convolves exactly this tensor in NCDHW layout.
+    reference = F.conv3d(
+        assembled.contiguous(), conv.weight, conv.bias, conv.stride, conv.padding, conv.dilation, conv.groups
+    )
+    verdict = _bitwise_equal(out, reference)
+    _record_verdict(verdicts, key, verdict)
+    if not verdict:
+        logger.info(
+            "channels-last conv_out input is not bitwise identical to the channels-first one for %s; "
+            "using the channels-first path for this shape",
+            tuple(assembled.shape),
+        )
+        return reference
+    return out
+
+
+def _residual_add(
+    x: torch.Tensor,
+    bias_x: torch.Tensor | None,
+    h: torch.Tensor,
+    bias_h: torch.Tensor | None,
+) -> torch.Tensor:
+    """``(x + bias_x) + (h + bias_h)`` where the biases are the un-added conv biases."""
+    if bias_x is None and bias_h is None:
+        return x + h
+    if _kernels_allowed(x):
+        out = dm.add_bias_residual(x, bias_x, h, bias_h)
+        if out is not None:
+            return out
+    # Out-of-place ``x + bias`` rounds exactly like ATen's in-place ``add_``.
+    if bias_x is not None:
+        x = x + bias_x.view(1, -1, 1, 1, 1)
+    if bias_h is not None:
+        h = h + bias_h.view(1, -1, 1, 1, 1)
+    return x + h
+
+
+# --------------------------------------------------------------------------- #
+# WanResample / WanUpsample
+# --------------------------------------------------------------------------- #
+
+
+def _interleave_time(x: torch.Tensor, batch: int, channels: int, frames: int, height: int, width: int) -> torch.Tensor:
+    """Upstream ``reshape(b, 2, c, t, h, w)`` + ``stack(..., 3)`` + ``reshape`` as one strided copy.
+
+    ``out[:, :, 2i] = x[:, :c, i]`` and ``out[:, :, 2i + 1] = x[:, c:, i]``; pure
+    data movement that also preserves a channels_last_3d layout (``torch.stack``
+    on the 6-D view always produced a contiguous tensor).
+    """
+    memory_format = torch.contiguous_format
+    if channels > 1 and x.is_contiguous(memory_format=torch.channels_last_3d):
+        memory_format = torch.channels_last_3d
+    out = torch.empty(
+        (batch, channels, frames * 2, height, width),
+        dtype=x.dtype,
+        device=x.device,
+        memory_format=memory_format,
+    )
+    out[:, :, 0::2].copy_(x[:, :channels])
+    out[:, :, 1::2].copy_(x[:, channels:])
+    return out
+
+
+def _merge_batch_and_frames(x: torch.Tensor) -> torch.Tensor:
+    """``(b, c, t, h, w)`` -> ``(b * t, c, h, w)`` view, keeping a channels-last layout recognizable.
+
+    For a channels_last_3d tensor with ``b * t == 1``, ``permute().reshape()`` hands
+    the size-1 batch dimension an arbitrary stride that PyTorch's layout heuristic
+    (``suggest_memory_format``) does not accept as channels_last, so ops without a
+    weight to vote for the layout, such as the nearest upsample, would allocate
+    NCHW output and force the next convolution to transpose it. Building the same
+    view with canonical strides avoids that; the values are untouched.
+    """
+    b, c, t, h, w = x.shape
+    if c > 1 and x.is_contiguous(memory_format=torch.channels_last_3d):
+        return x.as_strided((b * t, c, h, w), (h * w * c, 1, w * c, c))
+    return x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
+
+
+def _split_batch_and_frames(x: torch.Tensor, batch: int, frames: int) -> torch.Tensor:
+    """``(b * t, c, h, w)`` -> ``(b, c, t, h, w)`` view; see :func:`_merge_batch_and_frames`."""
+    _, c, h, w = x.shape
+    if c > 1 and x.is_contiguous(memory_format=torch.channels_last):
+        return x.as_strided((batch, c, frames, h, w), (frames * h * w * c, 1, h * w * c, w * c, c))
+    return x.view(batch, frames, c, h, w).permute(0, 2, 1, 3, 4)
+
+
+def _is_upsample_conv_pair(resample: nn.Module) -> bool:
+    """``nn.Sequential(WanUpsample, Conv2d)`` with a bias and zero padding, as the up-samplers are built."""
+    if type(resample) is not nn.Sequential or len(resample) != 2:
+        return False
+    upsample, conv = resample[0], resample[1]
+    return (
+        isinstance(upsample, nn.Upsample)
+        and type(conv) is nn.Conv2d
+        and conv.bias is not None
+        and conv.padding_mode == "zeros"
+    )
+
+
+def resample_forward(
+    self: WanResample,
+    x: torch.Tensor,
+    feat_cache: list[Any] | None = None,
+    feat_idx: list[int] | None = None,
+    *,
+    return_bias: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:
+    """``WanResample.forward``; with ``return_bias=True`` the upsampler's ``Conv2d`` bias may be left un-added.
+
+    The caller then receives ``(output, bias)`` and must fold ``bias`` into its
+    consumer (or add it); ``(output, None)`` means the bias is already applied.
+    """
+    if feat_idx is None:
+        feat_idx = [0]
+    batch, channels, frames, height, width = x.size()
+    if self.mode == "upsample3d" and feat_cache is not None:
+        index = feat_idx[0]
+        if feat_cache[index] is None:
+            feat_cache[index] = "Rep"
+            feat_idx[0] += 1
+        else:
+            x = _run_cached_causal_conv(self.time_conv, x, feat_cache, index)
+            feat_idx[0] += 1
+            x = _interleave_time(x, batch, channels, frames, height, width)
+
+    frames = x.shape[2]
+    x = _merge_batch_and_frames(x)
+    pending_bias = None
+    if return_bias and _is_upsample_conv_pair(self.resample) and _kernels_allowed(x):
+        upsample, conv = self.resample[0], self.resample[1]
+        x = F.conv2d(upsample(x), conv.weight, None, conv.stride, conv.padding, conv.dilation, conv.groups)
+        pending_bias = _deferred_conv_bias(conv, x)
+    else:
+        x = self.resample(x)
+    x = _split_batch_and_frames(x, batch, frames)
+
+    if self.mode == "downsample3d" and feat_cache is not None:
+        index = feat_idx[0]
+        if feat_cache[index] is None:
+            feat_cache[index] = x.clone()
+            feat_idx[0] += 1
+        else:
+            cache_x = x[:, :, -1:, :, :].clone()
+            x = self.time_conv(torch.cat([feat_cache[index][:, :, -1:, :, :], x], dim=2))
+            feat_cache[index] = cache_x
+            feat_idx[0] += 1
+    return (x, pending_bias) if return_bias else x
+
+
+def _is_nearest_2x(upsample: nn.Upsample) -> bool:
+    """True when ``upsample`` doubles height and width by nearest-neighbour selection."""
+    if upsample.size is not None or upsample.mode not in ("nearest", "nearest-exact"):
+        return False
+    scale = upsample.scale_factor
+    if isinstance(scale, (tuple, list)):
+        return len(scale) == 2 and all(float(factor) == 2.0 for factor in scale)
+    return scale is not None and float(scale) == 2.0
+
+
+def upsample_forward(self: WanUpsample, x: torch.Tensor) -> torch.Tensor:
+    # nearest-exact only replicates input values, so ``x.float()`` + ``type_as``
+    # is the identity; run the interpolation in the activation dtype.
+    if x.dtype in _NORM_DTYPES:
+        if _kernels_allowed(x) and _is_nearest_2x(self):
+            out = up.upsample_nearest_2x(x)
+            if out is not None:
+                return out
+        return nn.Upsample.forward(self, x)
+    return WanUpsample.forward(self, x)
+
+
+# --------------------------------------------------------------------------- #
+# Blocks and decoder
+# --------------------------------------------------------------------------- #
+
+
+def residual_block_forward(
+    self: WanResidualBlock,
+    x: torch.Tensor,
+    feat_cache: list[Any] | None = None,
+    feat_idx: list[int] | None = None,
+) -> torch.Tensor:
+    if feat_idx is None:
+        feat_idx = [0]
+    shortcut = self.conv_shortcut
+    residual_bias = None
+    if (
+        type(shortcut) is WanCausalConv3d
+        and shortcut.bias is not None
+        and not any(shortcut._padding)
+        and _kernels_allowed(x)
+    ):
+        residual = _conv_without_bias(shortcut, x)
+        residual_bias = _deferred_conv_bias(shortcut, residual)
+    else:
+        residual = shortcut(x)
+    x = _norm_act(self.norm1, self.nonlinearity, x)
+    # conv1 feeds only norm2. At the channels_last level the single-pass norm
+    # kernel adds the conv bias itself (rounded like ATen's ``add_``), so the
+    # convolution runs without it; the lossless level keeps the separate add
+    # because ``vector_norm`` must see the same bytes as upstream.
+    cfg = getattr(self, CFG_ATTR, None)
+    conv1_bias = None
+    if feat_cache is not None:
+        index = feat_idx[0]
+        if cfg is not None and cfg.channels_last:
+            x, conv1_bias = _run_cached_causal_conv(self.conv1, x, feat_cache, index, return_bias=True)
+        else:
+            x = _run_cached_causal_conv(self.conv1, x, feat_cache, index)
+        feat_idx[0] += 1
+    else:
+        x = self.conv1(x)
+
+    x = self.dropout(_norm_act(self.norm2, self.nonlinearity, x, pending_bias=conv1_bias))
+    conv2_bias = None
+    if feat_cache is not None:
+        index = feat_idx[0]
+        x, conv2_bias = _run_cached_causal_conv(self.conv2, x, feat_cache, index, return_bias=True)
+        feat_idx[0] += 1
+    else:
+        x = self.conv2(x)
+    return _residual_add(x, conv2_bias, residual, residual_bias)
+
+
+def residual_up_block_forward(
+    self: WanResidualUpBlock,
+    x: torch.Tensor,
+    feat_cache: list[Any] | None = None,
+    feat_idx: list[int] | None = None,
+    first_chunk: bool = False,
+) -> torch.Tensor:
+    if feat_idx is None:
+        feat_idx = [0]
+    # Upstream clones ``x`` for the shortcut; nothing below mutates it in place.
+    shortcut_source = x
+    for resnet in self.resnets:
+        if feat_cache is None:
+            x = resnet(x)
+        else:
+            x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx)
+
+    shortcut = self.avg_shortcut
+    # The upsampler's Conv2d output feeds only ``x + DupUp3D(...)``, so its bias
+    # is folded into that kernel (exact: rounded like ATen's ``add_`` first).
+    fuse_shortcut = type(shortcut) is DupUp3D and _kernels_allowed(x) and x.dtype == shortcut_source.dtype
+    conv_bias = None
+    if self.upsampler is not None:
+        if feat_cache is None:
+            x = self.upsampler(x)
+        elif fuse_shortcut and type(self.upsampler) is WanResample and hasattr(self.upsampler, CFG_ATTR):
+            x, conv_bias = resample_forward(self.upsampler, x, feat_cache, feat_idx, return_bias=True)
+        else:
+            x = self.upsampler(x, feat_cache=feat_cache, feat_idx=feat_idx)
+
+    if shortcut is None:
+        return x
+    if fuse_shortcut:
+        fused = dm.dup_up3d_add(
+            x,
+            shortcut_source,
+            shortcut.factor_t,
+            shortcut.factor_s,
+            shortcut.repeats,
+            first_chunk,
+            main_bias=conv_bias,
+        )
+        if fused is not None:
+            return fused
+    if conv_bias is not None:
+        x = _add_channel_bias(x, conv_bias)
+    return x + shortcut(shortcut_source, first_chunk=first_chunk)
+
+
+def decoder_forward(
+    self: WanDecoder3d,
+    x: torch.Tensor,
+    feat_cache: list[Any] | None = None,
+    feat_idx: list[int] | None = None,
+    first_chunk: bool = False,
+) -> torch.Tensor:
+    if feat_idx is None:
+        feat_idx = [0]
+    if feat_cache is not None:
+        index = feat_idx[0]
+        x = _run_cached_causal_conv(self.conv_in, x, feat_cache, index)
+        feat_idx[0] += 1
+    else:
+        x = self.conv_in(x)
+
+    x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx)
+    for up_block in self.up_blocks:
+        x = up_block(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
+
+    x = _norm_act(self.norm_out, self.nonlinearity, x)
+    if feat_cache is not None:
+        index = feat_idx[0]
+        # Channels-first (lossless) only: channels-last activations already take
+        # the pixels kernel and transpose-free cuDNN path inside the standard call.
+        out = _run_conv_out_channels_last(self.conv_out, x, feat_cache, index)
+        if out is None:
+            out = _run_cached_causal_conv(self.conv_out, x, feat_cache, index)
+        x = out
+        feat_idx[0] += 1
+    else:
+        x = self.conv_out(x)
+    return x
+
+
+__all__ = [
+    "CFG_ATTR",
+    "FastPathConfig",
+    "causal_conv_forward",
+    "decoder_forward",
+    "is_diffusers_rms_norm",
+    "resample_forward",
+    "residual_block_forward",
+    "residual_up_block_forward",
+    "rms_norm_fastpath",
+    "rms_norm_forward",
+    "upsample_forward",
+]

vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/install.py

@@ -0,0 +1,286 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Instance-level installer for the Wan VAE decoder fast path."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from contextlib import ExitStack
+from dataclasses import dataclass, field
+from functools import partial
+from types import MethodType
+from typing import Any
+
+import torch
+import torch.nn as nn
+from diffusers.models.autoencoders import AutoencoderKLWan
+from diffusers.models.autoencoders.autoencoder_kl_wan import (
+    DupUp3D,
+    WanCausalConv3d,
+    WanDecoder3d,
+    WanResample,
+    WanResidualBlock,
+    WanResidualUpBlock,
+    WanUpsample,
+)
+from vllm.logger import init_logger
+
+from . import forwards
+from . import triton_rms_norm as rn
+
+logger = init_logger(__name__)
+
+VAE_FAST_PATH_LEVELS: tuple[str, ...] = ("off", "lossless", "channels_last")
+
+REPORT_ATTR = "_vllm_omni_wan_fastpath_report"
+_UNDO_ATTR = "_vllm_omni_wan_fastpath_undo"
+
+_REPLACEMENT_FORWARDS = {
+    WanDecoder3d: forwards.decoder_forward,
+    WanCausalConv3d: forwards.causal_conv_forward,
+    WanResidualBlock: forwards.residual_block_forward,
+    WanResidualUpBlock: forwards.residual_up_block_forward,
+    WanResample: forwards.resample_forward,
+    WanUpsample: forwards.upsample_forward,
+}
+
+
+@dataclass(frozen=True)
+class WanVaeFastPathReport:
+    """What :func:`install_wan_vae_fastpath` did to one VAE instance."""
+
+    level: str
+    installed: bool
+    reason: str | None = None
+    patched: Mapping[str, int] = field(default_factory=dict)
+    fused_silu_dtypes: tuple[str, ...] = ()
+    channels_last: bool = False
+
+
+def is_installed(vae: nn.Module) -> bool:
+    report = getattr(vae, REPORT_ATTR, None)
+    return report is not None and report.installed
+
+
+def _skip(level: str, reason: str) -> WanVaeFastPathReport:
+    logger.info("Wan VAE fast path (%s) not installed: %s", level, reason)
+    return WanVaeFastPathReport(level=level, installed=False, reason=reason)
+
+
+def _convert_conv_memory_format(modules: list[nn.Module], *, channels_last: bool) -> int:
+    """Convert conv weights to (or back from) channels-last layouts.
+
+    Only 4D/5D parameters change (``nn.Module.to(memory_format=...)`` skips the
+    rest), so parameter names, dtypes and devices are untouched. Inputs are not
+    converted: cuDNN selects the channels-last algorithm when either operand
+    suggests it, so the layout propagates from the first convolution.
+    """
+    count = 0
+    for module in modules:
+        if isinstance(module, nn.Conv3d):
+            module.to(memory_format=torch.channels_last_3d if channels_last else torch.contiguous_format)
+            count += 1
+        elif isinstance(module, (nn.Conv2d, nn.ConvTranspose2d)):
+            module.to(memory_format=torch.channels_last if channels_last else torch.contiguous_format)
+            count += 1
+    return count
+
+
+@torch.no_grad()
+def _restore_module_layouts(
+    module: nn.Module,
+    tensor_strides: Mapping[str, tuple[int, ...]],
+    grad_strides: Mapping[str, tuple[int, ...]],
+) -> None:
+    """Restore layouts of the current tensors, retaining any updates made while installed."""
+
+    def restore(tensor: torch.Tensor, strides: tuple[int, ...]) -> None:
+        if tensor.stride() != strides:
+            restored = torch.empty_strided(tensor.shape, strides, device=tensor.device, dtype=tensor.dtype)
+            restored.copy_(tensor)
+            tensor.data = restored
+
+    for name, strides in tensor_strides.items():
+        restore(getattr(module, name), strides)
+    for name, strides in grad_strides.items():
+        grad = getattr(module, name).grad
+        if grad is not None:
+            restore(grad, strides)
+
+
+def install_wan_vae_fastpath(vae: nn.Module, *, level: str = "lossless") -> WanVaeFastPathReport:
+    """Install the Wan decoder fast path on one loaded ``AutoencoderKLWan``. Idempotent.
+
+    ``level``:
+      * ``"off"``: do nothing.
+      * ``"lossless"``: bind the bit-exact replacement forwards (Tier 1).
+      * ``"channels_last"``: Tier 1 plus channels-last conv weights (not bit-exact).
+
+    The installer refuses (and reports why) when the VAE is not a diffusers Wan
+    VAE, or when spatially-sharded decode is configured or already installed:
+    ``wan_spatial_shard`` rebinds the same forwards and replaces the causal
+    convolutions with halo-exchanging variants.
+    """
+    if level not in VAE_FAST_PATH_LEVELS:
+        raise ValueError(f"vae_fast_path must be one of {list(VAE_FAST_PATH_LEVELS)}, got {level!r}")
+    existing = getattr(vae, REPORT_ATTR, None)
+    if existing is not None:
+        if existing.level != level:
+            logger.warning(
+                "Wan VAE fast path already installed at level %r; ignoring request for level %r",
+                existing.level,
+                level,
+            )
+        return existing
+    if level == "off":
+        return WanVaeFastPathReport(level=level, installed=False, reason="disabled")
+    if not isinstance(vae, AutoencoderKLWan):
+        return _skip(level, f"{type(vae).__name__} is not a diffusers AutoencoderKLWan")
+    decoder = getattr(vae, "decoder", None)
+    if type(decoder) is not WanDecoder3d:
+        return _skip(level, f"decoder is {type(decoder).__name__}, not diffusers WanDecoder3d")
+    if getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False):
+        return _skip(level, "spatial-shard decode is already installed on this VAE")
+    parallel_mode = getattr(getattr(vae, "distributed_executor", None), "parallel_mode", "tile")
+    if isinstance(parallel_mode, str) and parallel_mode.startswith("spatial_shard"):
+        return _skip(level, f"vae_parallel_mode={parallel_mode!r} is not supported")
+
+    modules = [(f"decoder.{name}" if name else "decoder", module) for name, module in decoder.named_modules()]
+    post_quant_conv = getattr(vae, "post_quant_conv", None)
+    if type(post_quant_conv) is WanCausalConv3d:
+        modules.append(("post_quant_conv", post_quant_conv))
+    bindings: list[tuple[nn.Module, Callable[..., Any]]] = []
+    convs: list[nn.Module] = []
+    bypassed: set[nn.Module] = set()
+    for _, module in modules:
+        replacement = _REPLACEMENT_FORWARDS.get(type(module))
+        if forwards.is_diffusers_rms_norm(module):
+            replacement = forwards.rms_norm_forward
+            bypassed.add(module)
+        if replacement is not None:
+            bindings.append((module, replacement))
+        if isinstance(module, (nn.Conv2d, nn.Conv3d)):
+            convs.append(module)
+        if type(module) in (WanResample, DupUp3D, nn.SiLU) or (
+            type(module) is WanCausalConv3d and module is not post_quant_conv
+        ):
+            bypassed.add(module)
+        if type(module) is WanResample and forwards._is_upsample_conv_pair(module.resample):
+            bypassed.update((module.resample, module.resample[1]))
+
+    replaced = {module for module, _ in bindings}
+    for name, module in modules:
+        if module in replaced or module in bypassed:
+            current_forward = module.forward
+            # An explicit alias of the class method is safe and will be restored
+            # as the same instance attribute. Arbitrary wrappers cannot be
+            # composed with replacement forwards or direct functional calls.
+            if not (
+                isinstance(current_forward, MethodType)
+                and current_forward.__self__ is module
+                and current_forward.__func__ is type(module).forward
+            ):
+                return _skip(level, f"{name} has a custom forward that the fast path would replace or bypass")
+        if module in bypassed and (module._forward_pre_hooks or module._forward_hooks):
+            return _skip(level, f"{name} has forward hooks that the fast path would bypass")
+
+    device = next((p.device for p in decoder.parameters()), torch.device("cpu"))
+    fused_silu_dtypes: frozenset[torch.dtype] = frozenset()
+    if device.type == "cuda":
+        fused_silu_dtypes = frozenset(
+            dtype for dtype in (torch.bfloat16, torch.float16) if rn.silu_epilogue_is_exact(device, dtype)
+        )
+    cfg = forwards.FastPathConfig(fused_silu_dtypes=fused_silu_dtypes, channels_last=level == "channels_last")
+
+    with ExitStack() as rollback:
+        undo: list[Callable[[], None]] = []
+        patched: dict[str, int] = {}
+
+        def bind(module: nn.Module, forward: Callable[..., Any]) -> None:
+            original = {
+                name: module.__dict__[name] for name in ("forward", forwards.CFG_ATTR) if name in module.__dict__
+            }
+
+            def restore() -> None:
+                module.__dict__.pop("forward", None)
+                module.__dict__.pop(forwards.CFG_ATTR, None)
+                module.__dict__.update(original)
+
+            rollback.callback(restore)
+            undo.append(restore)
+            module.forward = MethodType(forward, module)
+            setattr(module, forwards.CFG_ATTR, cfg)
+            patched[type(module).__name__] = patched.get(type(module).__name__, 0) + 1
+
+        for module, replacement in bindings:
+            bind(module, replacement)
+
+        if cfg.channels_last:
+            # Keep the original storage only until installation commits. An OOM
+            # may interrupt Module.to partway through a module, so rollback must
+            # restore its tensors without allocating or calling Module.to again.
+            for module in convs:
+                for owner in module.modules():
+                    tensor_strides = {
+                        name: tensor.stride()
+                        for name, tensor in (
+                            *owner.named_parameters(recurse=False),
+                            *owner.named_buffers(recurse=False),
+                        )
+                    }
+                    grad_strides = {
+                        name: parameter.grad.stride()
+                        for name, parameter in owner.named_parameters(recurse=False)
+                        if parameter.grad is not None
+                    }
+                    undo.append(partial(_restore_module_layouts, owner, tensor_strides, grad_strides))
+                    for name, parameter in owner.named_parameters(recurse=False):
+                        rollback.callback(owner._parameters.__setitem__, name, parameter)
+                        rollback.callback(setattr, parameter, "data", parameter.data)
+                        rollback.callback(setattr, parameter, "grad", parameter.grad)
+                        if parameter.grad is not None:
+                            rollback.callback(setattr, parameter.grad, "data", parameter.grad.data)
+                    for name, buffer in owner.named_buffers(recurse=False):
+                        rollback.callback(owner._buffers.__setitem__, name, buffer)
+                        rollback.callback(setattr, buffer, "data", buffer.data)
+            converted = _convert_conv_memory_format(convs, channels_last=True)
+            logger.info("Wan VAE decoder: %d convolution weights converted to channels-last layout", converted)
+
+        report = WanVaeFastPathReport(
+            level=level,
+            installed=True,
+            patched=dict(patched),
+            fused_silu_dtypes=tuple(str(dtype).removeprefix("torch.") for dtype in sorted(fused_silu_dtypes, key=str)),
+            channels_last=cfg.channels_last,
+        )
+        rollback.callback(vae.__dict__.pop, REPORT_ATTR, None)
+        rollback.callback(vae.__dict__.pop, _UNDO_ATTR, None)
+        setattr(vae, _UNDO_ATTR, undo)
+        setattr(vae, REPORT_ATTR, report)
+        logger.info(
+            "Wan VAE fast path (%s) installed: patched=%s fused_silu=%s channels_last=%s",
+            level,
+            report.patched,
+            report.fused_silu_dtypes or "off",
+            report.channels_last,
+        )
+        rollback.pop_all()
+        return report
+
+
+def uninstall_wan_vae_fastpath(vae: nn.Module) -> None:
+    """Restore original forwards and tensor layouts, retaining current weight values."""
+    undo = vae.__dict__.pop(_UNDO_ATTR, None)
+    vae.__dict__.pop(REPORT_ATTR, None)
+    for restore in reversed(undo or []):
+        restore()
+
+
+__all__ = [
+    "REPORT_ATTR",
+    "VAE_FAST_PATH_LEVELS",
+    "WanVaeFastPathReport",
+    "install_wan_vae_fastpath",
+    "is_installed",
+    "uninstall_wan_vae_fastpath",
+]

vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_data_movement.py

@@ -0,0 +1,1000 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# ruff: noqa: N803
+"""Bit-exact data-movement kernels for the diffusers Wan causal VAE decoder.
+
+Adapted from SGLang PR #34125 (https://github.com/sgl-project/sglang/pull/34125).
+
+The kernels in this module only move values, fill zeros, or perform the same
+two-operand addition as the PyTorch reference path, so the decoder numerics are
+unchanged. They preserve the input memory layout (contiguous or
+channels_last_3d) instead of selecting a new one, so enabling them does not
+change the cuDNN algorithm or accumulation order of the surrounding
+convolutions.
+
+Every public wrapper returns ``None`` when it declines an input (rank, dtype,
+device, layout mismatch, or too many elements); callers then run the reference
+PyTorch expression.
+"""
+
+from __future__ import annotations
+
+import torch
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+from ._utils import _pick_block_width
+
+_MAX_INT32 = 2**31 - 1
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+_BLOCK_SIZE = 512
+# Elements handled per program by the row/pixel-oriented cat+pad kernels: with
+# 4 warps this is 32 elements per thread, enough independent 16-byte loads to
+# keep HBM busy without register pressure.
+_TILE_ELEMENTS = 4096
+_NUM_WARPS = 4
+# The channels-first tile holds int64 pointer tiles for source, output and cache
+# refresh; 8 warps halve the per-thread register footprint versus 4.
+_ROWS_BLOCK_W = 256
+_ROWS_NUM_WARPS = 8
+# Plane copies: 4096 elements per program on 8 warps = 16 elements (two
+# 16-byte vectors) per thread.
+_PLANE_BLOCK = 4096
+# Square tile edge for assembling channels-first input into channels-last output.
+_TRANSPOSE_BLOCK = 64
+_HAS_INTERLEAVE = HAS_TRITON and hasattr(tl, "interleave")
+
+if HAS_TRITON:
+
+    @triton.jit
+    def _cat_time_planes_kernel(
+        x_ptr,
+        cache_ptr,
+        out_ptr,
+        keep_ptr,
+        plane_size,
+        blocks_per_plane,
+        channels_outer,
+        out_frames,
+        cache_frames,
+        zero_front_frames,
+        keep_start,
+        keep_frames,
+        stride_x_batch,
+        stride_x_channel,
+        stride_x_frame,
+        stride_cache_batch,
+        stride_cache_channel,
+        stride_cache_frame,
+        has_cache: tl.constexpr,
+        has_keep: tl.constexpr,
+        BLOCK: tl.constexpr,
+    ):
+        """Temporal concat ``[zeros | cache | x]`` as whole-plane copies.
+
+        A plane is one frame of one channel (channels-first) or one frame of all
+        channels (channels-last): a contiguous run of ``plane_size`` elements
+        whose start is aligned, so every load and store vectorizes. No spatial
+        padding is written here; the consumer lets cuDNN pad inside the
+        convolution.
+        """
+        pid = tl.program_id(0).to(tl.int64)
+        plane = pid // blocks_per_plane
+        offsets = (pid % blocks_per_plane) * BLOCK + tl.arange(0, BLOCK)
+        mask = offsets < plane_size
+        frame = plane % out_frames
+        rest = plane // out_frames
+        channel = rest % channels_outer
+        batch = rest // channels_outer
+        source_frame = frame - zero_front_frames
+
+        values = tl.zeros([BLOCK], dtype=x_ptr.dtype.element_ty)
+        if source_frame >= cache_frames:
+            x_plane = (
+                x_ptr
+                + batch * stride_x_batch
+                + channel * stride_x_channel
+                + (source_frame - cache_frames) * stride_x_frame
+            )
+            values = tl.load(x_plane + offsets, mask=mask, other=0)
+        else:
+            if has_cache:
+                if source_frame >= 0:
+                    cache_plane = (
+                        cache_ptr
+                        + batch * stride_cache_batch
+                        + channel * stride_cache_channel
+                        + source_frame * stride_cache_frame
+                    )
+                    values = tl.load(cache_plane + offsets, mask=mask, other=0)
+        tl.store(out_ptr + plane * plane_size + offsets, values, mask=mask)
+        if has_keep:
+            if frame >= keep_start:
+                keep_plane = (batch * channels_outer + channel) * keep_frames + (frame - keep_start)
+                tl.store(keep_ptr + keep_plane * plane_size + offsets, values, mask=mask)
+
+    @triton.jit
+    def _cat_pad_rows_kernel(
+        x_ptr,
+        cache_ptr,
+        out_ptr,
+        keep_ptr,
+        channels,
+        frames,
+        height,
+        width,
+        cache_frames,
+        out_frames,
+        out_height,
+        out_width,
+        zero_front_frames,
+        pad_height,
+        pad_width,
+        keep_start,
+        keep_frames,
+        row_blocks,
+        stride_x_batch,
+        stride_x_channel,
+        stride_x_frame,
+        stride_x_height,
+        stride_x_width,
+        stride_cache_batch,
+        stride_cache_channel,
+        stride_cache_frame,
+        stride_cache_height,
+        stride_cache_width,
+        has_cache: tl.constexpr,
+        has_keep: tl.constexpr,
+        ROWS: tl.constexpr,
+        BLOCK_W: tl.constexpr,
+    ):
+        """Channels-first output: one program fills ``ROWS`` consecutive rows of one ``(batch, channel, frame)`` plane.
+
+        All rows of a program share one source (x, cache or causal zero padding),
+        selected once with scalar control flow, so each ``[ROWS, BLOCK_W]`` tile is a
+        single masked load and a single store; zero frames are stored without
+        loading. Row and column coordinates come from the program id and
+        ``tl.arange``; there is no per-element division.
+        """
+        pid = tl.program_id(0).to(tl.int64)
+        row_block = pid % row_blocks
+        rest = pid // row_blocks
+        frame = rest % out_frames
+        rest = rest // out_frames
+        channel = rest % channels
+        batch = rest // channels
+
+        out_rows = row_block * ROWS + tl.arange(0, ROWS)
+        row_mask = out_rows < out_height
+        in_rows = out_rows - pad_height
+        row_valid = row_mask & (in_rows >= 0) & (in_rows < height)
+        source_frame = frame - zero_front_frames
+        plane = (batch * channels + channel) * out_frames + frame
+        out_row_offsets = (plane * out_height + out_rows) * out_width
+        if has_keep:
+            keep_plane = (batch * channels + channel) * keep_frames + (frame - keep_start)
+            keep_row_offsets = (keep_plane * height + in_rows) * width
+
+        for w0 in range(0, out_width, BLOCK_W):
+            cols = w0 + tl.arange(0, BLOCK_W)
+            col_mask = cols < out_width
+            in_cols = cols - pad_width
+            interior = row_valid[:, None] & (col_mask & (in_cols >= 0) & (in_cols < width))[None, :]
+            values = tl.zeros([ROWS, BLOCK_W], dtype=x_ptr.dtype.element_ty)
+            if source_frame >= cache_frames:
+                x_plane = (
+                    x_ptr
+                    + batch * stride_x_batch
+                    + channel * stride_x_channel
+                    + (source_frame - cache_frames) * stride_x_frame
+                )
+                values = tl.load(
+                    x_plane + in_rows[:, None] * stride_x_height + in_cols[None, :] * stride_x_width,
+                    mask=interior,
+                    other=0,
+                )
+            else:
+                if has_cache:
+                    if source_frame >= 0:
+                        cache_plane = (
+                            cache_ptr
+                            + batch * stride_cache_batch
+                            + channel * stride_cache_channel
+                            + source_frame * stride_cache_frame
+                        )
+                        cache_tile = cache_plane + in_rows[:, None] * stride_cache_height
+                        values = tl.load(cache_tile + in_cols[None, :] * stride_cache_width, mask=interior, other=0)
+            store_mask = row_mask[:, None] & col_mask[None, :]
+            tl.store(out_ptr + out_row_offsets[:, None] + cols[None, :], values, mask=store_mask)
+            if has_keep:
+                if frame >= keep_start:
+                    tl.store(keep_ptr + keep_row_offsets[:, None] + in_cols[None, :], values, mask=interior)
+
+    @triton.jit
+    def _cat_pad_pixels_kernel(
+        x_ptr,
+        cache_ptr,
+        out_ptr,
+        keep_ptr,
+        channels,
+        frames,
+        height,
+        width,
+        cache_frames,
+        out_frames,
+        out_height,
+        out_width,
+        zero_front_frames,
+        pad_height,
+        pad_width,
+        keep_start,
+        keep_frames,
+        width_blocks,
+        stride_x_batch,
+        stride_x_channel,
+        stride_x_frame,
+        stride_x_height,
+        stride_x_width,
+        stride_cache_batch,
+        stride_cache_channel,
+        stride_cache_frame,
+        stride_cache_height,
+        stride_cache_width,
+        stride_keep_batch,
+        stride_keep_channel,
+        stride_keep_frame,
+        stride_keep_height,
+        stride_keep_width,
+        has_cache: tl.constexpr,
+        has_keep: tl.constexpr,
+        BLOCK_W: tl.constexpr,
+        BLOCK_C: tl.constexpr,
+    ):
+        """Channels-last output: one program assembles ``BLOCK_W`` pixels x ``BLOCK_C`` channels of one output row.
+
+        The output is written as contiguous channel spans per pixel. The input and
+        the cache refresh go through explicit strides, so a channels-first ``x`` can
+        be assembled straight into a channels-last output (the tile is loaded
+        pixel-contiguous and stored channel-contiguous, a fused transpose) while the
+        cache keeps ``x``'s layout.
+        """
+        pid = tl.program_id(0).to(tl.int64)
+        width_block = pid % width_blocks
+        rest = pid // width_blocks
+        out_row = rest % out_height
+        rest = rest // out_height
+        frame = rest % out_frames
+        batch = rest // out_frames
+        in_row = out_row - pad_height
+        row_valid = (in_row >= 0) & (in_row < height)
+        source_frame = frame - zero_front_frames
+
+        cols = width_block * BLOCK_W + tl.arange(0, BLOCK_W)
+        col_mask = cols < out_width
+        in_col = cols - pad_width
+        interior = col_mask & row_valid & (in_col >= 0) & (in_col < width)
+        out_pixels = (((batch * out_frames + frame) * out_height + out_row) * out_width + cols) * channels
+        if has_keep:
+            keep_pixels = (
+                batch * stride_keep_batch
+                + (frame - keep_start) * stride_keep_frame
+                + in_row * stride_keep_height
+                + in_col * stride_keep_width
+            )
+        for c0 in range(0, channels, BLOCK_C):
+            chans = c0 + tl.arange(0, BLOCK_C)
+            chan_mask = chans < channels
+            load_mask = interior[:, None] & chan_mask[None, :]
+            values = tl.zeros([BLOCK_W, BLOCK_C], dtype=x_ptr.dtype.element_ty)
+            if source_frame >= cache_frames:
+                x_frame = source_frame - cache_frames
+                x_base = x_ptr + batch * stride_x_batch + x_frame * stride_x_frame + in_row * stride_x_height
+                values = tl.load(
+                    x_base + in_col[:, None] * stride_x_width + chans[None, :] * stride_x_channel,
+                    mask=load_mask,
+                    other=0,
+                )
+            else:
+                if has_cache:
+                    if source_frame >= 0:
+                        cache_base = (
+                            cache_ptr
+                            + batch * stride_cache_batch
+                            + source_frame * stride_cache_frame
+                            + in_row * stride_cache_height
+                        )
+                        values = tl.load(
+                            cache_base + in_col[:, None] * stride_cache_width + chans[None, :] * stride_cache_channel,
+                            mask=load_mask,
+                            other=0,
+                        )
+            store_mask = col_mask[:, None] & chan_mask[None, :]
+            tl.store(out_ptr + out_pixels[:, None] + chans[None, :], values, mask=store_mask)
+            if has_keep:
+                if frame >= keep_start:
+                    keep_offsets = keep_pixels[:, None] + chans[None, :] * stride_keep_channel
+                    tl.store(keep_ptr + keep_offsets, values, mask=load_mask)
+
+    @triton.jit
+    def _dup_up3d_add_rows_kernel(
+        main_ptr,
+        source_ptr,
+        bias_ptr,
+        out_ptr,
+        out_channels,
+        out_frames,
+        out_height,
+        out_width,
+        source_width,
+        frame_offset,
+        row_blocks,
+        stride_main_batch,
+        stride_main_channel,
+        stride_main_frame,
+        stride_main_height,
+        stride_source_batch,
+        stride_source_channel,
+        stride_source_frame,
+        stride_source_height,
+        FT: tl.constexpr,
+        FS: tl.constexpr,
+        REPEATS: tl.constexpr,
+        HAS_BIAS: tl.constexpr,
+        ROWS: tl.constexpr,
+        BLOCK_W: tl.constexpr,
+    ):
+        """Channels-first ``main + DupUp3D(source)``: one program owns ``ROWS`` rows of one output plane.
+
+        The output plane ``(batch, channel, frame)`` maps to one source plane per row
+        (``repeats >= FS`` makes the source channel independent of the column), so
+        each source row is loaded once as a contiguous span and interleaved with
+        itself ``FS`` times; the main tile is loaded and the result stored with
+        contiguous rows. Rows are contiguous in ``main``/``out``/``source`` but the
+        other strides are arbitrary (frame-major views included).
+        """
+        pid = tl.program_id(0).to(tl.int64)
+        row_block = pid % row_blocks
+        rest = pid // row_blocks
+        frame = rest % out_frames
+        rest = rest // out_frames
+        channel = rest % out_channels
+        batch = rest // out_channels
+
+        unsliced_frame = frame + frame_offset
+        source_frame = unsliced_frame // FT
+        temporal_remainder = unsliced_frame % FT
+        out_rows = row_block * ROWS + tl.arange(0, ROWS)
+        row_mask = out_rows < out_height
+        source_rows = out_rows // FS
+        height_remainder = out_rows % FS
+        source_channel = (((channel * FT + temporal_remainder) * FS + height_remainder) * FS) // REPEATS
+
+        main_rows = (
+            batch * stride_main_batch
+            + channel * stride_main_channel
+            + frame * stride_main_frame
+            + out_rows * stride_main_height
+        )
+        source_row_offsets = (
+            batch * stride_source_batch
+            + source_channel * stride_source_channel
+            + source_frame * stride_source_frame
+            + source_rows * stride_source_height
+        )
+        dtype = main_ptr.dtype.element_ty
+        if HAS_BIAS:
+            bias = tl.load(bias_ptr + channel).to(tl.float32)
+        for w0 in range(0, out_width, BLOCK_W):
+            cols = w0 + tl.arange(0, BLOCK_W)
+            mask = row_mask[:, None] & (cols < out_width)[None, :]
+            offsets = main_rows[:, None] + cols[None, :]
+            main = tl.load(main_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+            if HAS_BIAS:
+                # ATen's in-place conv bias ``add_``: fp32 opmath, rounded once.
+                main = (main + bias).to(dtype).to(tl.float32)
+            source_cols = w0 // FS + tl.arange(0, BLOCK_W // FS)
+            source_mask = row_mask[:, None] & (source_cols < source_width)[None, :]
+            source = tl.load(
+                source_ptr + (source_row_offsets[:, None] + source_cols[None, :]), mask=source_mask, other=0.0
+            )
+            if FS >= 2:
+                source = tl.interleave(source, source)
+            if FS >= 4:
+                source = tl.interleave(source, source)
+            value = main + source.to(tl.float32)
+            tl.store(out_ptr + offsets, value.to(dtype), mask=mask)
+
+    @triton.jit
+    def _dup_up3d_add_pixels_kernel(
+        main_ptr,
+        source_ptr,
+        bias_ptr,
+        out_ptr,
+        out_channels,
+        out_frames,
+        out_height,
+        out_width,
+        source_width,
+        frame_offset,
+        width_blocks,
+        stride_main_batch,
+        stride_main_frame,
+        stride_main_height,
+        stride_main_width,
+        stride_source_batch,
+        stride_source_frame,
+        stride_source_height,
+        stride_source_width,
+        FT: tl.constexpr,
+        FS: tl.constexpr,
+        REPEATS: tl.constexpr,
+        HAS_BIAS: tl.constexpr,
+        BLOCK_W: tl.constexpr,
+        BLOCK_C: tl.constexpr,
+    ):
+        """Channels-last ``main + DupUp3D(source)``: one program owns ``BLOCK_W`` pixels x all channels of one row.
+
+        Every output pixel reads one source pixel (``col // FS``) whose channel
+        vector is gathered with the per-channel map ``source_channel(channel)``
+        (a contiguous span when in/out channels match, stride 2 when the block
+        halves the channel count); main and out are contiguous along channels.
+        """
+        pid = tl.program_id(0).to(tl.int64)
+        width_block = pid % width_blocks
+        rest = pid // width_blocks
+        out_row = rest % out_height
+        rest = rest // out_height
+        frame = rest % out_frames
+        batch = rest // out_frames
+
+        unsliced_frame = frame + frame_offset
+        source_frame = unsliced_frame // FT
+        temporal_remainder = unsliced_frame % FT
+        source_row = out_row // FS
+        height_remainder = out_row % FS
+        cols = width_block * BLOCK_W + tl.arange(0, BLOCK_W)
+        col_mask = cols < out_width
+        source_cols = cols // FS
+
+        main_pixels = (
+            batch * stride_main_batch
+            + frame * stride_main_frame
+            + out_row * stride_main_height
+            + cols * stride_main_width
+        )
+        source_pixels = (
+            batch * stride_source_batch
+            + source_frame * stride_source_frame
+            + source_row * stride_source_height
+            + source_cols * stride_source_width
+        )
+        dtype = main_ptr.dtype.element_ty
+        for c0 in range(0, out_channels, BLOCK_C):
+            chans = c0 + tl.arange(0, BLOCK_C)
+            chan_mask = chans < out_channels
+            mask = col_mask[:, None] & chan_mask[None, :]
+            offsets = main_pixels[:, None] + chans[None, :]
+            main = tl.load(main_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+            if HAS_BIAS:
+                bias = tl.load(bias_ptr + chans, mask=chan_mask, other=0.0).to(tl.float32)
+                main = (main + bias[None, :]).to(dtype).to(tl.float32)
+            source_channels = (((chans * FT + temporal_remainder) * FS + height_remainder) * FS) // REPEATS
+            source = tl.load(source_ptr + (source_pixels[:, None] + source_channels[None, :]), mask=mask, other=0.0).to(
+                tl.float32
+            )
+            tl.store(out_ptr + offsets, (main + source).to(dtype), mask=mask)
+
+    @triton.jit
+    def _add_bias_residual_kernel(
+        x_ptr,
+        bias_x_ptr,
+        h_ptr,
+        bias_h_ptr,
+        out_ptr,
+        total,
+        channels,
+        spatial,
+        has_bias_x: tl.constexpr,
+        has_bias_h: tl.constexpr,
+        channels_inner: tl.constexpr,
+        index_64bit: tl.constexpr,
+        block_size_const: tl.constexpr,
+    ):
+        """``round(round(x + bias_x) + round(h + bias_h))`` for dense ``(N, C, *spatial)`` tensors.
+
+        Reproduces ATen's per-conv ``output.add_(bias)`` (fp32 opmath, one
+        rounding) followed by the residual ``x + h`` (fp32 opmath, one rounding)
+        without materializing the biased tensors.
+        """
+        if index_64bit:
+            offsets = tl.program_id(0).to(tl.int64) * block_size_const + tl.arange(0, block_size_const).to(tl.int64)
+        else:
+            offsets = tl.program_id(0) * block_size_const + tl.arange(0, block_size_const)
+        mask = offsets < total
+        if channels_inner:
+            channel = offsets % channels
+        else:
+            channel = (offsets // spatial) % channels
+        dtype = x_ptr.dtype.element_ty
+
+        x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        if has_bias_x:
+            bias_x = tl.load(bias_x_ptr + channel, mask=mask, other=0.0).to(tl.float32)
+            x = (x + bias_x).to(dtype).to(tl.float32)
+        h = tl.load(h_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        if has_bias_h:
+            bias_h = tl.load(bias_h_ptr + channel, mask=mask, other=0.0).to(tl.float32)
+            h = (h + bias_h).to(dtype).to(tl.float32)
+        tl.store(out_ptr + offsets, (x + h).to(dtype), mask=mask)
+
+
+def _uses_channels_last_3d(tensor: torch.Tensor) -> bool:
+    return tensor.shape[1] > 1 and tensor.stride(1) == 1
+
+
+def _next_power_of_2(value: int) -> int:
+    return 1 << max(0, value - 1).bit_length()
+
+
+def cat_pad_5d(
+    x: torch.Tensor,
+    cache_x: torch.Tensor | None,
+    padding: list[int] | tuple[int, ...],
+    keep_cache_frames: int = 0,
+    *,
+    channels_last_output: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None:
+    """Fuse temporal concat, constant padding, and optional cache refresh.
+
+    ``padding`` follows the Wan causal Conv3d convention
+    ``(w_left, w_right, h_top, h_bottom, t_front, t_back)``. The output uses
+    the same dense layout that the reference ``cat``/``pad`` path selects, or
+    ``channels_last_3d`` when ``channels_last_output`` is set (a fused transpose
+    for channels-first ``x``; the cache refresh still follows ``x``'s layout).
+    With ``keep_cache_frames > 0`` the last frames of the assembled input
+    (without spatial padding) are also written to a second tensor that
+    replaces the reference ``x[:, :, -CACHE_T:].clone()`` cache refresh.
+    ``cache_x`` may use either layout. Unsupported inputs return ``None`` so
+    callers can execute the reference implementation.
+    """
+    if not HAS_TRITON or len(padding) != 6:
+        return None
+    pad_width_left, pad_width_right, pad_height_top, pad_height_bottom, pad_front, pad_back = padding
+    if any(value < 0 for value in padding):
+        return None
+    if pad_width_left != pad_width_right or pad_height_top != pad_height_bottom or pad_back != 0:
+        return None
+    if x.dim() != 5 or not x.is_cuda or x.dtype not in _SUPPORTED_DTYPES:
+        return None
+    if keep_cache_frames < 0:
+        return None
+
+    channels_inner = _uses_channels_last_3d(x)
+    out_channels_inner = channels_inner or channels_last_output
+    cache_frames = 0
+    if cache_x is not None:
+        if (
+            cache_x.dim() != 5
+            or cache_x.dtype != x.dtype
+            or cache_x.device != x.device
+            or cache_x.shape[0] != x.shape[0]
+            or cache_x.shape[1] != x.shape[1]
+            or cache_x.shape[3:] != x.shape[3:]
+        ):
+            return None
+        cache_frames = cache_x.shape[2]
+
+    zero_front_frames = pad_front - cache_frames
+    if zero_front_frames < 0:
+        return None
+
+    batch, channels, frames, height, width = x.shape
+    out_frames = pad_front + frames
+    out_height = height + 2 * pad_height_top
+    out_width = width + 2 * pad_width_left
+    keep_frames = min(keep_cache_frames, out_frames)
+    out = torch.empty(
+        (batch, channels, out_frames, out_height, out_width),
+        device=x.device,
+        dtype=x.dtype,
+        memory_format=torch.channels_last_3d if out_channels_inner else torch.contiguous_format,
+    )
+    if out.numel() == 0:
+        return None
+
+    if keep_frames:
+        keep_arg = torch.empty(
+            (batch, channels, keep_frames, height, width),
+            device=x.device,
+            dtype=x.dtype,
+            memory_format=torch.channels_last_3d if channels_inner else torch.contiguous_format,
+        )
+    else:
+        keep_arg = out
+
+    if cache_x is None:
+        cache_arg = x
+        cache_strides = (0, 0, 0, 0, 0)
+    else:
+        cache_arg = cache_x
+        cache_strides = cache_x.stride()
+
+    common = (
+        channels,
+        frames,
+        height,
+        width,
+        cache_frames,
+        out_frames,
+        out_height,
+        out_width,
+        zero_front_frames,
+        pad_height_top,
+        pad_width_left,
+        out_frames - keep_frames,
+        keep_frames,
+    )
+    with torch.get_device_module().device(x.device):
+        if out_channels_inner:
+            if channels_inner:
+                block_c = min(_next_power_of_2(channels), 1024)
+                block_w = max(1, _TILE_ELEMENTS // block_c)
+            else:
+                # Transposing tile: 64 pixels x 64 channels reads 128-byte pixel runs
+                # per channel and writes 128-byte channel runs per pixel.
+                block_c = min(_next_power_of_2(channels), _TRANSPOSE_BLOCK)
+                block_w = _TRANSPOSE_BLOCK
+            width_blocks = triton.cdiv(out_width, block_w)
+            grid = (batch * out_frames * out_height * width_blocks,)
+            _cat_pad_pixels_kernel[grid](
+                x,
+                cache_arg,
+                out,
+                keep_arg,
+                *common,
+                width_blocks,
+                *x.stride(),
+                *cache_strides,
+                *keep_arg.stride(),
+                has_cache=cache_x is not None,
+                has_keep=keep_frames > 0,
+                BLOCK_W=block_w,
+                BLOCK_C=block_c,
+                num_warps=_NUM_WARPS,
+            )
+        else:
+            # 256-wide column chunks keep 84% of the lanes useful on the 642-wide
+            # 720p rows (a 1024-wide block would waste 37%); ROWS rows share the
+            # chunk so a program still moves a 4096-element tile per iteration.
+            block_w = min(_next_power_of_2(out_width), _ROWS_BLOCK_W)
+            rows_per_program = max(1, _TILE_ELEMENTS // block_w)
+            row_blocks = triton.cdiv(out_height, rows_per_program)
+            grid = (batch * channels * out_frames * row_blocks,)
+            _cat_pad_rows_kernel[grid](
+                x,
+                cache_arg,
+                out,
+                keep_arg,
+                *common,
+                row_blocks,
+                *x.stride(),
+                *cache_strides,
+                has_cache=cache_x is not None,
+                has_keep=keep_frames > 0,
+                ROWS=rows_per_program,
+                BLOCK_W=block_w,
+                num_warps=_ROWS_NUM_WARPS,
+            )
+
+    if keep_cache_frames:
+        return out, keep_arg
+    return out
+
+
+def _plane_layout(x: torch.Tensor) -> tuple[int, int, int] | None:
+    """``(channels_outer, plane_size, channel_stride)`` if ``x``'s frames are contiguous planes.
+
+    Channels-first (including the frame-major view ``WanResample`` emits): each
+    ``(batch, channel, frame)`` plane is ``H * W`` contiguous elements.
+    Channels-last: each ``(batch, frame)`` plane is ``H * W * C`` contiguous elements.
+    """
+    _, channels, _, height, width = x.shape
+    if x.stride(4) == 1 and x.stride(3) == width:
+        return channels, height * width, x.stride(1)
+    if channels > 1 and x.stride(1) == 1 and x.stride(4) == channels and x.stride(3) == width * channels:
+        return 1, height * width * channels, 0
+    return None
+
+
+def cat_time_5d(
+    x: torch.Tensor,
+    cache_x: torch.Tensor | None,
+    pad_front: int,
+    keep_cache_frames: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor] | torch.Tensor | None:
+    """Assemble ``[zeros(pad_front - Tc) | cache | x]`` along time, without spatial padding.
+
+    Returns the assembled tensor (same memory format as ``x``) and, with
+    ``keep_cache_frames > 0``, its last frames as the next feature-cache entry.
+    Because only whole aligned planes move, this runs at memory bandwidth in
+    both layouts; the spatial zero padding is left to the convolution
+    (``padding=(0, ph, pw)``), whose bitwise agreement with a pre-padded input
+    the caller verifies. Unsupported inputs return ``None``.
+    """
+    if not HAS_TRITON or x.dim() != 5 or not x.is_cuda or x.dtype not in _SUPPORTED_DTYPES:
+        return None
+    if pad_front < 0 or keep_cache_frames < 0:
+        return None
+    layout = _plane_layout(x)
+    if layout is None:
+        return None
+    channels_outer, plane_size, x_channel_stride = layout
+    channels_last = channels_outer == 1
+
+    cache_frames = 0
+    cache_strides = (0, 0, 0)
+    if cache_x is not None:
+        cache_layout = _plane_layout(cache_x) if cache_x.dim() == 5 else None
+        if (
+            cache_layout is None
+            or cache_layout[:2] != (channels_outer, plane_size)
+            or cache_x.dtype != x.dtype
+            or cache_x.device != x.device
+            or cache_x.shape[0] != x.shape[0]
+            or cache_x.shape[1] != x.shape[1]
+            or cache_x.shape[3:] != x.shape[3:]
+        ):
+            return None
+        cache_frames = cache_x.shape[2]
+        cache_strides = (cache_x.stride(0), cache_layout[2], cache_x.stride(2))
+    zero_front_frames = pad_front - cache_frames
+    if zero_front_frames < 0:
+        return None
+
+    batch, channels, frames, height, width = x.shape
+    out_frames = pad_front + frames
+    keep_frames = min(keep_cache_frames, out_frames)
+    memory_format = torch.channels_last_3d if channels_last else torch.contiguous_format
+    out = torch.empty(
+        (batch, channels, out_frames, height, width), device=x.device, dtype=x.dtype, memory_format=memory_format
+    )
+    if out.numel() == 0:
+        return None
+    if keep_frames:
+        keep_arg = torch.empty(
+            (batch, channels, keep_frames, height, width), device=x.device, dtype=x.dtype, memory_format=memory_format
+        )
+    else:
+        keep_arg = out
+
+    # Flatten the plane/block grid: large channels-last planes exceed CUDA's
+    # 65,535-block limit on grid-y.
+    blocks_per_plane = triton.cdiv(plane_size, _PLANE_BLOCK)
+    grid = (batch * channels_outer * out_frames * blocks_per_plane,)
+    with torch.get_device_module().device(x.device):
+        _cat_time_planes_kernel[grid](
+            x,
+            x if cache_x is None else cache_x,
+            out,
+            keep_arg,
+            plane_size,
+            blocks_per_plane,
+            channels_outer,
+            out_frames,
+            cache_frames,
+            zero_front_frames,
+            out_frames - keep_frames,
+            keep_frames,
+            x.stride(0),
+            x_channel_stride,
+            x.stride(2),
+            *cache_strides,
+            has_cache=cache_x is not None,
+            has_keep=keep_frames > 0,
+            BLOCK=_PLANE_BLOCK,
+            num_warps=_ROWS_NUM_WARPS,
+        )
+    if keep_cache_frames:
+        return out, keep_arg
+    return out
+
+
+def _rows_layout(x: torch.Tensor) -> bool:
+    """Rows are contiguous: channels-first tensors and frame-major views."""
+    return x.stride(4) == 1
+
+
+def _pixels_layout(x: torch.Tensor) -> bool:
+    """Channels are contiguous and pixels are whole channel vectors (channels_last_3d)."""
+    return x.shape[1] > 1 and x.stride(1) == 1 and x.stride(4) == x.shape[1]
+
+
+def dup_up3d_add(
+    main: torch.Tensor,
+    source: torch.Tensor,
+    factor_temporal: int,
+    factor_spatial: int,
+    repeats: int,
+    drop_first_frames: bool,
+    main_bias: torch.Tensor | None = None,
+) -> torch.Tensor | None:
+    """Evaluate ``(main + main_bias) + DupUp3D(source)`` without materializing the shortcut.
+
+    ``DupUp3D`` is ``repeat_interleave`` + an 8-D ``permute().contiguous()``
+    + an optional leading-frame slice; every output element is one input
+    element at a closed-form index, so the gather is fused with the residual
+    add. ``main_bias`` is the un-added per-channel bias of the convolution that
+    produced ``main`` (the up block's resample ``Conv2d``); it is applied with
+    ATen's rounding before the add. The output keeps ``main``'s strides, as
+    the reference add would. Declines (``None``) exotic factors, ``repeats <
+    factor_spatial`` (the source channel would vary along a row) and mixed
+    layouts.
+    """
+    if not HAS_TRITON or main.dim() != 5 or source.dim() != 5:
+        return None
+    if factor_temporal <= 0 or factor_spatial not in (1, 2, 4) or repeats <= 0:
+        return None
+    if factor_spatial > 1 and not _HAS_INTERLEAVE:
+        return None
+    if factor_temporal & (factor_temporal - 1) or repeats & (repeats - 1) or repeats < factor_spatial:
+        return None
+    if not main.is_cuda or not source.is_cuda:
+        return None
+    if main.dtype != source.dtype or main.dtype not in _SUPPORTED_DTYPES or main.device != source.device:
+        return None
+
+    batch, source_channels, frames, height, width = source.shape
+    frame_offset = factor_temporal - 1 if drop_first_frames else 0
+    numerator = source_channels * repeats
+    denominator = factor_temporal * factor_spatial * factor_spatial
+    if numerator % denominator:
+        return None
+    out_channels = numerator // denominator
+    expected_shape = (
+        batch,
+        out_channels,
+        frames * factor_temporal - frame_offset,
+        height * factor_spatial,
+        width * factor_spatial,
+    )
+    if tuple(main.shape) != expected_shape or main.numel() == 0:
+        return None
+    if main_bias is not None and (
+        main_bias.dim() != 1
+        or main_bias.numel() != out_channels
+        or main_bias.device != main.device
+        or main_bias.dtype not in _SUPPORTED_DTYPES
+    ):
+        return None
+
+    out = torch.empty_like(main)
+    if out.stride() != main.stride():
+        return None
+    bias_arg = main if main_bias is None else main_bias.contiguous()
+    _, _, out_frames, out_height, out_width = expected_shape
+    factors = dict(FT=factor_temporal, FS=factor_spatial, REPEATS=repeats, HAS_BIAS=main_bias is not None)
+
+    with torch.get_device_module().device(main.device):
+        if _rows_layout(main) and _rows_layout(source):
+            block_w = _pick_block_width(out_width)
+            rows_per_program = max(1, _TILE_ELEMENTS // block_w)
+            row_blocks = triton.cdiv(out_height, rows_per_program)
+            grid = (batch * out_channels * out_frames * row_blocks,)
+            _dup_up3d_add_rows_kernel[grid](
+                main,
+                source,
+                bias_arg,
+                out,
+                out_channels,
+                out_frames,
+                out_height,
+                out_width,
+                width,
+                frame_offset,
+                row_blocks,
+                *main.stride()[:4],
+                *source.stride()[:4],
+                **factors,
+                ROWS=rows_per_program,
+                BLOCK_W=block_w,
+                num_warps=_NUM_WARPS,
+            )
+            return out
+        if _pixels_layout(main) and _pixels_layout(source):
+            block_c = min(_next_power_of_2(out_channels), 1024)
+            block_w = max(1, _TILE_ELEMENTS // block_c)
+            width_blocks = triton.cdiv(out_width, block_w)
+            grid = (batch * out_frames * out_height * width_blocks,)
+            main_strides = (main.stride(0), main.stride(2), main.stride(3), main.stride(4))
+            source_strides = (source.stride(0), source.stride(2), source.stride(3), source.stride(4))
+            _dup_up3d_add_pixels_kernel[grid](
+                main,
+                source,
+                bias_arg,
+                out,
+                out_channels,
+                out_frames,
+                out_height,
+                out_width,
+                width,
+                frame_offset,
+                width_blocks,
+                *main_strides,
+                *source_strides,
+                **factors,
+                BLOCK_W=block_w,
+                BLOCK_C=block_c,
+                num_warps=_NUM_WARPS,
+            )
+            return out
+    return None
+
+
+def _bias_ok(bias: torch.Tensor | None, x: torch.Tensor) -> bool:
+    if bias is None:
+        return True
+    return (
+        bias.device == x.device
+        and bias.dtype in _SUPPORTED_DTYPES
+        and bias.dim() == 1
+        and bias.numel() == x.shape[1]
+        and bias.is_contiguous()
+    )
+
+
+def add_bias_residual(
+    x: torch.Tensor,
+    bias_x: torch.Tensor | None,
+    h: torch.Tensor,
+    bias_h: torch.Tensor | None,
+) -> torch.Tensor | None:
+    """``(x + bias_x) + (h + bias_h)`` with ATen's rounding points, or ``None`` to decline.
+
+    Folds the per-channel convolution bias adds (``output.add_(bias)``) of a
+    residual block's ``conv2`` and ``conv_shortcut`` into the residual add.
+    ``x`` and ``h`` must share shape, dtype and a dense layout (both contiguous,
+    or both channels-last); either bias may be ``None`` when that operand
+    already carries its bias.
+    """
+    if not HAS_TRITON or not x.is_cuda or x.dim() not in (4, 5) or x.dtype not in _SUPPORTED_DTYPES:
+        return None
+    if h.shape != x.shape or h.dtype != x.dtype or h.device != x.device:
+        return None
+    if x.is_contiguous() and h.is_contiguous():
+        channels_inner = False
+    else:
+        memory_format = torch.channels_last_3d if x.dim() == 5 else torch.channels_last
+        channels_last = x.is_contiguous(memory_format=memory_format) and h.is_contiguous(memory_format=memory_format)
+        if x.shape[1] < 2 or not channels_last:
+            return None
+        channels_inner = True
+    if not (_bias_ok(bias_x, x) and _bias_ok(bias_h, x)):
+        return None
+    total = x.numel()
+    if total == 0 or total > _MAX_INT32 * 4:
+        return None
+    channels = x.shape[1]
+    spatial = total // (x.shape[0] * channels)
+
+    out = torch.empty_like(x)
+    grid = (triton.cdiv(total, _BLOCK_SIZE),)
+    with torch.get_device_module().device(x.device):
+        _add_bias_residual_kernel[grid](
+            x,
+            x if bias_x is None else bias_x,
+            h,
+            x if bias_h is None else bias_h,
+            out,
+            total,
+            channels,
+            spatial,
+            has_bias_x=bias_x is not None,
+            has_bias_h=bias_h is not None,
+            channels_inner=channels_inner,
+            index_64bit=total >= _MAX_INT32,
+            block_size_const=_BLOCK_SIZE,
+            enable_fp_fusion=False,
+        )
+    return out
+
+
+__all__ = ["add_bias_residual", "cat_pad_5d", "cat_time_5d", "dup_up3d_add"]

vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_rms_norm.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa: N803
+"""Bit-exact fused RMSNorm epilogue for the diffusers Wan VAE decoder.
+
+``WanRMS_norm.forward`` evaluates, for fp16/bf16 activations::
+
+    normalized = F.normalize(x.float(), dim=1).to(x.dtype)
+    out = normalized * scale * gamma + 0.0
+
+That is an fp32 copy of the activation, an fp32 reduction, an fp32 divide, a
+narrowing cast and three low-precision elementwise kernels: six or seven full
+passes over an activation that reaches 472 MB in the Cosmos3 720p decoder.
+
+The fast path keeps the reduction in ATen (``torch.linalg.vector_norm`` with
+``dtype=torch.float32`` reads the low-precision tensor once and runs the same
+reduction kernel ``F.normalize(x.float())`` runs, so the fp32 denominators are
+identical) and folds everything after it into one Triton kernel that reads
+``x`` once and writes the result once. The kernel is bit-exact with ATen
+rather than merely close:
+
+* ``tl.math.div_rn`` is IEEE round-to-nearest fp32 division, the operation
+  ATen's ``div`` performs in fp32 opmath. Triton's ``/`` lowers to an
+  approximate reciprocal-multiply and is not usable here.
+* Every value ATen materializes in ``x.dtype`` is rounded here at the same
+  point with ``.to(dtype).to(tl.float32)``, including the final ``+ 0.0``
+  (which turns ``-0.0`` into ``+0.0`` exactly like upstream's ``+ self.bias``).
+* ``enable_fp_fusion=False`` stops LLVM from contracting a multiply and its
+  neighbour into an FMA, which would skip one of those roundings.
+
+The optional SiLU epilogue reproduces ATen's ``x / (1 + expf(-x))`` with the
+libdevice ``expf`` and an IEEE divide. Whether that is bit-identical depends
+on the toolkit, so :func:`silu_epilogue_is_exact` verifies it exhaustively
+over all 65,536 bf16/fp16 values before callers enable it.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from vllm.triton_utils import HAS_TRITON, tl, tldevice, triton
+
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+# These kernels are HBM bound; a wide block gives every thread more independent
+# loads to overlap. 2048/4 measured ~6% faster than 1024/4 on the decoder's
+# real shapes and is bit-identical (the kernel is elementwise).
+_BLOCK = 2048
+_NUM_WARPS = 4
+_MAX_GRID_Y = 65535
+
+if HAS_TRITON:
+
+    @triton.jit
+    def _rms_norm_scale_kernel(
+        X,
+        DENOM,
+        GAMMA,
+        OUT,
+        spatial,
+        channels,
+        scale,
+        SILU: tl.constexpr,
+        BLOCK: tl.constexpr,
+    ):
+        """``out = ((x / denom) * scale * gamma + 0.0)`` [then SiLU] with ATen's rounding points.
+
+        ``X``/``OUT`` are dense ``(rows, spatial)`` with ``rows = batch * channels``;
+        ``DENOM`` is fp32 ``(batch, spatial)``; ``GAMMA`` is ``(channels,)``.
+        """
+        row = tl.program_id(0)
+        col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+        mask = col < spatial
+
+        channel = row % channels
+        batch = row // channels
+        base = row.to(tl.int64) * spatial + col
+        denom_base = batch.to(tl.int64) * spatial + col
+
+        dtype = X.dtype.element_ty
+        x = tl.load(X + base, mask=mask, other=0).to(tl.float32)
+        d = tl.load(DENOM + denom_base, mask=mask, other=1.0)
+        g = tl.load(GAMMA + channel).to(tl.float32)
+
+        v = tl.math.div_rn(x, d).to(dtype).to(tl.float32)
+        v = (v * scale).to(dtype).to(tl.float32)
+        v = (v * g).to(dtype).to(tl.float32)
+        v = (v + 0.0).to(dtype)
+        if SILU:
+            v32 = v.to(tl.float32)
+            v = tl.math.div_rn(v32, 1.0 + tldevice.exp(-v32)).to(dtype)
+        tl.store(OUT + base, v, mask=mask)
+
+
+def rms_norm_scale(
+    x: torch.Tensor,
+    denom: torch.Tensor,
+    gamma: torch.Tensor,
+    scale: float,
+    *,
+    silu: bool = False,
+) -> torch.Tensor | None:
+    """``(x / denom) * scale * gamma + 0.0`` (optionally followed by SiLU), or ``None`` to decline.
+
+    ``x`` must be a contiguous ``(N, C, S)`` tensor, ``denom`` a contiguous fp32
+    ``(N, 1, S)`` tensor and ``gamma`` a contiguous ``(C,)`` tensor of ``x``'s dtype.
+    """
+    if not HAS_TRITON or not x.is_cuda:
+        return None
+    if x.dim() != 3 or x.dtype not in _SUPPORTED_DTYPES:
+        return None
+    if denom.dtype is not torch.float32 or gamma.dtype is not x.dtype:
+        return None
+    # Raw pointers go to one kernel, which does no cross-device checking.
+    if denom.device != x.device or gamma.device != x.device:
+        return None
+    if not (x.is_contiguous() and denom.is_contiguous() and gamma.is_contiguous()):
+        return None
+    batch, channels, spatial = x.shape
+    if denom.shape != (batch, 1, spatial) or gamma.shape != (channels,):
+        return None
+    if spatial == 0 or batch * channels == 0:
+        return None
+    spatial_blocks = triton.cdiv(spatial, _BLOCK)
+    if spatial_blocks > _MAX_GRID_Y:
+        return None
+
+    out = torch.empty_like(x)
+    grid = (batch * channels, spatial_blocks)
+    with torch.get_device_module().device(x.device):
+        _rms_norm_scale_kernel[grid](
+            x,
+            denom,
+            gamma,
+            out,
+            spatial,
+            channels,
+            float(scale),
+            SILU=silu,
+            BLOCK=_BLOCK,
+            num_warps=_NUM_WARPS,
+            # FMA contraction would swallow one of the low-precision roundings.
+            enable_fp_fusion=False,
+        )
+    return out
+
+
+def _bitwise_equal_allow_nan(a: torch.Tensor, b: torch.Tensor) -> bool:
+    if a.shape != b.shape or a.dtype != b.dtype:
+        return False
+    int_dtype = torch.int16 if a.element_size() == 2 else torch.int32
+    same_bits = a.view(int_dtype) == b.view(int_dtype)
+    both_nan = torch.isnan(a) & torch.isnan(b)
+    return bool((same_bits | both_nan).all().item())
+
+
+_SILU_EXACT_CACHE: dict[tuple[str, torch.dtype], bool] = {}
+
+
+def silu_epilogue_is_exact(device: torch.device, dtype: torch.dtype) -> bool:
+    """Exhaustively verify the fused SiLU epilogue against ``F.silu`` for one dtype.
+
+    bf16 and fp16 have 65,536 values each, so the check is complete and costs
+    about a millisecond. The result is cached per device. fp32 cannot be
+    checked exhaustively and always returns ``False``.
+    """
+    if dtype not in (torch.bfloat16, torch.float16):
+        return False
+    key = (str(device), dtype)
+    cached = _SILU_EXACT_CACHE.get(key)
+    if cached is not None:
+        return cached
+    result = False
+    if HAS_TRITON and device.type == "cuda":
+        try:
+            bits = torch.arange(0, 65536, dtype=torch.int32, device=device).to(torch.int16)
+            x = bits.view(dtype).reshape(1, 1, -1)
+            denom = torch.ones((1, 1, x.shape[-1]), dtype=torch.float32, device=device)
+            gamma = torch.ones((1,), dtype=dtype, device=device)
+            out = rms_norm_scale(x, denom, gamma, 1.0, silu=True)
+            if out is not None:
+                # The reference epilogue also ends with ``+ 0.0`` before SiLU.
+                reference = F.silu(x + 0.0)
+                result = _bitwise_equal_allow_nan(out, reference)
+        except Exception:  # pragma: no cover - depends on the installed toolkit
+            result = False
+    _SILU_EXACT_CACHE[key] = result
+    return result
+
+
+__all__ = ["rms_norm_scale", "silu_epilogue_is_exact"]

vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_rms_norm_cl.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa: N803
+"""Single-pass channels-last RMSNorm(+SiLU) for the Wan VAE decoder (Tier 2, not bit-exact).
+
+With channels-last activations every spatial position owns a contiguous vector
+of ``C`` channels, so one program can load a ``(BLOCK_P, C)`` tile, reduce over
+the channels in registers and write the normalized (and optionally SiLU'd)
+result: the activation is read once and written once instead of the two-pass
+scheme :mod:`.triton_rms_norm` needs for channels-first layouts.
+
+This kernel is only dispatched at the tolerance-based ``channels_last`` level,
+so it uses fast math: one reciprocal per position, ``scale`` folded into
+``gamma``, the approximate ``exp`` and divide for SiLU, and a single rounding
+to the activation dtype at the end. The bit-exact epilogue sequence (IEEE
+``div_rn`` twice, libdevice ``exp``, six intermediate roundings) measured
+ALU-bound at 1.6 TB/s on the 512-channel stage.
+
+The kernel can also absorb the bias of the convolution that produced ``x``:
+``WanResidualBlock.conv1`` feeds only ``norm2``, so at this level the
+convolution runs without its bias and the norm adds it per channel in fp32
+before reducing, which removes ATen's separate strided ``add_`` pass (about
+0.6 ms per 720p chunk per block). The sum is rounded to the activation dtype
+exactly as ATen's ``add_`` does, so folding the bias changes no numerics.
+"""
+
+from __future__ import annotations
+
+import torch
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+_MAX_CHANNELS = 1024
+# Positions x channels per program. Autotuning on the first call picked one
+# position per program (512 bytes in flight per CTA, 0.8 TB/s on the 256-channel
+# stage); a fixed ~4096-element tile keeps enough bytes in flight to approach
+# HBM bandwidth, as measured for the data-movement kernels.
+_TILE_ELEMENTS = 4096
+_NUM_WARPS = 4
+
+if HAS_TRITON:
+
+    @triton.jit
+    def _rms_norm_cl_kernel(
+        X,
+        GAMMA,
+        BIAS,
+        OUT,
+        rows,
+        scale,
+        eps,
+        C: tl.constexpr,
+        BLOCK_C: tl.constexpr,
+        SILU: tl.constexpr,
+        HAS_BIAS: tl.constexpr,
+        BLOCK_P: tl.constexpr,
+    ):
+        row = tl.program_id(0).to(tl.int64) * BLOCK_P + tl.arange(0, BLOCK_P)
+        col = tl.arange(0, BLOCK_C)
+        row_mask = row < rows
+        col_mask = col < C
+        mask = row_mask[:, None] & col_mask[None, :]
+        offsets = row[:, None] * C + col[None, :]
+
+        dtype = X.dtype.element_ty
+        x = tl.load(X + offsets, mask=mask, other=0.0).to(tl.float32)
+        if HAS_BIAS:
+            # ATen's in-place conv bias ``add_``: fp32 opmath, rounded once.
+            bias = tl.load(BIAS + col, mask=col_mask, other=0.0).to(tl.float32)
+            x = (x + bias[None, :]).to(dtype).to(tl.float32)
+        inv_norm = 1.0 / tl.maximum(tl.sqrt(tl.sum(x * x, axis=1)), eps)
+        gamma = tl.load(GAMMA + col, mask=col_mask, other=0.0).to(tl.float32) * scale
+        v = x * inv_norm[:, None] * gamma[None, :]
+        if SILU:
+            v = v / (1.0 + tl.exp(-v))
+        tl.store(OUT + offsets, v.to(dtype), mask=mask)
+
+
+def _rows_view(x: torch.Tensor) -> torch.Tensor | None:
+    """``(positions, C)`` view of a channels-last tensor, or ``None`` if it is not one."""
+    channels = x.shape[1]
+    if x.dim() == 5:
+        if not x.is_contiguous(memory_format=torch.channels_last_3d):
+            return None
+        rows = x.permute(0, 2, 3, 4, 1).reshape(-1, channels)
+    elif x.dim() == 4:
+        if not x.is_contiguous(memory_format=torch.channels_last):
+            return None
+        rows = x.permute(0, 2, 3, 1).reshape(-1, channels)
+    else:
+        return None
+    if rows.data_ptr() != x.data_ptr() or not rows.is_contiguous():
+        return None
+    return rows
+
+
+def rms_norm_channels_last(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    scale: float,
+    *,
+    silu: bool = False,
+    bias: torch.Tensor | None = None,
+    eps: float = 1e-12,
+) -> torch.Tensor | None:
+    """``F.normalize(x + bias, dim=1) * scale * gamma`` [then SiLU] for channels-last ``x`` (fast math), or ``None``.
+
+    ``bias`` is an optional per-channel vector (the un-added bias of the
+    convolution that produced ``x``); ``None`` skips the add.
+    """
+    if not HAS_TRITON or not x.is_cuda or x.dtype not in _SUPPORTED_DTYPES:
+        return None
+    channels = x.shape[1]
+    if channels < 2 or channels > _MAX_CHANNELS or x.numel() == 0:
+        return None
+    if gamma.numel() != channels or gamma.dtype is not x.dtype or gamma.device != x.device:
+        return None
+    if bias is not None and (
+        bias.dim() != 1 or bias.numel() != channels or bias.device != x.device or bias.dtype not in _SUPPORTED_DTYPES
+    ):
+        return None
+    x_rows = _rows_view(x)
+    if x_rows is None:
+        return None
+    out = torch.empty_like(x)
+    out_rows = _rows_view(out)
+    if out_rows is None:
+        return None
+    rows = x_rows.shape[0]
+    block_c = triton.next_power_of_2(channels)
+    block_p = max(1, _TILE_ELEMENTS // block_c)
+    grid = (triton.cdiv(rows, block_p),)
+
+    with torch.get_device_module().device(x.device):
+        _rms_norm_cl_kernel[grid](
+            x_rows,
+            gamma.reshape(-1).contiguous(),
+            x_rows if bias is None else bias.contiguous(),
+            out_rows,
+            rows,
+            float(scale),
+            float(eps),
+            C=channels,
+            BLOCK_C=block_c,
+            SILU=silu,
+            HAS_BIAS=bias is not None,
+            BLOCK_P=block_p,
+            num_warps=_NUM_WARPS,
+            enable_fp_fusion=False,
+        )
+    return out
+
+
+__all__ = ["rms_norm_channels_last"]

vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_upsample.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# ruff: noqa: N803
+"""Bit-exact nearest-neighbour 2x spatial upsampling for the Wan VAE decoder.
+
+``WanResample`` doubles the height and width of every frame with
+``nn.Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact")`` before its
+``Conv2d``. With an integer factor of two both ``nearest`` and ``nearest-exact``
+map output index ``o`` to input index ``o // 2``, so the operation is a pure
+gather: every input element lands in a 2x2 block of the output and no
+arithmetic happens. ATen reads channels-first input through a contiguous copy
+(the decoder hands it the strided frame-major view ``WanResample`` builds) and
+runs its channels-last kernel at about a third of memory bandwidth on GB200; the
+kernels here read the input once, in place, and write every output element once.
+
+Channels-first: one program owns ``ROWS`` input rows, loads each as a
+contiguous span, interleaves it with itself (``[a, a, b, b, ...]``) and stores
+the result to the two identical output rows. Channels-last: one program owns
+``BLOCK_W`` pixels x ``BLOCK_C`` channels of one input row and stores each
+pixel's channel vector to its four output pixels as contiguous spans.
+
+The wrapper returns ``None`` for anything it does not handle (rank, dtype,
+device, layout) so the caller runs ``nn.Upsample.forward``.
+"""
+
+from __future__ import annotations
+
+import torch
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+from ._utils import _pick_block_width
+
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+# Input elements loaded per program iteration; each is stored four times, so a
+# 2048-element load tile keeps 8192 stores in flight per program.
+_TILE_ELEMENTS = 2048
+_NUM_WARPS = 4
+_MAX_BLOCK_C = 1024
+_HAS_INTERLEAVE = HAS_TRITON and hasattr(tl, "interleave")
+
+if HAS_TRITON:
+
+    @triton.jit
+    def _upsample2x_rows_kernel(
+        x_ptr,
+        out_ptr,
+        rows,
+        channels,
+        height,
+        width,
+        out_width,
+        stride_x_batch,
+        stride_x_channel,
+        stride_x_height,
+        ROWS: tl.constexpr,
+        BLOCK_W: tl.constexpr,
+    ):
+        """Channels-first: ``ROWS`` input rows -> two identical, width-doubled output rows each.
+
+        Input rows are contiguous spans (``stride(3) == 1``, ``stride(2) == width``)
+        but batch and channel strides are arbitrary, so the frame-major view the
+        decoder produces is read without a copy. The output is contiguous, so
+        flattened input row ``r`` maps to output rows ``2r`` and ``2r + 1``.
+        """
+        in_rows = tl.program_id(0).to(tl.int64) * ROWS + tl.arange(0, ROWS)
+        row_mask = in_rows < rows
+        in_height = in_rows % height
+        rest = in_rows // height
+        channel = rest % channels
+        batch = rest // channels
+        x_row_offsets = batch * stride_x_batch + channel * stride_x_channel + in_height * stride_x_height
+        out_row_offsets = (in_rows * 2) * out_width
+        for w0 in range(0, width, BLOCK_W):
+            cols = w0 + tl.arange(0, BLOCK_W)
+            mask = row_mask[:, None] & (cols < width)[None, :]
+            values = tl.load(x_ptr + (x_row_offsets[:, None] + cols[None, :]), mask=mask, other=0)
+            doubled = tl.interleave(values, values)
+            out_cols = 2 * w0 + tl.arange(0, 2 * BLOCK_W)
+            out_mask = row_mask[:, None] & (out_cols < out_width)[None, :]
+            out_offsets = out_row_offsets[:, None] + out_cols[None, :]
+            tl.store(out_ptr + out_offsets, doubled, mask=out_mask)
+            tl.store(out_ptr + (out_offsets + out_width), doubled, mask=out_mask)
+
+    @triton.jit
+    def _upsample2x_pixels_kernel(
+        x_ptr,
+        out_ptr,
+        channels,
+        height,
+        width,
+        width_blocks,
+        stride_x_batch,
+        stride_x_height,
+        stride_x_width,
+        BLOCK_W: tl.constexpr,
+        BLOCK_C: tl.constexpr,
+    ):
+        """Channels-last: ``BLOCK_W`` pixels of one input row -> their 2x2 output pixels.
+
+        Channels are the contiguous dimension, so every pixel's channel vector is
+        one contiguous load and four contiguous stores.
+        """
+        pid = tl.program_id(0).to(tl.int64)
+        width_block = pid % width_blocks
+        rest = pid // width_blocks
+        in_height = rest % height
+        batch = rest // height
+        cols = width_block * BLOCK_W + tl.arange(0, BLOCK_W)
+        col_mask = cols < width
+        x_row = batch * stride_x_batch + in_height * stride_x_height
+        out_row_stride = (2 * width) * channels
+        out_row = (batch * (2 * height) + 2 * in_height) * out_row_stride
+        out_pixels = out_row + (2 * cols) * channels
+        for c0 in range(0, channels, BLOCK_C):
+            chans = c0 + tl.arange(0, BLOCK_C)
+            mask = col_mask[:, None] & (chans < channels)[None, :]
+            values = tl.load(x_ptr + (x_row + cols[:, None] * stride_x_width + chans[None, :]), mask=mask, other=0)
+            first = out_pixels[:, None] + chans[None, :]
+            tl.store(out_ptr + first, values, mask=mask)
+            tl.store(out_ptr + (first + channels), values, mask=mask)
+            tl.store(out_ptr + (first + out_row_stride), values, mask=mask)
+            tl.store(out_ptr + (first + out_row_stride + channels), values, mask=mask)
+
+
+def upsample_nearest_2x(x: torch.Tensor) -> torch.Tensor | None:
+    """``F.interpolate(x, scale_factor=(2.0, 2.0), mode="nearest-exact")`` for 4D ``x``, or ``None``.
+
+    Identical to ``mode="nearest"`` as well (both select input index ``o // 2``
+    for a factor of two). The output takes the memory format ATen would choose:
+    contiguous for channels-first input (including the strided frame-major view
+    ``WanResample`` produces), channels-last for channels-last input.
+    """
+    if not _HAS_INTERLEAVE or x.dim() != 4 or not x.is_cuda or x.dtype not in _SUPPORTED_DTYPES:
+        return None
+    batch, channels, height, width = x.shape
+    if x.numel() == 0:
+        return None
+    out_shape = (batch, channels, 2 * height, 2 * width)
+
+    if x.stride(3) == 1 and x.stride(2) == width:
+        out = torch.empty(out_shape, device=x.device, dtype=x.dtype)
+        rows = batch * channels * height
+        block_w = _pick_block_width(width)
+        rows_per_program = max(1, _TILE_ELEMENTS // block_w)
+        grid = (triton.cdiv(rows, rows_per_program),)
+        with torch.get_device_module().device(x.device):
+            _upsample2x_rows_kernel[grid](
+                x,
+                out,
+                rows,
+                channels,
+                height,
+                width,
+                2 * width,
+                x.stride(0),
+                x.stride(1),
+                x.stride(2),
+                ROWS=rows_per_program,
+                BLOCK_W=block_w,
+                num_warps=_NUM_WARPS,
+            )
+        return out
+
+    if channels > 1 and x.is_contiguous(memory_format=torch.channels_last):
+        out = torch.empty(out_shape, device=x.device, dtype=x.dtype, memory_format=torch.channels_last)
+        block_c = min(triton.next_power_of_2(channels), _MAX_BLOCK_C)
+        block_w = max(1, _TILE_ELEMENTS // block_c)
+        width_blocks = triton.cdiv(width, block_w)
+        grid = (batch * height * width_blocks,)
+        with torch.get_device_module().device(x.device):
+            _upsample2x_pixels_kernel[grid](
+                x,
+                out,
+                channels,
+                height,
+                width,
+                width_blocks,
+                x.stride(0),
+                x.stride(2),
+                x.stride(3),
+                BLOCK_W=block_w,
+                BLOCK_C=block_c,
+                num_warps=_NUM_WARPS,
+            )
+        return out
+
+    return None
+
+
+__all__ = ["upsample_nearest_2x"]

vllm_omni/diffusion/registry.py

@@ -487,11 +487,58 @@ def initialize_model(
         # at model loading time, not inside individual model files
         _apply_sequence_parallel_if_enabled(model, od_config)
 
+        # Install the Wan VAE decoder fast path once the VAE parallel mode is
+        # known. This runs inside the weight-loading memory pool, so any weight
+        # re-layout it performs stays offloadable by sleep mode.
+        _apply_wan_vae_fastpath_if_enabled(model, od_config)
+
         return model
     else:
         raise ValueError(f"Model class {od_config.model_class_name} not found in diffusion model registry.")
 
 
+def _apply_wan_vae_fastpath_if_enabled(model, od_config: OmniDiffusionConfig) -> None:
+    """Install the Wan VAE decoder fast path on every diffusers Wan VAE of the pipeline.
+
+    Controlled by ``od_config.vae_fast_path`` (``off`` / ``lossless`` /
+    ``channels_last``). Only CUDA is supported; other platforms keep the
+    reference diffusers decoder. Failures never abort model loading.
+    """
+    level = getattr(od_config, "vae_fast_path", "lossless")
+    if level == "off":
+        return
+    if not current_omni_platform.is_cuda():
+        logger.debug("Wan VAE fast path is only supported on CUDA; skipping")
+        return
+
+    from diffusers.models.autoencoders import AutoencoderKLWan
+
+    from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import install_wan_vae_fastpath
+    from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+
+    candidates: list[tuple[str, nn.Module]] = []
+    try:
+        discovered = ModuleDiscovery.discover(model)
+        candidates = list(zip(discovered.vae_names, discovered.vaes, strict=True))
+    except Exception:
+        logger.debug("Wan VAE fast path: component discovery failed; falling back to `vae`", exc_info=True)
+    if not candidates and isinstance(getattr(model, "vae", None), nn.Module):
+        candidates = [("vae", model.vae)]
+
+    for name, vae in candidates:
+        if not isinstance(vae, AutoencoderKLWan):
+            continue
+        try:
+            report = install_wan_vae_fastpath(vae, level=level)
+        except Exception:
+            logger.warning(
+                "Failed to install the Wan VAE fast path on %s; using the reference decoder", name, exc_info=True
+            )
+            continue
+        if report.installed:
+            logger.info("Wan VAE fast path (%s) active on %s", level, name)
+
+
 def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -> None:
     """Apply sequence parallelism hooks if SP is enabled.
 

vllm_omni/engine/arg_utils.py

@@ -587,6 +587,7 @@ class OrchestratorArgs:
     step_execution: bool = False
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
+    vae_fast_path: str = "lossless"
     enable_multithread_weight_load: bool = True
     enable_broadcast_weight_load: bool = False
     num_weight_load_threads: int = 4

vllm_omni/entrypoints/cli/serve.py

@@ -824,6 +824,17 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).",
         )
+        omni_config_group.add_argument(
+            "--vae-fast-path",
+            choices=("off", "lossless", "channels_last"),
+            default="lossless",
+            help=(
+                "Wan VAE decoder fast path. 'lossless' (default) installs bit-exact fused kernels; "
+                "'channels_last' additionally switches decoder convolutions to channels-last memory "
+                "format and fuses RMSNorm+SiLU (faster, not bit-exact); 'off' keeps the reference "
+                "diffusers implementation."
+            ),
+        )
 
         # Parallel weight loading (faster diffusion startup)
         omni_config_group.add_argument(