From 8b2e18efdd5d22b27da16d46b192b82445604a8d Mon Sep 17 00:00:00 2001 From: Mike Lothian Date: Sun, 27 Sep 2026 15:06:40 +0100 Subject: [PATCH 06/18] server : keep speculative checkpoints on device Assisted-by: Claude Code (Claude Opus 5.5) --- tools/server/server-context.cpp | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index e95fb63ab..9aeb4a74d 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -255,6 +255,8 @@ struct server_slot { llama_tokens spec_prompt; std::vector spec_i_batch; common_prompt_checkpoint spec_ckpt; + // one speculative checkpoint per slot, so its tensor data can stay on the device + static constexpr llama_state_seq_flags spec_ckpt_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE; bool spec_is_replay = false; std::mt19937 spec_synth_rng; @@ -3020,7 +3022,7 @@ private: 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, slot.spec_ckpt_flags); } slot.spec_prompt = slot.prompt.tokens.get_text_tokens(); @@ -3059,7 +3061,7 @@ private: 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, slot.spec_ckpt_flags); } if (!llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, ckpt.pos_max + 1, -1)) { @@ -3078,7 +3080,7 @@ private: 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, slot.spec_ckpt_flags); //const int64_t t_total = ggml_time_us() - t_start; //printf("checkpoint total: %f ms\n", t_total / 1000.0); @@ -3090,7 +3092,7 @@ private: } if (use_ckpt_dft) { - ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_dft(ctx_dft, slot.id, slot.spec_ckpt_flags); } } }); @@ -3942,10 +3944,10 @@ private: 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, slot.spec_ckpt_flags); 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, slot.spec_ckpt_flags); } slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1); -- 2.55.0