ggml-vulkan: pin mmap CPU weights for faster H2D uploads
Export register_host_buffer/unregister via the backend reg so the existing GGML_CUDA_REGISTER_HOST path in llama-model-loader pins mmap'd expert weights. Imports the host pages through VK_EXT_external_memory_host, which populates device->pinned_memory and lets the existing pinned fast path in ggml_vk_buffer_write_2d_async DMA straight from system RAM instead of bouncing through the staging buffer + blocking host memcpy. Assisted-by: opencode
This commit is contained in:
@@ -17917,11 +17917,74 @@ static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg,
|
|||||||
return devices[device];
|
return devices[device];
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Import an mmap-backed host region as a Vulkan pinned buffer via
|
||||||
|
// VK_EXT_external_memory_host so H2D uploads DMA straight from system RAM
|
||||||
|
// instead of bouncing through the staging buffer + host memcpy. Mirrors the
|
||||||
|
// GGML_CUDA_REGISTER_HOST path; populates device->pinned_memory, which the
|
||||||
|
// existing pinned fast path in ggml_vk_buffer_write_2d_async looks up.
|
||||||
|
static bool ggml_backend_vk_register_host_buffer(void * buffer, size_t size) {
|
||||||
|
if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr && getenv("GGML_VK_REGISTER_HOST") == nullptr) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
if (size == 0) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
bool success = false;
|
||||||
|
for (size_t i = 0; i < GGML_VK_MAX_DEVICES; i++) {
|
||||||
|
vk_device& device = vk_instance.devices[i];
|
||||||
|
if (!device || !device->external_memory_host) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
vk_buffer buf = ggml_vk_buffer_from_host_ptr(device, buffer, size);
|
||||||
|
if (!buf || !buf->buffer) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
{
|
||||||
|
std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
|
||||||
|
device->pinned_memory.emplace_back(buffer, size, buf);
|
||||||
|
}
|
||||||
|
success = true;
|
||||||
|
}
|
||||||
|
return success;
|
||||||
|
}
|
||||||
|
|
||||||
|
static void ggml_backend_vk_unregister_host_buffer(void * buffer) {
|
||||||
|
for (size_t i = 0; i < GGML_VK_MAX_DEVICES; i++) {
|
||||||
|
vk_device& device = vk_instance.devices[i];
|
||||||
|
if (!device) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
|
||||||
|
for (auto it = device->pinned_memory.begin(); it != device->pinned_memory.end(); ++it) {
|
||||||
|
if (std::get<0>(*it) == buffer) {
|
||||||
|
vk_buffer buf = std::get<2>(*it);
|
||||||
|
device->pinned_memory.erase(it);
|
||||||
|
ggml_vk_destroy_buffer(buf);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
static void * ggml_backend_vk_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) {
|
||||||
|
GGML_UNUSED(reg);
|
||||||
|
if (strcmp(name, "ggml_backend_register_host_buffer") == 0) {
|
||||||
|
return (void *) ggml_backend_vk_register_host_buffer;
|
||||||
|
}
|
||||||
|
if (strcmp(name, "ggml_backend_unregister_host_buffer") == 0) {
|
||||||
|
return (void *) ggml_backend_vk_unregister_host_buffer;
|
||||||
|
}
|
||||||
|
return nullptr;
|
||||||
|
}
|
||||||
|
|
||||||
static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = {
|
static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = {
|
||||||
/* .get_name = */ ggml_backend_vk_reg_get_name,
|
/* .get_name = */ ggml_backend_vk_reg_get_name,
|
||||||
/* .get_device_count = */ ggml_backend_vk_reg_get_device_count,
|
/* .get_device_count = */ ggml_backend_vk_reg_get_device_count,
|
||||||
/* .get_device = */ ggml_backend_vk_reg_get_device,
|
/* .get_device = */ ggml_backend_vk_reg_get_device,
|
||||||
/* .get_proc_address = */ NULL,
|
/* .get_proc_address = */ ggml_backend_vk_reg_get_proc_address,
|
||||||
};
|
};
|
||||||
|
|
||||||
ggml_backend_reg_t ggml_backend_vk_reg() {
|
ggml_backend_reg_t ggml_backend_vk_reg() {
|
||||||
|
|||||||
Reference in New Issue
Block a user