Commit e95426f16 for llama.cpp
commit e95426f16982054c379ca7bdc8186996dd22f0fe
Author: Aman Gupta <amangupta052@gmail.com>
Date: Sun Nov 8 00:20:58 2026 +0800
llama_moe_cache: run missed experts on CPU, add tuning script (#30966)
* moe-cache : compute part of the misses on the CPU
--moe-cache-cpu-share F computes a share F of the MoE cache misses on the CPU
from the host experts instead of uploading them. The graph builds the expert
FFN twice: on the device from the cache slots, where the misses left to the CPU
read zero slots, and on the CPU from the host experts, where the experts
computed on the device have id -1. The two results are added before the down
bias.
- ggml-cpu: MUL_MAT_ID skips negative ids and writes zero rows
- sched: the copy callback is called for all the inputs of a split, and before
the host sync when there are no events
- the slot lookup callback reads the inputs of the CPU part together with the
selected experts, so the CPU part does not wait for the device, and its
result is copied back asynchronously
- each token uses distinct zero slots, mul_mat_id does not support an expert
selected twice by a token
- LLAMA_MOE_CACHE_TRACE writes the selected experts of each lookup, for a
simulator of the cache
* scripts : add moe-cache-tune.py to pick the MoE cache arguments
The script wraps a llama-server command: it starts the server at two cache
sizes and three CPU shares, fits the upload and CPU speeds of the machine, and
replays the routing of the runs through a copy of the cache logic to predict
other sizes and shares. It prints the recommended --moe-cache-mib and
--moe-cache-cpu-share.
docs/moe-cache.md describes the MoE cache arguments and the script.
* add am17an as codeowner for MoE cache
* moe-cache : upload all the misses until the cache is full
With a high --moe-cache-cpu-share, few misses are uploaded, so an empty cache
fills slowly and more misses than needed go to the CPU. The CPU share now
starts once all the slots of a cache group hold an expert. moe-cache-tune.py
replays the same rule.
* add backend sync for async transfer queue
* moe-cache : let the scheduler copy the results of the CPU part
* Fix MSVC compilation
* remove sync change
* llama-graph : extract the CPU share condition into can_cpu_share
* llama-graph : make llm_graph_params::moe_cache const
The graph construction no longer updates the MoE cache directly: the CPU
parts of the layers are accumulated in the graph result and applied to
the cache in llama_context::process_ubatch, after the graph inputs are
set and before the graph is computed.
Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (Mac Studio, M5 Ultra, 256GB)
* scripts : replay the ClassLRU policy in moe-cache-tune.py
* moe-cache : pick CPU misses by CRF instead of random (#31853)
* moe-cache : remove the backend logic of the CPU part from the graph
The CPU part reads the host experts, so the scheduler places it on the
CPU without set_tensor_backend. The cache gives no CPU map for batches
that a device would offload. Graphs with a CPU part look up their own
slot map, so the graph does not register the part with the cache.
ggml-backend : the inputs of the next split that are already computed
are copied to host memory before the current split runs, so a CPU split
does not wait for the device split before it. This replaces the stash
of the cache, and the copy callback is again called for host weights
only.
* llama-graph : return the MoE cache slots and CPU ids in a struct
The scheduler prefetch no longer checks that the next split runs on
another backend: the inputs of a split are never on its own backend,
so the input loop already skips that case.
---------
Co-authored-by: Gaurav Garg <gaugarg@nvidia.com>
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
Co-authored-by: Alexey Kopytko <alexey@kopytko.com>
diff --git a/CODEOWNERS b/CODEOWNERS
index a1e27dec3..f34a07f02 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -87,6 +87,7 @@
/src/llama-graph.* @CISC
/src/llama-model.* @CISC
/src/llama-vocab.* @CISC
+/src/llama-moe-cache.* @am17an
/src/models/ @ggml-org/llama-models
/tests/ @ggerganov
/tests/test-chat.* @pwilkin
diff --git a/common/arg.cpp b/common/arg.cpp
index 51f226e96..69be4a61e 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -2754,6 +2754,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
params.moe_cache_size = (size_t) value*1024*1024;
}
).set_env("LLAMA_ARG_MOE_CACHE_MIB"));
+ add_opt(common_arg(
+ {"--moe-cache-cpu-share"}, "F",
+ string_format("share of the MoE cache misses that the CPU computes from the host experts instead of uploading them to the GPU once the cache is full, from 0 to 1 (default: %.2f)", (double) params.moe_cache_cpu_share),
+ [](common_params & params, const std::string & value) {
+ const float share = std::stof(value);
+ if (share < 0.0f || share > 1.0f) {
+ throw std::invalid_argument("invalid value");
+ }
+ params.moe_cache_cpu_share = share;
+ }
+ ).set_env("LLAMA_ARG_MOE_CACHE_CPU_SHARE"));
add_opt(common_arg(
{"-ncffn", "--n-cpu-ffn"}, "N",
"keep the dense FFN weights of the first N layers in the CPU\n"
diff --git a/common/common.cpp b/common/common.cpp
index a2801b82f..75fd39665 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1758,7 +1758,8 @@ struct llama_context_params common_context_params_to_llama(const common_params &
cparams.type_k = params.cache_type_k;
cparams.type_v = params.cache_type_v;
- cparams.moe_cache_size = params.moe_cache_size;
+ cparams.moe_cache_size = params.moe_cache_size;
+ cparams.moe_cache_cpu_share = params.moe_cache_cpu_share;
return cparams;
}
diff --git a/common/common.h b/common/common.h
index 4005275ca..84abddc76 100644
--- a/common/common.h
+++ b/common/common.h
@@ -597,7 +597,8 @@ struct common_params {
ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K
ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V
- size_t moe_cache_size = 0; // GPU cache size in bytes for the MoE experts kept in the CPU, split among the GPUs like the layers
+ size_t moe_cache_size = 0; // GPU cache size in bytes for the MoE experts kept in the CPU, split among the GPUs like the layers
+ float moe_cache_cpu_share = 0.0f; // share of the MoE cache misses computed on the CPU instead of uploaded
common_conversation_mode conversation_mode = COMMON_CONVERSATION_MODE_AUTO;
diff --git a/docs/moe-cache.md b/docs/moe-cache.md
new file mode 100644
index 000000000..7491a2a31
--- /dev/null
+++ b/docs/moe-cache.md
@@ -0,0 +1,26 @@
+# MoE cache
+
+When the experts of MoE layers stay in host memory (`-cmoe`, `-ncmoe`, `-ot`), the MoE cache keeps the most used ones in device memory for small batches (decode, speculative decoding). Large batches still upload the experts they use.
+
+| Argument | |
+|---|---|
+| `--moe-cache-mib N` | cache size in MiB, split among the GPUs like the layers (default: 0, disabled) |
+| `--moe-cache-cpu-share F` | share of the cache misses computed on the CPU from the host experts instead of uploaded, once the cache is full (default: 0) |
+
+The cache hits run on the device. With `--moe-cache-cpu-share F` the CPU computes a share `F` of the misses while the device computes the hits and uploads the other misses into the cache. This is faster when the CPU reads the host experts faster than the device can upload them, for example with fast RAM, a PCIe 4 link or experts that are not in pinned memory. The host experts are pinned with `-lm none`, `-lm mlock` and `-lm dio`, unless `--no-host`, but not when they are mmap'd (the default). The device uploads pinned experts at the full speed of the link, so the CPU share often does not help then. Until the cache is full every miss is uploaded, so the cache fills as fast as without the CPU share. The best share depends on the machine, the model and the cache size.
+
+## Picking the arguments
+
+Put the script in front of your server command:
+
+```bash
+python3 scripts/moe-cache-tune.py -- build/bin/llama-server -m model.gguf -cmoe -t 8 -c 32768
+```
+
+It starts the server 6 times, at two cache sizes × shares 0, 0.5 and 0.95, and times the generation of a few thousand tokens. That takes a few minutes, mostly model loading.
+
+From these runs it fits the upload and CPU speeds of the machine, then replays the routing of all the runs through a copy of the cache logic for other sizes and shares. It prints the measured speeds, the predicted t/s and the recommended arguments: the lowest share within 5% of the fastest.
+
+- **Server command:** keep your own arguments (`-t`, `-c`, expert placement, load mode); the speeds and the free device memory depend on them.
+- **`--max-mib N`:** tunes a smaller cache than what fits on the device. The runs in `--out` are reused while the command stays the same, so trying another value is instant.
+- **`--markdown`:** prints Markdown tables to paste in a PR.
diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h
index 5a5ba966b..973aea25b 100644
--- a/ggml/include/ggml.h
+++ b/ggml/include/ggml.h
@@ -1519,6 +1519,7 @@ extern "C" {
// note: this op has undefined behaviour if a given `a` is referenced more than once for a given `b`
// (i.e. duplicate experts per token are not guaranteed to be supported in all backends)
// TODO: add backend support: https://github.com/ggml-org/llama.cpp/pull/29998#issuecomment-6153742533
+ // the CPU backend skips negative ids, their rows of the result are zero
GGML_API struct ggml_tensor * ggml_mul_mat_id(
struct ggml_context * ctx,
struct ggml_tensor * as,
diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp
index 3520b4349..92a53bf38 100644
--- a/ggml/src/ggml-backend.cpp
+++ b/ggml/src/ggml-backend.cpp
@@ -1246,6 +1246,14 @@ struct ggml_backend_sched {
int prev_graph_size;
int debug_realloc;
+
+ // inputs of the next split read from the backend of the current split before it computes, see ggml_backend_sched_prefetch_inputs
+ ggml_backend_event_t prefetch_events[GGML_SCHED_MAX_BACKENDS]; // created on first use, NULL if the backend has no events
+ bool prefetch_tried [GGML_SCHED_MAX_BACKENDS];
+ struct ggml_tensor * prefetched[GGML_SCHED_MAX_SPLIT_INPUTS];
+ int n_prefetched;
+ int prefetch_split; // the split that reads the prefetched inputs
+ int prefetch_backend_id; // the backend that copies them
};
#define hash_id(tensor) ggml_hash_find_or_insert(&sched->hash_set, tensor)
@@ -2086,6 +2094,16 @@ static void ggml_backend_sched_copy_input(ggml_backend_sched_t sched, struct ggm
ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, input);
struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy);
+ // read before the previous split, wait for the copy only
+ if (split == &sched->splits[sched->prefetch_split]) {
+ for (int i = 0; i < sched->n_prefetched; i++) {
+ if (sched->prefetched[i] == input) {
+ ggml_backend_event_synchronize(sched->prefetch_events[sched->prefetch_backend_id]);
+ return;
+ }
+ }
+ }
+
if (input->flags & GGML_TENSOR_FLAG_INPUT) {
// inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done
if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
@@ -2122,6 +2140,62 @@ static void ggml_backend_sched_copy_input(ggml_backend_sched_t sched, struct ggm
}
}
+// the inputs of the next split that are computed on the backend of this split before it, and copied to host memory,
+// are copied before this split computes, so that the next split waits for these copies instead of this split
+// e.g. a CPU split runs while a GPU split computes, if it reads only the results of the earlier splits
+static void ggml_backend_sched_prefetch_inputs(ggml_backend_sched_t sched, int split_id) {
+ sched->n_prefetched = 0;
+ if (split_id + 1 >= sched->n_splits) {
+ return;
+ }
+ struct ggml_backend_sched_split * split = &sched->splits[split_id];
+ struct ggml_backend_sched_split * next = &sched->splits[split_id + 1];
+ const int backend_id = split->backend_id;
+ ggml_backend_t backend = sched->backends[backend_id];
+ if (backend->iface.get_tensor_async == NULL) {
+ return;
+ }
+
+ for (int i = 0; i < next->n_inputs && sched->n_prefetched < GGML_SCHED_MAX_SPLIT_INPUTS; i++) {
+ struct ggml_tensor * input = next->inputs[i];
+ if ((input->flags & GGML_TENSOR_FLAG_INPUT) || tensor_backend_id(input) != backend_id) {
+ continue;
+ }
+ struct ggml_tensor * input_cpy = tensor_copy(input, next->backend_id, sched->cur_copy);
+ if (input_cpy->buffer == NULL || !ggml_backend_buffer_is_host(input_cpy->buffer)) {
+ continue;
+ }
+
+ // not if this split computes or writes the input
+ const struct ggml_tensor * base = input->view_src ? input->view_src : input;
+ bool written = false;
+ for (int j = 0; j < split->graph.n_nodes && !written; j++) {
+ const struct ggml_tensor * node = split->graph.nodes[j];
+ written = !ggml_is_view_op(node->op) && (node == base || node->view_src == base);
+ }
+ if (written) {
+ continue;
+ }
+
+ if (!sched->prefetch_tried[backend_id]) {
+ sched->prefetch_tried[backend_id] = true;
+ sched->prefetch_events[backend_id] = ggml_backend_event_new(backend->device);
+ }
+ if (sched->prefetch_events[backend_id] == NULL) {
+ return;
+ }
+
+ ggml_backend_tensor_get_async(backend, input, input_cpy->data, 0, ggml_nbytes(input));
+ sched->prefetched[sched->n_prefetched++] = input;
+ }
+
+ if (sched->n_prefetched > 0) {
+ sched->prefetch_split = split_id + 1;
+ sched->prefetch_backend_id = backend_id;
+ ggml_backend_event_record(sched->prefetch_events[backend_id], backend);
+ }
+}
+
static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) {
GGML_ASSERT(sched);
struct ggml_backend_sched_split * splits = sched->splits;
@@ -2150,6 +2224,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
ggml_backend_sched_copy_input(sched, split, split->inputs[input_id]);
}
}
+ // after the copies above, which may read host memory reused by the copies of the next split
+ // before the weights, whose copy callback may queue large uploads
+ ggml_backend_sched_prefetch_inputs(sched, split_id);
for (int input_id = 0; input_id < split->n_inputs; input_id++) {
if (ggml_backend_sched_is_host_weight(split->inputs[input_id])) {
ggml_backend_sched_copy_input(sched, split, split->inputs[input_id]);
@@ -2286,6 +2363,7 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
for (int c = 0; c < sched->n_copies; c++) {
ggml_backend_event_free(sched->events[b][c]);
}
+ ggml_backend_event_free(sched->prefetch_events[b]);
}
ggml_gallocr_free(sched->galloc);
ggml_free(sched->ctx);
diff --git a/include/llama.h b/include/llama.h
index 3877b69e2..428f51f8d 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -397,7 +397,8 @@ extern "C" {
enum ggml_type type_k; // data type for K cache [EXPERIMENTAL]
enum ggml_type type_v; // data type for V cache [EXPERIMENTAL]
- size_t moe_cache_size; // device cache in bytes for the experts kept in host memory, split among the devices like the layers, 0 = disabled [EXPERIMENTAL]
+ size_t moe_cache_size; // device cache in bytes for the experts kept in host memory, split among the devices like the layers, 0 = disabled [EXPERIMENTAL]
+ float moe_cache_cpu_share; // share of the MoE cache misses computed on the CPU instead of uploaded, 0 = all uploaded [EXPERIMENTAL]
// Abort callback
// if it returns true, execution of llama_decode() will be aborted
diff --git a/scripts/moe-cache-tune.py b/scripts/moe-cache-tune.py
new file mode 100755
index 000000000..d0680fa74
--- /dev/null
+++ b/scripts/moe-cache-tune.py
@@ -0,0 +1,630 @@
+#!/usr/bin/env python3
+# Picks --moe-cache-mib and --moe-cache-cpu-share for a llama-server command (see docs/moe-cache.md):
+# python3 scripts/moe-cache-tune.py [--max-mib N] -- llama-server -m model.gguf -cmoe [args]
+# Runs the server at 2 cache sizes x 3 CPU shares, fits the time per token, and replays the recorded
+# routing through a copy of the cache logic (src/llama-moe-cache.cpp) to predict other sizes and shares.
+
+import argparse
+import json
+import logging
+import math
+import os
+import re
+import shlex
+import signal
+import socket
+import struct
+import subprocess
+import sys
+import time
+import urllib.request
+
+logger = logging.getLogger("moe-cache-tune")
+
+MAX_BATCH = 32 # llama_moe_cache::impl::max_batch
+SHARES = [0.0, 0.5, 0.8, 0.9, 0.95, 0.98]
+CALIB_SHARES = [0.0, 0.5, 0.95]
+N_FIRST = 2048 # tokens per request, first run
+N_CALIB = 512 # tokens per request, other runs
+PROMPT = ("The lighthouse keeper had kept a logbook for thirty-one years, and in all that time he had never once written down a lie. "
+ "So when the ship appeared on the horizon on the first calm morning of October, with no flag, no lights and no one on deck, "
+ "he sat at his small wooden table for a long while before he picked up the pen. Below him the sea was flat and grey, and the "
+ "gulls had gone quiet. He wrote the date, the hour and the direction of the wind, and then he stopped, because the next line "
+ "would have to say what he had seen, and he was no longer sure that he had seen it at all. This is what he finally wrote:")
+
+
+def f32(x):
+ return struct.unpack("<f", struct.pack("<f", x))[0]
+
+
+def pad(x, n):
+ return (x + n - 1) // n * n
+
+
+# routing trace (LLAMA_MOE_CACHE_TRACE)
+
+class Group:
+ n_expert: int
+ layers: list
+ split: float
+ split_sum: float
+ host_bytes: int
+ dev_host_bytes: int
+ alignment: int
+ nb2: list
+ expert_bytes: int
+
+
+def read_trace(path):
+ with open(path, "rb") as f:
+ data = f.read()
+ if data[:4] != b"MCT1":
+ sys.exit(f"{path}: not a MoE cache trace")
+ pos = 4
+
+ def take(fmt):
+ nonlocal pos
+ v = struct.unpack_from("<" + fmt, data, pos)
+ pos += struct.calcsize("<" + fmt)
+ return v
+
+ n_layer, n_expert_used, n_group = take("3I")
+ groups = []
+ for _ in range(n_group):
+ g = Group()
+ _, g.n_expert, n_layers = take("3I")
+ g.layers = list(take(f"{n_layers}I"))
+ g.split, g.split_sum = take("2d")
+ g.host_bytes, g.dev_host_bytes, g.alignment = take("3Q")
+ (n_tensors,) = take("I")
+ g.nb2 = list(take(f"{n_tensors}Q"))
+ g.expert_bytes = sum(g.nb2)
+ groups.append(g)
+
+ records = []
+ while pos < len(data):
+ il, n_used, n_ids = struct.unpack_from("<3i", data, pos)
+ pos += 12
+ records.append((il, n_used, struct.unpack_from(f"<{n_ids}i", data, pos)))
+ pos += 4 * n_ids
+ return n_layer, n_expert_used, groups, records
+
+
+def concat_traces(paths):
+ traces = [read_trace(p) for p in paths]
+ n_layer, n_expert_used, groups, _ = traces[0]
+ return n_layer, n_expert_used, groups, [rec for t in traces for rec in t[3]]
+
+
+def trace_tokens(trace):
+ first = min(il for il, _, _ in trace[3])
+ return sum(len(ids) // n_used for il, n_used, ids in trace[3] if il == first)
+
+
+# cache simulation, mirrors src/llama-moe-cache.cpp
+
+def n_slots_of(g, size_mib, n_zero, n_expert_used):
+ budget = int(float(size_mib * 1024 * 1024) * g.split / g.split_sum * float(g.host_bytes) / float(g.dev_host_bytes))
+
+ def alloc(n):
+ return sum(pad(nb2 * (n + n_zero + 1), g.alignment) for nb2 in g.nb2)
+
+ if alloc(0) > budget:
+ return 0
+ lo, hi = 0, len(g.layers) * g.n_expert
+ while lo < hi:
+ mid = (lo + hi + 1) // 2
+ if alloc(mid) <= budget:
+ lo = mid
+ else:
+ hi = mid - 1
+ return lo if lo >= n_expert_used else 0
+
+
+def full_mib(trace, share):
+ """smallest cache that holds all experts"""
+ _, n_expert_used, groups, _ = trace
+ n_zero = min(n_expert_used, 64) if f32(share) > 0.0 else 0
+ res = 0.0
+ for g in groups:
+ alloc = sum(pad(nb2 * (len(g.layers) * g.n_expert + n_zero + 1), g.alignment) for nb2 in g.nb2)
+ res = max(res, alloc / (g.split / g.split_sum * g.host_bytes / g.dev_host_bytes))
+ return math.ceil(res / 2**20) + 1
+
+
+class ClassLRU:
+ """moe_cache_lru, with its float32 rounding of the CRF"""
+ N_CLASS = 8
+ CRF_HALF_LIFE = 128.0 # tokens
+
+ def __init__(self, layers, n_expert, n_slots):
+ self.n_expert = n_expert
+ self.n_slots = n_slots
+ self.slot_map = {il: [n_slots] * n_expert for il in layers}
+ self.key_of = [-1] * n_slots
+ self.prev = [s - 1 for s in range(n_slots)]
+ self.next = [s + 1 if s + 1 < n_slots else -1 for s in range(n_slots)]
+ self.head = [0] + [-1] * (self.N_CLASS - 1)
+ self.tail = [n_slots - 1] + [-1] * (self.N_CLASS - 1)
+ self.list_of = [0] * n_slots
+ self.n_filled = 0
+ self.crf = {} # key -> CRF, 0 before the first use
+ self.last_use = {} # key -> layer clock at the last use
+ self.clock = {il: 0 for il in layers}
+ self.seen = [0] * n_expert
+ self.seen_gen = 0
+
+ def move(self, s, c):
+ frm = self.list_of[s]
+ if frm == c and s == self.tail[frm]:
+ return
+ prev, nxt, head, tail = self.prev, self.next, self.head, self.tail
+ if prev[s] >= 0:
+ nxt[prev[s]] = nxt[s]
+ else:
+ head[frm] = nxt[s]
+ if nxt[s] >= 0:
+ prev[nxt[s]] = prev[s]
+ else:
+ tail[frm] = prev[s]
+ prev[s] = tail[c]
+ nxt[s] = -1
+ if tail[c] >= 0:
+ nxt[tail[c]] = s
+ else:
+ head[c] = s
+ tail[c] = s
+ self.list_of[s] = c
+
+ def class_of(self, key):
+ return min(self.N_CLASS - 1, int(f32(math.log2(self.crf[key])))) # (int) log2f
+
+ def decayed_crf(self, il, key):
+ crf = self.crf.get(key, 0.0)
+ if crf == 0.0:
+ return 0.0
+ x = f32(f32(-(self.clock[il] - self.last_use[key])) / self.CRF_HALF_LIFE)
+ return f32(crf * f32(2.0 ** x))
+
+ def use(self, il, key):
+ self.crf[key] = f32(self.decayed_crf(il, key) + 1.0)
+ self.last_use[key] = self.clock[il]
+
+ def is_selected(self, il, key):
+ return key >= 0 and key // self.n_expert == il and self.seen[key % self.n_expert] == self.seen_gen
+
+ def victim(self, il):
+ res, res_score = -1, -1.0
+ for c in range(self.N_CLASS):
+ s = self.head[c]
+ while s >= 0 and self.is_selected(il, self.key_of[s]):
+ s = self.next[s]
+ if s < 0:
+ continue
+ key = self.key_of[s]
+ if key < 0:
+ return s # unused slot
+ score = math.ldexp(float(self.clock[key // self.n_expert] - self.last_use[key]), -c) # age/2^class
+ if score > res_score:
+ res, res_score = s, score
+ if res < 0:
+ raise RuntimeError("no MoE cache victim")
+ return res
+
+ def plan(self, il, ids, n_used, n_zero, n_cpu):
+ """returns (uploaded, computed on the CPU)"""
+ self.seen_gen += 1
+ seen, gen = self.seen, self.seen_gen
+ uniq = []
+ for e in ids:
+ if seen[e] != gen:
+ seen[e] = gen
+ uniq.append(e)
+ if len(uniq) > self.n_slots:
+ raise RuntimeError(f"layer {il}: {len(uniq)} experts for {self.n_slots} slots")
+
+ slots = self.slot_map[il]
+ n_slots = self.n_slots
+ n_tokens = len(ids) // n_used
+ misses = []
+ for e in uniq:
+ key = il * self.n_expert + e
+ self.use(il, key)
+ if slots[e] < n_slots:
+ self.move(slots[e], self.class_of(key))
+ else:
+ misses.append(e)
+
+ # the misses with the lowest CRF go to the CPU
+ n_c = min(n_cpu(len(misses)), len(misses))
+ n_cpu_done = 0
+ if n_c > 0:
+ misses.sort(key=lambda e: (self.crf[il * self.n_expert + e], e))
+ # distinct zero slots per token, a miss without one is uploaded
+ token_zero = [0] * n_tokens
+ n_keep = n_c
+ i = 0
+ while i < n_keep:
+ e = misses[i]
+ toks = [it for it in range(n_tokens) if e in ids[it * n_used:(it + 1) * n_used]]
+ used = 0
+ for it in toks:
+ used |= token_zero[it]
+ z = ((~used) & (used + 1)).bit_length() - 1
+ if z >= n_zero:
+ n_keep -= 1
+ misses[i], misses[n_keep] = misses[n_keep], misses[i]
+ continue
+ for it in toks:
+ token_zero[it] |= 1 << z
+ slots[e] = n_slots + z
+ n_cpu_done += 1
+ i += 1
+ del misses[:n_keep]
+
+ misses.sort()
+ for e in misses:
+ s = self.victim(il)
+ key = self.key_of[s]
+ if key >= 0:
+ self.slot_map[key // self.n_expert][key % self.n_expert] = n_slots
+ else:
+ self.n_filled += 1
+ key = il * self.n_expert + e
+ self.key_of[s] = key
+ slots[e] = s
+ self.move(s, self.class_of(key))
+ self.clock[il] += n_tokens
+ return len(misses), n_cpu_done
+
+
+def replay(trace, size_mib, share, skip_tokens):
+ """GiB of experts uploaded and computed on the CPU per token, after skip_tokens"""
+ n_layer, n_expert_used, groups, records = trace
+ share = f32(share)
+ share_q = int(math.floor(f32(share * 65536.0) + 0.5)) # std::lround(cpu_share*65536.0f)
+ n_zero = min(n_expert_used, 64) if share > 0.0 else 0
+ caches, group_of = [], {}
+ for ig, g in enumerate(groups):
+ n = n_slots_of(g, size_mib, n_zero, n_expert_used)
+ caches.append(ClassLRU(g.layers, g.n_expert, n) if n > 0 else None)
+ group_of.update((il, ig) for il in g.layers)
+
+ first = min(il for il, _, _ in records)
+ acc = [0] * n_layer
+ tokens = seen = up = cpu = 0
+ for il, n_used, ids in records:
+ n_tokens = len(ids) // n_used
+ if il == first:
+ seen += n_tokens
+ if seen > skip_tokens:
+ tokens += n_tokens
+ count = seen > skip_tokens
+ g, c = groups[group_of[il]], caches[group_of[il]]
+ if c is None or n_tokens > MAX_BATCH or min(n_tokens * n_used, g.n_expert) > c.n_slots:
+ if count:
+ cpu += len(set(ids)) * g.expert_bytes # not cached
+ continue
+
+ def n_cpu(n_miss):
+ v = n_miss * (share_q if c.n_filled == c.n_slots else 0) + acc[il] # the share applies once the cache is full
+ acc[il] = v & 0xFFFF
+ return v >> 16
+
+ n_fill, n_c = c.plan(il, ids, n_used, n_zero, n_cpu)
+ if count:
+ up += n_fill * g.expert_bytes
+ cpu += n_c * g.expert_bytes
+ tokens = max(tokens, 1)
+ return up / 2**30 / tokens, cpu / 2**30 / tokens
+
+
+# time model: ms/token = base + up*up_ms + cpu*cpu_ms + split*(share > 0) + overlap*min(up, cpu), up and cpu in GiB/token
+
+def model_ms(m, up, cpu, share):
+ return m["base"] + m["up"] * up + m["cpu"] * cpu + (m["split"] if share > 0 else 0.0) + m["overlap"] * min(up, cpu)
+
+
+def solve(A, b):
+ """least squares by the normal equations"""
+ n = len(A[0])
+ M: list = [[sum(A[k][i] * A[k][j] for k in range(len(A))) for j in range(n)] + [sum(A[k][i] * b[k] for k in range(len(A)))] for i in range(n)]
+ for i in range(n):
+ p = max(range(i, n), key=lambda r: abs(M[r][i]))
+ M[i], M[p] = M[p], M[i]
+ if abs(M[i][i]) < 1e-12:
+ sys.exit("the runs do not determine the time model")
+ for r in range(n):
+ if r != i:
+ f = M[r][i] / M[i][i]
+ M[r] = [x - f * y for x, y in zip(M[r], M[i])]
+ return [M[i][n] / M[i][i] for i in range(n)]
+
+
+def fit(runs):
+ """returns the model, the relative error per run and whether the model is physical"""
+ y = [r["ms"] / r["tokens"] for r in runs]
+ rows = [[1.0, r["up_gib"], r["cpu_gib"], float(r["share"] > 0), min(r["up_gib"], r["cpu_gib"])] for r in runs]
+
+ def solve_model(cols):
+ coef = dict(zip(cols, solve([[row[i] for i in cols.values()] for row in rows], y)))
+ if coef.get("split", 0) < 0: # noise
+ del cols["split"]
+ return solve_model(cols)
+ return {k: coef.get(k, 0.0) for k in ("base", "up", "cpu", "split", "overlap")}
+
+ def physical(m):
+ # overlap hides at most the cheaper path, the rest of the token is a good part of the fastest run
+ return m["up"] > 0 and m["cpu"] > 0 and m["base"] >= 0.2 * min(y) and -m["overlap"] <= min(m["up"], m["cpu"])
+
+ m = solve_model({"base": 0, "up": 1, "cpu": 2, "split": 3, "overlap": 4})
+ if not physical(m):
+ m = solve_model({"base": 0, "up": 1, "cpu": 2, "split": 3})
+ return m, [(model_ms(m, r["up_gib"], r["cpu_gib"], r["share"]) - yy) / yy for r, yy in zip(runs, y)], physical(m)
+
+
+# output
+
+FMT = {"markdown": False, "tty": False}
+
+
+def say(text=""):
+ print(text) # noqa: NP100
+
+
+def best(text):
+ if FMT["markdown"]:
+ return f"**{text}**"
+ return f"\033[1m{text}\033[0m" if FMT["tty"] else f"{text}*"
+
+
+def code(text):
+ return f"`{text}`" if FMT["markdown"] else text
+
+
+def table(header, rows):
+ rows = [[str(c) for c in row] for row in rows]
+ if FMT["markdown"]:
+ return "\n".join(["| " + " | ".join(header) + " |", "|" + "---:|" * len(header)] + ["| " + " | ".join(row) + " |" for row in rows])
+
+ def width(text):
+ return len(re.sub(r"\033\[[0-9;]*m", "", text))
+
+ w = [max(width(c) for c in col) for col in zip(header, *rows)]
+ return "\n".join(" " + " ".join(" " * (wi - width(c)) + c for c, wi in zip(row, w)) for row in [header] + rows)
+
+
+# model file
+
+GGUF_SCALAR = {0: "B", 1: "b", 2: "H", 3: "h", 4: "I", 5: "i", 6: "f", 7: "?", 10: "Q", 11: "q", 12: "d"}
+EXPS_RE = re.compile(r"\.ffn_(up|down|gate|gate_up)_(ch|)exps") # LLM_FFN_EXPS_REGEX
+
+
+def model_files(path):
+ m = re.match(r"(.*)-(\d{5})-of-(\d{5})\.gguf$", path)
+ return [f"{m.group(1)}-{i:05d}-of-{m.group(3)}.gguf" for i in range(1, int(m.group(3)) + 1)] if m else [path]
+
+
+def gguf_tensor_sizes(path):
+ """tensor sizes from the data offsets"""
+ with open(path, "rb") as f:
+ def take(fmt):
+ return struct.unpack("<" + fmt, f.read(struct.calcsize("<" + fmt)))
+
+ def string():
+ (n,) = take("Q")
+ return f.read(n).decode("utf-8", "replace")
+
+ def value(t):
+ if t == 8:
+ return string()
+ if t == 9:
+ et, n = take("IQ")
+ if et in (8, 9):
+ for _ in range(n):
+ value(et)
+ else:
+ f.seek(n * struct.calcsize("<" + GGUF_SCALAR[et]), 1)
+ return None
+ return take(GGUF_SCALAR[t])[0]
+
+ if f.read(4) != b"GGUF":
+ sys.exit(f"{path}: not a GGUF file")
+ take("I")
+ n_tensors, n_kv = take("QQ")
+ kv = {}
+ for _ in range(n_kv):
+ key = string()
+ kv[key] = value(take("I")[0])
+ infos = []
+ for _ in range(n_tensors):
+ name = string()
+ take(f"{take('I')[0]}Q")
+ take("I")
+ infos.append((take("Q")[0], name))
+ data = pad(f.tell(), kv.get("general.alignment", 32))
+ end = os.path.getsize(path) - data
+ infos.sort()
+ return {name: (infos[i + 1][0] if i + 1 < len(infos) else end) - off for i, (off, name) in enumerate(infos)}
+
+
+def expert_mib(path):
+ sizes = {}
+ for p in model_files(path):
+ sizes.update(gguf_tensor_sizes(p))
+ return sum(v for k, v in sizes.items() if EXPS_RE.search(k)) / 2**20
+
+
+# server runs
+
+def http(port, endpoint, body=None, timeout=3600):
+ req = urllib.request.Request(f"http://127.0.0.1:{port}{endpoint}", data=json.dumps(body).encode() if body else None,
+ headers={"Content-Type": "application/json"})
+ with urllib.request.urlopen(req, timeout=timeout) as resp:
+ return json.loads(resp.read())
+
+
+def cache_stats(log):
+ """the statistics the cache logs at exit (-lv 4)"""
+ r = {"hits": 0, "misses": 0, "cpu": 0, "up_mib": 0.0}
+ for m in re.finditer(r"llama_moe_cache: ubatch .. 8: hits = (\d+), misses = (\d+)(?: \((\d+) on the CPU\))?, hit rate = [0-9.]+%, uploaded = ([0-9.]+) MiB", log):
+ r["hits"] += int(m.group(1))
+ r["misses"] += int(m.group(2))
+ r["cpu"] += int(m.group(3) or 0)
+ r["up_mib"] += float(m.group(4))
+ return r
+
+
+def device_room(log):
+ """host experts, cache device, its free memory at load and the memory the model uses on it"""
+ host = sum(float(x) for x in re.findall(r"MoE cache size =\s+[0-9.]+ MiB for ([0-9.]+) MiB of host experts", log))
+ m = re.search(r"(\S+) MoE cache size =", log)
+ if host == 0 or not m:
+ sys.exit("no MoE cache: the experts must be in host memory (-cmoe, -ncmoe or -ot)")
+ free = {d: int(x) for d, x in re.findall(r"using device (\S+) \(.*- (\d+) MiB free", log)}
+ used = sum(float(x) for d, x in re.findall(r"(\S+) +(?:model|KV|RS|compute|output) +buffer size =\s+([0-9.]+) MiB", log) if d == m.group(1))
+ if m.group(1) not in free:
+ sys.exit(f"the free memory of {m.group(1)} is not in the log, pass --max-mib")
+ return host, m.group(1), free[m.group(1)], used
+
+
+def run_server(cmd, out, mib, share, n_predict, seed, label, warm):
+ """two generations with the cache, reused if already in out"""
+ base = os.path.join(out, f"run-{mib}-{share:g}")
+ if os.path.isfile(base + ".json") and os.path.isfile(base + ".trace"):
+ r = json.load(open(base + ".json"))
+ else:
+ warm()
+ with socket.socket() as sock:
+ sock.bind(("127.0.0.1", 0))
+ port = sock.getsockname()[1]
+ args = cmd + ["--host", "127.0.0.1", "--port", str(port), "-lv", "4", "--moe-cache-mib", str(mib), "--moe-cache-cpu-share", str(share)]
+ timings: list = []
+ for attempt in range(2):
+ with open(base + ".log", "w") as log:
+ proc = subprocess.Popen(args, stdout=log, stderr=subprocess.STDOUT, env=dict(os.environ, LLAMA_MOE_CACHE_TRACE=base + ".trace"))
+ try:
+ while True:
+ if proc.poll() is not None:
+ sys.exit(f"llama-server exited ({proc.returncode}), see {base}.log")
+ try:
+ if http(port, "/health", timeout=5).get("status") == "ok":
+ break
+ except OSError:
+ pass
+ time.sleep(1)
+ n = min(n_predict, http(port, "/props")["default_generation_settings"]["n_ctx"] - 256)
+ timings = [http(port, "/completion", {"prompt": PROMPT, "n_predict": n, "ignore_eos": True, "cache_prompt": False,
+ "temperature": 0.8, "top_p": 0.95, "seed": seed + k})["timings"] for k in (0, 1)]
+ finally:
+ proc.send_signal(signal.SIGINT)
+ try:
+ proc.wait(timeout=120)
+ except subprocess.TimeoutExpired:
+ proc.kill()
+ # the second generation has a warmer cache, it is much slower only if the machine is busy
+ if timings[1]["predicted_per_second"] > timings[0]["predicted_per_second"] / 1.5 or attempt == 1:
+ break
+ logger.info(" slower second generation, machine busy? running again")
+ r = cache_stats(open(base + ".log", errors="replace").read())
+ r.update(mib=mib, share=share, ms=sum(t["predicted_ms"] for t in timings), tokens=sum(t["predicted_n"] for t in timings))
+ json.dump(r, open(base + ".json", "w"))
+ logger.info(f"{label} {mib} MiB, share {share:g}: {1000 * r['tokens'] / r['ms']:.1f} t/s")
+ r["hit"] = 100 * r["hits"] / max(r["hits"] + r["misses"], 1)
+ return r
+
+
+def tune(cmd, out, max_mib):
+ model = next((cmd[i + 1] for i, a in enumerate(cmd[:-1]) if a in ("-m", "--model")), None)
+ if model is None:
+ sys.exit("the server command needs -m")
+ os.makedirs(out, exist_ok=True)
+ command = os.path.join(out, "command.txt")
+ if os.path.isfile(command) and open(command).read() != shlex.join(cmd):
+ for f in os.listdir(out):
+ os.remove(os.path.join(out, f))
+ open(command, "w").write(shlex.join(cmd))
+
+ load_mode = next((cmd[i + 1] for i, a in reversed(list(enumerate(cmd[:-1]))) if a in ("-lm", "--load-mode")), os.environ.get("LLAMA_ARG_LOAD_MODE", "auto"))
+ warmed = []
+
+ def warm():
+ # read an mmap'd model once, so that no run reads it from disk
+ if not warmed and load_mode in ("auto", "mmap", "mmap+mlock"):
+ for p in model_files(model):
+ with open(p, "rb") as f:
+ while f.read(64 << 20):
+ pass
+ warmed.append(True)
+
+ def rnd(x):
+ return max(64, int(x) // 64 * 64)
+
+ n = 2 * len(CALIB_SHARES)
+ small = rnd(0.10 * expert_mib(model))
+ runs = [run_server(cmd, out, small, 0.0, N_FIRST, 1, f"1/{n}", warm)]
+ host, dev, free, used = device_room(open(os.path.join(out, f"run-{small}-0.log"), errors="replace").read())
+ room = int(free - used - 512)
+ cap = min(max_mib, room) if max_mib else room
+ large = rnd(min(0.30 * host, room))
+ if large <= small:
+ sys.exit(f"no room on {dev} for a cache larger than {small} MiB")
+ for i, (mib, share) in enumerate([(small, s) for s in CALIB_SHARES[1:]] + [(large, s) for s in CALIB_SHARES]):
+ runs.append(run_server(cmd, out, mib, share, N_CALIB, 2 * i + 3, f"{i + 2}/{n}", warm))
+
+ trace = concat_traces([os.path.join(out, f"run-{r['mib']}-{r['share']:g}.trace") for r in runs])
+ groups = trace[2]
+ gib = sum(g.expert_bytes * len(g.layers) for g in groups) / sum(len(g.layers) for g in groups) / 2**30
+ for r in runs:
+ r["up_gib"], r["cpu_gib"] = r["up_mib"] / 1024 / r["tokens"], r["cpu"] * gib / r["tokens"]
+ m, errs, ok = fit(runs)
+
+ full = full_mib(trace, max(SHARES))
+ top = min(cap, full)
+ sizes = sorted({rnd(min(f * host, top)) for f in (0.05, 0.10, 0.20, 0.30, 0.40, 0.50)} | {small, large, full if cap >= full else rnd(top)})
+ skip = trace_tokens(trace) // 2 # warm-up
+ rows, pick = [], {}
+ for mib in [s for s in sizes if s <= top]:
+ ms = [model_ms(m, *replay(trace, mib, share, skip), share) for share in SHARES]
+ k = min(i for i in range(len(ms)) if ms[i] <= 1.05 * min(ms)) # lowest share within 5% of the fastest
+ rows.append([mib] + [best(f"{1000 / x:.1f}") if i == k else f"{1000 / x:.1f}" for i, x in enumerate(ms)])
+ pick[mib] = (SHARES[k], 1000 / ms[k], 1000 / ms[0])
+
+ say(f"{code(os.path.basename(model))}: {host:.0f} MiB of experts, {dev} {free} MiB free, cache up to {cap} MiB")
+ say()
+ say(table(["cache MiB", "share", "t/s", "hits"], [[r["mib"], f"{r['share']:g}", f"{1000 * r['tokens'] / r['ms']:.1f}", f"{r['hit']:.0f}%"] for r in runs]))
+ say()
+ say(f"machine: uploads {1000 / m['up']:.1f} GiB/s, CPU {1000 / m['cpu']:.1f} GiB/s, CPU part {m['split']:.2f} ms, overlap {m['overlap']:.1f} ms/GiB,"
+ f" rest {m['base']:.2f} ms per token, fit within {100 * max(abs(e) for e in errs):.1f}%")
+ say()
+ say("predicted t/s, warm cache:")
+ say(table(["cache MiB"] + [f"{s:g}" for s in SHARES], rows))
+ say()
+ if not ok or max(abs(e) for e in errs) > 0.10:
+ sys.exit(f"the runs do not fit the time model, was the machine busy? delete {out} and run again")
+ if cap >= full:
+ say(f"all experts fit on {dev} ({full} MiB): drop {code('-cmoe')}, or tune a smaller cache with {code('--max-mib')}")
+ return
+ share, tps, tps0 = pick[max(pick)]
+ say(f"recommended: {code(f'--moe-cache-mib {max(pick)} --moe-cache-cpu-share {share:g}')} ({tps:.1f} t/s" + (f", {tps0:.1f} with share 0)" if share > 0 else ")"))
+
+
+def main():
+ ap = argparse.ArgumentParser(description="picks --moe-cache-mib and --moe-cache-cpu-share for a llama-server command",
+ usage="%(prog)s [--max-mib N] [--out DIR] [--markdown] -- llama-server -m model.gguf -cmoe [args]")
+ ap.add_argument("--max-mib", type=int, help="largest cache in MiB (default: what fits)")
+ ap.add_argument("--out", default="moe-cache-tune", help="run directory, reused for the same command (default: moe-cache-tune)")
+ ap.add_argument("--markdown", action="store_true", help="Markdown tables")
+ ap.add_argument("cmd", nargs=argparse.REMAINDER, help="llama-server command")
+ args = ap.parse_args()
+ cmd = args.cmd[1:] if args.cmd[:1] == ["--"] else args.cmd
+ if not cmd:
+ ap.error("pass the llama-server command after --")
+ logging.basicConfig(level=logging.INFO, format="%(message)s")
+ FMT["markdown"] = args.markdown
+ FMT["tty"] = sys.stdout.isatty() and not args.markdown
+ tune(cmd, args.out, args.max_mib)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index f6e7e9ca2..a450d8257 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -391,7 +391,8 @@ llama_context::llama_context(
cparams.op_offload = params.op_offload;
cparams.kv_unified = params.kv_unified;
- cparams.moe_cache_size = params.moe_cache_size;
+ cparams.moe_cache_size = params.moe_cache_size;
+ cparams.moe_cache_cpu_share = params.moe_cache_cpu_share;
// resolved after flash attention, see resolve_fused_ops()
{
@@ -612,7 +613,7 @@ llama_context::llama_context(
}
if (cparams.moe_cache_size > 0) {
- moe_cache = std::make_unique<llama_moe_cache>(model, model.no_alloc(), backend_ptrs, backend_buft, cparams.moe_cache_size);
+ moe_cache = std::make_unique<llama_moe_cache>(model, model.no_alloc(), backend_ptrs, backend_buft, cparams.moe_cache_size, cparams.moe_cache_cpu_share);
}
sched_reserve();
@@ -4273,6 +4274,7 @@ llama_context_params llama_context_default_params() {
/*.type_k =*/ GGML_TYPE_F16,
/*.type_v =*/ GGML_TYPE_F16,
/*.moe_cache_size =*/ 0,
+ /*.moe_cache_cpu_share =*/ 0.0f,
/*.abort_callback =*/ nullptr,
/*.abort_callback_data =*/ nullptr,
/*.embeddings =*/ false,
diff --git a/src/llama-cparams.h b/src/llama-cparams.h
index 9141cba3e..9d5fe2b3a 100644
--- a/src/llama-cparams.h
+++ b/src/llama-cparams.h
@@ -60,6 +60,7 @@ struct llama_cparams {
bool training; // set by llama_opt_init()
size_t moe_cache_size;
+ float moe_cache_cpu_share;
std::vector<bool> embeddings_layer_inp; // [n_layer()] extract input embeddings for layer
bool embeddings_layer_inp_dev; // write enabled layer inputs to a device buffer instead of host
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index d1ee7c090..9f99ba83e 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -19,6 +19,8 @@
#include <cassert>
#include <cmath>
+#include <algorithm>
+#include <array>
#include <cstring>
#include <numeric>
#include <sstream>
@@ -1753,13 +1755,14 @@ ggml_tensor * llm_graph_context::build_lora_mm_id(
ggml_tensor * cur, // ggml_tensor * b
ggml_tensor * ids,
ggml_tensor * w_s,
- ggml_tensor * slots) const {
+ ggml_tensor * slots,
+ bool cached) const {
// as in build_lora_mm: the folded weight reads the transformed input, the LoRA delta reads `cur`
ggml_tensor * cur_mm = build_hadamard_input(w, cur);
// the experts in the MoE cache are selected by their slots
ggml_tensor * res = slots == nullptr ?
ggml_mul_mat_id(ctx0, w, cur_mm, ids) :
- ggml_mul_mat_id(ctx0, moe_cache->get_experts(w), cur_mm, slots);
+ ggml_mul_mat_id(ctx0, cached ? moe_cache->get_experts(w) : w, cur_mm, slots);
if (prec_policy) {
prec_policy->apply(res);
@@ -2410,8 +2413,12 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
//call early so that topk-moe can be used
ggml_build_forward_expand(gf, weights);
- // the experts of host-resident layers may be read from the MoE cache
- ggml_tensor * slots = build_moe_cache_slots(selected_experts, up_exps, gate_exps, down_exps, gate_up_exps, il);
+ // 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);
cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens);
@@ -2422,145 +2429,146 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
cb(cur, "ffn_moe_weighted", il);
}
- ggml_tensor * up = nullptr;
- ggml_tensor * experts = nullptr;
+ // the experts selected by mm_ids, up to the down projection, from the MoE cache if cached
+ // ids are the selected experts, for the biases and scales
+ auto build_experts = [&](ggml_tensor * cur, ggml_tensor * ids, ggml_tensor * mm_ids, bool cached) {
+ ggml_tensor * up = nullptr;
- if (gate_up_exps) {
- // merged gate_up path: one mul_mat_id, then split into gate and up views
- ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s, slots); // [n_ff*2, n_expert_used, n_tokens]
- cb(gate_up, "ffn_moe_gate_up", il);
+ if (gate_up_exps) {
+ // merged gate_up path: one mul_mat_id, then split into gate and up views
+ ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, ids, up_exps_s, mm_ids, cached); // [n_ff*2, n_expert_used, n_tokens]
+ cb(gate_up, "ffn_moe_gate_up", il);
- if (up_exps_s) {
- cb(gate_up, "ffn_moe_gate_up_scaled", il);
- }
+ if (up_exps_s) {
+ cb(gate_up, "ffn_moe_gate_up_scaled", il);
+ }
- if (gate_up_exps_b) {
- gate_up = ggml_add_id(ctx0, gate_up, gate_up_exps_b, selected_experts);
- cb(gate_up, "ffn_moe_gate_up_biased", il);
- }
+ if (gate_up_exps_b) {
+ gate_up = ggml_add_id(ctx0, gate_up, gate_up_exps_b, ids);
+ cb(gate_up, "ffn_moe_gate_up_biased", il);
+ }
- const int64_t n_ff = gate_up->ne[0] / 2;
- cur = ggml_view_3d(ctx0, gate_up, n_ff, gate_up->ne[1], gate_up->ne[2], gate_up->nb[1], gate_up->nb[2], 0);
- cb(cur, "ffn_moe_gate", il);
- up = ggml_view_3d(ctx0, gate_up, n_ff, gate_up->ne[1], gate_up->ne[2], gate_up->nb[1], gate_up->nb[2], n_ff * gate_up->nb[0]);
- cb(up, "ffn_moe_up", il);
- } else {
- // separate gate and up path
- up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s, slots); // [n_ff, n_expert_used, n_tokens]
- cb(up, "ffn_moe_up", il);
+ const int64_t n_ff = gate_up->ne[0] / 2;
+ cur = ggml_view_3d(ctx0, gate_up, n_ff, gate_up->ne[1], gate_up->ne[2], gate_up->nb[1], gate_up->nb[2], 0);
+ cb(cur, "ffn_moe_gate", il);
+ up = ggml_view_3d(ctx0, gate_up, n_ff, gate_up->ne[1], gate_up->ne[2], gate_up->nb[1], gate_up->nb[2], n_ff * gate_up->nb[0]);
+ cb(up, "ffn_moe_up", il);
+ } else {
+ // separate gate and up path
+ up = build_lora_mm_id(up_exps, cur, ids, up_exps_s, mm_ids, cached); // [n_ff, n_expert_used, n_tokens]
+ cb(up, "ffn_moe_up", il);
- if (up_exps_s) {
- cb(up, "ffn_moe_up_scaled", il);
- }
+ if (up_exps_s) {
+ cb(up, "ffn_moe_up_scaled", il);
+ }
- if (up_exps_b) {
- up = ggml_add_id(ctx0, up, up_exps_b, selected_experts);
- cb(up, "ffn_moe_up_biased", il);
- }
+ if (up_exps_b) {
+ up = ggml_add_id(ctx0, up, up_exps_b, ids);
+ cb(up, "ffn_moe_up_biased", il);
+ }
- if (gate_exps) {
- cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s, slots); // [n_ff, n_expert_used, n_tokens]
- cb(cur, "ffn_moe_gate", il);
- } else {
- cur = up;
- }
+ if (gate_exps) {
+ cur = build_lora_mm_id(gate_exps, cur, ids, gate_exps_s, mm_ids, cached); // [n_ff, n_expert_used, n_tokens]
+ cb(cur, "ffn_moe_gate", il);
+ } else {
+ cur = up;
+ }
- if (gate_exps_s) {
- cb(cur, "ffn_moe_gate_scaled", il);
- }
+ if (gate_exps_s) {
+ cb(cur, "ffn_moe_gate_scaled", il);
+ }
- if (gate_exps_b) {
- cur = ggml_add_id(ctx0, cur, gate_exps_b, selected_experts);
- cb(cur, "ffn_moe_gate_biased", il);
+ if (gate_exps_b) {
+ cur = ggml_add_id(ctx0, cur, gate_exps_b, ids);
+ cb(cur, "ffn_moe_gate_biased", il);
+ }
}
- }
- const bool has_gate = gate_exps || gate_up_exps;
-
- switch (type_op) {
- case LLM_FFN_SILU:
- if (gate_exps) {
- if (il >= 0) {
- const float limit = hparams.swiglu_clamp_exp[il];
- constexpr float eps = 1e-6f;
- if (limit > eps) {
- if (arch == LLM_ARCH_MAPLE || arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) {
- cur = ggml_swiglu_clamp(ctx0, cur, up, limit);
- } else {
- up = ggml_clamp(ctx0, up, -limit, limit);
- cb(up, "ffn_moe_up_clamped", il);
- ggml_tensor * gate_act = ggml_silu(ctx0, cur);
- cb(gate_act, "ffn_moe_silu", il);
- gate_act = ggml_clamp(ctx0, gate_act, -INFINITY, limit);
- cb(gate_act, "ffn_moe_silu_clamped", il);
- cur = ggml_mul(ctx0, gate_act, up);
+ const bool has_gate = gate_exps || gate_up_exps;
+
+ switch (type_op) {
+ case LLM_FFN_SILU:
+ if (gate_exps) {
+ if (il >= 0) {
+ const float limit = hparams.swiglu_clamp_exp[il];
+ constexpr float eps = 1e-6f;
+ if (limit > eps) {
+ if (arch == LLM_ARCH_MAPLE || arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) {
+ cur = ggml_swiglu_clamp(ctx0, cur, up, limit);
+ } else {
+ up = ggml_clamp(ctx0, up, -limit, limit);
+ cb(up, "ffn_moe_up_clamped", il);
+ ggml_tensor * gate_act = ggml_silu(ctx0, cur);
+ cb(gate_act, "ffn_moe_silu", il);
+ gate_act = ggml_clamp(ctx0, gate_act, -INFINITY, limit);
+ cb(gate_act, "ffn_moe_silu_clamped", il);
+ cur = ggml_mul(ctx0, gate_act, up);
+ }
+ cb(cur, "ffn_moe_swiglu_limited", il);
+ break;
}
- cb(cur, "ffn_moe_swiglu_limited", il);
- break;
}
}
- }
- if (has_gate) {
- cur = ggml_swiglu_split(ctx0, cur, up);
- cb(cur, "ffn_moe_swiglu", il);
- } else {
- cur = ggml_silu(ctx0, cur);
- cb(cur, "ffn_moe_silu", il);
- } break;
- case LLM_FFN_SITU:
- {
- // situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * lb*tanh(up/lb)
- GGML_ASSERT(has_gate);
- const float beta = hparams.situ_beta;
- const float lb = hparams.situ_linear_beta;
-
- ggml_tensor * act = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, cur, 1.0f/beta)), beta);
- act = ggml_mul(ctx0, act, ggml_sigmoid(ctx0, cur));
- if (lb > 0.0f) {
- up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/lb)), lb);
- }
- cur = ggml_mul(ctx0, act, up);
- cb(cur, "ffn_moe_situ", il);
- } break;
- case LLM_FFN_GELU:
- if (has_gate) {
- cur = ggml_geglu_split(ctx0, cur, up);
- cb(cur, "ffn_moe_geglu", il);
- } else {
- cur = ggml_gelu(ctx0, cur);
- cb(cur, "ffn_moe_gelu", il);
- } break;
- case LLM_FFN_SWIGLU_OAI_MOE:
- {
- // TODO: move to hparams?
- constexpr float alpha = 1.702f;
- constexpr float limit = 7.0f;
- cur = ggml_swiglu_oai(ctx0, cur, up, alpha, limit);
- cb(cur, "ffn_moe_swiglu_oai", il);
- } break;
- case LLM_FFN_RELU:
- if (has_gate) {
- cur = ggml_reglu_split(ctx0, cur, up);
- cb(cur, "ffn_moe_reglu", il);
- } else {
- cur = ggml_relu(ctx0, cur);
- cb(cur, "ffn_moe_relu", il);
- } break;
- case LLM_FFN_RELU_SQR:
- if (has_gate) {
- // TODO: add support for gated squared relu
- GGML_ABORT("fatal error: gated squared relu not implemented");
- } else {
- cur = ggml_relu(ctx0, cur);
- cur = ggml_sqr(ctx0, cur);
- cb(cur, "ffn_moe_relu_sqr", il);
- } break;
- default:
- GGML_ABORT("fatal error");
- }
+ if (has_gate) {
+ cur = ggml_swiglu_split(ctx0, cur, up);
+ cb(cur, "ffn_moe_swiglu", il);
+ } else {
+ cur = ggml_silu(ctx0, cur);
+ cb(cur, "ffn_moe_silu", il);
+ } break;
+ case LLM_FFN_SITU:
+ {
+ // situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * lb*tanh(up/lb)
+ GGML_ASSERT(has_gate);
+ const float beta = hparams.situ_beta;
+ const float lb = hparams.situ_linear_beta;
+
+ ggml_tensor * act = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, cur, 1.0f/beta)), beta);
+ act = ggml_mul(ctx0, act, ggml_sigmoid(ctx0, cur));
+ if (lb > 0.0f) {
+ up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/lb)), lb);
+ }
+ cur = ggml_mul(ctx0, act, up);
+ cb(cur, "ffn_moe_situ", il);
+ } break;
+ case LLM_FFN_GELU:
+ if (has_gate) {
+ cur = ggml_geglu_split(ctx0, cur, up);
+ cb(cur, "ffn_moe_geglu", il);
+ } else {
+ cur = ggml_gelu(ctx0, cur);
+ cb(cur, "ffn_moe_gelu", il);
+ } break;
+ case LLM_FFN_SWIGLU_OAI_MOE:
+ {
+ // TODO: move to hparams?
+ constexpr float alpha = 1.702f;
+ constexpr float limit = 7.0f;
+ cur = ggml_swiglu_oai(ctx0, cur, up, alpha, limit);
+ cb(cur, "ffn_moe_swiglu_oai", il);
+ } break;
+ case LLM_FFN_RELU:
+ if (has_gate) {
+ cur = ggml_reglu_split(ctx0, cur, up);
+ cb(cur, "ffn_moe_reglu", il);
+ } else {
+ cur = ggml_relu(ctx0, cur);
+ cb(cur, "ffn_moe_relu", il);
+ } break;
+ case LLM_FFN_RELU_SQR:
+ if (has_gate) {
+ // TODO: add support for gated squared relu
+ GGML_ABORT("fatal error: gated squared relu not implemented");
+ } else {
+ cur = ggml_relu(ctx0, cur);
+ cur = ggml_sqr(ctx0, cur);
+ cb(cur, "ffn_moe_relu_sqr", il);
+ } break;
+ default:
+ GGML_ABORT("fatal error");
+ }
- {
// TODO: this graph-level rescaling should only be applied when the corresponding backend does not support
// ggml_prec_set_src(experts, GGML_PREC_BF16, 1);
// the amax op cannot be split across devices, so these architectures are denied -sm tensor [TAG_MOE_AMAX_SCALE]
@@ -2572,13 +2580,24 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
cur = amax_src1.res;
}
- experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s, slots); // [n_embd, n_expert_used, n_tokens]
+ ggml_tensor * experts = build_lora_mm_id(down_exps, cur, ids, down_exps_s, mm_ids, cached); // [n_embd, n_expert_used, n_tokens]
//ggml_prec_set_src(experts, GGML_PREC_BF16, 1); // (see TODO above)
cb(experts, "ffn_moe_down", il);
if (amax_src1.scale_s) {
experts = ggml_mul(ctx0, experts, amax_src1.scale_s);
}
+
+ 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);
}
if (down_exps_s) {
@@ -2638,24 +2657,26 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
return moe_out;
}
-ggml_tensor * llm_graph_context::build_moe_cache_slots(
+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;
+ return { nullptr, nullptr };
}
- ggml_tensor * slot_map = moe_cache->get_slot_map(il, selected_experts->ne[1], selected_experts->ne[0]);
+ 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;
+ return { nullptr, nullptr };
}
for (ggml_tensor * w : { up_exps, gate_exps, down_exps, gate_up_exps }) {
if (w != nullptr && moe_cache->get_experts(w) == nullptr) {
- return nullptr;
+ return { nullptr, nullptr };
}
}
@@ -2669,7 +2690,7 @@ ggml_tensor * llm_graph_context::build_moe_cache_slots(
// 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;
+ return { nullptr, nullptr };
}
ggml_backend_sched_set_tensor_backend(sched, slots, moe_cache->backend(il));
cb(slots, "ffn_moe_slots", il);
@@ -2677,7 +2698,14 @@ ggml_tensor * llm_graph_context::build_moe_cache_slots(
// 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);
- return ggml_reshape_2d(ctx0, slots, selected_experts->ne[0], selected_experts->ne[1]); // [n_expert_used, n_tokens]
+ // 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
diff --git a/src/llama-graph.h b/src/llama-graph.h
index 6d25dab8c..9867c9540 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -1023,6 +1023,11 @@ 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;
@@ -1121,13 +1126,14 @@ 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 read from the MoE cache at these slots (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 build_moe_cache_slots)
ggml_tensor * build_lora_mm_id(
ggml_tensor * w, // ggml_tensor * as
ggml_tensor * cur, // ggml_tensor * b
ggml_tensor * ids,
- ggml_tensor * w_s = nullptr,
- ggml_tensor * slots = nullptr) const;
+ ggml_tensor * w_s = nullptr,
+ ggml_tensor * slots = nullptr,
+ bool cached = true) const;
// rescale `cur` (the src1 of a matmul) by a power of two, computed from its maximum absolute
// keep both individual values and Q8 activation block sums within F16 range
@@ -1234,12 +1240,14 @@ struct llm_graph_context {
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
- ggml_tensor * build_moe_cache_slots(
+ // 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;
//
diff --git a/src/llama-moe-cache.cpp b/src/llama-moe-cache.cpp
index b736930f9..965f6a177 100644
--- a/src/llama-moe-cache.cpp
+++ b/src/llama-moe-cache.cpp
@@ -8,6 +8,9 @@
#include <algorithm>
#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
#include <stdexcept>
#include <unordered_map>
#include <vector>
@@ -15,7 +18,7 @@
namespace {
// 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
+// slot_map[il][e] is the slot of expert e of layer il, n_slots or later 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
@@ -25,8 +28,11 @@ struct moe_cache_lru {
int32_t n_expert = 0;
int32_t n_slots = 0;
+ int32_t empty = 0; // slot of the experts that are not cached
- std::vector<int32_t *> slot_map; // [n_layer] data of the slot maps, -1 if the expert is not cached
+ // [n_layer] data of the slot maps, the experts at slot n_slots or later are not cached
+ // the experts computed on the CPU read zero slots after the cached ones on the device
+ std::vector<int32_t *> slot_map;
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
@@ -35,6 +41,7 @@ struct moe_cache_lru {
int32_t head[n_class];
int32_t tail[n_class];
std::vector<int8_t> list_of; // [n_slots] class list of the slot
+ int32_t n_filled = 0; // slots that hold an expert
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
@@ -43,11 +50,14 @@ struct moe_cache_lru {
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; // the distinct selected experts
+ std::vector<int32_t> misses;
+ std::vector<uint64_t> token_zero; // [n_tokens] zero slots used by each token
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;
+ this->empty = n_slots;
slot_map.assign(n_layer, nullptr);
key_of.assign(n_slots, -1);
prev.resize(n_slots);
@@ -170,9 +180,14 @@ struct moe_cache_lru {
};
// give a slot to each expert selected by ids (n_tokens tokens) in layer il
+ // n_cpu(n) of the n misses are left to the CPU and read a zero slot on the device, they are not cached
+ // a zero slot is used at most once per token, a miss that does not get one is cached instead
// 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, int64_t n_tokens, std::vector<fill> & fills, size_t & n_hit) {
+ template <typename F>
+ bool plan(int32_t il, const int32_t * ids, size_t n_ids, int64_t n_tokens, int32_t n_used, int32_t n_zero, F && n_cpu,
+ std::vector<fill> & fills, std::vector<fill> & cpu, size_t & n_hit) {
fills.clear();
+ cpu.clear();
n_hit = 0;
if (++seen_gen == 0) {
@@ -194,24 +209,71 @@ struct moe_cache_lru {
int32_t * slots = slot_map[il];
// update the CRF of all selected experts and promote hits before choosing victims for misses
+ misses.clear();
for (int32_t e : uniq) {
const int32_t key = make_key(il, e);
use(il, key);
- if (slots[e] >= 0) {
+ if (slots[e] < n_slots) {
move(slots[e], class_of(key));
n_hit++;
+ } else {
+ misses.push_back(e);
+ }
+ }
+
+ // the misses with the lowest CRF are left to the CPU, a miss that keeps being used gains CRF and enters the cache
+ // all misses are used at the same clock, so their CRF compare without decay
+ const size_t n_c = std::min(n_cpu(misses.size()), misses.size());
+ if (n_c > 0) {
+ std::sort(misses.begin(), misses.end(), [&](int32_t a, int32_t b) {
+ const float crf_a = crf[make_key(il, a)];
+ const float crf_b = crf[make_key(il, b)];
+ return crf_a != crf_b ? crf_a < crf_b : a < b;
+ });
+ GGML_ASSERT((size_t) n_tokens*n_used == n_ids);
+ token_zero.assign(n_tokens, 0);
+ size_t n_keep = n_c;
+ for (size_t i = 0; i < n_keep; ) {
+ const int32_t e = misses[i];
+ uint64_t used = 0;
+ for (int64_t it = 0; it < n_tokens; ++it) {
+ for (int32_t k = 0; k < n_used; ++k) {
+ if (ids[it*n_used + k] == e) {
+ used |= token_zero[it];
+ }
+ }
+ }
+ int32_t z = 0;
+ while (z < 64 && (used >> z & 1)) {
+ z++;
+ }
+ if (z >= n_zero) {
+ std::swap(misses[i], misses[--n_keep]);
+ continue;
+ }
+ for (int64_t it = 0; it < n_tokens; ++it) {
+ for (int32_t k = 0; k < n_used; ++k) {
+ if (ids[it*n_used + k] == e) {
+ token_zero[it] |= 1ull << z;
+ }
+ }
+ }
+ slots[e] = n_slots + z;
+ cpu.push_back({ e, n_slots + z });
+ i++;
}
+ misses.erase(misses.begin(), misses.begin() + n_keep);
}
+
// 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) {
- continue;
- }
+ std::sort(misses.begin(), misses.end());
+ for (int32_t e : misses) {
const int32_t s = victim(il);
if (key_of[s] >= 0) {
- slot_map[key_layer(key_of[s])][key_expert(key_of[s])] = -1;
+ slot_map[key_layer(key_of[s])][key_expert(key_of[s])] = empty;
+ } else {
+ n_filled++;
}
const int32_t key = make_key(il, e);
key_of[s] = key;
@@ -283,6 +345,10 @@ struct llama_moe_cache::impl {
struct layer {
int32_t ig = -1; // -1 if the layer is not cached
ggml_tensor * slot_map = nullptr; // I32 [1, n_expert] in host memory
+ ggml_tensor * slot_map_cpu = nullptr; // looked up in place of slot_map by the graphs with a CPU part, if cpu_share > 0, its data is not used
+ ggml_tensor * cpu_map = nullptr; // I32 [1, n_expert] in host memory, if cpu_share > 0
+ std::vector<int32_t> cpu_experts; // experts set in cpu_map
+ uint32_t cpu_acc = 0; // fraction of a miss carried over to the next plan, in 1/65536
std::vector<ggml_tensor *> experts; // host expert tensors, in the order of the banks
};
@@ -295,6 +361,7 @@ struct llama_moe_cache::impl {
struct stats {
size_t hits = 0;
size_t misses = 0;
+ size_t cpu = 0; // misses computed on the CPU
size_t bytes = 0;
};
@@ -310,10 +377,18 @@ struct llama_moe_cache::impl {
std::vector<group> groups;
std::vector<layer> layers;
std::unordered_map<const ggml_tensor *, binding> bindings; // host experts -> cached experts
- std::unordered_map<const ggml_tensor *, int32_t> layer_of; // slot map -> layer
+ std::unordered_map<const ggml_tensor *, std::pair<int32_t, bool>> layer_of; // slot map -> layer, and if the graph that looks it up has a CPU part
std::vector<int32_t> ids;
std::vector<moe_cache_lru::fill> fills;
+ std::vector<moe_cache_lru::fill> cpu;
+
+ float cpu_share = 0.0f;
+ uint32_t cpu_share_q = 0; // cpu_share in 1/65536, integer so that a simulator can reproduce the rounding
+ int32_t n_zero = 0; // zero slots per group
+
+ // LLAMA_MOE_CACHE_TRACE: the selected experts of each slot lookup, for a simulator of the cache
+ FILE * trace = nullptr;
// slot maps in host memory
ggml_context_ptr ctx_host;
@@ -323,8 +398,8 @@ struct llama_moe_cache::impl {
// views used by copy_experts
ggml_context_ptr ctx_views;
- impl(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) :
- n_expert_used(model.hparams.n_expert_used_max()), layers(model.layers.size()) {
+ impl(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) :
+ n_expert_used(model.hparams.n_expert_used_max()), layers(model.layers.size()), cpu_share(std::clamp(cpu_share, 0.0f, 1.0f)) {
for (size_t i = 0; i < backends.size(); ++i) {
const auto dev_type = ggml_backend_dev_type(ggml_backend_get_device(backends[i]));
if (dev_type == GGML_BACKEND_DEVICE_TYPE_GPU || dev_type == GGML_BACKEND_DEVICE_TYPE_IGPU) {
@@ -372,12 +447,14 @@ struct llama_moe_cache::impl {
return;
}
+ // with misses on the CPU, the n_zero slots after the cached ones stay zero: the device computes these misses as zero
// one extra slot at the end, CUDA MMQ can read past the last expert
+ n_zero = cpu_share > 0.0f ? std::min<int32_t>(n_expert_used, 64) : 0;
auto alloc_size = [&](const group & g, int32_t n_slots) {
const size_t alignment = ggml_backend_buft_get_alignment(devices[g.id].buft);
size_t res = 0;
for (const ggml_tensor * t : g.ref) {
- res += GGML_PAD(t->nb[2]*(n_slots + 1), alignment);
+ res += GGML_PAD(t->nb[2]*(n_slots + n_zero + 1), alignment);
}
return res;
};
@@ -413,6 +490,14 @@ struct llama_moe_cache::impl {
}
// within a device the budget is split by the size of the experts, so each group caches the same fraction of its experts
+ cpu_share_q = (uint32_t) std::lround(cpu_share*65536.0f);
+ if (const char * path = getenv("LLAMA_MOE_CACHE_TRACE"); path != nullptr && *path != '\0') {
+ trace = fopen(path, "wb");
+ if (trace == nullptr) {
+ throw std::runtime_error(std::string("failed to open the MoE cache trace ") + path);
+ }
+ write_trace_header(split_sum);
+ }
std::vector<size_t> n_tensors(devices.size(), 0);
size_t n_tensors_host = 0;
for (group & g : groups) {
@@ -430,7 +515,7 @@ struct llama_moe_cache::impl {
}
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();
+ n_tensors_host += g.layers.size()*(n_zero > 0 ? 3 : 1); // slot maps and CPU map
}
if (n_tensors_host == 0) {
throw std::runtime_error("MoE cache is too small to hold the experts of one token");
@@ -466,7 +551,7 @@ struct llama_moe_cache::impl {
}
ggml_context * ctx = devices[g.id].ctx.get();
for (const ggml_tensor * t : g.ref) {
- ggml_tensor * bank = ggml_new_tensor_3d(ctx, t->type, t->ne[0], t->ne[1], g.n_slots + 1);
+ ggml_tensor * bank = ggml_new_tensor_3d(ctx, t->type, t->ne[0], t->ne[1], g.n_slots + n_zero + 1);
GGML_ASSERT(bank->nb[2] == t->nb[2]);
ggml_format_name(bank, "moe_cache.%zu.%s", ig, t->name);
g.banks.push_back(bank);
@@ -477,14 +562,23 @@ struct llama_moe_cache::impl {
l.experts = llama_moe_cache_layer_experts(model.layers[il]);
for (size_t ip = 0; ip < l.experts.size(); ++ip) {
ggml_tensor * bank = g.banks[ip];
- ggml_tensor * cached = ggml_view_3d(ctx, bank, bank->ne[0], bank->ne[1], g.n_slots, bank->nb[1], bank->nb[2], 0);
+ ggml_tensor * cached = ggml_view_3d(ctx, bank, bank->ne[0], bank->ne[1], g.n_slots + n_zero, bank->nb[1], bank->nb[2], 0);
ggml_format_name(cached, "moe_cache.%s", l.experts[ip]->name);
bindings[l.experts[ip]] = { il, (int32_t) ip, cached };
}
l.slot_map = ggml_new_tensor_2d(ctx_host.get(), GGML_TYPE_I32, 1, g.ref[0]->ne[2]);
ggml_format_name(l.slot_map, "moe_cache.slot_map-%d", il);
- layer_of[l.slot_map] = il;
+ layer_of[l.slot_map] = { il, false };
buf_host_size += GGML_PAD(ggml_nbytes(l.slot_map), alignment_host);
+ if (n_zero > 0) {
+ l.slot_map_cpu = ggml_new_tensor_2d(ctx_host.get(), GGML_TYPE_I32, 1, g.ref[0]->ne[2]);
+ ggml_format_name(l.slot_map_cpu, "moe_cache.slot_map_cpu-%d", il);
+ layer_of[l.slot_map_cpu] = { il, true };
+ buf_host_size += GGML_PAD(ggml_nbytes(l.slot_map_cpu), alignment_host);
+ l.cpu_map = ggml_new_tensor_2d(ctx_host.get(), GGML_TYPE_I32, 1, g.ref[0]->ne[2]);
+ ggml_format_name(l.cpu_map, "moe_cache.cpu_map-%d", il);
+ buf_host_size += GGML_PAD(ggml_nbytes(l.cpu_map), alignment_host);
+ }
}
devices[g.id].buf_size += alloc_size(g, g.n_slots);
}
@@ -502,13 +596,15 @@ struct llama_moe_cache::impl {
}
buf_host = llama_memory_bufs::alloc_buf(ctx_host.get(), buft_host, no_alloc, "MoE cache slot maps");
if (!no_alloc) {
- ggml_backend_buffer_clear(buf_host.get(), 0xff); // all slots are -1
+ ggml_backend_buffer_clear(buf_host.get(), 0xff); // the CPU maps are -1
buf_host_size = ggml_backend_buffer_get_size(buf_host.get());
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;
+ int32_t * slots = (int32_t *) layers[il].slot_map->data;
+ std::fill(slots, slots + layers[il].slot_map->ne[1], g.lru.empty);
+ g.lru.slot_map[il] = slots;
}
}
}
@@ -527,8 +623,8 @@ struct llama_moe_cache::impl {
if (d.host_bytes == 0) {
continue;
}
- LLAMA_LOG_INFO("%s: %10s MoE cache size = %8.2f MiB for %.2f MiB of host experts\n", __func__,
- ggml_backend_buft_name(d.buft), d.buf_size/1024.0/1024.0, d.host_bytes/1024.0/1024.0);
+ LLAMA_LOG_INFO("%s: %10s MoE cache size = %8.2f MiB for %.2f MiB of host experts, %.0f%% of the misses on the CPU\n", __func__,
+ ggml_backend_buft_name(d.buft), d.buf_size/1024.0/1024.0, d.host_bytes/1024.0/1024.0, 100.0*cpu_share);
for (const group & g : groups) {
if (g.id == (int32_t) id) {
LLAMA_LOG_INFO("%s: %2zu layers, %s: %5d slots (%.1f%%, %.1f tokens worth)\n", __func__,
@@ -543,11 +639,52 @@ struct llama_moe_cache::impl {
return devices[groups[layers[il].ig].id].backend;
}
+ // trace format, little endian: "MCT1", then u32 n_layer, n_expert_used, n_group
+ // per group: u32 device, n_expert, n_layers, the layers (u32), f64 split, split_sum, u64 group host bytes, device host bytes, alignment,
+ // u32 n_tensors, the bytes of one expert of each tensor (u64)
+ // then one record per slot lookup: i32 layer, experts per token, n_ids, the ids (i32)
+ void write_trace_header(double split_sum) {
+ auto u32 = [&](uint32_t v) { fwrite(&v, sizeof(v), 1, trace); };
+ auto u64 = [&](uint64_t v) { fwrite(&v, sizeof(v), 1, trace); };
+ auto f64 = [&](double v) { fwrite(&v, sizeof(v), 1, trace); };
+ fwrite("MCT1", 1, 4, trace);
+ u32(layers.size());
+ u32(n_expert_used);
+ u32(groups.size());
+ for (const group & g : groups) {
+ const device & d = devices[g.id];
+ u32(g.id);
+ u32(g.ref[0]->ne[2]);
+ u32(g.layers.size());
+ for (int32_t il : g.layers) {
+ u32(il);
+ }
+ f64(d.split);
+ f64(split_sum);
+ u64(g.host_bytes);
+ u64(d.host_bytes);
+ u64(ggml_backend_buft_get_alignment(d.buft));
+ u32(g.ref.size());
+ for (const ggml_tensor * t : g.ref) {
+ u64(t->nb[2]);
+ }
+ }
+ }
+
+ void write_trace(int32_t il, int32_t n_used) {
+ const int32_t rec[3] = { il, n_used, (int32_t) ids.size() };
+ fwrite(rec, sizeof(int32_t), 3, trace);
+ fwrite(ids.data(), sizeof(int32_t), ids.size(), trace);
+ }
+
~impl() {
log_stats();
+ if (trace != nullptr) {
+ fclose(trace);
+ }
}
- ggml_tensor * get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const {
+ 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;
}
@@ -557,7 +694,28 @@ struct llama_moe_cache::impl {
if (n_tokens == 0 || n_tokens > max_batch || std::min(n_tokens*n_expert_used, l.slot_map->ne[1]) > groups[l.ig].n_slots) {
return nullptr;
}
- return l.slot_map;
+ GGML_ASSERT(!cpu_part || l.slot_map_cpu != nullptr);
+ return cpu_part ? l.slot_map_cpu : l.slot_map;
+ }
+
+ 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;
+ }
+ // the scheduler would run the CPU part on a device that offloads such batches from host weights
+ ggml_tensor probe = {};
+ probe.op = GGML_OP_MUL_MAT_ID;
+ probe.type = GGML_TYPE_F32;
+ probe.ne[0] = 1;
+ probe.ne[1] = n_expert_used;
+ probe.ne[2] = n_tokens;
+ probe.ne[3] = 1;
+ for (const device & d : devices) {
+ if (ggml_backend_offload_op(d.backend, &probe)) {
+ return nullptr;
+ }
+ }
+ return layers[il].cpu_map;
}
ggml_tensor * get_experts(const ggml_tensor * w) const {
@@ -578,11 +736,11 @@ 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];
- if (slots == nullptr || slots[e] < 0) {
+ if (slots == nullptr || slots[e] >= g.n_slots) {
return 0;
}
int64_t n = 1;
- while (e + n <= last && slots[e + n] == slots[e] + n) {
+ while (e + n <= last && slots[e + n] < g.n_slots && slots[e + n] == slots[e] + n) {
n++;
}
@@ -605,8 +763,9 @@ struct llama_moe_cache::impl {
if (it == layer_of.end()) {
return false;
}
- const int32_t il = it->second;
- const layer & l = layers[il];
+ const int32_t il = it->second.first;
+ const bool has_cpu = it->second.second;
+ layer & l = layers[il];
group & g = groups[l.ig];
GGML_ASSERT(backend == devices[g.id].backend);
@@ -634,14 +793,38 @@ 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
+ const int32_t n_used = (int32_t) sel->src[0]->ne[0];
+
ids.resize(ggml_nelements(sel));
ggml_backend_tensor_get_async(backend, sel, ids.data(), 0, ggml_nbytes(sel));
ggml_backend_synchronize(backend);
+ if (trace != nullptr) {
+ write_trace(il, n_used);
+ }
+
+ // part of the misses is computed on the CPU once the cache is full, the carry keeps the share right on average
+ const uint32_t share_q = has_cpu && g.lru.n_filled == g.lru.n_slots ? cpu_share_q : 0;
+ auto n_cpu = [&](size_t n_miss) {
+ const uint64_t acc = (uint64_t) n_miss*share_q + l.cpu_acc;
+ l.cpu_acc = acc & 0xffff;
+ return (size_t) (acc >> 16);
+ };
size_t n_hit = 0;
- if (!g.lru.plan(il, ids.data(), ids.size(), sel->src[0]->ne[1], fills, n_hit)) {
+ if (!g.lru.plan(il, ids.data(), ids.size(), sel->src[0]->ne[1], n_used, n_zero, n_cpu, fills, cpu, n_hit)) {
GGML_ABORT("the MoE cache is too small for the experts selected in layer %d", il);
}
+ if (has_cpu) {
+ int32_t * cpu_map = (int32_t *) l.cpu_map->data;
+ for (int32_t e : l.cpu_experts) {
+ cpu_map[e] = -1;
+ }
+ l.cpu_experts.clear();
+ for (const auto & f : cpu) {
+ cpu_map[f.expert] = f.expert;
+ l.cpu_experts.push_back(f.expert);
+ }
+ }
// upload the missing experts, consecutive experts going to consecutive slots are uploaded together
size_t bytes = 0;
@@ -662,11 +845,12 @@ struct llama_moe_cache::impl {
stats & st = ids.size() <= (size_t) 8*n_expert_used ? stats_small : stats_large;
st.hits += n_hit;
- st.misses += fills.size();
+ st.misses += fills.size() + cpu.size();
+ st.cpu += cpu.size();
st.bytes += bytes;
// the next copy synchronizes the backend before it changes the slot map again
- ggml_backend_tensor_set_async(backend, dst, src->data, 0, ggml_nbytes(src));
+ ggml_backend_tensor_set_async(backend, dst, l.slot_map->data, 0, ggml_nbytes(l.slot_map));
return true;
}
@@ -677,8 +861,8 @@ struct llama_moe_cache::impl {
if (n == 0) {
return;
}
- LLAMA_LOG_INFO("llama_moe_cache: %s: hits = %zu, misses = %zu, hit rate = %.2f%%, uploaded = %.2f MiB\n",
- name, st.hits, st.misses, 100.0*st.hits/n, st.bytes/1024.0/1024.0);
+ LLAMA_LOG_INFO("llama_moe_cache: %s: hits = %zu, misses = %zu (%zu on the CPU), hit rate = %.2f%%, uploaded = %.2f MiB\n",
+ name, st.hits, st.misses, st.cpu, 100.0*st.hits/n, st.bytes/1024.0/1024.0);
};
log("ubatch <= 8", stats_small);
log("ubatch > 8", stats_large);
@@ -688,8 +872,8 @@ struct llama_moe_cache::impl {
}
};
-llama_moe_cache::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) :
- pimpl(new impl(model, no_alloc, backends, bufts, size)) {
+llama_moe_cache::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) :
+ pimpl(new impl(model, no_alloc, backends, bufts, size, cpu_share)) {
}
llama_moe_cache::~llama_moe_cache() = default;
@@ -698,14 +882,18 @@ ggml_backend_t llama_moe_cache::backend(int32_t il) const {
return pimpl->backend(il);
}
-ggml_tensor * llama_moe_cache::get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const {
- return pimpl->get_slot_map(il, n_tokens, n_expert_used);
+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);
}
ggml_tensor * llama_moe_cache::get_experts(const ggml_tensor * w) const {
return pimpl->get_experts(w);
}
+ggml_tensor * llama_moe_cache::get_cpu_map(int32_t il, int64_t n_tokens) const {
+ return pimpl->get_cpu_map(il, n_tokens);
+}
+
bool llama_moe_cache::copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph) {
return pimpl->copy(backend, src, dst, graph);
}
diff --git a/src/llama-moe-cache.h b/src/llama-moe-cache.h
index 6c8e36582..636bc4559 100644
--- a/src/llama-moe-cache.h
+++ b/src/llama-moe-cache.h
@@ -13,18 +13,24 @@ struct llama_model;
class llama_moe_cache {
public:
// backends are all the backends of the context, each GPU gets its own cache of the given size for the layers assigned to it
- 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);
+ // cpu_share is the share of the misses computed on the CPU instead of uploaded
+ 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
- ggml_tensor * get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const;
+ // 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;
+
// 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);
diff --git a/tools/cli/README.md b/tools/cli/README.md
index f1f6e6ed4..8e9f057f8 100644
--- a/tools/cli/README.md
+++ b/tools/cli/README.md
@@ -64,6 +64,7 @@
| `-cmoe, --cpu-moe` | keep all Mixture of Experts (MoE) weights in the CPU<br/>(env: LLAMA_ARG_CPU_MOE) |
| `-ncmoe, --n-cpu-moe N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU<br/>(env: LLAMA_ARG_N_CPU_MOE) |
| `--moe-cache-mib N` | GPU cache size in MiB for the MoE experts kept in the CPU. with multiple GPUs, it is split among them like the layers (--tensor-split) (default: 0, disabled)<br/>(env: LLAMA_ARG_MOE_CACHE_MIB) |
+| `--moe-cache-cpu-share F` | share of the MoE cache misses that the CPU computes from the host experts instead of uploading them to the GPU once the cache is full, from 0 to 1 (default: 0.00)<br/>(env: LLAMA_ARG_MOE_CACHE_CPU_SHARE) |
| `-ncffn, --n-cpu-ffn N` | keep the dense FFN weights of the first N layers in the CPU<br/>(dense models; for MoE expert weights use --n-cpu-moe)<br/>(env: LLAMA_ARG_N_CPU_FFN) |
| `-ngl, --gpu-layers, --n-gpu-layers N` | max. number of layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)<br/>(env: LLAMA_ARG_N_GPU_LAYERS) |
| `-sm, --split-mode {none,layer,row,tensor}` | how to split the model across multiple GPUs, one of:<br/>- none: use one GPU only<br/>- layer (default): split layers and KV across GPUs (pipelined)<br/>- row: split weight across GPUs by rows (parallelized)<br/>- tensor: split weights and KV across GPUs (parallelized, EXPERIMENTAL)<br/>(env: LLAMA_ARG_SPLIT_MODE) |
diff --git a/tools/completion/README.md b/tools/completion/README.md
index 7174a04ff..9b824f968 100644
--- a/tools/completion/README.md
+++ b/tools/completion/README.md
@@ -147,6 +147,7 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1
| `-cmoe, --cpu-moe` | keep all Mixture of Experts (MoE) weights in the CPU<br/>(env: LLAMA_ARG_CPU_MOE) |
| `-ncmoe, --n-cpu-moe N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU<br/>(env: LLAMA_ARG_N_CPU_MOE) |
| `--moe-cache-mib N` | GPU cache size in MiB for the MoE experts kept in the CPU. with multiple GPUs, it is split among them like the layers (--tensor-split) (default: 0, disabled)<br/>(env: LLAMA_ARG_MOE_CACHE_MIB) |
+| `--moe-cache-cpu-share F` | share of the MoE cache misses that the CPU computes from the host experts instead of uploading them to the GPU once the cache is full, from 0 to 1 (default: 0.00)<br/>(env: LLAMA_ARG_MOE_CACHE_CPU_SHARE) |
| `-ncffn, --n-cpu-ffn N` | keep the dense FFN weights of the first N layers in the CPU<br/>(dense models; for MoE expert weights use --n-cpu-moe)<br/>(env: LLAMA_ARG_N_CPU_FFN) |
| `-ngl, --gpu-layers, --n-gpu-layers N` | max. number of layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)<br/>(env: LLAMA_ARG_N_GPU_LAYERS) |
| `-sm, --split-mode {none,layer,row,tensor}` | how to split the model across multiple GPUs, one of:<br/>- none: use one GPU only<br/>- layer (default): split layers and KV across GPUs (pipelined)<br/>- row: split weight across GPUs by rows (parallelized)<br/>- tensor: split weights and KV across GPUs (parallelized, EXPERIMENTAL)<br/>(env: LLAMA_ARG_SPLIT_MODE) |
diff --git a/tools/server/README.md b/tools/server/README.md
index afa448b8a..c249e05d4 100644
--- a/tools/server/README.md
+++ b/tools/server/README.md
@@ -81,6 +81,7 @@ For the full list of features, please refer to [server's changelog](https://gith
| `-cmoe, --cpu-moe` | keep all Mixture of Experts (MoE) weights in the CPU<br/>(env: LLAMA_ARG_CPU_MOE) |
| `-ncmoe, --n-cpu-moe N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU<br/>(env: LLAMA_ARG_N_CPU_MOE) |
| `--moe-cache-mib N` | GPU cache size in MiB for the MoE experts kept in the CPU. with multiple GPUs, it is split among them like the layers (--tensor-split) (default: 0, disabled)<br/>(env: LLAMA_ARG_MOE_CACHE_MIB) |
+| `--moe-cache-cpu-share F` | share of the MoE cache misses that the CPU computes from the host experts instead of uploading them to the GPU once the cache is full, from 0 to 1 (default: 0.00)<br/>(env: LLAMA_ARG_MOE_CACHE_CPU_SHARE) |
| `-ncffn, --n-cpu-ffn N` | keep the dense FFN weights of the first N layers in the CPU<br/>(dense models; for MoE expert weights use --n-cpu-moe)<br/>(env: LLAMA_ARG_N_CPU_FFN) |
| `-ngl, --gpu-layers, --n-gpu-layers N` | max. number of layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)<br/>(env: LLAMA_ARG_N_GPU_LAYERS) |
| `-sm, --split-mode {none,layer,row,tensor}` | how to split the model across multiple GPUs, one of:<br/>- none: use one GPU only<br/>- layer (default): split layers and KV across GPUs (pipelined)<br/>- row: split weight across GPUs by rows (parallelized)<br/>- tensor: split weights and KV across GPUs (parallelized, EXPERIMENTAL)<br/>(env: LLAMA_ARG_SPLIT_MODE) |