Add supports_float64() to OmniPlatform and clean up MPS (#2488)

Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>

vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py

@@ -681,13 +681,8 @@ class FluxKontextPipeline(
                     neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
                     noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
 
-                latents_dtype = latents.dtype
                 latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
 
-                if latents.dtype != latents_dtype:
-                    if torch.backends.mps.is_available():
-                        latents = latents.to(latents_dtype)
-
                 if callback_on_step_end is not None:
                     callback_kwargs = {}
                     for k in callback_on_step_end_tensor_inputs:

vllm_omni/diffusion/models/flux2/pipeline_flux2.py

@@ -1062,12 +1062,8 @@ class Flux2Pipeline(nn.Module, SupportImageInput, ProgressBarMixin, DiffusionPip
                 noise_pred = noise_pred[:, : latents.size(1) :]
 
                 # compute the previous noisy sample x_t -> x_t-1
-                latents_dtype = latents.dtype
                 latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
 
-                if latents.dtype != latents_dtype and torch.backends.mps.is_available():
-                    latents = latents.to(latents_dtype)
-
                 if callback_on_step_end is not None:
                     callback_kwargs = {}
                     for k in callback_on_step_end_tensor_inputs:

vllm_omni/diffusion/models/mammoth_moda2/rope_real.py

@@ -18,6 +18,8 @@ from diffusers.models.embeddings import get_1d_rotary_pos_embed as _get_1d_rotar
 from einops import repeat
 from torch import nn
 
+from vllm_omni.platforms import current_omni_platform
+
 
 def apply_real_rotary_emb(x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor) -> torch.Tensor:
     """
@@ -119,7 +121,7 @@ class RotaryPosEmbedReal(nn.Module):
         axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], theta: int
     ) -> list[tuple[torch.Tensor, torch.Tensor]]:
         freqs_real = []
-        freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
+        freqs_dtype = torch.float64 if current_omni_platform.supports_float64() else torch.float32
         for i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
             cos_emb, sin_emb = get_1d_rotary_pos_embed_real(d, e, theta=theta, freqs_dtype=freqs_dtype)
             freqs_real.append((cos_emb, sin_emb))

vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py

@@ -19,6 +19,7 @@ from vllm.model_executor.layers.linear import (
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
 
@@ -411,7 +412,7 @@ class OmniGen2RotaryPosEmbed(nn.Module):
         axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], theta: int
     ) -> list[torch.Tensor]:
         freqs_cis = []
-        freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
+        freqs_dtype = torch.float64 if current_omni_platform.supports_float64() else torch.float32
         for i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
             emb = get_1d_rotary_pos_embed(d, e, theta=theta, freqs_dtype=freqs_dtype)
             freqs_cis.append(emb)

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -29,6 +29,7 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelOutput,
 )
 from vllm_omni.diffusion.forward_context import get_forward_context
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -171,7 +172,7 @@ class WanRotaryPosEmbed(nn.Module):
         # Split dimensions for temporal, height, width
         h_dim = w_dim = 2 * (attention_head_dim // 6)
         t_dim = attention_head_dim - h_dim - w_dim
-        freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
+        freqs_dtype = torch.float64 if current_omni_platform.supports_float64() else torch.float32
 
         freqs_cos = []
         freqs_sin = []

vllm_omni/platforms/interface.py

@@ -117,6 +117,10 @@ class OmniPlatform(Platform):
     def supports_cpu_offload(cls) -> bool:
         return True
 
+    @classmethod
+    def supports_float64(cls) -> bool:
+        return True
+
     @classmethod
     def set_device_control_env_var(cls, devices: str | int | None) -> None:
         import os

vllm_omni/platforms/musa/platform.py

@@ -81,6 +81,11 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
         """MUSA supports torch.compile with inductor backend."""
         return True
 
+    @classmethod
+    def supports_float64(cls) -> bool:
+        """MUSA does not support float64 yet."""
+        return False
+
     @classmethod
     def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
         """Get the torch device for MUSA platform.