| |
| |
| @@ -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 |
| |
| |
| @@ -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(); |
| |
| |
| @@ -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; |
| |
| |
| @@ -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}}, |
| |
| |
| @@ -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, |
| |
| |
| @@ -10,6 +10,7 @@ |
| #include <array> |
| #include <cinttypes> |
| #include <cstdint> |
| +#include <cstdlib> |
| #include <cstring> |
| #include <future> |
| #include <regex> |
| @@ -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++) { |
| |
| |
| @@ -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; |
| |
| |
| @@ -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 |
| |
| |
| @@ -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 <algorithm> |
| +#include <array> |
| +#include <atomic> |
| +#include <cerrno> |
| +#include <chrono> |
| +#include <condition_variable> |
| +#include <cstdio> |
| +#include <cstdlib> |
| +#include <cstring> |
| +#include <fstream> |
| +#include <memory> |
| +#include <mutex> |
| +#include <stdexcept> |
| +#include <thread> |
| +#include <vector> |
| +#include <fcntl.h> |
| +#include <sys/mman.h> |
| +#include <sys/stat.h> |
| +#include <unistd.h> |
| + |
| +namespace { |
| +struct range { uint64_t off, len; }; |
| +using expert = std::array<range, 3>; |
| + |
| +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<std::vector<expert>> layers; |
| + std::vector<uint64_t> calls; |
| + std::vector<uint8_t> ids; |
| + std::vector<char> resident_pages; |
| + std::vector<std::thread> workers; |
| + std::vector<range> queue; |
| + std::mutex batch_mutex, queue_mutex; |
| + std::condition_variable work_cv, done_cv; |
| + size_t next = 0, active = 0; |
| + std::atomic<uint64_t> bytes{0}, errors{0}; |
| + uint64_t skipped = 0, ranges = 0, callback_us = 0; |
| + |
| + template<class T> static T read(std::ifstream & f) { |
| + T value{}; |
| + if (!f.read(reinterpret_cast<char *>(&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<uint64_t>(f), dev = read<uint64_t>(f), ino = read<uint64_t>(f), mtime = read<uint64_t>(f); |
| + const uint32_t count = read<uint32_t>(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<uint32_t>(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<uint64_t>(f); r.len = read<uint64_t>(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<uint8_t *>(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<uint8_t> scratch(advice ? 0 : 1u << 20); |
| + std::unique_lock<std::mutex> 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<uint64_t>(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<std::mutex> 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<int32_t> 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<range> 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<range> 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<std::mutex> 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::microseconds>(std::chrono::steady_clock::now()-start).count(); |
| + } |
| + |
| + ~prefetcher() { |
| + { std::lock_guard<std::mutex> 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<prefetcher> 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<prefetcher>(); |
| + 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; |
| +} |
| |
| |
| @@ -0,0 +1,6 @@ |
| +#pragma once |
| +#include "ggml.h" |
| +#include <string> |
| + |
| +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); |
| |
| |
| @@ -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. |
|
|
