mirror of
https://github.com/ikawrakow/ik_llama.cpp.git
synced 2026-07-21 02:05:35 +00:00
Add --prefetch-experts to stream mmap'd MoE experts into page cache (#2101)
* Add --prefetch-experts to stream mmap'd MoE experts into page cache * Drop fds, fault experts in with MADV_POPULATE_READ instead of pread * Remove stale note about pread workers * Move MoE prefetch behind ggml_backend_prefetch_* wrappers * Cleanup stale comments * Add --prefetch-experts-threads, drop GGML_MOE_PREFETCH_THREADS env var
This commit is contained in:
dmaivel
@@ -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);
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
//
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
#include "ggml-alloc.h"
|
||||
#include "ggml-impl.h"
|
||||
#include "ggml-rpc.h"
|
||||
#include "ggml-moe-prefetch.h"
|
||||
|
||||
#include <cassert>
|
||||
#include <climits>
|
||||
@@ -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<uint32_t> 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<ggml_tensor *> host_weights; // originals (host mmap)
|
||||
std::vector<ggml_tensor *> nodes; // MoE nodes computed on a host buffer
|
||||
};
|
||||
std::vector<moe_split_info> 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) {
|
||||
|
||||
@@ -0,0 +1,470 @@
|
||||
#include "ggml-moe-prefetch.h"
|
||||
|
||||
#if defined(__linux__)
|
||||
|
||||
#include <sys/mman.h>
|
||||
#include <unistd.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <atomic>
|
||||
#include <condition_variable>
|
||||
#include <cstdlib>
|
||||
#include <cstdio>
|
||||
#include <deque>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <thread>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
// 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<int> 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<std::pair<uintptr_t, uint32_t>> read_chunks;
|
||||
};
|
||||
|
||||
struct job {
|
||||
uintptr_t addr;
|
||||
uint32_t len;
|
||||
std::shared_ptr<ticket> 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<job> queue;
|
||||
std::vector<std::thread> workers;
|
||||
bool shutdown = false;
|
||||
|
||||
std::unordered_map<const void *, std::shared_ptr<ticket>> tickets;
|
||||
|
||||
// cumulative observability counters (GGML_MOE_PREFETCH_DEBUG)
|
||||
std::atomic<uint64_t> n_jobs{0};
|
||||
std::atomic<uint64_t> n_skipped{0};
|
||||
std::atomic<uint64_t> bytes_populated{0};
|
||||
std::atomic<uint64_t> n_colded{0};
|
||||
|
||||
~prefetch_pool() { stop(); }
|
||||
|
||||
void stop() {
|
||||
{
|
||||
std::lock_guard<std::mutex> 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<std::mutex> 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<std::mutex> 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<std::mutex> 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<mapping_entry> mappings;
|
||||
|
||||
std::atomic<uint64_t> epoch{1};
|
||||
|
||||
std::mutex pool_mtx;
|
||||
std::shared_ptr<prefetch_pool> 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<std::mutex> 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<std::pair<size_t, size_t>> & 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<uint32_t> 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>((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<std::pair<size_t, size_t>> & ranges, bool urgent, bool track_reads) {
|
||||
auto & s = state();
|
||||
std::shared_ptr<prefetch_pool> pool_sp;
|
||||
{
|
||||
std::lock_guard<std::mutex> 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<std::mutex> lock(pool.mtx);
|
||||
if (pool.shutdown) return false;
|
||||
|
||||
auto & tk = pool.tickets[w];
|
||||
if (!tk) tk = std::make_shared<ticket>();
|
||||
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<job> 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<std::pair<size_t, size_t>> 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<std::mutex> 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<std::mutex> 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<std::mutex> 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<prefetch_pool>();
|
||||
s.pool->start(n_threads);
|
||||
}
|
||||
|
||||
bool ggml_moe_prefetch_enabled(void) {
|
||||
auto & s = state();
|
||||
std::lock_guard<std::mutex> 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<std::pair<size_t, size_t>> 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<std::pair<size_t, size_t>> 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<prefetch_pool> pool;
|
||||
{
|
||||
std::lock_guard<std::mutex> plock(s.pool_mtx);
|
||||
pool = s.pool;
|
||||
}
|
||||
if (!pool) return;
|
||||
std::vector<std::pair<uintptr_t, uint32_t>> chunks;
|
||||
{
|
||||
std::lock_guard<std::mutex> 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<prefetch_pool> pool;
|
||||
{
|
||||
std::lock_guard<std::mutex> plock(s.pool_mtx);
|
||||
pool = s.pool;
|
||||
}
|
||||
if (!pool) return;
|
||||
std::shared_ptr<ticket> tk;
|
||||
{
|
||||
std::lock_guard<std::mutex> 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<std::mutex> 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
|
||||
@@ -0,0 +1,52 @@
|
||||
#pragma once
|
||||
|
||||
#include "ggml.h"
|
||||
|
||||
#include <stddef.h>
|
||||
#include <stdbool.h>
|
||||
|
||||
#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
|
||||
@@ -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));
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
#include "llama-impl.h"
|
||||
|
||||
#include "ggml.h"
|
||||
#include "ggml-backend.h"
|
||||
|
||||
#include <cstring>
|
||||
#include <climits>
|
||||
@@ -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));
|
||||
|
||||
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user