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

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

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"]