From 4260e460832f5e95c41d8b6d8b2aea98007df524 Mon Sep 17 00:00:00 2001 From: Junchao Lyu Date: Fri, 3 Jul 2026 12:13:46 -0700 Subject: [PATCH 1/2] llama : stream MoE routed experts from disk (core + async pool + O_DIRECT) Run MoE models larger than RAM: routed expert weights (ffn_*_exps) are not materialized; each streamed layer keeps a small device-side cache of expert slots, filled on demand from the GGUF by a CPU id-remap op after the router top-k. Missing experts load via a pread thread pool while the op waits; eviction is decaying route hotness with an LRU tiebreak. Output is byte-identical to the unstreamed model. Options: --moe-stream, --moe-stream-cache , --moe-stream-io-threads N, and --moe-stream-direct (O_DIRECT expert reads, bypassing the page cache; falls back to buffered when the OS/filesystem does not support it, verified by a probe read at open time). Assisted-by: Claude --- common/arg.cpp | 44 +++ common/common.cpp | 13 + common/common.h | 6 + common/sampling.cpp | 2 + include/llama.h | 11 + src/CMakeLists.txt | 1 + src/llama-adapter.cpp | 3 + src/llama-context.cpp | 28 ++ src/llama-graph.cpp | 33 ++- src/llama-graph.h | 11 +- src/llama-model-loader.cpp | 15 +- src/llama-model-loader.h | 8 + src/llama-model.cpp | 150 +++++++++- src/llama-model.h | 6 + src/llama-moe-stream.cpp | 556 +++++++++++++++++++++++++++++++++++++ src/llama-moe-stream.h | 161 +++++++++++ src/llama.cpp | 7 + 17 files changed, 1043 insertions(+), 12 deletions(-) create mode 100644 src/llama-moe-stream.cpp create mode 100644 src/llama-moe-stream.h diff --git a/common/arg.cpp b/common/arg.cpp index fdf58614b67b..196807d48ef3 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2477,6 +2477,50 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } } ).set_env("LLAMA_ARG_N_CPU_MOE")); + add_opt(common_arg( + {"--moe-stream"}, + "stream Mixture of Experts (MoE) routed expert weights from disk on demand", + [](common_params & params) { + params.moe_stream = true; + } + ).set_env("LLAMA_ARG_MOE_STREAM")); + add_opt(common_arg( + {"--moe-stream-cache"}, "", + "expert cache for --moe-stream: memory budget in GiB (e.g. 40) or exact slots per layer with an 's' suffix (e.g. 64s); implies --moe-stream (default: auto)", + [](common_params & params, const std::string & value) { + params.moe_stream = true; + size_t pos = 0; + const uint64_t n = std::stoull(value, &pos); + std::string suffix = value.substr(pos); + for (auto & c : suffix) { + c = std::tolower(c); + } + if (suffix == "s" || suffix == "slot" || suffix == "slots") { + params.moe_stream_slots = n; + } else if (suffix.empty() || suffix == "g" || suffix == "gb" || suffix == "gib") { + params.moe_stream_budget = n * 1024ull * 1024ull * 1024ull; + } else { + throw std::invalid_argument("invalid value"); + } + } + ).set_env("LLAMA_ARG_MOE_STREAM_CACHE")); + add_opt(common_arg( + {"--moe-stream-io-threads"}, "N", + "I/O threads for --moe-stream expert loads; implies --moe-stream (default: auto)", + [](common_params & params, int value) { + params.moe_stream = true; + params.moe_stream_io_threads = value; + } + ).set_env("LLAMA_ARG_MOE_STREAM_IO_THREADS")); + add_opt(common_arg( + {"--moe-stream-direct"}, + "use O_DIRECT for --moe-stream expert reads (bypass the page cache); implies --moe-stream. " + "falls back to buffered reads if O_DIRECT is unsupported by the OS or filesystem", + [](common_params & params) { + params.moe_stream = true; + params.moe_stream_direct = true; + } + ).set_env("LLAMA_ARG_MOE_STREAM_DIRECT")); GGML_ASSERT(params.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0 add_opt(common_arg( {"-ngl", "--gpu-layers", "--n-gpu-layers"}, "N", diff --git a/common/common.cpp b/common/common.cpp index 0dd9ede5eb6e..55e6a5d74367 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1398,6 +1398,13 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode common_set_adapter_lora(lctx, params.lora_adapters); } + if (params.warmup && params.moe_stream) { + // the warmup graph routes every token through all experts at once, which cannot fit the + // streaming expert cache + COM_TRC("%s", "skipping warmup: not supported with MoE expert streaming\n"); + params.warmup = false; + } + if (params.warmup) { COM_TRC("%s", "warming up the model with an empty run - please wait ... (--no-warmup to disable)\n"); @@ -1545,6 +1552,12 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.use_extra_bufts = !params.no_extra_bufts; mparams.no_host = params.no_host; + mparams.moe_stream = params.moe_stream; + mparams.moe_stream_slots = params.moe_stream_slots; + mparams.moe_stream_budget = params.moe_stream_budget; + mparams.moe_stream_io_threads = params.moe_stream_io_threads; + mparams.moe_stream_direct = params.moe_stream_direct; + if (params.kv_overrides.empty()) { mparams.kv_overrides = NULL; } else { diff --git a/common/common.h b/common/common.h index 2adb310b83fe..8abd90574681 100644 --- a/common/common.h +++ b/common/common.h @@ -582,6 +582,12 @@ struct common_params { bool no_extra_bufts = false; // disable extra buffer types (used for weight repacking) bool no_host = false; // bypass host buffer allowing extra buffers to be used + bool moe_stream = false; // stream MoE routed expert weights from disk on demand + uint32_t moe_stream_slots = 0; // expert cache slots per streamed layer (0 = auto) + uint64_t moe_stream_budget = 0; // total expert cache byte budget, used when slots == 0 (0 = auto) + int32_t moe_stream_io_threads = 0; // expert load I/O threads (<= 0 = default) + bool moe_stream_direct = false; // use O_DIRECT for expert reads (bypass page cache) + bool single_turn = false; // single turn chat conversation ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K diff --git a/common/sampling.cpp b/common/sampling.cpp index 75a299e23ece..4968882a3f73 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -525,6 +525,8 @@ void common_perf_print(const struct llama_context * ctx, const struct common_sam LOG_INF("%s: unaccounted time = %10.2f ms / %5.1f %% (total - sampling - prompt eval - eval) / (total)\n", __func__, t_unacc_ms, t_unacc_pc); LOG_INF("%s: graphs reused = %10d\n", __func__, data.n_reused); + llama_moe_stream_print_stats(llama_get_model(ctx)); + common_memory_breakdown_print(ctx); } } diff --git a/include/llama.h b/include/llama.h index f723c9f60cfe..dc84713056e1 100644 --- a/include/llama.h +++ b/include/llama.h @@ -315,6 +315,13 @@ extern "C" { // override key-value pairs of the model meta data const struct llama_model_kv_override * kv_overrides; + // SSD streaming of MoE routed expert weights (experts are paged from the GGUF on demand + // into a per-layer cache of moe_stream_slots experts; requires moe_stream = true) + uint32_t moe_stream_slots; // expert cache slots per streamed layer (0 = auto) + uint64_t moe_stream_budget; // total cache byte budget, used when slots == 0 (0 = auto heuristic) + int32_t moe_stream_io_threads; // expert load I/O threads (<= 0 = default) + bool moe_stream_direct; // use O_DIRECT for expert reads (bypass page cache); falls back if unsupported + // Keep the booleans together to avoid misalignment during copy-by-value. bool vocab_only; // only load the vocabulary, no weights bool use_mmap; // use mmap if possible @@ -324,6 +331,7 @@ extern "C" { bool use_extra_bufts; // use extra buffer types (used for weight repacking) bool no_host; // bypass host buffer allowing extra buffers to be used bool no_alloc; // only load metadata and simulate memory allocations + bool moe_stream; // stream MoE routed expert weights from disk on demand }; struct llama_sampler_seq_config { @@ -1549,6 +1557,9 @@ extern "C" { LLAMA_API void llama_perf_sampler_print(const struct llama_sampler * chain); LLAMA_API void llama_perf_sampler_reset( struct llama_sampler * chain); + // print MoE expert streaming statistics (no-op when streaming is not enabled) + LLAMA_API void llama_moe_stream_print_stats(const struct llama_model * model); + // // training // diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 320784c3a8cc..2aaa8c39018f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -34,6 +34,7 @@ add_library(llama llama-model-loader.cpp llama-model-saver.cpp llama-model.cpp + llama-moe-stream.cpp llama-quant.cpp llama-sampler.cpp llama-vocab.cpp diff --git a/src/llama-adapter.cpp b/src/llama-adapter.cpp index 3e0fe66afff7..b679f7d7c182 100644 --- a/src/llama-adapter.cpp +++ b/src/llama-adapter.cpp @@ -329,6 +329,9 @@ static void llama_adapter_lora_init_impl(llama_model & model, const char * path_ // device buft and device ctx const auto * model_tensor = model.get_tensor(name.c_str()); if (!model_tensor) { + if (model.moe_stream() && name.find("_exps.") != std::string::npos) { + throw std::runtime_error("LoRA tensor '" + name + "' targets an SSD-streamed expert tensor, which is not supported"); + } throw std::runtime_error("LoRA tensor '" + name + "' does not exist in base model (hint: maybe wrong base model?)"); } diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0465430df43a..07c0f58f5393 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -9,6 +9,7 @@ #include "llama-memory.h" #include "llama-mmap.h" #include "llama-model.h" +#include "llama-moe-stream.h" #include "llama-ext.h" #include "llama.h" @@ -211,6 +212,28 @@ llama_context::llama_context( cparams.op_offload = params.op_offload; cparams.kv_unified = params.kv_unified; + if (model.moe_stream() && hparams.n_expert_used > 0) { + // one mul_mat_id needs every expert a ubatch touches resident at once, so a ubatch may + // not select more experts than the cache holds (worst case n_ubatch*n_expert_used) + const uint32_t n_ubatch_max = std::max(1u, model.moe_stream()->n_slots / hparams.n_expert_used); + if (cparams.n_ubatch > n_ubatch_max) { + LLAMA_LOG_WARN("%s: n_ubatch reduced from %u to %u so that a ubatch cannot select more experts than the %u-slot streaming cache\n", + __func__, cparams.n_ubatch, n_ubatch_max, model.moe_stream()->n_slots); + cparams.n_ubatch = n_ubatch_max; + } + + // op offload snapshots host weights to the device per graph split, which assumes they do + // not change during the graph - streamed caches are rewritten on demand + bool cache_on_host = false; + for (const auto & buf : model.moe_stream()->bufs) { + cache_on_host = cache_on_host || ggml_backend_buffer_is_host(buf.get()); + } + if (cache_on_host && cparams.op_offload) { + LLAMA_LOG_WARN("%s: disabling op offload: the expert streaming cache is in host memory\n", __func__); + cparams.op_offload = false; + } + } + // initialized later cparams.pipeline_parallel = false; @@ -2415,6 +2438,7 @@ llm_graph_params llama_context::graph_params( /*.loras =*/ loras.get(), /*.mctx =*/ mctx, /*.cross =*/ &cross, + /*.mstream =*/ model.moe_stream(), /*.samplers =*/ sampling.samplers, /*.n_outputs =*/ n_outputs, /*.cb =*/ graph_get_cb(), @@ -4094,6 +4118,10 @@ void llama_perf_context_print(const llama_context * ctx) { __func__, data.t_eval_ms, data.n_eval, data.t_eval_ms / data.n_eval, 1e3 / data.t_eval_ms * data.n_eval); LLAMA_LOG_INFO("%s: total time = %10.2f ms / %5d tokens\n", __func__, (t_end_ms - data.t_start_ms), (data.n_p_eval + data.n_eval)); LLAMA_LOG_INFO("%s: graphs reused = %10d\n", __func__, data.n_reused); + + if (const auto * mstream = ctx->get_model().moe_stream()) { + mstream->print_stats(); + } } void llama_perf_context_reset(llama_context * ctx) { diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 4c86e43c1f74..20487b5848f5 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -4,6 +4,7 @@ #include "llama-model.h" #include "llama-batch.h" #include "llama-cparams.h" +#include "llama-moe-stream.h" #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" @@ -1350,6 +1351,7 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) : loras (params.loras), mctx (params.mctx), cross (params.cross), + mstream (params.mstream), samplers (params.samplers), cb_func (params.cb), res (params.res), @@ -1405,15 +1407,20 @@ ggml_tensor * llm_graph_context::build_lora_mm_id( ggml_tensor * w, // ggml_tensor * as ggml_tensor * cur, // ggml_tensor * b ggml_tensor * ids, - ggml_tensor * w_s) const { + ggml_tensor * w_s, + ggml_tensor * ids_scale) const { ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids); if (w_s) { + // w_s always covers all experts, so index it with the original expert ids + // even when the GEMM ids are remapped cache slots (MoE streaming) + ggml_tensor * ids_s = ids_scale ? ids_scale : ids; + const int64_t n_expert = w_s->ne[0]; const int64_t n_tokens = cur->ne[2]; ggml_tensor * s = ggml_reshape_3d(ctx0, w_s, 1, n_expert, 1); s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); - s = ggml_get_rows(ctx0, s, ids); + s = ggml_get_rows(ctx0, s, ids_s); res = ggml_mul(ctx0, res, s); } for (const auto & lora : *loras) { @@ -1951,6 +1958,20 @@ ggml_tensor * llm_graph_context::build_moe_ffn( //call early so that topk-moe can be used ggml_build_forward_expand(gf, weights); + // MoE expert streaming: make the selected experts cache-resident and remap the ids fed to the + // expert GEMMs to cache slots; the routing itself (weights, scales, biases) keeps the original ids + llama_moe_stream_layer * msl = mstream ? mstream->layer(il) : nullptr; + if (msl && !msl->matches(gate_exps, up_exps, down_exps, gate_up_exps)) { + msl = nullptr; // a different expert group of the same layer (e.g. grovemoe chexps), not streamed + } + + ggml_tensor * ids_gemm = selected_experts; + if (msl) { + ggml_tensor * ids_cont = ggml_cont(ctx0, selected_experts); // top_k output is a view + ids_gemm = ggml_map_custom1(ctx0, ids_cont, llama_moe_stream_remap, 1, msl); + cb(ids_gemm, "ffn_moe_topk_stream", il); + } + cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens); if (weight_before_ffn) { @@ -1965,7 +1986,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn( if (gate_up_exps) { // merged gate_up path: one mul_mat_id, then split into gate and up views - ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s); // [n_ff*2, n_expert_used, n_tokens] + ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, ids_gemm, up_exps_s, selected_experts); // [n_ff*2, n_expert_used, n_tokens] cb(gate_up, "ffn_moe_gate_up", il); if (up_exps_s) { @@ -1984,7 +2005,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cb(up, "ffn_moe_up", il); } else { // separate gate and up path - up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s); // [n_ff, n_expert_used, n_tokens] + up = build_lora_mm_id(up_exps, cur, ids_gemm, up_exps_s, selected_experts); // [n_ff, n_expert_used, n_tokens] cb(up, "ffn_moe_up", il); if (up_exps_s) { @@ -1997,7 +2018,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn( } if (gate_exps) { - cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s); // [n_ff, n_expert_used, n_tokens] + cur = build_lora_mm_id(gate_exps, cur, ids_gemm, gate_exps_s, selected_experts); // [n_ff, n_expert_used, n_tokens] cb(cur, "ffn_moe_gate", il); } else { cur = up; @@ -2086,7 +2107,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn( GGML_ABORT("fatal error"); } - experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s); // [n_embd, n_expert_used, n_tokens] + experts = build_lora_mm_id(down_exps, cur, ids_gemm, down_exps_s, selected_experts); // [n_embd, n_expert_used, n_tokens] cb(experts, "ffn_moe_down", il); if (down_exps_s) { diff --git a/src/llama-graph.h b/src/llama-graph.h index 4b5b75c632ab..5d31375bc22c 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -660,6 +660,8 @@ using llm_graph_cb = std::function samplers; static bool samplers_equal( @@ -907,6 +912,8 @@ struct llm_graph_context { const llama_memory_context_i * mctx; const llama_cross * cross; + llama_moe_stream * mstream; + std::map samplers; const llm_graph_cb & cb_func; @@ -936,11 +943,13 @@ struct llm_graph_context { ggml_tensor * w_s = nullptr) const; // do mat_mul_id, while optionally apply lora and per-expert scale + // ids_scale: ids to use for the w_s gather, when the GEMM ids are remapped cache slots (MoE streaming) ggml_tensor * build_lora_mm_id( ggml_tensor * w, // ggml_tensor * as ggml_tensor * cur, // ggml_tensor * b ggml_tensor * ids, - ggml_tensor * w_s = nullptr) const; + ggml_tensor * w_s = nullptr, + ggml_tensor * ids_scale = nullptr) const; ggml_tensor * build_norm( ggml_tensor * cur, diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index e07b0d231347..3f156aa21f0a 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -557,6 +557,7 @@ llama_model_loader::llama_model_loader( llm_kv = LLM_KV(llm_arch_from_string(arch_name)); files.emplace_back(new llama_file(fname.c_str(), "rb", use_direct_io)); + file_paths.emplace_back(fname); contexts.emplace_back(ctx); if (use_mmap && use_direct_io) { @@ -639,6 +640,7 @@ llama_model_loader::llama_model_loader( } files.emplace_back(new llama_file(fname_split, "rb", use_direct_io)); + file_paths.emplace_back(fname_split); contexts.emplace_back(ctx); // Save tensors data offset info of the shard. @@ -683,6 +685,7 @@ llama_model_loader::llama_model_loader( llm_kv = LLM_KV(llm_arch_from_string(arch_name)); files.emplace_back(new llama_file(file)); + file_paths.emplace_back(); contexts.emplace_back(ctx); // Save tensors data offset info of the main file. @@ -895,7 +898,7 @@ const struct ggml_tensor * llama_model_loader::check_tensor_dims(const std::stri } // checks if the weight tensor can be used with the specified buffer type and device -static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w, ggml_op op, ggml_backend_buffer_type_t buft, ggml_backend_dev_t dev) { +bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w, ggml_op op, ggml_backend_buffer_type_t buft, ggml_backend_dev_t dev) { GGML_ASSERT(w != nullptr); if (op == GGML_OP_NONE) { @@ -1034,7 +1037,7 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w } // find the first buffer type in the list that can use the tensor -static ggml_backend_buffer_type_t select_weight_buft(const llama_hparams & hparams, ggml_tensor * tensor, ggml_op op, const buft_list_t * buft_list) { +ggml_backend_buffer_type_t select_weight_buft(const llama_hparams & hparams, ggml_tensor * tensor, ggml_op op, const buft_list_t * buft_list) { GGML_ASSERT(!buft_list->empty()); for (const auto & cur : *buft_list) { ggml_backend_dev_t cur_dev = cur.first; @@ -1104,9 +1107,13 @@ struct ggml_tensor * llama_model_loader::create_tensor( } // skip unused tensors - if (info.op == GGML_OP_NONE || (flags & TENSOR_SKIP)) { + if (info.op == GGML_OP_NONE || (flags & (TENSOR_SKIP | TENSOR_STREAMED))) { const size_t nbytes = ggml_nbytes(t_meta); - LLAMA_LOG_WARN("model has unused tensor %s (size = %zu bytes) -- ignoring\n", tn.str().c_str(), nbytes); + if (flags & TENSOR_STREAMED) { + LLAMA_LOG_DEBUG("tensor %s is SSD-streamed (size = %zu bytes)\n", tn.str().c_str(), nbytes); + } else { + LLAMA_LOG_WARN("model has unused tensor %s (size = %zu bytes) -- ignoring\n", tn.str().c_str(), nbytes); + } size_data -= nbytes; n_created++; diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index c476026d3e51..8493036b632e 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -28,6 +28,12 @@ enum llama_fver { const char * llama_file_version_name(llama_fver version); +// check if the buffer type of the given device supports op with the tensor as its weight +bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w, ggml_op op, ggml_backend_buffer_type_t buft, ggml_backend_dev_t dev); + +// find the first buffer type in the list that can use the tensor +ggml_backend_buffer_type_t select_weight_buft(const llama_hparams & hparams, ggml_tensor * tensor, ggml_op op, const buft_list_t * buft_list); + struct llama_model_loader { // Holds information on a model weight struct llama_tensor_weight { @@ -67,6 +73,7 @@ struct llama_model_loader { static const int TENSOR_DUPLICATED = 1 << 1; static const int TENSOR_SKIP = 1 << 2; static const int TENSOR_SKIP_IF_VIRTUAL = 1 << 3; + static const int TENSOR_STREAMED = 1 << 4; int n_kv = 0; int n_tensors = 0; @@ -81,6 +88,7 @@ struct llama_model_loader { bool no_alloc; llama_files files; + std::vector file_paths; // same order as files; empty string for FILE*-based loading llama_ftype ftype; llama_fver fver; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d58ebac28b9b..62819d7cbba0 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -7,6 +7,7 @@ #include "llama-mmap.h" #include "llama-cparams.h" #include "llama-model-loader.h" +#include "llama-moe-stream.h" #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" @@ -1010,6 +1011,9 @@ struct llama_model::impl { std::vector dev_layer; bool has_tensor_overrides; + + // MoE expert SSD streaming state, null when not enabled + std::unique_ptr moe_stream; }; llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique()) { @@ -1213,6 +1217,79 @@ void llama_model_base::load_vocab(llama_model_loader & ml) { vocab.load(ml, kv); } +static bool llama_moe_stream_is_exps(llm_tensor tensor) { + switch (tensor) { + case LLM_TENSOR_FFN_GATE_EXPS: + case LLM_TENSOR_FFN_UP_EXPS: + case LLM_TENSOR_FFN_DOWN_EXPS: + case LLM_TENSOR_FFN_GATE_UP_EXPS: + return true; + default: + return false; + } +} + +// like select_weight_buft, but only consider each device's default buffer type: extra buffer +// types (e.g. CPU weight repacking) may not support the partial writes that slot-wise expert +// cache updates require +static ggml_backend_buffer_type_t llama_moe_stream_select_buft(const llama_hparams & hparams, ggml_tensor * meta, const buft_list_t * buft_list) { + for (const auto & [dev, buft] : *buft_list) { + if (buft != ggml_backend_dev_buffer_type(dev)) { + continue; + } + if (weight_buft_supported(hparams, meta, GGML_OP_MUL_MAT_ID, buft, dev)) { + return buft; + } + } + return select_weight_buft(hparams, meta, GGML_OP_MUL_MAT_ID, buft_list); +} + +static bool llama_moe_stream_is_exps_name(const std::string & name) { + for (const char * suffix : { ".ffn_gate_exps.weight", ".ffn_up_exps.weight", ".ffn_down_exps.weight", ".ffn_gate_up_exps.weight" }) { + const size_t len = strlen(suffix); + if (name.size() >= len && name.compare(name.size() - len, len, suffix) == 0) { + return true; + } + } + return false; +} + +// resolve the per-layer expert cache slot count; returns 0 when streaming should not be enabled +// (not a MoE model, or the cache would hold every expert anyway) +static uint32_t llama_moe_stream_resolve_slots(const llama_model_params & params, const llama_hparams & hparams, const llama_model_loader & ml) { + if (hparams.n_expert == 0 || hparams.n_expert_used == 0) { + LLAMA_LOG_WARN("%s: MoE expert streaming requires a MoE model -- disabled\n", __func__); + return 0; + } + + uint32_t n_slots = params.moe_stream_slots; + + if (n_slots == 0 && params.moe_stream_budget > 0) { + // derive the per-layer slot count from the total byte budget + size_t nb_expert_sum = 0; // bytes of one expert across every streamable layer + for (const auto & [name, w] : ml.weights_map) { + if (llama_moe_stream_is_exps_name(name) && w.tensor->ne[2] == hparams.n_expert) { + nb_expert_sum += ggml_nbytes(w.tensor)/w.tensor->ne[2]; + } + } + if (nb_expert_sum > 0) { + n_slots = std::max(params.moe_stream_budget/nb_expert_sum, hparams.n_expert_used); + } + } + + if (n_slots == 0) { + n_slots = std::clamp(2*hparams.n_expert_used, 16, hparams.n_expert); + } + + if (n_slots >= hparams.n_expert) { + LLAMA_LOG_WARN("%s: MoE expert cache of %u slots covers all %u experts -- streaming disabled, loading normally\n", + __func__, n_slots, hparams.n_expert); + return 0; + } + + return n_slots; +} + bool llama_model_base::load_tensors(llama_model_loader & ml) { const auto & split_mode = params.split_mode; const auto & use_mlock = params.use_mlock; @@ -1302,6 +1379,18 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // assign the output layer pimpl->dev_output = get_layer_buft_list(n_layer_all); + if (params.moe_stream) { + const uint32_t n_slots = llama_moe_stream_resolve_slots(params, hparams, ml); + if (n_slots > 0) { + if (pimpl->has_tensor_overrides) { + LLAMA_LOG_WARN("%s: tensor buffer overrides (-ot/--cpu-moe) do not apply to SSD-streamed expert tensors\n", __func__); + } + pimpl->moe_stream = std::make_unique(n_layer_all, n_slots, params.moe_stream_io_threads, params.moe_stream_direct); + LLAMA_LOG_INFO("%s: MoE expert SSD streaming enabled, %u of %u experts cached per layer, %d I/O threads\n", + __func__, n_slots, hparams.n_expert, pimpl->moe_stream->n_io_threads); + } + } + const auto TENSOR_NOT_REQUIRED = llama_model_loader::TENSOR_NOT_REQUIRED; // create tensors for the weights @@ -1583,6 +1672,18 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { ctx_buf_maps.emplace_back(ctx, buf_map); } + if (pimpl->moe_stream) { + if (pimpl->moe_stream->ctxs.empty()) { + LLAMA_LOG_WARN("%s: no streamable expert tensors found -- MoE expert streaming disabled\n", __func__); + pimpl->moe_stream.reset(); + } else { + pimpl->moe_stream->alloc_bufs(ml.no_alloc); + if (!ml.no_alloc) { + pimpl->moe_stream->open_files(ml.file_paths); + } + } + } + if (llama_supports_gpu_offload()) { const int n_gpu = std::min(n_gpu_layers, n_layer_all); @@ -1629,6 +1730,37 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { ggml_tensor * llama_model_base::create_tensor(llama_model_loader & ml, const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags) { const buft_list_t * buft_list_layer = tn.bid == -1 ? nullptr : pimpl->dev_layer.at(tn.bid).buft_list; + + // route MoE routed-expert weights to the streaming expert cache instead of materializing them + if (pimpl->moe_stream && tn.bid >= 0 && buft_list_layer != nullptr && + (flags & (TENSOR_DUPLICATED | TENSOR_SKIP | TENSOR_SKIP_IF_VIRTUAL)) == 0 && + llama_moe_stream_is_exps(tn.tensor) && tn.suffix != nullptr && strcmp(tn.suffix, "weight") == 0) { + const std::string name = tn.str(); + const auto * w = ml.get_weight(name.c_str()); + + bool dims_ok = w != nullptr && w->tensor->ne[2] == (int64_t) hparams.n_expert; + if (dims_ok) { + size_t dim = 0; + for (int64_t d : ne) { + dims_ok = dims_ok && d == w->tensor->ne[dim++]; + } + } + + if (dims_ok) { + // register the skip so that the loader tensor accounting stays consistent; returns nullptr + ml.create_tensor( + hparams, &pimpl->cpu_buft_list, pimpl->dev_input.buft_list, pimpl->dev_output.buft_list, buft_list_layer, + tn, ne, flags | TENSOR_STREAMED); + + ggml_backend_buffer_type_t buft = llama_moe_stream_select_buft(hparams, w->tensor, buft_list_layer); + if (buft == nullptr) { + throw std::runtime_error(format("failed to find a buffer type for streamed tensor %s", name.c_str())); + } + + return pimpl->moe_stream->create_cache_tensor(tn.bid, buft, w->tensor, w->idx, w->offs); + } + } + return ml.create_tensor( hparams, &pimpl->cpu_buft_list, pimpl->dev_input.buft_list, pimpl->dev_output.buft_list, buft_list_layer, tn, ne, flags); @@ -1972,6 +2104,16 @@ bool llama_model::has_tensor_overrides() const { return pimpl->has_tensor_overrides; } +llama_moe_stream * llama_model::moe_stream() const { + return pimpl->moe_stream.get(); +} + +void llama_moe_stream_print_stats(const llama_model * model) { + if (model && model->moe_stream()) { + model->moe_stream()->print_stats(); + } +} + const ggml_tensor * llama_model::get_tensor(const char * name) const { auto it = std::find_if(tensors_by_name.begin(), tensors_by_name.end(), [name](const std::pair & it) { @@ -2291,6 +2433,10 @@ llama_model_params llama_model_default_params() { /*.progress_callback =*/ nullptr, /*.progress_callback_user_data =*/ nullptr, /*.kv_overrides =*/ nullptr, + /*.moe_stream_slots =*/ 0, + /*.moe_stream_budget =*/ 0, + /*.moe_stream_io_threads =*/ 0, + /*.moe_stream_direct =*/ false, /*.vocab_only =*/ false, /*.use_mmap =*/ true, /*.use_direct_io =*/ false, @@ -2299,6 +2445,7 @@ llama_model_params llama_model_default_params() { /*.use_extra_bufts =*/ true, /*.no_host =*/ false, /*.no_alloc =*/ false, + /*.moe_stream =*/ false, }; return result; @@ -2703,7 +2850,8 @@ llama_model_base::llama_model_base(const struct llama_model_params & params) : l TENSOR_DUPLICATED (llama_model_loader::TENSOR_DUPLICATED), TENSOR_NOT_REQUIRED (llama_model_loader::TENSOR_NOT_REQUIRED), TENSOR_SKIP (llama_model_loader::TENSOR_SKIP), - TENSOR_SKIP_IF_VIRTUAL(llama_model_loader::TENSOR_SKIP_IF_VIRTUAL) {} + TENSOR_SKIP_IF_VIRTUAL(llama_model_loader::TENSOR_SKIP_IF_VIRTUAL), + TENSOR_STREAMED (llama_model_loader::TENSOR_STREAMED) {} ggml_tensor * llama_model_base::create_tensor(const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags) { GGML_ASSERT(ml != nullptr); diff --git a/src/llama-model.h b/src/llama-model.h index 4800d2928c52..a6a20dc68a7a 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -17,6 +17,7 @@ struct llama_cparams; struct llama_ubatch; struct llama_model_loader; +struct llama_moe_stream; // available models enum llm_type { @@ -659,6 +660,10 @@ struct llama_model { bool has_tensor_overrides() const; + // MoE expert SSD streaming state, null when not enabled + // the pointee is mutable (residency changes during decode), only the pointer is owned here + llama_moe_stream * moe_stream() const; + const struct ggml_tensor * get_tensor(const char * name) const; float get_rope_freq_base (const llama_cparams & cparams, int il) const; @@ -703,6 +708,7 @@ struct llama_model_base : public llama_model { const int TENSOR_NOT_REQUIRED; const int TENSOR_SKIP; const int TENSOR_SKIP_IF_VIRTUAL; + const int TENSOR_STREAMED; explicit llama_model_base(const llama_model_params & params); virtual ~llama_model_base() = default; diff --git a/src/llama-moe-stream.cpp b/src/llama-moe-stream.cpp new file mode 100644 index 000000000000..19c755730f2f --- /dev/null +++ b/src/llama-moe-stream.cpp @@ -0,0 +1,556 @@ +#include "llama-moe-stream.h" + +#include "llama-impl.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include + +#ifdef _WIN32 +#include +#else +#include +#include +#endif + +static const uint32_t MOE_STREAM_IO_THREADS_DEFAULT = 9; +static const uint32_t MOE_STREAM_IO_THREADS_MAX = 18; +static const int64_t MOE_STREAM_HOT_DECAY_TOKENS = 64; + +// O_DIRECT alignment: 4096 is a multiple of any device logical block size (512/4096), so it is +// universally valid, and reading a few extra KB of head/tail padding per slab is negligible +static const size_t MOE_STREAM_DIRECT_ALIGN = 4096; + +// saturating increment - route-hotness counters accumulate over a whole run and must not wrap +static uint32_t sat_inc(uint32_t & c) { + if (c < UINT32_MAX - 1) { + c++; + } + return c; +} + +// page-aligned allocation, required both for O_DIRECT reads and for Metal private-buffer uploads +static void * moe_aligned_alloc(size_t n) { +#ifdef _WIN32 + return _aligned_malloc(n, MOE_STREAM_DIRECT_ALIGN); +#else + void * p = nullptr; + if (posix_memalign(&p, MOE_STREAM_DIRECT_ALIGN, n) != 0) { + p = nullptr; + } + return p; +#endif +} + +static void moe_aligned_free(void * p) { +#ifdef _WIN32 + _aligned_free(p); +#else + free(p); +#endif +} + +// read len bytes at file offset offs into staging (thread-safe positional read); staging must have +// room for len (+ 2*MOE_STREAM_DIRECT_ALIGN when direct). returns a pointer to the len bytes +// within staging, or nullptr on failure +static const uint8_t * llama_moe_stream_pread(llama_file & file, uint8_t * staging, size_t len, size_t offs, bool direct) { +#ifdef _WIN32 + GGML_UNUSED(direct); + // no positional read primitive; serialize the seek+read pairs + static std::mutex io_mtx; + std::lock_guard lock(io_mtx); + try { + file.seek(offs, SEEK_SET); + file.read_raw(staging, len); + return staging; + } catch (...) { + return nullptr; + } +#else + const int fd = file.file_id(); + + if (direct) { + // O_DIRECT requires the offset, length, and buffer all block-aligned + const size_t a = MOE_STREAM_DIRECT_ALIGN; + const size_t aoffs = offs & ~(a - 1); + const size_t head = offs - aoffs; + const size_t total = ((head + len + a - 1)/a)*a; + ssize_t r; + do { + r = pread(fd, staging, total, aoffs); + } while (r < 0 && errno == EINTR); + if (r < 0 || (size_t) r < head + len) { + return nullptr; + } + return staging + head; + } + + uint8_t * p = staging; + size_t left = len; + while (left > 0) { + const ssize_t r = pread(fd, p, left, offs); + if (r < 0) { + if (errno == EINTR) { + continue; + } + return nullptr; + } + if (r == 0) { + return nullptr; // unexpected EOF + } + p += r; + offs += (size_t) r; + left -= (size_t) r; + } + return staging; +#endif +} + +// true iff all of the given exps tensors are this layer's cache tensors - guards against a second, +// non-streamed expert group on the same layer index (e.g. grovemoe chexps) +bool llama_moe_stream_layer::matches(const ggml_tensor * gate, const ggml_tensor * up, + const ggml_tensor * down, const ggml_tensor * gate_up) const { + auto is_cache = [this](const ggml_tensor * t) { + for (const auto & w : weights) { + if (w.cache == t) { + return true; + } + } + return false; + }; + + size_t n = 0; + for (const ggml_tensor * t : { gate, up, down, gate_up }) { + if (t == nullptr) { + continue; + } + if (!is_cache(t)) { + return false; + } + n++; + } + + return n > 0 && n == weights.size(); +} + +// sizes the per-layer table and clamps the I/O thread count; workers are spawned lazily on first use +llama_moe_stream::llama_moe_stream(uint32_t n_layer, uint32_t n_slots, int32_t n_io_threads, bool direct) : n_slots(n_slots) { + layers.resize(n_layer); + + this->n_io_threads = n_io_threads <= 0 ? MOE_STREAM_IO_THREADS_DEFAULT : n_io_threads; + this->n_io_threads = std::min(this->n_io_threads, MOE_STREAM_IO_THREADS_MAX); + + debug = std::getenv("LLAMA_MOE_STREAM_DEBUG") != nullptr; + use_direct_io = direct; +} + +// stop and join the I/O workers before the cache buffers and files they use are destroyed +llama_moe_stream::~llama_moe_stream() { + { + std::lock_guard lock(mtx); + shutting_down = true; + q_demand.clear(); + } + cv_work.notify_all(); + for (auto & w : workers) { + w.join(); + } +} + +ggml_tensor * llama_moe_stream::create_cache_tensor( + int32_t il, ggml_backend_buffer_type_t buft, const ggml_tensor * meta, + uint16_t file_idx, size_t offs) { + GGML_ASSERT(il >= 0 && (size_t) il < layers.size()); + GGML_ASSERT(ggml_is_contiguous(meta)); + GGML_ASSERT(meta->ne[2] > 0 && meta->ne[3] == 1); + + const uint32_t n_expert = meta->ne[2]; + const size_t nb_expert = ggml_nbytes(meta) / n_expert; + GGML_ASSERT(nb_expert * n_expert == ggml_nbytes(meta)); + GGML_ASSERT(n_slots > 0 && n_slots < n_expert); + + ggml_context * ctx = nullptr; + for (auto & [cur_buft, cur_ctx] : ctxs) { + if (cur_buft == buft) { + ctx = cur_ctx.get(); + break; + } + } + if (ctx == nullptr) { + ggml_init_params params = { + /*.mem_size =*/ ggml_tensor_overhead()*(layers.size()*4 + 1), + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + ctx = ggml_init(params); + if (ctx == nullptr) { + throw std::runtime_error("failed to create ggml context for MoE expert streaming"); + } + ctxs.emplace_back(buft, ctx); + } + + ggml_tensor * cache = ggml_new_tensor_3d(ctx, meta->type, meta->ne[0], meta->ne[1], n_slots); + ggml_format_name(cache, "%s.stream_cache", meta->name); + GGML_ASSERT(ggml_nbytes(cache) == nb_expert * n_slots); + + auto & sl = layers[il]; + if (!sl) { + sl = std::make_unique(); + sl->mgr = this; + sl->il = il; + sl->n_expert = n_expert; + sl->n_slots = n_slots; + sl->slot_expert .resize(n_slots, -1); + sl->slot_state .resize(n_slots, LLAMA_MOE_STREAM_SLOT_EMPTY); + sl->slot_claimed .resize(n_slots, 0); + sl->slot_gen .resize(n_slots, 0); + sl->slot_last_use.resize(n_slots, 0); + sl->route_hotness.resize(n_expert, 0); + sl->seen .resize(n_expert, 0); + sl->keep .resize(n_slots, 0); + } + GGML_ASSERT(sl->n_expert == n_expert); + + sl->weights.push_back({ cache, file_idx, offs, nb_expert }); + + max_nb_expert = std::max(max_nb_expert, nb_expert); + + return cache; +} + +void llama_moe_stream::alloc_bufs(bool no_alloc) { + for (auto & [buft, ctx_ptr] : ctxs) { + ggml_context * ctx = ctx_ptr.get(); + if (ggml_get_first_tensor(ctx) == nullptr) { + continue; + } + + ggml_backend_buffer_t buf; + if (no_alloc) { + buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + t->buffer = buf; + } + } else { + buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft); + } + if (buf == nullptr) { + throw std::runtime_error(format("unable to allocate %s buffer for MoE expert streaming", ggml_backend_buft_name(buft))); + } + ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + bufs.emplace_back(buf); + + LLAMA_LOG_INFO("%s: %12s expert cache size = %8.2f MiB (%u slots per layer)\n", + __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf) / 1024.0 / 1024.0, n_slots); + } +} + +void llama_moe_stream::open_files(const std::vector & paths) { + for (const auto & path : paths) { + if (path.empty()) { + throw std::runtime_error("MoE expert streaming requires a file-based model (not a stream/file descriptor)"); + } + } + + auto open_all = [&](bool direct) { + files.clear(); + for (const auto & path : paths) { + files.emplace_back(new llama_file(path.c_str(), "rb", direct)); + } + }; + + open_all(use_direct_io); + + // fall back to buffered when O_DIRECT is unusable: either the open did not honor it (macOS, + // Windows, unsupported filesystems), or it opened but a probe read fails (some network/overlay + // filesystems accept the flag then reject aligned reads). reopening is needed because O_DIRECT + // is a property of the fd. done here, single-threaded, before any worker starts. + if (use_direct_io) { + bool ok = !files.empty() && files.front()->has_direct_io(); + if (ok) { + uint8_t * probe = (uint8_t *) moe_aligned_alloc(MOE_STREAM_DIRECT_ALIGN); + GGML_ASSERT(probe != nullptr); + ok = llama_moe_stream_pread(*files.front(), probe, MOE_STREAM_DIRECT_ALIGN, 0, /*direct =*/ true) != nullptr; + moe_aligned_free(probe); + } + if (!ok) { + LLAMA_LOG_WARN("%s: O_DIRECT not usable, falling back to buffered streaming reads\n", __func__); + use_direct_io = false; + open_all(false); + } + } + + if (use_direct_io) { + LLAMA_LOG_INFO("%s: MoE expert streaming uses O_DIRECT (page cache bypassed)\n", __func__); + } + + // one token drives ~one remap per streamed layer, so decaying every 64 tokens is + // 64 * n_streamed_layers remap calls (computed once here, off the hot path) + int64_t n_streamed = 0; + for (const auto & sl : layers) { + n_streamed += sl != nullptr; + } + hot_decay_interval = MOE_STREAM_HOT_DECAY_TOKENS * n_streamed; +} + +// spawn the I/O thread pool on first use (from the remap callback, under mtx) +void llama_moe_stream::start_workers_locked() { + if (workers_started) { + return; + } + workers_started = true; + workers.reserve(n_io_threads); + for (int32_t i = 0; i < n_io_threads; i++) { + workers.emplace_back([this]() { worker_loop(); }); + } +} + +// I/O worker: pops a reserved load, reads its expert slab(s) from the GGUF file into the cache +// slot, and marks the slot RESIDENT (or flags load_failed); stale/duplicate items are skipped +void llama_moe_stream::worker_loop() { + // page-aligned staging (Metal private buffers require page-aligned source + page-multiple + // length; O_DIRECT needs the extra head/tail slack for its aligned reads) + uint8_t * staging = (uint8_t *) moe_aligned_alloc(max_nb_expert + 2*MOE_STREAM_DIRECT_ALIGN); + GGML_ASSERT(staging != nullptr); + + std::unique_lock lk(mtx); + while (true) { + cv_work.wait(lk, [&]{ return shutting_down || !q_demand.empty(); }); + if (shutting_down) { + break; + } + + llama_moe_stream_work w = q_demand.front(); + q_demand.pop_front(); + + auto & sl = *w.sl; + if (w.gen != sl.slot_gen[w.slot] || + sl.slot_state[w.slot] != LLAMA_MOE_STREAM_SLOT_LOADING || + sl.slot_expert[w.slot] != w.expert || + sl.slot_claimed[w.slot]) { + continue; // stale or duplicate item + } + sl.slot_claimed[w.slot] = 1; + + lk.unlock(); + + bool ok = true; + for (const auto & wt : sl.weights) { + const uint8_t * data = llama_moe_stream_pread(*files[wt.file_idx], staging, wt.nb_expert, wt.offs + (size_t) w.expert*wt.nb_expert, use_direct_io); + if (data == nullptr) { + ok = false; + break; + } + ggml_backend_tensor_set(wt.cache, data, (size_t) w.slot*wt.nb_expert, wt.nb_expert); + } + + lk.lock(); + + sl.slot_claimed[w.slot] = 0; + if (!ok) { + load_failed = true; + } else { + sl.slot_state[w.slot] = LLAMA_MOE_STREAM_SLOT_RESIDENT; + } + cv_done.notify_all(); + } + lk.unlock(); + + moe_aligned_free(staging); +} + +// least valuable evictable slot: empty first, then coldest resident (min route hotness, oldest use +// as tiebreak); LOADING and keep slots are never candidates. returns -1 when no candidate exists +int32_t llama_moe_stream::pick_victim_locked(llama_moe_stream_layer & sl, const uint8_t * keep) const { + int32_t v = -1; + + for (uint32_t s = 0; s < sl.n_slots; s++) { + if ((keep && keep[s]) || sl.slot_state[s] == LLAMA_MOE_STREAM_SLOT_LOADING) { + continue; + } + if (sl.slot_state[s] == LLAMA_MOE_STREAM_SLOT_EMPTY) { + return s; + } + if (v < 0) { + v = s; + continue; + } + const uint32_t hs = sl.route_hotness[sl.slot_expert[s]]; + const uint32_t hv = sl.route_hotness[sl.slot_expert[v]]; + if (hs < hv || (hs == hv && sl.slot_last_use[s] < sl.slot_last_use[v])) { + v = s; + } + } + + return v; +} + +// bind expert -> slot and mark it LOADING: evict the slot's prior occupant, bump slot_gen (so any +// in-flight load for the old occupant is recognized as stale), and update the expert_slot index +void llama_moe_stream::reserve_slot_locked(llama_moe_stream_layer & sl, int32_t expert, int32_t slot) { + if (sl.slot_expert[slot] >= 0) { + if (debug) { + LLAMA_LOG_DEBUG("%s: layer %d: evict expert %d from slot %d\n", __func__, sl.il, sl.slot_expert[slot], slot); + } + sl.expert_slot.erase(sl.slot_expert[slot]); + } + + sl.slot_expert[slot] = expert; + sl.slot_state[slot] = LLAMA_MOE_STREAM_SLOT_LOADING; + sl.slot_gen[slot]++; + sl.slot_last_use[slot] = ++sl.use_counter; + sl.expert_slot[expert] = slot; + sl.seen[expert] = 1; +} + +size_t llama_moe_stream::size_bufs() const { + size_t size = 0; + for (const auto & buf : bufs) { + size += ggml_backend_buffer_get_size(buf.get()); + } + return size; +} + +void llama_moe_stream::print_stats() const { + std::lock_guard lock(mtx); + + const int64_t n_touched = stats.n_hit + stats.n_miss; + LLAMA_LOG_INFO("%s: moe stream: remap calls = %" PRId64 ", expert hits = %" PRId64 ", misses = %" PRId64 " (%" PRId64 " cold), hit rate = %.2f%%\n", + __func__, stats.n_calls, stats.n_hit, stats.n_miss, stats.n_miss_cold, + n_touched > 0 ? 100.0*stats.n_hit/n_touched : 0.0); + LLAMA_LOG_INFO("%s: moe stream: load stall = %.2f ms total (%.3f ms per remap call)\n", + __func__, stats.t_stall_us/1000.0, stats.n_calls > 0 ? stats.t_stall_us/1000.0/stats.n_calls : 0.0); +} + +// custom-op callback (single-threaded on ith 0): given the router's expert ids, ensure every touched +// expert is resident - reserving cache slots and demand-loading misses, stalling until they commit - +// then rewrite each id to its cache slot. this only relabels ids, so the same experts are computed +// in the same order; the result matches a non-streamed run (bit-exact when both paths use the same +// kernels, as on CUDA; a CPU build that repacks the non-streamed weights can differ in the last bits). +void llama_moe_stream_remap(ggml_tensor * dst, const ggml_tensor * a, int ith, int nth, void * userdata) { + GGML_UNUSED(nth); + if (ith != 0) { + return; + } + + auto * sl = (llama_moe_stream_layer *) userdata; + auto * mgr = sl->mgr; + + GGML_ASSERT(a->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_are_same_shape(a, dst)); + + const int64_t n = ggml_nelements(a); + + const int32_t * ids = (const int32_t *) a->data; + int32_t * out = (int32_t *) dst->data; + + std::unique_lock lk(mgr->mtx); + + if (mgr->load_failed) { + GGML_ABORT("MoE expert streaming: expert load failed (I/O error)"); + } + + mgr->stats.n_calls++; + mgr->start_workers_locked(); + + // distinct experts touched by this ubatch, in first-use order + sl->touched.assign(sl->n_expert, 0); + sl->uniq.clear(); + for (int64_t i = 0; i < n; i++) { + const int32_t e = ids[i]; + GGML_ASSERT(e >= 0 && (uint32_t) e < sl->n_expert); + if (!sl->touched[e]) { + sl->touched[e] = 1; + sl->uniq.push_back(e); + } + } + + if (sl->uniq.size() > sl->n_slots) { + GGML_ABORT("MoE expert streaming: layer %d needs %zu distinct experts but the cache has only %u slots; " + "increase --moe-stream-cache or reduce the ubatch size (-ub)", + sl->il, sl->uniq.size(), sl->n_slots); + } + + // route hotness for eviction; halved periodically so a formerly-hot expert ages out + for (const int32_t e : sl->uniq) { + sat_inc(sl->route_hotness[e]); + } + if (mgr->hot_decay_interval > 0 && mgr->stats.n_calls % mgr->hot_decay_interval == 0) { + for (auto & sl2 : mgr->layers) { + if (sl2) { + for (auto & h : sl2->route_hotness) { + h >>= 1; + } + } + } + } + + // classify the touched experts; reserve and enqueue demand loads in deterministic order + std::fill(sl->keep.begin(), sl->keep.end(), 0); + sl->demand_slots.clear(); + + bool waited = false; + for (const int32_t e : sl->uniq) { + const auto it = sl->expert_slot.find(e); + if (it != sl->expert_slot.end()) { + const int32_t s = it->second; + if (sl->slot_state[s] == LLAMA_MOE_STREAM_SLOT_LOADING) { + mgr->q_demand.push_back({ sl, e, s, sl->slot_gen[s] }); + mgr->cv_work.notify_one(); + waited = true; + } + mgr->stats.n_hit++; + sl->keep[s] = 1; + sl->demand_slots.push_back(s); + } else { + int32_t v; + while ((v = mgr->pick_victim_locked(*sl, sl->keep.data())) < 0) { + // every allowed slot is loading; wait for a commit and retry + mgr->cv_done.wait(lk); + if (mgr->load_failed) { + GGML_ABORT("MoE expert streaming: expert load failed (I/O error)"); + } + } + if (!sl->seen[e]) { + mgr->stats.n_miss_cold++; + } + mgr->reserve_slot_locked(*sl, e, v); + mgr->q_demand.push_back({ sl, e, v, sl->slot_gen[v] }); + mgr->cv_work.notify_one(); + mgr->stats.n_miss++; + waited = true; + sl->keep[v] = 1; + sl->demand_slots.push_back(v); + } + } + + if (waited) { + const int64_t t0 = ggml_time_us(); + mgr->cv_done.wait(lk, [&]{ + if (mgr->load_failed) { + return true; + } + for (const int32_t s : sl->demand_slots) { + if (sl->slot_state[s] != LLAMA_MOE_STREAM_SLOT_RESIDENT) { + return false; + } + } + return true; + }); + if (mgr->load_failed) { + GGML_ABORT("MoE expert streaming: expert load failed (I/O error)"); + } + mgr->stats.t_stall_us += ggml_time_us() - t0; + } + + for (int64_t i = 0; i < n; i++) { + const int32_t s = sl->expert_slot.at(ids[i]); + sl->slot_last_use[s] = ++sl->use_counter; + out[i] = s; + } +} diff --git a/src/llama-moe-stream.h b/src/llama-moe-stream.h new file mode 100644 index 000000000000..aebdb5cc4ee7 --- /dev/null +++ b/src/llama-moe-stream.h @@ -0,0 +1,161 @@ +#pragma once + +#include "llama-mmap.h" + +#include "ggml-cpp.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// SSD streaming of MoE routed expert weights +// +// Streamed layers do not materialize their ffn_*_exps tensors; instead each weight gets a +// device-side cache tensor of n_slots expert slabs, filled on demand from the GGUF file by an +// id-remapping custom op that runs on the CPU right after the router top-k. The remap only +// changes which cache slot an expert id resolves to - it never changes which experts the router +// selected, so streaming affects latency, not outputs. +// +// Missing experts are loaded by a pool of I/O threads while the remap op waits; eviction is by +// decaying route hotness with an LRU tiebreak. Reads are buffered by default, or O_DIRECT with +// LLAMA_MOE_STREAM_DIRECT=1 (bypasses the page cache; recommended when the model far exceeds RAM). +// +// note: multiple contexts decoding the same streamed model concurrently are not supported - +// one context can evict slots referenced by the other's in-flight graph. + +struct llama_moe_stream; + +enum llama_moe_stream_slot_state : uint8_t { + LLAMA_MOE_STREAM_SLOT_EMPTY = 0, + LLAMA_MOE_STREAM_SLOT_LOADING = 1, // reserved, load queued or in flight + LLAMA_MOE_STREAM_SLOT_RESIDENT = 2, +}; + +// one streamed weight tensor (gate/up/down or fused gate_up) of one layer +struct llama_moe_stream_weight { + ggml_tensor * cache = nullptr; // cache tensor {ne0, ne1, n_slots} + + uint16_t file_idx = 0; // GGUF split file index + size_t offs = 0; // file offset of the full exps tensor data + size_t nb_expert = 0; // bytes per expert slab +}; + +// per-layer streaming state - also the userdata of the id-remapping custom op +struct llama_moe_stream_layer { + llama_moe_stream * mgr = nullptr; + + int32_t il = -1; + uint32_t n_expert = 0; + uint32_t n_slots = 0; + + std::vector weights; // 2 (fused gate_up + down) or 3 entries + + // residency state, guarded by mgr->mtx + std::vector slot_expert; // [n_slots] expert id or -1 + std::vector slot_state; // [n_slots] llama_moe_stream_slot_state + std::vector slot_claimed; // [n_slots] a worker owns the load + std::vector slot_gen; // [n_slots] reservation generation + std::vector slot_last_use; // [n_slots] LRU stamps + std::unordered_map expert_slot; // RESIDENT and LOADING entries + + std::vector route_hotness; // [n_expert] decayed selection counts, for eviction + std::vector seen; // [n_expert] for cold-miss attribution + int64_t use_counter = 0; + + // scratch for the remap callback + std::vector uniq; + std::vector touched; + std::vector keep; // [n_slots] slots the current call must not evict + std::vector demand_slots; // slots the current call waits on + + // whether the exps tensors passed to build_moe_ffn are this layer's cache tensors + // (e.g. grovemoe evaluates a second, unstreamed expert group on the same layer index) + bool matches(const ggml_tensor * gate, const ggml_tensor * up, + const ggml_tensor * down, const ggml_tensor * gate_up) const; +}; + +// one queued expert load +struct llama_moe_stream_work { + llama_moe_stream_layer * sl = nullptr; + + int32_t expert = -1; + int32_t slot = -1; + uint64_t gen = 0; // stale unless it matches slot_gen[slot] +}; + +struct llama_moe_stream { + uint32_t n_slots = 0; // expert cache slots per streamed layer + int32_t n_io_threads = 0; + + std::vector> layers; // [n_layer], null = not streamed + + llama_moe_stream(uint32_t n_layer, uint32_t n_slots, int32_t n_io_threads, bool direct); + ~llama_moe_stream(); + + llama_moe_stream_layer * layer(int32_t il) const { + return il >= 0 && (size_t) il < layers.size() ? layers[il].get() : nullptr; + } + + // registers a streamed weight of layer il and returns its cache tensor + ggml_tensor * create_cache_tensor( + int32_t il, ggml_backend_buffer_type_t buft, const ggml_tensor * meta, + uint16_t file_idx, size_t offs); + + // allocate the cache tensor buffers (after all create_cache_tensor calls) + void alloc_bufs(bool no_alloc); + + // reopen the GGUF files for streaming reads + void open_files(const std::vector & paths); + + size_t size_bufs() const; + + void print_stats() const; + + bool use_direct_io = false; // O_DIRECT streaming reads (LLAMA_MOE_STREAM_DIRECT), no page cache + + llama_files files; // privately reopened GGUF files, same indices as the loader's + + size_t max_nb_expert = 0; + int64_t hot_decay_interval = 0; // remap calls between route-hotness halvings (0 = no decay) + + std::vector> ctxs; // one per buft + std::vector bufs; + + // load pool (queue and all layer residency state guarded by mtx) + mutable std::mutex mtx; + std::condition_variable cv_work; // queued work or shutdown + std::condition_variable cv_done; // a load committed or failed + + std::deque q_demand; + + std::vector workers; + bool workers_started = false; + bool shutting_down = false; + bool load_failed = false; + + bool debug = false; + + struct { + int64_t n_calls = 0; // remap invocations + int64_t n_hit = 0; // touched experts already resident or loading + int64_t n_miss = 0; // demand loads issued + int64_t n_miss_cold = 0; // first-ever touch of an expert + int64_t t_stall_us = 0; // wait time in miss handling + } stats; + + // internals + void start_workers_locked(); + void worker_loop(); + int32_t pick_victim_locked(llama_moe_stream_layer & sl, const uint8_t * keep) const; + void reserve_slot_locked(llama_moe_stream_layer & sl, int32_t expert, int32_t slot); +}; + +// callback of the id-remapping custom op inserted by build_moe_ffn +void llama_moe_stream_remap(ggml_tensor * dst, const ggml_tensor * a, int ith, int nth, void * userdata); diff --git a/src/llama.cpp b/src/llama.cpp index 0de6048f2820..c7e41ae24400 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -279,6 +279,13 @@ static bool llama_prepare_model_devices(const llama_model_params & params, llama static std::pair llama_model_load(struct gguf_context * metadata, llama_model_set_tensor_data_t set_tensor_data, void * set_tensor_data_ud, const std::string & fname, std::vector & splits, FILE * file, llama_model_params & params) { try { + if (params.moe_stream && params.use_mmap) { + // mmap prefetches the whole file into RAM (MAP_POPULATE / MADV_WILLNEED), which loads + // the streamed experts too and defeats streaming - forcing an OOM on a model >> RAM + LLAMA_LOG_WARN("%s: disabling mmap because MoE expert streaming is enabled\n", __func__); + params.use_mmap = false; + } + llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.use_mmap, params.use_direct_io, params.check_tensors, params.no_alloc, params.kv_overrides, params.tensor_buft_overrides); From db5268ffac4c3ddacd7edcf62e78e3a1ba8d8191 Mon Sep 17 00:00:00 2001 From: Junchao Lyu Date: Sat, 4 Jul 2026 10:21:20 -0700 Subject: [PATCH 2/2] support waved prefill Assisted-by: Claude --- src/llama-context.cpp | 35 +++-- src/llama-graph.cpp | 69 +++++++++- src/llama-moe-stream.cpp | 289 +++++++++++++++++++++++++++++++++++++++ src/llama-moe-stream.h | 40 ++++++ 4 files changed, 418 insertions(+), 15 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 07c0f58f5393..bf2d58205e0f 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -213,17 +213,13 @@ llama_context::llama_context( cparams.kv_unified = params.kv_unified; if (model.moe_stream() && hparams.n_expert_used > 0) { - // one mul_mat_id needs every expert a ubatch touches resident at once, so a ubatch may - // not select more experts than the cache holds (worst case n_ubatch*n_expert_used) - const uint32_t n_ubatch_max = std::max(1u, model.moe_stream()->n_slots / hparams.n_expert_used); - if (cparams.n_ubatch > n_ubatch_max) { - LLAMA_LOG_WARN("%s: n_ubatch reduced from %u to %u so that a ubatch cannot select more experts than the %u-slot streaming cache\n", - __func__, cparams.n_ubatch, n_ubatch_max, model.moe_stream()->n_slots); - cparams.n_ubatch = n_ubatch_max; - } + // ubatches that touch more experts than the streaming cache holds run the expert GEMMs in + // multiple waves, so no ubatch size restriction is needed + LLAMA_LOG_INFO("%s: MoE expert streaming with %u cache slots, n_ubatch = %u\n", + __func__, model.moe_stream()->n_slots, cparams.n_ubatch); // op offload snapshots host weights to the device per graph split, which assumes they do - // not change during the graph - streamed caches are rewritten on demand + // not change during the graph - streamed caches are rewritten between waves bool cache_on_host = false; for (const auto & buf : model.moe_stream()->bufs) { cache_on_host = cache_on_host || ggml_backend_buffer_is_host(buf.get()); @@ -2344,16 +2340,27 @@ void llama_context::output_reorder() { // uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { + uint32_t res; if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_DEEPSEEK4) { - return std::max(n_tokens * 40, 32u * model.n_tensors()); - } - uint32_t res = std::max(1024u, 8u*model.n_tensors()); - for (const auto & lora : model.loras) { - res += lora->get_n_nodes(); + res = std::max(n_tokens * 40, 32u * model.n_tensors()); + } else { + res = std::max(1024u, 8u*model.n_tensors()); + for (const auto & lora : model.loras) { + res += lora->get_n_nodes(); + } + } + if (const auto * mstream = model.moe_stream()) { + // multi-pass streamed prefill adds a bounded number of extra nodes per wave per streamed layer + const uint32_t n_eu = model.hparams.n_expert_used; + uint32_t cap = mstream->n_slots > n_eu ? (mstream->n_slots - n_eu)/2 : 0; + cap = std::max(cap, 1); + const uint32_t n_touch_max = std::min(model.hparams.n_expert, n_tokens*n_eu); + const uint32_t n_waves = (n_touch_max + cap - 1)/cap; + res += 24u*n_waves*(uint32_t) mstream->layers.size(); } return res; } diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 20487b5848f5..b2117bbb265f 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1965,8 +1965,33 @@ ggml_tensor * llm_graph_context::build_moe_ffn( msl = nullptr; // a different expert group of the same layer (e.g. grovemoe chexps), not streamed } - ggml_tensor * ids_gemm = selected_experts; + // a ubatch can touch more distinct experts than the cache holds; the expert GEMMs then run in + // waves of at most stream_wave_cap experts, the pairs of the other waves masked to zero and the + // wave outputs summed. n_touch_max caps at n_expert, so each expert is loaded at most once per + // ubatch regardless of batch size - wave splitting bounds prefill I/O to one sweep of them. + uint32_t n_stream_waves = 1; + uint32_t stream_wave_cap = 0; if (msl) { + // worst-case distinct experts: n_expert_used per token, but never more than n_expert total + const uint64_t n_touch_max = std::min((uint64_t) n_expert, (uint64_t) n_tokens*n_expert_used); + if (n_touch_max > msl->n_slots) { + // the cache must hold three sets at once: this wave's experts, the next wave's preloaded + // experts (so its loads overlap this wave's compute), and n_expert_used parking slots + // the masked-out pairs GEMM against (Metal needs a slot at most once per token row). + // so cap + cap + n_expert_used = n_slots -> cap = (n_slots - n_expert_used)/2 + stream_wave_cap = msl->n_slots > (uint32_t) n_expert_used ? (msl->n_slots - (uint32_t) n_expert_used)/2 : 0; + if (stream_wave_cap < (uint32_t) n_expert_used) { + // a wave must fit at least n_expert_used experts, i.e. n_slots >= 3*n_expert_used + GGML_ABORT("MoE expert streaming: multi-pass expert GEMMs need an expert cache of at least " + "3*n_expert_used slots (have %u, need %u); increase --moe-stream-cache or reduce -ub", + msl->n_slots, 3*(uint32_t) n_expert_used); + } + n_stream_waves = (uint32_t) ((n_touch_max + stream_wave_cap - 1)/stream_wave_cap); // ceil(n_touch_max/cap) + } + } + + ggml_tensor * ids_gemm = selected_experts; + if (msl && n_stream_waves == 1) { ggml_tensor * ids_cont = ggml_cont(ctx0, selected_experts); // top_k output is a view ids_gemm = ggml_map_custom1(ctx0, ids_cont, llama_moe_stream_remap, 1, msl); cb(ids_gemm, "ffn_moe_topk_stream", il); @@ -1981,6 +2006,9 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cb(cur, "ffn_moe_weighted", il); } + // the expert GEMM pipeline: run once normally, or once per wave under multi-pass prefill; + // biases and per-expert scales are always indexed by the original selected_experts + auto build_expert_gemms = [&](ggml_tensor * cur, ggml_tensor * ids_gemm) -> ggml_tensor * { ggml_tensor * up = nullptr; ggml_tensor * experts = nullptr; @@ -2119,6 +2147,45 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cb(experts, "ffn_moe_down_biased", il); } + return experts; + }; // build_expert_gemms + + ggml_tensor * experts = nullptr; + + if (msl && n_stream_waves > 1) { + ggml_tensor * ids_cont = ggml_cont(ctx0, selected_experts); + + for (uint32_t w = 0; w < n_stream_waves; w++) { + ggml_tensor * args[2] = { ids_cont, nullptr }; + int n_args = 1; + if (experts != nullptr) { + // ordering token: forces this wave's ids op to run after the previous wave's GEMMs + // consumed their slots; a 1-element view keeps the cross-backend copy tiny + args[1] = ggml_view_1d(ctx0, experts, 1, 0); + n_args = 2; + } + ggml_tensor * ids_w = ggml_custom_4d(ctx0, GGML_TYPE_I32, + ids_cont->ne[0], ids_cont->ne[1], 1, 1, + args, n_args, llama_moe_stream_wave_ids, 1, msl->wave_userdata(w, stream_wave_cap)); + cb(ids_w, "ffn_moe_wave_ids", il); + + ggml_tensor * e_w = build_expert_gemms(cur, ids_w); + + ggml_tensor * margs[2] = { ids_cont, ids_w }; + ggml_tensor * mask_w = ggml_custom_4d(ctx0, GGML_TYPE_F32, + 1, n_expert_used, n_tokens, 1, + margs, 2, llama_moe_stream_wave_mask, 1, msl->wave_userdata(w, stream_wave_cap)); + cb(mask_w, "ffn_moe_wave_mask", il); + + e_w = ggml_mul(ctx0, e_w, mask_w); // zero the pairs that belong to other waves + + experts = experts == nullptr ? e_w : ggml_add(ctx0, experts, e_w); + ggml_build_forward_expand(gf, experts); + } + } else { + experts = build_expert_gemms(cur, ids_gemm); + } + if (!weight_before_ffn) { experts = ggml_mul(ctx0, experts, weights); cb(experts, "ffn_moe_weighted", il); diff --git a/src/llama-moe-stream.cpp b/src/llama-moe-stream.cpp index 19c755730f2f..ce173960a492 100644 --- a/src/llama-moe-stream.cpp +++ b/src/llama-moe-stream.cpp @@ -424,6 +424,10 @@ void llama_moe_stream::print_stats() const { n_touched > 0 ? 100.0*stats.n_hit/n_touched : 0.0); LLAMA_LOG_INFO("%s: moe stream: load stall = %.2f ms total (%.3f ms per remap call)\n", __func__, stats.t_stall_us/1000.0, stats.n_calls > 0 ? stats.t_stall_us/1000.0/stats.n_calls : 0.0); + if (stats.n_wave_calls > 0) { + LLAMA_LOG_INFO("%s: moe stream: waves = %" PRId64 " (%" PRId64 " non-empty), preloads issued = %" PRId64 " (ready on arrival = %" PRId64 "), wave stall = %.2f ms\n", + __func__, stats.n_wave_calls, stats.n_waves_run, stats.n_preload_issued, stats.n_preload_ready, stats.t_stall_wave_us/1000.0); + } } // custom-op callback (single-threaded on ith 0): given the router's expert ids, ensure every touched @@ -554,3 +558,288 @@ void llama_moe_stream_remap(ggml_tensor * dst, const ggml_tensor * a, int ith, i out[i] = s; } } + +// stable per-wave userdata; grows lazily and records the per-wave expert capacity (set at build) +llama_moe_stream_wave * llama_moe_stream_layer::wave_userdata(int32_t wave, uint32_t capacity) { + GGML_ASSERT(capacity >= 1 && capacity <= n_slots); + plan_capacity = capacity; + while ((size_t) wave >= wave_ud.size()) { + auto ud = std::make_unique(); + ud->sl = this; + ud->wave = (int32_t) wave_ud.size(); + wave_ud.push_back(std::move(ud)); + } + return wave_ud[wave].get(); +} + +// wave 0 of a ubatch: record the distinct touched experts (sl.uniq, first-use order) and split them +// into consecutive groups of plan_capacity, one group per wave (sl.expert_wave[e] = e's wave) +void llama_moe_stream::plan_waves_locked(llama_moe_stream_layer & sl, const int32_t * ids, int64_t n) { + stats.n_calls++; + start_workers_locked(); + + sl.touched.assign(sl.n_expert, 0); + sl.uniq.clear(); + for (int64_t i = 0; i < n; i++) { + const int32_t e = ids[i]; + GGML_ASSERT(e >= 0 && (uint32_t) e < sl.n_expert); + if (!sl.touched[e]) { + sl.touched[e] = 1; + sl.uniq.push_back(e); + } + } + + GGML_ASSERT(sl.plan_capacity > 0); + sl.expert_wave.assign(sl.n_expert, 0xff); + for (size_t i = 0; i < sl.uniq.size(); i++) { + GGML_ASSERT(i/sl.plan_capacity < 0xff); + sl.expert_wave[sl.uniq[i]] = (uint8_t) (i/sl.plan_capacity); + } + sl.plan_n_waves = (uint32_t) ((sl.uniq.size() + sl.plan_capacity - 1)/sl.plan_capacity); + sl.plan_next_wave = 0; +} + +// make wave w's expert slice (uniq[w*cap .. +count)) resident, waiting for its loads, and best-effort +// preload the next wave so its loads overlap this wave's compute. leaves sl.demand_slots = this wave's +// slots and sl.plan_pool = the resident parking pool (>= n_ids slots) the emit draws masked pairs from +void llama_moe_stream::stage_wave_locked(std::unique_lock & lk, llama_moe_stream_layer & sl, int32_t w, uint32_t n_ids) { + const size_t first = (size_t) w*sl.plan_capacity; + const size_t count = first < sl.uniq.size() ? std::min(sl.plan_capacity, sl.uniq.size() - first) : 0; + + std::fill(sl.keep.begin(), sl.keep.end(), 0); + sl.demand_slots.clear(); + + // a small final wave has fewer than n_ids own slots; borrow the rest from the previous wave's + // pool so every token row has n_ids distinct resident parking slots for its masked pairs + std::vector borrowed; + if (count < n_ids) { + GGML_ASSERT(sl.plan_pool.size() >= n_ids - count); + for (size_t i = 0; i < n_ids - count; i++) { + borrowed.push_back(sl.plan_pool[i]); + sl.keep[sl.plan_pool[i]] = 1; // parking slots must survive this wave's loads + } + } + + // protect the next wave's already-resident experts so this wave's victims do not evict them + const size_t nfirst = first + sl.plan_capacity; + const size_t ncount = nfirst < sl.uniq.size() ? std::min(sl.plan_capacity, sl.uniq.size() - nfirst) : 0; + for (size_t i = nfirst; i < nfirst + ncount; i++) { + const auto it = sl.expert_slot.find(sl.uniq[i]); + if (it != sl.expert_slot.end()) { + sl.keep[it->second] = 1; + } + } + + // reserve and demand-load this wave's experts (per-expert, same path as the decode remap) + bool waited = false; + if (count > 0) { + stats.n_waves_run++; + for (size_t i = first; i < first + count; i++) { + const int32_t e = sl.uniq[i]; + const auto it = sl.expert_slot.find(e); + if (it != sl.expert_slot.end()) { + // already in the cache (resident, or still loading from the previous wave's preload) + const int32_t s = it->second; + if (sl.slot_state[s] == LLAMA_MOE_STREAM_SLOT_LOADING) { + q_demand.push_back({ &sl, e, s, sl.slot_gen[s] }); // promote to demand, wait for it + cv_work.notify_one(); + waited = true; + } else { + stats.n_preload_ready++; // resident from the previous wave's preload + } + stats.n_hit++; + sl.keep[s] = 1; + sl.demand_slots.push_back(s); + } else { + // miss: evict a non-kept slot and queue the load + int32_t v; + while ((v = pick_victim_locked(sl, sl.keep.data())) < 0) { + cv_done.wait(lk); + if (load_failed) { + GGML_ABORT("MoE expert streaming: expert load failed (I/O error)"); + } + } + if (!sl.seen[e]) { + stats.n_miss_cold++; + } + reserve_slot_locked(sl, e, v); + q_demand.push_back({ &sl, e, v, sl.slot_gen[v] }); + cv_work.notify_one(); + stats.n_miss++; + waited = true; + sl.keep[v] = 1; + sl.demand_slots.push_back(v); + } + } + } + + // best-effort preload of the next wave so its loads overlap this wave's compute; never waits, + // whatever cannot be reserved now simply becomes the next wave's demand load + if (std::getenv("LLAMA_MOE_STREAM_NO_PRELOAD") == nullptr) { + for (size_t i = nfirst; i < nfirst + ncount; i++) { + const int32_t e = sl.uniq[i]; + if (sl.expert_slot.find(e) != sl.expert_slot.end()) { + continue; + } + const int32_t v = pick_victim_locked(sl, sl.keep.data()); + if (v < 0) { + continue; + } + if (!sl.seen[e]) { + stats.n_miss_cold++; + } + reserve_slot_locked(sl, e, v); + sl.keep[v] = 1; + q_demand.push_back({ &sl, e, v, sl.slot_gen[v] }); + cv_work.notify_one(); + stats.n_preload_issued++; + } + } + + if (waited) { + const int64_t t0 = ggml_time_us(); + cv_done.wait(lk, [&]{ + if (load_failed) { + return true; + } + for (const int32_t s : sl.demand_slots) { + if (sl.slot_state[s] != LLAMA_MOE_STREAM_SLOT_RESIDENT) { + return false; + } + } + return true; + }); + if (load_failed) { + GGML_ABORT("MoE expert streaming: expert load failed (I/O error)"); + } + stats.t_stall_wave_us += ggml_time_us() - t0; + } + + // parking pool: this wave's own resident slots plus the borrowed ones (all keep-protected; + // the next same-layer reservation is ordered after this wave's GEMMs by the graph) + sl.plan_pool = sl.demand_slots; + sl.plan_pool.insert(sl.plan_pool.end(), borrowed.begin(), borrowed.end()); + GGML_ASSERT(sl.plan_pool.size() >= n_ids); +} + +// write out[i] = the cache slot the GEMM should index for each (token, expert) pair of wave w, one +// token row at a time: pairs whose expert is in this wave get its real slot; the rest park on distinct +// resident pool slots (pool_used prevents a repeat within the row, required by the Metal kernel) +void llama_moe_stream::emit_wave_slots(llama_moe_stream_layer & sl, const int32_t * ids, int32_t * out, + int32_t w, uint32_t n_ids, int64_t n_tok) { + for (int64_t t = 0; t < n_tok; t++) { + sl.pool_used.clear(); + + // pass 1: pairs whose expert belongs to this wave -> that expert's real (resident) slot + for (uint32_t kk = 0; kk < n_ids; kk++) { + const int64_t i = t*n_ids + kk; + const int32_t e = ids[i]; + GGML_ASSERT(sl.expert_wave[e] != 0xff); + if (sl.expert_wave[e] == (uint8_t) w) { + const int32_t s = sl.expert_slot.at(e); + GGML_ASSERT(sl.slot_state[s] == LLAMA_MOE_STREAM_SLOT_RESIDENT); + sl.slot_last_use[s] = ++sl.use_counter; + out[i] = s; + sl.pool_used.push_back(s); + } + } + + // pass 2: the remaining (masked) pairs -> the next pool slot not yet used in this row + size_t pi = 0; + for (uint32_t kk = 0; kk < n_ids; kk++) { + const int64_t i = t*n_ids + kk; + if (sl.expert_wave[ids[i]] == (uint8_t) w) { + continue; + } + while (std::find(sl.pool_used.begin(), sl.pool_used.end(), sl.plan_pool[pi]) != sl.pool_used.end()) { + pi++; + GGML_ASSERT(pi < sl.plan_pool.size()); + } + GGML_ASSERT(sl.slot_state[sl.plan_pool[pi]] == LLAMA_MOE_STREAM_SLOT_RESIDENT); + out[i] = sl.plan_pool[pi]; + sl.pool_used.push_back(sl.plan_pool[pi]); + pi++; + } + } +} + +// Custom-op callback for one pass of multi-pass prefill. When a ubatch touches more experts than the +// cache holds, build_moe_ffn runs the expert GEMMs in several waves; this runs once per wave (single- +// threaded on ith 0), in wave order. For wave w it makes that wave's expert slice resident (preloading +// the next wave), then writes the slot ids the GEMM indexes - see plan_waves_locked / stage_wave_locked +// / emit_wave_slots. The router's expert choice is untouched, so the output matches a non-streamed run. +void llama_moe_stream_wave_ids(ggml_tensor * dst, int ith, int nth, void * userdata) { + GGML_UNUSED(nth); + if (ith != 0) { + return; + } + + auto * ud = (llama_moe_stream_wave *) userdata; + auto * sl = ud->sl; + auto * mgr = sl->mgr; + + const int32_t w = ud->wave; + + const ggml_tensor * a = dst->src[0]; // contiguous selected ids + GGML_ASSERT(a->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(a)); + GGML_ASSERT(dst->data != a->data); // the emit must not clobber the ids other waves read + + const int64_t n = ggml_nelements(a); + const int32_t * ids = (const int32_t *) a->data; + int32_t * out = (int32_t *) dst->data; + + std::unique_lock lk(mgr->mtx); + + if (mgr->load_failed) { + GGML_ABORT("MoE expert streaming: expert load failed (I/O error)"); + } + + mgr->stats.n_wave_calls++; + + if (w == 0) { + mgr->plan_waves_locked(*sl, ids, n); + } + GGML_ASSERT(sl->plan_next_wave == w); // waves must run in order (enforced by the graph ordering token) + + const uint32_t n_ids = (uint32_t) a->ne[0]; // experts per token (n_expert_used) + + mgr->stage_wave_locked(lk, *sl, w, n_ids); // make this wave resident, preload the next, build the pool + sl->plan_next_wave = w + 1; + + mgr->emit_wave_slots(*sl, ids, out, w, n_ids, a->ne[1]); +} + +// multi-pass prefill: 1.0 for pairs whose expert belongs to wave w, 0.0 otherwise; multiplied into +// this wave's expert GEMM output so the masked-out (parked) pairs contribute nothing to the sum +void llama_moe_stream_wave_mask(ggml_tensor * dst, int ith, int nth, void * userdata) { + GGML_UNUSED(nth); + if (ith != 0) { + return; + } + + auto * ud = (llama_moe_stream_wave *) userdata; + auto * sl = ud->sl; + auto * mgr = sl->mgr; + + const int32_t w = ud->wave; + + const ggml_tensor * a = dst->src[0]; // contiguous selected ids + GGML_ASSERT(a->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(a)); + + const int64_t n = ggml_nelements(a); + const int32_t * ids = (const int32_t *) a->data; + float * out = (float *) dst->data; + + std::lock_guard lock(mgr->mtx); + + GGML_ASSERT(sl->plan_next_wave > w); // this wave's ids op has already run + + for (int64_t i = 0; i < n; i++) { + out[i] = sl->expert_wave[ids[i]] == (uint8_t) w ? 1.0f : 0.0f; + } +} diff --git a/src/llama-moe-stream.h b/src/llama-moe-stream.h index aebdb5cc4ee7..95032fa3d7f4 100644 --- a/src/llama-moe-stream.h +++ b/src/llama-moe-stream.h @@ -47,6 +47,14 @@ struct llama_moe_stream_weight { size_t nb_expert = 0; // bytes per expert slab }; +struct llama_moe_stream_layer; + +// userdata of one wave's custom ops (multi-pass prefill): identifies which pass this is +struct llama_moe_stream_wave { + llama_moe_stream_layer * sl = nullptr; + int32_t wave = -1; +}; + // per-layer streaming state - also the userdata of the id-remapping custom op struct llama_moe_stream_layer { llama_moe_stream * mgr = nullptr; @@ -75,6 +83,20 @@ struct llama_moe_stream_layer { std::vector keep; // [n_slots] slots the current call must not evict std::vector demand_slots; // slots the current call waits on + // wave plan for multi-pass prefill (guarded by mgr->mtx): the touched experts are split into + // plan_n_waves passes of at most plan_capacity experts each, run one pass at a time + uint32_t plan_capacity = 0; // experts per wave, set at graph build + uint32_t plan_n_waves = 0; // waves of the current call + int32_t plan_next_wave = -1; // wave expected to run next (ordering guard) + std::vector expert_wave; // [n_expert] wave each touched expert belongs to, 0xff = untouched + std::vector plan_pool; // resident slots the masked-out pairs of this wave park on + std::vector pool_used; // scratch: pool slots already used in the current token row + + std::vector> wave_ud; // stable per-wave op userdata + + // stable userdata for wave w (grows lazily); called at graph build time only + llama_moe_stream_wave * wave_userdata(int32_t wave, uint32_t capacity); + // whether the exps tensors passed to build_moe_ffn are this layer's cache tensors // (e.g. grovemoe evaluates a second, unstreamed expert group on the same layer index) bool matches(const ggml_tensor * gate, const ggml_tensor * up, @@ -148,6 +170,12 @@ struct llama_moe_stream { int64_t n_miss = 0; // demand loads issued int64_t n_miss_cold = 0; // first-ever touch of an expert int64_t t_stall_us = 0; // wait time in miss handling + + int64_t n_wave_calls = 0; // wave-ids invocations (>= n_calls under multi-pass prefill) + int64_t n_waves_run = 0; // non-empty waves + int64_t n_preload_issued = 0; // next-wave loads started during a wave's compute + int64_t n_preload_ready = 0; // wave experts already resident from the previous preload + int64_t t_stall_wave_us = 0; // wait time in wave miss handling } stats; // internals @@ -155,7 +183,19 @@ struct llama_moe_stream { void worker_loop(); int32_t pick_victim_locked(llama_moe_stream_layer & sl, const uint8_t * keep) const; void reserve_slot_locked(llama_moe_stream_layer & sl, int32_t expert, int32_t slot); + + // multi-pass prefill helpers (called by llama_moe_stream_wave_ids, all under mtx) + void plan_waves_locked(llama_moe_stream_layer & sl, const int32_t * ids, int64_t n); // wave 0: build the plan + void stage_wave_locked(std::unique_lock & lk, llama_moe_stream_layer & sl, int32_t w, uint32_t n_ids); // make wave w resident + preload next + void emit_wave_slots(llama_moe_stream_layer & sl, const int32_t * ids, int32_t * out, int32_t w, uint32_t n_ids, int64_t n_tok); // write the slot ids }; // callback of the id-remapping custom op inserted by build_moe_ffn void llama_moe_stream_remap(ggml_tensor * dst, const ggml_tensor * a, int ith, int nth, void * userdata); + +// callbacks of the multi-pass prefill custom ops inserted by build_moe_ffn when a ubatch touches +// more experts than the cache holds; each src[0] is the contiguous selected ids +// wave_ids: makes wave w's expert slice resident and emits slot ids (masked pairs park on a pool) +// wave_mask: emits 1.0 for pairs belonging to wave w, 0.0 otherwise +void llama_moe_stream_wave_ids (ggml_tensor * dst, int ith, int nth, void * userdata); +void llama_moe_stream_wave_mask(ggml_tensor * dst, int ith, int nth, void * userdata);