diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt index ec87b0807..719cd8c04 100644 --- a/ggml/CMakeLists.txt +++ b/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 diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index a128c0c06..ad329e791 100644 --- a/ggml/src/CMakeLists.txt +++ b/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 diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 6a097177f..00c83a5a9 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -18200,8 +18200,61 @@ static void ggml_compute_forward_mul_mat_id( } } + if (ith == 0) { + atomic_store(¶ms->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(¶ms->shared->current_chunk, 1); + } + return; + } +IQK_MulMat_Not_Available0:; + + // Reset counter for fallback path + if (ith == 0) { + atomic_store(¶ms->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(¶ms->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(¶ms->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 } diff --git a/ggml/src/iqk/iqk_mul_mat.cpp b/ggml/src/iqk/iqk_mul_mat.cpp index 14ed16348..bbfba28ca 100644 --- a/ggml/src/iqk/iqk_mul_mat.cpp +++ b/ggml/src/iqk/iqk_mul_mat.cpp @@ -1928,7 +1928,9 @@ 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)); } - barrier(barrier_data); + if (iq + 1 < q->ne[2]) { + barrier(barrier_data); + } } return true;