swap worker <--> main thread design

tools/server/server-context.cpp

@@ -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

@@ -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

@@ -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);
     }