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

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