Add --moe-mode sweep-prefill: streamed expert supply for CPU mul_mat_id

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<std::string> 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",

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;
 

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

ggml/src/ggml-cpu/ggml-cpu.c

@@ -37,6 +37,10 @@
 #if defined(__gnu_linux__)
 #include <syscall.h>
 #endif
+#if defined(__linux__)
+#include <sys/mman.h>
+#include <unistd.h>
+#endif
 
 #ifdef GGML_USE_OPENMP
 #include <omp.h>
@@ -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
 }
 
 /////////////////////////////////

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)
     };

src/llama-model.cpp

@@ -28,6 +28,7 @@
 #include <cstring>
 #include <cmath>
 #include <filesystem>
+#include <set>
 #include <fstream>
 #include <functional>
 #include <map>
@@ -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<const ggml_tensor *> flash_moe_sweep_tensors;
     std::string flash_moe_trace_file;
     std::unordered_map<std::string, llama_flash_moe_sidecar_entry> flash_moe_sidecar_entries;
     std::vector<llama_flash_moe_sparse_mapping> 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<int32_t>(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,
     };

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;

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;

src/models/qwen35moe.cpp

@@ -2,6 +2,8 @@
 
 #include "llama-memory-recurrent.h"
 
+#include <cstring>
+
 static bool qwen35moe_experimental_metal_split_glu_enabled() {
     static int enabled = -1;
     if (enabled == -1) {