From 2fb989b9e79bf4da8159855e24892c8f4c20300f Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Sat, 22 Aug 2026 16:16:06 +0200 Subject: [PATCH] fit: also take into account n_streams (#27496) * fit: also take into account n_streams * server: make the draft context follow the target context With a non-unified KV cache the target context now holds n_ctx_train tokens per sequence, while the draft context was still created with n_ctx = 0 and fell back to n_ctx_train / n_streams per sequence. A slot filled beyond that point makes the draft batch fail to decode, and the server answers 500 on the request. The draft context now takes its size from the target context, so both hold the same number of tokens per sequence. Contexts that share their cells with the target no longer need the kv_size override. The memory reserved for the draft model before fitting is measured at the largest context the target can take, since the draft context grows with the target and a fixed byte margin cannot express that. * fit: take an optional second model into account Illustrates the alternative discussed on the draft context fix. The memory of a draft or MTP context is currently handed to the fit as a fixed byte margin, which cannot express a memory that grows with the context the fit is still deciding on. common_fit_params now takes an optional second model that shares the devices of the main one. Its context follows the main context and its memory is measured again whenever that context changes, so the reduce path stays exact instead of conservative. A model that cannot be measured on its own, such as a shared cell MTP context, is skipped with a warning and the main model is fitted alone. This drops the reservation block in the server, which no longer has to probe the trained context size of the target to guess an upper bound. --------- Co-authored-by: Pascal --- common/common.cpp | 23 ++++++ common/fit.cpp | 122 +++++++++++++++++++++++++----- common/fit.h | 11 +++ common/speculative.cpp | 3 + tools/fit-params/fit-params.cpp | 1 + tools/llama-bench/llama-bench.cpp | 1 + tools/server/server-context.cpp | 57 +------------- 7 files changed, 145 insertions(+), 73 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 25ca838df..d84d57ac9 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1294,11 +1294,34 @@ common_init_result::common_init_result(common_params & params, bool model_only) if (params.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"); + + // the draft context is created from the same base params and follows the main context, fit both together + const bool has_draft = params.speculative.has_dft(); + const bool spec_mtp = std::find(params.speculative.types.begin(), params.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); + + common_params params_dft = common_base_params_to_speculative(params); + + auto mparams_dft = common_model_params_to_llama(params_dft); + auto cparams_dft = common_context_params_to_llama(params_dft); + if (spec_mtp) { + cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + } + cparams_dft.n_rs_seq = 0; + + const common_fit_extra_model extra = { + /*.path_model =*/ params_dft.model.path.c_str(), + /*.mparams =*/ &mparams_dft, + /*.cparams =*/ &cparams_dft, + /*.shares_model =*/ !has_draft, // an MTP context runs on the weights of the main model + }; + common_fit_params(params.model.path.c_str(), &mparams, &cparams, params.tensor_split, params.tensor_buft_overrides.data(), params.fit_params_target.data(), params.fit_params_min_ctx, + has_draft || spec_mtp ? &extra : nullptr, params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); } diff --git a/common/fit.cpp b/common/fit.cpp index dd1f3ef76..c601fe405 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -178,7 +178,7 @@ common_device_memory_data_vec common_get_device_memory_data( static void common_params_fit_impl( const char * path_model, struct llama_model_params * mparams, struct llama_context_params * cparams, float * tensor_split, struct llama_model_tensor_buft_override * tensor_buft_overrides, - size_t * margins_s, uint32_t n_ctx_min, enum ggml_log_level log_level) { + size_t * margins_s, uint32_t n_ctx_min, const common_fit_extra_model * extra, enum ggml_log_level log_level) { if (mparams->split_mode == LLAMA_SPLIT_MODE_TENSOR) { throw common_params_fit_exception("llama_params_fit is not implemented for SPLIT_MODE_TENSOR, abort"); } @@ -191,10 +191,92 @@ static void common_params_fit_impl( uint32_t hp_nct = 0; // hparams.n_ctx_train uint32_t hp_nex = 0; // hparams.n_expert + // with non-unified kv, we need to take into account n_streams + // for example, if memory can hold more than model's trained context size, we must extend the n_ctx to hold enough n_streams + const uint32_t n_streams = cparams->kv_unified ? 1 : std::max(1, cparams->n_seq_max); + const bool n_ctx_auto = cparams->n_ctx == 0; + + dmds_t dmds_extra; // memory of the extra model, laid out on the devices of the main model + uint32_t n_ctx_extra = 0; // context that memory was measured at + + // the extra model competes for the same memory as the main model, add it to every measurement + // its memory is measured again whenever the context it follows changes + auto add_extra_memory = [&](dmds_t & dmds) { + if (extra == nullptr) { + return; + } + + if (dmds_extra.empty() || n_ctx_extra != cparams->n_ctx) { + std::vector devs_extra; + uint32_t ngl_extra = 0; + uint32_t nct_extra = 0; + uint32_t nex_extra = 0; + + extra->cparams->n_ctx = cparams->n_ctx; + + LOG_TRC("%s: getting device memory data for the extra model at a context size of %" PRIu32 ":\n", + __func__, cparams->n_ctx); + + dmds_t measured; + try { + measured = common_get_device_memory_data_impl( + extra->path_model, extra->mparams, extra->cparams, devs_extra, ngl_extra, nct_extra, nex_extra, log_level); + } catch (const std::runtime_error & e) { + // the extra model is optional, fit the main model alone rather than giving up + LOG_WRN("%s: failed to measure the memory of the extra model, fitting without it: %s\n", __func__, e.what()); + dmds_extra = dmds_t(devs.size() + 1); + n_ctx_extra = cparams->n_ctx; + return; + } + + dmds_extra = dmds_t(devs.size() + 1); + dmds_extra.back().mb = measured.back().mb; + for (size_t je = 0; je < devs_extra.size(); je++) { + for (size_t id = 0; id < devs.size(); id++) { + if (devs_extra[je] == devs[id]) { + dmds_extra[id].mb.model += measured[je].mb.model; + dmds_extra[id].mb.context += measured[je].mb.context; + dmds_extra[id].mb.compute += measured[je].mb.compute; + break; + } + } + } + if (extra->shares_model) { + for (llama_device_memory_data & dmd : dmds_extra) { + dmd.mb.model = 0; + } + } + + n_ctx_extra = cparams->n_ctx; + } + + for (size_t id = 0; id < dmds.size(); id++) { + dmds[id].mb.model += dmds_extra[id].mb.model; + dmds[id].mb.context += dmds_extra[id].mb.context; + dmds[id].mb.compute += dmds_extra[id].mb.compute; + } + }; + // step 1: get data for default parameters and check whether any changes are necessary in the first place LOG_TRC("%s: getting device memory data for initial parameters:\n", __func__); - const dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + + // saturate instead of overflowing, this also preserves the UINT32_MAX sentinel of n_ctx_min: + const uint32_t n_ctx_max = (uint32_t) std::min(uint64_t(hp_nct) * n_streams, UINT32_MAX); + const uint32_t n_ctx_min_total = (uint32_t) std::min(uint64_t(n_ctx_min) * n_streams, UINT32_MAX); + + // llama_context would use only hp_nct in total for n_ctx == 0, resolve the context before measuring anything else: + if (n_ctx_auto) { + cparams->n_ctx = n_ctx_max; + if (n_streams > 1) { + LOG_TRC("%s: context size unset and KV cache not unified -> using %" PRIu32 " for %" PRIu32 " sequences:\n", + __func__, n_ctx_max, n_streams); + dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + } + } + add_extra_memory(dmds_full); + const size_t nd = devs.size(); // number of devices std::vector margins; // this function uses int64_t rather than size_t for memory sizes to more conveniently handle deficits @@ -307,8 +389,8 @@ static void common_params_fit_impl( "%s: cannot meet free memory targets on all devices, need to use %" PRId64 " MiB less in total\n", __func__, -global_surplus/MiB); } - if (cparams->n_ctx == 0) { - if (hp_nct > n_ctx_min) { + if (n_ctx_auto) { + if (n_ctx_max > n_ctx_min_total) { int64_t sum_used_target = sum_free; if (nd == 0) { sum_used_target -= margins[0]; @@ -328,8 +410,9 @@ static void common_params_fit_impl( } int64_t sum_projected_used_min_ctx = 0; - cparams->n_ctx = n_ctx_min; - const dmds_t dmds_min_ctx = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + cparams->n_ctx = n_ctx_min_total; + dmds_t dmds_min_ctx = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + add_extra_memory(dmds_min_ctx); if (nd == 0) { sum_projected_used_min_ctx = dmds_min_ctx.back().mb.total(); } else { @@ -339,14 +422,16 @@ static void common_params_fit_impl( } if (sum_used_target > sum_projected_used_min_ctx) { // linear interpolation between minimum and maximum context size: - cparams->n_ctx += (hp_nct - n_ctx_min) * (sum_used_target - sum_projected_used_min_ctx) + cparams->n_ctx += (n_ctx_max - n_ctx_min_total) * (sum_used_target - sum_projected_used_min_ctx) / (sum_projected_used - sum_projected_used_min_ctx); - cparams->n_ctx = std::max(cparams->n_ctx - cparams->n_ctx % 256, n_ctx_min); // round down context for CUDA backend + // round down context for CUDA backend, keep it divisible by the number of streams: + const uint32_t align = 256 * n_streams; + cparams->n_ctx = std::max(cparams->n_ctx - cparams->n_ctx % align, n_ctx_min_total); - const int64_t bytes_per_ctx = (sum_projected_used - sum_projected_used_min_ctx) / (hp_nct - n_ctx_min); - const int64_t memory_reduction = (hp_nct - cparams->n_ctx) * bytes_per_ctx; + const int64_t bytes_per_ctx = (sum_projected_used - sum_projected_used_min_ctx) / (n_ctx_max - n_ctx_min_total); + const int64_t memory_reduction = (n_ctx_max - cparams->n_ctx) * bytes_per_ctx; LOG_TRC("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n", - __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB); + __func__, n_ctx_max, cparams->n_ctx, memory_reduction/MiB); if (nd <= 1) { LOG_TRC("%s: entire model can be fit by reducing context\n", __func__); return; @@ -355,14 +440,14 @@ static void common_params_fit_impl( } else { const int64_t memory_reduction = sum_projected_used - sum_projected_used_min_ctx; LOG_TRC("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n", - __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB); + __func__, n_ctx_max, cparams->n_ctx, memory_reduction/MiB); } } else { if (n_ctx_min == UINT32_MAX) { - LOG_TRC("%s: user has requested full context size of %" PRIu32 " -> no change\n", __func__, hp_nct); + LOG_TRC("%s: user has requested full context size of %" PRIu32 " -> no change\n", __func__, n_ctx_max); } else { LOG_TRC("%s: default model context size is %" PRIu32 " which is <= the min. context size of %" PRIu32 " -> no change\n", - __func__, hp_nct, n_ctx_min); + __func__, n_ctx_max, n_ctx_min_total); } } } else { @@ -507,8 +592,9 @@ static void common_params_fit_impl( llama_model_params mparams_copy = *mparams; set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, mparams_copy); - const dmds_t dmd_nl = common_get_device_memory_data_impl( + dmds_t dmd_nl = common_get_device_memory_data_impl( path_model, &mparams_copy, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + add_extra_memory(dmd_nl); LOG_TRC("%s: memory for test allocation by device:\n", func_name); for (size_t id = 0; id < nd; id++) { @@ -535,8 +621,9 @@ static void common_params_fit_impl( mparams->tensor_buft_overrides = tensor_buft_overrides; LOG_TRC("%s: getting device memory data with all MoE tensors moved to system memory:\n", __func__); - const dmds_t dmds_cpu_moe = common_get_device_memory_data_impl( + dmds_t dmds_cpu_moe = common_get_device_memory_data_impl( path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + add_extra_memory(dmds_cpu_moe); for (size_t id = 0; id < nd; id++) { global_surplus_cpu_moe += dmds_cpu_moe[id].free; @@ -796,11 +883,12 @@ enum common_params_fit_status common_fit_params( llama_model_tensor_buft_override * tensor_buft_overrides, size_t * margins, uint32_t n_ctx_min, + const common_fit_extra_model * extra, ggml_log_level log_level) { const int64_t t0_us = llama_time_us(); common_params_fit_status status = COMMON_PARAMS_FIT_STATUS_SUCCESS; try { - common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, log_level); + common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, extra, log_level); LOG_TRC("%s: successfully fit params to free device memory\n", __func__); } catch (const common_params_fit_exception & e) { LOG_WRN("%s: failed to fit params to free device memory: %s\n", __func__, e.what()); diff --git a/common/fit.h b/common/fit.h index 208fc3069..824d386b0 100644 --- a/common/fit.h +++ b/common/fit.h @@ -11,6 +11,16 @@ enum common_params_fit_status { COMMON_PARAMS_FIT_STATUS_ERROR = 2, // a hard error occurred, e.g. because no model could be found at the specified path }; +// a second model that shares the devices of the main model, e.g. a draft model +// - its context follows the context of the main model, so its memory is measured again whenever that context changes +// - shares_model tells the fit that the weights are already counted in the main model, as for an MTP context +struct common_fit_extra_model { + const char * path_model; + llama_model_params * mparams; + llama_context_params * cparams; + bool shares_model; +}; + // fits mparams and cparams to free device memory (assumes system memory is unlimited) // - returns true if the parameters could be successfully modified to fit device memory // - this function is NOT thread safe because it modifies the global llama logger state @@ -24,6 +34,7 @@ common_params_fit_status common_fit_params( llama_model_tensor_buft_override * tensor_buft_overrides, // writable buffer for overrides, needs at least llama_max_tensor_buft_overrides elements size_t * margins, // margins of memory to leave per device in bytes uint32_t n_ctx_min, // minimum context size to set when trying to reduce memory use + const common_fit_extra_model * extra, // model to fit alongside the main one, nullptr if there is none ggml_log_level log_level); // minimum log level to print during fitting, lower levels go to debug log // print estimated memory to stdout diff --git a/common/speculative.cpp b/common/speculative.cpp index 8461e4107..4eef2212e 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2388,6 +2388,9 @@ common_speculative_init_result::common_speculative_init_result( cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; } + // the draft context holds as many tokens per sequence as the target context + cparams.n_ctx = llama_n_ctx(ctx_tgt); + // note: for small models maybe we can set this to the maximum possible draft from all speculative types // the extra memory for small models is likely negligible? cparams.n_rs_seq = 0; diff --git a/tools/fit-params/fit-params.cpp b/tools/fit-params/fit-params.cpp index 5d897bc46..3e78c8929 100644 --- a/tools/fit-params/fit-params.cpp +++ b/tools/fit-params/fit-params.cpp @@ -33,6 +33,7 @@ int llama_fit_params(int argc, char ** argv) { if (!params.fit_params_print) { const common_params_fit_status status = common_fit_params(params.model.path.c_str(), &mparams, &cparams, params.tensor_split, params.tensor_buft_overrides.data(), params.fit_params_target.data(), params.fit_params_min_ctx, + nullptr, params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); if (status != COMMON_PARAMS_FIT_STATUS_SUCCESS) { LOG_ERR("%s: failed to fit CLI arguments to free memory, exiting...\n", __func__); diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 03d59f08d..a2da93b9a 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -2294,6 +2294,7 @@ int llama_bench(int argc, char ** argv) { fit_overrides.data(), margins.data(), inst.fit_min_ctx, + nullptr, params.verbose ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 1293c8640..36d982832 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1040,62 +1040,7 @@ private: } } - // optionally reserve VRAM for the draft / MTP context before fitting the target model - if (params_base.fit_params) { - if (has_spec) { - // MTP draft context lives on the target model, only context+compute are new - bool measure_model_bytes = has_draft; - - common_params params_dft = common_base_params_to_speculative(params_base); - - auto mparams_dft = common_model_params_to_llama(params_dft); - auto cparams_dft = common_context_params_to_llama(params_dft); - if (spec_mtp) { - cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; - } - cparams_dft.n_rs_seq = 0; - - std::vector devs; - uint32_t hp_ngl = 0; - uint32_t hp_nct = 0; - uint32_t hp_nex = 0; - try { - auto dmd = common_get_device_memory_data( - params_dft.model.path.c_str(), &mparams_dft, &cparams_dft, - devs, hp_ngl, hp_nct, hp_nex, GGML_LOG_LEVEL_ERROR); - - GGML_ASSERT(!params_base.fit_params_target.empty()); - size_t total = 0; - - std::vector tgt_devices = params.devices; - - if (tgt_devices.empty()) { - for(size_t i = 0; i < ggml_backend_dev_count(); ++i) { - tgt_devices.push_back(ggml_backend_dev_get(i)); - } - } - - for (size_t j = 0; j < devs.size(); ++j) { - const size_t bytes = (measure_model_bytes ? dmd[j].model : 0) + dmd[j].context + dmd[j].compute; - total += bytes; - for (size_t i = 0; i < tgt_devices.size(); i++) { - if (tgt_devices[i] == devs[j]) { - SRV_DBG("[spec] adding %.2f MiB to fit_params_target for device %s\n", - bytes / (1024.0 * 1024.0), ggml_backend_dev_name(devs[j])); - params_base.fit_params_target[i] += bytes; - break; - } - } - } - SRV_TRC("[spec] estimated memory usage of %s is %.2f MiB\n", - has_draft ? "draft model" : "MTP context", - total / (1024.0 * 1024.0)); - } catch (const std::exception & e) { - SRV_WRN("[spec] failed to measure %s memory: %s\n", - has_draft ? "draft model" : "MTP context", e.what()); - } - } - } + // note: the draft / MTP context is fitted together with the target model, see common_fit_extra_model // attach a progress callback {