/src/llama.cpp/common/reasoning-budget.cpp

Source
#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_tokens> seqs;
    common_aho_corasick ac;
    size_t state = 0;

    token_matcher(const std::vector<llama_tokens> & seqs) : seqs(collect(seqs)), ac(build_trie(this->seqs)) {}

    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 res;
    }

    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 {
    const llama_vocab * vocab;

    token_matcher start_matcher;
    token_matcher end_matcher;
    llama_tokens forced_tokens;

    int32_t budget;           // maximum tokens in reasoning block
    int32_t remaining;        // tokens remaining in budget

    common_reasoning_budget_state state;

    // 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*/) {
    return "reasoning-budget";
}

static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_token token) {
    auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx;

    switch (ctx->state) {
        case REASONING_BUDGET_IDLE:
        {
            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);

                if (ctx->remaining <= 0) {
                    ctx->state = REASONING_BUDGET_FORCING;
                    ctx->force_pos = 0;
                    COM_TRC("%s", "budget=0, forcing immediately\n");
                }
            }
            break;
        }
        case REASONING_BUDGET_COUNTING:
        case REASONING_BUDGET_WAITING_UTF8:
        {
            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;
            }

            bool utf8_complete = true;
            if (ctx->vocab != nullptr) {
                const std::string piece = common_token_to_piece(ctx->vocab, token, false);
                utf8_complete = common_utf8_is_complete(piece);
            }

            if (ctx->state == REASONING_BUDGET_WAITING_UTF8) {
                if (utf8_complete) {
                    ctx->state = REASONING_BUDGET_FORCING;
                    ctx->force_pos = 0;
                    ctx->end_matcher.reset();
                    COM_TRC("%s", "UTF-8 complete, now forcing end sequence\n");
                }
            } else if (ctx->state == REASONING_BUDGET_COUNTING) {
                ctx->remaining--;
                if (ctx->remaining <= 0) {
                    if (utf8_complete) {
                        ctx->state = REASONING_BUDGET_FORCING;
                        ctx->force_pos = 0;
                        ctx->end_matcher.reset();
                        COM_TRC("%s", "budget exhausted, forcing end sequence\n");
                    } else {
                        ctx->state = REASONING_BUDGET_WAITING_UTF8;
                        ctx->end_matcher.reset();
                        COM_TRC("%s", "budget exhausted, waiting for UTF-8 completion\n");
                    }
                }
            }
            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) >= 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) {
                    ctx->state = REASONING_BUDGET_FORCING;
                    ctx->force_pos = 0;
                    COM_TRC("%s", "budget=0, forcing immediately\n");
                }
            }
            break;
    }
}

static void common_reasoning_budget_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
    auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx;

    if (ctx->state != REASONING_BUDGET_FORCING) {
        // passthrough — don't modify logits
        return;
    }

    if (ctx->force_pos >= ctx->forced_tokens.size()) {
        return;
    }

    const llama_token forced = ctx->forced_tokens[ctx->force_pos];

    // set all logits to -inf except the forced token
    for (size_t i = 0; i < cur_p->size; i++) {
        if (cur_p->data[i].id != forced) {
            cur_p->data[i].logit = -INFINITY;
        }
    }
}

static void common_reasoning_budget_reset(struct llama_sampler * smpl) {
    auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx;
    ctx->state = REASONING_BUDGET_IDLE;
    ctx->remaining = ctx->budget;
    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_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);

static void common_reasoning_budget_free(struct llama_sampler * smpl) {
    delete (common_reasoning_budget_ctx *) smpl->ctx;
}

static struct llama_sampler_i common_reasoning_budget_i = {
    /* .name              = */ common_reasoning_budget_name,
    /* .accept            = */ common_reasoning_budget_accept,
    /* .apply             = */ common_reasoning_budget_apply,
    /* .reset             = */ common_reasoning_budget_reset,
    /* .clone             = */ common_reasoning_budget_clone,
    /* .free              = */ common_reasoning_budget_free,
    /* .backend_init      = */ nullptr,
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ nullptr,
    /* .backend_set_input = */ nullptr,
    /* .backend_reset     = */ nullptr,
    /* .copy_state        = */ nullptr,
};

static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl) {
    const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx;

    return llama_sampler_init(
        /* .iface = */ &common_reasoning_budget_i,
        /* .ctx   = */ new common_reasoning_budget_ctx(*ctx)
    );
}

