ggml : --prefetch-experts-slots — lookahead H2D prefetch of host-resident MoE experts

common/arg.cpp

@@ -1768,6 +1768,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
                                    string_format("error: unknown value for --flash-attn: '%s'\n", value.c_str()));
                            }
                        }).set_env("LLAMA_ARG_FLASH_ATTN"));
+    add_opt(common_arg(
+        {"--prefetch-experts-slots"}, "N",
+        "MoE expert H2D staging slots (default 0 = off; N>=2 enables full-tensor prefetch with 1-deep lookahead for host-resident/ncmoe-offloaded expert weights during prefill; recommended 3; capped at 4). GPU memory cost = N x max_expert_tensor.",
+        [](common_params & params, const std::string & value) {
+            params.prefetch_experts_slots = std::stoi(value);
+        }
+    ));
     add_opt(common_arg(
         {"-p", "--prompt"}, "PROMPT",
         "prompt to start generation with; for system message, use -sys",

common/common.cpp

@@ -1746,6 +1746,7 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.offload_kqv       = !params.no_kv_offload;
     cparams.no_perf           = params.no_perf;
     cparams.op_offload        = !params.no_op_offload;
+    cparams.prefetch_experts_slots = params.prefetch_experts_slots;
     cparams.swa_full          = params.swa_full;
     cparams.kv_unified        = params.kv_unified;
 

common/common.h

@@ -497,6 +497,7 @@ struct common_params {
     enum llama_pooling_type      pooling_type      = LLAMA_POOLING_TYPE_UNSPECIFIED; // pooling type for embeddings
     enum llama_attention_type    attention_type    = LLAMA_ATTENTION_TYPE_UNSPECIFIED; // attention type for embeddings
     enum llama_flash_attn_type   flash_attn_type   = LLAMA_FLASH_ATTN_TYPE_AUTO; // whether to use Flash Attention
+    int prefetch_experts_slots = 0; // --prefetch-experts-slots: MoE expert H2D staging slots (0 = off, >=2 enables full-tensor lookahead prefetch of ncmoe-offloaded weights; recommended 3; capped at 4)
 
     struct common_params_sampling    sampling;
     struct common_params_speculative speculative;

ggml/include/ggml-backend.h

@@ -353,6 +353,15 @@ extern "C" {
     // Set a callback to be called for each resulting node during graph compute
     GGML_API void                 ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data);
 
+    // Configure full-tensor MoE expert prefetch on a scheduler (mindcontrol port of
+    // --prefetch-experts-slots). Only engages for splits whose MUL_MAT_ID weights are
+    // host-resident (GGML_BACKEND_BUFFER_USAGE_WEIGHTS, e.g. --n-cpu-moe) during large
+    // prefill batches; decode is unaffected (batch ids < 2*n_expert).
+    //   slots == 0  -> prefetch disabled (no memory overhead)
+    //   slots >= 2  -> prefetch enabled with 1-deep lookahead and per-split cross-stream wait;
+    //                  GPU staging cost = slots * max_expert_tensor. Capped at GGML_SCHED_MAX_PREFETCH_SLOTS.
+    GGML_API void                 ggml_backend_sched_set_prefetch_experts_slots(ggml_backend_sched_t sched, int slots);
+
     //
     // Meta backend
     //

include/llama.h

@@ -407,6 +407,12 @@ extern "C" {
                           // try to disable when n_seq_max > 1 for improved performance when the sequences do not share a large prefix
                           // ref: https://github.com/ggml-org/llama.cpp/pull/14363
 
+        // mindcontrol-port of --prefetch-experts-slots: MoE expert H2D staging slots.
+        // 0 = off (no memory overhead); >=2 enables full-tensor lookahead prefetch of
+        // host-resident (ncmoe) expert weights during prefill (GPU staging cost = slots
+        // * max expert tensor). Decode is unaffected.
+        int prefetch_experts_slots;  // set via llama_context_default_params() / llama_context_from_params
+
         // [EXPERIMENTAL]
         // backend sampler chain configuration (make sure the caller keeps the sampler chains alive)
         // note: the samplers must be sampler chains (i.e. use llama_sampler_chain_init)

src/llama-context.cpp

@@ -270,6 +270,7 @@ llama_context::llama_context(
     }
 
     cparams.op_offload = params.op_offload;
+    cparams.prefetch_experts_slots = params.prefetch_experts_slots;
     cparams.kv_unified = params.kv_unified;
 
     // initialized later
@@ -603,6 +604,7 @@ void llama_context::sched_reserve() {
     gf_res_reserve.reset(new llm_graph_result(max_nodes));
 
     sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload));
+    ggml_backend_sched_set_prefetch_experts_slots(sched.get(), cparams.prefetch_experts_slots);
 
     llama_memory_context_ptr mctx;
     if (memory) {
@@ -638,6 +640,7 @@ void llama_context::sched_reserve() {
                 LLAMA_LOG_WARN("%s: compute buffer allocation failed, retrying without pipeline parallelism\n", __func__);
                 cparams.pipeline_parallel = false;
                 sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, cparams.op_offload));
+                ggml_backend_sched_set_prefetch_experts_slots(sched.get(), cparams.prefetch_experts_slots);
                 gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get());
             }
             if (!gf) {
@@ -3649,6 +3652,7 @@ llama_context_params llama_context_default_params() {
         /*.op_offload                  =*/ true,
         /*.swa_full                    =*/ true,
         /*.kv_unified                  =*/ false,
+        /*.prefetch_experts_slots      =*/ 0,
         /*.sampler                     =*/ nullptr,
         /*.n_sampler                   =*/ 0,
         /*.ctx_other                   =*/ nullptr,

src/llama-cparams.h

@@ -54,6 +54,8 @@ struct llama_cparams {
     bool kv_unified;
     bool pipeline_parallel;
 
+    int prefetch_experts_slots = 0; // --prefetch-experts-slots: MoE expert H2D staging slots (0 = off, >=2 enables full-tensor lookahead prefetch; capped at 4)
+
     std::vector<bool> embeddings_layer_inp; // [n_layer()] extract input embeddings for layer
 
     enum llama_context_type ctx_type;