[MLX] Upgrade to Torch 2.13/MLX 0.32+ and redesign the Torch-MLX tensor bridge (#32984)

Co-authored-by: Alex Nails <alex.nails@radixark.ai>

python/pyproject_other.toml

@@ -167,14 +167,15 @@ diffusion_musa = [
   "vsa==0.0.4",
 ]
 
-# https://docs.sglang.io/platforms/mps.md
+# https://docs.sglang.io/hardware-platforms/apple_metal
 srt_mps = [
-    "mlx",
+    "mlx>=0.32.0",
     "mlx-lm",
     "sglang[runtime_common]",
-    "torch==2.11.0",
+    "torch==2.13.0",
     "torchaudio==2.11.0",
-    "torchvision",
+    "torchcodec==0.15.0",
+    "torchvision==0.28.0",
 ]
 
 diffusion_mps = [

python/sglang/_platform_stubs.py

@@ -431,9 +431,44 @@ def install_platform_stubs() -> None:
 
             pass
 
+        class _KernelInterface(_StubBase):
+            pass
+
         jit_mod.JITFunction = _JITFunction
+        jit_mod.KernelInterface = _KernelInterface
         runtime.jit = jit_mod
 
+        # Torch 2.13 imports these as classes while initializing Inductor, even on
+        # MPS where no Triton kernel is compiled. Define them explicitly so the
+        # catch-all meta-path finder does not materialize class names as modules.
+        autotuner = _make_mock("triton.runtime.autotuner")
+
+        class _OutOfResources(Exception):
+            pass
+
+        class _PTXASError(Exception):
+            pass
+
+        autotuner.OutOfResources = _OutOfResources
+        autotuner.PTXASError = _PTXASError
+        runtime.autotuner = autotuner
+
+        compiler_root = _make_mock("triton.compiler")
+
+        class _CompiledKernel(_StubBase):
+            pass
+
+        compiler_root.CompiledKernel = _CompiledKernel
+        compiler_impl = _make_mock("triton.compiler.compiler")
+
+        class _ASTSource(_StubBase):
+            pass
+
+        compiler_impl.ASTSource = _ASTSource
+        compiler_impl.triton_key = lambda: "triton-stub"
+        compiler_root.compiler = compiler_impl
+        triton.compiler = compiler_root
+
         # triton.runtime.driver
         driver = _make_mock("triton.runtime.driver")
         runtime.driver = driver
@@ -452,6 +487,11 @@ def install_platform_stubs() -> None:
         backends = _make_mock("triton.backends")
         triton.backends = backends
         compiler = _make_mock("triton.backends.compiler")
+
+        class _GPUTarget(_StubBase):
+            pass
+
+        compiler.GPUTarget = _GPUTarget
         backends.compiler = compiler
 
     mps = torch.mps

python/sglang/benchmark/one_batch.py

@@ -71,6 +71,7 @@ from sglang.srt.distributed.parallel_state import (
 )
 from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.entrypoints.engine import _set_envs_and_config
+from sglang.srt.hardware_backend.mlx.runtime import use_mlx
 from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
 from sglang.srt.layers.moe import initialize_moe_config
 from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config
@@ -96,7 +97,6 @@ from sglang.srt.utils import (
     suppress_other_loggers,
 )
 from sglang.srt.utils.hf_transformers_utils import get_tokenizer
-from sglang.srt.utils.tensor_bridge import use_mlx
 
 
 def start_profile(

python/sglang/kernels/ops/diffusion/README.md

@@ -21,7 +21,7 @@ free to move; the facade is not. `test_import_surface.py` enforces this, with
 a small allowlist for tests that deliberately exercise one backend.
 
 Resolution is lazy (PEP 562): the backends have disjoint heavy dependencies
-(Triton, CUTLASS/CuTe-DSL, FlyDSL on ROCm, MLX on Apple), so an eager
+(Triton, CUTLASS/CuTe-DSL, and FlyDSL on ROCm), so an eager
 re-export would make all of them import-time requirements everywhere.
 
 ## Layout

python/sglang/kernels/ops/diffusion/__init__.py

@@ -18,7 +18,7 @@ numerics and platform plumbing, ``sites`` the request-scoped mount policy, and
 ``README.md``: several norms look interchangeable and are not.
 
 Resolution is lazy (PEP 562).  The backends have disjoint, heavy dependencies
--- Triton, CUTLASS/CuTe-DSL, FlyDSL (ROCm), MLX (Apple) -- so an eager
+-- Triton, CUTLASS/CuTe-DSL, and FlyDSL (ROCm) -- so an eager
 re-export would turn every one of them into a hard import-time requirement on
 every platform.  ``_EXPORTS`` maps a symbol to its module and the import
 happens on first attribute access.

python/sglang/kernels/ops/diffusion/common/fallback_mps.py

@@ -1,122 +0,0 @@
-"""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 the live platform is MPS (see ``common.platform``).
-
-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.utils.tensor_bridge import mlx_to_torch, torch_to_mlx, use_mlx
-
-from .fallback_torch import (
-    apply_rotary_embedding_native as apply_rotary_embedding_native,
-)
-from .fallback_torch import (
-    fuse_scale_shift_kernel_native as fuse_scale_shift_kernel_native,
-)
-from .fallback_torch import (
-    norm_infer_native,
-    rms_norm_fn_native,
-    triton_one_pass_rms_norm_native,
-)
-
-_use_mlx = use_mlx()
-
-if _use_mlx:
-    import mlx.core as mx
-
-# 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."""
-        device = x.device
-        orig_dtype = x.dtype
-        x_mx = torch_to_mlx(x)
-        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)
-
-    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)."""
-        device = x.device
-        orig_dtype = x.dtype
-        if residual is not None:
-            x = x.float() + residual.float()
-            residual_out_val = x.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)
-        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)
-        if residual is not None and residual_out_val is not None:
-            return y, residual_out_val
-        return y

python/sglang/kernels/ops/diffusion/common/fallback_torch.py

@@ -8,6 +8,7 @@ implementations replace the Triton kernels
 from typing import Optional
 
 import torch
+import torch.nn.functional as F
 from torch import Tensor
 
 
@@ -69,20 +70,36 @@ def norm_infer_native(
     out: Optional[Tensor] = None,
 ) -> Tensor:
     """Native fallback for norm_infer (layer norm / rms norm inference)."""
-    orig_dtype = x.dtype
-    x = x.contiguous().float()
+    normalized_shape = (x.shape[-1],)
     if is_rms_norm:
-        variance = x.pow(2).mean(dim=-1, keepdim=True)
-        x_hat = x * torch.rsqrt(variance + eps)
+        # ``F.rms_norm`` returns in the input dtype before a separately added
+        # bias is applied. Promote the whole branch when a bias is present or
+        # parameters have a different dtype, retaining the old
+        # fp32-accumulate-then-cast contract. A bias triggers promotion even
+        # when its dtype matches the input; otherwise the normalized value is
+        # rounded before the affine add. Keep the native fast path for the
+        # common bias-free, same-dtype case.
+        needs_fp32 = bias is not None or (
+            weight is not None and weight.dtype != x.dtype
+        )
+        if needs_fp32:
+            result = F.rms_norm(
+                x.float(),
+                normalized_shape,
+                weight.float() if weight is not None else None,
+                eps,
+            )
+            if bias is not None:
+                result = result + bias.float()
+        else:
+            result = F.rms_norm(x, normalized_shape, weight, eps)
+            if bias is not None:
+                result = result + bias
     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)
+        result = F.layer_norm(x, normalized_shape, weight, bias, eps)
+    # Match the original fallback and Triton kernel contract even when a
+    # higher-precision weight or bias promotes PyTorch's intermediate result.
+    result = result.to(x.dtype)
     if out is not None:
         out.copy_(result)
         return out
@@ -93,12 +110,7 @@ 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)
+    return F.rms_norm(x, (x.shape[-1],), w, eps)
 
 
 def rms_norm_fn_native(
@@ -130,13 +142,13 @@ def rms_norm_fn_native(
         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
+    else:
+        w = None
+    x_hat = F.rms_norm(x, (x.shape[-1],), w, eps)
     if bias is not None:
         x_hat = x_hat + bias.float()
     final_dtype = out_dtype if out_dtype is not None else orig_dtype

python/sglang/kernels/ops/diffusion/common/platform.py

@@ -1,8 +1,8 @@
 """Platform predicates and the import-time fallback selector.
 
 Several diffusion Triton kernels have no Triton on the live device (Ascend
-NPU, Apple MPS, MUSA, CPU) and must resolve to a pure-``torch`` — or
-MLX-accelerated — implementation.  That choice is made once at import time,
+NPU, Apple MPS, MUSA, CPU) and must resolve to a pure-``torch`` implementation.
+That choice is made once at import time,
 which used to mean a hand-rolled four-branch ``if`` block repeated in every
 such module, each importing ``current_platform`` directly.
 
