ggml-vulkan: implement host buffer pinning for faster H2D uploads

Assisted-by: opencode

README.md

@@ -12,6 +12,34 @@
 
 LLM inference in C/C++
 
+## ⚡ This fork — Fable's MoE-offload prefill optimizations
+
+Two **opt-in** optimizations for large MoE models whose experts are offloaded to system RAM
+(`--n-cpu-moe`), found and implemented by Fable. Both are **off by default**, toggled via
+environment variables, and produce **token-identical** output to mainline.
+
+| Env var | What it does |
+| --- | --- |
+| `GGML_CUDA_REGISTER_HOST=1` | Page-locks (pins) the mmap'd CPU expert weights so host→device copies go straight over DMA instead of through the driver's hidden bounce buffer (~6–7 → ~20 GB/s). |
+| `GGML_SCHED_PREFETCH_EXPERTS=1` | Prefetches each layer's experts on a second CUDA stream, so the weight uploads overlap compute instead of stalling the GPU. |
+
+### Benchmark
+
+Measured on an **RTX 3060 12GB** with **Qwen3.6-35B-A3B** (`--n-cpu-moe 26`), prompt-processing at 2048 (`MODEL` = path to your `.gguf`):
+
+```bash
+# baseline (patches off):
+./build/bin/llama-bench -m MODEL -ngl 99 -ncmoe 26 -p 2048 -n 0 -r 5 -b 2048 -ub 2048
+
+# patched (both optimizations on):
+GGML_CUDA_REGISTER_HOST=1 GGML_SCHED_PREFETCH_EXPERTS=1 \
+./build/bin/llama-bench -m MODEL -ngl 99 -ncmoe 26 -p 2048 -n 0 -r 5 -b 2048 -ub 2048
+```
+
+Result: **~1143 → ~1880 t/s** prefill (**+64%**) — same GPU, same settings, token-identical.
+
+Branches: [`fable5/host-register`](https://github.com/thecodacus/llama.cpp/tree/fable5/host-register) (pinning only) · [`fable5/prefetch-experts`](https://github.com/thecodacus/llama.cpp/tree/fable5/prefetch-experts) (both — this branch).
+
 ## Recent API changes
 
 - [Changelog for `libllama` API](https://github.com/ggml-org/llama.cpp/issues/9289)

ggml/src/ggml-backend.cpp

@@ -761,6 +761,10 @@ static bool ggml_is_view_op(enum ggml_op op) {
 #define GGML_SCHED_MAX_COPIES 4
 #endif
 
+#ifndef GGML_SCHED_MAX_PREFETCH_SLOTS
+#define GGML_SCHED_MAX_PREFETCH_SLOTS 8
+#endif
+
 struct ggml_backend_sched_split {
     int backend_id;
     int i_start;
@@ -818,6 +822,19 @@ struct ggml_backend_sched {
 
     bool op_offload;
 
+    // full-tensor prefetch of offloaded MUL_MAT_ID weights (GGML_SCHED_PREFETCH_EXPERTS)
+    // with a large batch virtually every expert is used, so the routing ids are not worth
+    // waiting for; uploads run through a second backend instance on the same device so
+    // they overlap compute, alternating between two staging slots
+    bool prefetch_experts;
+    ggml_backend_t prefetch_backend;
+    int prefetch_n_slots;
+    ggml_backend_buffer_t prefetch_slots[GGML_SCHED_MAX_PREFETCH_SLOTS];
+    ggml_backend_event_t prefetch_ready[GGML_SCHED_MAX_PREFETCH_SLOTS];
+    ggml_backend_event_t prefetch_free[GGML_SCHED_MAX_PREFETCH_SLOTS];
+    bool prefetch_used[GGML_SCHED_MAX_PREFETCH_SLOTS];
+    int prefetch_cur;
+
     int debug;
 
     // used for debugging graph reallocations [GGML_SCHED_DEBUG_REALLOC]
@@ -1538,6 +1555,94 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) {
     return true;
 }
 
+static void ggml_backend_sched_prefetch_disable(ggml_backend_sched_t sched, ggml_backend_t split_backend) {
+    sched->prefetch_experts = false;
+    if (sched->prefetch_backend) {
+        ggml_backend_synchronize(split_backend);
+        ggml_backend_synchronize(sched->prefetch_backend);
+    }
+    for (int i = 0; i < sched->prefetch_n_slots; i++) {
+        ggml_backend_buffer_free(sched->prefetch_slots[i]);
+        sched->prefetch_slots[i] = NULL;
+        sched->prefetch_used[i] = false;
+    }
+}
+
+// slots are sized once for the largest offloaded expert tensor in the current graph so
+// that they never need to grow mid-eval
+static size_t ggml_backend_sched_prefetch_max_size(ggml_backend_sched_t sched) {
+    size_t max_size = 0;
+    for (int split_id = 0; split_id < sched->n_splits; split_id++) {
+        struct ggml_backend_sched_split * split = &sched->splits[split_id];
+        if (split->graph.n_nodes == 0 || split->graph.nodes[0]->op != GGML_OP_MUL_MAT_ID) {
+            continue;
+        }
+        for (int input_id = 0; input_id < split->n_inputs; input_id++) {
+            const ggml_tensor * input = split->inputs[input_id];
+            if (input->buffer &&
+                ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
+                ggml_backend_buffer_is_host(input->buffer)) {
+                max_size = std::max(max_size, ggml_nbytes(input));
+            }
+        }
+    }
+    return max_size;
+}
+
+static bool ggml_backend_sched_prefetch_init(ggml_backend_sched_t sched, ggml_backend_t split_backend, size_t size) {
+    if (sched->prefetch_backend == NULL) {
+        ggml_backend_dev_t dev = split_backend->device;
+        ggml_backend_dev_props props;
+        ggml_backend_dev_get_props(dev, &props);
+        if (!props.caps.async || !props.caps.events) {
+            sched->prefetch_experts = false;
+            return false;
+        }
+        sched->prefetch_backend = ggml_backend_dev_init(dev, NULL);
+        if (sched->prefetch_backend == NULL) {
+            sched->prefetch_experts = false;
+            return false;
+        }
+        for (int i = 0; i < sched->prefetch_n_slots; i++) {
+            sched->prefetch_ready[i] = ggml_backend_event_new(dev);
+            sched->prefetch_free[i]  = ggml_backend_event_new(dev);
+            if (sched->prefetch_ready[i] == NULL || sched->prefetch_free[i] == NULL) {
+                sched->prefetch_experts = false;
+                return false;
+            }
+        }
+    }
+
+    size = std::max(size, ggml_backend_sched_prefetch_max_size(sched));
+
+    ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(split_backend);
+    for (int i = 0; i < sched->prefetch_n_slots; i++) {
+        if (sched->prefetch_slots[i] == NULL || ggml_backend_buffer_get_size(sched->prefetch_slots[i]) < size) {
+            // allocate before freeing so a failure leaves the old slot intact
+            ggml_backend_buffer_t new_buf = ggml_backend_buft_alloc_buffer(buft, size);
+            if (new_buf == NULL) {
+                // overlap needs at least 2 slots, otherwise run with what fits
+                if (i >= 2 && sched->prefetch_slots[0] != NULL &&
+                    ggml_backend_buffer_get_size(sched->prefetch_slots[0]) >= size) {
+                    sched->prefetch_n_slots = i;
+                    sched->prefetch_cur = 0;
+                    return true;
+                }
+                ggml_backend_sched_prefetch_disable(sched, split_backend);
+                return false;
+            }
+            if (sched->prefetch_slots[i] != NULL) {
+                ggml_backend_synchronize(split_backend);
+                ggml_backend_synchronize(sched->prefetch_backend);
+                ggml_backend_buffer_free(sched->prefetch_slots[i]);
+            }
+            sched->prefetch_slots[i] = new_buf;
+            sched->prefetch_used[i] = false;
+        }
+    }
+    return true;
+}
+
 static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) {
     GGML_ASSERT(sched);
     struct ggml_backend_sched_split * splits = sched->splits;
@@ -1550,6 +1655,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
         struct ggml_backend_sched_split * split = &splits[split_id];
         int split_backend_id = split->backend_id;
         ggml_backend_t split_backend = sched->backends[split_backend_id];
+        int split_prefetch_slot = -1;
+        ggml_tensor * prefetch_input_cpy = NULL;
+        ggml_backend_buffer_t prefetch_saved_buffer = NULL;
+        void * prefetch_saved_data = NULL;
 
         // copy the input tensors to the split backend
         for (int input_id = 0; input_id < split->n_inputs; input_id++) {
@@ -1566,6 +1675,41 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                 }
                 ggml_backend_tensor_copy(input, input_cpy);
             } else {
+                // with a large batch virtually every expert is used, so instead of waiting
+                // for the routing ids, upload the full tensor through the prefetch backend
+                // and let the copy overlap compute of the previous split
+                if (sched->prefetch_experts && !sched->callback_eval && split_prefetch_slot == -1 && split->graph.n_nodes > 0) {
+                    ggml_tensor * node = split->graph.nodes[0];
+                    if (ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
+                        ggml_backend_buffer_is_host(input->buffer) &&
+                        node->op == GGML_OP_MUL_MAT_ID && node->src[0] == input_cpy) {
+                        const ggml_tensor * ids = node->src[2];
+                        const int64_t n_expert = input->ne[2];
+                        if (ids->ne[0]*ids->ne[1] >= 2*n_expert &&
+                            ggml_backend_sched_prefetch_init(sched, split_backend, ggml_nbytes(input))) {
+                            const int slot = sched->prefetch_cur;
+                            sched->prefetch_cur = (sched->prefetch_cur + 1) % sched->prefetch_n_slots;
+                            // wait for the previous user of this slot to finish computing
+                            if (sched->prefetch_used[slot]) {
+                                ggml_backend_event_wait(sched->prefetch_backend, sched->prefetch_free[slot]);
+                            }
+                            // point the staging copy at the slot only for the duration of
+                            // this split, so a fallback to the regular path on a later
+                            // eval can never see a dangling slot pointer
+                            prefetch_input_cpy    = input_cpy;
+                            prefetch_saved_buffer = input_cpy->buffer;
+                            prefetch_saved_data   = input_cpy->data;
+                            input_cpy->buffer = sched->prefetch_slots[slot];
+                            input_cpy->data   = ggml_backend_buffer_get_base(sched->prefetch_slots[slot]);
+                            ggml_backend_tensor_set_async(sched->prefetch_backend, input_cpy, input->data, 0, ggml_nbytes(input));
+                            ggml_backend_event_record(sched->prefetch_ready[slot], sched->prefetch_backend);
+                            ggml_backend_event_wait(split_backend, sched->prefetch_ready[slot]);
+                            split_prefetch_slot = slot;
+                            continue;
+                        }
+                    }
+                }
+
                 // wait for the split backend to finish using the input before overwriting it
                 if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
                     ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]);
@@ -1676,6 +1820,13 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
 
         if (!sched->callback_eval) {
             enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph);
+            if (split_prefetch_slot != -1) {
+                // the kernels have captured the slot address at launch, safe to restore
+                ggml_backend_event_record(sched->prefetch_free[split_prefetch_slot], split_backend);
+                sched->prefetch_used[split_prefetch_slot] = true;
+                prefetch_input_cpy->buffer = prefetch_saved_buffer;
+                prefetch_input_cpy->data   = prefetch_saved_data;
+            }
             if (ec != GGML_STATUS_SUCCESS) {
                 return ec;
             }
@@ -1788,6 +1939,15 @@ ggml_backend_sched_t ggml_backend_sched_new(
     sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends);
     sched->op_offload = op_offload;
 
+    // GGML_SCHED_PREFETCH_EXPERTS=1 enables the default slot count, higher values set it
+    // directly; more slots let uploads run further ahead of compute at the cost of one
+    // max-sized expert tensor of device memory per slot
+    const char * GGML_SCHED_PREFETCH_EXPERTS = getenv("GGML_SCHED_PREFETCH_EXPERTS");
+    const int prefetch_n_slots = GGML_SCHED_PREFETCH_EXPERTS ? atoi(GGML_SCHED_PREFETCH_EXPERTS) : 0;
+    sched->prefetch_experts = op_offload && prefetch_n_slots > 0;
+    // default of 3 covers the gate/up/down expert tensors of one MoE layer
+    sched->prefetch_n_slots = prefetch_n_slots <= 1 ? 3 : std::min(prefetch_n_slots, GGML_SCHED_MAX_PREFETCH_SLOTS);
+
     ggml_backend_sched_reset(sched);
 
     return sched;
@@ -1802,6 +1962,16 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
             ggml_backend_event_free(sched->events[b][c]);
         }
     }
