server : keep speculative recurrent-state checkpoints on-device

[vahpetr]

Summary

For recurrent / hybrid models whose target context is SEQ_RM_TYPE_FULL (Gated DeltaNet / Mamba-style hybrids, e.g. Qwen3-Next / qwen4exp), speculative decoding must checkpoint and restore the full recurrent state every round. The server takes those per-round snapshots with LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY, which serializes the whole state to host memory.

On AMD Strix Halo (gfx1151) this host round-trip is ~600 ms of each ~825 ms round — a constant ~73 % overhead that makes speculative decoding a net loss despite high draft acceptance.

Change

OR in LLAMA_STATE_SEQ_FLAGS_ON_DEVICE at the 8 live spec_ckpt calls (update_tgt / update_dft / load_tgt / load_dft) so the transient per-round speculative snapshots stay on-device. Left host-backed on purpose:

• the prompt-history checkpoint (cur.update_* next to update_pos) — a longer-lived snapshot that a later one would invalidate, and
• the disk / prompt-cache path (prompt_save, llama_state_seq_save_file) — needs host-accessible bytes.

8 lines, no API change; the library already supports the flag.

Measurements

AMD Strix Halo gfx1151 / Vulkan + RADV, Qwen3.8-Flash-Next Q4_K_M, -ctk q8_0 -ctv q8_0, temp 0, -np 1, MTP draft head:

config	decode t/s	draft acceptance
no draft	32.4	—
spec before (host checkpoint)	6.2	0.54
spec after (this PR), n-max 3, p-min 0.7	41.5	0.79
spec after, code (high-accept), n-max 6	56.6	0.91

Greedy output stays equivalent to no-draft (both diverge only through the backend's own non-deterministic parallel reductions — two no-draft runs already differ at temp 0).

Caveat for review

llama-memory-recurrent currently hard-aborts if cell_ranges.size() > 1 under ON_DEVICE. -np 1 yields one contiguous range, but fragmentation / recurrent-cache-wrap / truncation should be guarded (a controlled error, or a fall back to host checkpointing) rather than aborting. Happy to add a guard + a small regression test on the update→load invalidation semantics if maintainers prefer.

Root cause was first diagnosed by @JayToltTech in the qwen4exp MTP thread (#27836); this PR is the standalone server change (the touched code is already in master) plus independent gfx1151 measurements.

[ggml-gh-bot]

Hi @vahpetr, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

• AI-generated content: While code is allowed to be generated by AI, please write the PR description and commit messages on your own without the help of AI.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

[ggerganov]

A better fix would be to create such draft models with n_rs_seq > 1 to support rollback. This would avoid checkpointing all together.

[vahpetr]

Thanks, you're right. qwen3next and qwen4exp just weren't in
llm_arch_supports_rs_rollback, so they were falling back to the full
checkpoint. Adding them lets the recurrent state roll back natively, no
checkpointing needed.

Measured on a Strix Halo (gfx1151, Vulkan), MTP draft, temp 0, -np 1:

	decode
no draft	32 t/s
MTP, full checkpoint	6 t/s
MTP, on-device checkpoint (this PR)	41 t/s
MTP, rollback (#28120)	44 t/s

Output stays the same as the no-draft run. Opened #28120 for the 2-line change.
Should I close this one, or is the on-device checkpoint still worth keeping as a
fallback for recurrent arches that can't do rollback?

[ServeurpersoCom]

Opened #28123 for the qwen4exp side of this, it removes the checkpoint instead of moving it on device. MTP head and draft patch here: https://huggingface.co/dzannotti/Qwen3.8-Flash-Next-MTP-GGUF/

Your change is still worth pursuing for the recurrent architectures that have no rollback support yet, and for the spec types that leave n_rs_seq at zero, where the checkpoint stays on the hot path. The guard you mention in the caveat looks like the blocker to me though: llama-memory-recurrent.cpp:824 aborts the process rather than returning an error, so as written the flag turns a slowdown into a crash as soon as a sequence has more than one cell range.

[vahpetr]

This GGML_ABORT error in state_write (device state, range >1 cell) is still present in the main branch and causes a crash specifically in combination with the on-device-usage from this PR when -np > 1. I replaced it locally with a throw + host-fallback at the server-side checkpoint.
Is this the right approach?

[Spionkiller01]

state_seq_set_data abort on the second request (CUDA, 3-GPU split, qwen4exp)

Another CUDA data point, plus a caveat that I think explains why the speedup is hard to keep in practice.

Box: RTX 4070 Ti Super (sm_89) + Tesla V100 (sm_70) + Tesla P40 (sm_61), CUDA 12.6, driver 580.173, Ubuntu 24.04. Built unslothai#144 @ b76199698 (b10722) with this PR's 8-line change applied on top.

Target: unsloth/Qwen3.8-Flash-Next-UD-IQ4_XS + mmproj-F16, -c 131072 -np 1 -ctk q8_0 -ctv q8_0 -t 8, auto-fit (~68 GB across the three GPUs; the fitter leaves per_layer_token_embd, 28.8 GB, CPU-resident). Draft: the official standalone mtp-Qwen3.8-Flash-Next-Q8_0.gguf, --spec-type draft-mtp --spec-draft-n-max 2 -devd CUDA1.

Symptom: the first request succeeds, the second aborts.

E state_seq_set_data: error loading state: failed to restore kv cache
  ggml_abort()
  #5 llama_context::state_seq_set_data(int, unsigned char const*, unsigned long, unsigned int)

This reproduces with plain text requests, no image involved, and with -np 1 throughout — so it is not the cell_ranges.size() > 1 fragmentation case from your "Caveat for review". Reverting only this PR's 8 lines and rebuilding makes the identical sequence stable (text -> image -> text, repeated, all 200).

Why it may have been missed so far: it only triggers when the prompt cache actually restores. My throughput runs used /completion with cache_prompt: false and never hit it — three sequential requests, no abort — and that is exactly where the +33% figure below comes from. Every request through /v1/chat/completions, where prompt caching is on by default, aborts on the second call.

Numbers on this box (5 sequential /v1/chat/completions, 400 max tokens, temp 0.3):

arm	decode t/s	stable
no draft	29.6	yes
draft-mtp, this PR reverted	25.1	yes
draft-mtp + this PR	~40 (only measurable with cache_prompt: false)	no — aborts on request 2

Draft acceptance is 0.82 in both draft arms, so acceptance is not the problem here.

Two notes towards a fix:

• The root cause looks like Misc. bug: llama_state_seq_set_data_ext: invalid ON_DEVICE state can throw across the C API or abort #27439: on a failed restore the ON_DEVICE path aborts instead of returning the documented 0 = failed to load, so the server has no way to fall back. A guard that returns the error and lets the server fall back to a host checkpoint would make this PR usable here, which matches the guard you already offered in the caveat section.
• Without the on-device checkpoints, draft-mtp is a net loss on this hardware: 25.1 vs 29.6 t/s. The host round-trip is 112 MiB of recurrent state per round over PCIe across three cards, which matches your Strix Halo finding in direction, if not in magnitude.

Happy to rerun anything — the server target rebuilds in about 10 seconds here.