@@ -75,8 +75,8 @@ def lazy_fallback(kind: str, name: str) -> Callable:
     """Name a fallback without importing its module.
 
     ``select_impl`` is handed every candidate at once, so a plain import here
-    would pull in *all* fallback modules on every platform -- including MLX on
-    CUDA hosts.  The returned shim imports ``common.fallback_<kind>`` on its
+    would pull in *all* fallback modules on every platform.  The returned shim
+    imports ``common.fallback_<kind>`` on its
     first call instead, which for the unselected candidates never happens.
     """
 

python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py

@@ -736,7 +736,7 @@ def fuse_residual_layernorm_scale_shift_gate_select01_kernel(
 fuse_scale_shift_kernel = select_impl(
     fuse_scale_shift_kernel,
     npu=lazy_fallback("npu", "fuse_scale_shift_native"),
-    mps=lazy_fallback("mps", "fuse_scale_shift_kernel_native"),
+    mps=lazy_fallback("torch", "fuse_scale_shift_kernel_native"),
     musa=lazy_fallback("torch", "fuse_scale_shift_kernel_native"),
     cpu=lazy_fallback("torch", "fuse_scale_shift_kernel_native"),
 )

python/sglang/kernels/ops/diffusion/norm/norm_triton.py

@@ -649,11 +649,11 @@ def norm_infer(
 
 norm_infer = select_impl(
     norm_infer,
-    mps=lazy_fallback("mps", "norm_infer_native"),
+    mps=lazy_fallback("torch", "norm_infer_native"),
     cpu=lazy_fallback("torch", "norm_infer_native"),
 )
 rms_norm_fn = select_impl(
     rms_norm_fn,
-    mps=lazy_fallback("mps", "rms_norm_fn_native"),
+    mps=lazy_fallback("torch", "rms_norm_fn_native"),
     cpu=lazy_fallback("torch", "rms_norm_fn_native"),
 )

python/sglang/kernels/ops/diffusion/norm/rmsnorm_onepass_triton.py

@@ -72,6 +72,6 @@ def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6
 triton_one_pass_rms_norm = select_impl(
     triton_one_pass_rms_norm,
     # MPS keeps the api-logging wrapper the Triton entry point carries.
-    mps=debug_kernel_api(lazy_fallback("mps", "triton_one_pass_rms_norm_native")),
+    mps=debug_kernel_api(lazy_fallback("torch", "triton_one_pass_rms_norm_native")),
     cpu=lazy_fallback("torch", "triton_one_pass_rms_norm_native"),
 )

python/sglang/kernels/ops/diffusion/rope/rotary_triton.py

@@ -128,6 +128,6 @@ def apply_rotary_embedding(
 apply_rotary_embedding = select_impl(
     apply_rotary_embedding,
     npu=lazy_fallback("npu", "apply_rotary_embedding_native"),
-    mps=lazy_fallback("mps", "apply_rotary_embedding_native"),
+    mps=lazy_fallback("torch", "apply_rotary_embedding_native"),
     cpu=lazy_fallback("torch", "apply_rotary_embedding_native"),
 )

python/sglang/multimodal_gen/README.md

@@ -30,7 +30,7 @@ SGLang Diffusion supports Moore Threads GPUs (MTGPU) through the MUSA software s
 
 ### 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://docs.sglang.io/docs/sglang-diffusion/installation) for setup instructions.
+SGLang Diffusion supports Apple Silicon (M-series) via the MPS backend. Since Triton is Linux-only, Triton kernels are replaced with PyTorch-native fallbacks on MPS. See the [installation guide](https://docs.sglang.io/docs/sglang-diffusion/installation) for setup instructions.
 
 ## Getting Started
 

python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py

@@ -152,8 +152,11 @@ class TransformerLoader(ComponentLoader):
     def customized_load_kwargs_for_component(
         self, server_args: ServerArgs, component_name: str
     ) -> dict[str, bool]:
-        if current_platform.is_mps() and self._is_component_set_as_layerwise_load(
-            server_args, component_name
+        if (
+            current_platform.is_mps()
+            and server_args.should_configure_layerwise_offload_for_lazy_component(
+                component_name
+            )
         ):
             logger.info(
                 "Loading %s on CPU first for MPS layerwise offload", component_name

python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py

@@ -150,8 +150,11 @@ class VAELoader(ComponentLoader):
     def customized_load_kwargs_for_component(
         self, server_args: ServerArgs, component_name: str
     ) -> dict[str, bool]:
-        if current_platform.is_mps() and self._is_component_set_as_layerwise_load(
-            server_args, component_name
+        if (
+            current_platform.is_mps()
+            and server_args.should_configure_layerwise_offload_for_lazy_component(
+                component_name
+            )
         ):
             logger.info(
                 "Loading %s on CPU first for MPS layerwise offload", component_name

python/sglang/multimodal_gen/test/unit/test_diffusion_import_isolation.py

@@ -0,0 +1,37 @@
+"""Import smoke tests for the diffusion Torch path."""
+
+import os
+import subprocess
+import sys
+import unittest
+
+
+class TestDiffusionImportIsolation(unittest.TestCase):
+    def test_disabled_backend_does_not_import_mlx(self):
+        """Diffusion modules must remain usable without the optional MLX path."""
+        script = """
+import sys
+from sglang.kernels.ops.diffusion import norm_infer
+from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
+assert norm_infer is not None and RMSNorm is not None
+assert not any(name == "mlx" or name.startswith("mlx.") for name in sys.modules)
+"""
+        env = os.environ.copy()
+        env.pop("SGLANG_USE_MLX", None)
+        completed = subprocess.run(
+            [sys.executable, "-c", script],
+            capture_output=True,
+            text=True,
+            timeout=30,
+            check=False,
+            env=env,
+        )
+        self.assertEqual(
+            completed.returncode,
+            0,
+            msg=f"stdout={completed.stdout}\nstderr={completed.stderr}",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

python/sglang/multimodal_gen/test/unit/test_transformer_quant.py

@@ -67,6 +67,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
 )
 from sglang.multimodal_gen.runtime.loader.component_loaders import transformer_loader
 from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import (
+    TransformerLoader,
     _default_quantized_attention_backend,
     _resolve_checkpoint_load_device,
     _warn_if_expected_param_dtype_missing,
@@ -121,6 +122,29 @@ def _make_quant_config(name: str, **attrs):
 
 
 class TestTransformerQuantHelpers(unittest.TestCase):
+    def test_mps_layerwise_load_uses_residency_api(self):
+        server_args = SimpleNamespace(
+            should_configure_layerwise_offload_for_lazy_component=lambda name: (
+                name == "transformer"
+            )
+        )
+
+        with patch.object(
+            transformer_loader.current_platform, "is_mps", return_value=True
+        ):
+            self.assertEqual(
+                TransformerLoader().customized_load_kwargs_for_component(
+                    server_args, "transformer"
+                ),
+                {"cpu_offload_flag": True},
+            )
+            self.assertEqual(
+                TransformerLoader().customized_load_kwargs_for_component(
+                    server_args, "audio_dit"
+                ),
+                {},
+            )
+
     def _make_server_args(self, **overrides):
         defaults = dict(
             transformer_weights_path=None,

python/sglang/multimodal_gen/test/unit/test_vae_loader.py

@@ -38,6 +38,7 @@ class _FakeServerArgs:
         self.model_paths = {}
         self.revision = "test-revision"
         self.trust_remote_code = True
+        self.layerwise_components = set()
 
     def resolve_component_attention_backend(self, _component_name):
         return None, None
@@ -45,6 +46,9 @@ class _FakeServerArgs:
     def should_start_component_on_cpu(self, _component_name):
         return False
 
+    def should_configure_layerwise_offload_for_lazy_component(self, component_name):
+        return component_name in self.layerwise_components
+
 
 class TestKeepCheckpointMapped(unittest.TestCase):
     """The mapping is for hosts that cannot afford the whole deployment."""
@@ -92,6 +96,21 @@ class TestMatchCheckpointDtypes(unittest.TestCase):
 
 
 class TestVAELoader(unittest.TestCase):
+    def test_mps_layerwise_load_uses_residency_api(self):
+        loader = vae_loader.VAELoader()
+        server_args = _FakeServerArgs(QwenImagePipelineConfig())
+        server_args.layerwise_components.add("vae")
+
+        with patch.object(vae_loader.current_platform, "is_mps", return_value=True):
+            self.assertEqual(
+                loader.customized_load_kwargs_for_component(server_args, "vae"),
+                {"cpu_offload_flag": True},
+            )
+            self.assertEqual(
+                loader.customized_load_kwargs_for_component(server_args, "audio_vae"),
+                {},
+            )
+
     def test_quantized_vae_admission_leaves_plain_configs_unchanged(self):
         _require_native_loader_for_quantized_vae(
             {"_class_name": "AutoencoderKL"}, "vae"

python/sglang/srt/arg_groups/overrides.py

@@ -37,6 +37,7 @@ from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
 
 from sglang.srt.arg_groups.arg_utils import resolvable_fields
 from sglang.srt.environ import envs
+from sglang.srt.hardware_backend.mlx.runtime import use_mlx
 from sglang.srt.model_executor.cuda_graph_config import Backend
 from sglang.srt.utils.common import (
     cpu_has_amx_support,
@@ -62,7 +63,6 @@ from sglang.srt.utils.common import (
     is_xpu,
     xpu_has_xmx_support,
 )
-from sglang.srt.utils.tensor_bridge import use_mlx
 
 logger = logging.getLogger(__name__)
 

python/sglang/srt/hardware_backend/mlx/profiler.py

@@ -12,8 +12,8 @@ from typing import Any, Callable, Optional
 
 import torch
 
+from sglang.srt.hardware_backend.mlx.runtime import use_mlx
 from sglang.srt.managers.io_struct import ProfileReqOutput
-from sglang.srt.utils.tensor_bridge import use_mlx
 
 logger = logging.getLogger(__name__)
 

python/sglang/srt/hardware_backend/mlx/runtime.py

@@ -0,0 +1,69 @@
+"""Runtime gate for the opt-in MLX backend on Apple silicon."""
+
+from functools import lru_cache
+
+import torch
+from packaging.version import InvalidVersion, Version
+
+from sglang.srt.environ import envs
+
+_MIN_MLX_VERSION = Version("0.32.0")
+_SUPPORTED_TORCH_SERIES = (2, 13)
+
+
+def _is_stable_series(raw_version: object, series: tuple[int, int]) -> bool:
+    try:
+        version = Version(str(raw_version))
+    except InvalidVersion:
+        return False
+    return not version.is_prerelease and (version.major, version.minor) == series
+
+
+def _is_stable_at_least(raw_version: object, minimum: Version) -> bool:
+    try:
+        version = Version(str(raw_version))
+    except InvalidVersion:
+        return False
+    return not version.is_prerelease and version >= minimum
+
+
+@lru_cache(maxsize=1)
+def _validate_runtime() -> None:
+    try:
+        import mlx.core as mx
+    except ImportError:
+        raise RuntimeError(
+            "SGLANG_USE_MLX requires stable Torch 2.13.x and MLX >= 0.32.0, "
+            "but MLX is not installed; reinstall with "
+            "the srt_mps extra"
+        ) from None
+    mlx_version = getattr(mx, "__version__", None)
+    torch_version = getattr(torch, "__version__", None)
+    if not _is_stable_series(
+        torch_version, _SUPPORTED_TORCH_SERIES
+    ) or not _is_stable_at_least(mlx_version, _MIN_MLX_VERSION):
+        raise RuntimeError(
+            "SGLANG_USE_MLX requires stable Torch 2.13.x and MLX >= 0.32.0; "
+            "found "
+            f"Torch {torch_version or 'unknown'} + MLX {mlx_version or 'unknown'}; "
+            "reinstall with the srt_mps extra"
+        )
+
+    mps_backend = getattr(torch.backends, "mps", None)
+    is_mps_available = getattr(mps_backend, "is_available", None)
+    if not callable(is_mps_available) or not is_mps_available():
+        raise RuntimeError("SGLANG_USE_MLX requires an available PyTorch MPS device")
+
+    metal = getattr(mx, "metal", None)
+    is_available = getattr(metal, "is_available", None)
+    if not callable(is_available) or not is_available():
+        raise RuntimeError("SGLANG_USE_MLX requires an available MLX Metal device")
+
+
+@lru_cache(maxsize=1)
+def use_mlx() -> bool:
+    """Return whether the validated MLX backend was explicitly enabled."""
+    enabled = bool(envs.SGLANG_USE_MLX.get())
+    if enabled:
+        _validate_runtime()
+    return enabled

python/sglang/srt/managers/scheduler.py

@@ -103,6 +103,7 @@ from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.dllm.mixin.scheduler import SchedulerDllmMixin
 from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
+from sglang.srt.hardware_backend.mlx.runtime import use_mlx
 from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
 from sglang.srt.layers.moe import initialize_moe_config
 from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config
@@ -326,7 +327,6 @@ from sglang.srt.utils.hf_transformers_utils import (
 from sglang.srt.utils.msgspec_utils import msgspec_to_builtins
 from sglang.srt.utils.numa_utils import get_numa_node_if_available, numa_bind_to_node
 from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method
-from sglang.srt.utils.tensor_bridge import use_mlx
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 from sglang.utils import TypeBasedDispatcher, get_exception_traceback
 

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -32,6 +32,7 @@ from sglang.srt.configs.hybrid_arch import (
 )
 from sglang.srt.configs.model_config import ModelImpl, is_deepseek_dsa
 from sglang.srt.environ import envs
+from sglang.srt.hardware_backend.mlx.runtime import use_mlx
 from sglang.srt.managers.mm_schedule import init_mm_embedding_cache
 from sglang.srt.mem_cache.cache_init_params import CacheInitParams
 from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
@@ -46,7 +47,6 @@ from sglang.srt.runtime_context import (
     get_parallel,
     get_schedule,
 )
-from sglang.srt.utils.tensor_bridge import use_mlx
 
 if TYPE_CHECKING:
 

python/sglang/srt/mem_cache/registry.py

@@ -15,10 +15,10 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Callable, Optional
 
 from sglang.srt.environ import envs
+from sglang.srt.hardware_backend.mlx.runtime import use_mlx
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.mem_cache.cache_init_params import CacheInitParams
 from sglang.srt.runtime_context import get_disagg, get_memory
-from sglang.srt.utils.tensor_bridge import use_mlx
 
 if TYPE_CHECKING:
     from sglang.srt.configs.model_config import ModelConfig

python/sglang/srt/server_args.py

@@ -55,6 +55,7 @@ from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import
 )
 from sglang.srt.environ import envs
 from sglang.srt.function_call.function_call_parser import FunctionCallParser
+from sglang.srt.hardware_backend.mlx.runtime import use_mlx
 from sglang.srt.lora.lora_registry import LoRARef
 from sglang.srt.model_executor.cuda_graph_config import (
     ALLOWED_BACKENDS_PER_PHASE,
@@ -102,7 +103,6 @@ from sglang.srt.utils.common import (
 from sglang.srt.utils.hf_transformers_utils import check_gguf_file
 from sglang.srt.utils.network import NetworkAddress, get_free_port, wait_port_available
 from sglang.srt.utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
-from sglang.srt.utils.tensor_bridge import use_mlx
 from sglang.utils import is_in_ci
 
 logger = logging.getLogger(__name__)
@@ -3665,6 +3665,10 @@ class ServerArgs:
         self._handle_media_url_security()
         self._handle_hicache_ratio_default()
         self._validate_prefill_decode_interval()
+
+        # Reject an explicitly enabled but incompatible hardware runtime before
+        # model path resolution, downloads, or the dummy-model short circuit.
+        self._handle_hardware_runtime_validation()
         if self.model_path.lower() in ["none", "dummy"]:
             return
 
@@ -4401,6 +4405,13 @@ class ServerArgs:
                 )
             self.sampling_backend = "pytorch"
 
+    def _handle_hardware_runtime_validation(self):
+        # This is intentionally independent of self.device: setting
+        # SGLANG_USE_MLX opts into the MLX backend and must fail immediately if
+        # the environment cannot honor that request. With the flag unset,
+        # use_mlx() remains lazy and does not import MLX.
+        use_mlx()
+
     def _handle_npu_backends(self):
         if self.device == "npu":
             from sglang.srt.hardware_backend.npu.utils import set_default_server_args

python/sglang/srt/utils/tensor_bridge.py

@@ -1,78 +1,57 @@
 # Copied and adapted from: https://github.com/vllm-project/vllm-metal
 # SPDX-License-Identifier: Apache-2.0
-"""Tensor bridge between MLX and PyTorch.
+"""Tensor bridge between MLX and PyTorch on Apple silicon.
 
-Provides zero-copy conversion when possible using Apple Silicon's unified memory.
+The MLX backend requires MLX >= 0.32 and PyTorch >= 2.13.  Ordinary
+``torch_to_mlx`` conversion creates an independent MLX allocation.  The
+zero-copy ``mlx_call`` helper is available for a complete MLX operation and
+keeps all borrowed DLPack inputs alive until the result has been evaluated.
+This lifetime boundary matters because MLX may donate a borrowed input buffer
+to a lazy operation.
+
+Bridge entry points are serialized because Torch and MLX use different stream
+abstractions over the same Metal command queues.  The lock covers producer
+fencing, MLX evaluation, and DLPack import.  It cannot cover arbitrary MPS
+work outside the function, so callers must serialize any overlapping use or
+mutation of source and returned MPS tensors as well.
 """
 
 from __future__ import annotations
 
-import logging
-from functools import lru_cache
-from typing import TYPE_CHECKING, Literal
+from functools import lru_cache, wraps
+from threading import RLock
+from typing import TYPE_CHECKING, Any, Callable, 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
+_BRIDGE_LOCK = RLock()
 
 
-def is_mlx_available() -> bool:
-    """Return True when the ``mlx`` package can be imported."""
-    return _MLX_AVAILABLE
+def _serialized_bridge(function: Callable[..., Any]) -> Callable[..., Any]:
+    """Serialize one complete Torch/MLX crossing, including result export."""
+
+    @wraps(function)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        with _BRIDGE_LOCK:
+            return function(*args, **kwargs)
+
+    return wrapper
 
 
 @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
+def _mlx_core():
+    try:
+        import mlx.core as mx
+    except ImportError:
+        raise RuntimeError("The MLX tensor bridge requires MLX >= 0.32.0") from None
+    return mx
 
 
-# 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:
+def _get_torch_device() -> torch.device:
     """Get the PyTorch device for Metal/MPS.
 
     Returns:
@@ -83,146 +62,377 @@ def get_torch_device() -> torch.device:
     return torch.device("cpu")
 
 
-def _get_tensor_size_bytes(array: mx.array) -> int:
-    """Calculate the size of an MLX array in bytes.
+def _torch_to_mlx(
+    tensor: torch.Tensor,
+    *,
+    copy: bool,
+    synchronize: bool = True,
+) -> mx.array:
+    """Convert one tensor, optionally borrowing its MPS allocation."""
+    mx = _mlx_core()
+    tensor = tensor.detach()
 
-    Args:
-        array: MLX array
+    if tensor.device.type == "mps":
+        if synchronize:
+            # Torch and MLX do not share stream state on Metal.
+            torch.mps.synchronize()
+        return mx.asarray(tensor, copy=copy)
+    if tensor.device.type == "cpu":
+        # CPU tensors always get MLX-owned storage.  In particular, do not
+        # expose a NumPy/memoryview alias whose lifetime is controlled by the
+        # caller.
+        if tensor.dtype == torch.complex128:
+            raise ValueError(
+                "MLX 0.32 does not support complex128; convert the Torch tensor "
+                "to complex64 explicitly"
+            )
+        # MLX 0.32 does not support float64 on its default Metal stream.  Keep
+        # the dtype by constructing this uncommon CPU value on the CPU stream
+        # instead of silently downcasting it to float32.
+        if tensor.dtype == torch.float64:
+            with mx.stream(mx.cpu):
+                return mx.array(tensor, dtype=mx.float64)
+        return mx.array(tensor)
+    raise ValueError(
+        f"The MLX tensor bridge supports CPU and MPS tensors, got {tensor.device}"
+    )
 
-    Returns:
-        Size in bytes
+
+class MlxTensorView:
+    """A lifetime-bound, zero-copy MLX view of a Torch MPS tensor.
+
+    The view deliberately retains a detached Torch tensor *and* the imported
+    MLX array.  Holding only the array is insufficient: a later parameter
+    replacement or garbage collection could invalidate the borrowed storage
+    while MLX still has a lazy graph referring to it. This class is intended
+    for immutable inference weights; construct a new view after replacing the
+    source storage.
     """
-    return array.size * array.dtype.size
+
+    __slots__ = ("torch_tensor", "array")
+
+    def __init__(self, tensor: torch.Tensor, *, synchronize: bool = True):
+        with _BRIDGE_LOCK:
+            owner = tensor.detach()
+            if owner.device.type != "mps":
+                raise ValueError(
+                    f"MlxTensorView requires a Torch MPS tensor, got {owner.device}"
+                )
+            if synchronize:
+                torch.mps.synchronize()
+            self.torch_tensor = owner
+            self.array = _torch_to_mlx(owner, copy=False, synchronize=False)
+
+    @classmethod
+    def _from_synchronized(cls, tensor: torch.Tensor) -> MlxTensorView:
+        view = object.__new__(cls)
+        owner = tensor.detach()
+        if owner.device.type != "mps":
+            raise ValueError(
+                f"MlxTensorView requires a Torch MPS tensor, got {owner.device}"
+            )
+        view.torch_tensor = owner
+        view.array = _torch_to_mlx(owner, copy=False, synchronize=False)
+        return view
+
+    def matches(self, tensor: torch.Tensor) -> bool:
+        """Return whether ``tensor`` still refers to this borrowed storage."""
+        owner = tensor.detach()
+        return (
+            owner.device == self.torch_tensor.device
+            and owner.dtype == self.torch_tensor.dtype
+            and owner.shape == self.torch_tensor.shape
+            and owner.stride() == self.torch_tensor.stride()
+            and owner.data_ptr() == self.torch_tensor.data_ptr()
+        )
 
 
-def _is_safe_for_mps(array: mx.array) -> bool:
-    """Check if an array is safe to transfer to MPS without hitting size limits.
+@_serialized_bridge
+def borrow_torch_tensors(
+    *tensors: torch.Tensor, synchronize: bool = True
+) -> tuple[MlxTensorView, ...]:
+    """Borrow one or more Torch MPS tensors, optionally synchronizing once.
 
-    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
+    The returned views own the Torch tensor references for their entire
+    lifetime.  No data copy is made.  Set ``synchronize=False`` only when a
+    surrounding operation (such as :func:`mlx_call`) performs the producer
+    barrier immediately before consuming the views.  This helper is
+    intentionally separate from :func:`torch_to_mlx`, whose contract is an
+    independent MLX copy.
     """
-    return _get_tensor_size_bytes(array) < _MPS_SAFE_SIZE_BYTES
+    detached = tuple(tensor.detach() for tensor in tensors)
+    if any(tensor.device.type != "mps" for tensor in detached):
+        devices = ", ".join(str(tensor.device) for tensor in detached)
+        raise ValueError(f"borrow_torch_tensors requires MPS tensors, got {devices}")
+    if synchronize and detached:
+        torch.mps.synchronize()
+    return tuple(MlxTensorView._from_synchronized(tensor) for tensor in detached)
 
 
+@_serialized_bridge
 def torch_to_mlx(tensor: torch.Tensor) -> mx.array:
-    """Convert PyTorch tensor to MLX array.
+    """Convert a PyTorch tensor to an independent MLX array.
 
-    Uses numpy as an intermediate to enable zero-copy on unified memory.
+    MPS inputs are copied inside the unified Metal device.  Use ``mlx_call``
+    when a complete operation needs zero-copy MPS input imports; it owns the
+    borrowed MLX arrays for the complete lazy operation.
 
     Args:
-        tensor: PyTorch tensor (can be on any device)
+        tensor: PyTorch CPU or MPS tensor.
 
     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())
+    array = _torch_to_mlx(tensor, copy=True)
+    if tensor.device.type == "mps":
+        # Materialize the owned copy before the caller may mutate or release
+        # the Torch source.
+        _mlx_core().eval(array)
+    return array
 
 
-# TODO(perf): accept a list/batch of arrays and convert them in one pass
-# to reduce the Python ↔ MLX round-trip overhead.
+@_serialized_bridge
+def mlx_call(
+    operation: Callable[..., mx.array],
+    *tensors: torch.Tensor | MlxTensorView,
+    device: torch.device | Literal["mps", "cpu"] | None = None,
+) -> torch.Tensor:
+    """Run one MLX operation with zero-copy Torch MPS input imports.
+
+    The imported MLX arrays remain strongly referenced until
+    :func:`mlx_to_torch` evaluates and exports ``operation``'s result.  Keep
+    the operation inside this call; returning a lazy MLX result for later use
+    or stashing a borrowed input through a callback side effect would escape
+    the borrow scope.  The caller must also serialize any overlapping MPS work
+    outside this function, including use or mutation of source and returned
+    tensors.  The operation may allocate its own output normally.
+    """
+    mx = _mlx_core()
+    target_device = _get_torch_device() if device is None else torch.device(device)
+    if target_device.type not in {"cpu", "mps"}:
+        raise ValueError(
+            f"The MLX tensor bridge supports CPU and MPS targets, got {target_device}"
+        )
+    detached = tuple(
+        tensor.detach() for tensor in tensors if isinstance(tensor, torch.Tensor)
+    )
+    if any(tensor.device.type == "mps" for tensor in detached) or any(
+        isinstance(tensor, MlxTensorView) for tensor in tensors
+    ):
+        torch.mps.synchronize()
+    borrowed: tuple[Any, ...] = tuple(
+        (
+            tensor.array
+            if isinstance(tensor, MlxTensorView)
+            else _torch_to_mlx(tensor.detach(), copy=False, synchronize=False)
+        )
+        for tensor in tensors
+    )
+    # MLX does not support float64 on the Metal stream.  Keep an explicitly
+    # requested CPU call on the CPU stream when a borrowed input carries that
+    # dtype; otherwise even constructing the lazy result would fail before the
+    # export preparation below can move it.
+    if target_device.type == "cpu" and any(
+        array.dtype == mx.float64 for array in borrowed
+    ):
+        with mx.stream(mx.cpu):
+            result = operation(*borrowed)
+    else:
+        result = operation(*borrowed)
+    output = mlx_to_torch(result, device=target_device)
+    # Keep the imported MLX objects (and any MlxTensorView Torch owners) alive
+    # through lazy result evaluation and DLPack export.
+    _ = borrowed
+    return output
+
+
+def _prepare_mlx_export(
+    array: mx.array,
+    target_device: torch.device,
+    mx: Any,
+) -> mx.array:
+    """Prepare one lazy MLX result for the requested Torch target.
+
+    This intentionally does not evaluate the result.  Callers which export
+    several results should prepare every result first and then issue one
+    shared ``mx.eval`` boundary.
+    """
+    if target_device.type not in {"cpu", "mps"}:
+        raise ValueError(
+            f"The MLX tensor bridge supports CPU and MPS targets, got {target_device}"
+        )
+
+    if target_device.type == "mps" and array.dtype == mx.float64:
+        raise ValueError(
+            "MLX float64 arrays cannot be exported to a Torch MPS tensor; "
+            "use float32/bfloat16 or request device='cpu'"
+        )
+
+    return array
+
+
+def _has_negative_stride(array: mx.array) -> bool:
+    """Return whether an evaluated MLX array has a DLPack-incompatible view."""
+    # PyTorch's DLPack importer aborts the process for negative strides.  MLX
+    # exposes the evaluated layout through the buffer protocol, so inspect it
+    # before handing the capsule to PyTorch.
+    with memoryview(array) as view:
+        return any(stride < 0 for stride in (view.strides or ()))
+
+
+def _export_evaluated_mlx(
+    array: mx.array,
+    target_device: torch.device,
+    mx: Any,
+    *,
+    materialize_negative: bool = True,
+) -> torch.Tensor:
+    """Export an already-evaluated MLX result through one DLPack capsule.
+
+    Negative-stride results are materialized here as a safety fallback.  The
+    normal (contiguous/positive-stride) path performs no copy and no extra
+    evaluation; :func:`mlx_call_multi` batches any required materialization
+    evaluations for all outputs together.
+    """
+    if materialize_negative and _has_negative_stride(array):
+        materialize_stream = mx.cpu if target_device.type == "cpu" else mx.gpu
+        array = mx.contiguous(array, stream=materialize_stream)
+        mx.eval(array)
+
+    if target_device.type == "cpu":
+        # MLX owns CPU-accessible unified memory.  Request a CPU DLPack view
+        # explicitly rather than importing on MPS and copying back.
+        dlpack = array.__dlpack__(dl_device=(1, 0), copy=False)
+        return torch.utils.dlpack.from_dlpack(dlpack)
+    return torch.utils.dlpack.from_dlpack(array)
+
+
+@_serialized_bridge
+def mlx_call_multi(
+    operation: Callable[..., tuple[mx.array, ...]],
+    *tensors: torch.Tensor | MlxTensorView,
+    device: torch.device | Literal["mps", "cpu"] | None = None,
+) -> tuple[torch.Tensor, ...]:
+    """Run one MLX operation and export all of its outputs as Torch tensors.
+
+    ``operation`` must return a non-empty flat ``tuple`` or ``list`` of MLX
+    arrays.  All Torch/MPS inputs are fenced once before import, and all
+    ordinary outputs are evaluated with one ``mx.eval(*outputs)`` call before
+    being exported through DLPack.  The imported arrays and detached Torch
+    owners remain local until every output capsule has been consumed, which is
+    required when MLX lazily donates a borrowed input buffer.  Callers must
+    serialize any overlapping MPS work outside this function, including use
+    or mutation of source and returned tensors.
+
+    CPU targets retain the same float64 and negative-stride safeguards as
+    :func:`mlx_to_torch`.  A negative-stride output necessarily needs one
+    additional materialization evaluation; contiguous MPS model outputs take
+    the single-evaluation, zero-copy path.
+
+    """
+    mx = _mlx_core()
+    target_device = _get_torch_device() if device is None else torch.device(device)
+    if target_device.type not in {"cpu", "mps"}:
+        raise ValueError(
+            f"The MLX tensor bridge supports CPU and MPS targets, got {target_device}"
+        )
+
+    detached = tuple(
+        tensor.detach() for tensor in tensors if isinstance(tensor, torch.Tensor)
+    )
+    needs_mps_fence = any(tensor.device.type == "mps" for tensor in detached) or any(
+        isinstance(tensor, MlxTensorView) for tensor in tensors
+    )
+    if needs_mps_fence:
+        torch.mps.synchronize()
+
+    borrowed: tuple[Any, ...] = tuple(
+        (
+            tensor.array
+            if isinstance(tensor, MlxTensorView)
+            else _torch_to_mlx(tensor.detach(), copy=False, synchronize=False)
+        )
+        for tensor in tensors
+    )
+
+    if target_device.type == "cpu" and any(
+        array.dtype == mx.float64 for array in borrowed
+    ):
+        with mx.stream(mx.cpu):
+            result = operation(*borrowed)
+    else:
+        result = operation(*borrowed)
+    if not isinstance(result, (tuple, list)) or not result:
+        raise TypeError(
+            "mlx_call_multi operation must return a non-empty tuple or list of MLX arrays"
+        )
+    arrays = tuple(result)
+    if any(not isinstance(array, mx.array) for array in arrays):
+        raise TypeError("mlx_call_multi outputs must be MLX arrays")
+
+    # Prepare all outputs before crossing the one shared MLX evaluation
+    # boundary. This is the key difference from calling mlx_to_torch in a
+    # loop, which would fence/evaluate every result separately.
+    arrays = tuple(_prepare_mlx_export(array, target_device, mx) for array in arrays)
+    mx.eval(*arrays)
+
+    # DLPack cannot represent negative strides. Materialize all such outputs
+    # together so even this safety path has one additional evaluation boundary
+    # rather than one boundary per result.
+    negative = tuple(_has_negative_stride(array) for array in arrays)
+    if any(negative):
+        materialized = []
+        for array, needs_materialization in zip(arrays, negative):
+            if needs_materialization:
+                stream = mx.cpu if target_device.type == "cpu" else mx.gpu
+                array = mx.contiguous(array, stream=stream)
+            materialized.append(array)
+        arrays = tuple(materialized)
+        mx.eval(*(array for array, needs in zip(arrays, negative) if needs))
+
+    outputs = tuple(
+        _export_evaluated_mlx(array, target_device, mx, materialize_negative=False)
+        for array in arrays
+    )
+    # Keep both borrowed MLX views and their Torch owners alive through the
+    # final DLPack import.  (The local remains live until function return.)
+    _ = borrowed
+    return outputs
+
+
+@_serialized_bridge
 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.
+    MLX arrays with PyTorch-compatible strides share their unified-memory
+    allocation through DLPack, including explicit CPU views. Negative-stride
+    views are materialized because PyTorch's DLPack importer cannot represent
+    them safely. MLX is evaluated before the handoff because the frameworks do
+    not share stream state. Only CPU and MPS targets are supported; other
+    target devices are rejected.
 
     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()
+    mx = _mlx_core()
+    target_device = _get_torch_device() if device is None else torch.device(device)
+    array = _prepare_mlx_export(array, target_device, mx)
+    mx.eval(array)
+    return _export_evaluated_mlx(array, target_device, mx)
 
 
 __all__ = [
-    "is_mlx_available",
-    "use_mlx",
+    "MlxTensorView",
+    "borrow_torch_tensors",
+    "mlx_call",
+    "mlx_call_multi",
     "mlx_to_torch",
     "torch_to_mlx",
-    "get_torch_device",
 ]