static struct llama_sampler * common_reasoning_budget_init_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;
    }

    return llama_sampler_init(
        /* .iface = */ &common_reasoning_budget_i,
        /* .ctx   = */ new common_reasoning_budget_ctx {
            /* .vocab         = */ vocab,
            /* .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_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) {
    if (!smpl) {
        return REASONING_BUDGET_IDLE;
    }
    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;
    }

    auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx;

    // only a sampler that is actively counting down the budget may be forced;
    // any other state (idle, already forcing/waiting, or done) is left untouched
    if (ctx->state != REASONING_BUDGET_COUNTING) {
        return false;
    }

    ctx->state = REASONING_BUDGET_FORCING;
    ctx->force_pos = 0;
    ctx->end_matcher.reset();
    COM_TRC("%s", "forced into forcing state (manual transition)\n");

    return true;
}

Coverage Report

Created: 2026-08-22 07:18

Line	Count	Source
1		#include "reasoning-budget.h"
2		#include "common.h"
3		#include "trie.h"
4		#include "unicode.h"
5
6		#include "log.h"
7
8		#include <algorithm>
9		#include <cmath>
10		#include <cstdint>
11		#include <string>
12		#include <vector>
13
14		struct token_matcher {
15		std::vector<llama_tokens> seqs;
16		common_aho_corasick ac;
17		size_t state = 0;
18
19	0	token_matcher(const std::vector<llama_tokens> & seqs) : seqs(collect(seqs)), ac(build_trie(this->seqs)) {}
20
21	0	static std::vector<llama_tokens> collect(const std::vector<llama_tokens> & seqs) {
22	0	std::vector<llama_tokens> res;
23	0	for (const auto & seq : seqs) {
24	0	if (!seq.empty() && std::find(res.begin(), res.end(), seq) == res.end()) {
25	0	res.push_back(seq);
26	0	}
27	0	}
28	0	return res;
29	0	}
30
31	0	static common_trie build_trie(const std::vector<llama_tokens> & seqs) {
32	0	common_trie t;
33	0	for (const auto & seq : seqs) {
34	0	t.insert(std::vector<uint32_t>(seq.begin(), seq.end()));
35	0	}
36	0	return t;
37	0	}
38
39		// returns the index into seqs of the longest sequence ending at this token, or -1
40	0	int32_t advance(llama_token token) {
41	0	state = ac.next(state, (uint32_t) token);
42	0	const int32_t p = ac.match_pattern(state);
43	0	if (p >= 0) {
44	0	state = 0;
45	0	}
46	0	return p;
47	0	}
48
49	0	void reset() { state = 0; }
50		};
51
52		struct common_reasoning_budget_ctx {
53		const llama_vocab * vocab;
54
55		token_matcher start_matcher;
56		token_matcher end_matcher;
57		llama_tokens forced_tokens;
58
59		int32_t budget; // maximum tokens in reasoning block
60		int32_t remaining; // tokens remaining in budget
61
62		common_reasoning_budget_state state;
63
64		// for forcing
65		size_t force_pos; // next position in forced_tokens to force
66
67		int32_t end_match; // index into end_matcher.seqs of the sequence that transitioned to DONE, -1 if none
68		};
69
70	0	static const char * common_reasoning_budget_name(const struct llama_sampler * /smpl/) {
71	0	return "reasoning-budget";
72	0	}
73
74	0	static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_token token) {
75	0	auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx;
76
77	0	switch (ctx->state) {
78	0	case REASONING_BUDGET_IDLE:
79	0	{
80	0	if (ctx->start_matcher.advance(token) >= 0) {
81	0	ctx->state = REASONING_BUDGET_COUNTING;
82	0	ctx->remaining = ctx->budget;
83	0	COM_TRC("activated, budget=%d tokens\n", ctx->budget);
84
85	0	if (ctx->remaining <= 0) {
86	0	ctx->state = REASONING_BUDGET_FORCING;
87	0	ctx->force_pos = 0;
88	0	COM_TRC("%s", "budget=0, forcing immediately\n");
89	0	}
90	0	}
91	0	break;
92	0	}
93	0	case REASONING_BUDGET_COUNTING:
94	0	case REASONING_BUDGET_WAITING_UTF8:
95	0	{
96	0	const int32_t match = ctx->end_matcher.advance(token);
97	0	if (match >= 0) {
98	0	ctx->state = REASONING_BUDGET_DONE;
99	0	ctx->end_match = match;
100	0	COM_TRC("%s", "deactivated (natural end)\n");
101	0	break;
102	0	}
103
104	0	bool utf8_complete = true;
105	0	if (ctx->vocab != nullptr) {
106	0	const std::string piece = common_token_to_piece(ctx->vocab, token, false);
107	0	utf8_complete = common_utf8_is_complete(piece);
108	0	}
109
110	0	if (ctx->state == REASONING_BUDGET_WAITING_UTF8) {
111	0	if (utf8_complete) {
112	0	ctx->state = REASONING_BUDGET_FORCING;
113	0	ctx->force_pos = 0;
114	0	ctx->end_matcher.reset();
115	0	COM_TRC("%s", "UTF-8 complete, now forcing end sequence\n");
116	0	}
117	0	} else if (ctx->state == REASONING_BUDGET_COUNTING) {
118	0	ctx->remaining--;
119	0	if (ctx->remaining <= 0) {
120	0	if (utf8_complete) {
121	0	ctx->state = REASONING_BUDGET_FORCING;
122	0	ctx->force_pos = 0;
123	0	ctx->end_matcher.reset();
124	0	COM_TRC("%s", "budget exhausted, forcing end sequence\n");
125	0	} else {
126	0	ctx->state = REASONING_BUDGET_WAITING_UTF8;
127	0	ctx->end_matcher.reset();
128	0	COM_TRC("%s", "budget exhausted, waiting for UTF-8 completion\n");
129	0	}
130	0	}
131	0	}
132	0	break;
133	0	}
134	0	case REASONING_BUDGET_FORCING:
135	0	{
136		// track the end sequence within forced_tokens so it is also reported on DONE
137	0	const int32_t match = ctx->end_matcher.advance(token);
138	0	ctx->force_pos++;
139	0	if (ctx->force_pos >= ctx->forced_tokens.size()) {
140	0	ctx->state = REASONING_BUDGET_DONE;
141	0	ctx->end_match = match;
142	0	COM_TRC("%s", "forced sequence complete, done\n");
143	0	}
144	0	break;
145	0	}
146	0	case REASONING_BUDGET_DONE:
147		// Re-arm on a new start tag: some models emit multiple <think> blocks
148		// per response, and each should get a fresh budget window.
149	0	if (ctx->start_matcher.advance(token) >= 0) {
150	0	ctx->state = REASONING_BUDGET_COUNTING;
151	0	ctx->remaining = ctx->budget;
152	0	ctx->end_matcher.reset();
153	0	ctx->end_match = -1;
154	0	COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget);
155
156	0	if (ctx->remaining <= 0) {
157	0	ctx->state = REASONING_BUDGET_FORCING;
158	0	ctx->force_pos = 0;
159	0	COM_TRC("%s", "budget=0, forcing immediately\n");
160	0	}
161	0	}
162	0	break;
163	0	}
164	0	}
165
166	0	static void common_reasoning_budget_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
167	0	auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx;
168
169	0	if (ctx->state != REASONING_BUDGET_FORCING) {
170		// passthrough — don't modify logits
171	0	return;
172	0	}
173
174	0	if (ctx->force_pos >= ctx->forced_tokens.size()) {
175	0	return;
176	0	}
177
178	0	const llama_token forced = ctx->forced_tokens[ctx->force_pos];
179
180		// set all logits to -inf except the forced token
181	0	for (size_t i = 0; i < cur_p->size; i++) {
182	0	if (cur_p->data[i].id != forced) {
183	0	cur_p->data[i].logit = -INFINITY;
184	0	}
185	0	}
186	0	}
187
188	0	static void common_reasoning_budget_reset(struct llama_sampler * smpl) {
189	0	auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx;
190	0	ctx->state = REASONING_BUDGET_IDLE;
191	0	ctx->remaining = ctx->budget;
192	0	ctx->start_matcher.reset();
193	0	ctx->end_matcher.reset();
194	0	ctx->force_pos = 0;
195	0	ctx->end_match = -1;
196	0	}
197
198		static struct llama_sampler * common_reasoning_budget_init_state(
199		const struct llama_vocab * vocab, const std::vector<llama_tokens> & start_seqs,
200		const std::vector<llama_tokens> & end_seqs, const llama_tokens & forced_tokens,
201		int32_t budget, common_reasoning_budget_state initial_state);
202
203		static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl);
204
205	0	static void common_reasoning_budget_free(struct llama_sampler * smpl) {
206	0	delete (common_reasoning_budget_ctx *) smpl->ctx;
207	0	}
208
209		static struct llama_sampler_i common_reasoning_budget_i = {
210		/* .name = */ common_reasoning_budget_name,
211		/* .accept = */ common_reasoning_budget_accept,
212		/* .apply = */ common_reasoning_budget_apply,
213		/* .reset = */ common_reasoning_budget_reset,
214		/* .clone = */ common_reasoning_budget_clone,
215		/* .free = */ common_reasoning_budget_free,
216		/* .backend_init = */ nullptr,
217		/* .backend_accept = */ nullptr,
218		/* .backend_apply = */ nullptr,
219		/* .backend_set_input = */ nullptr,
220		/* .backend_reset = */ nullptr,
221		/* .copy_state = */ nullptr,
222		};
223
224	0	static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl) {
225	0	const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx;
226
227	0	return llama_sampler_init(
228	0	/* .iface = */ &common_reasoning_budget_i,
229	0	/* .ctx = / new common_reasoning_budget_ctx(ctx)
230	0	);
231	0	}
232
233		static struct llama_sampler * common_reasoning_budget_init_state(
234		const struct llama_vocab * vocab,
235		const std::vector<llama_tokens> & start_seqs,
236		const std::vector<llama_tokens> & end_seqs,
237		const llama_tokens & forced_tokens,
238		int32_t budget,
239	0	common_reasoning_budget_state initial_state) {
240		// promote COUNTING with budget <= 0 to FORCING
241	0	if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) {
242	0	initial_state = REASONING_BUDGET_FORCING;
243	0	}
244
245	0	return llama_sampler_init(
246	0	/* .iface = */ &common_reasoning_budget_i,
247	0	/* .ctx = */ new common_reasoning_budget_ctx {
248	0	/* .vocab = */ vocab,
249	0	/* .start_matcher = */ token_matcher(start_seqs),
250	0	/* .end_matcher = */ token_matcher(end_seqs),
251	0	/* .forced_tokens = */ forced_tokens,
252	0	/* .budget = */ budget,
253	0	/* .remaining = */ budget,
254	0	/* .state = */ initial_state,
255	0	/* .force_pos = */ 0,
256	0	/* .end_match = */ -1,
257	0	}
258	0	);
259	0	}
260
261		struct llama_sampler * common_reasoning_budget_init(
262		const struct llama_vocab * vocab,
263		const std::vector<llama_tokens> & start_seqs,
264		const std::vector<llama_tokens> & end_seqs,
265		const llama_tokens & forced_tokens,
266		int32_t budget,
267	0	common_reasoning_budget_state initial_state) {
268	0	return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, budget, initial_state);
269	0	}
270
271	0	common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) {
272	0	if (!smpl) {
273	0	return REASONING_BUDGET_IDLE;
274	0	}
275	0	return ((const common_reasoning_budget_ctx *)smpl->ctx)->state;
276	0	}
277
278	0	const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl) {
279	0	if (!smpl) {
280	0	return nullptr;
281	0	}
282
283	0	const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx;
284	0	if (ctx->end_match < 0) {
285	0	return nullptr;
286	0	}
287
288	0	return &ctx->end_matcher.seqs[ctx->end_match];
289	0	}
290
291	0	bool common_reasoning_budget_force(struct llama_sampler * smpl) {
292	0	if (!smpl) {
293	0	return false;
294	0	}
295
296	0	auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx;
297
298		// only a sampler that is actively counting down the budget may be forced;
299		// any other state (idle, already forcing/waiting, or done) is left untouched
300	0	if (ctx->state != REASONING_BUDGET_COUNTING) {
301	0	return false;
302	0	}
303
304	0	ctx->state = REASONING_BUDGET_FORCING;
305	0	ctx->force_pos = 0;
306	0	ctx->end_matcher.reset();
307	0	COM_TRC("%s", "forced into forcing state (manual transition)\n");
308
309	0	return true;
310	0	}