[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/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', "