Public/koboldcpp
1
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-08-26 15:15:16 +02:00
Code Issues Packages Projects Releases Wiki Activity

fit: also take into account n_streams (#27496)

Browse Source 
* 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 <admin@serveurperso.com>
This commit is contained in:
Xuan-Son Nguyen
2026-08-22 16:16:06 +02:00
committed by GitHub
parent 9fee29e943
commit 2fb989b9e7
7 changed files with 145 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/common.cpp
+23
View File
@@ -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);
}
common/fit.cpp
+105 -17
View File
@@ -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<uint32_t>(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<ggml_backend_dev_t> 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>(uint64_t(hp_nct) * n_streams, UINT32_MAX);
const uint32_t n_ctx_min_total = (uint32_t) std::min<uint64_t>(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<int64_t> 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());
common/fit.h
+11
View File
@@ -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
common/speculative.cpp
+3
View File
@@ -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;