fix: map mmapped weights through Metal buffers instead of CPU buffers (#2037)

docs/backend.md

@@ -161,6 +161,9 @@ resident allocations. Vulkan reports exceeding total memory are rejected because
 its heap-budget subtraction can underflow. Other backends use the cap instead of
 treating such reports as zero free memory. Failed checks log the reported free and
 total memory alongside tracked weight and runtime allocations.
+With `--mmap`, device-backed mappings count toward these budgets at their full
+mapped-file size, once per device buffer even when multiple parameter blocks
+share it. Mappings retained in the loader cache continue to count.
 
 Components are considered in `diffusion`, `te`, `vae` order so that repeatedly
 used diffusion weights have priority. Each component's weights use the first

src/core/ggml_extend_backend.cpp

@@ -13,6 +13,7 @@
 #endif
 
 #include "core/util.h"
+#include "ggml-backend-impl.h"
 #include "ggml-impl.h"
 #include "stable-diffusion.h"
 
@@ -433,6 +434,24 @@ bool sd_backend_is_cpu(ggml_backend_t backend) {
     return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU;
 }
 
+ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
+                                                          void* ptr,
+                                                          size_t size,
+                                                          size_t max_tensor_size) {
+    ggml_backend_buffer_t buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size);
+    if (buffer != nullptr && buffer->context == nullptr) {
+        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device);
+        if (reg != nullptr && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) {
+            // Metal can wrap a failed mapping in a non-null buffer. Its free callback also
+            // dereferences the missing context, so only release the outer buffer.
+            buffer->iface.free_buffer = nullptr;
+            ggml_backend_buffer_free(buffer);
+            return nullptr;
+        }
+    }
+    return buffer;
+}
+
 bool sd_backend_supports_cuda_mma(ggml_backend_t backend) {
 #ifdef SD_USE_CUDA
     if (!sd_backend_is(backend, "CUDA")) {

src/core/ggml_extend_backend.h

@@ -88,6 +88,10 @@ private:
 bool sd_backend_is(ggml_backend_t backend, const std::string& name);
 bool sd_backend_is_cpu(ggml_backend_t backend);
 bool sd_backend_supports_cuda_mma(ggml_backend_t backend);
+ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
+                                                          void* ptr,
+                                                          size_t size,
+                                                          size_t max_tensor_size);
 ggml_backend_t sd_backend_cpu_init();
 bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
 ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,

src/model_loader.cpp

@@ -874,7 +874,8 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
 
 std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
                                                        std::set<std::string> ignore_tensors,
-                                                       bool writable_mmap) {
+                                                       bool writable_mmap,
+                                                       ggml_backend_dev_t device) {
     std::set<std::string> names;
     for (const auto& entry : tensors) {
         names.insert(entry.first);
@@ -896,6 +897,39 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
         if (!fdata.mmbuffer)
             continue;
 
+        // Wrapped on first use: a device buffer makes the whole file resident on that device.
+        std::shared_ptr<struct ggml_backend_buffer> file_buffer = device == nullptr ? fdata.mmbuffer : nullptr;
+        bool file_unmappable                                    = false;
+
+        auto buffer_for_file = [&]() -> ggml_backend_buffer_t {
+            if (file_buffer || file_unmappable) {
+                return file_buffer.get();
+            }
+            auto cached = fdata.device_mmbuffers.find(device);
+            if (cached != fdata.device_mmbuffers.end()) {
+                file_buffer = cached->second;
+                return file_buffer.get();
+            }
+            size_t max_tensor_size = 0;
+            for (const auto& ts : fdata.tensors) {
+                max_tensor_size = std::max(max_tensor_size, static_cast<size_t>(ts.nbytes()));
+            }
+            ggml_backend_buffer_t buf = sd_backend_dev_buffer_from_host_ptr(device,
+                                                                            fdata.mmapped->writable_data(),
+                                                                            fdata.mmapped->size(),
+                                                                            max_tensor_size);
+            if (buf == nullptr) {
+                LOG_WARN("mmap: %s cannot map '%s', loading it instead",
+                         ggml_backend_dev_name(device), fdata.path.c_str());
+                file_unmappable = true;
+                return nullptr;
+            }
+            LOG_INFO("mmap: mapped '%s' for %s", fdata.path.c_str(), ggml_backend_dev_name(device));
+            file_buffer                    = std::shared_ptr<struct ggml_backend_buffer>(buf, ggml_backend_buffer_free);
+            fdata.device_mmbuffers[device] = file_buffer;
+            return file_buffer.get();
+        };
+
         const std::vector<TensorStorage>& file_tensors = fdata.tensors;
 
         size_t file_mapped_bytes   = 0;
@@ -944,7 +978,10 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
                 continue;
             }
 
-            ggml_backend_buffer_t buf_mmap = fdata.mmbuffer.get();
+            ggml_backend_buffer_t buf_mmap = buffer_for_file();
+            if (buf_mmap == nullptr) {
+                break;
+            }
             uint8_t* mmap_data = static_cast<uint8_t*>(ggml_backend_buffer_get_base(buf_mmap));
             dst_tensor->buffer = buf_mmap;
             dst_tensor->data   = mmap_data + tensor_offset;
@@ -956,7 +993,7 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
         if (file_mapped_bytes > 0) {
             mapped_tensors += file_mapped_tensors;
             mapped_bytes += file_mapped_bytes;
-            result.push_back({fdata.mmapped, fdata.mmbuffer});
+            result.push_back({fdata.mmapped, file_buffer});
         }
     }
 
@@ -972,6 +1009,16 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
     return result;
 }
 
