danielhanchen · Oct 12, 2026 · Oct 12, 2026 · Oct 12, 2026 · Oct 12, 2026 · Oct 12, 2026
diff --git a/common/common.cpp b/common/common.cpp
diff --git a/common/fit.cpp b/common/fit.cpp
diff --git a/docs/ops.md b/docs/ops.md
@@ -71,6 +71,7 @@ Legend:
 |                LIGHTNING_INDEXER | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ | ❌ |
+|                 MOE_CACHE_LOOKUP | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
 |                 MUL_MAT_HADAMARD | ❌ | ❌ | ✅ | ❌ | ✅ | 🟡 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |

diff --git a/docs/ops/CPU.csv b/docs/ops/CPU.csv
@@ -11445,6 +11445,7 @@
 "CPU","MUL_MAT_ID","type_a=iq4_xs,type_b=f32,n_mats=4,n_used=2,b=0,m=512,n=32,k=256,amax=1.000000,m_v=0","support","1","yes","CPU"
 "CPU","MUL_MAT_ID","type_a=bf16,type_b=f32,n_mats=4,n_used=2,b=0,m=512,n=1,k=256,amax=1.000000,m_v=0","support","1","yes","CPU"
 "CPU","MUL_MAT_ID","type_a=bf16,type_b=f32,n_mats=4,n_used=2,b=0,m=512,n=32,k=256,amax=1.000000,m_v=0","support","1","yes","CPU"
+"CPU","MOE_CACHE_LOOKUP","n_expert=8,n_used=2,n_tokens=4,n_slots=8,n_banks=3","support","0","no","CPU"
 "CPU","OUT_PROD","type_a=f32,type_b=f32,m=256,n=1,k=1,bs=[1,1],nr=[1,1],trans_b=0","support","1","yes","CPU"
 "CPU","OUT_PROD","type_a=f32,type_b=f32,m=256,n=1,k=1,bs=[1,1],nr=[1,2],trans_b=0","support","1","yes","CPU"
 "CPU","OUT_PROD","type_a=f32,type_b=f32,m=256,n=1,k=1,bs=[1,1],nr=[2,1],trans_b=0","support","1","yes","CPU"

diff --git a/docs/ops/CUDA.csv b/docs/ops/CUDA.csv
@@ -7613,6 +7613,7 @@
 "CUDA0","MUL_MAT_ID","type_a=iq4_xs,type_b=f32,n_mats=4,n_used=2,b=0,m=512,n=32,k=256","support","1","yes","CUDA"
 "CUDA0","MUL_MAT_ID","type_a=bf16,type_b=f32,n_mats=4,n_used=2,b=0,m=512,n=1,k=256","support","1","yes","CUDA"
 "CUDA0","MUL_MAT_ID","type_a=bf16,type_b=f32,n_mats=4,n_used=2,b=0,m=512,n=32,k=256","support","1","yes","CUDA"
+"CUDA0","MOE_CACHE_LOOKUP","n_expert=8,n_used=2,n_tokens=4,n_slots=8,n_banks=3","support","1","yes","CUDA"
 "CUDA0","OUT_PROD","type_a=f32,type_b=f32,m=256,n=1,k=1,bs=[1,1],nr=[1,1],trans_b=0","support","1","yes","CUDA"
 "CUDA0","OUT_PROD","type_a=f32,type_b=f32,m=256,n=1,k=1,bs=[1,1],nr=[1,2],trans_b=0","support","1","yes","CUDA"
 "CUDA0","OUT_PROD","type_a=f32,type_b=f32,m=256,n=1,k=1,bs=[1,1],nr=[2,1],trans_b=0","support","1","yes","CUDA"

diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h
@@ -117,6 +117,20 @@ extern "C" {
     // automatic fallback to sync copy if async is not supported
     GGML_API void ggml_backend_tensor_copy_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst);
 
