diff --git a/common/arg.cpp b/common/arg.cpp index 3d0183ed7026..86960c7495f2 100644 --- a/common/arg.cpp +++ b/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 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 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 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; diff --git a/common/common.cpp b/common/common.cpp index 16f78debd025..b3165ce24e5b 100644 --- a/common/common.cpp +++ b/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; } diff --git a/common/common.h b/common/common.h index 020b6a721ff9..da9d78db1c3b 100644 --- a/common/common.h +++ b/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 }; } // diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 4a8f6d4287da..7febc6db2e9b 100644 --- a/ggml/include/ggml-backend.h +++ b/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 // diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 11d3e8a81671..b35835e56da1 100644 --- a/ggml/include/ggml.h +++ b/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 { diff --git a/ggml/src/ggml-alloc.c b/ggml/src/ggml-alloc.c index a4b01ccf8a16..9fbc6d79aa15 100644 --- a/ggml/src/ggml-alloc.c +++ b/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; + } } } } diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 1a555bf2a4dc..94c8dbeaceb8 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -767,6 +767,8 @@ struct ggml_backend_sched_split { int i_end; struct ggml_tensor * inputs[GGML_SCHED_MAX_SPLIT_INPUTS]; int n_inputs; + struct ggml_tensor * writeback[GGML_SCHED_MAX_SPLIT_INPUTS]; + int n_writeback; // graph view of this split struct ggml_cgraph graph; }; @@ -784,6 +786,8 @@ struct ggml_backend_sched { // hash map of the nodes in the graph struct ggml_hash_set hash_set; int * hv_tensor_backend_ids; // [hash_set.size] + bool * hv_tensor_usr; // [hash_set.size] -- true if set by set_tensor_backend + bool * hv_tensor_writeback; // [hash_set.size] -- true if registered for pre/post-compute callbacks struct ggml_tensor ** hv_tensor_copies; // [hash_set.size][n_backends][n_copies] int * node_backend_ids; // [graph_size] @@ -813,6 +817,15 @@ struct ggml_backend_sched { ggml_backend_sched_eval_callback callback_eval; void * callback_eval_user_data; + bool prefetch_weights; + + // redirect_target[backend_id] is -1 for regular compute backends + // alias backends on the same physical device point to the regular compute backend id + int redirect_target[GGML_SCHED_MAX_BACKENDS]; + bool has_redirects; + + ggml_backend_t copy_backends[GGML_SCHED_MAX_BACKENDS]; + char * context_buffer; size_t context_buffer_size; @@ -848,6 +861,23 @@ static int ggml_backend_sched_backend_from_buffer(ggml_backend_sched_t sched, co return -1; } + // respect explicit set_tensor_backend assignments (walk view_src chain to root) + { + const struct ggml_tensor * cur = tensor; + while (cur) { + size_t idx = ggml_hash_find(&sched->hash_set, cur); + if (idx != SIZE_MAX && idx < sched->hash_set.size) { + int explicit_id = sched->hv_tensor_backend_ids[idx]; + if (explicit_id >= 0 && explicit_id < sched->n_backends) { + if (ggml_backend_supports_op(sched->backends[explicit_id], op)) { + return explicit_id; + } + } + } + cur = cur->view_src; + } + } + // find highest prio backend that supports the buffer type and the op for (int i = 0; i < sched->n_backends; i++) { if (ggml_backend_supports_buft(sched->backends[i], buffer->buft) && @@ -913,7 +943,7 @@ static int ggml_backend_sched_backend_id_from_cur(ggml_backend_sched_t sched, st } // skip ROPE since the rope freqs tensor is too small to choose a backend based on it // not an ideal solution - if (tensor->op != GGML_OP_ROPE && src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + if (tensor->op != GGML_OP_ROPE && src->buffer != NULL && (src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS || !ggml_backend_buffer_is_host(src->buffer))) { int src_backend_id = ggml_backend_sched_backend_from_buffer(sched, src, tensor); // check if a backend with higher prio wants to offload the op if (sched->op_offload && src_backend_id == sched->n_backends - 1 && ggml_backend_buffer_is_host(src->buffer)) { @@ -1183,8 +1213,10 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra } } } - } else { + } else if (sched->redirect_target[*node_backend_id] < 0 && !sched->hv_tensor_usr[hash_id(node)]) { // assigned node: upgrade to higher prio backend if possible + // skip for user-assigned nodes (set_tensor_backend) + // skip for nodes on redirected backends for (int b = 0; b < *node_backend_id; b++) { if (sched->bufts[b] == sched->bufts[*node_backend_id] && ggml_backend_supports_op(sched->backends[b], node)) { bool supported = true; @@ -1255,6 +1287,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra } split->i_start = 0; split->n_inputs = 0; + split->n_writeback = 0; int cur_backend_id = split->backend_id; for (; i < graph->n_nodes; i++) { struct ggml_tensor * node = graph->nodes[i]; @@ -1311,6 +1344,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra split->backend_id = node_backend_id; split->i_start = i; split->n_inputs = 0; + split->n_writeback = 0; cur_backend_id = node_backend_id; } @@ -1347,7 +1381,17 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra } } - if (src_backend_id != cur_backend_id && !ggml_backend_sched_buffer_supported(sched, src, cur_backend_id)) { + bool need_copy = src_backend_id != cur_backend_id && !ggml_backend_sched_buffer_supported(sched, src, cur_backend_id); + + if (!need_copy && sched->has_redirects && + cur_backend_id != sched->n_backends - 1 && + src->buffer != NULL && + ggml_backend_buffer_get_usage(src->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && + ggml_backend_buffer_is_host(src->buffer)) { + need_copy = true; + } + + if (need_copy) { // create a copy of the input in the split's backend if (tensor_id_copy(src_id, cur_backend_id, 0) == NULL) { ggml_backend_t backend = sched->backends[cur_backend_id]; @@ -1365,8 +1409,33 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra GGML_ASSERT(n_inputs < GGML_SCHED_MAX_SPLIT_INPUTS); split->inputs[n_inputs] = src; } + node->src[j] = tensor_id_copy(src_id, cur_backend_id, sched->cur_copy); } + + // collect registered writeback tensors for pre/post-compute callbacks + // walk view_src to root, check hv_tensor_writeback hash + { + struct ggml_tensor * wb = NULL; + struct ggml_tensor * root = src; + while (root) { + size_t root_id = ggml_hash_find(&sched->hash_set, root); + if (root_id != SIZE_MAX && root_id < sched->hash_set.size && sched->hv_tensor_writeback[root_id]) { + wb = root; + break; + } + root = root->view_src; + } + if (wb && split->n_writeback < GGML_SCHED_MAX_SPLIT_INPUTS) { + bool found = false; + for (int w = 0; w < split->n_writeback; w++) { + if (split->writeback[w] == wb) { found = true; break; } + } + if (!found) { + split->writeback[split->n_writeback++] = wb; + } + } + } } } split->i_end = graph->n_nodes; @@ -1388,7 +1457,11 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra sched->prev_leaf_backend_ids = tmp; } - int graph_size = std::max(graph->n_nodes, graph->n_leafs) + sched->n_splits*GGML_SCHED_MAX_SPLIT_INPUTS*2*sched->n_copies; + const int nodes_per_input = sched->prefetch_weights ? 3 : 2; + const int nodes_per_writeback = 3; + int graph_size = std::max(graph->n_nodes, graph->n_leafs) + + sched->n_splits*GGML_SCHED_MAX_SPLIT_INPUTS*nodes_per_input*sched->n_copies + + sched->n_splits*GGML_SCHED_MAX_SPLIT_INPUTS*nodes_per_writeback; // remember the actual graph_size for performing reallocation checks later [GGML_SCHED_DEBUG_REALLOC] sched->debug_prev_graph_size = sched->debug_graph_size; @@ -1433,11 +1506,83 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra graph_copy->nodes[graph_copy->n_nodes++] = input_cpy; } + // reserve writeback leaf slot at split start (avoid aliasing with activations freed mid-split) + for (int w = 0; w < split->n_writeback; w++) { + assert(graph_copy->size > graph_copy->n_nodes); + struct ggml_tensor * prealloc = ggml_view_tensor(sched->ctx, split->writeback[w]); + prealloc->op = GGML_OP_NONE; + prealloc->src[0] = split->writeback[w]; + sched->node_backend_ids[graph_copy->n_nodes] = split->backend_id; + graph_copy->nodes[graph_copy->n_nodes++] = prealloc; + } + + // prefetch: reserve immediate next sharded split's prefetch destinations + // (race-protect copy stream vs current compute, do not scan past CPU splits) + if (sched->prefetch_weights) { + struct ggml_backend_sched_split * next_gpu = NULL; + if (i + 1 < sched->n_splits) { + struct ggml_backend_sched_split * candidate = &sched->splits[i + 1]; + // only prefetch into redirected backends + // the main compute backend should use the normal input copy i.e. no prefetch + if (sched->redirect_target[candidate->backend_id] >= 0 && + sched->copy_backends[candidate->backend_id] != NULL) { + bool has_host_weights = false; + for (int j = 0; j < candidate->n_inputs; j++) { + if (candidate->inputs[j]->buffer != NULL && + ggml_backend_buffer_get_usage(candidate->inputs[j]->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && + ggml_backend_buffer_is_host(candidate->inputs[j]->buffer)) { + has_host_weights = true; + break; + } + } + if (has_host_weights) { + next_gpu = candidate; + } + } + } + if (next_gpu != NULL) { + for (int j = 0; j < next_gpu->n_inputs; j++) { + struct ggml_tensor * next_input = next_gpu->inputs[j]; + if (next_input->buffer != NULL && + ggml_backend_buffer_get_usage(next_input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && + ggml_backend_buffer_is_host(next_input->buffer)) { + const size_t id = hash_id(next_input); + struct ggml_tensor * next_cpy = tensor_id_copy(id, next_gpu->backend_id, sched->cur_copy); + assert(graph_copy->size > graph_copy->n_nodes); + struct ggml_tensor * keepalive = ggml_view_tensor(sched->ctx, next_cpy); + keepalive->src[0] = next_cpy; + sched->node_backend_ids[graph_copy->n_nodes] = next_gpu->backend_id; + graph_copy->nodes[graph_copy->n_nodes++] = keepalive; + } + } + + // reserve next split's writeback slot (double-buffer: prefetch copy stream vs. current compute) + for (int w = 0; w < next_gpu->n_writeback; w++) { + assert(graph_copy->size > graph_copy->n_nodes); + struct ggml_tensor * keepalive = ggml_view_tensor(sched->ctx, next_gpu->writeback[w]); + keepalive->op = GGML_OP_NONE; + keepalive->src[0] = next_gpu->writeback[w]; + sched->node_backend_ids[graph_copy->n_nodes] = next_gpu->backend_id; + graph_copy->nodes[graph_copy->n_nodes++] = keepalive; + } + } + } + for (int j = split->i_start; j < split->i_end; j++) { assert(graph_copy->size > graph_copy->n_nodes); sched->node_backend_ids[graph_copy->n_nodes] = tensor_backend_id(graph->nodes[j]); graph_copy->nodes[graph_copy->n_nodes++] = graph->nodes[j]; } + + // n_children fence for writeback (held through post_compute, required by FLAG_WRITEBACK semantics) + for (int w = 0; w < split->n_writeback; w++) { + assert(graph_copy->size > graph_copy->n_nodes); + struct ggml_tensor * keepalive = ggml_view_tensor(sched->ctx, split->writeback[w]); + keepalive->op = GGML_OP_NONE; + keepalive->src[0] = split->writeback[w]; + sched->node_backend_ids[graph_copy->n_nodes] = split->backend_id; + graph_copy->nodes[graph_copy->n_nodes++] = keepalive; + } } if (sched->n_copies > 1) { @@ -1542,13 +1687,20 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s for (int split_id = 0; split_id < sched->n_splits; split_id++) { struct ggml_backend_sched_split * split = &splits[split_id]; int split_backend_id = split->backend_id; + int copy_backend_id = split_backend_id; ggml_backend_t split_backend = sched->backends[split_backend_id]; + if (sched->redirect_target[split_backend_id] >= 0) { + ggml_backend_synchronize(sched->backends[split_backend_id]); + split_backend_id = sched->redirect_target[split_backend_id]; + split_backend = sched->backends[split_backend_id]; + } + // copy the input tensors to the split backend for (int input_id = 0; input_id < split->n_inputs; input_id++) { ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]); struct ggml_tensor * input = split->inputs[input_id]; - struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy); + struct ggml_tensor * input_cpy = tensor_copy(input, copy_backend_id, sched->cur_copy); if (input->flags & GGML_TENSOR_FLAG_INPUT) { // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done @@ -1661,7 +1813,12 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } else { ggml_backend_synchronize(split_backend); } - ggml_backend_tensor_copy(input, input_cpy); + if (sched->has_redirects && input->data != NULL && + ggml_backend_buffer_is_host(input->buffer)) { + ggml_backend_tensor_set_async(split_backend, input_cpy, input->data, 0, ggml_nbytes(input)); + } else { + ggml_backend_tensor_copy(input, input_cpy); + } } } } @@ -1747,10 +1904,14 @@ ggml_backend_sched_t ggml_backend_sched_new( // FIXME: needs to be size*2 to account for leafs (do it in graph_split instead) sched->hash_set = ggml_hash_set_new(graph_size); sched->hv_tensor_backend_ids = (int *) malloc(sched->hash_set.size * sizeof(sched->hv_tensor_backend_ids[0])); + sched->hv_tensor_usr = (bool *) calloc(sched->hash_set.size, sizeof(bool)); + sched->hv_tensor_writeback = (bool *) calloc(sched->hash_set.size, sizeof(bool)); sched->hv_tensor_copies = (ggml_tensor **) malloc(sched->hash_set.size * sched->n_backends * sched->n_copies * sizeof(struct ggml_tensor *)); const size_t ggml_sched_max_splits = graph_size; // at most there is one split for each node in the graph - const size_t nodes_size = graph_size + ggml_sched_max_splits*GGML_SCHED_MAX_SPLIT_INPUTS*2; + // per-split synthetic tensors: 4 per input (dep, cpy, prefetch keepalive, keepalive) + 3 per writeback + const size_t nodes_per_split = GGML_SCHED_MAX_SPLIT_INPUTS*(3+3); + const size_t nodes_size = graph_size + ggml_sched_max_splits*nodes_per_split; sched->node_backend_ids = (int *) calloc(nodes_size, sizeof(sched->node_backend_ids[0])); sched->leaf_backend_ids = (int *) calloc(nodes_size, sizeof(sched->leaf_backend_ids[0])); sched->prev_node_backend_ids = (int *) calloc(nodes_size, sizeof(sched->prev_node_backend_ids[0])); @@ -1759,7 +1920,7 @@ ggml_backend_sched_t ggml_backend_sched_new( sched->debug_graph_size = 0; sched->debug_prev_graph_size = 0; - sched->context_buffer_size = ggml_sched_max_splits*GGML_SCHED_MAX_SPLIT_INPUTS*2*sizeof(struct ggml_tensor) + ggml_graph_overhead_custom(graph_size, false); + sched->context_buffer_size = ggml_sched_max_splits*nodes_per_split*sizeof(struct ggml_tensor) + ggml_graph_overhead_custom(graph_size, false); sched->context_buffer = (char *) malloc(sched->context_buffer_size); const int initial_splits_capacity = 16; @@ -1778,6 +1939,30 @@ ggml_backend_sched_t ggml_backend_sched_new( } } + // auto-detect same-device backends for compute redirect + sched->has_redirects = false; + for (int i = 0; i < n_backends; i++) { + sched->redirect_target[i] = -1; + sched->copy_backends[i] = NULL; + ggml_backend_dev_t dev_i = ggml_backend_get_device(backends[i]); + if (ggml_backend_dev_type(dev_i) == GGML_BACKEND_DEVICE_TYPE_CPU) { + continue; + } + for (int j = 0; j < i; j++) { + if (ggml_backend_get_device(backends[j]) == dev_i) { + sched->redirect_target[i] = j; + sched->has_redirects = true; + break; + } + } + // create copy backend + events for compute/transfer overlap + ggml_backend_dev_props props; + ggml_backend_dev_get_props(dev_i, &props); + if (props.caps.copy_stream) { + sched->copy_backends[i] = ggml_backend_dev_init(dev_i, NULL); + } + } + sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends); sched->op_offload = op_offload; @@ -1794,12 +1979,17 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) { for (int c = 0; c < sched->n_copies; c++) { ggml_backend_event_free(sched->events[b][c]); } + if (sched->copy_backends[b] != NULL) { + ggml_backend_free(sched->copy_backends[b]); + } } ggml_gallocr_free(sched->galloc); ggml_free(sched->ctx); ggml_hash_set_free(&sched->hash_set); free(sched->splits); free(sched->hv_tensor_backend_ids); + free(sched->hv_tensor_usr); + free(sched->hv_tensor_writeback); free(sched->hv_tensor_copies); free(sched->node_backend_ids); free(sched->leaf_backend_ids); @@ -1817,6 +2007,8 @@ void ggml_backend_sched_reset(ggml_backend_sched_t sched) { if (!sched->is_reset) { ggml_hash_set_reset(&sched->hash_set); memset(sched->hv_tensor_backend_ids, -1, sched->hash_set.size * sizeof(sched->hv_tensor_backend_ids[0])); + memset(sched->hv_tensor_usr, 0, sched->hash_set.size * sizeof(bool)); + memset(sched->hv_tensor_writeback, 0, sched->hash_set.size * sizeof(bool)); memset(sched->hv_tensor_copies, 0, sched->hash_set.size * sched->n_backends * sched->n_copies * sizeof(struct ggml_tensor *)); sched->is_reset = true; } @@ -1913,6 +2105,95 @@ void ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backe sched->callback_eval_user_data = user_data; } +void ggml_backend_sched_set_prefetch_weights(ggml_backend_sched_t sched, bool enabled) { + GGML_ASSERT(sched); + sched->prefetch_weights = enabled; +} + +void ggml_backend_sched_add_writeback(ggml_backend_sched_t sched, struct ggml_tensor * tensor) { + GGML_ASSERT(sched); + size_t id = ggml_hash_find_or_insert(&sched->hash_set, tensor); + GGML_ASSERT(id != SIZE_MAX); + sched->hv_tensor_writeback[id] = true; + + // defer upfront alloc; keepalive view makes the leaf split-scoped + tensor->flags |= GGML_TENSOR_FLAG_WRITEBACK; +} + +bool ggml_backend_sched_get_split_info( + ggml_backend_sched_t sched, int split_id, + struct ggml_backend_sched_split_info * out) { + GGML_ASSERT(sched && out); + if (split_id < 0 || split_id >= sched->n_splits) return false; + + struct ggml_backend_sched_split * s = &sched->splits[split_id]; + out->graph = &s->graph; + out->backend_id = s->backend_id; + + auto effective_weight_copy_bytes = [&](struct ggml_tensor * inp) { + const size_t full_size = ggml_nbytes(inp); + + if (s->graph.n_nodes <= 0) { + return full_size; + } + + struct ggml_tensor * node = s->graph.nodes[0]; + if (node->op != GGML_OP_MUL_MAT_ID) { + return full_size; + } + + struct ggml_tensor * inp_cpy = tensor_copy(inp, s->backend_id, sched->cur_copy); + if (node->src[0] != inp_cpy) { + return full_size; + } + + const int64_t n_expert = inp->ne[2]; + if (n_expert <= 0 || node->src[2] == NULL) { + return full_size; + } + + const size_t expert_size = inp->nb[2]; + const int64_t n_routes = ggml_nelements(node->src[2]); + int64_t n_used = std::min(n_expert, std::max(1, n_routes)); + const int64_t n_expert_used = node->ne[1]; + // ids contains one route per token/top-k slot, not unique experts. + // For larger batches, estimate unique experts instead of charging every route. + if (n_used > n_expert_used) { + n_used = std::min(n_expert, std::max(n_expert_used, (n_expert + 3) / 4)); + } + const size_t padding = std::min(expert_size, 512); + + return std::min(full_size, (size_t) n_used * (expert_size + padding)); + }; + + out->input_weight_bytes = 0; + out->input_weight_copy_bytes = 0; + out->input_activ_bytes = 0; + for (int j = 0; j < s->n_inputs; j++) { + struct ggml_tensor * inp = s->inputs[j]; + if (inp->buffer != NULL && + ggml_backend_buffer_get_usage(inp->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && + ggml_backend_buffer_is_host(inp->buffer)) { + out->input_weight_bytes += ggml_nbytes(inp); + out->input_weight_copy_bytes += effective_weight_copy_bytes(inp); + } else { + out->input_activ_bytes += ggml_nbytes(inp); + } + } + + out->writeback_bytes = 0; + for (int w = 0; w < s->n_writeback; w++) { + out->writeback_bytes += ggml_nbytes(s->writeback[w]); + } + + out->can_prefetch_weights = + sched->prefetch_weights && + sched->redirect_target[s->backend_id] >= 0 && + sched->copy_backends[s->backend_id] != NULL; + + return true; +} + int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched) { GGML_ASSERT(sched); return sched->n_splits; @@ -1955,8 +2236,21 @@ void ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct gg int backend_index = ggml_backend_sched_backend_id(sched, backend); GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); tensor_backend_id(node) = backend_index; + sched->hv_tensor_usr[hash_id(node)] = true; SET_CAUSE(node, "usr"); - sched->is_reset = false; + + // not clearing is_reset here: split_graph() already sets is_reset=false, + // guaranteeing cleanup between graphs. clearing it here would cause + // reserve_size's sched_reset to discard these assignments before + // split_graph reads them. +} + +void ggml_backend_sched_set_tensor_backend_hint(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend) { + GGML_ASSERT(sched); + int backend_index = ggml_backend_sched_backend_id(sched, backend); + GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); + tensor_backend_id(node) = backend_index; + SET_CAUSE(node, "hint"); } ggml_backend_t ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node) { diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 49f840be2076..ff8bc27a1532 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/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, }; } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 3113de017f09..5c99c0c79662 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4658,6 +4658,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, }; } diff --git a/include/llama.h b/include/llama.h index ac267b5089ab..6044f12c7180 100644 --- a/include/llama.h +++ b/include/llama.h @@ -281,6 +281,7 @@ extern "C" { struct llama_model_tensor_buft_override { const char * pattern; ggml_backend_buffer_type_t buft; + int32_t backend_id; }; struct llama_model_params { @@ -319,6 +320,11 @@ extern "C" { bool use_extra_bufts; // use extra buffer types (used for weight repacking) bool no_host; // bypass host buffer allowing extra buffers to be used bool no_alloc; // only load metadata and simulate memory allocations + bool pshard; // enable pipelined sharding: weights on CPU host, pipelined to GPU per split + bool pshard_delegate_compute; // tensor overrides choose storage; scheduler chooses compute backends + bool pshard_cache_skip_load; // skip loading plan cache (rebuild from scratch, then overwrite) + size_t max_vram_alloc; // VRAM budget in MiB for the unified preload buffer (0 = auto from free VRAM) + struct llama_pshard_plan_registry * pshard_registry; // tier plan registry, caller-owned. populated by llama_params_fit_pshard }; struct llama_sampler_seq_config { @@ -380,6 +386,8 @@ extern "C" { // note: the samplers must be sampler chains (i.e. use llama_sampler_chain_init) struct llama_sampler_seq_config * samplers; size_t n_samplers; + + bool pshard; // enable pipelined sharding for this context (3 GPU backends, split callbacks) }; struct llama_model_tensor_override { @@ -532,6 +540,26 @@ extern "C" { uint32_t n_ctx_min, // minimum context size to set when trying to reduce memory use enum ggml_log_level log_level); // minimum log level to print during fitting, lower levels go to debug log + // Pipelined sharding planner: probes VRAM usage for each strategy/tier combination, + // selects the best plan (by TPS if benchmark data available, else by VRAM fit), + // and populates tensor_buft_overrides for model loading. + // Results are cached to .tensor_overrides.pshard_registry for fast subsequent launches. + // If all layers fit in VRAM, sets mparams->pshard=false and returns (baseline loading). + LLAMA_API void llama_params_fit_pshard( + const char * path_model, + struct llama_model_params * mparams, + struct llama_context_params * cparams, + struct 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 + + // Create/free a tier plan registry. Caller owns the pointer and passes it via + // mparams->pshard_registry before calling llama_params_fit_pshard. + // n_tier_max: largest batch size to probe (determines tier range, typically max(n_batch, 16384)). + // n_seq_max: for speculative decoding tiers. + LLAMA_API struct llama_pshard_plan_registry * llama_pshard_registry_create(uint32_t n_tier_max, uint32_t n_seq_max); + LLAMA_API void llama_pshard_registry_free(struct llama_pshard_plan_registry * registry); + LLAMA_API int64_t llama_time_us(void); LLAMA_API size_t llama_max_devices(void); diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 121c21fed957..ca100796556a 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -12,8 +12,10 @@ add_library(llama llama-adapter.cpp llama-arch.cpp llama-batch.cpp + llama-benchmark.cpp llama-chat.cpp llama-context.cpp + llama-context-pshard.cpp llama-cparams.cpp llama-grammar.cpp llama-graph.cpp @@ -25,11 +27,13 @@ add_library(llama llama-memory.cpp llama-memory-hybrid.cpp llama-memory-hybrid-iswa.cpp + llama-memory-pshard.cpp llama-memory-recurrent.cpp llama-mmap.cpp llama-model-loader.cpp llama-model-saver.cpp llama-model.cpp + llama-pshard-plan.cpp llama-quant.cpp llama-sampler.cpp llama-vocab.cpp diff --git a/src/llama-benchmark.cpp b/src/llama-benchmark.cpp new file mode 100644 index 000000000000..3a4e425672ef --- /dev/null +++ b/src/llama-benchmark.cpp @@ -0,0 +1,693 @@ +#include "llama-benchmark.h" +#include "llama-impl.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include + +llama_op_metrics llama_op_metrics_compute(const ggml_tensor * node) { + llama_op_metrics m = {}; + + switch (node->op) { + case GGML_OP_MUL_MAT: { + m.N = node->ne[0]; + m.M = node->ne[1]; + m.K = node->src[0]->ne[0]; + m.ops = 2.0 * m.N * m.K * m.M; + m.bytes = ggml_nbytes(node->src[0]) + ggml_nbytes(node->src[1]) + ggml_nbytes(node); + m.quant_type = ggml_type_name(node->src[0]->type); + break; + } + case GGML_OP_MUL_MAT_ID: { + m.N = node->ne[0]; + m.n_experts_used = node->ne[1]; + m.M = node->ne[2]; + m.K = node->src[0]->ne[0]; + const int64_t total_experts = node->src[0]->ne[2]; + m.ops = 2.0 * m.N * m.K * m.M * m.n_experts_used; + m.bytes = (ggml_nbytes(node->src[0]) * m.n_experts_used / total_experts) + + ggml_nbytes(node->src[1]) + ggml_nbytes(node); + m.quant_type = ggml_type_name(node->src[0]->type); + break; + } + case GGML_OP_FLASH_ATTN_EXT: { + m.head_dim = node->ne[0]; + m.n_q_heads = node->ne[1]; + m.n_tokens = node->ne[2]; + m.ctx_len = node->src[1]->ne[1]; + m.n_kv_heads = node->src[1]->ne[2]; + m.ops = 2.0 * m.n_tokens * m.head_dim * m.ctx_len * m.n_q_heads * 2.0; + m.bytes = ggml_nbytes(node->src[0]) + ggml_nbytes(node->src[1]) + + ggml_nbytes(node->src[2]) + ggml_nbytes(node); + break; + } + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: { + m.n_elements = ggml_nelements(node); + m.ops = 6.0 * m.n_elements; + m.bytes = ggml_nbytes(node->src[0]) + ggml_nbytes(node); + m.quant_type = ggml_type_name(node->type); + break; + } + case GGML_OP_RMS_NORM: { + m.n_elements = ggml_nelements(node); + m.ops = 3.0 * m.n_elements; + m.bytes = ggml_nbytes(node->src[0]) + ggml_nbytes(node); + m.quant_type = ggml_type_name(node->type); + break; + } + case GGML_OP_GLU: { + m.n_elements = ggml_nelements(node); + m.ops = 2.0 * m.n_elements; + m.bytes = ggml_nbytes(node->src[0]); + if (node->src[1]) { m.bytes += ggml_nbytes(node->src[1]); } + m.bytes += ggml_nbytes(node); + m.quant_type = ggml_type_name(node->type); + break; + } + case GGML_OP_ADD: + case GGML_OP_MUL: + case GGML_OP_SUB: + case GGML_OP_DIV: { + m.n_elements = ggml_nelements(node); + m.ops = m.n_elements; + m.bytes = ggml_nbytes(node->src[0]); + if (node->src[1]) { m.bytes += ggml_nbytes(node->src[1]); } + m.bytes += ggml_nbytes(node); + m.quant_type = ggml_type_name(node->type); + break; + } + case GGML_OP_GET_ROWS: { + m.n_elements = ggml_nelements(node); + m.bytes = ggml_nbytes(node); + m.quant_type = ggml_type_name(node->src[0]->type); + break; + } + case GGML_OP_SET_ROWS: + case GGML_OP_CPY: { + m.n_elements = ggml_nelements(node); + m.bytes = ggml_nbytes(node->src[0]) + ggml_nbytes(node); + m.quant_type = ggml_type_name(node->src[0]->type); + break; + } + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + case GGML_OP_NONE: + break; + default: { + m.n_elements = ggml_nelements(node); + m.bytes = ggml_nbytes(node); + if (node->src[0]) { m.bytes += ggml_nbytes(node->src[0]); } + m.quant_type = ggml_type_name(node->type); + break; + } + } + return m; +} + +std::string llama_benchmark_predictor::make_key( + const std::string & op, const std::string & quant, + int64_t N, int64_t K, int64_t batch) { + return op + "|" + quant + "|" + std::to_string(N) + "|" + + std::to_string(K) + "|" + std::to_string(batch); +} + +std::string llama_benchmark_predictor::make_attn_key( + int64_t ctx_len, int64_t n_kv_heads, int64_t n_tokens) { + return "FLASH_ATTN|" + std::to_string(ctx_len) + "|" + + std::to_string(n_kv_heads) + "|" + std::to_string(n_tokens); +} + +std::string llama_benchmark_predictor::make_elem_key( + const std::string & op, int64_t n_elements) { + return op + "|" + std::to_string(n_elements); +} + +static void build_entry_map( + std::unordered_map & map, + const std::vector & entries) { + map.clear(); + map.reserve(entries.size()); + for (const auto & e : entries) { + std::string key; + if (e.op_name == "MUL_MAT" || e.op_name == "MUL_MAT_ID") { + key = llama_benchmark_predictor::make_key(e.op_name, e.quant, e.N, e.K, e.B); + } else if (e.op_name.compare(0, 10, "FLASH_ATTN") == 0) { + // n_heads stores n_kv_heads for attention benchmarks (written by profiler) + key = llama_benchmark_predictor::make_attn_key(e.ctx_len, e.n_heads, e.n_tokens); + } else { + key = llama_benchmark_predictor::make_elem_key(e.op_name, e.n_elements); + } + map[key] = &e; + } +} + +void llama_benchmark_predictor::build_maps() { + build_entry_map(cpu_map, cpu_entries); + build_entry_map(gpu_map, gpu_entries); +} + +bool llama_benchmark_predictor::load_cpu(const char * filepath, int n_threads) { + cpu_entries.clear(); + cpu_map.clear(); + + FILE * f = fopen(filepath, "r"); + if (!f) { + LLAMA_LOG_WARN("%s: could not open CPU benchmark file: %s\n", __func__, filepath); + return false; + } + + double max_dram_bw = 0.0; + + char line[512]; + while (fgets(line, sizeof(line), f)) { + if (line[0] == '#') { + int tc = 0; + double dram_bw = 0.0, pcie_standalone = 0.0, pcie_concurrent = 0.0, cpu_eff = 0.0; + + if (sscanf(line, + "# Threads=%d: DRAM_BW=%lf GB/s, PCIe_Standalone=%lf GB/s," + " PCIe_Concurrent=%lf GB/s (CPU_Eff=%lf%%)", + &tc, &dram_bw, &pcie_standalone, &pcie_concurrent, &cpu_eff) == 5) { + if (dram_bw > max_dram_bw) { + max_dram_bw = dram_bw; + } + if (tc == n_threads) { + stats.peak_system_bw = dram_bw; + stats.peak_pcie_bw = pcie_standalone; + stats.eff_system_bw = dram_bw * (cpu_eff / 100.0); + stats.eff_pcie_bw = std::min(pcie_concurrent, dram_bw); + } + } else if (sscanf(line, "# Threads=%d: DRAM_BW=%lf GB/s", &tc, &dram_bw) == 2) { + if (dram_bw > max_dram_bw) { + max_dram_bw = dram_bw; + } + if (tc == n_threads) { + stats.peak_system_bw = dram_bw; + } + } + continue; + } + + if (line[0] == '\n') continue; + + llama_benchmark_entry e; + char op_name[64], quant[16]; + + int parsed = sscanf(line, + "%63s %15s %d %lf %lf %lf %lf %lf %lf" + " %lld %lld %lld %lld %lld %lld %lld %lld", + op_name, quant, &e.threads, + &e.ai, &e.bw_gb_s, &e.peak_gflops, &e.ridge, + &e.eff_gflops, &e.eff_pcie_bw, + &e.N, &e.K, &e.B, &e.n_tokens, &e.ctx_len, + &e.n_heads, &e.head_dim, &e.n_elements); + + if (parsed >= 9 && e.threads == n_threads) { + e.op_name = op_name; + e.quant = quant; + cpu_entries.push_back(std::move(e)); + } + } + fclose(f); + + // PCIe can't exceed system DRAM BW + if (max_dram_bw > 0.0) { + stats.peak_pcie_bw = std::min(stats.peak_pcie_bw, max_dram_bw); + } else { + stats.peak_pcie_bw = std::min(stats.peak_pcie_bw, stats.peak_system_bw); + } + + if (stats.eff_system_bw == 0.0) { stats.eff_system_bw = stats.peak_system_bw; } + if (stats.eff_pcie_bw == 0.0) { stats.eff_pcie_bw = stats.peak_pcie_bw; } + + build_maps(); + + LLAMA_LOG_INFO("%s: loaded %zu CPU entries for %d threads" + " (peak: DRAM=%.1f GB/s, PCIe=%.1f GB/s" + " | concurrent: DRAM=%.1f GB/s, PCIe=%.1f GB/s)\n", + __func__, cpu_entries.size(), n_threads, + stats.peak_system_bw, stats.peak_pcie_bw, + stats.eff_system_bw, stats.eff_pcie_bw); + + return !cpu_entries.empty(); +} + +bool llama_benchmark_predictor::load_gpu(const char * filepath) { + gpu_entries.clear(); + gpu_map.clear(); + + FILE * f = fopen(filepath, "r"); + if (!f) { + LLAMA_LOG_WARN("%s: could not open GPU benchmark file: %s\n", __func__, filepath); + return false; + } + + char line[512]; + while (fgets(line, sizeof(line), f)) { + if (line[0] == '#') { + const char * bw_str = strstr(line, "GPU_Memory_BW="); + const char * cp_str = strstr(line, "GPU_Peak_Compute="); + if (bw_str && cp_str) { + double mem_bw = 0.0, compute = 0.0; + sscanf(bw_str, "GPU_Memory_BW=%lf", &mem_bw); + sscanf(cp_str, "GPU_Peak_Compute=%lf", &compute); + stats.peak_gpu_mem_bw = mem_bw; + stats.peak_gpu_compute = compute; + } + continue; + } + + if (line[0] == '\n') continue; + + llama_benchmark_entry e; + char op_name[64], quant[16]; + + int parsed = sscanf(line, + "%63s %15s %lf %lf %lf %lf" + " %lld %lld %lld %lld %lld %lld %lld %lld", + op_name, quant, + &e.ai, &e.bw_gb_s, &e.peak_gflops, &e.ridge, + &e.N, &e.K, &e.B, &e.n_tokens, &e.ctx_len, + &e.n_heads, &e.head_dim, &e.n_elements); + + if (parsed == 14) { + e.op_name = op_name; + e.quant = quant; + e.threads = -1; + e.eff_gflops = e.peak_gflops; + e.eff_pcie_bw = 0.0; + gpu_entries.push_back(std::move(e)); + } + } + fclose(f); + + build_maps(); + + LLAMA_LOG_INFO("%s: loaded %zu GPU entries (mem_bw=%.1f GB/s, compute=%.1f GFLOP/s)\n", + __func__, gpu_entries.size(), + stats.peak_gpu_mem_bw, stats.peak_gpu_compute); + + return !gpu_entries.empty(); +} + +std::string llama_benchmark_predictor::make_timing_key( + bool is_gpu, bool async_copy, const char * op_name, + const llama_op_metrics & m, int32_t batch_size) { + // mode prefix: GPU vs CPU with/without concurrent PCIe (async_copy uses eff_* bandwidths) + std::string key = is_gpu ? "GPU|" : (async_copy ? "CPU_ASYNC|" : "CPU_SYNC|"); + key += op_name; + key += "|"; + + if (m.quant_type) { + key += m.quant_type; + key += "|"; + } + + if (strstr(op_name, "MUL_MAT_ID")) { + key += std::to_string(m.N) + "|" + std::to_string(m.K) + "|" + + std::to_string(batch_size) + "|" + std::to_string(m.M) + "|" + + std::to_string(m.n_experts_used); + } else if (strstr(op_name, "MUL_MAT")) { + key += std::to_string(m.N) + "|" + std::to_string(m.K) + "|" + + std::to_string(batch_size) + "|" + std::to_string(m.M); + } else if (strstr(op_name, "FLASH_ATTN")) { + key += std::to_string(m.head_dim) + "|" + std::to_string(m.n_q_heads) + "|" + + std::to_string(batch_size) + "|" + std::to_string(m.n_tokens) + "|" + + std::to_string(m.ctx_len) + "|" + std::to_string(m.n_kv_heads); + } else { + key += std::to_string(m.n_elements); + } + + return key; +} + +llama_split_timing llama_benchmark_predictor::predict_split( + struct ggml_tensor ** nodes, int n_nodes, + bool is_gpu, int32_t batch_size, bool async_copy, + timing_cache_t * timing_cache) const { + + llama_split_timing result = {}; + timing_cache_t local_timing_cache; + timing_cache_t & cache = timing_cache ? *timing_cache : local_timing_cache; + + const auto & map = is_gpu ? gpu_map : cpu_map; + const auto & entries = is_gpu ? gpu_entries : cpu_entries; + + double pcie_sum = 0.0; + + for (int i = 0; i < n_nodes; i++) { + ggml_tensor * node = nodes[i]; + llama_op_metrics m = llama_op_metrics_compute(node); + + if (m.ops == 0.0 && m.bytes == 0.0) continue; + + const char * op_name = ggml_op_name(node->op); + double op_time_ms = 0.0; + double pcie_contrib = 0.0; + + std::string tkey = make_timing_key(is_gpu, async_copy, op_name, m, batch_size); + auto cache_it = cache.find(tkey); + if (cache_it != cache.end()) { + op_time_ms = cache_it->second.first; + pcie_contrib = cache_it->second.second; + result.time_ms += op_time_ms; + pcie_sum += pcie_contrib; + result.n_cache_hit++; + result.op_count++; + continue; + } + + // build hash key for exact lookup + std::string hkey; + switch (node->op) { + case GGML_OP_MUL_MAT: + hkey = make_key("MUL_MAT", m.quant_type ? m.quant_type : "", m.N, m.K, batch_size); + break; + case GGML_OP_MUL_MAT_ID: + hkey = make_key("MUL_MAT_ID", m.quant_type ? m.quant_type : "", m.N, m.K, batch_size); + break; + case GGML_OP_FLASH_ATTN_EXT: + hkey = make_attn_key(m.ctx_len, m.n_kv_heads, batch_size); + break; + case GGML_OP_ROPE: + case GGML_OP_RMS_NORM: + case GGML_OP_GLU: + case GGML_OP_ADD: + case GGML_OP_MUL: + case GGML_OP_SUB: + case GGML_OP_DIV: + hkey = make_elem_key(op_name, m.n_elements); + break; + default: + break; + } + + const llama_benchmark_entry * match = nullptr; + bool exact = false; + + // try exact hash match + if (!hkey.empty()) { + auto it = map.find(hkey); + if (it != map.end()) { + match = it->second; + exact = true; + result.n_exact++; + } + } + + // fall back to nearest-neighbor + if (!match) { + switch (node->op) { + case GGML_OP_MUL_MAT: + match = find_nearest(entries, "MUL_MAT", m.quant_type, m.N, m.K, 0, 0, -1, batch_size); + break; + case GGML_OP_MUL_MAT_ID: + match = find_nearest(entries, "MUL_MAT_ID", m.quant_type, m.N, m.K, 0, 0, -1, batch_size); + break; + case GGML_OP_FLASH_ATTN_EXT: + match = find_nearest(entries, "FLASH_ATTN", nullptr, 0, 0, m.ctx_len, 0, m.n_kv_heads, batch_size); + break; + default: + match = find_nearest(entries, op_name, nullptr, 0, 0, 0, m.n_elements, -1, batch_size); + break; + } + if (match) { + result.n_nearest++; + } + } + + // compute timing from match or fall back to peak BW + if (match) { + // GPU always peak; CPU uses eff (concurrent) when async_copy, peak otherwise + const double gflops = is_gpu ? match->peak_gflops + : (async_copy ? match->eff_gflops : match->peak_gflops); + + if (exact) { + if (gflops > 0.0) { + op_time_ms = (m.ops / 1e9) / gflops * 1000.0; + } + } else { + const double ai = m.bytes > 0.0 ? (m.ops / m.bytes) : 0.0; + if (ai < match->ridge && match->bw_gb_s > 0.0) { + op_time_ms = (m.bytes / 1e9) / match->bw_gb_s * 1000.0; + } else if (gflops > 0.0) { + op_time_ms = (m.ops / 1e9) / gflops * 1000.0; + } + } + + pcie_contrib = op_time_ms * match->eff_pcie_bw; + } else { + result.n_fallback++; + if (m.bytes > 0.0) { + const double bw = is_gpu ? stats.peak_gpu_mem_bw + : (async_copy ? stats.eff_system_bw : stats.peak_system_bw); + if (bw > 0.0) { + op_time_ms = (m.bytes / 1e9) / bw * 1000.0; + } + if (!is_gpu && stats.eff_pcie_bw > 0.0) { + pcie_contrib = op_time_ms * stats.eff_pcie_bw; + } + } + } + + result.time_ms += op_time_ms; + pcie_sum += pcie_contrib; + result.op_count++; + + cache[tkey] = {op_time_ms, pcie_contrib}; + } + + // convert accumulated pcie_sum to time-weighted average BW + if (!is_gpu) { + result.eff_pcie_bw = (result.time_ms > 0.0) ? (pcie_sum / result.time_ms) + : stats.eff_pcie_bw; + } + + return result; +} + +static double get_bits_per_weight(const std::string & quant) { + static const std::unordered_map table = { + {"q2_K", 2.6}, {"q3_K", 3.4}, {"q4_0", 4.5}, {"q4_1", 5.0}, {"q4_K", 4.5}, + {"q5_0", 5.5}, {"q5_1", 6.0}, {"q5_K", 5.5}, {"q6_K", 6.6}, + {"q8_0", 8.5}, {"q8_1", 9.0}, {"f16", 16.0}, {"f32", 32.0}, + }; + auto it = table.find(quant); + return it != table.end() ? it->second : -1.0; +} + +const llama_benchmark_entry * llama_benchmark_predictor::find_nearest( + const std::vector & entries, + const char * op_name, const char * quant, + int64_t N, int64_t K, int64_t ctx_len, int64_t n_elements, + int64_t n_kv_heads, int64_t target_batch) { + + if (!op_name || entries.empty()) { + return nullptr; + } + + const llama_benchmark_entry * best = nullptr; + double best_score = 1e20; + + const double target_bpw = quant ? get_bits_per_weight(quant) : -1.0; + + // FLASH_ATTN entries have suffixed names (e.g. FLASH_ATTN_MHA, FLASH_ATTN_GQA-8) + const bool attn_query = (strncmp(op_name, "FLASH_ATTN", 10) == 0); + const bool matmul_query = (strcmp(op_name, "MUL_MAT") == 0 + || strcmp(op_name, "MUL_MAT_ID") == 0); + + for (const auto & b : entries) { + if (attn_query) { + if (b.op_name.compare(0, 10, "FLASH_ATTN") != 0) continue; + } else { + if (b.op_name != op_name) continue; + } + + double dim_score = 0.0; + double batch_score = 0.0; + double quant_score = 0.0; + + if (matmul_query) { + const double n_diff = std::abs((double)b.N - N) / std::max(N, (int64_t)1); + const double k_diff = std::abs((double)b.K - K) / std::max(K, (int64_t)1); + dim_score = n_diff + k_diff; + } else if (attn_query) { + const double ctx_diff = std::abs((double)b.ctx_len - ctx_len) / std::max(ctx_len, (int64_t)1); + double kv_diff = 0.0; + if (n_kv_heads > 0 && b.n_heads > 0) { + kv_diff = std::abs((double)b.n_heads - n_kv_heads) / std::max(n_kv_heads, (int64_t)1); + } + dim_score = ctx_diff + kv_diff * 0.5; + } else { + const double elem_diff = std::abs((double)b.n_elements - n_elements) / std::max(n_elements, (int64_t)1); + dim_score = elem_diff; + } + + if (target_batch > 0 && b.B > 0) { + batch_score = std::abs((double)b.B - target_batch) / std::max(target_batch, (int64_t)1); + } + + if (target_bpw > 0) { + const double bench_bpw = get_bits_per_weight(b.quant); + if (bench_bpw > 0) { + quant_score = std::abs(bench_bpw - target_bpw) / target_bpw; + } else { + quant_score = 1.0; + } + } else if (quant && b.quant != quant) { + continue; + } + + // batch most important (memory vs compute regime), then dims, then quant + const double score = batch_score * 1.0 + dim_score * 0.5 + quant_score * 0.3; + + if (score < best_score) { + best_score = score; + best = &b; + } + } + + return best; +} + +double llama_benchmark_predictor::predict_tps( + ggml_backend_sched_t sched, + int cpu_backend_id, + uint32_t kv_size, + int32_t batch_size, + uint32_t n_outputs, + bool has_rs) const { + + const int n_splits = ggml_backend_sched_get_n_splits(sched); + if (n_splits <= 0) return 0.0; + + LLAMA_LOG_DEBUG("%s: n_splits=%d, bs=%d, kv_size=%u, n_outputs=%u, has_rs=%d\n", + __func__, n_splits, batch_size, kv_size, n_outputs, (int)has_rs); + + const double pcie_bw = stats.peak_pcie_bw; + // KV transfer ratios vs full writeback_bytes the sched reports: + // download: only newly written cells -> batch_size/kv_size (decode bs=1/kv=1024 -> 0.1%) + // upload: all used cells (cache fills over time) -> 1.0 (conservative, assumes full cache) + // RS models use FULL mode (entire tensor every batch), so both are 1.0. + const double dl_ratio = has_rs ? 1.0 + : ((kv_size > 0) ? std::min(1.0, (double)batch_size / kv_size) : 1.0); + const double ul_ratio = 1.0; + double total_ms = 0.0; + bool copy_prefetched = false; + timing_cache_t timing_cache; + + for (int i = 0; i < n_splits; i++) { + struct ggml_backend_sched_split_info si = {}; + if (!ggml_backend_sched_get_split_info(sched, i, &si)) continue; + + const bool is_gpu = (si.backend_id != cpu_backend_id); + + double input_copy_ms = 0.0; + double input_copy_bytes = 0.0; + if (is_gpu && !copy_prefetched && pcie_bw > 0.0) { + input_copy_bytes = (double)si.input_weight_copy_bytes + + (double)si.writeback_bytes * ul_ratio; + input_copy_ms = (input_copy_bytes / 1e9 / pcie_bw) * 1000.0; + } + + // peek at next split to determine if async prefetch will overlap with this split + double prefetch_bytes = 0.0; + bool next_copy_prefetched = false; + if (i + 1 < n_splits) { + struct ggml_backend_sched_split_info next_si = {}; + if (ggml_backend_sched_get_split_info(sched, i + 1, &next_si) && + next_si.can_prefetch_weights) { + prefetch_bytes = (double)next_si.input_weight_bytes + + (double)next_si.writeback_bytes * ul_ratio; + next_copy_prefetched = prefetch_bytes > 0.0; + } + } + + const bool async_copy = (prefetch_bytes > 0.0); + + // compute cost (CPU splits use eff_gflops when async_copy due to PCIe contention) + struct ggml_tensor ** nodes = ggml_graph_nodes(si.graph); + int n_nodes = ggml_graph_n_nodes(si.graph); + llama_split_timing t = predict_split(nodes, n_nodes, is_gpu, batch_size, async_copy, &timing_cache); + + // prefetch cost: use concurrent PCIe BW for CPU splits (bus shared with DRAM), + // peak PCIe BW for GPU splits (GPU compute doesn't contend with PCIe DMA) + double prefetch_ms = 0.0; + if (prefetch_bytes > 0.0) { + const double eff_bw = (!is_gpu && t.eff_pcie_bw > 0.0) ? t.eff_pcie_bw : pcie_bw; + prefetch_ms = (prefetch_bytes / 1e9 / eff_bw) * 1000.0; + } + + // output scaling: use the output rows in the reserved graph, then scale to + // the runtime number of logits. This keeps the memory-probe graph as the + // source of truth and avoids double-scaling when it was already reduced. + if (i == n_splits - 1 && n_outputs > 0) { + for (int j = 0; j < n_nodes; j++) { + if (nodes[j]->name && strstr(nodes[j]->name, "result_output")) { + const llama_op_metrics out_m = llama_op_metrics_compute(nodes[j]); + const int32_t graph_outputs = (int32_t) std::max(1, out_m.M); + if (graph_outputs != batch_size || n_outputs < (uint32_t) graph_outputs) { + llama_split_timing out_included = predict_split(&nodes[j], 1, is_gpu, batch_size, async_copy, &timing_cache); + llama_split_timing out_graph = predict_split(&nodes[j], 1, is_gpu, graph_outputs, async_copy, &timing_cache); + if (out_included.time_ms > 0.0 || out_graph.time_ms > 0.0) { + const double scale = n_outputs < (uint32_t) graph_outputs + ? (double)n_outputs / (double)graph_outputs + : 1.0; + const double scaled_ms = out_graph.time_ms * scale; + t.time_ms -= out_included.time_ms; + t.time_ms += scaled_ms; + LLAMA_LOG_DEBUG("%s: out_t: included=%f graph=%f scaled=%f rows=%d->%u\n", + __func__, out_included.time_ms, out_graph.time_ms, scaled_ms, graph_outputs, n_outputs); + } + } + break; + } + } + } + + // KV/RS download cost (GPU splits, after compute, on compute stream) + double kv_dl_ms = 0.0; + if (is_gpu && si.writeback_bytes > 0 && pcie_bw > 0.0) { + double dl_bytes = (double)si.writeback_bytes * dl_ratio; + kv_dl_ms = (dl_bytes / 1e9 / pcie_bw) * 1000.0; + } + + // activation copy cost (synchronous, between splits on different backends) + double activ_copy_ms = 0.0; + if (si.input_activ_bytes > 0 && pcie_bw > 0.0) { + activ_copy_ms = ((double)si.input_activ_bytes / 1e9 / pcie_bw) * 1000.0; + } + + // Prefetch is enqueued after current inputs/precompute and before graph compute, + // so whatever remains can overlap the current split compute. CPU splits use the + // effective PCIe BW above because DMA contends with CPU memory traffic. + double split_ms = input_copy_ms + activ_copy_ms; + split_ms += std::max(t.time_ms + kv_dl_ms, prefetch_ms); + total_ms += split_ms; + copy_prefetched = next_copy_prefetched; + + double dl_bytes = is_gpu ? (double)si.writeback_bytes * dl_ratio : 0.0; + LLAMA_LOG_DEBUG("%s: split %d/%d [%s] input_copy=%.3f (%.2f MiB) compute=%.3f kv_dl=%.3f (%.2f MiB) prefetch=%.3f (%.2f MiB) activ=%.3f (%.2f MiB) -> %.3f ms" + " (exact=%d near=%d fall=%d cache=%d)\n", + __func__, i, n_splits, is_gpu ? "GPU" : "CPU", + input_copy_ms, input_copy_bytes / (1024.0 * 1024.0), + t.time_ms, + kv_dl_ms, dl_bytes / (1024.0 * 1024.0), + prefetch_ms, prefetch_bytes / (1024.0 * 1024.0), + activ_copy_ms, (double)si.input_activ_bytes / (1024.0 * 1024.0), + split_ms, + t.n_exact, t.n_nearest, t.n_fallback, t.n_cache_hit); + } + + double tps = (total_ms > 0.0) ? (batch_size * 1000.0 / total_ms) : 0.0; + LLAMA_LOG_DEBUG("%s: total=%.3f ms, dl_ratio=%.4f, ul_ratio=%.1f, pcie_bw=%.1f GB/s -> %.1f tps\n", + __func__, total_ms, dl_ratio, ul_ratio, pcie_bw, tps); + return tps; +} diff --git a/src/llama-benchmark.h b/src/llama-benchmark.h new file mode 100644 index 000000000000..8fbaf729d99c --- /dev/null +++ b/src/llama-benchmark.h @@ -0,0 +1,151 @@ +#pragma once + +#include +#include +#include +#include +#include + +struct ggml_tensor; +typedef struct ggml_backend_sched * ggml_backend_sched_t; + +// Per-op FLOPS/bytes metrics extracted from a ggml tensor node. +// Used by roofline prediction to classify ops as compute-bound or memory-bound. +struct llama_op_metrics { + double ops = 0.0; + double bytes = 0.0; + const char * quant_type = nullptr; + + int64_t N = 0, K = 0, M = 0; + int64_t ctx_len = 0; + int64_t n_elements = 0; + int64_t n_kv_heads = 0; + int64_t n_experts_used = 0; + int64_t head_dim = 0; + int64_t n_q_heads = 0; + int64_t n_tokens = 0; +}; + +// Extract ops/bytes metrics from a compute graph node. +// Zero-cost view ops (RESHAPE, VIEW, PERMUTE, TRANSPOSE, NONE) return {0,0}. +llama_op_metrics llama_op_metrics_compute(const ggml_tensor * node); + +// One row from profiler output (CPU or GPU). +// For FLASH_ATTN ops, n_heads stores n_kv_heads (written by profiler). +struct llama_benchmark_entry { + std::string op_name; + std::string quant; + int threads = 0; + + double ai = 0.0; // arithmetic intensity (FLOP/byte) + double bw_gb_s = 0.0; // effective memory bandwidth (GB/s) + double peak_gflops = 0.0; // standalone compute throughput (GFLOP/s) + double ridge = 0.0; // roofline ridge point (FLOP/byte) + double eff_gflops = 0.0; // concurrent compute (CPU: under PCIe load; GPU: == peak) + double eff_pcie_bw = 0.0; // effective PCIe BW during concurrent compute (CPU only) + + int64_t N = 0, K = 0, B = 0; + int64_t n_tokens = 0; + int64_t ctx_len = 0; + int64_t n_heads = 0; // n_kv_heads for FLASH_ATTN, unused for other ops + int64_t head_dim = 0; + int64_t n_elements = 0; +}; + +// Global bandwidth/compute stats parsed from profiler file headers. +struct llama_benchmark_stats { + // CPU (from cpu profiler header) + double peak_system_bw = 0.0; // peak DRAM BW (GB/s) + double peak_pcie_bw = 0.0; // peak PCIe BW standalone (GB/s) + double eff_system_bw = 0.0; // DRAM BW under concurrent PCIe load (GB/s) + double eff_pcie_bw = 0.0; // PCIe BW under concurrent CPU load (GB/s) + + // GPU (from gpu profiler header) + double peak_gpu_mem_bw = 0.0; // peak GPU memory BW (GB/s) + double peak_gpu_compute = 0.0; // peak GPU compute (GFLOP/s) +}; + +// Aggregate timing prediction for a set of graph nodes (typically one sched split). +struct llama_split_timing { + double time_ms = 0.0; + double eff_pcie_bw = 0.0; // time-weighted avg PCIe BW (CPU splits only) + int op_count = 0; + int n_cache_hit = 0; // timing cache reuses + int n_exact = 0; // exact benchmark hash hits + int n_nearest = 0; // nearest-neighbor matches + int n_fallback = 0; // memory BW fallback (no benchmark match) +}; + +// Benchmark data + O(1) hash maps + nearest-neighbor lookup + split timing prediction. +// Owns entry vectors; map pointers are stable after load. +struct llama_benchmark_predictor { + std::vector cpu_entries; + std::vector gpu_entries; + llama_benchmark_stats stats = {}; + + std::unordered_map cpu_map; + std::unordered_map gpu_map; + + // Timing cache: same shape + same mode = same timing across layers. + // Key: make_timing_key(), Value: {time_ms, pcie_contrib}. + using timing_cache_t = std::unordered_map>; + + llama_benchmark_predictor() = default; + llama_benchmark_predictor(const llama_benchmark_predictor &) = delete; + llama_benchmark_predictor & operator=(const llama_benchmark_predictor &) = delete; + + // Parse profiler output, populate entries + hash maps. + // CPU loader filters by n_threads; returns false if file missing or empty. + bool load_cpu(const char * filepath, int n_threads); + bool load_gpu(const char * filepath); + + bool has_cpu() const { return !cpu_entries.empty(); } + bool has_gpu() const { return !gpu_entries.empty(); } + + // Predict aggregate timing for a set of graph nodes. + // nodes/n_nodes: typically from one sched split (via ggml_graph_nodes/n_nodes). + // async_copy: true when PCIe transfers run concurrently (uses eff_* bandwidths for CPU). + llama_split_timing predict_split( + struct ggml_tensor ** nodes, int n_nodes, + bool is_gpu, int32_t batch_size, bool async_copy = true, + timing_cache_t * timing_cache = nullptr) const; + + // Nearest-neighbor search with weighted scoring (batch > dims > quant). + // For FLASH_ATTN queries, pass op_name="FLASH_ATTN" -- matches all FLASH_ATTN_* entries. + static const llama_benchmark_entry * find_nearest( + const std::vector & entries, + const char * op_name, const char * quant, + int64_t N, int64_t K, int64_t ctx_len, int64_t n_elements, + int64_t n_kv_heads = -1, int64_t target_batch = 1); + + // Key generators for hash map lookup (used by predict to build exact-match keys). + static std::string make_key( + const std::string & op, const std::string & quant, + int64_t N, int64_t K, int64_t batch); + + static std::string make_attn_key( + int64_t ctx_len, int64_t n_kv_heads, int64_t n_tokens); + + static std::string make_elem_key( + const std::string & op, int64_t n_elements); + + // Timing cache key: encodes backend type, async mode, op, shape, and batch. + static std::string make_timing_key( + bool is_gpu, bool async_copy, const char * op_name, + const llama_op_metrics & m, int32_t batch_size); + + // End-to-end TPS prediction across all sched splits. + // kv_size: total KV cells per layer (for partial upload ratio). + // has_rs: true if model has recurrent state layers (RS uses full-tensor transfer, not cell-granular). + // Returns 0 if no splits or no benchmark data. + double predict_tps( + ggml_backend_sched_t sched, + int cpu_backend_id, + uint32_t kv_size, + int32_t batch_size, + uint32_t n_outputs = 0, + bool has_rs = false) const; + +private: + void build_maps(); +}; diff --git a/src/llama-context-pshard.cpp b/src/llama-context-pshard.cpp new file mode 100644 index 000000000000..f2e1b7254be6 --- /dev/null +++ b/src/llama-context-pshard.cpp @@ -0,0 +1,35 @@ +#include "llama-context.h" + +#include "llama-model.h" +#include "llama-memory.h" +#include "llama-pipe-shard.h" +#include "llama-pshard-plan.h" +#include "llama-impl.h" + +#include "ggml-backend.h" + +#include +#include +#include + +void pshard_assign_tensors( + ggml_backend_sched_t sched, + const llama_model & model, + llama_memory_i * memory, + const std::vector & backends, + const pshard_dev_layout & layout) { + const auto & tbids = model.get_tensor_backend_ids(); + const auto & lbids = model.get_layer_backend_ids(); + + for (const auto & [tensor, bid] : tbids) { + if (bid >= 0 && bid < (int32_t)backends.size()) { + ggml_backend_sched_set_tensor_backend_hint(sched, tensor, backends[bid].get()); + } + } + + if (memory) { + for (auto * ps : memory->get_pipe_shards()) { + ps->assign_tensors(sched, lbids, backends, layout); + } + } +} diff --git a/src/llama-context.cpp b/src/llama-context.cpp index ee0c29235cd2..b33508d2649f 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -6,8 +6,10 @@ #include "llama-batch.h" #include "llama-io.h" #include "llama-memory.h" +#include "llama-pipe-shard.h" #include "llama-mmap.h" #include "llama-model.h" +#include "llama-pshard-plan.h" #include "llama-ext.h" #include "llama.h" @@ -23,8 +25,10 @@ llama_context::llama_context( const llama_model & model, - llama_context_params params) : + llama_context_params params, + llama_context_probe_reserve probe_reserve_) : model(model), + probe_reserve(probe_reserve_), cvec(std::make_unique()), loras(std::make_unique()), balloc(std::make_unique(model.hparams.n_pos_per_embd())) { @@ -53,6 +57,7 @@ llama_context::llama_context( cparams.no_perf = params.no_perf; cparams.pooling_type = params.pooling_type; cparams.warmup = false; + cparams.pshard = params.pshard; cparams.n_ctx = params.n_ctx == 0 ? hparams.n_ctx_train : params.n_ctx; cparams.rope_freq_base = params.rope_freq_base == 0.0f ? hparams.rope_freq_base_train : params.rope_freq_base; @@ -157,6 +162,12 @@ llama_context::llama_context( cparams.fused_gdn_ch = true; cparams.auto_fgdn = true; + // pshard remaps kv by layer so auto_fa cannot use device matching + if (cparams.pshard) { + cparams.auto_fa = false; + cparams.auto_fgdn = false; + } + // with causal attention, the batch size is limited by the context size cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch; @@ -219,12 +230,24 @@ llama_context::llama_context( if (!hparams.vocab_only) { // GPU backends - for (const auto & dev : model.devices) { - ggml_backend_t backend = ggml_backend_dev_init(dev.dev, nullptr); - if (backend == nullptr) { - throw std::runtime_error(format("failed to initialize %s backend", ggml_backend_dev_name(dev.dev))); + if (cparams.pshard && !model.devices.empty()) { + for (const auto & dev : model.devices) { + for (int i = 0; i < PSHARD_BACKENDS_PER_DEV; i++) { + ggml_backend_t backend = ggml_backend_dev_init(dev.dev, nullptr); + if (backend == nullptr) { + throw std::runtime_error(format("failed to initialize %s backend for pshard", ggml_backend_dev_name(dev.dev))); + } + backends.emplace_back(backend); + } + } + } else { + for (const auto & dev : model.devices) { + ggml_backend_t backend = ggml_backend_dev_init(dev.dev, nullptr); + if (backend == nullptr) { + throw std::runtime_error(format("failed to initialize %s backend", ggml_backend_dev_name(dev.dev))); + } + backends.emplace_back(backend); } - backends.emplace_back(backend); } // add ACCEL backends (such as BLAS) @@ -246,6 +269,11 @@ llama_context::llama_context( } backends.emplace_back(backend_cpu); + if (cparams.pshard) { + cparams.cpu_backend_id = (int32_t)(backends.size() - 1); + pshard_layout = pshard_dev_layout::for_device(0, cparams.cpu_backend_id); + } + // create a list of the set_n_threads functions in the backends for (auto & backend : backends) { ggml_backend_dev_t dev = ggml_backend_get_device(backend.get()); @@ -400,7 +428,15 @@ void llama_context::sched_reserve() { const int64_t t_start_us = ggml_time_us(); const uint32_t n_seqs = cparams.n_seq_max; - const uint32_t n_tokens = std::min(cparams.n_ctx, cparams.n_ubatch); + uint32_t n_tokens = std::min(cparams.n_ctx, cparams.n_ubatch); + + // pshard planner probes graphs up to the largest tier batch, which can exceed n_ubatch + if (cparams.pshard) { + const auto * registry = model.get_plan_registry(); + if (registry && !registry->tier_sizes.empty()) { + n_tokens = std::max(n_tokens, registry->tier_sizes.back()); + } + } const size_t max_nodes = this->graph_max_nodes(n_tokens); @@ -411,6 +447,10 @@ void llama_context::sched_reserve() { sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload)); + if (cparams.pshard) { + ggml_backend_sched_set_prefetch_weights(sched.get(), true); + } + llama_memory_context_ptr mctx; if (memory) { LLAMA_LOG_DEBUG("%s: reserving full memory module\n", __func__); @@ -466,7 +506,7 @@ void llama_context::sched_reserve() { cparams.auto_fa = false; } - if (cparams.auto_fgdn) { + if (cparams.auto_fgdn && !(cparams.pshard && model.hparams.no_alloc)) { LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", __func__); if (cparams.fused_gdn_ar) { @@ -555,73 +595,83 @@ void llama_context::sched_reserve() { int n_splits_tg = -1; int n_nodes_tg = -1; - // reserve pp (prompt processing) graph first so that buffers are only allocated once - { - auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get(), - model.hparams.no_alloc, model.hparams.no_alloc ? backend_buf_exp_size.data() : nullptr); + if (cparams.pshard && model.hparams.no_alloc) { + const uint32_t reserve_tokens = probe_reserve.n_tokens ? probe_reserve.n_tokens : n_tokens; + const uint32_t reserve_outputs = probe_reserve.n_outputs ? probe_reserve.n_outputs : reserve_tokens; + auto * gf = graph_reserve(reserve_tokens, n_seqs, reserve_outputs, mctx.get(), + true, backend_buf_exp_size.data()); if (!gf) { - if (cparams.pipeline_parallel) { - 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)); - gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get()); - } + throw std::runtime_error("failed to measure compute buffers for pshard probe"); + } + } else if (cparams.pshard) { + LLAMA_LOG_INFO("%s: pshard enabled, skipping baseline reserves (deferred to first plan apply)\n", __func__); + } else { + + // reserve pp (prompt processing) graph first so that buffers are only allocated once + { + auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get(), + model.hparams.no_alloc, model.hparams.no_alloc ? backend_buf_exp_size.data() : nullptr); if (!gf) { - throw std::runtime_error("failed to allocate compute pp buffers"); + if (cparams.pipeline_parallel) { + 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)); + gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get()); + } + if (!gf) { + throw std::runtime_error("failed to allocate compute pp buffers"); + } } + + n_splits_pp = ggml_backend_sched_get_n_splits(sched.get()); + n_nodes_pp = ggml_graph_n_nodes(gf); } - n_splits_pp = ggml_backend_sched_get_n_splits(sched.get()); - n_nodes_pp = ggml_graph_n_nodes(gf); - } + // reserve with tg (token generation) graph to get the number of splits and nodes + { + auto * gf = graph_reserve(n_seqs, n_seqs, n_seqs, mctx.get(), model.hparams.no_alloc); + if (!gf) { + throw std::runtime_error("failed to allocate compute tg buffers"); + } - // reserve with tg (token generation) graph to get the number of splits and nodes - { - auto * gf = graph_reserve(n_seqs, n_seqs, n_seqs, mctx.get(), model.hparams.no_alloc); - if (!gf) { - throw std::runtime_error("failed to allocate compute tg buffers"); + n_splits_tg = ggml_backend_sched_get_n_splits(sched.get()); + n_nodes_tg = ggml_graph_n_nodes(gf); } - n_splits_tg = ggml_backend_sched_get_n_splits(sched.get()); - n_nodes_tg = ggml_graph_n_nodes(gf); - } - - // reserve again with pp graph to avoid ggml-alloc reallocations during inference - { - // TODO: not sure if the following graph would be worst case for multi-stream KV caches: - // - // auto * gf = graph_reserve(n_tokens, 1, n_tokens, mctx.get()); - // - auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get(), model.hparams.no_alloc); - if (!gf) { - throw std::runtime_error("failed to allocate compute pp buffers"); + // reserve again with pp graph to avoid ggml-alloc reallocations during inference + { + auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get(), model.hparams.no_alloc); + if (!gf) { + throw std::runtime_error("failed to allocate compute pp buffers"); + } } - } - for (size_t i = 0; i < backend_ptrs.size(); ++i) { - ggml_backend_t backend = backend_ptrs[i]; - ggml_backend_buffer_type_t buft = backend_buft[i]; - if (!model.hparams.no_alloc) { - backend_buf_exp_size[i] = ggml_backend_sched_get_buffer_size(sched.get(), backend); + for (size_t i = 0; i < backend_ptrs.size(); ++i) { + ggml_backend_t backend = backend_ptrs[i]; + ggml_backend_buffer_type_t buft = backend_buft[i]; + if (!model.hparams.no_alloc) { + backend_buf_exp_size[i] = ggml_backend_sched_get_buffer_size(sched.get(), backend); + } + if (backend_buf_exp_size[i] > 1) { + LLAMA_LOG_INFO("%s: %10s compute buffer size = %8.2f MiB\n", __func__, + ggml_backend_buft_name(buft), + backend_buf_exp_size[i] / 1024.0 / 1024.0); + } } - if (backend_buf_exp_size[i] > 1) { - LLAMA_LOG_INFO("%s: %10s compute buffer size = %8.2f MiB\n", __func__, - ggml_backend_buft_name(buft), - backend_buf_exp_size[i] / 1024.0 / 1024.0); + + if (n_nodes_pp == n_nodes_tg) { + LLAMA_LOG_INFO("%s: graph nodes = %d\n", __func__, n_nodes_pp); + } else { + LLAMA_LOG_INFO("%s: graph nodes = %d (with bs=%d), %d (with bs=1)\n", __func__, n_nodes_pp, n_tokens, n_nodes_tg); } - } - if (n_nodes_pp == n_nodes_tg) { - LLAMA_LOG_INFO("%s: graph nodes = %d\n", __func__, n_nodes_pp); - } else { - LLAMA_LOG_INFO("%s: graph nodes = %d (with bs=%d), %d (with bs=1)\n", __func__, n_nodes_pp, n_tokens, n_nodes_tg); - } + if (n_splits_pp == n_splits_tg) { + LLAMA_LOG_INFO("%s: graph splits = %d\n", __func__, n_splits_pp); + } else { + LLAMA_LOG_INFO("%s: graph splits = %d (with bs=%d), %d (with bs=1)\n", __func__, n_splits_pp, n_tokens, n_splits_tg); + } - if (n_splits_pp == n_splits_tg) { - LLAMA_LOG_INFO("%s: graph splits = %d\n", __func__, n_splits_pp); - } else { - LLAMA_LOG_INFO("%s: graph splits = %d (with bs=%d), %d (with bs=1)\n", __func__, n_splits_pp, n_tokens, n_splits_tg); - } + } // !cparams.pshard const int64_t t_end_us = ggml_time_us(); @@ -1211,6 +1261,10 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll return nullptr; } + if (cparams.pshard) { + pshard_assign_tensors(sched.get(), model, memory.get(), backends, pshard_layout); + } + if (!ggml_backend_sched_alloc_graph(sched.get(), gf)) { LLAMA_LOG_ERROR("%s: failed to allocate graph\n", __func__); ret = GGML_STATUS_ALLOC_FAILED; @@ -2129,6 +2183,10 @@ ggml_cgraph * llama_context::graph_reserve( this->n_outputs = save_n_outputs; + if (cparams.pshard) { + pshard_assign_tensors(sched.get(), model, memory.get(), backends, pshard_layout); + } + // initialize scheduler with the specified graph if (split_only) { if (sizes) { @@ -2199,17 +2257,131 @@ ggml_status llama_context::graph_compute( } llm_graph_cb llama_context::graph_get_cb() const { - return [&](const llama_ubatch & ubatch, ggml_tensor * cur, const char * name, int il) { + // keep weight islands on last_weight_bid and return mixed tails to layer_bid + const bool delegate_compute = model.pshard_delegates_compute(); + int32_t last_weight_bid = -1; + int last_il = -2; + return [&, delegate_compute, last_weight_bid, last_il](const llama_ubatch & ubatch, ggml_tensor * cur, const char * name, int il) mutable { if (il >= 0) { ggml_format_name(cur, "%s-%d", name, il); } else { ggml_set_name(cur, name); } + // In delegated mode, layer nodes stay with the scheduler; non-layer heads/tails still honor exact tensor overrides. + if (cparams.pshard && delegate_compute && cur->view_src == nullptr) { + const auto & lbids = model.get_layer_backend_ids(); + const bool has_layer_backend = il >= 0 && lbids.find(il) != lbids.end(); + + if (!has_layer_backend) { + const auto & tbids = model.get_tensor_backend_ids(); + for (int j = 0; j < GGML_MAX_SRC; j++) { + ggml_tensor * src = cur->src[j]; + if (!src) { + continue; + } + + while (src->view_src) { + src = src->view_src; + } + + auto it = tbids.find(src); + if (it != tbids.end() && it->second >= 0 && it->second < (int32_t) backends.size()) { + ggml_backend_sched_set_tensor_backend(sched.get(), cur, backends[it->second].get()); + break; + } + } + } + } + + if (cparams.pshard && !delegate_compute) { + auto sched_backend_id = [&](ggml_backend_t backend) -> int32_t { + if (backend == nullptr) { + return -1; + } + for (int32_t i = 0; i < (int32_t) backends.size(); i++) { + if (backends[i].get() == backend) { + return i; + } + } + return -1; + }; + + int32_t bid = -1; + bool has_tensor_backend = false; + + if (il != last_il) { + last_il = il; + last_weight_bid = -1; + } + + const auto & tbids = model.get_tensor_backend_ids(); + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (!cur->src[j]) { + continue; + } + auto it = tbids.find(cur->src[j]); + if (it != tbids.end() && it->second >= 0) { + bid = it->second; + has_tensor_backend = true; + break; + } + } + + const auto & lbids = model.get_layer_backend_ids(); + auto lit = lbids.find(il); + const int32_t layer_bid = (lit != lbids.end()) ? lit->second : -1; + + bool mixed_src_backends = false; + if (!has_tensor_backend) { + int32_t src_bid = -1; + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (!cur->src[j]) { + continue; + } + + ggml_tensor * src = cur->src[j]; + ggml_backend_t src_backend = ggml_backend_sched_get_tensor_backend(sched.get(), src); + while (src_backend == nullptr && src->view_src != nullptr) { + src = src->view_src; + src_backend = ggml_backend_sched_get_tensor_backend(sched.get(), src); + } + + const int32_t cur_src_bid = sched_backend_id(src_backend); + if (cur_src_bid < 0) { + continue; + } + + if (src_bid < 0) { + src_bid = cur_src_bid; + } else if (src_bid != cur_src_bid) { + mixed_src_backends = true; + break; + } + } + } + + if (has_tensor_backend) { + last_weight_bid = bid; + } else if (mixed_src_backends && last_weight_bid >= 0 && layer_bid >= 0) { + bid = layer_bid; + last_weight_bid = bid; + } else if (last_weight_bid >= 0) { + bid = last_weight_bid; + } else { + bid = layer_bid; + } + + // views keep following their source storage. + if (bid >= 0 && bid < (int32_t)backends.size() && cur->view_src == nullptr) { + ggml_backend_sched_set_tensor_backend(sched.get(), cur, backends[bid].get()); + } + } + // norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends // FIXME: fix in ggml_backend_sched const bool full_offload = model.n_gpu_layers() > model.hparams.n_layer; - if (ubatch.n_tokens < 32 || full_offload) { + if (!cparams.pshard && (ubatch.n_tokens < 32 || full_offload)) { if (il != -1 && strcmp(name, "norm") == 0) { const auto & dev_layer = model.dev_layer(il); for (const auto & backend : backends) { @@ -2918,14 +3090,16 @@ llama_context_params llama_context_default_params() { /*.kv_unified =*/ false, /*.sampler =*/ nullptr, /*.n_sampler =*/ 0, + /*.pshard =*/ false, }; return result; } -llama_context * llama_init_from_model( +llama_context * llama_init_from_model_internal( llama_model * model, - llama_context_params params) { + llama_context_params params, + llama_context_probe_reserve probe_reserve) { if (!model) { LLAMA_LOG_ERROR("%s: model cannot be NULL\n", __func__); return nullptr; @@ -2996,7 +3170,7 @@ llama_context * llama_init_from_model( } try { - auto * ctx = new llama_context(*model, params); + auto * ctx = new llama_context(*model, params, probe_reserve); return ctx; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: failed to initialize the context: %s\n", __func__, err.what()); @@ -3005,6 +3179,12 @@ llama_context * llama_init_from_model( return nullptr; } +llama_context * llama_init_from_model( + llama_model * model, + llama_context_params params) { + return llama_init_from_model_internal(model, params, {}); +} + // deprecated llama_context * llama_new_context_with_model( llama_model * model, diff --git a/src/llama-context.h b/src/llama-context.h index e0d0085c1c3f..f263bba2bc83 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -22,6 +22,11 @@ class llama_io_write_i; struct llama_memory_i; struct llama_memory_context_i; +struct llama_context_probe_reserve { + uint32_t n_tokens = 0; + uint32_t n_outputs = 0; +}; + // "memory" as in physical memory for a buffer type, in bytes struct llama_memory_breakdown_data { size_t model = 0; // memory allocated for the model @@ -37,7 +42,8 @@ struct llama_context { // init scheduler and compute buffers, reserve worst-case graphs llama_context( const llama_model & model, - llama_context_params params); + llama_context_params params, + llama_context_probe_reserve probe_reserve = {}); ~llama_context(); @@ -257,6 +263,7 @@ struct llama_context { const llama_model & model; llama_cparams cparams; + llama_context_probe_reserve probe_reserve; llama_adapter_cvec_ptr cvec; llama_adapter_loras_ptr loras; @@ -313,6 +320,8 @@ struct llama_context { bool sched_need_reserve = true; + pshard_dev_layout pshard_layout = {}; + ggml_backend_t backend_cpu = nullptr; std::vector backends; @@ -357,3 +366,18 @@ struct llama_context { mutable int32_t n_reused = 0; // number of times the previous graph was reused }; + +llama_context * llama_init_from_model_internal( + llama_model * model, + llama_context_params params, + llama_context_probe_reserve probe_reserve = {}); + +// pshard free functions (implemented in llama-context-pshard.cpp) +struct llama_memory_i; + +void pshard_assign_tensors( + ggml_backend_sched_t sched, + const llama_model & model, + llama_memory_i * memory, + const std::vector & backends, + const pshard_dev_layout & layout); diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 9d359474132f..8ac9802da8d4 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -1,6 +1,7 @@ #pragma once #include "llama.h" +#include "ggml-backend.h" #include @@ -44,4 +45,30 @@ struct llama_cparams { ggml_backend_sched_eval_callback cb_eval; void * cb_eval_user_data; + bool pshard = false; + int32_t cpu_backend_id = -1; +}; + +inline constexpr int32_t PSHARD_BACKENDS_PER_DEV = 3; + +struct pshard_dev_layout { + int32_t compute; + int32_t shard_a; + int32_t shard_b; + int32_t cpu; + + int32_t shard(uint32_t il) const { return shard_a + (il % 2); } + + static pshard_dev_layout for_device(size_t dev_idx, int32_t cpu_backend_id) { + const int32_t base = (int32_t)(dev_idx * PSHARD_BACKENDS_PER_DEV); + return { base, base + 1, base + 2, cpu_backend_id }; + } + + static int32_t compute_cpu_backend_id(size_t n_devices) { + int32_t n_accel = 0; + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + if (ggml_backend_dev_type(ggml_backend_dev_get(i)) == GGML_BACKEND_DEVICE_TYPE_ACCEL) n_accel++; + } + return (int32_t)(n_devices * PSHARD_BACKENDS_PER_DEV) + n_accel; + } }; diff --git a/src/llama-kv-cache-iswa.cpp b/src/llama-kv-cache-iswa.cpp index 26e2cb4270b0..35fdfa8cc964 100644 --- a/src/llama-kv-cache-iswa.cpp +++ b/src/llama-kv-cache-iswa.cpp @@ -247,6 +247,13 @@ llama_kv_cache * llama_kv_cache_iswa::get_swa() const { return kv_swa.get(); } +std::vector llama_kv_cache_iswa::get_pipe_shards() { + std::vector result; + if (kv_base && kv_base->get_pipe_shard()) result.push_back(kv_base->get_pipe_shard()); + if (kv_swa && kv_swa->get_pipe_shard()) result.push_back(kv_swa->get_pipe_shard()); + return result; +} + // // llama_kv_cache_iswa_context // diff --git a/src/llama-kv-cache-iswa.h b/src/llama-kv-cache-iswa.h index 70ab22f0d608..b49e37060f5c 100644 --- a/src/llama-kv-cache-iswa.h +++ b/src/llama-kv-cache-iswa.h @@ -58,6 +58,8 @@ class llama_kv_cache_iswa : public llama_memory_i { std::map memory_breakdown() const override; + std::vector get_pipe_shards() override; + // state write/load void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 09102f549c8e..70621c2b8462 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -158,6 +158,75 @@ llama_kv_cache::llama_kv_cache( const bool is_mla = hparams.is_mla(); + if (model.is_pshard()) { + std::vector specs; + for (uint32_t il = 0; il < hparams.n_layer; il++) { + if (!hparams.has_kv(il)) continue; + if (filter && !filter(il)) continue; + + // MLA caches only the compressed latent + rope part in K; V is derived at attn time via wv_b + const uint32_t n_embd_k_gqa = is_mla + ? (hparams.n_lora_kv + hparams.n_rot()) + : hparams.n_embd_k_gqa(il); + const uint32_t n_embd_v_gqa = is_mla + ? 0u // MLA: no separate V cache; dim_t2=0 signals "skip t2" + : (!v_trans ? hparams.n_embd_v_gqa(il) : hparams.n_embd_v_gqa_max()); + + specs.push_back({ + /*.il =*/ il, + /*.type_t1 =*/ type_k, + /*.type_t2 =*/ type_v, + /*.dim_t1 =*/ n_embd_k_gqa, + /*.dim_t2 =*/ n_embd_v_gqa, + /*.seq_len =*/ kv_size, + /*.n_stream =*/ n_stream, + /*.is_1d =*/ false, + /*.name_t1 =*/ "cache_k", + /*.name_t2 =*/ "cache_v", + }); + } + + pipe_shard_kv = std::make_unique(); + auto * ps = pipe_shard_kv.get(); + + pipe_shard_kv->on_activate_gpu = [this, ps](int32_t il, ggml_tensor * k, ggml_tensor * v) { + auto idx = map_layer_ids[il]; + layers[idx].k = k; + layers[idx].v = v; + const auto & sv = ps->get_stream(idx); + layers[idx].k_stream = sv.t1_stream_gpu; + layers[idx].v_stream = sv.t2_stream_gpu; + }; + + pipe_shard_kv->on_activate_cpu = [this, ps](int32_t il, ggml_tensor * k, ggml_tensor * v) { + auto idx = map_layer_ids[il]; + layers[idx].k = k; + layers[idx].v = v; + const auto & sv = ps->get_stream(idx); + layers[idx].k_stream = sv.t1_stream_cpu; + layers[idx].v_stream = sv.t2_stream_cpu; + }; + + const int32_t cpu_bid = pshard_dev_layout::compute_cpu_backend_id(model.devices.size()); + if (!pipe_shard_kv->init(specs, model.get_layer_backend_ids(), cpu_bid, hparams.no_alloc)) { + throw std::runtime_error("failed to initialize KV pipe shard"); + } + + const auto & ps_layers = pipe_shard_kv->get_layers(); + layers.resize(ps_layers.size()); + for (size_t i = 0; i < ps_layers.size(); ++i) { + layers[i].il = ps_layers[i].il; + const bool cpu = pipe_shard_kv->is_cpu_only(ps_layers[i].il); + layers[i].k = cpu ? ps_layers[i].t1_cpu : ps_layers[i].t1_gpu; + layers[i].v = cpu ? ps_layers[i].t2_cpu : ps_layers[i].t2_gpu; + const auto & sv = ps->get_stream(i); + layers[i].k_stream = cpu ? sv.t1_stream_cpu : sv.t1_stream_gpu; + layers[i].v_stream = cpu ? sv.t2_stream_cpu : sv.t2_stream_gpu; + map_layer_ids[ps_layers[i].il] = (int32_t)i; + } + return; + } + for (uint32_t il = 0; il < hparams.n_layer; il++) { if (!hparams.has_kv(il)) { LLAMA_LOG_DEBUG("%s: layer %3d: does not have KV cache\n", __func__, il); @@ -620,9 +689,24 @@ std::map llama_kv_cache::memory_breakdown() } } + if (pipe_shard_kv) { + const auto & bufs = pipe_shard_kv->get_bufs(); + const auto & sizes = pipe_shard_kv->get_bufs_planned_sizes(); + for (size_t i = 0; i < bufs.size(); i++) { + if (!bufs[i]) continue; + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(bufs[i].get()); + ret[buft] += sizes[i]; + } + } + return ret; } +std::vector llama_kv_cache::get_pipe_shards() { + if (pipe_shard_kv) return { pipe_shard_kv.get() }; + return {}; +} + llama_memory_context_ptr llama_kv_cache::init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 0b62dc7b2320..69460133d006 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -4,6 +4,7 @@ #include "llama-graph.h" #include "llama-kv-cells.h" #include "llama-memory.h" +#include "llama-memory-pshard.h" #include #include @@ -152,6 +153,10 @@ class llama_kv_cache : public llama_memory_i { bool get_has_shift() const; + // pshard dual-buffer KV + llama_memory_pshard * get_pipe_shard() { return pipe_shard_kv.get(); } + std::vector get_pipe_shards() override; + ggml_type type_k() const; ggml_type type_v() const; @@ -270,6 +275,8 @@ class llama_kv_cache : public llama_memory_i { std::vector layers; + std::unique_ptr pipe_shard_kv; + // model layer id -> KV cache layer id std::unordered_map map_layer_ids; diff --git a/src/llama-memory-hybrid-iswa.cpp b/src/llama-memory-hybrid-iswa.cpp index 10e6b4597972..fd44dd632774 100644 --- a/src/llama-memory-hybrid-iswa.cpp +++ b/src/llama-memory-hybrid-iswa.cpp @@ -182,6 +182,13 @@ std::map llama_memory_hybrid_iswa::memory_br return mb; } +std::vector llama_memory_hybrid_iswa::get_pipe_shards() { + std::vector result; + for (auto * ps : mem_attn->get_pipe_shards()) { result.push_back(ps); } + for (auto * ps : mem_recr->get_pipe_shards()) { result.push_back(ps); } + return result; +} + void llama_memory_hybrid_iswa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { mem_attn->state_write(io, seq_id, flags); mem_recr->state_write(io, seq_id, flags); diff --git a/src/llama-memory-hybrid-iswa.h b/src/llama-memory-hybrid-iswa.h index 807c8aac96c9..090eff3351ce 100644 --- a/src/llama-memory-hybrid-iswa.h +++ b/src/llama-memory-hybrid-iswa.h @@ -70,6 +70,8 @@ class llama_memory_hybrid_iswa : public llama_memory_i { std::map memory_breakdown() const override; + std::vector get_pipe_shards() override; + // state write/load void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 4ce1af592c15..4468c4f30050 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -176,6 +176,13 @@ std::map llama_memory_hybrid::memory_breakdo return mb; } +std::vector llama_memory_hybrid::get_pipe_shards() { + std::vector result; + for (auto * ps : mem_attn->get_pipe_shards()) { result.push_back(ps); } + for (auto * ps : mem_recr->get_pipe_shards()) { result.push_back(ps); } + return result; +} + void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { mem_attn->state_write(io, seq_id, flags); diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 558cafdf984c..6d92940d4f1f 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -70,6 +70,8 @@ class llama_memory_hybrid : public llama_memory_i { std::map memory_breakdown() const override; + std::vector get_pipe_shards() override; + // state write/load void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; diff --git a/src/llama-memory-pshard.cpp b/src/llama-memory-pshard.cpp new file mode 100644 index 000000000000..eb47f2b654ee --- /dev/null +++ b/src/llama-memory-pshard.cpp @@ -0,0 +1,253 @@ +#include "llama-memory-pshard.h" + +#include "llama-impl.h" +#include "llama-cparams.h" + +#include "ggml-backend.h" + +#include + +bool llama_memory_pshard::init( + const std::vector & specs, + const std::unordered_map & layer_backend_ids, + int32_t cpu_backend_id, + bool no_alloc) { + + layers.clear(); + streams.clear(); + ctxs.clear(); + bufs.clear(); + bufs_planned_sizes.clear(); + map_layer_ids.clear(); + cpu_bid_ = cpu_backend_id; + layer_bids_ = layer_backend_ids; + + if (specs.empty()) return true; + + auto * main_gpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU); + auto * buft_gpu = ggml_backend_dev_buffer_type(main_gpu_dev); + auto * buft_cpu_host = ggml_backend_dev_host_buffer_type(main_gpu_dev); + + auto is_sharded = [&](uint32_t il) -> bool { + auto it = layer_backend_ids.find(il); + return it == layer_backend_ids.end() || it->second != 0; + }; + + size_t n_pinned = 0; + size_t n_sharded = 0; + for (const auto & sp : specs) { + if (is_sharded(sp.il)) { n_sharded++; } else { n_pinned++; } + } + + uint32_t max_n_stream = 1; + for (const auto & sp : specs) { max_n_stream = std::max(max_n_stream, sp.n_stream); } + + ggml_context * ctx_gpu_pinned = nullptr; + ggml_context * ctx_gpu_sharded = nullptr; + + auto make_ctx = [&](size_t n, uint32_t ns) -> ggml_context * { + size_t overhead = 2u * (1 + ns) * n * ggml_tensor_overhead(); + ggml_init_params params = { overhead, NULL, true }; + ggml_context * c = ggml_init(params); + if (c) ctxs.emplace_back(c); + return c; + }; + + if (n_pinned > 0) { ctx_gpu_pinned = make_ctx(n_pinned, max_n_stream); } + if (n_sharded > 0) { ctx_gpu_sharded = make_ctx(n_sharded, max_n_stream); } + if ((n_pinned > 0 && !ctx_gpu_pinned) || (n_sharded > 0 && !ctx_gpu_sharded)) { + LLAMA_LOG_ERROR("%s: failed to create GPU tensor contexts\n", __func__); + return false; + } + + for (const auto & sp : specs) { + const bool sharded = is_sharded(sp.il); + ggml_context * ctx = sharded ? ctx_gpu_sharded : ctx_gpu_pinned; + + // dim_t2 == 0 means "no second tensor" (MLA caches K-only) + ggml_tensor * t1 = nullptr; + ggml_tensor * t2 = nullptr; + if (sp.is_1d) { + t1 = ggml_new_tensor_1d(ctx, sp.type_t1, sp.dim_t1); + if (sp.dim_t2 > 0) { + t2 = ggml_new_tensor_1d(ctx, sp.type_t2, sp.dim_t2); + } + } else { + t1 = ggml_new_tensor_3d(ctx, sp.type_t1, sp.dim_t1, sp.seq_len, sp.n_stream); + if (sp.dim_t2 > 0) { + t2 = ggml_new_tensor_3d(ctx, sp.type_t2, sp.dim_t2, sp.seq_len, sp.n_stream); + } + } + + ggml_format_name(t1, "%s_l%d", sp.name_t1, sp.il); + if (t2) ggml_format_name(t2, "%s_l%d", sp.name_t2, sp.il); + + stream_views sv; + if (!sp.is_1d && sp.n_stream > 0) { + for (uint32_t s = 0; s < sp.n_stream; ++s) { + sv.t1_stream_gpu.push_back(ggml_view_2d(ctx, t1, sp.dim_t1, sp.seq_len, t1->nb[1], s * t1->nb[2])); + if (t2) { + sv.t2_stream_gpu.push_back(ggml_view_2d(ctx, t2, sp.dim_t2, sp.seq_len, t2->nb[1], s * t2->nb[2])); + } + } + } + + map_layer_ids[sp.il] = (int32_t)layers.size(); + + layer l; + l.il = sp.il; + l.t1_gpu = t1; + l.t2_gpu = t2; + l.t1_cpu = nullptr; + l.t2_cpu = nullptr; + layers.push_back(std::move(l)); + streams.push_back(std::move(sv)); + } + + if (ctx_gpu_pinned) { + const size_t planned_size = ggml_backend_alloc_ctx_tensors_from_buft_size(ctx_gpu_pinned, buft_gpu); + + ggml_backend_buffer_t buf; + if (no_alloc) { + // dummy 0-size buffer; size accounted via bufs_planned_sizes for memory_breakdown. + buf = ggml_backend_buft_alloc_buffer(buft_gpu, 0); + for (ggml_tensor * t = ggml_get_first_tensor(ctx_gpu_pinned); t != nullptr; t = ggml_get_next_tensor(ctx_gpu_pinned, t)) { + t->buffer = buf; + } + } else { + buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx_gpu_pinned, buft_gpu); + } + if (!buf) { + LLAMA_LOG_ERROR("%s: failed to allocate GPU buffer for pinned layers\n", __func__); + return false; + } + if (!no_alloc) { + ggml_backend_buffer_clear(buf, 0); + } + + LLAMA_LOG_INFO("%s: %10s pinned buffer = %8.2f MiB (%zu layers)%s\n", + __func__, ggml_backend_buffer_name(buf), + planned_size / 1024.0 / 1024.0, n_pinned, + no_alloc ? " (no_alloc)" : ""); + + bufs.emplace_back(buf); + bufs_planned_sizes.push_back(planned_size); + } + + { + size_t overhead = 2u * (1 + max_n_stream) * specs.size() * ggml_tensor_overhead(); + ggml_init_params params = { overhead, NULL, true }; + ggml_context * ctx_cpu = ggml_init(params); + if (!ctx_cpu) { + LLAMA_LOG_ERROR("%s: failed to create CPU context\n", __func__); + return false; + } + ctxs.emplace_back(ctx_cpu); + + for (size_t i = 0; i < specs.size(); ++i) { + const auto & sp = specs[i]; + ggml_tensor * t1_cpu = nullptr; + ggml_tensor * t2_cpu = nullptr; + + if (sp.is_1d) { + t1_cpu = ggml_new_tensor_1d(ctx_cpu, sp.type_t1, sp.dim_t1); + if (sp.dim_t2 > 0) { + t2_cpu = ggml_new_tensor_1d(ctx_cpu, sp.type_t2, sp.dim_t2); + } + } else { + t1_cpu = ggml_new_tensor_3d(ctx_cpu, sp.type_t1, sp.dim_t1, sp.seq_len, sp.n_stream); + if (sp.dim_t2 > 0) { + t2_cpu = ggml_new_tensor_3d(ctx_cpu, sp.type_t2, sp.dim_t2, sp.seq_len, sp.n_stream); + } + } + + ggml_format_name(t1_cpu, "%s_cpu_l%d", sp.name_t1, sp.il); + if (t2_cpu) ggml_format_name(t2_cpu, "%s_cpu_l%d", sp.name_t2, sp.il); + + layers[i].t1_cpu = t1_cpu; + layers[i].t2_cpu = t2_cpu; + + if (!sp.is_1d && sp.n_stream > 0) { + for (uint32_t s = 0; s < sp.n_stream; ++s) { + streams[i].t1_stream_cpu.push_back(ggml_view_2d(ctx_cpu, t1_cpu, sp.dim_t1, sp.seq_len, t1_cpu->nb[1], s * t1_cpu->nb[2])); + if (t2_cpu) { + streams[i].t2_stream_cpu.push_back(ggml_view_2d(ctx_cpu, t2_cpu, sp.dim_t2, sp.seq_len, t2_cpu->nb[1], s * t2_cpu->nb[2])); + } + } + } + } + + const size_t planned_size_cpu = ggml_backend_alloc_ctx_tensors_from_buft_size(ctx_cpu, buft_cpu_host); + + ggml_backend_buffer_t buf_cpu; + if (no_alloc) { + buf_cpu = ggml_backend_buft_alloc_buffer(buft_cpu_host, 0); + for (ggml_tensor * t = ggml_get_first_tensor(ctx_cpu); t != nullptr; t = ggml_get_next_tensor(ctx_cpu, t)) { + t->buffer = buf_cpu; + } + } else { + buf_cpu = ggml_backend_alloc_ctx_tensors_from_buft(ctx_cpu, buft_cpu_host); + } + if (!buf_cpu) { + LLAMA_LOG_ERROR("%s: failed to allocate CPU-pinned buffer\n", __func__); + return false; + } + if (!no_alloc) { + ggml_backend_buffer_clear(buf_cpu, 0); + } + bufs.emplace_back(buf_cpu); + bufs_planned_sizes.push_back(planned_size_cpu); + } + + LLAMA_LOG_INFO("%s: %zu pinned, %zu sharded\n", __func__, n_pinned, n_sharded); + return true; +} + +bool llama_memory_pshard::is_cpu_only(int32_t il) const { + auto it = layer_bids_.find((int)il); + return it != layer_bids_.end() && it->second == cpu_bid_; +} + +void llama_memory_pshard::activate_gpu(int32_t il) { + auto it = map_layer_ids.find(il); + if (it == map_layer_ids.end() || !on_activate_gpu) return; + auto & l = layers[it->second]; + on_activate_gpu(il, l.t1_gpu, l.t2_gpu); +} + +void llama_memory_pshard::activate_cpu(int32_t il) { + auto it = map_layer_ids.find(il); + if (it == map_layer_ids.end() || !on_activate_cpu) return; + auto & l = layers[it->second]; + on_activate_cpu(il, l.t1_cpu, l.t2_cpu); +} + +void llama_memory_pshard::assign_tensors( + ggml_backend_sched_t sched, + const std::unordered_map & layer_bids, + const std::vector & backends, + const pshard_dev_layout & layout) { + for (auto & l : layers) { + auto it = layer_bids.find((int)l.il); + const bool has_bid = it != layer_bids.end() && it->second >= 0 && it->second < (int32_t)backends.size(); + if (!has_bid) { + continue; + } + const int32_t bid = it->second; + if (bid == layout.compute) { + activate_gpu(l.il); + } else if (bid == layout.cpu) { + activate_cpu(l.il); + } else { + activate_gpu(l.il); + l.t1_gpu->data = NULL; l.t1_gpu->buffer = NULL; + ggml_backend_sched_set_tensor_backend(sched, l.t1_gpu, backends[bid].get()); + ggml_backend_sched_add_writeback(sched, l.t1_gpu); + if (l.t2_gpu) { + l.t2_gpu->data = NULL; l.t2_gpu->buffer = NULL; + ggml_backend_sched_set_tensor_backend(sched, l.t2_gpu, backends[bid].get()); + ggml_backend_sched_add_writeback(sched, l.t2_gpu); + } + } + } +} diff --git a/src/llama-memory-pshard.h b/src/llama-memory-pshard.h new file mode 100644 index 000000000000..8c1093e13c3b --- /dev/null +++ b/src/llama-memory-pshard.h @@ -0,0 +1,76 @@ +#pragma once + +#include "llama-pipe-shard.h" + +#include +#include +#include + +struct ggml_tensor; + +struct llama_memory_pshard : llama_memory_pipe_shard_i { + struct stream_views { + std::vector t1_stream_gpu; + std::vector t2_stream_gpu; + std::vector t1_stream_cpu; + std::vector t2_stream_cpu; + }; + + // parent cache callbacks set before init + using activate_fn_t = std::function; + + activate_fn_t on_activate_gpu; + activate_fn_t on_activate_cpu; + + struct tensor_spec { + uint32_t il; + ggml_type type_t1; + ggml_type type_t2; + uint32_t dim_t1; // ne[0] for t1 + uint32_t dim_t2; // ne[0] for t2 + uint32_t seq_len; // ne[1] (kv_size or flat for 1D) + uint32_t n_stream; // ne[2] (1 for RS) + bool is_1d; // true = ggml_new_tensor_1d (RS), false = 3D (KV) + const char * name_t1 = "cache_k"; // e.g. "cache_k" or "cache_r" + const char * name_t2 = "cache_v"; // e.g. "cache_v" or "cache_s" + }; + + bool init( + const std::vector & specs, + const std::unordered_map & layer_backend_ids, + int32_t cpu_backend_id, + bool no_alloc = false); + + bool is_cpu_only(int32_t il) const; + + ~llama_memory_pshard() override = default; + + // llama_memory_pipe_shard_i interface + const std::vector & get_layers() const override { return layers; } + + void activate_gpu(int32_t il) override; + void activate_cpu(int32_t il) override; + + void assign_tensors( + ggml_backend_sched_t sched, + const std::unordered_map & layer_bids, + const std::vector & backends, + const pshard_dev_layout & layout) override; + + // stream accessors (KV only, empty for RS) + const stream_views & get_stream(size_t idx) const { return streams[idx]; } + const std::vector & get_bufs() const { return bufs; } + + // planned buffer sizes including dummy no_alloc buffers + const std::vector & get_bufs_planned_sizes() const { return bufs_planned_sizes; } + +private: + std::vector layers; + std::vector streams; + std::unordered_map map_layer_ids; + std::vector ctxs; + std::vector bufs; + std::vector bufs_planned_sizes; + int32_t cpu_bid_ = -1; + std::unordered_map layer_bids_; +}; diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 9287fe45e963..9cb24ec31e1a 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -68,6 +68,50 @@ llama_memory_recurrent::llama_memory_recurrent( r_l.resize(n_layer); s_l.resize(n_layer); + if (model.is_pshard()) { + std::vector specs; + for (int i = 0; i < n_layer; i++) { + if (filter && !filter(i)) continue; + specs.push_back({ + /*.il =*/ (uint32_t)i, + /*.type_t1 =*/ type_r, + /*.type_t2 =*/ type_s, + /*.dim_t1 =*/ (uint32_t)(hparams.n_embd_r() * mem_size), + /*.dim_t2 =*/ (uint32_t)(hparams.n_embd_s() * mem_size), + /*.seq_len =*/ 0, + /*.n_stream =*/ 1, + /*.is_1d =*/ true, + /*.name_t1 =*/ "cache_r", + /*.name_t2 =*/ "cache_s", + }); + } + + pipe_shard_rs = std::make_unique(); + + pipe_shard_rs->on_activate_gpu = [this](int32_t il, ggml_tensor * r, ggml_tensor * s) { + r_l[il] = r; + s_l[il] = s; + }; + + pipe_shard_rs->on_activate_cpu = [this](int32_t il, ggml_tensor * r, ggml_tensor * s) { + r_l[il] = r; + s_l[il] = s; + }; + + const int32_t cpu_bid = pshard_dev_layout::compute_cpu_backend_id(model.devices.size()); + if (!pipe_shard_rs->init(specs, model.get_layer_backend_ids(), cpu_bid, hparams.no_alloc)) { + throw std::runtime_error("failed to initialize RS pipe shard"); + } + + const auto & ps_layers = pipe_shard_rs->get_layers(); + for (size_t i = 0; i < ps_layers.size(); ++i) { + const bool cpu = pipe_shard_rs->is_cpu_only(ps_layers[i].il); + r_l[ps_layers[i].il] = cpu ? ps_layers[i].t1_cpu : ps_layers[i].t1_gpu; + s_l[ps_layers[i].il] = cpu ? ps_layers[i].t2_cpu : ps_layers[i].t2_gpu; + } + return; + } + for (int i = 0; i < n_layer; i++) { if (filter && !filter(i)) { LLAMA_LOG_DEBUG("%s: layer %3d: skipped\n", __func__, i); @@ -373,9 +417,24 @@ std::map llama_memory_recurrent::memory_brea for (const auto & [_, buf] : ctxs_bufs) { ret[ggml_backend_buffer_get_type(buf.get())] += ggml_backend_buffer_get_size(buf.get()); } + + if (pipe_shard_rs) { + const auto & bufs = pipe_shard_rs->get_bufs(); + const auto & sizes = pipe_shard_rs->get_bufs_planned_sizes(); + for (size_t i = 0; i < bufs.size(); i++) { + if (!bufs[i]) continue; + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(bufs[i].get()); + ret[buft] += sizes[i]; + } + } return ret; } +std::vector llama_memory_recurrent::get_pipe_shards() { + if (pipe_shard_rs) return { pipe_shard_rs.get() }; + return {}; +} + llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { do { balloc.split_reset(); diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index 47f01d739124..4e6c22cd6e00 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -3,6 +3,7 @@ #include "llama-batch.h" #include "llama-graph.h" #include "llama-memory.h" +#include "llama-memory-pshard.h" #include #include @@ -103,6 +104,10 @@ class llama_memory_recurrent : public llama_memory_i { std::vector r_l; std::vector s_l; + std::unique_ptr pipe_shard_rs; + llama_memory_pshard * get_pipe_shard() { return pipe_shard_rs.get(); } + std::vector get_pipe_shards() override; + private: //const llama_model & model; const llama_hparams & hparams; diff --git a/src/llama-memory.h b/src/llama-memory.h index 4a157b91fdbd..d24f9d7bf668 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -117,6 +117,9 @@ struct llama_memory_i { virtual void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const = 0; virtual void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) = 0; + + // pipe-shard: return all pipe-shard objects for this memory (empty if not pshard) + virtual std::vector get_pipe_shards() { return {}; } }; using llama_memory_ptr = std::unique_ptr; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 4e65a45a50d8..3a3793eecbad 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1154,7 +1154,11 @@ struct ggml_tensor * llama_model_loader::create_tensor( // check overrides if (tensor_buft_overrides) { - std::string tensor_name = tn.str(); + // when a tensor is remapped (e.g. token_embd duplicated as output), + // match overrides against the remapped name so each copy can be + // placed independently (e.g. token_embd on CPU, output on GPU) + std::string tensor_name = (tn_tensor != tn.tensor) + ? LLM_TN_IMPL(tn.arch, tn_tensor, tn.suffix, tn.bid, tn.xid).str() : tn.str(); for (const auto * overrides = tensor_buft_overrides; overrides->pattern != nullptr; ++overrides) { std::regex pattern(overrides->pattern); if (std::regex_search(tensor_name, pattern)) { @@ -1256,7 +1260,8 @@ struct ggml_tensor * llama_model_loader::create_tensor( ggml_context * ctx = ctx_for_buft(buft); // if duplicated, check if the original tensor was allocated in the same buffer type context and avoid creating a new one - if (flags & TENSOR_DUPLICATED) { + // pshard needs separate tensors for output vs token_embd (different GPU/CPU placement) + if ((flags & TENSOR_DUPLICATED) && !force_duplicate_tied) { ggml_tensor * t = ggml_get_tensor(ctx, tn.str().c_str()); if (t) { return t; @@ -1273,7 +1278,18 @@ struct ggml_tensor * llama_model_loader::create_tensor( const bool duplicated = flags & TENSOR_DUPLICATED; struct ggml_tensor * tensor = ggml_dup_tensor(ctx, cur); - ggml_set_name(tensor, ggml_get_name(cur)); + + if (duplicated && force_duplicate_tied) { + ggml_set_name(tensor, "output.weight"); + if (weights_map.find("output.weight") == weights_map.end()) { + auto it = weights_map.find(ggml_get_name(cur)); + if (it != weights_map.end()) { + weights_map.emplace("output.weight", it->second); + } + } + } else { + ggml_set_name(tensor, ggml_get_name(cur)); + } if (duplicated) { size_data += ggml_nbytes(cur); diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 7b3d6703c03d..59bd931d8865 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -79,6 +79,7 @@ struct llama_model_loader { bool use_direct_io = false; bool check_tensors; bool no_alloc; + bool force_duplicate_tied = false; llama_files files; llama_ftype ftype; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index b265394ef736..409be66b2108 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -6,6 +6,7 @@ #include "llama-mmap.h" #include "llama-cparams.h" #include "llama-model-loader.h" +#include "llama-pshard-plan.h" #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" @@ -661,6 +662,11 @@ struct llama_model::impl { std::vector dev_layer; bool has_tensor_overrides; + + std::unordered_map tensor_backend_ids; + std::unordered_map layer_backend_ids; + + llama_pshard_plan_registry * plan_registry = nullptr; }; llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique()) { @@ -8115,6 +8121,38 @@ bool llama_model::load_tensors(llama_model_loader & ml) { } } + // build tensor -> backend_id map from overrides (for pshard scheduling) + if (params.tensor_buft_overrides) { + for (const auto & [name, tensor] : tensors_by_name) { + for (const auto * ov = params.tensor_buft_overrides; ov->pattern; ++ov) { + if (ov->backend_id >= 0 && std::regex_search(name, std::regex(ov->pattern))) { + pimpl->tensor_backend_ids[tensor] = ov->backend_id; + break; + } + } + } + if (!pimpl->tensor_backend_ids.empty()) { + LLAMA_LOG_INFO("%s: built tensor backend_id map: %zu tensors\n", + __func__, pimpl->tensor_backend_ids.size()); + } + + // build layer -> backend_id map from override patterns (blk\.N\..*) + for (const auto * ov = params.tensor_buft_overrides; ov->pattern; ++ov) { + if (ov->backend_id >= 0) { + std::smatch m; + std::string pat(ov->pattern); + if (std::regex_search(pat, m, std::regex(R"(blk\\\.(\d+)\\\.)"))) { + int layer = std::stoi(m[1].str()); + pimpl->layer_backend_ids[layer] = ov->backend_id; + } + } + } + if (!pimpl->layer_backend_ids.empty()) { + LLAMA_LOG_INFO("%s: built layer backend_id map: %zu layers\n", + __func__, pimpl->layer_backend_ids.size()); + } + } + ml.init_mappings(true, use_mlock ? &pimpl->mlock_mmaps : nullptr); pimpl->mappings.reserve(ml.mappings.size()); @@ -8250,6 +8288,8 @@ bool llama_model::load_tensors(llama_model_loader & ml) { } } + pimpl->plan_registry = params.pshard_registry; + return true; } @@ -8574,6 +8614,26 @@ bool llama_model::has_tensor_overrides() const { return pimpl->has_tensor_overrides; } +bool llama_model::is_pshard() const { + return params.pshard; +} + +bool llama_model::pshard_delegates_compute() const { + return params.pshard_delegate_compute; +} + +llama_pshard_plan_registry * llama_model::get_plan_registry() const { + return pimpl->plan_registry; +} + +const std::unordered_map & llama_model::get_tensor_backend_ids() const { + return pimpl->tensor_backend_ids; +} + +const std::unordered_map & llama_model::get_layer_backend_ids() const { + return pimpl->layer_backend_ids; +} + const ggml_tensor * llama_model::get_tensor(const char * name) const { auto it = std::find_if(tensors_by_name.begin(), tensors_by_name.end(), [name](const std::pair & it) { @@ -9304,6 +9364,11 @@ llama_model_params llama_model_default_params() { /*.use_extra_bufts =*/ true, /*.no_host =*/ false, /*.no_alloc =*/ false, + /*.pshard =*/ false, + /*.pshard_delegate_compute =*/ false, + /*.pshard_cache_skip_load =*/ false, + /*.max_vram_alloc =*/ 0, + /*.pshard_registry =*/ nullptr, }; return result; diff --git a/src/llama-model.h b/src/llama-model.h index bba70012e118..762f5187caa4 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -17,6 +17,8 @@ struct llama_cparams; struct llama_ubatch; struct llama_model_loader; +struct llama_pshard_plan; +struct llama_pshard_plan_registry; // available models enum llm_type { @@ -616,6 +618,14 @@ struct llama_model { bool has_tensor_overrides() const; + bool is_pshard() const; + bool pshard_delegates_compute() const; + + llama_pshard_plan_registry * get_plan_registry() const; + + const std::unordered_map & get_tensor_backend_ids() const; + const std::unordered_map & get_layer_backend_ids() const; + const struct ggml_tensor * get_tensor(const char * name) const; float get_rope_freq_base (const llama_cparams & cparams, int il) const; diff --git a/src/llama-pipe-shard.h b/src/llama-pipe-shard.h new file mode 100644 index 000000000000..ae7dbdad3c2c --- /dev/null +++ b/src/llama-pipe-shard.h @@ -0,0 +1,36 @@ +#pragma once + +#include "ggml-cpp.h" + +#include +#include +#include + +struct ggml_tensor; +struct pshard_dev_layout; +typedef struct ggml_backend_sched * ggml_backend_sched_t; + +// pipe-shard memory interface for KV cache or recurrent state. +struct llama_memory_pipe_shard_i { + // t1/t2 are k/v for KV cache, r/s for recurrent state + struct layer { + uint32_t il; + ggml_tensor * t1_gpu; + ggml_tensor * t2_gpu; + ggml_tensor * t1_cpu; + ggml_tensor * t2_cpu; + }; + + virtual ~llama_memory_pipe_shard_i() = default; + + virtual const std::vector & get_layers() const = 0; + + virtual void activate_gpu(int32_t il) = 0; + virtual void activate_cpu(int32_t il) = 0; + + virtual void assign_tensors( + ggml_backend_sched_t sched, + const std::unordered_map & layer_bids, + const std::vector & backends, + const pshard_dev_layout & layout) = 0; +}; diff --git a/src/llama-pshard-plan.cpp b/src/llama-pshard-plan.cpp new file mode 100644 index 000000000000..913ce6743802 --- /dev/null +++ b/src/llama-pshard-plan.cpp @@ -0,0 +1,1988 @@ +#include "llama-pshard-plan.h" + +#include "llama-benchmark.h" +#include "llama-impl.h" +#include "llama-memory.h" +#include "llama-model.h" +#include "llama-model-loader.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +static std::mutex g_probe_mutex; + +static std::vector llama_get_device_memory_data_safe( + const char * path_model, const llama_model_params * mparams, + const llama_context_params * cparams, + std::vector & devs, uint32_t & hp_ngl, + uint32_t & hp_n_ctx_train, uint32_t & hp_n_expert, uint32_t & hp_n_embd_r, + ggml_log_level log_level, + llama_probe_hook_t probe_hook = nullptr, + void * probe_hook_data = nullptr, + uint32_t probe_n_tokens = 0, + uint32_t probe_n_outputs = 0) { + std::lock_guard lock(g_probe_mutex); + return llama_get_device_memory_data(path_model, mparams, cparams, devs, + hp_ngl, hp_n_ctx_train, hp_n_expert, hp_n_embd_r, log_level, probe_hook, probe_hook_data, + probe_n_tokens, probe_n_outputs); +} + +const char * llama_get_overflow_pattern(size_t il, llama_layer_fraction lf) { + constexpr size_t n_strings = 1000; + GGML_ASSERT(il < n_strings); + switch (lf) { + case LLAMA_LAYER_FRACTION_ATTN: { + static std::array p; + if (p[il].empty()) { p[il] = "blk\\." + std::to_string(il) + "\\.ffn_(up|gate|down).*"; } + return p[il].c_str(); + } + case LLAMA_LAYER_FRACTION_UP: { + static std::array p; + if (p[il].empty()) { p[il] = "blk\\." + std::to_string(il) + "\\.ffn_(gate|down).*"; } + return p[il].c_str(); + } + case LLAMA_LAYER_FRACTION_GATE: { + static std::array p; + if (p[il].empty()) { p[il] = "blk\\." + std::to_string(il) + "\\.ffn_down.*"; } + return p[il].c_str(); + } + case LLAMA_LAYER_FRACTION_MOE: { + static std::array p; + if (p[il].empty()) { p[il] = "blk\\." + std::to_string(il) + "\\.ffn_(up|down|gate)_(ch|)exps"; } + return p[il].c_str(); + } + default: + return nullptr; + } +} + +static void llama_pshard_generate_overrides( + uint32_t n_pinned, + uint32_t n_layers, + ggml_backend_buffer_type_t gpu_buft, + ggml_backend_buffer_type_t host_buft, + struct llama_model_tensor_buft_override * tensor_buft_overrides, + llama_layer_fraction overflow_type, + llama_pshard_strategy strategy, + const pshard_dev_layout & layout, + bool pin_from_back = false, + bool output_on_gpu = false, + uint32_t n_attn_pinned = 0) { + thread_local std::array patterns_layer; + thread_local std::array patterns_layer_attn; + thread_local std::array patterns_layer_ffn; + thread_local std::string pat_output = "^output"; + + const uint32_t il_pin_start = pin_from_back ? (n_layers - n_pinned) : 0; + GGML_ASSERT(n_layers <= 1000 && "pshard: n_layers exceeds thread_local pattern array capacity"); + const uint32_t il_pin_end = pin_from_back ? n_layers : n_pinned; + const uint32_t il_boundary_raw = pin_from_back ? (il_pin_start > 0 ? il_pin_start - 1 : UINT32_MAX) : il_pin_end; + const uint32_t il_boundary = (overflow_type != LLAMA_LAYER_FRACTION_NONE && il_boundary_raw < n_layers) ? il_boundary_raw : UINT32_MAX; + const bool output_on_cpu = !output_on_gpu; + + static constexpr size_t OVERRIDE_CAP = 4096; + size_t itbo = 0; + + auto emit = [&](const char * pat, ggml_backend_buffer_type_t buft, int32_t bid) { + GGML_ASSERT(itbo + 1 < OVERRIDE_CAP && "override array overflow"); + tensor_buft_overrides[itbo] = { pat, buft, bid }; + itbo++; + }; + + { + thread_local std::string pat_tok_embd = "^token_embd"; + const int32_t out_bid = output_on_cpu ? layout.cpu : layout.compute; + emit(pat_output.c_str(), output_on_cpu ? host_buft : gpu_buft, out_bid); + emit(pat_tok_embd.c_str(), host_buft, layout.cpu); + } + + // one pattern cache per thread + for (uint32_t il = 0; il < n_layers; il++) { + if (patterns_layer[il].empty()) { patterns_layer[il] = "blk\\." + std::to_string(il) + "\\..*"; } + if (patterns_layer_attn[il].empty()) { patterns_layer_attn[il] = "blk\\." + std::to_string(il) + "\\.attn_(q|k|v|output|q_norm|k_norm).*"; } + if (patterns_layer_ffn[il].empty()) { patterns_layer_ffn[il] = "blk\\." + std::to_string(il) + "\\.ffn_((up|gate|down)\\.|(up|down|gate|gate_up)_(ch|)exps).*"; } + + if (il == il_boundary) { + const char * overflow_pat = llama_get_overflow_pattern(il, overflow_type); + if (overflow_pat) { + emit(overflow_pat, host_buft, layout.shard(il)); + } + emit(patterns_layer[il].c_str(), gpu_buft, layout.compute); + } else if (il >= il_pin_start && il < il_pin_end) { + emit(patterns_layer[il].c_str(), gpu_buft, layout.compute); + } else { + const bool use_alternating_shards = strategy == LLAMA_PSHARD_GPUONLY_LAYERPIN_LAYERSTREAM; + const int32_t shard_bid = use_alternating_shards ? layout.shard(il) : layout.shard_a; + + switch (strategy) { + case LLAMA_PSHARD_GPUONLY_LAYERPIN_LAYERSTREAM: + emit(patterns_layer[il].c_str(), host_buft, shard_bid); + break; + + case LLAMA_PSHARD_GPUONLY_ATTNPIN_FFNSTREAM: + emit(patterns_layer_ffn[il].c_str(), host_buft, shard_bid); + emit(patterns_layer[il].c_str(), gpu_buft, layout.compute); + break; + + case LLAMA_PSHARD_DYNAMIC_FFNCPU_ATTNSTREAM: + emit(patterns_layer_ffn[il].c_str(), host_buft, layout.cpu); + emit(patterns_layer[il].c_str(), host_buft, shard_bid); + break; + + case LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS: + if (n_attn_pinned > 0 && il < n_attn_pinned) { + emit(patterns_layer_ffn[il].c_str(), host_buft, layout.cpu); + emit(patterns_layer[il].c_str(), gpu_buft, layout.compute); + } else { + emit(patterns_layer[il].c_str(), host_buft, layout.cpu); + } + break; + + default: break; + } + } + } + tensor_buft_overrides[itbo] = { nullptr, nullptr, -1 }; + LLAMA_LOG_DEBUG("%s: %zu overrides emitted\n", __func__, itbo); +} + +// overflow names by enum value +static const char * const PSHARD_FRAC_NAMES[] = { "NONE", "ATTN", "UP", "GATE", "MOE" }; + +struct llama_pshard_search_ctx { + const char * path_model; + const struct llama_model_params * mparams; + const struct llama_context_params * cparams; + struct llama_model_tensor_buft_override * overrides; + uint32_t n_layers; + size_t vram_free; + ggml_backend_buffer_type_t gpu_buft; + ggml_backend_buffer_type_t host_buft; + pshard_dev_layout layout; + bool is_moe; + bool has_rs = false; + + // optional TPS predictor + const llama_benchmark_predictor * predictor = nullptr; + uint32_t kv_size = 0; + uint32_t cache_ubatch = 0; +}; + +struct llama_pshard_tps_hook_data { + const llama_benchmark_predictor * predictor; + int cpu_backend_id; + uint32_t kv_size; + int32_t batch_size; + uint32_t n_outputs; + bool has_rs; + float * out_tps; +}; + +static void pshard_tps_probe_hook(llama_context * ctx, void * user_data) { + auto * d = (llama_pshard_tps_hook_data *) user_data; + if (!d || !d->predictor || !ctx) return; + + double tps = d->predictor->predict_tps(ctx->get_sched(), d->cpu_backend_id, d->kv_size, d->batch_size, d->n_outputs, d->has_rs); + if (d->out_tps) { + *d->out_tps = (float)tps; + } +} + +static std::vector llama_pshard_probe_memory( + const llama_pshard_search_ctx & ctx, + const llama_model_params & mparams, + const llama_context_params & cparams, + ggml_log_level log_level, + llama_probe_hook_t probe_hook = nullptr, + void * probe_hook_data = nullptr) { + std::vector devs; + uint32_t hp_ngl = 0, hp_n_ctx_train = 0, hp_n_expert = 0, hp_n_embd_r = 0; + + const uint32_t probe_n_tokens = std::max(1, cparams.n_batch ? cparams.n_batch : cparams.n_ubatch); + const uint32_t probe_n_outputs = probe_n_tokens; + + llama_context_params cparams_probe = cparams; + if (ctx.cache_ubatch != 0) { + cparams_probe.n_batch = std::max(cparams_probe.n_batch, ctx.cache_ubatch); + cparams_probe.n_ubatch = ctx.cache_ubatch; + } + + return llama_get_device_memory_data_safe( + ctx.path_model, &mparams, &cparams_probe, devs, + hp_ngl, hp_n_ctx_train, hp_n_expert, hp_n_embd_r, + log_level, probe_hook, probe_hook_data, + probe_n_tokens, probe_n_outputs); +} + +struct llama_pshard_tier_prune { + uint32_t hi_pinned[LLAMA_PSHARD_COUNT]; + uint32_t hi_attn; + bool skip[LLAMA_PSHARD_COUNT]; + + void init(uint32_t n_layers) { + for (int s = 0; s < LLAMA_PSHARD_COUNT; s++) { + hi_pinned[s] = n_layers; + skip[s] = false; + } + hi_attn = n_layers; + } + + void update(int s, const llama_pshard_plan & plan) { + if (!plan.is_viable) { + return; + } + if (s == LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS) { + if (plan.n_attn_pinned > 0) { + hi_attn = plan.n_attn_pinned; + hi_pinned[s] = plan.n_pinned; + } + } else { + if (plan.n_pinned > 0) { + hi_pinned[s] = plan.n_pinned; + } + } + } +}; + +static llama_pshard_plan llama_pshard_search_strategy( + const llama_pshard_search_ctx & ctx, + llama_pshard_strategy strategy, + uint32_t hi_hint = UINT32_MAX, + uint32_t lo_hint = 0) { + + const auto * mparams = ctx.mparams; + const auto * cparams = ctx.cparams; + auto * tensor_buft_overrides = ctx.overrides; + const auto n_layers = ctx.n_layers; + const auto vram_free = ctx.vram_free; + const auto gpu_buft = ctx.gpu_buft; + const auto host_buft = ctx.host_buft; + const auto & layout = ctx.layout; + const auto is_moe = ctx.is_moe; + + llama_pshard_plan plan; + plan.strategy = strategy; + plan.batch_size = cparams->n_batch; + + const bool delegate_compute = llama_pshard_strategy_delegates_compute(strategy); + + const uint32_t hi_default = n_layers - 1; + uint32_t lo = lo_hint, hi = (hi_hint < hi_default) ? hi_hint : hi_default; + uint32_t best_n_pinned = lo_hint; + int64_t mem_lo = 0, mem_hi = (int64_t)vram_free * 2; + + // try the upper bound first + { + llama_pshard_generate_overrides(hi, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, LLAMA_LAYER_FRACTION_NONE, strategy, layout, false, false); + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = delegate_compute; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + try { + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR); + const int64_t gpu_used = d[0].mb.total(); + LLAMA_LOG_INFO("%s: [%s] n_pinned=%u -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s (hi-first)\n", + __func__, llama_pshard_strategy_name(strategy), hi, + gpu_used / (1024.0 * 1024.0), + d[0].mb.model / (1024.0 * 1024.0), d[0].mb.context / (1024.0 * 1024.0), d[0].mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + if (gpu_used <= (int64_t)vram_free) { + best_n_pinned = hi; + lo = hi + 1; + } else { + mem_hi = gpu_used; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [%s] hi-first probe failed (n_pinned=%u)\n", __func__, llama_pshard_strategy_name(strategy), hi); + } + } + + while (lo <= hi) { + uint32_t mid; + if (mem_hi > mem_lo && mem_hi > (int64_t)vram_free) { + int64_t target = (int64_t)vram_free; + mid = lo + (uint32_t)((double)(target - mem_lo) * (hi - lo) / (mem_hi - mem_lo)); + if (mid <= lo) mid = lo + 1; + if (mid > hi) mid = hi; + } else { + mid = (lo + hi) / 2; + } + + llama_pshard_generate_overrides(mid, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, LLAMA_LAYER_FRACTION_NONE, strategy, layout, false, false); + + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = delegate_compute; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + + try { + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR); + const int64_t gpu_used = d[0].mb.total(); + + LLAMA_LOG_INFO("%s: [%s] n_pinned=%u -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s\n", + __func__, llama_pshard_strategy_name(strategy), mid, + gpu_used / (1024.0 * 1024.0), + d[0].mb.model / (1024.0 * 1024.0), + d[0].mb.context / (1024.0 * 1024.0), + d[0].mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + + if (gpu_used <= (int64_t)vram_free) { + best_n_pinned = mid; + lo = mid + 1; + mem_lo = gpu_used; + } else { + if (mid == 0) break; + hi = mid - 1; + mem_hi = gpu_used; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [%s] probe failed (n_pinned=%u)\n", __func__, llama_pshard_strategy_name(strategy), mid); + if (mid == 0) break; + hi = mid - 1; + } + } + + llama_layer_fraction best_overflow = LLAMA_LAYER_FRACTION_NONE; + const uint32_t fallback_n_pinned = best_n_pinned; // known fit + const llama_layer_fraction fallback_overflow = LLAMA_LAYER_FRACTION_NONE; + if (best_n_pinned < n_layers - 1) { + const uint32_t frac_n_pinned = best_n_pinned + 1; // pin one more layer, partially + auto try_frac = [&](llama_layer_fraction frac) -> bool { + llama_pshard_generate_overrides(frac_n_pinned, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, frac, strategy, layout, false, false); + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = delegate_compute; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + try { + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR); + return d[0].mb.total() <= (int64_t)vram_free; + } catch (...) { + LLAMA_LOG_WARN("%s: [%s] overflow probe failed (frac=%d)\n", __func__, llama_pshard_strategy_name(strategy), (int)frac); + return false; + } + }; + // try one partial boundary layer + if (try_frac(LLAMA_LAYER_FRACTION_ATTN)) { + best_n_pinned = frac_n_pinned; + best_overflow = LLAMA_LAYER_FRACTION_ATTN; + if (try_frac(LLAMA_LAYER_FRACTION_UP)) { best_overflow = LLAMA_LAYER_FRACTION_UP; + if (try_frac(LLAMA_LAYER_FRACTION_GATE)) { best_overflow = LLAMA_LAYER_FRACTION_GATE; + if (try_frac(LLAMA_LAYER_FRACTION_MOE)) { best_overflow = LLAMA_LAYER_FRACTION_MOE; }}} + } + } + + plan.n_pinned = best_n_pinned; + plan.overflow = best_overflow; + plan.output_on_gpu = false; + + llama_pshard_generate_overrides(best_n_pinned, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, best_overflow, strategy, layout, false, plan.output_on_gpu); + { + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = delegate_compute; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + llama_pshard_tps_hook_data tps_data = { ctx.predictor, layout.cpu, ctx.kv_size, (int32_t)cparams->n_batch, cparams->n_seq_max, ctx.has_rs, &plan.tps }; + auto * hook = ctx.predictor ? pshard_tps_probe_hook : nullptr; + auto * hookdata = ctx.predictor ? (void *)&tps_data : nullptr; + + try { + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR, hook, hookdata); + plan.total_vram_req = d[0].mb.total(); + plan.scratch_measured = d[0].mb.compute; + plan.cache_measured = d[0].mb.context; + plan.is_viable = ((int64_t)plan.total_vram_req <= (int64_t)vram_free); + } catch (...) { + LLAMA_LOG_WARN("%s: [%s] final measurement probe failed (n_pinned=%u)\n", __func__, llama_pshard_strategy_name(strategy), best_n_pinned); + plan.is_viable = false; + } + } + + // drop the partial boundary layer if the final probe exceeds budget + if (!plan.is_viable && best_overflow != LLAMA_LAYER_FRACTION_NONE) { + best_n_pinned = fallback_n_pinned; + best_overflow = fallback_overflow; + plan.n_pinned = best_n_pinned; + plan.overflow = best_overflow; + llama_pshard_generate_overrides(best_n_pinned, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, best_overflow, strategy, layout, false, plan.output_on_gpu); + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = delegate_compute; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + llama_pshard_tps_hook_data tps_data = { ctx.predictor, layout.cpu, ctx.kv_size, (int32_t)cparams->n_batch, cparams->n_seq_max, ctx.has_rs, &plan.tps }; + auto * hook = ctx.predictor ? pshard_tps_probe_hook : nullptr; + auto * hookdata = ctx.predictor ? (void *)&tps_data : nullptr; + try { + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR, hook, hookdata); + plan.total_vram_req = d[0].mb.total(); + plan.scratch_measured = d[0].mb.compute; + plan.cache_measured = d[0].mb.context; + plan.is_viable = ((int64_t)plan.total_vram_req <= (int64_t)vram_free); + } catch (...) { + plan.is_viable = false; + } + } + + for (const auto * ov = tensor_buft_overrides; ov->pattern; ++ov) { + plan.overrides.push_back({ov->pattern, ov->buft, ov->backend_id}); + } + + return plan; +} + +static llama_pshard_plan llama_pshard_search_attn_pin( + const llama_pshard_search_ctx & ctx, + uint32_t hi_attn_hint = UINT32_MAX, + uint32_t hi_full_hint = UINT32_MAX, + uint32_t lo_full_hint = 0) { + + const auto * mparams = ctx.mparams; + const auto * cparams = ctx.cparams; + auto * tensor_buft_overrides = ctx.overrides; + const auto n_layers = ctx.n_layers; + const auto vram_free = ctx.vram_free; + const auto gpu_buft = ctx.gpu_buft; + const auto host_buft = ctx.host_buft; + const auto & layout = ctx.layout; + const auto is_moe = ctx.is_moe; + + llama_pshard_plan plan; + plan.strategy = LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS; + + auto measure_vram = [&](uint32_t n_full, uint32_t n_attn, bool out_gpu) -> llama_memory_breakdown_data { + llama_pshard_generate_overrides(n_full, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, LLAMA_LAYER_FRACTION_NONE, LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS, + layout, false, out_gpu, n_attn); + + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = true; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR); + return d[0].mb; + }; + + // phase 1: maximize attention layers on GPU + uint32_t n_attn = 0; + { + uint32_t lo = 0, hi = (hi_attn_hint < n_layers) ? hi_attn_hint : n_layers; + int64_t mem_lo = 0, mem_hi = (int64_t)vram_free * 2; + + try { + auto mb = measure_vram(0, hi, false); + int64_t gpu_used = mb.total(); + LLAMA_LOG_INFO("%s: [STATIC_ATTNPRIO_ALLMODELS p1] n_attn=%u -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s (hi-first)\n", + __func__, hi, gpu_used / (1024.0 * 1024.0), + mb.model / (1024.0 * 1024.0), mb.context / (1024.0 * 1024.0), mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + if (gpu_used <= (int64_t)vram_free) { + n_attn = hi; + lo = hi + 1; + } else { + mem_hi = gpu_used; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS p1] hi-first probe failed (n_attn=%u)\n", __func__, hi); + } + + while (lo <= hi) { + uint32_t mid; + if (mem_hi > mem_lo && mem_hi > (int64_t)vram_free) { + mid = lo + (uint32_t)((double)((int64_t)vram_free - mem_lo) * (hi - lo) / (mem_hi - mem_lo)); + if (mid <= lo) mid = lo + 1; + if (mid > hi) mid = hi; + } else { + mid = (lo + hi) / 2; + } + + try { + auto mb = measure_vram(0, mid, false); + int64_t gpu_used = mb.total(); + LLAMA_LOG_INFO("%s: [STATIC_ATTNPRIO_ALLMODELS p1] n_attn=%u -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s\n", + __func__, mid, gpu_used / (1024.0 * 1024.0), + mb.model / (1024.0 * 1024.0), mb.context / (1024.0 * 1024.0), mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + + if (gpu_used <= (int64_t)vram_free) { + n_attn = mid; + lo = mid + 1; + mem_lo = gpu_used; + } else { + if (mid == 0) break; + hi = mid - 1; + mem_hi = gpu_used; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS p1] probe failed (n_attn=%u)\n", __func__, mid); + if (mid == 0) break; + hi = mid - 1; + } + } + } + + // moe tries output on gpu before ffn pinning + // dense tries output on gpu after ffn pinning + bool output_on_gpu = false; + + if (is_moe && n_attn >= n_layers) { + try { + auto mb = measure_vram(0, n_attn, true); + int64_t gpu_used = mb.total(); + LLAMA_LOG_INFO("%s: [STATIC_ATTNPRIO_ALLMODELS p1b] output_on_gpu probe -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s\n", + __func__, gpu_used / (1024.0 * 1024.0), + mb.model / (1024.0 * 1024.0), mb.context / (1024.0 * 1024.0), mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + if (gpu_used <= (int64_t)vram_free) { + output_on_gpu = true; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS p1b] output_on_gpu probe failed\n", __func__); + } + } + + // phase 2: maximize fully pinned layers + uint32_t n_full = 0; + { + uint32_t hi_full_max = (hi_full_hint < n_attn) ? hi_full_hint : n_attn; + uint32_t lo = lo_full_hint, hi = hi_full_max; + int64_t mem_lo = 0, mem_hi = (int64_t)vram_free * 2; + + try { + auto mb = measure_vram(hi, n_attn, output_on_gpu); + int64_t gpu_used = mb.total(); + LLAMA_LOG_INFO("%s: [STATIC_ATTNPRIO_ALLMODELS p2] n_full=%u n_attn=%u -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s (hi-first)\n", + __func__, hi, n_attn, gpu_used / (1024.0 * 1024.0), + mb.model / (1024.0 * 1024.0), mb.context / (1024.0 * 1024.0), mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + if (gpu_used <= (int64_t)vram_free) { + n_full = hi; + lo = hi + 1; + } else { + mem_hi = gpu_used; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS p2] hi-first probe failed (n_full=%u, n_attn=%u)\n", __func__, hi, n_attn); + } + + while (lo <= hi) { + uint32_t mid; + if (mem_hi > mem_lo && mem_hi > (int64_t)vram_free) { + mid = lo + (uint32_t)((double)((int64_t)vram_free - mem_lo) * (hi - lo) / (mem_hi - mem_lo)); + if (mid <= lo) mid = lo + 1; + if (mid > hi) mid = hi; + } else { + mid = (lo + hi) / 2; + } + + try { + auto mb = measure_vram(mid, n_attn, output_on_gpu); + int64_t gpu_used = mb.total(); + LLAMA_LOG_INFO("%s: [STATIC_ATTNPRIO_ALLMODELS p2] n_full=%u n_attn=%u -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s\n", + __func__, mid, n_attn, gpu_used / (1024.0 * 1024.0), + mb.model / (1024.0 * 1024.0), mb.context / (1024.0 * 1024.0), mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + + if (gpu_used <= (int64_t)vram_free) { + n_full = mid; + lo = mid + 1; + mem_lo = gpu_used; + } else { + if (mid == 0) break; + hi = mid - 1; + mem_hi = gpu_used; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS p2] probe failed (n_full=%u, n_attn=%u)\n", __func__, mid, n_attn); + if (mid == 0) break; + hi = mid - 1; + } + } + } + + // dense tries output on gpu only after all layers fit + if (!is_moe && !output_on_gpu && n_full >= n_layers) { + try { + auto mb = measure_vram(n_full, n_attn, true); + int64_t gpu_used = mb.total(); + LLAMA_LOG_INFO("%s: [STATIC_ATTNPRIO_ALLMODELS p3] output_on_gpu probe (n_full=%u) -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s\n", + __func__, n_full, gpu_used / (1024.0 * 1024.0), + mb.model / (1024.0 * 1024.0), mb.context / (1024.0 * 1024.0), mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + if (gpu_used <= (int64_t)vram_free) { + output_on_gpu = true; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS p3] output_on_gpu probe failed\n", __func__); + } + } + + plan.n_pinned = n_full; + plan.n_attn_pinned = n_attn; + plan.output_on_gpu = output_on_gpu; + + llama_pshard_generate_overrides(n_full, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, LLAMA_LAYER_FRACTION_NONE, LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS, + layout, false, plan.output_on_gpu, n_attn); + { + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = true; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + llama_pshard_tps_hook_data tps_data = { ctx.predictor, layout.cpu, ctx.kv_size, (int32_t)cparams->n_batch, cparams->n_seq_max, ctx.has_rs, &plan.tps }; + auto * hook = ctx.predictor ? pshard_tps_probe_hook : nullptr; + auto * hookdata = ctx.predictor ? (void *)&tps_data : nullptr; + + try { + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR, hook, hookdata); + plan.total_vram_req = d[0].mb.total(); + plan.scratch_measured = d[0].mb.compute; + plan.cache_measured = d[0].mb.context; + plan.is_viable = ((int64_t)plan.total_vram_req <= (int64_t)vram_free); + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS] final measurement probe failed (n_full=%u, n_attn=%u)\n", __func__, n_full, n_attn); + plan.is_viable = false; + } + } + + for (const auto * ov = tensor_buft_overrides; ov->pattern; ++ov) { + plan.overrides.push_back({ov->pattern, ov->buft, ov->backend_id}); + } + + return plan; +} + +llama_pshard_plan_registry * llama_pshard_registry_create(uint32_t n_tier_max, uint32_t n_seq_max) { + auto * registry = new llama_pshard_plan_registry(); + registry->init(n_tier_max, n_seq_max); + return registry; +} + +void llama_pshard_registry_free(llama_pshard_plan_registry * registry) { + delete registry; +} + +// plan cache serialization + +uint64_t pshard_registry_fingerprint( + const struct llama_model_params * mparams, + const struct llama_context_params * cparams, + int64_t model_file_size) { + + (void) mparams; + + uint64_t h = 0xcbf29ce484222325ULL; + auto mix = [&](uint64_t v) { h ^= v; h *= 0x100000001b3ULL; }; + + mix(cparams->n_ctx); + mix(cparams->n_seq_max); + mix(cparams->n_threads); + mix((uint64_t)cparams->flash_attn_type); + mix((uint64_t)cparams->type_k); + mix((uint64_t)cparams->type_v); + mix((uint64_t)model_file_size); + mix((uint64_t)pshard_strategy_from_env()); + + return h; +} + +static std::string pshard_plan_to_ot(const llama_pshard_plan & plan, ggml_backend_buffer_type_t host_buft) { + std::string ot; + const char * buft_name = ggml_backend_buft_name(host_buft); + for (size_t i = 0; i < plan.overrides.size(); i++) { + if (i > 0) ot += ','; + ot += plan.overrides[i].pattern; + ot += '='; + ot += buft_name; + ot += ':'; + ot += std::to_string(plan.overrides[i].backend_id); + } + return ot; +} + +static const char * pshard_overflow_name(int overflow) { + static const char * const names[] = { "NONE", "ATTN", "UP", "GATE", "MOE" }; + return (overflow >= 0 && overflow < 5) ? names[overflow] : "NONE"; +} + +static int pshard_overflow_from_name(const char * name) { + if (!name) return 0; + if (strcmp(name, "ATTN") == 0) return 1; + if (strcmp(name, "UP") == 0) return 2; + if (strcmp(name, "GATE") == 0) return 3; + if (strcmp(name, "MOE") == 0) return 4; + return 0; +} + +static const size_t PSHARD_MIB = 1024ULL * 1024ULL; +static const int PSHARD_CACHE_MAX_SECTIONS = 32; +static const int PSHARD_CACHE_MAX_VARIANTS = 16; + +static uint32_t pshard_bytes_to_mib_ceil(size_t bytes) { + return (uint32_t)((bytes + PSHARD_MIB - 1) / PSHARD_MIB); +} + +static size_t pshard_mib_to_bytes(uint32_t mib) { + return (size_t)mib * PSHARD_MIB; +} + +static size_t pshard_mib_to_bytes(double mib) { + return (size_t)(mib * (double)PSHARD_MIB + 0.5); +} + +static bool pshard_parse_variant_header(const std::string & line, uint32_t & budget_mib, uint32_t & cache_ubatch) { + cache_ubatch = 0; + if (sscanf(line.c_str(), "[variant budget=%u cache_ubatch=%u]", &budget_mib, &cache_ubatch) == 2) { + return true; + } + return sscanf(line.c_str(), "[variant budget=%u]", &budget_mib) == 1; +} + +static bool pshard_plan_is_better(const llama_pshard_plan & candidate, const llama_pshard_plan & current); + +bool pshard_registry_save( + const llama_pshard_plan_registry * registry, uint64_t fingerprint, + const char * cache_path, ggml_backend_buffer_type_t host_buft, + const llama_context_params * cparams) { + if (!registry || !cache_path) return false; + + struct cache_section { + std::string header; + std::vector lines; + }; + std::vector sections; + + FILE * existing = fopen(cache_path, "r"); + if (existing) { + char line[8192]; + cache_section * cur = nullptr; + while (fgets(line, sizeof(line), existing)) { + std::string s = line; + while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) s.pop_back(); + if (s.compare(0, 13, "[fingerprint=") == 0) { + sections.push_back({s, {}}); + cur = §ions.back(); + } else if (cur) { + cur->lines.push_back(s); + } + } + fclose(existing); + } + + std::vector preserved_sections; + std::vector> preserved_variants; + + size_t inferred_budget = registry->pshard_disabled ? registry->baseline_vram_req : 0; + if (!registry->pshard_disabled) { + for (const auto & plan : registry->best_plans) { + if (plan.is_viable) { + inferred_budget = std::max(inferred_budget, plan.total_vram_req); + } + } + } + const uint32_t budget_mib = registry->budget_mib + ? registry->budget_mib + : pshard_bytes_to_mib_ceil(inferred_budget); + const uint32_t cache_ubatch = registry->cache_ubatch + ? registry->cache_ubatch + : (registry->tier_sizes.empty() ? 0 : registry->tier_sizes.back()); + + for (const auto & sec : sections) { + uint64_t fp = 0; + sscanf(sec.header.c_str(), "[fingerprint=0x%" SCNx64, &fp); + if (fp != fingerprint) { + preserved_sections.push_back(sec); + continue; + } + + std::vector cur_variant; + uint32_t cur_budget = 0; + uint32_t cur_cache_ubatch = 0; + bool in_variant = false; + auto flush_variant = [&]() { + if (in_variant && (cur_budget != budget_mib || cur_cache_ubatch != cache_ubatch)) { + preserved_variants.push_back(cur_variant); + } + cur_variant.clear(); + cur_budget = 0; + cur_cache_ubatch = 0; + in_variant = false; + }; + + for (const auto & ln : sec.lines) { + uint32_t parsed_budget = 0; + uint32_t parsed_cache_ubatch = 0; + if (pshard_parse_variant_header(ln, parsed_budget, parsed_cache_ubatch)) { + flush_variant(); + in_variant = true; + cur_budget = parsed_budget; + cur_cache_ubatch = parsed_cache_ubatch; + cur_variant.push_back(ln); + } else if (in_variant) { + cur_variant.push_back(ln); + } + } + flush_variant(); + } + + while ((int)preserved_sections.size() >= PSHARD_CACHE_MAX_SECTIONS) { + preserved_sections.erase(preserved_sections.begin()); + } + while ((int)preserved_variants.size() >= PSHARD_CACHE_MAX_VARIANTS) { + preserved_variants.erase(preserved_variants.begin()); + } + + FILE * f = fopen(cache_path, "w"); + if (!f) { + LLAMA_LOG_WARN("%s: could not write plan cache: %s\n", __func__, cache_path); + return false; + } + + fprintf(f, "# Generated file. Edit at your own risk.\n"); + + for (const auto & sec : preserved_sections) { + fprintf(f, "\n%s\n", sec.header.c_str()); + for (const auto & ln : sec.lines) { + fprintf(f, "%s\n", ln.c_str()); + } + } + + fprintf(f, "\n[fingerprint=0x%016" PRIx64 "]\n", fingerprint); + if (cparams) { + const char * fa_str = "unknown"; + switch (cparams->flash_attn_type) { + case LLAMA_FLASH_ATTN_TYPE_DISABLED: fa_str = "off"; break; + case LLAMA_FLASH_ATTN_TYPE_ENABLED: fa_str = "on"; break; + case LLAMA_FLASH_ATTN_TYPE_AUTO: fa_str = "auto"; break; + } + const int forced_strategy = pshard_strategy_from_env(); + fprintf(f, "# n_ctx=%u n_seq_max=%u n_threads=%d fa=%s type_k=%d type_v=%d strategy=%s\n", + cparams->n_ctx, cparams->n_seq_max, cparams->n_threads, + fa_str, (int)cparams->type_k, (int)cparams->type_v, + forced_strategy >= 0 ? llama_pshard_strategy_name((llama_pshard_strategy)forced_strategy) : "auto"); + } + + for (const auto & variant : preserved_variants) { + fprintf(f, "\n"); + for (const auto & ln : variant) { + fprintf(f, "%s\n", ln.c_str()); + } + } + + fprintf(f, "\n[variant budget=%u cache_ubatch=%u]\n", budget_mib, cache_ubatch); + if (registry->pshard_disabled) { + fprintf(f, "pshard_disabled=1 baseline_vram=%.1f\n", registry->baseline_vram_req / (1024.0 * 1024.0)); + } else { + for (size_t t = 0; t < registry->tier_sizes.size(); t++) { + const auto & plan = registry->best_plans[t]; + if (!plan.is_viable) { + fprintf(f, "[tier %zu bs=%u] not_viable\n", t, registry->tier_sizes[t]); + continue; + } + fprintf(f, "[tier %zu bs=%u]\n", t, registry->tier_sizes[t]); + fprintf(f, "strategy=%s n_pinned=%u n_attn_pinned=%u overflow=%s tps=%.2f vram=%.1f output_on_gpu=%d pin_from_back=%d\n", + llama_pshard_strategy_name(plan.strategy), + plan.n_pinned, plan.n_attn_pinned, + pshard_overflow_name(plan.overflow), + plan.tps, plan.total_vram_req / (1024.0 * 1024.0), + (int)plan.output_on_gpu, (int)plan.pin_from_back); + fprintf(f, "ot=%s\n", pshard_plan_to_ot(plan, host_buft).c_str()); + } + } + + fclose(f); + LLAMA_LOG_INFO("%s: saved budget=%u MiB cache_ubatch=%u variant with %zu tier plans to %s\n", + __func__, budget_mib, cache_ubatch, registry->tier_sizes.size(), cache_path); + return true; +} + +bool pshard_registry_load( + llama_pshard_plan_registry * registry, uint64_t fingerprint, + const char * cache_path, ggml_backend_buffer_type_t host_buft, + size_t current_budget, bool require_exact_budget) { + if (!registry || !cache_path) return false; + + FILE * f = fopen(cache_path, "r"); + if (!f) return false; + + char line[8192]; + bool in_section = false; + + char fp_header[64]; + snprintf(fp_header, sizeof(fp_header), "[fingerprint=0x%016" PRIx64 "]", fingerprint); + + struct tier_data { + uint32_t bs = 0; + bool viable = false; + llama_pshard_strategy strategy = LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS; + uint32_t n_pinned = 0; + uint32_t n_attn_pinned = 0; + int overflow = 0; + float tps = 0.0f; + double vram_mib = 0.0; + int output_on_gpu = 0; + int pin_from_back = 0; + std::string ot_line; + }; + struct variant_data { + uint32_t budget_mib = 0; + uint32_t cache_ubatch = 0; + bool pshard_disabled = false; + double baseline_vram_mib = 0.0; + std::vector tiers; + }; + std::vector variants; + variant_data * cur_variant = nullptr; + + while (fgets(line, sizeof(line), f)) { + std::string s = line; + while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) s.pop_back(); + + if (s.compare(0, 13, "[fingerprint=") == 0) { + if (in_section) break; // hit next section, stop + in_section = (s == fp_header); + continue; + } + if (!in_section) continue; + + uint32_t variant_budget = 0; + uint32_t variant_cache_ubatch = 0; + if (pshard_parse_variant_header(s, variant_budget, variant_cache_ubatch)) { + variants.push_back({}); + cur_variant = &variants.back(); + cur_variant->budget_mib = variant_budget; + cur_variant->cache_ubatch = variant_cache_ubatch; + continue; + } + if (!cur_variant) continue; + + if (s.rfind("pshard_disabled=1", 0) == 0) { + double baseline_mib = 0.0; + if (sscanf(s.c_str(), "pshard_disabled=1 baseline_vram=%lf", &baseline_mib) == 1) { + cur_variant->pshard_disabled = true; + cur_variant->baseline_vram_mib = baseline_mib; + } else { + LLAMA_LOG_WARN("%s: malformed pshard_disabled line: %s\n", __func__, s.c_str()); + } + continue; + } + + if (s.compare(0, 5, "[tier") == 0) { + tier_data td = {}; + size_t tier_idx = 0; + if (s.find("not_viable") != std::string::npos) { + if (sscanf(s.c_str(), "[tier %zu bs=%u]", &tier_idx, &td.bs) < 2) { + LLAMA_LOG_WARN("%s: malformed tier header (not_viable): %s\n", __func__, s.c_str()); + continue; + } + td.viable = false; + } else { + if (sscanf(s.c_str(), "[tier %zu bs=%u]", &tier_idx, &td.bs) < 2) { + LLAMA_LOG_WARN("%s: malformed tier header: %s\n", __func__, s.c_str()); + continue; + } + td.viable = true; + } + cur_variant->tiers.push_back(td); + } else if (s.compare(0, 9, "strategy=") == 0 && !cur_variant->tiers.empty()) { + auto & td = cur_variant->tiers.back(); + char strat_name[64] = {}, overflow_name[16] = {}; + if (sscanf(s.c_str(), "strategy=%63s n_pinned=%u n_attn_pinned=%u overflow=%15s tps=%f vram=%lf", + strat_name, &td.n_pinned, &td.n_attn_pinned, overflow_name, &td.tps, &td.vram_mib) < 4) { + LLAMA_LOG_WARN("%s: malformed strategy line: %s\n", __func__, s.c_str()); + td.viable = false; + continue; + } + + const char * ogg = strstr(s.c_str(), "output_on_gpu="); + const char * pfb = strstr(s.c_str(), "pin_from_back="); + if (!ogg || !pfb) { + LLAMA_LOG_WARN("%s: missing output_on_gpu/pin_from_back, invalidating cache: %s\n", __func__, s.c_str()); + td.viable = false; + continue; + } + td.output_on_gpu = atoi(ogg + 14); + td.pin_from_back = atoi(pfb + 14); + + td.overflow = pshard_overflow_from_name(overflow_name); + bool found_strategy = false; + for (int i = 0; i < LLAMA_PSHARD_COUNT; i++) { + if (strcmp(strat_name, llama_pshard_strategy_name((llama_pshard_strategy)i)) == 0) { + td.strategy = (llama_pshard_strategy)i; + found_strategy = true; + break; + } + } + if (!found_strategy) { + LLAMA_LOG_WARN("%s: unknown strategy in cache, invalidating tier: %s\n", __func__, strat_name); + td.viable = false; + } + } else if (s.compare(0, 3, "ot=") == 0 && !cur_variant->tiers.empty()) { + cur_variant->tiers.back().ot_line = s.substr(3); + } + } + fclose(f); + + if (variants.empty() || !in_section) return false; + + auto make_plan = [&](const tier_data & td) { + llama_pshard_plan plan; + plan.strategy = td.strategy; + plan.batch_size = td.bs; + plan.n_pinned = td.n_pinned; + plan.n_attn_pinned = td.n_attn_pinned; + plan.overflow = td.overflow; + plan.tps = td.tps; + plan.total_vram_req = (size_t)(td.vram_mib * 1024 * 1024); + plan.is_viable = td.viable; + plan.output_on_gpu = (bool)td.output_on_gpu; + plan.pin_from_back = (bool)td.pin_from_back; + + if (!td.ot_line.empty()) { + std::string remaining = td.ot_line; + while (!remaining.empty()) { + size_t comma = remaining.find(','); + std::string token = (comma != std::string::npos) ? remaining.substr(0, comma) : remaining; + remaining = (comma != std::string::npos) ? remaining.substr(comma + 1) : ""; + + size_t eq = token.find('='); + if (eq == std::string::npos) continue; + + std::string pattern = token.substr(0, eq); + std::string buft_bid = token.substr(eq + 1); + + int32_t backend_id = -1; + size_t colon = buft_bid.rfind(':'); + if (colon != std::string::npos) { + backend_id = atoi(buft_bid.c_str() + colon + 1); + } + + plan.overrides.push_back({ pattern, host_buft, backend_id }); + } + } + return plan; + }; + + const uint32_t current_budget_mib = pshard_bytes_to_mib_ceil(current_budget); + const uint32_t requested_cache_ubatch = registry->cache_ubatch; + bool skipped_cache_ubatch = false; + auto variant_cache_ubatch = [&](const variant_data & variant) { + return variant.cache_ubatch ? variant.cache_ubatch : requested_cache_ubatch; + }; + auto cache_ubatch_ok = [&](const variant_data & variant) { + if (requested_cache_ubatch == 0) return true; + if (variant.cache_ubatch == 0) return true; + if (variant.cache_ubatch <= requested_cache_ubatch) return true; + skipped_cache_ubatch = true; + return false; + }; + + for (const auto & variant : variants) { + if (!variant.pshard_disabled) continue; + if (!cache_ubatch_ok(variant)) continue; + const size_t baseline_vram = pshard_mib_to_bytes(variant.baseline_vram_mib); + if (baseline_vram <= current_budget) { + registry->tier_sizes.clear(); + registry->best_plans.clear(); + registry->pshard_disabled = true; + registry->baseline_vram_req = baseline_vram; + registry->budget_mib = variant.budget_mib; + registry->cache_ubatch = variant_cache_ubatch(variant); + LLAMA_LOG_INFO("%s: loaded pshard_disabled variant budget=%u MiB cache_ubatch=%u baseline=%.1f MiB from %s\n", + __func__, variant.budget_mib, registry->cache_ubatch, variant.baseline_vram_mib, cache_path); + return true; + } + } + + const variant_data * best_whole = nullptr; + for (const auto & variant : variants) { + if (variant.pshard_disabled || variant.tiers.empty()) continue; + if (!cache_ubatch_ok(variant)) continue; + if (require_exact_budget) { + if (variant.budget_mib != current_budget_mib) continue; + } else if (pshard_mib_to_bytes(variant.budget_mib) > current_budget) { + continue; + } + if (!best_whole || + variant_cache_ubatch(variant) > variant_cache_ubatch(*best_whole) || + (variant_cache_ubatch(variant) == variant_cache_ubatch(*best_whole) && + variant.budget_mib > best_whole->budget_mib)) { + best_whole = &variant; + } + } + + std::vector> selected; + auto add_or_fill_plan = [&](const tier_data & td, bool allow_existing) { + llama_pshard_plan plan = make_plan(td); + auto it = std::find_if(selected.begin(), selected.end(), + [&](const auto & p) { return p.first == td.bs; }); + if (it == selected.end()) { + selected.push_back({td.bs, std::move(plan)}); + } else if (allow_existing || !it->second.is_viable || + (!best_whole && plan.is_viable && pshard_plan_is_better(plan, it->second))) { + it->second = std::move(plan); + } + }; + + if (best_whole) { + for (const auto & td : best_whole->tiers) { + add_or_fill_plan(td, true); + } + } + + if (!require_exact_budget) { + for (const auto & variant : variants) { + if (variant.pshard_disabled || pshard_mib_to_bytes(variant.budget_mib) <= current_budget) continue; + if (!cache_ubatch_ok(variant)) continue; + for (const auto & td : variant.tiers) { + if (!td.viable || td.ot_line.empty()) continue; + if (pshard_mib_to_bytes(td.vram_mib) > current_budget) continue; + add_or_fill_plan(td, false); + } + } + } + + const uint32_t selected_cache_ubatch = best_whole ? variant_cache_ubatch(*best_whole) : requested_cache_ubatch; + selected.erase(std::remove_if(selected.begin(), selected.end(), + [&](const auto & p) { + if (p.first == 0) return true; + return selected_cache_ubatch > 0 && p.first > selected_cache_ubatch; + }), selected.end()); + if (selected.empty()) { + if (require_exact_budget && !variants.empty()) { + LLAMA_LOG_INFO("%s: cache miss, no exact budget=%u MiB variant in %s\n", + __func__, current_budget_mib, cache_path); + } + if (skipped_cache_ubatch) { + LLAMA_LOG_INFO("%s: cache miss, no variant with cache_ubatch <= target cache_ubatch=%u in %s\n", + __func__, requested_cache_ubatch, cache_path); + } + return false; + } + + std::sort(selected.begin(), selected.end(), + [](const auto & a, const auto & b) { return a.first < b.first; }); + + registry->tier_sizes.clear(); + registry->best_plans.clear(); + registry->pshard_disabled = false; + registry->baseline_vram_req = 0; + registry->budget_mib = best_whole ? best_whole->budget_mib : 0; + registry->cache_ubatch = selected_cache_ubatch; + + for (auto & item : selected) { + const auto & p = item.second; + if (p.is_viable && p.overrides.empty()) { + LLAMA_LOG_WARN("%s: plan cache corrupt: tier bs=%u viable but has no overrides\n", __func__, item.first); + registry->tier_sizes.clear(); + registry->best_plans.clear(); + return false; + } + registry->tier_sizes.push_back(item.first); + registry->best_plans.push_back(std::move(item.second)); + } + + if (require_exact_budget) { + LLAMA_LOG_INFO("%s: loaded %zu tier plans from exact budget=%u MiB cache_ubatch=%u variant in %s\n", + __func__, registry->tier_sizes.size(), registry->budget_mib, registry->cache_ubatch, cache_path); + } else if (best_whole) { + LLAMA_LOG_INFO("%s: loaded %zu tier plans from budget=%u MiB cache_ubatch=%u variant for current budget=%u MiB in %s\n", + __func__, registry->tier_sizes.size(), registry->budget_mib, registry->cache_ubatch, current_budget_mib, cache_path); + } else { + LLAMA_LOG_INFO("%s: loaded %zu salvaged tier plans for current budget=%u MiB from %s\n", + __func__, registry->tier_sizes.size(), current_budget_mib, cache_path); + } + return true; +} + +static bool pshard_plan_is_better(const llama_pshard_plan & candidate, const llama_pshard_plan & current) { + if (!current.is_viable) return true; + const bool candidate_has_tps = candidate.tps > 0.0f; + const bool current_has_tps = current.tps > 0.0f; + if (candidate_has_tps || current_has_tps) { + if (candidate_has_tps != current_has_tps) return candidate_has_tps; + if (candidate.tps != current.tps) return candidate.tps > current.tps; + } + if (candidate.n_pinned != current.n_pinned) return candidate.n_pinned > current.n_pinned; + if (candidate.n_attn_pinned != current.n_attn_pinned) return candidate.n_attn_pinned > current.n_attn_pinned; + if (candidate.overflow != current.overflow) return candidate.overflow > current.overflow; + return candidate.total_vram_req < current.total_vram_req; +} + +// true when every layer runs on the compute backend with no per layer override +// token_embd on host is allowed +// callers use this to skip pshard when baseline already fits +static bool pshard_plan_is_baseline_fit( + const llama_pshard_plan & plan, + uint32_t n_layers, + int32_t compute_bid) { + if (!plan.is_viable || plan.n_pinned < n_layers) return false; + for (const auto & ov : plan.overrides) { + // token_embd routes to CPU host buffer + if (ov.pattern.find("token_embd") != std::string::npos) continue; + if (ov.backend_id != compute_bid) return false; + } + return true; +} + +static llama_pshard_plan llama_pshard_search_baseline_fit_tier( + const llama_pshard_search_ctx & ctx, + const std::vector & dmds) { + const auto * path_model = ctx.path_model; + const auto * mparams = ctx.mparams; + const auto * cparams = ctx.cparams; + const auto n_layers = ctx.n_layers; + const auto vram_free = ctx.vram_free; + const auto gpu_buft = ctx.gpu_buft; + const auto host_buft = ctx.host_buft; + const auto & layout = ctx.layout; + + llama_pshard_plan plan; + plan.batch_size = cparams->n_batch; + + llama_model_tensor_buft_override local_overrides[4096]; + local_overrides[0] = { nullptr, nullptr, -1 }; + + llama_model_params mp_copy = *mparams; + mp_copy.pshard = false; + mp_copy.pshard_delegate_compute = false; + mp_copy.tensor_buft_overrides = nullptr; + mp_copy.max_vram_alloc = std::max(1, pshard_bytes_to_mib_ceil(vram_free)); + + llama_context_params cp_copy = *cparams; + cp_copy.pshard = false; + + float ts[16] = {}; + size_t margins[16] = {}; + + try { + llama_params_fit_impl(path_model, &mp_copy, &cp_copy, ts, local_overrides, margins, 0, + GGML_LOG_LEVEL_ERROR); + + const uint32_t ngl = (mp_copy.n_gpu_layers < 0) + ? (n_layers + 1) + : std::min((uint32_t)mp_copy.n_gpu_layers, n_layers + 1); + plan.n_pinned = (ngl > 0) ? (ngl - 1) : 0; + plan.pin_from_back = false; + plan.output_on_gpu = (ngl > 0); + + for (auto * ov = local_overrides; ov->pattern; ++ov) { + if (ov->buft == gpu_buft) { + ov->backend_id = layout.compute; + } else { + ov->buft = host_buft; + ov->backend_id = layout.cpu; + } + plan.overrides.push_back({ov->pattern, ov->buft, ov->backend_id}); + } + + mp_copy.tensor_buft_overrides = local_overrides[0].pattern ? local_overrides : nullptr; + const auto d = llama_pshard_probe_memory(ctx, mp_copy, cp_copy, GGML_LOG_LEVEL_ERROR); + plan.total_vram_req = d[0].mb.total(); + plan.scratch_measured = d[0].mb.compute; + plan.cache_measured = d[0].mb.context; + plan.is_viable = ((int64_t)plan.total_vram_req <= (int64_t)vram_free); + } catch (const std::exception & e) { + LLAMA_LOG_WARN("%s: baseline probe failed (bs=%u): %s\n", + __func__, cparams->n_batch, e.what()); + plan.is_viable = false; + } catch (...) { + LLAMA_LOG_WARN("%s: baseline probe failed (bs=%u): unknown exception\n", + __func__, cparams->n_batch); + plan.is_viable = false; + } + + LLAMA_LOG_INFO("%s: [bs=%-4u baseline] n_pinned=%2u/%2u, vram=%7.1f MiB, %s%s\n", + __func__, cparams->n_batch, plan.n_pinned, n_layers, + plan.total_vram_req / (1024.0 * 1024.0), + plan.is_viable ? "VIABLE" : "NOT VIABLE", + pshard_plan_is_baseline_fit(plan, n_layers, layout.compute) ? " (full fit)" : ""); + + return plan; +} + +// use attention priority when a forced strategy cannot fit the tier +// caller must set ctx.cparams for the target tier +static llama_pshard_plan llama_pshard_attn_pin_fallback( + const llama_pshard_search_ctx & ctx, + int force_strategy, + uint32_t hi_attn = UINT32_MAX, + uint32_t hi_pinned = UINT32_MAX) { + llama_pshard_plan fallback = llama_pshard_search_attn_pin(ctx, hi_attn, hi_pinned); + fallback.batch_size = ctx.cparams->n_batch; + LLAMA_LOG_INFO("llama_params_fit_pshard: [bs=%-4u %-10s] forced %s non-viable, STATIC_ATTNPRIO_ALLMODELS fallback: n_pinned=%2u/%2u, %s\n", + ctx.cparams->n_batch, llama_pshard_strategy_name(LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS), + llama_pshard_strategy_name((llama_pshard_strategy)force_strategy), + fallback.n_pinned, fallback.n_attn_pinned, + fallback.is_viable ? "VIABLE" : "NOT VIABLE"); + return fallback; +} + +static llama_pshard_plan llama_pshard_search_tier( + const llama_pshard_search_ctx & ctx, + int force_strategy, + const std::vector & dmds, + llama_pshard_tier_prune & prune) { + + const auto * path_model = ctx.path_model; + const auto * mparams = ctx.mparams; + const auto * cparams = ctx.cparams; + auto * tensor_buft_overrides = ctx.overrides; + const auto n_layers = ctx.n_layers; + const auto vram_free = ctx.vram_free; + const auto gpu_buft = ctx.gpu_buft; + const auto host_buft = ctx.host_buft; + const auto & layout = ctx.layout; + const auto is_moe = ctx.is_moe; + + llama_pshard_plan best; + + for (int s = 0; s < LLAMA_PSHARD_COUNT; s++) { + if (force_strategy >= 0 && force_strategy != s) continue; + if (prune.skip[s]) continue; + + llama_pshard_strategy strategy = (llama_pshard_strategy)s; + llama_pshard_plan plan; + + if (strategy == LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS) { + plan = llama_pshard_search_attn_pin(ctx, prune.hi_attn, prune.hi_pinned[s]); + } else { + plan = llama_pshard_search_strategy(ctx, strategy, prune.hi_pinned[s]); + } + + plan.batch_size = cparams->n_batch; + + { + const char * status = plan.is_viable ? "VIABLE" : "NOT VIABLE"; + char tps_buf[32] = ""; + if (plan.tps > 0.0f) { snprintf(tps_buf, sizeof(tps_buf), ", tps=%.1f", plan.tps); } + + if (plan.n_attn_pinned > 0) { + LLAMA_LOG_INFO("%s: [bs=%-4u %-10s] n_pinned=%2u (attn=%2u), overflow=%-4s, vram=%7.1f MiB, %s%s\n", + __func__, cparams->n_batch, llama_pshard_strategy_name(strategy), + plan.n_pinned, plan.n_attn_pinned, PSHARD_FRAC_NAMES[plan.overflow], + plan.total_vram_req / (1024.0 * 1024.0), status, tps_buf); + } else { + LLAMA_LOG_INFO("%s: [bs=%-4u %-10s] n_pinned=%2u, overflow=%-4s, vram=%7.1f MiB, %s%s\n", + __func__, cparams->n_batch, llama_pshard_strategy_name(strategy), + plan.n_pinned, PSHARD_FRAC_NAMES[plan.overflow], + plan.total_vram_req / (1024.0 * 1024.0), status, tps_buf); + } + } + + prune.update(s, plan); + + if (plan.is_viable && pshard_plan_is_better(plan, best)) { + best = plan; + } + } + + if (!best.is_viable && force_strategy >= 0 && force_strategy != LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS) { + llama_pshard_plan fallback = llama_pshard_attn_pin_fallback( + ctx, force_strategy, prune.hi_attn, prune.hi_pinned[LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS]); + prune.update(LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS, fallback); + if (fallback.is_viable) { + best = fallback; + } + } + + return best; +} + +// probe one strategy across tiers +static void llama_pshard_strategy_sweep( + int strategy, + const llama_pshard_search_ctx & ctx_template, + const llama_context_params & cparams_base, + const std::vector & dmds, + const llama_pshard_plan_registry & registry, + int force_strategy, + llama_pshard_plan * out_plans, // write out_plans[tier] + size_t n_tiers, + size_t first_tier) { + + if (force_strategy >= 0 && force_strategy != strategy) return; + + llama_model_tensor_buft_override local_overrides[4096]; + llama_pshard_search_ctx ctx = ctx_template; + ctx.overrides = local_overrides; + + llama_pshard_tier_prune prune; + prune.init(ctx.n_layers); + + // smaller tiers start from the previous fit + // smaller batches usually need less scratch + // hybrid SSM graphs are not monotonic in batch size + uint32_t prev_n_pinned = 0; + + for (int t = (int)n_tiers - 1; t >= (int)first_tier; t--) { + if (prune.skip[strategy]) break; + + llama_context_params cp_tier = cparams_base; + cp_tier.n_batch = registry.tier_sizes[t]; + cp_tier.n_ubatch = cp_tier.n_batch; + ctx.cparams = &cp_tier; + + llama_pshard_strategy strat = (llama_pshard_strategy)strategy; + llama_pshard_plan plan; + + if (out_plans[t].is_viable && out_plans[t].strategy == strat) { + LLAMA_LOG_INFO("%s: === tier %d (bs=%u) [cached %s] ===\n", + __func__, t, cp_tier.n_batch, llama_pshard_strategy_name(strat)); + plan = out_plans[t]; + } else { + LLAMA_LOG_INFO("%s: === tier %d (bs=%u) [parallel %s] ===\n", + __func__, t, cp_tier.n_batch, llama_pshard_strategy_name(strat)); + + const uint32_t lo_hint = ctx.has_rs ? 0 : prev_n_pinned; + + if (strat == LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS) { + plan = llama_pshard_search_attn_pin(ctx, prune.hi_attn, UINT32_MAX, lo_hint); + plan.batch_size = cp_tier.n_batch; + } else { + plan = llama_pshard_search_strategy(ctx, strat, UINT32_MAX, lo_hint); + plan.batch_size = cp_tier.n_batch; + } + } + + prune.update(strategy, plan); + + if (plan.is_viable && plan.n_pinned > prev_n_pinned) { + prev_n_pinned = plan.n_pinned; + } + + { + const char * status = plan.is_viable ? "VIABLE" : "NOT VIABLE"; + char tps_buf[32] = ""; + if (plan.tps > 0.0f) { snprintf(tps_buf, sizeof(tps_buf), ", tps=%.1f", plan.tps); } + + if (plan.n_attn_pinned > 0) { + LLAMA_LOG_INFO("%s: [tier=%d bs=%-4u %-10s] n_pinned=%2u (attn=%2u), overflow=%-4s, vram=%7.1f MiB, %s%s\n", + __func__, t, cp_tier.n_batch, llama_pshard_strategy_name(strat), + plan.n_pinned, plan.n_attn_pinned, PSHARD_FRAC_NAMES[plan.overflow], + plan.total_vram_req / (1024.0 * 1024.0), status, tps_buf); + } else { + LLAMA_LOG_INFO("%s: [tier=%d bs=%-4u %-10s] n_pinned=%2u, overflow=%-4s, vram=%7.1f MiB, %s%s\n", + __func__, t, cp_tier.n_batch, llama_pshard_strategy_name(strat), + plan.n_pinned, PSHARD_FRAC_NAMES[plan.overflow], + plan.total_vram_req / (1024.0 * 1024.0), status, tps_buf); + } + } + + out_plans[t] = plan; + } +} + +static void llama_pshard_parallel_worker( + std::atomic & next_strategy, + const llama_pshard_search_ctx & ctx_template, + const llama_context_params & cparams_base, + const std::vector & dmds, + const llama_pshard_plan_registry & registry, + int force_strategy, + llama_pshard_plan * all_plans, + size_t n_tiers, + size_t first_tier) { + + while (true) { + int s = next_strategy.fetch_add(1); + if (s >= LLAMA_PSHARD_COUNT) return; + + llama_pshard_plan * out = all_plans + s * n_tiers; + llama_pshard_strategy_sweep(s, ctx_template, cparams_base, dmds, registry, force_strategy, out, n_tiers, first_tier); + } +} + +static bool llama_pshard_params_supported( + const struct llama_model_params * mparams, + const struct llama_context_params * cparams) { + const llama_model_params default_mparams = llama_model_default_params(); + + auto disable = [](const char * reason) { + LLAMA_LOG_WARN("%s: %s, disabling pshard\n", "llama_params_fit_pshard", reason); + return false; + }; + + if (!cparams->offload_kqv) { + return disable("offload_kqv=false is not supported"); + } + if (mparams->split_mode == LLAMA_SPLIT_MODE_TENSOR) { + return disable("SPLIT_MODE_TENSOR is not supported"); + } + if (mparams->split_mode == LLAMA_SPLIT_MODE_ROW) { + return disable("SPLIT_MODE_ROW is not supported"); + } + if (mparams->n_gpu_layers != default_mparams.n_gpu_layers) { + return disable("n_gpu_layers is already set by the user"); + } + if (mparams->tensor_split) { + for (size_t i = 0; i < llama_max_devices(); i++) { + if (mparams->tensor_split[i] != 0.0f) { + return disable("tensor_split is already set by the user"); + } + } + } + if (mparams->tensor_buft_overrides && + (mparams->tensor_buft_overrides->pattern || mparams->tensor_buft_overrides->buft)) { + return disable("tensor_buft_overrides are already set by the user"); + } + + return true; +} + +void llama_params_fit_pshard( + const char * path_model, + struct llama_model_params * mparams, + struct llama_context_params * cparams, + struct llama_model_tensor_buft_override * tensor_buft_overrides, + size_t max_vram_mb, + size_t fit_target_mb) { + const int64_t t0_us = llama_time_us(); + + if (!llama_pshard_params_supported(mparams, cparams)) { + mparams->pshard = false; + mparams->pshard_delegate_compute = false; + cparams->pshard = false; + return; + } + + mparams->pshard = true; + mparams->pshard_delegate_compute = false; + cparams->pshard = true; + + // step 1: probe device memory and model parameters + std::vector devs; + uint32_t hp_ngl = 0, hp_nct = 0, hp_nex = 0, hp_nr = 0; + + llama_model_params mparams_probe = *mparams; + mparams_probe.pshard = false; + const auto dmds = llama_get_device_memory_data( + path_model, &mparams_probe, cparams, devs, hp_ngl, hp_nct, hp_nex, hp_nr, GGML_LOG_LEVEL_ERROR); + + if (devs.empty()) { + LLAMA_LOG_ERROR("%s: no GPU devices found\n", __func__); + return; + } + + const uint32_t n_layers = hp_ngl; + const size_t mib = 1024ULL * 1024ULL; + const size_t actual_vram_free = dmds[0].free; + const size_t fit_target_bytes = fit_target_mb * mib; + const size_t vram_free = max_vram_mb > 0 + ? max_vram_mb * mib + : (actual_vram_free > fit_target_bytes ? actual_vram_free - fit_target_bytes : 0); + + mparams->max_vram_alloc = std::max(1, pshard_bytes_to_mib_ceil(vram_free)); + + LLAMA_LOG_INFO("%s: probing pshard plans: %u layers, %.1f MiB VRAM free, %.1f MiB budget%s\n", + __func__, n_layers, + actual_vram_free / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + max_vram_mb > 0 ? " (-mva)" : " (free - fit target)"); + + // step 2: read forced strategy from env (PSHARD_STRATEGY) + const int force_strategy = pshard_strategy_from_env(); + if (force_strategy >= 0) { + LLAMA_LOG_INFO("%s: forcing strategy %s (PSHARD_STRATEGY=%s)\n", + __func__, llama_pshard_strategy_name((llama_pshard_strategy)force_strategy), + getenv("PSHARD_STRATEGY")); + } else if (getenv("PSHARD_STRATEGY")) { + LLAMA_LOG_WARN("%s: invalid PSHARD_STRATEGY='%s', ignoring\n", + __func__, getenv("PSHARD_STRATEGY")); + } + + // step 3: derive layout, buftypes, optional benchmark predictor, search ctx + ggml_backend_buffer_type_t gpu_buft = ggml_backend_dev_buffer_type(devs[0].dev); + ggml_backend_buffer_type_t host_buft = ggml_backend_dev_host_buffer_type(devs[0].dev); + if (!host_buft) { + host_buft = ggml_backend_cpu_buffer_type(); + } + const int32_t cpu_bid = pshard_dev_layout::compute_cpu_backend_id(devs.size()); + const pshard_dev_layout layout = pshard_dev_layout::for_device(0, cpu_bid); + + std::unique_ptr predictor; + { + const char * env_cpu = getenv("PSHARD_CPU_PROFILE"); + const char * env_gpu = getenv("PSHARD_GPU_PROFILE"); + const char * cpu_path = env_cpu ? env_cpu : "cpu_profile.txt"; + const char * gpu_path = env_gpu ? env_gpu : "gpu_profile.txt"; + + auto p = std::make_unique(); + const bool has_cpu = p->load_cpu(cpu_path, cparams->n_threads); + const bool has_gpu = p->load_gpu(gpu_path); + if (has_cpu || has_gpu) { + predictor = std::move(p); + LLAMA_LOG_INFO("%s: benchmark predictor loaded (cpu=%s gpu=%s)\n", + __func__, has_cpu ? "yes" : "no", has_gpu ? "yes" : "no"); + } + } + + const uint32_t n_ctx_plan = cparams->n_ctx > 0 ? cparams->n_ctx : hp_nct; + + llama_pshard_search_ctx ctx = { + path_model, mparams, cparams, tensor_buft_overrides, + n_layers, vram_free, gpu_buft, host_buft, layout, + /*is_moe=*/(hp_nex > 0), /*has_rs=*/(hp_nr > 0), + predictor.get(), n_ctx_plan, 0, + }; + + // step 4: registry lookup -- try to load .tensor_overrides.pshard_registry; merge cached plans by tier + const std::string cache_path = std::string(path_model) + ".tensor_overrides.pshard_registry"; + + int64_t model_file_size = 0; + if (FILE * mf = fopen(path_model, "rb")) { +#ifdef _WIN32 + _fseeki64(mf, 0, SEEK_END); + model_file_size = _ftelli64(mf); +#else + fseeko(mf, 0, SEEK_END); + model_file_size = ftello(mf); +#endif + fclose(mf); + } + const uint64_t fp = pshard_registry_fingerprint(mparams, cparams, model_file_size); + + llama_pshard_plan_registry * registry = mparams->pshard_registry; + bool needs_probe = true; + + if (registry) { + const uint32_t requested_tier_max = registry->cache_ubatch; + const uint32_t tier_max_auto = std::min(std::max(cparams->n_batch, (uint32_t) 16384), n_ctx_plan); + const uint32_t tier_max = std::min(requested_tier_max > 0 ? requested_tier_max : tier_max_auto, n_ctx_plan); + + if (registry->cache_ubatch != tier_max) { + registry->init(tier_max, cparams->n_seq_max); + } + + registry->budget_mib = pshard_bytes_to_mib_ceil(vram_free); + registry->cache_ubatch = registry->tier_sizes.empty() ? 0 : registry->tier_sizes.back(); + ctx.cache_ubatch = registry->cache_ubatch; + + // save requested tiers before load (load overwrites tier_sizes) + std::vector requested_tiers = registry->tier_sizes; + + llama_pshard_plan_registry cached; + cached.budget_mib = registry->budget_mib; + cached.cache_ubatch = registry->cache_ubatch; + if (!mparams->pshard_cache_skip_load && + pshard_registry_load(&cached, fp, cache_path.c_str(), host_buft, vram_free, true)) { + // merge cached plans into registry by matching batch size + std::unordered_map cache_map; + for (size_t i = 0; i < cached.tier_sizes.size(); i++) { + cache_map[cached.tier_sizes[i]] = cached.best_plans[i]; + } + + registry->tier_sizes = requested_tiers; + registry->best_plans.resize(requested_tiers.size()); + registry->pshard_disabled = cached.pshard_disabled; + registry->baseline_vram_req = cached.baseline_vram_req; + registry->budget_mib = cached.budget_mib; + registry->cache_ubatch = cached.cache_ubatch; + + size_t n_hit = 0; + for (size_t i = 0; i < requested_tiers.size(); i++) { + auto it = cache_map.find(requested_tiers[i]); + if (it != cache_map.end()) { + registry->best_plans[i] = it->second; + n_hit++; + } + } + + if (n_hit == requested_tiers.size()) { + needs_probe = false; + LLAMA_LOG_INFO("%s: loaded all %zu tiers from exact budget=%u MiB cache variant\n", + __func__, n_hit, registry->budget_mib); + } else if (n_hit > 0) { + LLAMA_LOG_INFO("%s: loaded %zu/%zu tiers from exact budget=%u MiB cache variant, %zu need probing\n", + __func__, n_hit, requested_tiers.size(), registry->budget_mib, requested_tiers.size() - n_hit); + } + } + + // common setup for cache hit and miss + } + + // step 4b: skip pshard when a cached baseline variant fits this budget + if (registry && registry->pshard_disabled) { + LLAMA_LOG_INFO("%s: cache says baseline %.1f MiB fits this budget (variant budget=%u MiB cache_ubatch=%u), using baseline loading\n", + __func__, registry->baseline_vram_req / (1024.0 * 1024.0), registry->budget_mib, registry->cache_ubatch); + mparams->pshard = false; + mparams->pshard_delegate_compute = false; + cparams->pshard = false; + mparams->n_gpu_layers = n_layers + 1; + tensor_buft_overrides[0] = { nullptr, nullptr, -1 }; + mparams->tensor_buft_overrides = nullptr; + + const int64_t t1_us = llama_time_us(); + LLAMA_LOG_INFO("%s: best strategy: baseline (cached), all %u layers on GPU, took %.2f s\n", + __func__, n_layers, (t1_us - t0_us) * 1e-6); + return; + } + + // step 5: probe tiers largest first and skip pshard when baseline already fits + if (registry && needs_probe) { + const size_t n_tiers = registry->tier_sizes.size(); + + // step 5a: baseline off-ramp only for global tiers + static constexpr uint32_t GLOBAL_FIT_MIN_BATCH = 512; + size_t first_probe_tier = 0; + size_t min_global_tier = n_tiers; + const bool run_baseline_offramp = force_strategy < 0; + if (run_baseline_offramp) { + for (size_t t = 0; t < n_tiers; t++) { + if (registry->tier_sizes[t] >= GLOBAL_FIT_MIN_BATCH) { + min_global_tier = t; + break; + } + } + + if (min_global_tier < n_tiers) { + size_t global_fit_tier = n_tiers; + + for (size_t t = n_tiers; t-- > min_global_tier; ) { + llama_pshard_plan baseline_plan; + if (pshard_plan_is_baseline_fit(registry->best_plans[t], n_layers, layout.compute)) { + LLAMA_LOG_INFO("%s: === tier %zu (bs=%u) [cached global-fit check] ===\n", + __func__, t, registry->tier_sizes[t]); + baseline_plan = registry->best_plans[t]; + } else { + const llama_context_params * saved = ctx.cparams; + llama_context_params cp_tier = *cparams; + cp_tier.n_batch = registry->tier_sizes[t]; + cp_tier.n_ubatch = cp_tier.n_batch; + ctx.cparams = &cp_tier; + + LLAMA_LOG_INFO("%s: === tier %zu (bs=%u) [global-fit check] ===\n", + __func__, t, registry->tier_sizes[t]); + baseline_plan = llama_pshard_search_baseline_fit_tier(ctx, dmds); + + ctx.cparams = saved; + } + + if (pshard_plan_is_baseline_fit(baseline_plan, n_layers, layout.compute)) { + registry->best_plans[t] = baseline_plan; + global_fit_tier = t; + break; + } + } + + if (global_fit_tier < n_tiers) { + const llama_pshard_plan & global = registry->best_plans[global_fit_tier]; + const uint32_t global_ubatch = registry->tier_sizes[global_fit_tier]; + const uint32_t global_cache_ubatch = std::min(global_ubatch, GLOBAL_FIT_MIN_BATCH); + const size_t global_vram_req = global.total_vram_req; + + registry->pshard_disabled = true; + registry->baseline_vram_req = global_vram_req; + registry->cache_ubatch = global_cache_ubatch; + registry->tier_sizes.clear(); + registry->best_plans.clear(); + pshard_registry_save(registry, fp, cache_path.c_str(), host_buft, cparams); + + mparams->pshard = false; + mparams->pshard_delegate_compute = false; + cparams->pshard = false; + mparams->n_gpu_layers = n_layers + 1; + tensor_buft_overrides[0] = { nullptr, nullptr, -1 }; + mparams->tensor_buft_overrides = nullptr; + + const int64_t t1_us = llama_time_us(); + LLAMA_LOG_INFO("%s: full baseline fit at tier %zu (bs=%u, %.1f MiB); using baseline loading with cache_ubatch=%u, took %.2f s\n", + __func__, global_fit_tier, global_ubatch, + global_vram_req / (1024.0 * 1024.0), + global_cache_ubatch, (t1_us - t0_us) * 1e-6); + return; + } + } + + LLAMA_LOG_INFO("%s: no full-fit global plan found down to bs=%u\n", + __func__, GLOBAL_FIT_MIN_BATCH); + } else { + LLAMA_LOG_INFO("%s: skipping baseline global-fit check for forced strategy %s\n", + __func__, llama_pshard_strategy_name((llama_pshard_strategy) force_strategy)); + } + + // step 5b: probe remaining tiers + if (first_probe_tier < n_tiers) + { + std::vector all_plans(LLAMA_PSHARD_COUNT * n_tiers); + for (size_t t = 0; t < n_tiers; t++) { + const llama_pshard_plan & plan = registry->best_plans[t]; + if (plan.is_viable && plan.strategy >= 0 && plan.strategy < LLAMA_PSHARD_COUNT) { + all_plans[(int) plan.strategy * n_tiers + t] = plan; + } + } + + int n_workers = std::min((int)cparams->n_threads, (int)LLAMA_PSHARD_COUNT); + n_workers = std::max(n_workers, 1); + + int n_pshard_strategies = LLAMA_PSHARD_COUNT; + n_workers = std::min(n_workers, n_pshard_strategies); + + LLAMA_LOG_INFO("%s: parallel planning with %d workers (%d pshard strategies)\n", + __func__, n_workers, n_pshard_strategies); + + std::atomic next_strategy{0}; + + if (n_workers <= 1) { + llama_pshard_parallel_worker(next_strategy, ctx, *cparams, dmds, + *registry, force_strategy, all_plans.data(), n_tiers, first_probe_tier); + } else { + std::vector threads; + for (int w = 1; w < n_workers; w++) { + threads.emplace_back(llama_pshard_parallel_worker, + std::ref(next_strategy), std::cref(ctx), std::cref(*cparams), + std::cref(dmds), std::cref(*registry), force_strategy, + all_plans.data(), n_tiers, first_probe_tier); + } + llama_pshard_parallel_worker(next_strategy, ctx, *cparams, dmds, + *registry, force_strategy, all_plans.data(), n_tiers, first_probe_tier); + + for (auto & t : threads) t.join(); + } + + for (size_t t = first_probe_tier; t < n_tiers; t++) { + llama_pshard_plan best; + for (int s = 0; s < LLAMA_PSHARD_COUNT; s++) { + auto & p = all_plans[s * n_tiers + t]; + if (p.is_viable && pshard_plan_is_better(p, best)) { + best = p; + } + } + registry->best_plans[t] = best; + } + + // step 5d: try attention priority fallback for forced strategies that cannot fit + // forced strategy sweeps can leave smaller tiers unfilled + if (force_strategy >= 0 && force_strategy != LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS) { + for (size_t t = first_probe_tier; t < n_tiers; t++) { + if (registry->best_plans[t].is_viable) continue; + + const llama_context_params * saved = ctx.cparams; + llama_context_params cp_tier = *cparams; + cp_tier.n_batch = registry->tier_sizes[t]; + cp_tier.n_ubatch = cp_tier.n_batch; + ctx.cparams = &cp_tier; + + llama_pshard_plan fallback = llama_pshard_attn_pin_fallback(ctx, force_strategy); + if (fallback.is_viable) { + registry->best_plans[t] = fallback; + } + + ctx.cparams = saved; + } + } + } + + pshard_registry_save(registry, fp, cache_path.c_str(), host_buft, cparams); + } + + // step 6: pick the active plan + if (registry) { + for (size_t t = registry->tier_sizes.size(); t-- > 0; ) { + auto * best = registry->get_best(t); + if (best && best->is_viable) { + registry->active_plan = best; + break; + } + } + } + + llama_pshard_plan best_plan; + if (registry && registry->active_plan) { + best_plan = *registry->active_plan; + } else { + llama_pshard_tier_prune prune_single; + prune_single.init(n_layers); + best_plan = llama_pshard_search_tier(ctx, force_strategy, dmds, prune_single); + if (best_plan.is_viable && registry) { + for (size_t t = 0; t < registry->tier_sizes.size(); t++) { + if (registry->tier_sizes[t] == best_plan.batch_size) { + registry->best_plans[t] = best_plan; + registry->active_plan = ®istry->best_plans[t]; + pshard_registry_save(registry, fp, cache_path.c_str(), host_buft, cparams); + break; + } + } + } + } + + // step 7: fall back to all cpu when no plan is viable + if (!best_plan.is_viable) { + LLAMA_LOG_WARN("%s: no viable plan found, falling back to STATIC_ATTNPRIO_ALLMODELS with n_pinned=0\n", __func__); + llama_pshard_generate_overrides(0, n_layers, gpu_buft, host_buft, tensor_buft_overrides, + LLAMA_LAYER_FRACTION_NONE, LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS, layout, + /*pin_from_back=*/false, /*output_on_gpu=*/false, /*n_attn_pinned=*/0); + mparams->pshard_delegate_compute = true; + mparams->n_gpu_layers = n_layers + 1; + mparams->tensor_buft_overrides = tensor_buft_overrides; + return; + } + + // step 8: apply best plan to tensor_buft_overrides + llama_pshard_generate_overrides(best_plan.n_pinned, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, (llama_layer_fraction)best_plan.overflow, best_plan.strategy, layout, + best_plan.pin_from_back, best_plan.output_on_gpu, best_plan.n_attn_pinned); + + for (size_t i = 0; tensor_buft_overrides[i].pattern; i++) { + if (tensor_buft_overrides[i].backend_id == layout.compute) { + tensor_buft_overrides[i].buft = host_buft; + } + } + + mparams->pshard_delegate_compute = llama_pshard_strategy_delegates_compute(best_plan.strategy); + mparams->n_gpu_layers = n_layers + 1; + mparams->tensor_buft_overrides = tensor_buft_overrides; + + const int64_t t1_us = llama_time_us(); + LLAMA_LOG_INFO("%s: best strategy: %s, n_pinned=%u, n_attn=%u/%u%s%s, took %.2f s\n", + __func__, llama_pshard_strategy_name(best_plan.strategy), + best_plan.n_pinned, best_plan.n_attn_pinned, n_layers, + best_plan.overflow ? " (partial: " : "", + best_plan.overflow ? PSHARD_FRAC_NAMES[best_plan.overflow] : "", + (t1_us - t0_us) * 1e-6); + + { + int n_bid0 = 0, n_total = 0; + for (const auto * ov = tensor_buft_overrides; ov->pattern; ++ov) { + n_total++; + if (ov->backend_id == layout.compute) n_bid0++; + LLAMA_LOG_DEBUG("%s: override: %-25s -> %-15s backend_id=%d\n", + __func__, ov->pattern, ggml_backend_buft_name(ov->buft), ov->backend_id); + } + LLAMA_LOG_INFO("%s: %d overrides, %d with bid=%d (compute)\n", + __func__, n_total, n_bid0, layout.compute); + } +} diff --git a/src/llama-pshard-plan.h b/src/llama-pshard-plan.h new file mode 100644 index 000000000000..45d5d288bce4 --- /dev/null +++ b/src/llama-pshard-plan.h @@ -0,0 +1,200 @@ +#pragma once + +#include "llama.h" +#include "llama-context.h" +#include "llama-cparams.h" +#include "llama-model.h" + +#include +#include +#include +#include +#include +#include + +enum llama_pshard_strategy { + LLAMA_PSHARD_GPUONLY_LAYERPIN_LAYERSTREAM = 0, + LLAMA_PSHARD_GPUONLY_ATTNPIN_FFNSTREAM = 1, + LLAMA_PSHARD_DYNAMIC_FFNCPU_ATTNSTREAM = 2, + LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS = 3, + LLAMA_PSHARD_COUNT +}; + +// strategy name for logging +inline const char * llama_pshard_strategy_name(llama_pshard_strategy s) { + switch (s) { + case LLAMA_PSHARD_GPUONLY_LAYERPIN_LAYERSTREAM: return "GPUONLY_LAYERPIN_LAYERSTREAM"; + case LLAMA_PSHARD_GPUONLY_ATTNPIN_FFNSTREAM: return "GPUONLY_ATTNPIN_FFNSTREAM"; + case LLAMA_PSHARD_DYNAMIC_FFNCPU_ATTNSTREAM: return "DYNAMIC_FFNCPU_ATTNSTREAM"; + case LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS: return "STATIC_ATTNPRIO_ALLMODELS"; + default: return "UNKNOWN"; + } +} + +inline bool llama_pshard_strategy_delegates_compute(llama_pshard_strategy s) { + return s == LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS; +} + +// PSHARD_STRATEGY accepts a name or numeric id +inline int pshard_strategy_from_env() { + const char * env = getenv("PSHARD_STRATEGY"); + if (!env || !*env) return -1; + for (int i = 0; i < LLAMA_PSHARD_COUNT; i++) { + if (strcmp(env, llama_pshard_strategy_name((llama_pshard_strategy)i)) == 0) { + return i; + } + } + char * end = nullptr; + long v = strtol(env, &end, 10); + if (end != env && *end == '\0' && v >= 0 && v < LLAMA_PSHARD_COUNT) { + return (int)v; + } + return -1; +} + +// cached tensor override entry +struct llama_pshard_override { + std::string pattern; + ggml_backend_buffer_type_t buft; + int32_t backend_id; +}; + +struct llama_pshard_plan { + llama_pshard_strategy strategy = LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS; + uint32_t batch_size = 0; + uint32_t n_pinned = 0; // fully pinned layers (all tensors on GPU) + uint32_t n_attn_pinned = 0; // attention priority layers on GPU (>= n_pinned) + int overflow = 0; // llama_layer_fraction + bool pin_from_back = false; + bool output_on_gpu = false; + + std::vector overrides; + + size_t total_vram_req = 0; + size_t scratch_measured = 0; + size_t cache_measured = 0; + float tps = 0.0f; // predicted tokens/sec (0 = no benchmark data) + bool is_viable = false; +}; + +enum llama_layer_fraction { + LLAMA_LAYER_FRACTION_NONE = 0, + LLAMA_LAYER_FRACTION_ATTN = 1, + LLAMA_LAYER_FRACTION_UP = 2, + LLAMA_LAYER_FRACTION_GATE = 3, + LLAMA_LAYER_FRACTION_MOE = 4, +}; + +const char * llama_get_overflow_pattern(size_t il, llama_layer_fraction lf); + +struct llama_device_memory_data { + int64_t total; + int64_t free; + llama_memory_breakdown_data mb; +}; + +// probe hook runs before context teardown +// used by TPS prediction to inspect scheduler splits +typedef void (*llama_probe_hook_t)(llama_context * ctx, void * user_data); + +std::vector llama_get_device_memory_data( + const char * path_model, const struct llama_model_params * mparams, + const struct llama_context_params * cparams, + std::vector & devs, uint32_t & hp_ngl, + uint32_t & hp_n_ctx_train, uint32_t & hp_n_expert, uint32_t & hp_n_embd_r, + enum ggml_log_level log_level, + llama_probe_hook_t probe_hook = nullptr, + void * probe_hook_data = nullptr, + uint32_t probe_n_tokens = 0, + uint32_t probe_n_outputs = 0); + +// fit params entry point used by pshard planning +void llama_params_fit_impl( + const char * path_model, struct llama_model_params * mparams, struct llama_context_params * cparams, + float * tensor_split, struct llama_model_tensor_buft_override * tensor_buft_overrides, + size_t * margins_s, uint32_t n_ctx_min, enum ggml_log_level log_level); + +// plan cache serialization. fingerprint covers only runtime plan-compatibility params +// so the planner binary and the runtime binary can share the same cache file. +// keep this in sync with planner save +uint64_t pshard_registry_fingerprint( + const struct llama_model_params * mparams, + const struct llama_context_params * cparams, + int64_t model_file_size); + +bool pshard_registry_save( + const struct llama_pshard_plan_registry * registry, uint64_t fingerprint, + const char * cache_path, ggml_backend_buffer_type_t host_buft, + const struct llama_context_params * cparams = nullptr); + +bool pshard_registry_load( + struct llama_pshard_plan_registry * registry, uint64_t fingerprint, + const char * cache_path, ggml_backend_buffer_type_t host_buft, + size_t current_budget, bool require_exact_budget = false); + +struct llama_pshard_plan_registry { + std::vector tier_sizes; + std::vector best_plans; // one best plan per tier + llama_pshard_plan * active_plan = nullptr; + uint32_t budget_mib = 0; + uint32_t cache_ubatch = 0; + + // variant marker for a baseline load that fits + // runtime still checks baseline_vram_req against the current budget + bool pshard_disabled = false; + size_t baseline_vram_req = 0; + + void init(uint32_t n_ubatch, uint32_t n_parallel = 1, uint32_t n_draft = 0) { + tier_sizes.clear(); + best_plans.clear(); + + if (n_ubatch == 0) { + cache_ubatch = 0; + return; + } + + // decode tiers + if (n_parallel <= 1) { + tier_sizes.push_back(1); + tier_sizes.push_back(16); + if (n_draft > 0) { + uint32_t verify_tier = n_draft + 1; + if (verify_tier > 16) { + tier_sizes.push_back(verify_tier); + } + } + } else { + for (uint32_t t = 1; t <= 64 && t < 512; t *= 4) { + tier_sizes.push_back(t); + if (t == 16) { + tier_sizes.push_back(32); + } + } + } + + // prefill tiers: x2 growth from 512 + for (uint32_t t = 512; t < n_ubatch; t *= 2) { + if (tier_sizes.empty() || tier_sizes.back() < t) { + tier_sizes.push_back(t); + } + } + + if (tier_sizes.empty() || tier_sizes.back() != n_ubatch) { + tier_sizes.push_back(n_ubatch); + } + + cache_ubatch = tier_sizes.empty() ? 0 : tier_sizes.back(); + best_plans.resize(tier_sizes.size()); + } + + size_t tier_index(uint32_t batch_size) const { + for (size_t i = 0; i < tier_sizes.size(); i++) { + if (tier_sizes[i] >= batch_size) return i; + } + return tier_sizes.size() - 1; + } + + llama_pshard_plan * get_best(size_t tier) { + return best_plans[tier].is_viable ? &best_plans[tier] : nullptr; + } +}; diff --git a/src/llama.cpp b/src/llama.cpp index 484372d8d106..8ae62e381de9 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -9,6 +9,7 @@ #include "llama-model-loader.h" #include "llama-model-saver.h" #include "llama-model.h" +#include "llama-pshard-plan.h" #include "ggml.h" #include "ggml-cpp.h" @@ -23,6 +24,8 @@ #include #include #include +#include +#include #include #include @@ -46,16 +49,13 @@ const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_ty GGML_ABORT("fatal error"); } -struct llama_device_memory_data { - int64_t total; - int64_t free; - llama_memory_breakdown_data mb; -}; - -static std::vector llama_get_device_memory_data( +std::vector llama_get_device_memory_data( const char * path_model, const llama_model_params * mparams, const llama_context_params * cparams, std::vector & devs, uint32_t & hp_ngl, uint32_t & hp_n_ctx_train, uint32_t & hp_n_expert, - const ggml_log_level log_level) { + uint32_t & hp_n_embd_r, + enum ggml_log_level log_level, + llama_probe_hook_t probe_hook, void * probe_hook_data, + uint32_t probe_n_tokens, uint32_t probe_n_outputs) { struct user_data_t { struct { ggml_log_callback callback; @@ -63,6 +63,9 @@ static std::vector llama_get_device_memory_data( } original_logger; ggml_log_level min_level; // prints below this log level go to debug log }; + static std::mutex log_mutex; + std::unique_lock log_lock(log_mutex); + user_data_t ud; llama_log_get(&ud.original_logger.callback, &ud.original_logger.user_data); ud.min_level = log_level; @@ -73,6 +76,8 @@ static std::vector llama_get_device_memory_data( ud->original_logger.callback(level_eff, text, ud->original_logger.user_data); }, &ud); + log_lock.unlock(); + llama_model_params mparams_copy = *mparams; mparams_copy.no_alloc = true; mparams_copy.use_mmap = false; @@ -84,7 +89,8 @@ static std::vector llama_get_device_memory_data( throw std::runtime_error("failed to load model"); } - llama_context * ctx = llama_init_from_model(model, *cparams); + llama_context_probe_reserve probe_reserve = { probe_n_tokens, probe_n_outputs }; + llama_context * ctx = llama_init_from_model_internal(model, *cparams, probe_reserve); if (ctx == nullptr) { llama_model_free(model); llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); @@ -136,30 +142,30 @@ static std::vector llama_get_device_memory_data( hp_ngl = model->hparams.n_layer; hp_n_ctx_train = model->hparams.n_ctx_train; hp_n_expert = model->hparams.n_expert; + hp_n_embd_r = model->hparams.n_embd_r(); llama_memory_breakdown_print(ctx); // goes to debug log + if (probe_hook) { + probe_hook(ctx, probe_hook_data); + } + llama_free(ctx); llama_model_free(model); - llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); + + { + std::lock_guard lg(log_mutex); + llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); + } + return ret; } -// enum to identify part of a layer for distributing its tensors: -enum layer_fraction_t { - LAYER_FRACTION_NONE = 0, // nothing - LAYER_FRACTION_ATTN = 1, // attention - LAYER_FRACTION_UP = 2, // attention + up - LAYER_FRACTION_GATE = 3, // attention + up + gate - LAYER_FRACTION_MOE = 4, // everything but sparse MoE weights -}; -// this enum is only used in llama_params_fit_impl but needs to be defined outside of it to fix a Windows compilation issue - class llama_params_fit_exception : public std::runtime_error { using std::runtime_error::runtime_error; }; -static void llama_params_fit_impl( +void llama_params_fit_impl( const char * path_model, struct llama_model_params * mparams, struct llama_context_params * cparams, float * tensor_split, struct llama_model_tensor_buft_override * tensor_buft_overrides, size_t * margins_s, uint32_t n_ctx_min, enum ggml_log_level log_level) { @@ -174,11 +180,12 @@ static void llama_params_fit_impl( uint32_t hp_ngl = 0; // hparams.n_gpu_layers uint32_t hp_nct = 0; // hparams.n_ctx_train uint32_t hp_nex = 0; // hparams.n_expert + uint32_t hp_nr = 0; // hparams.n_embd_r // step 1: get data for default parameters and check whether any changes are necessary in the first place LLAMA_LOG_DEBUG("%s: getting device memory data for initial parameters:\n", __func__); - const dmds_t dmds_full = llama_get_device_memory_data(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + const dmds_t dmds_full = llama_get_device_memory_data(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, hp_nr, log_level); const size_t nd = devs.size(); // number of devices if (nd == 0) { LLAMA_LOG_INFO("%s: no devices with dedicated memory found\n", __func__); @@ -191,6 +198,13 @@ static void llama_params_fit_impl( margins.push_back(margins_s[id]); } + if (mparams->max_vram_alloc > 0 && nd > 0) { + const int64_t budget = (int64_t)mparams->max_vram_alloc * 1024LL * 1024LL; + margins[0] = (dmds_full[0].free > budget) ? (dmds_full[0].free - budget) : 0; + LLAMA_LOG_INFO("%s: max_vram_alloc %" PRId64 " MiB budget, margin %" PRId64 " MiB (%" PRId64 " MiB free)\n", + __func__, budget / (1024*1024), margins[0] / (1024*1024), dmds_full[0].free / (1024*1024)); + } + std::vector dev_names; { dev_names.reserve(nd); @@ -292,7 +306,7 @@ static void llama_params_fit_impl( int64_t sum_projected_used_min_ctx = 0; cparams->n_ctx = n_ctx_min; - const dmds_t dmds_min_ctx = llama_get_device_memory_data(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + const dmds_t dmds_min_ctx = llama_get_device_memory_data(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, hp_nr, log_level); for (const auto & dmd : dmds_min_ctx) { sum_projected_used_min_ctx += dmd.mb.total(); } @@ -359,52 +373,12 @@ static void llama_params_fit_impl( // - for a dense model simply fill full layers, giving each device a contiguous slice of the model // - for a MoE model, same as dense model but with all MoE tensors in system memory - // utility function that returns a static C string matching the tensors for a specific layer index and layer fraction: - auto get_overflow_pattern = [&](const size_t il, const layer_fraction_t lf) -> const char * { - constexpr size_t n_strings = 1000; - if (il >= n_strings) { - throw std::runtime_error("at most " + std::to_string(n_strings) + " model layers are supported"); - } - switch (lf) { - case LAYER_FRACTION_ATTN: { - static std::array patterns; - if (patterns[il].empty()) { - patterns[il] = "blk\\." + std::to_string(il) + "\\.ffn_(gate|up|gate_up|down).*"; - } - return patterns[il].c_str(); - } - case LAYER_FRACTION_UP: { - static std::array patterns; - if (patterns[il].empty()) { - patterns[il] = "blk\\." + std::to_string(il) + "\\.ffn_(gate|gate_up|down).*"; - } - return patterns[il].c_str(); - } - case LAYER_FRACTION_GATE: { - static std::array patterns; - if (patterns[il].empty()) { - patterns[il] = "blk\\." + std::to_string(il) + "\\.ffn_down.*"; - } - return patterns[il].c_str(); - } - case LAYER_FRACTION_MOE: { - static std::array patterns; - if (patterns[il].empty()) { - patterns[il] = "blk\\." + std::to_string(il) + "\\.ffn_(up|down|gate_up|gate)_(ch|)exps"; - } - return patterns[il].c_str(); - } - default: - GGML_ABORT("fatal error"); - } - }; - struct ngl_t { uint32_t n_layer = 0; // number of total layers uint32_t n_part = 0; // number of partial layers, <= n_layer - // for the first partial layer varying parts can overflow, all further layers use LAYER_FRACTION_MOE: - layer_fraction_t overflow_type = LAYER_FRACTION_MOE; + // for the first partial layer varying parts can overflow, all further layers use LLAMA_LAYER_FRACTION_MOE: + llama_layer_fraction overflow_type = LLAMA_LAYER_FRACTION_MOE; uint32_t n_full() const { assert(n_layer >= n_part); @@ -414,43 +388,49 @@ static void llama_params_fit_impl( const size_t ntbo = llama_max_tensor_buft_overrides(); - // utility function to set n_gpu_layers and tensor_split auto set_ngl_tensor_split_tbo = [&]( const std::vector & ngl_per_device, const std::vector & overflow_bufts, llama_model_params & mparams) { - mparams.n_gpu_layers = 0; + uint32_t total_ngl = 0; for (size_t id = 0; id < nd; id++) { - mparams.n_gpu_layers += ngl_per_device[id].n_layer; + total_ngl += ngl_per_device[id].n_layer; if (nd > 1) { tensor_split[id] = ngl_per_device[id].n_layer; } } - assert(uint32_t(mparams.n_gpu_layers) <= hp_ngl + 1); - uint32_t il0 = hp_ngl + 1 - mparams.n_gpu_layers; // start index for tensor buft overrides + assert(total_ngl <= hp_ngl + 1); mparams.tensor_split = tensor_split; size_t itbo = 0; + + mparams.n_gpu_layers = total_ngl; + uint32_t il0 = hp_ngl + 1 - total_ngl; + for (size_t id = 0; id < nd; id++) { 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 llama_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 = llama_get_overflow_pattern(il, il == il0 ? ngl_per_device[id].overflow_type : LLAMA_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; }; @@ -464,7 +444,7 @@ static void llama_params_fit_impl( set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, mparams_copy); const dmds_t dmd_nl = llama_get_device_memory_data( - path_model, &mparams_copy, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + path_model, &mparams_copy, cparams, devs, hp_ngl, hp_nct, hp_nex, hp_nr, log_level); LLAMA_LOG_DEBUG("%s: memory for test allocation by device:\n", func_name); for (size_t id = 0; id < nd; id++) { @@ -486,13 +466,13 @@ static void llama_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; LLAMA_LOG_DEBUG("%s: getting device memory data with all MoE tensors moved to system memory:\n", __func__); const dmds_t dmds_cpu_moe = llama_get_device_memory_data( - path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, hp_nr, log_level); for (size_t id = 0; id < nd; id++) { global_surplus_cpu_moe += dmds_cpu_moe[id].free; @@ -508,7 +488,7 @@ static void llama_params_fit_impl( } // reset - tensor_buft_overrides[0] = {nullptr, nullptr}; + tensor_buft_overrides[0] = {nullptr, nullptr, -1}; mparams->tensor_buft_overrides = tensor_buft_overrides; } @@ -686,12 +666,12 @@ static void llama_params_fit_impl( if (ngl_per_device_test[id_dense_start_test].n_part == 0) { id_dense_start_test++; } - ngl_per_device_test[id].overflow_type = LAYER_FRACTION_UP; + ngl_per_device_test[id].overflow_type = LLAMA_LAYER_FRACTION_UP; std::vector overflow_bufts_test = overflow_bufts; if (id < nd - 1) { overflow_bufts_test[id] = ggml_backend_dev_buffer_type(devs[id + 1].dev); } - LLAMA_LOG_DEBUG("%s: trying to fit one extra layer with overflow_type=LAYER_FRACTION_UP\n", __func__); + LLAMA_LOG_DEBUG("%s: trying to fit one extra layer with overflow_type=LLAMA_LAYER_FRACTION_UP\n", __func__); std::vector mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts_test); if (mem_test[id] < targets[id] && (id + 1 == nd || mem_test[id + 1] < targets[id + 1])) { ngl_per_device = ngl_per_device_test; @@ -701,8 +681,8 @@ static void llama_params_fit_impl( LLAMA_LOG_DEBUG("%s: set ngl_per_device[%zu].(n_layer, n_part, overflow_type)=(%" PRIu32 ", %" PRIu32 ", UP), id_dense_start=%zu\n", __func__, id, ngl_per_device[id].n_layer, ngl_per_device[id].n_part, id_dense_start); - ngl_per_device_test[id].overflow_type = LAYER_FRACTION_GATE; - LLAMA_LOG_DEBUG("%s: trying to fit one extra layer with overflow_type=LAYER_FRACTION_GATE\n", __func__); + ngl_per_device_test[id].overflow_type = LLAMA_LAYER_FRACTION_GATE; + LLAMA_LOG_DEBUG("%s: trying to fit one extra layer with overflow_type=LLAMA_LAYER_FRACTION_GATE\n", __func__); mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts_test); if (mem_test[id] < targets[id] && (id + 1 == nd || mem_test[id + 1] < targets[id + 1])) { ngl_per_device = ngl_per_device_test; @@ -713,8 +693,8 @@ static void llama_params_fit_impl( __func__, id, ngl_per_device[id].n_layer, ngl_per_device[id].n_part, id_dense_start); } } else { - ngl_per_device_test[id].overflow_type = LAYER_FRACTION_ATTN; - LLAMA_LOG_DEBUG("%s: trying to fit one extra layer with overflow_type=LAYER_FRACTION_ATTN\n", __func__); + ngl_per_device_test[id].overflow_type = LLAMA_LAYER_FRACTION_ATTN; + LLAMA_LOG_DEBUG("%s: trying to fit one extra layer with overflow_type=LLAMA_LAYER_FRACTION_ATTN\n", __func__); mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts_test); if (mem_test[id] < targets[id] && (id + 1 == nd || mem_test[id + 1] < targets[id + 1])) { ngl_per_device = ngl_per_device_test; @@ -876,6 +856,8 @@ static int llama_model_load(struct gguf_context * metadata, llama_model_set_tens return 0; } + ml.force_duplicate_tied = params.pshard; + if (!model.load_tensors(ml)) { return -2; } @@ -1274,4 +1256,3 @@ const char * llama_print_system_info(void) { return s.c_str(); } - diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index b433c91d85eb..4b7287e2e37d 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -38,5 +38,6 @@ else() add_subdirectory(export-lora) endif() add_subdirectory(fit-params) + add_subdirectory(pshard-plan-params) add_subdirectory(results) endif() diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index b15a26a987b3..386a8a7a52da 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -382,7 +382,7 @@ static const cmd_params cmd_params_defaults = { /* flash_attn */ { false }, /* devices */ { {} }, /* tensor_split */ { std::vector(llama_max_devices(), 0.0f) }, - /* tensor_buft_overrides*/ { std::vector{ { nullptr, nullptr } } }, + /* tensor_buft_overrides*/ { std::vector{ { nullptr, nullptr, -1 } } }, /* use_mmap */ { true }, /* use_direct_io */ { false }, /* embeddings */ { false }, @@ -922,13 +922,13 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { invalid_param = true; break; } - group_tensor_buft_overrides.push_back({tensor_name, buft_list.at(buffer_type)}); + group_tensor_buft_overrides.push_back({tensor_name, buft_list.at(buffer_type), -1}); override_span_len = std::strcspn(override_group, ";"); } if (invalid_param) { break; } - group_tensor_buft_overrides.push_back({nullptr,nullptr}); + group_tensor_buft_overrides.push_back({nullptr, nullptr, -1}); params.tensor_buft_overrides.push_back(group_tensor_buft_overrides); override_group_span_len = std::strcspn(value, ","); } while (!last_group); @@ -1193,7 +1193,7 @@ struct cmd_params_instance { ggml_backend_cpu_buffer_type() }); } - merged.push_back({ nullptr, nullptr }); + merged.push_back({ nullptr, nullptr, -1 }); mparams.tensor_buft_overrides = merged.data(); } diff --git a/tools/pshard-plan-params/CMakeLists.txt b/tools/pshard-plan-params/CMakeLists.txt new file mode 100644 index 000000000000..a788671cfa74 --- /dev/null +++ b/tools/pshard-plan-params/CMakeLists.txt @@ -0,0 +1,8 @@ +set(TARGET llama-pshard-plan-params) +add_executable(${TARGET} pshard-plan-params.cpp) +target_link_libraries(${TARGET} PRIVATE common llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) + +if(LLAMA_TOOLS_INSTALL) + install(TARGETS ${TARGET} RUNTIME) +endif() diff --git a/tools/pshard-plan-params/README.md b/tools/pshard-plan-params/README.md new file mode 100644 index 000000000000..171b060ba8cf --- /dev/null +++ b/tools/pshard-plan-params/README.md @@ -0,0 +1,124 @@ +# pshard-plan-params + +`llama-pshard-plan-params` writes a pshard registry for a model and VRAM budget. + +It probes batch-size tiers such as `bs=1`, `bs=16`, and `bs=512`, compares the available placement strategies, and writes the selected tensor overrides to: + +```text +.gguf.tensor_overrides.pshard_registry +``` + +When CPU/GPU profiler results are available, the planner uses them to estimate throughput for each candidate plan. At runtime, pshard selects the smallest tier that covers the current batch. + +The registry is plain text so it can be inspected when debugging planner output. + +## Example usage + +```bash +# 1. Generate roofline profiles once per machine +# (writes profile JSONs next to the binary) +./build/bin/llama-profiler-cpu +./build/bin/llama-profiler-gpu + +# 2. Plan placement for a model + VRAM budget +./build/bin/llama-pshard-plan-params \ + --model /opt/models/Qwen3.5-27B-Q4_0.gguf \ + --max-vram-alloc 12000 \ + -c 8192 \ + -fa on + +ggml_cuda_init: found 1 CUDA devices (Total VRAM: 16302 MiB): + Device 0: NVIDIA GeForce RTX 5070 Ti, compute capability 12.0, VMM: yes, VRAM: 16302 MiB +main: planning pshard tensor overrides... +llama_params_fit_pshard: probing pshard plans: 64 layers, 12000.0 MiB VRAM free +pshard_registry_save: saved budget=12000 MiB cache_ubatch=8192 variant with 7 tier plans to /opt/models/Qwen3.5-27B-Q4_0.gguf.tensor_overrides.pshard_registry +main: planning complete, registry written next to model file + +# 3. Inspect the registry +cat /opt/models/Qwen3.5-27B-Q4_0.gguf.tensor_overrides.pshard_registry +``` + +Example registry output: + +```text +# Generated file. Edit at your own risk. + +[fingerprint=0x9041bb5b253cf89f] +# n_ctx=8192 n_seq_max=1 n_threads=8 fa=on type_k=1 type_v=1 strategy=auto + +[variant budget=12000 cache_ubatch=8192] +[tier 0 bs=1] +strategy=DYNAMIC_FFNCPU_ATTNSTREAM n_pinned=54 n_attn_pinned=0 overflow=NONE tps=12.85 vram=11975.7 output_on_gpu=0 pin_from_back=0 +ot=^output=CUDA_Host:3,^token_embd=CUDA_Host:3,blk\.0\..*=CUDA_Host:0, ... ,blk\.53\..*=CUDA_Host:0,blk\.54\.ffn_(up|gate|down).*=CUDA_Host:3,blk\.54\..*=CUDA_Host:1,blk\.55\.ffn_(up|gate|down).*=CUDA_Host:3,blk\.55\..*=CUDA_Host:2, ... ,blk\.63\.ffn_(up|gate|down).*=CUDA_Host:3,blk\.63\..*=CUDA_Host:2 +[tier 1 bs=16] +strategy=GPUONLY_ATTNPIN_FFNSTREAM n_pinned=48 n_attn_pinned=0 overflow=NONE tps=208.27 vram=11964.1 output_on_gpu=0 pin_from_back=0 +ot=... +[tier 2 bs=512] +strategy=GPUONLY_LAYERPIN_LAYERSTREAM n_pinned=52 n_attn_pinned=0 overflow=NONE tps=1994.90 vram=11973.3 output_on_gpu=0 pin_from_back=0 +ot=... +[tier 3 bs=1024] +strategy=GPUONLY_LAYERPIN_LAYERSTREAM n_pinned=51 n_attn_pinned=0 overflow=NONE tps=2286.87 vram=11909.8 output_on_gpu=0 pin_from_back=0 +ot=... +[tier 4 bs=2048] +strategy=GPUONLY_ATTNPIN_FFNSTREAM n_pinned=44 n_attn_pinned=0 overflow=NONE tps=2249.44 vram=11914.3 output_on_gpu=0 pin_from_back=0 +ot=... +[tier 5 bs=4096] +strategy=GPUONLY_ATTNPIN_FFNSTREAM n_pinned=40 n_attn_pinned=0 overflow=NONE tps=2095.88 vram=11868.6 output_on_gpu=0 pin_from_back=0 +ot=... +[tier 6 bs=8192] +strategy=GPUONLY_ATTNPIN_FFNSTREAM n_pinned=33 n_attn_pinned=0 overflow=NONE tps=2148.21 vram=11937.2 output_on_gpu=0 pin_from_back=0 +ot=^output=CUDA_Host:3,^token_embd=CUDA_Host:3,blk\.0\..*=CUDA_Host:0, ... ,blk\.32\..*=CUDA_Host:0,blk\.33\.ffn_(up|gate|down).*=CUDA_Host:2,blk\.33\..*=CUDA_Host:0, ... ,blk\.63\.ffn_(up|gate|down).*=CUDA_Host:2,blk\.63\..*=CUDA_Host:0 +``` + +## Planning for llama-bench + +Use `--bench-plan` when the registry is intended for `llama-bench -pshard`. The planner accepts the bench shape flags `-p` / `--n-prompt`, `-n` / `--n-gen`, `-pg`, and `-d` / `--n-depth`, then plans each unique context that llama-bench can run: + +```bash +./build/bin/llama-pshard-plan-params \ + --model /opt/models/Qwen3.5-35B-A3B-Q8_0.gguf \ + --bench-plan \ + -pg 512,200 \ + -pg 2048,200 \ + -d 0,1024 +``` + +For each generated context, `n_ctx` is the number of tokens that can be resident during the test: + +- prompt-only: `n_ctx = p + d` +- generation-only: `n_ctx = n + d` +- prompt+generation: `n_ctx = p + n + d` + +The largest planned tier for a bench context is capped by prompt batch demand, not by depth. For example, `-pg 2048,200 -d 1024` plans `n_ctx=3272` with `tier_cap=2048`, so the registry includes decode tiers plus prompt tiers up to `bs=2048`. + +If `-fa` / `--flash-attn` is not provided and `LLAMA_ARG_FLASH_ATTN` is not set, `--bench-plan` uses Flash Attention off to match llama-bench defaults. Normal non-bench planning keeps the regular common parameter defaults. + +## Strategies + +In the strategy names, `ATTN` refers to the attention/dense side of the layer, as opposed to FFN/MoE weights. + +Static schedules run GPU-resident tensors on GPU and CPU-resident tensors on CPU, with no streamed GPU execution for host-resident weights. + +- `STATIC_ATTNPRIO_ALLMODELS`: static attention-priority placement. It pins the attention/dense side across as many layers as fit, then uses the remaining budget to pin full layers. FFN/MoE that does not fit remains on CPU. Unlike `llama_params_fit`, this attention-priority placement applies to dense models too. + +Dynamic schedules split the layer between CPU and GPU execution. Some host-resident tensors are streamed to GPU scratch for execution, while other parts of the layer remain on CPU. + +- `DYNAMIC_FFNCPU_ATTNSTREAM`: pin as many full layers as fit, keep FFN/MoE on CPU in the remaining layers, and stream the attention/dense side for GPU execution. + +GPU-only schedules execute repeating-layer compute on GPU. Weights that do not fit in VRAM stay resident in host memory and are streamed to GPU scratch before use. + +- `GPUONLY_LAYERPIN_LAYERSTREAM`: pin as many full layers as fit, then stream the remaining layers for GPU execution. +- `GPUONLY_ATTNPIN_FFNSTREAM`: pin the attention/dense side for all layers, pin as many complete layers as the remaining budget allows, and stream FFN/MoE weights for GPU execution. + +## Notes + +- `--max-vram-alloc` / `-mva ` sets the absolute planning budget. If it is `0` or omitted, the planner uses the free VRAM available at planning time minus `--fit-target` / `-fitt` (default 1024 MiB). When `-mva` is non-zero, `-fitt` is ignored for pshard. +- `--pshard-tier-max ` caps the largest batch size the planner probes. The default is `min(max(n_batch, 16384), n_ctx)`. +- The `[fingerprint=...]` line invalidates the registry when plan-compatible inputs change (`n_ctx`, `n_seq_max`, threads, FA mode, KV cache types, GGUF file size, or forced `PSHARD_STRATEGY`). The comment below the fingerprint lists those same inputs. +- A fingerprint can contain multiple `[variant budget=... cache_ubatch=...]` blocks. Re-running the planner with a new budget or cache ubatch replaces only that variant and keeps the others. +- `cache_ubatch` records the runtime ubatch used for context/KV/SWA cache sizing. Tier compute scratch is still measured with the tier batch size. +- `pshard_disabled=1 baseline_vram=` is variant-scoped. Runtime skips pshard only when the measured baseline VRAM fits the current budget. +- `backend_id` values: + - `0` = GPU pinned compute + - `1` / `2` = shard compute lanes used for pipeline overlap + - `3` = CPU diff --git a/tools/pshard-plan-params/pshard-plan-params.cpp b/tools/pshard-plan-params/pshard-plan-params.cpp new file mode 100644 index 000000000000..7862500ebb53 --- /dev/null +++ b/tools/pshard-plan-params/pshard-plan-params.cpp @@ -0,0 +1,267 @@ +#include "llama.h" + +#include "arg.h" +#include "common.h" +#include "log.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_MSC_VER) +#pragma warning(disable: 4244 4267) // possible loss of data +#endif + +struct bench_plan_params { + bool enabled = false; + + std::vector n_prompt; + std::vector n_gen; + std::vector> n_pg; + std::vector n_depth; + bool flash_attn_seen = false; +}; + +static std::vector parse_int_range(const std::string & s) { + // first[-last[(+|*)step]] + std::regex range_regex(R"(^(\d+)(?:-(\d+)(?:([\+|\*])(\d+))?)?(?:,|$))"); + + std::smatch match; + std::string::const_iterator search_start(s.cbegin()); + + std::vector result; + while (std::regex_search(search_start, s.cend(), match, range_regex)) { + int first = std::stoi(match[1]); + int last = match[2].matched ? std::stoi(match[2]) : first; + char op = match[3].matched ? match[3].str()[0] : '+'; + int step = match[4].matched ? std::stoi(match[4]) : 1; + + for (int i = first; i <= last;) { + result.push_back(i); + + int prev_i = i; + if (op == '+') { + i += step; + } else if (op == '*') { + i *= step; + } else { + throw std::invalid_argument("invalid range format"); + } + + if (i <= prev_i) { + throw std::invalid_argument("invalid range"); + } + } + + search_start = match.suffix().first; + } + + if (search_start != s.cend()) { + throw std::invalid_argument("invalid range format"); + } + + return result; +} + +static std::pair parse_pg(const std::string & s) { + const size_t comma = s.find(','); + if (comma == std::string::npos || s.find(',', comma + 1) != std::string::npos) { + throw std::invalid_argument("invalid -pg format"); + } + + return { std::stoi(s.substr(0, comma)), std::stoi(s.substr(comma + 1)) }; +} + +static bool is_bench_plan_arg_with_value(const std::string & arg) { + return arg == "-p" || arg == "--n-prompt" || + arg == "-n" || arg == "--n-gen" || + arg == "-pg" || + arg == "-d" || arg == "--n-depth"; +} + +static bool parse_bench_plan_args(int argc, char ** argv, bench_plan_params & bench, std::vector & forwarded) { + for (int i = 1; i < argc; i++) { + if (std::string(argv[i]) == "--bench-plan") { + bench.enabled = true; + break; + } + } + + forwarded.clear(); + forwarded.push_back(argv[0]); + + for (int i = 1; i < argc; i++) { + const std::string arg = argv[i]; + + if (arg == "--bench-plan") { + continue; + } + + if (bench.enabled && (arg == "-fa" || arg == "--flash-attn")) { + bench.flash_attn_seen = true; + } + + if (!bench.enabled || !is_bench_plan_arg_with_value(arg)) { + forwarded.push_back(arg); + continue; + } + + if (++i >= argc) { + fprintf(stderr, "%s: missing value for %s\n", __func__, arg.c_str()); + return false; + } + + try { + if (arg == "-p" || arg == "--n-prompt") { + auto vals = parse_int_range(argv[i]); + bench.n_prompt.insert(bench.n_prompt.end(), vals.begin(), vals.end()); + } else if (arg == "-n" || arg == "--n-gen") { + auto vals = parse_int_range(argv[i]); + bench.n_gen.insert(bench.n_gen.end(), vals.begin(), vals.end()); + } else if (arg == "-pg") { + bench.n_pg.push_back(parse_pg(argv[i])); + } else if (arg == "-d" || arg == "--n-depth") { + auto vals = parse_int_range(argv[i]); + bench.n_depth.insert(bench.n_depth.end(), vals.begin(), vals.end()); + } + } catch (const std::exception & e) { + fprintf(stderr, "%s: invalid value for %s: %s (%s)\n", __func__, arg.c_str(), argv[i], e.what()); + return false; + } + } + + if (bench.enabled) { + if (bench.n_prompt.empty()) { + bench.n_prompt.push_back(512); + } + if (bench.n_gen.empty()) { + bench.n_gen.push_back(128); + } + if (bench.n_depth.empty()) { + bench.n_depth.push_back(0); + } + } + + return true; +} + +static void bench_plan_add_case(std::map & result, uint32_t n_ctx, uint32_t tier_cap) { + if (n_ctx == 0) { + return; + } + + auto & max_tier_cap = result[n_ctx]; + max_tier_cap = std::max(max_tier_cap, tier_cap); +} + +static std::map bench_plan_context_tier_caps(const bench_plan_params & bench) { + std::map result; + + for (const int d : bench.n_depth) { + for (const int p : bench.n_prompt) { + if (p > 0) { + bench_plan_add_case(result, (uint32_t) (p + d), (uint32_t) p); + } + } + + for (const int n : bench.n_gen) { + if (n > 0) { + bench_plan_add_case(result, (uint32_t) (n + d), 0); + } + } + + for (const auto & pg : bench.n_pg) { + if (pg.first > 0 || pg.second > 0) { + bench_plan_add_case(result, (uint32_t) (pg.first + pg.second + d), (uint32_t) std::max(pg.first, 0)); + } + } + } + + return result; +} + +static void plan_pshard_context(common_params & params, uint32_t n_ctx, uint32_t bench_tier_cap = 0, bool bench_plan = false) { + auto mparams = common_model_params_to_llama(params); + auto cparams = common_context_params_to_llama(params); + + cparams.n_ctx = n_ctx; + + params.tensor_buft_overrides.assign(4096, {}); + + const uint32_t tier_max_auto = bench_plan ? + std::max(bench_tier_cap, (uint32_t) 16) : + std::min(std::max(cparams.n_batch, (uint32_t) 16384), cparams.n_ctx); + const uint32_t tier_max_user = params.pshard_tier_max > 0 ? std::min(params.pshard_tier_max, tier_max_auto) : tier_max_auto; + const uint32_t tier_max = bench_plan ? tier_max_user : std::min(tier_max_user, cparams.n_ctx); + + mparams.pshard_registry = llama_pshard_registry_create(tier_max, cparams.n_seq_max); + mparams.pshard_cache_skip_load = true; + + const size_t fit_target_mb = params.fit_params_target.empty() ? 0 : params.fit_params_target[0] / (1024 * 1024); + + if (bench_plan) { + LOG_INF("%s: planning pshard tensor overrides for n_ctx=%u tier_cap=%u tier_max=%u...\n", + __func__, n_ctx, bench_tier_cap, tier_max); + } else { + LOG_INF("%s: planning pshard tensor overrides for n_ctx=%u tier_max=%u...\n", __func__, n_ctx, tier_max); + } + llama_params_fit_pshard(params.model.path.c_str(), &mparams, &cparams, + params.tensor_buft_overrides.data(), params.max_vram_alloc, fit_target_mb); + + llama_pshard_registry_free(mparams.pshard_registry); +} + +int main(int argc, char ** argv) { + common_params params; + bench_plan_params bench; + std::vector forwarded; + + if (!parse_bench_plan_args(argc, argv, bench, forwarded)) { + return 1; + } + + std::vector forwarded_argv; + forwarded_argv.reserve(forwarded.size()); + for (std::string & arg : forwarded) { + forwarded_argv.push_back(arg.data()); + } + + if (!common_params_parse((int) forwarded_argv.size(), forwarded_argv.data(), params, LLAMA_EXAMPLE_COMMON)) { + return 1; + } + + if (bench.enabled && !bench.flash_attn_seen && getenv("LLAMA_ARG_FLASH_ATTN") == nullptr) { + params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; + } + + common_init(); + llama_backend_init(); + llama_numa_init(params.numa); + + if (bench.enabled) { + const std::map ctx_caps = bench_plan_context_tier_caps(bench); + if (ctx_caps.empty()) { + fprintf(stderr, "%s: --bench-plan produced no non-zero llama-bench contexts\n", __func__); + return 1; + } + + LOG_INF("%s: planning %zu unique llama-bench context(s)\n", __func__, ctx_caps.size()); + for (const auto & ctx_cap : ctx_caps) { + plan_pshard_context(params, ctx_cap.first, ctx_cap.second, true); + } + } else { + auto cparams = common_context_params_to_llama(params); + plan_pshard_context(params, cparams.n_ctx); + } + + LOG_INF("%s: planning complete, registry written next to model file\n", __func__); + + return 0; +}