More fixes

This commit is contained in:
Gaurav Garg
2026-07-05 23:15:39 +05:30
parent eab4e8444f
commit f1f008b2ee
5 changed files with 153 additions and 29 deletions
+11 -3
View File
@@ -1674,6 +1674,10 @@ int llama_context::decode(const llama_batch & batch_inp) {
for (int32_t s = 0; s < ns; ++s) {
const llama_seq_id seq_id = batch_inp.seq_id ? batch_inp.seq_id[i][s] : 0;
if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) {
continue;
}
seq_output_count[seq_id]++;
auto sampler = sampling.samplers.find(seq_id);
if (sampler != sampling.samplers.end() &&
@@ -2310,9 +2314,13 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
}
}
const uint32_t n_sampling_outputs_max = std::min<uint64_t>(
std::min(n_tokens, cparams.n_outputs_max),
(uint64_t) cparams.n_seq_max * cparams.n_sampling_outputs_per_seq_max);
res += n_sampling_nodes;
if (cparams.n_outputs_max > 1) {
res += (cparams.n_outputs_max - 1) * n_sampling_nodes_max;
if (n_sampling_outputs_max > 1) {
res += (n_sampling_outputs_max - 1) * n_sampling_nodes_max;
}
return res;
}
@@ -2361,7 +2369,7 @@ ggml_cgraph * llama_context::graph_reserve(
for (uint32_t s = 0; s < n_seqs && n_outputs_set < n_outputs; ++s) {
const auto sampler = sampling.samplers.find(s);
if (t > 0 && (sampler == sampling.samplers.end() ||
cparams.n_sampling_outputs_per_seq_max == 1)) {
t >= cparams.n_sampling_outputs_per_seq_max)) {
continue;
}
+13 -4
View File
@@ -603,6 +603,10 @@ static llama_sampler_backend_probe llama_sampler_backend_probe_graph(
}
}
if (sampler->iface->backend_reset) {
sampler->iface->backend_reset(sampler);
}
return { std::move(ctx_ptr), gf };
}
@@ -1085,6 +1089,8 @@ struct llama_sampler_dist : public llama_sampler_backend {
// inputs for the current sampling graph
std::vector<ggml_tensor *> inp_uniforms;
bool copy_candidates = false;
};
static const char * llama_sampler_dist_name(const struct llama_sampler * smpl) {
@@ -1193,7 +1199,8 @@ static bool llama_sampler_dist_backend_init(
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_per_seq_max) {
auto * sctx = (llama_sampler_dist *) smpl->ctx;
GGML_UNUSED(n_outputs_per_seq_max);
sctx->copy_candidates = n_outputs_per_seq_max > 1;
const bool res = llama_sampler_backend_support(smpl, buft);
@@ -1261,9 +1268,11 @@ static void llama_sampler_dist_backend_apply(
sampled_token = ggml_get_rows(ctx, candidates, idx);
ggml_set_name(sampled_token, "dist_sampled_token");
// candidates may be a view whose backing storage can be reused
data->candidates = ggml_cont(ctx, data->candidates);
ggml_set_name(data->candidates, "dist_candidates_out");
if (sctx->copy_candidates) {
// candidates may be a view whose backing storage can be reused
data->candidates = ggml_cont(ctx, data->candidates);
ggml_set_name(data->candidates, "dist_candidates_out");
}
}
data->sampled = sampled_token;
+106 -9
View File
@@ -1666,6 +1666,99 @@ static void test_backend_multi_output_greedy(const test_params & params) {
printf("backend multi-output greedy test PASSED\n");
}
static void test_backend_multi_sequence_multi_output_dist(const test_params & params) {
const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
const int32_t n_vocab = llama_vocab_n_tokens(vocab);
const uint32_t seeds[] = { 88, 1337 };
// reduce the chance that swapped random inputs select the same token
const float temp = 10.0f;
llama_sampler_ptr chain_0(llama_sampler_chain_init(llama_sampler_chain_default_params()));
llama_sampler_ptr chain_1(llama_sampler_chain_init(llama_sampler_chain_default_params()));
llama_sampler_chain_add(chain_0.get(), llama_sampler_init_temp(temp));
llama_sampler_chain_add(chain_0.get(), llama_sampler_init_dist(seeds[0]));
llama_sampler_chain_add(chain_1.get(), llama_sampler_init_temp(temp));
llama_sampler_chain_add(chain_1.get(), llama_sampler_init_dist(seeds[1]));
std::vector<llama_sampler_seq_config> configs = {
{ 0, chain_0.get() },
{ 1, chain_1.get() },
};
test_context test_ctx(params, configs, 2, 6, 0, 3);
std::vector<llama_sampler_seq_config> reference_configs;
test_context reference_ctx(params, reference_configs, 2, 6);
const llama_token seq_tokens[2][3] = {
{ llama_vocab_bos(vocab), llama_vocab_eos(vocab), llama_vocab_bos(vocab) },
{ llama_vocab_eos(vocab), llama_vocab_bos(vocab), llama_vocab_eos(vocab) },
};
llama_batch batch = llama_batch_init(6, 0, 1);
for (int pos = 0; pos < 3; ++pos) {
common_batch_add(batch, seq_tokens[0][pos], pos, { 0 }, true);
common_batch_add(batch, seq_tokens[1][pos], pos, { 1 }, true);
}
GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
GGML_ASSERT(llama_decode(reference_ctx.ctx.get(), batch) == 0);
std::mt19937 reference_rngs[] = {
std::mt19937(seeds[0]),
std::mt19937(seeds[1]),
};
std::uniform_real_distribution<double> reference_dist(0.0, 1.0);
int outputs_per_seq[] = { 0, 0 };
for (int i = 0; i < batch.n_tokens; ++i) {
const llama_seq_id seq_id = batch.seq_id[i][0];
GGML_ASSERT(seq_id == 0 || seq_id == 1);
outputs_per_seq[seq_id]++;
const llama_token backend_token = llama_get_sampled_token_ith(test_ctx.ctx.get(), i);
const float * sampled_logits = llama_get_sampled_logits_ith(test_ctx.ctx.get(), i);
const float * sampled_probs = llama_get_sampled_probs_ith(test_ctx.ctx.get(), i);
const uint32_t n_logits = llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), i);
const uint32_t n_probs = llama_get_sampled_probs_count_ith(test_ctx.ctx.get(), i);
const float * reference_logits = llama_get_logits_ith(reference_ctx.ctx.get(), i);
GGML_ASSERT(backend_token >= 0 && backend_token < n_vocab);
GGML_ASSERT(sampled_logits != nullptr);
GGML_ASSERT(sampled_probs != nullptr);
GGML_ASSERT(reference_logits != nullptr);
GGML_ASSERT(n_logits == (uint32_t) n_vocab);
GGML_ASSERT(n_probs == (uint32_t) n_vocab);
float prob_sum = 0.0f;
float cumsum_before = 0.0f;
for (llama_token token = 0; token < n_vocab; ++token) {
const float expected_logit = reference_logits[token] / temp;
const float tolerance = 1e-4f * std::max(1.0f, std::fabs(expected_logit));
GGML_ASSERT(std::fabs(sampled_logits[token] - expected_logit) <= tolerance);
GGML_ASSERT(std::isfinite(sampled_probs[token]));
GGML_ASSERT(sampled_probs[token] >= 0.0f);
prob_sum += sampled_probs[token];
if (token < backend_token) {
cumsum_before += sampled_probs[token];
}
}
GGML_ASSERT(std::fabs(prob_sum - 1.0f) <= 1e-3f);
const float rnd = reference_dist(reference_rngs[seq_id]);
const float cumsum_sampled = cumsum_before + sampled_probs[backend_token];
GGML_ASSERT(rnd >= cumsum_before - 1e-4f);
GGML_ASSERT(rnd <= cumsum_sampled + 1e-4f);
}
GGML_ASSERT(outputs_per_seq[0] == 3);
GGML_ASSERT(outputs_per_seq[1] == 3);
llama_batch_free(batch);
printf("backend multi-sequence multi-output dist test PASSED\n");
}
static void test_backend_multi_output_sampling_chain(const test_params & params) {
const llama_seq_id seq_id = 0;
const int32_t seed = 88;
@@ -1698,12 +1791,15 @@ static void test_backend_multi_output_sampling_chain(const test_params & params)
std::vector<llama_token_data> reference_data(n_vocab);
std::mt19937 reference_rng(seed);
std::uniform_real_distribution<double> reference_dist(0.0, 1.0);
int32_t n_reused_after_first_round = -1;
const int n_outputs_per_round[] = { 4, 3, 4, 4 };
int32_t n_reused_before_repeat = -1;
int32_t pos = 0;
for (int round = 0; round < 2; ++round) {
llama_batch batch = llama_batch_init(4, 0, 1);
for (int i = 0; i < 4; ++i) {
common_batch_add(batch, llama_vocab_bos(vocab), round * 4 + i, { seq_id }, true);
for (int round = 0; round < (int) (sizeof(n_outputs_per_round) / sizeof(n_outputs_per_round[0])); ++round) {
const int n_outputs = n_outputs_per_round[round];
llama_batch batch = llama_batch_init(n_outputs, 0, 1);
for (int i = 0; i < n_outputs; ++i) {
common_batch_add(batch, llama_vocab_bos(vocab), pos++, { seq_id }, true);
}
GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
@@ -1795,10 +1891,10 @@ static void test_backend_multi_output_sampling_chain(const test_params & params)
llama_batch_free(batch);
const int32_t n_reused = llama_perf_context(test_ctx.ctx.get()).n_reused;
if (round == 0) {
n_reused_after_first_round = n_reused;
} else {
GGML_ASSERT(n_reused > n_reused_after_first_round);
if (round == 2) {
n_reused_before_repeat = n_reused;
} else if (round == 3) {
GGML_ASSERT(n_reused > n_reused_before_repeat);
}
}
@@ -1884,6 +1980,7 @@ static const backend_test_case BACKEND_TESTS[] = {
{ "set_sampler", test_backend_set_sampler, true },
{ "multi_output_limit", test_backend_multi_output_limit, true },
{ "multi_output_greedy", test_backend_multi_output_greedy, true },
{ "multi_sequence_multi_output_dist", test_backend_multi_sequence_multi_output_dist, true },
{ "multi_output_sampling_chain", test_backend_multi_output_sampling_chain, true },
{ "multi_output_cpu", test_backend_multi_output_cpu_suffix, true },
{ "mixed", test_backend_mixed_sampling, true },
+9 -2
View File
@@ -222,6 +222,7 @@ struct server_slot {
std::vector<int32_t> spec_i_batch;
common_prompt_checkpoint spec_ckpt;
bool spec_is_replay = false;
common_sampler_ptr spec_smpl_save;
// TODO: move members that belong to the task (such as `generated_text`, `has_new_line`) to task_results_state
// see https://github.com/ggml-org/llama.cpp/pull/18283#issuecomment-3710175837
@@ -359,6 +360,7 @@ struct server_slot {
spec_draft.clear();
spec_i_batch.clear();
spec_ckpt.clear();
spec_smpl_save.reset();
}
generated_tokens.clear();
generated_token_probs.clear();
@@ -3108,6 +3110,11 @@ private:
// update the batch with the sampled/drafted tokens
iterate(generating, [&](server_slot & slot) {
GGML_ASSERT(!slot.spec_smpl_save);
if (!slot.spec_draft.empty()) {
// backend sampling advances the sampler during llama_decode()
slot.spec_smpl_save.reset(common_sampler_clone(slot.smpl.get()));
}
slot.handle_last_sampled_token(batch);
});
@@ -3886,8 +3893,8 @@ private:
// verify and try to accept the draft
{
// save the sampler sampler state in case we need to restore it
common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get()));
GGML_ASSERT(slot.spec_smpl_save);
common_sampler_ptr smpl_save = std::move(slot.spec_smpl_save);
GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1);
auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft);
+14 -11
View File
@@ -25,34 +25,37 @@ def fixture_create_server():
def test_with_and_without_draft():
global server
server.model_draft = None # disable draft model
server.spec_type = None
server.start()
res = server.make_request("POST", "/completion", data={
request = {
"prompt": "I believe the meaning of life is",
"temperature": 0.0,
"top_k": 1,
"seed": 4242,
"n_predict": 16,
})
"return_tokens": True,
}
server.model_draft = None # disable draft model
server.spec_type = None
server.backend_sampling = True
server.start()
res = server.make_request("POST", "/completion", data=request)
assert res.status_code == 200
content_no_draft = res.body["content"]
tokens_no_draft = res.body["tokens"]
server.stop()
# create new server with draft model
create_server()
server.backend_sampling = True
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"temperature": 0.0,
"top_k": 1,
"n_predict": 16,
})
res = server.make_request("POST", "/completion", data=request)
assert res.status_code == 200
assert res.body["timings"]["draft_n"] > 0
content_draft = res.body["content"]
tokens_draft = res.body["tokens"]
assert content_no_draft == content_draft
assert tokens_no_draft == tokens_draft
def test_different_draft_min_draft_max():