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 on Vulkan. Imports host pages through VK_EXT_external_memory_host into device->pinned_memory, letting 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. A single Vulkan buffer cannot cover a multi-GB mmap (capped at device->max_buffer_size), so the region is imported in page-aligned chunks; a bound check makes tensors straddling a chunk boundary fall back to staging instead of reading out of bounds. Opt-in via GGML_CUDA_REGISTER_HOST=1 or GGML_VK_REGISTER_HOST=1, no-op otherwise. Assisted-by: opencode
This commit is contained in:
@@ -3283,6 +3283,7 @@ static vk_buffer ggml_vk_create_buffer(vk_device& device, size_t size, const std
|
||||
import_info.setPNext(&mem_flags_info);
|
||||
buf->device_memory = device->device.allocateMemory({ size, memory_type_idx, &import_info });
|
||||
} catch (const vk::SystemError& e) {
|
||||
GGML_LOG_WARN("ggml_vulkan: host pointer memory import failed (%s)\n", e.what());
|
||||
}
|
||||
} else {
|
||||
for (auto it = req_flags_list.begin(); it != req_flags_list.end(); it++) {
|
||||
@@ -8128,25 +8129,32 @@ static bool ggml_vk_buffer_write_2d_async(vk_context subctx, vk_buffer& dst, siz
|
||||
ggml_vk_host_get(dst->device, src, buf, buf_offset);
|
||||
|
||||
if (buf != nullptr) {
|
||||
// Memory is pinned, use as staging buffer
|
||||
std::vector<vk::BufferCopy> slices(1);
|
||||
if (width == spitch && width == dpitch) {
|
||||
// Only do single write if stride is equal
|
||||
slices[0].srcOffset = buf_offset;
|
||||
slices[0].dstOffset = offset;
|
||||
slices[0].size = width * height;
|
||||
} else {
|
||||
slices.resize(height);
|
||||
for (size_t i = 0; i < height; i++) {
|
||||
slices[i].srcOffset = buf_offset + i * spitch;
|
||||
slices[i].dstOffset = offset + i * dpitch;
|
||||
slices[i].size = width;
|
||||
// extent of the read in pinned source memory; guard against tensors that
|
||||
// straddle a pinned-chunk boundary (they fall back to staging below)
|
||||
size_t src_extent = (width == spitch) ? (size_t) width * height
|
||||
: (height > 0 ? (height - 1) * spitch + width : 0);
|
||||
if (buf_offset + src_extent <= buf->size) {
|
||||
// Memory is pinned, use as staging buffer
|
||||
std::vector<vk::BufferCopy> slices(1);
|
||||
if (width == spitch && width == dpitch) {
|
||||
// Only do single write if stride is equal
|
||||
slices[0].srcOffset = buf_offset;
|
||||
slices[0].dstOffset = offset;
|
||||
slices[0].size = width * height;
|
||||
} else {
|
||||
slices.resize(height);
|
||||
for (size_t i = 0; i < height; i++) {
|
||||
slices[i].srcOffset = buf_offset + i * spitch;
|
||||
slices[i].dstOffset = offset + i * dpitch;
|
||||
slices[i].size = width;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
ggml_vk_sync_buffers(nullptr, subctx);
|
||||
subctx->s->buffer->buf.copyBuffer(buf->buffer, dst->buffer, slices);
|
||||
return true;
|
||||
ggml_vk_sync_buffers(nullptr, subctx);
|
||||
subctx->s->buffer->buf.copyBuffer(buf->buffer, dst->buffer, slices);
|
||||
return true;
|
||||
}
|
||||
// straddles a chunk boundary: fall through to staging
|
||||
}
|
||||
VK_LOG_DEBUG("STAGING");
|
||||
|
||||
@@ -18243,11 +18251,96 @@ static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg,
|
||||
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.
|
||||
//
|
||||
// A single Vulkan buffer cannot cover a whole multi-GB mmap (it is capped at
|
||||
// device->max_buffer_size), so the region is imported in page-aligned chunks.
|
||||
// ggml_vk_host_get resolves a tensor pointer to the chunk that contains it,
|
||||
// and ggml_vk_buffer_write_2d_async falls back to staging for any tensor that
|
||||
// straddles a chunk boundary, so correctness is preserved.
|
||||
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 || device->max_buffer_size == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const size_t align = device->min_imported_host_pointer_alignment;
|
||||
size_t chunk = device->max_buffer_size & ~(align - 1);
|
||||
if (chunk == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
uint8_t * p = static_cast<uint8_t *>(buffer);
|
||||
size_t remaining = size;
|
||||
bool dev_success = false;
|
||||
while (remaining > 0) {
|
||||
size_t cur = std::min(remaining, chunk);
|
||||
vk_buffer buf = ggml_vk_buffer_from_host_ptr(device, p, cur);
|
||||
if (!buf || !buf->buffer) {
|
||||
break;
|
||||
}
|
||||
{
|
||||
std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
|
||||
device->pinned_memory.emplace_back(p, cur, buf);
|
||||
}
|
||||
dev_success = true;
|
||||
p += cur;
|
||||
remaining -= cur;
|
||||
}
|
||||
if (dev_success) {
|
||||
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 = {
|
||||
/* .get_name = */ ggml_backend_vk_reg_get_name,
|
||||
/* .get_device_count = */ ggml_backend_vk_reg_get_device_count,
|
||||
/* .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() {
|
||||
|
||||
Reference in New Issue
Block a user