spec: Add benchmark-only synthetic speculative acceptance options (#27711)

* Add benchmark-only synthetic speculative acceptance to llama-server and llama-cli

* Address review comments

* Address review comments

* Add some comments in the code
This commit is contained in:
Gaurav Garg
2026-08-27 13:53:42 +03:00
committed by GitHub
parent deae5ee133
commit 2bb9bddafa
11 changed files with 427 additions and 17 deletions
+2
View File
@@ -259,6 +259,8 @@ For the full list of features, please refer to [server's changelog](https://gith
| `--spec-draft-n-cpu-moe, --spec-draft-ncmoe, -ncmoed, --n-cpu-moe-draft N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU for the draft model<br/>(env: LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE) |
| `--spec-draft-n-max N` | number of tokens to draft for speculative decoding (default: 3)<br/>(env: LLAMA_ARG_SPEC_DRAFT_N_MAX) |
| `--spec-draft-n-min N` | minimum number of draft tokens to use for speculative decoding (default: 0)<br/>(env: LLAMA_ARG_SPEC_DRAFT_N_MIN) |
| `--spec-synth-len L` | target mean synthetic acceptance length, including the target token (benchmarking only)<br/>(env: LLAMA_ARG_SPEC_SYNTH_LEN) |
| `--spec-synth-rates P0,P1,...` | comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)<br/>(env: LLAMA_ARG_SPEC_SYNTH_RATES) |
| `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)<br/>(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) |
| `--spec-draft-p-min, --draft-p-min P` | minimum speculative decoding probability (greedy) (default: 0.00)<br/>(env: LLAMA_ARG_SPEC_DRAFT_P_MIN) |
| `--spec-draft-backend-sampling, --no-spec-draft-backend-sampling` | offload draft sampling to the backend (default: enabled)<br/>(env: LLAMA_ARG_SPEC_DRAFT_BACKEND_SAMPLING) |
+68 -2
View File
@@ -23,6 +23,7 @@
#include <exception>
#include <memory>
#include <filesystem>
#include <random>
#include <utility>
#include <fstream>
@@ -51,6 +52,50 @@ static common_speculative_output_limits server_output_limits(const common_params
return result;
}
// synthetic draft verification for benchmarking - accept draft tokens at random instead of by match with the target
// on replay the draft was already accepted before a context checkpoint restore, so repeat the same decisions
static std::vector<llama_token> server_sample_and_accept_synth(
common_sampler * smpl,
llama_context * ctx,
const std::vector<int32_t> & idxs,
const llama_tokens & draft,
const std::vector<double> & synth_probs,
std::mt19937 & rng,
bool is_replay) {
GGML_ASSERT(idxs.size() == draft.size() + 1);
GGML_ASSERT(synth_probs.size() >= draft.size());
std::vector<llama_token> result;
result.reserve(idxs.size());
const llama_vocab * vocab = llama_model_get_vocab(llama_get_model(ctx));
std::uniform_real_distribution<double> dist(0.0, 1.0);
for (size_t i = 0; i < draft.size(); ++i) {
const llama_token id = common_sampler_sample(smpl, ctx, idxs[i]);
const bool accept = is_replay || dist(rng) < synth_probs[i];
// do not accept a drafted EOG token - it would end the generation early
// on replay the last token is from the target and can be EOG, so skip this check
if (accept && (is_replay || !llama_vocab_is_eog(vocab, draft[i]))) {
// synthetic draft tokens do not advance grammar or reasoning state
// the last replay token is from the target and must advance both
const bool is_replay_target = is_replay && i + 1 == draft.size();
common_sampler_accept(smpl, draft[i], is_replay_target);
result.push_back(draft[i]);
continue;
}
common_sampler_accept(smpl, id, true);
result.push_back(id);
return result;
}
const llama_token id = common_sampler_sample(smpl, ctx, idxs[draft.size()]);
common_sampler_accept(smpl, id, true);
result.push_back(id);
return result;
}
// state diagram: https://github.com/ggml-org/llama.cpp/pull/9283
enum slot_state {
SLOT_STATE_IDLE,
@@ -211,6 +256,7 @@ struct server_slot {
std::vector<int32_t> spec_i_batch;
common_prompt_checkpoint spec_ckpt;
bool spec_is_replay = false;
std::mt19937 spec_synth_rng;
// 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
@@ -1194,6 +1240,9 @@ private:
spec.reset(common_speculative_init(params_base.speculative, params_base.n_parallel));
} catch (const std::exception & e) {
SRV_ERR("failed to initialize speculative decoding context: %s\n", e.what());
if (params_base.speculative.has_synth()) {
return false;
}
}
}
@@ -1209,6 +1258,11 @@ private:
model_dft = nullptr;
}
if (!spec && params_base.speculative.has_synth()) {
SRV_ERR("%s", "synthetic acceptance requires an initialized speculative decoding context\n");
return false;
}
for (int i = 0; i < params_base.n_parallel; i++) {
server_slot & slot = slots[i];
@@ -1717,6 +1771,13 @@ private:
SLT_TRC(slot, "sampler chain: %s\n", common_sampler_print(slot.smpl.get()).c_str());
SLT_TRC(slot, "sampler params: \n%s\n", task.params.sampling.print().c_str());
if (spec && !common_speculative_get_synth_probs(spec.get()).empty()) {
const uint32_t seed = task.params.sampling.seed == LLAMA_DEFAULT_SEED
? std::random_device{}()
: task.params.sampling.seed;
slot.spec_synth_rng.seed(seed);
}
} else {
slot.smpl.reset();
}
@@ -3802,7 +3863,12 @@ private:
common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get()));
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);
const auto & synth_probs = common_speculative_get_synth_probs(spec.get());
auto accepted = synth_probs.empty()
? common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft)
: server_sample_and_accept_synth(
slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft,
synth_probs, slot.spec_synth_rng, slot.spec_is_replay);
slot.spec_i_batch.clear();
GGML_ASSERT(accepted.size() >= 1);
@@ -3868,7 +3934,7 @@ private:
auto & n_accepted_per_pos = slot.n_accepted_per_pos;
if (n_accepted_per_pos.empty()) {
n_accepted_per_pos.resize(common_speculative_n_max(&params_base.speculative), 0);
n_accepted_per_pos.resize(common_speculative_n_max(spec.get()), 0);
}
for (size_t i = 0; i < n_accepted && i < n_accepted_per_pos.size(); ++i) {
n_accepted_per_pos[i]++;
@@ -52,6 +52,18 @@ def test_with_and_without_draft():
assert tokens_no_draft == tokens_draft
server.stop()
create_server()
assert server.spec_draft_n_max is not None
server.spec_synth_rates = [0.0] * server.spec_draft_n_max
server.start()
res = server.make_request("POST", "/completion", data=request)
assert res.status_code == 200
assert res.body["timings"]["draft_n"] > 0
assert res.body["timings"]["draft_n_accepted"] == 0
assert res.body["tokens"] == tokens_no_draft
def test_different_draft_min_draft_max():
global server
@@ -80,6 +92,66 @@ def test_different_draft_min_draft_max():
last_content = res.body["content"]
def test_synth_is_deterministic():
global server
assert server.spec_draft_n_max is not None
server.spec_synth_rates = [0.75 ** (i + 1) for i in range(server.spec_draft_n_max)]
server.start()
request = {
"prompt": "I believe the meaning of life is",
"temperature": 0.2,
"top_k": 5,
"seed": 4242,
"n_predict": 32,
}
responses = [server.make_request("POST", "/completion", data=request) for _ in range(2)]
for res in responses:
assert res.status_code == 200
assert res.body["timings"]["draft_n"] > 0
assert responses[0].body["timings"]["draft_n"] == responses[1].body["timings"]["draft_n"]
assert responses[0].body["timings"]["draft_n_accepted"] == responses[1].body["timings"]["draft_n_accepted"]
def test_synth_ignores_target_tokens():
global server
assert server.spec_draft_n_max is not None
server.spec_synth_rates = [1.0] * server.spec_draft_n_max
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"temperature": 0.0,
"seed": 4242,
"n_predict": 32,
})
assert res.status_code == 200
assert res.body["timings"]["draft_n"] > 0
assert res.body["timings"]["draft_n_accepted"] == res.body["timings"]["draft_n"]
res = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"temperature": 0.0,
"seed": 4242,
"n_predict": 6,
"grammar": 'root ::= "a"{5,5}',
})
assert res.status_code == 200, res.body
res = server.make_request("POST", "/completion", data={
"prompt": "Respond with only: OK",
"temperature": 0.0,
"seed": 4242,
"n_predict": 64,
"ignore_eos": True,
})
assert res.status_code == 200, res.body
assert res.body["tokens_predicted"] == 64
assert res.body["stop_type"] == "limit"
def test_slot_ctx_not_exceeded():
global server
server.n_ctx = 256
+7
View File
@@ -99,6 +99,8 @@ class ServerProcess:
spec_type: str | None = None
spec_draft_n_min: int | None = None
spec_draft_n_max: int | None = None
spec_synth_len: float | None = None
spec_synth_rates: List[float] | None = None
no_ui: bool | None = None
jinja: bool | None = None
reasoning_format: Literal['deepseek', 'none', 'nothink'] | None = None
@@ -245,6 +247,11 @@ class ServerProcess:
server_args.extend(["--spec-draft-n-max", self.spec_draft_n_max])
if self.spec_draft_n_min:
server_args.extend(["--spec-draft-n-min", self.spec_draft_n_min])
if self.spec_synth_len is not None:
server_args.extend(["--spec-synth-len", self.spec_synth_len])
if self.spec_synth_rates is not None:
rates = ",".join(str(rate) for rate in self.spec_synth_rates)
server_args.extend(["--spec-synth-rates", rates])
if self.no_ui:
server_args.append("--no-ui")
if self.no_models_autoload: