feat: overlap CPU cold-expert and GPU hot-expert computation in the MoE cache

common/arg.cpp

@@ -2805,6 +2805,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.no_op_offload = !value;
         }
     ));
+    add_opt(common_arg(
+        {"--sched-async-cpu"},
+        {"--no-sched-async-cpu"},
+        string_format("whether to run CPU graph splits on a worker thread so independent GPU splits overlap them (default: %s)", params.sched_async_cpu ? "true" : "false"),
+        [](common_params & params, bool value) {
+            params.sched_async_cpu = value;
+        }
+    ));
     add_opt(common_arg(
         {"--lora"}, "FNAME",
         "path to LoRA adapter (use comma-separated values to load multiple adapters)",

common/common.cpp

@@ -1619,6 +1619,7 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.offload_kqv       = !params.no_kv_offload;
     cparams.no_perf           = params.no_perf;
     cparams.op_offload        = !params.no_op_offload;
+    cparams.sched_async_cpu   = params.sched_async_cpu;
     cparams.swa_full          = params.swa_full;
     cparams.kv_unified        = params.kv_unified;
 

common/common.h

@@ -580,6 +580,7 @@ struct common_params {
     bool warmup            = true;  // warmup run
     bool check_tensors     = false; // validate tensor data
     bool no_op_offload     = false; // globally disable offload host tensor operations to device
+    bool sched_async_cpu   = true;  // run CPU graph splits on a worker thread (overlaps independent GPU splits)
     bool no_extra_bufts    = false; // disable extra buffer types (used for weight repacking)
     bool no_host           = false; // bypass host buffer allowing extra buffers to be used
 

ggml/include/ggml-backend.h

@@ -332,6 +332,10 @@ extern "C" {
     GGML_API size_t                     ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend);
 
     GGML_API void                 ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend);
+
+    // when enabled, CPU splits run on a worker thread so that independent splits
+    // on other backends execute concurrently with them
+    GGML_API void                 ggml_backend_sched_set_async_cpu(ggml_backend_sched_t sched, bool enable);
     GGML_API ggml_backend_t       ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node);
 
     // Split graph without allocating it

ggml/src/ggml-backend.cpp

@@ -20,6 +20,9 @@
 #include <stdlib.h>
 #include <string.h>
 #include <algorithm>
+#include <condition_variable>
+#include <mutex>
+#include <thread>
 #include <vector>
 
 #ifdef __APPLE__
@@ -775,6 +778,78 @@ struct ggml_backend_sched_split {
     struct ggml_cgraph graph;
 };
 
+// async execution of CPU splits (GGML_SCHED_ASYNC_CPU): a persistent worker
+// computes a CPU split while the main thread keeps launching later splits that
+// do not depend on it, so an independent GPU split overlaps the CPU compute
+struct ggml_sched_cpu_async {
+    std::thread             worker;
+    std::mutex              mtx;
+    std::condition_variable cv;
+    ggml_backend_t          job_backend = nullptr;
+    struct ggml_cgraph *    job_graph   = nullptr;
+    enum ggml_status        job_status  = GGML_STATUS_SUCCESS;
+    bool                    job_ready   = false;
+    bool                    job_done    = false;
+    bool                    stop        = false;
+    bool                    pending     = false; // main-thread view: a job is queued or running
+
+    ggml_sched_cpu_async() {
+        worker = std::thread([this]() {
+            for (;;) {
+                std::unique_lock<std::mutex> lock(mtx);
+                cv.wait(lock, [this]() { return job_ready || stop; });
+                if (stop) {
+                    return;
+                }
+                job_ready = false;
+                ggml_backend_t   backend = job_backend;
+                ggml_cgraph *    graph   = job_graph;
+                lock.unlock();
+
+                enum ggml_status status = ggml_backend_graph_compute_async(backend, graph);
+
+                lock.lock();
+                job_status = status;
+                job_done   = true;
+                cv.notify_all();
+            }
+        });
+    }
+
+    ~ggml_sched_cpu_async() {
+        {
+            std::lock_guard<std::mutex> lock(mtx);
+            stop = true;
+        }
+        cv.notify_all();
+        worker.join();
+    }
+
+    void launch(ggml_backend_t backend, struct ggml_cgraph * graph) {
+        {
+            std::lock_guard<std::mutex> lock(mtx);
+            job_backend = backend;
+            job_graph   = graph;
+            job_status  = GGML_STATUS_SUCCESS;
+            job_done    = false;
+            job_ready   = true;
+        }
+        cv.notify_all();
+        pending = true;
+    }
+
+    // wait for the in-flight job (if any); returns its status
+    enum ggml_status join() {
+        if (!pending) {
+            return GGML_STATUS_SUCCESS;
+        }
+        std::unique_lock<std::mutex> lock(mtx);
+        cv.wait(lock, [this]() { return job_done; });
+        pending = false;
+        return job_status;
+    }
+};
+
 struct ggml_backend_sched {
     bool is_reset; // true if the scheduler has been reset since the last graph split
     bool is_alloc;
@@ -835,6 +910,9 @@ struct ggml_backend_sched {
     bool prefetch_used[GGML_SCHED_MAX_PREFETCH_SLOTS];
     int prefetch_cur;
 
+    // async CPU split execution (GGML_SCHED_ASYNC_CPU); NULL when disabled
+    struct ggml_sched_cpu_async * cpu_async;
+
     int debug;
 
     // used for debugging graph reallocations [GGML_SCHED_DEBUG_REALLOC]
@@ -1647,6 +1725,11 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
     GGML_ASSERT(sched);
     struct ggml_backend_sched_split * splits = sched->splits;
 
+    if (sched->cpu_async) {
+        // a job left over from an aborted eval references stale split memory - drain it
+        sched->cpu_async->join();
+    }
+
     ggml_tensor * prev_ids_tensor = nullptr;
     std::vector<int32_t> ids;
     std::vector<ggml_bitset_t> used_ids;
@@ -1660,6 +1743,22 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
         ggml_backend_buffer_t prefetch_saved_buffer = NULL;
         void * prefetch_saved_data = NULL;
 
+        // an async CPU split may still be computing; join before anything that
+        // depends on it: another CPU split, or a split reading a CPU tensor
+        if (sched->cpu_async && sched->cpu_async->pending) {
+            bool must_join = split_backend_id == sched->n_backends - 1;
+            for (int input_id = 0; !must_join && input_id < split->n_inputs; input_id++) {
+                ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]);
+                must_join = input_backend == sched->backends[sched->n_backends - 1];
+            }
+            if (must_join) {
+                enum ggml_status ec = sched->cpu_async->join();
+                if (ec != GGML_STATUS_SUCCESS) {
+                    return ec;
+                }
+            }
+        }
+
         // copy the input tensors to the split backend
         for (int input_id = 0; input_id < split->n_inputs; input_id++) {
             ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]);
@@ -1822,6 +1921,12 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
         }
 
         if (!sched->callback_eval) {
+            if (sched->cpu_async && split_backend_id == sched->n_backends - 1 && split_prefetch_slot == -1) {
+                // run the CPU split on the worker; the loop continues launching
+                // later splits until one depends on this split's outputs
+                sched->cpu_async->launch(split_backend, &split->graph);
+                continue;
+            }
             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
@@ -1875,6 +1980,13 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
         }
     }
 
+    if (sched->cpu_async) {
+        enum ggml_status ec = sched->cpu_async->join();
+        if (ec != GGML_STATUS_SUCCESS) {
+            return ec;
+        }
+    }
+
     return GGML_STATUS_SUCCESS;
 }
 
@@ -1951,6 +2063,7 @@ ggml_backend_sched_t ggml_backend_sched_new(
     // 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;
@@ -1975,6 +2088,7 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
         }
         ggml_backend_free(sched->prefetch_backend);
     }
+    delete sched->cpu_async;
     ggml_gallocr_free(sched->galloc);
     ggml_free(sched->ctx);
     ggml_hash_set_free(&sched->hash_set);
@@ -2074,8 +2188,22 @@ enum ggml_status ggml_backend_sched_graph_compute_async(ggml_backend_sched_t sch
     return ggml_backend_sched_compute_splits(sched);
 }
 
