This commit is contained in:
Xuan Son Nguyen
2026-08-14 00:42:29 +02:00
parent a35d2af15f
commit d809c0f3ab
3 changed files with 151 additions and 152 deletions
+102 -110
View File
@@ -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<common_speculative *>(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<common_speculative *>(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<server_prompt_cache> 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];
+30 -27
View File
@@ -123,7 +123,7 @@ void server_queue::terminate() {
condition_tasks.notify_all();
}
bool server_queue::process_new_tasks(std::deque<server_task> * unhandled) {
bool server_queue::process_new_tasks(bool is_yielding) {
while (true) {
std::unique_lock<std::mutex> lock(mutex_tasks);
if (!running) {
@@ -138,11 +138,12 @@ bool server_queue::process_new_tasks(std::deque<server_task> * 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<void()> work;
{
std::unique_lock<std::mutex> lock(worker.mutex);
std::unique_lock<std::mutex> 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<std::mutex> 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<std::mutex> lock(worker.mutex);
std::unique_lock<std::mutex> 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<void()> && work) {
QUE_DBG("%s", "yielding to queue\n");
// tasks declined while the work is running
std::deque<server_task> unhandled;
{
std::unique_lock<std::mutex> lock(mutex_tasks);
GGML_ASSERT(!worker.busy && "yield_to_queue() cannot be nested");
worker.busy = true;
}
{
std::unique_lock<std::mutex> 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<std::mutex> 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<void()> && work) {
}
}
std::exception_ptr exception;
{
std::unique_lock<std::mutex> 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());
}
//
+19 -15
View File
@@ -22,26 +22,28 @@ private:
// queues
std::deque<server_task> queue_tasks;
std::deque<server_task> 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<server_task> 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<void()> 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<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; // guarded by mutex_tasks (not by worker_t::mutex)
bool busy = false;
};
worker_t worker;
// callback functions
std::function<bool(server_task &&)> callback_new_task;
std::function<void(void)> callback_update_slots;
std::function<void(bool)> callback_sleeping_state;
std::function<bool(server_task &&, bool)> callback_new_task;
std::function<void(void)> callback_update_slots;
std::function<void(bool)> 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<void()> && 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<bool(server_task &&)> 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<bool(server_task &&, bool)> 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<server_task> * 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();