fix two bugs in grammar memoization for reject_candidates_for_stack (#2431)

TOKEN/TOKEN_NOT early return: operator[] pre-inserts an empty candidates
vector at line 1123, setting cache_target. The TOKEN/TOKEN_NOT branch
returns without writing to *cache_target, leaving an empty vector in the
map. Subsequent lookups return that empty vector, reporting "reject
nothing" for what should be a non-empty reject set. Fix: write to
cache_target before the early return.

Pointer aliasing in candidate key: candidate lists were hashed as raw
struct bytes, which includes the code_points pointer value. glibc can
reuse heap addresses across calls, so the same pointer value can point
to different content, causing false cache hits. Fix: hash actual content
(index, id, partial_utf8 fields, and the decoded code_points sequence)
rather than raw bytes including the pointer.
This commit is contained in:
Reithan
2026-09-04 01:03:42 -07:00
committed by GitHub
parent 4d5779d1bb
commit 6637d25971
+23 -3
View File
@@ -1114,9 +1114,24 @@ llama_grammar_candidates llama_grammar_reject_candidates_for_stack(
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<const char *>(candidates.data());
auto candidates_hash = std::hash<bytes>{}({ candidates_hash_start, candidates_hash_size });
auto & candidates_memos = cache_hit->second;
// Hash candidate content (index, id, partial_utf8, decoded code_points sequence)
// rather than raw struct bytes, which include the code_points pointer value.
// Pointer values can be reused by the allocator across calls pointing to different
// content, causing false cache hits that return stale reject sets.
size_t candidates_hash = candidates.size();
auto combine = [&](size_t v) {
candidates_hash ^= v + 0x9e3779b9 + (candidates_hash << 6) + (candidates_hash >> 2);
};
for (const auto & c : candidates) {
combine(std::hash<size_t>{}(c.index));
combine(std::hash<llama_token>{}(c.id));
combine(std::hash<uint32_t>{}(c.partial_utf8.value));
combine(std::hash<int>{}(c.partial_utf8.n_remain));
for (const uint32_t * cp = c.code_points; *cp != 0; ++cp) {
combine(std::hash<uint32_t>{}(*cp));
}
}
if (auto cache_hit2 = candidates_memos.find(candidates_hash); cache_hit2 != candidates_memos.end()) {
return cache_hit2->second;
} else {
@@ -1142,6 +1157,11 @@ llama_grammar_candidates llama_grammar_reject_candidates_for_stack(
rejects.push_back(tok);
}
}
// cache_target was set by operator[] which pre-inserted an empty vector; write the
// result here so subsequent lookups don't return an empty reject set.
if (cache_target) {
*cache_target = rejects;
}
return rejects;
}