From dbff279dcc4f5cb1e4fa58e1405da933b3132563 Mon Sep 17 00:00:00 2001 From: Reithan Date: Sat, 21 Jun 2025 18:42:22 -0700 Subject: [PATCH 1/5] Add memoized cache to llama_grammar_reject_candidates_for_stack --- src/llama-grammar.cpp | 53 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/src/llama-grammar.cpp b/src/llama-grammar.cpp index bed706bb248d..799b47804720 100644 --- a/src/llama-grammar.cpp +++ b/src/llama-grammar.cpp @@ -8,10 +8,26 @@ #include #include +#include +#include +#include + // // helpers // +using bytes = std::pair; +using hash_entry_size = std::pair; + +template <> +struct std::hash +{ + std::size_t operator()(const bytes& x) const noexcept + { + return std::hash{}({x.first, x.second}); + } +}; + // NOTE: assumes valid utf8 (but checks for overrun) static std::pair decode_utf8(const char * src) { static const int lookup[] = { 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 4 }; @@ -864,6 +880,40 @@ llama_grammar_candidates llama_grammar_reject_candidates_for_stack( } return rejects; } + + + using candidates_memos = std::unordered_map; + using stack_memos = std::unordered_map; + static stack_memos memo_cache; + + auto stack_hash_start = reinterpret_cast(stack.data()); + auto stack_hash_size = sizeof(stack[0]) * stack.size(); + auto stack_hash = std::hash{}({ stack_hash_start, stack_hash_size }); + + llama_grammar_candidates* cache_target = nullptr; + + // Only check stash hash first - these are usually ~24b, and almost always under 64b + if (auto& cache_hit = memo_cache.find(stack_hash); cache_hit != memo_cache.end()) { + // Only attempt to memoize candidate lists under 1280 + // >75% of candidate lists are under 1280 and 50% are under 640b. Most 'problem' loops are under 24b + // However, candidate lists can be over 72k, so we need to limit our checks + auto candidates_hash_size = sizeof(candidates[0]) * candidates.size(); + if (candidates_hash_size < 1280) { + auto& candidates_memos = cache_hit->second; + auto candidates_hash_start = reinterpret_cast(candidates.data()); + auto candidates_hash = std::hash{}({ candidates_hash_start, candidates_hash_size }); + if (auto& cache_hit2 = candidates_memos.find(candidates_hash); cache_hit2 != candidates_memos.end()) { + return cache_hit2->second; + } + else { + cache_target = &(candidates_memos[candidates_hash]); + } + } + } + else { + memo_cache[stack_hash]; + } + const llama_grammar_element * stack_pos = stack.back(); @@ -900,6 +950,9 @@ llama_grammar_candidates llama_grammar_reject_candidates_for_stack( rejects.push_back({ tok.index, tok.code_points - 1, tok.partial_utf8 }); } + if (cache_target) { + *cache_target = rejects; + } return rejects; } From 4510328596b53b170436f24b262c6dd5788882dc Mon Sep 17 00:00:00 2001 From: Reithan Date: Sun, 22 Jun 2025 19:31:31 -0700 Subject: [PATCH 2/5] make size cutoff more aggressive and move to outer branch --- src/llama-grammar.cpp | 39 ++++++++++++++++++++------------------- 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/src/llama-grammar.cpp b/src/llama-grammar.cpp index 799b47804720..4ab7653a3f4a 100644 --- a/src/llama-grammar.cpp +++ b/src/llama-grammar.cpp @@ -890,30 +890,31 @@ llama_grammar_candidates llama_grammar_reject_candidates_for_stack( auto stack_hash_size = sizeof(stack[0]) * stack.size(); auto stack_hash = std::hash{}({ stack_hash_start, stack_hash_size }); - llama_grammar_candidates* cache_target = nullptr; - - // Only check stash hash first - these are usually ~24b, and almost always under 64b - if (auto& cache_hit = memo_cache.find(stack_hash); cache_hit != memo_cache.end()) { - // Only attempt to memoize candidate lists under 1280 - // >75% of candidate lists are under 1280 and 50% are under 640b. Most 'problem' loops are under 24b - // However, candidate lists can be over 72k, so we need to limit our checks - auto candidates_hash_size = sizeof(candidates[0]) * candidates.size(); - if (candidates_hash_size < 1280) { - auto& candidates_memos = cache_hit->second; - auto candidates_hash_start = reinterpret_cast(candidates.data()); - auto candidates_hash = std::hash{}({ candidates_hash_start, candidates_hash_size }); - if (auto& cache_hit2 = candidates_memos.find(candidates_hash); cache_hit2 != candidates_memos.end()) { + llama_grammar_candidates * cache_target = nullptr; + + // Only attempt to memoize candidate lists under 1280 + // >75% of candidate lists are under 1280 and 50% are under 640b. Most 'problem' loops are under 24b + // However, candidate lists can be over 72k, so we need to limit our checks + // Doing an over-aggressive size cutoff first befor any other processing 'saves' easy cases + // extra processing but still rescues 'hard' cases from slow down or hangs. + // This leads to a speed up of both easy and hard cases. + const size_t hash_cutoff = 80; + auto candidates_hash_size = sizeof(candidates[0]) * candidates.size(); + if (candidates_hash_size < hash_cutoff) { + // Only check stash hash first - these are usually ~24b, and almost always under 64b + if (auto & cache_hit = memo_cache.find(stack_hash); cache_hit != memo_cache.end()) { + auto & candidates_memos = cache_hit->second; + auto candidates_hash_start = reinterpret_cast(candidates.data()); + auto candidates_hash = std::hash{}({ candidates_hash_start, candidates_hash_size }); + if (auto & cache_hit2 = candidates_memos.find(candidates_hash); cache_hit2 != candidates_memos.end()) { return cache_hit2->second; - } - else { + } else { cache_target = &(candidates_memos[candidates_hash]); } + } else { + memo_cache[stack_hash]; } } - else { - memo_cache[stack_hash]; - } - const llama_grammar_element * stack_pos = stack.back(); From 241414e7482aae31e6095ba4ddcb7fc255aab4c3 Mon Sep 17 00:00:00 2001 From: Reithan Date: Sun, 22 Jun 2025 19:33:01 -0700 Subject: [PATCH 3/5] update comment --- src/llama-grammar.cpp | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/llama-grammar.cpp b/src/llama-grammar.cpp index 4ab7653a3f4a..39410e077d5e 100644 --- a/src/llama-grammar.cpp +++ b/src/llama-grammar.cpp @@ -892,9 +892,12 @@ llama_grammar_candidates llama_grammar_reject_candidates_for_stack( llama_grammar_candidates * cache_target = nullptr; - // Only attempt to memoize candidate lists under 1280 - // >75% of candidate lists are under 1280 and 50% are under 640b. Most 'problem' loops are under 24b - // However, candidate lists can be over 72k, so we need to limit our checks + // Tests show that >75% of candidate lists are under 1280 and 50% are under 640b. + // Most 'problem' loops are under 24b. However, candidate lists can be over 72k, + // so we need to limit our checks. + + + // We'll only attempt to memoize candidate lists under 80b // Doing an over-aggressive size cutoff first befor any other processing 'saves' easy cases // extra processing but still rescues 'hard' cases from slow down or hangs. // This leads to a speed up of both easy and hard cases. From 0fc4a88452be11418010a8f06ee1ee6e7f6d60a9 Mon Sep 17 00:00:00 2001 From: Reithan Date: Mon, 23 Jun 2025 01:58:09 -0700 Subject: [PATCH 4/5] add cache reset whenever grammar is reloaded --- gpttype_adapter.cpp | 1 + src/llama-grammar.cpp | 13 ++++++++----- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/gpttype_adapter.cpp b/gpttype_adapter.cpp index ad9d7f92255c..26cf048183e1 100644 --- a/gpttype_adapter.cpp +++ b/gpttype_adapter.cpp @@ -1773,6 +1773,7 @@ static void load_grammar(const std::string & gammarstr) { if(grammar!=nullptr) //on demand free when next grammar is loaded { + llama_grammar_reset_memos(); llama_grammar_free_impl(grammar); grammar = nullptr; } diff --git a/src/llama-grammar.cpp b/src/llama-grammar.cpp index 39410e077d5e..cdf6f9921820 100644 --- a/src/llama-grammar.cpp +++ b/src/llama-grammar.cpp @@ -28,6 +28,14 @@ struct std::hash } }; +using candidates_memos = std::unordered_map; +using stack_memos = std::unordered_map; +static stack_memos memo_cache; + +static void llama_grammar_reset_memos() { + memo_cache.clear(); +} + // NOTE: assumes valid utf8 (but checks for overrun) static std::pair decode_utf8(const char * src) { static const int lookup[] = { 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 4 }; @@ -881,11 +889,6 @@ llama_grammar_candidates llama_grammar_reject_candidates_for_stack( return rejects; } - - using candidates_memos = std::unordered_map; - using stack_memos = std::unordered_map; - static stack_memos memo_cache; - auto stack_hash_start = reinterpret_cast(stack.data()); auto stack_hash_size = sizeof(stack[0]) * stack.size(); auto stack_hash = std::hash{}({ stack_hash_start, stack_hash_size }); From 8ade80329423f7ae44165faf57633ddec0134ef5 Mon Sep 17 00:00:00 2001 From: Reithan Date: Tue, 24 Jun 2025 20:26:09 -0700 Subject: [PATCH 5/5] remove explicit reference types for compiler transportability --- src/llama-grammar.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/llama-grammar.cpp b/src/llama-grammar.cpp index cdf6f9921820..ead5d3f3a0fb 100644 --- a/src/llama-grammar.cpp +++ b/src/llama-grammar.cpp @@ -899,7 +899,6 @@ llama_grammar_candidates llama_grammar_reject_candidates_for_stack( // Most 'problem' loops are under 24b. However, candidate lists can be over 72k, // so we need to limit our checks. - // We'll only attempt to memoize candidate lists under 80b // Doing an over-aggressive size cutoff first befor any other processing 'saves' easy cases // extra processing but still rescues 'hard' cases from slow down or hangs. @@ -908,11 +907,11 @@ llama_grammar_candidates llama_grammar_reject_candidates_for_stack( auto candidates_hash_size = sizeof(candidates[0]) * candidates.size(); if (candidates_hash_size < hash_cutoff) { // Only check stash hash first - these are usually ~24b, and almost always under 64b - if (auto & cache_hit = memo_cache.find(stack_hash); cache_hit != memo_cache.end()) { + if (auto cache_hit = memo_cache.find(stack_hash); cache_hit != memo_cache.end()) { auto & candidates_memos = cache_hit->second; auto candidates_hash_start = reinterpret_cast(candidates.data()); auto candidates_hash = std::hash{}({ candidates_hash_start, candidates_hash_size }); - if (auto & cache_hit2 = candidates_memos.find(candidates_hash); cache_hit2 != candidates_memos.end()) { + if (auto cache_hit2 = candidates_memos.find(candidates_hash); cache_hit2 != candidates_memos.end()) { return cache_hit2->second; } else { cache_target = &(candidates_memos[candidates_hash]);