ggml : persistent expert cache for --n-cpu-moe (RFC #20757)

common/arg.cpp

@@ -2328,6 +2328,18 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             }
         }
     ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_N_CPU_MOE_DRAFT"));
+    add_opt(common_arg(
+        {"--expert-cache-slots"}, "N",
+        "N-slot LRU expert cache: N GPU slots per MoE layer (saves VRAM vs full GPU copy); use with --cpu-moe. "
+        "N=0 disables (default), N=-1 dedup-only, N=-2 or 'auto' fills available VRAM",
+        [](common_params & params, const std::string & value) {
+            if (value == "auto") {
+                params.expert_cache_n_slots = -2;
+            } else {
+                params.expert_cache_n_slots = std::stoi(value);
+            }
+        }
+    ).set_env("LLAMA_ARG_EXPERT_CACHE_SLOTS"));
     GGML_ASSERT(params.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0
     add_opt(common_arg(
         {"-ngl", "--gpu-layers", "--n-gpu-layers"}, "N",

common/common.cpp

@@ -1474,6 +1474,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.expert_cache_n_slots = params.expert_cache_n_slots;
     cparams.swa_full          = params.swa_full;
     cparams.kv_unified        = params.kv_unified;
 

common/common.h

@@ -543,6 +543,7 @@ struct common_params {
     bool no_op_offload     = false; // globally disable offload host tensor operations to device
     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
+    int32_t expert_cache_n_slots = 0; // N-slot LRU expert cache: 0=off, -1=dedup-only, N>0=N GPU slots per layer (--expert-cache-slots)
 
     bool single_turn       = false; // single turn chat conversation
 

ggml/include/ggml-backend.h

@@ -340,6 +340,20 @@ extern "C" {
     // Set a callback to be called for each resulting node during graph compute
     GGML_API void                 ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data);
 
+    // Expert cache for MoE weight offloading: enable/disable and query stats
+    GGML_API void ggml_backend_sched_set_expert_cache(
+        ggml_backend_sched_t sched, int32_t n_slots);
+
+    GGML_API void ggml_backend_sched_get_expert_cache_stats(
+        ggml_backend_sched_t sched,
+        int64_t * n_hits,
+        int64_t * n_misses,
+        int64_t * n_fate_hits,
+        int64_t * bytes_saved,
+        int64_t * bytes_copied);
+
+    GGML_API void ggml_backend_sched_reset_expert_cache_stats(ggml_backend_sched_t sched);
+
     //
     // Utils
     //

include/llama.h

@@ -367,6 +367,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
+        int32_t expert_cache_n_slots; // N-slot LRU expert cache: 0=off, -1=dedup-only, -2=auto (fill available VRAM), N>0=N GPU slots per layer
         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
@@ -991,6 +992,18 @@ extern "C" {
     // and is not necessary to call it explicitly in most cases
     LLAMA_API void llama_synchronize(struct llama_context * ctx);
 
+    // Expert cache statistics for MoE weight offloading (--expert-cache-slots / --n-cpu-moe).
+    // All output pointers are optional (may be NULL). Reset counters with llama_expert_cache_stats_reset().
+    LLAMA_API void llama_expert_cache_stats(
+            const struct llama_context * ctx,
+            int64_t * n_hits,
+            int64_t * n_misses,
+            int64_t * n_fate_hits,
+            int64_t * bytes_saved,
+            int64_t * bytes_copied);
+
+    LLAMA_API void llama_expert_cache_stats_reset(struct llama_context * ctx);
+
     // Token logits obtained from the last call to llama_decode()
     // The logits for which llama_batch.logits[i] != 0 are stored contiguously
     // in the order they have appeared in the batch.

src/llama-context.cpp

@@ -160,8 +160,9 @@ llama_context::llama_context(
 
     cparams.n_ubatch = std::min(cparams.n_batch, params.n_ubatch == 0 ? params.n_batch : params.n_ubatch);
 
-    cparams.op_offload = params.op_offload;
-    cparams.kv_unified = params.kv_unified;
+    cparams.op_offload          = params.op_offload;
+    cparams.expert_cache_n_slots = params.expert_cache_n_slots;
+    cparams.kv_unified          = params.kv_unified;
 
     // initialized later
     cparams.pipeline_parallel = false;
@@ -365,6 +366,15 @@ llama_context::llama_context(
 }
 
 llama_context::~llama_context() {
+    if (cparams.expert_cache_n_slots != 0 && sched) {
+        int64_t hits, misses, fate, saved, copied;
+        ggml_backend_sched_get_expert_cache_stats(sched.get(), &hits, &misses, &fate, &saved, &copied);
+        if (hits + misses > 0) {
+            LLAMA_LOG_INFO("%s: expert-cache: hits=%" PRId64 " (%.1f%%) misses=%" PRId64 " fate=%" PRId64 " saved=%.1fMB copied=%.1fMB\n",
+                __func__, hits, 100.0 * hits / (hits + misses), misses, fate,
+                saved / 1048576.0, copied / 1048576.0);
+        }
+    }
     if (!model.hparams.no_alloc) {
         for (size_t i = 0; i < backend_ptrs.size(); ++i) {
             ggml_backend_t             backend = backend_ptrs[i];
@@ -409,6 +419,65 @@ void llama_context::sched_reserve() {
 
     sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload));
 
+    // Auto-size expert cache: resolve -2 to a concrete slot count from available GPU VRAM.
+    if (cparams.expert_cache_n_slots == -2) {
+        const auto & hparams = model.hparams;
+        if (hparams.n_expert > 0) {
+            // Find first GPU backend to query free memory.
+            bool resolved = false;
+            for (const auto & backend : backends) {
+                auto * dev = ggml_backend_get_device(backend.get());
+                if (!dev || ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {
+                    continue;
+                }
+                size_t free_mem = 0, total_mem = 0;
+                ggml_backend_dev_memory(dev, &free_mem, &total_mem);
+
+                // Estimate per-expert size: gate_up fused = n_embd * 2 * n_ff_exp * type_size.
+                // Use bf16 (2 bytes) as conservative fallback.
+                const uint32_t n_ff_exp = hparams.n_ff_exp > 0
+                    ? hparams.n_ff_exp
+                    : (hparams.n_expert > 0 ? hparams.n_ff(0) / hparams.n_expert : 0);
+                if (n_ff_exp == 0) { break; }
+
+                const size_t expert_size_est = (size_t)hparams.n_embd * 2 * n_ff_exp * 2; // 2 bytes (bf16)
+
+                // Count MoE layers (layers with experts).
+                uint32_t n_moe_layers = 0;
+                for (uint32_t il = 0; il < hparams.n_layer; il++) {
+                    if (hparams.n_expert > 0) { n_moe_layers++; }
+                }
+                if (n_moe_layers == 0) { break; }
+
+                // Each MoE layer has ~2 expert weight tensors (gate_up + down).
+                const size_t per_slot_cost = (size_t)n_moe_layers * 2 * expert_size_est;
+
+                // Use 80% of free VRAM, bounded by [4, n_expert].
+                const size_t budget = (size_t)(free_mem * 0.8);
+                int32_t auto_slots = per_slot_cost > 0 ? (int32_t)(budget / per_slot_cost) : 0;
+                auto_slots = std::max(auto_slots, (int32_t)4);
+                auto_slots = std::min(auto_slots, (int32_t)hparams.n_expert);
+
+                cparams.expert_cache_n_slots = auto_slots;
+                LLAMA_LOG_INFO("%s: expert-cache auto-sized to %d slots (%.0f MiB free, %.1f MiB/slot)\n",
+                    __func__, auto_slots, free_mem / 1048576.0, per_slot_cost / 1048576.0);
+                resolved = true;
+                break;
+            }
+            if (!resolved) {
+                LLAMA_LOG_WARN("%s: expert-cache auto: no GPU found or invalid model params, disabling\n", __func__);
+                cparams.expert_cache_n_slots = 0;
+            }
+        } else {
+            LLAMA_LOG_WARN("%s: expert-cache auto: model has no experts, disabling\n", __func__);
+            cparams.expert_cache_n_slots = 0;
+        }
+    }
+
+    if (cparams.expert_cache_n_slots != 0) {
+        ggml_backend_sched_set_expert_cache(sched.get(), cparams.expert_cache_n_slots);
+    }
+
     llama_memory_context_ptr mctx;
     if (memory) {
         LLAMA_LOG_DEBUG("%s: reserving full memory module\n", __func__);
@@ -562,6 +631,9 @@ 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));
+                if (cparams.expert_cache_n_slots != 0) {
+                    ggml_backend_sched_set_expert_cache(sched.get(), cparams.expert_cache_n_slots);
+                }
                 gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get());
             }
             if (!gf) {
@@ -2217,6 +2289,27 @@ llm_graph_cb llama_context::graph_get_cb() const {
                 }
             }
         }
+
+        // Expert cache: when expert weights are on CPU (via --n-cpu-moe / tensor_buft_overrides)
+        // the scheduler normally assigns MUL_MAT_ID to CPU, preventing cross-backend copies
+        // and leaving the expert weight cache inactive.  Force MUL_MAT_ID to the layer's GPU
+        // backend so the scheduler creates a CPU→GPU copy that the cache can intercept.
+        if (cparams.expert_cache_n_slots != 0 &&
+                cur->op == GGML_OP_MUL_MAT_ID &&
+                il >= 0 &&
+                cur->src[0] != nullptr &&
+                cur->src[0]->buffer != nullptr &&
+                ggml_backend_buffer_is_host(cur->src[0]->buffer) &&
+                ggml_backend_buffer_get_usage(cur->src[0]->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
+            const auto & dev_layer = model.dev_layer(il);
+            for (const auto & backend : backends) {
+                if (ggml_backend_get_device(backend.get()) == dev_layer &&
+                        ggml_backend_supports_op(backend.get(), cur)) {
+                    ggml_backend_sched_set_tensor_backend(sched.get(), cur, backend.get());
+                    break;
+                }
+            }
+        }
     };
 }
 
@@ -2910,6 +3003,7 @@ llama_context_params llama_context_default_params() {
         /*.offload_kqv                 =*/ true,
         /*.no_perf                     =*/ true,
         /*.op_offload                  =*/ true,
+        /*.expert_cache_n_slots         =*/ 0,
         /*.swa_full                    =*/ true,
         /*.kv_unified                  =*/ false,
         /*.sampler                     =*/ nullptr,
@@ -3068,6 +3162,21 @@ void llama_synchronize(llama_context * ctx) {
     ctx->synchronize();
 }
 
+void llama_expert_cache_stats(
+        const llama_context * ctx,
+        int64_t * n_hits,
+        int64_t * n_misses,
+        int64_t * n_fate_hits,
+        int64_t * bytes_saved,
+        int64_t * bytes_copied) {
+    ggml_backend_sched_get_expert_cache_stats(
+        ctx->get_sched(), n_hits, n_misses, n_fate_hits, bytes_saved, bytes_copied);
+}
+
+void llama_expert_cache_stats_reset(llama_context * ctx) {
+    ggml_backend_sched_reset_expert_cache_stats(ctx->get_sched());
+}
+
 float * llama_get_logits(llama_context * ctx) {
     ctx->synchronize();
 

src/llama-cparams.h

@@ -37,6 +37,7 @@ struct llama_cparams {
     bool no_perf;
     bool warmup;
     bool op_offload;
+    int32_t expert_cache_n_slots;
     bool kv_unified;
     bool pipeline_parallel;
 

tests/CMakeLists.txt

@@ -209,6 +209,7 @@ llama_build_and_test(
     peg-parser/tests.h
 )
 llama_build_and_test(test-regex-partial.cpp)
+llama_build_and_test(test-expert-cache.cpp)
 
 if (NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "s390x")
     set(MODEL_NAME "tinyllamas/stories15M-q4_0.gguf")