From c482626450e20f35e5b59df52f51a766829bfe8a Mon Sep 17 00:00:00 2001 From: MaciejBalaNV Date: Tue, 22 Sep 2026 20:37:49 +0200 Subject: [PATCH] [optimization] Added Wan 2.2 VAE decode optimizations (#7056) Signed-off-by: Maciej Bala Signed-off-by: MaciejBalaNV --- benchmarks/diffusion/bench_wan_vae_decode.py | 324 ++++++ docs/.nav.yml | 1 + docs/features/README.md | 1 + docs/user_guide/diffusion/vae_fast_path.md | 102 ++ docs/user_guide/diffusion_features.md | 1 + pyproject.toml | 1 + recipes/cosmos3/Cosmos3-Nano.md | 4 + recipes/cosmos3/Cosmos3-Super.md | 4 + .../test_wan_vae_fastpath_install.py | 619 ++++++++++ .../test_wan_vae_fastpath_kernels.py | 903 +++++++++++++++ .../test_diffusion_config_propagation.py | 11 + tests/entrypoints/test_serve.py | 17 + vllm_omni/config/omni_config.py | 1 + vllm_omni/config/stage_config.py | 1 + vllm_omni/diffusion/data.py | 11 + .../autoencoders/autoencoder_kl_wan.py | 10 + .../autoencoders/wan_vae_fastpath/__init__.py | 36 + .../autoencoders/wan_vae_fastpath/_utils.py | 17 + .../autoencoders/wan_vae_fastpath/decode.py | 62 + .../autoencoders/wan_vae_fastpath/forwards.py | 735 ++++++++++++ .../autoencoders/wan_vae_fastpath/install.py | 286 +++++ .../wan_vae_fastpath/triton_data_movement.py | 1000 +++++++++++++++++ .../wan_vae_fastpath/triton_rms_norm.py | 190 ++++ .../wan_vae_fastpath/triton_rms_norm_cl.py | 154 +++ .../wan_vae_fastpath/triton_upsample.py | 190 ++++ vllm_omni/diffusion/registry.py | 47 + vllm_omni/engine/arg_utils.py | 1 + vllm_omni/entrypoints/cli/serve.py | 11 + 28 files changed, 4740 insertions(+) create mode 100644 benchmarks/diffusion/bench_wan_vae_decode.py create mode 100644 docs/user_guide/diffusion/vae_fast_path.md create mode 100644 tests/diffusion/distributed/test_wan_vae_fastpath_install.py create mode 100644 tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/__init__.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/_utils.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/decode.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/forwards.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/install.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_data_movement.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_rms_norm.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_rms_norm_cl.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_upsample.py diff --git a/benchmarks/diffusion/bench_wan_vae_decode.py b/benchmarks/diffusion/bench_wan_vae_decode.py new file mode 100644 index 000000000..fec4f8029 --- /dev/null +++ b/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() diff --git a/docs/.nav.yml b/docs/.nav.yml index ec428d76b..123de61e9 100644 --- a/docs/.nav.yml +++ b/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 diff --git a/docs/features/README.md b/docs/features/README.md index 7c3e89676..5d624a34b 100644 --- a/docs/features/README.md +++ b/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 | diff --git a/docs/user_guide/diffusion/vae_fast_path.md b/docs/user_guide/diffusion/vae_fast_path.md new file mode 100644 index 000000000..9830b9449 --- /dev/null +++ b/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 +``` diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index b27e80ccc..8ab33541d 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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) | diff --git a/pyproject.toml b/pyproject.toml index c463d3917..3f7c0b243 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -331,6 +331,7 @@ extend-ignore-identifiers-re = [ ".*arange", ".*MoBA", ".*temperal_downsample", + ".*temperal_upsample", ".*_nd", ".*nd$", ".*EmbedND$", diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index 29088e94c..856d05b5a 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/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 diff --git a/recipes/cosmos3/Cosmos3-Super.md b/recipes/cosmos3/Cosmos3-Super.md index 82ae1f20a..1e5dbf7f0 100644 --- a/recipes/cosmos3/Cosmos3-Super.md +++ b/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). diff --git a/tests/diffusion/distributed/test_wan_vae_fastpath_install.py b/tests/diffusion/distributed/test_wan_vae_fastpath_install.py new file mode 100644 index 000000000..abf36f092 --- /dev/null +++ b/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) diff --git a/tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py b/tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py new file mode 100644 index 000000000..424ac4a6c --- /dev/null +++ b/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}" diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index 7421dfb79..9a9406f98 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/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", diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index ad5e33dc1..48382d98b 100644 --- a/tests/entrypoints/test_serve.py +++ b/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.""" diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index d64372b84..e5ac4505d 100644 --- a/vllm_omni/config/omni_config.py +++ b/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 diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 164f7cd6b..584f6af26 100644 --- a/vllm_omni/config/stage_config.py +++ b/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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 9dd69502f..8f26ebe36 100644 --- a/vllm_omni/diffusion/data.py +++ b/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', " diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index 9ab28d09c..700bfcc79 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/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: diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/__init__.py b/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/__init__.py new file mode 100644 index 000000000..33dd5b4f1 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/_utils.py b/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/_utils.py new file mode 100644 index 000000000..53277b143 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/decode.py b/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/decode.py new file mode 100644 index 000000000..e4ea91773 --- /dev/null +++ b/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"] diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/forwards.py b/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/forwards.py new file mode 100644 index 000000000..d3d0bb9b6 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/install.py b/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/install.py new file mode 100644 index 000000000..1ea41b65f --- /dev/null +++ b/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", +] diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_data_movement.py b/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_data_movement.py new file mode 100644 index 000000000..ee2d7a63f --- /dev/null +++ b/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"] diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_rms_norm.py b/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_rms_norm.py new file mode 100644 index 000000000..c12e74b63 --- /dev/null +++ b/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"] diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_rms_norm_cl.py b/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_rms_norm_cl.py new file mode 100644 index 000000000..f60380909 --- /dev/null +++ b/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"] diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_upsample.py b/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/triton_upsample.py new file mode 100644 index 000000000..5013a970e --- /dev/null +++ b/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"] diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 17d8216db..a9cb0e411 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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. diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 8fb6c80e9..4d44354d0 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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 diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 57514bbe9..c26ba9375 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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(