kv-cache : hide padding details in the implementation

ggml-ci

Author: ggerganov

src/llama-context.cpp

@@ -1285,24 +1285,10 @@ int llama_context::decode(llama_batch & inp_batch) {
         }
 
         // find KV slot
-        {
-            if (!kv_self->find_slot(ubatch)) {
-                LLAMA_LOG_WARN("%s: failed to find KV cache slot for ubatch of size %d\n", __func__, ubatch.n_tokens);
-
-                return 1;
-            }
+        if (!kv_self->find_slot(ubatch)) {
+            LLAMA_LOG_WARN("%s: failed to find KV cache slot for ubatch of size %d\n", __func__, ubatch.n_tokens);
 
-            if (!is_recurrent) {
-                auto * kv = static_cast<llama_kv_cache_unified *>(kv_self);
-
-                // 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
-                const uint32_t pad = kv->get_padding(cparams);
-                kv->n = std::min(kv->size, std::max(pad, GGML_PAD(kv->cell_max(), pad)));
-
-                //printf("kv.n = %5d, kv.used = %5d, kv.head = %5d\n", kv->n, kv->used, kv->head);
-            }
+            return 1;
         }
 
         ggml_backend_sched_reset(sched.get());

src/llama-kv-cache.cpp

@@ -21,14 +21,17 @@ llama_kv_cache_unified::llama_kv_cache_unified(
                   ggml_type   type_k,
                   ggml_type   type_v,
                        bool   v_trans,
-                   uint32_t   kv_size) : hparams(hparams), cbs(std::move(cbs)), v_trans(v_trans) {
+                   uint32_t   kv_size,
+                   uint32_t   padding) : hparams(hparams), cbs(std::move(cbs)), v_trans(v_trans), padding(padding) {
     const int32_t n_layer = hparams.n_layer;
 
     has_shift = false;
     can_shift = true;
 
-    LLAMA_LOG_INFO("%s: kv_size = %d, type_k = '%s', type_v = '%s', n_layer = %d, can_shift = %d\n",
-            __func__, kv_size, ggml_type_name(type_k), ggml_type_name(type_v), n_layer, can_shift);
+    LLAMA_LOG_INFO("%s: kv_size = %d, type_k = '%s', type_v = '%s', n_layer = %d, can_shift = %d, padding = %d\n",
+            __func__, kv_size, ggml_type_name(type_k), ggml_type_name(type_v), n_layer, can_shift, padding);
+
+    GGML_ASSERT(kv_size % padding == 0 && "kv_size must be a multiple of padding");
 
     head = 0;
     size = kv_size;
@@ -463,6 +466,13 @@ bool llama_kv_cache_unified::find_slot(
 
     pending.ranges.push_back({head, head + n_tokens});
 
+    // 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);
+
     return true;
 }
 

src/llama-kv-cache.h

@@ -101,7 +101,8 @@ class llama_kv_cache_unified : public llama_kv_cache {
                       ggml_type   type_k,
                       ggml_type   type_v,
                            bool   v_trans,
-                       uint32_t   kv_size);
+                       uint32_t   kv_size,
+                       uint32_t   padding);
 
     ~llama_kv_cache_unified() = default;
 
@@ -196,6 +197,9 @@ class llama_kv_cache_unified : public llama_kv_cache {
     // computed before each graph build
     uint32_t n = 0;
 
+    // required padding
+    uint32_t padding = 1;
+
     std::vector<llama_kv_cell> cells;
 
     std::vector<ggml_tensor *> k_l; // per layer

src/llama-model.cpp

@@ -12808,7 +12808,9 @@ llama_memory_i * llama_model::create_memory(llama_cparams & cparams, const llama
             } break;
         default:
             {
-                cparams.n_ctx = GGML_PAD(cparams.n_ctx, llama_kv_cache_unified::get_padding(cparams));
+                const auto padding = llama_kv_cache_unified::get_padding(cparams);
+
+                cparams.n_ctx = GGML_PAD(cparams.n_ctx, padding);
 
                 LLAMA_LOG_DEBUG("%s: n_ctx = %u (padded)\n", __func__, cparams.n_ctx);
 
@@ -12832,7 +12834,8 @@ llama_memory_i * llama_model::create_memory(llama_cparams & cparams, const llama
                         params.type_k,
                         params.type_v,
                         !cparams.flash_attn,
-                        cparams.n_ctx);
+                        cparams.n_ctx,
+                        padding);
             }
     }