llama : stream MoE routed experts from disk

[freedomljc]

Overview

• Optional SSD streaming of MoE routed-expert weights so a model larger than RAM can run.
• Streamed layers keep a small device-side cache of n_slots expert slabs per layer.
• A CPU id-remap custom op after the router top-k maps expert ids -> cache slots; misses are
  demand-loaded from the GGUF by an async I/O worker pool.
• Reads use O_DIRECT to bypass the OS page cache - the key I/O optimization
  when the model far exceeds RAM: the page cache cannot help and otherwise thrashes competing for the
  same memory.
• Wave-Partitional Prefill: when a ubatch touches more experts than the cache holds, the
  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.
• Output matches a non-streamed run (bit-exact when both paths use the same kernels/ubatch).

Additional information

Testing / validation

• Streamed vs non-streamed output verified: bit-exact under matched kernels + ubatch.
• Validated on OLMoE Q4_K_M (fits in RAM; CPU + Metal) and GLM-5.2 (>>RAM; CUDA), incl. llama-server.

Usage / CLI

• --moe-stream-cache <NG|Ns> (GiB budget, or s suffix = slots; implies --moe-stream).
• --moe-stream-io-threads N, --moe-stream-direct (O_DIRECT).
• Enabling streaming AUTO-DISABLES mmap (with a warning) - mmap prefetch would page the whole
  model into RAM and defeat streaming.
• Requires a file-based MoE model (not stdin/fd).

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.

expert cache	prefill (no waves)	prefill (waved)	decode	cache hit rate
64 slots (~55 GB)	2.28 tok/s	5.65 tok/s	~1.83 tok/s	73%
90 slots (~79 GB)	2.88 tok/s	5.69 tok/s	~2.20 tok/s	79%

• Latency form: prefill ~625-637 ms/tok, decode ~430-507 ms/tok.

Known limitations

• Single-context only: concurrent decoding of multiple llama_context from the same streamed
  model shares one cache and can corrupt output. --parallel N within a SINGLE context is safe
  (all sequences batched into one graph; remap reserves all needed slots first).

Related discussion

• Feature Request: SSD streaming of MoE routed expert weights (run MoE models larger than system RAM) #25257

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES

[ggml-gh-bot]

Hi @freedomljc, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• Large PR: Large changes require prior discussion (e.g. an issue or RFC) and maintainers may not be able to review this PR as-is. Consider splitting it into smaller, focused PRs.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

[Midaychi]

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.

[freedomljc]

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.

[freedomljc]

Hi @CISC @ggerganov , when you get chance, could you take a look?

[Green-Sky]

What is the performance versus mmap ?

[rankaiyx]

Can the number of experts cached per layer be configured based on the amount of RAM?

[rankaiyx]

VRAM-RAM-NVMe—this three-tier caching setup will be very interesting.

[freedomljc]

What is the performance versus mmap ?

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 --cpu-moe + mmap.
Here's the perf comparision:

	streaming (O_DIRECT, cache 90s)	mmap + cpu-moe	ratio
prefill	5.58 tok/s	1.06 tok/s	5.3x
decode (512 tok)	2.08 tok/s	0.87 tok/s	2.4x

[freedomljc]

Can the number of experts cached per layer be configured based on the amount of RAM?

It's supported by setting --moe-stream-cache <N>G, and G is for GB.

VRAM-RAM-NVMe—this three-tier caching setup will be very interesting.

Yes, it's an interesting idea, but I'd like to defer it in the separate PR.

[freedomljc]

Gentle Ping. @CISC @ggerganov @pwilkin Can any one of maintainers take a look?
ssd-streaming on expert layper has been a popular idea for a while, e.g. https://github.com/danveloper/flash-moe, https://github.com/antirez/ds4

[lee-b]

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.

[Kuberwastaken]

This is actually really good

[Helldez]

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:

https://github.com/Helldez/BigMoeOnEdge

[xashr]

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.

[crusaderky]

Wave-Partitional Prefill: when a ubatch touches more experts than the cache holds, the
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.

This sounds like it could be applied to the current system that loads host RAM->VRAM during prefill?

[aldubl]

Предварительное заполнение с разделением по волнам: когда в одном пакете запросов затрагивается больше экспертов, чем вмещает кэш,
GEMM-ы экспертов запускаются волнами по W, каждая из которых содержит не более (n_slots - n_expert_used)/2 экспертов; пары других
волн обнуляются, а выходные данные волн суммируются. Каждый затронутый эксперт загружается один раз за
один пакет запросов, поэтому ограничение в n_ubatch снимается, и предварительное заполнение длинных подсказок происходит гораздо быстрее.

Похоже, это можно применить к существующей системе, которая загружает данные из оперативной памяти хоста в видеопамять во время предварительного заполнения?

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.
PR from the respected freedomljc: gen 1.90 t/s, pp 2.06 t/s.
My fork with 5 GB RAM: gen 1.98 t/s, pp 2.34 t/s.
My fork with 75 GB RAM: gen 2.52 t/s, pp 3.71 t/s.
Tested on Depseek v4 Flash Q8_K_XL.

My fork is just a rough implementation to test a hypothesis, but if you're interested:
https://github.com/aldubl/llama.cpp/tree/ssd-moe

[serdavid7]

Предварительное заполнение с разделением по волнам: когда в одном пакете запросов затрагивается больше экспертов, чем вмещает кэш,
GEMM-ы экспертов запускаются волнами по W, каждая из которых содержит не более (n_slots - n_expert_used)/2 экспертов; пары других
волн обнуляются, а выходные данные волн суммируются. Каждый затронутый эксперт загружается один раз за
один пакет запросов, поэтому ограничение в n_ubatch снимается, и предварительное заполнение длинных подсказок происходит гораздо быстрее.

Похоже, это можно применить к существующей системе, которая загружает данные из оперативной памяти хоста в видеопамять во время предварительного заполнения?

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. PR from the respected freedomljc: gen 1.90 t/s, pp 2.06 t/s. My fork with 5 GB RAM: gen 1.98 t/s, pp 2.34 t/s. My fork with 75 GB RAM: gen 2.52 t/s, pp 3.71 t/s. Tested on Depseek v4 Flash Q8_K_XL.

My fork is just a rough implementation to test a hypothesis, but if you're interested: https://github.com/aldubl/llama.cpp/tree/ssd-moe

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