cuda: MoE expert cache, adaptive VRAM caching of CPU-resident experts

common/arg.cpp

@@ -26,6 +26,7 @@
 #include <cinttypes>
 #include <climits>
 #include <cstdarg>
+#include <cstdlib>
 #include <fstream>
 #include <list>
 #include <regex>
@@ -2348,6 +2349,31 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             }
         }
     ).set_env("LLAMA_ARG_N_CPU_MOE"));
+    add_opt(common_arg(
+        {"--moe-cache"}, "N",
+        "adaptively cache the hottest CPU-resident MoE experts in spare VRAM "
+        "(default: auto; 0 = off; N = VRAM budget in MiB per device)",
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("invalid value");
+            }
+            // the cache lives in the CUDA backend and configures itself from the
+            // environment at backend-registration time
+            auto set_env_var = [](const char * name, const char * val) {
+#if defined(_WIN32)
+                _putenv_s(name, val);
+#else
+                setenv(name, val, 1);
+#endif
+            };
+            if (value == 0) {
+                set_env_var("GGML_CUDA_MOE_CACHE", "0");
+            } else {
+                set_env_var("GGML_CUDA_MOE_CACHE_BUDGET_MB", std::to_string(value).c_str());
+            }
+            (void) params;
+        }
+    ).set_env("LLAMA_ARG_MOE_CACHE"));
     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/fit.cpp

@@ -1,4 +1,8 @@
 #include "fit.h"
+#include "gguf.h"
+
+#include <cstring>
+#include <sys/stat.h>
 
 #include "log.h"
 
@@ -150,6 +154,88 @@ std::vector<llama_device_memory_data> common_get_device_memory_data(
     return ret;
 }
 
