Change max context length for auto-fitting with unified KV (#28849)

This commit is contained in:
Gaurav Garg
2026-09-16 16:08:50 +05:30
committed by GitHub
parent 37b53fd454
commit b04d4e567c
+7 -7
View File
@@ -192,9 +192,9 @@ 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);
// size the context for all sequences, but keep minimums and alignment per KV stream
const uint32_t n_seq_max = std::max<uint32_t>(1, cparams->n_seq_max);
const uint32_t n_streams = cparams->kv_unified ? 1 : 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
@@ -264,15 +264,15 @@ static void common_params_fit_impl(
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_max = (uint32_t) std::min<uint64_t>(uint64_t(hp_nct) * n_seq_max, 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);
if (n_seq_max > 1) {
LOG_TRC("%s: context size unset -> using %" PRIu32 " for %" PRIu32 " sequences:\n",
__func__, n_ctx_max, n_seq_max);
dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
}
}