+void ggml_backend_sched_set_async_cpu(ggml_backend_sched_t sched, bool enable) {
+    GGML_ASSERT(sched);
+    if (enable && sched->cpu_async == NULL) {
+        sched->cpu_async = new ggml_sched_cpu_async();
+    } else if (!enable && sched->cpu_async != NULL) {
+        sched->cpu_async->join();
+        delete sched->cpu_async;
+        sched->cpu_async = NULL;
+    }
+}
+
 void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) {
     GGML_ASSERT(sched);
+    if (sched->cpu_async) {
+        sched->cpu_async->join();
+    }
     for (int i = 0; i < sched->n_backends; i++) {
         ggml_backend_synchronize(sched->backends[i]);
     }

include/llama.h

@@ -392,6 +392,7 @@ extern "C" {
         bool offload_kqv; // offload the KQV ops (including the KV cache) to GPU
         bool no_perf;     // measure performance timings
         bool op_offload;  // offload host tensor operations to device
+        bool sched_async_cpu; // run CPU graph splits on a worker thread so independent GPU splits overlap them
         bool swa_full;    // use full-size SWA cache (https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)
                           // NOTE: setting to false when n_seq_max > 1 can cause bad performance in some cases
                           //       ref: https://github.com/ggml-org/llama.cpp/pull/13845#issuecomment-2924800573

src/llama-context.cpp

@@ -117,6 +117,7 @@ llama_context::llama_context(
     cparams.embeddings_nextn        = false;
     cparams.embeddings_nextn_masked = false;
     cparams.offload_kqv             = params.offload_kqv;
+    cparams.sched_async_cpu         = params.sched_async_cpu;
     cparams.no_perf                 = params.no_perf;
     cparams.warmup                  = false;
 
@@ -594,6 +595,7 @@ void llama_context::sched_reserve() {
     gf_res_reserve.reset(new llm_graph_result(max_nodes));
 
     sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload));
+    ggml_backend_sched_set_async_cpu(sched.get(), cparams.sched_async_cpu);
 
     llama_memory_context_ptr mctx;
     if (memory) {
@@ -629,6 +631,7 @@ void llama_context::sched_reserve() {
                 LLAMA_LOG_WARN("%s: compute buffer allocation failed, retrying without pipeline parallelism\n", __func__);
                 cparams.pipeline_parallel = false;
                 sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, cparams.op_offload));
+                ggml_backend_sched_set_async_cpu(sched.get(), cparams.sched_async_cpu);
                 gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get());
             }
             if (!gf) {
@@ -3494,6 +3497,7 @@ llama_context_params llama_context_default_params() {
         /*.offload_kqv                 =*/ true,
         /*.no_perf                     =*/ true,
         /*.op_offload                  =*/ true,
+        /*.sched_async_cpu             =*/ true,
         /*.swa_full                    =*/ true,
         /*.kv_unified                  =*/ false,
         /*.sampler                     =*/ nullptr,

src/llama-cparams.h

@@ -36,6 +36,7 @@ struct llama_cparams {
     bool embeddings_nextn_masked; // extract for only rows where batch.logits != 0
     bool causal_attn;
     bool offload_kqv;
+    bool sched_async_cpu;
     bool flash_attn;
     bool auto_fa;
     bool fused_gdn_ar;       // use fused gated delta net (autoregressive)

src/llama-graph.cpp

@@ -2009,14 +2009,25 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
             return down;
         };
 
-        ggml_tensor * hot = build_pack_chain(moe_cache->ffn_gate_exps_hot, moe_cache->ffn_up_exps_hot, moe_cache->ffn_down_exps_hot, ids_hot);
-        cb(hot, "ffn_moe_down_hot", il);
-
+        // cold chain is built first so its nodes precede the hot chain in the
+        // graph: the scheduler then emits [cold split, hot split] and, with
+        // async CPU splits, computes the cold chain on a worker while the
+        // hot chain (which has no CPU inputs) runs concurrently on the GPU.
+        // pinning the merge to CPU keeps it out of the hot split so the hot
+        // split stays free of cross-backend inputs.
         ggml_tensor * cold = build_pack_chain(gate_exps, up_exps, down_exps, ids_cold);
         cb(cold, "ffn_moe_down_cold", il);
 
-        experts = ggml_add(ctx0, hot, cold);
+        ggml_tensor * hot = build_pack_chain(moe_cache->ffn_gate_exps_hot, moe_cache->ffn_up_exps_hot, moe_cache->ffn_down_exps_hot, ids_hot);
+        cb(hot, "ffn_moe_down_hot", il);
+
+        experts = ggml_add(ctx0, cold, hot);
         cb(experts, "ffn_moe_down", il);
+        // decode-size batches only: for large (prefill) batches the CPU merge
+        // and its per-layer activation copies cost more than the overlap hides
+        if (cparams.sched_async_cpu && n_tokens <= 8) {
+            ggml_backend_sched_set_tensor_backend(sched, experts, backend_cpu);
+        }
     } else if (gate_up_exps) {
         // merged gate_up path: one mul_mat_id, then split into gate and up views
         ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s); // [n_ff*2, n_expert_used, n_tokens]

src/llama-model.cpp

@@ -1764,6 +1764,10 @@ void llama_model_base::init_moe_expert_cache() {
         }
         return;
     }
+    // weights usage pins the pack tensors to their backend during graph
+    // assignment - without it a CPU-assigned consumer can drag the hot
+    // matmuls (and a per-layer weight copy) onto the CPU
+    ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
 
     // fill packs (expert dim is outermost: one contiguous slab per expert)
     std::vector<uint8_t> slab;