+    // a copy of `size` bytes from `src` at byte offset `src_offset` to `dst` at byte offset `dst_offset`
+    struct ggml_backend_copy_region {
+        const struct ggml_tensor * src;
+        size_t                     src_offset;
+        struct ggml_tensor *       dst;
+        size_t                     dst_offset;
+        size_t                     size;
+    };
+
+    // queue the copies on `backend` as one batch in stream order, the copies of a batch may run in any order and must not overlap
+    // dst: device memory of the backend; src: device memory of the same device, or host memory it can read async (e.g. pinned)
+    // returns false (nothing copied) if the backend cannot batch them; with n_regions == 0, returns whether batched copies are supported
+    GGML_API bool ggml_backend_tensor_copy_regions_async(ggml_backend_t backend, const struct ggml_backend_copy_region * regions, size_t n_regions);
+
     GGML_API ggml_backend_dev_t ggml_backend_get_device(ggml_backend_t backend);
 
     //
@@ -334,6 +348,18 @@ extern "C" {
     // `graph_size_changed` is true if the graph size changed since the last allocation
     typedef void (*ggml_backend_sched_realloc_callback)(ggml_backend_sched_t sched, struct ggml_cgraph * graph, bool backend_ids_changed, bool graph_size_changed, void * user_data);
 
+    // How an input weight in a host buffer is copied to the backend of the split that uses it
+    enum ggml_backend_sched_prefetch_mode {
+        GGML_BACKEND_SCHED_PREFETCH_NONE  = 0, // copied when the split starts, the copy is allocated by the graph allocator (no prefetch)
+        GGML_BACKEND_SCHED_PREFETCH_SYNC  = 1, // copied when the split starts (the copy needs results of earlier splits), no weight of the split is prefetched, the copy uses a staging slot if the backend has an arena
+        GGML_BACKEND_SCHED_PREFETCH_ASYNC = 2, // copied into a staging slot on a separate copy queue, ahead of the split
+    };
+
+    // Callback selecting the prefetch mode of the host weight `src` (`dst` is its copy, `graph` the split graph), return NONE for host tensors that change during the graph computation
+    // ASYNC copies go through the copy callback with the copy queue backend (whole if it returns false) and may run before earlier splits, so they must not depend on graph results
+    // with a copy callback but no prefetch callback nothing is prefetched; env: GGML_SCHED_PREFETCH (0 = off), GGML_SCHED_PREFETCH_DEPTH
+    typedef enum ggml_backend_sched_prefetch_mode (*ggml_backend_sched_prefetch_callback)(const struct ggml_tensor * src, const struct ggml_tensor * dst, struct ggml_cgraph * graph, void * user_data);
+
     // Initialize a backend scheduler, backends with low index are given priority over backends with high index
     GGML_API ggml_backend_sched_t ggml_backend_sched_new(ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, int n_backends, size_t graph_size, bool parallel, bool op_offload);
     GGML_API void                 ggml_backend_sched_free(ggml_backend_sched_t sched);
@@ -351,6 +377,9 @@ extern "C" {
 
     GGML_API ggml_backend_buffer_type_t ggml_backend_sched_get_buffer_type(ggml_backend_sched_t sched, ggml_backend_t backend);
     GGML_API size_t                     ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend);
+    // a reserve that leaves less than `headroom` bytes free on a device with a staging arena drops the host weight prefetch (0 = never),
+    // default GGML_SCHED_PREFETCH_HEADROOM_MIB (512): for users whose memory measurement does not count the staging
+    GGML_API void                       ggml_backend_sched_set_prefetch_headroom(ggml_backend_sched_t sched, size_t headroom);
 
     GGML_API void                 ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend);
     GGML_API ggml_backend_t       ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node);
