diff --git a/pkgs-set/pkgs/llama-cpp-bonsai-fairness.patch b/pkgs-set/pkgs/llama-cpp-bonsai-fairness.patch new file mode 100644 index 0000000..0f79fcb --- /dev/null +++ b/pkgs-set/pkgs/llama-cpp-bonsai-fairness.patch @@ -0,0 +1,321 @@ +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; + diff --git a/pkgs-set/pkgs/llama-cpp-bonsai.nix b/pkgs-set/pkgs/llama-cpp-bonsai.nix index a8c60cc..de8d542 100644 --- a/pkgs-set/pkgs/llama-cpp-bonsai.nix +++ b/pkgs-set/pkgs/llama-cpp-bonsai.nix @@ -109,6 +109,11 @@ cudaPackages.backendStdenv.mkDerivation (finalAttrs: { # warp_reduce_sum's pairing and only log2(warps) steps go through smem, with CTA geometry per shape. # Bit-identical to the 1-column kernel; MTP decode +8.7% (two slots verifying +15%). ./llama-cpp-bonsai-decode-verify.patch + # Two-slot fairness: while a prompt is pending beside another active request, each server step is either a + # decode step or one prompt chunk (cut at the solo ubatch boundaries), picked by owed GPU time, so chat keeps + # ~22 t/s during a long prefill instead of ~1. No contention runs the old path unchanged. + # --fair-chunk (default n_ubatch, 0 = off), --fair-decode-weight (default 1). + ./llama-cpp-bonsai-fairness.patch ]; nativeBuildInputs = [