Repository: ggerganov/llama.cpp
File: src/llama-sampling.cpp (C++)

Code before:
    // compute the adapted target probability for the current sampling step
    auto target = std::clamp(ctx->target, 0.0f, 1.0f);
    float adapted_target = std::clamp(
        ctx->total_weight == 0.0f ? target : 2.0f * target - (ctx->weighted_sum / ctx->total_weight),
        0.0f, 1.0f
    );

    // adaptive probability transform
    //
    // quadratic near target for fine differentiation, transitioning to linear decay in the
    // tails. unbounded negative logits ensure proper suppression of far-from-target tokens
    // after the softmax.
    //
    for (size_t i = 0; i < cur_p->size; ++i) {
        float dist = std::abs((cur_p->data[i].p - adapted_target) * INV_WIDTH);
        cur_p->data[i].logit = PEAK_LOGIT_VALUE - SHARPNESS * dist * dist / (1.0f + dist);
    }

    // softmax and sample from the transformed distribution
    llama_sampler_softmax_impl(cur_p, false);
    const int idx   = llama_sample_dist(cur_p, ctx->rng);
    cur_p->selected = idx;

    // update EMA with the original probability of the selected token
    ctx->weighted_sum = ctx->original_probs[idx] + ctx->decay * ctx->weighted_sum;
    ctx->total_weight = 1.0f + ctx->decay * ctx->total_weight;
}

static void llama_sampler_adaptive_p_reset(struct llama_sampler * smpl) {
    auto * ctx = (llama_sampler_adaptive_p *) smpl->ctx;
    // ctx->target and ctx->decay never change after init, so it's safe to keep them as is.
    // original_probs is completely overwritten on every call to _apply.
    // so we only need to reset the EMA state.
    ctx->weighted_sum = ctx->target / (1.0f - ctx->decay);
    ctx->total_weight = 1.0f / (1.0f - ctx->decay);
}

static struct llama_sampler * llama_sampler_adaptive_p_clone(const struct llama_sampler * smpl) {
    const auto * ctx  = (const llama_sampler_adaptive_p *) smpl->ctx;
    auto * result     = llama_sampler_init_adaptive_p(ctx->target, ctx->decay, ctx->seed);
    auto * result_ctx = (llama_sampler_adaptive_p *) result->ctx;

    // copy everything (target, decay, and seed are already set)
    result_ctx->original_probs = ctx->original_probs;
    result_ctx->weighted_sum   = ctx->weighted_sum;
    result_ctx->total_weight   = ctx->total_weight;
    result_ctx->rng            = ctx->rng;

    return result;
}


Review comment (suggestion):
Shouldn't this logic be moved into `accept`? Some sampled tokens could get discarded, so they should not affect the sampling state.

Code after:
        0.0f, 1.0f
    );

    // adaptive probability transform
    //
    // quadratic near target for fine differentiation, transitioning to linear decay in the
    // tails. unbounded negative logits ensure proper suppression of far-from-target tokens
    // after the softmax.
    //
    for (size_t i = 0; i < cur_p->size; ++i) {
        if (cur_p->data[i].logit == -INFINITY) {
            // don't transform logits that are -INFINITY
            // (as masked out by e.g. min-p and top-p when using backend sampling)
            continue;
        }
        float dist = std::abs((cur_p->data[i].p - adapted_target) * INV_WIDTH);
        cur_p->data[i].logit = PEAK_LOGIT_VALUE - SHARPNESS * dist * dist / (1.0f + dist);
    }

    // softmax and sample from the transformed distribution
    llama_sampler_softmax_impl(cur_p, false);
    const int idx   = llama_sample_dist(cur_p, ctx->rng);
    cur_p->selected = idx;

    // store the selected token ID for acceptance later
    ctx->pending_token_id  = cur_p->data[idx].id;
    ctx->pending_token_idx = idx;
}

static void llama_sampler_adaptive_p_accept(struct llama_sampler * smpl, llama_token token) {
    auto * ctx = (llama_sampler_adaptive_p *) smpl->ctx;
    if (ctx->pending_token_id == token) {
        GGML_ASSERT(ctx->pending_token_id != LLAMA_TOKEN_NULL);
        GGML_ASSERT(ctx->pending_token_idx != -1);
        // update EMA with the original probability of the selected token
        ctx->weighted_sum = ctx->original_probs[ctx->pending_token_idx] + ctx->decay * ctx->weighted_sum;
        ctx->total_weight = 1.0f + ctx->decay * ctx->total_weight;
    }
    ctx->pending_token_id = LLAMA_TOKEN_NULL;
    ctx->pending_token_idx = -1;
}

static void llama_sampler_adaptive_p_reset(struct llama_sampler * smpl) {
    auto * ctx = (llama_sampler_adaptive_p *) smpl->ctx;
    // ctx->target and ctx->decay never change after init, so it's safe to keep them as is.
    // original_probs is completely overwritten on every call to _apply.
    // so we only need to reset the EMA state and pending token.
    ctx->weighted_sum      = ctx->target / (1.0f - ctx->decay);
    ctx->total_weight      = 1.0f / (1.0f - ctx->decay);
    ctx->pending_token_id  = LLAMA_TOKEN_NULL;
    ctx->pending_token_idx = -1;