Commit 0f663d247 for llama.cpp

commit 0f663d2471bb470c1979999daf7a05ec8fba5673
Author: Aman Gupta <amangupta052@gmail.com>
Date:   Tue Nov 3 14:08:42 2026 +0800

    dsv4.1: fix llama-context tests with partial offload (#31292)

diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp
index d5506f7a7..ef1be94a2 100644
--- a/src/llama-kv-cache-dsv4.cpp
+++ b/src/llama-kv-cache-dsv4.cpp
@@ -428,25 +428,6 @@ static void dsv4_state_read_k_cache(

         dsv4_state_read_tensor_streams(io, kv->get_k_storage(il), kv_size, n_rows_ref, s0, ns);
     }
-
-    if (n_rows_ref == 0) {
-        return;
-    }
-
-    // refill the mirrors from their sources; a mirror gets the same rows as its source, so a whole copy is exact
-    for (uint32_t il : kv->get_layer_ids()) {
-        const int32_t src = kv_src[il];
-        if (src == (int32_t) il) {
-            continue;
-        }
-        ggml_tensor * t_src = kv->get_k_storage(src);
-        ggml_tensor * t_dst = kv->get_k_storage(il);
-        if (t_src->type != t_dst->type || !ggml_are_same_shape(t_src, t_dst)) {
-            throw std::runtime_error("DSV4 K-cache mirror layout mismatch");
-        }
-
-        ggml_backend_tensor_copy(t_src, t_dst);
-    }
 }

 static std::string dsv4_plan_positions(const std::vector<int32_t> & values) {
@@ -1797,6 +1778,9 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id,
         dsv4_state_read_k_cache(io, kv_csa.get(), seq_id, flags, comp_il_src);
         dsv4_state_read_k_cache(io, kv_hca.get(), seq_id, flags, comp_il_src);
         dsv4_state_read_k_cache(io, kv_lid.get(), seq_id, flags, comp_il_src);
+
+        // the io can defer the tensor reads until it is destroyed, so the mirrors are refilled before the next graph
+        comp_mirrors_stale = true;
     }

     csa_state->state_read(io, seq_id, flags);
@@ -1882,6 +1866,30 @@ void llama_kv_cache_dsv4::update_rs_for_ubatches(const std::vector<llama_ubatch>
     }
 }

+void llama_kv_cache_dsv4::refill_comp_mirrors() {
+    if (!comp_mirrors_stale) {
+        return;
+    }
+
+    comp_mirrors_stale = false;
+
+    // a mirror gets the same rows as its source, so a whole copy is exact
+    for (const llama_kv_cache * kv : { kv_csa.get(), kv_hca.get(), kv_lid.get() }) {
+        for (uint32_t il : kv->get_layer_ids()) {
+            const int32_t src = comp_il_src[il];
+            if (src == (int32_t) il) {
+                continue;
+            }
+
+            ggml_tensor * t_src = kv->get_k_storage(src);
+            ggml_tensor * t_dst = kv->get_k_storage(il);
+            GGML_ASSERT(t_src->type == t_dst->type && ggml_are_same_shape(t_src, t_dst));
+
+            ggml_backend_tensor_copy(t_src, t_dst);
+        }
+    }
+}
+
 void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) {
     if (seq_id < 0) {
         kv_csa->clear(data);
@@ -2281,6 +2289,8 @@ bool llama_kv_cache_dsv4_context::apply() {
         lid_state->apply_copies(sc_info_lid);
     }

+    kv->refill_comp_mirrors();
+
     return res;
 }

diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h
index 1c8d23ed2..50cb11b62 100644
--- a/src/llama-kv-cache-dsv4.h
+++ b/src/llama-kv-cache-dsv4.h
@@ -157,6 +157,9 @@ public:
     const std::vector<uint32_t> & get_rs_idx() const;
     void update_rs_for_ubatches(const std::vector<llama_ubatch> & ubatches);

+    // copy the sources into their mirrors after a state read
+    void refill_comp_mirrors();
+
 private:
     llama_hparams hparams_raw;
     llama_hparams hparams_csa;
@@ -178,6 +181,8 @@ private:
     // [n_layer] layer -> layer whose compressed-KV storage it reads (itself, its V4.1 kv source or a mirror), -1 = none
     std::vector<int32_t> comp_read;

+    bool comp_mirrors_stale = false;
+
     std::unique_ptr<llama_kv_cache_iswa> kv_raw;
     std::unique_ptr<llama_kv_cache>      kv_csa;
     std::unique_ptr<llama_kv_cache>      kv_hca;
@@ -422,5 +427,5 @@ private:

     const llama_memory_status status;

-    const llama_kv_cache_dsv4 * kv = nullptr;
+    llama_kv_cache_dsv4 * kv = nullptr;
 };