diff --git a/common/speculative.cpp b/common/speculative.cpp index c922a3f59..25b94b397 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,54 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { } } + llama_token_data_array draft_fast_topk_sample(int idx) { + llama_synchronize(params.ctx_dft); + + const float * logits = llama_get_logits_ith(params.ctx_dft, idx); + 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,11 +1358,17 @@ 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); + llama_token_data_array fast_cur_p; + const llama_token_data_array * cur_p = nullptr; + if (draft_fast_topk) { + fast_cur_p = draft_fast_topk_sample(i_last[seq_id]); + cur_p = &fast_cur_p; + } else { + common_sampler_sample(smpl, ctx_dft, i_last[seq_id], true); + cur_p = common_sampler_get_candidates(smpl, 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); - 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", seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p, @@ -1247,6 +1378,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 +1398,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;