Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pwilkin merged 8 commits into from
Jul 25, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
22c51b6
common : extract trie/ac to a separate file
aldehir Jul 8, 2026
c2321fa
common : support multiple token sequences in the reasoning budget sam…
aldehir Jul 8, 2026
ea9aa31
common/trie : return matched word index
aldehir Jul 9, 2026
186241a
common/trie : rename "word" to "pattern"
aldehir Jul 9, 2026
d44a6c8
common/reasoning-budget : expose matched end sequence
aldehir Jul 9, 2026
e77709a
common/sampling : replay end sequence when reasoning budget is done
aldehir Jul 9, 2026
fc1c28b
cont : update to use multiple end sequences
aldehir Jul 10, 2026
72f4ebd
cont : clean up
aldehir Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions common/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,8 @@ add_library(${TARGET}
sampling.h
speculative.cpp
speculative.h
trie.cpp
trie.h
unicode.cpp
unicode.h
jinja/lexer.cpp
Expand Down
22 changes: 14 additions & 8 deletions common/chat.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 = {
Expand Down Expand Up @@ -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 = "<channel|>";
data.thinking_end_tags = {"<channel|>"};

data.preserved_tokens = {
"<|channel>",
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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 = "<think>";
data.thinking_end_tag = "</think>";
data.thinking_end_tags = {"</think>"};
data.preserved_tokens = {
"｜DSML｜",
"<think>",
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -2418,7 +2421,7 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem
};

data.thinking_start_tag = "<think>";
data.thinking_end_tag = "</think>";
data.thinking_end_tags = {"</think>"};

data.message_delimiters = {
{ COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant" },
Expand Down Expand Up @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion common/chat.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -274,7 +274,7 @@ struct common_chat_params {
std::string generation_prompt;
bool supports_thinking = false;
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<std::string> preserved_tokens;
std::vector<std::string> additional_stops;
Expand Down
12 changes: 6 additions & 6 deletions common/common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<llama_token> reasoning_budget_start; // start tag token sequence
std::vector<llama_token> reasoning_budget_end; // end tag token sequence
std::vector<llama_token> 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<llama_token> reasoning_budget_start; // start 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 + 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;

Expand Down
158 changes: 5 additions & 153 deletions common/peg-parser.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,10 @@
#include "common.h"
#include "json-schema-to-grammar.h"
#include "log.h"
#include "trie.h"
#include "unicode.h"

#include <algorithm>
#include <deque>
#include <initializer_list>
#include <map>
#include <memory>
Expand All @@ -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<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) {
if (pos + hex_count > str.length()) {
return {0, 0};
Expand Down Expand Up @@ -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;
Expand All @@ -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);
}
Expand Down Expand Up @@ -1559,7 +1411,7 @@ static std::string gbnf_ac_grammar(
const std::map<size_t, std::vector<uint32_t>> &,
const std::vector<uint32_t> &,
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 {
if (s == 0) {
Expand Down
Loading
Loading