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?