ggml-meta, sched : allocate a meta transport ring at each device's share

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.

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);
 

ggml/src/ggml-backend-meta.cpp

@@ -468,6 +468,10 @@ struct ggml_backend_meta_buffer_context {
     int stc_compute_index_next = 0;
     std::vector<ggml_backend_buffer_ptr> 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<uint32_t> 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<ggml_tensor *> 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<ggml_backend_buffer_t> 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<ggml_backend_buffer_t> 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<uint64_t>(std::max<uint64_t>(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);