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

docs/docs/hardware-platforms/apple_metal.mdx

@@ -4,20 +4,31 @@ metatags:
     description: "Run SGLang on Apple Silicon using the Metal backend."
 ---
 
-This document describes how run SGLang on Apple Silicon using [Metal (MLX)](https://opensource.apple.com/projects/mlx/). If you encounter issues or have questions, please [open an issue](https://github.com/sgl-project/sglang/issues).
+This document describes how to run the SGLang serving runtime on Apple Silicon
+using [MLX](https://opensource.apple.com/projects/mlx/). SGLang Diffusion uses
+PyTorch MPS instead; see its [installation guide](/docs/sglang-diffusion/installation#platform-specific-apple-mps).
+If you encounter issues or have questions, please [open an issue](https://github.com/sgl-project/sglang/issues).
 
 ## Prerequisites
 
-Building the native Metal kernels in `sgl-kernel` requires the Apple
+The MLX runtime requires Apple Silicon with macOS 14 or newer, stable PyTorch
-toolchain (`clang++`, the Metal framework headers, and `xcrun`). These ship
+2.13.x, and stable MLX 0.32.0 or newer. The `srt_mps` extra installs PyTorch
-with the **Xcode Command Line Tools**, which cannot be installed via `pip`:
+2.13.0 and MLX 0.32.0 or newer; startup accepts stable PyTorch 2.13 patch
+releases and newer stable MLX releases.
+
+With `SGLANG_USE_MLX=1`, SGLang validates both framework versions and Metal
+availability during argument initialization and stops before resolving or
+downloading a model when the runtime is incompatible.
+
+Building the optional native Metal kernels in `sgl-kernel` requires the Metal
+shader compiler from the full Xcode application. The standalone Xcode Command
+Line Tools are not sufficient. After installing Xcode, select it with:
 
 ```bash
-xcode-select --install
+sudo xcode-select -s /Applications/Xcode.app/Contents/Developer
 ```
 
-If you have the full Xcode app installed, the Command Line Tools are already
+Verify the compiler with `xcrun -sdk macosx metal --version`.
-available. You can verify with `xcode-select -p && xcrun --find metal`.
 
 ## Install SGLang
 
@@ -89,15 +100,15 @@ The MLX backend supports two quantization paths on Apple Silicon:
 
 ## Benchmarking with Requests
 
-`sglang.benchmark_one_batch` calls the synchronous prefill/decode methods directly without going through the scheduler and the overlap code path.
+`sglang.benchmark.one_batch` calls the synchronous prefill/decode methods directly without going through the scheduler and the overlap code path.
 
-`sglang.benchmark_offline_throughput` can toggle overlap scheduling as it uses the scheduler and the overlap code path by using the flag `--disable-overlap-schedule`.
+`sglang.benchmark.offline_throughput` can toggle overlap scheduling as it uses the scheduler and the overlap code path by using the flag `--disable-overlap-schedule`.
 
 ### Throughput Testing
 
 Basic synchronous one batch throughput:
 ```bash
-SGLANG_USE_MLX=1 python -m sglang.bench_one_batch \
+SGLANG_USE_MLX=1 python -m sglang.benchmark.one_batch \
   --model-path <MODEL_ID_OR_PATH> \
   --disable-cuda-graph \
   --tp-size 1 \
@@ -108,7 +119,7 @@ SGLANG_USE_MLX=1 python -m sglang.bench_one_batch \
 
 Synchronous offline throughput:
 ```bash
-SGLANG_USE_MLX=1 python -m sglang.bench_offline_throughput \
+SGLANG_USE_MLX=1 python -m sglang.benchmark.offline_throughput \
   --model-path <MODEL_ID_OR_PATH> \
   --disable-cuda-graph \
   --num-prompts 1 \
@@ -117,7 +128,7 @@ SGLANG_USE_MLX=1 python -m sglang.bench_offline_throughput \
 
 Asynchronous offline throughput:
 ```bash
-SGLANG_USE_MLX=1 python -m sglang.bench_offline_throughput \
+SGLANG_USE_MLX=1 python -m sglang.benchmark.offline_throughput \
   --model-path <MODEL_ID_OR_PATH> \
   --disable-cuda-graph \
   --num-prompts 1

docs/docs/references/environment_variables.mdx

@@ -1891,6 +1891,9 @@ SGLang supports various environment variables that can be used to configure its
 
 ## Apple Silicon (MLX / MPS)
 
+These variables configure the SRT MLX backend. SGLang Diffusion uses PyTorch
+MPS and does not read them.
+
 <table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
   <colgroup>
     <col style={{width: "33.3%"}} />

docs/docs/sglang-diffusion/environment_variables.mdx

@@ -127,7 +127,7 @@ description: "Configure SGLang diffusion behavior with environment variables."
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_USE_MLX</code></td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>not set</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Set to <code>1</code> to enable MLX fused Metal kernels for norm ops on MPS</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>SRT only: enables the MLX serving backend. It has no effect on SGLang Diffusion, which uses PyTorch MPS.</td>
     </tr>
   </tbody>
 </table>

docs/docs/sglang-diffusion/fused_kernels.mdx

@@ -183,7 +183,7 @@ Runtime code imports from the package, never from a submodule:
 from sglang.kernels.ops.diffusion import fused_rmsnorm_scale_shift_bitexact
 ```
 
-Resolution is lazy: the backends have disjoint, heavy dependencies (Triton, CUTLASS/CuTe-DSL, FlyDSL on ROCm, MLX on Apple), so an eager re-export would make every one of them an import-time requirement on every platform. Each public kernel is a predicate-plus-kernel pair — call `can_use_<op>(...)` first and fall back to the reference chain when it returns `False`; the kernel raises on an unsupported input rather than silently returning `None`.
+Resolution is lazy: the backends have disjoint, heavy dependencies (Triton, CUTLASS/CuTe-DSL, and FlyDSL on ROCm), so an eager re-export would make every one of them an import-time requirement on every platform. Each public kernel is a predicate-plus-kernel pair — call `can_use_<op>(...)` first and fall back to the reference chain when it returns `False`; the kernel raises on an unsupported input rather than silently returning `None`.
 
 The package `README.md` carries a selection matrix for the cases where several kernels look interchangeable and are not. The normalization domain alone holds more than a dozen implementations that differ by numerical contract, activation layout, and backend rather than by speed.
 

docs/docs/sglang-diffusion/installation.mdx

@@ -126,3 +126,6 @@ mv python/pyproject.toml python/pyproject.toml.bak
 cp python/pyproject_other.toml python/pyproject.toml
 uv pip install -e "python[all_mps]"
 ```
+
+SGLang Diffusion uses PyTorch MPS. The `all_mps` extra also installs the SRT
+MLX backend dependencies; `SGLANG_USE_MLX` applies only to SRT serving.

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
+    normalized_shape = (x.shape[-1],)
-    x = x.contiguous().float()
     if is_rms_norm:
-        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        # ``F.rms_norm`` returns in the input dtype before a separately added
-        x_hat = x * torch.rsqrt(variance + eps)
+        # bias is applied. Promote the whole branch when a bias is present or
-    else:
+        # parameters have a different dtype, retaining the old
-        mean = x.mean(dim=-1, keepdim=True)
+        # fp32-accumulate-then-cast contract. A bias triggers promotion even
-        variance = (x - mean).pow(2).mean(dim=-1, keepdim=True)
+        # when its dtype matches the input; otherwise the normalized value is
-        x_hat = (x - mean) * torch.rsqrt(variance + eps)
+        # rounded before the affine add. Keep the native fast path for the
-    if weight is not None:
+        # common bias-free, same-dtype case.
-        x_hat = x_hat * weight.float()
+        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:
-        x_hat = x_hat + bias.float()
+                result = result + bias.float()
-    result = x_hat.to(orig_dtype)
+        else:
+            result = F.rms_norm(x, normalized_shape, weight, eps)
+            if bias is not None:
+                result = result + bias
+    else:
+        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
+    return F.rms_norm(x, (x.shape[-1],), w, eps)
-    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(
@@ -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
+NPU, Apple MPS, MUSA, CPU) and must resolve to a pure-``torch`` implementation.
-MLX-accelerated — implementation.  That choice is made once at import time,
+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
+    would pull in *all* fallback modules on every platform.  The returned shim
-    CUDA hosts.  The returned shim imports ``common.fallback_<kind>`` on its
+    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(
+        if (
-            server_args, component_name
+            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(
+        if (
-            server_args, component_name
+            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, wraps
-from functools import lru_cache
+from threading import RLock
-from typing import TYPE_CHECKING, Literal
+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
+_BRIDGE_LOCK = RLock()
-try:
-    import mlx.core as mx  # noqa: F811
-
-    _MLX_AVAILABLE = True
-except ImportError:
-    pass
 
 
-def is_mlx_available() -> bool:
+def _serialized_bridge(function: Callable[..., Any]) -> Callable[..., Any]:
-    """Return True when the ``mlx`` package can be imported."""
+    """Serialize one complete Torch/MLX crossing, including result export."""
-    return _MLX_AVAILABLE
+
+    @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:
+def _mlx_core():
-    """Return True when the user opted-in via ``SGLANG_USE_MLX=1`` **and** MLX is importable."""
+    try:
-    return bool(envs.SGLANG_USE_MLX.get()) and _MLX_AVAILABLE
+        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.
+def _get_torch_device() -> torch.device:
-# 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:
@@ -83,146 +62,377 @@ def get_torch_device() -> torch.device:
     return torch.device("cpu")
 
 
-def _get_tensor_size_bytes(array: mx.array) -> int:
+def _torch_to_mlx(
-    """Calculate the size of an MLX array in bytes.
+    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:
+    if tensor.device.type == "mps":
-        array: MLX array
+        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:
+@_serialized_bridge
-    """Check if an array is safe to transfer to MPS without hitting size limits.
+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
+    The returned views own the Torch tensor references for their entire
-    multiple temporary buffers internally. We use a conservative threshold.
+    lifetime.  No data copy is made.  Set ``synchronize=False`` only when a
-
+    surrounding operation (such as :func:`mlx_call`) performs the producer
-    Args:
+    barrier immediately before consuming the views.  This helper is
-        array: MLX array to check
+    intentionally separate from :func:`torch_to_mlx`, whose contract is an
-
+    independent MLX copy.
-    Returns:
-        True if safe to transfer to MPS, False if should stay on CPU
     """
-    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
+    array = _torch_to_mlx(tensor, copy=True)
-    if tensor.device.type != "cpu":
+    if tensor.device.type == "mps":
-        tensor = tensor.cpu()
+        # Materialize the owned copy before the caller may mutate or release
-
+        # the Torch source.
-    tensor = tensor.detach()
+        _mlx_core().eval(array)
-
+    return array
-    # 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
+@_serialized_bridge
-# to reduce the Python ↔ MLX round-trip overhead.
+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:
+    mx = _mlx_core()
-        device = get_torch_device()
+    target_device = _get_torch_device() if device is None else torch.device(device)
-    elif isinstance(device, str):
+    array = _prepare_mlx_export(array, target_device, mx)
-        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)
+    return _export_evaluated_mlx(array, target_device, mx)
-        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",
+    "MlxTensorView",
-    "use_mlx",
+    "borrow_torch_tensors",
+    "mlx_call",
+    "mlx_call_multi",
     "mlx_to_torch",
     "torch_to_mlx",
-    "get_torch_device",
 ]

test/registered/kernels/ops/diffusion/test_import_surface.py

@@ -37,6 +37,8 @@ _DEEP_IMPORT_ALLOWLIST = {
     "python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py",
     "test/registered/kernels/ops/diffusion/test_model_fast_paths.py",
     "test/registered/kernels/ops/diffusion/test_sites.py",
+    # This test exercises the pure-Torch fallback implementation directly.
+    "test/registered/unit/utils/test_diffusion_torch_fallback.py",
 }
 
 
@@ -141,7 +143,7 @@ def test_importing_the_package_does_not_import_any_leaf_module():
     """The reason ``__getattr__`` is lazy rather than a block of re-exports.
 
     The backends have disjoint, heavy, mutually-exclusive dependencies --
-    Triton (CUDA/ROCm), CUTLASS/CuTe-DSL, FlyDSL (gfx950), MLX (Apple).  If
+    Triton (CUDA/ROCm), CUTLASS/CuTe-DSL, and FlyDSL (gfx950).  If
     ``_EXPORTS`` ever degrades into eager ``from .norm.x import y`` lines, all
     of them become import-time requirements on every platform, which is how a
     CPU-only or Apple install starts failing at ``import sglang``.

test/registered/unit/hardware_backend/mlx/test_runtime.py

@@ -0,0 +1,207 @@
+"""Tests for the opt-in MLX runtime gate."""
+
+import importlib.util
+import os
+import subprocess
+import sys
+import types
+import unittest
+from importlib.metadata import PackageNotFoundError, version
+from unittest import mock
+
+import torch
+from packaging.version import Version
+
+from sglang.srt.hardware_backend.mlx import runtime
+from sglang.test.ci.ci_register import register_mlx_ci
+
+register_mlx_ci(est_time=1, suite="stage-a-unit-test-mlx")
+
+
+def _fake_mlx(version: str | None, *, metal_available: bool = True):
+    fake_mlx = types.ModuleType("mlx")
+    fake_core = types.ModuleType("mlx.core")
+    if version is not None:
+        fake_core.__version__ = version
+    fake_core.metal = types.SimpleNamespace(is_available=lambda: metal_available)
+    fake_mlx.core = fake_core
+    return fake_mlx, fake_core
+
+
+def _has_supported_mlx() -> bool:
+    try:
+        installed = Version(version("mlx"))
+    except (PackageNotFoundError, ValueError):
+        return False
+    return not installed.is_prerelease and installed >= Version("0.32.0")
+
+
+class TestMlxRuntime(unittest.TestCase):
+    def test_disabled_backend_does_not_import_mlx(self):
+        script = """
+import sys
+from sglang.srt.hardware_backend.mlx.runtime import use_mlx
+from sglang.srt.server_args import ServerArgs
+assert use_mlx() is False
+ServerArgs(model_path="dummy")
+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}",
+        )
+
+    def test_version_gates(self):
+        self.assertTrue(runtime._is_stable_series("2.13.7", (2, 13)))
+        self.assertFalse(runtime._is_stable_series("2.14.0", (2, 13)))
+        self.assertFalse(runtime._is_stable_series("2.13.1rc1", (2, 13)))
+
+        minimum = Version("0.32.0")
+        self.assertTrue(runtime._is_stable_at_least("0.32.0", minimum))
+        self.assertTrue(runtime._is_stable_at_least("0.32.0+local", minimum))
+        self.assertTrue(runtime._is_stable_at_least("0.32.1.post1", minimum))
+        self.assertTrue(runtime._is_stable_at_least("0.33.0", minimum))
+        self.assertFalse(runtime._is_stable_at_least("0.31.9", minimum))
+        self.assertFalse(runtime._is_stable_at_least("0.33.0rc1", minimum))
+        self.assertFalse(runtime._is_stable_at_least("0.33.0.dev1", minimum))
+        self.assertFalse(runtime._is_stable_at_least("unknown", minimum))
+
+    def test_unvalidated_runtime_pairs_are_rejected(self):
+        cases = (
+            ("2.14.0", "0.32.0", "stable Torch 2.13.x"),
+            ("2.13.0", "0.31.9", "MLX >= 0.32.0"),
+            ("2.13.1rc1", "0.32.0", "stable Torch 2.13.x"),
+            ("2.13.0", "0.33.0rc1", "MLX >= 0.32.0"),
+            ("2.13.0", None, "MLX unknown"),
+        )
+        for torch_version, mlx_version, message in cases:
+            with self.subTest(torch=torch_version, mlx=mlx_version):
+                fake_mlx, fake_core = _fake_mlx(mlx_version)
+                runtime._validate_runtime.cache_clear()
+                try:
+                    with (
+                        mock.patch.dict(
+                            sys.modules, {"mlx": fake_mlx, "mlx.core": fake_core}
+                        ),
+                        mock.patch.object(torch, "__version__", torch_version),
+                        mock.patch.object(
+                            torch.backends.mps, "is_available", return_value=True
+                        ),
+                        self.assertRaisesRegex(RuntimeError, message),
+                    ):
+                        runtime._validate_runtime()
+                finally:
+                    runtime._validate_runtime.cache_clear()
+
+    def test_validated_runtime_accepts_supported_stable_releases(self):
+        for mlx_version in ("0.32.9", "0.33.0", "1.0.0"):
+            with self.subTest(mlx=mlx_version):
+                fake_mlx, fake_core = _fake_mlx(mlx_version)
+                runtime._validate_runtime.cache_clear()
+                try:
+                    with (
+                        mock.patch.dict(
+                            sys.modules, {"mlx": fake_mlx, "mlx.core": fake_core}
+                        ),
+                        mock.patch.object(torch, "__version__", "2.13.7"),
+                        mock.patch.object(
+                            torch.backends.mps, "is_available", return_value=True
+                        ),
+                    ):
+                        self.assertIsNone(runtime._validate_runtime())
+                finally:
+                    runtime._validate_runtime.cache_clear()
+
+    def test_missing_mlx_has_an_actionable_error(self):
+        runtime._validate_runtime.cache_clear()
+        try:
+            with mock.patch.dict(sys.modules, {"mlx": None, "mlx.core": None}):
+                with self.assertRaisesRegex(RuntimeError, "MLX is not installed"):
+                    runtime._validate_runtime()
+        finally:
+            runtime._validate_runtime.cache_clear()
+
+    def test_unavailable_metal_devices_have_actionable_errors(self):
+        fake_mlx = types.ModuleType("mlx")
+        fake_core = types.ModuleType("mlx.core")
+        fake_core.__version__ = "0.32.0"
+        fake_mlx.core = fake_core
+
+        cases = (
+            (False, True, "PyTorch MPS device"),
+            (True, False, "MLX Metal device"),
+        )
+        for torch_mps_available, mlx_metal_available, message in cases:
+            with self.subTest(message=message):
+                fake_core.metal = types.SimpleNamespace(
+                    is_available=lambda: mlx_metal_available
+                )
+                runtime._validate_runtime.cache_clear()
+                try:
+                    with mock.patch.dict(
+                        sys.modules, {"mlx": fake_mlx, "mlx.core": fake_core}
+                    ):
+                        with mock.patch.object(torch, "__version__", "2.13.0"):
+                            with mock.patch.object(
+                                torch.backends.mps,
+                                "is_available",
+                                return_value=torch_mps_available,
+                            ):
+                                with self.assertRaisesRegex(RuntimeError, message):
+                                    runtime._validate_runtime()
+                finally:
+                    runtime._validate_runtime.cache_clear()
+
+    @unittest.skipUnless(
+        importlib.util.find_spec("mlx") is not None
+        and torch.backends.mps.is_available(),
+        "requires MLX and MPS",
+    )
+    def test_incompatible_runtime_aborts_server_args_before_dummy_shortcut(self):
+        import mlx.core as mx
+
+        runtime.use_mlx.cache_clear()
+        runtime._validate_runtime.cache_clear()
+        try:
+            with mock.patch.dict(os.environ, {"SGLANG_USE_MLX": "1"}):
+                with mock.patch.object(torch, "__version__", "2.12.1"):
+                    with self.assertRaisesRegex(RuntimeError, "stable Torch 2.13.x"):
+                        from sglang.srt.server_args import ServerArgs
+
+                        ServerArgs(model_path="dummy")
+
+                runtime.use_mlx.cache_clear()
+                runtime._validate_runtime.cache_clear()
+                with mock.patch.object(mx, "__version__", "0.31.0"):
+                    with self.assertRaisesRegex(RuntimeError, "MLX >= 0.32.0"):
+                        ServerArgs(model_path="dummy")
+        finally:
+            runtime.use_mlx.cache_clear()
+            runtime._validate_runtime.cache_clear()
+
+    @unittest.skipUnless(
+        importlib.util.find_spec("mlx") is not None
+        and torch.backends.mps.is_available()
+        and not Version(torch.__version__).is_prerelease
+        and Version(torch.__version__).release[:2] == (2, 13)
+        and _has_supported_mlx(),
+        "requires the supported MLX runtime",
+    )
+    def test_current_runtime_is_supported(self):
+        runtime._validate_runtime.cache_clear()
+        runtime._validate_runtime()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/platforms/test_mps_triton_stub.py

@@ -0,0 +1,45 @@
+"""Compatibility tests for the macOS Triton import stub."""
+
+import platform
+import subprocess
+import sys
+import unittest
+
+import torch
+from packaging.version import Version
+
+from sglang.test.ci.ci_register import register_mlx_ci
+
+register_mlx_ci(est_time=1, suite="stage-a-unit-test-mlx")
+
+
+@unittest.skipUnless(
+    sys.platform == "darwin"
+    and platform.machine() == "arm64"
+    and torch.backends.mps.is_available()
+    and Version(torch.__version__) >= Version("2.13.0"),
+    "requires Torch >= 2.13 on Apple silicon",
+)
+class TestMpsTritonStub(unittest.TestCase):
+    def test_torch_inductor_imports_after_sglang_installs_stub(self):
+        script = """
+import sglang
+from torch._inductor.runtime.triton_heuristics import _KernelType
+assert _KernelType is not None
+"""
+        completed = subprocess.run(
+            [sys.executable, "-c", script],
+            capture_output=True,
+            text=True,
+            timeout=30,
+            check=False,
+        )
+        self.assertEqual(
+            completed.returncode,
+            0,
+            msg=f"stdout={completed.stdout}\nstderr={completed.stderr}",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/utils/test_diffusion_torch_fallback.py

@@ -0,0 +1,171 @@
+"""Correctness tests for the PyTorch diffusion fallbacks used on MPS."""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.diffusion.common.fallback_torch import (
+    apply_rotary_embedding_native,
+    fuse_scale_shift_kernel_native,
+    norm_infer_native,
+    rms_norm_fn_native,
+    triton_one_pass_rms_norm_native,
+)
+from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+register_cpu_ci(est_time=1, suite="base-b-test-cpu")
+register_cpu_ci(est_time=1, suite="base-b-test-cpu-arm64")
+register_mlx_ci(est_time=1, suite="stage-a-unit-test-mlx")
+
+
+class TestDiffusionTorchFallback(unittest.TestCase):
+    @property
+    def device(self):
+        return torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+
+    def test_norm_infer_matches_reference(self):
+        dtypes = (
+            (torch.float32, torch.float16, torch.bfloat16)
+            if self.device.type == "mps"
+            else (torch.float32, torch.bfloat16)
+        )
+        for dtype in dtypes:
+            with self.subTest(dtype=dtype):
+                x = torch.randn(4, 32, device=self.device, dtype=dtype)
+                weight = torch.randn(32, device=self.device, dtype=dtype)
+                bias = torch.randn(32, device=self.device, dtype=dtype)
+
+                rms = norm_infer_native(x, weight, bias, 1e-5, is_rms_norm=True)
+                x_fp32 = x.float()
+                rms_ref = (
+                    x_fp32
+                    * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + 1e-5)
+                    * weight.float()
+                    + bias.float()
+                ).to(dtype)
+
+                layer = norm_infer_native(x, weight, bias, 1e-5)
+                mean = x_fp32.mean(-1, keepdim=True)
+                layer_ref = (
+                    (x_fp32 - mean)
+                    * torch.rsqrt(
+                        (x_fp32 - mean).square().mean(-1, keepdim=True) + 1e-5
+                    )
+                    * weight.float()
+                    + bias.float()
+                ).to(dtype)
+
+                tolerance = 0 if dtype != torch.float32 else 2e-5
+                torch.testing.assert_close(
+                    rms.cpu(), rms_ref.cpu(), rtol=tolerance, atol=tolerance
+                )
+                torch.testing.assert_close(layer.cpu(), layer_ref.cpu())
+
+                out = torch.empty_like(x)
+                returned = norm_infer_native(x, weight, bias, 1e-5, out=out)
+                self.assertIs(returned, out)
+                torch.testing.assert_close(out.cpu(), layer_ref.cpu())
+
+    def test_norm_infer_preserves_input_dtype_with_fp32_parameters(self):
+        torch.manual_seed(0)
+        for dtype in (torch.float16, torch.bfloat16):
+            with self.subTest(dtype=dtype):
+                x = torch.randn(4, 32, device=self.device, dtype=dtype)
+                weight = torch.randn(32, device=self.device, dtype=torch.float32)
+                bias = torch.randn(32, device=self.device, dtype=torch.float32)
+
+                result = norm_infer_native(x, weight, bias, 1e-5, is_rms_norm=True)
+                x_fp32 = x.float()
+                reference = (
+                    x_fp32
+                    * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + 1e-5)
+                    * weight
+                    + bias
+                ).to(dtype)
+
+                self.assertEqual(result.dtype, dtype)
+                torch.testing.assert_close(
+                    result.cpu(), reference.cpu(), rtol=0, atol=0
+                )
+
+    def test_scale_shift_matches_broadcast_reference(self):
+        x = torch.randn(2, 6, 8, device=self.device)
+        scale = torch.randn(2, 8, device=self.device)
+        shift = torch.randn(2, 8, device=self.device)
+
+        result = fuse_scale_shift_kernel_native(x, scale, shift, scale_constant=0.5)
+        reference = x * (0.5 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        torch.testing.assert_close(result.cpu(), reference.cpu())
+
+        frame_scale = torch.randn(2, 3, 1, 8, device=self.device)
+        frame_shift = torch.randn(2, 3, 1, 8, device=self.device)
+        result = fuse_scale_shift_kernel_native(x, frame_scale, frame_shift)
+        expanded_scale = (
+            frame_scale.squeeze(2).unsqueeze(2).expand(-1, -1, 2, -1).reshape_as(x)
+        )
+        expanded_shift = (
+            frame_shift.squeeze(2).unsqueeze(2).expand(-1, -1, 2, -1).reshape_as(x)
+        )
+        reference = x * (1.0 + expanded_scale) + expanded_shift
+        torch.testing.assert_close(result.cpu(), reference.cpu())
+
+    def test_rotary_embedding_matches_reference(self):
+        x = torch.randn(4, 3, 8, device=self.device)
+        cos = torch.randn(4, 4, device=self.device)
+        sin = torch.randn(4, 4, device=self.device)
+
+        result = apply_rotary_embedding_native(x, cos, sin)
+        cos_expanded = cos.unsqueeze(-2)
+        sin_expanded = sin.unsqueeze(-2)
+        x1 = x[..., ::2]
+        x2 = x[..., 1::2]
+        reference = torch.stack(
+            (
+                x1 * cos_expanded - x2 * sin_expanded,
+                x2 * cos_expanded + x1 * sin_expanded,
+            ),
+            dim=-1,
+        ).flatten(-2)
+        torch.testing.assert_close(result.cpu(), reference.cpu())
+
+        full_cos = torch.repeat_interleave(cos, 2, dim=-1)
+        full_sin = torch.repeat_interleave(sin, 2, dim=-1)
+        interleaved = apply_rotary_embedding_native(
+            x, full_cos, full_sin, interleaved=True
+        )
+        torch.testing.assert_close(interleaved.cpu(), reference.cpu())
+
+    def test_one_pass_rms_norm_matches_reference(self):
+        x = torch.randn(8, 128, device=self.device, dtype=torch.float32)
+        weight = torch.randn(128, device=self.device, dtype=torch.float32)
+        result = triton_one_pass_rms_norm_native(x, weight, 1e-6)
+        reference = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-6) * weight
+        torch.testing.assert_close(result.cpu(), reference.cpu(), rtol=2e-5, atol=2e-5)
+
+    def test_rms_norm_fn_preserves_residual_contract(self):
+        x = torch.randn(4, 32, device=self.device, dtype=torch.float32)
+        residual = torch.randn_like(x)
+        weight = torch.randn(32, device=self.device)
+        bias = torch.randn(32, device=self.device)
+
+        result, residual_out = rms_norm_fn_native(
+            x,
+            weight,
+            bias,
+            residual=residual,
+            residual_in_fp32=True,
+            zero_centered_weight=True,
+        )
+
+        combined = x.float() + residual.float()
+        reference = combined * torch.rsqrt(
+            combined.square().mean(-1, keepdim=True) + 1e-6
+        )
+        reference = reference * (weight.float() + 1.0) + bias.float()
+        torch.testing.assert_close(result.cpu(), reference.cpu(), rtol=2e-5, atol=2e-5)
+        torch.testing.assert_close(residual_out.cpu(), combined.cpu())
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/utils/test_tensor_bridge.py

@@ -0,0 +1,603 @@
+"""Correctness and sharing tests for the PyTorch/MLX tensor bridge."""
+
+import gc
+import os
+import subprocess
+import sys
+import unittest
+from importlib.metadata import PackageNotFoundError, version
+from unittest import mock
+
+import torch
+from packaging.version import Version
+
+from sglang.srt.utils.tensor_bridge import (
+    MlxTensorView,
+    borrow_torch_tensors,
+    mlx_call,
+    mlx_call_multi,
+    mlx_to_torch,
+    torch_to_mlx,
+)
+from sglang.test.ci.ci_register import register_mlx_ci
+
+register_mlx_ci(est_time=2, suite="stage-a-unit-test-mlx")
+
+
+def _has_stable_version_at_least(distribution: str, minimum: Version) -> bool:
+    try:
+        installed = Version(version(distribution))
+    except (PackageNotFoundError, ValueError):
+        return False
+    return not installed.is_prerelease and installed >= minimum
+
+
+_HAS_MLX = _has_stable_version_at_least("mlx", Version("0.32.0"))
+_HAS_SUPPORTED_RUNTIME = (
+    _HAS_MLX
+    and torch.backends.mps.is_available()
+    and not Version(torch.__version__).is_prerelease
+    and Version(torch.__version__).release[:2] == (2, 13)
+)
+
+
+class TestTensorBridgeImport(unittest.TestCase):
+    def test_import_does_not_eagerly_import_mlx(self):
+        script = """
+import sys
+from sglang.srt.utils.tensor_bridge import mlx_to_torch, torch_to_mlx
+assert mlx_to_torch is not None and torch_to_mlx 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}",
+        )
+
+
+@unittest.skipUnless(_HAS_MLX, "requires MLX >= 0.32")
+class TestTensorBridgeCpu(unittest.TestCase):
+    def test_mlx_call_multi_preserves_cpu_float64(self):
+        import mlx.core as mx
+
+        source = torch.tensor([1.25, -2.5, 4.0, 8.0], dtype=torch.float64)
+        with mock.patch.object(mx, "eval", wraps=mx.eval) as evaluate:
+            first, second = mlx_call_multi(
+                lambda x: (x + 1, x * 2),
+                source,
+                device="cpu",
+            )
+
+        evaluate.assert_called_once()
+        self.assertEqual(first.dtype, torch.float64)
+        self.assertEqual(second.dtype, torch.float64)
+        torch.testing.assert_close(first, source + 1)
+        torch.testing.assert_close(second, source * 2)
+
+    def test_mlx_call_multi_materializes_cpu_negative_strides_safely(self):
+        import mlx.core as mx
+
+        source = torch.arange(8, dtype=torch.float32)
+        with mock.patch.object(mx, "eval", wraps=mx.eval) as evaluate:
+            first, reversed_ = mlx_call_multi(
+                lambda x: (x + 1, x[::-1]), source, device="cpu"
+            )
+
+        # The ordinary graph is evaluated once; all negative-stride results
+        # share one additional materialization boundary before DLPack export.
+        self.assertEqual(evaluate.call_count, 2)
+        torch.testing.assert_close(first, source + 1)
+        torch.testing.assert_close(reversed_, source.flip(0))
+
+    def test_mlx_call_multi_rejects_invalid_target_before_work(self):
+        operation = mock.Mock()
+        with self.assertRaisesRegex(ValueError, "CPU and MPS targets"):
+            mlx_call_multi(
+                operation,
+                torch.ones(1),
+                device="cuda",
+            )
+        operation.assert_not_called()
+
+    def test_mlx_call_rejects_invalid_target_before_work(self):
+        operation = mock.Mock()
+        with self.assertRaisesRegex(ValueError, "CPU and MPS targets"):
+            mlx_call(
+                operation,
+                torch.ones(1),
+                device="cuda",
+            )
+        operation.assert_not_called()
+
+    def test_mlx_call_multi_rejects_non_mlx_outputs(self):
+        with self.assertRaisesRegex(TypeError, "outputs must be MLX arrays"):
+            mlx_call_multi(
+                lambda _x: (torch.ones(1),),
+                torch.ones(1),
+                device="cpu",
+            )
+
+
+@unittest.skipUnless(_HAS_SUPPORTED_RUNTIME, "requires MLX >= 0.32 and Torch MPS")
+class TestTensorBridgeMetalSharing(unittest.TestCase):
+    def test_common_inference_dtypes_round_trip_losslessly(self):
+        import mlx.core as mx
+
+        cases = [
+            (torch.float32, mx.float32, [0.0, 1.0, -2.0]),
+            (torch.float16, mx.float16, [0.0, 1.0, -2.0]),
+            (torch.bfloat16, mx.bfloat16, [0.0, 1.0, -2.0]),
+            (torch.int32, mx.int32, [0, 1, -2]),
+            (torch.bool, mx.bool_, [False, True, False]),
+        ]
+        for torch_dtype, mlx_dtype, values in cases:
+            with self.subTest(dtype=torch_dtype):
+                source = torch.tensor(values, device="mps", dtype=torch_dtype)
+                array = torch_to_mlx(source)
+                round_tripped = mlx_to_torch(array)
+
+                self.assertEqual(array.dtype, mlx_dtype)
+                self.assertEqual(round_tripped.dtype, torch_dtype)
+                self.assertEqual(round_tripped.device.type, "mps")
+                self.assertTrue(torch.equal(round_tripped.cpu(), source.cpu()))
+
+    def test_torch_mps_to_mlx_is_an_explicit_copy(self):
+        import mlx.core as mx
+
+        tensor = torch.arange(24, device="mps", dtype=torch.float32)
+        tensor = tensor.to(torch.bfloat16).reshape(4, 6).T
+        expected = tensor.cpu().clone()
+        array = torch_to_mlx(tensor)
+
+        tensor.zero_()
+        torch.mps.synchronize()
+        mx.eval(array)
+        round_tripped = mlx_to_torch(array, device="cpu")
+        self.assertTrue(torch.equal(round_tripped, expected))
+
+        del tensor
+        gc.collect()
+        self.assertTrue(torch.equal(round_tripped, expected))
+
+    def test_mlx_to_torch_mps_shares_storage_and_lifetime(self):
+        import mlx.core as mx
+
+        array = mx.arange(16, dtype=mx.float32).reshape(4, 4)[:, 1:3]
+        tensor = mlx_to_torch(array)
+
+        self.assertEqual(tensor.device.type, "mps")
+        tensor.zero_()
+        torch.mps.synchronize()
+        self.assertTrue(mx.all(array == 0).item())
+
+        del array
+        gc.collect()
+        self.assertEqual(torch.count_nonzero(tensor).item(), 0)
+
+    def test_mps_round_trip_uses_independent_input_storage(self):
+        tensor = torch.arange(16, device="mps", dtype=torch.float32)
+
+        round_tripped = mlx_to_torch(torch_to_mlx(tensor))
+
+        self.assertNotEqual(round_tripped.data_ptr(), tensor.data_ptr())
+
+    def test_mlx_call_keeps_zero_copy_borrows_alive(self):
+        import mlx.core as mx
+
+        tensor = torch.randn(2, 8, device="mps", dtype=torch.float32)
+        weight = torch.randn(8, device="mps", dtype=torch.float32)
+        before = tensor.cpu().clone()
+        weight_before = weight.cpu().clone()
+        reference = torch.nn.functional.rms_norm(before, (8,), weight_before, 1e-6)
+        result = mlx_call(lambda x, w: mx.fast.rms_norm(x, w, 1e-6), tensor, weight)
+
+        torch.mps.synchronize()
+        self.assertTrue(torch.equal(tensor.cpu(), before))
+        self.assertTrue(torch.equal(weight.cpu(), weight_before))
+        self.assertNotEqual(result.data_ptr(), tensor.data_ptr())
+
+        del tensor, weight
+        gc.collect()
+        torch.testing.assert_close(result.cpu(), reference)
+
+    def test_persistent_view_keeps_torch_storage_alive(self):
+        import mlx.core as mx
+
+        source = torch.arange(16, device="mps", dtype=torch.float32).reshape(4, 4)
+        view = MlxTensorView(source)
+        self.assertTrue(view.matches(source))
+        del source
+        gc.collect()
+
+        result = mlx_call(lambda x: x + 1, view, device="mps")
+        torch.testing.assert_close(
+            result.cpu(), torch.arange(1, 17, dtype=torch.float32).reshape(4, 4)
+        )
+        # The view is still the owner after the result has been exported.
+        self.assertEqual(view.array.shape, (4, 4))
+        mx.eval(view.array)
+
+    def test_mlx_call_synchronizes_persistent_view_producers(self):
+        source = torch.zeros(8, device="mps", dtype=torch.float32)
+        view = MlxTensorView(source)
+        source.fill_(3)
+        with mock.patch.object(
+            torch.mps, "synchronize", wraps=torch.mps.synchronize
+        ) as synchronize:
+            result = mlx_call(lambda x: x + 1, view, device="mps")
+        synchronize.assert_called_once_with()
+        torch.testing.assert_close(result.cpu(), torch.full((8,), 4.0))
+
+    def test_batch_borrow_syncs_once_and_preserves_sources(self):
+        first = torch.randn(4, 8, device="mps", dtype=torch.bfloat16)
+        second = torch.randn(8, 8, device="mps", dtype=torch.bfloat16)
+        first_before = first.cpu().clone()
+        second_before = second.cpu().clone()
+        with mock.patch.object(
+            torch.mps, "synchronize", wraps=torch.mps.synchronize
+        ) as synchronize:
+            views = borrow_torch_tensors(first, second)
+        synchronize.assert_called_once_with()
+        self.assertTrue(torch.equal(first.cpu(), first_before))
+        self.assertTrue(torch.equal(second.cpu(), second_before))
+        self.assertTrue(views[0].matches(first))
+        self.assertTrue(views[1].matches(second))
+
+    def test_invalid_batch_borrow_does_not_synchronize(self):
+        mps_tensor = torch.ones(1, device="mps")
+        cpu_tensor = torch.ones(1)
+        with mock.patch.object(torch.mps, "synchronize") as synchronize:
+            with self.assertRaisesRegex(ValueError, "requires MPS tensors"):
+                borrow_torch_tensors(mps_tensor, cpu_tensor)
+        synchronize.assert_not_called()
+
+    def test_mlx_call_borrows_noncontiguous_view_for_call_scope(self):
+        import mlx.core as mx
+
+        base = torch.randn(4, 6, device="mps", dtype=torch.bfloat16)
+        tensor = base.T
+        weight = torch.randn(4, device="mps", dtype=torch.bfloat16)
+        base_before = base.cpu().clone()
+        tensor_before = tensor.cpu().clone()
+
+        result = mlx_call(lambda x, w: mx.fast.rms_norm(x, w, 1e-6), tensor, weight)
+
+        torch.mps.synchronize()
+        self.assertTrue(torch.equal(base.cpu(), base_before))
+        reference = torch.nn.functional.rms_norm(
+            tensor_before, (4,), weight.cpu(), 1e-6
+        )
+        torch.testing.assert_close(result.cpu(), reference)
+
+    def test_mlx_call_multi_fences_and_evaluates_once(self):
+        """A multi-output island must not evaluate each result independently."""
+        import mlx.core as mx
+
+        source = torch.arange(8, device="mps", dtype=torch.float32)
+        source_before = source.cpu().clone()
+        captured = {}
+        events = []
+        real_synchronize = torch.mps.synchronize
+        real_eval = mx.eval
+        real_from_dlpack = torch.utils.dlpack.from_dlpack
+
+        def synchronize_then_record():
+            real_synchronize()
+            events.append("fence returned")
+
+        def operation(x):
+            events.append("operation")
+            captured["arrays"] = (x + 1, x * 2)
+            return list(captured["arrays"])
+
+        def evaluate_and_record(*arrays):
+            events.append("eval")
+            return real_eval(*arrays)
+
+        def import_and_record(*args, **kwargs):
+            events.append("dlpack")
+            return real_from_dlpack(*args, **kwargs)
+
+        with (
+            mock.patch.object(
+                torch.mps, "synchronize", side_effect=synchronize_then_record
+            ) as synchronize,
+            mock.patch.object(mx, "eval", side_effect=evaluate_and_record) as evaluate,
+            mock.patch.object(
+                torch.utils.dlpack,
+                "from_dlpack",
+                side_effect=import_and_record,
+            ) as from_dlpack,
+        ):
+            first, second = mlx_call_multi(
+                operation,
+                source,
+                device="mps",
+            )
+
+        synchronize.assert_called_once_with()
+        self.assertEqual(
+            events,
+            ["fence returned", "operation", "eval", "dlpack", "dlpack"],
+        )
+        evaluate.assert_called_once()
+        self.assertEqual(len(evaluate.call_args.args), 2)
+        self.assertEqual(from_dlpack.call_count, 2)
+        torch.testing.assert_close(first.cpu(), source_before + 1)
+        torch.testing.assert_close(second.cpu(), source_before * 2)
+        self.assertEqual(first.device.type, "mps")
+        self.assertEqual(second.device.type, "mps")
+
+        # Mutation through the Torch result remains visible from the original
+        # MLX result allocation, proving that the positive-stride export did
+        # not insert a copy.
+        first.fill_(7)
+        torch.mps.synchronize()
+        self.assertTrue(mx.all(captured["arrays"][0] == 7).item())
+
+    def test_mlx_call_multi_keeps_borrowed_inputs_alive_until_all_exports(self):
+        source = torch.arange(8, device="mps", dtype=torch.float32)
+        view = MlxTensorView(source)
+        expected = source.cpu()
+
+        first, second = mlx_call_multi(
+            lambda x: (x + 3, x - 3),
+            view,
+            device="mps",
+        )
+        del source, view
+        gc.collect()
+
+        torch.testing.assert_close(first.cpu(), expected + 3)
+        torch.testing.assert_close(second.cpu(), expected - 3)
+
+    def test_concurrent_bridge_calls_are_serialized(self):
+        """Concurrent bridge entry points must not race Metal command buffers."""
+        script = """
+from concurrent.futures import ThreadPoolExecutor
+import torch
+from sglang.srt.utils.tensor_bridge import mlx_call
+
+source = torch.arange(8, device="mps", dtype=torch.float32)
+
+def worker(iterations):
+    for _ in range(iterations):
+        result = mlx_call(lambda x: x + 1, source, device="mps")
+        assert result.device.type == "mps"
+        del result
+    return True
+
+with ThreadPoolExecutor(max_workers=2) as pool:
+    futures = [pool.submit(worker, 64), pool.submit(worker, 64)]
+    assert all(future.result() for future in futures)
+"""
+        completed = subprocess.run(
+            [sys.executable, "-c", script],
+            capture_output=True,
+            text=True,
+            timeout=30,
+            check=False,
+        )
+        self.assertEqual(
+            completed.returncode,
+            0,
+            msg=f"stdout={completed.stdout}\nstderr={completed.stderr}",
+        )
+
+    def test_mlx_call_multi_rejects_non_sequence_output(self):
+        source = torch.ones(2, device="mps", dtype=torch.float32)
+        with self.assertRaisesRegex(TypeError, "non-empty tuple or list"):
+            mlx_call_multi(lambda x: x + 1, source, device="mps")
+
+    def test_mlx_call_multi_cpu_input_does_not_fence_mps(self):
+        source = torch.arange(8, dtype=torch.float32)
+
+        with mock.patch.object(torch.mps, "synchronize") as synchronize:
+            (result,) = mlx_call_multi(
+                lambda x: (x + 1,),
+                source,
+                device="mps",
+            )
+
+        synchronize.assert_not_called()
+        torch.testing.assert_close(result.cpu(), source + 1)
+
+    def test_mlx_call_multi_sync_failure_precedes_graph_build(self):
+        source = torch.ones(1, device="mps")
+        operation = mock.Mock()
+
+        with (
+            mock.patch.object(
+                torch.mps,
+                "synchronize",
+                side_effect=RuntimeError("producer fence failed"),
+            ),
+            self.assertRaisesRegex(RuntimeError, "producer fence failed"),
+        ):
+            mlx_call_multi(
+                operation,
+                source,
+                device="mps",
+            )
+
+        operation.assert_not_called()
+
+    def test_mlx_call_multi_propagates_operation_failure(self):
+        source = torch.ones(1, device="mps")
+
+        def operation(_source):
+            raise RuntimeError("graph build failed")
+
+        with self.assertRaisesRegex(RuntimeError, "graph build failed"):
+            mlx_call_multi(
+                operation,
+                source,
+                device="mps",
+            )
+
+    def test_bridge_detaches_autograd_and_synchronizes_producers(self):
+        import mlx.core as mx
+
+        tensor = torch.arange(8, device="mps", dtype=torch.float32)
+        tensor.requires_grad_()
+        with mock.patch.object(
+            torch.mps, "synchronize", wraps=torch.mps.synchronize
+        ) as synchronize:
+            array = torch_to_mlx(tensor)
+        synchronize.assert_called_once_with()
+
+        with mock.patch.object(mx, "eval", wraps=mx.eval) as evaluate:
+            round_tripped = mlx_to_torch(array)
+        evaluate.assert_called_once_with(array)
+        self.assertFalse(round_tripped.requires_grad)
+
+    def test_torch_cpu_input_is_an_explicit_copy(self):
+        import mlx.core as mx
+
+        tensor = torch.arange(8, dtype=torch.bfloat16)
+        array = torch_to_mlx(tensor)
+        tensor.zero_()
+        mx.eval(array)
+
+        self.assertEqual(
+            array.astype(mx.float32).tolist(),
+            [float(value) for value in range(8)],
+        )
+
+    def test_torch_cpu_float64_does_not_silently_downcast(self):
+        import mlx.core as mx
+
+        tensor = torch.tensor([1.25, -2.5], dtype=torch.float64)
+        array = torch_to_mlx(tensor)
+        mx.eval(array)
+        self.assertEqual(array.dtype, mx.float64)
+        round_tripped = mlx_to_torch(array, device="cpu")
+        self.assertEqual(round_tripped.dtype, torch.float64)
+        torch.testing.assert_close(round_tripped, tensor)
+
+    def test_unsupported_cpu_dtype_fails_instead_of_narrowing(self):
+        tensor = torch.tensor([1 + 2j], dtype=torch.complex128)
+        with self.assertRaisesRegex(ValueError, "complex128"):
+            torch_to_mlx(tensor)
+
+    def test_cpu_float64_export_is_materialized_on_cpu(self):
+        import mlx.core as mx
+
+        with mx.stream(mx.cpu):
+            array = mx.array([1.25, -2.5], dtype=mx.float64)
+        tensor = mlx_to_torch(array, device="cpu")
+        self.assertEqual(tensor.device.type, "cpu")
+        self.assertEqual(tensor.dtype, torch.float64)
+        torch.testing.assert_close(
+            tensor, torch.tensor([1.25, -2.5], dtype=torch.float64)
+        )
+
+    def test_cpu_float64_positive_stride_export_remains_zero_copy(self):
+        import mlx.core as mx
+
+        with mx.stream(mx.cpu):
+            base = mx.arange(8).astype(mx.float64)
+            array = base[::2]
+        tensor = mlx_to_torch(array, device="cpu")
+        self.assertEqual(tensor.stride(), (2,))
+
+        tensor.fill_(11)
+        torch.testing.assert_close(
+            torch.utils.dlpack.from_dlpack(array.__dlpack__(dl_device=(1, 0))),
+            torch.full((4,), 11, dtype=torch.float64),
+        )
+
+    def test_cpu_export_consumes_the_dlpack_capsule_once(self):
+        """A DLPack capsule is single-use and must not be imported twice."""
+        import mlx.core as mx
+
+        with mx.stream(mx.cpu):
+            array = mx.array([1.25, -2.5], dtype=mx.float32)
+
+        with mock.patch.object(
+            torch.utils.dlpack,
+            "from_dlpack",
+            wraps=torch.utils.dlpack.from_dlpack,
+        ) as from_dlpack:
+            tensor = mlx_to_torch(array, device="cpu")
+
+        self.assertEqual(from_dlpack.call_count, 1)
+        torch.testing.assert_close(
+            tensor, torch.tensor([1.25, -2.5], dtype=torch.float32)
+        )
+
+    def test_explicit_cpu_target_shares_storage(self):
+        import mlx.core as mx
+
+        array = mx.arange(16, dtype=mx.float32).reshape(4, 4)[:, ::2]
+        tensor = mlx_to_torch(array, device="cpu")
+
+        self.assertEqual(tensor.device.type, "cpu")
+        self.assertEqual(tensor.stride(), (4, 2))
+        tensor.zero_()
+        self.assertTrue(mx.all(array == 0).item())
+
+        del array
+        gc.collect()
+        self.assertEqual(torch.count_nonzero(tensor).item(), 0)
+
+    def test_negative_stride_views_materialize_without_aborting(self):
+        script = """
+import mlx.core as mx
+import torch
+from sglang.srt.utils.tensor_bridge import mlx_call, mlx_to_torch
+
+expected = torch.arange(15, -1, -1, dtype=torch.float32)
+for target in ("cpu", "mps"):
+    array = mx.arange(16, dtype=mx.float32)[::-1]
+    tensor = mlx_to_torch(array, device=target)
+    torch.testing.assert_close(tensor.cpu(), expected)
+    tensor.zero_()
+    if target == "mps":
+        torch.mps.synchronize()
+    assert mx.array_equal(array, mx.arange(16, dtype=mx.float32)[::-1]).item()
+
+with mx.stream(mx.cpu):
+    array = mx.arange(16).astype(mx.float64)[::-1]
+tensor = mlx_to_torch(array, device="cpu")
+torch.testing.assert_close(
+    tensor, torch.arange(15, -1, -1, dtype=torch.float64)
+)
+try:
+    mlx_to_torch(array, device="mps")
+except ValueError as exc:
+    assert "float64" in str(exc)
+else:
+    raise AssertionError("float64 MLX export to MPS must fail explicitly")
+
+source = torch.arange(16, device="mps", dtype=torch.float32)
+result = mlx_call(lambda x: x[::-1], source, device="mps")
+torch.testing.assert_close(result.cpu(), expected)
+"""
+        completed = subprocess.run(
+            [sys.executable, "-c", script],
+            capture_output=True,
+            text=True,
+            timeout=30,
+            check=False,
+        )
+        self.assertEqual(
+            completed.returncode,
+            0,
+            msg=f"stdout={completed.stdout}\nstderr={completed.stderr}",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()