ggml-meta : split a host-resident KV cache by head

CLAUDE.md

@@ -0,0 +1 @@
+AGENTS.md

ci/run.sh

@@ -312,6 +312,8 @@ function gg_run_test_llama_archs_tensor_split {
         GGML_CUDA_DEVICES=2 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
         GGML_CUDA_DEVICES=3 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
         GGML_CUDA_DEVICES=4 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
+        # the scheduler names a copy after its source, and 8 device names fill the name field
+        GGML_CUDA_DEVICES=8 ./build-ci-release/bin/test-llama-archs -s 1 -a llama 2>&1
     fi
 
     if [ ! -z ${GG_BUILD_METAL} ]; then
@@ -327,7 +329,7 @@ function gg_run_test_llama_archs_tensor_split {
 function gg_sum_test_llama_archs_tensor_split {
     gg_printf '### %s\n\n' "${ci}"
 
-    gg_printf 'Runs test-llama-archs with 1 to 4 devices\n'
+    gg_printf 'Runs test-llama-archs with 1 to 4 and 8 devices\n'
     gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
     gg_printf '```\n'
     gg_printf '%s\n' "$(cat $OUT/${ci}.log)"

common/arg.cpp

@@ -2464,9 +2464,9 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     ).set_env("LLAMA_ARG_LIVE_CONTEXT_WORKSPACE"));
     add_opt(common_arg(
         {"--kv-gpu-layers"}, "N",
-        string_format("with --no-kv-offload, keep the first N independently owned attention KV layers device-resident "
-                      "for standard and direct hybrid caches. Unsupported specialized caches ignore this option "
-                      "(default: %d)", params.kv_gpu_layers),
+        string_format("with --no-kv-offload, keep N independently owned attention KV layers device-resident. "
+                      "The layers are spent on the devices with the slowest host link first. Unsupported "
+                      "specialized caches ignore this option (default: %d)", params.kv_gpu_layers),
         [](common_params & params, int value) {
             if (value < 0) {
                 throw std::invalid_argument("--kv-gpu-layers must not be negative");
@@ -2935,6 +2935,26 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             }
         }
     ).set_env("LLAMA_ARG_TENSOR_SPLIT"));
+    add_opt(common_arg(
+        {"-as", "--attn-split"}, "N0,N1,N2,...",
+        "fraction of the attention heads to give each GPU under --split-mode tensor, comma-separated, "
+        "e.g. 3,1. Only the heads move, the rest of the model still follows --tensor-split. The share is "
+        "rounded to whole heads, so a ratio the head count cannot express takes the nearest one it can. "
+        "Useful when the GPUs differ in host bandwidth or in speed (default: follow --tensor-split)",
+        [](common_params & params, const std::string & value) {
+            const std::regex regex{ R"([,/]+)" };
+            std::sregex_token_iterator it{ value.begin(), value.end(), regex, -1 };
+            std::vector<std::string> split_arg{ it, {} };
+            if (split_arg.size() > llama_max_devices()) {
+                throw std::invalid_argument(
+                    string_format("got %zu input configs, but system only has %zu devices", split_arg.size(), llama_max_devices())
+                );
+            }
+            for (size_t i = 0; i < llama_max_devices(); ++i) {
+                params.attn_split[i] = i < split_arg.size() ? std::stof(split_arg[i]) : 0.0f;
+            }
+        }
+    ).set_env("LLAMA_ARG_ATTN_SPLIT"));
     add_opt(common_arg(
         {"-mg", "--main-gpu"}, "INDEX",
         string_format("the GPU to use for the model (with split-mode = none), or for intermediate results and KV (with split-mode = row) (default: %d)", params.main_gpu),

common/common.cpp

@@ -1690,6 +1690,8 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
     mparams.load_mode       = params.load_mode;
     mparams.lazy_mode = params.lazy_mode;
     mparams.tensor_split    = params.tensor_split;
+    mparams.attn_split      = std::any_of(params.attn_split, params.attn_split + llama_max_devices(),
+            [](float f) { return f != 0.0f; }) ? params.attn_split : nullptr;
     mparams.check_tensors   = params.check_tensors;
     mparams.use_extra_bufts = !params.no_extra_bufts;
     mparams.no_host         = params.no_host;

common/common.h

@@ -486,6 +486,7 @@ struct common_params {
     int32_t n_gpu_layers       = -1;    // number of layers to store in VRAM, -1 is auto, <= -2 is all
     int32_t main_gpu           = 0;     // the GPU that is used for scratch and small tensors
     float   tensor_split[128]  = {0};   // how split tensors should be distributed across GPUs
+    float   attn_split[128]    = {0};   // how attention heads should be distributed across GPUs
     bool    fit_params         = true;  // whether to fit unset model/context parameters to free device memory
     bool    fit_params_print   = false; // print the estimated required memory to run the model
     int32_t fit_params_min_ctx = 4096;  // minimum context size to set when trying to reduce memory use

ggml/src/ggml-backend-meta.cpp

@@ -485,6 +485,28 @@ static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct
     return it->second[index];
 }
 
+// A scheduler copy is named "<backend>#<source>#<copy>", where <source> also carries the suffixes
+// ggml appends for views. Recover the graph name the copy was made from.
+static std::string ggml_backend_meta_copy_source_name(const char * name) {
+    std::string ret = name;
+    const size_t first = ret.find('#');
+    if (first == std::string::npos) {
+        return ret;
+    }
+    ret.erase(0, first + 1);
+    // ggml writes a view suffix as " (...)", a graph name has no spaces
+    const size_t suffix = ret.find(" (");
+    if (suffix != std::string::npos) {
+        ret.erase(suffix);
+        return ret;
+    }
+    const size_t copy = ret.rfind('#');
+    if (copy != std::string::npos && ret.find_first_not_of("0123456789", copy + 1) == std::string::npos) {
+        ret.erase(copy);
+    }
+    return ret;
+}
+
 static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync);
 
 static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
@@ -831,10 +853,26 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
         if (ggml_nelements(tensor) == 0) {
             return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1};
         }
-        if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE && tensor->view_src == nullptr) {
+        // A host-resident KV cache reaches the graph as a copied-in leaf of the compute buffer.
+        // Mirroring it while the queries stay split by head makes each device attend the wrong
+        // heads, so ask the callback; it still answers MIRRORED for names it does not know.
+        const bool copied_in_leaf =
+            ggml_backend_buffer_get_usage(tensor->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE &&
+            tensor->op == GGML_OP_NONE && (tensor->flags & GGML_TENSOR_FLAG_INPUT) == 0;
+
+        if ((ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE || copied_in_leaf) &&
+                tensor->view_src == nullptr) {
             ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
             const ggml_backend_meta_device_context * dev_ctx = (const ggml_backend_meta_device_context *) dev->context;
-            ggml_backend_meta_split_state ret = dev_ctx->get_split_state(tensor, dev_ctx->get_split_state_ud);
+            // the callback classifies by name, so offer a copy under the name it was copied from
+            const ggml_tensor * tensor_query = tensor;
+            ggml_tensor tensor_named;
+            if (copied_in_leaf) {
+                tensor_named = *tensor;
+                ggml_set_name(&tensor_named, ggml_backend_meta_copy_source_name(tensor->name).c_str());
+                tensor_query = &tensor_named;
+            }
+            ggml_backend_meta_split_state ret = dev_ctx->get_split_state(tensor_query, dev_ctx->get_split_state_ud);
             if (ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) {
                 const int64_t granularity = ret.axis == GGML_BACKEND_SPLIT_AXIS_0 ? ggml_blck_size(tensor->type) : 1;
                 int64_t ne_sum = 0;
@@ -1416,6 +1454,33 @@ static void ggml_backend_meta_buffer_memset_tensor(
 static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
     const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer);
     const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false);
+
+    // A host-resident attention cache reaches this permuted, as [head_dim, n_kv, n_head_kv, 1]. The
+    // heads split, but interleaved per cell, so the chunk splice below cannot express the write.
+    // Each device's heads are one contiguous run per cell: ne[1] cells, from one stride to another.
+    const bool strided_head_split =
+        !ggml_is_contiguous(tensor) &&
+        split_state.axis == GGML_BACKEND_SPLIT_AXIS_2 &&
+        split_state.n_segments == 1 && split_state.nr[0] == 1 &&
+        tensor->ne[3] == 1 && tensor->nb[1] > tensor->nb[2] &&
+        offset == 0 && size == ggml_nbytes(tensor);
+
+    if (strided_head_split) {
+        size_t offset_data = 0;
+        for (size_t j = 0; j < n_bufs; j++) {
+            ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j);
+            const size_t nbytes = split_state.ne[j] * tensor->nb[2];
+            if (nbytes == 0) {
+                continue;
+            }
+            ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, 0, nbytes,
+                tensor->ne[1], simple_tensor->nb[1], tensor->nb[1]);
+            offset_data += nbytes;
+        }
+        GGML_ASSERT(offset_data == (size_t) tensor->ne[2] * tensor->nb[2]);
+        return;
+    }
+
     GGML_ASSERT(ggml_is_contiguous(tensor) || split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
 
     if (split_state.n_segments != 1 || split_state.nr[0] != 1) {
@@ -1544,6 +1609,33 @@ static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, gg
 static void ggml_backend_meta_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {
     const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer);
     const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false);
+
+    // A fused QKV puts Kcur and Vcur in a strided view, which a host-resident cache reads back here.
+    // The rows split, so the chunk splice below cannot express the read. Each device's part of a row
+    // is contiguous: ne[1] rows, from the device's own stride to the view's.
+    const bool strided_rows =
+        !ggml_is_contiguous(tensor) &&
+        split_state.axis == GGML_BACKEND_SPLIT_AXIS_0 &&
+        split_state.n_segments == 1 && split_state.nr[0] == 1 &&
+        tensor->ne[2] == 1 && tensor->ne[3] == 1 &&
+        offset == 0 && size == ggml_nbytes(tensor);
+
+    if (strided_rows) {
+        size_t offset_data = 0;
+        for (size_t j = 0; j < n_bufs; j++) {
+            const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j);
+            const size_t nbytes = ggml_row_size(tensor->type, split_state.ne[j]);
+            if (nbytes == 0) {
+                continue;
+            }
+            ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_data, 0, nbytes,
+                tensor->ne[1], simple_tensor->nb[1], tensor->nb[1]);
+            offset_data += nbytes;
+        }
+        GGML_ASSERT(offset_data == ggml_row_size(tensor->type, tensor->ne[0]));
+        return;
+    }
+
     GGML_ASSERT(ggml_is_contiguous(tensor) || split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
 
     if (split_state.n_segments != 1 || split_state.nr[0] != 1) {
@@ -1700,9 +1792,15 @@ static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_bac
     std::vector<ggml_backend_buffer_t> bufs;
     bufs.reserve(n_simple_bufts);
     for (size_t i = 0; i < n_simple_bufts; i++) {
-        bufs.push_back(ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size));
-        GGML_ASSERT(bufs.back() != nullptr);
-        max_size = std::max(max_size, ggml_backend_buffer_get_size(bufs.back()));
+        ggml_backend_buffer_t buf = ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size);
+        if (buf == nullptr) {
+            for (ggml_backend_buffer_t b : bufs) {
+                ggml_backend_buffer_free(b);
+            }
+            return nullptr;
+        }
+        bufs.push_back(buf);
+        max_size = std::max(max_size, ggml_backend_buffer_get_size(buf));
     }
     ggml_backend_meta_buffer_context * buf_ctx = new ggml_backend_meta_buffer_context(stc_static, stc_compute_0, stc_compute_1, bufs);
 
