llama : stream MoE routed experts from disk - #25294
Conversation
…RECT) 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 <N|NGiB>, --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
|
Hi @freedomljc, thanks for your contribution! Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:
Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below. |
Assisted-by: Claude
|
This saves system ram, but since it overrides cpu-moe and has no cooperation or hybridization with it, using this also seems to mean you need to have a lot of Vram to make up for the heavy loss of ssd streaming. |
It's primarily for the PCs with unified memory (e.g.: mac and dgx spark), where vram and system ram are in the same pool. The hybridization idea of using all three tiers make sense, probably we can tackle it as a follow-up. |
|
Hi @CISC @ggerganov , when you get chance, could you take a look? |
|
What is the performance versus mmap ? |
|
Can the number of experts cached per layer be configured based on the amount of RAM? |
|
VRAM-RAM-NVMe—this three-tier caching setup will be very interesting. |
Pure mmap with full GPU offload is impossible to run this model in GB10 as 240 GB would be required. I did the test for
|
It's supported by setting
Yes, it's an interesting idea, but I'd like to defer it in the separate PR. |
|
Gentle Ping. @CISC @ggerganov @pwilkin Can any one of maintainers take a look? |
|
There are other patches out there, which set up expert RAM<>VRAM swapping according to "expert hotness", or even Disk<>RAM<>VRAM tiered swapping. Would strongly suggest not merging until the others are evaluated, and the best long-term solution or most flexible base swapping architecture for future expansion in this area has been identified. It's going to become more important in future as models quickly ramp up to 1T, 2T, 5T, and beyond. |
|
This is actually really good |
|
Nice work. I've been doing basically the same thing (dense resident, routed experts streamed with O_DIRECT, bit-exact) but on Android phones over UFS, out-of-tree on the public eval-callback so no fork. One thing from the mobile side that might be useful for @lee-b's point about picking a flexible base: the O_DIRECT-vs-page-cache call is even sharper on a phone, because the kernel reclaims hard and you can't pin the dense weights to defend them. Not suggesting my approach over in-tree, just a data point that the design holds down to phone-class hardware. Repo if useful: |
|
PR #24524 has a good overview of previous attempts at VRAM cache solutions and a link to the discussion #24528, where also a 3 tier solution is mentioned. I agree that VRAM cache and MoE Streaming from Disk ("RAM Cache") could be implemented separately, but I also agree with @lee-b that it makes sense to see the bigger picture first to have compatible solutions. I would love to see both implemented as they deliver nice performance gains for MoE models. |
This sounds like it could be applied to the current system that loads host RAM->VRAM during prefill? |
Just for fun, I made a fork of this PR using Vibe (sorry, I used C++ in my student days last time) that loads some of the layers onto an RAM. That is, SSD -> VRAM + RAM -> VRAM. On the vanilla version of llama.cpp I was getting about 1 t/s. My fork is just a rough implementation to test a hypothesis, but if you're interested: |
Dspark MTP would be interesting to see if it can increase the t/s, the max I can get is 3.5 t/s 32GB, RTX 5090, PCIe5 NvME 14 Gb/ps |
…etal/Apple adaptations) Upstream PR ggml-org#25294 by Junchao Lyu, rebased onto master and adapted for Metal/Apple, plus a prefix-reuse fix in the server. Streamed layers do not materialize their ffn_*_exps tensors. Each weight gets a device-side cache of n_slots expert slabs, filled on demand from the GGUF by a CPU custom op that runs right after the router top-k and remaps expert ids to cache slots. The remap never changes which experts the router picked, so streaming affects latency, not outputs. A pool of I/O threads loads misses while the op waits; eviction is decaying route hotness with an LRU tiebreak. This makes a 104 GB model usable on 64 GB: 49 t/s prefill, ~7-9 t/s decode. Apple/Metal adaptations on top of the PR: - F_NOCACHE instead of O_DIRECT (Darwin has no O_DIRECT). Kept behind --moe-stream-direct and OFF by default: measured 20% SLOWER decode, because prefill's buffered reads warm the page cache that decode then hits. Any scheme making prefill "considerate" of decode is backwards here. - staging buffers are page-aligned for Metal private-buffer uploads, not just for direct I/O - TENSOR_STREAMED moved to bit 5; upstream took bit 4 for TENSOR_ALLOW_RESHAPE Server: -cms/--checkpoint-min-step adds a third context checkpoint at 4 + checkpoint_min_step when cms < n_ubatch. Upstream keeps only {4 + n_ubatch, 4}, so prefix-reuse granularity equals the ubatch and a follow-up turn re-prefills up to a full ubatch no matter what. Opt-in: the default cms is 8192 >= n_ubatch, leaving upstream behaviour unchanged. Observability that the tuning depended on and that a rebase should keep: LLAMA_MOE_STREAM_STATS_MS dumps per-window stall / cpu-op / gpu shares and miss and preload counts. It is what distinguishes an I/O regression from a compute one, and it lives in new files, so it costs nothing at rebase time. Note llama-server filters library LLAMA_LOG_INFO - diagnostics that must be visible there are WARN. Co-Authored-By: Junchao Lyu <ljc0711@gmail.com>
|
Prefill performance is now much improved (from 118.6 to 279.3 t/s in my case). master...ServeurpersoCom:llama.cpp:moe-stream-partition Memory view: every expert lives in exactly one place, pinned in VRAM or mirrored in pinned RAM, so the SSD is only read once at load and eviction never writes anything back.
Timeline view: within each of the 43 layers of a token, the router names its experts right before the FFN needs them, so the miss copies slip in behind the shared expert and only the leftover wait costs time.
Each layer has its own router and its own resident set, so within a single token every layer is an independent draw of 6 experts out of 256, and with about half of them resident per layer, one layer can land 3 misses and the next one all hits, pure luck of the draw repeated 43 times. Config[MoE-DeepSeek-V4-Flash-0731-284B-A13B-Q8_K_XL-POC] My config on an RTX PRO 6000 (96 GB) with 96 GB of DDR5-6600: the only knob is the VRAM expert budget, the pinned host mirror sizes itself automatically to whatever the cache does not pin (about 82 GiB here), and an optional --moe-stream-ram only exists to cap it on shared machines; in my case this gives +77% prefill and +50% decode vs. --n-cpu-moe, the big open question is how it behaves on more modest setups, and since the design only assumes VRAM + RAM >= experts, it should scale down proportionally with smaller models on smaller memory, which would be a strong upgrade for any config currently relying on --n-cpu-moe. |
|
Correction: not directly, each layer's router consumes that layer's hidden state, not the predicted token itself, so an MTP token prediction alone cannot tell us which experts the 43 layers will route to. |
|
@ServeurpersoCom so its similar but different to #26824 ? |
|
I'll have to take a look; I'm not planning to submit a PR, though, since this is just an experiment out of curiosity unless we come up with something really simple and generalizable that could be useful to everyone. |
Similar goal, different mechanism: #26824 computes cold experts on the CPU (hence the PP drop testers report), my branch keeps all compute on the GPU and streams the missing weights over PCIe from a pinned host mirror, so prefill gets faster too (+77% vs --n-cpu-moe). And it stays static by design: no heatmap, no persistence, just a pinned partition plus a small dynamic pool. |
Port of upstream PR ggml-org#25294 (SSD expert streaming: slot state machine, id remap, I/O workers, O_DIRECT reads, prefill waves), extended with a pinned expert partition in VRAM that is never evicted and a pinned host mirror (--moe-stream-ram) serving the misses at full PCIe speed instead of the model file. Co-authored-by: freedomljc <freedomljc@users.noreply.github.com>
Port of upstream PR ggml-org#25294 (SSD expert streaming: slot state machine, id remap, I/O workers, O_DIRECT reads, prefill waves), extended with a pinned expert partition in VRAM that is never evicted and a pinned host mirror (--moe-stream-ram) serving the misses at full PCIe speed instead of the model file. Co-authored-by: freedomljc <freedomljc@users.noreply.github.com>
Port of upstream PR ggml-org#25294 (SSD expert streaming: slot state machine, id remap, I/O workers, O_DIRECT reads, prefill waves), extended with a pinned expert partition in VRAM that is never evicted and a pinned host mirror (--moe-stream-ram) serving the misses at full PCIe speed instead of the model file. Co-authored-by: freedomljc <freedomljc@users.noreply.github.com>
Port of upstream PR ggml-org#25294 (SSD expert streaming: slot state machine, id remap, I/O workers, O_DIRECT reads, prefill waves), extended with a pinned expert partition in VRAM that is never evicted and a pinned host mirror (--moe-stream-ram) serving the misses at full PCIe speed instead of the model file. Co-authored-by: freedomljc <freedomljc@users.noreply.github.com>
Port of upstream PR ggml-org#25294 (SSD expert streaming: slot state machine, id remap, I/O workers, O_DIRECT reads, prefill waves), extended with a pinned expert partition in VRAM that is never evicted and a pinned host mirror (--moe-stream-ram) serving the misses at full PCIe speed instead of the model file. Co-authored-by: freedomljc <freedomljc@users.noreply.github.com>
Port of upstream PR ggml-org#25294 (SSD expert streaming: slot state machine, id remap, I/O workers, O_DIRECT reads, prefill waves), extended with a pinned expert partition in VRAM that is never evicted and a pinned host mirror (--moe-stream-ram) serving the misses at full PCIe speed instead of the model file. Co-authored-by: freedomljc <freedomljc@users.noreply.github.com>
Port of upstream PR ggml-org#25294 (SSD expert streaming: slot state machine, id remap, I/O workers, O_DIRECT reads, prefill waves), extended with a pinned expert partition in VRAM that is never evicted and a pinned host mirror (--moe-stream-ram) serving the misses at full PCIe speed instead of the model file. Co-authored-by: freedomljc <freedomljc@users.noreply.github.com>
|
Now that PLE and N-gram Embedding can be streamed from disk (#27794), the next giant step would be Expert Streaming/Caching for sure. We've seen so many different implementations, branches and PRs. Maybe it's time to centralize the comparison of different architectures under one Issue? The maintainers can then pick the most maintainable/portable/extensible idea. Anyway, good job to everyone working on these things. You're all heroes. |
Adds --moe-stream and --cache-mb, and two hooks in the server's load_model, so a MoE model whose weights do not fit in RAM can be served by this server rather than by a second one written around the streamer. Everything the streamer needs already travels in common_params -- the load mode, the extra-buffer-types flag, n_gpu_layers, and the eval callback pair -- so model and context creation stay a single common_init_from_params call. Phase one fills those in; phase two runs once the context exists, because a streamer learns which tensors are the expert weights by watching a graph go past the eval callback and cannot be armed before there is something to decode with. Upstream's warm-up is turned off in that case: the streamer does its own, and upstream's would run first with the callback installed and nothing behind it. The whole thing is behind BMOE_STREAM_IN_SERVER. Without that define a build is byte-for-byte this branch's parent, and asking for --moe-stream fails at load with a message rather than being accepted and ignored. Measured with BigMoeOnEdge's streamer on Qwen3.8-Flash-Next UD-Q5_K_XL, 158 GB of weights on a 123 GB machine: the model loads, the server answers, and a tools request comes back with finish_reason tool_calls, empty content and the call in ten incremental deltas with the reasoning kept in reasoning_content -- all of which the server already did, and none of which the streamer had to know about. Note for whoever finds this branch later: upstream PR ggml-org#25294 has been implementing expert streaming inside llama.cpp since July, with waved prefill and an argument against mmap. If it lands, this commit and the external streamer it exists for are both unnecessary.
|
Adopted this PR on a private fork for multi-model testing (CPU inference, Qwen3-Coder-30B-A3B Q4_K_M, temp-0 parity protocol). Two findings:
|
Keeps the routed expert weights on disk and pages them into a fixed-size per-layer cache on demand, so a model whose experts do not fit in RAM still runs at a useful rate. Based on PR ggml-org#25294 with Metal/Apple adaptations. --moe-stream enable --moe-stream-cache budget in GiB, or exact slots per layer with an 's' --moe-stream-io-threads reader threads --moe-stream-direct O_DIRECT reads, falling back to buffered Rebase note: upstream took TENSOR_READ_LAZY on bit 5 for its own on-demand row reads, so TENSOR_STREAMED moves to bit 6. The two are not equivalent - READ_LAZY reads rows through mmap, STREAMED means moe-stream owns the tensor's I/O entirely and the loader must not allocate or read it.
Keeps the routed expert weights on disk and pages them into a fixed-size per-layer cache on demand, so a model whose experts do not fit in RAM still runs at a useful rate. Based on PR ggml-org#25294 with Metal/Apple adaptations. --moe-stream enable --moe-stream-cache budget in GiB, or exact slots per layer with an 's' --moe-stream-io-threads reader threads --moe-stream-direct O_DIRECT reads, falling back to buffered Rebase note: upstream took TENSOR_READ_LAZY on bit 5 for its own on-demand row reads, so TENSOR_STREAMED moves to bit 6. The two are not equivalent - READ_LAZY reads rows through mmap, STREAMED means moe-stream owns the tensor's I/O entirely and the loader must not allocate or read it.
Keeps the routed expert weights on disk and pages them into a fixed-size per-layer cache on demand, so a model whose experts do not fit in RAM still runs at a useful rate. Based on PR ggml-org#25294 with Metal/Apple adaptations. --moe-stream enable --moe-stream-cache budget in GiB, or exact slots per layer with an 's' --moe-stream-io-threads reader threads --moe-stream-direct O_DIRECT reads, falling back to buffered Rebase note: upstream took TENSOR_READ_LAZY on bit 5 for its own on-demand row reads, so TENSOR_STREAMED moves to bit 6. The two are not equivalent - READ_LAZY reads rows through mmap, STREAMED means moe-stream owns the tensor's I/O entirely and the loader must not allocate or read it.
|
@ServeurpersoCom can you rebase your patches? I wanted to test them with qwen 3.8 flash next. |
…gml-org#25294 ggml-org#26003 ggml-org#26414 ggml-org#27402 ggml-org#20596 ggml-org#22671) and the CPU-kernel cherry-picks Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ayer (GroveMoE chunk experts) build_moe_ffn decided pool participation by layer_pooled(il). GroveMoE builds its chunk experts (ffn_*_chexps) through the same path on the same layer, so on a pooled layer they would have been remapped onto the pool's slots and bound to the layer's ids leaves. tensor_pooled(il, down_exps) keeps the chunk group on the ordinary streaming path. (Found in the upstream ggml-org#25294 triage; no GroveMoE model at hand, so verified by construction and the q35 gates only.) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…rrent mainline Port of the streaming machinery from PR ggml-org#25294 (feat/moe-streaming-core), cherry-picked onto current master with conflicts resolved. Skeleton scope: - llama_moe_stream module verbatim (cache slots, id-remap callbacks, wave prefill helpers, eviction: decaying route hotness + LRU) - params/CLI plumbing (--moe-stream, --moe-stream-cache, --moe-stream-io-threads) - TENSOR_STREAMED loader routing + buft selection helpers - context guards (op_offload disable when cache host-resident, graph_max_nodes wave budget) Deferred (next commit): llama-graph.cpp/h injection of the remap + wave custom ops into build_moe_ffn (14 commits of drift vs the PR era; same function shape verified by inspection). Disk/O_DIRECT path of the PR will be replaced by a RAM->device copy pool in the adaptation commit.
…itioned prefill) Ports the graph side of PR ggml-org#25294 onto current mainline: - build_lora_mm_id gains ids_scale (w_s gathered with original expert ids when the GEMM ids are remapped cache slots) - build_moe_ffn: msl lookup + matches() guard, single-pass id-remap op (llama_moe_stream_remap via ggml_map_custom1), and multi-pass wave prefill (llama_moe_stream_wave_ids/mask via ggml_custom_4d, per-wave GEMMs summed) when a ubatch touches more experts than the cache holds; wave cap = (n_slots - n_expert_used)/2 with 3*n_expert_used slot minimum (abort) - expert GEMM pipeline wrapped in build_expert_gemms lambda; biases and per-expert scales keep original ids throughout Compiles against current mainline (282/282). Runtime behavior unchanged unless --moe-stream is set.
… documented) --moe-stream-window <W>: pool of W full-layer expert slots, prefetched deterministically (layer order) — no expert-slot floor, so big-expert models (122B-A10B) can engage GPU prefill on 12 GB. Status: builds and runs; copies verified byte-correct (memcmp) and ids verified identity; output is still wrong — root cause not isolated in-session. Top suspects: the scheduler's handling of pool tensors shared across layers, or the map_custom1 remap op interacting with the shared weights. Expert-slot mode (PR ggml-org#25294) is unaffected (re-verified byte-identical). Debug trail in run-log \u00a714: warmup-stale-slot theory disproven (identity load is routing-independent), remap bypass disproven (pool never loads), copy and ids both verified correct.
Port of upstream PR ggml-org#25294 (SSD expert streaming: slot state machine, id remap, I/O workers, O_DIRECT reads, prefill waves), extended with a pinned expert partition in VRAM that is never evicted and a pinned host mirror (--moe-stream-ram) serving the misses at full PCIe speed instead of the model file. Co-authored-by: freedomljc <freedomljc@users.noreply.github.com>
Rebased/updated on my fork |


