From c22562c7112cb745df7f75a21b81f90dd7ae2802 Mon Sep 17 00:00:00 2001 From: Nigel Bosch Date: Sun, 16 Aug 2026 11:14:30 -0500 Subject: [PATCH 1/2] return logprobs in spec path --- tools/server/server-context.cpp | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index fdc8f2b805fd..7944b70bc723 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3811,13 +3811,15 @@ struct server_context_impl { 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); - auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); - slot.spec_i_batch.clear(); + GGML_ASSERT(spec_i_batch.size() == n_draft + 1); + auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, spec_i_batch, slot.spec_draft); GGML_ASSERT(accepted.size() >= 1); @@ -3904,7 +3906,10 @@ struct server_context_impl { 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; From 62bddb423312ad046e486a5cad3adefc8985f9e4 Mon Sep 17 00:00:00 2001 From: Jeonghyun Lee Date: Tue, 25 Aug 2026 23:59:01 +0900 Subject: [PATCH 2/2] 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 Claude-Session: https://claude.ai/code/session_01KARBWpPcijcBXzJEvttuwk --- tools/server/tests/unit/test_speculative.py | 36 +++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tools/server/tests/unit/test_speculative.py b/tools/server/tests/unit/test_speculative.py index 22b523954ec7..27ba4149e1a5 100644 --- a/tools/server/tests/unit/test_speculative.py +++ b/tools/server/tests/unit/test_speculative.py @@ -203,3 +203,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)