diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f5477356d61..b4fcb778e3b 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3879,18 +3879,20 @@ 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); + GGML_ASSERT(spec_i_batch.size() == n_draft + 1); 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) + ? common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, spec_i_batch, slot.spec_draft) : 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); - slot.spec_i_batch.clear(); GGML_ASSERT(accepted.size() >= 1); @@ -3977,7 +3979,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; diff --git a/tools/server/tests/unit/test_speculative.py b/tools/server/tests/unit/test_speculative.py index 22b523954ec..27ba4149e1a 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)