From 8ff2d943a31b3d54440698db41e62ea121d661be Mon Sep 17 00:00:00 2001 From: SamuelOliveirads Date: Fri, 17 Apr 2026 13:08:04 -0300 Subject: [PATCH] server: snapshot recurrent state in tensor --- examples/server/server-context.cpp | 35 +++-- include/llama.h | 6 + src/llama-context.h | 34 ++++ src/llama.cpp | 240 +++++++++++++++++++++++++++-- 4 files changed, 294 insertions(+), 21 deletions(-) diff --git a/examples/server/server-context.cpp b/examples/server/server-context.cpp index 6b0462a70..2a3cfbf89 100644 --- a/examples/server/server-context.cpp +++ b/examples/server/server-context.cpp @@ -3616,8 +3616,13 @@ void server_context::speculative_decoding_accept() { // for hybrid models: if any drafts were rejected, restore recurrent state const bool any_rejected = (ids.size() - 1) < n_draft; + const bool gpu_ckpt = llama_kv_cache_checkpoint_supported(ctx); if (any_rejected && slot.spec_ckpt_valid) { - llama_state_seq_set_data(ctx, slot.spec_ckpt_data.data(), slot.spec_ckpt_data.size(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + if (gpu_ckpt) { + llama_kv_cache_checkpoint_restore(ctx); + } else { + llama_state_seq_set_data(ctx, slot.spec_ckpt_data.data(), slot.spec_ckpt_data.size(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } llama_kv_cache_seq_rm(ctx, slot.id, slot.spec_ckpt_n_past, -1); @@ -3656,15 +3661,20 @@ void server_context::speculative_decoding_accept() { } llama_batch_free(re_batch); - SLT_DBG(slot, "spec checkpoint restored: re-decoded %d accepted tokens (rejected %d)\n", - n_accepted, (int)(n_draft - (ids.size() - 1))); + SLT_DBG(slot, "spec checkpoint restored (gpu=%s): re-decoded %d accepted tokens (rejected %d)\n", + gpu_ckpt ? "yes" : "no", n_accepted, (int)(n_draft - (ids.size() - 1))); } slot.spec_ckpt_valid = false; + if (gpu_ckpt) { + llama_kv_cache_checkpoint_delete(ctx); + } } else { llama_kv_cache_seq_rm(ctx, slot.id, slot.n_past, -1); slot.spec_ckpt_valid = false; - // discard saved sampler on full acceptance + if (gpu_ckpt) { + llama_kv_cache_checkpoint_delete(ctx); + } if (slot.spec_ckpt_sampler) { common_sampler_free(slot.spec_ckpt_sampler); slot.spec_ckpt_sampler = nullptr; @@ -4247,10 +4257,16 @@ void server_context::update_slots() { continue; } slot.spec_ckpt_n_past = slot.n_past - (int32_t)(slot.drafted.size() + 1); - const size_t ckpt_size = llama_state_seq_get_size(ctx, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - slot.spec_ckpt_data.resize(ckpt_size); - const size_t written = llama_state_seq_get_data(ctx, slot.spec_ckpt_data.data(), ckpt_size, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - slot.spec_ckpt_valid = (written > 0); + + if (llama_kv_cache_checkpoint_supported(ctx)) { + slot.spec_ckpt_valid = llama_kv_cache_checkpoint_save(ctx); + } else { + const size_t ckpt_size = llama_state_seq_get_size(ctx, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + slot.spec_ckpt_data.resize(ckpt_size); + const size_t written = llama_state_seq_get_data(ctx, slot.spec_ckpt_data.data(), ckpt_size, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + slot.spec_ckpt_valid = (written > 0); + } + if (slot.spec_ckpt_valid) { // save sampler state so we can restore RNG/grammar on rejection if (slot.spec_ckpt_sampler) { @@ -4258,7 +4274,8 @@ void server_context::update_slots() { } slot.spec_ckpt_sampler = common_sampler_init(model, slot.sparams); common_sampler_clone(slot.ctx_sampling, slot.spec_ckpt_sampler); - SLT_DBG(slot, "spec checkpoint saved: %zu bytes, n_past_pre_spec=%d\n", written, slot.spec_ckpt_n_past); + SLT_DBG(slot, "spec checkpoint saved (gpu=%s), n_past_pre_spec=%d\n", + llama_kv_cache_checkpoint_supported(ctx) ? "yes" : "no", slot.spec_ckpt_n_past); } else { SLT_WRN(slot, "%s", "failed to save spec checkpoint\n"); } diff --git a/include/llama.h b/include/llama.h index 13d05a381..ba282f47f 100644 --- a/include/llama.h +++ b/include/llama.h @@ -798,6 +798,12 @@ extern "C" { LLAMA_API void llama_kv_cache_clear( struct llama_context * ctx); + // GPU-resident checkpoint for speculative decoding (exclusive from hybrid models) + LLAMA_API bool llama_kv_cache_checkpoint_save(struct llama_context * ctx); + LLAMA_API bool llama_kv_cache_checkpoint_restore(struct llama_context * ctx); + LLAMA_API void llama_kv_cache_checkpoint_delete(struct llama_context * ctx); + LLAMA_API bool llama_kv_cache_checkpoint_supported(struct llama_context * ctx); + // Removes all tokens that belong to the specified sequence and have positions in [p0, p1) // Returns false if a partial sequence cannot be removed. Removing a whole sequence never fails // seq_id < 0 : match any sequence diff --git a/src/llama-context.h b/src/llama-context.h index be563122b..dbd04b14a 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -74,6 +74,40 @@ struct llama_kv_cache { return size; } + // GPU-resident checkpoint for speculative decoding (exclusive from hybrid models) + struct gpu_checkpoint { + std::vector cells_snapshot; + uint32_t head_snapshot = 0; + uint32_t used_snapshot = 0; + + std::vector s_l_shadow; + + std::vector> split_s_l_shadow; + + std::vector shadow_ctxs; + std::vector shadow_bufs; + + bool allocated = false; + bool saved = false; + + ~gpu_checkpoint() { + for (struct ggml_context * ctx : shadow_ctxs) { + ggml_free(ctx); + } + for (ggml_backend_buffer_t buf : shadow_bufs) { + ggml_backend_buffer_free(buf); + } + } + }; + + gpu_checkpoint ckpt; + + bool checkpoint_alloc_shadows(); + bool checkpoint_supported() const; + bool checkpoint_save(); + bool checkpoint_restore(); + void checkpoint_delete(); + ~llama_kv_cache() { for (struct ggml_context * ctx : ctxs) { ggml_free(ctx); diff --git a/src/llama.cpp b/src/llama.cpp index de62849bd..2da6ba8a3 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -1204,6 +1204,207 @@ static uint32_t llama_kv_cache_cell_max(const struct llama_kv_cache & cache, uin return 0; } +bool llama_kv_cache::checkpoint_supported() const { + for (const auto * s : s_l) { + if (s != nullptr) { + return true; + } + } + return false; +} + +bool llama_kv_cache::checkpoint_alloc_shadows() { + if (ckpt.allocated) { + return true; + } + + const uint32_t n_layer = (uint32_t)s_l.size(); + ckpt.s_l_shadow.resize(n_layer, nullptr); + + struct tensor_entry { + ggml_tensor * primary; + uint32_t il; + int split_idx; // -1 for non-split + }; + + std::map> buft_entries; + + uint32_t split_s_idx = 0; + for (uint32_t il = 0; il < n_layer; ++il) { + if (s_l[il] == nullptr) { + continue; + } + + auto * extra = s_l[il]->extra; + if (extra != nullptr) { + // Split tensor — shadow each per-device split tensor individually + auto * split_info = (const ggml_split_tensor_t *)extra; + for (int d = 0; d < split_info->n_device; ++d) { + if (split_info->splits[d] == nullptr) continue; + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(split_info->splits[d]->buffer); + buft_entries[buft].push_back({split_info->splits[d], il, d}); + } + split_s_idx++; + } else { + // Non-split tensor + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(s_l[il]->buffer); + buft_entries[buft].push_back({s_l[il], il, -1}); + } + } + + // Allocate shadow tensors grouped by buffer type + for (auto & [buft, entries] : buft_entries) { + ggml_init_params params = { + /*.mem_size =*/ entries.size() * ggml_tensor_overhead(), + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + LLAMA_LOG_ERROR("%s: failed to create ggml context for shadow tensors\n", __func__); + return false; + } + + for (auto & entry : entries) { + ggml_tensor * shadow = ggml_dup_tensor(ctx, entry.primary); + ggml_format_name(shadow, "shadow_s_l%d_d%d", entry.il, entry.split_idx); + entry.primary = shadow; // temporarily store shadow in entry for buffer allocation + } + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft); + if (!buf) { + LLAMA_LOG_ERROR("%s: failed to allocate buffer for shadow tensors\n", __func__); + ggml_free(ctx); + return false; + } + ggml_backend_buffer_clear(buf, 0); + LLAMA_LOG_INFO("%s: %10s shadow buffer size = %8.2f MiB\n", __func__, + ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf) / 1024.0 / 1024.0); + ckpt.shadow_ctxs.push_back(ctx); + ckpt.shadow_bufs.push_back(buf); + } + + // Build shadow lookup: for non-split layers, store directly in s_l_shadow + for (auto & [buft, entries] : buft_entries) { + for (auto & entry : entries) { + if (entry.split_idx == -1) { + ckpt.s_l_shadow[entry.il] = entry.primary; + } + } + } + + // Build shadow split tensors + ckpt.split_s_l_shadow.resize(split_s_l.size()); + split_s_idx = 0; + for (uint32_t il = 0; il < n_layer; ++il) { + if (s_l[il] == nullptr || s_l[il]->extra == nullptr) { + continue; + } + + auto * split_info = (const ggml_split_tensor_t *)s_l[il]->extra; + auto & shadow_split = ckpt.split_s_l_shadow[split_s_idx]; + shadow_split.resize(split_info->n_device, nullptr); + + for (int d = 0; d < split_info->n_device; ++d) { + if (split_info->splits[d] == nullptr) continue; + // Find corresponding shadow tensor from buft_entries + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(split_info->splits[d]->buffer); + for (auto & entry : buft_entries[buft]) { + if (entry.il == il && entry.split_idx == d) { + shadow_split[d] = entry.primary; + break; + } + } + } + split_s_idx++; + } + + ckpt.allocated = true; + return true; +} + +bool llama_kv_cache::checkpoint_save() { + if (!checkpoint_alloc_shadows()) { + return false; + } + + const uint32_t n_layer = (uint32_t)s_l.size(); + + // 1. Snapshot cell metadata (small, host-side) + ckpt.cells_snapshot = cells; + ckpt.head_snapshot = head; + ckpt.used_snapshot = used; + + // 2. Copy tensors: primary -> shadow (GPU-to-GPU when on same device) + uint32_t split_s_idx = 0; + for (uint32_t il = 0; il < n_layer; ++il) { + if (s_l[il] == nullptr) { + continue; + } + + if (s_l[il]->extra != nullptr) { + // Split tensor — copy per-device splits individually + auto * split_info = (const ggml_split_tensor_t *)s_l[il]->extra; + auto & shadow_split = ckpt.split_s_l_shadow[split_s_idx]; + for (int d = 0; d < split_info->n_device; ++d) { + if (split_info->splits[d] && shadow_split[d]) { + ggml_backend_tensor_copy(split_info->splits[d], shadow_split[d]); + } + } + split_s_idx++; + } else { + // Non-split tensor + ggml_backend_tensor_copy(s_l[il], ckpt.s_l_shadow[il]); + } + } + + ckpt.saved = true; + return true; +} + +bool llama_kv_cache::checkpoint_restore() { + if (!ckpt.saved) { + LLAMA_LOG_ERROR("%s: no checkpoint saved\n", __func__); + return false; + } + + const uint32_t n_layer = (uint32_t)s_l.size(); + + // 1. Restore cell metadata + cells = ckpt.cells_snapshot; + head = ckpt.head_snapshot; + used = ckpt.used_snapshot; + + // 2. Copy tensors: shadow -> primary (GPU-to-GPU) + uint32_t split_s_idx = 0; + for (uint32_t il = 0; il < n_layer; ++il) { + if (s_l[il] == nullptr) { + continue; + } + + if (s_l[il]->extra != nullptr) { + // Split tensor — copy per-device splits individually + auto * split_info = (const ggml_split_tensor_t *)s_l[il]->extra; + auto & shadow_split = ckpt.split_s_l_shadow[split_s_idx]; + for (int d = 0; d < split_info->n_device; ++d) { + if (split_info->splits[d] && shadow_split[d]) { + ggml_backend_tensor_copy(shadow_split[d], split_info->splits[d]); + } + } + split_s_idx++; + } else { + // Non-split tensor + ggml_backend_tensor_copy(ckpt.s_l_shadow[il], s_l[il]); + } + } + + return true; +} + +void llama_kv_cache::checkpoint_delete() { + ckpt.saved = false; +} + static void llama_kv_cache_clear(struct llama_kv_cache & cache) { for (int32_t i = 0; i < (int32_t) cache.size; ++i) { cache.cells[i].pos = -1; @@ -6406,6 +6607,22 @@ void llama_kv_cache_clear(struct llama_context * ctx) { llama_kv_cache_clear(ctx->kv_self); } +bool llama_kv_cache_checkpoint_save(struct llama_context * ctx) { + return ctx->kv_self.checkpoint_save(); +} + +bool llama_kv_cache_checkpoint_restore(struct llama_context * ctx) { + return ctx->kv_self.checkpoint_restore(); +} + +void llama_kv_cache_checkpoint_delete(struct llama_context * ctx) { + ctx->kv_self.checkpoint_delete(); +} + +bool llama_kv_cache_checkpoint_supported(struct llama_context * ctx) { + return ctx->kv_self.checkpoint_supported(); +} + bool llama_kv_cache_seq_rm(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) { return llama_kv_cache_seq_rm(ctx->kv_self, seq_id, p0, p1); } @@ -9001,18 +9218,17 @@ void llama_sampler_dry_free(struct llama_sampler_dry* smpl) { } struct llama_sampler_dry* llama_sampler_dry_clone(struct llama_sampler_dry* smpl) { - // nullptr is passed as vocab because it is only needed for raw sequence breaker processing, which we have already done and will be copying - auto* result = llama_sampler_init_dry(nullptr, smpl->dry_multiplier, smpl->dry_base, smpl->dry_allowed_length, smpl->dry_penalty_last_n, NULL, 0); - // Copy the state, including the processed breakers - { - auto* result_ctx = smpl; - result_ctx->dry_processed_breakers = smpl->dry_processed_breakers; - result_ctx->dry_repeat_count = smpl->dry_repeat_count; - result_ctx->dry_max_token_repeat = smpl->dry_max_token_repeat; - result_ctx->last_tokens = smpl->last_tokens; - } - - return result; + return new llama_sampler_dry { + /* .total_context_size = */ smpl->total_context_size, + /* .dry_multiplier = */ smpl->dry_multiplier, + /* .dry_base = */ smpl->dry_base, + /* .dry_allowed_length = */ smpl->dry_allowed_length, + /* .dry_penalty_last_n = */ smpl->dry_penalty_last_n, + /* .dry_processed_breakers = */ smpl->dry_processed_breakers, + /* .dry_repeat_count = */ smpl->dry_repeat_count, + /* .dry_max_token_repeat = */ smpl->dry_max_token_repeat, + /* .last_tokens = */ smpl->last_tokens, + }; } void llama_sampler_dry_accept(struct llama_sampler_dry* smpl, llama_token token) {