ggml: allow prefetching tensor overrides

common/arg.cpp

@@ -2812,6 +2812,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.no_op_offload = !value;
         }
     ));
+    add_opt(common_arg(
+        {"-pw", "--prefetch-weights"}, "0|1",
+        string_format("prefetch weight transfers to overlap CPU->GPU copies with compute (default: %d)", (int) params.prefetch_weights),
+        [](common_params & params, int value) {
+            params.prefetch_weights = value != 0;
+        }
+    ));
     add_opt(common_arg(
         {"--lora"}, "FNAME",
         "path to LoRA adapter (use comma-separated values to load multiple adapters)",

common/common.cpp

@@ -1615,6 +1615,7 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.op_offload        = !params.no_op_offload;
     cparams.swa_full          = params.swa_full;
     cparams.kv_unified        = params.kv_unified;
+    cparams.prefetch_weights  = params.prefetch_weights;
 
     cparams.type_k = params.cache_type_k;
     cparams.type_v = params.cache_type_v;

common/common.h

@@ -580,6 +580,7 @@ struct common_params {
     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
+    bool prefetch_weights  = false; // prefetch weight transfers to overlap CPU->GPU copies with compute
     bool no_extra_bufts    = false; // disable extra buffer types (used for weight repacking)
     bool no_host           = false; // bypass host buffer allowing extra buffers to be used
 

ggml/include/ggml-backend.h

@@ -154,6 +154,8 @@ extern "C" {
         bool buffer_from_host_ptr;
         // event synchronization
         bool events;
+        // dedicated copy stream for compute/transfer overlap
+        bool copy_stream;
     };
 
     // all the device properties
@@ -351,6 +353,9 @@ extern "C" {
     // Set a callback to be called for each resulting node during graph compute
     GGML_API void                 ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data);
 
+    // Enable async weight prefetching to overlap CPU->GPU transfers with compute
+    GGML_API void                 ggml_backend_sched_set_prefetch_weights(ggml_backend_sched_t sched, bool enabled);
+
     //
     // Meta backend
     //

ggml/src/ggml-backend-meta.cpp

@@ -132,6 +132,7 @@ static void ggml_backend_meta_device_get_props(ggml_backend_dev_t dev, ggml_back
         /* .host_buffer           = */ false, // Not implemented.
         /* .buffer_from_host_ptr  = */ false, // Not implemented.
         /* .events                = */ false, // Not implemented.
+        /* .copy_stream           = */ false, // Not available
     };
     for (ggml_backend_dev_t simple_dev : meta_dev_ctx->simple_devs) {
         ggml_backend_dev_props tmp_props;

ggml/src/ggml-blas/ggml-blas.cpp

@@ -367,6 +367,7 @@ static void ggml_backend_blas_device_get_props(ggml_backend_dev_t dev, struct gg
         /* .host_buffer           = */ false,
         /* .buffer_from_host_ptr  = */ true,
         /* .events                = */ false,
+        /* .copy_stream           = */ false,
     };
 }
 

ggml/src/ggml-cann/ggml-cann.cpp

@@ -2815,6 +2815,7 @@ static void ggml_backend_cann_device_get_props(ggml_backend_dev_t dev, ggml_back
         /* .host_buffer           = */ host_buffer,
         /* .buffer_from_host_ptr  = */ false,
         /* .events                = */ true,
+        /* .copy_stream           = */ false,
     };
 }
 

ggml/src/ggml-cpu/ggml-cpu.cpp

@@ -397,6 +397,7 @@ static void ggml_backend_cpu_device_get_props(ggml_backend_dev_t dev, struct ggm
         /* .host_buffer           = */ false,
         /* .buffer_from_host_ptr  = */ true,
         /* .events                = */ false,
+        /* .copy_stream           = */ false,
     };
 }
 

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -4696,6 +4696,7 @@ static void ggml_backend_cuda_device_get_props(ggml_backend_dev_t dev, ggml_back
         /* .host_buffer           = */ host_buffer,
         /* .buffer_from_host_ptr  = */ false,
         /* .events                = */ events,
+        /* .copy_stream           = */ true,
     };
 }
 

ggml/src/ggml-hexagon/ggml-hexagon.cpp

@@ -3930,6 +3930,7 @@ static void ggml_backend_hexagon_device_get_props(ggml_backend_dev_t dev, struct
         /* .host_buffer           = */ (bool) opt_hostbuf,
         /* .buffer_from_host_ptr  = */ false,
         /* .events                = */ false,
+        /* .copy_stream           = */ false,
     };
 }
 

