speculative: fix sampler for checkpoints

common/sampling.cpp

@@ -182,7 +182,9 @@ void common_sampler_clone(common_sampler * src, common_sampler * dst) {
         llama_sampler_dry_free(dst->smpl);
         dst->smpl = nullptr;
     }
-    dst->smpl = llama_sampler_dry_clone(src->smpl);
+    if (src->smpl) {
+        dst->smpl = llama_sampler_dry_clone(src->smpl);
+    }
 
 }
 

common/speculative.cpp

@@ -1133,7 +1133,9 @@ void common_speculative_accept(common_speculative * spec, uint16_t n_accepted) {
 
     common_speculative_state * impl = spec->curr_impl;
 
-    GGML_ASSERT(impl);
+    if (!impl) {
+        return;
+    }
 
     {
         common_time_meas tm(impl->t_accept_us, !impl->gen_perf);

examples/server/server-context.cpp

@@ -43,6 +43,7 @@ static void discard_speculative_checkpoint(server_slot & slot, llama_context * c
 static bool save_speculative_checkpoint(server_slot & slot, llama_model * model, llama_context * ctx, bool gpu_ckpt) {
     slot.spec_ckpt.clear();
     slot.spec_ckpt.n_past = slot.n_past - (int32_t)(slot.drafted.size() + 1);
+    slot.spec_ckpt.sampled = slot.sampled;
 
     if (gpu_ckpt) {
         slot.spec_ckpt.valid = llama_kv_cache_checkpoint_save(ctx);
@@ -3675,11 +3676,13 @@ void server_context::speculative_decoding_accept() {
             }
 
             if (!ids.empty()) {
-                const int n_accepted = (int)ids.size();
-                llama_batch re_batch = llama_batch_init(n_accepted, 0, 1);
-                for (int j = 0; j < n_accepted; j++) {
-                    const bool is_last = (j == n_accepted - 1);
-                    common_batch_add(re_batch, ids[j], slot.spec_ckpt.n_past + j, { slot.id }, is_last);
+                // Re-decode to restore the recurrent state after checkpoint restore.
+                const int n_re = (int)ids.size();
+                llama_batch re_batch = llama_batch_init(n_re, 0, 1);
+                common_batch_add(re_batch, slot.spec_ckpt.sampled, slot.spec_ckpt.n_past, { slot.id }, n_re == 1);
+                for (int j = 0; j < n_re - 1; j++) {
+                    const bool is_last = (j == n_re - 2);
+                    common_batch_add(re_batch, ids[j], slot.spec_ckpt.n_past + 1 + j, { slot.id }, is_last);
                 }
 
                 if (slot.has_mtp) {
@@ -3706,8 +3709,8 @@ void server_context::speculative_decoding_accept() {
                 }
 
                 llama_batch_free(re_batch);
-                SLT_DBG(slot, "spec checkpoint restored (gpu=%s): re-decoded %d accepted tokens (rejected %d)\n",
-                    gpu_ckpt ? "yes" : "no", n_accepted, (int)(n_draft - (ids.size() - 1)));
+                SLT_DBG(slot, "spec checkpoint restored (gpu=%s): re-decoded %d tokens (rejected %d drafts)\n",
+                    gpu_ckpt ? "yes" : "no", n_re, (int)(n_draft - (ids.size() - 1)));
             }
 
             discard_speculative_checkpoint(slot, ctx, gpu_ckpt);

examples/server/server-context.h

@@ -25,6 +25,7 @@ enum slot_command {
 struct server_speculative_checkpoint {
     bool valid = false;
     llama_pos n_past = 0;
+    llama_token sampled = LLAMA_TOKEN_NULL;
     std::vector<uint8_t> data;
     common_sampler * sampler = nullptr;
 

src/llama.cpp

@@ -9271,6 +9271,9 @@ void llama_sampler_dry_free(struct llama_sampler_dry* smpl) {
 }
 
 struct llama_sampler_dry* llama_sampler_dry_clone(struct llama_sampler_dry* smpl) {
+    if (!smpl) {
+        return nullptr;
+    }
     return new llama_sampler_dry {
         /* .total_context_size     = */ smpl->total_context_size,
         /* .dry_multiplier         = */ smpl->dry_multiplier,