sampling: index penalties by token id instead of scanning every candidate

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;
         }
 
-        cur_p->data[i].logit -= float(count) * ctx->penalty_freq + float(count > 0) * ctx->penalty_present;
+        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 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;
+
+    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;
+            }
+        }
+    }
+
+    // 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;

tests/test-sampling.cpp

@@ -187,6 +187,76 @@ 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<float> & probs, const std::vector<llama_token> & last_tokens,
+    float repeat_penalty, float alpha_frequency, float alpha_presence
+) {
+    auto run = [&](bool reversed) {
+        std::vector<llama_token_data> 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<float> 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<float> by_index = run(false);
+    const std::vector<float> by_lookup = run(true);
+
+    for (size_t i = 0; i < by_index.size(); i++) {
+        GGML_ASSERT(by_index[i] == by_lookup[i]);
+    }
+}
+
+// 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<float> & probs, const std::vector<llama_token> & last_tokens,
+    float repeat_penalty, float alpha_frequency, float alpha_presence
+) {
+    GGML_ASSERT(probs.size() > 1);
+
+    std::vector<llama_token_data> 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<float> & probs, const std::vector<llama_token> & last_tokens,
     const std::vector<float> & expected_probs, float dry_multiplier, float dry_base,
@@ -384,6 +454,13 @@ 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_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, {});