kv : change structures in kv-cells for improve operations asymptotic

ggml/src/ggml-impl.h

@@ -11,6 +11,9 @@
 #include <stdbool.h>
 #include <stdint.h>
 #include <string.h>
+#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__)
+    #include <immintrin.h>
+#endif
 
 #ifdef __ARM_FEATURE_SVE
 #include <arm_sve.h>
@@ -382,6 +385,12 @@ static inline uint32_t fp32_to_bits(float f) {
 }
 
 static inline float ggml_compute_fp16_to_fp32(ggml_fp16_t h) {
+#ifdef __F16C__
+    return _cvtsh_ss(h);
+#elif defined(__aarch64__) && defined(__ARM_FP) && (__ARM_FP & 2)
+    union { uint16_t u; __fp16 f; } u = { .u = h };
+    return (float)u.f;
+#else
     const uint32_t w = (uint32_t) h << 16;
     const uint32_t sign = w & UINT32_C(0x80000000);
     const uint32_t two_w = w + w;
@@ -402,6 +411,7 @@ static inline float ggml_compute_fp16_to_fp32(ggml_fp16_t h) {
     const uint32_t result = sign |
         (two_w < denormalized_cutoff ? fp32_to_bits(denormalized_value) : fp32_to_bits(normalized_value));
     return fp32_from_bits(result);
+#endif
 }
 
 static inline ggml_fp16_t ggml_compute_fp32_to_fp16(float f) {

src/CMakeLists.txt

@@ -27,6 +27,7 @@ add_library(llama
             llama-kv-cache-dsa.cpp
             llama-kv-cache-msa.cpp
             llama-kv-cache-dsv4.cpp
+            llama-kv-cells.cpp
             llama-memory.cpp
             llama-memory-hybrid.cpp
             llama-memory-hybrid-iswa.cpp

src/llama-batch.cpp

@@ -757,7 +757,6 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
     const int64_t n_pos_all  =              (int64_t) n_tokens*n_pos_per_embd;
 
     udata->token     .resize(n_tokens);
-    udata->embd      .resize(n_embd_all);
     udata->pos       .resize(n_pos_all);
     udata->n_seq_id  .resize(n_tokens);
     udata->seq_id    .resize(n_tokens);
@@ -766,6 +765,12 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
     udata->output    .resize(n_tokens);
 
     udata->seq_id_data.reserve(n_tokens);
+    if (batch.embd) {
+        udata->embd.clear();
+        udata->embd.reserve(n_embd_all);
+    } else {
+        udata->embd.resize(n_embd_all); // fill all size..new_size elems by 0.0f
+    }
 
     seq_set_t seq_set_unq;
 
@@ -775,7 +780,10 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
         }
 
         if (batch.embd) {
-            memcpy(udata->embd.data() + i*n_embd, batch.embd + (int64_t) idxs[i]*n_embd, n_embd*sizeof(float));
+            auto src = batch.embd + (int64_t) idxs[1i] * n_embd;
+            // use safe method for auto increase size
+            // next improvements - write own vector without automatic filling float)
+            udata->embd.insert(udata->embd.end(), src, src + n_embd);
         }
 
         for (size_t j = 0; j < (size_t)n_pos_per_embd; ++j) {

src/llama-kv-cache.cpp

@@ -396,19 +396,7 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
         auto & cells = v_cells[seq_to_stream[seq_id]];
         auto & head  = v_heads[seq_to_stream[seq_id]];
 
-        uint32_t new_head = cells.size();
-
-        for (uint32_t i = 0; i < cells.size(); ++i) {
-            if (!cells.pos_in(i, p0, p1)) {
-                continue;
-            }
-
-            if (cells.seq_has(i, seq_id) && cells.seq_rm(i, seq_id)) {
-                if (new_head == cells.size()) {
-                    new_head = i;
-                }
-            }
-        }
+        uint32_t new_head = cells.next_head(seq_id, p0, p1);
 
         // If we freed up a slot, set head to it so searching can start there.
         if (new_head != cells.size() && new_head < head) {
@@ -1121,7 +1109,7 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
 
                 seq_pos_max_rm[seq_id] = std::max(seq_pos_max_rm[seq_id], pos);
 
-                cells.rm(idx);
+                cells.rm_single(idx, seq_id);
             }
 
             cells.pos_set(idx, ubatch.pos[i]);
@@ -1133,10 +1121,7 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
                 };
                 cells.ext_set(idx, ext);
             }
-
-            for (int32_t s = 0; s < ubatch.n_seq_id[i]; s++) {
-                cells.seq_add(idx, ubatch.seq_id[i][s]);
-            }
+            cells.add_sequences(idx, ubatch.n_seq_id[i], ubatch.seq_id[i]);
         }
     }
 
@@ -1157,6 +1142,8 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
                     __func__, cells.seq_pos_min(s), seq_pos_max_rm[s], s);
 
             seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1);
+        } else {
+            cells.compact(s); // compact after seq_rm
         }
     }
 

src/llama-kv-cells.cpp

