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<bool> 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<double> 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;