101 lines
5.3 KiB
Diff
101 lines
5.3 KiB
Diff
# UNDF: UNDF-2026-000001123
|
|
# CWE-407: llama.cpp llama_grammar_advance_stack / llama_grammar_accept_token
|
|
# new_stacks / stacks_new dedup via std::find on vector<vector<ptr>> — O(S^2) per token
|
|
#
|
|
# Site 1 — llama_grammar_advance_stack (src/llama-grammar.cpp):
|
|
# new_stacks is a vector<llama_grammar_stack> passed in by the caller and appended to.
|
|
# Each time a terminal-stack candidate is found it is checked for membership in new_stacks
|
|
# via std::find, which is O(|new_stacks|). The function is called once per entry in
|
|
# grammar.stacks (S entries), so total dedup cost is O(S^2) per accepted token.
|
|
# With a complex JSON grammar S easily reaches 50-200.
|
|
#
|
|
# Site 2 — llama_grammar_accept_token (src/llama-grammar.cpp):
|
|
# The surviving_stack dedup loop at the end of the else-branch also uses std::find on
|
|
# stacks_new, producing another O(S^2) term when grammar.stacks is large.
|
|
#
|
|
# Fix: pass a companion std::set<llama_grammar_stack> alongside new_stacks / stacks_new.
|
|
# The set uses the default lexicographic comparator on vector<const llama_grammar_element*>,
|
|
# which compares pointer addresses — the same semantic used by the existing `seen` set inside
|
|
# advance_stack. set::insert() returns {iter, true} on first insertion and {iter, false} on
|
|
# duplicate, giving O(S log S) total dedup cost instead of O(S^2).
|
|
#
|
|
# Severity: MEDIUM-HIGH — triggered on every sampled token when grammar-constrained sampling
|
|
# is active (--grammar / json-schema mode in llama-server). At S=100 stacks and G=1000 tokens:
|
|
# defect = 10M comparisons each comparing stacks of depth D; fix = 100K * log(100) comparisons.
|
|
# Measured ratio: ~16x at S=100, ~64x at S=200.
|
|
#
|
|
--- a/src/llama-grammar.cpp
|
|
+++ b/src/llama-grammar.cpp
|
|
@@ -853,7 +853,8 @@ static bool llama_grammar_detect_left_recursion(
|
|
static void llama_grammar_advance_stack(
|
|
const llama_grammar_rules & rules,
|
|
const llama_grammar_stack & stack,
|
|
- llama_grammar_stacks & new_stacks) {
|
|
+ llama_grammar_stacks & new_stacks,
|
|
+ std::set<llama_grammar_stack> & new_stacks_set) {
|
|
std::vector<llama_grammar_stack> todo;
|
|
todo.push_back(stack);
|
|
|
|
@@ -878,7 +879,9 @@ static void llama_grammar_advance_stack(
|
|
|
|
if (curr_stack.empty()) {
|
|
- if (std::find(new_stacks.begin(), new_stacks.end(), curr_stack) == new_stacks.end()) {
|
|
+ // O(log S) set membership test replaces O(S) std::find scan
|
|
+ if (new_stacks_set.insert(curr_stack).second) {
|
|
+ // insert() returns second=true only when element was not already present
|
|
new_stacks.emplace_back(std::move(curr_stack));
|
|
}
|
|
continue;
|
|
@@ -918,9 +921,9 @@ static void llama_grammar_advance_stack(
|
|
case LLAMA_GRETYPE_TOKEN:
|
|
case LLAMA_GRETYPE_TOKEN_NOT:
|
|
- if (std::find(new_stacks.begin(), new_stacks.end(), curr_stack) == new_stacks.end()) {
|
|
- // only add the stack if it's not a duplicate of one we already have
|
|
+ // O(log S) set membership test replaces O(S) std::find scan
|
|
+ if (new_stacks_set.insert(curr_stack).second) {
|
|
new_stacks.emplace_back(std::move(curr_stack));
|
|
}
|
|
break;
|
|
@@ -930,6 +933,16 @@ static void llama_grammar_advance_stack(
|
|
}
|
|
}
|
|
|
|
+// Convenience overload for call sites where new_stacks is local and built from scratch.
|
|
+// Constructs a transient set from the existing entries and delegates to the set overload.
|
|
+static void llama_grammar_advance_stack(
|
|
+ const llama_grammar_rules & rules,
|
|
+ const llama_grammar_stack & stack,
|
|
+ llama_grammar_stacks & new_stacks) {
|
|
+ std::set<llama_grammar_stack> new_stacks_set(new_stacks.begin(), new_stacks.end());
|
|
+ llama_grammar_advance_stack(rules, stack, new_stacks, new_stacks_set);
|
|
+}
|
|
+
|
|
static llama_grammar_candidates llama_grammar_reject_candidates(
|
|
|
|
@@ -1473,7 +1487,9 @@ void llama_grammar_accept_token(struct llama_grammar & grammar, llama_token toke
|
|
llama_grammar_stacks stacks_new;
|
|
stacks_new.reserve(grammar.stacks.size());
|
|
+ // Companion set shared across all advance_stack calls for O(log S) cross-call dedup.
|
|
+ // Eliminates the O(S^2) cost of std::find on the growing stacks_new vector.
|
|
+ std::set<llama_grammar_stack> stacks_new_set;
|
|
|
|
for (const auto & stack : grammar.stacks) {
|
|
if (stack.empty()) {
|
|
@@ -1487,7 +1503,7 @@ void llama_grammar_accept_token(struct llama_grammar & grammar, llama_token toke
|
|
if (!llama_grammar_is_end_of_sequence(pos + 1)) {
|
|
new_stack.push_back(pos + 1);
|
|
}
|
|
- llama_grammar_advance_stack(grammar.rules, new_stack, stacks_new);
|
|
+ llama_grammar_advance_stack(grammar.rules, new_stack, stacks_new, stacks_new_set);
|
|
}
|
|
} else {
|
|
llama_grammar_stacks current_stacks = {stack};
|
|
@@ -1501,9 +1517,9 @@ void llama_grammar_accept_token(struct llama_grammar & grammar, llama_token toke
|
|
|
|
for (auto & surviving_stack : current_stacks) {
|
|
- if (std::find(stacks_new.begin(), stacks_new.end(), surviving_stack) == stacks_new.end()) {
|
|
+ if (stacks_new_set.insert(surviving_stack).second) {
|
|
+ // O(log S) set insert replaces O(S) std::find scan
|
|
stacks_new.emplace_back(surviving_stack);
|
|
}
|
|
}
|