Overview
n_slotsexpert slabs per layer.demand-loaded from the GGUF by an async I/O worker pool.
when the model far exceeds RAM: the page cache cannot help and otherwise thrashes competing for the
same memory.
expert GEMMs run in W waves of at most (n_slots - n_expert_used)/2 experts each; the pairs of the
other waves are masked to zero and the wave outputs summed. Each touched expert is loaded once per
ubatch, so the previous n_ubatch clamp is removed and long prompts prefill much faster.
Additional information
Testing / validation
Usage / CLI
--moe-stream-cache <NG|Ns>(GiB budget, orssuffix = slots; implies--moe-stream).--moe-stream-io-threads N,--moe-stream-direct(O_DIRECT).model into RAM and defeat streaming.
Benchmark data
GB10 (Grace-Blackwell, 128 GB unified, PCIe 4.0 SSD), GLM-5.2-UD-Q2_K_XL (~254 GB file, ~754 B params,
256 experts),
-ngl 99 --moe-stream-cache <#cache> --moe-stream-direct -c 4096, greedy, 512-token generation.Known limitations
llama_contextfrom the same streamedmodel shares one cache and can corrupt output.
--parallel Nwithin a SINGLE context is safe(all sequences batched into one graph; remap reserves all needed slots first).
Related discussion
Requirements