cont : minor fixes

ggml-ci

Author: ggerganov

src/llama-context.cpp

@@ -450,9 +450,22 @@ bool llama_context::kv_self_update(bool optimize) {
         memory_force_optimize = false;
 
         const auto kv_state = kv_self->init_update(this, optimize);
-        if (kv_state->get_status() == LLAMA_MEMORY_STATUS_NO_UPDATE) {
-            // no updates need to be performed
-            return false;
+        switch (kv_state->get_status()) {
+            case LLAMA_MEMORY_STATUS_SUCCESS:
+                {
+                    // noop
+                } break;
+            case LLAMA_MEMORY_STATUS_NO_UPDATE:
+                {
+                    // no updates need to be performed
+                    return false;
+                }
+            case LLAMA_MEMORY_STATUS_FAILED_PREPARE:
+            case LLAMA_MEMORY_STATUS_FAILED_COMPUTE:
+                {
+                    LLAMA_LOG_ERROR("%s: failed to prepare memory update\n", __func__);
+                    return false;
+                }
         }
 
         if (!kv_state->apply()) {

src/llama-kv-cache-unified.cpp

@@ -435,14 +435,16 @@ bool llama_kv_cache_unified::update(llama_context * lctx, bool do_shift, const d
         cells.reset_shift();
     }
 
-    if (!dinfo.ids.empty()) {
+    if (!dinfo.empty()) {
         LLAMA_LOG_DEBUG("%s: defragmenting KV cache\n", __func__);
 
         // apply moves:
         {
             const auto n_kv = dinfo.ids.size();
 
             for (uint32_t i = 0; i < n_kv; ++i) {
+                assert(dinfo.ids[i] <= n_kv);
+
                 if (dinfo.ids[i] == n_kv) {
                     continue;
                 }
@@ -1657,7 +1659,7 @@ llama_kv_cache_unified_state::llama_kv_cache_unified_state(
         llama_context * lctx,
         bool do_shift,
         defrag_info dinfo) : status(LLAMA_MEMORY_STATUS_SUCCESS), kv(kv), lctx(lctx), do_shift(do_shift), dinfo(std::move(dinfo)) {
-    if (!do_shift && dinfo.ids.empty()) {
+    if (!do_shift && dinfo.empty()) {
         status = LLAMA_MEMORY_STATUS_NO_UPDATE;
     }
 }
@@ -1684,7 +1686,7 @@ bool llama_kv_cache_unified_state::next() {
 bool llama_kv_cache_unified_state::apply() {
     assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
 
-    // this is a KV cache update
+    // no ubatches -> this is a KV cache update
     if (ubatches.empty()) {
         kv->update(lctx, do_shift, dinfo);
 

src/llama-kv-cache-unified.h

@@ -27,6 +27,10 @@ class llama_kv_cache_unified : public llama_kv_cache {
     using ubatch_heads = std::vector<uint32_t>;
 
     struct defrag_info {
+        bool empty() const {
+            return ids.empty();
+        }
+
         // contains information about which cell moves where:
         //  - cell i moves to ids[i]
         //  - if ids[i] == i || ids[i] == ids.size(), then cell i is not moved