openpangu: support server context checkpoints and prompt reuse (#2245)

Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com>

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)) {

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,

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);

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,

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

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) {

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<std::pair<uint32_t, uint32_t>> & 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<std::pair<uint32_t, uint32_t>> no_ranges;
+            write_kv_cache_data(ctx, no_ranges, seq_id, flags);
+            return;
+        }
+
         std::vector<std::pair<uint32_t, uint32_t>> 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<uint64_t> s_size_rows(n_layer, 0);
+        std::vector<uint32_t> 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);