[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: