add args -k and --force

common/arg.cpp

@@ -737,6 +737,20 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
             params.gpu_mem = value; // in GiB
         }
     ).set_env("LLAMA_ARG_CUDA_MEM"));
+    add_opt(llama_arg(
+        {"-k", "--n-cycles"}, "N",
+        format("number of cycles to output one token (default: %d)", params.n_cycles),
+        [](gpt_params & params, int value) {
+            params.n_cycles = value;
+        }
+    ).set_env("LLAMA_ARG_K"));
+    add_opt(llama_arg(
+        {"--force"},
+        format("force to start prefetching after computation (default: %s)", params.force ? "true" : "false"),
+        [](gpt_params & params) {
+            params.force = true;
+        }
+    ).set_env("LLAMA_ARG_FORCE"));
 // #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.

common/common.cpp

@@ -1012,7 +1012,7 @@ static bool assign_layers_to_device(
     };
 
     // get valid factors
-    std::vector<int> valid_k = find_factors(n_layer);
+    std::vector<int> valid_k = cparams.n_cycles > 0 ? {(int)cparams.n_cycles} : find_factors(n_layer);
 
     // assign devices to sets M1, M2, M3, and M4
     // M1: devices running on macOS without Metal, and with insufficient memory
@@ -1801,8 +1801,10 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param
     cparams.n_world           = params.n_world;
     cparams.rank              = params.rank;
     cparams.prefetch          = params.prefetch;
+    cparams.force             = params.force;
     cparams.keep_out_in_metal = params.keep_out_in_metal;
     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);
 
     if (cparams.master_ip != nullptr) {

common/common.h

@@ -149,7 +149,9 @@ struct gpt_params {
     std::string next_node_ip      = "localhost"; // ip address of my next node
     bool    prefetch              = false; // prefetch layer weights
     bool    keep_out_in_metal     =  true; // whether to keep output weights in metal memory, true by default
+    bool    force                 = false; // force to start prefetching after computation
     int32_t gpu_mem               = 999.0; // gpu memory to use, in GiB
+    int32_t n_cycles              =     0; // number of cycles to output one token
     int32_t n_predict             =    -1; // new tokens to predict
     int32_t n_ctx                 =     0; // context size
     int32_t n_batch               =  2048; // logical batch size for prompt processing (must be >=32 to use BLAS)