diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 61695b6d89..d88d23337a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -688,97 +688,99 @@ struct server_slot { other.prompt = prompt.clone(); other.init_sampler(); } - - // returns 0 on success - // caller need to update prompt.tokens after a successful call to keep track of the processing progress - int process_mtmd_chunk(size_t idx, size_t & n_tokens_out) { - GGML_ASSERT(mctx); - const auto & input_tokens = task->tokens; - const auto & chunk = input_tokens.find_chunk(idx); - int32_t res = 0; - - auto try_decode = [&]() -> int32_t { - if (mbatch) { - float * embd = mtmd_batch_get_output_embd(mbatch.get(), chunk.get()); - if (embd) { - void * cb_data = spec; - static auto cb = [](llama_batch batch, void * user_data) { - common_speculative * spec = static_cast(user_data); - if (!common_speculative_process(spec, batch)) { - return 1; - } - return 0; - }; - - llama_pos new_n_past; // unused for now - res = mtmd_helper_decode_image_chunk( - mctx, - ctx_tgt, - chunk.get(), - embd, - prompt.tokens.pos_next(), - id, - llama_n_batch(ctx_tgt), - &new_n_past, - cb, - cb_data - ); - if (res != 0) { - SLT_ERR(*this, "failed to decode mtmd chunk, idx = %zu, res = %d\n", idx, res); - return -1; - } - n_tokens_out = mtmd_input_chunk_get_n_tokens(chunk.get()); - return 0; // success - } - } - return 1; // (non-error) need to create & encode batch - }; - - // if the batch is already exist, try searching & encode - res = try_decode(); - if (res == 0) { - return 0; - } - if (res < 0) { - // fatal error - return res; - } - - // otherwise, the batch is either uninitialized or is used up - // we need to create & encode a new batch - mbatch.reset(mtmd_batch_init(mctx)); - res = mtmd_batch_add_chunk(mbatch.get(), chunk.get()); - GGML_ASSERT(res == 0); // we should never have an empty batch - - // try batching as much as possible - int n_added = 1; - size_t idx_cur = idx; - while (res == 0) { - auto [next_chunk, next_idx] = input_tokens.find_next_media_chunk(idx_cur); - if (next_chunk == nullptr) { - break; - } - res = mtmd_batch_add_chunk(mbatch.get(), next_chunk->get()); - n_added += (res == 0 ? 1 : 0); - idx_cur = next_idx; - SLT_DBG(*this, "try adding media chunk idx = %zu to batch, res = %d\n", next_idx, res); - // if res != 0, batch is full or chunk is not compatible -> this loop breaks - } - - // TODO @ngxson : move this log line to debug when it become more stable - SLT_TRC(*this, "encoding mtmd batch from idx = %zu, n_chunks = %d\n", idx, n_added); - - res = mtmd_batch_encode(mbatch.get()); - if (res != 0) { - SLT_ERR(*this, "failed to encode mtmd batch for chunk idx = %zu, res = %d\n", idx, res); - return -1; - } - - return try_decode(); - } }; +// returns 0 on success +// caller need to update prompt.tokens after a successful call to keep track of the processing progress +// note: this is not a member of server_slot because we want to run it inside yield_to_queue +// slot is passed as const to avoid accidental modification of the slot state +// only llama_context / mtmd_context / mbatch are allowed to be used inside +static int process_mtmd_chunk(const server_slot & slot, mtmd::batch_ptr & mbatch, size_t idx, size_t & n_tokens_out) { + GGML_ASSERT(slot.mctx); + const auto & mctx = slot.mctx; + const auto & input_tokens = slot.task->tokens; + const auto & chunk = input_tokens.find_chunk(idx); + int32_t res = 0; + auto try_decode = [&]() -> int32_t { + if (mbatch) { + float * embd = mtmd_batch_get_output_embd(mbatch.get(), chunk.get()); + if (embd) { + void * cb_data = slot.spec; + static auto cb = [](llama_batch batch, void * user_data) { + common_speculative * spec = static_cast(user_data); + if (!common_speculative_process(spec, batch)) { + return 1; + } + return 0; + }; + + llama_pos new_n_past; // unused for now + res = mtmd_helper_decode_image_chunk( + mctx, + slot.ctx_tgt, + chunk.get(), + embd, + slot.prompt.tokens.pos_next(), + slot.id, + llama_n_batch(slot.ctx_tgt), + &new_n_past, + cb, + cb_data + ); + if (res != 0) { + SLT_ERR(slot, "failed to decode mtmd chunk, idx = %zu, res = %d\n", idx, res); + return -1; + } + n_tokens_out = mtmd_input_chunk_get_n_tokens(chunk.get()); + return 0; // success + } + } + return 1; // (non-error) need to create & encode batch + }; + + // if the batch is already exist, try searching & encode + res = try_decode(); + if (res == 0) { + return 0; + } + if (res < 0) { + // fatal error + return res; + } + + // otherwise, the batch is either uninitialized or is used up + // we need to create & encode a new batch + mbatch.reset(mtmd_batch_init(mctx)); + res = mtmd_batch_add_chunk(mbatch.get(), chunk.get()); + GGML_ASSERT(res == 0); // we should never have an empty batch + + // try batching as much as possible + int n_added = 1; + size_t idx_cur = idx; + while (res == 0) { + auto [next_chunk, next_idx] = input_tokens.find_next_media_chunk(idx_cur); + if (next_chunk == nullptr) { + break; + } + res = mtmd_batch_add_chunk(mbatch.get(), next_chunk->get()); + n_added += (res == 0 ? 1 : 0); + idx_cur = next_idx; + SLT_DBG(slot, "try adding media chunk idx = %zu to batch, res = %d\n", next_idx, res); + // if res != 0, batch is full or chunk is not compatible -> this loop breaks + } + + // TODO @ngxson : move this log line to debug when it become more stable + SLT_TRC(slot, "encoding mtmd batch from idx = %zu, n_chunks = %d\n", idx, n_added); + + res = mtmd_batch_encode(mbatch.get()); + if (res != 0) { + SLT_ERR(slot, "failed to encode mtmd batch for chunk idx = %zu, res = %d\n", idx, res); + return -1; + } + + return try_decode(); +} // // server_context_impl (private implementation) @@ -854,9 +856,6 @@ private: int slots_debug = 0; int n_empty_consecutive = 0; - // true while a decode runs on the yield_to_queue() worker thread - bool is_decoding = false; - std::unique_ptr prompt_cache; server_metrics metrics; @@ -1357,8 +1356,8 @@ private: GGML_ASSERT(!sleeping); // wiring up server queues - queue_tasks.on_new_task([this](server_task && task) { - return process_single_task(std::move(task)); + queue_tasks.on_new_task([this](server_task && task, bool is_yielding) { + return process_single_task(std::move(task), is_yielding); }); queue_tasks.on_update_slots([this]() { update_slots(); @@ -2290,9 +2289,9 @@ private: } // returns false to decline the task, it is offered again after the decode is done - bool process_single_task(server_task && task) { - // during encoding / decoding, only accessing metrics is safe - if (is_decoding && task.type != SERVER_TASK_TYPE_METRICS) { + bool process_single_task(server_task && task, bool is_yielding) { + // while yielding, an encode / decode is running and only accessing metrics is safe + if (is_yielding && task.type != SERVER_TASK_TYPE_METRICS) { SRV_DBG("decoding, decline task, id_task = %d\n", task.id); return false; } @@ -3397,11 +3396,9 @@ private: // encode on the worker thread, so we can still handle metrics tasks size_t n_tokens_out = 0; int32_t res = 0; - is_decoding = true; queue_tasks.yield_to_queue([&]() { - res = slot.process_mtmd_chunk(cur_token_idx, n_tokens_out); + res = process_mtmd_chunk(slot, slot.mbatch, cur_token_idx, n_tokens_out); }); - is_decoding = false; if (res != 0) { SLT_ERR(slot, "failed to process mtmd chunk, res = %d\n", res); @@ -3582,15 +3579,14 @@ private: } // 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 int ret = 0; - is_decoding = true; queue_tasks.yield_to_queue([&]() { ret = llama_decode(ctx_tgt, batch_view); if (ret == 0 && has_output) { llama_synchronize(ctx_tgt); } }); - is_decoding = false; if (ret != 0) { { @@ -3642,7 +3638,7 @@ private: return false; // retry with the updated n_batch } else { // success, apply batch metrics - metrics_post_decode(off, batch_view.n_tokens); + metrics_post_decode(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] @@ -3955,7 +3951,8 @@ private: n_prompt_queued = 0; } - void metrics_post_decode(int32_t off, int32_t n_tokens) { + // 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()) { @@ -3968,13 +3965,10 @@ private: // note: a slot can be released before we get here, which clears its stats // the tokens were still computed, counted in the global metrics, not in slot uint64_t n_prompt_tokens = 0; - bool has_output = false; for (int i = off; i < off + n_tokens; ++i) { const auto & t = batch.tokens[i]; - has_output |= t.output; - if (!t.is_prompt) { continue; // generated tokens are handled after sampling } @@ -3990,14 +3984,12 @@ private: metrics_queue_prompt(n_prompt_tokens); if (has_output) { - // sync if we have at least one output in batch - // so that we can calculate the timings correctly - llama_synchronize(ctx_tgt); + // the context is already synchronized, so the timings are correct metrics_flush_prompt(); } // advance the prompt timing of the slots that had tokens in this batch - // note: a second pass, it must run after the sync above to reflect the compute + // note: a second pass, it must run after the sync to reflect the compute const int64_t t_now = ggml_time_us(); for (int i = off; i < off + n_tokens; ++i) { const auto & t = batch.tokens[i]; diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 0c82a73f99..588ff8d9bd 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -123,7 +123,7 @@ void server_queue::terminate() { condition_tasks.notify_all(); } -bool server_queue::process_new_tasks(std::deque * unhandled) { +bool server_queue::process_new_tasks(bool is_yielding) { while (true) { std::unique_lock lock(mutex_tasks); if (!running) { @@ -138,11 +138,12 @@ bool server_queue::process_new_tasks(std::deque * unhandled) { lock.unlock(); QUE_DBG("processing task, id = %d\n", task.id); - if (!callback_new_task(std::move(task))) { + if (!callback_new_task(std::move(task), is_yielding)) { // set it aside, do not put it back in the queue, else we offer it again in a loop - GGML_ASSERT(unhandled && "a task can only be declined while yielding"); + GGML_ASSERT(is_yielding && "a task can only be declined while yielding"); QUE_DBG("task declined, id = %d\n", task.id); - unhandled->push_back(std::move(task)); + lock.lock(); + queue_tasks_unhandled.push_back(std::move(task)); } } } @@ -151,7 +152,7 @@ void server_queue::worker_loop() { while (true) { std::function work; { - std::unique_lock lock(worker.mutex); + std::unique_lock lock(mutex_tasks); worker.cv.wait(lock, [&]{ return worker.stop || worker.work != nullptr; }); @@ -163,14 +164,16 @@ void server_queue::worker_loop() { } // note: do not hold any lock here, work() may post new tasks + std::exception_ptr exception; try { work(); } catch (...) { - worker.exception = std::current_exception(); + exception = std::current_exception(); } // signal completion to yield_to_queue() std::unique_lock lock(mutex_tasks); + worker.exception = std::move(exception); worker.busy = false; condition_tasks.notify_all(); } @@ -181,10 +184,10 @@ void server_queue::worker_stop() { return; } { - std::unique_lock lock(worker.mutex); + std::unique_lock lock(mutex_tasks); worker.stop = true; - worker.cv.notify_one(); } + worker.cv.notify_one(); worker.thread.join(); } @@ -193,26 +196,20 @@ void server_queue::yield_to_queue(std::function && work) { QUE_DBG("%s", "yielding to queue\n"); - // tasks declined while the work is running - std::deque unhandled; - { std::unique_lock lock(mutex_tasks); GGML_ASSERT(!worker.busy && "yield_to_queue() cannot be nested"); worker.busy = true; - } - { - std::unique_lock lock(worker.mutex); worker.work = std::move(work); - worker.cv.notify_one(); } + 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(&unhandled); + process_new_tasks(true); std::unique_lock lock(mutex_tasks); - // declined tasks are kept in unhandled, so a non-empty queue always has something new + // 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()); }); @@ -221,25 +218,27 @@ void server_queue::yield_to_queue(std::function && work) { } } + std::exception_ptr exception; { std::unique_lock lock(mutex_tasks); // put the declined tasks back, keeping their order - while (!unhandled.empty()) { - queue_tasks.push_front(std::move(unhandled.back())); - unhandled.pop_back(); + while (!queue_tasks_unhandled.empty()) { + queue_tasks.push_front(std::move(queue_tasks_unhandled.back())); + queue_tasks_unhandled.pop_back(); } // 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); } QUE_DBG("%s", "done yielding to queue\n"); - // the worker is idle now, safe to read worker.exception - if (worker.exception) { - std::exception_ptr exception = nullptr; - std::swap(exception, worker.exception); + // note: rethrow only after the declined tasks are back in the queue, so they are not lost + if (exception) { std::rethrow_exception(exception); } } @@ -265,7 +264,7 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { while (true) { QUE_DBG("%s", "processing new tasks\n"); - if (process_new_tasks()) { + if (process_new_tasks(false)) { break; // terminate } @@ -329,11 +328,15 @@ void server_queue::cleanup_pending_task(int id_target) { return task.id == id_target; }; queue_tasks.erase( - std::remove_if(queue_tasks.begin(), queue_tasks.end(), rm_func), + std::remove_if(queue_tasks.begin(), queue_tasks.end(), rm_func), queue_tasks.end()); queue_tasks_deferred.erase( - std::remove_if(queue_tasks_deferred.begin(), queue_tasks_deferred.end(), rm_func), + std::remove_if(queue_tasks_deferred.begin(), queue_tasks_deferred.end(), rm_func), queue_tasks_deferred.end()); + // a task declined while yielding is not in queue_tasks yet, but it can still be cancelled + queue_tasks_unhandled.erase( + std::remove_if(queue_tasks_unhandled.begin(), queue_tasks_unhandled.end(), rm_func), + queue_tasks_unhandled.end()); } // diff --git a/tools/server/server-queue.h b/tools/server/server-queue.h index b207971a82..f154262067 100644 --- a/tools/server/server-queue.h +++ b/tools/server/server-queue.h @@ -22,26 +22,28 @@ private: // queues std::deque queue_tasks; std::deque queue_tasks_deferred; + // tasks declined while yielding, put back in queue_tasks once the yield is done + // note: kept as a member so that cleanup_pending_task() can also reach them + std::deque queue_tasks_unhandled; std::mutex mutex_tasks; std::condition_variable condition_tasks; - // used by yield_to_queue + // used by yield_to_queue, all fields are guarded by mutex_tasks struct worker_t { - std::thread thread; - std::mutex mutex; - std::condition_variable cv; - std::function work; // pending work, picked up by the thread - std::exception_ptr exception; // exception thrown by work(), if any + std::thread thread; + std::condition_variable cv; // the worker sleeps on this until there is work + std::function work; // pending work, picked up by the thread + std::exception_ptr exception; // exception thrown by work(), if any bool stop = false; - bool busy = false; // guarded by mutex_tasks (not by worker_t::mutex) + bool busy = false; }; worker_t worker; // callback functions - std::function callback_new_task; - std::function callback_update_slots; - std::function callback_sleeping_state; + std::function callback_new_task; + std::function callback_update_slots; + std::function callback_sleeping_state; public: ~server_queue() { worker_stop(); } @@ -94,6 +96,7 @@ public: // 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 + // ref: https://github.com/ggml-org/llama.cpp/pull/27041 // // tasks declined by callback_new_task are put back in the queue once this returns void yield_to_queue(std::function && work); @@ -109,9 +112,10 @@ public: // // Register function to process a new task - // it returns false to decline the task, which is only allowed while yielding - // on decline, the task must be left untouched, it is put back in the queue later - void on_new_task(std::function callback) { + // 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 + void on_new_task(std::function callback) { callback_new_task = std::move(callback); } @@ -140,8 +144,8 @@ private: // process all pending tasks in the queue // returns true if the queue is terminated, false if there is no more task to process - // declined tasks are moved to unhandled, which must be set when called while yielding - bool process_new_tasks(std::deque * unhandled = nullptr); + // while yielding, declined tasks are moved to queue_tasks_unhandled + bool process_new_tasks(bool is_yielding); // for worker_t void worker_loop();