Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-08-26 15:16:23 +02:00
Code Issues Packages Projects Releases Wiki Activity

llama : support multi-output backend sampling (#25532)

Browse Source 
* Enable backend sampling with token speculation

* Clamp the mask sum before converting it into the sampled index

* Add a numeric context parameter declaring the maximum outputs one sequence

* More fixes

* Don't reuse memory for output views.

* Match dist between CPU and GPU

* Fix CPU and backend sampling mismatches

* Simpify some of the changes

* Fix tests on Vulkan

* More test fixes

* Rebase changes

* Rebase and address review comments

* Address review comments

* Address review comments

* Update src/llama-sampler.cpp

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
Gaurav Garg
2026-08-10 19:28:56 +05:30
committed by GitHub
parent d2f83055d6
commit dd1ea52433
24 changed files with 1305 additions and 391 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/server/server-context.cpp
+16 -20
View File
@@ -39,19 +39,18 @@ using json = nlohmann::ordered_json;
constexpr int HTTP_POLLING_SECONDS = 1;
static uint32_t server_n_outputs_max(const common_params & params) {
const uint32_t n_batch = params.n_batch;
static common_speculative_output_limits server_output_limits(const common_params & params) {
if (params.embedding ||
(params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) {
return n_batch;
return { params.n_batch, 1 };
}
const uint32_t n_outputs_per_seq = 1 + common_speculative_n_max(&params.speculative);
auto result = common_speculative_get_output_limits(
params.n_batch, params.n_parallel, common_speculative_n_max(&params.speculative));
const uint64_t n_outputs = (uint64_t) params.n_parallel * n_outputs_per_seq;
return std::max<uint32_t>(1, std::min<uint64_t>(n_batch, n_outputs));
result.total = std::max<int32_t>(1, result.total);
result.per_seq = std::max<int32_t>(1, result.per_seq);
return result;
}
// state diagram: https://github.com/ggml-org/llama.cpp/pull/9283
@@ -1063,7 +1062,9 @@ private:
const bool is_resume = sleeping;
params_base = params;
params_base.n_outputs_max = server_n_outputs_max(params_base);
const auto output_limits = server_output_limits(params_base);
params_base.n_outputs_max = output_limits.total;
params_base.n_outputs_max_per_seq = output_limits.per_seq;
const bool has_mmproj = !params.mmproj.path.empty();
const bool has_draft = params.speculative.has_dft();
@@ -1832,18 +1833,13 @@ private:
const bool need_pre_sample_logits = task.params.sampling.n_probs > 0 && !task.params.post_sampling_probs;
bool backend_sampling = true;
backend_sampling &= task.params.sampling.backend_sampling;
// TODO: speculative decoding requires multiple samples per batch - not supported yet
backend_sampling &= !(slot.can_speculate());
bool use_backend_sampling = task.params.sampling.backend_sampling;
// TODO: getting pre sampling logits is not yet supported with backend sampling
backend_sampling &= !need_pre_sample_logits;
use_backend_sampling &= !need_pre_sample_logits;
// TODO: tmp until backend sampling is fully implemented
if (backend_sampling) {
if (use_backend_sampling) {
llama_set_sampler(ctx_tgt, slot.id, common_sampler_get(slot.smpl.get()));
} else {
llama_set_sampler(ctx_tgt, slot.id, nullptr);
@@ -3865,7 +3861,8 @@ private:
// speculative decoding - main model sample and accept
iterate(slots, [&](server_slot & slot) {
if (slot.state != SLOT_STATE_GENERATING || !slot.can_speculate() || slot.spec_draft.empty()) {
if (slot.state != SLOT_STATE_GENERATING || !slot.can_speculate() ||
slot.spec_draft.empty() || slot.spec_i_batch.empty()) {
return;
}
@@ -3876,7 +3873,6 @@ 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_i_batch.size() == n_draft + 1);
@@ -3915,7 +3911,7 @@ private:
slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1);
slot.prompt.tokens.keep_first(ckpt.n_tokens);
slot.smpl = std::move(smpl_save);
common_sampler_copy(smpl_save.get(), slot.smpl.get());
return;
}
tools/server/tests/unit/test_speculative.py
+16 -15
View File
@@ -25,33 +25,34 @@ def fixture_create_server():
def test_with_and_without_draft():
global server
request = {
"prompt": "I believe the meaning of life is",
"temperature": 0.8,
"top_k": 40,
"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={
"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
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():