diff --git a/common/arg.cpp b/common/arg.cpp index 9753441313a7..d3abe1411533 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -267,6 +267,13 @@ static void parse_tensor_buffer_overrides(const std::string & value, std::vector std::string tensor_name = override.substr(0, pos); std::string buffer_type = override.substr(pos + 1); + int32_t backend_id = -1; + auto colon_pos = buffer_type.rfind(':'); + if (colon_pos != std::string::npos) { + backend_id = std::stoi(buffer_type.substr(colon_pos + 1)); + buffer_type = buffer_type.substr(0, colon_pos); + } + if (buft_list.find(buffer_type) == buft_list.end()) { printf("Available buffer types:\n"); for (const auto & it : buft_list) { @@ -277,7 +284,7 @@ static void parse_tensor_buffer_overrides(const std::string & value, std::vector // keep strings alive and avoid leaking memory by storing them in a static vector static std::list buft_overrides; buft_overrides.push_back(tensor_name); - overrides.push_back({buft_overrides.back().c_str(), buft_list.at(buffer_type)}); + overrides.push_back({buft_overrides.back().c_str(), buft_list.at(buffer_type), backend_id}); } } @@ -859,11 +866,11 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context // pad tensor_buft_overrides for llama_params_fit: const size_t ntbo = llama_max_tensor_buft_overrides(); while (params.tensor_buft_overrides.size() < ntbo) { - params.tensor_buft_overrides.push_back({nullptr, nullptr}); + params.tensor_buft_overrides.push_back({nullptr, nullptr, -1}); } if (!params.speculative.draft.tensor_buft_overrides.empty()) { - params.speculative.draft.tensor_buft_overrides.push_back({nullptr, nullptr}); + params.speculative.draft.tensor_buft_overrides.push_back({nullptr, nullptr, -1}); } if (!params.chat_template.empty() && !common_chat_verify_template(params.chat_template, params.use_jinja)) { @@ -1846,6 +1853,27 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.warmup = value; } ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_MTMD, LLAMA_EXAMPLE_EMBEDDING, LLAMA_EXAMPLE_RETRIEVAL, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_DEBUG})); + add_opt(common_arg( + {"-pshard", "--pshard"}, + "enable pipelined sharding (weights on CPU, pipelined to GPU per split)", + [](common_params & params) { + params.pshard = true; + } + )); + add_opt(common_arg( + {"-mva", "--max-vram-alloc"}, "N", + "VRAM budget in MB for pshard (0 = use actual free VRAM minus --fit-target)", + [](common_params & params, int value) { + params.max_vram_alloc = value; + } + )); + add_opt(common_arg( + {"--pshard-tier-max"}, "N", + "max tier batch size for pshard planning (0 = auto, capped by n_ctx)", + [](common_params & params, int value) { + params.pshard_tier_max = value; + } + )); add_opt(common_arg( {"--spm-infill"}, string_format( @@ -2615,7 +2643,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex // keep strings alive and avoid leaking memory by storing them in a static vector static std::list buft_overrides; buft_overrides.push_back(llm_ffn_exps_block_regex(i)); - params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()}); + params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type(), -1}); } } ).set_env("LLAMA_ARG_N_CPU_MOE")); @@ -2729,8 +2757,8 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).set_examples({LLAMA_EXAMPLE_FIT_PARAMS}).set_env("LLAMA_ARG_FIT_ESTIMATE")); add_opt(common_arg( { "-fitt", "--fit-target" }, "MiB0,MiB1,MiB2,...", - string_format("target margin per device for --fit, comma-separated list of values, " - "single value is broadcast across all devices, default: %zu", params.fit_params_target[0]/(1024*1024)), + string_format("target margin per device for --fit and pshard auto budget, comma-separated list of values, " + "single value is broadcast across all devices, default: %zu", params.fit_params_target[0]/(1024*1024)), [](common_params & params, const std::string & value) { std::string arg_next = value; @@ -3939,7 +3967,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex for (int i = 0; i < value; ++i) { static std::list buft_overrides_draft; buft_overrides_draft.push_back(llm_ffn_exps_block_regex(i)); - params.speculative.draft.tensor_buft_overrides.push_back({buft_overrides_draft.back().c_str(), ggml_backend_cpu_buffer_type()}); + params.speculative.draft.tensor_buft_overrides.push_back({buft_overrides_draft.back().c_str(), ggml_backend_cpu_buffer_type(), -1}); } } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE")); diff --git a/common/common.cpp b/common/common.cpp index 82dd780fd8b3..58482d721134 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1219,7 +1219,7 @@ common_init_result::common_init_result(common_params & params, bool model_only) auto mparams = common_model_params_to_llama(params); auto cparams = common_context_params_to_llama(params); - if (params.fit_params) { + auto fit_params = [&]() { COM_TRC("%s", "fitting params to device memory ...\n"); COM_TRC("%s", "(for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n"); common_fit_params(params.model.path.c_str(), &mparams, &cparams, @@ -1228,6 +1228,30 @@ common_init_result::common_init_result(common_params & params, bool model_only) params.fit_params_target.data(), params.fit_params_min_ctx, params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); + }; + + if (params.pshard) { + LOG_INF("%s: pshard enabled, probing and loading plan cache\n", __func__); + params.tensor_buft_overrides.resize(4096); + mparams.pshard_registry = common_pshard_registry_create(params.pshard_tier_max, cparams.n_seq_max); + const size_t fit_target_mb = params.fit_params_target.empty() ? 0 : params.fit_params_target[0] / (1024 * 1024); + common_fit_params_pshard(params.model.path.c_str(), &mparams, &cparams, + params.tensor_buft_overrides.data(), params.max_vram_alloc, fit_target_mb); + if (!mparams.pshard) { + LOG_WRN("%s: pshard not active for this configuration\n", __func__); + common_pshard_registry_free(mparams.pshard_registry); + mparams.pshard_registry = nullptr; + if (params.fit_params) { + fit_params(); + } + } else { + params.n_batch = (int32_t) cparams.n_batch; + params.n_ubatch = (int32_t) cparams.n_ubatch; + LOG_INF("%s: pshard runtime batch/ubatch set to selected cache_ubatch=%u\n", + __func__, cparams.n_ubatch); + } + } else if (params.fit_params) { + fit_params(); } llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams); @@ -1581,6 +1605,8 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.progress_callback_user_data = params.load_progress_callback_user_data; mparams.no_alloc = params.no_alloc; + mparams.pshard = params.pshard; + mparams.max_vram_alloc = params.max_vram_alloc; return mparams; } @@ -1618,6 +1644,7 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.type_k = params.cache_type_k; cparams.type_v = params.cache_type_v; + cparams.pshard = params.pshard; return cparams; } diff --git a/common/common.h b/common/common.h index 2792521836ae..95db7707c8ff 100644 --- a/common/common.h +++ b/common/common.h @@ -742,6 +742,10 @@ struct common_params { llama_progress_callback load_progress_callback = NULL; void * load_progress_callback_user_data = NULL; bool no_alloc = false; // Don't allocate model buffers + + bool pshard = false; + size_t max_vram_alloc = 0; + uint32_t pshard_tier_max = 0; }; // call once at the start of a program if it uses libcommon @@ -1080,7 +1084,7 @@ inline std::string llm_ffn_exps_block_regex(int idx) { } inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() { - return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() }; + return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type(), -1 }; } // diff --git a/common/fit.cpp b/common/fit.cpp index c79221cb00fa..bef690c176f0 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -477,21 +477,24 @@ static void common_params_fit_impl( il0 += ngl_per_device[id].n_full(); for (uint32_t il = il0; il < il0 + ngl_per_device[id].n_part; il++) { if (itbo + 1 >= ntbo) { - tensor_buft_overrides[itbo].pattern = nullptr; - tensor_buft_overrides[itbo].buft = nullptr; + tensor_buft_overrides[itbo].pattern = nullptr; + tensor_buft_overrides[itbo].buft = nullptr; + tensor_buft_overrides[itbo].backend_id = -1; itbo++; mparams.tensor_buft_overrides = tensor_buft_overrides; throw common_params_fit_exception("llama_max_tensor_buft_overrides() == " + std::to_string(ntbo) + " is insufficient for model"); } - tensor_buft_overrides[itbo].pattern = get_overflow_pattern(il, il == il0 ? ngl_per_device[id].overflow_type : LAYER_FRACTION_MOE); - tensor_buft_overrides[itbo].buft = il == il0 ? overflow_bufts[id] : ggml_backend_cpu_buffer_type(); + tensor_buft_overrides[itbo].pattern = get_overflow_pattern(il, il == il0 ? ngl_per_device[id].overflow_type : LAYER_FRACTION_MOE); + tensor_buft_overrides[itbo].buft = il == il0 ? overflow_bufts[id] : ggml_backend_cpu_buffer_type(); + tensor_buft_overrides[itbo].backend_id = -1; itbo++; } il0 += ngl_per_device[id].n_part; } - tensor_buft_overrides[itbo].pattern = nullptr; - tensor_buft_overrides[itbo].buft = nullptr; + tensor_buft_overrides[itbo].pattern = nullptr; + tensor_buft_overrides[itbo].buft = nullptr; + tensor_buft_overrides[itbo].backend_id = -1; itbo++; mparams.tensor_buft_overrides = tensor_buft_overrides; }; @@ -527,8 +530,8 @@ static void common_params_fit_impl( if (hp_nex > 0) { const static std::string pattern_moe_all = "blk\\.\\d+\\.ffn_(up|down|gate_up|gate)_(ch|)exps"; // matches all MoE tensors ggml_backend_buffer_type_t cpu_buft = ggml_backend_cpu_buffer_type(); - tensor_buft_overrides[0] = {pattern_moe_all.c_str(), cpu_buft}; - tensor_buft_overrides[1] = {nullptr, nullptr}; + tensor_buft_overrides[0] = {pattern_moe_all.c_str(), cpu_buft, -1}; + tensor_buft_overrides[1] = {nullptr, nullptr, -1}; mparams->tensor_buft_overrides = tensor_buft_overrides; LOG_TRC("%s: getting device memory data with all MoE tensors moved to system memory:\n", __func__); @@ -549,7 +552,7 @@ static void common_params_fit_impl( } // reset - tensor_buft_overrides[0] = {nullptr, nullptr}; + tensor_buft_overrides[0] = {nullptr, nullptr, -1}; mparams->tensor_buft_overrides = tensor_buft_overrides; } @@ -811,6 +814,47 @@ enum common_params_fit_status common_fit_params( return status; } +llama_pshard_plan_registry * common_pshard_registry_create(uint32_t n_tier_max, uint32_t n_seq_max) { + return llama_pshard_registry_create(n_tier_max, n_seq_max); +} + +void common_pshard_registry_free(llama_pshard_plan_registry * registry) { + llama_pshard_registry_free(registry); +} + +// Adapter handed to the pshard planner in libllama, which cannot call into +// common/ directly. Throws on failure, same as common_params_fit_impl. +static void common_fit_pshard_baseline( + const char * path_model, llama_model_params * mparams, llama_context_params * cparams, + float * tensor_split, llama_model_tensor_buft_override * tensor_buft_overrides, + size_t * margins, uint32_t n_ctx_min, ggml_log_level log_level, + void * /*user_data*/) { + common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, + margins, n_ctx_min, log_level); +} + +void common_fit_params_pshard( + const char * path_model, + llama_model_params * mparams, + llama_context_params * cparams, + llama_model_tensor_buft_override * tensor_buft_overrides, + size_t max_vram_mb, + size_t fit_target_mb) { + llama_params_fit_pshard_inference(path_model, mparams, cparams, tensor_buft_overrides, + max_vram_mb, fit_target_mb); +} + +void common_pshard_plan( + const char * path_model, + llama_model_params * mparams, + llama_context_params * cparams, + llama_model_tensor_buft_override * tensor_buft_overrides, + size_t max_vram_mb, + size_t fit_target_mb) { + llama_params_fit_pshard_planning(path_model, mparams, cparams, tensor_buft_overrides, + max_vram_mb, fit_target_mb, &common_fit_pshard_baseline, /*fit_fn_ud =*/ nullptr); +} + void common_memory_breakdown_print(const struct llama_context * ctx) { //const auto & devices = ctx->get_model().devices; const auto * model = llama_get_model(ctx); diff --git a/common/fit.h b/common/fit.h index 208fc30694e0..5c17d60e8bd5 100644 --- a/common/fit.h +++ b/common/fit.h @@ -26,6 +26,47 @@ common_params_fit_status common_fit_params( uint32_t n_ctx_min, // minimum context size to set when trying to reduce memory use ggml_log_level log_level); // minimum log level to print during fitting, lower levels go to debug log +// Pipelined-sharding (pshard) variant of common_fit_params. +// +// Loads the plan registry cached next to the model as +// .tensor_overrides.pshard_registry (written by llama-fit-params --pshard) +// and fills tensor_buft_overrides for model loading. If no usable plan applies, +// clears mparams->pshard and leaves baseline loading in place. This is the +// RUNTIME path - it never plans, so startup stays fast. +// +// mparams->pshard_registry must be set by the caller (common_pshard_registry_create). +// +// Create/free the tier plan registry. The caller owns the pointer and assigns it +// to mparams->pshard_registry before calling common_fit_params_pshard. +struct llama_pshard_plan_registry * common_pshard_registry_create(uint32_t n_tier_max, uint32_t n_seq_max); +void common_pshard_registry_free (struct llama_pshard_plan_registry * registry); + +void common_fit_params_pshard( + const char * path_model, + llama_model_params * mparams, + llama_context_params * cparams, + llama_model_tensor_buft_override * tensor_buft_overrides, + size_t max_vram_mb, // 0 = use actual free VRAM minus fit_target_mb + size_t fit_target_mb); // ignored when max_vram_mb > 0 + +// PLANNER path (llama-fit-params --pshard). +// +// Probes VRAM for each strategy/tier combination, picks the best plan per tier and +// writes/updates .tensor_overrides.pshard_registry, then fills +// tensor_buft_overrides. Loads any already-cached tiers first and only plans what is +// missing. If everything already fits in VRAM it clears mparams->pshard and falls back +// to baseline loading. +// +// Supplies the baseline params-fitting engine (common_fit_params) to the planner, which +// lives in libllama and cannot call into common/ directly. +void common_pshard_plan( + const char * path_model, + llama_model_params * mparams, + llama_context_params * cparams, + llama_model_tensor_buft_override * tensor_buft_overrides, + size_t max_vram_mb, + size_t fit_target_mb); + // print estimated memory to stdout void common_fit_print( const char * path_model, diff --git a/docs/pshard.md b/docs/pshard.md new file mode 100644 index 000000000000..f18130f01dc5 --- /dev/null +++ b/docs/pshard.md @@ -0,0 +1,136 @@ +# Pipeline Sharding (pshard) + +pshard enables model inference at user-specified VRAM budget. + +The overflow weights are spilled to the pinned host memory which are streamed into device memory (VRAM) for GPU execution during runtime as per different strategies. + +An offline planner probes candidate placements (host/device) against measured hardware characteristics and writes plans for different tiers; the runtime loads the appropriate plan. This gives three stages: + + +| Stage | Tool | Artifact | Scope | +| ------- | ------------------------------------------ | ----------------------------------------------- | -------------------------- | +| Profile | `llama-profiler-gpu`, `llama-profiler-cpu` | `gpu_profile.txt`, `cpu_profile.txt` | per machine | +| Plan | `llama-fit-params --pshard` | `.gguf.tensor_overrides.pshard_registry` | per model + context config | +| Run | `llama-cli`, `llama-server`, `llama-bench` | — | per invocation | + + +A plan holds one entry per batch-size *tier*; the runtime selects the smallest tier covering +the current batch. Each entry is a set of tensor→backend assignments: GPU-resident, one of two +shard lanes used for pipelining, or host-resident. + +## Limitations + +- Single CUDA device. Multi-GPU is not supported. +- Architectures using the DSA or DSV4 KV cache (DeepSeek family) are not supported; pshard +currently declines without a diagnostic. +- Host RAM must hold the full model. + + + +## Build + +Standard CUDA build — see [docs/build.md](build.md). The pshard tools are part of the default +target set. + +```bash +cmake -B build -DGGML_CUDA=ON +cmake --build build --config Release -j +``` + +--- + + + +## 1. Profile the machine + +The planner's cost model needs measured bandwidth and throughput, not vendor specs. The +profilers write `gpu_profile.txt` and `cpu_profile.txt` to the working directory, and the +planner resolves them from *its* working directory — run both from the same place. + +```bash +./build/bin/llama-profiler-gpu +./build/bin/llama-profiler-cpu +``` + +`--fast` (default) sweeps a reduced configuration set; `--full` is exhaustive and considerably +slower. `--output` redirects, in which case point the planner at the results via +`PSHARD_CPU_PROFILE` / `PSHARD_GPU_PROFILE`. + +The CPU profiler measures concurrent DRAM and PCIe bandwidth under load, so it takes a few +minutes. Results characterise the machine and are model-independent — generate once. + +## 2. Plan a placement + +```bash +./build/bin/llama-fit-params --pshard \ + --model models/your-model.gguf \ + --max-vram-alloc 8000 \ + -c 8192 -t 8 -fa on +``` + +The planner probes each supported strategy across the tier range, estimates throughput per +candidate from the profile data, and keeps the best viable plan per tier. Output is a plain-text +registry written next to the model. + +`--max-vram-alloc` sets the planning budget in MiB. Omitted, the planner uses VRAM free at +planning time minus `--fit-target`, which makes the resulting plan dependent on whatever else +was resident on the device. Set it explicitly for reproducible plans. + +Without profile data the planner degrades to VRAM-fit selection only, ignoring throughput. It +still produces a usable plan, but not necessarily the fastest one. + +Strategy descriptions, the registry format, and planning for `llama-bench` shapes are covered +in [tools/fit-params/README-pshard.md](../tools/fit-params/README-pshard.md). + +## 3. Run + +```bash +./build/bin/llama-cli --pshard \ + --model models/your-model.gguf \ + --max-vram-alloc 8000 \ + -c 8192 -t 8 -fa on +``` + +`llama-server` and `llama-bench` accept the same flags. + +### Plan/runtime fingerprint + +The registry is keyed by a hash of the parameters that affect plan validity: + +``` +n_ctx n_seq_max n_threads flash_attn_type type_k type_v model_file_size $PSHARD_STRATEGY +``` + +If the runtime's values hash differently there is no matching entry, and **pshard disables +itself and proceeds with a conventional load** — which then requires full VRAM. The fallback is +logged at warning level only, so it is easy to miss. + +Consequently `-c`, `-np`, `-t`, `-fa`, `-ctk`, `-ctv` and the model must be identical between +planning and running. `-t` is the most common oversight, since its default is derived and +varies by tool; pass it explicitly to both. + +`--max-vram-alloc` is deliberately excluded from the fingerprint. A single fingerprint can hold +several budget variants, and re-planning at a new budget replaces only that variant. + +### Confirming pshard is active + +``` +common_init_result: pshard enabled, probing and loading plan cache +llama_params_fit_pshard_inference: loaded 7 tier plans from cache (variant budget=8000 MiB ...) +common_init_result: pshard runtime batch/ubatch set to selected cache_ubatch=8192 +``` + +`no matching plan cache ... disabling pshard` or `pshard not active for this configuration` +indicates the fallback described above. `llama-bench` reports state directly in its `psh` +column. + +--- + + + +## See also + +- [docs/build.md](build.md) +- [tools/fit-params/README-pshard.md](../tools/fit-params/README-pshard.md) — planner +reference: strategies, registry format, `--bench-plan` + diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 39f802d250e1..49ca2fff4e36 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -34,6 +34,7 @@ else() add_subdirectory(gen-docs) add_subdirectory(training) add_subdirectory(diffusion) + add_subdirectory(llama-profiler) if (NOT GGML_BACKEND_DL) add_subdirectory(convert-llama2c-to-ggml) # these examples use the backends directly and cannot be built with dynamic loading diff --git a/examples/llama-profiler/CMakeLists.txt b/examples/llama-profiler/CMakeLists.txt new file mode 100644 index 000000000000..0e4a4ab686ad --- /dev/null +++ b/examples/llama-profiler/CMakeLists.txt @@ -0,0 +1,19 @@ +set(TARGET_CPU llama-profiler-cpu) +add_executable(${TARGET_CPU} profiler-cpu.cpp) +install(TARGETS ${TARGET_CPU} RUNTIME) +target_link_libraries(${TARGET_CPU} PRIVATE ggml ggml-base llama-common ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET_CPU} PRIVATE cxx_std_17) +if (GGML_CUDA) + target_link_libraries(${TARGET_CPU} PRIVATE ggml-cuda) + target_compile_definitions(${TARGET_CPU} PRIVATE GGML_USE_CUDA=1) +endif() + +set(TARGET_GPU llama-profiler-gpu) +add_executable(${TARGET_GPU} profiler-gpu.cpp) +install(TARGETS ${TARGET_GPU} RUNTIME) +target_link_libraries(${TARGET_GPU} PRIVATE ggml ggml-base ${CMAKE_THREAD_LIBS_INIT}) +if (GGML_CUDA) + target_link_libraries(${TARGET_GPU} PRIVATE ggml-cuda) + target_compile_definitions(${TARGET_GPU} PRIVATE GGML_USE_CUDA=1) +endif() +target_compile_features(${TARGET_GPU} PRIVATE cxx_std_17) diff --git a/examples/llama-profiler/profiler-common.h b/examples/llama-profiler/profiler-common.h new file mode 100644 index 000000000000..ae22f7c32f79 --- /dev/null +++ b/examples/llama-profiler/profiler-common.h @@ -0,0 +1,161 @@ +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +struct bench_timer { + using clk = std::chrono::high_resolution_clock; + clk::time_point t0; + void start() { t0 = clk::now(); } + double stop() { return std::chrono::duration(clk::now() - t0).count(); } +}; + +struct bench_result { + std::string op_name; + std::string quant_type; + + int N = 0, K = 0, B = 0; + int n_tokens = 0, ctx_len = 0, n_heads = 0, head_dim = 0; + int64_t n_elements = 0; + + double ops = 0.0; + double bytes = 0.0; + double time_s = 0.0; + + float arithmetic_intensity = 0.0f; + float effective_gflops = 0.0f; + float effective_bw_gb_s = 0.0f; + + void calculate_derived() { + arithmetic_intensity = (bytes > 0.0) ? (float)(ops / bytes) : 0.0f; + effective_gflops = (time_s > 0.0) ? (float)(ops / time_s / 1e9) : 0.0f; + effective_bw_gb_s = (time_s > 0.0) ? (float)(bytes / time_s / 1e9) : 0.0f; + } + + void print_dims() const { + if (op_name.find("MUL_MAT_ID") != std::string::npos) { + printf(" [N=%d K=%d experts=%d/%d B=%d]", N, K, n_tokens, ctx_len, B); + } else if (op_name.find("MUL_MAT") != std::string::npos) { + printf(" [N=%d K=%d B=%d]", N, K, B); + } else if (op_name.find("FLASH_ATTN") != std::string::npos) { + printf(" [tokens=%d ctx=%d heads=%d dim=%d]", n_tokens, ctx_len, n_heads, head_dim); + } else { + printf(" [n=%lld]", (long long)n_elements); + } + } +}; + +struct matmul_size { int32_t N; int32_t K; }; +struct moe_config { int32_t N; int32_t K; int32_t n_experts; int32_t n_experts_used; }; +struct attn_config { const char * name; int32_t n_q_heads; int32_t n_kv_heads; int32_t head_dim; }; + +inline std::vector get_matmul_sizes(bool fast) { + if (fast) { + return { + { 1024, 1024}, { 2048, 2048}, { 4096, 4096}, { 8192, 8192}, + { 512, 2048}, { 8192, 4096}, { 14336, 4096}, { 4096, 14336}, + {128256, 4096}, + }; + } + return { + { 1024, 1024}, { 2048, 2048}, { 4096, 4096}, { 8192, 8192}, { 16384, 16384}, + { 2048, 1024}, { 4096, 2048}, { 8192, 4096}, { 14336, 4096}, { 16384, 8192}, + { 22016, 4096}, { 28672, 8192}, + { 1024, 2048}, { 2048, 4096}, { 4096, 8192}, { 4096, 14336}, { 4096, 22016}, + { 8192, 16384}, { 8192, 28672}, + { 1024, 512}, { 1024, 4096}, + { 1536, 4096}, { 2048, 7168}, { 8192, 5120}, + { 32000, 4096}, {128256, 4096}, {128256, 8192}, {151936, 1024}, {151936, 8192}, + {248320, 4096}, + }; +} + +inline std::vector get_moe_configs(bool fast) { + if (fast) { + return { + { 2048, 7168, 256, 8}, { 1536, 4096, 128, 8}, {14336, 4096, 8, 2}, + { 8192, 5120, 16, 1}, { 1024, 1024, 128, 8}, { 2048, 1024, 128, 8}, + { 4096, 4096, 128, 8}, + }; + } + return { + { 1024, 4096, 512, 10}, { 1536, 4096, 128, 8}, { 2048, 7168, 256, 8}, + { 8192, 5120, 16, 1}, { 8192, 5120, 128, 1}, {14336, 4096, 8, 2}, + { 1024, 1024, 128, 8}, { 2048, 1024, 128, 8}, { 4096, 4096, 128, 8}, + }; +} + +inline std::vector get_attn_configs(bool fast) { + if (fast) { + return { + {"MHA", 32, 32, 128}, {"GQA-8", 32, 8, 128}, + {"GQA-4", 32, 4, 128}, {"MQA", 32, 1, 128}, + }; + } + return { + {"MHA", 32, 32, 128}, {"GQA-16", 32, 16, 128}, {"GQA-8", 32, 8, 128}, + {"GQA-4", 64, 4, 128}, {"GQA-2", 32, 2, 256}, {"MQA", 32, 1, 128}, + }; +} + +inline std::vector get_attn_ctx_lens(bool fast) { + if (fast) return { 1024, 4096, 8192, 16384 }; + return { 1024, 2048, 4096, 8192, 16384, 32768, 65536 }; +} + +inline std::vector get_matmul_quants(bool fast) { + if (fast) { + return { GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, GGML_TYPE_Q5_0, GGML_TYPE_Q2_K, GGML_TYPE_MXFP4 }; + } + return { GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q2_K, GGML_TYPE_MXFP4 }; +} + +inline bool parse_int_arg(const char * s, int32_t & out) { + char * end = nullptr; + long val = strtol(s, &end, 10); + if (end == s || *end != '\0' || val < INT32_MIN || val > INT32_MAX) { + return false; + } + out = (int32_t)val; + return true; +} + +struct ridge_result { + std::string key; + double peak_gflops; + double measured_bw; + double ridge; +}; + +template +inline std::vector compute_ridge_points( + const std::vector & results, double measured_bw) { + std::map groups; + + for (const auto & r : results) { + std::string key = r.op_name + "_" + r.quant_type; + groups[key] = std::max(groups[key], (double)r.effective_gflops); + } + + std::vector out; + for (const auto & [key, peak] : groups) { + double ridge = (measured_bw > 0.0) ? (peak / measured_bw) : 0.0; + out.push_back({key, peak, measured_bw, ridge}); + } + return out; +} + +inline std::vector create_quantized_data(ggml_type type, int64_t n_elements) { + size_t quant_size = ggml_row_size(type, n_elements); + return std::vector(quant_size, 0); +} diff --git a/examples/llama-profiler/profiler-cpu.cpp b/examples/llama-profiler/profiler-cpu.cpp new file mode 100644 index 000000000000..b465bad85433 --- /dev/null +++ b/examples/llama-profiler/profiler-cpu.cpp @@ -0,0 +1,595 @@ +#define CPU_WARMUP_ITERS 2 +#define CPU_TIMED_ITERS 2 + +#include "profiler-common.h" + +#include "common.h" +#include "ggml-cpu.h" + +#ifdef GGML_USE_CUDA +#include "ggml-cuda.h" +#endif + +#include +#include +#include +#include + +#if defined(_MSC_VER) +#include +#include +#endif + +static std::vector g_flush_buffer; + +static void flush_caches() { + if (g_flush_buffer.empty()) return; + + volatile char sum = 0; + for (size_t i = 0; i < g_flush_buffer.size(); i += 64) { + sum += g_flush_buffer[i]; + g_flush_buffer[i] = (char)(i & 0xFF); + } + g_flush_buffer[0] = sum; + +#if defined(_MSC_VER) + _mm_mfence(); +#elif defined(__GNUC__) || defined(__clang__) + __sync_synchronize(); +#endif +} + +static void init_flush_buffer() { + const size_t flush_size = 256 * 1024 * 1024; + g_flush_buffer.resize(flush_size); + for (size_t i = 0; i < flush_size; i += 4096) { + g_flush_buffer[i] = (char)(i & 0xFF); + } +} + +static double benchmark_cpu_dram_bandwidth(int threads) { + const size_t pool_bytes = 1024ULL * 1024 * 1024; + const size_t chunk_per_thread = pool_bytes / threads; + const int iterations = 10; + + std::vector pool(pool_bytes); + for (size_t i = 0; i < pool.size(); i += 4096) { + pool[i] = (uint8_t)(i & 0xFF); + } + + std::vector workers; + std::vector thread_bytes(threads, 0.0); + + bench_timer t; + t.start(); + + for (int tid = 0; tid < threads; ++tid) { + workers.emplace_back([&pool, &thread_bytes, tid, chunk_per_thread, iterations, pool_bytes]() { + const size_t start = tid * chunk_per_thread; + const size_t end_pos = (tid == (int)(pool_bytes / chunk_per_thread) - 1) ? pool_bytes : (start + chunk_per_thread); + volatile uint64_t local_sink = 0; + double local_bytes = 0.0; + for (int iter = 0; iter < iterations; ++iter) { + const size_t limit = end_pos - sizeof(uint64_t); + for (size_t offset = start; offset + 64 <= limit; offset += 64) { + local_sink += *(const uint64_t *)(pool.data() + offset); + local_bytes += 64.0; + } + } + thread_bytes[tid] = local_bytes; + (void)local_sink; + }); + } + for (auto & w : workers) w.join(); + + double elapsed = t.stop(); + double total_bytes = 0.0; + for (int i = 0; i < threads; i++) total_bytes += thread_bytes[i]; + return total_bytes / elapsed / 1e9; +} + +struct pcie_stress_ctx { + std::atomic active{false}; + std::atomic stop{false}; + + ggml_backend_t gpu_backend = nullptr; + ggml_backend_buffer_t host_buf = nullptr; + ggml_backend_buffer_t dev_buf = nullptr; + ggml_tensor * h_tensor = nullptr; + ggml_tensor * d_tensor = nullptr; + ggml_context * ctx = nullptr; + size_t transfer_size = 256 * 1024 * 1024; + double calibrated_bw_gb_s = 0.0; +}; + +static void pcie_stress_loop(pcie_stress_ctx * pcie) { + pcie->active.store(true, std::memory_order_release); + while (!pcie->stop.load(std::memory_order_acquire)) { + ggml_backend_tensor_set_async(pcie->gpu_backend, pcie->d_tensor, + pcie->h_tensor->data, 0, pcie->transfer_size); + ggml_backend_synchronize(pcie->gpu_backend); + ggml_backend_tensor_get_async(pcie->gpu_backend, pcie->d_tensor, + pcie->h_tensor->data, 0, pcie->transfer_size); + ggml_backend_synchronize(pcie->gpu_backend); + } + pcie->active.store(false, std::memory_order_release); +} + +static void calibrate_pcie(pcie_stress_ctx * pcie) { + printf("Calibrating standalone PCIe bandwidth...\n"); + bench_timer t; + t.start(); + const int cal_iterations = 20; + for (int i = 0; i < cal_iterations; ++i) { + ggml_backend_tensor_set_async(pcie->gpu_backend, pcie->d_tensor, + pcie->h_tensor->data, 0, pcie->transfer_size); + ggml_backend_synchronize(pcie->gpu_backend); + ggml_backend_tensor_get_async(pcie->gpu_backend, pcie->d_tensor, + pcie->h_tensor->data, 0, pcie->transfer_size); + ggml_backend_synchronize(pcie->gpu_backend); + } + double elapsed = t.stop(); + double bytes_moved = (double)cal_iterations * pcie->transfer_size * 2.0; + pcie->calibrated_bw_gb_s = bytes_moved / elapsed / 1e9; + printf(" Standalone PCIe BW: %.1f GB/s\n\n", pcie->calibrated_bw_gb_s); +} + +struct bench_result_cpu : bench_result { + int threads = 0; + float standalone_gflops = 0.0f; + float concurrent_gflops = 0.0f; + float concurrent_efficiency_pct = 0.0f; + float pcie_standalone_bw_gb_s = 0.0f; + + void print(double pcie_bw_ref = 0.0) const { + printf("%-20s quant=%-6s threads=%d AI=%.3f FLOP/byte BW=%.2f GB/s Perf=%.2f GFLOP/s", + op_name.c_str(), quant_type.c_str(), threads, + arithmetic_intensity, effective_bw_gb_s, effective_gflops); + if (concurrent_gflops > 0) { + printf(" | Concur=%.2f (%.1f%%)", concurrent_gflops, concurrent_efficiency_pct); + if (pcie_bw_ref > 0) { + double est_pcie_bw = pcie_bw_ref * (concurrent_efficiency_pct / 100.0) * 0.9; + printf(" PCIe~%.1f GB/s (%.1f%%)", est_pcie_bw, 100.0 * est_pcie_bw / pcie_bw_ref); + } + } + print_dims(); + printf("\n"); + } +}; + +static double benchmark_mul_mat_raw( + ggml_backend_t be, int N, int K, int batch_size, + ggml_type quant, int threads, + double * out_time_s, double * out_ops, double * out_bytes) { + + ggml_init_params params = { 4096ULL * 1024 * 1024, NULL, true }; + ggml_context * ctx = ggml_init(params); + ggml_backend_cpu_set_n_threads(be, threads); + + ggml_tensor * A = ggml_new_tensor_2d(ctx, quant, K, N); + ggml_tensor * B_tensor = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, batch_size); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_tensor * C = ggml_mul_mat(ctx, A, B_tensor); + ggml_build_forward_expand(gf, C); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, be); + if (!buffer) { + printf("SKIPPED: MUL_MAT N=%d K=%d B=%d %s (alloc failed)\n", N, K, batch_size, ggml_type_name(quant)); + ggml_free(ctx); return 0.0; + } + + std::vector A_data = create_quantized_data(quant, (int64_t)K * N); + std::vector B_data(K * batch_size, 1.0f); + ggml_backend_tensor_set(A, A_data.data(), 0, ggml_nbytes(A)); + ggml_backend_tensor_set(B_tensor, B_data.data(), 0, ggml_nbytes(B_tensor)); + + for (int i = 0; i < CPU_WARMUP_ITERS; ++i) { flush_caches(); ggml_backend_graph_compute_async(be, gf); } + + double total_time = 0.0; + bench_timer t; + for (int i = 0; i < CPU_TIMED_ITERS; ++i) { + flush_caches(); + t.start(); + ggml_backend_graph_compute_async(be, gf); + total_time += t.stop(); + } + + double time_per_iter = total_time / CPU_TIMED_ITERS; + double ops_total = 2.0 * N * K * batch_size; + double bytes_total = (double)(ggml_nbytes(A) + ggml_nbytes(B_tensor) + ggml_nbytes(C)); + + if (out_time_s) *out_time_s = time_per_iter; + if (out_ops) *out_ops = ops_total; + if (out_bytes) *out_bytes = bytes_total; + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return ops_total / time_per_iter / 1e9; +} + +static double benchmark_mul_mat_id_raw( + ggml_backend_t be, int N, int K, int n_experts, int n_experts_used, + int batch_size, ggml_type quant, int threads, + double * out_time_s, double * out_ops, double * out_bytes) { + + ggml_init_params params = { 8192ULL * 1024 * 1024, NULL, true }; + ggml_context * ctx = ggml_init(params); + ggml_backend_cpu_set_n_threads(be, threads); + + ggml_tensor * A = ggml_new_tensor_3d(ctx, quant, K, N, n_experts); + ggml_tensor * B = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, K, 1, batch_size); + ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_experts_used, batch_size); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, ggml_mul_mat_id(ctx, A, B, ids)); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, be); + if (!buffer) { + printf("SKIPPED: MUL_MAT_ID N=%d K=%d B=%d (alloc failed)\n", N, K, batch_size); + ggml_free(ctx); return 0.0; + } + + std::vector A_data = create_quantized_data(quant, (int64_t)K * N * n_experts); + ggml_backend_tensor_set(A, A_data.data(), 0, ggml_nbytes(A)); + std::vector B_data(K * batch_size, 1.0f); + ggml_backend_tensor_set(B, B_data.data(), 0, ggml_nbytes(B)); + std::vector ids_data(n_experts_used * batch_size); + for (int i = 0; i < n_experts_used * batch_size; i++) ids_data[i] = i % n_experts; + ggml_backend_tensor_set(ids, ids_data.data(), 0, ggml_nbytes(ids)); + + for (int i = 0; i < CPU_WARMUP_ITERS; ++i) { flush_caches(); ggml_backend_graph_compute_async(be, gf); } + + double total_time = 0.0; + bench_timer t; + for (int i = 0; i < CPU_TIMED_ITERS; ++i) { + flush_caches(); + t.start(); + ggml_backend_graph_compute_async(be, gf); + total_time += t.stop(); + } + + double time_per_iter = total_time / CPU_TIMED_ITERS; + double ops_total = 2.0 * N * K * batch_size * n_experts_used; + double bytes_total = (double)(ggml_nbytes(A) * n_experts_used / n_experts) + ggml_nbytes(B) + (double)(N * batch_size * n_experts_used * 4); + + if (out_time_s) *out_time_s = time_per_iter; + if (out_ops) *out_ops = ops_total; + if (out_bytes) *out_bytes = bytes_total; + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return ops_total / time_per_iter / 1e9; +} + +static double benchmark_flash_attn_raw( + ggml_backend_t be, int n_tokens, int ctx_len, + int n_q_heads, int n_kv_heads, int head_dim, + ggml_type kv_quant, int threads, + double * out_time_s, double * out_ops, double * out_bytes) { + + ggml_init_params params = { 8192ULL * 1024 * 1024, NULL, true }; + ggml_context * ctx = ggml_init(params); + ggml_backend_cpu_set_n_threads(be, threads); + + ggml_tensor * Q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, head_dim, n_tokens, n_q_heads, 1); + ggml_tensor * K = ggml_new_tensor_4d(ctx, kv_quant, head_dim, ctx_len, n_kv_heads, 1); + ggml_tensor * V = ggml_new_tensor_4d(ctx, kv_quant, head_dim, ctx_len, n_kv_heads, 1); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_tensor * out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, 1.0f / sqrtf((float)head_dim), 0.0f, 0.0f); + ggml_build_forward_expand(gf, out); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, be); + if (!buffer) { + printf("SKIPPED: FLASH_ATTN (alloc failed)\n"); + ggml_free(ctx); return 0.0; + } + + std::vector Q_data(ggml_nelements(Q), 1.0f); + ggml_backend_tensor_set(Q, Q_data.data(), 0, ggml_nbytes(Q)); + std::vector KV_data = create_quantized_data(kv_quant, ggml_nelements(K)); + ggml_backend_tensor_set(K, KV_data.data(), 0, ggml_nbytes(K)); + ggml_backend_tensor_set(V, KV_data.data(), 0, ggml_nbytes(V)); + + for (int i = 0; i < CPU_WARMUP_ITERS; ++i) { flush_caches(); ggml_backend_graph_compute_async(be, gf); } + + double total_time = 0.0; + bench_timer t; + for (int i = 0; i < CPU_TIMED_ITERS; ++i) { + flush_caches(); + t.start(); + ggml_backend_graph_compute_async(be, gf); + total_time += t.stop(); + } + + double time_per_iter = total_time / CPU_TIMED_ITERS; + double ops_total = 2.0 * n_tokens * head_dim * ctx_len * n_q_heads * 2; + double bytes_total = (double)(ggml_nbytes(Q) + ggml_nbytes(K) + ggml_nbytes(V) + ggml_nbytes(out)); + + if (out_time_s) *out_time_s = time_per_iter; + if (out_ops) *out_ops = ops_total; + if (out_bytes) *out_bytes = bytes_total; + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return ops_total / time_per_iter / 1e9; +} + +static bench_result_cpu run_concurrent( + std::function bench_fn, + const std::string & op_name, const char * quant_name, int threads, + pcie_stress_ctx * pcie) { + + bench_result_cpu result; + result.op_name = op_name; + result.quant_type = quant_name; + result.threads = threads; + result.pcie_standalone_bw_gb_s = pcie ? (float)pcie->calibrated_bw_gb_s : 0.0f; + + double standalone = bench_fn(&result.time_s, &result.ops, &result.bytes); + result.calculate_derived(); + result.standalone_gflops = result.effective_gflops; + + if (pcie && pcie->gpu_backend) { + pcie->stop.store(false, std::memory_order_release); + std::thread pcie_thread(pcie_stress_loop, pcie); + while (!pcie->active.load(std::memory_order_acquire)) std::this_thread::yield(); + + result.concurrent_gflops = (float)bench_fn(nullptr, nullptr, nullptr); + + pcie->stop.store(true, std::memory_order_release); + pcie_thread.join(); + } else { + result.concurrent_gflops = result.standalone_gflops; + } + + result.concurrent_efficiency_pct = (standalone > 0.0) + ? (float)(std::min)(100.0, 100.0 * result.concurrent_gflops / result.standalone_gflops) + : 100.0f; + + return result; +} + +static void run_matmul_benchmarks( + ggml_backend_t be, int threads, const std::vector & batch_sizes, + bool fast, pcie_stress_ctx * pcie, + std::vector & results) { + + auto sizes = get_matmul_sizes(fast); + auto quants = get_matmul_quants(fast); + + printf("=== MUL_MAT Operations ===\n\n"); + for (ggml_type qt : quants) { + printf("--- Quantization: %s ---\n", ggml_type_name(qt)); + for (int32_t bs : batch_sizes) { + printf(" [Batch=%d]\n", bs); + for (const auto & sz : sizes) { + if (qt == GGML_TYPE_Q2_K && (sz.K % 256 != 0)) continue; + + auto res = run_concurrent( + [&](double * t, double * o, double * b) { + return benchmark_mul_mat_raw(be, sz.N, sz.K, bs, qt, threads, t, o, b); + }, "MUL_MAT", ggml_type_name(qt), threads, pcie); + res.N = sz.N; res.K = sz.K; res.B = bs; + res.print(pcie ? pcie->calibrated_bw_gb_s : 0.0); + results.push_back(res); + } + } + printf("\n"); + } +} + +static void run_moe_benchmarks( + ggml_backend_t be, int threads, const std::vector & batch_sizes, + bool fast, pcie_stress_ctx * pcie, + std::vector & results) { + + auto configs = get_moe_configs(fast); + auto quants = get_matmul_quants(fast); + + printf("=== MUL_MAT_ID Operations (MoE) ===\n\n"); + for (ggml_type qt : quants) { + printf("--- MoE Quantization: %s ---\n", ggml_type_name(qt)); + for (int32_t bs : batch_sizes) { + printf(" [Batch=%d]\n", bs); + for (const auto & cfg : configs) { + if (qt == GGML_TYPE_Q2_K && (cfg.K % 256 != 0)) continue; + + auto res = run_concurrent( + [&](double * t, double * o, double * b) { + return benchmark_mul_mat_id_raw(be, cfg.N, cfg.K, cfg.n_experts, cfg.n_experts_used, bs, qt, threads, t, o, b); + }, "MUL_MAT_ID", ggml_type_name(qt), threads, pcie); + res.N = cfg.N; res.K = cfg.K; res.B = bs; + res.n_tokens = cfg.n_experts_used; res.ctx_len = cfg.n_experts; + res.print(pcie ? pcie->calibrated_bw_gb_s : 0.0); + results.push_back(res); + } + } + printf("\n"); + } +} + +static void run_attention_benchmarks( + ggml_backend_t be, int threads, const std::vector & batch_sizes, + bool fast, pcie_stress_ctx * pcie, + std::vector & results) { + + auto configs = get_attn_configs(fast); + auto ctx_lens = get_attn_ctx_lens(fast); + + printf("=== FLASH_ATTN Operations ===\n\n"); + for (const auto & cfg : configs) { + printf("--- %s (n_q=%d, n_kv=%d, head_dim=%d) ---\n", cfg.name, cfg.n_q_heads, cfg.n_kv_heads, cfg.head_dim); + for (int32_t n_tok : batch_sizes) { + printf(" [n_tokens=%d]\n", n_tok); + for (int32_t cl : ctx_lens) { + auto res = run_concurrent( + [&](double * t, double * o, double * b) { + return benchmark_flash_attn_raw(be, n_tok, cl, cfg.n_q_heads, cfg.n_kv_heads, cfg.head_dim, GGML_TYPE_F16, threads, t, o, b); + }, std::string("FLASH_ATTN_") + cfg.name, ggml_type_name(GGML_TYPE_F16), threads, pcie); + res.n_tokens = n_tok; res.ctx_len = cl; res.n_heads = cfg.n_kv_heads; res.head_dim = cfg.head_dim; + res.print(pcie ? pcie->calibrated_bw_gb_s : 0.0); + results.push_back(res); + } + } + printf("\n"); + } +} + +static void save_results_cpu( + const char * path, + const std::vector & results, + const std::vector & batch_sizes, + int threads, double dram_bw, double pcie_standalone_bw, double pcie_concurrent_bw, double cpu_eff, + bool has_gpu) { + + FILE * f = fopen(path, "w"); + if (!f) { fprintf(stderr, "Failed to open %s for writing\n", path); return; } + + fprintf(f, "# Concurrent Profiling (threads=%d, batch_sizes=[", threads); + for (size_t i = 0; i < batch_sizes.size(); i++) + fprintf(f, "%d%s", batch_sizes[i], i + 1 < batch_sizes.size() ? "," : ""); + fprintf(f, "])\n"); + + fprintf(f, "# Measured Bandwidths Per Thread Count:\n"); + if (has_gpu) { + fprintf(f, "# Threads=%d: DRAM_BW=%.1f GB/s, PCIe_Standalone=%.1f GB/s, PCIe_Concurrent=%.1f GB/s (CPU_Eff=%.1f%%)\n", + threads, dram_bw, pcie_standalone_bw, pcie_concurrent_bw, cpu_eff); + } else { + fprintf(f, "# Threads=%d: DRAM_BW=%.1f GB/s\n", threads, dram_bw); + } + + fprintf(f, "# op_name quant threads AI(FLOP/byte) BW(GB/s) GFLOP/s Ridge(FLOP/byte) Concurrent_GFLOP/s PCIe_Concurrent_BW N K B n_tokens ctx_len n_heads head_dim n_elements\n"); + + auto ridges = compute_ridge_points(results, dram_bw); + std::map ridge_map; + for (const auto & rr : ridges) ridge_map[rr.key] = rr.ridge; + + for (const auto & r : results) { + std::string key = r.op_name + "_" + r.quant_type; + double ridge = ridge_map.count(key) ? ridge_map[key] : 0.0; + double est_pcie = pcie_standalone_bw * (r.standalone_gflops > 0 ? r.concurrent_gflops / r.standalone_gflops : 1.0) * 0.9; + + fprintf(f, "%s %s %d %.4f %.2f %.2f %.4f %.2f %.2f %d %d %d %d %d %d %d %lld\n", + r.op_name.c_str(), r.quant_type.c_str(), r.threads, + r.arithmetic_intensity, r.effective_bw_gb_s, r.effective_gflops, ridge, + r.concurrent_gflops, est_pcie, + r.N, r.K, r.B, r.n_tokens, r.ctx_len, r.n_heads, r.head_dim, (long long)r.n_elements); + } + + fclose(f); + printf("Results saved to %s (%zu benchmarks)\n", path, results.size()); +} + +int main(int argc, char ** argv) { + int32_t fixed_threads = -1; + bool fast_mode = true; + const char * output_path = "cpu_profile.txt"; + + for (int i = 1; i < argc; ++i) { + if (!strcmp(argv[i], "--threads") && i + 1 < argc) { + if (!parse_int_arg(argv[++i], fixed_threads) || fixed_threads <= 0) { + fprintf(stderr, "Invalid --threads value: %s\n", argv[i]); + return 1; + } + } else if (!strcmp(argv[i], "--fast")) { + fast_mode = true; + } else if (!strcmp(argv[i], "--full")) { + fast_mode = false; + } else if (!strcmp(argv[i], "-h") || !strcmp(argv[i], "--help")) { + printf("usage: %s [options]\n", argv[0]); + printf("\n"); + printf("options:\n"); + printf(" -h, --help\n"); + printf(" --fast fast mode with fewer configs (default)\n"); + printf(" --full full mode with all configs\n"); + printf(" --threads number of CPU threads (default: auto)\n"); + printf(" --output output file (default: cpu_profile.txt)\n"); + return 0; + } else if (!strcmp(argv[i], "--output") && i + 1 < argc) { + output_path = argv[++i]; + } + } + + int32_t default_threads = common_cpu_get_num_math(); + int threads = (fixed_threads > 0) ? fixed_threads : default_threads; + std::vector batch_sizes = { 1, 64, 512 }; + + printf("=== CPU Profiler (cold-cache) ===\n"); + printf("Threads: %d%s\n", threads, fixed_threads > 0 ? " (user)" : " (auto)"); + printf("Mode: %s\n\n", fast_mode ? "FAST" : "FULL"); + + init_flush_buffer(); + + ggml_backend_t cpu_be = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + if (!cpu_be) { fprintf(stderr, "Failed to initialize CPU backend\n"); return 1; } + + ggml_quantize_init(GGML_TYPE_Q2_K); + ggml_quantize_init(GGML_TYPE_Q4_0); + ggml_quantize_init(GGML_TYPE_Q4_1); + ggml_quantize_init(GGML_TYPE_Q5_0); + ggml_quantize_init(GGML_TYPE_Q8_0); + ggml_quantize_init(GGML_TYPE_MXFP4); + + pcie_stress_ctx pcie; + pcie.gpu_backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_GPU, nullptr); + bool has_gpu = (pcie.gpu_backend != nullptr); + + if (has_gpu) { + ggml_backend_buffer_type_t host_buft = ggml_backend_dev_host_buffer_type( + ggml_backend_get_device(pcie.gpu_backend)); + if (host_buft) { + pcie.host_buf = ggml_backend_buft_alloc_buffer(host_buft, pcie.transfer_size); + pcie.dev_buf = ggml_backend_alloc_buffer(pcie.gpu_backend, pcie.transfer_size); + ggml_init_params p = { pcie.transfer_size + 8 * 1024 * 1024, NULL, true }; + pcie.ctx = ggml_init(p); + pcie.h_tensor = ggml_new_tensor_1d(pcie.ctx, GGML_TYPE_F32, pcie.transfer_size / 4); + pcie.d_tensor = ggml_new_tensor_1d(pcie.ctx, GGML_TYPE_F32, pcie.transfer_size / 4); + ggml_backend_tensor_alloc(pcie.host_buf, pcie.h_tensor, ggml_backend_buffer_get_base(pcie.host_buf)); + ggml_backend_tensor_alloc(pcie.dev_buf, pcie.d_tensor, ggml_backend_buffer_get_base(pcie.dev_buf)); + std::vector init_data(pcie.transfer_size / 4, 1.0f); + ggml_backend_tensor_set(pcie.h_tensor, init_data.data(), 0, pcie.transfer_size); + printf("GPU: %s\n", ggml_backend_name(pcie.gpu_backend)); + } else { + has_gpu = false; + } + } + if (!has_gpu) printf("No GPU — standalone mode\n\n"); + + printf("Measuring DRAM bandwidth...\n"); + double dram_bw = benchmark_cpu_dram_bandwidth(threads); + printf(" DRAM BW: %.1f GB/s\n", dram_bw); + if (has_gpu) calibrate_pcie(&pcie); + + std::vector all_results; + bench_timer overall; + overall.start(); + + run_matmul_benchmarks(cpu_be, threads, batch_sizes, fast_mode, has_gpu ? &pcie : nullptr, all_results); + run_moe_benchmarks(cpu_be, threads, batch_sizes, fast_mode, has_gpu ? &pcie : nullptr, all_results); + run_attention_benchmarks(cpu_be, threads, batch_sizes, fast_mode, has_gpu ? &pcie : nullptr, all_results); + + double pcie_concurrent_bw = 0.0, cpu_eff = 100.0; + if (has_gpu && !all_results.empty()) { + double sum_s = 0.0, sum_c = 0.0; + for (const auto & r : all_results) { sum_s += r.standalone_gflops; sum_c += r.concurrent_gflops; } + cpu_eff = (sum_s > 0) ? 100.0 * sum_c / sum_s : 100.0; + pcie_concurrent_bw = pcie.calibrated_bw_gb_s * (cpu_eff / 100.0) * 0.9; + } + + printf("\nTotal time: %.1f s, %zu benchmarks\n", overall.stop(), all_results.size()); + + save_results_cpu(output_path, all_results, batch_sizes, threads, dram_bw, + pcie.calibrated_bw_gb_s, pcie_concurrent_bw, cpu_eff, has_gpu); + + if (has_gpu) { + if (pcie.ctx) ggml_free(pcie.ctx); + if (pcie.host_buf) ggml_backend_buffer_free(pcie.host_buf); + if (pcie.dev_buf) ggml_backend_buffer_free(pcie.dev_buf); + ggml_backend_free(pcie.gpu_backend); + } + ggml_backend_free(cpu_be); + ggml_quantize_free(); + return 0; +} diff --git a/examples/llama-profiler/profiler-gpu.cpp b/examples/llama-profiler/profiler-gpu.cpp new file mode 100644 index 000000000000..127f4700d728 --- /dev/null +++ b/examples/llama-profiler/profiler-gpu.cpp @@ -0,0 +1,424 @@ +#define GPU_WARMUP_ITERS 3 +#define GPU_TIMED_ITERS 5 + +#include "profiler-common.h" + +#ifdef GGML_USE_CUDA +#include "ggml-cuda.h" +#endif + +#include +#include +#include + +static double benchmark_mul_mat_raw( + ggml_backend_t be, int N, int K, int batch_size, ggml_type quant, + double * out_time_s, double * out_ops, double * out_bytes) { + + ggml_init_params params = { 4096ULL * 1024 * 1024, NULL, true }; + ggml_context * ctx = ggml_init(params); + + ggml_tensor * A = ggml_new_tensor_2d(ctx, quant, K, N); + ggml_tensor * B_tensor = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, batch_size); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_tensor * C = ggml_mul_mat(ctx, A, B_tensor); + ggml_build_forward_expand(gf, C); + + if ((int64_t)N * batch_size > INT_MAX) { + printf("SKIPPED: MUL_MAT N=%d K=%d B=%d (output exceeds INT_MAX)\n", N, K, batch_size); + ggml_free(ctx); return 0.0; + } + + size_t gpu_free = 0, gpu_total = 0; + ggml_backend_dev_memory(ggml_backend_get_device(be), &gpu_free, &gpu_total); + size_t need = ggml_nbytes(A) + ggml_nbytes(B_tensor) + ggml_nbytes(C); + if (need > gpu_free * 0.9) { + printf("SKIPPED: MUL_MAT N=%d K=%d B=%d (need %.2f GB, free %.2f GB)\n", + N, K, batch_size, need / 1e9, gpu_free / 1e9); + ggml_free(ctx); return 0.0; + } + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, be); + if (!buffer) { fprintf(stderr, "Failed to allocate GPU buffer\n"); ggml_free(ctx); return 0.0; } + + std::vector A_float((size_t)K * N, 0.5f); + if (quant == GGML_TYPE_F32) { + ggml_backend_tensor_set(A, A_float.data(), 0, ggml_nbytes(A)); + } else { + std::vector A_q(ggml_nbytes(A)); + ggml_quantize_chunk(quant, A_float.data(), A_q.data(), 0, + (int64_t)K * N / ggml_blck_size(quant), 1, nullptr); + ggml_backend_tensor_set(A, A_q.data(), 0, ggml_nbytes(A)); + } + std::vector B_data((size_t)K * batch_size, 1.0f); + ggml_backend_tensor_set(B_tensor, B_data.data(), 0, ggml_nbytes(B_tensor)); + + for (int i = 0; i < GPU_WARMUP_ITERS; ++i) ggml_backend_graph_compute(be, gf); + ggml_backend_synchronize(be); + + bench_timer t; t.start(); + for (int i = 0; i < GPU_TIMED_ITERS; ++i) ggml_backend_graph_compute_async(be, gf); + ggml_backend_synchronize(be); + double total_time = t.stop(); + + double time_per_iter = total_time / GPU_TIMED_ITERS; + double ops_total = 2.0 * N * K * batch_size; + double bytes_total = (double)(ggml_nbytes(A) + ggml_nbytes(B_tensor) + ggml_nbytes(C)); + + if (out_time_s) *out_time_s = time_per_iter; + if (out_ops) *out_ops = ops_total; + if (out_bytes) *out_bytes = bytes_total; + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return ops_total / time_per_iter / 1e9; +} + +static double benchmark_mul_mat_id_raw( + ggml_backend_t be, int N, int K, int n_experts, int n_experts_used, + int batch_size, ggml_type quant, + double * out_time_s, double * out_ops, double * out_bytes) { + + ggml_init_params params = { 8192ULL * 1024 * 1024, NULL, true }; + ggml_context * ctx = ggml_init(params); + if (!ctx) { fprintf(stderr, "Failed to init context\n"); return 0.0; } + + ggml_tensor * A = ggml_new_tensor_3d(ctx, quant, K, N, n_experts); + ggml_tensor * B = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, K, 1, batch_size); + ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_experts_used, batch_size); + + if ((int64_t)N * batch_size * n_experts_used > INT_MAX) { + printf("SKIPPED: MUL_MAT_ID N=%d K=%d B=%d (output exceeds INT_MAX)\n", N, K, batch_size); + ggml_free(ctx); return 0.0; + } + + size_t gpu_free = 0, gpu_total = 0; + ggml_backend_dev_memory(ggml_backend_get_device(be), &gpu_free, &gpu_total); + size_t need = ggml_nbytes(A) + ggml_nbytes(B) + ggml_nbytes(ids); + if (need > gpu_free * 0.9) { + printf("SKIPPED: MUL_MAT_ID N=%d K=%d B=%d (need %.2f GB, free %.2f GB)\n", + N, K, batch_size, need / 1e9, gpu_free / 1e9); + ggml_free(ctx); return 0.0; + } + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_tensor * C = ggml_mul_mat_id(ctx, A, B, ids); + ggml_build_forward_expand(gf, C); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, be); + if (!buffer) { fprintf(stderr, "Failed to allocate GPU buffer\n"); ggml_free(ctx); return 0.0; } + + std::vector A_data(ggml_nbytes(A), 0); + ggml_backend_tensor_set(A, A_data.data(), 0, ggml_nbytes(A)); + std::vector B_data((size_t)K * batch_size, 1.0f); + ggml_backend_tensor_set(B, B_data.data(), 0, ggml_nbytes(B)); + std::vector ids_data(n_experts_used * batch_size); + for (int i = 0; i < n_experts_used * batch_size; i++) ids_data[i] = i % n_experts; + ggml_backend_tensor_set(ids, ids_data.data(), 0, ggml_nbytes(ids)); + + for (int i = 0; i < GPU_WARMUP_ITERS; ++i) ggml_backend_graph_compute(be, gf); + ggml_backend_synchronize(be); + + bench_timer t; t.start(); + for (int i = 0; i < GPU_TIMED_ITERS; ++i) ggml_backend_graph_compute_async(be, gf); + ggml_backend_synchronize(be); + double total_time = t.stop(); + + double time_per_iter = total_time / GPU_TIMED_ITERS; + double ops_total = 2.0 * N * K * batch_size * n_experts_used; + double bytes_total = (double)(ggml_nbytes(A) * n_experts_used / n_experts) + ggml_nbytes(B) + ggml_nbytes(C); + + if (out_time_s) *out_time_s = time_per_iter; + if (out_ops) *out_ops = ops_total; + if (out_bytes) *out_bytes = bytes_total; + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return ops_total / time_per_iter / 1e9; +} + +static double benchmark_flash_attn_raw( + ggml_backend_t be, int n_tokens, int ctx_len, + int n_q_heads, int n_kv_heads, int head_dim, ggml_type kv_quant, + double * out_time_s, double * out_ops, double * out_bytes) { + + ggml_init_params params = { 8192ULL * 1024 * 1024, NULL, true }; + ggml_context * ctx = ggml_init(params); + + ggml_tensor * Q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, head_dim, n_tokens, n_q_heads, 1); + ggml_tensor * K = ggml_new_tensor_4d(ctx, kv_quant, head_dim, ctx_len, n_kv_heads, 1); + ggml_tensor * V = ggml_new_tensor_4d(ctx, kv_quant, head_dim, ctx_len, n_kv_heads, 1); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_tensor * out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, 1.0f / sqrtf((float)head_dim), 0.0f, 0.0f); + ggml_build_forward_expand(gf, out); + + if ((int64_t)head_dim * n_tokens * n_q_heads > INT_MAX) { + printf("SKIPPED: FLASH_ATTN (output exceeds INT_MAX)\n"); + ggml_free(ctx); return 0.0; + } + + size_t gpu_free = 0, gpu_total = 0; + ggml_backend_dev_memory(ggml_backend_get_device(be), &gpu_free, &gpu_total); + size_t need = ggml_nbytes(Q) + ggml_nbytes(K) + ggml_nbytes(V) + ggml_nbytes(out); + if (need > gpu_free * 0.9) { + printf("SKIPPED: FLASH_ATTN (need %.2f GB, free %.2f GB)\n", need / 1e9, gpu_free / 1e9); + ggml_free(ctx); return 0.0; + } + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, be); + if (!buffer) { fprintf(stderr, "Failed to allocate GPU buffer\n"); ggml_free(ctx); return 0.0; } + + std::vector Q_data(ggml_nelements(Q), 1.0f); + ggml_backend_tensor_set(Q, Q_data.data(), 0, ggml_nbytes(Q)); + int64_t kv_elems = ggml_nelements(K); + std::vector KV_float(kv_elems, 0.5f); + if (kv_quant == GGML_TYPE_F32) { + ggml_backend_tensor_set(K, KV_float.data(), 0, ggml_nbytes(K)); + ggml_backend_tensor_set(V, KV_float.data(), 0, ggml_nbytes(V)); + } else { + std::vector Kq(ggml_nbytes(K)), Vq(ggml_nbytes(V)); + ggml_quantize_chunk(kv_quant, KV_float.data(), Kq.data(), 0, kv_elems / ggml_blck_size(kv_quant), 1, nullptr); + ggml_quantize_chunk(kv_quant, KV_float.data(), Vq.data(), 0, kv_elems / ggml_blck_size(kv_quant), 1, nullptr); + ggml_backend_tensor_set(K, Kq.data(), 0, ggml_nbytes(K)); + ggml_backend_tensor_set(V, Vq.data(), 0, ggml_nbytes(V)); + } + + for (int i = 0; i < GPU_WARMUP_ITERS; ++i) ggml_backend_graph_compute(be, gf); + ggml_backend_synchronize(be); + + bench_timer t; t.start(); + for (int i = 0; i < GPU_TIMED_ITERS; ++i) ggml_backend_graph_compute_async(be, gf); + ggml_backend_synchronize(be); + double total_time = t.stop(); + + double time_per_iter = total_time / GPU_TIMED_ITERS; + double ops_total = 2.0 * n_tokens * head_dim * ctx_len * n_q_heads * 2; + double bytes_total = (double)(ggml_nbytes(Q) + ggml_nbytes(K) + ggml_nbytes(V) + ggml_nbytes(out)); + + if (out_time_s) *out_time_s = time_per_iter; + if (out_ops) *out_ops = ops_total; + if (out_bytes) *out_bytes = bytes_total; + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return ops_total / time_per_iter / 1e9; +} + +static void run_matmul_benchmarks( + ggml_backend_t be, const std::vector & batch_sizes, + int32_t filter_batch, bool fast, + std::vector & results) { + + auto sizes = get_matmul_sizes(fast); + auto quants = get_matmul_quants(fast); + + printf("=== MUL_MAT Operations ===\n\n"); + for (ggml_type qt : quants) { + printf("--- Quantization: %s ---\n", ggml_type_name(qt)); + for (int32_t bs : batch_sizes) { + if (filter_batch >= 0 && bs != filter_batch) continue; + printf(" [Batch=%d]\n", bs); + for (const auto & sz : sizes) { + if (qt == GGML_TYPE_Q2_K && (sz.K % 256 != 0)) continue; + + bench_result res; + res.op_name = "MUL_MAT"; + res.quant_type = ggml_type_name(qt); + res.N = sz.N; res.K = sz.K; res.B = bs; + benchmark_mul_mat_raw(be, sz.N, sz.K, bs, qt, &res.time_s, &res.ops, &res.bytes); + res.calculate_derived(); + printf("%-20s quant=%-6s AI=%.3f BW=%.2f GB/s Perf=%.2f GFLOP/s", + res.op_name.c_str(), res.quant_type.c_str(), + res.arithmetic_intensity, res.effective_bw_gb_s, res.effective_gflops); + res.print_dims(); + printf("\n"); + results.push_back(res); + } + } + printf("\n"); + } +} + +static void run_moe_benchmarks( + ggml_backend_t be, const std::vector & batch_sizes, + int32_t filter_batch, bool fast, + std::vector & results) { + + auto configs = get_moe_configs(fast); + auto quants = get_matmul_quants(fast); + + printf("=== MUL_MAT_ID Operations (MoE) ===\n\n"); + for (ggml_type qt : quants) { + printf("--- MoE Quantization: %s ---\n", ggml_type_name(qt)); + for (int32_t bs : batch_sizes) { + if (filter_batch >= 0 && bs != filter_batch) continue; + printf(" [Batch=%d]\n", bs); + for (const auto & cfg : configs) { + if (qt == GGML_TYPE_Q2_K && (cfg.K % 256 != 0)) continue; + try { + bench_result res; + res.op_name = "MUL_MAT_ID"; + res.quant_type = ggml_type_name(qt); + res.N = cfg.N; res.K = cfg.K; res.B = bs; + res.n_tokens = cfg.n_experts_used; res.ctx_len = cfg.n_experts; + double gflops = benchmark_mul_mat_id_raw(be, cfg.N, cfg.K, cfg.n_experts, cfg.n_experts_used, bs, qt, + &res.time_s, &res.ops, &res.bytes); + if (gflops == 0.0) continue; + res.calculate_derived(); + printf("%-20s quant=%-6s AI=%.3f BW=%.2f GB/s Perf=%.2f GFLOP/s", + res.op_name.c_str(), res.quant_type.c_str(), + res.arithmetic_intensity, res.effective_bw_gb_s, res.effective_gflops); + res.print_dims(); + printf("\n"); + results.push_back(res); + } catch (const std::exception & e) { + printf("SKIPPED: MUL_MAT_ID N=%d K=%d %s (%s)\n", cfg.N, cfg.K, ggml_type_name(qt), e.what()); + } + } + } + printf("\n"); + } +} + +static void run_attention_benchmarks( + ggml_backend_t be, const std::vector & batch_sizes, + int32_t filter_batch, bool fast, + std::vector & results) { + + auto configs = get_attn_configs(fast); + auto ctx_lens = get_attn_ctx_lens(fast); + + printf("=== FLASH_ATTN Operations ===\n\n"); + for (const auto & cfg : configs) { + printf("--- %s (n_q=%d, n_kv=%d, head_dim=%d) ---\n", cfg.name, cfg.n_q_heads, cfg.n_kv_heads, cfg.head_dim); + for (int32_t n_tok : batch_sizes) { + if (filter_batch >= 0 && n_tok != filter_batch) continue; + printf(" [n_tokens=%d]\n", n_tok); + for (int32_t cl : ctx_lens) { + bench_result res; + res.op_name = std::string("FLASH_ATTN_") + cfg.name; + res.quant_type = ggml_type_name(GGML_TYPE_F16); + res.n_tokens = n_tok; res.ctx_len = cl; res.n_heads = cfg.n_kv_heads; res.head_dim = cfg.head_dim; + benchmark_flash_attn_raw(be, n_tok, cl, cfg.n_q_heads, cfg.n_kv_heads, cfg.head_dim, GGML_TYPE_F16, + &res.time_s, &res.ops, &res.bytes); + res.calculate_derived(); + printf("%-20s quant=%-6s AI=%.3f BW=%.2f GB/s Perf=%.2f GFLOP/s", + res.op_name.c_str(), res.quant_type.c_str(), + res.arithmetic_intensity, res.effective_bw_gb_s, res.effective_gflops); + res.print_dims(); + printf("\n"); + results.push_back(res); + } + } + printf("\n"); + } +} + +static void save_results_gpu( + const char * path, + const std::vector & results, + const char * backend_name, + const std::vector & batch_sizes, + double peak_gpu_bw, double peak_gpu_compute) { + + FILE * f = fopen(path, "w"); + if (!f) { fprintf(stderr, "Failed to open %s for writing\n", path); return; } + + fprintf(f, "# GPU Profiling (backend=%s, batch_sizes=[", backend_name); + for (size_t i = 0; i < batch_sizes.size(); i++) + fprintf(f, "%d%s", batch_sizes[i], i + 1 < batch_sizes.size() ? "," : ""); + fprintf(f, "], GPU_Memory_BW=%.1f GB/s, GPU_Peak_Compute=%.1f GFLOP/s)\n", peak_gpu_bw, peak_gpu_compute); + fprintf(f, "# op_name quant AI(FLOP/byte) BW(GB/s) GFLOP/s Ridge(FLOP/byte) N K B n_tokens ctx_len n_heads head_dim n_elements\n"); + + auto ridges = compute_ridge_points(results, peak_gpu_bw); + std::map ridge_map; + for (const auto & rr : ridges) ridge_map[rr.key] = rr.ridge; + + for (const auto & r : results) { + std::string key = r.op_name + "_" + r.quant_type; + double ridge = ridge_map.count(key) ? ridge_map[key] : 0.0; + fprintf(f, "%s %s %.4f %.2f %.2f %.4f %d %d %d %d %d %d %d %lld\n", + r.op_name.c_str(), r.quant_type.c_str(), + r.arithmetic_intensity, r.effective_bw_gb_s, r.effective_gflops, ridge, + r.N, r.K, r.B, r.n_tokens, r.ctx_len, r.n_heads, r.head_dim, (long long)r.n_elements); + } + + fclose(f); + printf("Results saved to %s (%zu benchmarks)\n", path, results.size()); +} + +int main(int argc, char ** argv) { + bool fast_mode = true; + int32_t filter_batch = -1; + const char * output_path = "gpu_profile.txt"; + + for (int i = 1; i < argc; ++i) { + if (!strcmp(argv[i], "--fast")) { + fast_mode = true; + } else if (!strcmp(argv[i], "--full")) { + fast_mode = false; + } else if (!strcmp(argv[i], "-h") || !strcmp(argv[i], "--help")) { + printf("usage: %s [options]\n", argv[0]); + printf("\n"); + printf("options:\n"); + printf(" -h, --help\n"); + printf(" --fast fast mode with fewer configs (default)\n"); + printf(" --full full mode with all configs\n"); + printf(" --batch only run batch size N\n"); + printf(" --output output file (default: gpu_profile.txt)\n"); + return 0; + } else if (!strcmp(argv[i], "--batch") && i + 1 < argc) { + if (!parse_int_arg(argv[++i], filter_batch)) { + fprintf(stderr, "Invalid --batch value: %s\n", argv[i]); + return 1; + } + } else if (!strcmp(argv[i], "--output") && i + 1 < argc) { + output_path = argv[++i]; + } + } + + std::vector batch_sizes = { 1, 64, 512, 1024, 2048, 4096, 8192, 16384 }; + + printf("=== GPU Profiler ===\n"); + printf("Mode: %s\n\n", fast_mode ? "FAST" : "FULL"); + + ggml_backend_t gpu_be = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_GPU, nullptr); + if (!gpu_be) { + fprintf(stderr, "No GPU backend available. Cannot run GPU profiler.\n"); + return 1; + } + printf("GPU: %s\n\n", ggml_backend_name(gpu_be)); + + ggml_quantize_init(GGML_TYPE_Q2_K); + ggml_quantize_init(GGML_TYPE_Q4_0); + ggml_quantize_init(GGML_TYPE_Q4_1); + ggml_quantize_init(GGML_TYPE_Q5_0); + ggml_quantize_init(GGML_TYPE_Q8_0); + ggml_quantize_init(GGML_TYPE_MXFP4); + + std::vector all_results; + bench_timer overall; overall.start(); + + run_matmul_benchmarks(gpu_be, batch_sizes, filter_batch, fast_mode, all_results); + run_moe_benchmarks(gpu_be, batch_sizes, filter_batch, fast_mode, all_results); + run_attention_benchmarks(gpu_be, batch_sizes, filter_batch, fast_mode, all_results); + + double peak_gpu_bw = 0.0, peak_gpu_compute = 0.0; + for (const auto & r : all_results) { + if (r.arithmetic_intensity < 2.0) + peak_gpu_bw = std::max(peak_gpu_bw, (double)r.effective_bw_gb_s); + if (r.arithmetic_intensity > 10.0 || (r.op_name == "MUL_MAT" && r.N >= 4096 && r.K >= 4096)) + peak_gpu_compute = std::max(peak_gpu_compute, (double)r.effective_gflops); + } + printf("\nEstimated GPU Memory BW: %.1f GB/s\n", peak_gpu_bw); + printf("Estimated GPU Compute: %.1f GFLOP/s\n", peak_gpu_compute); + printf("Total time: %.1f s, %zu benchmarks\n", overall.stop(), all_results.size()); + + save_results_gpu(output_path, all_results, ggml_backend_name(gpu_be), batch_sizes, peak_gpu_bw, peak_gpu_compute); + + ggml_backend_free(gpu_be); + ggml_quantize_free(); + return 0; +} diff --git a/ggml/include/ggml-alloc.h b/ggml/include/ggml-alloc.h index a7926a21a9a2..5a28fe1febd3 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 2924fdbe9884..4d674e2c3d33 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 @@ -332,6 +335,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 @@ -351,6 +355,55 @@ 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_weight_copy_bytes; + size_t input_activ_bytes; + size_t writeback_bytes; + bool can_prefetch_weights; + }; + + GGML_API bool ggml_backend_sched_get_split_info( + ggml_backend_sched_t sched, int split_id, + struct ggml_backend_sched_split_info * out); + // // Meta backend // diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 35f0c44ec421..171bf84f0f56 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -652,7 +652,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 3bda9abbe03d..7b962914052b 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 87615921c09b..03d459989ae2 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,22 @@ struct ggml_backend_sched { ggml_backend_sched_eval_callback callback_eval; void * callback_eval_user_data; + bool prefetch_weights; + + // redirect_target[backend_id] is -1 for regular compute backends + // alias backends on the same physical device point to the regular compute backend id + int redirect_target[GGML_SCHED_MAX_BACKENDS]; + bool has_redirects; + + ggml_backend_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 +882,23 @@ static int ggml_backend_sched_backend_from_buffer(ggml_backend_sched_t sched, co return -1; } + // respect explicit set_tensor_backend assignments (walk view_src chain to root) + { + const struct ggml_tensor * cur = tensor; + while (cur) { + size_t idx = ggml_hash_find(&sched->hash_set, cur); + if (idx != SIZE_MAX && idx < sched->hash_set.size) { + int explicit_id = sched->hv_tensor_backend_ids[idx]; + if (explicit_id >= 0 && explicit_id < sched->n_backends) { + if (ggml_backend_supports_op(sched->backends[explicit_id], op)) { + return explicit_id; + } + } + } + cur = cur->view_src; + } + } + // find highest prio backend that supports the buffer type and the op for (int i = 0; i < sched->n_backends; i++) { if (ggml_backend_supports_buft(sched->backends[i], buffer->buft) && @@ -913,7 +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)) { @@ -1185,8 +1236,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; @@ -1257,6 +1310,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]; @@ -1313,6 +1367,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; } @@ -1349,7 +1404,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]; @@ -1367,8 +1432,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; @@ -1390,7 +1480,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; @@ -1435,11 +1529,83 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra graph_copy->nodes[graph_copy->n_nodes++] = input_cpy; } + // reserve 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 immediate next sharded split's prefetch destinations + // (race-protect copy stream vs current compute, do not scan past CPU splits) + if (sched->prefetch_weights) { + struct ggml_backend_sched_split * next_gpu = NULL; + if (i + 1 < sched->n_splits) { + struct ggml_backend_sched_split * candidate = &sched->splits[i + 1]; + // only prefetch into redirected backends + // the main compute backend should use the normal input copy i.e. no prefetch + if (sched->redirect_target[candidate->backend_id] >= 0 && + sched->copy_backends[candidate->backend_id] != NULL) { + bool has_host_weights = false; + for (int j = 0; j < candidate->n_inputs; j++) { + if (candidate->inputs[j]->buffer != NULL && + ggml_backend_buffer_get_usage(candidate->inputs[j]->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && + ggml_backend_buffer_is_host(candidate->inputs[j]->buffer)) { + has_host_weights = true; + break; + } + } + if (has_host_weights) { + next_gpu = candidate; + } + } + } + if (next_gpu != NULL) { + for (int j = 0; j < next_gpu->n_inputs; j++) { + struct ggml_tensor * next_input = next_gpu->inputs[j]; + if (next_input->buffer != NULL && + ggml_backend_buffer_get_usage(next_input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && + ggml_backend_buffer_is_host(next_input->buffer)) { + const size_t id = hash_id(next_input); + struct ggml_tensor * next_cpy = tensor_id_copy(id, next_gpu->backend_id, sched->cur_copy); + assert(graph_copy->size > graph_copy->n_nodes); + struct ggml_tensor * keepalive = ggml_view_tensor(sched->ctx, next_cpy); + keepalive->src[0] = next_cpy; + sched->node_backend_ids[graph_copy->n_nodes] = next_gpu->backend_id; + graph_copy->nodes[graph_copy->n_nodes++] = keepalive; + } + } + + // reserve next split's writeback slot (double-buffer: prefetch copy stream vs. current compute) + for (int w = 0; w < next_gpu->n_writeback; w++) { + assert(graph_copy->size > graph_copy->n_nodes); + struct ggml_tensor * keepalive = ggml_view_tensor(sched->ctx, next_gpu->writeback[w]); + keepalive->op = GGML_OP_NONE; + keepalive->src[0] = next_gpu->writeback[w]; + sched->node_backend_ids[graph_copy->n_nodes] = next_gpu->backend_id; + graph_copy->nodes[graph_copy->n_nodes++] = keepalive; + } + } + } + for (int j = split->i_start; j < split->i_end; j++) { assert(graph_copy->size > graph_copy->n_nodes); sched->node_backend_ids[graph_copy->n_nodes] = tensor_backend_id(graph->nodes[j]); graph_copy->nodes[graph_copy->n_nodes++] = graph->nodes[j]; } + + // n_children fence for writeback (held through post_compute, required by FLAG_WRITEBACK semantics) + for (int w = 0; w < split->n_writeback; w++) { + assert(graph_copy->size > graph_copy->n_nodes); + struct ggml_tensor * keepalive = ggml_view_tensor(sched->ctx, split->writeback[w]); + keepalive->op = GGML_OP_NONE; + keepalive->src[0] = split->writeback[w]; + sched->node_backend_ids[graph_copy->n_nodes] = split->backend_id; + graph_copy->nodes[graph_copy->n_nodes++] = keepalive; + } } if (sched->n_copies > 1) { @@ -1526,6 +1692,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); @@ -1545,17 +1714,121 @@ 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; + 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]; 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]; + } + + const bool weights_prefetched = + prefetched_split_id == split_id && + prefetched_backend_id == copy_backend_id; + + if (prefetched_split_id == split_id && !weights_prefetched) { + prefetched_split_id = -1; + prefetched_backend_id = -1; + } + + 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; + next_gpu_id = nid; + } + } + } + if (next_gpu == NULL) { + return; + } + + 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]); + } + + 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 (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++) { 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,19 +1934,35 @@ 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]); } 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); + } + } + + 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) { @@ -1713,12 +2002,26 @@ 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) { ggml_backend_event_record(sched->events[split_backend_id][sched->cur_copy], split_backend); } } + } return GGML_STATUS_SUCCESS; @@ -1754,10 +2057,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])); @@ -1766,7 +2073,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; @@ -1785,6 +2092,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; @@ -1801,12 +2136,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); @@ -1824,6 +2166,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; } @@ -1920,6 +2264,195 @@ 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; + + auto effective_weight_copy_bytes = [&](struct ggml_tensor * inp) { + const size_t full_size = ggml_nbytes(inp); + + if (s->graph.n_nodes <= 0) { + return full_size; + } + + struct ggml_tensor * node = s->graph.nodes[0]; + if (node->op != GGML_OP_MUL_MAT_ID) { + return full_size; + } + + struct ggml_tensor * inp_cpy = tensor_copy(inp, s->backend_id, sched->cur_copy); + if (node->src[0] != inp_cpy) { + return full_size; + } + + const int64_t n_expert = inp->ne[2]; + if (n_expert <= 0 || node->src[2] == NULL) { + return full_size; + } + + const size_t expert_size = inp->nb[2]; + const int64_t n_routes = ggml_nelements(node->src[2]); + int64_t n_used = std::min(n_expert, std::max(1, n_routes)); + const int64_t n_expert_used = node->ne[1]; + // ids contains one route per token/top-k slot, not unique experts. + // For larger batches, estimate unique experts instead of charging every route. + if (n_used > n_expert_used) { + n_used = std::min(n_expert, std::max(n_expert_used, (n_expert + 3) / 4)); + } + const size_t padding = std::min(expert_size, 512); + + return std::min(full_size, (size_t) n_used * (expert_size + padding)); + }; + + out->input_weight_bytes = 0; + out->input_weight_copy_bytes = 0; + out->input_activ_bytes = 0; + for (int j = 0; j < s->n_inputs; j++) { + struct ggml_tensor * inp = s->inputs[j]; + if (inp->buffer != NULL && + ggml_backend_buffer_get_usage(inp->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && + ggml_backend_buffer_is_host(inp->buffer)) { + out->input_weight_bytes += ggml_nbytes(inp); + out->input_weight_copy_bytes += effective_weight_copy_bytes(inp); + } else { + out->input_activ_bytes += ggml_nbytes(inp); + } + } + + out->writeback_bytes = 0; + for (int w = 0; w < s->n_writeback; w++) { + out->writeback_bytes += ggml_nbytes(s->writeback[w]); + } + + out->can_prefetch_weights = + sched->prefetch_weights && + sched->redirect_target[s->backend_id] >= 0 && + sched->copy_backends[s->backend_id] != NULL; + + + return true; +} + int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched) { GGML_ASSERT(sched); return sched->n_splits; @@ -1962,8 +2495,21 @@ void ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct gg int backend_index = ggml_backend_sched_backend_id(sched, backend); GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); tensor_backend_id(node) = backend_index; + sched->hv_tensor_usr[hash_id(node)] = true; SET_CAUSE(node, "usr"); - sched->is_reset = false; + + // not clearing is_reset here: split_graph() already sets is_reset=false, + // guaranteeing cleanup between graphs. clearing it here would cause + // reserve_size's sched_reset to discard these assignments before + // split_graph reads them. +} + +void ggml_backend_sched_set_tensor_backend_hint(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend) { + GGML_ASSERT(sched); + int backend_index = ggml_backend_sched_backend_id(sched, backend); + GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); + tensor_backend_id(node) = backend_index; + SET_CAUSE(node, "hint"); } ggml_backend_t ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node) { diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 74631c2857ba..2bc4cab306ea 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 e73a7b8906ce..a5ff040d971b 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4425,6 +4425,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, @@ -4442,6 +4448,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() { @@ -4696,6 +4703,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 9fab69317006..3ce93f174eab 100644 --- a/include/llama.h +++ b/include/llama.h @@ -300,6 +300,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 { @@ -336,6 +337,11 @@ extern "C" { bool use_extra_bufts; // use extra buffer types (used for weight repacking) bool no_host; // bypass host buffer allowing extra buffers to be used bool no_alloc; // only load metadata and simulate memory allocations + bool pshard; // enable pipelined sharding: weights on CPU host, pipelined to GPU per split + bool pshard_delegate_compute; // tensor overrides choose storage; scheduler chooses compute backends + bool pshard_cache_skip_load; // skip loading plan cache (rebuild from scratch, then overwrite) + size_t max_vram_alloc; // VRAM budget in MiB for the unified preload buffer (0 = auto from free VRAM) + struct llama_pshard_plan_registry * pshard_registry; // tier plan registry, caller-owned. populated by llama_params_fit_pshard }; struct llama_sampler_seq_config { @@ -404,6 +410,8 @@ extern "C" { // a source/target/parent context // can be utilized in various ways, for example by sharing results or llama_memory between 2 contexts struct llama_context * ctx_other; + + bool pshard; // enable pipelined sharding for this context (3 GPU backends, split callbacks) }; struct llama_model_tensor_override { diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 320784c3a8cc..72167d134a5d 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -14,8 +14,10 @@ add_library(llama llama-adapter.cpp llama-arch.cpp llama-batch.cpp + llama-benchmark.cpp llama-chat.cpp llama-context.cpp + llama-context-pshard.cpp llama-cparams.cpp llama-grammar.cpp llama-graph.cpp @@ -29,11 +31,14 @@ 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-pshard-plan.cpp llama-quant.cpp llama-sampler.cpp llama-vocab.cpp diff --git a/src/llama-benchmark.cpp b/src/llama-benchmark.cpp new file mode 100644 index 000000000000..3a4e425672ef --- /dev/null +++ b/src/llama-benchmark.cpp @@ -0,0 +1,693 @@ +#include "llama-benchmark.h" +#include "llama-impl.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include + +llama_op_metrics llama_op_metrics_compute(const ggml_tensor * node) { + llama_op_metrics m = {}; + + switch (node->op) { + case GGML_OP_MUL_MAT: { + m.N = node->ne[0]; + m.M = node->ne[1]; + m.K = node->src[0]->ne[0]; + m.ops = 2.0 * m.N * m.K * m.M; + m.bytes = ggml_nbytes(node->src[0]) + ggml_nbytes(node->src[1]) + ggml_nbytes(node); + m.quant_type = ggml_type_name(node->src[0]->type); + break; + } + case GGML_OP_MUL_MAT_ID: { + m.N = node->ne[0]; + m.n_experts_used = node->ne[1]; + m.M = node->ne[2]; + m.K = node->src[0]->ne[0]; + const int64_t total_experts = node->src[0]->ne[2]; + m.ops = 2.0 * m.N * m.K * m.M * m.n_experts_used; + m.bytes = (ggml_nbytes(node->src[0]) * m.n_experts_used / total_experts) + + ggml_nbytes(node->src[1]) + ggml_nbytes(node); + m.quant_type = ggml_type_name(node->src[0]->type); + break; + } + case GGML_OP_FLASH_ATTN_EXT: { + m.head_dim = node->ne[0]; + m.n_q_heads = node->ne[1]; + m.n_tokens = node->ne[2]; + m.ctx_len = node->src[1]->ne[1]; + m.n_kv_heads = node->src[1]->ne[2]; + m.ops = 2.0 * m.n_tokens * m.head_dim * m.ctx_len * m.n_q_heads * 2.0; + m.bytes = ggml_nbytes(node->src[0]) + ggml_nbytes(node->src[1]) + + ggml_nbytes(node->src[2]) + ggml_nbytes(node); + break; + } + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: { + m.n_elements = ggml_nelements(node); + m.ops = 6.0 * m.n_elements; + m.bytes = ggml_nbytes(node->src[0]) + ggml_nbytes(node); + m.quant_type = ggml_type_name(node->type); + break; + } + case GGML_OP_RMS_NORM: { + m.n_elements = ggml_nelements(node); + m.ops = 3.0 * m.n_elements; + m.bytes = ggml_nbytes(node->src[0]) + ggml_nbytes(node); + m.quant_type = ggml_type_name(node->type); + break; + } + case GGML_OP_GLU: { + m.n_elements = ggml_nelements(node); + m.ops = 2.0 * m.n_elements; + m.bytes = ggml_nbytes(node->src[0]); + if (node->src[1]) { m.bytes += ggml_nbytes(node->src[1]); } + m.bytes += ggml_nbytes(node); + m.quant_type = ggml_type_name(node->type); + break; + } + case GGML_OP_ADD: + case GGML_OP_MUL: + case GGML_OP_SUB: + case GGML_OP_DIV: { + m.n_elements = ggml_nelements(node); + m.ops = m.n_elements; + m.bytes = ggml_nbytes(node->src[0]); + if (node->src[1]) { m.bytes += ggml_nbytes(node->src[1]); } + m.bytes += ggml_nbytes(node); + m.quant_type = ggml_type_name(node->type); + break; + } + case GGML_OP_GET_ROWS: { + m.n_elements = ggml_nelements(node); + m.bytes = ggml_nbytes(node); + m.quant_type = ggml_type_name(node->src[0]->type); + break; + } + case GGML_OP_SET_ROWS: + case GGML_OP_CPY: { + m.n_elements = ggml_nelements(node); + m.bytes = ggml_nbytes(node->src[0]) + ggml_nbytes(node); + m.quant_type = ggml_type_name(node->src[0]->type); + break; + } + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + case GGML_OP_NONE: + break; + default: { + m.n_elements = ggml_nelements(node); + m.bytes = ggml_nbytes(node); + if (node->src[0]) { m.bytes += ggml_nbytes(node->src[0]); } + m.quant_type = ggml_type_name(node->type); + break; + } + } + return m; +} + +std::string llama_benchmark_predictor::make_key( + const std::string & op, const std::string & quant, + int64_t N, int64_t K, int64_t batch) { + return op + "|" + quant + "|" + std::to_string(N) + "|" + + std::to_string(K) + "|" + std::to_string(batch); +} + +std::string llama_benchmark_predictor::make_attn_key( + int64_t ctx_len, int64_t n_kv_heads, int64_t n_tokens) { + return "FLASH_ATTN|" + std::to_string(ctx_len) + "|" + + std::to_string(n_kv_heads) + "|" + std::to_string(n_tokens); +} + +std::string llama_benchmark_predictor::make_elem_key( + const std::string & op, int64_t n_elements) { + return op + "|" + std::to_string(n_elements); +} + +static void build_entry_map( + std::unordered_map & map, + const std::vector & entries) { + map.clear(); + map.reserve(entries.size()); + for (const auto & e : entries) { + std::string key; + if (e.op_name == "MUL_MAT" || e.op_name == "MUL_MAT_ID") { + key = llama_benchmark_predictor::make_key(e.op_name, e.quant, e.N, e.K, e.B); + } else if (e.op_name.compare(0, 10, "FLASH_ATTN") == 0) { + // n_heads stores n_kv_heads for attention benchmarks (written by profiler) + key = llama_benchmark_predictor::make_attn_key(e.ctx_len, e.n_heads, e.n_tokens); + } else { + key = llama_benchmark_predictor::make_elem_key(e.op_name, e.n_elements); + } + map[key] = &e; + } +} + +void llama_benchmark_predictor::build_maps() { + build_entry_map(cpu_map, cpu_entries); + build_entry_map(gpu_map, gpu_entries); +} + +bool llama_benchmark_predictor::load_cpu(const char * filepath, int n_threads) { + cpu_entries.clear(); + cpu_map.clear(); + + FILE * f = fopen(filepath, "r"); + if (!f) { + LLAMA_LOG_WARN("%s: could not open CPU benchmark file: %s\n", __func__, filepath); + return false; + } + + double max_dram_bw = 0.0; + + char line[512]; + while (fgets(line, sizeof(line), f)) { + if (line[0] == '#') { + int tc = 0; + double dram_bw = 0.0, pcie_standalone = 0.0, pcie_concurrent = 0.0, cpu_eff = 0.0; + + if (sscanf(line, + "# Threads=%d: DRAM_BW=%lf GB/s, PCIe_Standalone=%lf GB/s," + " PCIe_Concurrent=%lf GB/s (CPU_Eff=%lf%%)", + &tc, &dram_bw, &pcie_standalone, &pcie_concurrent, &cpu_eff) == 5) { + if (dram_bw > max_dram_bw) { + max_dram_bw = dram_bw; + } + if (tc == n_threads) { + stats.peak_system_bw = dram_bw; + stats.peak_pcie_bw = pcie_standalone; + stats.eff_system_bw = dram_bw * (cpu_eff / 100.0); + stats.eff_pcie_bw = std::min(pcie_concurrent, dram_bw); + } + } else if (sscanf(line, "# Threads=%d: DRAM_BW=%lf GB/s", &tc, &dram_bw) == 2) { + if (dram_bw > max_dram_bw) { + max_dram_bw = dram_bw; + } + if (tc == n_threads) { + stats.peak_system_bw = dram_bw; + } + } + continue; + } + + if (line[0] == '\n') continue; + + llama_benchmark_entry e; + char op_name[64], quant[16]; + + int parsed = sscanf(line, + "%63s %15s %d %lf %lf %lf %lf %lf %lf" + " %lld %lld %lld %lld %lld %lld %lld %lld", + op_name, quant, &e.threads, + &e.ai, &e.bw_gb_s, &e.peak_gflops, &e.ridge, + &e.eff_gflops, &e.eff_pcie_bw, + &e.N, &e.K, &e.B, &e.n_tokens, &e.ctx_len, + &e.n_heads, &e.head_dim, &e.n_elements); + + if (parsed >= 9 && e.threads == n_threads) { + e.op_name = op_name; + e.quant = quant; + cpu_entries.push_back(std::move(e)); + } + } + fclose(f); + + // PCIe can't exceed system DRAM BW + if (max_dram_bw > 0.0) { + stats.peak_pcie_bw = std::min(stats.peak_pcie_bw, max_dram_bw); + } else { + stats.peak_pcie_bw = std::min(stats.peak_pcie_bw, stats.peak_system_bw); + } + + if (stats.eff_system_bw == 0.0) { stats.eff_system_bw = stats.peak_system_bw; } + if (stats.eff_pcie_bw == 0.0) { stats.eff_pcie_bw = stats.peak_pcie_bw; } + + build_maps(); + + LLAMA_LOG_INFO("%s: loaded %zu CPU entries for %d threads" + " (peak: DRAM=%.1f GB/s, PCIe=%.1f GB/s" + " | concurrent: DRAM=%.1f GB/s, PCIe=%.1f GB/s)\n", + __func__, cpu_entries.size(), n_threads, + stats.peak_system_bw, stats.peak_pcie_bw, + stats.eff_system_bw, stats.eff_pcie_bw); + + return !cpu_entries.empty(); +} + +bool llama_benchmark_predictor::load_gpu(const char * filepath) { + gpu_entries.clear(); + gpu_map.clear(); + + FILE * f = fopen(filepath, "r"); + if (!f) { + LLAMA_LOG_WARN("%s: could not open GPU benchmark file: %s\n", __func__, filepath); + return false; + } + + char line[512]; + while (fgets(line, sizeof(line), f)) { + if (line[0] == '#') { + const char * bw_str = strstr(line, "GPU_Memory_BW="); + const char * cp_str = strstr(line, "GPU_Peak_Compute="); + if (bw_str && cp_str) { + double mem_bw = 0.0, compute = 0.0; + sscanf(bw_str, "GPU_Memory_BW=%lf", &mem_bw); + sscanf(cp_str, "GPU_Peak_Compute=%lf", &compute); + stats.peak_gpu_mem_bw = mem_bw; + stats.peak_gpu_compute = compute; + } + continue; + } + + if (line[0] == '\n') continue; + + llama_benchmark_entry e; + char op_name[64], quant[16]; + + int parsed = sscanf(line, + "%63s %15s %lf %lf %lf %lf" + " %lld %lld %lld %lld %lld %lld %lld %lld", + op_name, quant, + &e.ai, &e.bw_gb_s, &e.peak_gflops, &e.ridge, + &e.N, &e.K, &e.B, &e.n_tokens, &e.ctx_len, + &e.n_heads, &e.head_dim, &e.n_elements); + + if (parsed == 14) { + e.op_name = op_name; + e.quant = quant; + e.threads = -1; + e.eff_gflops = e.peak_gflops; + e.eff_pcie_bw = 0.0; + gpu_entries.push_back(std::move(e)); + } + } + fclose(f); + + build_maps(); + + LLAMA_LOG_INFO("%s: loaded %zu GPU entries (mem_bw=%.1f GB/s, compute=%.1f GFLOP/s)\n", + __func__, gpu_entries.size(), + stats.peak_gpu_mem_bw, stats.peak_gpu_compute); + + return !gpu_entries.empty(); +} + +std::string llama_benchmark_predictor::make_timing_key( + bool is_gpu, bool async_copy, const char * op_name, + const llama_op_metrics & m, int32_t batch_size) { + // mode prefix: GPU vs CPU with/without concurrent PCIe (async_copy uses eff_* bandwidths) + std::string key = is_gpu ? "GPU|" : (async_copy ? "CPU_ASYNC|" : "CPU_SYNC|"); + key += op_name; + key += "|"; + + if (m.quant_type) { + key += m.quant_type; + key += "|"; + } + + if (strstr(op_name, "MUL_MAT_ID")) { + key += std::to_string(m.N) + "|" + std::to_string(m.K) + "|" + + std::to_string(batch_size) + "|" + std::to_string(m.M) + "|" + + std::to_string(m.n_experts_used); + } else if (strstr(op_name, "MUL_MAT")) { + key += std::to_string(m.N) + "|" + std::to_string(m.K) + "|" + + std::to_string(batch_size) + "|" + std::to_string(m.M); + } else if (strstr(op_name, "FLASH_ATTN")) { + key += std::to_string(m.head_dim) + "|" + std::to_string(m.n_q_heads) + "|" + + std::to_string(batch_size) + "|" + std::to_string(m.n_tokens) + "|" + + std::to_string(m.ctx_len) + "|" + std::to_string(m.n_kv_heads); + } else { + key += std::to_string(m.n_elements); + } + + return key; +} + +llama_split_timing llama_benchmark_predictor::predict_split( + struct ggml_tensor ** nodes, int n_nodes, + bool is_gpu, int32_t batch_size, bool async_copy, + timing_cache_t * timing_cache) const { + + llama_split_timing result = {}; + timing_cache_t local_timing_cache; + timing_cache_t & cache = timing_cache ? *timing_cache : local_timing_cache; + + const auto & map = is_gpu ? gpu_map : cpu_map; + const auto & entries = is_gpu ? gpu_entries : cpu_entries; + + double pcie_sum = 0.0; + + for (int i = 0; i < n_nodes; i++) { + ggml_tensor * node = nodes[i]; + llama_op_metrics m = llama_op_metrics_compute(node); + + if (m.ops == 0.0 && m.bytes == 0.0) continue; + + const char * op_name = ggml_op_name(node->op); + double op_time_ms = 0.0; + double pcie_contrib = 0.0; + + std::string tkey = make_timing_key(is_gpu, async_copy, op_name, m, batch_size); + auto cache_it = cache.find(tkey); + if (cache_it != cache.end()) { + op_time_ms = cache_it->second.first; + pcie_contrib = cache_it->second.second; + result.time_ms += op_time_ms; + pcie_sum += pcie_contrib; + result.n_cache_hit++; + result.op_count++; + continue; + } + + // build hash key for exact lookup + std::string hkey; + switch (node->op) { + case GGML_OP_MUL_MAT: + hkey = make_key("MUL_MAT", m.quant_type ? m.quant_type : "", m.N, m.K, batch_size); + break; + case GGML_OP_MUL_MAT_ID: + hkey = make_key("MUL_MAT_ID", m.quant_type ? m.quant_type : "", m.N, m.K, batch_size); + break; + case GGML_OP_FLASH_ATTN_EXT: + hkey = make_attn_key(m.ctx_len, m.n_kv_heads, batch_size); + break; + case GGML_OP_ROPE: + case GGML_OP_RMS_NORM: + case GGML_OP_GLU: + case GGML_OP_ADD: + case GGML_OP_MUL: + case GGML_OP_SUB: + case GGML_OP_DIV: + hkey = make_elem_key(op_name, m.n_elements); + break; + default: + break; + } + + const llama_benchmark_entry * match = nullptr; + bool exact = false; + + // try exact hash match + if (!hkey.empty()) { + auto it = map.find(hkey); + if (it != map.end()) { + match = it->second; + exact = true; + result.n_exact++; + } + } + + // fall back to nearest-neighbor + if (!match) { + switch (node->op) { + case GGML_OP_MUL_MAT: + match = find_nearest(entries, "MUL_MAT", m.quant_type, m.N, m.K, 0, 0, -1, batch_size); + break; + case GGML_OP_MUL_MAT_ID: + match = find_nearest(entries, "MUL_MAT_ID", m.quant_type, m.N, m.K, 0, 0, -1, batch_size); + break; + case GGML_OP_FLASH_ATTN_EXT: + match = find_nearest(entries, "FLASH_ATTN", nullptr, 0, 0, m.ctx_len, 0, m.n_kv_heads, batch_size); + break; + default: + match = find_nearest(entries, op_name, nullptr, 0, 0, 0, m.n_elements, -1, batch_size); + break; + } + if (match) { + result.n_nearest++; + } + } + + // compute timing from match or fall back to peak BW + if (match) { + // GPU always peak; CPU uses eff (concurrent) when async_copy, peak otherwise + const double gflops = is_gpu ? match->peak_gflops + : (async_copy ? match->eff_gflops : match->peak_gflops); + + if (exact) { + if (gflops > 0.0) { + op_time_ms = (m.ops / 1e9) / gflops * 1000.0; + } + } else { + const double ai = m.bytes > 0.0 ? (m.ops / m.bytes) : 0.0; + if (ai < match->ridge && match->bw_gb_s > 0.0) { + op_time_ms = (m.bytes / 1e9) / match->bw_gb_s * 1000.0; + } else if (gflops > 0.0) { + op_time_ms = (m.ops / 1e9) / gflops * 1000.0; + } + } + + pcie_contrib = op_time_ms * match->eff_pcie_bw; + } else { + result.n_fallback++; + if (m.bytes > 0.0) { + const double bw = is_gpu ? stats.peak_gpu_mem_bw + : (async_copy ? stats.eff_system_bw : stats.peak_system_bw); + if (bw > 0.0) { + op_time_ms = (m.bytes / 1e9) / bw * 1000.0; + } + if (!is_gpu && stats.eff_pcie_bw > 0.0) { + pcie_contrib = op_time_ms * stats.eff_pcie_bw; + } + } + } + + result.time_ms += op_time_ms; + pcie_sum += pcie_contrib; + result.op_count++; + + cache[tkey] = {op_time_ms, pcie_contrib}; + } + + // convert accumulated pcie_sum to time-weighted average BW + if (!is_gpu) { + result.eff_pcie_bw = (result.time_ms > 0.0) ? (pcie_sum / result.time_ms) + : stats.eff_pcie_bw; + } + + return result; +} + +static double get_bits_per_weight(const std::string & quant) { + static const std::unordered_map table = { + {"q2_K", 2.6}, {"q3_K", 3.4}, {"q4_0", 4.5}, {"q4_1", 5.0}, {"q4_K", 4.5}, + {"q5_0", 5.5}, {"q5_1", 6.0}, {"q5_K", 5.5}, {"q6_K", 6.6}, + {"q8_0", 8.5}, {"q8_1", 9.0}, {"f16", 16.0}, {"f32", 32.0}, + }; + auto it = table.find(quant); + return it != table.end() ? it->second : -1.0; +} + +const llama_benchmark_entry * llama_benchmark_predictor::find_nearest( + const std::vector & entries, + const char * op_name, const char * quant, + int64_t N, int64_t K, int64_t ctx_len, int64_t n_elements, + int64_t n_kv_heads, int64_t target_batch) { + + if (!op_name || entries.empty()) { + return nullptr; + } + + const llama_benchmark_entry * best = nullptr; + double best_score = 1e20; + + const double target_bpw = quant ? get_bits_per_weight(quant) : -1.0; + + // FLASH_ATTN entries have suffixed names (e.g. FLASH_ATTN_MHA, FLASH_ATTN_GQA-8) + const bool attn_query = (strncmp(op_name, "FLASH_ATTN", 10) == 0); + const bool matmul_query = (strcmp(op_name, "MUL_MAT") == 0 + || strcmp(op_name, "MUL_MAT_ID") == 0); + + for (const auto & b : entries) { + if (attn_query) { + if (b.op_name.compare(0, 10, "FLASH_ATTN") != 0) continue; + } else { + if (b.op_name != op_name) continue; + } + + double dim_score = 0.0; + double batch_score = 0.0; + double quant_score = 0.0; + + if (matmul_query) { + const double n_diff = std::abs((double)b.N - N) / std::max(N, (int64_t)1); + const double k_diff = std::abs((double)b.K - K) / std::max(K, (int64_t)1); + dim_score = n_diff + k_diff; + } else if (attn_query) { + const double ctx_diff = std::abs((double)b.ctx_len - ctx_len) / std::max(ctx_len, (int64_t)1); + double kv_diff = 0.0; + if (n_kv_heads > 0 && b.n_heads > 0) { + kv_diff = std::abs((double)b.n_heads - n_kv_heads) / std::max(n_kv_heads, (int64_t)1); + } + dim_score = ctx_diff + kv_diff * 0.5; + } else { + const double elem_diff = std::abs((double)b.n_elements - n_elements) / std::max(n_elements, (int64_t)1); + dim_score = elem_diff; + } + + if (target_batch > 0 && b.B > 0) { + batch_score = std::abs((double)b.B - target_batch) / std::max(target_batch, (int64_t)1); + } + + if (target_bpw > 0) { + const double bench_bpw = get_bits_per_weight(b.quant); + if (bench_bpw > 0) { + quant_score = std::abs(bench_bpw - target_bpw) / target_bpw; + } else { + quant_score = 1.0; + } + } else if (quant && b.quant != quant) { + continue; + } + + // batch most important (memory vs compute regime), then dims, then quant + const double score = batch_score * 1.0 + dim_score * 0.5 + quant_score * 0.3; + + if (score < best_score) { + best_score = score; + best = &b; + } + } + + return best; +} + +double llama_benchmark_predictor::predict_tps( + ggml_backend_sched_t sched, + int cpu_backend_id, + uint32_t kv_size, + int32_t batch_size, + uint32_t n_outputs, + bool has_rs) const { + + const int n_splits = ggml_backend_sched_get_n_splits(sched); + if (n_splits <= 0) return 0.0; + + LLAMA_LOG_DEBUG("%s: n_splits=%d, bs=%d, kv_size=%u, n_outputs=%u, has_rs=%d\n", + __func__, n_splits, batch_size, kv_size, n_outputs, (int)has_rs); + + const double pcie_bw = stats.peak_pcie_bw; + // KV transfer ratios vs full writeback_bytes the sched reports: + // download: only newly written cells -> batch_size/kv_size (decode bs=1/kv=1024 -> 0.1%) + // upload: all used cells (cache fills over time) -> 1.0 (conservative, assumes full cache) + // RS models use FULL mode (entire tensor every batch), so both are 1.0. + const double dl_ratio = has_rs ? 1.0 + : ((kv_size > 0) ? std::min(1.0, (double)batch_size / kv_size) : 1.0); + const double ul_ratio = 1.0; + double total_ms = 0.0; + bool copy_prefetched = false; + timing_cache_t timing_cache; + + for (int i = 0; i < n_splits; i++) { + struct ggml_backend_sched_split_info si = {}; + if (!ggml_backend_sched_get_split_info(sched, i, &si)) continue; + + const bool is_gpu = (si.backend_id != cpu_backend_id); + + double input_copy_ms = 0.0; + double input_copy_bytes = 0.0; + if (is_gpu && !copy_prefetched && pcie_bw > 0.0) { + input_copy_bytes = (double)si.input_weight_copy_bytes + + (double)si.writeback_bytes * ul_ratio; + input_copy_ms = (input_copy_bytes / 1e9 / pcie_bw) * 1000.0; + } + + // peek at next split to determine if async prefetch will overlap with this split + double prefetch_bytes = 0.0; + bool next_copy_prefetched = false; + if (i + 1 < n_splits) { + struct ggml_backend_sched_split_info next_si = {}; + if (ggml_backend_sched_get_split_info(sched, i + 1, &next_si) && + next_si.can_prefetch_weights) { + prefetch_bytes = (double)next_si.input_weight_bytes + + (double)next_si.writeback_bytes * ul_ratio; + next_copy_prefetched = prefetch_bytes > 0.0; + } + } + + const bool async_copy = (prefetch_bytes > 0.0); + + // compute cost (CPU splits use eff_gflops when async_copy due to PCIe contention) + struct ggml_tensor ** nodes = ggml_graph_nodes(si.graph); + int n_nodes = ggml_graph_n_nodes(si.graph); + llama_split_timing t = predict_split(nodes, n_nodes, is_gpu, batch_size, async_copy, &timing_cache); + + // prefetch cost: use concurrent PCIe BW for CPU splits (bus shared with DRAM), + // peak PCIe BW for GPU splits (GPU compute doesn't contend with PCIe DMA) + double prefetch_ms = 0.0; + if (prefetch_bytes > 0.0) { + const double eff_bw = (!is_gpu && t.eff_pcie_bw > 0.0) ? t.eff_pcie_bw : pcie_bw; + prefetch_ms = (prefetch_bytes / 1e9 / eff_bw) * 1000.0; + } + + // output scaling: use the output rows in the reserved graph, then scale to + // the runtime number of logits. This keeps the memory-probe graph as the + // source of truth and avoids double-scaling when it was already reduced. + if (i == n_splits - 1 && n_outputs > 0) { + for (int j = 0; j < n_nodes; j++) { + if (nodes[j]->name && strstr(nodes[j]->name, "result_output")) { + const llama_op_metrics out_m = llama_op_metrics_compute(nodes[j]); + const int32_t graph_outputs = (int32_t) std::max(1, out_m.M); + if (graph_outputs != batch_size || n_outputs < (uint32_t) graph_outputs) { + llama_split_timing out_included = predict_split(&nodes[j], 1, is_gpu, batch_size, async_copy, &timing_cache); + llama_split_timing out_graph = predict_split(&nodes[j], 1, is_gpu, graph_outputs, async_copy, &timing_cache); + if (out_included.time_ms > 0.0 || out_graph.time_ms > 0.0) { + const double scale = n_outputs < (uint32_t) graph_outputs + ? (double)n_outputs / (double)graph_outputs + : 1.0; + const double scaled_ms = out_graph.time_ms * scale; + t.time_ms -= out_included.time_ms; + t.time_ms += scaled_ms; + LLAMA_LOG_DEBUG("%s: out_t: included=%f graph=%f scaled=%f rows=%d->%u\n", + __func__, out_included.time_ms, out_graph.time_ms, scaled_ms, graph_outputs, n_outputs); + } + } + break; + } + } + } + + // KV/RS download cost (GPU splits, after compute, on compute stream) + double kv_dl_ms = 0.0; + if (is_gpu && si.writeback_bytes > 0 && pcie_bw > 0.0) { + double dl_bytes = (double)si.writeback_bytes * dl_ratio; + kv_dl_ms = (dl_bytes / 1e9 / pcie_bw) * 1000.0; + } + + // activation copy cost (synchronous, between splits on different backends) + double activ_copy_ms = 0.0; + if (si.input_activ_bytes > 0 && pcie_bw > 0.0) { + activ_copy_ms = ((double)si.input_activ_bytes / 1e9 / pcie_bw) * 1000.0; + } + + // Prefetch is enqueued after current inputs/precompute and before graph compute, + // so whatever remains can overlap the current split compute. CPU splits use the + // effective PCIe BW above because DMA contends with CPU memory traffic. + double split_ms = input_copy_ms + activ_copy_ms; + split_ms += std::max(t.time_ms + kv_dl_ms, prefetch_ms); + total_ms += split_ms; + copy_prefetched = next_copy_prefetched; + + double dl_bytes = is_gpu ? (double)si.writeback_bytes * dl_ratio : 0.0; + LLAMA_LOG_DEBUG("%s: split %d/%d [%s] input_copy=%.3f (%.2f MiB) compute=%.3f kv_dl=%.3f (%.2f MiB) prefetch=%.3f (%.2f MiB) activ=%.3f (%.2f MiB) -> %.3f ms" + " (exact=%d near=%d fall=%d cache=%d)\n", + __func__, i, n_splits, is_gpu ? "GPU" : "CPU", + input_copy_ms, input_copy_bytes / (1024.0 * 1024.0), + t.time_ms, + kv_dl_ms, dl_bytes / (1024.0 * 1024.0), + prefetch_ms, prefetch_bytes / (1024.0 * 1024.0), + activ_copy_ms, (double)si.input_activ_bytes / (1024.0 * 1024.0), + split_ms, + t.n_exact, t.n_nearest, t.n_fallback, t.n_cache_hit); + } + + double tps = (total_ms > 0.0) ? (batch_size * 1000.0 / total_ms) : 0.0; + LLAMA_LOG_DEBUG("%s: total=%.3f ms, dl_ratio=%.4f, ul_ratio=%.1f, pcie_bw=%.1f GB/s -> %.1f tps\n", + __func__, total_ms, dl_ratio, ul_ratio, pcie_bw, tps); + return tps; +} diff --git a/src/llama-benchmark.h b/src/llama-benchmark.h new file mode 100644 index 000000000000..8fbaf729d99c --- /dev/null +++ b/src/llama-benchmark.h @@ -0,0 +1,151 @@ +#pragma once + +#include +#include +#include +#include +#include + +struct ggml_tensor; +typedef struct ggml_backend_sched * ggml_backend_sched_t; + +// Per-op FLOPS/bytes metrics extracted from a ggml tensor node. +// Used by roofline prediction to classify ops as compute-bound or memory-bound. +struct llama_op_metrics { + double ops = 0.0; + double bytes = 0.0; + const char * quant_type = nullptr; + + int64_t N = 0, K = 0, M = 0; + int64_t ctx_len = 0; + int64_t n_elements = 0; + int64_t n_kv_heads = 0; + int64_t n_experts_used = 0; + int64_t head_dim = 0; + int64_t n_q_heads = 0; + int64_t n_tokens = 0; +}; + +// Extract ops/bytes metrics from a compute graph node. +// Zero-cost view ops (RESHAPE, VIEW, PERMUTE, TRANSPOSE, NONE) return {0,0}. +llama_op_metrics llama_op_metrics_compute(const ggml_tensor * node); + +// One row from profiler output (CPU or GPU). +// For FLASH_ATTN ops, n_heads stores n_kv_heads (written by profiler). +struct llama_benchmark_entry { + std::string op_name; + std::string quant; + int threads = 0; + + double ai = 0.0; // arithmetic intensity (FLOP/byte) + double bw_gb_s = 0.0; // effective memory bandwidth (GB/s) + double peak_gflops = 0.0; // standalone compute throughput (GFLOP/s) + double ridge = 0.0; // roofline ridge point (FLOP/byte) + double eff_gflops = 0.0; // concurrent compute (CPU: under PCIe load; GPU: == peak) + double eff_pcie_bw = 0.0; // effective PCIe BW during concurrent compute (CPU only) + + int64_t N = 0, K = 0, B = 0; + int64_t n_tokens = 0; + int64_t ctx_len = 0; + int64_t n_heads = 0; // n_kv_heads for FLASH_ATTN, unused for other ops + int64_t head_dim = 0; + int64_t n_elements = 0; +}; + +// Global bandwidth/compute stats parsed from profiler file headers. +struct llama_benchmark_stats { + // CPU (from cpu profiler header) + double peak_system_bw = 0.0; // peak DRAM BW (GB/s) + double peak_pcie_bw = 0.0; // peak PCIe BW standalone (GB/s) + double eff_system_bw = 0.0; // DRAM BW under concurrent PCIe load (GB/s) + double eff_pcie_bw = 0.0; // PCIe BW under concurrent CPU load (GB/s) + + // GPU (from gpu profiler header) + double peak_gpu_mem_bw = 0.0; // peak GPU memory BW (GB/s) + double peak_gpu_compute = 0.0; // peak GPU compute (GFLOP/s) +}; + +// Aggregate timing prediction for a set of graph nodes (typically one sched split). +struct llama_split_timing { + double time_ms = 0.0; + double eff_pcie_bw = 0.0; // time-weighted avg PCIe BW (CPU splits only) + int op_count = 0; + int n_cache_hit = 0; // timing cache reuses + int n_exact = 0; // exact benchmark hash hits + int n_nearest = 0; // nearest-neighbor matches + int n_fallback = 0; // memory BW fallback (no benchmark match) +}; + +// Benchmark data + O(1) hash maps + nearest-neighbor lookup + split timing prediction. +// Owns entry vectors; map pointers are stable after load. +struct llama_benchmark_predictor { + std::vector cpu_entries; + std::vector gpu_entries; + llama_benchmark_stats stats = {}; + + std::unordered_map cpu_map; + std::unordered_map gpu_map; + + // Timing cache: same shape + same mode = same timing across layers. + // Key: make_timing_key(), Value: {time_ms, pcie_contrib}. + using timing_cache_t = std::unordered_map>; + + llama_benchmark_predictor() = default; + llama_benchmark_predictor(const llama_benchmark_predictor &) = delete; + llama_benchmark_predictor & operator=(const llama_benchmark_predictor &) = delete; + + // Parse profiler output, populate entries + hash maps. + // CPU loader filters by n_threads; returns false if file missing or empty. + bool load_cpu(const char * filepath, int n_threads); + bool load_gpu(const char * filepath); + + bool has_cpu() const { return !cpu_entries.empty(); } + bool has_gpu() const { return !gpu_entries.empty(); } + + // Predict aggregate timing for a set of graph nodes. + // nodes/n_nodes: typically from one sched split (via ggml_graph_nodes/n_nodes). + // async_copy: true when PCIe transfers run concurrently (uses eff_* bandwidths for CPU). + llama_split_timing predict_split( + struct ggml_tensor ** nodes, int n_nodes, + bool is_gpu, int32_t batch_size, bool async_copy = true, + timing_cache_t * timing_cache = nullptr) const; + + // Nearest-neighbor search with weighted scoring (batch > dims > quant). + // For FLASH_ATTN queries, pass op_name="FLASH_ATTN" -- matches all FLASH_ATTN_* entries. + static const llama_benchmark_entry * find_nearest( + const std::vector & entries, + const char * op_name, const char * quant, + int64_t N, int64_t K, int64_t ctx_len, int64_t n_elements, + int64_t n_kv_heads = -1, int64_t target_batch = 1); + + // Key generators for hash map lookup (used by predict to build exact-match keys). + static std::string make_key( + const std::string & op, const std::string & quant, + int64_t N, int64_t K, int64_t batch); + + static std::string make_attn_key( + int64_t ctx_len, int64_t n_kv_heads, int64_t n_tokens); + + static std::string make_elem_key( + const std::string & op, int64_t n_elements); + + // Timing cache key: encodes backend type, async mode, op, shape, and batch. + static std::string make_timing_key( + bool is_gpu, bool async_copy, const char * op_name, + const llama_op_metrics & m, int32_t batch_size); + + // End-to-end TPS prediction across all sched splits. + // kv_size: total KV cells per layer (for partial upload ratio). + // has_rs: true if model has recurrent state layers (RS uses full-tensor transfer, not cell-granular). + // Returns 0 if no splits or no benchmark data. + double predict_tps( + ggml_backend_sched_t sched, + int cpu_backend_id, + uint32_t kv_size, + int32_t batch_size, + uint32_t n_outputs = 0, + bool has_rs = false) const; + +private: + void build_maps(); +}; diff --git a/src/llama-context-pshard.cpp b/src/llama-context-pshard.cpp new file mode 100644 index 000000000000..d064facb7299 --- /dev/null +++ b/src/llama-context-pshard.cpp @@ -0,0 +1,695 @@ +#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 +#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) { + const auto & tbids = model.get_tensor_backend_ids(); + const auto & lbids = model.get_layer_backend_ids(); + + for (const auto & [tensor, bid] : tbids) { + if (bid >= 0 && bid < (int32_t) backends.size()) { + ggml_backend_sched_set_tensor_backend_hint(sched, tensor, backends[bid].get()); + } + } + + if (memory) { + for (auto * ps : memory->get_pipe_shards()) { + ps->assign_tensors(sched, lbids, backends, layout); + } + } +} + +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; + const uint32_t n_outputs = n_tokens; + + // 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_outputs, 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_outputs, 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_plan_reserves() { + 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_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_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(); + + // 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++; + } + } + + 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); +} + +// 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()); +} + +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; + + 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) { + 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); + } + 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 c512477c0eab..13c5353d2d7e 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -6,12 +6,17 @@ #include "llama-impl.h" #include "llama-batch.h" #include "llama-io.h" +#include "llama-kv-cache.h" #include "llama-memory.h" +#include "llama-pipe-shard.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 @@ -81,8 +86,10 @@ static const llm_fused_op_probe llm_fused_op_dsv4_hc_post_probe = { llama_context::llama_context( const llama_model & model, - llama_context_params params) : + llama_context_params params, + llama_context_probe_reserve probe_reserve_) : model(model), + probe_reserve(probe_reserve_), cvec(std::make_unique()), loras(std::make_unique()), balloc(std::make_unique(model.hparams.n_pos_per_embd())) { @@ -125,6 +132,7 @@ llama_context::llama_context( cparams.ctx_type = params.ctx_type; cparams.pooling_type = params.pooling_type; + 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; @@ -258,6 +266,16 @@ llama_context::llama_context( cparams.fused_dsv4_hc_post = true; cparams.auto_fhc = true; + // pshard remaps kv by layer, so the device matching in resolve_fused_ops + // (which compares against model.dev_layer()) cannot be trusted - disable all + // automatic fused-op resolution + if (cparams.pshard) { + cparams.auto_fa = false; + cparams.auto_fgdn = false; + cparams.auto_flid = false; + cparams.auto_fhc = 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; @@ -324,12 +342,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) @@ -351,6 +381,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()); @@ -453,8 +488,18 @@ llama_context::llama_context( LLAMA_LOG_INFO("%s: pipeline parallelism enabled\n", __func__); } + if (cparams.pshard && !model.hparams.no_alloc) { + pshard_pack_cache_region(); + } + sched_reserve(); + if (cparams.pshard && !model.hparams.no_alloc) { + pshard_warmup_plan_reserves(); + const_cast(model).sync_dev_preload(); + pshard_apply_initial_plan(); + } + 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"); @@ -584,7 +629,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); @@ -595,6 +648,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__); @@ -618,75 +675,89 @@ void llama_context::sched_reserve() { int n_splits_tg = -1; int n_nodes_tg = -1; - const uint32_t n_outputs_pp = std::min(n_tokens, cparams.n_outputs_max); - - // reserve pp (prompt processing) graph first so that buffers are only allocated once - { - auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), - model.hparams.no_alloc, model.hparams.no_alloc ? backend_buf_exp_size.data() : nullptr); + if (cparams.pshard && model.hparams.no_alloc) { + const uint32_t reserve_tokens = probe_reserve.n_tokens ? probe_reserve.n_tokens : n_tokens; + const uint32_t reserve_outputs = probe_reserve.n_outputs ? probe_reserve.n_outputs : reserve_tokens; + auto * gf = graph_reserve(reserve_tokens, n_seqs, reserve_outputs, mctx.get(), + true, backend_buf_exp_size.data()); if (!gf) { - if (cparams.pipeline_parallel) { - LLAMA_LOG_WARN("%s: compute buffer allocation failed, retrying without pipeline parallelism\n", __func__); - cparams.pipeline_parallel = false; - sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, cparams.op_offload)); - gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get()); - } + throw std::runtime_error("failed to measure compute buffers for pshard probe"); + } + } else if (cparams.pshard) { + LLAMA_LOG_INFO("%s: pshard enabled, skipping baseline reserves (deferred to first plan apply)\n", __func__); + } else { + + const uint32_t n_outputs_pp = std::min(n_tokens, cparams.n_outputs_max); + + // reserve pp (prompt processing) graph first so that buffers are only allocated once + { + auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, 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_outputs_pp, 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_outputs_pp, 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 + { + // 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_outputs_pp, 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(); @@ -1339,11 +1410,14 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll } n_reused++; + } else { 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(); @@ -1357,11 +1431,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); + } + 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 @@ -1374,6 +1456,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); @@ -1771,6 +1857,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; @@ -1778,10 +1868,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; } @@ -2396,6 +2495,10 @@ ggml_cgraph * llama_context::graph_reserve( this->n_outputs = save_n_outputs; + if (cparams.pshard) { + pshard_assign_tensors(sched.get(), model, memory.get(), backends, pshard_layout); + } + // initialize scheduler with the specified graph if (split_only) { if (sizes) { @@ -2466,17 +2569,131 @@ ggml_status llama_context::graph_compute( } llm_graph_cb llama_context::graph_get_cb() const { - return [&](const llama_ubatch & ubatch, ggml_tensor * cur, const char * name, int il) { + // keep weight islands on last_weight_bid and return mixed tails to layer_bid + const bool delegate_compute = model.pshard_delegates_compute(); + int32_t last_weight_bid = -1; + int last_il = -2; + return [&, delegate_compute, last_weight_bid, last_il](const llama_ubatch & ubatch, ggml_tensor * cur, const char * name, int il) mutable { if (il >= 0) { ggml_format_name(cur, "%s-%d", name, il); } else { ggml_set_name(cur, name); } + // In delegated mode, layer nodes stay with the scheduler; non-layer heads/tails still honor exact tensor overrides. + if (cparams.pshard && delegate_compute && cur->view_src == nullptr) { + const auto & lbids = model.get_layer_backend_ids(); + const bool has_layer_backend = il >= 0 && lbids.find(il) != lbids.end(); + + if (!has_layer_backend) { + const auto & tbids = model.get_tensor_backend_ids(); + for (int j = 0; j < GGML_MAX_SRC; j++) { + ggml_tensor * src = cur->src[j]; + if (!src) { + continue; + } + + while (src->view_src) { + src = src->view_src; + } + + auto it = tbids.find(src); + if (it != tbids.end() && it->second >= 0 && it->second < (int32_t) backends.size()) { + ggml_backend_sched_set_tensor_backend(sched.get(), cur, backends[it->second].get()); + break; + } + } + } + } + + if (cparams.pshard && !delegate_compute) { + auto sched_backend_id = [&](ggml_backend_t backend) -> int32_t { + if (backend == nullptr) { + return -1; + } + for (int32_t i = 0; i < (int32_t) backends.size(); i++) { + if (backends[i].get() == backend) { + return i; + } + } + return -1; + }; + + int32_t bid = -1; + bool has_tensor_backend = false; + + if (il != last_il) { + last_il = il; + last_weight_bid = -1; + } + + const auto & tbids = model.get_tensor_backend_ids(); + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (!cur->src[j]) { + continue; + } + auto it = tbids.find(cur->src[j]); + if (it != tbids.end() && it->second >= 0) { + bid = it->second; + has_tensor_backend = true; + break; + } + } + + const auto & lbids = model.get_layer_backend_ids(); + auto lit = lbids.find(il); + const int32_t layer_bid = (lit != lbids.end()) ? lit->second : -1; + + bool mixed_src_backends = false; + if (!has_tensor_backend) { + int32_t src_bid = -1; + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (!cur->src[j]) { + continue; + } + + ggml_tensor * src = cur->src[j]; + ggml_backend_t src_backend = ggml_backend_sched_get_tensor_backend(sched.get(), src); + while (src_backend == nullptr && src->view_src != nullptr) { + src = src->view_src; + src_backend = ggml_backend_sched_get_tensor_backend(sched.get(), src); + } + + const int32_t cur_src_bid = sched_backend_id(src_backend); + if (cur_src_bid < 0) { + continue; + } + + if (src_bid < 0) { + src_bid = cur_src_bid; + } else if (src_bid != cur_src_bid) { + mixed_src_backends = true; + break; + } + } + } + + if (has_tensor_backend) { + last_weight_bid = bid; + } else if (mixed_src_backends && last_weight_bid >= 0 && layer_bid >= 0) { + bid = layer_bid; + last_weight_bid = bid; + } else if (last_weight_bid >= 0) { + bid = last_weight_bid; + } else { + bid = layer_bid; + } + + // views keep following their source storage. + if (bid >= 0 && bid < (int32_t)backends.size() && cur->view_src == nullptr) { + ggml_backend_sched_set_tensor_backend(sched.get(), cur, backends[bid].get()); + } + } + // norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends // FIXME: fix in ggml_backend_sched const bool full_offload = model.n_gpu_layers() > model.hparams.n_layer_all; - 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) { @@ -2900,7 +3117,7 @@ class llama_io_read_device : public llama_io_read_i { size_t llama_context::state_get_size() { llama_io_write_dummy io(false); 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; @@ -2910,7 +3127,7 @@ size_t llama_context::state_get_size() { size_t llama_context::state_get_data(uint8_t * dst, size_t size) { llama_io_write_host 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; @@ -2920,7 +3137,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_host 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; @@ -2935,7 +3152,7 @@ size_t llama_context::state_seq_get_size(llama_seq_id seq_id, llama_state_seq_fl io.write(&io_magic, sizeof(io_magic)); io.write(&seq_id, sizeof(seq_id)); - 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; @@ -2954,7 +3171,7 @@ size_t llama_context::state_seq_get_data(llama_seq_id seq_id, uint8_t * dst, siz io->write(&io_magic, sizeof(io_magic)); io->write(&seq_id, sizeof(seq_id)); - 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; @@ -2993,7 +3210,7 @@ size_t llama_context::state_seq_set_data(llama_seq_id seq_id, const uint8_t * sr llama_seq_id seq_id_read; io->read(&seq_id_read, sizeof(seq_id_read)); - 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; @@ -3032,7 +3249,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); @@ -3055,7 +3272,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; } @@ -3091,7 +3308,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; @@ -3115,7 +3332,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()); @@ -3123,72 +3340,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(); - return io.n_bytes(); + try { + if (memory) { + memory->state_write(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; + } } -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 // @@ -3500,14 +3768,16 @@ llama_context_params llama_context_default_params() { /*.sampler =*/ nullptr, /*.n_sampler =*/ 0, /*.ctx_other =*/ nullptr, + /*.pshard =*/ false, }; return result; } -llama_context * llama_init_from_model( +llama_context * llama_init_from_model_internal( llama_model * model, - llama_context_params params) { + llama_context_params params, + llama_context_probe_reserve probe_reserve) { if (!model) { LLAMA_LOG_ERROR("%s: model cannot be NULL\n", __func__); return nullptr; @@ -3580,7 +3850,7 @@ llama_context * llama_init_from_model( } try { - auto * ctx = new llama_context(*model, params); + auto * ctx = new llama_context(*model, params, probe_reserve); return ctx; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: failed to initialize the context: %s\n", __func__, err.what()); @@ -3589,6 +3859,12 @@ llama_context * llama_init_from_model( return nullptr; } +llama_context * llama_init_from_model( + llama_model * model, + llama_context_params params) { + return llama_init_from_model_internal(model, params, {}); +} + // deprecated llama_context * llama_new_context_with_model( llama_model * model, diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa8b562..41def7dd46eb 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -15,6 +15,7 @@ #include struct llama_model; +struct llama_pshard_plan; class llama_batch_allocr; class llama_io_read_i; @@ -39,11 +40,17 @@ struct llama_memory_buffer { using llama_memory_buffers = std::map; +struct llama_context_probe_reserve { + uint32_t n_tokens = 0; + uint32_t n_outputs = 0; +}; + struct llama_context { // init scheduler and compute buffers, reserve worst-case graphs llama_context( const llama_model & model, - llama_context_params params); + llama_context_params params, + llama_context_probe_reserve probe_reserve = {}); ~llama_context(); @@ -270,8 +277,31 @@ struct llama_context { size_t state_write_data(llama_io_write_i & io); size_t state_read_data (llama_io_read_i & io); - 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); + 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_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, 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, 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 @@ -280,6 +310,7 @@ struct llama_context { const llama_model & model; llama_cparams cparams; + llama_context_probe_reserve probe_reserve; llama_adapter_cvec_ptr cvec; llama_adapter_loras_ptr loras; @@ -345,6 +376,10 @@ 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; std::vector backends; @@ -392,3 +427,21 @@ struct llama_context { mutable int32_t n_reused = 0; // number of times the previous graph was reused }; + +llama_context * llama_init_from_model_internal( + llama_model * model, + llama_context_params params, + llama_context_probe_reserve probe_reserve = {}); + + +// pshard free functions (implemented in llama-context-pshard.cpp) +struct llama_memory_i; + +void pshard_assign_tensors( + ggml_backend_sched_t sched, + const llama_model & model, + llama_memory_i * memory, + const std::vector & backends, + const pshard_dev_layout & layout); + +void pshard_refresh_stream_views(llama_memory_i * memory); diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 5018170ed85e..7434a47fa715 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 #include @@ -62,4 +63,31 @@ struct llama_cparams { void * cb_eval_user_data; llama_context * ctx_other; + + 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-ext.h b/src/llama-ext.h index 348bbae95770..0abcc9b67ac3 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -90,6 +90,63 @@ LLAMA_API ggml_backend_dev_t llama_model_get_device(const struct llama_model * m LLAMA_API llama_memory_breakdown llama_get_memory_breakdown(const struct llama_context * ctx); +// +// pipeline sharding (pshard) +// + +struct llama_pshard_plan_registry; + +// Baseline params-fitting engine injected into the pshard planner. +// +// The engine itself lives in common/fit.cpp (common_fit_params), and common +// links against llama - so libllama cannot call it directly. The caller passes +// it in instead. Implementations must throw on failure. +typedef void (*llama_pshard_fit_fn)( + const char * path_model, + struct llama_model_params * mparams, + struct llama_context_params * cparams, + float * tensor_split, + struct llama_model_tensor_buft_override * tensor_buft_overrides, + size_t * margins_s, + uint32_t n_ctx_min, + enum ggml_log_level log_level, + void * user_data); + +// RUNTIME path. Loads the cached pshard plan registry written by the planner +// (llama-fit-params --pshard) and populates tensor_buft_overrides for model +// loading. Never plans, so startup stays fast. If no usable plan is found, +// clears mparams->pshard and returns, leaving baseline loading in place. +// Prefer common_fit_params_pshard() in common/fit.h. +LLAMA_API void llama_params_fit_pshard_inference( + const char * path_model, + struct llama_model_params * mparams, + struct llama_context_params * cparams, + struct llama_model_tensor_buft_override * tensor_buft_overrides, + size_t max_vram_mb, // 0 = use actual free VRAM minus fit_target_mb + size_t fit_target_mb); // ignored when max_vram_mb > 0 + +// PLANNER path. Probes VRAM usage for each strategy/tier combination, selects the best +// plan (by TPS if benchmark data is available, else by VRAM fit), writes/updates +// .tensor_overrides.pshard_registry and populates tensor_buft_overrides. +// If all layers fit in VRAM, sets mparams->pshard=false and returns (baseline loading). +// Prefer common_pshard_plan() in common/fit.h, which supplies fit_fn for you. +LLAMA_API void llama_params_fit_pshard_planning( + const char * path_model, + struct llama_model_params * mparams, + struct llama_context_params * cparams, + struct llama_model_tensor_buft_override * tensor_buft_overrides, + size_t max_vram_mb, // 0 = use actual free VRAM minus fit_target_mb + size_t fit_target_mb, // ignored when max_vram_mb > 0 + llama_pshard_fit_fn fit_fn, + void * fit_fn_ud); + +// Create/free a tier plan registry. Caller owns the pointer and passes it via +// mparams->pshard_registry before calling either entry point above. +// 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); + // Set whether the context outputs nextn embeddings or not // If masked == true, output the embeddings only for the tokens with batch.logits != 0 // If masked == false, output the embeddings for all tokens in the batch regardless of batch.logits diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 6d1c8f4e42a8..4a76599cb145 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1337,7 +1337,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), @@ -2675,8 +2675,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); } ggml_tensor * kq_mask = inp->get_kq_mask(); @@ -2767,7 +2771,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(); @@ -2925,13 +2931,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 e91866469adf..6f7fa80b8906 100644 --- a/src/llama-kv-cache-iswa.cpp +++ b/src/llama-kv-cache-iswa.cpp @@ -280,6 +280,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 7dab6eaa82c8..8b1637f99bed 100644 --- a/src/llama-kv-cache-iswa.h +++ b/src/llama-kv-cache-iswa.h @@ -78,6 +78,8 @@ class llama_kv_cache_iswa : public llama_memory_i { std::map memory_breakdown() const override; + std::vector get_pipe_shards() override; + // state write/load void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 44cb1668dacf..ff2936855ea4 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -160,6 +160,86 @@ 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 < 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()); + // no_alloc == planner probe: there is no runtime pshard_pack_cache_region() to assign + // addresses later, so pass no external buffer and let init() allocate eagerly. + ggml_backend_buffer_t preload_buf = hparams.no_alloc ? nullptr : model.get_dev_preload_buf(); + if (!pipe_shard_kv->init(specs, model.get_layer_backend_ids(), cpu_bid, hparams.no_alloc, 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 < n_layer; il++) { if (!hparams.has_kv(il)) { LLAMA_LOG_DEBUG("%s: layer %3d: does not have KV cache\n", __func__, il); @@ -752,9 +832,24 @@ std::map llama_kv_cache::memory_breakdown() } } + if (pipe_shard_kv) { + const auto & bufs = pipe_shard_kv->get_bufs(); + const auto & sizes = pipe_shard_kv->get_bufs_planned_sizes(); + for (size_t i = 0; i < bufs.size(); i++) { + if (!bufs[i]) continue; + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(bufs[i].get()); + ret[buft] += sizes[i]; + } + } + return ret; } +std::vector llama_kv_cache::get_pipe_shards() { + if (pipe_shard_kv) return { pipe_shard_kv.get() }; + return {}; +} + llama_memory_context_ptr llama_kv_cache::init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, @@ -919,6 +1014,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 @@ -2903,3 +3003,19 @@ 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; +} diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index d5a92f4405b5..bd1218f748fc 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 @@ -158,6 +159,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; @@ -291,6 +296,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; @@ -409,6 +416,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 06f7fd5428c4..7de83f25ae62 100644 --- a/src/llama-memory-hybrid-iswa.cpp +++ b/src/llama-memory-hybrid-iswa.cpp @@ -192,6 +192,13 @@ std::map llama_memory_hybrid_iswa::memory_br return mb; } +std::vector llama_memory_hybrid_iswa::get_pipe_shards() { + std::vector result; + for (auto * ps : mem_attn->get_pipe_shards()) { result.push_back(ps); } + for (auto * ps : mem_recr->get_pipe_shards()) { result.push_back(ps); } + return result; +} + void llama_memory_hybrid_iswa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { mem_attn->state_write(io, seq_id, flags); mem_recr->state_write(io, seq_id, flags); diff --git a/src/llama-memory-hybrid-iswa.h b/src/llama-memory-hybrid-iswa.h index c9d3f9f57c50..1d6210460b79 100644 --- a/src/llama-memory-hybrid-iswa.h +++ b/src/llama-memory-hybrid-iswa.h @@ -71,6 +71,8 @@ class llama_memory_hybrid_iswa : public llama_memory_i { std::map memory_breakdown() const override; + std::vector get_pipe_shards() override; + // state write/load void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6f..ca2616d9970b 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -187,6 +187,13 @@ std::map llama_memory_hybrid::memory_breakdo return mb; } +std::vector llama_memory_hybrid::get_pipe_shards() { + std::vector result; + for (auto * ps : mem_attn->get_pipe_shards()) { result.push_back(ps); } + for (auto * ps : mem_recr->get_pipe_shards()) { result.push_back(ps); } + return result; +} + void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { mem_attn->state_write(io, seq_id, flags); diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb74991..52308302294b 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -71,6 +71,8 @@ class llama_memory_hybrid : public llama_memory_i { std::map memory_breakdown() const override; + std::vector get_pipe_shards() override; + // state write/load void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; diff --git a/src/llama-memory-pshard.cpp b/src/llama-memory-pshard.cpp new file mode 100644 index 000000000000..2c3c0158921c --- /dev/null +++ b/src/llama-memory-pshard.cpp @@ -0,0 +1,636 @@ +#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); + } +} + +// The planning and inference stages DIVERGE here. Planning dispatches purely on the plan's +// backend_id map and has no notion of pinned layers; inference dispatches on l.is_pinned and +// requires GPU addresses that only exist after pshard_pack_cache_region(), which the planner +// never calls. Both variants are kept verbatim so the difference stays visible. +// +// external_buf is the discriminator: it is set only on the runtime path (it is what defers +// address assignment to pack_cache_region in init()), and null during planner probes. +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) { + if (external_buf) { + assign_tensors_inference(sched, layer_bids, backends, layout); + } else { + assign_tensors_planning(sched, layer_bids, backends, layout); + } +} + +// planning: bid-driven, no pinned-layer concept, no GPU-address precondition +void llama_memory_pshard::assign_tensors_planning( + ggml_backend_sched_t sched, + const std::unordered_map & layer_bids, + const std::vector & backends, + const pshard_dev_layout & layout) { + for (auto & l : layers) { + auto it = layer_bids.find((int)l.il); + const bool has_bid = it != layer_bids.end() && it->second >= 0 && it->second < (int32_t)backends.size(); + if (!has_bid) { + continue; + } + const int32_t bid = it->second; + if (bid == layout.compute) { + activate_gpu(l.il); + } else if (bid == layout.cpu) { + activate_cpu(l.il); + } else { + activate_gpu(l.il); + l.t1_gpu->data = NULL; l.t1_gpu->buffer = NULL; + ggml_backend_sched_set_tensor_backend(sched, l.t1_gpu, backends[bid].get()); + ggml_backend_sched_add_writeback(sched, l.t1_gpu); + if (l.t2_gpu) { + l.t2_gpu->data = NULL; l.t2_gpu->buffer = NULL; + ggml_backend_sched_set_tensor_backend(sched, l.t2_gpu, backends[bid].get()); + ggml_backend_sched_add_writeback(sched, l.t2_gpu); + } + } + } +} + +// inference: pinned-layer aware, expects addresses assigned by pshard_pack_cache_region() +void llama_memory_pshard::assign_tensors_inference( + 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..51cf3bdb5f12 --- /dev/null +++ b/src/llama-memory-pshard.h @@ -0,0 +1,132 @@ +#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; + + // dispatches to the _planning / _inference variant based on external_buf + void assign_tensors( + ggml_backend_sched_t sched, + const std::unordered_map & layer_bids, + const std::vector & backends, + const pshard_dev_layout & layout) override; + + void assign_tensors_planning( + ggml_backend_sched_t sched, + const std::unordered_map & layer_bids, + const std::vector & backends, + const pshard_dev_layout & layout); + + void assign_tensors_inference( + ggml_backend_sched_t sched, + const std::unordered_map & layer_bids, + const std::vector & backends, + const pshard_dev_layout & layout); + + 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 ef82eb976ca7..c1379b5bea98 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -72,6 +72,55 @@ 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()); + // no_alloc == planner probe: there is no runtime pshard_pack_cache_region() to assign + // addresses later, so pass no external buffer and let init() allocate eagerly. + ggml_backend_buffer_t preload_buf = hparams.no_alloc ? nullptr : model.get_dev_preload_buf(); + if (!pipe_shard_rs->init(specs, model.get_layer_backend_ids(), cpu_bid, hparams.no_alloc, 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); @@ -401,9 +450,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 b13b7b748f5e..edb6a3e0aa43 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 @@ -112,6 +113,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 db825396645e..45bb9d8bdecc 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -124,6 +124,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 43447f57d30b..11ec4f3d8e49 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1150,7 +1150,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)) { @@ -1252,7 +1256,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; @@ -1269,7 +1274,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); @@ -1404,7 +1420,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); @@ -1528,6 +1545,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); @@ -1696,3 +1717,372 @@ 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); } } + +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, + 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, + const std::vector * preload_order) { + + 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; + 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); + } + } + } + + { + 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()); + } + + 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) { + 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 75a3652d06bd..0e6248d8ce6d 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -14,9 +14,17 @@ #include #include #include +#include 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; + bool device_only_common = false; +}; + // lists of buffer types used for each layer using buft_list_t = std::vector>; @@ -79,6 +87,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; @@ -198,7 +207,27 @@ 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, + 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, + const std::vector * preload_order = nullptr); std::string ftype_name() const; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 51796921081f..704a5f206379 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -7,6 +7,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" @@ -20,12 +21,14 @@ #include "models/models.h" #include "ggml.h" +#include "ggml-alloc.h" #include "ggml-cpp.h" #include #include #include #include +#include #include #include #include @@ -35,6 +38,7 @@ #include #include #include +#include #include static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params & params) { @@ -337,6 +341,36 @@ llama_model * llama_model_create(llama_model_loader & ml, const llama_model_para return llama_model_create(arch, params); } +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; @@ -1034,6 +1068,26 @@ struct llama_model::impl { bool has_tensor_overrides; std::vector tensor_split_owned; + + 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; + + 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()) { @@ -1050,6 +1104,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_base::load_stats(llama_model_loader & ml) { @@ -1523,9 +1582,66 @@ bool llama_model_base::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()); + } + } + + 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()); @@ -1534,6 +1650,15 @@ bool llama_model_base::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(); @@ -1592,17 +1717,21 @@ bool llama_model_base::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); + } } } @@ -1647,7 +1776,9 @@ bool llama_model_base::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; } } @@ -1658,6 +1789,82 @@ bool llama_model_base::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 (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, + 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)); + } + } + return true; } @@ -2010,6 +2217,495 @@ bool llama_model::has_tensor_overrides() const { return pimpl->has_tensor_overrides; } +bool llama_model::is_pshard() const { + return params.pshard; +} + +bool llama_model::pshard_delegates_compute() const { + return params.pshard_delegate_compute; +} + +llama_pshard_plan_registry * llama_model::get_plan_registry() const { + return pimpl->plan_registry; +} + +const std::unordered_map & llama_model::get_tensor_backend_ids() const { + return pimpl->tensor_backend_ids; +} + +const std::unordered_map & llama_model::get_layer_backend_ids() const { + return pimpl->layer_backend_ids; +} + +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; + } +} + +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; + } + 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) { + 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", + __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); + 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()); + 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 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; + } + } + } 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); + 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 (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); + pimpl->weight_preload_map[tensor].gpu_addr = tensor->data; + } + scratch_off = talloc.offset; + + plan.cached_weight_offsets.clear(); + for (auto * tensor : pinned) { + 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); + } + } + 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); + 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; + } + } + } + + size_t n_uploaded = 0; + 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) { + ggml_backend_tensor_set_async(gpu, tensor, entry.cpu_addr, 0, ggml_nbytes(tensor)); + n_uploaded++; + bytes_uploaded += ggml_nbytes(tensor); + } + } + + 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); + } + + 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) { @@ -2336,6 +3032,11 @@ llama_model_params llama_model_default_params() { /*.use_extra_bufts =*/ true, /*.no_host =*/ false, /*.no_alloc =*/ false, + /*.pshard =*/ false, + /*.pshard_delegate_compute =*/ false, + /*.pshard_cache_skip_load =*/ false, + /*.max_vram_alloc =*/ 0, + /*.pshard_registry =*/ nullptr, }; return result; diff --git a/src/llama-model.h b/src/llama-model.h index 36d0480e5eb7..f31a72eb8255 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 { @@ -668,6 +671,29 @@ 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; + + 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); + 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); + + 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..3c71e59233ce --- /dev/null +++ b/src/llama-pshard-cache.cpp @@ -0,0 +1,772 @@ +#include "llama-pshard-plan.h" +#include "llama-impl.h" + +#include "ggml-backend.h" + +#include +#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_inference( + 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, 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_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_inference( + llama_pshard_plan_registry * registry, uint64_t fingerprint, + const char * cache_path, ggml_backend_buffer_type_t host_buft, + size_t current_budget, bool require_exact_budget) { + if (!registry || !cache_path) return false; + + FILE * f = fopen(cache_path, "r"); + if (!f) return false; + + char line[8192]; + bool in_section = false; + + char fp_header[64]; + snprintf(fp_header, sizeof(fp_header), "[fingerprint=0x%016" PRIx64 "]", fingerprint); + + struct tier_data { + uint32_t bs = 0; + bool viable = false; + llama_pshard_strategy strategy = LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS; + uint32_t n_pinned = 0; + uint32_t n_attn_pinned = 0; + int overflow = 0; + float tps = 0.0f; + double vram_mib = 0.0; + int output_on_gpu = 0; + int pin_from_back = 0; + std::string ot_line; + }; + struct variant_data { + uint32_t budget_mib = 0; + uint32_t cache_ubatch = 0; + bool pshard_disabled = false; + double baseline_vram_mib = 0.0; + std::vector tiers; + }; + std::vector variants; + variant_data * cur_variant = nullptr; + + while (fgets(line, sizeof(line), f)) { + std::string s = line; + while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) s.pop_back(); + + if (s.compare(0, 13, "[fingerprint=") == 0) { + if (in_section) break; + 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); + 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); + } + } + 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; +} + +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.load_mode = LLAMA_LOAD_MODE_NONE; + + 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_all; + 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) { + 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_inference", 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_inference( + const char * path_model, + struct llama_model_params * mparams, + struct llama_context_params * cparams, + struct llama_model_tensor_buft_override * tensor_buft_overrides, + size_t max_vram_mb, + size_t fit_target_mb) { + const 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; + } + + 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 (probe.devs.empty()) { + LLAMA_LOG_WARN("%s: no GPU devices found, disabling pshard\n", __func__); + mparams->pshard = false; + cparams->pshard = false; + return; + } + + const auto & devs = probe.devs; + const uint32_t n_layers = probe.n_layers; + const size_t vram_free = probe.vram_budget; + + 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 = 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)"); + + const uint32_t n_ctx_plan = cparams->n_ctx > 0 ? cparams->n_ctx : probe.n_ctx_train; + + auto * registry = mparams->pshard_registry; + if (!registry) { + LLAMA_LOG_ERROR("%s: pshard_registry is null\n", __func__); + mparams->pshard = false; + 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(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(n_ctx_plan, 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_inference(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(n_ctx_plan, 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]); + 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) { + 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); + llama_pshard_generate_overrides_inference( + 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.cpp b/src/llama-pshard-plan.cpp new file mode 100644 index 000000000000..069876d12c9e --- /dev/null +++ b/src/llama-pshard-plan.cpp @@ -0,0 +1,2055 @@ +#include "llama-pshard-plan.h" + +#include "llama-benchmark.h" +#include "llama-impl.h" +#include "llama-memory.h" +#include "llama-model.h" +#include "llama-model-loader.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +static std::mutex g_probe_mutex; + +// Loads the model + context with no_alloc and reports per-device memory use. +// +// This mirrors common_get_device_memory_data() in common/fit.cpp, but adds the +// probe hook and probe-reserve overrides that pshard tier probing needs - both +// require llama_init_from_model_internal(), which is internal to libllama, so +// this cannot live in common/. +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, + uint32_t & hp_n_embd_r, + enum ggml_log_level log_level, + llama_probe_hook_t probe_hook, void * probe_hook_data, + uint32_t probe_n_tokens, uint32_t probe_n_outputs) { + struct user_data_t { + struct { + ggml_log_callback callback; + void * user_data; + } original_logger; + ggml_log_level min_level; // prints below this log level go to debug log + }; + static std::mutex log_mutex; + std::unique_lock log_lock(log_mutex); + + user_data_t ud; + llama_log_get(&ud.original_logger.callback, &ud.original_logger.user_data); + ud.min_level = log_level; + + 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 >= ud->min_level ? level : GGML_LOG_LEVEL_DEBUG; + ud->original_logger.callback(level_eff, text, ud->original_logger.user_data); + }, &ud); + + log_lock.unlock(); + + llama_model_params mparams_copy = *mparams; + mparams_copy.no_alloc = true; + mparams_copy.load_mode = LLAMA_LOAD_MODE_NONE; + + llama_model * model = llama_model_load_from_file(path_model, mparams_copy); + if (model == nullptr) { + llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); + throw std::runtime_error("failed to load model"); + } + + llama_context_probe_reserve probe_reserve = { probe_n_tokens, probe_n_outputs }; + llama_context * ctx = llama_init_from_model_internal(model, *cparams, probe_reserve); + if (ctx == nullptr) { + llama_model_free(model); + llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); + throw std::runtime_error("failed to create llama_context from model"); + } + + std::vector ret(model->devices.size()); + + llama_memory_breakdown memory_breakdown = ctx->memory_breakdown(); + + for (const auto & [buft, mb] : memory_breakdown) { + if (ggml_backend_buft_is_host(buft)) { + continue; + } + + ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft); + if (!dev) { + continue; + } + for (size_t i = 0; i < ret.size(); i++) { + if (model->devices[i].dev == dev) { + ret[i].mb.model += mb.model; + ret[i].mb.context += mb.context; + ret[i].mb.compute += mb.compute; + break; + } + } + } + for (size_t i = 0; i < ret.size(); i++) { + size_t free; + size_t total; + ggml_backend_dev_memory(model->devices[i].dev, &free, &total); + + // devices can return 0 bytes for free and total memory if they do not + // have any to report. in this case, we will use the host memory as a fallback + // fixes: https://github.com/ggml-org/llama.cpp/issues/18577 + if (free == 0 && total == 0) { + ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); + if (cpu_dev == nullptr) { + throw std::runtime_error(format("%s: no CPU backend found", __func__)); + } + ggml_backend_dev_memory(cpu_dev, &free, &total); + } + ret[i].free = free; + ret[i].total = total; + } + + devs = model->devices; + hp_ngl = model->hparams.n_layer_all; + hp_n_ctx_train = model->hparams.n_ctx_train; + hp_n_expert = model->hparams.n_expert; + hp_n_embd_r = model->hparams.n_embd_r(); + + if (probe_hook) { + probe_hook(ctx, probe_hook_data); + } + + llama_free(ctx); + llama_model_free(model); + + { + std::lock_guard lg(log_mutex); + llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); + } + + return ret; +} + +static std::vector llama_get_device_memory_data_safe( + const char * path_model, const llama_model_params * mparams, + const llama_context_params * cparams, + std::vector & devs, uint32_t & hp_ngl, + uint32_t & hp_n_ctx_train, uint32_t & hp_n_expert, uint32_t & hp_n_embd_r, + ggml_log_level log_level, + llama_probe_hook_t probe_hook = nullptr, + void * probe_hook_data = nullptr, + uint32_t probe_n_tokens = 0, + uint32_t probe_n_outputs = 0) { + std::lock_guard lock(g_probe_mutex); + return llama_get_device_memory_data(path_model, mparams, cparams, devs, + hp_ngl, hp_n_ctx_train, hp_n_expert, hp_n_embd_r, log_level, probe_hook, probe_hook_data, + probe_n_tokens, probe_n_outputs); +} + + +static void llama_pshard_generate_overrides_planning( + uint32_t n_pinned, + uint32_t n_layers, + ggml_backend_buffer_type_t gpu_buft, + ggml_backend_buffer_type_t host_buft, + struct llama_model_tensor_buft_override * tensor_buft_overrides, + llama_layer_fraction overflow_type, + llama_pshard_strategy strategy, + const pshard_dev_layout & layout, + bool pin_from_back = false, + bool output_on_gpu = false, + uint32_t n_attn_pinned = 0) { + thread_local std::array patterns_layer; + thread_local std::array patterns_layer_attn; + thread_local std::array patterns_layer_ffn; + thread_local std::string pat_output = "^output"; + + const uint32_t il_pin_start = pin_from_back ? (n_layers - n_pinned) : 0; + GGML_ASSERT(n_layers <= 1000 && "pshard: n_layers exceeds thread_local pattern array capacity"); + const uint32_t il_pin_end = pin_from_back ? n_layers : n_pinned; + const uint32_t il_boundary_raw = pin_from_back ? (il_pin_start > 0 ? il_pin_start - 1 : UINT32_MAX) : il_pin_end; + const uint32_t il_boundary = (overflow_type != LLAMA_LAYER_FRACTION_NONE && il_boundary_raw < n_layers) ? il_boundary_raw : UINT32_MAX; + const bool output_on_cpu = !output_on_gpu; + + static constexpr size_t OVERRIDE_CAP = 4096; + size_t itbo = 0; + + auto emit = [&](const char * pat, ggml_backend_buffer_type_t buft, int32_t bid) { + GGML_ASSERT(itbo + 1 < OVERRIDE_CAP && "override array overflow"); + tensor_buft_overrides[itbo] = { pat, buft, bid }; + itbo++; + }; + + { + thread_local std::string pat_tok_embd = "^token_embd"; + const int32_t out_bid = output_on_cpu ? layout.cpu : layout.compute; + emit(pat_output.c_str(), output_on_cpu ? host_buft : gpu_buft, out_bid); + emit(pat_tok_embd.c_str(), host_buft, layout.cpu); + } + + // one pattern cache per thread + for (uint32_t il = 0; il < n_layers; il++) { + if (patterns_layer[il].empty()) { patterns_layer[il] = "blk\\." + std::to_string(il) + "\\..*"; } + if (patterns_layer_attn[il].empty()) { patterns_layer_attn[il] = "blk\\." + std::to_string(il) + "\\.attn_(q|k|v|output|q_norm|k_norm).*"; } + if (patterns_layer_ffn[il].empty()) { patterns_layer_ffn[il] = "blk\\." + std::to_string(il) + "\\.ffn_((up|gate|down)\\.|(up|down|gate|gate_up)_(ch|)exps).*"; } + + if (il == il_boundary) { + const char * overflow_pat = llama_get_overflow_pattern(il, overflow_type); + if (overflow_pat) { + emit(overflow_pat, host_buft, layout.shard(il)); + } + emit(patterns_layer[il].c_str(), gpu_buft, layout.compute); + } else if (il >= il_pin_start && il < il_pin_end) { + emit(patterns_layer[il].c_str(), gpu_buft, layout.compute); + } else { + const bool use_alternating_shards = strategy == LLAMA_PSHARD_GPUONLY_LAYERPIN_LAYERSTREAM; + const int32_t shard_bid = use_alternating_shards ? layout.shard(il) : layout.shard_a; + + switch (strategy) { + case LLAMA_PSHARD_GPUONLY_LAYERPIN_LAYERSTREAM: + emit(patterns_layer[il].c_str(), host_buft, shard_bid); + break; + + case LLAMA_PSHARD_GPUONLY_ATTNPIN_FFNSTREAM: + emit(patterns_layer_ffn[il].c_str(), host_buft, shard_bid); + emit(patterns_layer[il].c_str(), gpu_buft, layout.compute); + break; + + case LLAMA_PSHARD_DYNAMIC_FFNCPU_ATTNSTREAM: + emit(patterns_layer_ffn[il].c_str(), host_buft, layout.cpu); + emit(patterns_layer[il].c_str(), host_buft, shard_bid); + break; + + case LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS: + if (n_attn_pinned > 0 && il < n_attn_pinned) { + emit(patterns_layer_ffn[il].c_str(), host_buft, layout.cpu); + emit(patterns_layer[il].c_str(), gpu_buft, layout.compute); + } else { + emit(patterns_layer[il].c_str(), host_buft, layout.cpu); + } + break; + + default: break; + } + } + } + tensor_buft_overrides[itbo] = { nullptr, nullptr, -1 }; + LLAMA_LOG_DEBUG("%s: %zu overrides emitted\n", __func__, itbo); +} + +// overflow names by enum value +static const char * const PSHARD_FRAC_NAMES[] = { "NONE", "ATTN", "UP", "GATE", "MOE" }; + +struct llama_pshard_search_ctx { + const char * path_model; + const struct llama_model_params * mparams; + const struct llama_context_params * cparams; + struct llama_model_tensor_buft_override * overrides; + uint32_t n_layers; + size_t vram_free; + ggml_backend_buffer_type_t gpu_buft; + ggml_backend_buffer_type_t host_buft; + pshard_dev_layout layout; + bool is_moe; + bool has_rs = false; + + // baseline params-fitting engine, injected by the caller (lives in common/fit.cpp) + llama_pshard_fit_fn fit_fn = nullptr; + void * fit_fn_ud = nullptr; + + // optional TPS predictor + const llama_benchmark_predictor * predictor = nullptr; + uint32_t kv_size = 0; + uint32_t cache_ubatch = 0; +}; + +struct llama_pshard_tps_hook_data { + const llama_benchmark_predictor * predictor; + int cpu_backend_id; + uint32_t kv_size; + int32_t batch_size; + uint32_t n_outputs; + bool has_rs; + float * out_tps; +}; + +static void pshard_tps_probe_hook(llama_context * ctx, void * user_data) { + auto * d = (llama_pshard_tps_hook_data *) user_data; + if (!d || !d->predictor || !ctx) return; + + double tps = d->predictor->predict_tps(ctx->get_sched(), d->cpu_backend_id, d->kv_size, d->batch_size, d->n_outputs, d->has_rs); + if (d->out_tps) { + *d->out_tps = (float)tps; + } +} + +static std::vector llama_pshard_probe_memory( + const llama_pshard_search_ctx & ctx, + const llama_model_params & mparams, + const llama_context_params & cparams, + ggml_log_level log_level, + llama_probe_hook_t probe_hook = nullptr, + void * probe_hook_data = nullptr) { + std::vector devs; + uint32_t hp_ngl = 0, hp_n_ctx_train = 0, hp_n_expert = 0, hp_n_embd_r = 0; + + const uint32_t probe_n_tokens = std::max(1, cparams.n_batch ? cparams.n_batch : cparams.n_ubatch); + const uint32_t probe_n_outputs = probe_n_tokens; + + llama_context_params cparams_probe = cparams; + if (ctx.cache_ubatch != 0) { + cparams_probe.n_batch = std::max(cparams_probe.n_batch, ctx.cache_ubatch); + cparams_probe.n_ubatch = ctx.cache_ubatch; + } + + return llama_get_device_memory_data_safe( + ctx.path_model, &mparams, &cparams_probe, devs, + hp_ngl, hp_n_ctx_train, hp_n_expert, hp_n_embd_r, + log_level, probe_hook, probe_hook_data, + probe_n_tokens, probe_n_outputs); +} + +struct llama_pshard_tier_prune { + uint32_t hi_pinned[LLAMA_PSHARD_COUNT]; + uint32_t hi_attn; + bool skip[LLAMA_PSHARD_COUNT]; + + void init(uint32_t n_layers) { + for (int s = 0; s < LLAMA_PSHARD_COUNT; s++) { + hi_pinned[s] = n_layers; + skip[s] = false; + } + hi_attn = n_layers; + } + + void update(int s, const llama_pshard_plan & plan) { + if (!plan.is_viable) { + return; + } + if (s == LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS) { + if (plan.n_attn_pinned > 0) { + hi_attn = plan.n_attn_pinned; + hi_pinned[s] = plan.n_pinned; + } + } else { + if (plan.n_pinned > 0) { + hi_pinned[s] = plan.n_pinned; + } + } + } +}; + +static llama_pshard_plan llama_pshard_search_strategy( + const llama_pshard_search_ctx & ctx, + llama_pshard_strategy strategy, + uint32_t hi_hint = UINT32_MAX, + uint32_t lo_hint = 0) { + + const auto * mparams = ctx.mparams; + const auto * cparams = ctx.cparams; + auto * tensor_buft_overrides = ctx.overrides; + const auto n_layers = ctx.n_layers; + const auto vram_free = ctx.vram_free; + const auto gpu_buft = ctx.gpu_buft; + const auto host_buft = ctx.host_buft; + const auto & layout = ctx.layout; + const auto is_moe = ctx.is_moe; + + llama_pshard_plan plan; + plan.strategy = strategy; + plan.batch_size = cparams->n_batch; + + const bool delegate_compute = llama_pshard_strategy_delegates_compute(strategy); + + const uint32_t hi_default = n_layers - 1; + uint32_t lo = lo_hint, hi = (hi_hint < hi_default) ? hi_hint : hi_default; + uint32_t best_n_pinned = lo_hint; + int64_t mem_lo = 0, mem_hi = (int64_t)vram_free * 2; + + // try the upper bound first + { + llama_pshard_generate_overrides_planning(hi, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, LLAMA_LAYER_FRACTION_NONE, strategy, layout, false, false); + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = delegate_compute; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + try { + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR); + const int64_t gpu_used = d[0].mb.total(); + LLAMA_LOG_INFO("%s: [%s] n_pinned=%u -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s (hi-first)\n", + __func__, llama_pshard_strategy_name(strategy), hi, + gpu_used / (1024.0 * 1024.0), + d[0].mb.model / (1024.0 * 1024.0), d[0].mb.context / (1024.0 * 1024.0), d[0].mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + if (gpu_used <= (int64_t)vram_free) { + best_n_pinned = hi; + lo = hi + 1; + } else { + mem_hi = gpu_used; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [%s] hi-first probe failed (n_pinned=%u)\n", __func__, llama_pshard_strategy_name(strategy), hi); + } + } + + while (lo <= hi) { + uint32_t mid; + if (mem_hi > mem_lo && mem_hi > (int64_t)vram_free) { + int64_t target = (int64_t)vram_free; + mid = lo + (uint32_t)((double)(target - mem_lo) * (hi - lo) / (mem_hi - mem_lo)); + if (mid <= lo) mid = lo + 1; + if (mid > hi) mid = hi; + } else { + mid = (lo + hi) / 2; + } + + llama_pshard_generate_overrides_planning(mid, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, LLAMA_LAYER_FRACTION_NONE, strategy, layout, false, false); + + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = delegate_compute; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + + try { + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR); + const int64_t gpu_used = d[0].mb.total(); + + LLAMA_LOG_INFO("%s: [%s] n_pinned=%u -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s\n", + __func__, llama_pshard_strategy_name(strategy), mid, + gpu_used / (1024.0 * 1024.0), + d[0].mb.model / (1024.0 * 1024.0), + d[0].mb.context / (1024.0 * 1024.0), + d[0].mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + + if (gpu_used <= (int64_t)vram_free) { + best_n_pinned = mid; + lo = mid + 1; + mem_lo = gpu_used; + } else { + if (mid == 0) break; + hi = mid - 1; + mem_hi = gpu_used; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [%s] probe failed (n_pinned=%u)\n", __func__, llama_pshard_strategy_name(strategy), mid); + if (mid == 0) break; + hi = mid - 1; + } + } + + llama_layer_fraction best_overflow = LLAMA_LAYER_FRACTION_NONE; + const uint32_t fallback_n_pinned = best_n_pinned; // known fit + const llama_layer_fraction fallback_overflow = LLAMA_LAYER_FRACTION_NONE; + if (best_n_pinned < n_layers - 1) { + const uint32_t frac_n_pinned = best_n_pinned + 1; // pin one more layer, partially + auto try_frac = [&](llama_layer_fraction frac) -> bool { + llama_pshard_generate_overrides_planning(frac_n_pinned, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, frac, strategy, layout, false, false); + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = delegate_compute; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + try { + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR); + return d[0].mb.total() <= (int64_t)vram_free; + } catch (...) { + LLAMA_LOG_WARN("%s: [%s] overflow probe failed (frac=%d)\n", __func__, llama_pshard_strategy_name(strategy), (int)frac); + return false; + } + }; + // try one partial boundary layer + if (try_frac(LLAMA_LAYER_FRACTION_ATTN)) { + best_n_pinned = frac_n_pinned; + best_overflow = LLAMA_LAYER_FRACTION_ATTN; + if (try_frac(LLAMA_LAYER_FRACTION_UP)) { best_overflow = LLAMA_LAYER_FRACTION_UP; + if (try_frac(LLAMA_LAYER_FRACTION_GATE)) { best_overflow = LLAMA_LAYER_FRACTION_GATE; + if (try_frac(LLAMA_LAYER_FRACTION_MOE)) { best_overflow = LLAMA_LAYER_FRACTION_MOE; }}} + } + } + + plan.n_pinned = best_n_pinned; + plan.overflow = best_overflow; + plan.output_on_gpu = false; + + llama_pshard_generate_overrides_planning(best_n_pinned, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, best_overflow, strategy, layout, false, plan.output_on_gpu); + { + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = delegate_compute; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + llama_pshard_tps_hook_data tps_data = { ctx.predictor, layout.cpu, ctx.kv_size, (int32_t)cparams->n_batch, cparams->n_seq_max, ctx.has_rs, &plan.tps }; + auto * hook = ctx.predictor ? pshard_tps_probe_hook : nullptr; + auto * hookdata = ctx.predictor ? (void *)&tps_data : nullptr; + + try { + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR, hook, hookdata); + plan.total_vram_req = d[0].mb.total(); + plan.scratch_measured = d[0].mb.compute; + plan.cache_measured = d[0].mb.context; + plan.is_viable = ((int64_t)plan.total_vram_req <= (int64_t)vram_free); + } catch (...) { + LLAMA_LOG_WARN("%s: [%s] final measurement probe failed (n_pinned=%u)\n", __func__, llama_pshard_strategy_name(strategy), best_n_pinned); + plan.is_viable = false; + } + } + + // drop the partial boundary layer if the final probe exceeds budget + if (!plan.is_viable && best_overflow != LLAMA_LAYER_FRACTION_NONE) { + best_n_pinned = fallback_n_pinned; + best_overflow = fallback_overflow; + plan.n_pinned = best_n_pinned; + plan.overflow = best_overflow; + llama_pshard_generate_overrides_planning(best_n_pinned, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, best_overflow, strategy, layout, false, plan.output_on_gpu); + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = delegate_compute; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + llama_pshard_tps_hook_data tps_data = { ctx.predictor, layout.cpu, ctx.kv_size, (int32_t)cparams->n_batch, cparams->n_seq_max, ctx.has_rs, &plan.tps }; + auto * hook = ctx.predictor ? pshard_tps_probe_hook : nullptr; + auto * hookdata = ctx.predictor ? (void *)&tps_data : nullptr; + try { + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR, hook, hookdata); + plan.total_vram_req = d[0].mb.total(); + plan.scratch_measured = d[0].mb.compute; + plan.cache_measured = d[0].mb.context; + plan.is_viable = ((int64_t)plan.total_vram_req <= (int64_t)vram_free); + } catch (...) { + plan.is_viable = false; + } + } + + for (const auto * ov = tensor_buft_overrides; ov->pattern; ++ov) { + plan.overrides.push_back({ov->pattern, ov->buft, ov->backend_id}); + } + + return plan; +} + +static llama_pshard_plan llama_pshard_search_attn_pin( + const llama_pshard_search_ctx & ctx, + uint32_t hi_attn_hint = UINT32_MAX, + uint32_t hi_full_hint = UINT32_MAX, + uint32_t lo_full_hint = 0) { + + const auto * mparams = ctx.mparams; + const auto * cparams = ctx.cparams; + auto * tensor_buft_overrides = ctx.overrides; + const auto n_layers = ctx.n_layers; + const auto vram_free = ctx.vram_free; + const auto gpu_buft = ctx.gpu_buft; + const auto host_buft = ctx.host_buft; + const auto & layout = ctx.layout; + const auto is_moe = ctx.is_moe; + + llama_pshard_plan plan; + plan.strategy = LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS; + + auto measure_vram = [&](uint32_t n_full, uint32_t n_attn, bool out_gpu) -> llama_memory_breakdown_data { + llama_pshard_generate_overrides_planning(n_full, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, LLAMA_LAYER_FRACTION_NONE, LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS, + layout, false, out_gpu, n_attn); + + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = true; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR); + return d[0].mb; + }; + + // phase 1: maximize attention layers on GPU + uint32_t n_attn = 0; + { + uint32_t lo = 0, hi = (hi_attn_hint < n_layers) ? hi_attn_hint : n_layers; + int64_t mem_lo = 0, mem_hi = (int64_t)vram_free * 2; + + try { + auto mb = measure_vram(0, hi, false); + int64_t gpu_used = mb.total(); + LLAMA_LOG_INFO("%s: [STATIC_ATTNPRIO_ALLMODELS p1] n_attn=%u -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s (hi-first)\n", + __func__, hi, gpu_used / (1024.0 * 1024.0), + mb.model / (1024.0 * 1024.0), mb.context / (1024.0 * 1024.0), mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + if (gpu_used <= (int64_t)vram_free) { + n_attn = hi; + lo = hi + 1; + } else { + mem_hi = gpu_used; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS p1] hi-first probe failed (n_attn=%u)\n", __func__, hi); + } + + while (lo <= hi) { + uint32_t mid; + if (mem_hi > mem_lo && mem_hi > (int64_t)vram_free) { + mid = lo + (uint32_t)((double)((int64_t)vram_free - mem_lo) * (hi - lo) / (mem_hi - mem_lo)); + if (mid <= lo) mid = lo + 1; + if (mid > hi) mid = hi; + } else { + mid = (lo + hi) / 2; + } + + try { + auto mb = measure_vram(0, mid, false); + int64_t gpu_used = mb.total(); + LLAMA_LOG_INFO("%s: [STATIC_ATTNPRIO_ALLMODELS p1] n_attn=%u -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s\n", + __func__, mid, gpu_used / (1024.0 * 1024.0), + mb.model / (1024.0 * 1024.0), mb.context / (1024.0 * 1024.0), mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + + if (gpu_used <= (int64_t)vram_free) { + n_attn = mid; + lo = mid + 1; + mem_lo = gpu_used; + } else { + if (mid == 0) break; + hi = mid - 1; + mem_hi = gpu_used; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS p1] probe failed (n_attn=%u)\n", __func__, mid); + if (mid == 0) break; + hi = mid - 1; + } + } + } + + // moe tries output on gpu before ffn pinning + // dense tries output on gpu after ffn pinning + bool output_on_gpu = false; + + if (is_moe && n_attn >= n_layers) { + try { + auto mb = measure_vram(0, n_attn, true); + int64_t gpu_used = mb.total(); + LLAMA_LOG_INFO("%s: [STATIC_ATTNPRIO_ALLMODELS p1b] output_on_gpu probe -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s\n", + __func__, gpu_used / (1024.0 * 1024.0), + mb.model / (1024.0 * 1024.0), mb.context / (1024.0 * 1024.0), mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + if (gpu_used <= (int64_t)vram_free) { + output_on_gpu = true; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS p1b] output_on_gpu probe failed\n", __func__); + } + } + + // phase 2: maximize fully pinned layers + uint32_t n_full = 0; + { + uint32_t hi_full_max = (hi_full_hint < n_attn) ? hi_full_hint : n_attn; + uint32_t lo = lo_full_hint, hi = hi_full_max; + int64_t mem_lo = 0, mem_hi = (int64_t)vram_free * 2; + + try { + auto mb = measure_vram(hi, n_attn, output_on_gpu); + int64_t gpu_used = mb.total(); + LLAMA_LOG_INFO("%s: [STATIC_ATTNPRIO_ALLMODELS p2] n_full=%u n_attn=%u -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s (hi-first)\n", + __func__, hi, n_attn, gpu_used / (1024.0 * 1024.0), + mb.model / (1024.0 * 1024.0), mb.context / (1024.0 * 1024.0), mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + if (gpu_used <= (int64_t)vram_free) { + n_full = hi; + lo = hi + 1; + } else { + mem_hi = gpu_used; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS p2] hi-first probe failed (n_full=%u, n_attn=%u)\n", __func__, hi, n_attn); + } + + while (lo <= hi) { + uint32_t mid; + if (mem_hi > mem_lo && mem_hi > (int64_t)vram_free) { + mid = lo + (uint32_t)((double)((int64_t)vram_free - mem_lo) * (hi - lo) / (mem_hi - mem_lo)); + if (mid <= lo) mid = lo + 1; + if (mid > hi) mid = hi; + } else { + mid = (lo + hi) / 2; + } + + try { + auto mb = measure_vram(mid, n_attn, output_on_gpu); + int64_t gpu_used = mb.total(); + LLAMA_LOG_INFO("%s: [STATIC_ATTNPRIO_ALLMODELS p2] n_full=%u n_attn=%u -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s\n", + __func__, mid, n_attn, gpu_used / (1024.0 * 1024.0), + mb.model / (1024.0 * 1024.0), mb.context / (1024.0 * 1024.0), mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + + if (gpu_used <= (int64_t)vram_free) { + n_full = mid; + lo = mid + 1; + mem_lo = gpu_used; + } else { + if (mid == 0) break; + hi = mid - 1; + mem_hi = gpu_used; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS p2] probe failed (n_full=%u, n_attn=%u)\n", __func__, mid, n_attn); + if (mid == 0) break; + hi = mid - 1; + } + } + } + + // dense tries output on gpu only after all layers fit + if (!is_moe && !output_on_gpu && n_full >= n_layers) { + try { + auto mb = measure_vram(n_full, n_attn, true); + int64_t gpu_used = mb.total(); + LLAMA_LOG_INFO("%s: [STATIC_ATTNPRIO_ALLMODELS p3] output_on_gpu probe (n_full=%u) -> %.1f MiB (model=%.1f cache=%.1f compute=%.1f) budget %.1f %s\n", + __func__, n_full, gpu_used / (1024.0 * 1024.0), + mb.model / (1024.0 * 1024.0), mb.context / (1024.0 * 1024.0), mb.compute / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + gpu_used <= (int64_t)vram_free ? "FITS" : "OVER"); + if (gpu_used <= (int64_t)vram_free) { + output_on_gpu = true; + } + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS p3] output_on_gpu probe failed\n", __func__); + } + } + + plan.n_pinned = n_full; + plan.n_attn_pinned = n_attn; + plan.output_on_gpu = output_on_gpu; + + llama_pshard_generate_overrides_planning(n_full, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, LLAMA_LAYER_FRACTION_NONE, LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS, + layout, false, plan.output_on_gpu, n_attn); + { + llama_model_params mp = *mparams; + mp.pshard = true; + mp.pshard_delegate_compute = true; + mp.n_gpu_layers = n_layers + 1; + mp.tensor_buft_overrides = tensor_buft_overrides; + llama_pshard_tps_hook_data tps_data = { ctx.predictor, layout.cpu, ctx.kv_size, (int32_t)cparams->n_batch, cparams->n_seq_max, ctx.has_rs, &plan.tps }; + auto * hook = ctx.predictor ? pshard_tps_probe_hook : nullptr; + auto * hookdata = ctx.predictor ? (void *)&tps_data : nullptr; + + try { + const auto d = llama_pshard_probe_memory(ctx, mp, *cparams, GGML_LOG_LEVEL_ERROR, hook, hookdata); + plan.total_vram_req = d[0].mb.total(); + plan.scratch_measured = d[0].mb.compute; + plan.cache_measured = d[0].mb.context; + plan.is_viable = ((int64_t)plan.total_vram_req <= (int64_t)vram_free); + } catch (...) { + LLAMA_LOG_WARN("%s: [STATIC_ATTNPRIO_ALLMODELS] final measurement probe failed (n_full=%u, n_attn=%u)\n", __func__, n_full, n_attn); + plan.is_viable = false; + } + } + + for (const auto * ov = tensor_buft_overrides; ov->pattern; ++ov) { + plan.overrides.push_back({ov->pattern, ov->buft, ov->backend_id}); + } + + return plan; +} + + +// plan cache serialization + + +static std::string pshard_plan_to_ot(const llama_pshard_plan & plan, ggml_backend_buffer_type_t host_buft) { + std::string ot; + const char * buft_name = ggml_backend_buft_name(host_buft); + for (size_t i = 0; i < plan.overrides.size(); i++) { + if (i > 0) ot += ','; + ot += plan.overrides[i].pattern; + ot += '='; + ot += buft_name; + ot += ':'; + ot += std::to_string(plan.overrides[i].backend_id); + } + return ot; +} + +static const char * pshard_overflow_name(int overflow) { + static const char * const names[] = { "NONE", "ATTN", "UP", "GATE", "MOE" }; + return (overflow >= 0 && overflow < 5) ? names[overflow] : "NONE"; +} + +static int pshard_overflow_from_name(const char * name) { + if (!name) return 0; + if (strcmp(name, "ATTN") == 0) return 1; + if (strcmp(name, "UP") == 0) return 2; + if (strcmp(name, "GATE") == 0) return 3; + if (strcmp(name, "MOE") == 0) return 4; + return 0; +} + +static const size_t PSHARD_MIB = 1024ULL * 1024ULL; +static const int PSHARD_CACHE_MAX_SECTIONS = 32; +static const int PSHARD_CACHE_MAX_VARIANTS = 16; + +static uint32_t pshard_bytes_to_mib_ceil(size_t bytes) { + return (uint32_t)((bytes + PSHARD_MIB - 1) / PSHARD_MIB); +} + +static size_t pshard_mib_to_bytes(uint32_t mib) { + return (size_t)mib * PSHARD_MIB; +} + +static size_t pshard_mib_to_bytes(double mib) { + return (size_t)(mib * (double)PSHARD_MIB + 0.5); +} + +static bool pshard_parse_variant_header(const std::string & line, uint32_t & budget_mib, uint32_t & cache_ubatch) { + cache_ubatch = 0; + if (sscanf(line.c_str(), "[variant budget=%u cache_ubatch=%u]", &budget_mib, &cache_ubatch) == 2) { + return true; + } + return sscanf(line.c_str(), "[variant budget=%u]", &budget_mib) == 1; +} + +static bool pshard_plan_is_better(const llama_pshard_plan & candidate, const llama_pshard_plan & current); + +bool pshard_registry_save( + const llama_pshard_plan_registry * registry, uint64_t fingerprint, + const char * cache_path, ggml_backend_buffer_type_t host_buft, + const llama_context_params * cparams) { + if (!registry || !cache_path) return false; + + struct cache_section { + std::string header; + std::vector lines; + }; + std::vector sections; + + FILE * existing = fopen(cache_path, "r"); + if (existing) { + char line[8192]; + cache_section * cur = nullptr; + while (fgets(line, sizeof(line), existing)) { + std::string s = line; + while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) s.pop_back(); + if (s.compare(0, 13, "[fingerprint=") == 0) { + sections.push_back({s, {}}); + cur = §ions.back(); + } else if (cur) { + cur->lines.push_back(s); + } + } + fclose(existing); + } + + std::vector preserved_sections; + std::vector> preserved_variants; + + size_t inferred_budget = registry->pshard_disabled ? registry->baseline_vram_req : 0; + if (!registry->pshard_disabled) { + for (const auto & plan : registry->best_plans) { + if (plan.is_viable) { + inferred_budget = std::max(inferred_budget, plan.total_vram_req); + } + } + } + const uint32_t budget_mib = registry->budget_mib + ? registry->budget_mib + : pshard_bytes_to_mib_ceil(inferred_budget); + const uint32_t cache_ubatch = registry->cache_ubatch + ? registry->cache_ubatch + : (registry->tier_sizes.empty() ? 0 : registry->tier_sizes.back()); + + for (const auto & sec : sections) { + uint64_t fp = 0; + sscanf(sec.header.c_str(), "[fingerprint=0x%" SCNx64, &fp); + if (fp != fingerprint) { + preserved_sections.push_back(sec); + continue; + } + + std::vector cur_variant; + uint32_t cur_budget = 0; + uint32_t cur_cache_ubatch = 0; + bool in_variant = false; + auto flush_variant = [&]() { + if (in_variant && (cur_budget != budget_mib || cur_cache_ubatch != cache_ubatch)) { + preserved_variants.push_back(cur_variant); + } + cur_variant.clear(); + cur_budget = 0; + cur_cache_ubatch = 0; + in_variant = false; + }; + + for (const auto & ln : sec.lines) { + uint32_t parsed_budget = 0; + uint32_t parsed_cache_ubatch = 0; + if (pshard_parse_variant_header(ln, parsed_budget, parsed_cache_ubatch)) { + flush_variant(); + in_variant = true; + cur_budget = parsed_budget; + cur_cache_ubatch = parsed_cache_ubatch; + cur_variant.push_back(ln); + } else if (in_variant) { + cur_variant.push_back(ln); + } + } + flush_variant(); + } + + while ((int)preserved_sections.size() >= PSHARD_CACHE_MAX_SECTIONS) { + preserved_sections.erase(preserved_sections.begin()); + } + while ((int)preserved_variants.size() >= PSHARD_CACHE_MAX_VARIANTS) { + preserved_variants.erase(preserved_variants.begin()); + } + + FILE * f = fopen(cache_path, "w"); + if (!f) { + LLAMA_LOG_WARN("%s: could not write plan cache: %s\n", __func__, cache_path); + return false; + } + + fprintf(f, "# Generated file. Edit at your own risk.\n"); + + for (const auto & sec : preserved_sections) { + fprintf(f, "\n%s\n", sec.header.c_str()); + for (const auto & ln : sec.lines) { + fprintf(f, "%s\n", ln.c_str()); + } + } + + fprintf(f, "\n[fingerprint=0x%016" PRIx64 "]\n", fingerprint); + if (cparams) { + const char * fa_str = "unknown"; + switch (cparams->flash_attn_type) { + case LLAMA_FLASH_ATTN_TYPE_DISABLED: fa_str = "off"; break; + case LLAMA_FLASH_ATTN_TYPE_ENABLED: fa_str = "on"; break; + case LLAMA_FLASH_ATTN_TYPE_AUTO: fa_str = "auto"; break; + } + const int forced_strategy = pshard_strategy_from_env(); + fprintf(f, "# n_ctx=%u n_seq_max=%u n_threads=%d fa=%s type_k=%d type_v=%d strategy=%s\n", + cparams->n_ctx, cparams->n_seq_max, cparams->n_threads, + fa_str, (int)cparams->type_k, (int)cparams->type_v, + forced_strategy >= 0 ? llama_pshard_strategy_name((llama_pshard_strategy)forced_strategy) : "auto"); + } + + for (const auto & variant : preserved_variants) { + fprintf(f, "\n"); + for (const auto & ln : variant) { + fprintf(f, "%s\n", ln.c_str()); + } + } + + fprintf(f, "\n[variant budget=%u cache_ubatch=%u]\n", budget_mib, cache_ubatch); + if (registry->pshard_disabled) { + fprintf(f, "pshard_disabled=1 baseline_vram=%.1f\n", registry->baseline_vram_req / (1024.0 * 1024.0)); + } else { + for (size_t t = 0; t < registry->tier_sizes.size(); t++) { + const auto & plan = registry->best_plans[t]; + if (!plan.is_viable) { + fprintf(f, "[tier %zu bs=%u] not_viable\n", t, registry->tier_sizes[t]); + continue; + } + fprintf(f, "[tier %zu bs=%u]\n", t, registry->tier_sizes[t]); + fprintf(f, "strategy=%s n_pinned=%u n_attn_pinned=%u overflow=%s tps=%.2f vram=%.1f output_on_gpu=%d pin_from_back=%d\n", + llama_pshard_strategy_name(plan.strategy), + plan.n_pinned, plan.n_attn_pinned, + pshard_overflow_name(plan.overflow), + plan.tps, plan.total_vram_req / (1024.0 * 1024.0), + (int)plan.output_on_gpu, (int)plan.pin_from_back); + fprintf(f, "ot=%s\n", pshard_plan_to_ot(plan, host_buft).c_str()); + } + } + + fclose(f); + LLAMA_LOG_INFO("%s: saved budget=%u MiB cache_ubatch=%u variant with %zu tier plans to %s\n", + __func__, budget_mib, cache_ubatch, registry->tier_sizes.size(), cache_path); + return true; +} + +bool pshard_registry_load_planning( + llama_pshard_plan_registry * registry, uint64_t fingerprint, + const char * cache_path, ggml_backend_buffer_type_t host_buft, + size_t current_budget, bool require_exact_budget) { + if (!registry || !cache_path) return false; + + FILE * f = fopen(cache_path, "r"); + if (!f) return false; + + char line[8192]; + bool in_section = false; + + char fp_header[64]; + snprintf(fp_header, sizeof(fp_header), "[fingerprint=0x%016" PRIx64 "]", fingerprint); + + struct tier_data { + uint32_t bs = 0; + bool viable = false; + llama_pshard_strategy strategy = LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS; + uint32_t n_pinned = 0; + uint32_t n_attn_pinned = 0; + int overflow = 0; + float tps = 0.0f; + double vram_mib = 0.0; + int output_on_gpu = 0; + int pin_from_back = 0; + std::string ot_line; + }; + struct variant_data { + uint32_t budget_mib = 0; + uint32_t cache_ubatch = 0; + bool pshard_disabled = false; + double baseline_vram_mib = 0.0; + std::vector tiers; + }; + std::vector variants; + variant_data * cur_variant = nullptr; + + while (fgets(line, sizeof(line), f)) { + std::string s = line; + while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) s.pop_back(); + + if (s.compare(0, 13, "[fingerprint=") == 0) { + if (in_section) break; // hit next section, stop + in_section = (s == fp_header); + continue; + } + if (!in_section) continue; + + uint32_t variant_budget = 0; + uint32_t variant_cache_ubatch = 0; + if (pshard_parse_variant_header(s, variant_budget, variant_cache_ubatch)) { + variants.push_back({}); + cur_variant = &variants.back(); + cur_variant->budget_mib = variant_budget; + cur_variant->cache_ubatch = variant_cache_ubatch; + continue; + } + if (!cur_variant) continue; + + if (s.rfind("pshard_disabled=1", 0) == 0) { + double baseline_mib = 0.0; + if (sscanf(s.c_str(), "pshard_disabled=1 baseline_vram=%lf", &baseline_mib) == 1) { + cur_variant->pshard_disabled = true; + cur_variant->baseline_vram_mib = baseline_mib; + } else { + LLAMA_LOG_WARN("%s: malformed pshard_disabled line: %s\n", __func__, s.c_str()); + } + continue; + } + + if (s.compare(0, 5, "[tier") == 0) { + tier_data td = {}; + size_t tier_idx = 0; + if (s.find("not_viable") != std::string::npos) { + if (sscanf(s.c_str(), "[tier %zu bs=%u]", &tier_idx, &td.bs) < 2) { + LLAMA_LOG_WARN("%s: malformed tier header (not_viable): %s\n", __func__, s.c_str()); + continue; + } + td.viable = false; + } else { + if (sscanf(s.c_str(), "[tier %zu bs=%u]", &tier_idx, &td.bs) < 2) { + LLAMA_LOG_WARN("%s: malformed tier header: %s\n", __func__, s.c_str()); + continue; + } + td.viable = true; + } + cur_variant->tiers.push_back(td); + } else if (s.compare(0, 9, "strategy=") == 0 && !cur_variant->tiers.empty()) { + auto & td = cur_variant->tiers.back(); + char strat_name[64] = {}, overflow_name[16] = {}; + if (sscanf(s.c_str(), "strategy=%63s n_pinned=%u n_attn_pinned=%u overflow=%15s tps=%f vram=%lf", + strat_name, &td.n_pinned, &td.n_attn_pinned, overflow_name, &td.tps, &td.vram_mib) < 4) { + LLAMA_LOG_WARN("%s: malformed strategy line: %s\n", __func__, s.c_str()); + td.viable = false; + continue; + } + + const char * ogg = strstr(s.c_str(), "output_on_gpu="); + const char * pfb = strstr(s.c_str(), "pin_from_back="); + if (!ogg || !pfb) { + LLAMA_LOG_WARN("%s: missing output_on_gpu/pin_from_back, invalidating cache: %s\n", __func__, s.c_str()); + td.viable = false; + continue; + } + td.output_on_gpu = atoi(ogg + 14); + td.pin_from_back = atoi(pfb + 14); + + td.overflow = pshard_overflow_from_name(overflow_name); + bool found_strategy = false; + for (int i = 0; i < LLAMA_PSHARD_COUNT; i++) { + if (strcmp(strat_name, llama_pshard_strategy_name((llama_pshard_strategy)i)) == 0) { + td.strategy = (llama_pshard_strategy)i; + found_strategy = true; + break; + } + } + if (!found_strategy) { + LLAMA_LOG_WARN("%s: unknown strategy in cache, invalidating tier: %s\n", __func__, strat_name); + td.viable = false; + } + } else if (s.compare(0, 3, "ot=") == 0 && !cur_variant->tiers.empty()) { + cur_variant->tiers.back().ot_line = s.substr(3); + } + } + fclose(f); + + if (variants.empty() || !in_section) return false; + + auto make_plan = [&](const tier_data & td) { + llama_pshard_plan plan; + plan.strategy = td.strategy; + plan.batch_size = td.bs; + plan.n_pinned = td.n_pinned; + plan.n_attn_pinned = td.n_attn_pinned; + plan.overflow = td.overflow; + plan.tps = td.tps; + plan.total_vram_req = (size_t)(td.vram_mib * 1024 * 1024); + plan.is_viable = td.viable; + plan.output_on_gpu = (bool)td.output_on_gpu; + plan.pin_from_back = (bool)td.pin_from_back; + + if (!td.ot_line.empty()) { + std::string remaining = td.ot_line; + while (!remaining.empty()) { + size_t comma = remaining.find(','); + std::string token = (comma != std::string::npos) ? remaining.substr(0, comma) : remaining; + remaining = (comma != std::string::npos) ? remaining.substr(comma + 1) : ""; + + size_t eq = token.find('='); + if (eq == std::string::npos) continue; + + std::string pattern = token.substr(0, eq); + std::string buft_bid = token.substr(eq + 1); + + int32_t backend_id = -1; + size_t colon = buft_bid.rfind(':'); + if (colon != std::string::npos) { + backend_id = atoi(buft_bid.c_str() + colon + 1); + } + + plan.overrides.push_back({ pattern, host_buft, backend_id }); + } + } + return plan; + }; + + const uint32_t current_budget_mib = pshard_bytes_to_mib_ceil(current_budget); + const uint32_t requested_cache_ubatch = registry->cache_ubatch; + bool skipped_cache_ubatch = false; + auto variant_cache_ubatch = [&](const variant_data & variant) { + return variant.cache_ubatch ? variant.cache_ubatch : requested_cache_ubatch; + }; + auto cache_ubatch_ok = [&](const variant_data & variant) { + if (requested_cache_ubatch == 0) return true; + if (variant.cache_ubatch == 0) return true; + if (variant.cache_ubatch <= requested_cache_ubatch) return true; + skipped_cache_ubatch = true; + return false; + }; + + for (const auto & variant : variants) { + if (!variant.pshard_disabled) continue; + if (!cache_ubatch_ok(variant)) continue; + const size_t baseline_vram = pshard_mib_to_bytes(variant.baseline_vram_mib); + if (baseline_vram <= current_budget) { + registry->tier_sizes.clear(); + registry->best_plans.clear(); + registry->pshard_disabled = true; + registry->baseline_vram_req = baseline_vram; + registry->budget_mib = variant.budget_mib; + registry->cache_ubatch = variant_cache_ubatch(variant); + LLAMA_LOG_INFO("%s: loaded pshard_disabled variant budget=%u MiB cache_ubatch=%u baseline=%.1f MiB from %s\n", + __func__, variant.budget_mib, registry->cache_ubatch, variant.baseline_vram_mib, cache_path); + return true; + } + } + + const variant_data * best_whole = nullptr; + for (const auto & variant : variants) { + if (variant.pshard_disabled || variant.tiers.empty()) continue; + if (!cache_ubatch_ok(variant)) continue; + if (require_exact_budget) { + if (variant.budget_mib != current_budget_mib) continue; + } else if (pshard_mib_to_bytes(variant.budget_mib) > current_budget) { + continue; + } + if (!best_whole || + variant_cache_ubatch(variant) > variant_cache_ubatch(*best_whole) || + (variant_cache_ubatch(variant) == variant_cache_ubatch(*best_whole) && + variant.budget_mib > best_whole->budget_mib)) { + best_whole = &variant; + } + } + + std::vector> selected; + auto add_or_fill_plan = [&](const tier_data & td, bool allow_existing) { + llama_pshard_plan plan = make_plan(td); + auto it = std::find_if(selected.begin(), selected.end(), + [&](const auto & p) { return p.first == td.bs; }); + if (it == selected.end()) { + selected.push_back({td.bs, std::move(plan)}); + } else if (allow_existing || !it->second.is_viable || + (!best_whole && plan.is_viable && pshard_plan_is_better(plan, it->second))) { + it->second = std::move(plan); + } + }; + + if (best_whole) { + for (const auto & td : best_whole->tiers) { + add_or_fill_plan(td, true); + } + } + + if (!require_exact_budget) { + for (const auto & variant : variants) { + if (variant.pshard_disabled || pshard_mib_to_bytes(variant.budget_mib) <= current_budget) continue; + if (!cache_ubatch_ok(variant)) continue; + for (const auto & td : variant.tiers) { + if (!td.viable || td.ot_line.empty()) continue; + if (pshard_mib_to_bytes(td.vram_mib) > current_budget) continue; + add_or_fill_plan(td, false); + } + } + } + + const uint32_t selected_cache_ubatch = best_whole ? variant_cache_ubatch(*best_whole) : requested_cache_ubatch; + selected.erase(std::remove_if(selected.begin(), selected.end(), + [&](const auto & p) { + if (p.first == 0) return true; + return selected_cache_ubatch > 0 && p.first > selected_cache_ubatch; + }), selected.end()); + if (selected.empty()) { + if (require_exact_budget && !variants.empty()) { + LLAMA_LOG_INFO("%s: cache miss, no exact budget=%u MiB variant in %s\n", + __func__, current_budget_mib, cache_path); + } + if (skipped_cache_ubatch) { + LLAMA_LOG_INFO("%s: cache miss, no variant with cache_ubatch <= target cache_ubatch=%u in %s\n", + __func__, requested_cache_ubatch, cache_path); + } + return false; + } + + std::sort(selected.begin(), selected.end(), + [](const auto & a, const auto & b) { return a.first < b.first; }); + + registry->tier_sizes.clear(); + registry->best_plans.clear(); + registry->pshard_disabled = false; + registry->baseline_vram_req = 0; + registry->budget_mib = best_whole ? best_whole->budget_mib : 0; + registry->cache_ubatch = selected_cache_ubatch; + + for (auto & item : selected) { + const auto & p = item.second; + if (p.is_viable && p.overrides.empty()) { + LLAMA_LOG_WARN("%s: plan cache corrupt: tier bs=%u viable but has no overrides\n", __func__, item.first); + registry->tier_sizes.clear(); + registry->best_plans.clear(); + return false; + } + registry->tier_sizes.push_back(item.first); + registry->best_plans.push_back(std::move(item.second)); + } + + if (require_exact_budget) { + LLAMA_LOG_INFO("%s: loaded %zu tier plans from exact budget=%u MiB cache_ubatch=%u variant in %s\n", + __func__, registry->tier_sizes.size(), registry->budget_mib, registry->cache_ubatch, cache_path); + } else if (best_whole) { + LLAMA_LOG_INFO("%s: loaded %zu tier plans from budget=%u MiB cache_ubatch=%u variant for current budget=%u MiB in %s\n", + __func__, registry->tier_sizes.size(), registry->budget_mib, registry->cache_ubatch, current_budget_mib, cache_path); + } else { + LLAMA_LOG_INFO("%s: loaded %zu salvaged tier plans for current budget=%u MiB from %s\n", + __func__, registry->tier_sizes.size(), current_budget_mib, cache_path); + } + return true; +} + +static bool pshard_plan_is_better(const llama_pshard_plan & candidate, const llama_pshard_plan & current) { + if (!current.is_viable) return true; + const bool candidate_has_tps = candidate.tps > 0.0f; + const bool current_has_tps = current.tps > 0.0f; + if (candidate_has_tps || current_has_tps) { + if (candidate_has_tps != current_has_tps) return candidate_has_tps; + if (candidate.tps != current.tps) return candidate.tps > current.tps; + } + if (candidate.n_pinned != current.n_pinned) return candidate.n_pinned > current.n_pinned; + if (candidate.n_attn_pinned != current.n_attn_pinned) return candidate.n_attn_pinned > current.n_attn_pinned; + if (candidate.overflow != current.overflow) return candidate.overflow > current.overflow; + return candidate.total_vram_req < current.total_vram_req; +} + +// true when every layer runs on the compute backend with no per layer override +// token_embd on host is allowed +// callers use this to skip pshard when baseline already fits +static bool pshard_plan_is_baseline_fit( + const llama_pshard_plan & plan, + uint32_t n_layers, + int32_t compute_bid) { + if (!plan.is_viable || plan.n_pinned < n_layers) return false; + for (const auto & ov : plan.overrides) { + // token_embd routes to CPU host buffer + if (ov.pattern.find("token_embd") != std::string::npos) continue; + if (ov.backend_id != compute_bid) return false; + } + return true; +} + +static llama_pshard_plan llama_pshard_search_baseline_fit_tier( + const llama_pshard_search_ctx & ctx, + const std::vector & dmds) { + const auto * path_model = ctx.path_model; + const auto * mparams = ctx.mparams; + const auto * cparams = ctx.cparams; + const auto n_layers = ctx.n_layers; + const auto vram_free = ctx.vram_free; + const auto gpu_buft = ctx.gpu_buft; + const auto host_buft = ctx.host_buft; + const auto & layout = ctx.layout; + + llama_pshard_plan plan; + plan.batch_size = cparams->n_batch; + + llama_model_tensor_buft_override local_overrides[4096]; + local_overrides[0] = { nullptr, nullptr, -1 }; + + llama_model_params mp_copy = *mparams; + mp_copy.pshard = false; + mp_copy.pshard_delegate_compute = false; + mp_copy.tensor_buft_overrides = nullptr; + mp_copy.max_vram_alloc = std::max(1, pshard_bytes_to_mib_ceil(vram_free)); + + llama_context_params cp_copy = *cparams; + cp_copy.pshard = false; + + float ts[16] = {}; + size_t margins[16] = {}; + + try { + if (!ctx.fit_fn) { + throw std::runtime_error("pshard: no baseline fit function injected"); + } + ctx.fit_fn(path_model, &mp_copy, &cp_copy, ts, local_overrides, margins, 0, + GGML_LOG_LEVEL_ERROR, ctx.fit_fn_ud); + + const uint32_t ngl = (mp_copy.n_gpu_layers < 0) + ? (n_layers + 1) + : std::min((uint32_t)mp_copy.n_gpu_layers, n_layers + 1); + plan.n_pinned = (ngl > 0) ? (ngl - 1) : 0; + plan.pin_from_back = false; + plan.output_on_gpu = (ngl > 0); + + for (auto * ov = local_overrides; ov->pattern; ++ov) { + if (ov->buft == gpu_buft) { + ov->backend_id = layout.compute; + } else { + ov->buft = host_buft; + ov->backend_id = layout.cpu; + } + plan.overrides.push_back({ov->pattern, ov->buft, ov->backend_id}); + } + + mp_copy.tensor_buft_overrides = local_overrides[0].pattern ? local_overrides : nullptr; + const auto d = llama_pshard_probe_memory(ctx, mp_copy, cp_copy, GGML_LOG_LEVEL_ERROR); + plan.total_vram_req = d[0].mb.total(); + plan.scratch_measured = d[0].mb.compute; + plan.cache_measured = d[0].mb.context; + plan.is_viable = ((int64_t)plan.total_vram_req <= (int64_t)vram_free); + } catch (const std::exception & e) { + LLAMA_LOG_WARN("%s: baseline probe failed (bs=%u): %s\n", + __func__, cparams->n_batch, e.what()); + plan.is_viable = false; + } catch (...) { + LLAMA_LOG_WARN("%s: baseline probe failed (bs=%u): unknown exception\n", + __func__, cparams->n_batch); + plan.is_viable = false; + } + + LLAMA_LOG_INFO("%s: [bs=%-4u baseline] n_pinned=%2u/%2u, vram=%7.1f MiB, %s%s\n", + __func__, cparams->n_batch, plan.n_pinned, n_layers, + plan.total_vram_req / (1024.0 * 1024.0), + plan.is_viable ? "VIABLE" : "NOT VIABLE", + pshard_plan_is_baseline_fit(plan, n_layers, layout.compute) ? " (full fit)" : ""); + + return plan; +} + +// use attention priority when a forced strategy cannot fit the tier +// caller must set ctx.cparams for the target tier +static llama_pshard_plan llama_pshard_attn_pin_fallback( + const llama_pshard_search_ctx & ctx, + int force_strategy, + uint32_t hi_attn = UINT32_MAX, + uint32_t hi_pinned = UINT32_MAX) { + llama_pshard_plan fallback = llama_pshard_search_attn_pin(ctx, hi_attn, hi_pinned); + fallback.batch_size = ctx.cparams->n_batch; + LLAMA_LOG_INFO("llama_params_fit_pshard_planning: [bs=%-4u %-10s] forced %s non-viable, STATIC_ATTNPRIO_ALLMODELS fallback: n_pinned=%2u/%2u, %s\n", + ctx.cparams->n_batch, llama_pshard_strategy_name(LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS), + llama_pshard_strategy_name((llama_pshard_strategy)force_strategy), + fallback.n_pinned, fallback.n_attn_pinned, + fallback.is_viable ? "VIABLE" : "NOT VIABLE"); + return fallback; +} + +static llama_pshard_plan llama_pshard_search_tier( + const llama_pshard_search_ctx & ctx, + int force_strategy, + const std::vector & dmds, + llama_pshard_tier_prune & prune) { + + const auto * path_model = ctx.path_model; + const auto * mparams = ctx.mparams; + const auto * cparams = ctx.cparams; + auto * tensor_buft_overrides = ctx.overrides; + const auto n_layers = ctx.n_layers; + const auto vram_free = ctx.vram_free; + const auto gpu_buft = ctx.gpu_buft; + const auto host_buft = ctx.host_buft; + const auto & layout = ctx.layout; + const auto is_moe = ctx.is_moe; + + llama_pshard_plan best; + + for (int s = 0; s < LLAMA_PSHARD_COUNT; s++) { + if (force_strategy >= 0 && force_strategy != s) continue; + if (prune.skip[s]) continue; + + llama_pshard_strategy strategy = (llama_pshard_strategy)s; + llama_pshard_plan plan; + + if (strategy == LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS) { + plan = llama_pshard_search_attn_pin(ctx, prune.hi_attn, prune.hi_pinned[s]); + } else { + plan = llama_pshard_search_strategy(ctx, strategy, prune.hi_pinned[s]); + } + + plan.batch_size = cparams->n_batch; + + { + const char * status = plan.is_viable ? "VIABLE" : "NOT VIABLE"; + char tps_buf[32] = ""; + if (plan.tps > 0.0f) { snprintf(tps_buf, sizeof(tps_buf), ", tps=%.1f", plan.tps); } + + if (plan.n_attn_pinned > 0) { + LLAMA_LOG_INFO("%s: [bs=%-4u %-10s] n_pinned=%2u (attn=%2u), overflow=%-4s, vram=%7.1f MiB, %s%s\n", + __func__, cparams->n_batch, llama_pshard_strategy_name(strategy), + plan.n_pinned, plan.n_attn_pinned, PSHARD_FRAC_NAMES[plan.overflow], + plan.total_vram_req / (1024.0 * 1024.0), status, tps_buf); + } else { + LLAMA_LOG_INFO("%s: [bs=%-4u %-10s] n_pinned=%2u, overflow=%-4s, vram=%7.1f MiB, %s%s\n", + __func__, cparams->n_batch, llama_pshard_strategy_name(strategy), + plan.n_pinned, PSHARD_FRAC_NAMES[plan.overflow], + plan.total_vram_req / (1024.0 * 1024.0), status, tps_buf); + } + } + + prune.update(s, plan); + + if (plan.is_viable && pshard_plan_is_better(plan, best)) { + best = plan; + } + } + + if (!best.is_viable && force_strategy >= 0 && force_strategy != LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS) { + llama_pshard_plan fallback = llama_pshard_attn_pin_fallback( + ctx, force_strategy, prune.hi_attn, prune.hi_pinned[LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS]); + prune.update(LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS, fallback); + if (fallback.is_viable) { + best = fallback; + } + } + + return best; +} + +// probe one strategy across tiers +static void llama_pshard_strategy_sweep( + int strategy, + const llama_pshard_search_ctx & ctx_template, + const llama_context_params & cparams_base, + const std::vector & dmds, + const llama_pshard_plan_registry & registry, + int force_strategy, + llama_pshard_plan * out_plans, // write out_plans[tier] + size_t n_tiers, + size_t first_tier) { + + if (force_strategy >= 0 && force_strategy != strategy) return; + + llama_model_tensor_buft_override local_overrides[4096]; + llama_pshard_search_ctx ctx = ctx_template; + ctx.overrides = local_overrides; + + llama_pshard_tier_prune prune; + prune.init(ctx.n_layers); + + // smaller tiers start from the previous fit + // smaller batches usually need less scratch + // hybrid SSM graphs are not monotonic in batch size + uint32_t prev_n_pinned = 0; + + for (int t = (int)n_tiers - 1; t >= (int)first_tier; t--) { + if (prune.skip[strategy]) break; + + llama_context_params cp_tier = cparams_base; + cp_tier.n_batch = registry.tier_sizes[t]; + cp_tier.n_ubatch = cp_tier.n_batch; + ctx.cparams = &cp_tier; + + llama_pshard_strategy strat = (llama_pshard_strategy)strategy; + llama_pshard_plan plan; + + if (out_plans[t].is_viable && out_plans[t].strategy == strat) { + LLAMA_LOG_INFO("%s: === tier %d (bs=%u) [cached %s] ===\n", + __func__, t, cp_tier.n_batch, llama_pshard_strategy_name(strat)); + plan = out_plans[t]; + } else { + LLAMA_LOG_INFO("%s: === tier %d (bs=%u) [parallel %s] ===\n", + __func__, t, cp_tier.n_batch, llama_pshard_strategy_name(strat)); + + const uint32_t lo_hint = ctx.has_rs ? 0 : prev_n_pinned; + + if (strat == LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS) { + plan = llama_pshard_search_attn_pin(ctx, prune.hi_attn, UINT32_MAX, lo_hint); + plan.batch_size = cp_tier.n_batch; + } else { + plan = llama_pshard_search_strategy(ctx, strat, UINT32_MAX, lo_hint); + plan.batch_size = cp_tier.n_batch; + } + } + + prune.update(strategy, plan); + + if (plan.is_viable && plan.n_pinned > prev_n_pinned) { + prev_n_pinned = plan.n_pinned; + } + + { + const char * status = plan.is_viable ? "VIABLE" : "NOT VIABLE"; + char tps_buf[32] = ""; + if (plan.tps > 0.0f) { snprintf(tps_buf, sizeof(tps_buf), ", tps=%.1f", plan.tps); } + + if (plan.n_attn_pinned > 0) { + LLAMA_LOG_INFO("%s: [tier=%d bs=%-4u %-10s] n_pinned=%2u (attn=%2u), overflow=%-4s, vram=%7.1f MiB, %s%s\n", + __func__, t, cp_tier.n_batch, llama_pshard_strategy_name(strat), + plan.n_pinned, plan.n_attn_pinned, PSHARD_FRAC_NAMES[plan.overflow], + plan.total_vram_req / (1024.0 * 1024.0), status, tps_buf); + } else { + LLAMA_LOG_INFO("%s: [tier=%d bs=%-4u %-10s] n_pinned=%2u, overflow=%-4s, vram=%7.1f MiB, %s%s\n", + __func__, t, cp_tier.n_batch, llama_pshard_strategy_name(strat), + plan.n_pinned, PSHARD_FRAC_NAMES[plan.overflow], + plan.total_vram_req / (1024.0 * 1024.0), status, tps_buf); + } + } + + out_plans[t] = plan; + } +} + +static void llama_pshard_parallel_worker( + std::atomic & next_strategy, + const llama_pshard_search_ctx & ctx_template, + const llama_context_params & cparams_base, + const std::vector & dmds, + const llama_pshard_plan_registry & registry, + int force_strategy, + llama_pshard_plan * all_plans, + size_t n_tiers, + size_t first_tier) { + + while (true) { + int s = next_strategy.fetch_add(1); + if (s >= LLAMA_PSHARD_COUNT) return; + + llama_pshard_plan * out = all_plans + s * n_tiers; + llama_pshard_strategy_sweep(s, ctx_template, cparams_base, dmds, registry, force_strategy, out, n_tiers, first_tier); + } +} + +static bool llama_pshard_params_supported( + const struct llama_model_params * mparams, + const struct llama_context_params * cparams) { + const llama_model_params default_mparams = llama_model_default_params(); + + auto disable = [](const char * reason) { + LLAMA_LOG_WARN("%s: %s, disabling pshard\n", "llama_params_fit_pshard_planning", 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_planning( + const char * path_model, + struct llama_model_params * mparams, + struct llama_context_params * cparams, + struct llama_model_tensor_buft_override * tensor_buft_overrides, + size_t max_vram_mb, + size_t fit_target_mb, + llama_pshard_fit_fn fit_fn, + void * fit_fn_ud) { + const int64_t t0_us = llama_time_us(); + + if (!llama_pshard_params_supported(mparams, cparams)) { + mparams->pshard = false; + mparams->pshard_delegate_compute = false; + cparams->pshard = false; + return; + } + + mparams->pshard = true; + mparams->pshard_delegate_compute = false; + cparams->pshard = true; + + // step 1: probe device memory and model parameters + std::vector devs; + uint32_t hp_ngl = 0, hp_nct = 0, hp_nex = 0, hp_nr = 0; + + llama_model_params mparams_probe = *mparams; + mparams_probe.pshard = false; + const auto dmds = llama_get_device_memory_data( + path_model, &mparams_probe, cparams, devs, hp_ngl, hp_nct, hp_nex, hp_nr, GGML_LOG_LEVEL_ERROR); + + if (devs.empty()) { + LLAMA_LOG_ERROR("%s: no GPU devices found\n", __func__); + return; + } + + const uint32_t n_layers = hp_ngl; + const size_t mib = 1024ULL * 1024ULL; + const size_t actual_vram_free = dmds[0].free; + const size_t fit_target_bytes = fit_target_mb * mib; + const size_t vram_free = max_vram_mb > 0 + ? max_vram_mb * mib + : (actual_vram_free > fit_target_bytes ? actual_vram_free - fit_target_bytes : 0); + + mparams->max_vram_alloc = std::max(1, pshard_bytes_to_mib_ceil(vram_free)); + + LLAMA_LOG_INFO("%s: probing pshard plans: %u layers, %.1f MiB VRAM free, %.1f MiB budget%s\n", + __func__, n_layers, + actual_vram_free / (1024.0 * 1024.0), + vram_free / (1024.0 * 1024.0), + max_vram_mb > 0 ? " (-mva)" : " (free - fit target)"); + + // step 2: read forced strategy from env (PSHARD_STRATEGY) + const int force_strategy = pshard_strategy_from_env(); + if (force_strategy >= 0) { + LLAMA_LOG_INFO("%s: forcing strategy %s (PSHARD_STRATEGY=%s)\n", + __func__, llama_pshard_strategy_name((llama_pshard_strategy)force_strategy), + getenv("PSHARD_STRATEGY")); + } else if (getenv("PSHARD_STRATEGY")) { + LLAMA_LOG_WARN("%s: invalid PSHARD_STRATEGY='%s', ignoring\n", + __func__, getenv("PSHARD_STRATEGY")); + } + + // step 3: derive layout, buftypes, optional benchmark predictor, search ctx + ggml_backend_buffer_type_t gpu_buft = ggml_backend_dev_buffer_type(devs[0].dev); + ggml_backend_buffer_type_t host_buft = ggml_backend_dev_host_buffer_type(devs[0].dev); + if (!host_buft) { + host_buft = ggml_backend_cpu_buffer_type(); + } + const int32_t cpu_bid = pshard_dev_layout::compute_cpu_backend_id(devs.size()); + const pshard_dev_layout layout = pshard_dev_layout::for_device(0, cpu_bid); + + std::unique_ptr predictor; + { + const char * env_cpu = getenv("PSHARD_CPU_PROFILE"); + const char * env_gpu = getenv("PSHARD_GPU_PROFILE"); + const char * cpu_path = env_cpu ? env_cpu : "cpu_profile.txt"; + const char * gpu_path = env_gpu ? env_gpu : "gpu_profile.txt"; + + auto p = std::make_unique(); + const bool has_cpu = p->load_cpu(cpu_path, cparams->n_threads); + const bool has_gpu = p->load_gpu(gpu_path); + if (has_cpu || has_gpu) { + predictor = std::move(p); + LLAMA_LOG_INFO("%s: benchmark predictor loaded (cpu=%s gpu=%s)\n", + __func__, has_cpu ? "yes" : "no", has_gpu ? "yes" : "no"); + } + } + + const uint32_t n_ctx_plan = cparams->n_ctx > 0 ? cparams->n_ctx : hp_nct; + + llama_pshard_search_ctx ctx = { + path_model, mparams, cparams, tensor_buft_overrides, + n_layers, vram_free, gpu_buft, host_buft, layout, + /*is_moe=*/(hp_nex > 0), /*has_rs=*/(hp_nr > 0), + fit_fn, fit_fn_ud, + predictor.get(), n_ctx_plan, 0, + }; + + // step 4: registry lookup -- try to load .tensor_overrides.pshard_registry; merge cached plans by tier + const std::string cache_path = std::string(path_model) + ".tensor_overrides.pshard_registry"; + + int64_t model_file_size = 0; + if (FILE * mf = fopen(path_model, "rb")) { +#ifdef _WIN32 + _fseeki64(mf, 0, SEEK_END); + model_file_size = _ftelli64(mf); +#else + fseeko(mf, 0, SEEK_END); + model_file_size = ftello(mf); +#endif + fclose(mf); + } + const uint64_t fp = pshard_registry_fingerprint(mparams, cparams, model_file_size); + + llama_pshard_plan_registry * registry = mparams->pshard_registry; + bool needs_probe = true; + + if (registry) { + const uint32_t requested_tier_max = registry->cache_ubatch; + const uint32_t tier_max_auto = std::min(std::max(cparams->n_batch, (uint32_t) 16384), n_ctx_plan); + const uint32_t tier_max = std::min(requested_tier_max > 0 ? requested_tier_max : tier_max_auto, n_ctx_plan); + + if (registry->cache_ubatch != tier_max) { + registry->init(tier_max, cparams->n_seq_max); + } + + registry->budget_mib = pshard_bytes_to_mib_ceil(vram_free); + registry->cache_ubatch = registry->tier_sizes.empty() ? 0 : registry->tier_sizes.back(); + ctx.cache_ubatch = registry->cache_ubatch; + + // save requested tiers before load (load overwrites tier_sizes) + std::vector requested_tiers = registry->tier_sizes; + + llama_pshard_plan_registry cached; + cached.budget_mib = registry->budget_mib; + cached.cache_ubatch = registry->cache_ubatch; + if (!mparams->pshard_cache_skip_load && + pshard_registry_load_planning(&cached, fp, cache_path.c_str(), host_buft, vram_free, true)) { + // merge cached plans into registry by matching batch size + std::unordered_map cache_map; + for (size_t i = 0; i < cached.tier_sizes.size(); i++) { + cache_map[cached.tier_sizes[i]] = cached.best_plans[i]; + } + + registry->tier_sizes = requested_tiers; + registry->best_plans.resize(requested_tiers.size()); + registry->pshard_disabled = cached.pshard_disabled; + registry->baseline_vram_req = cached.baseline_vram_req; + registry->budget_mib = cached.budget_mib; + registry->cache_ubatch = cached.cache_ubatch; + + size_t n_hit = 0; + for (size_t i = 0; i < requested_tiers.size(); i++) { + auto it = cache_map.find(requested_tiers[i]); + if (it != cache_map.end()) { + registry->best_plans[i] = it->second; + n_hit++; + } + } + + if (n_hit == requested_tiers.size()) { + needs_probe = false; + LLAMA_LOG_INFO("%s: loaded all %zu tiers from exact budget=%u MiB cache variant\n", + __func__, n_hit, registry->budget_mib); + } else if (n_hit > 0) { + LLAMA_LOG_INFO("%s: loaded %zu/%zu tiers from exact budget=%u MiB cache variant, %zu need probing\n", + __func__, n_hit, requested_tiers.size(), registry->budget_mib, requested_tiers.size() - n_hit); + } + } + + // common setup for cache hit and miss + } + + // step 4b: skip pshard when a cached baseline variant fits this budget + if (registry && registry->pshard_disabled) { + LLAMA_LOG_INFO("%s: cache says baseline %.1f MiB fits this budget (variant budget=%u MiB cache_ubatch=%u), using baseline loading\n", + __func__, registry->baseline_vram_req / (1024.0 * 1024.0), registry->budget_mib, registry->cache_ubatch); + mparams->pshard = false; + mparams->pshard_delegate_compute = false; + cparams->pshard = false; + mparams->n_gpu_layers = n_layers + 1; + tensor_buft_overrides[0] = { nullptr, nullptr, -1 }; + mparams->tensor_buft_overrides = nullptr; + + const int64_t t1_us = llama_time_us(); + LLAMA_LOG_INFO("%s: best strategy: baseline (cached), all %u layers on GPU, took %.2f s\n", + __func__, n_layers, (t1_us - t0_us) * 1e-6); + return; + } + + // step 5: probe tiers largest first and skip pshard when baseline already fits + if (registry && needs_probe) { + const size_t n_tiers = registry->tier_sizes.size(); + + // step 5a: baseline off-ramp only for global tiers + static constexpr uint32_t GLOBAL_FIT_MIN_BATCH = 512; + size_t first_probe_tier = 0; + size_t min_global_tier = n_tiers; + const bool run_baseline_offramp = force_strategy < 0; + if (run_baseline_offramp) { + for (size_t t = 0; t < n_tiers; t++) { + if (registry->tier_sizes[t] >= GLOBAL_FIT_MIN_BATCH) { + min_global_tier = t; + break; + } + } + + if (min_global_tier < n_tiers) { + size_t global_fit_tier = n_tiers; + + for (size_t t = n_tiers; t-- > min_global_tier; ) { + llama_pshard_plan baseline_plan; + if (pshard_plan_is_baseline_fit(registry->best_plans[t], n_layers, layout.compute)) { + LLAMA_LOG_INFO("%s: === tier %zu (bs=%u) [cached global-fit check] ===\n", + __func__, t, registry->tier_sizes[t]); + baseline_plan = registry->best_plans[t]; + } else { + const llama_context_params * saved = ctx.cparams; + llama_context_params cp_tier = *cparams; + cp_tier.n_batch = registry->tier_sizes[t]; + cp_tier.n_ubatch = cp_tier.n_batch; + ctx.cparams = &cp_tier; + + LLAMA_LOG_INFO("%s: === tier %zu (bs=%u) [global-fit check] ===\n", + __func__, t, registry->tier_sizes[t]); + baseline_plan = llama_pshard_search_baseline_fit_tier(ctx, dmds); + + ctx.cparams = saved; + } + + if (pshard_plan_is_baseline_fit(baseline_plan, n_layers, layout.compute)) { + registry->best_plans[t] = baseline_plan; + global_fit_tier = t; + break; + } + } + + if (global_fit_tier < n_tiers) { + const llama_pshard_plan & global = registry->best_plans[global_fit_tier]; + const uint32_t global_ubatch = registry->tier_sizes[global_fit_tier]; + const uint32_t global_cache_ubatch = std::min(global_ubatch, GLOBAL_FIT_MIN_BATCH); + const size_t global_vram_req = global.total_vram_req; + + registry->pshard_disabled = true; + registry->baseline_vram_req = global_vram_req; + registry->cache_ubatch = global_cache_ubatch; + registry->tier_sizes.clear(); + registry->best_plans.clear(); + pshard_registry_save(registry, fp, cache_path.c_str(), host_buft, cparams); + + mparams->pshard = false; + mparams->pshard_delegate_compute = false; + cparams->pshard = false; + mparams->n_gpu_layers = n_layers + 1; + tensor_buft_overrides[0] = { nullptr, nullptr, -1 }; + mparams->tensor_buft_overrides = nullptr; + + const int64_t t1_us = llama_time_us(); + LLAMA_LOG_INFO("%s: full baseline fit at tier %zu (bs=%u, %.1f MiB); using baseline loading with cache_ubatch=%u, took %.2f s\n", + __func__, global_fit_tier, global_ubatch, + global_vram_req / (1024.0 * 1024.0), + global_cache_ubatch, (t1_us - t0_us) * 1e-6); + return; + } + } + + LLAMA_LOG_INFO("%s: no full-fit global plan found down to bs=%u\n", + __func__, GLOBAL_FIT_MIN_BATCH); + } else { + LLAMA_LOG_INFO("%s: skipping baseline global-fit check for forced strategy %s\n", + __func__, llama_pshard_strategy_name((llama_pshard_strategy) force_strategy)); + } + + // step 5b: probe remaining tiers + if (first_probe_tier < n_tiers) + { + std::vector all_plans(LLAMA_PSHARD_COUNT * n_tiers); + for (size_t t = 0; t < n_tiers; t++) { + const llama_pshard_plan & plan = registry->best_plans[t]; + if (plan.is_viable && plan.strategy >= 0 && plan.strategy < LLAMA_PSHARD_COUNT) { + all_plans[(int) plan.strategy * n_tiers + t] = plan; + } + } + + int n_workers = std::min((int)cparams->n_threads, (int)LLAMA_PSHARD_COUNT); + n_workers = std::max(n_workers, 1); + + int n_pshard_strategies = LLAMA_PSHARD_COUNT; + n_workers = std::min(n_workers, n_pshard_strategies); + + LLAMA_LOG_INFO("%s: parallel planning with %d workers (%d pshard strategies)\n", + __func__, n_workers, n_pshard_strategies); + + std::atomic next_strategy{0}; + + if (n_workers <= 1) { + llama_pshard_parallel_worker(next_strategy, ctx, *cparams, dmds, + *registry, force_strategy, all_plans.data(), n_tiers, first_probe_tier); + } else { + std::vector threads; + for (int w = 1; w < n_workers; w++) { + threads.emplace_back(llama_pshard_parallel_worker, + std::ref(next_strategy), std::cref(ctx), std::cref(*cparams), + std::cref(dmds), std::cref(*registry), force_strategy, + all_plans.data(), n_tiers, first_probe_tier); + } + llama_pshard_parallel_worker(next_strategy, ctx, *cparams, dmds, + *registry, force_strategy, all_plans.data(), n_tiers, first_probe_tier); + + for (auto & t : threads) t.join(); + } + + for (size_t t = first_probe_tier; t < n_tiers; t++) { + llama_pshard_plan best; + for (int s = 0; s < LLAMA_PSHARD_COUNT; s++) { + auto & p = all_plans[s * n_tiers + t]; + if (p.is_viable && pshard_plan_is_better(p, best)) { + best = p; + } + } + registry->best_plans[t] = best; + } + + // step 5d: try attention priority fallback for forced strategies that cannot fit + // forced strategy sweeps can leave smaller tiers unfilled + if (force_strategy >= 0 && force_strategy != LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS) { + for (size_t t = first_probe_tier; t < n_tiers; t++) { + if (registry->best_plans[t].is_viable) continue; + + const llama_context_params * saved = ctx.cparams; + llama_context_params cp_tier = *cparams; + cp_tier.n_batch = registry->tier_sizes[t]; + cp_tier.n_ubatch = cp_tier.n_batch; + ctx.cparams = &cp_tier; + + llama_pshard_plan fallback = llama_pshard_attn_pin_fallback(ctx, force_strategy); + if (fallback.is_viable) { + registry->best_plans[t] = fallback; + } + + ctx.cparams = saved; + } + } + } + + pshard_registry_save(registry, fp, cache_path.c_str(), host_buft, cparams); + } + + // step 6: pick the active plan + if (registry) { + for (size_t t = registry->tier_sizes.size(); t-- > 0; ) { + auto * best = registry->get_best(t); + if (best && best->is_viable) { + registry->active_plan = best; + break; + } + } + } + + llama_pshard_plan best_plan; + if (registry && registry->active_plan) { + best_plan = *registry->active_plan; + } else { + llama_pshard_tier_prune prune_single; + prune_single.init(n_layers); + best_plan = llama_pshard_search_tier(ctx, force_strategy, dmds, prune_single); + if (best_plan.is_viable && registry) { + for (size_t t = 0; t < registry->tier_sizes.size(); t++) { + if (registry->tier_sizes[t] == best_plan.batch_size) { + registry->best_plans[t] = best_plan; + registry->active_plan = ®istry->best_plans[t]; + pshard_registry_save(registry, fp, cache_path.c_str(), host_buft, cparams); + break; + } + } + } + } + + // step 7: fall back to all cpu when no plan is viable + if (!best_plan.is_viable) { + LLAMA_LOG_WARN("%s: no viable plan found, falling back to STATIC_ATTNPRIO_ALLMODELS with n_pinned=0\n", __func__); + llama_pshard_generate_overrides_planning(0, n_layers, gpu_buft, host_buft, tensor_buft_overrides, + LLAMA_LAYER_FRACTION_NONE, LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS, layout, + /*pin_from_back=*/false, /*output_on_gpu=*/false, /*n_attn_pinned=*/0); + mparams->pshard_delegate_compute = true; + mparams->n_gpu_layers = n_layers + 1; + mparams->tensor_buft_overrides = tensor_buft_overrides; + return; + } + + // step 8: apply best plan to tensor_buft_overrides + llama_pshard_generate_overrides_planning(best_plan.n_pinned, n_layers, gpu_buft, host_buft, + tensor_buft_overrides, (llama_layer_fraction)best_plan.overflow, best_plan.strategy, layout, + best_plan.pin_from_back, best_plan.output_on_gpu, best_plan.n_attn_pinned); + + for (size_t i = 0; tensor_buft_overrides[i].pattern; i++) { + if (tensor_buft_overrides[i].backend_id == layout.compute) { + tensor_buft_overrides[i].buft = host_buft; + } + } + + mparams->pshard_delegate_compute = llama_pshard_strategy_delegates_compute(best_plan.strategy); + mparams->n_gpu_layers = n_layers + 1; + mparams->tensor_buft_overrides = tensor_buft_overrides; + + const int64_t t1_us = llama_time_us(); + LLAMA_LOG_INFO("%s: best strategy: %s, n_pinned=%u, n_attn=%u/%u%s%s, took %.2f s\n", + __func__, llama_pshard_strategy_name(best_plan.strategy), + best_plan.n_pinned, best_plan.n_attn_pinned, n_layers, + best_plan.overflow ? " (partial: " : "", + best_plan.overflow ? PSHARD_FRAC_NAMES[best_plan.overflow] : "", + (t1_us - t0_us) * 1e-6); + + { + int n_bid0 = 0, n_total = 0; + for (const auto * ov = tensor_buft_overrides; ov->pattern; ++ov) { + n_total++; + if (ov->backend_id == layout.compute) n_bid0++; + LLAMA_LOG_DEBUG("%s: override: %-25s -> %-15s backend_id=%d\n", + __func__, ov->pattern, ggml_backend_buft_name(ov->buft), ov->backend_id); + } + LLAMA_LOG_INFO("%s: %d overrides, %d with bid=%d (compute)\n", + __func__, n_total, n_bid0, layout.compute); + } +} diff --git a/src/llama-pshard-plan.h b/src/llama-pshard-plan.h new file mode 100644 index 000000000000..527fa359ebf5 --- /dev/null +++ b/src/llama-pshard-plan.h @@ -0,0 +1,272 @@ +#pragma once + +#include "llama.h" +#include "llama-context.h" +#include "llama-cparams.h" +#include "llama-model.h" + +#include +#include +#include +#include +#include +#include + +enum llama_pshard_strategy { + LLAMA_PSHARD_GPUONLY_LAYERPIN_LAYERSTREAM = 0, + LLAMA_PSHARD_GPUONLY_ATTNPIN_FFNSTREAM = 1, + LLAMA_PSHARD_DYNAMIC_FFNCPU_ATTNSTREAM = 2, + LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS = 3, + LLAMA_PSHARD_COUNT +}; + +// strategy name for logging +inline const char * llama_pshard_strategy_name(llama_pshard_strategy s) { + switch (s) { + case LLAMA_PSHARD_GPUONLY_LAYERPIN_LAYERSTREAM: return "GPUONLY_LAYERPIN_LAYERSTREAM"; + case LLAMA_PSHARD_GPUONLY_ATTNPIN_FFNSTREAM: return "GPUONLY_ATTNPIN_FFNSTREAM"; + case LLAMA_PSHARD_DYNAMIC_FFNCPU_ATTNSTREAM: return "DYNAMIC_FFNCPU_ATTNSTREAM"; + case LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS: return "STATIC_ATTNPRIO_ALLMODELS"; + default: return "UNKNOWN"; + } +} + +inline bool llama_pshard_strategy_delegates_compute(llama_pshard_strategy s) { + return s == LLAMA_PSHARD_STATIC_ATTNPRIO_ALLMODELS; +} + +// PSHARD_STRATEGY accepts a name or numeric id +inline int pshard_strategy_from_env() { + const char * env = getenv("PSHARD_STRATEGY"); + if (!env || !*env) return -1; + for (int i = 0; i < LLAMA_PSHARD_COUNT; i++) { + if (strcmp(env, llama_pshard_strategy_name((llama_pshard_strategy)i)) == 0) { + return i; + } + } + char * end = nullptr; + long v = strtol(env, &end, 10); + if (end != env && *end == '\0' && v >= 0 && v < LLAMA_PSHARD_COUNT) { + return (int)v; + } + return -1; +} + +// cached tensor override entry +struct llama_pshard_override { + std::string pattern; + ggml_backend_buffer_type_t buft; + int32_t backend_id; +}; + +// 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_ATTNPRIO_ALLMODELS; + uint32_t batch_size = 0; + uint32_t n_pinned = 0; // fully pinned layers (all tensors on GPU) + uint32_t n_attn_pinned = 0; // attention priority layers on GPU (>= n_pinned) + int overflow = 0; // llama_layer_fraction + bool pin_from_back = false; + bool output_on_gpu = false; + + std::vector overrides; + + size_t total_vram_req = 0; + size_t scratch_measured = 0; + size_t cache_measured = 0; + float tps = 0.0f; // predicted tokens/sec (0 = no benchmark data) + bool is_viable = false; + + // 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); + +// Override generation. The planning and inference stages DIVERGE here: the planner +// emits gpu_buft for pinned layers, the runtime emits host_buft, and the planner adds +// override-array bounds assertions. Both variants are kept verbatim so the difference +// stays visible; reconciling them is deliberately a separate change. +// _planning -> llama-pshard-plan.cpp (file-local, used by the tier search) +// _inference -> llama-pshard-cache.cpp (used when applying a cached plan) +void llama_pshard_generate_overrides_inference( + 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); + +// probe hook runs before context teardown +// used by TPS prediction to inspect scheduler splits +typedef void (*llama_probe_hook_t)(llama_context * ctx, void * user_data); + +std::vector llama_get_device_memory_data( + const char * path_model, const struct llama_model_params * mparams, + const struct llama_context_params * cparams, + std::vector & devs, uint32_t & hp_ngl, + uint32_t & hp_n_ctx_train, uint32_t & hp_n_expert, uint32_t & hp_n_embd_r, + enum ggml_log_level log_level, + llama_probe_hook_t probe_hook = nullptr, + void * probe_hook_data = nullptr, + uint32_t probe_n_tokens = 0, + uint32_t probe_n_outputs = 0); + +// llama_device_memory_data / llama_memory_breakdown_data come from llama-ext.h +// (included via llama-context.h) +// +// llama_pshard_fit_fn, llama_params_fit_pshard_inference() (runtime), +// llama_params_fit_pshard_planning() (planner) and llama_pshard_registry_create/free() +// are declared in llama-ext.h so common/ and tools/ can reach them without +// touching the public llama.h + +// 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); + +// planner-only: writes/updates the registry file +bool pshard_registry_save( + const struct llama_pshard_plan_registry * registry, uint64_t fingerprint, + const char * cache_path, ggml_backend_buffer_type_t host_buft, + const struct llama_context_params * cparams = nullptr); + + +// Registry load. The two stages DIVERGE only in diagnostics and in where the +// cache_ubatch==0 fallback is applied (the runtime re-applies it in its entry point). +// Both variants are kept verbatim so the difference stays visible. +bool pshard_registry_load_planning( + 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); + +bool pshard_registry_load_inference( + struct llama_pshard_plan_registry * registry, uint64_t fingerprint, + const char * cache_path, ggml_backend_buffer_type_t host_buft, + size_t current_budget, bool require_exact_budget = false); + +struct llama_pshard_plan_registry { + std::vector tier_sizes; + std::vector best_plans; // one best plan per tier + llama_pshard_plan * active_plan = nullptr; + uint32_t budget_mib = 0; + uint32_t cache_ubatch = 0; + + // variant marker for a baseline load that fits + // runtime still checks baseline_vram_req against the current budget + bool pshard_disabled = false; + size_t baseline_vram_req = 0; + + void init(uint32_t n_ubatch, uint32_t n_parallel = 1, uint32_t n_draft = 0) { + tier_sizes.clear(); + best_plans.clear(); + + if (n_ubatch == 0) { + cache_ubatch = 0; + return; + } + + // decode tiers + if (n_parallel <= 1) { + tier_sizes.push_back(1); + tier_sizes.push_back(16); + if (n_draft > 0) { + uint32_t verify_tier = n_draft + 1; + if (verify_tier > 16) { + tier_sizes.push_back(verify_tier); + } + } + } else { + for (uint32_t t = 1; t <= 64 && t < 512; t *= 4) { + tier_sizes.push_back(t); + if (t == 16) { + tier_sizes.push_back(32); + } + } + } + + // prefill tiers: x2 growth from 512 + for (uint32_t t = 512; t < n_ubatch; t *= 2) { + if (tier_sizes.empty() || tier_sizes.back() < t) { + tier_sizes.push_back(t); + } + } + + if (tier_sizes.empty() || tier_sizes.back() != n_ubatch) { + tier_sizes.push_back(n_ubatch); + } + + cache_ubatch = tier_sizes.empty() ? 0 : tier_sizes.back(); + best_plans.resize(tier_sizes.size()); + } + + size_t tier_index(uint32_t batch_size) const { + for (size_t i = 0; i < tier_sizes.size(); i++) { + if (tier_sizes[i] >= batch_size) return i; + } + return tier_sizes.size() - 1; + } + + llama_pshard_plan * get_best(size_t tier) { + return best_plans[tier].is_viable ? &best_plans[tier] : nullptr; + } + + // 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 11ac9656d9f9..c492c33146a4 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -9,6 +9,7 @@ #include "llama-model-loader.h" #include "llama-model-saver.h" #include "llama-model.h" +#include "llama-pshard-plan.h" #include "ggml.h" #include "ggml-cpp.h" @@ -23,6 +24,8 @@ #include #include #include +#include +#include #include #include @@ -349,6 +352,8 @@ static std::pair llama_model_load(struct gguf_context * meta return {0, model_ptr.release()}; } + ml.force_duplicate_tied = params.pshard; + if (!model->load_tensors(ml)) { return {-2, nullptr}; } @@ -600,4 +605,3 @@ const char * llama_print_system_info(void) { return s.c_str(); } - diff --git a/tools/fit-params/CMakeLists.txt b/tools/fit-params/CMakeLists.txt index 8acdaef3712a..dd3805b7f509 100644 --- a/tools/fit-params/CMakeLists.txt +++ b/tools/fit-params/CMakeLists.txt @@ -2,7 +2,7 @@ set(TARGET llama-fit-params-impl) -add_library(${TARGET} fit-params.cpp) +add_library(${TARGET} fit-params.cpp fit-params-pshard.cpp) set_target_properties(${TARGET} PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS ON) target_include_directories(${TARGET} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}) diff --git a/tools/fit-params/README-pshard.md b/tools/fit-params/README-pshard.md new file mode 100644 index 000000000000..e19c61cdb132 --- /dev/null +++ b/tools/fit-params/README-pshard.md @@ -0,0 +1,124 @@ +# pshard planning (`llama-fit-params --pshard`) + +`llama-fit-params --pshard` writes a pshard registry for a model and VRAM budget. + +It probes batch-size tiers such as `bs=1`, `bs=16`, and `bs=512`, compares the available placement strategies, and writes the selected tensor overrides to: + +```text +.gguf.tensor_overrides.pshard_registry +``` + +When CPU/GPU profiler results are available, the planner uses them to estimate throughput for each candidate plan. At runtime, pshard selects the smallest tier that covers the current batch. + +The registry is plain text so it can be inspected when debugging planner output. + +## Example usage + +```bash +# 1. Generate roofline profiles once per machine +# (writes profile JSONs next to the binary) +./build/bin/llama-profiler-cpu +./build/bin/llama-profiler-gpu + +# 2. Plan placement for a model + VRAM budget +./build/bin/llama-fit-params --pshard \ + --model /opt/models/Qwen3.5-27B-Q4_0.gguf \ + --max-vram-alloc 12000 \ + -c 8192 \ + -fa on + +ggml_cuda_init: found 1 CUDA devices (Total VRAM: 16302 MiB): + Device 0: NVIDIA GeForce RTX 5070 Ti, compute capability 12.0, VMM: yes, VRAM: 16302 MiB +main: planning pshard tensor overrides... +llama_params_fit_pshard: probing pshard plans: 64 layers, 12000.0 MiB VRAM free +pshard_registry_save: saved budget=12000 MiB cache_ubatch=8192 variant with 7 tier plans to /opt/models/Qwen3.5-27B-Q4_0.gguf.tensor_overrides.pshard_registry +main: planning complete, registry written next to model file + +# 3. Inspect the registry +cat /opt/models/Qwen3.5-27B-Q4_0.gguf.tensor_overrides.pshard_registry +``` + +Example registry output: + +```text +# Generated file. Edit at your own risk. + +[fingerprint=0x9041bb5b253cf89f] +# n_ctx=8192 n_seq_max=1 n_threads=8 fa=on type_k=1 type_v=1 strategy=auto + +[variant budget=12000 cache_ubatch=8192] +[tier 0 bs=1] +strategy=DYNAMIC_FFNCPU_ATTNSTREAM n_pinned=54 n_attn_pinned=0 overflow=NONE tps=12.85 vram=11975.7 output_on_gpu=0 pin_from_back=0 +ot=^output=CUDA_Host:3,^token_embd=CUDA_Host:3,blk\.0\..*=CUDA_Host:0, ... ,blk\.53\..*=CUDA_Host:0,blk\.54\.ffn_(up|gate|down).*=CUDA_Host:3,blk\.54\..*=CUDA_Host:1,blk\.55\.ffn_(up|gate|down).*=CUDA_Host:3,blk\.55\..*=CUDA_Host:2, ... ,blk\.63\.ffn_(up|gate|down).*=CUDA_Host:3,blk\.63\..*=CUDA_Host:2 +[tier 1 bs=16] +strategy=GPUONLY_ATTNPIN_FFNSTREAM n_pinned=48 n_attn_pinned=0 overflow=NONE tps=208.27 vram=11964.1 output_on_gpu=0 pin_from_back=0 +ot=... +[tier 2 bs=512] +strategy=GPUONLY_LAYERPIN_LAYERSTREAM n_pinned=52 n_attn_pinned=0 overflow=NONE tps=1994.90 vram=11973.3 output_on_gpu=0 pin_from_back=0 +ot=... +[tier 3 bs=1024] +strategy=GPUONLY_LAYERPIN_LAYERSTREAM n_pinned=51 n_attn_pinned=0 overflow=NONE tps=2286.87 vram=11909.8 output_on_gpu=0 pin_from_back=0 +ot=... +[tier 4 bs=2048] +strategy=GPUONLY_ATTNPIN_FFNSTREAM n_pinned=44 n_attn_pinned=0 overflow=NONE tps=2249.44 vram=11914.3 output_on_gpu=0 pin_from_back=0 +ot=... +[tier 5 bs=4096] +strategy=GPUONLY_ATTNPIN_FFNSTREAM n_pinned=40 n_attn_pinned=0 overflow=NONE tps=2095.88 vram=11868.6 output_on_gpu=0 pin_from_back=0 +ot=... +[tier 6 bs=8192] +strategy=GPUONLY_ATTNPIN_FFNSTREAM n_pinned=33 n_attn_pinned=0 overflow=NONE tps=2148.21 vram=11937.2 output_on_gpu=0 pin_from_back=0 +ot=^output=CUDA_Host:3,^token_embd=CUDA_Host:3,blk\.0\..*=CUDA_Host:0, ... ,blk\.32\..*=CUDA_Host:0,blk\.33\.ffn_(up|gate|down).*=CUDA_Host:2,blk\.33\..*=CUDA_Host:0, ... ,blk\.63\.ffn_(up|gate|down).*=CUDA_Host:2,blk\.63\..*=CUDA_Host:0 +``` + +## Planning for llama-bench + +Use `--bench-plan` when the registry is intended for `llama-bench -pshard`. The planner accepts the bench shape flags `-p` / `--n-prompt`, `-n` / `--n-gen`, `-pg`, and `-d` / `--n-depth`, then plans each unique context that llama-bench can run: + +```bash +./build/bin/llama-fit-params --pshard \ + --model /opt/models/Qwen3.5-35B-A3B-Q8_0.gguf \ + --bench-plan \ + -pg 512,200 \ + -pg 2048,200 \ + -d 0,1024 +``` + +For each generated context, `n_ctx` is the number of tokens that can be resident during the test: + +- prompt-only: `n_ctx = p + d` +- generation-only: `n_ctx = n + d` +- prompt+generation: `n_ctx = p + n + d` + +The largest planned tier for a bench context is capped by prompt batch demand, not by depth. For example, `-pg 2048,200 -d 1024` plans `n_ctx=3272` with `tier_cap=2048`, so the registry includes decode tiers plus prompt tiers up to `bs=2048`. + +If `-fa` / `--flash-attn` is not provided and `LLAMA_ARG_FLASH_ATTN` is not set, `--bench-plan` uses Flash Attention off to match llama-bench defaults. Normal non-bench planning keeps the regular common parameter defaults. + +## Strategies + +In the strategy names, `ATTN` refers to the attention/dense side of the layer, as opposed to FFN/MoE weights. + +Static schedules run GPU-resident tensors on GPU and CPU-resident tensors on CPU, with no streamed GPU execution for host-resident weights. + +- `STATIC_ATTNPRIO_ALLMODELS`: static attention-priority placement. It pins the attention/dense side across as many layers as fit, then uses the remaining budget to pin full layers. FFN/MoE that does not fit remains on CPU. Unlike `llama_params_fit`, this attention-priority placement applies to dense models too. + +Dynamic schedules split the layer between CPU and GPU execution. Some host-resident tensors are streamed to GPU scratch for execution, while other parts of the layer remain on CPU. + +- `DYNAMIC_FFNCPU_ATTNSTREAM`: pin as many full layers as fit, keep FFN/MoE on CPU in the remaining layers, and stream the attention/dense side for GPU execution. + +GPU-only schedules execute repeating-layer compute on GPU. Weights that do not fit in VRAM stay resident in host memory and are streamed to GPU scratch before use. + +- `GPUONLY_LAYERPIN_LAYERSTREAM`: pin as many full layers as fit, then stream the remaining layers for GPU execution. +- `GPUONLY_ATTNPIN_FFNSTREAM`: pin the attention/dense side for all layers, pin as many complete layers as the remaining budget allows, and stream FFN/MoE weights for GPU execution. + +## Notes + +- `--max-vram-alloc` / `-mva ` sets the absolute planning budget. If it is `0` or omitted, the planner uses the free VRAM available at planning time minus `--fit-target` / `-fitt` (default 1024 MiB). When `-mva` is non-zero, `-fitt` is ignored for pshard. +- `--pshard-tier-max ` caps the largest batch size the planner probes. The default is `min(max(n_batch, 16384), n_ctx)`. +- The `[fingerprint=...]` line invalidates the registry when plan-compatible inputs change (`n_ctx`, `n_seq_max`, threads, FA mode, KV cache types, GGUF file size, or forced `PSHARD_STRATEGY`). The comment below the fingerprint lists those same inputs. +- A fingerprint can contain multiple `[variant budget=... cache_ubatch=...]` blocks. Re-running the planner with a new budget or cache ubatch replaces only that variant and keeps the others. +- `cache_ubatch` records the runtime ubatch used for context/KV/SWA cache sizing. Tier compute scratch is still measured with the tier batch size. +- `pshard_disabled=1 baseline_vram=` is variant-scoped. Runtime skips pshard only when the measured baseline VRAM fits the current budget. +- `backend_id` values: + - `0` = GPU pinned compute + - `1` / `2` = shard compute lanes used for pipeline overlap + - `3` = CPU diff --git a/tools/fit-params/fit-params-pshard.cpp b/tools/fit-params/fit-params-pshard.cpp new file mode 100644 index 000000000000..15d815346f0f --- /dev/null +++ b/tools/fit-params/fit-params-pshard.cpp @@ -0,0 +1,272 @@ +#include "llama.h" +#include "../src/llama-ext.h" + +#include "arg.h" +#include "common.h" +#include "fit.h" +#include "log.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_MSC_VER) +#pragma warning(disable: 4244 4267) // possible loss of data +#endif + +struct bench_plan_params { + bool enabled = false; + + std::vector n_prompt; + std::vector n_gen; + std::vector> n_pg; + std::vector n_depth; + bool flash_attn_seen = false; +}; + +static std::vector parse_int_range(const std::string & s) { + // first[-last[(+|*)step]] + std::regex range_regex(R"(^(\d+)(?:-(\d+)(?:([\+|\*])(\d+))?)?(?:,|$))"); + + std::smatch match; + std::string::const_iterator search_start(s.cbegin()); + + std::vector result; + while (std::regex_search(search_start, s.cend(), match, range_regex)) { + int first = std::stoi(match[1]); + int last = match[2].matched ? std::stoi(match[2]) : first; + char op = match[3].matched ? match[3].str()[0] : '+'; + int step = match[4].matched ? std::stoi(match[4]) : 1; + + for (int i = first; i <= last;) { + result.push_back(i); + + int prev_i = i; + if (op == '+') { + i += step; + } else if (op == '*') { + i *= step; + } else { + throw std::invalid_argument("invalid range format"); + } + + if (i <= prev_i) { + throw std::invalid_argument("invalid range"); + } + } + + search_start = match.suffix().first; + } + + if (search_start != s.cend()) { + throw std::invalid_argument("invalid range format"); + } + + return result; +} + +static std::pair parse_pg(const std::string & s) { + const size_t comma = s.find(','); + if (comma == std::string::npos || s.find(',', comma + 1) != std::string::npos) { + throw std::invalid_argument("invalid -pg format"); + } + + return { std::stoi(s.substr(0, comma)), std::stoi(s.substr(comma + 1)) }; +} + +static bool is_bench_plan_arg_with_value(const std::string & arg) { + return arg == "-p" || arg == "--n-prompt" || + arg == "-n" || arg == "--n-gen" || + arg == "-pg" || + arg == "-d" || arg == "--n-depth"; +} + +static bool parse_bench_plan_args(int argc, char ** argv, bench_plan_params & bench, std::vector & forwarded) { + for (int i = 1; i < argc; i++) { + if (std::string(argv[i]) == "--bench-plan") { + bench.enabled = true; + break; + } + } + + forwarded.clear(); + forwarded.push_back(argv[0]); + + for (int i = 1; i < argc; i++) { + const std::string arg = argv[i]; + + if (arg == "--bench-plan") { + continue; + } + + if (bench.enabled && (arg == "-fa" || arg == "--flash-attn")) { + bench.flash_attn_seen = true; + } + + if (!bench.enabled || !is_bench_plan_arg_with_value(arg)) { + forwarded.push_back(arg); + continue; + } + + if (++i >= argc) { + fprintf(stderr, "%s: missing value for %s\n", __func__, arg.c_str()); + return false; + } + + try { + if (arg == "-p" || arg == "--n-prompt") { + auto vals = parse_int_range(argv[i]); + bench.n_prompt.insert(bench.n_prompt.end(), vals.begin(), vals.end()); + } else if (arg == "-n" || arg == "--n-gen") { + auto vals = parse_int_range(argv[i]); + bench.n_gen.insert(bench.n_gen.end(), vals.begin(), vals.end()); + } else if (arg == "-pg") { + bench.n_pg.push_back(parse_pg(argv[i])); + } else if (arg == "-d" || arg == "--n-depth") { + auto vals = parse_int_range(argv[i]); + bench.n_depth.insert(bench.n_depth.end(), vals.begin(), vals.end()); + } + } catch (const std::exception & e) { + fprintf(stderr, "%s: invalid value for %s: %s (%s)\n", __func__, arg.c_str(), argv[i], e.what()); + return false; + } + } + + if (bench.enabled) { + if (bench.n_prompt.empty()) { + bench.n_prompt.push_back(512); + } + if (bench.n_gen.empty()) { + bench.n_gen.push_back(128); + } + if (bench.n_depth.empty()) { + bench.n_depth.push_back(0); + } + } + + return true; +} + +static void bench_plan_add_case(std::map & result, uint32_t n_ctx, uint32_t tier_cap) { + if (n_ctx == 0) { + return; + } + + auto & max_tier_cap = result[n_ctx]; + max_tier_cap = std::max(max_tier_cap, tier_cap); +} + +static std::map bench_plan_context_tier_caps(const bench_plan_params & bench) { + std::map result; + + for (const int d : bench.n_depth) { + for (const int p : bench.n_prompt) { + if (p > 0) { + bench_plan_add_case(result, (uint32_t) (p + d), (uint32_t) p); + } + } + + for (const int n : bench.n_gen) { + if (n > 0) { + bench_plan_add_case(result, (uint32_t) (n + d), 0); + } + } + + for (const auto & pg : bench.n_pg) { + if (pg.first > 0 || pg.second > 0) { + bench_plan_add_case(result, (uint32_t) (pg.first + pg.second + d), (uint32_t) std::max(pg.first, 0)); + } + } + } + + return result; +} + +static void plan_pshard_context(common_params & params, uint32_t n_ctx, uint32_t bench_tier_cap = 0, bool bench_plan = false) { + auto mparams = common_model_params_to_llama(params); + auto cparams = common_context_params_to_llama(params); + + cparams.n_ctx = n_ctx; + + params.tensor_buft_overrides.assign(4096, {}); + + const uint32_t tier_max_auto = bench_plan ? + std::max(bench_tier_cap, (uint32_t) 16) : + std::min(std::max(cparams.n_batch, (uint32_t) 16384), cparams.n_ctx); + const uint32_t tier_max_user = params.pshard_tier_max > 0 ? std::min(params.pshard_tier_max, tier_max_auto) : tier_max_auto; + const uint32_t tier_max = bench_plan ? tier_max_user : std::min(tier_max_user, cparams.n_ctx); + + mparams.pshard_registry = llama_pshard_registry_create(tier_max, cparams.n_seq_max); + mparams.pshard_cache_skip_load = true; + + const size_t fit_target_mb = params.fit_params_target.empty() ? 0 : params.fit_params_target[0] / (1024 * 1024); + + if (bench_plan) { + LOG_INF("%s: planning pshard tensor overrides for n_ctx=%u tier_cap=%u tier_max=%u...\n", + __func__, n_ctx, bench_tier_cap, tier_max); + } else { + LOG_INF("%s: planning pshard tensor overrides for n_ctx=%u tier_max=%u...\n", __func__, n_ctx, tier_max); + } + common_pshard_plan(params.model.path.c_str(), &mparams, &cparams, + params.tensor_buft_overrides.data(), params.max_vram_alloc, fit_target_mb); + + llama_pshard_registry_free(mparams.pshard_registry); +} + +// satisfies -Wmissing-declarations +int llama_fit_params_pshard(int argc, char ** argv); + +int llama_fit_params_pshard(int argc, char ** argv) { + common_params params; + bench_plan_params bench; + std::vector forwarded; + + if (!parse_bench_plan_args(argc, argv, bench, forwarded)) { + return 1; + } + + std::vector forwarded_argv; + forwarded_argv.reserve(forwarded.size()); + for (std::string & arg : forwarded) { + forwarded_argv.push_back(arg.data()); + } + + if (!common_params_parse((int) forwarded_argv.size(), forwarded_argv.data(), params, LLAMA_EXAMPLE_COMMON)) { + return 1; + } + + if (bench.enabled && !bench.flash_attn_seen && getenv("LLAMA_ARG_FLASH_ATTN") == nullptr) { + params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; + } + + common_init(); + llama_backend_init(); + llama_numa_init(params.numa); + + if (bench.enabled) { + const std::map ctx_caps = bench_plan_context_tier_caps(bench); + if (ctx_caps.empty()) { + fprintf(stderr, "%s: --bench-plan produced no non-zero llama-bench contexts\n", __func__); + return 1; + } + + LOG_INF("%s: planning %zu unique llama-bench context(s)\n", __func__, ctx_caps.size()); + for (const auto & ctx_cap : ctx_caps) { + plan_pshard_context(params, ctx_cap.first, ctx_cap.second, true); + } + } else { + auto cparams = common_context_params_to_llama(params); + plan_pshard_context(params, cparams.n_ctx); + } + + LOG_INF("%s: planning complete, registry written next to model file\n", __func__); + + return 0; +} diff --git a/tools/fit-params/fit-params.cpp b/tools/fit-params/fit-params.cpp index 5d897bc46699..2d2abe0ac4f3 100644 --- a/tools/fit-params/fit-params.cpp +++ b/tools/fit-params/fit-params.cpp @@ -7,6 +7,8 @@ #include "log.h" #include +#include +#include #if defined(_MSC_VER) #pragma warning(disable: 4244 4267) // possible loss of data @@ -14,10 +16,26 @@ // satisfies -Wmissing-declarations int llama_fit_params(int argc, char ** argv); +int llama_fit_params_pshard(int argc, char ** argv); int llama_fit_params(int argc, char ** argv) { common_params params; + // pipeline-sharding planner mode: separate argument handling (it pre-parses + // --bench-plan before forwarding the rest), so dispatch before common_init() + for (int i = 1; i < argc; i++) { + if (strcmp(argv[i], "-pshard") == 0 || strcmp(argv[i], "--pshard") == 0) { + std::vector fwd; + fwd.reserve(argc - 1); + for (int j = 0; j < argc; j++) { + if (j != i) { + fwd.push_back(argv[j]); + } + } + return llama_fit_params_pshard((int) fwd.size(), fwd.data()); + } + } + common_init(); if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_FIT_PARAMS)) { diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 29ad352d0cf3..ade0f6dfe428 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -338,6 +338,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 load_mode; @@ -350,6 +351,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; @@ -382,6 +385,7 @@ static const cmd_params cmd_params_defaults = { /* cpu_strict */ { false }, /* poll */ { 50 }, /* n_gpu_layers */ { -1 }, + /* n_gpu_layers_user */ false, /* n_cpu_moe */ { 0 }, /* split_mode */ { LLAMA_SPLIT_MODE_LAYER }, /* load_mode */ { LLAMA_LOAD_MODE_MMAP }, @@ -390,10 +394,12 @@ static const cmd_params cmd_params_defaults = { /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, /* 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 } } }, /* 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, @@ -468,6 +474,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" @@ -521,6 +529,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { params.progress = cmd_params_defaults.progress; params.no_warmup = cmd_params_defaults.no_warmup; params.offline = cmd_params_defaults.offline; + params.n_gpu_layers_user = cmd_params_defaults.n_gpu_layers_user; if (const char * env = getenv("HF_TOKEN")) { params.hf_token = env; @@ -722,6 +731,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } auto p = parse_int_range(argv[i], /*allow_negative=*/true); 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; @@ -911,6 +921,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; @@ -955,7 +976,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { do { if (override_group_span_len == 0) { // Adds an empty override-tensors for an empty span - params.tensor_buft_overrides.push_back({{}}); + params.tensor_buft_overrides.push_back({{nullptr, nullptr, -1}}); if (value[override_group_span_len] == '\0') { value = &value[override_group_span_len]; last_group = true; @@ -1005,13 +1026,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); @@ -1175,6 +1196,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; } @@ -1223,6 +1250,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; @@ -1238,6 +1267,8 @@ struct cmd_params_instance { mparams.main_gpu = main_gpu; mparams.tensor_split = tensor_split.data(); 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()) { @@ -1267,10 +1298,10 @@ struct cmd_params_instance { for (int i = 0; i < n_cpu_moe; ++i) { patterns.push_back(llm_ffn_exps_block_regex(i)); merged.push_back({ patterns.back().c_str(), - ggml_backend_cpu_buffer_type() }); + ggml_backend_cpu_buffer_type(), -1 }); } - merged.push_back({ nullptr, nullptr }); + merged.push_back({ nullptr, nullptr, -1 }); mparams.tensor_buft_overrides = merged.data(); } @@ -1283,6 +1314,9 @@ struct cmd_params_instance { split_mode == other.split_mode && main_gpu == other.main_gpu && tensor_split == other.tensor_split && load_mode == other.load_mode && 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); } @@ -1292,6 +1326,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; @@ -1299,6 +1335,7 @@ struct cmd_params_instance { cparams.embeddings = embeddings; cparams.op_offload = !no_op_offload; cparams.swa_full = false; + cparams.pshard = pshard; return cparams; } @@ -1312,6 +1349,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) @@ -1364,6 +1403,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, }; @@ -1400,6 +1441,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, }; @@ -1436,6 +1479,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, }; @@ -1477,6 +1522,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; @@ -1516,6 +1563,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; @@ -1575,7 +1624,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", "load_mode", "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" }; @@ -1589,11 +1639,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 == "flash_attn") { + field == "fit_target" || field == "fit_min_ctx" || field == "flash_attn" || + field == "max_vram_alloc") { return INT; } if (field == "f16_kv" || field == "no_kv_offload" || field == "cpu_strict" || - field == "embeddings" || field == "no_host") { + field == "embeddings" || field == "no_host" || field == "pshard") { return BOOL; } if (field == "avg_ts" || field == "stddev_ts") { @@ -1672,6 +1723,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), @@ -1865,6 +1918,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); @@ -1902,6 +1961,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"; } @@ -1992,6 +2057,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"); } @@ -2258,8 +2329,21 @@ int llama_bench(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) { + common_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 @@ -2275,19 +2359,22 @@ int llama_bench(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 common_fit_params can adjust them mparams.n_gpu_layers = llama_model_default_params().n_gpu_layers; @@ -2308,28 +2395,66 @@ int llama_bench(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 = common_pshard_registry_create(tier_max_auto, cparams.n_seq_max); + mparams.pshard_registry = pending_pshard_registry; + + common_fit_params_pshard(inst.model.c_str(), &mparams, &cparams, + pshard_overrides.data(), inst.max_vram_alloc, inst.fit_target); + + if (!mparams.pshard) { + common_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) { + common_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); @@ -2478,7 +2603,7 @@ int llama_bench(int argc, char ** argv) { ggml_threadpool_free_fn(threadpool); } - llama_model_free(lmodel); + free_active_model(); if (p) { p->print_footer();