+std::vector<ggml_backend_buffer_t> ModelLoader::get_device_mmap_buffers() const {
+    std::vector<ggml_backend_buffer_t> buffers;
+    for (const auto& fdata : file_data) {
+        for (const auto& entry : fdata.device_mmbuffers) {
+            buffers.push_back(entry.second.get());
+        }
+    }
+    return buffers;
+}
+
 bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                                bool enable_mmap,
                                const std::set<std::string>* target_tensor_names,
@@ -1115,6 +1162,11 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                     if (dst_tensor->buffer != nullptr && dst_tensor->buffer == fdata.mmbuffer.get()) {
                         continue;
                     }
+                    if (dst_tensor->buffer != nullptr &&
+                        std::any_of(fdata.device_mmbuffers.begin(), fdata.device_mmbuffers.end(),
+                                    [&](const auto& entry) { return entry.second.get() == dst_tensor->buffer; })) {
+                        continue;
+                    }
 
                     size_t nbytes_to_read = tensor_storage.nbytes_to_read();
 

src/model_loader.h

@@ -20,6 +20,8 @@ struct ModelFileData {
     std::vector<TensorStorage> tensors;
     std::shared_ptr<MmapWrapper> mmapped;
     std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
+    // mmapped wrapped by devices that can use host memory in place (buffer_from_host_ptr)
+    std::map<ggml_backend_dev_t, std::shared_ptr<struct ggml_backend_buffer>> device_mmbuffers;
     bool is_zip;
 };
 
@@ -120,7 +122,9 @@ public:
     void process_model_files(bool enable_mmap = false, bool writable_mmap = true);
     std::vector<MmapTensorStore> mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
                                               std::set<std::string> ignore_tensors = {},
-                                              bool writable                        = true);
+                                              bool writable                        = true,
+                                              ggml_backend_dev_t device            = nullptr);
+    std::vector<ggml_backend_buffer_t> get_device_mmap_buffers() const;
     bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                       bool use_mmap                                    = false,
                       const std::set<std::string>* target_tensor_names = nullptr,

src/model_manager.cpp

@@ -780,29 +780,42 @@ bool ModelManager::validate_tensor(const TensorState& state) const {
 
 bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
                                std::vector<ParamsStorageBlock*>& created_storage_blocks) {
-    std::map<std::string, ggml_tensor*> mmap_candidates;
+    // A GPU that computes on mmapped params in place cannot address a CPU buffer, and nothing
-    std::map<std::string, TensorState*> mmap_states;
+    // stages them for it, so they are mapped through a buffer of that GPU's device.
+    struct MmapGroup {
+        std::map<std::string, ggml_tensor*> candidates;
+        std::map<std::string, TensorState*> states;
+    };
+    std::map<ggml_backend_dev_t, MmapGroup> groups;
     for (TensorState* state : states) {
         if (state == nullptr || !can_mmap_storage(*state) || state->tensor == nullptr ||
             state->tensor->data != nullptr || state->tensor->view_src != nullptr) {
             continue;
         }
-        mmap_candidates[state->name] = state->tensor;
+        ggml_backend_dev_t device = nullptr;
-        mmap_states[state->name]     = state;
+        if (!sd_backend_is_cpu(state->compute_backend) && !sd_backend_is_cpu(state->params_backend)) {
+            device = ggml_backend_get_device(state->compute_backend);
         }
-    if (mmap_candidates.empty()) {
+        MmapGroup& group              = groups[device];
-        return true;
+        group.candidates[state->name] = state->tensor;
+        group.states[state->name]     = state;
     }
 
-    auto mmap_store = model_loader_.mmap_tensors(mmap_candidates, {}, writable_mmap_);
+    for (auto& [device, group] : groups) {
+        // Device buffers wrap read-only mappings only; params that LoRAs are merged into in place
+        // are loaded instead.
+        if (device != nullptr && writable_mmap_) {
+            continue;
+        }
+        auto mmap_store = model_loader_.mmap_tensors(group.candidates, {}, writable_mmap_, device);
         if (mmap_store.empty()) {
-        return true;
+            continue;
         }
 
         auto block                = std::make_unique<ParamsStorageBlock>();
         block->mmap_tensor_stores = std::move(mmap_store);
         ParamsStorageBlock* raw   = block.get();
-    for (const auto& pair : mmap_states) {
+        for (const auto& pair : group.states) {
             TensorState* state = pair.second;
             if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) {
                 block->states.push_back(state);
@@ -813,6 +826,7 @@ bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
             params_storage_blocks_.push_back(std::move(block));
             created_storage_blocks.push_back(raw);
         }
+    }
     return true;
 }
 
@@ -1353,8 +1367,9 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe
     }
 
     size_t total_size = 0;
+    std::unordered_set<ggml_backend_buffer_t> seen;
     auto add_buffer = [&](ggml_backend_buffer_t buffer) {
-        if (buffer == nullptr || ggml_backend_buffer_is_host(buffer)) {
+        if (buffer == nullptr || ggml_backend_buffer_is_host(buffer) || !seen.insert(buffer).second) {
             return;
         }
         ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(buffer);
@@ -1371,9 +1386,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe
         total_size               = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size;
     };
 
+    // The loader may retain device mappings after their parameter blocks are released.
+    for (ggml_backend_buffer_t buffer : model_loader_.get_device_mmap_buffers()) {
+        add_buffer(buffer);
+    }
     for (const auto& block : params_storage_blocks_) {
         if (block != nullptr) {
             add_buffer(block->buffer);
+            for (const auto& store : block->mmap_tensor_stores) {
+                add_buffer(store.mmbuffer.get());
+            }
         }
     }
     for (const auto& block : compute_staging_blocks_) {