From 099f9b553d63b3bbcd1120eee6b73288a4088e97 Mon Sep 17 00:00:00 2001 From: QianCyrus <101633534+QianCyrus@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:44:34 +0900 Subject: [PATCH] [Model] Optimize Wan VAE first-frame decoding with Conv2d (#8179) Signed-off-by: QianCyrus <101633534+QianCyrus@users.noreply.github.com> Signed-off-by: Hongsheng Liu --- benchmarks/diffusion/bench_wan_vae_decode.py | 106 +++++++++++++- docs/user_guide/diffusion/vae_fast_path.md | 37 ++++- .../test_wan_vae_fastpath_kernels.py | 130 ++++++++++++++++++ .../autoencoders/wan_vae_fastpath/forwards.py | 54 +++++++- 4 files changed, 314 insertions(+), 13 deletions(-) diff --git a/benchmarks/diffusion/bench_wan_vae_decode.py b/benchmarks/diffusion/bench_wan_vae_decode.py index fec4f8029..1af0af7af 100644 --- a/benchmarks/diffusion/bench_wan_vae_decode.py +++ b/benchmarks/diffusion/bench_wan_vae_decode.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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 @@ -30,10 +30,13 @@ return the full video on rank 0 only, so equality and PSNR are computed there. from __future__ import annotations import argparse +import json import math import os +import statistics import time from collections.abc import Iterable +from unittest.mock import patch import torch import torch.distributed as dist @@ -82,6 +85,10 @@ def parse_args() -> argparse.Namespace: 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( + "--first-frame-ablation", action="store_true", help="Compare channels_last with/without first-frame Conv2d" + ) + parser.add_argument("--latents", help="For ablation: saved VAE-input tensor (.pt), after latent mean/std scaling") parser.add_argument( "--vae-patch-parallel-size", type=int, @@ -254,8 +261,99 @@ def run_level( return output.detach(), stats +@torch.inference_mode() +def run_first_frame_ablation(args: argparse.Namespace, dtype: torch.dtype, device: torch.device) -> None: + """Same-model ABBA comparison; disable only the first-frame helper in the reference arm.""" + from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards + + if args.iters < 2 or args.warmup < 0: + raise ValueError("Ablation requires --iters >= 2 and --warmup >= 0") + vae = load_vae(args, dtype, device, parallel=False) + report = install_wan_vae_fastpath(vae, level="channels_last") + if not report.installed: + raise RuntimeError(f"channels_last installation failed: {report.reason}") + latents = make_latents(vae, args, dtype, device) + if args.latents: + saved = torch.load(args.latents, map_location="cpu", weights_only=True) + saved = saved[:, :, : latents.shape[2]] + if saved.shape != latents.shape: + raise ValueError(f"Saved latent shape {saved.shape} does not match {latents.shape}") + latents = saved.to(device=device, dtype=dtype).contiguous() + helpers = {"baseline": lambda conv, x, bias: None, "candidate": forwards._first_frame_conv2d} + optimized_calls = 0 + + def counted(conv, x, bias): + nonlocal optimized_calls + result = helpers["candidate"](conv, x, bias) + optimized_calls += result is not None + return result + + with patch.object(forwards, "_first_frame_conv2d", counted): + decode(vae, latents) + sync() + if optimized_calls == 0: + raise RuntimeError("The selected workload did not exercise first-frame Conv2d") + measurements = [] + outputs: dict[str, list[torch.Tensor]] = {arm: [] for arm in helpers} + for _ in range(args.warmup): + for helper in helpers.values(): + with patch.object(forwards, "_first_frame_conv2d", helper): + decode(vae, latents) + sync() + for iteration in range(args.iters): + # Alternating AB and BA pairs give ABBA order without dropping odd runs. + arms = ("baseline", "candidate") if iteration % 2 == 0 else ("candidate", "baseline") + for arm in arms: + with patch.object(forwards, "_first_frame_conv2d", helpers[arm]): + sync() + torch.accelerator.reset_peak_memory_stats() + start = time.perf_counter() + output = decode(vae, latents) + sync() + elapsed = (time.perf_counter() - start) * 1000 + peak = torch.accelerator.max_memory_allocated() + measurements.append({"arm": arm, "wall_ms": elapsed, "peak_allocated_bytes": peak}) + if len(outputs[arm]) < 2: + outputs[arm].append(output.cpu()) + del output + baseline, candidate = outputs["baseline"][0], outputs["candidate"][0] + difference = candidate.float() - baseline.float() + stats = {} + for arm in helpers: + times = [row["wall_ms"] for row in measurements if row["arm"] == arm] + stats[arm] = { + "median_ms": statistics.median(times), + "mean_ms": statistics.mean(times), + "std_ms": statistics.pstdev(times), + "min_ms": min(times), + "max_ms": max(times), + } + result = { + "scope": "VAE decode; same-head channels_last first-frame Conv2d ablation", + "args": vars(args), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "baseline": stats["baseline"], + "candidate": stats["candidate"], + "latency_reduction": 1 - stats["candidate"]["median_ms"] / stats["baseline"]["median_ms"], + "optimized_calls_per_decode": optimized_calls, + "outputs_finite": bool(torch.isfinite(baseline).all() and torch.isfinite(candidate).all()), + "self_repeat_equal": {arm: torch.equal(*values) for arm, values in outputs.items()}, + "candidate_vs_baseline": { + "torch_equal": torch.equal(candidate, baseline), + "max_abs": difference.abs().max().item(), + "relative_l2": (difference.norm() / baseline.float().norm().clamp_min(1e-30)).item(), + "psnr_db": psnr(candidate, baseline), + }, + "measurements": measurements, + } + print("FIRST_FRAME_ABLATION " + json.dumps(result)) + + def main() -> None: args = parse_args() + if args.latents and not args.first_frame_ablation: + raise SystemExit("--latents requires --first-frame-ablation") if not torch.cuda.is_available(): raise SystemExit("CUDA device required") rank, world_size = init_distributed() @@ -267,6 +365,12 @@ def main() -> None: ) device = torch.device("cuda", torch.accelerator.current_device_index()) if world_size > 1 else torch.device("cuda") dtype = DTYPES[args.dtype] + if args.first_frame_ablation: + if world_size != 1 or parallel: + raise SystemExit("--first-frame-ablation requires one GPU") + with torch.backends.cudnn.flags(benchmark=False, allow_tf32=False): + run_first_frame_ablation(args, dtype, device) + return 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: diff --git a/docs/user_guide/diffusion/vae_fast_path.md b/docs/user_guide/diffusion/vae_fast_path.md index 9830b9449..8361592c2 100644 --- a/docs/user_guide/diffusion/vae_fast_path.md +++ b/docs/user_guide/diffusion/vae_fast_path.md @@ -7,7 +7,8 @@ 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. +a handful of Triton kernels. The opt-in `channels_last` level also specializes +first-frame convolutions that have no history. ## Levels @@ -15,9 +16,9 @@ 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`. | +| `channels_last` | no | within tolerance (PSNR typically > 60 dB) | Everything in `lossless`, plus channels-last convolution weights, a single-pass channels-last RMSNorm+SiLU kernel that also absorbs the preceding `conv1` bias, and Conv2d for supported single-frame causal convolutions without history. | | `off` | no | bit-identical | Reference diffusers decoder. | ```bash @@ -61,6 +62,12 @@ stages: - 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 a supported causal convolution with one input frame and no history, + the temporal input is `[0, 0, x]`. The `channels_last` level uses Conv2d with + the last temporal weight slice, preserving the convolution formula for + finite inputs and the existing cache update. Calls with history retain + Conv3d. The weight slice is taken on each call, so weight updates and device + moves do not leave a stale packed copy. - 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 @@ -73,10 +80,10 @@ Cosmos3-Nano VAE, 1280x720 x 189 frames, bf16, one GB200 GPU, `bench_wan_vae_dec (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 | +| --- | ------------- | --------- | --- | +| `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. @@ -93,6 +100,22 @@ python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano --size 1280x720 --frames 189 --fast-path off,lossless,channels_last --profile ``` +For a single-GPU comparison that isolates first-frame Conv2d within +`channels_last`, use `--first-frame-ablation`. It alternates the two paths, +disables cuDNN TF32 and autotuning, and prints JSON with every timing sample, +summary statistics, peak allocated memory and output comparisons. Per-call +weight packing is included in the measured decode time. + +```bash +python benchmarks/diffusion/bench_wan_vae_decode.py \ + --model Wan-AI/Wan2.1-T2V-1.3B-Diffusers --size 832x480 --frames 1 \ + --dtype bf16 --first-frame-ablation --warmup 60 --iters 50 --seed 0 +``` + +The default input is a seeded latent tensor. `--latents /path/to/latents.pt` +can instead load a saved VAE input after latent mean/std scaling; the benchmark +selects the temporal prefix requested by `--frames` and checks its shape. + The same script benchmarks multi-GPU VAE decode when launched with `torchrun`; the decode is timed across all ranks and rank 0 reports: diff --git a/tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py b/tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py index 1b76fc87b..f0c0bb7be 100644 --- a/tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py +++ b/tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py @@ -22,6 +22,7 @@ from tests.diffusion.distributed.wan_vae_fastpath_helpers import ( # noqa: F401 original_wan_rms_norm, unpatched_wan_rms_norm, ) +from tests.helpers.mark import hardware_test 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, @@ -788,6 +789,135 @@ def test_cached_conv_chunk_loop_matches_reference_including_mixed_paths( assert torch.equal(out, ref), pattern +def _first_frame_test_conv(dtype: torch.dtype, **kwargs) -> WanCausalConv3d: + options = dict(kernel_size=3, padding=1) + options.update(kwargs) + conv = WanCausalConv3d(16, 32, **options).eval().to(device="cuda", dtype=dtype) + conv.to(memory_format=torch.channels_last_3d) + setattr(conv, fp.CFG_ATTR, fp.FastPathConfig(channels_last=True)) + return conv + + +@torch.no_grad() +@pytest.mark.parametrize("dtype", ALL_DTYPES) +@pytest.mark.parametrize("temporal_only", [False, True]) +@pytest.mark.parametrize("start", [None, "Rep"]) +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_first_frame_conv2d_preserves_following_cached_frame(dtype, temporal_only, start, monkeypatch) -> None: + torch.manual_seed(19) + kwargs = dict(kernel_size=(3, 1, 1), padding=(1, 0, 0)) if temporal_only else {} + conv = _first_frame_test_conv(dtype, **kwargs) + chunks = [_dense_5d((2, 16, 1, 7, 11), dtype, True) for _ in range(2)] + calls = [] + original = fp._first_frame_conv2d + + def spy(*args): + pair = original(*args) + calls.append(pair is not None) + return pair + + monkeypatch.setattr(fp, "_first_frame_conv2d", spy) + cache = [start] + reference_cache = None + # Disable TF32 so the FP32 comparison isolates Conv2d/Conv3d accumulation. + with torch.backends.cudnn.flags(enabled=True, allow_tf32=False): + for index, chunk in enumerate(chunks): + expected = WanCausalConv3d.forward(conv, chunk, reference_cache) + actual = fp._run_cached_causal_conv(conv, chunk, cache, 0) + assert actual.is_contiguous(memory_format=torch.channels_last_3d) + tol = dict(rtol=1e-4, atol=1e-5) if dtype is torch.float32 else dict(rtol=2e-2, atol=2e-2) + torch.testing.assert_close(actual, expected, **tol) + previous = torch.zeros_like(chunk) if index == 0 else chunks[index - 1] + assert _bits_equal(cache[0], torch.cat([previous, chunk], dim=2)) + assert cache[0].is_contiguous(memory_format=torch.channels_last_3d) + reference_cache = chunk.clone() if index == 0 and start is None else torch.cat([previous, chunk], dim=2) + assert calls == [True], "A later T=1 chunk with history must retain the temporal convolution" + + +@torch.no_grad() +@pytest.mark.parametrize("dtype", LOW_PRECISION) +@pytest.mark.parametrize("with_bias", [False, True]) +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_first_frame_conv2d_deferred_bias_matches_autocast(dtype, with_bias) -> None: + torch.manual_seed(20) + conv = _first_frame_test_conv(torch.float32) + if with_bias: + conv.bias.copy_(torch.randn_like(conv.bias)) + else: + conv.register_parameter("bias", None) + x = _dense_5d((1, 16, 1, 7, 11), torch.float32, True) + cache = [None] + with torch.autocast("cuda", dtype=dtype): + expected = WanCausalConv3d.forward(conv, x) + out, bias = fp._run_cached_causal_conv(conv, x, cache, 0, return_bias=True) + assert out.dtype == expected.dtype == dtype + if with_bias: + assert bias is not None and bias.dtype == dtype + assert _bits_equal(bias, conv.bias.to(dtype)) + out = out + bias.view(1, -1, 1, 1, 1) + else: + assert bias is None + assert isinstance(cache[0], torch.Tensor) + assert cache[0].dtype == x.dtype, "Autocast must not round the temporal cache" + assert _bits_equal(cache[0], torch.cat([torch.zeros_like(x), x], dim=2)) + torch.testing.assert_close(out, expected, rtol=2e-2, atol=2e-2) + + +@torch.no_grad() +@pytest.mark.parametrize( + "case,kwargs", + [ + ("lossless", {}), + ("multiframe", {}), + ("kernel", dict(kernel_size=(3, 5, 5), padding=(1, 2, 2))), + ("stride", dict(stride=(1, 2, 2))), + ("dilation", dict(padding=(1, 2, 2))), + ("groups", {}), + ("temporal_padding", dict(padding=(2, 1, 1))), + ], +) +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_first_frame_conv2d_declines_unsupported_calls(case, kwargs) -> None: + conv = _first_frame_test_conv(torch.float32, **kwargs) + if case == "lossless": + setattr(conv, fp.CFG_ATTR, fp.FastPathConfig()) + elif case == "dilation": + conv.dilation = (1, 2, 2) + elif case == "groups": + conv.groups = 2 + conv.weight = nn.Parameter(conv.weight[:, :8].contiguous(memory_format=torch.channels_last_3d)) + x = _dense_5d((1, 16, 2 if case == "multiframe" else 1, 7, 11), torch.float32, True) + assert fp._first_frame_conv2d(conv, x, conv.bias) is None + expected = WanCausalConv3d.forward(conv, x) + actual = fp._run_cached_causal_conv(conv, x, [None], 0) + assert _bits_equal(actual, expected), "Declined inputs must keep the existing exact fallback" + + +@torch.no_grad() +@pytest.mark.parametrize("attribute,value", [("padding_mode", "reflect"), ("padding", (0, 1, 1))]) +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_first_frame_conv2d_declines_nonzero_native_padding(attribute, value) -> None: + conv = _first_frame_test_conv(torch.float32) + setattr(conv, attribute, value) + x = _dense_5d((1, 16, 1, 7, 11), torch.float32, True) + assert fp._first_frame_conv2d(conv, x, conv.bias) is None + + +@torch.no_grad() +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_first_frame_conv2d_uses_reloaded_weights() -> None: + torch.manual_seed(21) + conv = _first_frame_test_conv(torch.bfloat16) + x = _dense_5d((1, 16, 1, 7, 11), torch.bfloat16, True) + before = fp._run_cached_causal_conv(conv, x, [None], 0) + replacement = {name: -value.clone() for name, value in conv.state_dict().items()} + conv.load_state_dict(replacement) + after = fp._run_cached_causal_conv(conv, x, [None], 0) + expected = WanCausalConv3d.forward(conv, x) + assert not torch.equal(before, after) + torch.testing.assert_close(after, expected, rtol=2e-2, atol=2e-2) + + # --------------------------------------------------------------------------- # # Whole decoder # --------------------------------------------------------------------------- # diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/forwards.py b/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/forwards.py index d3d0bb9b6..eaf698254 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/forwards.py +++ b/vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/forwards.py @@ -1,11 +1,11 @@ # 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. +"""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: +function at the ``lossless`` level 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 @@ -22,8 +22,10 @@ differs in how the bytes move: :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. +The opt-in ``channels_last`` level also changes convolution and normalization +arithmetic. Its first-frame causal Conv3d uses the mathematically equivalent +Conv2d when there is no history; floating-point outputs need not be bit-exact. +Whenever a kernel declines an input, the reference expression is used. """ from __future__ import annotations @@ -310,6 +312,42 @@ def _layout_tag(x: torch.Tensor) -> str: return "channels_last" if x.shape[1] > 1 and x.stride(1) == 1 else "channels_first" +def _first_frame_conv2d( + conv: WanCausalConv3d, x: torch.Tensor, bias: torch.Tensor | None +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Convolve ``[0, 0, x]`` with the last temporal weight slice, retaining the cache. + + Called only for CUDA inference without history. The identity holds for + finite operands, but Conv2d may round differently, so only the tolerance- + based channels_last level uses it. + """ + cfg = getattr(conv, CFG_ATTR, None) + if ( + cfg is None + or not cfg.channels_last + or x.shape[2] != 1 + or x.dtype not in _NORM_DTYPES + or conv.kernel_size not in ((3, 3, 3), (3, 1, 1)) + or conv.stride != (1, 1, 1) + or conv.dilation != (1, 1, 1) + or conv.groups != 1 + or conv.padding != (0, 0, 0) + or conv.padding_mode != "zeros" + or tuple(conv._padding) != ((1, 1, 1, 1, 2, 0) if conv.kernel_size == (3, 3, 3) else (0, 0, 0, 0, 2, 0)) + ): + return None + pair = dm.cat_time_5d(x, None, conv._padding[4], keep_cache_frames=CACHE_T) + if pair is None: + return None + _, next_cache = pair + # Materialize only the current weight slice; retaining a packed copy could + # become stale after weight loading or keep offloaded weights on the GPU. + weight = conv.weight[:, :, 2].contiguous(memory_format=torch.channels_last) + image = x[:, :, 0].contiguous(memory_format=torch.channels_last) + out = F.conv2d(image, weight, bias, padding=(conv._padding[2], conv._padding[0])) + return out.unsqueeze(2), next_cache + + def _run_cached_causal_conv( conv: nn.Module, x: torch.Tensor, @@ -349,6 +387,12 @@ def _run_cached_causal_conv( fold_bias = return_bias and conv.bias is not None bias = None if fold_bias else conv.bias + if payload is None: + first_frame = _first_frame_conv2d(conv, x, bias) + if first_frame is not None: + out, cache_list[index] = first_frame + return (out, _deferred_conv_bias(conv, out)) if return_bias else out + # Preferred: temporal concat only (aligned plane copies) + cuDNN spatial # padding, once verified bitwise for this (conv, shape). verdicts = _SPATIAL_PAD_VERDICTS.setdefault(conv, {})