diff --git a/common/arg.cpp b/common/arg.cpp index 3dc9a9420..ccaa090b0 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2541,20 +2541,6 @@ common_params_context common_params_parser_init(common_params & params, llama_ex LOG_WRN("DEPRECATED: --defrag-thold is deprecated and no longer necessary to specify\n"); } ).set_env("LLAMA_ARG_DEFRAG_THOLD")); - add_opt(common_arg( - {"--moe-expert-cache"}, "N", - string_format("GPU cache slots per host-resident MoE expert layer, 0 = disabled (default: %d)", params.n_moe_cache_slots), - [](common_params & params, int value) { - params.n_moe_cache_slots = value; - } - ).set_env("LLAMA_ARG_MOE_EXPERT_CACHE")); - add_opt(common_arg( - {"--moe-expert-cache-inserts"}, "N", - string_format("max expert uploads per layer per decode step for the MoE expert cache (default: %d)", params.n_moe_cache_inserts), - [](common_params & params, int value) { - params.n_moe_cache_inserts = value; - } - ).set_env("LLAMA_ARG_MOE_EXPERT_CACHE_INSERTS")); if (ex == LLAMA_EXAMPLE_SERVER) { // this is to make sure this option appears in the server-specific section of the help message add_opt(common_arg( diff --git a/common/common.cpp b/common/common.cpp index afd224c2d..085dade98 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1719,8 +1719,6 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.n_outputs_max_per_seq = std::max(params.n_outputs_max_per_seq, 0); cparams.n_batch = params.n_batch; cparams.n_ubatch = params.n_ubatch; - cparams.n_moe_cache_slots = params.n_moe_cache_slots; - cparams.n_moe_cache_inserts = params.n_moe_cache_inserts; cparams.n_threads = params.cpuparams.n_threads; cparams.n_threads_batch = params.cpuparams_batch.n_threads == -1 ? params.cpuparams.n_threads : params.cpuparams_batch.n_threads; diff --git a/common/common.h b/common/common.h index 52e6ceab2..0a85f11f9 100644 --- a/common/common.h +++ b/common/common.h @@ -455,8 +455,6 @@ struct common_params { int32_t n_ctx = 0; // context size, 0 == context the model was trained with int32_t n_batch = 2048; // logical batch size for prompt processing (must be >=32 to use BLAS) int32_t n_ubatch = 512; // physical batch size for prompt processing (must be >=32 to use BLAS) - int32_t n_moe_cache_slots = 0; // GPU cache slots per host-resident MoE expert layer (0 = disabled) - int32_t n_moe_cache_inserts = 2; // max expert uploads per layer per decode step int32_t n_keep = 0; // number of tokens to keep from initial prompt int32_t n_chunks = -1; // max number of chunks to process (-1 = unlimited) int32_t n_parallel = 1; // number of parallel sequences to decode diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 2354abfc2..224bdef92 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -352,15 +352,6 @@ extern "C" { // Returns the old callback for chaining GGML_API ggml_abort_callback_t ggml_set_abort_callback(ggml_abort_callback_t callback); - struct ggml_tensor; - - // MoE expert-routing observation callback: invoked by the CPU mul_mat_id - // with the op's expert-id tensor (I32 [n_expert_used, n_tokens]). Used by - // the llama MoE expert cache to drive LRU placement decisions. - typedef void (*ggml_moe_obs_cb_t)(const char * tensor_name, const struct ggml_tensor * ids, void * ud); - GGML_API void ggml_set_moe_obs_callback(ggml_moe_obs_cb_t cb, void * ud); - GGML_API ggml_moe_obs_cb_t ggml_get_moe_obs_callback(void ** ud); - GGML_NORETURN GGML_ATTRIBUTE_FORMAT(3, 4) GGML_API void ggml_abort(const char * file, int line, const char * fmt, ...); diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 0ff4c93d8..7c274c03b 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -1647,17 +1647,6 @@ static void ggml_compute_forward_mul_mat_id( // initialize matrix_row_counts memset(matrix_row_counts, 0, n_as*sizeof(int64_t)); - // llama MoE expert cache: when src[3] is set it is an I32 table mapping - // expert id -> device cache slot, with op_params[0] holding the "not - // cached" dummy value. Cached ids are served by the device-side cache - // chain, so this op skips them and zeroes their dst rows instead. - const int32_t * moe_tbl = NULL; - int32_t moe_dummy = 0; - if (dst->src[3]) { - moe_tbl = (const int32_t *) dst->src[3]->data; - moe_dummy = ggml_get_op_params_i32(dst, 0); - } - // group rows by src0 matrix for (int64_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) { for (int id = 0; id < n_ids; ++id) { @@ -1665,51 +1654,10 @@ static void ggml_compute_forward_mul_mat_id( assert(i02 >= 0 && i02 < n_as); - if (moe_tbl && moe_tbl[i02] != moe_dummy) { - memset((char *) dst->data + id*nb1 + iid1*nb2, 0, ne0*sizeof(float)); - continue; - } - MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = (struct mmid_row_mapping) {id, iid1}; matrix_row_counts[i02] += 1; } } - - // MoE routing observation for the llama expert cache - { - void * moe_obs_ud = NULL; - ggml_moe_obs_cb_t moe_obs_cb = ggml_get_moe_obs_callback(&moe_obs_ud); - if (moe_obs_cb && strstr(src0->name, "ffn_gate_exps")) { - moe_obs_cb(src0->name, ids, moe_obs_ud); - } - } - - // GGML_MOE_LOG: append the routed expert ids of every ffn_gate_exps - // mul_mat_id to the file named by the env var. Diagnostic only; the - // whole block is inert unless GGML_MOE_LOG is set at first use. - { - static FILE * moe_log_file = NULL; - static int moe_log_state = -1; - if (moe_log_state == -1) { - const char * moe_log_path = getenv("GGML_MOE_LOG"); - if (moe_log_path && moe_log_path[0]) { - moe_log_file = fopen(moe_log_path, "a"); - } - moe_log_state = moe_log_file ? 1 : 0; - } - if (moe_log_state == 1 && strstr(src0->name, "ffn_gate_exps")) { - flockfile(moe_log_file); - for (int64_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) { - fprintf(moe_log_file, "%s", src0->name); - for (int id = 0; id < n_ids; ++id) { - const int32_t i02 = *(const int32_t *) ((const char *) ids->data + iid1*ids->nb[1] + id*ids->nb[0]); - fprintf(moe_log_file, " %d", i02); - } - fputc('\n', moe_log_file); - } - funlockfile(moe_log_file); - } - } } // reset current_chunk diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 97bc07b32..9bcabdbd3 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -250,21 +250,6 @@ GGML_API ggml_abort_callback_t ggml_set_abort_callback(ggml_abort_callback_t cal return ret_val; } -static ggml_moe_obs_cb_t g_moe_obs_cb = NULL; -static void * g_moe_obs_ud = NULL; - -void ggml_set_moe_obs_callback(ggml_moe_obs_cb_t cb, void * ud) { - g_moe_obs_cb = cb; - g_moe_obs_ud = ud; -} - -ggml_moe_obs_cb_t ggml_get_moe_obs_callback(void ** ud) { - if (ud) { - *ud = g_moe_obs_ud; - } - return g_moe_obs_cb; -} - void ggml_abort(const char * file, int line, const char * fmt, ...) { fflush(stdout); diff --git a/include/llama.h b/include/llama.h index 890d86b53..c1ca6debd 100644 --- a/include/llama.h +++ b/include/llama.h @@ -393,10 +393,6 @@ extern "C" { uint32_t yarn_orig_ctx; // YaRN original context size float defrag_thold; // [DEPRECATED] defragment the KV cache if holes/size > thold, <= 0 disabled (default) - // GPU-resident LRU cache for host-offloaded MoE expert weights [EXPERIMENTAL] - int32_t n_moe_cache_slots; // cache slots per host-resident expert layer (0 = disabled) - int32_t n_moe_cache_inserts; // max expert uploads per layer per decode step - ggml_backend_sched_eval_callback cb_eval; void * cb_eval_user_data; diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 6a94f74be..97250c2e8 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -37,7 +37,6 @@ set(LLAMA_CORE_SOURCES llama-model-loader.cpp llama-model-saver.cpp llama-model.cpp - llama-moecache.cpp llama-quant.cpp llama-sampler.cpp llama-vocab.cpp diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 35ccb62e6..af37c6b53 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1,7 +1,5 @@ #include "llama-context.h" -#include "llama-moecache.h" - #include "ggml.h" #include "llama-arch.h" #include "llama-graph.h" @@ -97,8 +95,6 @@ llama_context::llama_context( // may need to be backend-dependent LLAMA_LOG_INFO("%s: constructing llama_context\n", __func__); - llama_moe_cache_init(model, params.n_moe_cache_slots, params.n_moe_cache_inserts); - t_start_us = model.t_start_us; t_load_us = model.t_load_us; @@ -2215,9 +2211,6 @@ int llama_context::decode(const llama_batch_ext & batch_inp) { // wait for the computation to finish (automatically done when obtaining the model output) //synchronize(); - // apply throttled MoE expert-cache updates between graph executions - llama_moe_cache_step(); - return 0; } @@ -3981,8 +3974,6 @@ llama_context_params llama_context_default_params() { /*.yarn_beta_slow =*/ -1.0f, /*.yarn_orig_ctx =*/ 0, /*.defrag_thold =*/ -1.0f, - /*.n_moe_cache_slots =*/ 0, - /*.n_moe_cache_inserts =*/ 2, /*.cb_eval =*/ nullptr, /*.cb_eval_user_data =*/ nullptr, /*.type_k =*/ GGML_TYPE_F16, diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8da956215..65d6dc7fb 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1,7 +1,5 @@ #include "llama-graph.h" -#include "llama-moecache.h" - #include "llama-impl.h" #include "llama-model.h" #include "llama-moe-cache.h" @@ -2218,9 +2216,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn( // the experts of host-resident layers may be read from the MoE cache ggml_tensor * slots = build_moe_cache_slots(selected_experts, up_exps, gate_exps, down_exps, gate_up_exps, il); - cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens); - ggml_tensor * mc_inp = cur; if (weight_before_ffn) { // repeat cur to [n_embd, n_expert_used, n_tokens] @@ -2256,11 +2252,6 @@ ggml_tensor * llm_graph_context::build_moe_ffn( up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s, slots); // [n_ff, n_expert_used, n_tokens] cb(up, "ffn_moe_up", il); - if (mcache) { - up->src[3] = mcache->host_table; - up->op_params[0] = mcache->n_slots; - } - if (up_exps_s) { cb(up, "ffn_moe_up_scaled", il); } @@ -2273,11 +2264,6 @@ ggml_tensor * llm_graph_context::build_moe_ffn( if (gate_exps) { cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s, slots); // [n_ff, n_expert_used, n_tokens] cb(cur, "ffn_moe_gate", il); - - if (mcache) { - cur->src[3] = mcache->host_table; - cur->op_params[0] = mcache->n_slots; - } } else { cur = up; } @@ -2384,43 +2370,6 @@ ggml_tensor * llm_graph_context::build_moe_ffn( } cb(experts, "ffn_moe_down", il); - if (mcache) { - experts->src[3] = mcache->host_table; - experts->op_params[0] = mcache->n_slots; - - // device-side chain over the cached experts, mirroring the LLM_FFN_SILU - // activation above (the only type_op the cache path is enabled for) - ggml_tensor * up_g = ggml_mul_mat_id(ctx0, mcache->up_c, mc_inp, mc_slot_ids); - ggml_tensor * gate_g = ggml_mul_mat_id(ctx0, mcache->gate_c, mc_inp, mc_slot_ids); - cb(up_g, "ffn_moe_cache_up", il); - cb(gate_g, "ffn_moe_cache_gate", il); - - ggml_tensor * act_g = nullptr; - { - const float limit = il >= 0 ? hparams.swiglu_clamp_exp[il] : 0.0f; - constexpr float eps = 1e-6f; - if (limit > eps) { - up_g = ggml_clamp(ctx0, up_g, -limit, limit); - if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { - gate_g = ggml_clamp(ctx0, gate_g, -INFINITY, limit); - act_g = ggml_swiglu_split(ctx0, gate_g, up_g); - } else { - ggml_tensor * ga = ggml_silu(ctx0, gate_g); - ga = ggml_clamp(ctx0, ga, -INFINITY, limit); - act_g = ggml_mul(ctx0, ga, up_g); - } - } else { - act_g = ggml_swiglu_split(ctx0, gate_g, up_g); - } - } - - ggml_tensor * down_g = ggml_mul_mat_id(ctx0, mcache->down_c, act_g, mc_slot_ids); - cb(down_g, "ffn_moe_cache_down", il); - - experts = ggml_add(ctx0, experts, down_g); - cb(experts, "ffn_moe_cache_merged", il); - } - if (down_exps_s) { cb(experts, "ffn_moe_down_scaled", il); } diff --git a/src/llama-moecache.cpp b/src/llama-moecache.cpp deleted file mode 100644 index 4e0732e80..000000000 --- a/src/llama-moecache.cpp +++ /dev/null @@ -1,407 +0,0 @@ -#include "llama-moecache.h" - -#include "llama-impl.h" -#include "llama-model.h" - -#include "ggml.h" -#include "ggml-backend.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace { - -struct layer_state { - llama_moe_cache_layer pub; - - // LRU bookkeeping (host side; the tables mirror expert_slot) - std::vector slot_expert; // slot -> expert id, -1 when empty - std::vector expert_slot; // expert id -> slot, -1 when uncached - std::vector slot_last_use; // slot -> lamport clock of last hit - std::vector pending; // uncached ids observed since last step (dedup, obs order) - - std::vector slot_in_flight; // slot has an upload pending - - uint64_t n_hit = 0; - uint64_t n_miss = 0; -}; - -struct upload_job { - size_t layer_idx; - int32_t expert; - int32_t slot; - bool done = false; -}; - -struct moe_cache { - int32_t n_slots = 0; - int32_t max_inserts = 2; - - uint64_t clock = 0; - uint64_t n_steps = 0; - - std::mutex mtx; // guards pending lists + clock (observe runs during graph exec) - - std::vector layers; - std::map by_up_src; - - std::vector ctxs; - std::vector bufs; - - // async upload worker: slices are copied to the device off the decode - // thread; the new table mapping is only published at a later step() once - // the upload has completed, so a running graph never reads a torn slot - std::thread worker; - std::mutex wmtx; - std::condition_variable wcv; - std::deque todo; - std::vector done; - bool stop = false; -}; - -moe_cache * g_cache = nullptr; -std::mutex g_init_mtx; -bool g_init_done = false; - -int parse_layer_from_name(const char * name) { - // "blk..ffn_gate_exps.weight" - if (strncmp(name, "blk.", 4) != 0) { - return -1; - } - return atoi(name + 4); -} - -void moe_obs_cb(const char * name, const struct ggml_tensor * ids, void * ud) { - moe_cache * mc = (moe_cache *) ud; - - const int64_t n_ids = ids->ne[0]; - const int64_t n_tokens = ids->ne[1]; - if (n_tokens > 4) { - return; // batch/prefill: the cache graph is not built there, don't pollute the LRU - } - - const int il = parse_layer_from_name(name); - if (il < 0) { - return; - } - - layer_state * ls = nullptr; - for (auto & l : mc->layers) { - if (l.pub.il == il) { ls = &l; break; } - } - if (!ls) { - return; - } - - std::lock_guard lock(mc->mtx); - for (int64_t t = 0; t < n_tokens; ++t) { - for (int64_t i = 0; i < n_ids; ++i) { - const int32_t id = *(const int32_t *) ((const char *) ids->data + t*ids->nb[1] + i*ids->nb[0]); - if (id < 0 || id >= (int32_t) ls->expert_slot.size()) { - continue; - } - const int32_t slot = ls->expert_slot[id]; - if (slot >= 0) { - ls->n_hit++; - ls->slot_last_use[slot] = ++mc->clock; - } else { - ls->n_miss++; - bool dup = false; - for (int32_t p : ls->pending) { - if (p == id) { dup = true; break; } - } - if (!dup) { - ls->pending.push_back(id); - } - } - } - } -} - -void upload_slice(ggml_tensor * dst_c, const ggml_tensor * src, int32_t expert, int32_t slot) { - const size_t sz = src->nb[2]; - if ((size_t) slot*dst_c->nb[2] + sz > ggml_nbytes(dst_c) || (size_t) expert*sz + sz > ggml_nbytes(src)) { - LLAMA_LOG_ERROR("moe-cache: bad upload %s <- %s expert=%d slot=%d sz=%zu dst_nb2=%zu dst_bytes=%zu src_bytes=%zu\n", - dst_c->name, src->name, expert, slot, sz, dst_c->nb[2], ggml_nbytes(dst_c), ggml_nbytes(src)); - return; - } - ggml_backend_tensor_set(dst_c, (const char *) src->data + (size_t) expert*sz, (size_t) slot*dst_c->nb[2], sz); -} - -void set_table_entry(llama_moe_cache_layer & pub, int32_t expert, int32_t slot_or_dummy) { - const int32_t v = slot_or_dummy; - ggml_backend_tensor_set(pub.dev_table, &v, (size_t) expert*sizeof(int32_t), sizeof(int32_t)); - ggml_backend_tensor_set(pub.host_table, &v, (size_t) expert*sizeof(int32_t), sizeof(int32_t)); -} - -} // namespace - -void llama_moe_cache_init(const llama_model & model, int32_t n_slots, int32_t max_inserts) { - std::lock_guard init_lock(g_init_mtx); - if (g_init_done) { - return; - } - [&]() { - if (n_slots <= 0) { - g_init_done = true; - return; - } - - auto * mc = new moe_cache(); - mc->n_slots = n_slots; - if (max_inserts > 0) { - mc->max_inserts = max_inserts; - } - - // collect the host-resident expert layers, grouped by the device buffer - // type of that layer's router (the cache lives next to the router) - struct cand { int il; const llama_layer * l; }; - std::map> groups; - - for (size_t il = 0; il < model.layers.size(); ++il) { - const auto & l = model.layers[il]; - if (!l.ffn_up_exps || !l.ffn_gate_exps || !l.ffn_down_exps || !l.ffn_gate_inp) { - continue; - } - if (!l.ffn_up_exps->data || !l.ffn_gate_exps->data || !l.ffn_down_exps->data) { - continue; // dry-run / memory-estimation model: weights not loaded, don't bind to it - } - if (!l.ffn_up_exps->buffer || !ggml_backend_buffer_is_host(l.ffn_up_exps->buffer)) { - continue; // experts already on a device: nothing to cache - } - if (!l.ffn_gate_inp->buffer || ggml_backend_buffer_is_host(l.ffn_gate_inp->buffer)) { - continue; // no device home for the cache - } - groups[ggml_backend_buffer_get_type(l.ffn_gate_inp->buffer)].push_back({(int) il, &l}); - } - - if (groups.empty()) { - LLAMA_LOG_INFO("%s: LLAMA_MOE_CACHE_SLOTS=%d but no host-resident expert layers found - disabled\n", __func__, n_slots); - delete mc; - return; - } - - // host buffer for the CPU-side tables - std::vector all; - for (auto & g : groups) { - all.insert(all.end(), g.second.begin(), g.second.end()); - } - - auto alloc_group = [&](ggml_backend_buffer_type_t buft, const std::vector & cands, bool tables_only) -> bool { - ggml_init_params ip = { - /*.mem_size =*/ ggml_tensor_overhead()*(cands.size()*4 + 8), - /*.mem_buffer=*/ nullptr, - /*.no_alloc =*/ true, - }; - ggml_context * ctx = ggml_init(ip); - if (!ctx) { - return false; - } - mc->ctxs.push_back(ctx); - - for (const auto & c : cands) { - layer_state * ls = nullptr; - for (auto & l : mc->layers) { - if (l.pub.il == c.il) { ls = &l; break; } - } - if (!ls) { - mc->layers.push_back({}); - ls = &mc->layers.back(); - ls->pub.il = c.il; - ls->pub.n_slots = n_slots; - ls->pub.up_src = c.l->ffn_up_exps; - ls->pub.gate_src = c.l->ffn_gate_exps; - ls->pub.down_src = c.l->ffn_down_exps; - } - - if (tables_only) { - ls->pub.host_table = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1, ls->pub.up_src->ne[2]); - ggml_format_name(ls->pub.host_table, "moe_cache_htbl.%d", c.il); - } else { - const ggml_tensor * u = c.l->ffn_up_exps; - const ggml_tensor * g = c.l->ffn_gate_exps; - const ggml_tensor * d = c.l->ffn_down_exps; - ls->pub.up_c = ggml_new_tensor_3d(ctx, u->type, u->ne[0], u->ne[1], n_slots + 1); - ls->pub.gate_c = ggml_new_tensor_3d(ctx, g->type, g->ne[0], g->ne[1], n_slots + 1); - ls->pub.down_c = ggml_new_tensor_3d(ctx, d->type, d->ne[0], d->ne[1], n_slots + 1); - ls->pub.dev_table = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1, u->ne[2]); - ggml_format_name(ls->pub.up_c, "moe_cache_up.%d", c.il); - ggml_format_name(ls->pub.gate_c, "moe_cache_gate.%d", c.il); - ggml_format_name(ls->pub.down_c, "moe_cache_down.%d", c.il); - ggml_format_name(ls->pub.dev_table, "moe_cache_tbl.%d", c.il); - } - } - - ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft); - if (!buf) { - LLAMA_LOG_WARN("%s: failed to allocate MoE cache buffer on %s - cache disabled\n", - __func__, ggml_backend_buft_name(buft)); - return false; - } - ggml_backend_buffer_clear(buf, 0); - mc->bufs.push_back(buf); - return true; - }; - - bool ok = alloc_group(ggml_backend_cpu_buffer_type(), all, /*tables_only=*/true); - for (auto & g : groups) { - if (!ok) { - break; - } - ok = alloc_group(g.first, g.second, /*tables_only=*/false); - } - - if (!ok) { - for (auto * b : mc->bufs) { ggml_backend_buffer_free(b); } - for (auto * c : mc->ctxs) { ggml_free(c); } - delete mc; - g_init_done = true; // a real model was seen and allocation failed: stay disabled - return; - } - - // init LRU state + tables (everything uncached -> dummy slot n_slots) - size_t vram = 0; - for (auto & ls : mc->layers) { - const int64_t n_expert = ls.pub.up_src->ne[2]; - ls.slot_expert.assign(n_slots, -1); - ls.expert_slot.assign(n_expert, -1); - ls.slot_last_use.assign(n_slots, 0); - ls.slot_in_flight.assign(n_slots, false); - - std::vector dummy(n_expert, n_slots); - ggml_backend_tensor_set(ls.pub.dev_table, dummy.data(), 0, n_expert*sizeof(int32_t)); - ggml_backend_tensor_set(ls.pub.host_table, dummy.data(), 0, n_expert*sizeof(int32_t)); - - mc->by_up_src[ls.pub.up_src] = &ls - mc->layers.data(); - vram += ggml_nbytes(ls.pub.up_c) + ggml_nbytes(ls.pub.gate_c) + ggml_nbytes(ls.pub.down_c); - LLAMA_LOG_DEBUG("moe-cache: init layer %d '%s' %zu bytes/expert\n", - ls.pub.il, ls.pub.up_src->name, ls.pub.up_src->nb[2]); - } - - mc->worker = std::thread([mc]() { - for (;;) { - upload_job j; - { - std::unique_lock lk(mc->wmtx); - mc->wcv.wait(lk, [mc]() { return mc->stop || !mc->todo.empty(); }); - if (mc->stop) { - return; - } - j = mc->todo.front(); - mc->todo.pop_front(); - } - auto & ls = mc->layers[j.layer_idx]; - upload_slice(ls.pub.up_c, ls.pub.up_src, j.expert, j.slot); - upload_slice(ls.pub.gate_c, ls.pub.gate_src, j.expert, j.slot); - upload_slice(ls.pub.down_c, ls.pub.down_src, j.expert, j.slot); - { - std::lock_guard lk(mc->wmtx); - j.done = true; - mc->done.push_back(j); - } - } - }); - - ggml_set_moe_obs_callback(moe_obs_cb, mc); - g_cache = mc; - g_init_done = true; - - LLAMA_LOG_INFO("%s: MoE expert cache enabled: %zu layers x %d slots, %d inserts/step, %.1f MiB device memory\n", - __func__, mc->layers.size(), n_slots, mc->max_inserts, vram/1024.0/1024.0); - }(); -} - -const llama_moe_cache_layer * llama_moe_cache_lookup(const ggml_tensor * up_exps) { - if (!g_cache) { - return nullptr; - } - auto it = g_cache->by_up_src.find(up_exps); - if (it == g_cache->by_up_src.end()) { - return nullptr; - } - return &g_cache->layers[it->second].pub; -} - -void llama_moe_cache_step() { - moe_cache * mc = g_cache; - if (!mc) { - return; - } - - // 1) publish completed uploads (sync point: no graph is executing) - { - std::lock_guard wlk(mc->wmtx); - std::lock_guard lk(mc->mtx); - for (const auto & j : mc->done) { - auto & ls = mc->layers[j.layer_idx]; - ls.slot_expert[j.slot] = j.expert; - ls.expert_slot[j.expert] = j.slot; - ls.slot_last_use[j.slot] = ++mc->clock; - ls.slot_in_flight[j.slot] = false; - set_table_entry(ls.pub, j.expert, j.slot); - } - mc->done.clear(); - } - - std::lock_guard lock(mc->mtx); - mc->n_steps++; - - // 2) schedule new uploads: evict at a sync point (clear the victim's table - // entry now), then hand the slice copies to the worker - for (size_t li = 0; li < mc->layers.size(); ++li) { - auto & ls = mc->layers[li]; - if (ls.pending.empty()) { - continue; - } - - int budget = mc->max_inserts; - for (auto it = ls.pending.rbegin(); it != ls.pending.rend() && budget > 0; ++it, --budget) { - const int32_t id = *it; - if (ls.expert_slot[id] >= 0) { - continue; - } - - // victim: an empty non-in-flight slot if any, else the LRU non-in-flight slot - int32_t slot = -1; - uint64_t best = UINT64_MAX; - for (int32_t s = 0; s < mc->n_slots; ++s) { - if (ls.slot_in_flight[s]) { - continue; - } - if (ls.slot_expert[s] < 0) { slot = s; break; } - if (ls.slot_last_use[s] < best) { best = ls.slot_last_use[s]; slot = s; } - } - if (slot < 0) { - break; // every slot is in flight; try again next step - } - - const int32_t victim = ls.slot_expert[slot]; - if (victim >= 0) { - ls.expert_slot[victim] = -1; - ls.slot_expert[slot] = -1; - set_table_entry(ls.pub, victim, mc->n_slots); - } - ls.slot_in_flight[slot] = true; - - std::lock_guard wlk(mc->wmtx); - mc->todo.push_back({li, id, slot}); - } - ls.pending.clear(); - } - mc->wcv.notify_one(); - - if (mc->n_steps % 512 == 0) { - uint64_t h = 0, m = 0; - for (auto & ls : mc->layers) { h += ls.n_hit; m += ls.n_miss; } - LLAMA_LOG_DEBUG("moe-cache: steps=%" PRIu64 " hits=%" PRIu64 " misses=%" PRIu64 " hit-rate=%.1f%%\n", - mc->n_steps, h, m, h + m ? 100.0*h/(h + m) : 0.0); - } -} diff --git a/src/llama-moecache.h b/src/llama-moecache.h deleted file mode 100644 index ed739f90c..000000000 --- a/src/llama-moecache.h +++ /dev/null @@ -1,61 +0,0 @@ -#pragma once - -// GPU-resident LRU cache for MoE expert weights that -ot pinned to host memory. -// -// Motivation (measured on Qwen3.8-Flash-Next, 512 experts / 10 routed): expert -// routing has strong temporal locality (LRU-64 hit rate ~67% over a mixed -// workload) even though the long-run distribution is near-uniform. Decode on a -// host-offloaded MoE layer is bound by host RAM bandwidth, so serving the hot -// experts from VRAM removes most of the per-token DIMM traffic. -// -// Mechanism (no custom kernels): -// - per cached layer, companion tensors up_c/gate_c/down_c of shape -// [ne0, ne1, n_slots+1] live in the device buffer of that layer's router; -// slot n_slots is permanently zero (the "dummy" slot). -// - an I32 table[512] maps expert id -> slot, or n_slots when uncached. -// One copy on device (read by get_rows to remap ids for the cache-side -// mul_mat_id chain) and one on host (read by the CPU mul_mat_id via -// src[3] to SKIP cached ids, zeroing their dst rows). -// - the two down-projection outputs are summed; uncached ids contribute 0 -// through the cache chain (zero slot) and cached ids contribute 0 through -// the CPU chain (skip), so the result is exact. -// - llama_moe_cache_step(), called at the end of llama_context::decode(), -// performs throttled LRU updates: at most LLAMA_MOE_CACHE_INSERTS expert -// uploads per layer per step via ggml_backend_tensor_set. -// -// Enabled via llama_context_params.n_moe_cache_slots (CLI: --moe-expert-cache). - -#include - -struct llama_model; -struct ggml_tensor; - -struct llama_moe_cache_layer { - int il = -1; - - int32_t n_slots = 0; - - // host-resident source weights (the authoritative experts) - ggml_tensor * up_src = nullptr; - ggml_tensor * gate_src = nullptr; - ggml_tensor * down_src = nullptr; - - // device-resident cache slots, ne[2] == n_slots + 1 (last slot all zeros) - ggml_tensor * up_c = nullptr; - ggml_tensor * gate_c = nullptr; - ggml_tensor * down_c = nullptr; - - // expert id -> slot (or n_slots when uncached); I32 [1, n_expert] - ggml_tensor * dev_table = nullptr; - ggml_tensor * host_table = nullptr; -}; - -// build the cache for every host-resident expert layer of the model. -// Safe to call more than once; only the first call does work. -void llama_moe_cache_init(const llama_model & model, int32_t n_slots, int32_t max_inserts); - -// nullptr when the cache is disabled or this tensor has no cached layer -const llama_moe_cache_layer * llama_moe_cache_lookup(const ggml_tensor * up_exps); - -// apply throttled LRU updates; call between graph executions only -void llama_moe_cache_step();