From f8668dd354c829a92e0e4d58a6c4eaec64ec8ddd 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 | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 7d636c7923b8..81dac470622e 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -4132,11 +4132,24 @@ struct server_context_impl { 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( + std::vector accepted; + if (!synth_probs.empty()) { + accepted = 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); + } else 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 { + accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); + } slot.spec_i_batch.clear(); GGML_ASSERT(accepted.size() >= 1);