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);