Skip to content

common/sampling: reset reasoning budget sampler state between generations #21594

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

Closed
cnsiva wants to merge 4 commits into from
Closed

common/sampling: reset reasoning budget sampler state between generations #21594

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c3b7e97
common/sampling: reset reasoning budget sampler state between generat…
cnsiva Apr 8, 2026
eb2f526
common/sampling: reset grammar sampler state between generations
cnsiva Apr 8, 2026
7807e30
common : fix sampler reset regression and prevent state leakage
cnsiva Apr 9, 2026
90907de
common : refactor reasoning-budget clone and add reset unit test
cnsiva Apr 14, 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
29 changes: 12 additions & 17 deletions common/reasoning-budget.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ struct common_reasoning_budget_ctx {
int32_t budget; // maximum tokens in reasoning block
int32_t remaining; // tokens remaining in budget

common_reasoning_budget_state initial_state; // state to restore on reset
common_reasoning_budget_state state;

// for forcing
Expand Down Expand Up @@ -150,29 +151,15 @@ static void common_reasoning_budget_apply(struct llama_sampler * smpl, llama_tok

static void common_reasoning_budget_reset(struct llama_sampler * smpl) {
auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx;
ctx->state = REASONING_BUDGET_IDLE;
ctx->state = ctx->initial_state;
ctx->remaining = ctx->budget;
ctx->start_matcher.reset();
ctx->end_matcher.reset();
ctx->force_pos = 0;
}

// forward declaration for use in clone
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);

static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl) {
const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx;
return common_reasoning_budget_init_state(
ctx->vocab,
ctx->start_matcher.tokens,
ctx->end_matcher.tokens,
ctx->forced_tokens,
ctx->budget,
ctx->state);
}
// forward declarations for use in vtable
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;
Expand All @@ -191,6 +178,13 @@ static struct llama_sampler_i common_reasoning_budget_i = {
/* .backend_set_input = */ 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(
&common_reasoning_budget_i,
new common_reasoning_budget_ctx(*ctx));
}

static struct llama_sampler * common_reasoning_budget_init_state(
const struct llama_vocab * vocab,
const std::vector<llama_token> & start_tokens,
Expand All @@ -212,6 +206,7 @@ static struct llama_sampler * common_reasoning_budget_init_state(
/* .forced_tokens = */ forced_tokens,
/* .budget = */ budget,
/* .remaining = */ budget,
/* .initial_state = */ initial_state,
/* .state = */ initial_state,
/* .force_pos = */ 0,
}
Expand Down
43 changes: 26 additions & 17 deletions common/sampling.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,13 +115,22 @@ struct common_sampler {

ring_buffer<llama_token> prev;

std::vector<llama_token> prefill_tokens;

std::vector<llama_token_data> cur;

llama_token_data_array cur_p;

void reset() {
prev.clear();

llama_sampler_reset(grmr);
if (grmr) {
for (const auto & token : prefill_tokens) {
llama_sampler_accept(grmr, token);
}
}
llama_sampler_reset(rbudget);
llama_sampler_reset(chain);
}

Expand Down Expand Up @@ -389,17 +398,16 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, st
params.backend_sampling = false;
}

auto * result = new common_sampler {
/* .params = */ params,
/* .grmr = */ grmr,
/* .rbudget = */ rbudget,
/* .chain = */ chain,
/* .prev = */ ring_buffer<llama_token>(std::max(32, params.n_prev)),
/* .cur = */ {},
/* .cur_p = */ {},
return new common_sampler {
/* .params = */ params,
/* .grmr = */ grmr,
/* .rbudget = */ rbudget,
/* .chain = */ chain,
/* .prev = */ ring_buffer<llama_token>(std::max(32, params.n_prev)),
/* .prefill_tokens = */ std::move(prefill_tokens),
/* .cur = */ {},
/* .cur_p = */ {},
};

return result;
}

void common_sampler_free(struct common_sampler * gsmpl) {
Expand Down Expand Up @@ -460,13 +468,14 @@ void common_sampler_reset(struct common_sampler * gsmpl) {

struct common_sampler * common_sampler_clone(common_sampler * gsmpl) {
return new common_sampler {
/* .params = */ gsmpl->params,
/* .grmr = */ llama_sampler_clone(gsmpl->grmr),
/* .rbudget = */ llama_sampler_clone(gsmpl->rbudget),
/* .chain = */ llama_sampler_clone(gsmpl->chain),
/* .prev = */ gsmpl->prev,
/* .cur = */ gsmpl->cur,
/* .cur_p = */ gsmpl->cur_p,
/* .params = */ gsmpl->params,
/* .grmr = */ llama_sampler_clone(gsmpl->grmr),
/* .rbudget = */ llama_sampler_clone(gsmpl->rbudget),
/* .chain = */ llama_sampler_clone(gsmpl->chain),
/* .prev = */ gsmpl->prev,
/* .prefill_tokens = */ gsmpl->prefill_tokens,
/* .cur = */ gsmpl->cur,
/* .cur_p = */ gsmpl->cur_p,
};
}

Expand Down
Loading