fix: three correctness bugs in --gdn-replay - #352
Open
baptisterajaut wants to merge 2 commits into
Open
Conversation
|
@baptisterajaut thank you for this, it was on my list todo, got so many side projects. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
--gdn-replaycorrupts generation on qwen35 (tested on RVN-IQ4_NL-multilingual-mtp, 27B, sm_86). Three independent bugs, all in the DRC phase 2 path. With the three fixed, greedy decoding is byte-identical to a run without the flag over 300 tokens, during which 117 draft tokens were rejected, so the replay path was exercised continuously rather than sitting idle.replay_len is never consumed.
get_replay_len()reads it and the only writers areseq_rm(setting it) and full sequence removal (zeroing it). The classic path does a consume-and-clear ofrs_idxins_copy_idx(), src/llama-memory-recurrent.cpp:1397. Without the equivalent, the first partial rejection latches a rollback that is then re-applied on every later decode, so the recurrent state permanently trails the token stream. Fixed with aconsume_replay_len()called once per decode at the end ofset_input, after every GDN layer has read the value during graph build.The hybrid can_reuse variants ignore replay_len.
llm_graph_input_rs::can_reusechecks it at src/llama-graph.cpp:396, butllm_graph_input_mem_hybrid::can_reuse(:1134) and its_kand_iswasiblings do not. qwen35 goes throughbuild_inp_mem_hybrid(), so the guard never fired for it, and a changed replay length means a differently shaped reconstruction subtree that reused topology cannot express. Upstream keeps the two bodies in sync, the divergence came in with the replay_len check.The conv state is never rolled back.
seq_rmrecordsreplay_leninstead of callingset_rs_idx(src/llama-memory-recurrent.cpp:239-243), sors_idxstays 0 for the whole run. Butrs_idxis the only rollback group selector, and the conv state keeps its(1 + n_rs_seq)layout:build_conv_statewrites all K snapshots every decode (src/models/delta-net-base.cpp:502-520,s_slot = K - t) and then always reads back group 0, the optimistic one. So the recurrent state was correctly rewound while the convolution window still held the rejected draft tokens, which is what produced the short range damage in the output (LeNorvégienwith no space,habveut,117.) around otherwise coherent text. The wanted depth is the samerollbackvalueseq_rmsaw, sobuild_conv_statenow selects the group explicitly.There is a fourth thing I did not fix, only flagged with a one-shot warning. The
elsebranch of the checkpoint update assumesbase_stateis still before the retained window, which only holds whenn_seq_tokens == n_rs_seq. For a strictly shorter batch the checkpoint lands inside the uncertain window, and since the ingredient ring only retains the lastn_rs_seqsteps there is no way to recover the true one. It looks unreachable with the current verify batch shape (n_draft + 1 > n_draft) and the warning never fired in my runs, so I left the behaviour alone.On the cost, in case you want it in the flag's help text: measured on a 3090 with the einstein logic prompt, 1200 tokens, fixed seed,
--spec-draft-n-max 3, acceptance is identical with and without the flag (0.60976 both, and the greedy output matches byte for byte), throughput goes from 64.0 to 57.1 t/s, and VRAM drops by 108 MiB. The replay is exact but costs about 11 percent, and the saving is on per sequence SSM state so it does not grow with context length.The default path should be untouched:
get_replay_len()returns 0 whengdn_replayis off, soconsume_replay_len()early-returns, the conv view is not taken, and the new can_reuse comparison is 0 == 0.🤖 Generated with Claude Code