Commit 63412d873 for llama.cpp
commit 63412d8737b3efa30748dd79ceb2c93487e19f2b
Author: Trunk Trick <2451421621@qq.com>
Date: Sun Oct 25 19:31:09 2026 +0800
llama: fix exhausting CPU and MEM caused by dry sequence breaker (#31035)
diff --git a/include/llama.h b/include/llama.h
index 3af6e2223..ae54e661b 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -1564,6 +1564,10 @@ extern "C" {
float penalty_freq, // must be finite, 0.0 = disabled
float penalty_present); // must be finite, 0.0 = disabled
+ /// maximum number of distinct DRY sequence breakers, duplicates not counted
+ /// processing each breaker scans the entire vocabulary, so the init cost is O(num_breakers x n_vocab)
+ #define LLAMA_DRY_MAX_SEQ_BREAKERS 64
+
/// @details DRY sampler, designed by p-e-w, as described in: https://github.com/oobabooga/text-generation-webui/pull/5677, porting Koboldcpp implementation authored by pi6am: https://github.com/LostRuins/koboldcpp/pull/982
LLAMA_API struct llama_sampler * llama_sampler_init_dry(
const struct llama_vocab * vocab,
diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp
index 521c5718e..9004efa20 100644
--- a/src/llama-sampler.cpp
+++ b/src/llama-sampler.cpp
@@ -18,6 +18,7 @@
#include <numeric>
#include <random>
#include <unordered_map>
+#include <unordered_set>
#include <stdexcept>
// the ring buffer works similarly to std::deque, but with a fixed capacity
@@ -3333,9 +3334,13 @@ struct llama_sampler_dry {
};
// Ported from Koboldcpp, original PR: https://github.com/LostRuins/koboldcpp/pull/982 (Original author: pi6am)
-static void get_overlapping_token_sequences(const llama_vocab & vocab, const std::string& str, std::unordered_multimap<llama_token, std::vector<llama_token>>& token_sequences, int max_tail_len = -1) {
- for (llama_token token_id = 0; token_id < (llama_token) vocab.n_tokens(); token_id++) {
- std::string word = vocab.detokenize({token_id}, true);
+static void get_overlapping_token_sequences(const llama_vocab & vocab, const std::vector<std::string> & vocab_words, const std::string& str, std::unordered_multimap<llama_token, std::vector<llama_token>>& token_sequences, int max_tail_len = -1) {
+ // the tail tokenization only depends on how many characters of str were matched, so compute each one at most once
+ std::vector<std::vector<llama_token>> tail_cache(str.size() + 1);
+ std::vector<bool> tail_cached(str.size() + 1, false);
+
+ for (llama_token token_id = 0; token_id < (llama_token) vocab_words.size(); token_id++) {
+ const std::string & word = vocab_words[token_id];
if (word.find(str) != std::string::npos) {
token_sequences.emplace(token_id, std::vector<llama_token>());
} else {
@@ -3352,10 +3357,14 @@ static void get_overlapping_token_sequences(const llama_vocab & vocab, const std
}
}
if (match) {
- std::vector<llama_token> tokenization = vocab.tokenize(str.substr(i), false, false);
- if (max_tail_len >= 0 && tokenization.size() > (size_t)max_tail_len) {
- tokenization.resize(max_tail_len);
+ if (!tail_cached[i]) {
+ tail_cache[i] = vocab.tokenize(str.substr(i), false, false);
+ if (max_tail_len >= 0 && tail_cache[i].size() > (size_t)max_tail_len) {
+ tail_cache[i].resize(max_tail_len);
+ }
+ tail_cached[i] = true;
}
+ const std::vector<llama_token> & tokenization = tail_cache[i];
// Ensure we don't already have a duplicate matching tokenization
auto its = token_sequences.equal_range(token_id);
@@ -3656,6 +3665,16 @@ struct llama_sampler * llama_sampler_init_dry(const struct llama_vocab * vocab,
if (dry_enabled && seq_breakers != nullptr && num_breakers > 0) {
// Process sequence breakers
+ // each distinct breaker costs a full vocabulary scan: drop duplicates and stop at the cap
+ std::unordered_set<std::string> seen_breakers;
+
+ // detokenize the vocabulary once instead of rescanning it for every breaker
+ std::vector<std::string> vocab_words;
+ vocab_words.reserve(vocab->n_tokens());
+ for (llama_token token_id = 0; token_id < (llama_token) vocab->n_tokens(); ++token_id) {
+ vocab_words.push_back(vocab->detokenize({token_id}, true));
+ }
+
for (size_t i = 0; i < num_breakers; ++i) {
if (seq_breakers[i] == nullptr || std::strlen(seq_breakers[i]) == 0) {
LLAMA_LOG_WARN("skipping null or empty DRY sequence breaker at index %zu\n", i);
@@ -3663,17 +3682,22 @@ struct llama_sampler * llama_sampler_init_dry(const struct llama_vocab * vocab,
}
std::string sequence_break(seq_breakers[i]);
- if (sequence_break.empty()) {
- LLAMA_LOG_WARN("skipping empty DRY sequence breaker\n");
- continue;
- }
-
if (sequence_break.size() > MAX_CHAR_LEN) {
LLAMA_LOG_WARN("truncating DRY sequence breaker to %d characters\n", MAX_CHAR_LEN);
sequence_break.resize(MAX_CHAR_LEN);
}
- get_overlapping_token_sequences(*vocab, sequence_break, processed_breakers, MAX_SEQ_LEN);
+ if (!seen_breakers.insert(sequence_break).second) {
+ continue;
+ }
+
+ if (seen_breakers.size() > LLAMA_DRY_MAX_SEQ_BREAKERS) {
+ LLAMA_LOG_WARN("ignoring DRY sequence breaker at index %zu: more than %d distinct breakers\n",
+ i, LLAMA_DRY_MAX_SEQ_BREAKERS);
+ break;
+ }
+
+ get_overlapping_token_sequences(*vocab, vocab_words, sequence_break, processed_breakers, MAX_SEQ_LEN);
}
}
diff --git a/tests/test-server-component.cpp b/tests/test-server-component.cpp
index 92eeec055..0d2329cf7 100644
--- a/tests/test-server-component.cpp
+++ b/tests/test-server-component.cpp
@@ -2,6 +2,7 @@
#include "../tools/server/server-common.h"
#include "../tools/server/server-decision.h"
+#include "../tools/server/server-schema.h"
#include <cstdlib>
#include <regex>
@@ -490,6 +491,42 @@ static void test_decision_depth(testing & t) {
});
}
+// DRY sequence breakers
+//
+
+static void test_dry_sequence_breakers(testing & t) {
+ common_params params_base;
+
+ const auto eval = [&](const json & data) {
+ return server_schema::eval_llama_cmpl_schema(nullptr, params_base, {}, data);
+ };
+
+ {
+ json data;
+ data["dry_sequence_breakers"] = json::array({"ab", "ab", "cd", "ab"});
+ const auto params = eval(data);
+ const std::vector<std::string> expected = {"ab", "cd"};
+ t.assert_true("duplicates removed, order kept", params.sampling.dry_sequence_breakers == expected);
+ }
+
+ {
+ std::vector<std::string> many;
+ for (int i = 0; i <= LLAMA_DRY_MAX_SEQ_BREAKERS; ++i) {
+ many.push_back("b" + std::to_string(i));
+ }
+ json data;
+ data["dry_sequence_breakers"] = many;
+
+ bool thrown = false;
+ try {
+ eval(data);
+ } catch (const std::exception &) {
+ thrown = true;
+ }
+ t.assert_true("more than the cap is rejected", thrown);
+ }
+}
+
int main(int argc, char ** argv) {
testing t;
@@ -511,6 +548,7 @@ int main(int argc, char ** argv) {
t.test("decision_escape", test_decision_escape);
t.test("rerank_template", test_rerank_template);
t.test("decision_depth", test_decision_depth);
+ t.test("dry_sequence_breakers", test_dry_sequence_breakers);
return t.summary();
}
diff --git a/tools/server/README.md b/tools/server/README.md
index 56d1ea5fe..8e27e37b8 100644
--- a/tools/server/README.md
+++ b/tools/server/README.md
@@ -563,7 +563,7 @@ These words will not be included in the completion, so make sure to add them to
`dry_penalty_last_n`: How many tokens to scan for repetitions. Default: `64`, where `0` is disabled.
-`dry_sequence_breakers`: Specify an array of sequence breakers for DRY sampling. Only a JSON array of strings is accepted. Default: `['\n', ':', '"', '*']`
+`dry_sequence_breakers`: Specify an array of sequence breakers for DRY sampling. Only a JSON array of strings is accepted. At most 64 distinct breakers, duplicates are ignored. Default: `['\n', ':', '"', '*']`
`xtc_probability`: Set the chance for token removal via XTC sampler. Default: `0.0`, which is disabled.
diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp
index 1e2dce888..0a14a9620 100644
--- a/tools/server/server-schema.cpp
+++ b/tools/server/server-schema.cpp
@@ -2,6 +2,8 @@
#include "json-schema-to-grammar.h"
+#include <unordered_set>
+
namespace server_schema {
//
@@ -246,6 +248,23 @@ std::vector<std::unique_ptr<field>> make_llama_cmpl_schema(const common_params &
if (ctx.params.sampling.dry_sequence_breakers.empty()) {
throw std::runtime_error("Error: dry_sequence_breakers must be a non-empty array of strings");
}
+
+ // each distinct breaker costs a full vocabulary scan at sampler init, so dedupe before counting
+ std::unordered_set<std::string> seen;
+ std::vector<std::string> unique_breakers;
+
+ for (const auto & breaker : ctx.params.sampling.dry_sequence_breakers) {
+ if (seen.insert(breaker).second) {
+ unique_breakers.push_back(breaker);
+ }
+ }
+
+ if (unique_breakers.size() > LLAMA_DRY_MAX_SEQ_BREAKERS) {
+ throw std::runtime_error(string_format(
+ "Error: dry_sequence_breakers must contain at most %d distinct entries", LLAMA_DRY_MAX_SEQ_BREAKERS));
+ }
+
+ ctx.params.sampling.dry_sequence_breakers = std::move(unique_breakers);
}));
// handle both "json_schema" and "grammar"