+// MoE expert cache placement preference: when a MoE model's experts cannot
+// all return to VRAM anyway, leaving ALL of them on the CPU and giving the
+// spare VRAM to the dynamic expert cache measures faster than any static
+// partial placement (754B: 19.2 vs 14.0 t/s; 397B: 33.3 vs 28.2). Only models
+// above the cache's expert-size gate qualify — below it the cache would stay
+// dormant and a static placement is strictly better.
+// scan one gguf file for a routed-expert tensor; returns -1 if none found,
+// else per-expert KiB (needs n_expert > 0)
+static long common_moe_cache_expert_kib_in_file(const char * path, int64_t n_expert) {
+    struct gguf_init_params ip = { /*no_alloc=*/ true, /*ctx=*/ nullptr };
+    struct gguf_context * gctx = gguf_init_from_file(path, ip);
+    if (!gctx) return -1;
+    long kib = -1;
+    const int64_t n_tensors = gguf_get_n_tensors(gctx);
+    for (int64_t i = 0; i < n_tensors; i++) {
+        const char * name = gguf_get_tensor_name(gctx, i);
+        if (!strstr(name, "ffn_up_exps") && !strstr(name, "ffn_gate_exps")) continue;
+        if (n_expert > 0) {
+            kib = (long)(gguf_get_tensor_size(gctx, i) / n_expert / 1024);
+        }
+        break;
+    }
+    gguf_free(gctx);
+    return kib;
+}
+
+static bool common_moe_cache_prefers_cpu_moe(const char * path_model, size_t usable_vram, size_t n_devices) {
+    const char * e = getenv("GGML_CUDA_MOE_CACHE");
+    if (e && atoi(e) == 0) return false;            // explicitly off
+    long min_kb = 256;
+    if (const char * m = getenv("GGML_CUDA_MOE_CACHE_MIN_EXPERT_KB")) min_kb = atol(m);
+
+    // expert count from part-1 metadata
+    int64_t n_expert = 0;
+    {
+        struct gguf_init_params ip = { /*no_alloc=*/ true, /*ctx=*/ nullptr };
+        struct gguf_context * gctx = gguf_init_from_file(path_model, ip);
+        if (!gctx) return false;
+        const int64_t kid = gguf_find_key(gctx, "general.architecture");
+        if (kid >= 0) {
+            std::string arch = gguf_get_val_str(gctx, kid);
+            const int64_t kex = gguf_find_key(gctx, (arch + ".expert_count").c_str());
+            if (kex >= 0) n_expert = gguf_get_val_u32(gctx, kex);
+        }
+        gguf_free(gctx);
+    }
+    if (n_expert <= 0) return false;
+
+    long kib = common_moe_cache_expert_kib_in_file(path_model, n_expert);
+    size_t model_bytes = 0;
+    {
+        // total bytes across split parts (placement economics needs the ratio)
+        const char * tag = strstr(path_model, "-00001-of-");
+        int n_parts = 1;
+        if (tag) n_parts = atoi(tag + strlen("-00001-of-"));
+        for (int part = 1; part <= (n_parts > 0 ? n_parts : 1); part++) {
+            std::string p(path_model);
+            if (tag) {
+                char rep_[16];
+                snprintf(rep_, sizeof(rep_), "-%05d-of-", part);
+                p.replace(tag - path_model, strlen("-00001-of-"), rep_);
+            }
+            struct stat st;
+            if (stat(p.c_str(), &st) == 0) model_bytes += (size_t)st.st_size;
+            if (kib < 0 && part >= 2) {
+                kib = common_moe_cache_expert_kib_in_file(p.c_str(), n_expert);
+            }
+        }
+    }
+    if (kib < min_kb) return false;
+
+    // Placement economics, calibrated on 6 measured configs (see
+    // MOE_CACHE_READINESS.md + matrix): the dynamic cache beats a static partial
+    // placement only for LARGE spills (static could fit < ~55% of the model),
+    // and on few devices only with big experts (per-device dispatch-chain
+    // serialization: 122B/MiniMax on one 3090 lose 6-18%, 4x 3090 wins).
+    const double spill_ratio = usable_vram > 0 ? (double)model_bytes / (double)usable_vram : 99.0;
+    if (spill_ratio < 1.8) return false;
+    if (n_devices < 2 && kib < 2048) return false;
+    return true;
+}
+
 static void common_params_fit_impl(
         const char * path_model, struct llama_model_params * mparams, struct llama_context_params * cparams,
         float * tensor_split, struct llama_model_tensor_buft_override * tensor_buft_overrides,
@@ -619,6 +705,12 @@ static void common_params_fit_impl(
         return;
     }
 
+    // expert-cache placement preference: snapshot the all-experts-on-CPU
+    // placement; if step 4 cannot return ALL experts to VRAM we prefer this
+    // snapshot (spare VRAM goes to the dynamic cache instead of a static mix)
+    const std::vector<ngl_t> ngl_all_cpu_moe = ngl_per_device;
+    const std::vector<ggml_backend_buffer_type_t> overflow_bufts_cpu_moe = overflow_bufts;
+
     // step 4: for a MoE model where all dense tensors fit,
     //     convert the dense-only layers in the back to full layers in the front until all devices are full
     // essentially the same procedure as for the dense-only layers except front-to-back
@@ -760,6 +852,23 @@ static void common_params_fit_impl(
             __func__, dev_names[id].c_str(), ngl_per_device[id].n_layer, ngl_per_device[id].n_part, mem[id]/MiB, projected_margin/MiB);
     }
 
+    {
+        uint32_t n_dense_only = 0;
+        for (size_t id = 0; id < nd; id++) {
+            n_dense_only += ngl_per_device[id].n_part;
+        }
+        size_t usable_vram = 0;
+        for (size_t id = 0; id < nd; id++) {
+            usable_vram += dmds_full[id].free > margins[id] ? dmds_full[id].free - margins[id] : 0;
+        }
+        if (n_dense_only > 0 && common_moe_cache_prefers_cpu_moe(path_model, usable_vram, nd)) {
+            LOG_INF("%s: experts cannot all fit in VRAM; expert cache active -> "
+                    "keeping ALL experts on CPU, spare VRAM goes to the dynamic cache "
+                    "(GGML_CUDA_MOE_CACHE=0 restores static placement)\n", __func__);
+            set_ngl_tensor_split_tbo(ngl_all_cpu_moe, overflow_bufts_cpu_moe, *mparams);
+            return;
+        }
+    }
     set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, *mparams);
 }
 

ggml/src/ggml-backend-moe-cache.h

