Add Metal DLPack zero-copy sharing

[XXXXRT666]

Proposed changes

This draft adds zero-copy Metal DLPack sharing for MLX arrays and PyTorch MPS tensors.

This PR builds on the merged DLPack import PR #3495 and requires nanobind support.

The main changes are:

• Import Metal DLPack arrays by wrapping the underlying Metal buffer instead of copying through CPU.
• Export MLX arrays to Metal DLPack using the MLX Metal buffer and DLPack byte_offset.
• Add mx.from_dlpack(..., copy=...) controls for Metal DLPack inputs.
• Keep mx.array(...) zero-copy for Metal DLPack inputs unless an explicit different dtype is requested.
• Document the explicit synchronization requirements between PyTorch MPS and MLX.

The shared lifetime is tied to the exported or imported buffer. Synchronization remains explicit: PyTorch writes require torch.mps.synchronize() before MLX reads, and MLX writes require mx.eval(...) before PyTorch reads.

For MLX arrays exported to PyTorch, later MLX updates may rebind the MLX array to a new buffer while the PyTorch tensor continues to reference the exported buffer.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[megacpp]

Hi @XXXXRT666 — read through this PR after @awni redirected us here from #3548. The nb::ndarray<nb::ro, nb::c_contig> approach over the in-flight nanobind PR (#1338) is materially cleaner than the manual capsule parsing we had in our downstream PoC, and lifting is_host_accessible() into mlx/allocator.h is the right level of abstraction. Closed the RFC; happy with this being the path forward for #2848.

Wanted to offer some testing help that complements the PyTorch MPS bring-up you have:

We maintain a downstream TileLang fork (https://github.com/DatasunriseOU/tilelang) whose TVM-FFI bridge exports kDLMetal DLPack capsules for tensors backed by id<MTLBuffer>. That gives a non-PyTorch Metal DLPack producer that exercises the same import path you're adding here. Specifically it covers:

• kDLMetal producers that do not require the PyTorch-MPS workaround for __dlpack__ — exercises the import path directly.
• Round-trip mx.array → DLPack → TVM-FFI Metal kernel → DLPack → mx.array zero-copy.
• Custom Metal kernels (via mlx.fast.metal_kernel) consuming an imported mx.array whose underlying MTLBuffer was allocated outside MLX.
• storageMode matrix (we hit Shared and Managed; Private is the obvious edge case the spec needs to nail down — your is_host_accessible() decision likely answers this implicitly but worth a sanity check from the producer side).
• byte_offset != 0 cases that we'd previously rejected outright in our PoC — your PR seems to handle these via byte_offset-aware import; happy to write a TileLang-side test for it.

If useful, once the PR converges I can:

1. Pull this branch into our TileLang test matrix and report back on any rough edges (CI on macOS Metal hardware).
2. Send a minimal standalone repro (no TileLang dependency) for any of the above scenarios if you'd like them as additions to python/tests/test_array.py.
3. Beta-test the mx.from_dlpack(..., copy=...) semantics against the dtype-mismatch-shares case (002360faa) once the API stabilizes.

Tag me here when you'd like input — no rush, just don't want this to slip past once it's review-ready.

(For the orthogonal mx.empty() piece that was also in our PoC, opened it as a separate issue per @awni's guidance.)

[McPatate]

This would be super cool if it landed for end to end "0-copy" support in safetensors! I'm working (safetensors/safetensors#767) on adding reading bytes from disk in raw MTLBuffers, which can then be handed to the framework via dlpack with 0-copy. Works well with torch, would be happy to see that land in mlx!

Also, support for byte_offset !=0 would be nice (already in the PR but commenting to notify it's useful) since we can go one step further: currently the mps path is pread -> MTLBuffer, but that goes through kernel pages before hitting userspace buffer. Having byte_offset non zero support would enable mmap-ing the file and creating MTLBuffers that reference specific slices of the mmap, which would demand-fault pages from disk into the page cache on first access and give userspace access directly, leaving only the disk -> kernel-side copy.

Quick question on the dl_tensor.data convention, torch's mps treats it as id<MTLBuffer>, passing the contents segfaults. Curious to know which direction MLX will be taking, as it impacts us downstream!

[XXXXRT666]

Quick question on the dl_tensor.data convention, torch's mps treats it as id<MTLBuffer>, passing the contents segfaults. Curious to know which direction MLX will be taking, as it impacts us downstream!

https://dmlc.github.io/dlpack/latest/c_api.html#c.DLTensor.data

The data pointer points to the allocated data. This will be CUDA device pointer, cl_mem handle in OpenCL, or id<MTLBuffer> for Metal.

[XXXXRT666]

One API question: should mx.array(...) always copy DLPack inputs, and should zero-copy / copy control live in mx.asarray(..., copy=...) instead?

That would match the mental model used by NumPy/PyTorch more closely: array creates a new array, while asarray may avoid a copy depending on copy. In that design, mx.from_dlpack(..., copy=...) could remain the explicit DLPack entry point, while mx.array(torch_mps_tensor) would not unexpectedly share the underlying Metal buffer by default.

[zcbenz]

Also there is no need to use std::move since shapes are no longer heap allocated, we don't migrate the old code but try to use the new pattern in new code.

[XXXXRT666]

Removed now

[XXXXRT666]

I think this should be handled explicitly rather than relying on argument evaluation order. Since a.shape() and {std::move(a)} are evaluated as separate function arguments, C++ may evaluate {std::move(a)} first. In that case, a.shape() would be called after a has been moved from.

Linux and Windows CPU CI failed because of this

[zcbenz]

One API question: should mx.array(...) always copy DLPack inputs, and should zero-copy / copy control live in mx.asarray(..., copy=...) instead?

That would match the mental model used by NumPy/PyTorch more closely: array creates a new array, while asarray may avoid a copy depending on copy. In that design, mx.from_dlpack(..., copy=...) could remain the explicit DLPack entry point, while mx.array(torch_mps_tensor) would not unexpectedly share the underlying Metal buffer by default.

I think this is very good design.

[XXXXRT666]

mlx 0.31.2

mlx_to_torch  (median µs / call, lower is better)
  shape               current → mps   dlpack → mps   dlpack stay-on-cpu
  16K  f32                      378            191                   16
  1M   f32                      404            402                   17
  16M  f32 (64MB)              2717           2501                   16
  1M  bf16                      276 n/a (TypeError)      n/a (TypeError)
  16M bf16 (32MB)              1302 n/a (TypeError)      n/a (TypeError)

torch_to_mlx  (median µs / call, lower is better)
  shape                current (.cpu+numpy)      dlpack (from MPS)
  16K  f32                              175                    172
  1M   f32                              691                    666
  16M  f32 (64MB)                      8216                   8470
  1M  bf16                              522                    422
  16M bf16 (32MB)                      2023                   2037

mlx 0.32.0.dev20260523+4e8decde9

mlx_to_torch  (median µs / call, lower is better)
  shape               current → mps   dlpack → mps   dlpack stay-on-cpu
  16K  f32                      237             18                   15
  1M   f32                      371             16                   17
  16M  f32 (64MB)              2688             15                   15
  1M  bf16                      279             19                   16
  16M bf16 (32MB)              1340             16                   14

torch_to_mlx  (median µs / call, lower is better)
  shape                current (.cpu+numpy)      dlpack (from MPS)
  16K  f32                              189                     16
  1M   f32                              676                     16
  16M  f32 (64MB)                      8806                     17
  1M  bf16                              484                     16
  16M bf16 (32MB)                      2106                     16

benchmark function

# --- mlx -> torch candidates ---------------------------------------------------


def mlx_to_torch_current(arr: mx.array, device: torch.device) -> torch.Tensor:
    arr = mx.contiguous(arr)
    mx.eval(arr)
    buf = memoryview(arr)
    dtype_map = {
        mx.float32: torch.float32,
        mx.float16: torch.float16,
        mx.bfloat16: torch.bfloat16,
    }
    t = torch.frombuffer(buf, dtype=dtype_map[arr.dtype]).reshape(arr.shape)
    if device.type == "mps":
        t = t.to(device)
    return t


def mlx_to_torch_dlpack_mps(arr: mx.array, device: torch.device) -> torch.Tensor:
    mx.eval(arr)
    t = torch.from_dlpack(arr)
    if device.type == "mps":
        t = t.to(device)
    return t


def mlx_to_torch_dlpack_cpu(arr: mx.array, device: torch.device) -> torch.Tensor:
    """Force a CPU-typed capsule via `dl_device=(kDLCPU, 0)` (Phase 2+).
    Falls back to the no-kwarg form for builds that don't accept it."""
    mx.eval(arr)
    try:
        cap = arr.__dlpack__(dl_device=(1, 0))
    except TypeError:
        # Older builds: zero-arg lambda. Capsule is already kDLCPU there.
        cap = arr.__dlpack__()
    return torch.from_dlpack(cap)


# --- torch -> mlx candidates ---------------------------------------------------


def torch_to_mlx_current(t: torch.Tensor) -> mx.array:
    if t.device.type != "cpu":
        t = t.cpu()
    t = t.detach()
    if t.dtype == torch.bfloat16:
        return mx.array(t)
    return mx.array(t.numpy())


def torch_to_mlx_dlpack(t: torch.Tensor) -> mx.array:
    """Use mx.from_dlpack when the API exists; old MLX falls back to CPU copy."""
    if hasattr(mx, "from_dlpack"):
        return mx.from_dlpack(t)
    return mx.array(t.detach().cpu())

[XXXXRT666]

I think this is very good design.

I updated the PR to follow this design: mx.array(...) now copies DLPack inputs, while mx.asarray(..., copy=...) and mx.from_dlpack(..., copy=...) provide the copy-control paths.

[zcbenz]

This basically looks good to me, thanks for the nice work!

[XXXXRT666]

I just realized that nanobind does not handle c_contig arrays very well. For torch.Tensor, it calls contiguous(), but it does not perform torch.accelerator.synchronization afterwards.

[zcbenz]

This is going to be really slow, would it be practical to preserve the original strides and and just do memcpy?

out.set_data(
    mx::allocator::malloc(data_size * itemsize),
    data_size,
    strides);
std::copy(src, src + data_size, dst);

[XXXXRT666]

Yes, I checked that PyTorch preserves the strides provided by DLPack on import, including for tensors that are not non_overlapping_and_dense. Strides only change if PyTorch performs a copy afterwards, for example when torch.as_tensor needs to change device or dtype, or when using torch.tensor, which always copies. In those cases, tensors that are not non_overlapping_and_dense are materialized into a compact layout.

I think we do not need to do that in MLX, and can simply preserve the original strides.

[XXXXRT666]

A non_overlapping_and_dense tensor refers to a tensor that is not like x[::2] or torch.tensor([1]).expand(3)

[zcbenz]

Can you collaborate on this? If the input is truly contiguous we shouldn't need this for copying, otherwise the flag was set for non-contiguous input.

[XXXXRT666]

This makes the astype copy force the array to be row_contiguous. However, I think this is no longer necessary if we preserve the strides and do not require imported arrays to always be row_contiguous

[XXXXRT666]

CPU imports now copy the underlying storage span and preserve the DLPack strides.

For Metal:

• copy=False preserves the DLPack strides by wrapping the original buffer.
• copy=True preserves strides for non_overlapping_and_dense layouts, since those stay on the vector copy path.
• Layouts that are not non_overlapping_and_dense may still go through the general copy path and be materialized to a compact layout.

[XXXXRT666]

The earlier CI failures came from three independent issues.

First, the new array::set_data(..., offset, deleter) signature accidentally broke existing CUDA call sites that still used the old argument order, causing CUDA compilation failures.

Second, astype had a use-after-move bug: the array was moved into the input list while its shape was also read in the same expression, and the C++ argument evaluation order made this crash on some platforms.

Third, the macOS Metal validation failure was from a test mutating the original PyTorch MPS tensor after exporting/importing it through DLPack; I think it's a problem with PyTorch

METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -c 'import torch; x=torch.tensor([1.0,2.0,3.0], device="mps"); x+=1; torch.mps.synchronize(); print(x.cpu())'

will cause the same error without MLX