kv-cache : rework recovery logic + restrict SWA batch params

ggerganov · ggerganov · commit 27057f24f1b2 · 2025-05-16T19:46:21.000+03:00
ggml-ci
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
@@ -93,6 +93,14 @@ llama_context::llama_context(
     }
 
     cparams.n_ubatch = std::min(cparams.n_batch, params.n_ubatch == 0 ? params.n_batch : params.n_ubatch);
+
+    // TODO: allowing this requires major refactor of the KV cache logic
+    //       ref: https://github.com/ggml-org/llama.cpp/pull/13194
+    if (hparams.n_swa > 0 && cparams.n_ubatch < cparams.n_batch) {
+        LLAMA_LOG_WARN("%s: SWA models currently do not support n_ubatch < n_batch - increasing n_ubatch to %d\n", __func__, cparams.n_batch);
+        cparams.n_ubatch = cparams.n_batch;
+    }
+
     cparams.op_offload = params.op_offload;
 
     const uint32_t n_ctx_per_seq = cparams.n_ctx / cparams.n_seq_max;
@@ -916,6 +924,9 @@ int llama_context::decode(llama_batch & inp_batch) {
         return -2;
     };
 
+    // =============================================================================================================
+    // TODO: refactor and simplify this
+
     // handle any pending defrags/shifts
     kv_self_update();
 
@@ -967,6 +978,8 @@ int llama_context::decode(llama_batch & inp_batch) {
         ubatches.clear();
     }
 
+    // =============================================================================================================
+
     // we now have prepared the ubatches for this llama_decode and are ready to start processing
 
     int64_t n_outputs_prev = 0;
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
@@ -333,44 +333,31 @@ llama_pos llama_kv_cache_unified::seq_pos_max(llama_seq_id seq_id) const {
 }
 
 void llama_kv_cache_unified::restore() {
-    if (pending.ubatches.empty()) {
-        return;
-    }
-
-    uint32_t new_head = size;
-
-    for (const auto & ubatch : pending.ubatches) {
-        for (uint32_t i = 0; i < ubatch.data.n_tokens; ++i) {
-            for (int s = 0; s < ubatch.data.n_seq_id[i]; ++s) {
-                const llama_seq_id seq_id = ubatch.data.seq_id[i][s];
-
-                cells[ubatch.head + i].seq_id.erase(seq_id);
-                if (cells[ubatch.head + i].seq_id.empty()) {
-                    used--;
-
-                    new_head = std::min(new_head, ubatch.head + i);
-                }
+    for (const auto & [id, cell] : recovery.cells) {
+        // TODO: move to new `struct kv_cells`
+        const bool is_empty0 = cells[id].is_empty();
+        const bool is_empty1 = cell.is_empty();
 
-                cells[ubatch.head + i].pos = -1;
-            }
+        if (!is_empty0 && is_empty1) {
+            used--;
+        } else if (is_empty0 && !is_empty1) {
+            used++;
         }
-    }
 
-    if (new_head != size && new_head < head) {
-        head = new_head;
+        cells[id] = cell;
     }
 
-    pending.clear();
+    recovery.clear();
 }
 
 void llama_kv_cache_unified::commit() {
-    if (pending.ubatches.empty()) {
-        LLAMA_LOG_WARN("%s: no pending KV cache updates to commit - might indicate a bug (ref: %s)\n",
-                __func__, "https://github.com/ggml-org/llama.cpp/pull/12695");
+    if (recovery.cells.empty()) {
+        LLAMA_LOG_WARN("%s: the recovery information upon a commit was empty - might indicate a bug (ref: %s)\n",
+                __func__, "https://github.com/ggml-org/llama.cpp/pull/13194");
         return;
     }
 
-    pending.clear();
+    recovery.clear();
 }
 
 bool llama_kv_cache_unified::update(llama_context & lctx) {
@@ -460,16 +447,11 @@ void llama_kv_cache_unified::set_full() {
     head = 0;
 }
 
-llama_sbatch llama_kv_cache_unified::sbatch_init(
-        const llama_batch & batch,
-        bool logits_all) {
+llama_sbatch llama_kv_cache_unified::sbatch_init(const llama_batch & batch, bool logits_all) {
     return llama_sbatch(batch, hparams.n_embd, true, logits_all);
 }
 
-llama_ubatch llama_kv_cache_unified::ubatch_next(
-        llama_sbatch & sbatch,
-        uint32_t n_ubatch,
-        bool embd_pooled) const {
+llama_ubatch llama_kv_cache_unified::ubatch_next(llama_sbatch & sbatch, uint32_t n_ubatch, bool embd_pooled) const {
     GGML_UNUSED(embd_pooled);
     return sbatch.split_simple(n_ubatch);
 }
@@ -490,6 +472,29 @@ bool llama_kv_cache_unified::find_slot(const llama_ubatch & ubatch) {
         return false;
     }
 
+//#define FIND_SLOT_DEBUG 1
+#if FIND_SLOT_DEBUG
+    LLAMA_LOG_WARN("begin: n = %5d, used = %5d, head = %5d, n_swa = %5d\n", n, used, head, n_swa);
+
+    // for debugging
+    {
+        std::string ss;
+        if (n_swa > 0) {
+            for (uint32_t i = 0; i < size; ++i) {
+                if (cells[i].pos == -1) {
+                    ss += '.';
+                } else {
+                    ss += std::to_string(*cells[i].seq_id.begin());
+                }
+                if (i%256 == 255) {
+                    ss += '\n';
+                }
+            }
+        }
+        LLAMA_LOG_WARN("\n%s\n", ss.c_str());
+    }
+#endif
+
     uint32_t n_tested = 0;
 
     while (true) {
@@ -520,6 +525,11 @@ bool llama_kv_cache_unified::find_slot(const llama_ubatch & ubatch) {
     }
 
     for (uint32_t i = 0; i < n_tokens; ++i) {
+        // remember the original state
+        if (recovery.cells.find(head + i) == recovery.cells.end()) {
+            recovery.cells[head + i] = cells[head + i];
+        }
+
         cells[head + i].pos = ubatch.pos[i];
 
         for (int32_t j = 0; j < ubatch.n_seq_id[i]; j++) {
@@ -529,14 +539,14 @@ bool llama_kv_cache_unified::find_slot(const llama_ubatch & ubatch) {
 
     used += n_tokens;
 
-    pending.ubatches.push_back({ head, ubatch });
-
     // a heuristic, to avoid attending the full cache if it is not yet utilized
     // after enough generations, the benefit from this heuristic disappears
     // if we start defragmenting the cache, the benefit from this will be more important
     n = std::min(size, std::max(padding, GGML_PAD(cell_max(), padding)));
 
-    //printf("n = %5d, used = %5d, head = %5d\n", n, used, head);
+#ifdef FIND_SLOT_DEBUG
+    LLAMA_LOG_WARN("end:   n = %5d, used = %5d, head = %5d, n_swa = %5d\n", n, used, head, n_swa);
+#endif
 
     return true;
 }
@@ -642,6 +652,34 @@ ggml_tensor * llama_kv_cache_unified::cpy_v(ggml_context * ctx, ggml_tensor * v_
     return ggml_cpy(ctx, v_cur, v_view);
 }
 
+void llama_kv_cache_unified::prune_swa(llama_seq_id seq_id, llama_pos p1) {
+    // no pruning is needed when the cache does not use SWA
+    GGML_ASSERT(swa_type != LLAMA_SWA_TYPE_NONE && "do not prune non-SWA cache");
+
+    for (uint32_t i = 0; i < size; ++i) {
+        const llama_pos p0 = cells[i].pos;
+
+        if (is_masked_swa(p0, p1)) {
+            if (seq_id < 0) {
+                cells[i].seq_id.clear();
+            } else if (cells[i].has_seq_id(seq_id)) {
+                cells[i].seq_id.erase(seq_id);
+            } else {
+                continue;
+            }
+
+            if (cells[i].is_empty()) {
+                // keep count of the number of used cells
+                if (cells[i].pos >= 0) {
+                    used--;
+                }
+
+                cells[i].pos = -1;
+            }
+        }
+    }
+}
+
 void llama_kv_cache_unified::set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const {
     const int64_t n_tokens     = ubatch->n_tokens;
     const int64_t n_seq_tokens = ubatch->n_seq_tokens;
@@ -1181,6 +1219,10 @@ uint32_t llama_kv_cache_unified::cell_max() const {
 }
 
 bool llama_kv_cache_unified::is_masked_swa(llama_pos p0, llama_pos p1) const {
+    if (p0 < 0) {
+        return true;
+    }
+
     switch (swa_type) {
         case LLAMA_SWA_TYPE_NONE:
             {
@@ -1589,13 +1631,13 @@ llama_kv_cache_unified_iswa::llama_kv_cache_unified_iswa(
                      bool   offload,
                  uint32_t   kv_size,
                  uint32_t   n_seq_max,
-                 uint32_t   n_batch,
+                 uint32_t   n_ubatch,
                  uint32_t   padding) : hparams(model.hparams) {
     llama_kv_cache_unified::layer_filter_cb filter_base = [&](int32_t il) { return !model.hparams.is_swa(il); };
     llama_kv_cache_unified::layer_filter_cb filter_swa  = [&](int32_t il) { return  model.hparams.is_swa(il); };
 
     const uint32_t kv_size_base = kv_size;
-    const uint32_t kv_size_swa  = std::min(kv_size, GGML_PAD(hparams.n_swa*n_seq_max + n_batch, padding));
+    const uint32_t kv_size_swa  = std::min(kv_size, GGML_PAD(hparams.n_swa*n_seq_max + n_ubatch, padding));
 
     LLAMA_LOG_INFO("%s: creating non-SWA KV cache, size = %u cells\n", __func__, kv_size_base);
 
@@ -1659,20 +1701,12 @@ void llama_kv_cache_unified_iswa::commit() {
     kv_base->commit();
     kv_swa ->commit();
 
-    if (pending.pos_max.empty()) {
-        return;
-    }
-
     // slide the attention window, forgetting/pruning old tokens that are outside the window
     for (const auto & [seq_id, pos_max] : pending.pos_max) {
-        if (pos_max <= (llama_pos) hparams.n_swa) {
-            continue;
-        }
-
-        kv_swa->seq_rm(seq_id, -1, pos_max - hparams.n_swa + 1);
+        kv_swa->prune_swa(seq_id, pos_max);
     }
 
-    pending.pos_max.clear();
+    pending.clear();
 }
 
 bool llama_kv_cache_unified_iswa::update(llama_context & lctx) {
@@ -1695,12 +1729,18 @@ void llama_kv_cache_unified_iswa::set_full() {
 }
 
 llama_sbatch llama_kv_cache_unified_iswa::sbatch_init(const llama_batch & batch, bool logits_all) {
+    pending.pos_max.clear();
+
     for (int i = 0; i < batch.n_tokens; ++i) {
         for (int s = 0; s < batch.n_seq_id[i]; ++s) {
             const llama_seq_id seq_id = batch.seq_id[i][s];
             const llama_pos    pos    = batch.pos[i];
 
-            pending.pos_max[seq_id] = std::max(pending.pos_max[seq_id], pos);
+            if (pending.pos_max.find(seq_id) == pending.pos_max.end()) {
+                pending.pos_max[seq_id] = pos;
+            } else {
+                pending.pos_max[seq_id] = std::max(pending.pos_max[seq_id], pos);
+            }
         }
     }
 
diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h
@@ -2,18 +2,19 @@
 
 #include "llama.h"
 #include "llama-io.h"
-#include "llama-batch.h"
 #include "llama-graph.h"
 #include "llama-memory.h"
 
 #include "ggml-cpp.h"
 
-#include <map>
 #include <set>
+#include <unordered_map>
 #include <vector>
 
 struct llama_cparams;
 struct llama_hparams;
+struct llama_ubatch;
+struct llama_sbatch;
 struct llama_model;
 struct llama_context;
 
@@ -171,6 +172,8 @@ class llama_kv_cache_unified : public llama_kv_cache {
     ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, int32_t il) const;
     ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, int32_t il) const;
 
+    void prune_swa(llama_seq_id seq_id, llama_pos p1);
+
     void set_input_kq_mask   (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const;
     void set_input_k_shift   (ggml_tensor * dst) const;
     void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const;
@@ -214,7 +217,7 @@ class llama_kv_cache_unified : public llama_kv_cache {
 
     uint32_t head = 0; // the location where the batch will be placed in the cache (see find_slot())
     uint32_t size = 0; // total number of cells, shared across all sequences
-    uint32_t used = 0; // used cells (i.e. at least one seq_id)
+    uint32_t used = 0; // used cells (i.e. at least one seq_id) (TODO: add `struct kv_cells` and keep track automaticallt)
 
     // computed before each graph build
     uint32_t n = 0;
@@ -233,27 +236,20 @@ class llama_kv_cache_unified : public llama_kv_cache {
     std::vector<ggml_context_ptr>        ctxs;
     std::vector<ggml_backend_buffer_ptr> bufs;
 
-    std::vector<kv_cell>  cells;
+    std::vector<kv_cell>  cells;  // TODO: replace with `struct kv_cells`
     std::vector<kv_layer> layers;
 
     // model layer id -> KV cache layer id
-    std::map<int32_t, int32_t> map_layer_ids;
-
-    struct ubatch_info {
-        uint32_t head;
+    std::unordered_map<int32_t, int32_t> map_layer_ids;
 
-        llama_ubatch data;
-    };
-
-    // pending cell updates that are not yet committed
+    // recovery information used to restore the KV cells to their original state in case of a failure
     struct {
         void clear() {
-            ubatches.clear();
+            cells.clear();
         }
 
-        // upon batch processing failure, we revert these ubatches from the KV cells
-        std::vector<ubatch_info> ubatches;
-    } pending;
+        std::unordered_map<uint32_t, kv_cell> cells;
+    } recovery;
 
     // defrag
     struct {
@@ -317,7 +313,7 @@ class llama_kv_cache_unified_iswa : public llama_kv_cache {
                          bool   offload,
                      uint32_t   kv_size,
                      uint32_t   n_seq_max,
-                     uint32_t   n_batch,
+                     uint32_t   n_ubatch,
                      uint32_t   padding);
 
     ~llama_kv_cache_unified_iswa() = default;
@@ -377,9 +373,12 @@ class llama_kv_cache_unified_iswa : public llama_kv_cache {
 private:
     const llama_hparams & hparams;
 
-    // pending cell updates that are not yet committed
     struct {
-        std::map<llama_seq_id, llama_pos> pos_max;
+        void clear() {
+            pos_max.clear();
+        }
+
+        std::unordered_map<llama_seq_id, llama_pos> pos_max;
     } pending;
 
     std::unique_ptr<llama_kv_cache_unified> kv_base;
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
@@ -13228,7 +13228,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                             cparams.offload_kqv,
                             cparams.n_ctx,
                             cparams.n_seq_max,
-                            cparams.n_batch,
+                            cparams.n_ubatch,
                             padding);
                 } else {
                     res = new llama_kv_cache_unified(
