Exploring Zero-Copy mmap Loading and KV Cache Pre-Allocation for MLX on Apple Silicon

Background

Apple Silicon’s Unified Memory Architecture (UMA) enables CPU and GPU to share the same physical memory pool with up to 800 GB/s bandwidth—eliminating PCIe bottlenecks and enabling zero-copy data sharing. This makes mmap an attractive loading strategy in theory.

The key question: would MLX benefit from adopting llama.cpp’s memory strategies? We implemented both techniques to find out.

Implementation

Zero-Copy mmap Loading

We implemented an MmapReader class in MLX’s C++ core that memory-maps safetensors files and exposes offset views via Metal buffers:

MmapReader::MmapReader(std::string file_path) {
  int fd = open(file_path.c_str(), O_RDONLY);
  struct stat st;
  fstat(fd, &st);
  file_size_ = st.st_size;
  mmap_ptr_ = mmap(nullptr, file_size_,
    PROT_READ, MAP_PRIVATE, fd, 0);
  close(fd);  // mmap survives fd closure
  madvise(mmap_ptr_, file_size_, MADV_SEQUENTIAL);
}

The critical zero-copy path creates offset views into the mmap buffer with a careful alignment check. Safetensors does not guarantee element-aligned offsets, so we verify offset % itemsize == 0 before allowing zero-copy access. Unaligned offsets fall back to memcpy.

if (reader_->is_mmap() && !swap_endianness_) {
  auto mmap_reader = dynamic_pointer_cast<MmapReader>(reader_);
  if (mmap_reader && (offset_ % out.itemsize() == 0)) {
    auto metal_buf = mmap_reader->get_metal_buffer();
    auto parent = array(metal_buf, {1}, uint8,
      [reader_ref](Buffer) { /* prevent dealloc */ });
    out.copy_shared_buffer(parent, strides, flags,
      out.size(), offset_ / out.itemsize());
    return;  // Zero-copy success
  }
}
// Fallback: allocate + memcpy for unaligned offsets

KV Cache Pre-Allocation

We added optional max_context_length parameter for upfront KV cache allocation. When specified, keys and values tensors are allocated for the full context window at startup, with mx.eval() forcing immediate physical allocation. If the sequence exceeds pre-allocated length, it gracefully falls back to dynamic expansion.

class KVCache(_BaseCache):
    step = 256
    def __init__(self, n_kv_heads=0, head_dim=0,
                 max_context_length=0, dtype=mx.float16):
        self.offset = 0
        if max_context_length > 0 and n_kv_heads > 0:
            L = ((max_context_length + 255) // 256) * 256
            self.keys = mx.zeros(
                (1, n_kv_heads, L, head_dim), dtype=dtype)
            self.values = mx.zeros(
                (1, n_kv_heads, L, head_dim), dtype=dtype)
            mx.eval(self.keys, self.values)  # Force alloc

Quantized dtype Bug Fix

During implementation, we discovered that QuantizedLinear.weight.dtype returns uint32 (the packed storage type) rather than the working precision. This caused the KV cache to be allocated in uint32 instead of float16. The fix uses scales.dtype instead, which stores dequantization factors in the correct working dtype.

Evaluation

Loading Speed: Standard vs. mmap

Results are highly inconsistent. Three out of eight models load slower with mmap. The 20.65x speedup on Qwen3-8B-8bit is an outlier—likely due to the standard loader hitting a pathological pread pattern. For most models in the 2–6 GB range, MLX’s standard loading is already efficient enough.

Inference Impact

* Memory nearly doubled—mmap region coexists with materialized copies due to lazy evaluation.

KV Cache Pre-Allocation

Pre-allocation does not improve throughput. Generation speed stays within <5% of dynamic mode. First token latency increases 5–7% due to forced mx.eval() at startup. Memory overhead is predictable: ~0.25 GB for 2048 tokens, ~0.5 GB for 4096 tokens.

Why MLX’s Existing Design Works Well

The 14B-6bit Anomaly

The largest model tested (11.18 GB) loaded slower with mmap (0.69x). At this size, the mmap approach must fault in an extremely large number of pages. The OS page fault handler becomes the bottleneck—each 16 KB page requires a kernel trap, page table update, and TLB insertion. The standard loader’s buffered reads, which batch these operations, prove more efficient at scale.

Recommendations

Conclusion

We explored whether llama.cpp’s memory management strategies—mmap loading and KV cache pre-allocation—could improve MLX’s performance on Apple Silicon. Our systematic benchmarks across eight models reveal inconsistent results: mmap helps for some larger models but hurts for small ones and can increase peak memory; KV pre-allocation achieves flat memory but offers no throughput benefit.

The conclusion: MLX’s existing memory management is already well-designed for its target use case. The design choices that differ from llama.cpp are not oversights but appropriate adaptations to MLX’s lazy evaluation model and the safetensors ecosystem. During this work, we also identified and fixed a quantized model dtype inference bug—arguably the most practically useful contribution.

Citation

@article{atomgradient2026optmlx,
  title={Exploring Zero-Copy mmap Loading and KV Cache
         Pre-Allocation for MLX on Apple Silicon},
  author={AtomGradient},
  year={2026},
  url={https://github.com/AtomGradient/OptMLX}
}