llama: add pshard runtime for plan switching and streamed weights

common/arg.cpp

@@ -259,6 +259,13 @@ static void parse_tensor_buffer_overrides(const std::string & value, std::vector
         std::string tensor_name = override.substr(0, pos);
         std::string buffer_type = override.substr(pos + 1);
 
+        int32_t backend_id = -1;
+        auto colon_pos = buffer_type.rfind(':');
+        if (colon_pos != std::string::npos) {
+            backend_id = std::stoi(buffer_type.substr(colon_pos + 1));
+            buffer_type = buffer_type.substr(0, colon_pos);
+        }
+
         if (buft_list.find(buffer_type) == buft_list.end()) {
             printf("Available buffer types:\n");
             for (const auto & it : buft_list) {
@@ -269,7 +276,7 @@ static void parse_tensor_buffer_overrides(const std::string & value, std::vector
         // keep strings alive and avoid leaking memory by storing them in a static vector
         static std::list<std::string> buft_overrides;
         buft_overrides.push_back(tensor_name);
-        overrides.push_back({buft_overrides.back().c_str(), buft_list.at(buffer_type)});
+        overrides.push_back({buft_overrides.back().c_str(), buft_list.at(buffer_type), backend_id});
     }
 }
 
@@ -624,11 +631,11 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
     // pad tensor_buft_overrides for llama_params_fit:
     const size_t ntbo = llama_max_tensor_buft_overrides();
     while (params.tensor_buft_overrides.size() < ntbo) {
-        params.tensor_buft_overrides.push_back({nullptr, nullptr});
+        params.tensor_buft_overrides.push_back({nullptr, nullptr, -1});
     }
 
     if (!params.speculative.tensor_buft_overrides.empty()) {
-        params.speculative.tensor_buft_overrides.push_back({nullptr, nullptr});
+        params.speculative.tensor_buft_overrides.push_back({nullptr, nullptr, -1});
     }
 
     if (!params.chat_template.empty() && !common_chat_verify_template(params.chat_template, params.use_jinja)) {
@@ -1557,6 +1564,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.warmup = value;
         }
     ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_MTMD, LLAMA_EXAMPLE_EMBEDDING, LLAMA_EXAMPLE_RETRIEVAL, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_DEBUG}));
+    add_opt(common_arg(
+        {"-pshard"},
+        "enable pipelined sharding (weights on CPU, pipelined to GPU per split)",
+        [](common_params & params) {
+            params.pshard = true;
+        }
+    ));
+    add_opt(common_arg(
+        {"-mva", "--max-vram-alloc"}, "N",
+        "VRAM budget in MB for pshard (0 = use actual free VRAM minus --fit-target)",
+        [](common_params & params, int value) {
+            params.max_vram_alloc = value;
+        }
+    ));
     add_opt(common_arg(
         {"--spm-infill"},
         string_format(
@@ -2306,7 +2327,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
                 // keep strings alive and avoid leaking memory by storing them in a static vector
                 static std::list<std::string> buft_overrides;
                 buft_overrides.push_back(llm_ffn_exps_block_regex(i));
-                params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()});
+                params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type(), -1});
             }
         }
     ).set_env("LLAMA_ARG_N_CPU_MOE"));
@@ -2327,7 +2348,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             for (int i = 0; i < value; ++i) {
                 static std::list<std::string> buft_overrides_draft;
                 buft_overrides_draft.push_back(llm_ffn_exps_block_regex(i));
-                params.speculative.tensor_buft_overrides.push_back({buft_overrides_draft.back().c_str(), ggml_backend_cpu_buffer_type()});
+                params.speculative.tensor_buft_overrides.push_back({buft_overrides_draft.back().c_str(), ggml_backend_cpu_buffer_type(), -1});
             }
         }
     ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_N_CPU_MOE_DRAFT"));