@@ -1737,27 +1835,35 @@ struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struc
         t->data = (void *) 0x2000000000000000; // FIXME
     }
     for (size_t i = 0; i < n_simple_bufts; i++) {
-        ggml_context * ctx = meta_buf_ctx->stc_static.ctxs[i].get();
+        ggml_context * ctx_simple = meta_buf_ctx->stc_static.ctxs[i].get();
         ggml_backend_buffer_type_t simple_buft = ggml_backend_meta_buft_simple_buft(buft, i);
 
         // If a ggml_context only has zero-sized tensors, ggml_backend_alloc_ctx_tensors_from_buft returns NULL.
         // For those edge cases, allocate a dummy buffer instead.
         bool any_nonzero_slice = false;
-        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx_simple); t != nullptr; t = ggml_get_next_tensor(ctx_simple, t)) {
             if (ggml_nelements(t) != 0) {
                 any_nonzero_slice = true;
                 break;
             }
         }
         if (any_nonzero_slice) {
-            meta_buf_ctx->bufs[i].reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx, simple_buft));
+            meta_buf_ctx->bufs[i].reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx_simple, simple_buft));
         } else {
             meta_buf_ctx->bufs[i].reset(ggml_backend_buft_alloc_buffer(simple_buft, 0));
-            for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            for (ggml_tensor * t = ggml_get_first_tensor(ctx_simple); t != nullptr; t = ggml_get_next_tensor(ctx_simple, t)) {
                 t->buffer = meta_buf_ctx->bufs[i].get();
             }
         }
-        GGML_ASSERT(meta_buf_ctx->bufs[i]);
+        if (meta_buf_ctx->bufs[i] == nullptr) {
+            // the loop above already pointed the tensors at the buffer that is freed here
+            for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+                t->buffer = nullptr;
+                t->data   = nullptr;
+            }
+            ggml_backend_buffer_free(meta_buf);
+            return nullptr;
+        }
         meta_buf->size = std::max(meta_buf->size, ggml_backend_buffer_get_size(meta_buf_ctx->bufs[i].get()));
     }
     return meta_buf;

ggml/src/ggml-backend.cpp

@@ -1034,6 +1034,19 @@ static void ggml_backend_sched_print_assignments(ggml_backend_sched_t sched, str
     }
 }
 
+// Name a copy of an input "<backend>#<source>#<copy>". The name field has a fixed size, so cut the
+// backend label rather than the source name - the source name is what identifies the copy.
+static void ggml_backend_sched_name_copy(
+        struct ggml_tensor * copy, const char * backend_name, const struct ggml_tensor * src, int c) {
+    const int n_tail = snprintf(NULL, 0, "#%s#%d", src->name, c);
+    const int n_max  = GGML_MAX_NAME - 1 - n_tail;
+    int n_head = (int) strlen(backend_name);
+    if (n_head > n_max) {
+        n_head = n_max > 0 ? n_max : 0;
+    }
+    ggml_format_name(copy, "%.*s#%s#%d", n_head, backend_name, src->name, c);
+}
+
 static bool ggml_backend_sched_buffer_supported(ggml_backend_sched_t sched, struct ggml_tensor * t, int backend_id) {
     ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer;
     ggml_backend_buffer_type_t buft = NULL;
@@ -1392,7 +1405,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
                                 tensor_copy = src; // use the original tensor as the current copy
                             } else {
                                 tensor_copy = ggml_dup_tensor_layout(sched->ctx, src);
-                                ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c);
+                                ggml_backend_sched_name_copy(tensor_copy, ggml_backend_name(backend), src, c);
                             }
                             ggml_set_input(tensor_copy);
                             ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor
@@ -1413,7 +1426,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
                         ggml_backend_t backend = sched->backends[cur_backend_id];
                         for (int c = 0; c < sched->n_copies; c++) {
                             struct ggml_tensor * tensor_copy = ggml_dup_tensor_layout(sched->ctx, src);
-                            ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c);
+                            ggml_backend_sched_name_copy(tensor_copy, ggml_backend_name(backend), src, c);
                             if (sched->n_copies > 1) {
                                 ggml_set_input(tensor_copy);
                                 ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor

include/llama.h

@@ -329,6 +329,10 @@ extern "C" {
         // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices()
         const float * tensor_split;
 
+        // proportion of the attention heads to give each GPU under split mode tensor, size: llama_max_devices()
+        // NULL follows tensor_split. Only the heads move - the rest of the model still follows tensor_split
+        const float * attn_split;
+
         // Called with a progress value between 0.0 and 1.0. Pass NULL to disable.
         // If the provided progress_callback returns true, model loading continues.
         // If it returns false, model loading is immediately aborted.