From 63412d8737b3efa30748dd79ceb2c93487e19f2b Mon Sep 17 00:00:00 2001 From: Trunk Trick <2451421621@qq.com> Date: Sun, 25 Oct 2026 19:31:09 +0800 Subject: [PATCH] llama: fix exhausting CPU and MEM caused by dry sequence breaker (#31035) --- include/llama.h | 4 +++ src/llama-sampler.cpp | 48 ++++++++++++++++++++++++--------- tests/test-server-component.cpp | 38 ++++++++++++++++++++++++++ tools/server/README.md | 2 +- tools/server/server-schema.cpp | 19 +++++++++++++ 5 files changed, 98 insertions(+), 13 deletions(-) diff --git a/include/llama.h b/include/llama.h index 3af6e22233..ae54e661be 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 521c5718e1..9004efa20d 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #include // 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>& 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 & vocab_words, const std::string& str, std::unordered_multimap>& 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> tail_cache(str.size() + 1); + std::vector 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()); } else { @@ -3352,10 +3357,14 @@ static void get_overlapping_token_sequences(const llama_vocab & vocab, const std } } if (match) { - std::vector 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 & 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 seen_breakers; + + // detokenize the vocabulary once instead of rescanning it for every breaker + std::vector 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 92eeec055b..0d2329cf73 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 #include @@ -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 expected = {"ab", "cd"}; + t.assert_true("duplicates removed, order kept", params.sampling.dry_sequence_breakers == expected); + } + + { + std::vector 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 56d1ea5fe8..8e27e37b8d 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 1e2dce8886..0a14a96209 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -2,6 +2,8 @@ #include "json-schema-to-grammar.h" +#include + namespace server_schema { // @@ -246,6 +248,23 @@ std::vector> 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 seen; + std::vector 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"