tools: add llama-pshard-plan-params for token-tiered placement planning

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(
+        {"-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(
+        {"--pshard-tier-max"}, "N",
+        "max tier batch size for pshard planning (0 = auto, capped by n_ctx)",
+        [](common_params & params, int value) {
+            params.pshard_tier_max = 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

@@ -1444,6 +1444,7 @@ 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.max_vram_alloc  = params.max_vram_alloc;
     return mparams;
 }
 

common/common.h

@@ -681,6 +681,9 @@ 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
+
+    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 +990,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-backend.h

@@ -154,6 +154,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 +331,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 +351,27 @@ 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);
+
+    // 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_weight_copy_bytes;
+        size_t               input_activ_bytes;
+        size_t               writeback_bytes;
+        bool                 can_prefetch_weights;
+    };
+
+    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 {

ggml/src/ggml-alloc.c

@@ -623,6 +623,13 @@ static void ggml_gallocr_allocate_node(ggml_gallocr_t galloc, struct ggml_tensor
     GGML_ASSERT(buffer_id >= 0);
     struct hash_node * hn = ggml_gallocr_hash_get(galloc, node);
 
+    // reach a deferred writeback leaf through its consumer view
+    if (ggml_impl_is_view(node) && node->view_src != NULL &&
+        (node->view_src->flags & GGML_TENSOR_FLAG_WRITEBACK)) {
+        ggml_gallocr_allocate_node(galloc, node->view_src, buffer_id);
+        return;
+    }
+
     if (!ggml_gallocr_is_allocated(galloc, node) && !ggml_impl_is_view(node)) {
         hn->allocated = true;
         assert(hn->addr.offset == 0);
@@ -723,6 +730,7 @@ static void ggml_gallocr_alloc_graph_impl(ggml_gallocr_t galloc, struct ggml_cgr
     // these may be tensors that the application is not using in the graph, but may still want to allocate for other purposes
     for (int i = 0; i < graph->n_leafs; i++) {
         struct ggml_tensor * leaf = graph->leafs[i];
+        if (leaf->flags & GGML_TENSOR_FLAG_WRITEBACK) {continue;} // writeback leafs: deferred, allocated via consumer
         ggml_gallocr_allocate_node(galloc, leaf, get_node_buffer_id(leaf_buffer_ids, i));
     }
 
@@ -737,7 +745,10 @@ static void ggml_gallocr_alloc_graph_impl(ggml_gallocr_t galloc, struct ggml_cgr
         // itself is never used and should not be considered a dependency
         if (ggml_impl_is_view(node) && node->op != GGML_OP_NONE) {
             struct ggml_tensor * view_src = node->view_src;
-            ggml_gallocr_hash_get(galloc, view_src)->n_views += 1;
+            // FLAG_WRITEBACK roots: lifetime is n_children-only, fenced by scheduler keepalives
+            if (!(view_src->flags & GGML_TENSOR_FLAG_WRITEBACK)) {
+                ggml_gallocr_hash_get(galloc, view_src)->n_views += 1;
+            }
         }
 
         if (node->flags & GGML_TENSOR_FLAG_INPUT) {
@@ -759,6 +770,23 @@ static void ggml_gallocr_alloc_graph_impl(ggml_gallocr_t galloc, struct ggml_cgr
         }
     }
 
+#ifndef NDEBUG
+    // every FLAG_WRITEBACK leaf must have an OP_NONE keepalive consumer
+    for (int i = 0; i < graph->n_leafs; i++) {
+        struct ggml_tensor * leaf = graph->leafs[i];
+        if (!(leaf->flags & GGML_TENSOR_FLAG_WRITEBACK)) continue;
+        bool fenced = false;
+        for (int n = 0; n < graph->n_nodes && !fenced; n++) {
+            struct ggml_tensor * node = graph->nodes[n];
+            if (node->op != GGML_OP_NONE) continue;
+            for (int s = 0; s < GGML_MAX_SRC; s++) {
+                if (node->src[s] == leaf) { fenced = true; break; }
+            }
+        }
+        GGML_ASSERT(fenced);
+    }
+#endif
+
     // allocate tensors
     for (int i = 0; i < graph->n_nodes; i++) {
         struct ggml_tensor * node = graph->nodes[i];
@@ -804,12 +832,15 @@ static void ggml_gallocr_alloc_graph_impl(ggml_gallocr_t galloc, struct ggml_cgr
             if (p_hn->n_children == 0 && p_hn->n_views == 0) {
                 if (ggml_impl_is_view(parent)) {
                     struct ggml_tensor * view_src = parent->view_src;
-                    struct hash_node * view_src_hn = ggml_gallocr_hash_get(galloc, view_src);
-                    view_src_hn->n_views -= 1;
-                    AT_PRINTF("view_src %s: %d children, %d views\n",
-                        view_src->name, view_src_hn->n_children, view_src_hn->n_views);
-                    if (view_src_hn->n_views == 0 && view_src_hn->n_children == 0 && view_src_hn->allocated) {
-                        ggml_gallocr_free_node(galloc, view_src);
+                    // FLAG_WRITEBACK roots: skip n_views accounting, leaf freed on keepalive
+                    if (!(view_src->flags & GGML_TENSOR_FLAG_WRITEBACK)) {
+                        struct hash_node * view_src_hn = ggml_gallocr_hash_get(galloc, view_src);
+                        view_src_hn->n_views -= 1;
+                        AT_PRINTF("view_src %s: %d children, %d views\n",
+                            view_src->name, view_src_hn->n_children, view_src_hn->n_views);
+                        if (view_src_hn->n_views == 0 && view_src_hn->n_children == 0 && view_src_hn->allocated) {
+                            ggml_gallocr_free_node(galloc, view_src);
+                        }
                     }
                 }
                 else if (p_hn->allocated) {
@@ -951,9 +982,19 @@ void ggml_gallocr_reserve_n_size(
         ggml_gallocr_t galloc, struct ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids, size_t * sizes) {
     GGML_ASSERT(ggml_gallocr_reserve_n_impl(galloc, graph, node_buffer_ids, leaf_buffer_ids, /*no_alloc =*/ true));
     for (int i = 0; i < galloc->n_buffers; i++) {
+        // mirror get_buffer_size: skip slots whose talloc is already counted by an earlier slot
+        bool shared = false;
+        for (int j = 0; j < i; j++) {
+            if (galloc->buf_tallocs[j] == galloc->buf_tallocs[i]) {
+                shared = true;
+                break;
+            }
+        }
         sizes[i] = 0;
-        for (int c = 0; c < galloc->buf_tallocs[i]->n_chunks; c++) {
-            sizes[i] += galloc->buf_tallocs[i]->chunks[c]->max_size;
+        if (!shared) {
+            for (int c = 0; c < galloc->buf_tallocs[i]->n_chunks; c++) {
+                sizes[i] += galloc->buf_tallocs[i]->chunks[c]->max_size;
+            }
         }
     }
 }