From d44e2cbe575e76fda9378ea9d0c6c28869e1ed3c Mon Sep 17 00:00:00 2001 From: Joel Farthing Date: Thu, 6 Aug 2026 03:00:14 -0500 Subject: [PATCH] openpangu: per-sequence state save/restore with --swa-compress (#2261) Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com> --- examples/server/server-context.cpp | 6 --- include/llama.h | 3 +- src/llama.cpp | 83 +++++++++++++++++++++++++----- 3 files changed, 72 insertions(+), 20 deletions(-) diff --git a/examples/server/server-context.cpp b/examples/server/server-context.cpp index 8abe97c1..38d22d2a 100644 --- a/examples/server/server-context.cpp +++ b/examples/server/server-context.cpp @@ -391,12 +391,6 @@ void server_context::init() { reuse_forced_off = true; } - if (params_base.cache_ram_mib != 0 && !llama_supports_full_state_io(ctx)) { - LLAMA_LOG_WARN("prompt cache is disabled: this context cannot save full sequence state (--swa-compress)\n"); - 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"); diff --git a/include/llama.h b/include/llama.h index eb7030c3..e6d2510c 100644 --- a/include/llama.h +++ b/include/llama.h @@ -722,8 +722,7 @@ extern "C" { // 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); - // False when full-state seq save/restore cannot work for this context (--swa-compress). - // PARTIAL_ONLY is unaffected. Context property, not a model one. + // true for every non-null context; none is excluded from full-state seq save/restore LLAMA_API bool llama_supports_full_state_io(const struct llama_context * ctx); LLAMA_API const char * llama_model_arch_string(const struct llama_model * model); diff --git a/src/llama.cpp b/src/llama.cpp index e5e228b7..9e8d6bd7 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -1070,6 +1070,7 @@ static llama_pos llama_kv_openpangu_state_pos(const llama_kv_cache & cache, llam static constexpr uint32_t LLAMA_OPENPANGU_PARTIAL_LAYOUT_MAGIC = 0x50414732u; // "PAG2" static constexpr uint32_t LLAMA_OPENPANGU_PARTIAL_STATE_MAGIC = 0x50414731u; // "PAG1" +static constexpr uint32_t LLAMA_SWA_COMPACT_LAYOUT_MAGIC = 0x53574143u; // "SWAC" 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); @@ -1513,6 +1514,7 @@ static bool llama_kv_cache_init( GGML_ABORT("fatal error"); } + // allocate tensors and initialize the buffers to avoid NaNs in the padding for (auto it : ctx_map) { ggml_backend_buffer_type_t buft = it.first; @@ -8595,7 +8597,7 @@ void llama_free(struct llama_context * ctx) { } bool llama_supports_full_state_io(const struct llama_context * ctx) { - return ctx != nullptr && !ctx->kv_self.any_compacted(); + return ctx != nullptr; } const struct llama_vocab* llama_model_get_vocab(const struct llama_model* model) { @@ -9716,6 +9718,15 @@ struct llama_data_write { continue; } + // only k_l is compacted; it holds the window at [sink_rows, sink_rows + live_swa()) + if (kv_self.is_compacted((int) il)) { + const size_t live = kv_self.live_swa(); + if (live) { + write_tensor_data(kv_self.k_l[il], kv_self.sink_rows * k_size_row, live * k_size_row, il); + } + continue; + } + // Read each range of cells of k_size length each into tmp_buf and write out for (const auto & range : cell_ranges) { const size_t range_size = range.second - range.first; @@ -9950,6 +9961,15 @@ struct llama_data_write { } GGML_ASSERT(cell_count == cell_count_check); + if (kv_self.any_compacted()) { + const uint32_t live_rows = kv_self.live_swa(); + write(&LLAMA_SWA_COMPACT_LAYOUT_MAGIC, sizeof(LLAMA_SWA_COMPACT_LAYOUT_MAGIC)); + write(&kv_self.size_swa, sizeof(kv_self.size_swa)); + write(&live_rows, sizeof(live_rows)); + write(&kv_self.pos_base_swa, sizeof(kv_self.pos_base_swa)); + write(&kv_self.head_swa, sizeof(kv_self.head_swa)); + } + write(&cell_count, sizeof(cell_count)); write_kv_cache_meta(kv_self, cell_ranges, seq_id); @@ -10063,6 +10083,11 @@ struct llama_data_read { return false; } + if (cell_count > kv_self.size) { + LLAMA_LOG_ERROR("%s: not enough cells in kv cache\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); @@ -10353,12 +10378,16 @@ struct llama_data_read { return false; } - if (cell_count) { + const bool k_compact = kv_self.is_compacted((int) il); + const uint32_t k_rows = k_compact ? kv_self.live_swa() : cell_count; + const uint32_t k_dst = k_compact ? kv_self.sink_rows : kv_self.head; + + if (k_rows) { // Read and set the keys for the whole cell range if (kv_self.k_l[il]->extra) { - read_kv_cache_data_split(ctx, kv_self.k_l[il], read(cell_count * k_size_row), kv_self.head, k_size_row, cell_count, il); + read_kv_cache_data_split(ctx, kv_self.k_l[il], read(k_rows * k_size_row), k_dst, k_size_row, k_rows, il); } else { - ggml_backend_tensor_set(kv_self.k_l[il], read(cell_count * k_size_row), kv_self.head * k_size_row, cell_count * k_size_row); + ggml_backend_tensor_set(kv_self.k_l[il], read(k_rows * k_size_row), k_dst * k_size_row, k_rows * k_size_row); } } } @@ -10697,7 +10726,44 @@ struct llama_data_read { return; } - bool res = read_kv_cache_meta(ctx, cell_count, seq_id) && read_kv_cache_data(ctx, cell_count, seq_id, flags); + struct llama_kv_cache & kv_self = ctx->kv_self; + + const bool compacted = kv_self.any_compacted(); + const bool compact_blob = cell_count == LLAMA_SWA_COMPACT_LAYOUT_MAGIC; + if (compacted != compact_blob) { + throw std::runtime_error("failed to restore kv cache: incompatible compacted sliding-window layout"); + } + + uint32_t restore_size_swa = 0; + uint32_t restore_live = 0; + llama_pos restore_pos_base = 0; + uint32_t restore_head_swa = 0; + if (compact_blob) { + read_to(&restore_size_swa, sizeof(restore_size_swa)); + read_to(&restore_live, sizeof(restore_live)); + read_to(&restore_pos_base, sizeof(restore_pos_base)); + read_to(&restore_head_swa, sizeof(restore_head_swa)); + if (restore_size_swa != kv_self.size_swa) { + throw std::runtime_error("failed to restore kv cache: compacted row count differs from this context"); + } + if (restore_head_swa < kv_self.sink_rows || restore_head_swa > kv_self.size_swa || + restore_head_swa - kv_self.sink_rows != restore_live) { + throw std::runtime_error("failed to restore kv cache: inconsistent compacted window state"); + } + read_to(&cell_count, sizeof(cell_count)); + if (restore_pos_base < 0 || + (uint64_t) restore_pos_base + restore_live != (uint64_t) cell_count) { + throw std::runtime_error("failed to restore kv cache: compacted window base inconsistent with cell count"); + } + } + + // scalars before rows: they set the placement offset, and seq_rm inside meta resets head_swa + bool res = read_kv_cache_meta(ctx, cell_count, seq_id); + if (res && compact_blob) { + kv_self.pos_base_swa = restore_pos_base; + kv_self.head_swa = restore_head_swa; + } + res = res && read_kv_cache_data(ctx, cell_count, seq_id, flags); if (!res) { if (seq_id == -1) { @@ -10967,13 +11033,6 @@ static bool llama_state_io_supported( const char * func, llama_state_seq_flags flags = 0, llama_seq_id seq_id = -1) { - if (ctx->kv_self.any_compacted() && flags == 0) { - LLAMA_LOG_ERROR("%s: full state save/restore is not supported with --swa-compress " - "(state I/O addresses cache rows by cell index; compacted layers hold a " - "translated subset)\n", func); - return false; - } - if (ctx->model.arch == LLM_ARCH_OPENPANGU) { if (seq_id >= 0 && llama_kv_qnext_seq_id_in_range(ctx->kv_self, seq_id) &&