+    if (sched->prefetch_backend) {
+        ggml_backend_synchronize(sched->prefetch_backend);
+        // the slot count may have been reduced after a failed allocation, free everything
+        for (int i = 0; i < GGML_SCHED_MAX_PREFETCH_SLOTS; i++) {
+            ggml_backend_event_free(sched->prefetch_ready[i]);
+            ggml_backend_event_free(sched->prefetch_free[i]);
+            ggml_backend_buffer_free(sched->prefetch_slots[i]);
+        }
+        ggml_backend_free(sched->prefetch_backend);
+    }
     ggml_gallocr_free(sched->galloc);
     ggml_free(sched->ctx);
     ggml_hash_set_free(&sched->hash_set);
@@ -1906,6 +2076,9 @@ void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) {
     for (int i = 0; i < sched->n_backends; i++) {
         ggml_backend_synchronize(sched->backends[i]);
     }
+    if (sched->prefetch_backend) {
+        ggml_backend_synchronize(sched->prefetch_backend);
+    }
     if (!sched->is_alloc) {
         // if the graph is not already allocated, always use copy 0 after a synchronization
         // this ensures that during generation the same copy is used every time,

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -17917,11 +17917,63 @@ static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg,
     return devices[device];
 }
 
+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;
+    }
+
+    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) {
+            continue;
+        }
+
+        std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
+        device->pinned_memory.push_back(std::make_tuple(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(); ) {
+            if (std::get<0>(*it) == buffer) {
+                vk_buffer buf = std::get<2>(*it);
+                ggml_vk_destroy_buffer(buf);
+                it = device->pinned_memory.erase(it);
+                break; // A buffer is registered once per device
+            } else {
+                ++it;
+            }
+        }
+    }
+}
+
+static void * ggml_backend_vk_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) {
+    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() {

src/llama-mmap.cpp

@@ -618,13 +618,47 @@ struct llama_mmap::impl {
 };
 
 llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa) : pimpl(std::make_unique<impl>(file, prefetch, numa)) {}
-llama_mmap::~llama_mmap() = default;
+
+llama_mmap::~llama_mmap() {
+    // unpin before the pages are unmapped by the impl destructor
+    if (host_reg_addr && host_unreg_fn) {
+        host_unreg_fn(host_reg_addr);
+    }
+}
 
 size_t llama_mmap::size() const { return pimpl->size; }
 void * llama_mmap::addr() const { return pimpl->addr; }
 
 void llama_mmap::unmap_fragment(size_t first, size_t last) { pimpl->unmap_fragment(first, last); }
 
+size_t llama_mmap::register_host(size_t first, size_t last, bool (*reg_fn)(void *, size_t), void (*unreg_fn)(void *)) {
+#ifdef _POSIX_MAPPED_FILES
+    if (host_reg_addr || !reg_fn || !unreg_fn || last <= first) {
+        return 0;
+    }
+
+    // expand outward to the page boundaries retained by unmap_fragment
+    const size_t page_size = sysconf(_SC_PAGESIZE);
+    first = first & ~(page_size - 1);
+    last  = (last + page_size - 1) & ~(page_size - 1);
+
+    void * reg_addr = (uint8_t *) pimpl->addr + first;
+    if (!reg_fn(reg_addr, last - first)) {
+        return 0;
+    }
+
+    host_reg_addr  = reg_addr;
+    host_unreg_fn  = unreg_fn;
+    return last - first;
+#else
+    GGML_UNUSED(first);
+    GGML_UNUSED(last);
+    GGML_UNUSED(reg_fn);
+    GGML_UNUSED(unreg_fn);
+    return 0;
+#endif
+}
+
 #if defined(_POSIX_MEMLOCK_RANGE) || defined(_WIN32)
 const bool llama_mmap::SUPPORTED  = true;
 #else

src/llama-mmap.h

@@ -50,11 +50,19 @@ struct llama_mmap {
 
     void unmap_fragment(size_t first, size_t last);
 
+    // pin the pages backing [first, last) with a backend allocator for faster H2D copies,
+    // unpinned in the destructor before the pages are unmapped
+    // returns the number of bytes registered, 0 on failure
+    size_t register_host(size_t first, size_t last, bool (*reg_fn)(void *, size_t), void (*unreg_fn)(void *));
+
     static const bool SUPPORTED;
 
 private:
     struct impl;
     std::unique_ptr<impl> pimpl;
+
+    void * host_reg_addr = nullptr;
+    void (*host_unreg_fn)(void *) = nullptr;
 };
 
 struct llama_mlock {

src/llama-model-loader.cpp

@@ -1673,6 +1673,15 @@ bool llama_model_loader::load_all_data(
     if (size_done >= size_data) {
         // unmap offloaded tensors and metadata
         if (use_mmap) {
+            // pin the pages backing the weights kept in system memory for faster H2D copies
+            bool (*reg_fn)(void *, size_t) = nullptr;
+            void (*unreg_fn)(void *) = nullptr;
+            for (size_t i = 0; i < ggml_backend_dev_count() && !reg_fn; i++) {
+                ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_dev_get(i));
+                reg_fn   = (bool (*)(void *, size_t)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_register_host_buffer");
+                unreg_fn = (void (*)(void *))         ggml_backend_reg_get_proc_address(reg, "ggml_backend_unregister_host_buffer");
+            }
+
             for (uint32_t idx = 0; idx < mappings.size(); idx++) {
                 const auto & mmap_used = mmaps_used.at(idx);
                 auto & mapping = mappings.at(idx);
@@ -1680,6 +1689,13 @@ bool llama_model_loader::load_all_data(
                 if (mmap_used.second != 0) {
                     mapping->unmap_fragment(mmap_used.second, mapping->size());
                 }
+                if (mmap_used.second > mmap_used.first) {
+                    size_t n_registered = mapping->register_host(mmap_used.first, mmap_used.second, reg_fn, unreg_fn);
+                    if (n_registered > 0) {
+                        LLAMA_LOG_INFO("%s: pinned %.2f MiB of mapped model memory for faster H2D transfers\n",
+                                __func__, n_registered / 1024.0 / 1024.0);
+                    }
+                }
             }
         }
         if (progress_callback) {