From e2e842a46990765f26cfb5b6c5f1f8e03c3601bc Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 11 Aug 2026 19:09:45 +0000 Subject: [PATCH 1/2] sampling : index penalties by token id instead of scanning every candidate llama_sampler_penalties_apply probed token_count once per candidate, so a model with a large vocab paid n_vocab hash lookups per token while at most penalty_last_n of them can ever match. Walk token_count instead and index cur_p directly. The direct index needs a token id to be its own index, so it is guarded and falls back to the old scan when an earlier sampler has dropped or reordered candidates. Qwen3.6-27B (248320 vocab, RTX/B200, presence_penalty 1.5): 130.3 -> 146.7 t/s. Output is unchanged: same arithmetic, applied to the same tokens. --- src/llama-sampler.cpp | 48 ++++++++++++++++++++++++++++++----------- tests/test-sampling.cpp | 46 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 13 deletions(-) diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index 34a7988262ea..96053cb1cec1 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -2954,26 +2954,48 @@ static void llama_sampler_penalties_apply(struct llama_sampler * smpl, llama_tok return; } - // Apply frequency and presence penalties to the cur_p - for (size_t i = 0; i < cur_p->size; ++i) { - const auto token_iter = ctx->token_count.find(cur_p->data[i].id); - if (token_iter == ctx->token_count.end()) { - continue; - } - - const int count = token_iter->second; - + auto penalize = [ctx](llama_token_data & cand, int count) { assert(count > 0 && count <= ctx->penalty_last_n); // The academic publication that described this technique actually just only divided, but that would cause tokens with negative logits to become more likely, which is obviously wrong. // This is common fix for this problem, which is to multiply by the penalty instead of dividing. - if (cur_p->data[i].logit <= 0) { - cur_p->data[i].logit *= ctx->penalty_repeat; + if (cand.logit <= 0) { + cand.logit *= ctx->penalty_repeat; } else { - cur_p->data[i].logit /= ctx->penalty_repeat; + cand.logit /= ctx->penalty_repeat; + } + + cand.logit -= float(count) * ctx->penalty_freq + float(count > 0) * ctx->penalty_present; + }; + + // token_count holds at most penalty_last_n entries, so walking it is much cheaper than probing it once per candidate. + // This needs a token id to be its own index in cur_p, which holds while no earlier sampler has dropped or reordered candidates. + bool by_index = cur_p->size == (size_t) ctx->n_vocab; + + for (const auto & it : ctx->token_count) { + if (!by_index) { + break; } - cur_p->data[i].logit -= float(count) * ctx->penalty_freq + float(count > 0) * ctx->penalty_present; + const llama_token token = it.first; + + by_index = token >= 0 && (size_t) token < cur_p->size && cur_p->data[token].id == token; + } + + // Apply frequency and presence penalties to the cur_p + if (by_index) { + for (const auto & it : ctx->token_count) { + penalize(cur_p->data[it.first], it.second); + } + } else { + for (size_t i = 0; i < cur_p->size; ++i) { + const auto token_iter = ctx->token_count.find(cur_p->data[i].id); + if (token_iter == ctx->token_count.end()) { + continue; + } + + penalize(cur_p->data[i], token_iter->second); + } } cur_p->sorted = false; diff --git a/tests/test-sampling.cpp b/tests/test-sampling.cpp index d727ab632afb..90774d09066d 100644 --- a/tests/test-sampling.cpp +++ b/tests/test-sampling.cpp @@ -187,6 +187,48 @@ static void test_penalties( tester.check(); } +// penalties indexes cur_p by token id when it can, so a reordered cur_p must still give the same logits +static void test_penalties_reordered( + const std::vector & probs, const std::vector & last_tokens, + float repeat_penalty, float alpha_frequency, float alpha_presence +) { + auto run = [&](bool reversed) { + std::vector cur; + cur.reserve(probs.size()); + for (llama_token token_id = 0; token_id < (llama_token) probs.size(); token_id++) { + cur.emplace_back(llama_token_data{token_id, logf(probs[token_id]), probs[token_id]}); + } + + if (reversed) { + std::reverse(cur.begin(), cur.end()); + } + + llama_token_data_array cur_p = { cur.data(), cur.size(), -1, false }; + + auto * sampler = llama_sampler_init_penalties((int32_t) probs.size(), (int32_t) last_tokens.size(), repeat_penalty, alpha_frequency, alpha_presence); + for (size_t i = 0; i < last_tokens.size(); i++) { + llama_sampler_accept(sampler, last_tokens[i]); + } + + llama_sampler_apply(sampler, &cur_p); + llama_sampler_free(sampler); + + std::vector logits(probs.size()); + for (size_t i = 0; i < cur_p.size; i++) { + logits[cur_p.data[i].id] = cur_p.data[i].logit; + } + + return logits; + }; + + const std::vector by_index = run(false); + const std::vector by_lookup = run(true); + + for (size_t i = 0; i < by_index.size(); i++) { + GGML_ASSERT(by_index[i] == by_lookup[i]); + } +} + static void test_dry( const std::vector & probs, const std::vector & last_tokens, const std::vector & expected_probs, float dry_multiplier, float dry_base, @@ -384,6 +426,10 @@ int main(void) { test_penalties({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2}, {0.000023f, 0.000023f, 0.000023f, 0.499966f, 0.499966f}, 1.0f, 5.0f, 5.0f); test_penalties({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2, 0, 0}, {0.000000f, 0.000023f, 0.000023f, 0.499977f, 0.499977f}, 1.0f, 5.0f, 5.0f); + test_penalties_reordered({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2, 0, 0}, 50.0f, 0.0f, 0.0f); + test_penalties_reordered({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2, 0, 0}, 1.0f, 0.0f, 1.5f); + test_penalties_reordered({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2, 0, 0}, 1.1f, 5.0f, 5.0f); + test_dry({0.25f, 0.25f, 0.25f, 0.25f}, {0, 1}, {0.25f, 0.25f, 0.25f, 0.25f}, 1.0f, 1.1f, 2, 4, {}); test_dry({0.25f, 0.25f, 0.25f, 0.25f}, {0, 1, 2, 0, 1}, {0.296923f, 0.296923f, 0.109232f, 0.296923f}, 1.0f, 1.1f, 2, 5, {}); From 3db8cb5b2e9bf291057b9f19960e8601a162da81 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 11 Aug 2026 20:20:55 +0000 Subject: [PATCH 2/2] sampling : require the full identity layout before indexing by token id The first version checked only that each penalized token sat at its own index. A cur_p holding the same id twice then took the by-index path and penalized one of the two entries, where the scan penalized both. llama.cpp's own samplers never build such an array, but the old behaviour should not change for anyone who does. Check the whole array is the identity layout instead, which also rules out duplicates. It is a compare per candidate against a map probe per candidate, and it stops at the first mismatch. Found by a differential fuzz over 33600 cases (vocab sizes, penalty settings, reversed / shuffled / truncated / duplicated / sorted layouts, inf and nan logits): 209 mismatches before, 0 after. --- src/llama-sampler.cpp | 16 ++++++++-------- tests/test-sampling.cpp | 31 +++++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index 96053cb1cec1..2b5b282b88b3 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -2969,17 +2969,17 @@ static void llama_sampler_penalties_apply(struct llama_sampler * smpl, llama_tok }; // token_count holds at most penalty_last_n entries, so walking it is much cheaper than probing it once per candidate. - // This needs a token id to be its own index in cur_p, which holds while no earlier sampler has dropped or reordered candidates. + // This needs cur_p to still be the untouched candidate array, where a token id is its own index and appears once. + // The check is a compare per candidate, against a map probe per candidate, and it stops at the first mismatch. bool by_index = cur_p->size == (size_t) ctx->n_vocab; - for (const auto & it : ctx->token_count) { - if (!by_index) { - break; + if (by_index) { + for (size_t i = 0; i < cur_p->size; ++i) { + if (cur_p->data[i].id != (llama_token) i) { + by_index = false; + break; + } } - - const llama_token token = it.first; - - by_index = token >= 0 && (size_t) token < cur_p->size && cur_p->data[token].id == token; } // Apply frequency and presence penalties to the cur_p diff --git a/tests/test-sampling.cpp b/tests/test-sampling.cpp index 90774d09066d..d73b392e4912 100644 --- a/tests/test-sampling.cpp +++ b/tests/test-sampling.cpp @@ -229,6 +229,34 @@ static void test_penalties_reordered( } } +// a cur_p that holds the same id twice must penalize both entries, so it cannot take the by-index path +static void test_penalties_duplicate_ids( + const std::vector & probs, const std::vector & last_tokens, + float repeat_penalty, float alpha_frequency, float alpha_presence +) { + GGML_ASSERT(probs.size() > 1); + + std::vector cur; + cur.reserve(probs.size()); + for (llama_token token_id = 0; token_id < (llama_token) probs.size(); token_id++) { + cur.emplace_back(llama_token_data{token_id, logf(probs[token_id]), probs[token_id]}); + } + cur.back() = cur.front(); + + llama_token_data_array cur_p = { cur.data(), cur.size(), -1, false }; + + auto * sampler = llama_sampler_init_penalties((int32_t) probs.size(), (int32_t) last_tokens.size(), repeat_penalty, alpha_frequency, alpha_presence); + for (size_t i = 0; i < last_tokens.size(); i++) { + llama_sampler_accept(sampler, last_tokens[i]); + } + + llama_sampler_apply(sampler, &cur_p); + llama_sampler_free(sampler); + + GGML_ASSERT(cur_p.data[0].id == cur_p.data[cur_p.size - 1].id); + GGML_ASSERT(cur_p.data[0].logit == cur_p.data[cur_p.size - 1].logit); +} + static void test_dry( const std::vector & probs, const std::vector & last_tokens, const std::vector & expected_probs, float dry_multiplier, float dry_base, @@ -430,6 +458,9 @@ int main(void) { test_penalties_reordered({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2, 0, 0}, 1.0f, 0.0f, 1.5f); test_penalties_reordered({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2, 0, 0}, 1.1f, 5.0f, 5.0f); + test_penalties_duplicate_ids({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2, 0, 0}, 50.0f, 0.0f, 0.0f); + test_penalties_duplicate_ids({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2, 0, 0}, 1.0f, 0.0f, 1.5f); + test_dry({0.25f, 0.25f, 0.25f, 0.25f}, {0, 1}, {0.25f, 0.25f, 0.25f, 0.25f}, 1.0f, 1.1f, 2, 4, {}); test_dry({0.25f, 0.25f, 0.25f, 0.25f}, {0, 1, 2, 0, 1}, {0.296923f, 0.296923f, 0.109232f, 0.296923f}, 1.0f, 1.1f, 2, 5, {});