Commit e467e5b5e for llama.cpp
commit e467e5b5e5ba43071d6151b3d881deb9300d000e
Author: Pascal <admin@serveurperso.com>
Date: Thu Oct 15 06:39:18 2026 +0200
kv-cache: honor PARTIAL_ONLY in the plain KV cache (#30514)
* llama: honor PARTIAL_ONLY in the plain KV cache
A non-SWA llama_kv_cache rolls back with seq_rm, so it holds no
partial state, yet state_write ignored the flag and serialized every
cell. A draft context made of a plain KV cache (MTP head, dense -md
draft) therefore copied its whole KV into each server checkpoint,
growing it by the draft KV size per token.
The cache now writes and reads nothing for PARTIAL_ONLY unless it is
SWA, and the iSWA wrapper no longer gates its base cache since it
filters itself.
* Update src/llama-kv-cache.cpp
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
* llama: delegate PARTIAL_ONLY to the caches of every memory wrapper
The hybrid, hybrid-idx and dsa-iswa wrappers skipped their attention
caches on PARTIAL_ONLY, which also dropped an SWA window that is
partial state. Each cache now applies the flag itself, so the
wrappers only delegate. The indexer is marked stale only when a
restore rewrites cells.
test-llama-context gains a case checking that the partial state is
not larger than the full state and, without SWA, does not grow with
the tokens.
---------
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
diff --git a/src/llama-kv-cache-dsa-iswa.cpp b/src/llama-kv-cache-dsa-iswa.cpp
index dc10342a1..784c78aa7 100644
--- a/src/llama-kv-cache-dsa-iswa.cpp
+++ b/src/llama-kv-cache-dsa-iswa.cpp
@@ -234,18 +234,12 @@ bool llama_kv_cache_dsa_iswa::get_can_shift() const {
}
void llama_kv_cache_dsa_iswa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
- if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
- kv_dsa->state_write(io, seq_id, flags);
- }
-
+ kv_dsa->state_write(io, seq_id, flags);
kv_swa->state_write(io, seq_id, flags);
}
void llama_kv_cache_dsa_iswa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
- if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
- kv_dsa->state_read(io, seq_id, flags);
- }
-
+ kv_dsa->state_read(io, seq_id, flags);
kv_swa->state_read(io, seq_id, flags);
}
diff --git a/src/llama-kv-cache-iswa.cpp b/src/llama-kv-cache-iswa.cpp
index a6ce547e6..da440f6a2 100644
--- a/src/llama-kv-cache-iswa.cpp
+++ b/src/llama-kv-cache-iswa.cpp
@@ -255,19 +255,13 @@ bool llama_kv_cache_iswa::get_can_shift() const {
}
void llama_kv_cache_iswa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
- if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
- kv_base->state_write(io, seq_id, flags);
- }
-
- kv_swa->state_write(io, seq_id, flags);
+ kv_base->state_write(io, seq_id, flags);
+ kv_swa ->state_write(io, seq_id, flags);
}
void llama_kv_cache_iswa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
- if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
- kv_base->state_read(io, seq_id, flags);
- }
-
- kv_swa->state_read(io, seq_id, flags);
+ kv_base->state_read(io, seq_id, flags);
+ kv_swa ->state_read(io, seq_id, flags);
}
llama_kv_cache * llama_kv_cache_iswa::get_base() const {
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index d5208d034..60a7e8554 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -2076,7 +2076,9 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla
return;
}
- GGML_UNUSED(flags);
+ if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) && swa_type == LLAMA_SWA_TYPE_NONE) {
+ return;
+ }
io.write(&n_stream, sizeof(n_stream));
@@ -2155,7 +2157,9 @@ const slot_info_vec_t * sinfos_in) {
return;
}
- GGML_UNUSED(flags);
+ if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) && swa_type == LLAMA_SWA_TYPE_NONE) {
+ return;
+ }
// TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG]
GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()));
diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp
index 21ec0b2fe..2c5a57680 100644
--- a/src/llama-memory-hybrid-idx.cpp
+++ b/src/llama-memory-hybrid-idx.cpp
@@ -257,13 +257,9 @@ void llama_memory_hybrid_idx::state_write(llama_io_write_i & io, llama_seq_id se
llama_memory_hybrid::state_write(io, seq_id, flags);
// [TAG_HYBRID_IDX_STATE] the indexer section goes last, so it is a pure suffix: an old reader stops early instead of misparsing it
- // The indexer mirrors the attention cache, so it uses the same PARTIAL_ONLY gate.
- if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
- if (mem_idx) {
- mem_idx->state_write(io, seq_id, flags);
- }
+ if (mem_idx) {
+ mem_idx->state_write(io, seq_id, flags);
}
-
}
void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
@@ -276,17 +272,16 @@ void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_
llama_kv_cache::slot_info_vec_t sinfos_attn;
try {
- if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
- get_mem_attn()->state_read_sinfo(io, seq_id, flags, mem_idx ? &sinfos_attn : nullptr, nullptr);
- }
+ // the indexer cache has the SWA type of the attention cache, so both read the same partial state
+ get_mem_attn()->state_read_sinfo(io, seq_id, flags, mem_idx ? &sinfos_attn : nullptr, nullptr);
get_mem_recr()->state_read(io, seq_id, flags);
// [TAG_HYBRID_IDX_STATE] must mirror the write order in state_write
- if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
- if (mem_idx) {
- mem_idx->state_read_sinfo(io, seq_id, flags, nullptr, &sinfos_attn);
- // the restore rewrites the cells behind the pool layout's back
+ if (mem_idx) {
+ mem_idx->state_read_sinfo(io, seq_id, flags, nullptr, &sinfos_attn);
+ // a restore that rewrites cells does it behind the pool layout's back
+ if (!sinfos_attn.empty()) {
mem_idx_stale_set(seq_id, 0);
}
}
diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp
index 62f12de95..bbd4bee0b 100644
--- a/src/llama-memory-hybrid.cpp
+++ b/src/llama-memory-hybrid.cpp
@@ -188,26 +188,18 @@ std::map<ggml_backend_buffer_type_t, size_t> llama_memory_hybrid::memory_breakdo
}
void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
- if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
- mem_attn->state_write(io, seq_id, flags);
- }
+ mem_attn->state_write(io, seq_id, flags);
mem_recr->state_write(io, seq_id, flags);
}
void llama_memory_hybrid::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
- const bool read_attn = (flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0;
-
- if (read_attn) {
- mem_attn->state_read(io, seq_id, flags);
- }
+ mem_attn->state_read(io, seq_id, flags);
try {
mem_recr->state_read(io, seq_id, flags);
} catch (...) {
- // the attention part is already restored - undo it
- if (read_attn) {
- mem_attn->state_clear(seq_id);
- }
+ // a failed restore leaves the sequence empty in both caches
+ mem_attn->state_clear(seq_id);
throw;
}
diff --git a/tests/llama-context/cases-kv.cpp b/tests/llama-context/cases-kv.cpp
index 152977468..27e6493e4 100644
--- a/tests/llama-context/cases-kv.cpp
+++ b/tests/llama-context/cases-kv.cpp
@@ -341,6 +341,49 @@ struct case_state_roundtrip : ctx_test {
}
};
+// the partial state holds what seq_rm cannot roll back: recurrent states and SWA windows
+struct case_state_partial_size : ctx_test {
+ const char * name() const override { return "ps"; }
+
+ const char * description() const override {
+ return "the partial state is not larger than the full state and, without SWA, does not grow with the tokens";
+ }
+
+ case_result run(model_run & mr) const override {
+ auto ctx = mr.ctx();
+ if (!ctx) {
+ return case_result::fail();
+ }
+
+ const llama_pos n_tokens = mr.tokens.size();
+ const llama_pos n_half = n_tokens / 2;
+
+ if (!decode_seq(ctx.get(), mr.tokens, 0, n_half, 0, true)) {
+ return case_result::fail();
+ }
+
+ const size_t partial_a = llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+ const size_t full_a = llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE);
+
+ if (!decode_seq(ctx.get(), mr.tokens, n_half, n_tokens, 0, true)) {
+ return case_result::fail();
+ }
+
+ const size_t partial_b = llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+ const size_t full_b = llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE);
+
+ TLOG_INF("%s: partial %zu -> %zu, full %zu -> %zu\n", __func__, partial_a, partial_b, full_a, full_b);
+
+ CTX_REQUIRE(partial_a <= full_a && partial_b <= full_b, "the partial state is larger than the full state");
+
+ if (llama_model_n_swa(mr.model) == 0) {
+ CTX_REQUIRE(partial_a == partial_b, "the partial state grows with the tokens");
+ }
+
+ return case_result::pass();
+ }
+};
+
// a failed restore must leave the sequence empty and must not change the logits of other sequences
struct case_state_restore_failure : ctx_test {
const char * name() const override { return "rf"; }
@@ -686,6 +729,7 @@ std::vector<ctx_test_ptr> make_kv_cases() {
cases.push_back(std::make_unique<case_seq_cp_scatter>(false));
cases.push_back(std::make_unique<case_seq_cp_scatter>(true));
cases.push_back(std::make_unique<case_state_roundtrip>());
+ cases.push_back(std::make_unique<case_state_partial_size>());
cases.push_back(std::make_unique<case_state_restore_failure>());
cases.push_back(std::make_unique<case_state_rotation>());
cases.push_back(std::make_unique<case_state_restore_fragmented>());