llama : overlap the AllReduce of split mode tensor with the compute of the next ubatch

common/arg.cpp

@@ -1769,6 +1769,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.flash_attn_native_quants = value;
         }
     ).set_env("LLAMA_ARG_FLASH_ATTN_NATIVE_QUANTS"));
+    add_opt(common_arg(
+        {"--split-overlap"},
+        {"--no-split-overlap"},
+        "with --split-mode tensor: alternate ubatches of at least 64 tokens between two schedulers, so that the "
+        "AllReduce of one overlaps with the compute of the other. Each scheduler has its own compute buffers: "
+        "half the --ubatch-size takes the same memory and is usually the fastest choice (default: disabled)",
+        [](common_params & params, bool value) {
+            params.split_overlap = value;
+        }
+    ).set_env("LLAMA_ARG_SPLIT_OVERLAP"));
     add_opt(common_arg(
         {"--cache-idle-slots"},
         {"--no-cache-idle-slots"},

common/common.cpp

@@ -1742,6 +1742,7 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.swa_full          = params.swa_full;
     cparams.kv_unified        = params.kv_unified;
     cparams.flash_attn_native_quants = params.flash_attn_native_quants;
+    cparams.split_overlap            = params.split_overlap;
 
     cparams.type_k = params.cache_type_k;
     cparams.type_v = params.cache_type_v;

common/common.h

@@ -507,6 +507,7 @@ struct common_params {
     enum llama_flash_attn_type   flash_attn_type   = LLAMA_FLASH_ATTN_TYPE_AUTO; // whether to use Flash Attention
 
     bool flash_attn_native_quants = false; // read a quantized K/V cache in place instead of casting it to F16
+    bool split_overlap            = false; // split mode tensor: overlap the AllReduce of a ubatch with the compute of the next
 
     struct common_params_sampling    sampling;
     struct common_params_speculative speculative;

docs/multi-gpu.md

@@ -37,6 +37,7 @@ Set with `--split-mode` / `-sm`.
 | `-sm` | `--split-mode` | `none` \| `layer` \| `tensor` | `layer` | See modes above. |
 | `-ts` | `--tensor-split` | comma-separated proportions, e.g. `3,1` | mode-dependent | How much of the model goes to each GPU. If omitted, `layer`/`row` use automatic splitting proportional to memory, while `tensor` splits tensor segments evenly. With `3,1` on two GPUs, GPU 0 gets 75 %, GPU 1 gets 25 %. The values follow the order in `--device`. |
 | `-as` | `--attn-split` | comma-separated proportions, e.g. `1,3` | follows `-ts` | With `--split-mode tensor`: the share of the attention heads per GPU. The heads take their weights and their KV cache along, the rest of the model follows `--tensor-split`. With GPUs of different speed, give the faster one more of `--tensor-split` and move KV heads to the other one to keep both in memory. The model load logs the KV heads each GPU got. |
+| | `--split-overlap` | - | off | With `--split-mode tensor`: alternate ubatches between two schedulers, so that the AllReduce of one overlaps with the compute of the other. See recipe 8. |
 | `-mg` | `--main-gpu` | integer device index | `0` | The single GPU used in `--split-mode none`. |
 | `-ngl` | `--n-gpu-layers` / `--gpu-layers` | integer \| `auto` \| `all` | `auto` | Maximum number of layers to keep in VRAM. Use `999` or `all` to push everything possible to the GPUs. |
 | `-dev` | `--device` | comma-separated device names, or `none` | auto | Restrict which devices llama.cpp may use. See `--list-devices` for names. |
@@ -135,6 +136,33 @@ Environment variables:
 
 The copy engines need 16 MiB of VRAM and 8 MiB of pinned memory per GPU and per context that has large reductions. They are allocated with the first large reduction. An MTP draft context processes the prompt batches as well, so it has its own buffers: 32 MiB of VRAM per GPU in total.
 
+### 8. Overlapping the AllReduce with compute (`--split-overlap`)
+
+```bash
+llama-server -m model.gguf -sm tensor -ub 256 --split-overlap
+```
+
+In `tensor` mode every layer ends with an AllReduce, and the GPUs wait for it. On a slow link this can take a large share of prompt processing: about 40% with an RTX 3060 on PCIe 3.0 x4.
+
+With `--split-overlap`, consecutive ubatches of at least 64 tokens alternate between two schedulers, also across calls. Each scheduler has its own backends, streams and compute buffers. A subgraph waits on each GPU for the subgraphs of the other ubatch that last read or wrote the same state tensors (KV cache, recurrent state). A subgraph without state, such as an FFN, does not wait. So while one ubatch transfers a layer's result, the other computes.
+
+It needs all layers, the KV cache and the experts on the GPUs (the CPU parts of a graph would run outside that order), and devices with events.
+
+The output is the same as without the option. Decode (small ubatches) is not affected.
+
+Notes:
+
+- Each scheduler reserves the compute buffers for a full ubatch. With half the `--ubatch-size`, the memory stays the same, and two ubatches are in flight at once. On the GPUs below, 256 per scheduler was faster than 128, 384 or 512.
+- Each scheduler captures its own CUDA graphs, which takes some more VRAM at run time (about 50 MiB here).
+- The overlap works best with the copy-engine AllReduce of recipe 7: NCCL kernels hold an SM each, and the compute of the other ubatch then runs slower.
+
+Qwen3.8-27B on an RTX 4070 (PCIe 4.0 x16) + RTX 3060 (PCIe 3.0 x4), `-ts 0.625,0.375 -as 1,3`, 262144 tokens of q5_1 cache in 2 slots, with MTP, prompt processing in t/s:
+
+| prompt | `-ub 512` | `-ub 256 --split-overlap` |
+|---|--:|--:|
+| 10k tokens | 733 | 947 |
+| 93k tokens | 570 | 734 |
+
 ---
 
 ## Troubleshooting

ggml/include/ggml-backend.h

@@ -430,6 +430,12 @@ extern "C" {
     GGML_API ggml_backend_dev_t ggml_backend_meta_device(
         ggml_backend_dev_t * devs, size_t n_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud);
 
+    // let two backends of the same meta device compute consecutive graphs as lanes: each lane runs on its own streams,
+    // and a subgraph waits for the subgraphs of the other lane that last touched the same state tensors (KV cache,
+    // recurrent state), so the AllReduce of one lane overlaps with the compute of the other. The graphs must not
+    // share intermediate buffers (one scheduler per lane). Returns false if a device has no events.
+    GGML_API bool ggml_backend_meta_share_lanes(ggml_backend_t meta_backend_a, ggml_backend_t meta_backend_b);
+
     //
     // Utils
     //

ggml/src/ggml-backend-meta.cpp

@@ -6,6 +6,7 @@
 #include "ggml-cpp.h"
 
 #include <algorithm>
+#include <unordered_map>
 #include <cassert>
 #include <cmath>
 #include <cstddef>
@@ -1836,6 +1837,20 @@ struct ggml_backend_meta_context {
     void *                               comm_ctx       = nullptr;
     ggml_backend_comm_allreduce_tensor_t comm_allreduce = nullptr;
 
+    // Lanes: meta backends of one device that compute consecutive graphs on their own streams, so that the AllReduce
+    // of one graph overlaps with the compute of the other, see ggml_backend_meta_share_lanes().
+    struct lane_group {
+        ggml_backend_meta_context * ctx[2] = { nullptr, nullptr };
+    };
+    std::shared_ptr<lane_group>                    lanes;
+    std::vector<std::vector<ggml_backend_event_t>> lane_events; // [subgraph][simple backend], recorded once it is computed
+
+    // the state tensors (KV cache, recurrent state: tensors of buffers that are neither weights nor compute buffers)
+    // that each subgraph of the current graph reads or writes, by the tensor at the root of their views
+    std::vector<std::vector<const ggml_tensor *>> lane_touch;
+    // the subgraph (index into lane_events) of this lane that last touched a state tensor, over all its graphs
+    std::unordered_map<const ggml_tensor *, size_t> lane_last_touch;
+
     ggml_backend_meta_context(ggml_backend_dev_t meta_dev, const char * params) {
         const size_t n_devs = ggml_backend_meta_dev_n_devs(meta_dev);
         n_reduce_steps = std::ceil(std::log2(n_devs));
@@ -1870,6 +1885,18 @@ struct ggml_backend_meta_context {
     }
 
     ~ggml_backend_meta_context() {
+        for (auto & events : lane_events) {
+            for (ggml_backend_event_t event : events) {
+                ggml_backend_event_free(event);
+            }
+        }
+        if (lanes) {
+            for (auto & ctx : lanes->ctx) {
+                if (ctx == this) {
+                    ctx = nullptr;
+                }
+            }
+        }
         if (comm_ctx != nullptr) {
             ggml_backend_comm_free_t comm_free = (ggml_backend_comm_free_t) ggml_backend_reg_get_proc_address(
                 ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_configs[0].backend)), "ggml_backend_comm_free");
@@ -2259,6 +2286,41 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend,
         backend_ctx->uid         = cgraph->uid;
         backend_ctx->n_subgraphs = n_subgraphs;
 
+        if (backend_ctx->lanes) {
+            auto state_root = [](const ggml_tensor * t) -> const ggml_tensor * {
+                while (t->view_src != nullptr) {
+                    t = t->view_src;
+                }
+                if (t->buffer == nullptr || !ggml_backend_buffer_is_meta(t->buffer)) {
+                    return nullptr;
+                }
+                const enum ggml_backend_buffer_usage usage = ggml_backend_buffer_get_usage(t->buffer);
+                return usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS || usage == GGML_BACKEND_BUFFER_USAGE_COMPUTE ? nullptr : t;
+            };
+            backend_ctx->lane_touch.assign(n_subgraphs, {});
+            for (size_t i_graph = 0; i_graph < n_subgraphs; i_graph++) {
+                auto & touch = backend_ctx->lane_touch[i_graph];
+                const auto & cfg = backend_ctx->backend_configs[0].cgraphs;
+                const int i_start_g = cfg[i_graph].offset;
+                const int i_stop_g  = i_graph + 1 < n_subgraphs ? cfg[i_graph + 1].offset : cgraph->n_nodes;
+                for (int i_node = i_start_g; i_node < i_stop_g; i_node++) {
+                    const ggml_tensor * node = cgraph->nodes[i_node];
+                    if (const ggml_tensor * root = state_root(node)) {
+                        touch.push_back(root);
+                    }
+                    for (int k = 0; k < GGML_MAX_SRC; k++) {
+                        if (node->src[k] != nullptr) {
+                            if (const ggml_tensor * root = state_root(node->src[k])) {
+                                touch.push_back(root);
+                            }
+                        }
+                    }
+                }
+                std::sort(touch.begin(), touch.end());
+                touch.erase(std::unique(touch.begin(), touch.end()), touch.end());
+            }
+        }
+
         if (max_tmp_size > backend_ctx->max_tmp_size) {
             for (size_t j = 0; j < n_backends; j++) {
                 auto & bcj = backend_ctx->backend_configs[j];
@@ -2466,13 +2528,56 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend,
     };
 
 
+    // With lanes, a subgraph waits on each device for the subgraphs of the other lane that last read or wrote the
+    // state tensors it reads or writes (KV cache, recurrent state), whichever graph and subgraph index that was:
+    // a layer may read the KV of another layer (shared KV), or read its state a subgraph before it writes it. The
+    // AllReduce after a subgraph is not waited for, it overlaps with the compute of the other lane. An event that
+    // the other lane recorded again since is later in its stream, so waiting for it waits for at least as much.
+    ggml_backend_meta_context * lane_other = nullptr;
+    if (backend_ctx->lanes) {
+        lane_other = backend_ctx->lanes->ctx[0] == backend_ctx ? backend_ctx->lanes->ctx[1] : backend_ctx->lanes->ctx[0];
+
+        while (backend_ctx->lane_events.size() < backend_ctx->n_subgraphs) {
+            std::vector<ggml_backend_event_t> events(n_backends);
+            for (size_t j = 0; j < n_backends; j++) {
+                events[j] = ggml_backend_event_new(ggml_backend_get_device(backend_ctx->backend_configs[j].backend));
+                GGML_ASSERT(events[j] != nullptr);
+            }
+            backend_ctx->lane_events.push_back(std::move(events));
+        }
+    }
+    std::vector<size_t> lane_waits;
+
     for (size_t i = 0; i < backend_ctx->n_subgraphs; i++) {
+        lane_waits.clear();
+        if (lane_other != nullptr) {
+            for (const ggml_tensor * t : backend_ctx->lane_touch[i]) {
+                auto it = lane_other->lane_last_touch.find(t);
+                if (it != lane_other->lane_last_touch.end()) {
+                    lane_waits.push_back(it->second);
+                }
+            }
+            std::sort(lane_waits.begin(), lane_waits.end());
+            lane_waits.erase(std::unique(lane_waits.begin(), lane_waits.end()), lane_waits.end());
+        }
+
         for (size_t j = 0; j < n_backends; j++) {
             auto & bcj = backend_ctx->backend_configs[j];
+            for (const size_t i_other : lane_waits) {
+                ggml_backend_event_wait(bcj.backend, lane_other->lane_events[i_other][j]);
+            }
             const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, bcj.cgraphs[i].cgraph_main);
             if (status != GGML_STATUS_SUCCESS) {
                 return status;
             }
+            if (backend_ctx->lanes) {
+                ggml_backend_event_record(backend_ctx->lane_events[i][j], bcj.backend);
+            }
+        }
+        if (backend_ctx->lanes) {
+            for (const ggml_tensor * t : backend_ctx->lane_touch[i]) {
+                backend_ctx->lane_last_touch[t] = i;
+            }
         }
 
         if (n_backends > 1 && i < backend_ctx->n_subgraphs - 1) {
@@ -2539,6 +2644,29 @@ size_t ggml_backend_meta_n_backends(ggml_backend_t meta_backend) {
     return backend_ctx->backend_configs.size();
 }
 
+bool ggml_backend_meta_share_lanes(ggml_backend_t meta_backend_a, ggml_backend_t meta_backend_b) {
+    GGML_ASSERT(ggml_backend_is_meta(meta_backend_a) && ggml_backend_is_meta(meta_backend_b));
+    GGML_ASSERT(meta_backend_a->device == meta_backend_b->device);
+    auto * ctx_a = (ggml_backend_meta_context *) meta_backend_a->context;
+    auto * ctx_b = (ggml_backend_meta_context *) meta_backend_b->context;
+    GGML_ASSERT(!ctx_a->lanes && !ctx_b->lanes);
+
+    // the lanes are ordered with events on each device
+    for (const auto & bc : ctx_a->backend_configs) {
+        ggml_backend_dev_props props;
+        ggml_backend_dev_get_props(ggml_backend_get_device(bc.backend), &props);
+        if (!props.caps.events) {
+            return false;
+        }
+    }
+
+    ctx_a->lanes = std::make_shared<ggml_backend_meta_context::lane_group>();
+    ctx_b->lanes = ctx_a->lanes;
+    ctx_a->lanes->ctx[0] = ctx_a;
+    ctx_a->lanes->ctx[1] = ctx_b;
+    return true;
+}
+
 ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, size_t index) {
     GGML_ASSERT(ggml_backend_is_meta(meta_backend));
     const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context;

include/llama.h

@@ -422,6 +422,9 @@ extern "C" {
         bool flash_attn_native_quants; // let Flash Attention read a quantized K/V cache in place instead of
                                        // casting it to F16 first; removes the transient F16 copy of the
                                        // attention window. Where no native loader exists the F16 path is used.
+        bool split_overlap;            // split mode tensor: alternate large ubatches between two schedulers, so that
+                                       // the AllReduce of one overlaps with the compute of the other; each has its own
+                                       // compute buffers, so a ubatch of half the size takes the same memory
 
         // [EXPERIMENTAL]
         // backend sampler chain configuration (make sure the caller keeps the sampler chains alive)