diff --git a/common/arg.cpp b/common/arg.cpp index d2c1838208dd..806f70578602 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2905,6 +2905,48 @@ 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. 1,3. The heads take their weights and their KV cache along, the rest of the model (FFN, " + "linear attention, output) still follows --tensor-split. Each layer gets whole KV heads, and the " + "remainder is spread over the layers, so the share over all layers follows the ratio within one " + "head (the log shows the result). 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 " + "(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()) + ); + } + 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), diff --git a/common/common.cpp b/common/common.cpp index d4fdd1b4c716..879d32a4fa39 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1681,6 +1681,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 21f29a1d5a6e..0959f5ca2919 100644 --- a/common/common.h +++ b/common/common.h @@ -480,6 +480,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 f82a68b107f2..882ec438070d 100644 --- a/include/llama.h +++ b/include/llama.h @@ -336,6 +336,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 5d48d702c5b8..1842f83b054d 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -28,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -445,6 +446,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str uint32_t il; size_t rotation; // when assigning tensor slices, rotate how the rounding is done for more even allocation + size_t il_eff = 0; // index of the layer among the previous layers of its kind, see get_il_eff }; auto get_tensor_config_impl = [&]( @@ -464,18 +466,21 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str uint32_t il; std::string prefix; size_t rotation; + size_t il_eff = 0; if (tensor_name.substr(0, 4) == "blk.") { const size_t length_prefix = tensor_name.find('.', 4); GGML_ASSERT(length_prefix != std::string::npos); prefix = tensor_name.substr(0, length_prefix + 1); il = std::stoull(tensor_name.substr(4, length_prefix)); - rotation = get_il_eff(il) % ud->n_devices; + il_eff = get_il_eff(il); + rotation = il_eff % ud->n_devices; } else if (tensor_name.substr(0, 6) == "cache_") { const size_t layer_index_start = tensor_name.find("_l", 6); GGML_ASSERT(layer_index_start != std::string::npos); il = std::stoull(tensor_name.substr(layer_index_start + 2)); prefix = "blk." + std::to_string(il) + "."; - rotation = get_il_eff(il) % ud->n_devices; + il_eff = get_il_eff(il); + rotation = il_eff % ud->n_devices; } else { il = 0; rotation = hparams.n_layer() % ud->n_devices; @@ -486,7 +491,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str tensor_axis_0 = ud->model->get_tensor((prefix + suffix_fallback).c_str()); } GGML_ASSERT(tensor_axis_0 != nullptr); - return {axis, tensor_axis_0, il, rotation}; + return {axis, tensor_axis_0, il, rotation, il_eff}; }; auto get_tensor_config = [&]() -> tensor_config { @@ -851,10 +856,32 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS) { const int64_t blck_size = ggml_blck_size(tc.tensor_axis_0->type); const float * tensor_split = ud->model->tensor_split(); - std::vector tensor_split_scan; + + // 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. + bool use_attn_split = false; + 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(); + use_attn_split = true; + } + } + + // The attention split spreads the rounding over the layers: the boundaries of layer il_eff are the + // differences of the rounded cumulative boundaries of the layers before it, so the heads over all layers + // of a kind follow the ratio within one head. The tensor split rotates the device that is rounded first. + const size_t rotation = use_attn_split ? 0 : tc.rotation; + + // in double: the tensors of a layer have different sizes, and a share that is inexact in binary must not + // round them to different numbers of heads + std::vector tensor_split_scan; tensor_split_scan.reserve(ud->n_devices); for (size_t j = 0; j < ud->n_devices; j++) { - tensor_split_scan.push_back(tensor_split == nullptr ? 0.0f : tensor_split[(j + tc.rotation) % ud->n_devices]); + tensor_split_scan.push_back(tensor_split == nullptr ? 0.0 : (double) tensor_split[(j + rotation) % ud->n_devices]); if (j > 0) { tensor_split_scan[j] += tensor_split_scan[j - 1]; } @@ -865,18 +892,32 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str const int64_t ne_s = segments[is].first; const uint32_t nr_s = segments[is].second; const int64_t g_s = granularity[is]; + + // the share of devices [0, j] in units of the granularity + const double n_units = (double) ne_s / g_s; + auto units_upto = [&](size_t j) -> double { + const double frac = tensor_split_scan.back() == 0.0 ? + double(j + 1)/ud->n_devices : tensor_split_scan[j]/tensor_split_scan.back(); + return n_units * frac; + }; + int64_t low = 0; size_t j = 0; for (; j < ud->n_devices - 1; j++) { - int64_t high = tensor_split_scan.back() == 0.0f ? - ne_s * (j+1)/ud->n_devices : ne_s * tensor_split_scan[j]/tensor_split_scan.back(); - if (high % g_s != 0) { - high -= high % g_s; + int64_t high_units; + if (use_attn_split) { + const double x = units_upto(j); + high_units = std::llround(x*(tc.il_eff + 1)) - std::llround(x*tc.il_eff); + } else { + // a share that is exact in decimal can come out just below the boundary in binary + high_units = (int64_t) std::floor(units_upto(j) + 1e-9); } - split_state.ne[is*ud->n_devices + (j + tc.rotation) % ud->n_devices] = high - low; + int64_t high = std::clamp(high_units, 0, (int64_t) n_units) * g_s; + high = std::max(high, low); + split_state.ne[is*ud->n_devices + (j + rotation) % ud->n_devices] = high - low; low = high; } - split_state.ne[is*ud->n_devices + (j + tc.rotation) % ud->n_devices] = ne_s - low; + split_state.ne[is*ud->n_devices + (j + rotation) % ud->n_devices] = ne_s - low; split_state.nr[is] = nr_s; } split_state.n_segments = segments.size(); @@ -1270,6 +1311,7 @@ struct llama_model::impl { bool has_tensor_overrides; std::vector tensor_split_owned; + std::vector attn_split_owned; }; bool llama_prec_policy::apply(ggml_tensor * res) const { @@ -1329,6 +1371,27 @@ 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 + // they are kept as fractions of their sum, so that large shares cannot overflow the rounding + double sum = 0.0; + 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 || !std::isfinite(sum) || !(sum > 0.0)) { + 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.resize(llama_max_devices()); + for (size_t i = 0; i < llama_max_devices(); i++) { + pimpl->attn_split_owned[i] = (float) (params.attn_split[i] / sum); + } + this->params.attn_split = pimpl->attn_split_owned.data(); + } + } pimpl->has_tensor_overrides = params.tensor_buft_overrides && params.tensor_buft_overrides[0].pattern; } @@ -1999,6 +2062,52 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } } + // the attention split gives each layer whole KV heads, rounding toward a different device from layer to layer: + // show what it came to over all layers + if (params.split_mode == LLAMA_SPLIT_MODE_TENSOR && params.attn_split != nullptr && get_split_state_ud.n_devices > 1) { + const size_t n_dev = get_split_state_ud.n_devices; + + std::vector n_head_kv_dev(n_dev, 0); + int64_t n_head_kv_all = 0; + + for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { + const std::string name = "blk." + std::to_string(il) + ".attn_output.weight"; + const ggml_tensor * attn_out = get_tensor(name.c_str()); + // layers left in host memory are not split, layers that reuse another layer's KV have none of their own + if (attn_out == nullptr || attn_out->buffer == nullptr || ggml_backend_buffer_is_host(attn_out->buffer) || !hparams.has_kv(il)) { + continue; + } + + const ggml_backend_meta_split_state split_state = llama_meta_device_get_split_state(attn_out, &get_split_state_ud); + if (split_state.axis != GGML_BACKEND_SPLIT_AXIS_0) { + continue; + } + + const int64_t n_embd_head_kv = (int64_t) hparams.n_embd_head_v(il) * hparams.n_gqa(il); + for (size_t j = 0; j < n_dev; ++j) { + n_head_kv_dev[j] += split_state.ne[j] / n_embd_head_kv; + } + n_head_kv_all += attn_out->ne[0] / n_embd_head_kv; + } + + float share_used = 0.0f; + for (size_t j = 0; j < n_dev; ++j) { + share_used += params.attn_split[j]; + } + if (share_used == 0.0f) { + LLAMA_LOG_WARN("%s: the attention split gives the %zu devices in use no share, the heads follow the tensor split\n", __func__, n_dev); + } + if (n_head_kv_all == 0) { + LLAMA_LOG_WARN("%s: the attention split has no effect, no layer of this model splits its attention heads on the devices\n", __func__); + } else { + std::string desc; + for (size_t j = 0; j < n_dev; ++j) { + desc += format("%sdevice %zu: %" PRId64 " (%.1f%%)", j > 0 ? ", " : "", j, n_head_kv_dev[j], 100.0*n_head_kv_dev[j]/n_head_kv_all); + } + LLAMA_LOG_INFO("%s: attention split, KV heads of all layers per device: %s\n", __func__, desc.c_str()); + } + } + if (ml.no_alloc) { return true; } @@ -2403,6 +2512,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; @@ -3342,6 +3455,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 6335dc2bc318..8b67b8b42484 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -868,6 +868,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 17f5870c4a98..d0250d5aa7d9 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -515,7 +515,8 @@ static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const float stdev, const std::vector & devs, const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, - const llama_model_tensor_buft_override * tensor_buft_overrides = nullptr, const size_t moe_cache_size = 0) { + const llama_model_tensor_buft_override * tensor_buft_overrides = nullptr, const size_t moe_cache_size = 0, + const std::vector & attn_split = {}) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; @@ -524,6 +525,11 @@ static std::pair get_model_and_ctx( model_params.devices = devs_copy.data(); model_params.split_mode = split_mode; model_params.tensor_buft_overrides = tensor_buft_overrides; + std::vector attn_split_all = attn_split; + if (!attn_split_all.empty()) { + attn_split_all.resize(llama_max_devices(), 0.0f); + model_params.attn_split = attn_split_all.data(); + } llama_context_params ctx_params = llama_context_default_params(); ctx_params.n_ctx = 0; @@ -952,9 +958,12 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con llama_split_mode split_mode; bool host_experts; // keep the experts in host memory, see host_experts_test size_t moe_cache_size; + std::vector attn_split; // empty: follow the tensor split - device_config(std::vector devs, std::string name, llama_split_mode split_mode, bool host_experts = false, size_t moe_cache_size = 0) - : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode), host_experts(host_experts), moe_cache_size(moe_cache_size) {} + device_config(std::vector devs, std::string name, llama_split_mode split_mode, bool host_experts = false, size_t moe_cache_size = 0, + std::vector attn_split = {}) + : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode), host_experts(host_experts), moe_cache_size(moe_cache_size), + attn_split(std::move(attn_split)) {} }; const llama_model_tensor_buft_override host_experts_overrides[] = { @@ -985,6 +994,18 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con if (target_backend == nullptr) { dev_configs.emplace_back(devices_meta, "Meta", LLAMA_SPLIT_MODE_TENSOR); + + // the attention heads and their KV cache get their own share: all on the first device, and an uneven one + if (devices_meta.size() > 1) { + std::vector attn_first(devices_meta.size(), 0.0f); + attn_first[0] = 1.0f; + dev_configs.emplace_back(devices_meta, "Meta -as 1/0", LLAMA_SPLIT_MODE_TENSOR, false, 0, attn_first); + + std::vector attn_uneven(devices_meta.size(), 3.0f); + attn_uneven[0] = 1.0f; + dev_configs.emplace_back(devices_meta, "Meta -as 1/3", LLAMA_SPLIT_MODE_TENSOR, false, 0, attn_uneven); + max_device_label_length = std::max(max_device_label_length, dev_configs.back().label.length()); + } } // the ops that use the experts are offloaded to the first device and the scheduler copies the used experts @@ -1090,7 +1111,7 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con } if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_arch_supports_sm_tensor(arch)) { test_executed = true; - model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, dc.devs, dc.split_mode, encode, overrides, dc.moe_cache_size); + model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, dc.devs, dc.split_mode, encode, overrides, dc.moe_cache_size, dc.attn_split); logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode); const double nmse_val = nmse(logits_cpu, logits_dev); snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val); diff --git a/tools/llama-bench/README.md b/tools/llama-bench/README.md index 376007212115..64c4daf083c9 100644 --- a/tools/llama-bench/README.md +++ b/tools/llama-bench/README.md @@ -72,6 +72,7 @@ test parameters: -lzm, --lazy-mode (default: auto) -embd, --embeddings <0|1> (default: 0) -ts, --tensor-split (default: 0) + -as, --attn-split (default: 0 = follow --tensor-split) -ot --override-tensor =;... (default: disabled) -nopo, --no-op-offload <0|1> (default: 0) diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index c3fc8b6da3c0..9e6069e174a1 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -361,6 +361,7 @@ struct cmd_params { std::vector flash_attn_native_quants; std::vector> devices; std::vector> tensor_split; + std::vector> attn_split; std::vector> tensor_buft_overrides; std::vector embeddings; std::vector no_op_offload; @@ -408,6 +409,7 @@ static const cmd_params cmd_params_defaults = { /* flash_attn_native_quants */ { false }, /* devices */ { {} }, /* tensor_split */ { std::vector(llama_max_devices(), 0.0f) }, + /* attn_split */ { std::vector(llama_max_devices(), 0.0f) }, /* tensor_buft_overrides*/ { std::vector{ { nullptr, nullptr } } }, /* embeddings */ { false }, /* no_op_offload */ { false }, @@ -484,6 +486,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -lzm, --lazy-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.lazy_mode, lazy_mode_str), ",").c_str()); printf(" -embd, --embeddings <0|1> (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str()); printf(" -ts, --tensor-split (default: 0)\n"); + printf(" -as, --attn-split (default: 0 = follow --tensor-split)\n"); printf(" -ot --override-tensor =;...\n"); printf(" (default: disabled)\n"); printf(" -nopo, --no-op-offload <0|1> (default: 0)\n"); @@ -926,11 +929,12 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } auto p = string_split(argv[i], split_delim); params.repack.insert(params.repack.end(), p.begin(), p.end()); - } else if (arg == "-ts" || arg == "--tensor-split") { + } else if (arg == "-ts" || arg == "--tensor-split" || arg == "-as" || arg == "--attn-split") { if (++i >= argc) { invalid_param = true; break; } + auto & dst = arg == "-ts" || arg == "--tensor-split" ? params.tensor_split : params.attn_split; for (auto ts : string_split(argv[i], split_delim)) { // split string by ; and / const std::regex regex{ R"([;/]+)" }; @@ -946,7 +950,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { tensor_split[i] = 0.0f; } } - params.tensor_split.push_back(tensor_split); + dst.push_back(tensor_split); } } else if (arg == "-ot" || arg == "--override-tensor") { if (++i >= argc) { @@ -1184,6 +1188,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { if (params.tensor_split.empty()) { params.tensor_split = cmd_params_defaults.tensor_split; } + if (params.attn_split.empty()) { + params.attn_split = cmd_params_defaults.attn_split; + } if (params.tensor_buft_overrides.empty()) { params.tensor_buft_overrides = cmd_params_defaults.tensor_buft_overrides; } @@ -1245,6 +1252,7 @@ struct cmd_params_instance { bool flash_attn_native_quants; std::vector devices; std::vector tensor_split; + std::vector attn_split; std::vector tensor_buft_overrides; bool embeddings; bool no_op_offload; @@ -1265,6 +1273,8 @@ struct cmd_params_instance { mparams.lazy_mode = lazy_mode; mparams.main_gpu = main_gpu; mparams.tensor_split = tensor_split.data(); + mparams.attn_split = std::any_of(attn_split.begin(), attn_split.end(), [](float f) { return f != 0.0f; }) ? + attn_split.data() : nullptr; mparams.no_host = no_host; mparams.use_extra_bufts = repack; @@ -1310,7 +1320,7 @@ struct cmd_params_instance { bool equal_mparams(const cmd_params_instance & other) const { return model == other.model && n_gpu_layers == other.n_gpu_layers && n_cpu_moe == other.n_cpu_moe && split_mode == other.split_mode && - main_gpu == other.main_gpu && tensor_split == other.tensor_split && + main_gpu == other.main_gpu && tensor_split == other.tensor_split && attn_split == other.attn_split && load_mode == other.load_mode && lazy_mode == other.lazy_mode && devices == other.devices && no_host == other.no_host && repack == other.repack && vec_tensor_buft_override_equal(tensor_buft_overrides, other.tensor_buft_overrides); @@ -1351,6 +1361,7 @@ static std::vector get_cmd_params_instances(const cmd_param for (const auto & mg : params.main_gpu) for (const auto & devs : params.devices) for (const auto & ts : params.tensor_split) + for (const auto & as : params.attn_split) for (const auto & ot : params.tensor_buft_overrides) for (const auto & noh : params.no_host) for (const auto & rpk : params.repack) @@ -1396,6 +1407,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .flash_attn_native_quants = */ fanq, /* .devices = */ devs, /* .tensor_split = */ ts, + /* .attn_split = */ as, /* .tensor_buft_overrides = */ ot, /* .embeddings = */ embd, /* .no_op_offload = */ nopo, @@ -1435,6 +1447,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .flash_attn_native_quants = */ fanq, /* .devices = */ devs, /* .tensor_split = */ ts, + /* .attn_split = */ as, /* .tensor_buft_overrides = */ ot, /* .embeddings = */ embd, /* .no_op_offload = */ nopo, @@ -1474,6 +1487,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .flash_attn_native_quants = */ fanq, /* .devices = */ devs, /* .tensor_split = */ ts, + /* .attn_split = */ as, /* .tensor_buft_overrides = */ ot, /* .embeddings = */ embd, /* .no_op_offload = */ nopo, @@ -1518,6 +1532,7 @@ struct test { bool flash_attn_native_quants; std::vector devices; std::vector tensor_split; + std::vector attn_split; std::vector tensor_buft_overrides; bool embeddings; bool no_op_offload; @@ -1560,6 +1575,7 @@ struct test { flash_attn_native_quants = inst.flash_attn_native_quants; devices = inst.devices; tensor_split = inst.tensor_split; + attn_split = inst.attn_split; tensor_buft_overrides = inst.tensor_buft_overrides; embeddings = inst.embeddings; no_op_offload = inst.no_op_offload; @@ -1623,7 +1639,7 @@ struct test { "n_ubatch", "n_threads", "cpu_mask", "cpu_strict", "poll", "type_k", "type_v", "n_gpu_layers", "n_cpu_moe", "split_mode", "main_gpu", "no_kv_offload", "flash_attn", "flash_attn_native_quants", "devices", "tensor_split", - "tensor_buft_overrides", "load_mode", "lazy_mode", + "attn_split", "tensor_buft_overrides", "load_mode", "lazy_mode", "embeddings", "no_op_offload", "no_host", "repack", "fit_target", "fit_min_ctx", "n_prompt", "n_gen", "n_depth", @@ -1655,23 +1671,29 @@ struct test { return STRING; } - std::vector get_values() const { - std::string tensor_split_str; - std::string tensor_buft_overrides_str; + static std::string split_to_string(const std::vector & split) { + std::string res; int max_nonzero = 0; for (size_t i = 0; i < llama_max_devices(); i++) { - if (tensor_split[i] > 0) { + if (split[i] > 0) { max_nonzero = i; } } for (int i = 0; i <= max_nonzero; i++) { char buf[32]; - snprintf(buf, sizeof(buf), "%.2f", tensor_split[i]); - tensor_split_str += buf; + snprintf(buf, sizeof(buf), "%.2f", split[i]); + res += buf; if (i < max_nonzero) { - tensor_split_str += "/"; + res += "/"; } } + return res; + } + + std::vector get_values() const { + const std::string tensor_split_str = split_to_string(tensor_split); + const std::string attn_split_str = split_to_string(attn_split); + std::string tensor_buft_overrides_str; if (tensor_buft_overrides.size() == 1) { // Last element of tensor_buft_overrides is always a null pattern // so if it is only one element long, it must be a null pattern. @@ -1718,6 +1740,7 @@ struct test { std::to_string(flash_attn_native_quants), devices_to_string(devices), tensor_split_str, + attn_split_str, tensor_buft_overrides_str, llama_load_mode_name(load_mode), lazy_mode_str(lazy_mode), @@ -1970,6 +1993,9 @@ struct markdown_printer : public printer { if (field == "tensor_split") { return "ts"; } + if (field == "attn_split") { + return "as"; + } if (field == "tensor_buft_overrides") { return "ot"; } @@ -2042,6 +2068,9 @@ struct markdown_printer : public printer { if (params.tensor_split.size() > 1 || params.tensor_split != cmd_params_defaults.tensor_split) { fields.emplace_back("tensor_split"); } + if (params.attn_split.size() > 1 || params.attn_split != cmd_params_defaults.attn_split) { + fields.emplace_back("attn_split"); + } if (params.tensor_buft_overrides.size() > 1 || !vec_vec_tensor_buft_override_equal(params.tensor_buft_overrides, cmd_params_defaults.tensor_buft_overrides)) { fields.emplace_back("tensor_buft_overrides"); }