llama.cpp
e95426f1 - llama_moe_cache: run missed experts on CPU, add tuning script (#30966)

Commit
3 days ago
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>
Author
Parents
Loading