Skip to content

Add Metal DLPack zero-copy sharing#3531

Draft
XXXXRT666 wants to merge 9 commits into
ml-explore:mainfrom
XXXXRT666:metal-dlpack-zero-copy-draft
Draft

Add Metal DLPack zero-copy sharing#3531
XXXXRT666 wants to merge 9 commits into
ml-explore:mainfrom
XXXXRT666:metal-dlpack-zero-copy-draft

Conversation

@XXXXRT666
Copy link
Copy Markdown
Contributor

@XXXXRT666 XXXXRT666 commented May 11, 2026

Proposed changes

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

This PR depends on #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)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant