Skip to content

fix: three correctness bugs in --gdn-replay - #352

Open
baptisterajaut wants to merge 2 commits into
TheTom:feature/turboquant-kv-cachefrom
baptisterajaut:fix/gdn-replay-correctness
Open

fix: three correctness bugs in --gdn-replay#352
baptisterajaut wants to merge 2 commits into
TheTom:feature/turboquant-kv-cachefrom
baptisterajaut:fix/gdn-replay-correctness

Conversation

@baptisterajaut

Copy link
Copy Markdown

--gdn-replay corrupts 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 are seq_rm (setting it) and full sequence removal (zeroing it). The classic path does a consume-and-clear of rs_idx in s_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 a consume_replay_len() called once per decode at the end of set_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_reuse checks it at src/llama-graph.cpp:396, but llm_graph_input_mem_hybrid::can_reuse (:1134) and its _k and _iswa siblings do not. qwen35 goes through build_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_rm records replay_len instead of calling set_rs_idx (src/llama-memory-recurrent.cpp:239-243), so rs_idx stays 0 for the whole run. But rs_idx is the only rollback group selector, and the conv state keeps its (1 + n_rs_seq) layout: build_conv_state writes 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égien with no space, habveut, 117.) around otherwise coherent text. The wanted depth is the same rollback value seq_rm saw, so build_conv_state now selects the group explicitly.

There is a fourth thing I did not fix, only flagged with a one-shot warning. The else branch of the checkpoint update assumes base_state is still before the retained window, which only holds when n_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 last n_rs_seq steps 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 when gdn_replay is off, so consume_replay_len() early-returns, the conv view is not taken, and the new can_reuse comparison is 0 == 0.

🤖 Generated with Claude Code

@github-actions github-actions Bot added the model label Sep 4, 2026
@giveen

giveen commented Sep 4, 2026

Copy link
Copy Markdown

@baptisterajaut thank you for this, it was on my list todo, got so many side projects.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants