From 9f4ed93dd8f0205774019193d2b6487c53417892 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Tue, 16 Dec 2025 22:17:33 +0800 Subject: [PATCH] [diffusion] multi-platform: use current_platform.device_type to replace hard-coded cuda device (#15232) Signed-off-by: Xiaodong Ye --- .../configs/sample/sampling_params.py | 4 ++-- .../runtime/loader/component_loader.py | 2 +- .../multimodal_gen/runtime/loader/fsdp_load.py | 2 +- .../runtime/models/dits/hunyuanvideo.py | 13 ++++++++++--- .../multimodal_gen/runtime/models/dits/zimage.py | 7 +++++-- .../runtime/models/encoders/stepllm.py | 5 ++++- .../runtime/pipelines_core/schedule_batch.py | 4 +++- .../runtime/pipelines_core/stages/base.py | 5 ++++- .../pipelines_core/stages/causal_denoising.py | 15 +++++++++++---- .../runtime/pipelines_core/stages/decoding.py | 5 ++++- .../runtime/pipelines_core/stages/denoising.py | 3 ++- .../pipelines_core/stages/denoising_dmd.py | 3 ++- .../runtime/pipelines_core/stages/encoding.py | 5 ++++- .../pipelines_core/stages/image_encoding.py | 5 ++++- .../pipelines_core/stages/input_validation.py | 3 ++- .../multimodal_gen/runtime/utils/distributed.py | 4 +++- 16 files changed, 62 insertions(+), 23 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 86e743995..b3f79d0a6 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/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", diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index 4fdc9cdf1..f93ae0ad2 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/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"), ) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 3ca0592be..9b0595e17 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/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"), diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py index ad9e10dd5..730067fa9 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py +++ b/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()], ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index 3aa2437bb..264f66492 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py b/python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py index dae653eac..26e3a8670 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index e962c21ba..c6d9b1da4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py index e79ea187b..1adecbfe0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/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): """ diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py index e06ee9bed..6145ac0a6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py index bf5932b14..d405b95c1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 8f30f5aeb..f3ef7d9a4 100755 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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, ): diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py index 538bab8a5..3cd1ea95b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py +++ b/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, ): diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py index 6fe00884c..9e30a8cb7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py +++ b/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() diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py index a86db6b46..4a6eb1ebc 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py +++ b/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() diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py index 6b17eb923..7d68a14a4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/utils/distributed.py b/python/sglang/multimodal_gen/runtime/utils/distributed.py index 24608e9e9..e63a046b8 100644 --- a/python/sglang/multimodal_gen/runtime/utils/distributed.py +++ b/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: