ikawrakow/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-03-09 21:40:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

This is very slightly better (#762)

Browse Source 
Co-authored-by: Iwan Kawrakow <iwan.kawrakow@gmail.com>
This commit is contained in:
Kawrakow
2025-09-05 21:31:02 +02:00
committed by GitHub
parent f74dd77143
commit 540a26514f
1 changed files with 47 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
ggml/src/ggml-backend.cpp
View File

@@ -1857,6 +1857,10 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) {
static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) {
struct ggml_backend_sched_split * splits = sched->splits;
std::vector<int32_t> ids;
std::vector<uint32_t> unique_ids;
ggml_tensor * last_ids_tensor = nullptr;
for (int i = 0; i < sched->n_splits; i++) {
#if IK_PRINT_TIMING
@@ -1865,14 +1869,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
struct ggml_backend_sched_split * split = &splits[i];
int split_backend_id = split->backend_id;
ggml_backend_t split_backend = sched->backends[split_backend_id];
ggml_backend_t last_input_backend = nullptr;
int cur_arg = 0;
std::vector<int32_t> ids;
std::set<int32_t> unique_ids;
//printf("Graph split %d has %d inputs:\n", i, split->n_inputs);
//for (int j = 0; j < split->n_inputs; j++) printf(" %s, %s\n", split->inputs[j]->name,
// split->inputs[j]->src[0] ? split->inputs[j]->src[0]->name : "none");
// copy the input tensors to the split backend
for (int j = 0; j < split->n_inputs; j++) {
@@ -1903,32 +1902,46 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
node->src[cur_arg] == input_cpy &&
(node->op == GGML_OP_MUL_MAT_ID || node->op == GGML_OP_MOE_FUSED_UP_GATE)) {
if (ids.empty()) {
// find the ids
ggml_tensor * ids_tensor = node->op == GGML_OP_MUL_MAT_ID ? node->src[2] : node->src[3];
ids.resize(ggml_nbytes(ids_tensor) / sizeof(int32_t));
if (input_backend != last_input_backend) {
ggml_backend_synchronize(input_backend);
last_input_backend = input_backend;
}
ggml_backend_tensor_get_async(split_backend, ids_tensor, ids.data(), 0, ggml_nbytes(ids_tensor));
//printf("node: %s have %d inputs, processing input %d\n", node->name, split->n_inputs, j);
ggml_tensor * ids_tensor = node->op == GGML_OP_MUL_MAT_ID ? node->src[2] : node->src[3];
auto ids_backend = split_backend;
ggml_backend_synchronize(split_backend);
// if the ids tensor is also an input of the split, it may not have been copied yet to the split backend
// in that case, we use the original ids tensor
for (int jj = j + 1; jj < split->n_inputs; ++jj) {
if (ids_tensor == tensor_copy(split->inputs[jj], split_backend_id, sched->cur_copy)) {
ids_tensor = split->inputs[jj];
ids_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[jj]);
break;
}
}
int n_expert = input->ne[2];
if (ids_tensor != last_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);
unique_ids.resize((n_expert + 31)/32);
std::memset(unique_ids.data(), 0, unique_ids.size()*sizeof(uint32_t));
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)];
unique_ids.insert(id);
unique_ids[id >> 5] |= (1u << (id & 31));
}
}
// group consecutive experts and copy them together
GGML_ASSERT(!unique_ids.empty());
last_ids_tensor = ids_tensor;
}
auto it = unique_ids.begin();
int32_t first_id = *it;
int32_t last_id = first_id;
auto copy_experts = [&](int32_t first_id, int32_t last_id) {
const size_t expert_size = (node->op == GGML_OP_MUL_MAT_ID || node->op == GGML_OP_MOE_FUSED_UP_GATE) ? input->nb[2] : input->nb[1];
const size_t expert_offset = first_id * expert_size;
@@ -1944,20 +1957,22 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
};
for (++it; it != unique_ids.end(); ++it) {
const int32_t id = *it;
auto next_on_id = [&unique_ids, n_expert] (int id) {
while (id < n_expert && (unique_ids[id >> 5] & (1u << (id & 31))) == 0) ++id;
return id;
};
auto next_off_id = [&unique_ids, n_expert] (int id) {
while (id < n_expert && (unique_ids[id >> 5] & (1u << (id & 31))) != 0) ++id;
return id;
};
if (id == last_id + 1) {
last_id = id;
continue;
}
copy_experts(first_id, last_id);
first_id = id;
last_id = id;
int first_id = next_on_id(0);
while (first_id < n_expert) {
int last_id = next_off_id(first_id+1);
copy_experts(first_id, last_id-1);
first_id = next_on_id(last_id);
}
copy_experts(first_id, last_id);
if (node->op == GGML_OP_MOE_FUSED_UP_GATE) ++cur_arg;
} else
// try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events