llama: GPU-resident LRU cache for host-offloaded MoE expert weights - #27861
llama: GPU-resident LRU cache for host-offloaded MoE expert weights#27861csantiago78 wants to merge 1 commit into
Conversation
…ights Measured on Qwen3.8-Flash-Next UD-Q4_K_XL (512 experts, 10 routed, 28 expert layers pinned to host RAM by -ot): expert routing has strong temporal locality (LRU-64 ~67% hit rate over a mixed workload) even though the long-run expert distribution is near-uniform, so a per-layer LRU cache of expert slices in VRAM removes most of the per-token host-RAM streaming that bounds decode. Mechanism (no custom CUDA kernels): - companion tensors [ne0, ne1, K+1] per cached layer in the device buffer of that layer's router; slot K stays all-zero (dummy) - I32 id->slot tables, one device copy (get_rows remaps ids for a second mul_mat_id chain over the cache) and one host copy (src[3] of the CPU mul_mat_id, which skips cached ids and zeroes their dst rows) - the two down outputs are summed; the split is exact by construction - decode-only (n_tokens == 1); batch/prefill builds the stock graph - throttled async uploads: evictions are published at a decode-boundary sync point, slices are copied by a worker thread, and the new mapping is only published after the upload completed, so a running graph never reads a torn slot Enable with LLAMA_MOE_CACHE_SLOTS=<K> (+ LLAMA_MOE_CACHE_INSERTS, _DEBUG). Inert without the env var. Measured decode, this box (2x3090, 1 RAM channel/socket, with numactl --interleave=all): 15.5 -> 19-20 tok/s warm at K=48-64. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A7KRyLzuBiGfVn3bdczjka
|
Hi @csantiago78, 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. |
|
Runtime verification of the CLI plumbing (previously compile-checked only): built this branch with CUDA and ran llama-server with |
|
On Qwen3.8-Flash-Next UD-Q4_K_XL on a single R9700 with DDR5 (at 3600MT/s), this increases tg from 12.8 to 14.8 (+15%) in my testcase with
Is this additional VRAM usage considered by edit It is not considered during |
|
Gave it some more testing.
That's 40% faster than the baseline and |
|
There's an interesting suboptimal behaviour though. When processing a bigger prompt, a lot of VRAM gets freed (32GB -> 22GB). Can be reproduced with my configuration above and I can't see anything relevant in the logs with |
| // map to the cache's zero slot. The two outputs sum to the exact result. | ||
| const llama_moe_cache_layer * mcache = nullptr; | ||
| ggml_tensor * mc_slot_ids = nullptr; | ||
| if (n_tokens == 1 && !gate_up_exps && gate_exps && down_exps && |
There was a problem hiding this comment.
I think the n_tokens == 1 here means that this will basically never trigger if there's a draft model. You'd (usually) validate N tokens per round and generate 1 in parallel.
There was a problem hiding this comment.
You actually mentioned that in the initial comment, nevermind
| const int64_t n_ids = ids->ne[0]; | ||
| const int64_t n_tokens = ids->ne[1]; | ||
| if (n_tokens > 4) { | ||
| return; // batch/prefill: the cache graph is not built there, don't pollute the LRU |
There was a problem hiding this comment.
In prefill/batch you could still update the LRU and once prefill is done, sort the prefill experts by frequency and upload the top N or the most recent ones. That improves cache-hit ratio a bit for me after bigger prefills (where otherwise the cache is mostly cold).
0001-moecache-warm-cache-from-prefill-expert-frequency.patch implements something like this if you want to take that as reference. It was mostly an experiment to see if it makes a bigger difference (it only makes a small difference but there's probably more fine-tuning possible).
| j = mc->todo.front(); | ||
| mc->todo.pop_front(); | ||
| } | ||
| auto & ls = mc->layers[j.layer_idx]; |
There was a problem hiding this comment.
There's an interesting suboptimal behaviour though. When processing a bigger prompt, a lot of VRAM gets freed (32GB -> 22GB). Can be reproduced with my configuration above and
llama-cli. When processing small prompts it's all fast, when e.g. pasting a lot of text or using/readon a bigger file, VRAM becomes a lot more empty and tg becomes slow.
That doesn't cause actual performance problems but the problem here is that with the MoE cache the expert weights are changing backends. Next time the graph needs to be re-allocated (big prompt), it might need less memory because of that, and then reallocation shrinks memory usage. And there's then repeated re-allocations happening regularly.
This can be avoided by only re-allocating if the new size doesn't fit in the previously allocated but that all doesn't seem ideal.
| if (ls.slot_last_use[s] < best) { best = ls.slot_last_use[s]; slot = s; } | ||
| } | ||
| if (slot < 0) { | ||
| break; // every slot is in flight; try again next step |
There was a problem hiding this comment.
AFAIU if uploading is constantly slower than decoding, at some point all slots are in flight and old experts are still uploaded while newer ones rarely end up in the cache, which should then reduce the hit rate considerably. Maybe this requires some other approach?
| // map to the cache's zero slot. The two outputs sum to the exact result. | ||
| const llama_moe_cache_layer * mcache = nullptr; | ||
| ggml_tensor * mc_slot_ids = nullptr; | ||
| if (n_tokens == 1 && !gate_up_exps && gate_exps && down_exps && |
There was a problem hiding this comment.
Also related to this, by doing this you have different graphs (and number of nodes) between decode and prefill, and that can cause reallocs all the time. When doing the first decode with n_tokens==1 with a different topology, a realloc happens that uses the new small n_kv. Then over time the KV cache fills up, QSA inputs that are depending on n_kv grow above their reserved limits and then there's a realloc every 250-300 decode steps (which wouldn't happen if the prefill allocation was kept, for example).
|
Some measurements on this, on Blackwell workstation cards with a 180B qwen4_exp model. I ran everything on the same file and the same card so the numbers can be compared to each other. Setup
Throughput is token generation only ( Throughput
Tuned is The stock number sits inside the range already reported here (12.8 → 18.0 on a single R9700), which I think is the useful part: on default settings this machine is not faster than anyone else's. My build carries changes beyond this PR, so the middle row exists to keep those separate from the tuning. Measured hit rate I captured a routing trace during normal decode (
The cache does better than the LRU-64 ≈ 67% / LRU-128 ≈ 81% figures used here so far, and it saturates around 384 slots — the residual 1.5% are cold-start misses. Batching does not reduce loads Miss experts per token:
Consecutive tokens share 44.2% of their experts, so batching looks like it should help. It does not, because the cache already covers that reuse over a much longer window: Two practical notes Slot counts are a memory budget, not a model property. 144 slots fit my own 119 GB conversion of this model; only 96 fit the 104 GB
That conversion does not carry the MTP head. My own conversion does, and MTP speculation is worth about +14% there. So the 41.31 above is the best I get on this file, not the best I get. Caveat: |
|
Tested this PR at commit bccbacd on Windows with an AMD Radeon RX 7600 8GB using the Vulkan backend. Hardware: Ryzen 9 5900X Build: GCC 15.2.0 I used --cpu-moe with -ngl 99, and compared decode throughput with and without --moe-expert-cache. Results: GigaChat-20B-A3B Q4_K_M I also tested Qwen3-30B-A3B with a 16K context, q8_0 KV cache and Flash Attention: 16K / q8 KV / 32 expert-cache slots: 16.5 tok/s So on this RX 7600 system, increasing the context from 4K to 16K did not reduce decode throughput in this configuration. I also confirmed that the same expert-cache options work through the built llama-server, with /v1/chat/completions producing ~16.2 tok/s on the Qwen3 30B setup. One build note: I needed two local Windows/toolchain workarounds: guard flockfile / funlockfile in the diagnostic MoE logging code on MinGW Neither change touches the MoE expert-cache logic itself. Overall, the cache appears to be working reliably on AMD RDNA3/Vulkan/Windows, and the speedup is significant on larger MoE models. Thanks for working on this. This did a great job. |
|
I tested the n_tokens == 1 limitation mentioned above with an embedded-MTP Qwen3.6 MoE model and hit the expected issue: speculative validation builds small multi-token decode graphs, so the expert cache is bypassed. The GGUF I tested uses the existing separate gate_exps / up_exps / down_exps path supported by this PR. This change does not add support for fused gate_up layouts or otherwise broaden the existing MoE-layout compatibility. I tried extending the existing cache remap to small decode batches (n_tokens <= 4) by making selected_experts contiguous, flattening it to n_expert_used * n_tokens for ggml_get_rows(), then reshaping the resulting slot ids back to [n_expert_used, n_tokens]. The core change is: ggml_tensor * mc_selected_experts = ggml_cont(ctx0, selected_experts); together with changing the existing cache guard from n_tokens == 1 to n_tokens > 0 && n_tokens <= 4. I built and runtime-tested this with Qwen3.6-35B-A3B using embedded MTP (--spec-type draft-mtp --spec-draft-n-max 2), and the cache path now runs successfully during MTP decode. If useful, I can provide the complete small patch. One additional data point: this wasn't only a runtime smoke test — combining MTP with the expert cache produced a substantial end-to-end speedup on this model. On the same Qwen3.6-35B-A3B setup, I measured roughly: plain decode, no MTP / no expert cache: ~17.1 tok/s So being able to use the cache during the small validation batches is material here; the combined configuration is roughly 1.8–1.9x the plain baseline. |
|
Duplicate slot ids per token break the CUDA batched mul_mat_id path (n_tokens > 1) It looks like every expert that hasn't been cached to VRAM has been mapped to the same dummy slot (n_slots), which means that a token's slot-id row can contain the same id multiple times. The CUDA batched kernels behind mul_mat_id assume that the ids of a token are going to be distinct, so the duplicates collapse into one entry. At batch 1 with quantized experts, this isn't shown because it uses a different kernel, mmvq, which simply asks for each output row which expert it needs and fetches it. Duplicates are harmless here Reproduction without a model: with the help of an agent, I added dup_ids to test_mul_mat_id in tests/test-backend-ops.cpp (llama.cpp's built-in test program) that writes the highest expert index into every other id slot of each token. In short, since the test normally gives each token ten different expert numbers, the goal was to overwrite every second one with the highest expert index, which mimics what the cache does when it points every uncached expert of a token at the same dummy slot. On this tree (2x RTX 3090, CUDA 12.0, driver 595.84): F16 cases fail already at n_tokens 1 to 8 (see the safe window section below: mmvq is not used on F16 for NVIDIA). Q4_K and Q6_K pass at n_tokens 1 to 8, and the run aborts at the first case that reaches mmq, the batched quantized kernel that takes over above 8 tokens. For reference, here is where the assumption lives:
The safe window, measured: quantized experts take mmvq iff n_tokens <= min(MMVQ_MAX_BATCH_SIZE, get_mmvq_mmid_max_batch()), which is 8 for Q4_K, Q5_1 and Q6_K on cc 860 and on cc 1200. With the helper fixed, dup-id cases pass 6/6 at n_tokens 8 (three types, broadcast and not) and fail 18/18 at 9, 12 and 16. F16 and BF16 experts never take mmvq on NVIDIA. The diff (gist link below) fixes the helper by:
The diff also guards flockfile/funlockfile in ggml-cpu.c with #ifndef _WIN32 so the GGML_MOE_LOG debug logging builds on Windows. After the diff there is no CUDA error and every pre-existing MUL_MAT_ID case passes, but 24 dup-id cases still fail because of 2 (the F16 small-batch kernel) and 3 (the block-count sizing). The duplicate cases that pass through those paths still produce wrong numbers; they just no longer crash and instead fail the CPU comparison. So the diff isn't a solution, but it is a step. There are a few options I see right now that could be solutions:
Gist link: https://gist.github.com/Inovello/9500167e5e8dc98fabe0df0c62ba3489 |
|
Ran this on a config I don't think has been covered yet: Vulkan, two GPUs, Qwen3.8-Flash-Next UD-Q2_K_XL, on 2x RX 6950XT 16GB over OCuLink Gen4 x4 from a Ryzen 7 7840U handheld (64 GB RAM), Windows 11 / MSVC. It works well — thanks for building this. Reporting back with numbers, two small fixes, and two bugs I hit when I tried to extend it. ResultsDecode after cache warm-up,
The cache slots do add up across cards, as the description says. At full 262144 context, Two fixes1. #if defined(_MSC_VER)
_lock_file(moe_log_file);
#else
flockfile(moe_log_file);
#endif2. Letting the cache serve small batches makes speculative decoding a win instead of a loss. With the static const int64_t moe_cache_max_tokens = []() {
const char * e = getenv("LLAMA_MOE_CACHE_MAX_TOKENS");
return e ? (int64_t) atoi(e) : (int64_t) 1;
}();
if (n_tokens >= 1 && n_tokens <= moe_cache_max_tokens && ...) {
mcache = llama_moe_cache_lookup(up_exps);
}
if (mcache) {
ggml_tensor * mc_tbl = mcache->dev_table; // [1, n_expert]
if (n_tokens > 1) {
mc_tbl = ggml_repeat_4d(ctx0, mc_tbl, 1, mc_tbl->ne[1], n_tokens, 1);
}
mc_slot_ids = ggml_get_rows(ctx0, mc_tbl, selected_experts);
mc_slot_ids = ggml_reshape_2d(ctx0, mc_slot_ids, n_expert_used, n_tokens);
...( I verified this rather than assuming: greedy decoding ( Both patches, with full diffs: https://github.com/ChangXiang-SCU/dual-egpu-moe-llm/tree/main/patches Two bugs, if you want to take the cache furtherRaising that limit past 16 produces wrong output — a short arithmetic prompt gets misread, longer prompts emit
Two independent causes, as far as I could isolate them: (a) The Vulkan topk-moe fusion misfires. (b) The expert tables mutate mid-prefill. A long prompt spans many ubatches, and Worth fixing if you're inclined, because the prize is large — with the cache serving prefill batches, prefill went 17.4 → 30.5 tok/s on a 3722-token prompt (+75 %, and the gain grows with prompt length). Prefill is by far the dominant cost in this setup, since 100 % of it currently runs on the host CPU. Happy to re-run anything on this hardware if it's useful. Full measurements, correctness methodology and the Windows/Vulkan build notes are at https://github.com/ChangXiang-SCU/dual-egpu-moe-llm |
|
Huge thank you to @csantiago78 and the entire llama.cpp community — this LRU expert cache is exactly the kind of optimization that makes CPU-offloaded MoE actually usable. The 31% throughput improvement (18.4 → 24.2 tok/s) on Qwen3.8-Flash-Next is impressive, and the fact that it uses a second As someone following the broader NVMe/MoE offloading landscape, I wanted to highlight a convergence I'm seeing:
The community now has multiple independent implementations of the same core idea (MoE expert caching/streaming). At this point it feels less like competing approaches and more like different modules of a single system. Would love to see these converge. Thank you for the thorough implementation and benchmarks. This work — and llama.cpp itself — is what keeps local AI accessible. 🙏 |
Summary
GPU-resident LRU cache for MoE expert weights that live in host memory (via
-ot ...exps=CPU,-ncmoe, etc.). Decode on a host-offloaded MoE layer is bound by host RAM bandwidth: every token streams the routed experts' weights from system RAM. This PR serves the recently used experts from VRAM instead.Opt-in via
--moe-expert-cache N(slots per host-resident expert layer;--moe-expert-cache-insertscaps uploads per layer per decode step). Fully inert when disabled.Motivation / measurements
Measured on Qwen3.8-Flash-Next UD-Q4_K_XL (512 experts, 10 routed/token, 28 expert layers pinned to host by
-ot), 2x RTX 3090 + dual-Xeon host (single populated RAM channel per socket):Mechanism (no new CUDA kernels)
[ne0, ne1, K+1]for up/gate/down in the device buffer of that layer's router. SlotKis permanently zero (the "dummy" slot).expert id -> slottable per layer, two copies:ggml_get_rowsremapsselected_expertsinto slot ids for a secondmul_mat_idchain over the cache tensors. Uncached ids map to the zero slot and contribute exactly 0.src[3]to the CPUmul_mat_id, which skips cached ids and zeroes their dst rows.n_tokens == 1); batches/prefill build the stock graph, so batch offload is untouched.ggml_backend_tensor_set, and the new mapping is only published at a later sync point after the copy completed - a running graph can never read a torn slot. (Synchronous uploads were measured to eat the entire win.)Known discussion points (why this is a draft)
up_expstensor pointer, becausebuild_moe_ffnhas no model access. Happy to rework the ownership (e.g. hang it offllama_model) per your preference.mul_mat_idobservation callback (ggml_set_moe_obs_callback) is a ggml -> llama upcall; suggestions for a cleaner layering welcome.LLM_FFN_SILUpath is wired; other MoE variants fall back to the stock graph.n_tokens == 1guard); extending the remap to small batches is straightforward if the approach is acceptable.Testing
test-arg-parserargument section passes (the URL/404 section fails in my sandbox for network reasons, unrelated).