Add native GLM-5.2 sparse execution

conversion/glm.py

@@ -237,6 +237,12 @@ def set_gguf_parameters(self):
         self.gguf_writer.add_indexer_head_count(self.hparams["index_n_heads"])
         self.gguf_writer.add_indexer_key_length(self.hparams["index_head_dim"])
         self.gguf_writer.add_indexer_top_k(self.hparams["index_topk"])
+        if (indexer_types := self.hparams.get("indexer_types")) is not None:
+            if len(indexer_types) != self.hparams["num_hidden_layers"]:
+                raise ValueError("indexer_types must contain one entry per target layer")
+            if invalid := set(indexer_types) - {"full", "shared"}:
+                raise ValueError(f"unsupported indexer_types values: {sorted(invalid)}")
+            self.gguf_writer.add_indexer_types([indexer_type == "full" for indexer_type in indexer_types])
 
 
 @ModelBase.register("SolarOpenForCausalLM")

ggml/src/ggml-cpu/ggml-cpu.c

@@ -1161,6 +1161,75 @@ void ggml_set_f32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2,
 
 // ggml_compute_forward_mul_mat
 
+static bool ggml_compute_forward_mul_mat_precise_ref_type(enum ggml_type type) {
+    return ggml_is_quantized(type) || type == GGML_TYPE_F16 || type == GGML_TYPE_BF16;
+}
+
+static float ggml_compute_forward_mul_mat_precise_ref_dot(const struct ggml_type_traits * traits,
+                                                          const void *                    weight,
+                                                          const float *                   activation,
+                                                          float *                         row,
+                                                          int64_t                         n_input) {
+    traits->to_float(weight, row, n_input);
+
+    double sum = 0.0;
+    for (int64_t input = 0; input < n_input; ++input) {
+        sum += (double) row[input] * (double) activation[input];
+    }
+    return (float) sum;
+}
+
+static bool ggml_compute_forward_mul_mat_precise_ref(const struct ggml_compute_params * params,
+                                                     struct ggml_tensor *               dst) {
+    const struct ggml_tensor * src0 = dst->src[0];
+    const struct ggml_tensor * src1 = dst->src[1];
+
+    if (!params->use_ref || src1->type != GGML_TYPE_F32 || !ggml_compute_forward_mul_mat_precise_ref_type(src0->type)) {
+        return false;
+    }
+
+    GGML_ASSERT(src0->ne[0] == src1->ne[0]);
+    GGML_ASSERT(dst->ne[0] == src0->ne[1]);
+    GGML_ASSERT(dst->ne[1] == src1->ne[1]);
+    GGML_ASSERT(dst->ne[2] == src1->ne[2]);
+    GGML_ASSERT(dst->ne[3] == src1->ne[3]);
+    GGML_ASSERT(src1->ne[2] % src0->ne[2] == 0);
+    GGML_ASSERT(src1->ne[3] % src0->ne[3] == 0);
+
+    const struct ggml_type_traits * traits = ggml_get_type_traits(src0->type);
+    GGML_ASSERT(traits->to_float != NULL);
+
+    const int64_t n_input   = src0->ne[0];
+    const int64_t n_output  = dst->ne[0];
+    const int64_t n_columns = dst->ne[1];
+    const int64_t n_tasks   = ggml_nelements(dst);
+    const int64_t r2        = src1->ne[2] / src0->ne[2];
+    const int64_t r3        = src1->ne[3] / src0->ne[3];
+
+    float * row = (float *) malloc(n_input * sizeof(float));
+    GGML_ASSERT(row != NULL);
+
+    for (int64_t task = params->ith; task < n_tasks; task += params->nth) {
+        const int64_t output = task % n_output;
+        int64_t       rest   = task / n_output;
+        const int64_t column = rest % n_columns;
+        rest /= n_columns;
+        const int64_t i2 = rest % dst->ne[2];
+        const int64_t i3 = rest / dst->ne[2];
+
+        const void * weight =
+            (const char *) src0->data + output * src0->nb[1] + (i2 / r2) * src0->nb[2] + (i3 / r3) * src0->nb[3];
+        const float * activation =
+            (const float *) ((const char *) src1->data + column * src1->nb[1] + i2 * src1->nb[2] + i3 * src1->nb[3]);
+        float * result = (float *) ((char *) dst->data + output * dst->nb[0] + column * dst->nb[1] + i2 * dst->nb[2] +
+                                    i3 * dst->nb[3]);
+        *result        = ggml_compute_forward_mul_mat_precise_ref_dot(traits, weight, activation, row, n_input);
+    }
+
+    free(row);
+    return true;
+}
+
 static void ggml_compute_forward_mul_mat_one_chunk(
     const struct ggml_compute_params * params,
     struct ggml_tensor * dst,
@@ -1254,6 +1323,9 @@ static void ggml_compute_forward_mul_mat_one_chunk(
 void ggml_compute_forward_mul_mat(
         const struct ggml_compute_params * params,
               struct ggml_tensor * dst) {
+    if (ggml_compute_forward_mul_mat_precise_ref(params, dst)) {
+        return;
+    }
 
     const struct ggml_tensor * src0 = dst->src[0];
     const struct ggml_tensor * src1 = dst->src[1];
@@ -1523,6 +1595,56 @@ static void ggml_compute_forward_mul_mat_id_one_chunk(
     }
 }
 
+static bool ggml_compute_forward_mul_mat_id_precise_ref(const struct ggml_compute_params * params,
+                                                        struct ggml_tensor *               dst) {
+    const struct ggml_tensor * src0 = dst->src[0];
+    const struct ggml_tensor * src1 = dst->src[1];
+    const struct ggml_tensor * ids  = dst->src[2];
+
+    if (!params->use_ref || src1->type != GGML_TYPE_F32 || !ggml_compute_forward_mul_mat_precise_ref_type(src0->type)) {
+        return false;
+    }
+
+    GGML_ASSERT(src0->ne[0] == src1->ne[0]);
+    GGML_ASSERT(dst->ne[0] == src0->ne[1]);
+    GGML_ASSERT(dst->ne[1] == ids->ne[0]);
+    GGML_ASSERT(dst->ne[2] == ids->ne[1]);
+    GGML_ASSERT(src1->ne[1] == 1 || src1->ne[1] == ids->ne[0]);
+    GGML_ASSERT(src1->ne[2] == ids->ne[1]);
+    GGML_ASSERT(src0->ne[3] == 1 && src1->ne[3] == 1 && dst->ne[3] == 1);
+
+    const struct ggml_type_traits * traits = ggml_get_type_traits(src0->type);
+    GGML_ASSERT(traits->to_float != NULL);
+
+    const int64_t n_input  = src0->ne[0];
+    const int64_t n_output = src0->ne[1];
+    const int64_t n_ids    = ids->ne[0];
+    const int64_t n_tokens = ids->ne[1];
+    const int64_t n_tasks  = n_output * n_ids * n_tokens;
+
+    float * row = (float *) malloc(n_input * sizeof(float));
+    GGML_ASSERT(row != NULL);
+
+    for (int64_t task = params->ith; task < n_tasks; task += params->nth) {
+        const int64_t output  = task % n_output;
+        const int64_t rest    = task / n_output;
+        const int64_t id      = rest % n_ids;
+        const int64_t token   = rest / n_ids;
+        const int64_t src1_id = src1->ne[1] == 1 ? 0 : id;
+        const int32_t expert  = *(const int32_t *) ((const char *) ids->data + id * ids->nb[0] + token * ids->nb[1]);
+        GGML_ASSERT(expert >= 0 && expert < src0->ne[2]);
+
+        const void *  weight = (const char *) src0->data + output * src0->nb[1] + expert * src0->nb[2];
+        const float * activation =
+            (const float *) ((const char *) src1->data + src1_id * src1->nb[1] + token * src1->nb[2]);
+        float * result = (float *) ((char *) dst->data + output * dst->nb[0] + id * dst->nb[1] + token * dst->nb[2]);
+        *result        = ggml_compute_forward_mul_mat_precise_ref_dot(traits, weight, activation, row, n_input);
+    }
+
+    free(row);
+    return true;
+}
+
 static void * incr_ptr_aligned(void ** p, size_t size, size_t align) {
 
     void * ptr = *p;
@@ -1534,6 +1656,9 @@ static void * incr_ptr_aligned(void ** p, size_t size, size_t align) {
 static void ggml_compute_forward_mul_mat_id(
         const struct ggml_compute_params * params,
               struct ggml_tensor * dst) {
+    if (ggml_compute_forward_mul_mat_id_precise_ref(params, dst)) {
+        return;
+    }
 
     const struct ggml_tensor * src0 = dst->src[0];
     const struct ggml_tensor * src1 = dst->src[1];

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

@@ -638,6 +638,24 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net(
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexer(ggml_metal_library_t lib,
+                                                                                  const ggml_tensor *  op) {
+    GGML_ASSERT(op->op == GGML_OP_LIGHTNING_INDEXER);
+
+    char base[256];
+    char name[256];
+
+    snprintf(base, sizeof(base), "kernel_lightning_indexer_%s", ggml_type_name(op->src[1]->type));
+    snprintf(name, sizeof(name), "%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);
+    }
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri(ggml_metal_library_t lib, const ggml_tensor * op) {
     char base[256];
     char name[256];

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

@@ -129,6 +129,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan          (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv              (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net   (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexer(ggml_metal_library_t       lib,
+                                                                                         const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext        (ggml_metal_library_t lib, const struct ggml_tensor * op, int nsg, int nxpsg, int r1ptg);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm            (ggml_metal_library_t lib, const struct ggml_tensor * op);

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

@@ -1272,6 +1272,28 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
             return true;
         case GGML_OP_GATED_DELTA_NET:
             return has_simdgroup_reduction && op->src[2]->ne[0] % 32 == 0;
+        case GGML_OP_LIGHTNING_INDEXER:
+            {
+                const enum ggml_type k_type = op->src[1]->type;
+                const bool k_type_supported =
+                    k_type == GGML_TYPE_F32 ||
+                    k_type == GGML_TYPE_F16 ||
+                    (k_type == GGML_TYPE_BF16 && has_bfloat) ||
+                    k_type == GGML_TYPE_Q8_0 ||
+                    k_type == GGML_TYPE_Q5_1 ||
+                    k_type == GGML_TYPE_Q5_0 ||
+                    k_type == GGML_TYPE_Q4_1 ||
+                    k_type == GGML_TYPE_Q4_0 ||
+                    k_type == GGML_TYPE_IQ4_NL;
+
+                return op->type == GGML_TYPE_F32 &&
+                    op->src[0]->type == GGML_TYPE_F32 &&
+                    k_type_supported &&
+                    op->src[2]->type == GGML_TYPE_F32 &&
+                    op->src[3]->type == GGML_TYPE_F16 &&
+                    ggml_is_contiguous_rows(op->src[0]) &&
+                    ggml_is_contiguous_rows(op->src[1]);
+            }
         case GGML_OP_SOLVE_TRI:
         case GGML_OP_MUL_MAT:
         case GGML_OP_MUL_MAT_ID:

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

@@ -927,6 +927,49 @@ typedef struct {
     uint64_t nb3;
 } ggml_metal_kargs_gated_delta_net;
 
+typedef struct {
+    int32_t  ne00;
+    int32_t  ne01;
+    int32_t  ne02;
+    int32_t  ne03;
+    uint64_t nb00;
+    uint64_t nb01;
+    uint64_t nb02;
+    uint64_t nb03;
+    int32_t  ne10;
+    int32_t  ne11;
+    int32_t  ne12;
+    int32_t  ne13;
+    uint64_t nb10;
+    uint64_t nb11;
+    uint64_t nb12;
+    uint64_t nb13;
+    int32_t  ne20;
+    int32_t  ne21;
+    int32_t  ne22;
+    int32_t  ne23;
+    uint64_t nb20;
+    uint64_t nb21;
+    uint64_t nb22;
+    uint64_t nb23;
+    int32_t  ne30;
+    int32_t  ne31;
+    int32_t  ne32;
+    int32_t  ne33;
+    uint64_t nb30;
+    uint64_t nb31;
+    uint64_t nb32;
+    uint64_t nb33;
+    int32_t  ne0;
+    int32_t  ne1;
+    int32_t  ne2;
+    int32_t  ne3;
+    uint64_t nb0;
+    uint64_t nb1;
+    uint64_t nb2;
+    uint64_t nb3;
+} ggml_metal_kargs_lightning_indexer;
+
 typedef struct {
     int32_t  ne00;
     int32_t  ne01;

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

@@ -337,6 +337,11 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
             {
                 n_fuse = ggml_metal_op_gated_delta_net(ctx, idx);
             } break;
+        case GGML_OP_LIGHTNING_INDEXER:
+            {
+                n_fuse = ggml_metal_op_lightning_indexer(ctx, idx);
+            }
+            break;
         case GGML_OP_SOLVE_TRI:
             {
                 n_fuse = ggml_metal_op_solve_tri(ctx, idx);
@@ -1674,6 +1679,83 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) {
     return 1;
 }
 
+int ggml_metal_op_lightning_indexer(ggml_metal_op_t ctx, int idx) {
+    ggml_tensor * op = ctx->node(idx);
+
+    ggml_metal_library_t lib = ctx->lib;
+    ggml_metal_encoder_t enc = ctx->enc;
+
+    GGML_TENSOR_LOCALS(int32_t, ne0, op->src[0], ne);
+    GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
+    GGML_TENSOR_LOCALS(int32_t, ne1, op->src[1], ne);
+    GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb);
+    GGML_TENSOR_LOCALS(int32_t, ne2, op->src[2], ne);
+    GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb);
+    GGML_TENSOR_LOCALS(int32_t, ne3, op->src[3], ne);
+    GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb);
+    GGML_TENSOR_LOCALS(int32_t, ne, op, ne);
+    GGML_TENSOR_LOCALS(uint64_t, nb, op, nb);
+
+    ggml_metal_kargs_lightning_indexer args = {
+        /*.ne00 =*/ne00,
+        /*.ne01 =*/ne01,
+        /*.ne02 =*/ne02,
+        /*.ne03 =*/ne03,
+        /*.nb00 =*/nb00,
+        /*.nb01 =*/nb01,
+        /*.nb02 =*/nb02,
+        /*.nb03 =*/nb03,
+        /*.ne10 =*/ne10,
+        /*.ne11 =*/ne11,
+        /*.ne12 =*/ne12,
+        /*.ne13 =*/ne13,
+        /*.nb10 =*/nb10,
+        /*.nb11 =*/nb11,
+        /*.nb12 =*/nb12,
+        /*.nb13 =*/nb13,
+        /*.ne20 =*/ne20,
+        /*.ne21 =*/ne21,
+        /*.ne22 =*/ne22,
+        /*.ne23 =*/ne23,
+        /*.nb20 =*/nb20,
+        /*.nb21 =*/nb21,
+        /*.nb22 =*/nb22,
+        /*.nb23 =*/nb23,
+        /*.ne30 =*/ne30,
+        /*.ne31 =*/ne31,
+        /*.ne32 =*/ne32,
+        /*.ne33 =*/ne33,
+        /*.nb30 =*/nb30,
+        /*.nb31 =*/nb31,
+        /*.nb32 =*/nb32,
+        /*.nb33 =*/nb33,
+        /*.ne0  =*/ne0,
+        /*.ne1  =*/ne1,
+        /*.ne2  =*/ne2,
+        /*.ne3  =*/ne3,
+        /*.nb0  =*/nb0,
+        /*.nb1  =*/nb1,
+        /*.nb2  =*/nb2,
+        /*.nb3  =*/nb3,
+    };
+
+    auto      pipeline = ggml_metal_library_get_pipeline_lightning_indexer(lib, op);
+    const int nth      = std::min(64, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
+    int       ida      = 0;
+
+    ggml_metal_encoder_set_pipeline(enc, pipeline);
+    ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), ida++);
+    ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), ida++);
+    ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[1]), ida++);
+    ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[2]), ida++);
+    ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[3]), ida++);
+    ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), ida++);
+
+    ggml_metal_encoder_dispatch_threadgroups(enc, (ne0 + nth - 1) / nth, ne1, ne3, nth, 1, 1);
+
+    return 1;
+}
+
 int ggml_metal_op_solve_tri(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 

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

@@ -59,6 +59,7 @@ int ggml_metal_op_ssm_conv          (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_ssm_scan          (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_rwkv              (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_gated_delta_net   (ggml_metal_op_t ctx, int idx);
+int ggml_metal_op_lightning_indexer(ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_solve_tri         (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_set               (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_cpy               (ggml_metal_op_t ctx, int idx);