Adaptive p: bugfix + optimization + refactor (#1155)

* adaptive-p sampler: fix zeroed orig_probs bug and refactor - Fix bug where original probabilities were captured as zero by calculating them from logits in llama_prep_adaptive_p (new). - Replace vector with unordered_map to track candidate probabilities, filtering for relevance via logit delta (16.6f). - Standardize API naming: llama_<action/verb>_<focus/name/topic>_<extra/info> - Update function signatures to follow most other samplers. * resolve merge bug * adaptive-p: revert reordering function definitions
2026-02-05 14:00:10 +00:00 · 2026-01-18 01:26:06 -05:00
parent d71a3ec315
commit 6dfbef27ec
5 changed files with 121 additions and 58 deletions
--- a/common/sampling.cpp
+++ b/common/sampling.cpp
@@ -118,7 +118,7 @@ struct llama_sampling_context * common_sampler_init(const struct llama_vocab* vo
            }
            case llama_sampler_type::ADAPTIVE_P:
            {
-                result->adapt_p_ctx=llama_sampler_init_adaptive_p(params.adaptive_target, params.adaptive_decay, result->rng());
+                result->adapt_p_ctx = llama_init_adaptive_p(params.adaptive_target, params.adaptive_decay, result->rng());
                break;
            }
            default:
@@ -423,7 +423,7 @@ static void sampler_queue(
    }
    if (use_adaptive_p) {
        // adaptive p should be put to the last, so we ignore the order in the sampler
-        llama_sample_adaptive_p(ctx_main, ctx_sampling->adapt_p_ctx, &cur_p);
+        llama_sample_adaptive_p(ctx_main, &cur_p, ctx_sampling->adapt_p_ctx);
    }
 }

