upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-09-25 15:38:00 -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/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_) {