@@ -0,0 +1,95 @@
+#pragma once
+
+#include <stdint.h>
+#include <stddef.h>
+
+// MoE expert cache — dynamic VRAM cache for MoE expert weights on CPU-resident
+// MUL_MAT_ID. Integration point is the CPU mul_mat_id kernel itself: thread 0
+// dispatches cached expert rows to the GPU while the remaining threads compute
+// the uncached rows, then results are collected into dst before the node ends.
+//
+// This table is the bridge between ggml-cpu (which cannot link CUDA) and the
+// CUDA backend (which registers the implementations at backend-reg time).
+// begin/plan/dispatch/collect for one node are called from a single thread
+// (ith == 0); the implementation may assume no concurrent node processing.
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+struct ggml_moe_cache_api {
+    // Decide whether the cache engages for this MUL_MAT_ID node.
+    // Returns the device id to use (>= 0) or -1 to stay on the pure-CPU path.
+    // Performs lazy per-device initialization on first use and selects the
+    // internal slot pool matching (expert_size, wtype).
+    //   tensor_name: src0->name (stable cache key source, e.g. "blk.7.ffn_up_exps.weight")
+    //   host_base:   src0->data (source for async inserts)
+    //   expert_size: src0->nb[2] (bytes per expert)
+    //   n_in/n_out:  src0->ne[0] / src0->ne[1]
+    //   wtype:       src0->type
+    //   n_expert:    src0->ne[2]
+    //   n_tokens:    ids->ne[1] (cache engages only when == 1)
+    int (*begin)(const char * tensor_name, const void * host_base, size_t expert_size,
+                 int64_t n_in, int64_t n_out, int wtype, int64_t n_expert, int64_t n_tokens);
+
+    // For each of the n_ids expert ids: slot_idx[k] = cache slot index (hit) or
+    // -1 (miss; CPU computes the row, an async insert may be enqueued).
+    // Returns the number of hits.
+    int (*plan)(int dev, const int32_t * ids, int n_ids, int32_t * slot_idx);
+
+    // One batched GPU launch computing all n_hits rows:
+    //   out_row[i] (n_out floats) = W[slot_idx_compact[i]] . act_rows[i]
+    // act_rows are host fp32 pointers (they may all be the same row for
+    // gate/up-style nodes; distinct rows for down-style nodes).
+    // Asynchronous; results are pulled into dst by collect().
+    void (*dispatch)(int dev, int wtype, int64_t n_in, int64_t n_out, int n_hits,
+                     const int32_t * slot_idx_compact, const float * const * act_rows);
+
+    // Synchronize the device's compute stream and copy the n_hits result rows
+    // (in dispatch order) into dst_rows[0..n_hits-1] (n_out floats each).
+    void (*collect)(int dev, int n_hits, float * const * dst_rows, int64_t n_out);
+
+    // Periodic stats logging (rate-limited internally).
+    void (*stats)(void);
+
+    // ---- GPU-resident dst handoff (down-projection round-trip elimination) ----
+    // The scheduler offers the GPU-side copy tensor of a CPU MUL_MAT_ID dst
+    // BEFORE the CPU split runs. If the cache then computes that node, it
+    // scatters its GPU rows directly into gpu_copy_data (async) instead of
+    // pulling them to the host.
+    void (*redirect_offer)(const void * host_dst_data, size_t nb1, int64_t n_rows,
+                           void * gpu_copy_data, void * consumer_backend);
+    // Called by the scheduler at the consumer split's input-copy site, after
+    // the CPU split completed. Returns 1 if the cache fully populated the GPU
+    // copy (it uploads the CPU-computed miss rows here and installs a
+    // stream-order dependency on the consumer backend) — the scheduler must
+    // then SKIP its own copy. Returns 0 for the normal copy path.
+    int (*redirect_finalize)(const void * host_dst_data, void * consumer_backend);
+
+    // ---- fused gate+up+GLU path ----
+    // Called by the CPU GLU kernel (swiglu split variant) from EVERY thread:
+    // returns the bitmask of dst rows the cache computed on the GPU (fused
+    // silu(gate)*up) — those rows must be SKIPPED by the CPU loop. Thread 0
+    // (ith == 0) additionally synchronizes the fused chain and scatters the
+    // GPU rows into dst before returning. Returns 0 when the cache has no
+    // pending fused work for this (src0, src1) pair.
+    unsigned long long (*glu_hits)(const void * src0_data, const void * src1_data,
+                                   void * dst_data, size_t dst_nb1, int ith);
+
+    // Host weight buffer teardown notification (model unload): drops queued
+    // insert jobs sourced from the range and resets per-block tensor-base
+    // learning. Must be called before the memory is unmapped/freed.
+    void (*invalidate)(const void * base, size_t size);
+
+    // Node wall-time sample for the bail-out judge. code is begin()'s return
+    // value: -3 = pure-CPU baseline sample, >= 0 = cache-engaged sample.
+    void (*node_time)(int code, int64_t wall_us);
+};
+
+// Zero-initialized in ggml-backend.cpp; populated by the CUDA backend in
+// ggml_backend_cuda_reg() when the cache is enabled (GGML_CUDA_MOE_CACHE=1).
+extern struct ggml_moe_cache_api ggml_moe_cache;
+
+#ifdef __cplusplus
+}
+#endif

