Commit ead9f3205 for llama.cpp
commit ead9f3205848f81ca079b4e760b76a01a6359f4b
Author: Nigel Bosch <pnigelb@gmail.com>
Date: Tue Oct 20 12:36:14 2026 -0500
server: Support logprobs with speculative decoding (#27196)
* return logprobs in spec path
* server: add regression test for logprobs with speculative decoding
Asserts accepted draft tokens report the target model's logprobs,
comparing a draft-enabled run against a no-draft baseline (#24271).
Fails on master (top_logprobs empty, logprob 0.0), passes with the
fix from #27196.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KARBWpPcijcBXzJEvttuwk
---------
Co-authored-by: Jeonghyun Lee <sig.jhl@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index 84cd67768..d66aae426 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -4522,11 +4522,14 @@ private:
GGML_ASSERT(n_draft > 0);
+ // batch indices of the draft tokens, used to get the token probabilities below
+ auto spec_i_batch = std::move(slot.spec_i_batch);
+
// verify and try to accept the draft
{
common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get()));
- GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1);
+ GGML_ASSERT(spec_i_batch.size() == n_draft + 1);
GGML_ASSERT(slot.spec_draft_q.empty() || (slot.spec_draft_q.size() == slot.spec_draft.size()));
const auto & synth_probs = common_speculative_get_synth_probs(spec.get());
@@ -4537,16 +4540,15 @@ private:
if (!synth_probs.empty()) {
// synthetic acceptance replaces verification entirely, so it comes first
accepted = server_sample_and_accept_synth(
- slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft,
+ slot.smpl.get(), slot.ctx_tgt, spec_i_batch, slot.spec_draft,
synth_probs, slot.spec_synth_rng, slot.spec_is_replay);
} else if (slot.spec_is_replay && slot.use_spec_rejection()) {
- accepted = server_accept_replay(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft);
+ accepted = server_accept_replay(slot.smpl.get(), slot.ctx_tgt, spec_i_batch, slot.spec_draft);
} else if (use_rejection) {
- accepted = common_sampler_sample_and_accept_n_rejection(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft, slot.spec_draft_q);
+ accepted = common_sampler_sample_and_accept_n_rejection(slot.smpl.get(), slot.ctx_tgt, spec_i_batch, slot.spec_draft, slot.spec_draft_q);
} else {
- accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft);
+ accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, spec_i_batch, slot.spec_draft);
}
- slot.spec_i_batch.clear();
GGML_ASSERT(accepted.size() >= 1);
@@ -4633,7 +4635,10 @@ private:
result.text_to_send = common_token_to_piece(slot.ctx_tgt, result.tok, accept_special_token(slot, result.tok));
result.prob = 1.0f; // set later
- // TODO: set result.probs
+ // post_sampling_probs is not supported with speculative decoding
+ if (slot.task->params.sampling.n_probs > 0 && !slot.task->params.post_sampling_probs) {
+ populate_token_probs(slot, result, false, params_base.special, spec_i_batch[i]);
+ }
slot.stats.n_gen += 1;
diff --git a/tools/server/tests/unit/test_speculative.py b/tools/server/tests/unit/test_speculative.py
index 7a0354302..0d41b1173 100644
--- a/tools/server/tests/unit/test_speculative.py
+++ b/tools/server/tests/unit/test_speculative.py
@@ -229,3 +229,39 @@ def test_multi_requests_parallel(n_slots: int, n_requests: int):
for res in results:
assert res.status_code == 200
assert match_regex("(wise|kind|owl|answer)+", res.body["content"])
+
+
+def test_draft_token_probs():
+ # tokens accepted from the draft must carry the target model's probabilities,
+ # same as without speculative decoding (#24271)
+ global server
+ request = {
+ "prompt": "I believe the meaning of life is",
+ "temperature": 0.0,
+ "top_k": 1,
+ "n_predict": 16,
+ "n_probs": 4,
+ }
+
+ server.model_draft = None # disable draft model
+ server.spec_type = None
+ server.start()
+ res = server.make_request("POST", "/completion", data=request)
+ assert res.status_code == 200
+ probs_no_draft = res.body["completion_probabilities"]
+ server.stop()
+
+ create_server()
+ server.start()
+ res = server.make_request("POST", "/completion", data=request)
+ assert res.status_code == 200
+ assert res.body["timings"]["draft_n"] > 0
+ probs_draft = res.body["completion_probabilities"]
+
+ assert len(probs_draft) == len(probs_no_draft)
+ for tok_draft, tok_no_draft in zip(probs_draft, probs_no_draft):
+ assert tok_draft["id"] == tok_no_draft["id"]
+ assert len(tok_draft["top_logprobs"]) == len(tok_no_draft["top_logprobs"])
+ assert abs(tok_draft["logprob"] - tok_no_draft["logprob"]) < 0.25
+ # real logprobs are not all exactly 0.0 (prob 1.0)
+ assert any(tok["logprob"] < -1e-4 for tok in probs_draft)