chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
R0CKSTAR
2026-01-16 17:01:09 +08:00
committed by GitHub
parent d9ed80b9f1
commit a1dd3d48ac
13 changed files with 425 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
docker/rocm.Dockerfile
View File

@@ -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

2
docs/platforms/amd_gpu.md
View File

@@ -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)

2
python/pyproject.toml
View File

@@ -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"

70
python/pyproject_other.toml
View File

@@ -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"

8
python/sglang/multimodal_gen/runtime/distributed/parallel_state.py
View File

@@ -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,

12
python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py
View File

@@ -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,
)

6
python/sglang/multimodal_gen/runtime/models/dits/flux.py
View File

@@ -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]:

6
python/sglang/multimodal_gen/runtime/models/dits/flux_2.py
View File

@@ -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]:

2
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
View File

@@ -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

6
python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
View File

@@ -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"]

25
python/sglang/multimodal_gen/runtime/platforms/__init__.py
View File

@@ -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:

3
python/sglang/multimodal_gen/runtime/platforms/interface.py
View File

@@ -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:

312
python/sglang/multimodal_gen/runtime/platforms/musa.py Normal file
View File

@@ -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]))