@@ -2427,8 +2448,8 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     ).set_env("LLAMA_ARG_FIT"));
     add_opt(common_arg(
         { "-fitt", "--fit-target" }, "MiB0,MiB1,MiB2,...",
-        string_format("target margin per device for --fit, comma-separated list of values, "
-            "single value is broadcast across all devices, default: %zu", params.fit_params_target[0]/(1024*1024)),
+        string_format("target margin per device for --fit and pshard auto budget, comma-separated list of values, "
+                      "single value is broadcast across all devices, default: %zu", params.fit_params_target[0]/(1024*1024)),
         [](common_params & params, const std::string & value) {
             std::string arg_next = value;
 

common/common.cpp

@@ -1144,14 +1144,38 @@ common_init_result::common_init_result(common_params & params) :
     auto mparams = common_model_params_to_llama(params);
     auto cparams = common_context_params_to_llama(params);
 
-    if (params.fit_params) {
+    auto fit_params = [&]() {
         LOG_INF("%s: fitting params to device memory, for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on\n", __func__);
         llama_params_fit(params.model.path.c_str(), &mparams, &cparams,
             params.tensor_split,
             params.tensor_buft_overrides.data(),
             params.fit_params_target.data(),
             params.fit_params_min_ctx,
             params.verbosity >= 4 ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
+    };
+
+    if (params.pshard) {
+        LOG_INF("%s: pshard enabled, probing and loading plan cache\n", __func__);
+        params.tensor_buft_overrides.resize(4096);
+        mparams.pshard_registry = llama_pshard_registry_create(params.pshard_tier_max, cparams.n_seq_max);
+        const size_t fit_target_mb = params.fit_params_target.empty() ? 0 : params.fit_params_target[0] / (1024 * 1024);
+        llama_params_fit_pshard(params.model.path.c_str(), &mparams, &cparams,
+            params.tensor_buft_overrides.data(), params.max_vram_alloc, fit_target_mb);
+        if (!mparams.pshard) {
+            LOG_WRN("%s: pshard not active for this configuration\n", __func__);
+            llama_pshard_registry_free(mparams.pshard_registry);
+            mparams.pshard_registry = nullptr;
+            if (params.fit_params) {
+                fit_params();
+            }
+        } else {
+            params.n_batch  = (int32_t) cparams.n_batch;
+            params.n_ubatch = (int32_t) cparams.n_ubatch;
+            LOG_INF("%s: pshard runtime batch/ubatch set to selected cache_ubatch=%u\n",
+                __func__, cparams.n_ubatch);
+        }
+    } else if (params.fit_params) {
+        fit_params();
     }
 
     llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams);
@@ -1444,6 +1468,8 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
     mparams.progress_callback_user_data = params.load_progress_callback_user_data;
     mparams.no_alloc                    = params.no_alloc;
 
+    mparams.pshard          = params.pshard;
+    mparams.max_vram_alloc  = params.max_vram_alloc;
     return mparams;
 }
 
@@ -1479,6 +1505,7 @@ struct llama_context_params common_context_params_to_llama(const common_params &
 
     cparams.type_k = params.cache_type_k;
     cparams.type_v = params.cache_type_v;
+    cparams.pshard            = params.pshard;
 
     return cparams;
 }

common/common.h

@@ -681,6 +681,10 @@ struct common_params {
     llama_progress_callback load_progress_callback = NULL;
     void *                  load_progress_callback_user_data = NULL;
     bool no_alloc = false; // Don't allocate model buffers
+
+    bool     pshard           = false;
+    size_t   max_vram_alloc   = 0;
+    uint32_t pshard_tier_max  = 0;
 };
 
 // call once at the start of a program if it uses libcommon
@@ -987,7 +991,7 @@ inline std::string llm_ffn_exps_block_regex(int idx) {
 }
 
 inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() {
-    return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() };
+    return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type(), -1 };
 }
 
 //

ggml/include/ggml-alloc.h

@@ -73,6 +73,21 @@ GGML_API bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph
 
 GGML_API size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id);
 
+// per-chunk introspection (after reserve_n / reserve_n_size)
+GGML_API int    ggml_gallocr_get_n_chunks(ggml_gallocr_t galloc, int buffer_id);
+GGML_API size_t ggml_gallocr_get_chunk_max_size(ggml_gallocr_t galloc, int buffer_id, int chunk_id);
+
+// set an externally-owned buffer for a buffer slot (must be called before reserve)
+GGML_API void ggml_gallocr_set_buffer(ggml_gallocr_t galloc, int buffer_id, ggml_backend_buffer_t buffer, size_t alloc_offset, size_t alloc_size);
+
+// update the allocation range for an external buffer (e.g. on plan switch)
+GGML_API void ggml_gallocr_set_alloc_range(ggml_gallocr_t galloc, int buffer_id, size_t alloc_offset, size_t alloc_size);
+
+// save/restore allocator state (for plan switch without re-reserve)
+GGML_API void ggml_gallocr_get_state_sizes(ggml_gallocr_t galloc, size_t * node_size, size_t * leaf_size);
+GGML_API void ggml_gallocr_save_state(ggml_gallocr_t galloc, void * node_buf, void * leaf_buf, int * n_nodes, int * n_leafs);
+GGML_API void ggml_gallocr_restore_state(ggml_gallocr_t galloc, const void * node_buf, size_t node_size, const void * leaf_buf, size_t leaf_size, int n_nodes, int n_leafs);
+
 // Utils
 // Create a buffer and allocate all the tensors in a ggml_context
 // ggml_backend_alloc_ctx_tensors_from_buft_size returns the size of the buffer that would be allocated by ggml_backend_alloc_ctx_tensors_from_buft

ggml/include/ggml-backend.h

