[diffusion][llm] macOS support (#19549)

Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

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"

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 (

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

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"<triton-stub {self.__name__!r}>"
+
+
+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

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

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

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

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

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

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

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?

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}")

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)

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.

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

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,

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"
 

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)