sched, ggml-meta : pipeline the host KV delivery under split mode tensor

ci/run.sh

@@ -312,6 +312,8 @@ function gg_run_test_llama_archs_tensor_split {
         GGML_CUDA_DEVICES=2 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
         GGML_CUDA_DEVICES=3 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
         GGML_CUDA_DEVICES=4 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
+        # the scheduler names a copy after its source, and 8 device names fill the name field
+        GGML_CUDA_DEVICES=8 ./build-ci-release/bin/test-llama-archs -s 1 -a llama 2>&1
     fi
 
     if [ ! -z ${GG_BUILD_METAL} ]; then
@@ -327,7 +329,7 @@ function gg_run_test_llama_archs_tensor_split {
 function gg_sum_test_llama_archs_tensor_split {
     gg_printf '### %s\n\n' "${ci}"
 
-    gg_printf 'Runs test-llama-archs with 1 to 4 devices\n'
+    gg_printf 'Runs test-llama-archs with 1 to 4 and 8 devices\n'
     gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
     gg_printf '```\n'
     gg_printf '%s\n' "$(cat $OUT/${ci}.log)"

common/arg.cpp

@@ -30,6 +30,7 @@
 #include <filesystem>
 #include <fstream>
 #include <list>
+#include <limits>
 #include <numeric>
 #include <regex>
 #include <set>
@@ -2429,6 +2430,40 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.kv_cpu_pinned = value;
         }
     ).set_env("LLAMA_ARG_KV_CPU_PINNED"));
+    add_opt(common_arg(
+        {"--kv-pipeline-depth"}, "N",
+        string_format("how many splits ahead the scheduler delivers a host-resident KV cache to the accelerator, so "
+                      "that the transfer runs while the previous split computes. 0 keeps the ordered path, where a "
+                      "decode token pays the transfer and the attention kernels in series; any other value turns the "
+                      "pipeline on, and 1 is the value that measures best. Only takes effect with a host-resident "
+                      "cache, e.g. --no-kv-offload or --kv-cpu-pinned, and costs (N + 2) * (largest staged split) of "
+                      "device memory (default: %d)", params.kv_pipeline_depth),
+        [](common_params & params, int value) {
+            if (value < 0 || value > LLAMA_KV_PIPELINE_DEPTH_MAX) {
+                throw std::invalid_argument(string_format("--kv-pipeline-depth must be between 0 and %d", LLAMA_KV_PIPELINE_DEPTH_MAX));
+            }
+            params.kv_pipeline_depth = value;
+        }
+    ).set_env("LLAMA_ARG_KV_PIPELINE_DEPTH"));
+    add_opt(common_arg(
+        {"--kv-pipeline-budget"}, "N",
+        string_format("hard cap, in MiB, on the device memory that pipelined delivery of a host-resident KV cache "
+                      "may use. A staging slot holds one attention layer's K or V over the whole context, so the "
+                      "requirement grows with the context; past this cap the scheduler declines and keeps the "
+                      "ordered path, so a host-resident cache never quietly trades away the device memory it exists "
+                      "to save. 0 removes the cap, %d is the largest accepted (default: %d)",
+                      LLAMA_KV_PIPELINE_BUDGET_MIB_MAX, params.kv_pipeline_budget_mib),
+        [](common_params & params, int value) {
+            constexpr size_t mib = 1024u*1024u;
+            if (value < 0 || value > LLAMA_KV_PIPELINE_BUDGET_MIB_MAX) {
+                throw std::invalid_argument(string_format("--kv-pipeline-budget must be between 0 and %d MiB", LLAMA_KV_PIPELINE_BUDGET_MIB_MAX));
+            }
+            if ((size_t) value > std::numeric_limits<size_t>::max()/mib) {
+                throw std::invalid_argument("--kv-pipeline-budget is out of range for this platform");
+            }
+            params.kv_pipeline_budget_mib = value;
+        }
+    ).set_env("LLAMA_ARG_KV_PIPELINE_BUDGET"));
     add_opt(common_arg(
         {"--recurrent-state-offload"},
         {"--no-recurrent-state-offload"},

common/common.cpp

@@ -1745,6 +1745,8 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.cb_eval_user_data = params.cb_eval_user_data;
     cparams.offload_kqv       = !params.no_kv_offload;
     cparams.kv_cpu_pinned     = params.kv_cpu_pinned;
+    cparams.kv_pipeline_depth      = (uint32_t) params.kv_pipeline_depth;
+    cparams.kv_pipeline_budget_mib = (uint32_t) params.kv_pipeline_budget_mib;
     cparams.recurrent_state_offload = params.recurrent_state_offload;
     cparams.kv_gpu_layers     = (uint32_t) std::max(0, params.kv_gpu_layers);
     cparams.phase_aware_workspace = params.phase_aware_workspace;

common/common.h

@@ -594,6 +594,8 @@ struct common_params {
     int32_t kv_gpu_layers  = 0;     // with no_kv_offload, keep this many attention KV layers device-resident
     bool phase_aware_workspace = false; // resize compute schedulers between prompt and generation phases
     bool live_context_workspace = false; // size supported attention workspaces from the padded live KV extent
+    int32_t kv_pipeline_depth = 0;  // splits of look-ahead for pipelined delivery of a host-resident KV cache (0 = off)
+    int32_t kv_pipeline_budget_mib = 128; // hard cap on the device memory that delivery may use (0 = uncapped)
     bool warmup            = true;  // warmup run
     bool check_tensors     = false; // validate tensor data
     bool no_op_offload     = false; // globally disable offload host tensor operations to device