--- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -86,6 +86,8 @@ llguidance.cpp log.cpp log.h + wee-expert-prefetch.cpp + wee-expert-prefetch.h ngram-cache.cpp ngram-cache.h ngram-map.cpp --- a/common/common.cpp +++ b/common/common.cpp @@ -5,6 +5,7 @@ #include "common.h" #include "fit.h" #include "log.h" +#include "wee-expert-prefetch.h" #include "llama.h" #include "sampling.h" #include "speculative.h" @@ -1434,6 +1435,11 @@ } common_init_result_ptr common_init_from_params(common_params & params, bool model_only) { + if (const char * map = getenv("WEE_EXPERT_MAP"); map && !model_only) { + if (params.cb_eval || params.n_gpu_layers != 0) throw std::runtime_error("Expert prefetch requires CPU-only execution without another eval callback"); + wee_expert_prefetch_init(map, params.model.path); + params.cb_eval = wee_expert_prefetch_eval; + } common_init_result_ptr res(new common_init_result(params, model_only)); llama_model * model = res->model(); --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -1628,7 +1628,12 @@ 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]); - assert(i02 >= 0 && i02 < n_as); + GGML_ASSERT(i02 >= -1 && i02 < n_as); + if (i02 == -1) { + // Inactive mixed-tier slots must not read any expert matrix. + 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; --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -441,6 +441,15 @@ { LLM_TENSOR_FFN_GATE_EXP, "blk.%d.ffn_gate.%d" }, { LLM_TENSOR_FFN_DOWN_EXP, "blk.%d.ffn_down.%d" }, { LLM_TENSOR_FFN_UP_EXP, "blk.%d.ffn_up.%d" }, + { LLM_TENSOR_FFN_GATE_EXPS_TIER1, "blk.%d.ffn_gate_exps_tier1" }, + { LLM_TENSOR_FFN_UP_EXPS_TIER1, "blk.%d.ffn_up_exps_tier1" }, + { LLM_TENSOR_FFN_DOWN_EXPS_TIER1, "blk.%d.ffn_down_exps_tier1" }, + { LLM_TENSOR_FFN_GATE_EXPS_TIER2, "blk.%d.ffn_gate_exps_tier2" }, + { LLM_TENSOR_FFN_UP_EXPS_TIER2, "blk.%d.ffn_up_exps_tier2" }, + { LLM_TENSOR_FFN_DOWN_EXPS_TIER2, "blk.%d.ffn_down_exps_tier2" }, + { LLM_TENSOR_FFN_GATE_EXPS_TIER3, "blk.%d.ffn_gate_exps_tier3" }, + { LLM_TENSOR_FFN_UP_EXPS_TIER3, "blk.%d.ffn_up_exps_tier3" }, + { LLM_TENSOR_FFN_DOWN_EXPS_TIER3, "blk.%d.ffn_down_exps_tier3" }, { LLM_TENSOR_FFN_GATE_EXPS, "blk.%d.ffn_gate_exps" }, { LLM_TENSOR_FFN_GATE_UP_EXPS, "blk.%d.ffn_gate_up_exps" }, { LLM_TENSOR_FFN_DOWN_EXPS, "blk.%d.ffn_down_exps" }, @@ -889,6 +898,15 @@ {LLM_TENSOR_DEC_ATTN_REL_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, {LLM_TENSOR_ENC_ATTN_REL_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, {LLM_TENSOR_FFN_DOWN_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_GATE_EXPS_TIER1, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_UP_EXPS_TIER1, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_DOWN_EXPS_TIER1, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_GATE_EXPS_TIER2, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_UP_EXPS_TIER2, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_DOWN_EXPS_TIER2, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_GATE_EXPS_TIER3, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_UP_EXPS_TIER3, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_DOWN_EXPS_TIER3, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_FFN_GATE_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_FFN_UP_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_FFN_GATE_UP_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -469,6 +469,15 @@ LLM_TENSOR_FFN_DOWN_EXPS, // merged experts LLM_TENSOR_FFN_GATE_EXPS, LLM_TENSOR_FFN_UP_EXPS, + LLM_TENSOR_FFN_GATE_EXPS_TIER1, + LLM_TENSOR_FFN_UP_EXPS_TIER1, + LLM_TENSOR_FFN_DOWN_EXPS_TIER1, + LLM_TENSOR_FFN_GATE_EXPS_TIER2, + LLM_TENSOR_FFN_UP_EXPS_TIER2, + LLM_TENSOR_FFN_DOWN_EXPS_TIER2, + LLM_TENSOR_FFN_GATE_EXPS_TIER3, + LLM_TENSOR_FFN_UP_EXPS_TIER3, + LLM_TENSOR_FFN_DOWN_EXPS_TIER3, LLM_TENSOR_FFN_GATE_UP_EXPS, LLM_TENSOR_FFN_DOWN_SHEXP, LLM_TENSOR_FFN_GATE_SHEXP, --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -1364,6 +1365,17 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps) { if (use_mmap) { + if (prefetch && std::getenv("WEE_LAZY_EXPERT_LOAD") && std::getenv("WEE_EXPERT_MAP")) { + size_t bytes = 0; + for (const auto & item : weights_map) { + if (item.first.find(".ffn_") == std::string::npos || item.first.find("_exps") == std::string::npos) continue; + const auto & w = item.second; + const size_t n = ggml_nbytes(w.tensor); + lazy_tensor_ranges[w.idx].emplace_back(w.offs, w.offs+n); + bytes += n; + } + LLAMA_LOG_INFO("%s: exact expert prefetch replaces startup preload of %.2f GiB\n", __func__, bytes/double(GiB)); + } mappings.reserve(files.size()); mmaps_used.reserve(files.size()); for (uint32_t idx = 0; idx < files.size(); idx++) { --- a/src/llama-model.h +++ b/src/llama-model.h @@ -329,6 +329,15 @@ struct ggml_tensor * ffn_gate_exps = nullptr; struct ggml_tensor * ffn_down_exps = nullptr; struct ggml_tensor * ffn_up_exps = nullptr; + struct ggml_tensor * ffn_gate_exps_tier1 = nullptr; + struct ggml_tensor * ffn_up_exps_tier1 = nullptr; + struct ggml_tensor * ffn_down_exps_tier1 = nullptr; + struct ggml_tensor * ffn_gate_exps_tier2 = nullptr; + struct ggml_tensor * ffn_up_exps_tier2 = nullptr; + struct ggml_tensor * ffn_down_exps_tier2 = nullptr; + struct ggml_tensor * ffn_gate_exps_tier3 = nullptr; + struct ggml_tensor * ffn_up_exps_tier3 = nullptr; + struct ggml_tensor * ffn_down_exps_tier3 = nullptr; struct ggml_tensor * ffn_gate_up_exps = nullptr; struct ggml_tensor * ffn_gate_inp_b = nullptr; struct ggml_tensor * ffn_gate_exps_b = nullptr; --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -198,9 +198,45 @@ layer.ple_conv1d = create_tensor(tn(LLM_TENSOR_PLE_CONV1D, "weight", il), { hparams.ple_conv_kernel, hc_dim }, 0); } - layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0); - layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0); - create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0); + const auto * router_meta = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il).str().c_str()); + if (!router_meta) throw std::runtime_error("missing expert router tensor"); + const auto * t0 = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il).str().c_str()); + const auto * t1 = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER1, "weight", il).str().c_str()); + const auto * t2 = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER2, "weight", il).str().c_str()); + const auto * t3 = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER3, "weight", il).str().c_str()); + const bool mixed = t1 || t2 || t3 || router_meta->ne[1] != n_expert; + const int64_t layer_experts = mixed ? router_meta->ne[1] : n_expert; + if (layer_experts < n_expert_used || layer_experts > n_expert) throw std::runtime_error("invalid layer expert count"); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, layer_experts }, 0); + if (mixed) { + if (!t0 || n_ff_exp != 640) throw std::runtime_error("missing mixed-tier base tensors"); + const int64_t n0 = t0->ne[2], n1 = t1 ? t1->ne[2] : 0, n2 = t2 ? t2->ne[2] : 0, n3 = t3 ? t3->ne[2] : 0; + if (n0 < 1 || n0+n1+n2+n3 != layer_experts) throw std::runtime_error("invalid mixed-tier bank sizes"); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n0 }, 0); + create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n0, 0); + if (t1) { + const auto * d1 = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_DOWN_EXPS_TIER1, "weight", il).str().c_str()); + if (!d1 || (d1->ne[0] != 640 && d1->ne[0] != 768)) throw std::runtime_error("invalid tier 1 down width"); + layer.ffn_gate_exps_tier1 = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER1, "weight", il), { n_embd, n_ff_exp, n1 }, 0); + layer.ffn_up_exps_tier1 = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS_TIER1, "weight", il), { n_embd, n_ff_exp, n1 }, 0); + layer.ffn_down_exps_tier1 = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS_TIER1, "weight", il), { d1->ne[0], n_embd, n1 }, 0); + } + if (t2) { + const auto * d2 = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_DOWN_EXPS_TIER2, "weight", il).str().c_str()); + if (!d2 || (d2->ne[0] != 640 && d2->ne[0] != 768)) throw std::runtime_error("invalid tier 2 down width"); + layer.ffn_gate_exps_tier2 = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER2, "weight", il), { n_embd, n_ff_exp, n2 }, 0); + layer.ffn_up_exps_tier2 = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS_TIER2, "weight", il), { n_embd, n_ff_exp, n2 }, 0); + layer.ffn_down_exps_tier2 = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS_TIER2, "weight", il), { d2->ne[0], n_embd, n2 }, 0); + } + if (t3) { + layer.ffn_gate_exps_tier3 = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER3, "weight", il), { n_embd, n_ff_exp, n3 }, 0); + layer.ffn_up_exps_tier3 = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS_TIER3, "weight", il), { n_embd, n_ff_exp, n3 }, 0); + layer.ffn_down_exps_tier3 = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS_TIER3, "weight", il), { n_ff_exp, n_embd, n3 }, 0); + } + } else { + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0); + create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0); + } layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, 0); layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0); @@ -904,8 +940,57 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_ffn(ggml_tensor * cur, const int il) { GGML_ASSERT(model.layers[il].ffn_gate_inp != nullptr); - ggml_tensor * moe_out = - build_moe_ffn(cur, + ggml_tensor * moe_out = nullptr; + if (model.layers[il].ffn_gate_exps_tier1 || model.layers[il].ffn_gate_exps_tier2 || model.layers[il].ffn_gate_exps_tier3 || model.layers[il].ffn_gate_inp->ne[1] != n_expert) { + const auto & layer = model.layers[il]; + const int64_t n_tokens = cur->ne[1]; + ggml_tensor * logits = build_lora_mm(layer.ffn_gate_inp, cur); + ggml_tensor * probs = ggml_soft_max(ctx0, logits); + cb(probs, "ffn_moe_probs", il); + ggml_tensor * ids = ggml_argsort_top_k(ctx0, probs, n_expert_used); + cb(ids, "ffn_moe_topk", il); + ggml_tensor * selected = ggml_get_rows(ctx0, ggml_reshape_3d(ctx0, probs, 1, layer.ffn_gate_inp->ne[1], n_tokens), ids); + selected = ggml_reshape_2d(ctx0, selected, n_expert_used, n_tokens); + ggml_tensor * denom = ggml_clamp(ctx0, ggml_sum_rows(ctx0, selected), 6.103515625e-5f, INFINITY); + selected = ggml_div(ctx0, selected, denom); + if (hparams.expert_weights_scale != 0.0f && hparams.expert_weights_scale != 1.0f) { + selected = ggml_scale(ctx0, selected, hparams.expert_weights_scale); + } + ggml_tensor * ids_f = ggml_cast(ctx0, ids, GGML_TYPE_F32); + ggml_tensor * inp = ggml_reshape_3d(ctx0, cur, cur->ne[0], 1, n_tokens); + ggml_tensor * weighted = nullptr; + ggml_tensor * ups[] = {layer.ffn_up_exps, layer.ffn_up_exps_tier1, layer.ffn_up_exps_tier2, layer.ffn_up_exps_tier3}; + ggml_tensor * gates[] = {layer.ffn_gate_exps, layer.ffn_gate_exps_tier1, layer.ffn_gate_exps_tier2, layer.ffn_gate_exps_tier3}; + ggml_tensor * downs[] = {layer.ffn_down_exps, layer.ffn_down_exps_tier1, layer.ffn_down_exps_tier2, layer.ffn_down_exps_tier3}; + int64_t tier_start = 0; + for (int tier = 0; tier < 4; ++tier) { + if (!ups[tier]) continue; + const int64_t tier_size = ups[tier]->ne[2]; + const float base = float(tier_start); + ggml_tensor * local_f = ggml_scale_bias(ctx0, ids_f, 1.0f, -base); + ggml_tensor * low = ggml_step(ctx0, ggml_scale_bias(ctx0, local_f, 1.0f, 0.5f)); + ggml_tensor * high = ggml_step(ctx0, ggml_scale_bias(ctx0, local_f, -1.0f, float(tier_size) - 0.5f)); + tier_start += tier_size; + ggml_tensor * mask = ggml_mul(ctx0, low, high); + // The CPU mul_mat_id backend treats -1 as an inactive expert slot. + ggml_tensor * local_f_clamped = ggml_clamp(ctx0, local_f, 0.0f, float(tier_size - 1)); + ggml_tensor * local = ggml_cast(ctx0, + ggml_add(ctx0, ggml_mul(ctx0, local_f_clamped, mask), ggml_scale_bias(ctx0, mask, 1.0f, -1.0f)), + GGML_TYPE_I32); + ggml_tensor * weights = ggml_reshape_3d(ctx0, ggml_mul(ctx0, selected, mask), 1, n_expert_used, n_tokens); + ggml_tensor * up = build_lora_mm_id(ups[tier], inp, local); + ggml_tensor * gate = build_lora_mm_id(gates[tier], inp, local); + ggml_tensor * act = ggml_swiglu_split(ctx0, gate, up); + if (downs[tier]->ne[0] > act->ne[0]) act = ggml_pad(ctx0, act, downs[tier]->ne[0] - act->ne[0], 0, 0, 0); + ggml_tensor * result = ggml_mul(ctx0, build_lora_mm_id(downs[tier], act, local), weights); + weighted = weighted ? ggml_add(ctx0, weighted, result) : result; + } + for (int i = 0; i < n_expert_used; ++i) { + ggml_tensor * piece = ggml_view_2d(ctx0, weighted, cur->ne[0], n_tokens, weighted->nb[2], i*weighted->nb[1]); + moe_out = moe_out ? ggml_add(ctx0, moe_out, piece) : piece; + } + } else { + moe_out = build_moe_ffn(cur, model.layers[il].ffn_gate_inp, model.layers[il].ffn_up_exps, model.layers[il].ffn_gate_exps, @@ -919,6 +1004,7 @@ model.layers[il].ffn_up_exps_s, model.layers[il].ffn_gate_exps_s, model.layers[il].ffn_down_exps_s); + } cb(moe_out, "ffn_moe_out", il); // shared experts, as in the Qwen3Next reference --- /dev/null +++ b/common/wee-expert-prefetch.cpp @@ -0,0 +1,208 @@ +// Selected-range I/O and worker pool adapted from mbolt, Copyright (c) 2026 Dor Amir. +// MIT license: wee-expert-prefetch.LICENSE. Variable mixed-tier maps are local additions. +#include "wee-expert-prefetch.h" +#include "ggml-backend.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +struct range { uint64_t off, len; }; +using expert = std::array; + +struct prefetcher { + int fd = -1; + uint8_t * map = nullptr; + size_t size = 0, page = 0; + bool advice = false, parallel_advice = false, stop = false; + std::vector> layers; + std::vector calls; + std::vector ids; + std::vector resident_pages; + std::vector workers; + std::vector queue; + std::mutex batch_mutex, queue_mutex; + std::condition_variable work_cv, done_cv; + size_t next = 0, active = 0; + std::atomic bytes{0}, errors{0}; + uint64_t skipped = 0, ranges = 0, callback_us = 0; + + template static T read(std::ifstream & f) { + T value{}; + if (!f.read(reinterpret_cast(&value), sizeof(value))) throw std::runtime_error("Truncated expert map"); + return value; + } + + void init(const char * path, const std::string & model_path) { + std::ifstream f(path, std::ios::binary); + char magic[8]; + if (!f.read(magic, 8) || memcmp(magic, "WEEPF001", 8)) throw std::runtime_error("Invalid expert map"); + const uint64_t expected_size = read(f), dev = read(f), ino = read(f), mtime = read(f); + const uint32_t count = read(f); + if (count != 48) throw std::runtime_error("This expert map requires the 48-layer model"); + fd = open(model_path.c_str(), O_RDONLY); + struct stat st{}; + if (fd < 0 || fstat(fd, &st)) throw std::runtime_error("Cannot inspect expert model"); +#if defined(__APPLE__) + const uint64_t actual_time = uint64_t(st.st_mtimespec.tv_sec) * 1000000000 + st.st_mtimespec.tv_nsec; +#else + const uint64_t actual_time = uint64_t(st.st_mtim.tv_sec) * 1000000000 + st.st_mtim.tv_nsec; +#endif + if (uint64_t(st.st_size) != expected_size || uint64_t(st.st_dev) != dev || uint64_t(st.st_ino) != ino || actual_time != mtime) { + throw std::runtime_error("Expert map does not match the loaded model; regenerate it"); + } + size = expected_size; + layers.resize(count); calls.resize(count); + for (auto & layer : layers) { + const uint32_t n = read(f); + if (n == 0 || n > 512) throw std::runtime_error("Invalid expert count"); + layer.resize(n); + for (auto & e : layer) for (auto & r : e) { + r.off = read(f); r.len = read(f); + if (!r.len || r.len > 8u*1024*1024 || r.off > size || r.len > size-r.off) throw std::runtime_error("Invalid expert range"); + } + } + if (f.peek() != std::ifstream::traits_type::eof()) throw std::runtime_error("Trailing expert map data"); + page = size_t(getpagesize()); + map = static_cast(mmap(nullptr, size, PROT_READ, MAP_SHARED, fd, 0)); + if (map == MAP_FAILED) { map = nullptr; throw std::runtime_error("Cannot map expert residency probe"); } + const char * method = getenv("WEE_EXPERT_METHOD"); + if (method && strcmp(method, "read") && strcmp(method, "advice") && strcmp(method, "parallel-advice")) throw std::runtime_error("Unknown expert prefetch method"); + parallel_advice = method && !strcmp(method, "parallel-advice"); + advice = parallel_advice || (method && !strcmp(method, "advice")); + const char * nt = getenv("WEE_EXPERT_THREADS"); + const int threads = nt ? atoi(nt) : 4; + if (threads < 1 || threads > 8) throw std::runtime_error("Expert I/O thread count must be 1..8"); + if (!advice || parallel_advice) for (int i = 0; i < threads; ++i) workers.emplace_back([this] { worker(); }); + fprintf(stderr, "expert-prefetch: method=%s workers=%zu scratch_MiB=%zu layers=%zu\n", parallel_advice ? "parallel-advice" : advice ? "advice" : "read", workers.size(), advice ? 0 : workers.size(), layers.size()); + } + + void issue_advice(range r) { +#if defined(__APPLE__) + radvisory ra{off_t(r.off), int(r.len)}; + if (fcntl(fd, F_RDADVISE, &ra)) ++errors; +#else + if (posix_fadvise(fd, off_t(r.off), off_t(r.len), POSIX_FADV_WILLNEED)) ++errors; +#endif + bytes += r.len; + } + + void worker() { + std::vector scratch(advice ? 0 : 1u << 20); + std::unique_lock lock(queue_mutex); + for (;;) { + work_cv.wait(lock, [&] { return stop || next < queue.size(); }); + if (stop) return; + const range r = queue[next++]; ++active; + lock.unlock(); + if (advice) issue_advice(r); + uint64_t pos = r.off; + while (!advice && pos < r.off+r.len) { + const size_t n = std::min(scratch.size(), r.off+r.len-pos); + const ssize_t got = pread(fd, scratch.data(), n, off_t(pos)); + if (got < 0 && errno == EINTR) continue; + if (got <= 0) { ++errors; break; } + pos += uint64_t(got); bytes += uint64_t(got); + } + lock.lock(); --active; + if (next == queue.size() && active == 0) done_cv.notify_all(); + } + } + + bool resident(const range & r) { + const uint64_t begin = r.off / page * page; + const uint64_t end = (r.off+r.len+page-1) / page * page; + resident_pages.resize((end-begin)/page); + if (mincore(map+begin, end-begin, resident_pages.data())) return false; + return std::all_of(resident_pages.begin(), resident_pages.end(), [](char c) { return (c & 1) != 0; }); + } + + void fetch(ggml_tensor * t, size_t layer_id) { + std::lock_guard batch_lock(batch_mutex); + const auto start = std::chrono::steady_clock::now(); + ++calls[layer_id]; + // Only the exact selected IDs are read. The tensors and routing stay untouched. + if (t->ne[0] != 10 || t->ne[1] > 64 || t->ne[2] != 1 || t->ne[3] != 1) return; + ids.resize(ggml_nbytes(t)); + ggml_backend_tensor_get(t, ids.data(), 0, ids.size()); + std::vector selected; + for (int64_t b = 0; b < t->ne[1]; ++b) for (int64_t k = 0; k < t->ne[0]; ++k) { + int32_t id; memcpy(&id, ids.data()+b*t->nb[1]+k*t->nb[0], sizeof(id)); + if (id < 0 || size_t(id) >= layers[layer_id].size()) { ++errors; return; } + selected.push_back(id); + } + std::sort(selected.begin(), selected.end()); + selected.erase(std::unique(selected.begin(), selected.end()), selected.end()); + std::vector missing; + for (int32_t id : selected) for (const auto & r : layers[layer_id][id]) { + if (resident(r)) ++skipped; + else missing.push_back(r); + } + std::sort(missing.begin(), missing.end(), [](range a, range b) { return a.off < b.off; }); + std::vector merged; + for (auto r : missing) { + if (!merged.empty() && r.off <= merged.back().off+merged.back().len+65536) { + auto & last = merged.back(); last.len = std::max(last.off+last.len, r.off+r.len)-last.off; + } else merged.push_back(r); + } + ranges += merged.size(); + if (advice && !parallel_advice) { + for (auto r : merged) { + issue_advice(r); + } + } else if (!merged.empty()) { + std::unique_lock lock(queue_mutex); + queue = std::move(merged); next = 0; + work_cv.notify_all(); + done_cv.wait(lock, [&] { return next == queue.size() && active == 0; }); + } + callback_us += std::chrono::duration_cast(std::chrono::steady_clock::now()-start).count(); + } + + ~prefetcher() { + { std::lock_guard lock(queue_mutex); stop = true; } + work_cv.notify_all(); + for (auto & thread : workers) if (thread.joinable()) thread.join(); + if (map) munmap(map, size); + if (fd >= 0) close(fd); + fprintf(stderr, "expert-prefetch: ranges=%llu bytes=%llu resident_slices=%llu errors=%llu callback_us=%llu layer_calls=", + (unsigned long long)ranges, (unsigned long long)bytes.load(), (unsigned long long)skipped, + (unsigned long long)errors.load(), (unsigned long long)callback_us); + for (size_t i = 0; i < calls.size(); ++i) fprintf(stderr, "%s%llu", i ? "," : "", (unsigned long long)calls[i]); + fprintf(stderr, "\n"); + } +}; +std::unique_ptr instance; +} + +void wee_expert_prefetch_init(const char * path, const std::string & model_path) { + if (instance) throw std::runtime_error("Expert prefetch is restricted to one model per process"); + instance = std::make_unique(); + instance->init(path, model_path); +} + +bool wee_expert_prefetch_eval(ggml_tensor * t, bool ask, void *) { + if (t->type != GGML_TYPE_I32 || strncmp(t->name, "ffn_moe_topk-", 13)) return !ask; + char * end = nullptr; + const long layer = strtol(t->name+13, &end, 10); + if (!instance || end == t->name+13 || *end || layer < 0 || size_t(layer) >= instance->layers.size()) return !ask; + if (ask) return true; + instance->fetch(t, size_t(layer)); + return true; +} --- /dev/null +++ b/common/wee-expert-prefetch.h @@ -0,0 +1,6 @@ +#pragma once +#include "ggml.h" +#include + +void wee_expert_prefetch_init(const char * map_path, const std::string & model_path); +bool wee_expert_prefetch_eval(ggml_tensor * tensor, bool ask, void * user_data); --- /dev/null +++ b/common/wee-expert-prefetch.LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Dor Amir + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE.