Commit 951d465f5 for llama.cpp
commit 951d465f5c7c16fb486ff45bb14d99984e68c6db
Author: Aman Gupta <amangupta052@gmail.com>
Date: Fri Nov 13 18:48:29 2026 +0800
llama: move moe-cache code out of llama-graph (#32157)
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 3cee72df0..a47cb1e5d 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -2425,12 +2425,9 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
// compute the expert input before the slot lookup, so that the scheduler prefetches it for the CPU part (the router can read a different tensor)
ggml_build_forward_expand(gf, cur);
- // the experts of host-resident layers may be read from the MoE cache, and part of the misses computed from the host experts
- // the CPU part does not support weights before the FFN, LoRAs or rotated experts
- const std::array<ggml_tensor *, 4> exps = { up_exps, gate_exps, down_exps, gate_up_exps };
- const bool can_cpu_share = !weight_before_ffn && loras->empty() &&
- std::none_of(exps.begin(), exps.end(), [&](ggml_tensor * w) { return w != nullptr && hdmd != nullptr && hdmd->rot.count(w) > 0; });
- const llm_graph_moe_cache_slots cache = build_moe_cache_slots(selected_experts, up_exps, gate_exps, down_exps, gate_up_exps, can_cpu_share, il);
+ // the experts of host-resident layers may be read from the MoE cache (see llama_moe_cache::build_lookup)
+ const llama_moe_cache_lookup cache = moe_cache == nullptr ? llama_moe_cache_lookup{} :
+ moe_cache->build_lookup(*this, { il, up_exps, gate_exps, down_exps, gate_up_exps, weight_before_ffn }, selected_experts);
cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens);
@@ -2603,14 +2600,12 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
return experts;
};
- ggml_tensor * experts = build_experts(cur, selected_experts, cache.slots, true);
- if (cache.cpu_ids != nullptr) {
- // the cached experts and the misses computed from the host experts, each reads the others as zero
- // the cached experts come first in the graph, so that they are computed while the misses are
- ggml_build_forward_expand(gf, experts);
- experts = ggml_add(ctx0, experts, build_experts(cur, selected_experts, cache.cpu_ids, false));
- cb(experts, "ffn_moe_down_merged", il);
- }
+ ggml_tensor * experts = cache.slots == nullptr ?
+ build_experts(cur, selected_experts, nullptr, true) :
+ moe_cache->build_experts(*this, cache, cur, selected_experts,
+ [](ggml_tensor * cur, ggml_tensor * ids, ggml_tensor * mm_ids, bool cached, void * user_data) {
+ return (*static_cast<decltype(build_experts) *>(user_data))(cur, ids, mm_ids, cached);
+ }, &build_experts);
if (down_exps_s) {
cb(experts, "ffn_moe_down_scaled", il);
@@ -2669,66 +2664,6 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
return moe_out;
}
-llm_graph_moe_cache_slots llm_graph_context::build_moe_cache_slots(
- ggml_tensor * selected_experts,
- ggml_tensor * up_exps,
- ggml_tensor * gate_exps,
- ggml_tensor * down_exps,
- ggml_tensor * gate_up_exps,
- bool cpu_part,
- int il) const {
- if (moe_cache == nullptr) {
- return { nullptr, nullptr };
- }
-
- ggml_tensor * cpu_map = cpu_part ? moe_cache->get_cpu_map(il, selected_experts->ne[1]) : nullptr;
- ggml_tensor * slot_map = moe_cache->get_slot_map(il, selected_experts->ne[1], selected_experts->ne[0], cpu_map != nullptr);
- if (slot_map == nullptr) {
- return { nullptr, nullptr };
- }
- for (ggml_tensor * w : { up_exps, gate_exps, down_exps, gate_up_exps }) {
- if (w == nullptr) {
- continue;
- }
- ggml_tensor * experts = moe_cache->get_experts(w);
- if (experts == nullptr) {
- return { nullptr, nullptr };
- }
- // the device of the cache must compute the matmul, else the scheduler copies the cache to another backend
- ggml_tensor * b = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, experts->ne[0], w == down_exps ? selected_experts->ne[0] : 1, selected_experts->ne[1]);
- if (!ggml_backend_supports_op(moe_cache->backend(il), ggml_mul_mat_id(ctx0, experts, b, selected_experts))) {
- return { nullptr, nullptr };
- }
- }
-
- ggml_tensor * ids = selected_experts;
- if (!ggml_is_contiguous(ids)) {
- ids = ggml_cont(ctx0, ids);
- }
- ids = ggml_reshape_1d(ctx0, ids, ggml_nelements(ids));
-
- // the slot map is a host weight, so the scheduler starts a new split here and copies it with the copy callback
- // the callback reads the selected experts, uploads the missing ones and updates the slot map
- ggml_tensor * slots = ggml_get_rows(ctx0, slot_map, ids); // [1, n_expert_used*n_tokens]
- if (!ggml_backend_supports_op(moe_cache->backend(il), slots)) {
- return { nullptr, nullptr };
- }
- ggml_backend_sched_set_tensor_backend(sched, slots, moe_cache->backend(il));
- cb(slots, "ffn_moe_slots", il);
-
- // add the lookup now: the transform of a Hadamard-folded expert reads host weights and would start its split first
- ggml_build_forward_expand(gf, slots);
-
- // the misses computed from the host experts, read after the lookup has picked them
- ggml_tensor * cpu_ids = nullptr;
- if (cpu_map != nullptr) {
- cpu_ids = ggml_get_rows(ctx0, cpu_map, ids); // [1, n_expert_used*n_tokens]
- cpu_ids = ggml_reshape_2d(ctx0, cpu_ids, selected_experts->ne[0], selected_experts->ne[1]);
- }
-
- return { ggml_reshape_2d(ctx0, slots, selected_experts->ne[0], selected_experts->ne[1]), cpu_ids };
-}
-
// input embeddings with optional lora
ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd, float tok_scale) const {
const int64_t n_embd_inp = hparams.n_embd_inp();
diff --git a/src/llama-graph.h b/src/llama-graph.h
index d3b527dce..559149508 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -1031,11 +1031,6 @@ struct llm_graph_amax_src1 {
ggml_tensor * scale_s; // scale to undo on the matmul result, or nullptr
};
-struct llm_graph_moe_cache_slots {
- ggml_tensor * slots; // [n_expert_used, n_tokens] slots of the selected experts in the MoE cache, or nullptr
- ggml_tensor * cpu_ids; // [n_expert_used, n_tokens] misses computed from the host experts, -1 for the other experts, or nullptr
-};
-
struct llm_graph_context {
const llama_model & model;
@@ -1136,7 +1131,7 @@ struct llm_graph_context {
ggml_tensor * w_s = nullptr) const;
// do mat_mul_id, while optionally apply lora and per-expert scale
- // if slots is set, the experts are selected by slots instead of ids, from the MoE cache if cached (see build_moe_cache_slots)
+ // if slots is set, the experts are selected by slots instead of ids, from the MoE cache if cached (see llama_moe_cache::build_lookup)
ggml_tensor * build_lora_mm_id(
ggml_tensor * w, // ggml_tensor * as
ggml_tensor * cur, // ggml_tensor * b
@@ -1249,17 +1244,6 @@ struct llm_graph_context {
ggml_tensor * post_exp_norm = nullptr,
float post_exp_norm_eps = 1e-5f) const;
- // the slots of the selected experts in the MoE cache, nullptr if the experts of the layer are not read from the cache
- // with cpu_part, also the misses computed from the host experts, if the cache leaves any to the CPU
- llm_graph_moe_cache_slots build_moe_cache_slots(
- ggml_tensor * selected_experts,
- ggml_tensor * up_exps,
- ggml_tensor * gate_exps,
- ggml_tensor * down_exps,
- ggml_tensor * gate_up_exps,
- bool cpu_part,
- int il) const;
-
//
// inputs
//
diff --git a/src/llama-moe-cache.cpp b/src/llama-moe-cache.cpp
index 965f6a177..1070b92b8 100644
--- a/src/llama-moe-cache.cpp
+++ b/src/llama-moe-cache.cpp
@@ -1,5 +1,6 @@
#include "llama-moe-cache.h"
+#include "llama-graph.h"
#include "llama-impl.h"
#include "llama-memory.h"
#include "llama-model.h"
@@ -635,6 +636,7 @@ struct llama_moe_cache::impl {
}
}
+ // the device that caches layer il
ggml_backend_t backend(int32_t il) const {
return devices[groups[layers[il].ig].id].backend;
}
@@ -684,6 +686,8 @@ struct llama_moe_cache::impl {
}
}
+ // the slot map of layer il, if its experts can be read from the cache for n_tokens tokens, nullptr otherwise
+ // cpu_part: the graph computes the misses in the CPU map of the layer (see get_cpu_map)
ggml_tensor * get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used, bool cpu_part) const {
if (il < 0 || il >= (int32_t) layers.size() || layers[il].ig < 0) {
return nullptr;
@@ -698,6 +702,8 @@ struct llama_moe_cache::impl {
return cpu_part ? l.slot_map_cpu : l.slot_map;
}
+ // the experts of layer il computed on the CPU: I32 [1, n_expert] in host memory, e for these experts, -1 for the others
+ // nullptr if no misses of layer il are computed on the CPU for ubatches of n_tokens tokens
ggml_tensor * get_cpu_map(int32_t il, int64_t n_tokens) const {
if (il < 0 || il >= (int32_t) layers.size() || layers[il].cpu_map == nullptr) {
return nullptr;
@@ -791,7 +797,7 @@ struct llama_moe_cache::impl {
}
}
GGML_ASSERT(ggml_is_contiguous(sel));
- GGML_ASSERT(sel->op == GGML_OP_RESHAPE); // of the [n_expert_used, n_tokens] ids, see build_moe_cache_slots
+ GGML_ASSERT(sel->op == GGML_OP_RESHAPE); // of the [n_expert_used, n_tokens] ids, see build_lookup
const int32_t n_used = (int32_t) sel->src[0]->ne[0];
@@ -878,20 +884,82 @@ llama_moe_cache::llama_moe_cache(const llama_model & model, bool no_alloc, const
llama_moe_cache::~llama_moe_cache() = default;
-ggml_backend_t llama_moe_cache::backend(int32_t il) const {
- return pimpl->backend(il);
+ggml_tensor * llama_moe_cache::get_experts(const ggml_tensor * w) const {
+ return pimpl->get_experts(w);
}
-ggml_tensor * llama_moe_cache::get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used, bool cpu_part) const {
- return pimpl->get_slot_map(il, n_tokens, n_expert_used, cpu_part);
-}
+llama_moe_cache_lookup llama_moe_cache::build_lookup(const llm_graph_context & g, const llama_moe_cache_layer & layer, ggml_tensor * selected_experts) const {
+ const int32_t il = layer.il;
+ ggml_context * ctx0 = g.ctx0;
-ggml_tensor * llama_moe_cache::get_experts(const ggml_tensor * w) const {
- return pimpl->get_experts(w);
+ // part of the misses may be computed from the host experts, but not with weights before the FFN, LoRAs or rotated experts
+ bool cpu_part = !layer.weight_before_ffn && g.loras->empty();
+ for (ggml_tensor * w : { layer.up_exps, layer.gate_exps, layer.down_exps, layer.gate_up_exps }) {
+ cpu_part = cpu_part && !(w != nullptr && g.hdmd != nullptr && g.hdmd->rot.count(w) > 0);
+ }
+
+ ggml_tensor * cpu_map = cpu_part ? pimpl->get_cpu_map(il, selected_experts->ne[1]) : nullptr;
+ ggml_tensor * slot_map = pimpl->get_slot_map(il, selected_experts->ne[1], selected_experts->ne[0], cpu_map != nullptr);
+ if (slot_map == nullptr) {
+ return {};
+ }
+ for (ggml_tensor * w : { layer.up_exps, layer.gate_exps, layer.down_exps, layer.gate_up_exps }) {
+ if (w == nullptr) {
+ continue;
+ }
+ ggml_tensor * experts = pimpl->get_experts(w);
+ if (experts == nullptr) {
+ return {};
+ }
+ // the device of the cache must compute the matmul, else the scheduler copies the cache to another backend
+ ggml_tensor * b = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, experts->ne[0], w == layer.down_exps ? selected_experts->ne[0] : 1, selected_experts->ne[1]);
+ if (!ggml_backend_supports_op(pimpl->backend(il), ggml_mul_mat_id(ctx0, experts, b, selected_experts))) {
+ return {};
+ }
+ }
+
+ ggml_tensor * ids = selected_experts;
+ if (!ggml_is_contiguous(ids)) {
+ ids = ggml_cont(ctx0, ids);
+ }
+ ids = ggml_reshape_1d(ctx0, ids, ggml_nelements(ids));
+
+ // the slot map is a host weight, so the scheduler starts a new split here and copies it with the copy callback
+ // the callback reads the selected experts, uploads the missing ones and updates the slot map
+ ggml_tensor * slots = ggml_get_rows(ctx0, slot_map, ids); // [1, n_expert_used*n_tokens]
+ if (!ggml_backend_supports_op(pimpl->backend(il), slots)) {
+ return {};
+ }
+ ggml_backend_sched_set_tensor_backend(g.sched, slots, pimpl->backend(il));
+ g.cb(slots, "ffn_moe_slots", il);
+
+ // add the lookup now: the transform of a Hadamard-folded expert reads host weights and would start its split first
+ ggml_build_forward_expand(g.gf, slots);
+
+ // the misses computed from the host experts, read after the lookup has picked them
+ ggml_tensor * cpu_ids = nullptr;
+ if (cpu_map != nullptr) {
+ cpu_ids = ggml_get_rows(ctx0, cpu_map, ids); // [1, n_expert_used*n_tokens]
+ cpu_ids = ggml_reshape_2d(ctx0, cpu_ids, selected_experts->ne[0], selected_experts->ne[1]);
+ }
+
+ return { il, ggml_reshape_2d(ctx0, slots, selected_experts->ne[0], selected_experts->ne[1]), cpu_ids };
}
-ggml_tensor * llama_moe_cache::get_cpu_map(int32_t il, int64_t n_tokens) const {
- return pimpl->get_cpu_map(il, n_tokens);
+ggml_tensor * llama_moe_cache::build_experts(const llm_graph_context & g, const llama_moe_cache_lookup & lookup,
+ ggml_tensor * cur, ggml_tensor * selected_experts, llama_moe_cache_build_experts_t build, void * user_data) const {
+ GGML_ASSERT(lookup.slots != nullptr);
+
+ ggml_tensor * experts = build(cur, selected_experts, lookup.slots, true, user_data);
+ if (lookup.cpu_ids != nullptr) {
+ // the cached experts and the misses computed from the host experts, each reads the others as zero
+ // the cached experts come first in the graph, so that they are computed while the misses are
+ ggml_build_forward_expand(g.gf, experts);
+ experts = ggml_add(g.ctx0, experts, build(cur, selected_experts, lookup.cpu_ids, false, user_data));
+ g.cb(experts, "ffn_moe_down_merged", lookup.il);
+ }
+
+ return experts;
}
bool llama_moe_cache::copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph) {
diff --git a/src/llama-moe-cache.h b/src/llama-moe-cache.h
index 636bc4559..833eac360 100644
--- a/src/llama-moe-cache.h
+++ b/src/llama-moe-cache.h
@@ -7,6 +7,28 @@
#include <vector>
struct llama_model;
+struct llm_graph_context;
+
+// the experts of a MoE layer, as llm_graph_context::build_moe_ffn passes them to the MoE cache
+struct llama_moe_cache_layer {
+ int32_t il;
+ ggml_tensor * up_exps;
+ ggml_tensor * gate_exps;
+ ggml_tensor * down_exps;
+ ggml_tensor * gate_up_exps;
+ bool weight_before_ffn; // the router weights scale the expert input
+};
+
+// the lookup of the selected experts of a layer in the MoE cache, see llama_moe_cache::build_lookup
+struct llama_moe_cache_lookup {
+ int32_t il = -1;
+ ggml_tensor * slots = nullptr; // [n_expert_used, n_tokens] slots of the selected experts in the MoE cache, nullptr if the layer does not read the cache
+ ggml_tensor * cpu_ids = nullptr; // [n_expert_used, n_tokens] misses computed from the host experts, -1 for the other experts, or nullptr
+};
+
+// builds the experts selected by mm_ids up to the down projection, from the MoE cache if cached: [n_embd, n_expert_used, n_tokens]
+// cur is the expert input and ids are the selected experts, for the biases and scales (see llm_graph_context::build_moe_ffn)
+using llama_moe_cache_build_experts_t = ggml_tensor * (*)(ggml_tensor * cur, ggml_tensor * ids, ggml_tensor * mm_ids, bool cached, void * user_data);
// keeps recently and repeatedly used experts of host-resident MoE layers in a device buffer
// each layer has a slot map in host memory: when the scheduler copies it to the device, the copy callback uploads the missing experts
@@ -17,19 +39,19 @@ public:
llama_moe_cache(const llama_model & model, bool no_alloc, const std::vector<ggml_backend_t> & backends, const std::vector<ggml_backend_buffer_type_t> & bufts, size_t size, float cpu_share);
~llama_moe_cache();
- // the device that caches layer il
- ggml_backend_t backend(int32_t il) const;
-
- // the slot map of layer il, if its experts can be read from the cache for n_tokens tokens, nullptr otherwise
- // cpu_part: the graph computes the misses in the CPU map of the layer (see get_cpu_map)
- ggml_tensor * get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used, bool cpu_part) const;
-
// the experts of w in the cache, nullptr if w is not cached
ggml_tensor * get_experts(const ggml_tensor * w) const;
- // the experts of layer il computed on the CPU: I32 [1, n_expert] in host memory, e for these experts, -1 for the others
- // nullptr if no misses of layer il are computed on the CPU for ubatches of n_tokens tokens
- ggml_tensor * get_cpu_map(int32_t il, int64_t n_tokens) const;
+ // the graph of a MoE layer, built by llm_graph_context::build_moe_ffn in two steps:
+
+ // 1. the lookup of the selected experts in the cache, built after the expert input and before the expert matmuls
+ // returns no slots if the layer does not read its experts from the cache for this ubatch
+ llama_moe_cache_lookup build_lookup(const llm_graph_context & g, const llama_moe_cache_layer & layer, ggml_tensor * selected_experts) const;
+
+ // 2. the experts of a layer with slots, from the cache and from the host experts: [n_embd, n_expert_used, n_tokens]
+ // calls build synchronously; does not retain user_data
+ ggml_tensor * build_experts(const llm_graph_context & g, const llama_moe_cache_lookup & lookup,
+ ggml_tensor * cur, ggml_tensor * selected_experts, llama_moe_cache_build_experts_t build, void * user_data) const;
// ggml_backend_sched copy callback, returns false if src is not a slot map
bool copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph);