From e14c6cfaa1fd35d491aa53fb7e95ec6cc9f75b0a Mon Sep 17 00:00:00 2001 From: Aditya Ukarande Date: Mon, 27 Apr 2026 01:17:37 -0700 Subject: [PATCH 1/2] feat: pipeshard plan switch/decode --- common/arg.cpp | 31 +- common/common.cpp | 22 +- common/common.h | 6 +- ggml/include/ggml-alloc.h | 15 + ggml/include/ggml-backend.h | 50 +++ ggml/include/ggml.h | 3 +- ggml/src/ggml-alloc.c | 266 +++++++++++- ggml/src/ggml-backend-impl.h | 3 + ggml/src/ggml-backend.cpp | 472 +++++++++++++++++++- ggml/src/ggml-cpu/ggml-cpu.cpp | 1 + ggml/src/ggml-cuda/ggml-cuda.cu | 8 + include/llama.h | 26 ++ src/CMakeLists.txt | 3 + src/llama-context-pshard.cpp | 646 +++++++++++++++++++++++++++ src/llama-context.cpp | 297 ++++++++++--- src/llama-context.h | 28 ++ src/llama-cparams.h | 27 ++ src/llama-graph.cpp | 22 +- src/llama-kv-cache-iswa.cpp | 7 + src/llama-kv-cache-iswa.h | 2 + src/llama-kv-cache.cpp | 121 ++++++ src/llama-kv-cache.h | 9 + src/llama-memory-hybrid-iswa.cpp | 7 + src/llama-memory-hybrid-iswa.h | 2 + src/llama-memory-hybrid.cpp | 7 + src/llama-memory-hybrid.h | 2 + src/llama-memory-pshard.cpp | 585 +++++++++++++++++++++++++ src/llama-memory-pshard.h | 119 +++++ src/llama-memory-recurrent.cpp | 60 +++ src/llama-memory-recurrent.h | 5 + src/llama-memory.h | 3 + src/llama-model-loader.cpp | 119 ++++- src/llama-model-loader.h | 16 + src/llama-model.cpp | 330 ++++++++++++++ src/llama-model.h | 18 + src/llama-pipe-shard.h | 69 +++ src/llama-pshard-cache.cpp | 701 ++++++++++++++++++++++++++++++ src/llama-pshard-plan.h | 233 ++++++++++ src/llama.cpp | 32 +- tools/llama-bench/llama-bench.cpp | 8 +- 40 files changed, 4249 insertions(+), 132 deletions(-) create mode 100644 src/llama-context-pshard.cpp create mode 100644 src/llama-memory-pshard.cpp create mode 100644 src/llama-memory-pshard.h create mode 100644 src/llama-pipe-shard.h create mode 100644 src/llama-pshard-cache.cpp create mode 100644 src/llama-pshard-plan.h diff --git a/common/arg.cpp b/common/arg.cpp index 3d0183ed7026..10d705ca9c22 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( + {"-pshard"}, + "enable pipelined sharding (weights on CPU, pipelined to GPU per split)", + [](common_params & params) { + params.pshard = true; + } + )); + add_opt(common_arg( + {"-mva", "--max-vram-alloc"}, "N", + "VRAM budget in MB for pshard (0 = use actual free VRAM)", + [](common_params & params, int value) { + params.max_vram_alloc = value; + } + )); add_opt(common_arg( {"--spm-infill"}, string_format( @@ -2306,7 +2327,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex // keep strings alive and avoid leaking memory by storing them in a static vector static std::list 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")); diff --git a/common/common.cpp b/common/common.cpp index 16f78debd025..8de41eb963e4 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1144,7 +1144,24 @@ common_init_result::common_init_result(common_params & params) : auto mparams = common_model_params_to_llama(params); auto cparams = common_context_params_to_llama(params); - if (params.fit_params) { + if (params.pshard) { + LOG_INF("%s: pshard enabled, probing and loading plan cache\n", __func__); + params.tensor_buft_overrides.resize(4096); + const uint32_t tier_max = std::min(std::max(cparams.n_batch, (uint32_t)16384), cparams.n_ctx); + mparams.pshard_registry = llama_pshard_registry_create(tier_max, cparams.n_seq_max); + llama_params_fit_pshard(params.model.path.c_str(), &mparams, &cparams, + params.tensor_buft_overrides.data(), params.max_vram_alloc); + if (!mparams.pshard) { + LOG_WRN("%s: pshard not active for this configuration\n", __func__); + llama_pshard_registry_free(mparams.pshard_registry); + mparams.pshard_registry = nullptr; + } else { + params.n_batch = (int32_t) cparams.n_batch; + params.n_ubatch = (int32_t) cparams.n_ubatch; + LOG_INF("%s: pshard runtime batch/ubatch set to selected cache_ubatch=%u\n", + __func__, cparams.n_ubatch); + } + } else if (params.fit_params) { LOG_INF("%s: fitting params to device memory, for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on\n", __func__); llama_params_fit(params.model.path.c_str(), &mparams, &cparams, params.tensor_split, @@ -1444,6 +1461,8 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.progress_callback_user_data = params.load_progress_callback_user_data; mparams.no_alloc = params.no_alloc; + mparams.pshard = params.pshard; + mparams.max_vram_alloc = params.max_vram_alloc; return mparams; } @@ -1479,6 +1498,7 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.type_k = params.cache_type_k; cparams.type_v = params.cache_type_v; + cparams.pshard = params.pshard; return cparams; } diff --git a/common/common.h b/common/common.h index 020b6a721ff9..1557f227faf8 100644 --- a/common/common.h +++ b/common/common.h @@ -681,6 +681,10 @@ struct common_params { llama_progress_callback load_progress_callback = NULL; void * load_progress_callback_user_data = NULL; bool no_alloc = false; // Don't allocate model buffers + + bool pshard = false; + size_t max_vram_alloc = 0; + uint32_t pshard_tier_max = 0; }; // call once at the start of a program if it uses libcommon @@ -987,7 +991,7 @@ inline std::string llm_ffn_exps_block_regex(int idx) { } inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() { - return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() }; + return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type(), -1 }; } // diff --git a/ggml/include/ggml-alloc.h b/ggml/include/ggml-alloc.h index 78aa059dde38..5cb34cfb3559 100644 --- a/ggml/include/ggml-alloc.h +++ b/ggml/include/ggml-alloc.h @@ -73,6 +73,21 @@ GGML_API bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph GGML_API size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id); +// per-chunk introspection (after reserve_n / reserve_n_size) +GGML_API int ggml_gallocr_get_n_chunks(ggml_gallocr_t galloc, int buffer_id); +GGML_API size_t ggml_gallocr_get_chunk_max_size(ggml_gallocr_t galloc, int buffer_id, int chunk_id); + +// set an externally-owned buffer for a buffer slot (must be called before reserve) +GGML_API void ggml_gallocr_set_buffer(ggml_gallocr_t galloc, int buffer_id, ggml_backend_buffer_t buffer, size_t alloc_offset, size_t alloc_size); + +// update the allocation range for an external buffer (e.g. on plan switch) +GGML_API void ggml_gallocr_set_alloc_range(ggml_gallocr_t galloc, int buffer_id, size_t alloc_offset, size_t alloc_size); + +// save/restore allocator state (for plan switch without re-reserve) +GGML_API void ggml_gallocr_get_state_sizes(ggml_gallocr_t galloc, size_t * node_size, size_t * leaf_size); +GGML_API void ggml_gallocr_save_state(ggml_gallocr_t galloc, void * node_buf, void * leaf_buf, int * n_nodes, int * n_leafs); +GGML_API void ggml_gallocr_restore_state(ggml_gallocr_t galloc, const void * node_buf, size_t node_size, const void * leaf_buf, size_t leaf_size, int n_nodes, int n_leafs); + // Utils // Create a buffer and allocate all the tensors in a ggml_context // ggml_backend_alloc_ctx_tensors_from_buft_size returns the size of the buffer that would be allocated by ggml_backend_alloc_ctx_tensors_from_buft diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 4a8f6d4287da..6e0e0f875ab5 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -94,6 +94,7 @@ extern "C" { GGML_API void ggml_backend_tensor_set_2d( struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data); GGML_API void ggml_backend_tensor_get_2d(const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data); GGML_API void ggml_backend_tensor_memset( struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size); + GGML_API void ggml_backend_tensor_memset_async(ggml_backend_t backend, struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size); GGML_API void ggml_backend_synchronize(ggml_backend_t backend); @@ -154,6 +155,8 @@ extern "C" { bool buffer_from_host_ptr; // event synchronization bool events; + // separate copy stream for compute/transfer overlap + bool copy_stream; }; // all the device properties @@ -348,6 +351,53 @@ extern "C" { // Set a callback to be called for each resulting node during graph compute GGML_API void ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data); + // set an externally-owned buffer for a backend (see ggml_gallocr_set_buffer) + GGML_API void ggml_backend_sched_set_buffer(ggml_backend_sched_t sched, ggml_backend_t backend, ggml_backend_buffer_t buffer, size_t alloc_offset, size_t alloc_size); + + // update the allocation range for an external buffer (see ggml_gallocr_set_alloc_range) + GGML_API void ggml_backend_sched_set_alloc_range(ggml_backend_sched_t sched, ggml_backend_t backend, size_t alloc_offset, size_t alloc_size); + + // per-chunk introspection (see ggml_gallocr_get_n_chunks / get_chunk_max_size) + GGML_API int ggml_backend_sched_get_n_chunks(ggml_backend_sched_t sched, ggml_backend_t backend); + GGML_API size_t ggml_backend_sched_get_chunk_max_size(ggml_backend_sched_t sched, ggml_backend_t backend, int chunk_id); + + // save/restore gallocr + backend_id state (for plan switch without re-reserve) + GGML_API ggml_gallocr_t ggml_backend_sched_get_galloc(ggml_backend_sched_t sched); + GGML_API void ggml_backend_sched_save_backend_ids(ggml_backend_sched_t sched, int * node_buf, int * leaf_buf, int * n_nodes, int * n_leafs); + GGML_API void ggml_backend_sched_restore_backend_ids(ggml_backend_sched_t sched, const int * node_buf, int n_nodes, const int * leaf_buf, int n_leafs); + + // Enable async weight prefetching to overlap CPU->GPU transfers with compute + GGML_API void ggml_backend_sched_set_prefetch_weights(ggml_backend_sched_t sched, bool enabled); + + // Per-split callbacks for stateful tensors (e.g. KV cache, recurrent state). + typedef void (*ggml_backend_sched_split_cb)(struct ggml_tensor * tensor, ggml_backend_t backend, void * user_data); + + GGML_API void ggml_backend_sched_set_split_callbacks( + ggml_backend_sched_t sched, + ggml_backend_sched_split_cb pre_compute, + ggml_backend_sched_split_cb post_compute, + void * user_data); + + GGML_API void ggml_backend_sched_set_prefetch_cb( + ggml_backend_sched_t sched, + ggml_backend_sched_split_cb prefetch_cb); + + // Register a tensor for pre/post-compute split callbacks. + GGML_API void ggml_backend_sched_add_writeback(ggml_backend_sched_t sched, struct ggml_tensor * tensor); + + // Per-split info snapshot for timing prediction. + struct ggml_backend_sched_split_info { + struct ggml_cgraph * graph; + int backend_id; + size_t input_weight_bytes; + size_t input_activ_bytes; + size_t writeback_bytes; + }; + + GGML_API bool ggml_backend_sched_get_split_info( + ggml_backend_sched_t sched, int split_id, + struct ggml_backend_sched_split_info * out); + // // Meta backend // 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..353098d8914f 100644 --- a/ggml/src/ggml-alloc.c +++ b/ggml/src/ggml-alloc.c @@ -362,6 +362,16 @@ static void ggml_dyn_tallocr_reset(struct ggml_dyn_tallocr * alloc) { #endif } +static void ggml_dyn_tallocr_reset_with_range(struct ggml_dyn_tallocr * alloc, size_t offset, size_t size) { + ggml_dyn_tallocr_reset(alloc); + ggml_dyn_tallocr_new_chunk(alloc, 0); + struct tallocr_chunk * c0 = alloc->chunks[0]; + c0->n_free_blocks = 1; + c0->free_blocks[0].offset = offset; + c0->free_blocks[0].size = size; + c0->max_size = 0; +} + static struct ggml_dyn_tallocr * ggml_dyn_tallocr_new(size_t alignment, size_t max_buffer_size) { struct ggml_dyn_tallocr * alloc = (struct ggml_dyn_tallocr *)malloc(sizeof(struct ggml_dyn_tallocr)); @@ -484,6 +494,10 @@ struct ggml_gallocr { struct ggml_dyn_tallocr ** buf_tallocs; // [n_buffers] int n_buffers; + bool * buf_external; // [n_buffers] + size_t * buf_alloc_offset; // [n_buffers] + size_t * buf_alloc_size; // [n_buffers] + struct ggml_hash_set hash_set; struct hash_node * hash_values; // [hash_set.size] @@ -507,6 +521,15 @@ ggml_gallocr_t ggml_gallocr_new_n(ggml_backend_buffer_type_t * bufts, int n_bufs galloc->buf_tallocs = calloc(n_bufs, sizeof(struct ggml_dyn_tallocr *)); GGML_ASSERT(galloc->buf_tallocs != NULL); + galloc->buf_external = calloc(n_bufs, sizeof(bool)); + GGML_ASSERT(galloc->buf_external != NULL); + + galloc->buf_alloc_offset = calloc(n_bufs, sizeof(size_t)); + GGML_ASSERT(galloc->buf_alloc_offset != NULL); + + galloc->buf_alloc_size = calloc(n_bufs, sizeof(size_t)); + GGML_ASSERT(galloc->buf_alloc_size != NULL); + for (int i = 0; i < n_bufs; i++) { galloc->bufts[i] = bufts[i]; galloc->buffers[i] = NULL; @@ -550,6 +573,9 @@ void ggml_gallocr_free(ggml_gallocr_t galloc) { } } if (!freed) { + if (galloc->buf_external[i]) { + galloc->buffers[i]->chunks[0] = NULL; + } ggml_vbuffer_free(galloc->buffers[i]); } } @@ -573,6 +599,9 @@ void ggml_gallocr_free(ggml_gallocr_t galloc) { free(galloc->bufts); free(galloc->buffers); free(galloc->buf_tallocs); + free(galloc->buf_external); + free(galloc->buf_alloc_offset); + free(galloc->buf_alloc_size); free(galloc->node_allocs); free(galloc->leaf_allocs); free(galloc); @@ -623,6 +652,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 +759,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 +774,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 +799,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 +861,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) { @@ -843,6 +903,24 @@ static bool ggml_gallocr_reserve_n_impl( ggml_dyn_tallocr_reset(galloc->buf_tallocs[i]); } + // constrain external buffers to their allocation range + for (int i = 0; i < galloc->n_buffers; i++) { + if (!galloc->buf_external[i]) { + continue; + } + bool already_seeded = false; + for (int j = 0; j < i; j++) { + if (galloc->buf_tallocs[j] == galloc->buf_tallocs[i]) { + already_seeded = true; + break; + } + } + if (already_seeded) { + continue; + } + ggml_dyn_tallocr_reset_with_range(galloc->buf_tallocs[i], galloc->buf_alloc_offset[i], galloc->buf_alloc_size[i]); + } + // allocate in hash table ggml_gallocr_alloc_graph_impl(galloc, graph, node_buffer_ids, leaf_buffer_ids); @@ -922,23 +1000,49 @@ static bool ggml_gallocr_reserve_n_impl( } } if (realloc) { + if (galloc->buf_external[i]) { + struct ggml_dyn_tallocr * talloc = galloc->buf_tallocs[i]; + if (!no_alloc) { + for (int c = 1; c < talloc->n_chunks; c++) { + size_t chunk_size = talloc->chunks[c]->max_size; + size_t cur_size = ggml_vbuffer_chunk_size(galloc->buffers[i], c); + if (chunk_size > cur_size) { + ggml_backend_buffer_free(galloc->buffers[i]->chunks[c]); + galloc->buffers[i]->chunks[c] = ggml_backend_buft_alloc_buffer(galloc->bufts[i], chunk_size); + if (galloc->buffers[i]->chunks[c] == NULL) { + GGML_LOG_ERROR("%s: failed to allocate overflow chunk %d of size %zu\n", __func__, c, chunk_size); + return false; + } + ggml_backend_buffer_set_usage(galloc->buffers[i]->chunks[c], GGML_BACKEND_BUFFER_USAGE_COMPUTE); + } + } + } + for (int c = talloc->n_chunks; c < GGML_VBUFFER_MAX_CHUNKS; c++) { + if (galloc->buffers[i]->chunks[c] == NULL) { + break; + } + ggml_backend_buffer_free(galloc->buffers[i]->chunks[c]); + galloc->buffers[i]->chunks[c] = NULL; + } + } else { #ifndef NDEBUG - { - size_t cur_size = galloc->buffers[i] ? ggml_vbuffer_size(galloc->buffers[i]) : 0; - if (cur_size > 0) { - GGML_LOG_DEBUG("%s: reallocating %s buffer from size %.02f MiB to %.02f MiB\n", - __func__, ggml_backend_buft_name(galloc->bufts[i]), cur_size / 1024.0 / 1024.0, new_size / 1024.0 / 1024.0); + { + size_t cur_size = galloc->buffers[i] ? ggml_vbuffer_size(galloc->buffers[i]) : 0; + if (cur_size > 0) { + GGML_LOG_DEBUG("%s: reallocating %s buffer from size %.02f MiB to %.02f MiB\n", + __func__, ggml_backend_buft_name(galloc->bufts[i]), cur_size / 1024.0 / 1024.0, new_size / 1024.0 / 1024.0); + } } - } #endif - ggml_vbuffer_free(galloc->buffers[i]); - if (no_alloc) { - galloc->buffers[i] = NULL; - } else { - galloc->buffers[i] = ggml_vbuffer_alloc(galloc->bufts[i], galloc->buf_tallocs[i], GGML_BACKEND_BUFFER_USAGE_COMPUTE); - if (galloc->buffers[i] == NULL) { - GGML_LOG_ERROR("%s: failed to allocate %s buffer of size %zu\n", __func__, ggml_backend_buft_name(galloc->bufts[i]), new_size); - return false; + ggml_vbuffer_free(galloc->buffers[i]); + if (no_alloc) { + galloc->buffers[i] = NULL; + } else { + galloc->buffers[i] = ggml_vbuffer_alloc(galloc->bufts[i], galloc->buf_tallocs[i], GGML_BACKEND_BUFFER_USAGE_COMPUTE); + if (galloc->buffers[i] == NULL) { + GGML_LOG_ERROR("%s: failed to allocate %s buffer of size %zu\n", __func__, ggml_backend_buft_name(galloc->bufts[i]), new_size); + return false; + } } } } @@ -951,9 +1055,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; + } } } } @@ -1114,6 +1228,112 @@ size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id) { return ggml_vbuffer_size(galloc->buffers[buffer_id]); } +int ggml_gallocr_get_n_chunks(ggml_gallocr_t galloc, int buffer_id) { + GGML_ASSERT(buffer_id >= 0 && buffer_id < galloc->n_buffers); + return galloc->buf_tallocs[buffer_id]->n_chunks; +} + +size_t ggml_gallocr_get_chunk_max_size(ggml_gallocr_t galloc, int buffer_id, int chunk_id) { + GGML_ASSERT(buffer_id >= 0 && buffer_id < galloc->n_buffers); + GGML_ASSERT(chunk_id >= 0 && chunk_id < galloc->buf_tallocs[buffer_id]->n_chunks); + return galloc->buf_tallocs[buffer_id]->chunks[chunk_id]->max_size; +} + +void ggml_gallocr_set_buffer(ggml_gallocr_t galloc, int buffer_id, + ggml_backend_buffer_t buffer, size_t alloc_offset, size_t alloc_size) { + GGML_ASSERT(buffer_id >= 0 && buffer_id < galloc->n_buffers); + GGML_ASSERT(galloc->buffers[buffer_id] == NULL && "buffer slot must be empty, call set_buffer before reserve"); + + struct vbuffer * vbuf = (struct vbuffer *)calloc(1, sizeof(struct vbuffer)); + GGML_ASSERT(vbuf != NULL); + vbuf->chunks[0] = buffer; + + size_t range = alloc_size > 0 ? alloc_size : ggml_backend_buffer_get_size(buffer) - alloc_offset; + + galloc->buffers[buffer_id] = vbuf; + galloc->buf_external[buffer_id] = true; + galloc->buf_alloc_offset[buffer_id] = alloc_offset; + galloc->buf_alloc_size[buffer_id] = range; + + for (int i = 0; i < galloc->n_buffers; i++) { + if (i != buffer_id && galloc->buf_tallocs[i] == galloc->buf_tallocs[buffer_id]) { + galloc->buffers[i] = vbuf; + galloc->buf_external[i] = true; + galloc->buf_alloc_offset[i] = alloc_offset; + galloc->buf_alloc_size[i] = range; + } + } +} + +void ggml_gallocr_set_alloc_range(ggml_gallocr_t galloc, int buffer_id, size_t alloc_offset, size_t alloc_size) { + GGML_ASSERT(buffer_id >= 0 && buffer_id < galloc->n_buffers); + GGML_ASSERT(galloc->buf_external[buffer_id] && "set_alloc_range requires an external buffer"); + + galloc->buf_alloc_offset[buffer_id] = alloc_offset; + galloc->buf_alloc_size[buffer_id] = alloc_size; + + for (int i = 0; i < galloc->n_buffers; i++) { + if (i != buffer_id && galloc->buf_tallocs[i] == galloc->buf_tallocs[buffer_id]) { + galloc->buf_alloc_offset[i] = alloc_offset; + galloc->buf_alloc_size[i] = alloc_size; + } + } +} + +void ggml_gallocr_get_state_sizes(ggml_gallocr_t galloc, size_t * node_size, size_t * leaf_size) { + if (!galloc) { + if (node_size) *node_size = 0; + if (leaf_size) *leaf_size = 0; + return; + } + if (node_size) *node_size = galloc->n_nodes * sizeof(struct node_alloc); + if (leaf_size) *leaf_size = galloc->n_leafs * sizeof(struct leaf_alloc); +} + +void ggml_gallocr_save_state(ggml_gallocr_t galloc, + void * node_buf, void * leaf_buf, + int * n_nodes, int * n_leafs) { + if (!galloc) { + if (n_nodes) *n_nodes = 0; + if (n_leafs) *n_leafs = 0; + return; + } + if (n_nodes) *n_nodes = galloc->n_nodes; + if (n_leafs) *n_leafs = galloc->n_leafs; + if (node_buf && galloc->n_nodes > 0) { + memcpy(node_buf, galloc->node_allocs, galloc->n_nodes * sizeof(struct node_alloc)); + } + if (leaf_buf && galloc->n_leafs > 0) { + memcpy(leaf_buf, galloc->leaf_allocs, galloc->n_leafs * sizeof(struct leaf_alloc)); + } +} + +void ggml_gallocr_restore_state(ggml_gallocr_t galloc, + const void * node_buf, size_t node_size, + const void * leaf_buf, size_t leaf_size, + int n_nodes, int n_leafs) { + if (!galloc || n_nodes == 0) return; + + if (galloc->n_nodes < n_nodes) { + free(galloc->node_allocs); + galloc->node_allocs = calloc(n_nodes, sizeof(struct node_alloc)); + GGML_ASSERT(galloc->node_allocs != NULL); + } + if (galloc->n_leafs < n_leafs) { + free(galloc->leaf_allocs); + galloc->leaf_allocs = calloc(n_leafs, sizeof(struct leaf_alloc)); + GGML_ASSERT(galloc->leaf_allocs != NULL); + } + galloc->n_nodes = n_nodes; + galloc->n_leafs = n_leafs; + if (node_buf && node_size > 0) { + memcpy(galloc->node_allocs, node_buf, node_size); + } + if (leaf_buf && leaf_size > 0) { + memcpy(galloc->leaf_allocs, leaf_buf, leaf_size); + } +} + // utils static void free_buffers(ggml_backend_buffer_t ** buffers, const size_t * n_buffers) { diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 9c56ec30c5f1..4eefdf405366 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -137,6 +137,9 @@ extern "C" { // (optional) sort/optimize the nodes in the graph void (*graph_optimize) (ggml_backend_t backend, struct ggml_cgraph * cgraph); + + // (optional) asynchronous tensor memset on the backend's stream + void (*memset_tensor_async)(ggml_backend_t backend, struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size); }; struct ggml_backend { diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 1a555bf2a4dc..d4307ea0cec3 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -411,6 +411,20 @@ void ggml_backend_tensor_memset(struct ggml_tensor * tensor, uint8_t value, size buf->iface.memset_tensor(buf, tensor, value, offset, size); } +void ggml_backend_tensor_memset_async(ggml_backend_t backend, struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { + GGML_ASSERT(backend); + GGML_ASSERT(tensor); + GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); + + if (backend->iface.memset_tensor_async == NULL) { + ggml_backend_synchronize(backend); + ggml_backend_tensor_memset(tensor, value, offset, size); + } else { + backend->iface.memset_tensor_async(backend, tensor, value, offset, size); + } +} + void ggml_backend_synchronize(ggml_backend_t backend) { GGML_ASSERT(backend); if (backend->iface.synchronize == NULL) { @@ -767,6 +781,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 +800,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 +831,20 @@ struct ggml_backend_sched { ggml_backend_sched_eval_callback callback_eval; void * callback_eval_user_data; + bool prefetch_weights; + + int redirect_target[GGML_SCHED_MAX_BACKENDS]; + bool has_redirects; + + ggml_backend_sched_split_cb split_pre_compute; + ggml_backend_sched_split_cb split_post_compute; + ggml_backend_sched_split_cb split_prefetch_cb; + void * split_cb_user_data; + + ggml_backend_t copy_backends[GGML_SCHED_MAX_BACKENDS]; + ggml_backend_event_t copy_events[GGML_SCHED_MAX_BACKENDS]; + ggml_backend_event_t compute_events[GGML_SCHED_MAX_BACKENDS]; + char * context_buffer; size_t context_buffer_size; @@ -848,6 +880,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) && @@ -1183,8 +1232,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 +1306,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 +1363,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 +1400,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 +1428,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 +1476,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 +1525,76 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra graph_copy->nodes[graph_copy->n_nodes++] = input_cpy; } + // reserve current split's 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 next GPU split's prefetch destinations (race-protect copy stream vs current compute) + if (sched->prefetch_weights) { + struct ggml_backend_sched_split * next_gpu = NULL; + for (int ni = i + 1; ni < sched->n_splits; ni++) { + struct ggml_backend_sched_split * candidate = &sched->splits[ni]; + if (sched->copy_backends[candidate->backend_id] == NULL) continue; + 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; break; } + } + 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) { @@ -1519,6 +1676,9 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { // synchronize without ggml_backend_sched_synchronize to avoid changing cur_copy for (int i = 0; i < sched->n_backends; i++) { ggml_backend_synchronize(sched->backends[i]); + if (sched->copy_backends[i] != NULL) { + ggml_backend_synchronize(sched->copy_backends[i]); + } } ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids); @@ -1538,17 +1698,97 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s ggml_tensor * prev_ids_tensor = nullptr; std::vector ids; std::vector used_ids; + bool next_weights_prefetched = false; 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]; + } + + bool weights_prefetched = next_weights_prefetched; + next_weights_prefetched = false; + + if (sched->prefetch_weights) { + ggml_backend_t copy_backend = sched->copy_backends[split_backend_id]; + if (copy_backend != NULL || weights_prefetched) { + if (weights_prefetched && sched->copy_events[copy_backend_id] != NULL) { + ggml_backend_event_wait(split_backend, sched->copy_events[copy_backend_id]); + } + + { + struct ggml_backend_sched_split * next_gpu = NULL; + for (int nid = split_id + 1; nid < sched->n_splits; nid++) { + struct ggml_backend_sched_split * candidate = &splits[nid]; + if (sched->copy_backends[candidate->backend_id] == NULL) continue; + 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; break; } + } + if (next_gpu != NULL) { + ggml_backend_t next_copy = sched->copy_backends[next_gpu->backend_id]; + + if (sched->compute_events[split_backend_id] != NULL) { + ggml_backend_event_wait(next_copy, sched->compute_events[split_backend_id]); + } + + for (int input_id = 0; input_id < next_gpu->n_inputs; input_id++) { + struct ggml_tensor * next_input = next_gpu->inputs[input_id]; + 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)) { + struct ggml_tensor * input_cpy = tensor_copy(next_input, next_gpu->backend_id, sched->cur_copy); + ggml_backend_tensor_set_async(next_copy, input_cpy, next_input->data, 0, ggml_nbytes(next_input)); + next_weights_prefetched = true; + } + } + + if (sched->split_prefetch_cb != NULL && next_gpu->n_writeback > 0) { + for (int w = 0; w < next_gpu->n_writeback; w++) { + sched->split_prefetch_cb(next_gpu->writeback[w], next_copy, sched->split_cb_user_data); + } + next_weights_prefetched = true; + } + + if (next_weights_prefetched) { + ggml_backend_event_record(sched->copy_events[next_gpu->backend_id], next_copy); + } + } + } + } + } + // 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]); + { + int input_bid = ggml_backend_sched_backend_id(sched, input_backend); + if (input_bid >= 0 && sched->redirect_target[input_bid] >= 0) { + input_backend = sched->backends[sched->redirect_target[input_bid]]; + } + } 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 (weights_prefetched && + input->buffer != NULL && + ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && + ggml_backend_buffer_is_host(input->buffer)) { + continue; + } 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,12 +1901,24 @@ 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); + } } } } } + // pre-compute: upload stateful cache (KV/RS) for WRITEBACK tensors in this split + if (sched->split_pre_compute != NULL && split->n_writeback > 0) { + for (int w = 0; w < split->n_writeback; w++) { + sched->split_pre_compute(split->writeback[w], split_backend, sched->split_cb_user_data); + } + } + if (!sched->callback_eval) { enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph); if (ec != GGML_STATUS_SUCCESS) { @@ -1706,6 +1958,19 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } + // post-compute: download stateful cache (KV/RS) for WRITEBACK tensors in this split + if (sched->split_post_compute != NULL && split->n_writeback > 0) { + ggml_backend_synchronize(split_backend); + for (int w = 0; w < split->n_writeback; w++) { + sched->split_post_compute(split->writeback[w], split_backend, sched->split_cb_user_data); + } + } + + // record compute done for copy stream sync + if (sched->compute_events[split_backend_id] != NULL) { + ggml_backend_event_record(sched->compute_events[split_backend_id], split_backend); + } + // record the event of this copy if (split->n_inputs > 0) { if (sched->events[split_backend_id][sched->cur_copy] != NULL) { @@ -1747,10 +2012,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: 3 per input (dep, cpy, prefetch 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 +2028,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 +2047,34 @@ 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; + sched->copy_events[i] = NULL; + sched->compute_events[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->copy_events[i] = ggml_backend_event_new(dev_i); + sched->compute_events[i] = ggml_backend_event_new(dev_i); + } + } + sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends); sched->op_offload = op_offload; @@ -1794,12 +2091,19 @@ 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]); } + ggml_backend_event_free(sched->copy_events[b]); + ggml_backend_event_free(sched->compute_events[b]); + 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 +2121,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 +2219,151 @@ 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; + + if (enabled) { + for (int b = 0; b < sched->n_backends; b++) { + if (sched->copy_backends[b] != NULL) continue; + ggml_backend_dev_t dev = ggml_backend_get_device(sched->backends[b]); + if (dev == NULL) continue; + struct ggml_backend_dev_props props; + ggml_backend_dev_get_props(dev, &props); + if (props.caps.copy_stream) { + sched->copy_backends[b] = ggml_backend_dev_init(dev, NULL); + sched->copy_events[b] = ggml_backend_event_new(dev); + sched->compute_events[b] = ggml_backend_event_new(dev); + } + } + } else { + for (int b = 0; b < sched->n_backends; b++) { + ggml_backend_event_free(sched->copy_events[b]); + sched->copy_events[b] = NULL; + ggml_backend_event_free(sched->compute_events[b]); + sched->compute_events[b] = NULL; + if (sched->copy_backends[b] != NULL) { + ggml_backend_free(sched->copy_backends[b]); + sched->copy_backends[b] = NULL; + } + } + } +} + +void ggml_backend_sched_set_split_callbacks( + ggml_backend_sched_t sched, + ggml_backend_sched_split_cb pre_compute, + ggml_backend_sched_split_cb post_compute, + void * user_data) { + GGML_ASSERT(sched); + sched->split_pre_compute = pre_compute; + sched->split_post_compute = post_compute; + sched->split_cb_user_data = user_data; +} + +void ggml_backend_sched_set_prefetch_cb( + ggml_backend_sched_t sched, + ggml_backend_sched_split_cb prefetch_cb) { + GGML_ASSERT(sched); + sched->split_prefetch_cb = prefetch_cb; +} + +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; +} + +void ggml_backend_sched_set_buffer(ggml_backend_sched_t sched, ggml_backend_t backend, ggml_backend_buffer_t buffer, size_t alloc_offset, size_t alloc_size) { + GGML_ASSERT(sched); + int backend_index = ggml_backend_sched_backend_id(sched, backend); + GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); + ggml_gallocr_set_buffer(sched->galloc, backend_index, buffer, alloc_offset, alloc_size); +} + +void ggml_backend_sched_set_alloc_range(ggml_backend_sched_t sched, ggml_backend_t backend, size_t alloc_offset, size_t alloc_size) { + GGML_ASSERT(sched); + int backend_index = ggml_backend_sched_backend_id(sched, backend); + GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); + ggml_gallocr_set_alloc_range(sched->galloc, backend_index, alloc_offset, alloc_size); +} + +int ggml_backend_sched_get_n_chunks(ggml_backend_sched_t sched, 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); + return ggml_gallocr_get_n_chunks(sched->galloc, backend_index); +} + +size_t ggml_backend_sched_get_chunk_max_size(ggml_backend_sched_t sched, ggml_backend_t backend, int chunk_id) { + GGML_ASSERT(sched); + int backend_index = ggml_backend_sched_backend_id(sched, backend); + GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); + return ggml_gallocr_get_chunk_max_size(sched->galloc, backend_index, chunk_id); +} + +ggml_gallocr_t ggml_backend_sched_get_galloc(ggml_backend_sched_t sched) { + GGML_ASSERT(sched); + return sched->galloc; +} + +void ggml_backend_sched_save_backend_ids(ggml_backend_sched_t sched, int * node_buf, int * leaf_buf, int * n_nodes, int * n_leafs) { + GGML_ASSERT(sched); + int nn = sched->graph.n_nodes; + int nl = sched->graph.n_leafs; + if (n_nodes) *n_nodes = nn; + if (n_leafs) *n_leafs = nl; + if (node_buf) memcpy(node_buf, sched->node_backend_ids, nn * sizeof(int)); + if (leaf_buf) memcpy(leaf_buf, sched->leaf_backend_ids, nl * sizeof(int)); +} + +void ggml_backend_sched_restore_backend_ids(ggml_backend_sched_t sched, const int * node_buf, int n_nodes, const int * leaf_buf, int n_leafs) { + if (!sched) return; + if (node_buf && n_nodes > 0) { + memcpy(sched->prev_node_backend_ids, node_buf, n_nodes * sizeof(int)); + memcpy(sched->node_backend_ids, node_buf, n_nodes * sizeof(int)); + } + if (leaf_buf && n_leafs > 0) { + memcpy(sched->prev_leaf_backend_ids, leaf_buf, n_leafs * sizeof(int)); + memcpy(sched->leaf_backend_ids, leaf_buf, n_leafs * sizeof(int)); + } +} + +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; + + out->input_weight_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); + } 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]); + } + + return true; +} + int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched) { GGML_ASSERT(sched); return sched->n_splits; @@ -1955,8 +2406,13 @@ 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. } 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..3d2670926c51 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4426,6 +4426,12 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph } } +static void ggml_backend_cuda_memset_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; + ggml_cuda_set_device(cuda_ctx->device); + CUDA_CHECK(cudaMemsetAsync((char *)tensor->data + offset, value, size, cuda_ctx->stream())); +} + static const ggml_backend_i ggml_backend_cuda_interface = { /* .get_name = */ ggml_backend_cuda_get_name, /* .free = */ ggml_backend_cuda_free, @@ -4443,6 +4449,7 @@ static const ggml_backend_i ggml_backend_cuda_interface = { /* .event_record = */ ggml_backend_cuda_event_record, /* .event_wait = */ ggml_backend_cuda_event_wait, /* .graph_optimize = */ ggml_backend_cuda_graph_optimize, + /* .memset_tensor_async = */ ggml_backend_cuda_memset_tensor_async, }; static ggml_guid_t ggml_backend_cuda_guid() { @@ -4658,6 +4665,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..60368e8bbf77 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,10 @@ 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_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 +385,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 +539,25 @@ 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 + + // 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..51e02760c3a9 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -14,6 +14,7 @@ add_library(llama llama-batch.cpp llama-chat.cpp llama-context.cpp + llama-context-pshard.cpp llama-cparams.cpp llama-grammar.cpp llama-graph.cpp @@ -25,11 +26,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-cache.cpp llama-quant.cpp llama-sampler.cpp llama-vocab.cpp diff --git a/src/llama-context-pshard.cpp b/src/llama-context-pshard.cpp new file mode 100644 index 000000000000..a459eb86eb71 --- /dev/null +++ b/src/llama-context-pshard.cpp @@ -0,0 +1,646 @@ +#include "llama-context.h" + +#include "llama-impl.h" +#include "llama-kv-cache.h" +#include "llama-kv-cache-iswa.h" +#include "llama-memory-hybrid.h" +#include "llama-memory-hybrid-iswa.h" +#include "llama-pipe-shard.h" +#include "llama-model.h" +#include "llama-pshard-plan.h" + +#include "ggml-backend.h" +#include "ggml-alloc.h" + +#include +#include +#include + +namespace { + struct llama_pshard_split_cb_context { + std::vector pipe_shards; + }; + + void pshard_pre_compute(ggml_tensor * tensor, ggml_backend_t backend, void * user_data) { + auto * ctx = (llama_pshard_split_cb_context *) user_data; + for (auto * ps : ctx->pipe_shards) { + if (ps->upload_if_owned(tensor, backend)) { + return; + } + } + } + + void pshard_prefetch(ggml_tensor * tensor, ggml_backend_t copy_backend, void * user_data) { + auto * ctx = (llama_pshard_split_cb_context *) user_data; + for (auto * ps : ctx->pipe_shards) { + if (ps->prefetch_if_owned(tensor, copy_backend)) { + return; + } + } + } + + void pshard_post_compute(ggml_tensor * tensor, ggml_backend_t backend, void * user_data) { + auto * ctx = (llama_pshard_split_cb_context *) user_data; + for (auto * ps : ctx->pipe_shards) { + if (ps->download_if_owned(tensor, backend)) { + return; + } + } + } + + thread_local llama_pshard_split_cb_context g_split_ctx; + thread_local std::vector> g_kv_write_cells; + thread_local std::vector> g_swa_write_cells; + + size_t total_pinned_cache_size(llama_memory_i * mem) { + size_t total = 0; + for (auto * ps : mem->get_pipe_shards()) { + total += ps->current_pinned_size(); + } + return total; + } + + void zero_pinned_layers(const std::vector & layers) { + for (const auto & l : layers) { + if (!l.is_pinned) continue; + size_t t1_bytes = ggml_nbytes(l.t1_gpu); + size_t t2_bytes = l.t2_gpu ? ggml_nbytes(l.t2_gpu) : 0; + std::vector zeros(std::max(t1_bytes, t2_bytes), 0); + ggml_backend_tensor_set(l.t1_gpu, zeros.data(), 0, t1_bytes); + if (l.t2_gpu) { + ggml_backend_tensor_set(l.t2_gpu, zeros.data(), 0, t2_bytes); + } + } + } + + void sync_pins(llama_memory_pipe_shard_i * ps, + const std::function & bid_for, + const pshard_dev_layout & layout, + ggml_backend_t gpu) { + for (const auto & l : ps->get_layers()) { + const int32_t bid = bid_for(l.il); + const bool want_pinned = (bid == layout.compute); + const bool want_cpu = (bid == layout.cpu); + + if (want_pinned && !l.is_pinned) { + ps->pin_layer(l.il); + ggml_backend_tensor_memset_async(gpu, l.t1_gpu, 0, 0, ggml_nbytes(l.t1_gpu)); + if (l.t2_gpu) { + ggml_backend_tensor_memset_async(gpu, l.t2_gpu, 0, 0, ggml_nbytes(l.t2_gpu)); + } + } else if (!want_pinned && l.is_pinned) { + ps->unpin_layer(l.il); + } + + if (bid >= 0) { + if (want_cpu) { + ps->activate_cpu(l.il); + } else { + ps->activate_gpu(l.il); + } + } + } + } +} // namespace + +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, + ggml_cgraph * gf) { + const auto & tbids = model.get_tensor_backend_ids(); + + for (const auto & [tensor, bid] : tbids) { + if (bid >= 0 && bid < (int32_t)backends.size()) { + ggml_backend_sched_set_tensor_backend(sched, tensor, backends[bid].get()); + } + } + + // catches compute nodes never named via cb() + // (e.g. wo matmul inside build_attn) + if (gf) { + const int n_nodes = ggml_graph_n_nodes(gf); + for (int i = 0; i < n_nodes; i++) { + ggml_tensor * cur = ggml_graph_node(gf, i); + if (cur->view_src) continue; + 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 && it->second < (int32_t)backends.size()) { + ggml_backend_sched_set_tensor_backend(sched, cur, backends[it->second].get()); + break; + } + } + } + } + + if (memory) { + const auto & lbids = model.get_layer_backend_ids(); + for (auto * ps : memory->get_pipe_shards()) { + ps->assign_tensors(sched, lbids, backends, layout); + } + } +} + +void pshard_refresh_stream_views(llama_memory_i * memory) { + if (!memory) return; + for (auto * ps : memory->get_pipe_shards()) { + for (const auto & l : ps->get_layers()) { + ps->refresh_stream_views(l.il); + } + } +} + +void llama_context::pshard_pack_cache_region() { + auto * buf = model.get_dev_preload_buf(); + if (!buf) return; + + auto pipe_shards = memory->get_pipe_shards(); + if (pipe_shards.empty()) return; + + size_t buf_total = ggml_backend_buffer_get_size(buf); + size_t alignment = ggml_backend_buffer_get_alignment(buf); + void * buf_base = ggml_backend_buffer_get_base(buf); + + const auto & lbids = model.get_layer_backend_ids(); + const auto & layout = pshard_layout; + + auto is_pinned = [&](uint32_t il) -> bool { + auto it = lbids.find(il); + return it != lbids.end() && it->second == layout.compute; + }; + + struct cache_entry { + uint32_t il; + llama_memory_pipe_shard_i * ps; + ggml_tensor * t1; + ggml_tensor * t2; + }; + + std::vector entries; + + for (auto * ps : pipe_shards) { + for (const auto & l : ps->get_layers()) { + entries.push_back({ l.il, ps, l.t1_gpu, l.t2_gpu }); + } + } + + std::sort(entries.begin(), entries.end(), [](const cache_entry & a, const cache_entry & b) { return a.il < b.il; }); + + size_t total_cache = 0; + + size_t offset_from_right = 0; + for (size_t i = 0; i < entries.size(); i++) { + auto & e = entries[i]; + + size_t t1_size = ((ggml_backend_buffer_get_alloc_size(buf, e.t1) + alignment - 1) / alignment) * alignment; + size_t t2_size = e.t2 ? ((ggml_backend_buffer_get_alloc_size(buf, e.t2) + alignment - 1) / alignment) * alignment : 0; + size_t layer_total = t1_size + t2_size; + + void * t2_addr = nullptr; + if (t2_size > 0) { + offset_from_right += t2_size; + t2_addr = (char *)buf_base + buf_total - offset_from_right; + } + + offset_from_right += t1_size; + void * t1_addr = (char *)buf_base + buf_total - offset_from_right; + + e.ps->set_external_addrs(e.il, t1_addr, t2_addr, layer_total); + total_cache += layer_total; + } + + size_t n_pinned_total = 0; + for (auto & e : entries) { + if (!is_pinned(e.il)) continue; + e.ps->pin_layer(e.il); + n_pinned_total++; + } + + for (auto * ps : pipe_shards) { + zero_pinned_layers(ps->get_layers()); + } + + size_t preloaded_size = model.get_dev_preloaded_size(); + size_t pinned_cache = total_pinned_cache_size(memory.get()); + + LLAMA_LOG_INFO("%s: %zu cache layers (%.2f MiB total), %zu pinned (%.2f MiB)\n", + __func__, entries.size(), total_cache / (1024.0 * 1024.0), + n_pinned_total, pinned_cache / (1024.0 * 1024.0)); + LLAMA_LOG_INFO("%s: layout: [weights 0..%.2f | scratch %.2f..%.2f | cache %.2f..%.2f MiB]\n", + __func__, + preloaded_size / (1024.0 * 1024.0), + preloaded_size / (1024.0 * 1024.0), + (buf_total - pinned_cache) / (1024.0 * 1024.0), + (buf_total - pinned_cache) / (1024.0 * 1024.0), + buf_total / (1024.0 * 1024.0)); +} + +void llama_context::pshard_setup_sched() { + ggml_backend_sched_set_prefetch_weights(sched.get(), true); + + g_split_ctx = {}; + g_split_ctx.pipe_shards = memory->get_pipe_shards(); + + if (!g_split_ctx.pipe_shards.empty()) { + ggml_backend_sched_set_split_callbacks(sched.get(), pshard_pre_compute, pshard_post_compute, &g_split_ctx); + ggml_backend_sched_set_prefetch_cb(sched.get(), pshard_prefetch); + } + + if (model.get_dev_preload_buf()) { + size_t preloaded_size = model.get_dev_preloaded_size(); + size_t buf_total = ggml_backend_buffer_get_size(model.get_dev_preload_buf()); + + size_t pinned_cache_size = total_pinned_cache_size(memory.get()); + + if (preloaded_size + pinned_cache_size > buf_total) { + LLAMA_LOG_ERROR("%s: weights = %.2f MiB, pinned cache = %.2f MiB, buffer = %.2f MiB, overshoot = %.2f MiB\n", + __func__, + preloaded_size / (1024.0 * 1024.0), + pinned_cache_size / (1024.0 * 1024.0), + buf_total / (1024.0 * 1024.0), + (preloaded_size + pinned_cache_size - buf_total) / (1024.0 * 1024.0)); + } + GGML_ASSERT(preloaded_size + pinned_cache_size <= buf_total && + "pshard: weights + pinned cache exceed VRAM buffer -- plan overshoots budget"); + + size_t scratch_size = buf_total - preloaded_size - pinned_cache_size; + + ggml_backend_sched_set_buffer(sched.get(), backends[pshard_layout.compute].get(), + model.get_dev_preload_buf(), preloaded_size, scratch_size); + + LLAMA_LOG_DEBUG("%s: single buffer: weights %.2f MiB, cache %.2f MiB, scratch %.2f MiB (total %.2f MiB)\n", + __func__, preloaded_size / (1024.0 * 1024.0), + pinned_cache_size / (1024.0 * 1024.0), + scratch_size / (1024.0 * 1024.0), buf_total / (1024.0 * 1024.0)); + } +} + +void llama_context::pshard_apply_plan(const llama_pshard_plan & plan, bool with_upload) { + ggml_backend_t gpu = backends[pshard_layout.compute].get(); + size_t scratch_off = const_cast(model).pshard_apply_plan(plan, with_upload ? gpu : nullptr); + + const auto & lbids = model.get_layer_backend_ids(); + const auto & layout = pshard_layout; + + auto bid_for = [&](uint32_t il) -> int32_t { + auto it = lbids.find(il); + return (it != lbids.end()) ? it->second : -1; + }; + + for (auto * ps : memory->get_pipe_shards()) { + sync_pins(ps, bid_for, layout, gpu); + } + + if (model.get_dev_preload_buf()) { + size_t buf_total = ggml_backend_buffer_get_size(model.get_dev_preload_buf()); + size_t pinned_cache_size = total_pinned_cache_size(memory.get()); + + if (scratch_off + pinned_cache_size > buf_total) { + LLAMA_LOG_ERROR("%s: scratch_off = %.2f MiB, pinned cache = %.2f MiB, buffer = %.2f MiB, overshoot = %.2f MiB\n", + __func__, + scratch_off / (1024.0 * 1024.0), + pinned_cache_size / (1024.0 * 1024.0), + buf_total / (1024.0 * 1024.0), + (scratch_off + pinned_cache_size - buf_total) / (1024.0 * 1024.0)); + } + GGML_ASSERT(scratch_off + pinned_cache_size <= buf_total && + "pshard: weights + pinned cache exceed VRAM buffer -- plan overshoots budget"); + + size_t scratch_size = buf_total - scratch_off - pinned_cache_size; + + ggml_backend_sched_set_alloc_range(sched.get(), backends[pshard_layout.compute].get(), + scratch_off, scratch_size); + + LLAMA_LOG_DEBUG("%s: scratch_off = %.2f MiB, cache = %.2f MiB, scratch = %.2f MiB\n", + __func__, scratch_off / (1024.0 * 1024.0), + pinned_cache_size / (1024.0 * 1024.0), scratch_size / (1024.0 * 1024.0)); + } + + if (plan.alloc_state.valid) { + auto * galloc = ggml_backend_sched_get_galloc(sched.get()); + ggml_gallocr_restore_state(galloc, + plan.alloc_state.node_allocs.data(), plan.alloc_state.node_allocs.size(), + plan.alloc_state.leaf_allocs.data(), plan.alloc_state.leaf_allocs.size(), + plan.alloc_state.n_nodes, plan.alloc_state.n_leafs); + ggml_backend_sched_restore_backend_ids(sched.get(), + plan.alloc_state.node_backend_ids.data(), (int)plan.alloc_state.node_backend_ids.size(), + plan.alloc_state.leaf_backend_ids.data(), (int)plan.alloc_state.leaf_backend_ids.size()); + } else { + pshard_reserve_and_save(plan); + } +} + +void llama_context::pshard_reserve_and_save(const llama_pshard_plan & plan) { + llama_memory_context_ptr mctx; + if (memory) { + mctx = memory->init_full(); + if (!mctx) { + LLAMA_LOG_ERROR("%s: failed to initialize memory context\n", __func__); + plan.alloc_state.valid = false; + return; + } + } + + const uint32_t n_seqs = cparams.n_seq_max; + const uint32_t n_tokens = plan.batch_size; + + // start with unconstrained scratch packing + ggml_backend_t gpu = backends[pshard_layout.compute].get(); + const bool external_buf = model.get_dev_preload_buf() != nullptr; + size_t scratch_off = 0; + size_t scratch_avail = 0; + + if (external_buf) { + const size_t buf_total = ggml_backend_buffer_get_size(model.get_dev_preload_buf()); + const size_t pinned_cache_size = total_pinned_cache_size(memory.get()); + scratch_off = plan.cached_scratch_off; + scratch_avail = buf_total - scratch_off - pinned_cache_size; + ggml_backend_sched_set_alloc_range(sched.get(), gpu, scratch_off, SIZE_MAX/2); + } + + auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get()); + + if (gf && external_buf) { + const int n_chunks = ggml_backend_sched_get_n_chunks(sched.get(), gpu); + const size_t chunk0_max = (n_chunks >= 1) ? ggml_backend_sched_get_chunk_max_size(sched.get(), gpu, 0) : 0; + const size_t chunk0_used = (chunk0_max > scratch_off) ? chunk0_max - scratch_off : 0; + + if (chunk0_used <= scratch_avail) { + ggml_backend_sched_set_alloc_range(sched.get(), gpu, scratch_off, chunk0_used); + pshard_save_alloc_state(plan); + return; + } + + LLAMA_LOG_WARN("%s: unconstrained packing %.2f MiB > scratch budget %.2f MiB; retrying constrained\n", + __func__, chunk0_used / (1024.0 * 1024.0), scratch_avail / (1024.0 * 1024.0)); + + ggml_backend_sched_set_alloc_range(sched.get(), gpu, scratch_off, scratch_avail); + gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get()); + } + + if (!gf) { + LLAMA_LOG_ERROR("%s: graph_reserve failed for plan %s n_pinned=%u; alloc state not saved\n", + __func__, llama_pshard_strategy_name(plan.strategy), plan.n_pinned); + plan.alloc_state.valid = false; + return; + } + + pshard_save_alloc_state(plan); +} + +void llama_context::pshard_save_alloc_state(const llama_pshard_plan & plan) { + auto * galloc = ggml_backend_sched_get_galloc(sched.get()); + + size_t node_size = 0, leaf_size = 0; + ggml_gallocr_get_state_sizes(galloc, &node_size, &leaf_size); + + plan.alloc_state.node_allocs.resize(node_size); + plan.alloc_state.leaf_allocs.resize(leaf_size); + ggml_gallocr_save_state(galloc, + plan.alloc_state.node_allocs.data(), + plan.alloc_state.leaf_allocs.data(), + &plan.alloc_state.n_nodes, + &plan.alloc_state.n_leafs); + + int sched_n_nodes = 0, sched_n_leafs = 0; + ggml_backend_sched_save_backend_ids(sched.get(), nullptr, nullptr, &sched_n_nodes, &sched_n_leafs); + + plan.alloc_state.node_backend_ids.resize(sched_n_nodes); + plan.alloc_state.leaf_backend_ids.resize(sched_n_leafs); + ggml_backend_sched_save_backend_ids(sched.get(), + plan.alloc_state.node_backend_ids.data(), + plan.alloc_state.leaf_backend_ids.data(), + nullptr, nullptr); + + plan.alloc_state.valid = (plan.alloc_state.n_nodes > 0); + + LLAMA_LOG_DEBUG("%s: saved alloc state: nodes=%d (%.1f KiB), leafs=%d (%.1f KiB), bids=%d/%d\n", + __func__, plan.alloc_state.n_nodes, node_size / 1024.0, + plan.alloc_state.n_leafs, leaf_size / 1024.0, + sched_n_nodes, sched_n_leafs); +} + +void llama_context::pshard_warmup_plans() { + auto * registry = model.get_plan_registry(); + if (!registry) return; + + LLAMA_LOG_INFO("%s: pre-computing scratch offsets for %zu tier plans ...\n", + __func__, registry->tier_sizes.size()); + const int64_t t0 = llama_time_us(); + + for (size_t t = 0; t < registry->tier_sizes.size(); t++) { + auto & plan = registry->best_plans[t]; + if (!plan.is_viable) continue; + + const_cast(model).pshard_compute_scratch_off(plan); // see pshard_apply_plan + + LLAMA_LOG_DEBUG("%s: tier %zu (bs=%u, %s, n_pinned=%u) scratch_off=%.2f MiB\n", + __func__, t, registry->tier_sizes[t], + llama_pshard_strategy_name(plan.strategy), plan.n_pinned, + plan.cached_scratch_off / (1024.0 * 1024.0)); + + pshard_apply_plan(plan, /*with_upload=*/false); + + if (!plan.alloc_state.valid) { + LLAMA_LOG_WARN("%s: tier %zu (bs=%u, %s, n_pinned=%u) reserve failed; marking unviable\n", + __func__, t, registry->tier_sizes[t], + llama_pshard_strategy_name(plan.strategy), plan.n_pinned); + plan.is_viable = false; + } + } + + llama_pshard_plan * initial = registry->get_best(0); + if (initial) { + pshard_apply_plan(*initial); + pshard_active_plan = initial; + } + + LLAMA_LOG_INFO("%s: tier summary:\n", __func__); + for (size_t t = 0; t < registry->tier_sizes.size(); t++) { + auto & plan = registry->best_plans[t]; + if (!plan.is_viable) { + LLAMA_LOG_INFO("%s: tier %zu bs=%-5u — no viable plan\n", __func__, t, registry->tier_sizes[t]); + continue; + } + + char attn_buf[32] = ""; + if (plan.n_attn_pinned > 0) { + snprintf(attn_buf, sizeof(attn_buf), " (attn=%u)", plan.n_attn_pinned); + } + + if (plan.tps > 0.0f) { + LLAMA_LOG_INFO("%s: tier %zu bs=%-5u %s n_pinned=%u%s tps=%.1f\n", + __func__, t, registry->tier_sizes[t], + llama_pshard_strategy_name(plan.strategy), plan.n_pinned, attn_buf, plan.tps); + } else { + LLAMA_LOG_INFO("%s: tier %zu bs=%-5u %s n_pinned=%u%s\n", + __func__, t, registry->tier_sizes[t], + llama_pshard_strategy_name(plan.strategy), plan.n_pinned, attn_buf); + } + } + + const int64_t t1 = llama_time_us(); + LLAMA_LOG_INFO("%s: pre-computed %zu tiers in %.1f ms\n", + __func__, registry->tier_sizes.size(), (t1 - t0) / 1000.0); +} + +void llama_context::pshard_switch_plan(const llama_pshard_plan & old_plan, const llama_pshard_plan & new_plan) { + ggml_backend_t gpu = backends[pshard_layout.compute].get(); + + auto pipe_shards = memory->get_pipe_shards(); + + // save old pin state + std::vector> old_pins(pipe_shards.size()); + for (size_t i = 0; i < pipe_shards.size(); i++) { + for (const auto & l : pipe_shards[i]->get_layers()) { + old_pins[i][l.il] = l.is_pinned; + } + } + + const_cast(model).pshard_set_backend_maps(new_plan); // see pshard_apply_plan + const auto & new_lbids = model.get_layer_backend_ids(); + auto will_be_pinned = [&](uint32_t il) -> bool { + auto it = new_lbids.find(il); + return it != new_lbids.end() && it->second == pshard_layout.compute; + }; + + // download layers moving off gpu + int n_down = 0, n_skip_down = 0; + for (auto * ps : pipe_shards) { + for (const auto & l : ps->get_layers()) { + if (!l.is_pinned) continue; + if (!will_be_pinned(l.il)) { + ps->download_for_switch(l.il, gpu); + n_down++; + } else { + n_skip_down++; + } + } + } + + // apply new plan + pshard_apply_plan(new_plan); + + // upload newly pinned layers + int n_up = 0, n_skip_up = 0; + for (size_t i = 0; i < pipe_shards.size(); i++) { + for (const auto & l : pipe_shards[i]->get_layers()) { + if (!l.is_pinned) continue; + auto it = old_pins[i].find(l.il); + if (it != old_pins[i].end() && it->second) { + n_skip_up++; + continue; + } + pipe_shards[i]->upload_for_switch(l.il, gpu); + n_up++; + } + } + + LLAMA_LOG_DEBUG("%s: %s (n_pinned=%u) -> %s (n_pinned=%u) | down=%d skip=%d up=%d skip=%d\n", + __func__, + llama_pshard_strategy_name(old_plan.strategy), old_plan.n_pinned, + llama_pshard_strategy_name(new_plan.strategy), new_plan.n_pinned, + n_down, n_skip_down, n_up, n_skip_up); +} + +// restore saved alloc state for the active plan +void llama_context::pshard_reapply_active_plan() { + if (!pshard_active_plan || !pshard_active_plan->alloc_state.valid) { + return; + } + const llama_pshard_plan & plan = *pshard_active_plan; + + if (model.get_dev_preload_buf()) { + const size_t buf_total = ggml_backend_buffer_get_size(model.get_dev_preload_buf()); + const size_t pinned_cache_size = total_pinned_cache_size(memory.get()); + const size_t scratch_off = plan.cached_scratch_off; + const size_t scratch_size = buf_total - scratch_off - pinned_cache_size; + + ggml_backend_sched_set_alloc_range(sched.get(), backends[pshard_layout.compute].get(), scratch_off, scratch_size); + } + + auto * galloc = ggml_backend_sched_get_galloc(sched.get()); + ggml_gallocr_restore_state(galloc, + plan.alloc_state.node_allocs.data(), plan.alloc_state.node_allocs.size(), + plan.alloc_state.leaf_allocs.data(), plan.alloc_state.leaf_allocs.size(), + plan.alloc_state.n_nodes, plan.alloc_state.n_leafs); + ggml_backend_sched_restore_backend_ids(sched.get(), + plan.alloc_state.node_backend_ids.data(), (int)plan.alloc_state.node_backend_ids.size(), + plan.alloc_state.leaf_backend_ids.data(), (int)plan.alloc_state.leaf_backend_ids.size()); +} + +void llama_context::pshard_maybe_switch(uint32_t n_tokens) { + auto * registry = model.get_plan_registry(); + if (!registry) return; + + size_t tier = registry->tier_index(n_tokens); + llama_pshard_plan * best = registry->get_best(tier); + if (!best) return; + + if (best != pshard_active_plan) { + if (pshard_active_plan) { + pshard_switch_plan(*pshard_active_plan, *best); + } else { + pshard_apply_plan(*best); + } + pshard_active_plan = best; + } else { + pshard_reapply_active_plan(); + } +} + +void llama_context::pshard_update_write_cells(llama_memory_context_i * mctx) { + g_kv_write_cells.clear(); + g_swa_write_cells.clear(); + for (auto * ps : g_split_ctx.pipe_shards) { + ps->set_write_cells(nullptr); + ps->clear_prefetch(); + } + + if (!mctx || g_split_ctx.pipe_shards.empty()) return; + + // bind write_cells to the matching pipe shard + // indices match get_pipe_shards order + // plain KV: [kv_ps] + // iSWA: [base_ps, swa_ps] + // hybrid: [kv_ps, rs_ps] + // hybrid_iswa: [base_ps, swa_ps, rs_ps] + auto assign_wc = [&](const llama_kv_cache_context * kv_ctx, + std::vector> & storage, size_t ps_idx) { + if (!kv_ctx || ps_idx >= g_split_ctx.pipe_shards.size()) return; + storage = kv_ctx->get_write_cells(); + bool has_any = false; + for (const auto & v : storage) { if (!v.empty()) { has_any = true; break; } } + if (has_any) { + g_split_ctx.pipe_shards[ps_idx]->set_write_cells(&storage); + LLAMA_LOG_DEBUG("%s: bound write_cells to pipe_shard[%zu] (%zu streams)\n", + __func__, ps_idx, storage.size()); + } + }; + + if (auto * kv_ctx = dynamic_cast(mctx)) { + assign_wc(kv_ctx, g_kv_write_cells, 0); + return; + } + + if (auto * iswa_ctx = dynamic_cast(mctx)) { + assign_wc(iswa_ctx->get_base(), g_kv_write_cells, 0); + assign_wc(iswa_ctx->get_swa(), g_swa_write_cells, 1); + return; + } + + if (auto * h = dynamic_cast(mctx)) { + assign_wc(h->get_attn(), g_kv_write_cells, 0); + return; + } + + if (auto * h = dynamic_cast(mctx)) { + auto * iswa_ctx = h->get_attn(); + assign_wc(iswa_ctx->get_base(), g_kv_write_cells, 0); + assign_wc(iswa_ctx->get_swa(), g_swa_write_cells, 1); + return; + } +} diff --git a/src/llama-context.cpp b/src/llama-context.cpp index ee0c29235cd2..bfe126cddf03 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -5,12 +5,16 @@ #include "llama-impl.h" #include "llama-batch.h" #include "llama-io.h" +#include "llama-kv-cache.h" #include "llama-memory.h" +#include "llama-memory-hybrid.h" #include "llama-mmap.h" #include "llama-model.h" +#include "llama-pshard-plan.h" #include "llama-ext.h" #include "llama.h" +#include #include #include #include @@ -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; @@ -153,6 +158,11 @@ llama_context::llama_context( cparams.flash_attn = params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED; cparams.auto_fa = params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO; + // pshard remaps kv by layer so auto_fa cannot use device matching + if (cparams.pshard) { + cparams.auto_fa = false; + } + cparams.fused_gdn_ar = true; cparams.fused_gdn_ch = true; cparams.auto_fgdn = true; @@ -219,12 +229,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 +268,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()); @@ -346,8 +373,17 @@ llama_context::llama_context( LLAMA_LOG_INFO("%s: pipeline parallelism enabled\n", __func__); } + if (cparams.pshard && !model.hparams.no_alloc) { + const_cast(model).sync_dev_preload(); + pshard_pack_cache_region(); + } + sched_reserve(); + if (cparams.pshard && !model.hparams.no_alloc) { + pshard_warmup_plans(); + } + if (!cparams.flash_attn) { if (ggml_is_quantized(params.type_v)) { throw std::runtime_error("quantized V cache was requested, but this requires Flash Attention"); @@ -400,7 +436,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 +455,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) { + pshard_setup_sched(); + } + llama_memory_context_ptr mctx; if (memory) { LLAMA_LOG_DEBUG("%s: reserving full memory module\n", __func__); @@ -555,73 +603,75 @@ 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 (!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()); - } + 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(); @@ -1197,7 +1247,9 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll res->reset(); ggml_backend_sched_reset(sched.get()); - ggml_backend_sched_set_eval_callback(sched.get(), cparams.cb_eval, cparams.cb_eval_user_data); + if (cparams.cb_eval) { + ggml_backend_sched_set_eval_callback(sched.get(), cparams.cb_eval, cparams.cb_eval_user_data); + } //const auto t_start_us = ggml_time_us(); @@ -1211,11 +1263,19 @@ 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, gf); + } + if (!ggml_backend_sched_alloc_graph(sched.get(), gf)) { LLAMA_LOG_ERROR("%s: failed to allocate graph\n", __func__); ret = GGML_STATUS_ALLOC_FAILED; return nullptr; } + + if (cparams.pshard) { + pshard_refresh_stream_views(memory.get()); + } } // set the input data for the input tensors @@ -1228,6 +1288,10 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll //LLAMA_LOG_INFO("graph set inputs time: %.3f ms\n", (ggml_time_us() - t_start_us)/1000.0); } + if (cparams.pshard) { + pshard_update_write_cells(mctx); + } + const auto status = graph_compute(res->get_gf(), ubatch.n_tokens > 1); if (status != GGML_STATUS_SUCCESS) { LLAMA_LOG_ERROR("%s: failed to compute graph, compute status: %d\n", __func__, status); @@ -1609,6 +1673,10 @@ int llama_context::decode(const llama_batch & batch_inp) { embd_seq.clear(); output_swaps.clear(); + if (cparams.pshard) { + pshard_maybe_switch(n_tokens_all); + } + sched_reserve(); bool did_optimize = false; @@ -1616,10 +1684,19 @@ int llama_context::decode(const llama_batch & batch_inp) { // handle any pending shifts/copies memory_update(false); + uint32_t n_ubatch_eff = cparams.n_ubatch; + if (cparams.pshard && n_tokens_all >= 512) { + auto * registry = model.get_plan_registry(); + if (registry && !registry->tier_sizes.empty()) { + const uint32_t max_ubatch = std::min(cparams.n_ubatch, registry->tier_sizes.back()); + n_ubatch_eff = registry->find_optimal_ubatch(n_tokens_all, max_ubatch); + } + } + llama_memory_context_ptr mctx; while (true) { - mctx = memory->init_batch(*balloc, cparams.n_ubatch, output_all); + mctx = memory->init_batch(*balloc, n_ubatch_eff, output_all); if (!mctx) { return -2; } @@ -2129,6 +2206,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, gf); + } + // initialize scheduler with the specified graph if (split_only) { if (sizes) { @@ -2199,17 +2280,104 @@ 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 + int32_t last_weight_bid = -1; + int last_il = -2; + return [&, 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); } + if (cparams.pshard) { + 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,6 +3086,7 @@ llama_context_params llama_context_default_params() { /*.kv_unified =*/ false, /*.sampler =*/ nullptr, /*.n_sampler =*/ 0, + /*.pshard =*/ false, }; return result; diff --git a/src/llama-context.h b/src/llama-context.h index e0d0085c1c3f..f1ca563965b9 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -13,6 +13,7 @@ #include struct llama_model; +struct llama_pshard_plan; class llama_batch_allocr; class llama_io_read_i; @@ -243,6 +244,17 @@ struct llama_context { llm_graph_cb graph_get_cb() const; + void pshard_setup_sched(); + void pshard_pack_cache_region(); + void pshard_apply_plan(const llama_pshard_plan & plan, bool with_upload = true); + void pshard_reapply_active_plan(); + void pshard_reserve_and_save(const llama_pshard_plan & plan); + void pshard_save_alloc_state(const llama_pshard_plan & plan); + void pshard_warmup_plans(); + void pshard_switch_plan(const llama_pshard_plan & old_plan, const llama_pshard_plan & new_plan); + void pshard_maybe_switch(uint32_t n_tokens); + void pshard_update_write_cells(llama_memory_context_i * mctx); + // TODO: read/write lora adapters and cvec size_t state_write_data(llama_io_write_i & io); size_t state_read_data (llama_io_read_i & io); @@ -313,6 +325,9 @@ struct llama_context { bool sched_need_reserve = true; + const llama_pshard_plan * pshard_active_plan = nullptr; + pshard_dev_layout pshard_layout = {}; + ggml_backend_t backend_cpu = nullptr; std::vector backends; @@ -357,3 +372,16 @@ struct llama_context { mutable int32_t n_reused = 0; // number of times the previous graph was reused }; + +// 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, + ggml_cgraph * gf); + +void pshard_refresh_stream_views(llama_memory_i * memory); 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-graph.cpp b/src/llama-graph.cpp index 8e2b6ab8e7e1..51c8a76b3ac8 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -925,7 +925,7 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) : n_embd_head_v (hparams.n_embd_head_v()), n_embd_v_gqa (hparams.n_embd_v_gqa()), n_expert (hparams.n_expert), - n_expert_used (cparams.warmup ? hparams.n_expert : hparams.n_expert_used), + n_expert_used (cparams.warmup && !cparams.pshard ? hparams.n_expert : hparams.n_expert_used), freq_base (cparams.rope_freq_base), freq_scale (cparams.rope_freq_scale), ext_factor (cparams.yarn_ext_factor), @@ -2128,8 +2128,12 @@ ggml_tensor * llm_graph_context::build_attn( const auto & k_idxs = inp->get_k_idxs(); const auto & v_idxs = inp->get_v_idxs(); - ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); - ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); + auto * set_k = mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il); + auto * set_v = mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il); + cb(set_k, "cache_k", il); + cb(set_v, "cache_v", il); + ggml_build_forward_expand(gf, set_k); + ggml_build_forward_expand(gf, set_v); } const auto & kq_mask = inp->get_kq_mask(); @@ -2214,7 +2218,9 @@ ggml_tensor * llm_graph_context::build_attn( { const auto & k_idxs = inp->get_k_idxs(); - ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + auto * set_k = mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il); + cb(set_k, "cache_k", il); + ggml_build_forward_expand(gf, set_k); } const auto & kq_mask = inp->get_kq_mask(); @@ -2290,13 +2296,17 @@ ggml_tensor * llm_graph_context::build_attn( if (k_cur) { const auto & k_idxs = is_swa ? inp->get_k_idxs_swa() : inp->get_k_idxs(); - ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + auto * set_k = mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il); + cb(set_k, "cache_k", il); + ggml_build_forward_expand(gf, set_k); } if (v_cur) { const auto & v_idxs = is_swa ? inp->get_v_idxs_swa() : inp->get_v_idxs(); - ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); + auto * set_v = mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il); + cb(set_v, "cache_v", il); + ggml_build_forward_expand(gf, set_v); } const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask(); 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..9ccbe433c270 100644 --- a/src/llama-kv-cache-iswa.h +++ b/src/llama-kv-cache-iswa.h @@ -70,6 +70,8 @@ class llama_kv_cache_iswa : public llama_memory_i { llama_kv_cache * get_base() const; llama_kv_cache * get_swa () const; + std::vector get_pipe_shards() override; + private: const llama_hparams & hparams; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 09102f549c8e..9d2c58f51eaf 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -158,6 +158,83 @@ 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; + + map_layer_ids[il] = (int32_t)specs.size(); + + // 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(); + pipe_shard_kv->mode = v_trans ? llama_memory_pshard::FULL : llama_memory_pshard::CELL_GRANULAR; + + 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; + }; + + pipe_shard_kv->on_cells_used = [this](uint32_t s) -> uint32_t { + return v_cells[s].used_max_p1(); + }; + + 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, model.get_dev_preload_buf())) { + 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,6 +697,16 @@ 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; } @@ -778,6 +865,11 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co GGML_ABORT("The current KV cache / model configuration does not support K-shift"); } + if (pipe_shard_kv) { + LLAMA_LOG_WARN("%s: pshard + KV shift -- testing pending, results may be incorrect\n", __func__); + pipe_shard_kv->prepare_for_host_access(); + } + LLAMA_LOG_DEBUG("%s: applying K-shift\n", __func__); // apply K-shift if needed @@ -1844,6 +1936,10 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { GGML_UNUSED(flags); + if (pipe_shard_kv) { + pipe_shard_kv->prepare_for_host_access(); + } + io.write(&n_stream, sizeof(n_stream)); for (uint32_t s = 0; s < n_stream; ++s) { @@ -1897,6 +1993,10 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { GGML_UNUSED(flags); + if (pipe_shard_kv) { + pipe_shard_kv->prepare_for_host_access(); + } + GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())); uint32_t n_stream_cur; @@ -2502,3 +2602,24 @@ void llama_kv_cache_context::set_input_k_rot(ggml_tensor * dst) const { void llama_kv_cache_context::set_input_v_rot(ggml_tensor * dst) const { kv->set_input_v_rot(dst); } + +std::vector> llama_kv_cache_context::get_write_cells() const { + if (i_cur >= sinfos.size()) { + return {}; + } + const auto & sinfo = sinfos[i_cur]; + const uint32_t ns = kv->get_n_stream(); + std::vector> result(ns); + for (size_t i = 0; i < sinfo.idxs.size(); ++i) { + uint32_t s = sinfo.s0 + (uint32_t)i; + if (s < ns) { + result[s] = sinfo.idxs[i]; + } + } + return result; +} + +std::vector llama_kv_cache::get_pipe_shards() { + if (pipe_shard_kv) return { pipe_shard_kv.get() }; + return {}; +} diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 0b62dc7b2320..6f1e2892f423 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; @@ -385,6 +392,8 @@ class llama_kv_cache_context : public llama_memory_context_i { void set_input_k_rot(ggml_tensor * dst) const; void set_input_v_rot(ggml_tensor * dst) const; + std::vector> get_write_cells() const; + private: llama_memory_status status; diff --git a/src/llama-memory-hybrid-iswa.cpp b/src/llama-memory-hybrid-iswa.cpp index 10e6b4597972..be2ab4f3a08d 100644 --- a/src/llama-memory-hybrid-iswa.cpp +++ b/src/llama-memory-hybrid-iswa.cpp @@ -200,6 +200,13 @@ llama_memory_recurrent * llama_memory_hybrid_iswa::get_mem_recr() const { return mem_recr.get(); } +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; +} + // // llama_memory_hybrid_iswa_context // diff --git a/src/llama-memory-hybrid-iswa.h b/src/llama-memory-hybrid-iswa.h index 807c8aac96c9..e0de0f663f10 100644 --- a/src/llama-memory-hybrid-iswa.h +++ b/src/llama-memory-hybrid-iswa.h @@ -82,6 +82,8 @@ class llama_memory_hybrid_iswa : public llama_memory_i { llama_kv_cache_iswa * get_mem_attn() const; llama_memory_recurrent * get_mem_recr() const; + std::vector get_pipe_shards() override; + private: const llama_hparams & hparams; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 4ce1af592c15..328b5a177f63 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -198,6 +198,13 @@ llama_memory_recurrent * llama_memory_hybrid::get_mem_recr() const { return mem_recr.get(); } +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; +} + llama_memory_hybrid_context::llama_memory_hybrid_context(llama_memory_status status) : status(status) {} llama_memory_hybrid_context::llama_memory_hybrid_context(llama_memory_hybrid * mem) : diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 558cafdf984c..1efccc5f1de9 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -82,6 +82,8 @@ class llama_memory_hybrid : public llama_memory_i { llama_kv_cache * get_mem_attn() const; llama_memory_recurrent * get_mem_recr() const; + std::vector get_pipe_shards() override; + private: const llama_hparams & hparams; diff --git a/src/llama-memory-pshard.cpp b/src/llama-memory-pshard.cpp new file mode 100644 index 000000000000..8e96759a98db --- /dev/null +++ b/src/llama-memory-pshard.cpp @@ -0,0 +1,585 @@ +#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, + ggml_backend_buffer_t preload_buf) { + + layers.clear(); + streams.clear(); + ctxs.clear(); + bufs.clear(); + bufs_planned_sizes.clear(); + map_layer_ids.clear(); + external_buf = preload_buf; + 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 (external_buf) { + LLAMA_LOG_INFO("%s: %zu layers created, addresses deferred to pack_cache_region\n", + __func__, layers.size()); + } else { + 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) { + 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)" : ""); + + for (auto & l : layers) { + if (!is_sharded(l.il)) l.is_pinned = true; + } + 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, external_buf=%s, mode=%s\n", + __func__, n_pinned, n_sharded, + external_buf ? "yes" : "no", + mode == FULL ? "full" : "cell_granular"); + + 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::prepare_for_host_access() { + uint32_t n_activated = 0; + for (const auto & l : layers) { + if (!l.is_pinned) { + activate_cpu(l.il); + n_activated++; + } + } + LLAMA_LOG_DEBUG("%s: activated %u/%zu sharded layers to CPU\n", __func__, n_activated, layers.size()); +} + +void llama_memory_pshard::pin_layer(int32_t il) { + auto it = map_layer_ids.find(il); + if (it == map_layer_ids.end() || !external_buf) return; + auto & l = layers[it->second]; + if (l.is_pinned) return; + GGML_ASSERT(l.t1_gpu_addr && "pin_layer: t1 external address not set"); + if (l.t2_gpu) GGML_ASSERT(l.t2_gpu_addr && "pin_layer: t2 external address not set"); + LLAMA_LOG_DEBUG("pin_layer: il=%d t1_addr=%p t2_addr=%p\n", il, l.t1_gpu_addr, l.t2_gpu_addr); + + l.t1_gpu->data = l.t1_gpu_addr; + l.t1_gpu->buffer = external_buf; + if (l.t2_gpu) { + l.t2_gpu->data = l.t2_gpu_addr; + l.t2_gpu->buffer = external_buf; + } + + auto & sv = streams[it->second]; + for (auto * v : sv.t1_stream_gpu) { v->data = (char *)l.t1_gpu->data + v->view_offs; } + if (l.t2_gpu) { + for (auto * v : sv.t2_stream_gpu) { v->data = (char *)l.t2_gpu->data + v->view_offs; } + } + + l.is_pinned = true; + activate_gpu(il); +} + +void llama_memory_pshard::unpin_layer(int32_t il) { + auto it = map_layer_ids.find(il); + if (it == map_layer_ids.end()) return; + LLAMA_LOG_DEBUG("unpin_layer: il=%d\n", il); + auto & l = layers[it->second]; + if (!l.is_pinned) return; + LLAMA_LOG_DEBUG("%s: layer %d\n", __func__, il); + + l.t1_gpu->data = NULL; + l.t1_gpu->buffer = NULL; + if (l.t2_gpu) { + l.t2_gpu->data = NULL; + l.t2_gpu->buffer = NULL; + } + + auto & sv = streams[it->second]; + for (auto * v : sv.t1_stream_gpu) { v->data = NULL; } + if (l.t2_gpu) { + for (auto * v : sv.t2_stream_gpu) { v->data = NULL; } + } + + l.is_pinned = false; + activate_cpu(il); +} + +void llama_memory_pshard::set_external_addrs(int32_t il, void * a1, void * a2, size_t sz) { + auto it = map_layer_ids.find(il); + if (it == map_layer_ids.end()) return; + auto & l = layers[it->second]; + l.t1_gpu_addr = a1; + l.t2_gpu_addr = a2; + l.alloc_size = sz; +} + +void llama_memory_pshard::refresh_stream_views(int32_t il) { + auto it = map_layer_ids.find(il); + if (it == map_layer_ids.end()) return; + auto & l = layers[it->second]; + auto & sv = streams[it->second]; + if (l.t1_gpu && l.t1_gpu->data) { + for (auto * v : sv.t1_stream_gpu) v->data = (char *)l.t1_gpu->data + v->view_offs; + } + if (l.t2_gpu && l.t2_gpu->data) { + for (auto * v : sv.t2_stream_gpu) v->data = (char *)l.t2_gpu->data + v->view_offs; + } +} + +size_t llama_memory_pshard::current_pinned_size() const { + size_t total = 0; + for (const auto & l : layers) { + if (l.is_pinned) total += l.alloc_size; + } + return total; +} + +void llama_memory_pshard::upload_full_one(ggml_tensor * t_gpu, ggml_tensor * t_cpu, ggml_backend_t gpu) { + if (!t_gpu || !t_cpu || !t_cpu->data) return; + LLAMA_LOG_DEBUG("upload_full_one: name=%s gpu_data=%p cpu_data=%p bytes=%zu\n", + t_gpu->name, t_gpu->data, t_cpu->data, ggml_nbytes(t_cpu)); + ggml_backend_tensor_set_async(gpu, t_gpu, t_cpu->data, 0, ggml_nbytes(t_cpu)); +} + +void llama_memory_pshard::download_full_one(ggml_tensor * t_gpu, ggml_tensor * t_cpu, ggml_backend_t be) { + if (!t_gpu || !t_cpu || !t_cpu->data) return; + LLAMA_LOG_DEBUG("download_full_one: name=%s gpu_data=%p cpu_data=%p bytes=%zu\n", + t_gpu->name, t_gpu->data, t_cpu->data, ggml_nbytes(t_cpu)); + ggml_backend_tensor_get_async(be, t_gpu, t_cpu->data, 0, ggml_nbytes(t_cpu)); +} + +void llama_memory_pshard::upload_full(int32_t il, ggml_backend_t gpu) { + auto it = map_layer_ids.find(il); + if (it == map_layer_ids.end()) return; + auto & l = layers[it->second]; + GGML_ASSERT(l.t1_cpu && l.t1_cpu->data && "upload_full: CPU t1 not allocated"); + upload_full_one(l.t1_gpu, l.t1_cpu, gpu); + upload_full_one(l.t2_gpu, l.t2_cpu, gpu); +} + +void llama_memory_pshard::download_full(int32_t il, ggml_backend_t be) { + auto it = map_layer_ids.find(il); + if (it == map_layer_ids.end()) return; + auto & l = layers[it->second]; + GGML_ASSERT(l.t1_gpu && l.t1_cpu && l.t1_cpu->data && "download_full: t1 not allocated"); + download_full_one(l.t1_gpu, l.t1_cpu, be); + download_full_one(l.t2_gpu, l.t2_cpu, be); +} + +std::vector +llama_memory_pshard::batch_ranges(const std::vector & sorted) { + std::vector ranges; + if (sorted.empty()) return ranges; + uint32_t start = sorted[0], count = 1; + for (size_t i = 1; i < sorted.size(); i++) { + if (sorted[i] == start + count) { + count++; + } else { + ranges.push_back({start, count}); + start = sorted[i]; + count = 1; + } + } + ranges.push_back({start, count}); + return ranges; +} + +void llama_memory_pshard::upload_cells_one(int32_t il, ggml_tensor * t_gpu, ggml_tensor * t_cpu, ggml_backend_t gpu, bool zero_tail) { + (void) il; + if (!t_gpu || !t_cpu || !t_cpu->data || !on_cells_used) return; + + const uint32_t ns = (uint32_t)t_gpu->ne[2]; + const size_t t_row = t_gpu->ne[0] * ggml_element_size(t_gpu); + const size_t seq_sz = t_gpu->ne[1]; + + LLAMA_LOG_DEBUG("upload_cells_one: il=%d name=%s gpu_data=%p cpu_data=%p ns=%u\n", + il, t_gpu->name, t_gpu->data, t_cpu->data, ns); + + for (uint32_t s = 0; s < ns; s++) { + if (write_cells && s < write_cells->size() && (*write_cells)[s].empty()) { + continue; + } + + uint32_t n_used = on_cells_used(s); + size_t base = s * seq_sz * t_row; + + if (n_used > 0) { + ggml_backend_tensor_set_async(gpu, t_gpu, (char *)t_cpu->data + base, base, n_used * t_row); + } + if (zero_tail && n_used < seq_sz) { + size_t tail = base + n_used * t_row; + size_t tail_n = seq_sz - n_used; + ggml_backend_tensor_memset_async(gpu, t_gpu, 0, tail, tail_n * t_row); + } + } +} + +void llama_memory_pshard::download_cells_one(int32_t il, ggml_tensor * t_gpu, ggml_tensor * t_cpu, ggml_backend_t be) { + (void) il; + if (!t_gpu || !t_cpu || !t_cpu->data || !on_cells_used) return; + + const uint32_t ns = (uint32_t)t_gpu->ne[2]; + const size_t t_row = t_gpu->ne[0] * ggml_element_size(t_gpu); + const size_t seq_sz = t_gpu->ne[1]; + + LLAMA_LOG_DEBUG("download_cells_one: il=%d name=%s gpu_data=%p cpu_data=%p ns=%u\n", + il, t_gpu->name, t_gpu->data, t_cpu->data, ns); + + for (uint32_t s = 0; s < ns; s++) { + uint32_t n_used = on_cells_used(s); + if (n_used == 0) continue; + + size_t base = s * seq_sz * t_row; + ggml_backend_tensor_get_async(be, t_gpu, (char *)t_cpu->data + base, base, n_used * t_row); + } +} + +void llama_memory_pshard::download_written_one( + int32_t il, ggml_tensor * t_gpu, ggml_tensor * t_cpu, + const std::vector> & wc_per_stream, ggml_backend_t be) { + (void) il; + if (!t_gpu || !t_cpu || !t_cpu->data) return; + + const size_t t_row = t_gpu->ne[0] * ggml_element_size(t_gpu); + const size_t seq_sz = t_gpu->ne[1]; + const uint32_t ns = (uint32_t)t_gpu->ne[2]; + const uint32_t ns_wc = (uint32_t)std::min((size_t)ns, wc_per_stream.size()); + + LLAMA_LOG_DEBUG("download_written_one: il=%d name=%s gpu_data=%p cpu_data=%p ns_wc=%u\n", + il, t_gpu->name, t_gpu->data, t_cpu->data, ns_wc); + + for (uint32_t s = 0; s < ns_wc; s++) { + if (wc_per_stream[s].empty()) continue; + + const size_t base = s * seq_sz * t_row; + + std::vector sorted = wc_per_stream[s]; + std::sort(sorted.begin(), sorted.end()); + sorted.erase(std::unique(sorted.begin(), sorted.end()), sorted.end()); + auto ranges = batch_ranges(sorted); + + for (const auto & r : ranges) { + GGML_ASSERT(r.start + r.count <= seq_sz && "download_written_one: cell range exceeds kv_size"); + size_t off = base + r.start * t_row; + ggml_backend_tensor_get_async(be, t_gpu, (char *)t_cpu->data + off, off, r.count * t_row); + } + } +} + +void llama_memory_pshard::upload_cells(int32_t il, ggml_backend_t gpu, bool zero_tail) { + auto it = map_layer_ids.find(il); + if (it == map_layer_ids.end()) return; + auto & l = layers[it->second]; + GGML_ASSERT(l.t1_cpu && l.t1_cpu->data && "upload_cells: CPU t1 not allocated"); + upload_cells_one(il, l.t1_gpu, l.t1_cpu, gpu, zero_tail); + upload_cells_one(il, l.t2_gpu, l.t2_cpu, gpu, zero_tail); +} + +void llama_memory_pshard::download_cells(int32_t il, ggml_backend_t be) { + auto it = map_layer_ids.find(il); + if (it == map_layer_ids.end()) return; + auto & l = layers[it->second]; + GGML_ASSERT(l.t1_cpu && l.t1_cpu->data && "download_cells: CPU t1 not allocated"); + download_cells_one(il, l.t1_gpu, l.t1_cpu, be); + download_cells_one(il, l.t2_gpu, l.t2_cpu, be); +} + +void llama_memory_pshard::download_written( + int32_t il, const std::vector> & wc_per_stream, ggml_backend_t be) { + auto it = map_layer_ids.find(il); + if (it == map_layer_ids.end()) return; + auto & l = layers[it->second]; + GGML_ASSERT(l.t1_cpu && l.t1_cpu->data && "download_written: CPU t1 not allocated"); + download_written_one(il, l.t1_gpu, l.t1_cpu, wc_per_stream, be); + download_written_one(il, l.t2_gpu, l.t2_cpu, wc_per_stream, be); +} + +void llama_memory_pshard::clear_prefetch() { + for (auto & l : layers) { l.prefetched_t1 = false; l.prefetched_t2 = false; } +} + +bool llama_memory_pshard::prefetch_if_owned(ggml_tensor * t, ggml_backend_t copy_backend) { + for (auto & l : layers) { + ggml_tensor * t_gpu = (t == l.t1_gpu) ? l.t1_gpu : (t == l.t2_gpu) ? l.t2_gpu : nullptr; + if (!t_gpu) continue; + ggml_tensor * t_cpu = (t == l.t1_gpu) ? l.t1_cpu : l.t2_cpu; + bool & flag = (t == l.t1_gpu) ? l.prefetched_t1 : l.prefetched_t2; + + if (mode == FULL) upload_full_one(t_gpu, t_cpu, copy_backend); + else upload_cells_one(l.il, t_gpu, t_cpu, copy_backend, true); + flag = true; + return true; + } + return false; +} + +bool llama_memory_pshard::upload_if_owned(ggml_tensor * t, ggml_backend_t backend) { + for (auto & l : layers) { + ggml_tensor * t_gpu = (t == l.t1_gpu) ? l.t1_gpu : (t == l.t2_gpu) ? l.t2_gpu : nullptr; + if (!t_gpu) continue; + ggml_tensor * t_cpu = (t == l.t1_gpu) ? l.t1_cpu : l.t2_cpu; + bool & flag = (t == l.t1_gpu) ? l.prefetched_t1 : l.prefetched_t2; + + if (flag) { flag = false; return true; } + if (mode == FULL) upload_full_one(t_gpu, t_cpu, backend); + else upload_cells_one(l.il, t_gpu, t_cpu, backend, true); + return true; + } + return false; +} + +bool llama_memory_pshard::download_if_owned(ggml_tensor * t, ggml_backend_t backend) { + for (auto & l : layers) { + ggml_tensor * t_gpu = (t == l.t1_gpu) ? l.t1_gpu : (t == l.t2_gpu) ? l.t2_gpu : nullptr; + if (!t_gpu) continue; + ggml_tensor * t_cpu = (t == l.t1_gpu) ? l.t1_cpu : l.t2_cpu; + + if (mode == FULL) download_full_one(t_gpu, t_cpu, backend); + else if (write_cells) download_written_one(l.il, t_gpu, t_cpu, *write_cells, backend); + else download_full_one(t_gpu, t_cpu, backend); + return true; + } + return false; +} + +void llama_memory_pshard::upload_for_switch(int32_t il, ggml_backend_t be) { + if (mode == FULL) { + upload_full(il, be); + } else { + upload_cells(il, be, false); + } +} + +void llama_memory_pshard::download_for_switch(int32_t il, ggml_backend_t be) { + if (mode == FULL) { + download_full(il, be); + } else { + download_cells(il, be); + } +} + +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 (const auto & l : layers) { + auto it = layer_bids.find((int)l.il); + if (l.is_pinned) { + GGML_ASSERT(l.t1_gpu->data != nullptr && "pinned layer missing GPU address"); + LLAMA_LOG_DEBUG("%s: layer %u -> pinned (GPU)\n", __func__, l.il); + activate_gpu(l.il); + } else if (it != layer_bids.end() && it->second >= 0 && it->second < (int32_t)backends.size()) { + if (it->second == layout.cpu) { + LLAMA_LOG_DEBUG("%s: layer %u -> CPU (bid=%d)\n", __func__, l.il, it->second); + activate_cpu(l.il); + } else { + LLAMA_LOG_DEBUG("%s: layer %u -> shard (bid=%d)\n", __func__, l.il, it->second); + 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[it->second].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[it->second].get()); + ggml_backend_sched_add_writeback(sched, l.t2_gpu); + } + } + } else if (!l.is_pinned) { + LLAMA_LOG_WARN("%s: layer %u has no backend_id in plan -- left unconfigured\n", __func__, l.il); + } + } +} diff --git a/src/llama-memory-pshard.h b/src/llama-memory-pshard.h new file mode 100644 index 000000000000..a4b77ffd739a --- /dev/null +++ b/src/llama-memory-pshard.h @@ -0,0 +1,119 @@ +#pragma once + +#include "llama-pipe-shard.h" + +#include +#include +#include + +struct llama_model; +struct llama_hparams; +struct ggml_tensor; + +struct llama_memory_pshard : llama_memory_pipe_shard_i { + enum transfer_mode { + FULL, // always transfer entire tensors (RS, or KV with v_trans) + CELL_GRANULAR, // transfer by cell ranges (KV without v_trans) + }; + + 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; + using cells_used_fn_t = std::function; + + transfer_mode mode = FULL; + activate_fn_t on_activate_gpu; + activate_fn_t on_activate_cpu; + cells_used_fn_t on_cells_used; // null for FULL mode + + 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, + ggml_backend_buffer_t preload_buf); + + 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 clear_prefetch() override; + bool prefetch_if_owned(ggml_tensor * t, ggml_backend_t be) override; + bool upload_if_owned(ggml_tensor * t, ggml_backend_t be) override; + bool download_if_owned(ggml_tensor * t, ggml_backend_t be) override; + + void upload_for_switch(int32_t il, ggml_backend_t be) override; + void download_for_switch(int32_t il, ggml_backend_t be) override; + + void activate_gpu(int32_t il) override; + void activate_cpu(int32_t il) override; + void prepare_for_host_access() override; + void pin_layer(int32_t il) override; + void unpin_layer(int32_t il) override; + void set_external_addrs(int32_t il, void * a1, void * a2, size_t sz) override; + + void refresh_stream_views(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; + + size_t current_pinned_size() const 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; + ggml_backend_buffer_t external_buf = nullptr; + int32_t cpu_bid_ = -1; + std::unordered_map layer_bids_; + + void upload_full(int32_t il, ggml_backend_t gpu); + void download_full(int32_t il, ggml_backend_t be); + void upload_cells(int32_t il, ggml_backend_t gpu, bool zero_tail); + void download_cells(int32_t il, ggml_backend_t be); + void download_written(int32_t il, const std::vector> & wc_per_stream, ggml_backend_t be); + + void upload_full_one(ggml_tensor * t_gpu, ggml_tensor * t_cpu, ggml_backend_t gpu); + void download_full_one(ggml_tensor * t_gpu, ggml_tensor * t_cpu, ggml_backend_t be); + void upload_cells_one(int32_t il, ggml_tensor * t_gpu, ggml_tensor * t_cpu, ggml_backend_t gpu, bool zero_tail); + void download_cells_one(int32_t il, ggml_tensor * t_gpu, ggml_tensor * t_cpu, ggml_backend_t be); + void download_written_one(int32_t il, ggml_tensor * t_gpu, ggml_tensor * t_cpu, const std::vector> & wc_per_stream, ggml_backend_t be); + + struct cell_range { uint32_t start; uint32_t count; }; + static std::vector batch_ranges(const std::vector & sorted); +}; diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 9287fe45e963..b0507b68d4eb 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -68,6 +68,52 @@ 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->mode = llama_memory_pshard::FULL; + + 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, model.get_dev_preload_buf())) { + 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 +419,23 @@ 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; + ret[ggml_backend_buffer_get_type(bufs[i].get())] += 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..cb1574d19129 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); @@ -1693,3 +1709,100 @@ void llama_model_loader::print_info() const { LLAMA_LOG_INFO("%s: file size = %.2f GiB (%.2f BPW) \n", __func__, n_bytes/1024.0/1024.0/1024.0, n_bytes*8.0/n_elements); } } + +bool llama_model_loader::preload_weights_to_device( + const std::unordered_map & tensor_backend_ids, + int target_backend_id, + size_t buf_size, + ggml_backend_buffer_t * out_buf, + ggml_backend_t * out_backend, + std::unordered_map * out_preload_map, + size_t * out_preloaded_size) { + + auto * gpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU); + if (!gpu_dev) { + return false; + } + + auto * buft = ggml_backend_dev_buffer_type(gpu_dev); + + std::vector pinned; + for (const auto & [tensor, bid] : tensor_backend_ids) { + if (bid == target_backend_id && tensor->data != nullptr) { + pinned.push_back(tensor); + } + } + + { + size_t pinned_bytes = 0; + for (const auto * t : pinned) { pinned_bytes += ggml_nbytes(t); } + LLAMA_LOG_INFO("%s: %zu tensors with bid=%d (%.2f MiB), %zu total in map\n", + __func__, pinned.size(), target_backend_id, + pinned_bytes / (1024.0 * 1024.0), tensor_backend_ids.size()); + } + + auto get_layer = [](const ggml_tensor * t) -> int { + const char * blk = strstr(ggml_get_name(t), "blk."); + return blk ? atoi(blk + 4) : 9999; + }; + auto get_cat = [](const ggml_tensor * t) -> int { + const char * name = ggml_get_name(t); + if (strstr(name, "attn_")) return 0; + if (strstr(name, "exps")) return 4; + if (strstr(name, "ffn_")) return 1; + if (strstr(name, "norm")) return 2; + return 3; + }; + std::sort(pinned.begin(), pinned.end(), + [&](const ggml_tensor * a, const ggml_tensor * b) { + int ca = get_cat(a), cb = get_cat(b); + if (ca != cb) return ca < cb; + int la = get_layer(a), lb = get_layer(b); + if (la != lb) return la < lb; + return strcmp(ggml_get_name(a), ggml_get_name(b)) < 0; + }); + + *out_buf = ggml_backend_buft_alloc_buffer(buft, buf_size); + if (!*out_buf) { + LLAMA_LOG_WARN("%s: failed to allocate %.2f MiB device buffer\n", + __func__, buf_size / (1024.0 * 1024.0)); + return false; + } + + ggml_backend_buffer_set_usage(*out_buf, GGML_BACKEND_BUFFER_USAGE_COMPUTE); + + for (const auto & [tensor, bid] : tensor_backend_ids) { + if (bid >= 0 && tensor->data != nullptr) { + (*out_preload_map)[tensor] = { tensor->data, nullptr, tensor->buffer }; + } + } + + struct ggml_tallocr talloc = ggml_tallocr_new(*out_buf); + ggml_backend_t gpu = ggml_backend_dev_init(gpu_dev, nullptr); + + size_t n_packed = 0; + for (auto * tensor : pinned) { + size_t tsize = ggml_backend_buffer_get_alloc_size(*out_buf, tensor); + if (talloc.offset + tsize > buf_size) { + break; + } + + tensor->buffer = NULL; + tensor->data = NULL; + ggml_tallocr_alloc(&talloc, tensor); + + (*out_preload_map)[tensor].gpu_addr = tensor->data; + ggml_backend_tensor_set_async(gpu, tensor, (*out_preload_map)[tensor].cpu_addr, 0, ggml_nbytes(tensor)); + n_packed++; + } + + if (n_packed < pinned.size()) { + LLAMA_LOG_WARN("%s: %zu/%zu pinned tensors did not fit in %.2f MiB buffer\n", + __func__, pinned.size() - n_packed, pinned.size(), buf_size / (1024.0 * 1024.0)); + } + + *out_preloaded_size = talloc.offset; + *out_backend = gpu; + + return true; +} diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 7b3d6703c03d..6f8255864529 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -17,6 +17,12 @@ using llama_buf_map = std::unordered_map; +struct weight_preload_entry { + void * cpu_addr = nullptr; + void * gpu_addr = nullptr; + ggml_backend_buffer_t host_buffer = nullptr; +}; + // lists of buffer types used for each layer using buft_list_t = std::vector>; @@ -79,6 +85,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; @@ -201,6 +208,15 @@ struct llama_model_loader { llama_progress_callback progress_callback, void * progress_callback_user_data); + bool preload_weights_to_device( + const std::unordered_map & tensor_backend_ids, + int target_backend_id, + size_t buf_size, + ggml_backend_buffer_t * out_buf, + ggml_backend_t * out_backend, + std::unordered_map * out_preload_map, + size_t * out_preloaded_size); + std::string ftype_name() const; void print_info() const; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index b265394ef736..d3dcf93d10b3 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" @@ -16,6 +17,7 @@ #include "models/models.h" #include "ggml.h" +#include "ggml-alloc.h" #include "ggml-cpp.h" #include @@ -661,6 +663,16 @@ 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; + + ggml_backend_buffer_t dev_preload_buf = nullptr; + ggml_backend_t dev_preload_backend = nullptr; + std::unordered_map weight_preload_map; + size_t dev_preloaded_size = 0; }; llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique()) { @@ -671,6 +683,11 @@ llama_model::~llama_model() { for (auto * lora : loras) { delete lora; } + if (pimpl->dev_preload_backend) { + ggml_backend_synchronize(pimpl->dev_preload_backend); + ggml_backend_free(pimpl->dev_preload_backend); + } + ggml_backend_buffer_free(pimpl->dev_preload_buf); } void llama_model::load_stats(llama_model_loader & ml) { @@ -8115,6 +8132,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 +8299,23 @@ bool llama_model::load_tensors(llama_model_loader & ml) { } } + if (params.pshard && params.max_vram_alloc > 0) { + size_t buf_size = params.max_vram_alloc * 1024ULL * 1024ULL; + if (ml.preload_weights_to_device( + pimpl->tensor_backend_ids, 0, buf_size, + &pimpl->dev_preload_buf, + &pimpl->dev_preload_backend, + &pimpl->weight_preload_map, + &pimpl->dev_preloaded_size)) { + LLAMA_LOG_INFO("%s: preloaded %zu weights (%.2f MiB) into %.2f MiB device buffer\n", + __func__, pimpl->weight_preload_map.size(), + pimpl->dev_preloaded_size / (1024.0 * 1024.0), + buf_size / (1024.0 * 1024.0)); + } + } + + pimpl->plan_registry = params.pshard_registry; + return true; } @@ -8574,6 +8640,266 @@ bool llama_model::has_tensor_overrides() const { return pimpl->has_tensor_overrides; } +bool llama_model::is_pshard() const { + return params.pshard; +} + +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; +} + +ggml_backend_buffer_t llama_model::get_dev_preload_buf() const { + return pimpl->dev_preload_buf; +} + +size_t llama_model::get_dev_preloaded_size() const { + return pimpl->dev_preloaded_size; +} + +void llama_model::sync_dev_preload() { + if (pimpl->dev_preload_backend) { + ggml_backend_synchronize(pimpl->dev_preload_backend); + ggml_backend_free(pimpl->dev_preload_backend); + pimpl->dev_preload_backend = nullptr; + } +} + +size_t llama_model::pshard_compute_scratch_off(const llama_pshard_plan & plan) { + if (plan.addrs_cached) { + return plan.cached_scratch_off; + } + if (!pimpl->dev_preload_buf) { + return 0; + } + + pshard_set_backend_maps(plan); + + void * buf_base = ggml_backend_buffer_get_base(pimpl->dev_preload_buf); + + std::vector>> saved; + saved.reserve(pimpl->weight_preload_map.size()); + for (auto & [tensor, entry] : pimpl->weight_preload_map) { + saved.push_back({tensor, {tensor->data, tensor->buffer}}); + } + + std::vector pinned; + for (auto & [tensor, entry] : pimpl->weight_preload_map) { + auto it = pimpl->tensor_backend_ids.find(tensor); + if (it != pimpl->tensor_backend_ids.end() && it->second == 0) { + pinned.push_back(tensor); + } + } + + auto get_layer = [](const ggml_tensor * t) -> int { + const char * blk = strstr(ggml_get_name(t), "blk."); + return blk ? atoi(blk + 4) : 9999; + }; + auto get_cat = [](const ggml_tensor * t) -> int { + const char * name = ggml_get_name(t); + if (strstr(name, "attn_")) return 0; + if (strstr(name, "exps")) return 4; + if (strstr(name, "ffn_")) return 1; + if (strstr(name, "norm")) return 2; + return 3; + }; + std::sort(pinned.begin(), pinned.end(), + [&](const ggml_tensor * a, const ggml_tensor * b) { + int ca = get_cat(a), cb = get_cat(b); + if (ca != cb) return ca < cb; + int la = get_layer(a), lb = get_layer(b); + if (la != lb) return la < lb; + return strcmp(ggml_get_name(a), ggml_get_name(b)) < 0; + }); + + struct ggml_tallocr talloc = ggml_tallocr_new(pimpl->dev_preload_buf); + size_t buf_size = ggml_backend_buffer_get_size(pimpl->dev_preload_buf); + + for (auto * tensor : pinned) { + size_t tsize = ggml_backend_buffer_get_alloc_size(pimpl->dev_preload_buf, tensor); + if (talloc.offset + tsize > buf_size) break; + tensor->buffer = NULL; + tensor->data = NULL; + ggml_tallocr_alloc(&talloc, tensor); + } + + size_t scratch_off = talloc.offset; + + plan.cached_weight_offsets.clear(); + for (auto * tensor : pinned) { + if (tensor->data) { + plan.cached_weight_offsets[std::string(ggml_get_name(tensor))] = + (size_t)((char *)tensor->data - (char *)buf_base); + } + } + plan.cached_scratch_off = scratch_off; + plan.addrs_cached = true; + + for (auto & [tensor, state] : saved) { + tensor->data = state.first; + tensor->buffer = state.second; + } + + return scratch_off; +} + +void llama_model::pshard_set_backend_maps(const llama_pshard_plan & plan) { + if (plan.maps_cached) { + pimpl->tensor_backend_ids.clear(); + for (const auto & [name, tensor] : tensors_by_name) { + auto it = plan.cached_tensor_bids.find(name); + if (it != plan.cached_tensor_bids.end()) { + pimpl->tensor_backend_ids[tensor] = it->second; + } + } + pimpl->layer_backend_ids = plan.cached_layer_bids; + } else { + pimpl->tensor_backend_ids.clear(); + for (const auto & [name, tensor] : tensors_by_name) { + for (const auto & ov : plan.overrides) { + if (ov.backend_id >= 0 && std::regex_search(name, std::regex(ov.pattern))) { + pimpl->tensor_backend_ids[tensor] = ov.backend_id; + plan.cached_tensor_bids[name] = ov.backend_id; + break; + } + } + } + + pimpl->layer_backend_ids.clear(); + for (const auto & ov : plan.overrides) { + if (ov.backend_id >= 0) { + std::smatch m; + if (std::regex_search(ov.pattern, m, std::regex(R"(blk\\\.(\d+)\\\.)"))) { + int layer = std::stoi(m[1].str()); + pimpl->layer_backend_ids[layer] = ov.backend_id; + } + } + } + plan.cached_layer_bids = pimpl->layer_backend_ids; + plan.maps_cached = true; + } +} + +size_t llama_model::pshard_apply_plan(const llama_pshard_plan & plan, ggml_backend_t gpu) { + pshard_set_backend_maps(plan); + + LLAMA_LOG_DEBUG("%s: rebuilt maps: %zu tensors, %zu layers (cached=%d)\n", + __func__, pimpl->tensor_backend_ids.size(), pimpl->layer_backend_ids.size(), + (int)plan.maps_cached); + + size_t scratch_off = 0; + if (pimpl->dev_preload_buf) { + void * buf_base = ggml_backend_buffer_get_base(pimpl->dev_preload_buf); + + std::unordered_map old_addrs; + old_addrs.reserve(pimpl->weight_preload_map.size()); + for (const auto & [tensor, entry] : pimpl->weight_preload_map) { + old_addrs[tensor] = tensor->data; + } + + if (plan.addrs_cached) { + scratch_off = plan.cached_scratch_off; + for (auto & [tensor, entry] : pimpl->weight_preload_map) { + auto it = plan.cached_weight_offsets.find(std::string(ggml_get_name(tensor))); + if (it != plan.cached_weight_offsets.end()) { + tensor->data = (char *)buf_base + it->second; + tensor->buffer = pimpl->dev_preload_buf; + } else { + tensor->data = entry.cpu_addr; + tensor->buffer = entry.host_buffer; + } + } + } else { + std::vector pinned; + for (auto & [tensor, entry] : pimpl->weight_preload_map) { + auto it = pimpl->tensor_backend_ids.find(tensor); + if (it != pimpl->tensor_backend_ids.end() && it->second == 0) { + pinned.push_back(tensor); + } + } + + auto get_layer = [](const ggml_tensor * t) -> int { + const char * blk = strstr(ggml_get_name(t), "blk."); + return blk ? atoi(blk + 4) : 9999; + }; + auto get_cat = [](const ggml_tensor * t) -> int { + const char * name = ggml_get_name(t); + if (strstr(name, "attn_")) return 0; + if (strstr(name, "exps")) return 4; + if (strstr(name, "ffn_")) return 1; + if (strstr(name, "norm")) return 2; + return 3; + }; + std::sort(pinned.begin(), pinned.end(), + [&](const ggml_tensor * a, const ggml_tensor * b) { + int ca = get_cat(a), cb = get_cat(b); + if (ca != cb) return ca < cb; + int la = get_layer(a), lb = get_layer(b); + if (la != lb) return la < lb; + return strcmp(ggml_get_name(a), ggml_get_name(b)) < 0; + }); + + struct ggml_tallocr talloc = ggml_tallocr_new(pimpl->dev_preload_buf); + size_t buf_size = ggml_backend_buffer_get_size(pimpl->dev_preload_buf); + + for (auto * tensor : pinned) { + size_t tsize = ggml_backend_buffer_get_alloc_size(pimpl->dev_preload_buf, tensor); + if (talloc.offset + tsize > buf_size) break; + tensor->buffer = NULL; + tensor->data = NULL; + ggml_tallocr_alloc(&talloc, tensor); + pimpl->weight_preload_map[tensor].gpu_addr = tensor->data; + } + scratch_off = talloc.offset; + + plan.cached_weight_offsets.clear(); + for (auto * tensor : pinned) { + if (tensor->data) { + plan.cached_weight_offsets[std::string(ggml_get_name(tensor))] = + (size_t)((char *)tensor->data - (char *)buf_base); + } + } + plan.cached_scratch_off = scratch_off; + plan.addrs_cached = true; + + for (auto & [tensor, entry] : pimpl->weight_preload_map) { + auto it = pimpl->tensor_backend_ids.find(tensor); + bool is_pinned = (it != pimpl->tensor_backend_ids.end() && it->second == 0); + if (!is_pinned) { + tensor->data = entry.cpu_addr; + tensor->buffer = entry.host_buffer; + } + } + } + + size_t n_uploaded = 0; + size_t bytes_uploaded = 0; + + for (auto & [tensor, entry] : pimpl->weight_preload_map) { + if (tensor->data == entry.cpu_addr) continue; + void * old = old_addrs[tensor]; + if (old != tensor->data && gpu) { + ggml_backend_tensor_set_async(gpu, tensor, entry.cpu_addr, 0, ggml_nbytes(tensor)); + n_uploaded++; + bytes_uploaded += ggml_nbytes(tensor); + } + } + + LLAMA_LOG_DEBUG("%s: scratch_off=%.2f MiB, %zu uploaded (%.2f MiB), cached=%d\n", + __func__, scratch_off / (1024.0 * 1024.0), + n_uploaded, bytes_uploaded / (1024.0 * 1024.0), (int)plan.addrs_cached); + } + + return scratch_off; +} + 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 +9630,10 @@ llama_model_params llama_model_default_params() { /*.use_extra_bufts =*/ true, /*.no_host =*/ false, /*.no_alloc =*/ false, + /*.pshard =*/ 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..8f75b88ba43d 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -17,6 +17,9 @@ struct llama_cparams; struct llama_ubatch; struct llama_model_loader; +struct weight_preload_entry; +struct llama_pshard_plan; +struct llama_pshard_plan_registry; // available models enum llm_type { @@ -616,6 +619,21 @@ struct llama_model { bool has_tensor_overrides() const; + bool is_pshard() const; + + const std::unordered_map & get_tensor_backend_ids() const; + const std::unordered_map & get_layer_backend_ids() const; + + ggml_backend_buffer_t get_dev_preload_buf() const; + size_t get_dev_preloaded_size() const; + void sync_dev_preload(); + + void pshard_set_backend_maps(const llama_pshard_plan & plan); + size_t pshard_compute_scratch_off(const llama_pshard_plan & plan); + size_t pshard_apply_plan(const llama_pshard_plan & plan, ggml_backend_t gpu = nullptr); + + llama_pshard_plan_registry * get_plan_registry() 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..371d3b83bba2 --- /dev/null +++ b/src/llama-pipe-shard.h @@ -0,0 +1,69 @@ +#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; + void * t1_gpu_addr = nullptr; + void * t2_gpu_addr = nullptr; + size_t alloc_size = 0; + bool is_pinned = false; + bool prefetched_t1 = false; + bool prefetched_t2 = false; + }; + + virtual ~llama_memory_pipe_shard_i() = default; + + virtual const std::vector & get_layers() const = 0; + + // per-batch write_cells (KV only). indexed by stream. + const std::vector> * write_cells = nullptr; + void set_write_cells(const std::vector> * wc) { write_cells = wc; } + + virtual void clear_prefetch() = 0; + + // sched split callbacks + virtual bool prefetch_if_owned(ggml_tensor * t, ggml_backend_t be) = 0; + virtual bool upload_if_owned(ggml_tensor * t, ggml_backend_t be) = 0; + virtual bool download_if_owned(ggml_tensor * t, ggml_backend_t be) = 0; + + // plan switch + virtual void upload_for_switch(int32_t il, ggml_backend_t be) = 0; + virtual void download_for_switch(int32_t il, ggml_backend_t be) = 0; + + virtual void activate_gpu(int32_t il) = 0; + virtual void activate_cpu(int32_t il) = 0; + + // ensure all layers are host-accessible (for state save/load, KV shift, defrag) + virtual void prepare_for_host_access() = 0; + + virtual void pin_layer(int32_t il) = 0; + virtual void unpin_layer(int32_t il) = 0; + + virtual void set_external_addrs(int32_t il, void * a1, void * a2, size_t sz) = 0; + + virtual void refresh_stream_views(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; + + virtual size_t current_pinned_size() const = 0; +}; diff --git a/src/llama-pshard-cache.cpp b/src/llama-pshard-cache.cpp new file mode 100644 index 000000000000..a80374464c26 --- /dev/null +++ b/src/llama-pshard-cache.cpp @@ -0,0 +1,701 @@ +#include "llama-pshard-plan.h" +#include "llama-impl.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +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; + } +} + +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, + bool output_on_gpu, + uint32_t n_attn_pinned) { + GGML_UNUSED(gpu_buft); + + 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); + 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; + + size_t itbo = 0; + + auto emit = [&](const char * pat, ggml_backend_buffer_type_t buft, int32_t bid) { + 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(), host_buft, out_bid); + emit(pat_tok_embd.c_str(), host_buft, layout.cpu); + } + + 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, (strategy == LLAMA_PSHARD_STATIC_FITPARAMS_DENSEPRIO_MOEONLY) ? layout.cpu : layout.shard(il)); + } + emit(patterns_layer[il].c_str(), host_buft, layout.compute); + } else if (il >= il_pin_start && il < il_pin_end) { + emit(patterns_layer[il].c_str(), host_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_STATIC_FITPARAMS_DENSEPRIO_MOEONLY: + emit(patterns_layer[il].c_str(), host_buft, layout.cpu); + break; + 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(), host_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(), host_buft, layout.compute); + } else { + emit(patterns_layer[il].c_str(), host_buft, layout.cpu); + } + break; + default: break; + } + } + } + tensor_buft_overrides[itbo] = { nullptr, nullptr, -1 }; +} + +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 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) { + 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; +} + +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; +} + +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_FITPARAMS_DENSEPRIO_MOEONLY; + 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; + 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; + } + 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) continue; + td.viable = false; + } else { + if (sscanf(s.c_str(), "[tier %zu bs=%u]", &tier_idx, &td.bs) < 2) 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_buf[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_buf, &td.tps, &td.vram_mib) < 4) { + 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) { + 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_buf); + td.strategy = LLAMA_PSHARD_STATIC_FITPARAMS_DENSEPRIO_MOEONLY; + 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; + break; + } + } + } 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; + LLAMA_LOG_INFO("%s: loaded pshard_disabled variant budget=%u MiB cache_ubatch=%u baseline=%.1f MiB from %s\n", + __func__, variant.budget_mib, variant.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 target cache_ubatch=%u in %s\n", + __func__, registry->tier_sizes.size(), registry->budget_mib, registry->cache_ubatch, current_budget_mib, requested_cache_ubatch, 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; +} + +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; +} + +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) { + const std::string cache_path = std::string(path_model) + ".tensor_overrides.pshard_registry"; + + if (!llama_pshard_params_supported(mparams, cparams)) { + mparams->pshard = false; + cparams->pshard = false; + return; + } + + std::vector devs; + uint32_t hp_ngl = 0, hp_nct = 0, hp_nex = 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, GGML_LOG_LEVEL_ERROR); + + if (devs.empty()) { + LLAMA_LOG_WARN("%s: no GPU devices found, disabling pshard\n", __func__); + mparams->pshard = false; + cparams->pshard = false; + return; + } + + const uint32_t n_layers = hp_ngl; + const size_t vram_free = (max_vram_mb > 0) ? max_vram_mb * 1024ULL * 1024ULL : dmds[0].free; + + if (mparams->max_vram_alloc == 0 && vram_free > 0) { + mparams->max_vram_alloc = vram_free / (1024 * 1024); + } + + 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(); + + LLAMA_LOG_INFO("%s: probe: %u layers, %.1f MiB VRAM free\n", + __func__, n_layers, vram_free / (1024.0 * 1024.0)); + + auto * registry = mparams->pshard_registry; + if (!registry) { + LLAMA_LOG_ERROR("%s: pshard_registry is null\n", __func__); + mparams->pshard = false; + cparams->pshard = false; + return; + } + registry->budget_mib = pshard_bytes_to_mib_ceil(vram_free); + const uint32_t runtime_n_batch = std::min(cparams->n_ctx, cparams->n_batch); + const uint32_t runtime_cache_ubatch = std::min(runtime_n_batch, cparams->n_ubatch == 0 ? runtime_n_batch : cparams->n_ubatch); + registry->cache_ubatch = registry->cache_ubatch ? std::min(cparams->n_ctx, registry->cache_ubatch) : runtime_cache_ubatch; + + int64_t model_file_size = 0; + { + FILE * mf = fopen(path_model, "rb"); + if (mf) { +#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); + + if (!pshard_registry_load(registry, fp, cache_path.c_str(), host_buft, vram_free, false)) { + LLAMA_LOG_WARN("%s: no matching plan cache at %s (fingerprint=0x%016" PRIx64 "), disabling pshard\n", + __func__, cache_path.c_str(), fp); + mparams->pshard = false; + cparams->pshard = false; + return; + } + + if (!registry->pshard_disabled) { + LLAMA_LOG_INFO("%s: loaded %zu tier plans from cache (variant budget=%u MiB cache_ubatch=%u)\n", + __func__, registry->tier_sizes.size(), registry->budget_mib, registry->cache_ubatch); + } + + // cached baseline fit for this budget + // use the normal load path + if (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; + cparams->pshard = false; + mparams->n_gpu_layers = n_layers + 1; + tensor_buft_overrides[0] = { nullptr, nullptr, -1 }; + mparams->tensor_buft_overrides = nullptr; + return; + } + + if (registry->cache_ubatch > 0) { + const uint32_t pshard_ubatch = std::min(cparams->n_ctx, registry->cache_ubatch); + cparams->n_batch = pshard_ubatch; + cparams->n_ubatch = pshard_ubatch; + } + + // pick the highest viable tier + size_t default_tier = registry->tier_sizes.size(); + llama_pshard_plan * best = nullptr; + for (size_t t = registry->tier_sizes.size(); t-- > 0; ) { + llama_pshard_plan * candidate = registry->get_best(t); + if (candidate && candidate->is_viable) { + default_tier = t; + best = candidate; + break; + } + } + if (!best) { + LLAMA_LOG_WARN("%s: no viable plan in cache, disabling pshard\n", __func__); + mparams->pshard = false; + cparams->pshard = false; + return; + } + if (default_tier < registry->tier_sizes.size() - 1) { + LLAMA_LOG_INFO("%s: highest tier (bs=%u) not viable, falling back to bs=%u\n", + __func__, registry->tier_sizes.back(), registry->tier_sizes[default_tier]); + } + + if (best->n_pinned > n_layers) { + LLAMA_LOG_WARN("%s: cache stale: n_pinned=%u > n_layers=%u, regenerate cache\n", + __func__, best->n_pinned, n_layers); + mparams->pshard = false; + cparams->pshard = false; + return; + } + + 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); + // baseline plans use the cached fit params dense only placement + // regenerating them from n_pinned and overflow would lose that placement + if (best->strategy == LLAMA_PSHARD_STATIC_FITPARAMS_DENSEPRIO_MOEONLY) { + thread_local std::vector shard_none_patterns; + shard_none_patterns.clear(); + shard_none_patterns.reserve(best->overrides.size()); + for (const auto & ov : best->overrides) { + shard_none_patterns.push_back(ov.pattern); + } + for (size_t i = 0; i < best->overrides.size(); i++) { + tensor_buft_overrides[i].pattern = shard_none_patterns[i].c_str(); + tensor_buft_overrides[i].buft = best->overrides[i].buft; + tensor_buft_overrides[i].backend_id = best->overrides[i].backend_id; + } + tensor_buft_overrides[best->overrides.size()] = { nullptr, nullptr, -1 }; + } else { + llama_pshard_generate_overrides( + best->n_pinned, n_layers, host_buft, host_buft, + tensor_buft_overrides, + (llama_layer_fraction)best->overflow, + best->strategy, layout, + best->pin_from_back, best->output_on_gpu, best->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->tensor_buft_overrides = tensor_buft_overrides; + mparams->n_gpu_layers = n_layers + 1; + + LLAMA_LOG_INFO("%s: plan: %s, n_pinned=%u/%u, vram=%zu MiB, n_gpu_layers=%d\n", + __func__, llama_pshard_strategy_name(best->strategy), + best->n_pinned, n_layers, mparams->max_vram_alloc, mparams->n_gpu_layers); +} diff --git a/src/llama-pshard-plan.h b/src/llama-pshard-plan.h new file mode 100644 index 000000000000..9c9f1c4429e3 --- /dev/null +++ b/src/llama-pshard-plan.h @@ -0,0 +1,233 @@ +#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_STATIC_FITPARAMS_DENSEPRIO_MOEONLY = 0, + LLAMA_PSHARD_GPUONLY_LAYERPIN_LAYERSTREAM = 1, + LLAMA_PSHARD_GPUONLY_ATTNPIN_FFNSTREAM = 2, + LLAMA_PSHARD_DYNAMIC_FFNCPU_ATTNSTREAM = 3, + LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS = 4, + LLAMA_PSHARD_COUNT +}; + +// strategy name for logging +inline const char * llama_pshard_strategy_name(llama_pshard_strategy s) { + static const char * const names[] = { + "STATIC_FITPARAMS_DENSEPRIO_MOEONLY", + "GPUONLY_LAYERPIN_LAYERSTREAM", + "GPUONLY_ATTNPIN_FFNSTREAM", + "DYNAMIC_FFNCPU_ATTNSTREAM", + "STATIC_ATTNPRIO_ALLMODELS", + }; + return (s >= 0 && s < LLAMA_PSHARD_COUNT) ? names[s] : "UNKNOWN"; +} + +// 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; +}; + +// saved allocator and backend ids for plan switches +struct llama_pshard_alloc_state { + std::vector node_allocs; + std::vector leaf_allocs; + std::vector node_backend_ids; + std::vector leaf_backend_ids; + int n_nodes = 0; + int n_leafs = 0; + bool valid = false; +}; + +struct llama_pshard_plan { + llama_pshard_strategy strategy = LLAMA_PSHARD_STATIC_FITPARAMS_DENSEPRIO_MOEONLY; + 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; + + // cached maps and offsets from first apply + mutable std::unordered_map cached_tensor_bids; + mutable std::unordered_map cached_layer_bids; + mutable std::unordered_map cached_weight_offsets; + mutable size_t cached_scratch_off = 0; + mutable bool maps_cached = false; + mutable bool addrs_cached = false; + + mutable llama_pshard_alloc_state alloc_state; +}; + +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); + +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); + +struct llama_device_memory_data { + int64_t total; + int64_t free; + llama_memory_breakdown_data mb; +}; + +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, + 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_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(); + + // 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; + } + + // pick the prefill ubatch with the lowest predicted ttft + // use max_ubatch when TPS data is missing + uint32_t find_optimal_ubatch(uint32_t n_prompt, uint32_t max_ubatch) const { + uint32_t best_ub = max_ubatch; + double best_time = 1e30; + + for (size_t t = 0; t < tier_sizes.size(); t++) { + uint32_t ts = tier_sizes[t]; + if (ts < 512 || ts > max_ubatch) continue; + + const auto & plan = best_plans[t]; + if (!plan.is_viable || plan.tps <= 0.0f) continue; + + double per_iter = (double)ts / (double)plan.tps; + uint32_t n_iters = (n_prompt + ts - 1) / ts; + double total = n_iters * per_iter; + + if (total < best_time) { + best_time = total; + best_ub = ts; + } + } + + return best_ub; + } +}; diff --git a/src/llama.cpp b/src/llama.cpp index 484372d8d106..e4162f9dc5bf 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" @@ -46,13 +47,7 @@ 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) { @@ -436,21 +431,24 @@ static void llama_params_fit_impl( il0 += ngl_per_device[id].n_full(); for (uint32_t il = il0; il < il0 + ngl_per_device[id].n_part; il++) { if (itbo + 1 >= ntbo) { - tensor_buft_overrides[itbo].pattern = nullptr; - tensor_buft_overrides[itbo].buft = nullptr; + tensor_buft_overrides[itbo].pattern = nullptr; + tensor_buft_overrides[itbo].buft = nullptr; + tensor_buft_overrides[itbo].backend_id = -1; itbo++; mparams.tensor_buft_overrides = tensor_buft_overrides; throw 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 = get_overflow_pattern(il, il == il0 ? ngl_per_device[id].overflow_type : LAYER_FRACTION_MOE); + tensor_buft_overrides[itbo].buft = il == il0 ? overflow_bufts[id] : ggml_backend_cpu_buffer_type(); + tensor_buft_overrides[itbo].backend_id = -1; itbo++; } il0 += ngl_per_device[id].n_part; } - tensor_buft_overrides[itbo].pattern = nullptr; - tensor_buft_overrides[itbo].buft = nullptr; + tensor_buft_overrides[itbo].pattern = nullptr; + tensor_buft_overrides[itbo].buft = nullptr; + tensor_buft_overrides[itbo].backend_id = -1; itbo++; mparams.tensor_buft_overrides = tensor_buft_overrides; }; @@ -486,8 +484,8 @@ 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__); @@ -508,7 +506,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; } @@ -876,6 +874,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; } 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(); } From c64c73db6dc3745a88ad3da14e2c17933c233435 Mon Sep 17 00:00:00 2001 From: Aditya Ukarande Date: Fri, 29 May 2026 12:35:38 -0700 Subject: [PATCH 2/2] Fix pshard prefetch tracking, state remaps, and common-weight preload --- common/arg.cpp | 6 +- common/common.cpp | 27 +- ggml/include/ggml-backend.h | 1 + ggml/src/ggml-backend.cpp | 172 ++++++++----- include/llama.h | 4 +- src/llama-context-pshard.cpp | 119 ++++++--- src/llama-context.cpp | 207 ++++++++++----- src/llama-context.h | 24 +- src/llama-kv-cache.cpp | 8 - src/llama-model-loader.cpp | 327 ++++++++++++++++++++++-- src/llama-model-loader.h | 17 +- src/llama-model.cpp | 411 ++++++++++++++++++++++++++++-- src/llama-model.h | 8 + src/llama-pshard-cache.cpp | 172 +++++++++---- src/llama-pshard-plan.h | 36 ++- tools/llama-bench/llama-bench.cpp | 163 ++++++++++-- 16 files changed, 1381 insertions(+), 321 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 10d705ca9c22..402df69003f4 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1573,7 +1573,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex )); add_opt(common_arg( {"-mva", "--max-vram-alloc"}, "N", - "VRAM budget in MB for pshard (0 = use actual free VRAM)", + "VRAM budget in MB for pshard (0 = use actual free VRAM minus --fit-target)", [](common_params & params, int value) { params.max_vram_alloc = value; } @@ -2448,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 8de41eb963e4..76dd44c971fd 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1144,17 +1144,30 @@ common_init_result::common_init_result(common_params & params) : auto mparams = common_model_params_to_llama(params); auto cparams = common_context_params_to_llama(params); + auto fit_params = [&]() { + LOG_INF("%s: fitting params to device memory, for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on\n", __func__); + llama_params_fit(params.model.path.c_str(), &mparams, &cparams, + params.tensor_split, + params.tensor_buft_overrides.data(), + params.fit_params_target.data(), + params.fit_params_min_ctx, + params.verbosity >= 4 ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); + }; + if (params.pshard) { LOG_INF("%s: pshard enabled, probing and loading plan cache\n", __func__); params.tensor_buft_overrides.resize(4096); - const uint32_t tier_max = std::min(std::max(cparams.n_batch, (uint32_t)16384), cparams.n_ctx); - mparams.pshard_registry = llama_pshard_registry_create(tier_max, cparams.n_seq_max); + mparams.pshard_registry = llama_pshard_registry_create(params.pshard_tier_max, cparams.n_seq_max); + const size_t fit_target_mb = params.fit_params_target.empty() ? 0 : params.fit_params_target[0] / (1024 * 1024); llama_params_fit_pshard(params.model.path.c_str(), &mparams, &cparams, - params.tensor_buft_overrides.data(), params.max_vram_alloc); + params.tensor_buft_overrides.data(), params.max_vram_alloc, fit_target_mb); if (!mparams.pshard) { LOG_WRN("%s: pshard not active for this configuration\n", __func__); llama_pshard_registry_free(mparams.pshard_registry); mparams.pshard_registry = nullptr; + if (params.fit_params) { + fit_params(); + } } else { params.n_batch = (int32_t) cparams.n_batch; params.n_ubatch = (int32_t) cparams.n_ubatch; @@ -1162,13 +1175,7 @@ common_init_result::common_init_result(common_params & params) : __func__, cparams.n_ubatch); } } else if (params.fit_params) { - LOG_INF("%s: fitting params to device memory, for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on\n", __func__); - llama_params_fit(params.model.path.c_str(), &mparams, &cparams, - params.tensor_split, - params.tensor_buft_overrides.data(), - params.fit_params_target.data(), - params.fit_params_min_ctx, - params.verbosity >= 4 ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); + fit_params(); } llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams); diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 6e0e0f875ab5..0bee90753beb 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -332,6 +332,7 @@ extern "C" { GGML_API size_t ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend); GGML_API void ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend); + GGML_API void ggml_backend_sched_set_tensor_backend_hint(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend); GGML_API ggml_backend_t ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node); // Split graph without allocating it diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index d4307ea0cec3..d9bc765988c5 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -833,6 +833,8 @@ struct ggml_backend_sched { 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; @@ -962,7 +964,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)) { @@ -1535,22 +1537,29 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra graph_copy->nodes[graph_copy->n_nodes++] = prealloc; } - // prefetch: reserve next GPU split's prefetch destinations (race-protect copy stream vs current compute) + // 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; - for (int ni = i + 1; ni < sched->n_splits; ni++) { - struct ggml_backend_sched_split * candidate = &sched->splits[ni]; - if (sched->copy_backends[candidate->backend_id] == NULL) continue; - 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 (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 (has_host_weights) { next_gpu = candidate; break; } } if (next_gpu != NULL) { for (int j = 0; j < next_gpu->n_inputs; j++) { @@ -1698,7 +1707,8 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s ggml_tensor * prev_ids_tensor = nullptr; std::vector ids; std::vector used_ids; - bool next_weights_prefetched = false; + int prefetched_split_id = -1; + int prefetched_backend_id = -1; for (int split_id = 0; split_id < sched->n_splits; split_id++) { struct ggml_backend_sched_split * split = &splits[split_id]; @@ -1712,64 +1722,87 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s split_backend = sched->backends[split_backend_id]; } - bool weights_prefetched = next_weights_prefetched; - next_weights_prefetched = false; + const bool weights_prefetched = + prefetched_split_id == split_id && + prefetched_backend_id == copy_backend_id; - if (sched->prefetch_weights) { - ggml_backend_t copy_backend = sched->copy_backends[split_backend_id]; - if (copy_backend != NULL || weights_prefetched) { - if (weights_prefetched && sched->copy_events[copy_backend_id] != NULL) { - ggml_backend_event_wait(split_backend, sched->copy_events[copy_backend_id]); - } + if (prefetched_split_id == split_id && !weights_prefetched) { + prefetched_split_id = -1; + prefetched_backend_id = -1; + } - { - struct ggml_backend_sched_split * next_gpu = NULL; - for (int nid = split_id + 1; nid < sched->n_splits; nid++) { - struct ggml_backend_sched_split * candidate = &splits[nid]; - if (sched->copy_backends[candidate->backend_id] == NULL) continue; - 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 (weights_prefetched && sched->copy_events[copy_backend_id] != NULL) { + ggml_backend_event_wait(split_backend, sched->copy_events[copy_backend_id]); + } + if (weights_prefetched) { + prefetched_split_id = -1; + prefetched_backend_id = -1; + } + + auto prefetch_next_split = [&]() { + if (!sched->prefetch_weights || prefetched_split_id >= 0) { + return; + } + + struct ggml_backend_sched_split * next_gpu = NULL; + int next_gpu_id = -1; + if (split_id + 1 < sched->n_splits) { + const int nid = split_id + 1; + struct ggml_backend_sched_split * candidate = &splits[nid]; + // 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; break; } } - if (next_gpu != NULL) { - ggml_backend_t next_copy = sched->copy_backends[next_gpu->backend_id]; + if (has_host_weights) { + next_gpu = candidate; + next_gpu_id = nid; + } + } + } + if (next_gpu == NULL) { + return; + } - if (sched->compute_events[split_backend_id] != NULL) { - ggml_backend_event_wait(next_copy, sched->compute_events[split_backend_id]); - } + ggml_backend_t next_copy = sched->copy_backends[next_gpu->backend_id]; - for (int input_id = 0; input_id < next_gpu->n_inputs; input_id++) { - struct ggml_tensor * next_input = next_gpu->inputs[input_id]; - 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)) { - struct ggml_tensor * input_cpy = tensor_copy(next_input, next_gpu->backend_id, sched->cur_copy); - ggml_backend_tensor_set_async(next_copy, input_cpy, next_input->data, 0, ggml_nbytes(next_input)); - next_weights_prefetched = true; - } - } + if (sched->compute_events[split_backend_id] != NULL) { + ggml_backend_event_wait(next_copy, sched->compute_events[split_backend_id]); + } - if (sched->split_prefetch_cb != NULL && next_gpu->n_writeback > 0) { - for (int w = 0; w < next_gpu->n_writeback; w++) { - sched->split_prefetch_cb(next_gpu->writeback[w], next_copy, sched->split_cb_user_data); - } - next_weights_prefetched = true; - } + bool did_prefetch = false; + for (int input_id = 0; input_id < next_gpu->n_inputs; input_id++) { + struct ggml_tensor * next_input = next_gpu->inputs[input_id]; + 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)) { + struct ggml_tensor * input_cpy = tensor_copy(next_input, next_gpu->backend_id, sched->cur_copy); + ggml_backend_tensor_set_async(next_copy, input_cpy, next_input->data, 0, ggml_nbytes(next_input)); + did_prefetch = true; + } + } - if (next_weights_prefetched) { - ggml_backend_event_record(sched->copy_events[next_gpu->backend_id], next_copy); - } - } + if (sched->split_prefetch_cb != NULL && next_gpu->n_writeback > 0) { + for (int w = 0; w < next_gpu->n_writeback; w++) { + sched->split_prefetch_cb(next_gpu->writeback[w], next_copy, sched->split_cb_user_data); } + did_prefetch = true; } - } + + if (did_prefetch) { + ggml_backend_event_record(sched->copy_events[next_gpu->backend_id], next_copy); + prefetched_split_id = next_gpu_id; + prefetched_backend_id = next_gpu->backend_id; + } + }; // copy the input tensors to the split backend for (int input_id = 0; input_id < split->n_inputs; input_id++) { @@ -1894,7 +1927,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } else { // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface - if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) { + bool copied = split_backend->iface.cpy_tensor_async && + split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy); + if (!copied) { ggml_backend_synchronize(input_backend); if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); @@ -1919,6 +1954,8 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } + prefetch_next_split(); + if (!sched->callback_eval) { enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph); if (ec != GGML_STATUS_SUCCESS) { @@ -1977,6 +2014,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s ggml_backend_event_record(sched->events[split_backend_id][sched->cur_copy], split_backend); } } + } return GGML_STATUS_SUCCESS; @@ -2415,6 +2453,14 @@ void ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct gg // 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) { GGML_ASSERT(sched); int backend_index = tensor_backend_id(node); diff --git a/include/llama.h b/include/llama.h index 60368e8bbf77..6044f12c7180 100644 --- a/include/llama.h +++ b/include/llama.h @@ -321,6 +321,7 @@ extern "C" { 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 @@ -549,7 +550,8 @@ extern "C" { 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 + 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. diff --git a/src/llama-context-pshard.cpp b/src/llama-context-pshard.cpp index a459eb86eb71..d064facb7299 100644 --- a/src/llama-context-pshard.cpp +++ b/src/llama-context-pshard.cpp @@ -13,7 +13,10 @@ #include "ggml-alloc.h" #include +#include +#include #include +#include #include namespace { @@ -101,6 +104,7 @@ namespace { } } } + } // namespace void pshard_assign_tensors( @@ -108,36 +112,17 @@ void pshard_assign_tensors( const llama_model & model, llama_memory_i * memory, const std::vector & backends, - const pshard_dev_layout & layout, - ggml_cgraph * gf) { + 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(sched, tensor, backends[bid].get()); - } - } - - // catches compute nodes never named via cb() - // (e.g. wo matmul inside build_attn) - if (gf) { - const int n_nodes = ggml_graph_n_nodes(gf); - for (int i = 0; i < n_nodes; i++) { - ggml_tensor * cur = ggml_graph_node(gf, i); - if (cur->view_src) continue; - 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 && it->second < (int32_t)backends.size()) { - ggml_backend_sched_set_tensor_backend(sched, cur, backends[it->second].get()); - break; - } - } + if (bid >= 0 && bid < (int32_t) backends.size()) { + ggml_backend_sched_set_tensor_backend_hint(sched, tensor, backends[bid].get()); } } if (memory) { - const auto & lbids = model.get_layer_backend_ids(); for (auto * ps : memory->get_pipe_shards()) { ps->assign_tensors(sched, lbids, backends, layout); } @@ -346,6 +331,7 @@ void llama_context::pshard_reserve_and_save(const llama_pshard_plan & plan) { const uint32_t n_seqs = cparams.n_seq_max; const uint32_t n_tokens = plan.batch_size; + const uint32_t n_outputs = n_tokens; // start with unconstrained scratch packing ggml_backend_t gpu = backends[pshard_layout.compute].get(); @@ -361,7 +347,7 @@ void llama_context::pshard_reserve_and_save(const llama_pshard_plan & plan) { ggml_backend_sched_set_alloc_range(sched.get(), gpu, scratch_off, SIZE_MAX/2); } - auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get()); + auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs, mctx.get()); if (gf && external_buf) { const int n_chunks = ggml_backend_sched_get_n_chunks(sched.get(), gpu); @@ -378,7 +364,7 @@ void llama_context::pshard_reserve_and_save(const llama_pshard_plan & plan) { __func__, chunk0_used / (1024.0 * 1024.0), scratch_avail / (1024.0 * 1024.0)); ggml_backend_sched_set_alloc_range(sched.get(), gpu, scratch_off, scratch_avail); - gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get()); + gf = graph_reserve(n_tokens, n_seqs, n_outputs, mctx.get()); } if (!gf) { @@ -423,7 +409,7 @@ void llama_context::pshard_save_alloc_state(const llama_pshard_plan & plan) { sched_n_nodes, sched_n_leafs); } -void llama_context::pshard_warmup_plans() { +void llama_context::pshard_warmup_plan_reserves() { auto * registry = model.get_plan_registry(); if (!registry) return; @@ -452,12 +438,6 @@ void llama_context::pshard_warmup_plans() { } } - llama_pshard_plan * initial = registry->get_best(0); - if (initial) { - pshard_apply_plan(*initial); - pshard_active_plan = initial; - } - LLAMA_LOG_INFO("%s: tier summary:\n", __func__); for (size_t t = 0; t < registry->tier_sizes.size(); t++) { auto & plan = registry->best_plans[t]; @@ -487,7 +467,29 @@ void llama_context::pshard_warmup_plans() { __func__, registry->tier_sizes.size(), (t1 - t0) / 1000.0); } -void llama_context::pshard_switch_plan(const llama_pshard_plan & old_plan, const llama_pshard_plan & new_plan) { +void llama_context::pshard_apply_initial_plan() { + auto * registry = model.get_plan_registry(); + if (!registry) return; + + size_t initial_tier = registry->tier_index(16); + llama_pshard_plan * initial = registry->get_best(initial_tier); + + if (!initial) { + initial = registry->get_best(0); + } + + if (initial) { + pshard_apply_plan(*initial); + pshard_active_plan = initial; + } +} + +void llama_context::pshard_switch_plan( + const llama_pshard_plan & old_plan, + const llama_pshard_plan & new_plan, + size_t old_tier, + size_t new_tier, + uint32_t n_tokens) { ggml_backend_t gpu = backends[pshard_layout.compute].get(); auto pipe_shards = memory->get_pipe_shards(); @@ -539,8 +541,14 @@ void llama_context::pshard_switch_plan(const llama_pshard_plan & old_plan, const } } - LLAMA_LOG_DEBUG("%s: %s (n_pinned=%u) -> %s (n_pinned=%u) | down=%d skip=%d up=%d skip=%d\n", + auto * registry = model.get_plan_registry(); + auto tier_bs = [&](size_t tier) -> uint32_t { + return registry && tier < registry->tier_sizes.size() ? registry->tier_sizes[tier] : 0; + }; + + LLAMA_LOG_DEBUG("%s: tokens=%u tier %zu(bs=%u) -> %zu(bs=%u): %s (n_pinned=%u) -> %s (n_pinned=%u) | down=%d skip=%d up=%d skip=%d\n", __func__, + n_tokens, old_tier, tier_bs(old_tier), new_tier, tier_bs(new_tier), llama_pshard_strategy_name(old_plan.strategy), old_plan.n_pinned, llama_pshard_strategy_name(new_plan.strategy), new_plan.n_pinned, n_down, n_skip_down, n_up, n_skip_up); @@ -572,7 +580,41 @@ void llama_context::pshard_reapply_active_plan() { plan.alloc_state.leaf_backend_ids.data(), (int)plan.alloc_state.leaf_backend_ids.size()); } +bool llama_context::pshard_prepare_host_access() { + if (!cparams.pshard || !memory) { + return false; + } + + auto pipe_shards = memory->get_pipe_shards(); + if (pipe_shards.empty()) { + return false; + } + + for (auto * ps : pipe_shards) { + ps->prepare_for_host_access(); + } + + pshard_memory_dirty = true; + return true; +} + +void llama_context::pshard_restore_after_host_access() { + if (!pshard_memory_dirty) { + return; + } + + if (pshard_active_plan) { + pshard_apply_plan(*pshard_active_plan); + } + + pshard_memory_dirty = false; +} + void llama_context::pshard_maybe_switch(uint32_t n_tokens) { + if (pshard_memory_dirty) { + pshard_restore_after_host_access(); + } + auto * registry = model.get_plan_registry(); if (!registry) return; @@ -582,7 +624,14 @@ void llama_context::pshard_maybe_switch(uint32_t n_tokens) { if (best != pshard_active_plan) { if (pshard_active_plan) { - pshard_switch_plan(*pshard_active_plan, *best); + size_t old_tier = registry->tier_sizes.size(); + for (size_t i = 0; i < registry->best_plans.size(); i++) { + if (®istry->best_plans[i] == pshard_active_plan) { + old_tier = i; + break; + } + } + pshard_switch_plan(*pshard_active_plan, *best, old_tier, tier, n_tokens); } else { pshard_apply_plan(*best); } diff --git a/src/llama-context.cpp b/src/llama-context.cpp index bfe126cddf03..97ebab317cf7 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -158,15 +158,16 @@ llama_context::llama_context( cparams.flash_attn = params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED; cparams.auto_fa = params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO; - // pshard remaps kv by layer so auto_fa cannot use device matching - if (cparams.pshard) { - cparams.auto_fa = false; - } - cparams.fused_gdn_ar = true; 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; @@ -374,14 +375,15 @@ llama_context::llama_context( } if (cparams.pshard && !model.hparams.no_alloc) { - const_cast(model).sync_dev_preload(); pshard_pack_cache_region(); } sched_reserve(); if (cparams.pshard && !model.hparams.no_alloc) { - pshard_warmup_plans(); + pshard_warmup_plan_reserves(); + const_cast(model).sync_dev_preload(); + pshard_apply_initial_plan(); } if (!cparams.flash_attn) { @@ -1243,6 +1245,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll } n_reused++; + } else { res->reset(); @@ -1264,7 +1267,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll } if (cparams.pshard) { - pshard_assign_tensors(sched.get(), model, memory.get(), backends, pshard_layout, gf); + pshard_assign_tensors(sched.get(), model, memory.get(), backends, pshard_layout); } if (!ggml_backend_sched_alloc_graph(sched.get(), gf)) { @@ -2207,7 +2210,7 @@ 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, gf); + pshard_assign_tensors(sched.get(), model, memory.get(), backends, pshard_layout); } // initialize scheduler with the specified graph @@ -2281,16 +2284,43 @@ ggml_status llama_context::graph_compute( llm_graph_cb llama_context::graph_get_cb() const { // 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 [&, last_weight_bid, last_il](const llama_ubatch & ubatch, ggml_tensor * cur, const char * name, int il) mutable { + 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); } - if (cparams.pshard) { + // 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; @@ -2533,7 +2563,7 @@ class llama_io_read_file : public llama_io_read_i { size_t llama_context::state_get_size() { llama_io_write_dummy io; try { - return state_write_data(io); + return state_write_data(io, false); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error getting state size: %s\n", __func__, err.what()); return 0; @@ -2543,7 +2573,7 @@ size_t llama_context::state_get_size() { size_t llama_context::state_get_data(uint8_t * dst, size_t size) { llama_io_write_buffer io(dst, size); try { - return state_write_data(io); + return state_write_data(io, true); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error saving state: %s\n", __func__, err.what()); return 0; @@ -2553,7 +2583,7 @@ size_t llama_context::state_get_data(uint8_t * dst, size_t size) { size_t llama_context::state_set_data(const uint8_t * src, size_t size) { llama_io_read_buffer io(src, size); try { - return state_read_data(io); + return state_read_data(io, true); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what()); return 0; @@ -2563,7 +2593,7 @@ size_t llama_context::state_set_data(const uint8_t * src, size_t size) { size_t llama_context::state_seq_get_size(llama_seq_id seq_id, llama_state_seq_flags flags) { llama_io_write_dummy io; try { - return state_seq_write_data(io, seq_id, flags); + return state_seq_write_data(io, seq_id, flags, false); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error getting state size: %s\n", __func__, err.what()); return 0; @@ -2573,7 +2603,7 @@ size_t llama_context::state_seq_get_size(llama_seq_id seq_id, llama_state_seq_fl size_t llama_context::state_seq_get_data(llama_seq_id seq_id, uint8_t * dst, size_t size, llama_state_seq_flags flags) { llama_io_write_buffer io(dst, size); try { - return state_seq_write_data(io, seq_id, flags); + return state_seq_write_data(io, seq_id, flags, true); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error saving state: %s\n", __func__, err.what()); return 0; @@ -2583,7 +2613,7 @@ size_t llama_context::state_seq_get_data(llama_seq_id seq_id, uint8_t * dst, siz size_t llama_context::state_seq_set_data(llama_seq_id seq_id, const uint8_t * src, size_t size, llama_state_seq_flags flags) { llama_io_read_buffer io(src, size); try { - return state_seq_read_data(io, seq_id, flags); + return state_seq_read_data(io, seq_id, flags, true); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what()); return 0; @@ -2622,7 +2652,7 @@ bool llama_context::state_load_file(const char * filepath, llama_token * tokens_ const size_t n_state_size_cur = file.size() - file.tell(); llama_io_read_file io( &file); - const size_t n_read = state_read_data(io); + const size_t n_read = state_read_data(io, true); if (n_read != n_state_size_cur) { LLAMA_LOG_ERROR("%s: did not read all of the session file data! size %zu, got %zu\n", __func__, n_state_size_cur, n_read); @@ -2645,7 +2675,7 @@ bool llama_context::state_save_file(const char * filepath, const llama_token * t // save the context state using stream saving llama_io_write_file io(&file); - state_write_data(io); + state_write_data(io, true); return true; } @@ -2681,7 +2711,7 @@ size_t llama_context::state_seq_load_file(llama_seq_id seq_id, const char * file { const size_t state_size = file.size() - file.tell(); llama_io_read_file io(&file); - const size_t nread = state_seq_read_data(io, seq_id, 0); + const size_t nread = state_seq_read_data(io, seq_id, 0, true); if (!nread) { LLAMA_LOG_ERROR("%s: failed to restore sequence state\n", __func__); return 0; @@ -2705,7 +2735,7 @@ size_t llama_context::state_seq_save_file(llama_seq_id seq_id, const char * file // save the context state using stream saving llama_io_write_file io(&file); - state_seq_write_data(io, seq_id, 0); + state_seq_write_data(io, seq_id, 0, true); const size_t res = file.tell(); GGML_ASSERT(res == sizeof(uint32_t) * 3 + sizeof(llama_token) * n_token_count + io.n_bytes()); @@ -2713,72 +2743,123 @@ size_t llama_context::state_seq_save_file(llama_seq_id seq_id, const char * file return res; } -size_t llama_context::state_write_data(llama_io_write_i & io) { - LLAMA_LOG_DEBUG("%s: writing state\n", __func__); +size_t llama_context::state_write_data(llama_io_write_i & io, bool pshard_host_access) { + const bool pshard_restore = pshard_host_access && pshard_prepare_host_access(); - // write model info - { - LLAMA_LOG_DEBUG("%s: - writing model info\n", __func__); + try { + LLAMA_LOG_DEBUG("%s: writing state\n", __func__); - const std::string arch_str = llm_arch_name(model.arch); - io.write_string(arch_str); - // TODO: add more model-specific info which should prevent loading the session file if not identical - } + // write model info + { + LLAMA_LOG_DEBUG("%s: - writing model info\n", __func__); - if (memory != nullptr) { - LLAMA_LOG_DEBUG("%s: - writing memory module\n", __func__); - memory->state_write(io); - } + const std::string arch_str = llm_arch_name(model.arch); + io.write_string(arch_str); + // TODO: add more model-specific info which should prevent loading the session file if not identical + } + + if (memory != nullptr) { + LLAMA_LOG_DEBUG("%s: - writing memory module\n", __func__); + memory->state_write(io); + } - return io.n_bytes(); + const size_t res = io.n_bytes(); + if (pshard_restore) { + pshard_restore_after_host_access(); + } + return res; + } catch (...) { + if (pshard_restore) { + pshard_restore_after_host_access(); + } + throw; + } } -size_t llama_context::state_read_data(llama_io_read_i & io) { - LLAMA_LOG_DEBUG("%s: reading state\n", __func__); +size_t llama_context::state_read_data(llama_io_read_i & io, bool pshard_host_access) { + const bool pshard_restore = pshard_host_access && pshard_prepare_host_access(); - // read model info - { - LLAMA_LOG_DEBUG("%s: - reading model info\n", __func__); + try { + LLAMA_LOG_DEBUG("%s: reading state\n", __func__); + + // read model info + { + LLAMA_LOG_DEBUG("%s: - reading model info\n", __func__); - const std::string cur_arch_str = llm_arch_name(model.arch); + const std::string cur_arch_str = llm_arch_name(model.arch); - std::string arch_str; - io.read_string(arch_str); - if (cur_arch_str != arch_str) { - throw std::runtime_error(format("wrong model arch: '%s' instead of '%s'", arch_str.c_str(), cur_arch_str.c_str())); + std::string arch_str; + io.read_string(arch_str); + if (cur_arch_str != arch_str) { + throw std::runtime_error(format("wrong model arch: '%s' instead of '%s'", arch_str.c_str(), cur_arch_str.c_str())); + } + // TODO: add more info which needs to be identical but which is not verified otherwise } - // TODO: add more info which needs to be identical but which is not verified otherwise - } - if (memory) { - LLAMA_LOG_DEBUG("%s: - reading memory module\n", __func__); + if (memory) { + LLAMA_LOG_DEBUG("%s: - reading memory module\n", __func__); - memory->state_read(io); - } + memory->state_read(io); + } - return io.n_bytes(); + const size_t res = io.n_bytes(); + if (pshard_restore) { + pshard_restore_after_host_access(); + } + return res; + } catch (...) { + if (pshard_restore) { + pshard_restore_after_host_access(); + } + throw; + } } -size_t llama_context::state_seq_write_data(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { +size_t llama_context::state_seq_write_data(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags, bool pshard_host_access) { GGML_UNUSED(seq_id); - if (memory) { - memory->state_write(io, seq_id, flags); - } + const bool pshard_restore = pshard_host_access && pshard_prepare_host_access(); + + try { + if (memory) { + memory->state_write(io, seq_id, flags); + } - return io.n_bytes(); + const size_t res = io.n_bytes(); + if (pshard_restore) { + pshard_restore_after_host_access(); + } + return res; + } catch (...) { + if (pshard_restore) { + pshard_restore_after_host_access(); + } + throw; + } } -size_t llama_context::state_seq_read_data(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { +size_t llama_context::state_seq_read_data(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags, bool pshard_host_access) { GGML_UNUSED(seq_id); - if (memory) { - memory->state_read(io, seq_id, flags); - } + const bool pshard_restore = pshard_host_access && pshard_prepare_host_access(); - return io.n_bytes(); -} + try { + if (memory) { + memory->state_read(io, seq_id, flags); + } + const size_t res = io.n_bytes(); + if (pshard_restore) { + pshard_restore_after_host_access(); + } + return res; + } catch (...) { + if (pshard_restore) { + pshard_restore_after_host_access(); + } + throw; + } +} // // perf // diff --git a/src/llama-context.h b/src/llama-context.h index f1ca563965b9..df1a657a87df 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -250,17 +250,25 @@ struct llama_context { void pshard_reapply_active_plan(); void pshard_reserve_and_save(const llama_pshard_plan & plan); void pshard_save_alloc_state(const llama_pshard_plan & plan); - void pshard_warmup_plans(); - void pshard_switch_plan(const llama_pshard_plan & old_plan, const llama_pshard_plan & new_plan); + void pshard_warmup_plan_reserves(); + void pshard_apply_initial_plan(); + void pshard_switch_plan( + const llama_pshard_plan & old_plan, + const llama_pshard_plan & new_plan, + size_t old_tier, + size_t new_tier, + uint32_t n_tokens); void pshard_maybe_switch(uint32_t n_tokens); void pshard_update_write_cells(llama_memory_context_i * mctx); + bool pshard_prepare_host_access(); + void pshard_restore_after_host_access(); // TODO: read/write lora adapters and cvec - size_t state_write_data(llama_io_write_i & io); - size_t state_read_data (llama_io_read_i & io); + size_t state_write_data(llama_io_write_i & io, bool pshard_host_access); + size_t state_read_data (llama_io_read_i & io, bool pshard_host_access); - size_t state_seq_write_data(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags); - size_t state_seq_read_data (llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags); + size_t state_seq_write_data(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags, bool pshard_host_access); + size_t state_seq_read_data (llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags, bool pshard_host_access); // // members @@ -326,6 +334,7 @@ struct llama_context { bool sched_need_reserve = true; const llama_pshard_plan * pshard_active_plan = nullptr; + bool pshard_memory_dirty = false; pshard_dev_layout pshard_layout = {}; ggml_backend_t backend_cpu = nullptr; @@ -381,7 +390,6 @@ void pshard_assign_tensors( const llama_model & model, llama_memory_i * memory, const std::vector & backends, - const pshard_dev_layout & layout, - ggml_cgraph * gf); + const pshard_dev_layout & layout); void pshard_refresh_stream_views(llama_memory_i * memory); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 9d2c58f51eaf..be51e028ba99 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1936,10 +1936,6 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { GGML_UNUSED(flags); - if (pipe_shard_kv) { - pipe_shard_kv->prepare_for_host_access(); - } - io.write(&n_stream, sizeof(n_stream)); for (uint32_t s = 0; s < n_stream; ++s) { @@ -1993,10 +1989,6 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { GGML_UNUSED(flags); - if (pipe_shard_kv) { - pipe_shard_kv->prepare_for_host_access(); - } - GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())); uint32_t n_stream_cur; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index cb1574d19129..8afc3cb8af8c 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1417,7 +1417,8 @@ bool llama_model_loader::load_all_data( llama_buf_map & bufs, llama_mlocks * lmlocks, llama_progress_callback progress_callback, - void * progress_callback_user_data) { + void * progress_callback_user_data, + const std::unordered_set * skip_tensors) { if (files.empty()) { for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { set_tensor_data(t, set_tensor_data_ud); @@ -1541,6 +1542,10 @@ bool llama_model_loader::load_all_data( } size_t n_size = ggml_nbytes(cur); + if (skip_tensors && skip_tensors->find(cur) != skip_tensors->end()) { + size_done += n_size; + continue; + } if (use_mmap) { const auto & mapping = mappings.at(weight->idx); @@ -1710,6 +1715,266 @@ void llama_model_loader::print_info() const { } } +struct llama_async_tensor_uploader { + static constexpr size_t n_buffers = 4; + + ggml_backend_t backend = nullptr; + std::vector host_buffers; + std::vector events; + std::vector host_ptrs; + size_t buffer_idx = 0; + size_t buffer_size = 0; + + ~llama_async_tensor_uploader() { + finish(); + } + + bool init(ggml_backend_buffer_t dst_buf, const char * func, size_t read_alignment = 1) { + if (!dst_buf) { + LLAMA_LOG_DEBUG("%s: no buffer found for async uploads\n", func); + return false; + } + + auto * buft = ggml_backend_buffer_get_type(dst_buf); + auto * dev = ggml_backend_buft_get_device(buft); + if (!dev) { + LLAMA_LOG_DEBUG("%s: no device found for buffer type %s for async uploads\n", func, + ggml_backend_buft_name(buft)); + return false; + } + + if (buft != ggml_backend_dev_buffer_type(dev)) { + LLAMA_LOG_DEBUG("%s: buffer type %s is not the default buffer type for device %s for async uploads\n", func, + ggml_backend_buft_name(buft), ggml_backend_dev_name(dev)); + return false; + } + + ggml_backend_dev_props props; + ggml_backend_dev_get_props(dev, &props); + if (!props.caps.async || !props.caps.host_buffer || !props.caps.events) { + LLAMA_LOG_DEBUG("%s: device %s does not support async, host buffers or events\n", func, + ggml_backend_dev_name(dev)); + return false; + } + + auto * host_buft = ggml_backend_dev_host_buffer_type(dev); + if (!host_buft) { + LLAMA_LOG_DEBUG("%s: no host buffer type found for device %s\n", func, + ggml_backend_dev_name(dev)); + return false; + } + + buffer_size = read_alignment != 1 ? 64 * 1024 * 1024 + 2 * read_alignment : 1024 * 1024; + host_buffers.reserve(n_buffers); + host_ptrs.reserve(n_buffers); + events.reserve(n_buffers); + + for (size_t idx = 0; idx < n_buffers; ++idx) { + auto * buf = ggml_backend_buft_alloc_buffer(host_buft, buffer_size); + if (!buf) { + LLAMA_LOG_DEBUG("%s: failed to allocate host buffer for async uploads for device %s\n", func, + ggml_backend_dev_name(dev)); + finish(); + return false; + } + + host_buffers.emplace_back(buf); + host_ptrs.emplace_back(ggml_backend_buffer_get_base(buf)); + + auto * event = ggml_backend_event_new(dev); + if (!event) { + LLAMA_LOG_DEBUG("%s: failed to create event for async uploads for device %s\n", func, + ggml_backend_dev_name(dev)); + finish(); + return false; + } + + events.emplace_back(event); + } + + backend = ggml_backend_dev_init(dev, nullptr); + if (!backend) { + LLAMA_LOG_DEBUG("%s: failed to initialize backend for device %s for async uploads\n", func, + ggml_backend_dev_name(dev)); + finish(); + return false; + } + + LLAMA_LOG_DEBUG("%s: using async uploads for device %s, buffer type %s, backend %s\n", func, + ggml_backend_dev_name(dev), + ggml_backend_buft_name(buft), + ggml_backend_name(backend)); + return true; + } + + bool active() const { + return backend != nullptr; + } + + void upload(llama_file & file, ggml_tensor * cur, size_t offset, size_t n_size) { + GGML_ASSERT(active()); + + size_t alignment = file.read_alignment(); + size_t aligned_offset = offset & ~(alignment - 1); + size_t offset_from_alignment = offset - aligned_offset; + file.seek(aligned_offset, SEEK_SET); + + size_t read_start = aligned_offset; + size_t read_end = (offset + n_size + alignment - 1) & ~(alignment - 1); + + size_t bytes_read = 0; + size_t data_read = 0; + + while (bytes_read < read_end - read_start) { + size_t read_size = std::min(buffer_size, read_end - read_start - bytes_read); + + uintptr_t ptr_dest_aligned = (reinterpret_cast(host_ptrs[buffer_idx]) + alignment - 1) & ~(alignment - 1); + + ggml_backend_event_synchronize(events[buffer_idx]); + + file.read_raw_unsafe(reinterpret_cast(ptr_dest_aligned), read_size); + + uintptr_t ptr_data = ptr_dest_aligned; + size_t data_to_copy = read_size; + + if (bytes_read == 0) { + ptr_data += offset_from_alignment; + data_to_copy -= offset_from_alignment; + } + + if (aligned_offset + bytes_read + read_size > offset + n_size) { + data_to_copy -= (read_end - (offset + n_size)); + } + + ggml_backend_tensor_set_async(backend, cur, + reinterpret_cast(ptr_data), data_read, data_to_copy); + ggml_backend_event_record(events[buffer_idx], backend); + + data_read += data_to_copy; + bytes_read += read_size; + + ++buffer_idx; + buffer_idx %= n_buffers; + } + } + + void finish() { + for (auto * event : events) { + ggml_backend_event_synchronize(event); + ggml_backend_event_free(event); + } + events.clear(); + + for (auto * buf : host_buffers) { + ggml_backend_buffer_free(buf); + } + host_buffers.clear(); + host_ptrs.clear(); + + ggml_backend_free(backend); + backend = nullptr; + buffer_idx = 0; + buffer_size = 0; + } +}; + +bool llama_model_loader::preload_common_weights_to_device( + const std::vector & preload_order, + size_t n_common, + size_t buf_size, + ggml_backend_buffer_t * out_buf, + std::unordered_map * out_preload_map, + size_t * out_preloaded_size, + std::unordered_set * out_device_only_tensors) { + if (n_common == 0) { + return false; + } + + auto * gpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU); + if (!gpu_dev) { + return false; + } + + auto * buft = ggml_backend_dev_buffer_type(gpu_dev); + *out_buf = ggml_backend_buft_alloc_buffer(buft, buf_size); + if (!*out_buf) { + LLAMA_LOG_WARN("%s: failed to allocate %.2f MiB device buffer\n", + __func__, buf_size / (1024.0 * 1024.0)); + return false; + } + + ggml_backend_buffer_set_usage(*out_buf, GGML_BACKEND_BUFFER_USAGE_COMPUTE); + + size_t alignment = 1; + for (const auto & file : files) { + alignment = std::max(file->read_alignment(), alignment); + } + + llama_async_tensor_uploader uploader; + const bool use_async = uploader.init(*out_buf, __func__, alignment); + std::vector> read_buf; + + struct ggml_tallocr talloc = ggml_tallocr_new(*out_buf); + size_t n_packed = 0; + + for (size_t i = 0; i < n_common; i++) { + ggml_tensor * tensor = preload_order[i]; + if (tensor == nullptr) { + continue; + } + + const auto * weight = get_weight(ggml_get_name(tensor)); + if (weight == nullptr) { + continue; + } + + const size_t tsize = ggml_backend_buffer_get_alloc_size(*out_buf, tensor); + if (talloc.offset + tsize > buf_size) { + LLAMA_LOG_ERROR("%s: common tensor %s does not fit in %.2f MiB preload buffer " + "(offset=%.2f MiB, size=%.2f MiB)\n", + __func__, ggml_get_name(tensor), + buf_size / (1024.0 * 1024.0), + talloc.offset / (1024.0 * 1024.0), + tsize / (1024.0 * 1024.0)); + GGML_ASSERT(false && "pshard: common canonical tensor does not fit preload buffer"); + } + + tensor->buffer = nullptr; + tensor->data = nullptr; + ggml_tallocr_alloc(&talloc, tensor); + + (*out_preload_map)[tensor] = { + /* cpu_addr = */ nullptr, + /* gpu_addr = */ tensor->data, + /* host_buffer = */ nullptr, + /* device_only_common = */ true, + }; + out_device_only_tensors->insert(tensor); + + auto & file = files.at(weight->idx); + const size_t n_size = ggml_nbytes(tensor); + if (use_async) { + uploader.upload(*file, tensor, weight->offs, n_size); + } else { + read_buf.resize(n_size); + file->seek(weight->offs, SEEK_SET); + file->read_raw(read_buf.data(), n_size); + ggml_backend_tensor_set(tensor, read_buf.data(), 0, n_size); + } + + n_packed++; + } + + uploader.finish(); + + *out_preloaded_size = talloc.offset; + + LLAMA_LOG_INFO("%s: preloaded %zu common tensors (%.2f MiB) into %.2f MiB device buffer\n", + __func__, n_packed, talloc.offset / (1024.0 * 1024.0), buf_size / (1024.0 * 1024.0)); + + return true; +} + bool llama_model_loader::preload_weights_to_device( const std::unordered_map & tensor_backend_ids, int target_backend_id, @@ -1717,7 +1982,8 @@ bool llama_model_loader::preload_weights_to_device( ggml_backend_buffer_t * out_buf, ggml_backend_t * out_backend, std::unordered_map * out_preload_map, - size_t * out_preloaded_size) { + size_t * out_preloaded_size, + const std::vector * preload_order) { auto * gpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU); if (!gpu_dev) { @@ -1727,9 +1993,18 @@ bool llama_model_loader::preload_weights_to_device( auto * buft = ggml_backend_dev_buffer_type(gpu_dev); std::vector pinned; - for (const auto & [tensor, bid] : tensor_backend_ids) { - if (bid == target_backend_id && tensor->data != nullptr) { - pinned.push_back(tensor); + if (preload_order != nullptr) { + pinned.reserve(preload_order->size()); + for (ggml_tensor * tensor : *preload_order) { + if (tensor != nullptr && tensor->data != nullptr) { + pinned.push_back(tensor); + } + } + } else { + for (const auto & [tensor, bid] : tensor_backend_ids) { + if (bid == target_backend_id && tensor->data != nullptr) { + pinned.push_back(tensor); + } } } @@ -1741,26 +2016,28 @@ bool llama_model_loader::preload_weights_to_device( pinned_bytes / (1024.0 * 1024.0), tensor_backend_ids.size()); } - auto get_layer = [](const ggml_tensor * t) -> int { - const char * blk = strstr(ggml_get_name(t), "blk."); - return blk ? atoi(blk + 4) : 9999; - }; - auto get_cat = [](const ggml_tensor * t) -> int { - const char * name = ggml_get_name(t); - if (strstr(name, "attn_")) return 0; - if (strstr(name, "exps")) return 4; - if (strstr(name, "ffn_")) return 1; - if (strstr(name, "norm")) return 2; - return 3; - }; - std::sort(pinned.begin(), pinned.end(), - [&](const ggml_tensor * a, const ggml_tensor * b) { - int ca = get_cat(a), cb = get_cat(b); - if (ca != cb) return ca < cb; - int la = get_layer(a), lb = get_layer(b); - if (la != lb) return la < lb; - return strcmp(ggml_get_name(a), ggml_get_name(b)) < 0; - }); + if (preload_order == nullptr) { + auto get_layer = [](const ggml_tensor * t) -> int { + const char * blk = strstr(ggml_get_name(t), "blk."); + return blk ? atoi(blk + 4) : 9999; + }; + auto get_cat = [](const ggml_tensor * t) -> int { + const char * name = ggml_get_name(t); + if (strstr(name, "attn_")) return 0; + if (strstr(name, "exps")) return 4; + if (strstr(name, "ffn_")) return 1; + if (strstr(name, "norm")) return 2; + return 3; + }; + std::sort(pinned.begin(), pinned.end(), + [&](const ggml_tensor * a, const ggml_tensor * b) { + int ca = get_cat(a), cb = get_cat(b); + if (ca != cb) return ca < cb; + int la = get_layer(a), lb = get_layer(b); + if (la != lb) return la < lb; + return strcmp(ggml_get_name(a), ggml_get_name(b)) < 0; + }); + } *out_buf = ggml_backend_buft_alloc_buffer(buft, buf_size); if (!*out_buf) { diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 6f8255864529..c29ae2f4a36d 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -14,6 +14,7 @@ #include #include #include +#include using llama_buf_map = std::unordered_map; @@ -21,6 +22,7 @@ struct weight_preload_entry { void * cpu_addr = nullptr; void * gpu_addr = nullptr; ggml_backend_buffer_t host_buffer = nullptr; + bool device_only_common = false; }; // lists of buffer types used for each layer @@ -206,7 +208,17 @@ struct llama_model_loader { llama_buf_map & bufs, llama_mlocks * lmlocks, llama_progress_callback progress_callback, - void * progress_callback_user_data); + void * progress_callback_user_data, + const std::unordered_set * skip_tensors = nullptr); + + bool preload_common_weights_to_device( + const std::vector & preload_order, + size_t n_common, + size_t buf_size, + ggml_backend_buffer_t * out_buf, + std::unordered_map * out_preload_map, + size_t * out_preloaded_size, + std::unordered_set * out_device_only_tensors); bool preload_weights_to_device( const std::unordered_map & tensor_backend_ids, @@ -215,7 +227,8 @@ struct llama_model_loader { ggml_backend_buffer_t * out_buf, ggml_backend_t * out_backend, std::unordered_map * out_preload_map, - size_t * out_preloaded_size); + size_t * out_preloaded_size, + const std::vector * preload_order = nullptr); std::string ftype_name() const; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d3dcf93d10b3..b8822a76d999 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include #include #include @@ -33,8 +34,39 @@ #include #include #include +#include #include +static int pshard_weight_layer(const ggml_tensor * t) { + const char * blk = strstr(ggml_get_name(t), "blk."); + return blk ? atoi(blk + 4) : 9999; +} + +static int pshard_weight_cat(const ggml_tensor * t) { + const char * name = ggml_get_name(t); + if (strstr(name, "attn_")) return 0; + if (strstr(name, "exps")) return 4; + if (strstr(name, "ffn_")) return 1; + if (strstr(name, "norm")) return 2; + return 3; +} + +static bool pshard_weight_less(const ggml_tensor * a, const ggml_tensor * b) { + const int ca = pshard_weight_cat(a); + const int cb = pshard_weight_cat(b); + if (ca != cb) { + return ca < cb; + } + + const int la = pshard_weight_layer(a); + const int lb = pshard_weight_layer(b); + if (la != lb) { + return la < lb; + } + + return strcmp(ggml_get_name(a), ggml_get_name(b)) < 0; +} + struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const struct ggml_tensor * tensor, void * userdata) { const llama_meta_device_get_split_state_userdata * ud = (const llama_meta_device_get_split_state_userdata *) userdata; const llama_hparams & hparams = ud->model->hparams; @@ -673,6 +705,16 @@ struct llama_model::impl { ggml_backend_t dev_preload_backend = nullptr; std::unordered_map weight_preload_map; size_t dev_preloaded_size = 0; + + struct pshard_weight_layout { + std::unordered_map offsets; + std::unordered_map tensors; + std::unordered_set common; + size_t common_end = 0; + bool ready = false; + }; + + pshard_weight_layout pshard_weight_layout; }; llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique()) { @@ -8164,9 +8206,34 @@ bool llama_model::load_tensors(llama_model_loader & ml) { } } + pimpl->plan_registry = params.pshard_registry; + + std::vector preload_order; + size_t n_common = 0; + std::unordered_map preload_tensor_backend_ids = pimpl->tensor_backend_ids; + const bool pshard_preload_requested = params.pshard && params.max_vram_alloc > 0; + const bool pshard_has_registry = + pshard_preload_requested && + pimpl->plan_registry != nullptr && + !pimpl->plan_registry->best_plans.empty(); + if (pshard_has_registry) { + preload_tensor_backend_ids = pshard_build_canonical_weight_order(preload_order, n_common); + } + ml.init_mappings(true, use_mlock ? &pimpl->mlock_mmaps : nullptr); pimpl->mappings.reserve(ml.mappings.size()); + std::unordered_set pshard_device_only_tensors; + if (pshard_has_registry && n_common > 0) { + const size_t buf_size = params.max_vram_alloc * 1024ULL * 1024ULL; + ml.preload_common_weights_to_device( + preload_order, n_common, buf_size, + &pimpl->dev_preload_buf, + &pimpl->weight_preload_map, + &pimpl->dev_preloaded_size, + &pshard_device_only_tensors); + } + // create the backend buffers std::vector> ctx_buf_maps; ctx_buf_maps.reserve(ml.ctx_map.size()); @@ -8175,6 +8242,15 @@ bool llama_model::load_tensors(llama_model_loader & ml) { const size_t n_max_backend_buffer = ml.ctx_map.size() * ml.files.size(); pimpl->ctxs_bufs.reserve(n_max_backend_buffer); + auto all_tensors_already_allocated = [](ggml_context * ctx) { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (t->data == nullptr && t->view_src == nullptr) { + return false; + } + } + return true; + }; + for (auto & [buft, ctx_ptr] : ml.ctx_map) { ggml_context * ctx = ctx_ptr.get(); @@ -8233,17 +8309,21 @@ bool llama_model::load_tensors(llama_model_loader & ml) { buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft); // real buffer } if (buf == nullptr) { - throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft))); + if (!all_tensors_already_allocated(ctx)) { + throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft))); + } } - if (use_mlock && ggml_backend_buffer_is_host(buf)) { + if (buf != nullptr && use_mlock && ggml_backend_buffer_is_host(buf)) { pimpl->mlock_bufs.emplace_back(new llama_mlock); auto & mlock_buf = pimpl->mlock_bufs.back(); mlock_buf->init (ggml_backend_buffer_get_base(buf)); mlock_buf->grow_to(ggml_backend_buffer_get_size(buf)); } - bufs.emplace_back(buf); - for (uint32_t idx = 0; idx < ml.files.size(); idx++) { - buf_map.emplace(idx, buf); + if (buf != nullptr) { + bufs.emplace_back(buf); + for (uint32_t idx = 0; idx < ml.files.size(); idx++) { + buf_map.emplace(idx, buf); + } } } @@ -8288,7 +8368,9 @@ bool llama_model::load_tensors(llama_model_loader & ml) { // load tensor data for (auto & [ctx, buf_map] : ctx_buf_maps) { - if (!ml.load_all_data(ctx, buf_map, use_mlock ? &pimpl->mlock_mmaps : NULL, params.progress_callback, params.progress_callback_user_data)) { + if (!ml.load_all_data(ctx, buf_map, use_mlock ? &pimpl->mlock_mmaps : NULL, + params.progress_callback, params.progress_callback_user_data, + pshard_device_only_tensors.empty() ? nullptr : &pshard_device_only_tensors)) { return false; } } @@ -8301,21 +8383,80 @@ bool llama_model::load_tensors(llama_model_loader & ml) { if (params.pshard && params.max_vram_alloc > 0) { size_t buf_size = params.max_vram_alloc * 1024ULL * 1024ULL; - if (ml.preload_weights_to_device( - pimpl->tensor_backend_ids, 0, buf_size, + + if (pimpl->dev_preload_buf != nullptr) { + for (const auto & [tensor, bid] : preload_tensor_backend_ids) { + if (tensor == nullptr || bid < 0 || tensor->data == nullptr) { + continue; + } + + auto it = pimpl->weight_preload_map.find(tensor); + if (it == pimpl->weight_preload_map.end()) { + pimpl->weight_preload_map[tensor] = { + /* cpu_addr = */ tensor->data, + /* gpu_addr = */ nullptr, + /* host_buffer = */ tensor->buffer, + /* device_only_common = */ false, + }; + } + } + + if (!preload_order.empty()) { + pshard_finalize_canonical_weight_layout(preload_order, n_common); + for (auto & plan : pimpl->plan_registry->best_plans) { + if (plan.is_viable) { + pshard_stamp_plan_offsets(plan); + } + } + } + + size_t n_preloaded = 0; + size_t preloaded_alloc_size = 0; + for (const auto & [tensor, entry] : pimpl->weight_preload_map) { + if (entry.gpu_addr != nullptr) { + n_preloaded++; + preloaded_alloc_size += ggml_backend_buffer_get_alloc_size(pimpl->dev_preload_buf, tensor); + } + } + + LLAMA_LOG_INFO("%s: preloaded %zu/%zu weights (packed=%.2f MiB, common=%.2f MiB) into %.2f MiB device buffer\n", + __func__, n_preloaded, pimpl->weight_preload_map.size(), + preloaded_alloc_size / (1024.0 * 1024.0), + pimpl->dev_preloaded_size / (1024.0 * 1024.0), + buf_size / (1024.0 * 1024.0)); + } else if (ml.preload_weights_to_device( + preload_tensor_backend_ids, 0, buf_size, &pimpl->dev_preload_buf, &pimpl->dev_preload_backend, &pimpl->weight_preload_map, - &pimpl->dev_preloaded_size)) { - LLAMA_LOG_INFO("%s: preloaded %zu weights (%.2f MiB) into %.2f MiB device buffer\n", - __func__, pimpl->weight_preload_map.size(), + &pimpl->dev_preloaded_size, + preload_order.empty() ? nullptr : &preload_order)) { + if (!preload_order.empty()) { + pshard_finalize_canonical_weight_layout(preload_order, n_common); + for (auto & plan : pimpl->plan_registry->best_plans) { + if (plan.is_viable) { + pshard_stamp_plan_offsets(plan); + } + } + } + + size_t n_preloaded = 0; + size_t preloaded_alloc_size = 0; + for (const auto & [tensor, entry] : pimpl->weight_preload_map) { + if (entry.gpu_addr != nullptr) { + n_preloaded++; + preloaded_alloc_size += ggml_backend_buffer_get_alloc_size(pimpl->dev_preload_buf, tensor); + } + } + + LLAMA_LOG_INFO("%s: preloaded %zu/%zu weights (packed=%.2f MiB, common=%.2f MiB) into %.2f MiB device buffer\n", + __func__, n_preloaded, pimpl->weight_preload_map.size(), + preloaded_alloc_size / (1024.0 * 1024.0), pimpl->dev_preloaded_size / (1024.0 * 1024.0), buf_size / (1024.0 * 1024.0)); } } - pimpl->plan_registry = params.pshard_registry; - return true; } @@ -8644,6 +8785,10 @@ 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; } @@ -8672,6 +8817,193 @@ void llama_model::sync_dev_preload() { } } +std::unordered_map llama_model::pshard_build_canonical_weight_order( + std::vector & preload_order, + size_t & n_common) { + std::unordered_map union_bids; + + auto * registry = pimpl->plan_registry; + if (registry == nullptr) { + preload_order.clear(); + n_common = 0; + return pimpl->tensor_backend_ids; + } + + std::unordered_map resident_count; + std::unordered_map resident_tensors; + int n_plans = 0; + + for (const auto & plan : registry->best_plans) { + if (!plan.is_viable) { + continue; + } + + pshard_set_backend_maps(plan); + + std::unordered_set resident; + for (const auto & [tensor, bid] : pimpl->tensor_backend_ids) { + if (tensor == nullptr || bid < 0) { + continue; + } + + auto uit = union_bids.find(tensor); + if (uit == union_bids.end() || bid == 0) { + union_bids[tensor] = bid; + } + + if (bid == 0) { + const std::string name = ggml_get_name(tensor); + resident.insert(name); + resident_tensors[name] = tensor; + } + } + + for (const auto & name : resident) { + resident_count[name]++; + } + n_plans++; + } + + if (n_plans == 0) { + preload_order.clear(); + n_common = 0; + return pimpl->tensor_backend_ids; + } + + std::vector common; + std::vector extras; + common.reserve(resident_tensors.size()); + extras.reserve(resident_tensors.size()); + + for (const auto & [name, tensor] : resident_tensors) { + if (resident_count[name] == n_plans) { + common.push_back(tensor); + } else { + extras.push_back(tensor); + } + } + + std::sort(common.begin(), common.end(), pshard_weight_less); + std::sort(extras.begin(), extras.end(), pshard_weight_less); + + preload_order.clear(); + preload_order.reserve(resident_tensors.size()); + preload_order.insert(preload_order.end(), common.begin(), common.end()); + n_common = preload_order.size(); + + // preload all possible extras once, but do not make their union layout part of + // the canonical plan offsets. Each plan packs only its own extras after common. + preload_order.insert(preload_order.end(), extras.begin(), extras.end()); + + LLAMA_LOG_INFO("%s: canonical layout from %d plans: %zu common, %zu extras\n", + __func__, n_plans, common.size(), extras.size()); + + return union_bids; +} + +void llama_model::pshard_finalize_canonical_weight_layout( + const std::vector & preload_order, + size_t n_common) { + if (!pimpl->dev_preload_buf) { + return; + } + + void * buf_base = ggml_backend_buffer_get_base(pimpl->dev_preload_buf); + pimpl->pshard_weight_layout.offsets.clear(); + pimpl->pshard_weight_layout.tensors.clear(); + pimpl->pshard_weight_layout.common.clear(); + pimpl->pshard_weight_layout.common_end = 0; + pimpl->pshard_weight_layout.ready = false; + + for (size_t i = 0; i < preload_order.size(); i++) { + ggml_tensor * tensor = preload_order[i]; + if (tensor == nullptr || tensor->data == nullptr || tensor->buffer != pimpl->dev_preload_buf) { + continue; + } + + const std::string name = ggml_get_name(tensor); + const size_t off = (size_t) ((char *) tensor->data - (char *) buf_base); + pimpl->pshard_weight_layout.offsets[name] = off; + pimpl->pshard_weight_layout.tensors[name] = tensor; + + if (i < n_common) { + pimpl->pshard_weight_layout.common.insert(name); + const size_t size = ggml_backend_buffer_get_alloc_size(pimpl->dev_preload_buf, tensor); + pimpl->pshard_weight_layout.common_end = std::max(pimpl->pshard_weight_layout.common_end, off + size); + } + } + + pimpl->pshard_weight_layout.ready = true; + pimpl->dev_preloaded_size = pimpl->pshard_weight_layout.common_end; + + LLAMA_LOG_INFO("%s: canonical common_end=%.2f MiB, packed=%zu/%zu tensors\n", + __func__, + pimpl->pshard_weight_layout.common_end / (1024.0 * 1024.0), + pimpl->pshard_weight_layout.offsets.size(), preload_order.size()); +} + +void llama_model::pshard_stamp_plan_offsets(const llama_pshard_plan & plan) { + if (!pimpl->pshard_weight_layout.ready || !pimpl->dev_preload_buf) { + return; + } + + pshard_set_backend_maps(plan); + + std::vector extras; + + plan.cached_weight_offsets.clear(); + size_t scratch_off = pimpl->pshard_weight_layout.common_end; + bool missing = false; + + for (const auto & [tensor, bid] : pimpl->tensor_backend_ids) { + if (tensor == nullptr || bid != 0) { + continue; + } + + const std::string name = ggml_get_name(tensor); + if (pimpl->pshard_weight_layout.common.find(name) != pimpl->pshard_weight_layout.common.end()) { + auto it = pimpl->pshard_weight_layout.offsets.find(name); + if (it == pimpl->pshard_weight_layout.offsets.end()) { + missing = true; + break; + } + const size_t size = ggml_backend_buffer_get_alloc_size(pimpl->dev_preload_buf, tensor); + plan.cached_weight_offsets[name] = it->second; + scratch_off = std::max(scratch_off, it->second + size); + } else { + extras.push_back(tensor); + } + } + + if (missing) { + plan.cached_weight_offsets.clear(); + plan.cached_scratch_off = 0; + plan.addrs_cached = false; + LLAMA_LOG_WARN("%s: plan %s bs=%u has resident tensors outside canonical layout; falling back to per-plan packing\n", + __func__, llama_pshard_strategy_name(plan.strategy), plan.batch_size); + return; + } + + std::sort(extras.begin(), extras.end(), pshard_weight_less); + + const size_t alignment = ggml_backend_buffer_get_alignment(pimpl->dev_preload_buf); + auto align_up = [alignment](size_t off) { + return ((off + alignment - 1) / alignment) * alignment; + }; + + scratch_off = align_up(scratch_off); + for (ggml_tensor * tensor : extras) { + const std::string name = ggml_get_name(tensor); + const size_t size = ggml_backend_buffer_get_alloc_size(pimpl->dev_preload_buf, tensor); + + plan.cached_weight_offsets[name] = scratch_off; + scratch_off = align_up(scratch_off + size); + } + + plan.cached_scratch_off = scratch_off; + plan.addrs_cached = true; +} + size_t llama_model::pshard_compute_scratch_off(const llama_pshard_plan & plan) { if (plan.addrs_cached) { return plan.cached_scratch_off; @@ -8788,6 +9120,7 @@ void llama_model::pshard_set_backend_maps(const llama_pshard_plan & plan) { } size_t llama_model::pshard_apply_plan(const llama_pshard_plan & plan, ggml_backend_t gpu) { + params.pshard_delegate_compute = llama_pshard_strategy_delegates_compute(plan.strategy); pshard_set_backend_maps(plan); LLAMA_LOG_DEBUG("%s: rebuilt maps: %zu tensors, %zu layers (cached=%d)\n", @@ -8797,6 +9130,7 @@ size_t llama_model::pshard_apply_plan(const llama_pshard_plan & plan, ggml_backe size_t scratch_off = 0; if (pimpl->dev_preload_buf) { void * buf_base = ggml_backend_buffer_get_base(pimpl->dev_preload_buf); + size_t buf_size = ggml_backend_buffer_get_size(pimpl->dev_preload_buf); std::unordered_map old_addrs; old_addrs.reserve(pimpl->weight_preload_map.size()); @@ -8811,6 +9145,9 @@ size_t llama_model::pshard_apply_plan(const llama_pshard_plan & plan, ggml_backe if (it != plan.cached_weight_offsets.end()) { tensor->data = (char *)buf_base + it->second; tensor->buffer = pimpl->dev_preload_buf; + } else if (entry.device_only_common) { + tensor->data = entry.gpu_addr; + tensor->buffer = pimpl->dev_preload_buf; } else { tensor->data = entry.cpu_addr; tensor->buffer = entry.host_buffer; @@ -8847,11 +9184,26 @@ size_t llama_model::pshard_apply_plan(const llama_pshard_plan & plan, ggml_backe }); struct ggml_tallocr talloc = ggml_tallocr_new(pimpl->dev_preload_buf); - size_t buf_size = ggml_backend_buffer_get_size(pimpl->dev_preload_buf); - for (auto * tensor : pinned) { + auto entry_it = pimpl->weight_preload_map.find(tensor); + if (entry_it == pimpl->weight_preload_map.end()) { + continue; + } size_t tsize = ggml_backend_buffer_get_alloc_size(pimpl->dev_preload_buf, tensor); - if (talloc.offset + tsize > buf_size) break; + if (entry_it->second.device_only_common) { + const size_t off = (size_t) ((char *) entry_it->second.gpu_addr - (char *) buf_base); + talloc.offset = std::max(talloc.offset, GGML_PAD(off + tsize, talloc.alignment)); + continue; + } + if (talloc.offset + tsize > buf_size) { + LLAMA_LOG_ERROR("%s: fallback packing cannot fit resident tensor %s " + "(offset=%.2f MiB, size=%.2f MiB, buffer=%.2f MiB)\n", + __func__, ggml_get_name(tensor), + talloc.offset / (1024.0 * 1024.0), + tsize / (1024.0 * 1024.0), + buf_size / (1024.0 * 1024.0)); + GGML_ASSERT(false && "pshard: fallback packing cannot fit resident tensor"); + } tensor->buffer = NULL; tensor->data = NULL; ggml_tallocr_alloc(&talloc, tensor); @@ -8861,7 +9213,13 @@ size_t llama_model::pshard_apply_plan(const llama_pshard_plan & plan, ggml_backe plan.cached_weight_offsets.clear(); for (auto * tensor : pinned) { - if (tensor->data) { + const size_t size = ggml_backend_buffer_get_alloc_size(pimpl->dev_preload_buf, tensor); + const bool in_preload_buf = + tensor->buffer == pimpl->dev_preload_buf && + tensor->data != nullptr && + (char *) tensor->data >= (char *) buf_base && + (char *) tensor->data + size <= (char *) buf_base + buf_size; + if (in_preload_buf) { plan.cached_weight_offsets[std::string(ggml_get_name(tensor))] = (size_t)((char *)tensor->data - (char *)buf_base); } @@ -8872,7 +9230,14 @@ size_t llama_model::pshard_apply_plan(const llama_pshard_plan & plan, ggml_backe for (auto & [tensor, entry] : pimpl->weight_preload_map) { auto it = pimpl->tensor_backend_ids.find(tensor); bool is_pinned = (it != pimpl->tensor_backend_ids.end() && it->second == 0); - if (!is_pinned) { + bool is_cached = plan.cached_weight_offsets.find(std::string(ggml_get_name(tensor))) != + plan.cached_weight_offsets.end(); + if (is_pinned) { + GGML_ASSERT(is_cached && "pshard: resident tensor missing cached preload offset"); + } else if (entry.device_only_common) { + tensor->data = entry.gpu_addr; + tensor->buffer = pimpl->dev_preload_buf; + } else { tensor->data = entry.cpu_addr; tensor->buffer = entry.host_buffer; } @@ -8883,6 +9248,7 @@ size_t llama_model::pshard_apply_plan(const llama_pshard_plan & plan, ggml_backe size_t bytes_uploaded = 0; for (auto & [tensor, entry] : pimpl->weight_preload_map) { + if (entry.device_only_common || entry.cpu_addr == nullptr) continue; if (tensor->data == entry.cpu_addr) continue; void * old = old_addrs[tensor]; if (old != tensor->data && gpu) { @@ -8892,8 +9258,12 @@ size_t llama_model::pshard_apply_plan(const llama_pshard_plan & plan, ggml_backe } } - LLAMA_LOG_DEBUG("%s: scratch_off=%.2f MiB, %zu uploaded (%.2f MiB), cached=%d\n", - __func__, scratch_off / (1024.0 * 1024.0), + if (gpu && n_uploaded > 0) { + ggml_backend_synchronize(gpu); + } + + LLAMA_LOG_DEBUG("%s: strategy=%s bs=%u scratch_off=%.2f MiB, %zu uploaded (%.2f MiB), cached=%d\n", + __func__, llama_pshard_strategy_name(plan.strategy), plan.batch_size, scratch_off / (1024.0 * 1024.0), n_uploaded, bytes_uploaded / (1024.0 * 1024.0), (int)plan.addrs_cached); } @@ -9631,6 +10001,7 @@ llama_model_params llama_model_default_params() { /*.no_host =*/ false, /*.no_alloc =*/ false, /*.pshard =*/ false, + /*.pshard_delegate_compute =*/ false, /*.pshard_cache_skip_load =*/ false, /*.max_vram_alloc =*/ 0, /*.pshard_registry =*/ nullptr, diff --git a/src/llama-model.h b/src/llama-model.h index 8f75b88ba43d..7e5507bc452d 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -620,6 +620,7 @@ struct llama_model { bool has_tensor_overrides() const; bool is_pshard() const; + bool pshard_delegates_compute() const; const std::unordered_map & get_tensor_backend_ids() const; const std::unordered_map & get_layer_backend_ids() const; @@ -629,6 +630,13 @@ struct llama_model { void sync_dev_preload(); void pshard_set_backend_maps(const llama_pshard_plan & plan); + std::unordered_map pshard_build_canonical_weight_order( + std::vector & preload_order, + size_t & n_common); + void pshard_finalize_canonical_weight_layout( + const std::vector & preload_order, + size_t n_common); + void pshard_stamp_plan_offsets(const llama_pshard_plan & plan); size_t pshard_compute_scratch_off(const llama_pshard_plan & plan); size_t pshard_apply_plan(const llama_pshard_plan & plan, ggml_backend_t gpu = nullptr); diff --git a/src/llama-pshard-cache.cpp b/src/llama-pshard-cache.cpp index a80374464c26..afb96ad6668c 100644 --- a/src/llama-pshard-cache.cpp +++ b/src/llama-pshard-cache.cpp @@ -11,6 +11,7 @@ #include #include #include +#include const char * llama_get_overflow_pattern(size_t il, llama_layer_fraction lf) { constexpr size_t n_strings = 1000; @@ -89,7 +90,7 @@ void llama_pshard_generate_overrides( if (il == il_boundary) { const char * overflow_pat = llama_get_overflow_pattern(il, overflow_type); if (overflow_pat) { - emit(overflow_pat, host_buft, (strategy == LLAMA_PSHARD_STATIC_FITPARAMS_DENSEPRIO_MOEONLY) ? layout.cpu : layout.shard(il)); + emit(overflow_pat, host_buft, layout.shard(il)); } emit(patterns_layer[il].c_str(), host_buft, layout.compute); } else if (il >= il_pin_start && il < il_pin_end) { @@ -98,9 +99,6 @@ void llama_pshard_generate_overrides( 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_STATIC_FITPARAMS_DENSEPRIO_MOEONLY: - emit(patterns_layer[il].c_str(), host_buft, layout.cpu); - break; case LLAMA_PSHARD_GPUONLY_LAYERPIN_LAYERSTREAM: emit(patterns_layer[il].c_str(), host_buft, shard_bid); break; @@ -212,7 +210,7 @@ bool pshard_registry_load( struct tier_data { uint32_t bs = 0; bool viable = false; - llama_pshard_strategy strategy = LLAMA_PSHARD_STATIC_FITPARAMS_DENSEPRIO_MOEONLY; + llama_pshard_strategy strategy = LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS; uint32_t n_pinned = 0; uint32_t n_attn_pinned = 0; int overflow = 0; @@ -293,13 +291,17 @@ bool pshard_registry_load( td.pin_from_back = atoi(pfb + 14); td.overflow = pshard_overflow_from_name(overflow_name_buf); - td.strategy = LLAMA_PSHARD_STATIC_FITPARAMS_DENSEPRIO_MOEONLY; + 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) { + td.viable = false; + } } else if (s.compare(0, 3, "ot=") == 0 && !cur_variant->tiers.empty()) { cur_variant->tiers.back().ot_line = s.substr(3); } @@ -488,6 +490,76 @@ void llama_pshard_registry_free(llama_pshard_plan_registry * registry) { delete registry; } +struct llama_pshard_cache_probe { + std::vector devs; + uint32_t n_layers = 0; + uint32_t n_ctx_train = 0; + uint32_t n_expert = 0; + size_t vram_free = 0; + size_t vram_budget = 0; + size_t vram_total = 0; + ggml_backend_buffer_type_t host_buft = nullptr; +}; + +static bool llama_pshard_probe_model_only( + const char * path_model, + const struct llama_model_params * mparams, + size_t max_vram_mb, + size_t fit_target_mb, + llama_pshard_cache_probe & probe) { + struct user_data_t { + struct { + ggml_log_callback callback; + void * user_data; + } original_logger; + }; + user_data_t ud; + llama_log_get(&ud.original_logger.callback, &ud.original_logger.user_data); + + llama_log_set([](ggml_log_level level, const char * text, void * user_data) { + const user_data_t * ud = (const user_data_t *) user_data; + const ggml_log_level level_eff = level >= GGML_LOG_LEVEL_ERROR ? level : GGML_LOG_LEVEL_DEBUG; + ud->original_logger.callback(level_eff, text, ud->original_logger.user_data); + }, &ud); + + llama_model_params mparams_probe = *mparams; + mparams_probe.no_alloc = true; + mparams_probe.pshard = false; + mparams_probe.use_mmap = false; + mparams_probe.use_mlock = false; + + llama_model * model = llama_model_load_from_file(path_model, mparams_probe); + llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); + + if (!model) { + return false; + } + + probe.devs = model->devices; + probe.n_layers = model->hparams.n_layer; + probe.n_ctx_train = model->hparams.n_ctx_train; + probe.n_expert = model->hparams.n_expert; + + if (!probe.devs.empty()) { + ggml_backend_dev_t dev = probe.devs[0].dev; + ggml_backend_dev_memory(dev, &probe.vram_free, &probe.vram_total); + + const size_t mib = 1024ULL * 1024ULL; + const size_t fit_target_bytes = fit_target_mb * mib; + probe.vram_budget = max_vram_mb > 0 + ? max_vram_mb * mib + : (probe.vram_free > fit_target_bytes ? probe.vram_free - fit_target_bytes : 0); + + probe.host_buft = ggml_backend_dev_host_buffer_type(dev); + if (!probe.host_buft) { + probe.host_buft = ggml_backend_cpu_buffer_type(); + } + } + + llama_model_free(model); + return true; +} + static bool llama_pshard_params_supported( const struct llama_model_params * mparams, const struct llama_context_params * cparams) { @@ -530,7 +602,8 @@ void llama_params_fit_pshard( 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 max_vram_mb, + size_t fit_target_mb) { const std::string cache_path = std::string(path_model) + ".tensor_overrides.pshard_registry"; if (!llama_pshard_params_supported(mparams, cparams)) { @@ -539,34 +612,38 @@ void llama_params_fit_pshard( return; } - std::vector devs; - uint32_t hp_ngl = 0, hp_nct = 0, hp_nex = 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, GGML_LOG_LEVEL_ERROR); + llama_pshard_cache_probe probe; + if (!llama_pshard_probe_model_only(path_model, mparams, max_vram_mb, fit_target_mb, probe)) { + LLAMA_LOG_WARN("%s: failed to probe model metadata, disabling pshard\n", __func__); + mparams->pshard = false; + cparams->pshard = false; + return; + } - if (devs.empty()) { + if (probe.devs.empty()) { LLAMA_LOG_WARN("%s: no GPU devices found, disabling pshard\n", __func__); mparams->pshard = false; cparams->pshard = false; return; } - const uint32_t n_layers = hp_ngl; - const size_t vram_free = (max_vram_mb > 0) ? max_vram_mb * 1024ULL * 1024ULL : dmds[0].free; + const auto & devs = probe.devs; + const uint32_t n_layers = probe.n_layers; + const size_t vram_free = probe.vram_budget; - if (mparams->max_vram_alloc == 0 && vram_free > 0) { - mparams->max_vram_alloc = vram_free / (1024 * 1024); + if (vram_free > 0) { + mparams->max_vram_alloc = std::max(1, pshard_bytes_to_mib_ceil(vram_free)); } - 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(); + ggml_backend_buffer_type_t host_buft = probe.host_buft; + + LLAMA_LOG_INFO("%s: probe: %u layers, %.1f MiB VRAM free, %.1f MiB budget%s\n", + __func__, n_layers, + probe.vram_free / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + max_vram_mb > 0 ? " (-mva)" : " (free - fit target)"); - LLAMA_LOG_INFO("%s: probe: %u layers, %.1f MiB VRAM free\n", - __func__, n_layers, vram_free / (1024.0 * 1024.0)); + const uint32_t n_ctx_plan = cparams->n_ctx > 0 ? cparams->n_ctx : probe.n_ctx_train; auto * registry = mparams->pshard_registry; if (!registry) { @@ -575,10 +652,18 @@ void llama_params_fit_pshard( cparams->pshard = false; return; } + 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); - const uint32_t runtime_n_batch = std::min(cparams->n_ctx, cparams->n_batch); + const uint32_t runtime_n_batch = std::min(n_ctx_plan, cparams->n_batch); const uint32_t runtime_cache_ubatch = std::min(runtime_n_batch, cparams->n_ubatch == 0 ? runtime_n_batch : cparams->n_ubatch); - registry->cache_ubatch = registry->cache_ubatch ? std::min(cparams->n_ctx, registry->cache_ubatch) : runtime_cache_ubatch; + registry->cache_ubatch = registry->cache_ubatch ? std::min(n_ctx_plan, registry->cache_ubatch) : runtime_cache_ubatch; int64_t model_file_size = 0; { @@ -625,7 +710,7 @@ void llama_params_fit_pshard( } if (registry->cache_ubatch > 0) { - const uint32_t pshard_ubatch = std::min(cparams->n_ctx, registry->cache_ubatch); + const uint32_t pshard_ubatch = std::min(n_ctx_plan, registry->cache_ubatch); cparams->n_batch = pshard_ubatch; cparams->n_ubatch = pshard_ubatch; } @@ -650,6 +735,10 @@ void llama_params_fit_pshard( if (default_tier < registry->tier_sizes.size() - 1) { LLAMA_LOG_INFO("%s: highest tier (bs=%u) not viable, falling back to bs=%u\n", __func__, registry->tier_sizes.back(), registry->tier_sizes[default_tier]); + const uint32_t tier_bs = registry->tier_sizes[default_tier]; + cparams->n_batch = std::min(cparams->n_batch, tier_bs); + cparams->n_ubatch = std::min(cparams->n_ubatch, tier_bs); + LLAMA_LOG_INFO("%s: clamped n_batch/n_ubatch to %u\n", __func__, tier_bs); } if (best->n_pinned > n_layers) { @@ -662,29 +751,12 @@ void llama_params_fit_pshard( 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); - // baseline plans use the cached fit params dense only placement - // regenerating them from n_pinned and overflow would lose that placement - if (best->strategy == LLAMA_PSHARD_STATIC_FITPARAMS_DENSEPRIO_MOEONLY) { - thread_local std::vector shard_none_patterns; - shard_none_patterns.clear(); - shard_none_patterns.reserve(best->overrides.size()); - for (const auto & ov : best->overrides) { - shard_none_patterns.push_back(ov.pattern); - } - for (size_t i = 0; i < best->overrides.size(); i++) { - tensor_buft_overrides[i].pattern = shard_none_patterns[i].c_str(); - tensor_buft_overrides[i].buft = best->overrides[i].buft; - tensor_buft_overrides[i].backend_id = best->overrides[i].backend_id; - } - tensor_buft_overrides[best->overrides.size()] = { nullptr, nullptr, -1 }; - } else { - llama_pshard_generate_overrides( - best->n_pinned, n_layers, host_buft, host_buft, - tensor_buft_overrides, - (llama_layer_fraction)best->overflow, - best->strategy, layout, - best->pin_from_back, best->output_on_gpu, best->n_attn_pinned); - } + llama_pshard_generate_overrides( + best->n_pinned, n_layers, host_buft, host_buft, + tensor_buft_overrides, + (llama_layer_fraction)best->overflow, + best->strategy, layout, + best->pin_from_back, best->output_on_gpu, best->n_attn_pinned); for (size_t i = 0; tensor_buft_overrides[i].pattern; i++) { if (tensor_buft_overrides[i].backend_id == layout.compute) { diff --git a/src/llama-pshard-plan.h b/src/llama-pshard-plan.h index 9c9f1c4429e3..55bb666c88d7 100644 --- a/src/llama-pshard-plan.h +++ b/src/llama-pshard-plan.h @@ -13,24 +13,26 @@ #include enum llama_pshard_strategy { - LLAMA_PSHARD_STATIC_FITPARAMS_DENSEPRIO_MOEONLY = 0, - LLAMA_PSHARD_GPUONLY_LAYERPIN_LAYERSTREAM = 1, - LLAMA_PSHARD_GPUONLY_ATTNPIN_FFNSTREAM = 2, - LLAMA_PSHARD_DYNAMIC_FFNCPU_ATTNSTREAM = 3, - LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS = 4, + 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) { - static const char * const names[] = { - "STATIC_FITPARAMS_DENSEPRIO_MOEONLY", - "GPUONLY_LAYERPIN_LAYERSTREAM", - "GPUONLY_ATTNPIN_FFNSTREAM", - "DYNAMIC_FFNCPU_ATTNSTREAM", - "STATIC_ATTNPRIO_ALLMODELS", - }; - return (s >= 0 && s < LLAMA_PSHARD_COUNT) ? names[s] : "UNKNOWN"; + 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 @@ -69,7 +71,7 @@ struct llama_pshard_alloc_state { }; struct llama_pshard_plan { - llama_pshard_strategy strategy = LLAMA_PSHARD_STATIC_FITPARAMS_DENSEPRIO_MOEONLY; + 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) @@ -159,6 +161,12 @@ struct llama_pshard_plan_registry { 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) { diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 386a8a7a52da..a3e0195018f5 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -331,6 +331,7 @@ struct cmd_params { std::vector cpu_strict; std::vector poll; std::vector n_gpu_layers; + bool n_gpu_layers_user; std::vector n_cpu_moe; std::vector split_mode; std::vector main_gpu; @@ -344,6 +345,8 @@ struct cmd_params { std::vector embeddings; std::vector no_op_offload; std::vector no_host; + std::vector pshard; + std::vector max_vram_alloc; std::vector fit_params_target; std::vector fit_params_min_ctx; ggml_numa_strategy numa; @@ -375,6 +378,7 @@ static const cmd_params cmd_params_defaults = { /* cpu_strict */ { false }, /* poll */ { 50 }, /* n_gpu_layers */ { 99 }, + /* n_gpu_layers_user */ false, /* n_cpu_moe */ { 0 }, /* split_mode */ { LLAMA_SPLIT_MODE_LAYER }, /* main_gpu */ { 0 }, @@ -388,6 +392,8 @@ static const cmd_params cmd_params_defaults = { /* embeddings */ { false }, /* no_op_offload */ { false }, /* no_host */ { false }, + /* pshard */ { false }, + /* max_vram_alloc */ { 0 }, /* fit_params_target */ { 0 }, /* fit_params_min_ctx */ { 0 }, /* numa */ GGML_NUMA_STRATEGY_DISABLED, @@ -459,6 +465,8 @@ static void print_usage(int /* argc */, char ** argv) { printf(" (default: disabled)\n"); printf(" -nopo, --no-op-offload <0|1> (default: 0)\n"); printf(" --no-host <0|1> (default: %s)\n", join(cmd_params_defaults.no_host, ",").c_str()); + printf(" -pshard enable pshard plan cache loading\n"); + printf(" -mva, --max-vram-alloc VRAM budget in MiB for pshard (0 = use actual free VRAM minus -fitt)\n"); printf("\n"); printf( "Multiple values can be given for each parameter by separating them with ','\n" @@ -511,6 +519,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { params.delay = cmd_params_defaults.delay; params.progress = cmd_params_defaults.progress; params.no_warmup = cmd_params_defaults.no_warmup; + params.n_gpu_layers_user = cmd_params_defaults.n_gpu_layers_user; if (const char * env = getenv("HF_TOKEN")) { params.hf_token = env; @@ -710,6 +719,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } auto p = parse_int_range(argv[i]); params.n_gpu_layers.insert(params.n_gpu_layers.end(), p.begin(), p.end()); + params.n_gpu_layers_user = true; } else if (arg == "-ncmoe" || arg == "--n-cpu-moe") { if (++i >= argc) { invalid_param = true; @@ -828,6 +838,17 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } auto p = string_split(argv[i], split_delim); params.no_host.insert(params.no_host.end(), p.begin(), p.end()); + } else if (arg == "-pshard" || arg == "--pshard") { + params.pshard.push_back(true); + } else if (arg == "-mva" || arg == "--max-vram-alloc") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = string_split(argv[i], split_delim); + for (const auto & v : p) { + params.max_vram_alloc.push_back(std::stoull(v)); + } } else if (arg == "-ts" || arg == "--tensor-split") { if (++i >= argc) { invalid_param = true; @@ -1096,6 +1117,12 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { if (params.no_host.empty()) { params.no_host = cmd_params_defaults.no_host; } + if (params.pshard.empty()) { + params.pshard = cmd_params_defaults.pshard; + } + if (params.max_vram_alloc.empty()) { + params.max_vram_alloc = cmd_params_defaults.max_vram_alloc; + } if (params.n_threads.empty()) { params.n_threads = cmd_params_defaults.n_threads; } @@ -1145,6 +1172,8 @@ struct cmd_params_instance { bool embeddings; bool no_op_offload; bool no_host; + bool pshard; + size_t max_vram_alloc; size_t fit_target; uint32_t fit_min_ctx; @@ -1161,6 +1190,8 @@ struct cmd_params_instance { mparams.use_mmap = use_mmap; mparams.use_direct_io = use_direct_io; mparams.no_host = no_host; + mparams.pshard = pshard; + mparams.max_vram_alloc = max_vram_alloc; if (n_cpu_moe <= 0) { if (tensor_buft_overrides.empty()) { @@ -1208,6 +1239,9 @@ struct cmd_params_instance { use_mmap == other.use_mmap && use_direct_io == other.use_direct_io && devices == other.devices && no_host == other.no_host && + pshard == other.pshard && + max_vram_alloc == other.max_vram_alloc && + (!pshard || (n_prompt + n_gen + n_depth) == (other.n_prompt + other.n_gen + other.n_depth)) && vec_tensor_buft_override_equal(tensor_buft_overrides, other.tensor_buft_overrides); } @@ -1217,6 +1251,8 @@ struct cmd_params_instance { cparams.n_ctx = n_prompt + n_gen + n_depth; cparams.n_batch = n_batch; cparams.n_ubatch = n_ubatch; + cparams.n_threads = n_threads; + cparams.n_threads_batch = n_threads; cparams.type_k = type_k; cparams.type_v = type_v; cparams.offload_kqv = !no_kv_offload; @@ -1224,6 +1260,7 @@ struct cmd_params_instance { cparams.embeddings = embeddings; cparams.op_offload = !no_op_offload; cparams.swa_full = false; + cparams.pshard = pshard; return cparams; } @@ -1237,6 +1274,8 @@ static std::vector get_cmd_params_instances(const cmd_param for (const auto & m : params.model) for (const auto & fpt : params.fit_params_target) for (const auto & fpc : params.fit_params_min_ctx) + for (const auto & ps : params.pshard) + for (const auto & mva : params.max_vram_alloc) for (const auto & nl : params.n_gpu_layers) for (const auto & ncmoe : params.n_cpu_moe) for (const auto & sm : params.split_mode) @@ -1291,6 +1330,8 @@ static std::vector get_cmd_params_instances(const cmd_param /* .embeddings = */ embd, /* .no_op_offload= */ nopo, /* .no_host = */ noh, + /* .pshard = */ ps, + /* .max_vram_alloc = */ mva, /* .fit_target = */ fpt, /* .fit_min_ctx = */ fpc, }; @@ -1328,6 +1369,8 @@ static std::vector get_cmd_params_instances(const cmd_param /* .embeddings = */ embd, /* .no_op_offload= */ nopo, /* .no_host = */ noh, + /* .pshard = */ ps, + /* .max_vram_alloc = */ mva, /* .fit_target = */ fpt, /* .fit_min_ctx = */ fpc, }; @@ -1365,6 +1408,8 @@ static std::vector get_cmd_params_instances(const cmd_param /* .embeddings = */ embd, /* .no_op_offload= */ nopo, /* .no_host = */ noh, + /* .pshard = */ ps, + /* .max_vram_alloc = */ mva, /* .fit_target = */ fpt, /* .fit_min_ctx = */ fpc, }; @@ -1407,6 +1452,8 @@ struct test { bool embeddings; bool no_op_offload; bool no_host; + bool pshard; + size_t max_vram_alloc; size_t fit_target; uint32_t fit_min_ctx; int n_prompt; @@ -1447,6 +1494,8 @@ struct test { embeddings = inst.embeddings; no_op_offload = inst.no_op_offload; no_host = inst.no_host; + pshard = inst.pshard; + max_vram_alloc = inst.max_vram_alloc; fit_target = inst.fit_target; fit_min_ctx = inst.fit_min_ctx; n_prompt = inst.n_prompt; @@ -1506,7 +1555,8 @@ struct test { "type_k", "type_v", "n_gpu_layers", "n_cpu_moe", "split_mode", "main_gpu", "no_kv_offload", "flash_attn", "devices", "tensor_split", "tensor_buft_overrides", "use_mmap", "use_direct_io", "embeddings", - "no_op_offload", "no_host", "fit_target", "fit_min_ctx", + "no_op_offload", "no_host", "pshard", "max_vram_alloc", + "fit_target", "fit_min_ctx", "n_prompt", "n_gen", "n_depth", "test_time", "avg_ns", "stddev_ns", "avg_ts", "stddev_ts" }; @@ -1520,11 +1570,12 @@ struct test { field == "poll" || field == "model_size" || field == "model_n_params" || field == "n_gpu_layers" || field == "main_gpu" || field == "n_prompt" || field == "n_gen" || field == "n_depth" || field == "avg_ns" || field == "stddev_ns" || field == "no_op_offload" || field == "n_cpu_moe" || - field == "fit_target" || field == "fit_min_ctx") { + field == "fit_target" || field == "fit_min_ctx" || field == "max_vram_alloc") { return INT; } if (field == "f16_kv" || field == "no_kv_offload" || field == "cpu_strict" || field == "flash_attn" || - field == "use_mmap" || field == "use_direct_io" || field == "embeddings" || field == "no_host") { + field == "use_mmap" || field == "use_direct_io" || field == "embeddings" || field == "no_host" || + field == "pshard") { return BOOL; } if (field == "avg_ts" || field == "stddev_ts") { @@ -1601,6 +1652,8 @@ struct test { std::to_string(embeddings), std::to_string(no_op_offload), std::to_string(no_host), + std::to_string(pshard), + std::to_string(max_vram_alloc), std::to_string(fit_target), std::to_string(fit_min_ctx), std::to_string(n_prompt), @@ -1797,6 +1850,12 @@ struct markdown_printer : public printer { if (field == "no_host") { return 4; } + if (field == "pshard") { + return 4; + } + if (field == "max_vram_alloc") { + return 6; + } int width = std::max((int) field.length(), 10); @@ -1837,6 +1896,12 @@ struct markdown_printer : public printer { if (field == "no_host") { return "noh"; } + if (field == "pshard") { + return "psh"; + } + if (field == "max_vram_alloc") { + return "mva"; + } if (field == "devices") { return "dev"; } @@ -1930,6 +1995,12 @@ struct markdown_printer : public printer { if (params.no_host.size() > 1 || params.no_host != cmd_params_defaults.no_host) { fields.emplace_back("no_host"); } + if (params.pshard.size() > 1 || params.pshard != cmd_params_defaults.pshard) { + fields.emplace_back("pshard"); + } + if (params.max_vram_alloc.size() > 1 || params.max_vram_alloc != cmd_params_defaults.max_vram_alloc) { + fields.emplace_back("max_vram_alloc"); + } if (params.fit_params_target.size() > 1 || params.fit_params_target != cmd_params_defaults.fit_params_target) { fields.emplace_back("fit_target"); } @@ -2193,8 +2264,21 @@ int main(int argc, char ** argv) { std::vector params_instances = get_cmd_params_instances(params); - llama_model * lmodel = nullptr; - const cmd_params_instance * prev_inst = nullptr; + llama_model * lmodel = nullptr; + llama_pshard_plan_registry * active_pshard_registry = nullptr; + const cmd_params_instance * prev_inst = nullptr; + + auto free_active_model = [&]() { + if (lmodel) { + llama_model_free(lmodel); + lmodel = nullptr; + } + if (active_pshard_registry) { + llama_pshard_registry_free(active_pshard_registry); + active_pshard_registry = nullptr; + } + prev_inst = nullptr; + }; // store the llama_context state at the previous depth that we performed a test // ref: https://github.com/ggml-org/llama.cpp/pull/16944#issuecomment-3478151721 @@ -2210,19 +2294,22 @@ int main(int argc, char ** argv) { auto mparams = inst.to_llama_mparams(); auto cparams = inst.to_llama_cparams(); - bool do_fit = inst.fit_target != cmd_params_defaults.fit_params_target[0] || - inst.fit_min_ctx != cmd_params_defaults.fit_params_min_ctx[0]; + const bool fit_target_changed = inst.fit_target != cmd_params_defaults.fit_params_target[0]; + const bool fit_ctx_changed = inst.fit_min_ctx != cmd_params_defaults.fit_params_min_ctx[0]; + const bool do_fit = fit_target_changed || fit_ctx_changed; + + if (inst.pshard && fit_ctx_changed) { + fprintf(stderr, "%s: error: -pshard cannot be combined with -fitc\n", __func__); + free_active_model(); + return 1; + } std::vector fit_tensor_split(llama_max_devices(), 0.0f); std::vector fit_overrides(llama_max_tensor_buft_overrides(), {nullptr, nullptr}); - if (do_fit) { + if (do_fit && !inst.pshard) { // free the previous model so fit sees full free VRAM - if (lmodel) { - llama_model_free(lmodel); - lmodel = nullptr; - prev_inst = nullptr; - } + free_active_model(); // use default n_gpu_layers and n_ctx so llama_params_fit can adjust them mparams.n_gpu_layers = llama_model_default_params().n_gpu_layers; @@ -2243,28 +2330,66 @@ int main(int argc, char ** argv) { params.verbose ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); } + std::vector pshard_overrides; + llama_pshard_plan_registry * pending_pshard_registry = nullptr; + if (inst.pshard) { + // free the previous model so pshard probes see the same free VRAM that the cached plan used + free_active_model(); + + pshard_overrides.resize(4096, { nullptr, nullptr, -1 }); + + if (!params.n_gpu_layers_user) { + mparams.n_gpu_layers = llama_model_default_params().n_gpu_layers; + } + + const uint32_t n_ctx_plan = cparams.n_ctx; + const uint32_t tier_max_auto = std::min(std::max(cparams.n_batch, (uint32_t) 16384), n_ctx_plan); + + pending_pshard_registry = llama_pshard_registry_create(tier_max_auto, cparams.n_seq_max); + mparams.pshard_registry = pending_pshard_registry; + + llama_params_fit_pshard(inst.model.c_str(), &mparams, &cparams, + pshard_overrides.data(), inst.max_vram_alloc, inst.fit_target); + + if (!mparams.pshard) { + llama_pshard_registry_free(pending_pshard_registry); + pending_pshard_registry = nullptr; + mparams.pshard_registry = nullptr; + } + } + + cmd_params_instance run_inst = inst; + run_inst.n_batch = (int) cparams.n_batch; + run_inst.n_ubatch = (int) cparams.n_ubatch; + run_inst.pshard = mparams.pshard && cparams.pshard; + run_inst.max_vram_alloc = mparams.max_vram_alloc; + // keep the same model between tests when possible if (!lmodel || !prev_inst || !inst.equal_mparams(*prev_inst)) { - if (lmodel) { - llama_model_free(lmodel); - } + free_active_model(); lmodel = llama_model_load_from_file(inst.model.c_str(), mparams); if (lmodel == NULL) { fprintf(stderr, "%s: error: failed to load model '%s'\n", __func__, inst.model.c_str()); + if (pending_pshard_registry) { + llama_pshard_registry_free(pending_pshard_registry); + pending_pshard_registry = nullptr; + } return 1; } + active_pshard_registry = pending_pshard_registry; + pending_pshard_registry = nullptr; prev_inst = &inst; } llama_context * ctx = llama_init_from_model(lmodel, cparams); if (ctx == NULL) { fprintf(stderr, "%s: error: failed to create context with model '%s'\n", __func__, inst.model.c_str()); - llama_model_free(lmodel); + free_active_model(); return 1; } - test t(inst, lmodel, ctx); + test t(run_inst, lmodel, ctx); llama_memory_clear(llama_get_memory(ctx), false); @@ -2413,7 +2538,7 @@ int main(int argc, char ** argv) { ggml_threadpool_free_fn(threadpool); } - llama_model_free(lmodel); + free_active_model(); if (p) { p->print_footer();