Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

common : add support for multiple end sequences in the reasoning budget sampler (#25544)

Browse Source 
* common : extract trie/ac to a separate file

* common : support multiple token sequences in the reasoning budget sampler

* common/trie : return matched word index

* common/trie : rename "word" to "pattern"

* common/reasoning-budget : expose matched end sequence

* common/sampling : replay end sequence when reasoning budget is done

* cont : update to use multiple end sequences

* cont : clean up
This commit is contained in:
Aldehir Rojas
2026-07-25 11:58:09 +02:00
committed by GitHub
parent d67c0b4107
commit 910196f6b3
14 changed files with 497 additions and 254 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/CMakeLists.txt
+2
View File
@@ -100,6 +100,8 @@ add_library(${TARGET}
sampling.h sampling.h
speculative.cpp speculative.cpp
speculative.h speculative.h
trie.cpp
trie.h
unicode.cpp unicode.cpp
unicode.h unicode.h
jinja/lexer.cpp jinja/lexer.cpp
common/chat.cpp
+14 -8
View File
@@ -1024,7 +1024,7 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_
data.supports_thinking = true; data.supports_thinking = true;
data.thinking_start_tag = "[THINK]"; 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.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.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, /* messages_override = */ adjusted_messages);
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
@@ -1150,6 +1150,9 @@ static common_chat_params common_chat_params_init_gpt_oss(const common_chat_temp
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
data.supports_thinking = true; 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 // These special tokens are required to parse properly, so we include them
// even if parse_tool_calls is false. // even if parse_tool_calls is false.
data.preserved_tokens = { data.preserved_tokens = {
@@ -1294,7 +1297,7 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ
data.format = COMMON_CHAT_FORMAT_PEG_GEMMA4; data.format = COMMON_CHAT_FORMAT_PEG_GEMMA4;
data.supports_thinking = true; data.supports_thinking = true;
data.thinking_start_tag = "<|channel>thought"; data.thinking_start_tag = "<|channel>thought";
data.thinking_end_tag = "<channel|>"; data.thinking_end_tags = {"<channel|>"};
data.preserved_tokens = { data.preserved_tokens = {
"<|channel>", "<|channel>",
@@ -1569,7 +1572,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|>"; const std::string GEN_PROMPT = "<|im_assistant|>assistant<|im_middle|>";
data.thinking_start_tag = THINK_START; data.thinking_start_tag = THINK_START;
data.thinking_end_tag = THINK_END; data.thinking_end_tags = {THINK_END};
if (inputs.has_continuation()) { if (inputs.has_continuation()) {
const auto & msg = inputs.continue_msg; const auto & msg = inputs.continue_msg;
@@ -1703,7 +1706,7 @@ static common_chat_params common_chat_params_init_lfm2(const common_chat_templat
} }
data.thinking_start_tag = THINK_START; 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_tools = inputs.tools.is_array() && !inputs.tools.empty();
auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object();
@@ -1943,7 +1946,7 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
data.supports_thinking = true; data.supports_thinking = true;
data.thinking_start_tag = "<think>"; data.thinking_start_tag = "<think>";
data.thinking_end_tag = "</think>"; data.thinking_end_tags = {"</think>"};
data.preserved_tokens = { data.preserved_tokens = {
"｜DSML｜", "｜DSML｜",
"<think>", "<think>",
@@ -2160,7 +2163,7 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
data.supports_thinking = true; data.supports_thinking = true;
data.thinking_start_tag = THINK_START; data.thinking_start_tag = THINK_START;
data.thinking_end_tag = THINK_END; data.thinking_end_tags = {THINK_END};
data.preserved_tokens = { data.preserved_tokens = {
TURN_START, TURN_END, CHATBOT, USER, SYSTEM, TURN_START, TURN_END, CHATBOT, USER, SYSTEM,
THINK_START, THINK_END, THINK_START, THINK_END,
@@ -2510,7 +2513,7 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem
}; };
data.thinking_start_tag = "<think>"; data.thinking_start_tag = "<think>";
data.thinking_end_tag = "</think>"; data.thinking_end_tags = {"</think>"};
data.message_delimiters = { data.message_delimiters = {
{ COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant" }, { COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant" },
@@ -2866,7 +2869,10 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_
auto_params.supports_thinking = autoparser.reasoning.mode != autoparser::reasoning_mode::NONE; auto_params.supports_thinking = autoparser.reasoning.mode != autoparser::reasoning_mode::NONE;
if (auto_params.supports_thinking) { if (auto_params.supports_thinking) {
auto_params.thinking_start_tag = trim_whitespace(autoparser.reasoning.start); 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; common_peg_arena arena;
arena.load(auto_params.parser); arena.load(auto_params.parser);
common/chat.h
+1 -1
View File
@@ -274,7 +274,7 @@ struct common_chat_params {
std::string generation_prompt; std::string generation_prompt;
bool supports_thinking = false; bool supports_thinking = false;
std::string thinking_start_tag; // e.g., "<think>" std::string thinking_start_tag; // e.g., "<think>"
std::string thinking_end_tag; // e.g., "</think>" std::vector<std::string> thinking_end_tags; // e.g., "</think>"
std::vector<common_grammar_trigger> grammar_triggers; std::vector<common_grammar_trigger> grammar_triggers;
std::vector<std::string> preserved_tokens; std::vector<std::string> preserved_tokens;
std::vector<std::string> additional_stops; std::vector<std::string> additional_stops;
common/common.h
+6 -6
View File
@@ -284,12 +284,12 @@ struct common_params_sampling {
// reasoning budget sampler parameters // reasoning budget sampler parameters
// these are populated by the server/CLI based on chat template params // these are populated by the server/CLI based on chat template params
int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget
std::vector<llama_token> reasoning_budget_start; // start tag token sequence std::vector<llama_token> reasoning_budget_start; // start tag token sequence
std::vector<llama_token> reasoning_budget_end; // end tag token sequence std::vector<llama_tokens> reasoning_budget_end; // end tag token sequences; the first tag is used as the forcing sequence
std::vector<llama_token> reasoning_budget_forced; // forced sequence (message + end tag) std::vector<llama_token> reasoning_budget_forced; // forced sequence (message + first end tag)
std::string reasoning_budget_message; // message injected before end tag when budget exhausted 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 reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime
bool backend_sampling = false; bool backend_sampling = false;
common/peg-parser.cpp
+5 -153
View File
@@ -3,10 +3,10 @@
#include "common.h" #include "common.h"
#include "json-schema-to-grammar.h" #include "json-schema-to-grammar.h"
#include "log.h" #include "log.h"
#include "trie.h"
#include "unicode.h" #include "unicode.h"
#include <algorithm> #include <algorithm>
#include <deque>
#include <initializer_list> #include <initializer_list>
#include <map> #include <map>
#include <memory> #include <memory>
@@ -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'); 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<uint32_t, size_t> children; // Use uint32_t to store Unicode codepoints
bool is_word;
};
std::vector<node> nodes;
trie(const std::vector<std::string> & 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<size_t> fail; // failure links
std::vector<size_t> order; // states in BFS order
std::vector<bool> terminal; // match states (directly or via a suffix link)
std::set<uint32_t> alphabet; // every character with a transition
aho_corasick(const std::vector<std::string> & strings) : t(strings) {
const auto & nodes = t.nodes;
const size_t n = nodes.size();
fail.assign(n, 0);
order.reserve(n);
std::deque<size_t> 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<uint32_t, size_t> parse_hex_escape(const std::string & str, size_t pos, int hex_count) { static std::pair<uint32_t, size_t> parse_hex_escape(const std::string & str, size_t pos, int hex_count) {
if (pos + hex_count > str.length()) { if (pos + hex_count > str.length()) {
return {0, 0}; return {0, 0};
@@ -797,7 +649,7 @@ struct parser_executor {
} }
common_peg_parse_result operator()(const common_peg_until_parser & p) const { 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 // Scan input and check for delimiters
size_t pos = start_pos; size_t pos = start_pos;
@@ -824,12 +676,12 @@ struct parser_executor {
// Check if a delimiter starts at this position // Check if a delimiter starts at this position
auto match = matcher.check_at(ctx.input, pos); 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 // Found a complete delimiter, return everything before it
return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); 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 // 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); 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<size_t, std::vector<uint32_t>> &, const std::map<size_t, std::vector<uint32_t>> &,
const std::vector<uint32_t> &, const std::vector<uint32_t> &,
const std::function<std::string(size_t)> &)> & build_rule) { const std::function<std::string(size_t)> &)> & build_rule) {
aho_corasick ac(strings); common_aho_corasick ac(strings);
auto state_name = [&](size_t s) -> std::string { auto state_name = [&](size_t s) -> std::string {
if (s == 0) { if (s == 0) {
common/reasoning-budget.cpp
+77 -39
View File
@@ -1,39 +1,52 @@
#include "reasoning-budget.h" #include "reasoning-budget.h"
#include "common.h" #include "common.h"
#include "trie.h"
#include "unicode.h" #include "unicode.h"
#include "log.h" #include "log.h"
#include <algorithm>
#include <cmath> #include <cmath>
#include <cstdint> #include <cstdint>
#include <string> #include <string>
#include <vector> #include <vector>
struct token_matcher { struct token_matcher {
std::vector<llama_token> tokens; std::vector<llama_tokens> seqs;
size_t pos = 0; common_aho_corasick ac;
size_t state = 0;
bool advance(llama_token token) { token_matcher(const std::vector<llama_tokens> & seqs) : seqs(collect(seqs)), ac(build_trie(this->seqs)) {}
if (tokens.empty()) {
return false;
}
if (token == tokens[pos]) { static std::vector<llama_tokens> collect(const std::vector<llama_tokens> & seqs) {
pos++; std::vector<llama_tokens> res;
if (pos >= tokens.size()) { for (const auto & seq : seqs) {
pos = 0; if (!seq.empty() && std::find(res.begin(), res.end(), seq) == res.end()) {
return true; res.push_back(seq);
}
} else {
pos = 0;
if (token == tokens[0]) {
pos = 1;
} }
} }
return false; return res;
} }
void reset() { pos = 0; } static common_trie build_trie(const std::vector<llama_tokens> & seqs) {
common_trie t;
for (const auto & seq : seqs) {
t.insert(std::vector<uint32_t>(seq.begin(), seq.end()));
}
return t;
}
// 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 p = ac.match_pattern(state);
if (p >= 0) {
state = 0;
}
return p;
}
void reset() { state = 0; }
}; };
struct common_reasoning_budget_ctx { struct common_reasoning_budget_ctx {
@@ -41,7 +54,7 @@ struct common_reasoning_budget_ctx {
token_matcher start_matcher; token_matcher start_matcher;
token_matcher end_matcher; token_matcher end_matcher;
std::vector<llama_token> forced_tokens; llama_tokens forced_tokens;
int32_t budget; // maximum tokens in reasoning block int32_t budget; // maximum tokens in reasoning block
int32_t remaining; // tokens remaining in budget int32_t remaining; // tokens remaining in budget
@@ -50,6 +63,8 @@ struct common_reasoning_budget_ctx {
// for forcing // for forcing
size_t force_pos; // next position in forced_tokens to force size_t force_pos; // next position in forced_tokens to force
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*/) { static const char * common_reasoning_budget_name(const struct llama_sampler * /*smpl*/) {
@@ -62,7 +77,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
switch (ctx->state) { switch (ctx->state) {
case REASONING_BUDGET_IDLE: case REASONING_BUDGET_IDLE:
{ {
if (ctx->start_matcher.advance(token)) { if (ctx->start_matcher.advance(token) >= 0) {
ctx->state = REASONING_BUDGET_COUNTING; ctx->state = REASONING_BUDGET_COUNTING;
ctx->remaining = ctx->budget; ctx->remaining = ctx->budget;
COM_TRC("activated, budget=%d tokens\n", ctx->budget); COM_TRC("activated, budget=%d tokens\n", ctx->budget);
@@ -78,8 +93,10 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
case REASONING_BUDGET_COUNTING: case REASONING_BUDGET_COUNTING:
case REASONING_BUDGET_WAITING_UTF8: case REASONING_BUDGET_WAITING_UTF8:
{ {
if (ctx->end_matcher.advance(token)) { const int32_t match = ctx->end_matcher.advance(token);
if (match >= 0) {
ctx->state = REASONING_BUDGET_DONE; ctx->state = REASONING_BUDGET_DONE;
ctx->end_match = match;
COM_TRC("%s", "deactivated (natural end)\n"); COM_TRC("%s", "deactivated (natural end)\n");
break; break;
} }
@@ -115,19 +132,25 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
break; break;
} }
case REASONING_BUDGET_FORCING: 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++; ctx->force_pos++;
if (ctx->force_pos >= ctx->forced_tokens.size()) { if (ctx->force_pos >= ctx->forced_tokens.size()) {
ctx->state = REASONING_BUDGET_DONE; ctx->state = REASONING_BUDGET_DONE;
ctx->end_match = match;
COM_TRC("%s", "forced sequence complete, done\n"); COM_TRC("%s", "forced sequence complete, done\n");
} }
break; break;
}
case REASONING_BUDGET_DONE: case REASONING_BUDGET_DONE:
// Re-arm on a new start tag: some models emit multiple <think> blocks // Re-arm on a new start tag: some models emit multiple <think> blocks
// per response, and each should get a fresh budget window. // 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->state = REASONING_BUDGET_COUNTING;
ctx->remaining = ctx->budget; ctx->remaining = ctx->budget;
ctx->end_matcher.reset(); ctx->end_matcher.reset();
ctx->end_match = -1;
COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget); COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget);
if (ctx->remaining <= 0) { if (ctx->remaining <= 0) {
@@ -169,11 +192,12 @@ static void common_reasoning_budget_reset(struct llama_sampler * smpl) {
ctx->start_matcher.reset(); ctx->start_matcher.reset();
ctx->end_matcher.reset(); ctx->end_matcher.reset();
ctx->force_pos = 0; ctx->force_pos = 0;
ctx->end_match = -1;
} }
static struct llama_sampler * common_reasoning_budget_init_state( static struct llama_sampler * common_reasoning_budget_init_state(
const struct llama_vocab * vocab, const std::vector<llama_token> & start_tokens, const struct llama_vocab * vocab, const std::vector<llama_tokens> & start_seqs,
const std::vector<llama_token> & end_tokens, const std::vector<llama_token> & forced_tokens, const std::vector<llama_tokens> & end_seqs, const llama_tokens & forced_tokens,
int32_t budget, common_reasoning_budget_state initial_state); int32_t budget, common_reasoning_budget_state initial_state);
static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl); static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl);
@@ -205,12 +229,12 @@ static struct llama_sampler * common_reasoning_budget_clone(const struct llama_s
} }
static struct llama_sampler * common_reasoning_budget_init_state( static struct llama_sampler * common_reasoning_budget_init_state(
const struct llama_vocab * vocab, const struct llama_vocab * vocab,
const std::vector<llama_token> & start_tokens, const std::vector<llama_tokens> & start_seqs,
const std::vector<llama_token> & end_tokens, const std::vector<llama_tokens> & end_seqs,
const std::vector<llama_token> & forced_tokens, const llama_tokens & forced_tokens,
int32_t budget, int32_t budget,
common_reasoning_budget_state initial_state) { common_reasoning_budget_state initial_state) {
// promote COUNTING with budget <= 0 to FORCING // promote COUNTING with budget <= 0 to FORCING
if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) { if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) {
initial_state = REASONING_BUDGET_FORCING; initial_state = REASONING_BUDGET_FORCING;
@@ -220,25 +244,26 @@ static struct llama_sampler * common_reasoning_budget_init_state(
/* .iface = */ &common_reasoning_budget_i, /* .iface = */ &common_reasoning_budget_i,
/* .ctx = */ new common_reasoning_budget_ctx { /* .ctx = */ new common_reasoning_budget_ctx {
/* .vocab = */ vocab, /* .vocab = */ vocab,
/* .start_matcher = */ { start_tokens, 0 }, /* .start_matcher = */ token_matcher(start_seqs),
/* .end_matcher = */ { end_tokens, 0 }, /* .end_matcher = */ token_matcher(end_seqs),
/* .forced_tokens = */ forced_tokens, /* .forced_tokens = */ forced_tokens,
/* .budget = */ budget, /* .budget = */ budget,
/* .remaining = */ budget, /* .remaining = */ budget,
/* .state = */ initial_state, /* .state = */ initial_state,
/* .force_pos = */ 0, /* .force_pos = */ 0,
/* .end_match = */ -1,
} }
); );
} }
struct llama_sampler * common_reasoning_budget_init( struct llama_sampler * common_reasoning_budget_init(
const struct llama_vocab * vocab, const struct llama_vocab * vocab,
const std::vector<llama_token> & start_tokens, const std::vector<llama_tokens> & start_seqs,
const std::vector<llama_token> & end_tokens, const std::vector<llama_tokens> & end_seqs,
const std::vector<llama_token> & forced_tokens, const llama_tokens & forced_tokens,
int32_t budget, int32_t budget,
common_reasoning_budget_state initial_state) { common_reasoning_budget_state initial_state) {
return common_reasoning_budget_init_state(vocab, start_tokens, end_tokens, forced_tokens, budget, 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) { common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) {
@@ -248,6 +273,19 @@ common_reasoning_budget_state common_reasoning_budget_get_state(const struct lla
return ((const common_reasoning_budget_ctx *)smpl->ctx)->state; return ((const common_reasoning_budget_ctx *)smpl->ctx)->state;
} }
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->end_match < 0) {
return nullptr;
}
return &ctx->end_matcher.seqs[ctx->end_match];
}
bool common_reasoning_budget_force(struct llama_sampler * smpl) { bool common_reasoning_budget_force(struct llama_sampler * smpl) {
if (!smpl) { if (!smpl) {
return false; return false;
common/reasoning-budget.h
+16 -10
View File
@@ -2,6 +2,8 @@
#include "llama.h" #include "llama.h"
#include "common.h"
#include <cstdint> #include <cstdint>
#include <vector> #include <vector>
@@ -17,30 +19,34 @@ enum common_reasoning_budget_state {
// reasoning block (e.g. between <think> and </think>). // reasoning block (e.g. between <think> and </think>).
// //
// State machine: IDLE -> COUNTING -> WAITING_UTF8 -> FORCING -> DONE // State machine: IDLE -> COUNTING -> WAITING_UTF8 -> FORCING -> DONE
// IDLE: passthrough, watching for start_tokens sequence // IDLE: passthrough, watching for a start sequence
// COUNTING: counting down remaining tokens, watching for natural end_tokens // COUNTING: counting down remaining tokens, watching for a natural end sequence
// WAITING_UTF8: budget exhausted, allowing tokens to complete a UTF-8 sequence // WAITING_UTF8: budget exhausted, allowing tokens to complete a UTF-8 sequence
// FORCING: forces forced_tokens token-by-token (all other logits -> -inf) // FORCING: forces forced_tokens token-by-token (all other logits -> -inf)
// DONE: passthrough forever // DONE: passthrough forever
// //
// Parameters: // Parameters:
// vocab - vocabulary (used for UTF-8 boundary detection; can be nullptr) // vocab - vocabulary (used for UTF-8 boundary detection; can be nullptr)
// start_tokens - token sequence that activates counting // start_seqs - token sequences, any of which activates counting
// end_tokens - token sequence for natural deactivation // end_seqs - token sequences, any of which naturally deactivates
// forced_tokens - token sequence forced when budget expires // forced_tokens - token sequence forced when budget expires
// budget - max tokens allowed in the reasoning block // budget - max tokens allowed in the reasoning block
// initial_state - initial state // initial_state - initial state
// //
struct llama_sampler * common_reasoning_budget_init( struct llama_sampler * common_reasoning_budget_init(
const struct llama_vocab * vocab, const struct llama_vocab * vocab,
const std::vector<llama_token> & start_tokens, const std::vector<llama_tokens> & start_seqs,
const std::vector<llama_token> & end_tokens, const std::vector<llama_tokens> & end_seqs,
const std::vector<llama_token> & forced_tokens, const llama_tokens & forced_tokens,
int32_t budget, int32_t budget,
common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE); common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE);
common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl); 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_end_match(const struct llama_sampler * smpl);
// Manually transition the reasoning budget sampler into the FORCING state. // Manually transition the reasoning budget sampler into the FORCING state.
// Returns true if the transition occurred. // Returns true if the transition occurred.
bool common_reasoning_budget_force(struct llama_sampler * smpl); bool common_reasoning_budget_force(struct llama_sampler * smpl);
common/sampling.cpp
+12 -1
View File
@@ -299,7 +299,7 @@ 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)) { 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( rbudget = common_reasoning_budget_init(
vocab, vocab,
params.reasoning_budget_start, {params.reasoning_budget_start},
params.reasoning_budget_end, params.reasoning_budget_end,
params.reasoning_budget_forced, params.reasoning_budget_forced,
params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens); params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens);
@@ -453,6 +453,17 @@ void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, boo
if (gsmpl->rbudget && is_generated) { if (gsmpl->rbudget && is_generated) {
llama_sampler_accept(gsmpl->rbudget, token); 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_end_match(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) { if (gsmpl->grmr && accept_grammar) {
common/trie.cpp
+123
View File
@@ -0,0 +1,123 @@
#include "trie.h"
#include "unicode.h"
#include <deque>
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].pattern >= 0) {
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};
}
int32_t common_trie::insert(const std::string & word) {
std::vector<uint32_t> 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;
}
return insert(symbols);
}
int32_t common_trie::insert(const std::vector<uint32_t> & 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;
}
}
if (nodes[current].pattern < 0) {
nodes[current].pattern = n_patterns++;
}
return nodes[current].pattern;
}
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<size_t> 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);
}
}
// 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].pattern >= 0 ? nodes[u].pattern : (u != 0 ? match[fail[u]] : -1);
}
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;
}
common/trie.h
+73
View File
@@ -0,0 +1,73 @@
#pragma once
#include <cstdint>
#include <map>
#include <set>
#include <string>
#include <string_view>
#include <vector>
// 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<uint32_t, size_t> children; // Use uint32_t to store Unicode codepoints
int32_t pattern = -1; // index of the pattern ending at this node, -1 if none
};
std::vector<node> nodes;
common_trie() {
create_node(); // root node
}
common_trie(const std::vector<std::string> & 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, returns its pattern index
int32_t insert(const std::string & word);
// Insert a raw symbol sequence, returns its pattern index (insertion order,
// duplicates keep the first index)
int32_t insert(const std::vector<uint32_t> & symbols);
private:
int32_t n_patterns = 0;
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<size_t> fail; // failure links
std::vector<size_t> order; // states in BFS order
std::vector<int32_t> match; // longest pattern ending at each state (directly or via a suffix link), -1 if none
std::set<uint32_t> alphabet; // every character with a transition
common_aho_corasick(common_trie trie);
common_aho_corasick(const std::vector<std::string> & strings)
: common_aho_corasick(common_trie(strings)) {}
size_t num_states() const { return t.nodes.size(); }
bool is_terminal(size_t s) const { return match[s] >= 0; }
// 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;
};
tests/test-chat.cpp
+9 -7
View File
@@ -1144,7 +1144,7 @@ static void test_peg_parser(common_chat_templates * tmpls,
// budget sampler inhibits grammar application while inside thinking blocks — // budget sampler inhibits grammar application while inside thinking blocks —
// triggers inside <think>...</think> are suppressed. // triggers inside <think>...</think> are suppressed.
bool use_reasoning_budget_path = false; 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; use_reasoning_budget_path = true;
for (const auto & trigger : parser.params_.grammar_triggers) { for (const auto & trigger : parser.params_.grammar_triggers) {
if (trigger.type != COMMON_GRAMMAR_TRIGGER_TYPE_WORD) { if (trigger.type != COMMON_GRAMMAR_TRIGGER_TYPE_WORD) {
@@ -1162,7 +1162,7 @@ static void test_peg_parser(common_chat_templates * tmpls,
// Walk through full_input tracking thinking state; only match triggers // Walk through full_input tracking thinking state; only match triggers
// when outside thinking blocks. // when outside thinking blocks.
const auto & think_start = parser.params_.thinking_start_tag; 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; bool in_thinking = false;
for (size_t i = 0; i < full_input.size(); ++i) { for (size_t i = 0; i < full_input.size(); ++i) {
@@ -1172,12 +1172,14 @@ static void test_peg_parser(common_chat_templates * tmpls,
i += think_start.size() - 1; i += think_start.size() - 1;
continue; 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) { 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; continue;
} }
// Outside thinking — check if any trigger word starts here // Outside thinking — check if any trigger word starts here
tests/test-reasoning-budget.cpp
+130 -19
View File
@@ -20,8 +20,8 @@
static void test_reasoning_budget( static void test_reasoning_budget(
const char * test_name, const char * test_name,
const std::vector<llama_token> & sequence, const std::vector<llama_token> & sequence,
const std::vector<llama_token> & start_tokens, const std::vector<llama_tokens> & start_seqs,
const std::vector<llama_token> & end_tokens, const std::vector<llama_tokens> & end_seqs,
const std::vector<llama_token> & forced_tokens, const std::vector<llama_token> & forced_tokens,
int32_t budget, int32_t budget,
common_reasoning_budget_state initial_state, common_reasoning_budget_state initial_state,
@@ -31,8 +31,12 @@ static void test_reasoning_budget(
// Find the maximum token ID to ensure our vocab covers all tokens // Find the maximum token ID to ensure our vocab covers all tokens
llama_token max_token = 0; llama_token max_token = 0;
for (auto t : sequence) max_token = std::max(max_token, t); for (auto t : sequence) max_token = std::max(max_token, t);
for (auto t : start_tokens) max_token = std::max(max_token, t); for (const auto & seq : start_seqs) {
for (auto t : end_tokens) max_token = std::max(max_token, t); 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); for (auto t : forced_tokens) max_token = std::max(max_token, t);
// Create a minimal sampler with mock vocabulary // Create a minimal sampler with mock vocabulary
@@ -40,8 +44,8 @@ static void test_reasoning_budget(
// The UTF-8 boundary check will treat all tokens as complete (safe fallback) // The UTF-8 boundary check will treat all tokens as complete (safe fallback)
auto * sampler = common_reasoning_budget_init( auto * sampler = common_reasoning_budget_init(
nullptr, // vocab - not used for basic state machine tests nullptr, // vocab - not used for basic state machine tests
start_tokens, start_seqs,
end_tokens, end_seqs,
forced_tokens, forced_tokens,
budget, budget,
initial_state initial_state
@@ -152,7 +156,7 @@ static void test_reasoning_budget_clone_mid_counting() {
const std::vector<llama_token> end = {101}; const std::vector<llama_token> end = {101};
const std::vector<llama_token> forced = {102, 101}; const std::vector<llama_token> 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, 100); // COUNTING, remaining=2
llama_sampler_accept(sampler, 50); // COUNTING, remaining=1 llama_sampler_accept(sampler, 50); // COUNTING, remaining=1
@@ -171,7 +175,7 @@ static void test_reasoning_budget_clone_mid_forcing() {
const std::vector<llama_token> end = {101}; const std::vector<llama_token> end = {101};
const std::vector<llama_token> forced = {102, 101}; const std::vector<llama_token> 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); GGML_ASSERT(get_forced_token(sampler, 102) == 102);
llama_sampler_accept(sampler, 102); // advance to the second forced token llama_sampler_accept(sampler, 102); // advance to the second forced token
@@ -191,7 +195,7 @@ static void test_reasoning_budget_force_manual() {
// if COUNTING, force() succeeds and begins forcing the end sequence from the start // 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, 100); // COUNTING, remaining=5
llama_sampler_accept(sampler, 50); // COUNTING, remaining=4 llama_sampler_accept(sampler, 50); // COUNTING, remaining=4
@@ -212,7 +216,7 @@ static void test_reasoning_budget_force_manual() {
// if IDLE, force() is a no-op // 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_force(sampler) && "force() must not transition from IDLE");
GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_IDLE); GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_IDLE);
@@ -222,7 +226,7 @@ static void test_reasoning_budget_force_manual() {
// if DONE, force() is a no-op // 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, 100); // COUNTING
llama_sampler_accept(sampler, 101); // natural end -> DONE llama_sampler_accept(sampler, 101); // natural end -> DONE
@@ -236,7 +240,7 @@ static void test_reasoning_budget_force_manual() {
// if FORCING, force() is a no-op and must not rewind the force position // 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); GGML_ASSERT(get_forced_token(sampler, 102) == 102);
llama_sampler_accept(sampler, 102); // advance to the second forced token (force_pos=1) llama_sampler_accept(sampler, 102); // advance to the second forced token (force_pos=1)
@@ -254,6 +258,81 @@ static void test_reasoning_budget_force_manual() {
fprintf(stderr, " Test 'manual force transition' passed\n"); fprintf(stderr, " Test 'manual force transition' passed\n");
} }
static void test_reasoning_budget_end_match() {
const std::vector<llama_tokens> start = {{100}};
const std::vector<llama_tokens> 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_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_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_end_match(sampler) == nullptr);
llama_sampler_free(sampler);
}
// overlapping end sequences: the longest one ending at the position wins
{
const std::vector<llama_tokens> 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_end_match(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_end_match(sampler) == nullptr);
llama_sampler_accept(sampler, 104); // forced sequence complete, DONE
const llama_tokens * matched = common_reasoning_budget_get_end_match(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_end_match(sampler) == nullptr);
llama_sampler_free(sampler);
}
// a null sampler is safely ignored
GGML_ASSERT(common_reasoning_budget_get_end_match(nullptr) == nullptr);
fprintf(stderr, " Test 'matched end sequence' passed\n");
}
// UTF-8 boundary detection unit test // UTF-8 boundary detection unit test
// Tests common_utf8_is_complete() from reasoning-budget.h // Tests common_utf8_is_complete() from reasoning-budget.h
static void test_utf8_boundary_detection() { static void test_utf8_boundary_detection() {
@@ -290,7 +369,7 @@ int main(void) {
const std::vector<llama_token> forced = {102}; // forced token (not used in this test) const std::vector<llama_token> forced = {102}; // forced token (not used in this test)
const std::vector<llama_token> sequence = {100, 50, 51, 101, 52}; // start, two tokens, end, one more const std::vector<llama_token> 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 5, // budget of 5 tokens
REASONING_BUDGET_IDLE, REASONING_BUDGET_IDLE,
SIZE_MAX, SIZE_MAX); // no forcing expected (natural end) SIZE_MAX, SIZE_MAX); // no forcing expected (natural end)
@@ -306,7 +385,7 @@ int main(void) {
const std::vector<llama_token> forced = {102, 101}; // forced message + end const std::vector<llama_token> forced = {102, 101}; // forced message + end
const std::vector<llama_token> sequence = {100, 50, 51, 52, 53}; // start + 4 tokens (budget=2) const std::vector<llama_token> 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 2, // budget of 2 tokens
REASONING_BUDGET_IDLE, REASONING_BUDGET_IDLE,
3, // forcing starts at i=3 (accept at i=2 depletes budget, apply at i=3 forces) 3, // forcing starts at i=3 (accept at i=2 depletes budget, apply at i=3 forces)
@@ -321,7 +400,7 @@ int main(void) {
const std::vector<llama_token> forced = {102, 101}; const std::vector<llama_token> forced = {102, 101};
const std::vector<llama_token> sequence = {100, 50, 51, 52}; // start token first, then 3 tokens const std::vector<llama_token> 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 0, // budget of 0 tokens
REASONING_BUDGET_COUNTING, // starts counting, promoted to FORCING since budget=0 REASONING_BUDGET_COUNTING, // starts counting, promoted to FORCING since budget=0
0, // forcing starts at i=0 (initialized in FORCING, apply forces immediately) 0, // forcing starts at i=0 (initialized in FORCING, apply forces immediately)
@@ -335,7 +414,7 @@ int main(void) {
const std::vector<llama_token> forced = {102}; const std::vector<llama_token> forced = {102};
const std::vector<llama_token> sequence = {50, 51, 52, 53}; const std::vector<llama_token> 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 2, // budget
REASONING_BUDGET_IDLE, REASONING_BUDGET_IDLE,
SIZE_MAX, SIZE_MAX); // no forcing (no start/end configured) SIZE_MAX, SIZE_MAX); // no forcing (no start/end configured)
@@ -350,7 +429,7 @@ int main(void) {
const std::vector<llama_token> forced = {102, 101}; const std::vector<llama_token> forced = {102, 101};
const std::vector<llama_token> sequence = {50, 51, 52, 53}; const std::vector<llama_token> 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 2, // budget of 2 tokens
REASONING_BUDGET_COUNTING, REASONING_BUDGET_COUNTING,
2, // forcing starts at i=2 (after 2 accepts deplete budget, apply at i=2 forces) 2, // forcing starts at i=2 (after 2 accepts deplete budget, apply at i=2 forces)
@@ -373,18 +452,50 @@ int main(void) {
const std::vector<llama_token> forced = {102, 101}; const std::vector<llama_token> forced = {102, 101};
const std::vector<llama_token> sequence = {100, 50, 101, 100, 60, 61, 62, 63}; const std::vector<llama_token> 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) 2, // budget of 2 tokens (per block)
REASONING_BUDGET_IDLE, REASONING_BUDGET_IDLE,
6, // forcing starts at i=6 (after second block exhausts at i=5) 6, // forcing starts at i=6 (after second block exhausts at i=5)
7); // forcing continues through i=7 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<llama_tokens> start = {{100}, {110, 111}};
const std::vector<llama_tokens> end = {{101}};
const std::vector<llama_token> forced = {102, 101};
const std::vector<llama_token> 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<llama_tokens> start = {{100}};
const std::vector<llama_tokens> end = {{101}, {103, 104}};
const std::vector<llama_token> forced = {102, 101};
const std::vector<llama_token> 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_counting();
test_reasoning_budget_clone_mid_forcing(); test_reasoning_budget_clone_mid_forcing();
test_reasoning_budget_force_manual(); test_reasoning_budget_force_manual();
test_reasoning_budget_end_match();
printf("OK (9 tests passed)\n"); printf("OK (12 tests passed)\n");
printf("Testing UTF-8 boundary detection... "); printf("Testing UTF-8 boundary detection... ");
test_utf8_boundary_detection(); test_utf8_boundary_detection();
tools/server/server-common.cpp
+2 -2
View File
@@ -1131,10 +1131,10 @@ json oaicompat_chat_params_parse(
reasoning_budget = opt.reasoning_budget; 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_tokens"] = reasoning_budget;
llama_params["reasoning_budget_start_tag"] = chat_params.thinking_start_tag; 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_budget_message"] = json_value(body, "reasoning_budget_message", opt.reasoning_budget_message);
llama_params["reasoning_control"] = json_value(body, "reasoning_control", false); llama_params["reasoning_control"] = json_value(body, "reasoning_control", false);
} }
tools/server/server-schema.cpp
+27 -8
View File
@@ -390,21 +390,40 @@ std::vector<std::unique_ptr<field>> make_llama_cmpl_schema(const common_params &
ctx.params.sampling.reasoning_budget_start = common_tokenize(ctx.vocab, data.at("reasoning_budget_start_tag").get<std::string>(), false, true); ctx.params.sampling.reasoning_budget_start = common_tokenize(ctx.vocab, data.at("reasoning_budget_start_tag").get<std::string>(), false, true);
})); }));
add((new field_str("reasoning_budget_end_tag")) add((new field_json("reasoning_budget_end_tags"))
->set_desc("Token string marking the end of the reasoning budget section") ->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) { ->set_handler([&](field_eval_context & ctx, const json & data) {
GGML_ASSERT(ctx.vocab != nullptr); GGML_ASSERT(ctx.vocab != nullptr);
std::string end_tag = data.at("reasoning_budget_end_tag").get<std::string>(); ctx.params.sampling.reasoning_budget_end.clear();
ctx.params.sampling.reasoning_budget_end = common_tokenize(ctx.vocab, end_tag, false, true); if (data.contains("reasoning_budget_end_tags")) {
for (const auto & t : data.at("reasoning_budget_end_tags")) {
std::string tag = t.get<std::string>();
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<std::string>();
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")) add((new field_str("reasoning_budget_message"))
->set_desc("Message to prepend to the reasoning budget end tag when forcing it") ->set_desc("Message to prepend to the reasoning budget end tag when forcing it")
->set_handler([&](field_eval_context & ctx, const json & data) { ->set_handler([&](field_eval_context & ctx, const json & data) {
GGML_ASSERT(ctx.vocab != nullptr); GGML_ASSERT(ctx.vocab != nullptr);
std::string end_tag = json_value(data, "reasoning_budget_end_tag", std::string()); if (!ctx.params.sampling.reasoning_budget_end.empty()) {
std::string message = data.at("reasoning_budget_message").get<std::string>(); llama_tokens end_tag = ctx.params.sampling.reasoning_budget_end.front();
ctx.params.sampling.reasoning_budget_forced = common_tokenize(ctx.vocab, message + end_tag, false, true); 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")) add((new field_json("logit_bias"))
@@ -546,7 +565,7 @@ task_params eval_llama_cmpl_schema(
// debugging // debugging
{ {
auto budget = params.sampling.reasoning_budget_tokens; 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(), budget, params.sampling.generation_prompt.c_str(),
params.sampling.reasoning_budget_start.size(), params.sampling.reasoning_budget_start.size(),
params.sampling.reasoning_budget_end.size(), params.sampling.reasoning_budget_end.size(),