Commit c4fb57969 for llama.cpp
commit c4fb579695320929b6a453f4f3a56951bda332ee
Author: Alexey Kopytko <alexey@kopytko.com>
Date: Fri Nov 6 18:32:54 2026 +0900
moe-cache : LRFU/MQ-like hybrid cache, ClassLRU (#31233)
* moe-cache : LRFU/MQ-inspired hybrid cache
* make the clock advance by N tokens: use the real token count (ne[1] is always 1, ids have n_expert_used entries per token)
* sel is a RESHAPE node, the orig input is in src[0], so sel->src[0]->ne[1] is the count
* c for the class index, last_use for last use time
* explain the roles of these
* compute the key once
* use the key directly instead of re-reading
* show how the score is computed from the age
* remove the branch and clarify exp2 and log2 relationship
* reword the header comment and elaborate on unused slots and sorting
* not a copy, not just skip
* std::array maybe
* Revert "std::array maybe"
This reverts commit 7649b1f172db4941a854ad115bec0c7e062b65b9.
* half life it is
* add key and iteration helpers
* extract decayed_crf
* bump classes to 8, document the formula, display the cache size in token worth count
* assert the computed number of classes matches the half life
diff --git a/src/llama-moe-cache.cpp b/src/llama-moe-cache.cpp
index 6a343a760..b736930f9 100644
--- a/src/llama-moe-cache.cpp
+++ b/src/llama-moe-cache.cpp
@@ -7,33 +7,45 @@
#include "ggml-cpp.h"
#include <algorithm>
+#include <cmath>
#include <stdexcept>
#include <unordered_map>
#include <vector>
namespace {
-// SLRU of the experts of a group of layers, the slot of each expert is kept in the slot map of its layer
-struct moe_cache_slru {
+// ClassLRU: one LRU list per frequency class, over the experts of a group of layers
+// slot_map[il][e] is the slot of expert e of layer il, -1 if not cached, the host slot map tensor
+// CRF (combined recency and frequency) = uses of an expert, decayed with a half-life of N tokens, retained across evictions
+// class = floor(log2(CRF)), at most n_class - 1; each slot is in the LRU list of the class, as of its last use
+// victim: an empty slot, or the least recently used slot of each list not in the current plan() with the max age/2^class
+struct moe_cache_lru {
+ static constexpr int n_class = 8; // floor(log2(1/(1 - 2^(-1/crf_half_life)))) = 7, class of CRF_max for half-life of 128
+ static constexpr float crf_half_life = 128.0f; // tokens
+
int32_t n_expert = 0;
int32_t n_slots = 0;
std::vector<int32_t *> slot_map; // [n_layer] data of the slot maps, -1 if the expert is not cached
- std::vector<int32_t> key_of; // [n_slots] il*n_expert + expert, -1 if empty
+ std::vector<int32_t> key_of; // [n_slots] key of the expert in the slot, -1 if empty
+
+ // one LRU list per class, the head is the least recently used end
+ std::vector<int32_t> prev; // [n_slots] -1 at the ends
+ std::vector<int32_t> next; // [n_slots]
+ int32_t head[n_class];
+ int32_t tail[n_class];
+ std::vector<int8_t> list_of; // [n_slots] class list of the slot
- // probation and protected LRU lists; hits promote slots to the protected half
- std::vector<int32_t> prev;
- std::vector<int32_t> next;
- int32_t head[2] = { -1, -1 };
- int32_t tail[2] = { -1, -1 };
- std::vector<bool> is_protected;
- int32_t n_protected = 0;
+ std::vector<float> crf; // [n_layer*n_expert]
+ std::vector<int64_t> last_use; // [n_layer*n_expert] layer clock at the last use, -1 if not used
+ std::vector<int64_t> clock; // [n_layer] tokens seen by the layer (as added by plan() for each ubatch)
- std::vector<uint32_t> seen; // [n_expert]
+ std::vector<uint32_t> seen; // [n_expert] seen[e] == seen_gen if expert e is selected by the current plan()
uint32_t seen_gen = 0;
- std::vector<int32_t> uniq;
+ std::vector<int32_t> uniq; // the distinct selected experts
void init(int32_t n_layer, int32_t n_expert, int32_t n_slots) {
+ GGML_ASSERT(n_class == (int) std::log2(1.0f/(1.0f - std::exp2(-1.0f/crf_half_life))) + 1);
this->n_expert = n_expert;
this->n_slots = n_slots;
slot_map.assign(n_layer, nullptr);
@@ -44,16 +56,23 @@ struct moe_cache_slru {
prev[s] = s - 1;
next[s] = s + 1 < n_slots ? s + 1 : -1;
}
+ for (int c = 0; c < n_class; ++c) {
+ head[c] = -1;
+ tail[c] = -1;
+ }
head[0] = 0;
tail[0] = n_slots - 1;
- is_protected.assign(n_slots, false);
+ list_of.assign(n_slots, 0);
+ crf.assign((size_t) n_layer*n_expert, 0.0f);
+ last_use.assign((size_t) n_layer*n_expert, -1);
+ clock.assign(n_layer, 0);
seen.assign(n_expert, 0);
}
- // move slot s to the most recently used end of a segment
- void move(int32_t s, bool protect) {
- const bool from = is_protected[s];
- if (from == protect && s == tail[from]) {
+ // move slot s to the most recently used end of list c
+ void move(int32_t s, int c) {
+ const int from = list_of[s];
+ if (from == c && s == tail[from]) {
return;
}
if (prev[s] >= 0) {
@@ -67,36 +86,82 @@ struct moe_cache_slru {
tail[from] = prev[s];
}
- prev[s] = tail[protect];
+ prev[s] = tail[c];
next[s] = -1;
- if (tail[protect] >= 0) {
- next[tail[protect]] = s;
+ if (tail[c] >= 0) {
+ next[tail[c]] = s;
} else {
- head[protect] = s;
+ head[c] = s;
+ }
+ tail[c] = s;
+ list_of[s] = c;
+ }
+
+ // key = index of an expert of a layer in crf and last_use
+ int32_t make_key(int32_t il, int32_t e) const { return il*n_expert + e; }
+ int32_t key_layer (int32_t key) const { return key / n_expert; }
+ int32_t key_expert(int32_t key) const { return key % n_expert; }
+
+ // for a key after use(): its CRF is >= 1, so log2 is >= 0 and the cast is the floor
+ int class_of(int32_t key) const {
+ return std::min(n_class - 1, (int) std::log2(crf[key]));
+ }
+
+ // the CRF of a key, decayed by the time since its last use
+ float decayed_crf(int32_t il, int32_t key) const {
+ return crf[key]*std::exp2(-(clock[il] - last_use[key])/crf_half_life);
+ }
+
+ // an unused key gets 1 because crf[key] is 0 prior to first use
+ void use(int32_t il, int32_t key) {
+ crf[key] = decayed_crf(il, key) + 1.0f;
+ last_use[key] = clock[il];
+ }
+
+ // the experts selected by this plan() must not be evicted, nor the slots this plan() filled already
+ bool is_selected(int32_t il, int32_t key) const {
+ if (key < 0) {
+ return false; // unused slot
}
- tail[protect] = s;
- is_protected[s] = protect;
- n_protected += int(protect) - int(from);
+ if (key_layer(key) != il) {
+ return false; // seen is per expert index, an expert of another layer can have the same index
+ }
+ return seen[key_expert(key)] == seen_gen;
}
- void touch(int32_t s) {
- move(s, true);
- if (n_protected > n_slots/2) {
- move(head[1], false);
+ // the least recently used slot of list c that this plan() does not select, -1 if none
+ int32_t candidate(int c, int32_t il) const {
+ for (int32_t s = head[c]; s >= 0; s = next[s]) {
+ if (!is_selected(il, key_of[s])) {
+ return s;
+ }
}
+ return -1;
}
int32_t victim(int32_t il) const {
- for (int segment = 0; segment < 2; ++segment) {
- for (int32_t s = head[segment]; s >= 0; s = next[s]) {
- const int32_t key = key_of[s];
- // keep every expert selected by this request, including earlier fills
- if (key < 0 || key / n_expert != il || seen[key % n_expert] != seen_gen) {
- return s;
- }
+ int32_t res = -1;
+ double res_score = -1.0;
+ for (int c = 0; c < n_class; ++c) {
+ const int32_t s = candidate(c, il);
+ if (s < 0) {
+ continue;
+ }
+ const int32_t key = key_of[s];
+ if (key < 0) {
+ return s; // unused slot
}
+ const int64_t age = clock[key_layer(key)] - last_use[key];
+ const double score = std::ldexp((double) age, -c); // age/2^class
+ if (score > res_score) {
+ res = s;
+ res_score = score;
+ }
+ }
+ if (res < 0) {
+ GGML_ABORT("no MoE cache victim");
}
- GGML_ABORT("no SLRU victim");
+ return res;
}
struct fill {
@@ -104,9 +169,9 @@ struct moe_cache_slru {
int32_t slot;
};
- // give a slot to each expert selected by ids in layer il, evicting from probation first
+ // give a slot to each expert selected by ids (n_tokens tokens) in layer il
// returns false if the ids select more distinct experts than there are slots
- bool plan(int32_t il, const int32_t * ids, size_t n_ids, std::vector<fill> & fills, size_t & n_hit) {
+ bool plan(int32_t il, const int32_t * ids, size_t n_ids, int64_t n_tokens, std::vector<fill> & fills, size_t & n_hit) {
fills.clear();
n_hit = 0;
@@ -128,14 +193,17 @@ struct moe_cache_slru {
int32_t * slots = slot_map[il];
- // promote hits before choosing victims for misses
+ // update the CRF of all selected experts and promote hits before choosing victims for misses
for (int32_t e : uniq) {
+ const int32_t key = make_key(il, e);
+ use(il, key);
if (slots[e] >= 0) {
- touch(slots[e]);
+ move(slots[e], class_of(key));
n_hit++;
}
}
- // sorted misses usually get consecutive slots, so the uploads can be merged
+ // fill the misses in expert order: victim() takes unused slots, and slots filled by one plan(), in list order,
+ // so consecutive experts often get consecutive slots and one upload
std::sort(uniq.begin(), uniq.end());
for (int32_t e : uniq) {
if (slots[e] >= 0) {
@@ -143,13 +211,15 @@ struct moe_cache_slru {
}
const int32_t s = victim(il);
if (key_of[s] >= 0) {
- slot_map[key_of[s] / n_expert][key_of[s] % n_expert] = -1;
+ slot_map[key_layer(key_of[s])][key_expert(key_of[s])] = -1;
}
- key_of[s] = il*n_expert + e;
- slots[e] = s;
- move(s, false);
+ const int32_t key = make_key(il, e);
+ key_of[s] = key;
+ slots[e] = s;
+ move(s, class_of(key));
fills.push_back({ e, s });
}
+ clock[il] += n_tokens;
return true;
}
};
@@ -207,7 +277,7 @@ struct llama_moe_cache::impl {
std::vector<ggml_tensor *> banks; // device storage of all slots, one per expert tensor
size_t host_bytes = 0;
int32_t n_slots = 0;
- moe_cache_slru slru;
+ moe_cache_lru lru;
};
struct layer {
@@ -243,7 +313,7 @@ struct llama_moe_cache::impl {
std::unordered_map<const ggml_tensor *, int32_t> layer_of; // slot map -> layer
std::vector<int32_t> ids;
- std::vector<moe_cache_slru::fill> fills;
+ std::vector<moe_cache_lru::fill> fills;
// slot maps in host memory
ggml_context_ptr ctx_host;
@@ -358,7 +428,7 @@ struct llama_moe_cache::impl {
g.n_slots = 0;
continue;
}
- g.slru.init(model.layers.size(), n_expert, g.n_slots);
+ g.lru.init(model.layers.size(), n_expert, g.n_slots);
n_tensors[g.id] += g.ref.size()*(1 + g.layers.size());
n_tensors_host += g.layers.size();
}
@@ -438,7 +508,7 @@ struct llama_moe_cache::impl {
for (group & g : groups) {
for (int32_t il : g.layers) {
if (layers[il].slot_map != nullptr) {
- g.slru.slot_map[il] = (int32_t *) layers[il].slot_map->data;
+ g.lru.slot_map[il] = (int32_t *) layers[il].slot_map->data;
}
}
}
@@ -461,8 +531,9 @@ struct llama_moe_cache::impl {
ggml_backend_buft_name(d.buft), d.buf_size/1024.0/1024.0, d.host_bytes/1024.0/1024.0);
for (const group & g : groups) {
if (g.id == (int32_t) id) {
- LLAMA_LOG_INFO("%s: %2zu layers, %s: %5d slots (%.1f%%)\n", __func__,
- g.layers.size(), ggml_type_name(g.ref.back()->type), g.n_slots, 100.0*g.n_slots/(g.layers.size()*g.ref[0]->ne[2]));
+ LLAMA_LOG_INFO("%s: %2zu layers, %s: %5d slots (%.1f%%, %.1f tokens worth)\n", __func__,
+ g.layers.size(), ggml_type_name(g.ref.back()->type), g.n_slots, 100.0*g.n_slots/(g.layers.size()*g.ref[0]->ne[2]),
+ (double) g.n_slots/(g.layers.size()*n_expert_used));
}
}
}
@@ -506,7 +577,7 @@ struct llama_moe_cache::impl {
}
// large batches only read the cache, so the experts used in generation stay in it
- const int32_t * slots = g.slru.slot_map[b.il];
+ const int32_t * slots = g.lru.slot_map[b.il];
if (slots == nullptr || slots[e] < 0) {
return 0;
}
@@ -561,13 +632,14 @@ 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
ids.resize(ggml_nelements(sel));
ggml_backend_tensor_get_async(backend, sel, ids.data(), 0, ggml_nbytes(sel));
ggml_backend_synchronize(backend);
size_t n_hit = 0;
- if (!g.slru.plan(il, ids.data(), ids.size(), fills, n_hit)) {
+ if (!g.lru.plan(il, ids.data(), ids.size(), sel->src[0]->ne[1], fills, n_hit)) {
GGML_ABORT("the MoE cache is too small for the experts selected in layer %d", il);
}