diff --git a/examples/server/server-context.cpp b/examples/server/server-context.cpp index a68565e5..b4bc3a28 100644 --- a/examples/server/server-context.cpp +++ b/examples/server/server-context.cpp @@ -18,23 +18,17 @@ #include #include -static void server_prompt_checkpoint_update(server_prompt_checkpoint & ckpt, llama_context * ctx, int id, int64_t n_tokens, llama_pos pos_min = -1, llama_pos pos_max = -1, int32_t offset = 0) { - if (pos_min == -1) { - pos_min = llama_kv_cache_seq_pos_min(ctx, id); - } - if (pos_max == -1) { - pos_max = llama_kv_cache_seq_pos_max(ctx, id); - } - const size_t checkpoint_size = llama_state_seq_get_size(ctx, id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - +static void server_prompt_checkpoint_update(server_prompt_checkpoint & ckpt, llama_context * ctx, int id, int64_t n_tokens, llama_pos pos_min, llama_pos pos_max, int32_t offset) { ckpt.pos_min = pos_min; ckpt.pos_max = pos_max; ckpt.pos_max_prompt = pos_max + offset; ckpt.pos_min_prompt = pos_min + offset; ckpt.n_tokens = n_tokens; + + const size_t checkpoint_size = llama_state_seq_get_size(ctx, id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); ckpt.data.resize(checkpoint_size); - const size_t n = llama_state_seq_get_data(ctx, ckpt.data.data(), checkpoint_size, id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + const size_t n = llama_state_seq_get_data(ctx, ckpt.data.data(), ckpt.data.size(), id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); if (n != checkpoint_size) { GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", checkpoint_size, n); } @@ -505,7 +499,7 @@ void server_slot::reset() { rewind_status = false; generated_token_probs.clear(); - checkpoint_pos = 0; + checkpoint_pos = -1; image_just_processed = false; do_checkpoint = false; if (spec != nullptr) { @@ -1788,14 +1782,14 @@ bool server_context::launch_slot_with_task(server_slot& slot, server_task& task) } while (false); slot.allow_ruless_prev = slot.allow_ruless; - if (llama_model_has_recurrent(llama_get_model(slot.ctx))) { + if (llama_model_has_recurrent(llama_get_model(slot.ctx)) || llama_model_is_deepseek4(llama_get_model(slot.ctx))) { params_base.can_ban_phrases = false; bool do_checkpoint = params_base.ctx_checkpoints_n > 0; // make checkpoints only for completion tasks do_checkpoint = do_checkpoint && task.type == SERVER_TASK_TYPE_COMPLETION; // make a checkpoint of the parts of the memory that cannot be rolled back. // checkpoints are created only if: - // - the model architecture is marked as recurrent or hybrid + // - the model architecture is marked as recurrent or hybrid, or has private per-position state (DSV4) // // TODO: try to make this conditional on the context or the memory module, instead of the model type params_base.do_checkpoint = do_checkpoint; @@ -2103,7 +2097,7 @@ bool server_context::system_prompt_set(const std::string& sys_prompt) { slot.n_kept_prompt = 0; slot.n_prompt_tokens_cache = 0; slot.server_cached_prompt.checkpoints.clear(); - slot.checkpoint_pos = 0; + slot.checkpoint_pos = -1; slot.do_checkpoint = false; if (slot.ctx_sampling != nullptr) { common_sampler_reset(slot.ctx_sampling); @@ -2721,7 +2715,7 @@ static size_t load_checkpoints_from_file(const std::string & filename, std::list size_t count; file.read(reinterpret_cast(&count), sizeof(count)); - for (int i = 0; i < count; i++) { + for (size_t i = 0; i < count; i++) { server_prompt_checkpoint checkpoint; file.read(reinterpret_cast(&checkpoint.pos_min), sizeof(checkpoint.pos_min)); file.read(reinterpret_cast(&checkpoint.pos_max), sizeof(checkpoint.pos_max)); @@ -3589,14 +3583,35 @@ void server_context::add_sampled_tokens() { } } -void server_context::create_checkpoint_at_interval(server_slot & slot, const gpt_params & params_base) { - if (params_base.do_checkpoint && params_base.ctx_checkpoints_interval > 0) { - auto pos = llama_kv_cache_seq_pos_max(slot.ctx, slot.id); - if (slot.checkpoint_pos + params_base.ctx_checkpoints_interval <= 1 + pos) { - bool created = create_checkpoint(slot); - if (created) { - slot.checkpoint_pos = pos; - } +// Verifies that a restored checkpoint reaches the expected cache position. +// Logs via SLT_ERR with an optional label prefix (e.g. "DSV4 "). +// Returns true if the position matches. +static bool verify_restored_checkpoint( + const server_prompt_checkpoint & ckpt, + server_slot & slot, + const char * label) +{ + const llama_pos check_pos = llama_kv_cache_seq_pos_max(slot.ctx, slot.id); + if (check_pos != ckpt.pos_max) { + SLT_ERR(slot, "%srestore position mismatch: cache pos_max=%d != checkpoint pos_max=%d — state corrupted\n", + label, check_pos, ckpt.pos_max); + return false; + } + return true; +} + +void server_context::create_checkpoint_at_interval(server_slot & slot) { + if (!this->params_base.do_checkpoint) { + return; + } + if (this->params_base.ctx_checkpoints_interval <= 0) { + return; + } + auto pos = llama_kv_cache_seq_pos_max(slot.ctx, slot.id); + if (slot.checkpoint_pos + this->params_base.ctx_checkpoints_interval <= pos) { + bool created = create_checkpoint(slot); + if (created) { + slot.checkpoint_pos = pos; } } } @@ -3604,10 +3619,12 @@ void server_context::create_checkpoint_at_interval(server_slot & slot, const gp void server_context::apply_checkpoint(server_slot & slot) { llama_pos pos_next = slot.cache_tokens.pos_next(slot.n_past); const auto pos_min_thold = std::max(0, pos_next - 1); + const bool is_dsv4 = llama_model_is_deepseek4(model); if (slot.n_past > 0 && slot.n_past < slot.cache_tokens.n_tokens()) { int32_t pos_min = llama_kv_cache_seq_pos_min(slot.ctx, slot.id); - if (pos_min >= pos_min_thold) { + // DSV4 has pos_min=0 (no eviction) so the guard always blocks it + if (pos_min >= pos_min_thold || is_dsv4) { SLT_WRN(slot, "n_past = %d, slot.prompt.tokens.size() = %d, seq_id = %d, pos_min = %d\n", slot.n_past, (int)slot.cache_tokens.size(), slot.id, pos_min); // search for a context checkpoint @@ -3615,7 +3632,7 @@ void server_context::apply_checkpoint(server_slot & slot) { slot.server_cached_prompt.checkpoints.rbegin(), slot.server_cached_prompt.checkpoints.rend(), [&](const auto & cur) { - return cur.pos_min < pos_min_thold || cur.pos_min == 0; + return cur.pos_max < (is_dsv4 ? pos_next : pos_min_thold); } ); @@ -3631,20 +3648,38 @@ void server_context::apply_checkpoint(server_slot & slot) { SLT_ERR(slot, "failed to restore context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n", it->pos_min, it->pos_max, it->n_tokens, (float)checkpoint_size / 1024 / 1024); do_reset = true; //printf("[DEBUG] `do_reset` was set to `true` after failing to restore a checkpoint"); - } else { - pos_next = std::min(pos_next, std::max(it->pos_min + 1, it->pos_max)); + } else if (!verify_restored_checkpoint(*it, slot, "")) { + do_reset = true; + } + + if (!do_reset) { + if (is_dsv4) { + pos_next = std::min(pos_next, it->pos_max + 1); + } else { + pos_next = std::min(pos_next, std::max(it->pos_min + 1, it->pos_max)); + } slot.n_past = slot.cache_tokens.size_up_to_pos(pos_next); - pos_next = slot.prompt_tokens.pos_next(slot.n_past_prompt); - pos_next = std::min(pos_next, std::max(it->pos_min_prompt + 1, it->pos_max_prompt)); - slot.n_past_prompt = slot.prompt_tokens.size_up_to_pos(pos_next); + { + const llama_pos pos_next_prompt = std::min( + slot.prompt_tokens.pos_next(slot.n_past_prompt), + it->pos_max_prompt + 1); + slot.n_past_prompt = slot.prompt_tokens.size_up_to_pos(pos_next_prompt); + } + + slot.checkpoint_pos = it->pos_max; + SLT_WRN(slot, "restored context checkpoint took %.2f ms (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_past = %d, size = %.3f MiB)\n", (ggml_time_us() - t_start) / 1000.0, it->pos_min, it->pos_max, it->n_tokens, slot.n_past, (float)checkpoint_size / 1024 / 1024); } } if (do_reset) { - SLT_WRN(slot, "forcing full prompt re-processing due to lack of cache data (likely due to SWA, see %s)\n", - "https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055"); + if (is_dsv4) { + SLT_WRN(slot, "%s", "no checkpoint before divergence point - reprocessing from scratch\n"); + } else { + SLT_WRN(slot, "forcing full prompt re-processing due to lack of cache data (likely due to SWA, see %s)\n", + "https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055"); + } slot.n_past = 0; slot.n_past_prompt = 0; slot.n_past_se = 0; @@ -3657,7 +3692,7 @@ void server_context::apply_checkpoint(server_slot & slot) { } { - // erase any checkpoints with pos_min > pos_min_thold + // erase checkpoints whose data extends at or past the next write position for (auto it = slot.server_cached_prompt.checkpoints.begin(); it != slot.server_cached_prompt.checkpoints.end();) { const auto & cur = *it; if (cur.pos_max > pos_min_thold) { @@ -4667,7 +4702,7 @@ void server_context::process_batch_tokens(int32_t & n_batch) { if (slot.do_checkpoint) { create_checkpoint(slot); } else { - create_checkpoint_at_interval(slot, params_base); + create_checkpoint_at_interval(slot); } } continue; // continue loop of slots @@ -4744,7 +4779,7 @@ void server_context::process_batch_tokens(int32_t & n_batch) { // create checkpoint during generation if (slot.n_decoded > 1) { - create_checkpoint_at_interval(slot, params_base); + create_checkpoint_at_interval(slot); } slot.t_token_generation = std::max(1, t_current - slot.t_start_generation) / 1e3; diff --git a/examples/server/server-context.h b/examples/server/server-context.h index c2936a41..c6f3d8ee 100644 --- a/examples/server/server-context.h +++ b/examples/server/server-context.h @@ -119,7 +119,7 @@ struct server_slot { void prompt_load(server_prompt_cache& prompt_cache, const server_tokens& tokens, float min_reusable_fraction); - size_t checkpoint_pos = 0; + llama_pos checkpoint_pos = -1; bool do_checkpoint = false; bool image_just_processed = false; @@ -393,7 +393,7 @@ struct server_context { void apply_checkpoint(server_slot & slot); - void create_checkpoint_at_interval(server_slot & slot, const gpt_params & params_base); + void create_checkpoint_at_interval(server_slot & slot); void release_slot_after_final_response(server_slot & slot); }; diff --git a/examples/server/server-task.h b/examples/server/server-task.h index 9685cd28..d04336a1 100644 --- a/examples/server/server-task.h +++ b/examples/server/server-task.h @@ -365,7 +365,7 @@ struct server_prompt_checkpoint { return data.size(); } - json to_json() { + json to_json() const { json j; j["pos_min"] = pos_min; j["pos_max"] = pos_max; @@ -412,7 +412,7 @@ struct server_prompt { }; } - json to_json() + json to_json() const { json j; j["tokens"] = tokens.to_json(); @@ -425,7 +425,6 @@ struct server_prompt { tokens.from_json(j.at("tokens")); n_kept_prompt = j.value("n_kept_prompt", 0); n_discarded_prompt = j.value("n_discarded_prompt", 0); - n_kept_prompt = j.value("n_kept_prompt", 0); } }; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 5014ab76..bb8f5b48 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2219,8 +2219,9 @@ bool llama_model_supports_ctx_shift(const struct llama_model * model) { } bool llama_model_supports_partial_kv_reuse(const struct llama_model * model) { - // These architectures cannot reconstruct their private per-position state after a mid-sequence rewind. - return model && model->arch != LLM_ARCH_OPENPANGU && model->arch != LLM_ARCH_DEEPSEEK4; + // OpenPangu has position-dependent private state outside the generic KV cache. + // DSV4 also has private per-position state, but uses state checkpoints to restore. + return model && model->arch != LLM_ARCH_OPENPANGU; } llm_tensor llm_tensor_type(llm_arch arch, const std::string & tensor_name, int il) { diff --git a/src/llama.cpp b/src/llama.cpp index 0ab709f4..0c1ac393 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -9296,6 +9296,23 @@ static inline ggml_tensor * get_kv_cache_split_tensor(const ggml_tensor * tensor return kv; } +// Compute per-stream byte offset and size for a DSV4 cache tensor. +// stream_idx >= 0 gives that stream's portion; use -1 for the full tensor. +// Tensors are laid out as [ne0, ne1, ...] with ne1 = per_stream_rows * n_stream. +static bool dsv4_stream_offset_size(const struct ggml_tensor * tensor, uint32_t n_stream, int32_t stream_idx, size_t & out_offset, size_t & out_size) { + if (stream_idx < 0 || (uint32_t)stream_idx >= n_stream || n_stream == 0 || tensor->ne[1] % n_stream != 0) { + LLAMA_LOG_ERROR("%s: invalid stream_idx=%d n_stream=%u ne[1]=%lld\n", __func__, stream_idx, n_stream, (long long)tensor->ne[1]); + out_offset = 0; + out_size = 0; + return false; + } + const size_t row_size = ggml_row_size(tensor->type, tensor->ne[0]); + const uint32_t rows_per_stream = (uint32_t)(tensor->ne[1] / n_stream); + out_offset = (size_t)stream_idx * rows_per_stream * row_size; + out_size = (size_t)rows_per_stream * row_size; + return true; +} + // TODO: replace all non-fatal assertions with returned errors or exceptions struct llama_data_write { virtual void write(const void * src, size_t size) = 0; @@ -9560,6 +9577,60 @@ struct llama_data_write { } } } + + // DSV4 compressed indexer cache (only for DSV4 models — preserves + // the old file layout for all other architectures) + if (ctx->model.arch == LLM_ARCH_DEEPSEEK4 && ctx->dsv4.cache.cache_ctx != nullptr) { + const uint32_t dsv4_n_layer = n_layer; + write(&dsv4_n_layer, sizeof(dsv4_n_layer)); + + // Per-sequence save: only the stream for this seq_id + // Full save: all streams + const uint32_t dsv4_single_stream = (seq_id != -1) ? 1 : 0; + write(&dsv4_single_stream, sizeof(dsv4_single_stream)); + const int32_t dsv4_stream_idx = (seq_id != -1) ? seq_id : -1; + write(&dsv4_stream_idx, sizeof(dsv4_stream_idx)); + write(&ctx->dsv4.cache.n_stream, sizeof(ctx->dsv4.cache.n_stream)); + + for (uint32_t il = 0; il < n_layer; ++il) { + uint32_t layer_type = 0; + if (il < ctx->dsv4.cache.csa_k.size() && ctx->dsv4.cache.csa_k[il] != nullptr) { + layer_type = 1; // CSA+LID layer + } else if (il < ctx->dsv4.cache.hca_k.size() && ctx->dsv4.cache.hca_k[il] != nullptr) { + layer_type = 2; // HCA layer + } + write(&layer_type, sizeof(layer_type)); + if (layer_type != 0) { + write_dsv4_cache(ctx, il, dsv4_stream_idx); + } + } + } + } + + void write_dsv4_cache(const struct llama_context * ctx, int il, int32_t stream_idx) { + const auto & cache = ctx->dsv4.cache; + const uint32_t n_stream = cache.n_stream; + auto write_tensor_stream = [&](const struct ggml_tensor * tensor, int layer_il) { + if (stream_idx < 0) { + write_tensor_data(tensor, 0, ggml_nbytes(tensor), layer_il); + } else { + size_t offset, size; + GGML_ASSERT(dsv4_stream_offset_size(tensor, n_stream, stream_idx, offset, size)); + write_tensor_data(tensor, offset, size, layer_il); + } + }; + if (il < (int)cache.csa_k.size() && cache.csa_k[il] != nullptr) { + write_tensor_stream(cache.csa_k[il], il); + write_tensor_stream(cache.lid_k[il], il); + write_tensor_stream(cache.csa_state_kv[il], il); + write_tensor_stream(cache.csa_state_score[il], il); + write_tensor_stream(cache.lid_state_kv[il], il); + write_tensor_stream(cache.lid_state_score[il], il); + } else if (il < (int)cache.hca_k.size() && cache.hca_k[il] != nullptr) { + write_tensor_stream(cache.hca_k[il], il); + write_tensor_stream(cache.hca_state_kv[il], il); + write_tensor_stream(cache.hca_state_score[il], il); + } } void write_kv_cache(const struct llama_context * ctx, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) { @@ -10131,6 +10202,85 @@ struct llama_data_read { } } } + + // DSV4 compressed indexer cache (only present for DSV4 models) + if (ctx->model.arch == LLM_ARCH_DEEPSEEK4) { + + auto & cache = ctx->dsv4.cache; + if (cache.cache_ctx == nullptr) { + LLAMA_LOG_ERROR("%s: DSV4 cache not initialized\n", __func__); + return false; + } + + uint32_t dsv4_n_layer; + read_to(&dsv4_n_layer, sizeof(dsv4_n_layer)); + if (dsv4_n_layer != n_layer) { + LLAMA_LOG_ERROR("%s: DSV4 cache layer count mismatch (%u != %u)\n", __func__, dsv4_n_layer, n_layer); + return false; + } + + uint32_t dsv4_single_stream; + read_to(&dsv4_single_stream, sizeof(dsv4_single_stream)); + + int32_t dsv4_stream_idx; + read_to(&dsv4_stream_idx, sizeof(dsv4_stream_idx)); + + uint32_t dsv4_n_stream; + read_to(&dsv4_n_stream, sizeof(dsv4_n_stream)); + if (dsv4_n_stream != cache.n_stream) { + LLAMA_LOG_ERROR("%s: DSV4 cache stream count mismatch (%u != %u)\n", __func__, dsv4_n_stream, cache.n_stream); + return false; + } + + // Consistency check: per-stream data requires a destination seq_id + if (dsv4_single_stream && seq_id == -1) { + LLAMA_LOG_ERROR("%s: per-stream DSV4 cache cannot be restored to full KV cache\n", __func__); + return false; + } + if (!dsv4_single_stream && seq_id != -1) { + LLAMA_LOG_ERROR("%s: full-stream DSV4 cache cannot be restored to single sequence\n", __func__); + return false; + } + + // Destination stream: when restoring per-stream, write to seq_id's slot + const int32_t dsv4_dst_stream = dsv4_single_stream ? (int32_t)seq_id : -1; + + for (uint32_t il = 0; il < n_layer; ++il) { + uint32_t layer_type; + read_to(&layer_type, sizeof(layer_type)); + + bool set_ok = true; + auto set_tensor_stream = [&](struct ggml_tensor * tensor) { + if (!set_ok) return; + if (dsv4_single_stream) { + size_t dst_offset, stream_size; + if (!dsv4_stream_offset_size(tensor, cache.n_stream, dsv4_dst_stream, dst_offset, stream_size)) { + set_ok = false; + return; + } + ggml_backend_tensor_set(tensor, read(stream_size), dst_offset, stream_size); + } else { + ggml_backend_tensor_set(tensor, read(ggml_nbytes(tensor)), 0, ggml_nbytes(tensor)); + } + }; + + if (layer_type == 1) { + set_tensor_stream(cache.csa_k[il]); + set_tensor_stream(cache.lid_k[il]); + set_tensor_stream(cache.csa_state_kv[il]); + set_tensor_stream(cache.csa_state_score[il]); + set_tensor_stream(cache.lid_state_kv[il]); + set_tensor_stream(cache.lid_state_score[il]); + } else if (layer_type == 2) { + set_tensor_stream(cache.hca_k[il]); + set_tensor_stream(cache.hca_state_kv[il]); + set_tensor_stream(cache.hca_state_score[il]); + } + if (!set_ok) { + return false; + } + } + } return true; } @@ -10403,9 +10553,8 @@ struct llama_data_read_file : llama_data_read { // Public state I/O excludes private DSV4 state, speculation uses an internal checkpoint. static bool llama_state_io_supported(const struct llama_context * ctx, const char * func) { - if (ctx->model.arch == LLM_ARCH_OPENPANGU || ctx->model.arch == LLM_ARCH_DEEPSEEK4) { - const char * arch = ctx->model.arch == LLM_ARCH_OPENPANGU ? "openPangu" : "DeepSeek4"; - LLAMA_LOG_ERROR("%s: state save/restore is not supported for %s (private cache and side state are not serialized)\n", func, arch); + if (ctx->model.arch == LLM_ARCH_OPENPANGU) { + LLAMA_LOG_ERROR("%s: state save/restore is not supported for openPangu (private cache and side state are not serialized)\n", func); return false; } return true;