From dbc449afd0b236997150209f8828dade6367907b Mon Sep 17 00:00:00 2001 From: Lumpiasty Date: Mon, 27 Jul 2026 23:56:14 +0200 Subject: [PATCH] sched: skip the MoE routing-ids readback when the batch uses every expert Offloaded expert weights are uploaded per split, and to decide which experts to upload the scheduler reads the routing ids back to the host. The ids are produced on the same device we are about to upload to, so the readback forces a full pipeline flush - 47 of them per eval on a 48-layer MoE. Once the batch draws enough experts the readback stops telling us anything. At 2048 tokens x 10 experts over 256 experts every expert comes back used, so the bitset is all ones and the copy is a single whole-tensor range anyway. Skip the readback when the batch guarantees that. Uploading an expert that no id selects cannot change the result, since mul_mat_id only reads the rows the ids point at, so this stays exact. Decode is unaffected: n_ids there is the number of experts per token, far below the threshold, so it keeps the bitset path. Measured on an RX 580 (Vulkan, Polaris) running Laguna-S-2.1 118B IQ2_M with all experts host-resident: pp2048 +4.7% at depth 32k, +0.9% at 16k, tg32 +3%. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-backend.cpp | 31 +++++++++++++++++++++++-------- 1 file changed, 23 insertions(+), 8 deletions(-) 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); + } } }