nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

openpangu: per-sequence state save/restore with --swa-compress (#2261) 

Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com>
This commit is contained in:
Joel Farthing 2026-08-06 03:00:14 -05:00 committed by GitHub
parent cf1aa57e1a
commit d44e2cbe57
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 72 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
examples/server/server-context.cpp
View File

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

3
include/llama.h
View File

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

83
src/llama.cpp
View File

@ -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) &&