From 73c900a97daf98ab662e07e1b1dbb89d01733ffc Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Fri, 4 Sep 2026 01:01:45 +0200 Subject: [PATCH] add synthid sampler --- include/llama.h | 14 ++++ src/llama-sampler.cpp | 180 ++++++++++++++++++++++++++++++++++++++++ tests/test-sampling.cpp | 42 ++++++++++ 3 files changed, 236 insertions(+) diff --git a/include/llama.h b/include/llama.h index ef7a012c43..4dda1f3778 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1394,6 +1394,20 @@ extern "C" { /// @details Top n sigma sampling as described in academic paper "Top-nσ: Not All Logits Are You Need" https://arxiv.org/pdf/2411.07641 LLAMA_API struct llama_sampler * llama_sampler_init_top_n_sigma(float n); + /// @details SynthID text watermarking, compatible with the SynthIDTextWatermarkLogitsProcessor from HF transformers + /// place it after truncation and temperature samplers and before the dist sampler + /// @param keys one watermarking key per tournament layer + /// @param sampling_table table of 0/1 values, maps a hashed ngram to a g-value + /// @param ngram_len number of tokens hashed together (context of ngram_len - 1 tokens plus the candidate) + /// @param context_history_size number of recent contexts to remember, a repeated context is not watermarked + LLAMA_API struct llama_sampler * llama_sampler_init_synthid( + const int64_t * keys, + size_t n_keys, + const uint8_t * sampling_table, + size_t sampling_table_size, + int32_t ngram_len, + int32_t context_history_size); + /// @details Mirostat 1.0 algorithm described in the paper https://arxiv.org/abs/2007.14966. Uses tokens instead of words. /// @param candidates A vector of `llama_token_data` containing the candidate tokens, their probabilities (p), and log-odds (logit) for the current position in the generated text. /// @param tau The target cross-entropy (or surprise) value you want to achieve for the generated text. A higher value corresponds to more surprising or less predictable text, while a lower value corresponds to less surprising or more predictable text. diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index 34a7988262..90a898cbec 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -3312,6 +3312,186 @@ struct llama_sampler * llama_sampler_init_top_n_sigma(float n) { ); } +// synthid + +struct llama_sampler_synthid { + const std::vector keys; + const std::vector sampling_table; + const int32_t ngram_len; + const int32_t context_history_size; + + // last ngram_len - 1 accepted tokens + std::vector context; + + ring_buffer context_history; + + std::vector hashes; + std::vector g_values; + + llama_sampler_synthid( + std::vector keys, + std::vector sampling_table, + int32_t ngram_len, + int32_t context_history_size) + : keys (std::move(keys)) + , sampling_table (std::move(sampling_table)) + , ngram_len (ngram_len) + , context_history_size(context_history_size) + , context (ngram_len - 1, 0) + , context_history (context_history_size) { + } + + // linear congruential generator, same constants as the HF implementation + static uint64_t hash(uint64_t h, uint64_t x) { + return (h + x) * 6364136223846793005ULL + 1; + } + + uint8_t g_value(uint64_t key) const { + const int64_t n = (int64_t) sampling_table.size(); + + // python-style modulo, the result is always in [0, n) + int64_t idx = (int64_t) key % n; + if (idx < 0) { + idx += n; + } + + return sampling_table[idx]; + } +}; + +static const char * llama_sampler_synthid_name(const struct llama_sampler * /*smpl*/) { + return "synthid"; +} + +static void llama_sampler_synthid_accept(struct llama_sampler * smpl, llama_token token) { + auto * ctx = (llama_sampler_synthid *) smpl->ctx; + + if (ctx->context.empty()) { + return; + } + + std::copy(ctx->context.begin() + 1, ctx->context.end(), ctx->context.begin()); + ctx->context.back() = token; +} + +static void llama_sampler_synthid_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) { + auto * ctx = (llama_sampler_synthid *) smpl->ctx; + + uint64_t h_context = 1; + for (const llama_token t : ctx->context) { + h_context = llama_sampler_synthid::hash(h_context, (uint64_t) (int64_t) t); + } + + bool is_repeated = false; + for (size_t i = 0; i < ctx->context_history.size(); ++i) { + if (ctx->context_history.rat(i) == h_context) { + is_repeated = true; + break; + } + } + + ctx->context_history.push_back(h_context); + + if (is_repeated) { + return; + } + + llama_sampler_softmax_impl(cur_p, false); + + ctx->hashes.resize(cur_p->size); + ctx->g_values.resize(cur_p->size); + + for (size_t i = 0; i < cur_p->size; ++i) { + ctx->hashes[i] = llama_sampler_synthid::hash(h_context, (uint64_t) (int64_t) cur_p->data[i].id); + } + + for (const int64_t key : ctx->keys) { + float g_mass = 0.0f; + + for (size_t i = 0; i < cur_p->size; ++i) { + ctx->g_values[i] = ctx->g_value(llama_sampler_synthid::hash(ctx->hashes[i], (uint64_t) key)); + g_mass += ctx->g_values[i] * cur_p->data[i].p; + } + + for (size_t i = 0; i < cur_p->size; ++i) { + cur_p->data[i].p *= 1.0f + ctx->g_values[i] - g_mass; + } + } + + for (size_t i = 0; i < cur_p->size; ++i) { + cur_p->data[i].logit = cur_p->data[i].p > 0.0f ? logf(cur_p->data[i].p) : -INFINITY; + } + + cur_p->sorted = false; +} + +static void llama_sampler_synthid_reset(struct llama_sampler * smpl) { + auto * ctx = (llama_sampler_synthid *) smpl->ctx; + std::fill(ctx->context.begin(), ctx->context.end(), 0); + ctx->context_history.clear(); +} + +static struct llama_sampler * llama_sampler_synthid_clone(const struct llama_sampler * smpl) { + const auto * ctx = (const llama_sampler_synthid *) smpl->ctx; + auto * result = llama_sampler_init_synthid( + ctx->keys.data(), + ctx->keys.size(), + ctx->sampling_table.data(), + ctx->sampling_table.size(), + ctx->ngram_len, + ctx->context_history_size); + + // copy the state + { + auto * result_ctx = (llama_sampler_synthid *) result->ctx; + + result_ctx->context = ctx->context; + result_ctx->context_history = ctx->context_history; + } + + return result; +} + +static void llama_sampler_synthid_free(struct llama_sampler * smpl) { + delete (llama_sampler_synthid *) smpl->ctx; +} + +static struct llama_sampler_i llama_sampler_synthid_i = { + /* .name = */ llama_sampler_synthid_name, + /* .accept = */ llama_sampler_synthid_accept, + /* .apply = */ llama_sampler_synthid_apply, + /* .reset = */ llama_sampler_synthid_reset, + /* .clone = */ llama_sampler_synthid_clone, + /* .free = */ llama_sampler_synthid_free, + /* .backend_init = */ nullptr, + /* .backend_accept = */ nullptr, + /* .backend_apply = */ nullptr, + /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, +}; + +struct llama_sampler * llama_sampler_init_synthid( + const int64_t * keys, + size_t n_keys, + const uint8_t * sampling_table, + size_t sampling_table_size, + int32_t ngram_len, + int32_t context_history_size) { + if (n_keys == 0 || sampling_table_size == 0 || ngram_len < 1 || context_history_size < 1) { + return llama_sampler_init_empty("?synthid"); + } + + return llama_sampler_init( + /* .iface = */ &llama_sampler_synthid_i, + /* .ctx = */ new llama_sampler_synthid( + std::vector(keys, keys + n_keys), + std::vector(sampling_table, sampling_table + sampling_table_size), + ngram_len, + context_history_size) + ); +} + // DRY struct llama_sampler_dry { diff --git a/tests/test-sampling.cpp b/tests/test-sampling.cpp index d727ab632a..9ae9277b26 100644 --- a/tests/test-sampling.cpp +++ b/tests/test-sampling.cpp @@ -221,6 +221,46 @@ static void test_top_n_sigma(const std::vector & probs, const std::vector tester.check(); } +static void test_synthid( + const std::vector & probs, const std::vector & last_tokens, const std::vector & probs_expected +) { + GGML_ASSERT(probs.size() == probs_expected.size()); + + // values generated with SynthIDTextWatermarkLogitsProcessor from HF transformers + const std::vector keys = { 654, 400, 836 }; + const std::vector sampling_table = { 1, 0, 0, 1, 1, 1, 0, 1, 0, 0, 1, 0, 1, 1, 0, 1 }; + + auto * sampler = llama_sampler_init_synthid(keys.data(), keys.size(), sampling_table.data(), sampling_table.size(), 3, 4); + + for (size_t i = 0; i < last_tokens.size(); i++) { + llama_sampler_accept(sampler, last_tokens[i]); + } + + { + sampler_tester tester(probs, probs_expected); + + DUMP(&tester.cur_p); + llama_sampler_apply(sampler, &tester.cur_p); + tester.apply(llama_sampler_init_dist(0)); + DUMP(&tester.cur_p); + + tester.check(); + } + + // same context again -> the watermark is not applied + { + sampler_tester tester(probs, probs); + + llama_sampler_apply(sampler, &tester.cur_p); + tester.apply(llama_sampler_init_dist(0)); + DUMP(&tester.cur_p); + + tester.check(); + } + + llama_sampler_free(sampler); +} + static void test_sampler_queue(const size_t n_vocab, const std::string & samplers_sequence, const int top_k, const float top_p, const float min_p ) { sampler_tester tester(n_vocab); @@ -395,6 +435,8 @@ int main(void) { test_top_n_sigma({0.1f, 0.2f, 0.3f, 0.4f}, {0.1f, 0.2f, 0.3f, 0.4f}, 0.00f); // top_n_sigma == 0 now represents a no-op rather than greedy decoding as of PR#13345 test_top_n_sigma({0.1f, 0.2f, 0.3f, 0.4f}, {0.1f, 0.2f, 0.3f, 0.4f}, 3.00f); + test_synthid({0.05f, 0.1f, 0.15f, 0.2f, 0.25f, 0.25f}, {3, 5}, {0.049753f, 0.008964f, 0.029852f, 0.294519f, 0.368149f, 0.248764f}); + test_sampler_queue(10000, "k", 10000, 1.0f, 1.0f); test_sampler_queue(10000, "k", 1, 1.0f, 1.0f); test_sampler_queue(10000, "p", 10000, 1.0f, 1.0f);