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>

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

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"

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;
 }

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;
 

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.

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,

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);

ggml/src/ggml-cpu/ops.cpp

@@ -9591,7 +9591,6 @@ void ggml_compute_forward_flash_attn_ext(
     switch (dst->op_params[3]) {
         case GGML_PREC_DEFAULT:
         case GGML_PREC_F32:
-        case GGML_PREC_BF16:
             {
                 // uses F32 accumulators
                 ggml_compute_forward_flash_attn_ext_f16(params, dst);

ggml/src/ggml-cpu/spacemit/ime.cpp

@@ -1180,7 +1180,7 @@ class tensor_traits_common : public tensor_traits_base {
         const int64_t DK = nek0;
         const int64_t DV = nev0;
 
-        const bool supported_prec  = (dst->op_params[3] == GGML_PREC_F32 || dst->op_params[3] == GGML_PREC_BF16 || dst->op_params[3] == GGML_PREC_DEFAULT);
+        const bool supported_prec  = (dst->op_params[3] == GGML_PREC_F32 || dst->op_params[3] == GGML_PREC_DEFAULT);
         const bool supported_types = (q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16);
         const bool supported_shape = (DK > 0 && DK <= 128 && DV > 0 && DV <= 128);
         const bool supported_vlen  = (__riscv_vlenb() == 128);

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -5978,7 +5978,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                 if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16 && !ggml_cuda_op_mul_mat_use_fwht(op)) {
                     return false;
                 }
-                if (op->op == GGML_OP_MUL_MAT_ID && (ggml_get_op_params_i32(op, 3) == GGML_PREC_BF16 || ggml_get_op_params_i32(op, 3) == GGML_PREC_F32)) {
+                if (op->op == GGML_OP_MUL_MAT_ID && ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) {
                     return false;
                 }
 #ifdef GGML_USE_MUSA

ggml/src/ggml-et/ggml-et-ops.cpp

@@ -1523,7 +1523,7 @@ bool ggml_et_op_flash_attn_ext(ggml_backend_et_device_context * dev_ctx, const g
     }
 
     const ggml_prec prec = ggml_flash_attn_ext_get_prec(node);
-    if (prec != GGML_PREC_F32 && prec != GGML_PREC_BF16 && prec != GGML_PREC_DEFAULT) {
+    if (prec != GGML_PREC_F32 && prec != GGML_PREC_DEFAULT) {
         GGML_LOG_ERROR("ET: FLASH_ATTN_EXT baseline kernel only supports F32 precision\n");
         return false;
     }

ggml/src/ggml-et/ggml-et.cpp

@@ -1299,7 +1299,7 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm
                 const bool me_eligible = op->src[1]->type == GGML_TYPE_F16 && op->src[2]->type == GGML_TYPE_F16 &&
                                          (op->src[0]->ne[0] % 32) == 0;
 
-                supported = me_eligible && mask_ok && (prec == GGML_PREC_F32 || prec == GGML_PREC_BF16 || prec == GGML_PREC_DEFAULT) &&
+                supported = me_eligible && mask_ok && (prec == GGML_PREC_F32 || prec == GGML_PREC_DEFAULT) &&
                             max_bias == 0.0f && logit_softcap == 0.0f && op->src[0]->nb[0] == sizeof(float) &&
                             op->src[1]->nb[0] == k_elem && op->src[2]->nb[0] == v_elem && op->nb[0] == sizeof(float) &&
                             op->src[0]->ne[0] == op->src[1]->ne[0] &&  // dk matches

ggml/src/ggml-hexagon/htp-kparams.cpp

@@ -944,7 +944,7 @@ static bool htp_matmul_check_types(
         if (src1->type != GGML_TYPE_F32 || !src2 || src2->type != GGML_TYPE_I32) {
             return false;
         }
-        if (ggml_get_op_params_i32(dst, 3) == GGML_PREC_BF16 || ggml_get_op_params_i32(dst, 3) == GGML_PREC_F32) {
+        if (ggml_get_op_params_i32(dst, 3) == GGML_PREC_F32) {
             return false;
         }
         switch (src0->type) {

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -2241,17 +2241,34 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_l2_norm(ggml_met
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_group_norm(ggml_metal_library_t lib, const ggml_tensor * op) {
     assert(op->op == GGML_OP_GROUP_NORM);
 
-    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
+    GGML_ASSERT(ggml_is_contiguous_rows(op->src[0]));
 
     char base[256];
     char name[256];
 
     snprintf(base, 256, "kernel_group_norm_f32");
-    snprintf(name, 256, "%s", base);
+
+    const ggml_tensor * src0 = op->src[0];
+
+    const bool flat = src0->nb[1] == src0->ne[0]*src0->nb[0] &&
+                      src0->nb[2] == src0->nb[1]*src0->ne[1] &&
+                      op->nb[1]   == op->ne[0] *op->nb[0]    &&
+                      op->nb[2]   == op->nb[1] *op->ne[1];
+
+    if (flat) {
+        snprintf(name, 256, "%s", base);
+    } else {
+        snprintf(name, 256, "%s_strided", base);
+    }
 
     ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
     if (!res.pipeline) {
-        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+        ggml_metal_cv_t cv = ggml_metal_cv_init();
+        ggml_metal_cv_set_bool(cv, flat, FC_GROUP_NORM + 0);
+
+        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
+
+        ggml_metal_cv_free(cv);
     }
 
     res.smem = 32*sizeof(float);

ggml/src/ggml-metal/ggml-metal-device.m

@@ -1681,9 +1681,13 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
         case GGML_OP_CUMSUM:
         case GGML_OP_MEAN:
         case GGML_OP_SOFT_MAX:
-        case GGML_OP_GROUP_NORM:
         case GGML_OP_L2_NORM:
             return has_simdgroup_reduction && ggml_is_contiguous_rows(op->src[0]);
+        case GGML_OP_GROUP_NORM:
+            return has_simdgroup_reduction &&
+                op->src[0]->type == GGML_TYPE_F32 &&
+                op->type == GGML_TYPE_F32 &&
+                ggml_is_contiguous_rows(op->src[0]);
         case GGML_OP_COUNT_EQUAL:
             return has_simdgroup_reduction &&
                 op->src[0]->type == GGML_TYPE_I32 &&

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -126,6 +126,7 @@
 #define FC_LIGHTNING_INDEXER           2200
 #define FC_MUL_MV_MMA                  2300
 #define FC_TOP_K_MASK                  2400
+#define FC_GROUP_NORM                  2500
 
 // op-specific constants
 #define OP_FLASH_ATTN_EXT_NQPSG 8
@@ -677,6 +678,10 @@ typedef struct {
     uint64_t nb00;
     uint64_t nb01;
     uint64_t nb02;
+    uint64_t nb03;
+    uint64_t nb1;
+    uint64_t nb2;
+    uint64_t nb3;
     int32_t  ngrp;
     float    eps;
 } ggml_metal_kargs_group_norm;

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -4610,6 +4610,10 @@ int ggml_metal_op_group_norm(ggml_metal_op_t ctx, int idx) {
         /*.nb00 =*/ nb00,
         /*.nb01 =*/ nb01,
         /*.nb02 =*/ nb02,
+        /*.nb03 =*/ nb03,
+        /*.nb1  =*/ nb1,
+        /*.nb2  =*/ nb2,
+        /*.nb3  =*/ nb3,
         /*.ngrp =*/ ngrp,
         /*.eps  =*/ eps,
     };
@@ -4633,7 +4637,7 @@ int ggml_metal_op_group_norm(ggml_metal_op_t ctx, int idx) {
 
     ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
 
-    ggml_metal_encoder_dispatch_threadgroups(enc, ngrp, 1, 1, nth, 1, 1);
+    ggml_metal_encoder_dispatch_threadgroups(enc, ngrp, 1, ne03, nth, 1, 1);
 
     return 1;
 }

ggml/src/ggml-metal/kernels/norm.metal

@@ -236,37 +236,59 @@ typedef decltype(kernel_l2_norm_impl<float, float>) kernel_l2_norm_t;
 template [[host_name("kernel_l2_norm_f32_f32")]]   kernel kernel_l2_norm_t kernel_l2_norm_impl<float,  float>;
 template [[host_name("kernel_l2_norm_f32_f32_4")]] kernel kernel_l2_norm_t kernel_l2_norm_impl<float4, float4>;
 
+constant bool FC_group_norm_flat [[function_constant(FC_GROUP_NORM + 0)]];
+
 kernel void kernel_group_norm_f32(
         constant ggml_metal_kargs_group_norm & args,
-        device const float * src0,
-        device       float * dst,
+        device const char * src0,
+        device       char * dst,
         threadgroup float  * buf [[threadgroup(0)]],
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint sgitg[[simdgroup_index_in_threadgroup]],
-        uint tiisg[[thread_index_in_simdgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    const int64_t ne = args.ne00*args.ne01*args.ne02;
-    const int64_t gs = args.ne00*args.ne01*((args.ne02 + args.ngrp - 1) / args.ngrp);
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const int64_t ne00 = args.ne00;
+    const int64_t ne01 = args.ne01;
+    const int64_t ne02 = args.ne02;
 
-    int start = tgpig * gs;
-    int end   = start + gs;
+    const int64_t ngrp = args.ngrp;
 
-    start += tpitg;
+    const int64_t i03 = tgpig.z; // batch index
+    const int64_t ig  = tgpig.x; // group index
 
-    if (end >= ne) {
-        end = ne;
-    }
+    const int64_t cpg = (ne02 + ngrp - 1)/ngrp;
+
+    const int64_t c0 =  ig*cpg;
+    const int64_t c1 = (c0 + cpg < ne02) ? c0 + cpg : ne02;
+
+    const int64_t npf  = ne00*ne01;      // elements per channel
+    const int64_t gs   = npf*(c1 - c0);  // elements in this group
+
+    device const char * x03 = src0 + i03*args.nb03 + c0*args.nb02;
+    device       char * y03 = dst  + i03*args.nb3  + c0*args.nb2;
 
     float tmp = 0.0f; // partial sum for thread in warp
 
-    for (int j = start; j < end; j += ntg) {
-        tmp += src0[j];
+    if (FC_group_norm_flat) {
+        device const float * x = (device const float *) x03;
+        for (int64_t e = tpitg.x; e < gs; e += ntg.x) {
+            tmp += x[e];
+        }
+    } else {
+        for (int64_t e = tpitg.x; e < gs; e += ntg.x) {
+            const int64_t i02 = c0 + e/npf;
+            const int64_t ir  = e - (i02 - c0)*npf;
+            const int64_t i01 = ir/ne00;
+            const int64_t i00 = ir - i01*ne00;
+
+            tmp += *(device const float *)(x03 + (i02 - c0)*args.nb02 + i01*args.nb01 + i00*sizeof(float));
+        }
     }
 
     threadgroup_barrier(mem_flags::mem_threadgroup);
     tmp = simd_sum(tmp);
-    if (ntg > N_SIMDWIDTH) {
+    if (ntg.x > N_SIMDWIDTH) {
         if (sgitg == 0) {
             buf[tiisg] = 0.0f;
         }
@@ -286,14 +308,31 @@ kernel void kernel_group_norm_f32(
     const float mean = tmp / gs;
     tmp = 0.0f;
 
-    for (int j = start; j < end; j += ntg) {
-        float xi = src0[j] - mean;
-        dst[j] = xi;
-        tmp += xi * xi;
+    if (FC_group_norm_flat) {
+        device const float * x = (device const float *) x03;
+        device       float * y = (device       float *) y03;
+        for (int64_t e = tpitg.x; e < gs; e += ntg.x) {
+            const float xi = x[e] - mean;
+            y[e] = xi;
+            tmp += xi*xi;
+        }
+    } else {
+        for (int64_t e = tpitg.x; e < gs; e += ntg.x) {
+            const int64_t i02 = c0 + e/npf;
+            const int64_t ir  = e - (i02 - c0)*npf;
+            const int64_t i01 = ir/ne00;
+            const int64_t i00 = ir - i01*ne00;
+
+            const float xi = *(device const float *)(x03 + (i02 - c0)*args.nb02 + i01*args.nb01 + i00*sizeof(float)) - mean;
+
+            *(device float *)(y03 + (i02 - c0)*args.nb2 + i01*args.nb1 + i00*sizeof(float)) = xi;
+
+            tmp += xi*xi;
+        }
     }
 
     tmp = simd_sum(tmp);
-    if (ntg > N_SIMDWIDTH) {
+    if (ntg.x > N_SIMDWIDTH) {
         if (sgitg == 0) {
             buf[tiisg] = 0.0f;
         }
@@ -312,7 +351,20 @@ kernel void kernel_group_norm_f32(
 
     const float variance = tmp / gs;
     const float scale = 1.0f/sqrt(variance + args.eps);
-    for (int j = start; j < end; j += ntg) {
-        dst[j] *= scale;
+
+    if (FC_group_norm_flat) {
+        device float * y = (device float *) y03;
+        for (int64_t e = tpitg.x; e < gs; e += ntg.x) {
+            y[e] *= scale;
+        }
+    } else {
+        for (int64_t e = tpitg.x; e < gs; e += ntg.x) {
+            const int64_t i02 = c0 + e/npf;
+            const int64_t ir  = e - (i02 - c0)*npf;
+            const int64_t i01 = ir/ne00;
+            const int64_t i00 = ir - i01*ne00;
+
+            *(device float *)(y03 + (i02 - c0)*args.nb2 + i01*args.nb1 + i00*sizeof(float)) *= scale;
+        }
     }
 }

ggml/src/ggml-opencl/ggml-opencl.cpp

@@ -12002,7 +12002,7 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
             return false;
         case GGML_OP_MUL_MAT_ID:
 #ifdef GGML_OPENCL_USE_ADRENO_KERNELS
-            if ((ggml_get_op_params_i32(op, 3) == GGML_PREC_BF16 || ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) &&
+            if (ggml_get_op_params_i32(op, 3) == GGML_PREC_F32 &&
                 use_adreno_moe_kernels(backend_ctx, op->src[0])) {
                 return false;
             }

ggml/src/ggml-sycl/element_wise.cpp

@@ -1183,9 +1183,7 @@ static void ggml_sycl_op_swiglu_oai(ggml_backend_sycl_context & ctx, ggml_tensor
     GGML_ASSERT(ggml_is_contiguous_1(src0));
     GGML_ASSERT(src0->nb[0] == ggml_element_size(src0));
     GGML_ASSERT(ggml_is_contiguous(dst));
-
-    GGML_ASSERT(src0->type == GGML_TYPE_F32);
-    GGML_ASSERT( dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
     GGML_ASSERT(src0->type == dst->type);
     GGML_ASSERT(dst->ne[0] == nc);
     GGML_ASSERT(ggml_nrows(dst) == ggml_nrows(src0));
@@ -1197,20 +1195,33 @@ static void ggml_sycl_op_swiglu_oai(ggml_backend_sycl_context & ctx, ggml_tensor
         GGML_ASSERT(src0->type == src1->type);
     }
 
-    //const int32_t swapped = ((const int32_t *) dst->op_params)[1];
     const int32_t swapped = ggml_get_op_params_i32(dst, 1);
-    const float alpha = ggml_get_op_params_f32(dst, 2);
-    const float limit = ggml_get_op_params_f32(dst, 3);
+    const float   alpha   = ggml_get_op_params_f32(dst, 2);
+    const float   limit   = ggml_get_op_params_f32(dst, 3);
 
-    float * src0_p = (float *) src0_d;
-    float * src1_p = (float *) src1_d;
+    if (src0->type == GGML_TYPE_F16) {
+        sycl::half * src0_p = (sycl::half *) src0_d;
+        sycl::half * src1_p = (sycl::half *) src1_d;
 
-    if (!src1) {
-        src0_p += swapped ? nc : 0;
-        src1_p += swapped ? 0 : nc;
+        if (!src1) {
+            src0_p += swapped ? nc : 0;
+            src1_p += swapped ? 0 : nc;
+        }
+
+        swiglu_oai_sycl(src0_p, src1_p, (sycl::half *) dst_d, ggml_nelements(dst), nc, src0_o / sizeof(sycl::half),
+                        src1_o / sizeof(sycl::half), alpha, limit, stream);
+    } else {
+        float * src0_p = (float *) src0_d;
+        float * src1_p = (float *) src1_d;
+
+        if (!src1) {
+            src0_p += swapped ? nc : 0;
+            src1_p += swapped ? 0 : nc;
+        }
+
+        swiglu_oai_sycl(src0_p, src1_p, (float *) dst_d, ggml_nelements(dst), nc, src0_o / sizeof(float),
+                        src1_o / sizeof(float), alpha, limit, stream);
     }
-
-    swiglu_oai_sycl(src0_p, src1_p, (float *)dst_d, ggml_nelements(dst), nc, src0_o / sizeof(float), src1_o / sizeof(float), alpha, limit, stream);
 }
 
 template <typename T>

ggml/src/ggml-vulkan/ggml-vulkan-comm.cpp

@@ -18,7 +18,7 @@ static bool ggml_vk_comm_needs_f32(const ggml_tensor * t, int depth = 0) {
     switch (t->op) {
         case GGML_OP_MUL_MAT:
         case GGML_OP_MUL_MAT_ID:
-            return ggml_get_op_params_i32(t, 0) == GGML_PREC_BF16 || ggml_get_op_params_i32(t, 0) == GGML_PREC_F32;
+            return ggml_get_op_params_i32(t, 0) == GGML_PREC_F32;
         case GGML_OP_ADD:
             return ggml_vk_comm_needs_f32(t->src[0], depth + 1) || ggml_vk_comm_needs_f32(t->src[1], depth + 1);
         case GGML_OP_MUL:

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -15696,7 +15696,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
                         return false;
                     }
                 }
-                if (ggml_prec_get_src(op, 1) == GGML_PREC_BF16 || ggml_prec_get_src(op, 1) == GGML_PREC_F32) {
+                if (ggml_prec_get_src(op, 1) == GGML_PREC_F32) {
                     return false;
                 }
                 switch (src0_type) {

ggml/src/ggml-webgpu/ggml-webgpu.cpp

@@ -4904,7 +4904,7 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
                 default:
                     break;
             }
-            if (ggml_get_op_params_i32(op, 3) == GGML_PREC_BF16 || ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) {
+            if (ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) {
                 supports_op = false;
             }
             break;

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

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()

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,

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

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);
@@ -2169,8 +2172,8 @@ ggml_tensor * llm_graph_context::build_ffn(
     if (down) {
         cur = build_lora_mm(down, cur);
         if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE || arch == LLM_ARCH_JAIS2) {
-            // GLM4, GLM4_MOE, and JAIS2 have numerical issues with F16 accumulators
-            ggml_prec_set_acc(cur, GGML_PREC_BF16);
+            // GLM4, GLM4_MOE, and JAIS2 seem to have numerical issues with half-precision accumulators
+            ggml_prec_set_acc(cur, GGML_PREC_F32);
         }
     }
 
@@ -2275,7 +2278,6 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
     if (probs_in == nullptr) {
         logits = build_lora_mm(gate_inp, cur); // [n_expert, n_tokens]
         if (gating_op == LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) {
-            // the reference computes the router in F32: https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash/blob/60d8d70770c6776ff598c94bb586a859a38244f1/inference/model.py#L565
             ggml_prec_set_acc(logits, GGML_PREC_F32);
         }
         cb(logits, "ffn_moe_logits", il);
@@ -2411,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);
 
@@ -2423,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 (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);
-        }
-
-        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);
-
-        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 (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_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);
-        }
-    }
-
-    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;
-                    }
-                }
+            if (up_exps_s) {
+                cb(gate_up, "ffn_moe_gate_up_scaled", il);
             }
 
-            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;
+            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);
+            }
 
-                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);
+            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_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, 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_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);
+                            }
+                            cb(cur, "ffn_moe_swiglu_limited", il);
+                            break;
+                        }
+                    }
                 }
-                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]
@@ -2573,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) {
@@ -2639,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 };
         }
     }
 
@@ -2670,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);
@@ -2678,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
@@ -3085,7 +3112,9 @@ ggml_tensor * llm_graph_context::build_attn_mha(
         ggml_tensor * kq = ggml_mul_mat(ctx0, k, q);
         cb(kq, "kq", il);
 
-        ggml_prec_set_acc(kq, GGML_PREC_BF16);
+        // note: this op tends to require high floating point range
+        //       while for some models F16 is enough, for others it is not, so we default to F32 here
+        ggml_prec_set_acc(kq, GGML_PREC_F32);
 
         if (model.arch == LLM_ARCH_GROK) {
             // need to do the following:
@@ -3336,9 +3365,9 @@ ggml_tensor * llm_graph_context::build_attn(
 
     if (wo) {
         if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE || arch == LLM_ARCH_JAIS2) {
-            // GLM4, GLM4_MOE, and JAIS2 have numerical issues with F16 accumulators
+            // GLM4, GLM4_MOE, and JAIS2 seem to have numerical issues with half-precision accumulators
             cur = build_lora_mm(wo, cur);
-            ggml_prec_set_acc(cur, GGML_PREC_BF16);
+            ggml_prec_set_acc(cur, GGML_PREC_F32);
             if (wo_s) {
                 cur = ggml_mul(ctx0, cur, wo_s);
             }
@@ -3425,9 +3454,9 @@ ggml_tensor * llm_graph_context::build_attn(
 
     if (wo) {
         if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE) {
-            // GLM4 and GLM4_MOE have numerical issues with F16 accumulators
+            // GLM4 and GLM4_MOE seem to have numerical issues with half-precision accumulators
             cur = build_lora_mm(wo, cur);
-            ggml_prec_set_acc(cur, GGML_PREC_BF16);
+            ggml_prec_set_acc(cur, GGML_PREC_F32);
             if (wo_s) {
                 cur = ggml_mul(ctx0, cur, wo_s);
             }

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;
 
     //

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;
     }
 
-    ~impl() {
-        log_stats();
+    // 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]);
+            }
+        }
     }
 
-    ggml_tensor * get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const {
+    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, 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);
 }

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);
 

src/models/deepseek41.cpp

@@ -875,7 +875,7 @@ llama_model_deepseek41::graph::graph(const llm_graph_params & params) :
         cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
         cb(cur, "ffn_norm", il);
 
-        // the reference computes the router in F32: https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash-1017/blob/1bae0e1088993b9b82de08f59de5be5a34e42170/inference/model.py#L813
+        // F32 router accumulation: expert selection is near-tie sensitive
         ggml_tensor * logits = build_lora_mm(layer.ffn_gate_inp, cur);
         ggml_prec_set_acc(logits, GGML_PREC_F32);
         cb(logits, "ffn_moe_logits", il);

src/models/minimax-m3.cpp

@@ -386,7 +386,6 @@ llama_model_minimax_m3::graph::graph(const llm_graph_params & params) : llm_grap
                     ggml_tensor * iq4 = ggml_reshape_4d(ctx0, iq, n_idx_dim, Hd, 1, ns);
                     ggml_tensor * sc  = ggml_mul_mat(ctx0,
                             ggml_reshape_4d(ctx0, ikp, n_idx_dim, n_ps, 1, ns), iq4);
-                    // the reference computes the indexer scores in F32: https://github.com/huggingface/transformers/blob/0f53828e83f2212060233d2ea552cc76617b8712/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py#L562
                     ggml_prec_set_acc(sc, GGML_PREC_F32);
                     // unmapped positions come out -inf, so they can never rank into the top-k
                     sc = ggml_add_inplace(ctx0, sc,
@@ -468,7 +467,7 @@ llama_model_minimax_m3::graph::graph(const llm_graph_params & params) : llm_grap
                         ggml_tensor * ikp = ggml_get_rows(ctx0, ik_s, psl_s);   // [n_idx_dim, n_ps]
                         ggml_tensor * sc = ggml_mul_mat(ctx0, ikp,
                                 ggml_reshape_2d(ctx0, iq_s, n_idx_dim, Hd*n_tps));
-                        // the reference computes the indexer scores in F32: https://github.com/huggingface/transformers/blob/0f53828e83f2212060233d2ea552cc76617b8712/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py#L562
+                        // indexer scores run in F32
                         ggml_prec_set_acc(sc, GGML_PREC_F32);
                         sc = ggml_reshape_3d(ctx0, sc, n_ps, Hd, n_tps);
                         // unmapped positions (holes, padding, empty cells) come out -inf

tests/test-backend-ops.cpp

@@ -442,11 +442,7 @@ static std::string var_to_str(ggml_type type) {
 }
 
 static std::string var_to_str(ggml_prec prec) {
-    switch (prec) {
-        case GGML_PREC_F32:  return "f32";
-        case GGML_PREC_BF16: return "bf16";
-        default:             return "def";
-    }
+    return prec == GGML_PREC_F32 ? "f32" : "def";
 }
 
 static std::string var_to_str(ggml_op_pool pool) {
@@ -12588,8 +12584,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
                                         for (int kv : { 113, 512, 1024, }) {
                                             if (nr2 != 1 && kv != 512) continue;
                                             for (int nb : { 1, 3, 32, 75, }) {
-                                                for (ggml_prec prec : {GGML_PREC_F32, GGML_PREC_BF16, GGML_PREC_DEFAULT}) {
-                                                    if (hsk != 128 && prec != GGML_PREC_F32) continue;
+                                                for (ggml_prec prec : {GGML_PREC_F32, GGML_PREC_DEFAULT}) {
+                                                    if (hsk != 128 && prec == GGML_PREC_DEFAULT) continue;
                                                     for (ggml_type type_KV : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL}) {
                                                         if (type_KV != GGML_TYPE_F16 && hsk != 64 && hsk != 72) continue;
                                                         // DeepSeek MLA: the V cache is a sub-view of the K cache

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) |

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) |

tools/mtmd/clip.cpp

@@ -831,8 +831,8 @@ ggml_tensor * clip_graph::build_attn(
         v = ggml_cont(ctx0, v);
 
         ggml_tensor * kq = ggml_mul_mat(ctx0, k, q);
-        // TODO: this most likely requires GGML_PREC_BF16, but will enable it later to avoid scope-creep
-        // ggml_prec_set_acc(kq, GGML_PREC_BF16);
+        // F32 may not needed for vision encoders?
+        // ggml_prec_set_acc(kq, GGML_PREC_F32);
 
         kq = ggml_soft_max_ext(ctx0, kq, kq_mask, kq_scale, 0.0f);
         if (sinks != nullptr) {

tools/mtmd/models/mimovl.cpp

@@ -2,7 +2,7 @@
 
 ggml_tensor * clip_graph_mimovl::build_mm(ggml_tensor * w, ggml_tensor * x) const {
     ggml_tensor * cur = ggml_mul_mat(ctx0, w, x);
-    ggml_prec_set_acc(cur, GGML_PREC_BF16);
+    ggml_prec_set_acc(cur, GGML_PREC_F32);
     return cur;
 }
 

tools/mtmd/models/qwen3tts-spkenc.cpp

@@ -27,7 +27,7 @@ ggml_tensor * clip_graph_qwen3tts_spkenc::conv1d_same(ggml_tensor * x, ggml_tens
 
     ggml_tensor * w2d = ggml_reshape_2d(ctx0, w, (int64_t) K * IC, OC);
     ggml_tensor * y   = ggml_mul_mat(ctx0, w2d, col); // [OC, T_out]
-    ggml_prec_set_acc(y, GGML_PREC_BF16);
+    ggml_prec_set_acc(y, GGML_PREC_F32);
 
     ggml_tensor * b2d = ggml_reshape_2d(ctx0, b, OC, 1);
     y = ggml_add(ctx0, y, b2d);

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) |

tools/server/server-context.cpp

@@ -3424,13 +3424,19 @@ private:
     }
 
     void iterate(std::vector<server_slot> & slots, std::function<void(server_slot &)> callback) {
-        for (auto & slot : slots) {
-            try {
-                callback(slot);
-            } catch (const std::exception & e) {
-                SLT_ERR(slot, "got exception: %s\n", e.what());
-                send_error(slot, std::string("got exception: ") + e.what(), ERROR_TYPE_SERVER);
-                slot.release();
+        for (int pass = 0; pass < 2; ++pass) {
+            for (auto & slot : slots) {
+                const bool is_child = slot.is_processing() && slot.task->is_child();
+                if (is_child != (pass == 1)) {
+                    continue;
+                }
+                try {
+                    callback(slot);
+                } catch (const std::exception & e) {
+                    SLT_ERR(slot, "got exception: %s\n", e.what());
+                    send_error(slot, std::string("got exception: ") + e.what(), ERROR_TYPE_SERVER);
+                    slot.release();
+                }
             }
         }
     }
@@ -4007,6 +4013,7 @@ private:
                 // check if this is a child slot
                 if (slot.state == SLOT_STATE_WAIT_OTHER) {
                     SLT_DBG(slot, "%s", "waiting for parent slot to complete\n");
+                    SLT_DBG(slot, "%s", "__TEST_TAG_WAIT_PARENT_SLOT__\n");
                     return;
                 }
 

tools/server/server-task.cpp

@@ -1836,7 +1836,7 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro
 }
 
 bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot) {
-    const int lcp_best = prompt.tokens.get_common_prefix(tokens_new);
+    int lcp_best = prompt.tokens.get_common_prefix(tokens_new);
 
     float f_keep_best = prompt.tokens.size() > 0 ? float(lcp_best) / prompt.tokens.size() : -1.0f; // empty slot: any cache entry wins
     float f_sim_best  = float(lcp_best) / tokens_new.size();
@@ -1862,6 +1862,7 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
         if (f_keep_best < f_keep_cur && f_sim_best < f_sim_cur) {
             f_keep_best = f_keep_cur;
             f_sim_best  = f_sim_cur;
+            lcp_best    = lcp_cur;
 
             it_best = it;
         }
@@ -1904,6 +1905,13 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
             }
         }
 
+        SRV_DBG("%s", "__TEST_TAG_PROMPT_CACHE_RESTORE__\n");
+
+        // the whole cached prompt is a prefix of the new one
+        if (lcp_best == (int) it_best->prompt.tokens.size()) {
+            SRV_DBG("%s", "__TEST_TAG_PROMPT_CACHE_FULL_RESTORE__\n");
+        }
+
         prompt = std::move(it_best->prompt);
 
         states.erase(it_best);

tools/server/tests/unit/test_systemone.py

@@ -122,10 +122,13 @@ def test_systemone_invalid_request(data: dict):
     assert "error" in res.body
 
 
-def test_systemone_shared_prompt():
+def test_systemone_shared_prompt(tmp_path):
     global server
     server = ServerPreset.tinyopenjev()
+    server.n_slots = len(TEST_QUESTIONS)
     server.server_metrics = True
+    server.debug = True
+    server.log_path = str(tmp_path / "server.log")
     server.start()
     res = server.make_request("POST", "/v1/systemone", data={
         "state": TEST_STATE,
@@ -139,8 +142,19 @@ def test_systemone_shared_prompt():
     assert n_cached % 2 == 0
     assert n_processed + n_cached == res.body["usage"]["input_tokens"]
 
-    # with one slot the prompt cannot be shared, the answers must be the same
+    # the next request takes an empty slot as parent, so child slot 0 has a lower id than the parent
+    res_next = server.make_request("POST", "/v1/systemone", data={
+        "state": "Unrelated follow-up case: " * 10,
+        "questions": TEST_QUESTIONS,
+    })
+    assert res_next.status_code == 200
+
+    # in each request, each child waits for the parent during 1 batch only
     server.stop()
+    with open(server.log_path) as f:
+        assert f.read().count("__TEST_TAG_WAIT_PARENT_SLOT__") == 2 * (len(TEST_QUESTIONS) - 1)
+
+    # with one slot the prompt cannot be shared, the answers must be the same
     server = ServerPreset.tinyopenjev()
     server.n_slots = 1
     server.start()
@@ -181,7 +195,7 @@ def test_systemone_prompt_cache(cache_ram, tmp_path):
         log = f.read()
     # the prompt cache is opt-in for decision models, and restores their states once enabled
     assert ("__TEST_TAG_CACHE_IDLE_SLOTS_ENABLED__" in log) == (cache_ram is not None)
-    assert ("found better prompt" in log) == (cache_ram is not None)
+    assert ("__TEST_TAG_PROMPT_CACHE_RESTORE__" in log) == (cache_ram is not None)
 
 
 def test_systemone_prompt_cache_child_slots(tmp_path):
@@ -210,7 +224,7 @@ def test_systemone_prompt_cache_child_slots(tmp_path):
     decide(TEST_STATE, TEST_QUESTIONS)
 
     with open(server.log_path) as f:
-        n_restored = f.read().count("found better prompt with f_keep = 1.000")
+        n_restored = f.read().count("__TEST_TAG_PROMPT_CACHE_FULL_RESTORE__")
 
     # every overwritten state comes back whole from the prompt cache
     for state in states:
@@ -219,7 +233,7 @@ def test_systemone_prompt_cache_child_slots(tmp_path):
     server.stop()
 
     with open(server.log_path) as f:
-        assert f.read().count("found better prompt with f_keep = 1.000") == n_restored + len(states)
+        assert f.read().count("__TEST_TAG_PROMPT_CACHE_FULL_RESTORE__") == n_restored + len(states)
 
 
 def test_systemone_images():