[Perf] Wan2.2 I2V optimization: convert datatype from FP32 to BF16 in vae (#2391)

Signed-off-by: KexiongYu <yukexiong1@huawei.com>
Signed-off-by: Canlin Guo <961750412@qq.com>
Co-authored-by: Canlin Guo <961750412@qq.com>

tests/diffusion/distributed/test_autoencoder_kl_wan.py

@@ -0,0 +1,43 @@
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_wan as wan_vae_module
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummyOmniAutoencoderKLWan(OmniAutoencoderKLWan):
+    def __init__(self, *, dtype: torch.dtype):
+        torch.nn.Module.__init__(self)
+        self.register_parameter("dummy_weight", torch.nn.Parameter(torch.ones(1, dtype=dtype)))
+
+
+def test_wan_vae_execution_context_handles_fp32():
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.float32)
+    with model._execution_context():
+        output = model.dummy_weight + 1
+    assert output.dtype == torch.float32
+
+
+def test_wan_vae_execution_context_handles_bf16():
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
+    with model._execution_context():
+        output = model.dummy_weight + 1
+    assert output.dtype == torch.bfloat16
+
+
+def test_wan_vae_execution_context_uses_platform_autocast(mocker):
+    sentinel = object()
+    platform = mocker.Mock()
+    platform.create_autocast_context.return_value = sentinel
+    mocker.patch.object(wan_vae_module, "current_omni_platform", platform)
+
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
+
+    assert model._execution_context() is sentinel
+    platform.create_autocast_context.assert_called_once_with(
+        device_type=model.dummy_weight.device.type,
+        dtype=torch.bfloat16,
+        enabled=True,
+    )

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from contextlib import nullcontext
 from typing import Any
 
 import torch
@@ -15,11 +16,38 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     GridSpec,
     TileTask,
 )
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
 
-class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
+class OmniAutoencoderKLWan(AutoencoderKLWan):
+    def _execution_context(self):
+        try:
+            first_param = next(self.parameters())
+        except StopIteration:
+            return nullcontext()
+
+        dtype = first_param.dtype
+        if dtype not in (torch.float16, torch.bfloat16):
+            return nullcontext()
+
+        return current_omni_platform.create_autocast_context(
+            device_type=first_param.device.type,
+            dtype=dtype,
+            enabled=True,
+        )
+
+    def encode(self, x: torch.Tensor, return_dict: bool = True):
+        with self._execution_context():
+            return super().encode(x, return_dict=return_dict)
+
+    def decode(self, z: torch.Tensor, return_dict: bool = True):
+        with self._execution_context():
+            return super().decode(z, return_dict=return_dict)
+
+
+class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
     @classmethod
     def from_pretrained(cls, *args: Any, **kwargs: Any):
         model = super().from_pretrained(*args, **kwargs)
@@ -84,11 +112,12 @@ class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
         """Decode a single latent tile into RGB space."""
         self.clear_cache()
         time = []
-        for k in range(len(task.tensor)):
-            self._conv_idx = [0]
-            tile = self.post_quant_conv(task.tensor[k])
-            decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
-            time.append(decoded)
+        with self._execution_context():
+            for k in range(len(task.tensor)):
+                self._conv_idx = [0]
+                tile = self.post_quant_conv(task.tensor[k])
+                decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
+                time.append(decoded)
         result = torch.cat(time, dim=2)
         return result
 

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -272,7 +272,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
             model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
         ).to(self.device)
         self.vae = DistributedAutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
         ).to(self.device)
 
         # Initialize transformers with correct config (weights loaded via load_weights)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -217,7 +217,7 @@ class Wan22I2VPipeline(
 
         # VAE
         self.vae = DistributedAutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
         ).to(self.device)
 
         # Transformers (weights loaded via load_weights)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py

@@ -24,13 +24,13 @@ from typing import Any, cast
 import numpy as np
 import PIL.Image
 import torch
-from diffusers import AutoencoderKLWan
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import AutoTokenizer, UMT5EncoderModel
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -174,8 +174,8 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress
         ).to(self.device)
 
         # VAE
-        self.vae = AutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+        self.vae = OmniAutoencoderKLWan.from_pretrained(
+            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
         ).to(self.device)
 
         # Single transformer (TI2V uses dense 5B model, not MoE)

vllm_omni/platforms/interface.py

@@ -1,12 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from contextlib import nullcontext
 from enum import Enum
 from typing import Any
 
 import torch
+from vllm.logger import init_logger
 from vllm.platforms import Platform
 
+logger = init_logger(__name__)
+
 
 class OmniPlatformEnum(Enum):
     """Enum for supported Omni platforms."""
@@ -113,6 +117,23 @@ class OmniPlatform(Platform):
     def get_free_memory(cls, device: torch.device | None = None) -> int:
         raise NotImplementedError
 
+    @classmethod
+    def create_autocast_context(
+        cls,
+        *,
+        device_type: str,
+        dtype: torch.dtype,
+        enabled: bool = True,
+    ):
+        if not enabled:
+            return nullcontext()
+
+        try:
+            return torch.autocast(device_type=device_type, dtype=dtype, enabled=True)
+        except (RuntimeError, TypeError, ValueError) as exc:
+            logger.warning("autocast unavailable for device_type=%s dtype=%s: %s", device_type, dtype, exc)
+            return nullcontext()
+
     @classmethod
     def supports_cpu_offload(cls) -> bool:
         return True

vllm_omni/platforms/npu/platform.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from contextlib import nullcontext
 from typing import Any
 
 import torch
@@ -106,6 +107,24 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
         device_props = torch.npu.get_device_properties(device_id)
         return device_props.total_memory
 
+    @classmethod
+    def create_autocast_context(cls, *, device_type, dtype, enabled=True):
+        if device_type != "npu":
+            return super().create_autocast_context(
+                device_type=device_type,
+                dtype=dtype,
+                enabled=enabled,
+            )
+        if not enabled:
+            return nullcontext()
+
+        # NPU-specific fallback
+        try:
+            return torch.npu.amp.autocast(dtype=dtype)
+        except (RuntimeError, TypeError, ValueError) as exc:
+            logger.warning("autocast unavailable for device_type=%s dtype=%s: %s", device_type, dtype, exc)
+        return nullcontext()
+
     @classmethod
     def get_profiler_cls(cls) -> str:
         return "vllm_omni.platforms.npu.profiler.NPUTorchProfilerWrapper"