From 88c3a208d7accaf176f5d2c5306e7364a2c8d9a5 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 11:35:05 +0200 Subject: [PATCH 1/3] llama : add an attention split separate from the tensor split The share of the attention heads a device gets decides both how much of a host- resident cache it receives and how much attention work it does, and neither has to follow the memory split. Add --attn-split (-as) to set it; unset it follows --tensor-split and nothing changes. The share is rounded to whole heads by the granularity rules that already keep GQA consistent. Assisted-by: Claude Opus 5 --- common/arg.cpp | 20 ++++++++++++++++++++ common/common.cpp | 2 ++ common/common.h | 1 + include/llama.h | 4 ++++ src/llama-model.cpp | 19 +++++++++++++++++++ src/llama-model.h | 1 + tests/test-llama-archs.cpp | 25 ++++++++++++++++++++++--- 7 files changed, 69 insertions(+), 3 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 3acfbab6e544..545783037795 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2904,6 +2904,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 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), diff --git a/common/common.cpp b/common/common.cpp index 05af52737426..966dcd077218 100644 --- a/common/common.cpp +++ b/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; diff --git a/common/common.h b/common/common.h index 8edd79f3d053..035261b2ee06 100644 --- a/common/common.h +++ b/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 diff --git a/include/llama.h b/include/llama.h index 386f0fd7ad16..4a1f379d9ee2 100644 --- a/include/llama.h +++ b/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. diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 66ae3b617fb6..55593a19e9fc 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -828,6 +828,15 @@ 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. A linear-attention layer is a + // different mechanism and keeps tensor_split. + if (ud->model->attn_split() != nullptr && !hparams.is_recr(tc.il) && + (std::regex_match(tensor_name, pattern_kv_cache) || + (tensor_name.substr(0, 4) == "blk." && tensor_name.find(".attn_") != std::string::npos))) { + tensor_split = ud->model->attn_split(); + } std::vector tensor_split_scan; tensor_split_scan.reserve(ud->n_devices); for (size_t j = 0; j < ud->n_devices; j++) { @@ -1221,6 +1230,7 @@ struct llama_model::impl { bool has_tensor_overrides; std::vector tensor_split_owned; + std::vector attn_split_owned; }; llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique()) { @@ -1230,6 +1240,10 @@ 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) { + 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 +1955,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 +2855,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, diff --git a/src/llama-model.h b/src/llama-model.h index 2cb74ea23273..76777c2f9e4c 100644 --- a/src/llama-model.h +++ b/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; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index cf193bf168e6..ff7d16191b3a 100644 --- a/tests/test-llama-archs.cpp +++ b/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 attn_split; }; static std::pair get_model_and_ctx( @@ -416,6 +417,11 @@ static std::pair get_model_and_ctx( devs_copy.push_back(nullptr); model_params.devices = devs_copy.data(); model_params.split_mode = split_mode; + std::vector 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()); } From aef345abd470c8156d84e15702b2f911bdcb3d15 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 08:18:16 +0200 Subject: [PATCH 2/3] llama : validate the attention split shares std::stof took a partial parse, a negative, a NaN and an all-zero set without a word, and those reach a cumulative division and a conversion to whole heads. Reject them at the command line and fall back to the tensor split at the API boundary. Also say what the rounding does instead of promising the nearest share. Assisted-by: Claude Opus 5 --- common/arg.cpp | 25 ++++++++++++++++++++++--- src/llama-model.cpp | 17 +++++++++++++++-- 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 545783037795..3d2fd10294b3 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2907,8 +2907,8 @@ common_params_context common_params_parser_init(common_params & params, llama_ex 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. " + "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"([,/]+)" }; @@ -2919,8 +2919,27 @@ common_params_context common_params_parser_init(common_params & params, llama_ex 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) { - params.attn_split[i] = i < split_arg.size() ? std::stof(split_arg[i]) : 0.0f; + 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")); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 55593a19e9fc..fcd78cf7c0da 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1241,8 +1241,21 @@ llama_model::llama_model(const llama_model_params & params) : params(params), pi this->params.tensor_split = pimpl->tensor_split_owned.data(); } if (params.attn_split != nullptr) { - pimpl->attn_split_owned.assign(params.attn_split, params.attn_split + llama_max_devices()); - this->params.attn_split = pimpl->attn_split_owned.data(); + // 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; } From 6fe42745f29a1fa1facac377a54debadb5eb6538 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Mon, 7 Sep 2026 21:42:17 +0200 Subject: [PATCH 3/3] llama : pick the attention split by what the tensor is measured against A substring of the name also caught tensors whose split does not count heads. Assisted-by: Claude Opus 5 --- src/llama-model.cpp | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index fcd78cf7c0da..652c72536f83 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -830,12 +830,15 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str 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. A linear-attention layer is a - // different mechanism and keeps tensor_split. - if (ud->model->attn_split() != nullptr && !hparams.is_recr(tc.il) && - (std::regex_match(tensor_name, pattern_kv_cache) || - (tensor_name.substr(0, 4) == "blk." && tensor_name.find(".attn_") != std::string::npos))) { - tensor_split = ud->model->attn_split(); + // 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 tensor_split_scan; tensor_split_scan.reserve(ud->n_devices);