Skip to content

reasoning-budget: do not re-fire the intro message in once mode #8

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
dzannotti merged 1 commit into from
Aug 29, 2026
Merged

reasoning-budget: do not re-fire the intro message in once mode #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
18 changes: 14 additions & 4 deletions common/reasoning-budget.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ struct common_reasoning_budget_ctx {

llama_tokens intro_forced_tokens;
size_t intro_force_pos;
bool intro_once;

int32_t grace_tokens;
int32_t grace_remaining;
Expand Down Expand Up @@ -154,6 +155,11 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
case REASONING_BUDGET_INTRO_FORCING:
ctx->intro_force_pos++;
if (ctx->intro_force_pos >= ctx->intro_forced_tokens.size()) {
if (ctx->intro_once) {
// intro-mode once: do not re-fire on later reasoning blocks
ctx->intro_forced_tokens.clear();
ctx->intro_force_pos = 0;
}
if (ctx->remaining <= 0) {
ctx->state = REASONING_BUDGET_FORCING;
ctx->force_pos = 0;
Expand Down Expand Up @@ -354,7 +360,8 @@ static struct llama_sampler * common_reasoning_budget_init_state(
const llama_tokens & intro_forced_tokens,
int32_t budget,
int32_t grace_tokens,
common_reasoning_budget_state initial_state);
common_reasoning_budget_state initial_state,
bool intro_once);

static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl);

Expand Down Expand Up @@ -395,7 +402,8 @@ static struct llama_sampler * common_reasoning_budget_init_state(
const llama_tokens & intro_forced_tokens,
int32_t budget,
int32_t grace_tokens,
common_reasoning_budget_state initial_state) {
common_reasoning_budget_state initial_state,
bool intro_once) {
if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) {
initial_state = REASONING_BUDGET_FORCING;
}
Expand Down Expand Up @@ -426,6 +434,7 @@ static struct llama_sampler * common_reasoning_budget_init_state(
/* .active_soft = */ 0,
/* .intro_forced_tokens = */ intro_forced_tokens,
/* .intro_force_pos = */ 0,
/* .intro_once = */ intro_once,
/* .grace_tokens = */ grace_tokens,
/* .grace_remaining = */ grace_tokens,
/* .hard_pending_prev_nl = */ false,
Expand All @@ -441,9 +450,10 @@ struct llama_sampler * common_reasoning_budget_init(
const llama_tokens & intro_forced_tokens,
int32_t budget,
int32_t grace_tokens,
common_reasoning_budget_state initial_state) {
common_reasoning_budget_state initial_state,
bool intro_once) {
return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, soft_points,
intro_forced_tokens, budget, grace_tokens, initial_state);
intro_forced_tokens, budget, grace_tokens, initial_state, intro_once);
}

common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) {
Expand Down
3 changes: 2 additions & 1 deletion common/reasoning-budget.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,8 @@ struct llama_sampler * common_reasoning_budget_init(
const llama_tokens & intro_forced_tokens,
int32_t budget,
int32_t grace_tokens = 0,
common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE);
common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE,
bool intro_once = false);

common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl);

Expand Down
3 changes: 2 additions & 1 deletion common/sampling.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -329,7 +329,8 @@ struct common_sampler * common_sampler_init(
reasoning_budget_active ? params.reasoning_budget_intro_forced : llama_tokens();
rbudget = common_reasoning_budget_init(vocab, { params.reasoning_budget_start }, params.reasoning_budget_end,
params.reasoning_budget_forced, soft_points, intro_tokens, budget_tokens,
params.reasoning_budget_grace_tokens);
params.reasoning_budget_grace_tokens, REASONING_BUDGET_IDLE,
params.reasoning_budget_intro_mode == "once");

for (const auto & token : prefill_tokens) {
llama_sampler_accept(rbudget, token);
Expand Down
24 changes: 24 additions & 0 deletions tests/test-reasoning-budget.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -438,6 +438,29 @@ static void test_reasoning_budget_intro_rearms_on_multiblock() {
llama_sampler_free(sampler);
}

static void test_reasoning_budget_intro_once_does_not_rearm() {
const std::vector<llama_token> start = { 100 };
const std::vector<llama_token> end = { 101 };
const std::vector<llama_token> forced = { 102, 101 };
const std::vector<llama_token> intro_forced = { 300, 301 };

auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, intro_forced, 5, 0, REASONING_BUDGET_IDLE, true);

llama_sampler_accept(sampler, 100);
GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING);
llama_sampler_accept(sampler, 300);
llama_sampler_accept(sampler, 301);
GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING);
llama_sampler_accept(sampler, 101);
GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE);

// second reasoning block: intro must not fire again in once mode
llama_sampler_accept(sampler, 100);
GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING);

llama_sampler_free(sampler);
}

static void test_reasoning_budget_hard_pending_grace_expires() {
const std::vector<llama_token> start = { 100 };
const std::vector<llama_token> end = { 101 };
Expand Down Expand Up @@ -769,6 +792,7 @@ int main(void) {
test_reasoning_budget_intro_forcing_budget_zero();
test_reasoning_budget_force_manual_from_intro();
test_reasoning_budget_intro_rearms_on_multiblock();
test_reasoning_budget_intro_once_does_not_rearm();
test_reasoning_budget_hard_pending_grace_expires();
test_reasoning_budget_hard_pending_natural_end();
test_reasoning_budget_force_manual_from_hard_pending();
Expand Down