From 6f9626996973d662057eccaa7f4d1c18415809e4 Mon Sep 17 00:00:00 2001 From: "T. M" Date: Fri, 28 Mar 2025 23:01:30 +0000 Subject: [PATCH 01/12] llama-server : implement universal assisted decoding --- common/speculative.cpp | 90 +++++++++++++++++++++++--------------- common/speculative.h | 6 ++- examples/server/server.cpp | 11 ++--- 3 files changed, 66 insertions(+), 41 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index ccad70fa9ed85..6b3e969fe88f9 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -11,20 +11,23 @@ #define SPEC_VOCAB_CHECK_START_TOKEN_ID 5 struct common_speculative { - struct llama_context * ctx; + struct llama_context * ctx_main; + struct llama_context * ctx_dft; struct common_sampler * smpl; llama_batch batch; - llama_tokens prompt; + llama_tokens prompt_dft; }; struct common_speculative * common_speculative_init( + struct llama_context * ctx_main, 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 = */ {}, + /* .ctx_main = */ ctx_main, + /* .ctx_dft = */ ctx_dft, + /* .smpl = */ nullptr, + /* .batch = */ llama_batch_init(llama_n_batch(ctx_dft), 0, 1), + /* .prompt = */ {}, }; // TODO: optimize or pass from outside? @@ -137,27 +140,39 @@ bool common_speculative_are_compatible( llama_tokens common_speculative_gen_draft( struct common_speculative * spec, struct common_speculative_params params, - const llama_tokens & prompt_tgt, + const llama_tokens & prompt_tgt_main_model, // target model tokens llama_token id_last) { auto & batch = spec->batch; - auto & ctx = spec->ctx; + auto & ctx_main = spec->ctx_main; + auto & ctx_dft = spec->ctx_dft; auto & smpl = spec->smpl; - auto & prompt = spec->prompt; + auto & prompt_dft = spec->prompt_dft; int reuse_i = 0; int reuse_n = 0; - const int n_ctx = llama_n_ctx(ctx) - params.n_draft; + const int n_ctx = llama_n_ctx(ctx_dft) - params.n_draft; + + llama_tokens prompt_tgt_draft_model; + if (!params.vocab_dft_compatible) { + // TODO: cache detokenized prompt string + std::string detokenized = common_detokenize(ctx_main, prompt_tgt_main_model, true); + LOG_DBG("main->draft detokenized string: '%s'\n", detokenized.c_str()); + prompt_tgt_draft_model = common_tokenize(ctx_dft, detokenized, false, false); + // FIXME: token healing + } + const llama_tokens &prompt_tgt = + params.vocab_dft_compatible ? prompt_tgt_main_model : prompt_tgt_draft_model; const int i_start = std::max(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) { + for (int i = 0; i < (int) prompt_dft.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]) { + i + cur < (int) prompt_dft.size() && + prompt_tgt[i_start + cur] == prompt_dft[i + cur]) { cur++; } @@ -167,21 +182,21 @@ llama_tokens common_speculative_gen_draft( } } - LOG_DBG("%s: reuse_i = %d, reuse_n = %d, prompt = %d\n", __func__, reuse_i, reuse_n, (int) prompt.size()); + LOG_DBG("%s: reuse_i = %d, reuse_n = %d, prompt = %d\n", __func__, reuse_i, reuse_n, (int) prompt_dft.size()); llama_tokens result; result.reserve(params.n_draft); if (reuse_n == 0) { - llama_kv_self_clear(ctx); + llama_kv_self_clear(ctx_dft); - prompt.clear(); + prompt_dft.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 (reuse_i + reuse_n < (int) prompt_dft.size() && prompt_dft[reuse_i + reuse_n] == id_last) { + for (int i = reuse_i + reuse_n + 1; i < (int) prompt_dft.size(); ++i) { + result.push_back(prompt_dft[i]); if (params.n_draft <= (int) result.size()) { break; @@ -192,16 +207,16 @@ llama_tokens common_speculative_gen_draft( } if (reuse_i > 0) { - llama_kv_self_seq_rm (ctx, 0, 0, reuse_i); - llama_kv_self_seq_add(ctx, 0, reuse_i, -1, -reuse_i); + llama_kv_self_seq_rm (ctx_dft, 0, 0, reuse_i); + llama_kv_self_seq_add(ctx_dft, 0, reuse_i, -1, -reuse_i); - prompt.erase(prompt.begin(), prompt.begin() + reuse_i); + prompt_dft.erase(prompt_dft.begin(), prompt_dft.begin() + reuse_i); } - if (reuse_n < (int) prompt.size()) { - llama_kv_self_seq_rm (ctx, 0, reuse_n, -1); + if (reuse_n < (int) prompt_dft.size()) { + llama_kv_self_seq_rm (ctx_dft, 0, reuse_n, -1); - prompt.erase(prompt.begin() + reuse_n, prompt.end()); + prompt_dft.erase(prompt_dft.begin() + reuse_n, prompt_dft.end()); } } @@ -212,28 +227,28 @@ llama_tokens common_speculative_gen_draft( //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]); common_batch_add(batch, prompt_tgt[i], i - i_start, { 0 }, false); - prompt.push_back(prompt_tgt[i]); + prompt_dft.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); + llama_decode(ctx_dft, batch); } - const llama_pos n_past = prompt.size(); + const llama_pos n_past = prompt_dft.size(); LOG_DBG("%s: n_past = %d\n", __func__, n_past); common_batch_clear(batch); common_batch_add (batch, id_last, n_past, { 0 }, true); - prompt.push_back(id_last); + prompt_dft.push_back(id_last); - //LOG_DBG("%s: draft prompt: %s\n", __func__, string_from(ctx, prompt).c_str()); + LOG_DBG("%s: draft prompt: %s\n", __func__, string_from(ctx_dft, prompt_dft).c_str()); - llama_decode(ctx, batch); + llama_decode(ctx_dft, batch); common_sampler_reset(smpl); @@ -241,13 +256,13 @@ llama_tokens common_speculative_gen_draft( for (int i = 0; i < params.n_draft; ++i) { common_batch_clear(batch); - common_sampler_sample(smpl, ctx, 0, true); + common_sampler_sample(smpl, ctx_dft, 0, true); const auto * cur_p = common_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, common_token_to_piece(ctx, cur_p->data[k].id).c_str()); + k, i, cur_p->data[k].id, cur_p->data[k].p, common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); } // add drafted token for each sequence @@ -269,10 +284,15 @@ llama_tokens common_speculative_gen_draft( common_batch_add(batch, id, n_past + i + 1, { 0 }, true); // evaluate the drafted tokens on the draft model - llama_decode(ctx, batch); + llama_decode(ctx_dft, batch); - prompt.push_back(id); + prompt_dft.push_back(id); } + if (!params.vocab_dft_compatible) { + std::string detokenized = common_detokenize(ctx_dft, result, true); + LOG_DBG("draft->main detokenized string: '%s'\n", detokenized.c_str()); + result = common_tokenize(ctx_main, detokenized, false, false); + } return result; } diff --git a/common/speculative.h b/common/speculative.h index 2b51a70ca1f72..efbcd66f33cc3 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -10,9 +10,13 @@ struct common_speculative_params { int n_reuse = 256; float p_min = 0.75f; // min probability required to accept a token in the draft + bool vocab_dft_compatible = true; // whether retokenization is needed }; -struct common_speculative * common_speculative_init(struct llama_context * ctx_dft); +struct common_speculative * common_speculative_init( + struct llama_context * ctx_main, + struct llama_context * ctx_dft +); void common_speculative_free(struct common_speculative * spec); diff --git a/examples/server/server.cpp b/examples/server/server.cpp index 17a292da153c1..75e385ce1edef 100644 --- a/examples/server/server.cpp +++ b/examples/server/server.cpp @@ -1830,6 +1830,7 @@ struct server_context { llama_context * ctx = nullptr; const llama_vocab * vocab = nullptr; + bool vocab_dft_compatible = true; llama_model * model_dft = nullptr; @@ -1924,10 +1925,9 @@ struct server_context { return false; } - if (!common_speculative_are_compatible(ctx, llama_init_dft.context.get())) { - SRV_ERR("the draft model '%s' is not compatible with the target model '%s'\n", params_base.speculative.model.c_str(), params_base.model.c_str()); - - return false; + vocab_dft_compatible = common_speculative_are_compatible(ctx, llama_init_dft.context.get()); + if (!vocab_dft_compatible) { + SRV_INF("the draft model '%s' is not compatible with the target model '%s'. tokens will be translated between the draft and target models.\n", params_base.speculative.model.c_str(), params_base.model.c_str()); } const int n_ctx_dft = llama_n_ctx(llama_init_dft.context.get()); @@ -1973,7 +1973,7 @@ struct server_context { return; } - slot.spec = common_speculative_init(slot.ctx_dft); + slot.spec = common_speculative_init(slot.ctx, slot.ctx_dft); if (slot.spec == nullptr) { SRV_ERR("%s", "failed to create speculator\n"); return; @@ -3321,6 +3321,7 @@ struct server_context { params_spec.n_draft = n_draft_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; + params_spec.vocab_dft_compatible = vocab_dft_compatible; llama_tokens draft = common_speculative_gen_draft(slot.spec, params_spec, slot.cache_tokens, id); From ff9e062350e00aadfd50670fad8c8c28875144da Mon Sep 17 00:00:00 2001 From: "T. M." Date: Sat, 28 Jun 2025 16:43:42 +0000 Subject: [PATCH 02/12] Erase prompt tail for kv-cache --- common/speculative.cpp | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index d53e54a4bfa3f..f5e4e98c4ad4f 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -218,15 +218,9 @@ llama_tokens common_speculative_gen_draft( prompt_dft.erase(prompt_dft.begin(), prompt_dft.begin() + reuse_i); } - // FIXME - /* - if (reuse_n < (int) prompt.size()) { + if (reuse_n < (int) prompt_dft.size()) { llama_memory_seq_rm (mem, 0, reuse_n, -1); } - if (reuse_n < (int) prompt_dft.size()) { - llama_kv_self_seq_rm (ctx_dft, 0, reuse_n, -1); - prompt_dft.erase(prompt_dft.begin() + reuse_n, prompt_dft.end()); - } */ } // prepare a batch to evaluate any new tokens in the prompt From 39ca594a91b9a518bb4fff16979db8f8ced402f8 Mon Sep 17 00:00:00 2001 From: "T. M." Date: Sat, 28 Jun 2025 17:03:53 +0000 Subject: [PATCH 03/12] set vocab_dft_compatible in common_speculative --- common/speculative.cpp | 45 +++++++++++++++------------- common/speculative.h | 1 - examples/speculative/speculative.cpp | 43 -------------------------- 3 files changed, 24 insertions(+), 65 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index f5e4e98c4ad4f..e2f9b6515d7cf 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -17,17 +17,19 @@ struct common_speculative { llama_batch batch; llama_tokens prompt_dft; + bool vocab_dft_compatible = true; // whether retokenization is needed }; struct common_speculative * common_speculative_init( struct llama_context * ctx_main, struct llama_context * ctx_dft) { auto * result = new common_speculative { - /* .ctx_main = */ ctx_main, - /* .ctx_dft = */ ctx_dft, - /* .smpl = */ nullptr, - /* .batch = */ llama_batch_init(llama_n_batch(ctx_dft), 0, 1), - /* .prompt = */ {}, + /* .ctx_main = */ ctx_main, + /* .ctx_dft = */ ctx_dft, + /* .smpl = */ nullptr, + /* .batch = */ llama_batch_init(llama_n_batch(ctx_dft), 0, 1), + /* .prompt_dft = */ {}, + /* .vocab_dft_compatible = */ common_speculative_are_compatible(ctx_main, ctx_dft), }; // TODO: optimize or pass from outside? @@ -78,8 +80,8 @@ 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) { + 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); @@ -93,31 +95,32 @@ bool common_speculative_are_compatible( 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); + LOG_ERR("%s: draft model vocab type must match target model to use speculation but ", __func__); + LOG_ERR("vocab_type_dft = %d while vocab_type_tgt = %d\n", vocab_type_dft, vocab_type_tgt); return false; } - if (llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) || + if ( + llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) || llama_vocab_get_add_eos(vocab_tgt) != llama_vocab_get_add_eos(vocab_dft) || llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft) || - llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft)) { - LOG_ERR("%s: draft vocab special tokens must match target vocab to use speculation\n", __func__); - LOG_ERR("%s: tgt: bos = %d (%d), eos = %d (%d)\n", __func__, llama_vocab_bos(vocab_tgt), llama_vocab_get_add_bos(vocab_tgt), llama_vocab_eos(vocab_tgt), llama_vocab_get_add_eos(vocab_tgt)); - LOG_ERR("%s: dft: bos = %d (%d), eos = %d (%d)\n", __func__, llama_vocab_bos(vocab_dft), llama_vocab_get_add_bos(vocab_dft), llama_vocab_eos(vocab_dft), llama_vocab_get_add_eos(vocab_dft)); + llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_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_vocab_n_tokens(vocab_tgt); const int n_vocab_dft = llama_vocab_n_tokens(vocab_dft); - - const int vocab_diff = std::abs(n_vocab_tgt - n_vocab_dft); + const int vocab_diff = n_vocab_tgt > n_vocab_dft + ? n_vocab_tgt - n_vocab_dft + : n_vocab_dft - n_vocab_tgt; 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_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE); + LOG_ERR("%s: draft model vocab must closely match target model to use speculation but ", __func__); + LOG_ERR("target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n", + n_vocab_tgt, llama_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE); return false; } @@ -125,8 +128,8 @@ bool common_speculative_are_compatible( const char * token_text_tgt = llama_vocab_get_text(vocab_tgt, i); const char * token_text_dft = llama_vocab_get_text(vocab_dft, i); if (std::strcmp(token_text_tgt, token_text_dft) != 0) { - LOG_ERR("%s: draft vocab vocab must match target vocab to use speculation but " - "token %d content differs - target '%s', draft '%s'\n", __func__, i, + LOG_ERR("%s: draft model vocab must match target model to use speculation but ", __func__); + LOG_ERR("token %d content differs - target '%s', draft '%s'\n", i, common_token_to_piece(ctx_tgt, i).c_str(), common_token_to_piece(ctx_dft, i).c_str()); return false; diff --git a/common/speculative.h b/common/speculative.h index efbcd66f33cc3..450668f93fe86 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -10,7 +10,6 @@ struct common_speculative_params { int n_reuse = 256; float p_min = 0.75f; // min probability required to accept a token in the draft - bool vocab_dft_compatible = true; // whether retokenization is needed }; struct common_speculative * common_speculative_init( diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp index 0adffdb006bcf..a09784ac23172 100644 --- a/examples/speculative/speculative.cpp +++ b/examples/speculative/speculative.cpp @@ -99,49 +99,6 @@ int main(int argc, char ** argv) { const bool vocab_type_dft = llama_vocab_type(vocab_dft); LOG_DBG("vocab_type dft: %d\n", 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 ", __func__); - LOG_ERR("vocab_type_dft = %d while vocab_type_tgt = %d\n", vocab_type_dft, vocab_type_tgt); - return 1; - } - - if ( - llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) || - llama_vocab_get_add_eos(vocab_tgt) != llama_vocab_get_add_eos(vocab_dft) || - llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft) || - llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft) - ) { - LOG_ERR("%s: draft model special tokens must match target model to use speculation\n", __func__); - return 1; - } - - { - const int n_vocab_tgt = llama_vocab_n_tokens(vocab_tgt); - const int n_vocab_dft = llama_vocab_n_tokens(vocab_dft); - const int vocab_diff = n_vocab_tgt > n_vocab_dft - ? n_vocab_tgt - n_vocab_dft - : n_vocab_dft - n_vocab_tgt; - - if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) { - LOG_ERR("%s: draft model vocab must closely match target model to use speculation but ", __func__); - LOG_ERR("target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n", - n_vocab_tgt, llama_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE); - return 1; - } - - 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_vocab_get_text(vocab_tgt, i); - const char * token_text_dft = llama_vocab_get_text(vocab_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 ", __func__); - LOG_ERR("token %d content differs - target '%s', draft '%s'\n", i, - common_token_to_piece(ctx_tgt, i).c_str(), - common_token_to_piece(ctx_dft, i).c_str()); - return 1; - } - } - } - auto * mem_tgt = llama_get_memory(ctx_tgt); auto * mem_dft = llama_get_memory(ctx_dft); From eb424dd6ba9ef1773f5253b03e0e94e2fad11752 Mon Sep 17 00:00:00 2001 From: "T. M." Date: Sat, 28 Jun 2025 17:04:06 +0000 Subject: [PATCH 04/12] rename ctx_main to ctx_tgt --- common/speculative.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index e2f9b6515d7cf..28272467046a8 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -11,7 +11,7 @@ #define SPEC_VOCAB_CHECK_START_TOKEN_ID 5 struct common_speculative { - struct llama_context * ctx_main; + struct llama_context * ctx_tgt; struct llama_context * ctx_dft; struct common_sampler * smpl; @@ -146,7 +146,7 @@ llama_tokens common_speculative_gen_draft( const llama_tokens & prompt_tgt_main_model, // target model tokens llama_token id_last) { auto & batch = spec->batch; - auto & ctx_main = spec->ctx_main; + auto & ctx_main = spec->ctx_tgt; auto & ctx_dft = spec->ctx_dft; auto & smpl = spec->smpl; auto & prompt_dft = spec->prompt_dft; From 2550f11fe7f81248b2212fce75fd78675fd6210c Mon Sep 17 00:00:00 2001 From: "T. M." Date: Sat, 28 Jun 2025 17:17:22 +0000 Subject: [PATCH 05/12] move vocab_dft_compatible to spec struct --- common/speculative.cpp | 30 +++++++++++-------- common/speculative.h | 2 +- .../speculative-simple/speculative-simple.cpp | 4 --- tools/server/server.cpp | 1 - 4 files changed, 18 insertions(+), 19 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 28272467046a8..61b7b56851c3f 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -21,15 +21,15 @@ struct common_speculative { }; struct common_speculative * common_speculative_init( - struct llama_context * ctx_main, + struct llama_context * ctx_tgt, struct llama_context * ctx_dft) { auto * result = new common_speculative { - /* .ctx_main = */ ctx_main, + /* .ctx_main = */ ctx_tgt, /* .ctx_dft = */ ctx_dft, /* .smpl = */ nullptr, /* .batch = */ llama_batch_init(llama_n_batch(ctx_dft), 0, 1), /* .prompt_dft = */ {}, - /* .vocab_dft_compatible = */ common_speculative_are_compatible(ctx_main, ctx_dft), + /* .vocab_dft_compatible = */ false, }; // TODO: optimize or pass from outside? @@ -64,6 +64,9 @@ struct common_speculative * common_speculative_init( } #endif + result->vocab_dft_compatible = common_speculative_are_compatible(ctx_tgt, ctx_dft); + LOG_DBG("vocab_dft_compatible = %d\n", result->vocab_dft_compatible); + return result; } @@ -95,8 +98,8 @@ bool common_speculative_are_compatible( 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 ", __func__); - LOG_ERR("vocab_type_dft = %d while vocab_type_tgt = %d\n", vocab_type_dft, vocab_type_tgt); + LOG_DBG("%s: draft model vocab type must match target model to use speculation but ", __func__); + LOG_DBG("vocab_type_dft = %d while vocab_type_tgt = %d\n", vocab_type_dft, vocab_type_tgt); return false; } @@ -106,7 +109,7 @@ bool common_speculative_are_compatible( llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft) || llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft) ) { - LOG_ERR("%s: draft model special tokens must match target model to use speculation\n", __func__); + LOG_DBG("%s: draft model special tokens must match target model to use speculation\n", __func__); return false; } @@ -118,8 +121,8 @@ bool common_speculative_are_compatible( : n_vocab_dft - n_vocab_tgt; if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) { - LOG_ERR("%s: draft model vocab must closely match target model to use speculation but ", __func__); - LOG_ERR("target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n", + LOG_DBG("%s: draft model vocab must closely match target model to use speculation but ", __func__); + LOG_DBG("target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n", n_vocab_tgt, llama_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE); return false; } @@ -128,8 +131,8 @@ bool common_speculative_are_compatible( const char * token_text_tgt = llama_vocab_get_text(vocab_tgt, i); const char * token_text_dft = llama_vocab_get_text(vocab_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 ", __func__); - LOG_ERR("token %d content differs - target '%s', draft '%s'\n", i, + LOG_DBG("%s: draft model vocab must match target model to use speculation but ", __func__); + LOG_DBG("token %d content differs - target '%s', draft '%s'\n", i, common_token_to_piece(ctx_tgt, i).c_str(), common_token_to_piece(ctx_dft, i).c_str()); return false; @@ -160,7 +163,7 @@ llama_tokens common_speculative_gen_draft( const int n_ctx = llama_n_ctx(ctx_dft) - params.n_draft; llama_tokens prompt_tgt_draft_model; - if (!params.vocab_dft_compatible) { + if (!spec->vocab_dft_compatible) { // TODO: cache detokenized prompt string std::string detokenized = common_detokenize(ctx_main, prompt_tgt_main_model, true); LOG_DBG("main->draft detokenized string: '%s'\n", detokenized.c_str()); @@ -168,7 +171,7 @@ llama_tokens common_speculative_gen_draft( // FIXME: token healing } const llama_tokens &prompt_tgt = - params.vocab_dft_compatible ? prompt_tgt_main_model : prompt_tgt_draft_model; + spec->vocab_dft_compatible ? prompt_tgt_main_model : prompt_tgt_draft_model; const int i_start = std::max(0, (int) prompt_tgt.size() - n_ctx); @@ -223,6 +226,7 @@ llama_tokens common_speculative_gen_draft( if (reuse_n < (int) prompt_dft.size()) { llama_memory_seq_rm (mem, 0, reuse_n, -1); + prompt_dft.erase(prompt_dft.begin() + reuse_n, prompt_dft.end()); } } @@ -295,7 +299,7 @@ llama_tokens common_speculative_gen_draft( prompt_dft.push_back(id); } - if (!params.vocab_dft_compatible) { + if (!spec->vocab_dft_compatible) { std::string detokenized = common_detokenize(ctx_dft, result, true); LOG_DBG("draft->main detokenized string: '%s'\n", detokenized.c_str()); result = common_tokenize(ctx_main, detokenized, false, false); diff --git a/common/speculative.h b/common/speculative.h index 450668f93fe86..b7b0e82a5fea1 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -13,7 +13,7 @@ struct common_speculative_params { }; struct common_speculative * common_speculative_init( - struct llama_context * ctx_main, + struct llama_context * ctx_tgt, struct llama_context * ctx_dft ); diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index a739f236dbbd3..e3857e8c73ab8 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -64,10 +64,6 @@ int main(int argc, char ** argv) { //model_dft = llama_init_dft.model.get(); ctx_dft = llama_init_dft.context.get(); - if (!common_speculative_are_compatible(ctx_tgt, ctx_dft)) { - return 1; - } - // Tokenize the prompt std::vector inp; inp = common_tokenize(ctx_tgt, params.prompt, true, true); diff --git a/tools/server/server.cpp b/tools/server/server.cpp index da0812aa8d333..ac7d8c424ccf1 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -3554,7 +3554,6 @@ struct server_context { params_spec.n_draft = n_draft_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; - params_spec.vocab_dft_compatible = vocab_dft_compatible; const llama_tokens & cached_text_tokens = slot.cache_tokens.get_text_tokens(); llama_tokens draft = common_speculative_gen_draft(slot.spec, params_spec, cached_text_tokens, id); From 3c35c9d9457218d8d57dcba4898adb8c4533a6d8 Mon Sep 17 00:00:00 2001 From: "T. M." Date: Sat, 28 Jun 2025 17:22:56 +0000 Subject: [PATCH 06/12] clear mem_dft, remove mem --- common/speculative.cpp | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 61b7b56851c3f..692336e8bf967 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -11,7 +11,7 @@ #define SPEC_VOCAB_CHECK_START_TOKEN_ID 5 struct common_speculative { - struct llama_context * ctx_tgt; + struct llama_context * ctx_tgt; // only used for retokenizing from ctx_dft struct llama_context * ctx_dft; struct common_sampler * smpl; @@ -149,12 +149,11 @@ llama_tokens common_speculative_gen_draft( const llama_tokens & prompt_tgt_main_model, // target model tokens llama_token id_last) { auto & batch = spec->batch; - auto & ctx_main = spec->ctx_tgt; + auto & ctx_tgt = spec->ctx_tgt; auto & ctx_dft = spec->ctx_dft; auto & smpl = spec->smpl; auto & prompt_dft = spec->prompt_dft; - auto * mem = llama_get_memory(ctx_main); auto * mem_dft = llama_get_memory(ctx_dft); int reuse_i = 0; @@ -165,7 +164,7 @@ llama_tokens common_speculative_gen_draft( llama_tokens prompt_tgt_draft_model; if (!spec->vocab_dft_compatible) { // TODO: cache detokenized prompt string - std::string detokenized = common_detokenize(ctx_main, prompt_tgt_main_model, true); + std::string detokenized = common_detokenize(ctx_tgt, prompt_tgt_main_model, true); LOG_DBG("main->draft detokenized string: '%s'\n", detokenized.c_str()); prompt_tgt_draft_model = common_tokenize(ctx_dft, detokenized, false, false); // FIXME: token healing @@ -197,7 +196,6 @@ llama_tokens common_speculative_gen_draft( result.reserve(params.n_draft); if (reuse_n == 0) { - llama_memory_clear(mem, false); llama_memory_clear(mem_dft, false); prompt_dft.clear(); } else { @@ -216,16 +214,13 @@ llama_tokens common_speculative_gen_draft( } if (reuse_i > 0) { - llama_memory_seq_rm (mem, 0, 0, reuse_i); - llama_memory_seq_add(mem, 0, reuse_i, -1, -reuse_i); - llama_memory_seq_rm (mem_dft, 0, 0, reuse_i); llama_memory_seq_add(mem_dft, 0, reuse_i, -1, -reuse_i); prompt_dft.erase(prompt_dft.begin(), prompt_dft.begin() + reuse_i); } if (reuse_n < (int) prompt_dft.size()) { - llama_memory_seq_rm (mem, 0, reuse_n, -1); + llama_memory_seq_rm (mem_dft, 0, reuse_n, -1); prompt_dft.erase(prompt_dft.begin() + reuse_n, prompt_dft.end()); } } @@ -302,7 +297,7 @@ llama_tokens common_speculative_gen_draft( if (!spec->vocab_dft_compatible) { std::string detokenized = common_detokenize(ctx_dft, result, true); LOG_DBG("draft->main detokenized string: '%s'\n", detokenized.c_str()); - result = common_tokenize(ctx_main, detokenized, false, false); + result = common_tokenize(ctx_tgt, detokenized, false, false); } return result; } From 12751c9d4756f0e98b38cfa2d611d7a20d570a4c Mon Sep 17 00:00:00 2001 From: "T. M." Date: Sat, 28 Jun 2025 17:50:43 +0000 Subject: [PATCH 07/12] detokenize id_last for incompatible models --- common/speculative.cpp | 28 ++++++++++++++++++++++------ 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 692336e8bf967..37dc7430a3336 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1,5 +1,6 @@ #include "speculative.h" +#include "llama.h" #include "log.h" #include "common.h" #include "sampling.h" @@ -163,12 +164,26 @@ llama_tokens common_speculative_gen_draft( llama_tokens prompt_tgt_draft_model; if (!spec->vocab_dft_compatible) { - // TODO: cache detokenized prompt string - std::string detokenized = common_detokenize(ctx_tgt, prompt_tgt_main_model, true); - LOG_DBG("main->draft detokenized string: '%s'\n", detokenized.c_str()); - prompt_tgt_draft_model = common_tokenize(ctx_dft, detokenized, false, false); - // FIXME: token healing + const llama_model * model_tgt = llama_get_model(ctx_tgt); + + std::string text; + text = common_detokenize(ctx_tgt, prompt_tgt_main_model, false); + LOG_DBG("main->draft detokenized string: '%s'\n", text.c_str()); + prompt_tgt_draft_model = common_tokenize(ctx_dft, text, false, false); + + text.clear(); + const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt); + int32_t n_chars; + n_chars = llama_detokenize(vocab_tgt, &id_last, 1, &text[0], text.size(), false, false); + if (n_chars < 0) { + text.resize(-n_chars); + n_chars = llama_detokenize(vocab_tgt, &id_last, 1, &text[0], text.size(), false, false); + } + text.resize(n_chars); + LOG_DBG("main->draft detokenized id_last(%d): '%s'\n", id_last, text.c_str()); + id_last = common_tokenize(ctx_dft, text, false, false)[0]; } + // prompt_tgt's tokens will always be compatible with ctx_dft const llama_tokens &prompt_tgt = spec->vocab_dft_compatible ? prompt_tgt_main_model : prompt_tgt_draft_model; @@ -216,6 +231,7 @@ llama_tokens common_speculative_gen_draft( if (reuse_i > 0) { llama_memory_seq_rm (mem_dft, 0, 0, reuse_i); llama_memory_seq_add(mem_dft, 0, reuse_i, -1, -reuse_i); + prompt_dft.erase(prompt_dft.begin(), prompt_dft.begin() + reuse_i); } @@ -295,7 +311,7 @@ llama_tokens common_speculative_gen_draft( } if (!spec->vocab_dft_compatible) { - std::string detokenized = common_detokenize(ctx_dft, result, true); + std::string detokenized = common_detokenize(ctx_dft, result, false); LOG_DBG("draft->main detokenized string: '%s'\n", detokenized.c_str()); result = common_tokenize(ctx_tgt, detokenized, false, false); } From 84199317afb9cc4b1ec33451f1e86dab657101a1 Mon Sep 17 00:00:00 2001 From: "T. M." Date: Sat, 28 Jun 2025 17:51:54 +0000 Subject: [PATCH 08/12] update comment --- common/speculative.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 37dc7430a3336..784b82b41c17c 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -147,7 +147,7 @@ bool common_speculative_are_compatible( llama_tokens common_speculative_gen_draft( struct common_speculative * spec, struct common_speculative_params params, - const llama_tokens & prompt_tgt_main_model, // target model tokens + const llama_tokens & prompt_tgt_main_model, // specified in target model vocab llama_token id_last) { auto & batch = spec->batch; auto & ctx_tgt = spec->ctx_tgt; From b9fdf203f26abdd755033c1e2f5dc24c881bf7cf Mon Sep 17 00:00:00 2001 From: "T. M." Date: Sat, 28 Jun 2025 18:30:16 +0000 Subject: [PATCH 09/12] add --spec-replace flag --- common/arg.cpp | 7 ++++ common/common.h | 1 + common/speculative.cpp | 42 ++++++++++++++++++- common/speculative.h | 4 ++ .../speculative-simple/speculative-simple.cpp | 3 ++ tools/server/server.cpp | 3 ++ 6 files changed, 59 insertions(+), 1 deletion(-) diff --git a/common/arg.cpp b/common/arg.cpp index c4ad85c47b61b..c5b67515d8476 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3217,6 +3217,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.speculative.model.path = value; } ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MODEL_DRAFT")); + add_opt(common_arg( + {"--spec-replace"}, "TARGET", "DRAFT", + "translate the string in TARGET into DRAFT if the draft model and main model are not compatible", + [](common_params & params, const std::string & tgt, const std::string & dft) { + params.speculative.replacements.push_back({ tgt, dft }); + } + ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER})); add_opt(common_arg( {"-ctkd", "--cache-type-k-draft"}, "TYPE", string_format( diff --git a/common/common.h b/common/common.h index e08a59eae7543..4d19504ff93df 100644 --- a/common/common.h +++ b/common/common.h @@ -198,6 +198,7 @@ struct common_params_speculative { 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.75f; // minimum speculative decoding probability (greedy) + std::vector> replacements; // main to speculative model replacements ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V diff --git a/common/speculative.cpp b/common/speculative.cpp index 784b82b41c17c..8f97e7a056449 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -7,6 +7,7 @@ #include #include +#include #define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128 #define SPEC_VOCAB_CHECK_START_TOKEN_ID 5 @@ -19,6 +20,7 @@ struct common_speculative { llama_batch batch; llama_tokens prompt_dft; bool vocab_dft_compatible = true; // whether retokenization is needed + std::map tgt_dft_replacements = {}; }; struct common_speculative * common_speculative_init( @@ -144,6 +146,41 @@ bool common_speculative_are_compatible( return true; } +void common_speculative_add_replacement_tgt_dft( + struct common_speculative * spec, + const char *source, const char *dest) { + spec->tgt_dft_replacements[source] = dest; +} + +static std::string replace_to_dft( + struct common_speculative * spec, + const std::string& input) { + std::string result = input; + for (const auto& pair : spec->tgt_dft_replacements) { + size_t pos = result.find(pair.first); + while (pos != std::string::npos) { + result.replace(pos, pair.first.length(), pair.second); + pos = result.find(pair.first, pos + pair.second.length()); + } + } + return result; +} + +static std::string replace_to_tgt( + struct common_speculative * spec, + const std::string& input) { + std::string result = input; + for (const auto& pair : spec->tgt_dft_replacements) { + size_t pos = result.find(pair.second); + while (pos != std::string::npos) { + result.replace(pos, pair.second.length(), pair.first); + pos = result.find(pair.second, pos + pair.first.length()); + } + } + return result; +} + + llama_tokens common_speculative_gen_draft( struct common_speculative * spec, struct common_speculative_params params, @@ -168,10 +205,11 @@ llama_tokens common_speculative_gen_draft( std::string text; text = common_detokenize(ctx_tgt, prompt_tgt_main_model, false); + text = replace_to_dft(spec, text); LOG_DBG("main->draft detokenized string: '%s'\n", text.c_str()); prompt_tgt_draft_model = common_tokenize(ctx_dft, text, false, false); - text.clear(); + const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt); int32_t n_chars; n_chars = llama_detokenize(vocab_tgt, &id_last, 1, &text[0], text.size(), false, false); @@ -180,6 +218,7 @@ llama_tokens common_speculative_gen_draft( n_chars = llama_detokenize(vocab_tgt, &id_last, 1, &text[0], text.size(), false, false); } text.resize(n_chars); + text = replace_to_dft(spec, text); LOG_DBG("main->draft detokenized id_last(%d): '%s'\n", id_last, text.c_str()); id_last = common_tokenize(ctx_dft, text, false, false)[0]; } @@ -312,6 +351,7 @@ llama_tokens common_speculative_gen_draft( if (!spec->vocab_dft_compatible) { std::string detokenized = common_detokenize(ctx_dft, result, false); + detokenized = replace_to_tgt(spec, detokenized); LOG_DBG("draft->main detokenized string: '%s'\n", detokenized.c_str()); result = common_tokenize(ctx_tgt, detokenized, false, false); } diff --git a/common/speculative.h b/common/speculative.h index b7b0e82a5fea1..e69d7aaa1eb00 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -23,6 +23,10 @@ bool common_speculative_are_compatible( const struct llama_context * ctx_tgt, const struct llama_context * ctx_dft); +void common_speculative_add_replacement_tgt_dft( + struct common_speculative * spec, + const char *source, const char *dest); + // 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, diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index e3857e8c73ab8..321cf66c848ae 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -127,6 +127,9 @@ int main(int argc, char ** argv) { params_spec.p_min = p_min; struct common_speculative * spec = common_speculative_init(ctx_tgt, ctx_dft); + for (auto &pair : params.speculative.replacements) { + common_speculative_add_replacement_tgt_dft(spec, pair.first.c_str(), pair.second.c_str()); + } llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, 1); diff --git a/tools/server/server.cpp b/tools/server/server.cpp index ac7d8c424ccf1..44052227c08fe 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -2079,6 +2079,9 @@ struct server_context { SRV_ERR("%s", "failed to create speculator\n"); return; } + for (auto &pair : params_base.speculative.replacements) { + common_speculative_add_replacement_tgt_dft(slot.spec, pair.first.c_str(), pair.second.c_str()); + } } SLT_INF(slot, "new slot n_ctx_slot = %d\n", slot.n_ctx); From 160769def71f0a55d23538cb4863b3fc5d58220e Mon Sep 17 00:00:00 2001 From: "T. M." Date: Sat, 28 Jun 2025 18:37:10 +0000 Subject: [PATCH 10/12] accept special tokens when translating between draft/main models --- common/speculative.cpp | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 8f97e7a056449..99ad09af62a32 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -204,23 +204,23 @@ llama_tokens common_speculative_gen_draft( const llama_model * model_tgt = llama_get_model(ctx_tgt); std::string text; - text = common_detokenize(ctx_tgt, prompt_tgt_main_model, false); + text = common_detokenize(ctx_tgt, prompt_tgt_main_model, true); text = replace_to_dft(spec, text); LOG_DBG("main->draft detokenized string: '%s'\n", text.c_str()); - prompt_tgt_draft_model = common_tokenize(ctx_dft, text, false, false); + prompt_tgt_draft_model = common_tokenize(ctx_dft, text, false, true); text.clear(); const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt); int32_t n_chars; - n_chars = llama_detokenize(vocab_tgt, &id_last, 1, &text[0], text.size(), false, false); + n_chars = llama_detokenize(vocab_tgt, &id_last, 1, &text[0], text.size(), false, true); if (n_chars < 0) { text.resize(-n_chars); - n_chars = llama_detokenize(vocab_tgt, &id_last, 1, &text[0], text.size(), false, false); + n_chars = llama_detokenize(vocab_tgt, &id_last, 1, &text[0], text.size(), false, true); } text.resize(n_chars); text = replace_to_dft(spec, text); LOG_DBG("main->draft detokenized id_last(%d): '%s'\n", id_last, text.c_str()); - id_last = common_tokenize(ctx_dft, text, false, false)[0]; + id_last = common_tokenize(ctx_dft, text, false, true)[0]; } // prompt_tgt's tokens will always be compatible with ctx_dft const llama_tokens &prompt_tgt = @@ -350,10 +350,10 @@ llama_tokens common_speculative_gen_draft( } if (!spec->vocab_dft_compatible) { - std::string detokenized = common_detokenize(ctx_dft, result, false); + std::string detokenized = common_detokenize(ctx_dft, result, true); detokenized = replace_to_tgt(spec, detokenized); LOG_DBG("draft->main detokenized string: '%s'\n", detokenized.c_str()); - result = common_tokenize(ctx_tgt, detokenized, false, false); + result = common_tokenize(ctx_tgt, detokenized, false, true); } return result; } From d1f32abadf69ef5dab464d8811056bc9ab56af71 Mon Sep 17 00:00:00 2001 From: "T. M." Date: Sat, 28 Jun 2025 19:25:00 +0000 Subject: [PATCH 11/12] Escape spec-replace --- common/arg.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/common/arg.cpp b/common/arg.cpp index c5b67515d8476..c994f82ace110 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -977,6 +977,10 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context for (auto & seq_breaker : params.sampling.dry_sequence_breakers) { string_process_escapes(seq_breaker); } + for (auto & pair : params.speculative.replacements) { + string_process_escapes(pair.first); + string_process_escapes(pair.second); + } } if (!params.kv_overrides.empty()) { From d23892ec0fabee81c754bcd279a39843034f9d79 Mon Sep 17 00:00:00 2001 From: "T. M." Date: Wed, 2 Jul 2025 05:41:10 +0000 Subject: [PATCH 12/12] clamp draft result to size to params.n_draft --- common/speculative.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/common/speculative.cpp b/common/speculative.cpp index 99ad09af62a32..3a519de21b779 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -354,6 +354,9 @@ llama_tokens common_speculative_gen_draft( detokenized = replace_to_tgt(spec, detokenized); LOG_DBG("draft->main detokenized string: '%s'\n", detokenized.c_str()); result = common_tokenize(ctx_tgt, detokenized, false, true); + if (result.size() > (size_t)params.n_draft) { + result.resize(params.n_draft); + } } return result; }