memory : keep a pending rollback across seq_cp, and default MTP draft depth to 2 on CUDA

[bri-prism]

Overview

Two changes for hybrid recurrent/attention models that use speculative decoding with an MTP head.

1. seq_cp loses a pending partial rollback (bug fix)

With n_rs_seq != 0, a partial seq_rm inside the snapshot window does not change the recurrent state right away. It records the rollback in rs_idx[seq], and the next decode reads the matching snapshot row. llama_memory_recurrent::seq_cp shares the cell with the destination but does not copy rs_idx. A sequence forked right after a rejected speculative round then continues from the state before the rollback. The source sequence is not affected.

• llama_memory_recurrent::seq_cp: copy rs_idx[src] to rs_idx[dst] (one line).
• llama_kv_cache_dsv4::seq_cp has the same problem in a different form: it copies the current state plane of the source and resets rs_idx[dst]. It now copies the snapshot plane that the pending rollback points to. The compressor state exposes views for all planes so that apply_copies can copy from a snapshot plane. Both fixes are in one commit.

The server reaches this through copy_state_to (child tasks with n > 1), and so does any client that forks a sequence between speculative rounds.

2. Default MTP draft depth on CUDA

--spec-draft-n-max defaults to 3 for every draft type. With --spec-type draft-mtp on CUDA, a depth of 3 is clearly slower than 1 or 2: the verify batch cost grows faster with depth than on Metal. On Metal, 3 is still the fastest of 1, 2 and 3. When draft-mtp is enabled, the main model offloads to a CUDA device and --spec-draft-n-max is not given, the default is now 2. An explicit --spec-draft-n-max (flag or LLAMA_ARG_SPEC_DRAFT_N_MAX) always wins, and nothing changes for other backends or other draft types.

Additional information

Test for change 1

The new case test_seq_cp_pending_rollback goes into the existing test-recurrent-state-rollback. It does not add a new test file. It builds two contexts:

• fork: decode a prompt on seq 0, roll back 3 tokens with seq_rm, seq_cp(0, 1), decode one token on seq 1
• reference: the same prompt on seq 1, roll back 3 tokens, decode the same token on seq 1

It then compares the serialized recurrent state of seq 1 (LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) byte for byte. The batch shapes match, so a correct implementation is bit-exact.

It compares state bytes, not logits, on purpose. With the N(0, 0.01) weights from test-llama-archs, the recurrent contribution is below float resolution in the logits. Without the fix, the logits of the fork still matched the reference exactly, while its recurrent state was off by an NMSE of about 1.4. This also means the existing logit comparisons in this test cannot detect a wrong snapshot row.

Results, on the tiny generated models that this test already runs on (gated delta net hybrid dense and MoE, Nemotron-H, DeepSeek V4):

build	with both fixes	recurrent fix removed	DSV4 fix removed
CPU	4/4 pass	3 recurrent models fail, DSV4 passes	DSV4 fails, others pass
Metal	4/4 pass	all 4 fail (both poisoned together)
CUDA (H200)	4/4 pass	all 4 fail (both poisoned together)

The memory buffers were on the GPU in the Metal and CUDA runs (checked with -lv 4). The existing cases in the test still pass.

Note: on the current base, test-generate-models (test-llama-archs -o) stops with an error at the dspark arch, before it writes the Nemotron-H and DeepSeek V4 models. This is not caused by this PR. I generated those two with test-llama-archs -a <arch> -o <dir>.

Measurements for change 2

llama-server, --spec-type draft-mtp --spec-draft-n-min 0 --spec-draft-p-min 0, greedy, -np 1, -fa on, cache_prompt: false. A quantized hybrid recurrent/attention model with an MTP head. 24 prose prompts (256 tokens, fixed length) and 64 everyday task prompts (up to 512 tokens, stop at EOS). One server start per arm, ABBA order (plain, 1, 2, 3, 3, 2, 1, plain). Throughput is the sum of generated tokens over the sum of server decode time. The values are speedups over plain decode on the same server build:

backend	depth 1	depth 2	depth 3
H200, all prompts	1.33x	1.34x	1.14x
H200, prose	1.23x	1.23x	1.03x
H200, tasks	1.38x	1.39x	1.20x
M5 Metal, all prompts	1.20x	1.27x	1.31x
M5 Metal, prose	1.16x	1.18x	1.19x
M5 Metal, tasks	1.22x	1.31x	1.37x

A second H200 run gave the same ranking (1.33x / 1.32x / 1.15x). A separate standalone probe on consumer CUDA cards (RTX 4090, RTX PRO 4500 Blackwell) also found depth 2 at or above depth 1, and depth 3 clearly lower. Depth 2 is the safe choice across these CUDA cards. On Metal, 3 stays the default.

With this change on H200, --spec-type draft-mtp without --spec-draft-n-max gives the same tokens as an explicit depth of 2 and the same speed (1.33x vs 1.15x for the old default of 3). The new default applies only in that case. It is not used with an explicit --spec-draft-n-max (flag or env), with -ngl 0 or -dev none, or on Metal.

A confidence stop (--spec-draft-p-min 0.4 to 0.8 with n_max 2 or 3) did not beat a fixed depth on CUDA. It raised acceptance to 0.9 or more, but throughput fell to 0.94x to 1.05x of plain. Most of that loss is CUDA graph reuse: the verify batch size changes from round to round. With GGML_CUDA_DISABLE_GRAPHS=1, p-min 0.6 with n_max 2 recovered to 1.11x (prose) and 1.28x (tasks), but it was still below a fixed depth of 1 or 2 in the same run. So this PR does not change the p_min default.