llama : add an attention split separate from the tensor split

common/arg.cpp

@@ -2904,6 +2904,45 @@ 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. Each share is "
+        "rounded down to whole heads, so a ratio the head count cannot express is not matched exactly. "
+        "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())
+                );
+            }
+            float sum = 0.0f;
+            for (size_t i = 0; i < llama_max_devices(); ++i) {
+                float share = 0.0f;
+                if (i < split_arg.size()) {
+                    size_t n_read = 0;
+                    try {
+                        share = std::stof(split_arg[i], &n_read);
+                    } catch (const std::exception &) {
+                        n_read = 0;
+                    }
+                    if (n_read != split_arg[i].size() || !std::isfinite(share) || share < 0.0f) {
+                        throw std::invalid_argument(
+                            string_format("invalid attention split share '%s'", split_arg[i].c_str())
+                        );
+                    }
+                }
+                params.attn_split[i] = share;
+                sum += share;
+            }
+            if (sum <= 0.0f) {
+                throw std::invalid_argument("the attention split shares must add up to more than zero");
+            }
+        }
+    ).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

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.

src/llama-model.cpp

@@ -828,6 +828,18 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
         const int64_t ne_axis   = unfolded ? tensor->ne[cache_copy.axis]*unit : tensor->ne[split_state.axis];
         const int64_t blck_size = ggml_blck_size(tc.tensor_axis_0->type);
         const float * tensor_split = ud->model->tensor_split();
+
+        // The share of the attention heads decides how much cache a device gets and how much attention
+        // work it does, neither of which has to follow the memory split. Only a split that counts heads
+        // follows it, and those are the ones measured against attn_output.weight - a linear-attention
+        // layer measures against ssm_out.weight and keeps tensor_split.
+        if (ud->model->attn_split() != nullptr) {
+            const std::string attn_out_name = "blk." + std::to_string(tc.il) + ".attn_output.weight";
+            const ggml_tensor * attn_out = ud->model->get_tensor(attn_out_name.c_str());
+            if (attn_out != nullptr && tc.tensor_axis_0 == attn_out) {
+                tensor_split = ud->model->attn_split();
+            }
+        }
         std::vector<float> tensor_split_scan;
         tensor_split_scan.reserve(ud->n_devices);
         for (size_t j = 0; j < ud->n_devices; j++) {
@@ -1221,6 +1233,7 @@ struct llama_model::impl {
     bool has_tensor_overrides;
 
     std::vector<float> tensor_split_owned;
+    std::vector<float> attn_split_owned;
 };
 
 llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique<impl>()) {
@@ -1230,6 +1243,23 @@ llama_model::llama_model(const llama_model_params & params) : params(params), pi
         pimpl->tensor_split_owned.assign(params.tensor_split, params.tensor_split + llama_max_devices());
         this->params.tensor_split = pimpl->tensor_split_owned.data();
     }
+    if (params.attn_split != nullptr) {
+        // the shares reach a cumulative division and a conversion to whole heads, so a negative, a
+        // non-finite or an all-zero set has no meaning here - fall back to the tensor split
+        float sum = 0.0f;
+        bool  ok  = true;
+        for (size_t i = 0; i < llama_max_devices(); i++) {
+            ok = ok && std::isfinite(params.attn_split[i]) && params.attn_split[i] >= 0.0f;
+            sum += params.attn_split[i];
+        }
+        if (!ok || !(sum > 0.0f)) {
+            LLAMA_LOG_WARN("%s: the attention split is not a set of non-negative shares; ignoring it\n", __func__);
+            this->params.attn_split = nullptr;
+        } else {
+            pimpl->attn_split_owned.assign(params.attn_split, params.attn_split + llama_max_devices());
+            this->params.attn_split = pimpl->attn_split_owned.data();
+        }
+    }
     pimpl->has_tensor_overrides = params.tensor_buft_overrides && params.tensor_buft_overrides[0].pattern;
 }
 
@@ -1941,6 +1971,10 @@ const float * llama_model::tensor_split() const {
     return params.tensor_split;
 }
 
+const float * llama_model::attn_split() const {
+    return params.attn_split;
+}
+
 uint32_t llama_model::n_gpu_layers() const {
     // note: plus 1 for the "output" layer
     return params.n_gpu_layers >= 0 ? params.n_gpu_layers : hparams.n_layer_all + 1;
@@ -2837,6 +2871,7 @@ llama_model_params llama_model_default_params() {
         /*.lazy_mode                   =*/ LLAMA_LAZY_MODE_AUTO,
         /*.main_gpu                    =*/ 0,
         /*.tensor_split                =*/ nullptr,
+        /*.attn_split                  =*/ nullptr,
         /*.progress_callback           =*/ nullptr,
         /*.progress_callback_user_data =*/ nullptr,
         /*.kv_overrides                =*/ nullptr,

src/llama-model.h

@@ -729,6 +729,7 @@ struct llama_model {
     size_t n_tensors() const;
     size_t n_devices() const;
     const float * tensor_split() const;
+    const float * attn_split() const;
 
     uint32_t n_gpu_layers() const;
     llama_split_mode split_mode() const;

tests/test-llama-archs.cpp

@@ -400,10 +400,11 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/)
 }
 
 // with offload_kqv=false the cache lives in host memory
-// n_seq_max > 1 gives the cache one stream per sequence
+// n_seq_max > 1 gives the cache one stream per sequence, attn_split gives the heads their own share
 struct kv_config {
-    bool     offload_kqv = true;
-    uint32_t n_seq_max   = 1;
+    bool               offload_kqv = true;
+    uint32_t           n_seq_max   = 1;
+    std::vector<float> attn_split;
 };
 
 static std::pair<llama_model_ptr, llama_context_ptr> get_model_and_ctx(
@@ -416,6 +417,11 @@ static std::pair<llama_model_ptr, llama_context_ptr> get_model_and_ctx(
     devs_copy.push_back(nullptr);
     model_params.devices = devs_copy.data();
     model_params.split_mode = split_mode;
+    std::vector<float> attn_split = kvc.attn_split;
+    if (!attn_split.empty()) {
+        attn_split.resize(llama_max_devices(), 0.0f);
+        model_params.attn_split = attn_split.data();
+    }
 
     llama_context_params ctx_params = llama_context_default_params();
     ctx_params.n_ctx = 0;
@@ -1459,6 +1465,19 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in
         kvc_host_streams.n_seq_max = 2;
         dev_configs.emplace_back(devices_meta, "Meta -nkvo -np 2", LLAMA_SPLIT_MODE_TENSOR, kvc_host_streams);
 
+        // the same, with all attention heads on the first device
+        if (devices_meta.size() > 1) {
+            kv_config kvc_attn = kvc_host;
+            kvc_attn.attn_split.assign(devices_meta.size(), 0.0f);
+            kvc_attn.attn_split[0] = 1.0f;
+            dev_configs.emplace_back(devices_meta, "Meta -nkvo -as", LLAMA_SPLIT_MODE_TENSOR, kvc_attn);
+
+            // a custom head split must also hold across streams
+            kv_config kvc_attn_streams = kvc_attn;
+            kvc_attn_streams.n_seq_max = 2;
+            dev_configs.emplace_back(devices_meta, "Meta -nkvo -as -np 2", LLAMA_SPLIT_MODE_TENSOR, kvc_attn_streams);
+        }
+
         for (const device_config & dc : dev_configs) {
             max_device_label_length = std::max(max_device_label_length, dc.label.length());
         }