From 82bacc5475f08bbe7c879ffe54d26e1fe01c8eb0 Mon Sep 17 00:00:00 2001 From: Petr Vasilev Date: Mon, 31 Aug 2026 22:41:43 +0300 Subject: [PATCH] server : keep speculative recurrent-state checkpoints on-device MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For recurrent/hybrid models whose target context is SEQ_RM_TYPE_FULL (e.g. Gated DeltaNet / Mamba hybrids, Qwen3-Next / qwen4exp), speculative decoding must checkpoint and restore the full recurrent state every round. The server took those per-round snapshots with LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY, which serializes the whole state to host — on AMD Strix Halo (gfx1151) this was ~600 ms of each ~825 ms round, a constant ~73% overhead that makes speculative decoding a net loss despite high draft acceptance. OR in LLAMA_STATE_SEQ_FLAGS_ON_DEVICE at the eight live spec_ckpt calls (update_tgt/update_dft/load_tgt/load_dft) so the transient speculative snapshots stay on-device. The prompt-history checkpoint (cur.update_* near update_pos) and the disk/prompt-cache path (prompt_save / llama_state_seq_save_file) keep host serialization, since ON_DEVICE buffers are transient and host-inaccessible. Measured on 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: no draft: 32.4 t/s spec before (host checkpoint): 6.2 t/s (5x LOSS) spec after (this change), n=3: 41.5 t/s (+28%, accept 0.79) spec after, code, n=6: 56.6 t/s (accept 0.91) Greedy output stays equivalent to no-draft (both diverge only through the backend's own non-deterministic reductions). Note: llama-memory-recurrent hard-aborts if cell_ranges.size() > 1 under ON_DEVICE; -np 1 yields one contiguous range, but fragmentation / cache-wrap paths should be guarded or fall back to host checkpointing. --- tools/server/server-context.cpp | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f5477356d61d..ea96cddb1047 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2993,7 +2993,7 @@ struct server_context_impl { llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id)); if (use_ckpt_dft) { - slot.spec_ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + slot.spec_ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); } slot.spec_prompt = slot.prompt.tokens.get_text_tokens(); @@ -3032,7 +3032,7 @@ struct server_context_impl { if (ctx_dft) { if (use_ckpt_dft) { - ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); } if (!llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, ckpt.pos_max + 1, -1)) { @@ -3051,7 +3051,7 @@ struct server_context_impl { if (use_ckpt_tgt) { //const int64_t t_start = ggml_time_us(); - ckpt.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); //const int64_t t_total = ggml_time_us() - t_start; //printf("checkpoint total: %f ms\n", t_total / 1000.0); @@ -3063,7 +3063,7 @@ struct server_context_impl { } if (use_ckpt_dft) { - ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); } } }); @@ -3339,8 +3339,8 @@ struct server_context_impl { if (!do_reset) { // restore the context checkpoint - it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); + it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); // restore the draft's speculative state common_speculative_set_state(spec.get(), slot.id, it->data_spec); @@ -3915,10 +3915,10 @@ struct server_context_impl { SLT_DBG(slot, "restoring speculative checkpoint (pos_min = %d, pos_max = %d, size = %zu)\n", ckpt.pos_min, ckpt.pos_max, ckpt.size()); - ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); if (slot.ctx_dft) { - ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); } slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1);