fix: three correctness bugs in --gdn-replay

[baptisterajaut]

--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

[giveen]

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

[TheTom]

Reviewed against base. All three diagnoses check out: replay_len had no consumer (only seq_rm wrote it and full removal zeroed it), only llm_graph_input_rs::can_reuse looked at it while the three hybrid input classes did not, and build_conv_state never selected the rolled-back group. The fixes match the patterns you cite (s_copy()'s consume-and-clear of rs_idx), and I traced every added path with --gdn-replay off: consume_replay_len() early-returns, the can_reuse additions compare 0 to 0, and conv_rollback == 0 leaves the graph byte-identical. So the default configuration and the whole backend matrix are untouched. Merging.

Two follow-ups, neither blocking this:

1. get_replay_len() max-reduces over the ubatch's sequences, so one rollback depth is applied to every sequence in a split_equal ubatch. Pre-existing from GDN ingredient-replay: cut MTP rollback VRAM  #331 (its header comment hedges with "single, in the n_seq_max==1 case"), but this PR is what makes the flag usable, so it should either document the single-slot restriction or assert n_seq_max == 1 when --gdn-replay is set. A quick -np 2 run would show whether it does damage today.
2. tests/test-recurrent-state-rollback.cpp already does seq_rm followed by three decode_one() calls, which is exactly the shape bug 1 broke. Running it with the flag on and asserting decodes two and three match the checkpoint-restored context would make this a permanent gate.

Would you mind posting your degeneracy numbers (unique/words, gzip ratio, repeated 3-grams) on #345 with the flag on at this head, so the reporter can confirm against their own repro before that issue is closed?

Nit: the file's tag convention is [TAG_...]; the CBA: prefix on the new comments is unexplained.