diff --git a/common/arg.cpp b/common/arg.cpp index f1e2bf6908..e346863e51 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -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( diff --git a/common/common.h b/common/common.h index 82fed22092..a333f702ac 100644 --- a/common/common.h +++ b/common/common.h @@ -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. diff --git a/tools/server/README.md b/tools/server/README.md index 07e58fe916..3c2228f343 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -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).
when set without -c/--ctx-size, the shared KV pool is sized to n_parallel*N
(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)
(env: LLAMA_ARG_CTX_CHECKPOINTS) | | `-cms, --checkpoint-min-step N` | minimum spacing between context checkpoints in tokens (default: 8192, 0 = no minimum)
(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)
(env: LLAMA_ARG_CACHE_RAM) | diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index e6c991f7cf..f5477356d6 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -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, diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 5fe2729ba1..22378b38c5 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -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()) {