/*.probs = */ nullptr,
/*.sampled = */ nullptr,
/*.candidates = */ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n),
+ /*.n_vocab = */ n,
};
ggml_cgraph * gf = ggml_new_graph(ctx);
// penalties
-struct llama_sampler_penalties {
+struct llama_sampler_penalties : public llama_sampler_backend {
const int32_t penalty_last_n;
const float penalty_repeat;
const float penalty_freq;
// a frequency map to count token occurrences
std::unordered_map<llama_token, int> token_count;
+
+ // backend graph inputs
+ ggml_tensor * inp_token_ids = nullptr;
+ ggml_tensor * inp_counts = nullptr;
+
+ // backend helpers
+ int32_t n_vocab = 0;
+ int32_t n_max = 0;
+ bool has_candidates = false;
+
+ std::vector<int32_t> host_token_ids;
+ std::vector<int32_t> host_counts;
+
+ static bool is_disabled(
+ int32_t penalty_last_n,
+ float penalty_repeat,
+ float penalty_freq,
+ float penalty_present) {
+ return penalty_last_n == 0 ||
+ (penalty_repeat == 1.0f && penalty_freq == 0.0f && penalty_present == 0.0f);
+ }
+
+ bool is_disabled() const {
+ return is_disabled(penalty_last_n, penalty_repeat, penalty_freq, penalty_present);
+ }
+
+ llama_sampler_penalties(
+ int32_t penalty_last_n,
+ float penalty_repeat,
+ float penalty_freq,
+ float penalty_present)
+ : llama_sampler_backend("penalties")
+ , penalty_last_n (penalty_last_n)
+ , penalty_repeat (penalty_repeat)
+ , penalty_freq (penalty_freq)
+ , penalty_present (penalty_present)
+ , prev (penalty_last_n) {
+ }
};
-static const char * llama_sampler_penalties_name(const struct llama_sampler * /*smpl*/) {
- return "penalties";
+static const char * llama_sampler_penalties_name(const struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_penalties *) smpl->ctx;
+ return ctx->get_name();
}
static void llama_sampler_penalties_accept(struct llama_sampler * smpl, llama_token token) {
static void llama_sampler_penalties_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
auto * ctx = (llama_sampler_penalties *) smpl->ctx;
- if ((ctx->penalty_last_n == 0) ||
- (ctx->penalty_repeat == 1.0f && ctx->penalty_freq == 0.0f && ctx->penalty_present == 0.0f)) {
+ if (ctx->is_disabled()) {
return;
}
{
auto * result_ctx = (llama_sampler_penalties *) result->ctx;
- result_ctx->prev = ctx->prev;
+ result_ctx->prev = ctx->prev;
+ result_ctx->token_count = ctx->token_count;
}
return result;
delete (llama_sampler_penalties *) smpl->ctx;
}
+static bool llama_sampler_penalties_backend_init(
+ struct llama_sampler * smpl,
+ ggml_backend_buffer_type_t buft) {
+ auto * sctx = (llama_sampler_penalties *) smpl->ctx;
+
+ const bool res = llama_sampler_backend_support(smpl, buft);
+
+ sctx->init(res);
+
+ return res;
+}
+
+static void llama_sampler_penalties_backend_apply(
+ struct llama_sampler * smpl,
+ struct ggml_context * ctx,
+ struct ggml_cgraph * gf,
+ struct llama_sampler_data * data) {
+ GGML_UNUSED(gf);
+
+ auto * sctx = (llama_sampler_penalties *) smpl->ctx;
+
+ if (sctx->is_disabled()) {
+ return;
+ }
+
+ GGML_ASSERT(data->n_vocab > 0 && data->n_vocab <= INT32_MAX);
+
+ 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);
+ ggml_set_name(sctx->inp_token_ids, "penalties_token_ids");
+ ggml_set_input(sctx->inp_token_ids);
+
+ sctx->inp_counts = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sctx->n_max);
+ ggml_set_name(sctx->inp_counts, "penalties_counts");
+ ggml_set_input(sctx->inp_counts);
+
+ if ((int32_t) sctx->host_token_ids.size() != sctx->n_max) {
+ sctx->host_token_ids.assign(sctx->n_max, 0);
+ sctx->host_counts.assign(sctx->n_max, 0);
+ }
+
+ // flatten
+ ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
+ ggml_tensor * gathered = logits;
+ ggml_tensor * counts_f32 = ggml_cast(ctx, sctx->inp_counts, GGML_TYPE_F32);
+
+ if (sctx->has_candidates) {
+ ggml_tensor * candidates = ggml_reshape_1d(
+ ctx, data->candidates, ggml_nelements(data->candidates));
+ const int64_t n_candidates = candidates->ne[0];
+ GGML_ASSERT(n_candidates == ggml_nelements(logits));
+
+ ggml_tensor * counts_rows = ggml_fill(
+ ctx, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, sctx->n_vocab), 0.0f);
+ ggml_tensor * scatter_rows = ggml_reshape_2d(ctx, counts_f32, 1, sctx->n_max);
+ counts_rows = ggml_set_rows(ctx, counts_rows, scatter_rows, sctx->inp_token_ids);
+ counts_f32 = ggml_get_rows(ctx, counts_rows, candidates);
+ counts_f32 = ggml_reshape_1d(ctx, counts_f32, n_candidates);
+ } else {
+ ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, ggml_nelements(logits));
+ gathered = ggml_get_rows(ctx, logits_rows, sctx->inp_token_ids);
+ gathered = ggml_reshape_1d(ctx, gathered, sctx->n_max);
+ }
+
+ ggml_tensor * active_mask = ggml_step(ctx, counts_f32);
+ ggml_tensor * inactive_mask = ggml_sub(ctx, ggml_fill(ctx, active_mask, 1.0f), active_mask);
+
+ ggml_tensor * penalized = gathered;
+
+ if (sctx->penalty_repeat != 1.0f) {
+ ggml_tensor * pos_mask = ggml_step(ctx, penalized);
+ ggml_tensor * neg_mask = ggml_sub(ctx, ggml_fill(ctx, pos_mask, 1.0f), pos_mask);
+
+ ggml_tensor * pos_scale = ggml_scale(ctx, pos_mask, 1.0f/sctx->penalty_repeat);
+ ggml_tensor * neg_scale = ggml_scale(ctx, neg_mask, sctx->penalty_repeat);
+ ggml_tensor * repeat_scale = ggml_add(ctx, pos_scale, neg_scale);
+
+ // scale inactive entries with 1 to avoid -INF * 0 = NaN for values masked by top-p
+ repeat_scale = ggml_mul(ctx, repeat_scale, active_mask);
+ repeat_scale = ggml_add(ctx, repeat_scale, inactive_mask);
+ penalized = ggml_mul(ctx, gathered, repeat_scale);
+ }
+
+ if (sctx->penalty_freq != 0.0f) {
+ ggml_tensor * penalty_freq = ggml_scale(ctx, counts_f32, sctx->penalty_freq);
+ penalized = ggml_sub(ctx, penalized, penalty_freq);
+ }
+
+ if (sctx->penalty_present != 0.0f) {
+ ggml_tensor * penalty_present = ggml_scale(ctx, active_mask, sctx->penalty_present);
+ penalized = ggml_sub(ctx, penalized, penalty_present);
+ }
+
+ if (sctx->has_candidates) {
+ data->logits = penalized;
+ } else {
+ ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, ggml_nelements(logits));
+ ggml_tensor * scatter_rows = ggml_reshape_2d(ctx, penalized, 1, sctx->n_max);
+ logits_rows = ggml_set_rows(ctx, logits_rows, scatter_rows, sctx->inp_token_ids);
+ data->logits = ggml_reshape_1d(ctx, logits_rows, ggml_nelements(logits));
+ }
+}
+
+static void llama_sampler_penalties_backend_set_input(struct llama_sampler * smpl) {
+ auto * sctx = (llama_sampler_penalties *) smpl->ctx;
+
+ if (!sctx->inp_token_ids || !sctx->inp_counts || sctx->n_max <= 0 || sctx->n_vocab <= 0) {
+ return;
+ }
+
+ if (sctx->is_disabled()) {
+ return;
+ }
+
+ // fill active entries from the map
+ int32_t n_active = 0;
+
+ for (const auto & it : sctx->token_count) {
+ GGML_ASSERT(n_active < sctx->n_max);
+ sctx->host_token_ids[n_active] = it.first;
+ sctx->host_counts [n_active] = it.second;
+ ++n_active;
+ }
+
+ // Sorting is required because backend_apply uses ggml_set_rows (a scatter-back operation)
+ std::vector<std::pair<int32_t, int32_t>> entries;
+ entries.reserve(n_active);
+ for (int32_t i = 0; i < n_active; ++i) {
+ entries.emplace_back(sctx->host_token_ids[i], sctx->host_counts[i]);
+ }
+ std::sort(entries.begin(), entries.end(), [](const auto & a, const auto & b) {
+ return a.first < b.first;
+ });
+ for (int32_t i = 0; i < n_active; ++i) {
+ sctx->host_token_ids[i] = entries[i].first;
+ sctx->host_counts [i] = entries[i].second;
+ }
+
+ // Padding: Finds a filler token id that is not present in token_count.
+ // Use it to do padding for the arrays, it avoids resizing every time.
+ // The arrays must always have exactly n_max entries (the GPU tensor is a fixed size).
+ int32_t filler = 0;
+ if (n_active < sctx->n_max) {
+ while (sctx->token_count.find(filler) != sctx->token_count.end()) {
+ ++filler;
+ }
+ GGML_ASSERT(filler < sctx->n_vocab);
+ }
+
+ // Fill the rest of the arrays with the filler token id and count 0.
+ // Inactive slots are padded with a unique dummy token ID (count = 0).
+ // The uniqueness matters because ggml_set_rows with duplicate indices can produce non-deterministic or incorrect results.
+ // Using a filler token with count 0 that isn't in the active set is safe, because the active_mask step in backend_apply filters them out via ggml_step(counts_f32)
+ for (int32_t i = n_active; i < sctx->n_max; ++i) {
+ sctx->host_token_ids[i] = filler;
+ sctx->host_counts [i] = 0;
+ }
+
+ ggml_backend_tensor_set(sctx->inp_token_ids, sctx->host_token_ids.data(), 0, sctx->n_max * sizeof(int32_t));
+ ggml_backend_tensor_set(sctx->inp_counts, sctx->host_counts.data(), 0, sctx->n_max * sizeof(int32_t));
+}
+
static struct llama_sampler_i llama_sampler_penalties_i = {
/* .name = */ llama_sampler_penalties_name,
/* .accept = */ llama_sampler_penalties_accept,
/* .reset = */ llama_sampler_penalties_reset,
/* .clone = */ llama_sampler_penalties_clone,
/* .free = */ llama_sampler_penalties_free,
- /* .backend_init = */ nullptr,
+ /* .backend_init = */ llama_sampler_penalties_backend_init,
/* .backend_accept = */ nullptr,
- /* .backend_apply = */ nullptr,
- /* .backend_set_input = */ nullptr,
+ /* .backend_apply = */ llama_sampler_penalties_backend_apply,
+ /* .backend_set_input = */ llama_sampler_penalties_backend_set_input,
};
struct llama_sampler * llama_sampler_init_penalties(
float penalty_present) {
penalty_last_n = std::max(penalty_last_n, 0);
- const bool is_empty = (penalty_last_n == 0 || (penalty_repeat == 1.0f && penalty_freq == 0.0f && penalty_present == 0.0f));
-
- if (is_empty) {
+ if (llama_sampler_penalties::is_disabled(
+ penalty_last_n, penalty_repeat, penalty_freq, penalty_present)) {
return llama_sampler_init_empty("?penalties");
}
return llama_sampler_init(
/* .iface = */ &llama_sampler_penalties_i,
- /* .ctx = */ new llama_sampler_penalties {
- /* .penalty_last_n = */ penalty_last_n,
- /* .penalty_repeat = */ penalty_repeat,
- /* .penalty_freq = */ penalty_freq,
- /* .penalty_present = */ penalty_present,
- /* .prev = */ ring_buffer<llama_token>(penalty_last_n),
- /* .token_count = */ {},
- }
+ /* .ctx = */ new llama_sampler_penalties(
+ penalty_last_n,
+ penalty_repeat,
+ penalty_freq,
+ penalty_present)
);
}
#endif
#include <algorithm>
+#include <cmath>
#include <cstdlib>
#include <cstring>
#include <fstream>
+#include <functional>
#include <map>
#include <string>
#include <unordered_map>
+#include <unordered_set>
#include <vector>
struct test_args {
printf("backend logit bias sampling test PASSED\n");
}
+static void accept_prompt(llama_sampler * smpl, const llama_vocab * vocab, const std::string & prompt) {
+ const llama_token bos = llama_vocab_bos(vocab);
+ if (bos != LLAMA_TOKEN_NULL) {
+ llama_sampler_accept(smpl, bos);
+ }
+
+ std::vector<llama_token> tokens(64);
+ int32_t n_tokens = llama_tokenize(vocab, prompt.c_str(), (int32_t) prompt.size(),
+ tokens.data(), (int32_t) tokens.size(), false, false);
+ if (n_tokens < 0) {
+ tokens.resize(-n_tokens);
+ n_tokens = llama_tokenize(vocab, prompt.c_str(), (int32_t) prompt.size(),
+ tokens.data(), (int32_t) tokens.size(), false, false);
+ }
+
+ for (int32_t i = 0; i < n_tokens; ++i) {
+ llama_sampler_accept(smpl, tokens[i]);
+ }
+}
+
+static std::vector<float> decode_raw_logits(const test_params & params, const std::string & prompt) {
+ const int seq_id = 0;
+ const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(params.model.get()));
+ std::vector<llama_sampler_seq_config> empty_configs;
+ test_context ctx(params, empty_configs);
+
+ GGML_ASSERT(ctx.decode({{ seq_id, prompt }}));
+
+ float * logits = llama_get_logits_ith(ctx.ctx.get(), ctx.idx_for_seq(seq_id));
+ GGML_ASSERT(logits != nullptr);
+ return std::vector<float>(logits, logits + n_vocab);
+}
+
+static std::vector<llama_token_data> apply_cpu_sampler(
+ const std::vector<float> & raw_logits,
+ llama_sampler * sampler) {
+ std::vector<llama_token_data> data;
+ data.reserve(raw_logits.size());
+ for (llama_token token = 0; token < (llama_token) raw_logits.size(); ++token) {
+ data.push_back({ token, raw_logits[token], 0.0f });
+ }
+
+ llama_token_data_array cur_p = { data.data(), data.size(), -1, false };
+ llama_sampler_apply(sampler, &cur_p);
+ data.resize(cur_p.size);
+ return data;
+}
+
+using sampler_setup_fn = std::function<void(llama_sampler *)>;
+using sampler_init_fn = std::function<llama_sampler *()>;
+
+enum class penalties_position {
+ before_filter,
+ after_filter,
+};
+
+static void add_filter_and_penalties(
+ llama_sampler * chain,
+ const sampler_init_fn & init_filter,
+ int32_t penalty_last_n,
+ float penalty_repeat,
+ float penalty_freq,
+ float penalty_present,
+ 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));
+ };
+
+ if (position == penalties_position::before_filter) {
+ add_penalties();
+ llama_sampler_chain_add(chain, init_filter());
+ } else {
+ llama_sampler_chain_add(chain, init_filter());
+ add_penalties();
+ }
+}
+
+static llama_sampler_ptr make_sampler_chain(
+ const sampler_setup_fn & add_samplers,
+ const sampler_setup_fn & accept_history) {
+ llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
+ add_samplers(chain.get());
+ accept_history(chain.get());
+ return chain;
+}
+
+struct backend_sampler_output {
+ std::vector<float> logits;
+ std::vector<llama_token> candidates;
+};
+
+static backend_sampler_output run_backend_sampler(
+ const test_params & params,
+ const std::string & prompt,
+ llama_sampler * sampler) {
+ const int seq_id = 0;
+ std::vector<llama_sampler_seq_config> configs = {{ seq_id, sampler }};
+ test_context ctx(params, configs);
+
+ GGML_ASSERT(ctx.decode({{ seq_id, prompt }}));
+ llama_synchronize(ctx.ctx.get());
+
+ const int32_t idx = ctx.idx_for_seq(seq_id);
+ const uint32_t n_logits = llama_get_sampled_logits_count_ith(ctx.ctx.get(), idx);
+ const uint32_t n_candidates = llama_get_sampled_candidates_count_ith(ctx.ctx.get(), idx);
+ float * logits = llama_get_sampled_logits_ith(ctx.ctx.get(), idx);
+ llama_token * candidates = llama_get_sampled_candidates_ith(ctx.ctx.get(), idx);
+ GGML_ASSERT(logits != nullptr);
+
+ backend_sampler_output result;
+ result.logits.assign(logits, logits + n_logits);
+ result.candidates.resize(n_logits);
+
+ if (n_candidates == 0) {
+ for (uint32_t i = 0; i < n_logits; ++i) {
+ result.candidates[i] = (llama_token) i;
+ }
+ } else {
+ GGML_ASSERT(candidates != nullptr);
+ GGML_ASSERT(n_candidates == n_logits);
+ std::memcpy(result.candidates.data(), candidates, n_candidates * sizeof(llama_token));
+ }
+
+ return result;
+}
+
+struct sampler_comparison_output {
+ std::vector<llama_token_data> expected;
+ backend_sampler_output actual;
+};
+
+static sampler_comparison_output run_sampler_comparison(
+ const test_params & params,
+ const std::string & prompt,
+ const std::vector<float> & raw_logits,
+ const sampler_setup_fn & add_samplers,
+ const sampler_setup_fn & accept_history) {
+ llama_sampler_ptr cpu_chain = make_sampler_chain(add_samplers, accept_history);
+ llama_sampler_ptr backend_chain = make_sampler_chain(add_samplers, accept_history);
+ return {
+ apply_cpu_sampler(raw_logits, cpu_chain.get()),
+ run_backend_sampler(params, prompt, backend_chain.get()),
+ };
+}
+
+static std::unordered_map<llama_token, float> map_logits(const std::vector<llama_token_data> & data) {
+ std::unordered_map<llama_token, float> result;
+ result.reserve(data.size());
+ for (const auto & item : data) {
+ result[item.id] = item.logit;
+ }
+ return result;
+}
+
+struct sampler_comparison_stats {
+ int n_mismatch = 0;
+ int n_masked = 0;
+ float max_diff = 0.0f;
+};
+
+static sampler_comparison_stats compare_sampler_outputs(
+ const char * name,
+ const std::unordered_map<llama_token, float> & expected,
+ const backend_sampler_output & actual,
+ bool allow_extra_candidates = false) {
+ GGML_ASSERT(actual.logits.size() == actual.candidates.size());
+
+ sampler_comparison_stats result;
+ std::unordered_set<llama_token> seen;
+ seen.reserve(actual.candidates.size());
+
+ for (size_t i = 0; i < actual.logits.size(); ++i) {
+ const llama_token token = actual.candidates[i];
+ const float logit = actual.logits[i];
+ if (!seen.insert(token).second || std::isnan(logit)) {
+ if (result.n_mismatch < 5) {
+ printf("%s token %d has invalid backend output\n", name, token);
+ }
+ ++result.n_mismatch;
+ continue;
+ }
+
+ const auto it = expected.find(token);
+ if (it == expected.end()) {
+ if (std::isinf(logit) && logit < 0.0f) {
+ ++result.n_masked;
+ } else if (!allow_extra_candidates) {
+ if (result.n_mismatch < 5) {
+ printf("%s token %d was not masked\n", name, token);
+ }
+ ++result.n_mismatch;
+ }
+ continue;
+ }
+
+ const float diff = fabsf(it->second - logit);
+ result.max_diff = std::max(result.max_diff, diff);
+ if (!std::isfinite(logit) || diff > 1e-3f) {
+ if (result.n_mismatch < 5) {
+ printf("%s mismatch token %d: cpu=%.6f backend=%.6f diff=%.6f\n",
+ name, token, it->second, logit, diff);
+ }
+ ++result.n_mismatch;
+ }
+ }
+
+ for (const auto & item : expected) {
+ if (seen.find(item.first) == seen.end()) {
+ if (result.n_mismatch < 5) {
+ printf("%s missing backend token %d\n", name, item.first);
+ }
+ ++result.n_mismatch;
+ }
+ }
+
+ printf("%s logits: max_diff=%.6f n_masked=%d n_mismatch=%d\n",
+ name, result.max_diff, result.n_masked, result.n_mismatch);
+ return result;
+}
+
+static float find_backend_logit(const backend_sampler_output & output, llama_token token) {
+ for (size_t i = 0; i < output.candidates.size(); ++i) {
+ if (output.candidates[i] == token) {
+ return output.logits[i];
+ }
+ }
+ GGML_ABORT("backend token not found");
+}
+
+static sampler_comparison_output run_penalties_comparison(
+ const test_params & params,
+ int32_t penalty_last_n,
+ float penalty_repeat,
+ float penalty_freq,
+ float penalty_present,
+ const std::string & prompt,
+ const std::function<void(llama_sampler *)> & extra_accept = {}) {
+ const auto * vocab = llama_model_get_vocab(params.model.get());
+ 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));
+ };
+ const auto accept_history = [&](llama_sampler * chain) {
+ accept_prompt(chain, vocab, prompt);
+ if (extra_accept) {
+ extra_accept(chain);
+ }
+ };
+
+ return run_sampler_comparison(
+ params, prompt, raw_logits, add_samplers, accept_history);
+}
+
+static void compare_penalties_logits(
+ const test_params & params,
+ int32_t penalty_last_n,
+ float penalty_repeat,
+ float penalty_freq,
+ float penalty_present,
+ const std::string & prompt,
+ const std::function<void(llama_sampler *)> & extra_accept = {}) {
+ const sampler_comparison_output output = run_penalties_comparison(
+ params, penalty_last_n, penalty_repeat, penalty_freq, penalty_present, prompt, extra_accept);
+
+ GGML_ASSERT(output.expected.size() == output.actual.logits.size());
+
+ const sampler_comparison_stats stats = compare_sampler_outputs(
+ "penalties", map_logits(output.expected), output.actual);
+ GGML_ASSERT(stats.n_masked == 0);
+ GGML_ASSERT(stats.n_mismatch == 0);
+}
+
+static void test_penalty_parameter_values(const test_params & params) {
+ struct penalty_test_case {
+ const char * name;
+ float repeat;
+ float frequency;
+ float presence;
+ };
+
+ const penalty_test_case cases[] = {
+ { "frequency -1", 1.0f, -1.0f, 0.0f },
+ { "frequency 0", 1.0f, 0.0f, 0.0f },
+ { "frequency 1", 1.0f, 1.0f, 0.0f },
+ { "presence -1", 1.0f, 0.0f, -1.0f },
+ { "presence 0", 1.0f, 0.0f, 0.0f },
+ { "presence 1", 1.0f, 0.0f, 1.0f },
+ { "repeat 1", 1.0f, 0.0f, 0.0f },
+ };
+
+ int n_failed = 0;
+ for (const auto & test : cases) {
+ const sampler_comparison_output output = run_penalties_comparison(
+ params, 64, test.repeat, test.frequency, test.presence, "Hello Hello world");
+ GGML_ASSERT(output.expected.size() == output.actual.logits.size());
+ const sampler_comparison_stats stats = compare_sampler_outputs(
+ test.name, map_logits(output.expected), output.actual);
+ n_failed += stats.n_mismatch != 0;
+ }
+
+ GGML_ASSERT(n_failed == 0);
+}
+
+static void compare_top_k_penalties_logits(
+ const test_params & params,
+ int32_t k,
+ int32_t penalty_last_n,
+ float penalty_repeat,
+ float penalty_freq,
+ float penalty_present,
+ const std::string & prompt,
+ penalties_position position) {
+ const auto * vocab = llama_model_get_vocab(params.model.get());
+ const std::vector<float> raw_logits = decode_raw_logits(params, prompt);
+ const int n_vocab = (int) raw_logits.size();
+
+ GGML_ASSERT(n_vocab > k);
+
+ const sampler_init_fn init_top_k = [k]() {
+ return llama_sampler_init_top_k(k);
+ };
+ llama_sampler_ptr top_k(init_top_k());
+ const std::vector<llama_token_data> top_k_data = apply_cpu_sampler(raw_logits, top_k.get());
+ GGML_ASSERT(top_k_data.size() == (size_t) k);
+ const llama_token retained_history_token = top_k_data[0].id;
+
+ llama_token excluded_history_token = LLAMA_TOKEN_NULL;
+ for (llama_token token = 0; token < n_vocab; ++token) {
+ const auto it = std::find_if(top_k_data.begin(), top_k_data.end(), [token](const llama_token_data & data) {
+ return data.id == token;
+ });
+ if (it == top_k_data.end()) {
+ excluded_history_token = token;
+ break;
+ }
+ }
+ GGML_ASSERT(excluded_history_token != LLAMA_TOKEN_NULL);
+
+ const auto add_samplers = [&](llama_sampler * chain) {
+ add_filter_and_penalties(chain, init_top_k,
+ penalty_last_n, penalty_repeat, penalty_freq, penalty_present, position);
+ };
+
+ auto accept_history = [&](llama_sampler * smpl) {
+ accept_prompt(smpl, vocab, prompt);
+ llama_sampler_accept(smpl, excluded_history_token);
+ llama_sampler_accept(smpl, excluded_history_token);
+ llama_sampler_accept(smpl, retained_history_token);
+ llama_sampler_accept(smpl, retained_history_token);
+ };
+
+ const sampler_comparison_output output = run_sampler_comparison(
+ params, prompt, raw_logits, add_samplers, accept_history);
+
+ GGML_ASSERT(output.expected.size() == (size_t) k);
+ GGML_ASSERT(output.actual.logits.size() == (size_t) k);
+
+ const std::unordered_map<llama_token, float> expected_logits = map_logits(output.expected);
+
+ if (position == penalties_position::after_filter) {
+ GGML_ASSERT(expected_logits.find(retained_history_token) != expected_logits.end());
+ GGML_ASSERT(fabsf(expected_logits.at(retained_history_token) - raw_logits[retained_history_token]) > 1e-6f);
+ GGML_ASSERT(expected_logits.find(excluded_history_token) == expected_logits.end());
+ GGML_ASSERT(std::find(output.actual.candidates.begin(), output.actual.candidates.end(),
+ excluded_history_token) == output.actual.candidates.end());
+ } else {
+ const std::unordered_map<llama_token, float> unpenalized_logits = map_logits(top_k_data);
+ bool changed = false;
+ for (const auto & item : expected_logits) {
+ const auto it = unpenalized_logits.find(item.first);
+ if (it == unpenalized_logits.end() || fabsf(it->second - item.second) > 1e-6f) {
+ changed = true;
+ break;
+ }
+ }
+ GGML_ASSERT(changed);
+ }
+
+ const char * name = position == penalties_position::before_filter
+ ? "penalties top-k"
+ : "top-k penalties";
+ const sampler_comparison_stats stats = compare_sampler_outputs(
+ name, expected_logits, output.actual);
+ GGML_ASSERT(stats.n_masked == 0);
+ GGML_ASSERT(stats.n_mismatch == 0);
+}
+
+static void compare_masking_penalties_logits(
+ const test_params & params,
+ const char * filter_name,
+ const sampler_init_fn & init_filter,
+ int32_t penalty_last_n,
+ float penalty_repeat,
+ float penalty_freq,
+ float penalty_present,
+ const std::string & prompt,
+ penalties_position position,
+ bool allow_extra_candidates,
+ bool add_history = true) {
+ const auto * vocab = llama_model_get_vocab(params.model.get());
+ const std::vector<float> raw_logits = decode_raw_logits(params, prompt);
+ const int n_vocab = (int) raw_logits.size();
+ llama_sampler_ptr filter(init_filter());
+ const std::vector<llama_token_data> filtered_data = apply_cpu_sampler(raw_logits, filter.get());
+ GGML_ASSERT(!filtered_data.empty());
+ GGML_ASSERT(filtered_data.size() < (size_t) n_vocab);
+
+ const llama_token penalized_token = filtered_data[0].id;
+ std::unordered_set<llama_token> retained_tokens;
+ retained_tokens.reserve(filtered_data.size());
+ for (const auto & data : filtered_data) {
+ retained_tokens.insert(data.id);
+ }
+
+ llama_token masked_token = LLAMA_TOKEN_NULL;
+ for (llama_token token = 0; token < n_vocab; ++token) {
+ if (retained_tokens.find(token) == retained_tokens.end()) {
+ masked_token = token;
+ break;
+ }
+ }
+ GGML_ASSERT(masked_token != LLAMA_TOKEN_NULL);
+
+ const auto add_samplers = [&](llama_sampler * chain) {
+ add_filter_and_penalties(chain, init_filter,
+ penalty_last_n, penalty_repeat, penalty_freq, penalty_present, position);
+ };
+ auto accept_history = [&](llama_sampler * smpl) {
+ if (!add_history) {
+ return;
+ }
+ accept_prompt(smpl, vocab, prompt);
+ llama_sampler_accept(smpl, penalized_token);
+ llama_sampler_accept(smpl, penalized_token);
+ llama_sampler_accept(smpl, masked_token);
+ llama_sampler_accept(smpl, masked_token);
+ };
+
+ const sampler_comparison_output output = run_sampler_comparison(
+ params, prompt, raw_logits, add_samplers, accept_history);
+
+ GGML_ASSERT(output.actual.logits.size() == (size_t) n_vocab);
+
+ const std::unordered_map<llama_token, float> expected_logits = map_logits(output.expected);
+
+ GGML_ASSERT(expected_logits.find(masked_token) == expected_logits.end());
+ if (add_history) {
+ if (position == penalties_position::after_filter) {
+ GGML_ASSERT(expected_logits.find(penalized_token) != expected_logits.end());
+ 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));
+ accept_history(penalties.get());
+ const std::unordered_map<llama_token, float> penalized_logits =
+ map_logits(apply_cpu_sampler(raw_logits, penalties.get()));
+ GGML_ASSERT(fabsf(penalized_logits.at(penalized_token) - raw_logits[penalized_token]) > 1e-6f);
+ }
+ }
+
+ const std::string name = position == penalties_position::before_filter
+ ? "penalties " + std::string(filter_name)
+ : std::string(filter_name) + " penalties";
+ const sampler_comparison_stats stats = compare_sampler_outputs(
+ name.c_str(), expected_logits, output.actual, allow_extra_candidates);
+ const float masked_logit = find_backend_logit(output.actual, masked_token);
+ GGML_ASSERT(stats.n_masked > 0);
+ GGML_ASSERT(std::isinf(masked_logit) && masked_logit < 0.0f);
+ GGML_ASSERT(stats.n_mismatch == 0);
+}
+
+static void test_backend_penalties_sampling(const test_params & params) {
+ printf("Testing backend penalties (repeat + freq + presence)\n");
+ compare_penalties_logits(params, 64, 1.1f, 0.5f, 0.25f, "Hello Hello world");
+
+ printf("Testing backend penalties with penalty_last_n > 64\n");
+ const auto * vocab = llama_model_get_vocab(params.model.get());
+ std::vector<llama_token> tokens(8);
+ int32_t n_tok = llama_tokenize(vocab, "a", 1, tokens.data(), (int32_t) tokens.size(), false, false);
+ if (n_tok < 0) {
+ tokens.resize(-n_tok);
+ n_tok = llama_tokenize(vocab, "a", 1, tokens.data(), (int32_t) tokens.size(), false, false);
+ }
+ GGML_ASSERT(n_tok > 0);
+ const llama_token tok = tokens[0];
+
+ compare_penalties_logits(params, 80, 1.15f, 0.1f, 0.05f, "a", [tok](llama_sampler * smpl) {
+ // accept_prompt already accepted BOS + one 'a'; fill the ring to n=80
+ for (int i = 0; i < 78; ++i) {
+ llama_sampler_accept(smpl, tok);
+ }
+ });
+
+ printf("Testing backend penalties without filler entries\n");
+ compare_penalties_logits(params, 64, 1.1f, 0.5f, 0.25f, "Hello", [](llama_sampler * smpl) {
+ for (llama_token token = 0; token < 64; ++token) {
+ llama_sampler_accept(smpl, token);
+ }
+ });
+
+ printf("Testing backend top-k followed by penalties\n");
+ compare_top_k_penalties_logits(params, 8, 64, 1.1f, 0.5f, 0.25f, "Hello",
+ penalties_position::after_filter);
+
+ printf("Testing backend penalties followed by top-k\n");
+ compare_top_k_penalties_logits(params, 8, 64, 1.1f, 0.5f, 0.25f, "Hello",
+ penalties_position::before_filter);
+
+ printf("Testing backend top-p followed by penalties\n");
+ compare_masking_penalties_logits(params, "top-p", []() {
+ return llama_sampler_init_top_p(0.9f, 0);
+ }, 64, 1.1f, 0.5f, 0.25f, "Hello", penalties_position::after_filter, true);
+
+ printf("Testing backend top-p followed by penalties with a large history window\n");
+ compare_masking_penalties_logits(params, "top-p large-window", []() {
+ return llama_sampler_init_top_p(0.9f, 0);
+ }, 4096, 1.1f, 0.5f, 0.25f, "Hello", penalties_position::after_filter, true);
+
+ printf("Testing backend penalties followed by top-p\n");
+ compare_masking_penalties_logits(params, "top-p", []() {
+ return llama_sampler_init_top_p(0.9f, 0);
+ }, 64, 1.1f, 0.5f, 0.25f, "Hello", penalties_position::before_filter, true);
+
+ printf("Testing backend min-p followed by penalties\n");
+ compare_masking_penalties_logits(params, "min-p", []() {
+ return llama_sampler_init_min_p(0.1f, 0);
+ }, 64, 1.1f, 0.5f, 0.25f, "Hello", penalties_position::after_filter, false);
+
+ printf("Testing backend penalties followed by min-p\n");
+ compare_masking_penalties_logits(params, "min-p", []() {
+ return llama_sampler_init_min_p(0.1f, 0);
+ }, 64, 1.1f, 0.5f, 0.25f, "Hello", penalties_position::before_filter, false);
+
+ printf("Testing backend top-p followed by penalties with empty history\n");
+ compare_masking_penalties_logits(params, "top-p empty", []() {
+ return llama_sampler_init_top_p(0.9f, 0);
+ }, 64, 1.1f, 0.5f, 0.25f, "Hello", penalties_position::after_filter, true, false);
+
+ printf("Testing backend top-p followed by individual penalties\n");
+ compare_masking_penalties_logits(params, "top-p repeat", []() {
+ return llama_sampler_init_top_p(0.9f, 0);
+ }, 64, 1.1f, 0.0f, 0.0f, "Hello", penalties_position::after_filter, true);
+ compare_masking_penalties_logits(params, "top-p frequency", []() {
+ return llama_sampler_init_top_p(0.9f, 0);
+ }, 64, 1.0f, 0.5f, 0.0f, "Hello", penalties_position::after_filter, true);
+ compare_masking_penalties_logits(params, "top-p presence", []() {
+ return llama_sampler_init_top_p(0.9f, 0);
+ }, 64, 1.0f, 0.0f, 0.25f, "Hello", penalties_position::after_filter, true);
+
+ printf("Testing backend penalty parameter values\n");
+ test_penalty_parameter_values(params);
+
+ printf("backend penalties sampling test PASSED\n");
+}
+
// This test verifies that it is possible to have two different backend samplers,
// one that uses the backend dist sampler, and another that uses CPU dist sampler.
static void test_backend_mixed_sampling(const test_params & params) {
static const backend_test_case BACKEND_TESTS[] = {
{ "greedy", test_backend_greedy_sampling, true },
{ "logit_bias", test_backend_logit_bias_sampling, true },
+ { "penalties", test_backend_penalties_sampling, true },
{ "temp", test_backend_temp_sampling, true },
{ "temp_ext", test_backend_temp_ext_sampling, true },
{ "top_k", test_backend_top_k_sampling, true },