@@ -0,0 +1,180 @@
+#include "llama-kv-cache.h"
+#include <cstdint>
+
+void llama_kv_cells::add_sequences(uint32_t i, int32_t n, llama_seq_id *_seq) {
+    llama_seq_id seq_id{};
+    assert(i < pos.size());
+    assert(pos[i] != -1);
+    for (int32_t s = 0; s < n; s++) {
+        seq_id = _seq[s];
+        assert(!seq[i].test(seq_id));
+        seq[i].set(seq_id);
+        seq_pos_inc(seq_id, pos[i]);
+    }
+}
+
+void llama_kv_cells::compact(llama_seq_id s) {
+    auto & v = seq_pos[s];
+
+    if (v.total == 0) {
+        v.clear();
+        return;
+    }
+
+    const uint32_t h = v.head;
+    const uint32_t t = v.tail;
+
+    if (t + 1 < v.cnt.size()) {
+        v.cnt.resize(t + 1);
+    }
+
+    if (h == 0 && t + 1 == v.cnt.size()) {
+        return;
+    }
+
+    if (h > 0 && h > v.cnt.size() / 4) {
+        v.cnt.erase(v.cnt.begin(), v.cnt.begin() + h);
+        v.base += (llama_pos)h;
+        v.tail  = t - h;
+        v.head  = 0;
+    }
+}
+
+uint32_t llama_kv_cells::next_head(int32_t seq_id, llama_pos p0, llama_pos p1) {
+    uint32_t new_head = size();
+
+    for (size_t w = 0; w < used_bits.size(); ++w) {
+        uint64_t mask = used_bits[w];
+        while (mask) {
+            const int bit = llama_bits::countr_zero64(mask);
+            const uint32_t i = (uint32_t)(w * 64 + bit);
+            mask &= mask - 1;
+
+            const llama_pos p = pos[i];
+            if (p < p0 || p >= p1) continue;
+
+            if (seq_has(i, seq_id) && seq_rm(i, seq_id)) {
+                if (new_head == size()) {
+                    new_head = i;
+                }
+            }
+        }
+    }
+    compact(seq_id);
+    return new_head;
+}
+
+void llama_kv_cells::set(const std::vector<uint32_t> & idxs, const llama_kv_cells & other) {
+    assert(idxs.size() == other.pos.size());
+
+    for (uint32_t j = 0; j < other.pos.size(); ++j) {
+        const auto idx = idxs[j];
+
+        if (pos[idx] == other.pos[j] && seq[idx] == other.seq[j]) {
+            ext[idx] = other.ext[j];
+            assert(shift[idx] == 0);
+            continue;
+        }
+
+        if (pos[idx] == -1 && other.pos[j] != -1) {
+            used_insert(idx);
+        }
+        if (pos[idx] != -1 && other.pos[j] == -1) {
+            used_erase(idx);
+        }
+        if (pos[idx] != -1) {
+            seq_pos_rm(idx);
+        }
+
+        pos[idx] = other.pos[j];
+        ext[idx] = other.ext[j];
+        seq[idx] = other.seq[j];
+
+        if (pos[idx] != -1) {
+            seq_pos_add(idx);
+        }
+
+        assert(shift[idx] == 0);
+    }
+}
+
+void llama_kv_cells::seq_pos_dec(llama_seq_id s, llama_pos p) {
+    auto & v = seq_pos[s];
+
+    assert(v.total > 0);
+    const uint32_t idx = (uint32_t) (p - v.base);
+    assert(idx < v.cnt.size() && v.cnt[idx] > 0);
+
+    --v.cnt[idx];
+    --v.total;
+
+    if (v.total == 0) {
+        v.clear();
+        return;
+    }
+
+    if (idx == v.head) {
+        while (v.cnt[v.head] == 0) {
+            ++v.head;
+        }
+    } else if (idx == v.tail) {
+        while (v.cnt[v.tail] == 0) {
+            --v.tail;
+        }
+    }
+}
+
+void llama_kv_cells::seq_pos_inc(llama_seq_id s, llama_pos p) {
+    auto & v = seq_pos[s];
+
+    if (v.total == 0) {
+        v.base = p;
+        v.cnt.assign(1, 1);
+        v.head = v.tail = 0;
+        v.total         = 1;
+        return;
+    }
+
+    if (p >= v.base) {
+        const uint32_t idx = (uint32_t) (p - v.base);
+        if (idx >= v.cnt.size()) {
+            v.cnt.resize(idx + 1, 0);
+        }
+        if (++v.cnt[idx] == 1) {
+            if (idx < v.head) {
+                v.head = idx;
+            }
+            if (idx > v.tail) {
+                v.tail = idx;
+            }
+        }
+    } else {
+        // rary
+        const uint32_t pre = (uint32_t) (v.base - p);
+        v.cnt.insert(v.cnt.begin(), pre, 0);
+        v.cnt[0] = 1;
+        v.base   = p;
+        v.head   = 0;
+        v.tail += pre;
+    }
+
+    ++v.total;
+}
+
+void llama_kv_cells::rm_single(uint32_t i, llama_seq_id seq_id) {  // need compact after some seq_pos_dec
+    assert(i < pos.size());
+    assert(pos[i] != -1);
+    assert(seq[i].count() == 1);
+    assert(seq[i].test(seq_id));
+
+    seq_pos_dec(seq_id, pos[i]);
+
+    seq[i].reset();
+
+    pos[i] = -1;
+    ext[i].reset();
+    shift[i] = 0;
+
+    used_erase(i);
+}
+