spec : disacard last drafted token with low prob (#22506)

diff --git a/common/speculative.cpp b/common/speculative.cpp
index c114bccde1aebfb43ae3cdba70818e508d1b2afb..bda9993b159d2e041a6355cfa156cf6d67e80879 100644 (file)
--- a/common/speculative.cpp
+++ b/common/speculative.cpp
@@ -467,7 +467,7 @@ struct common_speculative_state_draft : public common_speculative_state {
 
         prompt_dft.push_back(id_last);
 
-        LOG_DBG("%s: draft prompt: %s\n", __func__, string_from(ctx_dft, prompt_dft).c_str());
+        //LOG_DBG("%s: draft prompt: %s\n", __func__, string_from(ctx_dft, prompt_dft).c_str());
 
         int ret = llama_decode(ctx_dft, batch);
         if (ret != 0 && ret != 1) {
@@ -495,14 +495,14 @@ struct common_speculative_state_draft : public common_speculative_state {
 
             common_sampler_accept(smpl, id, true);
 
-            result.push_back(id);
-
-            if (sparams.n_max <= (int) result.size()) {
+            // only collect very high-confidence draft tokens
+            if (cur_p->data[0].p < sparams.p_min) {
                 break;
             }
 
-            // only collect very high-confidence draft tokens
-            if (cur_p->data[0].p < sparams.p_min) {
+            result.push_back(id);
+
+            if (sparams.n_max <= (int) result.size()) {
                 break;
             }
 

diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index e3822225bdbc263ad2f56e67497a35805582966c..ee8366d28c27f346291b16b21525e9cade4c9ee5 100644 (file)
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -354,6 +354,7 @@ struct server_slot {
 
                 // generate a new draft
                 spec_draft = common_speculative_draft(spec.get(), params_spec, tokens, sampled);
+                n_draft_total += spec_draft.size();
 
                 if (spec_draft.size() > (size_t) n_draft_max) {
                     SLT_WRN(*this, "draft size %d exceeds max %d, truncating\n", (int) spec_draft.size(), n_draft_max);
@@ -3019,7 +3020,6 @@ private:
 
                 // update how many tokens out of those tested were accepted
                 slot.n_draft_accepted += ids.size() - 1;
-                slot.n_draft_total += n_draft;
 
                 // add accepted tokens to the prompt
                 slot.prompt.tokens.keep_first(slot.prompt.n_tokens() - n_draft);