Skip to content

[Mirror] server: improve slots scheduling for n_cmpl #80

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ngxson wants to merge 13 commits into
base: master
Choose a base branch
from
+194 −103
Open

[Mirror] server: improve slots scheduling for n_cmpl #80

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
b55964a
server : make sure children tasks are scheduled to launch with parent
ngxson Jan 12, 2026
e32545f
fix
ngxson Jan 12, 2026
821e329
add comment pointing to this PR
ngxson Jan 12, 2026
25702ba
fix
ngxson Jan 12, 2026
9481b9d
clean up
ngxson Jan 12, 2026
f0349e4
more debug messages
ngxson Jan 13, 2026
da6e2ba
add pop_deferred_task with specific ID version
ngxson Jan 13, 2026
d6b0d23
improve the logic
ngxson Jan 13, 2026
ba86ad9
simple approach
ngxson Jan 14, 2026
79c1967
no double move
ngxson Jan 14, 2026
d5505b1
Merge branch 'master' into xsn/n_cmpl_sync_barrier
ngxson Jan 15, 2026
8b5474a
correct return type of launch_slots_with_parent_task
ngxson Jan 15, 2026
64b48eb
Merge branch 'master' into xsn/n_cmpl_sync_barrier
ngxson Jan 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
175 changes: 126 additions & 49 deletions tools/server/server-context.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,10 @@ struct server_slot {
std::unique_ptr<const server_task> task;
std::unique_ptr<const server_task> task_prev; // used for debugging

// if set, we only accept this task next (in other words, the task reserves this slot)
// used for scheduling n_children tasks after the parent
int task_id_next = -1;

// used to determine the slot that has been used the longest
int64_t t_last_used = -1;

Expand Down Expand Up @@ -176,6 +180,7 @@ struct server_slot {
void reset() {
SLT_DBG(*this, "%s", "\n");

task_id_next = -1;
n_prompt_tokens_cache = 0;

last_nl_pos = 0;
Expand Down Expand Up @@ -325,17 +330,6 @@ struct server_slot {
return n_draft_max;
}

// note: a slot can also be either a parent or a child
// TODO: move to server_task
bool is_parent() const {
return task->n_children > 0;
}

// TODO: move to server_task
bool is_child() const {
return task->id_parent >= 0;
}

void release() {
if (is_processing()) {
GGML_ASSERT(task);
Expand All @@ -348,7 +342,7 @@ struct server_slot {
state = SLOT_STATE_IDLE;

// do not keep context of the child slots - the parent's context is enough
if (is_child()) {
if (task->is_child()) {
clear(false);
}

Expand Down Expand Up @@ -1223,12 +1217,11 @@ struct server_context_impl {

slot.task = std::make_unique<const server_task>(std::move(task));

slot.state = slot.is_child()
slot.state = slot.task->is_child()
? SLOT_STATE_WAIT_OTHER // wait for the parent to process prompt
: SLOT_STATE_STARTED;

SLT_INF(slot, "processing task, is_child = %d\n", slot.is_child());

SLT_INF(slot, "processing task, is_child = %d\n", slot.task->is_child());
return true;
}

Expand Down Expand Up @@ -1623,9 +1616,7 @@ struct server_context_impl {

// tokenize the input if it's set by CLI, return false on error
bool tokenize_cli_input(server_task & task) {
if (task.cli_input == nullptr) {
return true; // nothing to do
}
GGML_ASSERT(task.cli_input != nullptr);
try {
auto & opt = oai_parser_opt;
common_chat_templates_inputs inputs;
Expand Down Expand Up @@ -1659,35 +1650,141 @@ struct server_context_impl {
return true;
}

// if N slots are reserved AND they are all available, return true
// if not, leave them in the reserved state and return false
// ref: https://github.com/ggml-org/llama.cpp/pull/18789
bool try_reserve_n_slots(const size_t n_required, const int task_id) {
size_t n_reserved = 0;
size_t n_available = 0;
for (auto & slot : slots) {
if (n_reserved >= n_required) {
break;
} else if (slot.task_id_next == task_id || slot.task_id_next == -1) {
// already reserved to this task OR not reserved by any other tasks
slot.task_id_next = task_id;
n_reserved++;
if (!slot.is_processing()) {
n_available++;
}
}
}
return n_available >= n_required;
}

void unreserve_slots(const int task_id) {
for (auto & slot : slots) {
if (slot.task_id_next == task_id) {
slot.task_id_next = -1;
}
}
}

// launch multiple slots for parent + child tasks
bool launch_slots_with_child_tasks(server_slot & parent_slot, server_task && parent_task) {
GGML_ASSERT(parent_task.is_parent());
SRV_INF("launching slots for parent task id_task = %d with %d child tasks\n",
parent_task.id, parent_task.n_children);

// launch all child tasks first
int i_child = 0;
for (auto & slot : slots) {
if (slot.id == parent_slot.id) {
continue;
}
if (i_child >= parent_task.n_children) {
break;
}
int id_child = parent_task.child_tasks[i_child].id;
if (!launch_slot_with_task(slot, std::move(parent_task.child_tasks[i_child]))) {
SRV_ERR("failed to launch slot with child task, id_task = %d\n", id_child);
return false;
}
i_child++;
}
parent_task.child_tasks.clear();

// finally, launch the parent task
int id_parent = parent_task.id;
if (!launch_slot_with_task(parent_slot, std::move(parent_task))) {
SRV_ERR("failed to launch slot with task, id_task = %d\n", id_parent);
return false;
}

// all slots have been successfully launched, unreserve them
unreserve_slots(id_parent);

return true;
}
Comment on lines +1635 to +1677

@coderabbitai coderabbitai Bot Jan 14, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Critical: Return values are inverted due to bool/int mismatch.

The function returns bool but uses integer return values (return 0 for success, return 2 for failure). Due to implicit conversion:

  • return 0 → false (but should indicate success)
  • return 2 → true (but should indicate failure)

At the call site (line 1755), !launch_slots_with_parent_task(...) will:

  • On success (returns 0/false): !false = true → enters error branch incorrectly
  • On failure (returns 2/true): !true = false → skips error branch incorrectly
🐛 Proposed fix 
     // launch multiple slots for parent + child tasks
     bool launch_slots_with_parent_task(server_slot & parent_slot, std::vector<server_slot *> & child_slots, server_task && parent_task) {
         GGML_ASSERT(!parent_slot.is_processing());
         GGML_ASSERT(parent_task.is_parent());
         GGML_ASSERT(child_slots.size() == parent_task.child_tasks.size());

         int id_parent = parent_task.id;

         SRV_INF("launching slots for parent task id_task = %d with %zu child tasks\n", id_parent, parent_task.child_tasks.size());

         // to be called in case of failure to release all launched slots
         auto release_slots = [this, id_parent]() {
             for (auto & slot : slots) {
                 if (slot.is_processing() && (
                         slot.task->id == id_parent ||
                         slot.task->id_parent == id_parent
                 )) {
                     slot.release();
                 }
             }
         };

         // launch all child tasks first
         size_t idx = 0;
         GGML_ASSERT(child_slots.size() == parent_task.child_tasks.size());
         for (auto * slot : child_slots) {
             int id_child = parent_task.child_tasks[idx].id;
             if (!launch_slot_with_task(*slot, std::move(parent_task.child_tasks[idx]))) {
                 SRV_ERR("failed to launch slot with child task, id_task = %d\n", id_child);
                 release_slots();
-                return 2;
+                return false;
             }
             idx++;
         }

         // finally, launch the parent task
         if (!launch_slot_with_task(parent_slot, std::move(parent_task))) {
             SRV_ERR("failed to launch slot with task, id_task = %d\n", id_parent);
             release_slots();
-            return 2;
+            return false;
         }

-        return 0;
+        return true;
     }
🤖 Prompt for AI Agents 
In `@tools/server/server-context.cpp` around lines 1662 - 1704, The function
launch_slots_with_parent_task currently returns bool but uses integer return
codes (0 for success, 2 for failure), causing inverted logic; change its
signature from bool to int and keep the existing return values (0 = success,
non-zero = failure), and update all call sites that invoke
launch_slots_with_parent_task (and any checks like
!launch_slots_with_parent_task(...)) to interpret the int correctly (e.g., check
for != 0 on failure or == 0 on success); no change needed to
launch_slot_with_task or release_slots logic besides adjusting callers to use
the int result.


void process_single_task(server_task && task) {
switch (task.type) {
case SERVER_TASK_TYPE_COMPLETION:
case SERVER_TASK_TYPE_INFILL:
case SERVER_TASK_TYPE_EMBEDDING:
case SERVER_TASK_TYPE_RERANK:
{
if (!tokenize_cli_input(task)) {
break;
// special case: if input is provided via CLI, tokenize it first
// otherwise, no need to tokenize as it's already done inside the HTTP thread
if (task.cli_input != nullptr) {
if (!tokenize_cli_input(task)) {
break;
}
}

const int id_slot = task.id_slot;

server_slot * slot = id_slot != -1 ? get_slot_by_id(id_slot) : get_available_slot(task);

//
// slot scheduling logic
//

if (slot == nullptr) {
// if no slot is available, we defer this task for processing later
SRV_DBG("no slot is available, defer task, id_task = %d\n", task.id);
queue_tasks.defer(std::move(task));
break;
}

if (slot->task_id_next != -1 && slot->task_id_next != task.id) {
// if the slot is reserved for another task, we defer this task for processing later
SRV_DBG("requested slot is reserved for another task (task_id_next = %d), defer task, id_task = %d\n", slot->task_id_next, task.id);
queue_tasks.defer(std::move(task));
break;
}

if (slot->is_processing()) {
// if requested slot is unavailable, we defer this task for processing later
SRV_DBG("requested slot is unavailable, defer task, id_task = %d\n", task.id);
queue_tasks.defer(std::move(task));
break;
}

if (task.is_parent()) {
// if this is a parent task, we want to make sure parent + all child tasks can be launched at the same time
// the current slot must be either reserved for this task, or free
GGML_ASSERT(slot->task_id_next == -1 || slot->task_id_next == task.id);
slot->task_id_next = task.id;

// need to reserve n_children more slots
if (try_reserve_n_slots(task.n_children, task.id)) {
// all required slots have been reserved, safe to proceed
int task_id = task.id;
if (!launch_slots_with_child_tasks(*slot, std::move(task))) {
SRV_ERR("failed to launch slots with child tasks, id_task = %d\n", task_id);
unreserve_slots(task_id); // task is cancelled, unreserve all slots
}
break;
} else {
// failed to reserve all required slots, we defer this task for processing later
SRV_DBG("failed to reserve %d slots, defer task, id_task = %d\n", task.n_children + 1, task.id);
// clear task_id_next for the current slot (while keeping other slots reserved)
slot->task_id_next = -1;
queue_tasks.defer(std::move(task));
break;
}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

if (!launch_slot_with_task(*slot, std::move(task))) {
SRV_ERR("failed to launch slot with task, id_task = %d\n", task.id);
break;
Expand All @@ -1702,6 +1799,7 @@ struct server_context_impl {
break;
}
}
unreserve_slots(task.id_target);
} break;
case SERVER_TASK_TYPE_NEXT_RESPONSE:
{
Expand Down Expand Up @@ -1959,7 +2057,7 @@ struct server_context_impl {
GGML_ABORT("not supported by multimodal");
}

if (slot.is_parent() || slot.is_child()) {
if (slot.task->is_parent() || slot.task->is_child()) {
send_error(slot, "context shift cannot be used for shared prompt", ERROR_TYPE_SERVER);
slot.release();
continue;
Expand Down Expand Up @@ -2106,21 +2204,6 @@ struct server_context_impl {

// this slot still has a prompt to be processed
if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_STARTED) {
// wait for all children to be launched
if (slot.is_parent()) {
int n_launched = 0;
for (auto & other : slots) {
if (other.is_processing() && other.is_child() && other.task->id_parent == slot.task->id) {
++n_launched;
}
}

if (n_launched < slot.task->n_children) {
SLT_DBG(slot, "waiting for children to be launched, n_children = %d, n_launched = %d\n", slot.task->n_children, n_launched);
continue;
}
}

const auto & input_tokens = slot.task->tokens;

// TODO: maybe move branch to outside of this loop in the future
Expand Down Expand Up @@ -2674,7 +2757,7 @@ struct server_context_impl {

// handle `n_cmpl > 1` tasks - when the main prompt is processed, activate all child tasks too
for (auto & slot : slots) {
if (slot.state == SLOT_STATE_DONE_PROMPT && slot.is_parent()) {
if (slot.state == SLOT_STATE_DONE_PROMPT && slot.task->is_parent()) {
SLT_INF(slot, "parent task prompt done, n_children = %d\n", slot.task->n_children);

std::vector<server_slot *> children;
Expand Down Expand Up @@ -2968,7 +3051,9 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
// Everything else, including multimodal completions.
inputs = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true);
}
tasks.reserve(inputs.size());

// tasks.reserve(inputs.size()); // TODO: this is inaccurate due to child tasks

for (size_t i = 0; i < inputs.size(); i++) {
server_task task = server_task(type);

Expand All @@ -2988,24 +3073,16 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
task.params.oaicompat_model = meta->model_name;

// prepare child tasks
std::vector<server_task> child_tasks;
if (task.params.n_cmpl > 1) {
task.n_children = task.params.n_cmpl - 1;

for (int j = 0; j < task.n_children; j++) {
server_task child = task.create_child(task.id, rd.get_new_id());

// use different sampling seed for each child
// note: https://github.com/ggml-org/llama.cpp/pull/18700#discussion_r2675115723
if (child.params.sampling.seed != LLAMA_DEFAULT_SEED) {
child.params.sampling.seed += j + 1;
}

tasks.push_back(std::move(child));
task.add_child(task.id, rd.get_new_id());
}
}

// note: the parent task always launches first
tasks.insert(tasks.begin(), std::move(task));
tasks.push_back(std::move(task));
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

rd.post_tasks(std::move(tasks));
Expand Down
24 changes: 16 additions & 8 deletions tools/server/server-queue.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,12 +217,12 @@ void server_response::add_waiting_task_id(int id_task) {
waiting_task_ids.insert(id_task);
}

void server_response::add_waiting_tasks(const std::vector<server_task> & tasks) {
void server_response::add_waiting_task_ids(const std::unordered_set<int> & id_tasks) {
std::unique_lock<std::mutex> lock(mutex_results);

for (const auto & task : tasks) {
RES_DBG("add task %d to waiting list. current waiting = %d (before add)\n", task.id, (int) waiting_task_ids.size());
waiting_task_ids.insert(task.id);
for (const auto & id_task : id_tasks) {
RES_DBG("add task %d to waiting list. current waiting = %d (before add)\n", id_task, (int) waiting_task_ids.size());
waiting_task_ids.insert(id_task);
}
}

Expand Down Expand Up @@ -327,6 +327,7 @@ void server_response::terminate() {

void server_response_reader::post_task(server_task && task, bool front) {
GGML_ASSERT(id_tasks.empty() && "post_task() can only be called once per reader");
GGML_ASSERT(!task.is_parent() && "not supported, use post_tasks() instead");
task.index = 0;
id_tasks.insert(task.id);
states.push_back(task.create_state());
Expand All @@ -338,11 +339,18 @@ void server_response_reader::post_tasks(std::vector<server_task> && tasks, bool
GGML_ASSERT(id_tasks.empty() && "post_tasks() can only be called once per reader");
id_tasks = server_task::get_list_id(tasks);
states.reserve(tasks.size());
for (size_t i = 0; i < tasks.size(); i++) {
tasks[i].index = i;
states.push_back(tasks[i].create_state());
size_t index = 0;
for (auto & task : tasks) {
task.index = index++;
states.push_back(task.create_state());
// for child tasks
for (auto & child_task : task.child_tasks) {
child_task.index = index++;
states.push_back(child_task.create_state());
}
}
queue_results.add_waiting_tasks(tasks);
GGML_ASSERT(states.size() == id_tasks.size());
queue_results.add_waiting_task_ids(id_tasks);
queue_tasks.post(std::move(tasks), front);
}

Expand Down
2 changes: 1 addition & 1 deletion tools/server/server-queue.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,7 @@ struct server_response {
// add the id_task to the list of tasks waiting for response
void add_waiting_task_id(int id_task);

void add_waiting_tasks(const std::vector<server_task> & tasks);
void add_waiting_task_ids(const std::unordered_set<int> & id_tasks);

// when the request is finished, we can remove task associated with it
void remove_waiting_task_id(int id_task);
Expand Down
Loading
Loading