From c238f827b2f52654bbc5ea7babe94572751c39ad Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 03:28:50 -0700 Subject: [PATCH 1/6] server: add --pipeline-groups to run the slots over several contexts A layer split over two nodes is a two-stage pipeline that a single llama_context feeds one batch at a time, so each stage sits idle while the other one computes. With --pipeline-groups N the server creates N llama_contexts from the one model, partitions its slots between them and gives each group its own batch and its own decode thread, so there are N batches in flight and both stages have work. Each context is created with n_seq_max = n_parallel / N and n_ctx = n_ctx / N, so the per-slot context and the total KV memory are unchanged. Slots are partitioned contiguously and carry the sequence id they use inside their own context. Slot selection for a new task still runs over all slots, so prompt cache similarity, the slot endpoints and the KV prefix reuse behave exactly as before. The model weights, the task queue, the results queue and the HTTP layer are shared. Task processing pauses the decode loops for the moment it looks at the slots. Speculative decoding, multimodal and idle sleeping are refused with N > 1 rather than half supported. With the default N = 1 there is one context, one batch and one update loop on the main thread, no locks and no extra threads. --- tools/server/README.md | 34 ++ tools/server/server-context.cpp | 577 ++++++++++++++++++++++++++------ tools/server/server-context.h | 5 + tools/server/server.cpp | 49 +++ 4 files changed, 556 insertions(+), 109 deletions(-) diff --git a/tools/server/README.md b/tools/server/README.md index 93736c3edfa9..bef9088f39c4 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -2074,6 +2074,40 @@ Note that the following endpoints are exempt from being considered as incoming t - `GET /models` - `GET /metrics` +## Pipeline groups + +`--pipeline-groups N` (default `1`) runs the server's slots over `N` independent `llama_context` +objects created from the same model. Each group has its own batch, its own sampling and its own +decode thread; the model weights, the task queue, the results queue and the HTTP layer are shared. + +This is meant for a layer split across two machines, e.g. + +```sh +llama-server -m model.gguf -c 32768 --parallel 16 \ + --rpc peer:50052 --device CUDA0,RPC0 -sm layer -ngl 99 \ + --pipeline-groups 2 +``` + +With one context, a layer split is a two-stage pipeline that is fed one batch at a time, so each +stage is idle while the other one computes. With two groups there are two batches in flight, so +while group A is being computed on the second stage, group B is being computed on the first one. + +Details: + +- The slots are partitioned contiguously: with `--parallel P` and `--pipeline-groups N`, group `g` + owns slots `[g*P/N, (g+1)*P/N)`. `--parallel` must be a positive multiple of `--pipeline-groups`. +- Each context is created with `n_seq_max = P/N` and `n_ctx = C/N`, so the per-slot context and the + total KV memory over all groups are the same as with a single context. `-c` must be given + explicitly and must be a multiple of `N`. +- Slot selection for an incoming request still runs over *all* slots, so prompt cache similarity and + the slot save / restore endpoints work exactly as before: a returning conversation lands on the + slot that still holds its prefix, whichever group that slot belongs to. +- Task processing briefly pauses the decode loops, so `/slots`, `/metrics` and cancellations are + answered after the in-flight decode of each group finishes rather than during it. +- `N > 1` is refused at startup together with speculative decoding (`--model-draft`, MTP), + multimodal (`--mmproj`) and `--sleep-idle-seconds`. +- With `N = 1` nothing changes: one context, one batch and one update loop on the main thread. + ## More examples ### Interactive mode diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b4..fdfac7121621 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -35,6 +35,12 @@ #include #endif +// used by the --pipeline-groups decode threads +#include +#include +#include +#include + constexpr int HTTP_POLLING_SECONDS = 1; static common_speculative_output_limits server_output_limits(const common_params & params) { @@ -68,7 +74,8 @@ struct server_batch { bool batch_rendered = false; struct token { - int32_t id_slot; + int32_t id_slot; // global slot id, used to attribute stats + int32_t seq_id; // sequence id inside the context of the slot's pipeline group llama_token token; llama_pos pos; bool output; @@ -108,22 +115,22 @@ struct server_batch { tokens.reserve(n_tokens_alloc); } - bool add(int32_t id_slot, llama_token token, llama_pos pos, bool output, bool is_prompt) { + bool add(int32_t id_slot, int32_t seq_id, llama_token token, llama_pos pos, bool output, bool is_prompt) { GGML_ASSERT(!has_embd); // cannot mix tokens + embd in same batch GGML_ASSERT(batch.pos != nullptr); if ((int32_t)tokens.size() >= n_tokens_alloc) { return false; } - tokens.push_back({ id_slot, token, pos, output, is_prompt }); + tokens.push_back({ id_slot, seq_id, token, pos, output, is_prompt }); return true; } - bool add(int32_t id_slot, const std::vector & embd_in, llama_pos pos, bool output, bool is_prompt) { + bool add(int32_t id_slot, int32_t seq_id, const std::vector & embd_in, llama_pos pos, bool output, bool is_prompt) { GGML_ASSERT(batch.pos != nullptr); if ((int32_t)tokens.size() >= n_tokens_alloc) { return false; } - tokens.push_back({ id_slot, LLAMA_TOKEN_NULL, pos, output, is_prompt }); + tokens.push_back({ id_slot, seq_id, LLAMA_TOKEN_NULL, pos, output, is_prompt }); has_embd = true; embd.insert(embd.end(), embd_in.begin(), embd_in.end()); return true; @@ -159,7 +166,7 @@ struct server_batch { common_batch_clear(batch); for (int32_t i = 0; i < size(); i++) { const auto & t = tokens[i]; - common_batch_add(batch, t.token, t.pos, { t.id_slot }, t.output); + common_batch_add(batch, t.token, t.pos, { t.seq_id }, t.output); } if (has_embd) { batch.token = nullptr; // will be restored on clear() @@ -194,6 +201,14 @@ struct server_batch { struct server_slot { int id; + // pipeline group that owns this slot, i.e. the index of ctx_tgt in server_context_impl::groups + // always 0 unless --pipeline-groups > 1 + int id_group = 0; + + // sequence id of this slot inside ctx_tgt / ctx_dft + // equal to id unless --pipeline-groups > 1, where each context only holds n_parallel/N sequences + int seq_id = 0; + llama_context * ctx_tgt = nullptr; llama_context * ctx_dft = nullptr; @@ -255,8 +270,8 @@ struct server_slot { return false; } - const size_t cur_size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); - const size_t cur_size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + const size_t cur_size_tgt = llama_state_seq_get_size_ext(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE); + const size_t cur_size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; const size_t cur_size = cur_size_tgt + cur_size_dft; @@ -268,16 +283,16 @@ struct server_slot { return false; } - llama_state_seq_get_data_ext(ctx_tgt, cur->data.main.data(), cur_size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); + llama_state_seq_get_data_ext(ctx_tgt, cur->data.main.data(), cur_size_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE); if (ctx_dft) { - llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE); + llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE); } return true; } bool prompt_load(server_prompt_cache & prompt_cache, const server_tokens & tokens) { - bool res = prompt_cache.load(prompt, tokens, ctx_tgt, ctx_dft, id); + bool res = prompt_cache.load(prompt, tokens, ctx_tgt, ctx_dft, seq_id); if (!res) { SLT_WRN(*this, "%s", "failed to load prompt from cache\n"); } @@ -288,7 +303,7 @@ struct server_slot { void prompt_clear() { SLT_TRC(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size()); - mem.seq_rm(id, -1, -1); + mem.seq_rm(seq_id, -1, -1); prompt.clear(); } @@ -351,7 +366,7 @@ struct server_slot { n_predict_max = -1; - llama_set_sampler(ctx_tgt, id, nullptr); + llama_set_sampler(ctx_tgt, seq_id, nullptr); // clear alora start alora_invocation_start = -1; @@ -463,9 +478,9 @@ struct server_slot { i_batch = batch.size(); if (!inp_embd.empty()) { - add_ok &= batch.add(id, inp_embd, prompt.tokens.pos_next(), true, false); + add_ok &= batch.add(id, seq_id, inp_embd, prompt.tokens.pos_next(), true, false); } else { - add_ok &= batch.add(id, sampled, prompt.tokens.pos_next(), true, false); + add_ok &= batch.add(id, seq_id, sampled, prompt.tokens.pos_next(), true, false); } SLT_DBG(*this, "slot decode token, id=%d, n_ctx = %d, n_tokens = %d, truncated = %d\n", @@ -483,9 +498,9 @@ struct server_slot { auto pos0 = prompt.tokens.pos_next(); - add_ok &= batch.add(id, sampled, pos0++, true, false); + add_ok &= batch.add(id, seq_id, sampled, pos0++, true, false); for (auto token : spec_draft) { - add_ok &= batch.add(this->id, token, pos0++, true, false); + add_ok &= batch.add(this->id, seq_id, token, pos0++, true, false); } } @@ -676,8 +691,9 @@ struct server_slot { void copy_state_to(server_slot & other) const { GGML_ASSERT(state == SLOT_STATE_DONE_PROMPT); - mem.seq_rm(other.id, -1, -1); - mem.seq_cp(id, other.id, -1, -1); + // note: parent and child slots are always in the same pipeline group, see get_free_slots() + mem.seq_rm(other.seq_id, -1, -1); + mem.seq_cp(seq_id, other.seq_id, -1, -1); other.i_batch = i_batch; @@ -780,6 +796,29 @@ static int process_mtmd_chunk(const server_slot & slot, mtmd::batch_ptr & mbatch return try_decode(); } +// A pipeline group is one llama_context with its own batch, its own decode loop and its own +// contiguous range of slots. With --pipeline-groups 1 (the default) there is exactly one group: +// it owns ctx_tgt and every slot, and its update loop runs on the main thread, as before. +// +// With N > 1 the point is that while group A's batch is being computed on the second stage of a +// layer split (the RPC peer), group B's batch can be computed on the first stage (the local GPU), +// so both devices are busy instead of each idling half of every decode step. +struct server_group { + int id = 0; + + llama_context * ctx = nullptr; + + server_batch batch; + + // slots owned by this group, in slot id order (slots are partitioned contiguously) + std::vector slots; + + // only used when n_groups > 1, all guarded by server_context_impl::mtx_engine + std::thread thread; + bool busy = false; // a decode is in flight, no one may touch ctx + int n_pause_req = 0; // someone wants the engine stopped, do not start a new iteration +}; + // // server_context_impl (private implementation) // @@ -804,6 +843,9 @@ struct server_context_impl { server_state_callback_t callback_state = [](server_state, json) -> void {}; + // number of pipeline groups requested via --pipeline-groups, must be set before load_model() + int n_pipeline_groups_req = 1; + server_context_impl() { mtmd_helper_log_set(common_log_default_callback, nullptr); } @@ -835,7 +877,18 @@ struct server_context_impl { llama_context * ctx_tgt = nullptr; - server_batch batch; + // pipeline groups, see --pipeline-groups and struct server_group + // groups[0]->ctx is always ctx_tgt; n_groups == 1 unless the user asked for more + int n_groups = 1; + std::vector> groups; + + // number of sequences per context, == params_base.n_parallel when n_groups == 1 + int n_seq_per_group = 1; + + // the following are only ever touched when n_groups > 1 + std::mutex mtx_engine; + std::condition_variable cv_engine; + bool groups_stop = false; llama_model * model_dft = nullptr; llama_context * ctx_dft = nullptr; @@ -894,6 +947,15 @@ struct server_context_impl { ctx_dft = nullptr; model_dft = nullptr; + // groups[0]->ctx is owned by llama_init, the rest were created by llama_init_from_model() + for (size_t g = 1; g < groups.size(); ++g) { + if (groups[g]->ctx != nullptr) { + llama_free(groups[g]->ctx); + groups[g]->ctx = nullptr; + } + } + groups.clear(); + llama_init.reset(); ctx_tgt = nullptr; @@ -1048,11 +1110,44 @@ struct server_context_impl { params_base.load_progress_callback_user_data = &load_progress_text; } - llama_init = common_init_from_params(params_base); + // --pipeline-groups: run the slots over N independent contexts of one model, so that the + // stages of a layer split can be busy at the same time. N == 1 is the default and keeps + // every code path below exactly as it was. + n_groups = std::max(1, n_pipeline_groups_req); + + if (n_groups > 1 && !validate_pipeline_groups(params_base, has_spec, has_mmproj)) { + return false; + } + + n_seq_per_group = params_base.n_parallel / n_groups; + + // note: with a single group this reference IS params_base, so nothing changes + common_params params_grp = n_groups > 1 ? params_base : common_params{}; + common_params & params_ctx = n_groups > 1 ? params_grp : params_base; + + if (n_groups > 1) { + // each context gets 1/N of the sequences and 1/N of the total context, so the per-slot + // context (n_ctx / n_seq_max) and the total KV memory over all contexts are unchanged + params_ctx.n_parallel = n_seq_per_group; + params_ctx.n_ctx = params_base.n_ctx / n_groups; + } + + llama_init = common_init_from_params(params_ctx); model_tgt = llama_init->model(); ctx_tgt = llama_init->context(); + if (n_groups > 1) { + // pick up whatever the parameter fitting resolved, but keep the totals the user asked for + const int n_parallel_total = params_base.n_parallel; + const int n_ctx_total = params_base.n_ctx; + + params_base = params_ctx; + + params_base.n_parallel = n_parallel_total; + params_base.n_ctx = n_ctx_total; + } + if (model_tgt == nullptr) { SRV_ERR("failed to load model, '%s'\n", params_base.model.path.c_str()); return false; @@ -1065,7 +1160,38 @@ struct server_context_impl { vocab = llama_model_get_vocab(model_tgt); - n_ctx = llama_n_ctx(ctx_tgt); + // the remaining contexts of the pipeline are created from the same model + { + groups.clear(); + groups.reserve(n_groups); + + for (int g = 0; g < n_groups; ++g) { + groups.emplace_back(new server_group()); + groups[g]->id = g; + } + + groups[0]->ctx = ctx_tgt; + + for (int g = 1; g < n_groups; ++g) { + llama_context_params cparams = common_context_params_to_llama(params_ctx); + + // make sure the extra contexts are identical to the one common_init_from_params made + cparams.n_ctx = llama_n_ctx(ctx_tgt); + cparams.n_seq_max = llama_n_seq_max(ctx_tgt); + + groups[g]->ctx = llama_init_from_model(model_tgt, cparams); + if (groups[g]->ctx == nullptr) { + SRV_ERR("failed to create llama_context for pipeline group %d\n", g); + return false; + } + + SRV_INF("created llama_context for pipeline group %d, n_ctx = %d, n_seq_max = %d\n", + g, (int) llama_n_ctx(groups[g]->ctx), (int) llama_n_seq_max(groups[g]->ctx)); + } + } + + // the total context over all groups, as requested by the user + n_ctx = llama_n_ctx(ctx_tgt) * n_groups; add_bos_token = llama_vocab_get_add_bos(vocab); @@ -1182,7 +1308,9 @@ struct server_context_impl { } // try speculative decoding - if (ctx_tgt_seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_NO) { + // note: a common_speculative is bound to one target context, and its code paths yield to + // the task queue, which only one thread may do - so it is off with several groups + if (ctx_tgt_seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_NO && n_groups == 1) { try { spec.reset(common_speculative_init(params_base.speculative, params_base.n_parallel)); } catch (const std::exception & e) { @@ -1205,10 +1333,17 @@ struct server_context_impl { for (int i = 0; i < params_base.n_parallel; i++) { server_slot & slot = slots[i]; - slot.id = i; - slot.ctx_tgt = ctx_tgt; - slot.ctx_dft = ctx_dft; - slot.mem.init(ctx_tgt, ctx_dft); + // slots are partitioned contiguously: group g owns slots [g*S, (g+1)*S) + server_group & grp = *groups[i / n_seq_per_group]; + + slot.id = i; + slot.id_group = grp.id; + slot.seq_id = i % n_seq_per_group; + slot.ctx_tgt = grp.ctx; + slot.ctx_dft = ctx_dft; + slot.mem.init(grp.ctx, ctx_dft); + + grp.slots.push_back(&slot); slot.spec = spec.get(); slot.n_ctx = n_ctx_slot; @@ -1263,7 +1398,9 @@ struct server_context_impl { { const int32_t n_batch = llama_n_batch(ctx_tgt); const int32_t n_embd = llama_model_n_embd_inp(model_tgt); - batch.init(std::max(n_batch, params_base.n_parallel), n_embd); + for (auto & grp : groups) { + grp->batch.init(std::max(n_batch, n_seq_per_group), n_embd); + } } if (params_base.cache_ram_mib != 0) { @@ -1315,6 +1452,48 @@ struct server_context_impl { return true; } + // refuse everything we cannot make safe with more than one context, rather than half-support it + bool validate_pipeline_groups(const common_params & params, bool has_spec, bool has_mmproj) const { + auto refuse = [](const char * what) { + SRV_ERR("--pipeline-groups > 1 is not supported together with %s\n", what); + return false; + }; + + if (params.n_parallel < n_groups || params.n_parallel % n_groups != 0) { + SRV_ERR("--parallel (%d) must be a positive multiple of --pipeline-groups (%d)\n", + params.n_parallel, n_groups); + return false; + } + + if (params.n_ctx <= 0) { + SRV_ERR("%s", "--pipeline-groups > 1 requires an explicit context size, pass -c N\n"); + return false; + } + + if (params.n_ctx % n_groups != 0) { + SRV_ERR("--ctx-size (%d) must be a multiple of --pipeline-groups (%d)\n", params.n_ctx, n_groups); + return false; + } + + // a common_speculative and its draft context are bound to one target context + if (has_spec) { + return refuse("speculative decoding (--model-draft / MTP)"); + } + + // mtmd_context is bound to one llama_context + if (has_mmproj) { + return refuse("multimodal (--mmproj)"); + } + + // entering / leaving the sleeping state destroys and rebuilds the contexts under the + // running group threads + if (params.sleep_idle_seconds >= 0) { + return refuse("--sleep-idle"); + } + + return true; + } + // unlike load_model(), this is only called once during initialization bool init() { GGML_ASSERT(ctx_tgt != nullptr); @@ -1327,7 +1506,11 @@ struct server_context_impl { return process_single_task(std::move(task), is_yielding); }); queue_tasks.on_update_slots([this]() { - update_slots(); + if (n_groups > 1) { + // each pipeline group runs its own update loop on its own thread + return; + } + update_slots(*groups[0]); }); queue_tasks.on_sleeping_state([this](bool sleeping) { handle_sleeping_state(sleeping); @@ -1424,6 +1607,107 @@ struct server_context_impl { return true; } + // Stops every pipeline group so that the caller can touch slot and context state safely. + // Constructing this is a no-op when there is a single group: the single update loop and the + // task processing then run on the same thread, exactly as before. + struct engine_guard { + server_context_impl * srv = nullptr; + std::unique_lock lk; + + explicit engine_guard(server_context_impl * srv_) { + if (srv_->n_groups <= 1) { + return; + } + + srv = srv_; + lk = std::unique_lock(srv->mtx_engine); + + // ask every group to stop at the start of its next iteration, then wait for the + // decodes that are already in flight + for (auto & grp : srv->groups) { + grp->n_pause_req++; + } + + srv->cv_engine.wait(lk, [&] { + for (auto & grp : srv->groups) { + if (grp->busy) { + return false; + } + } + return true; + }); + } + + ~engine_guard() { + if (srv == nullptr) { + return; + } + + for (auto & grp : srv->groups) { + grp->n_pause_req--; + } + + lk.unlock(); + srv->cv_engine.notify_all(); + } + + engine_guard(const engine_guard &) = delete; + engine_guard & operator=(const engine_guard &) = delete; + }; + + // the decode loop of one pipeline group, only used when n_groups > 1 + void group_loop(server_group & grp) { + while (true) { + { + std::unique_lock lk(mtx_engine); + if (groups_stop) { + return; + } + } + + if (update_slots(grp)) { + continue; + } + + // nothing to do for this group, wait for a task to be assigned to one of its slots + std::unique_lock lk(mtx_engine); + cv_engine.wait_for(lk, std::chrono::milliseconds(5), [&] { return groups_stop; }); + } + } + + void start_groups() { + if (n_groups <= 1) { + return; + } + + groups_stop = false; + + for (auto & grp : groups) { + server_group * g = grp.get(); + g->thread = std::thread([this, g]() { group_loop(*g); }); + } + + SRV_INF("started %d pipeline group decode threads\n", n_groups); + } + + void stop_groups() { + if (n_groups <= 1) { + return; + } + + { + std::unique_lock lk(mtx_engine); + groups_stop = true; + } + cv_engine.notify_all(); + + for (auto & grp : groups) { + if (grp->thread.joinable()) { + grp->thread.join(); + } + } + } + server_slot * get_slot_by_id(int id_slot) { // note: allow id_slot to be out of bounds (wrap around) id_slot = id_slot % slots.size(); @@ -1571,14 +1855,17 @@ struct server_context_impl { // - smarter decision which slot to clear (LRU or longest prompt?) // - move slot to level 2 cache instead of removing? // - instead of purging, try to store and resume later? - bool try_clear_idle_slots() { + bool try_clear_idle_slots(server_group & grp) { bool res = false; if (!params_base.kv_unified) { return res; } - for (auto & slot : slots) { + // only slots of this group, their KV lives in this group's context + for (auto * slot_ptr : grp.slots) { + auto & slot = *slot_ptr; + if (slot.is_processing()) { continue; } @@ -1703,9 +1990,9 @@ struct server_context_impl { // TODO: tmp until backend sampling is fully implemented if (use_backend_sampling) { - llama_set_sampler(ctx_tgt, slot.id, common_sampler_get(slot.smpl.get())); + llama_set_sampler(slot.ctx_tgt, slot.seq_id, common_sampler_get(slot.smpl.get())); } else { - llama_set_sampler(ctx_tgt, slot.id, nullptr); + llama_set_sampler(slot.ctx_tgt, slot.seq_id, nullptr); } SLT_TRC(slot, "sampler chain: %s\n", common_sampler_print(slot.smpl.get()).c_str()); @@ -1893,7 +2180,7 @@ struct server_context_impl { }); } } else { - std::vector cur = get_token_probabilities(ctx_tgt, idx, n_probs_request); + std::vector cur = get_token_probabilities(slot.ctx_tgt, idx, n_probs_request); const size_t max_probs = cur.size(); const size_t n_probs = std::min(max_probs, n_probs_request); @@ -2061,7 +2348,7 @@ struct server_context_impl { std::vector embd_res(n_embd_out, 0.0f); for (int i = 0; i < batch.n_tokens; ++i) { - if (!batch.logits[i] || batch.seq_id[i][0] != slot.id) { + if (!batch.logits[i] || batch.seq_id[i][0] != slot.seq_id) { continue; } @@ -2101,13 +2388,13 @@ struct server_context_impl { res->n_tokens = slot.task->n_tokens(); for (int i = 0; i < batch.n_tokens; ++i) { - if (!batch.logits[i] || batch.seq_id[i][0] != slot.id) { + if (!batch.logits[i] || batch.seq_id[i][0] != slot.seq_id) { continue; } - const float * embd = llama_get_embeddings_seq(ctx_tgt, batch.seq_id[i][0]); + const float * embd = llama_get_embeddings_seq(slot.ctx_tgt, batch.seq_id[i][0]); if (embd == NULL) { - embd = llama_get_embeddings_ith(ctx_tgt, i); + embd = llama_get_embeddings_ith(slot.ctx_tgt, i); } if (embd == NULL) { @@ -2147,9 +2434,13 @@ struct server_context_impl { return true; } - std::vector get_free_slots(size_t n_slots_needed, int exclude_id_slot) { + std::vector get_free_slots(size_t n_slots_needed, int exclude_id_slot, int id_group) { std::vector free_slots; for (auto & slot : slots) { + // the parent copies its KV into the children, so they must live in the same context + if (slot.id_group != id_group) { + continue; + } if (!slot.is_processing() && slot.id != exclude_id_slot) { free_slots.push_back(&slot); } @@ -2244,8 +2535,8 @@ struct server_context_impl { // this is not true for SWA models: https://github.com/ggml-org/llama.cpp/pull/24411#issuecomment-4677983225 cur.update_pos(slot.prompt.n_tokens() - n_tokens_cur, pos_min, pos_max); - cur.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - cur.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + cur.update_tgt(slot.ctx_tgt, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + cur.update_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); // stash the draft's speculative state with the checkpoint common_speculative_get_state(spec.get(), slot.id, cur.data_spec); @@ -2263,6 +2554,10 @@ struct server_context_impl { return false; } + // with more than one group the update loops run on their own threads, so pause them while + // we look at and modify the slots. no-op with a single group. + engine_guard guard(this); + switch (task.type) { case SERVER_TASK_TYPE_COMPLETION: case SERVER_TASK_TYPE_INFILL: @@ -2302,7 +2597,7 @@ struct server_context_impl { if (task.is_parent()) { // try getting free slots for all child tasks size_t n_child_tasks = task.child_tasks.size(); - std::vector child_slots = get_free_slots(n_child_tasks, slot->id); + std::vector child_slots = get_free_slots(n_child_tasks, slot->id, slot->id_group); if (child_slots.size() < n_child_tasks) { SRV_DBG("not enough free slots for child tasks, n_free = %zu, n_children = %zu, defer task, id_task = %d\n", child_slots.size(), n_child_tasks, id_task); queue_tasks.defer(std::move(task)); @@ -2456,7 +2751,7 @@ struct server_context_impl { GGML_ASSERT(packed.size() % sizeof(llama_token) == 0); const size_t nwrite = llama_state_seq_save_file( - ctx_tgt, filepath.c_str(), slot->id, + slot->ctx_tgt, filepath.c_str(), slot->seq_id, reinterpret_cast(packed.data()), packed.size() / sizeof(llama_token)); if (nwrite == 0) { send_error(task, "Unable to save slot", ERROR_TYPE_SERVER); @@ -2500,10 +2795,10 @@ struct server_context_impl { try { size_t n_packed = 0; llama_tokens packed; - nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, nullptr, 0, &n_packed); + nread = llama_state_seq_load_file(slot->ctx_tgt, filepath.c_str(), slot->seq_id, nullptr, 0, &n_packed); if (nread != 0) { packed.resize(std::max(1, n_packed)); - nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, packed.data(), packed.size(), &n_packed); + nread = llama_state_seq_load_file(slot->ctx_tgt, filepath.c_str(), slot->seq_id, packed.data(), packed.size(), &n_packed); } if (nread == 0) { throw std::runtime_error("No available space in KV cache or invalid slot save file"); @@ -2516,7 +2811,7 @@ struct server_context_impl { throw std::runtime_error("Restored prompt does not fit in the slot context"); } - if (!restored.validate(ctx_tgt)) { + if (!restored.validate(slot->ctx_tgt)) { throw std::runtime_error("Invalid tokens in slot save file"); } @@ -2635,11 +2930,11 @@ struct server_context_impl { } } - void abort_all_slots(const std::string & reason) { - for (auto & slot : slots) { - if (slot.is_processing()) { - send_error(slot, reason, ERROR_TYPE_SERVER); - slot.release(); + void abort_all_slots(server_group & grp, const std::string & reason) { + for (auto * slot : grp.slots) { + if (slot->is_processing()) { + send_error(*slot, reason, ERROR_TYPE_SERVER); + slot->release(); } } } @@ -2674,7 +2969,24 @@ struct server_context_impl { }; #endif - void update_slots() { + // runs one iteration of the decode loop of a single pipeline group + // returns true if the group had work to do + bool update_slots(server_group & grp) { + // shadow the single-context members - everything below operates on this group only + auto * ctx_tgt = grp.ctx; + auto & batch = grp.batch; + auto & slots = grp.slots; + + // when there is only one group there is only one thread and this lock is never engaged + std::unique_lock lk; + if (n_groups > 1) { + lk = std::unique_lock(mtx_engine); + cv_engine.wait(lk, [&]{ return groups_stop || grp.n_pause_req == 0; }); + if (groups_stop) { + return false; + } + } + #ifdef DEBUG_TIMINGS static int64_t t_prev = 0; int64_t t_start = ggml_time_us(); @@ -2692,8 +3004,8 @@ struct server_context_impl { { bool all_idle = true; - for (auto & slot : slots) { - if (slot.is_processing()) { + for (auto * slot : slots) { + if (slot->is_processing()) { all_idle = false; break; } @@ -2702,29 +3014,31 @@ struct server_context_impl { if (all_idle) { SRV_TRC("%s", "all slots are idle\n"); - metrics_flush_idle(); + metrics_flush_idle(grp); - return; // skip further processing + return false; // skip further processing - } else { + } else if (n_groups == 1) { SRV_DBG("%s", "posting NEXT_RESPONSE\n"); server_task task(SERVER_TASK_TYPE_NEXT_RESPONSE); task.id = queue_tasks.get_new_id(); queue_tasks.post(std::move(task)); } + // note: with more than one group each group drives its own loop, so there is no need + // to keep the shared task loop spinning } try { scoped_timer t(t_pre_decode, n_pre_decode); - pre_decode(); + pre_decode(grp); batch.render(); } catch (const std::exception & e) { SRV_ERR("pre_decode() failed: %s\n", e.what()); - abort_all_slots("pre_decode() failed: " + std::string(e.what())); + abort_all_slots(grp, "pre_decode() failed: " + std::string(e.what())); // the batch is half-built and not rendered, skip now to avoid UB - return; + return true; } GGML_ASSERT(batch.slot_batched || batch.size() == 0); @@ -2758,7 +3072,7 @@ struct server_context_impl { // TODO @ngxson : maybe handle n_batch == 1 here instead of inside decode() batch_view = batch.get_view(off, n_tokens); - bool ok = decode(n_batch, off, batch_view); + bool ok = decode(grp, lk, n_batch, off, batch_view); #ifdef DEBUG_TIMINGS llama_synchronize(ctx_tgt); #endif @@ -2775,22 +3089,28 @@ struct server_context_impl { } } catch (const std::exception & e) { SRV_ERR("decode() failed: %s\n", e.what()); - abort_all_slots("decode() failed: " + std::string(e.what())); + abort_all_slots(grp, "decode() failed: " + std::string(e.what())); break; // stop any further processing } try { scoped_timer t(t_post_decode, n_post_decode); - post_decode(n_tokens, off, batch_view); + post_decode(grp, n_tokens, off, batch_view); } catch (const std::exception & e) { SRV_ERR("post_decode() failed: %s\n", e.what()); - abort_all_slots("post_decode() failed: " + std::string(e.what())); + abort_all_slots(grp, "post_decode() failed: " + std::string(e.what())); break; // stop any further processing } } + + return true; } - void pre_decode() { + void pre_decode(server_group & grp) { + auto * ctx_tgt = grp.ctx; + auto & batch = grp.batch; + auto & slots = grp.slots; + (void) ctx_tgt; // apply context-shift if needed // TODO: simplify and improve iterate(slots, [&](server_slot & slot) { @@ -2832,8 +3152,8 @@ struct server_context_impl { SLT_WRN(slot, "slot context shift, n_keep = %d, n_left = %d, n_discard = %d\n", n_keep, n_left, n_discard); - slot.mem.seq_rm (slot.id, n_keep , n_keep + n_discard); - slot.mem.seq_add(slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); + slot.mem.seq_rm (slot.seq_id, n_keep , n_keep + n_discard); + slot.mem.seq_add(slot.seq_id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); // add generated tokens to cache // ref: https://github.com/ggml-org/llama.cpp/pull/16818#discussion_r2473269481 @@ -2900,11 +3220,11 @@ struct server_context_impl { slot.spec_ckpt.update_pos( slot.prompt.n_tokens(), - llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), slot.id), - llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id)); + llama_memory_seq_pos_min(llama_get_memory(slot.ctx_tgt), slot.seq_id), + llama_memory_seq_pos_max(llama_get_memory(slot.ctx_tgt), slot.seq_id)); if (use_ckpt_dft) { - slot.spec_ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + slot.spec_ckpt.update_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } slot.spec_prompt = slot.prompt.tokens.get_text_tokens(); @@ -2943,11 +3263,11 @@ struct server_context_impl { if (ctx_dft) { if (use_ckpt_dft) { - ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } - if (!llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, ckpt.pos_max + 1, -1)) { - GGML_ABORT("failed to remove sequence %d\n", slot.id); + if (!llama_memory_seq_rm(llama_get_memory(slot.ctx_dft), slot.seq_id, ckpt.pos_max + 1, -1)) { + GGML_ABORT("failed to remove sequence %d\n", slot.seq_id); } } @@ -2962,7 +3282,7 @@ struct server_context_impl { if (use_ckpt_tgt) { //const int64_t t_start = ggml_time_us(); - ckpt.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_tgt(slot.ctx_tgt, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); //const int64_t t_total = ggml_time_us() - t_start; //printf("checkpoint total: %f ms\n", t_total / 1000.0); @@ -2974,7 +3294,7 @@ struct server_context_impl { } if (use_ckpt_dft) { - ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } } }); @@ -3150,8 +3470,8 @@ struct server_context_impl { const int64_t kv_shift = (int64_t) head_p - (int64_t) head_c; - slot.mem.seq_rm (slot.id, head_p, head_c); - slot.mem.seq_add(slot.id, head_c, head_c + n_match, kv_shift); + slot.mem.seq_rm (slot.seq_id, head_p, head_c); + slot.mem.seq_add(slot.seq_id, head_c, head_c + n_match, kv_shift); for (size_t i = 0; i < n_match; i++) { slot.prompt.tokens.set_token(head_p + i, slot.prompt.tokens[head_c + i]); @@ -3181,9 +3501,9 @@ struct server_context_impl { const auto pos_min_thold = std::max(0, pos_next - n_swa - (has_new_tokens ? 0 : 1)); if (n_past > 0 && n_past <= slot.prompt.n_tokens()) { - const auto pos_min = llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), slot.id); + const auto pos_min = llama_memory_seq_pos_min(llama_get_memory(slot.ctx_tgt), slot.seq_id); if (pos_min == -1) { - SLT_ERR(slot, "n_past = %d, slot.prompt.tokens.size() = %d, seq_id = %d, pos_min = %d\n", n_past, (int) slot.prompt.tokens.size(), slot.id, pos_min); + SLT_ERR(slot, "n_past = %d, slot.prompt.tokens.size() = %d, seq_id = %d, pos_min = %d\n", n_past, (int) slot.prompt.tokens.size(), slot.seq_id, pos_min); GGML_ABORT("pos_min == -1, but n_past > 0 - should not happen: https://github.com/ggml-org/llama.cpp/pull/13833#discussion_r2116181237"); } @@ -3250,8 +3570,8 @@ struct server_context_impl { if (!do_reset) { // restore the context checkpoint - it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + it->load_tgt(slot.ctx_tgt, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + it->load_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); // restore the draft's speculative state common_speculative_set_state(spec.get(), slot.id, it->data_spec); @@ -3325,7 +3645,7 @@ struct server_context_impl { SLT_TRC(slot, "cached n_tokens = %d, memory_seq_rm [%d, end)\n", slot.prompt.n_tokens(), p0); - slot.mem.seq_rm(slot.id, p0, -1); + slot.mem.seq_rm(slot.seq_id, p0, -1); // If using an alora, there may be uncached tokens that come // before the invocation sequence. When this happens, the @@ -3423,7 +3743,7 @@ struct server_context_impl { // embedding requires all tokens in the batch to be output; // MTP also wants logits at every prompt position so the // streaming hook can mirror t_h_nextn into ctx_dft. - add_ok &= batch.add(slot.id, + add_ok &= batch.add(slot.id, slot.seq_id, cur_tok, /* pos = */ slot.prompt.tokens.pos_next(), /* output = */ slot.need_embd(), @@ -3493,8 +3813,8 @@ struct server_context_impl { } } - const auto pos_min = llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), slot.id); - const auto pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id); + const auto pos_min = llama_memory_seq_pos_min(llama_get_memory(slot.ctx_tgt), slot.seq_id); + const auto pos_max = llama_memory_seq_pos_max(llama_get_memory(slot.ctx_tgt), slot.seq_id); // nothing to checkpoint yet // TODO: is this check needed? @@ -3528,7 +3848,11 @@ struct server_context_impl { // returns true = success ; false = retry with smaller batch size // throw std::runtime_error on fatal error - bool decode(int32_t & n_batch, int32_t off, llama_batch & batch_view) { + bool decode(server_group & grp, std::unique_lock & lk, int32_t & n_batch, int32_t off, llama_batch & batch_view) { + auto * ctx_tgt = grp.ctx; + auto & batch = grp.batch; + auto & slots = grp.slots; + SRV_DBG("n_batch (effective) = %d, off = %d\n", n_batch, off); metrics_pre_decode(); @@ -3559,15 +3883,32 @@ struct server_context_impl { has_output |= batch.tokens[i].output; } - // yield to the queue, so we can still handle metrics tasks while decoding - // note: the sync is done here too, so that the wait is also covered by the yield int ret = 0; - queue_tasks.yield_to_queue([&]() { + if (n_groups > 1) { + // release the engine for the duration of the compute - this is the whole point of the + // feature: while this group is on one stage of the layer split, the other group can + // run its own pre_decode / post_decode and submit its batch to the other stage + grp.busy = true; + lk.unlock(); + ret = llama_decode(ctx_tgt, batch_view); if (ret == 0 && has_output) { llama_synchronize(ctx_tgt); } - }); + + lk.lock(); + grp.busy = false; + cv_engine.notify_all(); + } else { + // yield to the queue, so we can still handle metrics tasks while decoding + // note: the sync is done here too, so that the wait is also covered by the yield + queue_tasks.yield_to_queue([&]() { + ret = llama_decode(ctx_tgt, batch_view); + if (ret == 0 && has_output) { + llama_synchronize(ctx_tgt); + } + }); + } if (ret != 0) { { @@ -3593,14 +3934,14 @@ struct server_context_impl { if (!err.empty()) { SRV_ERR("%s off = %d, n_batch = %d, ret = %d\n", err.c_str(), off, n_batch, ret); - for (auto & slot : slots) { - if (slot.is_processing()) { - send_error(slot, err); - slot.release(); + for (auto * slot : slots) { + if (slot->is_processing()) { + send_error(*slot, err); + slot->release(); // note: it's complicated to keep track of how much of the current batch has been // processed before the error occurred, so we simply clear the entire context - slot.prompt_clear(); + slot->prompt_clear(); } } @@ -3610,7 +3951,7 @@ struct server_context_impl { } // retry with half the batch size to try to find a free slot in the KV cache - if (!try_clear_idle_slots()) { + if (!try_clear_idle_slots(grp)) { n_batch /= 2; } @@ -3619,12 +3960,14 @@ struct server_context_impl { return false; // retry with the updated n_batch } else { // success, apply batch metrics - metrics_post_decode(off, batch_view.n_tokens, has_output); + metrics_post_decode(grp, off, batch_view.n_tokens, has_output); } // TODO: avoid restoring the draft context and re-evaluating the drafted tokens when not needed [TAG_SPEC_AVOID_DRAFT_REEVAL] // for now, always re-evaluate for simplicity // ref: https://github.com/ggml-org/llama.cpp/pull/22728#issuecomment-4400925384 + // note: speculative decoding is refused with more than one group, so this always runs on + // the main thread and yield_to_queue() is safe here if (spec) { bool ok = true; queue_tasks.yield_to_queue([&]() { @@ -3640,12 +3983,14 @@ struct server_context_impl { } // handle `n_cmpl > 1` tasks - when the main prompt is processed, activate all child tasks too - for (auto & slot : slots) { + // note: children are always in the same group as the parent, see get_free_slots() + for (auto * slot_ptr : slots) { + auto & slot = *slot_ptr; if (slot.state == SLOT_STATE_DONE_PROMPT && slot.task->is_parent()) { std::vector children; - for (auto & other : slots) { - if (other.state == SLOT_STATE_WAIT_OTHER && slot.task->id == other.task->id_parent) { - children.push_back(&other); + for (auto * other : slots) { + if (other->state == SLOT_STATE_WAIT_OTHER && slot.task->id == other->task->id_parent) { + children.push_back(other); } } @@ -3665,7 +4010,9 @@ struct server_context_impl { return true; } - void post_decode(int32_t n_batch_tokens, int32_t off, llama_batch & batch_view) { + void post_decode(server_group & grp, int32_t n_batch_tokens, int32_t off, llama_batch & batch_view) { + auto & slots = grp.slots; + // for checking if a given batch index is inside batch_view auto is_inside_view = [&](int32_t idx) { return idx >= off && idx < off + n_batch_tokens; @@ -3821,13 +4168,13 @@ struct server_context_impl { SLT_DBG(slot, "restoring speculative checkpoint (pos_min = %d, pos_max = %d, size = %zu)\n", ckpt.pos_min, ckpt.pos_max, ckpt.size()); - ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_tgt(slot.ctx_tgt, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); if (slot.ctx_dft) { - ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } - slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1); + slot.mem.seq_rm(slot.seq_id, ckpt.pos_max + 1, -1); slot.prompt.tokens.keep_first(ckpt.n_tokens); common_sampler_copy(smpl_save.get(), slot.smpl.get()); @@ -3874,7 +4221,7 @@ struct server_context_impl { slot.sampled = ids.back(); // last accepted token SLT_DBG(slot, "add accepted tokens: sampled=%d, ids.size=%zu, n_draft=%zu\n", slot.sampled, ids.size(), n_draft); - slot.mem.seq_rm(slot.id, slot.prompt.tokens.pos_next(), -1); + slot.mem.seq_rm(slot.seq_id, slot.prompt.tokens.pos_next(), -1); for (size_t i = 0; i < ids.size(); ++i) { completion_token_output result; @@ -3940,7 +4287,9 @@ struct server_context_impl { } // has_output is computed by the caller, which also already synchronized the context if it is set - void metrics_post_decode(int32_t off, int32_t n_tokens, bool has_output) { + void metrics_post_decode(server_group & grp, int32_t off, int32_t n_tokens, bool has_output) { + auto & batch = grp.batch; + metrics.n_decode++; for (const auto & slot : slots) { if (slot.is_processing()) { @@ -3989,12 +4338,12 @@ struct server_context_impl { } // flush any queued prompt metrics if all slots are now idle - void metrics_flush_idle() { + void metrics_flush_idle(server_group & grp) { if (n_prompt_queued == 0) { return; } - llama_synchronize(ctx_tgt); + llama_synchronize(grp.ctx); metrics_flush_prompt(); } @@ -4035,7 +4384,17 @@ bool server_context::load_model(common_params & params) { void server_context::start_loop() { auto & params = impl->params_base; + + // no-op unless --pipeline-groups > 1 + impl->start_groups(); + impl->queue_tasks.start_loop(params.sleep_idle_seconds * 1000); + + impl->stop_groups(); +} + +void server_context::set_pipeline_groups(int n_groups) { + impl->n_pipeline_groups_req = n_groups; } void server_context::terminate() { diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 5d464b8e8cb7..315c96faa85f 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -110,6 +110,11 @@ struct server_context { // note: must be set before load_model() is called void set_state_callback(server_state_callback_t callback); + + // number of pipeline groups, i.e. independent llama_contexts over the one model, each with its + // own slots, batch and decode thread (--pipeline-groups, default 1) + // note: must be set before load_model() is called + void set_pipeline_groups(int n_groups); }; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 5fe2729ba1b2..45c8d7f7005b 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -14,8 +14,11 @@ #include #include +#include +#include #include #include +#include #include // for std::thread::hardware_concurrency #if defined(_WIN32) @@ -25,6 +28,48 @@ static std::function shutdown_handler; static std::atomic_flag is_terminating = ATOMIC_FLAG_INIT; +// --pipeline-groups N: run the slots over N independent llama_contexts of the same model, each +// with its own batch and decode thread. Useful with a layer split over two nodes (--rpc), where a +// single context leaves each stage idle for half of every decode step. +// The option is parsed here instead of in common/arg.cpp because it only means anything for the +// server; everything it changes lives under tools/server. +static int g_pipeline_groups = 1; + +static void server_take_pipeline_groups(int & argc, char ** argv) { + static const char * opt = "--pipeline-groups"; + const size_t opt_len = strlen(opt); + + int n_kept = 1; + + for (int i = 1; i < argc; i++) { + const std::string arg = argv[i]; + + if (arg == opt) { + if (i + 1 >= argc) { + fprintf(stderr, "error: %s requires a value\n", opt); + exit(1); + } + g_pipeline_groups = std::atoi(argv[++i]); + continue; + } + + if (arg.size() > opt_len + 1 && arg.compare(0, opt_len, opt) == 0 && arg[opt_len] == '=') { + g_pipeline_groups = std::atoi(arg.c_str() + opt_len + 1); + continue; + } + + argv[n_kept++] = argv[i]; + } + + argc = n_kept; + argv[n_kept] = nullptr; + + if (g_pipeline_groups < 1) { + fprintf(stderr, "error: %s must be >= 1\n", opt); + exit(1); + } +} + static inline void signal_handler(int signal) { if (is_terminating.test_and_set()) { // in case it hangs, we can force terminate the server by hitting Ctrl+C twice @@ -96,6 +141,9 @@ int llama_server(int argc, char ** argv) { // own arguments required by this example common_params params; + // strip the server-only --pipeline-groups before the common parser sees it + server_take_pipeline_groups(argc, argv); + common_init(); // start the stream session manager GC right after common init, before any HTTP route can @@ -168,6 +216,7 @@ int llama_server(common_params & params, int argc, char ** argv) { // struct that contains llama context and inference server_context ctx_server; + ctx_server.set_pipeline_groups(g_pipeline_groups); server_http_context ctx_http; if (!ctx_http.init(params)) { From f96bb35ffa0a6cf45996e60b70315ba918fee282 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 04:18:18 -0700 Subject: [PATCH 2/6] server: harden the pipeline group decode loop - the unlock around llama_decode is now RAII, so a throwing decode cannot leave a group marked busy (which would wedge every later task) nor return to the error handler without the engine lock - n_cmpl is rejected when it exceeds the slots of one group, instead of being deferred forever: the child slots take their KV from the parent, so they have to live in the parent's context - refuse --control-vector with more than one group, common_init_from_params only applies it to the context it creates - the queued prompt stats and the empty batch kill switch move into the group, they were shared counters flushed per group - post_decode uses the group's context, and the detokenize calls in the result path use the slot's own context - free the contexts already created if a later one fails, and do not index groups[0] when no model is loaded --- tools/server/server-context.cpp | 115 +++++++++++++++++++++----------- 1 file changed, 76 insertions(+), 39 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index fdfac7121621..5dc05194957e 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -813,6 +813,14 @@ struct server_group { // slots owned by this group, in slot id order (slots are partitioned contiguously) std::vector slots; + // queued prompt stats - llama_decode() is async, so the timing is only valid after a sync + // note: kept out of server_metrics, which is copied as-is into the task result + int64_t t_decode_start = 0; // start of the last submitted decode of this group + int64_t t_prompt_start = 0; // start of the oldest queued prompt decode of this group + uint64_t n_prompt_queued = 0; + + int n_empty_consecutive = 0; + // only used when n_groups > 1, all guarded by server_context_impl::mtx_engine std::thread thread; bool busy = false; // a decode is in flight, no one may touch ctx @@ -915,18 +923,10 @@ struct server_context_impl { int slots_debug = 0; // env: LLAMA_SERVER_SLOTS_DEBUG int slots_n_diff = 0; // env: LLAMA_SERVER_SLOTS_N_DIFF - int n_empty_consecutive = 0; - std::unique_ptr prompt_cache; server_metrics metrics; - // queued prompt stats - llama_decode() is async, so the timing is only valid after a sync - // note: kept out of server_metrics, which is copied as-is into the task result - int64_t t_decode_start = 0; // start of the last submitted decode - int64_t t_prompt_start = 0; // start of the oldest queued prompt decode - uint64_t n_prompt_queued = 0; - json json_ui_settings = json::object(); // Necessary similarity of prompt for slot selection @@ -1182,6 +1182,11 @@ struct server_context_impl { groups[g]->ctx = llama_init_from_model(model_tgt, cparams); if (groups[g]->ctx == nullptr) { SRV_ERR("failed to create llama_context for pipeline group %d\n", g); + for (int j = 1; j < g; ++j) { + llama_free(groups[j]->ctx); + groups[j]->ctx = nullptr; + } + groups.clear(); return false; } @@ -1485,6 +1490,12 @@ struct server_context_impl { return refuse("multimodal (--mmproj)"); } + // common_init_from_params() applies the control vector to the context it creates and only + // to that one, so the extra contexts would silently run without it + if (!params.control_vectors.empty()) { + return refuse("--control-vector"); + } + // entering / leaving the sleeping state destroys and rebuilds the contexts under the // running group threads if (params.sleep_idle_seconds >= 0) { @@ -1510,6 +1521,9 @@ struct server_context_impl { // each pipeline group runs its own update loop on its own thread return; } + if (groups.empty()) { + return; // no model loaded + } update_slots(*groups[0]); }); queue_tasks.on_sleeping_state([this](bool sleeping) { @@ -2011,7 +2025,7 @@ struct server_context_impl { : SLOT_STATE_STARTED; // reset server kill-switch counter - n_empty_consecutive = 0; + groups[slot.id_group]->n_empty_consecutive = 0; SLT_INF(slot, "processing task, is_child = %d\n", slot.task->is_child()); return true; @@ -2175,7 +2189,7 @@ struct server_context_impl { result.probs.push_back({ cur_p->data[i].id, - common_token_to_piece(ctx_tgt, cur_p->data[i].id, special), + common_token_to_piece(slot.ctx_tgt, cur_p->data[i].id, special), cur_p->data[i].p }); } @@ -2198,7 +2212,7 @@ struct server_context_impl { for (size_t i = 0; i < n_probs; i++) { result.probs.push_back({ cur[i].id, - common_token_to_piece(ctx_tgt, cur[i].id, special), + common_token_to_piece(slot.ctx_tgt, cur[i].id, special), cur[i].p }); } @@ -2295,7 +2309,7 @@ struct server_context_impl { res->tokens = std::move(slot.generated_tokens); } res->stats = slot.stats; - res->prompt = slot.task->tokens.detokenize(ctx_tgt, true); + res->prompt = slot.task->tokens.detokenize(slot.ctx_tgt, true); res->response_fields = std::move(slot.task->params.response_fields); res->truncated = slot.truncated; @@ -2318,7 +2332,7 @@ struct server_context_impl { // populate res.probs_output if (slot.task->params.sampling.n_probs > 0) { if (!slot.task->params.stream && slot.stop == STOP_TYPE_WORD) { - const llama_tokens stop_word_toks = common_tokenize(ctx_tgt, slot.stopping_word, false); + const llama_tokens stop_word_toks = common_tokenize(slot.ctx_tgt, slot.stopping_word, false); size_t safe_offset = std::min(slot.generated_token_probs.size(), stop_word_toks.size()); res->probs_output = std::vector( @@ -2597,6 +2611,15 @@ struct server_context_impl { if (task.is_parent()) { // try getting free slots for all child tasks size_t n_child_tasks = task.child_tasks.size(); + // the children take their KV from the parent, so they must fit in the + // parent's group. with a single group this is the limit the request + // schema already enforces, so nothing changes there. + if ((int) n_child_tasks + 1 > n_seq_per_group) { + send_error(task, string_format( + "n_cmpl must not exceed the number of slots per pipeline group (%d)", n_seq_per_group), + ERROR_TYPE_INVALID_REQUEST); + break; + } std::vector child_slots = get_free_slots(n_child_tasks, slot->id, slot->id_group); if (child_slots.size() < n_child_tasks) { SRV_DBG("not enough free slots for child tasks, n_free = %zu, n_children = %zu, defer task, id_task = %d\n", child_slots.size(), n_child_tasks, id_task); @@ -3691,7 +3714,7 @@ struct server_context_impl { // process the mtmd chunk // note: it submits its own decode, potentially be async // so the timing is queued and flushed on the next sync - metrics_pre_decode(); + metrics_pre_decode(grp); // encode on the worker thread, so we can still handle metrics tasks size_t n_tokens_out = 0; @@ -3707,7 +3730,7 @@ struct server_context_impl { return; // the slot is done, skip it entirely } - metrics_queue_prompt(n_tokens_out); + metrics_queue_prompt(grp, n_tokens_out); slot.stats.n_prompt_processed += n_tokens_out; slot.stats.update_prompt_last(); @@ -3855,18 +3878,18 @@ struct server_context_impl { SRV_DBG("n_batch (effective) = %d, off = %d\n", n_batch, off); - metrics_pre_decode(); + metrics_pre_decode(grp); if (batch.size() == 0) { SRV_WRN("%s", "no tokens to decode\n"); - if (++n_empty_consecutive > 3) { + if (++grp.n_empty_consecutive > 3) { GGML_ABORT("fatal error - please provide logs and repro in %s\n", "https://github.com/ggml-org/llama.cpp/pull/20277"); } return true; // nothing to decode } else { - n_empty_consecutive = 0; + grp.n_empty_consecutive = 0; } // TODO @ngxson : dft model may have different n_embd than the tgt model, so we check & reject if that's the case @@ -3888,17 +3911,28 @@ struct server_context_impl { // release the engine for the duration of the compute - this is the whole point of the // feature: while this group is on one stage of the layer split, the other group can // run its own pre_decode / post_decode and submit its batch to the other stage - grp.busy = true; - lk.unlock(); + // note: RAII, so a throwing decode cannot leave the group marked busy forever, nor + // return to the caller's error handling without the engine lock held + struct decode_window { + server_context_impl * srv; + server_group * grp; + std::unique_lock * lk; + decode_window(server_context_impl * srv, server_group * grp, std::unique_lock * lk) + : srv(srv), grp(grp), lk(lk) { + grp->busy = true; + lk->unlock(); + } + ~decode_window() { + lk->lock(); + grp->busy = false; + srv->cv_engine.notify_all(); + } + } window(this, &grp, &lk); ret = llama_decode(ctx_tgt, batch_view); if (ret == 0 && has_output) { llama_synchronize(ctx_tgt); } - - lk.lock(); - grp.busy = false; - cv_engine.notify_all(); } else { // yield to the queue, so we can still handle metrics tasks while decoding // note: the sync is done here too, so that the wait is also covered by the yield @@ -4011,7 +4045,10 @@ struct server_context_impl { } void post_decode(server_group & grp, int32_t n_batch_tokens, int32_t off, llama_batch & batch_view) { - auto & slots = grp.slots; + // shadow the single-context members, as update_slots() does + auto * ctx_tgt = grp.ctx; + auto & slots = grp.slots; + (void) ctx_tgt; // for checking if a given batch index is inside batch_view auto is_inside_view = [&](int32_t idx) { @@ -4262,28 +4299,28 @@ struct server_context_impl { // // call before submitting a decode, so that the queued prompt stats can be timed - void metrics_pre_decode() { - t_decode_start = ggml_time_us(); + void metrics_pre_decode(server_group & grp) { + grp.t_decode_start = ggml_time_us(); } // the batch is submitted, but its compute may not be done yet - void metrics_queue_prompt(uint64_t n_tokens) { + void metrics_queue_prompt(server_group & grp, uint64_t n_tokens) { if (n_tokens == 0) { return; } - if (n_prompt_queued == 0) { - t_prompt_start = t_decode_start; + if (grp.n_prompt_queued == 0) { + grp.t_prompt_start = grp.t_decode_start; } - n_prompt_queued += n_tokens; + grp.n_prompt_queued += n_tokens; } // call only after the context is synchronized, otherwise the time is meaningless - void metrics_flush_prompt() { - if (n_prompt_queued == 0) { + void metrics_flush_prompt(server_group & grp) { + if (grp.n_prompt_queued == 0) { return; } - metrics.add_prompt(n_prompt_queued, ggml_time_us() - t_prompt_start); - n_prompt_queued = 0; + metrics.add_prompt(grp.n_prompt_queued, ggml_time_us() - grp.t_prompt_start); + grp.n_prompt_queued = 0; } // has_output is computed by the caller, which also already synchronized the context if it is set @@ -4318,11 +4355,11 @@ struct server_context_impl { } } - metrics_queue_prompt(n_prompt_tokens); + metrics_queue_prompt(grp, n_prompt_tokens); if (has_output) { // the context is already synchronized, so the timings are correct - metrics_flush_prompt(); + metrics_flush_prompt(grp); } // advance the prompt timing of the slots that had tokens in this batch @@ -4339,12 +4376,12 @@ struct server_context_impl { // flush any queued prompt metrics if all slots are now idle void metrics_flush_idle(server_group & grp) { - if (n_prompt_queued == 0) { + if (grp.n_prompt_queued == 0) { return; } llama_synchronize(grp.ctx); - metrics_flush_prompt(); + metrics_flush_prompt(grp); } void metrics_on_prediction(const server_slot & slot) { From 2bd135942fbdb3aa1d583dd9a826467095db0e7b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 09:54:11 -0700 Subject: [PATCH 3/6] rpc: serialise the client connection and track the stored graph per connection One socket is cached per endpoint and is therefore shared by every backend of that endpoint, including the backends of different llama_contexts. A message is written as three unlocked send_data calls, so two contexts interleave their command streams and the server sees a malformed request within seconds. Make a whole message atomic on the wire, and hand the responses out in request order with a ticket, so a thread waiting for its response does not hold the send lock and the other contexts can keep submitting. last_graph_uid was kept per endpoint device while the graph it refers to is stored by the server per connection, and it was read and written without a lock, so two contexts on one connection could make RPC_CMD_GRAPH_RECOMPUTE re-run the other one's graph. Track it per connection and device and check it under the send lock. server: pause only the group that owns the slot a task touches process_single_task stopped every pipeline group for every task and waited for all the in-flight decodes. Holding the engine is already enough to keep the slot state stable, so only wait for the group whose context the task touches: the owning group for completions, cancel, control and slot save / restore / erase, every group for --cache-idle-slots and SET_LORA, none for metrics, /slots and get-lora. --- ggml/src/ggml-rpc/ggml-rpc.cpp | 89 +++++++++++++++++++++++++++------ ggml/src/ggml-rpc/transport.h | 24 +++++++++ tools/server/server-context.cpp | 60 ++++++++++++++++++++-- 3 files changed, 153 insertions(+), 20 deletions(-) diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 69a8a08ae172..ad7395906740 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -212,7 +212,7 @@ struct ggml_backend_rpc_device_context { uint32_t device; std::string name; std::string description; - uint64_t last_graph_uid; + // note: the uid of the last graph stored on the server is tracked per connection, see socket_t }; struct ggml_backend_rpc_buffer_type_context { @@ -300,7 +300,8 @@ static bool parse_endpoint(const std::string & endpoint, std::string & host, int // RPC request : | rpc_cmd (1 byte) | request_size (8 bytes) | request_data (request_size bytes) | // No response -static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size) { +// writes one whole message; the caller must hold sock->conn.mtx_send +static bool send_rpc_cmd_locked(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size) { uint8_t cmd_byte = cmd; if (!sock->send_data(&cmd_byte, sizeof(cmd_byte))) { return false; @@ -314,12 +315,61 @@ static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, return sock->flush(); } +static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size) { + std::lock_guard lock(sock->conn.mtx_send); + return send_rpc_cmd_locked(sock, cmd, input, input_size); +} + +// Reserves this thread's place in the response order of a connection. The server answers the +// commands of one connection strictly in the order it received them, so the n-th response +// belongs to the n-th response-bearing request that was written to the socket. The ticket is +// taken while mtx_send is still held by the sender, and always released, so a failed send +// cannot leave the later waiters stuck. +struct rpc_response_ticket { + rpc_conn_state & conn; + uint64_t seq; + + // must be constructed with conn.mtx_send held + explicit rpc_response_ticket(rpc_conn_state & conn) : conn(conn) { + std::lock_guard lock(conn.mtx_seq); + seq = conn.seq_next++; + } + + void wait() { + std::unique_lock lock(conn.mtx_seq); + conn.cv_seq.wait(lock, [this] { return conn.seq_serving == seq; }); + } + + ~rpc_response_ticket() { + std::lock_guard lock(conn.mtx_seq); + conn.seq_serving = seq + 1; + conn.cv_seq.notify_all(); + } +}; + // RPC request : | rpc_cmd (1 byte) | request_size (8 bytes) | request_data (request_size bytes) | // RPC response: | response_size (8 bytes) | response_data (response_size bytes) | static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size, void * output, size_t output_size) { - if (!send_rpc_cmd(sock, cmd, input, input_size)) { + std::unique_ptr ticket; + bool failed = false; + { + std::lock_guard lock(sock->conn.mtx_send); + ticket.reset(new rpc_response_ticket(sock->conn)); + if (!send_rpc_cmd_locked(sock, cmd, input, input_size)) { + // still take our turn, so the ticket is released in order and no later waiter is + // woken with a response that is not theirs + failed = true; + } + } + + if (failed) { + ticket->wait(); return false; } + + // the response is read outside mtx_send, so the other threads can keep submitting + ticket->wait(); + uint64_t out_size; if (!sock->recv_data(&out_size, sizeof(out_size))) { return false; @@ -731,21 +781,31 @@ static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, g ggml_backend_rpc_device_context * rpc_dev_ctx = (ggml_backend_rpc_device_context *)rpc_dev->context; GGML_ASSERT(cgraph->n_nodes > 0); - bool reuse = cgraph->uid != 0 && rpc_dev_ctx->last_graph_uid == cgraph->uid; - if (reuse) { + GGML_UNUSED(rpc_dev_ctx); + + auto sock = get_socket(rpc_ctx->endpoint); + + // The graph stored by RPC_CMD_GRAPH_COMPUTE lives on the server per connection and device, + // and one connection is shared by every backend of this endpoint - including the backends of + // other llama_contexts. So the uid of the last graph sent has to be tracked per connection, + // and the check has to happen under the same lock as the send, or a RECOMPUTE could re-run + // the graph another context stored in between. + std::unique_lock lock(sock->conn.mtx_send); + + auto & last_uid = sock->conn.last_graph_uid[rpc_ctx->device]; + if (cgraph->uid != 0 && last_uid == cgraph->uid) { rpc_msg_graph_recompute_req request; request.device = rpc_ctx->device; - auto sock = get_socket(rpc_ctx->endpoint); - bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_RECOMPUTE, &request, sizeof(request)); - RPC_STATUS_ASSERT(status); - } else { - rpc_dev_ctx->last_graph_uid = cgraph->uid; - std::vector input; - serialize_graph(rpc_ctx->device, cgraph, input); - auto sock = get_socket(rpc_ctx->endpoint); - bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size()); + bool status = send_rpc_cmd_locked(sock, RPC_CMD_GRAPH_RECOMPUTE, &request, sizeof(request)); RPC_STATUS_ASSERT(status); + return GGML_STATUS_SUCCESS; } + + last_uid = cgraph->uid; + std::vector input; + serialize_graph(rpc_ctx->device, cgraph, input); + bool status = send_rpc_cmd_locked(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size()); + RPC_STATUS_ASSERT(status); return GGML_STATUS_SUCCESS; } @@ -2044,7 +2104,6 @@ ggml_backend_reg_t ggml_backend_rpc_add_server(const char * endpoint) { /* .device = */ ind, /* .name = */ dev_name, /* .description = */ dev_desc, - /* .last_graph_uid = */ 0, }; ggml_backend_dev_t dev = new ggml_backend_device { diff --git a/ggml/src/ggml-rpc/transport.h b/ggml/src/ggml-rpc/transport.h index 3f747ecffd97..2705befbc19b 100644 --- a/ggml/src/ggml-rpc/transport.h +++ b/ggml/src/ggml-rpc/transport.h @@ -1,8 +1,11 @@ #pragma once +#include #include #include #include +#include +#include struct socket_t; typedef std::shared_ptr socket_ptr; @@ -10,9 +13,30 @@ typedef std::shared_ptr socket_ptr; static constexpr size_t MAX_CHUNK_SIZE = 1024ull * 1024ull * 1024ull; // 1 GiB static constexpr size_t RPC_CONN_CAPS_SIZE = 24; +// State shared by every client thread that uses one connection. A connection is looked up by +// endpoint and is therefore shared by all backends of that endpoint, including the backends of +// different llama_contexts, so all of it has to be serialised: +// - mtx_send makes a whole RPC message atomic on the wire +// - seq_* hands the responses out in request order (the server answers strictly in order), +// without holding mtx_send while waiting, so another thread can keep submitting work +// - last_graph_uid mirrors the server's per-connection stored graph for a device, so that +// RPC_CMD_GRAPH_RECOMPUTE can never re-run a graph submitted by another context +struct rpc_conn_state { + std::mutex mtx_send; + std::mutex mtx_seq; + std::condition_variable cv_seq; + uint64_t seq_next = 0; + uint64_t seq_serving = 0; + + std::unordered_map last_graph_uid; +}; + struct socket_t { ~socket_t(); + // guarded by conn.mtx_send / conn.mtx_seq, see rpc_conn_state + rpc_conn_state conn; + bool send_data(const void * data, size_t size); bool recv_data(void * data, size_t size); // Must be called at every message boundary: the RDMA transport coalesces diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 5dc05194957e..2b1f2e32157b 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1621,9 +1621,15 @@ struct server_context_impl { return true; } - // Stops every pipeline group so that the caller can touch slot and context state safely. + // Holds the engine so that the caller can look at the slot state safely, and, on request, + // waits for the in-flight decode of the groups whose context the caller is going to touch. // Constructing this is a no-op when there is a single group: the single update loop and the // task processing then run on the same thread, exactly as before. + // + // Taking mtx_engine already keeps every group out of a new iteration, so the slot state is + // stable as soon as the guard exists. Only touching a llama_context needs more than that, + // and only for the group that owns it: wait_for() blocks until that group's decode is done + // while the other groups keep computing. Tasks that touch no context wait for nobody. struct engine_guard { server_context_impl * srv = nullptr; std::unique_lock lk; @@ -1636,11 +1642,30 @@ struct server_context_impl { srv = srv_; lk = std::unique_lock(srv->mtx_engine); - // ask every group to stop at the start of its next iteration, then wait for the - // decodes that are already in flight + // ask every group to stop at the start of its next iteration, so that the slot state + // cannot change under us while wait_for() releases the lock for (auto & grp : srv->groups) { grp->n_pause_req++; } + } + + // wait until this group is not inside llama_decode, so its context can be touched + void wait_for(int id_group) { + if (srv == nullptr) { + return; + } + + GGML_ASSERT(id_group >= 0 && id_group < (int) srv->groups.size()); + server_group * grp = srv->groups[id_group].get(); + + srv->cv_engine.wait(lk, [&] { return !grp->busy; }); + } + + // wait for every group, for tasks that are not tied to one slot + void wait_for_all() { + if (srv == nullptr) { + return; + } srv->cv_engine.wait(lk, [&] { for (auto & grp : srv->groups) { @@ -2568,8 +2593,10 @@ struct server_context_impl { return false; } - // with more than one group the update loops run on their own threads, so pause them while - // we look at and modify the slots. no-op with a single group. + // with more than one group the update loops run on their own threads. Holding the engine + // is enough to look at and modify the slot state; the cases below additionally wait for + // the in-flight decode of the group whose context they touch, and only for that group. + // no-op with a single group. engine_guard guard(this); switch (task.type) { @@ -2608,6 +2635,10 @@ struct server_context_impl { break; } + // from here on the slot's context is touched (prompt cache, KV), so the group + // that owns it has to finish its decode. the other groups keep computing. + guard.wait_for(slot->id_group); + if (task.is_parent()) { // try getting free slots for all child tasks size_t n_child_tasks = task.child_tasks.size(); @@ -2636,6 +2667,9 @@ struct server_context_impl { } if (params_base.cache_idle_slots) { + // this walks every slot of every group + guard.wait_for_all(); + for (auto & slot : slots) { if (!slot.is_processing()) { SLT_TRC(slot, "%s", "saving idle slot to prompt cache\n"); @@ -2658,6 +2692,7 @@ struct server_context_impl { // release slot linked with the task id for (auto & slot : slots) { if (slot.task && slot.task->id == task.id_target) { + guard.wait_for(slot.id_group); slot.release(); break; } @@ -2678,6 +2713,9 @@ struct server_context_impl { break; } + // the sampler of this slot is used by its group between decodes + guard.wait_for(slot->id_group); + if (task.params.control_action == "reasoning_end") { // the budget sampler only exists when reasoning control was armed if (!slot->task->params.sampling.reasoning_control) { @@ -2759,6 +2797,9 @@ struct server_context_impl { break; } + // reads this slot's KV out of its context + guard.wait_for(slot->id_group); + const int64_t t_start = ggml_time_us(); std::string filename = task.slot_action.filename; @@ -2809,6 +2850,9 @@ struct server_context_impl { break; } + // writes this slot's KV into its context + guard.wait_for(slot->id_group); + const int64_t t_start = ggml_time_us(); std::string filename = task.slot_action.filename; @@ -2874,6 +2918,9 @@ struct server_context_impl { break; } + // prompt_clear() drops this slot's KV from its context + guard.wait_for(slot->id_group); + // Erase token cache const size_t n_erased = slot->prompt.tokens.size(); @@ -2913,6 +2960,9 @@ struct server_context_impl { } break; case SERVER_TASK_TYPE_SET_LORA: { + // the adapters are applied to every context + guard.wait_for_all(); + auto new_loras = construct_lora_list(task.set_lora); // logging for (size_t i = 0; i < new_loras.size(); ++i) { From 8de432016311347864040c3b8fa19abd6e792e3a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 16:17:42 -0700 Subject: [PATCH 4/6] server: sample the pipeline groups in parallel and give each group its own engine lock At 32 concurrent clients on a two-node layer split, post_decode cost 4.3 ms per slot with --pipeline-groups 2 against 0.85 ms per slot with one context, and it sits on each group's critical path between llama_synchronize and the next submit. Per-group timings show it is common_sampler_sample: 0.60 ms per row with one context, 3.6 to 4.6 ms per row with two, because the candidate array of a 248320-token vocabulary is about 4 MB of memory traffic per row and with a second group the pass runs against the other group's GPU work instead of in the gap when both GPUs are idle. - each group now samples its rows over a small worker pool. The rows are independent, so the tokens are the ones the serial pass would have produced; greedy output is byte-identical at N = 1, 2 and 4. The thread budget is divided by the number of groups, so a pipeline-groups run is not simply given more CPU. LLAMA_SERVER_SAMPLE_THREADS=1 turns it off. - the engine lock is now per group, so the host path of one group no longer excludes the other's. Measured on the pair this is worth nothing on its own (74.7 against 73.5 to 75.3 tok/s), which is reported as a falsification, but it is what the feature is supposed to guarantee and it is needed before the sampling pool can overlap anything. - get_available_slot() called prompt_save / prompt_load, which read and write the slot's sequence KV, before the guard waited for the owning group's decode. With --cache-ram 0 the cache is null so it never fired; with the cache on it is a live race against a running decode. The cache update now happens after the wait. - server_metrics is no longer written from several group threads at once, and each group only counts its own slots instead of every slot of the server. - LLAMA_SERVER_PIPE_PROF=1 prints the per-group host path every five seconds. That is how the cost above was found. README: document that a layer split must list the RPC device first and the local device last, so the output layer and its logits stay local. --- tools/server/README.md | 31 +- tools/server/server-context.cpp | 502 +++++++++++++++++++++++++++----- 2 files changed, 459 insertions(+), 74 deletions(-) diff --git a/tools/server/README.md b/tools/server/README.md index bef9088f39c4..b33c84f5f8f6 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -2084,10 +2084,25 @@ This is meant for a layer split across two machines, e.g. ```sh llama-server -m model.gguf -c 32768 --parallel 16 \ - --rpc peer:50052 --device CUDA0,RPC0 -sm layer -ngl 99 \ + --rpc peer:50052 --device RPC0,CUDA0 -sm layer -ngl 99 \ --pipeline-groups 2 ``` +Note the device order: **list the remote device first and the local one last**. With `-sm layer` +the devices are filled in the order they are given, so the last one holds the output layer. Put +the local GPU last and the logits are produced locally, which removes a `n_vocab * n_rows * 4` +byte transfer from every decode step (31.8 MB per step at 32 rows on a 248320-token vocabulary) +and lets the sampler read them out of local memory. On a pair of DGX Sparks with +Qwen3.8-27B-UD-Q4_K_XL at 32 concurrent clients this is worth more than the pipeline groups +themselves, and the two compound: + +| device order | groups | tok/s | TPOT ms | GPU busy, local / remote | +|---|---|---|---|---| +| `CUDA0,RPC0` | 1 | 94.9 | 310 | 44 / 43 pc | +| `CUDA0,RPC0` | 2 | 75.5 | 395 | 43 / 44 pc | +| `RPC0,CUDA0` | 1 | 99.7 | 295 | 41 / 43 pc | +| `RPC0,CUDA0` | 2 | **130.4** | 223 | 76 / 79 pc | + With one context, a layer split is a two-stage pipeline that is fed one batch at a time, so each stage is idle while the other one computes. With two groups there are two batches in flight, so while group A is being computed on the second stage, group B is being computed on the first one. @@ -2107,6 +2122,20 @@ Details: - `N > 1` is refused at startup together with speculative decoding (`--model-draft`, MTP), multimodal (`--mmproj`) and `--sleep-idle-seconds`. - With `N = 1` nothing changes: one context, one batch and one update loop on the main thread. +- Each group samples its own rows over a small worker pool, because a serial pass over the slots + costs tens of milliseconds per step between the decode and the next submit and, at `N > 1`, + competes for memory bandwidth with the other group's GPU work. The total number of sampling + threads is the same however many groups there are; `LLAMA_SERVER_SAMPLE_THREADS=1` turns the + pool off. The sampled tokens do not depend on how the pass is scheduled. +- Use `--cache-ram 0` with a layer split. The RAM prompt cache moves a whole slot state on every + slot handover, and on a split most of that state lives on the remote node, so it goes over the + wire; at 32 concurrent clients on the pair it costs about a quarter of the throughput with one + group and much more with two, because the handover runs on the single task thread and the other + group starves while it does. +- `LLAMA_SERVER_PIPE_PROF=1` prints, every five seconds and per group, how long an iteration + spends waiting for the engine, in `pre_decode`, in `llama_decode`, in `llama_synchronize`, and + in `post_decode`, and splits `post_decode` per token into sampling, detokenization, stop-string + handling and the result queue. That is how the numbers above were found. ## More examples diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 2b1f2e32157b..30778811dd7a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -318,6 +318,10 @@ struct server_slot { llama_token sampled; // in speculative mode, this is the last accepted token + // token produced by the parallel sampling pass of post_decode, LLAMA_TOKEN_NULL if that pass + // did not run for this slot (then the sequential path samples it as before) + llama_token pre_sampled = LLAMA_TOKEN_NULL; + // for TTS models, this is the embd generated from prev step, decode this to generate next hidden state // corresponding to one token position (size = n_embd) std::vector inp_embd; @@ -336,6 +340,7 @@ struct server_slot { int32_t n_gen_last = 0; void reset() { + pre_sampled = LLAMA_TOKEN_NULL; SLT_DBG(*this, "%s", "\n"); spec_is_replay = false; @@ -803,6 +808,144 @@ static int process_mtmd_chunk(const server_slot & slot, mtmd::batch_ptr & mbatch // With N > 1 the point is that while group A's batch is being computed on the second stage of a // layer split (the RPC peer), group B's batch can be computed on the first stage (the local GPU), // so both devices are busy instead of each idling half of every decode step. + +// ----------------------------------------------------------------------------- +// per-group host-path profiling, enabled with LLAMA_SERVER_PIPE_PROF=1 +// ----------------------------------------------------------------------------- + +static bool pipe_prof_enabled() { + const char * e = getenv("LLAMA_SERVER_PIPE_PROF"); + return e != nullptr && atoi(e) != 0; +} + +struct server_group_prof { + int64_t n_iter = 0; + int64_t n_tok = 0; + + int64_t t_lock = 0; // waiting for the engine lock at the top of the iteration + int64_t t_pre = 0; // pre_decode + render + int64_t t_submit = 0; // llama_decode + int64_t t_sync = 0; // llama_synchronize + int64_t t_relock = 0; // re-taking the engine lock after the decode + int64_t t_post = 0; // post_decode + int64_t t_sampl = 0; // of which: common_sampler_sample + int64_t t_sampl_par = 0; // of which: the parallel sampling pass + int64_t t_piece = 0; // of which: common_token_to_piece + int64_t t_proc = 0; // of which: process_token (stop strings, streaming) + int64_t t_send = 0; // of which: queue_results.send + int64_t t_iter = 0; // the whole iteration +}; + +struct prof_timer { + int64_t * acc; + int64_t t0; + prof_timer(int64_t * acc_, bool on) : acc(on ? acc_ : nullptr) { + if (acc) { t0 = ggml_time_us(); } + } + ~prof_timer() { + if (acc) { *acc += ggml_time_us() - t0; } + } + prof_timer(const prof_timer &) = delete; + prof_timer & operator=(const prof_timer &) = delete; +}; + + +// A tiny fixed worker pool used to sample the slots of one group in parallel. +// +// Sampling one row of a 250k-token vocabulary costs about 0.6 ms on this hardware, and it costs +// six times that while the other pipeline group is driving the GPUs, so a serial pass over the +// slots is tens of milliseconds sitting on the critical path between the decode and the next +// submit. The rows are independent - each slot has its own sampler and reads its own row of the +// logits - so they can be done at the same time. The output is identical either way. +struct server_par_for { + std::vector workers; + std::mutex mtx; + std::condition_variable cv_work; + std::condition_variable cv_done; + const std::function * fn = nullptr; + int n = 0; + int next = 0; + int n_running = 0; + uint64_t gen = 0; + bool stop = false; + + void start(int n_threads) { + for (int i = 0; i < n_threads; ++i) { + workers.emplace_back([this]() { worker(); }); + } + } + + // run fn(0..n_-1), the calling thread takes its share too + void run(int n_, const std::function & f) { + if (workers.empty() || n_ <= 1) { + for (int i = 0; i < n_; ++i) { + f(i); + } + return; + } + + std::unique_lock lk(mtx); + + fn = &f; + n = n_; + next = 0; + gen++; + + cv_work.notify_all(); + + take_jobs(lk); + + cv_done.wait(lk, [&] { return next >= n && n_running == 0; }); + } + + ~server_par_for() { + { + std::lock_guard lk(mtx); + stop = true; + } + cv_work.notify_all(); + for (auto & t : workers) { + if (t.joinable()) { + t.join(); + } + } + } + +private: + void take_jobs(std::unique_lock & lk) { + while (next < n) { + const int i = next++; + n_running++; + lk.unlock(); + // note: the job itself must not throw, the callers wrap it + (*fn)(i); + lk.lock(); + n_running--; + } + if (n_running == 0) { + cv_done.notify_all(); + } + } + + void worker() { + std::unique_lock lk(mtx); + uint64_t seen = 0; + while (true) { + cv_work.wait(lk, [&] { return stop || gen != seen; }); + if (stop) { + return; + } + seen = gen; + take_jobs(lk); + } + } +}; + +struct server_group; + +// the group whose decode loop is running on this thread, used by the deep call sites +static thread_local server_group * tls_group = nullptr; + struct server_group { int id = 0; @@ -821,10 +964,22 @@ struct server_group { int n_empty_consecutive = 0; - // only used when n_groups > 1, all guarded by server_context_impl::mtx_engine + // host-path profiling, only filled in when LLAMA_SERVER_PIPE_PROF=1 + server_group_prof prof; + + // sampling of this group's slots, run over several threads (see server_par_for) + server_par_for pool; + std::vector to_sample; + + // only used when n_groups > 1 + // note: the lock is per group on purpose - the whole point of the feature is that the host + // path of one group (pre_decode, sampling, streaming, post_decode) runs while the other + // group is on the GPU, so nothing here may be shared between groups std::thread thread; + std::mutex mtx; // guards this group's slots, batch and the two fields below + std::condition_variable cv; bool busy = false; // a decode is in flight, no one may touch ctx - int n_pause_req = 0; // someone wants the engine stopped, do not start a new iteration + int n_pause_req = 0; // someone wants this group stopped, do not start a new iteration }; // @@ -890,13 +1045,19 @@ struct server_context_impl { int n_groups = 1; std::vector> groups; + // LLAMA_SERVER_PIPE_PROF=1: time the host path of each group separately + const bool prof_on = pipe_prof_enabled(); + // number of sequences per context, == params_base.n_parallel when n_groups == 1 int n_seq_per_group = 1; // the following are only ever touched when n_groups > 1 - std::mutex mtx_engine; - std::condition_variable cv_engine; - bool groups_stop = false; + std::atomic groups_stop { false }; + + // server_metrics and the prompt cache are shared by every group, so they get their own locks + // instead of riding on a global engine lock. Both are off the per-token path. + std::mutex mtx_metrics; + std::mutex mtx_prompt_cache; llama_model * model_dft = nullptr; llama_context * ctx_dft = nullptr; @@ -1408,6 +1569,28 @@ struct server_context_impl { } } + // sampling threads. The budget is the same however many groups there are, so that a + // pipeline-groups run is not simply given more CPU than the single-context run. + { + int n_sampling_threads = 8; + + if (const char * e = getenv("LLAMA_SERVER_SAMPLE_THREADS")) { + n_sampling_threads = atoi(e); + } + + n_sampling_threads = std::min(n_sampling_threads, (int) std::thread::hardware_concurrency()); + n_sampling_threads = std::max(0, n_sampling_threads / n_groups); + + // the calling thread takes a share too, so this many extra workers + const int n_workers = std::max(0, n_sampling_threads - 1); + + for (auto & grp : groups) { + grp->pool.start(n_workers); + } + + SRV_INF("sampling threads per pipeline group: %d\n", n_sampling_threads); + } + 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"); @@ -1626,13 +1809,13 @@ struct server_context_impl { // Constructing this is a no-op when there is a single group: the single update loop and the // task processing then run on the same thread, exactly as before. // - // Taking mtx_engine already keeps every group out of a new iteration, so the slot state is + // Taking every group's lock keeps every group out of a new iteration, so the slot state is // stable as soon as the guard exists. Only touching a llama_context needs more than that, - // and only for the group that owns it: wait_for() blocks until that group's decode is done - // while the other groups keep computing. Tasks that touch no context wait for nobody. + // and only for the group that owns it: wait_for() drops the other groups' locks first, so + // their host path keeps running, then blocks until that group's decode is done. struct engine_guard { server_context_impl * srv = nullptr; - std::unique_lock lk; + std::vector> lks; explicit engine_guard(server_context_impl * srv_) { if (srv_->n_groups <= 1) { @@ -1640,25 +1823,35 @@ struct server_context_impl { } srv = srv_; - lk = std::unique_lock(srv->mtx_engine); - // ask every group to stop at the start of its next iteration, so that the slot state - // cannot change under us while wait_for() releases the lock - for (auto & grp : srv->groups) { - grp->n_pause_req++; + // take every group's lock, in group order, so that the slot state of the whole server + // is stable while the task is being routed. This blocks the host path of the groups, + // not their decodes, and it is released again as soon as the task knows which group + // it needs. + lks.resize(srv->groups.size()); + + for (size_t g = 0; g < srv->groups.size(); ++g) { + lks[g] = std::unique_lock(srv->groups[g]->mtx); + srv->groups[g]->n_pause_req++; } } - // wait until this group is not inside llama_decode, so its context can be touched + // let every group except id_group go, then wait until this one is not inside llama_decode void wait_for(int id_group) { if (srv == nullptr) { return; } GGML_ASSERT(id_group >= 0 && id_group < (int) srv->groups.size()); - server_group * grp = srv->groups[id_group].get(); - srv->cv_engine.wait(lk, [&] { return !grp->busy; }); + for (size_t g = 0; g < srv->groups.size(); ++g) { + if ((int) g != id_group) { + release(g); + } + } + + server_group * grp = srv->groups[id_group].get(); + grp->cv.wait(lks[id_group], [&] { return !grp->busy; }); } // wait for every group, for tasks that are not tied to one slot @@ -1667,14 +1860,10 @@ struct server_context_impl { return; } - srv->cv_engine.wait(lk, [&] { - for (auto & grp : srv->groups) { - if (grp->busy) { - return false; - } - } - return true; - }); + for (size_t g = 0; g < srv->groups.size(); ++g) { + server_group * grp = srv->groups[g].get(); + grp->cv.wait(lks[g], [&] { return !grp->busy; }); + } } ~engine_guard() { @@ -1682,26 +1871,30 @@ struct server_context_impl { return; } - for (auto & grp : srv->groups) { - grp->n_pause_req--; + for (size_t g = 0; g < srv->groups.size(); ++g) { + release(g); } - - lk.unlock(); - srv->cv_engine.notify_all(); } engine_guard(const engine_guard &) = delete; engine_guard & operator=(const engine_guard &) = delete; + + private: + void release(size_t g) { + if (!lks[g].owns_lock()) { + return; + } + srv->groups[g]->n_pause_req--; + lks[g].unlock(); + srv->groups[g]->cv.notify_all(); + } }; // the decode loop of one pipeline group, only used when n_groups > 1 void group_loop(server_group & grp) { while (true) { - { - std::unique_lock lk(mtx_engine); - if (groups_stop) { - return; - } + if (groups_stop.load(std::memory_order_relaxed)) { + return; } if (update_slots(grp)) { @@ -1709,8 +1902,9 @@ struct server_context_impl { } // nothing to do for this group, wait for a task to be assigned to one of its slots - std::unique_lock lk(mtx_engine); - cv_engine.wait_for(lk, std::chrono::milliseconds(5), [&] { return groups_stop; }); + std::unique_lock lk(grp.mtx); + grp.cv.wait_for(lk, std::chrono::milliseconds(5), + [&] { return groups_stop.load(std::memory_order_relaxed); }); } } @@ -1719,7 +1913,7 @@ struct server_context_impl { return; } - groups_stop = false; + groups_stop.store(false); for (auto & grp : groups) { server_group * g = grp.get(); @@ -1734,11 +1928,11 @@ struct server_context_impl { return; } - { - std::unique_lock lk(mtx_engine); - groups_stop = true; + for (auto & grp : groups) { + std::unique_lock lk(grp->mtx); + groups_stop.store(true); + grp->cv.notify_all(); } - cv_engine.notify_all(); for (auto & grp : groups) { if (grp->thread.joinable()) { @@ -1774,7 +1968,7 @@ struct server_context_impl { return nullptr; } - server_slot * get_available_slot(const server_task & task) { + server_slot * get_available_slot(const server_task & task, bool & need_cache_update) { server_slot * ret = nullptr; bool update_cache = false; @@ -1868,25 +2062,34 @@ struct server_context_impl { // cache prompts only for completion tasks update_cache = update_cache && task.type == SERVER_TASK_TYPE_COMPLETION; + } else { + update_cache = false; + } - if (update_cache) { - SRV_TRC("%s", "updating prompt cache\n"); + // note: the caller runs update_prompt_cache() once it knows the slot is free and the group + // that owns it is not decoding - reading and writing the sequence KV of a context + // while that context is computing is not allowed + need_cache_update = update_cache; - const int64_t t_start = ggml_time_us(); + return ret; + } - ret->prompt_save(*prompt_cache); + // moves the slot's current prompt into the RAM cache and loads the best prefix for the new + // task. Touches the slot's context, so the owning group must be out of llama_decode. + void update_prompt_cache(server_slot & slot, const server_task & task) { + SRV_TRC("%s", "updating prompt cache\n"); - if (!ret->prompt_load(*prompt_cache, task.tokens)) { - ret->prompt_clear(); - } + const int64_t t_start = ggml_time_us(); - prompt_cache->update(); + slot.prompt_save(*prompt_cache); - SRV_TRC("prompt cache update took %.2f ms\n", (ggml_time_us() - t_start) / 1000.0); - } + if (!slot.prompt_load(*prompt_cache, task.tokens)) { + slot.prompt_clear(); } - return ret; + prompt_cache->update(); + + SRV_TRC("prompt cache update took %.2f ms\n", (ggml_time_us() - t_start) / 1000.0); } // return true if at least one slot has been cleared @@ -2309,7 +2512,10 @@ struct server_context_impl { res->stats = slot.stats; } - queue_results.send(std::move(res)); + { + prof_timer psnd(tls_group ? &tls_group->prof.t_send : nullptr, prof_on && tls_group != nullptr); + queue_results.send(std::move(res)); + } } void send_final_response(server_slot & slot) { @@ -2615,7 +2821,8 @@ struct server_context_impl { const int id_task = task.id; - server_slot * slot = get_available_slot(task); + bool need_cache_update = false; + server_slot * slot = get_available_slot(task, need_cache_update); // // slot scheduling logic @@ -2639,6 +2846,10 @@ struct server_context_impl { // that owns it has to finish its decode. the other groups keep computing. guard.wait_for(slot->id_group); + if (need_cache_update) { + update_prompt_cache(*slot, task); + } + if (task.is_parent()) { // try getting free slots for all child tasks size_t n_child_tasks = task.child_tasks.size(); @@ -3042,6 +3253,47 @@ struct server_context_impl { }; #endif + // LLAMA_SERVER_PIPE_PROF=1: dump the host path of every group every 5 s and start a new window + // note: called with the group's own lock held when n_groups > 1 + std::atomic t_prof_last { 0 }; + + void prof_report() { + const int64_t t_now = ggml_time_us(); + int64_t t_last = t_prof_last.load(); + + if (t_last != 0 && t_now - t_last < 5 * 1000 * 1000) { + return; + } + if (!t_prof_last.compare_exchange_strong(t_last, t_now)) { + return; + } + if (t_last == 0) { + return; // first call only arms the window + } + + const double t_win = (t_now - t_last) / 1000.0; // ms + + for (auto & g : groups) { + auto & pr = g->prof; + + const double n_it = std::max(1, pr.n_iter); + const double n_tk = std::max(1, pr.n_tok); + auto per_it = [&](int64_t v) { return v / 1000.0 / n_it; }; + auto per_tk = [&](int64_t v) { return v / 1000.0 / n_tk; }; + + SRV_INF("PROF g%d win %.0f ms iters %" PRId64 " toks %" PRId64 + " | per iter ms: lock %.2f pre %.2f submit %.2f sync %.2f relock %.2f post %.2f iter %.2f" + " | per tok ms: sampl %.3f sampl_par %.3f piece %.3f proc %.3f send %.3f post %.3f\n", + g->id, t_win, pr.n_iter, pr.n_tok, + per_it(pr.t_lock), per_it(pr.t_pre), per_it(pr.t_submit), per_it(pr.t_sync), + per_it(pr.t_relock), per_it(pr.t_post), per_it(pr.t_iter), + per_tk(pr.t_sampl), per_tk(pr.t_sampl_par), per_tk(pr.t_piece), per_tk(pr.t_proc), + per_tk(pr.t_send), per_tk(pr.t_post)); + + pr = server_group_prof(); + } + } + // runs one iteration of the decode loop of a single pipeline group // returns true if the group had work to do bool update_slots(server_group & grp) { @@ -3050,16 +3302,25 @@ struct server_context_impl { auto & batch = grp.batch; auto & slots = grp.slots; + tls_group = &grp; + // when there is only one group there is only one thread and this lock is never engaged std::unique_lock lk; if (n_groups > 1) { - lk = std::unique_lock(mtx_engine); - cv_engine.wait(lk, [&]{ return groups_stop || grp.n_pause_req == 0; }); - if (groups_stop) { + prof_timer tl(&grp.prof.t_lock, prof_on); + lk = std::unique_lock(grp.mtx); + grp.cv.wait(lk, [&]{ return groups_stop.load(std::memory_order_relaxed) || grp.n_pause_req == 0; }); + if (groups_stop.load(std::memory_order_relaxed)) { return false; } } + prof_timer ti(&grp.prof.t_iter, prof_on); + if (prof_on) { + grp.prof.n_iter++; + prof_report(); + } + #ifdef DEBUG_TIMINGS static int64_t t_prev = 0; int64_t t_start = ggml_time_us(); @@ -3104,6 +3365,7 @@ struct server_context_impl { try { scoped_timer t(t_pre_decode, n_pre_decode); + prof_timer tp(&grp.prof.t_pre, prof_on); pre_decode(grp); batch.render(); } catch (const std::exception & e) { @@ -3168,6 +3430,7 @@ struct server_context_impl { try { scoped_timer t(t_post_decode, n_post_decode); + prof_timer tp(&grp.prof.t_post, prof_on); post_decode(grp, n_tokens, off, batch_view); } catch (const std::exception & e) { SRV_ERR("post_decode() failed: %s\n", e.what()); @@ -3686,7 +3949,10 @@ struct server_context_impl { slot.stats.n_prompt_cached = n_past; slot.stats.n_prompt_processed = 0; - metrics.add_prompt_cached(n_past); + { + std::lock_guard lk(mtx_metrics); + metrics.add_prompt_cached(n_past); + } slot.prompt.tokens.keep_first(n_past); @@ -3973,22 +4239,33 @@ struct server_context_impl { lk->unlock(); } ~decode_window() { - lk->lock(); + { + prof_timer tr(&grp->prof.t_relock, srv->prof_on); + lk->lock(); + } grp->busy = false; - srv->cv_engine.notify_all(); + grp->cv.notify_all(); } } window(this, &grp, &lk); - ret = llama_decode(ctx_tgt, batch_view); + { + prof_timer ts(&grp.prof.t_submit, prof_on); + ret = llama_decode(ctx_tgt, batch_view); + } if (ret == 0 && has_output) { + prof_timer ts(&grp.prof.t_sync, prof_on); llama_synchronize(ctx_tgt); } } else { // yield to the queue, so we can still handle metrics tasks while decoding // note: the sync is done here too, so that the wait is also covered by the yield queue_tasks.yield_to_queue([&]() { - ret = llama_decode(ctx_tgt, batch_view); + { + prof_timer ts(&grp.prof.t_submit, prof_on); + ret = llama_decode(ctx_tgt, batch_view); + } if (ret == 0 && has_output) { + prof_timer ts(&grp.prof.t_sync, prof_on); llama_synchronize(ctx_tgt); } }); @@ -4120,6 +4397,61 @@ struct server_context_impl { slot.task->params.sampling.preserved_tokens.find(token) != slot.task->params.sampling.preserved_tokens.end(); }; + // sample the rows of this sub-batch in parallel, before the sequential pass below walks + // the slots. Each row has its own sampler and its own row of the logits, so the tokens + // are exactly the ones the sequential path would have produced. + { + auto & to_sample = grp.to_sample; + + to_sample.clear(); + + for (auto * slot : slots) { + if (!is_inside_view(slot->i_batch)) { + continue; + } + if (slot->state == SLOT_STATE_DONE_PROMPT) { + if (slot->task->type == SERVER_TASK_TYPE_EMBEDDING || + slot->task->type == SERVER_TASK_TYPE_RERANK) { + continue; + } + } else if (slot->state != SLOT_STATE_GENERATING) { + continue; + } + if (slot->can_speculate()) { + continue; // the speculative pass owns the sampler of this slot + } + if (slot->task->params.sampling.backend_sampling) { + continue; // the token comes from the device, leave it to the sequential path + } + if (slot->task->params.sampling.n_probs > 0) { + continue; // populate_token_probs() reads the sampler state right after + } + + slot->pre_sampled = LLAMA_TOKEN_NULL; + to_sample.push_back(slot); + } + + if (to_sample.size() > 1) { + prof_timer ps(&grp.prof.t_sampl_par, prof_on); + + // resolve the first row on this thread: the first call after a decode may have to + // un-permute the output rows, which mutates the context + llama_get_logits_ith(grp.ctx, to_sample[0]->i_batch - off); + + grp.pool.run((int) to_sample.size(), [&](int i) { + server_slot * slot = to_sample[i]; + try { + slot->pre_sampled = common_sampler_sample(slot->smpl.get(), slot->ctx_tgt, slot->i_batch - off); + } catch (const std::exception & e) { + // leave it unsampled, the sequential pass below will hit the same error + // in the place that knows how to report it + SLT_ERR(*slot, "parallel sampling failed: %s\n", e.what()); + slot->pre_sampled = LLAMA_TOKEN_NULL; + } + }); + } + } + iterate(slots, [&](server_slot & slot) { // optionally send prompt processing progress if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_DONE_PROMPT) { @@ -4169,10 +4501,15 @@ struct server_context_impl { const int tok_idx = slot.i_batch - off; llama_token id; - { + if (slot.pre_sampled != LLAMA_TOKEN_NULL) { + id = slot.pre_sampled; + slot.pre_sampled = LLAMA_TOKEN_NULL; + } else { scoped_timer timer(t_sampl, n_sampl); + prof_timer ps(&grp.prof.t_sampl, prof_on); id = common_sampler_sample(slot.smpl.get(), slot.ctx_tgt, tok_idx); } + if (prof_on) { grp.prof.n_tok++; } slot.i_batch = -1; @@ -4193,14 +4530,22 @@ struct server_context_impl { completion_token_output result; result.tok = id; - result.text_to_send = common_token_to_piece(slot.ctx_tgt, result.tok, accept_special_token(slot, result.tok)); + { + prof_timer pp(&grp.prof.t_piece, prof_on); + result.text_to_send = common_token_to_piece(slot.ctx_tgt, result.tok, accept_special_token(slot, result.tok)); + } result.prob = 1.0f; // TODO: set it here instead of doing inside populate_token_probs if (slot.task->params.sampling.n_probs > 0) { populate_token_probs(slot, result, slot.task->params.post_sampling_probs, params_base.special, tok_idx); } - if (!process_token(result, slot)) { + bool keep_going; + { + prof_timer pt(&grp.prof.t_proc, prof_on); + keep_going = process_token(result, slot); + } + if (!keep_going) { // release slot because of stop condition slot.print_timings(); send_final_response(slot); @@ -4369,7 +4714,10 @@ struct server_context_impl { if (grp.n_prompt_queued == 0) { return; } - metrics.add_prompt(grp.n_prompt_queued, ggml_time_us() - grp.t_prompt_start); + { + std::lock_guard lk(mtx_metrics); + metrics.add_prompt(grp.n_prompt_queued, ggml_time_us() - grp.t_prompt_start); + } grp.n_prompt_queued = 0; } @@ -4377,12 +4725,18 @@ struct server_context_impl { void metrics_post_decode(server_group & grp, int32_t off, int32_t n_tokens, bool has_output) { auto & batch = grp.batch; - metrics.n_decode++; - for (const auto & slot : slots) { - if (slot.is_processing()) { - metrics.n_busy_slots++; + { + // note: only this group's slots - the other groups count their own, and their state + // may not be read from here + std::lock_guard lk(mtx_metrics); + + metrics.n_decode++; + for (const auto * slot : grp.slots) { + if (slot->is_processing()) { + metrics.n_busy_slots++; + } + metrics.n_tokens_max = std::max(metrics.n_tokens_max, (uint64_t) slot->prompt.n_tokens()); } - metrics.n_tokens_max = std::max(metrics.n_tokens_max, (uint64_t) slot.prompt.n_tokens()); } // apply enqueued prompt tokens stats @@ -4439,6 +4793,8 @@ struct server_context_impl { const uint64_t n = slot.stats.n_gen; const uint64_t n_steps = slot.stats.n_gen_steps(); + std::lock_guard lk(mtx_metrics); + metrics.predict .add(n, n_steps, t_us); metrics.predict_bucket.add(n, n_steps, t_us); From a1dd7c5e8d6fcc587529c9ae2390aed301e74265 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sat, 5 Sep 2026 21:38:36 -0700 Subject: [PATCH 5/6] server: give each pipeline group its own speculative decoding state A common_speculative and its draft (or MTP) context are bound to one target context, so with --pipeline-groups > 1 the server refused every drafter. The speculative state now lives in struct server_group: every group creates its own draft / MTP context against its own target context, sizes its own common_speculative for the group's sequences and addresses it by the slot's sequence id inside the group. The draft batches go through the group's draft context on the group's decode thread; the task-queue yields around the drafter are only taken on the single-group main-thread path, like llama_decode already does. Slot save / restore, checkpoints and the prompt cache carry the draft state per slot as before. With one group seq_id == id and the path is the same as before, spelled through groups[0]. With --model-draft the sidecar model is loaded once per group. validate_pipeline_groups no longer refuses --model-draft / --spec-type. --- tools/server/README.md | 9 +- tools/server/server-context.cpp | 165 ++++++++++++++++++++------------ 2 files changed, 113 insertions(+), 61 deletions(-) diff --git a/tools/server/README.md b/tools/server/README.md index b33c84f5f8f6..2374b76a9926 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -2119,8 +2119,13 @@ Details: slot that still holds its prefix, whichever group that slot belongs to. - Task processing briefly pauses the decode loops, so `/slots`, `/metrics` and cancellations are answered after the in-flight decode of each group finishes rather than during it. -- `N > 1` is refused at startup together with speculative decoding (`--model-draft`, MTP), - multimodal (`--mmproj`) and `--sleep-idle-seconds`. +- Speculative decoding works per group: every group owns a draft or MTP context bound to its own + target context and a `common_speculative` of its own, so `--spec-type draft-mtp` and + `--model-draft` combine with `N > 1` (with `--model-draft` the draft model is loaded once per + group). Slot save / restore, checkpoints and the prompt cache carry the draft state exactly as + with one group. +- `N > 1` is refused at startup together with multimodal (`--mmproj`), `--control-vector` and + `--sleep-idle-seconds`. - With `N = 1` nothing changes: one context, one batch and one update loop on the main thread. - Each group samples its own rows over a small worker pool, because a serial pass over the slots costs tens of milliseconds per step between the decode and the next submit and, at `N > 1`, diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 30778811dd7a..1723a55766e1 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -956,6 +956,19 @@ struct server_group { // slots owned by this group, in slot id order (slots are partitioned contiguously) std::vector slots; + // speculative decoding state of this group + // note: a common_speculative and its draft (or MTP) context are bound to one target context, + // so each group owns its own set, sized for the group's sequences and indexed by + // slot.seq_id (which is slot.id with a single group) + common_speculative_init_result_ptr spec_init; + + llama_model * model_dft = nullptr; + llama_context * ctx_dft = nullptr; + + common_speculative_ptr spec; + + common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; + // queued prompt stats - llama_decode() is async, so the timing is only valid after a sync // note: kept out of server_metrics, which is copied as-is into the task result int64_t t_decode_start = 0; // start of the last submitted decode of this group @@ -1059,15 +1072,9 @@ struct server_context_impl { std::mutex mtx_metrics; std::mutex mtx_prompt_cache; - llama_model * model_dft = nullptr; - llama_context * ctx_dft = nullptr; - - common_speculative_init_result_ptr spec_init; - + // note: the speculative decoding state (draft / MTP context, common_speculative) lives in + // the groups, see struct server_group common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; - common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; - - common_speculative_ptr spec; bool add_bos_token = true; @@ -1102,11 +1109,14 @@ struct server_context_impl { int64_t t_last_load_progress_ms = 0; void destroy() { - spec.reset(); - spec_init.reset(); + // the draft / MTP context of a group refers to the group's context, so it goes first + for (auto & grp : groups) { + grp->spec.reset(); + grp->spec_init.reset(); - ctx_dft = nullptr; - model_dft = nullptr; + grp->ctx_dft = nullptr; + grp->model_dft = nullptr; + } // groups[0]->ctx is owned by llama_init, the rest were created by llama_init_from_model() for (size_t g = 1; g < groups.size(); ++g) { @@ -1276,7 +1286,7 @@ struct server_context_impl { // every code path below exactly as it was. n_groups = std::max(1, n_pipeline_groups_req); - if (n_groups > 1 && !validate_pipeline_groups(params_base, has_spec, has_mmproj)) { + if (n_groups > 1 && !validate_pipeline_groups(params_base, has_mmproj)) { return false; } @@ -1366,31 +1376,41 @@ struct server_context_impl { load_progress_callback(0.0f, &load_progress_spec); load_progress_spec.t_last_load_progress_ms = 0; // reset so internal cbs aren't delayed - { - common_params params_dft = common_base_params_to_speculative(params_base); + // one draft / MTP context per group, each bound to the context of its group and sized + // for the group's sequences (with a single group params_ctx is params_base, as before) + // note: with --model-draft the draft model is loaded once per group + for (int g = 0; g < n_groups; ++g) { + server_group & grp = *groups[g]; + + common_params params_dft = common_base_params_to_speculative(params_ctx); // progress callback params_dft.load_progress_callback = load_progress_callback; params_dft.load_progress_callback_user_data = &load_progress_spec; - spec_init = common_speculative_init_from_params(params_dft, model_tgt, ctx_tgt); - model_dft = spec_init->model(); - ctx_dft = spec_init->context(); + grp.spec_init = common_speculative_init_from_params(params_dft, model_tgt, grp.ctx); + grp.model_dft = grp.spec_init->model(); + grp.ctx_dft = grp.spec_init->context(); - if (has_draft && model_dft == nullptr) { + if (has_draft && grp.model_dft == nullptr) { SRV_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str()); return false; } - if (ctx_dft == nullptr) { + if (grp.ctx_dft == nullptr) { SRV_ERR("%s", "failed to create MTP context\n"); return false; } - params_base.speculative.draft.ctx_tgt = ctx_tgt; - params_base.speculative.draft.ctx_dft = ctx_dft; + if (n_groups > 1) { + SRV_INF("created draft context for pipeline group %d, n_ctx = %d, n_seq_max = %d\n", + g, (int) llama_n_ctx(grp.ctx_dft), (int) llama_n_seq_max(grp.ctx_dft)); + } } + params_base.speculative.draft.ctx_tgt = ctx_tgt; + params_base.speculative.draft.ctx_dft = groups[0]->ctx_dft; + load_progress_callback(1.0f, &load_progress_spec); } @@ -1474,26 +1494,33 @@ struct server_context_impl { } // try speculative decoding - // note: a common_speculative is bound to one target context, and its code paths yield to - // the task queue, which only one thread may do - so it is off with several groups - if (ctx_tgt_seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_NO && n_groups == 1) { - try { - spec.reset(common_speculative_init(params_base.speculative, params_base.n_parallel)); - } catch (const std::exception & e) { - SRV_ERR("failed to initialize speculative decoding context: %s\n", e.what()); + // note: a common_speculative is bound to one target context, so each group gets its own, + // sized for the group's sequences (n_seq_per_group == n_parallel with one group) + for (auto & grp : groups) { + if (ctx_tgt_seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_NO) { + common_params_speculative params_spec = params_base.speculative; + + params_spec.draft.ctx_tgt = grp->ctx; + params_spec.draft.ctx_dft = grp->ctx_dft; + + try { + grp->spec.reset(common_speculative_init(params_spec, n_seq_per_group)); + } catch (const std::exception & e) { + SRV_ERR("failed to initialize speculative decoding context: %s\n", e.what()); + } } - } - if (ctx_dft) { - ctx_dft_seq_rm_type = common_context_can_seq_rm(ctx_dft); - } + if (grp->ctx_dft) { + grp->ctx_dft_seq_rm_type = common_context_can_seq_rm(grp->ctx_dft); + } - if (spec) { - SRV_TRC("%s", "speculative decoding context initialized\n"); - } else { - spec_init.reset(); - ctx_dft = nullptr; - model_dft = nullptr; + if (grp->spec) { + SRV_TRC("%s", "speculative decoding context initialized\n"); + } else { + grp->spec_init.reset(); + grp->ctx_dft = nullptr; + grp->model_dft = nullptr; + } } for (int i = 0; i < params_base.n_parallel; i++) { @@ -1506,11 +1533,11 @@ struct server_context_impl { slot.id_group = grp.id; slot.seq_id = i % n_seq_per_group; slot.ctx_tgt = grp.ctx; - slot.ctx_dft = ctx_dft; - slot.mem.init(grp.ctx, ctx_dft); + slot.ctx_dft = grp.ctx_dft; + slot.mem.init(grp.ctx, grp.ctx_dft); grp.slots.push_back(&slot); - slot.spec = spec.get(); + slot.spec = grp.spec.get(); slot.n_ctx = n_ctx_slot; slot.mctx = mctx; @@ -1641,7 +1668,7 @@ struct server_context_impl { } // refuse everything we cannot make safe with more than one context, rather than half-support it - bool validate_pipeline_groups(const common_params & params, bool has_spec, bool has_mmproj) const { + bool validate_pipeline_groups(const common_params & params, bool has_mmproj) const { auto refuse = [](const char * what) { SRV_ERR("--pipeline-groups > 1 is not supported together with %s\n", what); return false; @@ -1663,10 +1690,8 @@ struct server_context_impl { return false; } - // a common_speculative and its draft context are bound to one target context - if (has_spec) { - return refuse("speculative decoding (--model-draft / MTP)"); - } + // note: speculative decoding is fine - every group owns a draft / MTP context and a + // common_speculative of its own, see struct server_group // mtmd_context is bound to one llama_context if (has_mmproj) { @@ -2783,7 +2808,7 @@ struct server_context_impl { cur.update_tgt(slot.ctx_tgt, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); cur.update_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); // stash the draft's speculative state with the checkpoint - common_speculative_get_state(spec.get(), slot.id, cur.data_spec); + common_speculative_get_state(slot.spec, slot.seq_id, cur.data_spec); SLT_TRC(slot, "created context checkpoint %d of %d (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n", @@ -3447,6 +3472,11 @@ struct server_context_impl { auto & batch = grp.batch; auto & slots = grp.slots; (void) ctx_tgt; + + // the speculative state of this group + auto & spec = grp.spec; + auto * ctx_dft = grp.ctx_dft; + const auto ctx_dft_seq_rm_type = grp.ctx_dft_seq_rm_type; // apply context-shift if needed // TODO: simplify and improve iterate(slots, [&](server_slot & slot) { @@ -3536,7 +3566,7 @@ struct server_context_impl { generating.push_back(&slot); if (spec) { - common_speculative_get_draft_params(spec.get(), slot.id).drafting = false; + common_speculative_get_draft_params(spec.get(), slot.seq_id).drafting = false; const bool use_ckpt_tgt = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; const bool use_ckpt_dft = ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; @@ -3565,7 +3595,7 @@ struct server_context_impl { slot.spec_prompt = slot.prompt.tokens.get_text_tokens(); - common_speculative_get_draft_params(spec.get(), slot.id) = { + common_speculative_get_draft_params(spec.get(), slot.seq_id) = { /* .drafting = */ true, /* .n_max = */ n_draft_max, /* .n_past = */ slot.prompt.n_tokens(), @@ -3581,10 +3611,16 @@ struct server_context_impl { }); // generate the actual drafts (if any) + // note: only the main thread may yield to the task queue, and with several groups each + // group drafts on its own thread and against its own draft context if (!drafting.empty()) { - queue_tasks.yield_to_queue([&]() { + if (n_groups > 1) { common_speculative_draft(spec.get()); - }); + } else { + queue_tasks.yield_to_queue([&]() { + common_speculative_draft(spec.get()); + }); + } } // make checkpoints if needed @@ -3909,7 +3945,7 @@ struct server_context_impl { it->load_tgt(slot.ctx_tgt, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); it->load_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); // restore the draft's speculative state - common_speculative_set_state(spec.get(), slot.id, it->data_spec); + common_speculative_set_state(spec.get(), slot.seq_id, it->data_spec); pos_next = std::min(pos_next, std::max(it->pos_min + 1, it->pos_max)); n_past = std::min(slot.prompt.tokens.size_up_to_pos(pos_next), (size_t) it->n_tokens); @@ -4192,6 +4228,10 @@ struct server_context_impl { auto & batch = grp.batch; auto & slots = grp.slots; + // the speculative state of this group + auto & spec = grp.spec; + auto * model_dft = grp.model_dft; + SRV_DBG("n_batch (effective) = %d, off = %d\n", n_batch, off); metrics_pre_decode(grp); @@ -4327,13 +4367,17 @@ struct server_context_impl { // TODO: avoid restoring the draft context and re-evaluating the drafted tokens when not needed [TAG_SPEC_AVOID_DRAFT_REEVAL] // for now, always re-evaluate for simplicity // ref: https://github.com/ggml-org/llama.cpp/pull/22728#issuecomment-4400925384 - // note: speculative decoding is refused with more than one group, so this always runs on - // the main thread and yield_to_queue() is safe here + // note: only the main thread may yield to the task queue; with several groups the batch + // goes through this group's draft context on this group's thread if (spec) { bool ok = true; - queue_tasks.yield_to_queue([&]() { + if (n_groups > 1) { ok = common_speculative_process(spec.get(), batch_view); - }); + } else { + queue_tasks.yield_to_queue([&]() { + ok = common_speculative_process(spec.get(), batch_view); + }); + } if (!ok) { SRV_ERR("%s", "failed to process speculative batch\n"); @@ -4375,6 +4419,9 @@ struct server_context_impl { // shadow the single-context members, as update_slots() does auto * ctx_tgt = grp.ctx; auto & slots = grp.slots; + + // the speculative state of this group + auto & spec = grp.spec; (void) ctx_tgt; // for checking if a given batch index is inside batch_view @@ -4487,7 +4534,7 @@ struct server_context_impl { slot.state = SLOT_STATE_GENERATING; if (slot.can_speculate()) { - common_speculative_begin(spec.get(), slot.id, slot.prompt.tokens.get_text_tokens()); + common_speculative_begin(spec.get(), slot.seq_id, slot.prompt.tokens.get_text_tokens()); } } else if (slot.state != SLOT_STATE_GENERATING) { return; @@ -4619,7 +4666,7 @@ struct server_context_impl { SLT_INF(slot, "accepted %2zu/%2zu draft tokens\n", accepted.size() - 1, n_draft); } - common_speculative_accept(spec.get(), slot.id, accepted.size() - 1); + common_speculative_accept(spec.get(), slot.seq_id, accepted.size() - 1); slot.spec_draft = std::move(accepted); } From d98d90dd8f8074fe168e3ee7b5cbf8e3b9136e2d Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 8 Sep 2026 01:12:31 -0700 Subject: [PATCH 6/6] server: trim comments in the pipeline-groups changes --- ggml/src/ggml-rpc/ggml-rpc.cpp | 20 +--- ggml/src/ggml-rpc/transport.h | 12 +-- tools/server/server-context.cpp | 167 +++++--------------------------- tools/server/server-context.h | 3 +- tools/server/server.cpp | 8 +- 5 files changed, 37 insertions(+), 173 deletions(-) diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index ad7395906740..ea0e63ce8591 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -212,7 +212,6 @@ struct ggml_backend_rpc_device_context { uint32_t device; std::string name; std::string description; - // note: the uid of the last graph stored on the server is tracked per connection, see socket_t }; struct ggml_backend_rpc_buffer_type_context { @@ -300,7 +299,7 @@ static bool parse_endpoint(const std::string & endpoint, std::string & host, int // RPC request : | rpc_cmd (1 byte) | request_size (8 bytes) | request_data (request_size bytes) | // No response -// writes one whole message; the caller must hold sock->conn.mtx_send +// the caller must hold sock->conn.mtx_send static bool send_rpc_cmd_locked(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size) { uint8_t cmd_byte = cmd; if (!sock->send_data(&cmd_byte, sizeof(cmd_byte))) { @@ -320,11 +319,7 @@ static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, return send_rpc_cmd_locked(sock, cmd, input, input_size); } -// Reserves this thread's place in the response order of a connection. The server answers the -// commands of one connection strictly in the order it received them, so the n-th response -// belongs to the n-th response-bearing request that was written to the socket. The ticket is -// taken while mtx_send is still held by the sender, and always released, so a failed send -// cannot leave the later waiters stuck. +// the server answers one connection strictly in request order struct rpc_response_ticket { rpc_conn_state & conn; uint64_t seq; @@ -356,8 +351,7 @@ static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, std::lock_guard lock(sock->conn.mtx_send); ticket.reset(new rpc_response_ticket(sock->conn)); if (!send_rpc_cmd_locked(sock, cmd, input, input_size)) { - // still take our turn, so the ticket is released in order and no later waiter is - // woken with a response that is not theirs + // still take our turn, or a later waiter is woken with a response that is not theirs failed = true; } } @@ -367,7 +361,6 @@ static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, return false; } - // the response is read outside mtx_send, so the other threads can keep submitting ticket->wait(); uint64_t out_size; @@ -785,11 +778,8 @@ static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, g auto sock = get_socket(rpc_ctx->endpoint); - // The graph stored by RPC_CMD_GRAPH_COMPUTE lives on the server per connection and device, - // and one connection is shared by every backend of this endpoint - including the backends of - // other llama_contexts. So the uid of the last graph sent has to be tracked per connection, - // and the check has to happen under the same lock as the send, or a RECOMPUTE could re-run - // the graph another context stored in between. + // the stored graph is per connection and device, and other llama_contexts share the connection: + // the uid check must stay under mtx_send, or RECOMPUTE re-runs a graph stored in between std::unique_lock lock(sock->conn.mtx_send); auto & last_uid = sock->conn.last_graph_uid[rpc_ctx->device]; diff --git a/ggml/src/ggml-rpc/transport.h b/ggml/src/ggml-rpc/transport.h index 2705befbc19b..779646081281 100644 --- a/ggml/src/ggml-rpc/transport.h +++ b/ggml/src/ggml-rpc/transport.h @@ -13,14 +13,9 @@ typedef std::shared_ptr socket_ptr; static constexpr size_t MAX_CHUNK_SIZE = 1024ull * 1024ull * 1024ull; // 1 GiB static constexpr size_t RPC_CONN_CAPS_SIZE = 24; -// State shared by every client thread that uses one connection. A connection is looked up by -// endpoint and is therefore shared by all backends of that endpoint, including the backends of -// different llama_contexts, so all of it has to be serialised: -// - mtx_send makes a whole RPC message atomic on the wire -// - seq_* hands the responses out in request order (the server answers strictly in order), -// without holding mtx_send while waiting, so another thread can keep submitting work -// - last_graph_uid mirrors the server's per-connection stored graph for a device, so that -// RPC_CMD_GRAPH_RECOMPUTE can never re-run a graph submitted by another context +// a connection is looked up by endpoint, so every backend of that endpoint shares it, including +// those of other llama_contexts: mtx_send makes a whole message atomic on the wire, seq_* hands +// the responses out in request order without holding mtx_send, last_graph_uid mirrors the server struct rpc_conn_state { std::mutex mtx_send; std::mutex mtx_seq; @@ -34,7 +29,6 @@ struct rpc_conn_state { struct socket_t { ~socket_t(); - // guarded by conn.mtx_send / conn.mtx_seq, see rpc_conn_state rpc_conn_state conn; bool send_data(const void * data, size_t size); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 1723a55766e1..d8e1634dfb71 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -35,7 +35,6 @@ #include #endif -// used by the --pipeline-groups decode threads #include #include #include @@ -201,12 +200,10 @@ struct server_batch { struct server_slot { int id; - // pipeline group that owns this slot, i.e. the index of ctx_tgt in server_context_impl::groups - // always 0 unless --pipeline-groups > 1 + // pipeline group that owns this slot, index into server_context_impl::groups int id_group = 0; - // sequence id of this slot inside ctx_tgt / ctx_dft - // equal to id unless --pipeline-groups > 1, where each context only holds n_parallel/N sequences + // sequence id of this slot inside ctx_tgt / ctx_dft, equal to id unless --pipeline-groups > 1 int seq_id = 0; llama_context * ctx_tgt = nullptr; @@ -318,8 +315,7 @@ struct server_slot { llama_token sampled; // in speculative mode, this is the last accepted token - // token produced by the parallel sampling pass of post_decode, LLAMA_TOKEN_NULL if that pass - // did not run for this slot (then the sequential path samples it as before) + // token from the parallel sampling pass of post_decode, LLAMA_TOKEN_NULL if it did not run llama_token pre_sampled = LLAMA_TOKEN_NULL; // for TTS models, this is the embd generated from prev step, decode this to generate next hidden state @@ -801,18 +797,6 @@ static int process_mtmd_chunk(const server_slot & slot, mtmd::batch_ptr & mbatch return try_decode(); } -// A pipeline group is one llama_context with its own batch, its own decode loop and its own -// contiguous range of slots. With --pipeline-groups 1 (the default) there is exactly one group: -// it owns ctx_tgt and every slot, and its update loop runs on the main thread, as before. -// -// With N > 1 the point is that while group A's batch is being computed on the second stage of a -// layer split (the RPC peer), group B's batch can be computed on the first stage (the local GPU), -// so both devices are busy instead of each idling half of every decode step. - -// ----------------------------------------------------------------------------- -// per-group host-path profiling, enabled with LLAMA_SERVER_PIPE_PROF=1 -// ----------------------------------------------------------------------------- - static bool pipe_prof_enabled() { const char * e = getenv("LLAMA_SERVER_PIPE_PROF"); return e != nullptr && atoi(e) != 0; @@ -850,13 +834,7 @@ struct prof_timer { }; -// A tiny fixed worker pool used to sample the slots of one group in parallel. -// -// Sampling one row of a 250k-token vocabulary costs about 0.6 ms on this hardware, and it costs -// six times that while the other pipeline group is driving the GPUs, so a serial pass over the -// slots is tens of milliseconds sitting on the critical path between the decode and the next -// submit. The rows are independent - each slot has its own sampler and reads its own row of the -// logits - so they can be done at the same time. The output is identical either way. +// each slot has its own sampler and its own row of the logits, so parallel sampling is exact struct server_par_for { std::vector workers; std::mutex mtx; @@ -943,7 +921,6 @@ struct server_par_for { struct server_group; -// the group whose decode loop is running on this thread, used by the deep call sites static thread_local server_group * tls_group = nullptr; struct server_group { @@ -953,13 +930,9 @@ struct server_group { server_batch batch; - // slots owned by this group, in slot id order (slots are partitioned contiguously) std::vector slots; - // speculative decoding state of this group - // note: a common_speculative and its draft (or MTP) context are bound to one target context, - // so each group owns its own set, sized for the group's sequences and indexed by - // slot.seq_id (which is slot.id with a single group) + // note: bound to one target context, so each group owns its own set, indexed by slot.seq_id common_speculative_init_result_ptr spec_init; llama_model * model_dft = nullptr; @@ -969,25 +942,19 @@ struct server_group { common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; - // queued prompt stats - llama_decode() is async, so the timing is only valid after a sync - // note: kept out of server_metrics, which is copied as-is into the task result + // note: async, so only valid after a sync; kept out of server_metrics, which is copied as-is int64_t t_decode_start = 0; // start of the last submitted decode of this group int64_t t_prompt_start = 0; // start of the oldest queued prompt decode of this group uint64_t n_prompt_queued = 0; int n_empty_consecutive = 0; - // host-path profiling, only filled in when LLAMA_SERVER_PIPE_PROF=1 server_group_prof prof; - // sampling of this group's slots, run over several threads (see server_par_for) server_par_for pool; std::vector to_sample; - // only used when n_groups > 1 - // note: the lock is per group on purpose - the whole point of the feature is that the host - // path of one group (pre_decode, sampling, streaming, post_decode) runs while the other - // group is on the GPU, so nothing here may be shared between groups + // note: per group on purpose - one group's host path runs while the other is on the GPU std::thread thread; std::mutex mtx; // guards this group's slots, batch and the two fields below std::condition_variable cv; @@ -1019,7 +986,6 @@ struct server_context_impl { server_state_callback_t callback_state = [](server_state, json) -> void {}; - // number of pipeline groups requested via --pipeline-groups, must be set before load_model() int n_pipeline_groups_req = 1; server_context_impl() { @@ -1053,27 +1019,18 @@ struct server_context_impl { llama_context * ctx_tgt = nullptr; - // pipeline groups, see --pipeline-groups and struct server_group - // groups[0]->ctx is always ctx_tgt; n_groups == 1 unless the user asked for more int n_groups = 1; std::vector> groups; - // LLAMA_SERVER_PIPE_PROF=1: time the host path of each group separately const bool prof_on = pipe_prof_enabled(); - // number of sequences per context, == params_base.n_parallel when n_groups == 1 int n_seq_per_group = 1; - // the following are only ever touched when n_groups > 1 std::atomic groups_stop { false }; - // server_metrics and the prompt cache are shared by every group, so they get their own locks - // instead of riding on a global engine lock. Both are off the per-token path. std::mutex mtx_metrics; std::mutex mtx_prompt_cache; - // note: the speculative decoding state (draft / MTP context, common_speculative) lives in - // the groups, see struct server_group common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; bool add_bos_token = true; @@ -1281,9 +1238,6 @@ struct server_context_impl { params_base.load_progress_callback_user_data = &load_progress_text; } - // --pipeline-groups: run the slots over N independent contexts of one model, so that the - // stages of a layer split can be busy at the same time. N == 1 is the default and keeps - // every code path below exactly as it was. n_groups = std::max(1, n_pipeline_groups_req); if (n_groups > 1 && !validate_pipeline_groups(params_base, has_mmproj)) { @@ -1297,8 +1251,7 @@ struct server_context_impl { common_params & params_ctx = n_groups > 1 ? params_grp : params_base; if (n_groups > 1) { - // each context gets 1/N of the sequences and 1/N of the total context, so the per-slot - // context (n_ctx / n_seq_max) and the total KV memory over all contexts are unchanged + // 1/N of the sequences and of the context each, so per-slot context and total KV hold params_ctx.n_parallel = n_seq_per_group; params_ctx.n_ctx = params_base.n_ctx / n_groups; } @@ -1331,7 +1284,6 @@ struct server_context_impl { vocab = llama_model_get_vocab(model_tgt); - // the remaining contexts of the pipeline are created from the same model { groups.clear(); groups.reserve(n_groups); @@ -1366,7 +1318,6 @@ struct server_context_impl { } } - // the total context over all groups, as requested by the user n_ctx = llama_n_ctx(ctx_tgt) * n_groups; add_bos_token = llama_vocab_get_add_bos(vocab); @@ -1376,8 +1327,6 @@ struct server_context_impl { load_progress_callback(0.0f, &load_progress_spec); load_progress_spec.t_last_load_progress_ms = 0; // reset so internal cbs aren't delayed - // one draft / MTP context per group, each bound to the context of its group and sized - // for the group's sequences (with a single group params_ctx is params_base, as before) // note: with --model-draft the draft model is loaded once per group for (int g = 0; g < n_groups; ++g) { server_group & grp = *groups[g]; @@ -1494,8 +1443,6 @@ struct server_context_impl { } // try speculative decoding - // note: a common_speculative is bound to one target context, so each group gets its own, - // sized for the group's sequences (n_seq_per_group == n_parallel with one group) for (auto & grp : groups) { if (ctx_tgt_seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_NO) { common_params_speculative params_spec = params_base.speculative; @@ -1596,8 +1543,7 @@ struct server_context_impl { } } - // sampling threads. The budget is the same however many groups there are, so that a - // pipeline-groups run is not simply given more CPU than the single-context run. + // the budget is split across the groups, so a pipeline run is not given more CPU { int n_sampling_threads = 8; @@ -1608,7 +1554,6 @@ struct server_context_impl { n_sampling_threads = std::min(n_sampling_threads, (int) std::thread::hardware_concurrency()); n_sampling_threads = std::max(0, n_sampling_threads / n_groups); - // the calling thread takes a share too, so this many extra workers const int n_workers = std::max(0, n_sampling_threads - 1); for (auto & grp : groups) { @@ -1667,7 +1612,6 @@ struct server_context_impl { return true; } - // refuse everything we cannot make safe with more than one context, rather than half-support it bool validate_pipeline_groups(const common_params & params, bool has_mmproj) const { auto refuse = [](const char * what) { SRV_ERR("--pipeline-groups > 1 is not supported together with %s\n", what); @@ -1690,22 +1634,17 @@ struct server_context_impl { return false; } - // note: speculative decoding is fine - every group owns a draft / MTP context and a - // common_speculative of its own, see struct server_group - // mtmd_context is bound to one llama_context if (has_mmproj) { return refuse("multimodal (--mmproj)"); } - // common_init_from_params() applies the control vector to the context it creates and only - // to that one, so the extra contexts would silently run without it + // common_init_from_params() applies it only to the context it creates if (!params.control_vectors.empty()) { return refuse("--control-vector"); } - // entering / leaving the sleeping state destroys and rebuilds the contexts under the - // running group threads + // entering / leaving it rebuilds the contexts under the running group threads if (params.sleep_idle_seconds >= 0) { return refuse("--sleep-idle"); } @@ -1829,15 +1768,8 @@ struct server_context_impl { return true; } - // Holds the engine so that the caller can look at the slot state safely, and, on request, - // waits for the in-flight decode of the groups whose context the caller is going to touch. - // Constructing this is a no-op when there is a single group: the single update loop and the - // task processing then run on the same thread, exactly as before. - // - // Taking every group's lock keeps every group out of a new iteration, so the slot state is - // stable as soon as the guard exists. Only touching a llama_context needs more than that, - // and only for the group that owns it: wait_for() drops the other groups' locks first, so - // their host path keeps running, then blocks until that group's decode is done. + // holds every group's lock, so the slot state is stable while the guard exists; a no-op with a + // single group. wait_for() then drops the other groups' locks and waits out one group's decode. struct engine_guard { server_context_impl * srv = nullptr; std::vector> lks; @@ -1849,10 +1781,7 @@ struct server_context_impl { srv = srv_; - // take every group's lock, in group order, so that the slot state of the whole server - // is stable while the task is being routed. This blocks the host path of the groups, - // not their decodes, and it is released again as soon as the task knows which group - // it needs. + // always in group order, so two guards cannot deadlock against each other lks.resize(srv->groups.size()); for (size_t g = 0; g < srv->groups.size(); ++g) { @@ -1861,7 +1790,6 @@ struct server_context_impl { } } - // let every group except id_group go, then wait until this one is not inside llama_decode void wait_for(int id_group) { if (srv == nullptr) { return; @@ -1879,7 +1807,6 @@ struct server_context_impl { grp->cv.wait(lks[id_group], [&] { return !grp->busy; }); } - // wait for every group, for tasks that are not tied to one slot void wait_for_all() { if (srv == nullptr) { return; @@ -1915,7 +1842,6 @@ struct server_context_impl { } }; - // the decode loop of one pipeline group, only used when n_groups > 1 void group_loop(server_group & grp) { while (true) { if (groups_stop.load(std::memory_order_relaxed)) { @@ -1926,7 +1852,6 @@ struct server_context_impl { continue; } - // nothing to do for this group, wait for a task to be assigned to one of its slots std::unique_lock lk(grp.mtx); grp.cv.wait_for(lk, std::chrono::milliseconds(5), [&] { return groups_stop.load(std::memory_order_relaxed); }); @@ -2091,16 +2016,12 @@ struct server_context_impl { update_cache = false; } - // note: the caller runs update_prompt_cache() once it knows the slot is free and the group - // that owns it is not decoding - reading and writing the sequence KV of a context - // while that context is computing is not allowed need_cache_update = update_cache; return ret; } - // moves the slot's current prompt into the RAM cache and loads the best prefix for the new - // task. Touches the slot's context, so the owning group must be out of llama_decode. + // reads and writes the slot's sequence KV, so the owning group must be out of llama_decode void update_prompt_cache(server_slot & slot, const server_task & task) { SRV_TRC("%s", "updating prompt cache\n"); @@ -2129,7 +2050,6 @@ struct server_context_impl { return res; } - // only slots of this group, their KV lives in this group's context for (auto * slot_ptr : grp.slots) { auto & slot = *slot_ptr; @@ -2707,7 +2627,6 @@ struct server_context_impl { std::vector get_free_slots(size_t n_slots_needed, int exclude_id_slot, int id_group) { std::vector free_slots; for (auto & slot : slots) { - // the parent copies its KV into the children, so they must live in the same context if (slot.id_group != id_group) { continue; } @@ -2824,10 +2743,6 @@ struct server_context_impl { return false; } - // with more than one group the update loops run on their own threads. Holding the engine - // is enough to look at and modify the slot state; the cases below additionally wait for - // the in-flight decode of the group whose context they touch, and only for that group. - // no-op with a single group. engine_guard guard(this); switch (task.type) { @@ -2867,8 +2782,7 @@ struct server_context_impl { break; } - // from here on the slot's context is touched (prompt cache, KV), so the group - // that owns it has to finish its decode. the other groups keep computing. + // from here the slot's context is touched, so its group must finish its decode guard.wait_for(slot->id_group); if (need_cache_update) { @@ -2878,9 +2792,7 @@ struct server_context_impl { if (task.is_parent()) { // try getting free slots for all child tasks size_t n_child_tasks = task.child_tasks.size(); - // the children take their KV from the parent, so they must fit in the - // parent's group. with a single group this is the limit the request - // schema already enforces, so nothing changes there. + // the children take their KV from the parent, so they must fit its group if ((int) n_child_tasks + 1 > n_seq_per_group) { send_error(task, string_format( "n_cmpl must not exceed the number of slots per pipeline group (%d)", n_seq_per_group), @@ -2903,7 +2815,6 @@ struct server_context_impl { } if (params_base.cache_idle_slots) { - // this walks every slot of every group guard.wait_for_all(); for (auto & slot : slots) { @@ -3033,7 +2944,6 @@ struct server_context_impl { break; } - // reads this slot's KV out of its context guard.wait_for(slot->id_group); const int64_t t_start = ggml_time_us(); @@ -3086,7 +2996,6 @@ struct server_context_impl { break; } - // writes this slot's KV into its context guard.wait_for(slot->id_group); const int64_t t_start = ggml_time_us(); @@ -3154,7 +3063,6 @@ struct server_context_impl { break; } - // prompt_clear() drops this slot's KV from its context guard.wait_for(slot->id_group); // Erase token cache @@ -3278,7 +3186,6 @@ struct server_context_impl { }; #endif - // LLAMA_SERVER_PIPE_PROF=1: dump the host path of every group every 5 s and start a new window // note: called with the group's own lock held when n_groups > 1 std::atomic t_prof_last { 0 }; @@ -3319,8 +3226,7 @@ struct server_context_impl { } } - // runs one iteration of the decode loop of a single pipeline group - // returns true if the group had work to do + // one iteration of the decode loop of a single group, returns true if it had work to do bool update_slots(server_group & grp) { // shadow the single-context members - everything below operates on this group only auto * ctx_tgt = grp.ctx; @@ -3329,7 +3235,6 @@ struct server_context_impl { tls_group = &grp; - // when there is only one group there is only one thread and this lock is never engaged std::unique_lock lk; if (n_groups > 1) { prof_timer tl(&grp.prof.t_lock, prof_on); @@ -3384,8 +3289,7 @@ struct server_context_impl { task.id = queue_tasks.get_new_id(); queue_tasks.post(std::move(task)); } - // note: with more than one group each group drives its own loop, so there is no need - // to keep the shared task loop spinning + // note: each group drives its own loop, so the shared task loop need not keep spinning } try { @@ -3473,7 +3377,6 @@ struct server_context_impl { auto & slots = grp.slots; (void) ctx_tgt; - // the speculative state of this group auto & spec = grp.spec; auto * ctx_dft = grp.ctx_dft; const auto ctx_dft_seq_rm_type = grp.ctx_dft_seq_rm_type; @@ -3611,8 +3514,7 @@ struct server_context_impl { }); // generate the actual drafts (if any) - // note: only the main thread may yield to the task queue, and with several groups each - // group drafts on its own thread and against its own draft context + // note: only the main thread may yield to the task queue if (!drafting.empty()) { if (n_groups > 1) { common_speculative_draft(spec.get()); @@ -4228,7 +4130,6 @@ struct server_context_impl { auto & batch = grp.batch; auto & slots = grp.slots; - // the speculative state of this group auto & spec = grp.spec; auto * model_dft = grp.model_dft; @@ -4264,11 +4165,8 @@ struct server_context_impl { int ret = 0; if (n_groups > 1) { - // release the engine for the duration of the compute - this is the whole point of the - // feature: while this group is on one stage of the layer split, the other group can - // run its own pre_decode / post_decode and submit its batch to the other stage - // note: RAII, so a throwing decode cannot leave the group marked busy forever, nor - // return to the caller's error handling without the engine lock held + // note: RAII, or a throwing decode leaves the group busy forever, or returns to the + // caller's error handling without the engine lock held struct decode_window { server_context_impl * srv; server_group * grp; @@ -4367,8 +4265,6 @@ struct server_context_impl { // TODO: avoid restoring the draft context and re-evaluating the drafted tokens when not needed [TAG_SPEC_AVOID_DRAFT_REEVAL] // for now, always re-evaluate for simplicity // ref: https://github.com/ggml-org/llama.cpp/pull/22728#issuecomment-4400925384 - // note: only the main thread may yield to the task queue; with several groups the batch - // goes through this group's draft context on this group's thread if (spec) { bool ok = true; if (n_groups > 1) { @@ -4388,7 +4284,6 @@ struct server_context_impl { } // handle `n_cmpl > 1` tasks - when the main prompt is processed, activate all child tasks too - // note: children are always in the same group as the parent, see get_free_slots() for (auto * slot_ptr : slots) { auto & slot = *slot_ptr; if (slot.state == SLOT_STATE_DONE_PROMPT && slot.task->is_parent()) { @@ -4416,12 +4311,9 @@ struct server_context_impl { } void post_decode(server_group & grp, int32_t n_batch_tokens, int32_t off, llama_batch & batch_view) { - // shadow the single-context members, as update_slots() does auto * ctx_tgt = grp.ctx; auto & slots = grp.slots; - - // the speculative state of this group - auto & spec = grp.spec; + auto & spec = grp.spec; (void) ctx_tgt; // for checking if a given batch index is inside batch_view @@ -4444,9 +4336,6 @@ struct server_context_impl { slot.task->params.sampling.preserved_tokens.find(token) != slot.task->params.sampling.preserved_tokens.end(); }; - // sample the rows of this sub-batch in parallel, before the sequential pass below walks - // the slots. Each row has its own sampler and its own row of the logits, so the tokens - // are exactly the ones the sequential path would have produced. { auto & to_sample = grp.to_sample; @@ -4481,8 +4370,7 @@ struct server_context_impl { if (to_sample.size() > 1) { prof_timer ps(&grp.prof.t_sampl_par, prof_on); - // resolve the first row on this thread: the first call after a decode may have to - // un-permute the output rows, which mutates the context + // the first call after a decode may un-permute the rows, which mutates the context llama_get_logits_ith(grp.ctx, to_sample[0]->i_batch - off); grp.pool.run((int) to_sample.size(), [&](int i) { @@ -4490,8 +4378,7 @@ struct server_context_impl { try { slot->pre_sampled = common_sampler_sample(slot->smpl.get(), slot->ctx_tgt, slot->i_batch - off); } catch (const std::exception & e) { - // leave it unsampled, the sequential pass below will hit the same error - // in the place that knows how to report it + // leave it unsampled, the sequential pass below reports the same error SLT_ERR(*slot, "parallel sampling failed: %s\n", e.what()); slot->pre_sampled = LLAMA_TOKEN_NULL; } @@ -4773,8 +4660,7 @@ struct server_context_impl { auto & batch = grp.batch; { - // note: only this group's slots - the other groups count their own, and their state - // may not be read from here + // note: only this group's slots - another group's slot state may not be read here std::lock_guard lk(mtx_metrics); metrics.n_decode++; @@ -4875,7 +4761,6 @@ bool server_context::load_model(common_params & params) { void server_context::start_loop() { auto & params = impl->params_base; - // no-op unless --pipeline-groups > 1 impl->start_groups(); impl->queue_tasks.start_loop(params.sleep_idle_seconds * 1000); diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 315c96faa85f..d22756f8d4f8 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -111,8 +111,7 @@ struct server_context { // note: must be set before load_model() is called void set_state_callback(server_state_callback_t callback); - // number of pipeline groups, i.e. independent llama_contexts over the one model, each with its - // own slots, batch and decode thread (--pipeline-groups, default 1) + // independent llama_contexts over the one model, each with its own slots (default 1) // note: must be set before load_model() is called void set_pipeline_groups(int n_groups); }; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 45c8d7f7005b..dd9e2a4cf4ff 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -28,11 +28,7 @@ static std::function shutdown_handler; static std::atomic_flag is_terminating = ATOMIC_FLAG_INIT; -// --pipeline-groups N: run the slots over N independent llama_contexts of the same model, each -// with its own batch and decode thread. Useful with a layer split over two nodes (--rpc), where a -// single context leaves each stage idle for half of every decode step. -// The option is parsed here instead of in common/arg.cpp because it only means anything for the -// server; everything it changes lives under tools/server. +// parsed here, not in common/arg.cpp: everything it changes lives under tools/server static int g_pipeline_groups = 1; static void server_take_pipeline_groups(int & argc, char ** argv) { @@ -141,7 +137,7 @@ int llama_server(int argc, char ** argv) { // own arguments required by this example common_params params; - // strip the server-only --pipeline-groups before the common parser sees it + // must run before the common parser, which rejects the unknown option server_take_pipeline_groups(argc, argv); common_init();