@@ -378,6 +407,9 @@ extern "C" {
     // Set a callback to be called when the compute buffers of a graph are being reallocated
     GGML_API void                 ggml_backend_sched_set_realloc_callback(ggml_backend_sched_t sched, ggml_backend_sched_realloc_callback callback, void * user_data);
 
+    // Set a callback that selects how the input weights in host buffers are prefetched
+    GGML_API void                 ggml_backend_sched_set_prefetch_callback(ggml_backend_sched_t sched, ggml_backend_sched_prefetch_callback callback, void * user_data);
+
     //
     // Meta backend
     //

diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h
@@ -8,10 +8,10 @@ extern "C" {
 
 #define RPC_PROTO_MAJOR_VERSION    8
 #define RPC_PROTO_MINOR_VERSION    0
-#define RPC_PROTO_PATCH_VERSION    0
+#define RPC_PROTO_PATCH_VERSION    1
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16

diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h
@@ -601,6 +601,8 @@ extern "C" {
 
         GGML_OP_GLU,
 
+        GGML_OP_MOE_CACHE_LOOKUP,
+
         GGML_OP_COUNT,
     };
 
@@ -1503,6 +1505,24 @@ extern "C" {
             struct ggml_tensor  * b,
             struct ggml_tensor  * ids);
 
+    // device-side lookup in a cache of MoE experts (see llama-moe-cache.cpp)
+    // ids:   I32 [n_expert_used, n_tokens], the experts selected in layer il
+    // map:   I32 [n_expert, n_layer], the slot of each expert of each layer, -1 if not cached (updated in place)
+    // state: I32 [8 + 3*n_slots], eviction state of the slots of the cache (updated in place)
+    // banks: the cached experts, [ne0, ne1, n_slots + 1] each (updated in place)
+    // host:  device-accessible pointers to the experts of layer il in host memory, with the layout of the banks
+    // the missing experts get a slot (the least recently used one, approximated with a CLOCK) and are copied from host memory
+    // result: I32 [n_expert_used, n_tokens] slots of the selected experts, then the copied experts; ggml_view_2d(ctx, res, n_expert_used, n_tokens, res->nb[1], 0) gives the ids for ggml_mul_mat_id with a bank
+    GGML_API struct ggml_tensor * ggml_moe_cache_lookup(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * ids,
+            struct ggml_tensor  * map,
+            struct ggml_tensor  * state,
+            struct ggml_tensor ** banks,
+            const void         ** host,
+            int                   n_banks,
+            int                   il);
+
     // A: m columns, n rows,
     // B: p columns, n rows,
     // result is m columns, p rows

diff --git a/ggml/src/ggml-alloc.c b/ggml/src/ggml-alloc.c
@@ -939,6 +939,9 @@ static bool ggml_gallocr_reserve_n_impl(
                 galloc->buffers[i] = ggml_vbuffer_alloc(galloc->bufts[i], galloc->buf_tallocs[i], GGML_BACKEND_BUFFER_USAGE_COMPUTE);
                 if (galloc->buffers[i] == NULL) {
                     GGML_LOG_ERROR("%s: failed to allocate %s buffer of size %zu\n", __func__, ggml_backend_buft_name(galloc->bufts[i]), new_size);
+                    // the assignments point into the freed buffer: the next ggml_gallocr_alloc_graph must reserve again instead of using them
+                    galloc->n_nodes = 0;
+                    galloc->n_leafs = 0;
                     return false;
                 }
             }

diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp
@@ -153,6 +153,10 @@ static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type(
 
 static bool ggml_backend_meta_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) {
     GGML_ASSERT(ggml_backend_dev_is_meta(dev));
+    // the MoE cache lookup reads and writes per-device slot state through raw pointers, it has no split state
+    if (op->op == GGML_OP_MOE_CACHE_LOOKUP) {
+        return false;
+    }
     const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context;
     return std::all_of(meta_dev_ctx->simple_devs.begin(), meta_dev_ctx->simple_devs.end(),
         [op](ggml_backend_dev_t simple_dev) { return ggml_backend_dev_supports_op(simple_dev, op); });