server: add ctx-per-slot (--kv-unified-per-slot) (#24124)

* Add ctx-per-slot argument for unifid KV cache

* Swap out ctx fractions for ctx pool slots

* Formatting cleanup

* Remove ctx-pool-slots, make ctx-per-slot an int

* refactor it

---------

Co-authored-by: Xuan Son Nguyen <son@huggingface.co>
This commit is contained in:
Bartowski
2026-08-27 16:39:14 -04:00
committed by GitHub
parent 32176338a6
commit 18443257a3
5 changed files with 59 additions and 9 deletions
+37 -9
View File
@@ -1208,10 +1208,31 @@ private:
const int n_ctx_train = llama_model_n_ctx_train(model_tgt);
int n_ctx_slot = llama_n_ctx_seq(ctx_tgt);
if (n_ctx_slot > n_ctx_train) {
SRV_WRN("the slot context (%d) exceeds the training context of the model (%d) - capping\n", n_ctx_slot, n_ctx_train);
n_ctx_slot = n_ctx_train;
{
// note: the capping itself is done in n_ctx_slot(), here we only report it
const int n_ctx_seq = llama_n_ctx_seq(ctx_tgt);
if (params_base.kv_unified_per_slot > 0) {
if (n_ctx_seq > params_base.kv_unified_per_slot) {
SRV_INF("capping per-slot context (%d) to --kv-unified-per-slot (%d)\n",
n_ctx_seq, params_base.kv_unified_per_slot);
} else if (params_base.kv_unified_per_slot > n_ctx_seq) {
// cap is above the per-slot pool capacity, so it can never bind
SRV_WRN(
"--kv-unified-per-slot (%d) exceeds the per-slot pool capacity (%d) - cap has no effect, "
"slots are limited to %d (raise the KV pool with -c, or unset -c to size it to "
"n_parallel * kv_unified_per_slot)\n",
params_base.kv_unified_per_slot, n_ctx_seq, n_ctx_seq);
}
}
const int n_ctx_capped = params_base.kv_unified_per_slot > 0 ?
std::min(n_ctx_seq, params_base.kv_unified_per_slot) : n_ctx_seq;
if (n_ctx_capped > n_ctx_train) {
SRV_WRN("the slot context (%d) exceeds the training context of the model (%d) - capping\n",
n_ctx_capped, n_ctx_train);
}
}
slots.clear();
@@ -1227,7 +1248,7 @@ private:
// setup slots
SRV_INF("initializing, n_slots = %d, n_ctx_slot = %d, kv_unified = '%s'\n",
params_base.n_parallel, n_ctx_slot, params_base.kv_unified ? "true" : "false");
params_base.n_parallel, n_ctx_slot(), params_base.kv_unified ? "true" : "false");
// initialize slots
for (int i = 0; i < params_base.n_parallel; i++) {
@@ -1271,7 +1292,7 @@ private:
slot.ctx_dft = ctx_dft;
slot.mem.init(ctx_tgt, ctx_dft);
slot.spec = spec.get();
slot.n_ctx = n_ctx_slot;
slot.n_ctx = n_ctx_slot();
slot.mctx = mctx;
slot.prompt.tokens.has_mtmd = mctx != nullptr;
@@ -3975,8 +3996,15 @@ private:
});
}
int get_slot_n_ctx() {
return slots.back().n_ctx;
// context size of a single slot, capped by --kv-unified-per-slot and by the training context of the model
int n_ctx_slot() const {
int res = llama_n_ctx_seq(ctx_tgt);
if (params_base.kv_unified_per_slot > 0) {
res = std::min(res, params_base.kv_unified_per_slot);
}
return std::min(res, llama_model_n_ctx_train(model_tgt));
}
server_response_reader get_response_reader() {
@@ -4142,7 +4170,7 @@ server_context_meta server_context::get_meta() const {
/* has_inp_audio */ impl->chat_params.allow_audio,
/* has_inp_video */ impl->chat_params.allow_video,
/* json_ui_settings */ impl->json_ui_settings,
/* slot_n_ctx */ impl->get_slot_n_ctx(),
/* slot_n_ctx */ impl->n_ctx_slot(),
/* pooling_type */ llama_pooling_type(impl->ctx_tgt),
/* chat_params */ impl->chat_params,