From 22c51b60962761a7464f15cee746769863bf9de5 Mon Sep 17 00:00:00 2001 From: Alde Rojas Date: Wed, 8 Jul 2026 17:33:17 -0500 Subject: [PATCH 1/8] common : extract trie/ac to a separate file --- common/CMakeLists.txt | 2 + common/peg-parser.cpp | 158 ++---------------------------------------- common/trie.cpp | 118 +++++++++++++++++++++++++++++++ common/trie.h | 67 ++++++++++++++++++ 4 files changed, 192 insertions(+), 153 deletions(-) create mode 100644 common/trie.cpp create mode 100644 common/trie.h diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index 4cf580a056c8..99688f53b87b 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -100,6 +100,8 @@ add_library(${TARGET} sampling.h speculative.cpp speculative.h + trie.cpp + trie.h unicode.cpp unicode.h jinja/lexer.cpp diff --git a/common/peg-parser.cpp b/common/peg-parser.cpp index 807e952d902c..ef290ed7c057 100644 --- a/common/peg-parser.cpp +++ b/common/peg-parser.cpp @@ -3,10 +3,10 @@ #include "common.h" #include "json-schema-to-grammar.h" #include "log.h" +#include "trie.h" #include "unicode.h" #include -#include #include #include #include @@ -32,154 +32,6 @@ static bool is_hex_digit(const char c) { return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F'); } -// Trie for matching multiple literals. -// This is used in common_peg_until_parser and to build a GBNF exclusion grammar -struct trie { - struct node { - std::map children; // Use uint32_t to store Unicode codepoints - bool is_word; - }; - - std::vector nodes; - - trie(const std::vector & words) { - create_node(); // root node - for (const auto & w : words) { - insert(w); - } - } - - enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH }; - - // Check if a delimiter starts at the given position - match_result check_at(std::string_view sv, size_t start_pos) const { - size_t current = 0; // Start at root - size_t pos = start_pos; - - // LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str()); - - while (pos < sv.size()) { - auto result = common_parse_utf8_codepoint(sv, pos); - if (result.status != utf8_parse_result::SUCCESS) { - break; - } - - auto it = nodes[current].children.find(result.codepoint); - if (it == nodes[current].children.end()) { - // Can't continue matching - return match_result{match_result::NO_MATCH}; - } - - current = it->second; - pos += result.bytes_consumed; - - // Check if we've matched a complete word - if (nodes[current].is_word) { - return match_result{match_result::COMPLETE_MATCH}; - } - } - - // Reached end of input while still in the trie (not at root) - if (current != 0) { - // We're in the middle of a potential match - return match_result{match_result::PARTIAL_MATCH}; - } - - // Reached end at root (no match) - return match_result{match_result::NO_MATCH}; - } - - private: - size_t create_node() { - size_t index = nodes.size(); - nodes.emplace_back(); - return index; - } - - void insert(const std::string & word) { - size_t current = 0; - size_t pos = 0; - while (pos < word.length()) { - auto result = common_parse_utf8_codepoint(word, pos); - if (result.status != utf8_parse_result::SUCCESS) { - break; - } - - uint32_t ch = result.codepoint; - pos += result.bytes_consumed; - - auto it = nodes[current].children.find(ch); - if (it == nodes[current].children.end()) { - size_t child = create_node(); - nodes[current].children[ch] = child; - current = child; - } else { - current = it->second; - } - } - nodes[current].is_word = true; - } -}; - -// Aho-Corasick automaton -struct aho_corasick { - trie t; - std::vector fail; // failure links - std::vector order; // states in BFS order - std::vector terminal; // match states (directly or via a suffix link) - std::set alphabet; // every character with a transition - - aho_corasick(const std::vector & strings) : t(strings) { - const auto & nodes = t.nodes; - const size_t n = nodes.size(); - - fail.assign(n, 0); - order.reserve(n); - - std::deque queue{ 0 }; - while (!queue.empty()) { - size_t u = queue.front(); - queue.pop_front(); - order.push_back(u); - for (const auto & [ch, v] : nodes[u].children) { - if (u != 0) { - size_t f = fail[u]; - while (f && nodes[f].children.find(ch) == nodes[f].children.end()) { - f = fail[f]; - } - auto it = nodes[f].children.find(ch); - fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0; - } - queue.push_back(v); - } - } - - terminal.assign(n, false); - for (size_t u : order) { - terminal[u] = nodes[u].is_word || (u != 0 && terminal[fail[u]]); - } - - for (const auto & node : nodes) { - for (const auto & [ch, v] : node.children) { - alphabet.insert(ch); - } - } - } - - size_t num_states() const { return t.nodes.size(); } - bool is_terminal(size_t s) const { return terminal[s]; } - - // follow failure links until a transition on `ch` exists. - size_t next(size_t state, uint32_t ch) const { - const auto & nodes = t.nodes; - while (state && nodes[state].children.find(ch) == nodes[state].children.end()) { - state = fail[state]; - } - auto it = nodes[state].children.find(ch); - return it != nodes[state].children.end() ? it->second : 0; - } -}; - static std::pair parse_hex_escape(const std::string & str, size_t pos, int hex_count) { if (pos + hex_count > str.length()) { return {0, 0}; @@ -797,7 +649,7 @@ struct parser_executor { } common_peg_parse_result operator()(const common_peg_until_parser & p) const { - trie matcher(p.delimiters); + common_trie matcher(p.delimiters); // Scan input and check for delimiters size_t pos = start_pos; @@ -824,12 +676,12 @@ struct parser_executor { // Check if a delimiter starts at this position auto match = matcher.check_at(ctx.input, pos); - if (match == trie::COMPLETE_MATCH) { + if (match == common_trie::COMPLETE_MATCH) { // Found a complete delimiter, return everything before it return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); } - if (match == trie::PARTIAL_MATCH) { + if (match == common_trie::PARTIAL_MATCH) { // Found a partial match extending to end of input, return everything before it return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); } @@ -1559,7 +1411,7 @@ static std::string gbnf_ac_grammar( const std::map> &, const std::vector &, const std::function &)> & build_rule) { - aho_corasick ac(strings); + common_aho_corasick ac(strings); auto state_name = [&](size_t s) -> std::string { if (s == 0) { diff --git a/common/trie.cpp b/common/trie.cpp new file mode 100644 index 000000000000..a5d03fb1c83e --- /dev/null +++ b/common/trie.cpp @@ -0,0 +1,118 @@ +#include "trie.h" + +#include "unicode.h" + +#include + +common_trie::match_result common_trie::check_at(std::string_view sv, size_t start_pos) const { + size_t current = 0; // Start at root + size_t pos = start_pos; + + // LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str()); + + while (pos < sv.size()) { + auto result = common_parse_utf8_codepoint(sv, pos); + if (result.status != utf8_parse_result::SUCCESS) { + break; + } + + auto it = nodes[current].children.find(result.codepoint); + if (it == nodes[current].children.end()) { + // Can't continue matching + return match_result{match_result::NO_MATCH}; + } + + current = it->second; + pos += result.bytes_consumed; + + // Check if we've matched a complete word + if (nodes[current].is_word) { + return match_result{match_result::COMPLETE_MATCH}; + } + } + + // Reached end of input while still in the trie (not at root) + if (current != 0) { + // We're in the middle of a potential match + return match_result{match_result::PARTIAL_MATCH}; + } + + // Reached end at root (no match) + return match_result{match_result::NO_MATCH}; +} + +void common_trie::insert(const std::string & word) { + std::vector symbols; + size_t pos = 0; + while (pos < word.length()) { + auto result = common_parse_utf8_codepoint(word, pos); + if (result.status != utf8_parse_result::SUCCESS) { + break; + } + + symbols.push_back(result.codepoint); + pos += result.bytes_consumed; + } + insert(symbols); +} + +void common_trie::insert(const std::vector & symbols) { + size_t current = 0; + for (uint32_t ch : symbols) { + auto it = nodes[current].children.find(ch); + if (it == nodes[current].children.end()) { + size_t child = create_node(); + nodes[current].children[ch] = child; + current = child; + } else { + current = it->second; + } + } + nodes[current].is_word = true; +} + +common_aho_corasick::common_aho_corasick(common_trie trie) : t(std::move(trie)) { + const auto & nodes = t.nodes; + const size_t n = nodes.size(); + + fail.assign(n, 0); + order.reserve(n); + + std::deque queue{ 0 }; + while (!queue.empty()) { + size_t u = queue.front(); + queue.pop_front(); + order.push_back(u); + for (const auto & [ch, v] : nodes[u].children) { + if (u != 0) { + size_t f = fail[u]; + while (f && nodes[f].children.find(ch) == nodes[f].children.end()) { + f = fail[f]; + } + auto it = nodes[f].children.find(ch); + fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0; + } + queue.push_back(v); + } + } + + terminal.assign(n, false); + for (size_t u : order) { + terminal[u] = nodes[u].is_word || (u != 0 && terminal[fail[u]]); + } + + for (const auto & node : nodes) { + for (const auto & [ch, v] : node.children) { + alphabet.insert(ch); + } + } +} + +size_t common_aho_corasick::next(size_t state, uint32_t ch) const { + const auto & nodes = t.nodes; + while (state && nodes[state].children.find(ch) == nodes[state].children.end()) { + state = fail[state]; + } + auto it = nodes[state].children.find(ch); + return it != nodes[state].children.end() ? it->second : 0; +} diff --git a/common/trie.h b/common/trie.h new file mode 100644 index 000000000000..bf03471e7e7c --- /dev/null +++ b/common/trie.h @@ -0,0 +1,67 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +// Trie for matching multiple literals. +// This is used in common_peg_until_parser and to build a GBNF exclusion grammar +struct common_trie { + struct node { + std::map children; // Use uint32_t to store Unicode codepoints + bool is_word; + }; + + std::vector nodes; + + common_trie() { + create_node(); // root node + } + + common_trie(const std::vector & words) : common_trie() { + for (const auto & w : words) { + insert(w); + } + } + + enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH }; + + // Check if a delimiter starts at the given position + match_result check_at(std::string_view sv, size_t start_pos) const; + + // Insert a word as a sequence of Unicode codepoints + void insert(const std::string & word); + + // Insert a raw symbol sequence + void insert(const std::vector & symbols); + + private: + size_t create_node() { + size_t index = nodes.size(); + nodes.emplace_back(); + return index; + } +}; + +// Aho-Corasick automaton +struct common_aho_corasick { + common_trie t; + std::vector fail; // failure links + std::vector order; // states in BFS order + std::vector terminal; // match states (directly or via a suffix link) + std::set alphabet; // every character with a transition + + common_aho_corasick(common_trie trie); + + common_aho_corasick(const std::vector & strings) + : common_aho_corasick(common_trie(strings)) {} + + size_t num_states() const { return t.nodes.size(); } + bool is_terminal(size_t s) const { return terminal[s]; } + + // follow failure links until a transition on `ch` exists. + size_t next(size_t state, uint32_t ch) const; +}; From c2321fab46150b54296a99726ef9f445bb88df80 Mon Sep 17 00:00:00 2001 From: Alde Rojas Date: Wed, 8 Jul 2026 17:33:39 -0500 Subject: [PATCH 2/8] common : support multiple token sequences in the reasoning budget sampler --- common/reasoning-budget.cpp | 72 +++++++++++++++++---------------- common/reasoning-budget.h | 22 +++++----- common/sampling.cpp | 4 +- tests/test-reasoning-budget.cpp | 69 ++++++++++++++++++++++--------- 4 files changed, 101 insertions(+), 66 deletions(-) diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index 7da0bb1c57ce..5b224cfb110b 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -1,5 +1,6 @@ #include "reasoning-budget.h" #include "common.h" +#include "trie.h" #include "unicode.h" #include "log.h" @@ -10,30 +11,31 @@ #include struct token_matcher { - std::vector tokens; - size_t pos = 0; + common_aho_corasick ac; + size_t state = 0; - bool advance(llama_token token) { - if (tokens.empty()) { - return false; - } + token_matcher(const std::vector & seqs) : ac(build_trie(seqs)) {} - if (token == tokens[pos]) { - pos++; - if (pos >= tokens.size()) { - pos = 0; - return true; - } - } else { - pos = 0; - if (token == tokens[0]) { - pos = 1; + static common_trie build_trie(const std::vector & seqs) { + common_trie t; + for (const auto & seq : seqs) { + if (!seq.empty()) { + t.insert(std::vector(seq.begin(), seq.end())); } } + return t; + } + + bool advance(llama_token token) { + state = ac.next(state, (uint32_t) token); + if (ac.is_terminal(state)) { + state = 0; + return true; + } return false; } - void reset() { pos = 0; } + void reset() { state = 0; } }; struct common_reasoning_budget_ctx { @@ -41,7 +43,7 @@ struct common_reasoning_budget_ctx { token_matcher start_matcher; token_matcher end_matcher; - std::vector forced_tokens; + llama_tokens forced_tokens; int32_t budget; // maximum tokens in reasoning block int32_t remaining; // tokens remaining in budget @@ -172,8 +174,8 @@ static void common_reasoning_budget_reset(struct llama_sampler * smpl) { } static struct llama_sampler * common_reasoning_budget_init_state( - const struct llama_vocab * vocab, const std::vector & start_tokens, - const std::vector & end_tokens, const std::vector & forced_tokens, + const struct llama_vocab * vocab, const std::vector & start_seqs, + const std::vector & end_seqs, const llama_tokens & forced_tokens, int32_t budget, common_reasoning_budget_state initial_state); static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl); @@ -205,12 +207,12 @@ static struct llama_sampler * common_reasoning_budget_clone(const struct llama_s } static struct llama_sampler * common_reasoning_budget_init_state( - const struct llama_vocab * vocab, - const std::vector & start_tokens, - const std::vector & end_tokens, - const std::vector & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state) { + const struct llama_vocab * vocab, + const std::vector & start_seqs, + const std::vector & end_seqs, + const llama_tokens & forced_tokens, + int32_t budget, + common_reasoning_budget_state initial_state) { // promote COUNTING with budget <= 0 to FORCING if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) { initial_state = REASONING_BUDGET_FORCING; @@ -220,8 +222,8 @@ static struct llama_sampler * common_reasoning_budget_init_state( /* .iface = */ &common_reasoning_budget_i, /* .ctx = */ new common_reasoning_budget_ctx { /* .vocab = */ vocab, - /* .start_matcher = */ { start_tokens, 0 }, - /* .end_matcher = */ { end_tokens, 0 }, + /* .start_matcher = */ token_matcher(start_seqs), + /* .end_matcher = */ token_matcher(end_seqs), /* .forced_tokens = */ forced_tokens, /* .budget = */ budget, /* .remaining = */ budget, @@ -232,13 +234,13 @@ static struct llama_sampler * common_reasoning_budget_init_state( } struct llama_sampler * common_reasoning_budget_init( - const struct llama_vocab * vocab, - const std::vector & start_tokens, - const std::vector & end_tokens, - const std::vector & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state) { - return common_reasoning_budget_init_state(vocab, start_tokens, end_tokens, forced_tokens, budget, initial_state); + const struct llama_vocab * vocab, + const std::vector & start_seqs, + const std::vector & end_seqs, + const llama_tokens & forced_tokens, + int32_t budget, + common_reasoning_budget_state initial_state) { + return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, budget, initial_state); } common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) { diff --git a/common/reasoning-budget.h b/common/reasoning-budget.h index 0cf689a56637..8b62b3face62 100644 --- a/common/reasoning-budget.h +++ b/common/reasoning-budget.h @@ -2,6 +2,8 @@ #include "llama.h" +#include "common.h" + #include #include @@ -17,27 +19,27 @@ enum common_reasoning_budget_state { // reasoning block (e.g. between and ). // // State machine: IDLE -> COUNTING -> WAITING_UTF8 -> FORCING -> DONE -// IDLE: passthrough, watching for start_tokens sequence -// COUNTING: counting down remaining tokens, watching for natural end_tokens +// IDLE: passthrough, watching for a start sequence +// COUNTING: counting down remaining tokens, watching for a natural end sequence // WAITING_UTF8: budget exhausted, allowing tokens to complete a UTF-8 sequence // FORCING: forces forced_tokens token-by-token (all other logits -> -inf) // DONE: passthrough forever // // Parameters: // vocab - vocabulary (used for UTF-8 boundary detection; can be nullptr) -// start_tokens - token sequence that activates counting -// end_tokens - token sequence for natural deactivation +// start_seqs - token sequences, any of which activates counting +// end_seqs - token sequences, any of which naturally deactivates // forced_tokens - token sequence forced when budget expires // budget - max tokens allowed in the reasoning block // initial_state - initial state // struct llama_sampler * common_reasoning_budget_init( - const struct llama_vocab * vocab, - const std::vector & start_tokens, - const std::vector & end_tokens, - const std::vector & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE); + const struct llama_vocab * vocab, + const std::vector & start_seqs, + const std::vector & end_seqs, + const llama_tokens & forced_tokens, + int32_t budget, + common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE); common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl); diff --git a/common/sampling.cpp b/common/sampling.cpp index 75a299e23ece..5673ffa472f8 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -299,8 +299,8 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, st if (!params.reasoning_budget_start.empty() && !params.reasoning_budget_end.empty() && (params.grammar_lazy || params.reasoning_budget_tokens >= 0 || params.reasoning_control)) { rbudget = common_reasoning_budget_init( vocab, - params.reasoning_budget_start, - params.reasoning_budget_end, + {params.reasoning_budget_start}, + {params.reasoning_budget_end}, params.reasoning_budget_forced, params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens); diff --git a/tests/test-reasoning-budget.cpp b/tests/test-reasoning-budget.cpp index f54cff4f8a27..717aa8180b13 100644 --- a/tests/test-reasoning-budget.cpp +++ b/tests/test-reasoning-budget.cpp @@ -20,8 +20,8 @@ static void test_reasoning_budget( const char * test_name, const std::vector & sequence, - const std::vector & start_tokens, - const std::vector & end_tokens, + const std::vector & start_seqs, + const std::vector & end_seqs, const std::vector & forced_tokens, int32_t budget, common_reasoning_budget_state initial_state, @@ -31,8 +31,8 @@ static void test_reasoning_budget( // Find the maximum token ID to ensure our vocab covers all tokens llama_token max_token = 0; for (auto t : sequence) max_token = std::max(max_token, t); - for (auto t : start_tokens) max_token = std::max(max_token, t); - for (auto t : end_tokens) max_token = std::max(max_token, t); + for (const auto & seq : start_seqs) for (auto t : seq) max_token = std::max(max_token, t); + for (const auto & seq : end_seqs) for (auto t : seq) max_token = std::max(max_token, t); for (auto t : forced_tokens) max_token = std::max(max_token, t); // Create a minimal sampler with mock vocabulary @@ -40,8 +40,8 @@ static void test_reasoning_budget( // The UTF-8 boundary check will treat all tokens as complete (safe fallback) auto * sampler = common_reasoning_budget_init( nullptr, // vocab - not used for basic state machine tests - start_tokens, - end_tokens, + start_seqs, + end_seqs, forced_tokens, budget, initial_state @@ -152,7 +152,7 @@ static void test_reasoning_budget_clone_mid_counting() { const std::vector end = {101}; const std::vector forced = {102, 101}; - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 2, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 2, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING, remaining=2 llama_sampler_accept(sampler, 50); // COUNTING, remaining=1 @@ -171,7 +171,7 @@ static void test_reasoning_budget_clone_mid_forcing() { const std::vector end = {101}; const std::vector forced = {102, 101}; - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING); GGML_ASSERT(get_forced_token(sampler, 102) == 102); llama_sampler_accept(sampler, 102); // advance to the second forced token @@ -191,7 +191,7 @@ static void test_reasoning_budget_force_manual() { // if COUNTING, force() succeeds and begins forcing the end sequence from the start { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING, remaining=5 llama_sampler_accept(sampler, 50); // COUNTING, remaining=4 @@ -212,7 +212,7 @@ static void test_reasoning_budget_force_manual() { // if IDLE, force() is a no-op { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); GGML_ASSERT(!common_reasoning_budget_force(sampler) && "force() must not transition from IDLE"); GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_IDLE); @@ -222,7 +222,7 @@ static void test_reasoning_budget_force_manual() { // if DONE, force() is a no-op { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING llama_sampler_accept(sampler, 101); // natural end -> DONE @@ -236,7 +236,7 @@ static void test_reasoning_budget_force_manual() { // if FORCING, force() is a no-op and must not rewind the force position { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING); GGML_ASSERT(get_forced_token(sampler, 102) == 102); llama_sampler_accept(sampler, 102); // advance to the second forced token (force_pos=1) @@ -290,7 +290,7 @@ int main(void) { const std::vector forced = {102}; // forced token (not used in this test) const std::vector sequence = {100, 50, 51, 101, 52}; // start, two tokens, end, one more - test_reasoning_budget("natural end before budget exhausted", sequence, start, end, forced, + test_reasoning_budget("natural end before budget exhausted", sequence, {start}, {end}, forced, 5, // budget of 5 tokens REASONING_BUDGET_IDLE, SIZE_MAX, SIZE_MAX); // no forcing expected (natural end) @@ -306,7 +306,7 @@ int main(void) { const std::vector forced = {102, 101}; // forced message + end const std::vector sequence = {100, 50, 51, 52, 53}; // start + 4 tokens (budget=2) - test_reasoning_budget("budget exhausted forcing", sequence, start, end, forced, + test_reasoning_budget("budget exhausted forcing", sequence, {start}, {end}, forced, 2, // budget of 2 tokens REASONING_BUDGET_IDLE, 3, // forcing starts at i=3 (accept at i=2 depletes budget, apply at i=3 forces) @@ -321,7 +321,7 @@ int main(void) { const std::vector forced = {102, 101}; const std::vector sequence = {100, 50, 51, 52}; // start token first, then 3 tokens - test_reasoning_budget("activate immediately budget=0", sequence, start, end, forced, + test_reasoning_budget("activate immediately budget=0", sequence, {start}, {end}, forced, 0, // budget of 0 tokens REASONING_BUDGET_COUNTING, // starts counting, promoted to FORCING since budget=0 0, // forcing starts at i=0 (initialized in FORCING, apply forces immediately) @@ -335,7 +335,7 @@ int main(void) { const std::vector forced = {102}; const std::vector sequence = {50, 51, 52, 53}; - test_reasoning_budget("no start/end configured", sequence, start, end, forced, + test_reasoning_budget("no start/end configured", sequence, {start}, {end}, forced, 2, // budget REASONING_BUDGET_IDLE, SIZE_MAX, SIZE_MAX); // no forcing (no start/end configured) @@ -350,7 +350,7 @@ int main(void) { const std::vector forced = {102, 101}; const std::vector sequence = {50, 51, 52, 53}; - test_reasoning_budget("activate immediately with budget", sequence, start, end, forced, + test_reasoning_budget("activate immediately with budget", sequence, {start}, {end}, forced, 2, // budget of 2 tokens REASONING_BUDGET_COUNTING, 2, // forcing starts at i=2 (after 2 accepts deplete budget, apply at i=2 forces) @@ -373,18 +373,49 @@ int main(void) { const std::vector forced = {102, 101}; const std::vector sequence = {100, 50, 101, 100, 60, 61, 62, 63}; - test_reasoning_budget("multi-block re-arms budget after DONE", sequence, start, end, forced, + test_reasoning_budget("multi-block re-arms budget after DONE", sequence, {start}, {end}, forced, 2, // budget of 2 tokens (per block) REASONING_BUDGET_IDLE, 6, // forcing starts at i=6 (after second block exhausts at i=5) 7); // forcing continues through i=7 } + // Test 7: Multiple start sequences - the second sequence activates counting + // Flow: i=0 accept(110), i=1 accept(111)->COUNTING rem=2; i=2 accept(50)->rem=1; + // i=3 accept(51)->rem=0->FORCING; i=4..5 apply() forces the end sequence + { + const std::vector start = {{100}, {110, 111}}; + const std::vector end = {{101}}; + const std::vector forced = {102, 101}; + const std::vector sequence = {110, 111, 50, 51, 52, 53}; + + test_reasoning_budget("multiple start sequences", sequence, start, end, forced, + 2, // budget of 2 tokens + REASONING_BUDGET_IDLE, + 4, // forcing starts at i=4 (accept at i=3 depletes budget) + 5); // forcing continues through i=5 + } + + // Test 8: Multiple end sequences - natural end via the second sequence + // Flow: i=0 accept(100)->COUNTING rem=5; i=1 accept(50)->rem=4; + // i=2 accept(103)->partial end, rem=3; i=3 accept(104)->end matched, DONE + { + const std::vector start = {{100}}; + const std::vector end = {{101}, {103, 104}}; + const std::vector forced = {102, 101}; + const std::vector sequence = {100, 50, 103, 104, 52}; + + test_reasoning_budget("multiple end sequences", sequence, start, end, forced, + 5, // budget of 5 tokens + REASONING_BUDGET_IDLE, + SIZE_MAX, SIZE_MAX); // no forcing expected (natural end) + } + test_reasoning_budget_clone_mid_counting(); test_reasoning_budget_clone_mid_forcing(); test_reasoning_budget_force_manual(); - printf("OK (9 tests passed)\n"); + printf("OK (11 tests passed)\n"); printf("Testing UTF-8 boundary detection... "); test_utf8_boundary_detection(); From ea9aa31b47456f8e69afaf4f94eef4e0f26f7d83 Mon Sep 17 00:00:00 2001 From: Alde Rojas Date: Wed, 8 Jul 2026 19:02:35 -0500 Subject: [PATCH 3/8] common/trie : return matched word index --- common/reasoning-budget.cpp | 35 ++++++++++++++++--------- common/trie.cpp | 19 +++++++++----- common/trie.h | 28 ++++++++++++-------- tests/test-reasoning-budget.cpp | 46 +++++++++++++++++++++++++++++++++ 4 files changed, 98 insertions(+), 30 deletions(-) diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index 5b224cfb110b..75e6cf3ec2ae 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -5,34 +5,45 @@ #include "log.h" +#include #include #include #include #include struct token_matcher { + std::vector words; // unique non-empty seqs, indices match the trie word indices common_aho_corasick ac; size_t state = 0; - token_matcher(const std::vector & seqs) : ac(build_trie(seqs)) {} + token_matcher(const std::vector & seqs) : words(collect(seqs)), ac(build_trie(words)) {} - static common_trie build_trie(const std::vector & seqs) { - common_trie t; + static std::vector collect(const std::vector & seqs) { + std::vector words; for (const auto & seq : seqs) { - if (!seq.empty()) { - t.insert(std::vector(seq.begin(), seq.end())); + if (!seq.empty() && std::find(words.begin(), words.end(), seq) == words.end()) { + words.push_back(seq); } } + return words; + } + + static common_trie build_trie(const std::vector & words) { + common_trie t; + for (const auto & w : words) { + t.insert(std::vector(w.begin(), w.end())); + } return t; } - bool advance(llama_token token) { + // returns the index into words of the longest sequence ending at this token, or -1 + int32_t advance(llama_token token) { state = ac.next(state, (uint32_t) token); - if (ac.is_terminal(state)) { + const int32_t w = ac.match_word(state); + if (w >= 0) { state = 0; - return true; } - return false; + return w; } void reset() { state = 0; } @@ -64,7 +75,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to switch (ctx->state) { case REASONING_BUDGET_IDLE: { - if (ctx->start_matcher.advance(token)) { + if (ctx->start_matcher.advance(token) >= 0) { ctx->state = REASONING_BUDGET_COUNTING; ctx->remaining = ctx->budget; COM_TRC("activated, budget=%d tokens\n", ctx->budget); @@ -80,7 +91,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to case REASONING_BUDGET_COUNTING: case REASONING_BUDGET_WAITING_UTF8: { - if (ctx->end_matcher.advance(token)) { + if (ctx->end_matcher.advance(token) >= 0) { ctx->state = REASONING_BUDGET_DONE; COM_TRC("%s", "deactivated (natural end)\n"); break; @@ -126,7 +137,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to case REASONING_BUDGET_DONE: // Re-arm on a new start tag: some models emit multiple blocks // per response, and each should get a fresh budget window. - if (ctx->start_matcher.advance(token)) { + if (ctx->start_matcher.advance(token) >= 0) { ctx->state = REASONING_BUDGET_COUNTING; ctx->remaining = ctx->budget; ctx->end_matcher.reset(); diff --git a/common/trie.cpp b/common/trie.cpp index a5d03fb1c83e..9b84748da4a0 100644 --- a/common/trie.cpp +++ b/common/trie.cpp @@ -26,7 +26,7 @@ common_trie::match_result common_trie::check_at(std::string_view sv, size_t star pos += result.bytes_consumed; // Check if we've matched a complete word - if (nodes[current].is_word) { + if (nodes[current].word >= 0) { return match_result{match_result::COMPLETE_MATCH}; } } @@ -41,7 +41,7 @@ common_trie::match_result common_trie::check_at(std::string_view sv, size_t star return match_result{match_result::NO_MATCH}; } -void common_trie::insert(const std::string & word) { +int32_t common_trie::insert(const std::string & word) { std::vector symbols; size_t pos = 0; while (pos < word.length()) { @@ -53,10 +53,10 @@ void common_trie::insert(const std::string & word) { symbols.push_back(result.codepoint); pos += result.bytes_consumed; } - insert(symbols); + return insert(symbols); } -void common_trie::insert(const std::vector & symbols) { +int32_t common_trie::insert(const std::vector & symbols) { size_t current = 0; for (uint32_t ch : symbols) { auto it = nodes[current].children.find(ch); @@ -68,7 +68,10 @@ void common_trie::insert(const std::vector & symbols) { current = it->second; } } - nodes[current].is_word = true; + if (nodes[current].word < 0) { + nodes[current].word = n_words++; + } + return nodes[current].word; } common_aho_corasick::common_aho_corasick(common_trie trie) : t(std::move(trie)) { @@ -96,9 +99,11 @@ common_aho_corasick::common_aho_corasick(common_trie trie) : t(std::move(trie)) } } - terminal.assign(n, false); + // fail[u] points to a strictly shorter suffix, so the first word found on + // the fail chain (including u itself) is the longest word ending at u + match.assign(n, -1); for (size_t u : order) { - terminal[u] = nodes[u].is_word || (u != 0 && terminal[fail[u]]); + match[u] = nodes[u].word >= 0 ? nodes[u].word : (u != 0 ? match[fail[u]] : -1); } for (const auto & node : nodes) { diff --git a/common/trie.h b/common/trie.h index bf03471e7e7c..87b339c063a2 100644 --- a/common/trie.h +++ b/common/trie.h @@ -12,7 +12,7 @@ struct common_trie { struct node { std::map children; // Use uint32_t to store Unicode codepoints - bool is_word; + int32_t word = -1; // index of the word ending at this node, -1 if none }; std::vector nodes; @@ -32,13 +32,16 @@ struct common_trie { // Check if a delimiter starts at the given position match_result check_at(std::string_view sv, size_t start_pos) const; - // Insert a word as a sequence of Unicode codepoints - void insert(const std::string & word); + // Insert a word as a sequence of Unicode codepoints, returns its word index + int32_t insert(const std::string & word); - // Insert a raw symbol sequence - void insert(const std::vector & symbols); + // Insert a raw symbol sequence, returns its word index (insertion order, + // duplicates keep the first index) + int32_t insert(const std::vector & symbols); private: + int32_t n_words = 0; + size_t create_node() { size_t index = nodes.size(); nodes.emplace_back(); @@ -48,11 +51,11 @@ struct common_trie { // Aho-Corasick automaton struct common_aho_corasick { - common_trie t; - std::vector fail; // failure links - std::vector order; // states in BFS order - std::vector terminal; // match states (directly or via a suffix link) - std::set alphabet; // every character with a transition + common_trie t; + std::vector fail; // failure links + std::vector order; // states in BFS order + std::vector match; // longest word ending at each state (directly or via a suffix link), -1 if none + std::set alphabet; // every character with a transition common_aho_corasick(common_trie trie); @@ -60,7 +63,10 @@ struct common_aho_corasick { : common_aho_corasick(common_trie(strings)) {} size_t num_states() const { return t.nodes.size(); } - bool is_terminal(size_t s) const { return terminal[s]; } + bool is_terminal(size_t s) const { return match[s] >= 0; } + + // word index of the longest word ending at this state, -1 if none + int32_t match_word(size_t s) const { return match[s]; } // follow failure links until a transition on `ch` exists. size_t next(size_t state, uint32_t ch) const; diff --git a/tests/test-reasoning-budget.cpp b/tests/test-reasoning-budget.cpp index 717aa8180b13..da8fce6eb746 100644 --- a/tests/test-reasoning-budget.cpp +++ b/tests/test-reasoning-budget.cpp @@ -1,4 +1,5 @@ #include "reasoning-budget.h" +#include "trie.h" #include "unicode.h" #include "llama.h" @@ -254,6 +255,47 @@ static void test_reasoning_budget_force_manual() { fprintf(stderr, " Test 'manual force transition' passed\n"); } +// Aho-Corasick match reporting unit test +// match_word() must report the longest word ending at the current position +static void test_aho_corasick_match_word() { + // duplicate inserts keep the first word index + { + common_trie t; + GGML_ASSERT(t.insert(std::string("ab")) == 0); + GGML_ASSERT(t.insert(std::string("cd")) == 1); + GGML_ASSERT(t.insert(std::string("ab")) == 0); + } + + auto feed = [](const common_aho_corasick & ac, const std::string & input) { + size_t state = 0; + for (char c : input) { + state = ac.next(state, (uint32_t) c); + } + return state; + }; + + // match via a suffix link: after "abc" the state is not a word, but "bc" ends there + { + common_aho_corasick ac({"bc", "abcd"}); + GGML_ASSERT(ac.match_word(feed(ac, "ab")) == -1); + GGML_ASSERT(ac.match_word(feed(ac, "abc")) == 0); + GGML_ASSERT(ac.match_word(feed(ac, "xbc")) == 0); + } + + // multiple words end at the same position: the longest wins + { + common_aho_corasick ac({"c", "abc"}); + GGML_ASSERT(ac.match_word(feed(ac, "abc")) == 1); + GGML_ASSERT(ac.match_word(feed(ac, "xc")) == 0); + } + + // an empty word makes the start state terminal (used by gbnf_ac_grammar) + { + common_aho_corasick ac(std::vector{""}); + GGML_ASSERT(ac.is_terminal(0)); + } +} + // UTF-8 boundary detection unit test // Tests common_utf8_is_complete() from reasoning-budget.h static void test_utf8_boundary_detection() { @@ -417,6 +459,10 @@ int main(void) { printf("OK (11 tests passed)\n"); + printf("Testing Aho-Corasick match reporting... "); + test_aho_corasick_match_word(); + printf("OK\n"); + printf("Testing UTF-8 boundary detection... "); test_utf8_boundary_detection(); printf("OK\n"); From 186241ab7f8206113107d1add6361baf3a788d9f Mon Sep 17 00:00:00 2001 From: Alde Rojas Date: Wed, 8 Jul 2026 19:03:38 -0500 Subject: [PATCH 4/8] common/trie : rename "word" to "pattern" --- common/reasoning-budget.cpp | 26 +++++++++++++------------- common/trie.cpp | 14 +++++++------- common/trie.h | 14 +++++++------- tests/test-reasoning-budget.cpp | 24 ++++++++++++------------ 4 files changed, 39 insertions(+), 39 deletions(-) diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index 75e6cf3ec2ae..130d61c6062e 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -12,38 +12,38 @@ #include struct token_matcher { - std::vector words; // unique non-empty seqs, indices match the trie word indices + std::vector seqs; common_aho_corasick ac; size_t state = 0; - token_matcher(const std::vector & seqs) : words(collect(seqs)), ac(build_trie(words)) {} + token_matcher(const std::vector & seqs) : seqs(collect(seqs)), ac(build_trie(this->seqs)) {} static std::vector collect(const std::vector & seqs) { - std::vector words; + std::vector res; for (const auto & seq : seqs) { - if (!seq.empty() && std::find(words.begin(), words.end(), seq) == words.end()) { - words.push_back(seq); + if (!seq.empty() && std::find(res.begin(), res.end(), seq) == res.end()) { + res.push_back(seq); } } - return words; + return res; } - static common_trie build_trie(const std::vector & words) { + static common_trie build_trie(const std::vector & seqs) { common_trie t; - for (const auto & w : words) { - t.insert(std::vector(w.begin(), w.end())); + for (const auto & seq : seqs) { + t.insert(std::vector(seq.begin(), seq.end())); } return t; } - // returns the index into words of the longest sequence ending at this token, or -1 + // returns the index into seqs of the longest sequence ending at this token, or -1 int32_t advance(llama_token token) { state = ac.next(state, (uint32_t) token); - const int32_t w = ac.match_word(state); - if (w >= 0) { + const int32_t p = ac.match_pattern(state); + if (p >= 0) { state = 0; } - return w; + return p; } void reset() { state = 0; } diff --git a/common/trie.cpp b/common/trie.cpp index 9b84748da4a0..b5c9666ba2ee 100644 --- a/common/trie.cpp +++ b/common/trie.cpp @@ -26,7 +26,7 @@ common_trie::match_result common_trie::check_at(std::string_view sv, size_t star pos += result.bytes_consumed; // Check if we've matched a complete word - if (nodes[current].word >= 0) { + if (nodes[current].pattern >= 0) { return match_result{match_result::COMPLETE_MATCH}; } } @@ -68,10 +68,10 @@ int32_t common_trie::insert(const std::vector & symbols) { current = it->second; } } - if (nodes[current].word < 0) { - nodes[current].word = n_words++; + if (nodes[current].pattern < 0) { + nodes[current].pattern = n_patterns++; } - return nodes[current].word; + return nodes[current].pattern; } common_aho_corasick::common_aho_corasick(common_trie trie) : t(std::move(trie)) { @@ -99,11 +99,11 @@ common_aho_corasick::common_aho_corasick(common_trie trie) : t(std::move(trie)) } } - // fail[u] points to a strictly shorter suffix, so the first word found on - // the fail chain (including u itself) is the longest word ending at u + // fail[u] points to a strictly shorter suffix, so the first pattern found on + // the fail chain (including u itself) is the longest pattern ending at u match.assign(n, -1); for (size_t u : order) { - match[u] = nodes[u].word >= 0 ? nodes[u].word : (u != 0 ? match[fail[u]] : -1); + match[u] = nodes[u].pattern >= 0 ? nodes[u].pattern : (u != 0 ? match[fail[u]] : -1); } for (const auto & node : nodes) { diff --git a/common/trie.h b/common/trie.h index 87b339c063a2..0f7b16a36ad2 100644 --- a/common/trie.h +++ b/common/trie.h @@ -12,7 +12,7 @@ struct common_trie { struct node { std::map children; // Use uint32_t to store Unicode codepoints - int32_t word = -1; // index of the word ending at this node, -1 if none + int32_t pattern = -1; // index of the pattern ending at this node, -1 if none }; std::vector nodes; @@ -32,15 +32,15 @@ struct common_trie { // Check if a delimiter starts at the given position match_result check_at(std::string_view sv, size_t start_pos) const; - // Insert a word as a sequence of Unicode codepoints, returns its word index + // Insert a word as a sequence of Unicode codepoints, returns its pattern index int32_t insert(const std::string & word); - // Insert a raw symbol sequence, returns its word index (insertion order, + // Insert a raw symbol sequence, returns its pattern index (insertion order, // duplicates keep the first index) int32_t insert(const std::vector & symbols); private: - int32_t n_words = 0; + int32_t n_patterns = 0; size_t create_node() { size_t index = nodes.size(); @@ -54,7 +54,7 @@ struct common_aho_corasick { common_trie t; std::vector fail; // failure links std::vector order; // states in BFS order - std::vector match; // longest word ending at each state (directly or via a suffix link), -1 if none + std::vector match; // longest pattern ending at each state (directly or via a suffix link), -1 if none std::set alphabet; // every character with a transition common_aho_corasick(common_trie trie); @@ -65,8 +65,8 @@ struct common_aho_corasick { size_t num_states() const { return t.nodes.size(); } bool is_terminal(size_t s) const { return match[s] >= 0; } - // word index of the longest word ending at this state, -1 if none - int32_t match_word(size_t s) const { return match[s]; } + // index of the longest pattern ending at this state, -1 if none + int32_t match_pattern(size_t s) const { return match[s]; } // follow failure links until a transition on `ch` exists. size_t next(size_t state, uint32_t ch) const; diff --git a/tests/test-reasoning-budget.cpp b/tests/test-reasoning-budget.cpp index da8fce6eb746..6becab749755 100644 --- a/tests/test-reasoning-budget.cpp +++ b/tests/test-reasoning-budget.cpp @@ -256,9 +256,9 @@ static void test_reasoning_budget_force_manual() { } // Aho-Corasick match reporting unit test -// match_word() must report the longest word ending at the current position -static void test_aho_corasick_match_word() { - // duplicate inserts keep the first word index +// match_pattern() must report the longest pattern ending at the current position +static void test_aho_corasick_match_pattern() { + // duplicate inserts keep the first pattern index { common_trie t; GGML_ASSERT(t.insert(std::string("ab")) == 0); @@ -274,22 +274,22 @@ static void test_aho_corasick_match_word() { return state; }; - // match via a suffix link: after "abc" the state is not a word, but "bc" ends there + // match via a suffix link: after "abc" the state is not a pattern, but "bc" ends there { common_aho_corasick ac({"bc", "abcd"}); - GGML_ASSERT(ac.match_word(feed(ac, "ab")) == -1); - GGML_ASSERT(ac.match_word(feed(ac, "abc")) == 0); - GGML_ASSERT(ac.match_word(feed(ac, "xbc")) == 0); + GGML_ASSERT(ac.match_pattern(feed(ac, "ab")) == -1); + GGML_ASSERT(ac.match_pattern(feed(ac, "abc")) == 0); + GGML_ASSERT(ac.match_pattern(feed(ac, "xbc")) == 0); } - // multiple words end at the same position: the longest wins + // multiple patterns end at the same position: the longest wins { common_aho_corasick ac({"c", "abc"}); - GGML_ASSERT(ac.match_word(feed(ac, "abc")) == 1); - GGML_ASSERT(ac.match_word(feed(ac, "xc")) == 0); + GGML_ASSERT(ac.match_pattern(feed(ac, "abc")) == 1); + GGML_ASSERT(ac.match_pattern(feed(ac, "xc")) == 0); } - // an empty word makes the start state terminal (used by gbnf_ac_grammar) + // an empty pattern makes the start state terminal (used by gbnf_ac_grammar) { common_aho_corasick ac(std::vector{""}); GGML_ASSERT(ac.is_terminal(0)); @@ -460,7 +460,7 @@ int main(void) { printf("OK (11 tests passed)\n"); printf("Testing Aho-Corasick match reporting... "); - test_aho_corasick_match_word(); + test_aho_corasick_match_pattern(); printf("OK\n"); printf("Testing UTF-8 boundary detection... "); From d44a6c80e7fcb8872c75e55b058578c02790be25 Mon Sep 17 00:00:00 2001 From: Alde Rojas Date: Wed, 8 Jul 2026 19:14:34 -0500 Subject: [PATCH 5/8] common/reasoning-budget : expose matched end sequence --- common/reasoning-budget.cpp | 27 +++++++++++- common/reasoning-budget.h | 4 ++ tests/test-reasoning-budget.cpp | 78 ++++++++++++++++++++++++++++++++- 3 files changed, 107 insertions(+), 2 deletions(-) diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index 130d61c6062e..c890e34e9792 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -63,6 +63,8 @@ struct common_reasoning_budget_ctx { // for forcing size_t force_pos; // next position in forced_tokens to force + + int32_t matched_end; // index into end_matcher.seqs of the sequence that transitioned to DONE, -1 if none }; static const char * common_reasoning_budget_name(const struct llama_sampler * /*smpl*/) { @@ -91,8 +93,10 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to case REASONING_BUDGET_COUNTING: case REASONING_BUDGET_WAITING_UTF8: { - if (ctx->end_matcher.advance(token) >= 0) { + const int32_t match = ctx->end_matcher.advance(token); + if (match >= 0) { ctx->state = REASONING_BUDGET_DONE; + ctx->matched_end = match; COM_TRC("%s", "deactivated (natural end)\n"); break; } @@ -128,12 +132,17 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to break; } case REASONING_BUDGET_FORCING: + { + // track the end sequence within forced_tokens so it is also reported on DONE + const int32_t match = ctx->end_matcher.advance(token); ctx->force_pos++; if (ctx->force_pos >= ctx->forced_tokens.size()) { ctx->state = REASONING_BUDGET_DONE; + ctx->matched_end = match; COM_TRC("%s", "forced sequence complete, done\n"); } break; + } case REASONING_BUDGET_DONE: // Re-arm on a new start tag: some models emit multiple blocks // per response, and each should get a fresh budget window. @@ -141,6 +150,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to ctx->state = REASONING_BUDGET_COUNTING; ctx->remaining = ctx->budget; ctx->end_matcher.reset(); + ctx->matched_end = -1; COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget); if (ctx->remaining <= 0) { @@ -182,6 +192,7 @@ static void common_reasoning_budget_reset(struct llama_sampler * smpl) { ctx->start_matcher.reset(); ctx->end_matcher.reset(); ctx->force_pos = 0; + ctx->matched_end = -1; } static struct llama_sampler * common_reasoning_budget_init_state( @@ -240,6 +251,7 @@ static struct llama_sampler * common_reasoning_budget_init_state( /* .remaining = */ budget, /* .state = */ initial_state, /* .force_pos = */ 0, + /* .matched_end = */ -1, } ); } @@ -261,6 +273,19 @@ common_reasoning_budget_state common_reasoning_budget_get_state(const struct lla return ((const common_reasoning_budget_ctx *)smpl->ctx)->state; } +const llama_tokens * common_reasoning_budget_get_matched_end(const struct llama_sampler * smpl) { + if (!smpl) { + return nullptr; + } + + const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx; + if (ctx->matched_end < 0) { + return nullptr; + } + + return &ctx->end_matcher.seqs[ctx->matched_end]; +} + bool common_reasoning_budget_force(struct llama_sampler * smpl) { if (!smpl) { return false; diff --git a/common/reasoning-budget.h b/common/reasoning-budget.h index 8b62b3face62..93b616409330 100644 --- a/common/reasoning-budget.h +++ b/common/reasoning-budget.h @@ -43,6 +43,10 @@ struct llama_sampler * common_reasoning_budget_init( common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl); +// The end sequence that transitioned the sampler to DONE, or nullptr if none +// was recorded. Cleared when a new start sequence re-arms the sampler. +const llama_tokens * common_reasoning_budget_get_matched_end(const struct llama_sampler * smpl); + // Manually transition the reasoning budget sampler into the FORCING state. // Returns true if the transition occurred. bool common_reasoning_budget_force(struct llama_sampler * smpl); diff --git a/tests/test-reasoning-budget.cpp b/tests/test-reasoning-budget.cpp index 6becab749755..6132a072c5e4 100644 --- a/tests/test-reasoning-budget.cpp +++ b/tests/test-reasoning-budget.cpp @@ -255,6 +255,81 @@ static void test_reasoning_budget_force_manual() { fprintf(stderr, " Test 'manual force transition' passed\n"); } +static void test_reasoning_budget_matched_end() { + const std::vector start = {{100}}; + const std::vector end = {{101}, {103, 104}}; + + // natural end records the sequence that matched; re-arming clears it + { + auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 101}, 5, REASONING_BUDGET_IDLE); + + GGML_ASSERT(common_reasoning_budget_get_matched_end(sampler) == nullptr); + + llama_sampler_accept(sampler, 100); // COUNTING + llama_sampler_accept(sampler, 50); + llama_sampler_accept(sampler, 103); + llama_sampler_accept(sampler, 104); // end matched via {103, 104}, DONE + + const llama_tokens * matched = common_reasoning_budget_get_matched_end(sampler); + GGML_ASSERT(matched != nullptr); + GGML_ASSERT(*matched == llama_tokens({103, 104})); + + llama_sampler_accept(sampler, 100); // re-arm, COUNTING + GGML_ASSERT(common_reasoning_budget_get_matched_end(sampler) == nullptr); + + llama_sampler_free(sampler); + } + + // overlapping end sequences: the longest one ending at the position wins + { + const std::vector end_overlap = {{104}, {103, 104}}; + + auto * sampler = common_reasoning_budget_init(nullptr, start, end_overlap, {102, 104}, 5, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); // COUNTING + llama_sampler_accept(sampler, 103); + llama_sampler_accept(sampler, 104); // both {104} and {103, 104} end here + + const llama_tokens * matched = common_reasoning_budget_get_matched_end(sampler); + GGML_ASSERT(matched != nullptr); + GGML_ASSERT(*matched == llama_tokens({103, 104})); + + llama_sampler_free(sampler); + } + + // forcing records the end sequence terminating forced_tokens + { + auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 103, 104}, 0, REASONING_BUDGET_FORCING); + + llama_sampler_accept(sampler, 102); + llama_sampler_accept(sampler, 103); + GGML_ASSERT(common_reasoning_budget_get_matched_end(sampler) == nullptr); + llama_sampler_accept(sampler, 104); // forced sequence complete, DONE + + const llama_tokens * matched = common_reasoning_budget_get_matched_end(sampler); + GGML_ASSERT(matched != nullptr); + GGML_ASSERT(*matched == llama_tokens({103, 104})); + + llama_sampler_free(sampler); + } + + // forced_tokens not ending with a known end sequence records nothing + { + auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102}, 0, REASONING_BUDGET_FORCING); + + llama_sampler_accept(sampler, 102); // forced sequence complete, DONE + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE); + GGML_ASSERT(common_reasoning_budget_get_matched_end(sampler) == nullptr); + + llama_sampler_free(sampler); + } + + // a null sampler is safely ignored + GGML_ASSERT(common_reasoning_budget_get_matched_end(nullptr) == nullptr); + + fprintf(stderr, " Test 'matched end sequence' passed\n"); +} + // Aho-Corasick match reporting unit test // match_pattern() must report the longest pattern ending at the current position static void test_aho_corasick_match_pattern() { @@ -456,8 +531,9 @@ int main(void) { test_reasoning_budget_clone_mid_counting(); test_reasoning_budget_clone_mid_forcing(); test_reasoning_budget_force_manual(); + test_reasoning_budget_matched_end(); - printf("OK (11 tests passed)\n"); + printf("OK (12 tests passed)\n"); printf("Testing Aho-Corasick match reporting... "); test_aho_corasick_match_pattern(); From e77709a34ca1b7088fa2987454597eeaae0f92d8 Mon Sep 17 00:00:00 2001 From: Alde Rojas Date: Wed, 8 Jul 2026 19:38:00 -0500 Subject: [PATCH 6/8] common/sampling : replay end sequence when reasoning budget is done --- common/sampling.cpp | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/common/sampling.cpp b/common/sampling.cpp index 5673ffa472f8..6d5be31a76b6 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -453,6 +453,17 @@ void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, boo if (gsmpl->rbudget && is_generated) { llama_sampler_accept(gsmpl->rbudget, token); + + // if done, replay end sequence which may contain a grammar trigger + const bool is_done = common_reasoning_budget_get_state(gsmpl->rbudget) == REASONING_BUDGET_DONE; + if (gsmpl->grmr && !accept_grammar && is_done) { + const llama_tokens * end_seq = common_reasoning_budget_get_matched_end(gsmpl->rbudget); + if (end_seq) { + for (const llama_token end_token : *end_seq) { + llama_sampler_accept(gsmpl->grmr, end_token); + } + } + } } if (gsmpl->grmr && accept_grammar) { From fc1c28b3091226093b060cfa2618a6a03df64d00 Mon Sep 17 00:00:00 2001 From: Alde Rojas Date: Fri, 10 Jul 2026 15:41:36 -0500 Subject: [PATCH 7/8] cont : update to use multiple end sequences --- common/chat.cpp | 22 +++++++++++++-------- common/chat.h | 2 +- common/common.h | 12 ++++++------ common/sampling.cpp | 2 +- tests/test-chat.cpp | 16 +++++++++------- tools/server/server-common.cpp | 4 ++-- tools/server/server-schema.cpp | 35 ++++++++++++++++++++++++++-------- 7 files changed, 60 insertions(+), 33 deletions(-) diff --git a/common/chat.cpp b/common/chat.cpp index 22d2ee4a2a11..bde12c57fa61 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -1022,7 +1022,7 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_ data.supports_thinking = true; data.thinking_start_tag = "[THINK]"; - data.thinking_end_tag = "[/THINK]"; + data.thinking_end_tags = {"[/THINK]"}; data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override = */ adjusted_messages); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, /* messages_override = */ adjusted_messages); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; @@ -1148,6 +1148,9 @@ static common_chat_params common_chat_params_init_gpt_oss(const common_chat_temp data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; + data.thinking_start_tag = "<|channel|>analysis<|message|>"; + data.thinking_end_tags = {"<|end|>"}; + // These special tokens are required to parse properly, so we include them // even if parse_tool_calls is false. data.preserved_tokens = { @@ -1292,7 +1295,7 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ data.format = COMMON_CHAT_FORMAT_PEG_GEMMA4; data.supports_thinking = true; data.thinking_start_tag = "<|channel>thought"; - data.thinking_end_tag = ""; + data.thinking_end_tags = {""}; data.preserved_tokens = { "<|channel>", @@ -1567,7 +1570,7 @@ static common_chat_params common_chat_params_init_kimi_k2(const common_chat_temp const std::string GEN_PROMPT = "<|im_assistant|>assistant<|im_middle|>"; data.thinking_start_tag = THINK_START; - data.thinking_end_tag = THINK_END; + data.thinking_end_tags = {THINK_END}; if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; @@ -1701,7 +1704,7 @@ static common_chat_params common_chat_params_init_lfm2(const common_chat_templat } data.thinking_start_tag = THINK_START; - data.thinking_end_tag = THINK_END; + data.thinking_end_tags = {THINK_END}; auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); @@ -1864,7 +1867,7 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; data.thinking_start_tag = ""; - data.thinking_end_tag = ""; + data.thinking_end_tags = {""}; data.preserved_tokens = { "｜DSML｜", "", @@ -2077,7 +2080,7 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; data.thinking_start_tag = THINK_START; - data.thinking_end_tag = THINK_END; + data.thinking_end_tags = {THINK_END}; data.preserved_tokens = { TURN_START, TURN_END, CHATBOT, USER, SYSTEM, THINK_START, THINK_END, @@ -2418,7 +2421,7 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem }; data.thinking_start_tag = ""; - data.thinking_end_tag = ""; + data.thinking_end_tags = {""}; data.message_delimiters = { { COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant" }, @@ -2772,7 +2775,10 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_ auto_params.supports_thinking = autoparser.reasoning.mode != autoparser::reasoning_mode::NONE; if (auto_params.supports_thinking) { auto_params.thinking_start_tag = trim_whitespace(autoparser.reasoning.start); - auto_params.thinking_end_tag = trim_whitespace(autoparser.reasoning.end); + auto end_tag = trim_whitespace(autoparser.reasoning.end); + if (!end_tag.empty()) { + auto_params.thinking_end_tags = {std::move(end_tag)}; + } } common_peg_arena arena; arena.load(auto_params.parser); diff --git a/common/chat.h b/common/chat.h index 7898f1623f54..d79f4ecd773c 100644 --- a/common/chat.h +++ b/common/chat.h @@ -274,7 +274,7 @@ struct common_chat_params { std::string generation_prompt; bool supports_thinking = false; std::string thinking_start_tag; // e.g., "" - std::string thinking_end_tag; // e.g., "" + std::vector thinking_end_tags; // e.g., "" std::vector grammar_triggers; std::vector preserved_tokens; std::vector additional_stops; diff --git a/common/common.h b/common/common.h index bffc1767a7d1..1b883322cd55 100644 --- a/common/common.h +++ b/common/common.h @@ -283,12 +283,12 @@ struct common_params_sampling { // reasoning budget sampler parameters // these are populated by the server/CLI based on chat template params - int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget - std::vector reasoning_budget_start; // start tag token sequence - std::vector reasoning_budget_end; // end tag token sequence - std::vector reasoning_budget_forced; // forced sequence (message + end tag) - std::string reasoning_budget_message; // message injected before end tag when budget exhausted - bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime + int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget + std::vector reasoning_budget_start; // start tag token sequence + std::vector reasoning_budget_end; // end tag token sequences; the first tag is used as the forcing sequence + std::vector reasoning_budget_forced; // forced sequence (message + first end tag) + std::string reasoning_budget_message; // message injected before end tag when budget exhausted + bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime bool backend_sampling = false; diff --git a/common/sampling.cpp b/common/sampling.cpp index 6d5be31a76b6..84954130b7c4 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -300,7 +300,7 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, st rbudget = common_reasoning_budget_init( vocab, {params.reasoning_budget_start}, - {params.reasoning_budget_end}, + params.reasoning_budget_end, params.reasoning_budget_forced, params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens); diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 93685ec8ff42..a1e7feecaa57 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -1135,7 +1135,7 @@ static void test_peg_parser(common_chat_templates * tmpls, // budget sampler inhibits grammar application while inside thinking blocks — // triggers inside ... are suppressed. bool use_reasoning_budget_path = false; - if (parser.params_.grammar_lazy && !parser.params_.thinking_end_tag.empty()) { + if (parser.params_.grammar_lazy && !parser.params_.thinking_end_tags.empty()) { use_reasoning_budget_path = true; for (const auto & trigger : parser.params_.grammar_triggers) { if (trigger.type != COMMON_GRAMMAR_TRIGGER_TYPE_WORD) { @@ -1153,7 +1153,7 @@ static void test_peg_parser(common_chat_templates * tmpls, // Walk through full_input tracking thinking state; only match triggers // when outside thinking blocks. const auto & think_start = parser.params_.thinking_start_tag; - const auto & think_end = parser.params_.thinking_end_tag; + const auto & think_ends = parser.params_.thinking_end_tags; bool in_thinking = false; for (size_t i = 0; i < full_input.size(); ++i) { @@ -1163,12 +1163,14 @@ static void test_peg_parser(common_chat_templates * tmpls, i += think_start.size() - 1; continue; } - if (in_thinking && full_input.compare(i, think_end.size(), think_end) == 0) { - in_thinking = false; - i += think_end.size() - 1; - continue; - } if (in_thinking) { + for (const auto & think_end : think_ends) { + if (full_input.compare(i, think_end.size(), think_end) == 0) { + in_thinking = false; + i += think_end.size() - 1; + break; + } + } continue; } // Outside thinking — check if any trigger word starts here diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 78b5c6819b03..0ba22cb27391 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -1123,10 +1123,10 @@ json oaicompat_chat_params_parse( reasoning_budget = opt.reasoning_budget; } - if (!chat_params.thinking_end_tag.empty()) { + if (!chat_params.thinking_end_tags.empty()) { llama_params["reasoning_budget_tokens"] = reasoning_budget; llama_params["reasoning_budget_start_tag"] = chat_params.thinking_start_tag; - llama_params["reasoning_budget_end_tag"] = chat_params.thinking_end_tag; + llama_params["reasoning_budget_end_tags"] = chat_params.thinking_end_tags; llama_params["reasoning_budget_message"] = json_value(body, "reasoning_budget_message", opt.reasoning_budget_message); llama_params["reasoning_control"] = json_value(body, "reasoning_control", false); } diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index 89026eb4e3f0..e880f4ca728d 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -390,21 +390,40 @@ std::vector> make_llama_cmpl_schema(const common_params & ctx.params.sampling.reasoning_budget_start = common_tokenize(ctx.vocab, data.at("reasoning_budget_start_tag").get(), false, true); })); - add((new field_str("reasoning_budget_end_tag")) - ->set_desc("Token string marking the end of the reasoning budget section") + add((new field_json("reasoning_budget_end_tags")) + ->add_alias("reasoning_budget_end_tag") + ->set_desc("Token strings marking the end of the reasoning budget section; the first is forced when the budget expires") ->set_handler([&](field_eval_context & ctx, const json & data) { GGML_ASSERT(ctx.vocab != nullptr); - std::string end_tag = data.at("reasoning_budget_end_tag").get(); - ctx.params.sampling.reasoning_budget_end = common_tokenize(ctx.vocab, end_tag, false, true); + ctx.params.sampling.reasoning_budget_end.clear(); + if (data.contains("reasoning_budget_end_tags")) { + for (const auto & t : data.at("reasoning_budget_end_tags")) { + std::string tag = t.get(); + if (!tag.empty()) { + ctx.params.sampling.reasoning_budget_end.push_back(common_tokenize(ctx.vocab, tag, false, true)); + } + } + } else if (data.contains("reasoning_budget_end_tag")) { + std::string tag = data.at("reasoning_budget_end_tag").get(); + if (!tag.empty()) { + ctx.params.sampling.reasoning_budget_end.push_back(common_tokenize(ctx.vocab, tag, false, true)); + } + } })); add((new field_str("reasoning_budget_message")) ->set_desc("Message to prepend to the reasoning budget end tag when forcing it") ->set_handler([&](field_eval_context & ctx, const json & data) { GGML_ASSERT(ctx.vocab != nullptr); - std::string end_tag = json_value(data, "reasoning_budget_end_tag", std::string()); - std::string message = data.at("reasoning_budget_message").get(); - ctx.params.sampling.reasoning_budget_forced = common_tokenize(ctx.vocab, message + end_tag, false, true); + if (!ctx.params.sampling.reasoning_budget_end.empty()) { + llama_tokens end_tag = ctx.params.sampling.reasoning_budget_end.front(); + std::string message = json_value(data, "reasoning_budget_message", std::string()); + if (!message.empty()) { + llama_tokens message_tokens = common_tokenize(ctx.vocab, message, false, true); + end_tag.insert(end_tag.begin(), message_tokens.begin(), message_tokens.end()); + } + ctx.params.sampling.reasoning_budget_forced = std::move(end_tag); + } })); add((new field_json("logit_bias")) @@ -546,7 +565,7 @@ task_params eval_llama_cmpl_schema( // debugging { auto budget = params.sampling.reasoning_budget_tokens; - SRV_DBG("reasoning budget: tokens=%d, generation_prompt='%s', start=%zu toks, end=%zu toks, forced=%zu toks\n", + SRV_DBG("reasoning budget: tokens=%d, generation_prompt='%s', start=%zu toks, end=%zu seqs, forced=%zu toks\n", budget, params.sampling.generation_prompt.c_str(), params.sampling.reasoning_budget_start.size(), params.sampling.reasoning_budget_end.size(), From 72f4ebdf9453f2fda368f8f332bcf2449040b268 Mon Sep 17 00:00:00 2001 From: Alde Rojas Date: Fri, 10 Jul 2026 16:15:39 -0500 Subject: [PATCH 8/8] cont : clean up --- common/reasoning-budget.cpp | 18 ++++---- common/reasoning-budget.h | 2 +- common/sampling.cpp | 2 +- tests/test-reasoning-budget.cpp | 74 +++++++-------------------------- 4 files changed, 27 insertions(+), 69 deletions(-) diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index c890e34e9792..1fe242d062d1 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -64,7 +64,7 @@ struct common_reasoning_budget_ctx { // for forcing size_t force_pos; // next position in forced_tokens to force - int32_t matched_end; // index into end_matcher.seqs of the sequence that transitioned to DONE, -1 if none + int32_t end_match; // index into end_matcher.seqs of the sequence that transitioned to DONE, -1 if none }; static const char * common_reasoning_budget_name(const struct llama_sampler * /*smpl*/) { @@ -96,7 +96,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to const int32_t match = ctx->end_matcher.advance(token); if (match >= 0) { ctx->state = REASONING_BUDGET_DONE; - ctx->matched_end = match; + ctx->end_match = match; COM_TRC("%s", "deactivated (natural end)\n"); break; } @@ -138,7 +138,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to ctx->force_pos++; if (ctx->force_pos >= ctx->forced_tokens.size()) { ctx->state = REASONING_BUDGET_DONE; - ctx->matched_end = match; + ctx->end_match = match; COM_TRC("%s", "forced sequence complete, done\n"); } break; @@ -150,7 +150,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to ctx->state = REASONING_BUDGET_COUNTING; ctx->remaining = ctx->budget; ctx->end_matcher.reset(); - ctx->matched_end = -1; + ctx->end_match = -1; COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget); if (ctx->remaining <= 0) { @@ -192,7 +192,7 @@ static void common_reasoning_budget_reset(struct llama_sampler * smpl) { ctx->start_matcher.reset(); ctx->end_matcher.reset(); ctx->force_pos = 0; - ctx->matched_end = -1; + ctx->end_match = -1; } static struct llama_sampler * common_reasoning_budget_init_state( @@ -251,7 +251,7 @@ static struct llama_sampler * common_reasoning_budget_init_state( /* .remaining = */ budget, /* .state = */ initial_state, /* .force_pos = */ 0, - /* .matched_end = */ -1, + /* .end_match = */ -1, } ); } @@ -273,17 +273,17 @@ common_reasoning_budget_state common_reasoning_budget_get_state(const struct lla return ((const common_reasoning_budget_ctx *)smpl->ctx)->state; } -const llama_tokens * common_reasoning_budget_get_matched_end(const struct llama_sampler * smpl) { +const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl) { if (!smpl) { return nullptr; } const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx; - if (ctx->matched_end < 0) { + if (ctx->end_match < 0) { return nullptr; } - return &ctx->end_matcher.seqs[ctx->matched_end]; + return &ctx->end_matcher.seqs[ctx->end_match]; } bool common_reasoning_budget_force(struct llama_sampler * smpl) { diff --git a/common/reasoning-budget.h b/common/reasoning-budget.h index 93b616409330..1b89a04c42e8 100644 --- a/common/reasoning-budget.h +++ b/common/reasoning-budget.h @@ -45,7 +45,7 @@ common_reasoning_budget_state common_reasoning_budget_get_state(const struct lla // The end sequence that transitioned the sampler to DONE, or nullptr if none // was recorded. Cleared when a new start sequence re-arms the sampler. -const llama_tokens * common_reasoning_budget_get_matched_end(const struct llama_sampler * smpl); +const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl); // Manually transition the reasoning budget sampler into the FORCING state. // Returns true if the transition occurred. diff --git a/common/sampling.cpp b/common/sampling.cpp index 84954130b7c4..7b241e34f77f 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -457,7 +457,7 @@ void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, boo // if done, replay end sequence which may contain a grammar trigger const bool is_done = common_reasoning_budget_get_state(gsmpl->rbudget) == REASONING_BUDGET_DONE; if (gsmpl->grmr && !accept_grammar && is_done) { - const llama_tokens * end_seq = common_reasoning_budget_get_matched_end(gsmpl->rbudget); + const llama_tokens * end_seq = common_reasoning_budget_get_end_match(gsmpl->rbudget); if (end_seq) { for (const llama_token end_token : *end_seq) { llama_sampler_accept(gsmpl->grmr, end_token); diff --git a/tests/test-reasoning-budget.cpp b/tests/test-reasoning-budget.cpp index 6132a072c5e4..3bcc77e1733c 100644 --- a/tests/test-reasoning-budget.cpp +++ b/tests/test-reasoning-budget.cpp @@ -1,5 +1,4 @@ #include "reasoning-budget.h" -#include "trie.h" #include "unicode.h" #include "llama.h" @@ -32,8 +31,12 @@ static void test_reasoning_budget( // Find the maximum token ID to ensure our vocab covers all tokens llama_token max_token = 0; for (auto t : sequence) max_token = std::max(max_token, t); - for (const auto & seq : start_seqs) for (auto t : seq) max_token = std::max(max_token, t); - for (const auto & seq : end_seqs) for (auto t : seq) max_token = std::max(max_token, t); + for (const auto & seq : start_seqs) { + for (auto t : seq) max_token = std::max(max_token, t); + } + for (const auto & seq : end_seqs) { + for (auto t : seq) max_token = std::max(max_token, t); + } for (auto t : forced_tokens) max_token = std::max(max_token, t); // Create a minimal sampler with mock vocabulary @@ -255,7 +258,7 @@ static void test_reasoning_budget_force_manual() { fprintf(stderr, " Test 'manual force transition' passed\n"); } -static void test_reasoning_budget_matched_end() { +static void test_reasoning_budget_end_match() { const std::vector start = {{100}}; const std::vector end = {{101}, {103, 104}}; @@ -263,19 +266,19 @@ static void test_reasoning_budget_matched_end() { { auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 101}, 5, REASONING_BUDGET_IDLE); - GGML_ASSERT(common_reasoning_budget_get_matched_end(sampler) == nullptr); + GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); llama_sampler_accept(sampler, 100); // COUNTING llama_sampler_accept(sampler, 50); llama_sampler_accept(sampler, 103); llama_sampler_accept(sampler, 104); // end matched via {103, 104}, DONE - const llama_tokens * matched = common_reasoning_budget_get_matched_end(sampler); + const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler); GGML_ASSERT(matched != nullptr); GGML_ASSERT(*matched == llama_tokens({103, 104})); llama_sampler_accept(sampler, 100); // re-arm, COUNTING - GGML_ASSERT(common_reasoning_budget_get_matched_end(sampler) == nullptr); + GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); llama_sampler_free(sampler); } @@ -290,7 +293,7 @@ static void test_reasoning_budget_matched_end() { llama_sampler_accept(sampler, 103); llama_sampler_accept(sampler, 104); // both {104} and {103, 104} end here - const llama_tokens * matched = common_reasoning_budget_get_matched_end(sampler); + const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler); GGML_ASSERT(matched != nullptr); GGML_ASSERT(*matched == llama_tokens({103, 104})); @@ -303,10 +306,10 @@ static void test_reasoning_budget_matched_end() { llama_sampler_accept(sampler, 102); llama_sampler_accept(sampler, 103); - GGML_ASSERT(common_reasoning_budget_get_matched_end(sampler) == nullptr); + GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); llama_sampler_accept(sampler, 104); // forced sequence complete, DONE - const llama_tokens * matched = common_reasoning_budget_get_matched_end(sampler); + const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler); GGML_ASSERT(matched != nullptr); GGML_ASSERT(*matched == llama_tokens({103, 104})); @@ -319,58 +322,17 @@ static void test_reasoning_budget_matched_end() { llama_sampler_accept(sampler, 102); // forced sequence complete, DONE GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE); - GGML_ASSERT(common_reasoning_budget_get_matched_end(sampler) == nullptr); + GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); llama_sampler_free(sampler); } // a null sampler is safely ignored - GGML_ASSERT(common_reasoning_budget_get_matched_end(nullptr) == nullptr); + GGML_ASSERT(common_reasoning_budget_get_end_match(nullptr) == nullptr); fprintf(stderr, " Test 'matched end sequence' passed\n"); } -// Aho-Corasick match reporting unit test -// match_pattern() must report the longest pattern ending at the current position -static void test_aho_corasick_match_pattern() { - // duplicate inserts keep the first pattern index - { - common_trie t; - GGML_ASSERT(t.insert(std::string("ab")) == 0); - GGML_ASSERT(t.insert(std::string("cd")) == 1); - GGML_ASSERT(t.insert(std::string("ab")) == 0); - } - - auto feed = [](const common_aho_corasick & ac, const std::string & input) { - size_t state = 0; - for (char c : input) { - state = ac.next(state, (uint32_t) c); - } - return state; - }; - - // match via a suffix link: after "abc" the state is not a pattern, but "bc" ends there - { - common_aho_corasick ac({"bc", "abcd"}); - GGML_ASSERT(ac.match_pattern(feed(ac, "ab")) == -1); - GGML_ASSERT(ac.match_pattern(feed(ac, "abc")) == 0); - GGML_ASSERT(ac.match_pattern(feed(ac, "xbc")) == 0); - } - - // multiple patterns end at the same position: the longest wins - { - common_aho_corasick ac({"c", "abc"}); - GGML_ASSERT(ac.match_pattern(feed(ac, "abc")) == 1); - GGML_ASSERT(ac.match_pattern(feed(ac, "xc")) == 0); - } - - // an empty pattern makes the start state terminal (used by gbnf_ac_grammar) - { - common_aho_corasick ac(std::vector{""}); - GGML_ASSERT(ac.is_terminal(0)); - } -} - // UTF-8 boundary detection unit test // Tests common_utf8_is_complete() from reasoning-budget.h static void test_utf8_boundary_detection() { @@ -531,14 +493,10 @@ int main(void) { test_reasoning_budget_clone_mid_counting(); test_reasoning_budget_clone_mid_forcing(); test_reasoning_budget_force_manual(); - test_reasoning_budget_matched_end(); + test_reasoning_budget_end_match(); printf("OK (12 tests passed)\n"); - printf("Testing Aho-Corasick match reporting... "); - test_aho_corasick_match_pattern(); - printf("OK\n"); - printf("Testing UTF-8 boundary detection... "); test_utf8_boundary_detection(); printf("OK\n");