sgl-project--sglang/python/sglang/srt/utils/tensor_bridge.py

# Copied and adapted from: https://github.com/vllm-project/vllm-metal
# SPDX-License-Identifier: Apache-2.0
"""Tensor bridge between MLX and PyTorch.

Provides zero-copy conversion when possible using Apple Silicon's unified memory.
"""

from __future__ import annotations

import logging
from functools import lru_cache
from typing import TYPE_CHECKING, Literal

import torch

from sglang.srt.environ import envs

if TYPE_CHECKING:
    import mlx.core as mx

logger = logging.getLogger(__name__)

_MLX_AVAILABLE: bool = False
try:
    import mlx.core as mx  # noqa: F811

    _MLX_AVAILABLE = True
except ImportError:
    pass


def is_mlx_available() -> bool:
    """Return True when the ``mlx`` package can be imported."""
    return _MLX_AVAILABLE


@lru_cache(maxsize=1)
def use_mlx() -> bool:
    """Return True when the user opted-in via ``SGLANG_USE_MLX=1`` **and** MLX is importable."""
    return bool(envs.SGLANG_USE_MLX.get()) and _MLX_AVAILABLE


# MPS has a 4GB (2^32 bytes) limit for MPSTemporaryNDArray allocations.
# Metal may allocate multiple temporary buffers internally, so we use a
# conservative threshold of 1GB to avoid hitting the limit.
# See: https://github.com/anthropics/vllm-metal/issues/43
_MPS_SAFE_SIZE_BYTES = 1 << 30  # 1GB

# MLX to PyTorch dtype mapping
# TODO(perf): float64 is CPU-only in MLX (see ml-explore/mlx#1843).
# When the target device is GPU/MPS we should auto-downcast float64 → float32
# to avoid a runtime error; when the target is CPU we can keep float64.
# For now float64 is omitted from the mapping so it hits the ValueError
# fallback in mlx_to_torch().
MLX_TO_TORCH_DTYPE = (
    {
        mx.float32: torch.float32,
        mx.float16: torch.float16,
        mx.bfloat16: torch.bfloat16,
        mx.int32: torch.int32,
        mx.int64: torch.int64,
        mx.int16: torch.int16,
        mx.int8: torch.int8,
        mx.uint8: torch.uint8,
        mx.bool_: torch.bool,
    }
    if _MLX_AVAILABLE
    else {}
)

# PyTorch to MLX dtype mapping
TORCH_TO_MLX_DTYPE = {v: k for k, v in MLX_TO_TORCH_DTYPE.items()}


def get_torch_device() -> torch.device:
    """Get the PyTorch device for Metal/MPS.

    Returns:
        torch.device for MPS if available, else CPU
    """
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def _get_tensor_size_bytes(array: mx.array) -> int:
    """Calculate the size of an MLX array in bytes.

    Args:
        array: MLX array

    Returns:
        Size in bytes
    """
    return array.size * array.dtype.size


def _is_safe_for_mps(array: mx.array) -> bool:
    """Check if an array is safe to transfer to MPS without hitting size limits.

    MPS has a 4GB limit for MPSTemporaryNDArray, but Metal may allocate
    multiple temporary buffers internally. We use a conservative threshold.

    Args:
        array: MLX array to check

    Returns:
        True if safe to transfer to MPS, False if should stay on CPU
    """
    return _get_tensor_size_bytes(array) < _MPS_SAFE_SIZE_BYTES


def torch_to_mlx(tensor: torch.Tensor) -> mx.array:
    """Convert PyTorch tensor to MLX array.

    Uses numpy as an intermediate to enable zero-copy on unified memory.

    Args:
        tensor: PyTorch tensor (can be on any device)

    Returns:
        MLX array with the same data
    """
    # Move to CPU if on MPS for numpy conversion
    if tensor.device.type != "cpu":
        tensor = tensor.cpu()

    tensor = tensor.detach()

    # Note: numpy does not support bfloat16.
    if tensor.dtype == torch.bfloat16:
        return mx.array(tensor)

    return mx.array(tensor.numpy())


# TODO(perf): accept a list/batch of arrays and convert them in one pass
# to reduce the Python ↔ MLX round-trip overhead.
def mlx_to_torch(
    array: mx.array,
    device: torch.device | Literal["mps", "cpu"] | None = None,
    already_contiguous: bool = False,
) -> torch.Tensor:
    """Convert MLX array to PyTorch tensor.

    Uses numpy as an intermediate to enable zero-copy on unified memory.

    Args:
        array: MLX array
        device: Target PyTorch device (default: MPS if available)
        already_contiguous: Skip contiguity check if array is known contiguous

    Returns:
        PyTorch tensor with the same data
    """
    if device is None:
        device = get_torch_device()
    elif isinstance(device, str):
        device = torch.device(device)

    # Use memoryview for zero-copy conversion (bypasses numpy for bfloat16)
    # reference: https://github.com/ml-explore/mlx/issues/403
    torch_dtype = MLX_TO_TORCH_DTYPE.get(array.dtype)
    if torch_dtype is not None:
        if already_contiguous:
            # Fast path: skip contiguity check, single eval
            mx.eval(array)
            buffer = memoryview(array)
        else:
            # MLX views / non-contiguous arrays expose a non-contiguous buffer (or
            # sometimes no usable buffer), which `torch.frombuffer` can't consume.
            # Make contiguous first, then eval once
            array = mx.contiguous(array)
            mx.eval(array)
            buffer = memoryview(array)

        tensor = torch.frombuffer(buffer, dtype=torch_dtype).reshape(array.shape)
    else:
        # Fallback to numpy path for unsupported dtypes
        raise ValueError(f"Unsupported MLX dtype: {array.dtype}")

    # Move to target device, but check for MPS size limits first
    if device.type == "mps":
        if _is_safe_for_mps(array):
            tensor = tensor.to(device)
        else:
            # Large tensor - keep on CPU to avoid MPS 4GB limit crash
            # See: https://github.com/anthropics/vllm-metal/issues/43
            logger.debug(
                "Tensor too large for MPS (%d bytes > %d limit), keeping on CPU",
                _get_tensor_size_bytes(array),
                _MPS_SAFE_SIZE_BYTES,
            )
    elif device.type != "cpu":
        tensor = tensor.to(device)

    return tensor


def sync_mlx() -> None:
    """Synchronize MLX operations.

    Call this before converting MLX arrays to ensure all operations complete.
    """
    # Prefer an explicit MLX barrier when available; otherwise force evaluation.
    # `mx.eval([])` is a no-op, so we evaluate a tiny scalar as a safe fallback.
    try:
        mx.synchronize()
    except (AttributeError, TypeError):
        mx.eval(mx.array(0, dtype=mx.int32))


def sync_torch() -> None:
    """Synchronize PyTorch MPS operations.

    Call this before converting PyTorch tensors to ensure all operations complete.
    """
    if torch.backends.mps.is_available():
        torch.mps.synchronize()


__all__ = [
    "is_mlx_available",
    "use_mlx",
    "mlx_to_torch",
    "torch_to_mlx",
    "get_torch_device",
]