llama-hot-experts: pin hottest MoE experts in RAM via --pin-hot-experts

common/arg.cpp

@@ -2576,13 +2576,15 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         "- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)\n"
         "- mlock: force system to keep model in RAM rather than swapping or compressing\n"
         "- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing\n"
-        "- dio: use DirectIO if available\n",
+        "- dio: use DirectIO if available\n"
+        "- mmap+pin: mmap without mlock; only the N hottest MoE experts are mlocked at runtime (use with --pin-hot-experts)\n",
         [](common_params & params, const std::string & value) {
             /**/ if (value == "none")       { params.load_mode = LLAMA_LOAD_MODE_NONE;       }
             else if (value == "mmap")       { params.load_mode = LLAMA_LOAD_MODE_MMAP;       }
             else if (value == "mlock")      { params.load_mode = LLAMA_LOAD_MODE_MLOCK;      }
             else if (value == "mmap+mlock") { params.load_mode = LLAMA_LOAD_MODE_MMAP_MLOCK; }
             else if (value == "dio")        { params.load_mode = LLAMA_LOAD_MODE_DIRECT_IO;  }
+            else if (value == "mmap+pin")   { params.load_mode = LLAMA_LOAD_MODE_MMAP_PIN;   }
             else { throw std::invalid_argument("invalid value"); }
         }
     ).set_env("LLAMA_ARG_LOAD_MODE"));
@@ -2645,6 +2647,57 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             }
         }
     ).set_env("LLAMA_ARG_N_CPU_MOE"));
+    add_opt(common_arg(
+        {"--pin-hot-experts"}, "N",
+        string_format(
+            "lock the N most frequently used MoE experts per layer into RAM in place\n"
+            "(mlock on their existing weight tensors, no copy) so the OS cannot evict\n"
+            "them; ranking is GLOBAL across all layers (total slots = N x num_moe_layers),\n"
+            "the hot set is tracked dynamically from actual router decisions and refreshed\n"
+            "on the fly. Only affects experts kept in host (CPU) memory\n"
+            "(default: %d, 0 = disabled; incompatible with a custom eval callback)",
+            params.n_pin_hot_experts
+        ),
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("error: --pin-hot-experts must be >= 0");
+            }
+            params.n_pin_hot_experts = value;
+        }
+    ).set_env("LLAMA_ARG_PIN_HOTEXPERTS"));
+    add_opt(common_arg(
+        {"--pin-hot-experts-budget-mib"}, "N",
+        string_format(
+            "hard cap, in MiB, on total memory locked by --pin-hot-experts across ALL layers\n"
+            "combined (default: %" PRIu64 ", 0 = unlimited). mlock() faults pages into RAM as\n"
+            "part of locking them, so leaving this unlimited on a large model/N can get the\n"
+            "process killed by the OOM killer instead of --pin-hot-experts simply having no\n"
+            "effect. Leave enough headroom for the KV cache and compute buffers, e.g. total\n"
+            "RAM minus model size minus expected KV cache / activation memory",
+            params.n_pin_hot_experts_budget_mib
+        ),
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("error: --pin-hot-experts-budget-mib must be >= 0");
+            }
+            params.n_pin_hot_experts_budget_mib = (uint64_t) value;
+        }
+    ).set_env("LLAMA_ARG_PIN_HOTEXPERTS_BUDGET_MIB"));
+    add_opt(common_arg(
+        {"--pin-hot-experts-stats-interval"}, "N",
+        string_format(
+            "print hot-expert pinning stats (bytes locked, global pin counts, per-layer\n"
+            "breakdown) directly to stderr every N router observations (default: %" PRIu64 ", 0 = disabled --\n"
+            "a final summary is still printed when the context is destroyed)",
+            params.n_pin_hot_experts_stats_interval
+        ),
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("error: --pin-hot-experts-stats-interval must be >= 0");
+            }
+            params.n_pin_hot_experts_stats_interval = (uint64_t) value;
+        }
+    ).set_env("LLAMA_ARG_PIN_HOTEXPERTS_STATS_INTERVAL"));
     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

