[diffusion] hardware: support diffusion (single GPU, 3/N) (#17105)

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

docker/rocm.Dockerfile

@@ -184,9 +184,9 @@ RUN git clone ${SGL_REPO} \
     && cd .. \
     && rm -rf python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml \
     && if [ "$BUILD_TYPE" = "srt" ]; then \
-         python -m pip --no-cache-dir install -e "python[srt_hip,diffusion]"; \
+         python -m pip --no-cache-dir install -e "python[srt_hip,diffusion_hip]"; \
        else \
-         python -m pip --no-cache-dir install -e "python[all_hip,diffusion]"; \
+         python -m pip --no-cache-dir install -e "python[all_hip,diffusion_hip]"; \
        fi
 
 RUN python -m pip cache purge

docs/platforms/amd_gpu.md

@@ -55,7 +55,7 @@ python setup_rocm.py install
 # Install sglang python package along with diffusion support
 cd ..
 rm -rf python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml
-pip install -e "python[all_hip,diffusion]"
+pip install -e "python[all_hip,diffusion_hip]"
 ```
 
 ### Install Using Docker (Recommended)

python/pyproject.toml

@@ -93,7 +93,7 @@ diffusion = [
   "moviepy>=2.0.0",
   "opencv-python-headless==4.10.0.84",
   "remote-pdb",
-  "st_attn ==0.0.7",
+  "st_attn==0.0.7",
   "vsa==0.0.4",
   "runai_model_streamer",
   "cache-dit==1.1.8"

python/pyproject_other.toml

@@ -66,6 +66,7 @@ runtime_common = [
   "grpcio==1.75.1", # keep it align with compile_proto.py
   "grpcio-tools==1.75.1", # keep it align with compile_proto.py
   "grpcio-reflection==1.75.1", # required by srt/entrypoints/grpc_server.py
+  "bidict",
 ]
 
 tracing = [
@@ -84,26 +85,7 @@ srt_hip = [
   "wave-lang==3.8.2",
 ]
 
-# https://docs.sglang.io/platforms/ascend_npu.html
+diffusion_hip = [
-srt_npu = ["sglang[runtime_common]"]
-
-# For Intel Gaudi(device : hpu) follow the installation guide
-# https://docs.vllm.ai/en/latest/getting_started/gaudi-installation.html
-srt_hpu = ["sglang[runtime_common]"]
-
-test = [
-  "accelerate",
-  "expecttest",
-  "gguf",
-  "jsonlines",
-  "matplotlib",
-  "pandas",
-  "peft",
-  "pytest",
-  "sentence_transformers",
-  "tabulate",
-]
-diffusion = [
   "diffusers @ git+https://github.com/huggingface/diffusers.git@6290fdfda40610ce7b99920146853614ba529c6e",
   "opencv-python-headless==4.10.0.84",
   "imageio==2.36.0",
@@ -117,13 +99,61 @@ diffusion = [
   "vsa==0.0.4",
   "runai_model_streamer",
 ]
+
+# https://docs.sglang.io/platforms/ascend_npu.html
+srt_npu = ["sglang[runtime_common]"]
+
+# For Intel Gaudi(device : hpu) follow the installation guide
+# https://docs.vllm.ai/en/latest/getting_started/gaudi-installation.html
+srt_hpu = ["sglang[runtime_common]"]
+
+# https://docs.sglang.io/platforms/mthreads_gpu.md
+srt_musa = [
+    "sglang[runtime_common]",
+    "torch",
+    "torch_musa",
+    "torchada>=0.1.15",
+    "mthreads-ml-py",
+    "numpy<2.0",
+]
+
+diffusion_musa = [
+  "PyYAML==6.0.1",
+  "cloudpickle",
+  "diffusers==0.36.0",
+  "imageio==2.36.0",
+  "imageio-ffmpeg==0.5.1",
+  "moviepy>=2.0.0",
+  "opencv-python-headless==4.10.0.84",
+  "remote-pdb",
+  "st_attn==0.0.7",
+  "vsa==0.0.4",
+  "runai_model_streamer",
+  "cache-dit==1.1.8"
+]
+
+test = [
+  "accelerate",
+  "expecttest",
+  "gguf",
+  "jsonlines",
+  "matplotlib",
+  "pandas",
+  "peft",
+  "pytest",
+  "sentence_transformers",
+  "tabulate",
+]
+
 all_hip = ["sglang[srt_hip]"]
 all_npu = ["sglang[srt_npu]"]
 all_hpu = ["sglang[srt_hpu]"]
+all_musa = ["sglang[srt_musa]", "sglang[diffusion_musa]"]
 
 dev_hip = ["sglang[all_hip]", "sglang[test]"]
 dev_npu = ["sglang[all_npu]", "sglang[test]"]
 dev_hpu = ["sglang[all_hpu]", "sglang[test]"]
+dev_musa = ["sglang[all_musa]", "sglang[test]"]
 
 [project.urls]
 "Homepage" = "https://github.com/sgl-project/sglang"

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

@@ -242,8 +242,12 @@ def init_distributed_environment(
             "distributed environment"
         )
 
-        # For MPS, don't pass device_id as it doesn't support device indices
+        # For MPS and MUSA, don't pass device_id as it doesn't support device indices
-        extra_args = {} if current_platform.is_mps() else dict(device_id=device_id)
+        extra_args = (
+            {}
+            if (current_platform.is_mps() or current_platform.is_musa())
+            else dict(device_id=device_id)
+        )
         torch.distributed.init_process_group(
             backend=backend,
             init_method=distributed_init_method,

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

@@ -633,7 +633,11 @@ class CausalWanTransformer3DModel(BaseDiT, OffloadableDiTMixin):
             self.hidden_size,
             self.num_attention_heads,
             rope_dim_list,
-            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
+            dtype=(
+                torch.float32
+                if current_platform.is_mps() or current_platform.is_musa()
+                else torch.float64
+            ),
             rope_theta=10000,
             start_frame=start_frame,  # Assume that start_frame is 0 when kv_cache is None
         )
@@ -761,7 +765,11 @@ class CausalWanTransformer3DModel(BaseDiT, OffloadableDiTMixin):
             self.hidden_size,
             self.num_attention_heads,
             rope_dim_list,
-            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
+            dtype=(
+                torch.float32
+                if current_platform.is_mps() or current_platform.is_musa()
+                else torch.float64
+            ),
             rope_theta=10000,
             start_frame=start_frame,
         )

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

@@ -397,7 +397,11 @@ class FluxPosEmbed(nn.Module):
             rope_theta=theta,
             use_real=False,
             repeat_interleave_real=False,
-            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
+            dtype=(
+                torch.float32
+                if current_platform.is_mps() or current_platform.is_musa()
+                else torch.float64
+            ),
         )
 
     def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

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

@@ -626,7 +626,11 @@ class Flux2PosEmbed(nn.Module):
             rope_theta=theta,
             use_real=False,
             repeat_interleave_real=False,
-            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
+            dtype=(
+                torch.float32
+                if current_platform.is_mps() or current_platform.is_musa()
+                else torch.float64
+            ),
         )
 
     def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

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

@@ -116,7 +116,7 @@ class QwenEmbedRope(nn.Module):
         #     rope_theta=theta,
         #     use_real=False,
         #     repeat_interleave_real=False,
-        #     dtype=torch.float32 if current_platform.is_mps() else torch.float64,
+        #     dtype=torch.float32 if current_platform.is_mps() or current_platform.is_musa() else torch.float64,
         # )
 
         # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART

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

@@ -774,7 +774,11 @@ class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin):
         self.rotary_emb = NDRotaryEmbedding(
             rope_dim_list=self.rope_dim_list,
             rope_theta=10000,
-            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
+            dtype=(
+                torch.float32
+                if current_platform.is_mps() or current_platform.is_musa()
+                else torch.float64
+            ),
         )
 
         self.layer_names = ["blocks"]

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

@@ -101,11 +101,31 @@ def rocm_platform_plugin() -> str | None:
     )
 
 
+def musa_platform_plugin() -> str | None:
+    is_musa = False
+
+    try:
+        import pymtml
+
+        pymtml.mtmlLibraryInit()
+        try:
+            is_musa = pymtml.mtmlLibraryCountDevice() > 0
+        finally:
+            pymtml.mtmlLibraryShutDown()
+    except Exception as e:
+        logger.info("MUSA platform is unavailable: %s", e)
+
+    return (
+        "sglang.multimodal_gen.runtime.platforms.musa.MusaPlatform" if is_musa else None
+    )
+
+
 builtin_platform_plugins = {
     "cuda": cuda_platform_plugin,
     "rocm": rocm_platform_plugin,
     "mps": mps_platform_plugin,
     "cpu": cpu_platform_plugin,
+    "musa": musa_platform_plugin,
 }
 
 
@@ -128,6 +148,11 @@ def resolve_current_platform_cls_qualname() -> str:
     if platform_cls_qualname is not None:
         return platform_cls_qualname
 
+    # Fall back to MUSA
+    platform_cls_qualname = musa_platform_plugin()
+    if platform_cls_qualname is not None:
+        return platform_cls_qualname
+
     # Fall back to CPU as last resort
     platform_cls_qualname = cpu_platform_plugin()
     if platform_cls_qualname is not None:

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

@@ -46,6 +46,7 @@ class PlatformEnum(enum.Enum):
     TPU = enum.auto()
     CPU = enum.auto()
     MPS = enum.auto()
+    MUSA = enum.auto()
     OOT = enum.auto()
     UNSPECIFIED = enum.auto()
 
@@ -155,7 +156,7 @@ class Platform:
     @lru_cache(maxsize=1)
     def is_cuda_alike(self) -> bool:
         """Stateless version of :func:`torch.cuda.is_available`."""
-        return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)
+        return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM, PlatformEnum.MUSA)
 
     @lru_cache(maxsize=1)
     def is_mps(self) -> bool:

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

@@ -0,0 +1,312 @@
+"""
+This file is a platform abstraction for MThreads (MUSA) GPUs,
+adjusted to match the structure and interface of `cuda.py`.
+"""
+
+import os
+from collections.abc import Callable
+from functools import lru_cache, wraps
+from typing import Any, TypeVar
+
+import psutil
+import pymtml
+
+# isort: off
+import torch
+import torchada  # noqa: F401
+
+# isort: on
+from typing_extensions import ParamSpec
+
+from sglang.multimodal_gen.runtime.platforms.interface import (
+    AttentionBackendEnum,
+    DeviceCapability,
+    Platform,
+    PlatformEnum,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+
+def device_id_to_physical_device_id(device_id: int) -> int:
+    if "MUSA_VISIBLE_DEVICES" in os.environ:
+        device_ids = os.environ["MUSA_VISIBLE_DEVICES"].split(",")
+        if device_ids == [""]:
+            msg = (
+                "MUSA_VISIBLE_DEVICES is set to empty string, which means"
+                " GPU support is disabled. If you are using ray, please unset"
+                " the environment variable `MUSA_VISIBLE_DEVICES` inside the"
+                " worker/actor. "
+                "Check https://github.com/vllm-project/vllm/issues/8402 for"
+                " more information."
+            )
+            raise RuntimeError(msg)
+        physical_device_id = device_ids[device_id]
+        return int(physical_device_id)
+    else:
+        return device_id
+
+
+def with_mtml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
+    @wraps(fn)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+        pymtml.nvmlInit()
+        try:
+            return fn(*args, **kwargs)
+        finally:
+            pymtml.nvmlShutdown()
+
+    return wrapper
+
+
+class MusaPlatformBase(Platform):
+    _enum = PlatformEnum.MUSA
+    device_name: str = "musa"
+    device_type: str = "musa"
+    dispatch_key: str = "MUSA"
+    device_control_env_var: str = "MUSA_VISIBLE_DEVICES"
+
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        raise NotImplementedError
+
+    @classmethod
+    def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
+        if enforce_eager:
+            logger.warning(
+                "To see benefits of async output processing, enable MUSA "
+                "graph. Since, enforce-eager is enabled, async output "
+                "processor cannot be used"
+            )
+            return False
+        return True
+
+    @classmethod
+    def is_full_mtlink(cls, device_ids: list[int]) -> bool:
+        raise NotImplementedError
+
+    @classmethod
+    def log_warnings(cls) -> None:
+        pass
+
+    @classmethod
+    def get_current_memory_usage(
+        cls, device: torch.types.Device | None = None
+    ) -> float:
+        torch.cuda.reset_peak_memory_stats(device)
+        return float(torch.cuda.max_memory_allocated(device))
+
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int = 0,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group: Any = None,
+    ) -> float:
+        if empty_cache:
+            torch.cuda.empty_cache()
+
+        device_props = torch.cuda.get_device_properties(device_id)
+        if device_props.is_integrated:
+            free_gpu_memory = psutil.virtual_memory().available
+        else:
+            free_gpu_memory, _ = torch.cuda.mem_get_info(device_id)
+
+        if distributed:
+            import torch.distributed as dist
+
+            tensor = torch.tensor(free_gpu_memory, dtype=torch.float32, device="musa")
+            dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group)
+            free_gpu_memory = float(tensor.item())
+
+        return free_gpu_memory / (1 << 30)
+
+    @classmethod
+    def get_attn_backend_cls_str(
+        cls,
+        selected_backend: AttentionBackendEnum | None,
+        head_size: int,
+        dtype: torch.dtype,
+    ) -> str:
+        logger.info("Using Torch SDPA backend.")
+        return (
+            "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+        )
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # noqa
+
+
+# MTML utils
+# Note that MTML is not affected by `MUSA_VISIBLE_DEVICES`,
+# all the related functions work on real physical device ids.
+# the major benefit of using MTML is that it will not initialize MUSA
+class MtmlMusaPlatform(MusaPlatformBase):
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_mtml_context
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
+        try:
+            physical_device_id = device_id_to_physical_device_id(device_id)
+            handle = pymtml.nvmlDeviceGetHandleByIndex(physical_device_id)
+            major, minor = pymtml.nvmlDeviceGetCudaComputeCapability(handle)
+            return DeviceCapability(major=major, minor=minor)
+        except RuntimeError:
+            return None
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_mtml_context
+    def has_device_capability(
+        cls,
+        capability: tuple[int, int] | int,
+        device_id: int = 0,
+    ) -> bool:
+        try:
+            return bool(super().has_device_capability(capability, device_id))
+        except RuntimeError:
+            return False
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_mtml_context
+    def get_device_name(cls, device_id: int = 0) -> str:
+        physical_device_id = device_id_to_physical_device_id(device_id)
+        return cls._get_physical_device_name(physical_device_id)
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_mtml_context
+    def get_device_uuid(cls, device_id: int = 0) -> str:
+        physical_device_id = device_id_to_physical_device_id(device_id)
+        handle = pymtml.nvmlDeviceGetHandleByIndex(physical_device_id)
+        return str(pymtml.nvmlDeviceGetUUID(handle))
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_mtml_context
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        physical_device_id = device_id_to_physical_device_id(device_id)
+        handle = pymtml.nvmlDeviceGetHandleByIndex(physical_device_id)
+        return int(pymtml.nvmlDeviceGetMemoryInfo(handle).total)
+
+    @classmethod
+    @with_mtml_context
+    def is_full_mtlink(cls, physical_device_ids: list[int]) -> bool:
+        """
+        query if the set of gpus are fully connected by mtlink (1 hop)
+        """
+        handles = [pymtml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids]
+        for i, handle in enumerate(handles):
+            for j, peer_handle in enumerate(handles):
+                if i < j:
+                    try:
+                        p2p_status = pymtml.nvmlDeviceGetP2PStatus(
+                            handle,
+                            peer_handle,
+                            pymtml.NVML_P2P_CAPS_INDEX_NVLINK,
+                        )
+                        if p2p_status != pymtml.NVML_P2P_STATUS_OK:
+                            return False
+                    except pymtml.NVMLError:
+                        logger.exception(
+                            "MTLink detection failed. This is normal if"
+                            " your machine has no MTLink equipped."
+                        )
+                        return False
+        return True
+
+    @classmethod
+    def _get_physical_device_name(cls, device_id: int = 0) -> str:
+        handle = pymtml.nvmlDeviceGetHandleByIndex(device_id)
+        return str(pymtml.nvmlDeviceGetName(handle))
+
+    @classmethod
+    @with_mtml_context
+    def log_warnings(cls) -> None:
+        device_ids: int = pymtml.nvmlDeviceGetCount()
+        if device_ids > 1:
+            device_names = [cls._get_physical_device_name(i) for i in range(device_ids)]
+            if (
+                len(set(device_names)) > 1
+                and os.environ.get("MUSA_DEVICE_ORDER") != "PCI_BUS_ID"
+            ):
+                logger.warning(
+                    "Detected different devices in the system: %s. Please"
+                    " make sure to set `MUSA_DEVICE_ORDER=PCI_BUS_ID` to "
+                    "avoid unexpected behavior.",
+                    ", ".join(device_names),
+                )
+
+
+class NonMtmlMusaPlatform(MusaPlatformBase):
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
+        major, minor = torch.cuda.get_device_capability(device_id)
+        return DeviceCapability(major=major, minor=minor)
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return str(torch.cuda.get_device_name(device_id))
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        device_props = torch.cuda.get_device_properties(device_id)
+        return int(device_props.total_memory)
+
+    @classmethod
+    def is_full_mtlink(cls, physical_device_ids: list[int]) -> bool:
+        logger.error(
+            "MTLink detection not possible, as context support was"
+            " not found. Assuming no MTLink available."
+        )
+        return False
+
+
+# Autodetect either MTML-enabled or non-MTML platform
+# based on whether MTML is available.
+mtml_available = False
+
+if "MUSA_DISABLE_MTML" not in os.environ:
+    try:
+        try:
+            pymtml.nvmlInit()
+            mtml_available = True
+        except Exception:
+            mtml_available = False
+    finally:
+        if mtml_available:
+            pymtml.nvmlShutdown()
+
+MusaPlatform = MtmlMusaPlatform if mtml_available else NonMtmlMusaPlatform
+
+try:
+    from sphinx.ext.autodoc.mock import _MockModule
+
+    if not isinstance(pymtml, _MockModule):
+        MusaPlatform.log_warnings()
+except ModuleNotFoundError:
+    MusaPlatform.log_warnings()
+
+if __name__ == "__main__":
+    print(MusaPlatform.__name__)
+    print(MusaPlatform.get_device_name())
+    print(MusaPlatform.get_device_capability())
+    print(MusaPlatform.get_device_total_memory())
+    print(MusaPlatform.is_full_mtlink([0, 1, 2, 3, 4, 5, 6, 7]))