From 87eeec9f74d501672a6813f3582cbea8770bb77c Mon Sep 17 00:00:00 2001 From: Joel Farthing Date: Mon, 3 Aug 2026 01:23:01 -0500 Subject: [PATCH] openpangu: support server context checkpoints and prompt reuse (#2245) Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com> --- common/speculative.cpp | 5 + examples/server/server-context.cpp | 54 ++++-- include/llama.h | 4 +- src/graphs/build_openpangu.cpp | 2 +- src/llama-context.h | 5 +- src/llama-model.cpp | 4 +- src/llama.cpp | 267 +++++++++++++++++++++++++++-- 7 files changed, 300 insertions(+), 41 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 9ccec66be..4df99657f 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2055,6 +2055,11 @@ bool common_speculative_finalize_startup( params.clear_dft(); } + if (params.has_dft() && model != nullptr && llama_model_is_openpangu(model)) { + LOG_ERR("%s: openPangu does not support an external draft model; it has its own MTP heads\n", __func__); + return false; + } + if (params.has_dft()) { LLAMA_LOG_INFO("\n\n==================================loading DRAFT model==================================\n\n"); if (!common_speculative_load_draft_model(params, params_base)) { diff --git a/examples/server/server-context.cpp b/examples/server/server-context.cpp index a9a595e15..b28de4ca6 100644 --- a/examples/server/server-context.cpp +++ b/examples/server/server-context.cpp @@ -275,8 +275,9 @@ bool server_context::load_model(const gpt_params& params_) { void server_context::init() { const int32_t n_ctx_slot = n_ctx / params_base.n_parallel; - if (!system_prompt.empty() && llama_model_is_deepseek4(model)) { - throw std::runtime_error("DeepSeek4 server system prompts are unsupported because seq_cp does not copy private cache state"); + if (!system_prompt.empty() && + (llama_model_is_deepseek4(model) || llama_model_is_openpangu(model))) { + throw std::runtime_error("server system prompts are unsupported for openPangu and DeepSeek4 because seq_cp does not copy private cache state"); } LOG_INFO("initializing slots", { {"n_slots", params_base.n_parallel} }); @@ -381,6 +382,15 @@ void server_context::init() { metrics.init(); + bool reuse_forced_off = false; + if (llama_model_is_openpangu(model) && params_base.has_mtp && + (params_base.ctx_checkpoints_n > 0 || params_base.cache_ram_mib != 0)) { + LLAMA_LOG_WARN("context checkpoints and prompt cache are disabled for openPangu while MTP is enabled: the MTP companion keeps its own conv slot, which no saved target state carries\n"); + params_base.ctx_checkpoints_n = 0; + params_base.cache_ram_mib = 0; + reuse_forced_off = true; + } + if (params_base.cache_ram_mib != 0 && llama_model_supports_partial_kv_reuse(model)) { if (params_base.cache_ram_mib < 0) { LLAMA_LOG_INFO("prompt cache is enabled, size limit: %s\n", "no limit"); @@ -395,7 +405,7 @@ void server_context::init() { else { if (params_base.cache_ram_mib != 0) { LLAMA_LOG_WARN("prompt cache is disabled because this model has private state outside the generic KV cache\n"); - } else { + } else if (!reuse_forced_off) { LLAMA_LOG_INFO("%s", "prompt cache is disabled - use `--cache-ram N` to enable it\n"); } } @@ -1782,7 +1792,9 @@ bool server_context::launch_slot_with_task(server_slot& slot, server_task& task) } while (false); slot.allow_rules_prev = slot.allow_rules; - if (llama_model_has_recurrent(llama_get_model(slot.ctx)) || llama_model_is_deepseek4(llama_get_model(slot.ctx))) { + if (llama_model_has_recurrent(llama_get_model(slot.ctx)) || + llama_model_is_openpangu(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 @@ -2075,8 +2087,9 @@ void server_context::system_prompt_update() { } bool server_context::system_prompt_set(const std::string& sys_prompt) { - if (!sys_prompt.empty() && llama_model_is_deepseek4(model)) { - LOG_ERROR("DeepSeek4 server system prompts are unsupported because seq_cp does not copy private cache state", {}); + if (!sys_prompt.empty() && + (llama_model_is_deepseek4(model) || llama_model_is_openpangu(model))) { + LOG_ERROR("server system prompts are unsupported for openPangu and DeepSeek4 because seq_cp does not copy private cache state", {}); return false; } @@ -2807,7 +2820,7 @@ void server_context::process_single_task(server_task&& task) { if (task.data.contains("system_prompt")) { std::string sys_prompt = json_value(task.data, "system_prompt", std::string()); if (!system_prompt_set(sys_prompt)) { - send_error(task, "DeepSeek4 server system prompts are unsupported", ERROR_TYPE_INVALID_REQUEST); + send_error(task, "server system prompts are unsupported for openPangu and DeepSeek4", ERROR_TYPE_INVALID_REQUEST); break; } @@ -2935,6 +2948,11 @@ void server_context::process_single_task(server_task&& task) { break; } + if (llama_model_is_openpangu(model)) { + send_error(task, "slot save is unsupported for openPangu because per-sequence file state is not implemented", ERROR_TYPE_NOT_SUPPORTED); + break; + } + const size_t token_count = slot->cache_tokens.size(); const int64_t t_start = ggml_time_us(); @@ -3620,11 +3638,12 @@ 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); + const bool is_openpangu = llama_model_is_openpangu(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); - // DSV4 has pos_min=0 (no eviction) so the guard always blocks it - if (pos_min >= pos_min_thold || is_dsv4) { + // DSV4 and openPangu have pos_min=0 (no eviction) so the guard always blocks them + if (pos_min >= pos_min_thold || is_dsv4 || is_openpangu) { 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 @@ -3632,7 +3651,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_max < (is_dsv4 ? pos_next : pos_min_thold); + return cur.pos_max < (is_dsv4 || is_openpangu ? pos_next : pos_min_thold); } ); @@ -3642,6 +3661,9 @@ void server_context::apply_checkpoint(server_slot & slot) { // restore the context checkpoint const int64_t t_start = ggml_time_us(); const size_t checkpoint_size = it->data.size(); + if (is_openpangu) { + llama_kv_cache_seq_rm(slot.ctx, slot.id, it->pos_max + 1, -1); + } const size_t n = llama_state_seq_set_data(ctx, it->data.data(), checkpoint_size, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); if (n != checkpoint_size) { @@ -3653,7 +3675,7 @@ void server_context::apply_checkpoint(server_slot & slot) { } if (!do_reset) { - if (is_dsv4) { + if (is_dsv4 || is_openpangu) { 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)); @@ -3674,7 +3696,12 @@ void server_context::apply_checkpoint(server_slot & slot) { } if (do_reset) { - if (is_dsv4) { + if (is_openpangu) { + common_speculative_clear_sequence_kv(slot.spec, ctx, slot.id); + slot.server_cached_prompt.checkpoints.clear(); + slot.checkpoint_pos = -1; + } + if (is_dsv4 || is_openpangu) { 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", @@ -3754,6 +3781,7 @@ bool server_context::create_checkpoint(server_slot & slot) { bool do_checkpoint = !slot.image_just_processed; int32_t pos_min = llama_kv_cache_seq_pos_min(slot.ctx, slot.id); const auto pos_max = llama_kv_cache_seq_pos_max(slot.ctx, slot.id); + const auto checkpoint_pos_min = llama_model_is_openpangu(model) ? pos_max : pos_min; // no need for empty or small checkpoints do_checkpoint = do_checkpoint && (pos_min >= 0 && slot.cache_tokens.n_tokens() >= 64); @@ -3777,7 +3805,7 @@ bool server_context::create_checkpoint(server_slot & slot) { } auto & cur = slot.server_cached_prompt.checkpoints.emplace_back(); - server_prompt_checkpoint_update(cur, ctx, slot.id, slot.cache_tokens.n_tokens(), pos_min, pos_max, slot.n_past_offset); + server_prompt_checkpoint_update(cur, ctx, slot.id, slot.cache_tokens.n_tokens(), checkpoint_pos_min, pos_max, slot.n_past_offset); SLT_WRN(slot, "created context checkpoint %d of %d (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB, took %.2f ms)\n", (int)slot.server_cached_prompt.checkpoints.size(), params_base.ctx_checkpoints_n, cur.pos_min, cur.pos_max, cur.n_tokens, (float)cur.data.size() / 1024 / 1024, diff --git a/include/llama.h b/include/llama.h index fe6e7cf12..95750a90f 100644 --- a/include/llama.h +++ b/include/llama.h @@ -717,9 +717,7 @@ extern "C" { // (K-shift / context shift / self-extend), e.g. openPangu's latent cache. LLAMA_API bool llama_model_supports_ctx_shift(const struct llama_model * model); - // Returns false for models that can only reuse a cached sequence as a pure extension: - // rewinding into the middle of a decoded sequence loses per-position side state - // (e.g. openPangu keeps only the current recurrent conv state). + // Currently true for every model; no architecture is excluded from partial KV reuse. LLAMA_API bool llama_model_supports_partial_kv_reuse(const struct llama_model * model); LLAMA_API const char * llama_model_arch_string(const struct llama_model * model); diff --git a/src/graphs/build_openpangu.cpp b/src/graphs/build_openpangu.cpp index f2065d32f..02600d5f3 100644 --- a/src/graphs/build_openpangu.cpp +++ b/src/graphs/build_openpangu.cpp @@ -185,7 +185,7 @@ static ggml_tensor * openpangu_build_swa_mask_for_graph(llm_build_context & llm, // this site's two taps packed at float offset 2*site_off. The buffer is zeroed at cache // allocation and reset at pos 0, preserving zero history at sequence start (pos-0 graphs // are discarded from reuse via reset_previous, so the baked reset never runs at pos > 0). -// Speculative rollback snapshots/restores the whole slot via the spec checkpoint. +// Speculative rollback restores the whole slot: GPU mode snapshots it, CPU mode uses the partial layout. static ggml_tensor * openpangu_causal_conv(ggml_context * ctx, ggml_cgraph * gf, ggml_tensor * x, ggml_tensor * w, ggml_tensor * state_all, int64_t site_off, diff --git a/src/llama-context.h b/src/llama-context.h index 1b56b8924..391bd39e8 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -64,9 +64,8 @@ struct llama_kv_cache { bool hybrid = false; bool v_trans = true; // the value tensor is transposed - // openPangu s_l holds position-strict MoME conv state, not per-sequence recurrent - // slots; Qwen3Next-style s_l handling (seq ops, state serialization, s_copy) must - // skip it. Speculative rollback snapshots/restores it via the whole-slot spec checkpoint. + // openPangu s_l holds position-strict MoME conv state, not per-sequence recurrent slots: qnext + // seq ops and generic serialization skip it, the openPangu state layouts carry it instead. bool s_l_position_strict = false; // Note: The value of head isn't only used to optimize searching diff --git a/src/llama-model.cpp b/src/llama-model.cpp index bb8f5b48f..fdcd75e2b 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2219,9 +2219,7 @@ bool llama_model_supports_ctx_shift(const struct llama_model * model) { } bool llama_model_supports_partial_kv_reuse(const struct llama_model * model) { - // 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; + return model != nullptr; } 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 d78920092..b7030db73 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -1027,14 +1027,36 @@ static inline uint32_t llama_kv_qnext_state_slots(const llama_kv_cache & cache) } static inline bool llama_kv_has_qnext_state_storage(const llama_kv_cache & cache) { - // openPangu s_l is position-strict conv state, not qnext per-sequence state; keep it - // out of qnext seq-copy and state serialization (rollback rides the spec checkpoint). if (cache.s_l_position_strict) { return false; } return llama_kv_qnext_state_slots(cache) > 0; } +static inline bool llama_kv_has_openpangu_partial_state( + const llama_kv_cache & cache, + llm_arch arch, + llama_state_seq_flags flags) { + return arch == LLM_ARCH_OPENPANGU && + (flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) != 0 && + cache.s_l_position_strict && + llama_kv_qnext_state_slots(cache) > 0; +} + +static llama_pos llama_kv_openpangu_state_pos(const llama_kv_cache & cache, llama_seq_id seq_id) { + // -1 rather than seq_pos_max's 0, so an empty sequence cannot match a checkpoint saved at position 0 + llama_pos result = -1; + for (uint32_t i = 0; i < cache.size; ++i) { + if (cache.cells[i].has_seq_id(seq_id)) { + result = std::max(result, cache.cells[i].pos); + } + } + return result; +} + +static constexpr uint32_t LLAMA_OPENPANGU_PARTIAL_LAYOUT_MAGIC = 0x50414732u; // "PAG2" +static constexpr uint32_t LLAMA_OPENPANGU_PARTIAL_STATE_MAGIC = 0x50414731u; // "PAG1" + static inline bool llama_kv_qnext_seq_id_in_range(const llama_kv_cache & cache, llama_seq_id seq_id) { const uint32_t n_slots = llama_kv_qnext_state_slots(cache); return n_slots > 0 && seq_id >= 0 && (uint32_t) seq_id < n_slots; @@ -1372,9 +1394,6 @@ static bool llama_kv_cache_init( // MoME conv state for ggml_ssm_conv. Each qnext-style slot packs the three // conv sites as two tap-contiguous floats per channel: // [qa 2*n_lora_q | compresskv 2*n_lora_kv | o 2*n_head*v_dim]. - // s_l_position_strict stays true so qnext seq ops and state serialization - // skip this slot; speculative rollback snapshots/restores it via the - // whole-slot spec checkpoint. const int64_t conv_col_ne = hparams.n_lora_q + hparams.n_lora_kv + (int64_t) hparams.n_head(i)*hparams.n_embd_head_v(i); ggml_tensor * s_conv = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2*conv_col_ne, qnext_state_slots); @@ -8935,9 +8954,13 @@ static bool spec_ckpt_try_per_step(llama_kv_cache & kv, const llama_model & mode return true; } -static size_t llama_spec_ckpt_cpu_state_reserve(const llama_context * ctx, llama_seq_id seq_id) { +static size_t llama_spec_ckpt_cpu_state_reserve(llama_context * ctx, llama_seq_id seq_id) { const auto & kv_self = ctx->kv_self; + if (ctx->model.arch == LLM_ARCH_OPENPANGU) { + return llama_state_seq_get_size(ctx, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } + size_t size = sizeof(uint32_t); // cell_count if (seq_id >= 0 && llama_kv_qnext_seq_id_in_range(kv_self, seq_id) && (uint32_t) seq_id < kv_self.size) { @@ -9160,18 +9183,29 @@ enum llama_spec_ckpt_restore_result llama_spec_ckpt_restore_ex( llama_kv_cache_seq_rm(kv, seq_id, n_past, -1); return LLAMA_SPEC_CKPT_RESTORE_BASE_REPLAY_REQUIRED; - case LLAMA_SPEC_CKPT_CPU: + case LLAMA_SPEC_CKPT_CPU: { if (ctx->model.arch == LLM_ARCH_DEEPSEEK4) { llama_kv_cache_seq_rm(kv, seq_id, n_past, -1); return llama_dsv4_spec_ckpt_restore(ctx, false, 0); } + const bool is_openpangu = ctx->model.arch == LLM_ARCH_OPENPANGU; + if (is_openpangu) { + // conv state is position-strict, so the tail goes before it is written, with or without state + llama_kv_cache_seq_rm(kv, seq_id, n_past, -1); + } if (!kv.ckpt.cpu_state_data.empty()) { - llama_state_seq_set_data(ctx, kv.ckpt.cpu_state_data.data(), - kv.ckpt.cpu_state_data.size(), seq_id, - LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + const size_t restored = llama_state_seq_set_data( + ctx, kv.ckpt.cpu_state_data.data(), kv.ckpt.cpu_state_data.size(), + seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + if (is_openpangu && restored != kv.ckpt.cpu_state_data.size()) { + return LLAMA_SPEC_CKPT_RESTORE_FAILED; + } + } + if (!is_openpangu) { + llama_kv_cache_seq_rm(kv, seq_id, n_past, -1); } - llama_kv_cache_seq_rm(kv, seq_id, n_past, -1); return LLAMA_SPEC_CKPT_RESTORE_BASE_REPLAY_REQUIRED; + } default: return LLAMA_SPEC_CKPT_RESTORE_FAILED; @@ -9401,6 +9435,49 @@ struct llama_data_write { } } + void write_openpangu_state( + const struct llama_context * ctx, + llama_seq_id seq_id, + bool write_dsa_marker) { + const struct llama_kv_cache & kv_self = ctx->kv_self; + const uint32_t n_layer = kv_self.k_l.size(); + const uint32_t state_magic = LLAMA_OPENPANGU_PARTIAL_STATE_MAGIC; + const llama_pos state_pos = llama_kv_openpangu_state_pos(kv_self, seq_id); + write(&state_magic, sizeof(state_magic)); + write(&state_pos, sizeof(state_pos)); + + for (uint32_t il = 0; il < n_layer; ++il) { + const bool has_s_cache = il < kv_self.s_l.size() && kv_self.s_l[il] != nullptr; + + const int32_t s_type_i = has_s_cache ? (int32_t) kv_self.s_l[il]->type : -1; + write(&s_type_i, sizeof(s_type_i)); + + const uint64_t s_size_row = has_s_cache ? ggml_row_size(kv_self.s_l[il]->type, kv_self.s_l[il]->ne[0]) : 0; + write(&s_size_row, sizeof(s_size_row)); + + const uint32_t n_slots = has_s_cache ? (uint32_t) kv_self.s_l[il]->ne[1] : 0; + const uint32_t s_rows = has_s_cache && seq_id >= 0 && (uint32_t) seq_id < n_slots ? 1 : 0; + write(&s_rows, sizeof(s_rows)); + } + + if (write_dsa_marker) { + const uint32_t dsa_indexer_state = !kv_self.kr_l.empty() ? 1 : 0; + write(&dsa_indexer_state, sizeof(dsa_indexer_state)); + } + + // seq_id indexes the conv row; openPangu refuses n_seq_max > 1, so this is row 0 as the graph assumes + for (uint32_t il = 0; il < n_layer; ++il) { + const bool has_s_cache = il < kv_self.s_l.size() && kv_self.s_l[il] != nullptr; + if (has_s_cache) { + const uint32_t n_slots = (uint32_t) kv_self.s_l[il]->ne[1]; + if (seq_id >= 0 && (uint32_t) seq_id < n_slots) { + const size_t s_size_row = ggml_row_size(kv_self.s_l[il]->type, kv_self.s_l[il]->ne[0]); + write_tensor_data(kv_self.s_l[il], (size_t) seq_id * s_size_row, s_size_row, il); + } + } + } + } + void write_kv_cache_data(const struct llama_context * ctx, const std::vector> & cell_ranges, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) { const struct llama_kv_cache & kv_self = ctx->kv_self; @@ -9508,9 +9585,15 @@ struct llama_data_write { } } - const uint32_t qnext_state = llama_kv_has_qnext_state_storage(kv_self) ? 1 : 0; + const bool openpangu_partial = llama_kv_has_openpangu_partial_state(kv_self, ctx->model.arch, flags); + const uint32_t qnext_state = (llama_kv_has_qnext_state_storage(kv_self) || openpangu_partial) ? 1 : 0; write(&qnext_state, sizeof(qnext_state)); + if (openpangu_partial) { + write_openpangu_state(ctx, seq_id, true); + return; + } + if (qnext_state != 0) { for (uint32_t il = 0; il < n_layer; ++il) { const bool has_s_cache = il < kv_self.s_l.size() && kv_self.s_l[il] != nullptr; @@ -9567,6 +9650,10 @@ struct llama_data_write { } } + if (ctx->model.arch == LLM_ARCH_OPENPANGU) { + write_openpangu_state(ctx, seq_id, false); + } + // 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) { @@ -9624,6 +9711,14 @@ struct llama_data_write { void write_kv_cache(const struct llama_context * ctx, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) { const struct llama_kv_cache & kv_self = ctx->kv_self; + + if (llama_kv_has_openpangu_partial_state(kv_self, ctx->model.arch, flags)) { + write(&LLAMA_OPENPANGU_PARTIAL_LAYOUT_MAGIC, sizeof(LLAMA_OPENPANGU_PARTIAL_LAYOUT_MAGIC)); + const std::vector> no_ranges; + write_kv_cache_data(ctx, no_ranges, seq_id, flags); + return; + } + std::vector> cell_ranges; // ranges, from inclusive, to exclusive uint32_t cell_count = 0; @@ -9763,6 +9858,11 @@ struct llama_data_read { if (dest_seq_id != -1) { // single sequence + if (cell_count == 0 && ctx->model.arch == LLM_ARCH_OPENPANGU) { + LLAMA_LOG_ERROR("%s: openPangu sequence state carries no kv cells\n", __func__); + return false; + } + llama_kv_cache_seq_rm(kv_self, dest_seq_id, -1, -1); llama_batch batch = llama_batch_init(cell_count, 0, 1); @@ -9797,6 +9897,7 @@ struct llama_data_read { GGML_ASSERT(kv_self.cells[kv_self.head + cell_count - 1].pos == batch.pos[cell_count - 1]); GGML_ASSERT(kv_self.cells[kv_self.head].has_seq_id(dest_seq_id)); GGML_ASSERT(kv_self.cells[kv_self.head + cell_count - 1].has_seq_id(dest_seq_id)); + GGML_ASSERT(ctx->model.arch != LLM_ARCH_OPENPANGU || kv_self.head == 0); // Cleanup llama_batch_free(batch); @@ -9889,6 +9990,97 @@ struct llama_data_read { GGML_ASSERT(sum_split_row_size == row_size); } + bool read_openpangu_state( + struct llama_context * ctx, + uint32_t n_layer, + llama_seq_id seq_id, + bool read_dsa_marker) { + struct llama_kv_cache & kv_self = ctx->kv_self; + + uint32_t state_magic_ref = 0; + llama_pos state_pos_ref = -1; + read_to(&state_magic_ref, sizeof(state_magic_ref)); + read_to(&state_pos_ref, sizeof(state_pos_ref)); + + if (state_magic_ref != LLAMA_OPENPANGU_PARTIAL_STATE_MAGIC) { + LLAMA_LOG_ERROR("%s: incompatible openPangu state geometry\n", __func__); + return false; + } + + const llama_pos state_pos = llama_kv_openpangu_state_pos(kv_self, seq_id); + if (state_pos_ref != state_pos) { + LLAMA_LOG_ERROR("%s: openPangu state position mismatch (saved at %d, restoring at %d)\n", + __func__, (int) state_pos_ref, (int) state_pos); + return false; + } + + std::vector s_size_rows(n_layer, 0); + std::vector s_rows(n_layer, 0); + + for (uint32_t il = 0; il < n_layer; ++il) { + const bool has_s_cache = il < kv_self.s_l.size() && kv_self.s_l[il] != nullptr; + + int32_t s_type_i_ref; + uint64_t s_size_row_ref; + uint32_t s_rows_ref; + read_to(&s_type_i_ref, sizeof(s_type_i_ref)); + read_to(&s_size_row_ref, sizeof(s_size_row_ref)); + read_to(&s_rows_ref, sizeof(s_rows_ref)); + + const int32_t s_type_i = has_s_cache ? (int32_t) kv_self.s_l[il]->type : -1; + const uint64_t s_size_row = has_s_cache + ? ggml_row_size(kv_self.s_l[il]->type, kv_self.s_l[il]->ne[0]) + : 0; + const uint32_t n_slots = has_s_cache ? (uint32_t) kv_self.s_l[il]->ne[1] : 0; + const uint32_t expected_rows = has_s_cache && seq_id >= 0 && (uint32_t) seq_id < n_slots ? 1 : 0; + + if (s_type_i_ref != s_type_i || + s_size_row_ref != s_size_row || + s_rows_ref != expected_rows) { + LLAMA_LOG_ERROR("%s: incompatible openPangu state geometry at layer %u\n", __func__, il); + return false; + } + + s_size_rows[il] = s_size_row; + s_rows[il] = expected_rows; + } + + if (read_dsa_marker) { + uint32_t dsa_indexer_state_ref = 0; + read_to(&dsa_indexer_state_ref, sizeof(dsa_indexer_state_ref)); + const uint32_t dsa_indexer_state = !kv_self.kr_l.empty() ? 1 : 0; + if (dsa_indexer_state_ref != dsa_indexer_state) { + LLAMA_LOG_ERROR("%s: incompatible openPangu state geometry\n", __func__); + return false; + } + } + + size_t payload_size = 0; + for (uint32_t il = 0; il < n_layer; ++il) { + if (s_rows[il] != 0) { + payload_size += (size_t) s_size_rows[il]; + } + } + const uint8_t * payload = read(payload_size); + + size_t payload_offset = 0; + for (uint32_t il = 0; il < n_layer; ++il) { + if (s_rows[il] == 0) { + continue; + } + + const size_t s_size_row = (size_t) s_size_rows[il]; + const uint32_t s_dst_row = (uint32_t) seq_id; + if (kv_self.s_l[il]->extra) { + read_kv_cache_data_split(ctx, kv_self.s_l[il], payload + payload_offset, s_dst_row, s_size_row, 1, il); + } else { + ggml_backend_tensor_set(kv_self.s_l[il], payload + payload_offset, (size_t) s_dst_row * s_size_row, s_size_row); + } + payload_offset += s_size_row; + } + return true; + } + bool read_kv_cache_data(struct llama_context * ctx, uint32_t cell_count, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) { const struct llama_hparams & hparams = ctx->model.hparams; struct llama_kv_cache & kv_self = ctx->kv_self; @@ -10088,12 +10280,17 @@ struct llama_data_read { uint32_t qnext_state_ref = 0; read_to(&qnext_state_ref, sizeof(qnext_state_ref)); - const bool has_qnext_state = llama_kv_has_qnext_state_storage(kv_self); + const bool openpangu_partial = llama_kv_has_openpangu_partial_state(kv_self, ctx->model.arch, flags); + const bool has_qnext_state = llama_kv_has_qnext_state_storage(kv_self) || openpangu_partial; if ((qnext_state_ref != 0) != has_qnext_state) { LLAMA_LOG_ERROR("%s: incompatible qwen3next state cache presence\n", __func__); return false; } + if (openpangu_partial) { + return read_openpangu_state(ctx, n_layer, seq_id, true); + } + if (qnext_state_ref != 0) { for (uint32_t il = 0; il < n_layer; ++il) { const bool has_s_cache = il < kv_self.s_l.size() && kv_self.s_l[il] != nullptr; @@ -10197,6 +10394,11 @@ struct llama_data_read { } } + if (ctx->model.arch == LLM_ARCH_OPENPANGU && + !read_openpangu_state(ctx, n_layer, seq_id, false)) { + return false; + } + // DSV4 compressed indexer cache (only present for DSV4 models) if (ctx->model.arch == LLM_ARCH_DEEPSEEK4) { @@ -10282,6 +10484,19 @@ struct llama_data_read { uint32_t cell_count; read_to(&cell_count, sizeof(cell_count)); + const bool openpangu_partial = llama_kv_has_openpangu_partial_state(ctx->kv_self, ctx->model.arch, flags); + const bool skip_meta_layout = cell_count == LLAMA_OPENPANGU_PARTIAL_LAYOUT_MAGIC; + if (openpangu_partial != skip_meta_layout) { + throw std::runtime_error("failed to restore kv cache: incompatible openPangu partial-state layout"); + } + + if (openpangu_partial) { + if (!read_kv_cache_data(ctx, 0, seq_id, flags)) { + throw std::runtime_error("failed to restore kv cache: openPangu partial state"); + } + return; + } + bool res = read_kv_cache_meta(ctx, cell_count, seq_id) && read_kv_cache_data(ctx, cell_count, seq_id, flags); if (!res) { @@ -10545,10 +10760,21 @@ 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) { +// Public state I/O excludes private DSV4 state, speculation uses an internal checkpoint; +// openPangu instead carries its private conv state in the per-sequence layouts admitted here. +static bool llama_state_io_supported( + const struct llama_context * ctx, + const char * func, + llama_state_seq_flags flags = 0, + llama_seq_id seq_id = -1) { 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); + if (seq_id >= 0 && + llama_kv_qnext_seq_id_in_range(ctx->kv_self, seq_id) && + (flags == 0 || + llama_kv_has_openpangu_partial_state(ctx->kv_self, ctx->model.arch, flags))) { + return true; + } + LLAMA_LOG_ERROR("%s: only per-sequence state save/restore is supported for openPangu (whole-context and file-session state are not)\n", func); return false; } return true; @@ -10727,7 +10953,12 @@ bool llama_state_save_file(struct llama_context * ctx, const char * path_session } static size_t llama_state_seq_get_data_internal(struct llama_context * ctx, llama_data_write & data_ctx, llama_seq_id seq_id, llama_state_seq_flags flags) { - if (!llama_state_io_supported(ctx, __func__)) { + if (!llama_state_io_supported(ctx, __func__, flags, seq_id)) { + return 0; + } + if (ctx->model.arch == LLM_ARCH_OPENPANGU && flags == 0 && + llama_kv_openpangu_state_pos(ctx->kv_self, seq_id) < 0) { + LLAMA_LOG_ERROR("%s: openPangu sequence %d has no kv cells to save\n", __func__, seq_id); return 0; } llama_synchronize(ctx); @@ -10753,7 +10984,7 @@ size_t llama_state_seq_get_data(struct llama_context * ctx, uint8_t * dst, size_ } static size_t llama_state_seq_set_data_internal(struct llama_context * ctx, llama_data_read & data_ctx, llama_seq_id dest_seq_id, llama_state_seq_flags flags) { - if (!llama_state_io_supported(ctx, __func__)) { + if (!llama_state_io_supported(ctx, __func__, flags, dest_seq_id)) { return SIZE_MAX; } llama_synchronize(ctx);