llama : stream MoE routed experts from disk

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"}, "<NG|Ns>",
+        "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",

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 {

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

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

include/llama.h

@@ -318,6 +318,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
@@ -327,6 +334,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 {
@@ -1555,6 +1563,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
     //

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

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

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,24 @@ 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) {
+        // 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 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());
+        }
+        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;
 
@@ -2321,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<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
+        res = std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
+    } else {
+        res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
+        for (const auto & lora : model.loras) {
+            res += lora->get_n_nodes();
+        }
     }
-    uint32_t res = std::max<uint32_t>(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<uint32_t>(cap, 1);
+        const uint32_t n_touch_max = std::min<uint32_t>(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;
 }
@@ -2415,6 +2445,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 +4125,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) {