From fb16b8f76fb1d4a1902aba0eff6810677a60aa69 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 03:54:53 -0700 Subject: [PATCH 1/2] server: deliver a result to the one thread waiting for it server_response kept every pending result in one vector behind one condition variable. Each result was pushed after a linear walk of the waiting id set, then notify_all woke every waiting HTTP thread, and each of them took the same mutex and scanned the whole vector before going back to sleep. With N slots generating that is N wakeups and N vector scans per token, so N^2 per decode step, all of it contending for the mutex the decode thread needs to send the next token. Results are now queued on a per reader waiter. Ids registered together share one waiter, so a send is an O(1) lookup followed by a push and a wakeup of exactly the thread that asked for that task. Order is preserved: the waiter holds a FIFO and recv() takes the front, which is what scanning the shared vector from the start did. A reader whose ids have already been removed from the waiting list still waits out the poll interval it asked for rather than returning at once, so a caller that keeps polling does not spin, and the blocking recv() re-checks the running flag on a bounded wait so terminate() cannot leave it parked. Measured with llama-server at 32 slots, one request per slot, 128 prompt and 256 generated tokens: queue_results.send() 119.0 us to 3.6 us per call, and the whole result path per decode step 3.97 ms to 0.34 ms. --- tools/server/server-queue.cpp | 139 ++++++++++++++++++++++------------ tools/server/server-queue.h | 36 +++++++-- 2 files changed, 121 insertions(+), 54 deletions(-) diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 78169e9a5d8..de663ddabf5 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -387,79 +387,114 @@ void server_queue::cleanup_pending_task(int id_target) { // void server_response::add_waiting_task_id(int id_task) { - RES_DBG("add task %d to waiting list. current waiting = %d (before add)\n", id_task, (int) waiting_task_ids.size()); - std::unique_lock lock(mutex_results); - waiting_task_ids.insert(id_task); + + RES_DBG("add task %d to waiting list. current waiting = %d (before add)\n", id_task, (int) waiting.size()); + + waiting.emplace(id_task, std::make_shared()); } void server_response::add_waiting_task_ids(const std::unordered_set & id_tasks) { std::unique_lock lock(mutex_results); + // one waiter for the whole set: these ids belong to a single reader, which waits for any + // of them at a time + auto w = std::make_shared(); + for (const auto & id_task : id_tasks) { - RES_DBG("add task %d to waiting list. current waiting = %d (before add)\n", id_task, (int) waiting_task_ids.size()); - waiting_task_ids.insert(id_task); + RES_DBG("add task %d to waiting list. current waiting = %d (before add)\n", id_task, (int) waiting.size()); + waiting.emplace(id_task, w); } } void server_response::remove_waiting_task_id(int id_task) { - RES_DBG("remove task %d from waiting list. current waiting = %d (before remove)\n", id_task, (int) waiting_task_ids.size()); - std::unique_lock lock(mutex_results); - waiting_task_ids.erase(id_task); - // make sure to clean up all pending results - queue_results.erase( - std::remove_if(queue_results.begin(), queue_results.end(), [id_task](const server_task_result_ptr & res) { + + RES_DBG("remove task %d from waiting list. current waiting = %d (before remove)\n", id_task, (int) waiting.size()); + + auto it = waiting.find(id_task); + if (it == waiting.end()) { + return; + } + + // make sure to clean up all pending results of this task, the waiter may still be held by + // the other ids of the same reader + auto & results = it->second->results; + results.erase( + std::remove_if(results.begin(), results.end(), [id_task](const server_task_result_ptr & res) { return res->id == id_task; }), - queue_results.end()); + results.end()); + + waiting.erase(it); } void server_response::remove_waiting_task_ids(const std::unordered_set & id_tasks) { std::unique_lock lock(mutex_results); for (const auto & id_task : id_tasks) { - RES_DBG("remove task %d from waiting list. current waiting = %d (before remove)\n", id_task, (int) waiting_task_ids.size()); - waiting_task_ids.erase(id_task); + RES_DBG("remove task %d from waiting list. current waiting = %d (before remove)\n", id_task, (int) waiting.size()); + waiting.erase(id_task); + } +} + +server_response::waiter_ptr server_response::find_waiter(const std::unordered_set & id_tasks) const { + for (const auto & id_task : id_tasks) { + auto it = waiting.find(id_task); + if (it != waiting.end()) { + return it->second; + } } + + return nullptr; } server_task_result_ptr server_response::recv(const std::unordered_set & id_tasks) { + std::unique_lock lock(mutex_results); + + auto w = find_waiter(id_tasks); + GGML_ASSERT(w && "recv() called for task ids that are not in the waiting list"); + while (true) { - std::unique_lock lock(mutex_results); - condition_results.wait(lock, [&]{ - if (!running) { - RES_DBG("%s : queue result stop\n", "recv"); - std::terminate(); // we cannot return here since the caller is HTTP code - } - return !queue_results.empty(); - }); + if (!running) { + RES_DBG("%s : queue result stop\n", "recv"); + std::terminate(); // we cannot return here since the caller is HTTP code + } - for (size_t i = 0; i < queue_results.size(); i++) { - if (id_tasks.find(queue_results[i]->id) != id_tasks.end()) { - server_task_result_ptr res = std::move(queue_results[i]); - queue_results.erase(queue_results.begin() + i); - return res; - } + if (!w->results.empty()) { + server_task_result_ptr res = std::move(w->results.front()); + w->results.pop_front(); + return res; } + + // bounded, so a terminate() that lands after the id was removed from the map still + // gets noticed here + w->cv.wait_for(lock, std::chrono::seconds(1)); } // should never reach here } server_task_result_ptr server_response::recv_with_timeout(const std::unordered_set & id_tasks, int timeout) { - while (true) { - std::unique_lock lock(mutex_results); + std::unique_lock lock(mutex_results); - for (int i = 0; i < (int) queue_results.size(); i++) { - if (id_tasks.find(queue_results[i]->id) != id_tasks.end()) { - server_task_result_ptr res = std::move(queue_results[i]); - queue_results.erase(queue_results.begin() + i); - return res; - } + auto w = find_waiter(id_tasks); + if (!w) { + // the tasks are no longer in the waiting list, so no result can arrive for them. + // wait out the timeout anyway, so the caller sees the poll interval it asked for + // instead of a busy loop + condition_gone.wait_for(lock, std::chrono::seconds(timeout)); + return nullptr; + } + + while (true) { + if (!w->results.empty()) { + server_task_result_ptr res = std::move(w->results.front()); + w->results.pop_front(); + return res; } - std::cv_status cr_res = condition_results.wait_for(lock, std::chrono::seconds(timeout)); + std::cv_status cr_res = w->cv.wait_for(lock, std::chrono::seconds(timeout)); if (!running) { RES_DBG("%s : queue result stop\n", __func__); std::terminate(); // we cannot return here since the caller is HTTP code @@ -481,31 +516,39 @@ void server_response::send(server_task_result_ptr && result) { RES_DBG("sending result for task id = %d\n", result->id); std::unique_lock lock(mutex_results); - for (const auto & id_task : waiting_task_ids) { - if (result->id == id_task) { - RES_DBG("task id = %d pushed to result queue\n", result->id); - queue_results.emplace_back(std::move(result)); - condition_results.notify_all(); - return; - } + auto it = waiting.find(result->id); + if (it == waiting.end()) { + return; } + + RES_DBG("task id = %d pushed to result queue\n", result->id); + + auto & w = *it->second; + + w.results.emplace_back(std::move(result)); + w.cv.notify_one(); } void server_response::broadcast(server_task_result_ptr && result) { std::unique_lock lock(mutex_results); - for (const auto & id_task : waiting_task_ids) { + for (const auto & [id_task, w] : waiting) { RES_DBG("task id = %d pushed to result queue\n", id_task); server_task_result_ptr res_copy(result->clone()); res_copy->id = id_task; // override id with target task id - queue_results.emplace_back(std::move(res_copy)); + w->results.emplace_back(std::move(res_copy)); + w->cv.notify_one(); } - condition_results.notify_all(); } void server_response::terminate() { + std::unique_lock lock(mutex_results); running = false; - condition_results.notify_all(); + for (const auto & [id_task, w] : waiting) { + (void) id_task; + w->cv.notify_all(); + } + condition_gone.notify_all(); } // diff --git a/tools/server/server-queue.h b/tools/server/server-queue.h index e17733a743f..10109d7c258 100644 --- a/tools/server/server-queue.h +++ b/tools/server/server-queue.h @@ -5,10 +5,12 @@ #include #include #include +#include #include #include -#include +#include #include +#include // struct for managing server tasks // in most cases, use server_response_reader to post new tasks and retrieve results @@ -155,14 +157,36 @@ struct server_response { private: bool running = true; - // for keeping track of all tasks waiting for the result - std::unordered_set waiting_task_ids; + // One waiter per reader, shared by every task id that reader registered in one call. + // Results are queued on the waiter that owns the id, so sending a result wakes only the + // thread that is waiting for it, and that thread finds its result without searching. + // + // Previously there was a single result vector and a single condition variable: every + // result woke every waiting HTTP thread, and each of them re-took the mutex and scanned + // the whole vector before going back to sleep. With N slots generating that is N wakeups + // and N scans per token, i.e. N^2 per decode step, all of it contending with the decode + // thread for the same mutex. + struct waiter { + std::condition_variable cv; + + // FIFO, so results are handed out in the order they were sent, as before + std::deque results; + }; - // the main result queue (using ptr for polymorphism) - std::vector queue_results; + using waiter_ptr = std::shared_ptr; + + // task id --> the waiter that is expecting its results + std::unordered_map waiting; std::mutex mutex_results; - std::condition_variable condition_results; + + // only used to park a reader whose ids are no longer in the waiting list, so that it + // still returns after the timeout it asked for rather than spinning + std::condition_variable condition_gone; + + // all ids registered together share one waiter, so the first hit is the right one + // must be called with mutex_results held + waiter_ptr find_waiter(const std::unordered_set & id_tasks) const; public: // add the id_task to the list of tasks waiting for response From 755fb62336b72bfab73bb0597f1666a6276052e7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 03:54:53 -0700 Subject: [PATCH 2/2] server: only keep per token probabilities when the request asked for them server_slot::generated_token_probs is read in exactly one place, send_final_response(), and only under n_probs > 0. Every other request still pushed a completion_token_output per token, each with a heap allocated string, into a list that grows for the whole generation and is then discarded. The output is unchanged: with n_probs <= 0 nothing ever reads the list. --- tools/server/server-context.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b..1b208631c79 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -431,6 +431,13 @@ struct server_slot { return; } + // the buffer is only ever read by send_final_response(), and only when the request + // asked for per-token probabilities. Without them every token still copied a string + // and a vector into a list that grows for the whole generation and is then dropped. + if (task->params.sampling.n_probs <= 0) { + return; + } + generated_token_probs.push_back(token); }