ggml-vulkan: async unpinned H2D uploads via pooled staging

ggml_backend_vk_set_tensor_2d_async uploaded unpinned host memory through a
single shared sync_staging buffer followed by a blocking ggml_vk_synchronize,
which forces a full device sync per tensor. For MoE expert weights streamed
to the GPU (-ncmoe) this serializes every upload and starves the GPU between
MUL_MAT_ID ops (35B pp64: 10.8ms GPU / ~504ms wall).

Use a per-upload staging buffer borrowed from a context-local pool: the host
memcpy runs immediately into a host-visible buffer the GPU only reads after
the command buffer is submitted later, so no device sync is needed. Buffers
are recycled once the GPU is known idle (after the next sync); a drain cap
(before acquiring the recording context) bounds host staging memory. The
drain is placed before cpy_ctx acquisition so the sync never invalidates the
context we record into.

Assisted-by: opencode
This commit is contained in:
2026-07-16 23:14:15 +02:00
parent 71b61d18bf
commit bf55d009f4
+62 -6
View File
@@ -2160,6 +2160,13 @@ struct ggml_backend_vk_context {
std::vector<int> query_node_idx;
int32_t num_queries {};
int32_t query_idx {};
// Per-upload staging buffers used to upload unpinned host memory without the
// blocking device sync that the shared sync_staging path requires. Buffers
// move to inflight at upload time and back to free once the GPU is known to
// be idle (after a sync), so MoE-expert uploads no longer serialize compute.
std::vector<vk_buffer> async_staging_free;
std::vector<vk_buffer> async_staging_inflight;
};
static void * const vk_ptr_base = (void *)(uintptr_t) 0x1000; // NOLINT
@@ -15197,6 +15204,13 @@ static void ggml_vk_graph_cleanup(ggml_backend_vk_context * ctx) {
ctx->gc.contexts.clear();
ctx->pipeline_descriptor_set_requirements = 0;
ctx->descriptor_set_idx = 0;
// The GPU is idle here (called after synchronize): all staging buffers used
// for async host->device uploads this eval are safe to reuse.
for (auto & b : ctx->async_staging_inflight) {
ctx->async_staging_free.push_back(b);
}
ctx->async_staging_inflight.clear();
}
// Clean up on backend free
@@ -15215,6 +15229,15 @@ static void ggml_vk_cleanup(ggml_backend_vk_context * ctx) {
ggml_vk_destroy_buffer(ctx->prealloc_add_rms_partials);
ggml_vk_destroy_buffer(ctx->sync_staging);
for (auto & b : ctx->async_staging_free) {
ggml_vk_destroy_buffer(b);
}
ctx->async_staging_free.clear();
for (auto & b : ctx->async_staging_inflight) {
ggml_vk_destroy_buffer(b);
}
ctx->async_staging_inflight.clear();
ctx->prealloc_y_last_pipeline_used = nullptr;
ctx->prealloc_y_last_tensor_used = nullptr;
ctx->prealloc_y_last_decode_vector_staging = false;
@@ -15558,6 +15581,18 @@ static void ggml_backend_vk_set_tensor_2d_async(ggml_backend_t backend, ggml_ten
return;
}
// Bound host staging memory: if too many async uploads are outstanding,
// drain now (before acquiring the recording context below, since the sync
// would otherwise submit and reset the context we record into).
constexpr int ASYNC_STAGING_DRAIN = 16;
if ((int) ctx->async_staging_inflight.size() >= ASYNC_STAGING_DRAIN) {
ggml_vk_synchronize(ctx);
for (auto & b : ctx->async_staging_inflight) {
ctx->async_staging_free.push_back(b);
}
ctx->async_staging_inflight.clear();
}
ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)tensor->buffer->context;
vk_context cpy_ctx;
@@ -15582,8 +15617,27 @@ static void ggml_backend_vk_set_tensor_2d_async(ggml_backend_t backend, ggml_ten
if (!ret) {
const size_t staging_size = size * n_copies;
ggml_vk_ensure_sync_staging_buffer(ctx, staging_size);
ggml_vk_sync_buffers(nullptr, cpy_ctx);
// Upload through a pooled, per-upload staging buffer instead of the shared
// sync_staging + blocking ggml_vk_synchronize. A full device sync per expert
// tensor (the old path) serializes every MoE weight upload and stops the
// GGML_SCHED_PREFETCH_EXPERTS overlap from helping. Each upload borrows a
// host-visible buffer from async_staging_free (allocated on first use, recycled
// at the next sync when the GPU is idle); the host memcpy runs now since the GPU
// only reads the staging buffer after this command buffer is submitted later.
vk_buffer staging = nullptr;
for (auto it = ctx->async_staging_free.begin(); it != ctx->async_staging_free.end(); ++it) {
if ((*it)->size >= staging_size) {
staging = *it;
ctx->async_staging_free.erase(it);
break;
}
}
if (!staging) {
staging = ggml_vk_create_buffer_check(ctx->device, staging_size,
vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached,
vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent);
}
std::vector<vk::BufferCopy> slices(1);
if (size == stride_tensor) {
@@ -15599,16 +15653,18 @@ static void ggml_backend_vk_set_tensor_2d_async(ggml_backend_t backend, ggml_ten
}
}
cpy_ctx->s->buffer->buf.copyBuffer(ctx->sync_staging->buffer, buf->buffer, slices);
ggml_vk_sync_buffers(nullptr, cpy_ctx);
cpy_ctx->s->buffer->buf.copyBuffer(staging->buffer, buf->buffer, slices);
if (size == stride_data) {
deferred_memcpy(ctx->sync_staging->ptr, data, staging_size, &cpy_ctx->in_memcpys);
memcpy(staging->ptr, data, staging_size);
} else {
for (size_t i = 0; i < n_copies; i++) {
deferred_memcpy((uint8_t *)ctx->sync_staging->ptr + i * size, (const uint8_t *)data + i * stride_data, size, &cpy_ctx->in_memcpys);
memcpy((uint8_t *)staging->ptr + i * size, (const uint8_t *)data + i * stride_data, size);
}
}
ggml_vk_synchronize(ctx);
ctx->async_staging_inflight.push_back(staging);
}
}