ggml : adjust logic for offloading ops to weight's backend (#25832)

* ggml : adjust logic for offloading ops to weight's backend

* llama : dsv4 graph fixes
This commit is contained in:
Georgi Gerganov
2026-07-27 14:54:46 +03:00
committed by GitHub
parent 7ef790f90a
commit dee2a846b8
3 changed files with 34 additions and 20 deletions
+26 -17
View File
@@ -906,26 +906,35 @@ static int ggml_backend_sched_backend_id_from_cur(ggml_backend_sched_t sched, st
}
// operations with weights are preferably run on the same backend as the weights
for (int i = 0; i < GGML_MAX_SRC; i++) {
const struct ggml_tensor * src = tensor->src[i];
if (src == NULL) {
continue;
}
// skip ROPE since the rope freqs tensor is too small to choose a backend based on it
// not an ideal solution
if (tensor->op != GGML_OP_ROPE && src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
int src_backend_id = ggml_backend_sched_backend_from_buffer(sched, src, tensor);
// check if a backend with higher prio wants to offload the op
if (sched->op_offload && src_backend_id == sched->n_backends - 1 && ggml_backend_buffer_is_host(src->buffer)) {
for (int b = 0; b < src_backend_id; b++) {
if (ggml_backend_supports_op(sched->backends[b], tensor) && ggml_backend_offload_op(sched->backends[b], tensor)) {
SET_CAUSE(tensor, "1.off");
return b;
// TODO: there are exceptions (see below) - not an ideal solution
bool allow = true;
// skip ROPE since the rope freqs tensor is too small to choose a backend based on it
allow = allow && tensor->op != GGML_OP_ROPE;
// skip FLASH_ATTN_EXT since the sinks tensor is too small to choose a based based on it
allow = allow && tensor->op != GGML_OP_FLASH_ATTN_EXT;
if (allow) {
for (int i = 0; i < GGML_MAX_SRC; i++) {
const struct ggml_tensor * src = tensor->src[i];
if (src == NULL) {
continue;
}
if (src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
int src_backend_id = ggml_backend_sched_backend_from_buffer(sched, src, tensor);
// check if a backend with higher prio wants to offload the op
if (sched->op_offload && src_backend_id == sched->n_backends - 1 && ggml_backend_buffer_is_host(src->buffer)) {
for (int b = 0; b < src_backend_id; b++) {
if (ggml_backend_supports_op(sched->backends[b], tensor) && ggml_backend_offload_op(sched->backends[b], tensor)) {
SET_CAUSE(tensor, "1.off");
return b;
}
}
}
SET_CAUSE(tensor, "1.wgt%d", i);
return src_backend_id;
}
SET_CAUSE(tensor, "1.wgt%d", i);
return src_backend_id;
}
}