From db97f193b7212cdcc2b462416f82a80e3d8fe579 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Wed, 11 Mar 2026 04:11:07 +0800 Subject: [PATCH] [diffusion][llm] macOS support (#19549) Signed-off-by: Xiaodong Ye Co-authored-by: Mick --- docs/diffusion/environment_variables.md | 6 + docs/diffusion/index.md | 6 +- docs/diffusion/installation.md | 27 +- docs/platforms/apple_metal.md | 20 ++ python/pyproject_other.toml | 25 ++ python/sglang/__init__.py | 24 ++ python/sglang/_mps_stub.py | 256 +++++++++++++++ python/sglang/_triton_stub.py | 230 +++++++++++++ .../diffusion/triton/mps_fallback.py | 308 ++++++++++++++++++ .../jit_kernel/diffusion/triton/norm.py | 9 + .../diffusion/triton/rmsnorm_onepass.py | 8 + .../jit_kernel/diffusion/triton/rotary.py | 5 + .../diffusion/triton/scale_shift.py | 9 + python/sglang/multimodal_gen/README.md | 10 +- .../pipelines_core/stages/denoising.py | 5 +- python/sglang/srt/configs/device_config.py | 4 +- python/sglang/srt/environ.py | 3 + .../srt/layers/rotary_embedding/base.py | 3 + python/sglang/srt/managers/scheduler.py | 8 +- .../sglang/srt/mem_cache/memory_pool_host.py | 5 +- python/sglang/srt/server_args.py | 10 +- python/sglang/srt/utils/common.py | 14 +- 22 files changed, 984 insertions(+), 11 deletions(-) create mode 100644 docs/platforms/apple_metal.md create mode 100644 python/sglang/_mps_stub.py create mode 100644 python/sglang/_triton_stub.py create mode 100644 python/sglang/jit_kernel/diffusion/triton/mps_fallback.py diff --git a/docs/diffusion/environment_variables.md b/docs/diffusion/environment_variables.md index b02d7beb7..1dec2aafb 100644 --- a/docs/diffusion/environment_variables.md +++ b/docs/diffusion/environment_variables.md @@ -1,3 +1,9 @@ +## Apple MPS + +| Environment Variable | Default | Description | +|----------------------|---------|--------------------------------------------------------------| +| `SGLANG_USE_MLX` | not set | Set to `1` to enable MLX fused Metal kernels for norm ops on MPS | + ## Caching Acceleration These variables configure caching acceleration for Diffusion Transformer (DiT) models. diff --git a/docs/diffusion/index.md b/docs/diffusion/index.md index 6637ed785..6d786a0c5 100644 --- a/docs/diffusion/index.md +++ b/docs/diffusion/index.md @@ -7,7 +7,11 @@ SGLang Diffusion is an inference framework for accelerated image and video gener - **Broad Model Support**: Wan series, FastWan series, Hunyuan, Qwen-Image, Qwen-Image-Edit, Flux, Z-Image, GLM-Image, and more - **Fast Inference**: Optimized kernels, efficient scheduler loop, and Cache-DiT acceleration - **Ease of Use**: OpenAI-compatible API, CLI, and Python SDK -- **Multi-Platform**: NVIDIA GPUs (H100, H200, A100, B200, 4090), AMD GPUs (MI300X, MI325X) and Ascend NPU (A2, A3) +- **Multi-Platform**: + - NVIDIA GPUs (H100, H200, A100, B200, 4090) + - AMD GPUs (MI300X, MI325X) + - Ascend NPU (A2, A3) + - Apple Silicon (M-series via MPS) --- diff --git a/docs/diffusion/installation.md b/docs/diffusion/installation.md index 4cd62b10a..130ef03cd 100644 --- a/docs/diffusion/installation.md +++ b/docs/diffusion/installation.md @@ -69,7 +69,7 @@ For detailed ROCm system configuration and installation from source, see [AMD GP ## Platform-Specific: MUSA (Moore Threads GPUs) -For Moore Threads GPUs (MTGPU) with the MUSA software stack: +For Moore Threads GPUs (MTGPU) with the MUSA software stack, please follow the instructions below to install from source: ```bash # Clone the repository @@ -93,3 +93,28 @@ sglang generate --model-path black-forest-labs/FLUX.1-dev \ --prompt "A logo With Bold Large text: SGL Diffusion" \ --save-output ``` + +## Platform-Specific: Apple MPS + +For Apple MPS, please follow the instructions below to install from source: + +```bash +# Install ffmpeg +brew install ffmpeg + +# Install uv +brew install uv + +# Clone the repository +git clone https://github.com/sgl-project/sglang.git +cd sglang + +# Create and activate a virtual environment +uv venv -p 3.11 sglang-diffusion +source sglang-diffusion/bin/activate + +# Install the Python packages +uv pip install --upgrade pip +rm -f python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml +uv pip install -e "python[all_mps]" +``` diff --git a/docs/platforms/apple_metal.md b/docs/platforms/apple_metal.md new file mode 100644 index 000000000..3bcb3698b --- /dev/null +++ b/docs/platforms/apple_metal.md @@ -0,0 +1,20 @@ +# Apple Silicon with Metal + +This document describes how run SGLang on Apple Silicon using [Metal](https://developer.apple.com/metal/). If you encounter issues or have questions, please [open an issue](https://github.com/sgl-project/sglang/issues). + +## Install SGLang + +You can install SGLang using one of the methods below. + +### Install from Source + +```bash +# Use the default branch +git clone https://github.com/sgl-project/sglang.git +cd sglang + +# Install sglang python package +pip install --upgrade pip +rm -f python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml +uv pip install -e "python[all_mps]" +``` diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index 83cfbba04..78db9a496 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -136,6 +136,29 @@ diffusion_musa = [ "xatlas", ] +# https://docs.sglang.io/platforms/mps.md +srt_mps = [ + "sglang[runtime_common]", + "torch==2.9.1", + "torchao==0.9.0", + "torchaudio==2.9.1", + "torchvision", +] + +diffusion_mps = [ + "PyYAML==6.0.1", + "cloudpickle==3.1.2", + "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==2.1.0", + "cache-dit==1.2.3", + "addict==2.4.0", + "av==16.1.0", +] + test = [ "accelerate", "expecttest", @@ -152,10 +175,12 @@ test = [ all_hip = ["sglang[srt_hip]", "sglang[diffusion_hip]"] all_hpu = ["sglang[srt_hpu]"] all_musa = ["sglang[srt_musa]", "sglang[diffusion_musa]"] +all_mps = ["sglang[srt_mps]", "sglang[diffusion_mps]"] dev_hip = ["sglang[all_hip]", "sglang[test]"] dev_hpu = ["sglang[all_hpu]", "sglang[test]"] dev_musa = ["sglang[all_musa]", "sglang[test]"] +dev_mps = ["sglang[all_mps]", "sglang[test]"] [project.urls] "Homepage" = "https://github.com/sgl-project/sglang" diff --git a/python/sglang/__init__.py b/python/sglang/__init__.py index 509b145a9..7b6756b39 100644 --- a/python/sglang/__init__.py +++ b/python/sglang/__init__.py @@ -1,5 +1,29 @@ # SGLang public APIs +# Install stubs early for platforms where certain dependencies are unavailable +# (e.g. macOS/MPS has no triton, and torch.mps lacks Stream / set_device / +# get_device_properties). This must run before any downstream imports. +import sys as _sys + +if _sys.platform == "darwin": + try: + import torch as _torch + + if _torch.backends.mps.is_available(): + from sglang._triton_stub import install as _install_triton_stub + + _install_triton_stub() + del _install_triton_stub + + from sglang._mps_stub import install as _install_mps_stub + + _install_mps_stub() + del _install_mps_stub + del _torch + except ImportError: + pass +del _sys + # Frontend Language APIs from sglang.global_config import global_config from sglang.lang.api import ( diff --git a/python/sglang/_mps_stub.py b/python/sglang/_mps_stub.py new file mode 100644 index 000000000..aaaebe27e --- /dev/null +++ b/python/sglang/_mps_stub.py @@ -0,0 +1,256 @@ +"""Stub implementations for APIs missing from ``torch.mps``. + +``torch.mps`` lacks several APIs that ``torch.cuda`` provides (``Stream``, +``set_device``, ``get_device_properties``, …). Rather than scattering +``hasattr`` / ``getattr`` guards throughout the codebase, we monkey-patch +``torch.mps`` once at startup so that generic device-agnostic code paths +just work. +""" + +from __future__ import annotations + +import functools +from dataclasses import dataclass, field +from typing import Any + + +class Stream: + """Minimal stand-in for ``torch.cuda.Stream``. + + MPS does not expose user-visible streams. Every method is a no-op so + that code written for CUDA's multi-stream model still runs. + """ + + def __init__(self, device: Any = None, priority: int = 0) -> None: + pass + + def synchronize(self) -> None: + pass + + def wait_stream(self, stream: Any) -> None: + pass + + def wait_event(self, event: Any) -> None: + pass + + def record_event(self, event: Any = None) -> Any: + return None + + def query(self) -> bool: + return True + + # context-manager protocol (``with stream:``) + def __enter__(self) -> "Stream": + return self + + def __exit__(self, *args: Any) -> None: + pass + + +class Event: + """Minimal stand-in for ``torch.cuda.Event``.""" + + def __init__(self, enable_timing: bool = False) -> None: + pass + + def record(self, stream: Any = None) -> None: + pass + + def wait(self, stream: Any = None) -> None: + pass + + def query(self) -> bool: + return True + + def synchronize(self) -> None: + pass + + def elapsed_time(self, end_event: Any) -> float: + return 0.0 + + +_default_stream = Stream() + + +def current_stream(device: Any = None) -> Stream: + """Return the default (and only) MPS stream.""" + return _default_stream + + +def stream(s: Any) -> Stream: + """Return a context manager that is a no-op on MPS.""" + return s if s is not None else _default_stream + + +def set_device(device: Any) -> None: # noqa: ARG001 + """Set the current device. This is a no-op for MPS as it has exactly one device.""" + pass + + +def current_device() -> int: + """Return the index of the current MPS device (always 0).""" + return 0 + + +def device_count() -> int: + """Return the number of available MPS devices (always 1).""" + return 1 + + +@dataclass +class _MPSDeviceProperties: + """Mimics the object returned by ``torch.cuda.get_device_properties``.""" + + name: str = "Apple MPS" + total_memory: int = 0 # populated at install time + multi_processor_count: int = 0 + warp_size: int = 32 + is_integrated: bool = True + major: int = 0 + minor: int = 0 + # Extra attrs some callers inspect + _extra: dict = field(default_factory=dict) + + def __getattr__(self, name: str) -> Any: + # Return a safe default for any attribute we didn't anticipate + try: + return self._extra[name] + except KeyError: + return None + + +_cached_props: _MPSDeviceProperties | None = None + + +def get_device_properties(device: Any = 0) -> _MPSDeviceProperties: # noqa: ARG001 + """Return the properties of the MPS device. Results are cached after first call.""" + global _cached_props + if _cached_props is None: + import psutil + + _cached_props = _MPSDeviceProperties( + total_memory=psutil.virtual_memory().total, + ) + return _cached_props + + +class _MPSMemoryTracker: + """Tracks peak memory values on top of ``torch.mps`` current-value APIs. + + * ``memory_allocated`` → ``torch.mps.current_allocated_memory()`` + * ``memory_reserved`` → ``torch.mps.driver_allocated_memory()`` + * ``max_memory_*`` → high-water marks of the above + """ + + def __init__(self) -> None: + self._peak_allocated: int = 0 + self._peak_reserved: int = 0 + + def memory_allocated(self, device: Any = None) -> int: # noqa: ARG002 + import torch + + val = torch.mps.current_allocated_memory() + if val > self._peak_allocated: + self._peak_allocated = val + return val + + def memory_reserved(self, device: Any = None) -> int: # noqa: ARG002 + import torch + + val = torch.mps.driver_allocated_memory() + if val > self._peak_reserved: + self._peak_reserved = val + return val + + def max_memory_allocated(self, device: Any = None) -> int: # noqa: ARG002 + self.memory_allocated() + return self._peak_allocated + + def max_memory_reserved(self, device: Any = None) -> int: # noqa: ARG002 + self.memory_reserved() + return self._peak_reserved + + def reset_peak_memory_stats(self, device: Any = None) -> None: # noqa: ARG002 + import torch + + self._peak_allocated = torch.mps.current_allocated_memory() + self._peak_reserved = torch.mps.driver_allocated_memory() + + +_memory_tracker = _MPSMemoryTracker() + + +def _patch_non_blocking() -> None: + """Force ``non_blocking=False`` for copies targeting the MPS device. + + Unlike CUDA, MPS does not guarantee that a subsequent kernel on the same + "stream" will wait for an async host-to-device transfer to finish. Reading + the tensor before the transfer completes yields uninitialised (garbage) + data. Patching ``Tensor.to`` and ``Tensor.copy_`` centrally avoids having + to sprinkle ``non_blocking=not is_mps()`` at every call-site. + """ + import torch + + _original_to = torch.Tensor.to + + @functools.wraps(_original_to) + def _patched_to(self, *args, **kwargs): + if kwargs.get("non_blocking"): + # Detect target device from positional or keyword args + device = None + if args and isinstance(args[0], (str, torch.device)): + device = torch.device(args[0]) if isinstance(args[0], str) else args[0] + elif "device" in kwargs: + d = kwargs["device"] + device = torch.device(d) if isinstance(d, str) else d + if device is not None and device.type == "mps": + kwargs = {**kwargs, "non_blocking": False} + return _original_to(self, *args, **kwargs) + + torch.Tensor.to = _patched_to + + _original_copy_ = torch.Tensor.copy_ + + @functools.wraps(_original_copy_) + def _patched_copy_(self, src, non_blocking=False): + if non_blocking and self.device.type == "mps": + non_blocking = False + return _original_copy_(self, src, non_blocking=non_blocking) + + torch.Tensor.copy_ = _patched_copy_ + + +_installed = False + + +def install() -> None: + """Patch ``torch.mps`` with the stubs above. Safe to call multiple times.""" + global _installed + if _installed: + return + + import torch + + mps = torch.mps + # Only patch attributes that are actually missing + for name, obj in [ + ("Stream", Stream), + ("Event", Event), + ("current_stream", current_stream), + ("stream", stream), + ("set_device", set_device), + ("current_device", current_device), + ("device_count", device_count), + ("get_device_properties", get_device_properties), + ("reset_peak_memory_stats", _memory_tracker.reset_peak_memory_stats), + ("memory_allocated", _memory_tracker.memory_allocated), + ("memory_reserved", _memory_tracker.memory_reserved), + ("max_memory_allocated", _memory_tracker.max_memory_allocated), + ("max_memory_reserved", _memory_tracker.max_memory_reserved), + ]: + if not hasattr(mps, name): + setattr(mps, name, obj) + + _patch_non_blocking() + + _installed = True diff --git a/python/sglang/_triton_stub.py b/python/sglang/_triton_stub.py new file mode 100644 index 000000000..d78cdfb01 --- /dev/null +++ b/python/sglang/_triton_stub.py @@ -0,0 +1,230 @@ +""" +Mock triton module for platforms where triton is not available (e.g., macOS/MPS). + +This module provides stub implementations of triton APIs so that modules which +import triton at the top level can be loaded without error. The actual triton +kernels are never executed on these platforms – alternative backends (e.g. SDPA +for MPS) are used instead. + +Usage – call ``install()`` **before** any ``import triton`` in the process: + + from sglang._triton_stub import install + install() +""" + +import sys +import types + + +class _StubBase: + """A base class that any mock attribute can safely be subclassed from. + + Used when external code does ``class Foo(triton.runtime.KernelInterface):``. + """ + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + + +class _MockModule(types.ModuleType): + """A module whose every attribute is itself a ``_MockModule``. + + When called (e.g. ``@triton.jit``), it acts as a pass-through decorator so + that kernel *definitions* are syntactically valid even though they will never + be compiled. + """ + + def __init__(self, name: str): + super().__init__(name) + self.__path__: list[str] = [] # make it look like a package + self.__package__ = name + self.__file__ = __file__ + self._children: dict[str, object] = {} + # Set __spec__ so that importlib.util.find_spec() works on cached modules + import importlib + + self.__spec__ = importlib.machinery.ModuleSpec(name, None, is_package=True) + + def __getattr__(self, name: str): + """Handle attribute access by creating and returning a child _MockModule.""" + if name.startswith("__") and name.endswith("__"): + raise AttributeError(name) + full = f"{self.__name__}.{name}" + if full in sys.modules: + return sys.modules[full] + # If the name looks like a class (CamelCase / uppercase), return a + # stub class that can be used as a base class for inheritance. + if name[0:1].isupper(): + stub_cls = type(name, (_StubBase,), {"__module__": self.__name__}) + self._children[name] = stub_cls + return stub_cls + child = _MockModule(full) + sys.modules[full] = child + self._children[name] = child + return child + + def __call__(self, *args, **kwargs): + # Direct decorator usage: @triton.jit (receives the function) + if len(args) == 1 and callable(args[0]) and not kwargs: + return args[0] + + # Parameterised decorator: @triton.jit(...) → returns a decorator + def _decorator(fn): + return fn + + return _decorator + + def __instancecheck__(self, instance): + """Return False for all instance checks against the mock.""" + return False + + def __contains__(self, item): + """Return False for all membership checks.""" + return False + + def __iter__(self): + return iter([]) + + def __len__(self): + return 0 + + def __bool__(self): + return False + + def __repr__(self): + return f"" + + +def _cdiv(a: int, b: int) -> int: + """Ceiling division – mirrors ``triton.cdiv``.""" + return -(a // -b) + + +def _next_power_of_2(n: int) -> int: + """Mirrors ``triton.next_power_of_2``.""" + return 1 << (n - 1).bit_length() if n > 0 else 1 + + +class _Config: + """Minimal stand-in for ``triton.Config`` used in ``@triton.autotune``.""" + + def __init__(self, kwargs=None, num_warps=4, num_stages=2, **extra): + self.kwargs = kwargs or {} + self.num_warps = num_warps + self.num_stages = num_stages + + +class _TritonFinder: + """A meta-path finder that intercepts all ``import triton.*`` statements. + + When Python encounters ``import triton.backends.compiler``, it walks the + dotted path and tries to import each component. Our mock module's + ``__getattr__`` handles *attribute* access, but the import machinery uses + ``importlib`` finders, not attribute access, for sub-module resolution. + This finder bridges that gap by creating ``_MockModule`` instances for any + ``triton.*`` sub-module that isn't already in ``sys.modules``. + """ + + def find_module(self, fullname, path=None): + if fullname == "triton" or fullname.startswith("triton."): + return self + return None + + def load_module(self, fullname): + if fullname in sys.modules: + return sys.modules[fullname] + mod = _MockModule(fullname) + sys.modules[fullname] = mod + # Wire up the parent relationship + parts = fullname.rsplit(".", 1) + if len(parts) == 2: + parent_name, child_name = parts + parent = sys.modules.get(parent_name) + if parent is not None: + setattr(parent, child_name, mod) + return mod + + +def _make_mock(name: str) -> _MockModule: + """Create a ``_MockModule`` and register it in ``sys.modules``.""" + mod = _MockModule(name) + sys.modules[name] = mod + return mod + + +def install() -> None: + """Register a mock ``triton`` package in *sys.modules*. + + This is a no-op if a real ``triton`` is already importable. + """ + if "triton" in sys.modules: + return + # Check whether a real triton exists before installing the stub. + import importlib.util + + if importlib.util.find_spec("triton") is not None: + return + + # Register the meta-path finder FIRST so that any ``import triton.X`` + # during the rest of install() (or later) is handled. + sys.meta_path.insert(0, _TritonFinder()) + + triton = _make_mock("triton") + triton.__version__ = "3.0.0" + triton.cdiv = _cdiv + triton.next_power_of_2 = _next_power_of_2 + triton.Config = _Config + + # triton.language (commonly imported as ``tl``) + tl = _make_mock("triton.language") + + class _constexpr: + """Stand-in for ``tl.constexpr`` – works as both annotation and value wrapper.""" + + def __init__(self, value=None): + self.value = value + + def __repr__(self): + return f"constexpr({self.value!r})" + + tl.constexpr = _constexpr + triton.language = tl + + # triton.language.extra.libdevice + extra = _make_mock("triton.language.extra") + tl.extra = extra + libdevice = _make_mock("triton.language.extra.libdevice") + extra.libdevice = libdevice + + # triton.runtime.jit (JITFunction used in isinstance checks) + runtime = _make_mock("triton.runtime") + triton.runtime = runtime + jit_mod = _make_mock("triton.runtime.jit") + + class _JITFunction: + """Dummy so ``isinstance(fn, triton.runtime.jit.JITFunction)`` works.""" + + pass + + jit_mod.JITFunction = _JITFunction + runtime.jit = jit_mod + + # triton.runtime.driver (used by fla/utils.py) + driver = _make_mock("triton.runtime.driver") + runtime.driver = driver + + # triton.testing + testing = _make_mock("triton.testing") + triton.testing = testing + + # triton.tools / triton.tools.tensor_descriptor + tools = _make_mock("triton.tools") + triton.tools = tools + td = _make_mock("triton.tools.tensor_descriptor") + tools.tensor_descriptor = td + + # triton.backends / triton.backends.compiler (used by torch._inductor) + backends = _make_mock("triton.backends") + triton.backends = backends + compiler = _make_mock("triton.backends.compiler") + backends.compiler = compiler diff --git a/python/sglang/jit_kernel/diffusion/triton/mps_fallback.py b/python/sglang/jit_kernel/diffusion/triton/mps_fallback.py new file mode 100644 index 000000000..4ed5b9b36 --- /dev/null +++ b/python/sglang/jit_kernel/diffusion/triton/mps_fallback.py @@ -0,0 +1,308 @@ +"""MPS (Apple Silicon) fallbacks for Triton diffusion kernels. + +Triton is not available on macOS / Metal, so these pure-PyTorch (and +optionally MLX-accelerated) implementations replace the Triton kernels +at import time when ``current_platform.is_mps()`` is True. + +MLX acceleration (opt-in via ``SGLANG_USE_MLX=1``): + Norm ops use ``mx.fast.rms_norm`` / ``mx.fast.layer_norm`` — single fused + Metal kernels that are 1.4x–2.9x faster than the multi-step PyTorch MPS + decomposition for medium-to-large tensors. +""" + +from typing import Optional + +import torch +from torch import Tensor + +from sglang.srt.environ import envs + +# MLX acceleration – opt-in via SGLANG_USE_MLX=1 +_MLX_AVAILABLE = False +try: + import mlx.core as mx + + _MLX_AVAILABLE = True +except ImportError: + pass + +_USE_MLX = envs.SGLANG_USE_MLX.get() and _MLX_AVAILABLE + +# Dtype mapping for torch <-> MLX tensor bridge +_TORCH_TO_MLX_DTYPE = ( + { + torch.float32: mx.float32, + torch.float16: mx.float16, + torch.bfloat16: mx.bfloat16, + } + if _MLX_AVAILABLE + else {} +) + +_MLX_TO_TORCH_DTYPE = {v: k for k, v in _TORCH_TO_MLX_DTYPE.items()} + + +def _torch_to_mlx(tensor: torch.Tensor) -> "mx.array": + """Convert a PyTorch tensor to an MLX array (via numpy on CPU).""" + t = tensor.cpu().detach() + if t.dtype == torch.bfloat16: + return mx.array(t.float().numpy(), dtype=mx.bfloat16) + return mx.array(t.numpy()) + + +def _mlx_to_torch(array: "mx.array", device: torch.device) -> torch.Tensor: + """Convert an MLX array to a PyTorch tensor (zero-copy via memoryview).""" + torch_dtype = _MLX_TO_TORCH_DTYPE.get(array.dtype, torch.float32) + array = mx.contiguous(array) + mx.eval(array) + tensor = torch.frombuffer(memoryview(array), dtype=torch_dtype).reshape(array.shape) + if device.type == "mps": + tensor = tensor.to(device) + return tensor + + +def fuse_scale_shift_kernel_native( + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + scale_constant: float = 1.0, + block_l: int = 128, + block_c: int = 128, +): + """Native fallback for fuse_scale_shift_kernel with scale_constant support.""" + B, L, C = x.shape + + def _expand(t: torch.Tensor) -> torch.Tensor: + if t.dim() == 4: + # [B, F, 1, C] -> [B, L, C] + num_frames = t.shape[1] + frame_seqlen = L // num_frames + return ( + t.squeeze(2) + .unsqueeze(2) + .expand(-1, -1, frame_seqlen, -1) + .reshape(B, L, C) + ) + elif t.dim() == 2: + # [B, C] -> [B, 1, C] + return t.unsqueeze(1) + return t + + scale = _expand(scale) + shift = _expand(shift) + + return x * (scale_constant + scale) + shift + + +def fuse_scale_shift_gate_select01_kernel_native( + x: torch.Tensor, + scale0: torch.Tensor, + shift0: torch.Tensor, + gate0: torch.Tensor, + scale1: torch.Tensor, + shift1: torch.Tensor, + gate1: torch.Tensor, + index: torch.Tensor, + block_l: int = 128, + block_c: int = 128, +): + """Native fallback for fuse_scale_shift_gate_select01_kernel.""" + idx = index.unsqueeze(-1).bool() + scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1)) + shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1)) + gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1)) + y = x * (1 + scale) + shift + return y, gate + + +def apply_rotary_embedding_native( + x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False +) -> torch.Tensor: + """Native fallback for rotary embedding (shared with NPU implementation).""" + cos = cos.unsqueeze(-2).to(x.dtype) + sin = sin.unsqueeze(-2).to(x.dtype) + x1 = x[..., ::2] + x2 = x[..., 1::2] + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + return torch.stack((o1, o2), dim=-1).flatten(-2) + + +def norm_infer_native( + x: Tensor, + weight: Optional[Tensor], + bias: Optional[Tensor], + eps: float, + is_rms_norm: bool = False, + out: Optional[Tensor] = None, +) -> Tensor: + """Native fallback for norm_infer (layer norm / rms norm inference).""" + orig_dtype = x.dtype + x = x.contiguous().float() + if is_rms_norm: + variance = x.pow(2).mean(dim=-1, keepdim=True) + x_hat = x * torch.rsqrt(variance + eps) + else: + mean = x.mean(dim=-1, keepdim=True) + variance = (x - mean).pow(2).mean(dim=-1, keepdim=True) + x_hat = (x - mean) * torch.rsqrt(variance + eps) + if weight is not None: + x_hat = x_hat * weight.float() + if bias is not None: + x_hat = x_hat + bias.float() + result = x_hat.to(orig_dtype) + if out is not None: + out.copy_(result) + return out + return result + + +def triton_one_pass_rms_norm_native( + x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6 +) -> torch.Tensor: + """Native fallback for triton_one_pass_rms_norm.""" + shape = x.shape + orig_dtype = x.dtype + x = x.contiguous().float() + variance = x.pow(2).mean(dim=-1, keepdim=True) + x_hat = x * torch.rsqrt(variance + eps) + return (x_hat * w.float()).to(orig_dtype).view(shape) + + +def rms_norm_fn_native( + x, + weight, + bias, + residual=None, + x1=None, + weight1=None, + bias1=None, + eps=1e-6, + dropout_p=0.0, + rowscale=None, + prenorm=False, + residual_in_fp32=False, + zero_centered_weight=False, + return_dropout_mask=False, + out_dtype=None, + out=None, + residual_out=None, +): + """Native fallback for rms_norm_fn (inference only, no dropout/x1 support).""" + x_shape_og = x.shape + orig_dtype = x.dtype + x = x.reshape(-1, x.shape[-1]).float() + if residual is not None: + residual = residual.reshape(-1, residual.shape[-1]).float() + x = x + residual + residual_out_val = x.to(torch.float32 if residual_in_fp32 else orig_dtype) + else: + residual_out_val = None + variance = x.pow(2).mean(dim=-1, keepdim=True) + x_hat = x * torch.rsqrt(variance + eps) + if weight is not None: + w = weight.float() + if zero_centered_weight: + w = w + 1.0 + x_hat = x_hat * w + if bias is not None: + x_hat = x_hat + bias.float() + final_dtype = out_dtype if out_dtype is not None else orig_dtype + y = x_hat.to(final_dtype).reshape(x_shape_og) + if residual is not None and residual_out_val is not None: + return y, residual_out_val.reshape(x_shape_og) + return y + + +# MLX-accelerated norm ops (1.4x–2.9x faster than torch native on MPS) +# Uses mx.fast.rms_norm / mx.fast.layer_norm — single fused Metal kernels +# instead of 7+ separate PyTorch MPS kernel launches. + +if _USE_MLX: + + def norm_infer_native( # noqa: F811 + x: Tensor, + weight: Optional[Tensor], + bias: Optional[Tensor], + eps: float, + is_rms_norm: bool = False, + out: Optional[Tensor] = None, + ) -> Tensor: + """MLX-accelerated norm_infer (layer norm / rms norm inference).""" + device = x.device + orig_dtype = x.dtype + x_mx = _torch_to_mlx(x) + if is_rms_norm: + w_mx = ( + _torch_to_mlx(weight) if weight is not None else mx.ones(x_mx.shape[-1]) + ) + result_mx = mx.fast.rms_norm(x_mx, w_mx, eps) + else: + w_mx = _torch_to_mlx(weight) if weight is not None else None + b_mx = _torch_to_mlx(bias) if bias is not None else None + result_mx = mx.fast.layer_norm(x_mx, w_mx, b_mx, eps) + result = _mlx_to_torch(result_mx, device).to(orig_dtype) + if out is not None: + out.copy_(result) + return out + return result + + def triton_one_pass_rms_norm_native( # noqa: F811 + x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6 + ) -> torch.Tensor: + """MLX-accelerated triton_one_pass_rms_norm.""" + shape = x.shape + device = x.device + orig_dtype = x.dtype + x_mx = _torch_to_mlx(x.reshape(-1, x.shape[-1])) + w_mx = _torch_to_mlx(w) + result_mx = mx.fast.rms_norm(x_mx, w_mx, eps) + return _mlx_to_torch(result_mx, device).to(orig_dtype).view(shape) + + def rms_norm_fn_native( # noqa: F811 + x, + weight, + bias, + residual=None, + x1=None, + weight1=None, + bias1=None, + eps=1e-6, + dropout_p=0.0, + rowscale=None, + prenorm=False, + residual_in_fp32=False, + zero_centered_weight=False, + return_dropout_mask=False, + out_dtype=None, + out=None, + residual_out=None, + ): + """MLX-accelerated rms_norm_fn (inference only, no dropout/x1 support).""" + x_shape_og = x.shape + device = x.device + orig_dtype = x.dtype + x_flat = x.reshape(-1, x.shape[-1]) + if residual is not None: + residual = residual.reshape(-1, residual.shape[-1]).float() + x_flat = x_flat.float() + residual + residual_out_val = x_flat.to( + torch.float32 if residual_in_fp32 else orig_dtype + ) + else: + residual_out_val = None + if weight is not None and zero_centered_weight: + w = weight.float() + 1.0 + else: + w = weight + x_mx = _torch_to_mlx(x_flat) + w_mx = _torch_to_mlx(w) if w is not None else mx.ones(x_mx.shape[-1]) + result_mx = mx.fast.rms_norm(x_mx, w_mx, eps) + x_hat = _mlx_to_torch(result_mx, device) + if bias is not None: + x_hat = x_hat + bias.to(x_hat.device, x_hat.dtype) + final_dtype = out_dtype if out_dtype is not None else orig_dtype + y = x_hat.to(final_dtype).reshape(x_shape_og) + if residual is not None and residual_out_val is not None: + return y, residual_out_val.reshape(x_shape_og) + return y diff --git a/python/sglang/jit_kernel/diffusion/triton/norm.py b/python/sglang/jit_kernel/diffusion/triton/norm.py index 17a5bb1ca..53b1a8255 100644 --- a/python/sglang/jit_kernel/diffusion/triton/norm.py +++ b/python/sglang/jit_kernel/diffusion/triton/norm.py @@ -618,3 +618,12 @@ def rms_norm_fn( out, residual_out, ) + + +from sglang.multimodal_gen.runtime.platforms import current_platform + +if current_platform.is_mps(): + from .mps_fallback import norm_infer_native, rms_norm_fn_native + + norm_infer = norm_infer_native + rms_norm_fn = rms_norm_fn_native diff --git a/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py b/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py index df6f7bf21..c6992ac72 100644 --- a/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py +++ b/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py @@ -56,3 +56,11 @@ def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6 BLOCK_SIZE_SEQ=BLOCK_SIZE_SEQ, ) return y + + +from sglang.multimodal_gen.runtime.platforms import current_platform + +if current_platform.is_mps(): + from .mps_fallback import triton_one_pass_rms_norm_native + + triton_one_pass_rms_norm = triton_one_pass_rms_norm_native diff --git a/python/sglang/jit_kernel/diffusion/triton/rotary.py b/python/sglang/jit_kernel/diffusion/triton/rotary.py index 067a6ceb4..011bf7ad0 100644 --- a/python/sglang/jit_kernel/diffusion/triton/rotary.py +++ b/python/sglang/jit_kernel/diffusion/triton/rotary.py @@ -111,3 +111,8 @@ if current_platform.is_npu(): from .npu_fallback import apply_rotary_embedding_native apply_rotary_embedding = apply_rotary_embedding_native + +if current_platform.is_mps(): + from .mps_fallback import apply_rotary_embedding_native + + apply_rotary_embedding = apply_rotary_embedding_native diff --git a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py index b3a27b64c..8be4c11bd 100644 --- a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py +++ b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py @@ -411,3 +411,12 @@ if current_platform.is_npu(): from .npu_fallback import fuse_scale_shift_native fuse_scale_shift_kernel = fuse_scale_shift_native + +if current_platform.is_mps(): + from .mps_fallback import ( + fuse_scale_shift_gate_select01_kernel_native, + fuse_scale_shift_kernel_native, + ) + + fuse_scale_shift_kernel = fuse_scale_shift_kernel_native + fuse_scale_shift_gate_select01_kernel = fuse_scale_shift_gate_select01_kernel_native diff --git a/python/sglang/multimodal_gen/README.md b/python/sglang/multimodal_gen/README.md index 313567fde..b7f5892fa 100644 --- a/python/sglang/multimodal_gen/README.md +++ b/python/sglang/multimodal_gen/README.md @@ -12,7 +12,11 @@ SGLang Diffusion has the following features: - Broad model support: Wan series, FastWan series, Hunyuan, Qwen-Image, Qwen-Image-Edit, Flux, Z-Image, GLM-Image - Fast inference speed: enpowered by highly optimized kernel from sgl-kernel and efficient scheduler loop - Ease of use: OpenAI-compatible api, CLI, and python sdk support - - Multi-platform support: NVIDIA GPUs (H100, H200, A100, B200, 4090) and AMD GPUs (MI300X, MI325X) + - Multi-platform support: + - NVIDIA GPUs (H100, H200, A100, B200, 4090) + - AMD GPUs (MI300X, MI325X) + - Ascend NPU (A2, A3) + - Apple Silicon (M-series via MPS) ### AMD/ROCm Support @@ -22,6 +26,10 @@ SGLang Diffusion supports AMD Instinct GPUs through ROCm. On AMD platforms, we u SGLang Diffusion supports Moore Threads GPUs (MTGPU) through the MUSA software stack. On MUSA platforms, we use the Torch SDPA backend for attention. See the [installation guide](https://github.com/sgl-project/sglang/tree/main/docs/diffusion/installation.md) for setup instructions. +### Apple MPS Support + +SGLang Diffusion supports Apple Silicon (M-series) via the MPS backend. Since Triton is Linux-only, all Triton kernels are replaced with PyTorch-native fallbacks on MPS. Norm operations can be optionally accelerated with MLX fused Metal kernels (`SGLANG_USE_MLX=1`). See the [installation guide](https://github.com/sgl-project/sglang/tree/main/docs/diffusion/installation.md) for setup instructions. + ## Getting Started ```bash 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 f06343174..27da4c92b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -755,6 +755,9 @@ class DenoisingStage(PipelineStage): ): self.save_sta_search_results(batch) + # Capture references before potential deletion on MPS + dits = list(filter(None, [self.transformer, self.transformer_2])) + # deallocate transformer if on mps pipeline = self.pipeline() if self.pipeline else None if torch.backends.mps.is_available() and not is_warmup: @@ -772,7 +775,7 @@ class DenoisingStage(PipelineStage): ) # reset offload managers with prefetching first layer for next forward - for dit in filter(None, [self.transformer, self.transformer_2]): + for dit in dits: if isinstance(dit, OffloadableDiTMixin): # release all DiT weights to avoid peak VRAM usage, which may increasing the latency for next req # TODO: should be make this an option? diff --git a/python/sglang/srt/configs/device_config.py b/python/sglang/srt/configs/device_config.py index 8ddcfd108..9836f935c 100644 --- a/python/sglang/srt/configs/device_config.py +++ b/python/sglang/srt/configs/device_config.py @@ -5,13 +5,15 @@ import torch logger = logging.getLogger(__name__) +SUPPORTED_DEVICES = ["cuda", "xpu", "hpu", "cpu", "npu", "musa", "mps"] + class DeviceConfig: device: Optional[torch.device] gpu_id: Optional[int] def __init__(self, device: str = "cuda", gpu_id: int = -1) -> None: - if device in ["cuda", "xpu", "hpu", "cpu", "npu", "musa"]: + if device in SUPPORTED_DEVICES: self.device_type = device else: raise RuntimeError(f"Not supported device type: {device}") diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d1388d159..1433977ef 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -307,6 +307,9 @@ class Envs: SGLANG_ROCM_FUSED_DECODE_MLA = EnvBool(False) SGLANG_ROCM_DISABLE_LINEARQUANT = EnvBool(False) + # MPS (Apple Silicon) + SGLANG_USE_MLX = EnvBool(False) + # NPU SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT = EnvBool(False) SGLANG_NPU_USE_MULTI_STREAM = EnvBool(False) diff --git a/python/sglang/srt/layers/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py index e98007325..943fe8558 100644 --- a/python/sglang/srt/layers/rotary_embedding/base.py +++ b/python/sglang/srt/layers/rotary_embedding/base.py @@ -15,6 +15,7 @@ from sglang.srt.utils import ( is_cpu, is_cuda, is_hip, + is_mps, is_musa, is_npu, is_xpu, @@ -31,6 +32,7 @@ _is_cpu_amx_available = cpu_has_amx_support() _is_cpu = is_cpu() _is_xpu = is_xpu() _is_musa = is_musa() +_is_mps = is_mps() if _is_cuda: from sglang.jit_kernel.rope import apply_rope_with_cos_sin_cache_inplace @@ -70,6 +72,7 @@ class RotaryEmbedding(MultiPlatformOp): and not (_is_xpu) and not (_is_npu) and not (_is_musa) + and not (_is_mps) ): # rotary_embedding from sglang.jit_kernel.rope and vllm._custom_ops has the same implementation. # TODO: Test on different devices and remove this conditional. diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ded080a28..10615e072 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -20,6 +20,7 @@ import signal import sys import time from collections import deque +from contextlib import nullcontext from dataclasses import dataclass from http import HTTPStatus from typing import Any, Deque, Dict, List, Optional, Tuple, Union @@ -30,7 +31,6 @@ import torch import torch.distributed import zmq from torch.cuda import Stream as CudaStream -from torch.cuda import StreamContext as CudaStreamContext from torch.distributed import barrier from sglang.jit_kernel.ngram_embedding import update_token_table @@ -202,6 +202,7 @@ from sglang.srt.utils import ( get_int_env_var, get_numa_node, get_zmq_socket, + is_mps, kill_itself_when_parent_died, numa_bind_to_node, point_to_point_pyobj, @@ -219,6 +220,11 @@ from sglang.srt.utils.hf_transformers_utils import ( from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter from sglang.utils import TypeBasedDispatcher, get_exception_traceback +if is_mps(): + CudaStreamContext = nullcontext +else: + from torch.cuda import StreamContext as CudaStreamContext + logger = logging.getLogger(__name__) # Test retract decode for debugging purposes diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 276070b51..d71fd21cb 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -23,12 +23,13 @@ from sglang.srt.mem_cache.memory_pool import ( MLATokenToKVPool, NSATokenToKVPool, ) -from sglang.srt.utils import is_cuda, is_npu, is_xpu +from sglang.srt.utils import is_cuda, is_mps, is_npu, is_xpu _is_cuda = is_cuda() _is_npu = is_npu() _is_xpu = is_xpu() -if not (_is_npu or _is_xpu): +_is_mps = is_mps() +if not (_is_npu or _is_xpu or _is_mps): from sgl_kernel.kvcacheio import ( transfer_kv_all_layer, transfer_kv_all_layer_direct_lf_pf, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6bfa8ecb3..adcd1b521 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -51,6 +51,8 @@ from sglang.srt.utils.common import ( is_flashinfer_available, is_hip, is_hopper_with_cuda_12_3, + is_mps, + is_musa, is_no_spec_infer_or_topk_one, is_npu, is_remote_url, @@ -1045,8 +1047,8 @@ class ServerArgs: # 5. Pipeline parallelism if self.pp_size > 1: self.disable_piecewise_cuda_graph = True - # 6. Non-CUDA hardware (AMD, NPU, CPU, etc.) - if is_hip() or is_npu() or is_cpu(): + # 6. Non-CUDA hardware (AMD, NPU, CPU, MPS, MUSA, etc.) + if is_hip() or is_npu() or is_cpu() or is_mps() or is_musa(): self.disable_piecewise_cuda_graph = True # 7. MoE A2A backend if self.moe_a2a_backend != "none": @@ -2097,6 +2099,8 @@ class ServerArgs: return "trtllm_mha" elif is_hip(): return "aiter" + elif is_mps(): + return "torch_native" else: return "flashinfer" if is_flashinfer_available() else "triton" else: @@ -2112,6 +2116,8 @@ class ServerArgs: return "aiter" else: return "triton" + elif is_mps(): + return "torch_native" else: return "triton" diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 5bd48478e..ca2d72a46 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -194,6 +194,11 @@ def is_musa() -> bool: return hasattr(torch.version, "musa") and torch.version.musa is not None +@lru_cache(maxsize=1) +def is_mps() -> bool: + return torch.backends.mps.is_available() + + def is_float4_e2m1fn_x2(dtype) -> bool: """Check if dtype is float4_e2m1fn_x2 and CUDA is available.""" target_dtype = getattr(torch, "float4_e2m1fn_x2", None) @@ -596,6 +601,8 @@ def get_available_gpu_memory( # memory metric instead. free_gpu_memory = psutil.virtual_memory().available free_gpu_memory, total_gpu_memory = torch.musa.mem_get_info() + elif device == "mps": + free_gpu_memory = psutil.virtual_memory().available if distributed: tensor = torch.tensor(free_gpu_memory, dtype=torch.float32) @@ -2099,7 +2106,12 @@ def get_device(device_id: Optional[int] = None) -> str: return "musa" return "musa:{}".format(device_id) - raise RuntimeError("No accelerator (CUDA, XPU, HPU, NPU, MUSA) is available.") + if is_mps(): + if device_id is None: + return "mps" + return "mps:{}".format(device_id) + + raise RuntimeError("No accelerator (CUDA, XPU, HPU, NPU, MUSA, MPS) is available.") @lru_cache(maxsize=1)