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

Code before:

// adaptive probability transformation constants
static constexpr float DISTRIBUTION_WIDTH =  0.3f;
static constexpr float PEAK_LOGIT_VALUE   =  5.0f;
static constexpr float SHARPNESS          = 10.0f;
static constexpr float INV_WIDTH          =  1.0f / DISTRIBUTION_WIDTH;

static const char * llama_sampler_adaptive_p_name(const struct llama_sampler * /*smpl*/) {
    return "adaptive-p";
}

static void llama_sampler_adaptive_p_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
    auto * ctx = (llama_sampler_adaptive_p *) smpl->ctx;

    if (ctx->target < 0.0f) {
        // at negative target values, adaptive-p is no-op
        // we simply sample from the existing distribution
        llama_sampler_softmax_impl(cur_p, false);
        cur_p->selected = llama_sample_dist(cur_p, ctx->rng);
        return;
    }

    // softmax and store the original probabilities
    llama_sampler_softmax_impl(cur_p, false);
    ctx->original_probs.resize(cur_p->size);
    for (size_t i = 0; i < cur_p->size; ++i) {
        ctx->original_probs[i] = cur_p->data[i].p;
    }

    // using the EMA, 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) {
        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);
    }

Review comment (refactor):
The `llama_sampler_softmax_impl(cur_p, false);` call can be deduplicated here - call once before the `if`

Code after:
};

// adaptive probability transformation constants
static constexpr float DISTRIBUTION_WIDTH =  0.3f;
static constexpr float PEAK_LOGIT_VALUE   =  5.0f;
static constexpr float SHARPNESS          = 10.0f;
static constexpr float INV_WIDTH          =  1.0f / DISTRIBUTION_WIDTH;

static const char * llama_sampler_adaptive_p_name(const struct llama_sampler * /*smpl*/) {
    return "adaptive-p";
}

static void llama_sampler_adaptive_p_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
    auto * ctx = (llama_sampler_adaptive_p *) smpl->ctx;

    llama_sampler_softmax_impl(cur_p, false);

    if (ctx->target < 0.0f) {
        // at negative target values, adaptive-p is no-op
        // we simply sample from the existing distribution
        cur_p->selected = llama_sample_dist(cur_p, ctx->rng);
        return;
    }

    // store the original probabilities
    ctx->original_probs.resize(cur_p->size);
    for (size_t i = 0; i < cur_p->size; ++i) {
        ctx->original_probs[i] = cur_p->data[i].p;
    }

    // using the EMA, 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) {
        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);