From 65543471622cdb9f8b8b681782552325af190f0f Mon Sep 17 00:00:00 2001 From: piggidragon Date: Mon, 14 Sep 2026 19:45:28 +0200 Subject: [PATCH] ggml-meta, sched : allocate a meta transport ring at each device's share A meta buffer sizes every simple buffer for the whole tensor, so under split mode tensor each device allocated the full ring although it holds only its share of the heads. ggml_backend_meta_get_shares returns, per device, the largest share of a set of compute leaves, and ggml_backend_meta_alloc_buffer_shares allocates each simple buffer at that share, placing a tensor at X*share rounded up to the alignment. The share includes one alignment per entry, so rounded neighbours do not overlap, and binding asserts that a tensor fits. The scheduler applies the budget to the largest share and the headroom check to each device's own share, and grows the ring when a plan needs a larger share than it holds. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 23 +++-- ggml/src/ggml-backend-impl.h | 7 ++ ggml/src/ggml-backend-meta.cpp | 146 +++++++++++++++++++++++++++----- ggml/src/ggml-backend.cpp | 143 ++++++++++++++++++++++--------- 4 files changed, 254 insertions(+), 65 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index c39df7df6a66..0695b239650c 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -312,7 +312,7 @@ A device-resident KV run is unaffected, and was measured to confirm it: 38.5612 - **One ring per accelerator.** A layer-split model pipelines on every device that qualifies; a device with no room within the budget falls back to the ordered path on its own without disabling the others. The exception is a graph that cannot be allocated next to the rings: there every device that was holding one gives it back for good, because the allocator does not say which of them it competed with. - **The producer of a staged input must be the CPU or the consumer itself.** Neither part of a staged delivery is ordered against a third device: the stable prefix goes on the transfer stream and the rest on the consumer's own stream, where the ordered path would have synchronized the producer first. An input a second accelerator writes keeps the ordered path. - **It turns graph-level pipeline parallelism off while it is delivering.** A graph that delivered has to block the host on its consumer before the next graph writes the host cache, because the host source of a delivery is read long after the call that issued it returned. That block is what `n_copies > 1` exists to avoid, so the two do not overlap: with `-sm layer` over several GPUs and `--kv-cpu-pinned`, `llama_context` enables both and the ring wins. Use `--kv-pipeline-depth 0` to keep the graph-level pipelining instead. -- **Tensor parallelism pipelines through the meta backend**, and each device allocates the whole ring. See [Tensor parallelism](#tensor-parallelism). +- **Tensor parallelism pipelines through the meta backend**, and each device allocates its share of the ring. See [Tensor parallelism](#tensor-parallelism). - **A host write to the cache waits for the delivery.** `llama_memory_clear(mem, true)` waits for the scheduler before it clears the buffers, because a delivery the last decode issued can still be reading them. This was already needed without the transport: with a device-resident cache the same call cleared the buffers under the running graph, and `llama_decode` followed by that clear changed the logits of that decode on every trial. - The scheduler must be configured with the device's own default buffer type. A scheduler built on a split or host buffer type keeps the ordered path. - `GGML_KV_PIPELINE_DEPTH` and `GGML_KV_PIPELINE_BUDGET_MIB` set the defaults of a scheduler that nothing else configures. `llama_context` always configures its own from the context parameters, so under `llama-server` and `llama-bench` use `--kv-pipeline-depth` / `LLAMA_ARG_KV_PIPELINE_DEPTH` and `--kv-pipeline-budget` / `LLAMA_ARG_KV_PIPELINE_BUDGET` instead. @@ -330,7 +330,20 @@ What the meta backend adds: - **A transfer backend without a communicator.** The transfer backend never computes, so `ggml_backend_meta_init_transfer` builds its streams without starting a second NCCL context. - **The ring is a meta buffer.** A copy in it gets its per-device tensors from the same split-state callback as the copy the graph allocator would have made. The meta graph compute also rotates the compute containers of buffers that appear only as sources, or the ring's per-device tensors would accumulate one set per plan. -**Each device allocates the whole slot.** A meta buffer places a tensor at the same offset in every device's buffer and sizes each of those buffers for the whole tensor, so a device that holds half the heads still allocates the full slot. The meta compute buffers already work this way. The budget and the headroom check are applied per device, against the device with the least free memory, so a ring that fits the budget costs that much on every device. +**Each device allocates its share of the ring.** A meta buffer normally sizes every device's buffer for the whole tensor and places a tensor at the same offset in each, so a device that holds half the heads would allocate the full slot; the meta compute buffers still work this way. The ring is allocated with `ggml_backend_meta_alloc_buffer_shares` instead: + +- `ggml_backend_meta_get_shares` asks the split state of each staged copy and returns, per device, the largest fraction of an entry that device holds, in units of 1/65536. +- Device `j` allocates that share of the ring, and an entry at meta offset `X` lands at `X * share_j` there, rounded up to the device's alignment. The share includes one alignment per entry, so two neighbouring entries never overlap after rounding; binding a tensor asserts that it fits. +- The budget caps the ring on the device with the largest share, and the headroom check is applied on each device against its own share. +- A plan whose entries need a larger share than the ring was allocated with grows the ring, the same as one that needs a larger slot. + +A ring that holds a mirrored entry has a share of 1 on every device, which is the same as a plain meta buffer. The share is the largest over the entries, so a layer that puts all of its KV heads on one device costs that device the whole slot: gemma-4 at `-ts 55,45` allocates 8,192 KiB per slot on the first device and 4,097 KiB on the second. + +With `-ts 50,50` on Qwen3.8-27B-UD-IQ2_M each device allocates 4,097 KiB of an 8,192 KiB slot. That moves where the default budget declines. At 32,768 a slot is 70.7 MiB, so the ring is 212 MiB on each device if every device holds all of it, over the 128 MiB default, and 103.6 MiB per device at its share, under it. `llama-bench` at the default budget, one pass each: + +| depth | ordered | pipelined | ring per device | +|---:|---:|---:|---:| +| 32,768 | 4.07 | 4.88 | 3 x 34.5 MiB | ### Measurements @@ -361,7 +374,7 @@ Per decode graph at 16,384, `GGML_SCHED_TRANSPORT_DEBUG=2`: | blocked in the ordered copy | 103.81 ms | 12.16 ms | | blocked waiting for the consumer | 31.06 ms | 95.43 ms | | bytes delivered early / late | 0 / 0 MiB | 549.3 / 3.2 MiB | -| ring | - | 3 slots x 35 MiB, per device | +| ring | - | 3 slots x 35 MiB | The copy is three times the compute here, where on the single RTX 4070 above the two were about equal. The 553 MiB cross in 104 ms, about 5.3 GB/s, and the 3060's half of them crosses a gen3 x4 link. The pipeline hides the compute behind the copy, and the consumer wait now contains the rest of the transfer, so the token is bounded by the slower link rather than by the order of the work. The ceiling is `max(copy, compute)` plus the work outside the split loop, the same as on one device, and at 16,384 the pipeline is within a few milliseconds of it. @@ -376,7 +389,3 @@ On the same two devices: - Greedy `llama-completion`, 64 tokens behind a 3k prompt: Qwen3.8-27B-UD-IQ2_M gives `64e86551f7ef1638` with a device-resident cache and at `N = 0`, `1` and `4` with a host one. gemma-4-26B-A4B gives `5525e3f5ac7337d7` at `N = 0` and `1` with `-ts 50,50`, and `4f8986fb3655a567` at both with `-ts 55,45`. - `test-llama-archs` adds a `Meta -nkvo -np 2 -kvpd 1` configuration, which stages the cache on the meta ring and delivers it through the ranged head-split write. It passes on 2, 3 and 4 CUDA devices. It evaluates one ubatch, so it delivers only the late part. - `test-alloc` passes. Its meta test now covers a device of the meta type that is not the ggml meta backend, which stays ordered. - -## Future work - -- Allocate each device's part of a meta ring at its own share of the heads, instead of the whole slot on every device. diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 27c7340a605a..29aeb315172c 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -104,6 +104,13 @@ extern "C" { // a meta backend without a communicator, for moving data on streams of its own: a graph it computes reduces through copies GGML_API ggml_backend_t ggml_backend_meta_init_transfer(ggml_backend_dev_t meta_dev); + // A meta buffer whose simple buffer j holds shares[j]/65536 of `size`: a tensor at meta offset X lands at X*shares[j] there, rounded up to the alignment. + // Every tensor placed in it must take at most that share on each device, which ggml_backend_meta_get_shares returns for a set of tensors. + GGML_API ggml_backend_buffer_t ggml_backend_meta_alloc_buffer_shares(ggml_backend_buffer_type_t buft, size_t size, const uint32_t * shares); + + // for each simple buffer type of buft, the smallest share that holds every one of these compute leaves, 65536 for all of it + GGML_API void ggml_backend_meta_get_shares(ggml_backend_buffer_type_t buft, const struct ggml_tensor * const * tensors, size_t n_tensors, uint32_t * shares); + // temporary workaround to statically allocate tensors from a context in a deduplicated way: GGML_API struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft); diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index c192c8a966cb..6b4411b04edd 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -468,6 +468,10 @@ struct ggml_backend_meta_buffer_context { int stc_compute_index_next = 0; std::vector bufs; + // Share of the meta size each simple buffer holds, in units of 1/65536, empty when every simple buffer holds all of it. + // A tensor at meta offset X lands at X*share on simple buffer j, see ggml_backend_meta_alloc_buffer_shares. + std::vector shares; + // FIXME // The size of the split state cache is unbounded and can theoretically grow infinitely large. // However, it is also expensive to build and clearing it on every rebuild in ggml_backend_meta_graph_compute is too expensive. @@ -1241,6 +1245,38 @@ static void * ggml_backend_meta_buffer_get_base(ggml_backend_buffer_t buffer) { return (void *) 0x1000000000000000; // FIXME } +// ne and nb of the part of a tensor on simple buffer j +static void ggml_backend_meta_simple_shape(const ggml_tensor * tensor, const ggml_backend_meta_split_state & split_state, + size_t n_simple_bufs, size_t j, int64_t * ne, size_t * nb) { + for (size_t k = 0; k < GGML_MAX_DIMS; k++) { + ne[k] = tensor->ne[k]; + nb[k] = tensor->nb[k]; + } + const int split_dim = split_state.axis; + if (split_dim >= 0 && split_dim < GGML_MAX_DIMS) { + // TODO: the following assert fails for llama-parallel even though the results are correct: + // GGML_ASSERT(ggml_is_contiguously_allocated(tensor)); + ne[split_dim] = 0; + for (size_t s = 0; s < split_state.n_segments; s++) { + ne[split_dim] += split_state.ne[s*n_simple_bufs + j] * split_state.nr[s]; + } + for (int i = 0; i < GGML_MAX_DIMS; i++) { + if (tensor->nb[i] > tensor->nb[split_dim]) { + nb[i] = tensor->nb[i] * ne[split_dim]/tensor->ne[split_dim]; + } + } + } +} + +// offset on simple buffer j of what sits at meta offset `offset` +static size_t ggml_backend_meta_simple_offset(const ggml_backend_meta_buffer_context * buf_ctx, size_t j, size_t offset) { + if (buf_ctx->shares.empty() || buf_ctx->shares[j] == 65536) { + return offset; + } + const size_t alignment = ggml_backend_buffer_get_alignment(buf_ctx->bufs[j].get()); + return GGML_PAD((size_t) (((uint64_t) offset * buf_ctx->shares[j]) >> 16), alignment); +} + static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_meta_simple_tensor_container & stc, ggml_tensor * tensor) { GGML_ASSERT(ggml_backend_buffer_is_meta(tensor->buffer)); ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; @@ -1253,10 +1289,6 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m int split_dim = split_state.axis; int64_t ne[GGML_MAX_DIMS]; size_t nb[GGML_MAX_DIMS]; - for (size_t k = 0; k < GGML_MAX_DIMS; k++) { - ne[k] = tensor->ne[k]; - nb[k] = tensor->nb[k]; - } std::vector simple_tensors; simple_tensors.reserve(n_simple_bufs); @@ -1269,19 +1301,7 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m GGML_ABORT("multi buffers are not supported by the meta backend"); } - if (split_dim >= 0 && split_dim < GGML_MAX_DIMS) { - // TODO: the following assert fails for llama-parallel even though the results are correct: - // GGML_ASSERT(ggml_is_contiguously_allocated(tensor)); - ne[split_dim] = 0; - for (size_t s = 0; s < split_state.n_segments; s++) { - ne[split_dim] += split_state.ne[s*n_simple_bufs + j] * split_state.nr[s]; - } - for (int i = 0; i < GGML_MAX_DIMS; i++) { - if (tensor->nb[i] > tensor->nb[split_dim]) { - nb[i] = tensor->nb[i] * ne[split_dim]/tensor->ne[split_dim]; - } - } - } + ggml_backend_meta_simple_shape(tensor, split_state, n_simple_bufs, j, ne, nb); ggml_tensor * t_ij = ggml_new_tensor(simple_ctx, tensor->type, GGML_MAX_DIMS, ne); t_ij->op = tensor->op; @@ -1319,8 +1339,11 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m if (t_ij->view_src != nullptr) { t_ij->data = (char *) t_ij->view_src->data + t_ij->view_offs; } else if (simple_buf != nullptr) { - t_ij->data = (char *) ggml_backend_buffer_get_base(simple_buf) - + size_t(tensor->data) - size_t(ggml_backend_buffer_get_base(tensor->buffer)); + const size_t offset = ggml_backend_meta_simple_offset(buf_ctx, j, + size_t(tensor->data) - size_t(ggml_backend_buffer_get_base(tensor->buffer))); + GGML_ASSERT(buf_ctx->shares.empty() || + offset + ggml_backend_buffer_get_alloc_size(simple_buf, t_ij) <= ggml_backend_buffer_get_size(simple_buf)); + t_ij->data = (char *) ggml_backend_buffer_get_base(simple_buf) + offset; } if (simple_buf) { @@ -1858,6 +1881,91 @@ static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_bac return ggml_backend_buffer_init(buft, ggml_backend_meta_buffer_iface, buf_ctx, max_size); } +ggml_backend_buffer_t ggml_backend_meta_alloc_buffer_shares(ggml_backend_buffer_type_t buft, size_t size, const uint32_t * shares) { + GGML_ASSERT(ggml_backend_buft_is_meta(buft)); + const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); + + const ggml_init_params params = { + /*.mem_size =*/ 1024*1024*ggml_tensor_overhead(), // FIXME + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_backend_meta_simple_tensor_container stc_static; + ggml_backend_meta_simple_tensor_container stc_compute_0(params, n_simple_bufts); + ggml_backend_meta_simple_tensor_container stc_compute_1(params, n_simple_bufts); + + std::vector bufs; + bufs.reserve(n_simple_bufts); + for (size_t i = 0; i < n_simple_bufts; i++) { + GGML_ASSERT(shares[i] > 0 && shares[i] <= 65536); + const size_t size_i = (size_t) (((uint64_t) size * shares[i] + 65535) >> 16); + bufs.push_back(ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size_i)); + if (bufs.back() == nullptr) { + for (ggml_backend_buffer_t buf : bufs) { + ggml_backend_buffer_free(buf); + } + return nullptr; + } + } + ggml_backend_meta_buffer_context * buf_ctx = new ggml_backend_meta_buffer_context(stc_static, stc_compute_0, stc_compute_1, bufs); + buf_ctx->shares.assign(shares, shares + n_simple_bufts); + + // the meta size stays the one asked for: tensors are placed against it, and each simple buffer holds its share + return ggml_backend_buffer_init(buft, ggml_backend_meta_buffer_iface, buf_ctx, size); +} + +void ggml_backend_meta_get_shares(ggml_backend_buffer_type_t buft, const struct ggml_tensor * const * tensors, size_t n_tensors, uint32_t * shares) { + GGML_ASSERT(ggml_backend_buft_is_meta(buft)); + const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); + + // the split state is asked of a compute leaf, so ask it through an empty compute buffer of this type + ggml_backend_meta_simple_tensor_container stc_static; + ggml_backend_meta_simple_tensor_container stc_compute_0; + ggml_backend_meta_simple_tensor_container stc_compute_1; + std::vector bufs(n_simple_bufts, nullptr); + ggml_backend_meta_buffer_context * buf_ctx = new ggml_backend_meta_buffer_context(stc_static, stc_compute_0, stc_compute_1, bufs); + ggml_backend_buffer_t probe = ggml_backend_buffer_init(buft, ggml_backend_meta_buffer_iface, buf_ctx, 0); + probe->usage = GGML_BACKEND_BUFFER_USAGE_COMPUTE; + + for (size_t j = 0; j < n_simple_bufts; j++) { + shares[j] = 1; + } + + for (size_t i = 0; i < n_tensors; i++) { + ggml_tensor t = *tensors[i]; + GGML_ASSERT(t.view_src == nullptr); + t.buffer = probe; + t.data = ggml_backend_buffer_get_base(probe); + + const size_t alloc_size = ggml_backend_buft_get_alloc_size(buft, &t); + if (alloc_size == 0) { + continue; + } + + const ggml_backend_meta_split_state split_state = + ggml_backend_meta_get_split_state(buf_ctx->stc_compute[0], &t, /*assume_sync =*/ true); + for (size_t j = 0; j < n_simple_bufts; j++) { + int64_t ne[GGML_MAX_DIMS]; + size_t nb[GGML_MAX_DIMS]; + ggml_backend_meta_simple_shape(&t, split_state, n_simple_bufts, j, ne, nb); + + ggml_tensor t_j = t; + for (int k = 0; k < GGML_MAX_DIMS; k++) { + t_j.ne[k] = ne[k]; + t_j.nb[k] = nb[k]; + } + ggml_backend_buffer_type_t simple_buft = ggml_backend_meta_buft_simple_buft(buft, j); + + // with the alignment in it, the next tensor never starts inside this one after both offsets are rounded, see ggml_backend_meta_simple_offset + const uint64_t need = ggml_backend_buft_get_alloc_size(simple_buft, &t_j) + ggml_backend_buft_get_alignment(simple_buft); + const uint64_t share = (need*65536 + alloc_size - 1) / alloc_size; + shares[j] = (uint32_t) std::min(std::max(shares[j], share), 65536); + } + } + + ggml_backend_buffer_free(probe); +} + struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft) { const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index acd1ab108dbe..5a71a6ddbe24 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -817,6 +817,7 @@ struct ggml_backend_sched_transport_ring { ggml_backend_buffer_t buffer; // the ring itself size_t slot_size; size_t alignment; + uint32_t shares[GGML_BACKEND_META_MAX_DEVICES]; // share of the ring each device of the buffer holds, in units of 1/65536 struct ggml_backend_sched_transport_slot slots[GGML_SCHED_MAX_TRANSPORT_SLOTS]; @@ -2165,24 +2166,57 @@ static size_t ggml_backend_sched_transport_slot_alloc(size_t need, size_t limit, return std::max(size, need); } -// Free memory of the device a ring lands on, 0 when unknown. -// A meta buffer allocates the whole ring on every simple device, so the device with the least free memory decides; the meta device reports the sum. -static size_t ggml_backend_sched_transport_dev_free(ggml_backend_t backend) { - size_t dev_free = 0, dev_total = 0; - if (ggml_backend_is_meta(backend)) { - for (size_t j = 0; j < ggml_backend_meta_n_backends(backend); j++) { - ggml_backend_dev_t dev = ggml_backend_get_device(ggml_backend_meta_simple_backend(backend, j)); - size_t simple_free = 0, simple_total = 0; - ggml_backend_dev_memory(dev, &simple_free, &simple_total); - dev_free = j == 0 ? simple_free : std::min(dev_free, simple_free); +// A share of a size, in units of 1/65536, rounded up. +static size_t ggml_backend_sched_share_size(size_t size, uint32_t share) { + return size/65536*share + ((size%65536)*share + 65535)/65536; +} + +// The largest size whose share is at most `bytes`. +static size_t ggml_backend_sched_share_limit(size_t bytes, uint32_t share) { + GGML_ASSERT(share > 0); + const size_t q = bytes/share; + return q > SIZE_MAX/65536 ? SIZE_MAX : q*65536 + (bytes%share)*65536/share; +} + +// The devices a ring lands on, and the share of the ring each of them holds: a meta backend holds only each device's part of its entries. +// Returns how many devices there are. +static int ggml_backend_sched_transport_ring_shares(ggml_backend_sched_t sched, int backend_id, uint32_t * shares) { + const struct ggml_backend_sched_transport * tr = &sched->transport; + ggml_backend_t backend = sched->backends[backend_id]; + + if (!ggml_backend_is_meta(backend)) { + shares[0] = 65536; + return 1; + } + + std::vector entries; + for (int i = 0; i < sched->n_splits; i++) { + const struct ggml_backend_sched_split * split = &sched->splits[i]; + if (split->backend_id != backend_id || tr->split_order[i] < 0) { + continue; + } + for (int j = 0; j < split->n_inputs; j++) { + if (tr->input_staged[tr->split_input_ofs[i] + j]) { + entries.push_back(tensor_copy(split->inputs[j], backend_id, sched->cur_copy)); + } } - return dev_free; } - ggml_backend_dev_t dev = ggml_backend_get_device(backend); - if (dev != NULL) { - ggml_backend_dev_memory(dev, &dev_free, &dev_total); + ggml_backend_meta_get_shares(sched->bufts[backend_id], entries.data(), entries.size(), shares); + + return (int) ggml_backend_meta_n_backends(backend); +} + +// Free memory of each device a ring lands on, 0 where unknown. +static void ggml_backend_sched_transport_dev_free(ggml_backend_t backend, size_t * dev_free, int n_devs) { + for (int j = 0; j < n_devs; j++) { + ggml_backend_t simple = ggml_backend_is_meta(backend) ? ggml_backend_meta_simple_backend(backend, j) : backend; + ggml_backend_dev_t dev = ggml_backend_get_device(simple); + size_t total = 0; + dev_free[j] = 0; + if (dev != NULL) { + ggml_backend_dev_memory(dev, &dev_free[j], &total); + } } - return dev_free; } // Created on demand, so a backend that never gets to stage anything does not carry a second device context for nothing. @@ -2478,13 +2512,22 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { ring_size_max = SIZE_MAX; } - if (tr->budget > 0 && ring_size > tr->budget) { + // the budget caps what the ring costs on the device that holds the most of it + uint32_t shares[GGML_BACKEND_META_MAX_DEVICES]; + const int n_devs = ggml_backend_sched_transport_ring_shares(sched, bid, shares); + uint32_t share_max = 0; + for (int j = 0; j < n_devs; j++) { + share_max = std::max(share_max, shares[j]); + } + const size_t ring_bytes = ggml_backend_sched_share_size(ring_size, share_max); + + if (tr->budget > 0 && ring_bytes > tr->budget) { if (!r->reported_no_room) { GGML_LOG_WARN("%s: transport ring on %s needs %zu MiB now and %zu MiB at the full " "context, against a %zu MiB budget, staying on the ordered path (raise " "--kv-pipeline-budget to spend more device memory on it)\n", __func__, - ggml_backend_name(sched->backends[bid]), ring_size >> 20, - ring_size_max >> 20, tr->budget >> 20); + ggml_backend_name(sched->backends[bid]), ring_bytes >> 20, + ggml_backend_sched_share_size(ring_size_max, share_max) >> 20, tr->budget >> 20); r->reported_no_room = true; } ggml_backend_sched_transport_decline_backend(sched, bid); @@ -2499,41 +2542,59 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { continue; } - if (r->buffer == NULL || r->slot_size < slot_size[bid]) { + bool grow = r->buffer == NULL || r->slot_size < slot_size[bid]; + for (int j = 0; j < n_devs && !grow; j++) { + grow = r->shares[j] < shares[j]; + } + + if (grow) { ggml_backend_sched_transport_free_ring(sched, bid); ggml_backend_buffer_type_t buft = sched->bufts[bid]; + // grow past what this graph needs, but never past the full context, the budget, or the room left on each device + size_t slot_limit = std::min(slot_size_max[bid], SIZE_MAX/tr->n_slots); + if (tr->budget > 0) { + slot_limit = std::min(slot_limit, ggml_backend_sched_share_limit(tr->budget, share_max)/tr->n_slots); + } + // the graph allocator reserved before this, so leave it the room its buffers may still grow into - const size_t dev_free = ggml_backend_sched_transport_dev_free(sched->backends[bid]); - if (dev_free > 0 && (dev_free <= GGML_SCHED_TRANSPORT_HEADROOM || ring_size > dev_free - GGML_SCHED_TRANSPORT_HEADROOM)) { - if (!r->reported_no_room) { - GGML_LOG_WARN("%s: transport ring on %s would need %zu MiB and leave less than " - "%u MiB of the %zu MiB free, staying on the ordered path\n", __func__, - ggml_backend_name(sched->backends[bid]), ring_size >> 20, - GGML_SCHED_TRANSPORT_HEADROOM >> 20, dev_free >> 20); - r->reported_no_room = true; + size_t dev_free[GGML_BACKEND_META_MAX_DEVICES]; + ggml_backend_sched_transport_dev_free(sched->backends[bid], dev_free, n_devs); + bool no_room = false; + for (int j = 0; j < n_devs && !no_room; j++) { + if (dev_free[j] == 0) { + continue; } + const size_t need = ggml_backend_sched_share_size(ring_size, shares[j]); + if (dev_free[j] <= GGML_SCHED_TRANSPORT_HEADROOM || need > dev_free[j] - GGML_SCHED_TRANSPORT_HEADROOM) { + if (!r->reported_no_room) { + GGML_LOG_WARN("%s: transport ring on %s would need %zu MiB on device %d and leave less than " + "%u MiB of the %zu MiB free, staying on the ordered path\n", __func__, + ggml_backend_name(sched->backends[bid]), need >> 20, j, + GGML_SCHED_TRANSPORT_HEADROOM >> 20, dev_free[j] >> 20); + r->reported_no_room = true; + } + no_room = true; + break; + } + slot_limit = std::min(slot_limit, + ggml_backend_sched_share_limit(dev_free[j] - GGML_SCHED_TRANSPORT_HEADROOM, shares[j])/tr->n_slots); + } + if (no_room) { ggml_backend_sched_transport_decline_backend(sched, bid); continue; } - // grow past what this graph needs, but never past the full context, the budget, or what the headroom check just approved - size_t slot_limit = std::min(slot_size_max[bid], SIZE_MAX/tr->n_slots); - if (tr->budget > 0) { - slot_limit = std::min(slot_limit, tr->budget/tr->n_slots); - } - if (dev_free > GGML_SCHED_TRANSPORT_HEADROOM) { - slot_limit = std::min(slot_limit, (dev_free - GGML_SCHED_TRANSPORT_HEADROOM)/tr->n_slots); - } const size_t slot_alloc = ggml_backend_sched_transport_slot_alloc(slot_size[bid], slot_limit, r->alignment); const size_t alloc_size = slot_alloc*tr->n_slots; - ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, alloc_size); + ggml_backend_buffer_t buffer = ggml_backend_is_meta(sched->backends[bid]) ? + ggml_backend_meta_alloc_buffer_shares(buft, alloc_size, shares) : ggml_backend_buft_alloc_buffer(buft, alloc_size); if (buffer == NULL) { // the headroom check passed, so the device is out of memory for reasons this cannot see; a retry per graph costs a context per token GGML_LOG_WARN("%s: failed to allocate %zu MiB for the transport ring on %s, " - "pipelining disabled there\n", __func__, alloc_size >> 20, + "pipelining disabled there\n", __func__, ggml_backend_sched_share_size(alloc_size, share_max) >> 20, ggml_backend_name(sched->backends[bid])); ggml_backend_sched_transport_disable_backend(sched, bid); continue; @@ -2544,10 +2605,14 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { r->buffer = buffer; r->slot_size = slot_alloc; + memcpy(r->shares, shares, n_devs*sizeof(shares[0])); if (tr->debug > 0) { - GGML_LOG_INFO("%s: transport ring on %s: %d slots x %zu KiB\n", __func__, - ggml_backend_name(sched->backends[bid]), tr->n_slots, slot_alloc >> 10); + for (int j = 0; j < n_devs; j++) { + GGML_LOG_INFO("%s: transport ring on %s, device %d: %d slots x %zu KiB\n", __func__, + ggml_backend_name(sched->backends[bid]), j, tr->n_slots, + ggml_backend_sched_share_size(slot_alloc, shares[j]) >> 10); + } } }