diff --git a/common/arg.cpp b/common/arg.cpp index 86f8610a56d0..5d13a42f8b11 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2539,6 +2539,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.n_parallel = value; } ).set_env("LLAMA_ARG_N_PARALLEL").set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--pipeline-groups"}, "N", + string_format("run the server slots over N independent contexts of one model, so one " + "group decodes while another is between steps (default: %d)", params.n_pipeline_groups), + [](common_params & params, int value) { + if (value < 1) { + throw std::invalid_argument("error: --pipeline-groups must be >= 1\n"); + } + params.n_pipeline_groups = value; + } + ).set_env("LLAMA_ARG_PIPELINE_GROUPS").set_examples({LLAMA_EXAMPLE_SERVER})); } else { add_opt(common_arg( {"-np", "--parallel"}, "N", diff --git a/common/common.h b/common/common.h index de49dac9f63a..aea7016fa992 100644 --- a/common/common.h +++ b/common/common.h @@ -465,6 +465,8 @@ struct common_params { int32_t n_gpu_layers = -1; // number of layers to store in VRAM, -1 is auto, <= -2 is all int32_t main_gpu = 0; // the GPU that is used for scratch and small tensors float tensor_split[128] = {0}; // how split tensors should be distributed across GPUs + int32_t n_pipeline_groups = 1; // llama-server: run the slots over N independent contexts + bool fit_params = true; // whether to fit unset model/context parameters to free device memory bool fit_params_print = false; // print the estimated required memory to run the model int32_t fit_params_min_ctx = 4096; // minimum context size to set when trying to reduce memory use diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 69a8a08ae172..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; - uint64_t last_graph_uid; }; struct ggml_backend_rpc_buffer_type_context { @@ -300,7 +299,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) { +// 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 +314,55 @@ 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); +} + +// the server answers one connection strictly in request order +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, or a later waiter is woken with a response that is not theirs + failed = true; + } + } + + if (failed) { + ticket->wait(); return false; } + + ticket->wait(); + uint64_t out_size; if (!sock->recv_data(&out_size, sizeof(out_size))) { return false; @@ -731,21 +774,28 @@ 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 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]; + 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 +2094,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..779646081281 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,24 @@ 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; +// 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; + 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(); + 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/README.md b/tools/server/README.md index 93736c3edfa9..4db2f204f03c 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -2074,6 +2074,91 @@ 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 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. + +Details: + +- The option is registered with the normal argument parser, so it appears in `--help` and accepts + `LLAMA_ARG_PIPELINE_GROUPS` from the environment like any other option, with an explicit flag + winning over the environment. That environment path is how router mode reaches its children, + which are spawned rather than given a command line. +- 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 the **full** `n_ctx = C`. `n_ctx` is never + divided by the group count: `-c` is what every request should be able to reach, and splitting the + server into groups is an internal arrangement that should not redefine it. `-c` must still be + given explicitly, but it no longer has to be a multiple of `N`. + + What that gives a single request depends on the KV mode, because `llama_context` derives the + per-request context differently: + - **Unified KV** (`--kv-unified`): the per-request context *is* `n_ctx`, so every request can + reach the full `C`. This is the intended arrangement. + - **Split KV**: the per-request context is `n_ctx / n_seq_max`, which here is `C*N/P`. That is + `N` times what a slot got when `n_ctx` was divided as well, so slots gain context rather than + lose it, but a single request still cannot reach `C` unless the KV is unified. + + The cost either way is that the KV memory over all groups is roughly `N` times that of a single + context, rather than equal to it. That is included in the parameter fit, so if it does not fit, + the fitter lowers `n_ctx` and reports it rather than the server quietly serving less than asked. +- 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. +- 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 rows serially, on its own decode thread. Sampling them over a worker pool + was tried and removed: `common_sampler_sample` begins with `llama_synchronize`, which does + non-atomic read-modify-writes on the context's timing counters, and `set_logits` then re-enters + the same context through six more getters, so the workers raced on a context they shared. Making + that safe means duplicating a large part of the context API. +- 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 ### Interactive mode diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b4..d8c44894f5f6 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -35,6 +35,11 @@ #include #endif +#include +#include +#include +#include + constexpr int HTTP_POLLING_SECONDS = 1; static common_speculative_output_limits server_output_limits(const common_params & params) { @@ -68,7 +73,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 +114,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 +165,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 +200,12 @@ struct server_batch { struct server_slot { int id; + // 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 + int seq_id = 0; + llama_context * ctx_tgt = nullptr; llama_context * ctx_dft = nullptr; @@ -255,8 +267,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 +280,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 +300,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 +363,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 +475,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 +495,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 +688,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 +793,96 @@ static int process_mtmd_chunk(const server_slot & slot, mtmd::batch_ptr & mbatch return try_decode(); } +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_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 + + int64_t t_prof_last = 0; // start of this group's reporting window +}; + +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; +}; + + +struct server_group; + +static thread_local server_group * tls_group = nullptr; + +struct server_group { + int id = 0; + + llama_context * ctx = nullptr; + + // groups after the first are made by llama_init_from_model(), which does not attach the + // pool common_init_from_params() builds for group 0. Without one the CPU backend makes a + // throwaway pool per graph and ignores --cpu-mask, --prio, --poll and --cpu-strict. + // Declared after ctx so it outlives the llama_free() above groups.clear(). + std::unique_ptr threadpools; + + // common_speculative_init_result owns only the model and the context and never attaches a + // pool, so the draft context falls back to a throwaway pool per graph and ignores the draft + // --cpu-mask-draft, --prio, --poll and --cpu-strict settings. Declared before spec_init so + // it is destroyed after the draft context it is attached to. + std::unique_ptr threadpools_dft; + + server_batch batch; + + std::vector slots; + + // 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; + 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; + + // 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; + + server_group_prof prof; + + // 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; + bool busy = false; // a decode is in flight, no one may touch ctx + int n_pause_req = 0; // someone wants this group stopped, do not start a new iteration +}; + // // server_context_impl (private implementation) // @@ -804,6 +907,8 @@ struct server_context_impl { server_state_callback_t callback_state = [](server_state, json) -> void {}; + int n_pipeline_groups_req = 1; + server_context_impl() { mtmd_helper_log_set(common_log_default_callback, nullptr); } @@ -835,17 +940,19 @@ struct server_context_impl { llama_context * ctx_tgt = nullptr; - server_batch batch; + int n_groups = 1; + std::vector> groups; - llama_model * model_dft = nullptr; - llama_context * ctx_dft = nullptr; + const bool prof_on = pipe_prof_enabled(); - common_speculative_init_result_ptr spec_init; + int n_seq_per_group = 1; - 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; + std::atomic groups_stop { false }; - common_speculative_ptr spec; + std::mutex mtx_metrics; + std::mutex mtx_prompt_cache; + + common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; bool add_bos_token = true; @@ -862,18 +969,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 @@ -888,11 +987,23 @@ 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) { + if (groups[g]->ctx != nullptr) { + llama_free(groups[g]->ctx); + groups[g]->ctx = nullptr; + } + } + groups.clear(); llama_init.reset(); @@ -1048,11 +1159,67 @@ struct server_context_impl { params_base.load_progress_callback_user_data = &load_progress_text; } - llama_init = common_init_from_params(params_base); + n_groups = std::max(1, n_pipeline_groups_req); + + if (n_groups > 1 && !validate_pipeline_groups(params_base, 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) { + params_ctx.n_parallel = n_seq_per_group; + + // n_ctx is deliberately NOT divided, in either KV mode: -c is what the user asked + // every request to be able to reach, and splitting the server into groups is an + // internal arrangement that should not quietly redefine it. Each group is built with + // the full n_ctx. + // + // What that means per slot differs by mode, because llama_context derives the + // per-request context differently. Unified KV sets n_ctx_seq = n_ctx, so every + // request can reach C, which is the point. Split KV sets n_ctx_seq = n_ctx / + // n_seq_max, and n_seq_max here is P/N, so a slot gets C*N/P: N times what it had when + // n_ctx was divided as well, not the same. That is more context per slot than before, + // never less, so nothing that used to fit stops fitting. + // + // The cost is aggregate KV of roughly N*C rather than C. That is real, and it is what + // reserve_extra_group_memory() below exists to account for: it measures each extra + // context and charges it to the fit margins, so a configuration that no longer fits + // comes back as an n_ctx the fitter lowered and reported rather than a silent + // shrinking of what was asked for. + + // common_init_from_params fits ONE context and fixes the model placement from that + // estimate, but the n_groups - 1 other contexts are created afterwards, once placement + // can no longer change. Fitting the undivided n_ctx instead would not be enough: + // measured with common_get_device_memory_data on a 1.5B, KV tracks n_ctx (112, 224 and + // 448 MiB at 4k, 8k and 16k) but the compute buffer does not (536 MiB across the same + // range), so that would budget the KV and still miss n_groups - 1 compute buffers. + // Reserve what the fitter will not see in its per-device margin. + if (params_ctx.fit_params) { + reserve_extra_group_memory(params_ctx, has_draft, spec_mtp); + } + } + + 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 +1232,44 @@ struct server_context_impl { vocab = llama_model_get_vocab(model_tgt); - n_ctx = llama_n_ctx(ctx_tgt); + { + 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); + for (int j = 1; j < g; ++j) { + llama_free(groups[j]->ctx); + groups[j]->ctx = nullptr; + } + groups.clear(); + return false; + } + + groups[g]->threadpools = std::make_unique(); + groups[g]->threadpools->init(groups[g]->ctx, params_ctx); + + 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)); + } + } + + n_ctx = llama_n_ctx(ctx_tgt) * n_groups; add_bos_token = llama_vocab_get_add_bos(vocab); @@ -1074,31 +1278,44 @@ 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); + // 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; + // params_dft, not params_ctx: the draft has its own cpu params and this is the + // only place they can reach the draft context + grp.threadpools_dft = std::make_unique(); + grp.threadpools_dft->init(grp.ctx_dft, params_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); } @@ -1182,34 +1399,48 @@ struct server_context_impl { } // try speculative decoding - if (ctx_tgt_seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_NO) { - 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()); + 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++) { 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); - slot.spec = spec.get(); + // 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 = grp.ctx_dft; + slot.mem.init(grp.ctx, grp.ctx_dft); + + grp.slots.push_back(&slot); + slot.spec = grp.spec.get(); slot.n_ctx = n_ctx_slot; slot.mctx = mctx; @@ -1263,7 +1494,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 +1548,41 @@ struct server_context_impl { return true; } + 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; + }; + + 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; + } + + // mtmd_context is bound to one llama_context + if (has_mmproj) { + return refuse("multimodal (--mmproj)"); + } + + // common_init_from_params() applies it only to the context it creates + if (!params.control_vectors.empty()) { + return refuse("--control-vector"); + } + + // entering / leaving it 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 +1595,14 @@ 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; + } + if (groups.empty()) { + return; // no model loaded + } + update_slots(*groups[0]); }); queue_tasks.on_sleeping_state([this](bool sleeping) { handle_sleeping_state(sleeping); @@ -1424,6 +1699,222 @@ struct server_context_impl { return true; } + // Add what the extra pipeline groups will allocate to the fitter's per-device margin. Only the + // context and compute buffers for the target contexts, plus the draft contexts when speculation + // is on; the target weights are shared and already counted, and an MTP draft shares them too, so + // only a separate --model-draft adds weights per group. + void reserve_extra_group_memory(common_params & params_ctx, bool has_draft, bool spec_mtp) const { + const size_t n_extra = n_groups - 1; + + std::vector added(params_ctx.fit_params_target.size(), 0); + + // device order the margins are indexed by, taken from the target measurement below + std::vector devs_tgt; + + auto reserve = [&](common_params p, bool as_mtp, bool count_weights, bool is_target) { + auto mparams = common_model_params_to_llama(p); + auto cparams = common_context_params_to_llama(p); + if (as_mtp) { + cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + } + + std::vector devs; + uint32_t ngl = 0, n_ctx_train = 0, n_expert = 0; + + const auto mem = common_get_device_memory_data(p.model.path.c_str(), &mparams, &cparams, + devs, ngl, n_ctx_train, n_expert, GGML_LOG_LEVEL_ERROR); + + if (is_target) { + devs_tgt = devs; + } + + auto charge = [&](size_t id, const common_device_memory_data & md) { + if (id >= params_ctx.fit_params_target.size()) { + return; + } + size_t per_group = md.context + md.compute; + if (count_weights) { + per_group += md.model; + } + params_ctx.fit_params_target[id] += n_extra * per_group; + added[id] += n_extra * per_group; + }; + + // mem holds one entry per device plus a host entry at the back, while fit_params_target + // is indexed by device alone (common/fit.cpp builds margins over the nd devices). With + // no device at all that single margin is the host one, so only then is the host entry + // charged; otherwise charging it would bill host memory to a device's margin. + if (devs_tgt.empty()) { + if (!mem.empty()) { + charge(0, mem.back()); + } + return; + } + + for (size_t j = 0; j + 1 < mem.size() && j < devs.size(); ++j) { + // --device-draft may order the draft devices differently from the target, so match + // on device identity rather than position, the same way common_params_fit_impl maps + // its extra model. By position a draft allocation would be billed to another + // device's margin, approving a placement that then fails on the second context. + for (size_t id = 0; id < devs_tgt.size(); ++id) { + if (devs[j] == devs_tgt[id]) { + charge(id, mem[j]); + break; + } + } + } + }; + + reserve(params_ctx, false, false, true); + + if (has_draft || spec_mtp) { + reserve(common_base_params_to_speculative(params_ctx), spec_mtp, has_draft, false); + } + + for (size_t i = 0; i < added.size(); ++i) { + if (added[i] > 0) { + SRV_INF("pipeline groups: reserved %.0f MiB on device %zu for %zu extra context(s)\n", + added[i] / 1048576.0, i, n_extra); + } + } + } + + // 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; + + explicit engine_guard(server_context_impl * srv_) { + if (srv_->n_groups <= 1) { + return; + } + + srv = srv_; + + // 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) { + lks[g] = std::unique_lock(srv->groups[g]->mtx); + srv->groups[g]->n_pause_req++; + } + } + + void wait_for(int id_group) { + if (srv == nullptr) { + return; + } + + GGML_ASSERT(id_group >= 0 && id_group < (int) srv->groups.size()); + + 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; }); + } + + void wait_for_all() { + if (srv == nullptr) { + return; + } + + // a preceding wait_for() released every other group, and waiting on an unowned + // unique_lock is undefined. drop what is left and retake all of them in group order, + // the same order the constructor uses, so two guards still cannot deadlock. + for (size_t g = 0; g < srv->groups.size(); ++g) { + release(g); + } + + for (size_t g = 0; g < srv->groups.size(); ++g) { + lks[g].lock(); + srv->groups[g]->n_pause_req++; + } + + // a group cannot set busy without its own lock, so one already waited out stays idle + 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() { + if (srv == nullptr) { + return; + } + + for (size_t g = 0; g < srv->groups.size(); ++g) { + release(g); + } + } + + 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(); + } + }; + + void group_loop(server_group & grp) { + while (true) { + if (groups_stop.load(std::memory_order_relaxed)) { + return; + } + + if (update_slots(grp)) { + continue; + } + + std::unique_lock lk(grp.mtx); + grp.cv.wait_for(lk, std::chrono::milliseconds(5), + [&] { return groups_stop.load(std::memory_order_relaxed); }); + } + } + + void start_groups() { + if (n_groups <= 1) { + return; + } + + groups_stop.store(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; + } + + for (auto & grp : groups) { + std::unique_lock lk(grp->mtx); + groups_stop.store(true); + grp->cv.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(); @@ -1451,7 +1942,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; @@ -1545,25 +2036,30 @@ 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"); + need_cache_update = update_cache; - const int64_t t_start = ggml_time_us(); + return ret; + } - ret->prompt_save(*prompt_cache); + // 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"); - 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 @@ -1571,14 +2067,16 @@ 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) { + for (auto * slot_ptr : grp.slots) { + auto & slot = *slot_ptr; + if (slot.is_processing()) { continue; } @@ -1703,9 +2201,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()); @@ -1724,7 +2222,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; @@ -1888,12 +2386,12 @@ 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 }); } } 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); @@ -1911,7 +2409,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 }); } @@ -1983,7 +2481,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) { @@ -2008,7 +2509,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; @@ -2031,7 +2532,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( @@ -2061,7 +2562,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 +2602,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 +2648,12 @@ 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) { + if (slot.id_group != id_group) { + continue; + } if (!slot.is_processing() && slot.id != exclude_id_slot) { free_slots.push_back(&slot); } @@ -2170,14 +2674,17 @@ struct server_context_impl { SRV_TRC("launching slots for parent task id_task = %d with %zu child tasks\n", id_parent, parent_task.child_tasks.size()); - // to be called in case of failure to release all launched slots - auto release_slots = [this, id_parent]() { - for (auto & slot : slots) { - if (slot.is_processing() && ( - slot.task->id == id_parent || - slot.task->id_parent == id_parent + // to be called in case of failure to release all launched slots. + // only this group's slots: the parent and its children are guaranteed to share a group, and + // wait_for() released every other group's lock, so a global scan would read state and task + // while their workers write them + auto release_slots = [this, id_parent, &parent_slot]() { + for (auto * slot : groups[parent_slot.id_group]->slots) { + if (slot->is_processing() && ( + slot->task->id == id_parent || + slot->task->id_parent == id_parent )) { - slot.release(); + slot->release(); } } }; @@ -2244,10 +2751,10 @@ 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); + 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", @@ -2263,6 +2770,8 @@ struct server_context_impl { return false; } + engine_guard guard(this); + switch (task.type) { case SERVER_TASK_TYPE_COMPLETION: case SERVER_TASK_TYPE_INFILL: @@ -2279,7 +2788,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 @@ -2299,10 +2809,24 @@ struct server_context_impl { break; } + // 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) { + 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(); - std::vector child_slots = get_free_slots(n_child_tasks, slot->id); + // 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), + 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); queue_tasks.defer(std::move(task)); @@ -2318,6 +2842,8 @@ struct server_context_impl { } if (params_base.cache_idle_slots) { + guard.wait_for_all(); + for (auto & slot : slots) { if (!slot.is_processing()) { SLT_TRC(slot, "%s", "saving idle slot to prompt cache\n"); @@ -2340,6 +2866,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; } @@ -2360,6 +2887,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) { @@ -2441,6 +2971,8 @@ struct server_context_impl { break; } + guard.wait_for(slot->id_group); + const int64_t t_start = ggml_time_us(); std::string filename = task.slot_action.filename; @@ -2456,7 +2988,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); @@ -2491,6 +3023,8 @@ struct server_context_impl { break; } + guard.wait_for(slot->id_group); + const int64_t t_start = ggml_time_us(); std::string filename = task.slot_action.filename; @@ -2500,10 +3034,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 +3050,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"); } @@ -2556,6 +3090,8 @@ struct server_context_impl { break; } + guard.wait_for(slot->id_group); + // Erase token cache const size_t n_erased = slot->prompt.tokens.size(); @@ -2595,6 +3131,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) { @@ -2635,11 +3174,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 +3213,69 @@ struct server_context_impl { }; #endif - void update_slots() { + // each group reports and resets only its own counters, under its own lock: a reporter elected + // across groups would read and clear the others while their workers are still updating them + void prof_report(server_group & grp) { + auto & pr = grp.prof; + + const int64_t t_now = ggml_time_us(); + const int64_t t_last = pr.t_prof_last; + + if (t_last != 0 && t_now - t_last < 5 * 1000 * 1000) { + return; + } + + pr.t_prof_last = t_now; + + if (t_last == 0) { + return; // first call only arms the window + } + + const double t_win = (t_now - t_last) / 1000.0; // ms + + 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 piece %.3f proc %.3f send %.3f post %.3f\n", + grp.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_piece), per_tk(pr.t_proc), + per_tk(pr.t_send), per_tk(pr.t_post)); + + pr = server_group_prof(); + pr.t_prof_last = t_now; + } + + // 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; + auto & batch = grp.batch; + auto & slots = grp.slots; + + tls_group = &grp; + + std::unique_lock lk; + if (n_groups > 1) { + 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(grp); + } + #ifdef DEBUG_TIMINGS static int64_t t_prev = 0; int64_t t_start = ggml_time_us(); @@ -2692,8 +3293,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 +3303,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: each group drives its own loop, so the shared task loop need not keep spinning } try { scoped_timer t(t_pre_decode, n_pre_decode); - pre_decode(); + prof_timer tp(&grp.prof.t_pre, prof_on); + 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 +3361,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 +3378,33 @@ 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); + 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()); - 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; + + 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) { @@ -2832,8 +3446,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 @@ -2880,7 +3494,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; @@ -2900,16 +3514,16 @@ 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(); - 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(), @@ -2925,10 +3539,15 @@ struct server_context_impl { }); // generate the actual drafts (if any) + // note: only the main thread may yield to the task queue 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 @@ -2943,11 +3562,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 +3581,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 +3593,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 +3769,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 +3800,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,10 +3869,10 @@ 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); + 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); @@ -3293,7 +3912,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); @@ -3325,7 +3947,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 @@ -3371,7 +3993,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; @@ -3387,7 +4009,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(); @@ -3423,7 +4045,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 +4115,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,21 +4150,28 @@ 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; + + 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(); + 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 @@ -3559,15 +4188,53 @@ 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([&]() { - ret = llama_decode(ctx_tgt, batch_view); - if (ret == 0 && has_output) { + if (n_groups > 1) { + // 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; + 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() { + { + prof_timer tr(&grp->prof.t_relock, srv->prof_on); + lk->lock(); + } + grp->busy = false; + grp->cv.notify_all(); + } + } window(this, &grp, &lk); + + { + prof_timer ts(&grp.prof.t_submit, prof_on); + ret = llama_decode(ctx_tgt, batch_view); + } + // sync even with no output to read: ~decode_window clears busy, and a task thread that + // takes the guard must not touch ctx while the decode is still in flight + if (ret == 0) { + 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([&]() { + { + 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); + } + }); + } if (ret != 0) { { @@ -3593,14 +4260,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 +4277,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,7 +4286,7 @@ 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] @@ -3627,9 +4294,13 @@ struct server_context_impl { // ref: https://github.com/ggml-org/llama.cpp/pull/22728#issuecomment-4400925384 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"); @@ -3640,12 +4311,13 @@ struct server_context_impl { } // handle `n_cmpl > 1` tasks - when the main prompt is processed, activate all child tasks too - for (auto & slot : 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 +4337,12 @@ 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 * ctx_tgt = grp.ctx; + auto & slots = grp.slots; + auto & spec = grp.spec; + (void) ctx_tgt; + // 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; @@ -3721,7 +4398,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; @@ -3737,8 +4414,10 @@ struct server_context_impl { llama_token id; { 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; @@ -3759,14 +4438,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); @@ -3821,13 +4508,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()); @@ -3840,7 +4527,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); } @@ -3874,7 +4561,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; @@ -3915,38 +4602,48 @@ 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; + { + 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; } // 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) { - metrics.n_decode++; - for (const auto & slot : slots) { - if (slot.is_processing()) { - metrics.n_busy_slots++; + void metrics_post_decode(server_group & grp, int32_t off, int32_t n_tokens, bool has_output) { + auto & batch = grp.batch; + + { + // 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++; + 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 @@ -3969,11 +4666,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 @@ -3989,13 +4686,13 @@ struct server_context_impl { } // flush any queued prompt metrics if all slots are now idle - void metrics_flush_idle() { - if (n_prompt_queued == 0) { + void metrics_flush_idle(server_group & grp) { + if (grp.n_prompt_queued == 0) { return; } - llama_synchronize(ctx_tgt); - metrics_flush_prompt(); + llama_synchronize(grp.ctx); + metrics_flush_prompt(grp); } void metrics_on_prediction(const server_slot & slot) { @@ -4003,6 +4700,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); @@ -4035,7 +4734,16 @@ bool server_context::load_model(common_params & params) { void server_context::start_loop() { auto & params = impl->params_base; + + 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..d22756f8d4f8 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -110,6 +110,10 @@ struct server_context { // note: must be set before load_model() is called void set_state_callback(server_state_callback_t callback); + + // 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-models.cpp b/tools/server/server-models.cpp index db0fac99527b..251462451a5b 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -400,6 +400,8 @@ void server_model_meta::update_caps() { } } +int server_get_pipeline_groups(); + // // server_models // @@ -1026,6 +1028,19 @@ void server_models::load(const std::string & name, const load_options & opts) { std::vector child_args = inst.meta.args; // copy std::vector child_env = base_env; // copy child_env.push_back("LLAMA_SERVER_ROUTER_PORT=" + std::to_string(base_params.port)); + // the router strips --pipeline-groups before base_preset is built, so it cannot ride + // the preset like other options; hand it over explicitly. base_env is a copy of our own + // environment, so an inherited LLAMA_ARG_PIPELINE_GROUPS would otherwise reach the child + // unchanged: appending cannot override it, because execve leaves duplicates in place and + // getenv() returns the first entry, and an explicit --pipeline-groups 1 would append + // nothing at all. Drop any inherited entry first, then always set the resolved value. + { + static const std::string pg_prefix = "LLAMA_ARG_PIPELINE_GROUPS="; + child_env.erase(std::remove_if(child_env.begin(), child_env.end(), + [](const std::string & e) { return e.rfind(pg_prefix, 0) == 0; }), + child_env.end()); + child_env.push_back(pg_prefix + std::to_string(server_get_pipeline_groups())); + } if (opts.mode == SERVER_CHILD_MODE_DOWNLOAD) { inst.meta.status = SERVER_MODEL_STATUS_DOWNLOADING; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 5fe2729ba1b2..07dd7309dd42 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,12 @@ static std::function shutdown_handler; static std::atomic_flag is_terminating = ATOMIC_FLAG_INIT; +// set from params right after parsing, read by server_models::spawn to pass the setting on +static int g_pipeline_groups = 1; + +int server_get_pipeline_groups(); +int server_get_pipeline_groups() { return g_pipeline_groups; } + 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 @@ -106,6 +115,8 @@ int llama_server(int argc, char ** argv) { return 1; } + g_pipeline_groups = params.n_pipeline_groups; + llama_backend_init(); llama_numa_init(params.numa); @@ -168,6 +179,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(params.n_pipeline_groups); server_http_context ctx_http; if (!ctx_http.init(params)) {