common : refactor common_sampler + grammar logic changes (#17937)

* common : refactor common_sampler + grammar logic changes

* tests : increase max_tokens to get needed response

* batched : fix uninitialized samplers
This commit is contained in:
Georgi Gerganov
2025-12-14 10:11:13 +02:00
committed by GitHub
parent 3238b1400c
commit 254098a279
27 changed files with 372 additions and 293 deletions
+7 -13
View File
@@ -141,13 +141,15 @@ int main(int argc, char ** argv) {
// load the model and apply lora adapter, if any
LOG_INF("%s: load the model and apply lora adapter, if any\n", __func__);
common_init_result llama_init = common_init_from_params(params);
model = llama_init.model.get();
ctx = llama_init.context.get();
auto llama_init = common_init_from_params(params);
if (model == NULL) {
LOG_ERR("%s: error: unable to load model\n", __func__);
ctx = llama_init->context();
model = llama_init->model();
smpl = llama_init->sampler(0);
if (ctx == NULL) {
LOG_ERR("%s: error: unable to create context\n", __func__);
return 1;
}
@@ -474,12 +476,6 @@ int main(int argc, char ** argv) {
}
}
smpl = common_sampler_init(model, sparams);
if (!smpl) {
LOG_ERR("%s: failed to initialize sampling subsystem\n", __func__);
return 1;
}
LOG_INF("sampler seed: %u\n", common_sampler_get_seed(smpl));
LOG_INF("sampler params: \n%s\n", sparams.print().c_str());
LOG_INF("sampler chain: %s\n", common_sampler_print(smpl).c_str());
@@ -993,8 +989,6 @@ int main(int argc, char ** argv) {
LOG("\n\n");
common_perf_print(ctx, smpl);
common_sampler_free(smpl);
llama_backend_free();
ggml_threadpool_free_fn(threadpool);
@@ -419,10 +419,10 @@ int main(int argc, char ** argv) {
llama_numa_init(params.numa);
// load the model to get hparams
common_init_result llama_init = common_init_from_params(params);
auto llama_init = common_init_from_params(params);
llama_model * model = llama_init.model.get();
llama_context * ctx = llama_init.context.get();
auto * model = llama_init->model();
auto * ctx = llama_init->context();
// int n_ctx = llama_n_ctx(ctx);
int n_layers = llama_model_n_layer(model);
+3 -3
View File
@@ -1265,10 +1265,10 @@ int main(int argc, char ** argv) {
params.warmup = false;
// init
common_init_result llama_init = common_init_from_params(params);
auto llama_init = common_init_from_params(params);
llama_model * model = llama_init.model.get();
llama_context * ctx = llama_init.context.get();
auto * model = llama_init->model();
auto * ctx = llama_init->context();
if (model == nullptr || ctx == nullptr) {
LOG_ERR("%s : failed to init\n", __func__);
+3 -3
View File
@@ -65,7 +65,7 @@ static void sigint_handler(int signo) {
struct mtmd_cli_context {
mtmd::context_ptr ctx_vision;
common_init_result llama_init;
common_init_result_ptr llama_init;
llama_model * model;
llama_context * lctx;
@@ -89,8 +89,8 @@ struct mtmd_cli_context {
llama_pos n_past = 0;
mtmd_cli_context(common_params & params) : llama_init(common_init_from_params(params)) {
model = llama_init.model.get();
lctx = llama_init.context.get();
model = llama_init->model();
lctx = llama_init->context();
vocab = llama_model_get_vocab(model);
smpl = common_sampler_init(model, params.sampling);
n_threads = params.cpuparams.n_threads;
+3 -3
View File
@@ -2024,10 +2024,10 @@ int main(int argc, char ** argv) {
llama_numa_init(params.numa);
// load the model and apply lora adapter, if any
common_init_result llama_init = common_init_from_params(params);
auto llama_init = common_init_from_params(params);
llama_model * model = llama_init.model.get();
llama_context * ctx = llama_init.context.get();
auto * model = llama_init->model();
auto * ctx = llama_init->context();
if (model == NULL) {
LOG_ERR("%s: unable to load model\n", __func__);
+20 -27
View File
@@ -153,7 +153,7 @@ struct server_slot {
// sampling
json json_schema;
struct common_sampler * smpl = nullptr;
common_sampler_ptr smpl;
llama_token sampled; // in speculative mode, this is the last accepted token
llama_tokens drafted;
@@ -510,8 +510,8 @@ struct server_context_impl {
common_params params_base;
// note: keep these alive - they determine the lifetime of the model, context, etc.
common_init_result llama_init;
common_init_result llama_init_dft;
common_init_result_ptr llama_init;
common_init_result_ptr llama_init_dft;
llama_model * model = nullptr;
llama_context * ctx = nullptr;
@@ -557,9 +557,6 @@ struct server_context_impl {
// Clear any sampling context
for (server_slot & slot : slots) {
common_sampler_free(slot.smpl);
slot.smpl = nullptr;
llama_free(slot.ctx_dft);
slot.ctx_dft = nullptr;
@@ -580,8 +577,8 @@ struct server_context_impl {
llama_init = common_init_from_params(params_base);
model = llama_init.model.get();
ctx = llama_init.context.get();
model = llama_init->model();
ctx = llama_init->context();
if (model == nullptr) {
SRV_ERR("failed to load model, '%s'\n", params_base.model.path.c_str());
@@ -613,25 +610,25 @@ struct server_context_impl {
llama_init_dft = common_init_from_params(params_dft);
model_dft = llama_init_dft.model.get();
model_dft = llama_init_dft->model();
if (model_dft == nullptr) {
SRV_ERR("failed to load draft model, '%s'\n", params_base.speculative.model.path.c_str());
return false;
}
vocab_dft_compatible = common_speculative_are_compatible(ctx, llama_init_dft.context.get());
vocab_dft_compatible = common_speculative_are_compatible(ctx, llama_init_dft->context());
if (!vocab_dft_compatible) {
SRV_INF("the draft model '%s' is not compatible with the target model '%s'. tokens will be translated between the draft and target models.\n", params_base.speculative.model.path.c_str(), params_base.model.path.c_str());
}
const int n_ctx_dft = llama_n_ctx(llama_init_dft.context.get());
const int n_ctx_dft = llama_n_ctx(llama_init_dft->context());
cparams_dft = common_context_params_to_llama(params_dft);
cparams_dft.n_batch = n_ctx_dft;
// the context is not needed - we will create one for each slot
llama_init_dft.context.reset();
llama_init_dft->free_context();
}
chat_templates = common_chat_templates_init(model, params_base.chat_template);
@@ -1051,18 +1048,15 @@ struct server_context_impl {
// initialize samplers
{
if (slot.smpl != nullptr) {
common_sampler_free(slot.smpl);
}
slot.smpl.reset(common_sampler_init(model, task.params.sampling));
slot.smpl = common_sampler_init(model, task.params.sampling);
if (slot.smpl == nullptr) {
// for now, the only error that may happen here is invalid grammar
send_error(task, "Failed to parse grammar", ERROR_TYPE_INVALID_REQUEST);
return false;
}
SLT_INF(slot, "sampler chain: %s\n", common_sampler_print(slot.smpl).c_str());
SLT_INF(slot, "sampler chain: %s\n", common_sampler_print(slot.smpl.get()).c_str());
}
// initialize draft batch
@@ -1216,11 +1210,10 @@ struct server_context_impl {
}
void populate_token_probs(const server_slot & slot, completion_token_output & result, bool post_sampling, bool special, int idx) const {
size_t n_probs = slot.task->params.sampling.n_probs;
size_t n_vocab = llama_vocab_n_tokens(vocab);
const size_t n_probs = slot.task->params.sampling.n_probs;
if (post_sampling) {
const auto * cur_p = common_sampler_get_candidates(slot.smpl, true);
const auto * cur_p = common_sampler_get_candidates(slot.smpl.get(), true);
const size_t max_probs = cur_p->size;
// set probability for sampled token
@@ -1245,7 +1238,7 @@ struct server_context_impl {
std::vector<llama_token_data> cur = get_token_probabilities(ctx, idx);
// set probability for sampled token
for (size_t i = 0; i < n_vocab; i++) {
for (size_t i = 0; i < cur.size(); i++) {
// set probability for sampled token
if (cur[i].id == result.tok) {
result.prob = cur[i].p;
@@ -1255,7 +1248,7 @@ struct server_context_impl {
// set probability for top n_probs tokens
result.probs.reserve(n_probs);
for (size_t i = 0; i < std::min(n_vocab, n_probs); i++) {
for (size_t i = 0; i < std::min(cur.size(), n_probs); i++) {
result.probs.push_back({
cur[i].id,
common_token_to_piece(ctx, cur[i].id, special),
@@ -2301,13 +2294,13 @@ struct server_context_impl {
GGML_ASSERT(batch.n_tokens > 0);
common_sampler_reset(slot.smpl);
common_sampler_reset(slot.smpl.get());
// Process all prompt tokens through sampler system
for (int i = 0; i < slot.task->n_tokens(); ++i) {
llama_token id = input_tokens[i];
if (id != LLAMA_TOKEN_NULL) {
common_sampler_accept(slot.smpl, id, false);
common_sampler_accept(slot.smpl.get(), id, false);
}
}
@@ -2525,11 +2518,11 @@ struct server_context_impl {
const int tok_idx = slot.i_batch - i;
llama_token id = common_sampler_sample(slot.smpl, ctx, tok_idx);
llama_token id = common_sampler_sample(slot.smpl.get(), ctx, tok_idx);
slot.i_batch = -1;
common_sampler_accept(slot.smpl, id, true);
common_sampler_accept(slot.smpl.get(), id, true);
slot.n_decoded += 1;
@@ -2570,7 +2563,7 @@ struct server_context_impl {
size_t n_draft = slot.drafted.size();
// the accepted tokens from the speculation
const auto ids = common_sampler_sample_and_accept_n(slot.smpl, ctx, slot.i_batch_dft, slot.drafted);
const auto ids = common_sampler_sample_and_accept_n(slot.smpl.get(), ctx, slot.i_batch_dft, slot.drafted);
slot.i_batch_dft.clear();
slot.drafted.clear();
@@ -684,7 +684,7 @@ def test_anthropic_streaming_content_block_indices():
# Request that might produce both text and tool use
res = server.make_stream_request("POST", "/v1/messages", data={
"model": "test",
"max_tokens": 200,
"max_tokens": 400,
"stream": True,
"tools": [{
"name": "test_tool",
+6 -6
View File
@@ -568,10 +568,10 @@ int main(int argc, char ** argv) {
llama_context * ctx_ttc = NULL;
llama_context * ctx_cts = NULL;
common_init_result llama_init_ttc = common_init_from_params(params);
auto llama_init_ttc = common_init_from_params(params);
model_ttc = llama_init_ttc.model.get();
ctx_ttc = llama_init_ttc.context.get();
model_ttc = llama_init_ttc->model();
ctx_ttc = llama_init_ttc->context();
if (model_ttc == nullptr || ctx_ttc == nullptr) {
return ENOENT;
@@ -583,10 +583,10 @@ int main(int argc, char ** argv) {
params.embedding = true;
params.n_ubatch = params.n_batch;
common_init_result llama_init_cts = common_init_from_params(params);
auto llama_init_cts = common_init_from_params(params);
model_cts = llama_init_cts.model.get();
ctx_cts = llama_init_cts.context.get();
model_cts = llama_init_cts->model();
ctx_cts = llama_init_cts->context();
if (model_cts == nullptr || ctx_cts == nullptr) {
return ENOENT;