More fixes

tools/server/server-context.cpp

@@ -222,6 +222,7 @@ struct server_slot {
     std::vector<int32_t> spec_i_batch;
     common_prompt_checkpoint spec_ckpt;
     bool spec_is_replay = false;
+    common_sampler_ptr spec_smpl_save;
 
     // 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
@@ -359,6 +360,7 @@ struct server_slot {
             spec_draft.clear();
             spec_i_batch.clear();
             spec_ckpt.clear();
+            spec_smpl_save.reset();
         }
         generated_tokens.clear();
         generated_token_probs.clear();
@@ -3108,6 +3110,11 @@ private:
 
         // update the batch with the sampled/drafted tokens
         iterate(generating, [&](server_slot & slot) {
+            GGML_ASSERT(!slot.spec_smpl_save);
+            if (!slot.spec_draft.empty()) {
+                // backend sampling advances the sampler during llama_decode()
+                slot.spec_smpl_save.reset(common_sampler_clone(slot.smpl.get()));
+            }
             slot.handle_last_sampled_token(batch);
         });
 
@@ -3886,8 +3893,8 @@ private:
 
             // verify and try to accept the draft
             {
-                // save the sampler sampler state in case we need to restore it
-                common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get()));
+                GGML_ASSERT(slot.spec_smpl_save);
+                common_sampler_ptr smpl_save = std::move(slot.spec_smpl_save);
 
                 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);

tools/server/tests/unit/test_speculative.py

@@ -25,34 +25,37 @@ def fixture_create_server():
 
 def test_with_and_without_draft():
     global server
-    server.model_draft = None  # disable draft model
-    server.spec_type = None
-    server.start()
-    res = server.make_request("POST", "/completion", data={
+    request = {
         "prompt": "I believe the meaning of life is",
         "temperature": 0.0,
         "top_k": 1,
+        "seed": 4242,
         "n_predict": 16,
-    })
+        "return_tokens": True,
+    }
+
+    server.model_draft = None  # disable draft model
+    server.spec_type = None
+    server.backend_sampling = True
+    server.start()
+    res = server.make_request("POST", "/completion", data=request)
     assert res.status_code == 200
     content_no_draft = res.body["content"]
+    tokens_no_draft = res.body["tokens"]
     server.stop()
 
     # create new server with draft model
     create_server()
     server.backend_sampling = True
     server.start()
-    res = server.make_request("POST", "/completion", data={
-        "prompt": "I believe the meaning of life is",
-        "temperature": 0.0,
-        "top_k": 1,
-        "n_predict": 16,
-    })
+    res = server.make_request("POST", "/completion", data=request)
     assert res.status_code == 200
     assert res.body["timings"]["draft_n"] > 0
     content_draft = res.body["content"]
+    tokens_draft = res.body["tokens"]
 
     assert content_no_draft == content_draft
+    assert tokens_no_draft == tokens_draft
 
 
 def test_different_draft_min_draft_max():