[diffusion] Add is_float64_supported to Platform (#22112)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>

python/sglang/multimodal_gen/runtime/layers/linear.py

@@ -158,7 +158,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
             F.linear(x, layer.weight, bias)
             if current_platform.is_amp_supported() or bias is None
             else F.linear(x, layer.weight, bias.to(x.dtype))
-        )  # NOTE: this line assumes that we are using amp when using cuda and is needed to account for the fact that amp isn't supported in mps
+        )  # NOTE: explicit dtype cast for bias is needed on platforms where amp isn't supported
         return output
 
 

python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py

@@ -644,9 +644,9 @@ class CausalWanTransformer3DModel(BaseDiT, OffloadableDiTMixin):
             self.num_attention_heads,
             rope_dim_list,
             dtype=(
-                torch.float32
+                torch.float64
-                if current_platform.is_mps() or current_platform.is_musa()
+                if current_platform.is_float64_supported()
-                else torch.float64
+                else torch.float32
             ),
             rope_theta=10000,
             start_frame=start_frame,  # Assume that start_frame is 0 when kv_cache is None
@@ -776,9 +776,9 @@ class CausalWanTransformer3DModel(BaseDiT, OffloadableDiTMixin):
             self.num_attention_heads,
             rope_dim_list,
             dtype=(
-                torch.float32
+                torch.float64
-                if current_platform.is_mps() or current_platform.is_musa()
+                if current_platform.is_float64_supported()
-                else torch.float64
+                else torch.float32
             ),
             rope_theta=10000,
             start_frame=start_frame,

python/sglang/multimodal_gen/runtime/models/dits/flux.py

@@ -722,9 +722,9 @@ class FluxPosEmbed(nn.Module):
             use_real=False,
             repeat_interleave_real=False,
             dtype=(
-                torch.float32
+                torch.float64
-                if current_platform.is_mps() or current_platform.is_musa()
+                if current_platform.is_float64_supported()
-                else torch.float64
+                else torch.float32
             ),
         )
 

python/sglang/multimodal_gen/runtime/models/dits/flux_2.py

@@ -815,9 +815,9 @@ class Flux2PosEmbed(nn.Module):
             use_real=False,
             repeat_interleave_real=False,
             dtype=(
-                torch.float32
+                torch.float64
-                if current_platform.is_mps() or current_platform.is_musa()
+                if current_platform.is_float64_supported()
-                else torch.float64
+                else torch.float32
             ),
         )
 

python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py

@@ -915,9 +915,9 @@ class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin):
             rope_dim_list=self.rope_dim_list,
             rope_theta=10000,
             dtype=(
-                torch.float32
+                torch.float64
-                if current_platform.is_mps() or current_platform.is_musa()
+                if current_platform.is_float64_supported()
-                else torch.float64
+                else torch.float32
             ),
         )
 
@@ -1090,7 +1090,7 @@ class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin):
             encoder_hidden_states.to(orig_dtype)
             if not current_platform.is_amp_supported()
             else encoder_hidden_states
-        )  # cast to orig_dtype for MPS
+        )  # cast to orig_dtype if amp is not supported
 
         assert encoder_hidden_states.dtype == orig_dtype
 

python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_common_utils.py

@@ -149,8 +149,8 @@ class WanCausalConv3d(nn.Conv3d):
             padding[4] -= cache_x.shape[2]
         x = F.pad(x, padding)
         x = (
-            x.to(self.weight.dtype) if current_platform.is_mps() else x
+            x if current_platform.is_amp_supported() else x.to(self.weight.dtype)
-        )  # casting needed for mps since amp isn't supported
+        )  # casting needed if amp isn't supported
         return super().forward(x)
 
 

python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_dist_utils.py

@@ -304,8 +304,8 @@ class WanDistCausalConv3d(nn.Conv3d):
         x = F.pad(x, padding)
 
         x = (
-            x.to(self.weight.dtype) if current_platform.is_mps() else x
+            x if current_platform.is_amp_supported() else x.to(self.weight.dtype)
-        )  # casting needed for mps since amp isn't supported
+        )  # casting needed if amp isn't supported
 
         x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
             x,

python/sglang/multimodal_gen/runtime/platforms/interface.py

@@ -199,6 +199,11 @@ class Platform:
     def is_amp_supported(cls) -> bool:
         return True
 
+    @classmethod
+    @lru_cache(maxsize=1)
+    def is_float64_supported(cls) -> bool:
+        return True
+
     @classmethod
     def get_modelopt_fp4_quantize_op(cls) -> Callable | None:
         return None

python/sglang/multimodal_gen/runtime/platforms/mps.py

@@ -31,6 +31,11 @@ class MpsPlatform(Platform):
     def is_amp_supported(cls) -> bool:
         return False
 
+    @classmethod
+    @lru_cache(maxsize=1)
+    def is_float64_supported(cls) -> bool:
+        return False
+
     @classmethod
     def get_local_torch_device(cls) -> torch.device:
         return torch.device("mps")

python/sglang/multimodal_gen/runtime/platforms/musa.py

@@ -71,6 +71,11 @@ class MusaPlatformBase(Platform):
     dispatch_key: str = "MUSA"
     device_control_env_var: str = "MUSA_VISIBLE_DEVICES"
 
+    @classmethod
+    @lru_cache(maxsize=1)
+    def is_float64_supported(cls) -> bool:
+        return False
+
     @classmethod
     def get_local_torch_device(cls) -> torch.device:
         return torch.device(f"musa:{envs.LOCAL_RANK}")