@@ -1651,6 +1651,9 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.flash_attn_type   = params.flash_attn_type;
     cparams.cb_eval           = params.cb_eval;
     cparams.cb_eval_user_data = params.cb_eval_user_data;
+    cparams.n_pin_hot_experts  = params.n_pin_hot_experts;
+    cparams.n_pin_hot_experts_budget_bytes = params.n_pin_hot_experts_budget_mib * 1024ull * 1024ull;
+    cparams.n_pin_hot_experts_stats_interval = params.n_pin_hot_experts_stats_interval;
     cparams.offload_kqv       = !params.no_kv_offload;
     cparams.no_perf           = params.no_perf;
     cparams.op_offload        = !params.no_op_offload;

common/common.h

@@ -466,6 +466,17 @@ struct common_params {
     float   yarn_beta_slow        = -1.0f; // YaRN high correction dim
     int32_t yarn_orig_ctx         =     0; // YaRN original context length
 
+    // number of hottest MoE experts to mlock in place per layer, based on observed
+    // router usage, ranked GLOBALLY across all layers (total slots = N x num_moe_layers,
+    // 0 = disabled). See --pin-hot-experts.
+    int32_t n_pin_hot_experts      =     0;
+    // hard cap in MiB on total memory locked by n_pin_hot_experts, across all layers
+    // combined (0 = unlimited, NOT recommended -- see --pin-hot-experts-budget-mib).
+    uint64_t n_pin_hot_experts_budget_mib = 0;
+    // print hot-expert pinning stats to stderr every N router observations (0 = only
+    // at teardown). See --pin-hot-experts-stats-interval.
+    uint64_t n_pin_hot_experts_stats_interval = 200;
+
     // offload params
     std::vector<ggml_backend_dev_t> devices; // devices to use for offloading
 

include/llama.h

@@ -208,6 +208,7 @@ extern "C" {
         LLAMA_LOAD_MODE_MLOCK      = 2, // force system to keep model in RAM rather than swapping or compressing
         LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing
         LLAMA_LOAD_MODE_DIRECT_IO  = 4, // use direct I/O if available
+        LLAMA_LOAD_MODE_MMAP_PIN   = 5, // mmap, do NOT mlock all weights; only hot experts are mlocked at runtime (use with n_pin_hot_experts)
     };
 
     LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode);
