unify profiling, planning and inference

common/arg.cpp

@@ -267,6 +267,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) {
@@ -277,7 +284,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});
     }
 }
 
@@ -859,11 +866,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.draft.tensor_buft_overrides.empty()) {
-        params.speculative.draft.tensor_buft_overrides.push_back({nullptr, nullptr});
+        params.speculative.draft.tensor_buft_overrides.push_back({nullptr, nullptr, -1});
     }
 
     if (!params.chat_template.empty() && !common_chat_verify_template(params.chat_template, params.use_jinja)) {
@@ -1846,6 +1853,27 @@ 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", "--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(
+        {"--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(
@@ -2615,7 +2643,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"));
@@ -2729,8 +2757,8 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     ).set_examples({LLAMA_EXAMPLE_FIT_PARAMS}).set_env("LLAMA_ARG_FIT_ESTIMATE"));
     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;
 
@@ -3939,7 +3967,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.draft.tensor_buft_overrides.push_back({buft_overrides_draft.back().c_str(), ggml_backend_cpu_buffer_type()});
+                params.speculative.draft.tensor_buft_overrides.push_back({buft_overrides_draft.back().c_str(), ggml_backend_cpu_buffer_type(), -1});
             }
         }
     ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE"));

common/common.cpp

@@ -1219,7 +1219,7 @@ common_init_result::common_init_result(common_params & params, bool model_only)
     auto mparams = common_model_params_to_llama(params);
     auto cparams = common_context_params_to_llama(params);
 
-    if (params.fit_params) {
+    auto fit_params = [&]() {
         COM_TRC("%s", "fitting params to device memory ...\n");
         COM_TRC("%s", "(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");
         common_fit_params(params.model.path.c_str(), &mparams, &cparams,
@@ -1228,6 +1228,30 @@ common_init_result::common_init_result(common_params & params, bool model_only)
             params.fit_params_target.data(),
             params.fit_params_min_ctx,
             params.verbosity >= LOG_LEVEL_DEBUG ? 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 = common_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);
+        common_fit_params_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__);
+            common_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);
@@ -1581,6 +1605,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;
 }
 
@@ -1618,6 +1644,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

@@ -742,6 +742,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
@@ -1080,7 +1084,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 };
 }
 
 //

common/fit.cpp

@@ -477,21 +477,24 @@ static void common_params_fit_impl(
             il0 += ngl_per_device[id].n_full();
             for (uint32_t il = il0; il < il0 + ngl_per_device[id].n_part; il++) {
                 if (itbo + 1 >= ntbo) {
-                    tensor_buft_overrides[itbo].pattern = nullptr;
-                    tensor_buft_overrides[itbo].buft    = nullptr;
+                    tensor_buft_overrides[itbo].pattern    = nullptr;
+                    tensor_buft_overrides[itbo].buft       = nullptr;
+                    tensor_buft_overrides[itbo].backend_id = -1;
                     itbo++;
                     mparams.tensor_buft_overrides = tensor_buft_overrides;
                     throw common_params_fit_exception("llama_max_tensor_buft_overrides() == "
                         + std::to_string(ntbo) + " is insufficient for model");
                 }
-                tensor_buft_overrides[itbo].pattern = get_overflow_pattern(il, il == il0 ? ngl_per_device[id].overflow_type : LAYER_FRACTION_MOE);
-                tensor_buft_overrides[itbo].buft = il == il0 ? overflow_bufts[id] : ggml_backend_cpu_buffer_type();
+                tensor_buft_overrides[itbo].pattern    = get_overflow_pattern(il, il == il0 ? ngl_per_device[id].overflow_type : LAYER_FRACTION_MOE);
+                tensor_buft_overrides[itbo].buft       = il == il0 ? overflow_bufts[id] : ggml_backend_cpu_buffer_type();
+                tensor_buft_overrides[itbo].backend_id = -1;
                 itbo++;
             }
             il0 += ngl_per_device[id].n_part;
         }
-        tensor_buft_overrides[itbo].pattern = nullptr;
-        tensor_buft_overrides[itbo].buft    = nullptr;
+        tensor_buft_overrides[itbo].pattern    = nullptr;
+        tensor_buft_overrides[itbo].buft       = nullptr;
+        tensor_buft_overrides[itbo].backend_id = -1;
         itbo++;
         mparams.tensor_buft_overrides = tensor_buft_overrides;
     };
