From 7e4cb24e0c462abd2b8c9fe8a7ac8a6bb990d61f Mon Sep 17 00:00:00 2001 From: lexasub Date: Fri, 14 Aug 2026 05:58:48 +0400 Subject: [PATCH 1/9] improve kv cache performance --- src/CMakeLists.txt | 1 + src/llama-kv-cache.cpp | 23 +- src/llama-kv-cells.cpp | 178 ++++++ src/llama-kv-cells.h | 322 ++++++----- tests/CMakeLists.txt | 1 + tests/test-kv-cells.cpp | 1151 +++++++++++++++++++++++++++++++++++++++ 6 files changed, 1512 insertions(+), 164 deletions(-) create mode 100644 src/llama-kv-cells.cpp create mode 100644 tests/test-kv-cells.cpp diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 39ba3061f704..f0e381e955d9 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-kv-cache.cpp b/src/llama-kv-cache.cpp index 5382cd7266f8..d501c7d9ad1a 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.nextHead(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.seqS_add(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 000000000000..2212f1b34fb6 --- /dev/null +++ b/src/llama-kv-cells.cpp @@ -0,0 +1,178 @@ +#include "llama-kv-cache.h" +#include + +void llama_kv_cells::seqS_add(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 (h == 0 && t + 1 == v.cnt.size()) { + return; + } + + if (h > 0) { + v.cnt.erase(v.cnt.begin(), v.cnt.begin() + h); + } + v.cnt.resize(t + 1 - h); + + v.base += (llama_pos)h; + v.head = 0; + v.tail = t - h; +} + +uint32_t llama_kv_cells::nextHead(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 fddd31a0b219..e7f1d7f1231a 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 k = 0; k < N_SEQ_WORDS; ++k) { + if (seq[i].w[k]) return k * 64 + llama_bits::countr_zero64(seq[i].w[k]); } - 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 seqS_add(uint32_t i, int32_t n, llama_seq_id *_seq); + void compact(llama_seq_id s); + uint32_t nextHead(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,52 +482,83 @@ 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 k = 0; k < N_SEQ_WORDS; ++k) { + for (auto m = seq[i].w[k]; m; m &= m - 1) { + seq_pos_dec(k * 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 k = 0; k < N_SEQ_WORDS; ++k) { + for (auto m = seq[i].w[k]; m; m &= m - 1) { + seq_pos_inc(k * 64 + llama_bits::countr_zero64(m), pos[i]); + } + } } }; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 419e1eba4c2c..02376032e2cf 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 000000000000..04ddb1685bad --- /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("seqS_add", [&](testing & t) { + llama_kv_cells cells; + cells.resize(4); + + cells.pos_set(0, 7); + + llama_seq_id seqs[3] = {0, 3, 7}; + cells.seqS_add(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.nextHead(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.nextHead(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.nextHead(0, 0, 5)); + t.assert_equal(cells.size(), cells.nextHead(0, 8, 100)); + // no cell has seq 2 + t.assert_equal(cells.size(), cells.nextHead(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: { // seqS_add + 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.seqS_add(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: { // nextHead + 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.nextHead(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(); +} From 6f1152d05a957656171a70696504068870fa96ae Mon Sep 17 00:00:00 2001 From: lexasub Date: Fri, 14 Aug 2026 09:25:50 +0400 Subject: [PATCH 2/9] Revert "z" This reverts commit 57187887804305b2fdd63f2b0e0d6b96b53d30c5. --- src/llama-batch.cpp | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 2b98a552f48f..98d0cf76527b 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -752,20 +752,24 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u assert(n_tokens%n_seqs == 0); auto udata = std::make_shared(); + udata->seq_id_unq.resize(0); const int64_t n_embd_all = batch.embd ? (int64_t) n_tokens*n_embd : 0; 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); - udata->seq_id_unq.resize(0); udata->seq_idx .resize(LLAMA_MAX_SEQ, -1); udata->output .resize(n_tokens); udata->seq_id_data.reserve(n_tokens); + udata->token .resize(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,10 +779,13 @@ 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[i] * 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) { + for (size_t j = 0; j < (size_t) n_pos_per_embd; ++j) { // if we are using M-RoPE // if the current batch is text, we need to broadcast the same position across all RoPE sections // otherwise, the input batch is image embeddings, we copy the positions as-is @@ -802,6 +809,7 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u } } + udata->seq_id.resize(n_tokens); llama_seq_id * seq_id_ptr = udata->seq_id_data.data(); for (size_t i = 0; i < idxs.size(); ++i) { udata->seq_id[i] = seq_id_ptr; From 50ff3e30c429a2fafe9ee442f455dbc94bd4fc05 Mon Sep 17 00:00:00 2001 From: lexasub Date: Wed, 12 Aug 2026 05:57:15 +0400 Subject: [PATCH 3/9] tokenizer: very fast tokenizing --- src/llama-vocab.cpp | 66 +++++----- src/openhashmap.h | 289 ++++++++++++++++++++++++++++++++++++++++++++ src/unicode.cpp | 87 +++++++------ 3 files changed, 362 insertions(+), 80 deletions(-) create mode 100644 src/openhashmap.h diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 4a01dfd4cab6..f130319edbd3 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -4,6 +4,7 @@ #include "gguf.h" #include "llama-impl.h" #include "llama-model-loader.h" +#include "openhashmap.h" #include "unicode.h" @@ -644,7 +645,7 @@ struct llm_tokenizer_bpe_session { // build token(s) while (!work_queue.empty()) { - auto bigram = work_queue.pop_move(); + auto bigram = work_queue.pop_move();//cache miss 1.97, br miss 10.7 auto & left_symbol = symbols[bigram.left]; auto & right_symbol = symbols[bigram.right]; @@ -652,10 +653,12 @@ struct llm_tokenizer_bpe_session { if (left_symbol.n == 0 || right_symbol.n == 0) { continue; } - std::string left_token = std::string(left_symbol.text, left_symbol.n); - std::string right_token = std::string(right_symbol.text, right_symbol.n); - if (left_token + right_token != bigram.text) { - continue; // Skip this bigram if it's outdated + if (left_symbol.n + right_symbol.n != bigram.text.size()) { + continue; // Skip this bigram if it's outdated + } + if (memcmp(left_symbol.text, bigram.text.data(), left_symbol.n) != 0 || + memcmp(right_symbol.text, bigram.text.data() + left_symbol.n, right_symbol.n) != 0) { + continue; // Skip this bigram if it's outdated } // merge the right sym into the left one @@ -742,8 +745,9 @@ struct llm_tokenizer_bpe_session { bigram.left = left; bigram.right = right; - bigram.text = left_token + right_token; bigram.size = left_token.size() + right_token.size(); + bigram.text = std::move(left_token); + bigram.text += right_token; bigram.rank = rank_found; work_queue.push(bigram); @@ -754,7 +758,7 @@ struct llm_tokenizer_bpe_session { std::vector symbols; std::vector symbols_final; - llm_bigram_bpe::queue work_queue; + llm_bigram_bpe::queue work_queue; //TODO - improve algo (now pop_move, push eat cpy cycles) }; // @@ -1824,18 +1828,12 @@ struct llama_vocab::impl { // BertNormalizer options llama_vocab::normalizer_options normalizer_opts; - std::unordered_map token_to_id; + OpenHashMap token_to_id; // 262144 - conservative inital capacity std::vector id_to_token; std::vector cache_special_tokens; std::vector cache_token_to_piece; // llama_token_to_piece(special = true); - struct pair_hash { - size_t operator()(const std::pair & p) const { - return std::hash{}(p.first) ^ //create some hash for pair - (std::hash{}(p.second) << 1); - } - }; - std::unordered_map, int, pair_hash> bpe_ranks; + OpenHashMap, int, 262144, PairStringHash> bpe_ranks; // 262144 - conservative inital capacity // set of all tokens that cause "end of generation" std::set special_eog_ids; @@ -2004,7 +2002,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { second = word.substr(pos + 1); } - bpe_ranks.emplace(std::make_pair(first, second), i); + bpe_ranks.insert(std::make_pair(first, second), i); } } @@ -2095,7 +2093,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { second = word.substr(pos + 1); } - bpe_ranks.emplace(std::make_pair(first, second), i); + bpe_ranks.insert(std::make_pair(first, second), i); } } @@ -2438,7 +2436,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { word = "[EMPTY_" + std::to_string(i) + "]"; } - token_to_id[word] = i; + token_to_id.insert_or_assign(word, i); // operator[] semantics: last duplicate wins max_token_len = std::max(max_token_len, (int) word.size()); auto & token_data = id_to_token[i]; @@ -2466,8 +2464,8 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { // k-mers are the block right after , so only scan from there. if (tokenizer_model == "hybriddna") { const auto idx = token_to_id.find(""); - if (idx != token_to_id.end()) { - auto it = id_to_token.begin() + idx->second + 1; + if (idx != nullptr) { + auto it = id_to_token.begin() + (*idx) + 1; for (; it != id_to_token.end(); ++it) { std::string & text = it->text; if (text.size() > dna_kmer_marker.size() @@ -3001,7 +2999,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { }; auto _set_token_attr = [&] (const std::string & token, llama_token_attr attr, bool value) { - _set_tokenid_attr(token_to_id.at(token), attr, value); + auto tok_id = token_to_id.find(token); + if (tok_id != nullptr) { + _set_tokenid_attr(*tok_id, attr, value); + } }; std::string model_name; @@ -3024,7 +3025,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { || _contains_any(tokenizer_pre, {"jina-v2-de", "jina-v2-es", "jina-v2-code"}) || _contains_any(general_arch, {"nomic-bert-moe", "jina-bert-v3"}) ) { - if (token_to_id.count("") == 0) { + if (token_to_id.find("") == nullptr) { LLAMA_LOG_WARN("%s: Mask token is missing in vocab, please reconvert model!\n", __func__); } else { _set_token_attr("", LLAMA_TOKEN_ATTR_LSTRIP, true); @@ -3040,7 +3041,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { _set_token_attr(token, LLAMA_TOKEN_ATTR_RSTRIP, false); } } else if (_contains_any(model_name, {"modern-bert"})) { - if (token_to_id.count("[MASK]") == 0 ) { + if (token_to_id.find("[MASK]") == nullptr) { LLAMA_LOG_WARN("%s: Mask token missing in vocab!\n", __func__); } else { @@ -3875,8 +3876,8 @@ llama_token llama_vocab::byte_to_token(uint8_t ch) const { case LLAMA_VOCAB_TYPE_UGM: { const char buf[7] = { '<', '0', 'x', hex[ch >> 4], hex[ch & 15], '>', 0 }; auto token = pimpl->token_to_id.find(buf); - if (token != pimpl->token_to_id.end()) { - return (*token).second; + if (token != nullptr) { + return *token; } // Try to fall back to just the byte as a string const char buf2[2] = { (char)ch, 0 }; @@ -3899,9 +3900,8 @@ llama_token llama_vocab::byte_to_token(uint8_t ch) const { llama_token llama_vocab::text_to_token(const std::string & text) const { GGML_ASSERT(pimpl->type != LLAMA_VOCAB_TYPE_NONE); - auto it = pimpl->token_to_id.find(text); - if (it != pimpl->token_to_id.end()) { - return (*it).second; + if (auto it = pimpl->token_to_id.find(text)) { + return *it; } return LLAMA_TOKEN_NULL; } @@ -4046,15 +4046,11 @@ int llama_vocab::max_token_len() const { } int llama_vocab::find_bpe_rank(const std::string & token_left, const std::string & token_right) const { - GGML_ASSERT(token_left.find(' ') == std::string::npos); + GGML_ASSERT(token_left.find(' ') == std::string::npos); // TODO optimize find for small strings GGML_ASSERT(token_right.find(' ') == std::string::npos); - auto it = pimpl->bpe_ranks.find(std::make_pair(token_left, token_right)); - if (it == pimpl->bpe_ranks.end()) { - return -1; - } - - return it->second; + auto it = pimpl->bpe_ranks.find(std::make_pair(token_left, token_right)); // TODO compute hash from string,string without creating pair + return it ? *it : -1; } std::vector llama_vocab::get_bpe_merges() const { diff --git a/src/openhashmap.h b/src/openhashmap.h new file mode 100644 index 000000000000..0d3d97e9f532 --- /dev/null +++ b/src/openhashmap.h @@ -0,0 +1,289 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +struct StringHash { + size_t operator()(const std::string& s) const { + uint64_t h = 1469598103934665603ull; + for (unsigned char c : s) { + h ^= c; + h *= 1099511628211ull; + } + return h; + } +}; + +struct PairStringHash { + size_t operator()(const std::pair& p) const { + uint64_t h = 1469598103934665603ull; // FNV offset basis + for (unsigned char c : p.first) { + h ^= c; + h *= 1099511628211ull; // FNV prime + } + h *= 1099511628211ull; // mix между first и second + for (unsigned char c : p.second) { + h ^= c; + h *= 1099511628211ull; + } + return h; + } +}; + +// Open-addressing hash map with linear probing, cached 32-bit hashes and +// automatic growth. Built once, read many times by the tokenizer hot path, +// so there is deliberately no erasure support. +// +// Capacity is the *initial* capacity (power of 2). The table grows 2x when +// the load factor exceeds 0.7, so insert() can never silently drop data even +// if the number of entries exceeds Capacity (e.g. a vocab with more merges +// than 262144). +template +class OpenHashMap { + static_assert((Capacity & (Capacity - 1)) == 0, "Capacity must be power of 2"); + static_assert(Capacity >= 2, "Capacity must be >= 2"); + + struct Entry { + uint32_t hash = 0; // 0 == empty slot + Key first; + Value second{}; + }; + + std::vector entries_; + size_t capacity_ = Capacity; + size_t size_ = 0; + + // grow when size_ * GROW_NUM >= capacity_ * GROW_DEN (i.e. load >= 0.7) + static constexpr size_t GROW_NUM = 7; + static constexpr size_t GROW_DEN = 10; + + // murmur3 finalizer: good avalanche in the low bits that index the table + static uint32_t mix64(uint64_t h) { + h ^= h >> 33; + h *= 0xff51afd7ed558ccdull; + h ^= h >> 33; + h *= 0xc4ceb9fe1a85ec53ull; + h ^= h >> 33; + uint32_t r = (uint32_t) h; + return r == 0 ? 1u : r; // 0 is reserved for "empty" + } + + static uint32_t hash_key(const Key& key) { + return mix64((uint64_t) Hasher{}(key)); + } + + bool needs_grow() const { + return size_ * GROW_DEN >= capacity_ * GROW_NUM; + } + + void grow() { + const size_t new_cap = capacity_ * 2; + std::vector next(new_cap); + const uint32_t mask = (uint32_t) new_cap - 1; + size_t new_size = 0; + for (auto & e : entries_) { + if (e.hash == 0) { + continue; + } + size_t idx = e.hash & mask; + while (next[idx].hash != 0) { + idx = (idx + 1) & mask; + } + next[idx] = std::move(e); + ++new_size; + } + entries_ = std::move(next); + capacity_ = new_cap; + size_ = new_size; + } + +public: + OpenHashMap() : entries_(Capacity) {} + + // insert if the key is not present; returns pointer to the (existing or new) value. + // keeps the existing value on a key collision (emplace semantics) + Value* insert(const Key& key, const Value& value) { + if (needs_grow()) { + grow(); + } + const uint32_t h = hash_key(key); + const uint32_t mask = (uint32_t) capacity_ - 1; + size_t idx = h & mask; + for (size_t i = 0; i < capacity_; ++i) { + Entry & e = entries_[idx]; + if (e.hash == 0) { + e.hash = h; + e.first = key; + e.second = value; + ++size_; + return &e.second; + } + if (e.hash == h && e.first == key) { + return &e.second; // already present: keep the first value + } + idx = (idx + 1) & mask; + } + return nullptr; // unreachable: growth keeps the load factor < 1 + } + + // insert or overwrite the value if the key is already present (operator[] semantics) + Value* insert_or_assign(const Key& key, const Value& value) { + if (needs_grow()) { + grow(); + } + const uint32_t h = hash_key(key); + const uint32_t mask = (uint32_t) capacity_ - 1; + size_t idx = h & mask; + for (size_t i = 0; i < capacity_; ++i) { + Entry & e = entries_[idx]; + if (e.hash == 0) { + e.hash = h; + e.first = key; + e.second = value; + ++size_; + return &e.second; + } + if (e.hash == h && e.first == key) { + e.second = value; + return &e.second; + } + idx = (idx + 1) & mask; + } + return nullptr; // unreachable + } + + const Value* find(const Key& key) const { + if (size_ == 0) { + return nullptr; + } + const uint32_t h = hash_key(key); + const uint32_t mask = (uint32_t) capacity_ - 1; + size_t idx = h & mask; + for (size_t i = 0; i < capacity_; ++i) { + const Entry & e = entries_[idx]; + if (e.hash == 0) { + return nullptr; // linear probing: key cannot sit past the first empty slot + } + if (e.hash == h && e.first == key) { + return &e.second; + } + idx = (idx + 1) & mask; + } + return nullptr; + } + + const Value& at(const Key& key) const { + const Value* v = find(key); + if (!v) throw std::out_of_range("OpenHashMap::at: key not found"); + return *v; + } + + constexpr size_t size() const { return size_; } + + // --- Iterators --- + class Iterator; + class ConstIterator; + Iterator begin() { return Iterator(entries_.data(), 0, capacity_); } + Iterator end() { return Iterator(entries_.data(), capacity_, capacity_); } + ConstIterator begin() const { return ConstIterator(entries_.data(), 0, capacity_); } + ConstIterator end() const { return ConstIterator(entries_.data(), capacity_, capacity_); } + ConstIterator cbegin() const { return begin(); } + ConstIterator cend() const { return end(); } + + class Iterator { + public: + using iterator_category = std::forward_iterator_tag; + using value_type = Entry; + using difference_type = std::ptrdiff_t; + using pointer = Entry*; + using reference = Entry&; + + Iterator() : entries_(nullptr), index_(0), capacity_(0) {} + + reference operator*() { return entries_[index_]; } + pointer operator->() { return &entries_[index_]; } + + Iterator& operator++() { + ++index_; + advance_to_occupied(); + return *this; + } + + Iterator operator++(int) { + Iterator tmp = *this; + ++*this; + return tmp; + } + + bool operator==(const Iterator& other) const { return index_ == other.index_; } + bool operator!=(const Iterator& other) const { return index_ != other.index_; } + + private: + friend class OpenHashMap; + Iterator(Entry* entries, size_t index, size_t capacity) + : entries_(entries), index_(index), capacity_(capacity) { + advance_to_occupied(); + } + + void advance_to_occupied() { + while (index_ < capacity_ && entries_[index_].hash == 0) { + ++index_; + } + } + + Entry* entries_; + size_t index_; + size_t capacity_; + }; + + class ConstIterator { + public: + using iterator_category = std::forward_iterator_tag; + using value_type = Entry; + using difference_type = std::ptrdiff_t; + using pointer = const Entry*; + using reference = const Entry&; + + ConstIterator() : entries_(nullptr), index_(0), capacity_(0) {} + ConstIterator(const Iterator& it) + : entries_(it.entries_), index_(it.index_), capacity_(it.capacity_) {} + + reference operator*() const { return entries_[index_]; } + pointer operator->() const { return &entries_[index_]; } + + ConstIterator& operator++() { + ++index_; + advance_to_occupied(); + return *this; + } + + ConstIterator operator++(int) { + ConstIterator tmp = *this; + ++*this; + return tmp; + } + + bool operator==(const ConstIterator& other) const { return index_ == other.index_; } + bool operator!=(const ConstIterator& other) const { return index_ != other.index_; } + + private: + friend class OpenHashMap; + ConstIterator(const Entry* entries, size_t index, size_t capacity) + : entries_(entries), index_(index), capacity_(capacity) { + advance_to_occupied(); + } + + void advance_to_occupied() { + while (index_ < capacity_ && entries_[index_].hash == 0) { + ++index_; + } + } + + const Entry* entries_; + size_t index_; + size_t capacity_; + }; +}; diff --git a/src/unicode.cpp b/src/unicode.cpp index b02ecdc930fa..810c3fc477e6 100644 --- a/src/unicode.cpp +++ b/src/unicode.cpp @@ -2,6 +2,7 @@ #include "unicode-data.h" #include +#include #include #include #include @@ -145,28 +146,21 @@ static std::vector unicode_cpt_flags_array() { return cpt_flags; } -static std::unordered_map unicode_byte_to_utf8_map() { - std::unordered_map map; - for (int ch = 0x21; ch <= 0x7E; ++ch) { // u'!' to u'~' - assert(0 <= ch && ch < 256); - map[ch] = unicode_cpt_to_utf8(ch); - } - for (int ch = 0xA1; ch <= 0xAC; ++ch) { // u'¡' to u'¬' - assert(0 <= ch && ch < 256); - map[ch] = unicode_cpt_to_utf8(ch); - } - for (int ch = 0xAE; ch <= 0xFF; ++ch) { // u'®' to u'ÿ' - assert(0 <= ch && ch < 256); - map[ch] = unicode_cpt_to_utf8(ch); - } - auto n = 0; +static std::array unicode_byte_to_utf8_map() { + std::array mapping; + int n = 0; for (int ch = 0; ch < 256; ++ch) { - if (map.find(ch) == map.end()) { - map[ch] = unicode_cpt_to_utf8(256 + n); + // Printable ASCII: 0x21-0x7E or Latin-1 supplement: 0xA1-0xAC, 0xAE-0xFF + if ((ch >= 0x21 && ch <= 0x7E) || (ch >= 0xA1 && ch <= 0xAC) || (ch >= 0xAE && ch <= 0xFF)) { + mapping[ch] = unicode_cpt_to_utf8(ch); + } + // Everything else maps to U+0100+ + else { + mapping[ch] = unicode_cpt_to_utf8(256 + n); ++n; } } - return map; + return mapping; } static std::unordered_map unicode_utf8_to_byte_map() { @@ -738,13 +732,19 @@ static std::vector unicode_regex_split_custom_qwen35(const std::string & template static std::vector unicode_regex_split_stl(const std::basic_string & text, const std::basic_string & regex, const std::vector & offsets) { using BidirIt = typename std::basic_string::const_iterator; + thread_local std::basic_string cached_regex; + thread_local std::basic_regex cached_expr; // cache compiled regex #ifdef _MSC_VER // Bypass bug in MSVC: https://github.com/ggml-org/llama.cpp/issues/17830 constexpr auto regex_flags = std::regex_constants::ECMAScript; #else constexpr auto regex_flags = std::regex_constants::optimize | std::regex_constants::nosubs; #endif - std::basic_regex expr(regex, regex_flags); + if (cached_regex != regex) { // [[unlikely]] + cached_regex = regex; + cached_expr = std::basic_regex(regex, regex_flags); + } + const auto & expr = cached_expr; std::vector bpe_offsets; // store the offset of each word bpe_offsets.reserve(offsets.size()); // Reserve memory for the approximate size size_t start = 0; @@ -754,7 +754,7 @@ static std::vector unicode_regex_split_stl(const std::basic_string match = *it; + const std::match_results &match = *it; if (match.position() > start_idx) { bpe_offsets.emplace_back(match.position() - start_idx); } @@ -1086,32 +1086,29 @@ static std::vector unicode_regex_split_custom(const std::string & text, // std::string unicode_cpt_to_utf8(uint32_t cpt) { - std::string result; - - if (/* 0x00 <= cpt && */ cpt <= 0x7f) { - result.push_back(cpt); - return result; - } - if (0x80 <= cpt && cpt <= 0x7ff) { - result.push_back(0xc0 | ((cpt >> 6) & 0x1f)); - result.push_back(0x80 | (cpt & 0x3f)); - return result; - } - if (0x800 <= cpt && cpt <= 0xffff) { - result.push_back(0xe0 | ((cpt >> 12) & 0x0f)); - result.push_back(0x80 | ((cpt >> 6) & 0x3f)); - result.push_back(0x80 | (cpt & 0x3f)); - return result; + if (cpt > 0x10ffff) { + throw std::invalid_argument("invalid codepoint"); } - if (0x10000 <= cpt && cpt <= 0x10ffff) { - result.push_back(0xf0 | ((cpt >> 18) & 0x07)); - result.push_back(0x80 | ((cpt >> 12) & 0x3f)); - result.push_back(0x80 | ((cpt >> 6) & 0x3f)); - result.push_back(0x80 | (cpt & 0x3f)); - return result; + if (cpt <= 0x7f) { + return std::string(1, (char) cpt); } - throw std::invalid_argument("invalid codepoint"); + static const uint8_t lead_prefixes[] = {0xc0, 0xe0, 0xf0}; + static const uint8_t lead_shifts[] = {6, 12, 18}; + + size_t len = 1 + (cpt > 0x7f) + (cpt > 0x7ff) + (cpt > 0xffff); + size_t idx = len - 2; + size_t start = 4 - len; + + char buf[4]; + buf[3] = 0x80 | (cpt & 0x3f); + buf[2] = 0x80 | ((cpt >> 6) & 0x3f); + buf[1] = 0x80 | ((cpt >> 12) & 0x3f); + buf[0] = 0x80 | ((cpt >> 18) & 0x3f); + + buf[start] = lead_prefixes[idx] | ((cpt >> lead_shifts[idx]) & 0x3f); + + return std::string(buf + start, len); } std::vector unicode_cpts_normalize_nfd(const std::vector & cpts) { @@ -1160,8 +1157,8 @@ unicode_cpt_flags unicode_cpt_flags_from_utf8(const std::string & utf8) { } std::string unicode_byte_to_utf8(uint8_t byte) { - static std::unordered_map map = unicode_byte_to_utf8_map(); - return map.at(byte); + static std::array map = unicode_byte_to_utf8_map(); + return map[byte]; } uint8_t unicode_utf8_to_byte(const std::string & utf8) { From 60ed1065bf9d172686a83cb229c73d7ea4dae81b Mon Sep 17 00:00:00 2001 From: lexasub Date: Wed, 12 Aug 2026 19:28:24 +0400 Subject: [PATCH 4/9] improove set_input_kq_mask_impl --- src/llama-kv-cache.cpp | 219 +++++++++++++++++++++++++++-------------- 1 file changed, 146 insertions(+), 73 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index d501c7d9ad1a..ff5da654ee07 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -13,6 +13,11 @@ #include #include +#if defined(__AVX2__) && defined(__F16C__) +#include +#endif + + static bool ggml_is_power_of_2(int n) { return (n & (n - 1)) == 0; } @@ -1521,10 +1526,99 @@ struct args_set_input_kq_mask { int64_t n_tps; }; +template +static std::pair compute_position_window( + llama_pos p1, + uint32_t n_swa, + llama_swa_type swa_type +) { + int64_t lo = INT32_MIN; + int64_t hi = INT32_MAX; + + if constexpr (causal) { + hi = std::min(hi, (int64_t) p1); + } + + if constexpr (swa) { + switch (swa_type) { + case LLAMA_SWA_TYPE_STANDARD: + lo = std::max(lo, (int64_t) p1 - n_swa + 1); + break; + case LLAMA_SWA_TYPE_CHUNKED: + lo = std::max(lo, (int64_t) ((p1 / n_swa) * n_swa)); + break; + case LLAMA_SWA_TYPE_SYMMETRIC: { + const int64_t h = n_swa / 2; + lo = std::max(lo, (int64_t) p1 - h); + hi = std::min(hi, (int64_t) p1 + h); + break; + } + case LLAMA_SWA_TYPE_NONE: + break; + } + } + + return { (llama_pos) lo, (llama_pos) hi }; +} + +template +static void set_mask_cell_scalar( + const llama_kv_cells & cells, + uint32_t j, + llama_seq_id seq_id, + llama_pos p1, + llama_pos p1_x, llama_pos p1_y, + llama_pos wlo, llama_pos whi, + T * dst_ptr, + T mask_drop, + T mask_keep, + const llama_swa_type swa_type +) { + if constexpr (alibi) { + if (cells.is_empty(j) || !cells.seq_has(j, seq_id)) { + *dst_ptr = mask_drop; + return; + } + llama_pos p0 = cells.pos_get(j); + + if ((causal && p0 > p1) || + (swa && llama_hparams::is_masked_swa(0, swa_type, p0, p1))) { + *dst_ptr = mask_drop; + return; + } + if constexpr (is_2d) { + if (p0 == p1 && cells.ext_get(j).is_2d_gt(p1_x, p1_y)) { + *dst_ptr = mask_drop; + return; + } + } + *dst_ptr = llama_cast(static_cast(-std::abs(p0 - p1))); + return; + } + + bool nonempty = (cells.pos_get(j) >= 0); + bool keep = nonempty && cells.seq_has(j, seq_id); + + if (keep && (causal || swa)) { + llama_pos p0 = cells.pos_get(j); + keep = (p0 >= wlo && p0 <= whi); + } + + if constexpr (is_2d) { + if (keep) { + llama_pos p0 = cells.pos_get(j); + if (p0 == p1 && cells.ext_get(j).is_2d_gt(p1_x, p1_y)) { + keep = false; + } + } + } + + *dst_ptr = keep ? mask_keep : mask_drop; +} + template static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data) { - //const auto & hparams = args.hparams; - const auto & ubatch = args.ubatch; + const auto & ubatch = args.ubatch; const auto & v_cells = args.v_cells; const auto & seq_to_stream = args.seq_to_stream; @@ -1535,6 +1629,7 @@ static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data const int64_t n_kv = args.n_kv; const int64_t n_stream = args.n_stream; const int64_t n_tps = args.n_tps; + const uint32_t n_kv32 = (uint32_t) n_kv; const T mask_keep = llama_cast(0.0f); const T mask_drop = llama_cast(-INFINITY); @@ -1549,26 +1644,29 @@ static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data seq_pos_min[seq_id] = std::min(seq_pos_min[seq_id], ubatch->pos[i]); } + uint32_t seq_srct[LLAMA_MAX_SEQ]; + std::vector seq_idxs[LLAMA_MAX_SEQ]; + for (uint32_t s = 0; s < n_stream; ++s) { // bookkeeping of the KQ mask cells that could change for other tokens of the same sequence - std::unordered_map seq_srct; - std::unordered_map> seq_idxs; + std::fill_n(seq_srct, LLAMA_MAX_SEQ, UINT32_MAX); + for (auto & v : seq_idxs) v.clear(); for (uint32_t ii = 0; ii < n_tps; ++ii) { - const uint32_t i = s*n_tps + ii; + const uint32_t i = s * n_tps + ii; const llama_seq_id seq_id = ubatch->seq_id[i][0]; const auto & cells = v_cells.at(seq_to_stream[seq_id]); - llama_pos p0 = -1; - const llama_pos p1 = ubatch->pos[i]; + const llama_pos p1 = ubatch->pos[i]; // for M-RoPE - const llama_pos p1_x = is_2d ? ubatch->pos[i + ubatch->n_tokens*2] : 0; - const llama_pos p1_y = is_2d ? ubatch->pos[i + ubatch->n_tokens] : 0; + const llama_pos p1_x = is_2d ? ubatch->pos[i + ubatch->n_tokens * 2] : 0; + const llama_pos p1_y = is_2d ? ubatch->pos[i + ubatch->n_tokens] : 0; - const uint64_t idst = n_kv*i; + const uint64_t idst = n_kv * i; + const auto [wlo, whi] = compute_position_window(p1, n_swa, swa_type); // for tokens of the same sequence, the mask is mostly the same, so we can reuse it // the only cells that could change are the ones that are with similar positions as the @@ -1580,91 +1678,66 @@ static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data auto & idxs = seq_idxs[seq_id]; - if (!alibi) { - if (seq_srct.find(seq_id) != seq_srct.end()) { - const uint32_t srct = seq_srct[seq_id]; - - const uint64_t idst_prev = n_kv*srct; - - std::copy(data + idst_prev, data + idst_prev + n_kv, data + idst); - + if constexpr (!alibi) { + if (seq_srct[seq_id] != UINT32_MAX) { + const uint64_t idst_prev = n_kv * seq_srct[seq_id]; + memcpy(data + idst, data + idst_prev, (size_t) n_kv32 * sizeof(T)); prev = true; } else { - idxs.clear(); idxs.reserve(ubatch->n_tokens + n_swa + 32); seq_srct[seq_id] = i; } } - for (uint32_t jj = 0; jj < n_kv; ++jj) { - uint32_t j = jj; - - // we have an exiting mask for this sequence -> update just seq_idxs - if (!alibi) { - if (prev) { - if (jj >= idxs.size()) { - break; - } + if (prev) { + const llama_pos * cpos = cells.pos_data(); + for (uint32_t jj = 0; jj < (uint32_t) idxs.size(); ++jj) { + const uint32_t j = idxs[jj]; + const llama_pos p0 = cpos[j]; - j = idxs[jj]; + bool keep = (p0 >= wlo && p0 <= whi); + if (keep && is_2d && p0 == p1) { + keep = !cells.ext_get(j).is_2d_gt(p1_x, p1_y); } + data[idst + j] = keep ? mask_keep : mask_drop; } + continue; + } - if (cells.is_empty(j)) { - goto skip; - } - - // mask the token if not the same sequence - if (!cells.seq_has(j, seq_id)) { - goto skip; - } - - p0 = cells.pos_get(j); - - if (!alibi) { - if (!prev) { + if constexpr (alibi) { + for (uint32_t j = 0; j < n_kv32; ++j) { + set_mask_cell_scalar( + cells, j, seq_id, p1, p1_x, p1_y, wlo, whi, + &data[idst + j], mask_drop, mask_keep, swa_type + ); + if (!cells.is_empty(j) && cells.seq_has(j, seq_id)) { + llama_pos p0 = cells.pos_get(j); // record all cells for which: p0 >= seq_pos_min[seq_id] - n_swa - 32 - if (p0 + (int32_t) (n_swa + 32) >= seq_pos_min[seq_id]) { + if (p0 >= seq_pos_min[seq_id] - (int32_t)(n_swa + 32)) { idxs.push_back(j); } } } + continue; + } - if (causal) { - // mask future tokens - if (p0 > p1) { - goto skip; - } + const llama_pos * cpos = cells.pos_data(); + const int32_t rec_lo = seq_pos_min[seq_id] - (int32_t)(n_swa + 32); - // M-RoPE causal mask - if (is_2d) { - if (p0 == p1) { - const auto & p0_ext = cells.ext_get(j); + for (uint32_t j = 0; j < n_kv32; ++j) { + const llama_pos p0 = cpos[j]; + const bool nonempty = (p0 >= 0); + const bool has_seq = cells.seq_has(j, seq_id); - if (p0_ext.is_2d_gt(p1_x, p1_y)) { - goto skip; - } - } - } - } + bool keep = nonempty && has_seq; + if (keep && (causal || swa)) keep = (p0 >= wlo && p0 <= whi); + if (keep && is_2d && p0 == p1) keep = !cells.ext_get(j).is_2d_gt(p1_x, p1_y); + data[idst + j] = keep ? mask_keep : mask_drop; - // apply SWA if any - if (swa) { - if (llama_hparams::is_masked_swa(n_swa, swa_type, p0, p1)) { - goto skip; - } + if (nonempty && has_seq && p0 >= rec_lo) { + idxs.push_back(j); } - - if (alibi) { - data[idst + j] = llama_cast(static_cast(-std::abs(p0 - p1))); - } else { - data[idst + j] = mask_keep; - } - - continue; -skip: - data[idst + j] = mask_drop; } } } From dcadaba454c069f062f64fc642c355a52943e444 Mon Sep 17 00:00:00 2001 From: lexasub Date: Wed, 12 Aug 2026 20:44:12 +0400 Subject: [PATCH 5/9] add cache in unicode_regex_split_stl --- src/unicode.cpp | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/unicode.cpp b/src/unicode.cpp index 810c3fc477e6..74369ac3de3c 100644 --- a/src/unicode.cpp +++ b/src/unicode.cpp @@ -732,19 +732,16 @@ static std::vector unicode_regex_split_custom_qwen35(const std::string & template static std::vector unicode_regex_split_stl(const std::basic_string & text, const std::basic_string & regex, const std::vector & offsets) { using BidirIt = typename std::basic_string::const_iterator; - thread_local std::basic_string cached_regex; - thread_local std::basic_regex cached_expr; // cache compiled regex + thread_local std::unordered_map, std::basic_regex> cached_regex; // cache compiled regex #ifdef _MSC_VER // Bypass bug in MSVC: https://github.com/ggml-org/llama.cpp/issues/17830 constexpr auto regex_flags = std::regex_constants::ECMAScript; #else constexpr auto regex_flags = std::regex_constants::optimize | std::regex_constants::nosubs; #endif - if (cached_regex != regex) { // [[unlikely]] - cached_regex = regex; - cached_expr = std::basic_regex(regex, regex_flags); - } - const auto & expr = cached_expr; + if (cached_regex.size() > 1000) cached_regex.clear(); // basic force limit, without ttl + auto [_it, inserted] = cached_regex.try_emplace(regex, regex, regex_flags); + const auto & expr = _it->second; std::vector bpe_offsets; // store the offset of each word bpe_offsets.reserve(offsets.size()); // Reserve memory for the approximate size size_t start = 0; From 4461db05c8e2fddc351bc0853491669d7fe1ab86 Mon Sep 17 00:00:00 2001 From: lexasub Date: Thu, 13 Aug 2026 04:56:28 +0400 Subject: [PATCH 6/9] add hardware fp16_to_fp32 --- ggml/src/ggml-impl.h | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/ggml/src/ggml-impl.h b/ggml/src/ggml-impl.h index 62b76abbcec9..7189f9539f07 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) { From e0c6a72af30a05cd6d1e0270d7b1dff0d4f435ac Mon Sep 17 00:00:00 2001 From: lexasub Date: Thu, 13 Aug 2026 04:57:06 +0400 Subject: [PATCH 7/9] unroll dequantize_row_q4_K --- ggml/src/ggml-quants.c | 74 +++++++++++++++++++++++++++++++++++------- 1 file changed, 62 insertions(+), 12 deletions(-) diff --git a/ggml/src/ggml-quants.c b/ggml/src/ggml-quants.c index 1ebc50a763f1..e09dd02ebe42 100644 --- a/ggml/src/ggml-quants.c +++ b/ggml/src/ggml-quants.c @@ -1527,26 +1527,76 @@ void quantize_row_q4_K_ref(const float * GGML_RESTRICT x, block_q4_K * GGML_REST } void dequantize_row_q4_K(const block_q4_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { +#define Q4K_ONE() \ + *y1 = d1 * (*q & 0xF) - m1; \ + *y2 = d2 * (*q >> 4) - m2; \ + ++q; ++y1; ++y2; +// manual unroll, because compiler unrolling is different +#define Q4K_FOUR() Q4K_ONE() Q4K_ONE() Q4K_ONE() Q4K_ONE() +#define Q4K_16() Q4K_FOUR() Q4K_FOUR() Q4K_FOUR() Q4K_FOUR() +#define Q4K_32() Q4K_16() Q4K_16() + +#define DEQUANTIZE_Q4_K_BLOCK(scale1_expr, min1_expr, scale2_expr, min2_expr) \ + do { \ + const float d1 = d * (scale1_expr); \ + const float m1 = min * (min1_expr); \ + const float d2 = d * (scale2_expr); \ + const float m2 = min * (min2_expr); \ + Q4K_32() \ + } while (0) + assert(k % QK_K == 0); const int nb = k / QK_K; for (int i = 0; i < nb; i++) { - const uint8_t * q = x[i].qs; + if (i + 1 < nb) { + const char * next = (const char *) &x[i+1]; + __builtin_prefetch(next, 0, 1); + __builtin_prefetch(next + 64, 0, 1); + __builtin_prefetch(next + 128, 0, 1); + } const float d = GGML_FP16_TO_FP32(x[i].d); const float min = GGML_FP16_TO_FP32(x[i].dmin); + const uint8_t * GGML_RESTRICT scales = x[i].scales; - int is = 0; - uint8_t sc, m; - for (int j = 0; j < QK_K; j += 64) { - get_scale_min_k4(is + 0, x[i].scales, &sc, &m); - const float d1 = d * sc; const float m1 = min * m; - get_scale_min_k4(is + 1, x[i].scales, &sc, &m); - const float d2 = d * sc; const float m2 = min * m; - for (int l = 0; l < 32; ++l) *y++ = d1 * (q[l] & 0xF) - m1; - for (int l = 0; l < 32; ++l) *y++ = d2 * (q[l] >> 4) - m2; - q += 32; is += 2; - } + const uint8_t s0 = scales[0], s1 = scales[1], s2 = scales[2], s3 = scales[3]; + const uint8_t s4 = scales[4], s5 = scales[5], s6 = scales[6], s7 = scales[7]; + const uint8_t s8 = scales[8], s9 = scales[9], s10 = scales[10], s11 = scales[11]; + + const uint8_t * GGML_RESTRICT q = x[i].qs; + float * GGML_RESTRICT y1 = y; + float * GGML_RESTRICT y2 = y + 32; + + // --- Block j = 0 (is = 0, 1) --- + DEQUANTIZE_Q4_K_BLOCK( + (s0 & 63), (s4 & 63), + (s1 & 63), (s5 & 63) + ); + y1 += 32; y2 += 32; + + // --- Block j = 64 (is = 2, 3) --- + DEQUANTIZE_Q4_K_BLOCK( + (s2 & 63), (s6 & 63), + (s3 & 63), (s7 & 63) + ); + y1 += 32; y2 += 32; + + // --- Block j = 128 (is = 4, 5) --- + // (s >> 6) << 4 === (s & 0xC0) >> 2 + DEQUANTIZE_Q4_K_BLOCK( + ((s8 & 0xF) | ((s0 & 0xC0) >> 2)), ((s8 >> 4) | ((s4 & 0xC0) >> 2)), + ((s9 & 0xF) | ((s1 & 0xC0) >> 2)), ((s9 >> 4) | ((s5 & 0xC0) >> 2)) + ); + y1 += 32; y2 += 32; + + // --- Block j = 192 (is = 6, 7) --- + DEQUANTIZE_Q4_K_BLOCK( + ((s10 & 0xF) | ((s2 & 0xC0) >> 2)), ((s10 >> 4) | ((s6 & 0xC0) >> 2)), + ((s11 & 0xF) | ((s3 & 0xC0) >> 2)), ((s11 >> 4) | ((s7 & 0xC0) >> 2)) + ); + + y = y2; } } From 79c2cb9b9bc90298471d9faf43d8f7801ab9cef7 Mon Sep 17 00:00:00 2001 From: lexasub Date: Thu, 13 Aug 2026 19:24:05 +0400 Subject: [PATCH 8/9] opt1 --- ggml/src/ggml-cpu/ops.cpp | 91 +++++++++++++++++++++++++++++++++++---- ggml/src/ggml-quants.c | 37 +++++++++++++--- src/llama-kv-cache.cpp | 25 +++++++---- 3 files changed, 131 insertions(+), 22 deletions(-) diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 25bb7438389d..1de86b6d20f2 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -4877,17 +4877,90 @@ static void ggml_compute_forward_get_rows_q( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - for (int64_t i = ir0; i < ir1; ++i) { - const int64_t i12 = i/(ne11*ne10); - const int64_t i11 = (i - i12*ne11*ne10)/ne10; - const int64_t i10 = (i - i12*ne11*ne10 - i11*ne10); - const int64_t i01 = *(int32_t *) ((char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12); + const bool is_1d = (ne12 == 1 && ne11 == 1); + if (is_1d) { + const char * src1_ptr = (const char *) src1->data + ir0 * nb10; + char * dst_ptr = (char *) dst->data + ir0 * nb1; - GGML_ASSERT(i01 >= 0 && i01 < ne01); + const int64_t src1_step = nb10; + const int64_t dst_step = nb1; - dequantize_row_q( - (const void *) ((char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03), - (float *) ((char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3), nc); + int64_t last_i01 = -1; + const float * last_dst_ptr = nullptr; + + for (int64_t i = ir0; i < ir1; ++i) { + const int64_t i01 = *(const int32_t *) src1_ptr; + GGML_ASSERT(i01 >= 0 && i01 < ne01); + + // Cache hit (with prev) + if (i01 == last_i01 && last_dst_ptr != nullptr) { + memcpy(dst_ptr, last_dst_ptr, nc * sizeof(float)); + } else { + // Prefetch next row + if (i + 1 < ir1) { + const int64_t next_i01 = *(const int32_t *)(src1_ptr + src1_step); + const char * next_src0_row = (const char *) src0->data + next_i01 * nb01; + __builtin_prefetch(next_src0_row, 0, 1); + __builtin_prefetch(next_src0_row + 64, 0, 1); + __builtin_prefetch(next_src0_row + 128, 0, 1); + } + + dequantize_row_q( + (const void *) ((const char *) src0->data + i01 * nb01), + (float *) dst_ptr, + nc + ); + + last_i01 = i01; + last_dst_ptr = (const float *) dst_ptr; + } + + src1_ptr += src1_step; + dst_ptr += dst_step; + } + } + else { + int64_t i10 = ir0 % ne10; + int64_t i11 = (ir0 / ne10) % ne11; + int64_t i12 = ir0 / (ne10 * ne11); + + const char * ptr_src1 = (const char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12; + char * ptr_dst = (char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3; + const char * base_src0 = (const char *) src0->data + i11*nb02 + i12*nb03; + + for (int64_t i = ir0; i < ir1; ++i) { + const int64_t i01 = *(const int32_t *) ptr_src1; + GGML_ASSERT(i01 >= 0 && i01 < ne01); + + dequantize_row_q( + (const void *) (base_src0 + i01*nb01), + (float *) ptr_dst, + nc + ); + + ptr_src1 += nb10; + ptr_dst += nb1; + i10++; + + if (__builtin_expect(i10 == ne10, 0)) { + i10 = 0; + + ptr_src1 += nb11 - (ptrdiff_t)ne10 * (ptrdiff_t)nb10; + ptr_dst += nb2 - (ptrdiff_t)ne10 * (ptrdiff_t)nb1; + + i11++; + if (__builtin_expect(i11 == ne11, 0)) { + i11 = 0; + ptr_src1 += nb12 - (ptrdiff_t)ne11 * nb11; + ptr_dst += nb3 - (ptrdiff_t)ne11 * nb2; + + i12++; + base_src0 = (const char *) src0->data + i12*nb03; + } else { + base_src0 += nb02; + } + } + } } } diff --git a/ggml/src/ggml-quants.c b/ggml/src/ggml-quants.c index e09dd02ebe42..f41c81b56cd5 100644 --- a/ggml/src/ggml-quants.c +++ b/ggml/src/ggml-quants.c @@ -13,6 +13,14 @@ #include // for qsort #include // for GGML_ASSERT +#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || defined(_M_IX86) + #include + #define GGML_HAS_X86_INTRINSICS +#elif defined(__ARM_NEON) + #include + #define GGML_HAS_ARM_NEON +#endif + #ifdef GGML_USE_OPENMP #include #endif @@ -1549,12 +1557,12 @@ void dequantize_row_q4_K(const block_q4_K * GGML_RESTRICT x, float * GGML_RESTRI const int nb = k / QK_K; for (int i = 0; i < nb; i++) { - if (i + 1 < nb) { + /*if (i + 1 < nb) { const char * next = (const char *) &x[i+1]; __builtin_prefetch(next, 0, 1); __builtin_prefetch(next + 64, 0, 1); __builtin_prefetch(next + 128, 0, 1); - } + }*/ const float d = GGML_FP16_TO_FP32(x[i].d); const float min = GGML_FP16_TO_FP32(x[i].dmin); @@ -1565,8 +1573,9 @@ void dequantize_row_q4_K(const block_q4_K * GGML_RESTRICT x, float * GGML_RESTRI const uint8_t s8 = scales[8], s9 = scales[9], s10 = scales[10], s11 = scales[11]; const uint8_t * GGML_RESTRICT q = x[i].qs; - float * GGML_RESTRICT y1 = y; - float * GGML_RESTRICT y2 = y + 32; + float buf[256]; + float * GGML_RESTRICT y1 = buf; + float * GGML_RESTRICT y2 = buf + 32; // --- Block j = 0 (is = 0, 1) --- DEQUANTIZE_Q4_K_BLOCK( @@ -1596,7 +1605,25 @@ void dequantize_row_q4_K(const block_q4_K * GGML_RESTRICT x, float * GGML_RESTRI ((s11 & 0xF) | ((s3 & 0xC0) >> 2)), ((s11 >> 4) | ((s7 & 0xC0) >> 2)) ); - y = y2; + #if defined(GGML_HAS_X86_INTRINSICS) || defined(GGML_HAS_ARM_NEON) + if (((uintptr_t)y & 31) == 0) { + #if defined(GGML_HAS_X86_INTRINSICS) && defined(__AVX__) + for (int j = 0; j < 256; j += 8) _mm256_stream_ps(y + j, _mm256_loadu_ps(buf + j)); + #elif defined(GGML_HAS_X86_INTRINSICS) && defined(__SSE2__) + for (int j = 0; j < 256; j += 4) _mm_stream_ps(y + j, _mm_loadu_ps(buf + j)); + #elif defined(GGML_HAS_ARM_NEON) + for (int j = 0; j < 256; j += 4) vst1q_f32(y + j, vld1q_f32(buf + j)); + #else + memcpy(y, buf, 256 * sizeof(float)); + #endif + } else { + memcpy(y, buf, 256 * sizeof(float)); + } + #else + memcpy(y, buf, 256 * sizeof(float)); + #endif + + y += 256; } } diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index ff5da654ee07..1ba9d221d5d1 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1725,17 +1725,26 @@ static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data const llama_pos * cpos = cells.pos_data(); const int32_t rec_lo = seq_pos_min[seq_id] - (int32_t)(n_swa + 32); - for (uint32_t j = 0; j < n_kv32; ++j) { + // Reserve space for idxs + idxs.reserve(idxs.size() + n_kv32 / 4); + T* GGML_RESTRICT out = data + idst; + + for (uint32_t j = 0; j < n_kv32; ++j, ++out) { const llama_pos p0 = cpos[j]; - const bool nonempty = (p0 >= 0); - const bool has_seq = cells.seq_has(j, seq_id); + const bool alive = (p0 >= 0) && cells.seq_has(j, seq_id); + bool visible = alive; + + if constexpr (causal || swa) { + visible = visible && (p0 >= wlo && p0 <= whi); + } + + if constexpr (is_2d) { + if (visible && p0 == p1) visible = !cells.ext_get(j).is_2d_gt(p1_x, p1_y); + } - bool keep = nonempty && has_seq; - if (keep && (causal || swa)) keep = (p0 >= wlo && p0 <= whi); - if (keep && is_2d && p0 == p1) keep = !cells.ext_get(j).is_2d_gt(p1_x, p1_y); - data[idst + j] = keep ? mask_keep : mask_drop; + *out = visible ? mask_keep : mask_drop; - if (nonempty && has_seq && p0 >= rec_lo) { + if (alive && p0 >= rec_lo) { idxs.push_back(j); } } From 76c47519fb2eb954efba7c51f059806006dcfcad Mon Sep 17 00:00:00 2001 From: lexasub Date: Thu, 13 Aug 2026 21:53:33 +0400 Subject: [PATCH 9/9] opt2 --- ggml/src/ggml-cpu/ggml-cpu.c | 15 ++++ ggml/src/ggml-cpu/ops.cpp | 129 +++++++++++++++++++++++++++++++++++ ggml/src/ggml-cpu/ops.h | 1 + 3 files changed, 145 insertions(+) diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 7918845cca08..1a9e8e50843d 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -3053,6 +3053,21 @@ static int ggml_cpu_try_fuse_ops( } } } + if (node->op == GGML_OP_GET_ROWS) { + // GET_ROWS + ADD fusion + const enum ggml_op fuse_ops[] = { GGML_OP_GET_ROWS, GGML_OP_ADD }; + if (ggml_can_fuse(cgraph, node_n, fuse_ops, 2)) { + struct ggml_tensor * add_node = cgraph->nodes[node_n + 1]; + struct ggml_tensor * pos_embd = (add_node->src[0] == node) ? add_node->src[1] : add_node->src[0]; + + if (node->src[1]->type == GGML_TYPE_I32 && + pos_embd->type == GGML_TYPE_F32 && + add_node->type == GGML_TYPE_F32) { + + return ggml_compute_forward_get_rows_add_fused(params, node, add_node); + } + } + } return 0; } diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 1de86b6d20f2..231d0973c988 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -4964,6 +4964,94 @@ static void ggml_compute_forward_get_rows_q( } } +static void ggml_compute_forward_get_rows_add_fused_q( + const ggml_compute_params * params, + const ggml_tensor * dst_get_rows, + ggml_tensor * dst_add) { + + const ggml_tensor * src0 = dst_get_rows->src[0]; + const ggml_tensor * src1 = dst_get_rows->src[1]; + const ggml_tensor * add_src = (dst_add->src[0] == dst_get_rows) + ? dst_add->src[1] + : dst_add->src[0]; + auto dst = dst_add; + GGML_TENSOR_BINARY_OP_LOCALS + + const int64_t nc = ne00; + const int64_t nr = ggml_nelements(src1); + + const ggml_type type = src0->type; + ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float; + + const int ith = params->ith; + const int nth = params->nth; + + const int dr = (nr + nth - 1)/nth; + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); + + GGML_ASSERT(ne0 == nc); + GGML_ASSERT(ne1 == nc); + GGML_ASSERT(nb00 == ggml_type_size(type)); + GGML_ASSERT(add_src->type == GGML_TYPE_F32); + + const int64_t add_nelem = ggml_nelements(add_src); + const bool is_broadcast_single = (add_nelem == nc); + const bool is_positional = (add_nelem >= nr * nc); + + const char * src1_ptr = (const char *) src1->data + ir0 * nb10; + char * dst_ptr = (char *) dst_add->data + ir0 * nb1; + const int64_t src1_step = nb10; + const int64_t dst_step = nb1; + + int64_t last_i01 = -1; + int64_t last_add_i = -1; + const float * last_dst_ptr = nullptr; + + for (int64_t i = ir0; i < ir1; ++i) { + const int64_t i01 = *(const int32_t *) src1_ptr; + GGML_ASSERT(i01 >= 0 && i01 < ne01); + + int64_t add_i{0}; + if (!is_broadcast_single) { + if (is_positional) { + add_i = i; + } else { + add_i = i % (add_nelem / nc); + } + } + + const float * add_row = (const float *)((const char *)add_src->data + add_i * nb11); + + if (i01 == last_i01 && add_i == last_add_i && last_dst_ptr != nullptr) { + memcpy(dst_ptr, last_dst_ptr, nc * sizeof(float)); + } else { + if (i + 1 < ir1) { + const int64_t next_i01 = *(const int32_t *)(src1_ptr + src1_step); + const char * next_src0_row = (const char *) src0->data + next_i01 * nb01; + __builtin_prefetch(next_src0_row, 0, 1); + __builtin_prefetch(next_src0_row + 64, 0, 1); + __builtin_prefetch(next_src0_row + 128, 0, 1); + } + + dequantize_row_q( + (const void *) ((const char *) src0->data + i01 * nb01), + (float *) dst_ptr, + nc + ); + + ggml_vec_acc_f32(nc, (float *) dst_ptr, add_row); + + last_i01 = i01; + last_add_i = add_i; + last_dst_ptr = (const float *) dst_ptr; + } + + src1_ptr += src1_step; + dst_ptr += dst_step; + } +} + static void ggml_compute_forward_get_rows_f16( const ggml_compute_params * params, ggml_tensor * dst) { @@ -5161,6 +5249,47 @@ void ggml_compute_forward_get_rows( //} } +int ggml_compute_forward_get_rows_add_fused(const struct ggml_compute_params * params, struct ggml_tensor * dst_get_rows, struct ggml_tensor * dst_add) { + const ggml_tensor * src0 = dst_get_rows->src[0]; + + switch (src0->type) { + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q8_1: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_TQ1_0: + case GGML_TYPE_TQ2_0: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ2_S: + { + ggml_compute_forward_get_rows_add_fused_q(params, dst_get_rows, dst_add); + } break; + default: + { + return 0; + } + } + return 1; +} + template static void ggml_compute_forward_set_rows_impl( const ggml_compute_params * params, diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 4c1642a67603..7245d61bded7 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -54,6 +54,7 @@ void ggml_compute_forward_set(const struct ggml_compute_params * params, struct void ggml_compute_forward_cpy(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_cont(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_get_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst); +int ggml_compute_forward_get_rows_add_fused(const struct ggml_compute_params * params, struct ggml_tensor * dst_get_rows, struct ggml_tensor * dst_add); void ggml_compute_forward_get_rows_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_set_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_diag(const struct ggml_compute_params * params, struct ggml_tensor * dst);