server : do not re-verify replayed draft tokens after a checkpoint restore - #28061
server : do not re-verify replayed draft tokens after a checkpoint restore#28061JayToltTech wants to merge 1 commit into
Conversation
…store 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 9c5d899)
|
Hi @JayToltTech, thanks for your contribution! Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:
Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below. |
|
Don't submit other people's work, unless they asked you to or explicitely allowed it. |
|
Understood, and you are right — I should have asked first. Crediting is not consent, and that was my mistake, not a grey area. Apologies for the noise. @Nathanw1014 — the fix is yours ( Leaving this closed. The bug report at #28060 stays open on its own merits, and I have converted #28104 to a draft pending your permission. |
Fixes #28060.
What
With full-checkpoint rollback, a partial draft acceptance restores the pre-round state and re-decodes the accepted tokens to rebuild it. That replay currently runs through the same verification as a fresh draft. On backends where logits depend on batch shape or memory layout (Vulkan), the 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.
This accepts the replayed tokens without re-verifying, and samples only the continuation from the final position. The replayed prefix was already accepted by the verification that triggered the restore — the replay exists to rebuild state, not to re-decide.
Credit
The diagnosis and the patch are @Nathanw1014's, from
strix-halo-vulkancommit9c5d899ff, preserved here with authorship intact. I am submitting it because I hit the surrounding area while porting MTP onto current master and thought the fix should not stay on a fork. @Nathanw1014 — happy to close this if you would rather submit it yourself.Symptom independently reported by @flobob45 in #27836.
Conflict resolution
The commit does not apply to current master unchanged: master has since added the
synth_probssampling path, and the original also carried aspec_distsbranch that master does not have (spec_distshas no references intools/server/;spec_is_replayhas 6). I kept master'ssynth_probspath, grafted on only the replay short-circuit, and dropped thespec_distsbranch. That resolution is mine and is the part most worth reviewing.Testing
9723942(Vulkan,-DGGML_VULKAN=ON).Not verified by me: I have not reproduced the livelock on unmodified master — see the issue for why. The original author reports the change is bit-identical on CPU (800-token greedy pair, 118 restore rounds), which matches the expectation that only batch-shape-sensitive backends are affected. Reviewers with a CUDA or Metal setup that exercises checkpoint rollback may want to confirm no behaviour change there.