[diffusion] multi-platform: use current_platform.device_type to replace hard-coded cuda device (#15232)

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

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -430,8 +430,8 @@ class SamplingParams:
             "--generator-device",
             type=str,
             default=SamplingParams.generator_device,
-            choices=["cuda", "cpu"],
-            help="Device for random generator (cuda or cpu). Default: cuda",
+            choices=["cuda", "musa", "cpu"],
+            help="Device for random generator (cuda, musa or cpu). Default: cuda",
         )
         parser.add_argument(
             "--num-frames",

python/sglang/multimodal_gen/runtime/loader/component_loader.py

@@ -448,7 +448,7 @@ class TextEncoderLoader(ComponentLoader):
                     )
                 else:
                     mesh = init_device_mesh(
-                        "cuda",
+                        current_platform.device_type,
                         mesh_shape=(1, dist.get_world_size()),
                         mesh_dim_names=("offload", "replicate"),
                     )

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -116,7 +116,7 @@ def maybe_load_fsdp_model(
             hsdp_shard_dim = 1
 
         device_mesh = init_device_mesh(
-            "cuda",
+            current_platform.device_type,
             # (Replicate(), Shard(dim=0))
             mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
             mesh_dim_names=("replicate", "shard"),

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

@@ -36,7 +36,10 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import (
 from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.models.utils import modulate
-from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
 
 
 class MMDoubleStreamBlock(nn.Module):
@@ -677,7 +680,9 @@ class HunyuanVideoTransformer3DModel(CachableDiT):
         vec_ = torch.distributed.tensor.DTensor.from_local(
             vec_,
             torch.distributed.DeviceMesh(
-                "cuda", list(range(get_sp_world_size())), mesh_dim_names=("dp",)
+                current_platform.device_type,
+                list(range(get_sp_world_size())),
+                mesh_dim_names=("dp",),
             ),
             [torch.distributed.tensor.Replicate()],
         )
@@ -685,7 +690,9 @@ class HunyuanVideoTransformer3DModel(CachableDiT):
         inp = torch.distributed.tensor.DTensor.from_local(
             inp,
             torch.distributed.DeviceMesh(
-                "cuda", list(range(get_sp_world_size())), mesh_dim_names=("dp",)
+                current_platform.device_type,
+                list(range(get_sp_world_size())),
+                mesh_dim_names=("dp",),
             ),
             [torch.distributed.tensor.Replicate()],
         )

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

@@ -11,7 +11,10 @@ from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
 from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import _apply_rotary_emb
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
-from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
@@ -46,7 +49,7 @@ class TimestepEmbedder(nn.Module):
 
     @staticmethod
     def timestep_embedding(t, dim, max_period=10000):
-        with torch.amp.autocast("cuda", enabled=False):
+        with torch.amp.autocast(current_platform.device_type, enabled=False):
             half = dim // 2
             freqs = torch.exp(
                 -math.log(max_period)

python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py

@@ -24,6 +24,7 @@ from einops import rearrange
 from transformers import PretrainedConfig, PreTrainedModel
 
 from sglang.multimodal_gen.runtime.models.dits.stepvideo import StepVideoRMSNorm
+from sglang.multimodal_gen.runtime.platforms import current_platform
 
 
 class EmptyInitOnDevice(torch.overrides.TorchFunctionMode):
@@ -591,7 +592,9 @@ class STEP1TextEncoder(torch.nn.Module):
     def forward(self, prompts, with_mask=True, max_length=None):
         self.device = next(self.text_encoder.parameters()).device
 
-        with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        with torch.no_grad(), torch.amp.autocast(
+            current_platform.device_type, dtype=torch.bfloat16
+        ):
             if type(prompts) is str:
                 prompts = [prompts]
             txt_tokens = self.text_tokenizer(

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -89,7 +89,9 @@ class Req:
     num_outputs_per_prompt: int = 1
     seed: int | None = None
     seeds: list[int] | None = None
-    generator_device: str = "cuda"  # Device for random generator: "cuda" or "cpu"
+    generator_device: str = (
+        "cuda"  # Device for random generator: "cuda", "musa" or "cpu"
+    )
 
     # Tracking if embeddings are already processed
     is_prompt_processed: bool = False

python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py

@@ -17,6 +17,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
     VerificationResult,
 )
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
@@ -147,7 +148,9 @@ class PipelineStage(ABC):
     @property
     def device(self) -> torch.device:
         """Get the device for this stage."""
-        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        return torch.device(
+            current_platform.device_type,
+        )
 
     def set_logging(self, enable: bool):
         """

python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py

@@ -13,6 +13,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
 from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
     VerificationResult,
 )
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -167,7 +168,9 @@ class CausalDMDDenoisingStage(DenoisingStage):
                     image_latent[:, :, :1, :, :].to(target_dtype).permute(0, 2, 1, 3, 4)
                 )
                 with torch.autocast(
-                    device_type="cuda", dtype=target_dtype, enabled=autocast_enabled
+                    device_type=current_platform.device_type,
+                    dtype=target_dtype,
+                    enabled=autocast_enabled,
                 ):
                     _ = self.transformer(
                         image_first_btchw,
@@ -195,7 +198,9 @@ class CausalDMDDenoisingStage(DenoisingStage):
                     .permute(0, 2, 1, 3, 4)
                 )
                 with torch.autocast(
-                    device_type="cuda", dtype=target_dtype, enabled=autocast_enabled
+                    device_type=current_platform.device_type,
+                    dtype=target_dtype,
+                    enabled=autocast_enabled,
                 ):
                     _ = self.transformer(
                         ref_btchw,
@@ -295,7 +300,7 @@ class CausalDMDDenoisingStage(DenoisingStage):
 
                     with (
                         torch.autocast(
-                            device_type="cuda",
+                            device_type=current_platform.device_type,
                             dtype=target_dtype,
                             enabled=autocast_enabled,
                         ),
@@ -372,7 +377,9 @@ class CausalDMDDenoisingStage(DenoisingStage):
                 context_bcthw = current_latents.to(target_dtype)
                 with (
                     torch.autocast(
-                        device_type="cuda", dtype=target_dtype, enabled=autocast_enabled
+                        device_type=current_platform.device_type,
+                        dtype=target_dtype,
+                        enabled=autocast_enabled,
                     ),
                     set_forward_context(
                         current_timestep=0,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py

@@ -20,6 +20,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
 from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
     VerificationResult,
 )
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
@@ -133,7 +134,9 @@ class DecodingStage(PipelineStage):
 
         # Decode latents
         with torch.autocast(
-            device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled
+            device_type=current_platform.device_type,
+            dtype=vae_dtype,
+            enabled=vae_autocast_enabled,
         ):
             try:
                 # TODO: make it more specific

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -56,6 +56,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
 from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
     VerificationResult,
 )
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.platforms.interface import AttentionBackendEnum
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -963,7 +964,7 @@ class DenoisingStage(PipelineStage):
         timesteps_cpu = timesteps.cpu()
         num_timesteps = timesteps_cpu.shape[0]
         with torch.autocast(
-            device_type=("cuda" if torch.cuda.is_available() else "cpu"),
+            device_type=current_platform.device_type,
             dtype=target_dtype,
             enabled=autocast_enabled,
         ):

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py

@@ -12,6 +12,7 @@ from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler
 from sglang.multimodal_gen.runtime.models.utils import pred_noise_to_pred_video
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages import DenoisingStage
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
@@ -127,7 +128,7 @@ class DmdDenoisingStage(DenoisingStage):
 
                     # Predict noise residual
                     with torch.autocast(
-                        device_type="cuda",
+                        device_type=current_platform.device_type,
                         dtype=target_dtype,
                         enabled=autocast_enabled,
                     ):

python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py

@@ -17,6 +17,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
 from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
     VerificationResult,
 )
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
@@ -83,7 +84,9 @@ class EncodingStage(PipelineStage):
 
         # Encode image to latents
         with torch.autocast(
-            device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled
+            device_type=current_platform.device_type,
+            dtype=vae_dtype,
+            enabled=vae_autocast_enabled,
         ):
             if server_args.pipeline_config.vae_tiling:
                 self.vae.enable_tiling()

python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py

@@ -30,6 +30,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
 from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
     VerificationResult,
 )
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
@@ -253,7 +254,9 @@ class ImageVAEEncodingStage(PipelineStage):
 
             # Encode Image
             with torch.autocast(
-                device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled
+                device_type=current_platform.device_type,
+                dtype=vae_dtype,
+                enabled=vae_autocast_enabled,
             ):
                 if server_args.pipeline_config.vae_tiling:
                     self.vae.enable_tiling()

python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py

@@ -18,6 +18,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
     StageValidators,
     VerificationResult,
 )
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import best_output_size
@@ -61,7 +62,7 @@ class InputValidationStage(PipelineStage):
         if generator_device == "cpu":
             device_str = "cpu"
         else:
-            device_str = "cuda" if torch.cuda.is_available() else "cpu"
+            device_str = current_platform.device_type
 
         batch.generator = [
             torch.Generator(device_str).manual_seed(seed) for seed in seeds

python/sglang/multimodal_gen/runtime/utils/distributed.py

@@ -7,6 +7,8 @@ import numpy as np
 import torch
 import torch.distributed as dist
 
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
 
 def broadcast_pyobj(
     data: List[Any],
@@ -20,7 +22,7 @@ def broadcast_pyobj(
     of dist_group argument).
     """
     device = torch.device(
-        "cuda" if torch.cuda.is_available() and not force_cpu_device else "cpu"
+        current_platform.device_type if not force_cpu_device else "cpu"
     )
 
     if rank == src: