Add mmap_weights: zero-copy file-backed weight arrays

docs/src/python/ops.rst

@@ -132,6 +132,7 @@ Operations
    mean
    median
    meshgrid
+   mmap_weights
    min
    minimum
    moveaxis

mlx/io.h

@@ -34,6 +34,15 @@ load(std::shared_ptr<io::Reader> in_stream, StreamOrDevice s = {});
 /** Load array from file in .npy format */
 MLX_API array load(std::string file, StreamOrDevice s = {});
 
+/** Zero-copy view of a tensor stored in a file: the data is an mmap'd,
+ * page-cache-backed region wrapped via the allocator's no-copy path.
+ * The returned array is read-only by contract. */
+MLX_API array mmap_weights(
+    const std::string& file,
+    int64_t byte_offset,
+    Shape shape,
+    Dtype dtype);
+
 /** Load array map from .safetensors file format */
 MLX_API SafetensorsLoad
 load_safetensors(std::shared_ptr<io::Reader> in_stream, StreamOrDevice s = {});

mlx/io/CMakeLists.txt

@@ -1,4 +1,5 @@
-target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp)
+target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp
+                           ${CMAKE_CURRENT_SOURCE_DIR}/mmap.cpp)
 
 if(MLX_BUILD_SAFETENSORS)
   target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/safetensors.cpp)

mlx/io/mmap.cpp

@@ -0,0 +1,109 @@
+// Copyright © 2026 Apple Inc.
+
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <unistd.h>
+
+#include <stdexcept>
+
+#include "mlx/allocator.h"
+#include "mlx/io.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+array mmap_weights(
+    const std::string& file,
+    int64_t byte_offset,
+    Shape shape,
+    Dtype dtype) {
+  size_t nelem = 1;
+  for (auto s : shape) {
+    if (s < 0) {
+      throw std::invalid_argument("[mmap_weights] negative dimension");
+    }
+    nelem *= static_cast<size_t>(s);
+  }
+  size_t nbytes = nelem * size_of(dtype);
+  if (nbytes == 0) {
+    throw std::invalid_argument("[mmap_weights] empty tensor");
+  }
+
+  int fd = open(file.c_str(), O_RDONLY);
+  if (fd < 0) {
+    throw std::invalid_argument("[mmap_weights] cannot open: " + file);
+  }
+  struct stat st;
+  if (fstat(fd, &st) != 0 || byte_offset < 0 ||
+      static_cast<size_t>(byte_offset) + nbytes >
+          static_cast<size_t>(st.st_size)) {
+    close(fd);
+    throw std::invalid_argument("[mmap_weights] range out of bounds: " + file);
+  }
+
+  // Map from the enclosing page boundary; the tensor begins `delta` bytes in.
+  const size_t page = static_cast<size_t>(getpagesize());
+  const size_t base_off = (static_cast<size_t>(byte_offset) / page) * page;
+  const size_t delta = static_cast<size_t>(byte_offset) - base_off;
+  const size_t map_len = ((delta + nbytes + page - 1) / page) * page;
+
+  void* base = mmap(
+      nullptr,
+      map_len,
+      PROT_READ,
+      MAP_SHARED,
+      fd,
+      static_cast<off_t>(base_off));
+  close(fd); // the mapping keeps its own reference
+  if (base == MAP_FAILED) {
+    throw std::runtime_error("[mmap_weights] mmap failed: " + file);
+  }
+
+  // Metal setBuffer offsets must be aligned to the element size (mlx's own
+  // sliced arrays rely on the same). NOTE: stock safetensors gives NO
+  // alignment guarantee — real checkpoints put tensors at odd offsets — so
+  // callers typically need an aligned store.
+  if (delta % size_of(dtype) != 0) {
+    munmap(base, map_len);
+    throw std::invalid_argument(
+        "[mmap_weights] byte_offset must be aligned to the element size");
+  }
+
+  auto buf = allocator::make_buffer(base, map_len);
+  if (buf.ptr() == nullptr) {
+    munmap(base, map_len);
+    throw std::runtime_error(
+        "[mmap_weights] make_buffer failed (Metal unavailable or mapping "
+        "rejected)");
+  }
+
+  array out(shape, dtype, nullptr, {});
+  Strides strides(shape.size());
+  int64_t acc = 1;
+  for (int i = static_cast<int>(shape.size()) - 1; i >= 0; --i) {
+    strides[i] = acc;
+    acc *= shape[i];
+  }
+  array::Flags flags{};
+  flags.contiguous = true;
+  flags.row_contiguous = true;
+  flags.col_contiguous = shape.size() <= 1;
+  out.set_data(
+      buf,
+      nelem,
+      std::move(strides),
+      flags,
+      static_cast<int64_t>(delta),
+      // The buffer wraps an mmap'd file region: release the Metal wrapper
+      // (never recycle into the allocator cache) and drop the mapping. Runs
+      // on a Metal completion-handler thread — pure C++ only.
+      [base, map_len](allocator::Buffer b) {
+        allocator::release(b);
+        munmap(base, map_len);
+      });
+  out.set_status(array::Status::available);
+  return out;
+}
+
+} // namespace mlx::core

python/src/ops.cpp

@@ -4429,6 +4429,32 @@ void init_ops(nb::module_& m) {
             **kwargs (arrays): Arrays to be saved. Each array will be saved
               with the associated keyword as the output file name.
       )pbdoc");
+  m.def(
+      "mmap_weights",
+      [](const std::string& file,
+         int64_t offset,
+         const nb::object& shape,
+         const mx::Dtype& dtype) {
+        return mx::mmap_weights(file, offset, to_shape(shape), dtype);
+      },
+      "file"_a,
+      "offset"_a,
+      "shape"_a,
+      "dtype"_a,
+      R"pbdoc(
+        Zero-copy view of a tensor stored in a file (Metal only).
+
+        The array's data is an mmap'd, page-cache-backed region wrapped via
+        newBufferWithBytesNoCopy: no bytes are copied at creation, pages are
+        faulted in by first access (CPU or GPU) and stay wired while any
+        reference to the array lives. Read-only by contract.
+
+        Args:
+            file (str): Path to the file.
+            offset (int): Byte offset of the tensor data (4-byte aligned).
+            shape (list(int)): Tensor shape (row-major, contiguous).
+            dtype (Dtype): Element type.
+      )pbdoc");
   m.def(
       "load",
       &mlx_load_helper,

python/tests/test_mmap_weights.py

@@ -0,0 +1,126 @@
+# Copyright © 2026 Apple Inc.
+
+import gc
+import json
+import os
+import struct
+import tempfile
+import unittest
+
+import mlx.core as mx
+import mlx_tests
+
+
+def _write_aligned(tmpdir, tensors, align=64):
+    """Pack tensors into a flat 64-byte-aligned bin via a scratch
+    safetensors (gives exact serialized bytes per tensor)."""
+    st = os.path.join(tmpdir, "src.safetensors")
+    mx.save_safetensors(st, tensors)
+    with open(st, "rb") as f:
+        n = struct.unpack("<Q", f.read(8))[0]
+        hdr = json.loads(f.read(n))
+        blob = f.read()
+    path = os.path.join(tmpdir, "packed.bin")
+    index = {}
+    with open(path, "wb") as out:
+        for name, spec in hdr.items():
+            if name == "__metadata__":
+                continue
+            b0, b1 = spec["data_offsets"]
+            out.write(b"\0" * ((-out.tell()) % align))
+            index[name] = (out.tell(), spec["shape"])
+            out.write(blob[b0:b1])
+    return path, index
+
+
+class TestMmapWeights(mlx_tests.MLXTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls._tmp = tempfile.TemporaryDirectory()
+        mx.random.seed(7)
+        cls.tensors = {
+            "f32": mx.random.normal((123, 77)),
+            "bf16": mx.random.normal((300, 400)).astype(mx.bfloat16),
+            "f16": mx.random.normal((64, 32)).astype(mx.float16),
+            "u32": (mx.random.uniform(shape=(50, 9)) * 1e6).astype(mx.uint32),
+        }
+        mx.eval(cls.tensors)
+        cls.path, cls.index = _write_aligned(cls._tmp.name, cls.tensors)
+        cls.mx_dtypes = {
+            "f32": mx.float32,
+            "bf16": mx.bfloat16,
+            "f16": mx.float16,
+            "u32": mx.uint32,
+        }
+
+    @classmethod
+    def tearDownClass(cls):
+        cls._tmp.cleanup()
+
+    def _view(self, name):
+        off, shape = self.index[name]
+        return mx.mmap_weights(self.path, off, shape, self.mx_dtypes[name])
+
+    def _assert_bit_equal(self, a, b):
+        if a.dtype in (mx.bfloat16, mx.float16):
+            a, b = a.view(mx.uint16), b.view(mx.uint16)
+        elif a.dtype == mx.float32:
+            a, b = a.view(mx.uint32), b.view(mx.uint32)
+        self.assertTrue(mx.array_equal(a, b).item())
+
+    def test_bit_equality_all_dtypes(self):
+        for name, ref in self.tensors.items():
+            with self.subTest(dtype=name):
+                self._assert_bit_equal(self._view(name), ref)
+
+    def test_cpu_and_gpu_backends(self):
+        for dev in (mx.cpu,) + ((mx.gpu,) if mx.metal.is_available() else ()):
+            with self.subTest(device=dev):
+                with mx.stream(dev):
+                    v = self._view("f32")
+                    s = (v * 2).sum()
+                    mx.eval(s)
+                    self.assertAlmostEqual(
+                        s.item(), (self.tensors["f32"] * 2).sum().item(), places=3
+                    )
+
+    def test_ops_through_view(self):
+        v = self._view("f32")
+        ref = self.tensors["f32"]
+        out_v = mx.softmax(v @ v.T, axis=-1)
+        out_r = mx.softmax(ref @ ref.T, axis=-1)
+        mx.eval(out_v, out_r)
+        self._assert_bit_equal(out_v, out_r)
+
+    def test_donation_window_never_mutates_mapping(self):
+        # Drop the view's last reference before eval — if the buffer were
+        # donated, the graph would write into the read-only mapping.
+        expect = mx.array(self._view("bf16"))
+        mx.eval(expect)
+        v = self._view("bf16")
+        out = mx.abs(-(v + mx.array(1.0, dtype=mx.bfloat16)))
+        del v
+        mx.eval(out)
+        gc.collect()
+        self._assert_bit_equal(self._view("bf16"), expect)
+
+    def test_lifecycle_many_cycles(self):
+        for _ in range(200):
+            v = self._view("f32")
+            mx.eval(v.sum())
+            del v
+        gc.collect()
+        self._assert_bit_equal(self._view("f32"), self.tensors["f32"])
+
+    def test_errors(self):
+        off, shape = self.index["f32"]
+        with self.assertRaises(Exception):
+            mx.mmap_weights(self.path + ".missing", 0, [4], mx.float32)
+        with self.assertRaises(Exception):  # out of bounds
+            mx.mmap_weights(self.path, off, [10**6, 10**6], mx.float32)
+        with self.assertRaises(Exception):  # misaligned for element size
+            mx.mmap_weights(self.path, off + 1, shape, mx.float32)
+
+
+if __name__ == "__main__":
+    unittest.main()