@@ -94,6 +94,7 @@ extern "C" {
     GGML_API void ggml_backend_tensor_set_2d(      struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
     GGML_API void ggml_backend_tensor_get_2d(const struct ggml_tensor * tensor,       void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
     GGML_API void ggml_backend_tensor_memset(      struct ggml_tensor * tensor,     uint8_t value, size_t offset, size_t size);
+    GGML_API void ggml_backend_tensor_memset_async(ggml_backend_t backend, struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size);
 
     GGML_API void ggml_backend_synchronize(ggml_backend_t backend);
 
@@ -154,6 +155,8 @@ extern "C" {
         bool buffer_from_host_ptr;
         // event synchronization
         bool events;
+        // separate copy stream for compute/transfer overlap
+        bool copy_stream;
     };
 
     // all the device properties
@@ -329,6 +332,7 @@ extern "C" {
     GGML_API size_t                     ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend);
 
     GGML_API void                 ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend);
+    GGML_API void                 ggml_backend_sched_set_tensor_backend_hint(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend);
     GGML_API ggml_backend_t       ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node);
 
     // Split graph without allocating it
@@ -348,6 +352,53 @@ 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);
 
+    // set an externally-owned buffer for a backend (see ggml_gallocr_set_buffer)
+    GGML_API void                 ggml_backend_sched_set_buffer(ggml_backend_sched_t sched, ggml_backend_t backend, ggml_backend_buffer_t buffer, size_t alloc_offset, size_t alloc_size);
+
+    // update the allocation range for an external buffer (see ggml_gallocr_set_alloc_range)
+    GGML_API void                 ggml_backend_sched_set_alloc_range(ggml_backend_sched_t sched, ggml_backend_t backend, size_t alloc_offset, size_t alloc_size);
+
+    // per-chunk introspection (see ggml_gallocr_get_n_chunks / get_chunk_max_size)
+    GGML_API int                  ggml_backend_sched_get_n_chunks(ggml_backend_sched_t sched, ggml_backend_t backend);
+    GGML_API size_t               ggml_backend_sched_get_chunk_max_size(ggml_backend_sched_t sched, ggml_backend_t backend, int chunk_id);
+
+    // save/restore gallocr + backend_id state (for plan switch without re-reserve)
+    GGML_API ggml_gallocr_t       ggml_backend_sched_get_galloc(ggml_backend_sched_t sched);
+    GGML_API void                 ggml_backend_sched_save_backend_ids(ggml_backend_sched_t sched, int * node_buf, int * leaf_buf, int * n_nodes, int * n_leafs);
+    GGML_API void                 ggml_backend_sched_restore_backend_ids(ggml_backend_sched_t sched, const int * node_buf, int n_nodes, const int * leaf_buf, int n_leafs);
+
+    // 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);
+
+    // Per-split callbacks for stateful tensors (e.g. KV cache, recurrent state).
+    typedef void (*ggml_backend_sched_split_cb)(struct ggml_tensor * tensor, ggml_backend_t backend, void * user_data);
+
+    GGML_API void ggml_backend_sched_set_split_callbacks(
+        ggml_backend_sched_t sched,
+        ggml_backend_sched_split_cb pre_compute,
+        ggml_backend_sched_split_cb post_compute,
+        void * user_data);
+
+    GGML_API void ggml_backend_sched_set_prefetch_cb(
+        ggml_backend_sched_t sched,
+        ggml_backend_sched_split_cb prefetch_cb);
+
+    // Register a tensor for pre/post-compute split callbacks.
+    GGML_API void ggml_backend_sched_add_writeback(ggml_backend_sched_t sched, struct ggml_tensor * tensor);
+
+    // Per-split info snapshot for timing prediction.
+    struct ggml_backend_sched_split_info {
+        struct ggml_cgraph * graph;
+        int                  backend_id;
+        size_t               input_weight_bytes;
+        size_t               input_activ_bytes;
+        size_t               writeback_bytes;
+    };
+
+    GGML_API bool ggml_backend_sched_get_split_info(
+        ggml_backend_sched_t sched, int split_id,
+        struct ggml_backend_sched_split_info * out);
+
     //
     // Meta backend
     //

ggml/include/ggml.h

@@ -639,7 +639,8 @@ extern "C" {
         GGML_TENSOR_FLAG_OUTPUT  =  2, // ...is an output for the GGML compute graph
         GGML_TENSOR_FLAG_PARAM   =  4, // ...contains trainable parameters
         GGML_TENSOR_FLAG_LOSS    =  8, // ...defines loss for numerical optimization (multiple loss tensors add up)
-        GGML_TENSOR_FLAG_COMPUTE = 16, // ...must be computed
+        GGML_TENSOR_FLAG_COMPUTE   = 16, // ...must be computed
+        GGML_TENSOR_FLAG_WRITEBACK = 32, // ...is stateful cache (KV/RS) needing GPU->CPU writeback
     };
 
     enum ggml_tri_type {