common : store the ngram caches in unordered_dense maps

common/CMakeLists.txt

@@ -134,7 +134,7 @@ set_target_properties(${TARGET} PROPERTIES
 )
 
 target_include_directories(${TARGET} PUBLIC .)
-target_link_libraries     (${TARGET} PUBLIC vendor::nlohmann vendor::sheredom)
+target_link_libraries     (${TARGET} PUBLIC vendor::nlohmann vendor::sheredom vendor::unordered_dense)
 target_compile_features   (${TARGET} PUBLIC cxx_std_17)
 
 if (LLAMA_SUBPROCESS)

common/ngram-cache.h

@@ -2,8 +2,11 @@
 
 #include "llama.h"
 
-#include <unordered_map>
+#include <ankerl/unordered_dense.h>
+
+#include <algorithm>
 #include <string>
+#include <utility>
 #include <vector>
 
 #define LLAMA_NGRAM_MIN    1
@@ -55,10 +58,52 @@ struct common_ngram_hash_function {
 };
 
 // token -> number of times token has been seen
-typedef std::unordered_map<llama_token, int32_t> common_ngram_cache_part;
+// Most n-grams are followed by only a few distinct tokens, so the pairs live in one vector sorted by token
+// instead of a hash map of their own. Lookups use a binary search, which stays fast for the n-grams that
+// are followed by thousands of distinct tokens.
+struct common_ngram_cache_part {
+    typedef std::pair<llama_token, int32_t>          value_type;
+    typedef std::vector<value_type>::iterator       iterator;
+    typedef std::vector<value_type>::const_iterator const_iterator;
+
+    std::vector<value_type> entries;
+
+    iterator find(const llama_token token) {
+        const iterator it = lower_bound(token);
+        return it != entries.end() && it->first == token ? it : entries.end();
+    }
+
+    const_iterator find(const llama_token token) const {
+        const const_iterator it = lower_bound(token);
+        return it != entries.end() && it->first == token ? it : entries.end();
+    }
+
+    // Callers emplace a token only after find did not return it.
+    void emplace(const llama_token token, const int32_t count) {
+        entries.insert(lower_bound(token), value_type(token, count));
+    }
+
+    iterator       begin()       { return entries.begin(); }
+    iterator       end()         { return entries.end(); }
+    const_iterator begin() const { return entries.begin(); }
+    const_iterator end()   const { return entries.end(); }
+    size_t         size()  const { return entries.size(); }
+    bool           empty() const { return entries.empty(); }
+
+private:
+    iterator lower_bound(const llama_token token) {
+        return std::lower_bound(entries.begin(), entries.end(), token,
+            [](const value_type & entry, const llama_token t) { return entry.first < t; });
+    }
+
+    const_iterator lower_bound(const llama_token token) const {
+        return std::lower_bound(entries.begin(), entries.end(), token,
+            [](const value_type & entry, const llama_token t) { return entry.first < t; });
+    }
+};
 
 // n-gram -> empirical distribution of following tokens
-typedef std::unordered_map<common_ngram, common_ngram_cache_part, common_ngram_hash_function> common_ngram_cache;
+typedef ankerl::unordered_dense::map<common_ngram, common_ngram_cache_part, common_ngram_hash_function> common_ngram_cache;
 
 
 // Update an ngram cache with tokens.

scripts/sync_vendor.py

@@ -6,6 +6,7 @@
 import subprocess
 
 HTTPLIB_VERSION = "refs/tags/v0.57.1"
+UNORDERED_DENSE_VERSION = "refs/tags/v5.0.1"
 
 # used by examples/gguf-hash, these repos have no release tag, so we pin a commit
 XXHASH_COMMIT      = "9f465f1ea932d6ad9a26cd77496311ffa544cd68"
@@ -27,6 +28,10 @@
     f"https://raw.githubusercontent.com/yhirose/cpp-httplib/{HTTPLIB_VERSION}/split.py":  "split.py",
     f"https://raw.githubusercontent.com/yhirose/cpp-httplib/{HTTPLIB_VERSION}/LICENSE":   "vendor/cpp-httplib/LICENSE",
 
+    f"https://raw.githubusercontent.com/martinus/unordered_dense/{UNORDERED_DENSE_VERSION}/include/ankerl/unordered_dense.h": "vendor/ankerl/unordered_dense.h",
+    f"https://raw.githubusercontent.com/martinus/unordered_dense/{UNORDERED_DENSE_VERSION}/include/ankerl/stl.h":             "vendor/ankerl/stl.h",
+    f"https://raw.githubusercontent.com/martinus/unordered_dense/{UNORDERED_DENSE_VERSION}/LICENSE":                           "vendor/ankerl/LICENSE",
+
     "https://raw.githubusercontent.com/sheredom/subprocess.h/0dccaa9aa176dd6d7ef8afeca3c18d6e80a32795/subprocess.h": "vendor/sheredom/subprocess.h",
 
     f"https://raw.githubusercontent.com/Cyan4973/xxHash/{XXHASH_COMMIT}/xxhash.c":      "vendor/hash/xxhash/xxhash.c",

