Chunked experts (CPU) (#2202)

* Chunked experts

* Option to turn it off at compile time

ggml/CMakeLists.txt

@@ -95,6 +95,7 @@ option(GGML_LASX        "ggml: enable lasx"             ON)
 option(GGML_LSX         "ggml: enable lsx"              ON)
 option(GGML_SVE         "ggml: enable SVE"              OFF)
 option(GGML_NCCL        "ggml: enable NCCL"             ON)
+option(GGML_EXPERT_CHUNKING "ggml: enable expert chunking" ON)
 
 if (WIN32)
     # Default to Windows 10 (0x0A00). Windows 8 / 8.1 reached end of support in

ggml/src/CMakeLists.txt

@@ -219,6 +219,10 @@ if (GGML_IQK_MUL_MAT)
     endif()
 endif()
 
+if (GGML_EXPERT_CHUNKING)
+    add_compile_definitions(GGML_EXPERT_CHUNKING)
+endif()
+
 if (GGML_CUDA)
     cmake_minimum_required(VERSION 3.18)  # for CMAKE_CUDA_ARCHITECTURES
 

ggml/src/ggml.c

@@ -18200,8 +18200,61 @@ static void ggml_compute_forward_mul_mat_id(
         }
     }
 
+    if (ith == 0) {
+        atomic_store(&params->shared->current_chunk, nth);
+    }
+
     ggml_barrier(params->shared);
 
+#if GGML_USE_IQK_MULMAT
+#if defined GGML_EXPERT_CHUNKING
+    if (ne13 == 1 && dst->type == GGML_TYPE_F32) {
+        const void * wdata_mm    = (src1->type == vec_dot_type) ? src1->data : params->wdata;
+        const size_t row_size_mm = ggml_row_size(vec_dot_type, ne10);
+
+        const int chunks_per_expert = MAX(1, MIN(nth, (int)(ne01 / 32)));
+
+        int total_chunks = 0;
+        for (int a = 0; a < n_as; a++) {
+            if (matrix_row_counts[a] > 0) total_chunks += chunks_per_expert;
+        }
+
+        int chunk_id = ith;
+        while (chunk_id < total_chunks) {
+            // Map global chunk_id to (expert_index, local_chunk_index)
+            int acc = 0, cur_a = -1, local_chunk = 0;
+            for (int a = 0; a < n_as; a++) {
+                if (matrix_row_counts[a] == 0) continue;
+                if (chunk_id < acc + chunks_per_expert) {
+                    cur_a = a;
+                    local_chunk = chunk_id - acc;
+                    break;
+                }
+                acc += chunks_per_expert;
+            }
+
+            const char * src0_cur = (const char *) src0->data + cur_a*nb02;
+
+            if (!iqk_mul_mat_moe(ne01, matrix_row_counts[cur_a], ne00, ne11,
+                        src0->type, src0_cur, nb01,
+                        vec_dot_type, (const char *)wdata_mm, row_size_mm,
+                        (float *)dst->data, nb1, nb2,
+                        matrix_rows + cur_a*ne12, local_chunk, chunks_per_expert)) goto IQK_MulMat_Not_Available0;
+
+            chunk_id = atomic_fetch_add(&params->shared->current_chunk, 1);
+        }
+        return;
+    }
+IQK_MulMat_Not_Available0:;
+
+    // Reset counter for fallback path
+    if (ith == 0) {
+        atomic_store(&params->shared->current_chunk, 0);
+    }
+    ggml_barrier(params->shared);
+#endif
+#endif
+
     // compute each matrix multiplication in sequence
     for (int cur_a = 0; cur_a < n_as; ++cur_a) {
         const int64_t cne1 = matrix_row_counts[cur_a];
@@ -18476,6 +18529,75 @@ static void ggml_compute_forward_mul_mat_id_up_gate(
         }
     }
 
