diff --git a/common/common.cpp b/common/common.cpp index 387cf280b..c05eec9df 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2141,6 +2141,15 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa params.warmup = false; return true; } + if (arg == "--prefetch-experts") { + params.prefetch_experts = true; + return true; + } + if (arg == "--prefetch-experts-threads") { + CHECK_ARG; + params.prefetch_experts_threads = std::stoi(argv[i]); + return true; + } if (arg == "--fit-margin") { CHECK_ARG; int32_t margin = std::stoi(argv[i]); @@ -3247,6 +3256,9 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param options.push_back({ "*", " --cpu-moe", "keep all MoE weights in CPU memory"}); options.push_back({ "*", " --n-cpu-moe N", "keep MoE weights of the first N layers in CPU memory"}); options.push_back({ "*", " --defer-experts", "defer expert mmap residency on Linux to reduce model load time"}); + options.push_back({ "*", " --prefetch-experts", "stream mmap'd MoE expert weights into the page cache on Linux"}); + options.push_back({ "*", " --prefetch-experts-threads N", + "number of expert prefetch workers, tune to drive speed/type (default: auto)"}); options.push_back({ "*", " --fit-margin N", "safety margin in MiB when auto-fitting model offloading"}); options.push_back({ "*", "-wgt, --worst-graph-tokens N", "number of tokens to use for worst-case graph"}); options.push_back({ "*", " --fit", "automatically determine which tensors to offload to the GPU(s)"}); @@ -4289,6 +4301,8 @@ struct llama_context_params common_context_params_to_llama(const gpt_params & pa cparams.min_experts = params.min_experts; cparams.thresh_experts = params.thresh_experts; cparams.only_active_experts = params.only_active_exps; + cparams.prefetch_experts = params.prefetch_experts; + cparams.prefetch_experts_threads = params.prefetch_experts_threads; cparams.max_extra_alloc = params.max_extra_alloc_MiB; cparams.mtp = params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP); cparams.mtp_op_type = MTP_OP_NONE; @@ -5286,6 +5300,8 @@ void yaml_dump_non_result_info(FILE * stream, const gpt_params & params, const l fprintf(stream, "merge_qkv: %s # default: false\n", params.merge_qkv ? "true" : "false"); fprintf(stream, "merge_up_gate_exps: %s # default: false\n", params.merge_up_gate_exps ? "true" : "false"); fprintf(stream, "defer_experts: %s # default: false\n", params.defer_experts ? "true" : "false"); + fprintf(stream, "prefetch_experts: %s # default: false\n", params.prefetch_experts ? "true" : "false"); + fprintf(stream, "prefetch_experts_threads: %d # default: 0 (auto)\n", params.prefetch_experts_threads); fprintf(stream, "max_extra_alloc: %d # default: 256\n", params.max_extra_alloc_MiB); fprintf(stream, "penalize_nl: %s # default: false\n", sparams.penalize_nl ? "true" : "false"); fprintf(stream, "ppl_output_type: %d # default: 0\n", params.ppl_output_type); diff --git a/common/common.h b/common/common.h index 2194aef19..c590dcdd7 100644 --- a/common/common.h +++ b/common/common.h @@ -443,6 +443,8 @@ struct gpt_params { bool merge_qkv = false; // if true, merge separate Q, K, V tensors into a single, contiguous tensor bool merge_up_gate_exps= false; // if true, merge ffn_up_exps and ffn_gate_exps into a single, contiguous tensor bool defer_experts = false; // if true, defer expert mmap residency to speed up model loading (Linux only) + bool prefetch_experts = false; // if true, stream mmap'd MoE expert weights into the page cache (Linux only) + int prefetch_experts_threads = 0; // number of expert prefetch workers (<=0 = auto) bool k_cache_hadamard = false; // if true, use Hadamard transform for the K-cache (only makes sense with quantized cache) bool v_cache_hadamard = false; // if true, use Hadamard transform for the V-cache (only makes sense with quantized cache, which requires FA) bool split_mode_graph_scheduling = false; // if true, force split mode graph scheduling diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index e820bcd18..2fb92dd90 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -102,6 +102,7 @@ extern "C" { GGML_API GGML_CALL bool ggml_backend_is_cpu (ggml_backend_t backend); GGML_API void ggml_backend_cpu_set_n_threads (ggml_backend_t backend_cpu, int n_threads); + GGML_API void ggml_backend_cpu_set_moe_expert_prefetch(ggml_backend_t backend_cpu, bool enable); GGML_API void ggml_backend_cpu_set_abort_callback(ggml_backend_t backend_cpu, ggml_abort_callback abort_callback, void * abort_callback_data); // Create a backend buffer from an existing pointer @@ -215,6 +216,11 @@ extern "C" { GGML_API void ggml_backend_sched_set_split_mode_graph(ggml_backend_sched_t sched, bool on_or_off, bool async); GGML_API void ggml_backend_sched_set_max_extra_alloc(ggml_backend_sched_t sched, int extra_alloc_MiB); + // prefetch mmap'd MoE expert weights into the page cache + GGML_API bool ggml_backend_prefetch_init(int n_threads); + GGML_API void ggml_backend_prefetch_register_mapping(const void * addr, size_t size); + GGML_API void ggml_backend_prefetch_unregister_mapping(const void * addr); + // // Utils // diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index a24692edb..0f0d08ec2 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -831,6 +831,9 @@ extern "C" { // abort ggml_graph_compute when true ggml_abort_callback abort_callback; void * abort_callback_data; + + // read-ahead selected MoE expert weights in the CPU matmul-id kernels + bool moe_expert_prefetch; }; enum ggml_cgraph_eval_order { diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 290d7f30f..48848e4c6 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -1297,6 +1297,8 @@ add_library(ggml ggml.c ggml-alloc.c ggml-backend.cpp + ggml-moe-prefetch.cpp + ggml-moe-prefetch.h ggml-quants.c ggml-quants.h ${GGML_SOURCES_CUDA} ${GGML_HEADERS_CUDA} diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 19cb57348..3d065d7e6 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -2,6 +2,7 @@ #include "ggml-alloc.h" #include "ggml-impl.h" #include "ggml-rpc.h" +#include "ggml-moe-prefetch.h" #include #include @@ -803,6 +804,8 @@ struct ggml_backend_cpu_context { ggml_abort_callback abort_callback; void * abort_callback_data; + + bool moe_expert_prefetch; }; GGML_CALL static const char * ggml_backend_cpu_name(ggml_backend_t backend) { @@ -847,6 +850,7 @@ GGML_CALL static ggml_backend_graph_plan_t ggml_backend_cpu_graph_plan_create(gg cpu_plan->cplan.abort_callback = cpu_ctx->abort_callback; cpu_plan->cplan.abort_callback_data = cpu_ctx->abort_callback_data; + cpu_plan->cplan.moe_expert_prefetch = cpu_ctx->moe_expert_prefetch; return cpu_plan; } @@ -886,6 +890,7 @@ GGML_CALL static enum ggml_status ggml_backend_cpu_graph_compute(ggml_backend_t cplan.abort_callback = cpu_ctx->abort_callback; cplan.abort_callback_data = cpu_ctx->abort_callback_data; + cplan.moe_expert_prefetch = cpu_ctx->moe_expert_prefetch; return ggml_graph_compute(cgraph, &cplan); } @@ -953,6 +958,7 @@ ggml_backend_t ggml_backend_cpu_init(void) { ctx->work_size = 0; ctx->abort_callback = NULL; ctx->abort_callback_data = NULL; + ctx->moe_expert_prefetch = false; ggml_backend_t cpu_backend = (ggml_backend_t)malloc(sizeof(struct ggml_backend)); if (cpu_backend == NULL) { @@ -979,6 +985,13 @@ void ggml_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads) { ctx->n_threads = n_threads; } +void ggml_backend_cpu_set_moe_expert_prefetch(ggml_backend_t backend_cpu, bool enable) { + GGML_ASSERT(ggml_backend_is_cpu(backend_cpu)); + + struct ggml_backend_cpu_context * ctx = (struct ggml_backend_cpu_context *)backend_cpu->context; + ctx->moe_expert_prefetch = enable; +} + void ggml_backend_cpu_set_abort_callback(ggml_backend_t backend_cpu, ggml_abort_callback abort_callback, void * abort_callback_data) { GGML_ASSERT(ggml_backend_is_cpu(backend_cpu)); @@ -1216,6 +1229,22 @@ void ggml_backend_sched_set_max_extra_alloc(ggml_backend_sched_t sched, int extr } } +bool ggml_backend_prefetch_init(int n_threads) { + if (n_threads <= 0) { + n_threads = std::max(1, std::min(8, (int) std::thread::hardware_concurrency())); + } + ggml_moe_prefetch_set_n_threads(n_threads); + return ggml_moe_prefetch_enabled(); +} + +void ggml_backend_prefetch_register_mapping(const void * addr, size_t size) { + ggml_moe_prefetch_register_mapping(addr, size); +} + +void ggml_backend_prefetch_unregister_mapping(const void * addr) { + ggml_moe_prefetch_unregister_mapping(addr); +} + static inline bool ggml_backend_sched_offload_enabled(ggml_backend_sched_t sched, enum ggml_op op) { int int_op = (int)op; if (!sched || op < 0 || op >= GGML_OP_COUNT) return false; @@ -2049,6 +2078,11 @@ static void ggml_backend_sched_copy_inputs(ggml_backend_sched_t sched, ggml_back last_ids_tensor = ids_tensor; } + // when the expert prefetch engine streamed this tensor ahead + // (see the lookahead in compute_splits), wait for it so the + // host-side reads below hit warm page cache instead of faulting + ggml_moe_prefetch_wait(input); + const size_t expert_size = input->ne[2] > 1 ? input->nb[2] : input->nb[1]; if (input->ne[2] > 1) { @@ -2365,6 +2399,67 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s std::vector unique_ids; ggml_tensor * last_ids_tensor = nullptr; + // MoE expert prefetch; pre-scan the splits for expert-weight matmuls whose + // weights live in a host mmap. Two behaviors, both page-cache warmers. + // - lookahead streams the next splits' expert tensors in full while split + // i computes (batch/PP graphs only, where most experts are hit) + // - selective enqueues just the selected expert slices of a split's + // up/gate/down tensors once its ids have been copied to the host + struct moe_split_info { + int split; + int64_t n_tokens; + std::vector host_weights; // originals (host mmap) + std::vector nodes; // MoE nodes computed on a host buffer + }; + std::vector moe_infos; + const bool moe_prefetch = ggml_moe_prefetch_enabled(); + static const size_t moe_ahead = [] { + const char * env = getenv("GGML_MOE_PREFETCH_AHEAD"); + return env ? (size_t) std::max(0, atoi(env)) : (size_t) 3; + }(); + if (moe_prefetch) { + ggml_moe_prefetch_new_epoch(); + for (int i = 0; i < sched->n_splits; i++) { + moe_split_info info; + info.split = i; + info.n_tokens = 0; + for (int n = 0; n < splits[i].graph.n_nodes; ++n) { + ggml_tensor * node = splits[i].graph.nodes[n]; + if (node->op != GGML_OP_MUL_MAT_ID && node->op != GGML_OP_MOE_FUSED_UP_GATE) { + continue; + } + ggml_tensor * node_ids = node->op == GGML_OP_MUL_MAT_ID ? node->src[2] : node->src[3]; + info.n_tokens = std::max(info.n_tokens, node_ids ? node_ids->ne[1] : 0); + ggml_tensor * ws[2] = { node->src[0], node->op == GGML_OP_MOE_FUSED_UP_GATE ? node->src[1] : nullptr }; + bool node_on_host = false; + for (ggml_tensor * w : ws) { + if (w && w->buffer && ggml_backend_buffer_is_host(w->buffer) && + ggml_backend_buffer_get_usage(w->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + info.host_weights.push_back(w); + node_on_host = true; + } + } + if (node_on_host) { + info.nodes.push_back(node); + } + } + // in the offloaded case the node's weight srcs were rewritten to + // device copies; the host originals arrive as split inputs + for (int j = 0; j < splits[i].n_inputs; ++j) { + ggml_tensor * input = splits[i].inputs[j]; + if (input->ne[2] > 1 && input->buffer && ggml_backend_buffer_is_host(input->buffer) && + ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + info.host_weights.push_back(input); + } + } + if (!info.host_weights.empty()) { + moe_infos.push_back(std::move(info)); + } + } + } + size_t moe_next = 0; // first moe_infos entry with split >= current split + size_t moe_enq = 0; // moe_infos entries already enqueued for lookahead + for (int i = 0; i < sched->n_splits; i++) { #if IK_PRINT_TIMING int64_t tim1 = ggml_time_us(); @@ -2373,9 +2468,34 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s int split_backend_id = split->backend_id; ggml_backend_t split_backend = sched->backends[split_backend_id]; + if (moe_prefetch && !moe_infos.empty()) { + while (moe_next < moe_infos.size() && moe_infos[moe_next].split < i) { + moe_next++; + } + // keep the next `moe_ahead` MoE-bearing splits streaming in + const size_t want_end = std::min(moe_next + moe_ahead, moe_infos.size()); + for (size_t k = std::max(moe_enq, moe_next); k < want_end; ++k) { + if (moe_infos[k].n_tokens >= 32) { // batch graphs touch most experts (min batch offload) + for (ggml_tensor * w : moe_infos[k].host_weights) { + ggml_moe_prefetch_tensor(w); + } + } + moe_enq = k + 1; + } + } + // copy the input tensors to the split backend ggml_backend_sched_copy_inputs(sched, split, sched->needs_sync, ids, unique_ids, last_ids_tensor); + // ids are now final and host-visible; enqueue the selected expert + // slices of this split's host-computed MoE matmuls (up/gate and down + // share one ids tensor, so the down weights stream in during up/gate) + if (moe_prefetch && moe_next < moe_infos.size() && moe_infos[moe_next].split == i) { + for (ggml_tensor * node : moe_infos[moe_next].nodes) { + ggml_moe_prefetch_node(node); + } + } + if (split->n_inputs > 0 && !sched->own_cpy[split_backend_id]) { sched->needs_sync[split_backend_id] = true; } else { @@ -2390,6 +2510,15 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s return ec; } + // the pages the lookahead streamer just read for this split are one-shot + // streaming traffic; MADV_COLD them so the decode working set survives + if (moe_prefetch && moe_next < moe_infos.size() && moe_infos[moe_next].split == i && + moe_infos[moe_next].n_tokens >= 32) { + for (ggml_tensor * w : moe_infos[moe_next].host_weights) { + ggml_moe_prefetch_cold(w); + } + } + // record the event of this copy if (split->n_inputs > 0) { if (sched->events[split_backend_id][sched->cur_copy] != NULL) { diff --git a/ggml/src/ggml-moe-prefetch.cpp b/ggml/src/ggml-moe-prefetch.cpp new file mode 100644 index 000000000..52075f6ed --- /dev/null +++ b/ggml/src/ggml-moe-prefetch.cpp @@ -0,0 +1,470 @@ +#include "ggml-moe-prefetch.h" + +#if defined(__linux__) + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// chunk granularity; big enough to amortize syscall cost, +// small enough to spread one expert (~MBs) across several workers. +static constexpr size_t GGML_MOE_PREFETCH_CHUNK = 2u*1024u*1024u; +// lookahead jobs are dropped beyond this queue depth so a stalled consumer +// cannot accumulate unbounded work. +static constexpr size_t GGML_MOE_PREFETCH_MAX_QUEUE = 65536; + +// experts are faulted in with MADV_POPULATE_READ, which brings pages into the +// page cache AND this mm's page tables, so consumers take no minor faults. +// Unsupported kernels leave the engine off (the callers fall back to +// madvise(MADV_WILLNEED) hints) +#ifndef MADV_POPULATE_READ +#define MADV_POPULATE_READ 22 +#endif + +namespace { + +struct mapping_entry { + uintptr_t base; + size_t size; +}; + +struct ticket { + std::atomic pending{0}; + uint64_t epoch = 0; + bool track_reads = false; + // chunks the workers actually read from storage this epoch (i.e. pages + // that were cold before the sweep); guarded by prefetch_pool::mtx. + // Used by ggml_moe_prefetch_cold() to deactivate streaming traffic. + std::vector> read_chunks; +}; + +struct job { + uintptr_t addr; + uint32_t len; + std::shared_ptr tk; +}; + +struct prefetch_pool { + std::mutex mtx; + std::condition_variable cv_work; // workers sleep here + std::condition_variable cv_done; // waiters sleep here + std::deque queue; + std::vector workers; + bool shutdown = false; + + std::unordered_map> tickets; + + // cumulative observability counters (GGML_MOE_PREFETCH_DEBUG) + std::atomic n_jobs{0}; + std::atomic n_skipped{0}; + std::atomic bytes_populated{0}; + std::atomic n_colded{0}; + + ~prefetch_pool() { stop(); } + + void stop() { + { + std::lock_guard lock(mtx); + shutdown = true; + // complete outstanding tickets so no waiter blocks forever + for (auto & j : queue) { + if (j.tk) { + j.tk->pending.fetch_sub(1, std::memory_order_acq_rel); + } + } + queue.clear(); + cv_work.notify_all(); + cv_done.notify_all(); + } + for (auto & w : workers) { + if (w.joinable()) w.join(); + } + workers.clear(); + if (getenv("GGML_MOE_PREFETCH_DEBUG") && n_jobs.load() > 0) { + fprintf(stderr, "%s: jobs=%llu skipped_resident=%llu bytes_populated=%.2f GiB colded=%llu\n", __func__, + (unsigned long long) n_jobs.load(), (unsigned long long) n_skipped.load(), + (double) bytes_populated.load()/(1024.0*1024.0*1024.0), + (unsigned long long) n_colded.load()); + } + } + + void start(int n_threads) { + stop(); + std::lock_guard lock(mtx); + shutdown = false; + workers.reserve(n_threads); + for (int i = 0; i < n_threads; ++i) { + workers.emplace_back([this] { run(); }); + } + } + + static bool chunk_resident(uintptr_t addr, size_t len) { + const long page = sysconf(_SC_PAGESIZE); + const uintptr_t astart = addr & ~(uintptr_t)(page - 1); + const size_t alen = (addr + len) - astart; + const size_t npages = (alen + page - 1)/page; + if (npages > GGML_MOE_PREFETCH_CHUNK/4096 + 2) { + return false; + } + unsigned char vec[GGML_MOE_PREFETCH_CHUNK/4096 + 2]; + if (mincore((void *)astart, alen, vec) != 0) { + return false; + } + for (size_t i = 0; i < npages; ++i) { + if (!(vec[i] & 1)) return false; + } + return true; + } + + void run() { + const long page = sysconf(_SC_PAGESIZE); + for (;;) { + job j; + { + std::unique_lock lock(mtx); + cv_work.wait(lock, [this] { return shutdown || !queue.empty(); }); + if (shutdown) return; + j = std::move(queue.front()); + queue.pop_front(); + } + n_jobs.fetch_add(1, std::memory_order_relaxed); + bool done_read = false; + // skip when every page is already resident; keeps the drain rate + // high when the cache is warm + if (chunk_resident(j.addr, j.len)) { + n_skipped.fetch_add(1, std::memory_order_relaxed); + } else { + const uintptr_t astart = j.addr & ~(uintptr_t)(page - 1); + const size_t alen = ((j.addr + j.len + page - 1) & ~(uintptr_t)(page - 1)) - astart; + if (madvise((void *)astart, alen, MADV_POPULATE_READ) == 0) { + bytes_populated.fetch_add(alen, std::memory_order_relaxed); + done_read = true; + } // on failure the fault path takes over + } + if (j.tk) { + const int left = j.tk->pending.fetch_sub(1, std::memory_order_acq_rel); + std::lock_guard lock(mtx); + if (done_read && j.tk->track_reads) { + j.tk->read_chunks.emplace_back(j.addr, j.len); + } + if (left == 1) { + cv_done.notify_all(); + } + } + } + } +}; + +struct prefetch_state { + std::mutex reg_mtx; + std::vector mappings; + + std::atomic epoch{1}; + + std::mutex pool_mtx; + std::shared_ptr pool; +}; + +static prefetch_state & state() { + static prefetch_state s; + return s; +} + +// true when [p, p+len) lies inside a registered mmap +static bool is_mapped(const void * p, size_t len) { + auto & s = state(); + std::lock_guard lock(s.reg_mtx); + const uintptr_t a = (uintptr_t)p; + for (const auto & m : s.mappings) { + if (a >= m.base && a + len <= m.base + m.size) return true; + } + return false; +} + +// collect [offset, offset+len) ranges of the selected experts of weight tensor w +static void collect_selected_ranges(const ggml_tensor * w, const ggml_tensor * ids, + std::vector> & ranges) { + const int64_t n_as = w->ne[2]; + const size_t stride = w->nb[2]; + const size_t wbytes = ggml_nbytes(w); + + if (n_as <= 1 || stride == 0) { + ranges.emplace_back(0, wbytes); + return; + } + + std::vector seen((n_as + 31)/32, 0); + for (int64_t i1 = 0; i1 < ids->ne[1]; ++i1) { + for (int64_t i0 = 0; i0 < ids->ne[0]; ++i0) { + const int32_t id = *(const int32_t *)((const char *)ids->data + i1*ids->nb[1] + i0*ids->nb[0]); + if (id < 0 || id >= n_as) continue; // ids may hold -1 sentinels (ggml_top_k_thresh) + seen[id >> 5] |= 1u << (id & 31); + } + } + // coalesce consecutive experts into single ranges + int64_t id = 0; + while (id < n_as) { + while (id < n_as && !(seen[id >> 5] & (1u << (id & 31)))) ++id; + if (id >= n_as) break; + int64_t first = id; + while (id < n_as && (seen[id >> 5] & (1u << (id & 31)))) ++id; + const size_t off = (size_t)first*stride; + const size_t len = std::min((size_t)(id - first)*stride, wbytes - off); + if (off < wbytes && len > 0) { + ranges.emplace_back(off, len); + } + } +} + +// enqueue ranges of tensor w; returns false when the engine is off or w is not mmap-registered +static bool enqueue_ranges(const ggml_tensor * w, const std::vector> & ranges, bool urgent, bool track_reads) { + auto & s = state(); + std::shared_ptr pool_sp; + { + std::lock_guard plock(s.pool_mtx); + pool_sp = s.pool; + } + if (!pool_sp) return false; + + if (!is_mapped(w->data, ggml_nbytes(w))) return false; + + const uint64_t cur_epoch = s.epoch.load(std::memory_order_relaxed); + + prefetch_pool & pool = *pool_sp; + std::lock_guard lock(pool.mtx); + if (pool.shutdown) return false; + + auto & tk = pool.tickets[w]; + if (!tk) tk = std::make_shared(); + if (tk->epoch == cur_epoch) { + return true; // already enqueued for this scheduler pass + } + if (!urgent && pool.queue.size() > GGML_MOE_PREFETCH_MAX_QUEUE) { + return true; // drop lookahead under backpressure; do not mark the epoch + } + tk->epoch = cur_epoch; + tk->track_reads = track_reads; + tk->read_chunks.clear(); + + std::vector jobs; + for (const auto & r : ranges) { + for (size_t o = r.first; o < r.first + r.second; o += GGML_MOE_PREFETCH_CHUNK) { + const size_t len = std::min(GGML_MOE_PREFETCH_CHUNK, r.first + r.second - o); + jobs.push_back({(uintptr_t)w->data + o, (uint32_t)len, tk}); + } + } + if (jobs.empty()) return true; + + tk->pending.fetch_add((int)jobs.size(), std::memory_order_acq_rel); + if (urgent) { + pool.queue.insert(pool.queue.begin(), std::make_move_iterator(jobs.begin()), std::make_move_iterator(jobs.end())); + } else { + pool.queue.insert(pool.queue.end(), std::make_move_iterator(jobs.begin()), std::make_move_iterator(jobs.end())); + } + pool.cv_work.notify_all(); + return true; +} + +static void node_weights_and_ids(const ggml_tensor * node, const ggml_tensor * & w0, const ggml_tensor * & w1, const ggml_tensor * & ids) { + w0 = nullptr; w1 = nullptr; ids = nullptr; + if (node->op == GGML_OP_MUL_MAT_ID) { + w0 = node->src[0]; + ids = node->src[2]; + } else if (node->op == GGML_OP_MOE_FUSED_UP_GATE) { + w0 = node->src[0]; + w1 = node->src[1]; // NULL when up/gate are packed into one tensor + ids = node->src[3]; + } +} + +// legacy fallback; hints kernel readahead for the selected experts when the +// read pool is unavailable +static void legacy_madvise(const ggml_tensor * w, const ggml_tensor * ids) { + if (!w || !w->data || !ids || !ids->data) return; + std::vector> ranges; + collect_selected_ranges(w, ids, ranges); + const uintptr_t page_mask = (uintptr_t)sysconf(_SC_PAGESIZE) - 1; + const char * base = (const char *)w->data; + for (const auto & r : ranges) { + const uintptr_t start = (uintptr_t)(base + r.first); + const uintptr_t astart = start & ~page_mask; + (void) madvise((void *)astart, (size_t)(start + r.second - astart), MADV_WILLNEED); + } +} + +} // namespace + +void ggml_moe_prefetch_register_mapping(const void * base, size_t size) { + if (!base || size == 0) return; + auto & s = state(); + std::lock_guard lock(s.reg_mtx); + for (const auto & m : s.mappings) { + if (m.base == (uintptr_t)base) return; // contexts may re-register the same model + } + s.mappings.push_back({(uintptr_t)base, size}); +} + +void ggml_moe_prefetch_unregister_mapping(const void * base) { + auto & s = state(); + std::lock_guard lock(s.reg_mtx); + // a queued job may still point into this range; its madvise then fails + // (ENOMEM once unmapped) and the worker skips the chunk + s.mappings.erase(std::remove_if(s.mappings.begin(), s.mappings.end(), + [base](const mapping_entry & m) { return m.base == (uintptr_t)base; }), + s.mappings.end()); +} + +static bool populate_read_supported() { + const long page = sysconf(_SC_PAGESIZE); + void * p = mmap(nullptr, (size_t)page, PROT_READ, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (p == MAP_FAILED) return false; + const bool ok = madvise(p, (size_t)page, MADV_POPULATE_READ) == 0; + munmap(p, (size_t)page); + return ok; +} + +void ggml_moe_prefetch_set_n_threads(int n_threads) { + auto & s = state(); + std::lock_guard lock(s.pool_mtx); + if (n_threads <= 0) { + s.pool.reset(); + return; + } + if (!populate_read_supported()) { + fprintf(stderr, "%s: MADV_POPULATE_READ not supported; MoE prefetch disabled\n", __func__); + s.pool.reset(); + return; + } + if (s.pool && s.pool->workers.size() == (size_t)n_threads) { + return; + } + s.pool.reset(); + s.pool = std::make_shared(); + s.pool->start(n_threads); +} + +bool ggml_moe_prefetch_enabled(void) { + auto & s = state(); + std::lock_guard lock(s.pool_mtx); + return s.pool != nullptr; +} + +void ggml_moe_prefetch_new_epoch(void) { + state().epoch.fetch_add(1, std::memory_order_relaxed); +} + +void ggml_moe_prefetch_node(const struct ggml_tensor * node) { + const ggml_tensor * w0; const ggml_tensor * w1; const ggml_tensor * ids; + node_weights_and_ids(node, w0, w1, ids); + if (!w0 || !ids || !ids->data) return; + + std::vector> ranges; + collect_selected_ranges(w0, ids, ranges); + enqueue_ranges(w0, ranges, /*urgent =*/ true, /*track_reads =*/ false); + if (w1) { + ranges.clear(); + collect_selected_ranges(w1, ids, ranges); + enqueue_ranges(w1, ranges, /*urgent =*/ true, /*track_reads =*/ false); + } +} + +void ggml_moe_prefetch_tensor(const struct ggml_tensor * w) { + if (!w || !w->data) return; + std::vector> ranges; + ranges.emplace_back(0, ggml_nbytes(w)); + enqueue_ranges(w, ranges, /*urgent =*/ false, /*track_reads =*/ true); +} + +#ifndef MADV_COLD +#define MADV_COLD 20 +#endif +void ggml_moe_prefetch_cold(const struct ggml_tensor * w) { + auto & s = state(); + std::shared_ptr pool; + { + std::lock_guard plock(s.pool_mtx); + pool = s.pool; + } + if (!pool) return; + std::vector> chunks; + { + std::lock_guard lock(pool->mtx); + auto it = pool->tickets.find(w); + if (it == pool->tickets.end() || !it->second->track_reads) return; + chunks.swap(it->second->read_chunks); + } + const uintptr_t page_mask = (uintptr_t)sysconf(_SC_PAGESIZE) - 1; + for (const auto & c : chunks) { + const uintptr_t astart = c.first & ~page_mask; + (void) madvise((void *)astart, (size_t)(c.first + c.second - astart), MADV_COLD); + } + pool->n_colded.fetch_add(chunks.size(), std::memory_order_relaxed); +} + +void ggml_moe_prefetch_wait(const struct ggml_tensor * w) { + auto & s = state(); + std::shared_ptr pool; + { + std::lock_guard plock(s.pool_mtx); + pool = s.pool; + } + if (!pool) return; + std::shared_ptr tk; + { + std::lock_guard lock(pool->mtx); + auto it = pool->tickets.find(w); + if (it == pool->tickets.end()) return; + tk = it->second; + } + if (tk->pending.load(std::memory_order_acquire) <= 0) return; + // the shared_ptr keeps the pool alive; shutdown wakes cv_done + std::unique_lock lock(pool->mtx); + pool->cv_done.wait(lock, [&] { + return pool->shutdown || tk->pending.load(std::memory_order_acquire) <= 0; + }); +} + +void ggml_moe_prefetch_kernel_hook(const struct ggml_tensor * node, int ith) { + if (ith != 0) return; + const ggml_tensor * w0; const ggml_tensor * w1; const ggml_tensor * ids; + node_weights_and_ids(node, w0, w1, ids); + if (!w0 || !ids || !ids->data) return; + + if (!ggml_moe_prefetch_enabled()) { + legacy_madvise(w0, ids); + if (w1) legacy_madvise(w1, ids); + return; + } + // fire-and-forget enqueue, idempotent within the current epoch; a no-op + // when the scheduler hook already covered this node (the self-enqueue + // handles pure-CPU graphs). We deliberately do not wait; the compute + // threads' demand faults overlap with the workers' populates, which run ahead + // in the same expert-index order. + ggml_moe_prefetch_node(node); +} + +#else // !__linux__ + +void ggml_moe_prefetch_register_mapping(const void *, size_t) {} +void ggml_moe_prefetch_unregister_mapping(const void *) {} +void ggml_moe_prefetch_set_n_threads(int) {} +bool ggml_moe_prefetch_enabled(void) { return false; } +void ggml_moe_prefetch_new_epoch(void) {} +void ggml_moe_prefetch_node(const struct ggml_tensor *) {} +void ggml_moe_prefetch_tensor(const struct ggml_tensor *) {} +void ggml_moe_prefetch_wait(const struct ggml_tensor *) {} +void ggml_moe_prefetch_kernel_hook(const struct ggml_tensor *, int) {} +void ggml_moe_prefetch_cold(const struct ggml_tensor *) {} + +#endif diff --git a/ggml/src/ggml-moe-prefetch.h b/ggml/src/ggml-moe-prefetch.h new file mode 100644 index 000000000..5e323d705 --- /dev/null +++ b/ggml/src/ggml-moe-prefetch.h @@ -0,0 +1,52 @@ +#pragma once + +#include "ggml.h" + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// mapping registry; a registered range marks tensors as mmap-backed and +// eligible for prefetch. Nothing is opened or retained beyond (base, size). +void ggml_moe_prefetch_register_mapping (const void * base, size_t size); +void ggml_moe_prefetch_unregister_mapping(const void * base); + +// worker pool control; n_threads > 0 (re)creates the pool, <= 0 shuts it down +void ggml_moe_prefetch_set_n_threads(int n_threads); +bool ggml_moe_prefetch_enabled(void); + +// the epoch is bumped once per scheduler pass so per-tensor enqueues are +// idempotent within one graph execution but re-issued on the next. +void ggml_moe_prefetch_new_epoch(void); + +// selective enqueue for one MoE node (GGML_OP_MUL_MAT_ID or +// GGML_OP_MOE_FUSED_UP_GATE); reads the ids tensor bytes (must be final and +// host-visible) and enqueues the selected expert ranges of src weights. +void ggml_moe_prefetch_node(const struct ggml_tensor * node); + +// full-tensor lookahead enqueue (low priority), e.g. for prompt processing +// where the next layer's experts are needed in bulk. +void ggml_moe_prefetch_tensor(const struct ggml_tensor * w); + +// block until all pending prefetch jobs for tensor w are complete. +// Returns immediately when nothing is pending. +void ggml_moe_prefetch_wait(const struct ggml_tensor * w); + +// MADV_COLD the pages the streamer had to read from storage for tensor w, +// so prompt-processing streaming traffic is reclaimed ahead of the decode +// working set. Pages already resident before the sweep are left alone. +// Only applies to tensors enqueued via ggml_moe_prefetch_tensor. +void ggml_moe_prefetch_cold(const struct ggml_tensor * w); + +// kernel-entry hook, called at MoE matmul start (gated by +// cplan->moe_expert_prefetch). Thread 0 fire-and-forget enqueues when the +// scheduler hook did not run for this epoch (pure-CPU graphs); falls back to +// madvise(MADV_WILLNEED) when the engine is off. +void ggml_moe_prefetch_kernel_hook(const struct ggml_tensor * node, int ith); + +#ifdef __cplusplus +} +#endif diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index eb73b4ce6..c105ea37f 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -11,6 +11,7 @@ #include "ggml-quants.h" #include "ggml.h" #include "ggml-aarch64.h" +#include "ggml-moe-prefetch.h" #include "iqk/iqk_quantize.h" #include "iqk/iqk_cpu_ops.h" #if GGML_USE_IQK_MULMAT @@ -17354,6 +17355,11 @@ static void ggml_compute_forward_mul_mat_id( const int n_ids = ids->ne[0]; // n_expert_used const int n_as = ne02; // n_expert + // kick off read-ahead of the selected experts before src1 quantization, so storage reads overlap that work + if (params->shared->cplan && params->shared->cplan->moe_expert_prefetch) { + ggml_moe_prefetch_kernel_hook(dst, ith); + } + char * wdata_src1_end = (src1->type == vec_dot_type) ? (char *) params->wdata : (char *) params->wdata + GGML_PAD(ggml_row_size(vec_dot_type, src1->ne[0])*ggml_nrows(src1), sizeof(int64_t)); @@ -17621,6 +17627,12 @@ static void ggml_compute_forward_mul_mat_id_up_gate( const int n_ids = ids->ne[0]; // n_expert_used const int n_as = ne02; // n_expert + // read-ahead of the selected experts for both the up and gate weight tensors + // (gate is null when up/gate are merged into a single tensor) + if (params->shared->cplan && params->shared->cplan->moe_expert_prefetch) { + ggml_moe_prefetch_kernel_hook(dst, ith); + } + char * wdata_src1_end = (src1->type == vec_dot_type) ? (char *) params->wdata : (char *) params->wdata + GGML_PAD(ggml_row_size(vec_dot_type, src1->ne[0])*ggml_nrows(src1), sizeof(int64_t)); diff --git a/include/llama.h b/include/llama.h index 39a13ee4f..c2c183c15 100644 --- a/include/llama.h +++ b/include/llama.h @@ -497,6 +497,8 @@ extern "C" { int min_experts; float thresh_experts; bool only_active_experts; + bool prefetch_experts; // if true, stream mmap'd MoE expert weights into the page cache (Linux only) + int prefetch_experts_threads; // number of expert prefetch workers (<=0 = auto) bool k_cache_hadamard; // if true, apply Hadamard transfrom to K-cache bool v_cache_hadamard; // if true, apply Hadamard transfrom to V-cache (needs FA) bool split_mode_graph_scheduling; // if true, force split mode graph scheduling diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 7991df05f..b4fe0613f 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -39,6 +39,7 @@ struct llama_cparams { bool fused_mmad; bool rope_cache; bool graph_reuse; + bool prefetch_experts; bool k_cache_hadamard; bool v_cache_hadamard; bool dsa_indexer_hadamard = true; // apply Walsh-Hadamard rotation to DSA indexer q/k (precision) diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index 25dc8163a..bbea43058 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -3,6 +3,7 @@ #include "llama-impl.h" #include "ggml.h" +#include "ggml-backend.h" #include #include @@ -435,6 +436,7 @@ struct llama_mmap::impl { } ~impl() { + ggml_backend_prefetch_unregister_mapping(addr); for (const auto & frag : mapped_fragments) { if (munmap((char *) addr + frag.first, frag.second - frag.first)) { LLAMA_LOG_WARN("warning: munmap failed: %s\n", strerror(errno)); diff --git a/src/llama.cpp b/src/llama.cpp index e05e95ebf..fac695d4e 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -5137,6 +5137,7 @@ static void llama_graph_compute( if (lctx.backend_cpu != nullptr) { ggml_backend_cpu_set_n_threads(lctx.backend_cpu, n_threads); ggml_backend_cpu_set_abort_callback(lctx.backend_cpu, lctx.abort_callback, lctx.abort_callback_data); + ggml_backend_cpu_set_moe_expert_prefetch(lctx.backend_cpu, lctx.cparams.prefetch_experts); } ggml_backend_sched_graph_compute_async(lctx.sched, gf); @@ -5158,6 +5159,7 @@ static void llama_graph_compute_sched( if (lctx.backend_cpu != nullptr) { ggml_backend_cpu_set_n_threads(lctx.backend_cpu, n_threads); ggml_backend_cpu_set_abort_callback(lctx.backend_cpu, lctx.abort_callback, lctx.abort_callback_data); + ggml_backend_cpu_set_moe_expert_prefetch(lctx.backend_cpu, lctx.cparams.prefetch_experts); } ggml_backend_sched_graph_compute_async(sched, gf); @@ -6624,6 +6626,8 @@ struct llama_context_params llama_context_default_params() { /*.min_experts =*/ -1, /*.thtesh_experts =*/ 0.0f, /*.only_active_experts =*/ false, + /*.prefetch_experts =*/ false, + /*.prefetch_experts_threads =*/ 0, /*.k_cache_hadamard =*/ false, /*.v_cache_hadamard =*/ false, /*.split_mode_graph_scheduling =*/ false, @@ -7047,6 +7051,7 @@ struct llama_context * llama_init_from_model( if (cparams.dsa && (model->split_mode == LLAMA_SPLIT_MODE_GRAPH || model->split_mode == LLAMA_SPLIT_MODE_ATTN)) { LLAMA_LOG_WARN("%s: --dsa is not active under -sm graph/attn (tensor-parallel attention has no indexer); running dense MLA\n", __func__); } + cparams.prefetch_experts = params.prefetch_experts; cparams.k_cache_hadamard = params.k_cache_hadamard; cparams.v_cache_hadamard = params.v_cache_hadamard; // Folding H into wv_b/wk_b_pp permanently mutates the model; a later context @@ -7552,6 +7557,13 @@ struct llama_context * llama_init_from_model( LLAMA_LOG_INFO("%s: enabling only_active_experts scheduling\n", __func__); ggml_backend_sched_set_only_active_experts(ctx->sched, true); } + if (params.prefetch_experts) { + LLAMA_LOG_INFO("%s: enabling MoE expert read-ahead (prefetch_experts), %s\n", __func__, + ggml_backend_prefetch_init(params.prefetch_experts_threads) ? "threaded populate engine" : "madvise fallback"); + for (const auto & mapping : model->mappings) { + ggml_backend_prefetch_register_mapping(mapping->addr(), mapping->size()); + } + } if (model->split_mode == LLAMA_SPLIT_MODE_GRAPH && (!model->has_tensor_overrides() || cparams.split_mode_graph_scheduling)) { ggml_backend_sched_set_split_mode_graph(ctx->sched, true, cparams.scheduler_async); ggml_backend_sched_set_max_extra_alloc(ctx->sched, params.max_extra_alloc);