Support -fa auto in llama-bench (#23714)

* Support `-fa auto` in llama-bench

Make the default value of `-ngl` -1, similar to other tools.

Update README with latest usage and examples

* Address review comments
This commit is contained in:
Gaurav Garg
2026-05-31 02:03:57 +05:30
committed by GitHub
parent d6588daa80
commit aa46bda89b
2 changed files with 108 additions and 51 deletions
+38 -16
View File
@@ -19,6 +19,7 @@
#include <vector>
#include <unordered_set>
#include "arg.h"
#include "build-info.h"
#include "common.h"
#include "download.h"
@@ -275,9 +276,11 @@ static std::string pair_str(const std::pair<int, int> & p) {
return buf;
}
static std::vector<int> parse_int_range(const std::string & s) {
static std::vector<int> parse_int_range(const std::string & s, bool allow_negative = false) {
// first[-last[(+|*)step]]
std::regex range_regex(R"(^(\d+)(?:-(\d+)(?:([\+|\*])(\d+))?)?(?:,|$))");
std::regex range_regex(allow_negative
? R"(^(-?\d+)(?:-(\d+)(?:([\+|\*])(\d+))?)?(?:,|$))"
: R"(^(\d+)(?:-(\d+)(?:([\+|\*])(\d+))?)?(?:,|$))");
std::smatch match;
std::string::const_iterator search_start(s.cbegin());
@@ -337,7 +340,7 @@ struct cmd_params {
std::vector<llama_split_mode> split_mode;
std::vector<int> main_gpu;
std::vector<bool> no_kv_offload;
std::vector<bool> flash_attn;
std::vector<llama_flash_attn_type> flash_attn;
std::vector<std::vector<ggml_backend_dev_t>> devices;
std::vector<std::vector<float>> tensor_split;
std::vector<std::vector<llama_model_tensor_buft_override>> tensor_buft_overrides;
@@ -376,12 +379,12 @@ static const cmd_params cmd_params_defaults = {
/* cpu_mask */ { "0x0" },
/* cpu_strict */ { false },
/* poll */ { 50 },
/* n_gpu_layers */ { 99 },
/* n_gpu_layers */ { -1 },
/* n_cpu_moe */ { 0 },
/* split_mode */ { LLAMA_SPLIT_MODE_LAYER },
/* main_gpu */ { 0 },
/* no_kv_offload */ { false },
/* flash_attn */ { false },
/* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO },
/* devices */ { {} },
/* tensor_split */ { std::vector<float>(llama_max_devices(), 0.0f) },
/* tensor_buft_overrides*/ { std::vector<llama_model_tensor_buft_override>{ { nullptr, nullptr } } },
@@ -451,7 +454,7 @@ static void print_usage(int /* argc */, char ** argv) {
printf(" -sm, --split-mode <none|layer|row|tensor> (default: %s)\n", join(transform_to_str(cmd_params_defaults.split_mode, split_mode_str), ",").c_str());
printf(" -mg, --main-gpu <i> (default: %s)\n", join(cmd_params_defaults.main_gpu, ",").c_str());
printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str());
printf(" -fa, --flash-attn <0|1> (default: %s)\n", join(cmd_params_defaults.flash_attn, ",").c_str());
printf(" -fa, --flash-attn <on|off|auto> (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str());
printf(" -dev, --device <dev0/dev1/...> (default: auto)\n");
printf(" -mmp, --mmap <0|1> (default: %s)\n", join(cmd_params_defaults.use_mmap, ",").c_str());
printf(" -dio, --direct-io <0|1> (default: %s)\n", join(cmd_params_defaults.use_direct_io, ",").c_str());
@@ -710,7 +713,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
invalid_param = true;
break;
}
auto p = parse_int_range(argv[i]);
auto p = parse_int_range(argv[i], /*allow_negative=*/true);
params.n_gpu_layers.insert(params.n_gpu_layers.end(), p.begin(), p.end());
} else if (arg == "-ncmoe" || arg == "--n-cpu-moe") {
if (++i >= argc) {
@@ -793,8 +796,27 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
invalid_param = true;
break;
}
auto p = string_split<bool>(argv[i], split_delim);
params.flash_attn.insert(params.flash_attn.end(), p.begin(), p.end());
auto p = string_split<std::string>(argv[i], split_delim);
std::vector<llama_flash_attn_type> types;
for (const auto & v : p) {
llama_flash_attn_type type;
if (common_arg_utils::is_truthy(v)) {
type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
} else if (common_arg_utils::is_falsey(v)) {
type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
} else if (common_arg_utils::is_autoy(v)) {
type = LLAMA_FLASH_ATTN_TYPE_AUTO;
} else {
invalid_param = true;
break;
}
types.push_back(type);
}
if (invalid_param) {
break;
}
params.flash_attn.insert(params.flash_attn.end(), types.begin(), types.end());
} else if (arg == "-mmp" || arg == "--mmap") {
if (++i >= argc) {
invalid_param = true;
@@ -1138,7 +1160,7 @@ struct cmd_params_instance {
llama_split_mode split_mode;
int main_gpu;
bool no_kv_offload;
bool flash_attn;
llama_flash_attn_type flash_attn;
std::vector<ggml_backend_dev_t> devices;
std::vector<float> tensor_split;
std::vector<llama_model_tensor_buft_override> tensor_buft_overrides;
@@ -1222,7 +1244,7 @@ struct cmd_params_instance {
cparams.type_k = type_k;
cparams.type_v = type_v;
cparams.offload_kqv = !no_kv_offload;
cparams.flash_attn_type = flash_attn ? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED;
cparams.flash_attn_type = flash_attn;
cparams.embeddings = embeddings;
cparams.op_offload = !no_op_offload;
cparams.swa_full = false;
@@ -1400,7 +1422,7 @@ struct test {
llama_split_mode split_mode;
int main_gpu;
bool no_kv_offload;
bool flash_attn;
llama_flash_attn_type flash_attn;
std::vector<ggml_backend_dev_t> devices;
std::vector<float> tensor_split;
std::vector<llama_model_tensor_buft_override> tensor_buft_overrides;
@@ -1522,10 +1544,10 @@ struct test {
field == "poll" || field == "model_size" || field == "model_n_params" || field == "n_gpu_layers" ||
field == "main_gpu" || field == "n_prompt" || field == "n_gen" || field == "n_depth" || field == "avg_ns" ||
field == "stddev_ns" || field == "no_op_offload" || field == "n_cpu_moe" ||
field == "fit_target" || field == "fit_min_ctx") {
field == "fit_target" || field == "fit_min_ctx" || field == "flash_attn") {
return INT;
}
if (field == "f16_kv" || field == "no_kv_offload" || field == "cpu_strict" || field == "flash_attn" ||
if (field == "f16_kv" || field == "no_kv_offload" || field == "cpu_strict" ||
field == "use_mmap" || field == "use_direct_io" || field == "embeddings" || field == "no_host") {
return BOOL;
}
@@ -1594,7 +1616,7 @@ struct test {
split_mode_str(split_mode),
std::to_string(main_gpu),
std::to_string(no_kv_offload),
std::to_string(flash_attn),
std::to_string((int) flash_attn),
devices_to_string(devices),
tensor_split_str,
tensor_buft_overrides_str,
@@ -1779,7 +1801,7 @@ struct markdown_printer : public printer {
return 6;
}
if (field == "flash_attn") {
return 2;
return 3;
}
if (field == "devices") {
return -12;