upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-09-29 09:28:14 -05: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-24 00:53:49 +08:00
committed by GitHub
parent 2a4ebba818
commit 500ef5fa7c
6 changed files with 139 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/core/ggml_extend_backend.cpp
+19
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")) {
src/core/ggml_extend_backend.h
+4
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,
src/model_loader.cpp
+58 -6
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,10 +978,13 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
continue;
}
ggml_backend_buffer_t buf_mmap = fdata.mmbuffer.get();
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;
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;
file_mapped_bytes += tensor_size;
file_mapped_tensors++;
@@ -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
+5 -1
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,
src/model_manager.cpp
+50 -28
View File
@@ -780,38 +780,52 @@ 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;
}
if (mmap_candidates.empty()) {
return true;
}
auto mmap_store = model_loader_.mmap_tensors(mmap_candidates, {}, writable_mmap_);
if (mmap_store.empty()) {
return true;
}
auto block = std::make_unique<ParamsStorageBlock>();
block->mmap_tensor_stores = std::move(mmap_store);
ParamsStorageBlock* raw = block.get();
for (const auto& pair : mmap_states) {
TensorState* state = pair.second;
if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) {
block->states.push_back(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);
}
MmapGroup& group = groups[device];
group.candidates[state->name] = state->tensor;
group.states[state->name] = state;
}
if (!block->states.empty()) {
params_storage_blocks_.push_back(std::move(block));
created_storage_blocks.push_back(raw);
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()) {
continue;
}
auto block = std::make_unique<ParamsStorageBlock>();
block->mmap_tensor_stores = std::move(mmap_store);
ParamsStorageBlock* raw = block.get();
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);
}
}
if (!block->states.empty()) {
params_storage_blocks_.push_back(std::move(block));
created_storage_blocks.push_back(raw);
}
}
return true;
}
@@ -1353,15 +1367,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe
}
size_t total_size = 0;
auto add_buffer = [&](ggml_backend_buffer_t buffer) {
if (buffer == nullptr || ggml_backend_buffer_is_host(buffer)) {
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) || !seen.insert(buffer).second) {
return;
}
ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(buffer);
auto split_devices = split_buffer_devices_.find(buffer_type);
const bool on_device = split_devices == split_buffer_devices_.end()
? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device
: std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) {
? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device
: std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) {
return ggml_backend_get_device(entry.first) == compute_device;
});
if (!on_device) {
@@ -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_) {