wylab/llama.cpp
0
0
Fork 1
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-06 17:57:44 +02:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* 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 22:39:14 +02:00
committed by GitHub
co-authored by Xuan Son Nguyen
parent 32176338a6
commit 18443257a3
5 changed files with 59 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/arg.cpp
+8
View File
@@ -1643,6 +1643,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
}
}
).set_env("LLAMA_ARG_CTX_SIZE"));
add_opt(common_arg(
{ "--kv-unified-per-slot" }, "N",
"context limit per parallel slot (default: unset, behavior unchanged).\n"
"when set without -c/--ctx-size, the shared KV pool is sized to n_parallel*N",
[](common_params & params, int value) {
params.kv_unified_per_slot = value;
}
).set_env("LLAMA_ARG_KV_UNIFIED_PER_SLOT").set_examples({ LLAMA_EXAMPLE_SERVER }));
add_opt(common_arg(
{"-n", "--predict", "--n-predict"}, "N",
string_format(
common/common.h
+1
View File
@@ -627,6 +627,7 @@ struct common_params {
bool cache_prompt = true; // whether to enable prompt caching
bool cache_idle_slots = true; // save and clear idle slots upon starting a new task
int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot
int32_t kv_unified_per_slot = 0; // max context per parallel slot; 0 = unset
int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints
int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc.
tools/server/README.md
+1
View File
@@ -163,6 +163,7 @@ For the full list of features, please refer to [server's changelog](https://gith
| -------- | ----------- |
| `-lcs, --lookup-cache-static FNAME` | path to static lookup cache to use for lookup decoding (not updated by generation) |
| `-lcd, --lookup-cache-dynamic FNAME` | path to dynamic lookup cache to use for lookup decoding (updated by generation) |
| `--kv-unified-per-slot N` | context limit per parallel slot (default: unset, behavior unchanged).<br/>when set without -c/--ctx-size, the shared KV pool is sized to n_parallel*N<br/>(env: LLAMA_ARG_KV_UNIFIED_PER_SLOT) |
| `-ctxcp, --ctx-checkpoints, --swa-checkpoints N` | max number of context checkpoints to create per slot (default: 32)[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)<br/>(env: LLAMA_ARG_CTX_CHECKPOINTS) |
| `-cms, --checkpoint-min-step N` | minimum spacing between context checkpoints in tokens (default: 8192, 0 = no minimum)<br/>(env: LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT) |
| `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)<br/>(env: LLAMA_ARG_CACHE_RAM) |
tools/server/server-context.cpp
+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,
tools/server/server.cpp
+12
View File
@@ -157,6 +157,18 @@ int llama_server(common_params & params, int argc, char ** argv) {
}
}
// size the KV pool from --kv-unified-per-slot, unless the user pinned it with -c
// or with -c 0 for max context
const bool ctx_pool_auto_sized = params.kv_unified_per_slot > 0 &&
params.n_ctx == 0 &&
(uint32_t) params.fit_params_min_ctx != UINT32_MAX;
if (ctx_pool_auto_sized) {
params.n_ctx = params.n_parallel * params.kv_unified_per_slot;
SRV_INF("--kv-unified-per-slot: sizing KV pool to n_parallel * kv_unified_per_slot = %d * %d = %d\n", params.n_parallel,
params.kv_unified_per_slot, params.n_ctx);
}
// for consistency between server router mode and single-model mode, we set the same model name as alias
auto model_name = params.model.get_name();
if (params.model_alias.empty() && !model_name.empty()) {