Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ggml : size prefetch slots per layer and fix fallback use-after-free

Browse Source 
Profiling showed the 2-slot rotation stalls ~5.5ms per layer waiting for
the down_exps upload: 3 tensors per MoE layer need 3 slots for a full
layer of lookahead. Default is now 3 (one layer), configurable via the
env var value, degrading to however many slots fit on OOM.

Also fixes a use-after-free: on a failed slot regrow the repointed
staging tensors kept dangling pointers into freed device memory, which
graph reuse carried into later evals. Staging repoints are now restored
right after kernel launch, slots are sized once from the graph max, and
allocation happens before freeing.

Qwen3.6-35B-A3B pp2048 on RTX 3060: 1643 -> 1880 t/s (mainline: 1143).
This commit is contained in:
thecodacus
2026-07-26 21:58:43 +02:00
committed by Lumpiasty
parent 9a34ee7dbe
commit 31ba631b76
1 changed files with 90 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-backend.cpp
+90 -19
View File
@@ -761,6 +761,10 @@ static bool ggml_is_view_op(enum ggml_op op) {
#define GGML_SCHED_MAX_COPIES 4
#endif
#ifndef GGML_SCHED_MAX_PREFETCH_SLOTS
#define GGML_SCHED_MAX_PREFETCH_SLOTS 8
#endif
struct ggml_backend_sched_split {
int backend_id;
int i_start;
@@ -824,10 +828,11 @@ struct ggml_backend_sched {
// they overlap compute, alternating between two staging slots
bool prefetch_experts;
ggml_backend_t prefetch_backend;
ggml_backend_buffer_t prefetch_slots[2];
ggml_backend_event_t prefetch_ready[2];
ggml_backend_event_t prefetch_free[2];
bool prefetch_used[2];
int prefetch_n_slots;
ggml_backend_buffer_t prefetch_slots[GGML_SCHED_MAX_PREFETCH_SLOTS];
ggml_backend_event_t prefetch_ready[GGML_SCHED_MAX_PREFETCH_SLOTS];
ggml_backend_event_t prefetch_free[GGML_SCHED_MAX_PREFETCH_SLOTS];
bool prefetch_used[GGML_SCHED_MAX_PREFETCH_SLOTS];
int prefetch_cur;
int debug;
@@ -1550,6 +1555,40 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) {
return true;
}
static void ggml_backend_sched_prefetch_disable(ggml_backend_sched_t sched, ggml_backend_t split_backend) {
sched->prefetch_experts = false;
if (sched->prefetch_backend) {
ggml_backend_synchronize(split_backend);
ggml_backend_synchronize(sched->prefetch_backend);
}
for (int i = 0; i < sched->prefetch_n_slots; i++) {
ggml_backend_buffer_free(sched->prefetch_slots[i]);
sched->prefetch_slots[i] = NULL;
sched->prefetch_used[i] = false;
}
}
// slots are sized once for the largest offloaded expert tensor in the current graph so
// that they never need to grow mid-eval
static size_t ggml_backend_sched_prefetch_max_size(ggml_backend_sched_t sched) {
size_t max_size = 0;
for (int split_id = 0; split_id < sched->n_splits; split_id++) {
struct ggml_backend_sched_split * split = &sched->splits[split_id];
if (split->graph.n_nodes == 0 || split->graph.nodes[0]->op != GGML_OP_MUL_MAT_ID) {
continue;
}
for (int input_id = 0; input_id < split->n_inputs; input_id++) {
const ggml_tensor * input = split->inputs[input_id];
if (input->buffer &&
ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
ggml_backend_buffer_is_host(input->buffer)) {
max_size = std::max(max_size, ggml_nbytes(input));
}
}
}
return max_size;
}
static bool ggml_backend_sched_prefetch_init(ggml_backend_sched_t sched, ggml_backend_t split_backend, size_t size) {
if (sched->prefetch_backend == NULL) {
ggml_backend_dev_t dev = split_backend->device;
@@ -1564,7 +1603,7 @@ static bool ggml_backend_sched_prefetch_init(ggml_backend_sched_t sched, ggml_ba
sched->prefetch_experts = false;
return false;
}
for (int i = 0; i < 2; i++) {
for (int i = 0; i < sched->prefetch_n_slots; i++) {
sched->prefetch_ready[i] = ggml_backend_event_new(dev);
sched->prefetch_free[i] = ggml_backend_event_new(dev);
if (sched->prefetch_ready[i] == NULL || sched->prefetch_free[i] == NULL) {
@@ -1574,18 +1613,31 @@ static bool ggml_backend_sched_prefetch_init(ggml_backend_sched_t sched, ggml_ba
}
}
size = std::max(size, ggml_backend_sched_prefetch_max_size(sched));
ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(split_backend);
for (int i = 0; i < 2; i++) {
for (int i = 0; i < sched->prefetch_n_slots; i++) {
if (sched->prefetch_slots[i] == NULL || ggml_backend_buffer_get_size(sched->prefetch_slots[i]) < size) {
ggml_backend_synchronize(split_backend);
ggml_backend_synchronize(sched->prefetch_backend);
ggml_backend_buffer_free(sched->prefetch_slots[i]);
sched->prefetch_slots[i] = ggml_backend_buft_alloc_buffer(buft, size);
sched->prefetch_used[i] = false;
if (sched->prefetch_slots[i] == NULL) {
sched->prefetch_experts = false;
// allocate before freeing so a failure leaves the old slot intact
ggml_backend_buffer_t new_buf = ggml_backend_buft_alloc_buffer(buft, size);
if (new_buf == NULL) {
// overlap needs at least 2 slots, otherwise run with what fits
if (i >= 2 && sched->prefetch_slots[0] != NULL &&
ggml_backend_buffer_get_size(sched->prefetch_slots[0]) >= size) {
sched->prefetch_n_slots = i;
sched->prefetch_cur = 0;
return true;
}
ggml_backend_sched_prefetch_disable(sched, split_backend);
return false;
}
if (sched->prefetch_slots[i] != NULL) {
ggml_backend_synchronize(split_backend);
ggml_backend_synchronize(sched->prefetch_backend);
ggml_backend_buffer_free(sched->prefetch_slots[i]);
}
sched->prefetch_slots[i] = new_buf;
sched->prefetch_used[i] = false;
}
}
return true;
@@ -1604,6 +1656,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
int split_backend_id = split->backend_id;
ggml_backend_t split_backend = sched->backends[split_backend_id];
int split_prefetch_slot = -1;
ggml_tensor * prefetch_input_cpy = NULL;
ggml_backend_buffer_t prefetch_saved_buffer = NULL;
void * prefetch_saved_data = NULL;
// copy the input tensors to the split backend
for (int input_id = 0; input_id < split->n_inputs; input_id++) {
@@ -1633,11 +1688,17 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
if (ids->ne[0]*ids->ne[1] >= 2*n_expert &&
ggml_backend_sched_prefetch_init(sched, split_backend, ggml_nbytes(input))) {
const int slot = sched->prefetch_cur;
sched->prefetch_cur ^= 1;
sched->prefetch_cur = (sched->prefetch_cur + 1) % sched->prefetch_n_slots;
// wait for the previous user of this slot to finish computing
if (sched->prefetch_used[slot]) {
ggml_backend_event_wait(sched->prefetch_backend, sched->prefetch_free[slot]);
}
// point the staging copy at the slot only for the duration of
// this split, so a fallback to the regular path on a later
// eval can never see a dangling slot pointer
prefetch_input_cpy = input_cpy;
prefetch_saved_buffer = input_cpy->buffer;
prefetch_saved_data = input_cpy->data;
input_cpy->buffer = sched->prefetch_slots[slot];
input_cpy->data = ggml_backend_buffer_get_base(sched->prefetch_slots[slot]);
ggml_backend_tensor_set_async(sched->prefetch_backend, input_cpy, input->data, 0, ggml_nbytes(input));
@@ -1759,12 +1820,15 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
if (!sched->callback_eval) {
enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph);
if (ec != GGML_STATUS_SUCCESS) {
return ec;
}
if (split_prefetch_slot != -1) {
// the kernels have captured the slot address at launch, safe to restore
ggml_backend_event_record(sched->prefetch_free[split_prefetch_slot], split_backend);
sched->prefetch_used[split_prefetch_slot] = true;
prefetch_input_cpy->buffer = prefetch_saved_buffer;
prefetch_input_cpy->data = prefetch_saved_data;
}
if (ec != GGML_STATUS_SUCCESS) {
return ec;
}
} else {
// similar to ggml_backend_compare_graph_backend
@@ -1875,8 +1939,14 @@ ggml_backend_sched_t ggml_backend_sched_new(
sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends);
sched->op_offload = op_offload;
// GGML_SCHED_PREFETCH_EXPERTS=1 enables the default slot count, higher values set it
// directly; more slots let uploads run further ahead of compute at the cost of one
// max-sized expert tensor of device memory per slot
const char * GGML_SCHED_PREFETCH_EXPERTS = getenv("GGML_SCHED_PREFETCH_EXPERTS");
sched->prefetch_experts = op_offload && GGML_SCHED_PREFETCH_EXPERTS && atoi(GGML_SCHED_PREFETCH_EXPERTS);
const int prefetch_n_slots = GGML_SCHED_PREFETCH_EXPERTS ? atoi(GGML_SCHED_PREFETCH_EXPERTS) : 0;
sched->prefetch_experts = op_offload && prefetch_n_slots > 0;
// default of 3 covers the gate/up/down expert tensors of one MoE layer
sched->prefetch_n_slots = prefetch_n_slots <= 1 ? 3 : std::min(prefetch_n_slots, GGML_SCHED_MAX_PREFETCH_SLOTS);
ggml_backend_sched_reset(sched);
@@ -1894,7 +1964,8 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
}
if (sched->prefetch_backend) {
ggml_backend_synchronize(sched->prefetch_backend);
for (int i = 0; i < 2; i++) {
// the slot count may have been reduced after a failed allocation, free everything
for (int i = 0; i < GGML_SCHED_MAX_PREFETCH_SLOTS; i++) {
ggml_backend_event_free(sched->prefetch_ready[i]);
ggml_backend_event_free(sched->prefetch_free[i]);
ggml_backend_buffer_free(sched->prefetch_slots[i]);