@@ -527,8 +530,8 @@ static void common_params_fit_impl(
     if (hp_nex > 0) {
         const static std::string pattern_moe_all = "blk\\.\\d+\\.ffn_(up|down|gate_up|gate)_(ch|)exps"; // matches all MoE tensors
         ggml_backend_buffer_type_t cpu_buft = ggml_backend_cpu_buffer_type();
-        tensor_buft_overrides[0] = {pattern_moe_all.c_str(), cpu_buft};
-        tensor_buft_overrides[1] = {nullptr, nullptr};
+        tensor_buft_overrides[0] = {pattern_moe_all.c_str(), cpu_buft, -1};
+        tensor_buft_overrides[1] = {nullptr, nullptr, -1};
         mparams->tensor_buft_overrides = tensor_buft_overrides;
 
         LOG_TRC("%s: getting device memory data with all MoE tensors moved to system memory:\n", __func__);
@@ -549,7 +552,7 @@ static void common_params_fit_impl(
         }
 
         // reset
-        tensor_buft_overrides[0] = {nullptr, nullptr};
+        tensor_buft_overrides[0] = {nullptr, nullptr, -1};
         mparams->tensor_buft_overrides = tensor_buft_overrides;
     }
 
@@ -811,6 +814,47 @@ enum common_params_fit_status common_fit_params(
     return status;
 }
 
+llama_pshard_plan_registry * common_pshard_registry_create(uint32_t n_tier_max, uint32_t n_seq_max) {
+    return llama_pshard_registry_create(n_tier_max, n_seq_max);
+}
+
+void common_pshard_registry_free(llama_pshard_plan_registry * registry) {
+    llama_pshard_registry_free(registry);
+}
+
+// Adapter handed to the pshard planner in libllama, which cannot call into
+// common/ directly. Throws on failure, same as common_params_fit_impl.
+static void common_fit_pshard_baseline(
+        const char * path_model, llama_model_params * mparams, llama_context_params * cparams,
+        float * tensor_split, llama_model_tensor_buft_override * tensor_buft_overrides,
+        size_t * margins, uint32_t n_ctx_min, ggml_log_level log_level,
+        void * /*user_data*/) {
+    common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides,
+        margins, n_ctx_min, log_level);
+}
+
+void common_fit_params_pshard(
+        const char * path_model,
+        llama_model_params * mparams,
+        llama_context_params * cparams,
+        llama_model_tensor_buft_override * tensor_buft_overrides,
+        size_t max_vram_mb,
+        size_t fit_target_mb) {
+    llama_params_fit_pshard_inference(path_model, mparams, cparams, tensor_buft_overrides,
+        max_vram_mb, fit_target_mb);
+}
+
+void common_pshard_plan(
+        const char * path_model,
+        llama_model_params * mparams,
+        llama_context_params * cparams,
+        llama_model_tensor_buft_override * tensor_buft_overrides,
+        size_t max_vram_mb,
+        size_t fit_target_mb) {
+    llama_params_fit_pshard_planning(path_model, mparams, cparams, tensor_buft_overrides,
+        max_vram_mb, fit_target_mb, &common_fit_pshard_baseline, /*fit_fn_ud =*/ nullptr);
+}
+
 void common_memory_breakdown_print(const struct llama_context * ctx) {
     //const auto & devices = ctx->get_model().devices;
     const auto * model = llama_get_model(ctx);

common/fit.h

@@ -26,6 +26,47 @@ common_params_fit_status common_fit_params(
                            uint32_t   n_ctx_min,             // minimum context size to set when trying to reduce memory use
                      ggml_log_level   log_level);            // minimum log level to print during fitting, lower levels go to debug log
 
+// Pipelined-sharding (pshard) variant of common_fit_params.
+//
+// Loads the plan registry cached next to the model as
+// <model>.tensor_overrides.pshard_registry (written by llama-fit-params --pshard)
+// and fills tensor_buft_overrides for model loading. If no usable plan applies,
+// clears mparams->pshard and leaves baseline loading in place. This is the
+// RUNTIME path - it never plans, so startup stays fast.
+//
+// mparams->pshard_registry must be set by the caller (common_pshard_registry_create).
+//
+// Create/free the tier plan registry. The caller owns the pointer and assigns it
+// to mparams->pshard_registry before calling common_fit_params_pshard.
+struct llama_pshard_plan_registry * common_pshard_registry_create(uint32_t n_tier_max, uint32_t n_seq_max);
+void                                common_pshard_registry_free  (struct llama_pshard_plan_registry * registry);
+
+void common_fit_params_pshard(
+                         const char * path_model,
+                 llama_model_params * mparams,
+               llama_context_params * cparams,
+   llama_model_tensor_buft_override * tensor_buft_overrides,
+                             size_t   max_vram_mb,    // 0 = use actual free VRAM minus fit_target_mb
+                             size_t   fit_target_mb); // ignored when max_vram_mb > 0
+
+// PLANNER path (llama-fit-params --pshard).
+//
+// Probes VRAM for each strategy/tier combination, picks the best plan per tier and
+// writes/updates <model>.tensor_overrides.pshard_registry, then fills
+// tensor_buft_overrides. Loads any already-cached tiers first and only plans what is
+// missing. If everything already fits in VRAM it clears mparams->pshard and falls back
+// to baseline loading.
+//
+// Supplies the baseline params-fitting engine (common_fit_params) to the planner, which
+// lives in libllama and cannot call into common/ directly.
+void common_pshard_plan(
+                         const char * path_model,
+                 llama_model_params * mparams,
+               llama_context_params * cparams,
+   llama_model_tensor_buft_override * tensor_buft_overrides,
+                             size_t   max_vram_mb,
+                             size_t   fit_target_mb);
+
 // print estimated memory to stdout
 void common_fit_print(
                          const char * path_model,