+#if defined GGML_EXPERT_CHUNKING
+
+    if (ith == 0) {
+        atomic_store(&params->shared->current_chunk, nth);
+    }
+
+    ggml_barrier(params->shared);
+
+    const float limit = *(const float *)(dst->op_params + 1);
+
+    const void * wdata_ug    = (src1->type == vec_dot_type) ? src1->data : params->wdata;
+    const size_t row_size_ug = ggml_row_size(vec_dot_type, ne10);
+    const int64_t nr0_base = src0_2 ? ne01 : ne01/2;
+
+    const int chunks_per_expert_ug = MAX(1, MIN(nth, (int)(nr0_base / 32)));
+
+    int total_chunks_ug = 0;
+    for (int a = 0; a < n_as; a++) {
+        if (matrix_row_counts[a] > 0) total_chunks_ug += chunks_per_expert_ug;
+    }
+
+    int last_a = 0;
+    int last_acc = 0;
+    int chunk_id_ug = ith;
+    while (chunk_id_ug < total_chunks_ug) {
+        int acc = last_acc, cur_a = -1, local_chunk = 0;
+        for (int a = last_a; a < n_as; a++) {
+            if (matrix_row_counts[a] == 0) continue;
+            if (chunk_id_ug < acc + chunks_per_expert_ug) {
+                cur_a = a;
+                local_chunk = chunk_id_ug - acc;
+                break;
+            }
+            acc += chunks_per_expert_ug;
+        }
+        if (cur_a < 0) {
+            return;
+        }
+        last_a = cur_a;
+        last_acc = acc;
+
+        const char *src0_1_cur, *src0_2_cur, *up_b_cur = NULL, *gate_b_cur = NULL;
+        if (src0_2) {
+            src0_1_cur = (const char *) src0_1->data + cur_a*nb02;
+            src0_2_cur = (const char *) src0_2->data + cur_a*nb02;
+            up_b_cur   = up_b   ? (const char *)up_b->data + cur_a*nb41 : NULL;
+            gate_b_cur = gate_b ? (const char *)gate_b->data + cur_a*nb51 : NULL;
+        } else {
+            src0_2_cur = (const char *) src0_1->data + cur_a*nb02;
+            src0_1_cur = src0_2_cur + nb02/2;
+            if (up_b) {
+                GGML_ASSERT(!gate_b);
+                gate_b_cur = (const char *)up_b->data + cur_a*nb41;
+                up_b_cur   = gate_b_cur + nb41/2;
+            }
+        }
+
+        if (!iqk_moe_fused_up_gate(nr0_base, matrix_row_counts[cur_a], ne00, ne11, dst->op_params[0],
+                            type, src0_1_cur, src0_2_cur, nb01,
+                            vec_dot_type, (const char *)wdata_ug, row_size_ug,
+                            up_b_cur, gate_b_cur,
+                            (float *)dst->data, nb1, nb2,
+                            matrix_rows + cur_a*ne12, limit, local_chunk, chunks_per_expert_ug)) GGML_ABORT("fatal error");
+
+        chunk_id_ug = atomic_fetch_add(&params->shared->current_chunk, 1);
+    }
+
+#else
+
     ggml_barrier(params->shared);
 
     const float limit = *(const float *)(dst->op_params + 1);
@@ -18520,7 +18642,7 @@ static void ggml_compute_forward_mul_mat_id_up_gate(
                             matrix_rows + cur_a*ne12, limit, ith, nth)) GGML_ABORT("fatal error");
 
     }
-
+#endif
 #undef MMID_MATRIX_ROW
 }
 

ggml/src/iqk/iqk_mul_mat.cpp

@@ -1928,8 +1928,10 @@ bool iqk_indexer_topk(struct ggml_tensor * dst, void * work_buffer, barrier_t ba
             std::partial_sort(sorted, sorted + n_top_k, sorted + k->ne[1], [score] (int32_t l, int32_t r) -> bool { return score[l] > score[r]; });
             std::memcpy((char *)dst->data + dst->nb[1]*iq, sorted, n_top_k*sizeof(int32_t));
         }
+        if (iq + 1 < q->ne[2]) {
             barrier(barrier_data);
         }
+    }
 
     return true;
 }