Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

server: re-design yield_to_queue thread model (#27133)

Browse Source 
* run common_speculative_process in worker

* swap worker <--> main thread design
This commit is contained in:
Xuan-Son Nguyen
2026-08-15 16:48:40 +02:00
committed by GitHub
parent adb55e5148
commit 22b8e310b9
3 changed files with 78 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/server/server-context.cpp
+17 -8
View File
@@ -2947,8 +2947,10 @@ private:
});
// generate the actual drafts (if any)
{
common_speculative_draft(spec.get());
if (!drafting.empty()) {
queue_tasks.yield_to_queue([&]() {
common_speculative_draft(spec.get());
});
}
// make checkpoints if needed
@@ -3578,8 +3580,8 @@ private:
has_output |= batch.tokens[i].output;
}
// decode on the worker thread, so we can still handle metrics tasks while waiting
// note: the sync is done here too, so that the wait also happens off the main thread
// 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);
@@ -3644,11 +3646,18 @@ private:
// TODO: avoid restoring the draft context and re-evaluating the drafted tokens when not needed [TAG_SPEC_AVOID_DRAFT_REEVAL]
// for now, always re-evaluate for simplicity
// ref: https://github.com/ggml-org/llama.cpp/pull/22728#issuecomment-4400925384
if (!common_speculative_process(spec.get(), batch_view)) {
SRV_ERR("%s", "failed to process speculative batch\n");
if (spec) {
bool ok = true;
queue_tasks.yield_to_queue([&]() {
ok = common_speculative_process(spec.get(), batch_view);
});
// TODO: handle error
throw std::runtime_error("failed to process speculative batch");
if (!ok) {
SRV_ERR("%s", "failed to process speculative batch\n");
// TODO: handle error
throw std::runtime_error("failed to process speculative batch");
}
}
// handle `n_cmpl > 1` tasks - when the main prompt is processed, activate all child tasks too
tools/server/server-queue.cpp
+54 -32
View File
@@ -150,31 +150,46 @@ bool server_queue::process_new_tasks(bool is_yielding) {
void server_queue::worker_loop() {
while (true) {
std::function<void()> work;
{
std::unique_lock<std::mutex> lock(mutex_tasks);
// wait on busy instead of yielding - busy stays set even when the yield already ended
worker.cv.wait(lock, [&]{
return worker.stop || worker.work != nullptr;
return worker.stop || worker.busy;
});
if (worker.stop) {
return;
}
work = std::move(worker.work);
worker.work = nullptr;
}
// note: do not hold any lock here, work() may post new tasks
std::exception_ptr exception;
try {
work();
} catch (...) {
exception = std::current_exception();
// process tasks while the yield is active
while (true) {
bool terminated = false;
try {
// note: do not hold any lock here, the callback may post new tasks
terminated = process_new_tasks(true);
} catch (...) {
std::unique_lock<std::mutex> lock(mutex_tasks);
worker.exception = std::current_exception();
break;
}
std::unique_lock<std::mutex> lock(mutex_tasks);
if (terminated || worker.stop || !worker.yielding) {
break;
}
if (!queue_tasks.empty()) {
continue; // a new task arrived in the meantime
}
condition_tasks.wait(lock, [&]{
return worker.stop || !running || !worker.yielding || !queue_tasks.empty();
});
}
// signal completion to yield_to_queue()
std::unique_lock<std::mutex> lock(mutex_tasks);
worker.exception = std::move(exception);
worker.busy = false;
// signal to yield_to_queue() that no more tasks will be processed
{
std::unique_lock<std::mutex> lock(mutex_tasks);
worker.busy = false;
}
condition_tasks.notify_all();
}
}
@@ -188,6 +203,7 @@ void server_queue::worker_stop() {
worker.stop = true;
}
worker.cv.notify_one();
condition_tasks.notify_all();
worker.thread.join();
}
@@ -199,29 +215,29 @@ void server_queue::yield_to_queue(std::function<void()> && work) {
{
std::unique_lock<std::mutex> lock(mutex_tasks);
GGML_ASSERT(!worker.busy && "yield_to_queue() cannot be nested");
worker.busy = true;
worker.work = std::move(work);
worker.busy = true;
worker.yielding = true;
}
worker.cv.notify_one();
while (true) {
// note: on terminate this is a no-op, but we still wait for the work to finish
process_new_tasks(true);
std::unique_lock<std::mutex> lock(mutex_tasks);
// declined tasks are moved to queue_tasks_unhandled, so a non-empty queue always has something new
condition_tasks.wait(lock, [&]{
return !worker.busy || (running && !queue_tasks.empty());
});
if (!worker.busy) {
break;
}
// run the work on the current thread, so that all ggml compute stays on the same thread
std::exception_ptr exception;
try {
work();
} catch (...) {
exception = std::current_exception();
}
std::exception_ptr exception;
{
std::unique_lock<std::mutex> lock(mutex_tasks);
// the yield is over, wait for the worker to finish its current task
worker.yielding = false;
condition_tasks.notify_all();
condition_tasks.wait(lock, [&]{
return !worker.busy;
});
// put the declined tasks back, keeping their order
while (!queue_tasks_unhandled.empty()) {
queue_tasks.push_front(std::move(queue_tasks_unhandled.back()));
@@ -231,8 +247,12 @@ void server_queue::yield_to_queue(std::function<void()> && work) {
// make sure to avoid idle timeout here
time_last_task = ggml_time_ms();
// the worker is idle now, take the exception it may have left behind
std::swap(exception, worker.exception);
// an exception from work() takes precedence over the one from the worker
if (!exception) {
std::swap(exception, worker.exception);
} else {
worker.exception = nullptr;
}
}
QUE_DBG("%s", "done yielding to queue\n");
@@ -249,7 +269,9 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
// spawn the worker thread used by yield_to_queue()
GGML_ASSERT(!worker.thread.joinable() && "start_loop() is already running");
worker.stop = false;
worker.stop = false;
worker.busy = false;
worker.yielding = false;
worker.thread = std::thread([this]() { worker_loop(); });
constexpr auto max_wait_time = std::chrono::seconds(1);
tools/server/server-queue.h
+7 -6
View File
@@ -33,11 +33,11 @@ private:
// used by yield_to_queue, all fields are guarded by mutex_tasks
struct worker_t {
std::thread thread;
std::condition_variable cv; // the worker sleeps on this until there is work
std::function<void()> work; // pending work, picked up by the thread
std::exception_ptr exception; // exception thrown by work(), if any
bool stop = false;
bool busy = false;
std::condition_variable cv; // the worker sleeps on this until a yield starts
std::exception_ptr exception; // exception thrown while processing tasks, if any
bool stop = false;
bool busy = false; // set by yield_to_queue(), cleared by the worker once it is done processing tasks
bool yielding = false; // work() is still running on the start_loop() thread
};
worker_t worker;
@@ -93,7 +93,7 @@ public:
*/
void start_loop(int64_t idle_sleep_ms = -1);
// run work() on a separate thread, while the current thread calls process_new_tasks
// while waiting for work() to finish, run process_new_tasks on the worker thread
// returns once work() is done (may throw exceptions)
// must be called from start_loop() thread (ideally inside callback_update_slots)
// use case: return metrics while encode/decode is running
@@ -116,6 +116,7 @@ public:
// the second argument tells whether the queue is currently yielding (see yield_to_queue)
// only then may the callback return false to decline the task, and it must leave it
// untouched, so that it can be put back in the queue later
// note: while yielding, the callback runs on worker thread, not main thread
void on_new_task(std::function<bool(server_task &&, bool)> callback) {
callback_new_task = std::move(callback);
}