diff --git a/common/arg.cpp b/common/arg.cpp index 8b66d42fa..06d48e4b6 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2234,7 +2234,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } ).set_env("LLAMA_ARG_MOE_SIDECAR")); add_opt(common_arg( - {"--moe-mode"}, "{stock,resident,resident-bank,resident-slot-bank,slot-bank,oracle-all-hit,oracle-prefetch}", + {"--moe-mode"}, "{stock,resident,resident-bank,resident-slot-bank,slot-bank,sweep-prefill,oracle-all-hit,oracle-prefetch}", "Flash-MoE runtime mode", [](common_params & params, const std::string & value) { static const std::set valid = { @@ -2243,6 +2243,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex "resident-bank", "resident-slot-bank", "slot-bank", + "sweep-prefill", "oracle-all-hit", "oracle-prefetch", }; @@ -2252,6 +2253,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.moe_mode = value; } ).set_env("LLAMA_ARG_MOE_MODE")); + add_opt(common_arg( + {"--moe-sweep-min-tokens"}, "N", + string_format("minimum ubatch tokens to engage sweep-prefill expert streaming (default: %d)", params.moe_sweep_min_tokens), + [](common_params & params, int value) { + if (value < 1) { + throw std::invalid_argument("invalid value"); + } + params.moe_sweep_min_tokens = value; + } + ).set_env("LLAMA_ARG_MOE_SWEEP_MIN_TOKENS")); add_opt(common_arg( {"--moe-slot-bank"}, "N", "Flash-MoE slot-bank resident expert capacity per routed MoE layer", diff --git a/common/common.cpp b/common/common.cpp index e5f88209a..7f41b3468 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1357,6 +1357,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.moe_predict_prev_token = params.moe_predict_prev_token; mparams.moe_predict_top1_prev = params.moe_predict_top1_prev; mparams.moe_slot_bank = params.moe_slot_bank; + mparams.moe_sweep_min_tokens = params.moe_sweep_min_tokens; mparams.moe_topk_override = params.moe_topk_override; mparams.moe_cache_io_split = params.moe_cache_io_split; diff --git a/common/common.h b/common/common.h index a578f6352..f2be6e620 100644 --- a/common/common.h +++ b/common/common.h @@ -642,6 +642,7 @@ struct common_params { std::string moe_quant_map = ""; // reserved dynamic-quant policy path // NOLINT std::string oracle_dump = ""; // llama-cli tensor oracle dump directory // NOLINT int32_t moe_slot_bank = 0; // reserved slot-bank size // NOLINT + int32_t moe_sweep_min_tokens = 32; // min ubatch tokens to engage sweep-prefill streaming // NOLINT int32_t moe_topk_override = 0; // runtime reduction-only routed-expert override // NOLINT int32_t moe_cache_io_split = 4; // split each routed expert pread into N aligned chunks // NOLINT int32_t moe_force_expert = -1; // force routed selection to a single expert id // NOLINT diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index ea44e8af9..63dbf591c 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -37,6 +37,10 @@ #if defined(__gnu_linux__) #include #endif +#if defined(__linux__) +#include +#include +#endif #ifdef GGML_USE_OPENMP #include @@ -1506,6 +1510,20 @@ static void * incr_ptr_aligned(void ** p, size_t size, size_t align) { return ptr; } +#if defined(__linux__) +// Flash-MoE sweep-prefill: advise the kernel about the mmap'd slice of one +// expert (page-aligned superset of [data + expert*nb02, +nb02)) +static void ggml_mmid_sweep_advise(const struct ggml_tensor * src0, int64_t expert, size_t nb02, int advice) { + static long page_size = 0; + if (page_size == 0) { + page_size = sysconf(_SC_PAGESIZE); + } + const uintptr_t addr = (uintptr_t) src0->data + (uintptr_t) expert * nb02; + const uintptr_t aligned = addr & ~((uintptr_t) page_size - 1); + madvise((void *) aligned, (size_t) (addr - aligned) + nb02, advice); +} +#endif + static void ggml_compute_forward_mul_mat_id( const struct ggml_compute_params * params, struct ggml_tensor * dst) { @@ -1632,6 +1650,45 @@ static void ggml_compute_forward_mul_mat_id( ggml_barrier(params->threadpool); +#if defined(__linux__) + // Flash-MoE sweep-prefill: eligible mmap-backed expert tensors carry a + // marker in their (otherwise unused) leaf op_params, written at model + // load after verifying the data lies inside the file mapping. The + // expert-major loop below is already a physical-order sweep; thread 0 + // additionally streams the supply: readahead the next used expert slices + // (WILLNEED) and drop slices a few experts behind (DONTNEED). DONTNEED on + // the file-backed mapping is correctness-safe (clean pages refault), the + // lag only avoids refaults for threads still finishing earlier experts. + enum { + MMID_SWEEP_MAGIC = 0x53574550, // 'SWEP' + MMID_SWEEP_AHEAD = 2, + MMID_SWEEP_LAG = 2, + MMID_SWEEP_MAX_EXPERTS = 1024, + }; + const bool sweep_active = ith == 0 && + dst->src[0]->op_params[0] == MMID_SWEEP_MAGIC && + ids->ne[1] >= dst->src[0]->op_params[1] && + n_as <= MMID_SWEEP_MAX_EXPERTS; + + int32_t sweep_used[MMID_SWEEP_MAX_EXPERTS]; + int sweep_n_used = 0; + int sweep_pos = 0; // index into sweep_used tracking the current expert + int sweep_advised = 0; // count of used experts already WILLNEED'd + int64_t sweep_t0 = 0; + + if (sweep_active) { + for (int a = 0; a < n_as; ++a) { + if (matrix_row_counts[a] > 0) { + sweep_used[sweep_n_used++] = a; + } + } + sweep_t0 = ggml_time_us(); + for (; sweep_advised < sweep_n_used && sweep_advised < MMID_SWEEP_AHEAD; ++sweep_advised) { + ggml_mmid_sweep_advise(src0, sweep_used[sweep_advised], nb02, MADV_WILLNEED); + } + } +#endif + for (int cur_a = 0; cur_a < n_as; ++cur_a) { const int64_t cne1 = matrix_row_counts[cur_a]; @@ -1640,6 +1697,21 @@ static void ggml_compute_forward_mul_mat_id( } const char * src0_cur = (const char *) src0->data + cur_a * nb02; + +#if defined(__linux__) + if (sweep_active) { + while (sweep_pos < sweep_n_used && sweep_used[sweep_pos] < cur_a) { + sweep_pos++; + } + while (sweep_advised < sweep_n_used && sweep_advised <= sweep_pos + MMID_SWEEP_AHEAD) { + ggml_mmid_sweep_advise(src0, sweep_used[sweep_advised], nb02, MADV_WILLNEED); + sweep_advised++; + } + if (sweep_pos >= MMID_SWEEP_LAG + 1) { + ggml_mmid_sweep_advise(src0, sweep_used[sweep_pos - MMID_SWEEP_LAG - 1], nb02, MADV_DONTNEED); + } + } +#endif const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; const size_t row_size = ggml_row_size(vec_dot_type, ne10); @@ -1692,6 +1764,22 @@ static void ggml_compute_forward_mul_mat_id( current_chunk = atomic_fetch_add_explicit(current_chunk_ctr, 1, memory_order_relaxed); } } + +#if defined(__linux__) + if (sweep_active) { + static int sweep_stats = -1; + if (sweep_stats == -1) { + const char * val = getenv("GGML_SWEEP_STATS"); + sweep_stats = (val != NULL && val[0] != '\0' && strcmp(val, "0") != 0) ? 1 : 0; + } + if (sweep_stats) { + const double dt = (double) (ggml_time_us() - sweep_t0) / 1e6; + const double gb = (double) sweep_n_used * (double) nb02 / 1e9; + GGML_LOG_INFO("sweep: %s used=%d/%d bytes=%.3fGB wall=%.3fs bw_lb=%.2fGB/s\n", + dst->name, sweep_n_used, n_as, gb, dt, dt > 0.0 ? gb / dt : 0.0); + } + } +#endif } ///////////////////////////////// diff --git a/include/llama.h b/include/llama.h index dc97facf9..572b0866e 100644 --- a/include/llama.h +++ b/include/llama.h @@ -336,6 +336,7 @@ extern "C" { bool moe_predict_top1_prev; // prefetch only the first previous-token same-layer routed expert for the next token int32_t moe_slot_bank; // slot-bank resident expert capacity per routed MoE layer + int32_t moe_sweep_min_tokens; // min ubatch tokens to engage sweep-prefill expert streaming int32_t moe_topk_override; // runtime reduction-only override for routed experts per token (0 = model metadata) int32_t moe_cache_io_split; // split each routed expert pread into N page-aligned chunks (1 = disabled) }; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index bb4e2d968..653006437 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -28,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -471,6 +472,11 @@ struct llama_model::impl { int32_t flash_moe_slot_bank_size = 0; int32_t flash_moe_cache_io_split = 4; int32_t moe_n_expert_used = 0; + bool flash_moe_sweep_prefill_enabled = false; + int32_t flash_moe_sweep_min_tokens = 32; + // routed expert tensors whose data is verified to lie inside a file-backed + // mmap region (safe for madvise WILLNEED/DONTNEED streaming) + std::set flash_moe_sweep_tensors; std::string flash_moe_trace_file; std::unordered_map flash_moe_sidecar_entries; std::vector flash_moe_sparse_mappings; @@ -558,6 +564,8 @@ llama_model::llama_model(const llama_model_params & params) : params(params), pi pimpl->flash_moe_resident_source_enabled || pimpl->flash_moe_oracle_all_hit_enabled || pimpl->flash_moe_oracle_prefetch_enabled; + pimpl->flash_moe_sweep_prefill_enabled = llama_flash_moe_mode_is(params, "sweep-prefill"); + pimpl->flash_moe_sweep_min_tokens = std::max(1, params.moe_sweep_min_tokens); pimpl->flash_moe_trace_file = params.moe_trace_file ? params.moe_trace_file : ""; } @@ -8691,6 +8699,62 @@ const char * llama_model::flash_moe_trace_file() const { return pimpl->flash_moe_trace_file.empty() ? nullptr : pimpl->flash_moe_trace_file.c_str(); } +bool llama_model::flash_moe_sweep_prefill_enabled() const { + return pimpl->flash_moe_sweep_prefill_enabled; +} + +int32_t llama_model::flash_moe_sweep_min_tokens() const { + return pimpl->flash_moe_sweep_min_tokens; +} + +bool llama_model::flash_moe_sweep_tensor(const struct ggml_tensor * t) const { + return pimpl->flash_moe_sweep_tensors.count(t) > 0; +} + +void llama_model::flash_moe_register_sweep_tensors() { + if (!pimpl->flash_moe_sweep_prefill_enabled) { + return; + } + + // note: by the time load_tensors() returns, mapping ownership has moved + // from the loader into pimpl->mappings, so consult those + auto in_mapping = [&](const void * data, size_t size) { + for (const auto & mapping : pimpl->mappings) { + if (!mapping) { + continue; + } + const char * base = (const char *) mapping->addr(); + if ((const char *) data >= base && (const char *) data + size <= base + mapping->size()) { + return true; + } + } + return false; + }; + + size_t registered = 0; + for (const auto & layer : layers) { + for (ggml_tensor * t : { layer.ffn_gate_exps, layer.ffn_up_exps, layer.ffn_down_exps, layer.ffn_gate_up_exps }) { + // only tensors backed by the file mmap are safe to stream with + // madvise; repacked or heap-allocated tensors must stay untouched + if (t != nullptr && t->data != nullptr && in_mapping(t->data, ggml_nbytes(t))) { + // mark the weight leaf itself: op_params is unused on leaf + // tensors, and the CPU mul_mat_id reads this marker to engage + // the sweep (WILLNEED readahead + lagged DONTNEED) + t->op_params[0] = 0x53574550; // 'SWEP' + t->op_params[1] = pimpl->flash_moe_sweep_min_tokens; + pimpl->flash_moe_sweep_tensors.insert(t); + registered++; + } + } + } + + LLAMA_LOG_INFO("%s: sweep-prefill: %zu routed expert tensors eligible for streaming (min ubatch tokens: %d)\n", + __func__, registered, pimpl->flash_moe_sweep_min_tokens); + if (registered == 0) { + LLAMA_LOG_WARN("%s: sweep-prefill requested but no mmap-backed expert tensors found (is --no-mmap set?); falling back to stock behavior\n", __func__); + } +} + const llama_flash_moe_sidecar_entry * llama_model::flash_moe_sidecar_entry_for(const char * name) const { const auto it = pimpl->flash_moe_sidecar_entries.find(name); if (it == pimpl->flash_moe_sidecar_entries.end()) { @@ -9426,6 +9490,7 @@ llama_model_params llama_model_default_params() { /*.moe_predict_prev_token =*/ false, /*.moe_predict_top1_prev =*/ false, /*.moe_slot_bank =*/ 0, + /*.moe_sweep_min_tokens =*/ 32, /*.moe_topk_override =*/ 0, /*.moe_cache_io_split =*/ 4, }; diff --git a/src/llama-model.h b/src/llama-model.h index f593c500a..133240b23 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -596,6 +596,10 @@ struct llama_model { int32_t flash_moe_cache_io_split() const; int32_t moe_n_expert_used() const; const char * flash_moe_trace_file() const; + bool flash_moe_sweep_prefill_enabled() const; + int32_t flash_moe_sweep_min_tokens() const; + bool flash_moe_sweep_tensor(const struct ggml_tensor * t) const; + void flash_moe_register_sweep_tensors(); const llama_flash_moe_sidecar_entry * flash_moe_sidecar_entry_for(const char * name) const; float get_rope_freq_base (const llama_cparams & cparams, int il) const; diff --git a/src/llama.cpp b/src/llama.cpp index 0ea459354..c375cee55 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -897,15 +897,20 @@ static int llama_model_load(struct gguf_context * metadata, llama_model_set_tens const bool use_oracle_all_hit = moe_mode == "oracle-all-hit"; const bool use_oracle_prefetch = moe_mode == "oracle-prefetch"; const bool use_stock_resident = moe_mode == "stock" || moe_mode == "resident"; + const bool use_sweep_prefill = moe_mode == "sweep-prefill"; const bool use_slot_runtime = use_resident_slot_bank || use_slot_bank || use_oracle_all_hit || use_oracle_prefetch; const bool use_flash_moe_sidecar_runtime = use_resident_bank || use_slot_runtime; - if (!use_stock_resident && !use_resident_bank && !use_slot_runtime) { + if (!use_stock_resident && !use_sweep_prefill && !use_resident_bank && !use_slot_runtime) { throw std::runtime_error(format( - "Flash-MoE mode '%s' is not implemented in this build; supported modes are stock, resident, resident-bank, resident-slot-bank, slot-bank, oracle-all-hit, oracle-prefetch", + "Flash-MoE mode '%s' is not implemented in this build; supported modes are stock, resident, resident-bank, resident-slot-bank, slot-bank, sweep-prefill, oracle-all-hit, oracle-prefetch", moe_mode.c_str())); } + if (use_sweep_prefill && !params.use_mmap) { + throw std::runtime_error("Flash-MoE sweep-prefill mode requires mmap (remove --no-mmap)"); + } + if ((use_resident_bank || use_slot_runtime) && (params.moe_sidecar_path == nullptr || params.moe_sidecar_path[0] == '\0')) { throw std::runtime_error(format("Flash-MoE %s mode requires --moe-sidecar", moe_mode.c_str())); } @@ -983,6 +988,8 @@ static int llama_model_load(struct gguf_context * metadata, llama_model_set_tens if (!model.load_tensors(ml)) { return -2; } + + model.flash_moe_register_sweep_tensors(); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error loading model: %s\n", __func__, err.what()); return -1; diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp index b7a6877d4..7c92b444d 100644 --- a/src/models/qwen35moe.cpp +++ b/src/models/qwen35moe.cpp @@ -2,6 +2,8 @@ #include "llama-memory-recurrent.h" +#include + static bool qwen35moe_experimental_metal_split_glu_enabled() { static int enabled = -1; if (enabled == -1) {