diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 40e50c5c9..3f44d018d 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1716,18 +1716,33 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } if (ids_tensor != prev_ids_tensor) { - ids.resize(ggml_nbytes(ids_tensor) / sizeof(int32_t)); - ggml_backend_tensor_get_async(ids_backend, ids_tensor, ids.data(), 0, ggml_nbytes(ids_tensor)); - ggml_backend_synchronize(ids_backend); + // Reading the ids back forces a full pipeline flush, since the ids are + // produced on the same device we are about to upload to. Once the batch + // draws enough experts, essentially all of them come back used and the + // readback has bought nothing, so skip it and upload the whole tensor as one + // range. Uploading an expert that no id selects cannot change the result: + // mul_mat_id only reads the rows the ids point at, so this stays exact. + const bool all_experts_used = ids_tensor->ne[0] * ids_tensor->ne[1] >= 4 * n_expert; // find the used experts used_ids.clear(); used_ids.resize(ggml_bitset_size(n_expert)); - for (int64_t i1 = 0; i1 < ids_tensor->ne[1]; i1++) { - for (int64_t i0 = 0; i0 < ids_tensor->ne[0]; i0++) { - int32_t id = ids[i1 * ids_tensor->nb[1]/sizeof(int32_t) + i0 * ids_tensor->nb[0]/sizeof(int32_t)]; - GGML_ASSERT(id >= 0 && id < n_expert); - ggml_bitset_set(used_ids.data(), id); + + if (all_experts_used) { + for (int32_t i = 0; i < n_expert; i++) { + ggml_bitset_set(used_ids.data(), i); + } + } else { + ids.resize(ggml_nbytes(ids_tensor) / sizeof(int32_t)); + ggml_backend_tensor_get_async(ids_backend, ids_tensor, ids.data(), 0, ggml_nbytes(ids_tensor)); + ggml_backend_synchronize(ids_backend); + + for (int64_t i1 = 0; i1 < ids_tensor->ne[1]; i1++) { + for (int64_t i0 = 0; i0 < ids_tensor->ne[0]; i0++) { + int32_t id = ids[i1 * ids_tensor->nb[1]/sizeof(int32_t) + i0 * ids_tensor->nb[0]/sizeof(int32_t)]; + GGML_ASSERT(id >= 0 && id < n_expert); + ggml_bitset_set(used_ids.data(), id); + } } }