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
[email protected]
Co-authored-by: Georgi Gerganov
[email protected]
Co-authored-by: Alexey Kopytko
[email protected]