From 611953af315599579c023003baa4b12e9f4584c4 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sat, 22 Aug 2026 08:06:54 +0000 Subject: [PATCH] server : do not re-verify replayed draft tokens after a checkpoint restore With full-checkpoint rollback, a partial draft acceptance restores the pre-round state and re-decodes the accepted tokens to rebuild it. The replay went through the same verification as a fresh draft. On backends where logits change with batch shape or memory layout (Vulkan), that re-verification can reject a token the original verification accepted; the rejection restores the same checkpoint and replays again, and the slot loops on one position without emitting anything. qwen35moe with --spec-type draft-mtp stalled this way a few hundred tokens into long generations (the v0.6.8 MTP hang): the loop repeated "accepted 2/3, restore at pos 995" every 27 ms with the GPU at 90 percent. Accept the replayed tokens without re-verifying and sample only the continuation from the final position. The replayed prefix was accepted by the verification that triggered the restore; the replay exists to rebuild state. On backends with batch-shape invariant logits the re-verification always agreed, so behavior there is unchanged (verified bit-identical on CPU with and without this change, 800-token greedy pair, 118 restore rounds). Assisted-by: Claude Fable 5 (cherry picked from commit 9c5d899ff7966179f56e49edd5c7a57f7b6172e6) --- tools/server/server-context.cpp | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f5477356d61d..5ed439417f13 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3884,12 +3884,25 @@ struct server_context_impl { common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get())); GGML_ASSERT(slot.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) - : server_sample_and_accept_synth( - slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft, - synth_probs, slot.spec_synth_rng, slot.spec_is_replay); + std::vector accepted; + if (slot.spec_is_replay) { + // replayed tokens were accepted before the restore; re-verifying them can + // disagree when logits depend on batch shape, and each disagreement restores + // the same checkpoint again - the slot stops making progress + accepted = slot.spec_draft; + for (const llama_token id : accepted) { + common_sampler_accept(slot.smpl.get(), id, true); + } + accepted.push_back(common_sampler_sample(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch.back())); + common_sampler_accept(slot.smpl.get(), accepted.back(), true); + } else { + const auto & synth_probs = common_speculative_get_synth_probs(spec.get()); + accepted = synth_probs.empty() + ? common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft) + : server_sample_and_accept_synth( + slot.smpl.get(), slot.ctx_tgt, slot.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);