From 0ed847d3140baead542abe3e5e6fe841013e7340 Mon Sep 17 00:00:00 2001 From: dungquixote42 <62397442+dungquixote42@users.noreply.github.com> Date: Mon, 24 Aug 2026 12:49:49 -0400 Subject: [PATCH] Adaptive P Sampler: Quality Control (#2337) * adaptive p: do not transform masked tokens * adaptive p: reuse .cum_orig_prob for .cum_cur_p * adaptive p: add fallback for cum_orig_prob==0 * adaptive p: add reset in common_sampler_reset() * adaptive p: minor --- common/sampling.cpp | 9 ++++--- include/llama.h | 11 +++----- src/llama-sampling.cpp | 60 +++++++++++++++++++++++++++--------------- src/llama-sampling.h | 12 ++++----- src/llama.cpp | 9 ++++--- 5 files changed, 58 insertions(+), 43 deletions(-) diff --git a/common/sampling.cpp b/common/sampling.cpp index 102108ac..3899167e 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -164,9 +164,7 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, co case llama_sampler_type::ADAPTIVE_P: { if (params.adaptive_target >= 0.0f) { - GGML_ASSERT(vocab); - auto n_vocab = llama_vocab_n_tokens(vocab); - result->adapt_p_ctx = llama_init_adaptive_p(n_vocab, params.adaptive_target, params.adaptive_decay, params.adaptive_updt_w_cur, result->rng()); + result->adapt_p_ctx = llama_init_adaptive_p(params.adaptive_target, params.adaptive_decay, params.adaptive_updt_w_cur, result->rng()); } break; } @@ -219,6 +217,9 @@ void common_sampler_reset(common_sampler * ctx) { // llama_grammar_reset(ctx); ctx->prev.clear(); llama_sampler_dry_reset(ctx->smpl); + + llama_free_adaptive_p(ctx->adapt_p_ctx); + ctx->adapt_p_ctx = llama_init_adaptive_p(ctx->params.adaptive_target, ctx->params.adaptive_decay, ctx->params.adaptive_updt_w_cur, ctx->rng()); } void common_sampler_review(common_sampler * ctx, const size_t n_unsent, const bool rewind_status) { @@ -547,7 +548,7 @@ static llama_token llama_sampling_sample_impl( // adaptive p sampling llama_prep_adaptive_p(ctx_main, &cur_p, ctx_sampling->adapt_p_ctx); sampler_queue(ctx_main, params, ctx_sampling, cur_p, std::max(1, params.min_keep)); - id = llama_sample_token_adaptive_p(ctx_main, &cur_p, ctx_sampling->adapt_p_ctx); + id = llama_sample_token_adaptive_p(ctx_main, &cur_p, ctx_sampling->adapt_p_ctx, ctx_sampling->rng); } else { // temperature sampling size_t min_keep = std::max(1, params.min_keep); diff --git a/include/llama.h b/include/llama.h index 6ccb146a..b1e40e73 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1520,7 +1520,7 @@ LLAMA_API struct llama_grammar* llama_sampler_init_grammar_lazy_patterns( /// @details Adaptive p sampler initializer /// @param target Select tokens near this probability (valid range 0.0 to 1.0; <0 = disabled) /// @param decay Decay rate for target adaptation over time. lower values -> faster but less stable adaptation. (valid range 0.0 to 1.0; ≤0 = no adaptation) - LLAMA_API struct llama_sampler_adaptive_p * llama_init_adaptive_p(int n_vocab, + LLAMA_API struct llama_sampler_adaptive_p * llama_init_adaptive_p( const float target, const float decay, const bool updt_w_cur, @@ -1575,12 +1575,6 @@ LLAMA_API struct llama_grammar* llama_sampler_init_grammar_lazy_patterns( struct llama_context * ctx, llama_token_data_array * candidates); - /// @details Randomly selects a token from the candidates following adaptive p sampler. - llama_token llama_sample_token_adaptive_p( - struct llama_context * ctx, - llama_token_data_array * candidates, - struct llama_sampler_adaptive_p * adapt_p_ctx); - // // Model split // @@ -1659,6 +1653,9 @@ const std::vector> & llama_internal // This is a temporary workaround in order to fix race conditions when sampling with multiple sequences. llama_token llama_sample_token_with_rng(struct llama_context * ctx, llama_token_data_array * candidates, std::mt19937 & rng); +// Randomly selects a token from the candidates following adaptive p sampler. +llama_token llama_sample_token_adaptive_p(struct llama_context * ctx, llama_token_data_array * candidates, struct llama_sampler_adaptive_p * adapt_p_ctx, std::mt19937 & rng); + #endif // LLAMA_API_INTERNAL size_t llama_fill_from_utf8(void* utf8, void* cpts, void* scripts); diff --git a/src/llama-sampling.cpp b/src/llama-sampling.cpp index 4e0f4911..fc8bee7c 100644 --- a/src/llama-sampling.cpp +++ b/src/llama-sampling.cpp @@ -1096,7 +1096,13 @@ void llama_review_adaptive_p_impl(llama_sampler_adaptive_p * adapt_p_ctx, const llama_token llama_sample_token_adaptive_p_impl( struct llama_sampling * smpl, llama_token_data_array * candidates, - struct llama_sampler_adaptive_p * adapt_p_ctx) { + struct llama_sampler_adaptive_p * adapt_p_ctx, + std::mt19937 & rng) { + if (adapt_p_ctx->cum_orig_prob == 0.f) { + LLAMA_LOG_ERROR("%s[%d]: falling back to temperature sampling\n", __func__, __LINE__); + return llama_sample_token_with_rng_impl(smpl, candidates, rng); + } + GGML_ASSERT(candidates->size > 0); const int64_t t_start_sample_us = ggml_time_us(); @@ -1121,18 +1127,20 @@ llama_token llama_sample_token_adaptive_p_impl( GGML_ASSERT(id < int(ctx->orig_prob.size())); // update history - const float update_prob = ctx->updt_w_cur - ? candidates->data[idx].p / ctx->cum_cur_p - : ctx->orig_prob[id] / ctx->cum_orig_prob; + const float update_prob = (ctx->updt_w_cur ? candidates->data[idx].p + : ctx->orig_prob[id]) + / ctx->cum_orig_prob; if (update_prob > 0) { - ctx->history.push_back({ - ctx->decay * ctx->history.back().first + update_prob, // weighted_sum - ctx->decay * ctx->history.back().second + 1.0f }); // total_weight + const auto [weighted_sum, total_weight] = ctx->history.back(); + ctx->history.push_back({ ctx->decay * weighted_sum + update_prob, + ctx->decay * total_weight + 1.0f }); } smpl->t_sample_us += ggml_time_us() - t_start_sample_us; smpl->n_sample++; + // printf("%s[%d]: idx = %zu, id = %d\n", __func__, __LINE__, idx, id); + return id; } @@ -1159,7 +1167,15 @@ void llama_sample_adaptive_p_impl(struct llama_sampling * ctx, llama_token_data_ candidates->data[i].p = prob; cum_sum += prob; } - adapt_p_ctx->cum_cur_p = cum_sum; + + if (adapt_p_ctx->updt_w_cur) { + adapt_p_ctx->cum_orig_prob = cum_sum; + } + + if (adapt_p_ctx->cum_orig_prob == 0.f) { + LLAMA_LOG_ERROR("%s[%d]: cum_orig_prob=%f is invalid\n", __func__, __LINE__, adapt_p_ctx->cum_orig_prob); + return; + } // compute adapted target probability const float weighted_sum = adapt_p_ctx->history.back().first; @@ -1182,6 +1198,10 @@ void llama_sample_adaptive_p_impl(struct llama_sampling * ctx, llama_token_data_ // unbounded negative logits suppress far-from-target tokens after softmax float max_logit = -INFINITY; for (size_t i = 0; i < candidates->size; ++i) { + if (candidates->data[i].logit <= -999.f) { + // probably masked tokens + continue; + } const float dist = std::abs(candidates->data[i].p * fused_width - fused_target); const float logit = peak_logit_value - sharpness * dist * dist / (1.0f + dist); candidates->data[i].logit = logit; @@ -1193,17 +1213,19 @@ void llama_sample_adaptive_p_impl(struct llama_sampling * ctx, llama_token_data_ ctx->t_sample_us += ggml_time_us() - t_start; } -void llama_prep_adaptive_p_impl( - struct llama_sampling * smpl, - llama_token_data_array * candidates, - struct llama_sampler_adaptive_p * adapt_p_ctx) { +void llama_prep_adaptive_p_impl(struct llama_sampling * smpl, + llama_token_data_array * candidates, + struct llama_sampler_adaptive_p * adapt_p_ctx) { + auto & orig_prob = adapt_p_ctx->orig_prob; + orig_prob.resize(candidates->size); + if (adapt_p_ctx->updt_w_cur // update with current probability, original not needed || (adapt_p_ctx->target < 0.0f)) { // or disabled return; } + constexpr float kDelta = 30.0f; //16.6f; auto t_start = ggml_time_us(); - auto & orig_prob = adapt_p_ctx->orig_prob; if (candidates->size != orig_prob.size() || candidates->sorted) { LLAMA_LOG_ERROR("%s: this function must be called before any other sampler has been applied\n", __func__); LLAMA_LOG_ERROR("%s: the sampler has been initialized with a vocabulary of %zu, but is being called with %zu candidates\n", @@ -1221,12 +1243,10 @@ void llama_prep_adaptive_p_impl( if (smpl) smpl->t_sample_us += ggml_time_us() - t_start; } -struct llama_sampler_adaptive_p * llama_init_adaptive_p_impl(int n_vocab, - const float target, - const float decay, - const bool updt_w_cur, - const uint32_t seed) { - GGML_ASSERT(n_vocab > 0); +struct llama_sampler_adaptive_p * llama_init_adaptive_p_impl(const float target, + const float decay, + const bool updt_w_cur, + const uint32_t seed) { const float clamped_decay = std::clamp(decay, 0.0f, 0.99f); auto result = new llama_sampler_adaptive_p { /* .target = */ target, @@ -1236,14 +1256,12 @@ struct llama_sampler_adaptive_p * llama_init_adaptive_p_impl(int n_vocab, /* .history = */ {}, /* .orig_prob = */ {}, /* .cum_orig_prob = */ 1.0f, - /* .cum_cur_p = */ 1.0f, /* .max_xform_logit = */ -INFINITY, /* .cum_probs = */ {}, }; result->history.push_back({ target / (1.0f - clamped_decay), // weighted_sum 1.0f / (1.0f - clamped_decay) }); // total_weight - result->orig_prob.resize(n_vocab); return result; } diff --git a/src/llama-sampling.h b/src/llama-sampling.h index 3543df86..c581eace 100644 --- a/src/llama-sampling.h +++ b/src/llama-sampling.h @@ -77,18 +77,16 @@ struct llama_sampler_adaptive_p { float cum_orig_prob; // for normalizing orig_prob in sample_token // first referenced in sample - float cum_cur_p; // cumulative sum of current probabilities float max_xform_logit; // maximum logit found during transform // first referenced in sample_token std::vector cum_probs; // cumulative probability distribution }; -struct llama_sampler_adaptive_p * llama_init_adaptive_p_impl(int n_vocab, - const float target, - const float decay, - const bool updt_w_cur, - const uint32_t seed); +struct llama_sampler_adaptive_p * llama_init_adaptive_p_impl(const float target, + const float decay, + const bool updt_w_cur, + const uint32_t seed); struct llama_sampler_adaptive_p * llama_clone_adaptive_p(const struct llama_sampler_adaptive_p * adapt_p_ctx); @@ -127,6 +125,6 @@ llama_token llama_sample_token_mirostat_v2_impl(struct llama_sampling * smpl, ll llama_token llama_sample_token_greedy_impl (struct llama_sampling * smpl, llama_token_data_array * candidates); llama_token llama_sample_token_with_rng_impl (struct llama_sampling * smpl, llama_token_data_array * candidates, std::mt19937 & rng); llama_token llama_sample_token_impl (struct llama_sampling * smpl, llama_token_data_array * candidates); -llama_token llama_sample_token_adaptive_p_impl (struct llama_sampling * smpl, llama_token_data_array * candidates, struct llama_sampler_adaptive_p * adapt_p_ctx); +llama_token llama_sample_token_adaptive_p_impl (struct llama_sampling * smpl, llama_token_data_array * candidates, struct llama_sampler_adaptive_p * adapt_p_ctx, std::mt19937 & rng); diff --git a/src/llama.cpp b/src/llama.cpp index 7945d9db..46a17d74 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -12765,8 +12765,9 @@ llama_token llama_sample_token(struct llama_context * ctx, llama_token_data_arra llama_token llama_sample_token_adaptive_p( struct llama_context * ctx, llama_token_data_array * candidates, - struct llama_sampler_adaptive_p * adapt_p_ctx) { - return llama_sample_token_adaptive_p_impl(&ctx->sampling, candidates, adapt_p_ctx); + struct llama_sampler_adaptive_p * adapt_p_ctx, + std::mt19937 & rng) { + return llama_sample_token_adaptive_p_impl(&ctx->sampling, candidates, adapt_p_ctx, rng); } int llama_split_path(char * split_path, size_t maxlen, const char * path_prefix, int split_no, int split_count) { @@ -12823,8 +12824,8 @@ void llama_sampler_dry_accept(struct llama_sampler_dry* smpl, llama_token token) } -struct llama_sampler_adaptive_p * llama_init_adaptive_p(int n_vocab, const float target, const float decay, const bool updt_w_cur, const uint32_t seed) { - return llama_init_adaptive_p_impl(n_vocab, target, decay, updt_w_cur, seed); +struct llama_sampler_adaptive_p * llama_init_adaptive_p(const float target, const float decay, const bool updt_w_cur, const uint32_t seed) { + return llama_init_adaptive_p_impl(target, decay, updt_w_cur, seed); } struct llama_sampler_adaptive_p * llama_clone_adaptive_p(const struct llama_sampler_adaptive_p * adapt_p_ctx) {