From a4854f0349b6d30f3ee530a473371e8112683b31 Mon Sep 17 00:00:00 2001
From: Georgi Gerganov <ggerganov@gmail.com>
Date: Fri, 9 Jan 2026 15:30:39 +0200
Subject: [PATCH] cont : improve n_cmpl logic

- launch the parent task first so it finds the slot with best cache
- parent task waits for child tasks to be launched
- when a child task finishes - remove its cache
---
 tools/server/server-context.cpp | 46 +++++++++++++++++++++++++++------
 tools/server/server-task.h      |  2 ++
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index fe2b6ac5f5..51307dde8a 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -153,7 +153,7 @@ struct server_slot {
 
     common_sampler_ptr smpl;
 
-    llama_token sampled; // in speculative mode, this is the last accepted token
+    llama_token  sampled; // in speculative mode, this is the last accepted token
     llama_tokens drafted;
 
     // stats
@@ -201,6 +201,11 @@ struct server_slot {
         alora_invocation_start = -1;
     }
 
+    void clear() {
+        llama_memory_seq_rm(llama_get_memory(ctx), id, -1, -1);
+        prompt.tokens.clear();
+    }
+
     void init_sampler() const {
         const int64_t t_start = ggml_time_us();
 
@@ -321,10 +326,18 @@ struct server_slot {
 
             SLT_INF(*this, "stop processing: n_tokens = %d, truncated = %d\n", prompt.n_tokens(), truncated);
 
-            t_last_used = ggml_time_us();
+            t_last_used        =  ggml_time_us();
             t_token_generation = (ggml_time_us() - t_start_generation) / 1e3;
+
             state = SLOT_STATE_IDLE;
 
+            // do not keep context of the child slots - the parent's context is enough
+            if (is_child()) {
+                SLT_INF(*this, "clearing child slot with %zu tokens\n", prompt.tokens.size());
+
+                clear();
+            }
+
             task_prev = std::move(task);
             task.reset();
 
@@ -447,8 +460,8 @@ struct server_slot {
     void copy_state_to(server_slot & other) const {
         GGML_ASSERT(state == SLOT_STATE_DONE_PROMPT);
 
-        llama_memory_seq_rm(llama_get_memory(ctx), other.id,     0, -1);
-        llama_memory_seq_cp(llama_get_memory(ctx), id, other.id, 0, -1);
+        llama_memory_seq_rm(llama_get_memory(ctx), other.id,     -1, -1);
+        llama_memory_seq_cp(llama_get_memory(ctx), id, other.id, -1, -1);
 
         other.n_decoded   = n_decoded;
         other.n_remaining = n_remaining;
@@ -1033,15 +1046,14 @@ private:
         return ret;
     }
 
-    void clear_slot(server_slot & slot, bool allow_processing = false) const {
+    static void clear_slot(server_slot & slot, bool allow_processing = false) {
         if (!allow_processing) {
             GGML_ASSERT(!slot.is_processing());
         }
 
         SLT_WRN(slot, "clearing slot with %zu tokens\n", slot.prompt.tokens.size());
 
-        llama_memory_seq_rm(llama_get_memory(ctx), slot.id, -1, -1);
-        slot.prompt.tokens.clear();
+        slot.clear();
     }
 
     // return true if at least one slot has been cleared
@@ -2087,6 +2099,21 @@ private:
                     continue;
                 }
 
+                // 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;
+                    }
+                }
+
                 // this slot still has a prompt to be processed
                 if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_STARTED) {
                     const auto & input_tokens = slot.task->tokens;
@@ -2955,8 +2982,10 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
             task.params.oaicompat_cmpl_id = completion_id;
             task.params.oaicompat_model   = meta->model_name;
 
+            // prepare 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());
 
@@ -2970,7 +2999,8 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
                 }
             }
 
-            tasks.push_back(std::move(task));
+            // note: the parent task always launches first
+            tasks.insert(tasks.begin(), std::move(task));
         }
 
         rd.post_tasks(std::move(tasks));
diff --git a/tools/server/server-task.h b/tools/server/server-task.h
index 536a7727cc..cf08fced63 100644
--- a/tools/server/server-task.h
+++ b/tools/server/server-task.h
@@ -173,11 +173,13 @@ struct server_task {
 
     server_task create_child(int id_parent, int id_child) const {
         server_task copy;
+
         copy.id        = id_child;
         copy.id_parent = id_parent;
         copy.params    = params;
         copy.type      = type;
         copy.tokens    = tokens.clone();
+
         return copy;
     }