diff --git a/ggml/src/ggml-impl.h b/ggml/src/ggml-impl.h index 62b76abbcec..7189f9539f0 100644 --- a/ggml/src/ggml-impl.h +++ b/ggml/src/ggml-impl.h @@ -11,6 +11,9 @@ #include #include #include +#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) + #include +#endif #ifdef __ARM_FEATURE_SVE #include @@ -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) { diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 39ba3061f70..f0e381e955d 100644 --- a/src/CMakeLists.txt +++ b/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 diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 2b98a552f48..cd66a58aecd 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -757,7 +757,6 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & 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 & 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 & 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) { diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 5382cd7266f..39969a92c69 100644 --- a/src/llama-kv-cache.cpp +++ b/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 } } diff --git a/src/llama-kv-cells.cpp b/src/llama-kv-cells.cpp new file mode 100644 index 00000000000..a58567f12ae --- /dev/null +++ b/src/llama-kv-cells.cpp @@ -0,0 +1,180 @@ +#include "llama-kv-cache.h" +#include + +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 & 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); +} + diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index fddd31a0b21..9a4e265494a 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -3,13 +3,86 @@ #include "llama.h" #include "llama-cparams.h" -#include +#include #include +#include #include -#include -#include #include +#if defined(_MSC_VER) +# include +#endif +#if defined(__cpp_lib_bitops) && __cpp_lib_bitops >= 201907L +# define LLAMA_HAS_STD_BITOPS 1 +#elif defined(__has_include) && __has_include() && __cplusplus >= 202002L +# define LLAMA_HAS_STD_BITOPS 1 +#else +# define LLAMA_HAS_STD_BITOPS 0 +#endif + +namespace llama_bits { +inline int popcount64(uint64_t x) { +#if LLAMA_HAS_STD_BITOPS + return std::popcount(x); +#elif defined(__GNUC__) || defined(__clang__) + return __builtin_popcountll(x); +#elif defined(_MSC_VER) + return (int)__popcnt64(x); +#else + // Hacker's Delight + x = x - ((x >> 1) & UINT64_C(0x5555555555555555)); + x = (x & UINT64_C(0x3333333333333333)) + ((x >> 2) & UINT64_C(0x3333333333333333)); + x = (x + (x >> 4)) & UINT64_C(0x0F0F0F0F0F0F0F0F); + return (int)((x * UINT64_C(0x0101010101010101)) >> 56); +#endif +} + +inline int countr_zero64(uint64_t x) { + assert(x != 0); +#if LLAMA_HAS_STD_BITOPS + return std::countr_zero(x); +#elif defined(__GNUC__) || defined(__clang__) + return __builtin_ctzll(x); +#elif defined(_MSC_VER) + unsigned long idx; + _BitScanForward64(&idx, x); + return (int)idx; +#else + int r = 63; + if (x & UINT64_C(0x00000000FFFFFFFF)) r -= 32; else x >>= 32; + if (x & UINT64_C(0x000000000000FFFF)) r -= 16; else x >>= 16; + if (x & UINT64_C(0x00000000000000FF)) r -= 8; else x >>= 8; + if (x & UINT64_C(0x000000000000000F)) r -= 4; else x >>= 4; + if (x & UINT64_C(0x0000000000000003)) r -= 2; else x >>= 2; + if (x & UINT64_C(0x0000000000000001)) r -= 1; + return r; +#endif +} + +inline int countl_zero64(uint64_t x) { + assert(x != 0); +#if LLAMA_HAS_STD_BITOPS + return std::countl_zero(x); +#elif defined(__GNUC__) || defined(__clang__) + return __builtin_clzll(x); +#elif defined(_MSC_VER) + unsigned long idx; + _BitScanReverse64(&idx, x); + return 63 - (int)idx; +#else + int r = 0; + if (!(x & UINT64_C(0xFFFFFFFF00000000))) { r += 32; x <<= 32; } + if (!(x & UINT64_C(0xFFFF000000000000))) { r += 16; x <<= 16; } + if (!(x & UINT64_C(0xFF00000000000000))) { r += 8; x <<= 8; } + if (!(x & UINT64_C(0xF000000000000000))) { r += 4; x <<= 4; } + if (!(x & UINT64_C(0xC000000000000000))) { r += 2; x <<= 2; } + if (!(x & UINT64_C(0x8000000000000000))) { r += 1; } + return r; +#endif +} + +} // namespace llama_bits + struct llama_kv_cell_ext { // 2D spatial positions, typically used for M-RoPE llama_pos x = 0; @@ -41,7 +114,8 @@ class llama_kv_cells { has_shift = false; - used.clear(); + std::fill(used_bits.begin(), used_bits.end(), 0); + used_cnt = 0; for (uint32_t s = 0; s < LLAMA_MAX_SEQ; ++s) { seq_pos[s].clear(); @@ -66,6 +140,9 @@ class llama_kv_cells { shift.resize(n); seq.resize(n); + used_bits.assign((n + 63) / 64, 0); + used_cnt = 0; + reset(); } @@ -77,45 +154,35 @@ class llama_kv_cells { } uint32_t get_used() const { - return used.size(); + return used_cnt; } // the index of the first cell that is used // return 0 if no cells are used uint32_t used_min() const { - return used.empty() ? 0 : *used.begin(); + for (size_t w = 0; w < used_bits.size(); ++w) { + if (used_bits[w]) { + return (uint32_t)(w * 64 + llama_bits::countr_zero64(used_bits[w])); + } + } + return 0; } // the index of the last cell that is used + 1 // return 0 if no cells are used uint32_t used_max_p1() const { - return used.empty() ? 0 : *used.rbegin() + 1; + for (size_t w = used_bits.size(); w-- > 0;) { + if (used_bits[w]) { + return (uint32_t)(w * 64 + 64 - llama_bits::countl_zero64(used_bits[w])); + } + } + return 0; } bool get_has_shift() const { return has_shift; } - // move cell isrc to idst (used during defrag) - //void mv(uint32_t isrc, uint32_t idst) { - // assert(isrc < pos.size()); - // assert(idst < pos.size()); - - // assert(pos[idst] == -1); - // assert(pos[isrc] != -1); - - // pos [idst] = pos [isrc]; - // shift[idst] = shift[isrc]; - // seq [idst] = seq [isrc]; - - // pos [isrc] = -1; - // shift[isrc] = 0; - // seq [isrc].reset(); - - // used.erase (isrc); - // used.insert(idst); - //} - // copy the state of cells [i, i + n) (used for save/restore the state of the cells) llama_kv_cells cp(uint32_t i, uint32_t n) const { assert(i + n <= pos.size()); @@ -156,67 +223,8 @@ class llama_kv_cells { return res; } - // set the state of cells [i, i + other.pos.size()) (used for save/restore the state of the cells) - void set(uint32_t i, const llama_kv_cells & other) { - assert(i + other.pos.size() <= pos.size()); - - for (uint32_t j = 0; j < other.pos.size(); ++j) { - const auto idx = i + j; - - if (pos[idx] == -1 && other.pos[j] != -1) { - used.insert(i + j); - } - - if (pos[idx] != -1 && other.pos[j] == -1) { - used.erase(i + j); - } - - if (pos[idx] != -1) { - seq_pos_rm(i + j); - } - - pos[idx] = other.pos[j]; - ext[idx] = other.ext[j]; - seq[idx] = other.seq[j]; - - if (pos[idx] != -1) { - seq_pos_add(i + j); - } - - assert(shift[idx] == 0); - } - } - // set the state of cells [idxs[0], idxs[1], ..., idxs[idxs.size() - 1]) - void set(const std::vector & 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] == -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 set(const std::vector & idxs, const llama_kv_cells & other); // clear a non-empty cell void rm(uint32_t i) { @@ -230,12 +238,14 @@ class llama_kv_cells { ext[i].reset(); shift[i] = 0; - used.erase(i); + used_erase(i); } + void rm_single(uint32_t i, llama_seq_id seq_id); + // note: call only if the cell has seq_id // return true if the cell becomes empty - bool seq_rm(uint32_t i, llama_seq_id seq_id) { + bool seq_rm(uint32_t i, llama_seq_id seq_id) { // need compact after some seq_rm assert(i < pos.size()); assert(seq[i].test(seq_id)); assert(pos[i] != -1); @@ -249,7 +259,7 @@ class llama_kv_cells { ext[i].reset(); shift[i] = 0; - used.erase(i); + used_erase(i); return true; } @@ -279,7 +289,7 @@ class llama_kv_cells { ext[i].reset(); shift[i] = 0; - used.erase(i); + used_erase(i); return true; } @@ -319,13 +329,9 @@ class llama_kv_cells { // note: call only for cells with exactly one sequence llama_seq_id seq_get(uint32_t i) const { assert(seq[i].count() == 1); - - for (int s = 0; s < LLAMA_MAX_SEQ; ++s) { - if (seq[i].test(s)) { - return s; - } + for (int s = 0; s < N_SEQ_WORDS; ++s) { + if (seq[i].w[s]) return s * 64 + llama_bits::countr_zero64(seq[i].w[s]); } - return -1; } @@ -334,14 +340,8 @@ class llama_kv_cells { llama_pos seq_pos_min(llama_seq_id seq_id) const { assert(seq_id >= 0); assert(seq_id < LLAMA_MAX_SEQ); - - if (seq_pos[seq_id].empty()) { - return -1; - } - - assert(seq_pos[seq_id].begin()->second > 0); - - return seq_pos[seq_id].begin()->first; + const auto & v = seq_pos[seq_id]; + return v.total > 0 ? v.min() : -1; } // the maximum position of sequence seq_id currently present in any of the cells @@ -349,14 +349,8 @@ class llama_kv_cells { llama_pos seq_pos_max(llama_seq_id seq_id) const { assert(seq_id >= 0); assert(seq_id < LLAMA_MAX_SEQ); - - if (seq_pos[seq_id].empty()) { - return -1; - } - - assert(seq_pos[seq_id].rbegin()->second > 0); - - return seq_pos[seq_id].rbegin()->first; + const auto & v = seq_pos[seq_id]; + return v.total > 0 ? v.max() : -1; } // note: call only if the cell is not empty @@ -399,7 +393,7 @@ class llama_kv_cells { pos[i] = p; - used.insert(i); + used_insert(i); } void ext_set(uint32_t i, llama_kv_cell_ext p) { @@ -426,7 +420,7 @@ class llama_kv_cells { pos[i] = -1; shift[i] = 0; - used.erase(i); + used_erase(i); return true; } @@ -455,11 +449,16 @@ class llama_kv_cells { has_shift = true; } -private: + const llama_pos * pos_data() const { return pos.data(); } + void add_sequences(uint32_t i, int32_t n, llama_seq_id *_seq); + void compact(llama_seq_id s); + uint32_t next_head(int32_t seq_id, llama_pos p0, llama_pos p1); + private: bool has_shift = false; // set of indices of used cells (i.e. pos[i] != -1, allowed to not have any seq_id) - std::set used; + std::vector used_bits; + uint32_t used_cnt = 0; std::vector pos; @@ -483,50 +482,81 @@ class llama_kv_cells { // std::vector shift; - using seq_set_t = std::bitset; + static_assert(LLAMA_MAX_SEQ > 0 && (LLAMA_MAX_SEQ % 64) == 0, + "LLAMA_MAX_SEQ must be a multiple of 64"); + static constexpr int N_SEQ_WORDS = LLAMA_MAX_SEQ / 64; + + struct seq_set_t { + uint64_t w[N_SEQ_WORDS]{}; // zero-init + + void reset() { for (auto & x : w) x = 0; } + void reset(int s) { w[s >> 6] &= ~(1ull << (s & 63)); } + void set(int s) { w[s >> 6] |= 1ull << (s & 63); } + bool test(int s) const { return (w[s >> 6] >> (s & 63)) & 1; } + bool none() const { for (auto x : w) if (x) return false; return true; } + bool any() const { return !none(); } + int count() const { int c = 0; for (auto x : w) c += llama_bits::popcount64(x); return c; } + bool operator==(const seq_set_t & o) const { + for (int k = 0; k < N_SEQ_WORDS; ++k) if (w[k] != o.w[k]) return false; + return true; + } + bool operator!=(const seq_set_t & o) const { return !(*this == o); } + }; - // the bitset seq[i] tells us which sequences are currently occupying the i-th cell std::vector seq; - // the set seq_pos[s][p] tells us how many times the position p is currently present for sequence s - // if the position p is not present, seq_pos[s][p] is not set - // this way seq_pos[s].begin() and seq_pos[s].rbegin() give us the min/max positions currently in the cache - // - // note that we cannot a use an std::set because in some cases a position can occur more than once for the same seq: - // - during performing a cache reuse via (rm + add) - // - some vision models have input embeddings with repeating positions - // - std::map seq_pos[LLAMA_MAX_SEQ]; + struct seq_pos_t { + llama_pos base = 0; + std::vector cnt; + int64_t total = 0; + uint32_t head = 0; + uint32_t tail = 0; - // helper functions for updating `seq_pos`, once cell at a time: + void clear() { base = 0; cnt.clear(); total = 0; head = 0; tail = 0; } + llama_pos min() const { return base + (llama_pos)head; } + llama_pos max() const { return base + (llama_pos)tail; } + }; - void seq_pos_dec(llama_seq_id s, llama_pos p) { - auto it = seq_pos[s].find(p); - assert(it != seq_pos[s].end()); + seq_pos_t seq_pos[LLAMA_MAX_SEQ]; - if (--it->second == 0) { - seq_pos[s].erase(it); + void used_insert(uint32_t i) { + assert(i < pos.size()); + const uint64_t bit = 1ull << (i & 63); + if (!(used_bits[i >> 6] & bit)) { + used_bits[i >> 6] |= bit; + ++used_cnt; } } - void seq_pos_inc(llama_seq_id s, llama_pos p) { - seq_pos[s][p]++; + void used_erase(uint32_t i) { + assert(i < pos.size()); + const uint64_t bit = 1ull << (i & 63); + if (used_bits[i >> 6] & bit) { + used_bits[i >> 6] &= ~bit; + --used_cnt; + } } + // O(1) + void seq_pos_inc(llama_seq_id s, llama_pos p); + + // O(1) amort + void seq_pos_dec(llama_seq_id s, llama_pos p); + // remove cell i void seq_pos_rm(uint32_t i) { - for (int s = 0; s < LLAMA_MAX_SEQ; ++s) { - if (seq[i].test(s)) { - seq_pos_dec(s, pos[i]); + for (int s = 0; s < N_SEQ_WORDS; ++s) { + for (auto m = seq[i].w[s]; m; m &= m - 1) { + seq_pos_dec(s * 64 + llama_bits::countr_zero64(m), pos[i]); } } } // add cell i void seq_pos_add(uint32_t i) { - for (int s = 0; s < LLAMA_MAX_SEQ; ++s) { - if (seq[i].test(s)) { - seq_pos_inc(s, pos[i]); + for (int s = 0; s < N_SEQ_WORDS; ++s) { + for (auto m = seq[i].w[s]; m; m &= m - 1) { + seq_pos_inc(s * 64 + llama_bits::countr_zero64(m), pos[i]); } } } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 419e1eba4c2..02376032e2c 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -158,6 +158,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-grammar-integration.cpp) llama_build_and_test(test-llama-grammar.cpp) llama_build_and_test(test-batch-alloc.cpp) + llama_build_and_test(test-kv-cells.cpp) llama_build_and_test(test-chat.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) target_link_libraries(test-chat PRIVATE server-context) diff --git a/tests/test-kv-cells.cpp b/tests/test-kv-cells.cpp new file mode 100644 index 00000000000..517ebb75a01 --- /dev/null +++ b/tests/test-kv-cells.cpp @@ -0,0 +1,1151 @@ +#include "testing.h" + +#include "llama.h" + +#ifdef NDEBUG +#undef NDEBUG +#endif + +#include "../src/llama-kv-cells.h" + +#include +#include +#include +#include +#include +#include +#include + +static uint64_t ref_popcount(uint64_t x) { + uint64_t c = 0; + while (x) { + x &= x - 1; + ++c; + } + return c; +} + +static int ref_countr_zero(uint64_t x) { + int c = 0; + while ((x & 1) == 0) { + x >>= 1; + ++c; + } + return c; +} + +static int ref_countl_zero(uint64_t x) { + int c = 0; + while ((x & (1ull << 63)) == 0) { + x <<= 1; + ++c; + } + return c; +} + +static void test_bitops(testing & t) { + t.test("popcount64", [&](testing & t) { + t.assert_equal(0, llama_bits::popcount64(0)); + t.assert_equal(1, llama_bits::popcount64(1)); + t.assert_equal(2, llama_bits::popcount64(0x8000000000000001ull)); + t.assert_equal(8, llama_bits::popcount64(0xFF)); + t.assert_equal(64, llama_bits::popcount64(~0ull)); + + for (int i = 0; i < 64; ++i) { + t.assert_equal((int) ref_popcount(1ull << i), llama_bits::popcount64(1ull << i)); + } + + std::mt19937 rng(1); + for (int i = 0; i < 1000; ++i) { + const uint64_t v = ((uint64_t) rng() << 32) ^ rng(); + t.assert_equal((int) ref_popcount(v), llama_bits::popcount64(v)); + } + }); + + t.test("countr_zero64", [&](testing & t) { + t.assert_equal(0, llama_bits::countr_zero64(1)); + t.assert_equal(1, llama_bits::countr_zero64(2)); + t.assert_equal(3, llama_bits::countr_zero64(0x8)); + t.assert_equal(63, llama_bits::countr_zero64(0x8000000000000000ull)); + t.assert_equal(2, llama_bits::countr_zero64(0x4)); + + for (int i = 0; i < 64; ++i) { + t.assert_equal(ref_countr_zero(1ull << i), llama_bits::countr_zero64(1ull << i)); + } + + std::mt19937 rng(2); + for (int i = 0; i < 1000; ++i) { + uint64_t v = ((uint64_t) rng() << 32) ^ rng(); + if (v == 0) { + continue; + } + t.assert_equal(ref_countr_zero(v), llama_bits::countr_zero64(v)); + } + }); + + t.test("countl_zero64", [&](testing & t) { + t.assert_equal(63, llama_bits::countl_zero64(1)); + t.assert_equal(62, llama_bits::countl_zero64(2)); + t.assert_equal(60, llama_bits::countl_zero64(0x8)); + t.assert_equal(0, llama_bits::countl_zero64(0x8000000000000000ull)); + + for (int i = 0; i < 64; ++i) { + t.assert_equal(ref_countl_zero(1ull << i), llama_bits::countl_zero64(1ull << i)); + } + + std::mt19937 rng(3); + for (int i = 0; i < 1000; ++i) { + uint64_t v = ((uint64_t) rng() << 32) ^ rng(); + if (v == 0) { + continue; + } + t.assert_equal(ref_countl_zero(v), llama_bits::countl_zero64(v)); + } + }); +} + +static void test_ext(testing & t) { + t.test("is_2d_gt", [&](testing & t) { + llama_kv_cell_ext a{/*x=*/1, /*y=*/2}; + + // equal positions are not greater + t.assert_true(!a.is_2d_gt(1, 2)); + // equal y: larger x is greater, smaller x is not + t.assert_true(!a.is_2d_gt(3, 2)); + t.assert_true(a.is_2d_gt(0, 2)); + // y dominates x + t.assert_true(!a.is_2d_gt(1, 3)); + t.assert_true(a.is_2d_gt(1, 1)); + t.assert_true(!a.is_2d_gt(0, 3)); + t.assert_true(!a.is_2d_gt(5, 3)); + }); + + t.test("reset", [&](testing & t) { + llama_kv_cell_ext e{/*x=*/7, /*y=*/9}; + e.reset(); + t.assert_equal((llama_pos) 0, e.x); + t.assert_equal((llama_pos) 0, e.y); + }); +} + +static void test_basic(testing & t) { + t.test("resize_and_reset", [&](testing & t) { + llama_kv_cells cells; + cells.resize(10); + + t.assert_equal(10u, cells.size()); + for (uint32_t i = 0; i < 10; ++i) { + t.assert_true(cells.is_empty(i)); + } + t.assert_equal(0u, cells.get_used()); + t.assert_equal(0u, cells.used_min()); + t.assert_equal(0u, cells.used_max_p1()); + t.assert_true(!cells.get_has_shift()); + }); + + t.test("resize_resets_existing", [&](testing & t) { + llama_kv_cells cells; + cells.resize(10); + + cells.pos_set(3, 7); + cells.pos_add(3, 2); // also sets has_shift + t.assert_equal(1u, cells.get_used()); + t.assert_true(cells.get_has_shift()); + + cells.resize(10); + t.assert_equal(10u, cells.size()); + t.assert_true(cells.is_empty(3)); + t.assert_equal(0u, cells.get_used()); + t.assert_true(!cells.get_has_shift()); + }); + + t.test("pos_set_get", [&](testing & t) { + llama_kv_cells cells; + cells.resize(8); + + cells.pos_set(3, 42); + t.assert_equal(42, cells.pos_get(3)); + t.assert_true(!cells.is_empty(3)); + t.assert_equal(1u, cells.get_used()); + t.assert_equal(3u, cells.used_min()); + t.assert_equal(4u, cells.used_max_p1()); + }); + + t.test("used_tracking_noncontiguous", [&](testing & t) { + llama_kv_cells cells; + cells.resize(130); + + cells.pos_set(63, 1); + cells.pos_set(64, 2); + cells.pos_set(128, 3); + + t.assert_equal(3u, cells.get_used()); + t.assert_equal(63u, cells.used_min()); + t.assert_equal(129u, cells.used_max_p1()); + + cells.rm(63); + t.assert_equal(2u, cells.get_used()); + t.assert_equal(64u, cells.used_min()); + t.assert_equal(129u, cells.used_max_p1()); + + cells.rm(128); + t.assert_equal(1u, cells.get_used()); + t.assert_equal(64u, cells.used_min()); + t.assert_equal(65u, cells.used_max_p1()); + + cells.rm(64); + t.assert_equal(0u, cells.get_used()); + t.assert_equal(0u, cells.used_min()); + t.assert_equal(0u, cells.used_max_p1()); + }); + + t.test("pos_in", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + cells.pos_set(1, 5); + + t.assert_true(cells.pos_in(1, 5, 6)); + t.assert_true(cells.pos_in(1, 0, 6)); + t.assert_true(!cells.pos_in(1, 6, 7)); + t.assert_true(!cells.pos_in(1, 0, 5)); + + // empty cells never match ranges with p0 >= 0 (callers clamp p0) + t.assert_true(!cells.pos_in(0, 0, 1000)); + }); +} + +static void test_seq(testing & t) { + t.test("seq_add_has_count", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(1, 5); + t.assert_true(!cells.seq_has(1, 0)); + t.assert_equal(0, cells.seq_count(1)); + + cells.seq_add(1, 0); + cells.seq_add(1, 2); + cells.seq_add(1, 5); + + t.assert_true(cells.seq_has(1, 0)); + t.assert_true(cells.seq_has(1, 2)); + t.assert_true(cells.seq_has(1, 5)); + t.assert_true(!cells.seq_has(1, 1)); + t.assert_equal(3, cells.seq_count(1)); + + t.assert_equal(5, cells.seq_pos_min(0)); + t.assert_equal(5, cells.seq_pos_max(0)); + t.assert_equal(5, cells.seq_pos_min(2)); + t.assert_equal(5, cells.seq_pos_max(5)); + }); + + t.test("seq_get", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(0, 1); + cells.seq_add(0, 7); + t.assert_equal(7, cells.seq_get(0)); + + cells.seq_add(0, 3); + cells.seq_rm(0, 3); + t.assert_equal(7, cells.seq_get(0)); + }); + + t.test("seq_rm_partial", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(0, 5); + cells.seq_add(0, 0); + cells.seq_add(0, 1); + + // removing one of two seqs keeps the cell + t.assert_true(!cells.seq_rm(0, 0)); + t.assert_true(!cells.is_empty(0)); + t.assert_equal(1, cells.seq_count(0)); + t.assert_true(!cells.seq_has(0, 0)); + t.assert_true(cells.seq_has(0, 1)); + t.assert_equal(1u, cells.get_used()); + + t.assert_equal(-1, cells.seq_pos_min(0)); + t.assert_equal(-1, cells.seq_pos_max(0)); + t.assert_equal(5, cells.seq_pos_min(1)); + t.assert_equal(5, cells.seq_pos_max(1)); + }); + + t.test("seq_rm_full", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(2, 9); + cells.seq_add(2, 3); + + t.assert_true(cells.seq_rm(2, 3)); + t.assert_true(cells.is_empty(2)); + t.assert_equal(0u, cells.get_used()); + t.assert_equal(0u, cells.used_min()); + t.assert_equal(0u, cells.used_max_p1()); + t.assert_equal(-1, cells.seq_pos_min(3)); + t.assert_equal(-1, cells.seq_pos_max(3)); + }); + + t.test("seq_keep_only", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(0, 1); + cells.seq_add(0, 0); + cells.seq_add(0, 1); + cells.seq_add(0, 2); + + // keeping seq 1 drops the other seqs but keeps the cell + t.assert_true(!cells.seq_keep(0, 1)); + t.assert_true(!cells.is_empty(0)); + t.assert_equal(1, cells.seq_count(0)); + t.assert_equal(1, cells.seq_get(0)); + + t.assert_equal(-1, cells.seq_pos_min(0)); + t.assert_equal(-1, cells.seq_pos_min(2)); + t.assert_equal(1, cells.seq_pos_min(1)); + t.assert_equal(1, cells.seq_pos_max(1)); + }); + + t.test("seq_keep_absent_empties", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(1, 3); + cells.seq_add(1, 0); + cells.seq_add(1, 2); + + // the kept seq is not present: the whole cell is cleared + t.assert_true(cells.seq_keep(1, 5)); + t.assert_true(cells.is_empty(1)); + t.assert_equal(0u, cells.get_used()); + t.assert_equal(-1, cells.seq_pos_min(0)); + t.assert_equal(-1, cells.seq_pos_min(2)); + }); + + t.test("seq_keep_empty_noop", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + t.assert_true(!cells.seq_keep(0, 3)); + t.assert_true(cells.is_empty(0)); + t.assert_equal(0u, cells.get_used()); + }); + + t.test("add_sequences", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(0, 7); + + llama_seq_id seqs[3] = {0, 3, 7}; + cells.add_sequences(0, 3, seqs); + + t.assert_equal(3, cells.seq_count(0)); + for (auto s : seqs) { + t.assert_true(cells.seq_has(0, s)); + t.assert_equal(7, cells.seq_pos_min(s)); + t.assert_equal(7, cells.seq_pos_max(s)); + } + }); +} + +static void test_seq_pos(testing & t) { + t.test("min_max_across_cells", [&](testing & t) { + llama_kv_cells cells; + cells.resize(8); + + cells.pos_set(0, 10); + cells.seq_add(0, 0); + cells.pos_set(3, 5); + cells.seq_add(3, 0); + cells.pos_set(7, 20); + cells.seq_add(7, 0); + + t.assert_equal(5, cells.seq_pos_min(0)); + t.assert_equal(20, cells.seq_pos_max(0)); + + cells.rm(3); + t.assert_equal(10, cells.seq_pos_min(0)); + t.assert_equal(20, cells.seq_pos_max(0)); + + cells.rm(7); + t.assert_equal(10, cells.seq_pos_min(0)); + t.assert_equal(10, cells.seq_pos_max(0)); + + cells.rm(0); + t.assert_equal(-1, cells.seq_pos_min(0)); + t.assert_equal(-1, cells.seq_pos_max(0)); + }); + + t.test("duplicate_positions", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(0, 5); + cells.seq_add(0, 0); + cells.pos_set(1, 5); + cells.seq_add(1, 0); + + t.assert_equal(5, cells.seq_pos_min(0)); + t.assert_equal(5, cells.seq_pos_max(0)); + + // removing one of two cells at the same position keeps the position + cells.rm(0); + t.assert_equal(5, cells.seq_pos_min(0)); + t.assert_equal(5, cells.seq_pos_max(0)); + + cells.rm(1); + t.assert_equal(-1, cells.seq_pos_min(0)); + t.assert_equal(-1, cells.seq_pos_max(0)); + }); + + t.test("insert_before_base", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(0, 10); + cells.seq_add(0, 0); + + // a new cell with a smaller position exercises the prepend path + cells.pos_set(1, 5); + cells.seq_add(1, 0); + + t.assert_equal(5, cells.seq_pos_min(0)); + t.assert_equal(10, cells.seq_pos_max(0)); + + cells.rm(1); + t.assert_equal(10, cells.seq_pos_min(0)); + t.assert_equal(10, cells.seq_pos_max(0)); + + cells.rm(0); + t.assert_equal(-1, cells.seq_pos_min(0)); + }); + + t.test("compact", [&](testing & t) { + llama_kv_cells cells; + cells.resize(8); + + cells.pos_set(0, 10); + cells.seq_add(0, 0); + cells.pos_set(1, 11); + cells.seq_add(1, 0); + cells.pos_set(2, 12); + cells.seq_add(2, 0); + + cells.rm(0); + t.assert_equal(11, cells.seq_pos_min(0)); + cells.compact(0); + t.assert_equal(11, cells.seq_pos_min(0)); + t.assert_equal(12, cells.seq_pos_max(0)); + + // tracking still works after compaction: insert before the new base + cells.pos_set(0, 9); + cells.seq_add(0, 0); + t.assert_equal(9, cells.seq_pos_min(0)); + t.assert_equal(12, cells.seq_pos_max(0)); + + // and after the tail + cells.pos_set(3, 13); + cells.seq_add(3, 0); + t.assert_equal(9, cells.seq_pos_min(0)); + t.assert_equal(13, cells.seq_pos_max(0)); + }); + + t.test("absent_seq", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(0, 1); + cells.seq_add(0, 0); + + t.assert_equal(-1, cells.seq_pos_min(1)); + t.assert_equal(-1, cells.seq_pos_max(1)); + }); +} + +static void test_shift(testing & t) { + t.test("pos_add", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(0, 10); + cells.seq_add(0, 0); + cells.pos_set(1, 20); + cells.seq_add(1, 0); + + t.assert_true(!cells.get_has_shift()); + + t.assert_true(!cells.pos_add(0, 5)); + t.assert_equal(15, cells.pos_get(0)); + t.assert_equal(5, cells.get_shift(0)); + t.assert_true(cells.get_has_shift()); + t.assert_equal(15, cells.seq_pos_min(0)); + t.assert_equal(20, cells.seq_pos_max(0)); + + // untouched cell keeps a zero shift + t.assert_equal(0, cells.get_shift(1)); + + // shifts accumulate + t.assert_true(!cells.pos_add(0, -3)); + t.assert_equal(12, cells.pos_get(0)); + t.assert_equal(2, cells.get_shift(0)); + t.assert_equal(12, cells.seq_pos_min(0)); + }); + + t.test("pos_add_removes", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(0, 3); + cells.seq_add(0, 0); + cells.pos_set(1, 10); + cells.seq_add(1, 1); + + // 3 - 4 < 0 -> the cell is removed + t.assert_true(cells.pos_add(0, -4)); + t.assert_true(cells.is_empty(0)); + t.assert_equal(1u, cells.get_used()); + t.assert_equal(-1, cells.seq_pos_min(0)); + t.assert_equal(-1, cells.seq_pos_max(0)); + t.assert_true(cells.get_has_shift()); + + // the other cell is unaffected + t.assert_equal(10, cells.seq_pos_min(1)); + t.assert_equal(10, cells.seq_pos_max(1)); + }); + + t.test("pos_div", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(0, 100); + cells.seq_add(0, 0); + + cells.pos_div(0, 4); + t.assert_equal(25, cells.pos_get(0)); + t.assert_equal(75, cells.get_shift(0)); // 100 - 25 + t.assert_true(cells.get_has_shift()); + t.assert_equal(25, cells.seq_pos_min(0)); + t.assert_equal(25, cells.seq_pos_max(0)); + + // negative positions truncate toward zero + cells.pos_set(1, -7); + cells.seq_add(1, 0); + cells.pos_div(1, 2); + t.assert_equal(-3, cells.pos_get(1)); + t.assert_equal(-4, cells.get_shift(1)); // -7 - (-3) + t.assert_equal(-3, cells.seq_pos_min(0)); + t.assert_equal(25, cells.seq_pos_max(0)); + }); + + t.test("reset_shift", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(0, 10); + cells.seq_add(0, 0); + cells.pos_add(0, 5); // pos 15, shift 5 + cells.pos_div(0, 3); // pos 5, shift 5 + (15 - 5) = 15 + + t.assert_true(cells.get_has_shift()); + t.assert_equal(15, cells.get_shift(0)); + + cells.reset_shift(); + t.assert_true(!cells.get_has_shift()); + t.assert_equal(0, cells.get_shift(0)); + + // positions stay shifted + t.assert_equal(5, cells.pos_get(0)); + t.assert_equal(5, cells.seq_pos_min(0)); + }); +} + +static void test_remove(testing & t) { + t.test("rm", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(2, 7); + cells.seq_add(2, 0); + cells.ext_set(2, {/*x=*/3, /*y=*/4}); + + cells.rm(2); + t.assert_true(cells.is_empty(2)); + t.assert_equal(0u, cells.get_used()); + t.assert_equal(-1, cells.seq_pos_min(0)); + t.assert_equal(-1, cells.seq_pos_max(0)); + }); + + t.test("rm_single", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(1, 9); + cells.seq_add(1, 2); + + cells.rm_single(1, 2); + t.assert_true(cells.is_empty(1)); + t.assert_equal(0u, cells.get_used()); + t.assert_equal(-1, cells.seq_pos_min(2)); + t.assert_equal(-1, cells.seq_pos_max(2)); + }); +} + +static void test_save_restore(testing & t) { + t.test("cp_range_roundtrip", [&](testing & t) { + const uint32_t n = 16; + llama_kv_cells cells; + cells.resize(n); + + cells.pos_set(2, 10); + cells.seq_add(2, 0); + cells.seq_add(2, 1); + cells.ext_set(2, {/*x=*/5, /*y=*/6}); + + cells.pos_set(3, 11); + cells.seq_add(3, 1); + + cells.pos_set(4, 12); + cells.seq_add(4, 2); + + cells.pos_set(5, 13); + cells.seq_add(5, 0); + + // save the state of cells [2, 2 + 4) + const llama_kv_cells saved = cells.cp(2, 4); + + // the copy carries pos/ext/seq + t.assert_equal(4u, saved.size()); + t.assert_equal(10, saved.pos_get(0)); + t.assert_equal(11, saved.pos_get(1)); + t.assert_equal(12, saved.pos_get(2)); + t.assert_equal(13, saved.pos_get(3)); + t.assert_equal(5, saved.ext_get(0).x); + t.assert_equal(6, saved.ext_get(0).y); + t.assert_true(saved.seq_has(0, 0)); + t.assert_true(saved.seq_has(0, 1)); + t.assert_equal(2, saved.seq_count(0)); + t.assert_equal(1, saved.seq_count(1)); + t.assert_equal(1, saved.seq_count(3)); + + // wipe the original cells + cells.rm(2); + cells.rm(3); + cells.rm(4); + cells.rm(5); + t.assert_equal(0u, cells.get_used()); + + // restore + const std::vector idxs = {2, 3, 4, 5}; + cells.set(idxs, saved); + + t.assert_equal(4u, cells.get_used()); + t.assert_equal(2u, cells.used_min()); + t.assert_equal(6u, cells.used_max_p1()); + for (uint32_t j = 0; j < 4; ++j) { + t.assert_equal(saved.pos_get(j), cells.pos_get(2 + j)); + t.assert_equal(saved.seq_has(j, 0), cells.seq_has(2 + j, 0)); + t.assert_equal(saved.seq_has(j, 1), cells.seq_has(2 + j, 1)); + t.assert_equal(saved.seq_has(j, 2), cells.seq_has(2 + j, 2)); + t.assert_equal(saved.seq_count(j), cells.seq_count(2 + j)); + } + t.assert_equal(5, cells.ext_get(2).x); + t.assert_equal(6, cells.ext_get(2).y); + + // sequence position tracking is rebuilt + t.assert_equal(10, cells.seq_pos_min(0)); + t.assert_equal(13, cells.seq_pos_max(0)); + t.assert_equal(10, cells.seq_pos_min(1)); + t.assert_equal(11, cells.seq_pos_max(1)); + t.assert_equal(12, cells.seq_pos_min(2)); + t.assert_equal(12, cells.seq_pos_max(2)); + t.assert_equal(-1, cells.seq_pos_min(3)); + }); + + t.test("cp_idxs_roundtrip", [&](testing & t) { + llama_kv_cells cells; + cells.resize(8); + + cells.pos_set(1, 5); + cells.seq_add(1, 0); + cells.pos_set(6, 9); + cells.seq_add(6, 1); + + const std::vector idxs = {1, 6}; + const llama_kv_cells saved = cells.cp(idxs); + + t.assert_equal(2u, saved.size()); + t.assert_equal(5, saved.pos_get(0)); + t.assert_equal(9, saved.pos_get(1)); + t.assert_true(saved.seq_has(0, 0)); + t.assert_true(saved.seq_has(1, 1)); + + // restore into different cells (remap) + cells.rm(1); + cells.rm(6); + const std::vector dst = {3, 5}; + cells.set(dst, saved); + + t.assert_equal(2u, cells.get_used()); + t.assert_equal(5, cells.pos_get(3)); + t.assert_equal(9, cells.pos_get(5)); + t.assert_true(cells.seq_has(3, 0)); + t.assert_true(cells.seq_has(5, 1)); + t.assert_equal(5, cells.seq_pos_min(0)); + t.assert_equal(9, cells.seq_pos_min(1)); + }); + + t.test("set_replaces_existing", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(0, 1); + cells.seq_add(0, 0); + + // "other" describes a single empty cell + llama_kv_cells other; + other.resize(1); + + const std::vector idxs = {0}; + cells.set(idxs, other); + + t.assert_true(cells.is_empty(0)); + t.assert_equal(0u, cells.get_used()); + t.assert_equal(-1, cells.seq_pos_min(0)); + }); + + t.test("set_fast_path_identical", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(1, 5); + cells.seq_add(1, 0); + + // copy a cell with an identical state: set takes the fast path + const llama_kv_cells saved = cells.cp(1, 1); + + const std::vector idxs = {1}; + cells.set(idxs, saved); + + t.assert_equal(1u, cells.get_used()); + t.assert_equal(5, cells.pos_get(1)); + t.assert_true(cells.seq_has(1, 0)); + t.assert_equal(5, cells.seq_pos_min(0)); + t.assert_equal(5, cells.seq_pos_max(0)); + }); +} + +static void test_next_head(testing & t) { + t.test("removes_range_of_seq", [&](testing & t) { + llama_kv_cells cells; + cells.resize(8); + + cells.pos_set(0, 5); + cells.seq_add(0, 0); + cells.pos_set(1, 6); + cells.seq_add(1, 0); + cells.seq_add(1, 1); // shared cell + cells.pos_set(2, 7); + cells.seq_add(2, 1); + cells.pos_set(3, 8); + cells.seq_add(3, 0); // outside the range + + t.assert_equal(0u, cells.next_head(0, 5, 8)); + + t.assert_true(cells.is_empty(0)); + t.assert_equal(1, cells.seq_count(1)); // shared cell kept seq 1 + t.assert_true(cells.seq_has(2, 1)); + t.assert_true(cells.seq_has(3, 0)); + t.assert_equal(3u, cells.get_used()); + + t.assert_equal(8, cells.seq_pos_min(0)); + t.assert_equal(8, cells.seq_pos_max(0)); + t.assert_equal(6, cells.seq_pos_min(1)); + t.assert_equal(7, cells.seq_pos_max(1)); + }); + + t.test("returns_first_removed", [&](testing & t) { + llama_kv_cells cells; + cells.resize(8); + + cells.pos_set(2, 3); + cells.seq_add(2, 0); + cells.pos_set(4, 3); + cells.seq_add(4, 0); + + // both cells are freed; the lowest index is returned + t.assert_equal(2u, cells.next_head(0, 0, 100)); + + t.assert_true(cells.is_empty(2)); + t.assert_true(cells.is_empty(4)); + t.assert_equal(0u, cells.get_used()); + t.assert_equal(-1, cells.seq_pos_min(0)); + t.assert_equal(-1, cells.seq_pos_max(0)); + }); + + t.test("no_match", [&](testing & t) { + llama_kv_cells cells; + cells.resize(8); + + cells.pos_set(0, 5); + cells.seq_add(0, 0); + cells.pos_set(1, 7); + cells.seq_add(1, 1); + + // ranges do not cover any seq-0 cell + t.assert_equal(cells.size(), cells.next_head(0, 0, 5)); + t.assert_equal(cells.size(), cells.next_head(0, 8, 100)); + // no cell has seq 2 + t.assert_equal(cells.size(), cells.next_head(2, 0, 100)); + t.assert_equal(2u, cells.get_used()); + t.assert_equal(5, cells.seq_pos_min(0)); + t.assert_equal(7, cells.seq_pos_min(1)); + }); +} + +// reference model for the randomized test: same operations, naive O(n) state +struct cells_ref { + struct cell_t { + llama_pos pos = -1; + std::set seqs; + llama_pos shift = 0; + llama_kv_cell_ext ext = {}; + }; + + std::vector cells; + bool has_shift = false; + + void resize(uint32_t n) { + cells.assign(n, cell_t{}); + has_shift = false; + } + + uint32_t size() const { + return (uint32_t) cells.size(); + } + + uint32_t get_used() const { + uint32_t c = 0; + for (const auto & cl : cells) { + if (cl.pos != -1) { + ++c; + } + } + return c; + } + + uint32_t used_min() const { + for (uint32_t i = 0; i < cells.size(); ++i) { + if (cells[i].pos != -1) { + return i; + } + } + return 0; + } + + uint32_t used_max_p1() const { + for (uint32_t i = (uint32_t) cells.size(); i-- > 0;) { + if (cells[i].pos != -1) { + return i + 1; + } + } + return 0; + } + + llama_pos seq_pos_min(llama_seq_id s) const { + llama_pos res = -1; + for (const auto & cl : cells) { + if (cl.pos != -1 && cl.seqs.count(s)) { + res = res == -1 ? cl.pos : std::min(res, cl.pos); + } + } + return res; + } + + llama_pos seq_pos_max(llama_seq_id s) const { + llama_pos res = -1; + for (const auto & cl : cells) { + if (cl.pos != -1 && cl.seqs.count(s)) { + res = std::max(res, cl.pos); + } + } + return res; + } +}; + +static void test_random(testing & t) { + t.test("ops_vs_reference_model", [&](testing & t) { + std::mt19937 rng(1234); + + const uint32_t n = 130; + const uint32_t n_seq = 8; + + llama_kv_cells cells; + cells.resize(n); + + cells_ref ref; + ref.resize(n); + + auto check = [&](const std::string & msg) { + t.assert_equal(ref.size(), cells.size()); + + for (uint32_t i = 0; i < n; ++i) { + const auto & cl = ref.cells[i]; + + t.assert_equal(cl.pos == -1, cells.is_empty(i)); + + if (cl.pos != -1) { + t.assert_equal(cl.pos, cells.pos_get(i)); + t.assert_equal(cl.shift, cells.get_shift(i)); + t.assert_equal(cl.ext.x, cells.ext_get(i).x); + t.assert_equal(cl.ext.y, cells.ext_get(i).y); + t.assert_equal((int) cl.seqs.size(), cells.seq_count(i)); + if (cl.seqs.size() == 1) { + t.assert_equal((llama_seq_id) *cl.seqs.begin(), cells.seq_get(i)); + } + } + + for (llama_seq_id s = 0; s < (llama_seq_id) n_seq; ++s) { + t.assert_equal(cl.seqs.count(s) > 0, cells.seq_has(i, s)); + } + } + + t.assert_equal(ref.get_used(), cells.get_used()); + t.assert_equal(ref.used_min(), cells.used_min()); + t.assert_equal(ref.used_max_p1(), cells.used_max_p1()); + t.assert_equal(ref.has_shift, cells.get_has_shift()); + + for (llama_seq_id s = 0; s < (llama_seq_id) n_seq; ++s) { + t.assert_equal(msg + " seq_pos_min", ref.seq_pos_min(s), cells.seq_pos_min(s)); + t.assert_equal(msg + " seq_pos_max", ref.seq_pos_max(s), cells.seq_pos_max(s)); + } + }; + + for (uint32_t step = 0; step < 3000; ++step) { + const uint32_t i = rng() % n; + const llama_seq_id s = (llama_seq_id) (rng() % n_seq); + + switch (rng() % 12) { + case 0: { // pos_set on an empty cell + if (ref.cells[i].pos == -1) { + // keep positions non-negative: is_empty asserts pos == -1 or pos >= 0 + const llama_pos p = (llama_pos) (rng() % 100); + cells.pos_set(i, p); + ref.cells[i].pos = p; + } + } break; + + case 1: { // seq_add + if (ref.cells[i].pos != -1 && !ref.cells[i].seqs.count(s)) { + cells.seq_add(i, s); + ref.cells[i].seqs.insert(s); + } + } break; + + case 2: { // seq_rm + if (ref.cells[i].seqs.count(s)) { + const bool empty = cells.seq_rm(i, s); + ref.cells[i].seqs.erase(s); + if (ref.cells[i].seqs.empty()) { + ref.cells[i].pos = -1; + ref.cells[i].shift = 0; + ref.cells[i].ext = {}; + } + t.assert_equal(ref.cells[i].seqs.empty(), empty); + } + } break; + + case 3: { // seq_keep (skip used cells without seqs: the API asserts on them) + if (ref.cells[i].pos == -1 || !ref.cells[i].seqs.empty()) { + const bool empty = cells.seq_keep(i, s); + if (ref.cells[i].seqs.count(s)) { + ref.cells[i].seqs = {s}; + t.assert_true(!empty); + } else if (!ref.cells[i].seqs.empty()) { + ref.cells[i] = cells_ref::cell_t{}; + t.assert_true(empty); + } else { + t.assert_true(!empty); + } + } + } break; + + case 4: { // add_sequences + if (ref.cells[i].pos != -1) { + std::vector pick; + for (llama_seq_id k = 0; k < (llama_seq_id) n_seq; ++k) { + if (!ref.cells[i].seqs.count(k)) { + pick.push_back(k); + } + } + if (!pick.empty()) { + const size_t k = 1 + (rng() % pick.size()); + cells.add_sequences(i, (int32_t) k, pick.data()); + for (size_t j = 0; j < k; ++j) { + ref.cells[i].seqs.insert(pick[j]); + } + } + } + } break; + + case 5: { // rm + if (ref.cells[i].pos != -1) { + cells.rm(i); + ref.cells[i] = cells_ref::cell_t{}; + } + } break; + + case 6: { // rm_single + if (ref.cells[i].seqs.size() == 1) { + cells.rm_single(i, *ref.cells[i].seqs.begin()); + ref.cells[i] = cells_ref::cell_t{}; + } + } break; + + case 7: { // pos_add + if (ref.cells[i].pos != -1) { + const int d = (int) (rng() % 41) - 20; + const bool removed = cells.pos_add(i, d); + ref.has_shift = true; + if (ref.cells[i].pos + d < 0) { + // pos_add clears pos/shift/seq but not ext + ref.cells[i].seqs.clear(); + ref.cells[i].pos = -1; + ref.cells[i].shift = 0; + t.assert_true(removed); + } else { + ref.cells[i].pos += d; + ref.cells[i].shift += d; + t.assert_true(!removed); + } + } + } break; + + case 8: { // pos_div + if (ref.cells[i].pos != -1) { + const int d = 1 + (int) (rng() % 4); + const llama_pos p_old = ref.cells[i].pos; + cells.pos_div(i, d); + ref.cells[i].pos = p_old / d; + ref.cells[i].shift += p_old - ref.cells[i].pos; + ref.has_shift = true; + } + } break; + + case 9: { // next_head + const llama_pos p0 = (llama_pos) ((int) (rng() % 40) - 10); + const llama_pos p1 = p0 + (llama_pos) (rng() % 60); + + uint32_t new_head = n; + for (uint32_t k = 0; k < n; ++k) { + auto & cl = ref.cells[k]; + if (cl.pos == -1 || cl.pos < p0 || cl.pos >= p1 || !cl.seqs.count(s)) { + continue; + } + cl.seqs.erase(s); + if (cl.seqs.empty()) { + cl.pos = -1; + cl.shift = 0; + cl.ext = {}; + if (new_head == n) { + new_head = k; + } + } + } + + t.assert_equal(new_head, cells.next_head(s, p0, p1)); + } break; + + case 10: { // reset_shift + cells.reset_shift(); + for (auto & cl : ref.cells) { + cl.shift = 0; + } + ref.has_shift = false; + } break; + + case 11: { // ext_set + const llama_kv_cell_ext e{ + /*x=*/ (llama_pos) (rng() % 100), + /*y=*/ (llama_pos) (rng() % 100), + }; + cells.ext_set(i, e); + ref.cells[i].ext = e; + } break; + } + + // periodic save/restore roundtrip of a random contiguous range + if (step % 256 == 128) { + cells.reset_shift(); + for (auto & cl : ref.cells) { + cl.shift = 0; + } + ref.has_shift = false; + + const uint32_t a = rng() % (n - 8); + const uint32_t len = 1 + (rng() % 8); + + const llama_kv_cells saved = cells.cp(a, len); + + for (uint32_t k = a; k < a + len; ++k) { + if (ref.cells[k].pos != -1) { + cells.rm(k); + ref.cells[k] = cells_ref::cell_t{}; + } + } + + std::vector idxs; + for (uint32_t k = a; k < a + len; ++k) { + idxs.push_back(k); + } + cells.set(idxs, saved); + + // restore the reference model from the copy + for (uint32_t j = 0; j < len; ++j) { + auto & cl = ref.cells[a + j]; + if (saved.is_empty(j)) { + // empty cells are not touched by set(): it only copies ext from saved, + // which equals the current ext of the cell + continue; + } + cl.pos = saved.pos_get(j); + cl.ext = saved.ext_get(j); + cl.seqs.clear(); + for (llama_seq_id s2 = 0; s2 < (llama_seq_id) n_seq; ++s2) { + if (saved.seq_has(j, s2)) { + cl.seqs.insert(s2); + } + } + } + } + + check("step " + std::to_string(step)); + } + }); +} + +int main(int argc, char ** argv) { + testing t; + + if (argc > 1) { + t.set_filter(argv[1]); + } + + t.test("bitops", test_bitops); + t.test("ext", test_ext); + t.test("basic", test_basic); + t.test("seq", test_seq); + t.test("seq_pos", test_seq_pos); + t.test("shift", test_shift); + t.test("remove", test_remove); + t.test("save_restore", test_save_restore); + t.test("next_head", test_next_head); + t.test("random", test_random); + + return t.summary(); +}