@@ -376,6 +377,29 @@ extern "C" {
         ggml_backend_sched_eval_callback cb_eval;
         void * cb_eval_user_data;
 
+        // number of hottest MoE experts to mlock per layer (total slots = N x
+        // num_moe_layers), ranked GLOBALLY across all layers inside their
+        // existing host-memory weight tensors, based on observed router usage
+        // (0 = disabled). No effect on experts offloaded to a non-host buffer.
+        // incompatible with a caller-supplied cb_eval, since only one eval
+        // callback can be installed at a time (a warning is logged and pinning
+        // is skipped in that case).
+        int32_t n_pin_hot_experts;
+
+        // hard cap, in bytes, on total memory mlock'd by n_pin_hot_experts across
+        // ALL layers combined (0 = unlimited). Because mlock() faults pages into
+        // RAM as part of locking them, leaving this at 0 with a large model/N
+        // can cause the OS to kill the process for memory exhaustion rather than
+        // n_pin_hot_experts simply having no effect -- setting an explicit budget
+        // that leaves headroom for the KV cache and compute buffers is strongly
+        // recommended whenever n_pin_hot_experts > 0.
+        uint64_t n_pin_hot_experts_budget_bytes;
+
+        // print hot-expert pinning stats to stderr every N router observations
+        // (0 = disabled, a final summary is still printed when the context is
+        // destroyed). Bypasses the log callback and writes directly with fprintf.
+        uint64_t n_pin_hot_experts_stats_interval;
+
         enum ggml_type type_k; // data type for K cache [EXPERIMENTAL]
         enum ggml_type type_v; // data type for V cache [EXPERIMENTAL]
 

src/CMakeLists.txt

@@ -19,6 +19,7 @@ add_library(llama
             llama-cparams.cpp
             llama-grammar.cpp
             llama-graph.cpp
+            llama-hot-experts.cpp
             llama-hparams.cpp
             llama-impl.cpp
             llama-io.cpp

src/llama-context.cpp

@@ -3,6 +3,7 @@
 #include "ggml.h"
 #include "llama-arch.h"
 #include "llama-graph.h"
+#include "llama-hot-experts.h"
 #include "llama-impl.h"
 #include "llama-batch.h"
 #include "llama-io.h"
@@ -136,6 +137,22 @@ llama_context::llama_context(
 
     cparams.cb_eval           = params.cb_eval;
     cparams.cb_eval_user_data = params.cb_eval_user_data;
+    cparams.n_pin_hot_experts  = params.n_pin_hot_experts;
+    cparams.n_pin_hot_experts_budget_bytes = params.n_pin_hot_experts_budget_bytes;
+    cparams.n_pin_hot_experts_stats_interval = params.n_pin_hot_experts_stats_interval;
+
+    if (cparams.n_pin_hot_experts > 0) {
+        if (cparams.cb_eval != nullptr) {
+            LLAMA_LOG_WARN("%s: --pin-hot-experts requires the eval callback slot, but a custom cb_eval "
+                            "was already supplied; hot-expert pinning is disabled\n", __func__);
+        } else {
+            hot_experts = std::make_unique<llama_hot_expert_cache>(
+                model, cparams.n_pin_hot_experts, cparams.n_pin_hot_experts_budget_bytes,
+                cparams.n_pin_hot_experts_stats_interval);
+            cparams.cb_eval           = llama_hot_expert_cache::eval_callback;
+            cparams.cb_eval_user_data = hot_experts.get();
+        }
+    }
 
     cparams.ctx_other = nullptr;
 
@@ -3501,6 +3518,9 @@ llama_context_params llama_context_default_params() {
         /*.defrag_thold                =*/ -1.0f,
         /*.cb_eval                     =*/ nullptr,
         /*.cb_eval_user_data           =*/ nullptr,
+        /*.n_pin_hot_experts            =*/ 0,
+        /*.n_pin_hot_experts_budget_bytes=*/ 0,
+        /*.n_pin_hot_experts_stats_interval=*/ 200,
         /*.type_k                      =*/ GGML_TYPE_F16,
         /*.type_v                      =*/ GGML_TYPE_F16,
         /*.abort_callback              =*/ nullptr,

src/llama-context.h

@@ -16,6 +16,7 @@
 
 struct llama_model;
 class llama_batch_allocr;
+class llama_hot_expert_cache;
 
 class llama_io_read_i;
 class llama_io_write_i;
@@ -281,6 +282,10 @@ struct llama_context {
 
     llama_cparams cparams;
 
+    // --pin-hot-experts N: mlocks the N hottest MoE experts per layer in place, in RAM
+    // (null when disabled, i.e. n_pin_hot_experts <= 0 or a custom cb_eval was supplied)
+    std::unique_ptr<llama_hot_expert_cache> hot_experts;
+
     llama_adapter_cvec_ptr  cvec;
     llama_adapter_loras_ptr loras;
 

src/llama-cparams.h

@@ -61,5 +61,13 @@ struct llama_cparams {
     ggml_backend_sched_eval_callback cb_eval;
     void * cb_eval_user_data;
 
+    // --pin-hot-experts N: number of hottest MoE experts to mlock per layer, ranked
+    // GLOBALLY across all layers (total slots = N x num_moe_layers, 0 = disabled)
+    int32_t n_pin_hot_experts;
+    // hard cap in bytes on total memory locked by n_pin_hot_experts, across all layers (0 = unlimited)
+    uint64_t n_pin_hot_experts_budget_bytes;
+    // print pinning stats to stderr every N router observations (0 = only at context teardown)
+    uint64_t n_pin_hot_experts_stats_interval;
+
     llama_context * ctx_other;
 };