Moe expert residency

common/arg.cpp

@@ -2301,6 +2301,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.devices = parse_device_list(value);
         }
     ).set_env("LLAMA_ARG_DEVICE"));
+    add_opt(common_arg(
+        {"--moe-n-slots"}, "N",
+        "number of expert slots in GPU memory for MoE disk offloading (0 = disabled)",
+        [](common_params & params, const std::string & value) {
+            params.moe.n_slots = std::stoi(value);
+        }
+    ).set_env("LLAMA_ARG_MOE_N_SLOTS"));
+    add_opt(common_arg(
+        {"--moe-n-layers"}, "N",
+        "number of MoE layers to disk-offload (0 = all layers)",
+        [](common_params & params, const std::string & value) {
+            params.moe.n_layers = std::stoi(value);
+        }
+    ).set_env("LLAMA_ARG_MOE_N_LAYERS"));
     add_opt(common_arg(
         {"--list-devices"},
         "print list of available devices and exit",

common/common.cpp

@@ -1553,6 +1553,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
     mparams.progress_callback           = params.load_progress_callback;
     mparams.progress_callback_user_data = params.load_progress_callback_user_data;
     mparams.no_alloc                    = params.no_alloc;
+    mparams.moe                         = params.moe;
 
     return mparams;
 }
@@ -1568,6 +1569,7 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.n_threads         = params.cpuparams.n_threads;
     cparams.n_threads_batch   = params.cpuparams_batch.n_threads == -1 ?
                                 params.cpuparams.n_threads : params.cpuparams_batch.n_threads;
+    cparams.moe               = params.moe;
     cparams.embeddings        = params.embedding;
     cparams.rope_scaling_type = params.rope_scaling_type;
     cparams.rope_freq_base    = params.rope_freq_base;

common/common.h

@@ -444,6 +444,7 @@ struct common_params {
 
     // offload params
     std::vector<ggml_backend_dev_t> devices; // devices to use for offloading
+    llama_moe_params                moe = { 0, 0 };
 
     int32_t n_gpu_layers       = -1;    // number of layers to store in VRAM, -1 is auto, <= -2 is all
     int32_t main_gpu           = 0;     // the GPU that is used for scratch and small tensors

ggml/include/ggml-metal.h

@@ -56,6 +56,41 @@ GGML_BACKEND_API void ggml_backend_metal_capture_next_compute(ggml_backend_t bac
 
 GGML_BACKEND_API ggml_backend_reg_t ggml_backend_metal_reg(void);
 
+typedef struct ggml_backend_metal_event * ggml_backend_metal_event_t;
+
+GGML_BACKEND_API ggml_backend_metal_event_t ggml_backend_metal_event_new(ggml_backend_t backend);
+GGML_BACKEND_API void                       ggml_backend_metal_event_free(ggml_backend_metal_event_t event);
+GGML_BACKEND_API void   ggml_backend_metal_event_signal(ggml_backend_metal_event_t event, uint64_t value);
+GGML_BACKEND_API void * ggml_backend_metal_event_raw(ggml_backend_metal_event_t event);
+
+// Shared memory layout (MTLStorageModeShared).
+const int    MOE_MAX_IDS      = 4096;
+const size_t MOE_OFF_REQ      = 0;                                   // atomic_uint: request seq
+const size_t MOE_OFF_N        = 8;                                   // int32:       id count
+const size_t MOE_OFF_SELECTED = 16;                                  // int32[n]:    expert ids (GPU writes)
+const size_t MOE_OFF_REMAPPED = MOE_OFF_SELECTED + MOE_MAX_IDS * 4;  // int32[n]:    slot ids (CPU writes)
+const size_t MOE_MSG_NBYTES   = MOE_OFF_REMAPPED + MOE_MAX_IDS * 4;
+
+struct ggml_metal_moe_intercept {
+    int                        n;
+    uint32_t                   seq;
+    const struct ggml_tensor * msg_tensor;
+    ggml_backend_metal_event_t event;
+    bool                       reuse;
+};
+
+typedef bool (*ggml_metal_moe_query_fn)(void *                            user_data,
+                                        const struct ggml_tensor *        src0,
+                                        const struct ggml_tensor *        ids,
+                                        struct ggml_metal_moe_intercept * out);
+
+struct ggml_metal_moe_handler {
+    ggml_metal_moe_query_fn fn;
+    void *                  user_data;
+};
+
+GGML_BACKEND_API void ggml_backend_metal_set_moe_handler(ggml_backend_t backend, struct ggml_metal_moe_handler handler);
+
 #ifdef __cplusplus
 }
 #endif

ggml/src/ggml-metal/ggml-metal-context.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "ggml-metal-device.h"
+#include "ggml-metal.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -32,6 +33,7 @@ void ggml_metal_event_wait  (ggml_metal_t ctx, ggml_metal_event_t ev);
 ggml_metal_event_t ggml_metal_get_ev_cpy(ggml_metal_t ctx);
 
 void ggml_metal_set_n_cb            (ggml_metal_t ctx, int n_cb);
+void ggml_metal_set_moe_handler     (ggml_metal_t ctx, struct ggml_metal_moe_handler moe_handler);
 void ggml_metal_set_abort_callback  (ggml_metal_t ctx, ggml_abort_callback abort_callback, void * user_data);
 bool ggml_metal_supports_family     (ggml_metal_t ctx, int family);
 void ggml_metal_capture_next_compute(ggml_metal_t ctx);

ggml/src/ggml-metal/ggml-metal-context.m

@@ -79,6 +79,8 @@
     // error state - set when a command buffer fails during synchronize
     // once set, graph_compute will return GGML_STATUS_FAILED until the backend is recreated
     bool has_error;
+
+    struct ggml_metal_moe_handler moe_handler;
 };
 
 ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) {
@@ -660,6 +662,11 @@ ggml_metal_event_t ggml_metal_get_ev_cpy(ggml_metal_t ctx) {
     return ctx->ev_cpy;
 }
 
+void ggml_metal_set_moe_handler(ggml_metal_t ctx, struct ggml_metal_moe_handler moe_handler) {
+    GGML_ASSERT(moe_handler.fn == NULL || ctx->moe_handler.fn == NULL);
+    ctx->moe_handler = moe_handler;
+}
+
 void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) {
     if (ctx->n_cb != n_cb) {
         ctx->n_cb = MIN(n_cb, GGML_METAL_MAX_COMMAND_BUFFERS);
@@ -702,7 +709,8 @@ void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) {
             ctx->use_concurrency,
             ctx->capture_compute,
             ctx->debug_graph,
-            ctx->debug_fusion);
+            ctx->debug_fusion,
+            ctx->moe_handler);
 
         for (int idx = 0; idx < ggml_metal_op_n_nodes(ctx_op); ++idx) {
             const int res = ggml_metal_op_encode(ctx_op, idx);

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

@@ -176,6 +176,24 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows(ggml_me
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_moe_interceptor(ggml_metal_library_t lib) {
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_moe_interceptor");
+    snprintf(name, 256, "%s", 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);
+    }
+
+    res.nsg  = 1;
+    res.smem = 0;
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_diag(ggml_metal_library_t lib, const ggml_tensor * op) {
     char base[256];
     char name[256];

ggml/src/ggml-metal/ggml-metal-device.h

@@ -91,6 +91,8 @@ void ggml_metal_encoder_memory_barrier(ggml_metal_encoder_t encoder);
 
 void ggml_metal_encoder_end_encoding(ggml_metal_encoder_t encoder);
 
+void ggml_metal_encoder_wait_for_event(ggml_metal_encoder_t enc, void * event, uint64_t value);
+
 //
 // MTLLibrary wrapper
 //
@@ -161,6 +163,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_opt_step_
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_opt_step_sgd      (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_memset            (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_count_equal       (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_moe_interceptor   (ggml_metal_library_t lib);
 
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_pad(
         ggml_metal_library_t lib,
@@ -267,6 +270,8 @@ typedef struct ggml_metal_event * ggml_metal_event_t;
 
 void ggml_metal_event_encode_signal(ggml_metal_event_t ev, ggml_metal_cmd_buf_t cmd_buf);
 void ggml_metal_event_encode_wait  (ggml_metal_event_t ev, ggml_metal_cmd_buf_t cmd_buf);
+void ggml_metal_event_cpu_signal   (ggml_metal_event_t ev, uint64_t value);
+void * ggml_metal_event_get_obj    (ggml_metal_event_t ev);
 
 ggml_metal_device_t ggml_metal_device_init(int device);
 void ggml_metal_device_free(ggml_metal_device_t dev);

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

@@ -458,13 +458,19 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_compile_pipeline(ggml_
 
 struct ggml_metal_encoder {
     id<MTLComputeCommandEncoder> obj;
+    id<MTLCommandBuffer> cmd_buf;
+    bool concurrent;
 };
 
 ggml_metal_encoder_t ggml_metal_encoder_init(ggml_metal_cmd_buf_t cmd_buf_raw, bool concurrent) {
     ggml_metal_encoder_t res = calloc(1, sizeof(struct ggml_metal_encoder));
 
     id<MTLCommandBuffer> cmd_buf = (id<MTLCommandBuffer>) cmd_buf_raw;
 
+    res->cmd_buf = cmd_buf;
+    res->concurrent = concurrent;
+    [res->cmd_buf retain];
+
     if (concurrent) {
         res->obj = [cmd_buf computeCommandEncoderWithDispatchType: MTLDispatchTypeConcurrent];
     } else {
@@ -476,8 +482,23 @@ ggml_metal_encoder_t ggml_metal_encoder_init(ggml_metal_cmd_buf_t cmd_buf_raw, b
     return res;
 }
 
+void ggml_metal_encoder_wait_for_event(ggml_metal_encoder_t enc, void *event, uint64_t value) {
+    id<MTLSharedEvent> ev = (__bridge id<MTLSharedEvent>)event;
+    [enc->obj endEncoding];
+    [enc->obj release];
+    [enc->cmd_buf encodeWaitForEvent:ev value:value];
+    if (enc->concurrent) {
+        enc->obj = [enc->cmd_buf
+            computeCommandEncoderWithDispatchType:MTLDispatchTypeConcurrent];
+    } else {
+        enc->obj = [enc->cmd_buf computeCommandEncoder];
+    }
+    [enc->obj retain];
+}
+
 void ggml_metal_encoder_free(ggml_metal_encoder_t encoder) {
     [encoder->obj release];
+    [encoder->cmd_buf release];
     free(encoder);
 }
 
@@ -1004,6 +1025,14 @@ void ggml_metal_event_encode_wait(ggml_metal_event_t ev, ggml_metal_cmd_buf_t cm
     [cmd_buf encodeWaitForEvent:event value:atomic_load_explicit(&ev->value, memory_order_relaxed)];
 }
 
+void ggml_metal_event_cpu_signal(ggml_metal_event_t ev, uint64_t value) {
+    ((id<MTLSharedEvent>)ev->obj).signaledValue = value;
+}
+
+void *ggml_metal_event_get_obj(ggml_metal_event_t ev) {
+    return ev->obj;
+}
+
 ggml_metal_event_t ggml_metal_device_event_init(ggml_metal_device_t dev) {
     id<MTLSharedEvent> event = [dev->mtl_device newSharedEvent];
 

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

@@ -1,5 +1,6 @@
 #include "ggml-metal-ops.h"
 
+#include "ggml-metal.h"
 #include "ggml.h"
 #include "ggml-impl.h"
 #include "ggml-backend-impl.h"
@@ -36,7 +37,8 @@ struct ggml_metal_op {
         bool use_concurrency,
         bool use_capture,
         int  debug_graph,
-        int  debug_fusion) {
+        int  debug_fusion,
+        ggml_metal_moe_handler moe_handler) {
         this->dev             = dev;
         this->lib             = ggml_metal_device_get_library(dev);
         this->enc             = ggml_metal_encoder_init(cmd_buf, use_concurrency);
@@ -48,6 +50,7 @@ struct ggml_metal_op {
         this->use_capture     = use_capture;
         this->debug_graph     = debug_graph;
         this->debug_fusion    = debug_fusion;
+        this->moe_handler     = moe_handler;
         this->gf              = gf;
 
         idxs.reserve(gf->n_nodes);
@@ -100,6 +103,8 @@ struct ggml_metal_op {
     int debug_graph;
     int debug_fusion;
 
+    ggml_metal_moe_handler moe_handler;
+
 private:
     ggml_cgraph * gf;
 
@@ -120,7 +125,8 @@ ggml_metal_op_t ggml_metal_op_init(
         bool use_concurrency,
         bool use_capture,
         int debug_graph,
-        int debug_fusion) {
+        int debug_fusion,
+        ggml_metal_moe_handler moe_handler) {
     ggml_metal_op_t res = new ggml_metal_op(
         dev,
         cmd_buf,
@@ -131,7 +137,8 @@ ggml_metal_op_t ggml_metal_op_init(
         use_concurrency,
         use_capture,
         debug_graph,
-        debug_fusion);
+        debug_fusion,
+        moe_handler);
 
     return res;
 }
@@ -2310,6 +2317,35 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
     ggml_metal_buffer_id bid_src2 = ggml_metal_get_buffer_id(op->src[2]);
     ggml_metal_buffer_id bid_dst  = ggml_metal_get_buffer_id(op);
 
+    // moe interceptor
+    if (ctx->moe_handler.fn) {
+        ggml_metal_moe_intercept mi;
+        if (ctx->moe_handler.fn(ctx->moe_handler.user_data, op->src[0], op->src[2], &mi)) {
+            ggml_metal_buffer_id moe_base = ggml_metal_get_buffer_id(mi.msg_tensor);
+
+            if (!mi.reuse) {
+                auto pipeline = ggml_metal_library_get_pipeline_moe_interceptor(lib);
+
+                ggml_metal_buffer_id b_req  = { moe_base.metal, moe_base.offs + MOE_OFF_REQ };
+                ggml_metal_buffer_id b_msel = { moe_base.metal, moe_base.offs + MOE_OFF_SELECTED };
+                uint32_t             n_u    = (uint32_t) mi.n;
+
+                ggml_metal_encoder_set_pipeline(enc, pipeline);
+                ggml_metal_encoder_set_buffer(enc, bid_src2, 0);
+                ggml_metal_encoder_set_buffer(enc, b_req, 1);
+                ggml_metal_encoder_set_buffer(enc, b_msel, 2);
+                ggml_metal_encoder_set_bytes(enc, &n_u, sizeof(n_u), 3);
+                ggml_metal_encoder_set_bytes(enc, &mi.seq, sizeof(mi.seq), 4);
+                ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, 32, 1, 1);
+
+                ggml_metal_encoder_wait_for_event(enc, ggml_backend_metal_event_raw(mi.event), mi.seq);
+            }
+
+            bid_src2.metal = moe_base.metal;
+            bid_src2.offs  = moe_base.offs + MOE_OFF_REMAPPED;
+        }
+    }
+
     const uint32_t r2 = 1;
     const uint32_t r3 = 1;
 

ggml/src/ggml-metal/ggml-metal-ops.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "ggml-metal-device.h"
+#include "ggml-metal.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -18,7 +19,8 @@ ggml_metal_op_t ggml_metal_op_init(
         bool use_concurrency,
         bool use_capture,
         int  debug_graph,
-        int  debug_fusion);
+        int  debug_fusion,
+        struct ggml_metal_moe_handler);
 
 void ggml_metal_op_free(ggml_metal_op_t ctx);
 

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

@@ -565,6 +565,42 @@ static void ggml_backend_metal_set_n_cb(ggml_backend_t backend, int n_cb) {
     ggml_metal_set_n_cb(ctx, n_cb);
 }
 
+void ggml_backend_metal_set_moe_handler(ggml_backend_t backend, struct ggml_metal_moe_handler moe_handler) {
+    GGML_ASSERT(ggml_backend_is_metal(backend));
+
+    ggml_metal_t ctx = (ggml_metal_t) backend->context;
+
+    ggml_metal_set_moe_handler(ctx, moe_handler);
+}
+
+struct ggml_backend_metal_event {
+    ggml_metal_device_t dev;
+    ggml_metal_event_t  ev;
+};
+
+ggml_backend_metal_event_t ggml_backend_metal_event_new(ggml_backend_t backend) {
+    GGML_ASSERT(ggml_backend_is_metal(backend));
+    ggml_metal_device_t dev = (ggml_metal_device_t) backend->device->context;
+
+    auto * e = new ggml_backend_metal_event;
+    e->dev   = dev;
+    e->ev    = ggml_metal_device_event_init(dev);
+    return e;
+}
+
+void ggml_backend_metal_event_free(ggml_backend_metal_event_t event) {
+    ggml_metal_device_event_free(event->dev, event->ev);
+    delete event;
+}
+
+void ggml_backend_metal_event_signal(ggml_backend_metal_event_t event, uint64_t value) {
+    ggml_metal_event_cpu_signal(event->ev, value);
+}
+
+void * ggml_backend_metal_event_raw(ggml_backend_metal_event_t event) {
+    return ggml_metal_event_get_obj(event->ev);
+}
+
 static ggml_backend_i ggml_backend_metal_i = {
     /* .get_name                = */ ggml_backend_metal_name,
     /* .free                    = */ ggml_backend_metal_free,