vendor/CMakeLists.txt

@@ -7,5 +7,6 @@ add_subdirectory(stb)
 
 # only used by common
 if (LLAMA_BUILD_COMMON)
+    add_subdirectory(ankerl)
     add_subdirectory(cpp-httplib)
 endif()

vendor/ankerl/CMakeLists.txt

@@ -0,0 +1,6 @@
+# header-only: interface target exposing the vendor/ root so consumers
+# can include via <ankerl/unordered_dense.h>
+add_library(unordered_dense INTERFACE)
+add_library(vendor::unordered_dense ALIAS unordered_dense)
+
+target_include_directories(unordered_dense INTERFACE ..)

vendor/ankerl/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Martin Leitner-Ankerl
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

vendor/ankerl/stl.h

@@ -0,0 +1,85 @@
+///////////////////////// ankerl::unordered_dense::{map, set} /////////////////////////
+
+// A fast & densely stored hashmap and hashset based on robin-hood backward shift deletion.
+// Version 5.0.1
+// https://github.com/martinus/unordered_dense
+//
+// Licensed under the MIT License <http://opensource.org/licenses/MIT>.
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2022 Martin Leitner-Ankerl <martin.ankerl@gmail.com>
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+#ifndef ANKERL_STL_H
+#define ANKERL_STL_H
+
+#include <array>            // for array
+#include <cstdint>          // for uint64_t, uint32_t, std::uint8_t, UINT64_C
+#include <cstring>          // for size_t, memcpy, memset
+#include <functional>       // for equal_to, hash
+#include <initializer_list> // for initializer_list
+#include <iterator>         // for pair, distance
+#include <limits>           // for numeric_limits
+#include <memory>           // for allocator, allocator_traits, shared_ptr
+#include <optional>         // for optional
+#include <stdexcept>        // for out_of_range
+#include <string>           // for basic_string
+#include <string_view>      // for basic_string_view, hash
+#include <tuple>            // for forward_as_tuple
+#include <type_traits>      // for enable_if_t, declval, conditional_t, ena...
+#include <utility>          // for forward, exchange, pair, as_const, piece...
+#include <vector>           // for vector
+
+// <memory_resource> includes <mutex>, which fails to compile if
+// targeting GCC >= 13 with the (rewritten) win32 thread model, and
+// targeting Windows earlier than Vista (0x600).  GCC predefines
+// _REENTRANT when using the 'posix' model, and doesn't when using the
+// 'win32' model.
+#if defined __MINGW64__ && defined __GNUC__ && __GNUC__ >= 13 && !defined _REENTRANT
+// _WIN32_WINNT is guaranteed to be defined here because of the
+// <cstdint> inclusion above.
+#    ifndef _WIN32_WINNT
+#        error "_WIN32_WINNT not defined"
+#    endif
+#    if _WIN32_WINNT < 0x600
+#        define ANKERL_MEMORY_RESOURCE_IS_BAD() 1 // NOLINT(cppcoreguidelines-macro-usage)
+#    endif
+#endif
+#ifndef ANKERL_MEMORY_RESOURCE_IS_BAD
+#    define ANKERL_MEMORY_RESOURCE_IS_BAD() 0 // NOLINT(cppcoreguidelines-macro-usage)
+#endif
+
+#if defined(__has_include) && !defined(ANKERL_UNORDERED_DENSE_DISABLE_PMR)
+#    if __has_include(<memory_resource>) && !ANKERL_MEMORY_RESOURCE_IS_BAD()
+#        define ANKERL_UNORDERED_DENSE_PMR std::pmr // NOLINT(cppcoreguidelines-macro-usage)
+#        include <memory_resource>                  // for polymorphic_allocator
+#    elif __has_include(<experimental/memory_resource>)
+#        define ANKERL_UNORDERED_DENSE_PMR std::experimental::pmr // NOLINT(cppcoreguidelines-macro-usage)
+#        include <experimental/memory_resource>                   // for polymorphic_allocator
+#    endif
+#endif
+
+#if defined(_MSC_VER) && defined(_M_X64)
+#    include <intrin.h>
+#    if !defined(_M_ARM64EC)
+#        pragma intrinsic(_umul128)
+#    endif
+#endif
+
+#endif