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/reasoning-budget.cpp
+77 -39
View File
@@ -1,39 +1,52 @@
#include "reasoning-budget.h"
#include "common.h"
#include "trie.h"
#include "unicode.h"
#include "log.h"
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <string>
#include <vector>
struct token_matcher {
std::vector<llama_token> tokens;
size_t pos = 0;
std::vector<llama_tokens> seqs;
common_aho_corasick ac;
size_t state = 0;
bool advance(llama_token token) {
if (tokens.empty()) {
return false;
}
token_matcher(const std::vector<llama_tokens> & seqs) : seqs(collect(seqs)), ac(build_trie(this->seqs)) {}
if (token == tokens[pos]) {
pos++;
if (pos >= tokens.size()) {
pos = 0;
return true;
}
} else {
pos = 0;
if (token == tokens[0]) {
pos = 1;
static std::vector<llama_tokens> collect(const std::vector<llama_tokens> & seqs) {
std::vector<llama_tokens> res;
for (const auto & seq : seqs) {
if (!seq.empty() && std::find(res.begin(), res.end(), seq) == res.end()) {
res.push_back(seq);
}
}
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 {
@@ -41,7 +54,7 @@ struct common_reasoning_budget_ctx {
token_matcher start_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 remaining; // tokens remaining in budget
@@ -50,6 +63,8 @@ struct common_reasoning_budget_ctx {
// for forcing
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*/) {
@@ -62,7 +77,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);
@@ -78,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)) {
const int32_t match = ctx->end_matcher.advance(token);
if (match >= 0) {
ctx->state = REASONING_BUDGET_DONE;
ctx->end_match = match;
COM_TRC("%s", "deactivated (natural end)\n");
break;
}
@@ -115,19 +132,25 @@ 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->end_match = 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 <think> 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();
ctx->end_match = -1;
COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget);
if (ctx->remaining <= 0) {
@@ -169,11 +192,12 @@ static void common_reasoning_budget_reset(struct llama_sampler * smpl) {
ctx->start_matcher.reset();
ctx->end_matcher.reset();
ctx->force_pos = 0;
ctx->end_match = -1;
}
static struct llama_sampler * common_reasoning_budget_init_state(
const struct llama_vocab * vocab, const std::vector<llama_token> & start_tokens,
const std::vector<llama_token> & end_tokens, const std::vector<llama_token> & forced_tokens,
const struct llama_vocab * vocab, const std::vector<llama_tokens> & start_seqs,
const std::vector<llama_tokens> & 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 +229,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<llama_token> & start_tokens,
const std::vector<llama_token> & end_tokens,
const std::vector<llama_token> & forced_tokens,
int32_t budget,
common_reasoning_budget_state initial_state) {
const struct llama_vocab * vocab,
const std::vector<llama_tokens> & start_seqs,
const std::vector<llama_tokens> & 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,25 +244,26 @@ 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,
/* .state = */ initial_state,
/* .force_pos = */ 0,
/* .end_match = */ -1,
}
);
}
struct llama_sampler * common_reasoning_budget_init(
const struct llama_vocab * vocab,
const std::vector<llama_token> & start_tokens,
const std::vector<llama_token> & end_tokens,
const std::vector<llama_token> & 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<llama_tokens> & start_seqs,
const std::vector<llama_tokens> & 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) {
@@ -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;
}
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) {
if (!smpl) {
return false;