[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 <liuhongsheng4@huawei.com>

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:

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
 
@@ -17,7 +18,7 @@ The fast path is controlled by `--vae-fast-path` (engine argument
 | 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
@@ -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:
 

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
 # --------------------------------------------------------------------------- #

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, {})