ggml/src/ggml-backend.cpp

@@ -12,6 +12,11 @@
 #include "ggml-backend-impl.h"
 #include "ggml-alloc.h"
 #include "ggml-impl.h"
+#include "ggml-backend-moe-cache.h"
+
+// MoE expert cache function table; populated by the CUDA backend at registry
+// init when GGML_CUDA_MOE_CACHE=1, consumed by the CPU mul_mat_id kernel.
+struct ggml_moe_cache_api ggml_moe_cache = {};
 
 #include <assert.h>
 #include <limits.h>
@@ -109,6 +114,15 @@ void ggml_backend_buffer_free(ggml_backend_buffer_t buffer) {
         return;
     }
 
+    // MoE expert cache: host weight buffers can back queued cache-fill jobs;
+    // notify before the memory goes away (no-op when the cache is inactive)
+    if (ggml_moe_cache.invalidate && buffer->iface.get_base && ggml_backend_buffer_is_host(buffer)) {
+        void * base = ggml_backend_buffer_get_base(buffer);
+        if (base) {
+            ggml_moe_cache.invalidate(base, ggml_backend_buffer_get_size(buffer));
+        }
+    }
+
     if (buffer->iface.free_buffer != NULL) {
         buffer->iface.free_buffer(buffer);
     }
@@ -1659,6 +1673,15 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                     }
                     copy_experts(first_id, last_id);
                 } else {
+                    // MoE expert cache dst handoff: if the cache populated this
+                    // input's GPU copy directly (down-projection rows), skip the
+                    // round-trip copy AND the blocking sync — the cache installed
+                    // a stream-order dependency on this backend instead.
+                    if (ggml_moe_cache.redirect_finalize &&
+                        ggml_moe_cache.redirect_finalize(input->data, split_backend)) {
+                        continue;
+                    }
+
                     // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events
                     // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface
                     if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) {
@@ -1674,6 +1697,42 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
             }
         }
 
+        // MoE expert cache dst handoff: before running a CPU split that ends in a
+        // MUL_MAT_ID, offer the cache the GPU-side copy tensor of that dst so it
+        // can scatter its GPU-computed rows directly (skipping the host round
+        // trip). Engaged only with a single, unique CUDA consumer and one copy.
+        // note: the CPU backend is always the last one (asserted in sched_new);
+        // do NOT use ggml_backend_dev_type here — the CUDA implementation calls
+        // cudaGetDeviceProperties (~ms per call!) and this runs per split.
+        if (ggml_moe_cache.redirect_offer && sched->n_copies == 1 &&
+            split->graph.n_nodes > 0 &&
+            split_backend_id == sched->n_backends - 1) {
+            ggml_tensor * last = split->graph.nodes[split->graph.n_nodes - 1];
+            if (last->op == GGML_OP_MUL_MAT_ID && !(last->flags & GGML_TENSOR_FLAG_OUTPUT)) {
+                int consumer_split = -1;
+                int n_consumers = 0;
+                for (int j = split_id + 1; j < sched->n_splits && n_consumers < 2; j++) {
+                    for (int k = 0; k < splits[j].n_inputs; k++) {
+                        if (splits[j].inputs[k] == last) {
+                            n_consumers++;
+                            consumer_split = j;
+                            break;
+                        }
+                    }
+                }
+                if (n_consumers == 1) {
+                    ggml_backend_t cons_backend = sched->backends[splits[consumer_split].backend_id];
+                    ggml_tensor * cpy = tensor_copy(last, splits[consumer_split].backend_id, sched->cur_copy);
+                    if (cpy && cpy->data && splits[consumer_split].backend_id != sched->n_backends - 1 &&
+                        ggml_is_contiguous(last)) {
+                        ggml_moe_cache.redirect_offer(last->data, last->nb[1],
+                                                            last->ne[1] * last->ne[2],
+                                                            cpy->data, cons_backend);
+                    }
+                }
+            }
+        }
+
         if (!sched->callback_eval) {
             enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph);
             if (ec != GGML_STATUS_SUCCESS) {