From fb9b1f2b007ce3fc9c2a1bc5bae5e812fd0aab74 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Mon, 9 Jun 2025 13:04:22 +0400
Subject: [PATCH 01/42] reformat llama.cpp

---
 src/llama.cpp | 61 +++++++++++++++++++++++++++------------------------
 1 file changed, 32 insertions(+), 29 deletions(-)

diff --git a/src/llama.cpp b/src/llama.cpp
index 0e615b67..c03d92a5 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -7472,10 +7472,7 @@ static void llm_load_qwen2_tensors(
     uint32_t             n_world, 
     uint32_t             my_rank, 
     const uint32_t     * n_layer_window,
-    bool               * use_mmap_buffer,
     bool                 set_needed) {
-    (void)use_mmap_buffer; // unused in this function
-
     const auto tn = LLM_TN(model.arch);
 
     ggml_context * ctx_input        = nullptr;
@@ -7590,10 +7587,10 @@ static bool llm_load_tensors_impl(
             GGML_ASSERT(local_i != -1);
 
             if (local_i % window_size >= window_size - n_gpu_layers) {
-                // LLAMA_LOG_INFO("Layer %i assigned to gpu (cache index %i)\n", i, local_i);
+                LLAMA_LOG_DEBUG("Layer %i assigned to gpu (cache index %i)\n", i, local_i);
                 model.buft_layer[local_i] = llama_default_buffer_type_offload(model, main_gpu);
             } else {
-                // LLAMA_LOG_INFO("Layer %i assigned to cpu (cache index %i)\n", i, local_i);
+                LLAMA_LOG_DEBUG("Layer %i assigned to cpu (cache index %i)\n", i, local_i);
                 model.buft_layer[local_i] = llama_default_buffer_type_cpu(model, true);
             }
         }
@@ -7603,8 +7600,8 @@ static bool llm_load_tensors_impl(
     if (my_rank == 0) {
         model.buft_input  = llama_default_buffer_type_cpu(model, true);
         model.buft_output = llama_default_buffer_type_cpu(model, true);
-        // LLAMA_LOG_INFO("Layer input assigned to cpu\n");
-        // LLAMA_LOG_INFO("Layer output assigned to cpu\n");
+        LLAMA_LOG_DEBUG("Layer input assigned to cpu\n");
+        LLAMA_LOG_DEBUG("Layer output assigned to cpu\n");
     }
 
     // count used buffer types
@@ -8212,7 +8209,7 @@ static bool llm_load_tensors_impl(
                     }
                 } break;
             case LLM_ARCH_QWEN2:
-                llm_load_qwen2_tensors(ml, model, ctx_map, n_world, my_rank, n_layer_window, &use_mmap_buffer, true);
+                llm_load_qwen2_tensors(ml, model, ctx_map, n_world, my_rank, n_layer_window, true);
                 break;
             case LLM_ARCH_QWEN2MOE:
                 {
@@ -11182,8 +11179,6 @@ struct llm_build_context {
 
             cur = llm_build_out_embd(ctx0, lctx, hparams, cb);
 
-            // cur   = ggml_get_rows(ctx0,   cur, inp_out_ids);
-
             cur = llm_build_norm(ctx0, cur, hparams,
                     model.output_norm, NULL,
                     LLM_NORM_RMS, cb, -1);
@@ -12724,18 +12719,19 @@ struct llm_build_context {
     }
 
     std::vector<ggml_cgraph *> build_qwen2() {
+        // mutable variable, needed during the last layer of the computation to skip unused tokens
+        int32_t n_tokens = this->n_tokens;
 
         const int64_t n_embd_head = hparams.n_embd_head_v;
-        GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
-        GGML_ASSERT(n_embd_head == hparams.n_rot);
+        GGML_ASSERT(n_embd_head   == hparams.n_embd_head_k);
+        GGML_ASSERT(n_embd_head   == hparams.n_rot);
 
         // create a vector to hold the subgraphs
         std::vector<struct ggml_cgraph *> sub_gfs; 
         struct ggml_cgraph * sub_gf  = nullptr;
-        struct ggml_tensor * cur = nullptr;
-        struct ggml_tensor * inpL = nullptr;
-        struct ggml_tensor * inpB = nullptr;
-
+        struct ggml_tensor * cur     = nullptr;
+        struct ggml_tensor * inpL    = nullptr;
+        struct ggml_tensor * inpB    = nullptr;
         const  uint32_t      n_world = this->cparams.n_world;
         const  uint32_t      my_rank = this->cparams.rank;
         const  uint32_t    * n_layer_window = this->cparams.n_layer_window;
@@ -12751,7 +12747,7 @@ struct llm_build_context {
             sub_gfs.push_back(sub_gf);
 
             sub_gf = nullptr;
-            inpL = nullptr;
+            inpL   = nullptr;
         }
 
         // inpB - contains the output embedding from other nodes
@@ -12763,6 +12759,7 @@ struct llm_build_context {
         // KQ_mask (mask for 1 head, it will be broadcasted to all heads)
         struct ggml_tensor * KQ_mask = build_inp_KQ_mask();
 
+        const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
         for (int il = 0; il < n_layer; ++il) {
             if (!this_layer_is_mine(il, n_world, my_rank, n_layer_window)) {
                 // if we have an active sub-graph, add it to the list 
@@ -12771,7 +12768,6 @@ struct llm_build_context {
                     sub_gfs.push_back(sub_gf);
                     sub_gf = nullptr;
                 }
-                // synchronous input tensor
                 if (inpL != inpB) {
                     inpL = inpB;
                 }
@@ -12801,21 +12797,27 @@ struct llm_build_context {
                 // compute Q and K and RoPE them
                 struct ggml_tensor * Qcur = llm_build_lora_mm(lctx, ctx0, model.layers[local_il].wq, cur);
                 cb(Qcur, "Qcur", il);
-                Qcur = ggml_add(ctx0, Qcur, model.layers[local_il].bq);
-                cb(Qcur, "Qcur", il);
+                if (model.layers[local_il].bq) {
+                    Qcur = ggml_add(ctx0, Qcur, model.layers[local_il].bq);
+                    cb(Qcur, "Qcur", il);
+                }
 
                 struct ggml_tensor * Kcur = llm_build_lora_mm(lctx, ctx0, model.layers[local_il].wk, cur);
                 cb(Kcur, "Kcur", il);
-                Kcur = ggml_add(ctx0, Kcur, model.layers[local_il].bk);
-                cb(Kcur, "Kcur", il);
+                if (model.layers[local_il].bk) {
+                    Kcur = ggml_add(ctx0, Kcur, model.layers[local_il].bk);
+                    cb(Kcur, "Kcur", il);
+                }
 
                 struct ggml_tensor * Vcur = llm_build_lora_mm(lctx, ctx0, model.layers[local_il].wv, cur);
                 cb(Vcur, "Vcur", il);
-                Vcur = ggml_add(ctx0, Vcur, model.layers[local_il].bv);
-                cb(Vcur, "Vcur", il);
+                if (model.layers[local_il].bv) {
+                    Vcur = ggml_add(ctx0, Vcur, model.layers[local_il].bv);
+                    cb(Vcur, "Vcur", il);
+                }
 
                 Qcur = ggml_rope_ext(
-                    ctx0, ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head,    n_tokens), inp_pos, nullptr,
+                    ctx0, ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens), inp_pos, nullptr,
                     n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
                     ext_factor, attn_factor, beta_fast, beta_slow
                 );
@@ -12830,7 +12832,7 @@ struct llm_build_context {
 
                 cur = llm_build_kv(ctx0, lctx, kv_self, sub_gf,
                     model.layers[local_il].wo, model.layers[local_il].bo,
-                    Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, 1.0f/sqrtf(float(n_embd_head)), cb, il);
+                    Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, kq_scale, cb, il);
             }
 
             if (il == n_layer - 1) {
@@ -12840,7 +12842,7 @@ struct llm_build_context {
                 inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
             }
 
-            struct ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+            struct ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); // shortcut
             cb(ffn_inp, "ffn_inp", il);
 
             // feed-forward network
@@ -12858,6 +12860,8 @@ struct llm_build_context {
             cb(cur, "ffn_out", il);
 
             cur = ggml_add(ctx0, cur, ffn_inp); // shortcut
+            cb(cur, "ffn_out", il);
+
             cur = lctx.cvec.apply_to(ctx0, cur, il);
             cb(cur, "l_out", il);
 
@@ -17034,7 +17038,6 @@ static std::vector<struct ggml_cgraph *> llama_build_graph(
             } break;
         case LLM_ARCH_QWEN2:
             {
-                // result.push_back(llm.build_qwen2()); 
                 result = llm.build_qwen2();
             } break;
         case LLM_ARCH_QWEN2MOE:
@@ -21896,7 +21899,7 @@ void llama_model_n_flops(
                 llm_load_llama_tensors(*ml, *model, ctx_map, 1, 0, n_layer_window, &use_mmap_buffer, false);
                 break;
             case LLM_ARCH_QWEN2:
-                llm_load_qwen2_tensors(*ml, *model, ctx_map, 1, 0, n_layer_window, &use_mmap_buffer, false);
+                llm_load_qwen2_tensors(*ml, *model, ctx_map, 1, 0, n_layer_window, false);
                 break;
             default:
                 throw std::runtime_error("unsupported architecture\n");

From 3e6d83193086b881e65b3d3b65f0cbbd9d077fb2 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Wed, 11 Jun 2025 17:10:21 +0400
Subject: [PATCH 02/42] fix seq_id mismatch between head and worker devices

---
 examples/server/server.cpp | 10 +++----
 include/llama.h            |  3 ++-
 src/llama.cpp              | 53 ++++++++++++++++++++------------------
 3 files changed, 33 insertions(+), 33 deletions(-)

diff --git a/examples/server/server.cpp b/examples/server/server.cpp
index f67ba107..af39f1ac 100644
--- a/examples/server/server.cpp
+++ b/examples/server/server.cpp
@@ -1059,13 +1059,9 @@ struct server_context {
     }
 
     void kv_cache_clear() {
-        SRV_DBG("%s", "clearing KV cache\n");
-
-        // clear the entire KV cache
+        SRV_DBG("%s", "clearing all KV cache\n");
         llama_kv_cache_clear(ctx);
-
         llama_send_kv_cache_clear(ctx);
-
         clean_kv_cache = false;
     }
 
@@ -1090,7 +1086,7 @@ struct server_context {
                     llama_batch_add(batch, system_tokens[i + j], i + j, { 0 }, false);
                 }
 
-                if (llama_decode(ctx, batch) != 0) {
+                if (llama_decode(ctx, batch, true) != 0) {
                     SRV_ERR("%s", "llama_decode() failed\n");
                     return;
                 }
@@ -2311,7 +2307,7 @@ struct server_context {
                 0, 0, 0, // unused
             };
 
-            const int ret = llama_decode(ctx, batch_view);
+            const int ret = llama_decode(ctx, batch_view, true);
             metrics.on_decoded(slots);
 
             if (ret != 0) {
diff --git a/include/llama.h b/include/llama.h
index 8bb8ac50..86da593c 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -957,7 +957,8 @@ extern "C" {
     // < 0 - error
     LLAMA_API int32_t llama_decode(
             struct llama_context * ctx,
-              struct llama_batch   batch);
+              struct llama_batch   batch,
+                            bool   server_mode = false);
 
     // Set the number of threads used for decoding
     // n_threads is the number of threads used for generation (single token)
diff --git a/src/llama.cpp b/src/llama.cpp
index c03d92a5..af42f79d 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -17849,7 +17849,7 @@ struct sync_meta {
     int          div_factor    = 1;
 };
 
-static void llama_send_meta(zmq::socket_t & socket, struct sync_meta * meta) {
+static void llama_send_meta(zmq::socket_t & socket, struct sync_meta * meta, bool align_seq_ids = false) {
     GGML_ASSERT(meta != nullptr);
     try {
         std::vector<zmq::message_t> send_msgs;
@@ -17864,19 +17864,20 @@ static void llama_send_meta(zmq::socket_t & socket, struct sync_meta * meta) {
         }
 
         if (meta->n_seq_id != nullptr) {
-            GGML_ASSERT(meta->n_ctx > 0);
+            GGML_ASSERT(meta->n_tokens > 0);
             send_msgs.emplace_back("n_seq_id", strlen("n_seq_id"));
-            send_msgs.emplace_back(meta->n_seq_id, meta->n_ctx * sizeof(int32_t));
+            send_msgs.emplace_back(meta->n_seq_id, meta->n_tokens * sizeof(int32_t));
 
             // here we assume only a single seq_id per token is needed
             // pack all single seq_id values into a contiguous array
-            llama_seq_id * all_seq_ids = (llama_seq_id *) malloc(meta->n_ctx * sizeof(llama_seq_id));
-            for (uint32_t i = 0; i < meta->n_ctx; ++i) {
-                all_seq_ids[i] = meta->seq_id[i][0];
+            llama_seq_id * all_seq_ids = (llama_seq_id *) malloc(meta->n_tokens * sizeof(llama_seq_id));
+            int seq_id_offset = align_seq_ids ? 1 : 0;
+            for (int32_t i = 0; i < meta->n_tokens; ++i) {
+                all_seq_ids[i] = meta->seq_id[i][0] - seq_id_offset;
             }
             
             send_msgs.emplace_back("seq_id", strlen("seq_id"));
-            send_msgs.emplace_back(all_seq_ids, meta->n_ctx * sizeof(llama_seq_id));
+            send_msgs.emplace_back(all_seq_ids, meta->n_tokens * sizeof(llama_seq_id));
             free(all_seq_ids);
         }
 
@@ -17966,18 +17967,18 @@ static int llama_recv_meta(zmq::socket_t & socket, struct sync_meta * meta) {
         }
 
         if (key == "n_seq_id") {
-            GGML_ASSERT(meta->n_ctx > 0);
-            GGML_ASSERT(data_msg.size() == meta->n_ctx * sizeof(int32_t));
-            meta->n_seq_id = (int32_t *) malloc(meta->n_ctx * sizeof(int32_t));
-            std::memcpy(meta->n_seq_id, data_msg.data(), meta->n_ctx * sizeof(int32_t));
+            GGML_ASSERT(meta->n_tokens > 0);
+            GGML_ASSERT(data_msg.size() == meta->n_tokens * sizeof(int32_t));
+            meta->n_seq_id = (int32_t *) malloc(meta->n_tokens * sizeof(int32_t));
+            std::memcpy(meta->n_seq_id, data_msg.data(), meta->n_tokens * sizeof(int32_t));
         }
 
         if (key == "seq_id") {
-            GGML_ASSERT(meta->n_ctx > 0);
-            GGML_ASSERT(data_msg.size() == meta->n_ctx * sizeof(llama_seq_id));
+            GGML_ASSERT(meta->n_tokens > 0);
+            GGML_ASSERT(data_msg.size() == meta->n_tokens * sizeof(llama_seq_id));
             const llama_seq_id * all_seq_ids = (llama_seq_id *) data_msg.data();
-            meta->seq_id = (llama_seq_id **) malloc(meta->n_ctx * sizeof(llama_seq_id *));
-            for (uint32_t i = 0; i < meta->n_ctx; ++i) {
+            meta->seq_id = (llama_seq_id **) malloc(meta->n_tokens * sizeof(llama_seq_id *));
+            for (int32_t i = 0; i < meta->n_tokens; ++i) {
                 meta->seq_id[i] = (llama_seq_id *) malloc(sizeof(llama_seq_id));
                 meta->seq_id[i][0] = all_seq_ids[i];
             }
@@ -18203,7 +18204,8 @@ static void manage_graph_tensors(struct ggml_cgraph * cgraph, int advice, bool f
 //
 static int llama_decode_internal(
          llama_context & lctx,
-           llama_batch   batch_all) { // TODO: rename back to batch
+           llama_batch   batch_all,
+                  bool   server_mode) {
     const auto   & model   = lctx.model;
     const auto   & hparams = model.hparams;
     const auto   & cparams = lctx.cparams;
@@ -18275,16 +18277,16 @@ static int llama_decode_internal(
         if (meta.n_tokens > 0) {
             batch_all.n_tokens = meta.n_tokens;
             if (meta.pos != nullptr) {
-                batch_all.pos = (llama_pos *) malloc(cparams.n_ctx * sizeof(llama_pos));
-                std::memcpy(batch_all.pos, meta.pos, cparams.n_ctx * sizeof(llama_pos));
+                batch_all.pos = (llama_pos *) malloc(meta.n_ctx * sizeof(llama_pos));
+                std::memcpy(batch_all.pos, meta.pos, meta.n_ctx * sizeof(llama_pos));
             }
             if (meta.n_seq_id != nullptr) {
-                batch_all.n_seq_id = (int32_t *) malloc(cparams.n_ctx * sizeof(int32_t));
-                std::memcpy(batch_all.n_seq_id, meta.n_seq_id, cparams.n_ctx * sizeof(int32_t));
+                batch_all.n_seq_id = (int32_t *) malloc(meta.n_tokens * sizeof(int32_t));
+                std::memcpy(batch_all.n_seq_id, meta.n_seq_id, meta.n_tokens * sizeof(int32_t));
             }
             if (meta.seq_id != nullptr) {
-                batch_all.seq_id = (llama_seq_id **) malloc(cparams.n_ctx * sizeof(llama_seq_id *));
-                for (size_t i = 0; i < cparams.n_ctx; ++i) {
+                batch_all.seq_id = (llama_seq_id **) malloc(meta.n_tokens * sizeof(llama_seq_id *));
+                for (int32_t i = 0; i < meta.n_tokens; ++i) {
                     batch_all.seq_id[i] = (llama_seq_id *) malloc(sizeof(llama_seq_id));
                     batch_all.seq_id[i][0] = meta.seq_id[i][0];
                 }
@@ -18346,7 +18348,7 @@ static int llama_decode_internal(
         meta.logits    = batch_all.logits;
         meta.all_pos_0 = batch_all.all_pos_0;
         meta.all_pos_1 = batch_all.all_pos_1;
-        llama_send_meta(*lctx.send_socket, &meta);
+        llama_send_meta(*lctx.send_socket, &meta, server_mode);
     } 
     
     lctx.sbatch.from_batch(batch_all, n_embd,
@@ -23484,8 +23486,9 @@ int32_t llama_encode(
 
 int32_t llama_decode(
         struct llama_context * ctx,
-          struct llama_batch   batch) {
-    return llama_decode_internal(*ctx, batch);
+          struct llama_batch   batch,
+                        bool   server_mode) {
+    return llama_decode_internal(*ctx, batch, server_mode);
 }
 
 void llama_synchronize(struct llama_context * ctx) {

From d1b97f798e52ec7c8621a1e2264996f732776b80 Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@gmail.com>
Date: Fri, 23 May 2025 10:08:30 +0000
Subject: [PATCH 03/42] support reconnection

---
 common/common.cpp   |  17 +++-
 common/profiler.cpp |  30 +++++-
 common/profiler.h   |  14 ++-
 include/llama.h     |  13 ++-
 src/llama.cpp       | 221 ++++++++++++++++++++++++++++++++------------
 5 files changed, 226 insertions(+), 69 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index dff98506..18e0804f 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1717,6 +1717,8 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
         }
 
         // sychronize device profile to the master node
+        NodeType node_type;
+        char is_fowarder[32] = {0};
         if (my_rank == 0) {
             if (auto_schedule) {
                 std::vector<device_info> dev_info_set(n_world);
@@ -1743,14 +1745,14 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
             if (auto_schedule){
                 llama_send_device_info(lctx, &dev_info);
                 llama_recv_layer_setup(lctx, n_layer_window, n_gpu_layers);
-                llama_rebuild_topo    (lctx, n_layer_window, nullptr);
+                llama_rebuild_topo    (lctx, n_layer_window, nullptr, &node_type, is_fowarder);
             } else {
                 llama_recv_layer_setup(lctx, n_layer_window, n_gpu_layers);
             }
         }
 
         // if this is a weak device, then exit
-        if (n_layer_window[my_rank] <= 0) {
+        if (node_type == NodeType::NODE_TYPE_EXIT) {
             LOG_INF("No layer is assigned to me, exit.\n");
             llama_free(lctx);
             llama_free_model(model);
@@ -1762,7 +1764,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
         std::vector<uint32_t> n_layer_window_temp = {n_layer_window[0]}, n_gpu_layers_temp = {n_gpu_layers[0]};
 
         for (uint32_t i = 1; i < n_world; i++) {
-            if (n_layer_window[i] <= 0) {
+            if (n_layer_window[i] <= 0 && is_fowarder[i] == 0) {
                 continue;
             }
             if (i <= my_rank) {
@@ -1794,7 +1796,14 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
         n_world         = update_n_world;
 
         llama_update_context_with_rankworld(lctx, update_rank, update_n_world);
-        
+                
+        if(node_type == NodeType::NODE_TYPE_EXIT){
+            //just foward
+            while (true) {
+                llama_foward_messages(lctx);
+            }
+        }
+
         // update n_layer_window and n_gpu_layers
         std::copy(std::begin(n_layer_window), std::end(n_layer_window), params.n_layer_window);
         std::copy(std::begin(n_layer_window), std::end(n_layer_window), cparams.n_layer_window);
diff --git a/common/profiler.cpp b/common/profiler.cpp
index 0da13824..fa1a56b5 100644
--- a/common/profiler.cpp
+++ b/common/profiler.cpp
@@ -2621,7 +2621,7 @@ size_t serialize(const struct device_info * dev_info, char ** buffer) {
     return total_size;
 }
 
-void deserialize(const char * buffer, struct device_info * dev_info) {
+size_t deserialize(const char * buffer, struct device_info * dev_info) {
     const char * ptr = buffer;
 
     // rank
@@ -2821,6 +2821,32 @@ void deserialize(const char * buffer, struct device_info * dev_info) {
     ptr += sizeof(float);
 
     memcpy(&dev_info->gpu_props.cuda_mem_cpy_delay, ptr, sizeof(float));
+    ptr += sizeof(float);
 
     // no need to synchronize model flops and model params
-}
\ No newline at end of file
+    return ptr - buffer;
+}
+
+void TopoRebuildHelperInfo::deserialize(const char *buffer) {
+    size_t buffer_size = ::deserialize(buffer, &dev_info);
+    if (buffer_size == 0) {
+        LOG_ERR("%s: failed to deserialize device info\n", __func__);
+        return;
+    }
+    memcpy(&is_fowarder, buffer + buffer_size, 1);
+}
+
+size_t TopoRebuildHelperInfo::serialize(char **buffer) const{ 
+    size_t buffer_size = ::serialize(&dev_info, buffer);
+    char* buffer_ = (char*)malloc(buffer_size+1);
+    if (buffer_ == NULL) {
+        LOG_ERR("%s: failed to allocate %zu bytes for device info serialization\n", 
+                __func__, buffer_size);
+        return 0;
+    }
+    memcpy(buffer_, *buffer, buffer_size);
+    memcpy(buffer_ + buffer_size, &is_fowarder, 1);
+    free(*buffer);
+    *buffer = buffer_;
+    return buffer_size + 1;
+}
diff --git a/common/profiler.h b/common/profiler.h
index 06741d6c..5ac73a8c 100644
--- a/common/profiler.h
+++ b/common/profiler.h
@@ -346,6 +346,18 @@ struct device_info {
         model_bytes() {}
 };
 
+struct TopoRebuildHelperInfo{
+    struct device_info dev_info;
+    char               is_fowarder;
+    
+    TopoRebuildHelperInfo():
+        dev_info(),
+        is_fowarder(0){}
+    
+    void   deserialize(const char * buffer);
+    size_t serialize(char ** buffer) const;
+};
+
 enum profiler_backend_type {
     PROFILER_BACKEND_TYPE_CPU   = 0,
     PROFILER_BACKEND_TYPE_METAL = 1,
@@ -389,6 +401,6 @@ int      device_has_blas   (void);
 int      device_has_sycl   (void);
 
 size_t   serialize  (const struct device_info * dev_info, char ** buffer);
-void     deserialize(const char * buffer, struct device_info * dev_info);
+size_t   deserialize(const char * buffer, struct device_info * dev_info);
 
 #endif // PROFILER_H
diff --git a/include/llama.h b/include/llama.h
index 86da593c..c2e4d43c 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -448,6 +448,12 @@ extern "C" {
               struct llama_model_params   params);
 
     LLAMA_API void llama_free_model(struct llama_model * model);
+    
+    enum NodeType{
+        NODE_TYPE_WORKER,
+        NODE_TYPE_FOWARDER,
+        NODE_TYPE_EXIT,
+    };
 
     LLAMA_API void llama_init_sockets      (struct llama_context * ctx, uint32_t n_world, uint32_t my_rank);
     LLAMA_API void llama_free_sockets      (struct llama_context * ctx, char ** msg);
@@ -455,7 +461,12 @@ extern "C" {
     LLAMA_API int  llama_send_device_info  (struct llama_context * ctx, struct device_info * dev_info);
     LLAMA_API int  llama_bcast_startup_args(struct llama_context * ctx, uint32_t rank, struct startup_args * args);
     LLAMA_API int  llama_bcast_layer_setup (struct llama_context * ctx, uint32_t * n_layer_window, uint32_t * n_gpu_layers);
-    LLAMA_API int  llama_rebuild_topo      (struct llama_context * ctx, uint32_t * n_layer_window, struct device_info * dev_info_set);
+    LLAMA_API int  llama_rebuild_topo      (struct llama_context * ctx, 
+                                                        uint32_t * n_layer_window, 
+                                              struct device_info * dev_info_set, 
+                                                         NodeType* node_type,
+                                                         char    * is_fowarder);
+    LLAMA_API int  llama_foward_messages   (struct llama_context * ctx);
     LLAMA_API int  llama_recv_layer_setup  (struct llama_context * ctx, uint32_t * n_layer_window, uint32_t * n_gpu_layers);
 
     LLAMA_API int llm_load_tensors(
diff --git a/src/llama.cpp b/src/llama.cpp
index af42f79d..dd7b0c82 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -173,12 +173,12 @@ static void zeros(std::ofstream & file, size_t n) {
 }
 
 // zmq helpers
-static std::vector<zmq::message_t> dev_infos_to_messages(const device_info* infos, 
-                                                                   uint32_t n_world){
+static std::vector<zmq::message_t> topohelper_to_messages(const TopoRebuildHelperInfo* infos, 
+                                                                uint32_t n_world){
     std::vector<zmq::message_t> res;
     for (uint32_t i = 0; i < n_world; ++i) {
         char * buffer = nullptr;
-        size_t buffer_size = serialize(&infos[i], &buffer);
+        size_t buffer_size = infos[i].serialize(&buffer);
         res.emplace_back(buffer, buffer_size);
         free(buffer);
     }
@@ -20428,6 +20428,39 @@ static uint32_t map_rank_to_port(uint32_t rank, uint32_t data_port) {
     return data_port + rank;
 }
 
+static std::string try_connect(llama_context *ctx, uint32_t rank, TopoRebuildHelperInfo* infos, uint32_t n_world, zmq::socket_t** socket){
+    auto prv_rank = (rank - 1 + n_world) % n_world;
+    std::string ip = infos[prv_rank].dev_info.next_ip;
+    std::string send_endp = "tcp://" + ip + ":" + std::to_string(map_rank_to_port(rank, ctx->data_port));
+    *socket = new zmq::socket_t(*ctx->sock_context, zmq::socket_type::push);
+    int events = 0;
+    try {
+        (*socket)->set(zmq::sockopt::linger, 0);
+        (*socket)->set(zmq::sockopt::sndtimeo, 500);
+
+        (*socket)->connect(send_endp);
+
+        std::this_thread::sleep_for(std::chrono::milliseconds(500));
+
+        size_t events_size = sizeof(events);
+        (*socket)->getsockopt(ZMQ_EVENTS, &events, &events_size);
+
+    } catch (const zmq::error_t& e) {
+        delete *socket;
+        *socket = nullptr;
+        return "";
+    }
+
+    if((events & ZMQ_POLLOUT) != 0){
+        return ip;
+    }else{
+        delete *socket;
+        *socket = nullptr;
+        return "";
+    }
+    
+}
+
 void llama_init_sockets(struct llama_context * ctx, uint32_t n_world, uint32_t my_rank) {
     if (n_world == 1) {
         return; 
@@ -20602,95 +20635,161 @@ int llama_bcast_layer_setup(struct llama_context * ctx, uint32_t * n_layer_windo
     return 0;
 }
 
-int llama_rebuild_topo(llama_context * ctx, uint32_t * n_layer_window, device_info * dev_info_set) {
+int llama_rebuild_topo(llama_context * ctx,
+                            uint32_t * n_layer_window,
+                         device_info * dev_info_set,
+                            NodeType * node_type,
+                            char     * is_fowarder) {
     uint32_t n_world = ctx->cparams.n_world;
     uint32_t my_rank = ctx->cparams.rank;
-    device_info * dev_info_ptr = nullptr;
+    TopoRebuildHelperInfo* topo_helper = new TopoRebuildHelperInfo[n_world];
 
-    if (dev_info_set == nullptr) {
+    if (dev_info_set == nullptr){
+        // for rank!=0, recv all devices info
         std::vector<zmq::message_t> msgs;
         if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) {
             return -1;
         }
-        dev_info_ptr = new device_info[n_world];
         for (size_t i = 0; i < msgs.size(); i++) {
-            deserialize((const char *)msgs[i].data(), &dev_info_ptr[i]);
+            topo_helper[i].deserialize((char *)msgs[i].data());
         }
         GGML_ASSERT(msgs.size() == n_world);
     } else {
-        dev_info_ptr = dev_info_set;
+        for (size_t i = 0; i < n_world; i++) {
+            topo_helper[i].dev_info = dev_info_set[i];
+            topo_helper[i].is_fowarder =  0;
+        }
     }
 
     GGML_ASSERT(ctx != nullptr && ctx->send_socket != nullptr);
 
-    // notify next rank
     auto next_rank = (my_rank + 1) % n_world;
-    if (n_layer_window[next_rank] <= 0 && next_rank != 0) {
+    auto next_connect_rank = (my_rank + 1) % n_world;
+    zmq::socket_t* socket_to_close = nullptr;
+    bool is_not_exit = n_layer_window[my_rank] > 0 || topo_helper[my_rank].is_fowarder == 1;
+    if (is_not_exit){
+        // reconstruct socket to the next valid rank
+        auto current_rank = my_rank;
+        std::vector<uint32_t> nodes;
+        auto next_rank_ = next_rank;
+        while (next_rank_ != my_rank) {
+            nodes.push_back(next_rank_);
+            if (n_layer_window[next_rank_] > 0) {
+                break;
+            }
+            next_rank_   = (next_rank_ + 1) % n_world;
+            current_rank = (current_rank + 1) % n_world;
+        }
+        if (next_rank_ == my_rank) {
+            // only one node
+            ctx->next_node_ip = "";
+            socket_to_close = ctx->send_socket;
+            ctx->send_socket = nullptr;
+        } else {
+            // iterate node reverse
+            zmq::socket_t* socket = nullptr;
+            std::string ip;
+            for (int i = nodes.size() - 1; i > 0; --i) {
+                auto rank = nodes[i];
+                ip = try_connect(ctx, rank, topo_helper, n_world, &socket);
+                if(!ip.empty()){
+                    topo_helper[rank].is_fowarder = 1;
+                    next_connect_rank = rank;
+                    break;
+                }
+            }
+            if(next_connect_rank != next_rank){
+                // reset socket
+                GGML_ASSERT(socket != nullptr);
+                GGML_ASSERT(!ip.empty());
+                socket_to_close = ctx->send_socket;
+                ctx->send_socket = socket;
+                ctx->next_node_ip = ip;
+                ctx->cparams.original_next_rank = next_connect_rank;
+            }
+        }
+    }else if(n_layer_window[next_rank] <= 0 && topo_helper[my_rank].is_fowarder == 0){
+        socket_to_close = ctx->send_socket;
+    }
+
+    // notify next exiting node
+    if (socket_to_close != nullptr) {
+        GGML_ASSERT(n_layer_window[next_rank] <= 0 && topo_helper[next_rank].is_fowarder == 0);
         try {
-            auto msgs = dev_infos_to_messages(dev_info_ptr, n_world);
-            ctx->send_socket->set(zmq::sockopt::linger, 3500);
-            zmq::send_multipart(*ctx->send_socket, msgs);
+            auto msgs = topohelper_to_messages(topo_helper, n_world);
+            socket_to_close->set(zmq::sockopt::linger, 3500);
+            zmq::send_multipart(*socket_to_close, msgs);
         } catch (const zmq::error_t& e) {
             LLAMA_LOG_INFO("Failed to send data: %s\n", e.what());
-            if(!dev_info_set){
-                delete[] dev_info_ptr;
-            }
             return -1;
         }
     }
-
-    zmq::socket_t * socket_to_close = nullptr;
-    if (n_layer_window[my_rank] > 0) {
-        // reconstruct socket to the next valid rank
-        std::string next_ip;
-        auto current_rank = my_rank;
-
-        while (next_rank != my_rank) {
-            if (n_layer_window[next_rank] > 0) {
-                next_ip = dev_info_ptr[current_rank].next_ip;
-                break;
-            }
-            next_rank    = (next_rank    + 1) % n_world;
-            current_rank = (current_rank + 1) % n_world;
-        }
-
-        if (!next_ip.empty()) {
-            if ((my_rank + 1) % n_world != next_rank) {   
-                socket_to_close = ctx->send_socket;  
-                ctx->send_socket = new zmq::socket_t(*ctx->sock_context, zmq::socket_type::push);
-                std::string send_endp = "tcp://" + next_ip + ":" + std::to_string(map_rank_to_port(next_rank, ctx->data_port));
-                ctx->send_socket->connect(send_endp);
-                ctx->next_node_ip = next_ip;
-                ctx->cparams.original_next_rank = next_rank;
-            }
-
-            if (next_rank != 0) {
-                try {
-                    auto msgs = dev_infos_to_messages(dev_info_ptr, n_world);
-                    zmq::send_multipart(*ctx->send_socket, msgs);
-                } catch (const zmq::error_t &e) {
-                    LLAMA_LOG_INFO("Error binding/connecting recv socket to endpoint: %s", e.what());
-                    if(!dev_info_set){
-                        delete[] dev_info_ptr;
-                    }
-                    return -1;
-                }
-            }
-        } else {
-            // only one node
-            ctx->next_node_ip = "";
+    
+    // notify next connect node
+    if(!ctx->next_node_ip.empty() && is_not_exit){
+        GGML_ASSERT(ctx->send_socket != nullptr);
+        try {
+            auto msgs = topohelper_to_messages(topo_helper, n_world);
+            zmq::send_multipart(*ctx->send_socket, msgs);
+        } catch (const zmq::error_t& e) {
+            LLAMA_LOG_INFO("Failed to send data: %s\n", e.what());
+            return -1;
         }
     }
-
-    if (!dev_info_set) {
-        delete[] dev_info_ptr;
+    
+    if(n_layer_window[my_rank] > 0){
+        *node_type = NodeType::NODE_TYPE_WORKER;
+    }else if (topo_helper[my_rank].is_fowarder == 1){
+        *node_type = NodeType::NODE_TYPE_FOWARDER;
+    }else{
+        *node_type = NodeType::NODE_TYPE_EXIT;
     }
+    
+    if(ctx->send_socket != nullptr && *node_type!=NodeType::NODE_TYPE_EXIT){
+        // recv the whole view of all nodes
+        std::vector<zmq::message_t> msgs;
+        if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) {
+            return -1;
+        }
+        GGML_ASSERT(msgs.size() == n_world);
+        for (size_t i = 0; i < msgs.size(); i++) {
+            topo_helper[i].deserialize((char *)msgs[i].data());
+        }
+        // broadcast the whole view
+        if(next_connect_rank!=0){
+            try {
+                zmq::send_multipart(*ctx->send_socket, msgs);
+            } catch (const zmq::error_t& e) {
+                LLAMA_LOG_INFO("Failed to send data: %s\n", e.what());
+                return -1;
+            }
+        }
+    }
+    for(size_t i = 0; i < n_world; i++) {
+        is_fowarder[i] = topo_helper[i].is_fowarder;
+    }
+
 
     if(socket_to_close != nullptr){
         socket_to_close->close();
         delete socket_to_close;
     }
+    delete [] topo_helper;
+    return 0;
+}
 
+LLAMA_API int llama_foward_messages(llama_context *ctx) {
+    zmq::message_t message;
+    bool more = true;
+    
+    while (more) {
+        ctx->recv_socket->recv(message, zmq::recv_flags::none);
+        size_t more_size = sizeof(more);
+        ctx->recv_socket->getsockopt(ZMQ_RCVMORE, &more, &more_size);
+        
+        ctx->send_socket->send(message, 
+            more ? zmq::send_flags::sndmore : zmq::send_flags::none);
+    }
     return 0;
 }
 

From d6c8d322cdc011cec86521356f9f9f37329bee38 Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@foxmail.com>
Date: Tue, 3 Jun 2025 15:02:59 +0800
Subject: [PATCH 04/42] fix try_connect

---
 Makefile              |  9 ++++++-
 common/common.cpp     | 12 ++++-----
 common/profiler.cpp   |  4 +--
 common/profiler.h     |  4 +--
 include/llama.h       |  8 +++---
 src/llama.cpp         | 58 +++++++++++++++++--------------------------
 src/network-utils.cpp | 26 +++++++++++++++++++
 src/network-utils.h   |  7 ++++++
 8 files changed, 78 insertions(+), 50 deletions(-)
 create mode 100644 src/network-utils.cpp
 create mode 100644 src/network-utils.h

diff --git a/Makefile b/Makefile
index 06d91984..0d794eb8 100644
--- a/Makefile
+++ b/Makefile
@@ -950,7 +950,8 @@ OBJ_LLAMA = \
 	src/llama-grammar.o \
 	src/llama-sampling.o \
 	src/unicode.o \
-	src/unicode-data.o
+	src/unicode-data.o \
+	src/network-utils.o \
 
 OBJ_COMMON = \
 	common/profiler.o \
@@ -1139,6 +1140,11 @@ src/unicode-data.o: \
 	src/unicode-data.cpp \
 	src/unicode-data.h
 	$(CXX) $(CXXFLAGS) -c $< -o $@
+	
+src/network-utils.o: \
+	src/network-utils.cpp \
+	src/network-utils.h
+	$(CXX) $(CXXFLAGS) -c $< -o $@
 
 src/llama.o: \
 	src/llama.cpp \
@@ -1147,6 +1153,7 @@ src/llama.o: \
 	src/llama-grammar.h \
 	src/llama-sampling.h \
 	src/unicode.h \
+	src/network-utils.h \
 	include/llama.h \
 	ggml/include/ggml-cuda.h \
 	ggml/include/ggml-metal.h \
diff --git a/common/common.cpp b/common/common.cpp
index 18e0804f..d40647ec 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1718,7 +1718,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
 
         // sychronize device profile to the master node
         NodeType node_type;
-        char is_fowarder[32] = {0};
+        char is_forwarder[32] = {0};
         if (my_rank == 0) {
             if (auto_schedule) {
                 std::vector<device_info> dev_info_set(n_world);
@@ -1735,7 +1735,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
                     return iparams;
                 }
                 llama_bcast_layer_setup(lctx, n_layer_window, n_gpu_layers);
-                llama_rebuild_topo(lctx, n_layer_window, dev_info_set.data());
+                llama_rebuild_topo(lctx, n_layer_window, dev_info_set.data(), &node_type, is_forwarder);
             } else {
                 // use the user-defined n_layer_window
                 std::copy(std::begin(params.n_layer_window), std::end(params.n_layer_window), n_layer_window);
@@ -1745,7 +1745,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
             if (auto_schedule){
                 llama_send_device_info(lctx, &dev_info);
                 llama_recv_layer_setup(lctx, n_layer_window, n_gpu_layers);
-                llama_rebuild_topo    (lctx, n_layer_window, nullptr, &node_type, is_fowarder);
+                llama_rebuild_topo    (lctx, n_layer_window, nullptr, &node_type, is_forwarder);
             } else {
                 llama_recv_layer_setup(lctx, n_layer_window, n_gpu_layers);
             }
@@ -1764,7 +1764,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
         std::vector<uint32_t> n_layer_window_temp = {n_layer_window[0]}, n_gpu_layers_temp = {n_gpu_layers[0]};
 
         for (uint32_t i = 1; i < n_world; i++) {
-            if (n_layer_window[i] <= 0 && is_fowarder[i] == 0) {
+            if (n_layer_window[i] <= 0 && is_forwarder[i] == 0) {
                 continue;
             }
             if (i <= my_rank) {
@@ -1797,10 +1797,10 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
 
         llama_update_context_with_rankworld(lctx, update_rank, update_n_world);
                 
-        if(node_type == NodeType::NODE_TYPE_EXIT){
+        if(node_type == NodeType::NODE_TYPE_FORWARDER){
             //just foward
             while (true) {
-                llama_foward_messages(lctx);
+                llama_forward_messages(lctx);
             }
         }
 
diff --git a/common/profiler.cpp b/common/profiler.cpp
index fa1a56b5..18fe795d 100644
--- a/common/profiler.cpp
+++ b/common/profiler.cpp
@@ -2833,7 +2833,7 @@ void TopoRebuildHelperInfo::deserialize(const char *buffer) {
         LOG_ERR("%s: failed to deserialize device info\n", __func__);
         return;
     }
-    memcpy(&is_fowarder, buffer + buffer_size, 1);
+    memcpy(&is_forwarder, buffer + buffer_size, 1);
 }
 
 size_t TopoRebuildHelperInfo::serialize(char **buffer) const{ 
@@ -2845,7 +2845,7 @@ size_t TopoRebuildHelperInfo::serialize(char **buffer) const{
         return 0;
     }
     memcpy(buffer_, *buffer, buffer_size);
-    memcpy(buffer_ + buffer_size, &is_fowarder, 1);
+    memcpy(buffer_ + buffer_size, &is_forwarder, 1);
     free(*buffer);
     *buffer = buffer_;
     return buffer_size + 1;
diff --git a/common/profiler.h b/common/profiler.h
index 5ac73a8c..ff69a454 100644
--- a/common/profiler.h
+++ b/common/profiler.h
@@ -348,11 +348,11 @@ struct device_info {
 
 struct TopoRebuildHelperInfo{
     struct device_info dev_info;
-    char               is_fowarder;
+    char               is_forwarder;
     
     TopoRebuildHelperInfo():
         dev_info(),
-        is_fowarder(0){}
+        is_forwarder(0){}
     
     void   deserialize(const char * buffer);
     size_t serialize(char ** buffer) const;
diff --git a/include/llama.h b/include/llama.h
index c2e4d43c..21f77288 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -451,7 +451,7 @@ extern "C" {
     
     enum NodeType{
         NODE_TYPE_WORKER,
-        NODE_TYPE_FOWARDER,
+        NODE_TYPE_FORWARDER,
         NODE_TYPE_EXIT,
     };
 
@@ -463,10 +463,10 @@ extern "C" {
     LLAMA_API int  llama_bcast_layer_setup (struct llama_context * ctx, uint32_t * n_layer_window, uint32_t * n_gpu_layers);
     LLAMA_API int  llama_rebuild_topo      (struct llama_context * ctx, 
                                                         uint32_t * n_layer_window, 
-                                              struct device_info * dev_info_set, 
+                                              struct device_info * desv_info_set, 
                                                          NodeType* node_type,
-                                                         char    * is_fowarder);
-    LLAMA_API int  llama_foward_messages   (struct llama_context * ctx);
+                                                         char    * is_forwarder);
+    LLAMA_API int  llama_forward_messages   (struct llama_context * ctx);
     LLAMA_API int  llama_recv_layer_setup  (struct llama_context * ctx, uint32_t * n_layer_window, uint32_t * n_gpu_layers);
 
     LLAMA_API int llm_load_tensors(
diff --git a/src/llama.cpp b/src/llama.cpp
index dd7b0c82..281c7360 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -11,6 +11,7 @@
 #include "ggml-backend.h"
 
 #include "profiler.h"
+#include "network-utils.h"
 
 #ifdef GGML_USE_RPC
 #  include "ggml-rpc.h"
@@ -20431,34 +20432,22 @@ static uint32_t map_rank_to_port(uint32_t rank, uint32_t data_port) {
 static std::string try_connect(llama_context *ctx, uint32_t rank, TopoRebuildHelperInfo* infos, uint32_t n_world, zmq::socket_t** socket){
     auto prv_rank = (rank - 1 + n_world) % n_world;
     std::string ip = infos[prv_rank].dev_info.next_ip;
-    std::string send_endp = "tcp://" + ip + ":" + std::to_string(map_rank_to_port(rank, ctx->data_port));
+    auto port = map_rank_to_port(rank, ctx->data_port);
+    
+    if(!isPortOpen(ip, port)){
+        *socket = nullptr;
+        return "";
+    }
+    std::string send_endp = "tcp://" + ip + ":" + std::to_string(port);
     *socket = new zmq::socket_t(*ctx->sock_context, zmq::socket_type::push);
-    int events = 0;
     try {
-        (*socket)->set(zmq::sockopt::linger, 0);
-        (*socket)->set(zmq::sockopt::sndtimeo, 500);
-
         (*socket)->connect(send_endp);
-
-        std::this_thread::sleep_for(std::chrono::milliseconds(500));
-
-        size_t events_size = sizeof(events);
-        (*socket)->getsockopt(ZMQ_EVENTS, &events, &events_size);
-
     } catch (const zmq::error_t& e) {
         delete *socket;
         *socket = nullptr;
         return "";
     }
-
-    if((events & ZMQ_POLLOUT) != 0){
-        return ip;
-    }else{
-        delete *socket;
-        *socket = nullptr;
-        return "";
-    }
-    
+    return ip;
 }
 
 void llama_init_sockets(struct llama_context * ctx, uint32_t n_world, uint32_t my_rank) {
@@ -20639,7 +20628,7 @@ int llama_rebuild_topo(llama_context * ctx,
                             uint32_t * n_layer_window,
                          device_info * dev_info_set,
                             NodeType * node_type,
-                            char     * is_fowarder) {
+                            char     * is_forwarder) {
     uint32_t n_world = ctx->cparams.n_world;
     uint32_t my_rank = ctx->cparams.rank;
     TopoRebuildHelperInfo* topo_helper = new TopoRebuildHelperInfo[n_world];
@@ -20657,7 +20646,7 @@ int llama_rebuild_topo(llama_context * ctx,
     } else {
         for (size_t i = 0; i < n_world; i++) {
             topo_helper[i].dev_info = dev_info_set[i];
-            topo_helper[i].is_fowarder =  0;
+            topo_helper[i].is_forwarder =  0;
         }
     }
 
@@ -20666,7 +20655,7 @@ int llama_rebuild_topo(llama_context * ctx,
     auto next_rank = (my_rank + 1) % n_world;
     auto next_connect_rank = (my_rank + 1) % n_world;
     zmq::socket_t* socket_to_close = nullptr;
-    bool is_not_exit = n_layer_window[my_rank] > 0 || topo_helper[my_rank].is_fowarder == 1;
+    bool is_not_exit = n_layer_window[my_rank] > 0 || topo_helper[my_rank].is_forwarder == 1;
     if (is_not_exit){
         // reconstruct socket to the next valid rank
         auto current_rank = my_rank;
@@ -20692,13 +20681,13 @@ int llama_rebuild_topo(llama_context * ctx,
             for (int i = nodes.size() - 1; i > 0; --i) {
                 auto rank = nodes[i];
                 ip = try_connect(ctx, rank, topo_helper, n_world, &socket);
-                if(!ip.empty()){
-                    topo_helper[rank].is_fowarder = 1;
+                if (!ip.empty()) {
                     next_connect_rank = rank;
                     break;
                 }
             }
-            if(next_connect_rank != next_rank){
+            topo_helper[next_connect_rank].is_forwarder = 1;
+            if (next_connect_rank != next_rank) {
                 // reset socket
                 GGML_ASSERT(socket != nullptr);
                 GGML_ASSERT(!ip.empty());
@@ -20708,13 +20697,13 @@ int llama_rebuild_topo(llama_context * ctx,
                 ctx->cparams.original_next_rank = next_connect_rank;
             }
         }
-    }else if(n_layer_window[next_rank] <= 0 && topo_helper[my_rank].is_fowarder == 0){
+    }else if (n_layer_window[next_rank] <= 0 && topo_helper[next_rank].is_forwarder == 0) {
         socket_to_close = ctx->send_socket;
     }
 
     // notify next exiting node
     if (socket_to_close != nullptr) {
-        GGML_ASSERT(n_layer_window[next_rank] <= 0 && topo_helper[next_rank].is_fowarder == 0);
+        GGML_ASSERT(n_layer_window[next_rank] <= 0 && topo_helper[next_rank].is_forwarder == 0);
         try {
             auto msgs = topohelper_to_messages(topo_helper, n_world);
             socket_to_close->set(zmq::sockopt::linger, 3500);
@@ -20739,8 +20728,8 @@ int llama_rebuild_topo(llama_context * ctx,
     
     if(n_layer_window[my_rank] > 0){
         *node_type = NodeType::NODE_TYPE_WORKER;
-    }else if (topo_helper[my_rank].is_fowarder == 1){
-        *node_type = NodeType::NODE_TYPE_FOWARDER;
+    }else if (topo_helper[my_rank].is_forwarder == 1){
+        *node_type = NodeType::NODE_TYPE_FORWARDER;
     }else{
         *node_type = NodeType::NODE_TYPE_EXIT;
     }
@@ -20766,11 +20755,10 @@ int llama_rebuild_topo(llama_context * ctx,
         }
     }
     for(size_t i = 0; i < n_world; i++) {
-        is_fowarder[i] = topo_helper[i].is_fowarder;
+        is_forwarder[i] = topo_helper[i].is_forwarder;
     }
 
-
-    if(socket_to_close != nullptr){
+    if (socket_to_close != nullptr) {
         socket_to_close->close();
         delete socket_to_close;
     }
@@ -20778,9 +20766,9 @@ int llama_rebuild_topo(llama_context * ctx,
     return 0;
 }
 
-LLAMA_API int llama_foward_messages(llama_context *ctx) {
+int llama_forward_messages(llama_context *ctx) {
     zmq::message_t message;
-    bool more = true;
+    int more = true;
     
     while (more) {
         ctx->recv_socket->recv(message, zmq::recv_flags::none);
diff --git a/src/network-utils.cpp b/src/network-utils.cpp
new file mode 100644
index 00000000..b960153e
--- /dev/null
+++ b/src/network-utils.cpp
@@ -0,0 +1,26 @@
+#include "network-utils.h"
+
+#include <sys/socket.h>
+#include <netinet/in.h>
+#include <arpa/inet.h>
+#include <unistd.h>
+
+bool isPortOpen(const std::string& ip, uint32_t port, int timeout_sec) {
+    int sock = socket(AF_INET, SOCK_STREAM, 0);
+    if (sock < 0) return false;
+    
+    struct timeval tv;
+    tv.tv_sec = timeout_sec;
+    tv.tv_usec = 0;
+    setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
+    setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));
+    
+    struct sockaddr_in server;
+    server.sin_addr.s_addr = inet_addr(ip.c_str());
+    server.sin_family = AF_INET;
+    server.sin_port = htons(port);
+    
+    int res = connect(sock, (struct sockaddr*)&server, sizeof(server));
+    close(sock);
+    return res == 0;
+}
\ No newline at end of file
diff --git a/src/network-utils.h b/src/network-utils.h
new file mode 100644
index 00000000..8f0921e4
--- /dev/null
+++ b/src/network-utils.h
@@ -0,0 +1,7 @@
+#pragma once
+
+#include <string>
+
+typedef unsigned int uint32_t;
+
+bool isPortOpen(const std::string& ip, uint32_t port, int timeout_sec = 2);
\ No newline at end of file

From 2039e3b0c1bc9eaa957256b85352e01daf7aa15f Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@foxmail.com>
Date: Wed, 11 Jun 2025 21:05:31 +0800
Subject: [PATCH 05/42] fix: send and recv meta

---
 common/common.cpp | 10 +++++++++-
 include/llama.h   |  4 +++-
 src/llama.cpp     | 17 +++++++++++++++--
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index d40647ec..a21146b7 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1761,6 +1761,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
 
         // update my rank and n_world
         uint32_t update_rank = 0, update_n_world = 1;
+        uint32_t worker_rank = 0, n_worker       = 1;
         std::vector<uint32_t> n_layer_window_temp = {n_layer_window[0]}, n_gpu_layers_temp = {n_gpu_layers[0]};
 
         for (uint32_t i = 1; i < n_world; i++) {
@@ -1773,6 +1774,13 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
             update_n_world++;
             n_layer_window_temp.push_back(n_layer_window[i]);
             n_gpu_layers_temp.push_back(n_gpu_layers[i]);
+
+            if (n_layer_window[i] > 0) {
+                if (i <= my_rank) {
+                    worker_rank++;
+                }
+                n_worker++;
+            }
         }
 
         memset(n_layer_window, 0, n_world * sizeof(uint32_t));
@@ -1795,7 +1803,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
         params.n_world  = update_n_world;
         n_world         = update_n_world;
 
-        llama_update_context_with_rankworld(lctx, update_rank, update_n_world);
+        llama_update_context_with_rankworld(lctx, update_rank, update_n_world, worker_rank, n_worker);
                 
         if(node_type == NodeType::NODE_TYPE_FORWARDER){
             //just foward
diff --git a/include/llama.h b/include/llama.h
index 21f77288..7b63e96f 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -477,7 +477,9 @@ extern "C" {
     LLAMA_API void llama_update_context_with_rankworld(
                    struct llama_context * ctx, 
                                uint32_t   rank, 
-                               uint32_t   n_world);
+                               uint32_t   n_world,
+                               uint32_t   worker_rank,
+                               uint32_t   n_worker);
     
     LLAMA_API struct llama_context * llama_new_context_with_model(
                      struct llama_model * model,
diff --git a/src/llama.cpp b/src/llama.cpp
index 281c7360..dd0540cf 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -2597,6 +2597,9 @@ static_assert(std::is_trivially_copyable<llama_hparams>::value, "llama_hparams m
 struct llama_cparams {
     uint32_t  n_world;
     uint32_t  rank;
+    NodeType  node_type;
+    uint32_t  n_worker;
+    uint32_t  worker_rank;
     uint32_t  original_next_rank; // original rank of the next node
     uint32_t  n_layer_window[32];
     bool      prefetch;
@@ -18213,6 +18216,9 @@ static int llama_decode_internal(
     const uint32_t n_world = cparams.n_world;
     const uint32_t my_rank = cparams.rank;
 
+    const uint32_t n_worker    = cparams.n_worker;
+    const uint32_t worker_rank = cparams.worker_rank;
+
     lctx.is_encoding = false;
     const uint32_t n_tokens_all = batch_all.n_tokens;
     if (my_rank != 0) {
@@ -18268,7 +18274,7 @@ static int llama_decode_internal(
 
     sync_meta meta;
     meta.n_ctx = cparams.n_ctx;
-    bool is_last_dev = (my_rank == n_world - 1);
+    bool is_last_dev = (worker_rank == n_worker - 1);
 
     if (my_rank != 0) {
         if (llama_recv_meta(*lctx.recv_socket, &meta) == -1) {
@@ -20757,6 +20763,7 @@ int llama_rebuild_topo(llama_context * ctx,
     for(size_t i = 0; i < n_world; i++) {
         is_forwarder[i] = topo_helper[i].is_forwarder;
     }
+    ctx->cparams.node_type = *node_type;
 
     if (socket_to_close != nullptr) {
         socket_to_close->close();
@@ -20842,10 +20849,16 @@ void llama_free_sockets(struct llama_context * ctx, char ** msg) {
     }
 }
 
-void llama_update_context_with_rankworld(struct llama_context * ctx, uint32_t rank, uint32_t n_world) {
+void llama_update_context_with_rankworld(struct llama_context * ctx, 
+                                                       uint32_t rank, 
+                                                       uint32_t n_world, 
+                                                       uint32_t worker_rank, 
+                                                       uint32_t n_worker) {
     if (ctx) {
         ctx->cparams.rank    = rank;
         ctx->cparams.n_world = n_world;
+        ctx->cparams.worker_rank = worker_rank;
+        ctx->cparams.n_worker = n_worker;
     }
 }
 

From d4618de9918140f9b07a98683a8ff4820879d8af Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@foxmail.com>
Date: Wed, 11 Jun 2025 21:47:58 +0800
Subject: [PATCH 06/42] fix: block when free socket

---
 common/common.cpp | 19 +++++++++++++++----
 src/llama.cpp     |  8 ++++++--
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index a21146b7..699a2dd7 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -28,6 +28,7 @@
 #include <unordered_set>
 #include <vector>
 #include <thread>
+#include <atomic>
 
 #if defined(__APPLE__) && defined(__MACH__)
 #include <sys/types.h>
@@ -1806,10 +1807,20 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
         llama_update_context_with_rankworld(lctx, update_rank, update_n_world, worker_rank, n_worker);
                 
         if(node_type == NodeType::NODE_TYPE_FORWARDER){
-            //just foward
-            while (true) {
-                llama_forward_messages(lctx);
-            }
+            //just forward
+            std::atomic<bool> should_exit{false};
+            auto t = std::thread([lctx, &should_exit]() {
+                while(!should_exit) {
+                    llama_forward_messages(lctx);
+                }
+            });
+            char * stop_signal = nullptr;
+            llama_free_sockets(lctx, &stop_signal); // this will block until receive stop signal
+
+            should_exit = true;
+            t.join();
+
+            exit(0);
         }
 
         // update n_layer_window and n_gpu_layers
diff --git a/src/llama.cpp b/src/llama.cpp
index dd0540cf..4880d029 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -20776,9 +20776,13 @@ int llama_rebuild_topo(llama_context * ctx,
 int llama_forward_messages(llama_context *ctx) {
     zmq::message_t message;
     int more = true;
-    
+    int timeout_ms = 10;
+    ctx->recv_socket->setsockopt(ZMQ_RCVTIMEO, &timeout_ms, sizeof(timeout_ms));
     while (more) {
-        ctx->recv_socket->recv(message, zmq::recv_flags::none);
+        auto recv_result = ctx->recv_socket->recv(message, zmq::recv_flags::none);
+        if (!recv_result) {
+            return -1;
+        }
         size_t more_size = sizeof(more);
         ctx->recv_socket->getsockopt(ZMQ_RCVMORE, &more, &more_size);
         

From dc875bbef9d832e30cc910f44f89e85bd0aeb84a Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Fri, 13 Jun 2025 08:18:12 +0400
Subject: [PATCH 07/42] fix speculative decoding

---
 Makefile                             |  2 +
 examples/speculative/speculative.cpp | 62 ++++++++++++++++------------
 include/llama.h                      |  5 +++
 src/llama.cpp                        | 34 ++++++++++++++-
 4 files changed, 75 insertions(+), 28 deletions(-)

diff --git a/Makefile b/Makefile
index 06d91984..8d9f7410 100644
--- a/Makefile
+++ b/Makefile
@@ -2,6 +2,8 @@
 BUILD_TARGETS = \
 	llama-server \
 	llama-cli \
+	llama-speculative \
+	llama-gguf-split \
 	profile-tool
 
 # BUILD_TARGETS = \
diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp
index adf6255e..3716579e 100644
--- a/examples/speculative/speculative.cpp
+++ b/examples/speculative/speculative.cpp
@@ -12,7 +12,7 @@
 #include <string>
 #include <vector>
 
-#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE  100
+#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE  128
 #define SPEC_VOCAB_CHECK_START_TOKEN_ID 5
 
 struct seq_draft {
@@ -65,23 +65,29 @@ int main(int argc, char ** argv) {
     llama_context * ctx_tgt = NULL;
     llama_context * ctx_dft = NULL;
 
+    // load the draft model
+    // make a hard copy of params to use for the draft model
+    gpt_params params_draft   = params;
+    params_draft.model        = params_draft.model_draft;
+    params_draft.n_gpu_layers = params_draft.n_gpu_layers_draft;
+    params_draft.n_world      = 1;  // do not split the draft model across devices
+    params_draft.rank         = 0;  // always load the draft model on the head device
+    std::fill_n(params_draft.n_layer_window, params.n_world, 0);
+
+    if (params_draft.draft_cpuparams.n_threads > 0) {
+        params_draft.cpuparams.n_threads = params_draft.draft_cpuparams.n_threads;
+    }
+
+    params_draft.cpuparams_batch.n_threads = params_draft.draft_cpuparams_batch.n_threads;
+    llama_init_result llama_init_dft = llama_init_from_gpt_params(params_draft);
+    model_dft = llama_init_dft.model;
+    ctx_dft = llama_init_dft.context;
+
     // load the target model
     llama_init_result llama_init_tgt = llama_init_from_gpt_params(params);
     model_tgt = llama_init_tgt.model;
     ctx_tgt = llama_init_tgt.context;
 
-    // load the draft model
-    params.model = params.model_draft;
-    params.n_gpu_layers = params.n_gpu_layers_draft;
-    if (params.draft_cpuparams.n_threads > 0) {
-        params.cpuparams.n_threads = params.draft_cpuparams.n_threads;
-    }
-
-    params.cpuparams_batch.n_threads = params.draft_cpuparams_batch.n_threads;
-    llama_init_result llama_init_dft = llama_init_from_gpt_params(params);
-    model_dft = llama_init_dft.model;
-    ctx_dft = llama_init_dft.context;
-
     const bool vocab_type_tgt = llama_vocab_type(model_tgt);
     LOG_DBG("vocab_type tgt: %d\n", vocab_type_tgt);
 
@@ -161,9 +167,6 @@ int main(int argc, char ** argv) {
 
     const auto t_enc_end = ggml_time_us();
 
-    // the 2 models should have the same vocab
-    //GGML_ASSERT(n_vocab == llama_n_vocab(model_dft));
-
     // how many tokens to draft each time
     int n_draft = params.n_draft;
 
@@ -180,8 +183,6 @@ int main(int argc, char ** argv) {
     // target model sampling context (reuse the llama_context's sampling instance)
     struct gpt_sampler * smpl = gpt_sampler_init(model_tgt, params.sparams);
 
-    struct llama_sampler * softmax = llama_sampler_init_softmax();
-
     // draft sequence data
     std::vector<seq_draft> drafts(n_seq_dft);
 
@@ -258,10 +259,13 @@ int main(int argc, char ** argv) {
                         float r = u_dist(rng);
                         llama_token_data_array dist_dft = { drafts[s].dists[i_dft].data() , drafts[s].dists[i_dft].size(), LLAMA_TOKEN_NULL, true };
 
-                        //GGML_ASSERT(dist_tgt.size <= dist_dft.size);
+                        // if (dist_tgt.size > dist_dft.size) {
+                        //     LOG_ERR("dist_tgt.size (%zu) must be less than or equal to dist_dft.size (%zu)\n", dist_tgt.size, dist_dft.size);
+                        //     GGML_ASSERT(dist_tgt.size <= dist_dft.size);
+                        // }
 
                         // acquire the token probabilities assigned by the draft and target models
-                        for (size_t i = 0; i < dist_tgt.size; i++) {
+                        for (size_t i = 0; i < dist_tgt.size && i < dist_dft.size; i++) {
                             if (dist_tgt.data[i].id == drafts[s].tokens[i_dft]) {
                                 p_tgt = dist_tgt.data[i].p;
                             }
@@ -406,7 +410,6 @@ int main(int argc, char ** argv) {
         {
             LOG_DBG("the sampled target token (%d, '%s') did not match, or we ran out of drafted tokens\n", token_id, token_str.c_str());
 
-            // TODO: simplify
             {
                 LOG_DBG("keeping sequence %d, n_past_tgt = %d, n_past_dft = %d\n", s_keep, n_past_tgt, n_past_dft);
 
@@ -418,6 +421,12 @@ int main(int argc, char ** argv) {
                 llama_kv_cache_seq_keep(ctx_tgt, s_keep);
                 llama_kv_cache_seq_cp  (ctx_tgt, s_keep, 0, -1, -1);
                 llama_kv_cache_seq_keep(ctx_tgt, 0);
+
+                // notify other devices to manage the KV cache in the same way
+                llama_send_kv_cache_seq_rm  (ctx_tgt, s_keep, n_past_tgt, -1);
+                llama_send_kv_cache_seq_keep(ctx_tgt, s_keep);
+                llama_send_kv_cache_seq_cp  (ctx_tgt, s_keep, 0, -1, -1);
+                llama_send_kv_cache_seq_keep(ctx_tgt, 0);
             }
 
             for (int s = 0; s < n_seq_dft; ++s) {
@@ -435,7 +444,6 @@ int main(int argc, char ** argv) {
             llama_batch_add  (batch_dft, token_id, n_past_dft, { 0 }, true);
 
             llama_kv_cache_seq_rm(ctx_dft, 0, n_past_dft, -1);
-            // LOG_DBG("dft batch: %s\n", LOG_BATCH_TOSTR_PRETTY(ctx_dft, batch_dft).c_str());
             llama_decode(ctx_dft, batch_dft);
 
             ++n_past_dft;
@@ -575,12 +583,13 @@ int main(int argc, char ** argv) {
 
         // evaluate the target model on the drafted tokens
         {
-            llama_kv_cache_seq_keep(ctx_tgt, 0);
+            llama_kv_cache_seq_keep     (ctx_tgt, 0);
+            llama_send_kv_cache_seq_keep(ctx_tgt, 0);
             for (int s = 1; s < n_seq_dft; ++s) {
-                llama_kv_cache_seq_cp(ctx_tgt, 0, s, -1, -1);
+                llama_kv_cache_seq_cp     (ctx_tgt, 0, s, -1, -1);
+                llama_send_kv_cache_seq_cp(ctx_tgt, 0, s, -1, -1);
             }
 
-            // LOG_DBG("target batch: %s\n", LOG_BATCH_TOSTR_PRETTY(ctx_tgt, batch_tgt).c_str());
             llama_decode(ctx_tgt, batch_tgt);
             ++n_past_tgt;
         }
@@ -612,7 +621,7 @@ int main(int argc, char ** argv) {
 
     LOG_INF("\n");
     LOG_INF("draft:\n\n");
-    // TODO: print sampling/grammar timings for all drafts
+    
     llama_perf_context_print(ctx_dft);
 
     LOG_INF("\n");
@@ -624,7 +633,6 @@ int main(int argc, char ** argv) {
         gpt_sampler_free(drafts[s].smpl);
     }
 
-    llama_sampler_free(softmax);
     llama_batch_free(batch_dft);
 
     llama_free(ctx_tgt);
diff --git a/include/llama.h b/include/llama.h
index 86da593c..4c39b063 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -759,6 +759,11 @@ extern "C" {
     LLAMA_API void llama_kv_cache_seq_keep(
             struct llama_context * ctx,
                     llama_seq_id   seq_id);
+    
+    // Notify other nodes to keep only the specified sequence in their KV cache
+    LLAMA_API void llama_send_kv_cache_seq_keep(
+            struct llama_context * ctx,
+                    llama_seq_id   seq_id);
 
     // Adds relative position "delta" to all tokens that belong to the specified sequence and have positions in [p0, p1)
     // If the KV cache is RoPEd, the KV data is updated accordingly:
diff --git a/src/llama.cpp b/src/llama.cpp
index af42f79d..7cd74983 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -17841,6 +17841,9 @@ struct sync_meta {
     llama_pos    cp_p0         = 0;
     llama_pos    cp_p1         = 0;
 
+    bool kv_seq_keep           = false;
+    llama_seq_id keep_seq_id   = 0;
+
     // signal to divide the kv cache range
     bool kv_seq_div            = false;
     llama_seq_id div_seq_id    = 0;
@@ -17943,8 +17946,14 @@ static int llama_recv_meta(zmq::socket_t & socket, struct sync_meta * meta) {
         return 0;
     }
 
+    if (cmd == "kv_seq_keep" && recv_msgs.size() == 2) {
+        meta->kv_seq_keep = true;
+        std::memcpy(&meta->keep_seq_id,   recv_msgs[idx++].data(), sizeof(meta->keep_seq_id));
+        return 0;
+    }
+
     if (cmd == "kv_seq_div" && recv_msgs.size() == 5) {
-        meta->kv_seq_div   = true;
+        meta->kv_seq_div = true;
         std::memcpy(&meta->div_seq_id,    recv_msgs[idx++].data(), sizeof(meta->div_seq_id));
         std::memcpy(&meta->div_p0,        recv_msgs[idx++].data(), sizeof(meta->div_p0));
         std::memcpy(&meta->div_p1,        recv_msgs[idx++].data(), sizeof(meta->div_p1));
@@ -18331,6 +18340,14 @@ static int llama_decode_internal(
             return -1;
         }
 
+        if (kv_cache_op(meta.kv_seq_keep,
+                    [&]{ llama_kv_cache_seq_keep     (&lctx, meta.keep_seq_id); },
+                    [&]{ llama_send_kv_cache_seq_keep(&lctx, meta.keep_seq_id); },
+                    is_last_dev)) {
+            LLAMA_LOG_DEBUG("%s: received signal kv_cache_seq_keep\n", __func__);
+            return -1;
+        }
+
         if (kv_cache_op(meta.kv_seq_div,
                     [&]{ llama_kv_cache_seq_div     (&lctx, meta.div_seq_id, meta.div_p0, meta.div_p1, meta.div_factor); },
                     [&]{ llama_send_kv_cache_seq_div(&lctx, meta.div_seq_id, meta.div_p0, meta.div_p1, meta.div_factor); },
@@ -22349,6 +22366,21 @@ void llama_kv_cache_seq_keep(struct llama_context * ctx, llama_seq_id seq_id) {
     llama_kv_cache_seq_keep(ctx->kv_self, seq_id);
 }
 
+void llama_send_kv_cache_seq_keep(struct llama_context * ctx, llama_seq_id seq_id) {
+    if (ctx->send_socket == nullptr) {
+        return;
+    }
+
+    try {
+        std::vector<zmq::message_t> msgs;
+        msgs.emplace_back("kv_seq_keep", strlen("kv_seq_keep"));
+        msgs.emplace_back(&seq_id, sizeof(seq_id));
+        zmq::send_multipart(*ctx->send_socket, msgs);
+    } catch (const zmq::error_t & e) {
+        LLAMA_LOG_WARN("Failed to send kv_seq_keep: %s\n", e.what());
+    }
+}
+
 void llama_kv_cache_seq_add(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) {
     if (delta == 0) {
         return;

From 2687ef3126fcb120b014b9c80e98238ebc049781 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Fri, 13 Jun 2025 11:25:42 +0400
Subject: [PATCH 08/42] speculative: free sockets and send stop signal when
 inference ends

---
 examples/speculative/speculative.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp
index 3716579e..f4ab8262 100644
--- a/examples/speculative/speculative.cpp
+++ b/examples/speculative/speculative.cpp
@@ -628,6 +628,8 @@ int main(int argc, char ** argv) {
     LOG_INF("target:\n\n");
     gpt_perf_print(ctx_tgt, smpl);
 
+    llama_free_sockets(ctx_tgt, nullptr);
+
     gpt_sampler_free(smpl);
     for (int s = 0; s < n_seq_dft; ++s) {
         gpt_sampler_free(drafts[s].smpl);

From c9cae626cf214eab07c9f18b358b46e2bf8bf66a Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Fri, 13 Jun 2025 13:30:29 +0400
Subject: [PATCH 09/42] speculative: free sockets and send stop signal when
 inference ends

---
 examples/speculative/speculative.cpp | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp
index 3716579e..8eb2fa1c 100644
--- a/examples/speculative/speculative.cpp
+++ b/examples/speculative/speculative.cpp
@@ -628,6 +628,9 @@ int main(int argc, char ** argv) {
     LOG_INF("target:\n\n");
     gpt_perf_print(ctx_tgt, smpl);
 
+    char * stop_signal = nullptr;
+    llama_free_sockets(ctx_tgt, &stop_signal);
+
     gpt_sampler_free(smpl);
     for (int s = 0; s < n_seq_dft; ++s) {
         gpt_sampler_free(drafts[s].smpl);

From b5ccd62135b18ccd16aa759d2660e4e42c0aed96 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Sat, 14 Jun 2025 18:55:53 +0400
Subject: [PATCH 10/42] fix n_gpu_layers allocation errors

---
 common/common.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index dff98506..573e8107 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1421,8 +1421,9 @@ static bool assign_layers_to_device(
                 if (n_m < static_cast<uint32_t>(std::floor(W * vec_z_gpu[m]))) {
                     // if there is still free GPU memory
                     has_free_gpu_memory = true;
-                } else if (w_m > n_m) {
-                    // if the GPU is overloaded
+                }
+                if (w_m > n_m) {
+                    // if layers are offloaded to CPU
                     has_gpu_overload = true;
                 }
             } else if (!in_set(m, M4)) {

From dfb1feb54e6159009d5cda04f0a4e86dc66ed866 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Mon, 16 Jun 2025 12:09:07 +0400
Subject: [PATCH 11/42] update README

---
 README.md                            | 34 +++++++++++++++++-----------
 common/common.cpp                    |  4 ++++
 examples/speculative/speculative.cpp |  5 ++--
 3 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/README.md b/README.md
index d336ce88..c69194f0 100644
--- a/README.md
+++ b/README.md
@@ -34,26 +34,26 @@ And, if your devices are more powerful, you could unlock even more possibilities
 
 > Device D4 runs inside a Termux-simulated Linux. Device D1 reads disk data in random mode and D2~D4 read in sequential mode.
 
-**Table 2:** Token latency for Llama models (w/o device selection).
+**Table 2:** Token latency for Llama models (with device selection).
 | **Model**      | **llama.cpp** | **exo**   | **dllama** | **prima.cpp** |
-|-----------------|---------------|-----------|------------|---------------|
-| Llama 3-8B     | **15 ms**     | 263 ms    | 459 ms     | 54 ms         |
-| Llama 3-14B    | **20 ms**     | -         | -          | 65 ms         |
+|----------------|---------------|-----------|------------|---------------|
+| Llama 3-8B     | 15 ms         | 263 ms    | 459 ms     | **15 ms**     |
+| Llama 3-14B    | 20 ms         | -         | -          | **20 ms**     |
 | Llama 1-30B    | 202 ms        | -         | -          | **72 ms**     |
 | Llama 3-45B    | 328 ms        | -         | -          | **233 ms**    |
 | Llama 3-60B    | 7965 ms       | -         | -          | **468 ms**    |
 | Llama 1-65B    | 8807 ms       | -         | -          | **569 ms**    |
 | Llama 3-70B    | 10120 ms      | OOM       | OOM        | **674 ms**    |
 
-**Table 3:** Token latency for Qwen 2.5, QwQ, and DeepSeek R1 models (w/o device selection).
+**Table 3:** Token latency for Qwen 2.5, QwQ, and DeepSeek R1 models (with device selection).
 
 | **Model**                        | **llama.cpp** | **exo**       | **dllama** | **prima.cpp** |
 |-----------------------------------|---------------|---------------|------------|---------------|
-| Qwen-2.5-7B                      | **14 ms**     | 86 ms         | -          | 44 ms         |
-| DeepSeek-R1-Distill-Qwen-7B      | **14 ms**     | 68 ms       | -          | 52 ms         |
-| DeepSeek-R1-Distill-Llama-8B     | **14 ms**     | 77 ms       | 435 ms     | 59 ms         |
-| Qwen-2.5-14B                     | **23 ms**     | 31710 ms  | -          | 65 ms         |
-| DeepSeek-R1-Distill-Qwen-14B     | **24 ms**     | 23475 ms  | -          | 76 ms         |
+| Qwen-2.5-7B                      | 14 ms     | 86 ms         | -          | **14 ms**         |
+| DeepSeek-R1-Distill-Qwen-7B      | 14 ms     | 68 ms       | -          | **14 ms**         |
+| DeepSeek-R1-Distill-Llama-8B     | 14 ms     | 77 ms       | 435 ms     | **14 ms**         |
+| Qwen-2.5-14B                     | 23 ms     | 31710 ms  | -          | **23 ms**         |
+| DeepSeek-R1-Distill-Qwen-14B     | 24 ms     | 23475 ms  | -          | **24 ms**         |
 | Qwen-2.5-32B and QwQ-32B         | 224 ms        | OOM           | -          | **89 ms**     |
 | DeepSeek-R1-Distill-Qwen-32B     | 232 ms        | OOM           | -          | **93 ms**     |
 | DeepSeek-R1-Distill-Llama-70B    | 10978 ms      | OOM           | -          | **724 ms**    |
@@ -61,9 +61,9 @@ And, if your devices are more powerful, you could unlock even more possibilities
 
 > As video recording consumes some RAM, prima.cpp proactively reduces memory usage, resulting in slightly higher latency in the video compared to the table.
 
-> In the old version (w/o device selection), each device is assigned at least one model layer. This would lead to a 1:1:29:1 split for Llama 3-8B, which makes prima.cpp slower than llama.cpp.
+> ~~In the old version (w/o device selection), each device is assigned at least one model layer. This would lead to a 1:1:29:1 split for Llama 3-8B, which makes prima.cpp slower than llama.cpp.~~
 > 
-> **New:** In the latest version (with device selection), we will have a 0:0:32:0 split and weak devices removed, then prima.cpp would become llama.cpp when serving small models.
+> In the current version (with device selection), we will have a 32:0:0:0 split and weak devices removed, then prima.cpp would become llama.cpp when serving small models.
 
 ## 🔑 Key Features
 
@@ -72,8 +72,10 @@ And, if your devices are more powerful, you could unlock even more possibilities
 - - **GPU & CPU Offloading:** If a device has a GPU, you can use both GPU and CPU for inference. For example, when VRAM is full, we can offload some model layers to RAM.
 - - **Piped-ring parallelism with prefetching:** Prefetch upcoming layer weights to overlap disk loading latency and use advanced piped-ring parallelism to prevent the "prefetch-release" effect. This new parallelism improves pipeline parallelism by using a ring structure and allows devices to run multiple cycles to predict a new token.
 - - **Heterogeneity-aware workload distribution:** A scheduler is designed to optimize workload distribution based on each device's computing power, disk speed, memory, and OS (the OS will affect the disk speed and the memory management strategy). It decides how many model layers a device should handle and how many should run on GPU (if available). 
-- - **Automatic device selection:** If there are weak devices and removing them would speed up inference, prima.cpp will automatically discover and remove them.
+- - **Automatic device selection:** If there are weak devices and removing them would speed up inference, prima.cpp will automatically discover and remove them. This may retain some devices as proxy to prevent the socket connection from being blocked.
 - - **Quantization:** We now support Q4K, Q6K, Q80 and IQ1 quantization (GGUF format) and are exploring a Q4K-IQ1 hybrid for a better balance between performance and speed.
+- - **Speculative decoding:** We now support speculative decoding, which can [further speed up by up to 80%.](https://github.com/Lizonghang/prima.cpp/discussions/29)
+- **Dynamic batching**: We now support concurrent requests from multiple users and batch decoding. 
 - **Support Models:** We now support hot models like the **Llama, Qwen (and QwQ), and DeepSeek series**. More will be added in future updates.
 - **Cross-Platform:** The cluster can consist of devices with different OSs, including macOS, Linux, Android, HarmonyOS, etc. Now, Android and HarmonyOS devices require Termux, and Windows support will be added in future update.
 
@@ -120,6 +122,7 @@ Before using this project, ensure you have the following dependencies installed:
 **Linux (e.g., Ubuntu):**
 
 ```shell
+# Use apt in Linux and pkg in Termux
 sudo apt update -y && sudo apt install -y gcc-9 make cmake fio git wget libzmq3-dev
 ```
 
@@ -279,6 +282,8 @@ You can run prima.cpp in server mode, by launching `llama-server` on the rank 0
 ./llama-cli -m download/qwq-32b-q4_k_m.gguf --world 2 --rank 1 --master 192.168.1.2 --next 192.168.1.2 --prefetch
 ```
 
+You can specify `-np 4 --cont-batching` when launching `llama-server` to enable concurrent requests.
+
 After that, you can interact with the rank 0 device by calling the Chat Completion API:
 
 ```shell
@@ -374,6 +379,9 @@ curl -X POST http://localhost:8080/v1/cancel \
      -d '{"task_id": 0}'
 ```
 
+**9. How to use speculative decoding?**
+Please see "[Power prima.cpp with speculative decoding: Further speeds up by up to 80%](https://github.com/Lizonghang/prima.cpp/discussions/29)".
+
 ## ❤️ Acknowledgment
 This project builds upon the incredible work from the open-source community, especially [ggml, gguf](https://github.com/ggml-org/ggml), and [llama.cpp](https://github.com/ggml-org/llama.cpp). We gratefully acknowledge their contributions.
 
diff --git a/common/common.cpp b/common/common.cpp
index 573e8107..79a60bb2 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1247,6 +1247,10 @@ static bool assign_layers_to_device(
                     return cost * k;
                 }
             );
+            // apply higher priority to the head device, here 0.99 is a heuristic value
+            // to ensure that small models in homogeneous clusters result in 32:0 partitioning,
+            // rather than 1:31.
+            model.lp_.col_cost_[0] *= 0.99;
 
             // define the variable bounds
             model.lp_.col_lower_ = std::vector<double>(n_world * 2, 0.0);
diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp
index 8eb2fa1c..8527f8b6 100644
--- a/examples/speculative/speculative.cpp
+++ b/examples/speculative/speculative.cpp
@@ -70,8 +70,9 @@ int main(int argc, char ** argv) {
     gpt_params params_draft   = params;
     params_draft.model        = params_draft.model_draft;
     params_draft.n_gpu_layers = params_draft.n_gpu_layers_draft;
-    params_draft.n_world      = 1;  // do not split the draft model across devices
-    params_draft.rank         = 0;  // always load the draft model on the head device
+    params_draft.n_world      = 1;    // do not split the draft model across devices
+    params_draft.rank         = 0;    // always load the draft model on the head device
+    params_draft.use_mlock    = true; // always use mlock for the draft model
     std::fill_n(params_draft.n_layer_window, params.n_world, 0);
 
     if (params_draft.draft_cpuparams.n_threads > 0) {

From 104e3b2356b88ceec2a533a8e0607c033f67687b Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@foxmail.com>
Date: Tue, 17 Jun 2025 11:27:58 +0800
Subject: [PATCH 12/42] fix: replace localhost to 127.0.0.1

---
 src/llama.cpp         | 30 +++++++++++++++---------------
 src/network-utils.cpp |  2 +-
 src/network-utils.h   |  2 +-
 3 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/src/llama.cpp b/src/llama.cpp
index 915278cc..c6f6d3b5 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -3438,8 +3438,8 @@ struct llama_context {
     struct ggml_tensor * inp_KQ_mask_cross; // F32 [n_outputs_enc, n_batch]
 
     // sockets
-    std::string      master_ip     = "localhost";
-    std::string      next_node_ip  = "localhost";
+    std::string      master_ip     = "127.0.0.1";
+    std::string      next_node_ip  = "127.0.0.1";
     uint32_t         data_port     = 9000;
     uint32_t         signal_port   = 10000;
     zmq::context_t * sock_context  = nullptr;
@@ -20452,12 +20452,12 @@ static uint32_t map_rank_to_port(uint32_t rank, uint32_t data_port) {
     return data_port + rank;
 }
 
-static std::string try_connect(llama_context *ctx, uint32_t rank, TopoRebuildHelperInfo* infos, uint32_t n_world, zmq::socket_t** socket){
-    auto prv_rank = (rank - 1 + n_world) % n_world;
-    std::string ip = infos[prv_rank].dev_info.next_ip;
+static std::string try_connect(llama_context * ctx, uint32_t rank, TopoRebuildHelperInfo * infos, uint32_t n_world, zmq::socket_t ** socket){
+    auto prev_rank = (rank - 1 + n_world) % n_world;
+    std::string ip = infos[prev_rank].dev_info.next_ip;
     auto port = map_rank_to_port(rank, ctx->data_port);
     
-    if(!isPortOpen(ip, port)){
+    if (!is_port_open(ip, port)) {
         *socket = nullptr;
         return "";
     }
@@ -20679,7 +20679,7 @@ int llama_rebuild_topo(llama_context * ctx,
     auto next_connect_rank = (my_rank + 1) % n_world;
     zmq::socket_t* socket_to_close = nullptr;
     bool is_not_exit = n_layer_window[my_rank] > 0 || topo_helper[my_rank].is_forwarder == 1;
-    if (is_not_exit){
+    if (is_not_exit) {
         // reconstruct socket to the next valid rank
         auto current_rank = my_rank;
         std::vector<uint32_t> nodes;
@@ -20738,7 +20738,7 @@ int llama_rebuild_topo(llama_context * ctx,
     }
     
     // notify next connect node
-    if(!ctx->next_node_ip.empty() && is_not_exit){
+    if (!ctx->next_node_ip.empty() && is_not_exit) {
         GGML_ASSERT(ctx->send_socket != nullptr);
         try {
             auto msgs = topohelper_to_messages(topo_helper, n_world);
@@ -20749,15 +20749,15 @@ int llama_rebuild_topo(llama_context * ctx,
         }
     }
     
-    if(n_layer_window[my_rank] > 0){
+    if (n_layer_window[my_rank] > 0) {
         *node_type = NodeType::NODE_TYPE_WORKER;
-    }else if (topo_helper[my_rank].is_forwarder == 1){
+    } else if (topo_helper[my_rank].is_forwarder == 1) {
         *node_type = NodeType::NODE_TYPE_FORWARDER;
-    }else{
+    } else {
         *node_type = NodeType::NODE_TYPE_EXIT;
     }
-    
-    if(ctx->send_socket != nullptr && *node_type!=NodeType::NODE_TYPE_EXIT){
+
+    if (ctx->send_socket != nullptr && *node_type != NodeType::NODE_TYPE_EXIT) {
         // recv the whole view of all nodes
         std::vector<zmq::message_t> msgs;
         if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) {
@@ -20768,7 +20768,7 @@ int llama_rebuild_topo(llama_context * ctx,
             topo_helper[i].deserialize((char *)msgs[i].data());
         }
         // broadcast the whole view
-        if(next_connect_rank!=0){
+        if (next_connect_rank!=0) {
             try {
                 zmq::send_multipart(*ctx->send_socket, msgs);
             } catch (const zmq::error_t& e) {
@@ -20777,7 +20777,7 @@ int llama_rebuild_topo(llama_context * ctx,
             }
         }
     }
-    for(size_t i = 0; i < n_world; i++) {
+    for (size_t i = 0; i < n_world; i++) {
         is_forwarder[i] = topo_helper[i].is_forwarder;
     }
     ctx->cparams.node_type = *node_type;
diff --git a/src/network-utils.cpp b/src/network-utils.cpp
index b960153e..e7fa5ab1 100644
--- a/src/network-utils.cpp
+++ b/src/network-utils.cpp
@@ -5,7 +5,7 @@
 #include <arpa/inet.h>
 #include <unistd.h>
 
-bool isPortOpen(const std::string& ip, uint32_t port, int timeout_sec) {
+bool is_port_open(const std::string& ip, uint32_t port, int timeout_sec) {
     int sock = socket(AF_INET, SOCK_STREAM, 0);
     if (sock < 0) return false;
     
diff --git a/src/network-utils.h b/src/network-utils.h
index 8f0921e4..7a35475a 100644
--- a/src/network-utils.h
+++ b/src/network-utils.h
@@ -4,4 +4,4 @@
 
 typedef unsigned int uint32_t;
 
-bool isPortOpen(const std::string& ip, uint32_t port, int timeout_sec = 2);
\ No newline at end of file
+bool is_port_open(const std::string& ip, uint32_t port, int timeout_sec = 2);
\ No newline at end of file

From 6ff38b2a0c1aba7c6c15960a0a6d0bdeed183a9a Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@foxmail.com>
Date: Tue, 17 Jun 2025 12:00:04 +0800
Subject: [PATCH 13/42] add args: data-port and signal-port

---
 common/arg.cpp    | 14 ++++++++++++++
 common/common.cpp |  2 ++
 common/common.h   |  6 ++++--
 include/llama.h   |  2 ++
 src/llama.cpp     |  4 ++++
 5 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/common/arg.cpp b/common/arg.cpp
index 47d3c5e6..e282c80d 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -675,6 +675,20 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
             params.rank = value;
         }
     ).set_env("LLAMA_ARG_RANK"));
+    add_opt(llama_arg(
+        {"--data-port"}, "N",
+        format("data port for distributed inference (default: %d)", params.data_port),
+        [](gpt_params & params, int value) {
+            params.data_port = value;
+        }
+    ).set_env("LLAMA_ARG_DATA_PORT"));
+    add_opt(llama_arg(
+        {"--signal-port"}, "N",
+        format("signal port for distributed inference (default: %d)", params.signal_port),
+        [](gpt_params & params, int value) {
+            params.signal_port = value;
+        }
+    ).set_env("LLAMA_ARG_SIGNAL_PORT"));
     add_opt(llama_arg(
         {"-lw", "--layer-window", "--n-layer-window"}, "N",
         format("number of layers to process in each compute (e.g., 16,16)"),
diff --git a/common/common.cpp b/common/common.cpp
index 374ae1f8..38828be3 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -2032,6 +2032,8 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param
     }
     cparams.master_ip         = new char[params.master_ip.length() + 1];
     std::strcpy(cparams.master_ip, params.master_ip.c_str());
+    cparams.data_port         = params.data_port;
+    cparams.signal_port       = params.signal_port;
 
     if (cparams.next_node_ip != nullptr) {
         delete[] cparams.next_node_ip;
diff --git a/common/common.h b/common/common.h
index 0a679213..c6ffe136 100644
--- a/common/common.h
+++ b/common/common.h
@@ -145,8 +145,10 @@ struct gpt_params {
     int32_t n_world               =     1; // number of devices to use
     int32_t rank                  =     0; // my rank for distributed inference
     uint32_t n_layer_window[32]   =   {0}; // layer window size on each node
-    std::string master_ip         = "localhost"; // ip address of the master node
-    std::string next_node_ip      = "localhost"; // ip address of my next node
+    std::string master_ip         = "127.0.0.1"; // ip address of the master node
+    std::string next_node_ip      = "127.0.0.1"; // ip address of my next node
+    uint32_t data_port            =  9000;  // data port for distributed inference
+    uint32_t signal_port          =  10000; // signal port for distributed inference
     bool    prefetch              = false; // prefetch layer weights
     bool    keep_out_in_metal     =  true; // whether to keep output weights in metal memory, true by default
     bool    force                 = false; // force to start prefetching after computation
diff --git a/include/llama.h b/include/llama.h
index 4d6dd80d..3c220562 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -330,6 +330,8 @@ extern "C" {
         bool        keep_out_in_metal; // whether to keep output weights in metal memory
         char *      master_ip;         // ip address of the master node
         char *      next_node_ip;      // ip address of the next node
+        uint32_t    data_port;         // data port for distributed inference
+        uint32_t    signal_port;       // signal port for distributed inference
         uint32_t    n_ctx;             // text context, 0 = from model
         uint32_t    n_predict;         // number of tokens to predict
         uint32_t    n_batch;           // logical maximum batch size that can be submitted to llama_decode
diff --git a/src/llama.cpp b/src/llama.cpp
index c6f6d3b5..8b5af567 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -20266,6 +20266,8 @@ struct llama_context_params llama_context_default_params() {
         /*.keep_out_in_metal           =*/ true,
         /*.master_ip                   =*/ nullptr,
         /*.next_node_ip                =*/ nullptr,
+        /*.data_port                   =*/ 9000,
+        /*.signal_port                 =*/ 10000,
         /*.n_ctx                       =*/ 512,
         /*.n_predict                   =*/ 512,
         /*.n_batch                     =*/ 2048,
@@ -20896,6 +20898,8 @@ struct llama_context * llama_new_context_with_model(
 
     ctx->master_ip       = params.master_ip;
     ctx->next_node_ip    = params.next_node_ip;
+    ctx->data_port       = params.data_port;
+    ctx->signal_port     = params.signal_port;
     ctx->cparams.n_world = params.n_world;
     ctx->cparams.rank    = params.rank;
     ctx->cparams.force   = params.force;

From 67c4f703574fb738192e016c49f8858103352b66 Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@foxmail.com>
Date: Tue, 17 Jun 2025 12:08:53 +0800
Subject: [PATCH 14/42] fix: add log when serving as a proxy

---
 common/common.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 38828be3..d1d40a9e 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1810,9 +1810,10 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
         n_world         = update_n_world;
 
         llama_update_context_with_rankworld(lctx, update_rank, update_n_world, worker_rank, n_worker);
-                
-        if(node_type == NodeType::NODE_TYPE_FORWARDER){
+
+        if (node_type == NodeType::NODE_TYPE_FORWARDER) {
             //just forward
+            LOG_INF("No layer is assigned to me, and I serve as a network proxy.\n");
             std::atomic<bool> should_exit{false};
             auto t = std::thread([lctx, &should_exit]() {
                 while(!should_exit) {

From deeec668b8657c57e7bcb94e4942470f5e9e1fd5 Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@foxmail.com>
Date: Tue, 17 Jun 2025 13:20:06 +0800
Subject: [PATCH 15/42] fix: n_worker in draft model

(cherry picked from commit 921ad2b453b24b715ad5db6a703fb3df65fdcb80)
---
 common/common.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/common/common.cpp b/common/common.cpp
index d1d40a9e..39b95d32 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1682,6 +1682,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
         cparams.n_layer_window[0] = n_layers;
         mparams.n_layer_window[0] = n_layers;
         llama_context_n_layer_window(lctx)[0] = n_layers;
+        llama_update_context_with_rankworld(lctx, 0, 1, 0, 1);
 
 #if defined(GGML_USE_METAL) || defined(GGML_USE_CUDA)
         params.n_gpu_layers = std::min((int32_t)n_layers, params.n_gpu_layers);
@@ -1723,7 +1724,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
         }
 
         // sychronize device profile to the master node
-        NodeType node_type;
+        NodeType node_type = NodeType::NODE_TYPE_WORKER;
         char is_forwarder[32] = {0};
         if (my_rank == 0) {
             if (auto_schedule) {

From dd589561b42a8a2ab49ef7526bcd942335c7807c Mon Sep 17 00:00:00 2001
From: Zonghang Li <zonghang.li@mbzuai.ac.ae>
Date: Thu, 19 Jun 2025 08:02:43 +0000
Subject: [PATCH 16/42] improve the computing buffer estimate

---
 README.md           |  1 +
 common/arg.cpp      |  7 +++++
 common/common.cpp   | 11 +++----
 common/common.h     |  1 +
 common/profiler.cpp |  4 +--
 common/profiler.h   | 16 ++++++++--
 include/llama.h     |  5 +--
 src/llama.cpp       | 76 ++++++++++++++++++++++++++++++++-------------
 8 files changed, 87 insertions(+), 34 deletions(-)

diff --git a/README.md b/README.md
index c69194f0..300ffa6a 100644
--- a/README.md
+++ b/README.md
@@ -135,6 +135,7 @@ mkdir build && cd build
 cmake ..
 make -j$(nproc)
 sudo make install
+sudo ldconfig
 ```
 
 **macOS:**
diff --git a/common/arg.cpp b/common/arg.cpp
index e282c80d..f1a33372 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -765,6 +765,13 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
             params.force = true;
         }
     ).set_env("LLAMA_ARG_FORCE"));
+    add_opt(llama_arg(
+        {"--master-priority"}, "N",
+        format("priority to assign workload to the master (default: %f, set 1.01 to use master first, and 0.99 to offload to other devices)", params.master_priority),
+        [](gpt_params & params, const std::string & value) {
+            params.master_priority = std::stof(value);
+        }
+    ).set_env("LLAMA_ARG_MASTER_PRIORITY"));
 // #ifdef GGML_USE_METAL
 //     // warn: if the output layer weights are not kept in metal shared memory, its mmap-ed weight data
 //     // could be released by the OS and reloaded repeatedly, which causes additional disk I/O latency.
diff --git a/common/common.cpp b/common/common.cpp
index 39b95d32..cc4536b6 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1053,7 +1053,7 @@ static bool assign_layers_to_device(
             GGML_ASSERT(!is_windows && "Windows is not tested yet\n");
 
             bool use_gpu = dev.gpu_support.metal || dev.gpu_support.cuda;
-            llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, use_gpu, m == 0, w[m] * k, n[m] * k);
+            llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, use_gpu, m == 0, dev_info_set[0].model_bytes, w[m] > n[m]);
 
             int  l_m          = w[m] * k;  // total number of layers assigned to device m
             int  l_m_gpu      = n[m] * k;  // number of layers assigned to device m that run on GPU
@@ -1248,10 +1248,8 @@ static bool assign_layers_to_device(
                     return cost * k;
                 }
             );
-            // apply higher priority to the head device, here 0.99 is a heuristic value
-            // to ensure that small models in homogeneous clusters result in 32:0 partitioning,
-            // rather than 1:31.
-            model.lp_.col_cost_[0] *= 0.99;
+            // apply priority to the head device
+            model.lp_.col_cost_[0] *= 1.0 / cparams.master_priority;
 
             // define the variable bounds
             model.lp_.col_lower_ = std::vector<double>(n_world * 2, 0.0);
@@ -1524,7 +1522,7 @@ static bool assign_layers_to_device(
     for (uint32_t m = 0; m < n_world; ++m) {
         const device_info & dev = dev_info_set[m];
         bool use_gpu = dev.gpu_support.metal || dev.gpu_support.cuda;
-        llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, use_gpu, m == 0, w[m], n[m]);
+        llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, use_gpu, m == 0, dev_info_set[0].model_bytes);
 
         if (dev.gpu_support.cuda || dev.gpu_support.metal) {
             int64_t required_mem = w[m] * b_prime;
@@ -2024,6 +2022,7 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param
     cparams.rank              = params.rank;
     cparams.prefetch          = params.prefetch;
     cparams.force             = params.force;
+    cparams.master_priority   = params.master_priority;
     cparams.keep_out_in_metal = params.keep_out_in_metal;
     cparams.n_gpu_layers      = params.n_gpu_layers;
     cparams.n_cycles          = params.n_cycles;
diff --git a/common/common.h b/common/common.h
index c6ffe136..cd78c173 100644
--- a/common/common.h
+++ b/common/common.h
@@ -152,6 +152,7 @@ struct gpt_params {
     bool    prefetch              = false; // prefetch layer weights
     bool    keep_out_in_metal     =  true; // whether to keep output weights in metal memory, true by default
     bool    force                 = false; // force to start prefetching after computation
+    float   master_priority       =  1.01; // priority to assign workload to the master (set 1.01 to use master first, and 0.99 to offload to other devices)
     int32_t gpu_mem               = 999.0; // gpu memory to use, in GiB
     int32_t n_cycles              =     0; // number of cycles to output one token
     int32_t n_predict             =    -1; // new tokens to predict
diff --git a/common/profiler.cpp b/common/profiler.cpp
index 18fe795d..292dc026 100644
--- a/common/profiler.cpp
+++ b/common/profiler.cpp
@@ -1603,10 +1603,10 @@ static float device_disk_access_delay(struct device_info & dev_info, struct llam
 
 #if defined(GGML_USE_METAL) || defined(GGML_USE_CUDA)
     llama_kv_size(&cpu_kv_size, &gpu_kv_size, model, cparams, true);
-    llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, true,  true, n_layers, n_gpu_layers);
+    llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, true,  true, n_bytes, n_layers > n_gpu_layers);
 #else
     llama_kv_size(&cpu_kv_size, &gpu_kv_size, model, cparams, false);
-    llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, false, true, n_layers, n_gpu_layers);
+    llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, false, true, n_bytes, n_layers > n_gpu_layers);
 #endif
 
     double cpu_kv_size_gib     = static_cast<double>(cpu_kv_size) / 1024.0 / 1024.0 / 1024.0;     // convert to GiB
diff --git a/common/profiler.h b/common/profiler.h
index ff69a454..c904ef98 100644
--- a/common/profiler.h
+++ b/common/profiler.h
@@ -293,10 +293,20 @@ struct model_bytes {
     int64_t nb_layer;
     int64_t nb_output;
 
+    // used to estimate the compute buffer size 
+    int64_t nb_output_w;
+    int64_t nb_attn_norm_w;
+    int64_t nb_ffn_gate_w;
+    int64_t nb_ffn_down_w;
+
     model_bytes() :
-        nb_input (0),
-        nb_layer (0),
-        nb_output(0) {}
+        nb_input      (0),
+        nb_layer      (0),
+        nb_output     (0), 
+        nb_output_w   (0),
+        nb_attn_norm_w(0),
+        nb_ffn_gate_w (0),
+        nb_ffn_down_w (0) {}
 };
 
 struct disk_props {
diff --git a/include/llama.h b/include/llama.h
index 3c220562..c61dd851 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -327,6 +327,7 @@ extern "C" {
         uint32_t    n_cycles;          // number of cycles to output one token
         bool        prefetch;          // whether to prefetch layer weights
         bool        force;             // force to start prefetching after computation
+        float       master_priority;   // priority to assign workload to the master (set 1.01 to use master first, and 0.99 to offload to other devices)
         bool        keep_out_in_metal; // whether to keep output weights in metal memory
         char *      master_ip;         // ip address of the master node
         char *      next_node_ip;      // ip address of the next node
@@ -575,8 +576,8 @@ extern "C" {
          const struct llama_context_params   cparams, 
                                       bool   use_gpu,
                                       bool   is_master,
-                                       int   n_layers,
-                                       int   n_gpu_layers);
+                        struct model_bytes   n_bytes,
+                                      bool   offload);
 
     // Return the size of KV cache in the model
     LLAMA_API void llama_total_kv_size(
diff --git a/src/llama.cpp b/src/llama.cpp
index 8b5af567..15255ea0 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -3679,6 +3679,8 @@ void llama_profile_device(
     dev_info->gpu_props.cuda_read_vram_bw   = device_cuda_read_vram_bw();
     dev_info->gpu_props.metal_mem_cpy_delay = device_metal_mem_copy(model);
     dev_info->gpu_props.cuda_mem_cpy_delay  = device_cuda_mem_copy(model);
+#else
+    (void)gpu_mem;
 #endif
 
     if (is_dtype_exist(n_params, GGML_TYPE_F32)) {
@@ -20263,6 +20265,7 @@ struct llama_context_params llama_context_default_params() {
         /*.n_cycles                    =*/ 0,
         /*.prefetch                    =*/ false,
         /*.force                       =*/ false,
+        /*.master_priority             =*/ 1.01,
         /*.keep_out_in_metal           =*/ true,
         /*.master_ip                   =*/ nullptr,
         /*.next_node_ip                =*/ nullptr,
@@ -21860,8 +21863,8 @@ void llama_model_compute_buf_size(
    const struct llama_context_params   cparams,
                                 bool   use_gpu,
                                 bool   is_master,
-                                 int   n_layers,
-                                 int   n_gpu_layers) {
+                  struct model_bytes   n_bytes,
+                                bool   offload) {
     const llama_hparams hparams = model->hparams;
 
     // input tensors
@@ -21872,34 +21875,61 @@ void llama_model_compute_buf_size(
     const int64_t n_bak_embd = hparams.n_embd  * cparams.n_ubatch;
     const int64_t n_inp_pos  = cparams.n_ubatch;
     const int64_t n_kq_mask  = cparams.n_ctx   * cparams.n_ubatch;
-    const int64_t n_inp_out_ids = cparams.n_ubatch;
     const int64_t n_norm     = hparams.n_embd  * cparams.n_ubatch;
-    const int64_t n_qcur     = hparams.n_embd  * cparams.n_ubatch * 2;
-    const int64_t n_kq       = cparams.n_ctx   * cparams.n_ubatch * hparams.n_head();
+    const int64_t n_qcur     = hparams.n_embd  * cparams.n_ubatch;
+    const int64_t n_ffn_gate = hparams.n_ff()  * cparams.n_ubatch;
+    const int64_t n_ffn_up   = hparams.n_ff()  * cparams.n_ubatch;
+    const int64_t n_inp_out_ids = cparams.n_ubatch;
 
     // outputs
     const int64_t n_out_embd = hparams.n_embd  * cparams.n_ubatch;
-    const int64_t n_output   = hparams.n_vocab * cparams.n_ubatch;
+    const int64_t n_result   = hparams.n_vocab * cparams.n_ubatch;
 
-    // compute buffer size for input, each layer, and output
-    const int64_t n_buf_inp  = (n_inp_toks + n_inp_embd) * ggml_type_size(GGML_TYPE_F32);
-    const int64_t n_buf_act  = (n_bak_embd + n_inp_pos + n_kq_mask + 
-                                 n_inp_out_ids + n_norm + n_qcur + n_kq
-                                ) * ggml_type_size(GGML_TYPE_F32);
-    const int64_t n_buf_out  = (n_out_embd + n_output) * ggml_type_size(GGML_TYPE_F32);
-
-    *cpu_buf = 0;
-    *gpu_buf = 0;
-    if (is_master) *cpu_buf = n_buf_inp + n_buf_out;
+    // weights
+    const int64_t nb_output_w    = n_bytes.nb_output_w;
+    const int64_t nb_attn_norm_w = n_bytes.nb_attn_norm_w;
+    const int64_t nb_ffn_gate_w  = n_bytes.nb_ffn_gate_w;
+    const int64_t nb_ffn_down_w  = n_bytes.nb_ffn_down_w;
+    
+    const int64_t nb_act_buf_base = (n_bak_embd + n_norm + n_inp_pos + n_ffn_gate) * ggml_type_size(GGML_TYPE_F32);
+    *gpu_buf = use_gpu ? nb_act_buf_base : 0;
+    *cpu_buf = nb_act_buf_base;
+    int64_t gpu_host_buf = 0;
 
+    // estimate GPU computing buffer and GPU-host computing buffer
     if (use_gpu) {
-        *gpu_buf += n_buf_act;
-        if (n_layers > n_gpu_layers) {
-            *cpu_buf += n_buf_act;
+        if (is_master) {
+            if (offload) {
+                *gpu_buf += (n_ffn_up + n_qcur) * ggml_type_size(GGML_TYPE_F32) + nb_attn_norm_w + nb_ffn_gate_w;
+            } else {
+                *gpu_buf += (n_ffn_up + n_qcur + n_kq_mask + n_inp_out_ids) * ggml_type_size(GGML_TYPE_F32);
+            }
+            *gpu_buf     += (n_out_embd + n_result) * ggml_type_size(GGML_TYPE_F32) + nb_output_w;
+            gpu_host_buf  = (n_inp_toks + n_inp_embd + n_bak_embd + n_inp_pos + n_kq_mask + n_inp_out_ids + n_out_embd) * ggml_type_size(GGML_TYPE_F32);
+        } else {
+            if (offload) {
+                *gpu_buf += n_qcur * ggml_type_size(GGML_TYPE_F32) + nb_attn_norm_w + nb_ffn_gate_w + nb_ffn_down_w;
+            } else {
+                *gpu_buf += (n_ffn_up + n_kq_mask) * ggml_type_size(GGML_TYPE_F32);
+            }
+            gpu_host_buf  = (n_bak_embd + n_inp_pos + n_kq_mask) * ggml_type_size(GGML_TYPE_F32);
         }
-    } else {
-        *cpu_buf += n_buf_act;
+    } 
+    
+    // estimate CPU computing buffer
+    {
+        if (is_master) {
+            *cpu_buf += (n_ffn_up + n_kq_mask + n_inp_out_ids + n_qcur + n_inp_toks + n_inp_embd + n_out_embd + n_result) * ggml_type_size(GGML_TYPE_F32);
+        } else {
+            *cpu_buf += (n_ffn_up + n_kq_mask) * ggml_type_size(GGML_TYPE_F32);
+        }
+        *cpu_buf += gpu_host_buf;
     }
+
+    LLAMA_LOG_INFO("%s: compute buffer size = %7.2f MiB (GPU) + %7.2f MiB (GPU-Host)\n", __func__,
+            *gpu_buf / (1024.0 * 1024.0), gpu_host_buf / (1024.0 * 1024.0));
+    LLAMA_LOG_INFO("%s: compute buffer size = %7.2f MiB (CPU)\n", __func__,
+            *cpu_buf / (1024.0 * 1024.0));
 }
 
 void llama_total_kv_size(
@@ -22045,6 +22075,7 @@ void llama_model_n_flops(
             if (blk_suffix == "attn_norm.weight" || blk_suffix == "ffn_norm.weight") {
                 count_n_flops (n_flops,  GGML_TYPE_F32, PROFILER_LAYER_BACKEND, 4 * n_embd + 1); // rms norm
                 count_n_flops (n_flops,  cur->type,     PROFILER_LAYER_BACKEND, n_embd); // norm weights
+                n_bytes->nb_attn_norm_w = std::max(n_bytes->nb_attn_norm_w, (int64_t)ggml_nbytes(cur));
             } else if (blk_suffix == "attn_q.weight") {
                 count_n_flops (n_flops,  cur->type,     PROFILER_LAYER_BACKEND, 2 * n_embd * n_embd);
                 count_n_flops (n_flops,  GGML_TYPE_F32, PROFILER_LAYER_BACKEND, 2.5 * n_embd);  // rope
@@ -22062,9 +22093,11 @@ void llama_model_n_flops(
             } else if (blk_suffix == "ffn_gate.weight") {
                 count_n_flops (n_flops,  cur->type,     PROFILER_LAYER_BACKEND, 2 * n_embd * n_ff);
                 count_n_flops (n_flops,  GGML_TYPE_F32, PROFILER_LAYER_BACKEND, 8 * n_ff); // SiLU
+                n_bytes->nb_ffn_gate_w = std::max(n_bytes->nb_ffn_gate_w, (int64_t)ggml_nbytes(cur));
             } else if (blk_suffix == "ffn_down.weight") {
                 count_n_flops (n_flops,  cur->type,     PROFILER_LAYER_BACKEND, 2 * n_embd * n_ff);
                 count_n_flops (n_flops,  GGML_TYPE_F32, PROFILER_LAYER_BACKEND, n_embd); // shortcut
+                n_bytes->nb_ffn_down_w = std::max(n_bytes->nb_ffn_down_w, (int64_t)ggml_nbytes(cur));
             } else if (blk_suffix == "ffn_up.weight") {
                 count_n_flops (n_flops,  cur->type,     PROFILER_LAYER_BACKEND, 2 * n_embd * n_ff);
                 count_n_flops (n_flops,  GGML_TYPE_F32, PROFILER_LAYER_BACKEND, n_ff); // silu(gate(x)) * up(x)
@@ -22097,6 +22130,7 @@ void llama_model_n_flops(
                 count_n_flops (n_flops,  GGML_TYPE_F32, PROFILER_LAYER_OUTPUT, 5 * n_vocab); // softmax
                 count_n_params(n_params, cur->type,     PROFILER_LAYER_OUTPUT, ggml_nelements(cur));
                 count_n_bytes (n_bytes,                 PROFILER_LAYER_OUTPUT, ggml_nbytes(cur));
+                n_bytes->nb_output_w = std::max(n_bytes->nb_output_w, (int64_t)ggml_nbytes(cur));
             } else if (tensor_name == "rope_freqs.weight") {
                 if (!rope_used) {
                     count_n_params(n_params, cur->type,     PROFILER_LAYER_BACKEND, ggml_nelements(cur));

From 80e5b71b489532956cba02d134efaba1336401e3 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Fri, 20 Jun 2025 13:43:55 +0400
Subject: [PATCH 17/42] fix compute buffer estimate: tested on metal

---
 common/arg.cpp      |  3 ++
 common/common.cpp   | 44 +++++++++++++++++++---
 common/profiler.cpp | 14 ++++++-
 include/llama.h     | 15 +++++++-
 src/llama.cpp       | 89 ++++++++++++++++++++++++++++++++++-----------
 5 files changed, 134 insertions(+), 31 deletions(-)

diff --git a/common/arg.cpp b/common/arg.cpp
index f1a33372..4039b19c 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -749,6 +749,9 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
         format("maximum GPU memory to use (default: %d)", params.gpu_mem),
         [](gpt_params & params, int value) {
             params.gpu_mem = value; // in GiB
+            if (value == 0) {
+                LOG_WRN("WARN: Set --gpu-mem to 0 may lead to errors during workload distribution.\n");
+            }
         }
     ).set_env("LLAMA_ARG_CUDA_MEM"));
     add_opt(llama_arg(
diff --git a/common/common.cpp b/common/common.cpp
index cc4536b6..4f1ddff3 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -847,6 +847,16 @@ static std::string vec_to_str(const std::vector<T> & vec) {
     return oss.str();
 }
 
+static backend_type get_backend_type(const gpu_support & support) {
+    if (support.cuda)    return BACKEND_CUDA;
+    if (support.metal)   return BACKEND_METAL;
+    if (support.vulkan)  return BACKEND_VULKAN;
+    if (support.kompute) return BACKEND_KOMPUTE;
+    if (support.gpublas) return BACKEND_GPUBLAS;
+    if (support.sycl)    return BACKEND_SYCL;
+    return BACKEND_CPU;
+}
+
 static bool assign_layers_to_device(
                                 uint32_t   n_world,
                        const device_info * dev_info_set, 
@@ -972,7 +982,7 @@ static bool assign_layers_to_device(
         bool is_android = strcmp(dev.device_os, "Android") == 0;
         bool is_windows = strcmp(dev.device_os, "Windows") == 0;
         GGML_ASSERT(!is_windows && "Windows is not tested yet\n");
-        
+
         if ((is_macos && !dev.gpu_support.metal) || is_linux) {
             mem_budget[m] = dev.memory.available_physical;
         } else if (is_macos && dev.gpu_support.metal) {
@@ -985,11 +995,21 @@ static bool assign_layers_to_device(
         }
     }
 
-    // initialize w_m proportionally to memory budget and n_m to 0
+    // initialize w_m proportionally to memory budget
     float total_mem_budget = std::accumulate(mem_budget.begin(), mem_budget.end(), 0.0f);
     for (uint32_t m = 0; m < n_world; ++m) {
         w[m] = std::round(mem_budget[m] / total_mem_budget * n_layer);
-        n[m] = 0;
+    }
+    // no 0 is allowed in w, it must be at least 1
+    for (uint32_t m = 0; m < n_world; ++m) {
+        if (w[m] == 0) {
+            w[m] = 1;
+            // find the maximum and decrease it by 1
+            auto max_it = std::max_element(w.begin(), w.end());
+            if (max_it != w.end() && *max_it > 1) {
+                *max_it -= 1;
+            }
+        }
     }
     // adjust w[m] to ensure L mod W = 0
     int diff = n_layer - std::accumulate(w.begin(), w.end(), 0);
@@ -997,6 +1017,15 @@ static bool assign_layers_to_device(
                              : std::min_element(mem_budget.begin(), mem_budget.end());
     w[std::distance(mem_budget.begin(), device)] += diff;
 
+    // initialize n_m to w_m (if there is GPU), assume all layers can run on GPU
+    for (uint32_t m = 0; m < n_world; ++m) {
+        if (dev_info_set[m].gpu_support.metal || dev_info_set[m].gpu_support.cuda) {
+            n[m] = w[m];
+        } else {
+            n[m] = 0;
+        }
+    }
+
     // stores the actual read bandwidth (GB/s) for each device
     std::vector<float> disk_speed(n_world, 0.0f);
     for (uint32_t m = 0; m < n_world; ++m) {
@@ -1052,8 +1081,7 @@ static bool assign_layers_to_device(
             bool is_windows = strcmp(dev.device_os, "Windows") == 0;
             GGML_ASSERT(!is_windows && "Windows is not tested yet\n");
 
-            bool use_gpu = dev.gpu_support.metal || dev.gpu_support.cuda;
-            llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, use_gpu, m == 0, dev_info_set[0].model_bytes, w[m] > n[m]);
+            llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, get_backend_type(dev.gpu_support), m == 0, dev_info_set[0].model_bytes, w[m] > n[m], n[m] > 0);
 
             int  l_m          = w[m] * k;  // total number of layers assigned to device m
             int  l_m_gpu      = n[m] * k;  // number of layers assigned to device m that run on GPU
@@ -1424,14 +1452,18 @@ static bool assign_layers_to_device(
                 if (n_m < static_cast<uint32_t>(std::floor(W * vec_z_gpu[m]))) {
                     // if there is still free GPU memory
                     has_free_gpu_memory = true;
+                    LOG_INF("Device %d still has free GPU memory: w_m = %d, n_m = %d, W * vec_z_gpu[m]) = %d\n", 
+                        m, w_m, n_m, static_cast<uint32_t>(std::floor(W * vec_z_gpu[m])));
                 }
                 if (w_m > n_m) {
                     // if layers are offloaded to CPU
                     has_gpu_overload = true;
+                    LOG_INF("Device %d has GPU overload: w_m = %d, n_m = %d\n", m, w_m, n_m);
                 }
             } else if (!in_set(m, M4)) {
                 // if the CPU is overloaded
                 has_cpu_overload = true;
+                LOG_INF("Device %d has CPU overload.\n", m);
             }
         }
 
@@ -1522,7 +1554,7 @@ static bool assign_layers_to_device(
     for (uint32_t m = 0; m < n_world; ++m) {
         const device_info & dev = dev_info_set[m];
         bool use_gpu = dev.gpu_support.metal || dev.gpu_support.cuda;
-        llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, use_gpu, m == 0, dev_info_set[0].model_bytes);
+        llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, get_backend_type(dev.gpu_support), m == 0, dev_info_set[0].model_bytes, true);
 
         if (dev.gpu_support.cuda || dev.gpu_support.metal) {
             int64_t required_mem = w[m] * b_prime;
diff --git a/common/profiler.cpp b/common/profiler.cpp
index 292dc026..c788a3e6 100644
--- a/common/profiler.cpp
+++ b/common/profiler.cpp
@@ -1603,10 +1603,20 @@ static float device_disk_access_delay(struct device_info & dev_info, struct llam
 
 #if defined(GGML_USE_METAL) || defined(GGML_USE_CUDA)
     llama_kv_size(&cpu_kv_size, &gpu_kv_size, model, cparams, true);
-    llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, true,  true, n_bytes, n_layers > n_gpu_layers);
+
+    enum backend_type backend;
+#if GGML_USE_METAL
+    backend = BACKEND_METAL;
+#elif GGML_USE_CUDA
+    backend = BACKEND_CUDA;
+#endif
+    llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, backend,  true, n_bytes, n_layers > n_gpu_layers, n_gpu_layers > 0);
+
 #else
     llama_kv_size(&cpu_kv_size, &gpu_kv_size, model, cparams, false);
-    llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, false, true, n_bytes, n_layers > n_gpu_layers);
+
+    enum backend_type backend = BACKEND_CPU;
+    llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, backend, true, n_bytes, n_layers > n_gpu_layers, n_gpu_layers > 0);
 #endif
 
     double cpu_kv_size_gib     = static_cast<double>(cpu_kv_size) / 1024.0 / 1024.0 / 1024.0;     // convert to GiB
diff --git a/include/llama.h b/include/llama.h
index c61dd851..fc42856c 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -67,6 +67,16 @@ extern "C" {
     typedef int32_t llama_token;
     typedef int32_t llama_seq_id;
 
+    enum backend_type {
+        BACKEND_CPU     = 0,
+        BACKEND_CUDA    = 1,
+        BACKEND_METAL   = 2,
+        BACKEND_VULKAN  = 3,
+        BACKEND_KOMPUTE = 4,
+        BACKEND_GPUBLAS = 5,
+        BACKEND_SYCL    = 6
+    };
+
     enum llama_vocab_type {
         LLAMA_VOCAB_TYPE_NONE = 0, // For models without vocab
         LLAMA_VOCAB_TYPE_SPM  = 1, // LLaMA tokenizer based on byte-level BPE with byte fallback
@@ -574,10 +584,11 @@ extern "C" {
                                    int64_t * gpu_buf,
                   const struct llama_model * model, 
          const struct llama_context_params   cparams, 
-                                      bool   use_gpu,
+                         enum backend_type   backend,
                                       bool   is_master,
                         struct model_bytes   n_bytes,
-                                      bool   offload);
+                                      bool   offload,
+                                      bool   has_gpu_layers);
 
     // Return the size of KV cache in the model
     LLAMA_API void llama_total_kv_size(
diff --git a/src/llama.cpp b/src/llama.cpp
index 15255ea0..eb551280 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -3673,6 +3673,7 @@ void llama_profile_device(
     // reserved/limit memory to avoid potential OOM, default to 300 MiB
     dev_info->gpu_props.memory_free         = round(gpu_props.memory_free  / (double)(1 << 30) * 100) / 100;
     dev_info->gpu_props.memory_free         = std::min((float)gpu_mem, dev_info->gpu_props.memory_free) - 0.3;
+    dev_info->gpu_props.memory_free         = std::max(dev_info->gpu_props.memory_free, 0.0f);
 
     dev_info->gpu_props.memory_total        = round(gpu_props.memory_total / (double)(1 << 30) * 100) / 100;
     dev_info->gpu_props.metal_read_vram_bw  = device_metal_read_vram_bw();
@@ -21861,10 +21862,11 @@ void llama_model_compute_buf_size(
                              int64_t * gpu_buf, 
             const struct llama_model * model, 
    const struct llama_context_params   cparams,
-                                bool   use_gpu,
+                   enum backend_type   backend,
                                 bool   is_master,
                   struct model_bytes   n_bytes,
-                                bool   offload) {
+                                bool   offload,
+                                bool   has_gpu_layers) {
     const llama_hparams hparams = model->hparams;
 
     // input tensors
@@ -21879,6 +21881,9 @@ void llama_model_compute_buf_size(
     const int64_t n_qcur     = hparams.n_embd  * cparams.n_ubatch;
     const int64_t n_ffn_gate = hparams.n_ff()  * cparams.n_ubatch;
     const int64_t n_ffn_up   = hparams.n_ff()  * cparams.n_ubatch;
+    const int64_t n_ffn_out  = hparams.n_embd  * cparams.n_ubatch;
+    const int64_t n_ffn_inp  = hparams.n_embd  * cparams.n_ubatch;
+    const int64_t n_kq       = cparams.n_ctx   * cparams.n_ubatch * hparams.n_head();
     const int64_t n_inp_out_ids = cparams.n_ubatch;
 
     // outputs
@@ -21890,40 +21895,82 @@ void llama_model_compute_buf_size(
     const int64_t nb_attn_norm_w = n_bytes.nb_attn_norm_w;
     const int64_t nb_ffn_gate_w  = n_bytes.nb_ffn_gate_w;
     const int64_t nb_ffn_down_w  = n_bytes.nb_ffn_down_w;
-    
-    const int64_t nb_act_buf_base = (n_bak_embd + n_norm + n_inp_pos + n_ffn_gate) * ggml_type_size(GGML_TYPE_F32);
-    *gpu_buf = use_gpu ? nb_act_buf_base : 0;
-    *cpu_buf = nb_act_buf_base;
+
+    const int64_t type_size_f32 = ggml_type_size(GGML_TYPE_F32);
+
+    bool use_gpu = backend != BACKEND_CPU && has_gpu_layers;
+    *gpu_buf = 0;
+    *cpu_buf = 0;
     int64_t gpu_host_buf = 0;
 
-    // estimate GPU computing buffer and GPU-host computing buffer
-    if (use_gpu) {
+    if (backend == BACKEND_CUDA) {
+        const int64_t nb_act_buf_base = (n_bak_embd + n_norm + n_inp_pos + n_ffn_gate) * type_size_f32;
+        *gpu_buf = use_gpu ? nb_act_buf_base : 0;
+
+        // CUDA computing buffer and CUDA-host buffer
         if (is_master) {
             if (offload) {
-                *gpu_buf += (n_ffn_up + n_qcur) * ggml_type_size(GGML_TYPE_F32) + nb_attn_norm_w + nb_ffn_gate_w;
+                *gpu_buf += (n_ffn_up + n_qcur) * type_size_f32 + nb_attn_norm_w + nb_ffn_gate_w;
             } else {
-                *gpu_buf += (n_ffn_up + n_qcur + n_kq_mask + n_inp_out_ids) * ggml_type_size(GGML_TYPE_F32);
+                *gpu_buf += (n_ffn_up + n_qcur + n_kq_mask + n_inp_out_ids) * type_size_f32;
             }
-            *gpu_buf     += (n_out_embd + n_result) * ggml_type_size(GGML_TYPE_F32) + nb_output_w;
-            gpu_host_buf  = (n_inp_toks + n_inp_embd + n_bak_embd + n_inp_pos + n_kq_mask + n_inp_out_ids + n_out_embd) * ggml_type_size(GGML_TYPE_F32);
+            *gpu_buf     += (n_out_embd + n_result) * type_size_f32 + nb_output_w;
+            gpu_host_buf  = (n_inp_toks + n_inp_embd + n_bak_embd + n_inp_pos + n_kq_mask + n_inp_out_ids + n_out_embd) * type_size_f32;
         } else {
             if (offload) {
-                *gpu_buf += n_qcur * ggml_type_size(GGML_TYPE_F32) + nb_attn_norm_w + nb_ffn_gate_w + nb_ffn_down_w;
+                *gpu_buf += n_qcur * type_size_f32 + nb_attn_norm_w + nb_ffn_gate_w + nb_ffn_down_w;
             } else {
-                *gpu_buf += (n_ffn_up + n_kq_mask) * ggml_type_size(GGML_TYPE_F32);
+                *gpu_buf += (n_ffn_up + n_kq_mask) * type_size_f32;
             }
-            gpu_host_buf  = (n_bak_embd + n_inp_pos + n_kq_mask) * ggml_type_size(GGML_TYPE_F32);
+            gpu_host_buf  = (n_bak_embd + n_inp_pos + n_kq_mask) * type_size_f32;
         }
     } 
     
-    // estimate CPU computing buffer
-    {
-        if (is_master) {
-            *cpu_buf += (n_ffn_up + n_kq_mask + n_inp_out_ids + n_qcur + n_inp_toks + n_inp_embd + n_out_embd + n_result) * ggml_type_size(GGML_TYPE_F32);
+    else if (backend == BACKEND_METAL) {
+        const int64_t nb_act_buf_base = (n_inp_pos + n_kq_mask) * type_size_f32;
+        *gpu_buf = nb_act_buf_base;
+        *cpu_buf = nb_act_buf_base;
+
+        if (use_gpu) {
+            if (is_master) {
+                *cpu_buf += (n_inp_toks + n_inp_embd + n_bak_embd + n_inp_out_ids + n_out_embd + n_result) * type_size_f32;
+
+                if (offload) {
+                    *gpu_buf += (n_ffn_out + n_ffn_inp + n_inp_out_ids) * type_size_f32;
+                    *gpu_buf += std::max(n_ffn_up + n_ffn_gate, n_qcur + n_qcur + n_kq) * type_size_f32;
+                    *cpu_buf += n_norm * type_size_f32;
+                    *cpu_buf += std::max(n_ffn_up + n_ffn_gate, n_qcur + n_qcur + n_kq) * type_size_f32;
+                } else {
+                    *gpu_buf += (n_bak_embd + n_inp_out_ids + n_norm) * type_size_f32;
+                    *gpu_buf += std::max(n_ffn_up + n_ffn_gate, n_qcur + n_qcur + n_kq) * type_size_f32;
+                }
+            } else {
+                *gpu_buf += (n_ffn_out + n_ffn_inp) * type_size_f32;
+                *gpu_buf += std::max(n_ffn_up + n_ffn_gate, n_qcur + n_qcur + n_kq) * type_size_f32;
+
+                *cpu_buf += n_bak_embd * type_size_f32;
+                if (offload) {
+                    *cpu_buf += n_norm * type_size_f32;
+                    *cpu_buf += std::max(n_ffn_up + n_ffn_gate, n_qcur + n_qcur + n_kq) * type_size_f32;
+                }
+            }
         } else {
-            *cpu_buf += (n_ffn_up + n_kq_mask) * ggml_type_size(GGML_TYPE_F32);
+            *gpu_buf = 0;
+            *cpu_buf = 0;
         }
-        *cpu_buf += gpu_host_buf;
+    }
+
+    else if (backend != BACKEND_CPU) {
+        GGML_ASSERT(false && "Unsupported backend type for compute buffer estimation.\n");
+    }
+
+    // CPU computing buffer
+    if (*cpu_buf == 0) {
+        *cpu_buf = (n_inp_pos + n_kq_mask + n_bak_embd + n_norm) * type_size_f32;
+        if (is_master) {
+            *cpu_buf += (n_inp_toks + n_inp_embd + n_inp_out_ids + n_out_embd + n_result) * type_size_f32;   
+        }
+        *cpu_buf += std::max(n_ffn_gate + n_ffn_up, n_qcur + n_qcur + n_kq) * type_size_f32;
     }
 
     LLAMA_LOG_INFO("%s: compute buffer size = %7.2f MiB (GPU) + %7.2f MiB (GPU-Host)\n", __func__,

From 45e8b0420c0ad1020106d782517e1f9de9c3ab14 Mon Sep 17 00:00:00 2001
From: Zonghang Li <zonghang.li@mbzuai.ac.ae>
Date: Sun, 22 Jun 2025 08:10:57 +0000
Subject: [PATCH 18/42] fix compute buffer estimate: tested on cuda

---
 common/common.cpp   |  5 ++--
 common/profiler.cpp |  6 ++--
 common/profiler.h   | 18 +++++++-----
 include/llama.h     |  2 +-
 src/llama.cpp       | 67 +++++++++++++++++++++++++++++++++------------
 5 files changed, 68 insertions(+), 30 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 4f1ddff3..c241b54f 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1081,7 +1081,7 @@ static bool assign_layers_to_device(
             bool is_windows = strcmp(dev.device_os, "Windows") == 0;
             GGML_ASSERT(!is_windows && "Windows is not tested yet\n");
 
-            llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, get_backend_type(dev.gpu_support), m == 0, dev_info_set[0].model_bytes, w[m] > n[m], n[m] > 0);
+            llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, get_backend_type(dev.gpu_support), m, dev_info_set[0].model_bytes, w[m] > n[m], n[m] > 0);
 
             int  l_m          = w[m] * k;  // total number of layers assigned to device m
             int  l_m_gpu      = n[m] * k;  // number of layers assigned to device m that run on GPU
@@ -1242,6 +1242,7 @@ static bool assign_layers_to_device(
                 if (dev.gpu_support.metal && m == 0 && cparams.keep_out_in_metal) {
                     vec_z_gpu[m] -= (double)bo / (double)(n_layer * b_prime);
                 }
+                vec_z_gpu[m] = std::max(vec_z_gpu[m], 0.0f);
             }
         }
 
@@ -1554,7 +1555,7 @@ static bool assign_layers_to_device(
     for (uint32_t m = 0; m < n_world; ++m) {
         const device_info & dev = dev_info_set[m];
         bool use_gpu = dev.gpu_support.metal || dev.gpu_support.cuda;
-        llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, get_backend_type(dev.gpu_support), m == 0, dev_info_set[0].model_bytes, true);
+        llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, get_backend_type(dev.gpu_support), m, dev_info_set[0].model_bytes, true);
 
         if (dev.gpu_support.cuda || dev.gpu_support.metal) {
             int64_t required_mem = w[m] * b_prime;
diff --git a/common/profiler.cpp b/common/profiler.cpp
index c788a3e6..afc427f1 100644
--- a/common/profiler.cpp
+++ b/common/profiler.cpp
@@ -924,7 +924,7 @@ static void check_env_path() {
     setenv("PATH", update_env_path.c_str(), 1);
 }
 
-static void external_fio_impl(float * read_bw, float * write_bw, bool op_rand, int n_threads) {
+static void external_fio_impl(float * read_bw, float * write_bw, bool op_rand, int n_threads) {    
     pid_t pid = getpid(); // avoid conflict with other processes
 
     std::string test_file   = "fio_test_"   + std::to_string(pid);
@@ -1610,13 +1610,13 @@ static float device_disk_access_delay(struct device_info & dev_info, struct llam
 #elif GGML_USE_CUDA
     backend = BACKEND_CUDA;
 #endif
-    llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, backend,  true, n_bytes, n_layers > n_gpu_layers, n_gpu_layers > 0);
+    llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, backend, 0, n_bytes, n_layers > n_gpu_layers, n_gpu_layers > 0);
 
 #else
     llama_kv_size(&cpu_kv_size, &gpu_kv_size, model, cparams, false);
 
     enum backend_type backend = BACKEND_CPU;
-    llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, backend, true, n_bytes, n_layers > n_gpu_layers, n_gpu_layers > 0);
+    llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, backend, 0, n_bytes, n_layers > n_gpu_layers, n_gpu_layers > 0);
 #endif
 
     double cpu_kv_size_gib     = static_cast<double>(cpu_kv_size) / 1024.0 / 1024.0 / 1024.0;     // convert to GiB
diff --git a/common/profiler.h b/common/profiler.h
index c904ef98..00407c96 100644
--- a/common/profiler.h
+++ b/common/profiler.h
@@ -295,18 +295,22 @@ struct model_bytes {
 
     // used to estimate the compute buffer size 
     int64_t nb_output_w;
+    int64_t nb_output_norm_w;
     int64_t nb_attn_norm_w;
+    int64_t nb_attn_q_w;
     int64_t nb_ffn_gate_w;
     int64_t nb_ffn_down_w;
 
     model_bytes() :
-        nb_input      (0),
-        nb_layer      (0),
-        nb_output     (0), 
-        nb_output_w   (0),
-        nb_attn_norm_w(0),
-        nb_ffn_gate_w (0),
-        nb_ffn_down_w (0) {}
+        nb_input        (0),
+        nb_layer        (0),
+        nb_output       (0), 
+        nb_output_w     (0),
+        nb_output_norm_w(0),
+        nb_attn_norm_w  (0),
+        nb_attn_q_w     (0),
+        nb_ffn_gate_w   (0),
+        nb_ffn_down_w   (0) {}
 };
 
 struct disk_props {
diff --git a/include/llama.h b/include/llama.h
index fc42856c..2faecdfe 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -585,7 +585,7 @@ extern "C" {
                   const struct llama_model * model, 
          const struct llama_context_params   cparams, 
                          enum backend_type   backend,
-                                      bool   is_master,
+                                       int   my_rank,
                         struct model_bytes   n_bytes,
                                       bool   offload,
                                       bool   has_gpu_layers);
diff --git a/src/llama.cpp b/src/llama.cpp
index eb551280..de49cdce 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -21863,11 +21863,12 @@ void llama_model_compute_buf_size(
             const struct llama_model * model, 
    const struct llama_context_params   cparams,
                    enum backend_type   backend,
-                                bool   is_master,
+                                 int   my_rank,
                   struct model_bytes   n_bytes,
                                 bool   offload,
                                 bool   has_gpu_layers) {
     const llama_hparams hparams = model->hparams;
+    bool is_master = my_rank == 0;
 
     // input tensors
     const int64_t n_inp_toks = cparams.n_ubatch;
@@ -21884,6 +21885,8 @@ void llama_model_compute_buf_size(
     const int64_t n_ffn_out  = hparams.n_embd  * cparams.n_ubatch;
     const int64_t n_ffn_inp  = hparams.n_embd  * cparams.n_ubatch;
     const int64_t n_kq       = cparams.n_ctx   * cparams.n_ubatch * hparams.n_head();
+    const int64_t n_k        = cparams.n_ctx   * hparams.n_embd_head_k * hparams.n_head_kv();
+    const int64_t n_v        = cparams.n_ctx   * hparams.n_embd_head_v * hparams.n_head_kv();
     const int64_t n_inp_out_ids = cparams.n_ubatch;
 
     // outputs
@@ -21893,34 +21896,60 @@ void llama_model_compute_buf_size(
     // weights
     const int64_t nb_output_w    = n_bytes.nb_output_w;
     const int64_t nb_attn_norm_w = n_bytes.nb_attn_norm_w;
-    const int64_t nb_ffn_gate_w  = n_bytes.nb_ffn_gate_w;
-    const int64_t nb_ffn_down_w  = n_bytes.nb_ffn_down_w;
+    const int64_t nb_attn_q_w    = n_bytes.nb_attn_q_w;
 
+    // format bytes
     const int64_t type_size_f32 = ggml_type_size(GGML_TYPE_F32);
+    const int64_t type_size_f16 = ggml_type_size(GGML_TYPE_F16);
 
     bool use_gpu = backend != BACKEND_CPU && has_gpu_layers;
     *gpu_buf = 0;
     *cpu_buf = 0;
     int64_t gpu_host_buf = 0;
 
+    // GPU compute buffer
+    // estimate the GPU compute buffer, here we can only estimate the upper bound of various models,
+    // but cannot estimate the exact value.
     if (backend == BACKEND_CUDA) {
-        const int64_t nb_act_buf_base = (n_bak_embd + n_norm + n_inp_pos + n_ffn_gate) * type_size_f32;
-        *gpu_buf = use_gpu ? nb_act_buf_base : 0;
+        *gpu_buf = (n_bak_embd + n_norm) * type_size_f32;
 
-        // CUDA computing buffer and CUDA-host buffer
         if (is_master) {
-            if (offload) {
-                *gpu_buf += (n_ffn_up + n_qcur) * type_size_f32 + nb_attn_norm_w + nb_ffn_gate_w;
+            if (has_gpu_layers) {
+                if (offload) {
+                    *gpu_buf += std::max<int64_t>({
+                        (n_qcur + n_inp_pos + n_kq_mask + n_inp_out_ids) * type_size_f32 + nb_attn_norm_w,
+                        (n_qcur + n_inp_pos + n_norm) * type_size_f32 + nb_attn_norm_w,
+                        (n_qcur + n_qcur + n_kq_mask + n_inp_pos) * type_size_f32,
+                        (n_qcur + n_qcur + n_inp_pos) * type_size_f32 + nb_attn_q_w,
+                        n_inp_pos * type_size_f32 + (n_k + n_v) * type_size_f16 + nb_attn_q_w
+                    });
+                } else {
+                    *gpu_buf += (n_qcur + n_inp_pos + n_kq_mask + n_inp_out_ids) * type_size_f32;
+                }
+                *gpu_buf += (n_qcur + n_kq) * type_size_f32;
             } else {
-                *gpu_buf += (n_ffn_up + n_qcur + n_kq_mask + n_inp_out_ids) * type_size_f32;
+                *gpu_buf += (n_qcur + n_kq) * type_size_f32;
+                *gpu_buf += std::max<int64_t>({
+                    (n_kq_mask + n_qcur + n_inp_pos) * type_size_f32 + nb_attn_norm_w,
+                    (n_inp_pos + n_kq_mask) * type_size_f32 + n_v * type_size_f16 + nb_attn_norm_w,
+                });
             }
             *gpu_buf     += (n_out_embd + n_result) * type_size_f32 + nb_output_w;
             gpu_host_buf  = (n_inp_toks + n_inp_embd + n_bak_embd + n_inp_pos + n_kq_mask + n_inp_out_ids + n_out_embd) * type_size_f32;
         } else {
-            if (offload) {
-                *gpu_buf += n_qcur * type_size_f32 + nb_attn_norm_w + nb_ffn_gate_w + nb_ffn_down_w;
+            if (has_gpu_layers) {
+                if (offload) {
+                    *gpu_buf += (n_kq + n_qcur) * type_size_f32;
+                    *gpu_buf += std::max<int64_t>({
+                        (n_inp_pos + n_norm + n_kq_mask) * type_size_f32 + nb_attn_norm_w,
+                        (n_inp_pos + n_norm + n_qcur)    * type_size_f32 + nb_attn_norm_w,
+                        n_inp_pos * type_size_f32 + (n_k + n_v) * type_size_f16 + nb_attn_q_w,
+                    });
+                } else {
+                    *gpu_buf += (n_inp_pos + n_kq_mask + n_qcur + n_qcur + n_kq) * type_size_f32;
+                }
             } else {
-                *gpu_buf += (n_ffn_up + n_kq_mask) * type_size_f32;
+                *gpu_buf += (n_qcur + n_kq + n_kq_mask + n_qcur + n_inp_pos) * type_size_f32 + nb_attn_norm_w;
             }
             gpu_host_buf  = (n_bak_embd + n_inp_pos + n_kq_mask) * type_size_f32;
         }
@@ -21964,19 +21993,22 @@ void llama_model_compute_buf_size(
         GGML_ASSERT(false && "Unsupported backend type for compute buffer estimation.\n");
     }
 
-    // CPU computing buffer
+    // CPU compute buffer
     if (*cpu_buf == 0) {
         *cpu_buf = (n_inp_pos + n_kq_mask + n_bak_embd + n_norm) * type_size_f32;
         if (is_master) {
             *cpu_buf += (n_inp_toks + n_inp_embd + n_inp_out_ids + n_out_embd + n_result) * type_size_f32;   
         }
         *cpu_buf += std::max(n_ffn_gate + n_ffn_up, n_qcur + n_qcur + n_kq) * type_size_f32;
+        *cpu_buf += gpu_host_buf;
     }
 
-    LLAMA_LOG_INFO("%s: compute buffer size = %7.2f MiB (GPU) + %7.2f MiB (GPU-Host)\n", __func__,
-            *gpu_buf / (1024.0 * 1024.0), gpu_host_buf / (1024.0 * 1024.0));
-    LLAMA_LOG_INFO("%s: compute buffer size = %7.2f MiB (CPU)\n", __func__,
-            *cpu_buf / (1024.0 * 1024.0));
+    LLAMA_LOG_INFO("\n");
+    LLAMA_LOG_INFO("%s: here the compute buffer size is a predicted upper bound, not an exact value\n", __func__);
+    LLAMA_LOG_INFO("%s: (rank %d) compute buffer size = %7.2f MiB (GPU) + %7.2f MiB (GPU-Host)\n", __func__,
+            my_rank, *gpu_buf / (1024.0 * 1024.0), gpu_host_buf / (1024.0 * 1024.0));
+    LLAMA_LOG_INFO("%s: (rank %d) compute buffer size = %7.2f MiB (CPU and GPU-Host buffer)\n", __func__,
+            my_rank, *cpu_buf / (1024.0 * 1024.0));
 }
 
 void llama_total_kv_size(
@@ -22126,6 +22158,7 @@ void llama_model_n_flops(
             } else if (blk_suffix == "attn_q.weight") {
                 count_n_flops (n_flops,  cur->type,     PROFILER_LAYER_BACKEND, 2 * n_embd * n_embd);
                 count_n_flops (n_flops,  GGML_TYPE_F32, PROFILER_LAYER_BACKEND, 2.5 * n_embd);  // rope
+                n_bytes->nb_attn_q_w = std::max(n_bytes->nb_attn_q_w, (int64_t)ggml_nbytes(cur));
             } else if (blk_suffix == "attn_k.weight") {
                 count_n_flops (n_flops,  cur->type,     PROFILER_LAYER_BACKEND, 2 * n_embd * n_head_kv * n_embd_head_k);
                 count_n_flops (n_flops,  GGML_TYPE_F32, PROFILER_LAYER_BACKEND, 2.5 * n_embd_head_k * n_head_kv); // rope

From c926088d6af0e39b84888bc0a2848bd208bba067 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Sun, 22 Jun 2025 16:27:55 +0400
Subject: [PATCH 19/42] fix compute buffer estimate: test without highs

---
 common/common.cpp | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index c241b54f..472205b9 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1016,7 +1016,6 @@ static bool assign_layers_to_device(
     auto device = (diff > 0) ? std::max_element(mem_budget.begin(), mem_budget.end()) 
                              : std::min_element(mem_budget.begin(), mem_budget.end());
     w[std::distance(mem_budget.begin(), device)] += diff;
-
     // initialize n_m to w_m (if there is GPU), assume all layers can run on GPU
     for (uint32_t m = 0; m < n_world; ++m) {
         if (dev_info_set[m].gpu_support.metal || dev_info_set[m].gpu_support.cuda) {
@@ -1540,13 +1539,24 @@ static bool assign_layers_to_device(
     float total_mem_budget = std::accumulate(mem_budget.begin(), mem_budget.end(), 0.0f);
     for (uint32_t m = 0; m < n_world; ++m) {
         w[m] = std::round(mem_budget[m] / total_mem_budget * n_layer);
-        n[m] = 0;
+    }
+    // no 0 is allowed in w, it must be at least 1
+    for (uint32_t m = 0; m < n_world; ++m) {
+        if (w[m] == 0) {
+            w[m] = 1;
+            // find the maximum and decrease it by 1
+            auto max_it = std::max_element(w.begin(), w.end());
+            if (max_it != w.end() && *max_it > 1) {
+                *max_it -= 1;
+            }
+        }
     }
     // adjust w[m] to ensure L mod W = 0
     int diff = n_layer - std::accumulate(w.begin(), w.end(), 0);
     auto device = (diff > 0) ? std::max_element(mem_budget.begin(), mem_budget.end()) 
                             : std::min_element(mem_budget.begin(), mem_budget.end());
     w[std::distance(mem_budget.begin(), device)] += diff;
+
     std::copy(w.begin(), w.end(), n_layer_window);
 
     std::vector<float> vec_z_gpu(n_world, 0.0f);
@@ -1554,8 +1564,7 @@ static bool assign_layers_to_device(
 
     for (uint32_t m = 0; m < n_world; ++m) {
         const device_info & dev = dev_info_set[m];
-        bool use_gpu = dev.gpu_support.metal || dev.gpu_support.cuda;
-        llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, get_backend_type(dev.gpu_support), m, dev_info_set[0].model_bytes, true);
+        llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, get_backend_type(dev.gpu_support), m, dev_info_set[0].model_bytes, false, true);
 
         if (dev.gpu_support.cuda || dev.gpu_support.metal) {
             int64_t required_mem = w[m] * b_prime;

From 16ba3564cecda24ba283bbfad2aac064807795b2 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Tue, 24 Jun 2025 16:09:59 +0400
Subject: [PATCH 20/42] fix compute_buffer estimate: add context GPU usage

---
 common/common.cpp |  2 +-
 src/llama.cpp     | 26 +++++++++++++++-----------
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 472205b9..7520b703 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -986,7 +986,7 @@ static bool assign_layers_to_device(
         if ((is_macos && !dev.gpu_support.metal) || is_linux) {
             mem_budget[m] = dev.memory.available_physical;
         } else if (is_macos && dev.gpu_support.metal) {
-            mem_budget[m] = dev.gpu_props.memory_free;
+            mem_budget[m] = dev.gpu_props.memory_free + 1e-4; // to avoid division by zero
         } else if (is_android) {
             mem_budget[m] = dev.memory.available_physical + dev.memory.used_can_swap;
         } else {
diff --git a/src/llama.cpp b/src/llama.cpp
index de49cdce..dd232fe0 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -3669,12 +3669,8 @@ void llama_profile_device(
 #if defined(GGML_USE_METAL) || defined(GGML_USE_CUDA)
     dev_info->gpu_props.name                = gpu_props.name;
     dev_info->gpu_props.description         = gpu_props.description;
-
-    // reserved/limit memory to avoid potential OOM, default to 300 MiB
-    dev_info->gpu_props.memory_free         = round(gpu_props.memory_free  / (double)(1 << 30) * 100) / 100;
-    dev_info->gpu_props.memory_free         = std::min((float)gpu_mem, dev_info->gpu_props.memory_free) - 0.3;
-    dev_info->gpu_props.memory_free         = std::max(dev_info->gpu_props.memory_free, 0.0f);
-
+    dev_info->gpu_props.memory_free         = round(gpu_props.memory_free / (double)(1 << 30) * 100) / 100;
+    dev_info->gpu_props.memory_free         = std::min((float)gpu_mem, dev_info->gpu_props.memory_free);
     dev_info->gpu_props.memory_total        = round(gpu_props.memory_total / (double)(1 << 30) * 100) / 100;
     dev_info->gpu_props.metal_read_vram_bw  = device_metal_read_vram_bw();
     dev_info->gpu_props.cuda_read_vram_bw   = device_cuda_read_vram_bw();
@@ -21324,6 +21320,9 @@ void * llama_context_setup_backend(
             std::vector<ggml_cgraph *> gf = llama_build_graph(*ctx, ubatch, true);
 
             GGML_ASSERT(gf.size() <= MAX_SCHEDULERS && "Number of subgraphs exceeds the maximum number of schedulers\n");
+            for (size_t i = gf.size(); i < ctx->sched.size(); ++i) {
+                ggml_backend_sched_free(ctx->sched[i]);
+            }
             ctx->sched.resize(gf.size());
 
             // initialize scheduler with the worst-case graph
@@ -21993,7 +21992,7 @@ void llama_model_compute_buf_size(
         GGML_ASSERT(false && "Unsupported backend type for compute buffer estimation.\n");
     }
 
-    // CPU compute buffer
+    // CPU compute buffer for NUMA system or Metal with ngl=0
     if (*cpu_buf == 0) {
         *cpu_buf = (n_inp_pos + n_kq_mask + n_bak_embd + n_norm) * type_size_f32;
         if (is_master) {
@@ -22005,10 +22004,15 @@ void llama_model_compute_buf_size(
 
     LLAMA_LOG_INFO("\n");
     LLAMA_LOG_INFO("%s: here the compute buffer size is a predicted upper bound, not an exact value\n", __func__);
-    LLAMA_LOG_INFO("%s: (rank %d) compute buffer size = %7.2f MiB (GPU) + %7.2f MiB (GPU-Host)\n", __func__,
-            my_rank, *gpu_buf / (1024.0 * 1024.0), gpu_host_buf / (1024.0 * 1024.0));
-    LLAMA_LOG_INFO("%s: (rank %d) compute buffer size = %7.2f MiB (CPU and GPU-Host buffer)\n", __func__,
-            my_rank, *cpu_buf / (1024.0 * 1024.0));
+    LLAMA_LOG_INFO("%s: (rank %d) compute buffer size = %7.2f MiB (GPU) + %7.2f MiB (CPU / GPU-host buffer)\n", __func__,
+            my_rank, *gpu_buf / (1024.0 * 1024.0), *cpu_buf / (1024.0 * 1024.0));
+    
+    if (backend == BACKEND_CUDA) {
+        // context GPU memory usage, i.e. the initial memory cost of creating a CUDA context, 
+        // even before you launch any kernels or allocate your own buffers.
+        // this value may vary by GPU and CUDA version, but it's lower than 400 MiB in most cases.
+        *gpu_buf += 400 * 1024 * 1024;
+    }
 }
 
 void llama_total_kv_size(

From 90b1079d78715350731545d03d90c968e862ebb8 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Tue, 24 Jun 2025 16:37:16 +0400
Subject: [PATCH 21/42] fix compute_buffer estimate: remove unused memory for
 CUDA device

---
 src/llama.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/llama.cpp b/src/llama.cpp
index dd232fe0..45b15223 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -21994,11 +21994,13 @@ void llama_model_compute_buf_size(
 
     // CPU compute buffer for NUMA system or Metal with ngl=0
     if (*cpu_buf == 0) {
-        *cpu_buf = (n_inp_pos + n_kq_mask + n_bak_embd + n_norm) * type_size_f32;
         if (is_master) {
             *cpu_buf += (n_inp_toks + n_inp_embd + n_inp_out_ids + n_out_embd + n_result) * type_size_f32;   
         }
-        *cpu_buf += std::max(n_ffn_gate + n_ffn_up, n_qcur + n_qcur + n_kq) * type_size_f32;
+        if (offload) {
+            *cpu_buf = (n_inp_pos + n_kq_mask + n_bak_embd + n_norm) * type_size_f32;
+            *cpu_buf += std::max(n_ffn_gate + n_ffn_up, n_qcur + n_qcur + n_kq) * type_size_f32;
+        }
         *cpu_buf += gpu_host_buf;
     }
 

From 4dde8458cf2e4e13dae84ae074a945f3d011dd01 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Tue, 24 Jun 2025 19:29:10 +0400
Subject: [PATCH 22/42] fix compute buffer estimate: reserve 100 MiB VRAM to
 avoid potential OOM

---
 src/llama.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/llama.cpp b/src/llama.cpp
index 45b15223..a97e4ae1 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -22012,8 +22012,9 @@ void llama_model_compute_buf_size(
     if (backend == BACKEND_CUDA) {
         // context GPU memory usage, i.e. the initial memory cost of creating a CUDA context, 
         // even before you launch any kernels or allocate your own buffers.
-        // this value may vary by GPU and CUDA version, but it's lower than 400 MiB in most cases.
-        *gpu_buf += 400 * 1024 * 1024;
+        // this value may vary by GPU and CUDA version, but it's lower than 400 MiB in most cases,
+        // another 100 MiB is used to prevent accidental OOM.
+        *gpu_buf += 500 * 1024 * 1024;
     }
 }
 

From 72701ae8724fc2bfe7e86884b26b6762281e746f Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Tue, 24 Jun 2025 20:39:49 +0400
Subject: [PATCH 23/42] fix compute buffer estimate: reserve 200 MiB VRAM to
 avoid potential OOM

---
 src/llama.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/llama.cpp b/src/llama.cpp
index a97e4ae1..de548a66 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -22013,8 +22013,8 @@ void llama_model_compute_buf_size(
         // context GPU memory usage, i.e. the initial memory cost of creating a CUDA context, 
         // even before you launch any kernels or allocate your own buffers.
         // this value may vary by GPU and CUDA version, but it's lower than 400 MiB in most cases,
-        // another 100 MiB is used to prevent accidental OOM.
-        *gpu_buf += 500 * 1024 * 1024;
+        // another 200 MiB is used to prevent accidental OOM.
+        *gpu_buf += 600 * 1024 * 1024;
     }
 }
 

From 50807fd4e1ec79ebd820071bc718335ca28863a1 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Thu, 26 Jun 2025 08:56:31 +0400
Subject: [PATCH 24/42] halda: handle infeasible solution with weak device

---
 common/common.cpp | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 7520b703..70979382 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1443,11 +1443,17 @@ static bool assign_layers_to_device(
         }
 
         // check the solution
-        bool has_free_gpu_memory = false, has_gpu_overload = false, has_cpu_overload = false;
+        bool has_free_gpu_memory = false, has_gpu_overload = false, has_cpu_overload = false, has_weak_device = false;
         for (uint32_t m = 0; m < n_world; ++m) {
             // if (!dev_gpu[m]) continue;
             uint32_t w_m = best_solution[m], n_m = best_solution[m + n_world];
 
+            if (w_m == 1 && n_m == 0) {
+                // if the device is weak
+                has_weak_device = true;
+                LOG_INF("Device %d is weak, need to be removed: w_m = %d, n_m = %d\n", m, w_m, n_m);
+            }
+
             if (dev_gpu[m]) {
                 if (n_m < static_cast<uint32_t>(std::floor(W * vec_z_gpu[m]))) {
                     // if there is still free GPU memory
@@ -1467,7 +1473,7 @@ static bool assign_layers_to_device(
             }
         }
 
-        if (has_free_gpu_memory && (has_gpu_overload || has_cpu_overload)) {
+        if (!has_weak_device && has_free_gpu_memory && (has_gpu_overload || has_cpu_overload)) {
             int worst_device = -1;
             float worst_speed = std::numeric_limits<float>::max();
 

From 729870fcd7b2ecae6988107c0c2079f6c954458f Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Thu, 26 Jun 2025 14:47:34 +0400
Subject: [PATCH 25/42] topo rebuild: add a delay to avoid packet interleaving

---
 common/common.cpp   |  2 ++
 common/profiler.cpp |  8 ++++---
 src/llama.cpp       | 51 +++++++++++++++++++++++++++++----------------
 3 files changed, 40 insertions(+), 21 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 70979382..1aba9933 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1788,6 +1788,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
                     return iparams;
                 }
                 llama_bcast_layer_setup(lctx, n_layer_window, n_gpu_layers);
+                std::this_thread::sleep_for(std::chrono::milliseconds(200));  // add a delay to avoid packet interleaving
                 llama_rebuild_topo(lctx, n_layer_window, dev_info_set.data(), &node_type, is_forwarder);
             } else {
                 // use the user-defined n_layer_window
@@ -1798,6 +1799,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
             if (auto_schedule){
                 llama_send_device_info(lctx, &dev_info);
                 llama_recv_layer_setup(lctx, n_layer_window, n_gpu_layers);
+                std::this_thread::sleep_for(std::chrono::milliseconds(200));  // add a delay to avoid packet interleaving
                 llama_rebuild_topo    (lctx, n_layer_window, nullptr, &node_type, is_forwarder);
             } else {
                 llama_recv_layer_setup(lctx, n_layer_window, n_gpu_layers);
diff --git a/common/profiler.cpp b/common/profiler.cpp
index afc427f1..e9862bfb 100644
--- a/common/profiler.cpp
+++ b/common/profiler.cpp
@@ -2837,7 +2837,7 @@ size_t deserialize(const char * buffer, struct device_info * dev_info) {
     return ptr - buffer;
 }
 
-void TopoRebuildHelperInfo::deserialize(const char *buffer) {
+void TopoRebuildHelperInfo::deserialize(const char * buffer) {
     size_t buffer_size = ::deserialize(buffer, &dev_info);
     if (buffer_size == 0) {
         LOG_ERR("%s: failed to deserialize device info\n", __func__);
@@ -2846,14 +2846,16 @@ void TopoRebuildHelperInfo::deserialize(const char *buffer) {
     memcpy(&is_forwarder, buffer + buffer_size, 1);
 }
 
-size_t TopoRebuildHelperInfo::serialize(char **buffer) const{ 
+size_t TopoRebuildHelperInfo::serialize(char ** buffer) const{ 
     size_t buffer_size = ::serialize(&dev_info, buffer);
-    char* buffer_ = (char*)malloc(buffer_size+1);
+    char * buffer_ = (char *)malloc(buffer_size + 1);
+
     if (buffer_ == NULL) {
         LOG_ERR("%s: failed to allocate %zu bytes for device info serialization\n", 
                 __func__, buffer_size);
         return 0;
     }
+    
     memcpy(buffer_, *buffer, buffer_size);
     memcpy(buffer_ + buffer_size, &is_forwarder, 1);
     free(*buffer);
diff --git a/src/llama.cpp b/src/llama.cpp
index de548a66..b64fbd78 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -3633,6 +3633,10 @@ void llama_profile_device(
 
     dev_info->memory.total_physical     = round(device_physical_memory(false) / (double)(1 << 30) * 100) / 100;
     dev_info->memory.available_physical = round(device_physical_memory(true)  / (double)(1 << 30) * 100) / 100;
+
+    GGML_ASSERT(dev_info->memory.total_physical     > 0, "Failed to parse total physical memory\n");
+    GGML_ASSERT(dev_info->memory.available_physical > 0, "Failed to parse available physical memory\n");
+
     dev_info->memory.used_can_swap      = round(device_swappable_memory()     / (double)(1 << 30) * 100) / 100;
     dev_info->memory.total_swap         = round(device_swap_memory(false)     / (double)(1 << 30) * 100) / 100;
     dev_info->memory.available_swap     = round(device_swap_memory(true)      / (double)(1 << 30) * 100) / 100;
@@ -20658,8 +20662,8 @@ int llama_rebuild_topo(llama_context * ctx,
     uint32_t my_rank = ctx->cparams.rank;
     TopoRebuildHelperInfo* topo_helper = new TopoRebuildHelperInfo[n_world];
 
-    if (dev_info_set == nullptr){
-        // for rank!=0, recv all devices info
+    if (dev_info_set == nullptr) {
+        // for rank != 0, recv all devices info
         std::vector<zmq::message_t> msgs;
         if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) {
             return -1;
@@ -20671,7 +20675,7 @@ int llama_rebuild_topo(llama_context * ctx,
     } else {
         for (size_t i = 0; i < n_world; i++) {
             topo_helper[i].dev_info = dev_info_set[i];
-            topo_helper[i].is_forwarder =  0;
+            topo_helper[i].is_forwarder = 0;
         }
     }
 
@@ -20679,29 +20683,32 @@ int llama_rebuild_topo(llama_context * ctx,
 
     auto next_rank = (my_rank + 1) % n_world;
     auto next_connect_rank = (my_rank + 1) % n_world;
-    zmq::socket_t* socket_to_close = nullptr;
+    zmq::socket_t * socket_to_close = nullptr;
     bool is_not_exit = n_layer_window[my_rank] > 0 || topo_helper[my_rank].is_forwarder == 1;
+    
     if (is_not_exit) {
         // reconstruct socket to the next valid rank
         auto current_rank = my_rank;
         std::vector<uint32_t> nodes;
         auto next_rank_ = next_rank;
+
         while (next_rank_ != my_rank) {
             nodes.push_back(next_rank_);
             if (n_layer_window[next_rank_] > 0) {
                 break;
             }
-            next_rank_   = (next_rank_ + 1) % n_world;
+            next_rank_   = (next_rank_   + 1) % n_world;
             current_rank = (current_rank + 1) % n_world;
         }
+
         if (next_rank_ == my_rank) {
             // only one node
             ctx->next_node_ip = "";
-            socket_to_close = ctx->send_socket;
-            ctx->send_socket = nullptr;
+            socket_to_close   = ctx->send_socket;
+            ctx->send_socket  = nullptr;
         } else {
             // iterate node reverse
-            zmq::socket_t* socket = nullptr;
+            zmq::socket_t * socket = nullptr;
             std::string ip;
             for (int i = nodes.size() - 1; i > 0; --i) {
                 auto rank = nodes[i];
@@ -20716,13 +20723,13 @@ int llama_rebuild_topo(llama_context * ctx,
                 // reset socket
                 GGML_ASSERT(socket != nullptr);
                 GGML_ASSERT(!ip.empty());
-                socket_to_close = ctx->send_socket;
-                ctx->send_socket = socket;
+                socket_to_close   = ctx->send_socket;
+                ctx->send_socket  = socket;
                 ctx->next_node_ip = ip;
                 ctx->cparams.original_next_rank = next_connect_rank;
             }
         }
-    }else if (n_layer_window[next_rank] <= 0 && topo_helper[next_rank].is_forwarder == 0) {
+    } else if (n_layer_window[next_rank] <= 0 && topo_helper[next_rank].is_forwarder == 0) {
         socket_to_close = ctx->send_socket;
     }
 
@@ -20733,7 +20740,7 @@ int llama_rebuild_topo(llama_context * ctx,
             auto msgs = topohelper_to_messages(topo_helper, n_world);
             socket_to_close->set(zmq::sockopt::linger, 3500);
             zmq::send_multipart(*socket_to_close, msgs);
-        } catch (const zmq::error_t& e) {
+        } catch (const zmq::error_t & e) {
             LLAMA_LOG_INFO("Failed to send data: %s\n", e.what());
             return -1;
         }
@@ -20745,7 +20752,7 @@ int llama_rebuild_topo(llama_context * ctx,
         try {
             auto msgs = topohelper_to_messages(topo_helper, n_world);
             zmq::send_multipart(*ctx->send_socket, msgs);
-        } catch (const zmq::error_t& e) {
+        } catch (const zmq::error_t & e) {
             LLAMA_LOG_INFO("Failed to send data: %s\n", e.what());
             return -1;
         }
@@ -20770,18 +20777,20 @@ int llama_rebuild_topo(llama_context * ctx,
             topo_helper[i].deserialize((char *)msgs[i].data());
         }
         // broadcast the whole view
-        if (next_connect_rank!=0) {
+        if (next_connect_rank != 0) {
             try {
                 zmq::send_multipart(*ctx->send_socket, msgs);
-            } catch (const zmq::error_t& e) {
+            } catch (const zmq::error_t & e) {
                 LLAMA_LOG_INFO("Failed to send data: %s\n", e.what());
                 return -1;
             }
         }
     }
+
     for (size_t i = 0; i < n_world; i++) {
         is_forwarder[i] = topo_helper[i].is_forwarder;
     }
+    
     ctx->cparams.node_type = *node_type;
 
     if (socket_to_close != nullptr) {
@@ -20816,8 +20825,14 @@ int llama_recv_layer_setup(struct llama_context * ctx, uint32_t * n_layer_window
     uint32_t my_rank = ctx->cparams.rank;
 
     std::vector<zmq::message_t> recv_msgs;
-    if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(recv_msgs))) {
-        return -1;
+    while (true) {
+        recv_msgs.clear();
+        if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(recv_msgs))) {
+            return -1;
+        }
+        if (!recv_msgs.empty() && recv_msgs[0].to_string() == "n_layer_window") {
+            break;
+        }
     }
 
     GGML_ASSERT(recv_msgs[0].to_string() == "n_layer_window");
@@ -20827,7 +20842,7 @@ int llama_recv_layer_setup(struct llama_context * ctx, uint32_t * n_layer_window
     if (recv_msgs.size() > 2) {
         GGML_ASSERT(recv_msgs[2].to_string() == "n_gpu_layers");
         GGML_ASSERT(recv_msgs[3].size() == sizeof(uint32_t) * 32);
-        memcpy(n_gpu_layers,   recv_msgs[3].data(), sizeof(uint32_t) * 32);
+        memcpy(n_gpu_layers, recv_msgs[3].data(), sizeof(uint32_t) * 32);
     }
 
     if (my_rank != n_world - 1) {

From 3f27a253403505059b012883b11554424d54b1fa Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Thu, 26 Jun 2025 14:50:58 +0400
Subject: [PATCH 26/42] topo rebuild: add a delay to avoid packet interleaving

---
 src/llama.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/llama.cpp b/src/llama.cpp
index b64fbd78..7b07a291 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -3634,8 +3634,8 @@ void llama_profile_device(
     dev_info->memory.total_physical     = round(device_physical_memory(false) / (double)(1 << 30) * 100) / 100;
     dev_info->memory.available_physical = round(device_physical_memory(true)  / (double)(1 << 30) * 100) / 100;
 
-    GGML_ASSERT(dev_info->memory.total_physical     > 0, "Failed to parse total physical memory\n");
-    GGML_ASSERT(dev_info->memory.available_physical > 0, "Failed to parse available physical memory\n");
+    GGML_ASSERT(dev_info->memory.total_physical     > 0 && "Failed to parse total physical memory\n");
+    GGML_ASSERT(dev_info->memory.available_physical > 0 && "Failed to parse available physical memory\n");
 
     dev_info->memory.used_can_swap      = round(device_swappable_memory()     / (double)(1 << 30) * 100) / 100;
     dev_info->memory.total_swap         = round(device_swap_memory(false)     / (double)(1 << 30) * 100) / 100;
@@ -20790,7 +20790,7 @@ int llama_rebuild_topo(llama_context * ctx,
     for (size_t i = 0; i < n_world; i++) {
         is_forwarder[i] = topo_helper[i].is_forwarder;
     }
-    
+
     ctx->cparams.node_type = *node_type;
 
     if (socket_to_close != nullptr) {

From a05022c05ae836d599c1835a1ab84f73c3bbb5e2 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Thu, 26 Jun 2025 17:30:47 +0400
Subject: [PATCH 27/42] communication: use barrier instead of manually adding
 delay

---
 common/common.cpp |  22 +++----
 src/llama.cpp     | 159 +++++++++++++++++++++++++---------------------
 2 files changed, 95 insertions(+), 86 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 1aba9933..f5b4cffb 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1778,7 +1778,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
                 dev_info_set[0] = dev_info;
 
                 llama_gather_device_info(lctx, dev_info_set.data());
-                device_print_props(dev_info_set.data(), n_world, model, cparams);
+                device_print_props      (dev_info_set.data(), n_world, model, cparams);
 
                 // assign layers to devices and remove weak devices
                 if (!assign_layers_and_select_devices(n_world, dev_info_set, n_layer_window, n_gpu_layers, model, cparams)) {
@@ -1788,8 +1788,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
                     return iparams;
                 }
                 llama_bcast_layer_setup(lctx, n_layer_window, n_gpu_layers);
-                std::this_thread::sleep_for(std::chrono::milliseconds(200));  // add a delay to avoid packet interleaving
-                llama_rebuild_topo(lctx, n_layer_window, dev_info_set.data(), &node_type, is_forwarder);
+                llama_rebuild_topo     (lctx, n_layer_window, dev_info_set.data(), &node_type, is_forwarder);
             } else {
                 // use the user-defined n_layer_window
                 std::copy(std::begin(params.n_layer_window), std::end(params.n_layer_window), n_layer_window);
@@ -1797,12 +1796,11 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
             }
         } else {
             if (auto_schedule){
-                llama_send_device_info(lctx, &dev_info);
-                llama_recv_layer_setup(lctx, n_layer_window, n_gpu_layers);
-                std::this_thread::sleep_for(std::chrono::milliseconds(200));  // add a delay to avoid packet interleaving
-                llama_rebuild_topo    (lctx, n_layer_window, nullptr, &node_type, is_forwarder);
+                llama_send_device_info (lctx, &dev_info);
+                llama_recv_layer_setup (lctx, n_layer_window, n_gpu_layers);
+                llama_rebuild_topo     (lctx, n_layer_window, nullptr, &node_type, is_forwarder);
             } else {
-                llama_recv_layer_setup(lctx, n_layer_window, n_gpu_layers);
+                llama_recv_layer_setup (lctx, n_layer_window, n_gpu_layers);
             }
         }
 
@@ -1823,17 +1821,13 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
             if (n_layer_window[i] <= 0 && is_forwarder[i] == 0) {
                 continue;
             }
-            if (i <= my_rank) {
-                update_rank++;
-            }
+            if (i <= my_rank) update_rank++;
             update_n_world++;
             n_layer_window_temp.push_back(n_layer_window[i]);
             n_gpu_layers_temp.push_back(n_gpu_layers[i]);
 
             if (n_layer_window[i] > 0) {
-                if (i <= my_rank) {
-                    worker_rank++;
-                }
+                if (i <= my_rank) worker_rank++;
                 n_worker++;
             }
         }
diff --git a/src/llama.cpp b/src/llama.cpp
index 7b07a291..b70f338b 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -20518,12 +20518,11 @@ void llama_init_sockets(struct llama_context * ctx, uint32_t n_world, uint32_t m
 
 int llama_gather_device_info(struct llama_context * ctx, struct device_info * dev_info_set) {
     uint32_t n_world = ctx->cparams.n_world;
-    if (n_world == 1) {
-        return 0;
-    }
-
+    if (n_world == 1) return 0;
     GGML_ASSERT(dev_info_set != nullptr);
     GGML_ASSERT(ctx != nullptr && ctx->send_socket != nullptr);
+
+    // master rank sends its device info to its next rank
     try {
         char * buffer = nullptr;
         size_t buffer_size = serialize(&dev_info_set[0], &buffer);
@@ -20538,6 +20537,7 @@ int llama_gather_device_info(struct llama_context * ctx, struct device_info * de
         return -1;
     }
 
+    // master rank receives aggregated device info from its previous rank (a barrier op)
     std::vector<zmq::message_t> recv_msgs;
     if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(recv_msgs))) {
         return -1;
@@ -20551,6 +20551,7 @@ int llama_gather_device_info(struct llama_context * ctx, struct device_info * de
 }
 
 int llama_send_device_info(struct llama_context * ctx, struct device_info * dev_info) {
+    // non-master ranks receive device info from their previous rank
     std::vector<zmq::message_t> recv_msgs;
     if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(recv_msgs))) {
         return -1;
@@ -20559,6 +20560,7 @@ int llama_send_device_info(struct llama_context * ctx, struct device_info * dev_
     GGML_ASSERT(dev_info != nullptr);
     GGML_ASSERT(ctx != nullptr && ctx->send_socket != nullptr);
 
+    // non-master ranks aggregate their device info and forward it to their next rank
     try {
         char * buffer = nullptr;
         size_t buffer_size = serialize(dev_info, &buffer);
@@ -20577,79 +20579,95 @@ int llama_send_device_info(struct llama_context * ctx, struct device_info * dev_
 
 int llama_bcast_startup_args(llama_context * ctx, uint32_t rank, startup_args * args) {
     int32_t n_world = ctx->cparams.n_world;
-    GGML_ASSERT(n_world > 0);
+    if (n_world == 1) return 0;
     GGML_ASSERT(ctx != nullptr && ctx->send_socket != nullptr);
 
+    std::vector<zmq::message_t> msgs;
+
     if (rank == 0){
-        // send
+        // master rank sends its startup args to its next rank
         try {
-            std::vector<zmq::message_t> send_msgs;
-
-            send_msgs.emplace_back("should_profile", strlen("should_profile"));
-            send_msgs.emplace_back(&args->should_profile, sizeof(args->should_profile));
-
-            send_msgs.emplace_back("n_ctx", strlen("n_ctx"));
-            send_msgs.emplace_back(&args->n_ctx, sizeof(args->n_ctx));
-
-            zmq::send_multipart(*ctx->send_socket, send_msgs);
+            msgs.emplace_back("should_profile", strlen("should_profile"));
+            msgs.emplace_back(&args->should_profile, sizeof(args->should_profile));
+            msgs.emplace_back("n_ctx", strlen("n_ctx"));
+            msgs.emplace_back(&args->n_ctx, sizeof(args->n_ctx));
+            zmq::send_multipart(*ctx->send_socket, msgs);
         } catch (const zmq::error_t& e) {
             LLAMA_LOG_INFO("Failed to send data: %s\n", e.what());
             return -1;
         }
-    } else {
-        // receive
-        std::vector<zmq::message_t> recv_msgs;
-        if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(recv_msgs))) {
+
+        // master rank receives ack from its previous rank (a barrier op)
+        msgs.clear();
+        if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) {
             return -1;
         }
 
-        GGML_ASSERT(recv_msgs[0].to_string() == "should_profile");
-        GGML_ASSERT(recv_msgs[1].size() == sizeof(bool));
-        bool should_profile = *static_cast<bool*>(recv_msgs[1].data());
-        args->should_profile = should_profile;
+        GGML_ASSERT(msgs.size() == 1);
+        GGML_ASSERT(msgs[0].to_string() == "ACK");
+    } else {
+        // non-master ranks receive startup args from their previous rank
+        if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) {
+            return -1;
+        }
 
-        GGML_ASSERT(recv_msgs[2].to_string() == "n_ctx");
-        GGML_ASSERT(recv_msgs[3].size() == sizeof(uint32_t));
-        uint32_t n_ctx = *static_cast<uint32_t*>(recv_msgs[3].data());        
-        args->n_ctx = n_ctx;
+        GGML_ASSERT(msgs[0].to_string() == "should_profile");
+        GGML_ASSERT(msgs[1].size() == sizeof(bool));
+        GGML_ASSERT(msgs[2].to_string() == "n_ctx");
+        GGML_ASSERT(msgs[3].size() == sizeof(uint32_t));
 
-        if ((int)rank != (int)n_world - 1){
-            // send
-            try {
-                zmq::send_multipart(*ctx->send_socket, recv_msgs);
-            } catch (const zmq::error_t & e) {
-                LLAMA_LOG_INFO("Failed to send data: %s\n", e.what());
-                return -1;
+        args->should_profile = *static_cast<bool*>(msgs[1].data());      
+        args->n_ctx = *static_cast<uint32_t*>(msgs[3].data());
+
+        // non-master ranks forward the startup args to their next rank
+        try {
+            // last rank just sends an ack to pass the barrier
+            if ((int)rank == (int)n_world - 1) {
+                msgs.clear();
+                msgs.emplace_back("ACK", strlen("ACK"));
             }
+            zmq::send_multipart(*ctx->send_socket, msgs);
+        } catch (const zmq::error_t & e) {
+            LLAMA_LOG_INFO("Failed to send data: %s\n", e.what());
+            return -1;
         }
     }
+
     return 0;
 }
 
 int llama_bcast_layer_setup(struct llama_context * ctx, uint32_t * n_layer_window, uint32_t * n_gpu_layers) {
     uint32_t n_world = ctx->cparams.n_world;
-    if (n_world == 1) {
-        return 0;
-    }
-
+    if (n_world == 1) return 0;
     GGML_ASSERT(ctx != nullptr && ctx->send_socket != nullptr);
-    try {
-        std::vector<zmq::message_t> send_msgs;
 
-        send_msgs.emplace_back("n_layer_window", strlen("n_layer_window"));
-        send_msgs.emplace_back(n_layer_window, sizeof(uint32_t) * 32);
+    std::vector<zmq::message_t> msgs;
+
+    // master rank sends its layer setup to its next rank
+    try {
+        msgs.emplace_back("n_layer_window", strlen("n_layer_window"));
+        msgs.emplace_back(n_layer_window, sizeof(uint32_t) * 32);
 
         if (n_gpu_layers != nullptr) {
-            send_msgs.emplace_back("n_gpu_layers", strlen("n_gpu_layers"));
-            send_msgs.emplace_back(n_gpu_layers, sizeof(uint32_t) * 32);
+            msgs.emplace_back("n_gpu_layers", strlen("n_gpu_layers"));
+            msgs.emplace_back(n_gpu_layers, sizeof(uint32_t) * 32);
         }
 
-        zmq::send_multipart(*ctx->send_socket, send_msgs);
-    } catch (const zmq::error_t& e) {
+        zmq::send_multipart(*ctx->send_socket, msgs);
+    } catch (const zmq::error_t & e) {
         LLAMA_LOG_INFO("Failed to send data: %s\n", e.what());
         return -1;
     }
 
+    // master rank receives ack from its previous rank (a barrier op)
+    msgs.clear();
+    if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) {
+        return -1;
+    }
+
+    GGML_ASSERT(msgs.size() == 1);
+    GGML_ASSERT(msgs[0].to_string() == "ACK");
+
     return 0;
 }
 
@@ -20663,7 +20681,7 @@ int llama_rebuild_topo(llama_context * ctx,
     TopoRebuildHelperInfo* topo_helper = new TopoRebuildHelperInfo[n_world];
 
     if (dev_info_set == nullptr) {
-        // for rank != 0, recv all devices info
+        // for non-master ranks, receive devices info from their previous rank
         std::vector<zmq::message_t> msgs;
         if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) {
             return -1;
@@ -20824,34 +20842,33 @@ int llama_recv_layer_setup(struct llama_context * ctx, uint32_t * n_layer_window
     uint32_t n_world = ctx->cparams.n_world;
     uint32_t my_rank = ctx->cparams.rank;
 
-    std::vector<zmq::message_t> recv_msgs;
-    while (true) {
-        recv_msgs.clear();
-        if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(recv_msgs))) {
-            return -1;
-        }
-        if (!recv_msgs.empty() && recv_msgs[0].to_string() == "n_layer_window") {
-            break;
-        }
+    // non-master ranks receive data from their previous rank
+    std::vector<zmq::message_t> msgs;
+    if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) {
+        return -1;
     }
 
-    GGML_ASSERT(recv_msgs[0].to_string() == "n_layer_window");
-    GGML_ASSERT(recv_msgs[1].size() == sizeof(uint32_t) * 32);
-    memcpy(n_layer_window, recv_msgs[1].data(), sizeof(uint32_t) * 32);
+    GGML_ASSERT(msgs[0].to_string() == "n_layer_window");
+    GGML_ASSERT(msgs[1].size() == sizeof(uint32_t) * 32);
+    memcpy(n_layer_window, msgs[1].data(), sizeof(uint32_t) * 32);
 
-    if (recv_msgs.size() > 2) {
-        GGML_ASSERT(recv_msgs[2].to_string() == "n_gpu_layers");
-        GGML_ASSERT(recv_msgs[3].size() == sizeof(uint32_t) * 32);
-        memcpy(n_gpu_layers, recv_msgs[3].data(), sizeof(uint32_t) * 32);
+    if (msgs.size() > 2) {
+        GGML_ASSERT(msgs[2].to_string() == "n_gpu_layers");
+        GGML_ASSERT(msgs[3].size() == sizeof(uint32_t) * 32);
+        memcpy(n_gpu_layers, msgs[3].data(), sizeof(uint32_t) * 32);
     }
 
-    if (my_rank != n_world - 1) {
-        try {
-            zmq::send_multipart(*ctx->send_socket, recv_msgs);
-        } catch (const zmq::error_t& e) {
-            LLAMA_LOG_INFO("Failed to send data: %s\n", e.what());
-            return -1;
+    // non-master ranks forward the received message to their next rank
+    try {
+        if (my_rank == n_world - 1) {
+            // last rank just sends an ack to pass the barrier
+            msgs.clear();
+            msgs.emplace_back("ACK", strlen("ACK"));
         }
+        zmq::send_multipart(*ctx->send_socket, msgs);
+    } catch (const zmq::error_t& e) {
+        LLAMA_LOG_INFO("Failed to send data: %s\n", e.what());
+        return -1;
     }
     
     return 0;
@@ -20863,9 +20880,7 @@ void llama_free_sockets(struct llama_context * ctx, char ** msg) {
     // to adapt to the new topology, use old next_rank
     const uint32_t next_rank = ctx->cparams.original_next_rank;
 
-    if (n_world == 1) {
-        return;
-    }
+    if (n_world == 1) return;
 
     zmq::socket_t signal_sender(*ctx->sock_context, zmq::socket_type::push);
     std::string endp = "tcp://" + ctx->next_node_ip + ":" + std::to_string(map_rank_to_port(next_rank, ctx->signal_port));

From aacfa8a231f36014c329ff799c02033460edbd96 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Thu, 26 Jun 2025 20:45:45 +0400
Subject: [PATCH 28/42] fix compute buffer estimate: reserve 300 MiB VRAM to
 avoid potential OOM

---
 src/llama.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/llama.cpp b/src/llama.cpp
index b70f338b..abda2327 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -22043,8 +22043,8 @@ void llama_model_compute_buf_size(
         // context GPU memory usage, i.e. the initial memory cost of creating a CUDA context, 
         // even before you launch any kernels or allocate your own buffers.
         // this value may vary by GPU and CUDA version, but it's lower than 400 MiB in most cases,
-        // another 200 MiB is used to prevent accidental OOM.
-        *gpu_buf += 600 * 1024 * 1024;
+        // another 300 MiB is used to prevent accidental OOM.
+        *gpu_buf += 700 * 1024 * 1024;
     }
 }
 

From ba59a1a07a22efb00aa9b2bf387602812d73633c Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Thu, 26 Jun 2025 22:33:28 +0400
Subject: [PATCH 29/42] update README

---
 README.md | 32 +++++++++++++++-----------------
 1 file changed, 15 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index 300ffa6a..77cfd12d 100644
--- a/README.md
+++ b/README.md
@@ -123,7 +123,7 @@ Before using this project, ensure you have the following dependencies installed:
 
 ```shell
 # Use apt in Linux and pkg in Termux
-sudo apt update -y && sudo apt install -y gcc-9 make cmake fio git wget libzmq3-dev
+sudo apt update -y && sudo apt install -y gcc-9 make cmake fio git wget libzmq3-dev curl
 ```
 
 For HiGHS, download and install from [source](https://github.com/ERGO-Code/HiGHS):
@@ -141,7 +141,7 @@ sudo ldconfig
 **macOS:**
 
 ```shell
-brew install gcc make cmake fio git wget highs zeromq
+brew install gcc make cmake fio git wget highs zeromq curl
 ```
 
 ### Build, Download, and Test
@@ -205,7 +205,7 @@ graph LR;
 
 > **NOTE:** This ring communication is a communication overlay, not the physical topology. These devices are physically fully connected because they all connect to the same Wi-Fi.
 
-> If possible, disable the firewall to prevent the ports needed (e.g., 9000, 10000) been blocked.
+> If possible, disable the firewall to prevent the ports needed (e.g., 9000, 10000) been blocked, or you can use `--data-port` (9000, by default) and `--signal-port` (10000, by default) to customize the ports used.
 
 
 Take QwQ-32B as an example, run the following commands on the devices to launch distributed inference:
@@ -229,7 +229,7 @@ Once started, prima.cpp will profile each device and decide how much workload to
 ### (Optional) Run with Prebuilt Docker Image
 Assume we have a host machine with at least 32 CPU cores, 32 GiB RAM, and 32 GiB VRAM. We simulate 4 homogeneous nodes using Docker containers, with each node allocated 8 CPU cores, 8 GiB RAM, and 8 GiB VRAM. Follow the below steps to get started:
 
-1. Pull our prebuilt Docker image (e.g., [`prima.cpp:1.0.1-cuda`](https://hub.docker.com/repository/docker/lizonghango00o1/prima.cpp/general)) and run 4 containers:
+1. Pull our prebuilt Docker image (e.g., [`prima.cpp:1.0.2-cuda`](https://hub.docker.com/repository/docker/lizonghango00o1/prima.cpp/general)) and run 4 containers:
 
 ```shell
 sudo docker run -dit --name prima-v1 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="0-7"   --network host --gpus all prima.cpp:1.0.1-cuda
@@ -238,9 +238,7 @@ sudo docker run -dit --name prima-v3 --memory=8gb --memory-swap=8gb --cpus 8 --c
 sudo docker run -dit --name prima-v4 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="24-31" --network host --gpus all prima.cpp:1.0.1-cuda
 ```
 
-> If your host machine does not have a GPU, ignore the `--gpus all` option.
-
-2. Download the model file [`qwq-32b-q4_k_m.gguf`](https://huggingface.co/Qwen/QwQ-32B-GGUF) and copy it into each container:
+1. Download the model file [`qwq-32b-q4_k_m.gguf`](https://huggingface.co/Qwen/QwQ-32B-GGUF) and copy it into each container:
 
 ```shell
 cd prima.cpp/download
@@ -250,27 +248,27 @@ sudo docker cp qwq-32b-q4_k_m.gguf prima-v3:/root/prima.cpp/download/
 sudo docker cp qwq-32b-q4_k_m.gguf prima-v4:/root/prima.cpp/download/
 ```
 
-3. (Optional) Enter each container, rebuild prima.cpp if your host machine does not have a GPU:
+1. Enter each container and build prima.cpp:
 
 ```shell
-cd /root/prima.cpp && make clean
-make -j$(nproc)  # If not rank 0
-make USE_HIGHS=1 -j$(nproc)  # If rank 0
+cd /root/prima.cpp
+make GGML_CUDA=1 USE_HIGHS=1 -j$(nproc)  # For rank 0
+make GGML_CUDA=1 -j$(nproc)  # For other ranks
 ``` 
 
 4. Enter each container and launch the distributed inference:
 
 ```shell
 cd /root/prima.cpp
-(prima-v1) ./llama-cli -m download/qwq-32b-q4_k_m.gguf -c 1024 -n 256 -p "what is edge AI?" --world 4 --rank 0 --prefetch --gpu-mem 8
-(prima-v2) ./llama-cli -m download/qwq-32b-q4_k_m.gguf -c 1024 --world 4 --rank 1 --prefetch --gpu-mem 8
-(prima-v3) ./llama-cli -m download/qwq-32b-q4_k_m.gguf -c 1024 --world 4 --rank 2 --prefetch --gpu-mem 8
-(prima-v4) ./llama-cli -m download/qwq-32b-q4_k_m.gguf -c 1024 --world 4 --rank 3 --prefetch --gpu-mem 8
+(prima-v1) ./llama-cli -m download/qwq-32b-q4_k_m.gguf --world 4 --rank 0 --prefetch --gpu-mem 8 -c 4096 -n 256 -p "what is edge AI?"
+(prima-v2) ./llama-cli -m download/qwq-32b-q4_k_m.gguf --world 4 --rank 1 --prefetch --gpu-mem 8
+(prima-v3) ./llama-cli -m download/qwq-32b-q4_k_m.gguf --world 4 --rank 2 --prefetch --gpu-mem 8
+(prima-v4) ./llama-cli -m download/qwq-32b-q4_k_m.gguf --world 4 --rank 3 --prefetch --gpu-mem 8
 ``` 
 
-> If your host machine does not have a GPU, ignore the `--gpu-mem` option.
+> You can ignore `--gpu-mem` if you don't want to limit VRAM usage.
 
-> If you update to the latest code, non-rank 0 nodes can omit `-c 1024`.
+> Always use `git fetch` to update the local repository.
 
 ### Run in Server Mode
 You can run prima.cpp in server mode, by launching `llama-server` on the rank 0 device (with `--host` and `--port` specified) and `llama-cli` on the others. Here is an example with 2 devices:

From 3a03549fed89ee0b1833e207c9d51c6c6b04240f Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Thu, 26 Jun 2025 22:37:08 +0400
Subject: [PATCH 30/42] update README

---
 README.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 77cfd12d..d312aa11 100644
--- a/README.md
+++ b/README.md
@@ -238,7 +238,7 @@ sudo docker run -dit --name prima-v3 --memory=8gb --memory-swap=8gb --cpus 8 --c
 sudo docker run -dit --name prima-v4 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="24-31" --network host --gpus all prima.cpp:1.0.1-cuda
 ```
 
-1. Download the model file [`qwq-32b-q4_k_m.gguf`](https://huggingface.co/Qwen/QwQ-32B-GGUF) and copy it into each container:
+2. Download the model file [`qwq-32b-q4_k_m.gguf`](https://huggingface.co/Qwen/QwQ-32B-GGUF) and copy it into each container:
 
 ```shell
 cd prima.cpp/download
@@ -248,7 +248,7 @@ sudo docker cp qwq-32b-q4_k_m.gguf prima-v3:/root/prima.cpp/download/
 sudo docker cp qwq-32b-q4_k_m.gguf prima-v4:/root/prima.cpp/download/
 ```
 
-1. Enter each container and build prima.cpp:
+3. Enter each container and build prima.cpp:
 
 ```shell
 cd /root/prima.cpp

From 11ce0d58f7dbe1f3fa328fec9a7c68a810464e2d Mon Sep 17 00:00:00 2001
From: Zonghang Li <zonghang.li@mbzuai.ac.ae>
Date: Fri, 27 Jun 2025 12:42:16 +0000
Subject: [PATCH 31/42] fix compute buffer estimate: don't reverse CUDA VRAM
 for output layer

---
 src/llama.cpp | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/llama.cpp b/src/llama.cpp
index abda2327..01cf82a2 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -20976,7 +20976,7 @@ void * llama_context_setup_backend(
     auto       & cparams = ctx->cparams;
 
     std::copy(std::begin(params.n_layer_window), std::end(params.n_layer_window), cparams.n_layer_window);
-    cparams.prefetch           = params.prefetch;
+    cparams.prefetch         = params.prefetch;
     cparams.n_seq_max        = std::max(1u, params.n_seq_max);
     cparams.n_threads        = params.n_threads;
     cparams.n_threads_batch  = params.n_threads_batch;
@@ -21359,6 +21359,16 @@ void * llama_context_setup_backend(
             bool ok = true;
             GGML_ASSERT(ctx->sched.size() == gf.size());
             for (size_t i = 0; i < gf.size(); ++i) {
+
+#if defined(GGML_USE_CUDA)
+                if ((cparams.rank == 0 && (i == 0 || i == gf.size() - 1)) 
+                        || model->n_gpu_layers == 0) {
+                    continue;
+                }
+#elif defined(GGML_USE_METAL)
+
+#endif
+
                 ok = ok & ggml_backend_sched_reserve(ctx->sched[i], gf[i]);
             }
             if (!ok) {
@@ -21963,7 +21973,8 @@ void llama_model_compute_buf_size(
                     (n_inp_pos + n_kq_mask) * type_size_f32 + n_v * type_size_f16 + nb_attn_norm_w,
                 });
             }
-            *gpu_buf     += (n_out_embd + n_result) * type_size_f32 + nb_output_w;
+            // we run the output layer on CPU by default
+            // *gpu_buf     += (n_out_embd + n_result) * type_size_f32 + nb_output_w;
             gpu_host_buf  = (n_inp_toks + n_inp_embd + n_bak_embd + n_inp_pos + n_kq_mask + n_inp_out_ids + n_out_embd) * type_size_f32;
         } else {
             if (has_gpu_layers) {
@@ -22036,7 +22047,7 @@ void llama_model_compute_buf_size(
 
     LLAMA_LOG_INFO("\n");
     LLAMA_LOG_INFO("%s: here the compute buffer size is a predicted upper bound, not an exact value\n", __func__);
-    LLAMA_LOG_INFO("%s: (rank %d) compute buffer size = %7.2f MiB (GPU) + %7.2f MiB (CPU / GPU-host buffer)\n", __func__,
+    LLAMA_LOG_INFO("%s: (rank %d) compute buffer size = %7.2f MiB (GPU) + %7.2f MiB (CPU & GPU-host buffer)\n", __func__,
             my_rank, *gpu_buf / (1024.0 * 1024.0), *cpu_buf / (1024.0 * 1024.0));
     
     if (backend == BACKEND_CUDA) {

From e8d3e5a6314646ae1bb4b7ff9bcae27d0d8e2214 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Fri, 27 Jun 2025 20:16:30 +0400
Subject: [PATCH 32/42] update README

---
 README.md     | 8 ++++----
 src/llama.cpp | 6 ++----
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index d312aa11..35191f3e 100644
--- a/README.md
+++ b/README.md
@@ -232,10 +232,10 @@ Assume we have a host machine with at least 32 CPU cores, 32 GiB RAM, and 32 GiB
 1. Pull our prebuilt Docker image (e.g., [`prima.cpp:1.0.2-cuda`](https://hub.docker.com/repository/docker/lizonghango00o1/prima.cpp/general)) and run 4 containers:
 
 ```shell
-sudo docker run -dit --name prima-v1 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="0-7"   --network host --gpus all prima.cpp:1.0.1-cuda
-sudo docker run -dit --name prima-v2 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="8-15"  --network host --gpus all prima.cpp:1.0.1-cuda
-sudo docker run -dit --name prima-v3 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="16-23" --network host --gpus all prima.cpp:1.0.1-cuda
-sudo docker run -dit --name prima-v4 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="24-31" --network host --gpus all prima.cpp:1.0.1-cuda
+sudo docker run -dit --name prima-v1 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="0-7"   --network host --gpus all prima.cpp:1.0.2-cuda
+sudo docker run -dit --name prima-v2 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="8-15"  --network host --gpus all prima.cpp:1.0.2-cuda
+sudo docker run -dit --name prima-v3 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="16-23" --network host --gpus all prima.cpp:1.0.2-cuda
+sudo docker run -dit --name prima-v4 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="24-31" --network host --gpus all prima.cpp:1.0.2-cuda
 ```
 
 2. Download the model file [`qwq-32b-q4_k_m.gguf`](https://huggingface.co/Qwen/QwQ-32B-GGUF) and copy it into each container:
diff --git a/src/llama.cpp b/src/llama.cpp
index 01cf82a2..871daac7 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -21365,8 +21365,6 @@ void * llama_context_setup_backend(
                         || model->n_gpu_layers == 0) {
                     continue;
                 }
-#elif defined(GGML_USE_METAL)
-
 #endif
 
                 ok = ok & ggml_backend_sched_reserve(ctx->sched[i], gf[i]);
@@ -21933,10 +21931,10 @@ void llama_model_compute_buf_size(
     const int64_t n_result   = hparams.n_vocab * cparams.n_ubatch;
 
     // weights
-    const int64_t nb_output_w    = n_bytes.nb_output_w;
     const int64_t nb_attn_norm_w = n_bytes.nb_attn_norm_w;
     const int64_t nb_attn_q_w    = n_bytes.nb_attn_q_w;
-
+    // const int64_t nb_output_w    = n_bytes.nb_output_w;
+    
     // format bytes
     const int64_t type_size_f32 = ggml_type_size(GGML_TYPE_F32);
     const int64_t type_size_f16 = ggml_type_size(GGML_TYPE_F16);

From 1ea2d61a97b3551df2266176f3a0dc9ef773d859 Mon Sep 17 00:00:00 2001
From: Zonghang Li <zonghang.li@mbzuai.ac.ae>
Date: Sat, 28 Jun 2025 05:59:19 +0000
Subject: [PATCH 33/42] speedup: add arg --keep-out-in-cuda to run the output
 layer on CUDA

---
 README.md         |  2 ++
 common/arg.cpp    | 12 ++++++++++++
 common/common.cpp | 22 +++++++++++++---------
 common/common.h   |  1 +
 include/llama.h   |  2 ++
 src/llama.cpp     | 43 ++++++++++++++++++++++++++++++++++++-------
 6 files changed, 66 insertions(+), 16 deletions(-)

diff --git a/README.md b/README.md
index 35191f3e..a1a18538 100644
--- a/README.md
+++ b/README.md
@@ -226,6 +226,8 @@ Take QwQ-32B as an example, run the following commands on the devices to launch
 
 Once started, prima.cpp will profile each device and decide how much workload to assign, e.g., how many model layers each device should handle, and how many of them should run on GPU (if available).
 
+> By default, the output layer runs on the CPU. However, if you have enough total VRAM, add `--keep-out-in-cuda` to the master to run it on the GPU.
+
 ### (Optional) Run with Prebuilt Docker Image
 Assume we have a host machine with at least 32 CPU cores, 32 GiB RAM, and 32 GiB VRAM. We simulate 4 homogeneous nodes using Docker containers, with each node allocated 8 CPU cores, 8 GiB RAM, and 8 GiB VRAM. Follow the below steps to get started:
 
diff --git a/common/arg.cpp b/common/arg.cpp
index 4039b19c..c971bb49 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -775,6 +775,7 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
             params.master_priority = std::stof(value);
         }
     ).set_env("LLAMA_ARG_MASTER_PRIORITY"));
+
 // #ifdef GGML_USE_METAL
 //     // warn: if the output layer weights are not kept in metal shared memory, its mmap-ed weight data
 //     // could be released by the OS and reloaded repeatedly, which causes additional disk I/O latency.
@@ -787,6 +788,17 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
 //         }
 //     ).set_env("LLAMA_ARG_KEEP_INP_OUT_IN_METAL"));
 // #endif
+
+#ifdef GGML_USE_CUDA
+    add_opt(llama_arg(
+        {"--keep-out-in-cuda"},
+        format("whether to compute the output layer on CUDA (default: %s)", params.keep_out_in_cuda ? "true" : "false"),
+        [](gpt_params & params) {
+            params.keep_out_in_cuda = true;
+        }
+    ).set_env("LLAMA_ARG_KEEP_INP_OUT_IN_CUDA"));
+#endif
+
     add_opt(llama_arg(
         {"-n", "--predict", "--n-predict"}, "N",
         format("number of tokens to predict (default: %d, -1 = infinity, -2 = until context filled)", params.n_predict),
diff --git a/common/common.cpp b/common/common.cpp
index f5b4cffb..92c95b3b 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -2017,16 +2017,19 @@ struct llama_model_params llama_model_params_from_gpt_params(const gpt_params &
     if (params.n_gpu_layers != -1) {
         mparams.n_gpu_layers = params.n_gpu_layers;
     }
-    mparams.n_world         = params.n_world;
-    mparams.rank            = params.rank;
-    mparams.rpc_servers     = params.rpc_servers.c_str();
-    mparams.main_gpu        = params.main_gpu;
-    mparams.split_mode      = params.split_mode;
-    mparams.tensor_split    = params.tensor_split;
-    mparams.use_mmap        = params.use_mmap;
-    mparams.use_mlock       = params.use_mlock;
-    mparams.check_tensors   = params.check_tensors;
+
+    mparams.n_world           = params.n_world;
+    mparams.rank              = params.rank;
+    mparams.rpc_servers       = params.rpc_servers.c_str();
+    mparams.main_gpu          = params.main_gpu;
+    mparams.split_mode        = params.split_mode;
+    mparams.tensor_split      = params.tensor_split;
+    mparams.use_mmap          = params.use_mmap;
+    mparams.use_mlock         = params.use_mlock;
+    mparams.check_tensors     = params.check_tensors;
     mparams.keep_out_in_metal = params.keep_out_in_metal;
+    mparams.keep_out_in_cuda  = params.keep_out_in_cuda;
+
     std::copy(std::begin(params.n_layer_window), std::end(params.n_layer_window), mparams.n_layer_window);
     if (params.kv_overrides.empty()) {
         mparams.kv_overrides = NULL;
@@ -2068,6 +2071,7 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param
     cparams.force             = params.force;
     cparams.master_priority   = params.master_priority;
     cparams.keep_out_in_metal = params.keep_out_in_metal;
+    cparams.keep_out_in_cuda  = params.keep_out_in_cuda;
     cparams.n_gpu_layers      = params.n_gpu_layers;
     cparams.n_cycles          = params.n_cycles;
     std::copy(std::begin(params.n_layer_window), std::end(params.n_layer_window), cparams.n_layer_window);
diff --git a/common/common.h b/common/common.h
index cd78c173..54eb7e8b 100644
--- a/common/common.h
+++ b/common/common.h
@@ -151,6 +151,7 @@ struct gpt_params {
     uint32_t signal_port          =  10000; // signal port for distributed inference
     bool    prefetch              = false; // prefetch layer weights
     bool    keep_out_in_metal     =  true; // whether to keep output weights in metal memory, true by default
+    bool    keep_out_in_cuda      = false; // whether to run the output layer on CUDA, false by default
     bool    force                 = false; // force to start prefetching after computation
     float   master_priority       =  1.01; // priority to assign workload to the master (set 1.01 to use master first, and 0.99 to offload to other devices)
     int32_t gpu_mem               = 999.0; // gpu memory to use, in GiB
diff --git a/include/llama.h b/include/llama.h
index 2faecdfe..cf7fafa0 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -325,6 +325,7 @@ extern "C" {
         bool use_mlock;     // force system to keep model in RAM
         bool check_tensors; // validate model tensor data
         bool keep_out_in_metal; // whether to keep output weights in metal memory
+        bool keep_out_in_cuda;  // whether to run the output layer on CUDA
     };
 
     // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations
@@ -339,6 +340,7 @@ extern "C" {
         bool        force;             // force to start prefetching after computation
         float       master_priority;   // priority to assign workload to the master (set 1.01 to use master first, and 0.99 to offload to other devices)
         bool        keep_out_in_metal; // whether to keep output weights in metal memory
+        bool        keep_out_in_cuda;  // whether to run the output layer on CUDA
         char *      master_ip;         // ip address of the master node
         char *      next_node_ip;      // ip address of the next node
         uint32_t    data_port;         // data port for distributed inference
diff --git a/src/llama.cpp b/src/llama.cpp
index 871daac7..e5965fd8 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -7562,6 +7562,7 @@ static bool llm_load_tensors_impl(
         int                     main_gpu,
         bool                    use_mlock,
         bool                    keep_out_in_metal,
+        bool                    keep_out_in_cuda,
         llama_progress_callback progress_callback,
         void                  * progress_callback_user_data) {
     auto & hparams = model.hparams;
@@ -7606,9 +7607,15 @@ static bool llm_load_tensors_impl(
     // assign the input and output layers on CPU by default
     if (my_rank == 0) {
         model.buft_input  = llama_default_buffer_type_cpu(model, true);
-        model.buft_output = llama_default_buffer_type_cpu(model, true);
         LLAMA_LOG_DEBUG("Layer input assigned to cpu\n");
-        LLAMA_LOG_DEBUG("Layer output assigned to cpu\n");
+
+        if (keep_out_in_cuda) {
+            model.buft_output = llama_default_buffer_type_offload(model, main_gpu);
+            LLAMA_LOG_DEBUG("Layer output assigned to gpu\n");
+        } else {
+            model.buft_output = llama_default_buffer_type_cpu(model, true);
+            LLAMA_LOG_DEBUG("Layer output assigned to cpu\n");
+        }
     }
 
     // count used buffer types
@@ -9535,7 +9542,8 @@ int llm_load_tensors(
     try {
         if (!llm_load_tensors_impl(
             *ml, *model, params.n_world, params.rank, params.n_layer_window, params.n_gpu_layers, params.split_mode, 
-            params.main_gpu, params.use_mlock, params.keep_out_in_metal, params.progress_callback, params.progress_callback_user_data
+            params.main_gpu, params.use_mlock, params.keep_out_in_metal, params.keep_out_in_cuda, params.progress_callback, 
+            params.progress_callback_user_data
         )) {
             return -2;
         }
@@ -20247,6 +20255,7 @@ struct llama_model_params llama_model_default_params() {
         /*.use_mlock                   =*/ false,
         /*.check_tensors               =*/ false,
         /*.keep_out_in_metal           =*/ true,
+        /*.keep_out_in_cuda            =*/ false,
     };
 
 #ifdef GGML_USE_METAL
@@ -20268,6 +20277,7 @@ struct llama_context_params llama_context_default_params() {
         /*.force                       =*/ false,
         /*.master_priority             =*/ 1.01,
         /*.keep_out_in_metal           =*/ true,
+        /*.keep_out_in_cuda            =*/ false,
         /*.master_ip                   =*/ nullptr,
         /*.next_node_ip                =*/ nullptr,
         /*.data_port                   =*/ 9000,
@@ -21361,14 +21371,31 @@ void * llama_context_setup_backend(
             for (size_t i = 0; i < gf.size(); ++i) {
 
 #if defined(GGML_USE_CUDA)
-                if ((cparams.rank == 0 && (i == 0 || i == gf.size() - 1)) 
-                        || model->n_gpu_layers == 0) {
+                // output layer
+                if (!params.keep_out_in_cuda && cparams.rank == 0 && i == gf.size() - 1) {
+                    continue;
+                }
+
+                // input layer
+                if (cparams.rank == 0 && i == 0) {
+                    continue;
+                }
+
+                // ignore all backend layers if n_gpu_layers is 0
+                if (model->n_gpu_layers == 0) {
+                    continue;
+                }
+
+                // don't reserve for repeated backend layers
+                if ((cparams.rank == 0 && i > 1 && i < gf.size() - 1) 
+                    || (cparams.rank > 0 && i > 0)) {
                     continue;
                 }
 #endif
 
                 ok = ok & ggml_backend_sched_reserve(ctx->sched[i], gf[i]);
             }
+            
             if (!ok) {
                 LLAMA_LOG_ERROR("%s: failed to allocate compute buffers\n", __func__);
                 llama_free(ctx);
@@ -21933,7 +21960,7 @@ void llama_model_compute_buf_size(
     // weights
     const int64_t nb_attn_norm_w = n_bytes.nb_attn_norm_w;
     const int64_t nb_attn_q_w    = n_bytes.nb_attn_q_w;
-    // const int64_t nb_output_w    = n_bytes.nb_output_w;
+    const int64_t nb_output_w    = n_bytes.nb_output_w;
     
     // format bytes
     const int64_t type_size_f32 = ggml_type_size(GGML_TYPE_F32);
@@ -21972,7 +21999,9 @@ void llama_model_compute_buf_size(
                 });
             }
             // we run the output layer on CPU by default
-            // *gpu_buf     += (n_out_embd + n_result) * type_size_f32 + nb_output_w;
+            if (cparams.keep_out_in_cuda) { 
+                *gpu_buf += (n_out_embd + n_result) * type_size_f32 + nb_output_w;
+            }
             gpu_host_buf  = (n_inp_toks + n_inp_embd + n_bak_embd + n_inp_pos + n_kq_mask + n_inp_out_ids + n_out_embd) * type_size_f32;
         } else {
             if (has_gpu_layers) {

From 2e8e42a5ad2949fa8b7f4171fddb0d88a684750b Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@foxmail.com>
Date: Mon, 23 Jun 2025 20:36:32 +0800
Subject: [PATCH 34/42] Add speculative decoding support to the server and
 command-line interfaces

---
 Makefile                             |   8 +
 common/arg.cpp                       |   4 +-
 common/common.cpp                    |   2 +-
 common/common.h                      |  18 +-
 common/sampling.cpp                  |  39 ++++
 common/sampling.h                    |  21 +++
 common/speculative.cpp               | 271 +++++++++++++++++++++++++++
 common/speculative.h                 |  28 +++
 examples/server/server.cpp           | 221 +++++++++++++++++++---
 examples/speculative/speculative.cpp |   4 +-
 src/llama.cpp                        |   6 +
 11 files changed, 591 insertions(+), 31 deletions(-)
 create mode 100644 common/speculative.cpp
 create mode 100644 common/speculative.h

diff --git a/Makefile b/Makefile
index 39ae0b9f..b1bbc0ed 100644
--- a/Makefile
+++ b/Makefile
@@ -963,6 +963,7 @@ OBJ_COMMON = \
 	common/console.o \
 	common/ngram-cache.o \
 	common/sampling.o \
+	common/speculative.o \
 	common/train.o \
 	common/build-info.o \
 	common/json-schema-to-grammar.o
@@ -1239,6 +1240,13 @@ common/json-schema-to-grammar.o: \
 	common/json-schema-to-grammar.h
 	$(CXX) $(CXXFLAGS) -c $< -o $@
 
+# speculative 
+common/speculative.o: \
+	common/speculative.cpp \
+	common/speculative.h \
+	include/llama.h
+	$(CXX) $(CXXFLAGS) -c $< -o $@
+
 common/train.o: \
 	common/train.cpp \
 	common/train.h
diff --git a/common/arg.cpp b/common/arg.cpp
index c971bb49..813f87e8 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -1704,9 +1704,9 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
         {"-md", "--model-draft"}, "FNAME",
         "draft model for speculative decoding (default: unused)",
         [](gpt_params & params, const std::string & value) {
-            params.model_draft = value;
+            params.speculative.model = value;
         }
-    ).set_examples({LLAMA_EXAMPLE_SPECULATIVE}));
+    ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER}));
     add_opt(llama_arg(
         {"-mu", "--model-url"}, "MODEL_URL",
         "model download url (default: unused)",
diff --git a/common/common.cpp b/common/common.cpp
index 92c95b3b..b6182b22 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -3111,7 +3111,7 @@ void yaml_dump_non_result_info(FILE * stream, const gpt_params & params, const l
     fprintf(stream, "mirostat_lr: %f # default: 0.1\n", sparams.mirostat_eta);
     fprintf(stream, "mlock: %s # default: false\n", params.use_mlock ? "true" : "false");
     fprintf(stream, "model: %s # default: %s\n", params.model.c_str(), DEFAULT_MODEL_PATH);
-    fprintf(stream, "model_draft: %s # default:\n", params.model_draft.c_str());
+    fprintf(stream, "model_draft: %s # default:\n", params.speculative.model.c_str());
     fprintf(stream, "multiline_input: %s # default: false\n", params.multiline_input ? "true" : "false");
     fprintf(stream, "n_gpu_layers: %d # default: -1\n", params.n_gpu_layers);
     fprintf(stream, "n_predict: %d # default: -1 (unlimited)\n", params.n_predict);
diff --git a/common/common.h b/common/common.h
index 54eb7e8b..044dfdf5 100644
--- a/common/common.h
+++ b/common/common.h
@@ -33,6 +33,8 @@ struct llama_lora_adapter_container : llama_lora_adapter_info {
     struct llama_lora_adapter * adapter;
 };
 
+using llama_tokens = std::vector<llama_token>;
+
 // build info
 extern int LLAMA_BUILD_NUMBER;
 extern char const * LLAMA_COMMIT;
@@ -141,6 +143,20 @@ struct gpt_sampler_params {
     std::string print() const;
 };
 
+struct common_params_speculative {
+    int32_t n_ctx        =     0; // draft context size
+    int32_t n_max        =    16; // maximum number of tokens to draft during speculative decoding
+    int32_t n_min        =     5; // minimum number of draft tokens to use for speculative decoding
+    int32_t n_gpu_layers =    -1; // number of layers to store in VRAM for the draft model (-1 - use default)
+    float   p_split      =  0.1f; // speculative decoding split probability
+    float   p_min        =  0.9f; // minimum speculative decoding probability (greedy)
+
+    struct cpu_params cpuparams;
+    struct cpu_params cpuparams_batch;
+
+    std::string model = ""; // draft model for speculative decoding                          // NOLINT
+};
+
 struct gpt_params {
     int32_t n_world               =     1; // number of devices to use
     int32_t rank                  =     0; // my rank for distributed inference
@@ -198,9 +214,9 @@ struct gpt_params {
     enum llama_attention_type    attention_type    = LLAMA_ATTENTION_TYPE_UNSPECIFIED; // attention type for embeddings
 
     struct gpt_sampler_params sparams;
+    struct common_params_speculative speculative;
 
     std::string model                = ""; // model path                                                    // NOLINT
-    std::string model_draft          = ""; // draft model for speculative decoding                          // NOLINT
     std::string model_alias          = "unknown"; // model alias                                            // NOLINT
     std::string model_url            = ""; // model url to download                                         // NOLINT
     std::string hf_token             = ""; // HF token                                                      // NOLINT
diff --git a/common/sampling.cpp b/common/sampling.cpp
index 3dc7f112..b8db920a 100644
--- a/common/sampling.cpp
+++ b/common/sampling.cpp
@@ -318,6 +318,45 @@ llama_token gpt_sampler_sample(struct gpt_sampler * gsmpl, struct llama_context
     return cur_p.data[cur_p.selected].id;
 }
 
+std::vector<llama_token> gpt_sampler_sample_and_accept_n(struct gpt_sampler * gsmpl, struct llama_context * ctx, const std::vector<int> & idxs, const llama_tokens & draft, bool grammar_first) {
+    GGML_ASSERT(idxs.size() == draft.size() + 1 && "idxs.size() must be draft.size() + 1");
+
+    std::vector<llama_token> result;
+    result.reserve(idxs.size());
+
+    size_t i = 0;
+    for (; i < draft.size(); i++) {
+        const llama_token id = gpt_sampler_sample(gsmpl, ctx, idxs[i], grammar_first);
+
+        gpt_sampler_accept(gsmpl, id, true);
+
+        result.push_back(id);
+
+        if (draft[i] != id) {
+            break;
+        }
+    }
+
+    if (i == draft.size()) {
+        const llama_token id = gpt_sampler_sample(gsmpl, ctx, idxs[i], grammar_first);
+
+        gpt_sampler_accept(gsmpl, id, true);
+
+        result.push_back(id);
+    }
+
+    return result;
+}
+
+std::vector<llama_token> gpt_sampler_sample_and_accept_n(struct gpt_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first) {
+    std::vector<int> idxs(draft.size() + 1);
+    for (size_t i = 0; i < idxs.size(); ++i) {
+        idxs[i] = i;
+    }
+
+    return gpt_sampler_sample_and_accept_n(gsmpl, ctx, idxs, draft, grammar_first);
+}
+
 uint32_t gpt_sampler_get_seed(const struct gpt_sampler * gsmpl) {
     return llama_sampler_get_seed(gsmpl->chain);
 }
diff --git a/common/sampling.h b/common/sampling.h
index d0e1a920..82d6b023 100644
--- a/common/sampling.h
+++ b/common/sampling.h
@@ -60,6 +60,27 @@ void gpt_perf_print(const struct llama_context * ctx, const struct gpt_sampler *
 //
 llama_token gpt_sampler_sample(struct gpt_sampler * gsmpl, struct llama_context * ctx, int idx, bool grammar_first = false);
 
+// generalized version of gpt_sampler_sample
+//
+// will cross-reference the sampled tokens with a batch of draft tokens and accept those that match
+// if the sampler disagrees at some point, we stop and return the accepted tokens up to now
+//
+//      gpt_sampler_sample_n(gsmpl, ctx, { idx }, {});
+//
+// is equivalent to
+//
+//      gpt_sampler_sample(gsmpl, ctx, idx);
+//      gpt_sampler_accept(gsmpl, token, true);
+//
+// requires: idxs.size() == draft.size() + 1
+//
+// returns at least 1 token, up to idxs.size()
+//
+std::vector<llama_token> gpt_sampler_sample_and_accept_n(struct gpt_sampler * gsmpl, struct llama_context * ctx, const std::vector<int> & idxs, const llama_tokens & draft, bool grammar_first = false);
+
+// assume idxs == [ 0, 1, 2, ..., draft.size() ]
+std::vector<llama_token> gpt_sampler_sample_and_accept_n(struct gpt_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first = false);
+
 uint32_t gpt_sampler_get_seed(const struct gpt_sampler * gsmpl);
 
 // helpers
diff --git a/common/speculative.cpp b/common/speculative.cpp
new file mode 100644
index 00000000..8a03d08c
--- /dev/null
+++ b/common/speculative.cpp
@@ -0,0 +1,271 @@
+#include "speculative.h"
+
+#include "log.h"
+#include "common.h"
+#include "sampling.h"
+
+#include <cstring>
+
+#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE  128
+#define SPEC_VOCAB_CHECK_START_TOKEN_ID 5
+
+struct common_speculative {
+    struct llama_context * ctx;
+    struct gpt_sampler * smpl;
+
+    llama_batch batch;
+    llama_tokens prompt;
+};
+
+struct common_speculative * common_speculative_init(
+        struct llama_context * ctx_dft) {
+    auto * result = new common_speculative {
+        /* .ctx    = */ ctx_dft,
+        /* .smpl   = */ nullptr,
+        /* .batch  = */ llama_batch_init(llama_n_batch(ctx_dft), 0, 1),
+        /* .prompt = */ {},
+    };
+
+    // TODO: optimize or pass from outside?
+#if 0
+    {
+        common_params_sampling params;
+        params.no_perf = false;
+
+        params.top_k = 40;
+        params.top_p = 0.9;
+
+        params.samplers = {
+            COMMON_SAMPLER_TYPE_TOP_K,
+            COMMON_SAMPLER_TYPE_TOP_P,
+            COMMON_SAMPLER_TYPE_INFILL,
+        };
+
+        result->smpl = gpt_sampler_init(llama_get_model(ctx_dft), params);
+    }
+#else
+    {
+        gpt_sampler_params params;
+        params.no_perf = false;
+
+        params.top_k = 10;
+
+        params.samplers = {
+            GPT_SAMPLER_TYPE_TOP_K,
+        };
+
+        result->smpl = gpt_sampler_init(llama_get_model(ctx_dft), params);
+    }
+#endif
+
+    llama_update_context_with_rankworld(result->ctx, 0, 1, 0, 1);
+
+    return result;
+}
+
+void common_speculative_free(struct common_speculative * spec) {
+    gpt_sampler_free(spec->smpl);
+
+    llama_batch_free(spec->batch);
+
+    delete spec;
+}
+
+bool common_speculative_are_compatible(
+        const struct llama_context * ctx_tgt,
+        const struct llama_context * ctx_dft) {
+    const struct llama_model * model_tgt = llama_get_model(ctx_tgt);
+    const struct llama_model * model_dft = llama_get_model(ctx_dft);
+
+    const bool vocab_type_tgt = llama_vocab_type(model_tgt);
+    LOG_DBG("%s: vocab_type tgt: %d\n", __func__, vocab_type_tgt);
+
+    const bool vocab_type_dft = llama_vocab_type(model_dft);
+    LOG_DBG("%s: vocab_type dft: %d\n", __func__, vocab_type_dft);
+
+    if (vocab_type_tgt != vocab_type_dft) {
+        LOG_ERR("%s: draft model vocab type must match target model to use speculation but "
+                     "vocab_type_dft = %d while vocab_type_tgt = %d\n", __func__, vocab_type_dft, vocab_type_tgt);
+        return false;
+    }
+
+    if (llama_add_bos_token(model_tgt) != llama_add_bos_token(model_dft) ||
+        llama_add_eos_token(model_tgt) != llama_add_eos_token(model_dft) ||
+        llama_token_bos(model_tgt) != llama_token_bos(model_dft) ||
+        llama_token_eos(model_tgt) != llama_token_eos(model_dft)
+    ) {
+        LOG_ERR("%s: draft model special tokens must match target model to use speculation\n", __func__);
+        return false;
+    }
+
+    {
+        const int n_vocab_tgt = llama_n_vocab(model_tgt);
+        const int n_vocab_dft = llama_n_vocab(model_dft);
+
+        const int vocab_diff = std::abs(n_vocab_tgt - n_vocab_dft);
+
+        if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) {
+            LOG_ERR("%s: draft model vocab must closely match target model to use speculation but "
+                         "target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n",
+                    __func__, n_vocab_tgt, llama_n_vocab(model_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE);
+            return false;
+        }
+
+        for (int i = SPEC_VOCAB_CHECK_START_TOKEN_ID; i < std::min(n_vocab_tgt, n_vocab_dft); ++i) {
+            const char * token_text_tgt = llama_token_get_text(model_tgt, i);
+            const char * token_text_dft = llama_token_get_text(model_dft, i);
+            if (std::strcmp(token_text_tgt, token_text_dft) != 0) {
+                LOG_ERR("%s: draft model vocab must match target model to use speculation but "
+                             "token %d content differs - target '%s', draft '%s'\n", __func__, i,
+                        llama_token_to_piece(ctx_tgt, i).c_str(),
+                        llama_token_to_piece(ctx_dft, i).c_str());
+                return false;
+            }
+        }
+    }
+
+    return true;
+}
+
+llama_tokens common_speculative_gen_draft(
+        struct common_speculative * spec,
+        struct common_speculative_params params,
+        const llama_tokens & prompt_tgt,
+        llama_token id_last) {
+    auto & batch  = spec->batch;
+    auto & ctx    = spec->ctx;
+    auto & smpl   = spec->smpl;
+    auto & prompt = spec->prompt;
+
+    int reuse_i = 0;
+    int reuse_n = 0;
+
+    const int n_ctx = llama_n_ctx(ctx) - params.n_draft;
+
+    const int i_start = std::max<int>(0, (int) prompt_tgt.size() - n_ctx);
+
+    // reuse as much as possible from the old draft context
+    // ideally, the draft context should be as big as the target context and we will always reuse the entire prompt
+    for (int i = 0; i < (int) prompt.size(); ++i) {
+        int cur = 0;
+        while (i_start + cur < (int) prompt_tgt.size() &&
+               i       + cur < (int) prompt.size() &&
+               prompt_tgt[i_start + cur] == prompt[i + cur]) {
+            cur++;
+        }
+
+        if ((cur >= params.n_reuse || n_ctx >= (int) prompt_tgt.size()) && cur > reuse_n) {
+            reuse_i = i;
+            reuse_n = cur;
+        }
+    }
+
+    LOG_DBG("%s: reuse_i = %d, reuse_n = %d, prompt = %d\n", __func__, reuse_i, reuse_n, (int) prompt.size());
+
+    llama_tokens result;
+    result.reserve(params.n_draft);
+
+    if (reuse_n == 0) {
+        llama_kv_cache_clear(ctx);
+
+        prompt.clear();
+    } else {
+        // this happens when a previous draft has been discarded (for example, due to being too small), but the
+        // target model agreed with it. in this case, we simply pass back the previous results to save compute
+        if (reuse_i + reuse_n < (int) prompt.size() && prompt[reuse_i + reuse_n] == id_last) {
+            for (int i = reuse_i + reuse_n + 1; i < (int) prompt.size(); ++i) {
+                result.push_back(prompt[i]);
+
+                if (params.n_draft <= (int) result.size()) {
+                    break;
+                }
+            }
+
+            return result;
+        }
+
+        if (reuse_i > 0) {
+            llama_kv_cache_seq_rm (ctx, 0, 0, reuse_i);
+            llama_kv_cache_seq_add(ctx, 0, reuse_i, -1, -reuse_i);
+
+            prompt.erase(prompt.begin(), prompt.begin() + reuse_i);
+        }
+
+        if (reuse_n < (int) prompt.size()) {
+            llama_kv_cache_seq_rm (ctx, 0, reuse_n, -1);
+
+            prompt.erase(prompt.begin() + reuse_n, prompt.end());
+        }
+    }
+
+    // prepare a batch to evaluate any new tokens in the prompt
+    llama_batch_clear(batch);
+
+    for (size_t i = i_start + reuse_n; i < prompt_tgt.size(); ++i) {
+        //LOG_DBG("i = %d, i_start = %d, reuse_n = %d, i - i_start = %d, id = %6d\n", i, i_start, reuse_n, i - i_start, prompt_tgt[i]);
+        llama_batch_add(batch, prompt_tgt[i], i - i_start, { 0 }, false);
+
+        prompt.push_back(prompt_tgt[i]);
+    }
+
+    // we should rarely end-up here during normal decoding
+    if (batch.n_tokens > 0) {
+        //LOG_DBG("%s: draft prompt batch: %s\n", __func__, string_from(ctx, batch).c_str());
+
+        llama_decode(ctx, batch);
+    }
+
+    const llama_pos n_past = prompt.size();
+
+    LOG_DBG("%s: n_past = %d\n", __func__, n_past);
+
+    llama_batch_clear(batch);
+    llama_batch_add  (batch, id_last, n_past, { 0 }, true);
+
+    prompt.push_back(id_last);
+
+    //LOG_DBG("%s: draft prompt: %s\n", __func__, string_from(ctx, prompt).c_str());
+
+    llama_decode(ctx, batch);
+
+    gpt_sampler_reset(smpl);
+
+    // sample n_draft tokens from the draft model
+    for (int i = 0; i < params.n_draft; ++i) {
+        llama_batch_clear(batch);
+
+        gpt_sampler_sample(smpl, ctx, 0, true);
+
+        const auto * cur_p = gpt_sampler_get_candidates(smpl);
+
+        for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
+            LOG_DBG(" - draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
+                    k, i, cur_p->data[k].id, cur_p->data[k].p, llama_token_to_piece(ctx, cur_p->data[k].id).c_str());
+        }
+
+        // add drafted token for each sequence
+        const llama_token id = cur_p->data[0].id;
+
+        // only collect very high-confidence draft tokens
+        if (cur_p->data[0].p < params.p_min) {
+            break;
+        }
+
+        gpt_sampler_accept(smpl, id, true);
+
+        result.push_back(id);
+
+        if (params.n_draft <= (int) result.size()) {
+            break;
+        }
+
+        llama_batch_add(batch, id, n_past + i + 1, { 0 }, true);
+
+        // evaluate the drafted tokens on the draft model
+        llama_decode(ctx, batch);
+
+        prompt.push_back(id);
+    }
+
+    return result;
+}
\ No newline at end of file
diff --git a/common/speculative.h b/common/speculative.h
new file mode 100644
index 00000000..0af10ea4
--- /dev/null
+++ b/common/speculative.h
@@ -0,0 +1,28 @@
+#pragma once
+
+#include "llama.h"
+#include "common.h"
+
+struct common_speculative;
+
+struct common_speculative_params {
+    int n_draft = 16;  // max drafted tokens
+    int n_reuse = 256;
+
+    float p_min = 0.9f; // min probabiliy required to accept a token in the draft
+};
+
+struct common_speculative * common_speculative_init(struct llama_context * ctx_dft);
+
+void common_speculative_free(struct common_speculative * spec);
+
+bool common_speculative_are_compatible(
+        const struct llama_context * ctx_tgt,
+        const struct llama_context * ctx_dft);
+
+// sample up to n_draft tokens and add them to the batch using the draft model
+llama_tokens common_speculative_gen_draft(
+               struct common_speculative * spec,
+        struct common_speculative_params   params,
+                      const llama_tokens & prompt,
+                             llama_token   id_last);
\ No newline at end of file
diff --git a/examples/server/server.cpp b/examples/server/server.cpp
index af39f1ac..d086eaf3 100644
--- a/examples/server/server.cpp
+++ b/examples/server/server.cpp
@@ -6,6 +6,7 @@
 #include "sampling.h"
 #include "json-schema-to-grammar.h"
 #include "llama.h"
+#include "speculative.h"
 
 // Change JSON_ASSERT from assert() to GGML_ASSERT:
 #define JSON_ASSERT GGML_ASSERT
@@ -133,6 +134,9 @@ struct slot_params {
     int32_t  n_predict = -1; // new tokens to predict
 
     std::vector<std::string> antiprompt;
+    
+    struct gpt_sampler_params        sampling;
+    struct common_params_speculative speculative;
 
     json input_prefix;
     json input_suffix;
@@ -142,6 +146,12 @@ struct server_slot {
     int id;
     int id_task = -1;
 
+    llama_batch batch_spec;
+
+    llama_context * ctx_dft = nullptr;
+
+    common_speculative * spec = nullptr;
+
     // the index relative to completion multi-task request
     size_t index = 0;
 
@@ -231,7 +241,7 @@ struct server_slot {
         generated_token_probs.clear();
     }
 
-    bool has_budget(gpt_params &global_params) {
+    bool has_budget(const gpt_params &global_params) {
         if (params.n_predict == -1 && global_params.n_predict == -1) {
             return true; // limitless
         }
@@ -251,6 +261,10 @@ struct server_slot {
         return state != SLOT_STATE_IDLE;
     }
 
+    bool can_speculate() const {
+        return ctx_dft && params.speculative.n_max > 0 && params.cache_prompt;
+    }
+
     void add_token(const completion_token_output & token) {
         if (!is_processing()) {
             SLT_WRN(*this, "%s", "slot is not processing\n");
@@ -615,6 +629,9 @@ struct server_context {
 
     gpt_params params;
 
+    llama_model * model_dft = nullptr;
+    llama_context_params cparams_dft;
+
     llama_batch batch = {};
 
     bool clean_kv_cache = true;
@@ -652,17 +669,33 @@ struct server_context {
             model = nullptr;
         }
 
+        if (model_dft) {
+            llama_free_model(model_dft);
+            model_dft = nullptr;
+        }
+
         // Clear any sampling context
         for (server_slot & slot : slots) {
             if (slot.smpl != nullptr) {
                 gpt_sampler_free(slot.smpl);
             }
+            slot.smpl = nullptr;
+
+            llama_free(slot.ctx_dft);
+            slot.ctx_dft = nullptr;
+
+            common_speculative_free(slot.spec);
+            slot.spec = nullptr;
+
+            llama_batch_free(slot.batch_spec);
         }
 
         llama_batch_free(batch);
     }
 
     bool load_model(const gpt_params & params_) {
+        SRV_INF("loading model '%s'\n", params.model.c_str());
+
         params = params_;
 
         // dedicate one sequence to the system prompt
@@ -685,6 +718,44 @@ struct server_context {
 
         add_bos_token = llama_add_bos_token(model);
         has_eos_token = !llama_add_eos_token(model);
+        
+        if (!params.speculative.model.empty()) {
+            SRV_INF("loading draft model '%s'\n", params.speculative.model.c_str());
+
+            auto params_dft = params;
+
+            params_dft.model        = params.speculative.model;
+            params_dft.n_ctx        = params.speculative.n_ctx;
+            params_dft.n_gpu_layers = params.speculative.n_gpu_layers;
+            params_dft.n_world      = 1;  // do not split the draft model across devicesAdd commentMore actions
+            params_dft.rank         = 0;  // always load the draft model on the head device
+
+            std::fill_n(params_dft.n_layer_window, params.n_world, 0);
+            
+            llama_init_result llama_init_dft = llama_init_from_gpt_params(params_dft);
+
+            model_dft = llama_init_dft.model;
+
+            if (model_dft == nullptr) {
+                SRV_ERR("failed to load draft model, '%s'\n", params.speculative.model.c_str());
+                return false;
+            }
+
+            if (!common_speculative_are_compatible(ctx, llama_init_dft.context)) {
+                SRV_ERR("the draft model '%s' is not compatible with the target model '%s'\n", params.speculative.model.c_str(), params.model.c_str());
+
+                llama_free      (llama_init_dft.context);
+                llama_free_model(llama_init_dft.model);
+
+                return false;
+            }
+
+            cparams_dft = llama_context_params_from_gpt_params(params);
+            cparams_dft.n_batch = llama_n_ctx(llama_init_dft.context);
+
+            // the context is not needed - we will create one for each slot
+            llama_free(llama_init_dft.context);
+        }
 
         return true;
     }
@@ -708,6 +779,30 @@ struct server_context {
             slot.id = i;
             slot.n_ctx = n_ctx_slot;
             slot.n_predict = params.n_predict;
+            
+            if (model_dft) {
+                slot.batch_spec = llama_batch_init(params.speculative.n_max + 1, 0, 1);
+
+                slot.ctx_dft = llama_new_context_with_model(model_dft, cparams_dft);
+                
+                if (llama_context_setup_backend(model, cparams_dft, slot.ctx_dft) == nullptr) {
+                    SRV_ERR("%s: failed to setup context with model '%s'\n", __func__, params.model.c_str());
+                    llama_free(slot.ctx_dft);
+                    llama_free_model(model);
+                    return;
+                }
+                
+                if (slot.ctx_dft == nullptr) {
+                    SRV_ERR("%s", "failed to create draft context\n");
+                    return;
+                }
+
+                slot.spec = common_speculative_init(slot.ctx_dft);
+                if (slot.spec == nullptr) {
+                    SRV_ERR("%s", "failed to create speculator\n");
+                    return;
+                }
+            }
 
             SLT_INF(slot, "new slot n_ctx_slot = %d\n", slot.n_ctx);
 
@@ -875,6 +970,8 @@ struct server_context {
         slot_params default_params;
         // Sampling parameter defaults are loaded from the global server context (but individual requests can still override them)
         auto default_sparams = params.sparams;
+        default_params.speculative = params.speculative;
+        
         const auto & data = task.data;
 
         if (data.count("__oaicompat") != 0) {
@@ -909,6 +1006,12 @@ struct server_context {
         slot.sparams.seed              = json_value(data, "seed",              default_sparams.seed);
         slot.sparams.n_probs           = json_value(data, "n_probs",           default_sparams.n_probs);
         slot.sparams.min_keep          = json_value(data, "min_keep",          default_sparams.min_keep);
+        
+        slot.params.speculative.n_min = json_value(data, "speculative.n_min", default_params.speculative.n_min);
+        slot.params.speculative.n_max = json_value(data, "speculative.n_max", default_params.speculative.n_max);
+        slot.params.speculative.p_min = json_value(data, "speculative.p_min", default_params.speculative.p_min);
+
+        slot.params.speculative.n_min = std::min(slot.params.speculative.n_max, slot.params.speculative.n_min);
 
         // process "json_schema" and "grammar"
         if (data.contains("json_schema") && !data.at("json_schema").is_null() && data.contains("grammar") && !data.at("grammar").is_null()) {
@@ -1049,6 +1152,12 @@ struct server_context {
                 return false;
             }
         }
+        
+        if (slot.ctx_dft) {
+            llama_batch_free(slot.batch_spec);
+
+            slot.batch_spec = llama_batch_init(slot.params.speculative.n_max + 1, 0, 1);
+        }
 
         slot.state = SLOT_STATE_PROCESSING_PROMPT;
         slot.prompt_tokens.clear();
@@ -2357,38 +2466,100 @@ struct server_context {
                     continue; // continue loop of slots
                 }
 
-                completion_token_output result;
-                const llama_token id = gpt_sampler_sample(slot.smpl, ctx, slot.i_batch - i);
+                llama_token id;
 
-                gpt_sampler_accept(slot.smpl, id, true);
+                {
+                    completion_token_output result;
 
-                slot.n_decoded += 1;
-                if (slot.n_decoded == 1) {
-                    slot.t_start_generation = ggml_time_us();
-                    slot.t_prompt_processing = (slot.t_start_generation - slot.t_start_process_prompt) / 1e3;
-                    metrics.on_prompt_eval(slot);
+                    id = gpt_sampler_sample(slot.smpl, ctx, slot.i_batch - i);
+
+                    slot.i_batch = -1;
+
+                    gpt_sampler_accept(slot.smpl, id, true);
+
+                    slot.n_decoded += 1;
+                    if (slot.n_decoded == 1) {
+                        slot.t_start_generation = ggml_time_us();
+                        slot.t_prompt_processing = (slot.t_start_generation - slot.t_start_process_prompt) / 1e3;
+                        metrics.on_prompt_eval(slot);
+                    }
+
+                    result.tok = id;
+
+                    const auto * cur_p = gpt_sampler_get_candidates(slot.smpl);
+
+                    for (size_t i = 0; i < (size_t) slot.params.sampling.n_probs; ++i) {
+                        result.probs.push_back({
+                            cur_p->data[i].id,
+                                i >= cur_p->size ? 0.0f : cur_p->data[i].p,
+                        });
+                    }
+
+                    if (!process_token(result, slot)) {
+                        // release slot because of stop condition
+                        slot.release();
+                        slot.print_timings();
+                        send_final_response(slot);
+                        metrics.on_prediction(slot);
+                        continue;
+                    }
                 }
 
-                result.tok = id;
-
-                const auto * cur_p = gpt_sampler_get_candidates(slot.smpl);
-
-                for (size_t i = 0; i < (size_t) slot.sparams.n_probs; ++i) {
-                    result.probs.push_back({
-                        cur_p->data[i].id,
-                        i >= cur_p->size ? 0.0f : cur_p->data[i].p,
-                    });
+                // check if the slot supports speculative decoding
+                if (!slot.can_speculate()) {
+                    continue;
                 }
 
-                if (!process_token(result, slot)) {
-                    // release slot because of stop condition
-                    slot.release();
-                    slot.print_timings();
-                    send_final_response(slot);
-                    metrics.on_prediction(slot);
+                struct common_speculative_params params_spec;
+                params_spec.n_draft   = slot.params.speculative.n_max;
+                params_spec.n_reuse   = llama_n_ctx(slot.ctx_dft) - slot.params.speculative.n_max;
+                params_spec.p_min     = slot.params.speculative.p_min;
+
+                llama_tokens draft = common_speculative_gen_draft(slot.spec, params_spec, slot.cache_tokens, id);
+
+                // ignore small drafts
+                if (slot.params.speculative.n_min > (int) draft.size()) {
+                    continue;
                 }
 
-                slot.i_batch = -1;
+                // construct the speculation batch
+                llama_batch_clear(slot.batch_spec);
+                llama_batch_add  (slot.batch_spec, id, slot.n_past, { slot.id }, true);
+
+                for (size_t i = 0; i < draft.size(); ++i) {
+                    llama_batch_add(slot.batch_spec, draft[i], slot.n_past + 1 + i, { slot.id }, true);
+                }
+
+                llama_decode(ctx, slot.batch_spec);
+
+                // the accepted tokens from the speculation
+                const auto ids = gpt_sampler_sample_and_accept_n(slot.smpl, ctx, draft);
+
+                slot.n_past    += ids.size();
+                slot.n_decoded += ids.size();
+
+                slot.cache_tokens.push_back(id);
+                slot.cache_tokens.insert(slot.cache_tokens.end(), ids.begin(), ids.end() - 1);
+
+                llama_kv_cache_seq_rm(ctx, slot.id, slot.n_past, -1);
+
+                for (size_t i = 0; i < ids.size(); ++i) {
+                    completion_token_output result;
+
+                    result.tok = ids[i];
+
+                    if (!process_token(result, slot)) {
+                        // release slot because of stop condition
+                        slot.release();
+                        slot.print_timings();
+                        send_final_response(slot);
+                        metrics.on_prediction(slot);
+                        break;
+                    }
+                }
+
+                SRV_DBG("accepted %d/%d draft tokens\n", (int) ids.size() - 1, (int) draft.size());
+                
             }
         }
 
diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp
index 8527f8b6..0a6c4701 100644
--- a/examples/speculative/speculative.cpp
+++ b/examples/speculative/speculative.cpp
@@ -41,7 +41,7 @@ int main(int argc, char ** argv) {
 
     gpt_init();
 
-    if (params.model_draft.empty()) {
+    if (params.speculative.model.empty()) {
         LOG_ERR("%s: --model-draft is required\n", __func__);
         return 1;
     }
@@ -68,7 +68,7 @@ int main(int argc, char ** argv) {
     // load the draft model
     // make a hard copy of params to use for the draft model
     gpt_params params_draft   = params;
-    params_draft.model        = params_draft.model_draft;
+    params_draft.model        = params_draft.speculative.model;
     params_draft.n_gpu_layers = params_draft.n_gpu_layers_draft;
     params_draft.n_world      = 1;    // do not split the draft model across devices
     params_draft.rank         = 0;    // always load the draft model on the head device
diff --git a/src/llama.cpp b/src/llama.cpp
index e5965fd8..a57715e3 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -20944,6 +20944,12 @@ struct llama_context * llama_new_context_with_model(
     ctx->cparams.rank    = params.rank;
     ctx->cparams.force   = params.force;
     ctx->cparams.original_next_rank = (params.rank + 1) % params.n_world;
+    
+    auto &hparams = model->hparams;
+    auto &cparams = ctx->cparams;
+    cparams.n_ctx            = params.n_ctx           == 0    ? hparams.n_ctx_train           : params.n_ctx;
+    ctx->logits_all = params.logits_all;
+
     return ctx;
 }
 

From d248f3c40e127897eb8a981ed38b90dcc2c51037 Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@gmail.com>
Date: Fri, 27 Jun 2025 06:07:47 +0000
Subject: [PATCH 35/42] fix: some fields in cparams_draft

---
 common/arg.cpp             |  5 +++--
 examples/server/server.cpp | 14 ++++++++++----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/common/arg.cpp b/common/arg.cpp
index 813f87e8..2f999dcc 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -1555,13 +1555,14 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
         {"-ngld", "--gpu-layers-draft", "--n-gpu-layers-draft"}, "N",
         "number of layers to store in VRAM for the draft model",
         [](gpt_params & params, int value) {
-            params.n_gpu_layers_draft = value;
+            params.n_gpu_layers_draft = value; // TODO: remove
+            params.speculative.n_gpu_layers = value;
             if (!llama_supports_gpu_offload()) {
                 fprintf(stderr, "warning: not compiled with GPU offload support, --gpu-layers-draft option will be ignored\n");
                 fprintf(stderr, "warning: see main README.md for information on enabling GPU BLAS support\n");
             }
         }
-    ).set_examples({LLAMA_EXAMPLE_SPECULATIVE}));
+    ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER}));
     add_opt(llama_arg(
         {"-sm", "--split-mode"}, "{none,layer,row}",
         "how to split the model across multiple GPUs, one of:\n"
diff --git a/examples/server/server.cpp b/examples/server/server.cpp
index d086eaf3..93a9cf33 100644
--- a/examples/server/server.cpp
+++ b/examples/server/server.cpp
@@ -727,6 +727,7 @@ struct server_context {
             params_dft.model        = params.speculative.model;
             params_dft.n_ctx        = params.speculative.n_ctx;
             params_dft.n_gpu_layers = params.speculative.n_gpu_layers;
+            params_dft.use_mlock    = true;
             params_dft.n_world      = 1;  // do not split the draft model across devicesAdd commentMore actions
             params_dft.rank         = 0;  // always load the draft model on the head device
 
@@ -749,9 +750,14 @@ struct server_context {
 
                 return false;
             }
-
+            
             cparams_dft = llama_context_params_from_gpt_params(params);
             cparams_dft.n_batch = llama_n_ctx(llama_init_dft.context);
+            cparams_dft.n_world = 1;
+            cparams_dft.rank    = 0;
+            std::fill_n(cparams_dft.n_layer_window, 32, 0);
+            cparams_dft.n_layer_window[0] = llama_n_layer(model_dft);
+            cparams_dft.n_gpu_layers      = params.speculative.n_gpu_layers;
 
             // the context is not needed - we will create one for each slot
             llama_free(llama_init_dft.context);
@@ -785,10 +791,10 @@ struct server_context {
 
                 slot.ctx_dft = llama_new_context_with_model(model_dft, cparams_dft);
                 
-                if (llama_context_setup_backend(model, cparams_dft, slot.ctx_dft) == nullptr) {
-                    SRV_ERR("%s: failed to setup context with model '%s'\n", __func__, params.model.c_str());
+                if (llama_context_setup_backend(model_dft, cparams_dft, slot.ctx_dft) == nullptr) {
+                    SRV_ERR("%s: failed to setup context with model '%s'\n", __func__, params.speculative.model.c_str());
                     llama_free(slot.ctx_dft);
-                    llama_free_model(model);
+                    llama_free_model(model_dft);
                     return;
                 }
                 

From 9bf6565df4b94a43414a1821220be55247abee15 Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@gmail.com>
Date: Fri, 27 Jun 2025 06:30:57 +0000
Subject: [PATCH 36/42] fix: load draft model first

---
 examples/server/server.cpp | 66 ++++++++++++++++++++------------------
 1 file changed, 35 insertions(+), 31 deletions(-)

diff --git a/examples/server/server.cpp b/examples/server/server.cpp
index 93a9cf33..5b970e4f 100644
--- a/examples/server/server.cpp
+++ b/examples/server/server.cpp
@@ -700,6 +700,40 @@ struct server_context {
 
         // dedicate one sequence to the system prompt
         params.n_parallel += 1;
+        
+        // load draft model first
+        llama_init_result llama_init_dft;
+        if (!params.speculative.model.empty()) {
+            SRV_INF("loading draft model '%s'\n", params.speculative.model.c_str());
+
+            auto params_dft = params;
+
+            params_dft.model        = params.speculative.model;
+            params_dft.n_ctx        = params.speculative.n_ctx;
+            params_dft.n_gpu_layers = params.speculative.n_gpu_layers;
+            params_dft.use_mlock    = true;
+            params_dft.n_world      = 1;  // do not split the draft model across devicesAdd commentMore actions
+            params_dft.rank         = 0;  // always load the draft model on the head device
+
+            std::fill_n(params_dft.n_layer_window, params.n_world, 0);
+            
+            llama_init_dft = llama_init_from_gpt_params(params_dft);
+
+            model_dft = llama_init_dft.model;
+
+            if (model_dft == nullptr) {
+                SRV_ERR("failed to load draft model, '%s'\n", params.speculative.model.c_str());
+                return false;
+            }
+            
+            cparams_dft = llama_context_params_from_gpt_params(params);
+            cparams_dft.n_batch = llama_n_ctx(llama_init_dft.context);
+            cparams_dft.n_world = 1;
+            cparams_dft.rank    = 0;
+            std::fill_n(cparams_dft.n_layer_window, 32, 0);
+            cparams_dft.n_layer_window[0] = llama_n_layer(model_dft);
+            cparams_dft.n_gpu_layers      = params.speculative.n_gpu_layers;
+        }
 
         llama_init_result llama_init = llama_init_from_gpt_params(params);
 
@@ -719,29 +753,7 @@ struct server_context {
         add_bos_token = llama_add_bos_token(model);
         has_eos_token = !llama_add_eos_token(model);
         
-        if (!params.speculative.model.empty()) {
-            SRV_INF("loading draft model '%s'\n", params.speculative.model.c_str());
-
-            auto params_dft = params;
-
-            params_dft.model        = params.speculative.model;
-            params_dft.n_ctx        = params.speculative.n_ctx;
-            params_dft.n_gpu_layers = params.speculative.n_gpu_layers;
-            params_dft.use_mlock    = true;
-            params_dft.n_world      = 1;  // do not split the draft model across devicesAdd commentMore actions
-            params_dft.rank         = 0;  // always load the draft model on the head device
-
-            std::fill_n(params_dft.n_layer_window, params.n_world, 0);
-            
-            llama_init_result llama_init_dft = llama_init_from_gpt_params(params_dft);
-
-            model_dft = llama_init_dft.model;
-
-            if (model_dft == nullptr) {
-                SRV_ERR("failed to load draft model, '%s'\n", params.speculative.model.c_str());
-                return false;
-            }
-
+        if (!params.speculative.model.empty()){
             if (!common_speculative_are_compatible(ctx, llama_init_dft.context)) {
                 SRV_ERR("the draft model '%s' is not compatible with the target model '%s'\n", params.speculative.model.c_str(), params.model.c_str());
 
@@ -750,14 +762,6 @@ struct server_context {
 
                 return false;
             }
-            
-            cparams_dft = llama_context_params_from_gpt_params(params);
-            cparams_dft.n_batch = llama_n_ctx(llama_init_dft.context);
-            cparams_dft.n_world = 1;
-            cparams_dft.rank    = 0;
-            std::fill_n(cparams_dft.n_layer_window, 32, 0);
-            cparams_dft.n_layer_window[0] = llama_n_layer(model_dft);
-            cparams_dft.n_gpu_layers      = params.speculative.n_gpu_layers;
 
             // the context is not needed - we will create one for each slot
             llama_free(llama_init_dft.context);

From b4929d510aee558f4cd5015daee9dfa6de0d0b50 Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@gmail.com>
Date: Mon, 30 Jun 2025 04:35:59 +0000
Subject: [PATCH 37/42] fix: args in speculative

---
 common/arg.cpp                       | 22 ++++++++++++++++++----
 common/common.h                      |  1 -
 examples/speculative/speculative.cpp |  2 +-
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/common/arg.cpp b/common/arg.cpp
index 2f999dcc..a338f613 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -627,12 +627,19 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
         }
     ).set_examples({LLAMA_EXAMPLE_SPECULATIVE}));
     add_opt(llama_arg(
-        {"--draft"}, "N",
-        format("number of tokens to draft for speculative decoding (default: %d)", params.n_draft),
+        {"--draft-max", "--draft", "--draft-n"}, "N",
+        format("number of tokens to draft for speculative decoding (default: %d)", params.speculative.n_max),
         [](gpt_params & params, int value) {
-            params.n_draft = value;
+            params.speculative.n_max = value;
         }
-    ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP}));
+    ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER}));
+    add_opt(llama_arg(
+        {"--draft-min", "--draft-n-min"}, "N",
+        format("minimum number of draft tokens to use for speculative decoding (default: %d)", params.speculative.n_min),
+        [](gpt_params & params, int value) {
+            params.speculative.n_min = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER}));
     add_opt(llama_arg(
         {"-ps", "--p-split"}, "N",
         format("speculative decoding split probability (default: %.1f)", (double)params.p_split),
@@ -640,6 +647,13 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
             params.p_split = std::stof(value);
         }
     ).set_examples({LLAMA_EXAMPLE_SPECULATIVE}));
+    add_opt(llama_arg(
+        {"--draft-p-min"}, "P",
+        format("minimum speculative decoding probability (greedy) (default: %.1f)", (double)params.speculative.p_min),
+        [](gpt_params & params, const std::string & value) {
+            params.speculative.p_min = std::stof(value);
+        }
+    ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER}));
     add_opt(llama_arg(
         {"-lcs", "--lookup-cache-static"}, "FNAME",
         "path to static lookup cache to use for lookup decoding (not updated by generation)",
diff --git a/common/common.h b/common/common.h
index 044dfdf5..b454f799 100644
--- a/common/common.h
+++ b/common/common.h
@@ -177,7 +177,6 @@ struct gpt_params {
     int32_t n_batch               =  2048; // logical batch size for prompt processing (must be >=32 to use BLAS)
     int32_t n_ubatch              =   512; // physical batch size for prompt processing (must be >=32 to use BLAS)
     int32_t n_keep                =     0; // number of tokens to keep from initial prompt
-    int32_t n_draft               =     5; // number of tokens to draft during speculative decoding
     int32_t n_chunks              =    -1; // max number of chunks to process (-1 = unlimited)
     int32_t n_parallel            =     1; // number of parallel sequences to decode
     int32_t n_sequences           =     1; // number of sequences to decode
diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp
index 0a6c4701..ad610ac8 100644
--- a/examples/speculative/speculative.cpp
+++ b/examples/speculative/speculative.cpp
@@ -169,7 +169,7 @@ int main(int argc, char ** argv) {
     const auto t_enc_end = ggml_time_us();
 
     // how many tokens to draft each time
-    int n_draft = params.n_draft;
+    int n_draft = params.speculative.n_max;
 
     int n_predict = 0;
     int n_drafted = 0;

From ca5996e7a6fbf900fa881f9d08ced2749abb8702 Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@gmail.com>
Date: Mon, 30 Jun 2025 07:31:05 +0000
Subject: [PATCH 38/42] fix: slot id

---
 examples/server/server.cpp | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/examples/server/server.cpp b/examples/server/server.cpp
index 5b970e4f..1acf1421 100644
--- a/examples/server/server.cpp
+++ b/examples/server/server.cpp
@@ -2534,10 +2534,10 @@ struct server_context {
 
                 // construct the speculation batch
                 llama_batch_clear(slot.batch_spec);
-                llama_batch_add  (slot.batch_spec, id, slot.n_past, { slot.id }, true);
+                llama_batch_add  (slot.batch_spec, id, slot.n_past, { slot.id + 1 }, true);
 
                 for (size_t i = 0; i < draft.size(); ++i) {
-                    llama_batch_add(slot.batch_spec, draft[i], slot.n_past + 1 + i, { slot.id }, true);
+                    llama_batch_add(slot.batch_spec, draft[i], slot.n_past + 1 + i, { slot.id + 1 }, true);
                 }
 
                 llama_decode(ctx, slot.batch_spec);
@@ -2551,7 +2551,8 @@ struct server_context {
                 slot.cache_tokens.push_back(id);
                 slot.cache_tokens.insert(slot.cache_tokens.end(), ids.begin(), ids.end() - 1);
 
-                llama_kv_cache_seq_rm(ctx, slot.id, slot.n_past, -1);
+                llama_kv_cache_seq_rm     (ctx, slot.id + 1, slot.n_past, -1);
+                llama_send_kv_cache_seq_rm(ctx, slot.id    , slot.n_past, -1);
 
                 for (size_t i = 0; i < ids.size(); ++i) {
                     completion_token_output result;

From 0cf87c88377b8e2bdcd853f201d6669ca7681ffa Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@foxmail.com>
Date: Sun, 6 Jul 2025 10:05:01 +0800
Subject: [PATCH 39/42] fix: set cache_prompt default to true

---
 examples/server/server.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/examples/server/server.cpp b/examples/server/server.cpp
index 1acf1421..5cbd7f38 100644
--- a/examples/server/server.cpp
+++ b/examples/server/server.cpp
@@ -127,7 +127,7 @@ struct server_task_result {
 
 struct slot_params {
     bool stream       = true;
-    bool cache_prompt = false; // remember the prompt to avoid reprocessing all prompt
+    bool cache_prompt = true; // remember the prompt to avoid reprocessing all prompt
 
     int32_t  n_keep    =  0; // number of tokens to keep from initial prompt
     int32_t  n_discard =  0; // number of tokens after n_keep that may be discarded when shifting context, 0 defaults to half
@@ -993,7 +993,7 @@ struct server_context {
         }
 
         slot.params.stream             = json_value(data, "stream",            false);
-        slot.params.cache_prompt       = json_value(data, "cache_prompt",      false);
+        slot.params.cache_prompt       = json_value(data, "cache_prompt",      true);
         slot.params.n_predict          = json_value(data, "n_predict",         json_value(data, "max_tokens", default_params.n_predict));
         slot.sparams.top_k             = json_value(data, "top_k",             default_sparams.top_k);
         slot.sparams.top_p             = json_value(data, "top_p",             default_sparams.top_p);

From b019a707b8c4480f9e4d7607fc15bcb62d76ccc0 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Sun, 13 Jul 2025 13:42:24 +0800
Subject: [PATCH 40/42] server: fix bugs

---
 common/arg.cpp             |  1 +
 examples/server/server.cpp |  4 ++++
 src/llama.cpp              | 13 ++++++++-----
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/common/arg.cpp b/common/arg.cpp
index a338f613..45954b52 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -673,6 +673,7 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
         format("size of the prompt context (default: %d, 0 = loaded from model)", params.n_ctx),
         [](gpt_params & params, int value) {
             params.n_ctx = value;
+            params.speculative.n_ctx = value;
         }
     ).set_env("LLAMA_ARG_CTX_SIZE"));
     add_opt(llama_arg(
diff --git a/examples/server/server.cpp b/examples/server/server.cpp
index 5cbd7f38..3844c886 100644
--- a/examples/server/server.cpp
+++ b/examples/server/server.cpp
@@ -760,6 +760,8 @@ struct server_context {
                 llama_free      (llama_init_dft.context);
                 llama_free_model(llama_init_dft.model);
 
+                model_dft = nullptr;
+
                 return false;
             }
 
@@ -3566,6 +3568,8 @@ int main(int argc, char ** argv) {
     LOG_INF("%s: loading model\n", __func__);
 
     if (!ctx_server.load_model(params)) {
+        char * stop_signal = nullptr;
+        llama_free_sockets(ctx_server.ctx, &stop_signal);
         clean_up();
         t.join();
         LOG_ERR("%s: exiting due to model loading error\n", __func__);
diff --git a/src/llama.cpp b/src/llama.cpp
index a57715e3..9aa9cd82 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -17878,7 +17878,7 @@ static void llama_send_meta(zmq::socket_t & socket, struct sync_meta * meta, boo
 
         if (meta->pos != nullptr) {
             send_msgs.emplace_back("pos", strlen("pos"));
-            send_msgs.emplace_back(meta->pos, meta->n_ctx * sizeof(llama_pos));
+            send_msgs.emplace_back(meta->pos, meta->n_tokens * sizeof(llama_pos));
         }
 
         if (meta->n_seq_id != nullptr) {
@@ -17986,8 +17986,8 @@ static int llama_recv_meta(zmq::socket_t & socket, struct sync_meta * meta) {
         }
 
         if (key == "pos") {
-            meta->pos = (llama_pos *) malloc(meta->n_ctx * sizeof(llama_pos));
-            std::memcpy(meta->pos, data_msg.data(), meta->n_ctx * sizeof(llama_pos));
+            meta->pos = (llama_pos *) malloc(meta->n_tokens * sizeof(llama_pos));
+            std::memcpy(meta->pos, data_msg.data(), meta->n_tokens * sizeof(llama_pos));
         }
 
         if (key == "n_seq_id") {
@@ -18304,8 +18304,8 @@ static int llama_decode_internal(
         if (meta.n_tokens > 0) {
             batch_all.n_tokens = meta.n_tokens;
             if (meta.pos != nullptr) {
-                batch_all.pos = (llama_pos *) malloc(meta.n_ctx * sizeof(llama_pos));
-                std::memcpy(batch_all.pos, meta.pos, meta.n_ctx * sizeof(llama_pos));
+                batch_all.pos = (llama_pos *) malloc(meta.n_tokens * sizeof(llama_pos));
+                std::memcpy(batch_all.pos, meta.pos, meta.n_tokens * sizeof(llama_pos));
             }
             if (meta.n_seq_id != nullptr) {
                 batch_all.n_seq_id = (int32_t *) malloc(meta.n_tokens * sizeof(int32_t));
@@ -22089,6 +22089,9 @@ void llama_model_compute_buf_size(
         // this value may vary by GPU and CUDA version, but it's lower than 400 MiB in most cases,
         // another 300 MiB is used to prevent accidental OOM.
         *gpu_buf += 700 * 1024 * 1024;
+    } else if (backend == BACKEND_METAL) {
+        // 300 MiB is used to prevent accidental OOM, e.g., automatic quantization conversion.
+        *gpu_buf += 300 * 1024 * 1024;
     }
 }
 

From 86ca21e49c1ad24fc41c513e4e13315e6726387a Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Sun, 13 Jul 2025 21:52:59 +0800
Subject: [PATCH 41/42] server: fix bugs when running speculative decoding

---
 README.md                  | 1 +
 examples/server/server.cpp | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index a1a18538..d8e5d91c 100644
--- a/README.md
+++ b/README.md
@@ -381,6 +381,7 @@ curl -X POST http://localhost:8080/v1/cancel \
 ```
 
 **9. How to use speculative decoding?**
+
 Please see "[Power prima.cpp with speculative decoding: Further speeds up by up to 80%](https://github.com/Lizonghang/prima.cpp/discussions/29)".
 
 ## ❤️ Acknowledgment
diff --git a/examples/server/server.cpp b/examples/server/server.cpp
index 3844c886..a1cfa90c 100644
--- a/examples/server/server.cpp
+++ b/examples/server/server.cpp
@@ -2542,7 +2542,7 @@ struct server_context {
                     llama_batch_add(slot.batch_spec, draft[i], slot.n_past + 1 + i, { slot.id + 1 }, true);
                 }
 
-                llama_decode(ctx, slot.batch_spec);
+                llama_decode(ctx, slot.batch_spec, true);
 
                 // the accepted tokens from the speculation
                 const auto ids = gpt_sampler_sample_and_accept_n(slot.smpl, ctx, draft);

From bdf9d8e74bc6e3b9433f655ecb60f1e537bd41e1 Mon Sep 17 00:00:00 2001
From: "Li, Zonghang" <870644199@qq.com>
Date: Thu, 17 Jul 2025 21:03:41 +0800
Subject: [PATCH 42/42] llama-server: fix k-shift when output overlength

---
 common/arg.cpp             |  7 +++++++
 examples/server/server.cpp | 12 ++++++------
 src/llama.cpp              | 16 ++++++++--------
 3 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/common/arg.cpp b/common/arg.cpp
index 45954b52..5e4bb56c 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -2037,6 +2037,13 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
             params.chat_template = value;
         }
     ).set_examples({LLAMA_EXAMPLE_MAIN, LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CHAT_TEMPLATE"));
+    add_opt(llama_arg(
+        {"-sys", "--system-prompt"}, "PROMPT",
+        "system prompt to use with model (if applicable, depending on chat template)",
+        [](gpt_params & params, const std::string & value) {
+            params.system_prompt = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_MAIN, LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SYSTEM_PROMPT"));
     add_opt(llama_arg(
         {"-sps", "--slot-prompt-similarity"}, "SIMILARITY",
         format("how much the prompt of a request must match the prompt of a slot in order to use that slot (default: %.2f, 0.0 = disabled)\n", params.slot_prompt_similarity),
diff --git a/examples/server/server.cpp b/examples/server/server.cpp
index a1cfa90c..2a899ed5 100644
--- a/examples/server/server.cpp
+++ b/examples/server/server.cpp
@@ -1013,7 +1013,7 @@ struct server_context {
         slot.sparams.mirostat_tau      = json_value(data, "mirostat_tau",      default_sparams.mirostat_tau);
         slot.sparams.mirostat_eta      = json_value(data, "mirostat_eta",      default_sparams.mirostat_eta);
         slot.sparams.penalize_nl       = json_value(data, "penalize_nl",       default_sparams.penalize_nl);
-        slot.params.n_keep             = json_value(data, "n_keep",            slot.params.n_keep);
+        slot.params.n_keep             = json_value(data, "n_keep",            params.n_keep);
         slot.params.n_discard          = json_value(data, "n_discard",         default_params.n_discard);
         slot.sparams.seed              = json_value(data, "seed",              default_sparams.seed);
         slot.sparams.n_probs           = json_value(data, "n_probs",           default_sparams.n_probs);
@@ -1215,7 +1215,8 @@ struct server_context {
 
             // assign the system KV cache to all parallel sequences
             for (int32_t i = 1; i <= params.n_parallel; ++i) {
-                llama_kv_cache_seq_cp(ctx, 0, i, -1, -1);
+                llama_kv_cache_seq_cp     (ctx, 0, i,     -1, -1);
+                llama_send_kv_cache_seq_cp(ctx, 0, i - 1, -1, -1);
             }
         }
 
@@ -2029,7 +2030,6 @@ struct server_context {
         }
 
         // apply context-shift if needed
-        // TODO: simplify and improve
         for (server_slot & slot : slots) {
             if (slot.ga_n == 1) {
                 if (slot.is_processing() && (int) system_tokens.size() + slot.n_past >= slot.n_ctx - 1) {
@@ -2204,14 +2204,14 @@ struct server_context {
                         } else {
                             if (!params.ctx_shift) {
                                 // if context shift is disabled, we make sure prompt size is smaller than KV size
-                                if ((int) system_tokens.size() + slot.n_prompt_tokens >= slot.n_ctx) {
+                                if ((int)system_tokens.size() + slot.n_prompt_tokens >= slot.n_ctx) {
                                     slot.release();
                                     send_error(slot, "the request exceeds the available context size. try increasing the context size or enable context shift", ERROR_TYPE_INVALID_REQUEST);
                                     continue;
                                 }
                             }
                             if (slot.params.n_keep < 0) {
-                                slot.params.n_keep = slot.n_prompt_tokens;
+                                slot.params.n_keep = (int)system_tokens.size() + slot.n_prompt_tokens + 3; // +3 for <think> tag
                             }
                             slot.params.n_keep = std::min(slot.n_ctx - 4, slot.params.n_keep);
 
@@ -3590,7 +3590,7 @@ int main(int argc, char ** argv) {
     }
 
     // print sample chat example to make it clear which template is used
-    LOG_INF("%s: chat template, built_in: %d, chat_example: '%s'\n", __func__, params.chat_template.empty(), llama_chat_format_example(ctx_server.model, params.chat_template).c_str());
+    // LOG_INF("%s: chat template, built_in: %d, chat_example: '%s'\n", __func__, params.chat_template.empty(), llama_chat_format_example(ctx_server.model, params.chat_template).c_str());
 
     ctx_server.queue_tasks.on_new_task(std::bind(
                 &server_context::process_single_task, &ctx_server, std::placeholders::_1));
diff --git a/src/llama.cpp b/src/llama.cpp
index 9aa9cd82..9015c550 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -19108,19 +19108,19 @@ static void llama_kv_cache_update_internal(struct llama_context & lctx) {
             GGML_ABORT("Deepseek2 does not support K-shift");
         }
 
-        for (size_t i = 0; i < lctx.sched.size(); ++i) {
-            ggml_backend_sched_reset(lctx.sched[i]);
+        auto * sched = lctx.sched.at(0);
 
-            ggml_cgraph * gf = llama_build_graph_k_shift(lctx);
+        ggml_backend_sched_reset(sched);
 
-            ggml_backend_sched_alloc_graph(lctx.sched[i], gf);
+        ggml_cgraph * gf = llama_build_graph_k_shift(lctx);
 
-            llama_set_k_shift(lctx);
+        ggml_backend_sched_alloc_graph(sched, gf);
 
-            llama_graph_compute(lctx, gf, lctx.sched[i], lctx.cparams.n_threads, lctx.threadpool);
+        llama_set_k_shift(lctx);
 
-            need_reserve = true;
-        }
+        llama_graph_compute(lctx, gf, sched, lctx.cparams.n_threads, lctx.threadpool);
+        
+        need_reserve = true;
 
         {
             auto & kv_self = lctx.kv_self;