From 2913a33e9a54fd0969e3d3001bf67666f75fcb12 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Fri, 21 Aug 2026 13:34:31 +0200 Subject: [PATCH] fit: also take into account n_streams --- common/fit.cpp | 44 ++++++++++++++++++++++++++++++++------------ 1 file changed, 32 insertions(+), 12 deletions(-) diff --git a/common/fit.cpp b/common/fit.cpp index dd1f3ef766..757b4937c7 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -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(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(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, 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 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 {