speedup: add arg --keep-out-in-cuda to run the output layer on CUDA

This commit is contained in:
Zonghang Li
2025-06-28 10:58:18 +04:00
committed by Li, Zonghang
parent e8d3e5a631
commit 1ea2d61a97
6 changed files with 66 additions and 16 deletions
+12
View File
@@ -775,6 +775,7 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
params.master_priority = std::stof(value);
}
).set_env("LLAMA_ARG_MASTER_PRIORITY"));
// #ifdef GGML_USE_METAL
// // warn: if the output layer weights are not kept in metal shared memory, its mmap-ed weight data
// // could be released by the OS and reloaded repeatedly, which causes additional disk I/O latency.
@@ -787,6 +788,17 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
// }
// ).set_env("LLAMA_ARG_KEEP_INP_OUT_IN_METAL"));
// #endif
#ifdef GGML_USE_CUDA
add_opt(llama_arg(
{"--keep-out-in-cuda"},
format("whether to compute the output layer on CUDA (default: %s)", params.keep_out_in_cuda ? "true" : "false"),
[](gpt_params & params) {
params.keep_out_in_cuda = true;
}
).set_env("LLAMA_ARG_KEEP_INP_OUT_IN_CUDA"));
#endif
add_opt(llama_arg(
{"-n", "--predict", "--n-predict"}, "N",
format("number of tokens to predict (default: %d, -1 = infinity, -2 = until context filled)", params.n_predict),
+13 -9
View File
@@ -2017,16 +2017,19 @@ struct llama_model_params llama_model_params_from_gpt_params(const gpt_params &
if (params.n_gpu_layers != -1) {
mparams.n_gpu_layers = params.n_gpu_layers;
}
mparams.n_world = params.n_world;
mparams.rank = params.rank;
mparams.rpc_servers = params.rpc_servers.c_str();
mparams.main_gpu = params.main_gpu;
mparams.split_mode = params.split_mode;
mparams.tensor_split = params.tensor_split;
mparams.use_mmap = params.use_mmap;
mparams.use_mlock = params.use_mlock;
mparams.check_tensors = params.check_tensors;
mparams.n_world = params.n_world;
mparams.rank = params.rank;
mparams.rpc_servers = params.rpc_servers.c_str();
mparams.main_gpu = params.main_gpu;
mparams.split_mode = params.split_mode;
mparams.tensor_split = params.tensor_split;
mparams.use_mmap = params.use_mmap;
mparams.use_mlock = params.use_mlock;
mparams.check_tensors = params.check_tensors;
mparams.keep_out_in_metal = params.keep_out_in_metal;
mparams.keep_out_in_cuda = params.keep_out_in_cuda;
std::copy(std::begin(params.n_layer_window), std::end(params.n_layer_window), mparams.n_layer_window);
if (params.kv_overrides.empty()) {
mparams.kv_overrides = NULL;
@@ -2068,6 +2071,7 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param
cparams.force = params.force;
cparams.master_priority = params.master_priority;
cparams.keep_out_in_metal = params.keep_out_in_metal;
cparams.keep_out_in_cuda = params.keep_out_in_cuda;
cparams.n_gpu_layers = params.n_gpu_layers;
cparams.n_cycles = params.n_cycles;
std::copy(std::begin(params.n_layer_window), std::end(params.n_layer_window), cparams.n_layer_window);
+1
View File
@@ -151,6 +151,7 @@ struct gpt_params {
uint32_t signal_port = 10000; // signal port for distributed inference
bool prefetch = false; // prefetch layer weights
bool keep_out_in_metal = true; // whether to keep output weights in metal memory, true by default
bool keep_out_in_cuda = false; // whether to run the output layer on CUDA, false by default
bool force = false; // force to start prefetching after computation
float master_priority = 1.01; // priority to assign workload to the master (set 1.01 to use master first, and 0.99 to offload to other devices)
int32_t gpu_mem = 999.0; // gpu memory to use, in GiB