diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 6dfd5ec07..77c9df153 100644 --- a/docker/rocm.Dockerfile +++ b/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 diff --git a/docs/platforms/amd_gpu.md b/docs/platforms/amd_gpu.md index 068afa997..1759e9a73 100644 --- a/docs/platforms/amd_gpu.md +++ b/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) diff --git a/python/pyproject.toml b/python/pyproject.toml index f8ef96eca..32ab5aff8 100755 --- a/python/pyproject.toml +++ b/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" diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index d767b39a1..49e5ca687 100755 --- a/python/pyproject_other.toml +++ b/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 -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 = [ +diffusion_hip = [ "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" diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index 2f8f8d712..b0264adb3 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/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 - extra_args = {} if current_platform.is_mps() else dict(device_id=device_id) + # For MPS and MUSA, don't pass device_id as it doesn't support device indices + 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, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py index e4f4940e1..6244f7390 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py +++ b/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, ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index 7c98211c6..0ac092d9e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/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]: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index 533102690..a157b88d3 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/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]: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 73d478dcb..7d0e3847b 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 444e42472..3dcb8185b 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/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"] diff --git a/python/sglang/multimodal_gen/runtime/platforms/__init__.py b/python/sglang/multimodal_gen/runtime/platforms/__init__.py index ee515bb24..88d2f0155 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/__init__.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index 6c7297076..6f3a2ebe6 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/platforms/musa.py b/python/sglang/multimodal_gen/runtime/platforms/musa.py new file mode 100644 index 000000000..6a6de78e2 --- /dev/null +++ b/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]))