From 79a945f768c9c00612dab63d0b5a24892bcbb78c Mon Sep 17 00:00:00 2001 From: Ghimli Date: Sat, 1 Aug 2026 22:33:21 +0200 Subject: [PATCH] llama-hot-experts: pin hottest MoE experts in RAM via --pin-hot-experts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Problem: Large MoE models loaded via mmap suffer from OS page cache eviction of infrequently accessed experts, causing severe latency spikes when those experts are needed later. Solution: --pin-hot-experts N dynamically tracks expert usage at runtime and uses mlock() to pin the top-N most frequently used experts in RAM, preventing OS eviction. Cold experts remain paged normally. Dense parts are mlocked automatically: Before any hot experts are pinned, all dense (non-MoE-expert) tensors that reside in host memory are mlocked in place. Dense parts (embeddings, attention/FFN weights, RMSNorm, output projection, router weights, etc.) are used on every single token — they are the hottest data by definition. This happens unconditionally when --pin-hot-experts is enabled and mlock is supported, consuming the global budget first so that hot experts only ever get the leftover budget. Key features: - Real-time tracking via tensor callback on ffn_moe_topk tensors - On-the-fly eviction/replacement of cold experts with hot ones - Budget cap via --pin-hot-experts-budget-mib - Fail-safe rollback on mlock failure - Per-layer stats reporting at --pin-hot-experts-stats-interval N - New load mode --load-mode mmap+pin (mmap without global mlock) CLI arguments: --pin-hot-experts Number of hot experts to pin (0=off) --pin-hot-experts-budget-mib Max pinned memory in MiB (0=unlimited) --pin-hot-experts-stats-interval N Print stats every N tokens --load-mode mmap+pin mmap without global mlock Naming convention: C API: underscores (n_pin_hot_experts) CLI: hyphens (--pin-hot-experts) --- common/arg.cpp | 55 ++++- common/common.cpp | 3 + common/common.h | 11 + include/llama.h | 24 +++ src/CMakeLists.txt | 1 + src/llama-context.cpp | 20 ++ src/llama-context.h | 5 + src/llama-cparams.h | 8 + src/llama-hot-experts.cpp | 433 ++++++++++++++++++++++++++++++++++++++ src/llama-hot-experts.h | 172 +++++++++++++++ src/llama-mmap.cpp | 6 + src/llama-mmap.h | 8 + src/llama.cpp | 13 +- 13 files changed, 753 insertions(+), 6 deletions(-) create mode 100644 src/llama-hot-experts.cpp create mode 100644 src/llama-hot-experts.h diff --git a/common/arg.cpp b/common/arg.cpp index 79480e06f9d2..14993270021d 100644 --- a/common/arg.cpp +++ b/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", diff --git a/common/common.cpp b/common/common.cpp index ff27d392fb2e..9e6fb6e75f14 100644 --- a/common/common.cpp +++ b/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; diff --git a/common/common.h b/common/common.h index 919c0ea103a4..0c89262d74c3 100644 --- a/common/common.h +++ b/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 devices; // devices to use for offloading diff --git a/include/llama.h b/include/llama.h index 6e53e2297235..90e896871941 100644 --- a/include/llama.h +++ b/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] diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 320784c3a8cc..1d6953d3a804 100644 --- a/src/CMakeLists.txt +++ b/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 diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 5ef7becf6f29..a605e73e56b2 100644 --- a/src/llama-context.cpp +++ b/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( + 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, diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa8b562..f305c388edf0 100644 --- a/src/llama-context.h +++ b/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 hot_experts; + llama_adapter_cvec_ptr cvec; llama_adapter_loras_ptr loras; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 5018170ed85e..96c5bdea8cdf 100644 --- a/src/llama-cparams.h +++ b/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; }; diff --git a/src/llama-hot-experts.cpp b/src/llama-hot-experts.cpp new file mode 100644 index 000000000000..a0d92cdf0d74 --- /dev/null +++ b/src/llama-hot-experts.cpp @@ -0,0 +1,433 @@ +#include "llama-hot-experts.h" + +#include "ggml-backend.h" +#include "llama-impl.h" +#include "llama-model.h" + +#include +#include +#include +#include +#include + +llama_hot_expert_cache::llama_hot_expert_cache(const llama_model & model, + int32_t n_pin_experts, + uint64_t budget_bytes, + uint64_t stats_interval) : + model(model), + n_pin(n_pin_experts), + budget_bytes(budget_bytes), + stats_interval(stats_interval) { + // Count MoE layers by checking which layers have ffn_down_exps.weight + int32_t n_moe_layers = 0; + for (int32_t il = 0; il < (int32_t) model.hparams.n_layer(); il++) { + const std::string name = "blk." + std::to_string(il) + ".ffn_down_exps.weight"; + if (model.get_tensor(name.c_str()) != nullptr) { + n_moe_layers++; + } + } + n_pin_total = n_pin * n_moe_layers; + + if (n_moe_layers == 0) { + LLAMA_LOG_WARN("%s: no MoE layers detected in model, --pin-hot-experts has no effect\n", __func__); + } else if (budget_bytes > 0) { + LLAMA_LOG_INFO( + "%s: pinning (mlock) up to %d hottest MoE experts per layer (%d MoE layers, " + "%d total global slots) in place, budget %.2f MiB total, pin/evict on the fly\n", + __func__, n_pin, n_moe_layers, n_pin_total, budget_bytes / (1024.0 * 1024.0)); + } else { + LLAMA_LOG_WARN( + "%s: --pin-hot-experts has NO memory budget cap (--pin-hot-experts-budget-mib " + "was not set); with enough layers/experts this WILL try to lock more memory " + "than physically fits and can be killed by the OOM killer. Setting an explicit " + "budget that leaves headroom for the KV cache and compute buffers is strongly " + "recommended.\n", + __func__); + } + if (stats_interval > 0) { + LLAMA_LOG_INFO("%s: printing pinning stats to stderr every %" PRIu64 " router observations\n", __func__, + stats_interval); + } + if (!llama_mlock::SUPPORTED) { + LLAMA_LOG_WARN( + "%s: mlock is not supported on this platform, --pin-hot-experts will only " + "track usage statistics and will not actually lock any memory\n", + __func__); + } + + // Dense parts are used on every token, so they are the hottest by definition: + // mlock all dense tensors that didn't fit in VRAM (i.e. live in host memory) + // FIRST, consuming the global budget before any MoE expert is pinned. The + // remaining budget is what the hot-expert pin/evict logic may use below. + lock_dense_parts(); + if (llama_mlock::SUPPORTED && n_dense_bytes_locked > 0) { + LLAMA_LOG_INFO( + "%s: mlocked %.2f MiB of dense (non-MoE) tensors in RAM first (used on every token); " + "%.2f MiB of the pin budget remains for hot MoE experts\n", + __func__, n_dense_bytes_locked / (1024.0 * 1024.0), + (budget_bytes == 0 ? UINT64_MAX : (budget_bytes - n_dense_bytes_locked)) / (1024.0 * 1024.0)); + } else if (llama_mlock::SUPPORTED) { + LLAMA_LOG_INFO("%s: no dense (non-MoE) tensors in host RAM to mlock (all offloaded to VRAM or mlock unsupported)\n", + __func__); + } +} + +void llama_hot_expert_cache::lock_dense_parts() { + if (!llama_mlock::SUPPORTED) { + return; // stats-only mode, nothing can be locked + } + + // A "dense" tensor is any model tensor that is NOT one of the MoE expert + // tensors (ffn_*_exps.*). Dense parts are used on every token, so they are + // the hottest by definition and must be mlocked before any hot expert. + auto is_expert_tensor = [](const std::string & name) -> bool { + return name.find("_exps.") != std::string::npos; + }; + + for (const auto & [name, t] : llama_internal_get_tensor_map(&model)) { + if (is_expert_tensor(name)) { + continue; // handled by the hot-expert pin/evict mechanism below + } + if (!ggml_backend_buffer_is_host(t->buffer) || t->data == nullptr) { + continue; // offloaded to VRAM (did fit) or empty: nothing to lock in RAM + } + + const size_t nbytes = ggml_nbytes(t); + + // honor the global budget BEFORE touching memory -- mlock() faults pages + // in, so checking after the fact is too late to prevent an OOM. Dense + // tensors are locked first, so the remaining budget is what hot experts + // may consume afterwards. + if (budget_bytes > 0 && n_bytes_locked + nbytes > budget_bytes) { + LLAMA_LOG_DEBUG( + "%s: skipping dense tensor %s (%zu bytes), would exceed the %.2f MiB pin budget " + "(%.2f MiB already locked)\n", + __func__, name.c_str(), nbytes, budget_bytes / (1024.0 * 1024.0), + n_bytes_locked / (1024.0 * 1024.0)); + continue; + } + + auto lock = std::unique_ptr(new llama_mlock()); + lock->init(t->data); + lock->grow_to(nbytes); + + const size_t locked = lock->size(); + if (locked == 0) { + continue; + } + if (locked < nbytes) { + LLAMA_LOG_WARN( + "%s: only locked %zu/%zu bytes for dense tensor %s (system out of lockable " + "memory?)\n", + __func__, locked, nbytes, name.c_str()); + } + + n_bytes_locked += locked; + n_dense_bytes_locked += locked; + dense_locks.push_back(std::move(lock)); + } +} + +llama_hot_expert_cache::~llama_hot_expert_cache() { + print_stats(); +} + +void llama_hot_expert_cache::print_stats() const { + std::lock_guard lock(mu); + + size_t total_distinct_seen = counts.size(); + size_t total_pinned = pinned.size(); + uint64_t global_coldest_count = UINT64_MAX; + uint64_t global_hottest_count = 0; + + // Per-layer breakdown of pinned experts + std::unordered_map pinned_per_layer; + for (const auto & [key, pe] : pinned) { + pinned_per_layer[key.layer]++; + } + + if (!pinned_rank.empty()) { + global_coldest_count = std::get<0>(*pinned_rank.begin()); + global_hottest_count = std::get<0>(*pinned_rank.rbegin()); + } + + LLAMA_LOG_INFO("[pin-hot-experts] obs=%" PRIu64 + " | locked=%.2f MiB (dense_in_ram=%.2f MiB + experts=%.2f MiB) | moe_layers=%zu | " + "pinned=%zu/%d (global, N=%d x layers=%zu) | distinct (layer,expert) seen=%zu", + n_eval_calls, n_bytes_locked / (1024.0 * 1024.0), + n_dense_bytes_locked / (1024.0 * 1024.0), + (n_bytes_locked - n_dense_bytes_locked) / (1024.0 * 1024.0), + layers.size(), total_pinned, n_pin_total, n_pin, + layers.size(), total_distinct_seen); + + if (!pinned_rank.empty()) { + LLAMA_LOG_CONT(" | pinned count range=[%" PRIu64 ", %" PRIu64 "]", global_coldest_count, global_hottest_count); + } + + if (!pinned_per_layer.empty()) { + LLAMA_LOG_CONT(" | per-layer: {"); + // Sort by layer index for readable output + std::vector> sorted_layers(pinned_per_layer.begin(), pinned_per_layer.end()); + std::sort(sorted_layers.begin(), sorted_layers.end()); + for (size_t i = 0; i < sorted_layers.size(); i++) { + const auto & [il, cnt] = sorted_layers[i]; + LLAMA_LOG_CONT("L%d=%zu%s", il, cnt, (i + 1 < sorted_layers.size()) ? ", " : ""); + } + LLAMA_LOG_CONT("}"); + } + LLAMA_LOG_CONT("\n"); +} + +bool llama_hot_expert_cache::eval_callback(struct ggml_tensor * t, bool ask, void * user_data) { + auto * self = static_cast(user_data); + + // we only need the *values* of the top-k expert-selection tensor, named + // "ffn_moe_topk-" by llm_graph_context::cb() -> llama_context::graph_get_cb() + static const char prefix[] = "ffn_moe_topk-"; + if (strncmp(t->name, prefix, sizeof(prefix) - 1) != 0) { + return ask ? false : true; // not interested, let the scheduler carry on either way + } + + if (ask) { + // request that the scheduler makes this tensor's data readable on the host + return true; + } + + const int il = atoi(t->name + sizeof(prefix) - 1); + self->on_topk_tensor(il, t); + + return true; +} + +void llama_hot_expert_cache::on_topk_tensor(int il, const struct ggml_tensor * t) { + if (t->type != GGML_TYPE_I32) { + return; // unexpected, be defensive rather than misinterpret bytes + } + + const int64_t n_expert_used = t->ne[0]; + const int64_t n_tokens = t->ne[1]; + const int64_t n_ids = n_expert_used * n_tokens; + + std::vector ids(n_ids); + if (ggml_backend_buffer_is_host(t->buffer)) { + std::memcpy(ids.data(), t->data, n_ids * sizeof(int32_t)); + } else { + ggml_backend_tensor_get(t, ids.data(), 0, n_ids * sizeof(int32_t)); + } + + uint64_t eval_calls_now = 0; + { + std::lock_guard lock(mu); + + auto & ls = layers[il]; + if (!ls.resolved_tensors) { + resolve_tensors(il, ls); + } + + for (int32_t id : ids) { + if (id < 0) { + continue; // padding / unused slot + } + observe_expert(il, ls, id); + } + + n_eval_calls++; + eval_calls_now = n_eval_calls; + } // lock released here + + // print_stats() takes the same mutex itself, so this must run outside the + // scope above (the mutex is not recursive) + if (stats_interval > 0 && eval_calls_now % stats_interval == 0) { + print_stats(); + } +} + +void llama_hot_expert_cache::resolve_tensors(int il, layer_state & ls) { + const std::string base = "blk." + std::to_string(il) + "."; + + ls.t_gate = model.get_tensor((base + "ffn_gate_exps.weight").c_str()); + ls.t_up = model.get_tensor((base + "ffn_up_exps.weight").c_str()); + ls.t_down = model.get_tensor((base + "ffn_down_exps.weight").c_str()); + ls.t_gate_up = model.get_tensor((base + "ffn_gate_up_exps.weight").c_str()); + + ls.resolved_tensors = true; + + if (!ls.t_down || (!ls.t_gate && !ls.t_gate_up)) { + LLAMA_LOG_WARN( + "%s: layer %d does not look like a (supported) MoE FFN layer, " + "hot-expert pinning disabled for this layer\n", + __func__, il); + return; + } + + // pinning only makes sense (and is only safe to do in place) when the expert + // tensors actually live in host memory, e.g. all experts kept on the CPU + // while only the dense/router parts are offloaded to VRAM + const ggml_tensor * repr = ls.t_down; + ls.tensors_are_host = ggml_backend_buffer_is_host(repr->buffer); + + if (!ls.tensors_are_host) { + LLAMA_LOG_WARN( + "%s: layer %d's MoE experts are not in host memory (offloaded to a " + "device buffer), --pin-hot-experts has no effect for this layer\n", + __func__, il); + } +} + +void llama_hot_expert_cache::observe_expert(int il, layer_state & ls, int32_t expert_id) { + expert_key key{ il, expert_id }; + + uint64_t & count = counts[key]; // default-constructs to 0 + const uint64_t old_count = count; + count++; + const uint64_t new_count = count; + + if (!ls.tensors_are_host || n_pin <= 0) { + return; // stats-only mode, nothing to pin + } + + if (pinned.count(key)) { + // already pinned: keep its ordered-set position up to date + pinned_rank.erase({ old_count, il, expert_id }); + pinned_rank.insert({ new_count, il, expert_id }); + return; + } + + if ((int32_t) pinned.size() < n_pin_total) { + // a pin slot is still free: pin immediately, no eviction needed + if (pin_expert(il, ls, expert_id)) { + pinned_rank.insert({ new_count, il, expert_id }); + } + return; + } + + // all N slots are taken: only take over if we just overtook the coldest pinned expert + const auto coldest = pinned_rank.begin(); + if (coldest != pinned_rank.end() && new_count > std::get<0>(*coldest)) { + const uint64_t evict_count = std::get<0>(*coldest); + const int evict_layer = std::get<1>(*coldest); + const int32_t evict_id = std::get<2>(*coldest); + + pinned_rank.erase(coldest); + + auto & evict_ls = layers[evict_layer]; + if (!evict_ls.resolved_tensors) { + resolve_tensors(evict_layer, evict_ls); + } + unpin_expert(evict_layer, evict_ls, evict_id); + + if (pin_expert(il, ls, expert_id)) { + pinned_rank.insert({ new_count, il, expert_id }); + } else { + // New expert failed to lock anything (budget exhausted, mlock error, etc.). + // Roll back the eviction: re-pin the old expert to keep the data structures + // and actual locked pages consistent. + pin_expert(evict_layer, evict_ls, evict_id); + pinned_rank.insert({ evict_count, evict_layer, evict_id }); + } + } +} + +size_t llama_hot_expert_cache::lock_expert_row(const struct ggml_tensor * w, + int32_t expert_id, + std::unique_ptr & out_lock) { + if (!w || expert_id < 0 || expert_id >= w->ne[2] || !llama_mlock::SUPPORTED) { + return 0; + } + if (!ggml_backend_buffer_is_host(w->buffer) || w->data == nullptr) { + return 0; + } + + const size_t nbytes = ggml_nbytes(w) / (size_t) w->ne[2]; + + // enforce the global budget BEFORE touching any memory -- mlock() itself + // faults pages in, so checking after the fact is too late to prevent an OOM + if (budget_bytes > 0 && n_bytes_locked + nbytes > budget_bytes) { + LLAMA_LOG_DEBUG( + "%s: skipping expert %d, would exceed the %.2f MiB pin budget " + "(%.2f MiB already locked)\n", + __func__, expert_id, budget_bytes / (1024.0 * 1024.0), n_bytes_locked / (1024.0 * 1024.0)); + return 0; + } + + const size_t offset = (size_t) expert_id * w->nb[2]; + void * ptr = (uint8_t *) w->data + offset; + + // lock the expert's rows IN PLACE inside the model's own tensor -- this is the + // exact memory ggml_mul_mat_id() reads during build_moe_ffn(), so the lock + // directly protects the data the compute graph actually uses. + out_lock.reset(new llama_mlock()); + out_lock->init(ptr); + out_lock->grow_to(nbytes); + + // only count what was ACTUALLY locked -- grow_to() silently stops (and logs a + // warning) on failure rather than throwing, so size() may be less than nbytes + const size_t locked = out_lock->size(); + if (locked < nbytes) { + LLAMA_LOG_WARN( + "%s: only locked %zu/%zu bytes for expert %d (system out of lockable " + "memory?) -- consider lowering --pin-hot-experts N or its budget\n", + __func__, locked, nbytes, expert_id); + } + if (locked == 0) { + out_lock.reset(); + } + + n_bytes_locked += locked; // update the running total immediately so sibling + // tensors of the SAME expert also respect the budget + + return locked; +} + +bool llama_hot_expert_cache::pin_expert(int il, layer_state & ls, int32_t expert_id) { + if (!ls.tensors_are_host) { + return false; + } + + expert_key key{ il, expert_id }; + if (pinned.count(key)) { + return true; + } + + pinned_expert pe; + + if (ls.t_gate_up) { + pe.nbytes_locked += lock_expert_row(ls.t_gate_up, expert_id, pe.gate_up_lock); + } else if (ls.t_gate) { + pe.nbytes_locked += lock_expert_row(ls.t_gate, expert_id, pe.gate_lock); + } + if (ls.t_up) { + pe.nbytes_locked += lock_expert_row(ls.t_up, expert_id, pe.up_lock); + } + pe.nbytes_locked += lock_expert_row(ls.t_down, expert_id, pe.down_lock); + + // n_bytes_locked was already updated incrementally inside lock_expert_row() + // (so sibling tensors of this same expert see an up-to-date budget) + + if (pe.nbytes_locked == 0) { + // Nothing was actually locked (budget exhausted, mlock not supported, etc.). + // Do not insert a dead entry into the pinned map. + return false; + } + + pinned.emplace(key, std::move(pe)); + return true; +} + +void llama_hot_expert_cache::unpin_expert(int il, layer_state & /*ls*/, int32_t expert_id) { + expert_key key{ il, expert_id }; + auto it = pinned.find(key); + if (it == pinned.end()) { + return; + } + + n_bytes_locked -= it->second.nbytes_locked; + pinned.erase(it); // pinned_expert's destructor releases the mlock guards +} + +bool llama_hot_expert_cache::is_pinned(int il, int32_t expert_id) const { + std::lock_guard lock(mu); + + expert_key key{ il, expert_id }; + return pinned.count(key) != 0; +} diff --git a/src/llama-hot-experts.h b/src/llama-hot-experts.h new file mode 100644 index 000000000000..81160236b753 --- /dev/null +++ b/src/llama-hot-experts.h @@ -0,0 +1,172 @@ +#pragma once + +// --pin-hot-experts N +// +// Tracks, GLOBALLY across all MoE layers, how often each (layer, expert) +// pair gets selected by the router (via ggml_backend_sched's eval callback +// on the "ffn_moe_topk-" nodes) and keeps the top N experts **per layer** +// (total capacity = N * num_moe_layers) locked in RAM with mlock()/VirtualLock(), +// IN PLACE inside the model's own weight tensors (ffn_gate_exps / ffn_up_exps / +// ffn_down_exps). +// +// Rationale: with --pin-hot-experts, all MoE expert tensors are assumed to +// live in host (CPU) memory (only the dense/router parts are typically +// offloaded to VRAM). If the model was loaded via mmap (the default) and +// does not fit entirely in RAM, or --mlock was not used, the OS is free to +// evict a cold expert's pages and re-fault them in from disk the next time +// it is selected. This does NOT copy any data anywhere and does NOT change +// what ggml_mul_mat_id() reads: it locks the exact bytes that are already +// used by the compute graph, so the benefit is real and unconditional +// whenever those pages would otherwise be evictable. +// +// Pin/evict is fully ONLINE (no periodic "refresh" pass, no sorting): +// a single global ordered set tracks the hottest (layer, expert) pairs +// across ALL layers, keyed by usage count, with the coldest pinned expert +// at the front. The total pin budget is N * num_moe_layers. Every time an +// expert is selected: +// - if it's already pinned, its position in the ordered set is updated +// - else if there's still a free pin slot, it is pinned immediately +// - else if its count just overtook the coldest pinned expert, that +// expert is evicted (unlocked) and this one is pinned in its place +// This is O(log N) per observation and never needs to re-rank the whole +// model, so there is no "N most used experts of the last K decisions" +// window -- ranking is exact and always up to date. + +#include "ggml.h" +#include "llama-mmap.h" + +#include +#include +#include +#include +#include +#include +#include + +struct llama_model; + +class llama_hot_expert_cache { + public: + // n_pin_experts: number of hottest experts to keep mlock'd per layer (N in --pin-hot-experts N) + // total global capacity = N * num_moe_layers, ranked globally + // budget_bytes: hard cap on total bytes locked across ALL layers combined (0 = unlimited, NOT recommended) + // stats_interval: print_stats() is called automatically every `stats_interval` router + // observations (0 = disabled, only the destructor prints a final summary) + llama_hot_expert_cache(const llama_model & model, + int32_t n_pin_experts, + uint64_t budget_bytes, + uint64_t stats_interval = 200); + ~llama_hot_expert_cache(); + + llama_hot_expert_cache(const llama_hot_expert_cache &) = delete; + llama_hot_expert_cache & operator=(const llama_hot_expert_cache &) = delete; + + // ggml_backend_sched_eval_callback-compatible entry point. + // Pass `this` as user_data when installing. + static bool eval_callback(struct ggml_tensor * t, bool ask, void * user_data); + + // true if `expert_id` in layer `il` is currently mlock'd in place + bool is_pinned(int il, int32_t expert_id) const; + + // Prints a summary (bytes locked, per-layer breakdown, router observations) + // directly to stderr with fprintf. Deliberately bypasses LLAMA_LOG_* / the + // ggml log callback: at destruction time (process teardown, or a caller + // that already tore down its own log sink) those can silently swallow + // output, so this is a best-effort guaranteed-visible dump. + void print_stats() const; + + private: + // Unique key identifying a specific expert in a specific layer + struct expert_key { + int layer; + int32_t expert_id; + + bool operator==(const expert_key & o) const { return layer == o.layer && expert_id == o.expert_id; } + }; + + struct expert_key_hash { + std::size_t operator()(const expert_key & k) const { + return std::hash()(k.layer) ^ (std::hash()(k.expert_id) << 1); + } + }; + + struct mlock_deleter { + void operator()(llama_mlock * p) const { + if (p) { + p->unlock(); + delete p; + } + } + }; + + // holds the mlock guards keeping one expert's rows resident for each + // relevant weight tensor; unlock() is called via custom deleter on destruction + struct pinned_expert { + std::unique_ptr gate_lock; + std::unique_ptr up_lock; + std::unique_ptr down_lock; + std::unique_ptr gate_up_lock; + size_t nbytes_locked = 0; + }; + + struct layer_state { + const ggml_tensor * t_gate = nullptr; + const ggml_tensor * t_up = nullptr; + const ggml_tensor * t_down = nullptr; + const ggml_tensor * t_gate_up = nullptr; + + bool tensors_are_host = false; // false => experts live on a non-CPU backend, pinning is a no-op + bool resolved_tensors = false; + }; + + void on_topk_tensor(int il, const struct ggml_tensor * t); + void resolve_tensors(int il, layer_state & ls); + + // mlock() every dense (non-MoE-expert) tensor that lives in host memory + // (i.e. did not fit in VRAM) in place, BEFORE any hot expert is pinned, + // because dense parts are used on every token (hottest by definition). + // Consumes the global budget first, so hot experts only ever get the + // leftover budget. Stats-only (no-op) when mlock is unsupported. + void lock_dense_parts(); + + // called once per (layer, selected expert) observation; updates global counts and + // pins/evicts on the fly against the global top-N set + void observe_expert(int il, layer_state & ls, int32_t expert_id); + + // Returns true if at least some bytes were actually locked. + bool pin_expert(int il, layer_state & ls, int32_t expert_id); + void unpin_expert(int il, layer_state & ls, int32_t expert_id); + + // locks the byte range of `expert_id`'s row within tensor `w` in place, honoring + // the remaining global budget; returns bytes ACTUALLY locked (0 on failure/skip/no + // budget left). + size_t lock_expert_row(const struct ggml_tensor * w, + int32_t expert_id, + std::unique_ptr & out_lock); + + const llama_model & model; + + const int32_t n_pin; // N experts per layer + int32_t n_pin_total = 0; // N * num_moe_layers (global cap) + const uint64_t budget_bytes; // 0 = unlimited + const uint64_t stats_interval; // 0 = disabled periodic printing + + mutable std::mutex mu; + std::unordered_map layers; + + // Global tracking across all layers + std::unordered_map counts; // (layer, expert_id) -> times selected + + // Global pinned set: (count, layer, expert_id) ordered ascending by count + // begin() is always the coldest pinned expert globally + std::set> pinned_rank; + std::unordered_map pinned; + + uint64_t n_eval_calls = 0; + uint64_t n_bytes_locked = 0; // sum of llama_mlock::size(), i.e. bytes ACTUALLY locked, across all layers + + // mlock guards keeping the dense (non-MoE-expert) tensors resident. These are + // locked once in the constructor and released when the cache is destroyed. + std::vector> dense_locks; + uint64_t n_dense_bytes_locked = 0; // bytes ACTUALLY locked for dense tensors +}; diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index ed572da7fb54..60182998a00d 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -767,6 +767,12 @@ llama_mlock::~llama_mlock() = default; void llama_mlock::init(void * ptr) { pimpl->init(ptr); } void llama_mlock::grow_to(size_t target_size) { pimpl->grow_to(target_size); } +size_t llama_mlock::size() const { return pimpl->size; } +void llama_mlock::unlock() { + if (pimpl->addr != NULL && pimpl->size > 0) { + pimpl->raw_unlock(pimpl->addr, pimpl->size); + } +} #if defined(_POSIX_MEMLOCK_RANGE) || defined(_WIN32) const bool llama_mlock::SUPPORTED = true; diff --git a/src/llama-mmap.h b/src/llama-mmap.h index b7d5c61e95ff..e056bfee6a39 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -64,6 +64,14 @@ struct llama_mlock { void init(void * ptr); void grow_to(size_t target_size); + // bytes actually locked so far (may be less than the last grow_to() target + // if locking failed partway through -- see failed_already in the impl) + size_t size() const; + + // explicitly unlock the pages (NOT called by the destructor -- the default + // behavior is to leave pages locked until process exit, matching mlock(2)) + void unlock(); + static const bool SUPPORTED; private: diff --git a/src/llama.cpp b/src/llama.cpp index d6e0bbfefa72..d1a16442b7c2 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -58,16 +58,19 @@ const char * llama_load_mode_name(enum llama_load_mode load_mode) { return "mmap+mlock"; case LLAMA_LOAD_MODE_DIRECT_IO: return "dio"; + case LLAMA_LOAD_MODE_MMAP_PIN: + return "mmap+pin"; } GGML_ABORT("fatal error"); } enum llama_load_mode llama_load_mode_from_str(const char * str) { - if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; } - if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; } - if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; } - if (std::strcmp(str, "mmap+mlock") == 0) { return LLAMA_LOAD_MODE_MMAP_MLOCK; } - if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; } + if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; } + if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; } + if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; } + if (std::strcmp(str, "mmap+mlock") == 0) { return LLAMA_LOAD_MODE_MMAP_MLOCK; } + if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; } + if (std::strcmp(str, "mmap+pin") == 0) { return LLAMA_LOAD_MODE_MMAP_PIN; } throw std::invalid_argument(std::string("unknown load mode: ") + str); }