ggml/src/ggml-metal/ggml-metal.cpp

@@ -681,6 +681,7 @@ static void ggml_backend_metal_device_get_props(ggml_backend_dev_t dev, ggml_bac
         /* .host_buffer          = */ false,
         /* .buffer_from_host_ptr = */ true,
         /* .events               = */ true,
+        /* .copy_stream          = */ false,
     };
 }
 

ggml/src/ggml-opencl/ggml-opencl.cpp

@@ -10768,6 +10768,7 @@ static void ggml_backend_opencl_device_get_props(ggml_backend_dev_t dev, struct
         /* .host_buffer           = */ false,
         /* .buffer_from_host_ptr  = */ false,
         /* .events                = */ false,
+        /* .copy_stream           = */ false,
     };
 }
 

ggml/src/ggml-openvino/ggml-openvino.cpp

@@ -763,6 +763,7 @@ static void ggml_backend_openvino_device_get_props(ggml_backend_dev_t dev, ggml_
         /* .host_buffer           = */ false,
         /* .buffer_from_host_ptr  = */ false,
         /* .events                = */ false,
+        /* .copy_stream           = */ false,
     };
 }
 

ggml/src/ggml-rpc/ggml-rpc.cpp

@@ -1800,6 +1800,7 @@ static void ggml_backend_rpc_device_get_props(ggml_backend_dev_t dev, struct ggm
         /* .host_buffer           = */ false,
         /* .buffer_from_host_ptr  = */ false,
         /* .events                = */ false,
+        /* .copy_stream           = */ false,
     };
 }
 

ggml/src/ggml-sycl/ggml-sycl.cpp

@@ -5612,6 +5612,7 @@ static void ggml_backend_sycl_device_get_props(ggml_backend_dev_t dev, ggml_back
         /* .host_buffer           = */ host_buffer,
         /* .buffer_from_host_ptr  = */ false,
         /* .events                = */ events,
+        /* .copy_stream           = */ false,
     };
 }
 

ggml/src/ggml-virtgpu/ggml-backend-device.cpp

@@ -70,6 +70,7 @@ static void ggml_backend_remoting_device_get_props(ggml_backend_dev_t dev, ggml_
     props->caps.buffer_from_host_ptr = false;
     props->caps.async                = false;
     props->caps.events               = false;
+    props->caps.copy_stream          = false;
 }
 
 ggml_backend_buffer_type_t ggml_backend_remoting_device_get_buffer_type(ggml_backend_dev_t dev) {

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -17447,6 +17447,7 @@ static void ggml_backend_vk_device_get_props(ggml_backend_dev_t dev, struct ggml
         /* .host_buffer           = */ true,
         /* .buffer_from_host_ptr  = */ false,
         /* .events                = */ true,
+        /* .copy_stream           = */ false,
     };
 }
 

ggml/src/ggml-webgpu/ggml-webgpu.cpp

@@ -3829,6 +3829,7 @@ static void ggml_backend_webgpu_device_get_props(ggml_backend_dev_t dev, struct
         /* .host_buffer           = */ false,
         /* .buffer_from_host_ptr  = */ false,
         /* .events                = */ false,
+        /* .copy_stream           = */ false,
     };
 }
 

ggml/src/ggml-zdnn/ggml-zdnn.cpp

@@ -419,22 +419,22 @@ static enum ggml_status ggml_backend_zdnn_graph_compute(ggml_backend_t backend,
 }
 
 static ggml_backend_i ggml_backend_zdnn_i = {
-    /* .get_name               = */ ggml_backend_zdnn_name,
-    /* .free                   = */ ggml_backend_zdnn_free,
-    /* .set_tensor_async       = */ NULL,
-    /* .get_tensor_async       = */ NULL,
-    /* .set_tensor_2d_async    = */ NULL,
-    /* .get_tensor_2d_async    = */ NULL,
-    /* .cpy_tensor_async       = */ NULL,
-    /* .synchronize            = */ NULL,
-    /* .graph_plan_create      = */ NULL,
-    /* .graph_plan_free        = */ NULL,
-    /* .graph_plan_update      = */ NULL,
-    /* .graph_plan_compute     = */ NULL,
-    /* .graph_compute          = */ ggml_backend_zdnn_graph_compute,
-    /* .event_record           = */ NULL,
-    /* .event_wait             = */ NULL,
-    /* .graph_optimize         = */ NULL,
+    /* .get_name                = */ ggml_backend_zdnn_name,
+    /* .free                    = */ ggml_backend_zdnn_free,
+    /* .set_tensor_async        = */ NULL,
+    /* .get_tensor_async        = */ NULL,
+    /* .set_tensor_2d_async     = */ NULL,
+    /* .get_tensor_2d_async     = */ NULL,
+    /* .cpy_tensor_async        = */ NULL,
+    /* .synchronize             = */ NULL,
+    /* .graph_plan_create       = */ NULL,
+    /* .graph_plan_free         = */ NULL,
+    /* .graph_plan_update       = */ NULL,
+    /* .graph_plan_compute      = */ NULL,
+    /* .graph_compute           = */ ggml_backend_zdnn_graph_compute,
+    /* .event_record            = */ NULL,
+    /* .event_wait              = */ NULL,
+    /* .graph_optimize          = */ NULL,
 };
 
 static ggml_guid_t ggml_backend_zdnn_guid(void) {
@@ -487,7 +487,8 @@ static void ggml_backend_zdnn_device_get_props(ggml_backend_dev_t dev, ggml_back
         /* .async                = */ false,
         /* .host_buffer          = */ false,
         /* .buffer_from_host_ptr = */ false,
-        /* .events               = */ false
+        /* .events               = */ false,
+        /* .copy_stream          = */ false,
     };
 }
 

ggml/src/ggml-zendnn/ggml-zendnn.cpp

@@ -524,7 +524,8 @@ static void ggml_backend_zendnn_device_get_props(ggml_backend_dev_t dev, struct
         /* .async                = */ false,
         /* .host_buffer          = */ false,
         /* .buffer_from_host_ptr = */ true,
-        /* .events               = */ false
+        /* .events               = */ false,
+        /* .copy_stream          = */ false,
     };
 }
 

include/llama.h

@@ -396,6 +396,8 @@ extern "C" {
                           // try to disable when n_seq_max > 1 for improved performance when the sequences do not share a large prefix
                           // ref: https://github.com/ggml-org/llama.cpp/pull/14363
 
+        bool prefetch_weights; // prefetch weight transfers to overlap CPU->GPU copies with compute
+
         // [EXPERIMENTAL]
         // backend sampler chain configuration (make sure the caller keeps the sampler chains alive)
         // note: the samplers must be sampler chains (i.e. use llama_sampler_chain_init)

src/llama-context.cpp

@@ -265,8 +265,9 @@ llama_context::llama_context(
 
     cparams.n_outputs_max = params.n_outputs_max == 0 || llama_model_has_encoder(&model) ? cparams.n_batch : params.n_outputs_max;
 
-    cparams.op_offload = params.op_offload;
-    cparams.kv_unified = params.kv_unified;
+    cparams.op_offload       = params.op_offload;
+    cparams.kv_unified       = params.kv_unified;
+    cparams.prefetch_weights = params.prefetch_weights;
 
     // initialized later
     cparams.pipeline_parallel = false;
@@ -594,6 +595,7 @@ void llama_context::sched_reserve() {
     gf_res_reserve.reset(new llm_graph_result(max_nodes));
 
     sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload));
+    ggml_backend_sched_set_prefetch_weights(sched.get(), cparams.prefetch_weights);
 
     llama_memory_context_ptr mctx;
     if (memory) {
@@ -629,6 +631,7 @@ void llama_context::sched_reserve() {
                 LLAMA_LOG_WARN("%s: compute buffer allocation failed, retrying without pipeline parallelism\n", __func__);
                 cparams.pipeline_parallel = false;
                 sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, cparams.op_offload));
+                ggml_backend_sched_set_prefetch_weights(sched.get(), cparams.prefetch_weights);
                 gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get());
             }
             if (!gf) {
@@ -3499,6 +3502,7 @@ llama_context_params llama_context_default_params() {
         /*.op_offload                  =*/ true,
         /*.swa_full                    =*/ true,
         /*.kv_unified                  =*/ false,
+        /*.prefetch_weights            =*/ false,
         /*.sampler                     =*/ nullptr,
         /*.n_sampler                   =*/ 0,
         /*.ctx_other                   =*/ nullptr,

src/llama-cparams.h

@@ -52,6 +52,7 @@ struct llama_cparams {
     bool op_offload;
     bool kv_unified;
     bool pipeline_parallel;
+    bool prefetch_weights;
 
     std::vector<bool> embeddings_layer_inp; // [n_layer()] extract input embeddings for layer