Add --n-cpu-moe to llama_bench (#937)

* Add --n-cpu-moe to llama_banch

* Add usage

---------

Co-authored-by: Iwan Kawrakow <iwan.kawrakow@gmail.com>
This commit is contained in:
Kawrakow
2025-11-11 08:44:59 +02:00
committed by GitHub
parent 121ed91165
commit 25cd985c9b

View File

@@ -283,7 +283,7 @@ static const cmd_params cmd_params_defaults = {
/* type_k */ {GGML_TYPE_F16},
/* type_v */ {GGML_TYPE_F16},
/* n_threads */ {{cpu_get_num_math(), cpu_get_num_math()}},
/* n_gpu_layers */ {99},
/* n_gpu_layers */ {999},
/* rpc_servers */ {""},
/* split_mode */ {LLAMA_SPLIT_MODE_LAYER},
/* main_gpu */ {0},
@@ -330,6 +330,7 @@ static void print_usage(int /* argc */, char ** argv) {
printf(" -t, --threads <n> (default: %s)\n", join(cmd_params_defaults.n_threads, ",").c_str());
printf(" -tgb, --threads-gen-batch <n1,n2> (default: %s)\n", join(cmd_params_defaults.n_threads, ",").c_str());
printf(" -ngl, --n-gpu-layers <n> (default: %s)\n", join(cmd_params_defaults.n_gpu_layers, ",").c_str());
printf(" --n-cpu-moe <n> (default: none)\n");
printf(" -rpc, --rpc <rpc_servers> (default: %s)\n", join(cmd_params_defaults.rpc_servers, ",").c_str());
printf(" -sm, --split-mode <none|layer|row> (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());
@@ -428,6 +429,19 @@ bool parse_buft_overrides(const std::string& value, std::vector<llama_model_tens
}
return true;
}
bool add_cpu_buft_overrides(const char * arg, std::vector<llama_model_tensor_buft_override>& overrides) {
int n_layers = std::stoi(arg);
if (n_layers < 0) {
fprintf(stderr, "error: Invalid value for --n-cpu-moe: %s\n", arg);
return false;
}
for (int32_t l = 0; l < n_layers; ++l) {
std::string pattern = "blk\\." + std::to_string(l) + "\\.(ffn_(up|down|gate)_exps\\.weight)";
overrides.push_back({strdup(pattern.c_str()), ggml_backend_cpu_buffer_type()});
}
return true;
}
template<class T1, class T2>
std::vector<std::pair<T1,T2>> string_split_pairs(const std::string & str, char delim) {
std::vector<std::pair<T1,T2>> values;
@@ -800,6 +814,15 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
invalid_param = true;
break;
}
} else if (arg == "--n-cpu-moe") {
if (++i >= argc) {
invalid_param = true;
break;
}
if (!add_cpu_buft_overrides(argv[i], params.buft_overrides)) {
invalid_param = true;
break;
}
} else {
invalid_param = true;
break;