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:
@@ -200,6 +200,8 @@
|
||||
| `--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) |
|
||||
|
||||
@@ -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) |
|
||||
|
||||
@@ -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(¶ms_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
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user