fit: also take into account n_streams

This commit is contained in:
Xuan Son Nguyen
2026-08-21 13:34:31 +02:00
parent 5b6ddc9675
commit 2913a33e9a
+32 -12
View File
@@ -191,10 +191,28 @@ 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;
// 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, redo the estimate with the full context:
if (n_ctx_auto && n_streams > 1) {
cparams->n_ctx = n_ctx_max;
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);
}
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 +325,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,7 +346,7 @@ static void common_params_fit_impl(
}
int64_t sum_projected_used_min_ctx = 0;
cparams->n_ctx = n_ctx_min;
cparams->n_ctx = n_ctx_min_total;
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);
if (nd == 0) {
sum_projected_used_min_ctx = dmds_min_ctx.back().mb.total();
@@ -339,14 +357,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 +375,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 {