diff --git a/common/speculative.cpp b/common/speculative.cpp index c922a3f59..892ec8abb 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -13,8 +13,12 @@ #include #include +#include +#include +#include #include #include +#include #include #include @@ -45,6 +49,61 @@ static std::string common_speculative_get_devices_str(const std::vector std::numeric_limits::max()) { + LOG_WRN("%s: invalid %s='%s'; using default %d\n", + __func__, env_name, env, default_top_k); + return default_top_k; + } + + return (int32_t) value; +} + +static float common_speculative_mtp_draft_logit_gap_min_from_env() { + const char * env = std::getenv("LLAMA_MTP_DRAFT_LOGIT_GAP_MIN"); + if (env == nullptr || env[0] == '\0') { + return 0.0f; + } + + char * end = nullptr; + errno = 0; + const float value = std::strtof(env, &end); + if (errno != 0 || end == env || *end != '\0' || value < 0.0f) { + LOG_WRN("%s: invalid LLAMA_MTP_DRAFT_LOGIT_GAP_MIN='%s'; disabling logit-gap gate\n", + __func__, env); + return 0.0f; + } + + return value; +} + +static bool common_speculative_mtp_draft_fast_topk_from_env() { + const char * env = std::getenv("LLAMA_MTP_DRAFT_FAST_TOPK"); + if (env == nullptr || env[0] == '\0') { + return false; + } + + return strcmp(env, "1") == 0 || strcmp(env, "true") == 0 || + strcmp(env, "TRUE") == 0 || strcmp(env, "yes") == 0 || + strcmp(env, "on") == 0; +} + struct common_speculative_config { common_speculative_type type; common_params_speculative params; @@ -928,6 +987,12 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { std::vector i_last; std::vector> chain_h; + std::vector fast_topk_candidates; + + int32_t draft_top_k = 10; + float draft_logit_gap_min = 0.0f; + bool draft_fast_topk = false; + int32_t n_vocab = 0; common_speculative_impl_draft_mtp(const common_params_speculative & params, uint32_t n_seq) : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq) @@ -958,11 +1023,23 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { // TODO: fix, how to call without malloc batch.token = (llama_token *) malloc(sizeof(llama_token) * n_b); + draft_top_k = common_speculative_mtp_draft_top_k_from_env(); + LOG_INF("%s: - draft_top_k=%d (LLAMA_MTP_DRAFT_TOP_K)\n", __func__, draft_top_k); + draft_logit_gap_min = common_speculative_mtp_draft_logit_gap_min_from_env(); + LOG_INF("%s: - draft_logit_gap_min=%.3f (LLAMA_MTP_DRAFT_LOGIT_GAP_MIN)\n", + __func__, draft_logit_gap_min); + draft_fast_topk = common_speculative_mtp_draft_fast_topk_from_env() && + !this->params.backend_sampling; + LOG_INF("%s: - draft_fast_topk=%d (LLAMA_MTP_DRAFT_FAST_TOPK)\n", + __func__, draft_fast_topk ? 1 : 0); + n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(llama_get_model(ctx_dft))); + fast_topk_candidates.resize(std::max(1, draft_top_k)); + smpls.resize(n_seq); for (auto & s : smpls) { common_params_sampling sparams; sparams.no_perf = false; - sparams.top_k = 10; + sparams.top_k = draft_top_k; sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K }; s.reset(common_sampler_init(llama_get_model(ctx_dft), sparams)); } @@ -972,7 +1049,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { if (this->params.backend_sampling) { for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { llama_sampler * chain = llama_sampler_chain_init(llama_sampler_chain_default_params()); - llama_sampler_chain_add(chain, llama_sampler_init_top_k(10)); + llama_sampler_chain_add(chain, llama_sampler_init_top_k(draft_top_k)); if (!llama_set_sampler(ctx_dft, seq_id, chain)) { LOG_WRN("%s: backend offload failed for seq_id=%d; using CPU sampler\n", __func__, (int) seq_id); @@ -1046,6 +1123,51 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { } } + llama_token_data_array draft_fast_topk_sample(const float * logits) { + GGML_ASSERT(logits != nullptr); + + const int32_t k_max = std::min(draft_top_k, n_vocab); + GGML_ASSERT(k_max > 0); + int32_t k_size = 0; + + for (llama_token token_id = 0; token_id < n_vocab; ++token_id) { + const float logit = logits[token_id]; + if (k_size < k_max) { + int32_t pos = k_size++; + while (pos > 0 && logit > fast_topk_candidates[pos - 1].logit) { + fast_topk_candidates[pos] = fast_topk_candidates[pos - 1]; + --pos; + } + fast_topk_candidates[pos] = llama_token_data{token_id, logit, 0.0f}; + continue; + } + + if (logit <= fast_topk_candidates[k_max - 1].logit) { + continue; + } + + int32_t pos = k_max - 1; + while (pos > 0 && logit > fast_topk_candidates[pos - 1].logit) { + fast_topk_candidates[pos] = fast_topk_candidates[pos - 1]; + --pos; + } + fast_topk_candidates[pos] = llama_token_data{token_id, logit, 0.0f}; + } + + const float max_logit = fast_topk_candidates[0].logit; + double denom = 0.0; + for (int32_t k = 0; k < k_size; ++k) { + denom += std::exp((double) fast_topk_candidates[k].logit - (double) max_logit); + } + const double inv_denom = denom > 0.0 ? 1.0 / denom : 0.0; + for (int32_t k = 0; k < k_size; ++k) { + fast_topk_candidates[k].p = + (float) (std::exp((double) fast_topk_candidates[k].logit - (double) max_logit) * inv_denom); + } + + return llama_token_data_array{fast_topk_candidates.data(), (size_t) k_size, 0, true}; + } + bool process(const llama_batch & batch_in) override { if (batch_in.n_tokens <= 0) { return true; @@ -1233,10 +1355,22 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { auto * smpl = smpls[seq_id].get(); - common_sampler_sample(smpl, ctx_dft, i_last[seq_id], true); - const float * h_row = llama_get_embeddings_nextn_ith(ctx_dft, i_last[seq_id]); - - const auto * cur_p = common_sampler_get_candidates(smpl, true); + llama_token_data_array fast_cur_p; + const llama_token_data_array * cur_p = nullptr; + const float * h_row = nullptr; + if (draft_fast_topk) { + const llama_nextn_output_row row = + llama_get_logits_and_embeddings_nextn_ith(ctx_dft, i_last[seq_id]); + GGML_ASSERT(row.logits != nullptr); + GGML_ASSERT(row.embd_nextn != nullptr); + + fast_cur_p = draft_fast_topk_sample(row.logits); + cur_p = &fast_cur_p; + h_row = row.embd_nextn; + } else { + common_sampler_sample(smpl, ctx_dft, i_last[seq_id], true); + cur_p = common_sampler_get_candidates(smpl, true); + } for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", @@ -1247,6 +1381,18 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { // add drafted token for each sequence const llama_token id = cur_p->data[0].id; + // Optional cheap confidence gate: reject ambiguous draft steps + // using only the already materialized top-k logits. This avoids + // the full-vocab softmax recomputation tried by the raw-prob + // patch while restoring a confidence signal when top_k is tiny. + if (draft_logit_gap_min > 0.0f && cur_p->size > 1 && + cur_p->data[0].logit - cur_p->data[1].logit < draft_logit_gap_min) { + drafting[seq_id] = false; + n_drafting--; + + continue; + } + // only collect very high-confidence draft tokens if (cur_p->data[0].p < params.p_min) { drafting[seq_id] = false; @@ -1255,7 +1401,9 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { continue; } - common_sampler_accept(smpl, id, true); + if (!draft_fast_topk) { + common_sampler_accept(smpl, id, true); + } auto & dp = dparams.at(seq_id); auto & result = *dp.result; @@ -1268,6 +1416,11 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { continue; } + if (h_row == nullptr) { + h_row = llama_get_embeddings_nextn_ith(ctx_dft, i_last[seq_id]); + GGML_ASSERT(h_row != nullptr); + } + if (chain_heads) { // ref: https://github.com/ggml-org/llama.cpp/pull/24340#discussion_r3448031546 chain_h[seq_id].insert(chain_h[seq_id].end(), h_row, h_row + n_embd); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 220240ea9..e505d51e4 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3727,6 +3727,20 @@ float * llama_get_embeddings_nextn_ith(llama_context * ctx, int32_t i) { return ctx->get_embeddings_nextn_ith(i); } +llama_nextn_output_row llama_get_logits_and_embeddings_nextn_ith(llama_context * ctx, int32_t i) { + ctx->synchronize(); + + float * logits = ctx->get_sampled_logits_ith(i); + if (!logits) { + logits = ctx->get_logits_ith(i); + } + + return llama_nextn_output_row{ + logits, + ctx->get_embeddings_nextn_ith(i), + }; +} + float * llama_get_embeddings_layer_inp(llama_context * ctx, uint32_t lid) { ctx->synchronize(); diff --git a/src/llama-ext.h b/src/llama-ext.h index 348bbae95..7e0c1bfce 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -107,6 +107,19 @@ LLAMA_API float * llama_get_embeddings_nextn(struct llama_context * ctx); // LLAMA_API float * llama_get_embeddings_ith(struct llama_context * ctx, int32_t i); LLAMA_API float * llama_get_embeddings_nextn_ith(struct llama_context * ctx, int32_t i); +struct llama_nextn_output_row { + float * logits; + float * embd_nextn; +}; + +// Staging helper for speculative NextN/MTP code: synchronize once and return +// both the logits row and the corresponding NextN embedding row. Equivalent to +// calling llama_get_logits_ith() and llama_get_embeddings_nextn_ith() for the +// same row, but avoids a second public-API synchronization. +LLAMA_API llama_nextn_output_row llama_get_logits_and_embeddings_nextn_ith( + struct llama_context * ctx, + int32_t i); + // Set whether the context outputs the input embeddings of a specific layer LLAMA_API void llama_set_embeddings_layer_inp(struct llama_context * ctx, uint32_t lid, bool value);