diff --git a/docs/docs/hardware-platforms/apple_metal.mdx b/docs/docs/hardware-platforms/apple_metal.mdx index 188bebe27..c835a7b5e 100644 --- a/docs/docs/hardware-platforms/apple_metal.mdx +++ b/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 -toolchain (`clang++`, the Metal framework headers, and `xcrun`). These ship -with the **Xcode Command Line Tools**, which cannot be installed via `pip`: +The MLX runtime requires Apple Silicon with macOS 14 or newer, stable PyTorch +2.13.x, and stable MLX 0.32.0 or newer. The `srt_mps` extra installs PyTorch +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 -available. You can verify with `xcode-select -p && xcrun --find metal`. +Verify the compiler with `xcrun -sdk macosx metal --version`. ## 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 \ --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 \ --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 \ --disable-cuda-graph \ --num-prompts 1 diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx index 2d79ff276..eaf5447c8 100644 --- a/docs/docs/references/environment_variables.mdx +++ b/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. + diff --git a/docs/docs/sglang-diffusion/environment_variables.mdx b/docs/docs/sglang-diffusion/environment_variables.mdx index 6c61676bf..89ece5419 100644 --- a/docs/docs/sglang-diffusion/environment_variables.mdx +++ b/docs/docs/sglang-diffusion/environment_variables.mdx @@ -127,7 +127,7 @@ description: "Configure SGLang diffusion behavior with environment variables." - +
SGLANG_USE_MLX not setSet to 1 to enable MLX fused Metal kernels for norm ops on MPSSRT only: enables the MLX serving backend. It has no effect on SGLang Diffusion, which uses PyTorch MPS.
diff --git a/docs/docs/sglang-diffusion/fused_kernels.mdx b/docs/docs/sglang-diffusion/fused_kernels.mdx index ee3a9ec3f..07a0facce 100644 --- a/docs/docs/sglang-diffusion/fused_kernels.mdx +++ b/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_(...)` 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_(...)` 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. diff --git a/docs/docs/sglang-diffusion/installation.mdx b/docs/docs/sglang-diffusion/installation.mdx index d949256b8..b33473346 100644 --- a/docs/docs/sglang-diffusion/installation.mdx +++ b/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. diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index e4952cecd..6cb15d633 100755 --- a/python/pyproject_other.toml +++ b/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 = [ diff --git a/python/sglang/_platform_stubs.py b/python/sglang/_platform_stubs.py index ccb3c2c15..8e20db7a1 100644 --- a/python/sglang/_platform_stubs.py +++ b/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 diff --git a/python/sglang/benchmark/one_batch.py b/python/sglang/benchmark/one_batch.py index 6f5e67360..8f06a93bf 100644 --- a/python/sglang/benchmark/one_batch.py +++ b/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( diff --git a/python/sglang/kernels/ops/diffusion/README.md b/python/sglang/kernels/ops/diffusion/README.md index ba8f053a2..8d017789e 100644 --- a/python/sglang/kernels/ops/diffusion/README.md +++ b/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 diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 59e1a9113..279f8aa70 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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. diff --git a/python/sglang/kernels/ops/diffusion/common/fallback_mps.py b/python/sglang/kernels/ops/diffusion/common/fallback_mps.py deleted file mode 100644 index 31298b434..000000000 --- a/python/sglang/kernels/ops/diffusion/common/fallback_mps.py +++ /dev/null @@ -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 diff --git a/python/sglang/kernels/ops/diffusion/common/fallback_torch.py b/python/sglang/kernels/ops/diffusion/common/fallback_torch.py index de115747b..87c59ddb7 100644 --- a/python/sglang/kernels/ops/diffusion/common/fallback_torch.py +++ b/python/sglang/kernels/ops/diffusion/common/fallback_torch.py @@ -8,6 +8,7 @@ implementations replace the Triton kernels from typing import Optional import torch +import torch.nn.functional as F from torch import Tensor @@ -69,20 +70,36 @@ def norm_infer_native( out: Optional[Tensor] = None, ) -> Tensor: """Native fallback for norm_infer (layer norm / rms norm inference).""" - orig_dtype = x.dtype - x = x.contiguous().float() + normalized_shape = (x.shape[-1],) if is_rms_norm: - variance = x.pow(2).mean(dim=-1, keepdim=True) - x_hat = x * torch.rsqrt(variance + eps) + # ``F.rms_norm`` returns in the input dtype before a separately added + # bias is applied. Promote the whole branch when a bias is present or + # parameters have a different dtype, retaining the old + # fp32-accumulate-then-cast contract. A bias triggers promotion even + # when its dtype matches the input; otherwise the normalized value is + # rounded before the affine add. Keep the native fast path for the + # common bias-free, same-dtype case. + needs_fp32 = bias is not None or ( + weight is not None and weight.dtype != x.dtype + ) + if needs_fp32: + result = F.rms_norm( + x.float(), + normalized_shape, + weight.float() if weight is not None else None, + eps, + ) + if bias is not None: + result = result + bias.float() + else: + result = F.rms_norm(x, normalized_shape, weight, eps) + if bias is not None: + result = result + bias else: - mean = x.mean(dim=-1, keepdim=True) - variance = (x - mean).pow(2).mean(dim=-1, keepdim=True) - x_hat = (x - mean) * torch.rsqrt(variance + eps) - if weight is not None: - x_hat = x_hat * weight.float() - if bias is not None: - x_hat = x_hat + bias.float() - result = x_hat.to(orig_dtype) + result = F.layer_norm(x, normalized_shape, weight, bias, eps) + # Match the original fallback and Triton kernel contract even when a + # higher-precision weight or bias promotes PyTorch's intermediate result. + result = result.to(x.dtype) if out is not None: out.copy_(result) return out @@ -93,12 +110,7 @@ def triton_one_pass_rms_norm_native( x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6 ) -> torch.Tensor: """Native fallback for triton_one_pass_rms_norm.""" - shape = x.shape - orig_dtype = x.dtype - x = x.contiguous().float() - variance = x.pow(2).mean(dim=-1, keepdim=True) - x_hat = x * torch.rsqrt(variance + eps) - return (x_hat * w.float()).to(orig_dtype).view(shape) + return F.rms_norm(x, (x.shape[-1],), w, eps) def rms_norm_fn_native( @@ -130,13 +142,13 @@ def rms_norm_fn_native( residual_out_val = x.to(torch.float32 if residual_in_fp32 else orig_dtype) else: residual_out_val = None - variance = x.pow(2).mean(dim=-1, keepdim=True) - x_hat = x * torch.rsqrt(variance + eps) if weight is not None: w = weight.float() if zero_centered_weight: w = w + 1.0 - x_hat = x_hat * w + else: + w = None + x_hat = F.rms_norm(x, (x.shape[-1],), w, eps) if bias is not None: x_hat = x_hat + bias.float() final_dtype = out_dtype if out_dtype is not None else orig_dtype diff --git a/python/sglang/kernels/ops/diffusion/common/platform.py b/python/sglang/kernels/ops/diffusion/common/platform.py index 01a744ad6..7b763a1b0 100644 --- a/python/sglang/kernels/ops/diffusion/common/platform.py +++ b/python/sglang/kernels/ops/diffusion/common/platform.py @@ -1,8 +1,8 @@ """Platform predicates and the import-time fallback selector. Several diffusion Triton kernels have no Triton on the live device (Ascend -NPU, Apple MPS, MUSA, CPU) and must resolve to a pure-``torch`` — or -MLX-accelerated — implementation. That choice is made once at import time, +NPU, Apple MPS, MUSA, CPU) and must resolve to a pure-``torch`` implementation. +That choice is made once at import time, which used to mean a hand-rolled four-branch ``if`` block repeated in every such module, each importing ``current_platform`` directly. @@ -75,8 +75,8 @@ def lazy_fallback(kind: str, name: str) -> Callable: """Name a fallback without importing its module. ``select_impl`` is handed every candidate at once, so a plain import here - would pull in *all* fallback modules on every platform -- including MLX on - CUDA hosts. The returned shim imports ``common.fallback_`` on its + would pull in *all* fallback modules on every platform. The returned shim + imports ``common.fallback_`` on its first call instead, which for the unselected candidates never happens. """ diff --git a/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py b/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py index 668cc2372..09b2e2a74 100644 --- a/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py +++ b/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"), ) diff --git a/python/sglang/kernels/ops/diffusion/norm/norm_triton.py b/python/sglang/kernels/ops/diffusion/norm/norm_triton.py index 22f6d3aa7..4a535e23a 100644 --- a/python/sglang/kernels/ops/diffusion/norm/norm_triton.py +++ b/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"), ) diff --git a/python/sglang/kernels/ops/diffusion/norm/rmsnorm_onepass_triton.py b/python/sglang/kernels/ops/diffusion/norm/rmsnorm_onepass_triton.py index 1aba75f51..23cd9663b 100644 --- a/python/sglang/kernels/ops/diffusion/norm/rmsnorm_onepass_triton.py +++ b/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"), ) diff --git a/python/sglang/kernels/ops/diffusion/rope/rotary_triton.py b/python/sglang/kernels/ops/diffusion/rope/rotary_triton.py index dc8830b91..f4fbd3645 100644 --- a/python/sglang/kernels/ops/diffusion/rope/rotary_triton.py +++ b/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"), ) diff --git a/python/sglang/multimodal_gen/README.md b/python/sglang/multimodal_gen/README.md index 05ad5cf35..d458b9fda 100644 --- a/python/sglang/multimodal_gen/README.md +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index 6f2e18949..935ece08e 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py @@ -152,8 +152,11 @@ class TransformerLoader(ComponentLoader): def customized_load_kwargs_for_component( self, server_args: ServerArgs, component_name: str ) -> dict[str, bool]: - if current_platform.is_mps() and self._is_component_set_as_layerwise_load( - server_args, component_name + if ( + current_platform.is_mps() + and server_args.should_configure_layerwise_offload_for_lazy_component( + component_name + ) ): logger.info( "Loading %s on CPU first for MPS layerwise offload", component_name diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 1d5cb93a3..94cda9011 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py @@ -150,8 +150,11 @@ class VAELoader(ComponentLoader): def customized_load_kwargs_for_component( self, server_args: ServerArgs, component_name: str ) -> dict[str, bool]: - if current_platform.is_mps() and self._is_component_set_as_layerwise_load( - server_args, component_name + if ( + current_platform.is_mps() + and server_args.should_configure_layerwise_offload_for_lazy_component( + component_name + ) ): logger.info( "Loading %s on CPU first for MPS layerwise offload", component_name diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_import_isolation.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_import_isolation.py new file mode 100644 index 000000000..2d7f7856e --- /dev/null +++ b/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() diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index 58b27439a..32100a825 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/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, diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py index 94c4b7367..bd0dbc828 100644 --- a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py +++ b/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" diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index eda2818be..8722782f3 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/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__) diff --git a/python/sglang/srt/hardware_backend/mlx/profiler.py b/python/sglang/srt/hardware_backend/mlx/profiler.py index c009ee149..f5c64a80f 100644 --- a/python/sglang/srt/hardware_backend/mlx/profiler.py +++ b/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__) diff --git a/python/sglang/srt/hardware_backend/mlx/runtime.py b/python/sglang/srt/hardware_backend/mlx/runtime.py new file mode 100644 index 000000000..aaea0bc94 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index b1bba1fdb..7c4e7166c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 097a2adba..612fb1a68 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/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: diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index 5f3980b8a..1f0764931 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 53316be92..dccfe8885 100644 --- a/python/sglang/srt/server_args.py +++ b/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 diff --git a/python/sglang/srt/utils/tensor_bridge.py b/python/sglang/srt/utils/tensor_bridge.py index f59c9d175..302b00396 100644 --- a/python/sglang/srt/utils/tensor_bridge.py +++ b/python/sglang/srt/utils/tensor_bridge.py @@ -1,78 +1,57 @@ # Copied and adapted from: https://github.com/vllm-project/vllm-metal # SPDX-License-Identifier: Apache-2.0 -"""Tensor bridge between MLX and PyTorch. +"""Tensor bridge between MLX and PyTorch on Apple silicon. -Provides zero-copy conversion when possible using Apple Silicon's unified memory. +The MLX backend requires MLX >= 0.32 and PyTorch >= 2.13. Ordinary +``torch_to_mlx`` conversion creates an independent MLX allocation. The +zero-copy ``mlx_call`` helper is available for a complete MLX operation and +keeps all borrowed DLPack inputs alive until the result has been evaluated. +This lifetime boundary matters because MLX may donate a borrowed input buffer +to a lazy operation. + +Bridge entry points are serialized because Torch and MLX use different stream +abstractions over the same Metal command queues. The lock covers producer +fencing, MLX evaluation, and DLPack import. It cannot cover arbitrary MPS +work outside the function, so callers must serialize any overlapping use or +mutation of source and returned MPS tensors as well. """ from __future__ import annotations -import logging -from functools import lru_cache -from typing import TYPE_CHECKING, Literal +from functools import lru_cache, wraps +from threading import RLock +from typing import TYPE_CHECKING, Any, Callable, Literal import torch -from sglang.srt.environ import envs - if TYPE_CHECKING: import mlx.core as mx -logger = logging.getLogger(__name__) -_MLX_AVAILABLE: bool = False -try: - import mlx.core as mx # noqa: F811 - - _MLX_AVAILABLE = True -except ImportError: - pass +_BRIDGE_LOCK = RLock() -def is_mlx_available() -> bool: - """Return True when the ``mlx`` package can be imported.""" - return _MLX_AVAILABLE +def _serialized_bridge(function: Callable[..., Any]) -> Callable[..., Any]: + """Serialize one complete Torch/MLX crossing, including result export.""" + + @wraps(function) + def wrapper(*args: Any, **kwargs: Any) -> Any: + with _BRIDGE_LOCK: + return function(*args, **kwargs) + + return wrapper @lru_cache(maxsize=1) -def use_mlx() -> bool: - """Return True when the user opted-in via ``SGLANG_USE_MLX=1`` **and** MLX is importable.""" - return bool(envs.SGLANG_USE_MLX.get()) and _MLX_AVAILABLE +def _mlx_core(): + try: + import mlx.core as mx + except ImportError: + raise RuntimeError("The MLX tensor bridge requires MLX >= 0.32.0") from None + return mx -# MPS has a 4GB (2^32 bytes) limit for MPSTemporaryNDArray allocations. -# Metal may allocate multiple temporary buffers internally, so we use a -# conservative threshold of 1GB to avoid hitting the limit. -# See: https://github.com/anthropics/vllm-metal/issues/43 -_MPS_SAFE_SIZE_BYTES = 1 << 30 # 1GB - -# MLX to PyTorch dtype mapping -# TODO(perf): float64 is CPU-only in MLX (see ml-explore/mlx#1843). -# When the target device is GPU/MPS we should auto-downcast float64 → float32 -# to avoid a runtime error; when the target is CPU we can keep float64. -# For now float64 is omitted from the mapping so it hits the ValueError -# fallback in mlx_to_torch(). -MLX_TO_TORCH_DTYPE = ( - { - mx.float32: torch.float32, - mx.float16: torch.float16, - mx.bfloat16: torch.bfloat16, - mx.int32: torch.int32, - mx.int64: torch.int64, - mx.int16: torch.int16, - mx.int8: torch.int8, - mx.uint8: torch.uint8, - mx.bool_: torch.bool, - } - if _MLX_AVAILABLE - else {} -) - -# PyTorch to MLX dtype mapping -TORCH_TO_MLX_DTYPE = {v: k for k, v in MLX_TO_TORCH_DTYPE.items()} - - -def get_torch_device() -> torch.device: +def _get_torch_device() -> torch.device: """Get the PyTorch device for Metal/MPS. Returns: @@ -83,146 +62,377 @@ def get_torch_device() -> torch.device: return torch.device("cpu") -def _get_tensor_size_bytes(array: mx.array) -> int: - """Calculate the size of an MLX array in bytes. +def _torch_to_mlx( + tensor: torch.Tensor, + *, + copy: bool, + synchronize: bool = True, +) -> mx.array: + """Convert one tensor, optionally borrowing its MPS allocation.""" + mx = _mlx_core() + tensor = tensor.detach() - Args: - array: MLX array + if tensor.device.type == "mps": + if synchronize: + # Torch and MLX do not share stream state on Metal. + torch.mps.synchronize() + return mx.asarray(tensor, copy=copy) + if tensor.device.type == "cpu": + # CPU tensors always get MLX-owned storage. In particular, do not + # expose a NumPy/memoryview alias whose lifetime is controlled by the + # caller. + if tensor.dtype == torch.complex128: + raise ValueError( + "MLX 0.32 does not support complex128; convert the Torch tensor " + "to complex64 explicitly" + ) + # MLX 0.32 does not support float64 on its default Metal stream. Keep + # the dtype by constructing this uncommon CPU value on the CPU stream + # instead of silently downcasting it to float32. + if tensor.dtype == torch.float64: + with mx.stream(mx.cpu): + return mx.array(tensor, dtype=mx.float64) + return mx.array(tensor) + raise ValueError( + f"The MLX tensor bridge supports CPU and MPS tensors, got {tensor.device}" + ) - Returns: - Size in bytes + +class MlxTensorView: + """A lifetime-bound, zero-copy MLX view of a Torch MPS tensor. + + The view deliberately retains a detached Torch tensor *and* the imported + MLX array. Holding only the array is insufficient: a later parameter + replacement or garbage collection could invalidate the borrowed storage + while MLX still has a lazy graph referring to it. This class is intended + for immutable inference weights; construct a new view after replacing the + source storage. """ - return array.size * array.dtype.size + + __slots__ = ("torch_tensor", "array") + + def __init__(self, tensor: torch.Tensor, *, synchronize: bool = True): + with _BRIDGE_LOCK: + owner = tensor.detach() + if owner.device.type != "mps": + raise ValueError( + f"MlxTensorView requires a Torch MPS tensor, got {owner.device}" + ) + if synchronize: + torch.mps.synchronize() + self.torch_tensor = owner + self.array = _torch_to_mlx(owner, copy=False, synchronize=False) + + @classmethod + def _from_synchronized(cls, tensor: torch.Tensor) -> MlxTensorView: + view = object.__new__(cls) + owner = tensor.detach() + if owner.device.type != "mps": + raise ValueError( + f"MlxTensorView requires a Torch MPS tensor, got {owner.device}" + ) + view.torch_tensor = owner + view.array = _torch_to_mlx(owner, copy=False, synchronize=False) + return view + + def matches(self, tensor: torch.Tensor) -> bool: + """Return whether ``tensor`` still refers to this borrowed storage.""" + owner = tensor.detach() + return ( + owner.device == self.torch_tensor.device + and owner.dtype == self.torch_tensor.dtype + and owner.shape == self.torch_tensor.shape + and owner.stride() == self.torch_tensor.stride() + and owner.data_ptr() == self.torch_tensor.data_ptr() + ) -def _is_safe_for_mps(array: mx.array) -> bool: - """Check if an array is safe to transfer to MPS without hitting size limits. +@_serialized_bridge +def borrow_torch_tensors( + *tensors: torch.Tensor, synchronize: bool = True +) -> tuple[MlxTensorView, ...]: + """Borrow one or more Torch MPS tensors, optionally synchronizing once. - MPS has a 4GB limit for MPSTemporaryNDArray, but Metal may allocate - multiple temporary buffers internally. We use a conservative threshold. - - Args: - array: MLX array to check - - Returns: - True if safe to transfer to MPS, False if should stay on CPU + The returned views own the Torch tensor references for their entire + lifetime. No data copy is made. Set ``synchronize=False`` only when a + surrounding operation (such as :func:`mlx_call`) performs the producer + barrier immediately before consuming the views. This helper is + intentionally separate from :func:`torch_to_mlx`, whose contract is an + independent MLX copy. """ - return _get_tensor_size_bytes(array) < _MPS_SAFE_SIZE_BYTES + detached = tuple(tensor.detach() for tensor in tensors) + if any(tensor.device.type != "mps" for tensor in detached): + devices = ", ".join(str(tensor.device) for tensor in detached) + raise ValueError(f"borrow_torch_tensors requires MPS tensors, got {devices}") + if synchronize and detached: + torch.mps.synchronize() + return tuple(MlxTensorView._from_synchronized(tensor) for tensor in detached) +@_serialized_bridge def torch_to_mlx(tensor: torch.Tensor) -> mx.array: - """Convert PyTorch tensor to MLX array. + """Convert a PyTorch tensor to an independent MLX array. - Uses numpy as an intermediate to enable zero-copy on unified memory. + MPS inputs are copied inside the unified Metal device. Use ``mlx_call`` + when a complete operation needs zero-copy MPS input imports; it owns the + borrowed MLX arrays for the complete lazy operation. Args: - tensor: PyTorch tensor (can be on any device) + tensor: PyTorch CPU or MPS tensor. Returns: MLX array with the same data """ - # Move to CPU if on MPS for numpy conversion - if tensor.device.type != "cpu": - tensor = tensor.cpu() - - tensor = tensor.detach() - - # Note: numpy does not support bfloat16. - if tensor.dtype == torch.bfloat16: - return mx.array(tensor) - - return mx.array(tensor.numpy()) + array = _torch_to_mlx(tensor, copy=True) + if tensor.device.type == "mps": + # Materialize the owned copy before the caller may mutate or release + # the Torch source. + _mlx_core().eval(array) + return array -# TODO(perf): accept a list/batch of arrays and convert them in one pass -# to reduce the Python ↔ MLX round-trip overhead. +@_serialized_bridge +def mlx_call( + operation: Callable[..., mx.array], + *tensors: torch.Tensor | MlxTensorView, + device: torch.device | Literal["mps", "cpu"] | None = None, +) -> torch.Tensor: + """Run one MLX operation with zero-copy Torch MPS input imports. + + The imported MLX arrays remain strongly referenced until + :func:`mlx_to_torch` evaluates and exports ``operation``'s result. Keep + the operation inside this call; returning a lazy MLX result for later use + or stashing a borrowed input through a callback side effect would escape + the borrow scope. The caller must also serialize any overlapping MPS work + outside this function, including use or mutation of source and returned + tensors. The operation may allocate its own output normally. + """ + mx = _mlx_core() + target_device = _get_torch_device() if device is None else torch.device(device) + if target_device.type not in {"cpu", "mps"}: + raise ValueError( + f"The MLX tensor bridge supports CPU and MPS targets, got {target_device}" + ) + detached = tuple( + tensor.detach() for tensor in tensors if isinstance(tensor, torch.Tensor) + ) + if any(tensor.device.type == "mps" for tensor in detached) or any( + isinstance(tensor, MlxTensorView) for tensor in tensors + ): + torch.mps.synchronize() + borrowed: tuple[Any, ...] = tuple( + ( + tensor.array + if isinstance(tensor, MlxTensorView) + else _torch_to_mlx(tensor.detach(), copy=False, synchronize=False) + ) + for tensor in tensors + ) + # MLX does not support float64 on the Metal stream. Keep an explicitly + # requested CPU call on the CPU stream when a borrowed input carries that + # dtype; otherwise even constructing the lazy result would fail before the + # export preparation below can move it. + if target_device.type == "cpu" and any( + array.dtype == mx.float64 for array in borrowed + ): + with mx.stream(mx.cpu): + result = operation(*borrowed) + else: + result = operation(*borrowed) + output = mlx_to_torch(result, device=target_device) + # Keep the imported MLX objects (and any MlxTensorView Torch owners) alive + # through lazy result evaluation and DLPack export. + _ = borrowed + return output + + +def _prepare_mlx_export( + array: mx.array, + target_device: torch.device, + mx: Any, +) -> mx.array: + """Prepare one lazy MLX result for the requested Torch target. + + This intentionally does not evaluate the result. Callers which export + several results should prepare every result first and then issue one + shared ``mx.eval`` boundary. + """ + if target_device.type not in {"cpu", "mps"}: + raise ValueError( + f"The MLX tensor bridge supports CPU and MPS targets, got {target_device}" + ) + + if target_device.type == "mps" and array.dtype == mx.float64: + raise ValueError( + "MLX float64 arrays cannot be exported to a Torch MPS tensor; " + "use float32/bfloat16 or request device='cpu'" + ) + + return array + + +def _has_negative_stride(array: mx.array) -> bool: + """Return whether an evaluated MLX array has a DLPack-incompatible view.""" + # PyTorch's DLPack importer aborts the process for negative strides. MLX + # exposes the evaluated layout through the buffer protocol, so inspect it + # before handing the capsule to PyTorch. + with memoryview(array) as view: + return any(stride < 0 for stride in (view.strides or ())) + + +def _export_evaluated_mlx( + array: mx.array, + target_device: torch.device, + mx: Any, + *, + materialize_negative: bool = True, +) -> torch.Tensor: + """Export an already-evaluated MLX result through one DLPack capsule. + + Negative-stride results are materialized here as a safety fallback. The + normal (contiguous/positive-stride) path performs no copy and no extra + evaluation; :func:`mlx_call_multi` batches any required materialization + evaluations for all outputs together. + """ + if materialize_negative and _has_negative_stride(array): + materialize_stream = mx.cpu if target_device.type == "cpu" else mx.gpu + array = mx.contiguous(array, stream=materialize_stream) + mx.eval(array) + + if target_device.type == "cpu": + # MLX owns CPU-accessible unified memory. Request a CPU DLPack view + # explicitly rather than importing on MPS and copying back. + dlpack = array.__dlpack__(dl_device=(1, 0), copy=False) + return torch.utils.dlpack.from_dlpack(dlpack) + return torch.utils.dlpack.from_dlpack(array) + + +@_serialized_bridge +def mlx_call_multi( + operation: Callable[..., tuple[mx.array, ...]], + *tensors: torch.Tensor | MlxTensorView, + device: torch.device | Literal["mps", "cpu"] | None = None, +) -> tuple[torch.Tensor, ...]: + """Run one MLX operation and export all of its outputs as Torch tensors. + + ``operation`` must return a non-empty flat ``tuple`` or ``list`` of MLX + arrays. All Torch/MPS inputs are fenced once before import, and all + ordinary outputs are evaluated with one ``mx.eval(*outputs)`` call before + being exported through DLPack. The imported arrays and detached Torch + owners remain local until every output capsule has been consumed, which is + required when MLX lazily donates a borrowed input buffer. Callers must + serialize any overlapping MPS work outside this function, including use + or mutation of source and returned tensors. + + CPU targets retain the same float64 and negative-stride safeguards as + :func:`mlx_to_torch`. A negative-stride output necessarily needs one + additional materialization evaluation; contiguous MPS model outputs take + the single-evaluation, zero-copy path. + + """ + mx = _mlx_core() + target_device = _get_torch_device() if device is None else torch.device(device) + if target_device.type not in {"cpu", "mps"}: + raise ValueError( + f"The MLX tensor bridge supports CPU and MPS targets, got {target_device}" + ) + + detached = tuple( + tensor.detach() for tensor in tensors if isinstance(tensor, torch.Tensor) + ) + needs_mps_fence = any(tensor.device.type == "mps" for tensor in detached) or any( + isinstance(tensor, MlxTensorView) for tensor in tensors + ) + if needs_mps_fence: + torch.mps.synchronize() + + borrowed: tuple[Any, ...] = tuple( + ( + tensor.array + if isinstance(tensor, MlxTensorView) + else _torch_to_mlx(tensor.detach(), copy=False, synchronize=False) + ) + for tensor in tensors + ) + + if target_device.type == "cpu" and any( + array.dtype == mx.float64 for array in borrowed + ): + with mx.stream(mx.cpu): + result = operation(*borrowed) + else: + result = operation(*borrowed) + if not isinstance(result, (tuple, list)) or not result: + raise TypeError( + "mlx_call_multi operation must return a non-empty tuple or list of MLX arrays" + ) + arrays = tuple(result) + if any(not isinstance(array, mx.array) for array in arrays): + raise TypeError("mlx_call_multi outputs must be MLX arrays") + + # Prepare all outputs before crossing the one shared MLX evaluation + # boundary. This is the key difference from calling mlx_to_torch in a + # loop, which would fence/evaluate every result separately. + arrays = tuple(_prepare_mlx_export(array, target_device, mx) for array in arrays) + mx.eval(*arrays) + + # DLPack cannot represent negative strides. Materialize all such outputs + # together so even this safety path has one additional evaluation boundary + # rather than one boundary per result. + negative = tuple(_has_negative_stride(array) for array in arrays) + if any(negative): + materialized = [] + for array, needs_materialization in zip(arrays, negative): + if needs_materialization: + stream = mx.cpu if target_device.type == "cpu" else mx.gpu + array = mx.contiguous(array, stream=stream) + materialized.append(array) + arrays = tuple(materialized) + mx.eval(*(array for array, needs in zip(arrays, negative) if needs)) + + outputs = tuple( + _export_evaluated_mlx(array, target_device, mx, materialize_negative=False) + for array in arrays + ) + # Keep both borrowed MLX views and their Torch owners alive through the + # final DLPack import. (The local remains live until function return.) + _ = borrowed + return outputs + + +@_serialized_bridge def mlx_to_torch( array: mx.array, device: torch.device | Literal["mps", "cpu"] | None = None, - already_contiguous: bool = False, ) -> torch.Tensor: """Convert MLX array to PyTorch tensor. - Uses numpy as an intermediate to enable zero-copy on unified memory. + MLX arrays with PyTorch-compatible strides share their unified-memory + allocation through DLPack, including explicit CPU views. Negative-stride + views are materialized because PyTorch's DLPack importer cannot represent + them safely. MLX is evaluated before the handoff because the frameworks do + not share stream state. Only CPU and MPS targets are supported; other + target devices are rejected. Args: array: MLX array device: Target PyTorch device (default: MPS if available) - already_contiguous: Skip contiguity check if array is known contiguous Returns: PyTorch tensor with the same data """ - if device is None: - device = get_torch_device() - elif isinstance(device, str): - device = torch.device(device) - - # Use memoryview for zero-copy conversion (bypasses numpy for bfloat16) - # reference: https://github.com/ml-explore/mlx/issues/403 - torch_dtype = MLX_TO_TORCH_DTYPE.get(array.dtype) - if torch_dtype is not None: - if already_contiguous: - # Fast path: skip contiguity check, single eval - mx.eval(array) - buffer = memoryview(array) - else: - # MLX views / non-contiguous arrays expose a non-contiguous buffer (or - # sometimes no usable buffer), which `torch.frombuffer` can't consume. - # Make contiguous first, then eval once - array = mx.contiguous(array) - mx.eval(array) - buffer = memoryview(array) - - tensor = torch.frombuffer(buffer, dtype=torch_dtype).reshape(array.shape) - else: - # Fallback to numpy path for unsupported dtypes - raise ValueError(f"Unsupported MLX dtype: {array.dtype}") - - # Move to target device, but check for MPS size limits first - if device.type == "mps": - if _is_safe_for_mps(array): - tensor = tensor.to(device) - else: - # Large tensor - keep on CPU to avoid MPS 4GB limit crash - # See: https://github.com/anthropics/vllm-metal/issues/43 - logger.debug( - "Tensor too large for MPS (%d bytes > %d limit), keeping on CPU", - _get_tensor_size_bytes(array), - _MPS_SAFE_SIZE_BYTES, - ) - elif device.type != "cpu": - tensor = tensor.to(device) - - return tensor - - -def sync_mlx() -> None: - """Synchronize MLX operations. - - Call this before converting MLX arrays to ensure all operations complete. - """ - # Prefer an explicit MLX barrier when available; otherwise force evaluation. - # `mx.eval([])` is a no-op, so we evaluate a tiny scalar as a safe fallback. - try: - mx.synchronize() - except (AttributeError, TypeError): - mx.eval(mx.array(0, dtype=mx.int32)) - - -def sync_torch() -> None: - """Synchronize PyTorch MPS operations. - - Call this before converting PyTorch tensors to ensure all operations complete. - """ - if torch.backends.mps.is_available(): - torch.mps.synchronize() + mx = _mlx_core() + target_device = _get_torch_device() if device is None else torch.device(device) + array = _prepare_mlx_export(array, target_device, mx) + mx.eval(array) + return _export_evaluated_mlx(array, target_device, mx) __all__ = [ - "is_mlx_available", - "use_mlx", + "MlxTensorView", + "borrow_torch_tensors", + "mlx_call", + "mlx_call_multi", "mlx_to_torch", "torch_to_mlx", - "get_torch_device", ] diff --git a/test/registered/kernels/ops/diffusion/test_import_surface.py b/test/registered/kernels/ops/diffusion/test_import_surface.py index 482524581..9eef11a19 100644 --- a/test/registered/kernels/ops/diffusion/test_import_surface.py +++ b/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``. diff --git a/test/registered/unit/hardware_backend/mlx/test_runtime.py b/test/registered/unit/hardware_backend/mlx/test_runtime.py new file mode 100644 index 000000000..5c56bb517 --- /dev/null +++ b/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() diff --git a/test/registered/unit/platforms/test_mps_triton_stub.py b/test/registered/unit/platforms/test_mps_triton_stub.py new file mode 100644 index 000000000..1fc8733bd --- /dev/null +++ b/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() diff --git a/test/registered/unit/utils/test_diffusion_torch_fallback.py b/test/registered/unit/utils/test_diffusion_torch_fallback.py new file mode 100644 index 000000000..624c77978 --- /dev/null +++ b/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() diff --git a/test/registered/unit/utils/test_tensor_bridge.py b/test/registered/unit/utils/test_tensor_bridge.py new file mode 100644 index 000000000..5b23cfc7c --- /dev/null +++ b/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()