DarthAffe/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-09-24 20:20:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Dong Wang 2026-09-23 18:53:49 +02:00 committed by GitHub
parent 2a4ebba818
commit 500ef5fa7c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 139 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
docs/backend.md
View File

@ -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

19
src/core/ggml_extend_backend.cpp
View File

@ -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")) {

4
src/core/ggml_extend_backend.h
View File

@ -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,

58
src/model_loader.cpp
View File

@ -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();

6
src/model_loader.h
View File

@ -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,

42
src/model_manager.cpp
View File

@ -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;
std::map<std::string, TensorState*> mmap_states;
// A GPU that computes on mmapped params in place cannot address a CPU buffer, and nothing
// 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;
mmap_states[state->name] = state;
ggml_backend_dev_t device = nullptr;
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()) {
return true;
MmapGroup& group = groups[device];
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_) {