diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index 8c65908b0fa2..538839a468b2 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -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; @@ -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; @@ -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); @@ -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; } @@ -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, @@ -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) { diff --git a/common/reasoning-budget.h b/common/reasoning-budget.h index 4151777e82c4..e3ababe0a8d6 100644 --- a/common/reasoning-budget.h +++ b/common/reasoning-budget.h @@ -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); diff --git a/common/sampling.cpp b/common/sampling.cpp index 8141dfa0f79f..0e40d86e893c 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -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); diff --git a/tests/test-reasoning-budget.cpp b/tests/test-reasoning-budget.cpp index e4c6d3c44e3e..45b7c3da35fe 100644 --- a/tests/test-reasoning-budget.cpp +++ b/tests/test-reasoning-budget.cpp @@ -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 start = { 100 }; + const std::vector end = { 101 }; + const std::vector forced = { 102, 101 }; + const std::vector 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 start = { 100 }; const std::vector end = { 101 }; @@ -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();