llama : move n_vocab from llama_sampler_data to penalty_sampler (#26520)

This matches how it is done for logit_bias and mirostat samplers, see
https://github.com/ggml-org/llama.cpp/pull/25262#discussion_r3703951151

diff --git a/common/sampling.cpp b/common/sampling.cpp
index 5698c0263b24d46d6eb9b16380f37e4e3288bd69..ba5504ed0118d813c98837c7fe9aeeb10594bcb0 100644 (file)
--- a/common/sampling.cpp
+++ b/common/sampling.cpp
@@ -383,7 +383,7 @@ struct common_sampler * common_sampler_init(
                     samplers.push_back(llama_sampler_init_infill(vocab));
                     break;
                 case COMMON_SAMPLER_TYPE_PENALTIES:
-                    samplers.push_back(llama_sampler_init_penalties(params.penalty_last_n, params.penalty_repeat, params.penalty_freq, params.penalty_present));
+                    samplers.push_back(llama_sampler_init_penalties(llama_vocab_n_tokens(vocab), params.penalty_last_n, params.penalty_repeat, params.penalty_freq, params.penalty_present));
                     break;
                 case COMMON_SAMPLER_TYPE_ADAPTIVE_P:
                     // the `adaptive-p` sampler is like `dist` and `mirostat` in that it selects

diff --git a/include/llama.h b/include/llama.h
index f2d7e38858378b2843833d21d31bae310c0adddf..fb2ca38cee4f8ba84bb6178f1e345e066b0d07e5 100644 (file)
--- a/include/llama.h
+++ b/include/llama.h
@@ -1256,7 +1256,6 @@ extern "C" {
         struct ggml_tensor * probs;
         struct ggml_tensor * sampled;
         struct ggml_tensor * candidates;
-        int64_t              n_vocab;
     };
 
     // user code can implement the interface below in order to create custom llama_sampler
@@ -1425,6 +1424,7 @@ extern "C" {
 
     /// NOTE: Avoid using on the full vocabulary as searching for repeated tokens can become slow. For example, apply top-k or top-p sampling first.
     LLAMA_API struct llama_sampler * llama_sampler_init_penalties(
+                             int32_t   n_vocab,
                              int32_t   penalty_last_n,   // last n tokens to penalize (0 = disable penalty, -1 = context size)
                                float   penalty_repeat,   // must be > 0.0, 1.0 = disabled
                                float   penalty_freq,     // must be finite, 0.0 = disabled

diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 9dde345df406bfe749b14fd0dc2b711762f01b50..2be3b75fb9825ccc9aa08cda294f46d6422c61ea 100644 (file)
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -3683,7 +3683,6 @@ void llm_graph_context::build_sampling() const {
             /*.probs       =*/ nullptr,
             /*.sampled     =*/ nullptr,
             /*.candidates  =*/ nullptr,
-            /*.n_vocab     =*/ logits_seq->ne[0],
         };
 
         assert(sampler->iface->backend_apply);

diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp
index b2f1abe7378512faaeb47a26718573ba03bd96a0..6cf2d27cf9ad7e5929da4ae7684492e36796b773 100644 (file)
--- a/src/llama-sampler.cpp
+++ b/src/llama-sampler.cpp
@@ -589,7 +589,6 @@ static bool llama_sampler_backend_support(
         /*.probs      = */ nullptr,
         /*.sampled    = */ nullptr,
         /*.candidates = */ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n),
-        /*.n_vocab    = */ n,
     };
 
     ggml_cgraph * gf = ggml_new_graph(ctx);
@@ -2640,6 +2639,7 @@ struct llama_sampler * llama_sampler_init_grammar_lazy_patterns(
 // penalties
 
 struct llama_sampler_penalties : public llama_sampler_backend {
+    const int32_t n_vocab;
     const int32_t penalty_last_n;
     const float   penalty_repeat;
     const float   penalty_freq;
@@ -2655,7 +2655,6 @@ struct llama_sampler_penalties : public llama_sampler_backend {
     ggml_tensor * inp_counts    = nullptr;
 
     // backend helpers
-    int32_t n_vocab = 0;
     int32_t n_max   = 0;
     bool has_candidates = false;
 
@@ -2676,11 +2675,13 @@ struct llama_sampler_penalties : public llama_sampler_backend {
     }
 
     llama_sampler_penalties(
+            int32_t n_vocab,
             int32_t penalty_last_n,
             float   penalty_repeat,
             float   penalty_freq,
             float   penalty_present)
         : llama_sampler_backend("penalties")
+        , n_vocab         (n_vocab)
         , penalty_last_n  (penalty_last_n)
         , penalty_repeat  (penalty_repeat)
         , penalty_freq    (penalty_freq)
@@ -2766,6 +2767,7 @@ static void llama_sampler_penalties_reset(struct llama_sampler * smpl) {
 static struct llama_sampler * llama_sampler_penalties_clone(const struct llama_sampler * smpl) {
     const auto * ctx = (const llama_sampler_penalties *) smpl->ctx;
     auto * result = llama_sampler_init_penalties(
+            ctx->n_vocab,
             ctx->penalty_last_n,
             ctx->penalty_repeat,
             ctx->penalty_freq,
@@ -2811,10 +2813,9 @@ static void llama_sampler_penalties_backend_apply(
         return;
     }
 
-    GGML_ASSERT(data->n_vocab > 0 && data->n_vocab <= INT32_MAX);
+    GGML_ASSERT(sctx->n_vocab > 0);
 
     sctx->has_candidates = data->candidates != nullptr;
-    sctx->n_vocab = (int32_t) data->n_vocab;
     sctx->n_max   = std::min(sctx->penalty_last_n, sctx->n_vocab);
 
     sctx->inp_token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sctx->n_max);
@@ -2965,6 +2966,7 @@ static struct llama_sampler_i llama_sampler_penalties_i = {
 };
 
 struct llama_sampler * llama_sampler_init_penalties(
+        int32_t n_vocab,
         int32_t penalty_last_n,
         float penalty_repeat,
         float penalty_freq,
@@ -2979,6 +2981,7 @@ struct llama_sampler * llama_sampler_init_penalties(
     return llama_sampler_init(
         /* .iface = */ &llama_sampler_penalties_i,
         /* .ctx   = */ new llama_sampler_penalties(
+            n_vocab,
             penalty_last_n,
             penalty_repeat,
             penalty_freq,

diff --git a/tests/test-backend-sampler.cpp b/tests/test-backend-sampler.cpp
index 1a46468ba2641a2cb4673bf5680696ee2c2947cd..1165f46f0c92a0988ccf3b6aa0a36e959e94737f 100644 (file)
--- a/tests/test-backend-sampler.cpp
+++ b/tests/test-backend-sampler.cpp
@@ -823,6 +823,7 @@ enum class penalties_position {
 static void add_filter_and_penalties(
         llama_sampler * chain,
         const sampler_init_fn & init_filter,
+        int32_t n_vocab,
         int32_t penalty_last_n,
         float penalty_repeat,
         float penalty_freq,
@@ -830,7 +831,7 @@ static void add_filter_and_penalties(
         penalties_position position) {
     const auto add_penalties = [&]() {
         llama_sampler_chain_add(chain, llama_sampler_init_penalties(
-                    penalty_last_n, penalty_repeat, penalty_freq, penalty_present));
+                    n_vocab, penalty_last_n, penalty_repeat, penalty_freq, penalty_present));
     };
 
     if (position == penalties_position::before_filter) {
@@ -1006,7 +1007,7 @@ static sampler_comparison_output run_penalties_comparison(
     const std::vector<float> raw_logits = decode_raw_logits(params, prompt);
     const auto add_samplers = [&](llama_sampler * chain) {
         llama_sampler_chain_add(chain, llama_sampler_init_penalties(
-                    penalty_last_n, penalty_repeat, penalty_freq, penalty_present));
+                    llama_vocab_n_tokens(vocab), penalty_last_n, penalty_repeat, penalty_freq, penalty_present));
     };
     const auto accept_history = [&](llama_sampler * chain) {
         accept_prompt(chain, vocab, prompt);
@@ -1105,7 +1106,7 @@ static void compare_top_k_penalties_logits(
     GGML_ASSERT(excluded_history_token != LLAMA_TOKEN_NULL);
 
     const auto add_samplers = [&](llama_sampler * chain) {
-        add_filter_and_penalties(chain, init_top_k,
+        add_filter_and_penalties(chain, init_top_k, n_vocab,
                 penalty_last_n, penalty_repeat, penalty_freq, penalty_present, position);
     };
 
@@ -1190,7 +1191,7 @@ static void compare_masking_penalties_logits(
     GGML_ASSERT(masked_token != LLAMA_TOKEN_NULL);
 
     const auto add_samplers = [&](llama_sampler * chain) {
-        add_filter_and_penalties(chain, init_filter,
+        add_filter_and_penalties(chain, init_filter, n_vocab,
                 penalty_last_n, penalty_repeat, penalty_freq, penalty_present, position);
     };
     auto accept_history = [&](llama_sampler * smpl) {
@@ -1218,7 +1219,7 @@ static void compare_masking_penalties_logits(
             GGML_ASSERT(fabsf(expected_logits.at(penalized_token) - raw_logits[penalized_token]) > 1e-6f);
         } else {
             llama_sampler_ptr penalties(llama_sampler_init_penalties(
-                        penalty_last_n, penalty_repeat, penalty_freq, penalty_present));
+                        n_vocab, penalty_last_n, penalty_repeat, penalty_freq, penalty_present));
             accept_history(penalties.get());
             const std::unordered_map<llama_token, float> penalized_logits =
                 map_logits(apply_cpu_sampler(raw_logits, penalties.get()));

diff --git a/tests/test-sampling.cpp b/tests/test-sampling.cpp
index 2aecff90e7bb4b8c09e32ae3dab24d41ca2138f0..297f760157df24f2344b39ce66af5efe987b1aea 100644 (file)
--- a/tests/test-sampling.cpp
+++ b/tests/test-sampling.cpp
@@ -144,7 +144,7 @@ static void test_penalties(
 
     sampler_tester tester(probs, probs_expected);
 
-    auto * sampler = llama_sampler_init_penalties(last_tokens.size(), repeat_penalty, alpha_frequency, alpha_presence);
+    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]);