@@ -471,15 +471,9 @@ static llama_token llama_sampling_sample_impl(
            id = llama_sample_token_mirostat_v2(ctx_main, &cur_p, mirostat_tau, mirostat_eta, &ctx_sampling->mirostat_mu);
        } else if (adaptive_target >= 0.0f && ctx_sampling->adapt_p_ctx!=nullptr) {
            // adaptive p sampling
-            static thread_local std::vector<float> orig_probs;
-            orig_probs.resize(cur_p.size);
-
-            // store original probabilities
-            for (size_t ii = 0; ii < cur_p.size; ++ii) {
-                orig_probs[ii] = cur_p.data[ii].p;
-            }
+            llama_prep_adaptive_p(&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, orig_probs.data());
+            id = llama_sample_token_adaptive_p(ctx_main, &cur_p, ctx_sampling->adapt_p_ctx);
        } else {
            // temperature sampling
            size_t min_keep = std::max(1, params.min_keep);
--- a/include/llama.h
+++ b/include/llama.h
@@ -1384,16 +1384,20 @@ 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_sampler_init_adaptive_p(
+    LLAMA_API struct llama_sampler_adaptive_p * llama_init_adaptive_p(
           const float target,
           const float decay,
        const uint32_t seed);

+    void llama_prep_adaptive_p(
+                 llama_token_data_array * candidates,
+        struct llama_sampler_adaptive_p * adapt_p_ctx);
+
    /// @details Adaptive p sampler described in https://github.com/MrJackSpade/adaptive-p-docs/blob/main/README.md
    void llama_sample_adaptive_p(
-            struct llama_context * ctx,
- struct llama_sampler_adaptive_p * adapt_p_ctx,
-          llama_token_data_array * candidates);
+                   struct llama_context * ctx,
+                 llama_token_data_array * candidates,
+        struct llama_sampler_adaptive_p * adapt_p_ctx);


    /// @details Mirostat 1.0 algorithm described in the paper https://arxiv.org/abs/2007.14966. Uses tokens instead of words.
@@ -1437,8 +1441,7 @@ LLAMA_API struct llama_grammar* llama_sampler_init_grammar_lazy_patterns(
    llama_token llama_sample_token_adaptive_p(
            struct llama_context * ctx,
          llama_token_data_array * candidates,
- struct llama_sampler_adaptive_p * adapt_p_ctx,
-                           float * orig_probs);
+ struct llama_sampler_adaptive_p * adapt_p_ctx);

    //
    // Model split
--- a/src/llama-sampling.cpp
+++ b/src/llama-sampling.cpp
@@ -1038,41 +1038,47 @@ struct llama_sampler_dry* llama_sampler_init_dry_impl(const struct llama_vocab&
 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,
-                              float * orig_probs)
+    struct llama_sampler_adaptive_p * adapt_p_ctx)
 {
    GGML_ASSERT(candidates->size > 0);
    const int64_t t_start_sample_us = ggml_time_us();

-    const size_t count = candidates->size;
-    adapt_p_ctx->probs.resize(count);
+    struct llama_sampler_adaptive_p * ctx = adapt_p_ctx;
+    ctx->cum_probs.resize(candidates->size);

-    // cumulative distribution
-    const float max_logit = adapt_p_ctx->max_logit;
+    // compute cumulative probability distribution
+    const float max_logit = ctx->max_xform_logit;
    float cum_prob = 0.0f;
-    for (size_t i = 0; i < count; ++i) {
+    for (size_t i = 0; i < candidates->size; ++i) {
        cum_prob += expf(candidates->data[i].logit - max_logit);
-        adapt_p_ctx->probs[i] = cum_prob;
+        ctx->cum_probs[i] = cum_prob;
    }
-    adapt_p_ctx->probs.back() += 1.0f;  // safety margin in case rng() ~= rng.max()
+    ctx->cum_probs.back() += 1.0f;  // safety margin in case rng() ~= rng.max()

-    // find token with cum_prob > target_cum_prob
-    const float target_cum_prob = cum_prob * (float)adapt_p_ctx->rng() / (float)adapt_p_ctx->rng.max();
-    auto iter = std::upper_bound(adapt_p_ctx->probs.begin(), adapt_p_ctx->probs.end(), target_cum_prob);
-    GGML_ASSERT(iter != adapt_p_ctx->probs.end());
-    llama_token id = candidates->data[std::distance(adapt_p_ctx->probs.begin(), iter)].id;
+    // select first token whose cum_prob > target_cum_prob
+    const float target_cum_prob = cum_prob * (float)ctx->rng() / (float)ctx->rng.max();
+    auto iter = std::upper_bound(ctx->cum_probs.begin(), ctx->cum_probs.end(), target_cum_prob);
+    GGML_ASSERT(iter != ctx->cum_probs.end());
+    const size_t idx = std::distance(ctx->cum_probs.begin(), iter);
+    llama_token id = candidates->data[idx].id;

    smpl->t_sample_us += ggml_time_us() - t_start_sample_us;
    smpl->n_sample++;

+    float update_prob = candidates->data[idx].p;    // not ideal
+    if (ctx->orig_prob_map.contains(id)) {
+        // selected token id is among tracked ids
+        update_prob = ctx->orig_prob_map[id] / ctx->cum_orig_prob;
+    }
+
    // update history with original probability of selected token
-    adapt_p_ctx->weighted_sum = adapt_p_ctx->decay * adapt_p_ctx->weighted_sum + orig_probs[id];
-    adapt_p_ctx->total_weight = adapt_p_ctx->decay * adapt_p_ctx->total_weight + 1.0f;
+    ctx->weighted_sum = ctx->decay * ctx->weighted_sum + update_prob;
+    ctx->total_weight = ctx->decay * ctx->total_weight + 1.0f;

    return id;
 }

-void llama_sampler_adaptive_p_apply(struct llama_sampler_adaptive_p * adapt_p_ctx, llama_token_data_array * candidates)
+void llama_sample_adaptive_p_impl(llama_token_data_array * candidates, struct llama_sampler_adaptive_p * adapt_p_ctx)
 {
    if (adapt_p_ctx->target < 0.0f) {
        // sampler is disabled
@@ -1082,14 +1088,16 @@ void llama_sampler_adaptive_p_apply(struct llama_sampler_adaptive_p * adapt_p_ct

    // incomplete softmax because final division can be fused
    float max_l = candidates->data[0].logit;
-    for (size_t i = 1; i < candidates->size; ++i) {
-        max_l = std::max(max_l, candidates->data[i].logit);
+    if (!candidates->sorted) {
+        for (size_t i = 1; i < candidates->size; ++i) {
+            max_l = std::max(max_l, candidates->data[i].logit);
+        }
    }
    float cum_sum = 0.0f;
    for (size_t i = 0; i < candidates->size; ++i) {
-        const float p = expf(candidates->data[i].logit - max_l);
-        candidates->data[i].p = p;
-        cum_sum += p;
+        const float prob = expf(candidates->data[i].logit - max_l);
+        candidates->data[i].p = prob;
+        cum_sum += prob;
    }

    // compute adapted target probability
@@ -1117,10 +1125,45 @@ void llama_sampler_adaptive_p_apply(struct llama_sampler_adaptive_p * adapt_p_ct
        max_logit = std::max(max_logit, logit);
    }
    candidates->sorted = false;
-    adapt_p_ctx->max_logit = max_logit;
+    adapt_p_ctx->max_xform_logit = max_logit;
 }

-struct llama_sampler_adaptive_p * llama_sampler_init_adaptive_p_impl(
+void llama_prep_adaptive_p_impl(
+             llama_token_data_array * candidates,
+    struct llama_sampler_adaptive_p * adapt_p_ctx)
+{
+    if (!candidates->sorted) {
+        std::sort(candidates->data, candidates->data + candidates->size,
+            [](const llama_token_data & a, const llama_token_data & b) {
+                return a.logit > b.logit;
+            });
+        candidates->sorted = true;
+    }
+    const float max_logit = candidates->data[0].logit;
+
+    // decide how many tokens to track based on logit delta
+    // i.e. do not track unlikely tokens
+    auto iter = std::lower_bound(
+        candidates->data,
+        candidates->data + candidates->size,
+        max_logit - 16.6f,  // delta
+        [](const llama_token_data & data, const float delta) {
+            return data.logit > delta;
+        });
+    const size_t n_track = std::distance(candidates->data, iter);
+
+    // store orig_prob_map and cum_orig_prob to estimate original probability later
+    float cum_prob = 0.0f;
+    adapt_p_ctx->orig_prob_map.clear();
+    for (size_t i = 0; i < n_track; ++i) {
+        const float prob = expf(candidates->data[i].logit - max_logit);
+        cum_prob += prob;
+        adapt_p_ctx->orig_prob_map[candidates->data[i].id] = prob;
+    }
+    adapt_p_ctx->cum_orig_prob = cum_prob;
+}
+
+struct llama_sampler_adaptive_p * llama_init_adaptive_p_impl(
       const float target,
       const float decay,
    const uint32_t seed)
@@ -1132,12 +1175,13 @@ struct llama_sampler_adaptive_p * llama_sampler_init_adaptive_p_impl(
        /* .rng             = */ std::mt19937(seed),
        /* .weighted_sum    = */ target / (1.0f - clamped_decay),
        /* .total_weight    = */ 1.0f / (1.0f - clamped_decay),
-        /* .max_logit       = */ 0.0f,
-        /* .probs           = */ {},
+        /* .orig_logit_map  = */ {},
+        /* .cum_orig_prob   = */ 0.0f,
+        /* .max_xform_logit = */ -INFINITY,
+        /* .cum_probs       = */ {},
    };
 }

-
 // grammar

 struct llama_sampler_grammar {
--- a/src/llama-sampling.h
+++ b/src/llama-sampling.h
@@ -61,6 +61,7 @@ struct llama_sampler_dry * llama_sampler_init_dry_impl(

 void llama_sampler_dry_apply(struct llama_sampler_dry* smpl, llama_token_data_array* cur_p);

+
 // maintains an exponential moving average of the *ORIGINAL* probabilities of selected tokens
 // used to compute an adapted target at each sampling step.
 // see llama.h for a full description of the sampler
@@ -70,15 +71,30 @@ struct llama_sampler_adaptive_p {
    std::mt19937 rng;       // RNG
    float weighted_sum;     // sum(p_n * decay^N)
    float total_weight;     // sum(decay^i), converges to 1/(1-decay)
-    float max_logit;        // maximum logit found during transform
-    std::vector<float> probs;   // cumulative probabilities
+
+    // first referenced in prep
+    std::unordered_map<llama_token, float> orig_prob_map;  // probabilities before sampler_queue
+    float cum_orig_prob;    // for normalizing orig_prob in sample_token
+
+    // first referenced in sample
+    float max_xform_logit;  // maximum logit found during transform
+
+    // first referenced in sample_token
+    std::vector<float> cum_probs;   // cumulative probability distribution
 };

-void llama_sampler_adaptive_p_apply(
-    struct llama_sampler_adaptive_p * adapt_p_ctx,
-             llama_token_data_array * candidates);
+struct llama_sampler_adaptive_p * llama_init_adaptive_p_impl(
+       const float target,
+       const float decay,
+    const uint32_t seed);

-struct llama_sampler_adaptive_p * llama_sampler_init_adaptive_p_impl(const float target, const float decay, const uint32_t seed);
+void llama_prep_adaptive_p_impl(
+             llama_token_data_array * candidates,
+    struct llama_sampler_adaptive_p * adapt_p_ctx);
+
+void llama_sample_adaptive_p_impl(
+             llama_token_data_array * candidates,
+    struct llama_sampler_adaptive_p * adapt_p_ctx);


 void llama_sample_repetition_penalties_impl(
@@ -101,6 +117,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, float * orig_probs);
+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);


--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -7689,11 +7689,18 @@ void llama_sample_dry([[maybe_unused]] struct llama_context* ctx, struct llama_s

 void llama_sample_adaptive_p(
    [[maybe_unused]] struct llama_context * ctx,
-          struct llama_sampler_adaptive_p * adapt_p_ctx,
-                   llama_token_data_array * candidates) {
-    llama_sampler_adaptive_p_apply(adapt_p_ctx, candidates);
+                   llama_token_data_array * candidates,
+          struct llama_sampler_adaptive_p * adapt_p_ctx)
+{
+    llama_sample_adaptive_p_impl(candidates, adapt_p_ctx);
 }

+void llama_prep_adaptive_p(llama_token_data_array * candidates, struct llama_sampler_adaptive_p * adapt_p_ctx)
+{
+    llama_prep_adaptive_p_impl(candidates, adapt_p_ctx);
+}
+
+
 void llama_sample_repetition_penalties(
            struct llama_context * ctx,
          llama_token_data_array * candidates,
@@ -7736,10 +7743,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,
-                              float * orig_probs)
+    struct llama_sampler_adaptive_p * adapt_p_ctx)
 {
-    return llama_sample_token_adaptive_p_impl(&ctx->sampling, candidates, adapt_p_ctx, orig_probs);
+    return llama_sample_token_adaptive_p_impl(&ctx->sampling, candidates, adapt_p_ctx);
 }

 int llama_split_path(char * split_path, size_t maxlen, const char * path_prefix, int split_no, int split_count) {
@@ -7794,9 +7800,9 @@ void llama_sampler_dry_accept(struct llama_sampler_dry* smpl, llama_token token)
 }


-struct llama_sampler_adaptive_p * llama_sampler_init_adaptive_p(const float target, const float decay, const uint32_t seed)
+struct llama_sampler_adaptive_p * llama_init_adaptive_p(const float target, const float decay, const uint32_t seed)
 {
-    return llama_sampler_init_adaptive_p_impl(target, decay, seed);
+    return llama_init_adaptive_p_impl(target, decay, seed);
 }