Commit 33d6135cb for llama.cpp
commit 33d6135cbe20e90e67d9db1c874fbb4a9dc64b19
Author: Evan Huus <eapache@gmail.com>
Date: Mon Oct 19 07:28:50 2026 -0400
moe-cache : replace MoE LRU cache with a 50/50 SLRU cache (#30637)
diff --git a/src/llama-moe-cache.cpp b/src/llama-moe-cache.cpp
index f22d986dc..4d3596b24 100644
--- a/src/llama-moe-cache.cpp
+++ b/src/llama-moe-cache.cpp
@@ -12,19 +12,21 @@
namespace {
-// LRU 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_lru {
+// 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 {
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
- // doubly linked list of the slots, head is the least recently used
+ // 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 = -1;
- int32_t tail = -1;
+ int32_t head[2] = { -1, -1 };
+ int32_t tail[2] = { -1, -1 };
+ std::vector<bool> is_protected;
+ int32_t n_protected = 0;
std::vector<uint32_t> seen; // [n_expert]
uint32_t seen_gen = 0;
@@ -41,27 +43,59 @@ struct moe_cache_lru {
prev[s] = s - 1;
next[s] = s + 1 < n_slots ? s + 1 : -1;
}
- head = 0;
- tail = n_slots - 1;
+ head[0] = 0;
+ tail[0] = n_slots - 1;
+ is_protected.assign(n_slots, false);
seen.assign(n_expert, 0);
}
- // move slot s to the tail (most recently used)
- void touch(int32_t s) {
- if (s == tail) {
+ // 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]) {
return;
}
if (prev[s] >= 0) {
next[prev[s]] = next[s];
} else {
- head = next[s];
+ head[from] = next[s];
+ }
+ if (next[s] >= 0) {
+ prev[next[s]] = prev[s];
+ } else {
+ tail[from] = prev[s];
}
- prev[next[s]] = prev[s];
- prev[s] = tail;
+ prev[s] = tail[protect];
next[s] = -1;
- next[tail] = s;
- tail = s;
+ if (tail[protect] >= 0) {
+ next[tail[protect]] = s;
+ } else {
+ head[protect] = s;
+ }
+ tail[protect] = s;
+ is_protected[s] = protect;
+ n_protected += int(protect) - int(from);
+ }
+
+ void touch(int32_t s) {
+ move(s, true);
+ if (n_protected > n_slots/2) {
+ move(head[1], false);
+ }
+ }
+
+ 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;
+ }
+ }
+ }
+ GGML_ABORT("no SLRU victim");
}
struct fill {
@@ -69,7 +103,7 @@ struct moe_cache_lru {
int32_t slot;
};
- // give a slot to each expert selected by ids in layer il, the misses evict the least recently used experts
+ // give a slot to each expert selected by ids in layer il, evicting from probation first
// 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) {
fills.clear();
@@ -93,7 +127,7 @@ struct moe_cache_lru {
int32_t * slots = slot_map[il];
- // hits go to the tail first, so the head can be evicted below
+ // promote hits before choosing victims for misses
for (int32_t e : uniq) {
if (slots[e] >= 0) {
touch(slots[e]);
@@ -106,13 +140,13 @@ struct moe_cache_lru {
if (slots[e] >= 0) {
continue;
}
- const int32_t s = head;
+ const int32_t s = victim(il);
if (key_of[s] >= 0) {
slot_map[key_of[s] / n_expert][key_of[s] % n_expert] = -1;
}
key_of[s] = il*n_expert + e;
slots[e] = s;
- touch(s);
+ move(s, false);
fills.push_back({ e, s });
}
return true;
@@ -172,7 +206,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_lru lru;
+ moe_cache_slru slru;
};
struct layer {
@@ -208,7 +242,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_lru::fill> fills;
+ std::vector<moe_cache_slru::fill> fills;
// slot maps in host memory
ggml_context_ptr ctx_host;
@@ -323,7 +357,7 @@ struct llama_moe_cache::impl {
g.n_slots = 0;
continue;
}
- g.lru.init(model.layers.size(), n_expert, g.n_slots);
+ g.slru.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();
}
@@ -421,7 +455,7 @@ struct llama_moe_cache::impl {
for (group & g : groups) {
for (int32_t il : g.layers) {
if (layers[il].slot_map != nullptr) {
- g.lru.slot_map[il] = (int32_t *) layers[il].slot_map->data;
+ g.slru.slot_map[il] = (int32_t *) layers[il].slot_map->data;
}
}
}
@@ -489,7 +523,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.lru.slot_map[b.il];
+ const int32_t * slots = g.slru.slot_map[b.il];
if (slots == nullptr || slots[e] < 0) {
return 0;
}
@@ -550,7 +584,7 @@ struct llama_moe_cache::impl {
ggml_backend_synchronize(backend);
size_t n_hit = 0;
- if (!g.lru.plan(il, ids.data(), ids.size(), fills, n_hit)) {
+ if (!g.slru.plan(il, ids.data(), ids.size(), fills, n_hit)) {
GGML_ABORT("the MoE cache is too small for the experts selected in layer %d", il);
}
diff --git a/src/llama-moe-cache.h b/src/llama-moe-cache.h
index 3b6726b86..9baf06390 100644
--- a/src/llama-moe-cache.h
+++ b/src/llama-moe-cache.h
@@ -8,7 +8,7 @@
struct llama_model;
-// keeps the most recently used experts of host-resident MoE layers in a device buffer
+// 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
class llama_moe_cache {
public: