server: run sampling in a threadpool

tools/server/server-common.cpp

@@ -1583,3 +1583,54 @@ server_tokens format_prompt_rerank(
 
     return result;
 }
+
+//
+// threadpool
+//
+
+server_threadpool::~server_threadpool() {
+    {
+        std::lock_guard<std::mutex> lock(mtx);
+        stop = true;
+    }
+    cv.notify_all();
+    for (auto & t : threads) t.join();
+}
+
+void server_threadpool::init(int n) {
+    for (int i = 0; i < n; i++) {
+        threads.emplace_back([this]() {
+            while (true) {
+                std::function<void()> task;
+                {
+                    std::unique_lock<std::mutex> lock(mtx);
+                    cv.wait(lock, [this]() { return stop || !tasks.empty(); });
+                    if (stop && tasks.empty()) return;
+                    task = std::move(tasks.front());
+                    tasks.pop();
+                }
+                task();
+                {
+                    std::lock_guard<std::mutex> lock(mtx);
+                    pending--;
+                }
+                cv_done.notify_one();
+            }
+        });
+    }
+}
+
+void server_threadpool::enqueue(std::function<void()> fn) {
+    {
+        std::lock_guard<std::mutex> lock(mtx);
+        GGML_ASSERT(!stop && !threads.empty());
+        tasks.push(std::move(fn));
+        pending++;
+    }
+    cv.notify_one();
+}
+
+void server_threadpool::wait_all() {
+    std::unique_lock<std::mutex> lock(mtx);
+    cv_done.wait(lock, [this]() { return pending == 0; });
+}

tools/server/server-common.h

@@ -12,6 +12,11 @@
 #include <string>
 #include <vector>
 #include <cinttypes>
+#include <thread>
+#include <mutex>
+#include <condition_variable>
+#include <queue>
+#include <functional>
 
 using json = nlohmann::ordered_json;
 
@@ -370,3 +375,35 @@ server_tokens format_prompt_rerank(
         mtmd_context * mctx,
         const std::string & query,
         const std::string & doc);
+
+//
+// threadpool utils
+// to be used for multi-threaded sampling
+//
+
+struct server_threadpool {
+    std::vector<std::thread> threads;
+    std::queue<std::function<void()>> tasks;
+    std::mutex mtx;
+    std::condition_variable cv;
+    std::condition_variable cv_done;
+    int pending = 0;
+    bool stop = false;
+
+    ~server_threadpool();
+    void init(int n);
+
+    template<typename T>
+    void run_all(std::vector<T> & tasks, std::function<void(T&)> handler) {
+        for (auto & item : tasks) {
+            enqueue([&handler, &item]() {
+                handler(item);
+            });
+        }
+        wait_all();
+    }
+
+private:
+    void enqueue(std::function<void()> fn);
+    void wait_all();
+};

tools/server/server-context.cpp

@@ -866,6 +866,9 @@ public:
     // note: chat_params must not be refreshed upon existing sleeping state
     server_chat_params chat_params;
 
+    // threadpool for parallel sampling
+    server_threadpool threadpool;
+
     server_state_callback_t callback_state = [](server_state, json) -> void {};
 
     server_context_impl() {
@@ -1457,6 +1460,8 @@ private:
 
         metrics.init();
 
+        threadpool.init(params_base.cpuparams.n_threads);
+
         if (params_base.cache_idle_slots) {
             if (params_base.cache_ram_mib == 0) {
                 SRV_WRN("%s", "--cache-idle-slots requires --cache-ram, disabling\n");
@@ -3699,6 +3704,12 @@ private:
         return true;
     }
 
+    struct sampling_task {
+        server_slot * slot;
+        int32_t tok_idx;
+        llama_token sampled_id; // result
+    };
+
     void post_decode(int32_t n_batch_tokens, int32_t off, llama_batch & batch_view) {
         // for checking if a given batch index is inside batch_view
         auto is_inside_view = [&](int32_t idx) {
@@ -3720,6 +3731,9 @@ private:
                 slot.task->params.sampling.preserved_tokens.find(token) != slot.task->params.sampling.preserved_tokens.end();
         };
 
+        std::vector<sampling_task> smpl_tasks;
+        smpl_tasks.reserve(slots.size());
+
         iterate(slots, [&](server_slot & slot) {
             // optionally send prompt processing progress
             if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_DONE_PROMPT) {
@@ -3767,6 +3781,17 @@ private:
 
             // shifted according to the current sub-batch
             const int tok_idx = slot.i_batch - off;
+            smpl_tasks.push_back({&slot, tok_idx, LLAMA_TOKEN_NULL});
+        });
+
+        // run multiple sampling tasks in parallel
+        threadpool.run_all<sampling_task>(smpl_tasks, [](sampling_task & task) {
+            task.sampled_id = common_sampler_sample(task.slot->smpl.get(), task.slot->ctx_tgt, task.tok_idx);
+        });
+
+        iterate(slots, [&](server_slot & slot) {
+            auto & smpl_task = smpl_tasks[slot.id];
+            auto tok_idx = smpl_task.tok_idx;
 
             llama_token id;
             {