llama-bench : add no-kv-offload parameter (#4812)

Author: slaren

examples/llama-bench/llama-bench.cpp

@@ -138,6 +138,7 @@ struct cmd_params {
     std::vector<int> n_threads;
     std::vector<int> n_gpu_layers;
     std::vector<int> main_gpu;
+    std::vector<bool> no_kv_offload;
     std::vector<bool> mul_mat_q;
     std::vector<std::array<float, LLAMA_MAX_DEVICES>> tensor_split;
     int reps;
@@ -155,6 +156,7 @@ static const cmd_params cmd_params_defaults = {
     /* n_threads     */ {get_num_physical_cores()},
     /* n_gpu_layers  */ {99},
     /* main_gpu      */ {0},
+    /* no_kv_offload */ {false},
     /* mul_mat_q     */ {true},
     /* tensor_split  */ {{}},
     /* reps          */ 5,
@@ -176,6 +178,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("  -ngl, --n-gpu-layers <n>          (default: %s)\n", join(cmd_params_defaults.n_gpu_layers, ",").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("  -mmq, --mul-mat-q <0|1>           (default: %s)\n", join(cmd_params_defaults.mul_mat_q, ",").c_str());
     printf("  -ts, --tensor_split <ts0/ts1/..>               \n");
     printf("  -r, --repetitions <n>             (default: %d)\n", cmd_params_defaults.reps);
@@ -309,6 +312,13 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
                 break;
             }
             params.main_gpu = split<int>(argv[i], split_delim);
+        } else if (arg == "-nkvo" || arg == "--no-kv-offload") {
+            if (++i >= argc) {
+                invalid_param = true;
+                break;
+            }
+            auto p = split<bool>(argv[i], split_delim);
+            params.no_kv_offload.insert(params.no_kv_offload.end(), p.begin(), p.end());
         } else if (arg == "-mmq" || arg == "--mul-mat-q") {
             if (++i >= argc) {
                 invalid_param = true;
@@ -383,6 +393,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
     if (params.type_v.empty())       { params.type_v = cmd_params_defaults.type_v; }
     if (params.n_gpu_layers.empty()) { params.n_gpu_layers = cmd_params_defaults.n_gpu_layers; }
     if (params.main_gpu.empty())     { params.main_gpu = cmd_params_defaults.main_gpu; }
+    if (params.no_kv_offload.empty()){ params.no_kv_offload = cmd_params_defaults.no_kv_offload; }
     if (params.mul_mat_q.empty())    { params.mul_mat_q = cmd_params_defaults.mul_mat_q; }
     if (params.tensor_split.empty()) { params.tensor_split = cmd_params_defaults.tensor_split; }
     if (params.n_threads.empty())    { params.n_threads = cmd_params_defaults.n_threads; }
@@ -400,6 +411,7 @@ struct cmd_params_instance {
     int n_threads;
     int n_gpu_layers;
     int main_gpu;
+    bool no_kv_offload;
     bool mul_mat_q;
     std::array<float, LLAMA_MAX_DEVICES> tensor_split;
 
@@ -428,6 +440,7 @@ struct cmd_params_instance {
         cparams.type_k = type_k;
         cparams.type_v = type_v;
         cparams.mul_mat_q = mul_mat_q;
+        cparams.offload_kqv = !no_kv_offload;
 
         return cparams;
     }
@@ -444,6 +457,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances_int(const cmd_p
     for (const auto & tk : params.type_k)
     for (const auto & tv : params.type_v)
     for (const auto & mmq : params.mul_mat_q)
+    for (const auto & nkvo : params.no_kv_offload)
     for (const auto & nt : params.n_threads) {
         cmd_params_instance instance = {
             /* .model        = */ m,
@@ -455,6 +469,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances_int(const cmd_p
             /* .n_threads    = */ nt,
             /* .n_gpu_layers = */ nl,
             /* .main_gpu     = */ mg,
+            /* .no_kv_offload= */ nkvo,
             /* .mul_mat_q    = */ mmq,
             /* .tensor_split = */ ts,
         };
@@ -476,6 +491,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
     for (const auto & tk : params.type_k)
     for (const auto & tv : params.type_v)
     for (const auto & mmq : params.mul_mat_q)
+    for (const auto & nkvo : params.no_kv_offload)
     for (const auto & nt : params.n_threads) {
         for (const auto & n_prompt : params.n_prompt) {
             if (n_prompt == 0) {
@@ -491,6 +507,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
                 /* .n_threads    = */ nt,
                 /* .n_gpu_layers = */ nl,
                 /* .main_gpu     = */ mg,
+                /* .no_kv_offload= */ nkvo,
                 /* .mul_mat_q    = */ mmq,
                 /* .tensor_split = */ ts,
             };
@@ -511,6 +528,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
                 /* .n_threads    = */ nt,
                 /* .n_gpu_layers = */ nl,
                 /* .main_gpu     = */ mg,
+                /* .no_kv_offload= */ nkvo,
                 /* .mul_mat_q    = */ mmq,
                 /* .tensor_split = */ ts,
             };
@@ -559,6 +577,7 @@ struct test {
     ggml_type type_v;
     int n_gpu_layers;
     int main_gpu;
+    bool no_kv_offload;
     bool mul_mat_q;
     std::array<float, LLAMA_MAX_DEVICES> tensor_split;
     int n_prompt;
@@ -579,6 +598,7 @@ struct test {
         type_v = inst.type_v;
         n_gpu_layers = inst.n_gpu_layers;
         main_gpu = inst.main_gpu;
+        no_kv_offload = inst.no_kv_offload;
         mul_mat_q = inst.mul_mat_q;
         tensor_split = inst.tensor_split;
         n_prompt = inst.n_prompt;
@@ -640,7 +660,8 @@ struct test {
             "cpu_info", "gpu_info",
             "model_filename", "model_type", "model_size", "model_n_params",
             "n_batch", "n_threads", "type_k", "type_v",
-            "n_gpu_layers", "main_gpu", "mul_mat_q", "tensor_split",
+            "n_gpu_layers", "main_gpu", "no_kv_offload",
+            "mul_mat_q", "tensor_split",
             "n_prompt", "n_gen", "test_time",
             "avg_ns", "stddev_ns",
             "avg_ts", "stddev_ts"
@@ -659,7 +680,7 @@ struct test {
             return INT;
         }
         if (field == "cuda" || field == "opencl" || field == "metal" || field == "gpu_blas" || field == "blas" ||
-            field == "f16_kv" || field == "mul_mat_q") {
+            field == "f16_kv" || field == "no_kv_offload" || field == "mul_mat_q") {
             return BOOL;
         }
         if (field == "avg_ts" || field == "stddev_ts") {
@@ -690,7 +711,8 @@ struct test {
             cpu_info, gpu_info,
             model_filename, model_type, std::to_string(model_size), std::to_string(model_n_params),
             std::to_string(n_batch), std::to_string(n_threads), ggml_type_name(type_k), ggml_type_name(type_v),
-            std::to_string(n_gpu_layers), std::to_string(main_gpu), std::to_string(mul_mat_q), tensor_split_str,
+            std::to_string(n_gpu_layers), std::to_string(main_gpu), std::to_string(no_kv_offload),
+            std::to_string(mul_mat_q), tensor_split_str,
             std::to_string(n_prompt), std::to_string(n_gen), test_time,
             std::to_string(avg_ns()), std::to_string(stdev_ns()),
             std::to_string(avg_ts()), std::to_string(stdev_ts())
@@ -851,6 +873,9 @@ struct markdown_printer : public printer {
         if (field == "mul_mat_q") {
             return "mmq";
         }
+        if (field == "no_kv_offload") {
+            return "nkvo";
+        }
         if (field == "tensor_split") {
             return "ts";
         }
@@ -885,6 +910,9 @@ struct markdown_printer : public printer {
         if (params.mul_mat_q.size() > 1 || params.mul_mat_q != cmd_params_defaults.mul_mat_q) {
             fields.push_back("mul_mat_q");
         }
+        if (params.no_kv_offload.size() > 1 || params.no_kv_offload != cmd_params_defaults.no_kv_offload) {
+            fields.push_back("no_kv_offload");
+        }
         if (params.tensor_split.size() > 1 || params.tensor_split != cmd_params_defaults.tensor_split) {
             fields.push_back("tensor_split");
         }