diff --git a/common/arg.cpp b/common/arg.cpp index 68920618..fd8f2931 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3630,6 +3630,28 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.sse_ping_interval = value; } ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SSE_PING_INTERVAL")); + add_opt(common_arg( + {"--fair-chunk"}, "N", + "while another request is active, process prompts in steps of N tokens that alternate with the other\n" + "requests' steps by equal time share, instead of one batch mixing decode and prompt tokens\n" + "(-1 = n_ubatch, 0 = off; default: -1)", + [](common_params & params, int value) { + if (value < -1) { + throw std::invalid_argument("fair-chunk must be -1, 0 or positive"); + } + params.fair_chunk = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_FAIR_CHUNK")); + add_opt(common_arg( + {"--fair-decode-weight"}, "F", + string_format("GPU time share of each decoding request relative to a prefilling one, see --fair-chunk (default: %.1f)", (double) params.fair_decode_weight), + [](common_params & params, const std::string & value) { + params.fair_decode_weight = std::stof(value); + if (!(params.fair_decode_weight > 0.0f)) { + throw std::invalid_argument("fair-decode-weight must be positive"); + } + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_FAIR_DECODE_WEIGHT")); add_opt(common_arg( {"--threads-http"}, "N", string_format("number of threads used to process HTTP requests (default: %d)", params.n_threads_http), diff --git a/common/common.h b/common/common.h index a7dc6224..5b288622 100644 --- a/common/common.h +++ b/common/common.h @@ -612,6 +612,8 @@ struct common_params { int32_t timeout_read = 3600; // http read timeout in seconds int32_t timeout_write = timeout_read; // http write timeout in seconds int32_t sse_ping_interval = 30; // SSE ping interval in seconds + int32_t fair_chunk = -1; // prompt tokens per step while other requests are active (-1 = n_ubatch, 0 = off) + float fair_decode_weight = 1.0f; // GPU share of each decoding request relative to a prefilling one int32_t n_threads_http = -1; // number of threads to process HTTP requests (TODO: support threadpool) int32_t n_cache_reuse = 0; // min chunk size to reuse from the cache via KV shifting bool cache_prompt = true; // whether to enable prompt caching diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f37bed58..6016f357 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -68,6 +68,33 @@ enum slot_state { struct server_slot; // forward declaration +// while a prompt is being processed beside another active request, each step runs either the decoding slots +// or one prompt chunk of one slot, never both; the choice goes to the member owed the most GPU time +// (equal processor sharing: every request weighs 1, a decoding one --fair-decode-weight) +struct server_fair { + enum step_kind { STEP_LEGACY, STEP_DECODE, STEP_PROMPT }; + + step_kind kind = STEP_LEGACY; + int id_slot = -1; // the slot of a STEP_PROMPT + + // membership and weights of the step being run + std::vector member; + double w_dec = 0.0; + double w_total = 0.0; + + // seconds of GPU time owed to the decoding slots as a group and to each prompt slot + double owed_dec = 0.0; + std::vector owed; + + int64_t t_step = 0; + + // the current contended window, for the log line when it ends + double t_dec = 0.0; + double t_prompt = 0.0; + int n_dec = 0; + int n_prompt = 0; +}; + struct server_batch { llama_batch batch; bool batch_rendered = false; @@ -153,6 +180,12 @@ struct server_batch { return (int32_t)tokens.size(); } + // drops the tokens from index n on (before render) + void truncate(int32_t n) { + GGML_ASSERT(!batch_rendered && !has_embd && n >= 0 && n <= size()); + tokens.resize(n); + } + void set_output(int32_t idx, bool output) { GGML_ASSERT(idx >= 0 && idx < (int32_t)tokens.size()); tokens[idx].output = output; @@ -916,6 +949,8 @@ private: server_batch batch; + server_fair fair; + llama_model * model_dft = nullptr; llama_context * ctx_dft = nullptr; @@ -1524,6 +1559,13 @@ private: batch.init(std::max(n_batch, params_base.n_parallel), n_embd); } + fair.member.assign(slots.size(), false); + fair.owed.assign(slots.size(), 0.0); + if (params_base.fair_chunk != 0) { + SRV_INF("fair GPU sharing: prompt steps of %d tokens while other requests are active, decode weight %.2f\n", + fair_chunk(), (double) params_base.fair_decode_weight); + } + if (params_base.cache_ram_mib != 0) { if (params_base.cache_ram_mib < 0) { SRV_TRC("prompt cache is enabled, size limit: %s\n", "no limit"); @@ -3452,6 +3494,111 @@ private: }; #endif + // prompt tokens per step while other requests are active + int32_t fair_chunk() const { + return params_base.fair_chunk < 0 ? (int32_t) llama_n_ubatch(ctx_tgt) : params_base.fair_chunk; + } + + // how many of the n prompt tokens gathered for a solo batch a prompt step keeps: the chunk, moved back so the + // trailing tokens llama_decode keeps together stay in one step, as split_equal() would cut that batch + int32_t fair_cut(int32_t n) const { + const int32_t chunk = fair_chunk(); + if (n <= chunk) { + return n; + } + const int32_t n_rs = llama_n_rs_seq(ctx_tgt); + const int32_t n_tail = n_rs > 0 ? n_rs + 1 : 0; + return n - chunk < n_tail ? n - n_tail : chunk; + } + + void fair_end_window() { + if (fair.n_dec + fair.n_prompt > 0) { + SRV_INF("fair: contended window over, decode %.2f s in %d steps, prompt %.2f s in %d steps\n", + fair.t_dec, fair.n_dec, fair.t_prompt, fair.n_prompt); + } + fair.t_dec = fair.t_prompt = 0.0; + fair.n_dec = fair.n_prompt = 0; + fair.owed_dec = 0.0; + std::fill(fair.owed.begin(), fair.owed.end(), 0.0); + std::fill(fair.member.begin(), fair.member.end(), false); + } + + // picks the kind of the next step: STEP_LEGACY unless a prompt is pending beside another active request + void fair_choose() { + fair.kind = server_fair::STEP_LEGACY; + fair.id_slot = -1; + + if (params_base.fair_chunk == 0 || !params_base.cont_batching) { + return; + } + + int n_dec = 0; + int n_prompt = 0; + for (auto & slot : slots) { + const bool prompting = slot.state == SLOT_STATE_STARTED || slot.state == SLOT_STATE_PROCESSING_PROMPT; + n_dec += slot.state == SLOT_STATE_GENERATING; + n_prompt += prompting; + // newcomers start even; a slot that left took its balance with it + if (!prompting) { + fair.owed[slot.id] = 0.0; + } + fair.member[slot.id] = prompting; + } + + if (n_prompt == 0 || n_dec + n_prompt < 2) { + fair_end_window(); + return; + } + if (n_dec == 0) { + fair.owed_dec = 0.0; + } + + fair.w_dec = n_dec * (double) params_base.fair_decode_weight; + fair.w_total = fair.w_dec + n_prompt; + + // the decoders win ties (their steps are short), then the prompt that started first + double best = -1e30; + if (n_dec > 0) { + fair.kind = server_fair::STEP_DECODE; + best = fair.owed_dec; + } + const server_slot * best_slot = nullptr; + for (auto & slot : slots) { + if (!fair.member[slot.id]) { + continue; + } + const double owed = fair.owed[slot.id]; + if (owed > best || (owed == best && best_slot && slot.task->id < best_slot->task->id)) { + best = owed; + best_slot = &slot; + fair.kind = server_fair::STEP_PROMPT; + fair.id_slot = slot.id; + } + } + } + + // charges a step of t seconds: every member earns its share, the one that ran pays for it + void fair_settle(double t) { + const auto clamp = [](double x) { return std::clamp(x, -30.0, 30.0); }; + + const bool ran_dec = fair.kind == server_fair::STEP_DECODE; + if (fair.w_dec > 0.0) { + fair.owed_dec = clamp(fair.owed_dec + t * fair.w_dec / fair.w_total - (ran_dec ? t : 0.0)); + } + for (size_t i = 0; i < fair.owed.size(); i++) { + if (fair.member[i]) { + fair.owed[i] = clamp(fair.owed[i] + t / fair.w_total - ((int) i == fair.id_slot ? t : 0.0)); + } + } + if (ran_dec) { + fair.t_dec += t; + fair.n_dec++; + } else { + fair.t_prompt += t; + fair.n_prompt++; + } + } + void update_slots() { telemetry.note_gpu(); @@ -3495,6 +3642,8 @@ private: } } + fair.t_step = ggml_time_us(); + try { scoped_timer t(t_pre_decode, n_pre_decode); pre_decode(); @@ -3568,6 +3717,12 @@ private: break; // stop any further processing } } + + if (fair.kind != server_fair::STEP_LEGACY) { + // decode is async: charge the step once the GPU is done with it + llama_synchronize(ctx_tgt); + fair_settle((ggml_time_us() - fair.t_step) / 1e6); + } } void pre_decode() { @@ -3638,6 +3793,8 @@ private: // start populating the batch for this iteration batch.clear(); + fair_choose(); + // track if given slot can be batched with slots already in the batch auto & slot_batched = batch.slot_batched; @@ -3646,7 +3803,7 @@ private: // determine which slots are generating and drafting iterate(slots, [&](server_slot & slot) { - if (slot.state != SLOT_STATE_GENERATING) { + if (slot.state != SLOT_STATE_GENERATING || fair.kind == server_fair::STEP_PROMPT) { return; } @@ -3772,7 +3929,7 @@ private: auto & alora_disabled_id = batch.alora_disabled_id; // next, batch any pending prompts without exceeding n_batch - if (params_base.cont_batching || batch.size() == 0) { + if ((params_base.cont_batching || batch.size() == 0) && fair.kind != server_fair::STEP_DECODE) { bool add_ok = true; // false means the batch is full, skip remaining slots iterate(slots, [&](server_slot & slot) { @@ -3780,6 +3937,10 @@ private: return; // batch is full, skip remaining slots } + if (fair.kind == server_fair::STEP_PROMPT && slot.id != fair.id_slot) { + return; + } + if (!slot.is_processing()) { return; } @@ -4215,6 +4376,11 @@ private: } has_mtmd = true; + + // a contended step encodes at most one image, then yields + if (fair.kind == server_fair::STEP_PROMPT) { + break; + } } const auto & spans = slot.task->params.message_spans; @@ -4290,6 +4456,16 @@ private: } } + // a contended step keeps one chunk of what a solo step would have gathered + if (fair.kind == server_fair::STEP_PROMPT) { + const int32_t n_gathered = batch.size() - n_tokens_prev; + const int32_t n_drop = n_gathered - fair_cut(n_gathered); + if (n_drop > 0) { + batch.truncate(batch.size() - n_drop); + slot.prompt.tokens.keep_first(slot.prompt.n_tokens() - n_drop); + } + } + // the number of tokens added to the batch for the current slot const auto n_tokens_cur = batch.size() - n_tokens_prev;