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>

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

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 {

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

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

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

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